feat: import CUDA kernels from xllm/CCCL/FLA upstream repos

Sources cloned and tree'd (no --depth):
  - jd-opensource/xllm: ILU kernels, CUDA kernels, MoE kernels
  - NVIDIA/cccl: CUB tuning/dispatch headers (block-level primitives)
  - fla-org/flash-linear-attention: Triton GDN kernels
  - NVIDIA/cutlass: grouped GEMM reference (read, not copied)
  - Dao-AILab/flash-attention: attention kernel reference (SM80+, read only)

New CUDA kernels (from xllm, SM-agnostic, portable to BI-V100):
  ex_engine/xllm_kernels/cuda/activation.cu    (188 lines) — silu_and_mul, gelu
  ex_engine/xllm_kernels/cuda/norm.cu          (600 lines) — rms_norm, fused_add_rms_norm
  ex_engine/xllm_kernels/cuda/rope.cu          (258 lines) — rotary_embedding
  ex_engine/xllm_kernels/cuda/block_copy.cu    (209 lines) — copy_blocks, swap_blocks
  ex_engine/xllm_kernels/cuda/reshape_paged_cache.cu (101 lines) — KV cache ops
  ex_engine/xllm_kernels/cuda/headers/         (5 headers for compilation)

ILU bridge kernel sources (from xllm, verified SAME as upstream):
  ex_engine/xllm_kernels/ilu/    (10 files, 925 lines total)
  — activation.cpp, attention.cpp, fused_moe.cpp, group_gemm.cpp,
    matmul.cpp, norm.cpp, rope.cpp, ilu_ops_api.h, ixformer.h, utils.h

FLA Triton GDN kernels (for GatedDeltaNet without SM90+ FlashQLA):
  ex_engine/fla_kernels/gated_delta_rule/  (7 files, 2370 lines)
  — chunk_fwd.py (428), chunk.py (487), wy_fast.py (409),
    fused_recurrent.py (392), naive.py (161), gate.py (380)

CCCL sync (12 tuning + 14 dispatch headers updated from NVIDIA/cccl):
  cccl_upstream/cub/cub/device/dispatch/tuning/ — 12 changed files synced
  cccl_upstream/cub/cub/device/dispatch/ — 14 changed dispatch files synced

Compilation targets for real machine (ivcore10):
  1. CUDA kernels: --cuda-gpu-arch=ivcore10 via corex clang/16
  2. ILU bridges: torch.utils.cpp_extension linking ixformer .so
  3. FLA kernels: Triton JIT (if Triton works on BI-V100)

ex_engine/xllm_kernels/ilu/activation.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 "ilu_ops_api.h"
+
+using namespace ixformer;
+
+namespace xllm::kernel::ilu {
+
+void act_and_mul(torch::Tensor out,
+                 torch::Tensor input,
+                 const std::string& act_mode) {
+  if (act_mode == "silu") {
+    infer::silu_and_mul(input, out);
+  } else {
+    LOG(FATAL) << "Unsupported act mode: " << act_mode
+               << ", only support silu, gelu, gelu_tanh";
+  }
+}
+}  // namespace xllm::kernel::ilu

ex_engine/xllm_kernels/ilu/attention.cpp

