feat: support Qwen3-ASR on NPU device.

xllm/core/framework/hf_model_loader.cpp

@@ -787,6 +787,13 @@ bool HFModelLoader::load_args(const std::string& model_weights_path) {
     return false;
   }
 
+  if (args_.has_feature_extractor() &&
+      !load_audio_preprocessor_args(model_weights_path)) {
+    LOG(ERROR) << "Failed to load audio preprocess args from "
+               << model_weights_path;
+    return false;
+  }
+
   // Some hacky logics to support loading of old models
   // always use float16 for quantization
   // TODO: support quantization for other data types
@@ -1256,4 +1263,44 @@ bool HFModelLoader::load_video_preprocessor_args(
   return true;
 }
 
+bool HFModelLoader::load_audio_preprocessor_args(
+    const std::string& model_weights_path) {
+  // audio preprocessor args
+  JsonReader audio_preprocess_reader;
+  const std::string audio_preprocess_file_path =
+      model_weights_path + "/preprocessor_config.json";
+  if (audio_preprocess_reader.parse(audio_preprocess_file_path)) {
+    LOG(INFO) << "Success to parse audio preprocess args file: "
+              << audio_preprocess_file_path;
+
+    if (audio_preprocess_reader.contains("feature_size")) {
+      args_.mm_audio_feature_size() = audio_preprocess_reader.value_or<int64_t>(
+          "feature_size", args_.mm_audio_feature_size());
+    }
+    if (audio_preprocess_reader.contains("sampling_rate")) {
+      args_.mm_audio_sampling_rate() =
+          audio_preprocess_reader.value_or<int64_t>(
+              "sampling_rate", args_.mm_audio_sampling_rate());
+    }
+    if (audio_preprocess_reader.contains("hop_length")) {
+      args_.mm_audio_hop_length() = audio_preprocess_reader.value_or<int64_t>(
+          "hop_length", args_.mm_audio_hop_length());
+    }
+    if (audio_preprocess_reader.contains("chunk_length")) {
+      args_.mm_audio_chunk_length() = audio_preprocess_reader.value_or<int64_t>(
+          "chunk_length", args_.mm_audio_chunk_length());
+    }
+    if (audio_preprocess_reader.contains("n_fft")) {
+      args_.mm_audio_n_fft() = audio_preprocess_reader.value_or<int64_t>(
+          "n_fft", args_.mm_audio_n_fft());
+    }
+    if (audio_preprocess_reader.contains("dither")) {
+      args_.mm_audio_dither() = audio_preprocess_reader.value_or<double>(
+          "dither", args_.mm_audio_dither());
+    }
+  }
+
+  return true;
+}
+
 }  // namespace xllm

xllm/core/framework/hf_model_loader.h

