[init] baseline7 from project_6

ex_engine/xllm_layers/common/activation.cpp

@@ -0,0 +1,38 @@
+/* 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.
+==============================================================================*/
+
+#include "activation.h"
+
+#include "kernels/ops_api.h"
+namespace xllm {
+namespace layer {
+
+ActivationImpl::ActivationImpl(const std::string& act_mode, bool is_gated)
+    : act_mode_(act_mode), is_gated_(is_gated) {}
+
+void ActivationImpl::forward(torch::Tensor& input, torch::Tensor& output) {
+  xllm::kernel::ActivationParams activation_params;
+  activation_params.input = input;
+  activation_params.output = output;
+  activation_params.act_mode = act_mode_;
+  activation_params.is_gated = is_gated_;
+  xllm::kernel::active(activation_params);
+  // Unified assignment: NPU returns new tensor, others modify in-place (no-op
+  // assignment)
+  output = activation_params.output;
+}
+
+}  // namespace layer
+}  // namespace xllm

ex_engine/xllm_layers/common/activation.h

@@ -0,0 +1,38 @@
+/* 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 <torch/torch.h>
+
+#include <string>
+
+namespace xllm {
+namespace layer {
+
+class ActivationImpl : public torch::nn::Module {
+ public:
+  ActivationImpl(const std::string& act_mode, bool is_gated);
+
+  void forward(torch::Tensor& input, torch::Tensor& output);
+
+ private:
+  std::string act_mode_;
+  bool is_gated_;
+};
+TORCH_MODULE(Activation);
+
+}  // namespace layer
+}  // namespace xllm

ex_engine/xllm_layers/common/dense_mlp.cpp

@@ -0,0 +1,141 @@
+/* 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.
+==============================================================================*/
+
+#include "dense_mlp.h"
+
+#include <glog/logging.h>
+
+#include "kernels/ops_api.h"
+#include "platform/device.h"
+
+namespace xllm {
+namespace layer {
+
+DenseMLPImpl::DenseMLPImpl(int64_t hidden_size,
+                           int64_t intermediate_size,
+                           bool is_gated,
+                           bool has_bias,
+                           const std::string& hidden_act,
+                           bool enable_result_reduction,
+                           const QuantArgs& quant_args,
+                           ProcessGroup* process_group,
+                           const torch::TensorOptions& options,
+                           const std::string& module_prefix)
+    : is_gated_(is_gated),
+      intermediate_size_(intermediate_size),
+      process_group_(process_group),
+      hidden_act_(hidden_act) {
+  // Check if using w8a8 smoothquant quantization
+  is_smoothquant_ = quant_args.quant_method() == kQuantMethodSmoothquant;
+
+  if (is_smoothquant_) {
+    // Safety check: only w8a8 smoothquant is supported
+    if (quant_args.bits() != 8 || !quant_args.activation_dynamic()) {
+      LOG(FATAL)
+          << "DenseMLP w8a8 mode only supports w8a8 smoothquant quantization. "
+          << "Got bits=" << quant_args.bits()
+          << ", activation_dynamic=" << quant_args.activation_dynamic();
+    }
+  }
+
+  // Determine extra args based on quantization mode
+  LinearExtraArgs gate_up_proj_extra_args("none", false);
+  LinearExtraArgs down_proj_extra_args("none", false);
+  if (is_smoothquant_) {
+    // For per-token smoothquant, use specific args
+    down_proj_extra_args = LinearExtraArgs(hidden_act_, is_gated_);
+  }
+
+  // 1. gate + up
+  int64_t out_feature = is_gated_ ? intermediate_size_ * 2 : intermediate_size_;
+  gate_up_proj_ =
+      register_module("gate_up_proj",
+                      ColumnParallelLinear(hidden_size,
+                                           out_feature,
+                                           /*bias=*/has_bias,
+                                           /*gather_output=*/false,
+                                           quant_args,
+                                           process_group_,
+                                           options,
+                                           gate_up_proj_extra_args));
+
+  act_ = register_module("act", Activation(hidden_act_, is_gated_));
+
+  // 2. down
+  const auto down_proj_quant_args =
+      module_prefix.empty()
+          ? quant_args
+          : quant_args.for_module(module_prefix + ".down_proj");
+  down_proj_ = register_module("down_proj",
+                               RowParallelLinear(intermediate_size_,
+                                                 hidden_size,
+                                                 /*bias=*/has_bias,
+                                                 /*input_is_parallelized=*/true,
+                                                 enable_result_reduction,
+                                                 down_proj_quant_args,
+                                                 process_group_,
+                                                 options,
+                                                 down_proj_extra_args));
+}
+
+torch::Tensor DenseMLPImpl::forward(const torch::Tensor& hidden_states) {
+  // input shape: [num_tokens, hidden_size]
+  auto gate_up = gate_up_proj_->forward(hidden_states);
+
+  if (is_smoothquant_) {
+    // For w8a8 quantization, the active operation is fused with the down_proj
+    return down_proj_->forward(gate_up);
+  } else {
+    torch::Tensor output;
+    if (Device::type_str() != "npu") {
+      int64_t batch_size = gate_up.sizes()[0];
+      output = torch::empty(
+          {batch_size, intermediate_size_ / process_group_->world_size()},
+          gate_up.options());
+    }
+
+    act_->forward(gate_up, output);
+    return down_proj_->forward(output);
+  }
+}
+
+void DenseMLPImpl::load_state_dict(const StateDict& state_dict) {
+  gate_up_proj_->load_state_dict(state_dict, {"gate_proj.", "up_proj."});
+  down_proj_->load_state_dict(state_dict.get_dict_with_prefix("down_proj."));
+}
+
+void DenseMLPImpl::load_state_dict(const StateDict& state_dict,
+                                   const std::vector<std::string>& gate_up_name,
+                                   const std::string& down_name) {
+  if (is_gated_) {
+    CHECK_EQ(gate_up_name.size(), 2);
+    gate_up_proj_->load_state_dict(state_dict, gate_up_name);
+  } else {
+    CHECK_EQ(gate_up_name.size(), 1);
+    gate_up_proj_->load_state_dict(
+        state_dict.get_dict_with_prefix(gate_up_name[0]));
+  }
+  down_proj_->load_state_dict(state_dict.get_dict_with_prefix(down_name));
+}
+
+std::optional<torch::Tensor> DenseMLPImpl::get_fp8_input_scale() const {
+  if (gate_up_proj_) {
+    return gate_up_proj_->get_input_scale();
+  }
+  return std::nullopt;
+}
+
+}  // namespace layer
+}  // namespace xllm

ex_engine/xllm_layers/common/dense_mlp.h

@@ -0,0 +1,67 @@
+/* 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 <torch/torch.h>
+
+#include "activation.h"
+#include "framework/model/model_args.h"
+#include "framework/parallel_state/parallel_args.h"
+#include "framework/quant_args.h"
+#include "framework/state_dict/state_dict.h"
+#include "linear.h"
+
+namespace xllm {
+namespace layer {
+
+class DenseMLPImpl : public torch::nn::Module {
+ public:
+  DenseMLPImpl() = default;
+  DenseMLPImpl(int64_t hidden_size,
+               int64_t intermediate_size,
+               bool is_gated,
+               bool has_bias,
+               const std::string& hidden_act,
+               bool enable_result_reduction,
+               const QuantArgs& quant_args,
+               ProcessGroup* process_group,
+               const torch::TensorOptions& options,
+               const std::string& module_prefix = "");
+
+  torch::Tensor forward(const torch::Tensor& hidden_states);
+
+  void load_state_dict(const StateDict& state_dict);
+  void load_state_dict(const StateDict& state_dict,
+                       const std::vector<std::string>& gate_up_name,
+                       const std::string& down_name);
+
+  // Get FP8 input scale from gate_up_proj for fused RMSNorm+FP8 quantization
+  std::optional<torch::Tensor> get_fp8_input_scale() const;
+
+ private:
+  bool is_gated_;
+  int64_t intermediate_size_;
+  ProcessGroup* process_group_;
+  ColumnParallelLinear gate_up_proj_{nullptr};
+  RowParallelLinear down_proj_{nullptr};
+  Activation act_{nullptr};
+  bool is_smoothquant_;
+  std::string hidden_act_;
+};
+TORCH_MODULE(DenseMLP);
+
+}  // namespace layer
+}  // namespace xllm

ex_engine/xllm_layers/common/fused_moe.cpp

@@ -0,0 +1,58 @@
+/* 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.
+==============================================================================*/
+
+#include "fused_moe.h"
+
+#include <glog/logging.h>
+
+namespace xllm {
+namespace layer {
+
+FusedMoEImpl::FusedMoEImpl(const ModelArgs& /*model_args*/,
+                           const FusedMoEArgs& /*moe_args*/,
+                           const QuantArgs& /*quant_args*/,
+                           const ParallelArgs& /*parallel_args*/,
+                           const torch::TensorOptions& /*options*/) {
+  NOT_IMPLEMENTED_WITH_MSG(
+      "FusedMoE is not supported for this backend. Please use CUDA, MLU or "
+      "ILU backend for MoE models.");
+}
+
+torch::Tensor FusedMoEImpl::forward_experts(
+    const torch::Tensor& /*hidden_states*/,
+    const torch::Tensor& /*router_logits*/,
+    bool /*enable_all2all_communication*/) {
+  NOT_IMPLEMENTED_WITH_MSG(
+      "FusedMoE is not supported for this backend. Please use CUDA, MLU or "
+      "ILU backend for MoE models.");
+  return torch::Tensor();
+}
+
+torch::Tensor FusedMoEImpl::forward(const torch::Tensor& /*hidden_states*/,
+                                    const ModelInputParams& /*input_params*/) {
+  NOT_IMPLEMENTED_WITH_MSG(
+      "FusedMoE is not supported for this backend. Please use CUDA, MLU or "
+      "ILU backend for MoE models.");
+  return torch::Tensor();
+}
+
+void FusedMoEImpl::load_state_dict(const StateDict& /*state_dict*/) {
+  NOT_IMPLEMENTED_WITH_MSG(
+      "FusedMoE is not supported for this backend. Please use CUDA, MLU or "
+      "ILU backend for MoE models.");
+}
+
+}  // namespace layer
+}  // namespace xllm

ex_engine/xllm_layers/common/fused_moe.h

@@ -0,0 +1,54 @@
+/* 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 <torch/torch.h>
+
+#include "dense_mlp.h"
+#include "framework/model/model_args.h"
+#include "framework/model/model_input_params.h"
+#include "framework/parallel_state/parallel_args.h"
+#include "framework/quant_args.h"
+#include "framework/state_dict/state_dict.h"
+#include "framework/state_dict/utils.h"
+#include "fused_moe_base.h"
+#include "linear.h"
+
+namespace xllm {
+namespace layer {
+
+// FusedMoE common implementation - placeholder for unsupported backends
+// Actual implementations are in backend-specific fused_moe.h files.
+class FusedMoEImpl : public torch::nn::Module {
+ public:
+  FusedMoEImpl() = default;
+  FusedMoEImpl(const ModelArgs& model_args,
+               const FusedMoEArgs& moe_args,
+               const QuantArgs& quant_args,
+               const ParallelArgs& parallel_args,
+               const torch::TensorOptions& options);
+
+  torch::Tensor forward_experts(const torch::Tensor& hidden_states,
+                                const torch::Tensor& router_logits,
+                                bool enable_all2all_communication);
+  torch::Tensor forward(const torch::Tensor& hidden_states,
+                        const ModelInputParams& input_params);
+  void load_state_dict(const StateDict& state_dict);
+};
+TORCH_MODULE(FusedMoE);
+
+}  // namespace layer
+}  // namespace xllm

ex_engine/xllm_layers/common/rms_norm.cpp

@@ -0,0 +1,144 @@
+/* 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.
+==============================================================================*/
+
+#include "rms_norm.h"
+
+#include <glog/logging.h>
+
+#include "kernels/ops_api.h"
+#include "platform/device.h"
+
+namespace xllm {
+namespace layer {
+
+const static std::string kLayerNormMode = "layernorm";
+const static std::string kRmsNormMode = "rmsnorm";
+
+RMSNormImpl::RMSNormImpl(int64_t dim,
+                         double eps,
+                         const torch::TensorOptions& options)
+    : norm_dim_(dim), eps_(eps), mode_(kRmsNormMode) {
+  weight_ = register_parameter("weight",
+                               torch::empty({dim}, options),
+                               /*requires_grad=*/false);
+}
+
+RMSNormImpl::RMSNormImpl(const ModelContext& context)
+    : RMSNormImpl(context.get_model_args().hidden_size(),
+                  context.get_model_args().rms_norm_eps(),
+                  context.get_tensor_options()) {}
+
+std::tuple<torch::Tensor, std::optional<torch::Tensor>> RMSNormImpl::forward(
+    torch::Tensor& input,
+    std::optional<torch::Tensor> residual,
+    std::optional<torch::Tensor> inplace_output) {
+  auto org_shape = input.sizes().vec();
+  input = input.reshape({-1, norm_dim_});
+
+  torch::Tensor output;
+  if (Device::type_str() != "npu") {
+    if (inplace_output.has_value()) {
+      output = inplace_output.value();
+      output = output.reshape({-1, norm_dim_});
+    } else {
+      output = torch::empty_like(input);
+    }
+  }
+
+  std::optional<torch::Tensor> residual_out;
+  if (residual.has_value()) {
+    residual.value() = residual.value().reshape({-1, norm_dim_});
+    if (Device::type_str() == "mlu" || Device::type_str() == "ilu") {
+      residual_out = residual.value();
+    }
+  }
+
+  xllm::kernel::FusedLayerNormParams fused_layernorm_params;
+  fused_layernorm_params.input = input;
+  fused_layernorm_params.residual = residual;
+  fused_layernorm_params.output = output;
+  fused_layernorm_params.residual_out = residual_out;
+  fused_layernorm_params.weight = weight_;
+  fused_layernorm_params.eps = eps_;
+  fused_layernorm_params.mode = mode_;
+  fused_layernorm_params.store_output_before_norm = residual_out.has_value();
+  if (bias_.defined()) {
+    fused_layernorm_params.beta = bias_;
+  }
+
+  xllm::kernel::fused_layernorm(fused_layernorm_params);
+
+  output = fused_layernorm_params.output;
+  residual_out = fused_layernorm_params.residual_out;
+
+  output = output.view(org_shape);
+  if (residual_out.has_value()) {
+    residual_out.value() = residual_out.value().view(org_shape);
+  }
+  return std::make_tuple(output, residual_out);
+}
+
+std::tuple<torch::Tensor, std::optional<torch::Tensor>>
+RMSNormImpl::forward_fp8(torch::Tensor& input,
+                         const torch::Tensor& fp8_scale,
+                         std::optional<torch::Tensor> residual) {
+  // Only supported on CUDA for now
+  CHECK(Device::type_str() == "cuda")
+      << "forward_fp8 is only supported on CUDA";
+  CHECK(mode_ == kRmsNormMode)
+      << "forward_fp8 only supports RMSNorm mode, not LayerNorm";
+
+  if (residual.has_value()) {
+    // Fused Add + RMSNorm + FP8 Quantization
+    xllm::kernel::FusedAddRmsNormStaticFp8QuantParams params;
+    params.input = input;
+    params.residual = residual.value();
+    params.weight = weight_;
+    params.scale = fp8_scale;
+    params.epsilon = eps_;
+
+    auto [output, updated_residual] =
+        xllm::kernel::fused_add_rms_norm_static_fp8_quant(params);
+
+    return std::make_tuple(output, updated_residual);
+  } else {
+    // RMSNorm + FP8 Quantization (no residual)
+    xllm::kernel::RmsNormStaticFp8QuantParams params;
+    params.input = input;
+    params.weight = weight_;
+    params.scale = fp8_scale;
+    params.epsilon = eps_;
+
+    auto output = xllm::kernel::rms_norm_static_fp8_quant(params);
+
+    return std::make_tuple(output, std::nullopt);
+  }
+}
+
+void RMSNormImpl::load_state_dict(const StateDict& state_dict) {
+  LOAD_WEIGHT(weight);
+  if (bias_.defined()) {
+    LOAD_WEIGHT(bias);
+  }
+}
+
+void RMSNormImpl::set_layernorm_mode() {
+  mode_ = kLayerNormMode;
+  bias_ = register_parameter(
+      "bias", torch::empty({norm_dim_}, weight_.options()), false);
+}
+
+}  // namespace layer
+}  // namespace xllm

ex_engine/xllm_layers/common/rms_norm.h

@@ -0,0 +1,64 @@
+/* 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 <torch/torch.h>
+
+#include "core/framework/model_context.h"
+#include "framework/state_dict/state_dict.h"
+#include "framework/state_dict/utils.h"
+
+namespace xllm {
+namespace layer {
+
+class RMSNormImpl : public torch::nn::Module {
+ public:
+  RMSNormImpl(int64_t dim, double eps, const torch::TensorOptions& options);
+  RMSNormImpl(const ModelContext& context);
+
+  // Standard forward: returns (normalized_output, updated_residual)
+  std::tuple<torch::Tensor, std::optional<torch::Tensor>> forward(
+      torch::Tensor& input,
+      std::optional<torch::Tensor> residual = std::nullopt,
+      std::optional<torch::Tensor> inplace_output = std::nullopt);
+
+  // Fused forward with FP8 quantization output (for static quantization)
+  // Returns: (fp8_quantized_output, updated_residual)
+  // This combines RMSNorm + FP8 quantization to reduce memory bandwidth
+  std::tuple<torch::Tensor, std::optional<torch::Tensor>> forward_fp8(
+      torch::Tensor& input,
+      const torch::Tensor& fp8_scale,
+      std::optional<torch::Tensor> residual = std::nullopt);
+
+  void set_layernorm_mode();
+
+  void load_state_dict(const StateDict& state_dict);
+
+  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_;
+  std::string mode_;
+};
+TORCH_MODULE(RMSNorm);
+
+}  // namespace layer
+}  // namespace xllm

ex_engine/xllm_layers/common/rotary_embedding.cpp

@@ -0,0 +1,307 @@
+/* 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.
+==============================================================================*/
+
+#include "rotary_embedding.h"
+
+#include "kernels/ops_api.h"
+#include "platform/device.h"
+
+namespace xllm {
+namespace layer {
+
+RotaryEmbeddingImpl::RotaryEmbeddingImpl(const ModelContext& context) {
+  LOG(FATAL) << "Not implement currently.";
+}
+
+RotaryEmbeddingImpl::RotaryEmbeddingImpl(int64_t rotary_dim,
+                                         int64_t max_position_embeddings,
+                                         int64_t rope_theta,
+                                         bool interleaved,
+                                         const torch::TensorOptions& options)
+    : interleaved_(interleaved) {
+  auto inv_freq = rotary::compute_inv_freq(rotary_dim, rope_theta, options);
+  const auto cos_sin = rotary::compute_cos_sin_cache(
+      rotary_dim, max_position_embeddings, interleaved, inv_freq, options);
+  cos_sin_cache_ = register_buffer("cos_sin_cache", cos_sin);
+
+  auto cos_sin_vec = cos_sin_cache_.chunk(2, /*dim=*/-1);
+  cos_ = cos_sin_vec[0].view({-1, rotary_dim});
+  sin_ = cos_sin_vec[1].view({-1, rotary_dim});
+
+  // Pre-compute [cos_half, sin_half] format used by the CUDA/ILU/MUSA kernels.
+  const auto dev = Device::type_str();
+  if (dev == "cuda" || dev == "ilu" || dev == "musa") {
+    auto chunks = cos_sin_cache_.chunk(4, -1);
+    precomputed_cos_sin_cache_ =
+        torch::cat({chunks[0], chunks[2]}, -1).contiguous();
+  }
+}
+
+void RotaryEmbeddingImpl::forward(torch::Tensor& q,
+                                  torch::Tensor& k,
+                                  const torch::Tensor& positions,
+                                  const torch::Tensor& cu_query_lens,
+                                  int64_t max_query_len,
+                                  bool is_prompt) {
+  bool discrete;
+  std::optional<torch::Tensor> position_ids;
+  if (is_prompt) {
+    discrete = false;
+    if (Device::type_str() == "cuda" || Device::type_str() == "npu" ||
+        Device::type_str() == "ilu" || Device::type_str() == "musa") {
+      position_ids = positions;
+    }
+  } else {
+    discrete = true;
+    position_ids = positions;
+  }
+
+  xllm::kernel::RotaryParams rotary_params;
+  rotary_params.q = q;
+  rotary_params.k = k;
+  rotary_params.sin = sin_;
+  rotary_params.cos = cos_;
+  rotary_params.cos_sin = cos_sin_cache_;
+  rotary_params.precomputed_cos_sin = precomputed_cos_sin_cache_;
+  rotary_params.position_ids = position_ids;
+  rotary_params.cu_query_lens = cu_query_lens;
+  rotary_params.interleaved = interleaved_;
+  rotary_params.discrete = discrete;
+  rotary_params.max_query_len = max_query_len;
+  xllm::kernel::apply_rotary(rotary_params);
+
+  q = rotary_params.q;
+  k = rotary_params.k;
+}
+
+// Single tensor forward for MLA architecture
+void RotaryEmbeddingImpl::forward(torch::Tensor& input,
+                                  const torch::Tensor& positions,
+                                  const torch::Tensor& cu_query_lens,
+                                  int64_t max_query_len,
+                                  bool is_prompt) {
+  bool discrete;
+  std::optional<torch::Tensor> position_ids;
+  if (is_prompt) {
+    discrete = false;
+    if (Device::type_str() == "cuda" || Device::type_str() == "npu" ||
+        Device::type_str() == "ilu") {
+      position_ids = positions;
+    }
+  } else {
+    discrete = true;
+    position_ids = positions;
+  }
+
+  xllm::kernel::RotaryParams rotary_params;
+  rotary_params.q = input;
+  rotary_params.sin = sin_;
+  rotary_params.cos = cos_;
+  rotary_params.cos_sin = cos_sin_cache_;
+  rotary_params.position_ids = position_ids;
+  rotary_params.cu_query_lens = cu_query_lens;
+  rotary_params.interleaved = interleaved_;
+  rotary_params.discrete = discrete;
+  rotary_params.max_query_len = max_query_len;
+  xllm::kernel::apply_rotary(rotary_params);
+
+  input = rotary_params.q;
+}
+
+MRotaryEmbeddingImpl::MRotaryEmbeddingImpl(
+    int64_t rotary_dim,
+    int64_t max_position_embeddings,
+    int64_t rope_theta,
+    bool interleaved,
+    const std::vector<int64_t>& rope_scaling_mrope_section,
+    const torch::TensorOptions& options)
+    : RotaryEmbeddingImpl(rotary_dim,
+                          max_position_embeddings,
+                          rope_theta,
+                          interleaved,
+                          options),
+      mrope_section_(rope_scaling_mrope_section) {
+  mrope_cu_seq_lens_ = torch::zeros(2, torch::kInt32).to(options.device());
+}
+
+void MRotaryEmbeddingImpl::forward(torch::Tensor& q,
+                                   torch::Tensor& k,
+                                   const torch::Tensor& positions,
+                                   const AttentionMetadata& attn_metadata) {
+  bool only_prefill =
+      (attn_metadata.is_prefill || attn_metadata.is_chunked_prefill);
+  if (!only_prefill || mrope_section_.empty()) {
+    torch::Tensor position_ids = positions;
+    if (positions.dim() == 2) {
+      position_ids = positions[0];
+    }
+    return RotaryEmbeddingImpl::forward(q,
+                                        k,
+                                        position_ids,
+                                        attn_metadata.q_cu_seq_lens,
+                                        attn_metadata.max_query_len,
+                                        attn_metadata.is_prefill);
+  }
+
+  int64_t num_tokens = positions.size(-1);
+  mrope_cu_seq_lens_[1] = num_tokens;
+  CHECK(attn_metadata.mrope_cos.defined() && attn_metadata.mrope_sin.defined());
+  xllm::kernel::RotaryParams rotary_params;
+  rotary_params.q = q;
+  rotary_params.k = k;
+  rotary_params.sin = attn_metadata.mrope_sin;
+  rotary_params.cos = attn_metadata.mrope_cos;
+  rotary_params.cos_sin = cos_sin_cache_;
+  rotary_params.precomputed_cos_sin = precomputed_cos_sin_cache_;
+  rotary_params.position_ids = std::nullopt;
+  rotary_params.cu_query_lens = mrope_cu_seq_lens_;
+  rotary_params.interleaved = interleaved_;
+  rotary_params.discrete = false;
+  rotary_params.max_query_len = num_tokens;
+  xllm::kernel::apply_rotary(rotary_params);
+
+  q = rotary_params.q;
+  k = rotary_params.k;
+}
+
+DeepseekScalingRotaryEmbeddingImpl::DeepseekScalingRotaryEmbeddingImpl(
+    int64_t head_size,
+    int64_t rotary_dim,
+    int64_t max_position_embeddings,
+    int64_t rope_scaling_original_max_position_embeddings,
+    int64_t rope_theta,
+    bool interleaved,
+    float scaling_factor,
+    float extrapolation_factor,
+    float attn_factor,
+    float beta_fast,
+    float beta_slow,
+    float mscale,
+    float mscale_all_dim,
+    const torch::TensorOptions& options)
+    : head_size_(head_size),
+      rotary_dim_(rotary_dim),
+      interleaved_(interleaved) {
+  auto inv_freq = rotary::apply_deepseek_yarn_rope_scaling(
+      scaling_factor,
+      extrapolation_factor,
+      beta_fast,
+      beta_slow,
+      rotary_dim,
+      rope_theta,
+      rope_scaling_original_max_position_embeddings);
+  const auto cos_sin = rotary::compute_cos_sin_cache(rotary_dim,
+                                                     max_position_embeddings,
+                                                     interleaved,
+                                                     scaling_factor,
+                                                     attn_factor,
+                                                     mscale,
+                                                     mscale_all_dim,
+                                                     inv_freq,
+                                                     options);
+  cos_sin_cache_ = register_buffer("cos_sin_cache", cos_sin);
+
+  auto cos_sin_vec = cos_sin_cache_.chunk(2, /*dim=*/-1);
+  cos_ = cos_sin_vec[0].view({-1, rotary_dim});
+  sin_ = cos_sin_vec[1].view({-1, rotary_dim});
+
+  // Pre-compute [cos_half, sin_half] format used by the CUDA/ILU/MUSA kernels.
+  const auto dev = Device::type_str();
+  if (dev == "cuda" || dev == "ilu" || dev == "musa") {
+    auto chunks = cos_sin_cache_.chunk(4, -1);
+    precomputed_cos_sin_cache_ =
+        torch::cat({chunks[0], chunks[2]}, -1).contiguous();
+  }
+}
+
+void DeepseekScalingRotaryEmbeddingImpl::forward(
+    torch::Tensor& input,
+    const torch::Tensor& positions,
+    const torch::Tensor& cu_query_lens,
+    int64_t max_query_len,
+    bool is_prompt) {
+  const int32_t dim = -1;
+  bool discrete;
+  std::optional<torch::Tensor> position_ids;
+  if (is_prompt) {
+    discrete = false;
+    position_ids = std::nullopt;
+  } else {
+    discrete = true;
+    position_ids = positions;
+    max_query_len = 1;
+  }
+  auto input_rot = input.slice(dim, 0, rotary_dim_);
+  torch::Tensor input_pass;
+  if (rotary_dim_ < head_size_) {
+    input_pass = input.slice(dim, rotary_dim_, head_size_);
+  }
+
+  xllm::kernel::RotaryParams rotary_params;
+  rotary_params.q = input_rot;
+  rotary_params.sin = sin_;
+  rotary_params.cos = cos_;
+  rotary_params.cos_sin = cos_sin_cache_;
+  rotary_params.precomputed_cos_sin = precomputed_cos_sin_cache_;
+  rotary_params.position_ids = position_ids;
+  rotary_params.cu_query_lens = cu_query_lens;
+  rotary_params.interleaved = interleaved_;
+  rotary_params.discrete = discrete;
+  rotary_params.max_query_len = max_query_len;
+  xllm::kernel::apply_rotary(rotary_params);
+  input_rot = rotary_params.q;
+
+  if (rotary_dim_ < head_size_) {
+    input = torch::cat({input_rot, input_pass}, dim);
+  } else {
+    input = input_rot;
+  }
+}
+
+// Factory function: creates the appropriate RoPE type based on model args
+std::shared_ptr<RotaryEmbeddingBase> create_mla_rotary_embedding(
+    const ModelArgs& args,
+    int64_t rotary_dim,
+    int64_t max_position_embeddings,
+    bool interleaved,
+    const torch::TensorOptions& options) {
+  if (args.rope_scaling_rope_type() == "deepseek_yarn") {
+    return std::make_shared<DeepseekScalingRotaryEmbeddingImpl>(
+        rotary_dim,  // head_size (same as rotary_dim for MLA)
+        rotary_dim,
+        max_position_embeddings,
+        args.rope_scaling_original_max_position_embeddings(),
+        args.rope_theta(),
+        interleaved,
+        args.rope_scaling_factor(),
+        args.rope_extrapolation_factor(),
+        args.rope_scaling_attn_factor(),
+        args.rope_scaling_beta_fast(),
+        args.rope_scaling_beta_slow(),
+        args.rope_scaling_mscale(),
+        args.rope_scaling_mscale_all_dim(),
+        options);
+  } else {
+    // default rope type
+    return std::make_shared<RotaryEmbeddingImpl>(rotary_dim,
+                                                 max_position_embeddings,
+                                                 args.rope_theta(),
+                                                 interleaved,
+                                                 options);
+  }
+}
+
+}  // namespace layer
+}  // namespace xllm