@@ -0,0 +1,163 @@
+
+/* 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 "ilu_ops_api.h"
+#include "ixinfer.h"
+#include "utils.h"
+
+using namespace ixformer;
+
+namespace xllm::kernel::ilu {
+
+void reshape_paged_cache(torch::Tensor& key,
+                         std::optional<torch::Tensor>& value,
+                         torch::Tensor& key_cache,
+                         std::optional<torch::Tensor>& value_cache,
+                         torch::Tensor& slot_mapping) {
+  auto value_ = value.value_or(torch::Tensor());
+  auto value_cache_ = value_cache.value_or(torch::Tensor());
+
+  int64_t key_token_stride = key.stride(0);
+  int64_t value_token_stride = 0;
+  if (value_.defined()) {
+    value_token_stride = value_.stride(0);
+  }
+  slot_mapping = slot_mapping.to(at::kLong);
+  infer::xllm_reshape_and_cache(key,
+                                value_,
+                                key_cache,
+                                value_cache_,
+                                slot_mapping,
+                                key_token_stride,
+                                value_token_stride);
+}
+
+void batch_prefill(torch::Tensor& query,
+                   const torch::Tensor& key,
+                   const std::optional<torch::Tensor>& value,
+                   torch::Tensor& output,
+                   std::optional<torch::Tensor>& output_lse,
+                   const std::optional<torch::Tensor>& q_cu_seq_lens,
+                   const std::optional<torch::Tensor>& kv_cu_seq_lens,
+                   const std::optional<torch::Tensor>& alibi_slope,
+                   const std::optional<torch::Tensor>& attn_bias,
+                   const std::optional<torch::Tensor>& q_quant_scale,
+                   const std::optional<torch::Tensor>& k_quant_scale,
+                   const std::optional<torch::Tensor>& v_quant_scale,
+                   const torch::Tensor& block_tables,
+                   int64_t max_query_len,
+                   int64_t max_seq_len,
+                   float scale,
+                   bool is_causal,
+                   int64_t window_size_left,
+                   int64_t window_size_right,
+                   const std::string& compute_dtype,
+                   bool return_lse) {
+  double softcap = 0.0;
+  bool sqrt_alibi = false;
+  auto q_cu_seq_lens_ = q_cu_seq_lens.value_or(torch::Tensor());
+  auto kv_cu_seq_lens_ = kv_cu_seq_lens.value_or(torch::Tensor());
+  auto q_quant_scale_ = q_quant_scale.value_or(torch::Tensor());
+  auto k_quant_scale_ = k_quant_scale.value_or(torch::Tensor());
+  auto v_quant_scale_ = v_quant_scale.value_or(torch::Tensor());
+  auto block_tables_ = block_tables;
+  auto key_ = key;
+  auto value_ = value.value();
+  infer::ixinfer_flash_attn_unpad_with_block_tables(query,
+                                                    key_,
+                                                    value_,
+                                                    output,
+                                                    block_tables_,
+                                                    q_cu_seq_lens_,
+                                                    kv_cu_seq_lens_,
+                                                    max_query_len,
+                                                    max_seq_len,
+                                                    is_causal,
+                                                    window_size_left,
+                                                    window_size_right,
+                                                    static_cast<double>(scale),
+                                                    softcap,
+                                                    sqrt_alibi,
+                                                    alibi_slope,
+                                                    c10::nullopt,
+                                                    output_lse);
+}
+
+void batch_decode(torch::Tensor& query,
+                  const torch::Tensor& k_cache,
+                  torch::Tensor& output,
+                  const torch::Tensor& block_table,
+                  const torch::Tensor& seq_lens,
+                  const std::optional<torch::Tensor>& v_cache,
+                  std::optional<torch::Tensor>& output_lse,
+                  const std::optional<torch::Tensor>& q_quant_scale,
+                  const std::optional<torch::Tensor>& k_cache_quant_scale,
+                  const std::optional<torch::Tensor>& v_cache_quant_scale,
+                  const std::optional<torch::Tensor>& out_quant_scale,
+                  const std::optional<torch::Tensor>& alibi_slope,
+                  const std::optional<torch::Tensor>& mask,
+                  const std::string& compute_dtype,
+                  int64_t max_seq_len,
+                  int64_t window_size_left,
+                  int64_t window_size_right,
+                  float scale,
+                  bool return_lse,
+                  bool is_causal,
+                  int64_t kv_cache_quant_bit_size) {
+  if (query.dim() == 4) {
+    query =
+        query
+            .view({query.size(0) * query.size(1), query.size(2), query.size(3)})
+            .contiguous();
+  }
+  if (output.dim() == 4) {
+    output = output
+                 .view({output.size(0) * output.size(1),
+                        output.size(2),
+                        output.size(3)})
+                 .contiguous();
+    ;
+  }
+  auto v_cache_ = v_cache.value_or(torch::Tensor());
+  int64_t num_kv_heads = k_cache.size(1);
+  int64_t page_block_size = k_cache.size(2);
+  double softcap = 0.0;
+  bool enable_cuda_graph = false;
+  bool use_sqrt_alibi = false;
+  auto block_table_ = block_table;
+  auto k_cache_ = k_cache;
+  auto seq_lens_ = seq_lens;
+  infer::xllm_paged_attention(output,
+                              query,
+                              k_cache_,
+                              v_cache_,
+                              num_kv_heads,
+                              scale,
+                              block_table_,
+                              seq_lens_,
+                              page_block_size,
+                              max_seq_len,
+                              alibi_slope,
+                              is_causal,
+                              (int32_t)window_size_left,
+                              (int32_t)window_size_right,
+                              softcap,
+                              enable_cuda_graph,
+                              use_sqrt_alibi,
+                              c10::nullopt);
+}
+
+}  // namespace xllm::kernel::ilu

ex_engine/xllm_kernels/ilu/fused_moe.cpp

@@ -0,0 +1,99 @@
+/* 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 <glog/logging.h>
+
+#include "ilu_ops_api.h"
+
+namespace xllm::kernel::ilu {
+
+std::tuple<torch::Tensor, torch::Tensor> moe_active_topk(
+    const torch::Tensor& input,
+    int64_t topk,
+    int64_t num_expert_group,
+    int64_t topk_group,
+    bool normalize,
+    const std::optional<torch::Tensor>& mask,
+    const std::string& normed_by,
+    const std::string& scoring_func,
+    double route_scale,
+    const std::optional<torch::Tensor>& e_score_correction_bias) {
+  torch::Tensor input_ = input.to(torch::kFloat32);
+  auto reduce_weight =
+      torch::empty({input.size(0), topk},
+                   torch::dtype(torch::kFloat).device(input.device()));
+  auto topk_indices =
+      torch::empty({input.size(0), topk},
+                   torch::dtype(torch::kInt32).device(input.device()));
+  auto token_expert_indices =
+      torch::empty({input.size(0), topk},
+                   torch::dtype(torch::kInt32).device(input.device()));
+
+  infer::topk_softmax(
+      reduce_weight, topk_indices, token_expert_indices, input_, false);
+
+  auto tt = reduce_weight.sum(-1);
+  if (normalize) {
+    reduce_weight = reduce_weight / reduce_weight.sum(-1).unsqueeze(-1);
+  }
+  return std::make_tuple(reduce_weight, topk_indices);
+}
+
+std::vector<torch::Tensor> moe_gen_idx(torch::Tensor& expert_id,
+                                       int64_t expert_num) {
+  auto src_dst = expert_id.new_empty({expert_id.numel()});
+  auto dst_src = torch::empty_like(src_dst);
+  auto expert_sizes_gpu = expert_id.new_empty({expert_num});
+  auto expert_sizes_gpu_cumsum = expert_id.new_zeros({expert_id.numel() + 1});
+  infer::moe_compute_token_index_api(expert_id,
+                                     src_dst,
+                                     dst_src,
+                                     expert_sizes_gpu,
+                                     /*expert_mask=*/std::nullopt,
+                                     /*expert_sizes_cpu*/ std::nullopt,
+                                     /*expert_sizes_gpu*/ std::nullopt,
+                                     0,
+                                     expert_num,
+                                     expert_num);
+
+  expert_sizes_gpu_cumsum = expert_sizes_gpu.cumsum(-1);
+  return {src_dst, dst_src, expert_sizes_gpu, expert_sizes_gpu_cumsum};
+}
+
+torch::Tensor moe_expand_input(const torch::Tensor& input,
+                               const torch::Tensor& gather_index,
+                               const torch::Tensor& combine_idx,
+                               int64_t topk) {
+  int64_t dst_tokens = input.size(0) * topk;
+  auto output = input.new_empty({dst_tokens, input.size(1)});
+  infer::moe_expand_input(
+      output, input, combine_idx, gather_index, dst_tokens, topk);
+
+  return output;
+}
+
+torch::Tensor moe_combine_result(torch::Tensor& input, torch::Tensor& weight) {
+  input = input.view({-1, weight.size(1), input.size(1)});
+  auto output = input.new_empty({input.size(0), input.size(2)});
+  infer::moe_output_reduce_sum(output,
+                               input,
+                               weight,
+                               /*mask=*/std::nullopt,
+                               /*extra_residual*/ std::nullopt,
+                               /*scaling_factor=*/1.0);
+  return output;
+}
+
+}  // namespace xllm::kernel::ilu

