feat: optimize qwen3.5 causal_conv1d and chunk_gated_delta_rule on dcu.

xllm/core/kernels/dcu/CMakeLists.txt

@@ -59,6 +59,11 @@ set(DCU_LOCAL_HIP_SOURCES
   ${CMAKE_SOURCE_DIR}/xllm/core/kernels/dcu/topk_gate.cpp
   ${CMAKE_SOURCE_DIR}/xllm/core/kernels/dcu/gated_layer_norm.cpp
   ${CMAKE_SOURCE_DIR}/xllm/core/kernels/dcu/gemma_rms_norm.cpp
+
+  ${CMAKE_SOURCE_DIR}/xllm/core/kernels/dcu/causal_conv1d_adapter.cpp
+  ${CMAKE_SOURCE_DIR}/xllm/core/kernels/dcu/causal_conv1d/causal_conv1d_fwd.hip
+  ${CMAKE_SOURCE_DIR}/xllm/core/kernels/dcu/causal_conv1d/causal_conv1d_update.hip
+  ${CMAKE_SOURCE_DIR}/xllm/core/kernels/dcu/causal_conv1d/causal_conv1d_new_interface.hip
 )
 
 hipify_sources_target(
@@ -135,9 +140,9 @@ target_include_directories(dcu_kernels
   PUBLIC
     ${CMAKE_CURRENT_LIST_DIR}
     ${CMAKE_CURRENT_BINARY_DIR}
+    ${AITER_CPP_API_LIB_DIR}/../../csrc/include
   PRIVATE
     ${CMAKE_CURRENT_BINARY_DIR}/..
-    ${AITER_CPP_API_LIB_DIR}/../../aiter_meta/csrc/include
 )
 
 target_compile_definitions(dcu_kernels PRIVATE

xllm/core/kernels/dcu/causal_conv1d/causal_conv1d.h

@@ -0,0 +1,88 @@
+/******************************************************************************
+ * Copyright (c) 2024, Tri Dao.
+ ******************************************************************************/
+
+#pragma once
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+struct ConvParamsBase {
+  using index_t = uint32_t;
+
+  int batch, dim, seqlen, width;
+  int max_seqlen;
+  bool silu_activation;
+
+  index_t x_batch_stride;
+  index_t x_c_stride;
+  index_t x_l_stride;
+  index_t weight_c_stride;
+  index_t weight_width_stride;
+  index_t out_batch_stride;
+  index_t out_c_stride;
+  index_t out_l_stride;
+
+  int conv_state_len;
+  index_t conv_state_batch_stride;
+  index_t conv_state_c_stride;
+  index_t conv_state_l_stride;
+
+  // Common data pointers.
+  void* __restrict__ x_ptr;
+  void* __restrict__ weight_ptr;
+  void* __restrict__ bias_ptr;
+  void* __restrict__ out_ptr;
+
+  int32_t* __restrict__ query_start_loc_ptr;
+  bool* __restrict__ has_initial_state_ptr;
+
+  void* __restrict__ conv_state_ptr;
+  int32_t* __restrict__ cache_seqlens;
+
+  // Only used if the elements of the batch are gathered from a larger buffer,
+  // which may happen for continuous batching.
+  int32_t* __restrict__ conv_state_indices_ptr;
+  index_t conv_state_indices_stride;
+
+  void* __restrict__ seq_idx_ptr;
+
+  // No __restrict__ since initial_states could be the same as final_states.
+  void* initial_states_ptr;
+  index_t initial_states_batch_stride;
+  index_t initial_states_l_stride;
+  index_t initial_states_c_stride;
+
+  void* final_states_ptr;
+  index_t final_states_batch_stride;
+  index_t final_states_l_stride;
+  index_t final_states_c_stride;
+
+  int pad_slot_id;
+};
+
+struct ConvParamsBwd : public ConvParamsBase {
+  index_t dx_batch_stride;
+  index_t dx_c_stride;
+  index_t dx_l_stride;
+  index_t dweight_c_stride;
+  index_t dweight_width_stride;
+  index_t dout_batch_stride;
+  index_t dout_c_stride;
+  index_t dout_l_stride;
+
+  // Common data pointers.
+  void* __restrict__ dx_ptr;
+  void* __restrict__ dweight_ptr;
+  void* __restrict__ dbias_ptr;
+  void* __restrict__ dout_ptr;
+
+  void* dinitial_states_ptr;
+  index_t dinitial_states_batch_stride;
+  index_t dinitial_states_l_stride;
+  index_t dinitial_states_c_stride;
+
+  void* dfinal_states_ptr;
+  index_t dfinal_states_batch_stride;
+  index_t dfinal_states_l_stride;
+  index_t dfinal_states_c_stride;
+};

xllm/core/kernels/dcu/causal_conv1d/causal_conv1d_common_hip.h

@@ -0,0 +1,104 @@
+// !!! This is a file automatically generated by hipify!!!
+#include <ATen/dtk_macros.h>
+/******************************************************************************
+ * Copyright (c) 2023, Tri Dao.
+ ******************************************************************************/
+
+#pragma once
+
+#ifndef USE_ROCM
+#include <hip/hip_bf16.h>
+
+template <typename T>
+__device__ inline T shuffle_xor(T val, int offset) {
+  return __shfl_xor_sync(uint32_t(-1), val, offset);
+}
+
+constexpr size_t custom_max(std::initializer_list<size_t> ilist) {
+  return std::max(ilist);
+}
+
+template <typename T>
+constexpr T constexpr_min(T a, T b) {
+  return std::min(a, b);
+}
+
+#else
+#include <hip/hip_bf16.h>
+
+template <typename T>
+__device__ inline T shuffle_xor(T val, int offset) {
+  return __shfl_xor(val, offset);
+}
+constexpr size_t custom_max(std::initializer_list<size_t> ilist) {
+  return *std::max_element(ilist.begin(), ilist.end());
+}
+
+template <typename T>
+constexpr T constexpr_min(T a, T b) {
+  return a < b ? a : b;
+}
+#endif
+#include <hip/hip_fp16.h>
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template <int BYTES>
+struct BytesToType {};
+
+template <>
+struct BytesToType<16> {
+  using Type = uint4;
+  static_assert(sizeof(Type) == 16);
+};
+
+template <>
+struct BytesToType<8> {
+  using Type = uint64_t;
+  static_assert(sizeof(Type) == 8);
+};
+
+template <>
+struct BytesToType<4> {
+  using Type = uint32_t;
+  static_assert(sizeof(Type) == 4);
+};
+
+template <>
+struct BytesToType<2> {
+  using Type = uint16_t;
+  static_assert(sizeof(Type) == 2);
+};
+
+template <>
+struct BytesToType<1> {
+  using Type = uint8_t;
+  static_assert(sizeof(Type) == 1);
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template <typename T>
+struct SumOp {
+  __device__ inline T operator()(T const& x, T const& y) { return x + y; }
+};
+
+template <int THREADS>
+struct Allreduce {
+  static_assert(THREADS == 32 || THREADS == 16 || THREADS == 8 || THREADS == 4);
+  template <typename T, typename Operator>
+  static __device__ inline T run(T x, Operator& op) {
+    constexpr int OFFSET = THREADS / 2;
+    x = op(x, shuffle_xor(x, OFFSET));
+    return Allreduce<OFFSET>::run(x, op);
+  }
+};
+
+template <>
+struct Allreduce<2> {
+  template <typename T, typename Operator>
+  static __device__ inline T run(T x, Operator& op) {
+    x = op(x, shuffle_xor(x, 1));
+    return x;
+  }
+};