diff --git a/python/sglang/srt/models/qwen3_omni_moe.py b/python/sglang/srt/models/qwen3_omni_moe.py index 3fc5d4c6f..badb6be08 100644 --- a/python/sglang/srt/models/qwen3_omni_moe.py +++ b/python/sglang/srt/models/qwen3_omni_moe.py @@ -70,8 +70,18 @@ class Qwen3OmniMoeAudioEncoderLayer(nn.Module): self.dropout = config.dropout self.activation_fn = ACT2FN[config.activation_function] self.activation_dropout = config.activation_dropout - self.fc1 = nn.Linear(self.embed_dim, config.encoder_ffn_dim) - self.fc2 = nn.Linear(config.encoder_ffn_dim, self.embed_dim) + self.fc1 = ColumnParallelLinear( + self.embed_dim, + config.encoder_ffn_dim, + bias=True, + prefix=f"{prefix}.fc1", + ) + self.fc2 = RowParallelLinear( + config.encoder_ffn_dim, + self.embed_dim, + bias=True, + prefix=f"{prefix}.fc2", + ) self.final_layer_norm = nn.LayerNorm(self.embed_dim) def forward( @@ -98,9 +108,9 @@ class Qwen3OmniMoeAudioEncoderLayer(nn.Module): hidden_states = residual + hidden_states residual = hidden_states hidden_states = self.final_layer_norm(hidden_states) - hidden_states = self.fc1(hidden_states) + hidden_states, _ = self.fc1(hidden_states) hidden_states = self.activation_fn(hidden_states) - hidden_states = self.fc2(hidden_states) + hidden_states, _ = self.fc2(hidden_states) hidden_states = residual + hidden_states if hidden_states.dtype == torch.float16: @@ -187,12 +197,10 @@ class Qwen3OmniMoeAudioEncoder(PreTrainedModel): 2, padding=1, ) - self.conv_out = nn.Linear( - config.downsample_hidden_size - * ((((config.num_mel_bins + 1) // 2 + 1) // 2 + 1) // 2), - config.d_model, - bias=False, + conv_out_dim = config.downsample_hidden_size * ( + (((config.num_mel_bins + 1) // 2 + 1) // 2 + 1) // 2 ) + self.conv_out = nn.Linear(conv_out_dim, config.d_model, bias=False) self.proj1 = nn.Linear(config.d_model, config.d_model) self.act = ACT2FN[config.activation_function] self.proj2 = nn.Linear(config.d_model, config.output_dim) @@ -238,23 +246,34 @@ class Qwen3OmniMoeAudioEncoder(PreTrainedModel): padded_feature = nn.utils.rnn.pad_sequence( chunk_list, batch_first=True ).transpose(1, 2) + + # Introduce vectorized mask to avoid many small tensors feature_lens_after_cnn = _get_feat_extract_output_lengths(chunk_lengths) - padded_mask_after_cnn = nn.utils.rnn.pad_sequence( - [ - torch.ones(length, dtype=torch.bool, device=padded_feature.device) - for length in feature_lens_after_cnn - ], - batch_first=True, + max_len_after_cnn = ( + int(feature_lens_after_cnn.max().item()) + if feature_lens_after_cnn.numel() + else 0 ) + + idx = torch.arange(max_len_after_cnn, device=padded_feature.device) + padded_mask_after_cnn = idx.unsqueeze(0) < feature_lens_after_cnn.unsqueeze(1) + padded_feature = padded_feature.unsqueeze(1) - # Split to chunk to avoid OOM during convolution - padded_embeds = [] - for chunk in padded_feature.split(self.conv_chunksize, dim=0): - padded_embed = F.gelu(self.conv2d1(chunk)) + + # Add fast path + chunk normal path + if padded_feature.size(0) <= self.conv_chunksize: + padded_embed = F.gelu(self.conv2d1(padded_feature)) padded_embed = F.gelu(self.conv2d2(padded_embed)) padded_embed = F.gelu(self.conv2d3(padded_embed)) - padded_embeds.append(padded_embed) - padded_embed = torch.cat(padded_embeds, dim=0) + else: + padded_embeds = [] + for chunk in padded_feature.split(self.conv_chunksize, dim=0): + x = F.gelu(self.conv2d1(chunk)) + x = F.gelu(self.conv2d2(x)) + x = F.gelu(self.conv2d3(x)) + padded_embeds.append(x) + padded_embed = torch.cat(padded_embeds, dim=0) + b, c, f, t = padded_embed.size() padded_embed = self.conv_out( padded_embed.permute(0, 3, 1, 2).contiguous().view(b, t, c * f) @@ -271,11 +290,13 @@ class Qwen3OmniMoeAudioEncoder(PreTrainedModel): window_aftercnn = padded_mask_after_cnn.shape[-1] * ( self.n_window_infer // (self.n_window * 2) ) - for cnn_len in aftercnn_lens: - cu_chunk_lens += [window_aftercnn] * (cnn_len // window_aftercnn) + # Use tolist() for efficient batch conversion from tensor to Python + for cnn_len in aftercnn_lens.tolist(): + num_full_chunks = cnn_len // window_aftercnn remainder = cnn_len % window_aftercnn - if remainder != 0: - cu_chunk_lens += [remainder] + cu_chunk_lens.extend([window_aftercnn] * num_full_chunks) + if remainder: + cu_chunk_lens.append(remainder) cu_seqlens = torch.tensor(cu_chunk_lens, device=aftercnn_lens.device).cumsum( -1, dtype=torch.int32 ) @@ -438,9 +459,12 @@ class Qwen3OmniMoeThinkerForConditionalGeneration(Qwen3VLMoeForConditionalGenera ) def get_audio_feature(self, items: List[MultimodalDataItem]): - feature_attention_mask = torch.cat( - [item.feature_attention_mask for item in items], dim=0 - ).type(torch.long) + device = next(self.audio_tower.parameters()).device + feature_attention_mask = ( + torch.cat([item.feature_attention_mask for item in items], dim=0) + .type(torch.long) + .to(device) + ) input_features = ( torch.cat([item.feature for item in items]) .type(self.audio_tower.dtype)