ex_engine/xllm_kernels/ilu/group_gemm.cpp

@@ -0,0 +1,39 @@
+/* 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 "ilu_ops_api.h"
+
+namespace xllm::kernel::ilu {
+
+torch::Tensor group_gemm(torch::Tensor& input,
+                         torch::Tensor& weight,
+                         torch::Tensor& tokens_per_experts,
+                         const std::optional<torch::Tensor>& dst_to_src,
+                         torch::Tensor& output) {
+  infer::moe_w16a16_group_gemm(
+      output,
+      input,
+      weight,
+      tokens_per_experts,
+      dst_to_src,
+      /*bias=*/std::nullopt,
+      /*format=*/"TN",
+      /*persistent=*/0,
+      /*output_n=*/tokens_per_experts.sum().item<int64_t>());
+
+  return output;
+}
+
+}  // namespace xllm::kernel::ilu

ex_engine/xllm_kernels/ilu/ilu_ops_api.h

@@ -0,0 +1,153 @@
+/* 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 <ATen/DynamicLibrary.h>
+#include <ATen/core/dispatch/Dispatcher.h>
+#include <cuda_runtime.h>
+#include <glog/logging.h>
+#include <torch/all.h>
+
+#include <optional>
+
+#include "ATen/Tensor.h"
+#include "ATen/cuda/CUDAEvent.h"
+#include "c10/core/Device.h"
+#include "c10/core/DeviceGuard.h"
+#include "c10/core/GradMode.h"
+#include "c10/core/InferenceMode.h"
+#include "c10/core/MemoryFormat.h"
+#include "c10/core/ScalarType.h"
+#include "c10/core/TensorOptions.h"
+#include "c10/cuda/CUDAFunctions.h"
+#include "c10/cuda/CUDAGuard.h"
+#include "c10/cuda/CUDAStream.h"
+#include "ixformer.h"
+#include "kernels/kernels.h"
+
+// #include "utils.h"
+using namespace ixformer;
+
+namespace xllm::kernel::ilu {
+
+void apply_rope_pos_ids_cos_sin_cache(torch::Tensor& query,
+                                      torch::Tensor& key,
+                                      torch::Tensor& cos_sin_cache,
+                                      torch::Tensor& positions,
+                                      bool interleave);
+
+// act_mode only support silu, gelu, gelu_tanh
+void act_and_mul(torch::Tensor out,
+                 torch::Tensor input,
+                 const std::string& act_mode);
+
+void reshape_paged_cache(
+    torch::Tensor& key,                   //  (num_tokens, num_heads, head_size)
+    std::optional<torch::Tensor>& value,  // (num_tokens, num_heads, head_size)
+    torch::Tensor& key_cache,  // (num_blocks, num_heads, block_size, head_size)
+    std::optional<torch::Tensor>&
+        value_cache,  // (num_blocks, num_heads, block_size, head_size)
+    torch::Tensor& slot_mapping);  //(num_tokens)
+
+void batch_prefill(torch::Tensor& query,
+                   const torch::Tensor& key,
+                   const std::optional<torch::Tensor>& value,
+                   torch::Tensor& output,
+                   std::optional<torch::Tensor>& output_lse,
+                   const std::optional<torch::Tensor>& q_cu_seq_lens,
+                   const std::optional<torch::Tensor>& kv_cu_seq_lens,
+                   const std::optional<torch::Tensor>& alibi_slope,
+                   const std::optional<torch::Tensor>& attn_bias,
+                   const std::optional<torch::Tensor>& q_quant_scale,
+                   const std::optional<torch::Tensor>& k_quant_scale,
+                   const std::optional<torch::Tensor>& v_quant_scale,
+                   const torch::Tensor& block_tables,
+                   int64_t max_query_len,
+                   int64_t max_seq_len,
+                   float scale,
+                   bool is_causal,
+                   int64_t window_size_left,
+                   int64_t window_size_right,
+                   const std::string& compute_dtype,
+                   bool return_lse);
+
+void batch_decode(torch::Tensor& query,
+                  const torch::Tensor& k_cache,
+                  torch::Tensor& output,
+                  const torch::Tensor& block_table,
+                  const torch::Tensor& seq_lens,
+                  const std::optional<torch::Tensor>& v_cache,
+                  std::optional<torch::Tensor>& output_lse,
+                  const std::optional<torch::Tensor>& q_quant_scale,
+                  const std::optional<torch::Tensor>& k_cache_quant_scale,
+                  const std::optional<torch::Tensor>& v_cache_quant_scale,
+                  const std::optional<torch::Tensor>& out_quant_scale,
+                  const std::optional<torch::Tensor>& alibi_slope,
+                  const std::optional<torch::Tensor>& mask,
+                  const std::string& compute_dtype,
+                  int64_t max_seq_len,
+                  int64_t window_size_left,
+                  int64_t window_size_right,
+                  float scale,
+                  bool return_lse,
+                  bool is_causal,
+                  int64_t kv_cache_quant_bit_size);
+
+void residual_layer_norm(torch::Tensor& input,
+                         torch::Tensor& output,
+                         std::optional<torch::Tensor>& residual,
+                         torch::Tensor& weight,
+                         std::optional<torch::Tensor>& bias,
+                         std::optional<torch::Tensor>& residual_out,
+                         double eps);
+
+void rms_norm(torch::Tensor& output,
+              torch::Tensor& input,
+              torch::Tensor& weight,
+              double eps);
+
+torch::Tensor matmul(torch::Tensor a,
+                     torch::Tensor b,
+                     std::optional<torch::Tensor> bias);
+
+std::tuple<torch::Tensor, torch::Tensor> moe_active_topk(
+    const torch::Tensor& input,
+    int64_t topk,
+    int64_t num_expert_group,
+    int64_t topk_group,
+    bool normalize,
+    const std::optional<torch::Tensor>& mask,
+    const std::string& normed_by,
+    const std::string& scoring_func,
+    double route_scale,
+    const std::optional<torch::Tensor>& e_score_correction_bias);
+
+std::vector<torch::Tensor> moe_gen_idx(torch::Tensor& expert_id,
+                                       int64_t expert_num);
+
+torch::Tensor moe_expand_input(const torch::Tensor& input,
+                               const torch::Tensor& gather_index,
+                               const torch::Tensor& combine_idx,
+                               int64_t topk);
+
+torch::Tensor group_gemm(torch::Tensor& input,
+                         torch::Tensor& weight,
+                         torch::Tensor& tokens_per_experts,
+                         const std::optional<torch::Tensor>& dst_to_src,
+                         torch::Tensor& output);
+
+torch::Tensor moe_combine_result(torch::Tensor& input, torch::Tensor& weight);
+}  // namespace xllm::kernel::ilu

ex_engine/xllm_kernels/ilu/ixformer.h

@@ -0,0 +1,147 @@
+/* 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 <torch/all.h>
+
+#include "ATen/Tensor.h"
+#include "utils.h"
+
+namespace ixformer::infer {
+torch::Tensor ixinfer_flash_attn_unpad_with_block_tables(
+    torch::Tensor& query,
+    torch::Tensor& key_cache,
+    torch::Tensor& value_cache,
+    torch::Tensor& out,
+    torch::Tensor& block_tables,
+    torch::Tensor& cu_seq_q,
+    torch::Tensor& cu_seq_k,
+    int64_t max_seq_q,
+    int64_t max_seq_k,
+    bool is_causal,
+    int64_t window_left,
+    int64_t window_right,
+    double scale,
+    double softcap,
+    bool sqrt_alibi,
+    const std::optional<torch::Tensor>& alibi_slopes,
+    const std::optional<torch::Tensor>& sinks,
+    std::optional<torch::Tensor>& lse);
+
+void silu_and_mul(torch::Tensor& input, torch::Tensor& output);
+
+torch::Tensor xllm_paged_attention(
+    torch::Tensor& out,
+    torch::Tensor& query,
+    torch::Tensor& key_cache,
+    torch::Tensor& value_cache,
+    int64_t num_kv_heads,
+    double scale,
+    torch::Tensor& block_tables,
+    torch::Tensor& context_lens,
+    int64_t block_size,
+    int64_t max_context_len,
+    const std::optional<torch::Tensor>& alibi_slopes,
+    bool causal,
+    int32_t window_left,
+    int32_t window_right,
+    double softcap,
+    bool enable_cuda_graph,
+    bool use_sqrt_alibi,
+    const std::optional<torch::Tensor>& sinks);
+
+torch::Tensor ixformer_linear(torch::Tensor& input,
+                              torch::Tensor& weight,
+                              int64_t act_type,
+                              const std::optional<torch::Tensor>& bias,
+                              const std::optional<torch::Tensor>& out,
+                              const std::optional<bool> persistent);
+
+torch::Tensor ixformer_linear_ex(torch::Tensor& input,
+                                 torch::Tensor& weight,
+                                 const c10::optional<torch::Tensor>& bias,
+                                 const c10::optional<torch::Tensor>& out);
+
+void xllm_reshape_and_cache(torch::Tensor& key,
+                            torch::Tensor& value,
+                            torch::Tensor& key_cache,
+                            torch::Tensor& value_cache,
+                            torch::Tensor& slot_mapping,
+                            int64_t key_token_stride,
+                            int64_t value_token_stride);
+
+void xllm_rotary_embedding(torch::Tensor& positions,
+                           torch::Tensor& query,
+                           torch::Tensor& key,
+                           int64_t head_size,
+                           torch::Tensor& cos_sin_cache,
+                           bool is_neox);
+
+void residual_rms_norm(torch::Tensor& input,
+                       torch::Tensor& residual,
+                       torch::Tensor& weight,
+                       torch::Tensor& output,
+                       torch::Tensor& residual_output,
+                       const std::optional<torch::Tensor>& fused_bias,
+                       double alpha,
+                       double eps,
+                       bool is_post);
+
+void rms_norm(torch::Tensor& input,
+              torch::Tensor& weight,
+              torch::Tensor& output,
+              const std::optional<torch::Tensor>& fused_bias,
+              double eps);
+
+void topk_softmax(torch::Tensor& topk_weights,
+                  torch::Tensor& topk_indices,
+                  torch::Tensor& token_expert_indices,
+                  torch::Tensor& gating_output,
+                  bool renormalize);
+
+void moe_compute_token_index_api(
+    torch::Tensor& topk_ids,
+    torch::Tensor& src_dst,
+    torch::Tensor& dst_src,
+    torch::Tensor& expert_sizes_gpu,
+    const c10::optional<torch::Tensor>& expert_mask,
+    const c10::optional<torch::Tensor>& expert_sizes_cpu,
+    const c10::optional<torch::Tensor>& expand_tokens_gpu,
+    int64_t start_expert_id,
+    int64_t end_expert_id,
+    int64_t num_experts);
+
+void moe_expand_input(torch::Tensor outputs,
+                      torch::Tensor inputs,
+                      torch::Tensor dst_to_src,
+                      const c10::optional<torch::Tensor>& src_to_dst,
+                      int64_t dst_tokens,
+                      int64_t expand_factor);
+
+void moe_w16a16_group_gemm(torch::Tensor output,
+                           torch::Tensor inputs,
+                           torch::Tensor weights,
+                           torch::Tensor tokens_per_experts,
+                           const c10::optional<torch::Tensor>& dst_to_src,
+                           const c10::optional<torch::Tensor>& bias,
+                           std::string format,
+                           int64_t persistent,
+                           int64_t output_n);
+
+void moe_output_reduce_sum(torch::Tensor outputs,
+                           torch::Tensor inputs,
+                           const c10::optional<torch::Tensor>& mul_weight,
+                           const c10::optional<torch::Tensor>& mask,
+                           const c10::optional<torch::Tensor>& extra_residual,
+                           double scaling_factor);
+}  // namespace ixformer::infer

ex_engine/xllm_kernels/ilu/matmul.cpp

@@ -0,0 +1,73 @@
+/* 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 "ilu_ops_api.h"
+#include "util/env_var.h"
+
+namespace xllm::kernel::ilu {
+
+bool gemv_conditions(const torch::Tensor& input,
+                     const torch::Tensor& weight,
+                     const torch::Tensor& bias,
+                     int64_t gemv_max_batch) {
+  // gemv input:[m,k] weight:[n,k]
+  // 1. m <= gemv_max_batch
+  // 2. k % 32 == 0 && n % 2 == 0
+  // 3. bias is None
+
+  torch::Tensor input_view = input.view({-1, input.size(-1)});
+  torch::Tensor weight_view = weight.view({-1, weight.size(-1)});
+
+  int64_t m = input_view.size(0);
+  int64_t k = input_view.size(1);
+  int64_t n = weight_view.size(0);
+
+  if (bias.defined() == false && m <= gemv_max_batch && k % 32 == 0 &&
+      n % 2 == 0) {
+    return true;
+  }
+  return false;
+}
+
+torch::Tensor matmul(torch::Tensor a,
+                     torch::Tensor b,
+                     std::optional<torch::Tensor> bias) {
+  int64_t act_type = -1;
+  bool persistent = false;
+  std::vector<int64_t> output_shape = a.sizes().vec();
+  if (!output_shape.empty()) {
+    output_shape[output_shape.size() - 1] = b.size(0);
+  }
+  torch::Tensor output = a.new_empty(output_shape);
+
+  bool use_gemv = true;
+  const int64_t gemv_max_batch = 1;
+  const bool disable_infer_gemm_ex =
+      xllm::util::get_bool_env("DISABLE_INFER_GEMM_EX", false);
+
+  use_gemv =
+      use_gemv &&
+      gemv_conditions(a, b, bias.value_or(at::Tensor()), gemv_max_batch) &&
+      !disable_infer_gemm_ex && (act_type == -1);
+
+  if (use_gemv) {
+    output = infer::ixformer_linear_ex(a, b, bias, output);
+  } else {
+    output = infer::ixformer_linear(a, b, act_type, bias, output, persistent);
+  }
+  return output;
+}
+
+}  // namespace xllm::kernel::ilu

ex_engine/xllm_kernels/ilu/norm.cpp

@@ -0,0 +1,51 @@
+/* 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 "ilu_ops_api.h"
+#include "utils.h"
+
+using namespace ixformer;
+
+namespace xllm::kernel::ilu {
+
+void residual_layer_norm(torch::Tensor& input,
+                         torch::Tensor& output,
+                         std::optional<torch::Tensor>& residual,
+                         torch::Tensor& weight,
+                         std::optional<torch::Tensor>& bias,
+                         std::optional<torch::Tensor>& residual_out,
+                         double eps) {
+  auto residual_ = residual.value_or(torch::zeros_like(input));
+  torch::Tensor residual_out_ = residual_out.value_or(torch::zeros_like(input));
+  infer::residual_rms_norm(input,
+                           residual_,
+                           weight,
+                           output,
+                           residual_out_,
+                           bias,
+                           /*alpha=*/1.0,
+                           eps,
+                           false);
+}
+
+void rms_norm(torch::Tensor& output,
+              torch::Tensor& input,
+              torch::Tensor& weight,
+              double eps) {
+  std::optional<torch::Tensor> fused_bias = std::nullopt;
+  infer::rms_norm(input, weight, output, fused_bias, eps);
+}
+
+}  // namespace xllm::kernel::ilu

