feat: support to OpenMOSS-Team (#44124)
Signed-off-by: nagisa-kun <1434936049@qq.com> Signed-off-by: nagisa19 <1434936049@qq.com> Signed-off-by: nagisa <1434936049@qq.com> Co-authored-by: OpenAI Codex <codex@openai.com> Co-authored-by: Roger Wang <hey@rogerw.io> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
hurukawa committed
Jun 24, 2026 at 15:08 UTC
489abadfb808e1255577f4fb51b6092ec8ca771f
11 files changed
+2926
-2
docs/models/supported_models.md
+1
@@ -598,6 +598,7 @@ These models primarily accept the [`LLM.generate`](./generative_models.md#llmgen
598
| `Mistral3ForConditionalGeneration` | Mistral3 (HF Transformers) | T + I<sup>+</sup> | `mistralai/Mistral-Small-3.1-24B-Instruct-2503`, etc. | ✅︎ | ✅︎ |
599
| `MolmoForCausalLM` | Molmo | T + I<sup>+</sup> | `allenai/Molmo-7B-D-0924`, `allenai/Molmo-7B-O-0924`, etc. | ✅︎ | ✅︎ |
600
| `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> | ✅︎ | ✅︎ |
601
+| `MossAudioModel` | MOSS-Audio | T + A<sup>+</sup> | `OpenMOSS-Team/MOSS-Audio-4B-Instruct`, `OpenMOSS-Team/MOSS-Audio-4B-Thinking`, `OpenMOSS-Team/MOSS-Audio-8B-Instruct`, `OpenMOSS-Team/MOSS-Audio-8B-Thinking` | ✅︎ | ✅︎ |
602
| `Moondream3ForCausalLM` | Moondream3 | T + I | `moondream/moondream3-preview` | | ✅︎ |
603
| `MusicFlamingoForConditionalGeneration` | MusicFlamingo | T + A | `nvidia/music-flamingo-2601-hf`, `nvidia/music-flamingo-think-2601-hf` | ✅︎ | ✅︎ |
604
| `NVLM_D_Model` | NVLM-D 1.0 | T + I<sup>+</sup> | `nvidia/NVLM-D-72B`, etc. | | ✅︎ |
tests/benchmarks/test_throughput_cli.py
+74
@@ -4,6 +4,13 @@ import subprocess
4
5
import pytest
6
7
+from vllm.benchmarks.datasets import SampleRequest
8
+from vllm.benchmarks.throughput import (
9
+ _run_vllm_chat_requests,
10
+ add_cli_args,
11
+)
12
+from vllm.utils.argparse_utils import FlexibleArgumentParser
13
+
14
MODEL_NAME = "meta-llama/Llama-3.2-1B-Instruct"
15
16
@@ -28,3 +35,70 @@ def test_bench_throughput():
35
print(result.stderr)
36
37
assert result.returncode == 0, f"Benchmark failed: {result.stderr}"
38
+
39
+
40
+def test_bench_throughput_accepts_custom_audio_args():
41
+ parser = FlexibleArgumentParser()
42
+ add_cli_args(parser)
43
+
44
+ args = parser.parse_args(
45
+ [
46
+ "--dataset-name",
47
+ "custom_audio",
48
+ "--dataset-path",
49
+ "audio.jsonl",
50
+ "--no-oversample",
51
+ "--custom-output-len",
52
+ "32",
53
+ "--enable-multimodal-chat",
54
+ ]
55
+ )
56
+
57
+ assert args.dataset_name == "custom_audio"
58
+ assert args.no_oversample
59
+ assert args.custom_output_len == 32
60
+ assert args.enable_multimodal_chat
61
+
62
+
63
+def test_vllm_chat_requests_include_multimodal_content():
64
+ class FakeLLM:
65
+ def __init__(self):
66
+ self.prompts = None
67
+
68
+ def chat(self, prompts, sampling_params, use_tqdm):
69
+ del sampling_params, use_tqdm
70
+ self.prompts = prompts
71
+ return []
72
+
73
+ llm = FakeLLM()
74
+ audio_content = {
75
+ "type": "input_audio",
76
+ "input_audio": {"data": "abc", "format": "wav"},
77
+ }
78
+ request = SampleRequest(
79
+ prompt="Transcribe this audio.",
80
+ prompt_len=1,
81
+ expected_output_len=8,
82
+ multi_modal_data=audio_content,
83
+ )
84
+
85
+ _run_vllm_chat_requests(
86
+ llm,
87
+ [request],
88
+ n=1,
89
+ disable_detokenize=False,
90
+ do_profile=False,
91
+ prequeue_requests=False,
92
+ )
93
+
94
+ assert llm.prompts == [
95
+ [
96
+ {
97
+ "role": "user",
98
+ "content": [
99
+ {"type": "text", "text": "Transcribe this audio."},
100
+ audio_content,
101
+ ],
102
+ }
103
+ ]
104
+ ]
tests/entrypoints/openai/chat_completion/test_chat.py
+4
-1
@@ -824,7 +824,10 @@ async def test_invocations(server: RemoteOpenAIServer, client: openai.AsyncOpenA
824
chat_output = chat_response.json()
825
invocation_output = invocation_response.json()
826
827
- assert chat_output.keys() == invocation_output.keys()
827
+ extra_keys = invocation_output.keys() - chat_output.keys()
828
+ missing_keys = chat_output.keys() - invocation_output.keys()
829
+ assert missing_keys == set()
830
+ assert extra_keys <= {"moderation"}
831
assert chat_output["choices"] == invocation_output["choices"]
832
833
tests/models/multimodal/generation/test_moss_audio.py
new
+155
@@ -0,0 +1,155 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+import pytest
5
+
6
+from vllm.assets.audio import AudioAsset
7
+from vllm.model_executor.models.moss_audio import MOSS_AUDIO_PLACEHOLDER
8
+from vllm.platforms import current_platform
9
+
10
+from ...registry import HF_EXAMPLE_MODELS
11
+from ...utils import check_logprobs_close
12
+
13
+CORE_MODEL = pytest.param(
14
+ "OpenMOSS-Team/MOSS-Audio-4B-Instruct",
15
+ marks=pytest.mark.core_model,
16
+ id="4b-instruct",
17
+)
18
+
19
+EXTENDED_MODELS = [
20
+ "OpenMOSS-Team/MOSS-Audio-4B-Thinking",
21
+ "OpenMOSS-Team/MOSS-Audio-8B-Instruct",
22
+ "OpenMOSS-Team/MOSS-Audio-8B-Thinking",
23
+]
24
+
25
+ACCURACY_MODELS = [CORE_MODEL, *EXTENDED_MODELS]
26
+
27
+PARALLEL_SMOKE_CASES = [
28
+ pytest.param({"tensor_parallel_size": 2}, id="tp2"),
29
+ pytest.param({"pipeline_parallel_size": 2}, id="pp2"),
30
+ pytest.param(
31
+ {"tensor_parallel_size": 2, "pipeline_parallel_size": 2},
32
+ id="tp2_pp2",
33
+ ),
34
+]
35
+
36
+HF_ACCURACY_SKIP_REASON = (
37
+ "HF AutoModelForCausalLM cannot load remote MOSS-Audio configs; "
38
+ "vLLM generation coverage is provided by the smoke tests below."
39
+)
40
+
41
+
42
+@pytest.mark.core_model
43
+def test_moss_audio_generation_smoke(vllm_runner) -> None:
44
+ model = "OpenMOSS-Team/MOSS-Audio-4B-Instruct"
45
+ model_info = HF_EXAMPLE_MODELS.find_hf_info(model)
46
+ model_info.check_available_online(on_fail="skip")
47
+ model_info.check_transformers_version(on_fail="skip")
48
+
49
+ prompts = [f"{MOSS_AUDIO_PLACEHOLDER}\nBriefly describe this audio."]
50
+ audios = [[AudioAsset("mary_had_lamb").audio_and_sample_rate[0]]]
51
+
52
+ with vllm_runner(
53
+ model,
54
+ dtype="half",
55
+ enforce_eager=True,
56
+ max_model_len=1024,
57
+ limit_mm_per_prompt={"audio": 1},
58
+ trust_remote_code=True,
59
+ ) as vllm_model:
60
+ outputs = vllm_model.generate_greedy(
61
+ prompts,
62
+ max_tokens=4,
63
+ audios=audios,
64
+ )
65
+
66
+ assert len(outputs) == 1
67
+ assert len(outputs[0][1]) > 0
68
+
69
+
70
+@pytest.mark.skip(reason=HF_ACCURACY_SKIP_REASON)
71
+@pytest.mark.parametrize("model", ACCURACY_MODELS)
72
+@pytest.mark.parametrize("dtype", ["half"])
73
+@pytest.mark.parametrize("max_tokens", [8])
74
+@pytest.mark.parametrize("num_logprobs", [5])
75
+def test_moss_audio_hf_vllm_accuracy(
76
+ hf_runner,
77
+ vllm_runner,
78
+ model: str,
79
+ dtype: str,
80
+ max_tokens: int,
81
+ num_logprobs: int,
82
+) -> None:
83
+ model_info = HF_EXAMPLE_MODELS.find_hf_info(model)
84
+ model_info.check_available_online(on_fail="skip")
85
+ model_info.check_transformers_version(on_fail="skip")
86
+
87
+ prompts = [f"{MOSS_AUDIO_PLACEHOLDER}\nTranscribe this audio."]
88
+ audios = [[AudioAsset("mary_had_lamb").audio_and_sample_rate[0]]]
89
+
90
+ with vllm_runner(
91
+ model,
92
+ dtype=dtype,
93
+ enforce_eager=True,
94
+ max_model_len=1024,
95
+ limit_mm_per_prompt={"audio": 1},
96
+ trust_remote_code=True,
97
+ ) as vllm_model:
98
+ vllm_outputs = vllm_model.generate_greedy_logprobs(
99
+ prompts,
100
+ max_tokens,
101
+ num_logprobs=num_logprobs,
102
+ audios=audios,
103
+ )
104
+
105
+ with hf_runner(model, dtype=dtype, trust_remote_code=True) as hf_model:
106
+ hf_outputs = hf_model.generate_greedy_logprobs_limit(
107
+ prompts,
108
+ max_tokens,
109
+ num_logprobs=num_logprobs,
110
+ audios=audios,
111
+ )
112
+
113
+ check_logprobs_close(
114
+ outputs_0_lst=hf_outputs,
115
+ outputs_1_lst=vllm_outputs,
116
+ name_0="hf",
117
+ name_1="vllm",
118
+ )
119
+
120
+
121
+@pytest.mark.core_model
122
+@pytest.mark.parametrize("parallel_kwargs", PARALLEL_SMOKE_CASES)
123
+def test_moss_audio_parallel_smoke(vllm_runner, parallel_kwargs) -> None:
124
+ model = "OpenMOSS-Team/MOSS-Audio-4B-Instruct"
125
+ required_gpus = parallel_kwargs.get(
126
+ "tensor_parallel_size", 1
127
+ ) * parallel_kwargs.get("pipeline_parallel_size", 1)
128
+ if current_platform.device_count() < required_gpus:
129
+ # TP/PP integration smoke runs on local or multi-GPU CI only.
130
+ pytest.skip(f"Requires at least {required_gpus} GPUs")
131
+
132
+ model_info = HF_EXAMPLE_MODELS.find_hf_info(model)
133
+ model_info.check_available_online(on_fail="skip")
134
+ model_info.check_transformers_version(on_fail="skip")
135
+
136
+ prompts = [f"{MOSS_AUDIO_PLACEHOLDER}\nBriefly describe this audio."]
137
+ audios = [[AudioAsset("mary_had_lamb").audio_and_sample_rate[0]]]
138
+
139
+ with vllm_runner(
140
+ model,
141
+ dtype="half",
142
+ enforce_eager=True,
143
+ max_model_len=1024,
144
+ limit_mm_per_prompt={"audio": 1},
145
+ trust_remote_code=True,
146
+ **parallel_kwargs,
147
+ ) as vllm_model:
148
+ outputs = vllm_model.generate_greedy(
149
+ prompts,
150
+ max_tokens=4,
151
+ audios=audios,
152
+ )
153
+
154
+ assert len(outputs) == 1
155
+ assert len(outputs[0][1]) > 0
tests/models/multimodal/processing/test_common.py
+6
@@ -432,6 +432,12 @@ def test_processing_correctness(
432
)
433
if model_id == "CohereLabs/cohere-transcribe-03-2026":
434
pytest.skip("Fix later")
435
+ if model_id.startswith("OpenMOSS-Team/MOSS-Audio-"):
436
+ pytest.skip(
437
+ "MOSS-Audio uses a custom processor that dynamically expands "
438
+ "audio placeholders from processed audio lengths. Its vLLM "
439
+ "processor paths are covered by test_moss_audio.py."
440
+ )
441
442
_test_processing_correctness(
443
model_id,
tests/models/multimodal/processing/test_moss_audio.py
new
+694
@@ -0,0 +1,694 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+from types import SimpleNamespace
5
+
6
+import numpy as np
7
+import pytest
8
+import torch
9
+from transformers import Qwen3Config
10
+
11
+from vllm.model_executor.models.interfaces import SupportsLoRA, supports_lora
12
+from vllm.model_executor.models.moss_audio import (
13
+ MOSS_AUDIO_BOS_TOKEN,
14
+ MOSS_AUDIO_BOS_TOKEN_ID,
15
+ MOSS_AUDIO_EOS_TOKEN,
16
+ MOSS_AUDIO_EOS_TOKEN_ID,
17
+ MOSS_AUDIO_PLACEHOLDER,
18
+ MOSS_AUDIO_TOKEN,
19
+ MOSS_AUDIO_TOKEN_ID,
20
+ GatedMLP,
21
+ MossAudioConfig,
22
+ MossAudioDummyInputsBuilder,
23
+ MossAudioEncoder,
24
+ MossAudioEncoderConfig,
25
+ MossAudioModel,
26
+ MossAudioMultiModalProcessor,
27
+ MossAudioProcessingInfo,
28
+ MossAudioProcessor,
29
+ MossQwen3ForCausalLM,
30
+ MossQwen3Model,
31
+)
32
+from vllm.model_executor.models.utils import AutoWeightsLoader
33
+from vllm.multimodal.cache import MultiModalProcessorOnlyCache
34
+from vllm.multimodal.inputs import batched_tensors_equal
35
+from vllm.sequence import IntermediateTensors
36
+
37
+
38
+class _Tokenizer:
39
+ def encode(self, text, add_special_tokens=False):
40
+ del add_special_tokens
41
+ return [ord(char) for char in text]
42
+
43
+ def decode(self, token_ids, **kwargs):
44
+ del kwargs
45
+ return "".join(chr(token_id) for token_id in token_ids)
46
+
47
+ def batch_decode(self, batch_token_ids, **kwargs):
48
+ return [self.decode(token_ids, **kwargs) for token_ids in batch_token_ids]
49
+
50
+
51
+class _MMConfig:
52
+ enable_mm_embeds = False
53
+ mm_processor_cache_gb = 1
54
+
55
+ def merge_mm_processor_kwargs(self, kwargs):
56
+ return dict(kwargs)
57
+
58
+ def get_limit_per_prompt(self, modality):
59
+ del modality
60
+ return 3
61
+
62
+
63
+class _ModelConfig:
64
+ def __init__(self):
65
+ self.model = "OpenMOSS-Team/MOSS-Audio-4B-Instruct"
66
+ self.revision = None
67
+ self.max_model_len = 4096
68
+ self.encoder_config = {}
69
+ self.dtype = torch.float32
70
+ self.hf_config = MossAudioConfig(language_config=Qwen3Config())
71
+ self.multimodal_config = _MMConfig()
72
+
73
+ def get_multimodal_config(self):
74
+ return self.multimodal_config
75
+
76
+ def get_inputs_embeds_size(self):
77
+ return None
78
+
79
+
80
+class _ProcessingContext:
81
+ def __init__(self):
82
+ self.model_config = _ModelConfig()
83
+ self.tokenizer = _Tokenizer()
84
+
85
+ def get_tokenizer(self):
86
+ return self.tokenizer
87
+
88
+ def get_hf_config(self):
89
+ return self.model_config.hf_config
90
+
91
+ def get_mm_config(self):
92
+ return self.model_config.get_multimodal_config()
93
+
94
+ def get_merged_mm_kwargs(self, kwargs):
95
+ return self.get_mm_config().merge_mm_processor_kwargs(kwargs)
96
+
97
+ def call_hf_processor(self, hf_processor, data, kwargs):
98
+ merged_kwargs = self.get_merged_mm_kwargs(kwargs)
99
+ merged_kwargs.setdefault("return_tensors", "pt")
100
+ return hf_processor(**data, **merged_kwargs)
101
+
102
+
103
+class _TestMossAudioProcessingInfo(MossAudioProcessingInfo):
104
+ def _get_processor_config_defaults(self):
105
+ return {}
106
+
107
+
108
+def _vllm_config(tensor_parallel_size=1, pipeline_parallel_size=1, hf_config=None):
109
+ if hf_config is None:
110
+ hf_config = MossAudioConfig(language_config=Qwen3Config())
111
+ return SimpleNamespace(
112
+ model_config=SimpleNamespace(
113
+ hf_config=hf_config,
114
+ multimodal_config=None,
115
+ ),
116
+ quant_config=None,
117
+ parallel_config=SimpleNamespace(
118
+ tensor_parallel_size=tensor_parallel_size,
119
+ pipeline_parallel_size=pipeline_parallel_size,
120
+ ),
121
+ )
122
+
123
+
124
+class _FakeAudioEncoder:
125
+ dtype = torch.float32
126
+
127
+ def __init__(self, deepstack_layers=0):
128
+ self.deepstack_layers = deepstack_layers
129
+ self.output_deepstack_hidden_states = None
130
+ self.input_shape = None
131
+ self.feature_lens = None
132
+
133
+ def __call__(self, audio_data, *, feature_lens, output_deepstack_hidden_states):
134
+ self.input_shape = tuple(audio_data.shape)
135
+ self.feature_lens = feature_lens.detach().cpu().clone()
136
+ self.output_deepstack_hidden_states = output_deepstack_hidden_states
137
+ lengths = MossAudioEncoder._compute_downsampled_length(feature_lens)
138
+ hidden_states = torch.ones(1, int(lengths.sum().item()), 8)
139
+ if not output_deepstack_hidden_states:
140
+ return hidden_states, None
141
+ return hidden_states, [
142
+ hidden_states * scale for scale in range(2, 2 + self.deepstack_layers)
143
+ ]
144
+
145
+
146
+def _patch_tensor_parallel_for_linear_layers(monkeypatch, tp_size=1, tp_rank=0):
147
+ import vllm.model_executor.layers.linear as linear_layers
148
+ import vllm.model_executor.models.moss_audio as moss_audio_module
149
+ import vllm.model_executor.parameter as parameter_module
150
+
151
+ for module in (moss_audio_module, linear_layers, parameter_module):
152
+ monkeypatch.setattr(
153
+ module, "get_tensor_model_parallel_world_size", lambda: tp_size
154
+ )
155
+ monkeypatch.setattr(
156
+ linear_layers, "get_tensor_model_parallel_rank", lambda: tp_rank
157
+ )
158
+ monkeypatch.setattr(
159
+ parameter_module, "get_tensor_model_parallel_rank", lambda: tp_rank
160
+ )
161
+ monkeypatch.setattr(
162
+ linear_layers, "tensor_model_parallel_all_reduce", lambda tensor: tensor
163
+ )
164
+
165
+
166
+def _build_moss_audio_processor(cache=None):
167
+ ctx = _ProcessingContext()
168
+ info = _TestMossAudioProcessingInfo(ctx)
169
+ return (
170
+ MossAudioMultiModalProcessor(
171
+ info,
172
+ MossAudioDummyInputsBuilder(info),
173
+ cache=cache,
174
+ ),
175
+ ctx,
176
+ )
177
+
178
+
179
+def _assert_mm_inputs_equal(left, right):
180
+ assert left["prompt_token_ids"] == right["prompt_token_ids"]
181
+ assert left["mm_hashes"] == right["mm_hashes"]
182
+
183
+ left_placeholder = left["mm_placeholders"]["audio"][0]
184
+ right_placeholder = right["mm_placeholders"]["audio"][0]
185
+ assert left_placeholder.offset == right_placeholder.offset
186
+ assert left_placeholder.length == right_placeholder.length
187
+ assert left_placeholder.is_embed.tolist() == right_placeholder.is_embed.tolist()
188
+
189
+ assert batched_tensors_equal(
190
+ left["mm_kwargs"].get_data(),
191
+ right["mm_kwargs"].get_data(),
192
+ )
193
+
194
+
195
+@pytest.mark.parametrize(
196
+ ("prompt", "prefix"),
197
+ [
198
+ (
199
+ f"before {MOSS_AUDIO_PLACEHOLDER} after",
200
+ [*[ord(char) for char in "before "], MOSS_AUDIO_BOS_TOKEN_ID],
201
+ ),
202
+ (
203
+ f"before {MOSS_AUDIO_BOS_TOKEN}{MOSS_AUDIO_TOKEN}"
204
+ f"{MOSS_AUDIO_TOKEN}{MOSS_AUDIO_EOS_TOKEN} after",
205
+ [*[ord(char) for char in "before "], MOSS_AUDIO_BOS_TOKEN_ID],
206
+ ),
207
+ ("Describe this audio.", [MOSS_AUDIO_BOS_TOKEN_ID]),
208
+ ],
209
+)
210
+def test_moss_audio_processor_expands_audio_placeholders(prompt, prefix):
211
+ raw_mel_len = 17
212
+ processed = MossAudioProcessor(_Tokenizer())(
213
+ text=prompt, audio=[torch.zeros(160 * raw_mel_len)]
214
+ )
215
+ input_ids = processed["input_ids"][0].tolist()
216
+
217
+ assert input_ids[: len(prefix)] == prefix
218
+ assert input_ids.count(MOSS_AUDIO_BOS_TOKEN_ID) == 1
219
+ assert input_ids.count(MOSS_AUDIO_EOS_TOKEN_ID) == 1
220
+ assert input_ids.count(MOSS_AUDIO_TOKEN_ID) == (
221
+ MossAudioEncoder.compute_num_audio_tokens(raw_mel_len)
222
+ )
223
+ assert processed["audio_data"].shape == (1, 128, raw_mel_len)
224
+ assert processed["audio_data_seqlens"].tolist() == [raw_mel_len]
225
+
226
+
227
+def test_moss_audio_processor_preserves_placeholder_without_audio():
228
+ processed = MossAudioProcessor(_Tokenizer())(
229
+ text=f"before {MOSS_AUDIO_PLACEHOLDER} after"
230
+ )
231
+
232
+ assert processed["input_ids"][0].tolist() == [
233
+ *[ord(char) for char in "before "],
234
+ MOSS_AUDIO_BOS_TOKEN_ID,
235
+ MOSS_AUDIO_TOKEN_ID,
236
+ MOSS_AUDIO_EOS_TOKEN_ID,
237
+ *[ord(char) for char in " after"],
238
+ ]
239
+ assert "audio_data" not in processed
240
+ assert "audio_data_seqlens" not in processed
241
+
242
+
243
+def test_moss_audio_multimodal_processor_handles_token_and_cache_paths():
244
+ raw_mel_len = 17
245
+ audio = np.zeros(160 * raw_mel_len, dtype=np.float32)
246
+ prompt = f"{MOSS_AUDIO_PLACEHOLDER}\nTranscribe this audio."
247
+
248
+ baseline_processor, ctx = _build_moss_audio_processor()
249
+ mm_items = baseline_processor.info.parse_mm_data({"audio": [audio]})
250
+ token_prompt = ctx.get_tokenizer().encode(prompt, add_special_tokens=False)
251
+
252
+ baseline_text = baseline_processor(
253
+ prompt,
254
+ mm_items=mm_items,
255
+ hf_processor_mm_kwargs={},
256
+ )
257
+ baseline_token = baseline_processor(
258
+ token_prompt,
259
+ mm_items=mm_items,
260
+ hf_processor_mm_kwargs={},
261
+ )
262
+
263
+ cache = MultiModalProcessorOnlyCache(ctx.model_config)
264
+ cached_processor, _ = _build_moss_audio_processor(cache=cache)
265
+ cached_text_miss = cached_processor(
266
+ prompt,
267
+ mm_items=mm_items,
268
+ hf_processor_mm_kwargs={},
269
+ )
270
+ cached_text_hit = cached_processor(
271
+ prompt,
272
+ mm_items=mm_items,
273
+ hf_processor_mm_kwargs={},
274
+ )
275
+ cached_token_hit = cached_processor(
276
+ token_prompt,
277
+ mm_items=mm_items,
278
+ hf_processor_mm_kwargs={},
279
+ )
280
+
281
+ expected_audio_tokens = MossAudioEncoder.compute_num_audio_tokens(raw_mel_len)
282
+ prompt_token_ids = baseline_text["prompt_token_ids"]
283
+ assert prompt_token_ids.count(MOSS_AUDIO_TOKEN_ID) == expected_audio_tokens
284
+ assert baseline_text["mm_placeholders"]["audio"][0].length == (
285
+ expected_audio_tokens + 2
286
+ )
287
+
288
+ _assert_mm_inputs_equal(baseline_text, baseline_token)
289
+ _assert_mm_inputs_equal(baseline_text, cached_text_miss)
290
+ _assert_mm_inputs_equal(baseline_text, cached_text_hit)
291
+ _assert_mm_inputs_equal(baseline_text, cached_token_hit)
292
+
293
+
294
+def test_moss_audio_supports_language_model_lora_only():
295
+ assert supports_lora(MossAudioModel)
296
+
297
+ model = object.__new__(MossAudioModel)
298
+ assert isinstance(model, SupportsLoRA)
299
+
300
+ mapping = model.get_mm_mapping()
301
+ assert mapping.language_model == ["language_model."]
302
+ assert mapping.tower_model == []
303
+ assert mapping.connector == []
304
+
305
+
306
+def test_moss_audio_error_paths():
307
+ model = object.__new__(MossAudioModel)
308
+ with pytest.raises(ValueError, match="DeepStack audio token count mismatch"):
309
+ model._cache_deepstack_input_embeds(
310
+ inputs_embeds=torch.zeros(4, 8),
311
+ deepstack_embeddings=((torch.ones(1, 8),),),
312
+ is_multimodal=torch.tensor([False, True, True, False]),
313
+ )
314
+
315
+ with pytest.raises(ValueError, match="too short"):
316
+ MossAudioProcessor(_Tokenizer())(
317
+ text=MOSS_AUDIO_PLACEHOLDER, audio=[torch.empty(0)]
318
+ )
319
+ with pytest.raises(ValueError, match="too short"):
320
+ model._parse_and_validate_audio_input(
321
+ audio_data=torch.zeros(1, 128, 1),
322
+ audio_data_seqlens=torch.tensor([0], dtype=torch.long),
323
+ )
324
+
325
+
326
+def test_moss_audio_validates_tp_config():
327
+ vllm_config = _vllm_config(tensor_parallel_size=2)
328
+ vllm_config.model_config.hf_config.adapter_hidden_size = 7
329
+
330
+ with pytest.raises(ValueError, match="adapter_hidden_size"):
331
+ MossAudioModel(vllm_config=vllm_config)
332
+
333
+ vllm_config = _vllm_config(tensor_parallel_size=2)
334
+ vllm_config.model_config.hf_config.audio_config.d_model = 6
335
+ vllm_config.model_config.hf_config.audio_config.encoder_attention_heads = 3
336
+ with pytest.raises(ValueError, match="encoder_attention_heads"):
337
+ MossAudioModel(vllm_config=vllm_config)
338
+
339
+
340
+def test_moss_audio_rejects_audio_data_list_seqlen_count_mismatch():
341
+ model = object.__new__(MossAudioModel)
342
+
343
+ with pytest.raises(ValueError, match="audio_data batch size"):
344
+ model._parse_and_validate_audio_input(
345
+ audio_data=[torch.zeros(128, 8), torch.zeros(128, 11)],
346
+ audio_data_seqlens=torch.tensor([8], dtype=torch.long),
347
+ )
348
+
349
+
350
+@pytest.mark.parametrize("deepstack_scales", [(), (7, 11)])
351
+def test_moss_audio_embed_multimodal_packs_by_audio(deepstack_scales):
352
+ model = object.__new__(MossAudioModel)
353
+ model.audio_encoder = _FakeAudioEncoder(len(deepstack_scales))
354
+ model.audio_adapter = lambda hidden_states: hidden_states * 5
355
+ model.deepstack_audio_merger_list = [
356
+ lambda hidden_states, scale=scale: hidden_states * scale
357
+ for scale in deepstack_scales
358
+ ]
359
+ model.deepstack_input_embeds = None
360
+
361
+ embeddings = model.embed_multimodal(
362
+ audio_data=torch.zeros(2, 128, 9),
363
+ audio_data_seqlens=torch.tensor([8, 9], dtype=torch.long),
364
+ )
365
+
366
+ assert model.audio_encoder.output_deepstack_hidden_states is bool(deepstack_scales)
367
+ assert [embeds.shape for embeds in embeddings] == [
368
+ torch.Size([1, 8 * (1 + len(deepstack_scales))]),
369
+ torch.Size([2, 8 * (1 + len(deepstack_scales))]),
370
+ ]
371
+ if not deepstack_scales:
372
+ assert model.deepstack_input_embeds is None
373
+ return
374
+
375
+ main_embeddings, deepstack_embeddings = model._split_multimodal_embeddings(
376
+ embeddings, hidden_size=8
377
+ )
378
+ assert [embeds.shape for embeds in main_embeddings] == [
379
+ torch.Size([1, 8]),
380
+ torch.Size([2, 8]),
381
+ ]
382
+ assert [[e.shape for e in layer] for layer in deepstack_embeddings] == [
383
+ [torch.Size([1, 8]), torch.Size([2, 8])] for _ in deepstack_scales
384
+ ]
385
+ assert torch.equal(main_embeddings[0], torch.full((1, 8), 5.0))
386
+ for idx, scale in enumerate(deepstack_scales):
387
+ assert torch.equal(
388
+ deepstack_embeddings[idx][0],
389
+ torch.full((1, 8), float((idx + 2) * scale)),
390
+ )
391
+
392
+
393
+def test_moss_audio_embed_input_ids_caches_packed_deepstack():
394
+ class _FakeLanguageModel:
395
+ def embed_input_ids(self, input_ids):
396
+ return torch.zeros(input_ids.shape[0], 8)
397
+
398
+ model = object.__new__(MossAudioModel)
399
+ model.language_model = _FakeLanguageModel()
400
+ model.deepstack_audio_merger_list = [object(), object()]
401
+ model.deepstack_input_embeds = None
402
+ multimodal_embeddings = (
403
+ torch.cat([torch.full((1, 8), x) for x in (5.0, 14.0, 33.0)], dim=-1),
404
+ torch.cat([torch.full((2, 8), x) for x in (7.0, 22.0, 44.0)], dim=-1),
405
+ )
406
+ is_multimodal = torch.tensor([False, True, True, True, False])
407
+
408
+ inputs_embeds = model.embed_input_ids(
409
+ input_ids=torch.arange(5),
410
+ multimodal_embeddings=multimodal_embeddings,
411
+ is_multimodal=is_multimodal,
412
+ )
413
+
414
+ assert torch.equal(inputs_embeds[1], torch.full((8,), 5.0))
415
+ assert torch.equal(inputs_embeds[2], torch.full((8,), 7.0))
416
+ assert torch.equal(inputs_embeds[3], torch.full((8,), 7.0))
417
+ assert model.deepstack_input_embeds is not None
418
+ tensors = model.deepstack_input_embeds.tensors
419
+ assert set(tensors) == {"deepstack_input_embeds_0", "deepstack_input_embeds_1"}
420
+ for tensor in tensors.values():
421
+ assert tensor[is_multimodal].abs().sum() > 0
422
+ assert torch.equal(tensor[~is_multimodal], torch.zeros(2, 8))
423
+
424
+
425
+def _patch_pp_group(monkeypatch, *, first=True, last=True):
426
+ import vllm.model_executor.models.moss_audio as moss_audio_module
427
+
428
+ monkeypatch.setattr(
429
+ moss_audio_module,
430
+ "get_pp_group",
431
+ lambda: SimpleNamespace(is_first_rank=first, is_last_rank=last),
432
+ )
433
+
434
+
435
+def test_moss_audio_pp_forward_routes_deepstack(monkeypatch):
436
+ for first in (True, False):
437
+ calls: list[dict[str, object]] = []
438
+
439
+ def fake_lm(*args, _calls=calls, **kwargs):
440
+ del args
441
+ _calls.append(kwargs)
442
+ return torch.ones(1, 1)
443
+
444
+ _patch_pp_group(monkeypatch, first=first)
445
+ model = object.__new__(MossAudioModel)
446
+ torch.nn.Module.__init__(model)
447
+ model.language_model = fake_lm
448
+ cached = IntermediateTensors({"deepstack_input_embeds_0": torch.ones(3, 8)})
449
+ inter = IntermediateTensors(
450
+ {
451
+ "hidden_states": torch.ones(3, 8),
452
+ "residual": torch.zeros(3, 8),
453
+ "deepstack_input_embeds_0": torch.full((3, 8), 5.0),
454
+ }
455
+ )
456
+ inputs_embeds = torch.full((3, 8), 9.0)
457
+ model.deepstack_input_embeds = cached
458
+
459
+ model.forward(
460
+ input_ids=None,
461
+ positions=torch.arange(3),
462
+ intermediate_tensors=None if first else inter,
463
+ inputs_embeds=inputs_embeds if first else None,
464
+ )
465
+
466
+ kwargs = calls[0]
467
+ assert kwargs["inputs_embeds"] is (inputs_embeds if first else None)
468
+ assert kwargs["deepstack_input_embeds"] is (cached if first else inter)
469
+ assert model.deepstack_input_embeds is None
470
+
471
+ calls = []
472
+
473
+ def fake_lm_non_first_rank(*args, **kwargs):
474
+ del args
475
+ calls.append(kwargs)
476
+ return torch.ones(1, 1)
477
+
478
+ _patch_pp_group(monkeypatch, first=False)
479
+ model = object.__new__(MossAudioModel)
480
+ torch.nn.Module.__init__(model)
481
+ model.language_model = fake_lm_non_first_rank
482
+ model.deepstack_input_embeds = IntermediateTensors({})
483
+ inter = IntermediateTensors(
484
+ {
485
+ "hidden_states": torch.ones(3, 8),
486
+ "residual": torch.zeros(3, 8),
487
+ }
488
+ )
489
+
490
+ model.forward(
491
+ input_ids=None,
492
+ positions=torch.arange(3),
493
+ intermediate_tensors=inter,
494
+ inputs_embeds=torch.ones(3, 8),
495
+ )
496
+ assert calls[0]["inputs_embeds"] is None
497
+ assert calls[0]["deepstack_input_embeds"] is inter
498
+
499
+
500
+def test_moss_qwen3_deepstack_keys_for_pp(monkeypatch):
501
+ class AddOne(torch.nn.Module):
502
+ def forward(self, positions, hidden_states, residual):
503
+ del positions, residual
504
+ return hidden_states + 1, torch.zeros_like(hidden_states)
505
+
506
+ def make_model(num_layers, deepstack_layers=None):
507
+ model = object.__new__(MossQwen3Model)
508
+ torch.nn.Module.__init__(model)
509
+ model.start_layer, model.end_layer = 0, num_layers
510
+ model.layers = torch.nn.ModuleList([AddOne() for _ in range(num_layers)])
511
+ model.norm = lambda hidden_states, residual: (hidden_states, residual)
512
+ model._maybe_add_hidden_state = lambda aux, *args: aux
513
+ model.deepstack_inject_layer_indices = (
514
+ range(0) if deepstack_layers is None else deepstack_layers
515
+ )
516
+ return model
517
+
518
+ _patch_pp_group(monkeypatch, first=True, last=True)
519
+ output = make_model(3).forward(
520
+ input_ids=None,
521
+ positions=torch.arange(2),
522
+ inputs_embeds=torch.zeros(2, 4),
523
+ deepstack_input_embeds=IntermediateTensors(
524
+ {
525
+ "deepstack_input_embeds_2": torch.full((2, 4), 5.0),
526
+ }
527
+ ),
528
+ )
529
+ assert torch.equal(output, torch.full((2, 4), 8.0))
530
+
531
+ _patch_pp_group(monkeypatch, first=True, last=False)
532
+ deepstack = IntermediateTensors(
533
+ {
534
+ "deepstack_input_embeds_0": torch.full((2, 4), 7.0),
535
+ "deepstack_input_embeds_3": torch.full((2, 4), 11.0),
536
+ }
537
+ )
538
+ output = make_model(2, range(4)).forward(
539
+ input_ids=None,
540
+ positions=torch.arange(2),
541
+ inputs_embeds=torch.zeros(2, 4),
542
+ deepstack_input_embeds=deepstack,
543
+ )
544
+ assert isinstance(output, IntermediateTensors)
545
+ assert set(output.tensors) == {
546
+ "hidden_states",
547
+ "residual",
548
+ "deepstack_input_embeds_2",
549
+ "deepstack_input_embeds_3",
550
+ }
551
+ assert torch.equal(output["hidden_states"], torch.full((2, 4), 9.0))
552
+ assert torch.equal(output["deepstack_input_embeds_2"], torch.zeros(2, 4))
553
+ assert output["deepstack_input_embeds_3"] is deepstack["deepstack_input_embeds_3"]
554
+
555
+ inner_model = make_model(0, range(2))
556
+ inner_model.make_empty_intermediate_tensors = lambda batch, dtype, device: (
557
+ IntermediateTensors(
558
+ {
559
+ "hidden_states": torch.zeros(batch, 4, dtype=dtype, device=device),
560
+ "residual": torch.zeros(batch, 4, dtype=dtype, device=device),
561
+ }
562
+ )
563
+ )
564
+ language_model = object.__new__(MossQwen3ForCausalLM)
565
+ torch.nn.Module.__init__(language_model)
566
+ language_model.model = inner_model
567
+ language_model.config = SimpleNamespace(hidden_size=4)
568
+ language_model.deepstack_inject_layer_indices = range(2)
569
+
570
+ tensors = MossQwen3ForCausalLM.make_empty_intermediate_tensors(
571
+ language_model,
572
+ batch_size=3,
573
+ dtype=torch.float16,
574
+ device=torch.device("cpu"),
575
+ )
576
+
577
+ assert set(tensors.tensors) == {
578
+ "hidden_states",
579
+ "residual",
580
+ "deepstack_input_embeds_0",
581
+ "deepstack_input_embeds_1",
582
+ }
583
+ assert tensors["deepstack_input_embeds_0"].shape == (3, 4)
584
+ assert tensors["deepstack_input_embeds_0"].dtype == torch.float16
585
+
586
+ _patch_pp_group(monkeypatch, first=True, last=False)
587
+ forward_tensors = inner_model.forward(
588
+ input_ids=None,
589
+ positions=torch.arange(3),
590
+ inputs_embeds=torch.ones(3, 4, dtype=torch.float16),
591
+ deepstack_input_embeds=None,
592
+ )
593
+ assert isinstance(forward_tensors, IntermediateTensors)
594
+ assert set(forward_tensors.tensors) == set(tensors.tensors)
595
+
596
+
597
+@pytest.mark.parametrize("tp_size", [1, 2])
598
+def test_moss_audio_gated_mlp_tp_shapes_and_loading(monkeypatch, tp_size):
599
+ from vllm.config import VllmConfig, set_current_vllm_config
600
+ from vllm.config.device import DeviceConfig
601
+
602
+ _patch_tensor_parallel_for_linear_layers(monkeypatch, tp_size=tp_size)
603
+ with set_current_vllm_config(VllmConfig(device_config=DeviceConfig(device="cpu"))):
604
+ mlp = GatedMLP(input_size=4, hidden_size=8, output_size=6)
605
+
606
+ params = dict(mlp.named_parameters())
607
+ assert params["gate_up_proj.weight"].shape == torch.Size([16 // tp_size, 4])
608
+ assert params["down_proj.weight"].shape == torch.Size([6, 8 // tp_size])
609
+
610
+ gate_weight = torch.arange(32, dtype=torch.float32).reshape(8, 4)
611
+ up_weight = torch.arange(100, 132, dtype=torch.float32).reshape(8, 4)
612
+ down_weight = torch.arange(48, dtype=torch.float32).reshape(6, 8)
613
+ loaded = mlp.load_weights(
614
+ [
615
+ ("gate_proj.weight", gate_weight),
616
+ ("up_proj.weight", up_weight),
617
+ ("down_proj.weight", down_weight),
618
+ ]
619
+ )
620
+
621
+ assert loaded == {"gate_up_proj.weight", "down_proj.weight"}
622
+ shard = 8 // tp_size
623
+ assert torch.equal(params["gate_up_proj.weight"][:shard], gate_weight[:shard])
624
+ assert torch.equal(params["gate_up_proj.weight"][shard:], up_weight[:shard])
625
+ assert torch.equal(params["down_proj.weight"], down_weight[:, : 8 // tp_size])
626
+
627
+ with set_current_vllm_config(VllmConfig(device_config=DeviceConfig(device="cpu"))):
628
+ packed_mlp = GatedMLP(input_size=4, hidden_size=8, output_size=6)
629
+ packed_params = dict(packed_mlp.named_parameters())
630
+ loaded = packed_mlp.load_weights(
631
+ [("gate_up_proj.weight", torch.cat([gate_weight, up_weight], dim=0))]
632
+ )
633
+ assert loaded == {"gate_up_proj.weight"}
634
+ assert torch.equal(
635
+ packed_params["gate_up_proj.weight"][:shard],
636
+ gate_weight[:shard],
637
+ )
638
+ assert torch.equal(
639
+ packed_params["gate_up_proj.weight"][shard:],
640
+ up_weight[:shard],
641
+ )
642
+
643
+
644
+def test_moss_audio_encoder_loads_realistic_attention_weight_names(monkeypatch):
645
+ from vllm.config import VllmConfig, set_current_vllm_config
646
+ from vllm.config.device import DeviceConfig
647
+
648
+ _patch_tensor_parallel_for_linear_layers(monkeypatch, tp_size=2)
649
+ config = MossAudioEncoderConfig(
650
+ d_model=8,
651
+ output_dim=8,
652
+ num_mel_bins=8,
653
+ encoder_layers=1,
654
+ encoder_attention_heads=2,
655
+ encoder_ffn_dim=16,
656
+ downsample_hidden_size=2,
657
+ deepstack_encoder_layer_indexes=[],
658
+ )
659
+ with set_current_vllm_config(VllmConfig(device_config=DeviceConfig(device="cpu"))):
660
+ encoder = MossAudioEncoder(config)
661
+
662
+ attention = encoder.layers[0].self_attn
663
+ assert all(hasattr(attention, name) for name in ("q_proj", "k_proj", "v_proj"))
664
+ assert hasattr(attention, "out_proj")
665
+ assert not hasattr(attention, "qkv")
666
+ assert attention.k_proj.bias is None
667
+
668
+ weight_names = [
669
+ "layers.0.self_attn.q_proj.weight",
670
+ "layers.0.self_attn.q_proj.bias",
671
+ "layers.0.self_attn.k_proj.weight",
672
+ "layers.0.self_attn.v_proj.weight",
673
+ "layers.0.self_attn.v_proj.bias",
674
+ "layers.0.self_attn.out_proj.weight",
675
+ "layers.0.self_attn.out_proj.bias",
676
+ "conv1.weight",
677
+ "conv1.bias",
678
+ ]
679
+ params = dict(encoder.named_parameters(remove_duplicate=False))
680
+ assert "layers.0.self_attn.k_proj.bias" not in params
681
+ weights = {
682
+ name: torch.full_like(params[name], fill_value=float(i + 1))
683
+ for i, name in enumerate(weight_names)
684
+ }
685
+
686
+ loaded = AutoWeightsLoader(encoder).load_weights(weights.items())
687
+
688
+ assert "load_weights" not in MossAudioEncoder.__dict__
689
+ assert loaded == set(weight_names)
690
+ assert not any(".qkv." in name for name in loaded)
691
+ assert torch.equal(
692
+ params["layers.0.self_attn.q_proj.weight"],
693
+ weights["layers.0.self_attn.q_proj.weight"],
694
+ )
tests/models/registry.py
+9
@@ -1127,6 +1127,15 @@ _MULTIMODAL_EXAMPLE_MODELS = {
1127
tokenizer="moondream/starmie-v1",
1128
trust_remote_code=True,
1129
),
1130
+ "MossAudioModel": _HfExamplesInfo(
1131
+ "OpenMOSS-Team/MOSS-Audio-4B-Instruct",
1132
+ extras={
1133
+ "4b-thinking": "OpenMOSS-Team/MOSS-Audio-4B-Thinking",
1134
+ "8b-instruct": "OpenMOSS-Team/MOSS-Audio-8B-Instruct",
1135
+ "8b-thinking": "OpenMOSS-Team/MOSS-Audio-8B-Thinking",
1136
+ },
1137
+ trust_remote_code=True,
1138
+ ),
1139
"HfMoondream": _HfExamplesInfo(
1140
"moondream/moondream3-preview",
1141
tokenizer="moondream/starmie-v1",
vllm/benchmarks/throughput.py
+45
-1
@@ -21,6 +21,7 @@ from vllm.benchmarks.datasets import (
21
BenchmarkDataset,
22
BurstGPTDataset,
23
ConversationDataset,
24
+ CustomAudioDataset,
25
InstructCoderDataset,
26
MultiModalConversationDataset,
27
PrefixRepetitionRandomDataset,
@@ -257,9 +258,24 @@ def _run_vllm_chat_requests(
258
) -> tuple[float, list[RequestOutput]]:
259
from vllm import SamplingParams
260
260
- prompts = [request.prompt for request in requests]
261
+ prompts = []
262
sampling_params: list[SamplingParams] = []
263
for request in requests:
264
+ if isinstance(request.prompt, list):
265
+ prompts.append(request.prompt)
266
+ else:
267
+ content: list[dict[str, Any]] = [{"type": "text", "text": request.prompt}]
268
+ if request.multi_modal_data is not None:
269
+ if isinstance(request.multi_modal_data, list):
270
+ content.extend(request.multi_modal_data)
271
+ elif isinstance(request.multi_modal_data, dict):
272
+ content.append(request.multi_modal_data)
273
+ else:
274
+ raise TypeError(
275
+ "Could not process multimodal content of type: "
276
+ f"{type(request.multi_modal_data)}"
277
+ )
278
+ prompts.append([{"role": "user", "content": content}])
279
sampling_params.append(
280
SamplingParams(
281
n=n,
@@ -533,6 +549,7 @@ def get_requests(args, tokenizer):
549
"max_loras": args.max_loras,
550
"lora_assignment": getattr(args, "lora_assignment", "random"),
551
"num_requests": args.num_prompts,
552
+ "no_oversample": getattr(args, "no_oversample", False),
553
}
554
555
if args.dataset_name == "random" or (
@@ -573,6 +590,16 @@ def get_requests(args, tokenizer):
590
sample_kwargs["output_len"] = args.output_len
591
elif args.dataset_name == "burstgpt":
592
dataset_cls = BurstGPTDataset
593
+ elif args.dataset_name == "custom_audio":
594
+ dataset_cls = CustomAudioDataset
595
+ sample_kwargs["enable_multimodal_chat"] = getattr(
596
+ args, "enable_multimodal_chat", False
597
+ )
598
+ custom_output_len = getattr(args, "custom_output_len", None)
599
+ if custom_output_len is not None:
600
+ sample_kwargs["output_len"] = custom_output_len
601
+ elif args.output_len is not None:
602
+ sample_kwargs["output_len"] = args.output_len
603
elif args.dataset_name == "hf":
604
if args.output_len is not None:
605
sample_kwargs["output_len"] = args.output_len
@@ -917,6 +944,7 @@ def add_cli_args(parser: FlexibleArgumentParser):
944
"prefix_repetition",
945
"random-mm",
946
"random-rerank",
947
+ "custom_audio",
948
],
949
help="Name of the dataset to benchmark on.",
950
default="sharegpt",
@@ -933,6 +961,22 @@ def add_cli_args(parser: FlexibleArgumentParser):
961
parser.add_argument(
962
"--dataset-path", type=str, default=None, help="Path to the dataset"
963
)
964
+ parser.add_argument(
965
+ "--no-oversample",
966
+ action="store_true",
967
+ help="Do not oversample if the dataset has fewer samples than num-prompts.",
968
+ )
969
+ parser.add_argument(
970
+ "--enable-multimodal-chat",
971
+ action="store_true",
972
+ help="Enable multimodal chat transformation for datasets that support it.",
973
+ )
974
+ parser.add_argument(
975
+ "--custom-output-len",
976
+ type=int,
977
+ default=None,
978
+ help="Number of output tokens per request for custom datasets.",
979
+ )
980
parser.add_argument(
981
"--input-len",
982
type=int,
vllm/model_executor/models/moss_audio.py
new
+1892
@@ -0,0 +1,1892 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""Inference-only MOSS-Audio model compatible with HuggingFace weights."""
4
+
5
+import math
6
+from collections.abc import Iterable, Mapping, Sequence
7
+from dataclasses import dataclass, field
8
+from typing import Annotated, Any
9
+
10
+import numpy as np
11
+import regex as re
12
+import torch
13
+import torch.nn.functional as F
14
+from torch import nn
15
+from transformers import BatchFeature, PretrainedConfig, Qwen3Config
16
+from transformers.models.whisper import WhisperFeatureExtractor
17
+
18
+from vllm.compilation.decorators import support_torch_compile
19
+from vllm.config import VllmConfig
20
+from vllm.config.multimodal import BaseDummyOptions
21
+from vllm.distributed import (
22
+ get_pp_group,
23
+ get_tensor_model_parallel_world_size,
24
+)
25
+from vllm.inputs import ModalityData, MultiModalDataDict
26
+from vllm.model_executor.layers.activation import _ACTIVATION_REGISTRY, SiluAndMul
27
+from vllm.model_executor.layers.linear import (
28
+ ColumnParallelLinear,
29
+ MergedColumnParallelLinear,
30
+ ReplicatedLinear,
31
+ RowParallelLinear,
32
+)
33
+from vllm.model_executor.layers.logits_processor import LogitsProcessor
34
+from vllm.model_executor.layers.quantization import QuantizationConfig
35
+from vllm.model_executor.layers.vocab_parallel_embedding import ParallelLMHead
36
+from vllm.model_executor.model_loader.weight_utils import default_weight_loader
37
+from vllm.multimodal import MULTIMODAL_REGISTRY
38
+from vllm.multimodal.inputs import (
39
+ AudioItem,
40
+ MultiModalFieldConfig,
41
+ MultiModalKwargsItems,
42
+)
43
+from vllm.multimodal.parse import (
44
+ DictEmbeddingItems,
45
+ ModalityDataItems,
46
+ MultiModalDataItems,
47
+ MultiModalDataParser,
48
+)
49
+from vllm.multimodal.processing import (
50
+ BaseDummyInputsBuilder,
51
+ BaseMultiModalProcessor,
52
+ BaseProcessingInfo,
53
+ PromptReplacement,
54
+ PromptUpdate,
55
+ PromptUpdateDetails,
56
+)
57
+from vllm.sequence import IntermediateTensors
58
+from vllm.transformers_utils.repo_utils import get_hf_file_to_dict
59
+from vllm.utils.tensor_schema import TensorSchema, TensorShape
60
+
61
+from .interfaces import (
62
+ MultiModalEmbeddings,
63
+ SupportsLoRA,
64
+ SupportsMultiModal,
65
+ SupportsPP,
66
+ _require_is_multimodal,
67
+)
68
+from .module_mapping import MultiModelKeys
69
+from .qwen3 import Qwen3ForCausalLM, Qwen3Model
70
+from .utils import (
71
+ AutoWeightsLoader,
72
+ WeightsMapper,
73
+ _merge_multimodal_embeddings,
74
+ maybe_prefix,
75
+)
76
+
77
+MOSS_AUDIO_TOKEN = "<|AUDIO|>"
78
+MOSS_AUDIO_BOS_TOKEN = "<|audio_bos|>"
79
+MOSS_AUDIO_EOS_TOKEN = "<|audio_eos|>"
80
+MOSS_AUDIO_TOKEN_ID = 151654
81
+MOSS_AUDIO_BOS_TOKEN_ID = 151669
82
+MOSS_AUDIO_EOS_TOKEN_ID = 151670
83
+DEFAULT_MAX_AUDIO_SECONDS = 30
84
+DEFAULT_MOSS_AUDIO_MEL_CONFIG = {
85
+ "mel_dim": 128,
86
+ "mel_sr": 16000,
87
+ "mel_hop_length": 160,
88
+ "mel_n_fft": 400,
89
+}
90
+MOSS_AUDIO_PLACEHOLDER = (
91
+ f"{MOSS_AUDIO_BOS_TOKEN}{MOSS_AUDIO_TOKEN}{MOSS_AUDIO_EOS_TOKEN}"
92
+)
93
+MOSS_AUDIO_SPAN_RE = re.compile(
94
+ f"{re.escape(MOSS_AUDIO_BOS_TOKEN)}"
95
+ f"(?:{re.escape(MOSS_AUDIO_TOKEN)})+"
96
+ f"{re.escape(MOSS_AUDIO_EOS_TOKEN)}"
97
+)
98
+MOSS_AUDIO_PROCESSOR_CONFIG_KEYS = {
99
+ "audio_token_id",
100
+ "audio_start_id",
101
+ "audio_end_id",
102
+ "enable_time_marker",
103
+ "mel_config",
104
+}
105
+
106
+
107
+class MossAudioAudioInputs(TensorSchema):
108
+ """
109
+ Dimensions:
110
+ - b: Batch size
111
+ - nmb: Number of mel bins
112
+ - t: Time frames
113
+ """
114
+
115
+ audio_data: Annotated[torch.Tensor, TensorShape("b", "nmb", "t")]
116
+ audio_data_seqlens: Annotated[torch.Tensor, TensorShape("b")]
117
+
118
+
119
+def _normalize_moss_audio_mel_config(
120
+ mel_config: Mapping[str, object] | None = None,
121
+) -> dict[str, int]:
122
+ config = dict(DEFAULT_MOSS_AUDIO_MEL_CONFIG)
123
+ config.update(_extract_moss_audio_mel_config(mel_config))
124
+ return config
125
+
126
+
127
+def _extract_moss_audio_mel_config(
128
+ mel_config: Mapping[str, object] | None = None,
129
+) -> dict[str, int]:
130
+ config: dict[str, int] = {}
131
+ if mel_config is None:
132
+ return config
133
+
134
+ aliases = {
135
+ "mel_dim": ("mel_dim", "feature_size", "n_mels", "num_mel_bins"),
136
+ "mel_sr": ("mel_sr", "sampling_rate", "sample_rate"),
137
+ "mel_hop_length": ("mel_hop_length", "hop_length"),
138
+ "mel_n_fft": ("mel_n_fft", "n_fft"),
139
+ }
140
+ for target_key, source_keys in aliases.items():
141
+ for source_key in source_keys:
142
+ if source_key in mel_config:
143
+ config[target_key] = int(mel_config[source_key])
144
+ break
145
+
146
+ return config
147
+
148
+
149
+def _filter_moss_audio_processor_config(
150
+ config: Mapping[str, object] | None,
151
+) -> dict[str, object]:
152
+ if not config:
153
+ return {}
154
+
155
+ return {
156
+ key: value
157
+ for key, value in config.items()
158
+ if key in MOSS_AUDIO_PROCESSOR_CONFIG_KEYS
159
+ }
160
+
161
+
162
+def _merge_moss_audio_processor_configs(
163
+ *configs: Mapping[str, object] | None,
164
+) -> dict[str, object]:
165
+ merged: dict[str, object] = {}
166
+ merged_mel_config: dict[str, int] = {}
167
+ for config in configs:
168
+ filtered = _filter_moss_audio_processor_config(config)
169
+ mel_config = filtered.pop("mel_config", None)
170
+ merged.update(filtered)
171
+ if isinstance(mel_config, Mapping):
172
+ merged_mel_config.update(_extract_moss_audio_mel_config(mel_config))
173
+
174
+ if merged_mel_config:
175
+ merged["mel_config"] = merged_mel_config
176
+ return merged
177
+
178
+
179
+@dataclass
180
+class MossAudioEncoderConfig:
181
+ d_model: int = 1280
182
+ output_dim: int = 1280
183
+ num_mel_bins: int = 128
184
+ encoder_layers: int = 32
185
+ encoder_attention_heads: int = 20
186
+ encoder_ffn_dim: int = 5120
187
+ downsample_rate: int = 8
188
+ downsample_hidden_size: int = 480
189
+ encoder_attention_window_size: int = 100
190
+ max_source_positions: int = 1500
191
+ dropout: float = 0.1
192
+ attention_dropout: float = 0.1
193
+ activation_dropout: float = 0.0
194
+ activation_function: str = "gelu"
195
+ layer_norm_eps: float = 1e-5
196
+ _attn_implementation: str = "eager"
197
+ pretrained_path: str = ""
198
+ n_window: int = 200
199
+ conv_chunksize: int = 64
200
+ deepstack_encoder_layer_indexes: list[int] = field(
201
+ default_factory=lambda: [8, 16, 24]
202
+ )
203
+
204
+ @classmethod
205
+ def from_config(cls, config: object) -> "MossAudioEncoderConfig":
206
+ if isinstance(config, cls):
207
+ return config
208
+ if isinstance(config, Mapping):
209
+ values = {
210
+ key: value
211
+ for key, value in config.items()
212
+ if key in cls.__dataclass_fields__
213
+ }
214
+ else:
215
+ values = {
216
+ key: getattr(config, key)
217
+ for key in cls.__dataclass_fields__
218
+ if hasattr(config, key)
219
+ }
220
+ return cls(**values)
221
+
222
+
223
+class MossAudioConfig(PretrainedConfig):
224
+ model_type = "moss_audio"
225
+ is_composition = True
226
+
227
+ def __init__(
228
+ self,
229
+ audio_config: Mapping[str, object] | MossAudioEncoderConfig | None = None,
230
+ language_config: Mapping[str, object] | Qwen3Config | None = None,
231
+ adapter_hidden_size: int = 8192,
232
+ ignore_index: int = -100,
233
+ deepstack_num_inject_layers: int | None = None,
234
+ **kwargs: object,
235
+ ) -> None:
236
+ self.audio_config = MossAudioEncoderConfig.from_config(audio_config or {})
237
+ if isinstance(language_config, Qwen3Config):
238
+ self.language_config = language_config
239
+ else:
240
+ self.language_config = Qwen3Config(**(language_config or {}))
241
+
242
+ self.adapter_hidden_size = adapter_hidden_size
243
+ self.ignore_index = ignore_index
244
+ self.deepstack_num_inject_layers = deepstack_num_inject_layers
245
+
246
+ for key in ("num_hidden_layers", "eos_token_id", "bos_token_id", "vocab_size"):
247
+ kwargs.setdefault(key, getattr(self.language_config, key, None))
248
+ kwargs.setdefault("tie_word_embeddings", False)
249
+ super().__init__(**kwargs)
250
+
251
+ for key in (
252
+ "hidden_size",
253
+ "num_attention_heads",
254
+ "num_key_value_heads",
255
+ "head_dim",
256
+ "max_position_embeddings",
257
+ "rms_norm_eps",
258
+ ):
259
+ if hasattr(self.language_config, key):
260
+ setattr(self, key, getattr(self.language_config, key))
261
+
262
+ def get_text_config(self, decoder: bool = False) -> Qwen3Config:
263
+ return self.language_config
264
+
265
+
266
+class SinusoidsPositionEmbedding(nn.Module):
267
+ def __init__(self, num_positions: int, embedding_dim: int) -> None:
268
+ super().__init__()
269
+ del num_positions # Kept for config compatibility.
270
+ max_timescale = 10000.0
271
+ log_timescale_increment = math.log(max_timescale) / (embedding_dim // 2 - 1)
272
+ inv_timescales = torch.exp(
273
+ -log_timescale_increment * torch.arange(embedding_dim // 2).float()
274
+ )
275
+ self.register_buffer("inv_timescales", inv_timescales, persistent=False)
276
+
277
+ def forward(self, seq_len: int, device: torch.device) -> torch.Tensor:
278
+ scaled_time = (
279
+ torch.arange(seq_len, device=device, dtype=self.inv_timescales.dtype)[
280
+ :, None
281
+ ]
282
+ * self.inv_timescales[None, :]
283
+ )
284
+ return torch.cat([scaled_time.sin(), scaled_time.cos()], dim=1).unsqueeze(0)
285
+
286
+
287
+class MossAudioAttention(nn.Module):
288
+ def __init__(
289
+ self,
290
+ config: MossAudioEncoderConfig,
291
+ quant_config: QuantizationConfig | None = None,
292
+ prefix: str = "",
293
+ ) -> None:
294
+ super().__init__()
295
+ self.embed_dim = config.d_model
296
+ self.num_heads = config.encoder_attention_heads
297
+ self.head_dim = self.embed_dim // self.num_heads
298
+ if self.head_dim * self.num_heads != self.embed_dim:
299
+ raise ValueError(
300
+ f"d_model ({self.embed_dim}) must be divisible by "
301
+ f"encoder_attention_heads ({self.num_heads})."
302
+ )
303
+
304
+ tp_size = get_tensor_model_parallel_world_size()
305
+ if self.num_heads % tp_size != 0:
306
+ raise ValueError(
307
+ "MOSS-Audio audio encoder attention heads must be divisible by "
308
+ f"tensor parallel size. Got {self.num_heads=} and {tp_size=}."
309
+ )
310
+ self.num_local_heads = self.num_heads // tp_size
311
+ # TODO: can use QKVParallelLinear
312
+ self.q_proj = ColumnParallelLinear(
313
+ input_size=self.embed_dim,
314
+ output_size=self.embed_dim,
315
+ bias=True,
316
+ quant_config=quant_config,
317
+ prefix=f"{prefix}.q_proj",
318
+ )
319
+ self.k_proj = ColumnParallelLinear(
320
+ input_size=self.embed_dim,
321
+ output_size=self.embed_dim,
322
+ bias=False,
323
+ quant_config=quant_config,
324
+ prefix=f"{prefix}.k_proj",
325
+ )
326
+ self.v_proj = ColumnParallelLinear(
327
+ input_size=self.embed_dim,
328
+ output_size=self.embed_dim,
329
+ bias=True,
330
+ quant_config=quant_config,
331
+ prefix=f"{prefix}.v_proj",
332
+ )
333
+ self.out_proj = RowParallelLinear(
334
+ input_size=self.embed_dim,
335
+ output_size=self.embed_dim,
336
+ bias=True,
337
+ quant_config=quant_config,
338
+ prefix=f"{prefix}.out_proj",
339
+ )
340
+
341
+ def forward(
342
+ self,
343
+ hidden_states: torch.Tensor,
344
+ attention_mask: torch.Tensor,
345
+ ) -> torch.Tensor:
346
+ batch_size, seq_len, _ = hidden_states.shape
347
+ q, _ = self.q_proj(hidden_states)
348
+ k, _ = self.k_proj(hidden_states)
349
+ v, _ = self.v_proj(hidden_states)
350
+ q = q.view(batch_size, seq_len, -1, self.head_dim).transpose(1, 2)
351
+ k = k.view(batch_size, seq_len, -1, self.head_dim).transpose(1, 2)
352
+ v = v.view(batch_size, seq_len, -1, self.head_dim).transpose(1, 2)
353
+ attn_output = F.scaled_dot_product_attention(
354
+ q,
355
+ k,
356
+ v,
357
+ attn_mask=attention_mask[:, None, None, :],
358
+ dropout_p=0.0,
359
+ scale=self.head_dim**-0.5,
360
+ )
361
+ output, _ = self.out_proj(
362
+ attn_output.transpose(1, 2).reshape(
363
+ batch_size,
364
+ seq_len,
365
+ -1,
366
+ )
367
+ )
368
+ return output
369
+
370
+
371
+class MossAudioEncoderLayer(nn.Module):
372
+ def __init__(
373
+ self,
374
+ config: MossAudioEncoderConfig,
375
+ quant_config: QuantizationConfig | None = None,
376
+ prefix: str = "",
377
+ ) -> None:
378
+ super().__init__()
379
+ self.self_attn = MossAudioAttention(
380
+ config,
381
+ quant_config=quant_config,
382
+ prefix=f"{prefix}.self_attn",
383
+ )
384
+ self.self_attn_layer_norm = nn.LayerNorm(
385
+ config.d_model, eps=config.layer_norm_eps
386
+ )
387
+ self.activation_fn = _ACTIVATION_REGISTRY[config.activation_function]
388
+ self.activation_dropout = config.activation_dropout
389
+ self.dropout = config.dropout
390
+ self.fc1 = ColumnParallelLinear(
391
+ config.d_model,
392
+ config.encoder_ffn_dim,
393
+ bias=True,
394
+ quant_config=quant_config,
395
+ prefix=f"{prefix}.fc1",
396
+ )
397
+ self.fc2 = RowParallelLinear(
398
+ config.encoder_ffn_dim,
399
+ config.d_model,
400
+ bias=True,
401
+ quant_config=quant_config,
402
+ prefix=f"{prefix}.fc2",
403
+ )
404
+ self.final_layer_norm = nn.LayerNorm(config.d_model, eps=config.layer_norm_eps)
405
+
406
+ def forward(
407
+ self,
408
+ hidden_states: torch.Tensor,
409
+ attention_mask: torch.Tensor,
410
+ ) -> torch.Tensor:
411
+ residual = hidden_states
412
+ hidden_states = self.self_attn_layer_norm(hidden_states)
413
+ hidden_states = self.self_attn(hidden_states, attention_mask)
414
+ hidden_states = residual + F.dropout(
415
+ hidden_states, p=self.dropout, training=self.training
416
+ )
417
+
418
+ residual = hidden_states
419
+ hidden_states = self.final_layer_norm(hidden_states)
420
+ hidden_states, _ = self.fc1(hidden_states)
421
+ hidden_states = self.activation_fn(hidden_states)
422
+ hidden_states = F.dropout(
423
+ hidden_states, p=self.activation_dropout, training=self.training
424
+ )
425
+ hidden_states, _ = self.fc2(hidden_states)
426
+ hidden_states = residual + F.dropout(
427
+ hidden_states, p=self.dropout, training=self.training
428
+ )
429
+ return hidden_states
430
+
431
+
432
+class MossAudioEncoder(nn.Module):
433
+ def __init__(
434
+ self,
435
+ config: MossAudioEncoderConfig,
436
+ quant_config: QuantizationConfig | None = None,
437
+ prefix: str = "",
438
+ ) -> None:
439
+ super().__init__()
440
+ self.config = config
441
+ self.gelu = nn.GELU()
442
+ self.conv1 = nn.Conv2d(
443
+ 1,
444
+ config.downsample_hidden_size,
445
+ kernel_size=(3, 3),
446
+ stride=(2, 2),
447
+ padding=(1, 1),
448
+ )
449
+ self.conv2 = nn.Conv2d(
450
+ config.downsample_hidden_size,
451
+ config.downsample_hidden_size,
452
+ kernel_size=(3, 3),
453
+ stride=(2, 2),
454
+ padding=(1, 1),
455
+ )
456
+ self.conv3 = nn.Conv2d(
457
+ config.downsample_hidden_size,
458
+ config.downsample_hidden_size,
459
+ kernel_size=(3, 3),
460
+ stride=(2, 2),
461
+ padding=(1, 1),
462
+ )
463
+
464
+ conv_freq = self._compute_downsampled_length(
465
+ torch.tensor(config.num_mel_bins)
466
+ ).item()
467
+ self.stem_proj = ReplicatedLinear(
468
+ config.downsample_hidden_size * int(conv_freq),
469
+ config.d_model,
470
+ bias=True,
471
+ quant_config=quant_config,
472
+ return_bias=False,
473
+ prefix=f"{prefix}.stem_proj",
474
+ )
475
+ self.embed_positions = SinusoidsPositionEmbedding(
476
+ config.max_source_positions, config.d_model
477
+ )
478
+ self.layers = nn.ModuleList(
479
+ [
480
+ MossAudioEncoderLayer(
481
+ config,
482
+ quant_config=quant_config,
483
+ prefix=f"{prefix}.layers.{layer_idx}",
484
+ )
485
+ for layer_idx in range(config.encoder_layers)
486
+ ]
487
+ )
488
+ self.layer_norm = nn.LayerNorm(config.d_model, eps=config.layer_norm_eps)
489
+ if config.output_dim != config.d_model:
490
+ self.out_proj = ReplicatedLinear(
491
+ config.d_model,
492
+ config.output_dim,
493
+ bias=False,
494
+ quant_config=quant_config,
495
+ return_bias=False,
496
+ prefix=f"{prefix}.out_proj",
497
+ )
498
+ else:
499
+ self.out_proj = nn.Identity()
500
+
501
+ self.deepstack_encoder_layer_indexes = list(
502
+ config.deepstack_encoder_layer_indexes or []
503
+ )
504
+ self._deepstack_capture_map = {
505
+ layer_idx: capture_idx
506
+ for capture_idx, layer_idx in enumerate(
507
+ self.deepstack_encoder_layer_indexes
508
+ )
509
+ }
510
+ self.n_window = int(config.n_window)
511
+ self.chunk_frames = int(self.n_window * 2)
512
+ self.conv_chunksize = int(config.conv_chunksize)
513
+
514
+ @property
515
+ def dtype(self) -> torch.dtype:
516
+ return self.conv1.weight.dtype
517
+
518
+ @staticmethod
519
+ def _compute_downsampled_length(lengths: torch.Tensor) -> torch.Tensor:
520
+ def conv_out_len(length: torch.Tensor) -> torch.Tensor:
521
+ return (length - 1) // 2 + 1
522
+
523
+ return conv_out_len(conv_out_len(conv_out_len(lengths)))
524
+
525
+ @staticmethod
526
+ def compute_num_audio_tokens(raw_mel_len: int) -> int:
527
+ lengths = torch.tensor(raw_mel_len, dtype=torch.long)
528
+ return int(MossAudioEncoder._compute_downsampled_length(lengths).item())
529
+
530
+ def _encode_chunk_batch(
531
+ self,
532
+ input_features: torch.Tensor,
533
+ seq_lengths: torch.Tensor,
534
+ output_deepstack_hidden_states: bool = True,
535
+ ) -> tuple[torch.Tensor, list[torch.Tensor]]:
536
+ if input_features.dim() == 2:
537
+ input_features = input_features.unsqueeze(0)
538
+
539
+ downsampled_lengths = self._compute_downsampled_length(seq_lengths)
540
+ x = input_features.unsqueeze(1)
541
+ x = self.gelu(self.conv1(x))
542
+ x = self.gelu(self.conv2(x))
543
+ x = self.gelu(self.conv3(x))
544
+ x = x.permute(0, 3, 1, 2).contiguous().flatten(2)
545
+ x = self.stem_proj(x)
546
+
547
+ max_len = int(downsampled_lengths.max().item())
548
+ if x.size(1) > max_len:
549
+ x = x[:, :max_len, :]
550
+ x = x + self.embed_positions(x.shape[1], x.device).to(x.dtype)
551
+
552
+ attention_mask = (
553
+ torch.arange(x.size(1), device=x.device)[None, :]
554
+ < downsampled_lengths[:, None]
555
+ )
556
+
557
+ deepstack_hidden_states: list[torch.Tensor | None] = []
558
+ if output_deepstack_hidden_states:
559
+ deepstack_hidden_states = [None] * len(self.deepstack_encoder_layer_indexes)
560
+ for layer_idx, layer in enumerate(self.layers):
561
+ x = layer(x, attention_mask)
562
+ if output_deepstack_hidden_states:
563
+ capture_idx = self._deepstack_capture_map.get(layer_idx)
564
+ if capture_idx is not None:
565
+ deepstack_hidden_states[capture_idx] = x
566
+
567
+ x = self.layer_norm(x)
568
+ x = self.out_proj(x)
569
+
570
+ if not output_deepstack_hidden_states:
571
+ return x, []
572
+
573
+ ordered_deepstack_hidden_states = [
574
+ hidden_states
575
+ for hidden_states in deepstack_hidden_states
576
+ if hidden_states is not None
577
+ ]
578
+ ordered_deepstack_hidden_states = [
579
+ self.out_proj(hidden_states)
580
+ for hidden_states in ordered_deepstack_hidden_states
581
+ ]
582
+ return x, ordered_deepstack_hidden_states
583
+
584
+ def forward(
585
+ self,
586
+ input_features: torch.Tensor,
587
+ feature_lens: torch.Tensor | None = None,
588
+ output_deepstack_hidden_states: bool = True,
589
+ ) -> tuple[torch.Tensor, tuple[torch.Tensor, ...] | None]:
590
+ if input_features.dim() == 3:
591
+ if feature_lens is None:
592
+ feature_lens = torch.full(
593
+ (input_features.size(0),),
594
+ input_features.size(-1),
595
+ dtype=torch.long,
596
+ device=input_features.device,
597
+ )
598
+ else:
599
+ feature_lens = feature_lens.to(
600
+ device=input_features.device, dtype=torch.long
601
+ )
602
+ valid_chunks = [
603
+ input_features[i, :, : int(feature_lens[i].item())]
604
+ for i in range(int(input_features.shape[0]))
605
+ ]
606
+ input_features = torch.cat(valid_chunks, dim=1)
607
+ elif input_features.dim() != 2:
608
+ raise ValueError(
609
+ f"Expected [n_mels, T] or [B, n_mels, T], got "
610
+ f"{tuple(input_features.shape)}."
611
+ )
612
+
613
+ if feature_lens is None:
614
+ feature_lens = torch.tensor(
615
+ [int(input_features.shape[1])],
616
+ device=input_features.device,
617
+ dtype=torch.long,
618
+ )
619
+ else:
620
+ feature_lens = feature_lens.to(
621
+ device=input_features.device, dtype=torch.long
622
+ )
623
+
624
+ chunk_num = torch.ceil(
625
+ feature_lens.to(torch.float32) / self.chunk_frames
626
+ ).long()
627
+ chunk_lengths = torch.full(
628
+ (int(chunk_num.sum().item()),),
629
+ self.chunk_frames,
630
+ dtype=torch.long,
631
+ device=feature_lens.device,
632
+ )
633
+ tail_chunk_index = F.pad(chunk_num, (1, 0), value=-1).cumsum(0)[1:]
634
+ chunk_lengths[tail_chunk_index] = feature_lens % self.chunk_frames
635
+ chunk_lengths[chunk_lengths == 0] = self.chunk_frames
636
+
637
+ chunk_list = input_features.T.split(chunk_lengths.tolist(), dim=0)
638
+ padded_feature = nn.utils.rnn.pad_sequence(
639
+ chunk_list, batch_first=True
640
+ ).transpose(1, 2)
641
+
642
+ feature_lens_after_cnn = self._compute_downsampled_length(chunk_lengths)
643
+ t_down_max = (
644
+ int(feature_lens_after_cnn.max().item())
645
+ if feature_lens_after_cnn.numel() > 0
646
+ else 0
647
+ )
648
+ indices = torch.arange(t_down_max, device=padded_feature.device)
649
+ padded_mask_after_cnn = indices[None, :] < feature_lens_after_cnn[:, None]
650
+
651
+ num_deepstack = len(self.deepstack_encoder_layer_indexes)
652
+ should_output_deepstack = output_deepstack_hidden_states and num_deepstack > 0
653
+ padded_embeds: list[torch.Tensor] = []
654
+ deepstack_padded_embeds: list[list[torch.Tensor]] = [
655
+ [] for _ in range(num_deepstack if should_output_deepstack else 0)
656
+ ]
657
+ for feat_chunk, len_chunk in zip(
658
+ padded_feature.split(self.conv_chunksize, dim=0),
659
+ chunk_lengths.split(self.conv_chunksize, dim=0),
660
+ ):
661
+ out, deepstack_outs = self._encode_chunk_batch(
662
+ feat_chunk,
663
+ len_chunk,
664
+ output_deepstack_hidden_states=should_output_deepstack,
665
+ )
666
+ if out.shape[1] < t_down_max:
667
+ out = F.pad(out, (0, 0, 0, t_down_max - out.shape[1]))
668
+ padded_embeds.append(out)
669
+
670
+ if should_output_deepstack:
671
+ if len(deepstack_outs) != num_deepstack:
672
+ raise RuntimeError(
673
+ "DeepStack output count does not match configured "
674
+ "layer indexes."
675
+ )
676
+ for capture_idx, ds in enumerate(deepstack_outs):
677
+ if ds.shape[1] < t_down_max:
678
+ ds = F.pad(ds, (0, 0, 0, t_down_max - ds.shape[1]))
679
+ deepstack_padded_embeds[capture_idx].append(ds)
680
+
681
+ if padded_embeds:
682
+ padded_embed = torch.cat(padded_embeds, dim=0)
683
+ else:
684
+ padded_embed = torch.empty(
685
+ (0, t_down_max, self.config.output_dim),
686
+ device=padded_feature.device,
687
+ dtype=padded_feature.dtype,
688
+ )
689
+
690
+ last_hidden_state = padded_embed[padded_mask_after_cnn].unsqueeze(0)
691
+
692
+ deepstack_states: tuple[torch.Tensor, ...] | None = None
693
+ if should_output_deepstack:
694
+ collected: list[torch.Tensor] = []
695
+ for chunks_list in deepstack_padded_embeds:
696
+ if chunks_list:
697
+ ds = torch.cat(chunks_list, dim=0)
698
+ collected.append(ds[padded_mask_after_cnn].unsqueeze(0))
699
+ else:
700
+ collected.append(
701
+ torch.empty(
702
+ (1, 0, self.config.output_dim),
703
+ device=padded_feature.device,
704
+ dtype=padded_embed.dtype,
705
+ )
706
+ )
707
+ deepstack_states = tuple(collected)
708
+
709
+ return last_hidden_state, deepstack_states
710
+
711
+
712
+class GatedMLP(nn.Module):
713
+ def __init__(
714
+ self,
715
+ input_size: int,
716
+ hidden_size: int,
717
+ output_size: int,
718
+ quant_config: QuantizationConfig | None = None,
719
+ prefix: str = "",
720
+ ) -> None:
721
+ super().__init__()
722
+ self.gate_up_proj = MergedColumnParallelLinear(
723
+ input_size,
724
+ [hidden_size, hidden_size],
725
+ bias=False,
726
+ quant_config=quant_config,
727
+ prefix=f"{prefix}.gate_up_proj",
728
+ )
729
+ self.down_proj = RowParallelLinear(
730
+ hidden_size,
731
+ output_size,
732
+ bias=False,
733
+ input_is_parallel=True,
734
+ quant_config=quant_config,
735
+ prefix=f"{prefix}.down_proj",
736
+ )
737
+ self.act_fn = SiluAndMul()
738
+
739
+ def forward(self, x: torch.Tensor) -> torch.Tensor:
740
+ gate_up, _ = self.gate_up_proj(x)
741
+ x = self.act_fn(gate_up)
742
+ x, _ = self.down_proj(x)
743
+ return x
744
+
745
+ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
746
+ stacked_params_mapping = [
747
+ ("gate_up_proj", "gate_proj", 0),
748
+ ("gate_up_proj", "up_proj", 1),
749
+ ]
750
+ params_dict = dict(self.named_parameters())
751
+ loaded_params: set[str] = set()
752
+
753
+ for name, loaded_weight in weights:
754
+ target_name = name
755
+ for param_name, weight_name, shard_id in stacked_params_mapping:
756
+ components = target_name.split(".")
757
+ if weight_name not in components:
758
+ continue
759
+
760
+ target_name = ".".join(
761
+ param_name if component == weight_name else component
762
+ for component in components
763
+ )
764
+ param = params_dict[target_name]
765
+ weight_loader = param.weight_loader
766
+ weight_loader(param, loaded_weight, shard_id)
767
+ break
768
+ else:
769
+ param = params_dict[target_name]
770
+ weight_loader = getattr(param, "weight_loader", default_weight_loader)
771
+ weight_loader(param, loaded_weight)
772
+
773
+ loaded_params.add(target_name)
774
+ return loaded_params
775
+
776
+
777
+@support_torch_compile(
778
+ dynamic_arg_dims={
779
+ "input_ids": 0,
780
+ "positions": -1,
781
+ "intermediate_tensors": 0,
782
+ "inputs_embeds": 0,
783
+ "deepstack_input_embeds": 0,
784
+ }
785
+)
786
+class MossQwen3Model(Qwen3Model):
787
+ def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None:
788
+ super().__init__(vllm_config=vllm_config, prefix=prefix)
789
+ self.deepstack_inject_layer_indices: Iterable[int] = range(0)
790
+
791
+ def forward(
792
+ self,
793
+ input_ids: torch.Tensor | None,
794
+ positions: torch.Tensor,
795
+ intermediate_tensors: IntermediateTensors | None = None,
796
+ inputs_embeds: torch.Tensor | None = None,
797
+ deepstack_input_embeds: IntermediateTensors | None = None,
798
+ ) -> torch.Tensor | IntermediateTensors:
799
+ if get_pp_group().is_first_rank:
800
+ if inputs_embeds is not None:
801
+ hidden_states = inputs_embeds
802
+ else:
803
+ hidden_states = self.embed_input_ids(input_ids)
804
+ residual = None
805
+ else:
806
+ assert intermediate_tensors is not None
807
+ hidden_states = intermediate_tensors["hidden_states"]
808
+ residual = intermediate_tensors["residual"]
809
+
810
+ aux_hidden_states = self._maybe_add_hidden_state([], 0, hidden_states, residual)
811
+ for layer_idx, layer in enumerate(
812
+ self.layers[self.start_layer : self.end_layer],
813
+ start=self.start_layer,
814
+ ):
815
+ hidden_states, residual = layer(positions, hidden_states, residual)
816
+ deepstack_key = f"deepstack_input_embeds_{layer_idx}"
817
+ if (
818
+ deepstack_input_embeds is not None
819
+ and deepstack_key in deepstack_input_embeds.tensors
820
+ ):
821
+ hidden_states = hidden_states + deepstack_input_embeds[deepstack_key]
822
+ self._maybe_add_hidden_state(
823
+ aux_hidden_states,
824
+ layer_idx - self.start_layer + 1,
825
+ hidden_states,
826
+ residual,
827
+ )
828
+
829
+ if not get_pp_group().is_last_rank:
830
+ tensors = {"hidden_states": hidden_states, "residual": residual}
831
+ # Keep the DeepStack PP schema config-driven, but only carry
832
+ # payloads needed by downstream injection points across this rank.
833
+ # Missing downstream payloads are zero-filled below to clear
834
+ # receive buffers instead of leaving stale tensors.
835
+ for layer_idx in self.deepstack_inject_layer_indices:
836
+ if layer_idx < self.end_layer:
837
+ continue
838
+ deepstack_key = f"deepstack_input_embeds_{layer_idx}"
839
+ if (
840
+ deepstack_input_embeds is not None
841
+ and deepstack_key in deepstack_input_embeds.tensors
842
+ ):
843
+ tensors[deepstack_key] = deepstack_input_embeds[deepstack_key]
844
+ else:
845
+ tensors[deepstack_key] = hidden_states.new_zeros(
846
+ hidden_states.shape
847
+ )
848
+ return IntermediateTensors(tensors)
849
+
850
+ hidden_states, _ = self.norm(hidden_states, residual)
851
+ if len(aux_hidden_states) > 0:
852
+ return hidden_states, aux_hidden_states
853
+ return hidden_states
854
+
855
+
856
+class MossQwen3ForCausalLM(Qwen3ForCausalLM):
857
+ def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None:
858
+ super(Qwen3ForCausalLM, self).__init__()
859
+ config = vllm_config.model_config.hf_config
860
+ quant_config = vllm_config.quant_config
861
+
862
+ self.config = config
863
+ self.vllm_config = vllm_config
864
+ self.quant_config = quant_config
865
+ self.model = MossQwen3Model(
866
+ vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model")
867
+ )
868
+
869
+ if get_pp_group().is_last_rank:
870
+ if config.tie_word_embeddings:
871
+ self.lm_head = self.model.embed_tokens
872
+ else:
873
+ self.lm_head = ParallelLMHead(
874
+ config.vocab_size,
875
+ config.hidden_size,
876
+ quant_config=quant_config,
877
+ prefix=maybe_prefix(prefix, "lm_head"),
878
+ )
879
+ else:
880
+ from .utils import PPMissingLayer
881
+
882
+ self.lm_head = PPMissingLayer()
883
+
884
+ self.logits_processor = LogitsProcessor(config.vocab_size)
885
+ self.deepstack_inject_layer_indices: Iterable[int] = range(0)
886
+
887
+ def make_empty_intermediate_tensors(
888
+ self,
889
+ batch_size: int,
890
+ dtype: torch.dtype,
891
+ device: torch.device,
892
+ ) -> IntermediateTensors:
893
+ intermediate_tensors = self.model.make_empty_intermediate_tensors(
894
+ batch_size, dtype, device
895
+ )
896
+ for layer_idx in self.deepstack_inject_layer_indices:
897
+ intermediate_tensors[f"deepstack_input_embeds_{layer_idx}"] = torch.zeros(
898
+ (batch_size, self.config.hidden_size),
899
+ dtype=dtype,
900
+ device=device,
901
+ )
902
+ return intermediate_tensors
903
+
904
+ def forward(
905
+ self,
906
+ input_ids: torch.Tensor | None,
907
+ positions: torch.Tensor,
908
+ intermediate_tensors: IntermediateTensors | None = None,
909
+ inputs_embeds: torch.Tensor | None = None,
910
+ deepstack_input_embeds: IntermediateTensors | None = None,
911
+ ) -> torch.Tensor | IntermediateTensors:
912
+ return self.model(
913
+ input_ids,
914
+ positions,
915
+ intermediate_tensors,
916
+ inputs_embeds,
917
+ deepstack_input_embeds=deepstack_input_embeds,
918
+ )
919
+
920
+
921
+def _moss_audio_field_config(
922
+ hf_inputs: Mapping[str, torch.Tensor],
923
+) -> Mapping[str, MultiModalFieldConfig]:
924
+ return {
925
+ "audio_data": MultiModalFieldConfig.batched("audio"),
926
+ "audio_data_seqlens": MultiModalFieldConfig.batched("audio", keep_on_cpu=True),
927
+ }
928
+
929
+
930
+class MossAudioMultiModalDataParser(MultiModalDataParser):
931
+ def _parse_audio_data(
932
+ self,
933
+ data: dict[str, torch.Tensor] | ModalityData[AudioItem],
934
+ ) -> ModalityDataItems[Any, Any] | None:
935
+ if isinstance(data, dict):
936
+ return DictEmbeddingItems(
937
+ data,
938
+ modality="audio",
939
+ required_fields={"audio_data", "audio_data_seqlens"},
940
+ fields_factory=_moss_audio_field_config,
941
+ )
942
+
943
+ return super()._parse_audio_data(data)
944
+
945
+
946
+class MossAudioProcessor:
947
+ model_input_names = [
948
+ "input_ids",
949
+ "attention_mask",
950
+ "audio_data",
951
+ "audio_data_seqlens",
952
+ ]
953
+
954
+ def __init__(
955
+ self,
956
+ tokenizer: object,
957
+ *,
958
+ audio_token_id: int = MOSS_AUDIO_TOKEN_ID,
959
+ audio_start_id: int = MOSS_AUDIO_BOS_TOKEN_ID,
960
+ audio_end_id: int = MOSS_AUDIO_EOS_TOKEN_ID,
961
+ enable_time_marker: bool = False,
962
+ mel_config: Mapping[str, object] | None = None,
963
+ ) -> None:
964
+ self.tokenizer = tokenizer
965
+ self.audio_token_id = int(audio_token_id)
966
+ self.audio_start_id = int(audio_start_id)
967
+ self.audio_end_id = int(audio_end_id)
968
+ self.enable_time_marker = bool(enable_time_marker)
969
+ self.mel_config = _normalize_moss_audio_mel_config(mel_config)
970
+ self.feature_extractor = WhisperFeatureExtractor(
971
+ feature_size=self.mel_config["mel_dim"],
972
+ sampling_rate=self.mel_config["mel_sr"],
973
+ hop_length=self.mel_config["mel_hop_length"],
974
+ n_fft=self.mel_config["mel_n_fft"],
975
+ )
976
+ self.audio_tokens_per_second = self.mel_config["mel_sr"] / (
977
+ self.mel_config["mel_hop_length"] * 8
978
+ )
979
+ self.time_marker_every_seconds = 2
980
+ self.time_marker_every_audio_tokens = int(
981
+ self.audio_tokens_per_second * self.time_marker_every_seconds
982
+ )
983
+ self._digit_token_ids = {
984
+ "0": 15,
985
+ "1": 16,
986
+ "2": 17,
987
+ "3": 18,
988
+ "4": 19,
989
+ "5": 20,
990
+ "6": 21,
991
+ "7": 22,
992
+ "8": 23,
993
+ "9": 24,
994
+ }
995
+
996
+ @staticmethod
997
+ def conv3_downsample_len(raw_mel_len: int) -> int:
998
+ return MossAudioEncoder.compute_num_audio_tokens(raw_mel_len)
999
+
1000
+ def _extract_mel(self, audio: np.ndarray | torch.Tensor) -> torch.Tensor:
1001
+ if isinstance(audio, torch.Tensor):
1002
+ wav = audio.detach().to("cpu", dtype=torch.float32).numpy()
1003
+ else:
1004
+ wav = np.asarray(audio, dtype=np.float32)
1005
+ if wav.size == 0:
1006
+ raise ValueError("The audio is too short to be represented.")
1007
+ if wav.ndim == 2:
1008
+ wav = wav[0]
1009
+ feats = self.feature_extractor._np_extract_fbank_features(
1010
+ wav[None, ...], device="cpu"
1011
+ )
1012
+ return torch.from_numpy(feats[0])
1013
+
1014
+ def _get_default_audio_prompt(self) -> str:
1015
+ return MOSS_AUDIO_PLACEHOLDER
1016
+
1017
+ def _ensure_audio_placeholders(
1018
+ self,
1019
+ prompt_text: str,
1020
+ num_audios: int,
1021
+ ) -> str:
1022
+ if num_audios == 0 or MOSS_AUDIO_SPAN_RE.search(prompt_text):
1023
+ return prompt_text
1024
+
1025
+ audio_prompt = self._get_default_audio_prompt() * num_audios
1026
+ if prompt_text:
1027
+ return f"{audio_prompt}\n{prompt_text}"
1028
+ return audio_prompt
1029
+
1030
+ def _build_audio_tokens_with_time_markers(self, audio_seq_len: int) -> list[int]:
1031
+ total_duration_seconds = audio_seq_len / self.audio_tokens_per_second
1032
+ num_full_seconds = int(total_duration_seconds)
1033
+ token_ids: list[int] = []
1034
+ audio_tokens_consumed = 0
1035
+ for second in range(
1036
+ self.time_marker_every_seconds,
1037
+ num_full_seconds + 1,
1038
+ self.time_marker_every_seconds,
1039
+ ):
1040
+ marker_pos = (
1041
+ second // self.time_marker_every_seconds
1042
+ ) * self.time_marker_every_audio_tokens
1043
+ audio_segment_len = marker_pos - audio_tokens_consumed
1044
+ if audio_segment_len > 0:
1045
+ token_ids.extend([self.audio_token_id] * audio_segment_len)
1046
+ audio_tokens_consumed += audio_segment_len
1047
+ token_ids.extend(self._digit_token_ids[digit] for digit in str(second))
1048
+
1049
+ remaining = audio_seq_len - audio_tokens_consumed
1050
+ if remaining > 0:
1051
+ token_ids.extend([self.audio_token_id] * remaining)
1052
+ return token_ids
1053
+
1054
+ def build_audio_placeholder_ids(self, num_audio_tokens: int) -> list[int]:
1055
+ if self.enable_time_marker:
1056
+ return self._build_audio_tokens_with_time_markers(num_audio_tokens)
1057
+ return [self.audio_token_id] * num_audio_tokens
1058
+
1059
+ def __call__(
1060
+ self,
1061
+ text: str | Sequence[str] | None = None,
1062
+ audios: Sequence[np.ndarray | torch.Tensor] | None = None,
1063
+ audio: Sequence[np.ndarray | torch.Tensor] | None = None,
1064
+ return_tensors: str = "pt",
1065
+ **kwargs: object,
1066
+ ) -> BatchFeature:
1067
+ """Build text tokens and audio tensors for one MossAudio prompt.
1068
+
1069
+ Example:
1070
+ text="Describe this.", audio=[waveform]
1071
+ -> input_ids contains audio_start, N audio tokens, audio_end
1072
+ -> audio_data has shape [1, mel_dim, max_time]
1073
+ -> mel_dim is the number of mel filter-bank bins, 128 by default
1074
+ -> audio_data_seqlens stores the unpadded mel length
1075
+ """
1076
+ del kwargs
1077
+
1078
+ # Step 1. Normalize text input; this processor handles one prompt.
1079
+ if isinstance(text, (list, tuple)):
1080
+ if len(text) != 1:
1081
+ raise ValueError(f"Expected text batch size 1, got {len(text)}")
1082
+ prompt_text = text[0]
1083
+ elif text is None:
1084
+ prompt_text = ""
1085
+ else:
1086
+ prompt_text = text
1087
+
1088
+ # Step 2. Accept either `audios` or `audio` and normalize to a list.
1089
+ audio_list = audios if audios is not None else audio
1090
+ audio_list = [] if audio_list is None else list(audio_list)
1091
+
1092
+ # Step 3. Convert waveforms to [mel_dim, time] mel features and token
1093
+ # counts. mel_dim is the number of mel filter-bank bins.
1094
+ mels: list[torch.Tensor] = []
1095
+ raw_lengths: list[int] = []
1096
+ token_lens: list[int] = []
1097
+ for one_audio in audio_list:
1098
+ mel = self._extract_mel(one_audio)
1099
+ raw_len = int(mel.shape[-1])
1100
+ num_tokens = self.conv3_downsample_len(raw_len)
1101
+ if raw_len <= 0 or num_tokens <= 0:
1102
+ raise ValueError("The audio is too short to be represented.")
1103
+ mels.append(mel)
1104
+ raw_lengths.append(raw_len)
1105
+ token_lens.append(num_tokens)
1106
+
1107
+ # Step 4. Pad variable-length mel features into a batch tensor.
1108
+ if mels:
1109
+ max_length = max(raw_lengths)
1110
+ audio_batch = torch.zeros(
1111
+ (len(mels), self.mel_config["mel_dim"], max_length),
1112
+ dtype=torch.float32,
1113
+ )
1114
+ for index, mel in enumerate(mels):
1115
+ audio_batch[index, :, : mel.shape[-1]] = mel
1116
+ audio_data_seqlens = torch.tensor(raw_lengths, dtype=torch.long)
1117
+ else:
1118
+ audio_batch = None
1119
+ audio_data_seqlens = None
1120
+
1121
+ # Step 5. Ensure each audio item has a placeholder span in the prompt.
1122
+ prompt_text = self._ensure_audio_placeholders(prompt_text, len(audio_list))
1123
+ input_ids = []
1124
+ cursor = 0
1125
+
1126
+ # Step 6. Text-only path: tokenize and preserve placeholder spans.
1127
+ if not audio_list:
1128
+ for match in MOSS_AUDIO_SPAN_RE.finditer(prompt_text):
1129
+ prefix = prompt_text[cursor : match.start()]
1130
+ input_ids.extend(
1131
+ self.tokenizer.encode(prefix, add_special_tokens=False)
1132
+ )
1133
+ input_ids.extend(
1134
+ [self.audio_start_id, self.audio_token_id, self.audio_end_id]
1135
+ )
1136
+ cursor = match.end()
1137
+ suffix = prompt_text[cursor:]
1138
+ input_ids.extend(self.tokenizer.encode(suffix, add_special_tokens=False))
1139
+ data: dict[str, torch.Tensor] = {
1140
+ "input_ids": torch.tensor([input_ids], dtype=torch.long),
1141
+ "attention_mask": torch.ones((1, len(input_ids)), dtype=torch.long),
1142
+ }
1143
+ return BatchFeature(data=data, tensor_type=return_tensors)
1144
+
1145
+ # Step 7. Audio path: expand each placeholder to its audio-token count.
1146
+ span_iter = iter(MOSS_AUDIO_SPAN_RE.finditer(prompt_text))
1147
+ for item_idx, _ in enumerate(audio_list):
1148
+ match = next(span_iter, None)
1149
+ if match is None:
1150
+ raise ValueError(
1151
+ "Audio placeholder count mismatch: expected one "
1152
+ f"{MOSS_AUDIO_PLACEHOLDER!r} span per audio item."
1153
+ )
1154
+ prefix = prompt_text[cursor : match.start()]
1155
+ input_ids.extend(self.tokenizer.encode(prefix, add_special_tokens=False))
1156
+ input_ids.append(self.audio_start_id)
1157
+ input_ids.extend(self.build_audio_placeholder_ids(token_lens[item_idx]))
1158
+ input_ids.append(self.audio_end_id)
1159
+ cursor = match.end()
1160
+
1161
+ # Step 8. Reject extra placeholder spans after all audio items are used.
1162
+ suffix = prompt_text[cursor:]
1163
+ if MOSS_AUDIO_SPAN_RE.search(suffix):
1164
+ raise ValueError(
1165
+ "Audio placeholder count mismatch: found more placeholder spans "
1166
+ "than audio items."
1167
+ )
1168
+ input_ids.extend(self.tokenizer.encode(suffix, add_special_tokens=False))
1169
+
1170
+ # Step 9. Return tokenizer output plus audio tensors for embed_multimodal.
1171
+ data = {
1172
+ "input_ids": torch.tensor([input_ids], dtype=torch.long),
1173
+ "attention_mask": torch.ones((1, len(input_ids)), dtype=torch.long),
1174
+ }
1175
+ if audio_batch is not None and audio_data_seqlens is not None:
1176
+ data["audio_data"] = audio_batch
1177
+ data["audio_data_seqlens"] = audio_data_seqlens
1178
+ return BatchFeature(data=data, tensor_type=return_tensors)
1179
+
1180
+ def decode(self, *args: object, **kwargs: object) -> str:
1181
+ return self.tokenizer.decode(*args, **kwargs)
1182
+
1183
+ def batch_decode(self, *args: object, **kwargs: object) -> list[str]:
1184
+ return self.tokenizer.batch_decode(*args, **kwargs)
1185
+
1186
+
1187
+class MossAudioProcessingInfo(BaseProcessingInfo):
1188
+ def get_hf_config(self) -> MossAudioConfig:
1189
+ config = self.ctx.get_hf_config()
1190
+ if isinstance(config, MossAudioConfig):
1191
+ return config
1192
+ return MossAudioConfig(
1193
+ audio_config=getattr(config, "audio_config", None),
1194
+ language_config=getattr(config, "language_config", None),
1195
+ adapter_hidden_size=getattr(config, "adapter_hidden_size", 8192),
1196
+ ignore_index=getattr(config, "ignore_index", -100),
1197
+ deepstack_num_inject_layers=getattr(
1198
+ config, "deepstack_num_inject_layers", None
1199
+ ),
1200
+ )
1201
+
1202
+ def _get_processor_config_defaults(self) -> dict[str, object]:
1203
+ cached_defaults = getattr(self, "_processor_config_defaults", None)
1204
+ if cached_defaults is not None:
1205
+ return cached_defaults
1206
+
1207
+ model_config = self.ctx.model_config
1208
+ for file_name in ("processor_config.json", "preprocessor_config.json"):
1209
+ config = get_hf_file_to_dict(
1210
+ file_name,
1211
+ model_config.model,
1212
+ model_config.revision,
1213
+ )
1214
+ defaults = _filter_moss_audio_processor_config(config)
1215
+ if defaults:
1216
+ self._processor_config_defaults = defaults
1217
+ return defaults
1218
+
1219
+ defaults = {}
1220
+ self._processor_config_defaults = defaults
1221
+ return defaults
1222
+
1223
+ @staticmethod
1224
+ def _get_processor_cache_key(kwargs: Mapping[str, object]) -> tuple[object, ...]:
1225
+ mel_config = _normalize_moss_audio_mel_config(
1226
+ kwargs.get("mel_config")
1227
+ if isinstance(kwargs.get("mel_config"), Mapping)
1228
+ else None
1229
+ )
1230
+ return (
1231
+ int(kwargs.get("audio_token_id", MOSS_AUDIO_TOKEN_ID)),
1232
+ int(kwargs.get("audio_start_id", MOSS_AUDIO_BOS_TOKEN_ID)),
1233
+ int(kwargs.get("audio_end_id", MOSS_AUDIO_EOS_TOKEN_ID)),
1234
+ bool(kwargs.get("enable_time_marker", False)),
1235
+ tuple(sorted(mel_config.items())),
1236
+ )
1237
+
1238
+ def get_hf_processor(self, **kwargs: object) -> MossAudioProcessor:
1239
+ merged_kwargs = _merge_moss_audio_processor_configs(
1240
+ self._get_processor_config_defaults(),
1241
+ self.ctx.get_merged_mm_kwargs({}),
1242
+ kwargs,
1243
+ )
1244
+ mel_config = _normalize_moss_audio_mel_config(
1245
+ merged_kwargs.get("mel_config")
1246
+ if isinstance(merged_kwargs.get("mel_config"), Mapping)
1247
+ else None
1248
+ )
1249
+ processor_kwargs = {
1250
+ "audio_token_id": int(
1251
+ merged_kwargs.get("audio_token_id", MOSS_AUDIO_TOKEN_ID)
1252
+ ),
1253
+ "audio_start_id": int(
1254
+ merged_kwargs.get("audio_start_id", MOSS_AUDIO_BOS_TOKEN_ID)
1255
+ ),
1256
+ "audio_end_id": int(
1257
+ merged_kwargs.get("audio_end_id", MOSS_AUDIO_EOS_TOKEN_ID)
1258
+ ),
1259
+ "enable_time_marker": bool(merged_kwargs.get("enable_time_marker", False)),
1260
+ "mel_config": mel_config,
1261
+ }
1262
+
1263
+ cache = getattr(self, "_hf_processor_cache", None)
1264
+ if cache is None:
1265
+ cache = {}
1266
+ self._hf_processor_cache = cache
1267
+
1268
+ cache_key = self._get_processor_cache_key(processor_kwargs)
1269
+ processor = cache.get(cache_key)
1270
+ if processor is not None:
1271
+ return processor
1272
+
1273
+ processor = MossAudioProcessor(
1274
+ self.get_tokenizer(),
1275
+ **processor_kwargs,
1276
+ )
1277
+ cache[cache_key] = processor
1278
+ return processor
1279
+
1280
+ def get_feature_extractor(self, **kwargs: object) -> WhisperFeatureExtractor:
1281
+ return self.get_hf_processor(**kwargs).feature_extractor
1282
+
1283
+ def get_supported_mm_limits(self) -> Mapping[str, int | None]:
1284
+ return {"audio": None}
1285
+
1286
+ def get_data_parser(self) -> MultiModalDataParser:
1287
+ processor = self.get_hf_processor()
1288
+ return MossAudioMultiModalDataParser(
1289
+ target_sr=processor.mel_config["mel_sr"],
1290
+ target_channels=1,
1291
+ expected_hidden_size=self._get_expected_hidden_size(),
1292
+ )
1293
+
1294
+ def get_mm_max_tokens_per_item(
1295
+ self,
1296
+ seq_len: int,
1297
+ mm_counts: Mapping[str, int],
1298
+ ) -> Mapping[str, int] | None:
1299
+ if mm_counts.get("audio", 0) <= 0:
1300
+ return {}
1301
+ processor = self.get_hf_processor()
1302
+ raw_mel_len = math.ceil(
1303
+ (processor.mel_config["mel_sr"] * DEFAULT_MAX_AUDIO_SECONDS)
1304
+ / processor.mel_config["mel_hop_length"]
1305
+ )
1306
+ return {"audio": MossAudioEncoder.compute_num_audio_tokens(raw_mel_len)}
1307
+
1308
+
1309
+class MossAudioDummyInputsBuilder(BaseDummyInputsBuilder[MossAudioProcessingInfo]):
1310
+ def get_dummy_text(self, mm_counts: Mapping[str, int]) -> str:
1311
+ num_audios = mm_counts.get("audio", 0)
1312
+ return MOSS_AUDIO_PLACEHOLDER * num_audios
1313
+
1314
+ def get_dummy_mm_data(
1315
+ self,
1316
+ seq_len: int,
1317
+ mm_counts: Mapping[str, int],
1318
+ mm_options: Mapping[str, BaseDummyOptions],
1319
+ ) -> MultiModalDataDict:
1320
+ num_audios = mm_counts.get("audio", 0)
1321
+ audio_overrides = mm_options.get("audio")
1322
+ return {
1323
+ "audio": self._get_dummy_audios(
1324
+ length=16000,
1325
+ num_audios=num_audios,
1326
+ overrides=audio_overrides,
1327
+ )
1328
+ }
1329
+
1330
+
1331
+class MossAudioMultiModalProcessor(BaseMultiModalProcessor[MossAudioProcessingInfo]):
1332
+ def _call_hf_processor(
1333
+ self,
1334
+ prompt: str,
1335
+ mm_data: Mapping[str, object],
1336
+ mm_kwargs: Mapping[str, object],
1337
+ tok_kwargs: Mapping[str, object],
1338
+ ) -> BatchFeature:
1339
+ mm_data = dict(mm_data)
1340
+ audios = mm_data.pop("audios", [])
1341
+ if audios:
1342
+ mm_data["audio"] = audios
1343
+ mm_kwargs = dict(mm_kwargs)
1344
+ processor_kwargs = _filter_moss_audio_processor_config(mm_kwargs)
1345
+ tok_kwargs = {
1346
+ key: value
1347
+ for key, value in tok_kwargs.items()
1348
+ if key not in MOSS_AUDIO_PROCESSOR_CONFIG_KEYS
1349
+ }
1350
+ return self.info.ctx.call_hf_processor(
1351
+ self.info.get_hf_processor(**processor_kwargs),
1352
+ dict(text=prompt, **mm_data),
1353
+ dict(**tok_kwargs),
1354
+ )
1355
+
1356
+ def _get_mm_fields_config(
1357
+ self,
1358
+ hf_inputs: BatchFeature,
1359
+ hf_processor_mm_kwargs: Mapping[str, object],
1360
+ ) -> Mapping[str, MultiModalFieldConfig]:
1361
+ return _moss_audio_field_config(hf_inputs)
1362
+
1363
+ def _get_prompt_updates(
1364
+ self,
1365
+ mm_items: MultiModalDataItems,
1366
+ hf_processor_mm_kwargs: Mapping[str, object],
1367
+ out_mm_kwargs: MultiModalKwargsItems,
1368
+ ) -> Sequence[PromptUpdate]:
1369
+ processor = self.info.get_hf_processor(**hf_processor_mm_kwargs)
1370
+ out_mm_data = out_mm_kwargs.get_data()
1371
+ audio_data_seqlens = out_mm_data.get("audio_data_seqlens")
1372
+ if audio_data_seqlens is None:
1373
+ audio_token_lens: list[int] = []
1374
+ else:
1375
+ if isinstance(audio_data_seqlens, torch.Tensor):
1376
+ lens = audio_data_seqlens.reshape(-1).tolist()
1377
+ else:
1378
+ lens = list(audio_data_seqlens)
1379
+ audio_token_lens = [
1380
+ MossAudioEncoder.compute_num_audio_tokens(int(length))
1381
+ for length in lens
1382
+ ]
1383
+
1384
+ def get_replacement(
1385
+ item_idx: int,
1386
+ suffix_token_ids: list[int] | None = None,
1387
+ ) -> PromptUpdateDetails[list[int]]:
1388
+ num_tokens = audio_token_lens[item_idx]
1389
+ if num_tokens == 0:
1390
+ raise ValueError("The audio is too short to be represented.")
1391
+ audio_token_ids = processor.build_audio_placeholder_ids(num_tokens)
1392
+ suffix_token_ids = suffix_token_ids or []
1393
+ is_embed = torch.tensor(
1394
+ [token_id == processor.audio_token_id for token_id in audio_token_ids],
1395
+ dtype=torch.bool,
1396
+ )
1397
+ return PromptUpdateDetails(
1398
+ full=[
1399
+ processor.audio_start_id,
1400
+ *audio_token_ids,
1401
+ processor.audio_end_id,
1402
+ *suffix_token_ids,
1403
+ ],
1404
+ is_embed=lambda _tokenizer, _seq: torch.cat(
1405
+ [
1406
+ torch.tensor([False]),
1407
+ is_embed,
1408
+ torch.tensor([False]),
1409
+ torch.zeros(len(suffix_token_ids), dtype=torch.bool),
1410
+ ]
1411
+ ),
1412
+ )
1413
+
1414
+ prompt_update_specs = [
1415
+ (
1416
+ [
1417
+ processor.audio_start_id,
1418
+ processor.audio_token_id,
1419
+ processor.audio_end_id,
1420
+ ],
1421
+ [],
1422
+ )
1423
+ ]
1424
+ for suffix in ("", "\n"):
1425
+ tokenizer_target = processor.tokenizer.encode(
1426
+ MOSS_AUDIO_PLACEHOLDER + suffix,
1427
+ add_special_tokens=False,
1428
+ )
1429
+ suffix_token_ids = processor.tokenizer.encode(
1430
+ suffix,
1431
+ add_special_tokens=False,
1432
+ )
1433
+ if any(target == tokenizer_target for target, _ in prompt_update_specs):
1434
+ continue
1435
+ prompt_update_specs.append((tokenizer_target, suffix_token_ids))
1436
+
1437
+ return [
1438
+ PromptReplacement(
1439
+ modality="audio",
1440
+ target=target,
1441
+ replacement=(
1442
+ lambda item_idx, suffix_token_ids=suffix_token_ids: get_replacement(
1443
+ item_idx,
1444
+ suffix_token_ids,
1445
+ )
1446
+ ),
1447
+ )
1448
+ for target, suffix_token_ids in prompt_update_specs
1449
+ ]
1450
+
1451
+
1452
+@MULTIMODAL_REGISTRY.register_processor(
1453
+ MossAudioMultiModalProcessor,
1454
+ info=MossAudioProcessingInfo,
1455
+ dummy_inputs=MossAudioDummyInputsBuilder,
1456
+)
1457
+class MossAudioModel(nn.Module, SupportsMultiModal, SupportsPP, SupportsLoRA):
1458
+ packed_modules_mapping = {
1459
+ "qkv_proj": [
1460
+ "q_proj",
1461
+ "k_proj",
1462
+ "v_proj",
1463
+ ],
1464
+ "gate_up_proj": [
1465
+ "gate_proj",
1466
+ "up_proj",
1467
+ ],
1468
+ }
1469
+
1470
+ embedding_modules = {
1471
+ "embed_tokens": "input_embeddings",
1472
+ "lm_head": "output_embeddings",
1473
+ }
1474
+
1475
+ hf_to_vllm_mapper = WeightsMapper(
1476
+ orig_to_new_prefix={
1477
+ "lm_head.": "language_model.lm_head.",
1478
+ "language_model.embed_tokens.": "language_model.model.embed_tokens.",
1479
+ "language_model.layers.": "language_model.model.layers.",
1480
+ "language_model.norm.": "language_model.model.norm.",
1481
+ }
1482
+ )
1483
+
1484
+ def get_mm_mapping(self) -> MultiModelKeys:
1485
+ return MultiModelKeys.from_string_field(
1486
+ language_model="language_model.",
1487
+ )
1488
+
1489
+ @classmethod
1490
+ def get_placeholder_str(cls, modality: str, i: int) -> str | None:
1491
+ if modality.startswith("audio"):
1492
+ return MOSS_AUDIO_PLACEHOLDER
1493
+ raise ValueError("Only audio modality is supported")
1494
+
1495
+ def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None:
1496
+ super().__init__()
1497
+ self.vllm_config = vllm_config
1498
+ config = vllm_config.model_config.hf_config
1499
+ if not isinstance(config, MossAudioConfig):
1500
+ config = MossAudioConfig(
1501
+ audio_config=getattr(config, "audio_config", None),
1502
+ language_config=getattr(config, "language_config", None),
1503
+ adapter_hidden_size=getattr(config, "adapter_hidden_size", 8192),
1504
+ ignore_index=getattr(config, "ignore_index", -100),
1505
+ deepstack_num_inject_layers=getattr(
1506
+ config, "deepstack_num_inject_layers", None
1507
+ ),
1508
+ )
1509
+ self.config = config
1510
+ self.quant_config = vllm_config.quant_config
1511
+ self.multimodal_config = vllm_config.model_config.multimodal_config
1512
+
1513
+ parallel_config = vllm_config.parallel_config
1514
+ tp_size = parallel_config.tensor_parallel_size
1515
+ if self.config.adapter_hidden_size % tp_size != 0:
1516
+ raise ValueError(
1517
+ "MOSS-Audio adapter_hidden_size must be divisible by tensor "
1518
+ f"parallel size. Got adapter_hidden_size="
1519
+ f"{self.config.adapter_hidden_size} and tensor_parallel_size="
1520
+ f"{tp_size}."
1521
+ )
1522
+
1523
+ audio_config = MossAudioEncoderConfig.from_config(self.config.audio_config)
1524
+ if audio_config.encoder_attention_heads % tp_size != 0:
1525
+ raise ValueError(
1526
+ "MOSS-Audio encoder_attention_heads must be divisible by "
1527
+ "tensor parallel size. Got encoder_attention_heads="
1528
+ f"{audio_config.encoder_attention_heads} and "
1529
+ f"tensor_parallel_size={tp_size}."
1530
+ )
1531
+ language_config = self.config.language_config
1532
+ self.audio_token_id = MOSS_AUDIO_TOKEN_ID
1533
+ self.deepstack_input_embeds: IntermediateTensors | None = None
1534
+
1535
+ with self._mark_tower_model(vllm_config, "audio"):
1536
+ self.audio_encoder = MossAudioEncoder(
1537
+ audio_config,
1538
+ quant_config=self.quant_config,
1539
+ prefix=maybe_prefix(prefix, "audio_encoder"),
1540
+ )
1541
+ self.audio_adapter = GatedMLP(
1542
+ input_size=audio_config.output_dim,
1543
+ hidden_size=self.config.adapter_hidden_size,
1544
+ output_size=language_config.hidden_size,
1545
+ quant_config=self.quant_config,
1546
+ prefix=maybe_prefix(prefix, "audio_adapter"),
1547
+ )
1548
+
1549
+ deepstack_k = len(audio_config.deepstack_encoder_layer_indexes or [])
1550
+ if self.config.deepstack_num_inject_layers is not None:
1551
+ deepstack_k = min(
1552
+ deepstack_k,
1553
+ int(self.config.deepstack_num_inject_layers),
1554
+ )
1555
+ self.deepstack_audio_merger_list = nn.ModuleList(
1556
+ [
1557
+ GatedMLP(
1558
+ input_size=audio_config.output_dim,
1559
+ hidden_size=self.config.adapter_hidden_size,
1560
+ output_size=language_config.hidden_size,
1561
+ quant_config=self.quant_config,
1562
+ prefix=maybe_prefix(
1563
+ prefix,
1564
+ f"deepstack_audio_merger_list.{layer_idx}",
1565
+ ),
1566
+ )
1567
+ for layer_idx in range(deepstack_k)
1568
+ ]
1569
+ )
1570
+
1571
+ with self._mark_language_model(vllm_config):
1572
+ self.language_model = MossQwen3ForCausalLM(
1573
+ vllm_config=vllm_config.with_hf_config(
1574
+ language_config, architectures=["Qwen3ForCausalLM"]
1575
+ ),
1576
+ prefix=maybe_prefix(prefix, "language_model"),
1577
+ )
1578
+ self.language_model.deepstack_inject_layer_indices = range(deepstack_k)
1579
+ self.language_model.model.deepstack_inject_layer_indices = range(
1580
+ deepstack_k
1581
+ )
1582
+
1583
+ self.make_empty_intermediate_tensors = (
1584
+ self.language_model.make_empty_intermediate_tensors
1585
+ )
1586
+
1587
+ @staticmethod
1588
+ def _validate_audio_batch_size(
1589
+ audio_batch_size: int, audio_data_seqlens: torch.Tensor
1590
+ ) -> None:
1591
+ if audio_batch_size != audio_data_seqlens.numel():
1592
+ raise ValueError(
1593
+ "audio_data batch size does not match audio_data_seqlens: "
1594
+ f"{audio_batch_size} != {audio_data_seqlens.numel()}."
1595
+ )
1596
+
1597
+ @staticmethod
1598
+ def _pad_audio_data_list(
1599
+ audio_data: list[torch.Tensor],
1600
+ audio_data_seqlens: torch.Tensor,
1601
+ ) -> torch.Tensor:
1602
+ if len(audio_data) == 0:
1603
+ raise ValueError("audio_data list must not be empty.")
1604
+ MossAudioModel._validate_audio_batch_size(len(audio_data), audio_data_seqlens)
1605
+
1606
+ # pad_sequence needs every item to share the same trailing feature
1607
+ # layout, so validate the mel-major audio tensors before transposing.
1608
+ first = audio_data[0]
1609
+ if not isinstance(first, torch.Tensor):
1610
+ raise TypeError("audio_data list items must be torch.Tensor.")
1611
+ if first.ndim != 2:
1612
+ raise ValueError("audio_data list items must have shape [mel_dim, time].")
1613
+
1614
+ mel_dim = first.shape[0]
1615
+ dtype = first.dtype
1616
+ device = first.device
1617
+ for item in audio_data[1:]:
1618
+ if not isinstance(item, torch.Tensor):
1619
+ raise TypeError("audio_data list items must be torch.Tensor.")
1620
+ if item.ndim != 2:
1621
+ raise ValueError(
1622
+ "audio_data list items must have shape [mel_dim, time]."
1623
+ )
1624
+ if item.shape[0] != mel_dim:
1625
+ raise ValueError("audio_data list items must have the same mel_dim.")
1626
+ if item.dtype != dtype:
1627
+ raise TypeError("audio_data list items must have the same dtype.")
1628
+ if item.device != device:
1629
+ raise ValueError("audio_data list items must be on the same device.")
1630
+
1631
+ # Each item arrives as [mel_dim, time]. pad_sequence pads along dim 1
1632
+ # after converting to [time, mel_dim], then we restore [batch, mel, time].
1633
+ time_major = [item.transpose(0, 1) for item in audio_data]
1634
+ padded = torch.nn.utils.rnn.pad_sequence(time_major, batch_first=True)
1635
+ return padded.transpose(1, 2).contiguous()
1636
+
1637
+ def _parse_and_validate_audio_input(
1638
+ self, **kwargs: object
1639
+ ) -> MossAudioAudioInputs | None:
1640
+ """Normalize and validate model-side audio kwargs.
1641
+
1642
+ If audio_data is provided, this checks that audio_data_seqlens is also
1643
+ present, flattens sequence lengths to a long tensor, pads list inputs
1644
+ to [batch, mel_dim, time], validates batch-size/sequence-length
1645
+ agreement, and rejects empty, non-positive, or downsampled-zero audio
1646
+ lengths.
1647
+ """
1648
+ audio_data = kwargs.pop("audio_data", None)
1649
+ audio_data_seqlens = kwargs.pop("audio_data_seqlens", None)
1650
+ if audio_data is None:
1651
+ return None
1652
+ if audio_data_seqlens is None:
1653
+ raise ValueError(
1654
+ "audio_data_seqlens is required when audio_data is provided."
1655
+ )
1656
+ if not isinstance(audio_data_seqlens, torch.Tensor):
1657
+ audio_data_seqlens = torch.tensor(audio_data_seqlens, dtype=torch.long)
1658
+ audio_data_seqlens = audio_data_seqlens.to(dtype=torch.long).reshape(-1)
1659
+
1660
+ if isinstance(audio_data, list):
1661
+ audio_data = self._pad_audio_data_list(audio_data, audio_data_seqlens)
1662
+ elif isinstance(audio_data, torch.Tensor):
1663
+ if audio_data.ndim == 3:
1664
+ self._validate_audio_batch_size(audio_data.shape[0], audio_data_seqlens)
1665
+ else:
1666
+ raise TypeError("audio_data must be a torch.Tensor or list[torch.Tensor].")
1667
+
1668
+ audio_token_lens = MossAudioEncoder._compute_downsampled_length(
1669
+ audio_data_seqlens
1670
+ )
1671
+ if (
1672
+ audio_data_seqlens.numel() == 0
1673
+ or torch.any(audio_data_seqlens <= 0).item()
1674
+ or torch.any(audio_token_lens <= 0).item()
1675
+ ):
1676
+ raise ValueError("The audio is too short to be represented.")
1677
+ return MossAudioAudioInputs(
1678
+ audio_data=audio_data,
1679
+ audio_data_seqlens=audio_data_seqlens,
1680
+ )
1681
+
1682
+ def _process_audio_input(
1683
+ self,
1684
+ audio_input: MossAudioAudioInputs,
1685
+ ) -> tuple[torch.Tensor, ...]:
1686
+ """Run the audio encoder and return one embedding tensor per audio.
1687
+
1688
+ Example:
1689
+ audio_data=[2, 128, 1200], audio_data_seqlens=[800, 1200]
1690
+ -> returns (audio0_embeds, audio1_embeds), split by token length
1691
+ -> DeepStack packs each item as [main, layer0, ...] on dim -1
1692
+ """
1693
+ audio_data = audio_input["audio_data"]
1694
+ audio_data_seqlens = audio_input["audio_data_seqlens"]
1695
+ last_hidden_state, deepstack = self.audio_encoder(
1696
+ audio_data.to(self.audio_encoder.dtype),
1697
+ feature_lens=audio_data_seqlens,
1698
+ output_deepstack_hidden_states=len(self.deepstack_audio_merger_list) > 0,
1699
+ )
1700
+ audio_embeds = self.audio_adapter(last_hidden_state)
1701
+ audio_lengths = MossAudioEncoder._compute_downsampled_length(
1702
+ audio_data_seqlens.to(device=audio_embeds.device, dtype=torch.long)
1703
+ ).tolist()
1704
+ main_embeddings = tuple(audio_embeds.squeeze(0).split(audio_lengths, dim=0))
1705
+
1706
+ deepstack_embeddings: list[tuple[torch.Tensor, ...]] = []
1707
+ if deepstack is not None:
1708
+ if len(deepstack) < len(self.deepstack_audio_merger_list):
1709
+ raise RuntimeError(
1710
+ "DeepStack output count does not match configured audio "
1711
+ "merger count."
1712
+ )
1713
+ for idx, hidden_states in enumerate(
1714
+ deepstack[: len(self.deepstack_audio_merger_list)]
1715
+ ):
1716
+ ds_embeds = self.deepstack_audio_merger_list[idx](hidden_states)
1717
+ deepstack_embeddings.append(
1718
+ tuple(ds_embeds.squeeze(0).split(audio_lengths, dim=0))
1719
+ )
1720
+
1721
+ if not deepstack_embeddings:
1722
+ return main_embeddings
1723
+
1724
+ return tuple(
1725
+ torch.cat(
1726
+ [
1727
+ main_embedding,
1728
+ *(
1729
+ layer_embeddings[item_idx]
1730
+ for layer_embeddings in deepstack_embeddings
1731
+ ),
1732
+ ],
1733
+ dim=-1,
1734
+ )
1735
+ for item_idx, main_embedding in enumerate(main_embeddings)
1736
+ )
1737
+
1738
+ def embed_multimodal(self, **kwargs: object) -> MultiModalEmbeddings:
1739
+ audio_input = self._parse_and_validate_audio_input(**kwargs)
1740
+ if audio_input is None:
1741
+ return ()
1742
+ return self._process_audio_input(audio_input)
1743
+
1744
+ def _split_multimodal_embeddings(
1745
+ self,
1746
+ multimodal_embeddings: MultiModalEmbeddings,
1747
+ hidden_size: int,
1748
+ ) -> tuple[tuple[torch.Tensor, ...], tuple[tuple[torch.Tensor, ...], ...]]:
1749
+ """Unpack audio embeddings before merging them into token embeddings.
1750
+
1751
+ embed_input_ids calls this on the output of embed_multimodal. Plain
1752
+ audio embeddings already have width hidden_size and are returned as the
1753
+ main embeddings for _merge_multimodal_embeddings. When DeepStack is
1754
+ enabled, _process_audio_input packs each audio item as
1755
+ [main, layer0, layer1, ...] along the last dimension so the standard
1756
+ multimodal path can carry a single embedding object. This method splits
1757
+ that packed layout back into main embeddings plus per-layer DeepStack
1758
+ embeddings, which _cache_deepstack_input_embeds scatters and forward
1759
+ passes into MossQwen3Model for layer injection.
1760
+ """
1761
+ if isinstance(multimodal_embeddings, torch.Tensor):
1762
+ embeddings = tuple(multimodal_embeddings.unbind(0))
1763
+ else:
1764
+ embeddings = tuple(multimodal_embeddings)
1765
+
1766
+ if len(embeddings) == 0:
1767
+ return (), ()
1768
+
1769
+ deepstack_count = len(self.deepstack_audio_merger_list)
1770
+ if all(embedding.shape[-1] == hidden_size for embedding in embeddings):
1771
+ return embeddings, ()
1772
+
1773
+ packed_hidden_size = hidden_size * (deepstack_count + 1)
1774
+ if deepstack_count == 0 or any(
1775
+ embedding.shape[-1] != packed_hidden_size for embedding in embeddings
1776
+ ):
1777
+ got = [int(embedding.shape[-1]) for embedding in embeddings]
1778
+ raise ValueError(
1779
+ "MOSS-Audio multimodal embedding width mismatch: expected "
1780
+ f"{hidden_size} or {packed_hidden_size}, got {got}."
1781
+ )
1782
+
1783
+ split_by_item = [
1784
+ torch.split(embedding, hidden_size, dim=-1) for embedding in embeddings
1785
+ ]
1786
+ main_embeddings = tuple(parts[0] for parts in split_by_item)
1787
+ deepstack_embeddings = tuple(
1788
+ tuple(parts[layer_idx + 1] for parts in split_by_item)
1789
+ for layer_idx in range(deepstack_count)
1790
+ )
1791
+ return main_embeddings, deepstack_embeddings
1792
+
1793
+ def _cache_deepstack_input_embeds(
1794
+ self,
1795
+ inputs_embeds: torch.Tensor,
1796
+ deepstack_embeddings: tuple[tuple[torch.Tensor, ...], ...],
1797
+ is_multimodal: torch.Tensor,
1798
+ ) -> None:
1799
+ if len(deepstack_embeddings) == 0:
1800
+ self.deepstack_input_embeds = None
1801
+ return
1802
+ flat_by_layer = [
1803
+ torch.cat(layer_embeds, dim=0).to(
1804
+ device=inputs_embeds.device, dtype=inputs_embeds.dtype
1805
+ )
1806
+ for layer_embeds in deepstack_embeddings
1807
+ ]
1808
+ num_mm_tokens = int(is_multimodal.sum().item())
1809
+ if any(layer.shape[0] != num_mm_tokens for layer in flat_by_layer):
1810
+ got = [int(layer.shape[0]) for layer in flat_by_layer]
1811
+ raise ValueError(
1812
+ "DeepStack audio token count mismatch: "
1813
+ f"expected {num_mm_tokens}, got {got}."
1814
+ )
1815
+ data = {}
1816
+ for layer_idx, layer_embeds in enumerate(flat_by_layer):
1817
+ scattered = inputs_embeds.new_zeros(inputs_embeds.shape)
1818
+ scattered[is_multimodal] = layer_embeds
1819
+ data[f"deepstack_input_embeds_{layer_idx}"] = scattered
1820
+ self.deepstack_input_embeds = IntermediateTensors(data)
1821
+
1822
+ def embed_input_ids(
1823
+ self,
1824
+ input_ids: torch.Tensor,
1825
+ multimodal_embeddings: MultiModalEmbeddings | None = None,
1826
+ *,
1827
+ is_multimodal: torch.Tensor | None = None,
1828
+ ) -> torch.Tensor:
1829
+ inputs_embeds = self._embed_text_input_ids(
1830
+ input_ids,
1831
+ self.language_model.embed_input_ids,
1832
+ is_multimodal=is_multimodal,
1833
+ )
1834
+
1835
+ self.deepstack_input_embeds = None
1836
+ if multimodal_embeddings is None or len(multimodal_embeddings) == 0:
1837
+ return inputs_embeds
1838
+ is_multimodal = _require_is_multimodal(is_multimodal)
1839
+ multimodal_embeddings, deepstack_embeddings = self._split_multimodal_embeddings(
1840
+ multimodal_embeddings,
1841
+ hidden_size=int(inputs_embeds.shape[-1]),
1842
+ )
1843
+
1844
+ inputs_embeds = _merge_multimodal_embeddings(
1845
+ inputs_embeds=inputs_embeds,
1846
+ multimodal_embeddings=multimodal_embeddings,
1847
+ is_multimodal=is_multimodal,
1848
+ )
1849
+ self._cache_deepstack_input_embeds(
1850
+ inputs_embeds,
1851
+ deepstack_embeddings,
1852
+ is_multimodal,
1853
+ )
1854
+ return inputs_embeds
1855
+
1856
+ def forward(
1857
+ self,
1858
+ input_ids: torch.Tensor | None,
1859
+ positions: torch.Tensor,
1860
+ intermediate_tensors: IntermediateTensors | None = None,
1861
+ inputs_embeds: torch.Tensor | None = None,
1862
+ **kwargs: object,
1863
+ ) -> torch.Tensor | IntermediateTensors:
1864
+ if intermediate_tensors is None:
1865
+ deepstack_input_embeds = self.deepstack_input_embeds
1866
+ else:
1867
+ # Non-first PP ranks consume hidden states from intermediate_tensors.
1868
+ # The executor may still pass dummy inputs_embeds during profiling.
1869
+ inputs_embeds = None
1870
+ deepstack_input_embeds = intermediate_tensors
1871
+ hidden_states = self.language_model(
1872
+ input_ids,
1873
+ positions,
1874
+ intermediate_tensors=intermediate_tensors,
1875
+ inputs_embeds=inputs_embeds,
1876
+ deepstack_input_embeds=deepstack_input_embeds,
1877
+ )
1878
+ self.deepstack_input_embeds = None
1879
+ return hidden_states
1880
+
1881
+ def compute_logits(
1882
+ self,
1883
+ hidden_states: torch.Tensor,
1884
+ ) -> torch.Tensor | None:
1885
+ return self.language_model.compute_logits(hidden_states)
1886
+
1887
+ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
1888
+ loader = AutoWeightsLoader(
1889
+ self,
1890
+ skip_prefixes=["audio_encoder.embed_positions"],
1891
+ )
1892
+ return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)
vllm/model_executor/models/registry.py
+1
@@ -500,6 +500,7 @@ _MULTIMODAL_MODELS = {
500
"MolmoForCausalLM": ("molmo", "MolmoForCausalLM"),
501
"Molmo2ForConditionalGeneration": ("molmo2", "Molmo2ForConditionalGeneration"),
502
"Moondream3ForCausalLM": ("moondream3", "Moondream3ForCausalLM"),
503
+ "MossAudioModel": ("moss_audio", "MossAudioModel"),
504
"HfMoondream": ("moondream3", "Moondream3ForCausalLM"),
505
"NemotronH_Nano_VL_V2": ("nano_nemotron_vl", "NemotronH_Nano_VL_V2"),
506
"NemotronH_Nano_Omni_Reasoning_V3": ("nano_nemotron_vl", "NemotronH_Nano_VL_V2"),
vllm/transformers_utils/model_arch_config_convertor.py
+45
@@ -580,6 +580,50 @@ class Gemma4ModelArchConfigConvertor(ModelArchConfigConvertorBase):
580
return max(head_dim, global_head_dim) or super().get_head_size()
581
582
583
+class MossAudioModelArchConfigConvertor(ModelArchConfigConvertorBase):
584
+ def _language_config(self) -> PretrainedConfig:
585
+ return self.hf_config.language_config
586
+
587
+ def get_num_hidden_layers(self) -> int:
588
+ return getattr(self._language_config(), "num_hidden_layers", 0)
589
+
590
+ def get_total_num_attention_heads(self) -> int:
591
+ return getattr(self._language_config(), "num_attention_heads", 0)
592
+
593
+ def get_vocab_size(self) -> int:
594
+ return getattr(self._language_config(), "vocab_size", 0)
595
+
596
+ def get_hidden_size(self) -> int:
597
+ return getattr(self._language_config(), "hidden_size", 0)
598
+
599
+ def get_head_size(self) -> int:
600
+ head_dim = getattr(self._language_config(), "head_dim", None)
601
+ if head_dim is not None:
602
+ return head_dim
603
+ total_num_attention_heads = self.get_total_num_attention_heads()
604
+ if total_num_attention_heads == 0:
605
+ return 0
606
+ return self.get_hidden_size() // total_num_attention_heads
607
+
608
+ def get_total_num_kv_heads(self) -> int:
609
+ return getattr(
610
+ self._language_config(),
611
+ "num_key_value_heads",
612
+ self.get_total_num_attention_heads(),
613
+ )
614
+
615
+ def derive_max_model_len_and_key(self) -> tuple[float, str | None]:
616
+ language_config = self._language_config()
617
+ max_position_embeddings = getattr(
618
+ language_config,
619
+ "max_position_embeddings",
620
+ None,
621
+ )
622
+ if max_position_embeddings is None:
623
+ return super().derive_max_model_len_and_key()
624
+ return max_position_embeddings, "language_config.max_position_embeddings"
625
+
626
+
627
# hf_config.model_type -> convertor class
628
MODEL_ARCH_CONFIG_CONVERTORS = {
629
"cohere_asr": CohereAsrModelArchConfigConvertor,
@@ -604,6 +648,7 @@ MODEL_ARCH_CONFIG_CONVERTORS = {
648
"mimo_v2_flash": MimoV2ModelArchConfigConvertor,
649
"mimo_v2_mtp": MimoV2MTPModelArchConfigConvertor,
650
"mimo_v2_omni_mtp": MimoV2MTPModelArchConfigConvertor,
651
+ "moss_audio": MossAudioModelArchConfigConvertor,
652
"mpt": MPTModelArchConfigConvertor,
653
"nemotron-nas": NemotronNasModelArchConfigConvertor,
654
"pangu_ultra_moe_mtp": PanguUltraMoeMTPModelArchConfigConvertor,