diff --git a/xllm/core/framework/hf_model_loader.cpp b/xllm/core/framework/hf_model_loader.cpp index 67b869a92d..097522216e 100644 --- a/xllm/core/framework/hf_model_loader.cpp +++ b/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( + "feature_size", args_.mm_audio_feature_size()); + } + if (audio_preprocess_reader.contains("sampling_rate")) { + args_.mm_audio_sampling_rate() = + audio_preprocess_reader.value_or( + "sampling_rate", args_.mm_audio_sampling_rate()); + } + if (audio_preprocess_reader.contains("hop_length")) { + args_.mm_audio_hop_length() = audio_preprocess_reader.value_or( + "hop_length", args_.mm_audio_hop_length()); + } + if (audio_preprocess_reader.contains("chunk_length")) { + args_.mm_audio_chunk_length() = audio_preprocess_reader.value_or( + "chunk_length", args_.mm_audio_chunk_length()); + } + if (audio_preprocess_reader.contains("n_fft")) { + args_.mm_audio_n_fft() = audio_preprocess_reader.value_or( + "n_fft", args_.mm_audio_n_fft()); + } + if (audio_preprocess_reader.contains("dither")) { + args_.mm_audio_dither() = audio_preprocess_reader.value_or( + "dither", args_.mm_audio_dither()); + } + } + + return true; +} + } // namespace xllm diff --git a/xllm/core/framework/hf_model_loader.h b/xllm/core/framework/hf_model_loader.h index 8dd0590f0f..32830ffb00 100644 --- a/xllm/core/framework/hf_model_loader.h +++ b/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_; } diff --git a/xllm/core/framework/model/model_args.h b/xllm/core/framework/model/model_args.h index c23983bc5d..3da220968f 100644 --- a/xllm/core/framework/model/model_args.h +++ b/xllm/core/framework/model/model_args.h @@ -239,6 +239,38 @@ struct ModelArgs { PROPERTY(float, partial_rotary_factor) = 0.0f; PROPERTY(std::vector, 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; } diff --git a/xllm/core/layers/npu/CMakeLists.txt b/xllm/core/layers/npu/CMakeLists.txt index b84aa3cf82..c659fefb08 100644 --- a/xllm/core/layers/npu/CMakeLists.txt +++ b/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 diff --git a/xllm/core/layers/npu/loader/qwen3_audio_encoder_loader.cpp b/xllm/core/layers/npu/loader/qwen3_audio_encoder_loader.cpp new file mode 100644 index 0000000000..861e1b03c7 --- /dev/null +++ b/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 +#include +#include +#include +#include + +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> 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 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({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 diff --git a/xllm/core/layers/npu/loader/qwen3_audio_encoder_loader.h b/xllm/core/layers/npu/loader/qwen3_audio_encoder_loader.h new file mode 100644 index 0000000000..fc881e75a3 --- /dev/null +++ b/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 + +#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 diff --git a/xllm/core/layers/npu/npu_qwen3_audio_encoder_layer_impl.cpp b/xllm/core/layers/npu/npu_qwen3_audio_encoder_layer_impl.cpp new file mode 100644 index 0000000000..d0d4ef83ed --- /dev/null +++ b/xllm/core/layers/npu/npu_qwen3_audio_encoder_layer_impl.cpp @@ -0,0 +1,158 @@ +/* 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/npu_qwen3_audio_encoder_layer_impl.h" + +#include + +#include + +#include "torch_npu/csrc/core/npu/NPUCachingAllocator.h" + +namespace xllm::layer { + +namespace { +constexpr int32_t kWeightCountPerLayer = 18; +constexpr int32_t kActivationInputCount = 2; +constexpr int32_t kInputCount = kWeightCountPerLayer + kActivationInputCount; +} // namespace + +void NpuQwen3AudioEncoderLayerImpl::param_from_args( + atb_speed::qwen::AudioEncoderLayerParam& param, + const ModelArgs& args, + const ParallelArgs& parallel_args) { + param.isBF16 = args.dtype() == "bfloat16"; + param.rmsNormEps = args.mm_audio_layer_norm_eps(); + param.worldSize = parallel_args.world_size(); + const bool padding_heads = + args.mm_audio_num_attention_heads() % param.worldSize > 0; + if (padding_heads) { + LOG(FATAL) << "You are running qwen3 audio encoder with " << param.worldSize + << " cards, but got attention heads num " + << args.mm_audio_num_attention_heads() + << ". The attention head count must be divisible by the " + "tensor parallel world size."; + } + param.numAttentionHeadsPerRank = + args.mm_audio_num_attention_heads() / param.worldSize; + param.hiddenSizePerAttentionHead = + args.mm_audio_hidden_size() / args.mm_audio_num_attention_heads(); + param.numKeyValueHeadsPerRank = + static_cast(args.mm_audio_num_attention_heads()) / + param.worldSize; + param.rank = parallel_args.rank(); + param.backend = "lccl"; + param.enableLogN = false; +} + +NpuQwen3AudioEncoderLayerImpl::NpuQwen3AudioEncoderLayerImpl( + const ModelContext& context) + : BaseLayer(context) { + const ModelArgs& model_args = context.get_model_args(); + const ParallelArgs& parallel_args = context.get_parallel_args(); + const torch::TensorOptions options = context.get_tensor_options(); + param_from_args(encode_param_, model_args, parallel_args); + atb_weight_tensors_.resize(kWeightCountPerLayer); + dtype_ = torch::typeMetaToScalarType(options.dtype()); + loader_ = + std::make_unique(kWeightCountPerLayer, context); +} + +void NpuQwen3AudioEncoderLayerImpl::merge_loaded_weights() { + loader_->merge_loaded_weights(); + std::vector& at_weight_tensors = + loader_->get_at_weight_tensors(); + c10_npu::NPUCachingAllocator::emptyCache(); + for (int32_t index = 0; index < kWeightCountPerLayer; ++index) { + atb_weight_tensors_[index] = + atb_speed::Utils::AtTensor2Tensor(at_weight_tensors[index]); + } + + init_layer(); +} + +int64_t NpuQwen3AudioEncoderLayerImpl::init_layer() { + name_ = "qwen3_audio_encoder_layer"; + model_name_ = "qwen3_audio"; + CHECK_OPERATION_STATUS_RETURN(init_node(encode_node_, encode_param_)); + return atb::NO_ERROR; +} + +int64_t NpuQwen3AudioEncoderLayerImpl::init_node( + atb_speed::Model::Node& node, + atb_speed::qwen::AudioEncoderLayerParam& param) { + atb::Operation* operation = nullptr; + atb_speed::qwen::Qwen3_Audio_EncoderLayer(param, &operation); + node.operation.reset(operation); + if (node.operation == nullptr) { + LOG(ERROR) << "node.operation is null"; + return -1; + } + if (node.operation->GetInputNum() < kInputCount) { + LOG(ERROR) << "Qwen3 audio encoder operation expects at least " + << kInputCount << " inputs, got " + << node.operation->GetInputNum(); + return -1; + } + node.inTensors.resize(node.operation->GetInputNum()); + node.outTensors.resize(1); + for (int32_t weight_tensor_id = 0; weight_tensor_id < kWeightCountPerLayer; + ++weight_tensor_id) { + node.inTensors.at(weight_tensor_id) = + &atb_weight_tensors_[weight_tensor_id]; + } + + node.variantPack.inTensors.reserve(node.inTensors.size()); + node.variantPack.inTensors.resize(node.inTensors.size()); + node.variantPack.outTensors.reserve(1); + node.variantPack.outTensors.resize(1); + return atb::NO_ERROR; +} + +torch::Tensor NpuQwen3AudioEncoderLayerImpl::forward( + torch::Tensor& x, + torch::Tensor& cu_seqlen, + std::vector& cu_seqlen_vec, + int32_t node_id) { + build_node_variant_pack(encode_node_, x, cu_seqlen, cu_seqlen_vec); + const atb::Status status = execute_node(encode_node_, node_id); + LOG_IF(FATAL, status != 0) + << model_name_ << " execute encode layer failed, error code: " << status; + return x; +} + +void NpuQwen3AudioEncoderLayerImpl::build_node_variant_pack( + atb_speed::Model::Node& node, + torch::Tensor& x, + torch::Tensor& cu_seqlen, + std::vector& cu_seqlen_vec) { + internal_tensors_ = atb_speed::Utils::AtTensor2Tensor(x); + + node.variantPack.inTensors.at(kWeightCountPerLayer) = internal_tensors_; + node.variantPack.inTensors.at(kWeightCountPerLayer + 1) = + atb_speed::Utils::AtTensor2Tensor(cu_seqlen); + node.variantPack.inTensors.at(kWeightCountPerLayer + 1).hostData = + cu_seqlen_vec.data(); + + for (int32_t index = 0; index < kWeightCountPerLayer; ++index) { + CHECK_THROW(node.inTensors.at(index) == nullptr, + model_name_ << " inTensor " << index << " is NULL"); + node.variantPack.inTensors.at(index) = *node.inTensors.at(index); + } + + node.variantPack.outTensors.at(0) = internal_tensors_; +} + +} // namespace xllm::layer diff --git a/xllm/core/layers/npu/npu_qwen3_audio_encoder_layer_impl.h b/xllm/core/layers/npu/npu_qwen3_audio_encoder_layer_impl.h new file mode 100644 index 0000000000..a25a32398b --- /dev/null +++ b/xllm/core/layers/npu/npu_qwen3_audio_encoder_layer_impl.h @@ -0,0 +1,69 @@ +/* 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 +#include +#include + +#include "atb/atb_infer.h" +#include "atb_speed/base/model.h" +#include "core/framework/model/model_args.h" +#include "core/framework/state_dict/state_dict.h" +#include "core/layers/npu/loader/qwen3_audio_encoder_loader.h" +#include "core/layers/npu/npu_base_layer.h" +#include "xllm_atb_layers/models/qwen3_audio/qwen3_audio.h" + +namespace xllm::layer { + +class NpuQwen3AudioEncoderLayerImpl final : public BaseLayer { + public: + explicit NpuQwen3AudioEncoderLayerImpl(const ModelContext& context); + + ~NpuQwen3AudioEncoderLayerImpl() override = default; + + void merge_loaded_weights() override; + + int64_t init_layer() override; + + torch::Tensor forward(torch::Tensor& x, + torch::Tensor& cu_seqlen, + std::vector& cu_seqlen_vec, + int32_t node_id = 0); + + private: + void build_node_variant_pack(atb_speed::Model::Node& node, + torch::Tensor& x, + torch::Tensor& cu_seqlen, + std::vector& cu_seqlen_vec); + + void param_from_args(atb_speed::qwen::AudioEncoderLayerParam& param, + const ModelArgs& args, + const ParallelArgs& parallel_args); + + int64_t init_node(atb_speed::Model::Node& node, + atb_speed::qwen::AudioEncoderLayerParam& param); + + atb_speed::Model::Node encode_node_; + std::string model_name_; + + atb_speed::qwen::AudioEncoderLayerParam encode_param_; + + atb::Tensor internal_tensors_; +}; +TORCH_MODULE(NpuQwen3AudioEncoderLayer); + +} // namespace xllm::layer diff --git a/xllm/core/util/CMakeLists.txt b/xllm/core/util/CMakeLists.txt index 2bc1c7ffb9..418de1904f 100644 --- a/xllm/core/util/CMakeLists.txt +++ b/xllm/core/util/CMakeLists.txt @@ -4,6 +4,7 @@ cc_library( NAME util HDRS + audio_utils.h concurrent_queue.h blocking_counter.h blockingconcurrentqueue.h diff --git a/xllm/core/util/audio_utils.h b/xllm/core/util/audio_utils.h new file mode 100644 index 0000000000..16da2cc708 --- /dev/null +++ b/xllm/core/util/audio_utils.h @@ -0,0 +1,153 @@ +/* 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 +#include + +#include +#include +#include + +namespace xllm::audio_utils { + +inline torch::Tensor hertz_to_mel(const torch::Tensor& freq, + const std::string& mel_scale = "htk") { + CHECK(mel_scale == "slaney" || mel_scale == "htk" || mel_scale == "kaldi") + << "mel_scale must be one of 'htk', 'slaney', or 'kaldi'."; + + if (mel_scale == "htk") { + return 2595.0 * torch::log10(1.0f + freq / 700.0); + } else if (mel_scale == "kaldi") { + const float kaldi_scale = 1127.0; + return kaldi_scale * torch::log1p(freq / 700.0); + } + + const double min_log_hertz = 1000.0; + const double min_log_mel = 15.0; + const double logstep = 27.0 / std::log(6.4); + + const torch::Tensor mels = 3.0 * freq / 200.0; + + const torch::Tensor result = + torch::where(freq >= min_log_hertz, + min_log_mel + torch::log(freq / min_log_hertz) * logstep, + mels); + + return result; +} + +inline torch::Tensor mel_to_hertz(const torch::Tensor& mels, + const std::string& mel_scale = "htk") { + CHECK(mel_scale == "slaney" || mel_scale == "htk" || mel_scale == "kaldi") + << "mel_scale must be one of 'htk', 'slaney', or 'kaldi'."; + + if (mel_scale == "htk") { + return 700.0 * (torch::pow(10.0, mels / 2595.0) - 1.0); + } else if (mel_scale == "kaldi") { + const float kaldi_scale = 1127.0; + return 700.0 * (torch::exp(mels / kaldi_scale) - 1.0); + } + + const double min_log_hertz = 1000.0; + const double min_log_mel = 15.0; + const double logstep = std::log(6.4) / 27.0; + + const torch::Tensor freq = 200 * mels / 3.0; + + const torch::Tensor result = + torch::where(mels >= min_log_mel, + min_log_hertz * torch::exp(logstep * (mels - min_log_mel)), + freq); + + return result; +} + +inline torch::Tensor create_triangular_filter_bank( + const torch::Tensor& fft_freqs, + const torch::Tensor& filter_freqs) { + // fft_freqs: [num_frequency_bins] + // filter_freqs: [num_mel_filters] + + const torch::Tensor filter_diff = torch::diff(filter_freqs); + + const torch::Tensor fft_freqs_expanded = fft_freqs.unsqueeze(1); + const torch::Tensor filter_freqs_expanded = filter_freqs.unsqueeze(0); + + const torch::Tensor slopes = filter_freqs_expanded - fft_freqs_expanded; + const torch::Tensor down_slopes = + -slopes.slice(1, 0, -2) / filter_diff.slice(0, 0, -1); + const torch::Tensor up_slopes = slopes.slice(1, 2) / filter_diff.slice(0, 1); + + auto mel_filters = torch::minimum(down_slopes, up_slopes); + mel_filters = torch::clamp_min(mel_filters, 0.0f); + + return mel_filters; +} + +inline torch::Tensor mel_filter_bank(int64_t num_frequency_bins, + int64_t num_mel_filters, + double min_frequency, + double max_frequency, + int64_t sampling_rate, + const std::string& norm = "", + const std::string& mel_scale = "htk", + bool triangularize_in_mel_space = false) { + CHECK(norm.empty() || norm == "slaney") << "norm must be empty or 'slaney'."; + CHECK_GE(num_frequency_bins, 2); + CHECK_LE(min_frequency, max_frequency); + + if (max_frequency > sampling_rate / 2.0) { + LOG(WARNING) << "max_frequency exceeds the Nyquist frequency."; + } + + const torch::Tensor mel_min_scalar = + hertz_to_mel(torch::tensor(min_frequency), mel_scale); + const torch::Tensor mel_max_scalar = + hertz_to_mel(torch::tensor(max_frequency), mel_scale); + const double mel_min = mel_min_scalar.item(); + const double mel_max = mel_max_scalar.item(); + + const torch::Tensor mel_freqs = + torch::linspace(mel_min, mel_max, num_mel_filters + 2); + torch::Tensor filter_freqs = mel_to_hertz(mel_freqs, mel_scale); + torch::Tensor fft_freqs; + + if (triangularize_in_mel_space) { + const float fft_bin_width = + static_cast(sampling_rate) / ((num_frequency_bins - 1) * 2); + const torch::Tensor indices = + torch::arange(num_frequency_bins, torch::kFloat32); + fft_freqs = hertz_to_mel(fft_bin_width * indices, mel_scale); + filter_freqs = mel_freqs; + } else { + fft_freqs = torch::linspace(0, sampling_rate / 2, num_frequency_bins); + } + + auto mel_filters = create_triangular_filter_bank(fft_freqs, filter_freqs); + + if (norm == "slaney") { + const torch::Tensor filter_widths = + filter_freqs.slice(0, 2, num_mel_filters + 2) - + filter_freqs.slice(0, 0, num_mel_filters); + const torch::Tensor enorm = 2.0 / filter_widths; + mel_filters *= enorm.unsqueeze(0); + } + + return mel_filters; +} + +} // namespace xllm::audio_utils diff --git a/xllm/models/llm/npu/llm_model_base.h b/xllm/models/llm/npu/llm_model_base.h index 1c354918d7..15b674fba5 100644 --- a/xllm/models/llm/npu/llm_model_base.h +++ b/xllm/models/llm/npu/llm_model_base.h @@ -23,6 +23,7 @@ limitations under the License. #include #include #include +#include #include #include "core/common/global_flags.h" @@ -469,6 +470,12 @@ class LlmForCausalLMImplBase : public torch::nn::Module { virtual void load_model( std::unique_ptr loader, std::string prefix = "model." /*llm model weight prefix*/) { + load_model(std::move(loader), std::move(prefix), "lm_head."); + } + + void load_model(std::unique_ptr loader, + std::string prefix, + std::string lm_head_prefix) { for (const auto& state_dict : loader->get_state_dicts()) { // The same model_type may come from checkpoints with different top-level // weight prefixes. Try these candidate prefixes in order to improve @@ -486,7 +493,7 @@ class LlmForCausalLMImplBase : public torch::nn::Module { prefix + "embed_tokens.", "embed_tokens."})); } else { npu_lm_head_->load_state_dict( - state_dict->get_dict_with_prefix("lm_head.")); + state_dict->get_dict_with_prefix(lm_head_prefix)); } } } @@ -497,7 +504,7 @@ class LlmForCausalLMImplBase : public torch::nn::Module { if (tie_word_embeddings) { npu_lm_head_->verify_loaded_weights("embed_tokens."); } else { - npu_lm_head_->verify_loaded_weights("lm_head."); + npu_lm_head_->verify_loaded_weights(lm_head_prefix); } } diff --git a/xllm/models/models.h b/xllm/models/models.h index a07ba8c0ad..fa4d8e1dea 100644 --- a/xllm/models/models.h +++ b/xllm/models/models.h @@ -64,6 +64,8 @@ limitations under the License. #include "vlm/npu/oxygen_vlm.h" // IWYU pragma: keep #include "vlm/npu/qwen2_5_vl.h" // IWYU pragma: keep #include "vlm/npu/qwen2_vl.h" // IWYU pragma: keep +#include "vlm/npu/qwen3_asr.h" // IWYU pragma: keep +#include "vlm/npu/qwen3_omni_moe.h" // IWYU pragma: keep #include "vlm/npu/qwen3_vl.h" // IWYU pragma: keep #include "vlm/npu/qwen3_vl_moe.h" // IWYU pragma: keep #include "vlm/qwen3_5.h" // IWYU pragma: keep diff --git a/xllm/models/vlm/npu/qwen3_asr.h b/xllm/models/vlm/npu/qwen3_asr.h new file mode 100644 index 0000000000..e05dfe4dd4 --- /dev/null +++ b/xllm/models/vlm/npu/qwen3_asr.h @@ -0,0 +1,275 @@ +/* Copyright 2025 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 +#include +#include +#include +#include +#include + +#include "models/llm/npu/qwen3.h" +#include "models/model_registry.h" +#include "models/vlm/mposition/mposition.h" +#include "models/vlm/npu/qwen3_audio_encoder.h" +#include "processors/qwen3_asr_processor.h" +#include "processors/qwen3_audio_common.h" + +namespace xllm::npu::model { + +inline void load_qwen3_asr_model_args(const JsonReader& json, ModelArgs* args) { + LOAD_ARG_OR(model_type, "model_type", "qwen3_asr"); + LOAD_ARG_OR(has_feature_extractor, "has_feature_extractor", true); + LOAD_ARG_OR(mm_audio_feature_size, "feature_size", 128); + LOAD_ARG_OR(mm_audio_sampling_rate, "sampling_rate", 16000); + LOAD_ARG_OR(mm_audio_hop_length, "hop_length", 160); + LOAD_ARG_OR(mm_audio_chunk_length, "chunk_length", 30); + LOAD_ARG_OR(mm_audio_n_fft, "n_fft", 400); + LOAD_ARG_OR(mm_audio_dither, "dither", 0.0); + LOAD_ARG_OR(mm_audio_truncation, "truncation", false); + LOAD_ARG_OR(mm_audio_do_normalize, "do_normalize", false); + + LOAD_ARG_OR(audio_token_id, "thinker_config.audio_token_id", 151676); + LOAD_ARG_OR( + audio_start_token_id, "thinker_config.audio_start_token_id", 151669); + LOAD_ARG_OR(audio_end_token_id, "thinker_config.audio_end_token_id", 151670); + LOAD_ARG_OR(dtype, "thinker_config.dtype", "bfloat16"); + + LOAD_ARG_OR( + attention_bias, "thinker_config.text_config.attention_bias", false); + LOAD_ARG_OR( + attention_dropout, "thinker_config.text_config.attention_dropout", 0.0); + LOAD_ARG_OR(bos_token_id, "thinker_config.text_config.bos_token_id", 151643); + LOAD_ARG_OR(eos_token_id, "thinker_config.text_config.eos_token_id", 151645); + LOAD_ARG_OR(hidden_act, "thinker_config.text_config.hidden_act", "silu"); + LOAD_ARG_OR(hidden_size, "thinker_config.text_config.hidden_size", 2048); + LOAD_ARG_OR( + intermediate_size, "thinker_config.text_config.intermediate_size", 6144); + LOAD_ARG_OR(max_position_embeddings, + "thinker_config.text_config.max_position_embeddings", + 65536); + LOAD_ARG_OR( + max_window_layers, "thinker_config.text_config.max_window_layers", 28); + LOAD_ARG_OR(n_heads, "thinker_config.text_config.num_attention_heads", 16); + LOAD_ARG_OR(n_layers, "thinker_config.text_config.num_hidden_layers", 28); + LOAD_ARG_OR(n_kv_heads, "thinker_config.text_config.num_key_value_heads", 8); + LOAD_ARG_OR(rms_norm_eps, "thinker_config.text_config.rms_norm_eps", 1e-06); + LOAD_ARG_OR( + sliding_window, "thinker_config.text_config.sliding_window", 32768); + LOAD_ARG_OR(tie_word_embeddings, + "thinker_config.text_config.tie_word_embeddings", + true); + LOAD_ARG_OR( + initializer_range, "thinker_config.text_config.initializer_range", 0.02); + LOAD_ARG_OR(use_sliding_window, + "thinker_config.text_config.use_sliding_window", + false); + LOAD_ARG_OR_FUNC(head_dim, "thinker_config.text_config.head_dim", [args] { + return args->hidden_size() / args->n_heads(); + }); + LOAD_ARG_OR(rope_scaling_rope_type, + "thinker_config.text_config.rope_scaling.type", + "mrope"); + LOAD_ARG(rope_scaling_mrope_section, + "thinker_config.text_config.rope_scaling.mrope_section"); + LOAD_ARG_OR(rope_theta, "thinker_config.text_config.rope_theta", 1000000.0f); + LOAD_ARG_OR(vocab_size, "thinker_config.text_config.vocab_size", 151936); + + if (args->rope_scaling_rope_type() == "default") { + args->rope_scaling_rope_type() = "mrope"; + } + SET_ARG(stop_token_ids, std::unordered_set({151643, 151645})); + + LOAD_ARG_OR(mm_audio_num_attention_heads, + "thinker_config.audio_config.encoder_attention_heads", + 16); + LOAD_ARG_OR( + mm_audio_hidden_size, "thinker_config.audio_config.d_model", 1024); + LOAD_ARG_OR(mm_audio_layer_norm_eps, + "thinker_config.audio_config.layer_norm_eps", + 1e-5); + LOAD_ARG_OR(mm_audio_downsample_hidden_size, + "thinker_config.audio_config.downsample_hidden_size", + 480); + LOAD_ARG_OR( + mm_audio_num_mel_bins, "thinker_config.audio_config.num_mel_bins", 128); + LOAD_ARG_OR(mm_audio_max_source_positions, + "thinker_config.audio_config.max_source_positions", + 1500); + LOAD_ARG_OR(mm_audio_n_window, "thinker_config.audio_config.n_window", 50); + LOAD_ARG_OR(mm_audio_n_window_infer, + "thinker_config.audio_config.n_window_infer", + 800); + LOAD_ARG_OR(mm_audio_conv_chunksize, + "thinker_config.audio_config.conv_chunksize", + 500); + LOAD_ARG_OR(mm_audio_encoder_layers, + "thinker_config.audio_config.encoder_layers", + 24); + LOAD_ARG_OR( + mm_audio_output_dim, "thinker_config.audio_config.output_dim", 2048); +} + +class Qwen3ASRForConditionalGenerationImpl final : public torch::nn::Module { + public: + explicit Qwen3ASRForConditionalGenerationImpl(const ModelContext& context) + : model_args_(context.get_model_args()), + options_(context.get_tensor_options()) { + audio_tower_ = register_module("audio_tower", Qwen3AudioEncoder(context)); + language_model_ = + register_module("language_model", QWen3ForCausalLM(context)); + } + + void prepare_encoder_input(const ModelInputParams& input_params, + std::optional& audio_inputs) { + const auto& mm_data = input_params.multimodal.mm_data; + torch::Tensor input_features; + if (std::optional res = + mm_data.get(qwen3_audio::kInputFeaturesKey)) { + input_features = res.value(); + } + + torch::Tensor feature_lengths; + if (std::optional res = + mm_data.get(qwen3_audio::kFeatureLengthKey)) { + feature_lengths = res.value(); + } + + torch::Tensor feature_origin_lengths; + if (std::optional res = + mm_data.get(qwen3_audio::kFeatureOriginLengthsKey)) { + feature_origin_lengths = res.value(); + } + + if (input_features.defined() && feature_lengths.defined() && + feature_origin_lengths.defined()) { + audio_inputs = Qwen3AudioInputs{ + input_features, feature_lengths, feature_origin_lengths}; + } + } + + MMDict get_multimodal_embeddings(const ModelInputParams& input_params) { + std::optional audio_input; + prepare_encoder_input(input_params, audio_input); + MMDict multimodal_embeds; + if (audio_input) { + const torch::Tensor feature_origin_lengths = + audio_input->feature_origin_lengths.to(options_.device(), + torch::kLong); + + const torch::Tensor input_features = + audio_input->input_features.permute({1, 0}).to(options_); + + const torch::Tensor audio_embeds = + audio_tower_->forward(input_features, feature_origin_lengths); + + const torch::Tensor audio_tokens = + audio_input->feature_lengths.cpu().contiguous().to(torch::kLong); + + std::vector feature_lens_vec( + audio_tokens.data_ptr(), + audio_tokens.data_ptr() + audio_tokens.numel()); + + multimodal_embeds[get_embedding_key(MMType::AUDIO)] = + audio_embeds.split(feature_lens_vec, /*dim=*/0); + } + return multimodal_embeds; + } + + torch::Tensor merge_multimodal_embeddings( + torch::Tensor inputs_embeds, + const torch::Tensor& multimodal_embeds, + const torch::Tensor& is_multimodal) { + inputs_embeds.index_put_({is_multimodal}, multimodal_embeds); + return inputs_embeds; + } + + torch::Tensor get_input_embeddings(const torch::Tensor input_ids, + const ModelInputParams& input_params) { + const auto& mm_data = input_params.multimodal.mm_data; + torch::Tensor inputs_embeds = + language_model_->get_input_embeddings(input_ids); + std::optional audio_embeds = + mm_data.get(get_embedding_key(MMType::AUDIO)); + std::optional audio_mask = + mm_data.get(qwen3_audio::kMaskKey); + if (audio_embeds.has_value() && audio_mask.has_value()) { + inputs_embeds = merge_multimodal_embeddings( + inputs_embeds, audio_embeds.value(), audio_mask.value()); + } + return inputs_embeds; + } + + ModelOutput forward(const torch::Tensor& tokens, + const torch::Tensor& positions, + std::vector& kv_caches, + const ModelInputParams& input_params) { + return language_model_(tokens, positions, kv_caches, input_params); + } + + torch::Tensor logits(const torch::Tensor& hidden_states, + const torch::Tensor& selected_indices) { + return language_model_->logits(hidden_states, selected_indices); + } + + void load_model(std::unique_ptr loader) { + for (const auto& state_dict : loader->get_state_dicts()) { + audio_tower_->load_state_dict( + state_dict->get_dict_with_prefix("thinker.audio_tower.")); + } + audio_tower_->verify_loaded_weights("thinker.audio_tower."); + audio_tower_->merge_loaded_weights(); + audio_tower_->to(options_.device(), + torch::typeMetaToScalarType(options_.dtype())); + + if (!model_args_.encoder_embedding_mode()) { + language_model_->load_model( + std::move(loader), "thinker.model.", "thinker.lm_head."); + } + } + + layer::NpuLmHead get_npu_lm_head() { + return language_model_->get_npu_lm_head(); + } + + void set_npu_lm_head(layer::NpuLmHead& head) { + language_model_->set_npu_lm_head(head); + } + + layer::NpuWordEmbedding get_npu_word_embedding() { + return language_model_->get_npu_word_embedding(); + } + + void set_npu_word_embedding(layer::NpuWordEmbedding& npu_word_embedding) { + language_model_->set_npu_word_embedding(npu_word_embedding); + } + + private: + ModelArgs model_args_; + torch::TensorOptions options_; + Qwen3AudioEncoder audio_tower_{nullptr}; + QWen3ForCausalLM language_model_{nullptr}; +}; +TORCH_MODULE(Qwen3ASRForConditionalGeneration); + +REGISTER_MULTIMODAL_PROCESSOR(qwen3_asr, Qwen3ASRMultimodalProcessor); +REGISTER_CAUSAL_VLM_MODEL(qwen3_asr, Qwen3ASRForConditionalGeneration); +REGISTER_MPOSITION_GENERATOR(qwen3_asr, xllm::Qwen3VLMPositionGenerator); + +REGISTER_MODEL_ARGS(qwen3_asr, [&] { load_qwen3_asr_model_args(json, args); }); + +} // namespace xllm::npu::model diff --git a/xllm/models/vlm/npu/qwen3_audio_encoder.h b/xllm/models/vlm/npu/qwen3_audio_encoder.h new file mode 100644 index 0000000000..6a8926ddac --- /dev/null +++ b/xllm/models/vlm/npu/qwen3_audio_encoder.h @@ -0,0 +1,432 @@ +/* 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 + +#include +#include +#include +#include +#include +#include +#include + +#include "core/framework/model_context.h" +#include "core/framework/state_dict/state_dict.h" +#include "core/framework/state_dict/utils.h" +#include "core/layers/npu/npu_qwen3_audio_encoder_layer_impl.h" +#include "processors/qwen3_audio_common.h" + +namespace xllm::npu::model { + +class Qwen3AudioSinusoidalPositionEmbeddingImpl final + : public torch::nn::Module { + public: + Qwen3AudioSinusoidalPositionEmbeddingImpl(int64_t length, + int64_t channels, + double max_timescale = 10000.0) { + CHECK_EQ(channels % 2, 0) + << "Qwen3 audio sinusoidal position embedding requires an even " + "channel count."; + + const double log_timescale_increment = + std::log(max_timescale) / (channels / 2 - 1); + const torch::Tensor inverse_timescales = + torch::exp(-log_timescale_increment * torch::arange(channels / 2)) + .to(torch::kFloat32); + const torch::Tensor scaled_time = + torch::arange(length).unsqueeze(1) * inverse_timescales.unsqueeze(0); + position_embedding_ = + torch::cat({torch::sin(scaled_time), torch::cos(scaled_time)}, 1); + } + + torch::Tensor forward(int64_t sequence_length) { + return position_embedding_.slice(0, 0, sequence_length); + } + + private: + torch::Tensor position_embedding_; +}; +TORCH_MODULE(Qwen3AudioSinusoidalPositionEmbedding); + +class Qwen3AudioEncoderBlockImpl final : public torch::nn::Module { + public: + explicit Qwen3AudioEncoderBlockImpl(const ModelContext& context) { + encoder_layer_ = register_module("encoder_layer", + layer::NpuQwen3AudioEncoderLayer(context)); + } + + torch::Tensor forward(torch::Tensor& hidden_states, + torch::Tensor& cumulative_sequence_lengths, + std::vector& cumulative_sequence_lengths_vec, + int32_t node_id) { + return encoder_layer_(hidden_states, + cumulative_sequence_lengths, + cumulative_sequence_lengths_vec, + node_id); + } + + void load_state_dict(const StateDict& state_dict) { + encoder_layer_->load_state_dict(state_dict); + } + + void verify_loaded_weights() const { + encoder_layer_->verify_loaded_weights(); + } + + void merge_loaded_weights() { encoder_layer_->merge_loaded_weights(); } + + private: + layer::NpuQwen3AudioEncoderLayer encoder_layer_{nullptr}; +}; +TORCH_MODULE(Qwen3AudioEncoderBlock); + +class Qwen3AudioEncoderImpl final : public torch::nn::Module { + public: + explicit Qwen3AudioEncoderImpl(const ModelContext& context) { + const ModelArgs& model_args = context.get_model_args(); + options_ = context.get_tensor_options(); + const int64_t downsample_hidden_size = + model_args.mm_audio_downsample_hidden_size(); + const int64_t num_mel_bins = model_args.mm_audio_num_mel_bins(); + const int64_t max_source_positions = + model_args.mm_audio_max_source_positions(); + embed_dim_ = model_args.mm_audio_hidden_size(); + n_window_ = model_args.mm_audio_n_window(); + n_window_infer_ = model_args.mm_audio_n_window_infer(); + conv_chunk_size_ = model_args.mm_audio_conv_chunksize(); + CHECK_GT(embed_dim_, 0); + CHECK_GT(n_window_, 0); + CHECK_GE(n_window_infer_, n_window_ * 2); + CHECK_EQ(n_window_infer_ % (n_window_ * 2), 0); + CHECK_GT(conv_chunk_size_, 0); + + positional_embedding_ = + register_module("positional_embedding", + Qwen3AudioSinusoidalPositionEmbedding( + max_source_positions, embed_dim_)); + + layers_ = register_module("layers", torch::nn::ModuleList()); + for (int64_t index = 0; index < model_args.mm_audio_encoder_layers(); + ++index) { + Qwen3AudioEncoderBlock layer(context); + layers_->push_back(layer); + } + + ln_post_ = register_module( + "ln_post", + torch::nn::LayerNorm(torch::nn::LayerNormOptions({embed_dim_}) + .elementwise_affine(true))); + + conv2d1_ = register_module( + "conv2d1", + torch::nn::Conv2d(torch::nn::Conv2dOptions(1, downsample_hidden_size, 3) + .stride(2) + .padding(1) + .bias(true))); + conv2d2_ = register_module( + "conv2d2", + torch::nn::Conv2d(torch::nn::Conv2dOptions( + downsample_hidden_size, downsample_hidden_size, 3) + .stride(2) + .padding(1) + .bias(true))); + conv2d3_ = register_module( + "conv2d3", + torch::nn::Conv2d(torch::nn::Conv2dOptions( + downsample_hidden_size, downsample_hidden_size, 3) + .stride(2) + .padding(1) + .bias(true))); + + const int64_t conv_output_dim = (((num_mel_bins + 1) / 2 + 1) / 2 + 1) / 2; + conv_out_ = register_module( + "conv_out", + torch::nn::Linear( + torch::nn::LinearOptions(downsample_hidden_size * conv_output_dim, + embed_dim_) + .bias(false))); + proj1_ = register_module( + "proj1", + torch::nn::Linear( + torch::nn::LinearOptions(embed_dim_, embed_dim_).bias(true))); + proj2_ = register_module( + "proj2", + torch::nn::Linear(torch::nn::LinearOptions( + embed_dim_, model_args.mm_audio_output_dim()) + .bias(true))); + } + + torch::Tensor forward(const torch::Tensor& input_features, + const torch::Tensor& feature_lengths) { + const torch::Tensor output_lengths = + qwen3_audio::get_feature_output_lengths(feature_lengths, n_window_ * 2); + const torch::Tensor chunk_counts = + torch::ceil(feature_lengths / (n_window_ * 2)).to(torch::kLong); + const int64_t total_chunks = chunk_counts.sum().item(); + + torch::Tensor chunk_lengths = + torch::full({total_chunks}, + n_window_ * 2, + torch::TensorOptions() + .dtype(torch::kLong) + .device(feature_lengths.device())); + const torch::Tensor padded_chunk_counts = torch::nn::functional::pad( + chunk_counts, torch::nn::functional::PadFuncOptions({1, 0}).value(-1)); + const torch::Tensor tail_chunk_indices = + padded_chunk_counts.cumsum(0).slice(0, 1); + const torch::Tensor remainder = feature_lengths % (n_window_ * 2); + chunk_lengths.index_put_({torch::indexing::TensorIndex(tail_chunk_indices)}, + remainder); + chunk_lengths.index_put_({torch::indexing::TensorIndex(chunk_lengths == 0)}, + n_window_ * 2); + + const torch::Tensor transposed_features = input_features.t(); + const torch::Tensor chunk_lengths_cpu = + chunk_lengths.to(torch::kCPU).contiguous(); + const c10::IntArrayRef split_sizes( + chunk_lengths_cpu.data_ptr(), + static_cast(chunk_lengths_cpu.size(0))); + const std::vector chunks = + transposed_features.split_with_sizes(split_sizes, 0); + torch::Tensor padded_features = + torch::nn::utils::rnn::pad_sequence(chunks, true).transpose(1, 2); + const torch::Tensor chunk_output_lengths = + qwen3_audio::get_feature_output_lengths(chunk_lengths.to(torch::kLong), + n_window_ * 2); + + std::vector mask_tensors; + mask_tensors.reserve(static_cast(chunk_output_lengths.size(0))); + for (int64_t index = 0; index < chunk_output_lengths.size(0); ++index) { + const int64_t length = chunk_output_lengths[index].item(); + mask_tensors.emplace_back( + torch::full({length}, + 1, + torch::TensorOptions() + .dtype(torch::kBool) + .device(padded_features.device()))); + } + const torch::Tensor padded_mask_after_cnn = + torch::nn::utils::rnn::pad_sequence(mask_tensors, true); + + padded_features = padded_features.unsqueeze(1); + std::vector padded_embeddings; + const int64_t batch_size = padded_features.size(0); + padded_embeddings.reserve(static_cast( + (batch_size + conv_chunk_size_ - 1) / conv_chunk_size_)); + for (int64_t start = 0; start < batch_size; start += conv_chunk_size_) { + const int64_t end = std::min(start + conv_chunk_size_, batch_size); + const torch::Tensor chunk = padded_features.slice(0, start, end); + torch::Tensor embedding = torch::gelu(conv2d1_(chunk)); + embedding = torch::gelu(conv2d2_(embedding)); + embedding = torch::gelu(conv2d3_(embedding)); + padded_embeddings.emplace_back(embedding); + } + + torch::Tensor padded_embedding = torch::cat(padded_embeddings, 0); + const int64_t batch = padded_embedding.size(0); + const int64_t channels = padded_embedding.size(1); + const int64_t frequency = padded_embedding.size(2); + const int64_t time = padded_embedding.size(3); + const torch::Tensor reshaped = + padded_embedding.permute({0, 3, 1, 2}) + .contiguous() + .view({batch, time, channels * frequency}); + padded_embedding = conv_out_(reshaped); + const torch::Tensor position_embedding = + positional_embedding_->forward(padded_embedding.size(1)) + .unsqueeze(0) + .to(options_.device(), padded_embedding.dtype()); + padded_embedding = padded_embedding + position_embedding; + + torch::Tensor hidden_states = + padded_embedding + .masked_select( + padded_mask_after_cnn.unsqueeze(-1).expand_as(padded_embedding)) + .view({-1, padded_embedding.size(-1)}); + const int64_t window_after_cnn = + padded_mask_after_cnn.size(-1) * (n_window_infer_ / (n_window_ * 2)); + + std::vector cumulative_chunk_lengths; + for (int64_t index = 0; index < output_lengths.size(0); ++index) { + const int64_t cnn_length = output_lengths[index].item(); + const int64_t full_windows = cnn_length / window_after_cnn; + for (int64_t window = 0; window < full_windows; ++window) { + cumulative_chunk_lengths.emplace_back( + static_cast(window_after_cnn)); + } + const int64_t tail_length = cnn_length % window_after_cnn; + if (tail_length != 0) { + cumulative_chunk_lengths.emplace_back( + static_cast(tail_length)); + } + } + + torch::Tensor cumulative_sequence_lengths = + torch::tensor(cumulative_chunk_lengths, + torch::TensorOptions() + .device(output_lengths.device()) + .dtype(torch::kInt32)); + const torch::Tensor cumulative_sequence_lengths_cpu = + cumulative_sequence_lengths.cpu().contiguous(); + std::vector cumulative_sequence_lengths_vec( + cumulative_sequence_lengths_cpu.data_ptr(), + cumulative_sequence_lengths_cpu.data_ptr() + + cumulative_sequence_lengths_cpu.numel()); + + const int32_t layer_count = static_cast(layers_->size()); + for (int32_t index = 0; index < layer_count; ++index) { + hidden_states = layers_[index]->as()->forward( + hidden_states, + cumulative_sequence_lengths, + cumulative_sequence_lengths_vec, + index); + } + + hidden_states = ln_post_(hidden_states); + hidden_states = proj1_(hidden_states); + hidden_states = torch::gelu(hidden_states); + return proj2_(hidden_states); + } + + void load_state_dict(const StateDict& state_dict) { + weight::load_weight(state_dict, + "conv_out.weight", + conv_out_->weight, + is_conv_out_weight_loaded_); + weight::load_weight( + state_dict, "proj1.weight", proj1_->weight, is_proj1_weight_loaded_); + weight::load_weight( + state_dict, "proj1.bias", proj1_->bias, is_proj1_bias_loaded_); + weight::load_weight( + state_dict, "proj2.weight", proj2_->weight, is_proj2_weight_loaded_); + weight::load_weight( + state_dict, "proj2.bias", proj2_->bias, is_proj2_bias_loaded_); + weight::load_weight(state_dict, + "conv2d1.weight", + conv2d1_->weight, + is_conv2d1_weight_loaded_); + weight::load_weight( + state_dict, "conv2d1.bias", conv2d1_->bias, is_conv2d1_bias_loaded_); + weight::load_weight(state_dict, + "conv2d2.weight", + conv2d2_->weight, + is_conv2d2_weight_loaded_); + weight::load_weight( + state_dict, "conv2d2.bias", conv2d2_->bias, is_conv2d2_bias_loaded_); + weight::load_weight(state_dict, + "conv2d3.weight", + conv2d3_->weight, + is_conv2d3_weight_loaded_); + weight::load_weight( + state_dict, "conv2d3.bias", conv2d3_->bias, is_conv2d3_bias_loaded_); + weight::load_weight(state_dict, + "ln_post.weight", + ln_post_->weight, + is_ln_post_weight_loaded_); + weight::load_weight( + state_dict, "ln_post.bias", ln_post_->bias, is_ln_post_bias_loaded_); + + const size_t layer_count = layers_->size(); + for (size_t index = 0; index < layer_count; ++index) { + const std::string prefix = "layers." + std::to_string(index) + "."; + layers_[index]->as()->load_state_dict( + state_dict.get_dict_with_prefix(prefix)); + } + } + + void verify_loaded_weights(const std::string& prefix) { + CHECK(is_conv_out_weight_loaded_) + << "weight is not loaded for " << prefix << "conv_out.weight"; + CHECK(is_proj1_weight_loaded_) + << "weight is not loaded for " << prefix << "proj1.weight"; + CHECK(is_proj1_bias_loaded_) + << "weight is not loaded for " << prefix << "proj1.bias"; + CHECK(is_proj2_weight_loaded_) + << "weight is not loaded for " << prefix << "proj2.weight"; + CHECK(is_proj2_bias_loaded_) + << "weight is not loaded for " << prefix << "proj2.bias"; + CHECK(is_conv2d1_weight_loaded_) + << "weight is not loaded for " << prefix << "conv2d1.weight"; + CHECK(is_conv2d1_bias_loaded_) + << "weight is not loaded for " << prefix << "conv2d1.bias"; + CHECK(is_conv2d2_weight_loaded_) + << "weight is not loaded for " << prefix << "conv2d2.weight"; + CHECK(is_conv2d2_bias_loaded_) + << "weight is not loaded for " << prefix << "conv2d2.bias"; + CHECK(is_conv2d3_weight_loaded_) + << "weight is not loaded for " << prefix << "conv2d3.weight"; + CHECK(is_conv2d3_bias_loaded_) + << "weight is not loaded for " << prefix << "conv2d3.bias"; + CHECK(is_ln_post_weight_loaded_) + << "weight is not loaded for " << prefix << "ln_post.weight"; + CHECK(is_ln_post_bias_loaded_) + << "weight is not loaded for " << prefix << "ln_post.bias"; + + const size_t layer_count = layers_->size(); + for (size_t index = 0; index < layer_count; ++index) { + layers_[index]->as()->verify_loaded_weights(); + } + } + + void merge_loaded_weights() { + const size_t layer_count = layers_->size(); + for (size_t index = 0; index < layer_count; ++index) { + layers_[index]->as()->merge_loaded_weights(); + } + } + + private: + int64_t embed_dim_ = 0; + int64_t n_window_ = 0; + int64_t n_window_infer_ = 0; + int64_t conv_chunk_size_ = 0; + torch::TensorOptions options_; + + Qwen3AudioSinusoidalPositionEmbedding positional_embedding_{nullptr}; + torch::nn::ModuleList layers_{nullptr}; + torch::nn::LayerNorm ln_post_{nullptr}; + torch::nn::Conv2d conv2d1_{nullptr}; + torch::nn::Conv2d conv2d2_{nullptr}; + torch::nn::Conv2d conv2d3_{nullptr}; + torch::nn::Linear conv_out_{nullptr}; + torch::nn::Linear proj1_{nullptr}; + torch::nn::Linear proj2_{nullptr}; + + bool is_conv_out_weight_loaded_ = false; + bool is_conv2d1_weight_loaded_ = false; + bool is_conv2d1_bias_loaded_ = false; + bool is_conv2d2_weight_loaded_ = false; + bool is_conv2d2_bias_loaded_ = false; + bool is_conv2d3_weight_loaded_ = false; + bool is_conv2d3_bias_loaded_ = false; + bool is_ln_post_weight_loaded_ = false; + bool is_ln_post_bias_loaded_ = false; + bool is_proj1_weight_loaded_ = false; + bool is_proj1_bias_loaded_ = false; + bool is_proj2_weight_loaded_ = false; + bool is_proj2_bias_loaded_ = false; +}; +TORCH_MODULE(Qwen3AudioEncoder); + +struct Qwen3AudioInputs { + torch::Tensor input_features; + torch::Tensor feature_lengths; + torch::Tensor feature_origin_lengths; +}; + +} // namespace xllm::npu::model diff --git a/xllm/models/vlm/npu/qwen3_omni_moe.h b/xllm/models/vlm/npu/qwen3_omni_moe.h new file mode 100644 index 0000000000..7af30199de --- /dev/null +++ b/xllm/models/vlm/npu/qwen3_omni_moe.h @@ -0,0 +1,87 @@ +/* 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 +#include +#include + +#include "models/vlm/npu/qwen3_omni_moe_thinker.h" +#include "processors/qwen3_omni_moe_processor.h" + +namespace xllm::npu::model { + +class Qwen3OmniMoeForConditionalGenerationImpl final + : public torch::nn::Module { + public: + explicit Qwen3OmniMoeForConditionalGenerationImpl( + const ModelContext& context) { + thinker_ = register_module( + "thinker", Qwen3OmniMoeThinkerForConditionalGeneration(context)); + } + + ModelOutput forward(const torch::Tensor& tokens, + const torch::Tensor& positions, + std::vector& kv_caches, + const ModelInputParams& input_params) { + torch::NoGradGuard no_grad; + return thinker_(tokens, positions, kv_caches, input_params); + } + + torch::Tensor logits(const torch::Tensor& hidden_states, + const torch::Tensor& selected_indices) { + return thinker_->logits(hidden_states, selected_indices); + } + + void load_model(std::unique_ptr loader) { + thinker_->load_model(std::move(loader)); + } + + torch::Tensor get_input_embeddings(const torch::Tensor input_ids, + const ModelInputParams& input_params) { + return thinker_->get_input_embeddings(input_ids, input_params); + } + + MMDict get_multimodal_embeddings(const ModelInputParams& input_params) { + return thinker_->get_multimodal_embeddings(input_params); + } + layer::NpuLmHead get_npu_lm_head() { return thinker_->get_npu_lm_head(); } + + void set_npu_lm_head(layer::NpuLmHead& head) { + thinker_->set_npu_lm_head(head); + } + + layer::NpuWordEmbedding get_npu_word_embedding() { + return thinker_->get_npu_word_embedding(); + } + + void set_npu_word_embedding(layer::NpuWordEmbedding& npu_word_embedding) { + thinker_->set_npu_word_embedding(npu_word_embedding); + } + + private: + Qwen3OmniMoeThinkerForConditionalGeneration thinker_{nullptr}; +}; +TORCH_MODULE(Qwen3OmniMoeForConditionalGeneration); + +REGISTER_MULTIMODAL_PROCESSOR(qwen3_omni_moe, Qwen3OmniMoeMultimodalProcessor); +REGISTER_CAUSAL_VLM_MODEL(qwen3_omni_moe, Qwen3OmniMoeForConditionalGeneration); +REGISTER_MPOSITION_GENERATOR(qwen3_omni_moe, xllm::Qwen3VLMPositionGenerator); + +REGISTER_MODEL_ARGS(qwen3_omni_moe, + [&] { load_qwen3_omni_moe_model_args(json, args); }); + +} // namespace xllm::npu::model diff --git a/xllm/models/vlm/npu/qwen3_omni_moe_thinker.h b/xllm/models/vlm/npu/qwen3_omni_moe_thinker.h new file mode 100644 index 0000000000..ef115fab43 --- /dev/null +++ b/xllm/models/vlm/npu/qwen3_omni_moe_thinker.h @@ -0,0 +1,1128 @@ +/* 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 +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "core/framework/config/model_config.h" +#include "core/framework/kv_cache/kv_cache.h" +#include "core/framework/model/model_input_params.h" +#include "core/framework/model_context.h" +#include "core/framework/state_dict/state_dict.h" +#include "core/layers/npu/npu_lm_head_impl.h" +#include "core/layers/npu/npu_qwen3_vision_encoder_layer_impl.h" +#include "models/llm/npu/qwen3_moe.h" +#include "models/model_registry.h" +#include "models/vlm/mposition/mposition.h" +#include "models/vlm/npu/qwen3_audio_encoder.h" +#include "models/vlm/npu/qwen3_vl.h" +#include "processors/qwen3_audio_common.h" +#include "processors/qwen3_omni_moe_processor.h" + +namespace xllm::npu::model { + +inline void load_qwen3_omni_moe_model_args(const JsonReader& json, + ModelArgs* args) { + LOAD_ARG_OR(model_type, "model_type", "qwen3_omni_moe"); + LOAD_ARG_OR(has_feature_extractor, "has_feature_extractor", true); + LOAD_ARG_OR(mm_audio_feature_size, "feature_size", 128); + LOAD_ARG_OR(mm_audio_sampling_rate, "sampling_rate", 16000); + LOAD_ARG_OR(mm_audio_hop_length, "hop_length", 160); + LOAD_ARG_OR(mm_audio_chunk_length, "chunk_length", 30); + LOAD_ARG_OR(mm_audio_n_fft, "n_fft", 400); + LOAD_ARG_OR(mm_audio_dither, "dither", 0.0); + LOAD_ARG_OR(mm_audio_truncation, "truncation", false); + LOAD_ARG_OR(mm_audio_do_normalize, "do_normalize", false); + SET_ARG(mm_use_audio_in_video, + ModelConfig::get_instance().use_audio_in_video()); + LOAD_ARG_OR(mm_position_id_per_seconds, "position_id_per_seconds", 13); + LOAD_ARG_OR(mm_fps, "fps", 1.0); + + LOAD_ARG_OR( + vision_start_token_id, "thinker_config.vision_start_token_id", 151652); + LOAD_ARG_OR( + vision_end_token_id, "thinker_config.vision_end_token_id", 151653); + LOAD_ARG_OR(vision_token_id, "thinker_config.vision_token_id", 151654); + LOAD_ARG_OR(image_token_id, "thinker_config.image_token_id", 151655); + LOAD_ARG_OR(video_token_id, "thinker_config.video_token_id", 151656); + LOAD_ARG_OR(audio_token_id, "thinker_config.audio_token_id", 151675); + LOAD_ARG_OR( + audio_start_token_id, "thinker_config.audio_start_token_id", 151669); + LOAD_ARG_OR(audio_end_token_id, "thinker_config.audio_end_token_id", 151670); + LOAD_ARG_OR(dtype, "thinker_config.dtype", "bfloat16"); + + LOAD_ARG_OR( + attention_bias, "thinker_config.text_config.attention_bias", false); + LOAD_ARG_OR( + attention_dropout, "thinker_config.text_config.attention_dropout", 0.0); + LOAD_ARG_OR( + decoder_sparse_step, "thinker_config.text_config.decoder_sparse_step", 1); + LOAD_ARG_OR(bos_token_id, "thinker_config.text_config.bos_token_id", 151643); + LOAD_ARG_OR(eos_token_id, "thinker_config.text_config.eos_token_id", 151645); + LOAD_ARG_OR(hidden_act, "thinker_config.text_config.hidden_act", "silu"); + LOAD_ARG_OR(hidden_size, "thinker_config.text_config.hidden_size", 2048); + LOAD_ARG_OR( + intermediate_size, "thinker_config.text_config.intermediate_size", 768); + LOAD_ARG_OR(max_position_embeddings, + "thinker_config.text_config.max_position_embeddings", + 65536); + LOAD_ARG_OR( + max_window_layers, "thinker_config.text_config.max_window_layers", 28); + LOAD_ARG_OR(n_heads, "thinker_config.text_config.num_attention_heads", 32); + LOAD_ARG_OR(n_layers, "thinker_config.text_config.num_hidden_layers", 48); + LOAD_ARG_OR(n_kv_heads, "thinker_config.text_config.num_key_value_heads", 4); + LOAD_ARG_OR(rms_norm_eps, "thinker_config.text_config.rms_norm_eps", 1e-06); + LOAD_ARG_OR( + sliding_window, "thinker_config.text_config.sliding_window", 32768); + LOAD_ARG_OR(tie_word_embeddings, + "thinker_config.text_config.tie_word_embeddings", + false); + LOAD_ARG_OR( + initializer_range, "thinker_config.text_config.initializer_range", 0.02); + LOAD_ARG_OR(use_sliding_window, + "thinker_config.text_config.use_sliding_window", + false); + LOAD_ARG_OR(moe_intermediate_size, + "thinker_config.text_config.moe_intermediate_size", + 768); + LOAD_ARG_OR( + norm_topk_prob, "thinker_config.text_config.norm_topk_prob", true); + LOAD_ARG_OR(num_experts, "thinker_config.text_config.num_experts", 128); + LOAD_ARG_OR( + num_experts_per_tok, "thinker_config.text_config.num_experts_per_tok", 8); + LOAD_ARG_OR_FUNC(head_dim, "thinker_config.text_config.head_dim", [args] { + return args->hidden_size() / args->n_heads(); + }); + LOAD_ARG_OR(output_router_logits, + "thinker_config.text_config.output_router_logits", + false); + LOAD_ARG_OR(router_aux_loss_coef, + "thinker_config.text_config.router_aux_loss_coef", + 0.001f); + LOAD_ARG_OR(mlp_only_layers, + "thinker_config.text_config.mlp_only_layers", + std::vector()); + LOAD_ARG_OR(rope_scaling_rope_type, + "thinker_config.text_config.rope_scaling.type", + "mrope"); + LOAD_ARG(rope_scaling_mrope_section, + "thinker_config.text_config.rope_scaling.mrope_section"); + LOAD_ARG_OR(rope_theta, "thinker_config.text_config.rope_theta", 1000000.0f); + LOAD_ARG_OR(vocab_size, "thinker_config.text_config.vocab_size", 152064); + + if (args->rope_scaling_rope_type() == "default") { + args->rope_scaling_rope_type() = "mrope"; + } + SET_ARG(stop_token_ids, std::unordered_set{args->eos_token_id()}); + + LOAD_ARG_OR(mm_num_hidden_layers, "thinker_config.vision_config.depth", 27); + LOAD_ARG_OR(mm_hidden_act, + "thinker_config.vision_config.hidden_act", + "gelu_pytorch_tanh"); + LOAD_ARG_OR(mm_hidden_size, "thinker_config.vision_config.hidden_size", 1152); + LOAD_ARG_OR(mm_intermediate_size, + "thinker_config.vision_config.intermediate_size", + 4304); + LOAD_ARG_OR( + mm_num_attention_heads, "thinker_config.vision_config.num_heads", 16); + LOAD_ARG_OR(mm_num_channels, "thinker_config.vision_config.in_channels", 3); + LOAD_ARG_OR( + mm_projection_dim, "thinker_config.vision_config.out_hidden_size", 2048); + LOAD_ARG_OR(mm_patch_size, "thinker_config.vision_config.patch_size", 16); + LOAD_ARG_OR(mm_num_position_embeddings, + "thinker_config.vision_config.num_position_embeddings", + 2304); + LOAD_ARG_OR(mm_spatial_merge_size, + "thinker_config.vision_config.spatial_merge_size", + 2); + LOAD_ARG(mm_deepstack_visual_indexes, + "thinker_config.vision_config.deepstack_visual_indexes"); + LOAD_ARG_OR(mm_temporal_patch_size, + "thinker_config.vision_config.temporal_patch_size", + 2); + LOAD_ARG_OR(mm_image_size, "thinker_config.vision_config.image_size", 768); + LOAD_ARG_OR_FUNC( + mm_head_dim, "thinker_config.vision_config.head_dim", [args] { + return args->mm_hidden_size() / args->mm_num_attention_heads(); + }); + + LOAD_ARG_OR(mm_audio_num_attention_heads, + "thinker_config.audio_config.encoder_attention_heads", + 20); + LOAD_ARG_OR( + mm_audio_hidden_size, "thinker_config.audio_config.d_model", 1280); + LOAD_ARG_OR(mm_audio_layer_norm_eps, + "thinker_config.audio_config.layer_norm_eps", + 1e-5); + LOAD_ARG_OR(mm_audio_downsample_hidden_size, + "thinker_config.audio_config.downsample_hidden_size", + 480); + LOAD_ARG_OR( + mm_audio_num_mel_bins, "thinker_config.audio_config.num_mel_bins", 128); + LOAD_ARG_OR(mm_audio_max_source_positions, + "thinker_config.audio_config.max_source_positions", + 1500); + LOAD_ARG_OR(mm_audio_n_window, "thinker_config.audio_config.n_window", 50); + LOAD_ARG_OR(mm_audio_n_window_infer, + "thinker_config.audio_config.n_window_infer", + 800); + LOAD_ARG_OR(mm_audio_conv_chunksize, + "thinker_config.audio_config.conv_chunksize", + 500); + LOAD_ARG_OR(mm_audio_encoder_layers, + "thinker_config.audio_config.encoder_layers", + 32); + LOAD_ARG_OR( + mm_audio_output_dim, "thinker_config.audio_config.output_dim", 2048); +} + +class Qwen3OmniMoeThinkerVisionPatchEmbedImpl final : public torch::nn::Module { + public: + explicit Qwen3OmniMoeThinkerVisionPatchEmbedImpl( + const ModelContext& context) { + const ModelArgs& model_args = context.get_model_args(); + const torch::TensorOptions options = context.get_tensor_options(); + + const int64_t in_features = + model_args.mm_num_channels() * model_args.mm_temporal_patch_size() * + model_args.mm_patch_size() * model_args.mm_patch_size(); + + const int64_t out_features = model_args.mm_hidden_size(); + + proj_ = register_module( + "proj", + torch::nn::Linear( + torch::nn::LinearOptions(in_features, out_features).bias(true))); + + proj_->weight.set_data(proj_->weight.to(options)); + proj_->bias.set_data(proj_->bias.to(options)); + } + + torch::Tensor forward(torch::Tensor hidden_states) { + return proj_(hidden_states); + } + + void load_state_dict(const StateDict& state_dict) { + auto weight = state_dict.get_tensor("proj.weight"); + if (weight.defined()) { + weight = weight.reshape({weight.size(0), -1}); + DCHECK_EQ(proj_->weight.sizes(), weight.sizes()) + << "proj weight size mismatch for " << name(); + proj_->weight.data().copy_(weight); + proj_weight_loaded_ = true; + } + auto bias = state_dict.get_tensor("proj.bias"); + if (bias.defined()) { + bias = bias.reshape({bias.size(0)}); + DCHECK_EQ(proj_->bias.sizes(), bias.sizes()) + << "proj bias size mismatch for " << name(); + proj_->bias.data().copy_(bias); + proj_bias_loaded_ = true; + } + } + + void verify_loaded_weights(const std::string& prefix) const { + CHECK(proj_weight_loaded_) + << "weight is not loaded for " << prefix + "proj.weight"; + CHECK(proj_bias_loaded_) + << "bias is not loaded for " << prefix + "proj.bias"; + } + + private: + bool proj_weight_loaded_ = false; + bool proj_bias_loaded_ = false; + torch::nn::Linear proj_{nullptr}; +}; +TORCH_MODULE(Qwen3OmniMoeThinkerVisionPatchEmbed); + +class Qwen3OmniMoeThinkerVisionBlockImpl final : public torch::nn::Module { + public: + explicit Qwen3OmniMoeThinkerVisionBlockImpl(const ModelContext& context) { + encoder_layer_ = register_module( + "encoder_layer", layer::NpuQwen3VisionEncoderLayer(context)); + } + + torch::Tensor forward(torch::Tensor& hidden_states, + torch::Tensor& m_cos_pos, + torch::Tensor& m_sin_pos, + torch::Tensor& cu_seq_len, + std::vector& cu_seq_len_vec, + int32_t node_id) { + return encoder_layer_(hidden_states, + m_cos_pos, + m_sin_pos, + cu_seq_len, + cu_seq_len_vec, + node_id); + } + + void load_state_dict(const StateDict& state_dict) { + encoder_layer_->load_state_dict(state_dict); + } + + void verify_loaded_weights(const std::string& prefix) const { + encoder_layer_->verify_loaded_weights(); + } + void merge_loaded_weights() { encoder_layer_->merge_loaded_weights(); } + + private: + layer::NpuQwen3VisionEncoderLayer encoder_layer_{nullptr}; +}; +TORCH_MODULE(Qwen3OmniMoeThinkerVisionBlock); + +class Qwen3OmniMoeThinkerVisionRotaryEmbeddingImpl final + : public torch::nn::Module { + public: + explicit Qwen3OmniMoeThinkerVisionRotaryEmbeddingImpl( + const ModelContext& context) { + const ModelArgs& model_args = context.get_model_args(); + const torch::TensorOptions options = context.get_tensor_options(); + + dim_ = model_args.mm_head_dim() / 2; + theta_ = 10000.0; + + const torch::TensorOptions float_options = options.dtype(torch::kFloat32); + const torch::Tensor inv_freq = + 1.0 / + torch::pow(theta_, torch::arange(0, dim_, 2, float_options) / dim_); + inv_freq_ = register_buffer("inv_freq", inv_freq); + } + + void update_freqs_cache(int64_t seqlen) { + if (seqlen <= seq_len_cached_) { + return; + } + + seqlen *= 2; + seq_len_cached_ = seqlen; + + const torch::TensorOptions options = torch::TensorOptions() + .dtype(torch::kFloat32) + .device(inv_freq_.device()); + inv_freq_ = + 1.0 / torch::pow(theta_, torch::arange(0, dim_, 2, options) / dim_); + auto seq = torch::arange(seqlen, options); + freqs_cached_ = torch::outer(seq, inv_freq_); + } + + torch::Tensor forward(int64_t seqlen) { + update_freqs_cache(seqlen); + return freqs_cached_.slice(0, 0, seqlen); + } + + private: + int64_t dim_ = 0; + double theta_ = 0.0; + + int64_t seq_len_cached_ = 0; + torch::Tensor inv_freq_; + torch::Tensor freqs_cached_; +}; +TORCH_MODULE(Qwen3OmniMoeThinkerVisionRotaryEmbedding); + +class Qwen3OmniMoeThinkerVisionPatchMergerImpl final + : public torch::nn::Module { + public: + explicit Qwen3OmniMoeThinkerVisionPatchMergerImpl( + const ModelContext& context, + bool use_postshuffle_norm = false) { + const ModelArgs& model_args = context.get_model_args(); + const torch::TensorOptions options = context.get_tensor_options(); + const int64_t output_size = model_args.mm_projection_dim(); + const int64_t context_size = model_args.mm_hidden_size(); + const int64_t spatial_merge_size = model_args.mm_spatial_merge_size(); + hidden_size_ = context_size * spatial_merge_size * spatial_merge_size; + use_postshuffle_norm_ = use_postshuffle_norm; + if (use_postshuffle_norm) { + norm_ = register_module( + "norm", + torch::nn::LayerNorm(torch::nn::LayerNormOptions({hidden_size_}) + .elementwise_affine(true) + .eps(1e-6))); + } else { + norm_ = register_module( + "norm", + torch::nn::LayerNorm(torch::nn::LayerNormOptions({context_size}) + .elementwise_affine(true) + .eps(1e-6))); + } + + norm_->weight.set_data(norm_->weight.to(options)); + norm_->bias.set_data(norm_->bias.to(options)); + + auto fc1 = torch::nn::Linear( + torch::nn::LinearOptions(hidden_size_, hidden_size_).bias(true)); + fc1->weight.set_data(fc1->weight.to(options)); + fc1->bias.set_data(fc1->bias.to(options)); + auto act = torch::nn::GELU(); + auto fc2 = torch::nn::Linear( + torch::nn::LinearOptions(hidden_size_, output_size).bias(true)); + fc2->weight.set_data(fc2->weight.to(options)); + fc2->bias.set_data(fc2->bias.to(options)); + mlp_ = register_module("mlp", torch::nn::Sequential(fc1, act, fc2)); + layers_ = std::make_tuple(fc1, act, fc2); + } + + torch::Tensor forward(torch::Tensor hidden_states) { + if (use_postshuffle_norm_) { + hidden_states = norm_(hidden_states.view({-1, hidden_size_})); + } else { + hidden_states = norm_(hidden_states).view({-1, hidden_size_}); + } + return mlp_->forward(hidden_states); + } + + void load_state_dict(const StateDict& state_dict) { + // norm + const auto& norm_dict = state_dict.get_dict_with_prefix("ln_q."); + const auto& norm_weight = norm_dict.get_tensor("weight"); + if (norm_weight.defined()) { + CHECK_EQ(norm_->weight.sizes(), norm_weight.sizes()) + << "weight size mismatch for " << name(); + norm_->weight.data().copy_(norm_weight); + is_norm_weight_loaded_ = true; + } + const auto norm_bias = norm_dict.get_tensor("bias"); + if (norm_bias.defined()) { + CHECK_EQ(norm_->bias.sizes(), norm_bias.sizes()) + << "bias size mismatch for " << name(); + norm_->bias.data().copy_(norm_bias); + is_norm_bias_loaded_ = true; + } + + const auto& fc1_dict = state_dict.get_dict_with_prefix("mlp.0."); + const auto& fc1_weight = fc1_dict.get_tensor("weight"); + if (fc1_weight.defined()) { + CHECK_EQ(std::get<0>(layers_)->weight.sizes(), fc1_weight.sizes()) + << "weight size mismatch for " << name(); + std::get<0>(layers_)->weight.data().copy_(fc1_weight); + is_fc1_weight_loaded_ = true; + } + const auto fc1_bias = fc1_dict.get_tensor("bias"); + if (fc1_bias.defined()) { + CHECK_EQ(std::get<0>(layers_)->bias.sizes(), fc1_bias.sizes()) + << "bias size mismatch for " << name(); + std::get<0>(layers_)->bias.data().copy_(fc1_bias); + is_fc1_bias_loaded_ = true; + } + + const auto& fc2_dict = state_dict.get_dict_with_prefix("mlp.2."); + const auto& fc2_weight = fc2_dict.get_tensor("weight"); + if (fc2_weight.defined()) { + CHECK_EQ(std::get<2>(layers_)->weight.sizes(), fc2_weight.sizes()) + << "weight size mismatch for " << name(); + std::get<2>(layers_)->weight.data().copy_(fc2_weight); + is_fc2_weight_loaded_ = true; + } + const auto fc2_bias = fc2_dict.get_tensor("bias"); + if (fc2_bias.defined()) { + CHECK_EQ(std::get<2>(layers_)->bias.sizes(), fc2_bias.sizes()) + << "bias size mismatch for " << name(); + std::get<2>(layers_)->bias.data().copy_(fc2_bias); + is_fc2_bias_loaded_ = true; + } + } + + void verify_loaded_weights(const std::string& prefix) const { + CHECK(is_fc1_weight_loaded_) + << "weight is not loaded for " << prefix + "mlp.0.weight"; + CHECK(is_fc1_bias_loaded_) + << "bias is not loaded for " << prefix + "mlp.0.bias"; + CHECK(is_fc2_weight_loaded_) + << "weight is not loaded for " << prefix + "mlp.2.weight"; + CHECK(is_fc2_bias_loaded_) + << "bias is not loaded for " << prefix + "mlp.2.bias"; + CHECK(is_norm_weight_loaded_) + << "weight is not loaded for " << prefix + "ln_q.weight"; + CHECK(is_norm_bias_loaded_) + << "bias is not loaded for " << prefix + "ln_q.bias"; + } + + private: + int64_t hidden_size_ = 0; + bool use_postshuffle_norm_ = false; + torch::nn::LayerNorm norm_{nullptr}; + torch::nn::Sequential mlp_{nullptr}; + std::tuple layers_ = { + nullptr, + nullptr, + nullptr}; + bool is_fc1_weight_loaded_ = false; + bool is_fc1_bias_loaded_ = false; + bool is_fc2_weight_loaded_ = false; + bool is_fc2_bias_loaded_ = false; + bool is_norm_weight_loaded_ = false; + bool is_norm_bias_loaded_ = false; +}; +TORCH_MODULE(Qwen3OmniMoeThinkerVisionPatchMerger); + +class Qwen3OmniMoeThinkerVisionTransformerImpl final + : public torch::nn::Module { + public: + explicit Qwen3OmniMoeThinkerVisionTransformerImpl(const ModelContext& context) + : options_(context.get_tensor_options()) { + const ModelArgs& model_args = context.get_model_args(); + hidden_size_ = model_args.mm_hidden_size(); + patch_size_ = model_args.mm_patch_size(); + spatial_merge_size_ = model_args.mm_spatial_merge_size(); + deepstack_visual_indexes_ = model_args.mm_deepstack_visual_indexes(); + image_size_ = model_args.mm_image_size(); + num_grid_per_side_ = image_size_ / patch_size_; + + patch_embed_ = register_module( + "patch_embed", Qwen3OmniMoeThinkerVisionPatchEmbed(context)); + rotary_pos_emb_ = register_module( + "rotary_pos_emb", Qwen3OmniMoeThinkerVisionRotaryEmbedding(context)); + + blocks_ = register_module("blocks", torch::nn::ModuleList()); + deepstack_mergers_ = + register_module("deepstack_mergers", torch::nn::ModuleList()); + + emb_ = register_module( + "embedding", + torch::nn::Embedding(num_grid_per_side_ * num_grid_per_side_, + hidden_size_)); + emb_->weight.set_data(emb_->weight.to(options_)); + + merger_ = register_module("merger", + Qwen3OmniMoeThinkerVisionPatchMerger(context)); + + layers_.reserve(static_cast(model_args.mm_num_hidden_layers())); + for (int32_t index = 0; index < model_args.mm_num_hidden_layers(); + ++index) { + Qwen3OmniMoeThinkerVisionBlock block(context); + blocks_->push_back(block); + layers_.emplace_back(block); + } + const size_t deepstack_count = deepstack_visual_indexes_.size(); + deepstack_merger_layers_.reserve(deepstack_count); + for (size_t index = 0; index < deepstack_count; ++index) { + Qwen3OmniMoeThinkerVisionPatchMerger merger( + context, /*use_postshuffle_norm=*/true); + deepstack_mergers_->push_back(merger); + deepstack_merger_layers_.emplace_back(merger); + } + } + + torch::Tensor rot_pos_emb(torch::Tensor grid_thw) { + std::vector pos_ids_vec; + const int64_t count = grid_thw.size(0); + pos_ids_vec.reserve(static_cast(count)); + + const torch::Tensor grid_thw_cpu = grid_thw.cpu(); + const torch::TensorOptions options = + torch::TensorOptions().dtype(torch::kLong).device(grid_thw.device()); + + for (int64_t index = 0; index < count; ++index) { + const int64_t temporal = grid_thw_cpu[index][0].item(); + const int64_t height = grid_thw_cpu[index][1].item(); + const int64_t width = grid_thw_cpu[index][2].item(); + + torch::Tensor height_position_ids = + torch::arange(height, options).unsqueeze(1).expand({-1, width}); + height_position_ids = height_position_ids + .reshape({height / spatial_merge_size_, + spatial_merge_size_, + width / spatial_merge_size_, + spatial_merge_size_}) + .permute({0, 2, 1, 3}) + .flatten(); + + torch::Tensor width_position_ids = + torch::arange(width, options).unsqueeze(0).expand({height, -1}); + width_position_ids = width_position_ids + .reshape({height / spatial_merge_size_, + spatial_merge_size_, + width / spatial_merge_size_, + spatial_merge_size_}) + .permute({0, 2, 1, 3}) + .flatten(); + + pos_ids_vec.emplace_back( + torch::stack({height_position_ids, width_position_ids}, /*dim=*/-1) + .repeat({temporal, 1})); + } + + const torch::Tensor position_ids = torch::cat(pos_ids_vec, /*dim=*/0); + const torch::Tensor max_grid_size = + grid_thw + .index({torch::indexing::Slice(), + torch::indexing::Slice(1, torch::indexing::None)}) + .max(); + + const torch::Tensor rotary_position_embedding = + rotary_pos_emb_(max_grid_size.item()) + .index({position_ids}) + .flatten(1); + + return rotary_position_embedding; + } + + torch::Tensor fast_pos_embed_interpolate(const torch::Tensor& grid_thw) { + const torch::Device device = grid_thw.device(); + const int64_t hidden_dim = hidden_size_; + const int64_t merge_size = spatial_merge_size_; + + const torch::Tensor grid_cpu = grid_thw.to(torch::kCPU); + const int64_t count = grid_thw.size(0); + + std::vector outputs; + outputs.reserve(static_cast(count)); + + for (int64_t index = 0; index < count; ++index) { + const int64_t temporal = grid_cpu[index][0].item(); + const int64_t height = grid_cpu[index][1].item(); + const int64_t width = grid_cpu[index][2].item(); + + auto h_idxs = torch::linspace(0, + static_cast(num_grid_per_side_ - 1), + height, + torch::kFloat32) + .to(device); + auto w_idxs = torch::linspace(0, + static_cast(num_grid_per_side_ - 1), + width, + torch::kFloat32) + .to(device); + + auto h_floor = h_idxs.to(torch::kLong); + auto w_floor = w_idxs.to(torch::kLong); + auto h_ceil = torch::clamp(h_floor + 1, 0, num_grid_per_side_ - 1); + auto w_ceil = torch::clamp(w_floor + 1, 0, num_grid_per_side_ - 1); + + auto dh = h_idxs - h_floor; + auto dw = w_idxs - w_floor; + + auto mesh_d = torch::meshgrid({dh, dw}, "ij"); + auto dh_grid = mesh_d[0], dw_grid = mesh_d[1]; + + auto mesh_floor = torch::meshgrid({h_floor, w_floor}, "ij"); + auto h_floor_grid = mesh_floor[0]; + auto w_floor_grid = mesh_floor[1]; + + auto mesh_ceil = torch::meshgrid({h_ceil, w_ceil}, "ij"); + auto h_ceil_grid = mesh_ceil[0]; + auto w_ceil_grid = mesh_ceil[1]; + + auto h_floor_grid_idx = h_floor_grid * num_grid_per_side_; + auto h_ceil_grid_idx = h_ceil_grid * num_grid_per_side_; + + auto w11 = dh_grid * dw_grid; + auto w10 = dh_grid - w11; + auto w01 = dw_grid - w11; + auto w00 = 1.0f - dh_grid - dw_grid + w11; + + auto idx00 = h_floor_grid_idx + w_floor_grid; + auto idx01 = h_floor_grid_idx + w_ceil_grid; + auto idx10 = h_ceil_grid_idx + w_floor_grid; + auto idx11 = h_ceil_grid_idx + w_ceil_grid; + + auto indices = torch::stack({idx00, idx01, idx10, idx11}, 0) + .reshape({4, -1}) + .to(torch::kLong); + auto weights = torch::stack({w00, w01, w10, w11}, 0) + .reshape({4, -1, 1}) + .to(options_); + + auto embeds = emb_(indices); + + const torch::Tensor combined = (embeds * weights).sum(/*dim=*/0); + + auto repeated = + combined.unsqueeze(0).expand({temporal, -1, -1}).contiguous(); + repeated = repeated.view({temporal, + height / merge_size, + merge_size, + width / merge_size, + merge_size, + hidden_dim}); + repeated = repeated.permute({0, 1, 3, 2, 4, 5}).reshape({-1, hidden_dim}); + + outputs.emplace_back(repeated); + } + + return torch::cat(outputs, 0); + } + + std::tuple> forward( + torch::Tensor hidden_states, + torch::Tensor grid_thw) { + hidden_states = patch_embed_(hidden_states); + const torch::Tensor pos_embeds = fast_pos_embed_interpolate(grid_thw); + hidden_states = hidden_states + pos_embeds; + const torch::Tensor rotary_pos_emb = rot_pos_emb(grid_thw); + torch::Tensor cu_seqlens = + torch::repeat_interleave( + grid_thw.index({torch::indexing::Slice(), 1}) * + grid_thw.index({torch::indexing::Slice(), 2}), + grid_thw.index({torch::indexing::Slice(), 0})) + .cumsum(/*dim=*/0, torch::kInt32); + cu_seqlens = + torch::nn::functional::pad(cu_seqlens, + torch::nn::functional::PadFuncOptions({1, 0}) + .mode(torch::kConstant) + .value(0)); + cu_seqlens = torch::diff(cu_seqlens); + + m_cos_ = rotary_pos_emb.cos().type_as(hidden_states); + m_cos_ = m_cos_.repeat({1, 2}); + m_sin_ = rotary_pos_emb.sin().type_as(hidden_states); + m_sin_ = m_sin_.repeat({1, 2}); + + torch::Tensor cu_seqlens_cpu = cu_seqlens.cpu(); + std::vector cu_seqlens_vec( + cu_seqlens_cpu.data_ptr(), + cu_seqlens_cpu.data_ptr() + cu_seqlens_cpu.numel()); + std::vector deepstack_feature_lists; + deepstack_feature_lists.reserve(deepstack_visual_indexes_.size()); + const int32_t layer_count = static_cast(blocks_->size()); + for (int32_t index = 0; index < layer_count; ++index) { + hidden_states = layers_[index]( + hidden_states, m_cos_, m_sin_, cu_seqlens, cu_seqlens_vec, index); + auto it = std::find(deepstack_visual_indexes_.begin(), + deepstack_visual_indexes_.end(), + index); + + if (it != deepstack_visual_indexes_.end()) { + const size_t merger_index = static_cast( + std::distance(deepstack_visual_indexes_.begin(), it)); + deepstack_feature_lists.emplace_back( + deepstack_merger_layers_[merger_index](hidden_states)); + } + } + hidden_states = merger_(hidden_states); + return std::make_tuple(hidden_states, deepstack_feature_lists); + } + + void load_state_dict(const StateDict& state_dict) { + patch_embed_->load_state_dict( + state_dict.get_dict_with_prefix("patch_embed.")); + const size_t layer_count = layers_.size(); + for (size_t index = 0; index < layer_count; ++index) { + layers_[index]->load_state_dict(state_dict.get_dict_with_prefix( + "blocks." + std::to_string(index) + ".")); + } + + merger_->load_state_dict(state_dict.get_dict_with_prefix("merger.")); + + const size_t merger_count = deepstack_merger_layers_.size(); + for (size_t index = 0; index < merger_count; ++index) { + deepstack_merger_layers_[index]->load_state_dict( + state_dict.get_dict_with_prefix("merger_list." + + std::to_string(index) + ".")); + } + + const auto& emb_dict = state_dict.get_dict_with_prefix("pos_embed."); + const auto& emb_weight = emb_dict.get_tensor("weight"); + if (emb_weight.defined()) { + CHECK_EQ(emb_->weight.sizes(), emb_weight.sizes()) + << "weight size mismatch for " << name(); + emb_->weight.data().copy_(emb_weight); + is_emb_weight_loaded_ = true; + } + } + + void verify_loaded_weights(const std::string& prefix) const { + patch_embed_->verify_loaded_weights(prefix + "patch_embed."); + const size_t layer_count = layers_.size(); + for (size_t index = 0; index < layer_count; ++index) { + layers_[index]->verify_loaded_weights(prefix + "blocks." + + std::to_string(index) + "."); + } + merger_->verify_loaded_weights(prefix + "merger."); + + const size_t merger_count = deepstack_merger_layers_.size(); + for (size_t index = 0; index < merger_count; ++index) { + deepstack_merger_layers_[index]->verify_loaded_weights( + prefix + "merger_list." + std::to_string(index) + "."); + } + CHECK(is_emb_weight_loaded_) + << "weight is not loaded for " << prefix + "pos_embed.weight"; + } + + void merge_loaded_weights() { + for (Qwen3OmniMoeThinkerVisionBlock& layer : layers_) { + layer->merge_loaded_weights(); + } + } + + private: + int64_t hidden_size_ = 0; + int64_t patch_size_ = 0; + int64_t spatial_merge_size_ = 0; + std::vector deepstack_visual_indexes_; + int64_t image_size_ = 0; + int64_t num_grid_per_side_ = 0; + + Qwen3OmniMoeThinkerVisionPatchEmbed patch_embed_{nullptr}; + Qwen3OmniMoeThinkerVisionRotaryEmbedding rotary_pos_emb_{nullptr}; + torch::nn::Embedding emb_{nullptr}; + + torch::nn::ModuleList blocks_{nullptr}; + std::vector layers_; + + torch::nn::ModuleList deepstack_mergers_{nullptr}; + std::vector deepstack_merger_layers_; + Qwen3OmniMoeThinkerVisionPatchMerger merger_{nullptr}; + + torch::Tensor m_cos_; + torch::Tensor m_sin_; + bool is_emb_weight_loaded_ = false; + torch::TensorOptions options_; +}; +TORCH_MODULE(Qwen3OmniMoeThinkerVisionTransformer); + +class Qwen3OmniMoeThinkerForConditionalGenerationImpl final + : public torch::nn::Module { + public: + explicit Qwen3OmniMoeThinkerForConditionalGenerationImpl( + const ModelContext& context) + : model_args_(context.get_model_args()), + options_(context.get_tensor_options()) { + visual_ = register_module("visual", + Qwen3OmniMoeThinkerVisionTransformer(context)); + audio_tower_ = register_module("audio_tower", Qwen3AudioEncoder(context)); + language_model_ = + register_module("language_model", Qwen3MoeForCausalLM(context)); + } + + void prepare_encoder_input(const ModelInputParams& input_params, + std::optional& image_inputs, + std::optional& video_inputs, + std::optional& audio_inputs) { + const auto& mm_data = input_params.multimodal.mm_data; + torch::Tensor pixel_values; + if (std::optional value = + mm_data.get("pixel_values")) { + pixel_values = value.value(); + } + + torch::Tensor image_grid_thw; + if (std::optional value = + mm_data.get("image_grid_thw")) { + image_grid_thw = value.value(); + } + + torch::Tensor pixel_values_videos; + if (std::optional value = + mm_data.get("pixel_values_videos")) { + pixel_values_videos = value.value(); + } + + torch::Tensor video_grid_thw; + if (std::optional value = + mm_data.get("video_grid_thw")) { + video_grid_thw = value.value(); + } + + torch::Tensor input_features; + if (std::optional res = + mm_data.get(qwen3_audio::kInputFeaturesKey)) { + input_features = res.value(); + } + + torch::Tensor feature_lengths; + if (std::optional res = + mm_data.get(qwen3_audio::kFeatureLengthKey)) { + feature_lengths = res.value(); + } + + torch::Tensor feature_origin_lengths; + if (std::optional res = + mm_data.get(qwen3_audio::kFeatureOriginLengthsKey)) { + feature_origin_lengths = res.value(); + } + + if (pixel_values.defined() && image_grid_thw.defined()) { + image_inputs = Qwen3_VLImageInputs{pixel_values, image_grid_thw}; + } + + if (pixel_values_videos.defined() && video_grid_thw.defined()) { + video_inputs = Qwen3_VLVideoInputs{pixel_values_videos, video_grid_thw}; + } + + if (input_features.defined() && feature_lengths.defined() && + feature_origin_lengths.defined()) { + audio_inputs = Qwen3AudioInputs{ + input_features, feature_lengths, feature_origin_lengths}; + } + } + + MMDict get_multimodal_embeddings(const ModelInputParams& input_params) { + std::optional image_input; + std::optional video_input; + std::optional audio_input; + prepare_encoder_input(input_params, image_input, video_input, audio_input); + MMDict multimodal_embeds; + const int64_t merge_size = model_args_.mm_image_merge_size(); + if (image_input) { + auto [image_embeds, deep_stacks] = + visual_(image_input->pixel_values.to(options_), + image_input->image_grid_thw.to(options_.device())); + + auto image_tokens = + (image_input->image_grid_thw.prod(-1) / merge_size / merge_size) + .cpu() + .contiguous() + .to(torch::kLong); + std::vector image_tokens_vec( + image_tokens.data_ptr(), + image_tokens.data_ptr() + image_tokens.numel()); + std::vector image_embedding_parts{image_embeds}; + image_embedding_parts.insert( + image_embedding_parts.end(), deep_stacks.begin(), deep_stacks.end()); + multimodal_embeds[get_embedding_key(MMType::IMAGE)] = + torch::cat(image_embedding_parts, /*dim=*/1) + .split(image_tokens_vec, /*dim=*/0); + } + if (video_input) { + auto [video_embeds, deep_stacks] = + visual_(video_input->pixel_values_videos.to(options_), + video_input->video_grid_thw.to(options_.device())); + auto video_tokens = + (video_input->video_grid_thw.prod(-1) / merge_size / merge_size) + .cpu() + .contiguous() + .to(torch::kLong); + std::vector video_tokens_vec( + video_tokens.data_ptr(), + video_tokens.data_ptr() + video_tokens.numel()); + std::vector video_embedding_parts{video_embeds}; + video_embedding_parts.insert( + video_embedding_parts.end(), deep_stacks.begin(), deep_stacks.end()); + multimodal_embeds[get_embedding_key(MMType::VIDEO)] = + torch::cat(video_embedding_parts, /*dim=*/1) + .split(video_tokens_vec, /*dim=*/0); + } + if (audio_input) { + const torch::Tensor feature_origin_lengths = + audio_input->feature_origin_lengths.to(options_.device(), + torch::kLong); + + const torch::Tensor input_features = + audio_input->input_features.permute({1, 0}).to(options_); + + const torch::Tensor audio_embeds = + audio_tower_->forward(input_features, feature_origin_lengths); + + const torch::Tensor audio_tokens = + audio_input->feature_lengths.cpu().contiguous().to(torch::kLong); + + std::vector feature_lens_vec( + audio_tokens.data_ptr(), + audio_tokens.data_ptr() + audio_tokens.numel()); + + multimodal_embeds[get_embedding_key(MMType::AUDIO)] = + audio_embeds.split(feature_lens_vec, /*dim=*/0); + } + if (model_args_.mm_use_audio_in_video() && video_input && audio_input) { + const std::vector origin_audio_embeds = + std::get>( + multimodal_embeds[get_embedding_key(MMType::AUDIO)]); + const std::vector origin_video_embeds = + std::get>( + multimodal_embeds[get_embedding_key(MMType::VIDEO)]); + CHECK_GE(origin_audio_embeds.size(), origin_video_embeds.size()); + CHECK(!origin_video_embeds.empty()); + std::vector audio_in_video_embeds; + audio_in_video_embeds.reserve(origin_video_embeds.size()); + std::vector scattered_audio_embeds; + scattered_audio_embeds.reserve(origin_audio_embeds.size()); + size_t audio_index = 0; + size_t video_index = 0; + + const auto& mm_data = input_params.multimodal.mm_data; + const std::vector& mm_data_vec = mm_data.mm_data_vec(); + for (const MMData& sequence_mm_data : mm_data_vec) { + const MMItemVec& mm_items = sequence_mm_data.items(); + for (size_t item_index = 0; item_index < mm_items.size(); + ++item_index) { + const MMDataItem& item = mm_items[item_index]; + if (item.is_type(MMType::AUDIO)) { + CHECK_LT(audio_index, origin_audio_embeds.size()); + scattered_audio_embeds.emplace_back( + origin_audio_embeds[audio_index++]); + } else if (item.is_type(MMType::VIDEO)) { + CHECK(item_index + 1 < mm_items.size()); + CHECK(mm_items[item_index + 1].is_type(MMType::AUDIO)); + CHECK_LT(video_index, origin_video_embeds.size()); + CHECK_LT(audio_index, origin_audio_embeds.size()); + const torch::Tensor video_embedding = + origin_video_embeds[video_index++]; + const torch::Tensor audio_embedding = + origin_audio_embeds[audio_index++]; + std::optional audio_in_video_token_ids = + item.get( + qwen3_omni_moe::kAudioInVideoTokenIdsKey); + CHECK(audio_in_video_token_ids.has_value()); + torch::Tensor audio_in_video_embedding = + torch::full({audio_in_video_token_ids->size(0), + origin_video_embeds[0].size(1)}, + 1, + options_); + const torch::Tensor video_mask = torch::isin( + audio_in_video_token_ids.value(), model_args_.video_token_id()); + const torch::Tensor audio_mask = torch::isin( + audio_in_video_token_ids.value(), model_args_.audio_token_id()); + audio_in_video_embedding.index_put_({video_mask}, video_embedding); + audio_in_video_embedding.index_put_({audio_mask}, audio_embedding); + audio_in_video_embeds.emplace_back(audio_in_video_embedding); + scattered_audio_embeds.emplace_back( + audio_embedding.slice(/*dim=*/0, /*start=*/0, /*end=*/0)); + ++item_index; + } + } + } + CHECK_EQ(audio_index, origin_audio_embeds.size()); + CHECK_EQ(video_index, origin_video_embeds.size()); + CHECK_EQ(scattered_audio_embeds.size(), origin_audio_embeds.size()); + multimodal_embeds[get_embedding_key(MMType::AUDIO)] = + scattered_audio_embeds; + multimodal_embeds[get_embedding_key(MMType::VIDEO)] = + audio_in_video_embeds; + } + return multimodal_embeds; + } + + torch::Tensor merge_multimodal_embeddings( + torch::Tensor inputs_embeds, + const torch::Tensor& multimodal_embeds, + const torch::Tensor& is_multimodal) { + inputs_embeds.index_put_({is_multimodal}, multimodal_embeds); + return inputs_embeds; + } + + torch::Tensor get_input_embeddings(const torch::Tensor input_ids, + const ModelInputParams& input_params) { + const auto& mm_data = input_params.multimodal.mm_data; + torch::Tensor inputs_embeds = + language_model_->get_input_embeddings(input_ids); + const size_t num_deepstacks = + model_args_.mm_deepstack_visual_indexes().size(); + std::vector deepstack_input_embeds( + num_deepstacks, torch::zeros_like(inputs_embeds)); + auto merge_visual_modality = [&](const std::string& embed_key, + const std::string& mask_key) { + std::optional embedding = + mm_data.get(embed_key); + std::optional mask = mm_data.get(mask_key); + if (!embedding.has_value() || !mask.has_value()) { + return; + } + const std::vector chunks = embedding.value().chunk( + static_cast(num_deepstacks + 1), /*dim=*/1); + inputs_embeds = + merge_multimodal_embeddings(inputs_embeds, chunks[0], mask.value()); + for (size_t index = 0; index < num_deepstacks; ++index) { + deepstack_input_embeds[index] = merge_multimodal_embeddings( + deepstack_input_embeds[index], chunks[index + 1], mask.value()); + } + }; + merge_visual_modality(get_embedding_key(MMType::IMAGE), "image|mask"); + merge_visual_modality(get_embedding_key(MMType::VIDEO), "video|mask"); + std::optional audio_embeds = + mm_data.get(get_embedding_key(MMType::AUDIO)); + std::optional audio_mask = + mm_data.get(qwen3_audio::kMaskKey); + if (audio_embeds.has_value() && audio_mask.has_value()) { + inputs_embeds = merge_multimodal_embeddings( + inputs_embeds, audio_embeds.value(), audio_mask.value()); + } + input_params.multimodal.deep_stacks = std::move(deepstack_input_embeds); + return inputs_embeds; + } + + ModelOutput forward(const torch::Tensor& tokens, + const torch::Tensor& positions, + std::vector& kv_caches, + const ModelInputParams& input_params) { + return language_model_(tokens, positions, kv_caches, input_params); + } + + torch::Tensor logits(const torch::Tensor& hidden_states, + const torch::Tensor& selected_indices) { + return language_model_->logits(hidden_states, selected_indices); + } + + void load_model(std::unique_ptr loader) { + for (const auto& state_dict : loader->get_state_dicts()) { + visual_->load_state_dict( + state_dict->get_dict_with_prefix("thinker.visual.")); + audio_tower_->load_state_dict( + state_dict->get_dict_with_prefix("thinker.audio_tower.")); + } + // verify + visual_->verify_loaded_weights("thinker.visual."); + visual_->merge_loaded_weights(); + audio_tower_->verify_loaded_weights("thinker.audio_tower."); + audio_tower_->merge_loaded_weights(); + audio_tower_->to(options_.device(), + torch::typeMetaToScalarType(options_.dtype())); + + if (!model_args_.encoder_embedding_mode()) { + language_model_->load_model( + std::move(loader), "thinker.model.", "thinker.lm_head."); + } + } + + layer::NpuLmHead get_npu_lm_head() { + return language_model_->get_npu_lm_head(); + } + + void set_npu_lm_head(layer::NpuLmHead& head) { + language_model_->set_npu_lm_head(head); + } + + layer::NpuWordEmbedding get_npu_word_embedding() { + return language_model_->get_npu_word_embedding(); + } + + void set_npu_word_embedding(layer::NpuWordEmbedding& npu_word_embedding) { + language_model_->set_npu_word_embedding(npu_word_embedding); + } + + private: + ModelArgs model_args_; + torch::TensorOptions options_; + Qwen3OmniMoeThinkerVisionTransformer visual_{nullptr}; + Qwen3AudioEncoder audio_tower_{nullptr}; + Qwen3MoeForCausalLM language_model_{nullptr}; +}; +TORCH_MODULE(Qwen3OmniMoeThinkerForConditionalGeneration); + +REGISTER_MULTIMODAL_PROCESSOR(qwen3_omni_moe_thinker, + Qwen3OmniMoeMultimodalProcessor); +REGISTER_CAUSAL_VLM_MODEL(qwen3_omni_moe_thinker, + Qwen3OmniMoeThinkerForConditionalGeneration); +REGISTER_MPOSITION_GENERATOR(qwen3_omni_moe_thinker, + xllm::Qwen3VLMPositionGenerator); + +REGISTER_MODEL_ARGS(qwen3_omni_moe_thinker, + [&] { load_qwen3_omni_moe_model_args(json, args); }); + +} // namespace xllm::npu::model diff --git a/xllm/processors/CMakeLists.txt b/xllm/processors/CMakeLists.txt index f063b2855e..3c6e7d23ab 100644 --- a/xllm/processors/CMakeLists.txt +++ b/xllm/processors/CMakeLists.txt @@ -54,6 +54,12 @@ cc_library( qwen2_vl_prompt_processor.h qwen2_vl_image_processor.h qwen2_vl_video_processor.h + qwen3_asr_processor.h + qwen3_audio_common.h + qwen3_audio_processor.h + qwen3_audio_prompt_processor.h + qwen3_omni_moe_processor.h + qwen3_omni_moe_prompt_processor.h qwen3_vl_prompt_processor.h qwen3_vl_video_processor.h video_processor.h @@ -75,6 +81,9 @@ cc_library( qwen2_vl_prompt_processor.cpp qwen2_vl_image_processor.cpp qwen2_vl_video_processor.cpp + qwen3_audio_processor.cpp + qwen3_audio_prompt_processor.cpp + qwen3_omni_moe_prompt_processor.cpp qwen3_vl_prompt_processor.cpp qwen3_vl_video_processor.cpp DEPS diff --git a/xllm/processors/qwen3_asr_processor.h b/xllm/processors/qwen3_asr_processor.h new file mode 100644 index 0000000000..ac1114916d --- /dev/null +++ b/xllm/processors/qwen3_asr_processor.h @@ -0,0 +1,30 @@ +/* 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 "processors/multimodal_processor.h" +#include "processors/qwen3_audio_processor.h" +#include "processors/qwen3_audio_prompt_processor.h" + +namespace xllm { + +using Qwen3ASRMultimodalProcessor = + MultimodalProcessor; + +} // namespace xllm diff --git a/xllm/processors/qwen3_audio_common.h b/xllm/processors/qwen3_audio_common.h new file mode 100644 index 0000000000..218940952b --- /dev/null +++ b/xllm/processors/qwen3_audio_common.h @@ -0,0 +1,47 @@ +/* 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 +#include + +#include + +namespace xllm::qwen3_audio { + +inline constexpr char kInputFeaturesKey[] = "input_features"; +inline constexpr char kFeatureLengthKey[] = "feat_length"; +inline constexpr char kFeatureOriginLengthsKey[] = "feat_origin_lens"; +inline constexpr char kMaskKey[] = "audio|mask"; + +inline torch::Tensor get_feature_output_lengths( + const torch::Tensor& input_lengths, + int64_t window_length) { + CHECK_GT(window_length, 0); + + torch::Tensor output_lengths = input_lengths % window_length; + int64_t full_window_output_length = window_length; + constexpr int32_t kConvolutionLayerCount = 3; + for (int32_t index = 0; index < kConvolutionLayerCount; ++index) { + output_lengths = torch::floor_divide(output_lengths - 1, 2) + 1; + full_window_output_length = (full_window_output_length - 1) / 2 + 1; + } + output_lengths += torch::floor_divide(input_lengths, window_length) * + full_window_output_length; + return output_lengths.to(torch::kInt64); +} + +} // namespace xllm::qwen3_audio diff --git a/xllm/processors/qwen3_audio_processor.cpp b/xllm/processors/qwen3_audio_processor.cpp new file mode 100644 index 0000000000..cd5e07d46c --- /dev/null +++ b/xllm/processors/qwen3_audio_processor.cpp @@ -0,0 +1,130 @@ +/* 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 "processors/qwen3_audio_processor.h" + +#include +#include + +#include "core/util/audio_utils.h" +#include "processors/qwen3_audio_common.h" + +namespace xllm { + +Qwen3AudioProcessor::Qwen3AudioProcessor(const ModelArgs& args) + : feature_size_(args.mm_audio_feature_size()), + sampling_rate_(args.mm_audio_sampling_rate()), + hop_length_(args.mm_audio_hop_length()), + chunk_length_(args.mm_audio_chunk_length()), + n_fft_(args.mm_audio_n_fft()), + window_length_(args.mm_audio_n_window() * 2), + dither_(args.mm_audio_dither()), + truncation_(args.mm_audio_truncation()), + do_normalize_(args.mm_audio_do_normalize()) { + CHECK_GT(feature_size_, 0); + CHECK_GT(sampling_rate_, 0); + CHECK_GT(hop_length_, 0); + CHECK_GT(chunk_length_, 0); + CHECK_GT(n_fft_, 0); + CHECK_GT(window_length_, 0); + CHECK_EQ(feature_size_, args.mm_audio_num_mel_bins()); + mel_filters_ = + audio_utils::mel_filter_bank(1 + n_fft_ / 2, + feature_size_, + /*min_frequency=*/0.0, + /*max_frequency=*/sampling_rate_ / 2.0, + sampling_rate_, + /*norm=*/"slaney", + /*mel_scale=*/"slaney"); +} + +torch::Tensor Qwen3AudioProcessor::extract_log_mel_features( + const torch::Tensor& waveform) const { + const torch::TensorOptions options = + torch::TensorOptions().dtype(torch::kFloat32).device(torch::kCPU); + torch::Tensor audio = waveform.to(options); + if (dither_ != 0.0) { + audio = audio + torch::randn_like(audio) * dither_; + } + const torch::Tensor window = torch::hann_window(n_fft_, true, options); + const torch::Tensor stft = torch::stft(audio, + n_fft_, + hop_length_, + n_fft_, + window, + /*center=*/true, + /*pad_mode=*/"reflect", + /*normalized=*/false, + /*onesided=*/std::nullopt, + /*return_complex=*/true); + const torch::Tensor magnitudes = stft.slice(-1, 0, -1).abs().pow(2); + const torch::Tensor mel_spec = torch::matmul(mel_filters_.t(), magnitudes); + torch::Tensor log_spec = torch::clamp(mel_spec, 1e-10).log10(); + const torch::Tensor max_value = log_spec.max(); + log_spec = torch::maximum(log_spec, max_value - 8.0); + return (log_spec + 4.0) / 4.0; +} + +bool Qwen3AudioProcessor::process(const torch::Tensor& origin_audio, + const AudioMetadata& metadata, + MMDataItem& output_item) const { + if (origin_audio.dim() != 1) { + LOG(ERROR) << "Qwen3 audio processor only supports mono audio, got shape " + << origin_audio.sizes(); + return false; + } + if (metadata.sample_rate > 0 && metadata.sample_rate != sampling_rate_) { + LOG(ERROR) << "Qwen3 audio processor expects " << sampling_rate_ + << " Hz audio, got " << metadata.sample_rate << " Hz."; + return false; + } + + torch::Tensor waveform = origin_audio.to(torch::kCPU, torch::kFloat32); + const int64_t max_samples = chunk_length_ * sampling_rate_; + if (truncation_ && waveform.size(0) > max_samples) { + waveform = waveform.slice(0, 0, max_samples); + } + if (waveform.numel() == 0) { + LOG(ERROR) << "Qwen3 audio processor received empty audio."; + return false; + } + if (do_normalize_) { + const torch::Tensor variance = waveform.var(false); + waveform = (waveform - waveform.mean()) / torch::sqrt(variance + 1e-7); + } + + torch::Tensor features = extract_log_mel_features(waveform).transpose(0, 1); + const int64_t valid_frames = waveform.size(0) / hop_length_; + if (valid_frames <= 0) { + LOG(ERROR) << "Qwen3 audio is shorter than one feature frame."; + return false; + } + features = features.slice(0, 0, std::min(valid_frames, features.size(0))) + .contiguous(); + + const torch::Tensor feature_origin_lengths = + torch::tensor({features.size(0)}, torch::dtype(torch::kLong)); + const torch::Tensor feature_lengths = qwen3_audio::get_feature_output_lengths( + feature_origin_lengths, window_length_); + output_item = MMDataItem( + MMType::AUDIO, + MMDict{{qwen3_audio::kInputFeaturesKey, features}, + {qwen3_audio::kFeatureLengthKey, feature_lengths}, + {qwen3_audio::kFeatureOriginLengthsKey, feature_origin_lengths}}, + metadata); + return true; +} + +} // namespace xllm diff --git a/xllm/processors/qwen3_audio_processor.h b/xllm/processors/qwen3_audio_processor.h new file mode 100644 index 0000000000..625703423b --- /dev/null +++ b/xllm/processors/qwen3_audio_processor.h @@ -0,0 +1,48 @@ +/* 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 + +#include "core/framework/model/model_args.h" +#include "processors/audio_processor.h" + +namespace xllm { + +class Qwen3AudioProcessor final : public AudioProcessor { + public: + explicit Qwen3AudioProcessor(const ModelArgs& args); + + bool process(const torch::Tensor& origin_audio, + const AudioMetadata& metadata, + MMDataItem& output_item) const override; + + private: + torch::Tensor extract_log_mel_features(const torch::Tensor& waveform) const; + + int64_t feature_size_ = 0; + int64_t sampling_rate_ = 0; + int64_t hop_length_ = 0; + int64_t chunk_length_ = 0; + int64_t n_fft_ = 0; + int64_t window_length_ = 0; + double dither_ = 0.0; + bool truncation_ = false; + bool do_normalize_ = false; + torch::Tensor mel_filters_; +}; + +} // namespace xllm diff --git a/xllm/processors/qwen3_audio_prompt_processor.cpp b/xllm/processors/qwen3_audio_prompt_processor.cpp new file mode 100644 index 0000000000..c30689ad5c --- /dev/null +++ b/xllm/processors/qwen3_audio_prompt_processor.cpp @@ -0,0 +1,126 @@ +/* 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 "processors/qwen3_audio_prompt_processor.h" + +#include +#include +#include +#include +#include +#include + +#include "processors/qwen3_audio_common.h" + +namespace xllm { + +namespace { + +void append_tokens(std::string& output, + const std::string& token, + int64_t count) { + CHECK_GE(count, 0); + for (int64_t index = 0; index < count; ++index) { + output.append(token); + } +} + +void set_item_span(MMDataItem& item, int32_t offset, int32_t length) { + CHECK_GE(length, 0); + item.mutable_state().mutable_token_pos() = {offset, length}; + item.mutable_state().mutable_mm_token_mask() = torch::ones( + {length}, torch::TensorOptions().dtype(torch::kBool).device(torch::kCPU)); + item.mutable_state().mutable_mm_token_num() = length; +} + +} // namespace + +Qwen3AudioPromptProcessor::Qwen3AudioPromptProcessor(const ModelArgs& args) + : audio_token_id_(args.audio_token_id()), + audio_start_token_id_(args.audio_start_token_id()), + audio_end_token_id_(args.audio_end_token_id()) {} + +void Qwen3AudioPromptProcessor::process(std::string& prompt, + const MMData& mm_data) { + torch::Tensor feature_lengths; + if (std::optional value = + mm_data.get(qwen3_audio::kFeatureLengthKey)) { + feature_lengths = value.value(); + } + if (!feature_lengths.defined()) { + return; + } + + int64_t total_audio_tokens = 0; + const int64_t audio_count = feature_lengths.size(0); + for (int64_t index = 0; index < audio_count; ++index) { + total_audio_tokens += feature_lengths[index].item(); + } + + std::string output; + output.reserve(prompt.size() + + static_cast(total_audio_tokens) * audio_token_.size()); + + size_t begin = 0; + int64_t audio_index = 0; + size_t audio_position = prompt.find(audio_token_, begin); + while (audio_position != std::string::npos) { + CHECK_LT(audio_index, audio_count) + << "The prompt contains more audio placeholders than audio inputs."; + output.append(prompt, begin, audio_position - begin); + append_tokens( + output, audio_token_, feature_lengths[audio_index].item()); + ++audio_index; + begin = audio_position + audio_token_.size(); + audio_position = prompt.find(audio_token_, begin); + } + output.append(prompt, begin, std::string::npos); + CHECK_EQ(audio_index, audio_count) + << "The number of audio placeholders does not match audio inputs."; + prompt = std::move(output); +} + +void Qwen3AudioPromptProcessor::find_mm_spans( + const std::vector& token_ids, + MMData& mm_data) { + auto search_begin = token_ids.begin(); + int32_t audio_index = 0; + MMItemVec& mm_items = mm_data.items(); + while (true) { + auto audio_start = + std::find(search_begin, token_ids.end(), audio_start_token_id_); + if (audio_start == token_ids.end()) { + break; + } + auto audio_end = + std::find(audio_start + 1, token_ids.end(), audio_end_token_id_); + CHECK(audio_end != token_ids.end()); + CHECK(audio_start + 1 != audio_end); + CHECK_EQ(*(audio_start + 1), audio_token_id_); + CHECK_LT(audio_index, static_cast(mm_items.size())); + CHECK(mm_items[audio_index].is_type(MMType::AUDIO)); + + const int32_t offset = + static_cast(std::distance(token_ids.begin(), audio_start + 1)); + const int32_t length = + static_cast(std::distance(audio_start + 1, audio_end)); + set_item_span(mm_items[audio_index], offset, length); + ++audio_index; + search_begin = std::next(audio_end); + } + CHECK_EQ(audio_index, static_cast(mm_items.size())); +} + +} // namespace xllm diff --git a/xllm/processors/qwen3_audio_prompt_processor.h b/xllm/processors/qwen3_audio_prompt_processor.h new file mode 100644 index 0000000000..7cf58fd363 --- /dev/null +++ b/xllm/processors/qwen3_audio_prompt_processor.h @@ -0,0 +1,43 @@ +/* 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 +#include +#include + +#include "core/framework/model/model_args.h" +#include "processors/prompt_processor.h" + +namespace xllm { + +class Qwen3AudioPromptProcessor final : public PromptProcessor { + public: + explicit Qwen3AudioPromptProcessor(const ModelArgs& args); + + void process(std::string& prompt, const MMData& mm_data) override; + void find_mm_spans(const std::vector& token_ids, + MMData& mm_data) override; + + private: + const std::string audio_token_ = "<|audio_pad|>"; + + int32_t audio_token_id_ = 0; + int32_t audio_start_token_id_ = 0; + int32_t audio_end_token_id_ = 0; +}; + +} // namespace xllm diff --git a/xllm/processors/qwen3_omni_moe_processor.h b/xllm/processors/qwen3_omni_moe_processor.h new file mode 100644 index 0000000000..21e6cfa07f --- /dev/null +++ b/xllm/processors/qwen3_omni_moe_processor.h @@ -0,0 +1,32 @@ +/* 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 "processors/multimodal_processor.h" +#include "processors/qwen2_vl_image_processor.h" +#include "processors/qwen3_audio_processor.h" +#include "processors/qwen3_omni_moe_prompt_processor.h" +#include "processors/qwen3_vl_video_processor.h" + +namespace xllm { + +using Qwen3OmniMoeMultimodalProcessor = + MultimodalProcessor; + +} // namespace xllm diff --git a/xllm/processors/qwen3_omni_moe_prompt_processor.cpp b/xllm/processors/qwen3_omni_moe_prompt_processor.cpp new file mode 100644 index 0000000000..1ca7791554 --- /dev/null +++ b/xllm/processors/qwen3_omni_moe_prompt_processor.cpp @@ -0,0 +1,365 @@ +/* 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 "processors/qwen3_omni_moe_prompt_processor.h" + +#include + +#include +#include +#include +#include +#include +#include +#include +#include + +#include "processors/qwen3_audio_common.h" + +namespace xllm { + +namespace { + +class ModalityCursor final { + public: + ModalityCursor(const std::string& token, const torch::Tensor& sizes) + : token_(token), sizes_(sizes) { + if (sizes_.defined()) { + count_ = sizes_.size(0); + } + } + + torch::Tensor next_size() { + CHECK(index_ < count_) << "The index of " << token_ + << " modality is out of range, have " << count_ + << " modality inputs but try to access index " + << index_; + return sizes_[index_++]; + } + + const std::string& token() const { return token_; } + + void verify_consumed() const { + CHECK_EQ(index_, count_) << "The number of " << token_ + << " placeholders does not match modality inputs."; + } + + private: + const std::string& token_; + const torch::Tensor& sizes_; + int64_t index_ = 0; + int64_t count_ = 0; +}; + +void append_tokens(std::string& output, + const std::string& token, + int64_t count) { + CHECK_GE(count, 0); + for (int64_t index = 0; index < count; ++index) { + output.append(token); + } +} + +void set_item_span(MMDataItem& item, int32_t offset, int32_t length) { + CHECK_GE(length, 0); + item.mutable_state().mutable_token_pos() = {offset, length}; + item.mutable_state().mutable_mm_token_mask() = torch::ones( + {length}, torch::TensorOptions().dtype(torch::kBool).device(torch::kCPU)); + item.mutable_state().mutable_mm_token_num() = length; +} + +} // namespace + +Qwen3OmniMoePromptProcessor::Qwen3OmniMoePromptProcessor(const ModelArgs& args) + : vision_start_token_id_(args.vision_start_token_id()), + vision_end_token_id_(args.vision_end_token_id()), + image_token_id_(args.image_token_id()), + video_token_id_(args.video_token_id()), + audio_token_id_(args.audio_token_id()), + audio_start_token_id_(args.audio_start_token_id()), + audio_end_token_id_(args.audio_end_token_id()), + merge_size_(static_cast(args.mm_image_merge_size())), + position_id_per_seconds_(args.mm_position_id_per_seconds()), + use_audio_in_video_(args.mm_use_audio_in_video()) { + CHECK_GT(merge_size_, 0); + CHECK_GT(args.mm_fps(), 0.0); + video_second_per_grid_ = + static_cast(args.mm_temporal_patch_size()) / args.mm_fps(); +} + +void Qwen3OmniMoePromptProcessor::process(std::string& prompt, + const MMData& mm_data) { + torch::Tensor image_grid_thw; + if (std::optional value = + mm_data.get("image_grid_thw")) { + image_grid_thw = value.value(); + } + + torch::Tensor video_grid_thw; + if (std::optional value = + mm_data.get("video_grid_thw")) { + video_grid_thw = value.value(); + } + + torch::Tensor feature_lengths; + if (std::optional value = + mm_data.get(qwen3_audio::kFeatureLengthKey)) { + feature_lengths = value.value(); + } + + if (!image_grid_thw.defined() && !video_grid_thw.defined() && + !feature_lengths.defined()) { + return; + } + + const int32_t merge_length = merge_size_ * merge_size_; + + int64_t total_audio_tokens = 0; + if (feature_lengths.defined()) { + const int64_t count = feature_lengths.size(0); + for (int64_t index = 0; index < count; ++index) { + total_audio_tokens += feature_lengths[index].item(); + } + } + + int64_t total_image_tokens = 0; + if (image_grid_thw.defined()) { + const int64_t count = image_grid_thw.size(0); + for (int64_t index = 0; index < count; ++index) { + total_image_tokens += + image_grid_thw[index].prod().item() / merge_length; + } + } + + int64_t total_video_tokens = 0; + if (video_grid_thw.defined()) { + const int64_t count = video_grid_thw.size(0); + for (int64_t index = 0; index < count; ++index) { + total_video_tokens += + video_grid_thw[index].prod().item() / merge_length; + } + } + + const size_t reserve_size = + prompt.size() + + static_cast(total_image_tokens) * image_token_.size() + + static_cast(total_video_tokens) * video_token_.size() + + static_cast(total_audio_tokens) * audio_token_.size(); + std::string output; + output.reserve(reserve_size); + + ModalityCursor audio_cursor(audio_token_, feature_lengths); + ModalityCursor image_cursor(image_token_, image_grid_thw); + ModalityCursor video_cursor(video_token_, video_grid_thw); + + size_t begin = 0; + std::pair special_token = + find_special_token(prompt, begin); + while (special_token.second != std::string::npos) { + output.append(prompt, begin, special_token.second - begin); + + ModalityCursor* modality = nullptr; + if (special_token.first == TokenType::AUDIO) { + modality = &audio_cursor; + } else if (special_token.first == TokenType::IMAGE) { + modality = &image_cursor; + } else if (special_token.first == TokenType::VIDEO) { + modality = &video_cursor; + } + CHECK(modality != nullptr); + const torch::Tensor modality_size = modality->next_size(); + const std::string& modality_token = modality->token(); + if (special_token.first == TokenType::AUDIO) { + append_tokens(output, modality_token, modality_size.item()); + } else if (special_token.first == TokenType::VIDEO && use_audio_in_video_) { + const torch::Tensor audio_size = audio_cursor.next_size(); + const std::string& audio_token = audio_cursor.token(); + const torch::Tensor audio_token_indices = + torch::arange(audio_size.item(), torch::kInt32); + + const int32_t temporal = modality_size[0].item(); + const int32_t height = modality_size[1].item() / merge_size_; + const int32_t width = modality_size[2].item() / merge_size_; + CHECK_GT(temporal, 0); + CHECK_GT(height, 0); + CHECK_GT(width, 0); + + torch::Tensor video_token_indices = + torch::arange(temporal, torch::kFloat32).view({temporal, 1, 1}); + video_token_indices = + video_token_indices.expand({temporal, height, width}).reshape({-1}); + video_token_indices = video_token_indices * video_second_per_grid_ * + position_id_per_seconds_; + auto video_indices = video_token_indices.accessor(); + auto audio_indices = audio_token_indices.accessor(); + + std::string placeholder = audio_start_token_; + size_t video_index = 0; + size_t audio_index = 0; + const size_t video_length = video_indices.size(0); + const size_t audio_length = audio_indices.size(0); + while (video_index < video_length && audio_index < audio_length) { + if (video_indices[video_index] <= audio_indices[audio_index]) { + placeholder.append(modality_token); + ++video_index; + } else { + placeholder.append(audio_token); + ++audio_index; + } + } + append_tokens(placeholder, + modality_token, + static_cast(video_length - video_index)); + append_tokens(placeholder, + audio_token, + static_cast(audio_length - audio_index)); + placeholder.append(audio_end_token_); + output.append(placeholder); + } else { + append_tokens(output, + modality_token, + modality_size.prod().item() / merge_length); + } + + begin = special_token.second + modality_token.size(); + special_token = find_special_token(prompt, begin); + } + + if (begin < prompt.size()) { + output.append(prompt, begin, std::string::npos); + } + audio_cursor.verify_consumed(); + image_cursor.verify_consumed(); + video_cursor.verify_consumed(); + prompt = std::move(output); +} + +void Qwen3OmniMoePromptProcessor::find_mm_spans( + const std::vector& token_ids, + MMData& mm_data) { + auto search_begin = token_ids.begin(); + int32_t global_mm_index = 0; + MMItemVec& mm_items = mm_data.items(); + while (true) { + auto vision_start = + std::find(search_begin, token_ids.end(), vision_start_token_id_); + auto vision_end = + vision_start == token_ids.end() + ? token_ids.end() + : std::find( + vision_start + 1, token_ids.end(), vision_end_token_id_); + auto audio_start = + std::find(search_begin, token_ids.end(), audio_start_token_id_); + auto audio_end = + audio_start == token_ids.end() + ? token_ids.end() + : std::find(audio_start + 1, token_ids.end(), audio_end_token_id_); + if (vision_start == token_ids.end() && audio_start == token_ids.end()) { + break; + } + + auto span_start = std::min(vision_start, audio_start); + auto span_end = std::min(vision_end, audio_end); + auto outer_span_end = std::max(vision_end, audio_end); + CHECK(span_start != token_ids.end()); + if (span_start == vision_start) { + CHECK(vision_end != token_ids.end()); + } else { + CHECK(audio_end != token_ids.end()); + } + CHECK(span_end != token_ids.end()); + CHECK(global_mm_index < static_cast(mm_items.size())); + + const int32_t offset = + static_cast(std::distance(token_ids.begin(), span_start)); + const int32_t length = + static_cast(std::distance(span_start + 1, span_end)); + MMDataItem& item = mm_items[global_mm_index]; + int32_t consumed_item_count = 1; + + if (*span_start == vision_start_token_id_ && + span_start + 1 != token_ids.end() && + *(span_start + 1) == audio_start_token_id_) { + CHECK(item.is_type(MMType::VIDEO)); + CHECK(global_mm_index + 1 < static_cast(mm_items.size())); + CHECK(mm_items[global_mm_index + 1].is_type(MMType::AUDIO)); + CHECK_GT(length, 1); + CHECK_EQ(*(span_start + 2), video_token_id_) + << "Audio-in-video placeholder must start with a video token."; + set_item_span(item, offset + 2, length - 1); + std::vector audio_in_video_token_ids( + token_ids.begin() + offset + 2, + token_ids.begin() + offset + 2 + length - 1); + item.add(qwen3_omni_moe::kAudioInVideoTokenIdsKey, + torch::tensor(audio_in_video_token_ids, torch::kInt32)); + set_item_span(mm_items[global_mm_index + 1], offset + 2, 0); + consumed_item_count = 2; + span_end = outer_span_end; + } else { + const int32_t first_token = *(span_start + 1); + if (*span_start == vision_start_token_id_) { + if (first_token == image_token_id_) { + CHECK(item.is_type(MMType::IMAGE)); + } else { + CHECK_EQ(first_token, video_token_id_); + CHECK(item.is_type(MMType::VIDEO)); + } + } else { + CHECK(item.is_type(MMType::AUDIO)); + CHECK_EQ(first_token, audio_token_id_); + } + set_item_span(item, offset + 1, length); + } + global_mm_index += consumed_item_count; + search_begin = std::next(span_end); + } + CHECK_EQ(global_mm_index, static_cast(mm_items.size())); +} + +std::pair +Qwen3OmniMoePromptProcessor::find_special_token(const std::string& prompt, + size_t begin) const { + struct TokenInfo { + const std::string& token; + TokenType type; + size_t position = std::string::npos; + }; + + std::array tokens = {{{image_token_, TokenType::IMAGE}, + {video_token_, TokenType::VIDEO}, + {audio_token_, TokenType::AUDIO}}}; + for (TokenInfo& token : tokens) { + token.position = prompt.find(token.token, begin); + } + + auto earliest = + std::min_element(tokens.begin(), + tokens.end(), + [](const TokenInfo& lhs, const TokenInfo& rhs) { + if (lhs.position == std::string::npos) { + return false; + } + if (rhs.position == std::string::npos) { + return true; + } + return lhs.position < rhs.position; + }); + if (earliest == tokens.end() || earliest->position == std::string::npos) { + return {TokenType::INVALID, std::string::npos}; + } + return {earliest->type, earliest->position}; +} + +} // namespace xllm diff --git a/xllm/processors/qwen3_omni_moe_prompt_processor.h b/xllm/processors/qwen3_omni_moe_prompt_processor.h new file mode 100644 index 0000000000..78e3aa66e8 --- /dev/null +++ b/xllm/processors/qwen3_omni_moe_prompt_processor.h @@ -0,0 +1,65 @@ +/* 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 +#include +#include +#include + +#include "core/framework/model/model_args.h" +#include "processors/prompt_processor.h" + +namespace xllm { + +namespace qwen3_omni_moe { +inline constexpr char kAudioInVideoTokenIdsKey[] = "audio_in_video_token_ids"; +} // namespace qwen3_omni_moe + +class Qwen3OmniMoePromptProcessor final : public PromptProcessor { + public: + explicit Qwen3OmniMoePromptProcessor(const ModelArgs& args); + + void process(std::string& prompt, const MMData& mm_data) override; + void find_mm_spans(const std::vector& token_ids, + MMData& mm_data) override; + + private: + enum class TokenType { INVALID, IMAGE, VIDEO, AUDIO }; + + std::pair find_special_token(const std::string& prompt, + size_t begin) const; + + const std::string image_token_ = "<|image_pad|>"; + const std::string video_token_ = "<|video_pad|>"; + const std::string audio_token_ = "<|audio_pad|>"; + const std::string audio_start_token_ = "<|audio_start|>"; + const std::string audio_end_token_ = "<|audio_end|>"; + + int32_t vision_start_token_id_ = 0; + int32_t vision_end_token_id_ = 0; + int32_t image_token_id_ = 0; + int32_t video_token_id_ = 0; + int32_t audio_token_id_ = 0; + int32_t audio_start_token_id_ = 0; + int32_t audio_end_token_id_ = 0; + int32_t merge_size_ = 0; + int32_t position_id_per_seconds_ = 0; + bool use_audio_in_video_ = false; + double video_second_per_grid_ = 0.0; +}; + +} // namespace xllm