ex_engine/xllm_kernels/ilu/rope.cpp

@@ -0,0 +1,31 @@
+/* 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 "ilu_ops_api.h"
+#include "utils.h"
+
+namespace xllm::kernel::ilu {
+
+void apply_rope_pos_ids_cos_sin_cache(torch::Tensor& query,
+                                      torch::Tensor& key,
+                                      torch::Tensor& cos_sin_cache,
+                                      torch::Tensor& positions,
+                                      bool interleave) {
+  const int64_t head_size = cos_sin_cache.size(-1);
+  infer::xllm_rotary_embedding(
+      positions, query, key, head_size, cos_sin_cache, !interleave);
+}
+
+}  // namespace xllm::kernel::ilu

ex_engine/xllm_kernels/ilu/utils.h

@@ -0,0 +1,63 @@
+/* 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
+namespace xllm::kernel::ilu {
+#undef check_tensor_contiguous
+#define check_tensor_contiguous(x, type) \
+  TORCH_CHECK(x.scalar_type() == type);  \
+  TORCH_CHECK(x.is_cuda());              \
+  TORCH_CHECK(x.is_contiguous());
+
+#undef check_tensor_half_bf_float
+#define check_tensor_half_bf_float(x)                       \
+  TORCH_CHECK(x.scalar_type() == at::ScalarType::Half ||    \
+              x.scalar_type() == at::ScalarType::Float ||   \
+              x.scalar_type() == at::ScalarType::BFloat16); \
+  TORCH_CHECK(x.is_cuda());
+
+// from torchCheckMsgImpl
+inline const char* ixformer_check_msg_impl(const char* msg) { return msg; }
+// // If there is just 1 user-provided C-string argument, use it.
+
+#define IXFORMER_CHECK_MSG(cond, type, ...)                 \
+  (ixformer_check_msg_impl(                                 \
+      "Expected " #cond                                     \
+      " to be true, but got false.  "                       \
+      "(Could this error message be improved?  If so, "     \
+      "please report an enhancement request to ixformer.)", \
+      ##__VA_ARGS__))
+
+#define IXFORMER_CHECK(cond, ...)                                            \
+  {                                                                          \
+    if (!(cond)) {                                                           \
+      std::cerr << __FILE__ << " (" << __LINE__ << ")"                       \
+                << "-" << __FUNCTION__ << " : "                              \
+                << IXFORMER_CHECK_MSG(cond, "", ##__VA_ARGS__) << std::endl; \
+      throw std::runtime_error("IXFORMER_CHECK ERROR");                      \
+    }                                                                        \
+  }
+
+#undef CUINFER_CHECK
+#define CUINFER_CHECK(func)                                                \
+  do {                                                                     \
+    cuinferStatus_t status = (func);                                       \
+    if (status != CUINFER_STATUS_SUCCESS) {                                \
+      std::cerr << "Error in file " << __FILE__ << " on line " << __LINE__ \
+                << ": " << cuinferGetErrorString(status) << std::endl;     \
+      throw std::runtime_error("CUINFER_CHECK ERROR");                     \
+    }                                                                      \
+  } while (0)
+
+}  // namespace xllm::kernel::ilu