@@ -51,6 +51,7 @@ class HFModelLoader : public ModelLoader {
   bool load_tokenizer_args(const std::string& model_weights_path);
   bool load_image_preprocessor_args(const std::string& model_weights_path);
   bool load_video_preprocessor_args(const std::string& model_weights_path);
+  bool load_audio_preprocessor_args(const std::string& model_weights_path);
   std::string model_weights_path() const override {
     return model_weights_path_;
   }

xllm/core/framework/model/model_args.h

@@ -239,6 +239,38 @@ struct ModelArgs {
   PROPERTY(float, partial_rotary_factor) = 0.0f;
   PROPERTY(std::vector<std::string>, layer_types) = {};
 
+  // Qwen3 Omni multimodal processor args.
+  PROPERTY(int32_t, mm_position_id_per_seconds) = 0;
+  PROPERTY(double, mm_fps) = 0;
+  PROPERTY(bool, mm_use_audio_in_video) = false;
+
+  // Audio processor args.
+  PROPERTY(bool, has_feature_extractor) = false;
+  PROPERTY(int64_t, mm_audio_feature_size) = 0;
+  PROPERTY(int64_t, mm_audio_sampling_rate) = 0;
+  PROPERTY(int64_t, mm_audio_hop_length) = 0;
+  PROPERTY(int64_t, mm_audio_chunk_length) = 0;
+  PROPERTY(int64_t, mm_audio_n_fft) = 0;
+  PROPERTY(double, mm_audio_dither) = 0.0;
+  PROPERTY(bool, mm_audio_truncation) = false;
+  PROPERTY(bool, mm_audio_do_normalize) = false;
+
+  // Qwen3 audio encoder args.
+  PROPERTY(int32_t, audio_token_id) = 0;
+  PROPERTY(int32_t, audio_start_token_id) = 0;
+  PROPERTY(int32_t, audio_end_token_id) = 0;
+  PROPERTY(int64_t, mm_audio_num_attention_heads) = 0;
+  PROPERTY(int64_t, mm_audio_hidden_size) = 0;
+  PROPERTY(double, mm_audio_layer_norm_eps) = 1e-5;
+  PROPERTY(int64_t, mm_audio_downsample_hidden_size) = 0;
+  PROPERTY(int64_t, mm_audio_num_mel_bins) = 0;
+  PROPERTY(int64_t, mm_audio_max_source_positions) = 0;
+  PROPERTY(int64_t, mm_audio_n_window) = 0;
+  PROPERTY(int64_t, mm_audio_n_window_infer) = 0;
+  PROPERTY(int64_t, mm_audio_conv_chunksize) = 0;
+  PROPERTY(int64_t, mm_audio_encoder_layers) = 0;
+  PROPERTY(int64_t, mm_audio_output_dim) = 0;
+
   // Vision model's dropout
   PROPERTY(float, mm_dropout) = 0.0f;
 
@@ -855,6 +887,33 @@ inline std::ostream& operator<<(std::ostream& os, const ModelArgs& args) {
   os << ", base_image_seq_len: " << args.base_image_seq_len();
   os << ", max_image_seq_len: " << args.max_image_seq_len();
   os << "]";
+  os << ", mm_position_id_per_seconds: " << args.mm_position_id_per_seconds();
+  os << ", mm_use_audio_in_video: " << args.mm_use_audio_in_video();
+  os << ", has_feature_extractor: " << args.has_feature_extractor();
+  os << ", mm_audio_feature_size: " << args.mm_audio_feature_size();
+  os << ", mm_audio_sampling_rate: " << args.mm_audio_sampling_rate();
+  os << ", mm_audio_hop_length: " << args.mm_audio_hop_length();
+  os << ", mm_audio_chunk_length: " << args.mm_audio_chunk_length();
+  os << ", mm_audio_n_fft: " << args.mm_audio_n_fft();
+  os << ", mm_audio_dither: " << args.mm_audio_dither();
+  os << ", mm_audio_truncation: " << args.mm_audio_truncation();
+  os << ", mm_audio_do_normalize: " << args.mm_audio_do_normalize();
+
+  os << ", audio_token_id: " << args.audio_token_id();
+  os << ", mm_audio_num_attention_heads: "
+     << args.mm_audio_num_attention_heads();
+  os << ", mm_audio_hidden_size: " << args.mm_audio_hidden_size();
+  os << ", mm_audio_layer_norm_eps: " << args.mm_audio_layer_norm_eps();
+  os << ", mm_audio_downsample_hidden_size: "
+     << args.mm_audio_downsample_hidden_size();
+  os << ", mm_audio_num_mel_bins: " << args.mm_audio_num_mel_bins();
+  os << ", mm_audio_max_source_positions: "
+     << args.mm_audio_max_source_positions();
+  os << ", mm_audio_n_window: " << args.mm_audio_n_window();
+  os << ", mm_audio_n_window_infer: " << args.mm_audio_n_window_infer();
+  os << ", mm_audio_conv_chunksize: " << args.mm_audio_conv_chunksize();
+  os << ", mm_audio_encoder_layers: " << args.mm_audio_encoder_layers();
+  os << ", mm_audio_output_dim: " << args.mm_audio_output_dim();
   return os;
 }
 

xllm/core/layers/npu/CMakeLists.txt

@@ -15,6 +15,7 @@ cc_library(
     npu_lm_head_impl.h
     npu_qwen2_vision_encoder_layer_impl.h
     npu_qwen2dot5_vision_encoder_layer_impl.h
+    npu_qwen3_audio_encoder_layer_impl.h
     npu_qwen3_vision_encoder_layer_impl.h
     npu_kimik25_vision_encoder_layer_impl.h
     npu_qwen3_moe_decoder_layer_impl.h
@@ -60,6 +61,7 @@ cc_library(
     loader/mistral_decoder_loader.h
     loader/qwen2_vision_encoder_loader.h
     loader/qwen2dot5_vision_encoder_loader.h
+    loader/qwen3_audio_encoder_loader.h
     loader/qwen3_vision_encoder_loader.h
     loader/kimik25_vision_encoder_loader.h
     loader/rms_norm_loader.h
@@ -74,6 +76,7 @@ cc_library(
     npu_lm_head_impl.cpp
     npu_qwen2_vision_encoder_layer_impl.cpp
     npu_qwen2dot5_vision_encoder_layer_impl.cpp
+    npu_qwen3_audio_encoder_layer_impl.cpp
     npu_qwen3_vision_encoder_layer_impl.cpp
     npu_kimik25_vision_encoder_layer_impl.cpp
     npu_qwen3_moe_decoder_layer_impl.cpp
@@ -119,6 +122,7 @@ cc_library(
     loader/mistral_decoder_loader.cpp
     loader/qwen2_vision_encoder_loader.cpp
     loader/qwen2dot5_vision_encoder_loader.cpp
+    loader/qwen3_audio_encoder_loader.cpp
     loader/qwen3_vision_encoder_loader.cpp
     loader/kimik25_vision_encoder_loader.cpp
     loader/rms_norm_loader.cpp

xllm/core/layers/npu/loader/qwen3_audio_encoder_loader.cpp

@@ -0,0 +1,150 @@
+/* Copyright 2026 The xLLM Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    https://github.com/jd-opensource/xllm/blob/main/LICENSE
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "core/layers/npu/loader/qwen3_audio_encoder_loader.h"
+
+#include <cstdint>
+#include <string>
+#include <unordered_map>
+#include <utility>
+#include <vector>
+
+namespace xllm::layer {
+
+namespace {
+
+constexpr int32_t kInputNormWeight = 0;
+constexpr int32_t kInputNormBias = 1;
+constexpr int32_t kPostNormWeight = 2;
+constexpr int32_t kPostNormBias = 3;
+constexpr int32_t kQkvWeight = 4;
+constexpr int32_t kQkvBias = 5;
+constexpr int32_t kAttentionOutWeight = 6;
+constexpr int32_t kAttentionOutBias = 7;
+constexpr int32_t kLinearFc1Weight = 8;
+constexpr int32_t kLinearFc1Bias = 9;
+constexpr int32_t kLinearFc2Weight = 10;
+constexpr int32_t kLinearFc2Bias = 11;
+constexpr int32_t kQueryWeight = 12;
+constexpr int32_t kQueryBias = 13;
+constexpr int32_t kKeyWeight = 14;
+constexpr int32_t kKeyBias = 15;
+constexpr int32_t kValueWeight = 16;
+constexpr int32_t kValueBias = 17;
+
+const std::vector<std::pair<int32_t, std::string>> kWeightMapping = {
+    {kInputNormWeight, "self_attn_layer_norm.weight"},
+    {kInputNormBias, "self_attn_layer_norm.bias"},
+    {kPostNormWeight, "final_layer_norm.weight"},
+    {kPostNormBias, "final_layer_norm.bias"},
+    {kAttentionOutWeight, "self_attn.out_proj.weight"},
+    {kAttentionOutBias, "self_attn.out_proj.bias"},
+    {kLinearFc1Weight, "fc1.weight"},
+    {kLinearFc1Bias, "fc1.bias"},
+    {kLinearFc2Weight, "fc2.weight"},
+    {kLinearFc2Bias, "fc2.bias"},
+    {kQueryWeight, "self_attn.q_proj.weight"},
+    {kQueryBias, "self_attn.q_proj.bias"},
+    {kKeyWeight, "self_attn.k_proj.weight"},
+    {kKeyBias, "self_attn.k_proj.bias"},
+    {kValueWeight, "self_attn.v_proj.weight"},
+    {kValueBias, "self_attn.v_proj.bias"}};
+
+const std::unordered_map<int32_t, int32_t> kWeightShard = {
+    {kAttentionOutWeight, 1},
+    {kLinearFc1Weight, 0},
+    {kLinearFc1Bias, 0},
+    {kLinearFc2Weight, 1},
+};
+
+}  // namespace
+
+Qwen3AudioEncoderLoader::Qwen3AudioEncoderLoader(uint64_t weight_count,
+                                                 const ModelContext& context)
+    : BaseLoader(weight_count, context) {
+  const ParallelArgs& parallel_args = context.get_parallel_args();
+  const torch::TensorOptions options = context.get_tensor_options();
+  encode_param_rank_ = parallel_args.rank();
+  encode_param_world_size_ = parallel_args.world_size();
+  at_weight_tensors_.resize(weight_count);
+  dtype_ = torch::typeMetaToScalarType(options.dtype());
+  for (uint64_t index = 0; index < weight_count; ++index) {
+    at_weight_tensors_[index] = torch::zeros({1}).to(options);
+  }
+}
+
+void Qwen3AudioEncoderLoader::load_state_dict(const StateDict& state_dict) {
+  for (const auto& [index, name] : kWeightMapping) {
+    auto shard = kWeightShard.find(index);
+    if (shard != kWeightShard.end()) {
+      set_weight(state_dict, name, index, shard->second);
+    } else {
+      set_weight(state_dict, name, index);
+    }
+  }
+}
+
+void Qwen3AudioEncoderLoader::verify_loaded_weights() const {
+  for (const auto& [index, name] : kWeightMapping) {
+    CHECK(at_weight_tensors_[index].sizes() != std::vector<int64_t>({1}))
+        << "weight is not loaded for " << name;
+  }
+}
+
+void Qwen3AudioEncoderLoader::merge_loaded_weights() {
+  // Split packed QKV weights when tensor parallelism is enabled.
+  get_weights_col_packed_qkv();
+
+  const torch::Tensor new_qkv_weight =
+      torch::cat({at_weight_tensors_[kQueryWeight],
+                  at_weight_tensors_[kKeyWeight],
+                  at_weight_tensors_[kValueWeight]},
+                 0)
+          .to(device_);
+  at_weight_tensors_[kQkvWeight] = new_qkv_weight;
+  at_weight_tensors_[kQueryWeight] = torch::zeros({1}).to(device_);
+  at_weight_tensors_[kKeyWeight] = torch::zeros({1}).to(device_);
+  at_weight_tensors_[kValueWeight] = torch::zeros({1}).to(device_);
+
+  const torch::Tensor new_qkv_bias =
+      torch::cat({at_weight_tensors_[kQueryBias],
+                  at_weight_tensors_[kKeyBias],
+                  at_weight_tensors_[kValueBias]},
+                 0)
+          .to(device_);
+  at_weight_tensors_[kQkvBias] = new_qkv_bias;
+  at_weight_tensors_[kQueryBias] = torch::zeros({1}).to(device_);
+  at_weight_tensors_[kKeyBias] = torch::zeros({1}).to(device_);
+  at_weight_tensors_[kValueBias] = torch::zeros({1}).to(device_);
+}
+
+void Qwen3AudioEncoderLoader::get_weights_col_packed_qkv() {
+  const int32_t rank = encode_param_rank_;
+  const int32_t world_size = encode_param_world_size_;
+  at_weight_tensors_[kQueryWeight] =
+      at_weight_tensors_[kQueryWeight].chunk(world_size, 0)[rank].to(device_);
+  at_weight_tensors_[kKeyWeight] =
+      at_weight_tensors_[kKeyWeight].chunk(world_size, 0)[rank].to(device_);
+  at_weight_tensors_[kValueWeight] =
+      at_weight_tensors_[kValueWeight].chunk(world_size, 0)[rank].to(device_);
+  at_weight_tensors_[kQueryBias] =
+      at_weight_tensors_[kQueryBias].chunk(world_size, 0)[rank].to(device_);
+  at_weight_tensors_[kKeyBias] =
+      at_weight_tensors_[kKeyBias].chunk(world_size, 0)[rank].to(device_);
+  at_weight_tensors_[kValueBias] =
+      at_weight_tensors_[kValueBias].chunk(world_size, 0)[rank].to(device_);
+}
+
+}  // namespace xllm::layer

xllm/core/layers/npu/loader/qwen3_audio_encoder_loader.h

@@ -0,0 +1,39 @@
+/* Copyright 2026 The xLLM Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    https://github.com/jd-opensource/xllm/blob/main/LICENSE
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#pragma once
+
+#include <cstdint>
+
+#include "core/layers/npu/loader/base_loader.h"
+
+namespace xllm::layer {
+
+class Qwen3AudioEncoderLoader final : public BaseLoader {
+ public:
+  Qwen3AudioEncoderLoader(uint64_t weight_count, const ModelContext& context);
+
+  void load_state_dict(const StateDict& state_dict) override;
+  void verify_loaded_weights() const override;
+  void merge_loaded_weights() override;
+
+ private:
+  void get_weights_col_packed_qkv();
+
+  int32_t encode_param_rank_ = 0;
+  int32_t encode_param_world_size_ = 1;
+};
+
+}  // namespace xllm::layer