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/cuda/headers/cuda_ops_api.h

@@ -0,0 +1,306 @@
+/* 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 <glog/logging.h>
+
+#include <optional>
+#include <tuple>
+#include <vector>
+
+#include "utils.h"
+
+namespace xllm::kernel::cuda {
+
+// TODO: add head_size parameter
+void rotary_embedding(torch::Tensor& positions,
+                      torch::Tensor& query,
+                      std::optional<torch::Tensor> key,
+                      torch::Tensor& cos_sin_cache,
+                      // int64_t head_size,
+                      bool is_neox);
+
+// 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 slot_ids,   // [n_tokens]
+    torch::Tensor keys,       // [n_tokens, n_kv_heads, head_dim]
+    torch::Tensor values,     // [n_tokens, n_kv_heads, head_dim]
+    torch::Tensor key_cache,  // [n_blocks, block_size, n_heads, head_dim]
+    torch::Tensor value_cache);
+
+void block_copy(torch::Tensor key_cache_ptrs,
+                torch::Tensor value_cache_ptrs,
+                torch::Tensor src_block_indices,
+                torch::Tensor dst_block_indices,
+                torch::Tensor cum_sum,
+                int64_t numel_per_block,
+                torch::ScalarType cache_dtype);
+#if !defined(USE_DCU)
+void batch_prefill(const std::string& uri,
+                   ffi::Array<int64_t> plan_info,
+                   torch::Tensor float_workspace_buffer,
+                   torch::Tensor int_workspace_buffer,
+                   torch::Tensor page_locked_int_workspace_buffer,
+                   torch::Tensor query,
+                   torch::Tensor key,
+                   torch::Tensor value,
+                   torch::Tensor q_cu_seq_lens,
+                   torch::Tensor kv_cu_seq_lens,
+                   int64_t window_left,
+                   double sm_scale,
+                   torch::Tensor output,
+                   std::optional<torch::Tensor>& output_lse,
+                   const std::optional<torch::Tensor>& mask = std::nullopt);
+
+// Wrapper function for batch_prefill that conditionally uses AttentionRunner
+// for piecewise CUDA Graph capture
+void batch_prefill_with_optional_piecewise_capture(
+    const std::string& uri,
+    ffi::Array<int64_t> plan_info,
+    torch::Tensor float_workspace_buffer,
+    torch::Tensor int_workspace_buffer,
+    torch::Tensor page_locked_int_workspace_buffer,
+    torch::Tensor query,
+    torch::Tensor key,
+    torch::Tensor value,
+    torch::Tensor q_cu_seq_lens,
+    torch::Tensor kv_cu_seq_lens,
+    int64_t window_left,
+    double sm_scale,
+    torch::Tensor output,
+    std::optional<torch::Tensor>& output_lse);
+
+void batch_prefill_non_causal(
+    const std::string& uri,
+    ffi::Array<int64_t> plan_info,
+    torch::Tensor float_workspace_buffer,
+    torch::Tensor int_workspace_buffer,
+    torch::Tensor page_locked_int_workspace_buffer,
+    torch::Tensor query,
+    torch::Tensor key,
+    torch::Tensor value,
+    torch::Tensor q_cu_seq_lens,
+    torch::Tensor kv_cu_seq_lens,
+    int64_t window_left,
+    double sm_scale,
+    torch::Tensor output,
+    std::optional<torch::Tensor>& output_lse,
+    const std::optional<torch::Tensor>& mask = std::nullopt);
+
+void batch_chunked_prefill(
+    const std::string& uri,
+    ffi::Array<int64_t> plan_info,
+    torch::Tensor float_workspace_buffer,
+    torch::Tensor int_workspace_buffer,
+    torch::Tensor page_locked_int_workspace_buffer,
+    torch::Tensor query,
+    torch::Tensor k_cache,
+    torch::Tensor v_cache,
+    torch::Tensor paged_kv_indptr,
+    torch::Tensor paged_kv_indices,
+    torch::Tensor paged_kv_last_page_len,
+    int64_t window_left,
+    double sm_scale,
+    torch::Tensor output,
+    std::optional<torch::Tensor>& output_lse,
+    std::optional<torch::Tensor> qo_indptr = std::nullopt,
+    bool causal = true);
+
+void batch_decode(const std::string& uri,
+                  ffi::Array<int64_t> plan_info,
+                  torch::Tensor float_workspace_buffer,
+                  torch::Tensor int_workspace_buffer,
+                  torch::Tensor page_locked_int_workspace_buffer,
+                  torch::Tensor query,
+                  torch::Tensor k_cache,
+                  torch::Tensor v_cache,
+                  torch::Tensor paged_kv_indptr,
+                  torch::Tensor paged_kv_indices,
+                  torch::Tensor paged_kv_last_page_len,
+                  int64_t window_left,
+                  double sm_scale,
+                  torch::Tensor output,
+                  std::optional<torch::Tensor>& output_lse,
+                  bool use_tensor_core,
+                  std::optional<torch::Tensor> qo_indptr = std::nullopt);
+#endif  // !defined(USE_DCU)
+void rms_norm(torch::Tensor output,
+              torch::Tensor input,
+              torch::Tensor weight,
+              double eps);
+
+void fused_add_rms_norm(torch::Tensor& input,     // [..., hidden_size]
+                        torch::Tensor& residual,  // [..., hidden_size]
+                        torch::Tensor& weight,    // [hidden_size]
+                        double epsilon);
+
+torch::Tensor matmul(torch::Tensor a,
+                     torch::Tensor b,
+                     std::optional<torch::Tensor> bias);
+
+void cutlass_scaled_mm(torch::Tensor& c,
+                       torch::Tensor const& a,
+                       torch::Tensor const& b,
+                       torch::Tensor const& a_scales,
+                       torch::Tensor const& b_scales,
+                       std::optional<torch::Tensor> const& bias);
+
+// Static scaled FP8 quantization
+// Quantizes input tensor to FP8 using a pre-computed scale factor
+void static_scaled_fp8_quant(torch::Tensor& out,           // [..., d]
+                             torch::Tensor const& input,   // [..., d]
+                             torch::Tensor const& scale);  // [1]
+
+// FP8 scaled quantize: quantizes input tensor to FP8 e4m3 format
+// Returns: (quantized_output, scale)
+std::tuple<torch::Tensor, torch::Tensor> fp8_scaled_quantize(
+    const torch::Tensor& input,
+    const std::optional<torch::Tensor>& output = std::nullopt,
+    const std::optional<torch::Tensor>& scale = std::nullopt);
+
+// ============================================================================
+// Fused RMSNorm + Static FP8 Quantization
+// ============================================================================
+// These functions combine RMSNorm and FP8 quantization to reduce memory
+// bandwidth by avoiding the intermediate write-back to global memory.
+
+// Fused RMSNorm + Static FP8 Quantization (without residual)
+// Combines RMSNorm normalization and FP8 quantization in a single kernel.
+// This is optimal for the first layer where no residual connection exists.
+void rms_norm_static_fp8_quant(
+    torch::Tensor& out,     // [..., hidden_size], FP8 output
+    torch::Tensor& input,   // [..., hidden_size], input tensor
+    torch::Tensor& weight,  // [hidden_size], RMSNorm weight
+    torch::Tensor& scale,   // [1], FP8 quantization scale
+    double epsilon);        // RMSNorm epsilon
+
+// Fused Add + RMSNorm + Static FP8 Quantization (with residual)
+// Combines residual addition, RMSNorm, and FP8 quantization in a single kernel.
+// The residual tensor is updated in-place with the sum of input and residual.
+void fused_add_rms_norm_static_fp8_quant(
+    torch::Tensor& out,       // [..., hidden_size], FP8 output
+    torch::Tensor& input,     // [..., hidden_size], input tensor
+    torch::Tensor& residual,  // [..., hidden_size], residual (updated in-place)
+    torch::Tensor& weight,    // [hidden_size], RMSNorm weight
+    torch::Tensor& scale,     // [1], FP8 quantization scale
+    double epsilon);          // RMSNorm epsilon
+
+// FP8 scaled matmul for W8A8 quantization using CUTLASS kernels
+// Performs: c = (a @ b.T) with scales applied
+torch::Tensor fp8_scaled_matmul(
+    const torch::Tensor& a,
+    const torch::Tensor& b,
+    const torch::Tensor& a_scale,
+    const torch::Tensor& b_scale,
+    torch::ScalarType output_dtype,
+    const std::optional<torch::Tensor>& bias = std::nullopt,
+    const std::optional<torch::Tensor>& output = std::nullopt);
+
+std::pair<torch::Tensor, torch::Tensor> compute_topk_for_beam_search(
+    torch::Tensor combined_probs,
+    uint32_t batch_size,
+    uint32_t beam_size,
+    uint32_t top_k,
+    torch::Device device);
+
+std::pair<torch::Tensor, torch::Tensor> compute_topk_general(
+    torch::Tensor input,
+    uint32_t batch_size,
+    uint32_t input_length,
+    uint32_t k,
+    torch::Device device);
+
+torch::Tensor air_log_softmax_last_dim(const torch::Tensor& input,
+                                       const torch::Tensor& temperatures);
+
+void fused_qk_norm_rope(
+    torch::Tensor& qkv,   // Combined QKV tensor [num_tokens,
+                          // (num_heads_q+num_heads_k+num_heads_v)*head_dim]
+    int64_t num_heads_q,  // Number of query heads
+    int64_t num_heads_k,  // Number of key heads
+    int64_t num_heads_v,  // Number of value heads
+    int64_t head_dim,     // Dimension per head
+    double eps,           // Epsilon for RMS normalization
+    const torch::Tensor& q_weight,  // RMSNorm weights for query [head_dim]
+    const torch::Tensor& k_weight,  // RMSNorm weights for key [head_dim]
+    const torch::Tensor&
+        cos_sin_cache,  // Cos/sin cache [max_position, rotary_dim]
+    bool interleaved,   // Whether RoPE is applied in interleaved style
+    const torch::Tensor& position_ids  // Position IDs for RoPE [num_tokens]
+);
+
+std::tuple<torch::Tensor, torch::Tensor> moe_fused_topk(
+    torch::Tensor& gating_output,
+    int64_t topk,
+    bool renormalize,
+    const std::optional<torch::Tensor>& correction_bias,
+    const std::string& scoring_func);
+
+torch::Tensor random_sample(const torch::Tensor& probs);
+
+torch::Tensor cutlass_fused_moe(
+    const torch::Tensor& input,                   // [num_tokens, hidden]
+    const torch::Tensor& token_selected_experts,  // [num_tokens, top_k]
+    const torch::Tensor& token_final_scales,      // [num_tokens, top_k]
+    const torch::Tensor&
+        fc1_expert_weights,  // [num_experts, inter_dim, hidden]
+    const torch::Tensor&
+        fc2_expert_weights,  // [num_experts, hidden, inter_dim]
+    torch::ScalarType output_dtype,
+    const std::vector<torch::Tensor>& quant_scales,
+    int32_t tp_size,
+    int32_t tp_rank,
+    int32_t ep_size,
+    int32_t ep_rank,
+    int32_t cluster_size,
+    int32_t cluster_rank,
+    const std::optional<torch::Tensor>& fc1_expert_biases = std::nullopt,
+    const std::optional<torch::Tensor>& fc2_expert_biases = std::nullopt,
+    const std::optional<torch::Tensor>& input_sf = std::nullopt,
+    const std::optional<torch::Tensor>& swiglu_alpha = std::nullopt,
+    const std::optional<torch::Tensor>& swiglu_beta = std::nullopt,
+    const std::optional<torch::Tensor>& swiglu_limit = std::nullopt,
+    const std::optional<torch::Tensor>& output = std::nullopt,
+    bool enable_alltoall = false,
+    bool use_deepseek_fp8_block_scale = false,
+    bool use_w4_group_scaling = false,
+    bool use_mxfp8_act_scaling = false,
+    bool min_latency_mode = false,
+    bool use_packed_weights = false,
+    int32_t tune_max_num_tokens = 8192,
+    ActivationType activation_type = ActivationType::SWIGLU);
+
+// ---- moe_compute_index (moe_compute_index.cu) ----
+// Fused routing index: bincount + argsort replacement.
+// Returns {src_dst, dst_src, expert_sizes}.
+std::tuple<torch::Tensor, torch::Tensor, torch::Tensor> moe_compute_index(
+    const torch::Tensor& expert_id,
+    int64_t num_experts);
+
+// ---- moe_combine_result (moe_combine.cu) ----
+// Fused combine: reorder + weighted sum in one pass.
+torch::Tensor moe_combine_result(const torch::Tensor& gemm2,
+                                 const torch::Tensor& reduce_weight,
+                                 int64_t N,
+                                 int32_t topk);
+
+}  // namespace xllm::kernel::cuda

ex_engine/xllm_kernels/cuda/headers/device_utils.cuh

@@ -0,0 +1,116 @@
+/* Copyright 2025-2026 The xLLM Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    https://github.com/jd-opensource/xllm/blob/main/LICENSE
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#pragma once
+
+#if defined(USE_DCU)
+#include <hip/amd_detail/amd_hip_bf16.h>
+
+#include <hipcub/hipcub.hpp>
+
+namespace cub = hipcub;
+#else
+#include <cub/cub.cuh>
+#if CUB_VERSION >= 200800
+#include <cuda/functional>
+#endif
+#endif
+
+namespace xllm::kernel::cuda {
+#if !defined(USE_DCU)
+using BFloat16Type = __nv_bfloat16;
+
+#define WARP_SIZE 32
+#define XLLM_KERNEL_ATTR(MAX_THREADS)
+#else
+using BFloat16Type = hip_bfloat16;
+
+#define WARP_SIZE 64
+#define XLLM_KERNEL_ATTR(MAX_THREADS) __launch_bounds__(MAX_THREADS, 1)
+#endif
+#define MAX(a, b) ((a) > (b) ? (a) : (b))
+#define MIN(a, b) ((a) < (b) ? (a) : (b))
+
+// Aligned array type
+template <typename T,
+          // Number of elements in the array
+          int N,
+          // Alignment requirement in bytes
+          int Alignment = sizeof(T) * N>
+class alignas(Alignment) AlignedArray {
+  T data[N];
+};
+
+#define XLLM_SHFL_XOR_SYNC(mask, var, lane_mask) \
+  __shfl_xor_sync((mask), (var), (lane_mask))
+#define XLLM_SHFL_XOR_SYNC_WIDTH(mask, var, lane_mask, width) \
+  __shfl_xor_sync((mask), (var), (lane_mask), (width))
+
+template <typename T>
+__device__ __forceinline__ T xllm_ldg(const T* ptr) {
+#if defined(USE_DCU)
+  return *ptr;
+#else
+  return __ldg(ptr);
+#endif
+}
+
+// Define reduction operators based on CUB version.
+#if defined(USE_DCU)
+using MaxReduceOp = hipcub::Max;
+using MinReduceOp = hipcub::Min;
+#elif CUB_VERSION >= 200800
+using MaxReduceOp = ::cuda::maximum<>;
+using MinReduceOp = ::cuda::minimum<>;
+#else
+using MaxReduceOp = cub::Max;
+using MinReduceOp = cub::Min;
+#endif
+
+template <typename T>
+__device__ float convert_to_float(T x) {
+  if constexpr (std::is_same_v<T, __half>) {
+    return __half2float(x);
+#if defined(USE_DCU)
+  } else if constexpr (std::is_same_v<T, hip_bfloat16>) {
+    return __bfloat162float(reinterpret_cast<const __hip_bfloat16&>(x));
+#else
+  } else if constexpr (std::is_same_v<T, __nv_bfloat16>) {
+    return __bfloat162float(x);
+#endif
+
+  } else if constexpr (std::is_same_v<T, float>) {
+    return x;
+  } else {
+    return static_cast<float>(x);
+  }
+}
+
+// Constructs some constants needed to partition the work across threads at
+// compile time.
+template <typename T, int EXPERTS, int BYTES_PER_LDG>
+struct TopkConstants {
+  static constexpr int ELTS_PER_LDG = BYTES_PER_LDG / sizeof(T);
+  static_assert(EXPERTS / (ELTS_PER_LDG * WARP_SIZE) == 0 ||
+                    EXPERTS % (ELTS_PER_LDG * WARP_SIZE) == 0,
+                "");
+  static constexpr int VECs_PER_THREAD =
+      MAX(1, EXPERTS / (ELTS_PER_LDG * WARP_SIZE));
+  static constexpr int VPT = VECs_PER_THREAD * ELTS_PER_LDG;
+  static constexpr int THREADS_PER_ROW = EXPERTS / VPT;
+  static constexpr int ROWS_PER_WARP = WARP_SIZE / THREADS_PER_ROW;
+};
+
+}  // namespace xllm::kernel::cuda

ex_engine/xllm_kernels/cuda/headers/fp8_quant_utils.cuh

@@ -0,0 +1,239 @@
+/* Copyright 2025-2026 The xLLM Authors. All Rights Reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     https://github.com/jd-opensource/xllm/blob/main/LICENSE
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ * ===========================================================================*/
+
+#pragma once
+// clang-format off
+#include <c10/util/Float8_e4m3fn.h>
+#include <cmath>
+#include <torch/types.h>
+// clang-format on
+namespace xllm {
+namespace kernel {
+namespace cuda {
+
+// FP8 type max value definitions
+template <typename T,
+          typename = std::enable_if_t<std::is_same_v<T, c10::Float8_e4m3fn> ||
+                                      std::is_same_v<T, int8_t>>>
+struct quant_type_max {
+  static constexpr T val() { return std::numeric_limits<T>::max(); }
+};
+
+template <typename T>
+__host__ __device__ static constexpr T quant_type_max_v =
+    quant_type_max<T>::val();
+
+// Minimum scaling factor for quantization types
+template <typename T,
+          typename = std::enable_if_t<std::is_same_v<T, c10::Float8_e4m3fn> ||
+                                      std::is_same_v<T, int8_t>>>
+struct min_scaling_factor {
+  __device__ __host__ static inline float val() {
+    return 1.0f / (quant_type_max_v<T> * 512.0f);
+  }
+};
+
+template <>
+struct min_scaling_factor<int8_t> {
+  __device__ __host__ static inline float val() {
+    return std::numeric_limits<float>::epsilon();
+  }
+};
+
+// Vectorization containers
+template <typename scalar_t, size_t vec_size>
+struct __align__(vec_size * sizeof(scalar_t)) vec_n_t {
+  scalar_t val[vec_size];
+};
+
+template <typename quant_type_t, size_t vec_size>
+struct __align__(vec_size * sizeof(quant_type_t)) q8_n_t {
+  static_assert(std::is_same_v<quant_type_t, int8_t> ||
+                std::is_same_v<quant_type_t, c10::Float8_e4m3fn>);
+  quant_type_t val[vec_size];
+};
+
+// Atomic max for float
+__device__ __forceinline__ float atomicMaxFloat(float* addr, float value) {
+  float old;
+  old = (value >= 0)
+            ? __int_as_float(atomicMax((int*)addr, __float_as_int(value)))
+            : __uint_as_float(
+                  atomicMin((unsigned int*)addr, __float_as_uint(value)));
+  return old;
+}
+
+// FP8 conversion functions
+namespace fp8 {
+
+#ifdef ENABLE_FP8
+
+#include <cuda_fp8.h>
+
+// float -> c10::Float8_e4m3fn conversion
+template <typename Tout, typename Tin>
+__inline__ __device__ Tout
+vec_conversion(const Tin& x,
+               const __nv_fp8_interpretation_t fp8_type = __NV_E4M3) {
+  return x;
+}
+
+template <>
+__inline__ __device__ c10::Float8_e4m3fn
+vec_conversion<c10::Float8_e4m3fn, float>(
+    const float& a,
+    const __nv_fp8_interpretation_t fp8_type) {
+#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ < 800
+  return static_cast<c10::Float8_e4m3fn>(a);
+#else
+  return c10::Float8_e4m3fn(__nv_cvt_float_to_fp8(a, __NV_SATFINITE, fp8_type),
+                            c10::Float8_e4m3fn::from_bits());
+#endif
+}
+
+#endif  // ENABLE_FP8
+
+}  // namespace fp8
+
+// Scaled FP8 conversion with saturation
+template <bool is_scale_inverted, typename fp8_type>
+__device__ __forceinline__ fp8_type scaled_fp8_conversion(float const val,
+                                                          float const scale) {
+  float x = 0.0f;
+  if constexpr (is_scale_inverted) {
+    x = val * scale;
+  } else {
+    x = val / scale;
+  }
+
+  float r =
+      fmaxf(-quant_type_max_v<fp8_type>, fminf(x, quant_type_max_v<fp8_type>));
+
+#ifdef ENABLE_FP8
+  // Use hardware cvt instruction for fp8 on nvidia
+  return fp8::vec_conversion<fp8_type, float>(r);
+#else
+  return static_cast<fp8_type>(r);
+#endif
+}
+
+// Vectorization utilities
+template <int VEC_SIZE, typename InT, typename OutT, typename ScaOp>
+struct DefaultVecOp {
+  ScaOp scalar_op;
+
+  __device__ __forceinline__ void operator()(
+      vec_n_t<OutT, VEC_SIZE>& dst,
+      const vec_n_t<InT, VEC_SIZE>& src) const {
+#pragma unroll
+    for (int i = 0; i < VEC_SIZE; ++i) {
+      scalar_op(dst.val[i], src.val[i]);
+    }
+  }
+};
+
+template <int VEC_SIZE,
+          typename InT,
+          typename OutT,
+          typename VecOp,
+          typename ScaOp>
+__device__ inline void vectorize_with_alignment(
+    const InT* in,
+    OutT* out,
+    int len,
+    int tid,
+    int stride,
+    VecOp&& vec_op,       // vec_n_t<InT,16> -> vec_n_t<OutT,16>
+    ScaOp&& scalar_op) {  // InT -> OutT
+  static_assert(VEC_SIZE > 0 && (VEC_SIZE & (VEC_SIZE - 1)) == 0,
+                "VEC_SIZE must be a positive power-of-two");
+  constexpr int WIDTH = VEC_SIZE * sizeof(InT);
+  uintptr_t addr = reinterpret_cast<uintptr_t>(in);
+
+  // Fast path when the whole region is already aligned
+  bool can_vec = ((addr & (WIDTH - 1)) == 0) && ((len & (VEC_SIZE - 1)) == 0);
+  if (can_vec) {
+    int num_vec = len / VEC_SIZE;
+
+    using vin_t = vec_n_t<InT, VEC_SIZE>;
+    using vout_t = vec_n_t<OutT, VEC_SIZE>;
+    auto* v_in = reinterpret_cast<const vin_t*>(in);
+    auto* v_out = reinterpret_cast<vout_t*>(out);
+
+    for (int i = tid; i < num_vec; i += stride) {
+      vout_t tmp;
+      vin_t src = v_in[i];
+      vec_op(tmp, src);
+      v_out[i] = tmp;
+    }
+    return;
+  }
+
+  int misalignment_offset = addr & (WIDTH - 1);
+  int alignment_bytes = WIDTH - misalignment_offset;
+  int prefix_elems = alignment_bytes & (WIDTH - 1);
+  prefix_elems /= sizeof(InT);
+  prefix_elems = min(prefix_elems, len);
+
+  // Prefix handling
+  for (int i = tid; i < prefix_elems; i += stride) {
+    scalar_op(out[i], in[i]);
+  }
+
+  in += prefix_elems;
+  out += prefix_elems;
+  len -= prefix_elems;
+
+  int num_vec = len / VEC_SIZE;
+  using vin_t = vec_n_t<InT, VEC_SIZE>;
+  using vout_t = vec_n_t<OutT, VEC_SIZE>;
+  auto* v_in = reinterpret_cast<const vin_t*>(in);
+  auto* v_out = reinterpret_cast<vout_t*>(out);
+
+  // Vectorized main part
+  for (int i = tid; i < num_vec; i += stride) {
+    vout_t tmp;
+    vin_t src = v_in[i];
+    vec_op(tmp, src);
+    v_out[i] = tmp;
+  }
+
+  // Tail handling
+  int tail_start = num_vec * VEC_SIZE;
+  for (int i = tid + tail_start; i < len; i += stride) {
+    scalar_op(out[i], in[i]);
+  }
+}
+
+template <int VEC_SIZE, typename InT, typename OutT, typename ScaOp>
+__device__ __forceinline__ void vectorize_with_alignment(const InT* in,
+                                                         OutT* out,
+                                                         int len,
+                                                         int tid,
+                                                         int stride,
+                                                         ScaOp&& scalar_op) {
+  using Vec = DefaultVecOp<VEC_SIZE, InT, OutT, std::decay_t<ScaOp>>;
+  vectorize_with_alignment<VEC_SIZE>(in,
+                                     out,
+                                     len,
+                                     tid,
+                                     stride,
+                                     Vec{scalar_op},
+                                     std::forward<ScaOp>(scalar_op));
+}
+
+}  // namespace cuda
+}  // namespace kernel
+}  // namespace xllm

