feat: support MiniMax-M2.7 Channel FP8 w8a8 on DCU.

xllm/core/kernels/cuda/rope.cu

@@ -169,9 +169,16 @@ void rotary_embedding(
     torch::Tensor& cos_sin_cache,  // [max_position, rot_dim]
     bool is_neox) {
   // num_tokens = batch_size * seq_len
-  int64_t head_size = cos_sin_cache.size(-1);
+  const int positions_ndim = positions.dim();
+  const int query_ndim = query.dim();
+  // For partial rotary models, e.g. MiniMax-M2 with head_dim=128 and
+  // rotary_dim=64, the cache width is the rotary dimension rather than the
+  // physical per-head stride. When query is already shaped as
+  // [*, num_heads, head_size], infer the real head_size from query itself.
+  int64_t head_size = (query_ndim == positions_ndim + 2)
+                          ? query.size(-1)
+                          : cos_sin_cache.size(-1);
   int64_t num_tokens = positions.numel();
-  int positions_ndim = positions.dim();
 
   // Make sure num_tokens dim is consistent across positions, query, and key
   CHECK(positions_ndim == 1 || positions_ndim == 2)
@@ -209,7 +216,6 @@ void rotary_embedding(
   // Determine head stride: for [*, heads, head_size] use stride of last dim;
   // for flat [*, heads*head_size], heads blocks are contiguous of size
   // head_size
-  int query_ndim = query.dim();
   int64_t head_stride =
       (query_ndim == positions_ndim + 2) ? query.stride(-2) : head_size;
 

xllm/core/kernels/dcu/aiter_quant_adapter.cpp

@@ -17,20 +17,7 @@ limitations under the License.
 
 #include <glog/logging.h>
 
-namespace aiter {
-namespace native {
-
-void dynamic_per_token_scaled_quant(
-    torch::Tensor& out,
-    const torch::Tensor& input,
-    torch::Tensor& scales,
-    const std::optional<at::Tensor>& scale_ub = std::nullopt,
-    bool shuffle_scale = false,
-    const std::optional<at::Tensor>& num_rows = std::nullopt,
-    int num_rows_factor = 1);
-
-}  // namespace native
-}  // namespace aiter
+#include "quant_api.h"
 
 namespace xllm {
 namespace kernel {

xllm/core/kernels/dcu/topk_gate.cpp

@@ -98,6 +98,10 @@ std::tuple<torch::Tensor, torch::Tensor> moe_grouped_topk(
         << "dcu::moe_grouped_topk: correction bias is supported only for "
            "sigmoid scoring";
     torch::Tensor bias = correction_bias.value();
+    // AITER grouped topk requires correction_bias to match gating_output dtype.
+    if (bias.scalar_type() != gating_output.scalar_type()) {
+      bias = bias.to(gating_output.scalar_type());
+    }
     aiter::native::biased_grouped_topk(gating_output,
                                        bias,
                                        topk_weights,

xllm/core/layers/CMakeLists.txt

@@ -47,6 +47,9 @@ cc_library(
     qwen3_vision_layer.h
     qwen3_decoder_layer.h
     qwen3_moe_decoder_layer.h
+    $<$<BOOL:${USE_DCU}>:minimax_m2_attention.h>
+    $<$<BOOL:${USE_DCU}>:minimax_m2_decoder_layer.h>
+    $<$<BOOL:${USE_DCU}>:minimax_m2_tensor_parallel_rms_norm.h>
   SRCS
     oxygen_vision_layer.cpp
     $<$<BOOL:${USE_NPU}>:deepseek_v4_decoder_layer.cpp>
@@ -55,13 +58,17 @@ cc_library(
     qwen2_5_vision_layer.cpp
     qwen3_vision_layer.cpp
     qwen3_moe_decoder_layer.cpp
+    $<$<BOOL:${USE_DCU}>:minimax_m2_attention.cpp>
+    $<$<BOOL:${USE_DCU}>:minimax_m2_decoder_layer.cpp>
+    $<$<BOOL:${USE_DCU}>:minimax_m2_tensor_parallel_rms_norm.cpp>
   DEPS
     $<$<BOOL:${USE_NPU}>:npu_layers>
     $<$<BOOL:${USE_NPU}>:npu_torch_layers>
     $<$<BOOL:${USE_MLU}>:mlu_layers>
-    $<$<BOOL:${USE_ILU}>:ilu_layers>
+    $<$<BOOL:${USE_ILU}>:ilu_layers>  
     $<$<AND:$<BOOL:${USE_CUDA}>,$<NOT:$<BOOL:${USE_MUSA}>>>:cuda_layers>
     $<$<BOOL:${USE_MUSA}>:musa_layers>
+    $<$<BOOL:${USE_DCU}>:dcu_layers>
     :common_layers
     :parallel_state
     :state_dict

xllm/core/layers/dcu/fused_moe.cpp

@@ -144,10 +144,10 @@ FusedMoEImpl::FusedMoEImpl(const ModelArgs& model_args,
   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),
-                           /*requires_grad*/ false);
+    e_score_correction_bias_ = register_parameter(
+        "e_score_correction_bias",
+        torch::empty({num_experts}, options.dtype(torch::kFloat32)),
+        /*requires_grad*/ false);
   }
 
   gate_ = register_module(
@@ -245,7 +245,10 @@ std::pair<torch::Tensor, torch::Tensor> FusedMoEImpl::route_experts(
     const torch::Tensor& router_logits_2d) {
   std::optional<torch::Tensor> e_score_correction_bias = std::nullopt;
   if (e_score_correction_bias_.defined()) {
-    e_score_correction_bias = e_score_correction_bias_;
+    e_score_correction_bias =
+        e_score_correction_bias_.scalar_type() == torch::kFloat32
+            ? e_score_correction_bias_
+            : e_score_correction_bias_.to(torch::kFloat32);
   }
 
   torch::Tensor reduce_weight;
@@ -501,33 +504,25 @@ torch::Tensor FusedMoEImpl::forward_experts(
     auto final_hidden_states = forward_fp8_channelwise_experts(
         hidden_states, hidden_states_2d, router_logits_2d);
 
-    auto current_stream = device_.current_stream();
-    routed_stream_->wait_stream(*current_stream);
-    {
-      torch::StreamGuard stream_guard = routed_stream_->set_stream_guard();
-      if (parallel_args_.ep_size() > 1) {
-        final_hidden_states = parallel_state::reduce(
-            final_hidden_states, parallel_args_.moe_ep_group_);
-      }
-      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 (tp_pg_->world_size() > 1) {
+      final_hidden_states = parallel_state::reduce(final_hidden_states, tp_pg_);
     }
 
     torch::Tensor shared_expert_output;
     if (n_shared_experts_ > 0) {
+      auto current_stream = device_.current_stream();
       shared_stream_->wait_stream(*current_stream);
       {
         torch::StreamGuard stream_guard = shared_stream_->set_stream_guard();
         shared_expert_output = shared_experts_(hidden_states);
         shared_expert_output =
             shared_expert_output.reshape({-1, shared_expert_output.size(-1)});
       }
-    }
 
-    current_stream->wait_stream(*routed_stream_);
-    if (n_shared_experts_ > 0) {
       current_stream->wait_stream(*shared_stream_);
       final_hidden_states += shared_expert_output.reshape(hidden_states_shape);
     }

xllm/core/layers/minimax_m2_attention.cpp

@@ -0,0 +1,162 @@
+/* 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 "minimax_m2_attention.h"
+
+#include <glog/logging.h>
+
+#include <cmath>
+#include <tuple>
+
+namespace xllm {
+namespace layer {
+
+MiniMaxM2AttentionImpl::MiniMaxM2AttentionImpl(const ModelContext& context) {
+  const auto& 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 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());
+
+  CHECK_EQ(total_num_heads % tp_size, 0);
+  num_heads_ = total_num_heads / tp_size;
+  if (total_num_kv_heads >= tp_size) {
+    CHECK_EQ(total_num_kv_heads % tp_size, 0);
+    num_kv_heads_ = total_num_kv_heads / tp_size;
+    num_kv_head_replicas_ = 1;
+  } else {
+    CHECK_EQ(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_;
+  CHECK_EQ(num_heads_ % num_kv_heads_, 0)
+      << "MiniMax-M2 local attention heads must be divisible by local KV "
+         "heads.";
+  scaling_ = std::sqrt(1.0f / static_cast<float>(head_dim_));
+
+  qkv_proj_ = register_module("qkv_proj",
+                              QKVParallelLinear(args.hidden_size(),
+                                                num_heads_,
+                                                num_kv_heads_,
+                                                head_dim_,
+                                                num_kv_head_replicas_,
+                                                /*bias=*/false,
+                                                /*gather_output=*/false,
+                                                parallel_args,
+                                                options,
+                                                quant_args));
+
+  o_proj_ = register_module("o_proj",
+                            RowParallelLinear(total_num_heads * head_dim_,
+                                              args.hidden_size(),
+                                              /*bias=*/false,
+                                              /*input_is_parallelized=*/true,
+                                              /*enable_result_reduction=*/true,
+                                              quant_args,
+                                              parallel_args.tp_group_,
+                                              options));
+
+  q_norm_tp_ = register_module(
+      "q_norm",
+      MiniMaxM2TensorParallelRMSNorm(q_size_,
+                                     total_num_heads * head_dim_,
+                                     /*replica_factor=*/1,
+                                     args.rms_norm_eps(),
+                                     parallel_args.tp_group_,
+                                     options));
+  k_norm_tp_ = register_module(
+      "k_norm",
+      MiniMaxM2TensorParallelRMSNorm(kv_size_,
+                                     total_num_kv_heads * head_dim_,
+                                     num_kv_head_replicas_,
+                                     args.rms_norm_eps(),
+                                     parallel_args.tp_group_,
+                                     options));
+
+  const int64_t rotary_dim =
+      args.rotary_dim() > 0 ? args.rotary_dim() : args.head_dim();
+  rotary_emb_ = register_module("rope",
+                                RotaryEmbedding(rotary_dim,
+                                                args.max_position_embeddings(),
+                                                args.rope_theta(),
+                                                /*interleaved=*/false,
+                                                options));
+
+  attn_ = register_module("attn",
+                          Attention(num_heads_,
+                                    head_dim_,
+                                    scaling_,
+                                    num_kv_heads_,
+                                    args.sliding_window()));
+}
+
+torch::Tensor MiniMaxM2AttentionImpl::forward(
+    const torch::Tensor& positions,
+    const torch::Tensor& hidden_states,
+    const AttentionMetadata& attn_metadata,
+    KVCache& kv_cache) {
+  if (attn_metadata.is_dummy) {
+    return torch::zeros_like(hidden_states);
+  }
+
+  torch::Tensor qkv = qkv_proj_->forward(hidden_states);
+  torch::Tensor q = qkv.slice(/*dim=*/-1, /*start=*/0, /*end=*/q_size_);
+  torch::Tensor k = qkv.slice(/*dim=*/-1,
+                              /*start=*/q_size_,
+                              /*end=*/q_size_ + kv_size_);
+  torch::Tensor v = qkv.slice(/*dim=*/-1, /*start=*/q_size_ + kv_size_);
+
+  const int64_t num_tokens = q.size(0);
+  std::tie(q, k) = forward_minimax_m2_qk_rms_norm(q_norm_tp_, k_norm_tp_, q, k);
+
+  torch::Tensor q_heads =
+      q.view({num_tokens, num_heads_, head_dim_}).contiguous();
+  torch::Tensor k_heads =
+      k.view({num_tokens, num_kv_heads_, head_dim_}).contiguous();
+  torch::Tensor position_ids = positions;
+  if (position_ids.dim() == 2) {
+    position_ids = position_ids[0];
+  }
+  rotary_emb_->forward(
+      q_heads,
+      k_heads,
+      position_ids,
+      attn_metadata.q_cu_seq_lens,
+      attn_metadata.max_query_len,
+      attn_metadata.is_prefill || attn_metadata.is_chunked_prefill);
+
+  torch::Tensor q_flat = q_heads.reshape({num_tokens, q_size_}).contiguous();
+  torch::Tensor k_flat = k_heads.reshape({num_tokens, kv_size_}).contiguous();
+  torch::Tensor v_flat = v.reshape({num_tokens, kv_size_}).contiguous();
+  torch::Tensor out = std::get<0>(
+      attn_->forward(attn_metadata, q_flat, k_flat, v_flat, kv_cache));
+  return o_proj_->forward(out);
+}
+
+void MiniMaxM2AttentionImpl::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."));
+  q_norm_tp_->load_state_dict(state_dict.get_dict_with_prefix("q_norm."));
+  k_norm_tp_->load_state_dict(state_dict.get_dict_with_prefix("k_norm."));
+}
+
+}  // namespace layer
+}  // namespace xllm

xllm/core/layers/minimax_m2_attention.h

@@ -0,0 +1,62 @@
+/* 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 "common/attention.h"
+#include "common/linear.h"
+#include "common/rotary_embedding.h"
+#include "framework/kv_cache/kv_cache.h"
+#include "framework/model_context.h"
+#include "framework/state_dict/state_dict.h"
+#include "minimax_m2_tensor_parallel_rms_norm.h"
+
+namespace xllm {
+namespace layer {
+
+class MiniMaxM2AttentionImpl : public torch::nn::Module {
+ public:
+  MiniMaxM2AttentionImpl() = default;
+  explicit MiniMaxM2AttentionImpl(const ModelContext& context);
+
+  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);
+
+ private:
+  int64_t num_heads_ = 0;
+  int64_t num_kv_heads_ = 0;
+  int64_t num_kv_head_replicas_ = 0;
+  int64_t head_dim_ = 0;
+  int64_t q_size_ = 0;
+  int64_t kv_size_ = 0;
+  float scaling_ = 1.0f;
+
+  QKVParallelLinear qkv_proj_{nullptr};
+  RowParallelLinear o_proj_{nullptr};
+  MiniMaxM2TensorParallelRMSNorm q_norm_tp_{nullptr};
+  MiniMaxM2TensorParallelRMSNorm k_norm_tp_{nullptr};
+  RotaryEmbedding rotary_emb_{nullptr};
+  Attention attn_{nullptr};
+};
+TORCH_MODULE(MiniMaxM2Attention);
+
+}  // namespace layer
+}  // namespace xllm