ex_engine/xllm_layers/common/rotary_embedding.h

@@ -0,0 +1,158 @@
+/* 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 <torch/torch.h>
+#include <torch/types.h>
+
+#include <memory>
+
+#include "attention_metadata.h"
+#include "core/framework/model_context.h"
+#include "framework/model/model_args.h"
+#include "rotary_embedding_util.h"
+
+namespace xllm {
+namespace layer {
+
+class RotaryEmbeddingBase : public torch::nn::Module {
+ public:
+  ~RotaryEmbeddingBase() override = default;
+
+  virtual void forward(torch::Tensor& input,
+                       const torch::Tensor& positions,
+                       const torch::Tensor& cu_query_lens,
+                       int64_t max_query_len,
+                       bool is_prompt) = 0;
+  virtual const torch::Tensor& get_sin_cache() const = 0;
+  virtual const torch::Tensor& get_cos_cache() const = 0;
+  virtual const bool get_interleaved() const = 0;
+};
+
+class RotaryEmbeddingImpl : public RotaryEmbeddingBase {
+ public:
+  RotaryEmbeddingImpl(int64_t rotary_dim,
+                      int64_t max_position_embeddings,
+                      int64_t rope_theta,
+                      bool interleaved,
+                      const torch::TensorOptions& options);
+  RotaryEmbeddingImpl(const ModelContext& context);
+
+  void forward(torch::Tensor& q,
+               torch::Tensor& k,
+               const torch::Tensor& positions,
+               const torch::Tensor& cu_query_lens,
+               int64_t max_query_len,
+               bool is_prompt);
+  // Single tensor forward for MLA architecture
+  void forward(torch::Tensor& input,
+               const torch::Tensor& positions,
+               const torch::Tensor& cu_query_lens,
+               int64_t max_query_len,
+               bool is_prompt) override;
+
+  const torch::Tensor& precomputed_cos_sin_cache() {
+    return precomputed_cos_sin_cache_;
+  }
+
+  torch::Tensor get_cos_sin_cache() { return cos_sin_cache_; }
+  const torch::Tensor& get_sin_cache() const override { return sin_; }
+  const torch::Tensor& get_cos_cache() const override { return cos_; }
+  const bool get_interleaved() const override { return interleaved_; }
+
+ protected:
+  bool interleaved_;
+  torch::Tensor cos_sin_cache_;
+  // Pre-formatted [cos_half, sin_half] cache for CUDA/MUSA/ILU kernels.
+  // Avoids chunk/cat operations on every forward call.
+  torch::Tensor precomputed_cos_sin_cache_;
+
+ private:
+  torch::Tensor sin_;
+  torch::Tensor cos_;
+};
+TORCH_MODULE(RotaryEmbedding);
+
+class MRotaryEmbeddingImpl : public RotaryEmbeddingImpl {
+ public:
+  MRotaryEmbeddingImpl(int64_t rotary_dim,
+                       int64_t max_position_embeddings,
+                       int64_t rope_theta,
+                       bool interleaved,
+                       const std::vector<int64_t>& rope_scaling_mrope_section,
+                       const torch::TensorOptions& options);
+
+  void forward(torch::Tensor& q,
+               torch::Tensor& k,
+               const torch::Tensor& positions,
+               const AttentionMetadata& attn_metadata);
+
+ private:
+  std::vector<int64_t> mrope_section_;
+  torch::Tensor mrope_cu_seq_lens_;
+};
+TORCH_MODULE(MRotaryEmbedding);
+
+class DeepseekScalingRotaryEmbeddingImpl : public RotaryEmbeddingBase {
+ public:
+  DeepseekScalingRotaryEmbeddingImpl(
+      int64_t head_size,
+      int64_t rotary_dim,
+      int64_t max_position_embeddings,
+      int64_t rope_scaling_original_max_position_embeddings,
+      int64_t rope_theta,
+      bool interleaved,
+      float scaling_factor,
+      float extrapolation_factor,
+      float attn_factor,
+      float beta_fast,
+      float beta_slow,
+      float mscale,
+      float mscale_all_dim,
+      const torch::TensorOptions& options);
+
+  void forward(torch::Tensor& input,
+               const torch::Tensor& positions,
+               const torch::Tensor& cu_query_lens,
+               int64_t max_query_len,
+               bool is_prompt) override;
+  const torch::Tensor& get_sin_cache() const override { return sin_; }
+  const torch::Tensor& get_cos_cache() const override { return cos_; }
+  const bool get_interleaved() const override { return interleaved_; }
+
+ private:
+  int64_t head_size_;
+  int64_t rotary_dim_;
+  bool interleaved_;
+  torch::Tensor sin_;
+  torch::Tensor cos_;
+  torch::Tensor cos_sin_cache_;
+  // Pre-formatted [cos_half, sin_half] cache for CUDA/MUSA/ILU kernels.
+  // Avoids chunk/cat operations on every forward call.
+  torch::Tensor precomputed_cos_sin_cache_;
+};
+TORCH_MODULE(DeepseekScalingRotaryEmbedding);
+
+// Factory function: creates the appropriate RoPE type based on model args
+std::shared_ptr<RotaryEmbeddingBase> create_mla_rotary_embedding(
+    const ModelArgs& args,
+    int64_t rotary_dim,
+    int64_t max_position_embeddings,
+    bool interleaved,
+    const torch::TensorOptions& options);
+
+}  // namespace layer
+}  // namespace xllm

ex_engine/xllm_layers/ilu/attention.cpp

@@ -0,0 +1,189 @@
+/* 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.
+==============================================================================*/
+
+#include "attention.h"
+
+#include "kernels/ilu/ilu_ops_api.h"
+#include "kernels/ops_api.h"
+
+namespace xllm {
+namespace layer {
+AttentionImpl::AttentionImpl(int64_t num_heads,
+                             int64_t head_size,
+                             float scale,
+                             int64_t num_kv_heads,
+                             int64_t sliding_window)
+    : num_heads_(num_heads),
+      head_size_(head_size),
+      scale_(scale),
+      num_kv_heads_(num_kv_heads),
+      v_head_dim_(head_size),
+      use_fused_mla_qkv_(false),
+      enable_lighting_indexer_(false),
+      enable_mla_(false),
+      sliding_window_(sliding_window) {
+  if (sliding_window_ > -1) {
+    sliding_window_ = sliding_window_ - 1;
+  }
+}
+
+AttentionImpl::AttentionImpl(int64_t num_heads,
+                             int64_t head_size,
+                             int64_t num_kv_heads,
+                             int64_t v_head_dim,
+                             int64_t sliding_window,
+                             float scale,
+                             bool use_fused_mla_qkv,
+                             bool enable_lighting_indexer,
+                             bool enable_mla)
+    : num_heads_(num_heads),
+      head_size_(head_size),
+      scale_(scale),
+      num_kv_heads_(num_kv_heads),
+      v_head_dim_(v_head_dim),
+      use_fused_mla_qkv_(use_fused_mla_qkv),
+      enable_lighting_indexer_(enable_lighting_indexer),
+      enable_mla_(enable_mla),
+      sliding_window_(sliding_window) {
+  if (sliding_window_ > -1) {
+    sliding_window_ = sliding_window_ - 1;
+  }
+}
+
+std::tuple<torch::Tensor, std::optional<torch::Tensor>> AttentionImpl::forward(
+    const AttentionMetadata& attn_metadata,
+    torch::Tensor& query,
+    torch::Tensor& key,
+    torch::Tensor& value,
+    KVCache& kv_cache) {
+  std::optional<torch::Tensor> output_lse = std::nullopt;
+  torch::Tensor output;
+  if (enable_mla_) {
+    output = torch::empty({query.size(0), num_heads_ * v_head_dim_},
+                          query.options());
+  } else {
+    output = torch::empty_like(query);
+  }
+  if (attn_metadata.is_dummy) {
+    return std::make_tuple(output, output_lse);
+  }
+
+  bool only_prefill =
+      attn_metadata.is_prefill || attn_metadata.is_chunked_prefill;
+  int64_t num_kv_heads = (enable_mla_ && !only_prefill) ? 1 : num_kv_heads_;
+  torch::Tensor k_cache = kv_cache.get_k_cache();
+  std::optional<torch::Tensor> v_cache;
+  std::optional<torch::Tensor> v;
+  if (!enable_mla_) {
+    v = value.view({-1, num_kv_heads, head_size_});
+    v_cache = kv_cache.get_v_cache();
+  }
+
+  bool skip_process_cache = enable_mla_ && (only_prefill || use_fused_mla_qkv_);
+  if (!skip_process_cache) {
+    xllm::kernel::ReshapePagedCacheParams reshape_paged_cache_params;
+    reshape_paged_cache_params.key = key.view({-1, num_kv_heads, head_size_});
+    reshape_paged_cache_params.value = v;
+    reshape_paged_cache_params.k_cache = k_cache;
+    reshape_paged_cache_params.v_cache = v_cache;
+    reshape_paged_cache_params.slot_mapping = attn_metadata.slot_mapping;
+    xllm::kernel::reshape_paged_cache(reshape_paged_cache_params);
+  }
+
+  if (enable_lighting_indexer_ || !only_prefill) {
+    decoder_forward(query, output, k_cache, v_cache, attn_metadata);
+  } else {
+    prefill_forward(query, key, value, output, k_cache, v_cache, attn_metadata);
+  }
+
+  int64_t head_size = enable_mla_ ? v_head_dim_ : head_size_;
+  output = output.view({-1, num_heads_ * head_size});
+  return {output, output_lse};
+}
+
+void AttentionImpl::prefill_forward(torch::Tensor& query,
+                                    torch::Tensor& key,
+                                    torch::Tensor& value,
+                                    torch::Tensor& output,
+                                    const torch::Tensor& k_cache,
+                                    const std::optional<torch::Tensor>& v_cache,
+                                    const AttentionMetadata& attn_metadata) {
+  int64_t head_size_v = enable_mla_ ? v_head_dim_ : head_size_;
+  std::optional<torch::Tensor> output_lse = std::nullopt;
+  query = query.view({-1, num_heads_, head_size_});
+  output = output.view({-1, num_heads_, head_size_v});
+  // torch::Tensor k_cache_ = k_cache;
+  // torch::Tensor v_cache_ = v_cache.value();
+  xllm::kernel::ilu::batch_prefill(query,
+                                   k_cache,
+                                   v_cache,
+                                   output,
+                                   output_lse,
+                                   attn_metadata.q_cu_seq_lens,
+                                   attn_metadata.kv_cu_seq_lens,
+                                   /*alibi_slope=*/std::nullopt,
+                                   /*attn_bias=*/std::nullopt,
+                                   /*q_quant_scale=*/std::nullopt,
+                                   /*k_quant_scale=*/std::nullopt,
+                                   /*v_quant_scale=*/std::nullopt,
+                                   attn_metadata.block_table,
+                                   attn_metadata.max_query_len,
+                                   attn_metadata.max_seq_len,
+                                   scale_,
+                                   attn_metadata.is_causal,
+                                   sliding_window_,
+                                   /*window_size_right=*/-1,
+                                   attn_metadata.compute_dtype,
+                                   /*return_lse=*/false);
+}
+
+void AttentionImpl::decoder_forward(torch::Tensor& query,
+                                    torch::Tensor& output,
+                                    const torch::Tensor& k_cache,
+                                    const std::optional<torch::Tensor>& v_cache,
+                                    const AttentionMetadata& attn_metadata) {
+  int64_t head_size_v = enable_mla_ ? v_head_dim_ : head_size_;
+  query = query.view({-1, 1, num_heads_, head_size_});
+  output = output.view({-1, 1, num_heads_, head_size_v});
+  std::optional<torch::Tensor> output_lse = std::nullopt;
+
+  int64_t block_aligned_max_seq_len =
+      attn_metadata.block_table.size(-1) * k_cache.size(2);
+
+  xllm::kernel::ilu::batch_decode(query,
+                                  k_cache,
+                                  output,
+                                  attn_metadata.block_table,
+                                  attn_metadata.kv_seq_lens,
+                                  v_cache,
+                                  output_lse,
+                                  /*q_quant_scale=*/std::nullopt,
+                                  /*k_quant_scale=*/std::nullopt,
+                                  /*v_quant_scale=*/std::nullopt,
+                                  /*out_quant_scale=*/std::nullopt,
+                                  /*alibi_slope=*/std::nullopt,
+                                  attn_metadata.attn_mask,
+                                  attn_metadata.compute_dtype,
+                                  block_aligned_max_seq_len,
+                                  sliding_window_,
+                                  /*window_size_right=*/-1,
+                                  scale_,
+                                  /*return_lse=*/false,
+                                  attn_metadata.is_causal,
+                                  /*kv_cache_quant_bit_size=*/-1);
+}
+
+}  // namespace layer
+}  // namespace xllm

ex_engine/xllm_layers/ilu/attention.h