ex_engine/xllm_kernels/cuda/headers/type_convert.cuh

@@ -0,0 +1,231 @@
+/* Copyright 2025 The vLLM Authors and 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 <cuda_bf16.h>
+#include <cuda_fp16.h>
+#include <torch/all.h>
+
+// ref to:
+// https://github.com/vllm-project/vllm/blob/main/csrc/type_convert.cuh
+
+/* Converter helpers for the conversion from torch types to HIP/CUDA types,
+   and the associated type conversions within HIP/CUDA. These helpers need
+   to be implemented for now because the relevant type conversion
+   operators/constructors are not consistently implemented by HIP/CUDA, so
+   a generic conversion via type casts cannot be implemented.
+
+   Each helper should have the member static constexpr bool `exists`:
+   If false, the optimized kernel is not used for the corresponding torch type.
+   If true, the helper should be fully defined as shown in the examples below.
+ */
+namespace xllm::kernel::cuda {
+template <typename torch_type>
+class _typeConvert {
+ public:
+  static constexpr bool exists = false;
+};
+
+template <>
+class _typeConvert<float> {
+ public:
+  static constexpr bool exists = true;
+  using hip_type = float;
+  using packed_hip_type = float2;
+  using packed_hip_type4 = float4;  // For 128-bit vectorization
+
+  __device__ static __forceinline__ float convert(hip_type x) { return x; }
+  __device__ static __forceinline__ float2 convert(packed_hip_type x) {
+    return x;
+  }
+  __device__ static __forceinline__ float4 convert(packed_hip_type4 x) {
+    return x;
+  }
+};
+
+#if defined(USE_DCU) || (defined(CUDA_VERSION) && (CUDA_VERSION >= 12000)) || \
+    defined(USE_MACA)
+// CUDA < 12.0 runs into issues with packed type conversion
+template <>
+class _typeConvert<c10::Half> {
+ public:
+  static constexpr bool exists = true;
+  using hip_type = __half;
+  using packed_hip_type = __half2;
+
+  __device__ static __forceinline__ float convert(hip_type x) {
+    return __half2float(x);
+  }
+  __device__ static __forceinline__ float2 convert(packed_hip_type x) {
+    return __half22float2(x);
+  }
+  __device__ static __forceinline__ hip_type convert(float x) {
+    return __float2half_rn(x);
+  }
+  __device__ static __forceinline__ packed_hip_type convert(float2 x) {
+    return __float22half2_rn(x);
+  }
+};
+#endif  // defined(USE_DCU) || CUDA_VERSION >= 12000
+
+#if defined(USE_DCU)
+template <>
+class _typeConvert<c10::BFloat16> {
+ public:
+  static constexpr bool exists = true;
+  using hip_type = __hip_bfloat16;
+  using packed_hip_type = __hip_bfloat162;
+
+  __device__ static __forceinline__ float convert(hip_type x) {
+    return __bfloat162float(x);
+  }
+  __device__ static __forceinline__ float2 convert(packed_hip_type x) {
+    return __bfloat1622float2(x);
+  }
+  __device__ static __forceinline__ hip_type convert(float x) {
+    return __float2bfloat16(x);
+  }
+  __device__ static __forceinline__ packed_hip_type convert(float2 x) {
+    return __float22bfloat162_rn(x);
+  }
+};
+#elif defined(CUDA_VERSION) && (CUDA_VERSION >= 12000) &&   \
+        defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 800) || \
+    defined(USE_MACA)
+
+// CUDA_ARCH < 800 does not have BF16 support.
+template <>
+class _typeConvert<c10::BFloat16> {
+ public:
+  static constexpr bool exists = true;
+  using hip_type = __nv_bfloat16;
+  using packed_hip_type = __nv_bfloat162;
+
+  __device__ static __forceinline__ float convert(hip_type x) {
+    return __bfloat162float(x);
+  }
+  __device__ static __forceinline__ float2 convert(packed_hip_type x) {
+    return __bfloat1622float2(x);
+  }
+  __device__ static __forceinline__ hip_type convert(float x) {
+    return __float2bfloat16(x);
+  }
+  __device__ static __forceinline__ packed_hip_type convert(float2 x) {
+    return __float22bfloat162_rn(x);
+  }
+};
+#endif
+
+/* Vector helper to generate vectorized and packed FP16/BF16 ops
+   for appropriate specializations of fused_add_rms_norm_kernel.
+   Only functions that are necessary in that kernel are implemented.
+   Alignment to 16 bytes is required to use 128-bit global memory ops.
+ */
+
+template <typename scalar_t, int width>
+class alignas(16) _f16Vec {
+ public:
+  /* Not theoretically necessary that width is a power of 2 but should
+     almost always be the case for optimization purposes */
+  static_assert(width > 0 && (width & (width - 1)) == 0,
+                "Width is not a positive power of 2!");
+  using Converter = _typeConvert<scalar_t>;
+  using T1 = typename Converter::hip_type;
+  using T2 = typename Converter::packed_hip_type;
+  T1 data[width];
+
+  __device__ _f16Vec& operator+=(const _f16Vec<scalar_t, width>& other) {
+    if constexpr (width % 2 == 0) {
+#pragma unroll
+      for (int i = 0; i < width; i += 2) {
+        if constexpr (std::is_same_v<T2, float2>) {
+          data[i] += other.data[i];
+          data[i + 1] += other.data[i + 1];
+        } else {
+          T2 temp{data[i], data[i + 1]};
+          temp += T2{other.data[i], other.data[i + 1]};
+          data[i] = temp.x;
+          data[i + 1] = temp.y;
+        }
+      }
+    } else {
+#pragma unroll
+      for (int i = 0; i < width; ++i) data[i] += other.data[i];
+    }
+    return *this;
+  }
+
+  __device__ _f16Vec& operator*=(const _f16Vec<scalar_t, width>& other) {
+    if constexpr (width % 2 == 0) {
+#pragma unroll
+      for (int i = 0; i < width; i += 2) {
+        if constexpr (std::is_same_v<T2, float2>) {
+          data[i] *= other.data[i];
+          data[i + 1] *= other.data[i + 1];
+        } else {
+          T2 temp{data[i], data[i + 1]};
+          temp *= T2{other.data[i], other.data[i + 1]};
+          data[i] = temp.x;
+          data[i + 1] = temp.y;
+        }
+      }
+    } else {
+#pragma unroll
+      for (int i = 0; i < width; ++i) data[i] *= other.data[i];
+    }
+    return *this;
+  }
+
+  __device__ _f16Vec& operator*=(const float scale) {
+    if constexpr (width % 2 == 0) {
+#pragma unroll
+      for (int i = 0; i < width; i += 2) {
+        float2 temp_f = Converter::convert(T2{data[i], data[i + 1]});
+        temp_f.x *= scale;
+        temp_f.y *= scale;
+        T2 temp = Converter::convert(temp_f);
+        data[i] = temp.x;
+        data[i + 1] = temp.y;
+      }
+    } else {
+#pragma unroll
+      for (int i = 0; i < width; ++i) {
+        float temp = Converter::convert(data[i]) * scale;
+        data[i] = Converter::convert(temp);
+      }
+    }
+    return *this;
+  }
+
+  __device__ float sum_squares() const {
+    float result = 0.0f;
+    if constexpr (width % 2 == 0) {
+#pragma unroll
+      for (int i = 0; i < width; i += 2) {
+        float2 z = Converter::convert(T2{data[i], data[i + 1]});
+        result += z.x * z.x + z.y * z.y;
+      }
+    } else {
+#pragma unroll
+      for (int i = 0; i < width; ++i) {
+        float x = Converter::convert(data[i]);
+        result += x * x;
+      }
+    }
+    return result;
+  }
+};
+}  // namespace xllm::kernel::cuda

