feat: add NPU fused norm operators for Wan2.2 DiT.

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

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<torch::Tensor> weight,
+                                  std::optional<torch::Tensor> 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

xllm/core/kernels/npu/npu_ops_api.h

@@ -143,6 +143,23 @@ std::tuple<torch::Tensor, torch::Tensor, torch::Tensor> 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<torch::Tensor> weight,
+                                  std::optional<torch::Tensor> bias,
+                                  double eps);
+
 void apply_rotary(torch::Tensor& q,
                   torch::Tensor& k,
                   const torch::Tensor& cos_sin_cache,

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<torch::Tensor, torch::Tensor> rms_norm_dynamic_quant(
     RmsNormDynamicQuantParams& params) {
 #if defined(USE_NPU)

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<torch::Tensor, torch::Tensor> rms_norm_dynamic_quant(
     RmsNormDynamicQuantParams& params);
 

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<torch::Tensor> weight;
+  // Optional affine bias (beta) [H]. Only used when elementwise_affine.
+  std::optional<torch::Tensor> 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;

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

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

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 <glog/logging.h>
+#include <torch/torch.h>
+
+#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