@@ -0,0 +1,82 @@
+/* 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 <torch/torch.h>
+
+#include <tuple>
+
+#include "framework/kv_cache/kv_cache.h"
+#include "framework/model/model_input_params.h"
+#include "layers/common/attention_metadata.h"
+
+namespace xllm {
+namespace layer {
+class AttentionImpl : public torch::nn::Module {
+ public:
+  AttentionImpl() = default;
+
+  AttentionImpl(int64_t num_heads,
+                int64_t head_size,
+                float scale,
+                int64_t num_kv_heads,
+                int64_t sliding_window);
+  AttentionImpl(int64_t num_heads,
+                int64_t head_size,
+                int64_t num_kv_heads,
+                int64_t v_head_dim,
+                int64_t sliding_window,
+                float scale,
+                bool use_fused_mla_qkv,
+                bool enable_lighting_indexer,
+                bool enable_mla);
+
+  std::tuple<torch::Tensor, std::optional<torch::Tensor>> forward(
+      const AttentionMetadata& attn_metadata,
+      torch::Tensor& query,
+      torch::Tensor& key,
+      torch::Tensor& value,
+      KVCache& kv_cache);
+
+  void prefill_forward(torch::Tensor& query,
+                       torch::Tensor& key,
+                       torch::Tensor& value,
+                       torch::Tensor& output,
+                       const torch::Tensor& k_cache,
+                       const std::optional<torch::Tensor>& v_cache,
+                       const AttentionMetadata& attn_metadata);
+
+  void decoder_forward(torch::Tensor& query,
+                       torch::Tensor& output,
+                       const torch::Tensor& k_cache,
+                       const std::optional<torch::Tensor>& v_cache,
+                       const AttentionMetadata& attn_metadata);
+
+ private:
+  int64_t num_heads_;
+  int64_t head_size_;
+  float scale_;
+  int64_t num_kv_heads_;
+  int64_t v_head_dim_;
+  bool use_fused_mla_qkv_;
+  bool enable_lighting_indexer_;
+  bool enable_mla_;
+  int64_t sliding_window_;
+};
+TORCH_MODULE(Attention);
+
+}  // namespace layer
+}  // namespace xllm

ex_engine/xllm_layers/ilu/fused_moe.cpp

@@ -0,0 +1,797 @@
+/* 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 "fused_moe.h"
+
+#include <glog/logging.h>
+
+#include <iomanip>
+
+#include "common/global_flags.h"
+#include "framework/parallel_state/parallel_state.h"
+#include "kernels/ops_api.h"
+#include "layers/common/dp_utils.h"
+#include "util/utils.h"
+
+namespace {
+
+int32_t get_dtype_size(torch::ScalarType dtype) {
+  return static_cast<int32_t>(torch::elementSize(dtype));
+}
+
+}  // namespace
+
+namespace xllm {
+namespace layer {
+
+FusedMoEImpl::FusedMoEImpl(const ModelArgs& model_args,
+                           const FusedMoEArgs& moe_args,
+                           const QuantArgs& quant_args,
+                           const ParallelArgs& parallel_args,
+                           const torch::TensorOptions& options)
+    : num_total_experts_(static_cast<int64_t>(model_args.n_routed_experts())),
+      topk_(model_args.num_experts_per_tok()),
+      num_expert_group_(model_args.n_group()),
+      topk_group_(model_args.topk_group()),
+      route_scale_(model_args.routed_scaling_factor()),
+      hidden_size_(model_args.hidden_size()),
+      n_shared_experts_(model_args.n_shared_experts()),
+      is_gated_(moe_args.is_gated),
+      renormalize_(model_args.norm_topk_prob() ? 1 : 0),
+      hidden_act_(model_args.hidden_act()),
+      scoring_func_(model_args.scoring_func()),
+      quant_args_(quant_args),
+      parallel_args_(parallel_args),
+      options_(options),
+      device_(options.device()) {
+  const int64_t num_experts = num_total_experts_;
+  const int64_t intermediate_size =
+      static_cast<int64_t>(model_args.moe_intermediate_size());
+  const std::string& topk_method = model_args.topk_method();
+  int64_t ep_size = parallel_args.ep_size();
+  int64_t ep_rank = 0;
+  tp_pg_ = parallel_args.tp_group_;
+  if (ep_size > 1) {
+    ep_rank = parallel_args.moe_ep_group_->rank();
+    tp_pg_ = parallel_args.moe_tp_group_;
+  }
+
+  // smoothquant check: If quant_method is not empty, only w8a8 smoothquant is
+  // supported
+  if (!quant_args.quant_method().empty()) {
+    if (quant_args.quant_method() != "smoothquant" || quant_args.bits() != 8 ||
+        !quant_args.activation_dynamic()) {
+      LOG(FATAL) << "FusedMoE only supports w8a8 smoothquant quantization when "
+                    "quant_method is set. "
+                 << "Got quant_method=" << quant_args.quant_method()
+                 << ", bits=" << quant_args.bits()
+                 << ", activation_dynamic=" << quant_args.activation_dynamic();
+    }
+    // If confirmed as smoothquant w8a8, set is_smoothquant_ to true
+    is_smoothquant_ = true;
+  } else {
+    is_smoothquant_ = false;
+  }
+
+  // Deep EP initialization check
+  enable_deep_ep_ = FLAGS_expert_parallel_degree == 2 && ep_size > 1;
+  if (enable_deep_ep_) {
+    // for now, we only implement the deep ep for decode stage.
+    // so we will assume the max_token_num is limited to max_batch_size * (1+K)
+    // K is the number of speculative tokens.
+    int64_t dispatch_token_size;
+    if (quant_args.quant_method() == "smoothquant") {
+      // float32 is for the scale of the quantized input
+      dispatch_token_size = hidden_size_ * get_dtype_size(torch::kInt8) +
+                            get_dtype_size(torch::kFloat32);
+    } else {
+      dispatch_token_size =
+          hidden_size_ * get_dtype_size(options_.dtype().toScalarType());
+    }
+    torch::ScalarType combine_dtype = options_.dtype().toScalarType();
+    int64_t combine_token_size = hidden_size_ * get_dtype_size(combine_dtype);
+    // Ensure calculation base is at least ep_size
+    int64_t effective_seqs =
+        std::max((int64_t)FLAGS_max_seqs_per_batch, (int64_t)ep_size);
+    // NOTE: FLAGS_max_seqs_per_batch represents the maximum total batch size,
+    // regardless of the dp size. To ensure robust scheduling and account
+    // for the worst-case scenario, we must guarantee that each rank is capable
+    // of handling the maximum possible number of tokens. Therefore, we define
+    // max_num_tokens_per_rank as the full maximum value, without dividing by
+    // either the rank count or the dp size.
+    int64_t max_num_tokens_per_rank =
+        (1 + FLAGS_num_speculative_tokens) * effective_seqs * topk_;
+
+    // make sure that all layers share the same deep ep instance
+    //  so that the memory footprint is minimized
+    deep_ep_ = DeepEPManager::get_instance(dispatch_token_size,
+                                           combine_token_size,
+                                           max_num_tokens_per_rank,
+                                           num_experts,
+                                           parallel_args,
+                                           options_);
+
+    // obtain the buffer and parameters of deep ep
+    deep_ep_buffer_ = deep_ep_->get_buffer();
+    deep_ep_params_ = deep_ep_->get_params();
+
+    // intermediate buffer that can be initialized once
+    // we place these tensor here in order to speed up forward pass
+    int64_t n_tokens_recv = deep_ep_params_.max_num_tokens_recv;
+    int64_t token_bytes = is_smoothquant_
+                              ? get_dtype_size(torch::kInt8)
+                              : get_dtype_size(options_.dtype().toScalarType());
+    token_bytes = token_bytes * hidden_size_;
+    int64_t head_size = n_tokens_recv * token_bytes;
+    dispatch_recv_token_tensor_head_ =
+        deep_ep_buffer_.combine_send_token_tensor.narrow(0, 0, head_size)
+            .view({n_tokens_recv, token_bytes});
+    // input scale in smoothquant
+    if (is_smoothquant_) {
+      int64_t tail_size = n_tokens_recv * get_dtype_size(torch::kFloat32);
+      dispatch_recv_token_tensor_tail_ =
+          deep_ep_buffer_.combine_send_token_tensor
+              .narrow(0, head_size, tail_size)
+              .view({n_tokens_recv, -1});
+    }
+  }
+
+  // calculate the number of experts per rank
+  num_experts_per_rank_ = num_experts / ep_size;
+  start_expert_id_ = ep_rank * num_experts_per_rank_;
+
+  if (topk_method == "noaux_tc") {
+    e_score_correction_bias_ = register_parameter(
+        "e_score_correction_bias", torch::empty({num_experts}, options), false);
+  }
+
+  gate_ = register_module(
+      "gate_proj",
+      ReplicatedLinear(hidden_size_, num_experts, false, quant_args, options));
+  if (n_shared_experts_ > 0) {
+    ProcessGroup* shared_expert_pg;
+    if (parallel_args_.ep_size() > 1) {
+      // we use tp=1 for shared experts computation in deep ep mode
+      CHECK(parallel_args_.ep_size() == parallel_args_.world_size())
+          << "Models with shared experts only support ep_size equal to "
+             "world size for now.";
+      shared_expert_pg = parallel_args.moe_tp_group_;
+    } else {
+      shared_expert_pg = parallel_args.process_group_;
+    }
+    // The shared experts computation can proceed in parallel with the
+    // final communication step during the MoE computation, as long as it
+    // remains independent of any communication operations. For optimal
+    // performance, ensure that the shared experts layer on each rank always
+    // maintains its own unique weights.
+    shared_experts_ =
+        register_module("shared_experts",
+                        DenseMLP(hidden_size_,
+                                 intermediate_size * n_shared_experts_,
+                                 is_gated_,
+                                 false,
+                                 hidden_act_,
+                                 /*enable_result_reduction=*/true,
+                                 quant_args,
+                                 shared_expert_pg,
+                                 options));
+  }
+
+  // create weight buffer
+  const int64_t world_size = tp_pg_->world_size();
+  int64_t local_intermediate_size = intermediate_size / world_size;
+  if (is_smoothquant_) {
+    auto quant_option = options_.dtype(torch::kInt8);
+    auto fp_option = options_.dtype(torch::kFloat32);
+    w13_ = register_parameter(
+        "w13",
+        torch::empty(
+            {num_experts_per_rank_, local_intermediate_size * 2, hidden_size_},
+            quant_option),
+        false);
+    w13_scale_ = register_parameter(
+        "w13_scale",
+        torch::empty({num_experts_per_rank_, local_intermediate_size * 2},
+                     fp_option),
+        false);
+    // Note: We do not check enable_deep_ep_ here, since smooth quantization
+    // information may be needed even when deep EP mode is disabled. This allows
+    // retrieving quantization parameters for any subset of experts as required.
+    input_smooth_ = register_parameter(
+        "input_smooth",
+        torch::empty({num_total_experts_, hidden_size_}, fp_option),
+        false);
+    w2_ = register_parameter(
+        "w2",
+        torch::empty(
+            {num_experts_per_rank_, hidden_size_, local_intermediate_size},
+            quant_option),
+        false);
+    w2_scale_ = register_parameter(
+        "w2_scale",
+        torch::empty({num_experts_per_rank_, hidden_size_}, fp_option),
+        false);
+    act_smooth_ = register_parameter(
+        "act_smooth",
+        torch::empty({num_experts_per_rank_, local_intermediate_size},
+                     fp_option),
+        false);
+
+  } else {
+    w13_ = register_parameter(
+        "w13",
+        torch::empty(
+            {num_experts_per_rank_, local_intermediate_size * 2, hidden_size_},
+            options_),
+        false);
+    w2_ = register_parameter(
+        "w2",
+        torch::empty(
+            {num_experts_per_rank_, hidden_size_, local_intermediate_size},
+            options_),
+        false);
+  }
+}
+
+torch::Tensor FusedMoEImpl::create_group_gemm_output(
+    const torch::Tensor& a,
+    const torch::Tensor& b,
+    const torch::Tensor& group_list,
+    torch::ScalarType dtype,
+    torch::Tensor& workspace) {
+  // unify shape logic: define the target shape once.
+  bool is_3d_weight = (b.dim() != 2);
+  int64_t num_tokens = a.size(0);
+  int64_t out_dim = is_3d_weight ? b.size(1) : b.size(0);
+
+  std::vector<int64_t> output_shape;
+  int64_t required_elements = num_tokens * out_dim;
+
+  if (is_3d_weight) {
+    output_shape = {num_tokens, out_dim};
+  } else {
+    output_shape = {group_list.size(0), num_tokens, out_dim};
+    required_elements *= group_list.size(0);
+  }
+
+  auto options = a.options().dtype(dtype);
+
+  // non-smoothquant: direct allocation
+  if (!is_smoothquant_) {
+    return torch::empty(output_shape, options);
+  }
+
+  // smoothquant: managed workspace logic
+  if (!workspace.defined()) {
+    // Lazy initialization: allocate max buffer for the lifecycle
+    // Note: accessing class members w13_ and w2_ directly for context
+    int64_t max_width = std::max(w13_.size(1), w2_.size(1));
+    workspace = torch::empty({num_tokens * max_width}, options);
+  }
+
+  // view construction
+  CHECK(workspace.numel() >= required_elements)
+      << "FusedMoE Workspace too small! Alloc: " << workspace.numel()
+      << ", Req: " << required_elements;
+
+  // utilize the pre-calculated output_shape
+  return workspace.slice(0, 0, required_elements).view(output_shape);
+}
+
+torch::Tensor FusedMoEImpl::select_experts(
+    const torch::Tensor& hidden_states_2d,
+    const torch::Tensor& router_logits_2d,
+    SelectedExpertInfo& selected_expert_info,
+    bool enable_all2all_communication) {
+  // prepare the parameters for select_experts
+  std::optional<torch::Tensor> e_score_correction_bias = std::nullopt;
+  if (e_score_correction_bias_.defined()) {
+    e_score_correction_bias = e_score_correction_bias_;
+  }
+  int64_t expert_size = w13_.size(0);
+
+  // Step 1: apply softmax topk or sigmoid topk / routing logic
+  torch::Tensor reduce_weight;
+  torch::Tensor expert_id;
+  {
+    xllm::kernel::MoeFusedTopkParams moe_active_topk_params;
+    moe_active_topk_params.input = router_logits_2d;
+    moe_active_topk_params.topk = topk_;
+    moe_active_topk_params.num_expert_group = num_expert_group_;
+    moe_active_topk_params.topk_group = topk_group_;
+    moe_active_topk_params.normalize = renormalize_;
+    moe_active_topk_params.normed_by = "topk_logit";
+    moe_active_topk_params.scoring_func = scoring_func_;
+    moe_active_topk_params.route_scale = route_scale_;
+    moe_active_topk_params.e_score_correction_bias = e_score_correction_bias;
+    std::tie(reduce_weight, expert_id) =
+        xllm::kernel::moe_active_topk(moe_active_topk_params);
+  }
+
+  // Step 2: generate expert ids
+  torch::Tensor gather_idx;
+  torch::Tensor combine_idx;
+  torch::Tensor token_count;
+  std::optional<torch::Tensor> cusum_token_count;
+  {
+    xllm::kernel::MoeGenIdxParams moe_gen_idx_params;
+    moe_gen_idx_params.expert_id = expert_id;
+    moe_gen_idx_params.expert_num = num_total_experts_;
+    std::vector<torch::Tensor> output_vec =
+        xllm::kernel::moe_gen_idx(moe_gen_idx_params);
+    gather_idx = output_vec[0];
+    combine_idx = output_vec[1];
+    token_count = output_vec[2];
+    // during all2all communication, we do not need cusum_token_count in the
+    // following computation
+    if (enable_all2all_communication) {
+      cusum_token_count = std::nullopt;
+    } else {
+      cusum_token_count = output_vec[3];
+    }
+  }
+
+  // Step 3: expand and quantize input if needed
+  torch::Tensor expand_hidden_states;
+  torch::Tensor hidden_states_scale;
+  torch::Tensor token_count_slice;
+  // all2all related variables
+  torch::Tensor dispatch_send_token_tensor;
+  // in all2all, the input is scattered, so there is no need to slice the token
+  // count, and we can use the dispatch buffer directly
+  if (enable_all2all_communication) {
+    token_count_slice = token_count;
+    int64_t num_token_expand = hidden_states_2d.size(0) * topk_;
+    int64_t dispatch_bytes =
+        num_token_expand * deep_ep_params_.dispatch_token_size;
+    dispatch_send_token_tensor =
+        deep_ep_buffer_.dispatch_send_token_tensor.slice(0, 0, dispatch_bytes)
+            .view({num_token_expand, deep_ep_params_.dispatch_token_size});
+  } else {
+    token_count_slice =
+        token_count.slice(0, start_expert_id_, start_expert_id_ + expert_size);
+  }
+
+  if (is_smoothquant_) {
+    xllm::kernel::ScaledQuantizeParams scaled_quantize_params;
+    scaled_quantize_params.x = hidden_states_2d;
+    // use dispatch_send_token_tensor buffer for input
+    //  to reduce memory footprint
+    if (enable_all2all_communication) {
+      scaled_quantize_params.smooth = input_smooth_;
+      scaled_quantize_params.output =
+          dispatch_send_token_tensor.slice(1, 0, hidden_size_);
+    } else {
+      scaled_quantize_params.smooth = input_smooth_.slice(
+          0, start_expert_id_, start_expert_id_ + expert_size);
+      scaled_quantize_params.gather_index_start_position =
+          cusum_token_count.value().index({start_expert_id_}).unsqueeze(0);
+    }
+    scaled_quantize_params.token_count = token_count_slice;
+    scaled_quantize_params.gather_index = gather_idx;
+    scaled_quantize_params.act_mode = "none";
+    scaled_quantize_params.active_coef = 1.0;
+    scaled_quantize_params.is_gated = false;
+    scaled_quantize_params.quant_type = torch::kChar;
+    std::tie(expand_hidden_states, hidden_states_scale) =
+        xllm::kernel::scaled_quantize(scaled_quantize_params);
+    if (enable_all2all_communication) {
+      // since view_as_dtype has not supported stride yet,
+      //  we need to copy the scale output to the dispatch buffer
+      torch::Tensor dispatch_scale_slice =
+          dispatch_send_token_tensor.slice(1, hidden_size_);
+      torch::Tensor hidden_states_scale_bytes =
+          view_as_dtype(hidden_states_scale, torch::kInt8)
+              .view_as(dispatch_scale_slice);
+      dispatch_scale_slice.copy_(hidden_states_scale_bytes);
+    }
+  } else {
+    xllm::kernel::MoeExpandInputParams moe_expand_input_params;
+    moe_expand_input_params.input = hidden_states_2d;
+    moe_expand_input_params.gather_index = gather_idx;
+    moe_expand_input_params.combine_idx = combine_idx;
+    moe_expand_input_params.topk = topk_;
+    expand_hidden_states =
+        xllm::kernel::moe_expand_input(moe_expand_input_params);
+    if (enable_all2all_communication) {
+      // use copy to place the output inside the dispatch buffer
+      torch::Tensor dispatch_tensor =
+          view_as_dtype(expand_hidden_states, torch::kChar);
+      dispatch_send_token_tensor.copy_(dispatch_tensor);
+    }
+  }
+
+  // collect the selected tensor
+  selected_expert_info.reduce_weight = reduce_weight;
+  selected_expert_info.combine_idx = combine_idx;
+  selected_expert_info.token_count_slice = token_count_slice;
+  selected_expert_info.cusum_token_count = cusum_token_count;
+  if (is_smoothquant_) {
+    selected_expert_info.input_scale = hidden_states_scale;
+  }
+
+  return expand_hidden_states;
+}
+
+torch::Tensor FusedMoEImpl::forward_experts(const torch::Tensor& hidden_states,
+                                            const torch::Tensor& router_logits,
+                                            bool enable_all2all_communication) {
+  if (!stream_initialized_) {
+    // update device record
+    device_ = xllm::Device(hidden_states.device());
+
+    // acquire streams from the pool again
+    routed_stream_ = device_.get_stream_from_pool();
+    shared_stream_ = device_.get_stream_from_pool();
+    stream_initialized_ = true;
+  }
+
+  std::optional<torch::Tensor> e_score_correction_bias = std::nullopt;
+  if (e_score_correction_bias_.defined()) {
+    e_score_correction_bias = e_score_correction_bias_;
+  }
+
+  // prepare the parameters for MoE computation
+  torch::Tensor shared_expert_output;
+  torch::IntArrayRef hidden_states_shape = hidden_states.sizes();
+  torch::ScalarType hidden_states_dtype = hidden_states.dtype().toScalarType();
+  torch::Tensor hidden_states_2d =
+      hidden_states.reshape({-1, hidden_states.size(-1)});
+  torch::Tensor router_logits_2d =
+      router_logits.reshape({-1, router_logits.size(-1)});
+  int64_t group_gemm_max_dim = enable_all2all_communication
+                                   ? deep_ep_params_.max_num_tokens_recv / topk_
+                                   : hidden_states_2d.size(0);
+  int64_t expert_size = w13_.size(0);
+
+  // Step 1-3: select experts
+  SelectedExpertInfo selected_expert_info;
+  torch::Tensor expand_hidden_states =
+      select_experts(hidden_states_2d,
+                     router_logits_2d,
+                     selected_expert_info,
+                     enable_all2all_communication);
+
+  // Communciation Step 1: Dipatch
+  // intermediate outputs that are used both in dispatch and combine
+  torch::Tensor gather_by_rank_index;
+  torch::Tensor token_sum;
+  if (enable_all2all_communication) {
+    int64_t dispatch_token_num = hidden_states_2d.size(0) * topk_;
+
+    // 1. Dispatch Step: Generate layout and send data
+    deep_ep_->dispatch_step(dispatch_token_num,
+                            selected_expert_info.token_count_slice);
+
+    // 2. Process Result: Generate indices and unpack to computation buffer
+    // use the buffer during initialization for the output
+    expand_hidden_states = dispatch_recv_token_tensor_head_;
+    std::optional<torch::Tensor> output_tail = std::nullopt;
+    if (is_smoothquant_) {
+      output_tail = dispatch_recv_token_tensor_tail_;
+      // update selected_expert_info with the tail (input scale)
+      selected_expert_info.input_scale = output_tail;
+    }
+
+    DeepEPMetaResult deep_ep_meta = deep_ep_->process_dispatch_result(
+        num_experts_per_rank_, expand_hidden_states, output_tail);
+
+    // Extract metadata for subsequent steps
+    gather_by_rank_index = deep_ep_meta.gather_rank_index;
+    selected_expert_info.token_count_slice = deep_ep_meta.token_count_slice;
+    token_sum = deep_ep_meta.token_sum;
+  }
+
+  // common gemm workspace for reduce memory footprint
+  torch::Tensor gemm_workspace;
+
+  // Step 4: group gemm 1
+  torch::Tensor gemm1_out =
+      create_group_gemm_output(expand_hidden_states,
+                               w13_,
+                               selected_expert_info.token_count_slice,
+                               hidden_states_dtype,
+                               gemm_workspace);
+  // ensure the lifespan of these parameters via brace
+  {
+    xllm::kernel::GroupGemmParams group_gemm_params;
+    torch::ScalarType a_dtype =
+        is_smoothquant_ ? torch::kInt8 : hidden_states_dtype;
+    group_gemm_params.a =
+        view_as_dtype(expand_hidden_states, a_dtype).view({-1, hidden_size_});
+    group_gemm_params.b = w13_;
+    group_gemm_params.token_count =
+        selected_expert_info.token_count_slice.to("cpu");
+    if (is_smoothquant_) {
+      torch::Tensor a_scale =
+          selected_expert_info.input_scale.value().flatten();
+      selected_expert_info.input_scale =
+          view_as_dtype(a_scale, torch::kFloat32);
+      group_gemm_params.a_scale = selected_expert_info.input_scale;
+      group_gemm_params.b_scale = w13_scale_;
+    }
+    group_gemm_params.max_dim = group_gemm_max_dim;
+    group_gemm_params.trans_a = false;
+    group_gemm_params.trans_b = true;
+    group_gemm_params.a_quant_bit = is_smoothquant_ ? 8 : -1;
+    group_gemm_params.output = gemm1_out;
+    group_gemm_params.combine_idx = std::nullopt;
+    gemm1_out = xllm::kernel::group_gemm(group_gemm_params);
+  }
+
+  // Step 5: activation or scaled quantization(fused with activation)
+  torch::Tensor act_out;
+  torch::Tensor act_out_scale;
+  if (is_smoothquant_) {
+    int64_t slice_dim = gemm1_out.size(1);
+    if (is_gated_) slice_dim /= 2;
+    // slice operation is a view, does not take up extra memory, but points to
+    // the same memory
+    act_out = expand_hidden_states.slice(1, 0, slice_dim);
+    act_out_scale =
+        selected_expert_info.input_scale.value().slice(0, 0, gemm1_out.size(0));
+    // call scaled quantization kernel (also fused with activation)
+    xllm::kernel::ScaledQuantizeParams scaled_quantize_params;
+    scaled_quantize_params.x = gemm1_out;
+    scaled_quantize_params.smooth = act_smooth_;
+    scaled_quantize_params.token_count = selected_expert_info.token_count_slice;
+    scaled_quantize_params.output = act_out;
+    scaled_quantize_params.output_scale = act_out_scale;
+    scaled_quantize_params.act_mode = hidden_act_;
+    scaled_quantize_params.active_coef = 1.0;
+    scaled_quantize_params.is_gated = is_gated_;
+    scaled_quantize_params.quant_type = torch::kChar;
+    std::tie(act_out, act_out_scale) =
+        xllm::kernel::scaled_quantize(scaled_quantize_params);
+  } else {
+    act_out = is_gated_
+                  ? gemm1_out.slice(1, 0, gemm1_out.size(1) / 2).contiguous()
+                  : gemm1_out;
+    // call activation kernel
+    xllm::kernel::ActivationParams activation_params;
+    activation_params.input = gemm1_out;
+    activation_params.output = act_out;
+    activation_params.cusum_token_count =
+        selected_expert_info.cusum_token_count;
+    activation_params.act_mode = hidden_act_;
+    activation_params.is_gated = is_gated_;
+    activation_params.start_expert_id = start_expert_id_;
+    activation_params.expert_size = expert_size;
+    xllm::kernel::active(activation_params);
+  }
+
+  // Step 6: group gemm 2
+  torch::Tensor gemm2_out =
+      create_group_gemm_output(act_out,
+                               w2_,
+                               selected_expert_info.token_count_slice,
+                               hidden_states_dtype,
+                               gemm_workspace);
+  // ensure the lifespan of these parameters via brace
+  {
+    xllm::kernel::GroupGemmParams group_gemm_params;
+    group_gemm_params.a = act_out;
+    group_gemm_params.b = w2_;
+    group_gemm_params.token_count =
+        selected_expert_info.token_count_slice.to("cpu");
+    if (is_smoothquant_) {
+      group_gemm_params.a_scale = act_out_scale;
+      group_gemm_params.b_scale = w2_scale_;
+    }
+    group_gemm_params.max_dim = group_gemm_max_dim;
+    group_gemm_params.trans_a = false;
+    group_gemm_params.trans_b = true;
+    group_gemm_params.a_quant_bit = is_smoothquant_ ? 8 : -1;
+    group_gemm_params.output = gemm2_out;
+    group_gemm_params.combine_idx = selected_expert_info.combine_idx;
+    gemm2_out = xllm::kernel::group_gemm(group_gemm_params);
+  }
+
+  // Communciation Step 2: Combine
+  if (enable_all2all_communication) {
+    int64_t num_token_expand = hidden_states_2d.size(0) * topk_;
+    // Delegate pack, layout generation and combine to DeepEP
+    torch::Tensor combine_send_layout =
+        deep_ep_->combine_step_pack(gemm2_out,
+                                    gather_by_rank_index,
+                                    token_sum,
+                                    hidden_size_,
+                                    hidden_states_dtype);
+
+    // create a wait event for the current stream to finish computation
+    auto current_stream = device_.current_stream();
+    routed_stream_->wait_stream(*current_stream);
+    // pure communciation kernel: dispatch
+    {
+      torch::StreamGuard stream_guard = routed_stream_->set_stream_guard();
+      gemm2_out = deep_ep_->combine_step_comm(combine_send_layout,
+                                              num_token_expand,
+                                              hidden_size_,
+                                              hidden_states_dtype);
+    }
+
+    // pure computation kernel: shared experts
+    if (n_shared_experts_ > 0) {
+      shared_stream_->wait_stream(*current_stream);
+      torch::StreamGuard stream_guard = shared_stream_->set_stream_guard();
+      shared_expert_output = shared_experts_(hidden_states);
+    }
+
+    // join for parallelization
+    current_stream->wait_stream(*routed_stream_);
+    if (n_shared_experts_ > 0) {
+      current_stream->wait_stream(*shared_stream_);
+    }
+  }
+
+  // After group gemm is finished, some tensors are no
+  // longer needed. We must explicitly release the memory.
+  expand_hidden_states = torch::Tensor();
+  selected_expert_info.input_scale = std::nullopt;
+  act_out = torch::Tensor();
+
+  // Step 7: combine the intermediate results and get the final hidden states
+  torch::Tensor final_hidden_states;
+  // ensure the lifespan of these parameters via brace
+  {
+    xllm::kernel::MoeCombineResultParams moe_combine_result_params;
+    moe_combine_result_params.input = gemm2_out;
+    moe_combine_result_params.reduce_weight =
+        selected_expert_info.reduce_weight;
+    moe_combine_result_params.gather_ids = selected_expert_info.combine_idx;
+    moe_combine_result_params.cusum_token_count =
+        selected_expert_info.cusum_token_count;
+    moe_combine_result_params.start_expert_id = start_expert_id_;
+    moe_combine_result_params.expert_size = expert_size;
+    moe_combine_result_params.bias = std::nullopt;
+    // if all2all communication is enabled and shared output is provided,
+    //  we will fused the add up to combine result
+    if (enable_all2all_communication && n_shared_experts_ > 0) {
+      moe_combine_result_params.residual =
+          shared_expert_output.reshape({-1, shared_expert_output.size(-1)});
+    }
+
+    final_hidden_states =
+        xllm::kernel::moe_combine_result(moe_combine_result_params);
+  }
+
+  // reshape the final hidden states to the original shape
+  final_hidden_states = final_hidden_states.reshape(hidden_states_shape);
+
+  if (enable_all2all_communication) {
+    return final_hidden_states;
+  }
+
+  // Communciation Step 3: AllReduce for non-all2all communication
+  // shared experts can be parallelized with the final communication step
+  // during moe computation.
+  auto current_stream = device_.current_stream();
+  routed_stream_->wait_stream(*current_stream);
+  {
+    torch::StreamGuard stream_guard = routed_stream_->set_stream_guard();
+    if (tp_pg_->world_size() > 1) {
+      final_hidden_states = parallel_state::reduce(final_hidden_states, tp_pg_);
+    }
+    if (parallel_args_.ep_size() > 1) {
+      final_hidden_states = parallel_state::reduce(
+          final_hidden_states, parallel_args_.moe_ep_group_);
+    }
+  }
+
+  if (n_shared_experts_ > 0) {
+    shared_stream_->wait_stream(*current_stream);
+    torch::StreamGuard stream_guard = shared_stream_->set_stream_guard();
+    // for non all2all, we compute the shared experts parallelized with the
+    // final communication step
+    shared_expert_output = shared_experts_(hidden_states);
+    shared_expert_output =
+        shared_expert_output.reshape({-1, shared_expert_output.size(-1)});
+  }
+
+  // join for parallelization
+  current_stream->wait_stream(*routed_stream_);
+  if (n_shared_experts_ > 0) {
+    current_stream->wait_stream(*shared_stream_);
+    final_hidden_states += shared_expert_output;
+  }
+
+  return final_hidden_states;
+}
+
+torch::Tensor FusedMoEImpl::forward(const torch::Tensor& hidden_states,
+                                    const ModelInputParams& input_params) {
+  // we only support all2all communication for decode stage for now
+  bool enable_all2all_communication =
+      enable_deep_ep_ && std::all_of(input_params.dp_is_decode.begin(),
+                                     input_params.dp_is_decode.end(),
+                                     [](int32_t val) { return val == 1; });
+
+  bool is_dp_ep_parallel =
+      parallel_args_.dp_size() > 1 && parallel_args_.ep_size() > 1;
+  // during all2all communication, the output has been
+  //  gathered and sliced by dispatch and combine steps,
+  //  so we do not need to gather input and slice output again
+  bool need_gather_and_slice =
+      is_dp_ep_parallel && !enable_all2all_communication;
+
+  auto input = hidden_states;
+  if (need_gather_and_slice) {
+    input = parallel_state::gather(input,
+                                   parallel_args_.dp_local_process_group_,
+                                   input_params.dp_global_token_nums);
+  }
+  // MoE Gate
+  auto router_logits = gate_(input);
+
+  // MoE Experts
+  auto output =
+      forward_experts(input, router_logits, enable_all2all_communication);
+
+  if (need_gather_and_slice) {
+    output = get_dp_local_slice(output, input_params, parallel_args_);
+  }
+
+  return output;
+}
+
+void FusedMoEImpl::load_e_score_correction_bias(const StateDict& state_dict) {
+  if (e_score_correction_bias_.defined() &&
+      !e_score_correction_bias_is_loaded_) {
+    LOAD_WEIGHT(e_score_correction_bias);
+  }
+}
+
+void FusedMoEImpl::load_experts(const StateDict& state_dict) {
+  const int64_t rank = tp_pg_->rank();
+  const int64_t world_size = tp_pg_->world_size();
+  const int64_t start_expert_id = start_expert_id_;
+  const int64_t num_experts_per_rank = num_experts_per_rank_;
+  const int64_t num_total_experts = num_total_experts_;
+  std::vector<std::string> prefixes = {"gate_proj.", "up_proj."};
+  if (is_smoothquant_) {
+    LOAD_MOE_FUSED_WEIGHT("qweight", w1, w3, w13);
+    LOAD_MOE_FUSED_WEIGHT("per_channel_scale", w1_scale, w3_scale, w13_scale);
+    // When supporting DeepEP All2All mode,
+    // we need to load the complete set of expert weights corresponding to
+    // "up_proj.smooth". Note that even if deep EP mode is not enabled, it
+    // remains possible to retrieve the smooth quantization information for a
+    // subset of experts. Therefore, we intentionally do not check whether
+    // deep_ep_ is enabled in this case.
+    LOAD_MOE_ALL_EXPERT_WEIGHT("up_proj.", "smooth", input_smooth, -1);
+    LOAD_MOE_WEIGHT("down_proj.", "qweight", w2, 1);
+    LOAD_MOE_WEIGHT("down_proj.", "per_channel_scale", w2_scale, -1);
+    LOAD_MOE_WEIGHT("down_proj.", "smooth", act_smooth, 0);
+  } else {
+    LOAD_MOE_FUSED_WEIGHT("weight", w1, w3, w13);
+    LOAD_MOE_WEIGHT("down_proj.", "weight", w2, 1);
+  }
+}
+
+void FusedMoEImpl::load_state_dict(const StateDict& state_dict) {
+  if (state_dict.size() == 0) {
+    return;
+  }
+
+  if (n_shared_experts_ > 0) {
+    shared_experts_->load_state_dict(
+        state_dict.get_dict_with_prefix("shared_experts."));
+  }
+  gate_->load_state_dict(state_dict.get_dict_with_prefix("gate."));
+  load_e_score_correction_bias(state_dict.get_dict_with_prefix("gate."));
+  load_experts(state_dict.get_dict_with_prefix("experts."));
+}
+
+}  // namespace layer
+}  // namespace xllm

ex_engine/xllm_layers/ilu/fused_moe.h