ex_engine/xllm_kernels/cuda/headers/utils.h

@@ -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.
+==============================================================================*/
+
+#pragma once
+
+#include <ATen/DynamicLibrary.h>
+#if defined(USE_DCU)
+#include <c10/hip/HIPGuard.h>
+#else
+#include <c10/cuda/CUDAGuard.h>
+#endif
+#include <glog/logging.h>
+#include <torch/torch.h>
+#if !defined(USE_DCU)
+#include <tvm/ffi/container/array.h>
+#include <tvm/ffi/container/tensor.h>
+#include <tvm/ffi/extra/c_env_api.h>
+#include <tvm/ffi/extra/module.h>
+#include <tvm/ffi/optional.h>
+#endif
+
+#include <string>
+#include <tuple>
+#include <type_traits>
+#include <unordered_map>
+
+#if defined(__CUDACC__) || defined(_NVHPC_CUDA) || defined(__HIPCC__)
+#define HOST_DEVICE_INLINE __host__ __device__ __forceinline__
+#define DEVICE_INLINE __device__ __forceinline__
+#define HOST_INLINE __host__ __forceinline__
+#else
+#define HOST_DEVICE_INLINE inline
+#define DEVICE_INLINE inline
+#define HOST_INLINE inline
+#endif
+
+#if !defined(USE_DCU)
+namespace ffi = tvm::ffi;
+#endif
+
+namespace xllm::kernel::cuda {
+
+template <typename T>
+HOST_DEVICE_INLINE constexpr std::enable_if_t<std::is_integral_v<T>, T>
+ceil_div(T a, T b) {
+  return (a + b - 1) / b;
+}
+
+enum class ActivationType : int8_t {
+  GELU = 0,
+  RELU = 1,
+  SILU = 2,
+  SWIGLU = 3,
+  GEGLU = 4,
+  SWIGLU_BIAS = 5,
+  RELU2 = 6,
+  IDENTITY = 7,
+  INVALID_TYPE = 8
+};
+
+// torch tensor is only on cpu
+torch::Tensor get_cache_buffer(const int32_t seq_len,
+                               const torch::Device& device);
+
+// NOLINTBEGIN(cppcoreguidelines-macro-usage)
+#define DISPATCH_CASE_FLOATING_TYPES(...)              \
+  AT_DISPATCH_CASE(at::ScalarType::Float, __VA_ARGS__) \
+  AT_DISPATCH_CASE(at::ScalarType::Half, __VA_ARGS__)  \
+  AT_DISPATCH_CASE(at::ScalarType::BFloat16, __VA_ARGS__)
+#define DISPATCH_FLOATING_TYPES(TYPE, NAME, ...) \
+  AT_DISPATCH_SWITCH(TYPE, NAME, DISPATCH_CASE_FLOATING_TYPES(__VA_ARGS__))
+#define DISPATCH_CASE_HALF_TYPES(...)                 \
+  AT_DISPATCH_CASE(at::ScalarType::Half, __VA_ARGS__) \
+  AT_DISPATCH_CASE(at::ScalarType::BFloat16, __VA_ARGS__)
+#define DISPATCH_HALF_TYPES(TYPE, NAME, ...) \
+  AT_DISPATCH_SWITCH(TYPE, NAME, DISPATCH_CASE_HALF_TYPES(__VA_ARGS__))
+// NOLINTEND(cppcoreguidelines-macro-usage)
+
+bool should_use_tensor_core(torch::ScalarType kv_cache_dtype,
+                            int64_t num_attention_heads,
+                            int64_t num_kv_heads);
+
+bool support_pdl();
+
+std::string path_to_uri_so_lib(const std::string& uri);
+
+std::string determine_attention_backend(int64_t pos_encoding_mode,
+                                        bool use_fp16_qk_reduction,
+                                        bool use_custom_mask);
+
+std::string get_batch_prefill_uri(const std::string& backend,
+                                  torch::ScalarType dtype_q,
+                                  torch::ScalarType dtype_kv,
+                                  torch::ScalarType dtype_o,
+                                  torch::ScalarType dtype_idx,
+                                  int64_t head_dim_qk,
+                                  int64_t head_dim_vo,
+                                  int64_t pos_encoding_mode,
+                                  bool use_sliding_window,
+                                  bool use_logits_soft_cap,
+                                  bool use_fp16_qk_reduction);
+
+std::string get_batch_decode_uri(torch::ScalarType dtype_q,
+                                 torch::ScalarType dtype_kv,
+                                 torch::ScalarType dtype_o,
+                                 torch::ScalarType dtype_idx,
+                                 int64_t head_dim_qk,
+                                 int64_t head_dim_vo,
+                                 int64_t pos_encoding_mode,
+                                 bool use_sliding_window,
+                                 bool use_logits_soft_cap);
+
+std::tuple<torch::Tensor, double> split_scale_param(const torch::Tensor& scale);
+
+#if !defined(USE_DCU)
+DLDataType to_dl_data_type(torch::ScalarType scalar_type);
+
+// below are tvm-ffi related functions
+ffi::Tensor to_ffi_tensor(const torch::Tensor& torch_tensor);
+
+ffi::Optional<ffi::Tensor> to_ffi_optional_tensor(
+    const std::optional<torch::Tensor>& optional);
+
+ffi::Array<ffi::Tensor> to_ffi_array_tensors(
+    const std::vector<torch::Tensor>& torch_tensors);
+
+ffi::Optional<ffi::Array<ffi::Tensor>> to_ffi_optional_array_tensors(
+    const std::optional<std::vector<torch::Tensor>>& optional);
+
+ffi::Module get_module(const std::string& uri);
+
+ffi::Function get_function(const std::string& uri,
+                           const std::string& func_name);
+
+inline void bind_tvmffi_stream_to_current_torch_stream(
+    const torch::Device& device) {
+  const auto cur = c10::cuda::getCurrentCUDAStream(device.index());
+  // DLPack device type for CUDA is 2 (kDLCUDA).
+  void* original_stream = nullptr;
+  const int rc = TVMFFIEnvSetStream(
+      /*device_type=*/2,
+      /*device_id=*/device.index(),
+      reinterpret_cast<void*>(cur.stream()),
+      &original_stream);
+  if (rc != 0) {
+    LOG(WARNING) << "[tvmffi.stream] failed to set stream, rc=" << rc
+                 << " dev=" << device.index();
+  }
+}
+#endif  // !defined(USE_DCU)
+}  // namespace xllm::kernel::cuda