[Model] Remove MiniMaxText01, MiniMaxVL01, MiniMaxForCausalLM (#45993)
Signed-off-by: Xianbao QIAN <xianbao.qian@gmail.com> Co-authored-by: Claude <noreply@anthropic.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
Tiezhen WANG committed
Jun 22, 2026 at 15:20 UTC
d14e551a5326c75f3213aa35ae4dbbed7a04ea02
20 files changed
+10
-3881
docs/contributing/model/basic.md
+2
-2
@@ -136,7 +136,7 @@ The model should also be added to the `MODELS_CONFIG_MAP` dictionary in [vllm/mo
136
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).
137
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).
138
139
-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.
139
+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`.
140
Please follow the same guidelines as case (2) for implementing these models.
141
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).
142
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
144
Please see [`LinearAttentionMetadata`](../../../vllm/v1/attention/backends/linear_attn.py) or [`ShortConvAttentionMetadata`](../../../vllm/v1/attention/backends/short_conv_attn.py) for examples of this.
145
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.
146
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.
147
-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.
147
+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.
148
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.
docs/features/tool_calling.md
-9
@@ -321,15 +321,6 @@ For Qwen2.5, the chat template in tokenizer_config.json has already included sup
321
322
Flags: `--tool-call-parser hermes`
323
324
-### MiniMax Models (`minimax_m1`)
325
-
326
-Supported models:
327
-
328
-* `MiniMaxAi/MiniMax-M1-40k` (use with [examples/tool_chat_template_minimax_m1.jinja](../../examples/tool_chat_template_minimax_m1.jinja))
329
-* `MiniMaxAi/MiniMax-M1-80k` (use with [examples/tool_chat_template_minimax_m1.jinja](../../examples/tool_chat_template_minimax_m1.jinja))
330
-
331
-Flags: `--tool-call-parser minimax --chat-template examples/tool_chat_template_minimax_m1.jinja`
332
-
324
### DeepSeek-V3 Models (`deepseek_v3`)
325
326
Supported models:
docs/models/supported_models.md
-4
@@ -441,7 +441,6 @@ th {
441
| `MiMoV2ForCausalLM` | MiMoV2Pro | `XiaomiMiMo/MiMo-V2.5-Pro`, etc. | | ✅︎ |
442
| `MiniCPMForCausalLM` | MiniCPM | `openbmb/MiniCPM-2B-sft-bf16`, `openbmb/MiniCPM-2B-dpo-bf16`, `openbmb/MiniCPM-S-1B-sft`, etc. | ✅︎ | ✅︎ |
443
| `MiniCPM3ForCausalLM` | MiniCPM3 | `openbmb/MiniCPM3-4B`, etc. | ✅︎ | ✅︎ |
444
-| `MiniMaxForCausalLM` | MiniMax-Text | `MiniMaxAI/MiniMax-Text-01-hf`, etc. | | |
444
| `MiniMaxM2ForCausalLM` | MiniMax-M2, MiniMax-M2.1 | `MiniMaxAI/MiniMax-M2`, etc. | ✅︎ | ✅︎ |
445
| `MistralForCausalLM` | Ministral-3, Mistral, Mistral-Instruct | `mistralai/Ministral-3-3B-Instruct-2512`, `mistralai/Mistral-7B-v0.1`, `mistralai/Mistral-7B-Instruct-v0.1`, etc. | ✅︎ | ✅︎ |
446
| `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 {
486
| `TeleChat2ForCausalLM` | TeleChat2 | `Tele-AI/TeleChat2-3B`, `Tele-AI/TeleChat2-7B`, `Tele-AI/TeleChat2-35B`, etc. | ✅︎ | ✅︎ |
487
| `TeleChat3ForCausalLM` | TeleChat3 | `Tele-AI/TeleChat3-36B-Thinking`, `Tele-AI/TeleChat3-Coder-36B-Thinking`, etc. | ✅︎ | ✅︎ |
488
| `TeleFLMForCausalLM` | TeleFLM | `CofeAI/FLM-2-52B-Instruct-2407`, `CofeAI/Tele-FLM`, etc. | ✅︎ | ✅︎ |
490
-| `MiniMaxM1ForCausalLM` | MiniMax-Text | `MiniMaxAI/MiniMax-M1-40k`, `MiniMaxAI/MiniMax-M1-80k`, etc. | | |
491
-| `MiniMaxText01ForCausalLM` | MiniMax-Text | `MiniMaxAI/MiniMax-Text-01`, etc. | | |
489
| `Zamba2ForCausalLM` | Zamba2 | `Zyphra/Zamba2-7B-instruct`, `Zyphra/Zamba2-2.7B-instruct`, `Zyphra/Zamba2-1.2B-instruct`, etc. | | |
490
491
!!! note
@@ -595,7 +592,6 @@ These models primarily accept the [`LLM.generate`](./generative_models.md#llmgen
592
| `MiMoV2OmniForCausalLM` | MiMo-V2.5-Omni | T + I<sup>E+</sup> + V<sup>E+</sup> + A<sup>+</sup> | `XiaomiMiMo/MiMo-V2.5-Omni` | | ✅︎ |
593
| `MiniCPMO` | MiniCPM-O | T + I<sup>E+</sup> + V<sup>E+</sup> + A<sup>E+</sup> | `openbmb/MiniCPM-o-2_6`, etc. | ✅︎ | ✅︎ |
594
| `MiniCPMV` | MiniCPM-V | T + I<sup>E+</sup> + V<sup>E+</sup> | `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. | ✅︎ | |
598
-| `MiniMaxVL01ForConditionalGeneration` | MiniMax-VL | T + I<sup>E+</sup> | `MiniMaxAI/MiniMax-VL-01`, etc. | | ✅︎ |
595
| `Mistral3ForConditionalGeneration` | Mistral3 (HF Transformers) | T + I<sup>+</sup> | `mistralai/Mistral-Small-3.1-24B-Instruct-2503`, etc. | ✅︎ | ✅︎ |
596
| `MolmoForCausalLM` | Molmo | T + I<sup>+</sup> | `allenai/Molmo-7B-D-0924`, `allenai/Molmo-7B-O-0924`, etc. | ✅︎ | ✅︎ |
597
| `Molmo2ForConditionalGeneration` | Molmo2 | T + I<sup>+</sup> / V | `allenai/Molmo2-4B`, `allenai/Molmo2-8B`, `allenai/Molmo2-O-7B`, `allenai/MolmoWeb-4B`<sup>^</sup>, `allenai/MolmoWeb-8B`<sup>^</sup> | ✅︎ | ✅︎ |
docs/usage/v1_guide.md
+1
-1
@@ -128,7 +128,7 @@ Models that use Mamba-2 and Mamba-1 layers (e.g., `Mamba2ForCausalLM`, `MambaFor
128
Hybrid models that combine Mamba-2 and Mamba-1 layers with standard attention layers are also supported (e.g., `BambaForCausalLM`,
129
`Zamba2ForCausalLM`, `NemotronHForCausalLM`, `FalconH1ForCausalLM` and `GraniteMoeHybridForCausalLM`, `JambaForCausalLM`, `Plamo2ForCausalLM`).
130
131
-Hybrid models with mechanisms different to Mamba are also supported (e.g, `MiniMaxText01ForCausalLM`, `MiniMaxM1ForCausalLM`, `Lfm2ForCausalLM`).
131
+Hybrid models with mechanisms different to Mamba are also supported (e.g, `Lfm2ForCausalLM`).
132
133
Please note that prefix caching is not yet supported for any of the above models.
134
examples/generate/multimodal/vision_language_offline.py
-34
@@ -1481,39 +1481,6 @@ def run_minicpmv(questions: list[str], modality: str) -> ModelRequestData:
1481
return run_minicpmv_base(questions, modality, "openbmb/MiniCPM-V-2_6")
1482
1483
1484
-def run_minimax_vl_01(questions: list[str], modality: str) -> ModelRequestData:
1485
- assert modality == "image"
1486
-
1487
- model_name = "MiniMaxAI/MiniMax-VL-01"
1488
-
1489
- engine_args = EngineArgs(
1490
- model=model_name,
1491
- max_num_seqs=2,
1492
- limit_mm_per_prompt={modality: 1},
1493
- trust_remote_code=True,
1494
- tensor_parallel_size=8,
1495
- )
1496
-
1497
- tokenizer = AutoTokenizer.from_pretrained(model_name)
1498
- messages = [
1499
- [
1500
- {
1501
- "role": "user",
1502
- "content": [{"type": "image"}, {"type": "text", "text": question}],
1503
- }
1504
- ]
1505
- for question in questions
1506
- ]
1507
- prompts = tokenizer.apply_chat_template(
1508
- messages, add_generation_prompt=True, tokenize=False
1509
- )
1510
-
1511
- return ModelRequestData(
1512
- engine_args=engine_args,
1513
- prompts=prompts,
1514
- )
1515
-
1516
-
1484
# Mistral-3 HF-format
1485
def run_mistral3(questions: list[str], modality: str) -> ModelRequestData:
1486
assert modality == "image"
@@ -2485,7 +2452,6 @@ model_example_map = {
2452
"mantis": run_mantis,
2453
"minicpmo": run_minicpmo,
2454
"minicpmv": run_minicpmv,
2488
- "minimax_vl_01": run_minimax_vl_01,
2455
"mistral3": run_mistral3,
2456
"molmo": run_molmo,
2457
"molmo2": run_molmo2,
examples/tool_chat_template_minimax_m1.jinja
deleted
-91
@@ -1,91 +0,0 @@
1
-{{ '<begin_of_document>' -}}
2
-{%- if custom_tools is defined %}
3
- {%- set tools = custom_tools %}
4
-{%- endif %}
5
-{%- if not tools is defined %}
6
- {%- set tools = none %}
7
-{%- endif %}
8
-
9
-{#- Extract system message #}
10
-{% set ns = namespace(system_prompt='') -%}
11
-{%- if messages[0]['role'] == 'system' %}
12
- {%- if messages[0]['content'] is string %}
13
- {%- set ns.system_prompt = messages[0]['content']|trim %}
14
- {%- else %}
15
- {%- set ns.system_prompt = messages[0]['content'][0]['text']|trim %}
16
- {%- endif %}
17
- {%- set messages = messages[1:] %}
18
-{%- else %}
19
- {%- if tools is not none %}
20
- {%- set ns.system_prompt = "You are a helpful assistant created by Minimax based on MiniMax-M1 model." %}
21
- {%- else %}
22
- {%- set ns.system_prompt = "You are a helpful assistant created by Minimax based on MiniMax-M1 model." %}
23
- {%- endif %}
24
-{%- endif %}
25
-
26
-{#- System message #}
27
-{%- if ns.system_prompt != '' %}
28
-{{ '<beginning_of_sentence>system ai_setting=assistant\n' + ns.system_prompt + '<end_of_sentence>\n' -}}
29
-{%- endif %}
30
-
31
-{#- Tools configuration #}
32
-{%- if tools is not none %}
33
-{{ '<beginning_of_sentence>system tool_setting=tools\nYou are provided with these tools:\n<tools>\n' -}}
34
-{%- for tool in tools %}
35
-{{ tool | tojson ~ '\n' -}}
36
-{%- endfor %}
37
-{{ '</tools>\n\nIf you need to call tools, please respond with <tool_calls></tool_calls> XML tags, and provide tool-name and json-object of arguments, following the format below:\n<tool_calls>\n{"name": <tool-name>, "arguments": <args-json-object>}\n...\n</tool_calls><end_of_sentence>\n' -}}
38
-{%- endif %}
39
-
40
-{#- Process messages #}
41
-{%- for message in messages %}
42
- {%- if not (message.role == 'ipython' or message.role == 'tool' or 'tool_calls' in message) %}
43
- {%- if message['role'] == 'user' %}
44
-{{ '<beginning_of_sentence>user name=user\n' -}}
45
-{%- if message['content'] is string %}
46
-{{ message['content']|trim -}}
47
-{%- else %}
48
-{%- for content in message['content'] %}
49
-{%- if content['type'] == 'text' %}
50
-{{ content['text']|trim -}}
51
-{%- endif %}
52
-{%- endfor %}
53
-{%- endif %}
54
-{{ '<end_of_sentence>\n' -}}
55
- {%- elif message['role'] == 'assistant' %}
56
-{{ '<beginning_of_sentence>ai name=assistant\n' -}}
57
-{%- if message['content'] is string %}
58
-{{ message['content']|trim -}}
59
-{%- else %}
60
-{%- for content in message['content'] | selectattr('type', 'equalto', 'text') %}
61
-{{ content['text']|trim -}}
62
-{%- endfor %}
63
-{%- endif %}
64
-{{ '<end_of_sentence>\n' -}}
65
- {%- endif %}
66
- {%- elif 'tool_calls' in message %}
67
-{{ '<beginning_of_sentence>ai name=assistant\n<tool_calls>\n' -}}
68
-{%- for tool_call in message.tool_calls %}
69
-{{ '{"name": "' + tool_call.function.name + '", "arguments": ' + tool_call.function.arguments | tojson + '}\n' -}}
70
-{%- endfor %}
71
-{{ '</tool_calls><end_of_sentence>\n' -}}
72
- {%- elif message.role == "tool" or message.role == "ipython" %}
73
-{{ '<beginning_of_sentence>tool name=tools\n' -}}
74
-{%- if message.content is string %}
75
-{{ 'tool result: ' + message.content + '\n\n' -}}
76
-{%- else %}
77
-{%- for content in message['content'] %}
78
-{%- if content['type'] == 'text' %}
79
-{{ 'tool result: ' + content['text'] + '\n\n' -}}
80
-{%- elif content.get('name') %}
81
-{{ 'tool name: ' + content['name'] + '\ntool result: ' + content['text'] + '\n\n' -}}
82
-{%- endif %}
83
-{%- endfor %}
84
-{%- endif %}
85
-{{ '<end_of_sentence>\n' -}}
86
- {%- endif %}
87
-{%- endfor %}
88
-
89
-{%- if add_generation_prompt %}
90
-{{ '<beginning_of_sentence>ai name=assistant\n' -}}
91
-{%- endif %}
\ No newline at end of file
rust/src/chat/src/renderer/hf/format.rs
-1
@@ -386,7 +386,6 @@ mod tests {
386
tool_chat_template_llama3.2_pythonic.jinja => String
387
tool_chat_template_llama4_json.jinja => OpenAi
388
tool_chat_template_llama4_pythonic.jinja => OpenAi
389
- tool_chat_template_minimax_m1.jinja => OpenAi
389
tool_chat_template_mistral.jinja => String
390
tool_chat_template_mistral3.jinja => OpenAi
391
tool_chat_template_mistral_parallel.jinja => String
rust/src/chat/tests/templates/vllm_examples/tool_chat_template_minimax_m1.jinja
deleted
-91
@@ -1,91 +0,0 @@
1
-{{ '<begin_of_document>' -}}
2
-{%- if custom_tools is defined %}
3
- {%- set tools = custom_tools %}
4
-{%- endif %}
5
-{%- if not tools is defined %}
6
- {%- set tools = none %}
7
-{%- endif %}
8
-
9
-{#- Extract system message #}
10
-{% set ns = namespace(system_prompt='') -%}
11
-{%- if messages[0]['role'] == 'system' %}
12
- {%- if messages[0]['content'] is string %}
13
- {%- set ns.system_prompt = messages[0]['content']|trim %}
14
- {%- else %}
15
- {%- set ns.system_prompt = messages[0]['content'][0]['text']|trim %}
16
- {%- endif %}
17
- {%- set messages = messages[1:] %}
18
-{%- else %}
19
- {%- if tools is not none %}
20
- {%- set ns.system_prompt = "You are a helpful assistant created by Minimax based on MiniMax-M1 model." %}
21
- {%- else %}
22
- {%- set ns.system_prompt = "You are a helpful assistant created by Minimax based on MiniMax-M1 model." %}
23
- {%- endif %}
24
-{%- endif %}
25
-
26
-{#- System message #}
27
-{%- if ns.system_prompt != '' %}
28
-{{ '<beginning_of_sentence>system ai_setting=assistant\n' + ns.system_prompt + '<end_of_sentence>\n' -}}
29
-{%- endif %}
30
-
31
-{#- Tools configuration #}
32
-{%- if tools is not none %}
33
-{{ '<beginning_of_sentence>system tool_setting=tools\nYou are provided with these tools:\n<tools>\n' -}}
34
-{%- for tool in tools %}
35
-{{ tool | tojson ~ '\n' -}}
36
-{%- endfor %}
37
-{{ '</tools>\n\nIf you need to call tools, please respond with <tool_calls></tool_calls> XML tags, and provide tool-name and json-object of arguments, following the format below:\n<tool_calls>\n{"name": <tool-name>, "arguments": <args-json-object>}\n...\n</tool_calls><end_of_sentence>\n' -}}
38
-{%- endif %}
39
-
40
-{#- Process messages #}
41
-{%- for message in messages %}
42
- {%- if not (message.role == 'ipython' or message.role == 'tool' or 'tool_calls' in message) %}
43
- {%- if message['role'] == 'user' %}
44
-{{ '<beginning_of_sentence>user name=user\n' -}}
45
-{%- if message['content'] is string %}
46
-{{ message['content']|trim -}}
47
-{%- else %}
48
-{%- for content in message['content'] %}
49
-{%- if content['type'] == 'text' %}
50
-{{ content['text']|trim -}}
51
-{%- endif %}
52
-{%- endfor %}
53
-{%- endif %}
54
-{{ '<end_of_sentence>\n' -}}
55
- {%- elif message['role'] == 'assistant' %}
56
-{{ '<beginning_of_sentence>ai name=assistant\n' -}}
57
-{%- if message['content'] is string %}
58
-{{ message['content']|trim -}}
59
-{%- else %}
60
-{%- for content in message['content'] | selectattr('type', 'equalto', 'text') %}
61
-{{ content['text']|trim -}}
62
-{%- endfor %}
63
-{%- endif %}
64
-{{ '<end_of_sentence>\n' -}}
65
- {%- endif %}
66
- {%- elif 'tool_calls' in message %}
67
-{{ '<beginning_of_sentence>ai name=assistant\n<tool_calls>\n' -}}
68
-{%- for tool_call in message.tool_calls %}
69
-{{ '{"name": "' + tool_call.function.name + '", "arguments": ' + tool_call.function.arguments | tojson + '}\n' -}}
70
-{%- endfor %}
71
-{{ '</tool_calls><end_of_sentence>\n' -}}
72
- {%- elif message.role == "tool" or message.role == "ipython" %}
73
-{{ '<beginning_of_sentence>tool name=tools\n' -}}
74
-{%- if message.content is string %}
75
-{{ 'tool result: ' + message.content + '\n\n' -}}
76
-{%- else %}
77
-{%- for content in message['content'] %}
78
-{%- if content['type'] == 'text' %}
79
-{{ 'tool result: ' + content['text'] + '\n\n' -}}
80
-{%- elif content.get('name') %}
81
-{{ 'tool name: ' + content['name'] + '\ntool result: ' + content['text'] + '\n\n' -}}
82
-{%- endif %}
83
-{%- endfor %}
84
-{%- endif %}
85
-{{ '<end_of_sentence>\n' -}}
86
- {%- endif %}
87
-{%- endfor %}
88
-
89
-{%- if add_generation_prompt %}
90
-{{ '<beginning_of_sentence>ai name=assistant\n' -}}
91
-{%- endif %}
\ No newline at end of file
tests/models/multimodal/generation/test_common.py
-23
@@ -810,29 +810,6 @@ VLM_TEST_SETTINGS = {
810
hf_output_post_proc=model_utils.minicpmv_trunc_hf_output,
811
patch_hf_runner=model_utils.minicpmv_26_patch_hf_runner,
812
),
813
- "minimax_vl_01": VLMTestInfo(
814
- models=["MiniMaxAI/MiniMax-VL-01"],
815
- prompt_formatter=lambda img_prompt: f"<beginning_of_sentence>user: {img_prompt} assistant:<end_of_sentence>", # noqa: E501
816
- img_idx_to_prompt=lambda _: "<image>",
817
- test_type=(VLMTestType.IMAGE, VLMTestType.MULTI_IMAGE),
818
- max_model_len=8192,
819
- max_num_seqs=4,
820
- dtype="bfloat16",
821
- hf_output_post_proc=model_utils.minimax_vl_01_hf_output,
822
- patch_hf_runner=model_utils.minimax_vl_01_patch_hf_runner,
823
- auto_cls=AutoModelForImageTextToText,
824
- marks=[
825
- large_gpu_mark(min_gb=80),
826
- # TODO: [ROCm] Fix pickle issue with ROCm spawn and tp>1
827
- pytest.mark.skipif(
828
- current_platform.is_rocm(),
829
- reason=(
830
- "ROCm: Model too large for single GPU; "
831
- "multi-GPU blocked by HF _LazyConfigMapping pickle issue with spawn"
832
- ),
833
- ),
834
- ],
835
- ),
813
"molmo": VLMTestInfo(
814
models=["allenai/Molmo-7B-D-0924"],
815
test_type=(VLMTestType.IMAGE, VLMTestType.MULTI_IMAGE),
tests/models/multimodal/generation/vlm_utils/model_utils.py
-18
@@ -245,13 +245,6 @@ def minicpmv_trunc_hf_output(hf_output: RunnerOutput, model: str) -> RunnerOutpu
245
return output_ids, output_str, out_logprobs
246
247
248
-def minimax_vl_01_hf_output(hf_output: RunnerOutput, model: str) -> RunnerOutput:
249
- output_ids, output_str, out_logprobs = hf_output
250
- if output_str.endswith("<end_of_sentence>"):
251
- output_str = output_str.split("<end_of_sentence>")[0]
252
- return output_ids, output_str, out_logprobs
253
-
254
-
248
def ultravox_trunc_hf_output(hf_output: RunnerOutput, model: str) -> RunnerOutput:
249
output_ids, output_str, out_logprobs = hf_output
250
@@ -1023,17 +1016,6 @@ def minicpmv_26_patch_hf_runner(hf_model: HfRunner) -> HfRunner:
1016
return hf_model
1017
1018
1026
-def minimax_vl_01_patch_hf_runner(hf_model: HfRunner) -> HfRunner:
1027
- orig_generate = hf_model.model.generate
1028
-
1029
- def _generate(self, *args, image_sizes=None, **kwargs):
1030
- return orig_generate(*args, decode_text=False, **kwargs)
1031
-
1032
- hf_model.model.generate = types.MethodType(_generate, hf_model.model)
1033
-
1034
- return hf_model
1035
-
1036
-
1019
def molmo_patch_hf_runner(hf_model: HfRunner) -> HfRunner:
1020
"""Patches and returns an instance of the HfRunner to use for Molmo."""
1021
hf_processor = hf_model.processor
tests/models/multimodal/processing/test_minimax_vl_01.py
deleted
-113
@@ -1,113 +0,0 @@
1
-# SPDX-License-Identifier: Apache-2.0
2
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
-
4
-import pytest
5
-from PIL import Image
6
-
7
-from vllm.multimodal import MULTIMODAL_REGISTRY
8
-from vllm.multimodal.parse import ImageSize
9
-from vllm.multimodal.processing import BaseMultiModalProcessor
10
-
11
-from ....conftest import ImageTestAssets
12
-from ...utils import build_model_context
13
-
14
-
15
-@pytest.mark.parametrize("model_id", ["MiniMaxAI/MiniMax-VL-01"])
16
-@pytest.mark.parametrize("num_imgs", [1, 2])
17
-def test_processor_override(
18
- image_assets: ImageTestAssets,
19
- model_id: str,
20
- num_imgs: int,
21
-):
22
- ctx = build_model_context(
23
- model_id,
24
- mm_processor_kwargs=None,
25
- limit_mm_per_prompt={"image": num_imgs},
26
- )
27
- processor = MULTIMODAL_REGISTRY.create_processor(ctx.model_config)
28
- prompt = "<image>" * num_imgs
29
- image = Image.new("RGB", size=(364, 364))
30
- mm_data = {"image": [image] * num_imgs}
31
-
32
- processed_inputs = processor(
33
- prompt,
34
- mm_items=processor.info.parse_mm_data(mm_data),
35
- hf_processor_mm_kwargs={},
36
- )
37
- image_placeholders = processed_inputs["mm_placeholders"]["image"]
38
-
39
- assert len(image_placeholders) == num_imgs
40
-
41
-
42
-def _validate_image_prompt_replacements_one(
43
- processor: BaseMultiModalProcessor,
44
- num_imgs: int,
45
- failed_size_excs: list[tuple[ImageSize, Exception]],
46
- image_size: ImageSize,
47
-) -> None:
48
- prompt = "<image>" * num_imgs
49
- image = Image.new("RGB", size=image_size)
50
- mm_data = {"image": [image] * num_imgs}
51
-
52
- try:
53
- processed_inputs = processor(
54
- prompt,
55
- mm_items=processor.info.parse_mm_data(mm_data),
56
- hf_processor_mm_kwargs={},
57
- )
58
-
59
- image_placeholders = processed_inputs["mm_placeholders"]["image"]
60
- assert len(image_placeholders) == num_imgs
61
-
62
- except Exception as exc:
63
- failed_size_excs.append((image_size, exc))
64
-
65
-
66
-def _test_image_prompt_replacements(
67
- processor,
68
- *,
69
- num_imgs: int,
70
- image_sizes: list[ImageSize],
71
-) -> None:
72
- failed_size_excs = list[tuple[ImageSize, Exception]]()
73
-
74
- for size in image_sizes:
75
- _validate_image_prompt_replacements_one(
76
- processor, num_imgs, failed_size_excs, size
77
- )
78
-
79
- if failed_size_excs:
80
- msg = "Found failing image sizes:" + "\n========\n".join(
81
- f"[{size}]\n{exc}" for size, exc in failed_size_excs
82
- )
83
- raise AssertionError(msg)
84
-
85
-
86
-@pytest.mark.parametrize("model_id", ["MiniMaxAI/MiniMax-VL-01"])
87
-@pytest.mark.parametrize("num_imgs", [1, 2])
88
-def test_processor_prompt_replacements_regression(model_id, num_imgs):
89
- ctx = build_model_context(
90
- model_id,
91
- mm_processor_kwargs=None,
92
- limit_mm_per_prompt={"image": num_imgs},
93
- )
94
- processor = MULTIMODAL_REGISTRY.create_processor(ctx.model_config)
95
-
96
- image_ratios = [
97
- (171, 152),
98
- (184, 161),
99
- (198, 176),
100
- (333, 296),
101
- (369, 328),
102
- (488, 183),
103
- (2560, 1669),
104
- ]
105
- image_sizes = [
106
- size for w, h in image_ratios for size in [ImageSize(w, h), ImageSize(h, w)]
107
- ]
108
-
109
- _test_image_prompt_replacements(
110
- processor,
111
- num_imgs=num_imgs,
112
- image_sizes=image_sizes,
113
- )
tests/models/registry.py
-13
@@ -421,15 +421,6 @@ _TEXT_GENERATION_EXAMPLE_MODELS = {
421
},
422
trust_remote_code=True,
423
),
424
- "MiniMaxForCausalLM": _HfExamplesInfo("MiniMaxAI/MiniMax-Text-01-hf"),
425
- "MiniMaxText01ForCausalLM": _HfExamplesInfo(
426
- "MiniMaxAI/MiniMax-Text-01",
427
- trust_remote_code=True,
428
- revision="a59aa9cbc53b9fb8742ca4e9e1531b9802b6fdc3",
429
- ),
430
- "MiniMaxM1ForCausalLM": _HfExamplesInfo(
431
- "MiniMaxAI/MiniMax-M1-40k", trust_remote_code=True
432
- ),
424
"MiniMaxM2ForCausalLM": _HfExamplesInfo(
425
"MiniMaxAI/MiniMax-M2",
426
trust_remote_code=True,
@@ -1113,10 +1104,6 @@ _MULTIMODAL_EXAMPLE_MODELS = {
1104
"openbmb/MiniCPM-V-4_6",
1105
min_transformers_version="5.7.0",
1106
),
1116
- "MiniMaxVL01ForConditionalGeneration": _HfExamplesInfo(
1117
- "MiniMaxAI/MiniMax-VL-01",
1118
- trust_remote_code=True,
1119
- ),
1107
"MiniMaxM3SparseForConditionalGeneration": _HfExamplesInfo(
1108
"MiniMaxAI/MiniMax-M3",
1109
trust_remote_code=True,
tests/models/test_initialization.py
-5
@@ -98,11 +98,6 @@ def can_initialize(
98
vllm_config.validate_block_size()
99
return scheduler_kv_cache_config
100
101
- if model_arch == "MiniMaxVL01ForConditionalGeneration":
102
- pytest.skip(
103
- "pickle error when loading `transformers.models.auto.CONFIG_MAPPING`"
104
- )
105
-
101
if model_arch == "MoonshotKimiaForCausalLM":
102
pytest.skip(
103
"Kimi-Audio requires SpeechToTextConfig "
tests/tool_parsers/test_minimax_tool_parser.py
deleted
-1227
@@ -1,1227 +0,0 @@
1
-# SPDX-License-Identifier: Apache-2.0
2
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
-# ruff: noqa: E501
4
-
5
-import json
6
-from typing import Any
7
-
8
-import pytest
9
-
10
-from vllm.entrypoints.openai.chat_completion.protocol import (
11
- ChatCompletionToolsParam,
12
-)
13
-from vllm.entrypoints.openai.engine.protocol import (
14
- FunctionCall,
15
- ToolCall,
16
-)
17
-from vllm.tokenizers import get_tokenizer
18
-from vllm.tool_parsers.minimax_tool_parser import MinimaxToolParser
19
-
20
-# Use a common model that is likely to be available
21
-MODEL = "MiniMaxAi/MiniMax-M1-40k"
22
-
23
-
24
-@pytest.fixture(scope="module")
25
-def minimax_tokenizer():
26
- return get_tokenizer(tokenizer_name=MODEL)
27
-
28
-
29
-@pytest.fixture
30
-def minimax_tool_parser(minimax_tokenizer):
31
- return MinimaxToolParser(minimax_tokenizer)
32
-
33
-
34
-@pytest.fixture
35
-def sample_tools():
36
- return [
37
- ChatCompletionToolsParam(
38
- type="function",
39
- function={
40
- "name": "get_current_weather",
41
- "description": "Get the current weather",
42
- "parameters": {
43
- "type": "object",
44
- "properties": {
45
- "city": {"type": "string", "description": "The city name"},
46
- "state": {"type": "string", "description": "The state code"},
47
- "unit": {"type": "string", "enum": ["fahrenheit", "celsius"]},
48
- },
49
- "required": ["city", "state"],
50
- },
51
- },
52
- ),
53
- ChatCompletionToolsParam(
54
- type="function",
55
- function={
56
- "name": "calculate_area",
57
- "description": "Calculate area of a shape",
58
- "parameters": {
59
- "type": "object",
60
- "properties": {
61
- "shape": {"type": "string"},
62
- "dimensions": {"type": "object"},
63
- "precision": {"type": "integer"},
64
- },
65
- },
66
- },
67
- ),
68
- ]
69
-
70
-
71
-def assert_tool_calls(
72
- actual_tool_calls: list[ToolCall], expected_tool_calls: list[ToolCall]
73
-):
74
- assert len(actual_tool_calls) == len(expected_tool_calls)
75
-
76
- for actual_tool_call, expected_tool_call in zip(
77
- actual_tool_calls, expected_tool_calls
78
- ):
79
- assert isinstance(actual_tool_call.id, str)
80
- assert len(actual_tool_call.id) > 16
81
-
82
- assert actual_tool_call.type == "function"
83
- assert actual_tool_call.function == expected_tool_call.function
84
-
85
-
86
-def test_extract_tool_calls_no_tools(minimax_tool_parser):
87
- model_output = "This is a test"
88
- extracted_tool_calls = minimax_tool_parser.extract_tool_calls(
89
- model_output, request=None
90
- ) # type: ignore[arg-type]
91
- assert not extracted_tool_calls.tools_called
92
- assert extracted_tool_calls.tool_calls == []
93
- assert extracted_tool_calls.content == model_output
94
-
95
-
96
-@pytest.mark.parametrize(
97
- ids=[
98
- "single_tool_call",
99
- "multiple_tool_calls",
100
- "tool_call_with_content_before",
101
- "tool_call_with_single_line_json",
102
- "tool_call_incomplete_tag",
103
- ],
104
- argnames=["model_output", "expected_tool_calls", "expected_content"],
105
- argvalues=[
106
- (
107
- """<tool_calls>
108
-{"name": "get_current_weather", "arguments": {"city": "Dallas", "state": "TX", "unit": "fahrenheit"}}
109
-</tool_calls>""",
110
- [
111
- ToolCall(
112
- function=FunctionCall(
113
- name="get_current_weather",
114
- arguments=json.dumps(
115
- {
116
- "city": "Dallas",
117
- "state": "TX",
118
- "unit": "fahrenheit",
119
- }
120
- ),
121
- )
122
- )
123
- ],
124
- None,
125
- ),
126
- (
127
- """<tool_calls>
128
-{"name": "get_current_weather", "arguments": {"city": "Dallas", "state": "TX", "unit": "fahrenheit"}}
129
-{"name": "get_current_weather", "arguments": {"city": "Orlando", "state": "FL", "unit": "fahrenheit"}}
130
-</tool_calls>""",
131
- [
132
- ToolCall(
133
- function=FunctionCall(
134
- name="get_current_weather",
135
- arguments=json.dumps(
136
- {
137
- "city": "Dallas",
138
- "state": "TX",
139
- "unit": "fahrenheit",
140
- }
141
- ),
142
- )
143
- ),
144
- ToolCall(
145
- function=FunctionCall(
146
- name="get_current_weather",
147
- arguments=json.dumps(
148
- {
149
- "city": "Orlando",
150
- "state": "FL",
151
- "unit": "fahrenheit",
152
- }
153
- ),
154
- )
155
- ),
156
- ],
157
- None,
158
- ),
159
- (
160
- """I'll help you check the weather. <tool_calls>
161
-{"name": "get_current_weather", "arguments": {"city": "Seattle", "state": "WA", "unit": "celsius"}}
162
-</tool_calls>""",
163
- [
164
- ToolCall(
165
- function=FunctionCall(
166
- name="get_current_weather",
167
- arguments=json.dumps(
168
- {
169
- "city": "Seattle",
170
- "state": "WA",
171
- "unit": "celsius",
172
- }
173
- ),
174
- )
175
- )
176
- ],
177
- "I'll help you check the weather.",
178
- ),
179
- (
180
- """<tool_calls>
181
-{"name": "get_current_weather", "arguments": {"city": "New York", "state": "NY", "unit": "celsius"}}
182
-</tool_calls>""",
183
- [
184
- ToolCall(
185
- function=FunctionCall(
186
- name="get_current_weather",
187
- arguments=json.dumps(
188
- {
189
- "city": "New York",
190
- "state": "NY",
191
- "unit": "celsius",
192
- }
193
- ),
194
- )
195
- )
196
- ],
197
- None,
198
- ),
199
- (
200
- """<tool_calls>
201
-{"name": "get_current_weather", "arguments": {"city": "Boston", "state": "MA"}}""",
202
- [
203
- ToolCall(
204
- function=FunctionCall(
205
- name="get_current_weather",
206
- arguments=json.dumps(
207
- {
208
- "city": "Boston",
209
- "state": "MA",
210
- }
211
- ),
212
- )
213
- )
214
- ],
215
- None,
216
- ),
217
- ],
218
-)
219
-def test_extract_tool_calls(
220
- minimax_tool_parser, model_output, expected_tool_calls, expected_content
221
-):
222
- extracted_tool_calls = minimax_tool_parser.extract_tool_calls(
223
- model_output, request=None
224
- ) # type: ignore[arg-type]
225
- assert extracted_tool_calls.tools_called
226
-
227
- assert_tool_calls(extracted_tool_calls.tool_calls, expected_tool_calls)
228
-
229
- assert extracted_tool_calls.content == expected_content
230
-
231
-
232
-def test_preprocess_model_output_with_thinking_tags(minimax_tool_parser):
233
- """Test that tool calls within thinking tags are removed during preprocessing."""
234
- model_output = """<think>Let me think about this. <tool_calls>
235
-{"name": "fake_tool", "arguments": {"param": "value"}}
236
-</tool_calls> This should be removed.</think>
237
-
238
-I'll help you with that. <tool_calls>
239
-{"name": "get_current_weather", "arguments": {"city": "Seattle", "state": "WA"}}
240
-</tool_calls>"""
241
-
242
- processed_output = minimax_tool_parser.preprocess_model_output(model_output)
243
-
244
- # The tool call within thinking tags should be removed
245
- assert "fake_tool" not in processed_output
246
- # But the thinking tag itself should remain
247
- assert "<think>" in processed_output
248
- assert "</think>" in processed_output
249
- # The actual tool call outside thinking tags should remain
250
- assert "get_current_weather" in processed_output
251
-
252
-
253
-def test_extract_tool_calls_with_thinking_tags(minimax_tool_parser):
254
- """Test tool extraction when thinking tags contain tool calls that should be ignored."""
255
- model_output = """<think>I should use a tool. <tool_calls>
256
-{"name": "ignored_tool", "arguments": {"should": "ignore"}}
257
-</tool_calls></think>
258
-
259
-Let me help you with the weather. <tool_calls>
260
-{"name": "get_current_weather", "arguments": {"city": "Miami", "state": "FL", "unit": "fahrenheit"}}
261
-</tool_calls>"""
262
-
263
- extracted_tool_calls = minimax_tool_parser.extract_tool_calls(
264
- model_output, request=None
265
- ) # type: ignore[arg-type]
266
-
267
- assert extracted_tool_calls.tools_called
268
- assert len(extracted_tool_calls.tool_calls) == 1
269
- assert extracted_tool_calls.tool_calls[0].function.name == "get_current_weather"
270
-
271
- # Content extraction is based on the position of the first <tool_calls> in the original model_output
272
- # Since preprocessing removes tool calls within thinking tags, the actual first <tool_calls> is the external one
273
- expected_content = """<think>I should use a tool. <tool_calls>
274
-{"name": "ignored_tool", "arguments": {"should": "ignore"}}
275
-</tool_calls></think>
276
-
277
-Let me help you with the weather."""
278
- assert extracted_tool_calls.content == expected_content
279
-
280
-
281
-def test_extract_tool_calls_invalid_json(minimax_tool_parser):
282
- """Test that invalid JSON in tool calls is handled gracefully."""
283
- model_output = """<tool_calls>
284
-{"name": "valid_tool", "arguments": {"city": "Seattle"}}
285
-{invalid json here}
286
-{"name": "another_valid_tool", "arguments": {"param": "value"}}
287
-</tool_calls>"""
288
-
289
- extracted_tool_calls = minimax_tool_parser.extract_tool_calls(
290
- model_output, request=None
291
- ) # type: ignore[arg-type]
292
-
293
- assert extracted_tool_calls.tools_called
294
- # Should extract only the valid JSON tool calls
295
- assert len(extracted_tool_calls.tool_calls) == 2
296
- assert extracted_tool_calls.tool_calls[0].function.name == "valid_tool"
297
- assert extracted_tool_calls.tool_calls[1].function.name == "another_valid_tool"
298
-
299
-
300
-def test_extract_tool_calls_missing_name_or_arguments(minimax_tool_parser):
301
- """Test that tool calls missing name or arguments are filtered out."""
302
- model_output = """<tool_calls>
303
-{"name": "valid_tool", "arguments": {"city": "Seattle"}}
304
-{"name": "missing_args"}
305
-{"arguments": {"city": "Portland"}}
306
-{"name": "another_valid_tool", "arguments": {"param": "value"}}
307
-</tool_calls>"""
308
-
309
- extracted_tool_calls = minimax_tool_parser.extract_tool_calls(
310
- model_output, request=None
311
- ) # type: ignore[arg-type]
312
-
313
- assert extracted_tool_calls.tools_called
314
- # Should extract only the valid tool calls with both name and arguments
315
- assert len(extracted_tool_calls.tool_calls) == 2
316
- assert extracted_tool_calls.tool_calls[0].function.name == "valid_tool"
317
- assert extracted_tool_calls.tool_calls[1].function.name == "another_valid_tool"
318
-
319
-
320
-def test_streaming_basic_functionality(minimax_tool_parser):
321
- """Test basic streaming functionality."""
322
- # Reset streaming state
323
- minimax_tool_parser.current_tool_name_sent = False
324
- minimax_tool_parser.prev_tool_call_arr = []
325
- minimax_tool_parser.current_tool_id = -1
326
- minimax_tool_parser.streamed_args_for_tool = []
327
-
328
- # Test with a simple tool call
329
- current_text = """<tool_calls>
330
-{"name": "get_current_weather", "arguments": {"city": "Seattle"}}
331
-</tool_calls>"""
332
-
333
- # First call should handle the initial setup
334
- result = minimax_tool_parser.extract_tool_calls_streaming(
335
- previous_text="",
336
- current_text=current_text,
337
- delta_text="</tool_calls>",
338
- previous_token_ids=[],
339
- current_token_ids=[],
340
- delta_token_ids=[],
341
- request=None,
342
- )
343
-
344
- # The result might be None or contain tool call information
345
- # This depends on the internal state management
346
- if result is not None and hasattr(result, "tool_calls") and result.tool_calls:
347
- assert len(result.tool_calls) >= 0
348
-
349
-
350
-def test_streaming_with_content_before_tool_calls(minimax_tool_parser):
351
- """Test streaming when there's content before tool calls."""
352
- # Reset streaming state
353
- minimax_tool_parser.current_tool_name_sent = False
354
- minimax_tool_parser.prev_tool_call_arr = []
355
- minimax_tool_parser.current_tool_id = -1
356
- minimax_tool_parser.streamed_args_for_tool = []
357
-
358
- current_text = "I'll help you with that. <tool_calls>"
359
-
360
- # When there's content before tool calls, it should be returned as content
361
- result = minimax_tool_parser.extract_tool_calls_streaming(
362
- previous_text="I'll help you",
363
- current_text=current_text,
364
- delta_text=" with that. <tool_calls>",
365
- previous_token_ids=[],
366
- current_token_ids=[],
367
- delta_token_ids=[],
368
- request=None,
369
- )
370
-
371
- if result is not None and hasattr(result, "content"):
372
- # Should contain some content
373
- assert result.content is not None
374
-
375
-
376
-def test_streaming_no_tool_calls(minimax_tool_parser):
377
- """Test streaming when there are no tool calls."""
378
- current_text = "This is just regular text without any tool calls."
379
-
380
- result = minimax_tool_parser.extract_tool_calls_streaming(
381
- previous_text="This is just regular text",
382
- current_text=current_text,
383
- delta_text=" without any tool calls.",
384
- previous_token_ids=[],
385
- current_token_ids=[],
386
- delta_token_ids=[],
387
- request=None,
388
- )
389
-
390
- # Should return the delta text as content
391
- assert result is not None
392
- assert hasattr(result, "content")
393
- assert result.content == " without any tool calls."
394
-
395
-
396
-def test_streaming_with_thinking_tags(minimax_tool_parser):
397
- """Test streaming with thinking tags that contain tool calls."""
398
- # Reset streaming state
399
- minimax_tool_parser.current_tool_name_sent = False
400
- minimax_tool_parser.prev_tool_call_arr = []
401
- minimax_tool_parser.current_tool_id = -1
402
- minimax_tool_parser.streamed_args_for_tool = []
403
-
404
- current_text = """<think><tool_calls>{"name": "ignored", "arguments": {}}</tool_calls></think><tool_calls>{"name": "real_tool", "arguments": {"param": "value"}}</tool_calls>"""
405
-
406
- result = minimax_tool_parser.extract_tool_calls_streaming(
407
- previous_text="",
408
- current_text=current_text,
409
- delta_text=current_text,
410
- previous_token_ids=[],
411
- current_token_ids=[],
412
- delta_token_ids=[],
413
- request=None,
414
- )
415
-
416
- # The preprocessing should remove tool calls from thinking tags
417
- # and only process the real tool call
418
- if result is not None and hasattr(result, "tool_calls") and result.tool_calls:
419
- for tool_call in result.tool_calls:
420
- assert tool_call.function.name != "ignored"
421
-
422
-
423
-def test_extract_tool_calls_multiline_json_not_supported(minimax_tool_parser):
424
- """Test that multiline JSON in tool calls is not currently supported."""
425
- model_output = """<tool_calls>
426
-{
427
- "name": "get_current_weather",
428
- "arguments": {
429
- "city": "New York",
430
- "state": "NY",
431
- "unit": "celsius"
432
- }
433
-}
434
-</tool_calls>"""
435
-
436
- extracted_tool_calls = minimax_tool_parser.extract_tool_calls(
437
- model_output, request=None
438
- ) # type: ignore[arg-type]
439
-
440
- # Multiline JSON is currently not supported, should return no tools called
441
- assert not extracted_tool_calls.tools_called
442
- assert extracted_tool_calls.tool_calls == []
443
- assert extracted_tool_calls.content is None
444
-
445
-
446
-def test_streaming_arguments_incremental_output(minimax_tool_parser):
447
- """Test that streaming arguments are returned incrementally, not cumulatively."""
448
- # Reset streaming state
449
- minimax_tool_parser.current_tool_name_sent = False
450
- minimax_tool_parser.prev_tool_call_arr = []
451
- minimax_tool_parser.current_tool_id = -1
452
- minimax_tool_parser.streamed_args_for_tool = []
453
-
454
- # Simulate progressive tool call building
455
- stages = [
456
- # Stage 1: Function name complete
457
- '<tool_calls>\n{"name": "get_current_weather", "arguments": ',
458
- # Stage 2: Arguments object starts with first key
459
- '<tool_calls>\n{"name": "get_current_weather", "arguments": {"city": ',
460
- # Stage 3: First parameter value added
461
- '<tool_calls>\n{"name": "get_current_weather", "arguments": {"city": "Seattle"',
462
- # Stage 4: Second parameter added
463
- '<tool_calls>\n{"name": "get_current_weather", "arguments": {"city": "Seattle", "state": "WA"',
464
- # Stage 5: Third parameter added, arguments complete
465
- '<tool_calls>\n{"name": "get_current_weather", "arguments": {"city": "Seattle", "state": "WA", "unit": "celsius"}}',
466
- # Stage 6: Tool calls closed
467
- '<tool_calls>\n{"name": "get_current_weather", "arguments": {"city": "Seattle", "state": "WA", "unit": "celsius"}}\n</tool',
468
- '<tool_calls>\n{"name": "get_current_weather", "arguments": {"city": "Seattle", "state": "WA", "unit": "celsius"}}\n</tool_calls>',
469
- ]
470
-
471
- function_name_sent = False
472
- previous_args_content = ""
473
-
474
- for i, current_text in enumerate(stages):
475
- previous_text = stages[i - 1] if i > 0 else ""
476
- delta_text = current_text[len(previous_text) :] if i > 0 else current_text
477
-
478
- result = minimax_tool_parser.extract_tool_calls_streaming(
479
- previous_text=previous_text,
480
- current_text=current_text,
481
- delta_text=delta_text,
482
- previous_token_ids=[],
483
- current_token_ids=[],
484
- delta_token_ids=[],
485
- request=None,
486
- )
487
-
488
- print(f"Stage {i}: Current text: {repr(current_text)}")
489
- print(f"Stage {i}: Delta text: {repr(delta_text)}")
490
-
491
- if result is not None and hasattr(result, "tool_calls") and result.tool_calls:
492
- tool_call = result.tool_calls[0]
493
-
494
- # Check if function name is sent (should happen only once)
495
- if tool_call.function and tool_call.function.name:
496
- assert tool_call.function.name == "get_current_weather"
497
- function_name_sent = True
498
- print(f"Stage {i}: Function name sent: {tool_call.function.name}")
499
-
500
- # Check if arguments are sent incrementally
501
- if tool_call.function and tool_call.function.arguments:
502
- args_fragment = tool_call.function.arguments
503
- print(f"Stage {i}: Got arguments fragment: {repr(args_fragment)}")
504
-
505
- # For incremental output, each fragment should be new content only
506
- # The fragment should not contain all previous content
507
- if i >= 2 and previous_args_content: # After we start getting arguments
508
- # The new fragment should not be identical to or contain all previous content
509
- assert args_fragment != previous_args_content, (
510
- f"Fragment should be incremental, not cumulative: {args_fragment}"
511
- )
512
-
513
- # If this is truly incremental, the fragment should be relatively small
514
- # compared to the complete arguments so far
515
- if len(args_fragment) > len(previous_args_content):
516
- print(
517
- "Warning: Fragment seems cumulative rather than incremental"
518
- )
519
-
520
- previous_args_content = args_fragment
521
-
522
- # Verify function name was sent at least once
523
- assert function_name_sent, "Function name should have been sent"
524
-
525
-
526
-def test_streaming_arguments_delta_only(minimax_tool_parser):
527
- """Test that each streaming call returns only the delta (new part) of arguments."""
528
- # Reset streaming state
529
- minimax_tool_parser.current_tool_name_sent = False
530
- minimax_tool_parser.prev_tool_call_arr = []
531
- minimax_tool_parser.current_tool_id = -1
532
- minimax_tool_parser.streamed_args_for_tool = []
533
-
534
- # Simulate two consecutive calls with growing arguments
535
- call1_text = (
536
- '<tool_calls>\n{"name": "test_tool", "arguments": {"param1": "value1"}}'
537
- )
538
- call2_text = '<tool_calls>\n{"name": "test_tool", "arguments": {"param1": "value1", "param2": "value2"}}'
539
-
540
- print(f"Call 1 text: {repr(call1_text)}")
541
- print(f"Call 2 text: {repr(call2_text)}")
542
-
543
- # First call - should get the function name and initial arguments
544
- result1 = minimax_tool_parser.extract_tool_calls_streaming(
545
- previous_text="",
546
- current_text=call1_text,
547
- delta_text=call1_text,
548
- previous_token_ids=[],
549
- current_token_ids=[],
550
- delta_token_ids=[],
551
- request=None,
552
- )
553
-
554
- print(f"Result 1: {result1}")
555
- if result1 and hasattr(result1, "tool_calls") and result1.tool_calls:
556
- for i, tc in enumerate(result1.tool_calls):
557
- print(f" Tool call {i}: {tc}")
558
-
559
- # Second call - should only get the delta (new part) of arguments
560
- result2 = minimax_tool_parser.extract_tool_calls_streaming(
561
- previous_text=call1_text,
562
- current_text=call2_text,
563
- delta_text=', "param2": "value2"}',
564
- previous_token_ids=[],
565
- current_token_ids=[],
566
- delta_token_ids=[],
567
- request=None,
568
- )
569
-
570
- print(f"Result 2: {result2}")
571
- if result2 and hasattr(result2, "tool_calls") and result2.tool_calls:
572
- for i, tc in enumerate(result2.tool_calls):
573
- print(f" Tool call {i}: {tc}")
574
-
575
- # Verify the second call only returns the delta
576
- if result2 is not None and hasattr(result2, "tool_calls") and result2.tool_calls:
577
- tool_call = result2.tool_calls[0]
578
- if tool_call.function and tool_call.function.arguments:
579
- args_delta = tool_call.function.arguments
580
- print(f"Arguments delta from second call: {repr(args_delta)}")
581
-
582
- # Should only contain the new part, not the full arguments
583
- # The delta should be something like ', "param2": "value2"}' or just '"param2": "value2"'
584
- assert (
585
- ', "param2": "value2"}' in args_delta
586
- or '"param2": "value2"' in args_delta
587
- ), f"Expected delta containing param2, got: {args_delta}"
588
-
589
- # Should NOT contain the previous parameter data
590
- assert '"param1": "value1"' not in args_delta, (
591
- f"Arguments delta should not contain previous data: {args_delta}"
592
- )
593
-
594
- # The delta should be relatively short (incremental, not cumulative)
595
- expected_max_length = len(', "param2": "value2"}') + 10 # Some tolerance
596
- assert len(args_delta) <= expected_max_length, (
597
- f"Delta seems too long (possibly cumulative): {args_delta}"
598
- )
599
-
600
- print("✓ Delta validation passed")
601
- else:
602
- print("No arguments in result2 tool call")
603
- else:
604
- print("No tool calls in result2 or result2 is None")
605
- # This might be acceptable if no incremental update is needed
606
- # But let's at least verify that result1 had some content
607
- assert result1 is not None, "At least the first call should return something"
608
-
609
-
610
-def test_streaming_openai_compatibility(minimax_tool_parser):
611
- """Test that streaming behavior with buffering works correctly."""
612
- # Reset streaming state
613
- minimax_tool_parser.current_tool_name_sent = False
614
- minimax_tool_parser.prev_tool_call_arr = []
615
- minimax_tool_parser.current_tool_id = -1
616
- minimax_tool_parser.streamed_args_for_tool = []
617
- # Reset buffering state
618
- minimax_tool_parser.pending_buffer = ""
619
- minimax_tool_parser.in_thinking_tag = False
620
- minimax_tool_parser.thinking_depth = 0
621
-
622
- # Test scenario: simple buffering without complex tool call context
623
- test_cases: list[dict[str, Any]] = [
624
- {
625
- "stage": "Token: <",
626
- "previous": "",
627
- "current": "<",
628
- "delta": "<",
629
- "expected_content": None, # Should be buffered
630
- },
631
- {
632
- "stage": "Token: tool_calls>",
633
- "previous": "<",
634
- "current": "<tool_calls>",
635
- "delta": "tool_calls>",
636
- "expected_content": None, # Complete tag, should not output
637
- },
638
- {
639
- "stage": "Regular content",
640
- "previous": "Hello",
641
- "current": "Hello world",
642
- "delta": " world",
643
- "expected_content": " world", # Normal content should pass through
644
- },
645
- {
646
- "stage": "Content with end tag start",
647
- "previous": "Text",
648
- "current": "Text content</tool_",
649
- "delta": " content</tool_",
650
- "expected_content": " content", # Content part output, </tool_ buffered
651
- },
652
- {
653
- "stage": "Complete end tag",
654
- "previous": "Text content</tool_",
655
- "current": "Text content</tool_calls>",
656
- "delta": "calls>",
657
- "expected_content": None, # Complete close tag, should not output
658
- },
659
- ]
660
-
661
- for i, test_case in enumerate(test_cases):
662
- print(f"\n--- Stage {i}: {test_case['stage']} ---")
663
- print(f"Previous: {repr(test_case['previous'])}")
664
- print(f"Current: {repr(test_case['current'])}")
665
- print(f"Delta: {repr(test_case['delta'])}")
666
-
667
- result = minimax_tool_parser.extract_tool_calls_streaming(
668
- previous_text=test_case["previous"],
669
- current_text=test_case["current"],
670
- delta_text=test_case["delta"],
671
- previous_token_ids=[],
672
- current_token_ids=[],
673
- delta_token_ids=[],
674
- request=None,
675
- )
676
-
677
- print(f"Result: {result}")
678
-
679
- # Check expected content
680
- if test_case["expected_content"] is None:
681
- assert result is None or not getattr(result, "content", None), (
682
- f"Stage {i}: Expected no content, got {result}"
683
- )
684
- print("✓ No content output as expected")
685
- else:
686
- assert result is not None and hasattr(result, "content"), (
687
- f"Stage {i}: Expected content, got {result}"
688
- )
689
- assert result.content == test_case["expected_content"], (
690
- f"Stage {i}: Expected content {test_case['expected_content']}, got {result.content}"
691
- )
692
- print(f"✓ Content matches: {repr(result.content)}")
693
-
694
- print("✓ Streaming test with buffering completed successfully")
695
-
696
-
697
-def test_streaming_thinking_tag_buffering(minimax_tool_parser):
698
- """Test that tool calls within thinking tags are properly handled during streaming."""
699
- # Reset streaming state
700
- minimax_tool_parser.current_tool_name_sent = False
701
- minimax_tool_parser.prev_tool_call_arr = []
702
- minimax_tool_parser.current_tool_id = -1
703
- minimax_tool_parser.streamed_args_for_tool = []
704
- # Reset buffering state
705
- minimax_tool_parser.pending_buffer = ""
706
- minimax_tool_parser.in_thinking_tag = False
707
- minimax_tool_parser.thinking_depth = 0
708
-
709
- # Test scenario: tool calls within thinking tags should be ignored
710
- test_cases: list[dict[str, Any]] = [
711
- {
712
- "stage": "Start thinking",
713
- "previous": "",
714
- "current": "<think>I need to use a tool. <tool_calls>",
715
- "delta": "<think>I need to use a tool. <tool_calls>",
716
- "expected_content": "<think>I need to use a tool. <tool_calls>", # Should pass through as content
717
- },
718
- {
719
- "stage": "Tool call in thinking",
720
- "previous": "<think>I need to use a tool. <tool_calls>",
721
- "current": '<think>I need to use a tool. <tool_calls>\n{"name": "ignored_tool", "arguments": {"param": "value"}}\n</tool_calls>',
722
- "delta": '\n{"name": "ignored_tool", "arguments": {"param": "value"}}\n</tool_calls>',
723
- "expected_content": '\n{"name": "ignored_tool", "arguments": {"param": "value"}}\n</tool_calls>', # </tool_calls> should be preserved in thinking tags
724
- },
725
- {
726
- "stage": "Real tool call after thinking",
727
- "previous": '<think>I need to use a tool. <tool_calls>\n{"name": "ignored_tool", "arguments": {"param": "value"}}\n</tool_calls></think>',
728
- "current": '<think>I need to use a tool. <tool_calls>\n{"name": "ignored_tool", "arguments": {"param": "value"}}\n</tool_calls></think>\n<tool_calls>',
729
- "delta": "\n<tool_calls>",
730
- "expected_content": "\n", # Should output '\n' and suppress <tool_calls>
731
- },
732
- ]
733
-
734
- for i, test_case in enumerate(test_cases):
735
- print(f"\n--- Stage {i}: {test_case['stage']} ---")
736
- print(f"Previous: {repr(test_case['previous'])}")
737
- print(f"Current: {repr(test_case['current'])}")
738
- print(f"Delta: {repr(test_case['delta'])}")
739
-
740
- result = minimax_tool_parser.extract_tool_calls_streaming(
741
- previous_text=test_case["previous"],
742
- current_text=test_case["current"],
743
- delta_text=test_case["delta"],
744
- previous_token_ids=[],
745
- current_token_ids=[],
746
- delta_token_ids=[],
747
- request=None,
748
- )
749
-
750
- print(f"Result: {result}")
751
-
752
- # Check expected content
753
- if "expected_content" in test_case:
754
- if test_case["expected_content"] is None:
755
- assert result is None or not getattr(result, "content", None), (
756
- f"Stage {i}: Expected no content, got {result}"
757
- )
758
- else:
759
- assert result is not None and hasattr(result, "content"), (
760
- f"Stage {i}: Expected content, got {result}"
761
- )
762
- assert result.content == test_case["expected_content"], (
763
- f"Stage {i}: Expected content {test_case['expected_content']}, got {result.content}"
764
- )
765
- print(f"✓ Content matches: {repr(result.content)}")
766
-
767
- # Check tool calls
768
- if test_case.get("expected_tool_call"):
769
- assert (
770
- result is not None
771
- and hasattr(result, "tool_calls")
772
- and result.tool_calls
773
- ), f"Stage {i}: Expected tool call, got {result}"
774
-
775
- tool_call = result.tool_calls[0]
776
- assert tool_call.function.name == "real_tool", (
777
- f"Expected real_tool, got {tool_call.function.name}"
778
- )
779
- print(f"✓ Real tool call detected: {tool_call.function.name}")
780
-
781
- print("✓ Thinking tag buffering test completed successfully")
782
-
783
-
784
-def reset_streaming_state(minimax_tool_parser):
785
- """Helper function to properly reset the streaming state for MinimaxToolParser."""
786
- # Reset minimax-specific state
787
- minimax_tool_parser._reset_streaming_state()
788
-
789
- # Reset base class state (these should still be reset for compatibility)
790
- minimax_tool_parser.prev_tool_call_arr = []
791
- minimax_tool_parser.current_tool_id = -1
792
- minimax_tool_parser.current_tool_name_sent = False
793
- minimax_tool_parser.streamed_args_for_tool = []
794
-
795
-
796
-def test_streaming_complex_scenario_with_multiple_tools(minimax_tool_parser):
797
- """Test complex streaming scenario: tools inside <think> tags and multiple tool calls in one group."""
798
- # Reset streaming state
799
- reset_streaming_state(minimax_tool_parser)
800
-
801
- # Complex scenario: tools inside thinking tags and multiple tools in one group
802
- test_stages: list[dict[str, Any]] = [
803
- {
804
- "stage": "Initial content",
805
- "previous": "",
806
- "current": "Let me help you with this task.",
807
- "delta": "Let me help you with this task.",
808
- "expected_content": "Let me help you with this task.",
809
- "expected_tool_calls": 0,
810
- },
811
- {
812
- "stage": "Start thinking tag",
813
- "previous": "Let me help you with this task.",
814
- "current": "Let me help you with this task.<think>I need to analyze this situation first.",
815
- "delta": "<think>I need to analyze this situation first.",
816
- "expected_content": "<think>I need to analyze this situation first.",
817
- "expected_tool_calls": 0,
818
- },
819
- {
820
- "stage": "Tool call inside thinking tag starts",
821
- "previous": "Let me help you with this task.<think>I need to analyze this situation first.",
822
- "current": "Let me help you with this task.<think>I need to analyze this situation first.<tool_calls>",
823
- "delta": "<tool_calls>",
824
- "expected_content": "<tool_calls>", # Inside thinking tags, tool tags should be preserved as content
825
- "expected_tool_calls": 0,
826
- },
827
- {
828
- "stage": "Complete tool call inside thinking tag",
829
- "previous": "Let me help you with this task.<think>I need to analyze this situation first.<tool_calls>",
830
- "current": 'Let me help you with this task.<think>I need to analyze this situation first.<tool_calls>\n{"name": "internal_analysis", "arguments": {"query": "analyze situation"}}\n</tool_calls>',
831
- "delta": '\n{"name": "internal_analysis", "arguments": {"query": "analyze situation"}}\n</tool_calls>',
832
- "expected_content": '\n{"name": "internal_analysis", "arguments": {"query": "analyze situation"}}\n</tool_calls>',
833
- "expected_tool_calls": 0, # Tools inside thinking tags should be ignored
834
- },
835
- {
836
- "stage": "End thinking tag",
837
- "previous": 'Let me help you with this task.<think>I need to analyze this situation first.<tool_calls>\n{"name": "internal_analysis", "arguments": {"query": "analyze situation"}}\n</tool_calls>',
838
- "current": 'Let me help you with this task.<think>I need to analyze this situation first.<tool_calls>\n{"name": "internal_analysis", "arguments": {"query": "analyze situation"}}\n</tool_calls></think>',
839
- "delta": "</think>",
840
- "expected_content": "</think>",
841
- "expected_tool_calls": 0,
842
- },
843
- {
844
- "stage": "Multiple tools group starts",
845
- "previous": 'Let me help you with this task.<think>I need to analyze this situation first.<tool_calls>\n{"name": "internal_analysis", "arguments": {"query": "analyze situation"}}\n</tool_calls></think>',
846
- "current": 'Let me help you with this task.<think>I need to analyze this situation first.<tool_calls>\n{"name": "internal_analysis", "arguments": {"query": "analyze situation"}}\n</tool_calls></think>\nNow I need to get weather information and calculate area.<tool_calls>',
847
- "delta": "\nNow I need to get weather information and calculate area.<tool_calls>",
848
- "expected_content": "\nNow I need to get weather information and calculate area.", # <tool_calls> should be filtered
849
- "expected_tool_calls": 0,
850
- },
851
- {
852
- "stage": "First tool in group",
853
- "previous": 'Let me help you with this task.<think>I need to analyze this situation first.<tool_calls>\n{"name": "internal_analysis", "arguments": {"query": "analyze situation"}}\n</tool_calls></think>\nNow I need to get weather information and calculate area.<tool_calls>',
854
- "current": 'Let me help you with this task.<think>I need to analyze this situation first.<tool_calls>\n{"name": "internal_analysis", "arguments": {"query": "analyze situation"}}\n</tool_calls></think>\nNow I need to get weather information and calculate area.<tool_calls>\n{"name": "get_current_weather", "arguments": {"city": "Seattle", "state": "WA", "unit": "celsius"}}',
855
- "delta": '\n{"name": "get_current_weather", "arguments": {"city": "Seattle", "state": "WA", "unit": "celsius"}}',
856
- "expected_content": None, # No content should be output when tool call is in progress
857
- "expected_tool_calls": 1,
858
- "expected_tool_name": "get_current_weather",
859
- },
860
- {
861
- "stage": "Second tool in group",
862
- "previous": 'Let me help you with this task.<think>I need to analyze this situation first.<tool_calls>\n{"name": "internal_analysis", "arguments": {"query": "analyze situation"}}\n</tool_calls></think>\nNow I need to get weather information and calculate area.<tool_calls>\n{"name": "get_current_weather", "arguments": {"city": "Seattle", "state": "WA", "unit": "celsius"}}',
863
- "current": 'Let me help you with this task.<think>I need to analyze this situation first.<tool_calls>\n{"name": "internal_analysis", "arguments": {"query": "analyze situation"}}\n</tool_calls></think>\nNow I need to get weather information and calculate area.<tool_calls>\n{"name": "get_current_weather", "arguments": {"city": "Seattle", "state": "WA", "unit": "celsius"}}\n{"name": "calculate_area", "arguments": {"shape": "rectangle", "dimensions": {"width": 10, "height": 5}}}',
864
- "delta": '\n{"name": "calculate_area", "arguments": {"shape": "rectangle", "dimensions": {"width": 10, "height": 5}}}',
865
- "expected_content": None,
866
- "expected_tool_calls": 1,
867
- "expected_tool_name": "calculate_area",
868
- },
869
- {
870
- "stage": "Complete tool calls group",
871
- "previous": 'Let me help you with this task.<think>I need to analyze this situation first.<tool_calls>\n{"name": "internal_analysis", "arguments": {"query": "analyze situation"}}\n</tool_calls></think>\nNow I need to get weather information and calculate area.<tool_calls>\n{"name": "get_current_weather", "arguments": {"city": "Seattle", "state": "WA", "unit": "celsius"}}\n{"name": "calculate_area", "arguments": {"shape": "rectangle", "dimensions": {"width": 10, "height": 5}}}',
872
- "current": 'Let me help you with this task.<think>I need to analyze this situation first.<tool_calls>\n{"name": "internal_analysis", "arguments": {"query": "analyze situation"}}\n</tool_calls></think>\nNow I need to get weather information and calculate area.<tool_calls>\n{"name": "get_current_weather", "arguments": {"city": "Seattle", "state": "WA", "unit": "celsius"}}\n{"name": "calculate_area", "arguments": {"shape": "rectangle", "dimensions": {"width": 10, "height": 5}}}</tool_calls>',
873
- "delta": "</tool_calls>",
874
- "expected_content": None,
875
- "expected_tool_calls": 0,
876
- },
877
- ]
878
-
879
- tool_calls_count = 0
880
-
881
- for i, test_case in enumerate(test_stages):
882
- print(f"\n--- Stage {i}: {test_case['stage']} ---")
883
- print(
884
- f"Previous: {repr(test_case['previous'][:100])}{'...' if len(test_case['previous']) > 100 else ''}"
885
- )
886
- print(f"Current: {repr(test_case['current'][-100:])}")
887
- print(f"Delta: {repr(test_case['delta'])}")
888
-
889
- result = minimax_tool_parser.extract_tool_calls_streaming(
890
- previous_text=test_case["previous"],
891
- current_text=test_case["current"],
892
- delta_text=test_case["delta"],
893
- previous_token_ids=[],
894
- current_token_ids=[],
895
- delta_token_ids=[],
896
- request=None,
897
- )
898
-
899
- print(f"Result: {result}")
900
-
901
- # Check expected content
902
- if test_case["expected_content"] is None:
903
- assert result is None or not getattr(result, "content", None), (
904
- f"Stage {i}: Expected no content output, got {result}"
905
- )
906
- print("✓ No content output as expected")
907
- else:
908
- assert result is not None and hasattr(result, "content"), (
909
- f"Stage {i}: Expected content output, got {result}"
910
- )
911
- assert result.content == test_case["expected_content"], (
912
- f"Stage {i}: Expected content {repr(test_case['expected_content'])}, got {repr(result.content)}"
913
- )
914
- print(f"✓ Content matches: {repr(result.content)}")
915
-
916
- # Check tool calls
917
- expected_tool_calls = test_case["expected_tool_calls"]
918
- actual_tool_calls = (
919
- len(result.tool_calls)
920
- if result and hasattr(result, "tool_calls") and result.tool_calls
921
- else 0
922
- )
923
-
924
- if expected_tool_calls > 0:
925
- assert actual_tool_calls >= expected_tool_calls, (
926
- f"Stage {i}: Expected at least {expected_tool_calls} tool calls, got {actual_tool_calls}"
927
- )
928
-
929
- if "expected_tool_name" in test_case:
930
- # Find the tool call with the expected name
931
- found_tool_call = None
932
- for tool_call in result.tool_calls:
933
- if tool_call.function.name == test_case["expected_tool_name"]:
934
- found_tool_call = tool_call
935
- break
936
-
937
- assert found_tool_call is not None, (
938
- 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]}"
939
- )
940
- print(f"✓ Tool call correct: {found_tool_call.function.name}")
941
-
942
- # Ensure tools inside thinking tags are not called
943
- assert found_tool_call.function.name != "internal_analysis", (
944
- f"Stage {i}: Tool 'internal_analysis' inside thinking tags should not be called"
945
- )
946
-
947
- tool_calls_count += actual_tool_calls
948
- print(f"✓ Detected {actual_tool_calls} tool calls")
949
- else:
950
- assert actual_tool_calls == 0, (
951
- f"Stage {i}: Expected no tool calls, got {actual_tool_calls}"
952
- )
953
-
954
- # Verify overall results
955
- print("\n=== Test Summary ===")
956
- print(f"Total tool calls count: {tool_calls_count}")
957
- assert tool_calls_count >= 2, (
958
- f"Expected at least 2 valid tool calls (outside thinking tags), but got {tool_calls_count}"
959
- )
960
-
961
- print("✓ Complex streaming test completed:")
962
- print(" - ✓ Tools inside thinking tags correctly ignored")
963
- print(" - ✓ Two tool groups outside thinking tags correctly parsed")
964
- print(" - ✓ Content and tool call streaming correctly handled")
965
- print(" - ✓ Buffering mechanism works correctly")
966
-
967
-
968
-def test_streaming_character_by_character_output(minimax_tool_parser):
969
- """Test character-by-character streaming output to simulate real streaming scenarios."""
970
- # Reset streaming state
971
- reset_streaming_state(minimax_tool_parser)
972
-
973
- # Complete text that will be streamed character by character
974
- complete_text = """I'll help you with the weather analysis. <think>Let me think about this. <tool_calls>
975
-{"name": "internal_analysis", "arguments": {"type": "thinking"}}
976
-</tool_calls>This tool should be ignored.</think>
977
-
978
-Now I'll get the weather information for you. <tool_calls>
979
-{"name": "get_current_weather", "arguments": {"city": "Seattle", "state": "WA", "unit": "celsius"}}
980
-{"name": "calculate_area", "arguments": {"shape": "rectangle", "dimensions": {"width": 10, "height": 5}}}
981
-</tool_calls>Here are the results."""
982
-
983
- print("\n=== Starting character-by-character streaming test ===")
984
- print(f"Complete text length: {len(complete_text)} characters")
985
-
986
- # Track the streaming results
987
- content_fragments = []
988
- tool_calls_detected = []
989
-
990
- # Stream character by character
991
- for i in range(1, len(complete_text) + 1):
992
- current_text = complete_text[:i]
993
- previous_text = complete_text[: i - 1] if i > 1 else ""
994
- delta_text = complete_text[i - 1 : i]
995
-
996
- # Show progress every 50 characters
997
- if i % 50 == 0 or i == len(complete_text):
998
- print(f"Progress: {i}/{len(complete_text)} characters")
999
-
1000
- # Call the streaming parser
1001
- result = minimax_tool_parser.extract_tool_calls_streaming(
1002
- previous_text=previous_text,
1003
- current_text=current_text,
1004
- delta_text=delta_text,
1005
- previous_token_ids=[],
1006
- current_token_ids=[],
1007
- delta_token_ids=[],
1008
- request=None,
1009
- )
1010
-
1011
- # Collect results
1012
- if result is not None:
1013
- if hasattr(result, "content") and result.content:
1014
- content_fragments.append(result.content)
1015
- # Log important content fragments
1016
- if any(
1017
- keyword in result.content
1018
- for keyword in [
1019
- "<think>",
1020
- "</think>",
1021
- "<tool_calls>",
1022
- "</tool_calls>",
1023
- ]
1024
- ):
1025
- print(f" Char {i}: Content fragment: {repr(result.content)}")
1026
-
1027
- if hasattr(result, "tool_calls") and result.tool_calls:
1028
- for tool_call in result.tool_calls:
1029
- tool_info = {
1030
- "character_position": i,
1031
- "function_name": tool_call.function.name
1032
- if tool_call.function
1033
- else None,
1034
- "arguments": tool_call.function.arguments
1035
- if tool_call.function
1036
- else None,
1037
- }
1038
- tool_calls_detected.append(tool_info)
1039
- print(f" Char {i}: Tool call detected: {tool_call.function.name}")
1040
- if tool_call.function.arguments:
1041
- print(f" Arguments: {repr(tool_call.function.arguments)}")
1042
-
1043
- # Verify results
1044
- print("\n=== Streaming Test Results ===")
1045
- print(f"Total content fragments: {len(content_fragments)}")
1046
- print(f"Total tool calls detected: {len(tool_calls_detected)}")
1047
-
1048
- # Reconstruct content from fragments
1049
- reconstructed_content = "".join(content_fragments)
1050
- print(f"Reconstructed content length: {len(reconstructed_content)}")
1051
-
1052
- # Verify thinking tags content is preserved
1053
- assert "<think>" in reconstructed_content, (
1054
- "Opening thinking tag should be preserved in content"
1055
- )
1056
- assert "</think>" in reconstructed_content, (
1057
- "Closing thinking tag should be preserved in content"
1058
- )
1059
-
1060
- # Verify that tool calls inside thinking tags are NOT extracted as actual tool calls
1061
- thinking_tool_calls = [
1062
- tc for tc in tool_calls_detected if tc["function_name"] == "internal_analysis"
1063
- ]
1064
- assert len(thinking_tool_calls) == 0, (
1065
- f"Tool calls inside thinking tags should be ignored, but found: {thinking_tool_calls}"
1066
- )
1067
-
1068
- # Verify that real tool calls outside thinking tags ARE extracted
1069
- weather_tool_calls = [
1070
- tc for tc in tool_calls_detected if tc["function_name"] == "get_current_weather"
1071
- ]
1072
- area_tool_calls = [
1073
- tc for tc in tool_calls_detected if tc["function_name"] == "calculate_area"
1074
- ]
1075
- print(tool_calls_detected)
1076
- assert len(weather_tool_calls) > 0, (
1077
- "get_current_weather tool call should be detected"
1078
- )
1079
- assert len(area_tool_calls) > 0, "calculate_area tool call should be detected"
1080
-
1081
- # Verify tool call arguments are properly streamed
1082
- weather_args_found = any(
1083
- tc["arguments"] for tc in weather_tool_calls if tc["arguments"]
1084
- )
1085
- area_args_found = any(tc["arguments"] for tc in area_tool_calls if tc["arguments"])
1086
-
1087
- print(f"Weather tool call with arguments: {weather_args_found}")
1088
- print(f"Area tool call with arguments: {area_args_found}")
1089
-
1090
- # Verify content before and after tool calls
1091
- assert "I'll help you with the weather analysis." in reconstructed_content, (
1092
- "Initial content should be preserved"
1093
- )
1094
- assert "Here are the results." in reconstructed_content, (
1095
- "Final content should be preserved"
1096
- )
1097
-
1098
- # Verify that <tool_calls> and </tool_calls> tags are not included in the final content
1099
- # (they should be filtered out when not inside thinking tags)
1100
- content_outside_thinking = reconstructed_content
1101
- # Remove thinking tag content to check content outside
1102
- if "<think>" in content_outside_thinking and "</think>" in content_outside_thinking:
1103
- start_think = content_outside_thinking.find("<think>")
1104
- end_think = content_outside_thinking.find("</think>") + len("</think>")
1105
- content_outside_thinking = (
1106
- content_outside_thinking[:start_think]
1107
- + content_outside_thinking[end_think:]
1108
- )
1109
-
1110
- # Outside thinking tags, tool_calls tags should be filtered
1111
- tool_calls_in_content = content_outside_thinking.count("<tool_calls>")
1112
- assert tool_calls_in_content == 0, (
1113
- f"<tool_calls> tags should be filtered from content outside thinking tags, but found {tool_calls_in_content}"
1114
- )
1115
-
1116
- print("\n=== Character-by-character streaming test completed successfully ===")
1117
- print("✓ Tool calls inside thinking tags correctly ignored")
1118
- print("✓ Tool calls outside thinking tags correctly detected")
1119
- print("✓ Content properly streamed and reconstructed")
1120
- print("✓ Tool call tags properly filtered from content")
1121
- print("✓ Character-level streaming works correctly")
1122
-
1123
-
1124
-def test_streaming_character_by_character_simple_tool_call(minimax_tool_parser):
1125
- """Test character-by-character streaming for a simple tool call scenario."""
1126
- # Reset streaming state
1127
- reset_streaming_state(minimax_tool_parser)
1128
-
1129
- # Simple tool call text
1130
- simple_text = 'Let me check the weather. <tool_calls>\n{"name": "get_weather", "arguments": {"city": "NYC"}}\n</tool_calls>'
1131
-
1132
- print("\n=== Simple character-by-character test ===")
1133
- print(f"Text: {repr(simple_text)}")
1134
-
1135
- content_parts = []
1136
- tool_name_sent = False
1137
- tool_args_sent = False
1138
-
1139
- for i in range(1, len(simple_text) + 1):
1140
- current_text = simple_text[:i]
1141
- previous_text = simple_text[: i - 1] if i > 1 else ""
1142
- delta_text = simple_text[i - 1 : i]
1143
-
1144
- result = minimax_tool_parser.extract_tool_calls_streaming(
1145
- previous_text=previous_text,
1146
- current_text=current_text,
1147
- delta_text=delta_text,
1148
- previous_token_ids=[],
1149
- current_token_ids=[],
1150
- delta_token_ids=[],
1151
- request=None,
1152
- )
1153
-
1154
- if result:
1155
- if hasattr(result, "content") and result.content:
1156
- content_parts.append(result.content)
1157
- print(
1158
- f" Char {i} ({repr(delta_text)}): Content: {repr(result.content)}"
1159
- )
1160
-
1161
- if hasattr(result, "tool_calls") and result.tool_calls:
1162
- for tool_call in result.tool_calls:
1163
- if tool_call.function and tool_call.function.name:
1164
- tool_name_sent = True
1165
- print(f" Char {i}: Tool name: {tool_call.function.name}")
1166
- if tool_call.function and tool_call.function.arguments:
1167
- tool_args_sent = True
1168
- print(
1169
- f" Char {i}: Tool args: {repr(tool_call.function.arguments)}"
1170
- )
1171
-
1172
- # Verify basic expectations
1173
- reconstructed_content = "".join(content_parts)
1174
- print(f"Final reconstructed content: {repr(reconstructed_content)}")
1175
-
1176
- assert tool_name_sent, "Tool name should be sent during streaming"
1177
- assert tool_args_sent, "Tool arguments should be sent during streaming"
1178
- assert "Let me check the weather." in reconstructed_content, (
1179
- "Initial content should be preserved"
1180
- )
1181
-
1182
- print("✓ Simple character-by-character test passed")
1183
-
1184
-
1185
-def test_streaming_character_by_character_with_buffering(minimax_tool_parser):
1186
- """Test character-by-character streaming with edge cases that trigger buffering."""
1187
- # Reset streaming state
1188
- reset_streaming_state(minimax_tool_parser)
1189
-
1190
- # Text that includes potential buffering scenarios
1191
- buffering_text = 'Hello world<tool_calls>\n{"name": "test"}\n</tool_calls>done'
1192
-
1193
- print("\n=== Buffering character-by-character test ===")
1194
- print(f"Text: {repr(buffering_text)}")
1195
-
1196
- all_content = []
1197
-
1198
- for i in range(1, len(buffering_text) + 1):
1199
- current_text = buffering_text[:i]
1200
- previous_text = buffering_text[: i - 1] if i > 1 else ""
1201
- delta_text = buffering_text[i - 1 : i]
1202
-
1203
- result = minimax_tool_parser.extract_tool_calls_streaming(
1204
- previous_text=previous_text,
1205
- current_text=current_text,
1206
- delta_text=delta_text,
1207
- previous_token_ids=[],
1208
- current_token_ids=[],
1209
- delta_token_ids=[],
1210
- request=None,
1211
- )
1212
-
1213
- if result and hasattr(result, "content") and result.content:
1214
- all_content.append(result.content)
1215
- print(f" Char {i} ({repr(delta_text)}): {repr(result.content)}")
1216
-
1217
- final_content = "".join(all_content)
1218
- print(f"Final content: {repr(final_content)}")
1219
-
1220
- # The parser should handle the edge case where </tool_calls> appears before <tool_calls>
1221
- assert "Hello" in final_content, "Initial 'Hello' should be preserved"
1222
- assert "world" in final_content, (
1223
- "Content after false closing tag should be preserved"
1224
- )
1225
- assert "done" in final_content, "Final content should be preserved"
1226
-
1227
- print("✓ Buffering character-by-character test passed")
tests/v1/attention/test_attention_backends_selection.py
+3
-1
@@ -6,10 +6,12 @@ from types import SimpleNamespace
6
7
import pytest
8
9
+from vllm.model_executor.layers.mamba.linear.minimax_linear_attn import (
10
+ MiniMaxText01LinearAttention,
11
+)
12
from vllm.model_executor.layers.mamba.mamba_mixer import MambaMixer
13
from vllm.model_executor.layers.mamba.mamba_mixer2 import MambaMixer2
14
from vllm.model_executor.layers.mamba.short_conv import ShortConv
12
-from vllm.model_executor.models.minimax_text_01 import MiniMaxText01LinearAttention
15
from vllm.v1.attention.backends.linear_attn import LinearAttentionBackend
16
from vllm.v1.attention.backends.mamba1_attn import Mamba1AttentionBackend
17
from vllm.v1.attention.backends.mamba2_attn import Mamba2AttentionBackend
vllm/model_executor/models/minimax_text_01.py
deleted
-1000
@@ -1,1000 +0,0 @@
1
-# SPDX-License-Identifier: Apache-2.0
2
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
-"""Inference-only MiniMaxText01 model."""
4
-
5
-from collections.abc import Iterable
6
-from itertools import islice
7
-from typing import TYPE_CHECKING
8
-
9
-if TYPE_CHECKING:
10
- pass
11
-
12
-import regex as re
13
-import torch
14
-from torch import nn
15
-from transformers import MiniMaxConfig
16
-
17
-from vllm.compilation.decorators import support_torch_compile
18
-from vllm.config import CacheConfig, VllmConfig
19
-from vllm.distributed.parallel_state import (
20
- get_pp_group,
21
- get_tensor_model_parallel_rank,
22
- get_tensor_model_parallel_world_size,
23
-)
24
-from vllm.forward_context import get_forward_context
25
-from vllm.model_executor.layers.activation import SiluAndMul
26
-from vllm.model_executor.layers.attention import Attention
27
-from vllm.model_executor.layers.fused_moe import (
28
- FusedMoE,
29
-)
30
-from vllm.model_executor.layers.layernorm import RMSNorm
31
-from vllm.model_executor.layers.linear import (
32
- MergedColumnParallelLinear,
33
- QKVParallelLinear,
34
- ReplicatedLinear,
35
- RowParallelLinear,
36
-)
37
-from vllm.model_executor.layers.logits_processor import LogitsProcessor
38
-from vllm.model_executor.layers.mamba.linear.minimax_linear_attn import (
39
- MiniMaxText01LinearAttention,
40
-)
41
-from vllm.model_executor.layers.mamba.mamba_utils import (
42
- MambaStateCopyFunc,
43
- MambaStateCopyFuncCalculator,
44
- MambaStateDtypeCalculator,
45
- MambaStateShapeCalculator,
46
-)
47
-from vllm.model_executor.layers.quantization import QuantizationConfig
48
-from vllm.model_executor.layers.rotary_embedding import get_rope
49
-from vllm.model_executor.layers.vocab_parallel_embedding import (
50
- ParallelLMHead,
51
- VocabParallelEmbedding,
52
-)
53
-from vllm.model_executor.model_loader.weight_utils import default_weight_loader
54
-from vllm.model_executor.models.utils import maybe_prefix
55
-from vllm.sequence import IntermediateTensors
56
-from vllm.v1.attention.backend import AttentionMetadata
57
-
58
-from .interfaces import HasInnerState, IsHybrid
59
-from .utils import (
60
- AutoWeightsLoader,
61
- PPMissingLayer,
62
- is_pp_missing_parameter,
63
- make_layers,
64
-)
65
-
66
-
67
-def replace_weight_name(
68
- name: str, key: str = None, to: str = None, count: int = None, prefix: str = None
69
-) -> str:
70
- name = name.replace(key, to) if count is None else name.replace(key, to, count)
71
- return name
72
-
73
-
74
-def weight_loader_with_alias(alias: str):
75
- def wrapper(func: callable):
76
- def inner_func(
77
- param: torch.Tensor,
78
- loaded_weight: torch.Tensor,
79
- *args,
80
- prefix: str = None,
81
- **kwargs,
82
- ):
83
- value = func(param, loaded_weight, *args, **kwargs)
84
- return value
85
-
86
- return inner_func
87
-
88
- return wrapper
89
-
90
-
91
-class MiniMaxText01MLP(nn.Module):
92
- def __init__(
93
- self,
94
- hidden_size: int,
95
- intermediate_size: int,
96
- quant_config: QuantizationConfig | None = None,
97
- layer_idx: int = None,
98
- prefix: str = "mlp",
99
- ) -> None:
100
- super().__init__()
101
- self.layer_idx = layer_idx
102
-
103
- self.gate_up_proj = MergedColumnParallelLinear(
104
- hidden_size,
105
- [intermediate_size] * 2,
106
- bias=False,
107
- quant_config=quant_config,
108
- prefix=f"{prefix}.gate_up_proj",
109
- )
110
- self.down_proj = RowParallelLinear(
111
- intermediate_size,
112
- hidden_size,
113
- bias=False,
114
- quant_config=quant_config,
115
- prefix=f"{prefix}.down_proj",
116
- )
117
- self.act_fn = SiluAndMul()
118
- return
119
-
120
- def forward(self, x: torch.Tensor) -> torch.Tensor:
121
- gate_up, _ = self.gate_up_proj(x)
122
- x = self.act_fn(gate_up)
123
- x, _ = self.down_proj(x)
124
- return x
125
-
126
-
127
-class MiniMaxText01MoE(nn.Module):
128
- def __init__(
129
- self,
130
- num_experts: int,
131
- top_k: int,
132
- hidden_size: int,
133
- intermediate_size: int,
134
- params_dtype: torch.dtype | None = None,
135
- layer_idx: int = None,
136
- quant_config: QuantizationConfig | None = None,
137
- prefix: str = "moe",
138
- ) -> None:
139
- super().__init__()
140
-
141
- self.layer_idx = layer_idx
142
- self.tp_size = get_tensor_model_parallel_world_size()
143
- self.num_total_experts = num_experts
144
- self.top_k = top_k
145
- self.hidden_size = hidden_size
146
- self.intermediate_size = intermediate_size // self.tp_size
147
- self.quant_config = quant_config
148
-
149
- if params_dtype is None:
150
- params_dtype = torch.get_default_dtype()
151
- self.params_dtype = params_dtype
152
-
153
- self.gate = ReplicatedLinear(
154
- self.hidden_size,
155
- self.num_total_experts,
156
- bias=False,
157
- params_dtype=torch.float32,
158
- quant_config=None,
159
- prefix=f"{prefix}.gate",
160
- )
161
- self.gate.weight.weight_loader = MiniMaxText01MoE.gate_weight_loader
162
-
163
- self.experts = FusedMoE(
164
- num_experts=self.num_total_experts,
165
- top_k=self.top_k,
166
- hidden_size=self.hidden_size,
167
- intermediate_size=self.intermediate_size * self.tp_size,
168
- params_dtype=self.params_dtype,
169
- renormalize=True,
170
- quant_config=self.quant_config,
171
- tp_size=self.tp_size,
172
- prefix=f"{prefix}.experts",
173
- )
174
- return
175
-
176
- @staticmethod
177
- def gate_weight_loader(param: nn.Parameter, loaded_weight: torch.Tensor) -> None:
178
- assert param.size() == loaded_weight.size()
179
- param.data.copy_(loaded_weight.to(torch.float32))
180
- return
181
-
182
- def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
183
- num_tokens, hidden_size = hidden_states.shape
184
- hidden_states = hidden_states.view(-1, self.hidden_size)
185
- router_logits_fp32, _ = self.gate(hidden_states.to(torch.float32))
186
- final_hidden_states = self.experts(
187
- hidden_states, router_logits_fp32.to(hidden_states.dtype)
188
- )
189
- final_hidden = final_hidden_states.view(num_tokens, hidden_size)
190
- return final_hidden
191
-
192
-
193
-class MiniMaxText01Attention(nn.Module):
194
- def __init__(
195
- self,
196
- hidden_size: int,
197
- num_heads: int,
198
- head_dim: int,
199
- num_kv_heads: int,
200
- max_position: int = 4096 * 32,
201
- rope_parameters: dict | None = None,
202
- sliding_window: int | None = None,
203
- quant_config: QuantizationConfig | None = None,
204
- layer_idx: int = None,
205
- cache_config: CacheConfig | None = None,
206
- prefix: str = "mha",
207
- ) -> None:
208
- super().__init__()
209
- self.layer_idx = layer_idx
210
-
211
- self.hidden_size = hidden_size
212
- tp_size = get_tensor_model_parallel_world_size()
213
- self.total_num_heads = num_heads
214
- assert self.total_num_heads % tp_size == 0
215
- self.num_heads = self.total_num_heads // tp_size
216
- self.total_num_kv_heads = num_kv_heads
217
- if self.total_num_kv_heads >= tp_size:
218
- assert self.total_num_kv_heads % tp_size == 0
219
- else:
220
- assert tp_size % self.total_num_kv_heads == 0
221
- self.num_kv_heads = max(1, self.total_num_kv_heads // tp_size)
222
- self.head_dim = head_dim
223
-
224
- self.q_size = self.num_heads * self.head_dim
225
- self.kv_size = self.num_kv_heads * self.head_dim
226
- self.scaling = self.head_dim**-0.5
227
- self.sliding_window = sliding_window
228
- self.prefix = prefix
229
-
230
- self.qkv_proj = QKVParallelLinear(
231
- hidden_size,
232
- self.head_dim,
233
- self.total_num_heads,
234
- self.total_num_kv_heads,
235
- bias=False,
236
- quant_config=quant_config,
237
- prefix=f"{prefix}.qkv_proj",
238
- )
239
- self.o_proj = RowParallelLinear(
240
- self.total_num_heads * self.head_dim,
241
- hidden_size,
242
- bias=False,
243
- quant_config=quant_config,
244
- prefix=f"{prefix}.o_proj",
245
- )
246
- self.attn = Attention(
247
- self.num_heads,
248
- self.head_dim,
249
- self.scaling,
250
- num_kv_heads=self.num_kv_heads,
251
- cache_config=cache_config,
252
- quant_config=quant_config,
253
- prefix=f"{prefix}.attn",
254
- )
255
- self.rotary_emb = get_rope(
256
- head_size=self.head_dim,
257
- max_position=max_position,
258
- rope_parameters=rope_parameters,
259
- is_neox_style=True,
260
- dtype=torch.float32,
261
- )
262
- return
263
-
264
- def forward(
265
- self,
266
- hidden_states: torch.Tensor,
267
- output: torch.Tensor,
268
- positions: torch.Tensor,
269
- **kwargs,
270
- ) -> None:
271
- qkv, _ = self.qkv_proj(hidden_states)
272
- q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)
273
- q, k = self.rotary_emb(positions, q, k)
274
- attn_output = self.attn(q, k, v)
275
- output[:], _ = self.o_proj(attn_output)
276
-
277
-
278
-class MiniMaxText01DecoderLayer(nn.Module):
279
- def __init__(
280
- self,
281
- config: MiniMaxConfig,
282
- vllm_config: VllmConfig,
283
- expert_num: int = 1,
284
- layer_id: int = None,
285
- linear_layer_id: int | None = None,
286
- prefix: str = "decoder",
287
- ) -> None:
288
- self._ilayer = layer_id
289
- self._irank = get_tensor_model_parallel_rank()
290
- self.prefix = prefix
291
- super().__init__()
292
-
293
- self.hidden_size = config.hidden_size
294
- self.expert_num = expert_num
295
-
296
- head_dim = getattr(config, "head_dim", None)
297
- if head_dim is None:
298
- head_dim = config.hidden_size // config.num_attention_heads
299
- rotary_dim = getattr(config, "rotary_dim", head_dim)
300
- config.rope_parameters["partial_rotary_factor"] = rotary_dim / head_dim
301
- if hasattr(config, "max_model_len") and isinstance(config.max_model_len, int):
302
- max_position_embeddings = min(
303
- config.max_position_embeddings, config.max_model_len
304
- )
305
- if config.attention_type == 0:
306
- self.self_attn = MiniMaxText01LinearAttention(
307
- config,
308
- vllm_config,
309
- prefix=prefix,
310
- )
311
- elif config.attention_type == 1:
312
- self.self_attn = MiniMaxText01Attention(
313
- hidden_size=self.hidden_size,
314
- num_heads=config.num_attention_heads,
315
- head_dim=head_dim,
316
- num_kv_heads=config.num_key_value_heads,
317
- max_position=max_position_embeddings,
318
- rope_parameters=config.rope_parameters,
319
- sliding_window=config.sliding_window,
320
- quant_config=vllm_config.quant_config,
321
- layer_idx=self._ilayer,
322
- cache_config=vllm_config.cache_config,
323
- prefix=prefix,
324
- )
325
- else:
326
- raise ValueError(
327
- f"Unsupported attention_type {self.config.attention_type}: "
328
- f"should be 0 (linear) or 1 (full)."
329
- )
330
-
331
- if expert_num == 1:
332
- self.mlp = MiniMaxText01MLP(
333
- hidden_size=self.hidden_size,
334
- intermediate_size=config.intermediate_size,
335
- quant_config=vllm_config.quant_config,
336
- layer_idx=self._ilayer,
337
- prefix=prefix,
338
- )
339
- else:
340
- self.block_sparse_moe = MiniMaxText01MoE(
341
- num_experts=expert_num,
342
- top_k=config.num_experts_per_tok,
343
- hidden_size=config.hidden_size,
344
- intermediate_size=config.intermediate_size,
345
- layer_idx=self._ilayer,
346
- quant_config=vllm_config.quant_config,
347
- prefix=prefix,
348
- )
349
-
350
- self.input_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
351
- self.post_attention_layernorm = RMSNorm(
352
- config.hidden_size, eps=config.rms_norm_eps
353
- )
354
- if config.attention_type == 0:
355
- self.layernorm_attention_alpha = getattr(
356
- config,
357
- "layernorm_linear_attention_alpha",
358
- getattr(config, "linear_attn_alpha_factor", 1),
359
- )
360
- self.layernorm_attention_beta = getattr(
361
- config,
362
- "layernorm_linear_attention_beta",
363
- getattr(config, "linear_attn_beta_factor", 1),
364
- )
365
- else:
366
- self.layernorm_attention_alpha = getattr(
367
- config,
368
- "layernorm_full_attention_alpha",
369
- getattr(config, "full_attn_alpha_factor", 1),
370
- )
371
- self.layernorm_attention_beta = getattr(
372
- config,
373
- "layernorm_full_attention_beta",
374
- getattr(config, "full_attn_beta_factor", 1),
375
- )
376
- self.layernorm_mlp_alpha = getattr(
377
- config, "layernorm_mlp_alpha", getattr(config, "mlp_alpha_factor", 1)
378
- )
379
- self.layernorm_mlp_beta = getattr(
380
- config, "layernorm_mlp_beta", getattr(config, "mlp_beta_factor", 1)
381
- )
382
- self.postnorm = getattr(config, "postnorm", False)
383
- self.shared_moe = False
384
-
385
- shared_intermediate = getattr(config, "shared_intermediate_size", 0)
386
- if isinstance(shared_intermediate, list):
387
- shared_intermediate = (
388
- shared_intermediate[layer_id]
389
- if layer_id < len(shared_intermediate)
390
- else 0
391
- )
392
- if shared_intermediate > 0:
393
- self.shared_moe = True
394
- self.shared_mlp = MiniMaxText01MLP(
395
- hidden_size=self.hidden_size,
396
- intermediate_size=shared_intermediate,
397
- quant_config=vllm_config.quant_config,
398
- layer_idx=self._ilayer,
399
- prefix=prefix,
400
- )
401
- self.coefficient = ReplicatedLinear(
402
- self.hidden_size,
403
- 1,
404
- bias=False,
405
- quant_config=vllm_config.quant_config,
406
- params_dtype=torch.float32,
407
- )
408
- self.coefficient.weight.weight_loader = self.shared_moe_coefficient_loader
409
- self.shared_moe_mode = getattr(config, "shared_moe_mode", "softmax")
410
- return
411
-
412
- def forward(
413
- self,
414
- hidden_states: torch.Tensor,
415
- positions: torch.Tensor,
416
- attn_metadata: AttentionMetadata,
417
- residual: torch.Tensor | None,
418
- is_warmup: bool = False,
419
- **kwargs,
420
- ) -> tuple[torch.Tensor, torch.Tensor]:
421
- layernorm_input = hidden_states
422
- layernorm_output = self.input_layernorm(layernorm_input)
423
- residual = layernorm_output if self.postnorm else layernorm_input
424
- self_attention_output = torch.empty_like(layernorm_output)
425
- self.self_attn(
426
- hidden_states=layernorm_output,
427
- output=self_attention_output,
428
- positions=positions,
429
- )
430
-
431
- residual = residual * self.layernorm_attention_alpha
432
- self_attention_output = self_attention_output * self.layernorm_attention_beta
433
-
434
- layernorm_input = residual + self_attention_output
435
- layernorm_output = self.post_attention_layernorm(layernorm_input)
436
- residual = layernorm_output if self.postnorm else layernorm_input
437
-
438
- if self.expert_num == 1:
439
- hidden_states = self.mlp(layernorm_output)
440
- else:
441
- moe_layernorm_output = layernorm_output.clone()
442
- moe_hidden_states = self.block_sparse_moe(moe_layernorm_output)
443
- if self.shared_moe:
444
- before_moe_dtype = layernorm_output.dtype
445
- moe_hidden_fp32 = moe_hidden_states.to(torch.float32)
446
- output_mlp = self.shared_mlp(layernorm_output).to(torch.float32)
447
-
448
- coef, _ = self.coefficient(layernorm_output.to(torch.float32))
449
-
450
- if self.shared_moe_mode == "softmax":
451
- coef = torch.nn.functional.softmax(coef, dim=-1)
452
- hidden_states = moe_hidden_fp32 * (1 - coef) + output_mlp * coef
453
- elif self.shared_moe_mode == "sigmoid":
454
- coef = torch.nn.functional.sigmoid(coef)
455
- hidden_states = moe_hidden_fp32 * (1 - coef) + output_mlp * coef
456
-
457
- hidden_states = hidden_states.to(before_moe_dtype)
458
- else:
459
- hidden_states = moe_hidden_states
460
-
461
- residual = residual * self.layernorm_mlp_alpha
462
- hidden_states = hidden_states * self.layernorm_mlp_beta
463
-
464
- hidden_states = residual + hidden_states
465
-
466
- return hidden_states, None
467
-
468
- @staticmethod
469
- def shared_moe_coefficient_loader(
470
- param: torch.Tensor, loaded_weight: torch.Tensor
471
- ) -> None:
472
- assert param.size() == loaded_weight.size()
473
-
474
- param.data.copy_(loaded_weight.to(torch.float32))
475
- return
476
-
477
-
478
-@support_torch_compile
479
-class MiniMaxText01Model(nn.Module):
480
- def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
481
- super().__init__()
482
- config: MiniMaxConfig = vllm_config.model_config.hf_config
483
- scheduler_config = vllm_config.scheduler_config
484
- self.config = config
485
- self.CONCAT_FFN = True
486
-
487
- self.vocab_size = config.vocab_size
488
-
489
- self.decoder_attention_types = getattr(
490
- config, "attn_type_list", False
491
- ) or getattr(config, "decoder_attention_types", False)
492
- # The HF format uses "layer_types" instead of "attn_type_list"
493
- # where "linear_attention" is 0 and "full_attention" is 1
494
- if not self.decoder_attention_types and hasattr(config, "layer_types"):
495
- self.decoder_attention_types = []
496
- for layer_type in config.layer_types:
497
- if layer_type == "linear_attention":
498
- self.decoder_attention_types.append(0)
499
- elif layer_type == "full_attention":
500
- self.decoder_attention_types.append(1)
501
- else:
502
- raise ValueError(f"Unsupported layer type: {layer_type}")
503
- # Default to full attention
504
- if not self.decoder_attention_types:
505
- self.decoder_attention_types = [1] * config.num_hidden_layers
506
- self.num_layers = config.num_hidden_layers
507
-
508
- self._layer_barrier = False
509
- if get_pp_group().is_first_rank:
510
- self.embed_tokens = VocabParallelEmbedding(
511
- self.vocab_size,
512
- config.hidden_size,
513
- org_num_embeddings=self.vocab_size,
514
- )
515
- else:
516
- self.embed_tokens = PPMissingLayer()
517
-
518
- def layer_fn(prefix):
519
- layer_idx = int(prefix.split(".")[-1])
520
- layer_config = config
521
- layer_config.attention_type = self.decoder_attention_types[layer_idx]
522
- layer_config.layer_idx = layer_idx
523
-
524
- decoder_kwargs = {
525
- "layer_id": layer_idx,
526
- "vllm_config": vllm_config,
527
- }
528
-
529
- if layer_config.attention_type == 0:
530
- decoder_kwargs["linear_layer_id"] = sum(
531
- 1 for i in range(layer_idx) if self.decoder_attention_types[i] == 0
532
- )
533
- else:
534
- decoder_kwargs["linear_layer_id"] = None
535
-
536
- if hasattr(config, "num_local_experts") and isinstance(
537
- config.num_local_experts, list
538
- ):
539
- decoder_kwargs["expert_num"] = config.num_local_experts[layer_idx]
540
- elif hasattr(config, "num_local_experts") and isinstance(
541
- config.num_local_experts, int
542
- ):
543
- decoder_kwargs["expert_num"] = config.num_local_experts
544
- else:
545
- decoder_kwargs["expert_num"] = 1
546
-
547
- return MiniMaxText01DecoderLayer(
548
- layer_config, **decoder_kwargs, prefix=prefix
549
- )
550
-
551
- self.start_layer, self.end_layer, self.layers = make_layers(
552
- config.num_hidden_layers, layer_fn, prefix=f"{prefix}.layers"
553
- )
554
-
555
- linear_layer_nums = sum(
556
- 1
557
- for i in range(config.num_hidden_layers)
558
- if self.decoder_attention_types[i] == 0
559
- )
560
- max_slots_number = scheduler_config.max_num_seqs
561
- self.cache_shape = (
562
- linear_layer_nums,
563
- max_slots_number,
564
- config.num_attention_heads // get_tensor_model_parallel_world_size(),
565
- config.head_dim,
566
- config.head_dim,
567
- )
568
- _dummy = torch.zeros(1)
569
- self._dtype = _dummy.dtype
570
- del _dummy
571
-
572
- norm_kwargs = {}
573
- if hasattr(config, "rms_norm_eps"):
574
- norm_kwargs["eps"] = config.rms_norm_eps
575
- if get_pp_group().is_last_rank:
576
- self.norm = RMSNorm(config.hidden_size, **norm_kwargs)
577
- else:
578
- self.norm = PPMissingLayer()
579
- self.embed_scale = 1.0
580
- return
581
-
582
- def _clear_prefill_cache(
583
- self, attn_metadata, minimax_cache_tensors: torch.Tensor, **kwargs
584
- ):
585
- seq_to_slot_maps = {}
586
- seq_id_map = sum(list(kwargs["request_ids_to_seq_ids"].values()), [])
587
- for _, seq_to_slot_map in self.minimax_cache.cache_indices_mapping.items():
588
- seq_to_slot_maps.update(seq_to_slot_map)
589
-
590
- slots_to_clear = []
591
- for _prefill_id in range(getattr(attn_metadata, "num_prefills", 0)):
592
- if _prefill_id >= len(seq_id_map):
593
- break
594
- seq_id = seq_id_map[_prefill_id]
595
- if (
596
- attn_metadata.context_lens_tensor[_prefill_id] == 0
597
- and seq_id in seq_to_slot_maps
598
- ):
599
- slots_to_clear.append(seq_to_slot_maps[seq_id])
600
-
601
- if slots_to_clear:
602
- slots_tensor = torch.tensor(
603
- slots_to_clear, device=minimax_cache_tensors.device, dtype=torch.long
604
- )
605
- minimax_cache_tensors[:, slots_tensor, ...] = 0
606
-
607
- def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
608
- return self.embed_tokens(input_ids)
609
-
610
- def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
611
- params_dict = dict(self.named_parameters())
612
- loaded_params: set[str] = set()
613
-
614
- def which_layer(name: str) -> int:
615
- if "layers" in name:
616
- after_layer = name.split("layers")[-1]
617
- return int(after_layer.split(".")[1])
618
- return None
619
-
620
- def is_linear_attn_layer(layer_idx: int) -> bool:
621
- if layer_idx is None or layer_idx >= len(self.decoder_attention_types):
622
- return False
623
- return self.decoder_attention_types[layer_idx] == 0
624
-
625
- def is_moe_weight(name: str) -> bool:
626
- return "block_sparse_moe" in name and not name.endswith(".bias")
627
-
628
- def get_expert_id(param_name):
629
- pattern = r"layers\.\d+\.block_sparse_moe\.experts\.(\d+)\."
630
- match = re.search(pattern, param_name)
631
- if match:
632
- return match.group(1)
633
- return None
634
-
635
- def load_sparse_moe_weight(
636
- name: str, loaded_weight: torch.Tensor, self
637
- ) -> None:
638
- if isinstance(self.config.num_local_experts, list):
639
- expert_params_mapping = [
640
- (
641
- "w13_weight" if weight_name in ["w1", "w3"] else "w2_weight",
642
- f"experts.{expert_id}.{weight_name}.weight",
643
- expert_id,
644
- )
645
- for expert_id in range(max(self.config.num_local_experts))
646
- for weight_name in ["w1", "w2", "w3"]
647
- ]
648
- else:
649
- expert_params_mapping = [
650
- (
651
- "w13_scale" if weight_name in ["w1", "w3"] else "w2_scale",
652
- f"{expert_id}.{weight_name}.weight_scale",
653
- expert_id,
654
- weight_name,
655
- )
656
- for expert_id in range(self.config.num_local_experts)
657
- for weight_name in ["w1", "w2", "w3"]
658
- ] + [
659
- (
660
- "w13_weight" if weight_name in ["w1", "w3"] else "w2_weight",
661
- f"{expert_id}.{weight_name}.weight",
662
- expert_id,
663
- weight_name,
664
- )
665
- for expert_id in range(self.config.num_local_experts)
666
- for weight_name in ["w1", "w2", "w3"]
667
- ]
668
- for param_name, weight_name, expert_id, shard_id in expert_params_mapping:
669
- name_expert_id = get_expert_id(name)
670
- if name_expert_id is not None and int(name_expert_id) != int(expert_id):
671
- continue
672
- if weight_name not in name:
673
- continue
674
- name = name.replace(weight_name, param_name)
675
- if is_pp_missing_parameter(name, self):
676
- return
677
- param = params_dict[name]
678
- weight_loader = param.weight_loader
679
- weight_loader = weight_loader_with_alias(name)(weight_loader)
680
- weight_loader(
681
- param,
682
- loaded_weight,
683
- weight_name,
684
- expert_id=expert_id,
685
- shard_id=shard_id,
686
- )
687
- loaded_params.add(name)
688
- break
689
- else:
690
- if is_pp_missing_parameter(name, self):
691
- return
692
- param = params_dict[name]
693
- weight_loader = getattr(param, "weight_loader", default_weight_loader)
694
- weight_loader = weight_loader_with_alias(name)(weight_loader)
695
- weight_loader(param, loaded_weight)
696
- loaded_params.add(name)
697
- return
698
-
699
- def is_shared_mlp_weight(name: str) -> bool:
700
- return "shared_mlp" in name and not name.endswith(".bias")
701
-
702
- def load_shared_mlp_weight(
703
- name: str, loaded_weight: torch.Tensor, self
704
- ) -> None:
705
- if not self.CONCAT_FFN:
706
- if "gate_proj" in name:
707
- name = name.replace("gate_proj", "w1", 1)
708
- elif "up_proj" in name:
709
- name = name.replace("up_proj", "w3", 1)
710
- elif "down_proj" in name:
711
- name = name.replace("down_proj", "w2", 1)
712
- else:
713
- if "gate_proj" in name:
714
- name = name.replace("gate_proj", "gate_up_proj", 1)
715
- loaded_shard_id = 0
716
- elif "up_proj" in name:
717
- name = name.replace("up_proj", "gate_up_proj", 1)
718
- loaded_shard_id = 1
719
- if is_pp_missing_parameter(name, self):
720
- return
721
- param = params_dict[name]
722
- weight_loader = getattr(param, "weight_loader", default_weight_loader)
723
- weight_loader = weight_loader_with_alias(name)(weight_loader)
724
- if not self.CONCAT_FFN:
725
- weight_loader(param, loaded_weight)
726
- else:
727
- if "gate_up_proj" in name:
728
- weight_loader(param, loaded_weight, loaded_shard_id)
729
- elif "down_proj" in name:
730
- weight_loader(param, loaded_weight)
731
- else:
732
- raise AssertionError("MLP weight not in [gate_up_proj, down_proj]")
733
- loaded_params.add(name)
734
- return
735
-
736
- def is_mha_weight(name: str) -> bool:
737
- return "self_attn" in name and not name.endswith(".bias")
738
-
739
- def load_linear_attn_weight(
740
- name: str, loaded_weight: torch.Tensor, self
741
- ) -> None:
742
- if is_pp_missing_parameter(name, self):
743
- return
744
- param = params_dict[name]
745
-
746
- weight_loader = getattr(
747
- param, "weight_loader", MiniMaxText01LinearAttention.weight_direct_load
748
- )
749
- weight_loader = weight_loader_with_alias(name)(weight_loader)
750
- weight_loader(param, loaded_weight)
751
- loaded_params.add(name)
752
- return
753
-
754
- def load_flash_attn_weight(
755
- name: str, loaded_weight: torch.Tensor, self
756
- ) -> None:
757
- flash_mha_params_mapping = [
758
- ("qkv_proj", "q_proj", "q"),
759
- ("qkv_proj", "k_proj", "k"),
760
- ("qkv_proj", "v_proj", "v"),
761
- ("gate_up_proj", "gate_proj", 0),
762
- ("gate_up_proj", "up_proj", 1),
763
- ]
764
- for param_name, weight_name, shard_id in flash_mha_params_mapping:
765
- if weight_name not in name:
766
- continue
767
- name = name.replace(weight_name, param_name)
768
- if is_pp_missing_parameter(name, self):
769
- return
770
- param = params_dict[name]
771
- weight_loader = getattr(param, "weight_loader", default_weight_loader)
772
- weight_loader = weight_loader_with_alias(name)(weight_loader)
773
- weight_loader(param, loaded_weight, shard_id)
774
- loaded_params.add(name)
775
- break
776
- else:
777
- if is_pp_missing_parameter(name, self):
778
- return
779
- param = params_dict[name]
780
-
781
- weight_loader = getattr(param, "weight_loader", default_weight_loader)
782
- weight_loader = weight_loader_with_alias(name)(weight_loader)
783
- weight_loader(param, loaded_weight)
784
- loaded_params.add(name)
785
- return
786
-
787
- def is_layer_norm_weight(name: str) -> bool:
788
- return "norm" in name and not name.endswith(".bias") and name in params_dict
789
-
790
- def load_layer_norm_weight(
791
- name: str, loaded_weight: torch.Tensor, self
792
- ) -> None:
793
- if is_pp_missing_parameter(name, self):
794
- return
795
- param = params_dict[name]
796
- weight_loader = getattr(param, "weight_loader", default_weight_loader)
797
- weight_loader = weight_loader_with_alias(name)(weight_loader)
798
- weight_loader(param, loaded_weight)
799
- loaded_params.add(name)
800
- return
801
-
802
- def load_basic_weight(name: str, loaded_weight: torch.Tensor, self) -> None:
803
- if is_pp_missing_parameter(name, self):
804
- return
805
- param = params_dict[name]
806
- weight_loader = getattr(param, "weight_loader", default_weight_loader)
807
- weight_loader = weight_loader_with_alias(name)(weight_loader)
808
- weight_loader(param, loaded_weight)
809
- loaded_params.add(name)
810
- return
811
-
812
- for name, loaded_weight in weights:
813
- weight_at_layer = which_layer(name)
814
- if weight_at_layer and weight_at_layer >= len(self.decoder_attention_types):
815
- continue
816
-
817
- if is_layer_norm_weight(name):
818
- load_layer_norm_weight(name, loaded_weight, self)
819
- continue
820
- if is_mha_weight(name):
821
- if is_linear_attn_layer(weight_at_layer):
822
- load_linear_attn_weight(name, loaded_weight, self)
823
- else:
824
- load_flash_attn_weight(name, loaded_weight, self)
825
- continue
826
- if is_moe_weight(name):
827
- load_sparse_moe_weight(name, loaded_weight, self)
828
- continue
829
- if is_shared_mlp_weight(name):
830
- load_shared_mlp_weight(name, loaded_weight, self)
831
- continue
832
-
833
- if "rotary_emb.inv_freq" in name:
834
- continue
835
-
836
- load_basic_weight(name, loaded_weight, self)
837
- return loaded_params
838
-
839
- def forward(
840
- self,
841
- input_ids: torch.Tensor | None,
842
- positions: torch.Tensor,
843
- intermediate_tensors: IntermediateTensors | None = None,
844
- inputs_embeds: torch.Tensor | None = None,
845
- **kwargs,
846
- ) -> torch.Tensor | IntermediateTensors:
847
- forward_context = get_forward_context()
848
- attn_metadata = forward_context.attn_metadata
849
-
850
- if get_pp_group().is_first_rank:
851
- if inputs_embeds is None:
852
- hidden_states = self.embed_scale * self.embed_tokens(input_ids)
853
- else:
854
- hidden_states = inputs_embeds
855
- residual = None
856
- else:
857
- assert intermediate_tensors is not None
858
- hidden_states = intermediate_tensors["hidden_states"]
859
- residual = intermediate_tensors["residual"]
860
-
861
- for layer in islice(self.layers, self.start_layer, self.end_layer):
862
- hidden_states, residual = layer(
863
- hidden_states=hidden_states,
864
- positions=positions,
865
- attn_metadata=attn_metadata,
866
- residual=residual,
867
- )
868
- if not get_pp_group().is_last_rank:
869
- return IntermediateTensors(
870
- {"hidden_states": hidden_states, "residual": residual}
871
- )
872
- if residual is not None:
873
- hidden_states, _ = self.norm(hidden_states, residual)
874
- else:
875
- hidden_states = self.norm(hidden_states)
876
-
877
- return hidden_states
878
-
879
-
880
-class MiniMaxText01ForCausalLM(nn.Module, HasInnerState, IsHybrid):
881
- def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None:
882
- super().__init__()
883
- config = vllm_config.model_config.hf_config
884
-
885
- self.config = config
886
-
887
- if not hasattr(config, "sliding_window"):
888
- config.sliding_window = None
889
-
890
- self.CONCAT_FFN = True
891
-
892
- if hasattr(vllm_config.model_config, "max_model_len"):
893
- self.config.max_model_len = vllm_config.model_config.max_model_len
894
- self.model = MiniMaxText01Model(
895
- vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model")
896
- )
897
- if get_pp_group().is_last_rank:
898
- self.lm_head = ParallelLMHead(
899
- config.vocab_size,
900
- self.config.hidden_size,
901
- prefix=maybe_prefix(prefix, "lm_head"),
902
- )
903
-
904
- self.logits_processor = LogitsProcessor(
905
- config.vocab_size, self.config.vocab_size
906
- )
907
-
908
- else:
909
- self.lm_head = PPMissingLayer()
910
- self.lm_head.float()
911
- flash_layer_count = sum(
912
- 1 for attn_type in self.model.decoder_attention_types if attn_type == 1
913
- )
914
- self.kv_cache = [torch.tensor([]) for _ in range(flash_layer_count)]
915
- return
916
-
917
- def copy_inputs_before_cuda_graphs(self, input_buffers, **kwargs):
918
- return self.model.minimax_cache.copy_inputs_before_cuda_graphs(
919
- input_buffers, **kwargs
920
- )
921
-
922
- def get_seqlen_agnostic_capture_inputs(self, batch_size: int):
923
- return self.model.minimax_cache.get_seqlen_agnostic_capture_inputs(batch_size)
924
-
925
- def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
926
- return self.model.embed_input_ids(input_ids)
927
-
928
- def forward(
929
- self,
930
- input_ids: torch.Tensor | None,
931
- positions: torch.Tensor,
932
- intermediate_tensors: IntermediateTensors | None = None,
933
- inputs_embeds: torch.Tensor | None = None,
934
- **kwargs,
935
- ) -> torch.Tensor:
936
- hidden_states = self.model(
937
- input_ids, positions, intermediate_tensors, inputs_embeds, **kwargs
938
- )
939
-
940
- return hidden_states
941
-
942
- def compute_logits(self, hidden_states: torch.Tensor) -> torch.Tensor:
943
- logits = self.logits_processor(self.lm_head, hidden_states.float())
944
-
945
- return logits
946
-
947
- def make_empty_intermediate_tensors(
948
- self, batch_size: int, dtype: torch.dtype, device: torch.device
949
- ) -> IntermediateTensors:
950
- return IntermediateTensors(
951
- {
952
- "hidden_states": torch.zeros(
953
- (batch_size, self.config.hidden_size), dtype=dtype, device=device
954
- ),
955
- "residual": torch.zeros(
956
- (batch_size, self.config.hidden_size), dtype=dtype, device=device
957
- ),
958
- }
959
- )
960
-
961
- @classmethod
962
- def get_mamba_state_dtype_from_config(
963
- cls,
964
- vllm_config: "VllmConfig",
965
- ) -> tuple[torch.dtype, torch.dtype]:
966
- return MambaStateDtypeCalculator.linear_attention_state_dtype(
967
- vllm_config.model_config.dtype,
968
- vllm_config.cache_config.mamba_cache_dtype,
969
- )
970
-
971
- @classmethod
972
- def get_mamba_state_shape_from_config(
973
- cls,
974
- vllm_config: "VllmConfig",
975
- ) -> tuple[tuple[int, ...], ...]:
976
- """Calculate shape for MiniMaxText01LinearAttention cache.
977
-
978
- Args:
979
- vllm_config: vLLM config
980
-
981
- Returns:
982
- Tuple containing:
983
- - state_shape: Shape of the cache
984
- """
985
- parallel_config = vllm_config.parallel_config
986
- hf_config = vllm_config.model_config.hf_config
987
-
988
- return MambaStateShapeCalculator.linear_attention_state_shape(
989
- num_heads=hf_config.num_attention_heads,
990
- tp_size=parallel_config.tensor_parallel_size,
991
- head_dim=hf_config.head_dim,
992
- )
993
-
994
- @classmethod
995
- def get_mamba_state_copy_func(cls) -> tuple[MambaStateCopyFunc]:
996
- return MambaStateCopyFuncCalculator.linear_attention_state_copy_func()
997
-
998
- def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
999
- loader = AutoWeightsLoader(self)
1000
- return loader.load_weights(weights)
vllm/model_executor/models/minimax_vl_01.py
deleted
-385
@@ -1,385 +0,0 @@
1
-# SPDX-License-Identifier: Apache-2.0
2
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
-from collections.abc import Iterable, Mapping
4
-from typing import Annotated, Literal, TypeAlias
5
-
6
-import torch
7
-import torch.nn as nn
8
-from transformers import BatchFeature, PretrainedConfig
9
-from transformers.models.llava_next.modeling_llava_next import (
10
- get_anyres_image_grid_shape,
11
- unpad_image,
12
-)
13
-
14
-from vllm.config import VllmConfig
15
-from vllm.model_executor.layers.activation import get_act_fn
16
-from vllm.model_executor.layers.linear import ColumnParallelLinear, RowParallelLinear
17
-from vllm.model_executor.layers.quantization import QuantizationConfig
18
-from vllm.multimodal import MULTIMODAL_REGISTRY
19
-from vllm.multimodal.inputs import MultiModalFieldConfig
20
-from vllm.sequence import IntermediateTensors
21
-from vllm.utils.tensor_schema import TensorSchema, TensorShape
22
-
23
-from .clip import CLIPVisionModel
24
-from .interfaces import MultiModalEmbeddings, SupportsMultiModal, SupportsPP
25
-from .llava import (
26
- BaseLlavaMultiModalProcessor,
27
- LlavaDummyInputsBuilder,
28
- init_vision_tower_for_llava,
29
-)
30
-from .llava_next import LlavaNextProcessingInfo
31
-from .pixtral import PixtralHFVisionModel
32
-from .siglip import SiglipVisionModel
33
-from .utils import (
34
- AutoWeightsLoader,
35
- init_vllm_registered_model,
36
- maybe_prefix,
37
-)
38
-
39
-
40
-class MiniMaxVL01ImagePixelInputs(TensorSchema):
41
- """
42
- Dimensions:
43
- - bn: Batch size * number of images
44
- - np: Number of patches + 1
45
- - c: Number of channels (3)
46
- - h: Height
47
- - w: Width
48
-
49
- Note that `num_patches` may be different per batch and image,
50
- in which case the data is passed as a list instead of a batched tensor.
51
- """
52
-
53
- type: Literal["pixel_values"] = "pixel_values"
54
- pixel_values: Annotated[
55
- torch.Tensor | list[torch.Tensor],
56
- TensorShape("bn", "np", 3, "h", "w", dynamic_dims={"np", "h", "w"}),
57
- ]
58
-
59
- image_sizes: Annotated[torch.Tensor | None, TensorShape("bn", 2)]
60
- # This should be in `(height, width)` format.
61
-
62
-
63
-class MiniMaxVL01ImageEmbeddingInputs(TensorSchema):
64
- """
65
- Dimensions:
66
- - bn: Batch size * number of images
67
- - ifs: Image feature size
68
- - hs: Hidden size (must match language model backbone)
69
- """
70
-
71
- type: Literal["image_embeds"] = "image_embeds"
72
- data: Annotated[torch.Tensor, TensorShape("bn", "ifs", "hs")]
73
-
74
-
75
-MiniMaxVL01ImageInputs: TypeAlias = (
76
- MiniMaxVL01ImagePixelInputs | MiniMaxVL01ImageEmbeddingInputs
77
-)
78
-
79
-
80
-class MiniMaxVL01MultiModalProjector(nn.Module):
81
- def __init__(
82
- self,
83
- vision_hidden_size: int,
84
- text_hidden_size: int,
85
- projector_hidden_act: str,
86
- multimodal_projector_bias: bool,
87
- quant_config: QuantizationConfig | None = None,
88
- prefix: str = "",
89
- ):
90
- super().__init__()
91
-
92
- self.linear_1 = ColumnParallelLinear(
93
- vision_hidden_size,
94
- text_hidden_size,
95
- bias=multimodal_projector_bias,
96
- quant_config=quant_config,
97
- prefix=f"{prefix}.linear_1",
98
- )
99
- self.act = get_act_fn(projector_hidden_act)
100
- self.linear_2 = RowParallelLinear(
101
- text_hidden_size,
102
- text_hidden_size,
103
- bias=multimodal_projector_bias,
104
- quant_config=quant_config,
105
- prefix=f"{prefix}.linear_2",
106
- )
107
-
108
- def forward(self, image_features: torch.Tensor) -> torch.Tensor:
109
- hidden_states, _ = self.linear_1(image_features)
110
- hidden_states = self.act(hidden_states)
111
- hidden_states, _ = self.linear_2(hidden_states)
112
- return hidden_states
113
-
114
-
115
-class MiniMaxVL01DummyInputsBuilder(LlavaDummyInputsBuilder):
116
- pass
117
-
118
-
119
-class MiniMaxVL01ProcessingInfo(LlavaNextProcessingInfo):
120
- def get_hf_config(self): # Need to override the config type
121
- return self.ctx.get_hf_config(PretrainedConfig)
122
-
123
- def get_hf_processor(self, **kwargs: object):
124
- hf_processor = self.ctx.get_hf_processor(**kwargs)
125
- image_processor = hf_processor.image_processor
126
- image_processor.anyres_preprocess = image_processor.anyres_for_vllm_preprocess
127
-
128
- return hf_processor
129
-
130
- def get_supported_mm_limits(self) -> Mapping[str, int | None]:
131
- return {"image": None}
132
-
133
-
134
-class MiniMaxVL01MultiModalProcessor(
135
- BaseLlavaMultiModalProcessor[MiniMaxVL01ProcessingInfo]
136
-):
137
- def _call_hf_processor(
138
- self,
139
- prompt: str,
140
- mm_data: Mapping[str, object],
141
- mm_kwargs: Mapping[str, object],
142
- tok_kwargs: Mapping[str, object],
143
- ) -> BatchFeature:
144
- processed_outputs = super()._call_hf_processor(
145
- prompt=prompt,
146
- mm_data=mm_data,
147
- mm_kwargs=mm_kwargs,
148
- tok_kwargs=tok_kwargs,
149
- )
150
-
151
- pixel_values = processed_outputs.get("pixel_values")
152
- if pixel_values is not None:
153
- # Avoid padding since we need the output for each image to be
154
- # independent of other images for the cache to work correctly
155
- image_sizes = processed_outputs["image_sizes"]
156
- assert len(pixel_values) == len(image_sizes)
157
-
158
- processed_outputs["pixel_values"] = [
159
- p[:, :h, :w] for p, (h, w) in zip(pixel_values, image_sizes)
160
- ]
161
-
162
- return processed_outputs
163
-
164
- def _get_mm_fields_config(
165
- self,
166
- hf_inputs: BatchFeature,
167
- hf_processor_mm_kwargs: Mapping[str, object],
168
- ) -> Mapping[str, MultiModalFieldConfig]:
169
- return {
170
- "pixel_values": MultiModalFieldConfig.batched("image"),
171
- "image_sizes": MultiModalFieldConfig.batched("image"),
172
- "image_embeds": MultiModalFieldConfig.batched("image"),
173
- }
174
-
175
-
176
-@MULTIMODAL_REGISTRY.register_processor(
177
- MiniMaxVL01MultiModalProcessor,
178
- info=MiniMaxVL01ProcessingInfo,
179
- dummy_inputs=MiniMaxVL01DummyInputsBuilder,
180
-)
181
-class MiniMaxVL01ForConditionalGeneration(nn.Module, SupportsMultiModal, SupportsPP):
182
- packed_modules_mapping = {
183
- "qkv_proj": ["q_proj", "k_proj", "v_proj"],
184
- "gate_up_proj": ["gate_proj", "up_proj"],
185
- }
186
-
187
- @classmethod
188
- def get_placeholder_str(cls, modality: str, i: int) -> str | None:
189
- if modality.startswith("image"):
190
- return "<image>"
191
-
192
- raise ValueError("Only image modality is supported")
193
-
194
- def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None:
195
- super().__init__()
196
-
197
- config = vllm_config.model_config.hf_config
198
- quant_config = vllm_config.quant_config
199
- multimodal_config = vllm_config.model_config.multimodal_config
200
-
201
- self.config = config
202
- self.multimodal_config = multimodal_config
203
-
204
- with self._mark_tower_model(vllm_config, "image"):
205
- self.vision_tower = init_vision_tower_for_llava(
206
- config,
207
- quant_config=quant_config,
208
- require_post_norm=False,
209
- prefix=maybe_prefix(prefix, "vision_tower"),
210
- )
211
- self.multi_modal_projector = MiniMaxVL01MultiModalProjector(
212
- vision_hidden_size=config.vision_config.hidden_size,
213
- text_hidden_size=config.text_config.hidden_size,
214
- projector_hidden_act=config.projector_hidden_act,
215
- multimodal_projector_bias=True,
216
- quant_config=quant_config,
217
- prefix=maybe_prefix(prefix, "multi_modal_projector"),
218
- )
219
- self.image_newline = nn.Parameter(
220
- torch.empty(config.text_config.hidden_size)
221
- )
222
-
223
- with self._mark_language_model(vllm_config):
224
- self.language_model = init_vllm_registered_model(
225
- vllm_config=vllm_config,
226
- hf_config=config.text_config,
227
- prefix=maybe_prefix(prefix, "language_model"),
228
- )
229
-
230
- self.vision_feature_layer = config.vision_feature_layer
231
- self.vocab_size = config.text_config.vocab_size
232
- self.pad_token_id = -1
233
- if self.config.text_config.pad_token_id is not None:
234
- self.pad_token_id = self.config.text_config.pad_token_id
235
-
236
- self.make_empty_intermediate_tensors = (
237
- self.language_model.make_empty_intermediate_tensors
238
- )
239
-
240
- def _image_pixels_to_features(
241
- self,
242
- vision_tower: CLIPVisionModel | SiglipVisionModel | PixtralHFVisionModel,
243
- pixel_values: torch.Tensor | list[torch.Tensor],
244
- ) -> torch.Tensor | tuple[torch.Tensor, ...]:
245
- # NOTE: we skip the step to select the vision feature layer since
246
- # this is already done inside the vision tower
247
- feature_select_strategy = self.config.vision_feature_select_strategy
248
- return tuple(
249
- vision_tower(p, feature_select_strategy=feature_select_strategy)
250
- for p in pixel_values
251
- )
252
-
253
- # adapted from https://huggingface.co/MiniMaxAI/MiniMax-VL-01/blob/main/modeling_minimax_vl_01.py#L616-L631
254
- def pack_image_features(
255
- self, image_features: list[torch.Tensor], image_sizes: torch.Tensor
256
- ):
257
- new_image_features = []
258
- for image_idx, image_feature in enumerate(image_features):
259
- if image_feature.shape[0] > 1:
260
- base_image_feature = image_feature[0]
261
- image_feature = image_feature[1:]
262
- height = width = (
263
- self.config.vision_config.image_size
264
- // self.config.vision_config.patch_size
265
- )
266
- if height * width != base_image_feature.shape[0]:
267
- raise ValueError(
268
- "The number of patches is not consistent with the image size."
269
- )
270
- num_patch_height, num_patch_width = get_anyres_image_grid_shape(
271
- image_sizes[image_idx],
272
- self.config.image_grid_pinpoints,
273
- self.config.vision_config.image_size,
274
- )
275
-
276
- image_feature = image_feature.view(
277
- num_patch_height, num_patch_width, height, width, -1
278
- )
279
- image_feature = image_feature.permute(4, 0, 2, 1, 3).contiguous()
280
- image_feature = image_feature.flatten(1, 2).flatten(2, 3)
281
- image_feature = unpad_image(image_feature, image_sizes[image_idx])
282
-
283
- image_feature = torch.cat(
284
- (
285
- image_feature,
286
- self.image_newline[:, None, None]
287
- .expand(*image_feature.shape[:-1], 1)
288
- .to(image_feature.dtype),
289
- ),
290
- dim=-1,
291
- )
292
- image_feature = image_feature.flatten(1, 2).transpose(0, 1)
293
- image_feature = torch.cat((base_image_feature, image_feature), dim=0)
294
- else:
295
- image_feature = image_feature[0]
296
- image_feature = torch.cat(
297
- (image_feature, self.image_newline[None].to(image_feature)), dim=0
298
- )
299
- new_image_features.append(image_feature)
300
- return new_image_features
301
-
302
- def _process_image_pixels(
303
- self,
304
- inputs: MiniMaxVL01ImagePixelInputs,
305
- ) -> torch.Tensor | tuple[torch.Tensor, ...]:
306
- pixel_values = inputs["pixel_values"]
307
- return self._image_pixels_to_features(self.vision_tower, pixel_values)
308
-
309
- def _process_image_input(
310
- self,
311
- image_input: MiniMaxVL01ImageInputs,
312
- ) -> torch.Tensor | tuple[torch.Tensor, ...]:
313
- if image_input["type"] == "image_embeds":
314
- return image_input["data"]
315
-
316
- image_features = self._process_image_pixels(image_input)
317
-
318
- if isinstance(image_features, torch.Tensor):
319
- return self.multi_modal_projector(image_features)
320
-
321
- feature_sizes = [image_feature.shape[0] for image_feature in image_features]
322
-
323
- image_embeds = self.multi_modal_projector(torch.cat(image_features))
324
- image_embeds = torch.split(image_embeds, feature_sizes)
325
- image_sizes = image_input.get("image_sizes")
326
- return self.pack_image_features(image_embeds, image_sizes)
327
-
328
- def _parse_and_validate_image_input(
329
- self, **kwargs: object
330
- ) -> MiniMaxVL01ImageInputs | None:
331
- pixel_values = kwargs.pop("pixel_values", None)
332
- image_sizes = kwargs.pop("image_sizes", None)
333
- image_embeds = kwargs.pop("image_embeds", None)
334
-
335
- if pixel_values is None and image_embeds is None:
336
- return None
337
-
338
- if pixel_values is not None and image_sizes is not None:
339
- return MiniMaxVL01ImagePixelInputs(
340
- type="pixel_values",
341
- pixel_values=pixel_values,
342
- image_sizes=image_sizes,
343
- )
344
-
345
- if image_embeds is not None:
346
- return MiniMaxVL01ImageEmbeddingInputs(
347
- type="image_embeds",
348
- data=image_embeds,
349
- )
350
-
351
- raise AssertionError("This line should be unreachable.")
352
-
353
- def embed_multimodal(self, **kwargs: object) -> MultiModalEmbeddings:
354
- image_input = self._parse_and_validate_image_input(**kwargs)
355
- if image_input is None:
356
- return []
357
-
358
- return self._process_image_input(image_input)
359
-
360
- def forward(
361
- self,
362
- input_ids: torch.Tensor | None,
363
- positions: torch.Tensor,
364
- intermediate_tensors: IntermediateTensors | None = None,
365
- inputs_embeds: torch.Tensor | None = None,
366
- **kwargs: object,
367
- ) -> torch.Tensor | IntermediateTensors:
368
- if intermediate_tensors is not None:
369
- inputs_embeds = None
370
-
371
- hidden_states = self.language_model.model(
372
- input_ids, positions, intermediate_tensors, inputs_embeds=inputs_embeds
373
- )
374
-
375
- return hidden_states
376
-
377
- def compute_logits(
378
- self,
379
- hidden_states: torch.Tensor,
380
- ) -> torch.Tensor | None:
381
- return self.language_model.compute_logits(hidden_states)
382
-
383
- def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
384
- loader = AutoWeightsLoader(self)
385
- return loader.load_weights(weights)
vllm/model_executor/models/registry.py
+4
-7
@@ -159,9 +159,6 @@ _TEXT_GENERATION_MODELS = {
159
"MellumForCausalLM": ("mellum", "MellumForCausalLM"),
160
"MiniCPMForCausalLM": ("minicpm", "MiniCPMForCausalLM"),
161
"MiniCPM3ForCausalLM": ("minicpm3", "MiniCPM3ForCausalLM"),
162
- "MiniMaxForCausalLM": ("minimax_text_01", "MiniMaxText01ForCausalLM"),
163
- "MiniMaxText01ForCausalLM": ("minimax_text_01", "MiniMaxText01ForCausalLM"),
164
- "MiniMaxM1ForCausalLM": ("minimax_text_01", "MiniMaxText01ForCausalLM"),
162
"MiniMaxM2ForCausalLM": ("minimax_m2", "MiniMaxM2ForCausalLM"),
163
"MiniMaxM3SparseForCausalLM": (
164
"vllm.models.minimax_m3",
@@ -490,10 +487,6 @@ _MULTIMODAL_MODELS = {
487
"vllm.models.minimax_m3",
488
"MiniMaxM3SparseForConditionalGeneration",
489
),
493
- "MiniMaxVL01ForConditionalGeneration": (
494
- "minimax_vl_01",
495
- "MiniMaxVL01ForConditionalGeneration",
496
- ),
490
"MiniCPMO": ("minicpmo", "MiniCPMO"),
491
"MiniCPMV": ("minicpmv", "MiniCPMV"),
492
"MiniCPMV4_6ForConditionalGeneration": (
@@ -735,6 +728,10 @@ _PREVIOUSLY_SUPPORTED_MODELS = {
728
"XverseForCausalLM": "0.23.0",
729
"Dots1ForCausalLM": "0.23.0",
730
"BambaForCausalLM": "0.23.0",
731
+ "MiniMaxForCausalLM": "0.23.0",
732
+ "MiniMaxText01ForCausalLM": "0.23.0",
733
+ "MiniMaxM1ForCausalLM": "0.23.0",
734
+ "MiniMaxVL01ForConditionalGeneration": "0.23.0",
735
}
736
737
_OOT_SUPPORTED_MODELS = {
vllm/tool_parsers/__init__.py
-4
@@ -130,10 +130,6 @@ _TOOL_PARSERS_TO_REGISTER = {
130
"minimax_m3_tool_parser",
131
"MinimaxM3ToolParser",
132
),
133
- "minimax": (
134
- "minimax_tool_parser",
135
- "MinimaxToolParser",
136
- ),
133
"minicpm5": (
134
"minicpm5xml_tool_parser",
135
"MiniCPM5XMLToolParser",
vllm/tool_parsers/minimax_tool_parser.py
deleted
-852
@@ -1,852 +0,0 @@
1
-# SPDX-License-Identifier: Apache-2.0
2
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
-
4
-import json
5
-from collections.abc import Sequence
6
-from typing import Any
7
-
8
-import regex as re
9
-
10
-from vllm.entrypoints.chat_utils import make_tool_call_id
11
-from vllm.entrypoints.openai.chat_completion.protocol import (
12
- ChatCompletionRequest,
13
-)
14
-from vllm.entrypoints.openai.engine.protocol import (
15
- DeltaFunctionCall,
16
- DeltaMessage,
17
- DeltaToolCall,
18
- ExtractedToolCallInformation,
19
- FunctionCall,
20
- ToolCall,
21
-)
22
-from vllm.logger import init_logger
23
-from vllm.tokenizers import TokenizerLike
24
-from vllm.tool_parsers.abstract_tool_parser import (
25
- Tool,
26
- ToolParser,
27
-)
28
-from vllm.tool_parsers.utils import extract_intermediate_diff
29
-
30
-logger = init_logger(__name__)
31
-
32
-
33
-class MinimaxToolParser(ToolParser):
34
- def __init__(self, tokenizer: TokenizerLike, tools: list[Tool] | None = None):
35
- super().__init__(tokenizer, tools)
36
-
37
- # Initialize streaming state for tracking tool call progress
38
- self.streaming_state: dict[str, Any] = {
39
- "current_tool_index": -1, # Index of current tool being processed
40
- "tool_ids": [], # List of tool call IDs
41
- "sent_tools": [], # List of tools that have been sent
42
- }
43
-
44
- # Define tool call tokens and patterns
45
- self.tool_call_start_token = "<tool_calls>"
46
- self.tool_call_end_token = "</tool_calls>"
47
- self.tool_call_regex = re.compile(
48
- r"<tool_calls>(.*?)</tool_calls>|<tool_calls>(.*)", re.DOTALL
49
- )
50
- self.thinking_tag_pattern = r"<think>(.*?)</think>"
51
- self.tool_name_pattern = re.compile(r'"name":\s*"([^"]+)"')
52
- self.tool_args_pattern = re.compile(r'"arguments":\s*')
53
-
54
- # Buffer for handling partial tool calls during streaming
55
- self.pending_buffer = ""
56
- self.in_thinking_tag = False
57
-
58
- if not self.model_tokenizer:
59
- raise ValueError(
60
- "The model tokenizer must be passed to the ToolParser "
61
- "constructor during construction."
62
- )
63
-
64
- # Get token IDs for tool call start/end tokens
65
- self.tool_call_start_token_id = self.vocab.get(self.tool_call_start_token)
66
- self.tool_call_end_token_id = self.vocab.get(self.tool_call_end_token)
67
-
68
- if self.tool_call_start_token_id is None or self.tool_call_end_token_id is None:
69
- logger.warning(
70
- "Minimax Tool parser could not locate tool call start/end "
71
- "tokens in the tokenizer. Falling back to string matching."
72
- )
73
-
74
- def preprocess_model_output(self, model_output: str) -> str:
75
- """
76
- Preprocess model output by removing tool calls from thinking tags.
77
-
78
- Args:
79
- model_output: Raw model output string
80
-
81
- Returns:
82
- Preprocessed model output with tool calls removed from thinking tags
83
- """
84
-
85
- def remove_tool_calls_from_think(match):
86
- think_content = match.group(1)
87
- cleaned_content = re.sub(
88
- r"<tool_calls>.*?</tool_calls>", "", think_content, flags=re.DOTALL
89
- )
90
- return f"<think>{cleaned_content}</think>"
91
-
92
- return re.sub(
93
- self.thinking_tag_pattern,
94
- remove_tool_calls_from_think,
95
- model_output,
96
- flags=re.DOTALL,
97
- )
98
-
99
- def _clean_duplicate_braces(self, args_text: str) -> str:
100
- """
101
- Clean duplicate closing braces from arguments text.
102
-
103
- Args:
104
- args_text: Raw arguments text
105
-
106
- Returns:
107
- Cleaned arguments text with proper JSON formatting
108
- """
109
- args_text = args_text.strip()
110
- if not args_text:
111
- return args_text
112
-
113
- try:
114
- json.loads(args_text)
115
- return args_text
116
- except json.JSONDecodeError:
117
- pass
118
-
119
- while args_text.endswith("}}"):
120
- candidate = args_text[:-1]
121
- try:
122
- json.loads(candidate)
123
- return candidate
124
- except json.JSONDecodeError:
125
- args_text = candidate
126
-
127
- return args_text
128
-
129
- def _clean_delta_braces(self, delta_text: str) -> str:
130
- """
131
- Clean delta text by removing excessive closing braces.
132
-
133
- Args:
134
- delta_text: Delta text to clean
135
-
136
- Returns:
137
- Cleaned delta text
138
- """
139
- if not delta_text:
140
- return delta_text
141
-
142
- delta_stripped = delta_text.strip()
143
-
144
- if delta_stripped and all(c in "}\n\r\t " for c in delta_stripped):
145
- brace_count = delta_stripped.count("}")
146
- if brace_count > 1:
147
- return "}\n" if delta_text.endswith("\n") else "}"
148
-
149
- return delta_text
150
-
151
- def extract_tool_calls(
152
- self,
153
- model_output: str,
154
- request: ChatCompletionRequest,
155
- ) -> ExtractedToolCallInformation:
156
- """
157
- Extract tool calls from model output for non-streaming mode.
158
-
159
- Args:
160
- model_output: Complete model output
161
- request: Chat completion request
162
-
163
- Returns:
164
- ExtractedToolCallInformation containing tool calls and content
165
- """
166
- processed_output = self.preprocess_model_output(model_output)
167
-
168
- if self.tool_call_start_token not in processed_output:
169
- return ExtractedToolCallInformation(
170
- tools_called=False, tool_calls=[], content=model_output
171
- )
172
-
173
- try:
174
- function_call_tuples = self.tool_call_regex.findall(processed_output)
175
-
176
- raw_function_calls = []
177
- for match in function_call_tuples:
178
- tool_call_content = match[0] if match[0] else match[1]
179
- if tool_call_content.strip():
180
- lines = tool_call_content.strip().split("\n")
181
- for line in lines:
182
- line = line.strip()
183
- if line and line.startswith("{") and line.endswith("}"):
184
- try:
185
- parsed_call = json.loads(line)
186
- raw_function_calls.append(parsed_call)
187
- except json.JSONDecodeError:
188
- continue
189
-
190
- tool_calls = []
191
- for function_call in raw_function_calls:
192
- if "name" in function_call and "arguments" in function_call:
193
- tool_calls.append(
194
- ToolCall(
195
- type="function",
196
- function=FunctionCall(
197
- name=function_call["name"],
198
- arguments=json.dumps(
199
- function_call["arguments"], ensure_ascii=False
200
- ),
201
- ),
202
- )
203
- )
204
-
205
- processed_pos = processed_output.find(self.tool_call_start_token)
206
- if processed_pos != -1:
207
- processed_content = processed_output[:processed_pos].strip()
208
-
209
- if processed_content:
210
- lines = processed_content.split("\n")
211
- for line in reversed(lines):
212
- line = line.strip()
213
- if line:
214
- pos = model_output.find(line)
215
- if pos != -1:
216
- content = model_output[: pos + len(line)]
217
- break
218
- else:
219
- content = ""
220
- else:
221
- content = ""
222
- else:
223
- content = model_output
224
-
225
- return ExtractedToolCallInformation(
226
- tools_called=len(tool_calls) > 0,
227
- tool_calls=tool_calls,
228
- content=content.strip() if content.strip() else None,
229
- )
230
-
231
- except Exception:
232
- logger.exception(
233
- "An unexpected error occurred during tool call extraction."
234
- )
235
- return ExtractedToolCallInformation(
236
- tools_called=False, tool_calls=[], content=model_output
237
- )
238
-
239
- def _update_thinking_state(self, text: str) -> None:
240
- """
241
- Update the thinking tag state based on text content.
242
-
243
- Args:
244
- text: Text to analyze for thinking tags
245
- """
246
- open_count = text.count("<think>")
247
- close_count = text.count("</think>")
248
- self.in_thinking_tag = open_count > close_count or (
249
- open_count == close_count and text.endswith("</think>")
250
- )
251
-
252
- def _is_potential_tag_start(self, text: str) -> bool:
253
- """
254
- Check if text might be the start of a tool call tag.
255
-
256
- Args:
257
- text: Text to check
258
-
259
- Returns:
260
- True if text could be the start of a tool call tag
261
- """
262
- for tag in [self.tool_call_start_token, self.tool_call_end_token]:
263
- if any(
264
- tag.startswith(text[-i:])
265
- for i in range(1, min(len(text) + 1, len(tag)))
266
- ):
267
- return True
268
- return False
269
-
270
- def _should_buffer_content(self, delta_text: str) -> bool:
271
- """
272
- Determine if content should be buffered for later processing.
273
-
274
- Args:
275
- delta_text: Delta text to check
276
-
277
- Returns:
278
- True if content should be buffered
279
- """
280
- if self.in_thinking_tag:
281
- return False
282
- return bool(
283
- self.pending_buffer
284
- or self.tool_call_start_token in delta_text
285
- or self.tool_call_end_token in delta_text
286
- or delta_text.startswith("<")
287
- )
288
-
289
- def _split_content_for_buffering(self, delta_text: str) -> tuple[str, str]:
290
- """
291
- Split delta text into safe content and potential tag content.
292
-
293
- Args:
294
- delta_text: Delta text to split
295
-
296
- Returns:
297
- Tuple of (safe_content, potential_tag_content)
298
- """
299
- if self.in_thinking_tag:
300
- return delta_text, ""
301
-
302
- for tag in [self.tool_call_start_token, self.tool_call_end_token]:
303
- for i in range(1, len(tag)):
304
- tag_prefix = tag[:i]
305
- pos = delta_text.rfind(tag_prefix)
306
- if pos != -1 and tag.startswith(delta_text[pos:]):
307
- return delta_text[:pos], delta_text[pos:]
308
- return delta_text, ""
309
-
310
- def _process_buffer(self, new_content: str) -> str:
311
- """
312
- Process buffered content and return output content.
313
-
314
- Args:
315
- new_content: New content to add to buffer
316
-
317
- Returns:
318
- Processed output content
319
- """
320
- self.pending_buffer += new_content
321
- output_content = ""
322
-
323
- if self.in_thinking_tag:
324
- output_content = self.pending_buffer
325
- self.pending_buffer = ""
326
- return output_content
327
-
328
- while self.pending_buffer:
329
- start_pos = self.pending_buffer.find(self.tool_call_start_token)
330
- end_pos = self.pending_buffer.find(self.tool_call_end_token)
331
-
332
- if start_pos != -1 and (end_pos == -1 or start_pos < end_pos):
333
- tag_pos, tag_len = start_pos, len(self.tool_call_start_token)
334
- elif end_pos != -1:
335
- tag_pos, tag_len = end_pos, len(self.tool_call_end_token)
336
- else:
337
- if self._is_potential_tag_start(self.pending_buffer):
338
- break
339
- output_content += self.pending_buffer
340
- self.pending_buffer = ""
341
- break
342
-
343
- output_content += self.pending_buffer[:tag_pos]
344
- self.pending_buffer = self.pending_buffer[tag_pos + tag_len :]
345
-
346
- return output_content
347
-
348
- def _reset_streaming_state(self) -> None:
349
- """Reset the streaming state to initial values."""
350
- self.streaming_state = {
351
- "current_tool_index": -1,
352
- "tool_ids": [],
353
- "sent_tools": [],
354
- }
355
-
356
- def _advance_to_next_tool(self) -> None:
357
- """Advance to the next tool in the streaming sequence."""
358
- self.streaming_state["current_tool_index"] = (
359
- int(self.streaming_state["current_tool_index"]) + 1
360
- )
361
-
362
- def _set_current_tool_index(self, index: int) -> None:
363
- """
364
- Set the current tool index.
365
-
366
- Args:
367
- index: Tool index to set
368
- """
369
- self.streaming_state["current_tool_index"] = index
370
-
371
- def _get_current_tool_index(self) -> int:
372
- """
373
- Get the current tool index.
374
-
375
- Returns:
376
- Current tool index
377
- """
378
- return int(self.streaming_state["current_tool_index"])
379
-
380
- def _get_next_unsent_tool_index(self, tool_count: int) -> int:
381
- """
382
- Get the index of the next unsent tool.
383
-
384
- Args:
385
- tool_count: Total number of tools
386
-
387
- Returns:
388
- Index of next unsent tool, or -1 if all tools sent
389
- """
390
- sent_tools = list(self.streaming_state["sent_tools"])
391
- for i in range(tool_count):
392
- if i < len(sent_tools):
393
- if not sent_tools[i]["sent_name"]:
394
- return i
395
- else:
396
- return i
397
- return -1
398
-
399
- def _ensure_state_arrays(self, tool_count: int) -> None:
400
- """
401
- Ensure state arrays have sufficient capacity for tool_count tools.
402
-
403
- Args:
404
- tool_count: Number of tools to prepare for
405
- """
406
- sent_tools = list(self.streaming_state["sent_tools"])
407
- tool_ids = list(self.streaming_state["tool_ids"])
408
-
409
- while len(sent_tools) < tool_count:
410
- sent_tools.append(
411
- {
412
- "sent_name": False,
413
- "sent_arguments": "",
414
- "id": make_tool_call_id(),
415
- }
416
- )
417
-
418
- while len(tool_ids) < tool_count:
419
- tool_ids.append(None)
420
-
421
- self.streaming_state["sent_tools"] = sent_tools
422
- self.streaming_state["tool_ids"] = tool_ids
423
-
424
- def _detect_tools_in_text(self, text: str) -> int:
425
- """
426
- Detect the number of tools in text by counting name patterns.
427
-
428
- Args:
429
- text: Text to analyze
430
-
431
- Returns:
432
- Number of tools detected
433
- """
434
- matches = self.tool_name_pattern.findall(text)
435
- return len(matches)
436
-
437
- def _find_tool_boundaries(self, text: str) -> list[tuple[int, int]]:
438
- """
439
- Find the boundaries of tool calls in text.
440
-
441
- Args:
442
- text: Text to analyze
443
-
444
- Returns:
445
- List of (start, end) positions for tool calls
446
- """
447
- boundaries = []
448
- i = 0
449
- while i < len(text):
450
- if text[i] == "{":
451
- start = i
452
- depth = 0
453
- has_name = False
454
- has_arguments = False
455
-
456
- while i < len(text):
457
- if text[i] == "{":
458
- depth += 1
459
- elif text[i] == "}":
460
- depth -= 1
461
- if depth == 0:
462
- end = i + 1
463
- segment = text[start:end]
464
- if '"name"' in segment and '"arguments"' in segment:
465
- boundaries.append((start, end))
466
- break
467
-
468
- if not has_name and '"name"' in text[start : i + 1]:
469
- has_name = True
470
- if not has_arguments and '"arguments"' in text[start : i + 1]:
471
- has_arguments = True
472
-
473
- i += 1
474
-
475
- if depth > 0 and has_name:
476
- boundaries.append((start, i))
477
- else:
478
- i += 1
479
- return boundaries
480
-
481
- def _extract_tool_args(self, tool_content: str, args_match: re.Match[str]) -> str:
482
- """
483
- Extract tool arguments from tool content.
484
-
485
- Args:
486
- tool_content: Tool call content
487
- args_match: Regex match for arguments pattern
488
-
489
- Returns:
490
- Extracted arguments as string
491
- """
492
- args_start_pos = args_match.end()
493
- remaining_content = tool_content[args_start_pos:]
494
-
495
- if remaining_content.strip().startswith("{"):
496
- depth = 0
497
- for i, char in enumerate(remaining_content):
498
- if char == "{":
499
- depth += 1
500
- elif char == "}":
501
- depth -= 1
502
- if depth == 0:
503
- return remaining_content[: i + 1]
504
- else:
505
- args_end = remaining_content.find("}")
506
- if args_end > 0:
507
- return remaining_content[:args_end].strip()
508
-
509
- return remaining_content.rstrip("}").strip()
510
-
511
- def _get_current_tool_content(
512
- self, text: str, tool_index: int
513
- ) -> tuple[str | None, str | None]:
514
- """
515
- Get the content of a specific tool by index.
516
-
517
- Args:
518
- text: Text containing tool calls
519
- tool_index: Index of tool to extract
520
-
521
- Returns:
522
- Tuple of (tool_name, tool_arguments) or (None, None) if not found
523
- """
524
- boundaries = self._find_tool_boundaries(text)
525
-
526
- if tool_index >= len(boundaries):
527
- return None, None
528
-
529
- start, end = boundaries[tool_index]
530
- tool_content = text[start:end]
531
-
532
- name_match = self.tool_name_pattern.search(tool_content)
533
- name = name_match.group(1) if name_match else None
534
-
535
- args_match = self.tool_args_pattern.search(tool_content)
536
- if args_match:
537
- try:
538
- args_text = self._extract_tool_args(tool_content, args_match)
539
- return name, args_text
540
- except Exception:
541
- remaining_content = tool_content[args_match.end() :]
542
- args_text = remaining_content.rstrip("}").strip()
543
- return name, args_text
544
-
545
- return name, None
546
-
547
- def _handle_tool_name_streaming(
548
- self, tool_content: str, tool_count: int
549
- ) -> DeltaMessage | None:
550
- """
551
- Handle streaming of tool names.
552
-
553
- Args:
554
- tool_content: Content containing tool calls
555
- tool_count: Total number of tools
556
-
557
- Returns:
558
- DeltaMessage with tool name or None if no tool to stream
559
- """
560
- next_idx = self._get_next_unsent_tool_index(tool_count)
561
-
562
- if next_idx == -1:
563
- return None
564
-
565
- boundaries = self._find_tool_boundaries(tool_content)
566
- if next_idx >= len(boundaries):
567
- return None
568
-
569
- tool_name, _ = self._get_current_tool_content(tool_content, next_idx)
570
- if not tool_name:
571
- return None
572
-
573
- self._set_current_tool_index(next_idx)
574
- sent_tools = list(self.streaming_state["sent_tools"])
575
- tool_ids = list(self.streaming_state["tool_ids"])
576
-
577
- tool_id = sent_tools[next_idx]["id"]
578
- tool_ids[next_idx] = tool_id
579
- sent_tools[next_idx]["sent_name"] = True
580
-
581
- self.streaming_state["sent_tools"] = sent_tools
582
- self.streaming_state["tool_ids"] = tool_ids
583
-
584
- return DeltaMessage(
585
- tool_calls=[
586
- DeltaToolCall(
587
- index=next_idx,
588
- type="function",
589
- id=tool_id,
590
- function=DeltaFunctionCall(name=tool_name).model_dump(
591
- exclude_none=True
592
- ),
593
- )
594
- ]
595
- )
596
-
597
- def _handle_tool_args_streaming(
598
- self, tool_content: str, tool_count: int
599
- ) -> DeltaMessage | None:
600
- """
601
- Handle streaming of tool arguments.
602
-
603
- Args:
604
- tool_content: Content containing tool calls
605
- tool_count: Total number of tools
606
-
607
- Returns:
608
- DeltaMessage with tool arguments or None if no arguments to stream
609
- """
610
- current_idx = self._get_current_tool_index()
611
-
612
- if current_idx < 0 or current_idx >= tool_count:
613
- return None
614
-
615
- tool_name, tool_args = self._get_current_tool_content(tool_content, current_idx)
616
- if not tool_name or tool_args is None:
617
- return None
618
-
619
- sent_tools = list(self.streaming_state["sent_tools"])
620
-
621
- if not sent_tools[current_idx]["sent_name"]:
622
- return None
623
-
624
- clean_args = self._clean_duplicate_braces(tool_args)
625
- sent_args = sent_tools[current_idx]["sent_arguments"]
626
-
627
- if clean_args != sent_args:
628
- if sent_args and clean_args.startswith(sent_args):
629
- args_delta = extract_intermediate_diff(clean_args, sent_args)
630
- if args_delta:
631
- args_delta = self._clean_delta_braces(args_delta)
632
- sent_tools[current_idx]["sent_arguments"] = clean_args
633
- self.streaming_state["sent_tools"] = sent_tools
634
-
635
- if clean_args.endswith("}"):
636
- self._advance_to_next_tool()
637
-
638
- return DeltaMessage(
639
- tool_calls=[
640
- DeltaToolCall(
641
- index=current_idx,
642
- function=DeltaFunctionCall(
643
- arguments=args_delta
644
- ).model_dump(exclude_none=True),
645
- )
646
- ]
647
- )
648
- elif not sent_args and clean_args:
649
- clean_args_delta = self._clean_delta_braces(clean_args)
650
- sent_tools[current_idx]["sent_arguments"] = clean_args
651
- self.streaming_state["sent_tools"] = sent_tools
652
-
653
- if clean_args.endswith("}"):
654
- self._advance_to_next_tool()
655
-
656
- return DeltaMessage(
657
- tool_calls=[
658
- DeltaToolCall(
659
- index=current_idx,
660
- function=DeltaFunctionCall(
661
- arguments=clean_args_delta
662
- ).model_dump(exclude_none=True),
663
- )
664
- ]
665
- )
666
-
667
- return None
668
-
669
- def _is_end_tool_calls(self, current_text: str) -> bool:
670
- if self.tool_call_end_token not in current_text:
671
- return False
672
-
673
- end_token_positions = []
674
- search_start = 0
675
- while True:
676
- pos = current_text.find(self.tool_call_end_token, search_start)
677
- if pos == -1:
678
- break
679
- end_token_positions.append(pos)
680
- search_start = pos + 1
681
-
682
- think_regions = []
683
- for match in re.finditer(
684
- self.thinking_tag_pattern, current_text, flags=re.DOTALL
685
- ):
686
- think_regions.append((match.start(), match.end()))
687
-
688
- for pos in end_token_positions:
689
- in_think = any(
690
- pos >= t_start and pos < t_end for t_start, t_end in think_regions
691
- )
692
- if not in_think:
693
- return True
694
-
695
- return False
696
-
697
- def extract_tool_calls_streaming(
698
- self,
699
- previous_text: str,
700
- current_text: str,
701
- delta_text: str,
702
- previous_token_ids: Sequence[int],
703
- current_token_ids: Sequence[int],
704
- delta_token_ids: Sequence[int],
705
- request: ChatCompletionRequest,
706
- ) -> DeltaMessage | None:
707
- self._update_thinking_state(current_text)
708
-
709
- if self.in_thinking_tag:
710
- return DeltaMessage(content=delta_text)
711
-
712
- if self._should_buffer_content(delta_text):
713
- buffered_output = self._process_buffer(delta_text)
714
- return DeltaMessage(content=buffered_output) if buffered_output else None
715
-
716
- if self._is_end_tool_calls(current_text):
717
- return DeltaMessage(content=delta_text)
718
-
719
- safe_content, potential_tag = self._split_content_for_buffering(delta_text)
720
- if potential_tag:
721
- self.pending_buffer += potential_tag
722
- return DeltaMessage(content=safe_content) if safe_content else None
723
-
724
- processed_current_text = self.preprocess_model_output(current_text)
725
-
726
- if self.tool_call_start_token not in processed_current_text:
727
- if (
728
- self.tool_call_end_token in delta_text
729
- and self.tool_call_start_token in current_text
730
- ):
731
- return None
732
- if delta_text.strip() == "" and self.tool_call_start_token in current_text:
733
- return None
734
- if (
735
- self._get_current_tool_index() != -1
736
- and self.tool_call_end_token in current_text
737
- ):
738
- self._reset_streaming_state()
739
- return DeltaMessage(content=delta_text)
740
-
741
- if (
742
- self.tool_call_start_token_id is not None
743
- and self.tool_call_start_token_id in delta_token_ids
744
- and len(delta_token_ids) == 1
745
- ):
746
- return None
747
-
748
- original_tool_start = self._find_tool_start_outside_thinking(current_text)
749
- if original_tool_start is None:
750
- return None
751
-
752
- content_before_tools = self._extract_content_before_tools(
753
- current_text, delta_text, original_tool_start
754
- )
755
- if content_before_tools:
756
- return DeltaMessage(content=content_before_tools)
757
-
758
- try:
759
- tool_content = self._extract_tool_content(current_text, original_tool_start)
760
- current_tools_count = self._detect_tools_in_text(tool_content)
761
-
762
- if current_tools_count == 0:
763
- return None
764
-
765
- if self._get_current_tool_index() == -1:
766
- self._reset_streaming_state()
767
-
768
- self._ensure_state_arrays(current_tools_count)
769
-
770
- return self._handle_tool_name_streaming(
771
- tool_content, current_tools_count
772
- ) or self._handle_tool_args_streaming(tool_content, current_tools_count)
773
-
774
- except Exception:
775
- logger.exception(
776
- "An unexpected error occurred ", "during streaming tool call handling."
777
- )
778
- return None
779
-
780
- def _find_tool_start_outside_thinking(self, current_text: str) -> int | None:
781
- """
782
- Find the start position of tool calls outside of thinking tags.
783
-
784
- Args:
785
- current_text: Current text to search
786
-
787
- Returns:
788
- Position of tool call start or None if not found
789
- """
790
- search_start = 0
791
- while True:
792
- pos = current_text.find(self.tool_call_start_token, search_start)
793
- if pos == -1:
794
- return None
795
-
796
- think_regions = [
797
- (m.start(), m.end())
798
- for m in re.finditer(
799
- r"<think>(.*?)</think>", current_text, flags=re.DOTALL
800
- )
801
- ]
802
- in_think = any(
803
- pos >= t_start and pos < t_end for t_start, t_end in think_regions
804
- )
805
-
806
- if not in_think:
807
- return pos
808
-
809
- search_start = pos + 1
810
-
811
- def _extract_content_before_tools(
812
- self, current_text: str, delta_text: str, tool_start: int
813
- ) -> str | None:
814
- """
815
- Extract content that appears before tool calls.
816
-
817
- Args:
818
- current_text: Current text
819
- delta_text: Delta text
820
- tool_start: Start position of tools
821
-
822
- Returns:
823
- Content before tools or None
824
- """
825
- if tool_start > 0:
826
- delta_start_pos = len(current_text) - len(delta_text)
827
- if delta_start_pos < tool_start:
828
- content_part = delta_text
829
- if delta_start_pos + len(delta_text) > tool_start:
830
- content_part = delta_text[: tool_start - delta_start_pos]
831
- return content_part if content_part else None
832
- return None
833
-
834
- def _extract_tool_content(self, current_text: str, tool_start: int) -> str:
835
- """
836
- Extract tool content from current text starting at tool_start.
837
-
838
- Args:
839
- current_text: Current text
840
- tool_start: Start position of tool calls
841
-
842
- Returns:
843
- Extracted tool content
844
- """
845
- tool_content_start = tool_start + len(self.tool_call_start_token)
846
- tool_content = current_text[tool_content_start:]
847
-
848
- end_pos = tool_content.find(self.tool_call_end_token)
849
- if end_pos != -1:
850
- tool_content = tool_content[:end_pos]
851
-
852
- return tool_content