@@ -0,0 +1,131 @@
+/* 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 <torch/torch.h>
+
+#include "framework/model/model_args.h"
+#include "framework/model/model_input_params.h"
+#include "framework/parallel_state/parallel_args.h"
+#include "framework/quant_args.h"
+#include "framework/state_dict/state_dict.h"
+#include "framework/state_dict/utils.h"
+#include "layers/common/deep_ep.h"
+#include "layers/common/dense_mlp.h"
+#include "layers/common/fused_moe_base.h"
+#include "layers/common/linear.h"
+#include "platform/device.h"
+#include "util/tensor_helper.h"
+
+namespace xllm {
+namespace layer {
+
+class FusedMoEImpl : public torch::nn::Module {
+ public:
+  FusedMoEImpl() = default;
+  FusedMoEImpl(const ModelArgs& model_args,
+               const FusedMoEArgs& moe_args,
+               const QuantArgs& quant_args,
+               const ParallelArgs& parallel_args,
+               const torch::TensorOptions& options);
+
+  torch::Tensor forward_experts(const torch::Tensor& hidden_states,
+                                const torch::Tensor& router_logits,
+                                bool enable_all2all_communication);
+  torch::Tensor forward(const torch::Tensor& hidden_states,
+                        const ModelInputParams& input_params);
+  void load_state_dict(const StateDict& state_dict);
+
+ private:
+  // struct to store the selected expert info
+  struct SelectedExpertInfo {
+    torch::Tensor reduce_weight;
+    torch::Tensor combine_idx;
+    torch::Tensor token_count_slice;
+    std::optional<torch::Tensor> cusum_token_count;
+    std::optional<torch::Tensor> input_scale;
+  };
+
+  // initial steps for MoE computation, select the experts for each token
+  torch::Tensor select_experts(const torch::Tensor& hidden_states_2d,
+                               const torch::Tensor& router_logits_2d,
+                               SelectedExpertInfo& selected_expert_info,
+                               bool enable_all2all_communication);
+
+ private:
+  int64_t num_total_experts_;
+  int64_t topk_;
+  int64_t num_expert_group_;
+  int64_t topk_group_;
+  double route_scale_;
+  int64_t hidden_size_;
+  int64_t n_shared_experts_;
+  bool is_gated_;
+  int64_t renormalize_;
+  std::string hidden_act_;
+  std::string scoring_func_;
+  bool is_smoothquant_;
+
+  int64_t num_experts_per_rank_;
+  int64_t start_expert_id_;
+
+  // Deep EP related parameters
+  bool enable_deep_ep_;
+  DeepEPBuffer deep_ep_buffer_;
+  DeepEPParams deep_ep_params_;
+  torch::Tensor dispatch_recv_token_tensor_head_;
+  torch::Tensor dispatch_recv_token_tensor_tail_;
+
+  // steams for parallel shared experts
+  std::unique_ptr<Stream> shared_stream_;
+  std::unique_ptr<Stream> routed_stream_;
+  xllm::Device device_;
+  bool stream_initialized_ = false;
+
+  ReplicatedLinear gate_{nullptr};
+  DenseMLP shared_experts_{nullptr};
+  DeepEP deep_ep_{nullptr};
+
+  QuantArgs quant_args_;
+  ParallelArgs parallel_args_;
+  torch::TensorOptions options_;
+  ProcessGroup* tp_pg_;
+
+  DEFINE_WEIGHT(w13);
+  DEFINE_FUSED_WEIGHT(w1);
+  DEFINE_FUSED_WEIGHT(w3);
+  DEFINE_FUSED_WEIGHT(w2);
+  DEFINE_WEIGHT(e_score_correction_bias);
+  DEFINE_WEIGHT(w13_scale);
+  DEFINE_FUSED_WEIGHT(w1_scale);
+  DEFINE_FUSED_WEIGHT(w3_scale);
+  DEFINE_FUSED_WEIGHT(w2_scale);
+  DEFINE_FUSED_WEIGHT(input_smooth);
+  DEFINE_FUSED_WEIGHT(act_smooth);
+
+  void load_e_score_correction_bias(const StateDict& state_dict);
+  void load_experts(const StateDict& state_dict);
+  // create the group gemm output tensor with the workspace
+  torch::Tensor create_group_gemm_output(const torch::Tensor& a,
+                                         const torch::Tensor& b,
+                                         const torch::Tensor& group_list,
+                                         torch::ScalarType dtype,
+                                         torch::Tensor& workspace);
+};
+TORCH_MODULE(FusedMoE);
+
+}  // namespace layer
+}  // namespace xllm

ex_engine/xllm_layers/mlu/qwen3_5_attention.cpp

@@ -0,0 +1,236 @@
+/* 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 "qwen3_5_attention.h"
+
+#include <glog/logging.h>
+
+#include <tuple>
+
+#include "kernels/ops_api.h"
+namespace xllm {
+namespace layer {
+
+Qwen3_5AttentionImpl::Qwen3_5AttentionImpl(const ModelArgs& args,
+                                           const QuantArgs& quant_args,
+                                           const ParallelArgs& parallel_args,
+                                           const torch::TensorOptions& options,
+                                           int32_t layer_id) {
+  const int64_t tp_size = parallel_args.tp_group_->world_size();
+  const int64_t total_num_heads = args.n_heads();
+  const int64_t total_num_kv_heads = args.n_kv_heads().value_or(args.n_heads());
+  layer_id_ = layer_id;
+  rank_ = parallel_args.tp_group_->rank();
+  CHECK(total_num_heads % tp_size == 0);
+  num_heads_ = total_num_heads / tp_size;
+
+  if (total_num_kv_heads >= tp_size) {
+    CHECK(total_num_kv_heads % tp_size == 0);
+    num_kv_heads_ = total_num_kv_heads / tp_size;
+    num_kv_head_replicas_ = 1;
+  } else {
+    CHECK(tp_size % total_num_kv_heads == 0);
+    num_kv_heads_ = 1;
+    num_kv_head_replicas_ = tp_size / total_num_kv_heads;
+  }
+
+  head_dim_ = args.head_dim();
+  q_size_ = num_heads_ * head_dim_;
+  kv_size_ = num_kv_heads_ * head_dim_;
+  scaling_ = 1.0f / std::sqrt(static_cast<float>(head_dim_));
+  attn_output_gate_ = args.attn_output_gate();
+  mrope_cu_seq_lens_ = torch::zeros(2, torch::kInt32).to(options.device());
+  // 1. QKV linear
+  qkv_proj_ = register_module(
+      "qkv_proj",
+      QKVParallelLinear(args.hidden_size(),
+                        attn_output_gate_ ? num_heads_ * 2 : num_heads_,
+                        num_kv_heads_,
+                        args.head_dim(),
+                        num_kv_head_replicas_,
+                        /*bias=*/args.attention_bias(),
+                        /*gather_output=*/false,
+                        parallel_args,
+                        options));
+
+  // 2. O proj
+  o_proj_ = register_module("o_proj",
+                            RowParallelLinear(total_num_heads * head_dim_,
+                                              args.hidden_size(),
+                                              /*bias=*/false,
+                                              /*input_is_parallelized=*/true,
+                                              /*if_reduce_results=*/true,
+                                              quant_args,
+                                              parallel_args.tp_group_,
+                                              options));
+
+  // 3. Q norm
+  q_norm_ = register_module(
+      "q_norm", Qwen3NextRMSNorm(head_dim_, args.rms_norm_eps(), options));
+
+  // 4. K norm
+  k_norm_ = register_module(
+      "k_norm", Qwen3NextRMSNorm(head_dim_, args.rms_norm_eps(), options));
+
+  // 5. Attention
+  attn_ = register_module("attn",
+                          Attention(num_heads_,
+                                    head_dim_,
+                                    scaling_,
+                                    num_kv_heads_,
+                                    args.sliding_window()));
+
+  // 6. Rotary embedding
+  const int32_t rotary_dim =
+      static_cast<int32_t>(head_dim_ * args.partial_rotary_factor());
+  rotary_emb_ =
+      register_module("rope",
+                      MRotaryEmbedding(rotary_dim,
+                                       args.max_position_embeddings(),
+                                       args.rope_theta(),
+                                       /*interleaved=*/false,
+                                       args.rope_scaling_mrope_section(),
+                                       options));
+}
+
+void Qwen3_5AttentionImpl::rotary_emb_forward(
+    torch::Tensor& q,
+    torch::Tensor& k,
+    const torch::Tensor& positions,
+    const AttentionMetadata& attn_metadata) {
+  auto q_shape = q.sizes();
+  auto k_shape = k.sizes();
+  auto num_tokens = positions.size(-1);
+  mrope_cu_seq_lens_[1] = num_tokens;
+
+  xllm::kernel::RotaryParams rotary_params;
+  bool only_prefill =
+      (attn_metadata.is_prefill || attn_metadata.is_chunked_prefill);
+  if (only_prefill) {
+    rotary_params.sin = attn_metadata.mrope_sin;
+    rotary_params.cos = attn_metadata.mrope_cos;
+    rotary_params.position_ids = std::nullopt;
+    rotary_params.cu_query_lens = mrope_cu_seq_lens_;
+    rotary_params.interleaved = false;
+    rotary_params.discrete = false;
+    rotary_params.max_query_len = num_tokens;
+
+    rotary_params.q = q.view({num_tokens, -1, head_dim_});
+    xllm::kernel::apply_rotary(rotary_params);
+    q = rotary_params.q.reshape(q_shape);
+
+    rotary_params.q = k.view({num_tokens, -1, head_dim_});
+    xllm::kernel::apply_rotary(rotary_params);
+    k = rotary_params.q.reshape(k_shape);
+  } else {
+    if (positions.dim() == 2) {
+      rotary_params.position_ids = positions[0];
+    } else {
+      rotary_params.position_ids = positions;
+    }
+    rotary_params.sin = rotary_emb_->get_sin_cache();
+    rotary_params.cos = rotary_emb_->get_cos_cache();
+
+    rotary_params.interleaved = false;
+    rotary_params.discrete = true;
+    rotary_params.max_query_len = num_tokens;
+    rotary_params.q = q.view({1, num_tokens, -1, head_dim_});
+    xllm::kernel::apply_rotary(rotary_params);
+    q = rotary_params.q.reshape(q_shape);
+
+    rotary_params.q = k.view({1, num_tokens, -1, head_dim_});
+    xllm::kernel::apply_rotary(rotary_params);
+    k = rotary_params.q.reshape(k_shape);
+  }
+}
+
+torch::Tensor Qwen3_5AttentionImpl::forward(
+    const torch::Tensor& positions,
+    const torch::Tensor& hidden_states,
+    const AttentionMetadata& attn_metadata,
+    KVCache& kv_cache) {
+  // 1. qkv projection
+  auto qkv = qkv_proj_->forward(hidden_states);
+  torch::Tensor q, k, v;
+  torch::Tensor gate;
+
+  if (attn_output_gate_) {
+    // Split qkv for attn_output_gate case: [q_size*2, kv_size, kv_size]
+    auto q_gate = qkv.slice(/*dim=*/-1, 0, q_size_ * 2);
+    k = qkv.slice(/*dim=*/-1, q_size_ * 2, q_size_ * 2 + kv_size_);
+    v = qkv.slice(
+        /*dim=*/-1, q_size_ * 2 + kv_size_, q_size_ * 2 + kv_size_ * 2);
+    v = v.contiguous();
+
+    std::vector<int64_t> orig_shape;
+    for (int64_t i = 0; i < q_gate.dim() - 1; i++) {
+      orig_shape.push_back(q_gate.size(i));
+    }
+    std::vector<int64_t> new_shape = orig_shape;
+    new_shape.push_back(num_heads_);
+    new_shape.push_back(-1);
+    torch::Tensor q_gate_reshaped = q_gate.reshape(new_shape);
+    auto chunks = torch::chunk(q_gate_reshaped, 2, /*dim=*/-1);
+    q = chunks[0];
+    gate = chunks[1];
+
+    std::vector<int64_t> q_new_shape = orig_shape;
+    q_new_shape.push_back(-1);
+    q = q.reshape(q_new_shape);
+
+    std::vector<int64_t> gate_new_shape = orig_shape;
+    gate_new_shape.push_back(-1);
+    gate = gate.reshape(gate_new_shape);
+  } else {
+    // Normal case: [q_size, kv_size, kv_size]
+    q = qkv.slice(/*dim=*/-1, 0, q_size_);
+    k = qkv.slice(/*dim=*/-1, q_size_, q_size_ + kv_size_);
+    v = qkv.slice(/*dim=*/-1, q_size_ + kv_size_, q_size_ + 2 * kv_size_);
+  }
+
+  const int64_t T = q.size(0);
+
+  auto q_reshaped = q.reshape({T, num_heads_, head_dim_});
+  auto q_normed = std::get<0>(q_norm_->forward(q_reshaped));
+  auto k_reshaped = k.reshape({T, num_kv_heads_, head_dim_});
+  auto k_normed = std::get<0>(k_norm_->forward(k_reshaped));
+
+  q = q_normed.view({T, q_size_});
+  k = k_normed.view({T, kv_size_});
+  rotary_emb_forward(q, k, positions, attn_metadata);
+  auto out = std::get<0>(attn_->forward(attn_metadata, q, k, v, kv_cache));
+
+  if (attn_output_gate_) {
+    gate = torch::sigmoid(gate);
+    out = out * gate;
+  }
+
+  out = o_proj_->forward(out);
+  return out;
+}
+
+void Qwen3_5AttentionImpl::load_state_dict(const StateDict& state_dict) {
+  qkv_proj_->load_state_dict(state_dict, {"q_proj.", "k_proj.", "v_proj."});
+  o_proj_->load_state_dict(state_dict.get_dict_with_prefix("o_proj."));
+  if (auto w = state_dict.get_tensor("q_norm.weight"); w.defined()) {
+    q_norm_->load_state_dict(StateDict({{"weight", w}}));
+  }
+  if (auto w = state_dict.get_tensor("k_norm.weight"); w.defined()) {
+    k_norm_->load_state_dict(StateDict({{"weight", w}}));
+  }
+}
+
+}  // namespace layer
+}  // namespace xllm

ex_engine/xllm_layers/mlu/qwen3_5_attention.h

@@ -0,0 +1,79 @@
+/* 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 <torch/torch.h>
+
+#include "attention.h"
+#include "framework/kv_cache/kv_cache.h"
+#include "framework/model/model_args.h"
+#include "framework/parallel_state/parallel_args.h"
+#include "framework/quant_args.h"
+#include "framework/state_dict/state_dict.h"
+#include "layers/common/linear.h"
+#include "layers/common/partial_rotary_embedding.h"
+#include "layers/common/qwen3_next_rms_norm.h"
+#include "layers/common/rotary_embedding.h"
+
+namespace xllm {
+namespace layer {
+
+class Qwen3_5AttentionImpl : public torch::nn::Module {
+ public:
+  Qwen3_5AttentionImpl() = default;
+  Qwen3_5AttentionImpl(const ModelArgs& args,
+                       const QuantArgs& quant_args,
+                       const ParallelArgs& parallel_args,
+                       const torch::TensorOptions& options,
+                       int32_t layer_id);
+
+  torch::Tensor forward(const torch::Tensor& positions,
+                        const torch::Tensor& hidden_states,
+                        const AttentionMetadata& attn_metadata,
+                        KVCache& kv_cache);
+
+  void load_state_dict(const StateDict& state_dict);
+  void rotary_emb_forward(torch::Tensor& q,
+                          torch::Tensor& k,
+                          const torch::Tensor& positions,
+                          const AttentionMetadata& attn_metadata);
+
+ private:
+  int64_t num_heads_;
+  int64_t num_kv_heads_;
+  int64_t num_kv_head_replicas_;
+  int64_t head_dim_;
+  int64_t q_size_;
+  int64_t kv_size_;
+  float scaling_;
+  bool attn_output_gate_;
+  int32_t layer_id_;
+  int32_t rank_;
+
+  QKVParallelLinear qkv_proj_{nullptr};
+  RowParallelLinear o_proj_{nullptr};
+
+  Qwen3NextRMSNorm q_norm_{nullptr};
+  Qwen3NextRMSNorm k_norm_{nullptr};
+
+  Attention attn_{nullptr};
+  MRotaryEmbedding rotary_emb_{nullptr};
+  torch::Tensor mrope_cu_seq_lens_;
+};
+TORCH_MODULE(Qwen3_5Attention);
+
+}  // namespace layer
+}  // namespace xllm

ex_engine/xllm_layers/mlu/qwen3_5_decoder_layer.cpp

@@ -0,0 +1,193 @@
+/* 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 "qwen3_5_decoder_layer.h"
+
+#include <glog/logging.h>
+
+#include "common/global_flags.h"
+#include "layers/common/dp_utils.h"
+
+namespace xllm {
+namespace layer {
+namespace {
+bool use_moe_all2all(bool enable_deep_ep,
+                     const ModelInputParams& input_params) {
+  return enable_deep_ep && all_dp_ranks_are_decode(input_params);
+}
+
+bool is_moe_layer(const ModelArgs& model_args, int32_t layer_id) {
+  const auto& mlp_only_layers = model_args.mlp_only_layers();
+  return std::count(mlp_only_layers.begin(), mlp_only_layers.end(), layer_id) ==
+             0 &&
+         model_args.n_routed_experts() > 0 &&
+         (layer_id + 1) % model_args.decoder_sparse_step() == 0;
+}
+}  // namespace
+
+Qwen3_5DecoderLayerImpl::Qwen3_5DecoderLayerImpl(const ModelContext& context,
+                                                 int32_t layer_id)
+    : parallel_args_(context.get_parallel_args()) {
+  const auto& model_args = context.get_model_args();
+  const auto& quant_args = context.get_quant_args();
+  const auto& options = context.get_tensor_options();
+
+  const bool use_moe = is_moe_layer(model_args, layer_id);
+
+  enable_deep_ep_ = use_moe && FLAGS_expert_parallel_degree == 2;
+  if (enable_deep_ep_) {
+    CHECK_EQ(parallel_args_.dp_size(), parallel_args_.world_size())
+        << "Qwen3.5 MoE only support deep ep all2all when dp_size == "
+           "world_size";
+    CHECK_EQ(parallel_args_.dp_size(), parallel_args_.ep_size())
+        << "Qwen3.5 MoE only support deep ep all2all when dp_size == ep_size";
+  }
+
+  auto layer_types = model_args.layer_types();
+  if (layer_types.empty()) {
+    int32_t interval = model_args.full_attention_interval();
+    for (int32_t i = 0; i < model_args.n_layers(); i++) {
+      layer_types.push_back((i + 1) % interval == 0 ? "full_attention"
+                                                    : "linear_attention");
+    }
+  }
+
+  if (layer_id >= 0 && layer_id < static_cast<int32_t>(layer_types.size())) {
+    layer_type_ = layer_types[layer_id];
+  } else {
+    layer_type_ = "full_attention";
+  }
+
+  if (layer_type_ == "linear_attention") {
+    // TODO: support linear attention
+  } else {
+    full_attention_ = register_module(
+        "self_attn",
+        Qwen3_5Attention(
+            model_args, quant_args, parallel_args_, options, layer_id));
+  }
+
+  input_norm_ = register_module(
+      "input_layernorm",
+      Qwen3NextRMSNorm(
+          model_args.hidden_size(), model_args.rms_norm_eps(), options));
+
+  post_norm_ = register_module(
+      "post_attention_layernorm",
+      Qwen3NextRMSNorm(
+          model_args.hidden_size(), model_args.rms_norm_eps(), options));
+
+  if (use_moe) {
+    moe_mlp_ = register_module("mlp",
+                               Qwen3_5FusedMoE(model_args,
+                                               FusedMoEArgs{.is_gated = true},
+                                               quant_args,
+                                               parallel_args_,
+                                               options));
+  } else {
+    mlp_ = register_module("mlp",
+                           DenseMLP(model_args.hidden_size(),
+                                    model_args.intermediate_size(),
+                                    true,
+                                    false,
+                                    model_args.hidden_act(),
+                                    /*enable_result_reduction=*/true,
+                                    quant_args,
+                                    parallel_args_.tp_group_,
+                                    options));
+  }
+}
+
+void Qwen3_5DecoderLayerImpl::load_state_dict(const StateDict& state_dict) {
+  if (layer_type_ == "linear_attention") {
+    // TODO: support linear attention
+  } else {
+    full_attention_->load_state_dict(
+        state_dict.get_dict_with_prefix("self_attn."));
+  }
+  input_norm_->load_state_dict(
+      state_dict.get_dict_with_prefix("input_layernorm."));
+  post_norm_->load_state_dict(
+      state_dict.get_dict_with_prefix("post_attention_layernorm."));
+  if (moe_mlp_) {
+    moe_mlp_->load_state_dict(state_dict.get_dict_with_prefix("mlp."));
+  } else {
+    mlp_->load_state_dict(state_dict.get_dict_with_prefix("mlp."));
+  }
+}
+
+torch::Tensor Qwen3_5DecoderLayerImpl::run_moe(
+    torch::Tensor x,
+    const ModelInputParams& input_params) {
+  const bool enable_moe_all2all =
+      use_moe_all2all(enable_deep_ep_, input_params);
+  if (need_dp_moe_gather(parallel_args_, enable_moe_all2all)) {
+    x = gather_dp_tokens(x, input_params, parallel_args_);
+    x = moe_mlp_->forward_experts(x, enable_moe_all2all);
+    return get_dp_local_slice(x, input_params, parallel_args_);
+  }
+  return moe_mlp_->forward_experts(x, enable_moe_all2all);
+}
+
+std::tuple<torch::Tensor, std::optional<torch::Tensor>>
+Qwen3_5DecoderLayerImpl::apply_norm(Qwen3NextRMSNorm& norm,
+                                    torch::Tensor& input,
+                                    std::optional<torch::Tensor>& residual) {
+  if (!residual.has_value()) {
+    auto new_residual = input;
+    auto output = std::get<0>(norm->forward(input));
+    return {output, new_residual};
+  }
+  auto orig_dtype = input.dtype();
+  input = input + residual.value();
+  auto new_residual = input;
+  input = input.to(orig_dtype);
+  auto output = std::get<0>(norm->forward(input));
+  return {output, new_residual};
+}
+
+torch::Tensor Qwen3_5DecoderLayerImpl::forward(
+    torch::Tensor& x,
+    std::optional<torch::Tensor>& residual,
+    torch::Tensor& positions,
+    const AttentionMetadata& attn_metadata,
+    KVCache& kv_cache,
+    const ModelInputParams& input_params) {
+  // Pre-attention norm
+  std::tie(x, residual) = apply_norm(input_norm_, x, residual);
+
+  // Attention
+  if (full_attention_) {
+    x = full_attention_->forward(positions, x, attn_metadata, kv_cache);
+  } else {
+    // TODO: support linear attention
+  }
+
+  auto orig_dtype = x.dtype();
+  // Post-attention norm
+  std::tie(x, residual) = apply_norm(post_norm_, x, residual);
+
+  // MLP/MoE
+  if (moe_mlp_) {
+    x = run_moe(x, input_params);
+  } else {
+    x = mlp_->forward(x);
+  }
+  x = x.to(orig_dtype);
+  return x;
+}
+
+}  // namespace layer
+}  // namespace xllm

ex_engine/xllm_layers/mlu/qwen3_5_decoder_layer.h

@@ -0,0 +1,73 @@
+/* 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 <torch/torch.h>
+
+#include <optional>
+#include <string>
+
+#include "framework/kv_cache/kv_cache.h"
+#include "framework/model/model_args.h"
+#include "framework/model/model_input_params.h"
+#include "framework/model_context.h"
+#include "framework/parallel_state/parallel_args.h"
+#include "framework/state_dict/state_dict.h"
+#include "layers/common/dense_mlp.h"
+#include "layers/common/qwen3_next_rms_norm.h"
+#include "layers/mlu/qwen3_5_attention.h"
+#include "layers/mlu/qwen3_5_fused_moe.h"
+
+namespace xllm {
+namespace layer {
+
+class Qwen3_5DecoderLayerImpl final : public torch::nn::Module {
+ public:
+  Qwen3_5DecoderLayerImpl(const ModelContext& context, int32_t layer_id);
+
+  void load_state_dict(const StateDict& state_dict);
+
+  torch::Tensor forward(torch::Tensor& x,
+                        std::optional<torch::Tensor>& residual,
+                        torch::Tensor& positions,
+                        const AttentionMetadata& attn_metadata,
+                        KVCache& kv_cache,
+                        const ModelInputParams& input_params);
+
+ private:
+  std::tuple<torch::Tensor, std::optional<torch::Tensor>> apply_norm(
+      Qwen3NextRMSNorm& norm,
+      torch::Tensor& input,
+      std::optional<torch::Tensor>& residual);
+
+  torch::Tensor run_moe(torch::Tensor x, const ModelInputParams& input_params);
+
+  std::string layer_type_;
+  Qwen3_5Attention full_attention_{nullptr};
+  // TODO: support linear attention
+  // Qwen3_5GatedDeltaNet linear_attention_{nullptr};
+  DenseMLP mlp_{nullptr};
+  Qwen3_5FusedMoE moe_mlp_{nullptr};
+  Qwen3NextRMSNorm input_norm_{nullptr};
+  Qwen3NextRMSNorm post_norm_{nullptr};
+  ParallelArgs parallel_args_;
+  bool enable_deep_ep_ = false;
+};
+
+TORCH_MODULE(Qwen3_5DecoderLayer);
+
+}  // namespace layer
+}  // namespace xllm

ex_engine/xllm_layers/mlu/qwen3_5_fused_moe.cpp

@@ -0,0 +1,209 @@
+/* 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 "qwen3_5_fused_moe.h"
+
+#include <glog/logging.h>
+
+#include "framework/parallel_state/parallel_state.h"
+#include "framework/state_dict/utils.h"
+
+namespace xllm {
+namespace layer {
+namespace {
+torch::Tensor get_tensor_with_weight_suffix(const StateDict& state_dict,
+                                            const std::string& tensor_name) {
+  auto tensor = state_dict.get_tensor(tensor_name);
+  if (!tensor.defined()) {
+    tensor = state_dict.get_tensor(tensor_name + ".weight");
+  }
+  return tensor;
+}
+
+torch::Tensor slice_expert_weights(const torch::Tensor& weight,
+                                   int64_t start_expert_id,
+                                   int64_t num_experts_per_rank) {
+  return weight
+      .slice(0, start_expert_id, start_expert_id + num_experts_per_rank)
+      .contiguous();
+}
+
+bool load_fused_gate_up_fallback(const StateDict& state_dict,
+                                 int64_t rank,
+                                 int64_t world_size,
+                                 int64_t start_expert_id,
+                                 int64_t num_experts_per_rank,
+                                 torch::Tensor& w13) {
+  auto fused_gate_up =
+      get_tensor_with_weight_suffix(state_dict, "gate_up_proj");
+  if (!fused_gate_up.defined()) {
+    return false;
+  }
+
+  if (world_size > 1) {
+    CHECK_EQ(fused_gate_up.size(1) % 2, 0)
+        << "gate_up_proj dim1 must be even, got " << fused_gate_up.size(1);
+    const int64_t full_intermediate = fused_gate_up.size(1) / 2;
+    CHECK_EQ(full_intermediate % world_size, 0)
+        << "gate_up_proj intermediate dim is not divisible by world_size";
+    const int64_t inter_shard = full_intermediate / world_size;
+
+    auto gate_full = fused_gate_up.slice(1, 0, full_intermediate);
+    auto up_full =
+        fused_gate_up.slice(1, full_intermediate, full_intermediate * 2);
+    auto gate_shard =
+        gate_full.slice(1, rank * inter_shard, (rank + 1) * inter_shard);
+    auto up_shard =
+        up_full.slice(1, rank * inter_shard, (rank + 1) * inter_shard);
+    fused_gate_up = torch::cat({gate_shard, up_shard}, 1);
+  }
+
+  auto gate_up_slice = slice_expert_weights(
+      fused_gate_up, start_expert_id, num_experts_per_rank);
+  CHECK_EQ(w13.sizes(), gate_up_slice.sizes())
+      << "weight size mismatch for " << state_dict.prefix()
+      << "experts.gate_up_proj";
+  w13.copy_(gate_up_slice);
+  return true;
+}
+
+bool load_fused_down_fallback(const StateDict& state_dict,
+                              int64_t rank,
+                              int64_t world_size,
+                              int64_t start_expert_id,
+                              int64_t num_experts_per_rank,
+                              torch::Tensor& w2) {
+  auto fused_down = get_tensor_with_weight_suffix(state_dict, "down_proj");
+  if (!fused_down.defined()) {
+    return false;
+  }
+
+  if (world_size > 1) {
+    CHECK_EQ(fused_down.size(2) % world_size, 0)
+        << "down_proj dim2 is not divisible by world_size";
+    const int64_t down_shard = fused_down.size(2) / world_size;
+    fused_down =
+        fused_down.slice(2, rank * down_shard, (rank + 1) * down_shard);
+  }
+
+  auto down_slice =
+      slice_expert_weights(fused_down, start_expert_id, num_experts_per_rank);
+  CHECK_EQ(w2.sizes(), down_slice.sizes())
+      << "weight size mismatch for " << state_dict.prefix()
+      << "experts.down_proj";
+  w2.copy_(down_slice);
+  return true;
+}
+}  // namespace
+
+Qwen3_5FusedMoEImpl::Qwen3_5FusedMoEImpl(const ModelArgs& model_args,
+                                         const FusedMoEArgs& moe_args,
+                                         const QuantArgs& quant_args,
+                                         const ParallelArgs& parallel_args,
+                                         const torch::TensorOptions& options)
+    : FusedMoEImpl(model_args, moe_args, quant_args, parallel_args, options) {
+  if (n_shared_experts_ > 0) {
+    shared_expert_gate_ = register_module(
+        "shared_expert_gate",
+        torch::nn::Linear(
+            torch::nn::LinearOptions(hidden_size_, 1).bias(false)));
+    shared_expert_gate_->weight.set_data(
+        shared_expert_gate_->weight.to(options));
+  }
+}
+
+void Qwen3_5FusedMoEImpl::load_experts(const StateDict& state_dict) {
+  FusedMoEImpl::load_experts(state_dict);
+
+  if (!is_smoothquant_) {
+    if (!w13_is_loaded_) {
+      w13_is_loaded_ = load_fused_gate_up_fallback(state_dict,
+                                                   tp_pg_->rank(),
+                                                   tp_pg_->world_size(),
+                                                   start_expert_id_,
+                                                   num_experts_per_rank_,
+                                                   w13_);
+    }
+
+    if (!w2_is_loaded_) {
+      w2_is_loaded_ = load_fused_down_fallback(state_dict,
+                                               tp_pg_->rank(),
+                                               tp_pg_->world_size(),
+                                               start_expert_id_,
+                                               num_experts_per_rank_,
+                                               w2_);
+    }
+  }
+}
+
+void Qwen3_5FusedMoEImpl::load_state_dict(const StateDict& state_dict) {
+  if (state_dict.size() == 0) {
+    return;
+  }
+
+  if (n_shared_experts_ > 0) {
+    shared_experts_->load_state_dict(
+        state_dict.get_dict_with_prefix("shared_expert."));
+    auto weight = state_dict.get_tensor("shared_expert_gate.weight");
+    if (weight.defined()) {
+      weight = weight.reshape({weight.size(0), -1});
+      DCHECK_EQ(shared_expert_gate_->weight.sizes(), weight.sizes())
+          << "proj weight size mismatch for " << name();
+      shared_expert_gate_->weight.data().copy_(weight);
+    }
+  }
+  gate_->load_state_dict(state_dict.get_dict_with_prefix("gate."));
+  load_experts(state_dict.get_dict_with_prefix("experts."));
+}
+
+void Qwen3_5FusedMoEImpl::final_comm_allreduce(
+    torch::Tensor& final_hidden_states,
+    const torch::Tensor& hidden_states,
+    torch::Tensor& shared_expert_output) {
+  auto current_stream = device_.current_stream();
+  routed_stream_->wait_stream(*current_stream);
+  {
+    torch::StreamGuard stream_guard = routed_stream_->set_stream_guard();
+    if (tp_pg_->world_size() > 1) {
+      final_hidden_states = parallel_state::reduce(final_hidden_states, tp_pg_);
+    }
+    if (parallel_args_.ep_size() > 1) {
+      final_hidden_states = parallel_state::reduce(
+          final_hidden_states, parallel_args_.moe_ep_group_);
+    }
+  }
+
+  if (n_shared_experts_ > 0) {
+    shared_stream_->wait_stream(*current_stream);
+    torch::StreamGuard stream_guard = shared_stream_->set_stream_guard();
+    shared_expert_output = shared_experts_(hidden_states);
+    if (shared_expert_gate_) {
+      auto gate = torch::sigmoid(shared_expert_gate_->forward(hidden_states));
+      shared_expert_output = gate * shared_expert_output;
+    }
+    shared_expert_output =
+        shared_expert_output.reshape({-1, shared_expert_output.size(-1)});
+  }
+
+  // join for parallelization
+  current_stream->wait_stream(*routed_stream_);
+  if (n_shared_experts_ > 0) {
+    current_stream->wait_stream(*shared_stream_);
+    final_hidden_states += shared_expert_output;
+  }
+}
+
+}  // namespace layer
+}  // namespace xllm

ex_engine/xllm_layers/mlu/qwen3_5_fused_moe.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 "layers/mlu/fused_moe.h"
+
+namespace xllm {
+namespace layer {
+
+class Qwen3_5FusedMoEImpl final : public FusedMoEImpl {
+ public:
+  Qwen3_5FusedMoEImpl() = default;
+
+  Qwen3_5FusedMoEImpl(const ModelArgs& model_args,
+                      const FusedMoEArgs& moe_args,
+                      const QuantArgs& quant_args,
+                      const ParallelArgs& parallel_args,
+                      const torch::TensorOptions& options);
+
+  void load_state_dict(const StateDict& state_dict) override;
+
+ protected:
+  void final_comm_allreduce(torch::Tensor& final_hidden_states,
+                            const torch::Tensor& hidden_states,
+                            torch::Tensor& shared_expert_output) override;
+
+ private:
+  void load_experts(const StateDict& state_dict);
+  torch::nn::Linear shared_expert_gate_{nullptr};
+};
+
+TORCH_MODULE(Qwen3_5FusedMoE);
+}  // namespace layer
+}  // namespace xllm

ex_engine/xllm_layers/npu_torch/CMakeLists.txt

@@ -0,0 +1,28 @@
+include(cc_library)
+
+cc_library(
+  NAME
+    npu_torch_layers
+  HDRS
+    fused_moe.h
+    attention.h
+    qwen3_gated_delta_net_base.h
+    qwen3_next_attention.h
+    qwen3_next_gated_delta_net.h
+    qwen3_5_gated_delta_net.h
+    qwen3_next_hybrid_decoder_layer_base.h
+    qwen3_next_decoder_layer_impl.h
+    qwen3_5_decoder_layer_impl.h
+  SRCS
+    fused_moe.cpp
+    attention.cpp
+    qwen3_gated_delta_net_base.cpp
+    qwen3_next_attention.cpp
+    qwen3_next_gated_delta_net.cpp
+    qwen3_next_hybrid_decoder_layer_base.cpp
+    qwen3_5_gated_delta_net.cpp
+    qwen3_next_decoder_layer_impl.cpp
+    qwen3_5_decoder_layer_impl.cpp
+  DEPS
+    :common_layers
+)

ex_engine/xllm_layers/npu_torch/attention.cpp

@@ -0,0 +1,152 @@
+/* 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.
+==============================================================================*/
+
+#include "attention.h"
+
+#include "kernels/npu/npu_ops_api.h"
+#include "kernels/ops_api.h"
+
+DECLARE_bool(enable_chunked_prefill);
+namespace xllm {
+namespace layer {
+
+AttentionImpl::AttentionImpl(int64_t num_heads,
+                             int64_t head_size,
+                             float scale,
+                             int64_t num_kv_heads,
+                             int64_t sliding_window)
+    : num_heads_(num_heads),
+      head_size_(head_size),
+      num_kv_heads_(num_kv_heads),
+      sliding_window_(sliding_window),
+      scale_(scale) {
+  if (sliding_window_ > -1) {
+    sliding_window_ = sliding_window_ - 1;
+  }
+}
+
+std::tuple<torch::Tensor, std::optional<torch::Tensor>> AttentionImpl::forward(
+    const AttentionMetadata& attn_metadata,
+    torch::Tensor& query,
+    torch::Tensor& key,
+    torch::Tensor& value,
+    KVCache& kv_cache) {
+  std::optional<torch::Tensor> output_lse = std::nullopt;
+  torch::Tensor output = torch::empty_like(query);
+
+  if (attn_metadata.is_dummy) {
+    return std::make_tuple(output, output_lse);
+  }
+
+  bool only_prefill =
+      attn_metadata.is_prefill || attn_metadata.is_chunked_prefill;
+
+  torch::Tensor k_cache = kv_cache.get_k_cache();
+  torch::Tensor v = value.view({-1, num_kv_heads_, head_size_});
+  std::optional<torch::Tensor> v_cache = kv_cache.get_v_cache();
+
+  // Reshape and cache key/value
+  xllm::kernel::ReshapePagedCacheParams reshape_paged_cache_params;
+  reshape_paged_cache_params.key = key.view({-1, num_kv_heads_, head_size_});
+  reshape_paged_cache_params.value = v;
+  reshape_paged_cache_params.k_cache = k_cache;
+  reshape_paged_cache_params.v_cache = v_cache;
+  reshape_paged_cache_params.slot_mapping = attn_metadata.slot_mapping;
+  xllm::kernel::reshape_paged_cache(reshape_paged_cache_params);
+
+  if (only_prefill) {
+    prefill_forward(query, key, value, output, k_cache, v_cache, attn_metadata);
+  } else {
+    decoder_forward(query, output, k_cache, v_cache, attn_metadata);
+  }
+
+  output = output.view({-1, num_heads_ * head_size_});
+  return {output, output_lse};
+}
+
+void AttentionImpl::prefill_forward(torch::Tensor& query,
+                                    torch::Tensor& key,
+                                    torch::Tensor& value,
+                                    torch::Tensor& output,
+                                    const torch::Tensor& k_cache,
+                                    const std::optional<torch::Tensor>& v_cache,
+                                    const AttentionMetadata& attn_metadata) {
+  query = query.view({-1, num_heads_, head_size_});
+  output = output.view({-1, num_heads_, head_size_});
+
+  if (attn_metadata.is_prefill) {
+    key = key.view({-1, num_kv_heads_, head_size_});
+    value = value.view({-1, num_kv_heads_, head_size_});
+
+    xllm::kernel::npu::batch_prefill(query,
+                                     key,
+                                     value,
+                                     attn_metadata.attn_mask,
+                                     attn_metadata.kv_seq_lens_host,
+                                     scale_,
+                                     output);
+  } else if (attn_metadata.is_chunked_prefill) {
+    xllm::kernel::npu::batch_prefill(query,
+                                     k_cache,
+                                     v_cache.value(),
+                                     attn_metadata.attn_mask,
+                                     attn_metadata.kv_seq_lens_host,
+                                     scale_,
+                                     output);
+  }
+}
+
+void AttentionImpl::decoder_forward(torch::Tensor& query,
+                                    torch::Tensor& output,
+                                    const torch::Tensor& k_cache,
+                                    const std::optional<torch::Tensor>& v_cache,
+                                    const AttentionMetadata& attn_metadata) {
+  query = query.view({-1, 1, num_heads_, head_size_});
+  output = output.view({-1, 1, num_heads_, head_size_});
+
+  torch::Tensor kv_seq_lens;
+  if (attn_metadata.kv_seq_lens_host.defined()) {
+    kv_seq_lens = attn_metadata.kv_seq_lens_host;
+  } else {
+    // Fallback if host tensor isn't prepared.
+    kv_seq_lens = attn_metadata.kv_seq_lens;
+  }
+
+  if (attn_metadata.paged_attention_tiling_data.defined()) {
+    // Use CustomPagedAttention for ACL graph mode to avoid .to(kCPU) operations
+
+    xllm::kernel::npu::batch_decode_acl_graph(
+        query,
+        k_cache,
+        v_cache.value_or(torch::Tensor()),
+        scale_,
+        attn_metadata.block_table,
+        kv_seq_lens,
+        attn_metadata.paged_attention_tiling_data,
+        output);
+  } else {
+    // Standard PagedAttention path
+    xllm::kernel::npu::batch_decode(query,
+                                    k_cache,
+                                    v_cache.value_or(torch::Tensor()),
+                                    scale_,
+                                    attn_metadata.block_table,
+                                    kv_seq_lens,
+                                    output);
+  }
+}
+
+}  // namespace layer
+}  // namespace xllm

