anodyine/vllm-omni
1
0 Fork 0
Code Issues 4 Pull Requests 2 Actions Packages Projects Releases Wiki Activity

Merge branch 'main' into dev/vllm-align

Browse Source 
Signed-off-by: tzhouam <tzhouam@connect.ust.hk>
This commit is contained in:
tzhouam committed 2026-01-30 02:53:03 +00:00
2 parents 62dbf07fb6 0e07eb6ff9
commit 4b895459e4
196 files changed
+9406 -4132
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

vllm_omni/model_executor/models/qwen2_5_omni/qwen2_5_omni.py
+3 -3
View File
@@ -22,7 +22,6 @@ from vllm.model_executor.models.qwen2_5_omni_thinker import (
Qwen2_5OmniThinkerProcessingInfo,
)
from vllm.model_executor.models.utils import init_vllm_registered_model, maybe_prefix
from vllm.model_executor.models.vision import get_llm_pos_ids_for_vision
# from vllm.model_executor.models.qwen2_code2wav_dit import Qwen2Code2wav
from vllm.multimodal import MULTIMODAL_REGISTRY
@@ -36,7 +35,8 @@ from vllm_omni.model_executor.custom_process_mixin import CustomProcessMixin
from vllm_omni.model_executor.model_loader.weight_utils import download_weights_from_hf_specific
from vllm_omni.model_executor.models.output_templates import OmniOutput
from vllm_omni.model_executor.models.utils import add_prefix_to_loaded_weights, split_list_into_ranges
from vllm_omni.utils.platform_utils import is_npu
from vllm_omni.model_executor.models.vision import get_llm_pos_ids_for_vision
from vllm_omni.platforms import current_omni_platform
TALKER_CODEC_EOS_TOKEN_ID = 8294
TALKER_CODEC_BOS_TOKEN_ID = 8293
@@ -247,7 +247,7 @@ class Qwen2_5OmniForConditionalGeneration(
if inputs_embeds is not None and inputs_embeds.device != thinker_dev:
inputs_embeds = inputs_embeds.to(thinker_dev)
if is_npu():
if current_omni_platform.is_npu():
# TODO: remove this hack when NPU supports batched inputs properly
thinker_input_ids = input_ids[0] if input_ids is not None and added_batch_dim else input_ids
# For MRoPE, positions shape is [3, num_tokens] (T/H/W), don't slice it
vllm_omni/model_executor/models/qwen2_5_omni/qwen2_5_omni_token2wav.py
+8 -4
View File
@@ -35,7 +35,7 @@ from vllm.v1.sample.metadata import SamplingMetadata
from vllm.v1.sample.sampler import Sampler
from vllm_omni.model_executor.models.qwen2_5_omni.audio_length import cap_and_align_mel_length, resolve_max_mel_frames
from vllm_omni.utils.platform_utils import is_npu
from vllm_omni.platforms import current_omni_platform
# Provide a no-op auto_docstring decorator to satisfy annotations if missing
@@ -732,10 +732,14 @@ def kaiser_sinc_filter1d(cutoff: float, half_width: float, kernel_size: int) ->
beta = 0.0
# TODO: When torch.kaiser_window supports NPU, remove the device="cpu" argument
if is_npu():
if current_omni_platform.is_npu():
kaiser_window = torch.kaiser_window(
kernel_size, beta=beta, periodic=False, dtype=torch.float32, device="cpu"
).to("npu")
elif current_omni_platform.is_xpu():
kaiser_window = torch.kaiser_window(
kernel_size, beta=beta, periodic=False, dtype=torch.float32, device="cpu"
).to("xpu")
else:
kaiser_window = torch.kaiser_window(kernel_size, beta=beta, periodic=False, dtype=torch.float32)
@@ -796,7 +800,7 @@ class UpSample1d(nn.Module):
def forward(self, hidden_states):
channels = hidden_states.shape[1]
if is_npu():
if current_omni_platform.is_npu():
# TODO: When F.pad supports replicate mode on NPU, remove this branch
input_dtype = hidden_states.dtype
# F.pad in NPU doesn't support BF16 when mode is replicate.
@@ -843,7 +847,7 @@ class DownSample1d(nn.Module):
def forward(self, hidden_states):
channels = hidden_states.shape[1]
if is_npu():
if current_omni_platform.is_npu():
input_dtype = hidden_states.dtype
# F.pad in NPU doesn't support BF16 when mode is replicate.
# To ensure the accuracy, manually pad the input tensor.
vllm_omni/model_executor/models/qwen3_omni/qwen3_omni.py
+3 -3
View File
@@ -35,7 +35,7 @@ from vllm.v1.sample.sampler import Sampler
from vllm_omni.model_executor.custom_process_mixin import CustomProcessMixin
from vllm_omni.model_executor.models.output_templates import OmniOutput
from vllm_omni.model_executor.models.utils import add_prefix_to_loaded_weights, safe_tensor_reshape
from vllm_omni.utils.platform_utils import is_npu
from vllm_omni.platforms import current_omni_platform
# Special token IDs for Qwen3 Omni MoE
# Reference: https://huggingface.co/Qwen/Qwen3-Omni-30B-A3B-Instruct/blob/main/tokenizer_config.json
@@ -269,7 +269,7 @@ class Qwen3OmniMoeForConditionalGeneration(
# ========== Stage 1: Thinker ==========
if self.model_stage == "thinker":
thinker_dev = self._module_device(self.thinker)
if is_npu():
if current_omni_platform.is_npu():
# Normalize to batched inputs if needed
_added_batch_dim = False
if input_ids is not None and input_ids.ndim == 1:
@@ -308,7 +308,7 @@ class Qwen3OmniMoeForConditionalGeneration(
"capture_layer_indices": [0, int(accept_layer)],
"return_hidden_states": True,
}
if is_npu():
if current_omni_platform.is_npu():
# TODO: remove this hack when NPU supports batched inputs properly
thinker_input_ids = input_ids[0] if input_ids is not None and _added_batch_dim else input_ids
thinker_inputs_embeds = (
vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_code2wav.py
+1 -1
View File
@@ -219,7 +219,7 @@ class Qwen3OmniMoeCode2Wav(nn.Module):
end_index = codes.shape[-1]
# TODO: need to optimize algorithms, current only support
# chunk_size = left_context_size = 25
if end_index == chunk_size:
if end_index <= chunk_size:
context_size = 0
else:
context_size = left_context_size
vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_moe_code_predictor_mtp.py
+141 -182
View File
@@ -13,9 +13,15 @@ import torch
import torch.nn as nn
import torch.nn.functional as F
from transformers import Cache, PretrainedConfig
from transformers.models.qwen3_omni_moe.modeling_qwen3_omni_moe import Qwen3OmniMoeRotaryEmbedding
from transformers.generation.logits_process import (
LogitsProcessorList,
TopKLogitsWarper,
TopPLogitsWarper,
)
from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS
from vllm.compilation.decorators import support_torch_compile
from vllm.config import CacheConfig, ModelConfig, VllmConfig
from vllm.config import CacheConfig, ModelConfig, VllmConfig, get_current_vllm_config
from vllm.forward_context import get_forward_context
from vllm.logger import init_logger
from vllm.model_executor.layers.layernorm import RMSNorm
from vllm.model_executor.layers.linear import (
@@ -24,81 +30,12 @@ from vllm.model_executor.layers.linear import (
RowParallelLinear,
)
from vllm.model_executor.layers.quantization import QuantizationConfig
from vllm.model_executor.layers.rotary_embedding import get_rope
from vllm.model_executor.layers.vocab_parallel_embedding import VocabParallelEmbedding
from vllm.utils.torch_utils import direct_register_custom_op
logger = init_logger(__name__)
# ============================================================================
# Rotary Embeddings and Helper Functions
# ============================================================================
class Qwen3OmniCodePredictorRotaryEmbedding(nn.Module):
"""Rotary positional embeddings for the code predictor."""
def __init__(self, dim, max_position_embeddings=2048, base=10000):
super().__init__()
self.dim = dim
self.max_position_embeddings = max_position_embeddings
self.base = base
inv_freq = 1.0 / (self.base ** (torch.arange(0, self.dim, 2, dtype=torch.int64).float() / self.dim))
self.register_buffer("inv_freq", inv_freq, persistent=False)
def forward(self, x: torch.Tensor, position_ids: torch.LongTensor) -> tuple[torch.Tensor, torch.Tensor]:
inv_freq = self.inv_freq[None, :, None].float().expand(position_ids.shape[0], -1, 1)
position = position_ids[:, None, :].float()
freqs = (inv_freq @ position).transpose(1, 2)
emb = torch.cat((freqs, freqs), dim=-1)
cos = emb.cos().to(x.dtype)
sin = emb.sin().to(x.dtype)
return cos, sin
def rotate_half(x):
"""Rotates half the hidden dims of the input."""
x1 = x[..., : x.shape[-1] // 2]
x2 = x[..., x.shape[-1] // 2 :]
return torch.cat((-x2, x1), dim=-1)
def apply_rotary_pos_emb(q, k, cos, sin, position_ids=None, unsqueeze_dim=1):
"""Applies Rotary Position Embedding to the query and key tensors.
Args:
q (`torch.Tensor`): The query tensor.
k (`torch.Tensor`): The key tensor.
cos (`torch.Tensor`): The cosine part of the rotary embedding.
sin (`torch.Tensor`): The sine part of the rotary embedding.
position_ids (`torch.Tensor`, *optional*):
Deprecated and unused.
unsqueeze_dim (`int`, *optional*, defaults to 1):
The 'unsqueeze_dim' argument specifies the dimension along which to unsqueeze cos[position_ids] and
sin[position_ids] so that they can be properly broadcasted to the dimensions of q and k. For example, note
that cos[position_ids] and sin[position_ids] have the shape [batch_size, seq_len, head_dim]. Then, if q and
k have the shape [batch_size, heads, seq_len, head_dim], then setting unsqueeze_dim=1 makes
cos[position_ids] and sin[position_ids] broadcastable to the shapes of q and k. Similarly, if q and k have
the shape [batch_size, seq_len, heads, head_dim], then set unsqueeze_dim=2.
Returns:
`tuple(torch.Tensor)` comprising of the query and key tensors rotated using the Rotary Position Embedding.
"""
cos = cos.unsqueeze(unsqueeze_dim)
sin = sin.unsqueeze(unsqueeze_dim)
q_embed = (q * cos) + (rotate_half(q) * sin)
k_embed = (k * cos) + (rotate_half(k) * sin)
return q_embed, k_embed
def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor:
"""Repeat key/value heads for grouped query attention."""
if n_rep == 1:
return hidden_states
bsz, num_kv_heads, seq_len, head_dim = hidden_states.shape
hidden_states = hidden_states[:, :, None, :, :].expand(bsz, num_kv_heads, n_rep, seq_len, head_dim)
return hidden_states.reshape(bsz, num_kv_heads * n_rep, seq_len, head_dim)
# ============================================================================
# Code Predictor Attention Layer
# ============================================================================
@@ -149,6 +86,12 @@ class Qwen3OmniCodePredictorAttention(nn.Module):
prefix=f"{prefix}.o_proj",
disable_tp=True,
)
self.rotary_emb = get_rope(
self.head_dim,
max_position=config.code_predictor_config.max_position_embeddings,
rope_parameters=None,
dual_chunk_attention_config=None,
)
self.q_size = self.num_heads * self.head_dim
self.kv_size = self.num_key_value_heads * self.head_dim
@@ -159,12 +102,19 @@ class Qwen3OmniCodePredictorAttention(nn.Module):
self.is_causal = True
self.config = config
self.attention_backends = ["flash_attention_2", "xformers", "eager", "sdpa"]
cudagraph_mode = get_current_vllm_config().compilation_config.cudagraph_mode
if "flash_attention_2" in ALL_ATTENTION_FUNCTIONS and cudagraph_mode.has_full_cudagraphs():
logger.warning(
f"CUDAGraphMode.{cudagraph_mode.name} is currently not supported "
f"with flash attention for Qwen3-Omni talker MTP."
f"removing flash attention from attention_backends"
)
self.attention_backends.remove("flash_attention_2")
def forward(
self,
hidden_states: torch.Tensor,
causal_mask: torch.Tensor,
cos: torch.Tensor,
sin: torch.Tensor,
past_key_values: Cache | None = None,
cache_position: torch.LongTensor | None = None,
use_cache: bool = False,
@@ -176,34 +126,39 @@ class Qwen3OmniCodePredictorAttention(nn.Module):
q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)
# Reshape for attention
q = q.view(bsz, seq_len, self.num_heads, self.head_dim)
k = k.view(bsz, seq_len, self.num_key_value_heads, self.head_dim)
v = v.view(bsz, seq_len, self.num_key_value_heads, self.head_dim)
q = q.reshape(bsz, seq_len, self.num_heads, self.head_dim)
k = k.reshape(bsz, seq_len, self.num_key_value_heads, self.head_dim)
v = v.reshape(bsz, seq_len, self.num_key_value_heads, self.head_dim)
# Apply normalization
q = self.q_norm(q)
k = self.k_norm(k)
q = self.q_norm(q).contiguous()
k = self.k_norm(k).contiguous()
q = q.reshape(-1, self.q_size)
k = k.reshape(-1, self.kv_size)
# Apply RoPE
q, k = self.rotary_emb(position_ids, q, k)
# Reshape for attention
q = q.reshape(bsz, seq_len, self.num_heads, self.head_dim)
k = k.reshape(bsz, seq_len, self.num_key_value_heads, self.head_dim)
v_heads = v.transpose(1, 2).contiguous()
q_heads = q.transpose(1, 2).contiguous()
k_heads = k.transpose(1, 2).contiguous()
q_heads, k_heads = apply_rotary_pos_emb(q_heads, k_heads, cos, sin)
v_heads = v.transpose(1, 2).contiguous()
if past_key_values is not None:
sin, cos = self.rotary_emb.get_cos_sin(seq_len)
# sin and cos are specific to RoPE models; cache_position needed for the static cache
cache_kwargs = {"sin": sin, "cos": cos, "cache_position": cache_position}
k_heads, v_heads = past_key_values.update(k_heads, v_heads, self.layer_idx, cache_kwargs)
from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS
# Try attention backends in order of preference, with runtime error handling
# This handles cases where the backend is registered but not actually available
attention_backends = ["flash_attention_2", "xformers", "eager", "sdpa"]
attn_output = None
last_error = None
for backend_name in attention_backends:
for backend_name in self.attention_backends:
if backend_name not in ALL_ATTENTION_FUNCTIONS:
continue
@@ -219,15 +174,10 @@ class Qwen3OmniCodePredictorAttention(nn.Module):
scaling=self.head_dim**-0.5,
sliding_window=None,
use_cache=use_cache,
position_ids=position_ids,
position_ids=position_ids[:seq_len].unsqueeze(0),
output_hidden_states=True,
output_attentions=False,
)
# Success - log fallback if not using flash_attention_2
if backend_name != "flash_attention_2":
logger.warning_once(
f"Using {backend_name} attention backend (flash_attention_2 not available or failed)"
)
break
except (ValueError, ImportError, RuntimeError, AttributeError) as e:
# Store error and try next backend
@@ -335,9 +285,6 @@ class Qwen3OmniCodePredictorMTPLayer(nn.Module):
def mtp_block(
self,
hidden_states: torch.Tensor,
causal_mask: torch.Tensor,
cos: torch.Tensor,
sin: torch.Tensor,
past_key_values: Cache | None = None,
cache_position: torch.LongTensor | None = None,
use_cache: bool = False,
@@ -346,9 +293,7 @@ class Qwen3OmniCodePredictorMTPLayer(nn.Module):
# Self-attention with residual
residual = hidden_states
hidden_states = self.input_layernorm(hidden_states)
hidden_states = self.self_attn(
hidden_states, causal_mask, cos, sin, past_key_values, cache_position, use_cache, position_ids
)
hidden_states = self.self_attn(hidden_states, past_key_values, cache_position, use_cache, position_ids)
hidden_states = residual + hidden_states
# MLP with residual
@@ -359,47 +304,6 @@ class Qwen3OmniCodePredictorMTPLayer(nn.Module):
return hidden_states
def forward(
self,
input_ids: torch.Tensor,
positions: torch.Tensor,
previous_hidden_states: torch.Tensor,
inputs_embeds: torch.Tensor | None = None,
spec_step_index: int = 0,
) -> torch.Tensor:
assert inputs_embeds is not None, "inputs_embeds required for MTP"
# Mask position 0 (not needed for MTP)
inputs_embeds[positions == 0] = 0
hidden_states = torch.cat([inputs_embeds, previous_hidden_states], dim=-1)
# Get position info for RoPE
batch_size, seq_len, _ = hidden_states.shape
device = hidden_states.device
position_ids = torch.arange(seq_len, dtype=torch.long, device=device).unsqueeze(0).expand(batch_size, -1)
# Get RoPE embeddings
head_dim = self.self_attn.head_dim
rotary_emb = Qwen3OmniCodePredictorRotaryEmbedding(
head_dim, max_position_embeddings=self.config.code_predictor_config.max_position_embeddings
)
rotary_emb = Qwen3OmniMoeRotaryEmbedding(self.config)
cos, sin = rotary_emb(hidden_states, position_ids)
# Create causal mask
causal_mask = (
torch.triu(torch.ones(seq_len, seq_len, dtype=torch.bool, device=device), diagonal=1)
.unsqueeze(0)
.unsqueeze(0)
)
causal_mask = causal_mask.masked_fill(causal_mask, float("-inf"))
# Forward through MTP block
hidden_states = self.mtp_block(hidden_states, causal_mask, cos, sin)
return hidden_states
class Qwen3OmniCodePredictorBaseModel(nn.Module):
"""
@@ -444,9 +348,6 @@ class Qwen3OmniCodePredictorBaseModel(nn.Module):
self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
# RoPE
self.rotary_emb = Qwen3OmniMoeRotaryEmbedding(config=config)
def forward(
self,
inputs_embeds: torch.Tensor,
@@ -462,7 +363,6 @@ class Qwen3OmniCodePredictorBaseModel(nn.Module):
Args:
inputs_embeds: [batch, seq_len, hidden_size]
attention_mask: Optional attention mask tensor
position_ids: Optional position IDs tensor
past_key_values: Optional cached key-value pairs
use_cache: Whether to use cache
@@ -473,37 +373,11 @@ class Qwen3OmniCodePredictorBaseModel(nn.Module):
Named tuple with .last_hidden_state and .past_key_values attributes
"""
batch_size, seq_len, _ = inputs_embeds.shape
# Create positions tensor if not provided
# positions must be [num_tokens] or [batch_size, seq_len]
if position_ids is None:
if cache_position is not None:
position_ids = cache_position # [num_tokens]
else:
position_ids = torch.arange(seq_len, device=inputs_embeds.device).unsqueeze(0) # [1, seq_len]
else:
position_ids = position_ids.flatten() # Ensure [num_tokens]
# Extract cos/sin from rotary_emb cache
# The cos_sin_cache is [max_pos, rotary_dim * 2]
cos, sin = self.rotary_emb(inputs_embeds, position_ids)
# Create causal mask
device = inputs_embeds.device
causal_mask = (
torch.triu(torch.ones(seq_len, seq_len, dtype=torch.bool, device=device), diagonal=1)
.unsqueeze(0)
.unsqueeze(0)
)
causal_mask = causal_mask.masked_fill(causal_mask, float("-inf"))
# Forward through decoder layers
hidden_states = inputs_embeds
for layer in self.layers:
hidden_states = layer.mtp_block(
hidden_states, causal_mask, cos, sin, past_key_values, cache_position, use_cache, position_ids
)
hidden_states = layer.mtp_block(hidden_states, past_key_values, cache_position, use_cache, position_ids)
# Final norm
hidden_states = self.norm(hidden_states)
@@ -517,6 +391,32 @@ class Qwen3OmniCodePredictorBaseModel(nn.Module):
return self.codec_embedding
def code_predictor_sample(
logits: torch.Tensor,
layer_name: str,
) -> torch.Tensor:
forward_context = get_forward_context()
self = forward_context.no_compile_layers[layer_name]
logits = self.logits_processors(None, logits[:, -1])
probs = F.softmax(logits, dim=-1)
code = torch.multinomial(probs.squeeze(1), num_samples=1) # [batch, 1]
return code
def code_predictor_sample_fake(
logits: torch.Tensor,
layer_name: str,
) -> torch.Tensor:
return torch.empty((logits.shape[0], 1), dtype=torch.int64, device=logits.device)
direct_register_custom_op(
op_name="qwen3_omni_code_predictor_sample",
op_func=code_predictor_sample,
fake_impl=code_predictor_sample_fake,
)
@support_torch_compile
class Qwen3OmniMoeTalkerCodePredictor(nn.Module):
"""
@@ -552,30 +452,89 @@ class Qwen3OmniMoeTalkerCodePredictor(nn.Module):
for _ in range(self.num_code_groups - 1)
]
)
self.logits_processors = LogitsProcessorList(
[
TopKLogitsWarper(top_k=50),
TopPLogitsWarper(top_p=0.8),
]
)
max_batch_size = vllm_config.scheduler_config.max_num_seqs
self.position_ids_buffer = []
for layer_idx in range(self.num_code_groups - 1):
self.position_ids_buffer.append(torch.arange(layer_idx + 2, dtype=torch.int64).repeat(max_batch_size))
compilation_config = get_current_vllm_config().compilation_config
if prefix in compilation_config.static_forward_context:
raise ValueError(f"Duplicate layer name: {prefix}")
compilation_config.static_forward_context[prefix] = self
self.layer_name = prefix
def forward(
self,
inputs_embeds: torch.Tensor, # [batch, seq_len, hidden_size]
layer_idx: int, # Which layer to predict (0-num_layers-2 for layers 1-num_layers-1)
layer0_code: torch.Tensor,
layer0_embed: torch.Tensor,
last_talker_hidden: torch.Tensor,
) -> tuple[torch.Tensor, torch.Tensor]:
"""
Forward pass for code predictor.
Args:
inputs_embeds: Input embeddings [batch_size, seq_len, hidden_size]
layer_idx: Which residual layer to predict (0-num_layers-2 for layers 1-num_layers-1)
layer0_code:
Code index for code-group (layer) 0.
Shape: [batch_size, 1], dtype typically int64.
last_talker_hidden:
Shape: [batch_size, hidden_size].
Returns:
logits: Predicted logits [batch_size, seq_len, vocab_size]
hidden_states: Output hidden states [batch_size, seq_len, hidden_size]
pos_all_layers:
Predicted codes for all code groups, including `layer0_code`.
Shape: [batch_size, num_code_groups, 1].
current_input:
The final input embedding sequence after appending embeddings of all
predicted codes (one token per predicted layer).
Shape: [batch_size, num_code_groups + 2, hidden_size].
"""
# Pass through base model
hidden_states = self.model(inputs_embeds)
pos_codes = [layer0_code] # Start with layer 0: [batch, 1]
try:
current_input = torch.cat([last_talker_hidden, layer0_embed], dim=1) # [batch, 2, hidden_size]
except Exception as e:
print(f"Error in current_input: {e}")
print(f"last_talker_hidden shape: {last_talker_hidden.shape}")
print(f"prev_embed shape: {layer0_embed.shape}")
raise e
batch_size = current_input.shape[0]
# Get logits from corresponding head
logits = self.lm_head[layer_idx](hidden_states)
# Predict all residual layers (layers 1 to num_code_groups-1) autoregressively
begin_pos = 0
for layer_idx in range(self.num_code_groups - 1):
# Input for this layer: [last_talker_hidden, prev_layer_embed]
return logits, hidden_states
# Forward through code_predictor model
position_ids = self.position_ids_buffer[layer_idx][: (layer_idx + 2) * batch_size] # [1, seq_len]
begin_pos += layer_idx + 2
outputs = self.model(
inputs_embeds=current_input,
attention_mask=None,
position_ids=position_ids,
past_key_values=None,
use_cache=False,
cache_position=None,
)
hidden_state = outputs.last_hidden_state # [batch, 2, hidden_size]
# Use the corresponding lm_head for this layer
logits = self.lm_head[layer_idx](hidden_state[:, -1:, :])
code = torch.ops.vllm.qwen3_omni_code_predictor_sample(logits, self.layer_name)
pos_codes.append(code)
# Update prev_embed for next layer (if not last layer)
# layer_idx=0 predicts layer 1, embed with codec_embedding[1]
new_embed = self.model.codec_embedding[layer_idx](code) # [batch, 1, hidden_size]
current_input = torch.cat([current_input, new_embed], dim=1) # [batch, 3~n, hidden_size]
pos_all_layers = torch.stack(pos_codes, dim=1) # [batch, num_code_groups, 1]
return pos_all_layers, current_input
def load_weights(self, weights: list[tuple[str, torch.Tensor]]) -> set[str]:
"""Load weights with mapping for fused QKV and gate_up projections.
vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_moe_talker.py
+8 -54
View File
@@ -4,11 +4,6 @@ from typing import Any
import torch
import torch.nn as nn
import torch.nn.functional as F
from transformers.generation.logits_process import (
LogitsProcessorList,
TopKLogitsWarper,
TopPLogitsWarper,
)
from transformers.models.qwen3_omni_moe.configuration_qwen3_omni_moe import (
Qwen3OmniMoeTalkerConfig,
)
@@ -140,7 +135,10 @@ class Qwen3OmniMoeTalkerForConditionalGeneration(
self.code_predictor = Qwen3OmniMoeTalkerCodePredictor(
vllm_config=vllm_config, prefix=maybe_prefix(prefix, "code_predictor")
)
self.empty_code = torch.empty((1, 0), dtype=torch.long)
self.layer0_embed_buffer = torch.zeros(
(vllm_config.scheduler_config.max_num_seqs, 1, self.config.text_config.hidden_size),
dtype=vllm_config.model_config.dtype,
)
def code_predictor_forward(
self,
@@ -192,62 +190,18 @@ class Qwen3OmniMoeTalkerForConditionalGeneration(
all_codes_per_position = []
middle_hidden_states = [] # Collect hidden states for each position
logits_processors = LogitsProcessorList(
[
TopKLogitsWarper(top_k=top_k),
TopPLogitsWarper(top_p=top_p),
]
)
# Generate residual layers for each position
for pos in range(seq_len):
layer0_code = input_ids[:, pos : pos + 1] # [batch, 1]
# Predict all residual layers (layers 1 to num_code_groups-1) autoregressively
pos_codes = [layer0_code] # Start with layer 0: [batch, 1]
# Initial input: [last_talker_hidden, layer0_embed]
layer0_embed = self.embed_input_ids(layer0_code)
prev_embed = layer0_embed # Track previous layer embedding
try:
current_input = torch.cat([last_talker_hidden, prev_embed], dim=1) # [batch, 2, hidden_size]
except Exception as e:
print(f"Error in current_input: {e}")
print(f"last_talker_hidden shape: {last_talker_hidden.shape}")
print(f"prev_embed shape: {prev_embed.shape}")
raise e
for layer_idx in range(self.num_code_groups - 1):
# Input for this layer: [last_talker_hidden, prev_layer_embed]
# Forward through code_predictor model
outputs = self.code_predictor.model(
inputs_embeds=current_input,
attention_mask=None,
position_ids=None,
past_key_values=None,
use_cache=False,
cache_position=None,
)
hidden_state = outputs.last_hidden_state # [batch, 2, hidden_size]
# Use the corresponding lm_head for this layer
logits = self.code_predictor.lm_head[layer_idx](hidden_state[:, -1:, :]) # [batch, 1, vocab_size]
logits = logits_processors(None, logits[:, -1])
# Sample from the filtered distribution
probs = F.softmax(logits, dim=-1)
code = torch.multinomial(probs.squeeze(1), num_samples=1) # [batch, 1]
pos_codes.append(code)
# Update prev_embed for next layer (if not last layer)
# layer_idx=0 predicts layer 1, embed with codec_embedding[1]
new_embed = self.code_predictor.model.codec_embedding[layer_idx](code) # [batch, 1, hidden_size]
current_input = torch.cat([current_input, new_embed], dim=1) # [batch, 3~n, hidden_size]
self.layer0_embed_buffer[:batch_size].copy_(layer0_embed)
pos_all_layers, current_input = self.code_predictor(
layer0_code, self.layer0_embed_buffer[:batch_size], last_talker_hidden
)
# Stack all layers for this position: [batch, num_code_groups, 1]
pos_all_layers = torch.stack(pos_codes, dim=1) # [batch, num_code_groups, 1]
all_codes_per_position.append(pos_all_layers)
middle_hidden_states.append(current_input[:, 2:-1, :])