data: port complete MoE + xllm layer call chains from upstream repos

MoE call chain from ds_vllm (vllm-project/vllm latest):
  ex_engine/moe/ — 20 files, 8736 lines
  - modular_kernel.py (1630 lines) — base classes for modular MoE
  - experts/fused_batched_moe.py (972 lines) — NaiveBatchedExperts
  - prepare_finalize/batched.py (171 lines) — token grouping by expert
  - topk_weight_and_reduce.py (176 lines) — scatter-add finalize
  - fused_moe.py (1740 lines) — main fused_moe dispatch
  - config.py (1407 lines) — FusedMoEQuantConfig
  - activation.py, utils.py, layer.py, etc.

xllm layer code (jd-opensource/xllm):
  ex_engine/xllm_layers/ — 39 files, 5859 lines
  - ilu/fused_moe.cpp (797 lines) — production ixformer 7-step MoE pipeline
  - ilu/attention.cpp (189 lines) — paged_attention + flash_attn bridge
  - npu_torch/qwen3_gated_delta_net_base.cpp (576 lines) — GDN reference
  - common/rms_norm.cpp, rotary_embedding.cpp, activation.cpp, dense_mlp.cpp

xllm ILU kernels — synced 10 files to upstream (diffs from prior edits)

These are reference implementations, NOT hand-written.
Source repos: vllm-project/vllm, jd-opensource/xllm

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