@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 974138751b
Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git

Synced Nov 10, 14:00 UTC

[Refactor] GLM-ASR Modeling (#31779)

Signed-off-by: JaredforReal <w13431838023@gmail.com> Co-authored-by: Isotr0py <mozf@mail2.sysu.edu.cn>

Jared Wen committed Jan 7, 2026 at 21:08 UTC 974138751bdb8e41743d0beac0ac4c5d89308c2f
2 files changed +672 -41
vllm/model_executor/models/glmasr.py
+644 -36
@@ -8,18 +8,22 @@ import numpy as np
8 import torch
9 import torch.nn as nn
10 from transformers import BatchFeature
11 -from transformers.models.glmasr import GlmAsrConfig, GlmAsrEncoder, GlmAsrProcessor
11 +from transformers.models.glmasr import GlmAsrConfig, GlmAsrProcessor
12 from transformers.models.whisper import WhisperFeatureExtractor
13
14 +from vllm.attention.layers.mm_encoder_attention import MMEncoderAttention
15 from vllm.config import ModelConfig, SpeechToTextConfig, VllmConfig
16 from vllm.config.multimodal import BaseDummyOptions
17 +from vllm.distributed.parallel_state import get_tensor_model_parallel_world_size
18 from vllm.inputs.data import PromptType
19 from vllm.model_executor.layers.activation import get_act_fn
20 from vllm.model_executor.layers.linear import (
21 ColumnParallelLinear,
22 + QKVParallelLinear,
23 RowParallelLinear,
24 )
25 from vllm.model_executor.layers.quantization import QuantizationConfig
26 +from vllm.model_executor.layers.rotary_embedding.common import ApplyRotaryEmb
27 from vllm.model_executor.models.module_mapping import MultiModelKeys
28 from vllm.multimodal import MULTIMODAL_REGISTRY
29 from vllm.multimodal.inputs import (
@@ -35,6 +39,8 @@ from vllm.multimodal.parse import (
39 MultiModalDataParser,
40 )
41 from vllm.multimodal.processing import (
42 + BaseMultiModalProcessor,
43 + BaseProcessingInfo,
44 PromptReplacement,
45 PromptUpdate,
46 PromptUpdateDetails,
@@ -45,21 +51,12 @@ from vllm.tokenizers import cached_tokenizer_from_config
51 from vllm.transformers_utils.processor import cached_processor_from_config
52 from vllm.utils.tensor_schema import TensorSchema, TensorShape
53
48 -from .audioflamingo3 import (
49 - AudioFlamingo3MultiModalDataParser,
50 - AudioFlamingo3MultiModalProcessor,
51 - AudioFlamingo3ProcessingInfo,
52 -)
53 -from .audioflamingo3 import (
54 - _audioflamingo3_field_config as _glmasr_field_config,
55 -)
54 from .glmasr_utils import (
55 DEFAULT_CONV_PARAMS,
56 DEFAULT_MAX_AUDIO_LEN_S,
57 DEFAULT_MERGE_FACTOR,
58 _flatten_audio_features_by_length,
59 _get_audio_output_lengths_for_tower,
62 - _get_num_features_for_item,
60 _group_audio_embeddings,
61 _normalize_chunk_counts,
62 )
@@ -74,6 +71,460 @@ from .utils import AutoWeightsLoader, init_vllm_registered_model, maybe_prefix
71 from .whisper import ISO639_1_SUPPORTED_LANGS
72
73
74 +class GlmAsrEncoderRotaryEmbedding(nn.Module):
75 + """
76 + Rotary Position Embedding for GLM-ASR encoder.
77 +
78 + Computes rotary position embeddings on-demand for efficiency.
79 + Only caches inv_freq as a buffer; cos/sin are computed during forward
80 + to avoid wasted computation during initialization and ensure correct
81 + device placement.
82 + """
83 +
84 + def __init__(self, config) -> None:
85 + super().__init__()
86 +
87 + # Compute inverse frequencies following transformers implementation
88 + head_dim = getattr(
89 + config, "head_dim", config.hidden_size // config.num_attention_heads
90 + )
91 +
92 + # Handle rope_parameters if present (for compatibility with transformers config)
93 + if hasattr(config, "rope_parameters") and config.rope_parameters:
94 + base = config.rope_parameters.get("rope_theta", 10000.0)
95 + partial_rotary_factor = config.rope_parameters.get(
96 + "partial_rotary_factor", 1.0
97 + )
98 + dim = int(head_dim * partial_rotary_factor)
99 + self.attention_scaling = config.rope_parameters.get(
100 + "attention_scaling", 1.0
101 + )
102 + else:
103 + base = getattr(config, "rope_theta", 10000.0)
104 + dim = head_dim
105 + self.attention_scaling = 1.0
106 +
107 + self.dim = dim
108 + self.head_dim = head_dim
109 +
110 + # Only cache inv_freq; cos/sin computed on-demand in correct device
111 + inv_freq = 1.0 / (base ** (torch.arange(0, dim, 2, dtype=torch.float) / dim))
112 + self.register_buffer("inv_freq", inv_freq, persistent=False)
113 +
114 + def forward(self, seq_len: int) -> torch.Tensor:
115 + """
116 + Compute rotary position frequencies for given sequence length.
117 +
118 + Args:
119 + seq_len: The sequence length to compute embeddings for.
120 +
121 + Returns:
122 + Frequency tensor with shape [seq_len, dim/2]. Use .cos() and
123 + .sin() to get the rotary embedding components.
124 + """
125 + # Compute on the same device as inv_freq (automatically correct after .to())
126 + seq = torch.arange(
127 + seq_len, device=self.inv_freq.device, dtype=self.inv_freq.dtype
128 + )
129 + freqs = torch.outer(seq, self.inv_freq)
130 + return freqs * self.attention_scaling
131 +
132 +
133 +class GlmAsrEncoderAttention(nn.Module):
134 + """
135 + Optimized Multi-headed Grouped Query Attention for GLM-ASR encoder.
136 +
137 + Uses vLLM's QKVParallelLinear for fused projections, ApplyRotaryEmb for
138 + rotary position embeddings, and MMEncoderAttention for hardware-optimized
139 + attention computation with automatic backend selection.
140 + """
141 +
142 + def __init__(
143 + self,
144 + config,
145 + quant_config: QuantizationConfig | None = None,
146 + prefix: str = "",
147 + ):
148 + super().__init__()
149 + self.config = config
150 + self.hidden_size = config.hidden_size
151 + self.num_heads = config.num_attention_heads
152 + self.num_kv_heads = getattr(
153 + config, "num_key_value_heads", config.num_attention_heads
154 + )
155 + self.head_dim = self.hidden_size // self.num_heads
156 +
157 + self.tp_size = get_tensor_model_parallel_world_size()
158 + self.num_heads_per_rank = self.num_heads // self.tp_size
159 + self.num_kv_heads_per_rank = max(1, self.num_kv_heads // self.tp_size)
160 +
161 + # Use QKVParallelLinear for fused QKV projection
162 + # Note: GLM-ASR uses bias on Q and V, but not K
163 + # For simplicity with QKVParallelLinear, we use bias=True for all
164 + self.qkv_proj = QKVParallelLinear(
165 + self.hidden_size,
166 + self.head_dim,
167 + self.num_heads,
168 + self.num_kv_heads,
169 + bias=True,
170 + quant_config=quant_config,
171 + prefix=f"{prefix}.qkv_proj",
172 + )
173 +
174 + self.o_proj = RowParallelLinear(
175 + self.hidden_size,
176 + self.hidden_size,
177 + bias=True,
178 + quant_config=quant_config,
179 + prefix=f"{prefix}.o_proj",
180 + )
181 +
182 + # Use vLLM's ApplyRotaryEmb CustomOp
183 + # enforce_enable=True ensures the op is always enabled (important for ViT)
184 + self.apply_rotary_emb = ApplyRotaryEmb(enforce_enable=True)
185 +
186 + # Use vLLM's MMEncoderAttention for hardware-optimized attention
187 + # Automatically selects Flash Attention, SDPA, or Pallas based on device
188 + self.attn = MMEncoderAttention(
189 + num_heads=self.num_heads_per_rank,
190 + head_size=self.head_dim,
191 + num_kv_heads=self.num_kv_heads_per_rank,
192 + prefix=f"{prefix}.attn",
193 + )
194 +
195 + def forward(
196 + self,
197 + hidden_states: torch.Tensor,
198 + rotary_pos_emb_cos: torch.Tensor,
199 + rotary_pos_emb_sin: torch.Tensor,
200 + ) -> torch.Tensor:
201 + """
202 + Args:
203 + hidden_states: [batch_size, seq_len, hidden_size]
204 + rotary_pos_emb_cos: [seq_len, rotary_dim/2] - cosine of rotary embeddings
205 + rotary_pos_emb_sin: [seq_len, rotary_dim/2] - sine of rotary embeddings
206 +
207 + Returns:
208 + [batch_size, seq_len, hidden_size]
209 + """
210 + batch_size, seq_len, _ = hidden_states.shape
211 +
212 + # QKV projection - fused for efficiency
213 + qkv, _ = self.qkv_proj(hidden_states)
214 +
215 + # Split into q, k, v
216 + q_size = self.num_heads_per_rank * self.head_dim
217 + kv_size = self.num_kv_heads_per_rank * self.head_dim
218 + q, k, v = qkv.split([q_size, kv_size, kv_size], dim=-1)
219 +
220 + # Reshape to [batch, seq, num_heads, head_dim] for ApplyRotaryEmb
221 + q = q.view(batch_size, seq_len, self.num_heads_per_rank, self.head_dim)
222 + k = k.view(batch_size, seq_len, self.num_kv_heads_per_rank, self.head_dim)
223 + v = v.view(batch_size, seq_len, self.num_kv_heads_per_rank, self.head_dim)
224 +
225 + # Apply rotary position embeddings using vLLM's ApplyRotaryEmb
226 + # ApplyRotaryEmb expects x: [batch, seq, heads, head_dim]
227 + # cos/sin: [seq_len, rotary_dim/2]
228 + q = self.apply_rotary_emb(q, rotary_pos_emb_cos, rotary_pos_emb_sin)
229 + k = self.apply_rotary_emb(k, rotary_pos_emb_cos, rotary_pos_emb_sin)
230 +
231 + # MMEncoderAttention expects [batch, seq, num_heads, head_dim]
232 + # It handles GQA internally via repeat_interleave
233 + attn_output = self.attn(q, k, v)
234 +
235 + # Reshape back to [batch, seq, hidden_size]
236 + attn_output = attn_output.view(batch_size, seq_len, -1)
237 +
238 + # Output projection
239 + output, _ = self.o_proj(attn_output)
240 + return output
241 +
242 +
243 +class GlmAsrEncoderMLP(nn.Module):
244 + """
245 + Optimized MLP for GLM-ASR encoder.
246 + Uses vLLM's parallel linear layers for better performance.
247 + """
248 +
249 + def __init__(
250 + self,
251 + config,
252 + quant_config: QuantizationConfig | None = None,
253 + prefix: str = "",
254 + ):
255 + super().__init__()
256 + self.config = config
257 + self.hidden_size = config.hidden_size
258 + self.intermediate_size = config.intermediate_size
259 +
260 + self.fc1 = ColumnParallelLinear(
261 + self.hidden_size,
262 + self.intermediate_size,
263 + bias=True,
264 + quant_config=quant_config,
265 + prefix=f"{prefix}.fc1",
266 + )
267 +
268 + self.act_fn = get_act_fn(config.hidden_act)
269 +
270 + self.fc2 = RowParallelLinear(
271 + self.intermediate_size,
272 + self.hidden_size,
273 + bias=True,
274 + quant_config=quant_config,
275 + prefix=f"{prefix}.fc2",
276 + )
277 +
278 + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
279 + hidden_states, _ = self.fc1(hidden_states)
280 + hidden_states = self.act_fn(hidden_states)
281 + hidden_states, _ = self.fc2(hidden_states)
282 + return hidden_states
283 +
284 +
285 +class GlmAsrEncoderLayer(nn.Module):
286 + """
287 + Optimized Transformer encoder layer for GLM-ASR.
288 + Combines attention and MLP with residual connections and layer norms.
289 + """
290 +
291 + def __init__(
292 + self,
293 + config,
294 + quant_config: QuantizationConfig | None = None,
295 + prefix: str = "",
296 + ):
297 + super().__init__()
298 + self.hidden_size = config.hidden_size
299 +
300 + self.self_attn = GlmAsrEncoderAttention(
301 + config,
302 + quant_config=quant_config,
303 + prefix=f"{prefix}.self_attn",
304 + )
305 +
306 + self.mlp = GlmAsrEncoderMLP(
307 + config,
308 + quant_config=quant_config,
309 + prefix=f"{prefix}.mlp",
310 + )
311 +
312 + layer_norm_eps = getattr(config, "layer_norm_eps", 1e-5)
313 + self.input_layernorm = nn.LayerNorm(self.hidden_size, eps=layer_norm_eps)
314 + self.post_attention_layernorm = nn.LayerNorm(
315 + self.hidden_size, eps=layer_norm_eps
316 + )
317 +
318 + def forward(
319 + self,
320 + hidden_states: torch.Tensor,
321 + rotary_pos_emb_cos: torch.Tensor,
322 + rotary_pos_emb_sin: torch.Tensor,
323 + ) -> torch.Tensor:
324 + """
325 + Args:
326 + hidden_states: [batch_size, seq_len, hidden_size]
327 + rotary_pos_emb_cos: [seq_len, rotary_dim/2] - cosine of rotary embeddings
328 + rotary_pos_emb_sin: [seq_len, rotary_dim/2] - sine of rotary embeddings
329 +
330 + Returns:
331 + [batch_size, seq_len, hidden_size]
332 + """
333 + # Self-attention with residual
334 + residual = hidden_states
335 + hidden_states = self.input_layernorm(hidden_states)
336 + hidden_states = self.self_attn(
337 + hidden_states=hidden_states,
338 + rotary_pos_emb_cos=rotary_pos_emb_cos,
339 + rotary_pos_emb_sin=rotary_pos_emb_sin,
340 + )
341 + hidden_states = residual + hidden_states
342 +
343 + # MLP with residual
344 + residual = hidden_states
345 + hidden_states = self.post_attention_layernorm(hidden_states)
346 + hidden_states = self.mlp(hidden_states)
347 + hidden_states = residual + hidden_states
348 +
349 + return hidden_states
350 +
351 +
352 +class _GlmAsrEncoderOutput:
353 + """
354 + Simple output container compatible with transformers' BaseModelOutput.
355 +
356 + This lightweight container holds the encoder output and is compatible
357 + with the transformers library's output format while being more efficient
358 + than a full dataclass.
359 +
360 + Attributes:
361 + last_hidden_state: Final layer hidden states from the encoder.
362 + Shape: [batch_size, seq_len, hidden_size]
363 + """
364 +
365 + __slots__ = ("last_hidden_state",)
366 +
367 + def __init__(self, last_hidden_state: torch.Tensor):
368 + self.last_hidden_state = last_hidden_state
369 +
370 +
371 +class GlmAsrEncoder(nn.Module):
372 + """
373 + Optimized GLM-ASR Audio Encoder with vLLM native implementation.
374 +
375 + This encoder processes audio features through convolutional layers
376 + followed by transformer layers with rotary position embeddings.
377 + Optimized for performance with:
378 + - QKVParallelLinear for fused attention projections
379 + - Tensor parallelism support via ColumnParallelLinear/RowParallelLinear
380 + - Quantization support
381 + - Flash Attention (SDPA)
382 + """
383 +
384 + # Mapping for weight loading: transformers uses separate q/k/v, we use fused qkv
385 + packed_modules_mapping = {
386 + "qkv_proj": ["q_proj", "k_proj", "v_proj"],
387 + }
388 +
389 + def __init__(
390 + self,
391 + config,
392 + quant_config: QuantizationConfig | None = None,
393 + prefix: str = "",
394 + ):
395 + super().__init__()
396 + self.config = config
397 +
398 + # Convolutional feature extraction layers
399 + self.conv1 = nn.Conv1d(
400 + config.num_mel_bins,
401 + config.hidden_size,
402 + kernel_size=3,
403 + padding=1,
404 + )
405 + self.conv2 = nn.Conv1d(
406 + config.hidden_size,
407 + config.hidden_size,
408 + kernel_size=3,
409 + stride=2,
410 + padding=1,
411 + )
412 +
413 + # Transformer encoder layers
414 + self.layers = nn.ModuleList(
415 + [
416 + GlmAsrEncoderLayer(
417 + config,
418 + quant_config=quant_config,
419 + prefix=f"{prefix}.layers.{layer_idx}",
420 + )
421 + for layer_idx in range(config.num_hidden_layers)
422 + ]
423 + )
424 +
425 + # Final layer norm
426 + layer_norm_eps = getattr(config, "layer_norm_eps", 1e-5)
427 + self.norm = nn.LayerNorm(config.hidden_size, eps=layer_norm_eps)
428 +
429 + # Rotary position embeddings
430 + self.rotary_emb = GlmAsrEncoderRotaryEmbedding(config)
431 +
432 + def _get_feat_extract_output_lengths(
433 + self, input_lengths: torch.Tensor
434 + ) -> tuple[torch.Tensor, torch.Tensor]:
435 + """
436 + Compute the output length after convolutions.
437 +
438 + Args:
439 + input_lengths: Input sequence lengths [batch_size]
440 +
441 + Returns:
442 + Tuple of (output after conv1, output after conv2)
443 + """
444 + # Conv1: kernel=3, stride=1, padding=1
445 + output_lengths_conv1 = (input_lengths + 2 * 1 - 3) // 1 + 1
446 +
447 + # Conv2: kernel=3, stride=2, padding=1
448 + output_lengths_conv2 = (output_lengths_conv1 + 2 * 1 - 3) // 2 + 1
449 +
450 + return output_lengths_conv1, output_lengths_conv2
451 +
452 + def forward(self, input_features: torch.Tensor) -> _GlmAsrEncoderOutput:
453 + """
454 + Forward pass through the encoder.
455 +
456 + Args:
457 + input_features: [batch_size, num_mel_bins, seq_len]
458 +
459 + Returns:
460 + _GlmAsrEncoderOutput: Object with .last_hidden_state attribute \
461 + containing [batch_size, seq_len', hidden_size] where seq_len' \
462 + is the sequence length after convolutions
463 + """
464 + # Apply convolutional layers with GELU activation
465 + hidden_states = torch.nn.functional.gelu(self.conv1(input_features))
466 + hidden_states = torch.nn.functional.gelu(self.conv2(hidden_states))
467 +
468 + # Transpose to [batch_size, seq_len, hidden_size]
469 + hidden_states = hidden_states.transpose(1, 2)
470 + output_seq_len = hidden_states.shape[1]
471 +
472 + # Compute rotary position embeddings on-demand
473 + rotary_pos_emb = self.rotary_emb(output_seq_len)
474 + rotary_pos_emb_cos = rotary_pos_emb.cos().to(dtype=hidden_states.dtype)
475 + rotary_pos_emb_sin = rotary_pos_emb.sin().to(dtype=hidden_states.dtype)
476 +
477 + # Apply transformer layers
478 + for encoder_layer in self.layers:
479 + hidden_states = encoder_layer(
480 + hidden_states, rotary_pos_emb_cos, rotary_pos_emb_sin
481 + )
482 +
483 + # Final layer norm
484 + hidden_states = self.norm(hidden_states)
485 +
486 + # Return in a format compatible with transformers' BaseModelOutput
487 + return _GlmAsrEncoderOutput(last_hidden_state=hidden_states)
488 +
489 + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
490 + """Custom weight loading to handle q_proj/k_proj/v_proj -> qkv_proj mapping."""
491 + from vllm.model_executor.model_loader.weight_utils import default_weight_loader
492 +
493 + stacked_params_mapping = [
494 + # (param_name, shard_name, shard_id)
495 + ("qkv_proj", "q_proj", "q"),
496 + ("qkv_proj", "k_proj", "k"),
497 + ("qkv_proj", "v_proj", "v"),
498 + ]
499 + params_dict = dict(self.named_parameters())
500 + loaded_params: set[str] = set()
501 +
502 + for name, loaded_weight in weights:
503 + for param_name, weight_name, shard_id in stacked_params_mapping:
504 + if weight_name not in name:
505 + continue
506 + name = name.replace(weight_name, param_name)
507 + # Skip loading extra bias for GPTQ models.
508 + if name.endswith(".bias") and name not in params_dict:
509 + continue
510 +
511 + param = params_dict[name]
512 + weight_loader = param.weight_loader
513 + weight_loader(param, loaded_weight, shard_id)
514 + break
515 + else:
516 + # Default weight loading for non-stacked params
517 + if name.endswith(".bias") and name not in params_dict:
518 + continue
519 + if name not in params_dict:
520 + continue
521 + param = params_dict[name]
522 + weight_loader = getattr(param, "weight_loader", default_weight_loader)
523 + weight_loader(param, loaded_weight)
524 + loaded_params.add(name)
525 + return loaded_params
526 +
527 +
528 class GlmAsrFeatureInputs(TensorSchema):
529 """
530 Dimensions:
@@ -117,6 +568,19 @@ GlmAsrInputs: TypeAlias = GlmAsrFeatureInputs | GlmAsrEmbeddingInputs
568
569
570 class GlmAsrMultiModalProjector(nn.Module):
571 + """
572 + Projects audio encoder outputs to language model hidden space.
573 +
574 + This projector uses a two-layer MLP to map audio features from the
575 + encoder's intermediate size to the language model's hidden size.
576 + Uses vLLM's parallel linear layers for tensor parallelism support.
577 +
578 + Architecture:
579 + - Linear layer: intermediate_size -> hidden_size * 2
580 + - Activation function (e.g., GELU)
581 + - Linear layer: hidden_size * 2 -> hidden_size
582 + """
583 +
584 def __init__(
585 self,
586 config: GlmAsrConfig,
@@ -145,7 +609,14 @@ class GlmAsrMultiModalProjector(nn.Module):
609 return hidden_states
610
611
148 -class GlmAsrProcessingInfo(AudioFlamingo3ProcessingInfo):
612 +class GlmAsrProcessingInfo(BaseProcessingInfo):
613 + """
614 + Processing information provider for GLM-ASR model.
615 +
616 + Provides access to model configuration, processor, and feature extractor
617 + needed for audio preprocessing and multimodal integration.
618 + """
619 +
620 def get_hf_config(self) -> GlmAsrConfig:
621 return self.ctx.get_hf_config(GlmAsrConfig)
622
@@ -153,13 +624,21 @@ class GlmAsrProcessingInfo(AudioFlamingo3ProcessingInfo):
624 return self.ctx.get_hf_processor(GlmAsrProcessor, **kwargs)
625
626 def get_feature_extractor(self, **kwargs: object) -> WhisperFeatureExtractor:
156 - # Reuse parent implementation, but add type annotation and assertion
157 - feature_extractor = super().get_feature_extractor(**kwargs)
158 - assert isinstance(feature_extractor, WhisperFeatureExtractor)
159 - return feature_extractor
627 + return self.get_hf_processor(**kwargs).feature_extractor
628 +
629 + def get_supported_mm_limits(self) -> Mapping[str, int | None]:
630 + return {"audio": None}
631
632
633 class GlmAsrDummyInputsBuilder(BaseDummyInputsBuilder[GlmAsrProcessingInfo]):
634 + """
635 + Builder for dummy inputs used in profiling and testing.
636 +
637 + Generates dummy text prompts and audio data that match the expected
638 + format for GLM-ASR model inputs. Used for memory profiling and
639 + performance benchmarking.
640 + """
641 +
642 def get_dummy_text(self, mm_counts: Mapping[str, int]) -> str:
643 num_audios = mm_counts.get("audio", 0)
644 hf_processor = self.info.get_hf_processor()
@@ -188,7 +667,51 @@ class GlmAsrDummyInputsBuilder(BaseDummyInputsBuilder[GlmAsrProcessingInfo]):
667 }
668
669
191 -class GlmAsrMultiModalDataParser(AudioFlamingo3MultiModalDataParser):
670 +def _glmasr_field_config(
671 + hf_inputs: Mapping[str, torch.Tensor],
672 +) -> dict[str, MultiModalFieldConfig]:
673 + """
674 + Configure multimodal field batching strategy for GLM-ASR.
675 +
676 + Determines how to batch audio inputs based on whether chunking is used.
677 + When chunk_counts is present, features are flattened across chunks;
678 + otherwise, they are batched normally.
679 +
680 + Args:
681 + hf_inputs: Dictionary of preprocessed inputs from HuggingFace processor.
682 +
683 + Returns:
684 + Dictionary mapping field names to MultiModalFieldConfig objects \
685 + that specify batching behavior.
686 + """
687 + chunk_counts = hf_inputs.get("chunk_counts")
688 + if chunk_counts is not None:
689 + return dict(
690 + audio_embeds=MultiModalFieldConfig.batched("audio"),
691 + input_features=MultiModalFieldConfig.flat_from_sizes(
692 + "audio", chunk_counts, dim=0
693 + ),
694 + feature_attention_mask=MultiModalFieldConfig.flat_from_sizes(
695 + "audio", chunk_counts, dim=0
696 + ),
697 + chunk_counts=MultiModalFieldConfig.batched("audio"),
698 + )
699 + return dict(
700 + audio_embeds=MultiModalFieldConfig.batched("audio"),
701 + input_features=MultiModalFieldConfig.batched("audio"),
702 + feature_attention_mask=MultiModalFieldConfig.batched("audio"),
703 + chunk_counts=MultiModalFieldConfig.batched("audio"),
704 + )
705 +
706 +
707 +class GlmAsrMultiModalDataParser(MultiModalDataParser):
708 + """
709 + Custom parser for GLM-ASR multimodal data.
710 +
711 + Extends the base parser to handle GLM-ASR specific audio data formats,
712 + including both pre-computed audio embeddings and raw audio features.
713 + """
714 +
715 def _parse_audio_data(
716 self,
717 data: dict[str, torch.Tensor] | ModalityData[Any],
@@ -203,7 +726,12 @@ class GlmAsrMultiModalDataParser(AudioFlamingo3MultiModalDataParser):
726 return super()._parse_audio_data(data)
727
728
206 -class GlmAsrMultiModalProcessor(AudioFlamingo3MultiModalProcessor):
729 +class GlmAsrMultiModalProcessor(BaseMultiModalProcessor["GlmAsrProcessingInfo"]):
730 + """
731 + GLM-ASR processor that inherits directly from BaseMultiModalProcessor
732 + for better performance and cleaner implementation.
733 + """
734 +
735 def _get_data_parser(self) -> MultiModalDataParser:
736 feature_extractor = self.info.get_feature_extractor()
737 return GlmAsrMultiModalDataParser(target_sr=feature_extractor.sampling_rate)
@@ -214,7 +742,6 @@ class GlmAsrMultiModalProcessor(AudioFlamingo3MultiModalProcessor):
742 feature_extractor: WhisperFeatureExtractor,
743 processor: GlmAsrProcessor,
744 ) -> list[int]:
217 - """Calculate chunk counts for each audio."""
745 sampling_rate = feature_extractor.sampling_rate
746 chunk_length = feature_extractor.chunk_length
747 max_audio_len = getattr(processor, "max_audio_len", DEFAULT_MAX_AUDIO_LEN_S)
@@ -248,10 +775,14 @@ class GlmAsrMultiModalProcessor(AudioFlamingo3MultiModalProcessor):
775 prompt_ids = self._apply_hf_processor_tokens_only(prompt_ids)
776 return BatchFeature(dict(input_ids=[prompt_ids]), tensor_type="pt")
777
251 - # Get processor for chunk counts calculation
252 - processor = self.info.get_hf_processor(**mm_kwargs)
778 + # Handle sampling_rate
779 + feature_extractor = self.info.get_feature_extractor(**mm_kwargs)
780 + mm_kwargs = dict(
781 + **mm_kwargs,
782 + sampling_rate=feature_extractor.sampling_rate,
783 + )
784
254 - # Call parent method (it will handle sampling_rate)
785 + # Call parent method
786 outputs = super()._call_hf_processor(
787 prompt=prompt,
788 mm_data=mm_data,
@@ -259,9 +790,24 @@ class GlmAsrMultiModalProcessor(AudioFlamingo3MultiModalProcessor):
790 tok_kwargs=tok_kwargs,
791 )
792
262 - # Postprocess: rename mask and add chunk counts.
263 - if "input_features_mask" in outputs:
264 - outputs["feature_attention_mask"] = outputs.pop("input_features_mask")
793 + # Postprocess: rename mask and add chunk counts
794 + # Handle different key names from different transformers versions
795 + if "input_feature_mask" in outputs:
796 + outputs["feature_attention_mask"] = outputs.pop("input_feature_mask")
797 + elif "feature_attention_mask" not in outputs and "input_features" in outputs:
798 + # If no mask is provided, create one from input_features
799 + input_features = outputs["input_features"]
800 + if isinstance(input_features, torch.Tensor):
801 + # Create a mask of all ones matching the sequence length
802 + mask = torch.ones(
803 + input_features.shape[0],
804 + input_features.shape[-1],
805 + dtype=torch.long,
806 + )
807 + outputs["feature_attention_mask"] = mask
808 +
809 + # Get processor for chunk counts calculation
810 + processor = self.info.get_hf_processor(**mm_kwargs)
811
812 # Override chunk counts calculation with GLM-ASR specific logic
813 chunk_counts = self._calculate_chunk_counts(
@@ -295,22 +841,58 @@ class GlmAsrMultiModalProcessor(AudioFlamingo3MultiModalProcessor):
841 audio_token_id = processor.audio_token_id
842
843 merge_factor = getattr(config, "merge_factor", DEFAULT_MERGE_FACTOR)
844 + conv_params = getattr(config, "conv_params", DEFAULT_CONV_PARAMS)
845 out_mm_data = out_mm_kwargs.get_data()
846 feature_attention_mask = out_mm_data.get("feature_attention_mask")
847 chunk_counts = out_mm_data.get("chunk_counts")
848
302 - def get_replacement_glmasr(item_idx: int):
303 - conv_params = getattr(config, "conv_params", DEFAULT_CONV_PARAMS)
304 - audio_embeds = out_mm_data.get("audio_embeds")
305 - num_features = _get_num_features_for_item(
306 - feature_attention_mask,
307 - chunk_counts,
308 - item_idx,
309 - audio_embeds,
310 - merge_factor,
311 - conv_params,
849 + # Pre-compute audio output lengths if feature_attention_mask is available
850 + audio_output_lengths: list[int] = []
851 + if feature_attention_mask is not None:
852 + # Compute output lengths for all audio items
853 + from .glmasr_utils import (
854 + _as_list_chunk_counts,
855 + _get_audio_output_lengths_from_mask,
856 )
857
858 + if chunk_counts is not None:
859 + start_idx = 0
860 + for count in _as_list_chunk_counts(chunk_counts):
861 + end_idx = start_idx + count
862 + mask = feature_attention_mask[start_idx:end_idx]
863 + if isinstance(mask, list):
864 + mask = torch.stack(mask)
865 +
866 + lengths = _get_audio_output_lengths_from_mask(
867 + mask, merge_factor, conv_params
868 + )
869 + audio_output_lengths.append(int(lengths.sum().item()))
870 + start_idx = end_idx
871 + else:
872 + # Single chunk per audio
873 + for idx in range(len(feature_attention_mask)):
874 + mask = feature_attention_mask[idx : idx + 1]
875 + if isinstance(mask, list):
876 + mask = torch.tensor(mask).unsqueeze(0)
877 + lengths = _get_audio_output_lengths_from_mask(
878 + mask, merge_factor, conv_params
879 + )
880 + audio_output_lengths.append(int(lengths.sum().item()))
881 +
882 + def get_replacement_glmasr(item_idx: int):
883 + # Use pre-computed lengths if available, otherwise fall back to audio_embeds
884 + if audio_output_lengths:
885 + num_features = audio_output_lengths[item_idx]
886 + else:
887 + audio_embeds = out_mm_data.get("audio_embeds")
888 + if audio_embeds is not None:
889 + embed = audio_embeds[item_idx]
890 + num_features = embed.shape[0]
891 + else:
892 + raise ValueError(
893 + "Either feature_attention_mask or audio_embeds must be provided"
894 + )
895 +
896 if num_features == 0:
897 raise ValueError("Audio is too short")
898
@@ -352,7 +934,12 @@ class GlmAsrForConditionalGeneration(
934 self.config = config
935 self.multimodal_config = multimodal_config
936
355 - self.audio_tower = GlmAsrEncoder(config.audio_config)
937 + # Use optimized vLLM native encoder
938 + self.audio_tower = GlmAsrEncoder(
939 + config.audio_config,
940 + quant_config=quant_config,
941 + prefix=maybe_prefix(prefix, "audio_tower"),
942 + )
943 self.multi_modal_projector = GlmAsrMultiModalProjector(
944 config,
945 quant_config=quant_config,
@@ -419,12 +1006,31 @@ class GlmAsrForConditionalGeneration(
1006 audio_input.get("chunk_counts"), num_chunks=num_chunks
1007 )
1008
1009 + # Convert input_features to model dtype (e.g., bfloat16) to match model weights
1010 + input_features = input_features.to(dtype=self.audio_tower.conv1.weight.dtype)
1011 +
1012 + # audio_tower returns [batch_size, seq_len, hidden_size] where hidden_size=1280
1013 audio_hidden_states = self.audio_tower(input_features).last_hidden_state
1014 +
1015 + # GLM-ASR merges consecutive frames: 4 frames with hidden_size=1280
1016 + # -> 1 frame with intermediate_size=5120
1017 + hidden_size = self.config.audio_config.hidden_size
1018 + intermediate_size = self.config.audio_config.intermediate_size
1019 + merge_ratio = intermediate_size // hidden_size
1020 +
1021 + # Truncate sequence length to be divisible by merge_ratio
1022 + seq_len = audio_hidden_states.shape[1]
1023 + seq_len_truncated = (seq_len // merge_ratio) * merge_ratio
1024 + if seq_len_truncated < seq_len:
1025 + audio_hidden_states = audio_hidden_states[:, :seq_len_truncated, :]
1026 +
1027 + # Reshape to merge consecutive frames
1028 audio_hidden_states = audio_hidden_states.reshape(
1029 num_chunks,
1030 -1,
426 - self.config.audio_config.intermediate_size,
1031 + intermediate_size,
1032 )
1033 +
1034 audio_features = self.multi_modal_projector(audio_hidden_states)
1035
1036 merge_factor = getattr(self.config, "merge_factor", DEFAULT_MERGE_FACTOR)
@@ -453,7 +1059,9 @@ class GlmAsrForConditionalGeneration(
1059 audio_input = self._parse_and_validate_audio_input(**kwargs)
1060 if audio_input is None:
1061 return []
1062 +
1063 masked_audio_features = self._process_audio_input(audio_input)
1064 +
1065 return masked_audio_features
1066
1067 def forward(
vllm/model_executor/models/glmasr_utils.py
+28 -5
@@ -71,14 +71,37 @@ def _get_audio_output_lengths_for_tower(
71 merge_factor: int,
72 conv_params: list[tuple[int, int, int]],
73 ) -> torch.Tensor:
74 + """
75 + Calculate the output lengths after audio processing.
76 +
77 + The output length accounts for:
78 + 1. Convolution layers (downsampling)
79 + 2. Merge factor (further downsampling during projection)
80 +
81 + Args:
82 + audio_tower: The audio encoder module
83 + audio_lengths: Input feature lengths [batch_size]
84 + merge_factor: Factor for merging adjacent features
85 + conv_params: List of (padding, kernel_size, stride) for each conv layer
86 +
87 + Returns:
88 + Output lengths after all processing [batch_size]
89 + """
90 + # First, calculate the output length after convolutions
91 if hasattr(audio_tower, "_get_feat_extract_output_lengths"):
75 - _, audio_output_lengths = audio_tower._get_feat_extract_output_lengths(
92 + _, conv_output_lengths = audio_tower._get_feat_extract_output_lengths(
93 audio_lengths
94 )
78 - return audio_output_lengths
79 - return _get_audio_output_lengths_from_lengths(
80 - audio_lengths, merge_factor, conv_params
81 - )
95 + else:
96 + conv_output_lengths = audio_lengths
97 + for padding, kernel_size, stride in conv_params:
98 + conv_output_lengths = _calculate_conv_output_length(
99 + conv_output_lengths, padding, kernel_size, stride
100 + )
101 +
102 + # Then, apply merge_factor to get final output length
103 + # Formula: (conv_output_lengths - merge_factor) // merge_factor + 1
104 + return (conv_output_lengths - merge_factor) // merge_factor + 1
105
106
107 def _flatten_audio_features_by_length(