ex_engine/xllm_layers/npu_torch/attention.h

@@ -0,0 +1,70 @@
+/* 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 <torch/torch.h>
+
+#include <tuple>
+
+#include "framework/kv_cache/kv_cache.h"
+#include "framework/model/model_input_params.h"
+#include "layers/common/attention_metadata.h"
+
+namespace xllm {
+namespace layer {
+
+class AttentionImpl : public torch::nn::Module {
+ public:
+  AttentionImpl() = default;
+
+  AttentionImpl(int64_t num_heads,
+                int64_t head_size,
+                float scale,
+                int64_t num_kv_heads,
+                int64_t sliding_window);
+
+  std::tuple<torch::Tensor, std::optional<torch::Tensor>> forward(
+      const AttentionMetadata& attn_metadata,
+      torch::Tensor& query,
+      torch::Tensor& key,
+      torch::Tensor& value,
+      KVCache& kv_cache);
+
+  void prefill_forward(torch::Tensor& query,
+                       torch::Tensor& key,
+                       torch::Tensor& value,
+                       torch::Tensor& output,
+                       const torch::Tensor& k_cache,
+                       const std::optional<torch::Tensor>& v_cache,
+                       const AttentionMetadata& attn_metadata);
+
+  void decoder_forward(torch::Tensor& query,
+                       torch::Tensor& output,
+                       const torch::Tensor& k_cache,
+                       const std::optional<torch::Tensor>& v_cache,
+                       const AttentionMetadata& attn_metadata);
+
+ private:
+  int64_t num_heads_;
+  int64_t head_size_;
+  float scale_;
+  int64_t num_kv_heads_;
+  int64_t sliding_window_;
+};
+TORCH_MODULE(Attention);
+
+}  // namespace layer
+}  // namespace xllm

ex_engine/xllm_layers/npu_torch/fused_moe.cpp

@@ -0,0 +1,513 @@
+/* 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.
+==============================================================================*/
+
+#include "fused_moe.h"
+
+#include <glog/logging.h>
+
+#include <numeric>
+#include <vector>
+
+#include "framework/parallel_state/parallel_state.h"
+#include "kernels/ops_api.h"
+
+namespace xllm {
+namespace layer {
+
+namespace {
+// Generic local tensor helpers.
+torch::Tensor create_group_gemm_output(
+    const torch::Tensor& a,
+    const torch::Tensor& b,
+    const torch::Tensor& group_list,
+    torch::ScalarType dtype = torch::ScalarType::BFloat16) {
+  torch::TensorOptions target_options = a.options().dtype(dtype);
+  if (b.dim() != 2) {
+    return torch::empty({a.size(0), b.size(1)}, target_options);
+  }
+  return torch::empty({group_list.size(0), a.size(0), b.size(0)},
+                      target_options);
+}
+
+torch::Tensor get_tensor_with_weight_suffix(const StateDict& state_dict,
+                                            const std::string& tensor_name) {
+  auto tensor = state_dict.get_tensor(tensor_name);
+  if (!tensor.defined()) {
+    tensor = state_dict.get_tensor(tensor_name + ".weight");
+  }
+  return tensor;
+}
+
+torch::Tensor slice_expert_weights(const torch::Tensor& weight,
+                                   int64_t start_expert_id,
+                                   int64_t num_experts_per_rank) {
+  return weight
+      .slice(0, start_expert_id, start_expert_id + num_experts_per_rank)
+      .contiguous();
+}
+
+// Qwen3.5-MoE fused checkpoint fallback helpers.
+bool load_fused_gate_up_fallback(const StateDict& state_dict,
+                                 int64_t rank,
+                                 int64_t world_size,
+                                 int64_t start_expert_id,
+                                 int64_t num_experts_per_rank,
+                                 torch::Tensor& w13) {
+  auto fused_gate_up =
+      get_tensor_with_weight_suffix(state_dict, "gate_up_proj");
+  if (!fused_gate_up.defined()) {
+    return false;
+  }
+
+  if (world_size > 1) {
+    CHECK_EQ(fused_gate_up.size(1) % 2, 0)
+        << "gate_up_proj dim1 must be even, got " << fused_gate_up.size(1);
+    const int64_t full_intermediate = fused_gate_up.size(1) / 2;
+    CHECK_EQ(full_intermediate % world_size, 0)
+        << "gate_up_proj intermediate dim is not divisible by world_size";
+    const int64_t inter_shard = full_intermediate / world_size;
+
+    auto gate_full = fused_gate_up.slice(1, 0, full_intermediate);
+    auto up_full =
+        fused_gate_up.slice(1, full_intermediate, full_intermediate * 2);
+    auto gate_shard =
+        gate_full.slice(1, rank * inter_shard, (rank + 1) * inter_shard);
+    auto up_shard =
+        up_full.slice(1, rank * inter_shard, (rank + 1) * inter_shard);
+    fused_gate_up = torch::cat({gate_shard, up_shard}, 1);
+  }
+
+  auto gate_up_slice = slice_expert_weights(
+      fused_gate_up, start_expert_id, num_experts_per_rank);
+  CHECK_EQ(w13.sizes(), gate_up_slice.sizes())
+      << "weight size mismatch for " << state_dict.prefix()
+      << "experts.gate_up_proj";
+  w13.copy_(gate_up_slice);
+  return true;
+}
+
+bool load_fused_down_fallback(const StateDict& state_dict,
+                              int64_t rank,
+                              int64_t world_size,
+                              int64_t start_expert_id,
+                              int64_t num_experts_per_rank,
+                              torch::Tensor& w2) {
+  auto fused_down = get_tensor_with_weight_suffix(state_dict, "down_proj");
+  if (!fused_down.defined()) {
+    return false;
+  }
+
+  if (world_size > 1) {
+    CHECK_EQ(fused_down.size(2) % world_size, 0)
+        << "down_proj dim2 is not divisible by world_size";
+    const int64_t down_shard = fused_down.size(2) / world_size;
+    fused_down =
+        fused_down.slice(2, rank * down_shard, (rank + 1) * down_shard);
+  }
+
+  auto down_slice =
+      slice_expert_weights(fused_down, start_expert_id, num_experts_per_rank);
+  CHECK_EQ(w2.sizes(), down_slice.sizes())
+      << "weight size mismatch for " << state_dict.prefix()
+      << "experts.down_proj";
+  w2.copy_(down_slice);
+  return true;
+}
+
+}  // namespace
+
+FusedMoEImpl::FusedMoEImpl(const ModelArgs& model_args,
+                           const FusedMoEArgs& moe_args,
+                           const QuantArgs& quant_args,
+                           const ParallelArgs& parallel_args,
+                           const torch::TensorOptions& options)
+    : num_total_experts_(model_args.n_routed_experts()),
+      topk_(model_args.num_experts_per_tok()),
+      hidden_size_(model_args.hidden_size()),
+      n_shared_experts_(model_args.n_shared_experts()),
+      is_gated_(moe_args.is_gated),
+      renormalize_(model_args.norm_topk_prob() ? 1 : 0),
+      hidden_act_(model_args.hidden_act()),
+      is_smoothquant_(false),
+      quant_args_(quant_args),
+      parallel_args_(parallel_args),
+      options_(options),
+      tp_pg_(parallel_args.tp_group_) {
+  const int64_t num_experts = num_total_experts_;
+  const int64_t intermediate_size =
+      static_cast<int64_t>(model_args.moe_intermediate_size());
+  const std::string& topk_method = model_args.topk_method();
+  int64_t ep_size = parallel_args.ep_size();
+  int64_t ep_rank = 0;
+  if (ep_size > 1) {
+    ep_rank = parallel_args.moe_ep_group_->rank();
+    tp_pg_ = parallel_args.moe_tp_group_;
+  }
+
+  // smoothquant check: If quant_method is not empty, only w8a8 smoothquant is
+  // supported
+  if (!quant_args.quant_method().empty()) {
+    if (quant_args.quant_method() != "smoothquant" || quant_args.bits() != 8 ||
+        !quant_args.activation_dynamic()) {
+      LOG(FATAL) << "FusedMoE only supports w8a8 smoothquant quantization when "
+                    "quant_method is set. "
+                 << "Got quant_method=" << quant_args.quant_method()
+                 << ", bits=" << quant_args.bits()
+                 << ", activation_dynamic=" << quant_args.activation_dynamic();
+    }
+    // If confirmed as smoothquant w8a8, set is_smoothquant_ to true
+    is_smoothquant_ = true;
+  } else {
+    is_smoothquant_ = false;
+  }
+
+  // calculate the number of experts per rank
+  num_experts_per_rank_ = num_experts / ep_size;
+  start_expert_id_ = ep_rank * num_experts_per_rank_;
+
+  if (topk_method == "noaux_tc") {
+    e_score_correction_bias_ = register_parameter(
+        "e_score_correction_bias", torch::empty({num_experts}, options), false);
+  }
+
+  gate_ = register_module(
+      "gate_proj",
+      ReplicatedLinear(hidden_size_, num_experts, false, quant_args, options));
+  if (n_shared_experts_ > 0) {
+    /*
+    The shared_experts are usually implemented using the RowParallelLinear
+    layer. Typically, this output serves as the enable_result_reduction results
+    for the module. If only tensor parallelism is applied, immediate
+    reduction of the shared_experts output isn't necessary; instead, we perform
+    the reduction once at the end of the MoE operation.
+    */
+    shared_experts_ =
+        register_module("shared_experts",
+                        DenseMLP(hidden_size_,
+                                 intermediate_size * n_shared_experts_,
+                                 is_gated_,
+                                 false,
+                                 hidden_act_,
+                                 /*enable_result_reduction=*/false,
+                                 quant_args,
+                                 tp_pg_,
+                                 options));
+    shared_expert_gate_ = register_module(
+        "shared_expert_gate",
+        torch::nn::Linear(
+            torch::nn::LinearOptions(hidden_size_, 1).bias(false)));
+    shared_expert_gate_->weight.set_data(
+        shared_expert_gate_->weight.to(options));
+  }
+
+  // create weight buffer
+  const int64_t world_size = tp_pg_->world_size();
+  int64_t local_intermediate_size = intermediate_size / world_size;
+  if (is_smoothquant_) {
+    auto quant_option = options_.dtype(torch::kInt8);
+    auto fp_option = options_.dtype(torch::kFloat32);
+    w13_ = register_parameter(
+        "w13",
+        torch::empty(
+            {num_experts_per_rank_, local_intermediate_size * 2, hidden_size_},
+            quant_option),
+        false);
+    w13_scale_ = register_parameter(
+        "w13_scale",
+        torch::empty({num_experts_per_rank_, local_intermediate_size * 2},
+                     fp_option),
+        false);
+    input_smooth_ = register_parameter(
+        "input_smooth",
+        torch::empty({num_experts_per_rank_, hidden_size_}, fp_option),
+        false);
+    w2_ = register_parameter(
+        "w2",
+        torch::empty(
+            {num_experts_per_rank_, hidden_size_, local_intermediate_size},
+            quant_option),
+        false);
+    w2_scale_ = register_parameter(
+        "w2_scale",
+        torch::empty({num_experts_per_rank_, hidden_size_}, fp_option),
+        false);
+    act_smooth_ = register_parameter(
+        "act_smooth",
+        torch::empty({num_experts_per_rank_, local_intermediate_size},
+                     fp_option),
+        false);
+
+  } else {
+    w13_ = register_parameter(
+        "w13",
+        torch::empty(
+            {num_experts_per_rank_, local_intermediate_size * 2, hidden_size_},
+            options_),
+        false);
+    w2_ = register_parameter(
+        "w2",
+        torch::empty(
+            {num_experts_per_rank_, hidden_size_, local_intermediate_size},
+            options_),
+        false);
+  }
+}
+
+torch::Tensor FusedMoEImpl::select_experts(
+    const torch::Tensor& hidden_states_2d,
+    const torch::Tensor& router_logits_2d,
+    SelectedExpertInfo& selected_expert_info) {
+  // prepare the parameters for select_experts
+  xllm::kernel::MoeFusedTopkParams moe_active_topk_params;
+  moe_active_topk_params.input = router_logits_2d;
+  moe_active_topk_params.finished = torch::Tensor();
+  moe_active_topk_params.topk = topk_;
+  moe_active_topk_params.scoring_func = "softmax";
+  auto [topk_weights, topk_ids] =
+      xllm::kernel::moe_active_topk(moe_active_topk_params);
+  topk_ids = topk_ids.to(torch::kInt32);
+  if (renormalize_) {
+    topk_weights = topk_weights / (topk_weights.sum(-1, true) + 1e-6);
+  }
+
+  xllm::kernel::MoeInitRoutingV2Params moe_init_routing_params;
+  moe_init_routing_params.x = hidden_states_2d;
+  moe_init_routing_params.expert_idx = topk_ids;
+  moe_init_routing_params.scale = std::nullopt;
+  moe_init_routing_params.offset = std::nullopt;
+  moe_init_routing_params.active_num = hidden_states_2d.size(0) * topk_;
+  moe_init_routing_params.expert_capacity = 0;
+  moe_init_routing_params.expert_num = num_experts_per_rank_;
+  moe_init_routing_params.drop_pad_mode = 0;
+  moe_init_routing_params.expert_tokens_num_type = 1;
+  moe_init_routing_params.expert_tokens_num_flag = true;
+  moe_init_routing_params.row_idx_type = 0;
+  std::vector<int64_t> expert_range = {
+      start_expert_id_, start_expert_id_ + num_experts_per_rank_};
+  moe_init_routing_params.active_expert_range = expert_range;
+  moe_init_routing_params.quant_mode = -1;
+  // TODO: NPU moe_init_routing_v2 is equivalent to moe_gen_idx +
+  // moe_expand_input (and the token_count/cusum outputs) on other backends.
+  auto [expand_hidden_states, expand_row_ids, group_list, dynamic_scale] =
+      xllm::kernel::moe_init_routing_v2(moe_init_routing_params);
+  (void)dynamic_scale;
+
+  // collect the selected tensor
+  selected_expert_info.reduce_weight = topk_weights;
+  selected_expert_info.combine_idx = expand_row_ids;
+  selected_expert_info.token_count_slice = group_list;
+  selected_expert_info.cusum_token_count = group_list;
+  return expand_hidden_states;
+}
+
+torch::Tensor FusedMoEImpl::forward_expert(
+    const torch::Tensor& hidden_states,
+    const torch::Tensor& router_logits,
+    const std::optional<torch::Tensor>& shared_output) {
+  // prepare the parameters for MoE computation
+  torch::IntArrayRef hidden_states_shape = hidden_states.sizes();
+  torch::ScalarType hidden_states_dtype = hidden_states.dtype().toScalarType();
+  torch::Tensor hidden_states_2d =
+      hidden_states.reshape({-1, hidden_states.size(-1)});
+  torch::Tensor router_logits_2d =
+      router_logits.reshape({-1, router_logits.size(-1)});
+
+  // Step 1-3: select experts
+  SelectedExpertInfo selected_expert_info;
+  torch::Tensor expand_hidden_states =
+      select_experts(hidden_states_2d, router_logits_2d, selected_expert_info);
+
+  // Step 4: group gemm 1
+  torch::Tensor gemm1_out =
+      create_group_gemm_output(expand_hidden_states,
+                               w13_,
+                               selected_expert_info.token_count_slice,
+                               hidden_states_dtype);
+
+  {
+    xllm::kernel::GroupGemmParams group_gemm_params;
+    group_gemm_params.a = expand_hidden_states;
+    if (w13_.size(1) != expand_hidden_states.size(1)) {
+      w13_ = w13_.transpose(1, 2);
+    }
+    group_gemm_params.b = w13_;
+    group_gemm_params.group_list = selected_expert_info.token_count_slice;
+    group_gemm_params.split_item = 2;
+    group_gemm_params.group_type = 0;
+    group_gemm_params.group_list_type = 1;
+    gemm1_out = xllm::kernel::group_gemm(group_gemm_params);
+  }
+
+  // Step 5: activation
+  torch::Tensor act_out;
+
+  xllm::kernel::ActivationParams activation_params;
+  activation_params.input = gemm1_out;
+  activation_params.output = act_out;
+  activation_params.act_mode = hidden_act_;
+  activation_params.is_gated = is_gated_;
+  xllm::kernel::active(activation_params);
+  act_out = activation_params.output;
+  // Step 6: group gemm 2
+  torch::Tensor gemm2_out =
+      create_group_gemm_output(act_out,
+                               w2_,
+                               selected_expert_info.token_count_slice,
+                               hidden_states_dtype);
+
+  {
+    xllm::kernel::GroupGemmParams group_gemm_params;
+    group_gemm_params.a = act_out;
+    if (w2_.size(1) != act_out.size(1)) {
+      w2_ = w2_.transpose(1, 2);
+    }
+    group_gemm_params.b = w2_;
+    group_gemm_params.group_list = selected_expert_info.token_count_slice;
+    group_gemm_params.split_item = 2;
+    group_gemm_params.group_type = 0;
+    group_gemm_params.group_list_type = 1;
+    gemm2_out = xllm::kernel::group_gemm(group_gemm_params);
+  }
+
+  // Step 7: combine the intermediate results and get the final hidden states
+  torch::Tensor final_hidden_states;
+  xllm::kernel::MoeCombineResultParams moe_combine_params;
+  moe_combine_params.input = gemm2_out;
+  moe_combine_params.reduce_weight = selected_expert_info.reduce_weight;
+  moe_combine_params.gather_ids = selected_expert_info.combine_idx;
+  final_hidden_states = xllm::kernel::moe_combine_result(moe_combine_params);
+  if (shared_output.has_value()) {
+    final_hidden_states = final_hidden_states + shared_output.value();
+  }
+  // reshape the final hidden states to the original shape
+  final_hidden_states = final_hidden_states.reshape(hidden_states_shape);
+
+  if (tp_pg_->world_size() > 1) {
+    final_hidden_states = parallel_state::reduce(final_hidden_states, tp_pg_);
+  }
+  if (parallel_args_.ep_size() > 1) {
+    final_hidden_states = parallel_state::reduce(final_hidden_states,
+                                                 parallel_args_.moe_ep_group_);
+  }
+  return final_hidden_states;
+}
+
+torch::Tensor FusedMoEImpl::forward(const torch::Tensor& hidden_states,
+                                    const ModelInputParams& input_params) {
+  auto input = hidden_states;
+  bool need_slice = false;
+  if (parallel_args_.dp_size() > 1 && parallel_args_.ep_size() > 1) {
+    input = parallel_state::gather(input,
+                                   parallel_args_.dp_local_process_group_,
+                                   input_params.dp_global_token_nums);
+    need_slice = true;
+  }
+
+  std::optional<torch::Tensor> shared_output = std::nullopt;
+  if (n_shared_experts_ > 0) {
+    shared_output = shared_experts_(input);
+    if (shared_expert_gate_) {
+      auto gate = torch::sigmoid(shared_expert_gate_->forward(input));
+      if (shared_output.has_value()) {
+        torch::Tensor res = gate * shared_output.value();
+        shared_output = res;
+      }
+    }
+  }
+  auto router_logits = gate_(input);
+  auto output = forward_expert(input, router_logits, shared_output);
+
+  if (need_slice) {
+    const auto& dp_tokens = input_params.dp_global_token_nums;
+    const int64_t dp_rank = parallel_args_.dp_local_process_group_->rank();
+    auto start =
+        std::accumulate(dp_tokens.begin(), dp_tokens.begin() + dp_rank, 0);
+    auto end = start + dp_tokens[dp_rank];
+    output = output.slice(0, start, end);
+  }
+  return output;
+}
+
+void FusedMoEImpl::load_e_score_correction_bias(const StateDict& state_dict) {
+  if (e_score_correction_bias_.defined() &&
+      !e_score_correction_bias_is_loaded_) {
+    LOAD_WEIGHT(e_score_correction_bias);
+  }
+}
+
+void FusedMoEImpl::load_experts(const StateDict& state_dict) {
+  const int64_t rank = tp_pg_->rank();
+  const int64_t world_size = tp_pg_->world_size();
+  const int64_t start_expert_id = start_expert_id_;
+  const int64_t num_experts_per_rank = num_experts_per_rank_;
+  std::vector<std::string> prefixes = {"gate_proj.", "up_proj."};
+  if (is_smoothquant_) {
+    LOAD_MOE_FUSED_WEIGHT("qweight", w1, w3, w13);
+    LOAD_MOE_FUSED_WEIGHT("per_channel_scale", w1_scale, w3_scale, w13_scale);
+    LOAD_MOE_WEIGHT("up_proj.", "smooth", input_smooth, -1);
+    LOAD_MOE_WEIGHT("down_proj.", "qweight", w2, 1);
+    LOAD_MOE_WEIGHT("down_proj.", "per_channel_scale", w2_scale, -1);
+    LOAD_MOE_WEIGHT("down_proj.", "smooth", act_smooth, 0);
+  } else {
+    LOAD_MOE_FUSED_WEIGHT("weight", w1, w3, w13);
+    LOAD_MOE_WEIGHT("down_proj.", "weight", w2, 1);
+
+    // Some Qwen3.5-MoE checkpoints store expert weights in fused tensors
+    // (gate_up_proj / down_proj). Fall back to this format when split
+    // gate_proj/up_proj tensors are absent.
+    if (!w13_is_loaded_) {
+      w13_is_loaded_ = load_fused_gate_up_fallback(state_dict,
+                                                   rank,
+                                                   world_size,
+                                                   start_expert_id,
+                                                   num_experts_per_rank,
+                                                   w13_);
+    }
+
+    if (!w2_is_loaded_) {
+      w2_is_loaded_ = load_fused_down_fallback(state_dict,
+                                               rank,
+                                               world_size,
+                                               start_expert_id,
+                                               num_experts_per_rank,
+                                               w2_);
+    }
+  }
+}
+
+void FusedMoEImpl::load_state_dict(const StateDict& state_dict) {
+  if (state_dict.size() == 0) {
+    return;
+  }
+
+  if (n_shared_experts_ > 0) {
+    shared_experts_->load_state_dict(
+        state_dict.get_dict_with_prefix("shared_expert."));
+    auto weight = state_dict.get_tensor("shared_expert_gate.weight");
+    if (weight.defined()) {
+      weight = weight.reshape({weight.size(0), -1});
+      DCHECK_EQ(shared_expert_gate_->weight.sizes(), weight.sizes())
+          << "proj weight size mismatch for " << name();
+      shared_expert_gate_->weight.data().copy_(weight);
+    }
+  }
+
+  gate_->load_state_dict(state_dict.get_dict_with_prefix("gate."));
+  load_e_score_correction_bias(state_dict.get_dict_with_prefix("gate."));
+  load_experts(state_dict.get_dict_with_prefix("experts."));
+}
+
+}  // namespace layer
+}  // namespace xllm

