diff --git a/xllm/core/kernels/npu/CMakeLists.txt b/xllm/core/kernels/npu/CMakeLists.txt index c647a19819..cf49301367 100644 --- a/xllm/core/kernels/npu/CMakeLists.txt +++ b/xllm/core/kernels/npu/CMakeLists.txt @@ -14,6 +14,7 @@ cc_library( active.cpp attention.cpp fused_layernorm.cpp + npu_adalayernorm.cpp matmul.cpp npu_causal_conv1d.cpp npu_fused_infer_attention.cpp diff --git a/xllm/core/kernels/npu/npu_adalayernorm.cpp b/xllm/core/kernels/npu/npu_adalayernorm.cpp new file mode 100644 index 0000000000..78b9dfc910 --- /dev/null +++ b/xllm/core/kernels/npu/npu_adalayernorm.cpp @@ -0,0 +1,109 @@ +/* Copyright 2026 The xLLM Authors. + +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/kernels/npu/aclnn/pytorch_npu_helper.hpp" +#include "core/kernels/npu/npu_ops_api.h" + +namespace xllm::kernel::npu { + +torch::Tensor fused_adalayer_norm(const torch::Tensor& input, + const torch::Tensor& scale, + const torch::Tensor& shift, + std::optional weight, + std::optional bias, + double eps) { + CHECK(input.dim() == 3) << "Input tensor must be 3D [B, S, H]."; + CHECK(input.numel() > 0) << "Input tensor should not be empty."; + CHECK(scale.dim() == 2 || scale.dim() == 3) + << "Scale tensor dim must be 2 or 3."; + CHECK(shift.dim() == 2 || shift.dim() == 3) + << "Shift tensor dim must be 2 or 3."; + CHECK(scale.size(-1) == input.size(-1)) + << "Scale last dim must equal input last dim."; + CHECK(shift.size(-1) == input.size(-1)) + << "Shift last dim must equal input last dim."; + if (weight.has_value()) { + CHECK(weight.value().dim() == 1 && weight.value().size(0) == input.size(-1)) + << "Weight dim must equal input last dim."; + } + if (bias.has_value()) { + CHECK(bias.value().dim() == 1 && bias.value().size(0) == input.size(-1)) + << "Bias dim must equal input last dim."; + } + + // aclnnAdaLayerNorm only supports broadcast modulation where scale/shift are + // [B, H] or [B, 1, H] (a single modulation vector broadcast over the + // sequence dimension). A genuine token-wise modulation [B, S, H] with S > 1 + // is not directly supported. In that case, mirror the mindiesd reference: + // fold the sequence dimension into the batch so each token becomes its own + // "row" -- input -> [B*S, 1, H], scale/shift -> [B*S, H] -- run the op, then + // restore the original [B, S, H] shape. + const int64_t batch_size = input.size(0); + const int64_t seq_len = input.size(1); + const int64_t hidden_size = input.size(2); + + auto is_tokenwise = [&](const torch::Tensor& t) { + return t.dim() == 3 && t.size(1) == seq_len && seq_len != 1; + }; + const bool tokenwise = is_tokenwise(scale) || is_tokenwise(shift); + + torch::Tensor x_arg = input; + // aclnnAdaLayerNorm's 2D path is faster than its 3D broadcast path. + // Squeeze [B, 1, H] -> [B, H] to use the faster 2D code path. + auto normalize_modulation = [](const torch::Tensor& t) { + if (t.dim() == 3 && t.size(1) == 1) { + return t.squeeze(1); // [B, 1, H] -> [B, H] + } + return t; + }; + torch::Tensor scale_arg = normalize_modulation(scale); + torch::Tensor shift_arg = normalize_modulation(shift); + + if (tokenwise) { + auto expand_modulation = [&](const torch::Tensor& t) { + torch::Tensor expanded; + if (t.dim() == 2) { + expanded = t.unsqueeze(1).expand({batch_size, seq_len, hidden_size}); + } else if (t.size(1) == 1) { + expanded = t.expand({batch_size, seq_len, hidden_size}); + } else { + expanded = t; + } + return expanded.reshape({batch_size * seq_len, hidden_size}); + }; + x_arg = input.reshape({batch_size * seq_len, 1, hidden_size}); + scale_arg = expand_modulation(scale); + shift_arg = expand_modulation(shift); + } + + auto output = torch::empty(x_arg.sizes().vec(), x_arg.options()); + + EXEC_NPU_CMD(aclnnAdaLayerNorm, + x_arg, + scale_arg, + shift_arg, + weight, + bias, + eps, + output); + + if (tokenwise) { + output = output.reshape({batch_size, seq_len, hidden_size}); + } + + return output; +} + +} // namespace xllm::kernel::npu diff --git a/xllm/core/kernels/npu/npu_ops_api.h b/xllm/core/kernels/npu/npu_ops_api.h index dea67faf54..c6da986340 100644 --- a/xllm/core/kernels/npu/npu_ops_api.h +++ b/xllm/core/kernels/npu/npu_ops_api.h @@ -143,6 +143,23 @@ std::tuple add_rms_norm( const torch::Tensor& gamma, double epsilon); +/// @brief Fused adaptive LayerNorm: LayerNorm(x, weight, bias) * (1 + scale) + +/// shift. Calls aclnnAdaLayerNorm under the hood. +/// The (1 + scale) is applied inside the kernel, so pass the raw scale. +/// @param input Input tensor [B, L, C] +/// @param scale Raw scale modulation factor [B, C] (kernel applies +1) +/// @param shift Shift modulation factor [B, C] +/// @param weight Optional affine weight [C] +/// @param bias Optional affine bias [C] +/// @param eps Epsilon for numerical stability +/// @return Normalized and modulated output [B, L, C] +torch::Tensor fused_adalayer_norm(const torch::Tensor& input, + const torch::Tensor& scale, + const torch::Tensor& shift, + std::optional weight, + std::optional bias, + double eps); + void apply_rotary(torch::Tensor& q, torch::Tensor& k, const torch::Tensor& cos_sin_cache, diff --git a/xllm/core/kernels/ops_api.cpp b/xllm/core/kernels/ops_api.cpp index b38addd7df..23797f8a0f 100644 --- a/xllm/core/kernels/ops_api.cpp +++ b/xllm/core/kernels/ops_api.cpp @@ -389,6 +389,20 @@ void fused_layernorm(FusedLayerNormParams& params) { #endif } +torch::Tensor fused_adalayer_norm(AdaLayerNormParams& params) { +#if defined(USE_NPU) + params.output = npu::fused_adalayer_norm(params.input, + params.scale, + params.shift, + params.weight, + params.bias, + params.eps); +#else + NOT_IMPLEMENTED(); +#endif + return params.output; +} + std::tuple rms_norm_dynamic_quant( RmsNormDynamicQuantParams& params) { #if defined(USE_NPU) diff --git a/xllm/core/kernels/ops_api.h b/xllm/core/kernels/ops_api.h index 1ead17c396..c1667b77bb 100644 --- a/xllm/core/kernels/ops_api.h +++ b/xllm/core/kernels/ops_api.h @@ -42,6 +42,8 @@ void dequant_from_paged_cache(ReshapeFromCacheParams& params); void fused_layernorm(FusedLayerNormParams& params); +torch::Tensor fused_adalayer_norm(AdaLayerNormParams& params); + std::tuple rms_norm_dynamic_quant( RmsNormDynamicQuantParams& params); diff --git a/xllm/core/kernels/param.h b/xllm/core/kernels/param.h index ec1b2ff182..12301cfa48 100644 --- a/xllm/core/kernels/param.h +++ b/xllm/core/kernels/param.h @@ -277,6 +277,27 @@ struct FusedLayerNormParams { bool dynamic_quant = false; }; +// Fused adaptive LayerNorm parameters. +// Computes: LayerNorm(x) * (1 + scale) + shift in a single fused kernel. +struct AdaLayerNormParams { + // Input tensor [B, S, H]. Last dimension is hidden_size. + torch::Tensor input; + // Scale modulation. Shape [B, H] or [B, 1, H] (broadcast over sequence) + // or [B, S, H] (token-wise). The (1 + scale) is applied inside the kernel, + // so pass the raw scale. + torch::Tensor scale; + // Shift modulation. Same shape constraints as scale. + torch::Tensor shift; + // Optional affine weight (gamma) [H]. Only used when elementwise_affine. + std::optional weight; + // Optional affine bias (beta) [H]. Only used when elementwise_affine. + std::optional bias; + // Output tensor. Same shape as input. Written back by the kernel. + torch::Tensor output; + // Epsilon for numerical stability. + double eps = 1e-6; +}; + struct RmsNormDynamicQuantParams { torch::Tensor input; torch::Tensor weight; diff --git a/xllm/core/layers/common/CMakeLists.txt b/xllm/core/layers/common/CMakeLists.txt index a34e89f0f2..a885373e67 100755 --- a/xllm/core/layers/common/CMakeLists.txt +++ b/xllm/core/layers/common/CMakeLists.txt @@ -10,6 +10,7 @@ cc_library( qwen3_next_rms_norm.h deepseek_v4_rotary_embedding.h partial_rotary_embedding.h + ada_layer_norm.h rms_norm_gated.h rms_norm.h rotary_embedding.h @@ -35,6 +36,7 @@ cc_library( qwen3_next_rms_norm.cpp deepseek_v4_rotary_embedding.cpp partial_rotary_embedding.cpp + ada_layer_norm.cpp rms_norm_gated.cpp rms_norm.cpp rotary_embedding.cpp diff --git a/xllm/core/layers/common/ada_layer_norm.cpp b/xllm/core/layers/common/ada_layer_norm.cpp new file mode 100644 index 0000000000..82ee2f2e93 --- /dev/null +++ b/xllm/core/layers/common/ada_layer_norm.cpp @@ -0,0 +1,62 @@ +/* Copyright 2026 The xLLM Authors. + +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 "ada_layer_norm.h" + +#include "kernels/ops_api.h" + +namespace xllm { +namespace layer { + +AdaLayerNormImpl::AdaLayerNormImpl(int64_t dim, + double eps, + bool elementwise_affine, + const torch::TensorOptions& options) + : norm_dim_(dim), eps_(eps), elementwise_affine_(elementwise_affine) { + if (elementwise_affine_) { + weight_ = register_parameter( + "weight", torch::ones({dim}, options), /*requires_grad=*/false); + bias_ = register_parameter( + "bias", torch::zeros({dim}, options), /*requires_grad=*/false); + } +} + +torch::Tensor AdaLayerNormImpl::forward(const torch::Tensor& input, + const torch::Tensor& scale, + const torch::Tensor& shift) { + xllm::kernel::AdaLayerNormParams params; + params.input = input; + params.scale = scale; + params.shift = shift; + if (weight_.defined()) { + params.weight = weight_; + } + if (bias_.defined()) { + params.bias = bias_; + } + params.eps = eps_; + + return xllm::kernel::fused_adalayer_norm(params); +} + +void AdaLayerNormImpl::load_state_dict(const StateDict& state_dict) { + if (elementwise_affine_) { + LOAD_WEIGHT(weight); + LOAD_WEIGHT(bias); + } +} + +} // namespace layer +} // namespace xllm diff --git a/xllm/core/layers/common/ada_layer_norm.h b/xllm/core/layers/common/ada_layer_norm.h new file mode 100644 index 0000000000..4155b6452c --- /dev/null +++ b/xllm/core/layers/common/ada_layer_norm.h @@ -0,0 +1,68 @@ +/* Copyright 2026 The xLLM Authors. + +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 "framework/state_dict/state_dict.h" +#include "framework/state_dict/utils.h" + +namespace xllm { +namespace layer { + +// Fused adaptive LayerNorm (AdaLayerNorm). +// Computes: LayerNorm(x) * (1 + scale) + shift in a single fused kernel. +// The caller passes the raw scale/shift (the +1 is applied by the kernel). +class AdaLayerNormImpl : public torch::nn::Module { + public: + AdaLayerNormImpl(int64_t dim, + double eps, + bool elementwise_affine, + const torch::TensorOptions& options); + + // input: [B, S, H] + // scale: [B, H], [B, 1, H], or [B, S, H] (raw scale, +1 applied internally) + // shift: same shape constraints as scale + torch::Tensor forward(const torch::Tensor& input, + const torch::Tensor& scale, + const torch::Tensor& shift); + + void load_state_dict(const StateDict& state_dict); + + void verify_loaded_weights(const std::string& prefix) const { + if (elementwise_affine_) { + CHECK(weight_is_loaded_) + << "weight is not loaded for " << prefix + "weight"; + CHECK(bias_is_loaded_) << "bias is not loaded for " << prefix + "bias"; + } + } + + torch::Tensor weight() const { return weight_; } + torch::Tensor bias() const { return bias_; } + double eps() const { return eps_; } + + private: + DEFINE_WEIGHT(weight); + DEFINE_WEIGHT(bias); + int64_t norm_dim_; + double eps_; + bool elementwise_affine_; +}; +TORCH_MODULE(AdaLayerNorm); + +} // namespace layer +} // namespace xllm diff --git a/xllm/models/dit/transformers/transformer_wan.h b/xllm/models/dit/transformers/transformer_wan.h index 9a702ee2af..ced8da7ea5 100644 --- a/xllm/models/dit/transformers/transformer_wan.h +++ b/xllm/models/dit/transformers/transformer_wan.h @@ -27,12 +27,14 @@ limitations under the License. #include #include +#include "core/framework/config/dit_config.h" #include "core/framework/config/load_config.h" #include "core/framework/config/parallel_config.h" #include "core/framework/dit_model_loader.h" #include "core/framework/model/model_input_params.h" #include "core/framework/state_dict/state_dict.h" #include "core/framework/state_dict/utils.h" +#include "core/layers/common/ada_layer_norm.h" #include "core/layers/common/add_matmul.h" #include "core/layers/common/rms_norm.h" #include "models/dit/utils/dit_parallel_linear.h" @@ -1268,10 +1270,13 @@ class WanTransformerBlockImpl : public torch::nn::Module { cross_attn_norm_ = model_args.cross_attn_norm(); qk_norm_ = model_args.qk_norm(); - norm1_ = - register_module("norm1", FP32LayerNorm(context, dim_, eps_, false)); + ada_norm1_ = register_module( + "ada_norm1", + layer::AdaLayerNorm( + dim_, eps_, /*elementwise_affine=*/false, options_)); attn1_ = register_module( "attn1", WanAttention(context, parallel_args, -1, rainfusion_config)); + attn2_ = register_module( "attn2", WanAttention( @@ -1291,8 +1296,10 @@ class WanTransformerBlockImpl : public torch::nn::Module { false, ffn_dim_, true)); - norm3_ = - register_module("norm3", FP32LayerNorm(context, dim_, eps_, false)); + ada_norm3_ = register_module( + "ada_norm3", + layer::AdaLayerNorm( + dim_, eps_, /*elementwise_affine=*/false, options_)); scale_shift_table_ = register_parameter("scale_shift_table", torch::randn({1, 6, dim_}, options_) / @@ -1332,11 +1339,15 @@ class WanTransformerBlockImpl : public torch::nn::Module { c_gate_msa = splits[5]; } - torch::Tensor norm1_result = norm1_->forward(hidden_states); + auto scale_msa_2d = + scale_msa.dim() == 3 ? scale_msa.select(1, 0) : scale_msa; + auto shift_msa_2d = + shift_msa.dim() == 3 ? shift_msa.select(1, 0) : shift_msa; torch::Tensor norm_hidden_states = - (norm1_result.to(hidden_states.dtype()) * (1 + scale_msa) + shift_msa); + ada_norm1_->forward(hidden_states, scale_msa_2d, shift_msa_2d); torch::Tensor attn_output = attn1_->forward( norm_hidden_states, norm_hidden_states, rotary_emb, rf_state); + hidden_states = hidden_states + attn_output * gate_msa; if (cross_attn_norm_) { @@ -1348,8 +1359,12 @@ class WanTransformerBlockImpl : public torch::nn::Module { attn_output = attn2_->forward( norm_hidden_states, encoder_hidden_states, std::nullopt, rf_state); hidden_states = hidden_states + attn_output; - torch::Tensor norm2_result = norm3_->forward(hidden_states); - norm_hidden_states = (norm2_result * (1 + c_scale_msa) + c_shift_msa); + auto c_scale_msa_2d = + c_scale_msa.dim() == 3 ? c_scale_msa.select(1, 0) : c_scale_msa; + auto c_shift_msa_2d = + c_shift_msa.dim() == 3 ? c_shift_msa.select(1, 0) : c_shift_msa; + norm_hidden_states = + ada_norm3_->forward(hidden_states, c_scale_msa_2d, c_shift_msa_2d); torch::Tensor ff_output = ff_->forward(norm_hidden_states); hidden_states = hidden_states + ff_output * c_gate_msa; @@ -1359,7 +1374,7 @@ class WanTransformerBlockImpl : public torch::nn::Module { void load_state_dict(const StateDict& state_dict) { attn1_->load_state_dict(state_dict.get_dict_with_prefix("attn1.")); attn2_->load_state_dict(state_dict.get_dict_with_prefix("attn2.")); - if (cross_attn_norm_ && norm2_) { + if (cross_attn_norm_) { norm2_->load_state_dict(state_dict.get_dict_with_prefix("norm2.")); } ff_->load_state_dict(state_dict.get_dict_with_prefix("ffn.")); @@ -1401,12 +1416,12 @@ class WanTransformerBlockImpl : public torch::nn::Module { int64_t block_idx_ = 0; std::string qk_norm_; - FP32LayerNorm norm1_{nullptr}; WanAttention attn1_{nullptr}; WanAttention attn2_{nullptr}; - FP32LayerNorm norm2_{nullptr}; WanFeedForward ff_{nullptr}; - FP32LayerNorm norm3_{nullptr}; + layer::AdaLayerNorm ada_norm1_{nullptr}; // self-attn pre-norm (fused) + FP32LayerNorm norm2_{nullptr}; // cross-attn pre-norm (bf16 LayerNorm) + layer::AdaLayerNorm ada_norm3_{nullptr}; // FFN pre-norm (fused) torch::Tensor scale_shift_table_; bool scale_shift_table_loaded_{false}; @@ -1470,8 +1485,10 @@ class WanTransformer3DModelImpl : public torch::nn::Module { transformer_layers_.push_back(block); } - norm_out_ = register_module( - "norm_out", FP32LayerNorm(context, inner_dim_, 1e-6, false)); + ada_norm_out_ = register_module( + "ada_norm_out", + layer::AdaLayerNorm( + inner_dim_, 1e-6, /*elementwise_affine=*/false, options_)); int64_t patch_prod = patch_size_[0] * patch_size_[1] * patch_size_[2]; proj_out_ = register_module( "proj_out", @@ -1605,12 +1622,15 @@ class WanTransformer3DModelImpl : public torch::nn::Module { shift = shift.to(hidden_states.device()); scale = scale.to(hidden_states.device()); - auto norm_result = norm_out_->forward(hidden_states, /*keep_fp32*/ true); - auto one_plus_scale = - (1 + scale.to(hidden_states.dtype())).to(torch::kFloat32); - auto shift_fp32 = shift.to(torch::kFloat32); - auto norm_out = norm_result * one_plus_scale + shift_fp32; - hidden_states = norm_out.to(hidden_states.dtype()); + auto hidden_states_dtype = hidden_states.dtype(); + + // Drop the redundant sequence dim so the fused kernel uses the fast 2D + // [B,H] path instead of the token-wise fold. + auto scale_2d = scale.dim() == 3 ? scale.select(1, 0) : scale; + auto shift_2d = shift.dim() == 3 ? shift.select(1, 0) : shift; + hidden_states = ada_norm_out_->forward(hidden_states, + scale_2d.to(hidden_states_dtype), + shift_2d.to(hidden_states_dtype)); if (::xllm::ParallelConfig::get_instance().sp_size() > 1 && seq_len != pad_seq_len) { @@ -1699,7 +1719,7 @@ class WanTransformer3DModelImpl : public torch::nn::Module { rope_->set_freqs_cos(freqs_cos_fp32.to(device)); rope_->set_freqs_sin(freqs_sin_fp32.to(device)); condition_embedder_->to(device); - norm_out_->to(device); + ada_norm_out_->to(device); proj_out_->to(device); scale_shift_table_.set_data(scale_shift_table_.to(device)); @@ -1743,7 +1763,7 @@ class WanTransformer3DModelImpl : public torch::nn::Module { WanRotaryPosEmbed rope_{nullptr}; torch::nn::ModuleList blocks_; std::vector transformer_layers_; - FP32LayerNorm norm_out_{nullptr}; + layer::AdaLayerNorm ada_norm_out_{nullptr}; // final norm (fused) layer::AddMatmul proj_out_{nullptr}; torch::Tensor scale_shift_table_; bool scale_shift_table_loaded_{false};