ex_engine/xllm_layers/npu_torch/fused_moe.h

@@ -0,0 +1,113 @@
+/* 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 <torch/torch.h>
+
+#include <optional>
+
+#include "framework/model/model_args.h"
+#include "framework/model/model_input_params.h"
+#include "framework/parallel_state/parallel_args.h"
+#include "framework/quant_args.h"
+#include "framework/state_dict/state_dict.h"
+#include "framework/state_dict/utils.h"
+#include "layers/common/dense_mlp.h"
+#include "layers/common/fused_moe_base.h"
+#include "layers/common/linear.h"
+
+namespace xllm {
+namespace layer {
+
+class FusedMoEImpl : public torch::nn::Module {
+ public:
+  FusedMoEImpl() = default;
+  FusedMoEImpl(const ModelArgs& model_args,
+               const FusedMoEArgs& moe_args,
+               const QuantArgs& quant_args,
+               const ParallelArgs& parallel_args,
+               const torch::TensorOptions& options);
+
+  torch::Tensor forward_expert(
+      const torch::Tensor& hidden_states,
+      const torch::Tensor& router_logits,
+      const std::optional<torch::Tensor>& shared_output);
+  torch::Tensor forward(const torch::Tensor& hidden_states,
+                        const ModelInputParams& input_params);
+  void load_state_dict(const StateDict& state_dict);
+
+ private:
+  // struct to store the selected expert info
+  struct SelectedExpertInfo {
+    torch::Tensor reduce_weight;
+    torch::Tensor combine_idx;
+    torch::Tensor token_count_slice;
+    torch::Tensor cusum_token_count;
+    std::optional<torch::Tensor> input_scale;
+  };
+
+  // initial steps for MoE computation, select the experts for each token
+  torch::Tensor select_experts(const torch::Tensor& hidden_states_2d,
+                               const torch::Tensor& router_logits_2d,
+                               SelectedExpertInfo& selected_expert_info);
+
+ private:
+  int64_t num_total_experts_;
+  int64_t topk_;
+  int64_t num_expert_group_;
+  int64_t topk_group_;
+  double route_scale_;
+  int64_t hidden_size_;
+  int64_t n_shared_experts_;
+  bool is_gated_;
+  bool has_score_bias_;
+  bool has_bias_;
+  bool skip_bias_add_;
+  int64_t renormalize_;
+  std::string hidden_act_;
+  std::string scoring_func_;
+  bool is_smoothquant_;
+
+  int64_t num_experts_per_rank_;
+  int64_t start_expert_id_;
+
+  ReplicatedLinear gate_{nullptr};
+  DenseMLP shared_experts_{nullptr};
+  torch::nn::Linear shared_expert_gate_{nullptr};
+  QuantArgs quant_args_;
+  ParallelArgs parallel_args_;
+  torch::TensorOptions options_;
+  ProcessGroup* tp_pg_;
+
+  DEFINE_WEIGHT(w13);
+  DEFINE_FUSED_WEIGHT(w1);
+  DEFINE_FUSED_WEIGHT(w3);
+  DEFINE_FUSED_WEIGHT(w2);
+  DEFINE_WEIGHT(e_score_correction_bias);
+  DEFINE_WEIGHT(w13_scale);
+  DEFINE_FUSED_WEIGHT(w1_scale);
+  DEFINE_FUSED_WEIGHT(w3_scale);
+  DEFINE_FUSED_WEIGHT(w2_scale);
+  DEFINE_FUSED_WEIGHT(input_smooth);
+  DEFINE_FUSED_WEIGHT(act_smooth);
+
+  void load_e_score_correction_bias(const StateDict& state_dict);
+  void load_experts(const StateDict& state_dict);
+};
+TORCH_MODULE(FusedMoE);
+
+}  // namespace layer
+}  // namespace xllm

ex_engine/xllm_layers/npu_torch/qwen3_5_decoder_layer_impl.cpp

@@ -0,0 +1,32 @@
+/* Copyright 2025-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 "qwen3_5_decoder_layer_impl.h"
+
+namespace xllm {
+namespace layer {
+
+Qwen3_5DecoderLayerImpl::Qwen3_5DecoderLayerImpl(const ModelContext& context,
+                                                 int32_t layer_id)
+    : Qwen3NextDecoderLayerImpl(context,
+                                layer_id,
+                                std::make_shared<Qwen3_5GatedDeltaNetImpl>(
+                                    context.get_model_args(),
+                                    context.get_quant_args(),
+                                    context.get_parallel_args(),
+                                    context.get_tensor_options())) {}
+
+}  // namespace layer
+}  // namespace xllm

ex_engine/xllm_layers/npu_torch/qwen3_5_decoder_layer_impl.h

@@ -0,0 +1,32 @@
+/* Copyright 2025-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 "layers/npu_torch/qwen3_5_gated_delta_net.h"
+#include "layers/npu_torch/qwen3_next_decoder_layer_impl.h"
+
+namespace xllm {
+namespace layer {
+
+class Qwen3_5DecoderLayerImpl : public Qwen3NextDecoderLayerImpl {
+ public:
+  explicit Qwen3_5DecoderLayerImpl(const ModelContext& context,
+                                   int32_t layer_id);
+};
+TORCH_MODULE(Qwen3_5DecoderLayer);
+
+}  // namespace layer
+}  // namespace xllm

ex_engine/xllm_layers/npu_torch/qwen3_5_gated_delta_net.cpp

@@ -0,0 +1,185 @@
+/* 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 "qwen3_5_gated_delta_net.h"
+
+#include <glog/logging.h>
+
+namespace xllm {
+namespace layer {
+
+Qwen3_5GatedDeltaNetImpl::Qwen3_5GatedDeltaNetImpl(
+    const ModelArgs& args,
+    const QuantArgs& quant_args,
+    const ParallelArgs& parallel_args,
+    const torch::TensorOptions& options)
+    : Qwen3NextGatedDeltaNetImpl(args,
+                                 quant_args,
+                                 parallel_args,
+                                 options,
+                                 /*init_projections=*/false) {
+  in_proj_qkv_ = register_module("in_proj_qkv",
+                                 ColumnParallelLinear(args.hidden_size(),
+                                                      k_size_ * 2 + v_size_,
+                                                      /*bias=*/false,
+                                                      /*gather_output=*/false,
+                                                      quant_args,
+                                                      parallel_args.tp_group_,
+                                                      options));
+  in_proj_z_ = register_module("in_proj_z",
+                               ColumnParallelLinear(args.hidden_size(),
+                                                    v_size_,
+                                                    /*bias=*/false,
+                                                    /*gather_output=*/false,
+                                                    quant_args,
+                                                    parallel_args.tp_group_,
+                                                    options));
+  in_proj_b_ = register_module("in_proj_b",
+                               ColumnParallelLinear(args.hidden_size(),
+                                                    num_v_heads_,
+                                                    /*bias=*/false,
+                                                    /*gather_output=*/false,
+                                                    quant_args,
+                                                    parallel_args.tp_group_,
+                                                    options));
+  in_proj_a_ = register_module("in_proj_a",
+                               ColumnParallelLinear(args.hidden_size(),
+                                                    num_v_heads_,
+                                                    /*bias=*/false,
+                                                    /*gather_output=*/false,
+                                                    quant_args,
+                                                    parallel_args.tp_group_,
+                                                    options));
+}
+
+torch::Tensor Qwen3_5GatedDeltaNetImpl::merge_qkvz_from_split_activations(
+    const torch::Tensor& qkv,
+    const torch::Tensor& z) const {
+  CHECK_EQ(qkv.dim(), 3) << "Expected qkv activation to be 3D, got "
+                         << qkv.sizes();
+  CHECK_EQ(z.dim(), 3) << "Expected z activation to be 3D, got " << z.sizes();
+  CHECK_EQ(qkv.size(0), z.size(0)) << "qkv/z batch size mismatch.";
+  CHECK_EQ(qkv.size(1), z.size(1)) << "qkv/z sequence size mismatch.";
+  CHECK_EQ(qkv.size(2), (2 * k_size_ + v_size_) / tp_size_)
+      << "Unexpected qkv hidden size for Qwen3.5.";
+  CHECK_EQ(z.size(2), v_size_ / tp_size_)
+      << "Unexpected z hidden size for Qwen3.5.";
+  CHECK_GT(num_k_heads_, 0) << "linear_num_key_heads must be positive.";
+  CHECK_EQ(num_v_heads_ % num_k_heads_, 0)
+      << "linear_num_value_heads must be divisible by linear_num_key_heads.";
+
+  const int64_t bs = qkv.size(0);
+  const int64_t seqlen = qkv.size(1);
+  const int64_t local_k_heads = num_k_heads_ / tp_size_;
+  const int64_t local_v_heads = num_v_heads_ / tp_size_;
+  const int64_t num_v_heads_per_k = num_v_heads_ / num_k_heads_;
+
+  auto qkv_split = torch::split(
+      qkv, {k_size_ / tp_size_, k_size_ / tp_size_, v_size_ / tp_size_}, 2);
+  auto q = qkv_split[0].view({bs, seqlen, local_k_heads, head_k_dim_});
+  auto k = qkv_split[1].view({bs, seqlen, local_k_heads, head_k_dim_});
+  auto v = qkv_split[2].view({bs, seqlen, local_v_heads, head_v_dim_});
+  auto z_view = z.view({bs, seqlen, local_v_heads, head_v_dim_});
+
+  v = v.view({bs, seqlen, local_k_heads, num_v_heads_per_k * head_v_dim_});
+  z_view =
+      z_view.view({bs, seqlen, local_k_heads, num_v_heads_per_k * head_v_dim_});
+
+  return torch::cat({q, k, v, z_view}, -1).view({bs, seqlen, -1}).contiguous();
+}
+
+torch::Tensor Qwen3_5GatedDeltaNetImpl::merge_ba_from_split_activations(
+    const torch::Tensor& b,
+    const torch::Tensor& a) const {
+  CHECK_EQ(b.dim(), 3) << "Expected b activation to be 3D, got " << b.sizes();
+  CHECK_EQ(a.dim(), 3) << "Expected a activation to be 3D, got " << a.sizes();
+  CHECK_EQ(b.size(0), a.size(0)) << "b/a batch size mismatch.";
+  CHECK_EQ(b.size(1), a.size(1)) << "b/a sequence size mismatch.";
+  CHECK_EQ(b.size(2), num_v_heads_ / tp_size_)
+      << "Unexpected b hidden size for Qwen3.5.";
+  CHECK_EQ(a.size(2), num_v_heads_ / tp_size_)
+      << "Unexpected a hidden size for Qwen3.5.";
+  CHECK_GT(num_k_heads_, 0) << "linear_num_key_heads must be positive.";
+  CHECK_EQ(num_v_heads_ % num_k_heads_, 0)
+      << "linear_num_value_heads must be divisible by linear_num_key_heads.";
+
+  const int64_t bs = b.size(0);
+  const int64_t seqlen = b.size(1);
+  const int64_t local_k_heads = num_k_heads_ / tp_size_;
+  const int64_t num_v_heads_per_k = num_v_heads_ / num_k_heads_;
+
+  auto b_view = b.view({bs, seqlen, local_k_heads, num_v_heads_per_k});
+  auto a_view = a.view({bs, seqlen, local_k_heads, num_v_heads_per_k});
+  return torch::cat({b_view, a_view}, -1).view({bs, seqlen, -1}).contiguous();
+}
+
+std::pair<torch::Tensor, torch::Tensor>
+Qwen3_5GatedDeltaNetImpl::project_padded_inputs(
+    const torch::Tensor& hidden_states,
+    const AttentionMetadata& attn_metadata) {
+  auto qkv = reshape_qkvz_with_pad(attn_metadata,
+                                   in_proj_qkv_->forward(hidden_states));
+  auto z_proj =
+      reshape_qkvz_with_pad(attn_metadata, in_proj_z_->forward(hidden_states));
+  auto b_proj =
+      reshape_qkvz_with_pad(attn_metadata, in_proj_b_->forward(hidden_states));
+  auto a_proj =
+      reshape_qkvz_with_pad(attn_metadata, in_proj_a_->forward(hidden_states));
+  return {merge_qkvz_from_split_activations(qkv, z_proj),
+          merge_ba_from_split_activations(b_proj, a_proj)};
+}
+
+void Qwen3_5GatedDeltaNetImpl::load_projection_state_dict(
+    const StateDict& state_dict) {
+  auto in_proj_qkv_state_dict = state_dict.get_dict_with_prefix("in_proj_qkv.");
+  if (in_proj_qkv_state_dict.size() > 0 && !in_proj_qkv_->is_weight_loaded()) {
+    in_proj_qkv_->load_state_dict(
+        in_proj_qkv_state_dict,
+        /*shard_tensor_count=*/3,
+        /*shard_sizes=*/
+        {k_size_ / tp_size_, k_size_ / tp_size_, v_size_ / tp_size_});
+  }
+
+  auto in_proj_z_state_dict = state_dict.get_dict_with_prefix("in_proj_z.");
+  if (in_proj_z_state_dict.size() > 0 && !in_proj_z_->is_weight_loaded()) {
+    in_proj_z_->load_state_dict(in_proj_z_state_dict);
+  }
+
+  auto in_proj_b_state_dict = state_dict.get_dict_with_prefix("in_proj_b.");
+  if (in_proj_b_state_dict.size() > 0 && !in_proj_b_->is_weight_loaded()) {
+    in_proj_b_->load_state_dict(in_proj_b_state_dict);
+  }
+
+  auto in_proj_a_state_dict = state_dict.get_dict_with_prefix("in_proj_a.");
+  if (in_proj_a_state_dict.size() > 0 && !in_proj_a_->is_weight_loaded()) {
+    in_proj_a_->load_state_dict(in_proj_a_state_dict);
+  }
+}
+
+void Qwen3_5GatedDeltaNetImpl::verify_projection_weights(
+    const std::string& prefix) const {
+  CHECK(in_proj_qkv_ && in_proj_qkv_->is_weight_loaded())
+      << "Missing required weight after all shards loaded: " << prefix
+      << "in_proj_qkv.weight";
+  CHECK(in_proj_z_ && in_proj_z_->is_weight_loaded())
+      << "Missing required weight after all shards loaded: " << prefix
+      << "in_proj_z.weight";
+  CHECK(in_proj_b_ && in_proj_b_->is_weight_loaded())
+      << "Missing required weight after all shards loaded: " << prefix
+      << "in_proj_b.weight";
+  CHECK(in_proj_a_ && in_proj_a_->is_weight_loaded())
+      << "Missing required weight after all shards loaded: " << prefix
+      << "in_proj_a.weight";
+}
+
+}  // namespace layer
+}  // namespace xllm

ex_engine/xllm_layers/npu_torch/qwen3_5_gated_delta_net.h

@@ -0,0 +1,58 @@
+/* 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 <torch/torch.h>
+
+#include <string>
+#include <utility>
+
+#include "qwen3_next_gated_delta_net.h"
+
+namespace xllm {
+namespace layer {
+
+class Qwen3_5GatedDeltaNetImpl : public Qwen3NextGatedDeltaNetImpl {
+ public:
+  Qwen3_5GatedDeltaNetImpl() = default;
+  Qwen3_5GatedDeltaNetImpl(const ModelArgs& args,
+                           const QuantArgs& quant_args,
+                           const ParallelArgs& parallel_args,
+                           const torch::TensorOptions& options);
+
+ protected:
+  std::pair<torch::Tensor, torch::Tensor> project_padded_inputs(
+      const torch::Tensor& hidden_states,
+      const AttentionMetadata& attn_metadata) override;
+
+  void load_projection_state_dict(const StateDict& state_dict) override;
+  void verify_projection_weights(const std::string& prefix) const override;
+
+ private:
+  torch::Tensor merge_qkvz_from_split_activations(const torch::Tensor& qkv,
+                                                  const torch::Tensor& z) const;
+  torch::Tensor merge_ba_from_split_activations(const torch::Tensor& b,
+                                                const torch::Tensor& a) const;
+
+  ColumnParallelLinear in_proj_qkv_{nullptr};
+  ColumnParallelLinear in_proj_z_{nullptr};
+  ColumnParallelLinear in_proj_b_{nullptr};
+  ColumnParallelLinear in_proj_a_{nullptr};
+};
+TORCH_MODULE(Qwen3_5GatedDeltaNet);
+
+}  // namespace layer
+}  // namespace xllm

ex_engine/xllm_layers/npu_torch/qwen3_gated_delta_net_base.cpp

@@ -0,0 +1,576 @@
+/* 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 "qwen3_gated_delta_net_base.h"
+
+#include <glog/logging.h>
+#include <torch/torch.h>
+
+#include <tuple>
+
+#include "xllm/core/kernels/ops_api.h"
+
+namespace xllm {
+namespace layer {
+
+namespace {
+torch::Tensor l2norm(const torch::Tensor& x, int64_t dim, double eps = 1e-6) {
+  auto norm = torch::sqrt(torch::sum(torch::square(x), dim, true) + eps);
+  return x / norm;
+}
+
+std::tuple<torch::Tensor, torch::Tensor> torch_recurrent_gated_delta_rule(
+    torch::Tensor query,
+    torch::Tensor key,
+    torch::Tensor value,
+    torch::Tensor g,
+    torch::Tensor beta,
+    std::optional<torch::Tensor> initial_state,
+    bool output_final_state = true,
+    bool use_qk_l2norm_in_kernel = true) {
+  auto initial_dtype = query.dtype();
+
+  if (use_qk_l2norm_in_kernel) {
+    query = l2norm(query, -1, 1e-6);
+    key = l2norm(key, -1, 1e-6);
+  }
+
+  auto to_float32_and_transpose = [](torch::Tensor x) {
+    return x.transpose(1, 2).contiguous().to(torch::kFloat32);
+  };
+  query = to_float32_and_transpose(query);
+  key = to_float32_and_transpose(key);
+  value = to_float32_and_transpose(value);
+  beta = to_float32_and_transpose(beta);
+  g = to_float32_and_transpose(g);
+
+  int64_t batch_size = key.size(0);
+  int64_t num_heads = key.size(1);
+  int64_t sequence_length = key.size(2);
+  int64_t k_head_dim = key.size(3);
+  int64_t v_head_dim = value.size(3);
+
+  float scale_val = 1.0 / std::sqrt(static_cast<float>(query.size(-1)));
+  torch::Tensor scale = torch::tensor(scale_val, query.options());
+  query = query * scale;
+  torch::Tensor core_attn_out = torch::zeros(
+      {batch_size, num_heads, sequence_length, v_head_dim},
+      torch::TensorOptions().dtype(torch::kFloat32).device(value.device()));
+  torch::Tensor last_recurrent_state;
+  if (!initial_state.has_value()) {
+    last_recurrent_state = torch::zeros(
+        {batch_size, num_heads, k_head_dim, v_head_dim},
+        torch::TensorOptions().dtype(torch::kFloat32).device(value.device()));
+  } else {
+    last_recurrent_state =
+        initial_state.value().to(value.device(), torch::kFloat32);
+  }
+
+  for (int64_t i = 0; i < sequence_length; ++i) {
+    torch::Tensor q_t = query.select(2, i);
+    torch::Tensor k_t = key.select(2, i);
+    torch::Tensor v_t = value.select(2, i);
+    torch::Tensor g_t = g.select(2, i).exp().unsqueeze(-1).unsqueeze(-1);
+    torch::Tensor beta_t = beta.select(2, i).unsqueeze(-1);
+    last_recurrent_state = last_recurrent_state * g_t;
+    torch::Tensor kv_mem =
+        torch::sum(last_recurrent_state * k_t.unsqueeze(-1), -2);
+    torch::Tensor delta = (v_t - kv_mem) * beta_t;
+    last_recurrent_state =
+        last_recurrent_state + k_t.unsqueeze(-1) * delta.unsqueeze(-2);
+    core_attn_out.select(2, i) =
+        torch::sum(last_recurrent_state * q_t.unsqueeze(-1), -2);
+  }
+
+  core_attn_out = core_attn_out.transpose(1, 2).contiguous().to(initial_dtype);
+  return std::make_tuple(core_attn_out, last_recurrent_state);
+}
+
+std::tuple<torch::Tensor, torch::Tensor> torch_chunk_gated_delta_rule(
+    torch::Tensor query,
+    torch::Tensor key,
+    torch::Tensor value,
+    torch::Tensor g,
+    torch::Tensor beta,
+    int64_t chunk_size = 64,
+    c10::optional<torch::Tensor> initial_state = c10::nullopt,
+    bool output_final_state = true,
+    bool use_qk_l2norm_in_kernel = true) {
+  auto initial_dtype = query.dtype();
+  if (use_qk_l2norm_in_kernel) {
+    query = l2norm(query, -1, 1e-6);
+    key = l2norm(key, -1, 1e-6);
+  }
+  auto to_float32 = [](torch::Tensor x) {
+    return x.transpose(1, 2).contiguous().to(torch::kFloat32);
+  };
+
+  query = to_float32(query);
+  key = to_float32(key);
+  value = to_float32(value);
+  beta = to_float32(beta);
+  g = to_float32(g);
+
+  auto batch_size = query.size(0);
+  auto num_heads = query.size(1);
+  auto sequence_length = query.size(2);
+  auto k_head_dim = key.size(-1);
+  auto v_head_dim = value.size(-1);
+
+  int64_t pad_size = (chunk_size - sequence_length % chunk_size) % chunk_size;
+  query = torch::nn::functional::pad(
+      query, torch::nn::functional::PadFuncOptions({0, 0, 0, pad_size}));
+  key = torch::nn::functional::pad(
+      key, torch::nn::functional::PadFuncOptions({0, 0, 0, pad_size}));
+  value = torch::nn::functional::pad(
+      value, torch::nn::functional::PadFuncOptions({0, 0, 0, pad_size}));
+  beta = torch::nn::functional::pad(
+      beta, torch::nn::functional::PadFuncOptions({0, pad_size}));
+  g = torch::nn::functional::pad(
+      g, torch::nn::functional::PadFuncOptions({0, pad_size}));
+
+  int64_t total_sequence_length = sequence_length + pad_size;
+  float scale = 1.0 / std::sqrt(static_cast<float>(query.size(-1)));
+  query = query * scale;
+  auto v_beta = value * beta.unsqueeze(-1);
+  auto k_beta = key * beta.unsqueeze(-1);
+  auto reshape_to_chunks = [chunk_size](torch::Tensor x) {
+    auto shape = x.sizes();
+    std::vector<int64_t> new_shape = {
+        shape[0], shape[1], shape[2] / chunk_size, chunk_size, shape[3]};
+    return x.reshape(new_shape);
+  };
+
+  query = reshape_to_chunks(query);
+  key = reshape_to_chunks(key);
+  value = reshape_to_chunks(value);
+  k_beta = reshape_to_chunks(k_beta);
+  v_beta = reshape_to_chunks(v_beta);
+
+  auto g_shape = g.sizes();
+  std::vector<int64_t> g_new_shape = {
+      g_shape[0], g_shape[1], g_shape[2] / chunk_size, chunk_size};
+  g = g.reshape(g_new_shape);
+  auto mask = torch::triu(
+      torch::ones(
+          {chunk_size, chunk_size},
+          torch::TensorOptions().dtype(torch::kBool).device(query.device())),
+      0);
+
+  g = g.cumsum(-1);
+  auto g_diff = g.unsqueeze(-1) - g.unsqueeze(-2);
+  auto decay_mask = g_diff.tril().exp().to(torch::kFloat32);
+  decay_mask = decay_mask.tril();
+  auto attn = -(torch::matmul(k_beta, key.transpose(-1, -2)) * decay_mask)
+                   .masked_fill(mask, 0.0);
+  for (int64_t i = 1; i < chunk_size; ++i) {
+    if (!attn.is_contiguous()) {
+      attn = attn.contiguous();
+    }
+    auto row = attn.slice(-2, i, i + 1)
+                   .slice(-1, 0, i)
+                   .squeeze(-2)
+                   .clone()
+                   .contiguous();
+    auto sub = attn.slice(-2, 0, i).slice(-1, 0, i).clone().contiguous();
+    auto row_unsq = row.unsqueeze(-1).contiguous();
+    auto row_sub_mul = (row_unsq * sub).contiguous();
+    auto row_sub_sum = row_sub_mul.sum(-2).contiguous();
+    auto row_final = (row + row_sub_sum).contiguous();
+    attn.index_put_({torch::indexing::Ellipsis,
+                     torch::indexing::Slice(i, i + 1),
+                     torch::indexing::Slice(0, i)},
+                    row_final.unsqueeze(-2));
+  }
+
+  attn = attn +
+         torch::eye(
+             chunk_size,
+             torch::TensorOptions().dtype(attn.dtype()).device(attn.device()));
+  value = torch::matmul(attn, v_beta);
+  auto k_cumdecay = torch::matmul(attn, (k_beta * g.exp().unsqueeze(-1)));
+  torch::Tensor last_recurrent_state;
+  if (!initial_state.has_value()) {
+    last_recurrent_state = torch::zeros(
+        {batch_size, num_heads, k_head_dim, v_head_dim},
+        torch::TensorOptions().dtype(value.dtype()).device(value.device()));
+  } else {
+    last_recurrent_state = initial_state.value().to(value);
+  }
+  auto core_attn_out = torch::zeros_like(value);
+  mask = torch::triu(
+      torch::ones(
+          {chunk_size, chunk_size},
+          torch::TensorOptions().dtype(torch::kBool).device(query.device())),
+      1);
+  int64_t num_chunks = total_sequence_length / chunk_size;
+  for (int64_t i = 0; i < num_chunks; ++i) {
+    auto q_i = query.select(2, i);
+    auto k_i = key.select(2, i);
+    auto v_i = value.select(2, i);
+    auto attn_i =
+        (torch::matmul(q_i, k_i.transpose(-1, -2)) * decay_mask.select(2, i))
+            .masked_fill_(mask, 0.0);
+    auto v_prime = torch::matmul(k_cumdecay.select(2, i), last_recurrent_state);
+    auto v_new = v_i - v_prime;
+    auto attn_inter = torch::matmul(q_i * g.select(2, i).unsqueeze(-1).exp(),
+                                    last_recurrent_state);
+    core_attn_out.select(2, i) = attn_inter + torch::matmul(attn_i, v_new);
+    auto g_i_last = g.select(2, i).select(-1, -1).unsqueeze(-1);
+    auto g_exp_term = (g_i_last - g.select(2, i)).exp().unsqueeze(-1);
+    auto k_g_exp = (k_i * g_exp_term).transpose(-1, -2).contiguous();
+    last_recurrent_state = last_recurrent_state * g_i_last.unsqueeze(-1).exp() +
+                           torch::matmul(k_g_exp, v_new);
+  }
+  auto core_attn_out_shape = core_attn_out.sizes();
+  std::vector<int64_t> reshape_shape = {
+      core_attn_out_shape[0],
+      core_attn_out_shape[1],
+      core_attn_out_shape[2] * core_attn_out_shape[3],
+      core_attn_out_shape[4]};
+  core_attn_out = core_attn_out.reshape(reshape_shape);
+  core_attn_out = core_attn_out.slice(2, 0, sequence_length);
+  core_attn_out = core_attn_out.transpose(1, 2).contiguous().to(initial_dtype);
+  return std::make_tuple(core_attn_out, last_recurrent_state);
+}
+}  // namespace
+
+Qwen3GatedDeltaNetBaseImpl::Qwen3GatedDeltaNetBaseImpl(
+    const ModelArgs& args,
+    const QuantArgs& quant_args,
+    const ParallelArgs& parallel_args,
+    const torch::TensorOptions& options) {
+  tp_size_ = parallel_args.tp_group_->world_size();
+  rank_ = parallel_args.tp_group_->rank();
+  num_k_heads_ = args.linear_num_key_heads();
+  num_v_heads_ = args.linear_num_value_heads();
+  head_k_dim_ = args.linear_key_head_dim();
+  head_v_dim_ = args.linear_value_head_dim();
+  k_size_ = num_k_heads_ * head_k_dim_;
+  v_size_ = num_v_heads_ * head_v_dim_;
+  conv_kernel_size_ = args.linear_conv_kernel_dim();
+
+  // Shared causal conv projection over mixed QKV states.
+  conv1d_ = register_module("conv1d",
+                            ColumnParallelLinear(args.linear_conv_kernel_dim(),
+                                                 k_size_ * 2 + v_size_,
+                                                 /*bias=*/false,
+                                                 /*gather_output=*/false,
+                                                 quant_args,
+                                                 parallel_args.tp_group_,
+                                                 options));
+
+  auto opts = options.dtype(torch::kFloat32);
+  dt_bias_ = register_parameter("dt_bias",
+                                torch::ones({num_v_heads_ / tp_size_}, opts),
+                                /*requires_grad=*/false);
+
+  A_log_ = register_parameter("A_log",
+                              torch::empty({num_v_heads_ / tp_size_}, opts),
+                              /*requires_grad=*/false);
+
+  // Output projection and gated RMSNorm shared by hybrid variants.
+  o_proj_ = register_module("out_proj",
+                            RowParallelLinear(v_size_,
+                                              args.hidden_size(),
+                                              /*bias=*/false,
+                                              /*input_is_parallelized=*/true,
+                                              /*if_reduce_results=*/true,
+                                              quant_args,
+                                              parallel_args.tp_group_,
+                                              options));
+
+  norm_ = register_module(
+      "norm", RmsNormGated(head_v_dim_, args.rms_norm_eps(), options));
+}
+
+void Qwen3GatedDeltaNetBaseImpl::load_common_state_dict(
+    const StateDict& state_dict) {
+  const int64_t rank = rank_;
+  const int64_t world_size = tp_size_;
+  const int32_t shard_tensor_count = 3;
+  const std::vector<int64_t> shard_sizes = {
+      k_size_ / tp_size_, k_size_ / tp_size_, v_size_ / tp_size_};
+
+  if (auto w = state_dict.get_tensor("conv1d.weight"); w.defined()) {
+    conv1d_->load_state_dict(
+        StateDict({{"weight", w.squeeze(1)}}), shard_tensor_count, shard_sizes);
+  }
+  o_proj_->load_state_dict(state_dict.get_dict_with_prefix("out_proj."));
+  if (auto w = state_dict.get_tensor("norm.weight"); w.defined()) {
+    norm_->load_state_dict(StateDict({{"weight", w}}));
+  }
+  LOAD_SHARDED_WEIGHT(dt_bias, 0);
+  LOAD_SHARDED_WEIGHT(A_log, 0);
+}
+
+void Qwen3GatedDeltaNetBaseImpl::verify_common_loaded_weights(
+    const std::string& prefix) const {
+  CHECK(dt_bias_is_loaded_)
+      << "Missing required weight after all shards loaded: " << prefix
+      << "dt_bias";
+  CHECK(A_log_is_loaded_) << "Missing required weight after all shards loaded: "
+                          << prefix << "A_log";
+}
+
+torch::Tensor Qwen3GatedDeltaNetBaseImpl::forward(
+    const torch::Tensor& hidden_states,
+    const AttentionMetadata& attn_metadata,
+    KVCache& kv_cache,
+    const ModelInputParams& input_params) {
+  auto [qkvz_padded, ba_padded] =
+      project_padded_inputs(hidden_states, attn_metadata);
+  int64_t batch_size = qkvz_padded.size(0);
+  int64_t seq_len = qkvz_padded.size(1);
+
+  torch::Tensor qkvz_flat =
+      qkvz_padded.view({batch_size * seq_len, qkvz_padded.size(-1)});
+  torch::Tensor ba_flat =
+      ba_padded.view({batch_size * seq_len, ba_padded.size(-1)});
+  xllm::kernel::FusedQkvzbaSplitReshapeParams fused_params;
+  fused_params.mixed_qkvz = qkvz_flat;
+  fused_params.mixed_ba = ba_flat;
+  fused_params.num_heads_qk = static_cast<int32_t>(num_k_heads_ / tp_size_);
+  fused_params.num_heads_v = static_cast<int32_t>(num_v_heads_ / tp_size_);
+  fused_params.head_qk = static_cast<int32_t>(head_k_dim_);
+  fused_params.head_v = static_cast<int32_t>(head_v_dim_);
+
+  torch::Tensor mixed_qkv, z, b, a;
+  std::tie(mixed_qkv, z, b, a) =
+      xllm::kernel::fused_qkvzba_split_reshape_cat(fused_params);
+
+  mixed_qkv = mixed_qkv.view({batch_size, seq_len, mixed_qkv.size(-1)});
+  z = z.view({batch_size, seq_len, num_v_heads_ / tp_size_, head_v_dim_});
+  b = b.view({batch_size, seq_len, num_v_heads_ / tp_size_});
+  a = a.view({batch_size, seq_len, num_v_heads_ / tp_size_});
+
+  torch::Tensor conv_cache = kv_cache.get_conv_cache();
+  torch::Tensor ssm_cache = kv_cache.get_ssm_cache();
+  torch::Tensor g, beta, core_attn_out, last_recurrent_state;
+  auto device = mixed_qkv.device();
+  auto conv_weight = conv1d_->weight();
+  auto linear_state_indices = get_linear_state_indices(input_params, device);
+
+  if (attn_metadata.is_prefill) {
+    mixed_qkv = mixed_qkv.transpose(1, 2);
+    torch::Tensor conv_state =
+        (seq_len < conv_kernel_size_ - 1)
+            ? torch::pad(mixed_qkv, {0, conv_kernel_size_ - 1 - seq_len})
+        : (seq_len > conv_kernel_size_ - 1)
+            ? mixed_qkv.narrow(
+                  -1, seq_len - conv_kernel_size_ + 1, conv_kernel_size_ - 1)
+            : mixed_qkv;
+    conv_state = conv_state.transpose(1, 2).contiguous();
+    conv_cache.index_put_({linear_state_indices},
+                          conv_state.to(conv_cache.dtype()));
+    torch::Tensor bias;
+    auto conv_output =
+        torch::conv1d(mixed_qkv,
+                      conv_weight.unsqueeze(1).to(device),
+                      bias,
+                      /*stride=*/std::vector<int64_t>{1},
+                      /*padding=*/std::vector<int64_t>{3},
+                      /*dilation=*/std::vector<int64_t>{1},
+                      /*groups=*/static_cast<int64_t>(mixed_qkv.size(1)));
+    mixed_qkv = torch::silu(conv_output.slice(2, 0, seq_len));
+
+  } else {
+    xllm::kernel::CausalConv1dUpdateParams conv1d_params;
+    conv1d_params.x = mixed_qkv.reshape({-1, mixed_qkv.size(-1)});
+    conv1d_params.conv_state = conv_cache;
+    conv1d_params.weight = conv_weight;
+    conv1d_params.conv_state_indices = linear_state_indices;
+    conv1d_params.block_idx_last_scheduled_token =
+        std::optional<torch::Tensor>();
+    conv1d_params.initial_state_idx = std::optional<torch::Tensor>();
+    conv1d_params.query_start_loc = attn_metadata.q_cu_seq_lens;
+    conv1d_params.max_query_len = attn_metadata.max_query_len;
+    mixed_qkv = xllm::kernel::causal_conv1d_update(conv1d_params);
+    // Reshape back to 3D [batch_size, dim, seq_len]
+    mixed_qkv =
+        mixed_qkv.view({batch_size, -1, mixed_qkv.size(-1)}).contiguous();
+    mixed_qkv = mixed_qkv.transpose(1, 2);
+  }
+
+  // Compute gated delta net decay and beta terms.
+  if (attn_metadata.is_prefill) {
+    xllm::kernel::FusedGdnGatingParams gdn_params;
+    gdn_params.A_log = A_log_;
+    gdn_params.a = a.contiguous().view({-1, a.size(-1)});
+    gdn_params.b = b.contiguous().view({-1, b.size(-1)});
+    gdn_params.dt_bias = dt_bias_;
+    gdn_params.beta = 1.0f;
+    gdn_params.threshold = 20.0f;
+    std::tie(g, beta) = xllm::kernel::fused_gdn_gating(gdn_params);
+    g = g.squeeze(0).contiguous().view({batch_size, seq_len, a.size(-1)});
+    beta = beta.squeeze(0).contiguous().view({batch_size, seq_len, b.size(-1)});
+  } else {
+    xllm::kernel::FusedGdnGatingParams gdn_params;
+    gdn_params.A_log = A_log_;
+    gdn_params.a = a.view({-1, a.size(-1)});
+    gdn_params.b = b.view({-1, b.size(-1)});
+    gdn_params.dt_bias = dt_bias_;
+    gdn_params.beta = 1.0f;
+    gdn_params.threshold = 20.0f;
+    std::tie(g, beta) = xllm::kernel::fused_gdn_gating(gdn_params);
+  }
+  auto [processed_q, processed_k, processed_v] = process_mixed_qkv(mixed_qkv);
+  // Apply chunked or recurrent gated-delta attention and update caches.
+  if (attn_metadata.is_prefill) {
+    xllm::kernel::ChunkGatedDeltaRuleParams chunk_gated_delta_params;
+    chunk_gated_delta_params.q = processed_q;
+    chunk_gated_delta_params.k = processed_k;
+    chunk_gated_delta_params.v = processed_v;
+    chunk_gated_delta_params.g = g;
+    chunk_gated_delta_params.beta = beta;
+    // Get initial state from ssm_cache for sequences with previous state
+    // Shape: [batch_size, num_heads, head_k_dim, head_v_dim]
+    torch::Tensor initial_state_tensor =
+        torch::index_select(ssm_cache, 0, linear_state_indices);
+    // Todo: chunked-prefill/prefix-cache use initial_state
+    initial_state_tensor.fill_(0.0);
+    chunk_gated_delta_params.initial_state = initial_state_tensor;
+    chunk_gated_delta_params.output_final_state = true;
+    chunk_gated_delta_params.cu_seqlens = attn_metadata.q_cu_seq_lens;
+    chunk_gated_delta_params.head_first = false;
+    chunk_gated_delta_params.use_qk_l2norm_in_kernel = true;
+    std::tie(core_attn_out, last_recurrent_state) =
+        xllm::kernel::chunk_gated_delta_rule(chunk_gated_delta_params);
+    ssm_cache.index_put_(
+        {linear_state_indices},
+        last_recurrent_state.transpose(-1, -2).to(ssm_cache.dtype()));
+  } else {
+    processed_q = xllm::kernel::l2_norm(processed_q, 1e-6);
+    processed_k = xllm::kernel::l2_norm(processed_k, 1e-6);
+    auto zero = torch::zeros({1}, attn_metadata.q_seq_lens.options());
+    torch::Tensor actual_seq_lengths =
+        torch::cat({zero, attn_metadata.q_seq_lens}, 0);
+    double scale = 1.0 / std::sqrt(static_cast<float>(processed_q.size(-1)));
+    core_attn_out = xllm::kernel::recurrent_gated_delta_rule(
+                        processed_q.reshape(
+                            {-1, processed_q.size(-2), processed_q.size(-1)}),
+                        processed_k.reshape(
+                            {-1, processed_k.size(-2), processed_k.size(-1)}),
+                        processed_v.reshape(
+                            {-1, processed_v.size(-2), processed_v.size(-1)}),
+                        ssm_cache,
+                        beta.squeeze(0).contiguous(),
+                        scale,
+                        actual_seq_lengths,
+                        linear_state_indices,
+                        c10::nullopt,
+                        g.squeeze(0).contiguous(),
+                        c10::nullopt)
+                        .unsqueeze(0)
+                        .contiguous();
+  }
+
+  auto z_reshaped = z.view({-1, z.size(-1)});
+  auto core_attn_out_reshaped =
+      core_attn_out.view({-1, core_attn_out.size(-1)});
+  auto norm_out = norm_->forward(core_attn_out_reshaped, z_reshaped);
+  auto z_shape_og = z.sizes().vec();
+  norm_out = norm_out.view(z_shape_og);
+  norm_out = norm_out.view({-1, norm_out.size(2), norm_out.size(3)});
+
+  // Project the normalized attention output back to hidden size.
+  auto rearranged_norm =
+      norm_out.reshape({norm_out.size(0), norm_out.size(1) * norm_out.size(2)});
+  rearranged_norm = reshape_qkvz_unpad(attn_metadata, rearranged_norm);
+  auto attn_output = o_proj_->forward(rearranged_norm);
+  return attn_output;
+}
+
+torch::Tensor Qwen3GatedDeltaNetBaseImpl::reshape_qkvz_unpad(
+    const AttentionMetadata& attn_metadata,
+    const torch::Tensor& padded_qkvz) const {
+  if (!attn_metadata.is_prefill) {
+    return padded_qkvz;
+  }
+  std::vector<torch::Tensor> valid_batches;
+  int64_t bs = attn_metadata.q_seq_lens.size(0);
+  int64_t max_len = attn_metadata.max_query_len;
+  const auto& ori_seq_lens = attn_metadata.q_seq_lens;
+  auto reshaped_qkvz = padded_qkvz.view({bs, max_len, -1});
+  for (int64_t b = 0; b < bs; ++b) {
+    int64_t ori_len = ori_seq_lens[b].template item<int64_t>();
+    torch::Tensor valid_batch = reshaped_qkvz[b].slice(0, 0, ori_len);
+    valid_batches.push_back(valid_batch);
+  }
+  return torch::cat(valid_batches, 0).contiguous();
+}
+
+torch::Tensor Qwen3GatedDeltaNetBaseImpl::get_linear_state_indices(
+    const ModelInputParams& input_params,
+    const torch::Device& device) const {
+  CHECK(!input_params.linear_state_ids.empty())
+      << "linear_state_ids must be populated for gated delta net";
+  if (input_params.linear_state_indices.defined()) {
+    return input_params.linear_state_indices;
+  }
+  return torch::tensor(
+      input_params.linear_state_ids,
+      torch::TensorOptions().dtype(torch::kInt).device(device));
+}
+
+torch::Tensor Qwen3GatedDeltaNetBaseImpl::reshape_qkvz_with_pad(
+    const AttentionMetadata& attn_metadata,
+    const torch::Tensor& qkvz) const {
+  int64_t bs = attn_metadata.q_seq_lens.size(0);
+  int64_t max_len = attn_metadata.max_query_len;
+  const auto& start_loc = attn_metadata.q_seq_lens;
+  if (!attn_metadata.is_prefill) {
+    return qkvz.view({qkvz.size(0), -1, qkvz.size(-1)});
+  }
+  std::vector<torch::Tensor> batches;
+  int64_t idx = 0;
+  for (int64_t b = 0; b < bs; ++b) {
+    int64_t cur_len = start_loc[b].template item<int64_t>();
+    torch::Tensor batch = qkvz.slice(0, idx, idx + cur_len).contiguous();
+    idx = idx + cur_len;
+    if (batch.size(0) != max_len) {
+      batch = batch.size(0) > max_len
+                  ? batch.slice(0, 0, max_len).contiguous()
+                  : torch::nn::functional::pad(
+                        batch,
+                        torch::nn::functional::PadFuncOptions(
+                            {0, 0, 0, max_len - batch.size(0)}))
+                        .contiguous();
+    }
+    batches.push_back(batch);
+  }
+  auto ret = torch::stack(batches, 0).contiguous();
+  return ret;
+}
+
+std::tuple<torch::Tensor, torch::Tensor, torch::Tensor>
+Qwen3GatedDeltaNetBaseImpl::process_mixed_qkv(torch::Tensor& mixed_qkv) const {
+  mixed_qkv = mixed_qkv.transpose(1, 2);
+  int64_t batch_size = mixed_qkv.size(0);
+  int64_t seq_len = mixed_qkv.size(1);
+  std::vector<int64_t> split_sizes = {
+      k_size_ / tp_size_, k_size_ / tp_size_, v_size_ / tp_size_};
+  auto processed_qkv = torch::split(mixed_qkv, split_sizes, 2);
+  auto processed_q = processed_qkv[0];
+  auto processed_k = processed_qkv[1];
+  auto processed_v = processed_qkv[2];
+  processed_q = processed_q.view(
+      {batch_size, seq_len, num_k_heads_ / tp_size_, head_k_dim_});
+  processed_k = processed_k.view(
+      {batch_size, seq_len, num_k_heads_ / tp_size_, head_k_dim_});
+  processed_v = processed_v.view(
+      {batch_size, seq_len, num_v_heads_ / tp_size_, head_v_dim_});
+  return std::make_tuple(processed_q, processed_k, processed_v);
+}
+
+}  // namespace layer
+}  // namespace xllm

ex_engine/xllm_layers/npu_torch/qwen3_gated_delta_net_base.h

@@ -0,0 +1,90 @@
+/* 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 <torch/torch.h>
+
+#include <string>
+#include <tuple>
+#include <utility>
+
+#include "attention.h"
+#include "framework/kv_cache/kv_cache.h"
+#include "framework/model/model_args.h"
+#include "framework/parallel_state/parallel_args.h"
+#include "framework/quant_args.h"
+#include "framework/state_dict/state_dict.h"
+#include "framework/state_dict/utils.h"
+#include "layers/common/linear.h"
+#include "layers/common/rms_norm_gated.h"
+
+namespace xllm {
+namespace layer {
+
+class Qwen3GatedDeltaNetBaseImpl : public torch::nn::Module {
+ public:
+  Qwen3GatedDeltaNetBaseImpl() = default;
+  Qwen3GatedDeltaNetBaseImpl(const ModelArgs& args,
+                             const QuantArgs& quant_args,
+                             const ParallelArgs& parallel_args,
+                             const torch::TensorOptions& options);
+
+  virtual void load_state_dict(const StateDict& state_dict) = 0;
+  virtual void verify_loaded_weights(const std::string& prefix) const = 0;
+
+  torch::Tensor forward(const torch::Tensor& hidden_states,
+                        const AttentionMetadata& attn_metadata,
+                        KVCache& kv_cache,
+                        const ModelInputParams& input_params);
+
+ protected:
+  virtual std::pair<torch::Tensor, torch::Tensor> project_padded_inputs(
+      const torch::Tensor& hidden_states,
+      const AttentionMetadata& attn_metadata) = 0;
+
+  void load_common_state_dict(const StateDict& state_dict);
+  void verify_common_loaded_weights(const std::string& prefix) const;
+
+  torch::Tensor reshape_qkvz_with_pad(const AttentionMetadata& attn_metadata,
+                                      const torch::Tensor& qkvz) const;
+  torch::Tensor reshape_qkvz_unpad(const AttentionMetadata& attn_metadata,
+                                   const torch::Tensor& padded_qkvz) const;
+  torch::Tensor get_linear_state_indices(const ModelInputParams& input_params,
+                                         const torch::Device& device) const;
+
+  std::tuple<torch::Tensor, torch::Tensor, torch::Tensor> process_mixed_qkv(
+      torch::Tensor& mixed_qkv) const;
+
+  int64_t num_k_heads_ = 0;
+  int64_t num_v_heads_ = 0;
+  int64_t head_k_dim_ = 0;
+  int64_t head_v_dim_ = 0;
+  int64_t k_size_ = 0;
+  int64_t v_size_ = 0;
+  int64_t tp_size_ = 1;
+  int64_t rank_ = 0;
+  int32_t conv_kernel_size_ = 0;
+
+  ColumnParallelLinear conv1d_{nullptr};
+  RowParallelLinear o_proj_{nullptr};
+  RmsNormGated norm_{nullptr};
+
+  DEFINE_WEIGHT(dt_bias);
+  DEFINE_WEIGHT(A_log);
+};
+
+}  // namespace layer
+}  // namespace xllm

ex_engine/xllm_layers/npu_torch/qwen3_next_attention.cpp

@@ -0,0 +1,291 @@
+/* Copyright 2025-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 "qwen3_next_attention.h"
+
+#include <glog/logging.h>
+
+#include <tuple>
+#include <vector>
+
+#include "common/flash_comm1_context.h"
+
+namespace xllm {
+namespace layer {
+
+Qwen3NextAttentionImpl::Qwen3NextAttentionImpl(
+    const ModelArgs& args,
+    const QuantArgs& quant_args,
+    const ParallelArgs& parallel_args,
+    const torch::TensorOptions& options,
+    int32_t layer_id) {
+  const int64_t tp_size = parallel_args.tp_group_->world_size();
+  const int64_t total_num_heads = args.n_heads();
+  const int64_t total_num_kv_heads = args.n_kv_heads().value_or(args.n_heads());
+  layer_id_ = layer_id;
+  rank_ = parallel_args.tp_group_->rank();
+  CHECK(total_num_heads % tp_size == 0);
+  num_heads_ = total_num_heads / tp_size;
+
+  if (total_num_kv_heads >= tp_size) {
+    CHECK(total_num_kv_heads % tp_size == 0);
+    num_kv_heads_ = total_num_kv_heads / tp_size;
+    num_kv_head_replicas_ = 1;
+  } else {
+    CHECK(tp_size % total_num_kv_heads == 0);
+    num_kv_heads_ = 1;
+    num_kv_head_replicas_ = tp_size / total_num_kv_heads;
+  }
+
+  head_dim_ = args.head_dim();
+  q_size_ = num_heads_ * head_dim_;
+  kv_size_ = num_kv_heads_ * head_dim_;
+  scaling_ = 1.0f / std::sqrt(static_cast<float>(head_dim_));
+  attn_output_gate_ = args.attn_output_gate();
+  // 1. QKV linear
+  qkv_proj_ = register_module(
+      "qkv_proj",
+      QKVParallelLinear(args.hidden_size(),
+                        attn_output_gate_ ? num_heads_ * 2 : num_heads_,
+                        num_kv_heads_,
+                        args.head_dim(),
+                        num_kv_head_replicas_,
+                        /*bias=*/args.attention_bias(),
+                        /*gather_output=*/false,
+                        parallel_args,
+                        options,
+                        quant_args));
+
+  // 2. O proj
+  o_proj_ = register_module("o_proj",
+                            RowParallelLinear(total_num_heads * head_dim_,
+                                              args.hidden_size(),
+                                              /*bias=*/false,
+                                              /*input_is_parallelized=*/true,
+                                              /*if_reduce_results=*/true,
+                                              quant_args,
+                                              parallel_args.tp_group_,
+                                              options));
+
+  // 3. Q norm
+  q_norm_ = register_module(
+      "q_norm", Qwen3NextRMSNorm(head_dim_, args.rms_norm_eps(), options));
+
+  // 4. K norm
+  k_norm_ = register_module(
+      "k_norm", Qwen3NextRMSNorm(head_dim_, args.rms_norm_eps(), options));
+
+  // 5. Rotary embedding
+  const int rotary_dim =
+      static_cast<int>(head_dim_ * args.partial_rotary_factor());
+  rotary_emb_ =
+      register_module("rotary_emb",
+                      PartialRotaryEmbedding(rotary_dim,
+                                             args.max_position_embeddings(),
+                                             args.rope_theta(),
+                                             head_dim_,
+                                             true,
+                                             false,
+                                             options));
+
+  // 6. Attention
+  attn_ = register_module("attn",
+                          Attention(num_heads_,
+                                    head_dim_,
+                                    scaling_,
+                                    num_kv_heads_,
+                                    args.sliding_window()));
+
+  // 7. Fused split_qkv_rmsnorm_mrope kernel setup
+  rotary_dim_ = static_cast<int64_t>(head_dim_ * args.partial_rotary_factor());
+  rms_norm_eps_ = args.rms_norm_eps();
+  mrope_section_ = args.rope_scaling_mrope_section();
+  is_interleaved_ = args.rope_scaling_mrope_interleaved();
+  use_fused_qkv_ = false;
+  if (attn_output_gate_ && !mrope_section_.empty() &&
+      mrope_section_.size() == 3 && rotary_dim_ > 0 &&
+      xllm::kernel::has_split_qkv_rmsnorm_mrope_specialization(
+          num_heads_, num_kv_heads_, head_dim_)) {
+    mrope_gather_pattern_ =
+        xllm::kernel::build_split_qkv_rmsnorm_mrope_gather_pattern(
+            rotary_dim_, mrope_section_, is_interleaved_, options.device());
+    use_fused_qkv_ = true;
+    LOG(INFO) << "Qwen3NextAttention layer " << layer_id_
+              << ": using fused split_qkv_rmsnorm_mrope kernel";
+  }
+}
+
+torch::Tensor Qwen3NextAttentionImpl::build_mrope_cos_sin(
+    const torch::Tensor& positions) const {
+  auto cos_sin_cache = rotary_emb_->get_cos_sin_cache();
+  if (positions.dim() == 1) {
+    return cos_sin_cache.index_select(0, positions).repeat({1, 3});
+  }
+  // positions is [3, T] for mRoPE (graph mode or VL)
+  // transpose from [3, T] to [T, 3]
+  auto positions_t = positions.permute({1, 0}).contiguous();
+  auto gathered = cos_sin_cache.index_select(0, positions_t.view({-1}));
+  // [T, 3, rope_dim]
+  return gathered.view({positions.size(1), -1});
+}
+
+torch::Tensor Qwen3NextAttentionImpl::forward(
+    const torch::Tensor& positions,
+    const torch::Tensor& hidden_states,
+    const AttentionMetadata& attn_metadata,
+    KVCache& kv_cache,
+    const torch::Tensor& mrope_cos_sin) {
+  const FlashComm1Context* fc1_ctx = get_current_flash_comm1_context();
+  torch::Tensor h = hidden_states;
+
+  if (fc1_ctx && is_sequence_sharded(*fc1_ctx)) {
+    h = gather_sequence(hidden_states, *fc1_ctx);
+  }
+
+  auto qkv = qkv_proj_->forward(h);
+
+  if (use_fused_qkv_) {
+    const int64_t T = qkv.size(0);
+    xllm::kernel::SplitQkvRmsnormMropeParams params;
+    params.qkvg = qkv;
+    params.q_weight = q_norm_->weight();
+    params.k_weight = k_norm_->weight();
+    params.cos_sin = mrope_cos_sin;
+    params.gather_pattern = mrope_gather_pattern_;
+    params.eps = rms_norm_eps_;
+    params.num_q_heads = num_heads_;
+    params.num_kv_heads = num_kv_heads_;
+    params.head_size = head_dim_;
+
+    auto [q, k, v, gate] = xllm::kernel::split_qkv_rmsnorm_mrope(params);
+
+    auto q_flat = q.view({T, q_size_});
+    auto k_flat = k.view({T, kv_size_});
+    auto v_flat = v.view({T, kv_size_});
+
+    auto out = std::get<0>(
+        attn_->forward(attn_metadata, q_flat, k_flat, v_flat, kv_cache));
+    out = out * torch::sigmoid(gate.view({T, q_size_}));
+
+    if (fc1_ctx && is_sequence_sharded(*fc1_ctx)) {
+      return o_proj_->forward(out, row_parallel_reduce_mode_for_fc1(*fc1_ctx));
+    }
+    return o_proj_->forward(out);
+  }
+
+  // Fallback path: weight-reordered layout [Q | G | K | V]
+  torch::Tensor q, k, v;
+  torch::Tensor gate;
+
+  if (attn_output_gate_) {
+    q = qkv.slice(-1, 0, q_size_);
+    gate = qkv.slice(-1, q_size_, q_size_ * 2);
+    k = qkv.slice(-1, q_size_ * 2, q_size_ * 2 + kv_size_);
+    v = qkv.slice(-1, q_size_ * 2 + kv_size_, q_size_ * 2 + kv_size_ * 2);
+  } else {
+    q = qkv.slice(-1, 0, q_size_);
+    k = qkv.slice(-1, q_size_, q_size_ + kv_size_);
+    v = qkv.slice(-1, q_size_ + kv_size_, q_size_ + 2 * kv_size_);
+  }
+
+  const int64_t T = q.size(0);
+  auto q_3d = q.view({T, num_heads_, head_dim_});
+  q = std::get<0>(q_norm_->forward(q_3d)).view({T, q_size_});
+  auto k_3d = k.view({T, num_kv_heads_, head_dim_});
+  k = std::get<0>(k_norm_->forward(k_3d)).view({T, kv_size_});
+
+  rotary_emb_->forward(positions, q, k);
+  auto out = std::get<0>(attn_->forward(attn_metadata, q, k, v, kv_cache));
+
+  if (attn_output_gate_) {
+    out = out * torch::sigmoid(gate);
+  }
+
+  if (fc1_ctx && is_sequence_sharded(*fc1_ctx)) {
+    return o_proj_->forward(out, row_parallel_reduce_mode_for_fc1(*fc1_ctx));
+  }
+  return o_proj_->forward(out);
+}
+
+void Qwen3NextAttentionImpl::load_state_dict(const StateDict& state_dict) {
+  qkv_proj_->load_state_dict(state_dict, {"q_proj.", "k_proj.", "v_proj."});
+
+  if (attn_output_gate_ && qkv_proj_->is_weight_loaded() &&
+      !qkv_weight_reordered_) {
+    // Rearrange q_proj rows from per-head interleaved [q0,g0,q1,g1,...]
+    // to grouped [q0,q1,...,g0,g1,...] so forward output is [Q|G|K|V].
+    auto w = qkv_proj_->weight();
+    auto qg_rows = w.slice(0, 0, q_size_ * 2);
+    const int64_t hidden = w.size(1);
+    auto qg_3d = qg_rows.view({num_heads_, 2 * head_dim_, hidden});
+    auto q_part = qg_3d.slice(1, 0, head_dim_);
+    auto g_part = qg_3d.slice(1, head_dim_, 2 * head_dim_);
+    auto reordered = torch::cat(
+        {q_part.reshape({q_size_, hidden}), g_part.reshape({q_size_, hidden})},
+        0);
+    qg_rows.copy_(reordered);
+
+    // Reorder weight_scale and weight_offset for W8A8 dynamic quantization.
+    // These are per-channel (per output row) tensors that must match the
+    // reordered weight layout for correct dequantization.
+    const int64_t qg_size = q_size_ * 2;
+    auto reorder_per_channel = [this, qg_size](torch::Tensor tensor) {
+      if (!tensor.defined() || tensor.numel() == 0) {
+        return;
+      }
+      auto qg_part = tensor.slice(0, 0, qg_size);
+      auto qg_2d = qg_part.view({num_heads_, 2 * head_dim_});
+      auto q_scale = qg_2d.slice(1, 0, head_dim_);
+      auto g_scale = qg_2d.slice(1, head_dim_, 2 * head_dim_);
+      auto reordered_scale = torch::cat(
+          {q_scale.reshape({q_size_}), g_scale.reshape({q_size_})}, 0);
+      qg_part.copy_(reordered_scale);
+    };
+
+    if (qkv_proj_->is_weight_scale_loaded()) {
+      reorder_per_channel(qkv_proj_->weight_scale());
+    }
+    if (qkv_proj_->is_weight_offset_loaded()) {
+      reorder_per_channel(qkv_proj_->weight_offset());
+    }
+
+    qkv_weight_reordered_ = true;
+  }
+
+  o_proj_->load_state_dict(state_dict.get_dict_with_prefix("o_proj."));
+  if (auto w = state_dict.get_tensor("q_norm.weight"); w.defined()) {
+    q_norm_->load_state_dict(StateDict({{"weight", w}}));
+  }
+  if (auto w = state_dict.get_tensor("k_norm.weight"); w.defined()) {
+    k_norm_->load_state_dict(StateDict({{"weight", w}}));
+  }
+
+  // Gemma RMSNorm uses (1 + w) as the scale factor, but the fused kernel
+  // uses standard RMSNorm (w only). Pre-add 1 so the fused kernel produces
+  // the same result as Qwen3NextRMSNorm (gemma_rms_norm).
+  if (use_fused_qkv_) {
+    if (q_norm_->is_weight_loaded() && !q_norm_weight_adjusted_) {
+      q_norm_->weight().add_(1.0);
+      q_norm_weight_adjusted_ = true;
+    }
+    if (k_norm_->is_weight_loaded() && !k_norm_weight_adjusted_) {
+      k_norm_->weight().add_(1.0);
+      k_norm_weight_adjusted_ = true;
+    }
+  }
+}
+
+}  // namespace layer
+}  // namespace xllm

ex_engine/xllm_layers/npu_torch/qwen3_next_attention.h

@@ -0,0 +1,88 @@
+/* Copyright 2025-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 <torch/torch.h>
+
+#include <vector>
+
+#include "attention.h"
+#include "framework/kv_cache/kv_cache.h"
+#include "framework/model/model_args.h"
+#include "framework/parallel_state/parallel_args.h"
+#include "framework/quant_args.h"
+#include "framework/state_dict/state_dict.h"
+#include "kernels/ops_api.h"
+#include "layers/common/linear.h"
+#include "layers/common/partial_rotary_embedding.h"
+#include "layers/common/qwen3_next_rms_norm.h"
+
+namespace xllm {
+namespace layer {
+
+class Qwen3NextAttentionImpl : public torch::nn::Module {
+ public:
+  Qwen3NextAttentionImpl() = default;
+  Qwen3NextAttentionImpl(const ModelArgs& args,
+                         const QuantArgs& quant_args,
+                         const ParallelArgs& parallel_args,
+                         const torch::TensorOptions& options,
+                         int32_t layer_id);
+
+  torch::Tensor forward(const torch::Tensor& positions,
+                        const torch::Tensor& hidden_states,
+                        const AttentionMetadata& attn_metadata,
+                        KVCache& kv_cache,
+                        const torch::Tensor& mrope_cos_sin);
+
+  torch::Tensor build_mrope_cos_sin(const torch::Tensor& positions) const;
+
+  void load_state_dict(const StateDict& state_dict);
+
+ private:
+  int64_t num_heads_;
+  int64_t num_kv_heads_;
+  int64_t num_kv_head_replicas_;
+  int64_t head_dim_;
+  int64_t q_size_;
+  int64_t kv_size_;
+  float scaling_;
+  bool attn_output_gate_;
+  int32_t layer_id_;
+  int32_t rank_;
+  int64_t rotary_dim_;
+  float rms_norm_eps_;
+  bool use_fused_qkv_;
+  bool is_interleaved_;
+  bool qkv_weight_reordered_ = false;
+  bool q_norm_weight_adjusted_ = false;
+  bool k_norm_weight_adjusted_ = false;
+  std::vector<int64_t> mrope_section_;
+  torch::Tensor mrope_gather_pattern_;
+
+  QKVParallelLinear qkv_proj_{nullptr};
+  RowParallelLinear o_proj_{nullptr};
+
+  Qwen3NextRMSNorm q_norm_{nullptr};
+  Qwen3NextRMSNorm k_norm_{nullptr};
+
+  Attention attn_{nullptr};
+  PartialRotaryEmbedding rotary_emb_{nullptr};
+};
+TORCH_MODULE(Qwen3NextAttention);
+
+}  // namespace layer
+}  // namespace xllm

ex_engine/xllm_layers/npu_torch/qwen3_next_decoder_layer_impl.cpp

@@ -0,0 +1,41 @@
+/* Copyright 2025-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 "qwen3_next_decoder_layer_impl.h"
+
+namespace xllm {
+namespace layer {
+
+Qwen3NextDecoderLayerImpl::Qwen3NextDecoderLayerImpl(
+    const ModelContext& context,
+    int32_t layer_id)
+    : Qwen3NextDecoderLayerImpl(context,
+                                layer_id,
+                                std::make_shared<Qwen3NextGatedDeltaNetImpl>(
+                                    context.get_model_args(),
+                                    context.get_quant_args(),
+                                    context.get_parallel_args(),
+                                    context.get_tensor_options())) {}
+
+Qwen3NextDecoderLayerImpl::Qwen3NextDecoderLayerImpl(
+    const ModelContext& context,
+    int32_t layer_id,
+    std::shared_ptr<Qwen3GatedDeltaNetBaseImpl> linear_attention_module)
+    : Qwen3HybridDecoderLayerImplBase(context,
+                                      layer_id,
+                                      std::move(linear_attention_module)) {}
+
+}  // namespace layer
+}  // namespace xllm

ex_engine/xllm_layers/npu_torch/qwen3_next_decoder_layer_impl.h

@@ -0,0 +1,38 @@
+/* Copyright 2025-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 "layers/npu_torch/qwen3_next_gated_delta_net.h"
+#include "layers/npu_torch/qwen3_next_hybrid_decoder_layer_base.h"
+
+namespace xllm {
+namespace layer {
+
+class Qwen3NextDecoderLayerImpl : public Qwen3HybridDecoderLayerImplBase {
+ public:
+  explicit Qwen3NextDecoderLayerImpl(const ModelContext& context,
+                                     int32_t layer_id);
+
+ protected:
+  Qwen3NextDecoderLayerImpl(
+      const ModelContext& context,
+      int32_t layer_id,
+      std::shared_ptr<Qwen3GatedDeltaNetBaseImpl> linear_attention_module);
+};
+TORCH_MODULE(Qwen3NextDecoderLayer);
+
+}  // namespace layer
+}  // namespace xllm

ex_engine/xllm_layers/npu_torch/qwen3_next_gated_delta_net.cpp

@@ -0,0 +1,118 @@
+/* Copyright 2025-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 "qwen3_next_gated_delta_net.h"
+
+#include <glog/logging.h>
+
+namespace xllm {
+namespace layer {
+
+Qwen3NextGatedDeltaNetImpl::Qwen3NextGatedDeltaNetImpl(
+    const ModelArgs& args,
+    const QuantArgs& quant_args,
+    const ParallelArgs& parallel_args,
+    const torch::TensorOptions& options)
+    : Qwen3NextGatedDeltaNetImpl(args,
+                                 quant_args,
+                                 parallel_args,
+                                 options,
+                                 /*init_projections=*/true) {}
+
+Qwen3NextGatedDeltaNetImpl::Qwen3NextGatedDeltaNetImpl(
+    const ModelArgs& args,
+    const QuantArgs& quant_args,
+    const ParallelArgs& parallel_args,
+    const torch::TensorOptions& options,
+    bool init_projections)
+    : Qwen3GatedDeltaNetBaseImpl(args, quant_args, parallel_args, options) {
+  if (init_projections) {
+    init_next_projections(args, quant_args, parallel_args, options);
+  }
+}
+
+void Qwen3NextGatedDeltaNetImpl::init_next_projections(
+    const ModelArgs& args,
+    const QuantArgs& quant_args,
+    const ParallelArgs& parallel_args,
+    const torch::TensorOptions& options) {
+  // QKVZ projection used by Qwen3-Next linear attention.
+  qkvz_proj_ = register_module("in_proj_qkvz",
+                               ColumnParallelLinear(args.hidden_size(),
+                                                    k_size_ * 2 + v_size_ * 2,
+                                                    /*bias=*/false,
+                                                    /*gather_output=*/false,
+                                                    quant_args,
+                                                    parallel_args.tp_group_,
+                                                    options));
+  // BA projection used to derive gating and beta terms.
+  ba_proj_ = register_module("in_proj_ba",
+                             ColumnParallelLinear(args.hidden_size(),
+                                                  num_v_heads_ * 2,
+                                                  /*bias=*/false,
+                                                  /*gather_output=*/false,
+                                                  quant_args,
+                                                  parallel_args.tp_group_,
+                                                  options));
+}
+
+std::pair<torch::Tensor, torch::Tensor>
+Qwen3NextGatedDeltaNetImpl::project_decode_inputs(
+    const torch::Tensor& hidden_states) {
+  auto qkvz = qkvz_proj_->forward(hidden_states);
+  auto ba = ba_proj_->forward(hidden_states);
+  return {qkvz.view({qkvz.size(0), -1, qkvz.size(-1)}),
+          ba.view({ba.size(0), -1, ba.size(-1)})};
+}
+
+std::pair<torch::Tensor, torch::Tensor>
+Qwen3NextGatedDeltaNetImpl::project_flat_inputs(
+    const torch::Tensor& hidden_states) {
+  return {qkvz_proj_->forward(hidden_states), ba_proj_->forward(hidden_states)};
+}
+
+void Qwen3NextGatedDeltaNetImpl::load_state_dict(const StateDict& state_dict) {
+  load_projection_state_dict(state_dict);
+  load_common_state_dict(state_dict);
+}
+
+void Qwen3NextGatedDeltaNetImpl::load_projection_state_dict(
+    const StateDict& state_dict) {
+  auto qkvz_state_dict = state_dict.get_dict_with_prefix("in_proj_qkvz.");
+  if (qkvz_state_dict.size() > 0 && !qkvz_proj_->is_weight_loaded()) {
+    qkvz_proj_->load_state_dict(qkvz_state_dict);
+  }
+
+  auto ba_state_dict = state_dict.get_dict_with_prefix("in_proj_ba.");
+  if (ba_state_dict.size() > 0 && !ba_proj_->is_weight_loaded()) {
+    ba_proj_->load_state_dict(ba_state_dict);
+  }
+}
+
+void Qwen3NextGatedDeltaNetImpl::verify_loaded_weights(
+    const std::string& prefix) const {
+  verify_projection_weights(prefix);
+  verify_common_loaded_weights(prefix);
+}
+
+void Qwen3NextGatedDeltaNetImpl::verify_projection_weights(
+    const std::string& prefix) const {
+  CHECK(qkvz_proj_ && qkvz_proj_->is_weight_loaded())
+      << "Missing required weight after all shards loaded: " << prefix
+      << "in_proj_qkvz.weight";
+  CHECK(ba_proj_ && ba_proj_->is_weight_loaded())
+      << "Missing required weight after all shards loaded: " << prefix
+      << "in_proj_ba.weight";
+}
+
+}  // namespace layer
+}  // namespace xllm

ex_engine/xllm_layers/npu_torch/qwen3_next_gated_delta_net.h

@@ -0,0 +1,66 @@
+/* Copyright 2025-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 <torch/torch.h>
+
+#include <string>
+#include <utility>
+
+#include "qwen3_gated_delta_net_base.h"
+
+namespace xllm {
+namespace layer {
+
+class Qwen3NextGatedDeltaNetImpl : public Qwen3GatedDeltaNetBaseImpl {
+ public:
+  Qwen3NextGatedDeltaNetImpl() = default;
+  Qwen3NextGatedDeltaNetImpl(const ModelArgs& args,
+                             const QuantArgs& quant_args,
+                             const ParallelArgs& parallel_args,
+                             const torch::TensorOptions& options);
+
+  void load_state_dict(const StateDict& state_dict) override;
+  void verify_loaded_weights(const std::string& prefix) const override;
+
+ protected:
+  Qwen3NextGatedDeltaNetImpl(const ModelArgs& args,
+                             const QuantArgs& quant_args,
+                             const ParallelArgs& parallel_args,
+                             const torch::TensorOptions& options,
+                             bool init_projections);
+
+  std::pair<torch::Tensor, torch::Tensor> project_decode_inputs(
+      const torch::Tensor& hidden_states) override;
+  std::pair<torch::Tensor, torch::Tensor> project_flat_inputs(
+      const torch::Tensor& hidden_states) override;
+
+  virtual void load_projection_state_dict(const StateDict& state_dict);
+  virtual void verify_projection_weights(const std::string& prefix) const;
+
+  void init_next_projections(const ModelArgs& args,
+                             const QuantArgs& quant_args,
+                             const ParallelArgs& parallel_args,
+                             const torch::TensorOptions& options);
+
+ private:
+  ColumnParallelLinear qkvz_proj_{nullptr};
+  ColumnParallelLinear ba_proj_{nullptr};
+};
+TORCH_MODULE(Qwen3NextGatedDeltaNet);
+
+}  // namespace layer
+}  // namespace xllm

ex_engine/xllm_layers/npu_torch/qwen3_next_hybrid_decoder_layer_base.cpp

@@ -0,0 +1,176 @@
+/* Copyright 2025-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 "qwen3_next_hybrid_decoder_layer_base.h"
+
+#include <algorithm>
+#include <optional>
+#include <tuple>
+
+#include "common/flash_comm1_context.h"
+
+namespace xllm {
+namespace layer {
+
+Qwen3HybridDecoderLayerImplBase::Qwen3HybridDecoderLayerImplBase(
+    const ModelContext& context,
+    int32_t layer_id,
+    std::shared_ptr<Qwen3GatedDeltaNetBaseImpl> linear_attention_module) {
+  const auto& model_args = context.get_model_args();
+  const auto& quant_args = context.get_quant_args();
+  const auto& parallel_args = context.get_parallel_args();
+  const auto& options = context.get_tensor_options();
+  const bool use_full_attention = is_full_attention_layer(model_args, layer_id);
+
+  // Initialize attention layers
+  if (use_full_attention) {
+    attention_ = register_module(
+        "self_attn",
+        Qwen3NextAttention(
+            model_args, quant_args, parallel_args, options, layer_id));
+  } else {
+    linear_attention_ =
+        register_module("linear_attn", std::move(linear_attention_module));
+  }
+
+  // Initialize norm layers
+  input_norm_ = register_module(
+      "input_layernorm",
+      Qwen3NextRMSNorm(
+          model_args.hidden_size(), model_args.rms_norm_eps(), options));
+
+  post_norm_ = register_module(
+      "post_attention_layernorm",
+      Qwen3NextRMSNorm(
+          model_args.hidden_size(), model_args.rms_norm_eps(), options));
+
+  // Initialize mlp
+  auto mlp_only_layers = model_args.mlp_only_layers();
+  if ((std::count(mlp_only_layers.begin(), mlp_only_layers.end(), layer_id) ==
+       0) &&
+      model_args.n_routed_experts() > 0 &&
+      (layer_id + 1) % model_args.decoder_sparse_step() == 0) {
+    moe_mlp_ = register_module("mlp",
+                               FusedMoE(model_args,
+                                        FusedMoEArgs{.is_gated = true},
+                                        quant_args,
+                                        parallel_args,
+                                        options));
+  } else {
+    mlp_ = register_module("mlp",
+                           DenseMLP(model_args.hidden_size(),
+                                    model_args.intermediate_size(),
+                                    true,
+                                    false,
+                                    model_args.hidden_act(),
+                                    /*enable_result_reduction=*/true,
+                                    quant_args,
+                                    parallel_args.tp_group_,
+                                    options));
+  }
+}
+
+void Qwen3HybridDecoderLayerImplBase::load_state_dict(
+    const StateDict& state_dict) {
+  if (attention_) {
+    attention_->load_state_dict(state_dict.get_dict_with_prefix("self_attn."));
+  } else {
+    linear_attention_->load_state_dict(
+        state_dict.get_dict_with_prefix("linear_attn."));
+  }
+  input_norm_->load_state_dict(
+      state_dict.get_dict_with_prefix("input_layernorm."));
+  post_norm_->load_state_dict(
+      state_dict.get_dict_with_prefix("post_attention_layernorm."));
+  if (moe_mlp_) {
+    moe_mlp_->load_state_dict(state_dict.get_dict_with_prefix("mlp."));
+  } else {
+    mlp_->load_state_dict(state_dict.get_dict_with_prefix("mlp."));
+  }
+}
+
+void Qwen3HybridDecoderLayerImplBase::verify_loaded_weights(
+    const std::string& prefix) const {
+  if (linear_attention_) {
+    linear_attention_->verify_loaded_weights(prefix + "linear_attn.");
+  }
+}
+
+torch::Tensor Qwen3HybridDecoderLayerImplBase::forward(
+    torch::Tensor& x,
+    std::optional<torch::Tensor>& residual,
+    torch::Tensor& positions,
+    const AttentionMetadata& attn_metadata,
+    KVCache& kv_cache,
+    const ModelInputParams& input_params,
+    const torch::Tensor& mrope_cos_sin) {
+  const FlashComm1Context* fc1_ctx = get_current_flash_comm1_context();
+  // Pre-attention norm
+  if (!residual.has_value()) {
+    residual = x;
+    x = std::get<0>(input_norm_->forward(x));
+  } else {
+    if (fc1_ctx && is_sequence_sharded(*fc1_ctx) &&
+        residual.value().size(0) != x.size(0)) {
+      residual = maybe_shard_residual(residual.value(), *fc1_ctx);
+    }
+    if (fc1_ctx && is_sequence_sharded(*fc1_ctx)) {
+      CHECK_EQ(residual.value().size(0), x.size(0))
+          << "FC1 input residual and hidden states must share the same "
+          << "padded local sequence layout.";
+    }
+    std::tie(x, residual) = input_norm_->forward(x, residual);
+  }
+
+  // Attention
+  if (attention_) {
+    x = attention_->forward(
+        positions, x, attn_metadata, kv_cache, mrope_cos_sin);
+  } else {
+    x = linear_attention_->forward(x, attn_metadata, kv_cache, input_params);
+  }
+
+  // Post-attention norm
+  // Ensure the residual layout matches the attention output before post_norm.
+  if (fc1_ctx && is_sequence_sharded(*fc1_ctx) && residual.has_value() &&
+      residual.value().size(0) != x.size(0)) {
+    residual = maybe_shard_residual(residual.value(), *fc1_ctx);
+    CHECK_EQ(residual.value().size(0), x.size(0))
+        << "FC1 post-attention residual and hidden states must share the same "
+        << "padded local sequence layout.";
+  }
+
+  std::tie(x, residual) = post_norm_->forward(x, residual);
+
+  // MLP forward
+  if (moe_mlp_) {
+    x = moe_mlp_(x, input_params);
+  } else {
+    x = mlp_(x);
+  }
+
+  return x;
+}
+
+torch::Tensor Qwen3HybridDecoderLayerImplBase::build_mrope_cos_sin(
+    const torch::Tensor& positions) const {
+  if (attention_) {
+    return attention_->build_mrope_cos_sin(positions);
+  }
+  return {};
+}
+
+}  // namespace layer
+}  // namespace xllm

ex_engine/xllm_layers/npu_torch/qwen3_next_hybrid_decoder_layer_base.h

@@ -0,0 +1,90 @@
+/* Copyright 2025-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 <torch/torch.h>
+
+#include <memory>
+#include <string>
+
+#include "framework/kv_cache/kv_cache.h"
+#include "framework/model/model_input_params.h"
+#include "framework/model_context.h"
+#include "framework/state_dict/state_dict.h"
+#include "layers/common/dense_mlp.h"
+#include "layers/common/qwen3_next_rms_norm.h"
+#include "layers/npu_torch/fused_moe.h"
+#include "layers/npu_torch/qwen3_gated_delta_net_base.h"
+#include "layers/npu_torch/qwen3_next_attention.h"
+
+namespace xllm {
+namespace layer {
+
+class Qwen3HybridDecoderLayerModule : public torch::nn::Module {
+ public:
+  virtual void load_state_dict(const StateDict& state_dict) = 0;
+  virtual void verify_loaded_weights(const std::string& prefix) const = 0;
+  virtual torch::Tensor forward(torch::Tensor& x,
+                                std::optional<torch::Tensor>& residual,
+                                torch::Tensor& positions,
+                                const AttentionMetadata& attn_metadata,
+                                KVCache& kv_cache,
+                                const ModelInputParams& input_params,
+                                const torch::Tensor& mrope_cos_sin = {}) = 0;
+  virtual torch::Tensor build_mrope_cos_sin(
+      const torch::Tensor& positions) const {
+    return {};
+  }
+};
+
+using Qwen3HybridDecoderLayerModulePtr =
+    std::shared_ptr<Qwen3HybridDecoderLayerModule>;
+
+class Qwen3HybridDecoderLayerImplBase : public Qwen3HybridDecoderLayerModule {
+ public:
+  explicit Qwen3HybridDecoderLayerImplBase(
+      const ModelContext& context,
+      int32_t layer_id,
+      std::shared_ptr<Qwen3GatedDeltaNetBaseImpl> linear_attention_module);
+
+  void load_state_dict(const StateDict& state_dict) override;
+
+  void verify_loaded_weights(const std::string& prefix) const override;
+
+  torch::Tensor forward(torch::Tensor& x,
+                        std::optional<torch::Tensor>& residual,
+                        torch::Tensor& positions,
+                        const AttentionMetadata& attn_metadata,
+                        KVCache& kv_cache,
+                        const ModelInputParams& input_params,
+                        const torch::Tensor& mrope_cos_sin = {}) override;
+
+  torch::Tensor build_mrope_cos_sin(
+      const torch::Tensor& positions) const override;
+
+ protected:
+  Qwen3NextAttention attention_{nullptr};
+  std::shared_ptr<Qwen3GatedDeltaNetBaseImpl> linear_attention_;
+
+  DenseMLP mlp_{nullptr};
+  FusedMoE moe_mlp_{nullptr};
+
+  Qwen3NextRMSNorm input_norm_{nullptr};
+  Qwen3NextRMSNorm post_norm_{nullptr};
+};
+
+}  // namespace layer
+}  // namespace xllm