From 0243057d2892ff5ea8732239110db2dc9f9d0565 Mon Sep 17 00:00:00 2001 From: "feng.ding" Date: Sat, 18 Jul 2026 04:14:15 +0000 Subject: [PATCH] feat: optimize qwen3.5 causal_conv1d and chunk_gated_delta_rule on dcu. --- xllm/core/kernels/dcu/CMakeLists.txt | 7 +- .../kernels/dcu/causal_conv1d/causal_conv1d.h | 88 +++ .../causal_conv1d/causal_conv1d_common_hip.h | 104 +++ .../dcu/causal_conv1d/causal_conv1d_fwd.hip | 610 +++++++++++++++++ .../causal_conv1d_new_interface.hip | 294 ++++++++ .../causal_conv1d/causal_conv1d_update.hip | 149 +++++ .../kernels/dcu/causal_conv1d/static_switch.h | 26 + .../kernels/dcu/causal_conv1d_adapter.cpp | 253 +++++++ xllm/core/kernels/dcu/causal_conv1d_adapter.h | 58 ++ xllm/core/layers/dcu/CMakeLists.txt | 4 + .../layers/dcu/qwen3_gated_delta_net_base.cpp | 626 +++++++----------- .../dcu/qwen3_gated_delta_net_helpers.h | 59 ++ .../layers/dcu/qwen3_gated_delta_net_impls.h | 146 ++++ .../dcu/qwen3_gated_delta_net_kernel_impl.cpp | 416 ++++++++++++ .../dcu/qwen3_gated_delta_net_torch_impl.cpp | 184 +++++ 15 files changed, 2637 insertions(+), 387 deletions(-) create mode 100644 xllm/core/kernels/dcu/causal_conv1d/causal_conv1d.h create mode 100644 xllm/core/kernels/dcu/causal_conv1d/causal_conv1d_common_hip.h create mode 100644 xllm/core/kernels/dcu/causal_conv1d/causal_conv1d_fwd.hip create mode 100644 xllm/core/kernels/dcu/causal_conv1d/causal_conv1d_new_interface.hip create mode 100644 xllm/core/kernels/dcu/causal_conv1d/causal_conv1d_update.hip create mode 100644 xllm/core/kernels/dcu/causal_conv1d/static_switch.h create mode 100644 xllm/core/kernels/dcu/causal_conv1d_adapter.cpp create mode 100644 xllm/core/kernels/dcu/causal_conv1d_adapter.h create mode 100644 xllm/core/layers/dcu/qwen3_gated_delta_net_helpers.h create mode 100644 xllm/core/layers/dcu/qwen3_gated_delta_net_impls.h create mode 100644 xllm/core/layers/dcu/qwen3_gated_delta_net_kernel_impl.cpp create mode 100644 xllm/core/layers/dcu/qwen3_gated_delta_net_torch_impl.cpp diff --git a/xllm/core/kernels/dcu/CMakeLists.txt b/xllm/core/kernels/dcu/CMakeLists.txt index 347e9fb9dd..c81297a23f 100755 --- a/xllm/core/kernels/dcu/CMakeLists.txt +++ b/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 diff --git a/xllm/core/kernels/dcu/causal_conv1d/causal_conv1d.h b/xllm/core/kernels/dcu/causal_conv1d/causal_conv1d.h new file mode 100644 index 0000000000..982475636b --- /dev/null +++ b/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; +}; diff --git a/xllm/core/kernels/dcu/causal_conv1d/causal_conv1d_common_hip.h b/xllm/core/kernels/dcu/causal_conv1d/causal_conv1d_common_hip.h new file mode 100644 index 0000000000..e87903c5e5 --- /dev/null +++ b/xllm/core/kernels/dcu/causal_conv1d/causal_conv1d_common_hip.h @@ -0,0 +1,104 @@ +// !!! This is a file automatically generated by hipify!!! +#include +/****************************************************************************** + * Copyright (c) 2023, Tri Dao. + ******************************************************************************/ + +#pragma once + +#ifndef USE_ROCM +#include + +template +__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 ilist) { + return std::max(ilist); +} + +template +constexpr T constexpr_min(T a, T b) { + return std::min(a, b); +} + +#else +#include + +template +__device__ inline T shuffle_xor(T val, int offset) { + return __shfl_xor(val, offset); +} +constexpr size_t custom_max(std::initializer_list ilist) { + return *std::max_element(ilist.begin(), ilist.end()); +} + +template +constexpr T constexpr_min(T a, T b) { + return a < b ? a : b; +} +#endif +#include + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template +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 +struct SumOp { + __device__ inline T operator()(T const& x, T const& y) { return x + y; } +}; + +template +struct Allreduce { + static_assert(THREADS == 32 || THREADS == 16 || THREADS == 8 || THREADS == 4); + template + static __device__ inline T run(T x, Operator& op) { + constexpr int OFFSET = THREADS / 2; + x = op(x, shuffle_xor(x, OFFSET)); + return Allreduce::run(x, op); + } +}; + +template <> +struct Allreduce<2> { + template + static __device__ inline T run(T x, Operator& op) { + x = op(x, shuffle_xor(x, 1)); + return x; + } +}; diff --git a/xllm/core/kernels/dcu/causal_conv1d/causal_conv1d_fwd.hip b/xllm/core/kernels/dcu/causal_conv1d/causal_conv1d_fwd.hip new file mode 100644 index 0000000000..1eb77e49f2 --- /dev/null +++ b/xllm/core/kernels/dcu/causal_conv1d/causal_conv1d_fwd.hip @@ -0,0 +1,610 @@ +// !!! This is a file automatically generated by hipify!!! +#include +#include "hip/hip_runtime.h" +/****************************************************************************** + * Copyright (c) 2024, Tri Dao. + ******************************************************************************/ + +#include +#include +#include // For C10_HIP_CHECK and C10_HIP_KERNEL_LAUNCH_CHECK + +#ifndef USE_ROCM + #include + #include +#else + #include + namespace cub = hipcub; +#endif + +#include "causal_conv1d.h" +#include "causal_conv1d_common_hip.h" +#include "static_switch.h" + +//typedef float v2f __attribute__((ext_vector_type(2))); + +template +__forceinline__ __device__ void buffer_load_lds_dwordx1_wait() { + asm volatile( + "s_waitcnt vmcnt(%0)\n\t" + "s_barrier\n" + :: "B"(COUNT) + :); +} + +using vec4_uint = __attribute__((__vector_size__(4 * sizeof(uint32_t)))) uint32_t; + +//template +template +__device__ __forceinline__ vec4_uint prepare_for_buffer_load(T* ptr) { + vec4_uint res; + *(uint64_t*)&res = reinterpret_cast(ptr); + //if constexpr (Do_CacheSwizzle) { + // if constexpr (kHeadDim == 128) { + // res[1] += 0x41000000; // 62 bit: cache swizzle; 48~61: Stride + // } else if constexpr (kHeadDim == 192) { + // res[1] += 0x41800000; // stride 192Bytes and change tagram + // } else if constexpr (kHeadDim == 64) { + // res[1] += 0x40800000; // stride 128Bytes and change tagram + // } + //} + res[2] = 0x80000000; + res[3] = 0x00020000; + return res; +} + +template +__forceinline__ __device__ void inline_buffer_load_dword_lds(DataType *const shared_addr, const vec4_uint global_addr, const int &lds_offset, const int &gvOffset_s, const int &gvOffset_v) { + + int ldsAddrPerWave = reinterpret_cast(shared_addr) + (lds_offset << shfl_count); + int offset_s = gvOffset_s << shfl_count; + int offset_v = gvOffset_v << shfl_count; + + asm volatile("s_mov_b32 m0, %1 \n\t" + "buffer_load_dword %0, %2, %3 ,offen offset:0, lds \n" + :: "v"(offset_v), "s"(ldsAddrPerWave), "s"(global_addr), "s"(offset_s) + :); +} + +template +__forceinline__ __device__ void inline_buffer_load_dwordx2_lds(DataType *const shared_addr, const vec4_uint global_addr, const int &lds_offset, const int &gvOffset_s, const int &gvOffset_v) { + + int ldsAddrPerWave = reinterpret_cast(shared_addr) + (lds_offset << shfl_count); + int offset_s = gvOffset_s << shfl_count; + int offset_v = gvOffset_v << shfl_count; + + asm volatile("s_mov_b32 m0, %1 \n\t" + "buffer_load_dwordx2 %0, %2, %3 ,offen offset:0, lds \n" + :: "v"(offset_v), "s"(ldsAddrPerWave), "s"(global_addr), "s"(offset_s) + :); +} + +template +__forceinline__ __device__ void inline_buffer_load_dwordx4_lds(DataType *const shared_addr, const vec4_uint global_addr, const int &lds_offset, const int &gvOffset_s, const int &gvOffset_v) { + + int ldsAddrPerWave = reinterpret_cast(shared_addr) + (lds_offset << shfl_count); + int offset_s = gvOffset_s << shfl_count; + int offset_v = gvOffset_v << shfl_count; + + asm volatile("s_mov_b32 m0, %1 \n\t" + "buffer_load_dwordx4 %0, %2, %3 ,offen offset:0, lds \n" + :: "v"(offset_v), "s"(ldsAddrPerWave), "s"(global_addr), "s"(offset_s) + :); +} +template +struct Causal_conv1d_fwd_kernel_traits { + using input_t = input_t_; + using weight_t = weight_t_; + static constexpr int kNThreads = kNThreads_; + static constexpr int kWidth = kWidth_; + static constexpr int kNBytes = sizeof(input_t); + static_assert(kNBytes == 2 || kNBytes == 4); + static constexpr int kNElts = kNBytes == 4 ? 4 : 8; + static_assert(kWidth <= kNElts); + static constexpr bool kIsVecLoad = kIsVecLoad_; + using vec_t = typename BytesToType::Type; + using BlockLoadT = hipcub::BlockLoad; + using BlockLoadVecT = hipcub::BlockLoad; + using BlockStoreT = hipcub::BlockStore; + using BlockStoreVecT = hipcub::BlockStore; + static constexpr int kSmemIOSize = kIsVecLoad + ? 0 + : custom_max({sizeof(typename BlockLoadT::TempStorage), sizeof(typename BlockStoreT::TempStorage)}); + static constexpr int kSmemExchangeSize = kNThreads * kNBytes * kNElts; + static constexpr int kSmemSize = kSmemIOSize + kSmemExchangeSize * 3; +}; + +template +__global__ __launch_bounds__(Ktraits::kNThreads) +void causal_conv1d_fwd_kernel(ConvParamsBase params) { + constexpr int kWidth = Ktraits::kWidth; + constexpr int kNThreads = Ktraits::kNThreads; + constexpr int kNElts = Ktraits::kNElts; + static constexpr bool kIsVecLoad = Ktraits::kIsVecLoad; + using input_t = typename Ktraits::input_t; + using vec_t = typename Ktraits::vec_t; + using weight_t = typename Ktraits::weight_t; + + // Shared memory. + extern __shared__ char smem_[]; + auto& smem_load = reinterpret_cast(smem_); + auto& smem_load_vec = reinterpret_cast(smem_); + auto& smem_store = reinterpret_cast(smem_); + auto& smem_store_vec = reinterpret_cast(smem_); + vec_t *smem_exchange1 = reinterpret_cast(smem_ + Ktraits::kSmemIOSize); + vec_t *smem_exchange2 = reinterpret_cast(smem_ + Ktraits::kSmemIOSize + Ktraits::kSmemExchangeSize); + vec_t *smem_exchange3 = reinterpret_cast(smem_ + Ktraits::kSmemIOSize + Ktraits::kSmemExchangeSize + Ktraits::kSmemExchangeSize); + input_t *smem1 = reinterpret_cast(smem_ + Ktraits::kSmemIOSize); + input_t *smem2 = reinterpret_cast(smem_ + Ktraits::kSmemIOSize + Ktraits::kSmemExchangeSize); + input_t *smem3 = reinterpret_cast(smem_ + Ktraits::kSmemIOSize + Ktraits::kSmemExchangeSize + Ktraits::kSmemExchangeSize); + + const int tidx = threadIdx.x; + const int wid = __builtin_amdgcn_readfirstlane(tidx / 64); + const int tidOfwave = tidx % 64; + const int batch_id = blockIdx.x; + const int channel_id = blockIdx.y; + const int Type_offset = 4/sizeof(input_t); + const int kNElts_offset = (kNElts*sizeof(input_t))/4; + input_t *x = reinterpret_cast(params.x_ptr) + batch_id * params.x_batch_stride + + channel_id * params.x_c_stride; + weight_t *weight = reinterpret_cast(params.weight_ptr) + channel_id * params.weight_c_stride; + input_t *out = reinterpret_cast(params.out_ptr) + batch_id * params.out_batch_stride + + channel_id * params.out_c_stride; + float bias_val = params.bias_ptr == nullptr ? 0.f : float(reinterpret_cast(params.bias_ptr)[channel_id]); + + vec4_uint global_addr = prepare_for_buffer_load(x); + // Thread 0 will load the last elements of the previous chunk, so we initialize those to 0. + input_t zeros[kNElts] = {0}; + smem_exchange1[tidx] = reinterpret_cast(zeros)[0]; + #if defined(__gfx926__) || defined(__gfx928__) + inline_buffer_load_dword_lds(smem2, global_addr, wid * 64 , wid * 64 , tidOfwave); + inline_buffer_load_dword_lds(smem2, global_addr, wid * 64 + kNThreads , wid * 64 + kNThreads , tidOfwave); + inline_buffer_load_dword_lds(smem2, global_addr, wid * 64 + kNThreads * 2, wid * 64 + kNThreads * 2, tidOfwave); + inline_buffer_load_dword_lds(smem2, global_addr, wid * 64 + kNThreads * 3, wid * 64 + kNThreads * 3, tidOfwave); + #else + inline_buffer_load_dwordx4_lds(smem2,global_addr,wid* 64*kNElts_offset,wid*64*kNElts_offset,tidOfwave*kNElts_offset); + #endif + __builtin_amdgcn_sched_barrier(0); + buffer_load_lds_dwordx1_wait<0>(); + + float weight_vals[kWidth]; + #pragma unroll + for (int i = 0; i < kWidth; ++i) { weight_vals[i] = float(weight[i * params.weight_width_stride]); } + + constexpr int kChunkSize = kNThreads * kNElts; + const int n_chunks = (params.seqlen + kChunkSize - 1) / kChunkSize; + for (int chunk = 0; chunk < n_chunks; ++chunk) { + input_t x_vals_load[2 * kNElts] = {0}; + //if constexpr(kIsVecLoad) { + #if defined(__gfx926__) || defined(__gfx928__) + inline_buffer_load_dword_lds(smem3, global_addr, wid * 64 , (chunk + 1) * kChunkSize / Type_offset + wid * 64 , tidOfwave); + inline_buffer_load_dword_lds(smem3, global_addr, wid * 64 + kNThreads , (chunk + 1) * kChunkSize / Type_offset + wid * 64 + kNThreads , tidOfwave); + inline_buffer_load_dword_lds(smem3, global_addr, wid * 64 + kNThreads * 2, (chunk + 1) * kChunkSize / Type_offset + wid * 64 + kNThreads * 2, tidOfwave); + inline_buffer_load_dword_lds(smem3, global_addr, wid * 64 + kNThreads * 3, (chunk + 1) * kChunkSize / Type_offset + wid * 64 + kNThreads * 3, tidOfwave); + #else + inline_buffer_load_dwordx4_lds(smem3,global_addr,wid* 64*kNElts_offset,(chunk+1)*kChunkSize/Type_offset+wid*64*kNElts_offset,tidOfwave*kNElts_offset); + #endif + //}else{ + // + // + // + //} + reinterpret_cast(x_vals_load)[1] = smem_exchange2[tidx]; + reinterpret_cast(x_vals_load)[0] = smem_exchange2[tidx - 1]; + if(tidx == 0) + { + reinterpret_cast(x_vals_load)[0] = smem_exchange1[kNThreads-1]; + } + /* + float out_vals[kNElts]; + v2f weight_v = {weight_vals[0],weight_vals[1]}; + v2f* out_vals_v = (v2f*)out_vals; + #pragma unroll + for (int i = 0; i < kNElts/2; i++) { + out_vals_v[i] = {bias_val,bias_val}; + v2f x_vals_v = {x_vals_load[kNElts-1+i*2],x_vals_load[kNElts+i*2]}; + out_vals_v[i] += weight_v * x_vals_v; + } + */ + float out_vals[kNElts]; + #pragma unroll + for (int i = 0; i < kNElts; ++i) { + out_vals[i] = bias_val; + #pragma unroll + for (int w = 0; w < kWidth; ++w) { + out_vals[i] += weight_vals[w] * x_vals_load[kNElts + i - (kWidth - w - 1)]; + } + //out_vals[i] += weight_vals[0] * x_vals_load[kNElts + i - 1]; + //out_vals[i] += weight_vals[1] * x_vals_load[kNElts + i - 0]; + } + if (params.silu_activation) { + #pragma unroll + for (int i = 0; i < kNElts; ++i) { + out_vals[i] = out_vals[i] / (1 + expf(-out_vals[i])); + //out_vals[i] = out_vals[i]*0.5*(1.0f+tanhf(0.5f*out_vals[i])); + } + } + + input_t out_vals_store[kNElts]; + #pragma unroll + for (int i = 0; i < kNElts; ++i) { out_vals_store[i] = out_vals[i]; } + + __builtin_amdgcn_sched_barrier(0); + buffer_load_lds_dwordx1_wait<0>(); + + + if constexpr(kIsVecLoad) { + typename Ktraits::BlockStoreVecT(smem_store_vec).Store(reinterpret_cast(out), reinterpret_cast(out_vals_store), (params.seqlen - chunk * kChunkSize) / kNElts); + } else { + typename Ktraits::BlockStoreT(smem_store).Store(out, out_vals_store, params.seqlen - chunk * kChunkSize); + } + out += kChunkSize; + vec_t* temp = smem_exchange1; + smem_exchange1 = smem_exchange2; + smem_exchange2 = smem_exchange3; + smem_exchange3 = temp; + input_t* temp_mem = smem1; + smem1 = smem2; + smem2 = smem3; + smem3 = temp_mem; + } +} + + +template +void causal_conv1d_fwd_launch(ConvParamsBase ¶ms, hipStream_t stream) { + static constexpr int kNElts = sizeof(input_t) == 4 ? 4 : 8; + BOOL_SWITCH(params.seqlen % kNElts == 0, kIsVecLoad, [&] { + using Ktraits = Causal_conv1d_fwd_kernel_traits; + constexpr int kSmemSize = Ktraits::kSmemSize; + dim3 grid(params.batch, params.dim); + + auto kernel = &causal_conv1d_fwd_kernel; + + if (kSmemSize >= 48 * 1024) { + #ifndef USE_ROCM + C10_HIP_CHECK(hipFuncSetAttribute( + kernel, hipFuncAttributeMaxDynamicSharedMemorySize, kSmemSize)); + #else + // There is a slight signature discrepancy in HIP and CUDA "FuncSetAttribute" function. + C10_HIP_CHECK(hipFuncSetAttribute( + (void *) kernel, hipFuncAttributeMaxDynamicSharedMemorySize, kSmemSize)); + std::cerr << "Warning (causal_conv1d fwd launch): attempting to set maxDynamicSharedMemorySize on an AMD GPU which is currently a non-op (in ROCm versions <= 6.1). This might lead to undefined behavior. \n" << std::endl; + #endif + } + hipLaunchKernelGGL(( kernel), dim3(grid), dim3(Ktraits::kNThreads), kSmemSize, stream, params); + /* + hipDeviceSynchronize(); + + hipEvent_t start, stop; + hipEventCreate(&start); + hipEventCreate(&stop); + + // 记录开始时间 + hipEventRecord(start, 0); + + for(int i=0; i<100;i++) + { + hipLaunchKernelGGL(( kernel), dim3(grid), dim3(Ktraits::kNThreads), kSmemSize, stream, params); + } + // 记录结束时间 + hipEventRecord(stop, 0); + hipEventSynchronize(stop); // 等待事件完成 + + // 计算运行时间（毫秒） + float run_time = 0; + hipEventElapsedTime(&run_time, start, stop); + + std::cout<<"------------------------the BMZ the seqlen is:"< +void causal_conv1d_fwd_cuda(ConvParamsBase ¶ms, hipStream_t stream) { + if (params.width == 2) { + causal_conv1d_fwd_launch<256, 2, input_t, weight_t>(params, stream); + } else if (params.width == 3) { + causal_conv1d_fwd_launch<256, 3, input_t, weight_t>(params, stream); + } else if (params.width == 4) { + causal_conv1d_fwd_launch<256, 4, input_t, weight_t>(params, stream); + } +} + +template +struct Causal_conv1d_channellast_fwd_kernel_traits { + // The cache line is 128 bytes, and we try to read 16 bytes per thread. + // So we have 8 threads per "row", so 32 or 64 elements in the channel dimension. + // That leaves 4 columns per warp, and so 16 columns per block (assuming each block has 128 + // threads). Each each load is 16 x 32|64 elements in the L x C dimensions. + using input_t = input_t_; + using weight_t = weight_t_; + static constexpr int kNThreads = kNThreads_; + static_assert(kNThreads % 64 == 0); + static constexpr int kNWarps = kNThreads / 64; + static constexpr int kWidth = kWidth_; + static constexpr int kChunkSizeL = kChunkSizeL_; + static constexpr int kNBytes = sizeof(input_t); + static_assert(kNBytes == 2 || kNBytes == 4); + static constexpr int kNElts = kNBytes == 4 ? 2 : 4; + static constexpr int kNEltsPerRow = 64 / kNBytes; + static constexpr int kNThreadsPerRow = kNEltsPerRow / kNElts; // Always 8 for now + static_assert(kNThreadsPerRow * kNBytes * kNElts == 64); + static constexpr int kNColsPerWarp = 64 / kNThreadsPerRow; // Always 4 for now + static_assert(kNColsPerWarp * kNThreadsPerRow == 64); + static constexpr int kNColsPerLoad = kNColsPerWarp * kNWarps; + static constexpr int kNLoads = kChunkSizeL / kNColsPerLoad; + static_assert(kNLoads * kNColsPerLoad == kChunkSizeL); + static constexpr bool kIsVecLoad = kIsVecLoad_; + using vec_t = typename BytesToType::Type; + /* + using input_t = input_t_; + using weight_t = weight_t_; + static constexpr int kNThreads = kNThreads_; + static_assert(kNThreads % 32 == 0); + static constexpr int kNWarps = kNThreads / 32; + static constexpr int kWidth = kWidth_; + static constexpr int kChunkSizeL = kChunkSizeL_; + static constexpr int kNBytes = sizeof(input_t); + static_assert(kNBytes == 2 || kNBytes == 4); + static constexpr int kNElts = kNBytes == 4 ? 4 : 8; + static constexpr int kNEltsPerRow = 128 / kNBytes; + static constexpr int kNThreadsPerRow = kNEltsPerRow / kNElts; // Always 8 for now + static_assert(kNThreadsPerRow * kNBytes * kNElts == 128); + static constexpr int kNColsPerWarp = 32 / kNThreadsPerRow; // Always 4 for now + static_assert(kNColsPerWarp * kNThreadsPerRow == 32); + static constexpr int kNColsPerLoad = kNColsPerWarp * kNWarps; + static constexpr int kNLoads = kChunkSizeL / kNColsPerLoad; + static_assert(kNLoads * kNColsPerLoad == kChunkSizeL); + static constexpr bool kIsVecLoad = kIsVecLoad_; + using vec_t = typename BytesToType::Type; + */ + // using BlockLoadT = hipcub::BlockLoad; + // using BlockStoreT = hipcub::BlockStore; + // static constexpr int kSmemSize = ::max({sizeof(typename BlockLoadT::TempStorage), + // sizeof(typename BlockStoreT::TempStorage)}); + // static constexpr int kSmemSize = kChunkSizeL * kNEltsPerRow * kNBytes; +}; + +template +__global__ __launch_bounds__(Ktraits::kNThreads) +void causal_conv1d_channellast_fwd_kernel(ConvParamsBase params) { + constexpr int kWidth = Ktraits::kWidth; + constexpr int kNThreads = Ktraits::kNThreads; + constexpr int kNElts = Ktraits::kNElts; + constexpr int kNWarp = Ktraits::kNWarps; + constexpr int kNThreadsPerC = Ktraits::kNThreadsPerRow; + constexpr int kLPerLoad = Ktraits::kNColsPerLoad; + constexpr int kChunkSizeL = Ktraits::kChunkSizeL; + constexpr int kChunkSizeC = Ktraits::kNEltsPerRow; + using input_t = typename Ktraits::input_t; + using vec_t = typename Ktraits::vec_t; + using weight_t = typename Ktraits::weight_t; + + // Shared memory. + __shared__ input_t x_smem[kWidth - 1 + kChunkSizeL][kChunkSizeC + kNElts]; + + const int batch_id = blockIdx.x; + const int chunk_l_id = blockIdx.y; + const int chunk_c_id = blockIdx.z; + const int tid = threadIdx.x; + const int l_idx = tid / kNThreadsPerC; + const int c_idx = tid % kNThreadsPerC; + input_t *x = reinterpret_cast(params.x_ptr) + batch_id * params.x_batch_stride + + (chunk_l_id * kChunkSizeL + l_idx) * params.x_l_stride + chunk_c_id * kChunkSizeC + c_idx * kNElts; + weight_t *weight = reinterpret_cast(params.weight_ptr) + + chunk_c_id * kChunkSizeC * params.weight_c_stride; + input_t *out = reinterpret_cast(params.out_ptr) + batch_id * params.out_batch_stride + + (chunk_l_id * kChunkSizeL + l_idx) * params.out_l_stride + chunk_c_id * kChunkSizeC + c_idx * kNElts; + int *seq_idx = !kHasSeqIdx ? nullptr : reinterpret_cast(params.seq_idx_ptr) + + batch_id * params.seqlen + chunk_l_id * kChunkSizeL; + input_t *initial_states = params.initial_states_ptr == nullptr || chunk_l_id > 0 ? nullptr + : reinterpret_cast(params.initial_states_ptr) + batch_id * params.initial_states_batch_stride + l_idx * params.initial_states_l_stride + chunk_c_id * kChunkSizeC + c_idx * kNElts; + // The last L-chunk will also have enough info to write to final states, since it also contain a few x values + // from the previous L-chunk. + input_t *final_states = params.final_states_ptr == nullptr || chunk_l_id < gridDim.y - 1 ? nullptr + : reinterpret_cast(params.final_states_ptr) + batch_id * params.final_states_batch_stride + l_idx * params.final_states_l_stride + chunk_c_id * kChunkSizeC + c_idx * kNElts; +/* + #pragma unroll + for (int l = 0; l < Ktraits::kNLoads; ++l) { + input_t x_vals_load[kNElts] = {0}; + if (chunk_l_id * kChunkSizeL + l * kLPerLoad + l_idx < params.seqlen + && chunk_c_id * kChunkSizeC + c_idx * kNElts < params.dim) { + reinterpret_cast(x_vals_load)[0] = *reinterpret_cast(x + l * kLPerLoad * params.x_l_stride); + } + reinterpret_cast(x_smem[kWidth - 1 + l * kLPerLoad + l_idx])[c_idx] = reinterpret_cast(x_vals_load)[0]; + } + */ + input_t x_vals_load[kNElts] = {0}; + if (chunk_l_id * kChunkSizeL + l_idx < params.seqlen + && chunk_c_id * kChunkSizeC + c_idx * kNElts < params.dim) { + reinterpret_cast(x_vals_load)[0] = *reinterpret_cast(x); + } + reinterpret_cast(x_smem[kWidth - 1 + l_idx])[c_idx] = reinterpret_cast(x_vals_load)[0]; + // Load the elements from the previous chunk that are needed for convolution. + if (l_idx < kWidth - 1) { + input_t x_vals_load[kNElts] = {0}; + if (chunk_l_id * kChunkSizeL + l_idx - (kWidth - 1) >= 0 + && chunk_l_id * kChunkSizeL + l_idx - (kWidth - 1) < params.seqlen + && chunk_c_id * kChunkSizeC + c_idx * kNElts < params.dim) { + reinterpret_cast(x_vals_load)[0] = *reinterpret_cast(x - (kWidth - 1) * params.x_l_stride); + } else if (initial_states != nullptr + && chunk_l_id * kChunkSizeL + l_idx - (kWidth - 1) < 0 + && chunk_c_id * kChunkSizeC + c_idx * kNElts < params.dim) { + reinterpret_cast(x_vals_load)[0] = *reinterpret_cast(initial_states); + } + reinterpret_cast(x_smem[l_idx])[c_idx] = reinterpret_cast(x_vals_load)[0]; + } + + __syncthreads(); + + if (final_states != nullptr + && l_idx < kWidth - 1 + && chunk_c_id * kChunkSizeC + c_idx * kNElts < params.dim) { + // x_smem[0] contains element at index chunk_l_id * kChunkSizeL - (kWidth - 1) + // So last few elements (index params.seqlen - kWidth + 1 + l_idx) are stored in x_smem[params.seqlen - kWidth + 1 + l_idx - (chunk_l_id * kChunkSizeL - kWidth + 1)][c_idx] + *reinterpret_cast(final_states) = reinterpret_cast(x_smem[params.seqlen + l_idx - chunk_l_id * kChunkSizeL])[c_idx]; + } + + constexpr int kLPerThread = constexpr_min(kChunkSizeL * kChunkSizeC / kNThreads, kChunkSizeL); + static_assert(kLPerThread * kNThreads == kChunkSizeL * kChunkSizeC); + constexpr int kNThreadsPerRow = kChunkSizeL / kLPerThread; + static_assert(kNThreadsPerRow * kLPerThread == kChunkSizeL); + // kChunkSizeL, kLPerThread, kNThreadsPerRow should be powers of 2 for simplicity + static_assert((kChunkSizeL & (kChunkSizeL - 1)) == 0); + static_assert((kLPerThread & (kLPerThread - 1)) == 0); + static_assert((kNThreadsPerRow & (kNThreadsPerRow - 1)) == 0); + static_assert(kNThreadsPerRow <= 32); + + const int row_idx = tid / kNThreadsPerRow; + const int col_idx = tid % kNThreadsPerRow; + + float bias_val = params.bias_ptr == nullptr || chunk_c_id * kChunkSizeC + row_idx >= params.dim ? 0.f : float(reinterpret_cast(params.bias_ptr)[chunk_c_id * kChunkSizeC + row_idx]); + float weight_vals[kWidth] = {0}; + if (chunk_c_id * kChunkSizeC + row_idx < params.dim) { + #pragma unroll + for (int w = 0; w < kWidth; ++w) { + weight_vals[w] = weight[row_idx * params.weight_c_stride + w * params.weight_width_stride]; + } + } + float x_vals[kWidth - 1 + kLPerThread]; + #pragma unroll + for (int i = 0; i < kWidth - 1 + kLPerThread; ++i) { + x_vals[i] = float(x_smem[col_idx * kLPerThread + i][row_idx]); + } + int seq_idx_thread[kWidth - 1 + kLPerThread]; + if constexpr (kHasSeqIdx) { + #pragma unroll + for (int i = 0; i < kWidth - 1 + kLPerThread; ++i) { + seq_idx_thread[i] = chunk_l_id * kChunkSizeL + col_idx * kLPerThread + i - (kWidth - 1) >= 0 ? seq_idx[col_idx * kLPerThread + i - (kWidth - 1)] : -1; + } + } + + float out_vals[kLPerThread]; + #pragma unroll + for (int i = 0; i < kLPerThread; ++i) { + out_vals[i] = bias_val; + const int seq_idx_cur = !kHasSeqIdx ? 0 : seq_idx_thread[i + kWidth - 1]; + // For padding tokens (seq_idx < 0), we skip the computation and set the output to 0. + if (seq_idx_cur < 0) { + out_vals[i] = 0.f; + continue; + } + #pragma unroll + for (int w = 0; w < kWidth; ++w) { + if constexpr (!kHasSeqIdx) { + out_vals[i] += weight_vals[w] * x_vals[i + w]; + } else { + out_vals[i] += seq_idx_thread[i + w] == seq_idx_cur ? weight_vals[w] * x_vals[i + w] : 0.f; + } + } + if (params.silu_activation) {out_vals[i] = out_vals[i] / (1 + expf(-out_vals[i])); } + } + + __syncthreads(); + #pragma unroll + for (int i = 0; i < kLPerThread; ++i) { x_smem[col_idx * kLPerThread + i][row_idx] = out_vals[i]; } + __syncthreads(); +/* + #pragma unroll + for (int l = 0; l < Ktraits::kNLoads; ++l) { + input_t out_vals_store[kNElts]; + reinterpret_cast(out_vals_store)[0] = reinterpret_cast(x_smem[l * kLPerLoad + l_idx])[c_idx]; + if (chunk_l_id * kChunkSizeL + l * kLPerLoad + l_idx < params.seqlen + && chunk_c_id * kChunkSizeC + c_idx * kNElts < params.dim) { + *reinterpret_cast(out + l * kLPerLoad * params.out_l_stride) = reinterpret_cast(out_vals_store)[0]; + } + } +*/ + input_t out_vals_store[kNElts]; + reinterpret_cast(out_vals_store)[0] = reinterpret_cast(x_smem[l_idx])[c_idx]; + if (chunk_l_id * kChunkSizeL + l_idx < params.seqlen + && chunk_c_id * kChunkSizeC + c_idx * kNElts < params.dim) { + *reinterpret_cast(out) = reinterpret_cast(out_vals_store)[0]; + } +} + +template +void causal_conv1d_channellast_fwd_launch(ConvParamsBase ¶ms, hipStream_t stream) { + BOOL_SWITCH(params.seq_idx_ptr != nullptr, kHasSeqIdx, [&] { + using Ktraits = Causal_conv1d_channellast_fwd_kernel_traits; + //using Ktraits = Causal_conv1d_channellast_fwd_kernel_traits; + // constexpr int kSmemSize = Ktraits::kSmemSize; + constexpr int kChunkSizeL = Ktraits::kChunkSizeL; + constexpr int kChunkSizeC = Ktraits::kNEltsPerRow; + const int n_chunks_L = (params.seqlen + kChunkSizeL - 1) / kChunkSizeL; + const int n_chunks_C = (params.dim + kChunkSizeC - 1) / kChunkSizeC; + dim3 grid(params.batch, n_chunks_L, n_chunks_C); + dim3 block(Ktraits::kNThreads); + auto kernel = &causal_conv1d_channellast_fwd_kernel; + hipLaunchKernelGGL(( kernel), dim3(grid), dim3(Ktraits::kNThreads), 0, stream, params); + /* + hipDeviceSynchronize(); + // if (kSmemSize >= 48 * 1024) { + // C10_HIP_CHECK(hipFuncSetAttribute( + // kernel, hipFuncAttributeMaxDynamicSharedMemorySize, kSmemSize)); + // } + //hipLaunchKernelGGL(( kernel), dim3(grid), dim3(Ktraits::kNThreads), kSmemSize, stream, params); + hipEvent_t start, stop; + hipEventCreate(&start); + hipEventCreate(&stop); + + // 记录开始时间 + hipEventRecord(start, 0); + for(int i =0; i<100;i++) + { + hipLaunchKernelGGL(( kernel), dim3(grid), dim3(Ktraits::kNThreads), 0, stream, params); + } + // 记录结束时间 + hipEventRecord(stop, 0); + hipEventSynchronize(stop); // 等待事件完成 + + // 计算运行时间（毫秒） + float run_time = 0; + hipEventElapsedTime(&run_time, start, stop); + + std::cout<<"------------------------the BMZ channellast the seqlen is:"< +void causal_conv1d_channellast_fwd_cuda(ConvParamsBase ¶ms, hipStream_t stream) { + if (params.width == 2) { + causal_conv1d_channellast_fwd_launch<256, 2, input_t, weight_t>(params, stream); + } else if (params.width == 3) { + causal_conv1d_channellast_fwd_launch<256, 3, input_t, weight_t>(params, stream); + } else if (params.width == 4) { + causal_conv1d_channellast_fwd_launch<256, 4, input_t, weight_t>(params, stream); + } +} + +template void causal_conv1d_fwd_cuda(ConvParamsBase ¶ms, hipStream_t stream); +template void causal_conv1d_fwd_cuda(ConvParamsBase ¶ms, hipStream_t stream); +template void causal_conv1d_fwd_cuda(ConvParamsBase ¶ms, hipStream_t stream); +template void causal_conv1d_fwd_cuda(ConvParamsBase ¶ms, hipStream_t stream); +template void causal_conv1d_fwd_cuda(ConvParamsBase ¶ms, hipStream_t stream); +template void causal_conv1d_fwd_cuda(ConvParamsBase ¶ms, hipStream_t stream); +template void causal_conv1d_fwd_cuda(ConvParamsBase ¶ms, hipStream_t stream); +template void causal_conv1d_fwd_cuda(ConvParamsBase ¶ms, hipStream_t stream); +template void causal_conv1d_fwd_cuda(ConvParamsBase ¶ms, hipStream_t stream); + +template void causal_conv1d_channellast_fwd_cuda(ConvParamsBase ¶ms, hipStream_t stream); +template void causal_conv1d_channellast_fwd_cuda(ConvParamsBase ¶ms, hipStream_t stream); +template void causal_conv1d_channellast_fwd_cuda(ConvParamsBase ¶ms, hipStream_t stream); +template void causal_conv1d_channellast_fwd_cuda(ConvParamsBase ¶ms, hipStream_t stream); +template void causal_conv1d_channellast_fwd_cuda(ConvParamsBase ¶ms, hipStream_t stream); +template void causal_conv1d_channellast_fwd_cuda(ConvParamsBase ¶ms, hipStream_t stream); +template void causal_conv1d_channellast_fwd_cuda(ConvParamsBase ¶ms, hipStream_t stream); +template void causal_conv1d_channellast_fwd_cuda(ConvParamsBase ¶ms, hipStream_t stream); +template void causal_conv1d_channellast_fwd_cuda(ConvParamsBase ¶ms, hipStream_t stream); diff --git a/xllm/core/kernels/dcu/causal_conv1d/causal_conv1d_new_interface.hip b/xllm/core/kernels/dcu/causal_conv1d/causal_conv1d_new_interface.hip new file mode 100644 index 0000000000..0068cfece1 --- /dev/null +++ b/xllm/core/kernels/dcu/causal_conv1d/causal_conv1d_new_interface.hip @@ -0,0 +1,294 @@ +// !!! This is a file automatically generated by hipify!!! +#include +#include +#include +#include +#include +#include "causal_conv1d.h" +#include "causal_conv1d_common_hip.h" +#include "static_switch.h" + +// #define BLOCK_M 64 +// #define BLOCK_N 256 + +namespace { + +template +__device__ __forceinline__ float to_float(T val) { return static_cast(val); } +template <> __device__ __forceinline__ float to_float(at::Half val) { return __half2float(val); } +template <> __device__ __forceinline__ float to_float(at::BFloat16 val) { return static_cast(val); } + +template +__device__ __forceinline__ T from_float(float val) { return static_cast(val); } +template <> __device__ __forceinline__ at::Half from_float(float val) { return __float2half(val); } +template <> __device__ __forceinline__ at::BFloat16 from_float(float val) { return static_cast(val); } + +template +struct Causal_conv1d_varlen_fwd_kernel_traits { + using input_t = input_t_; + using weight_t = weight_t_; + static constexpr int kBlockM = kBlockM_; + static constexpr int kBlockN = kBlockN_; + static constexpr int kNThreads = kNThreads_; + static constexpr int kWidth = kWidth_; + static constexpr int kNBytes = sizeof(input_t); + static_assert(kNBytes == 2 || kNBytes == 4); + static constexpr int kNElts = kNBytes == 4 ? 4 : 8; + static_assert(kWidth <= kNElts); + using vec_t = typename BytesToType::Type; +}; + + +template +__global__ __launch_bounds__(256, 1) +void causal_conv1d_fwd_kernel_hcu(ConvParamsBase params) { + using input_t = typename Ktraits::input_t; + using weight_t = typename Ktraits::weight_t; + constexpr int kWidth = Ktraits::kWidth; + constexpr int kNThreads = Ktraits::kNThreads; + constexpr int kBlockM = Ktraits::kBlockM; + constexpr int kBlockN = Ktraits::kBlockN; + + int batch_idx = blockIdx.x; + int chunk_idx = blockIdx.y; + int feat_block = blockIdx.z; + + int tid = threadIdx.x; + int feat = feat_block * kBlockN + tid; + bool feat_valid = feat < params.dim; + + if ((batch_idx == params.pad_slot_id) || (!feat_valid)) return; + + int2 read_query_2 = ((int2*)(params.query_start_loc_ptr + batch_idx))[0]; + int seq_start = read_query_2.x; + int seq_end = read_query_2.y; + int seqlen = seq_end - seq_start; + int token_offset = kBlockM * chunk_idx; + int segment_len = min(kBlockM, seqlen - token_offset); + if (segment_len <= 0) return; + + int cache_idx = params.conv_state_indices_ptr != nullptr ? params.conv_state_indices_ptr[batch_idx * params.conv_state_indices_stride] : batch_idx; + if (cache_idx == params.pad_slot_id) return; + + input_t* conv_states_ptr = reinterpret_cast(params.initial_states_ptr) + cache_idx * params.initial_states_batch_stride + feat * params.initial_states_c_stride; + input_t* x_ptr = reinterpret_cast(params.x_ptr) + feat * params.x_c_stride; + weight_t* w_ptr = reinterpret_cast(params.weight_ptr) + feat * params.weight_c_stride; + input_t* out_ptr = reinterpret_cast(params.out_ptr) + feat * params.out_c_stride + (seq_start + token_offset) * params.out_l_stride; + + // Reuse + int state_len = params.conv_state_len; + int stride_istate_token = params.initial_states_l_stride; + int stride_x_token = params.x_l_stride; + int stride_w_width = params.weight_width_stride; + int stride_o_token = params.out_l_stride; + + float col0 = 0.0f, col1 = 0.0f, col2 = 0.0f, col3 = 0.0f; + bool load_init_state = kHasInitialStates ? params.has_initial_state_ptr[batch_idx] : false; + + if (chunk_idx == 0) { + if (load_init_state) { + if constexpr (kWidth == 2) { + col0 = to_float(conv_states_ptr[(state_len - 1) * stride_istate_token]); + } else if constexpr (kWidth == 3) { + col1 = to_float(conv_states_ptr[(state_len - 1) * stride_istate_token]); + col0 = to_float(conv_states_ptr[(state_len - 2) * stride_istate_token]); + } else if constexpr (kWidth == 4) { + col2 = to_float(conv_states_ptr[(state_len - 1) * stride_istate_token]); + col1 = to_float(conv_states_ptr[(state_len - 2) * stride_istate_token]); + col0 = to_float(conv_states_ptr[(state_len - 3) * stride_istate_token]); + } else if constexpr (kWidth == 5) { + col3 = to_float(conv_states_ptr[(state_len - 1) * stride_istate_token]); + col2 = to_float(conv_states_ptr[(state_len - 2) * stride_istate_token]); + col1 = to_float(conv_states_ptr[(state_len - 3) * stride_istate_token]); + col0 = to_float(conv_states_ptr[(state_len - 4) * stride_istate_token]); + } + } + + // Update conv_state cache + if constexpr (kHasCache) { + if (state_len <= seqlen) { + int start_pos = seqlen - state_len; + for (int s = 0; s < state_len; s++) { + int x_idx = seq_start + start_pos + s; + input_t val = input_t(0); + if (x_idx < seq_end) { + val = x_ptr[x_idx * stride_x_token]; + } + conv_states_ptr[s * stride_istate_token] = val; + } + } else { + if (load_init_state) { + int shift = seqlen; + for (int s = 0; s < state_len - shift; s++) { + conv_states_ptr[s * stride_istate_token] = conv_states_ptr[(s + shift) * stride_istate_token]; + } + for (int s = 0; s < shift; s++) { + int x_idx = seq_start + s; + input_t val = input_t(0); + if (x_idx < seq_end) { + val = x_ptr[x_idx * stride_x_token]; + } + conv_states_ptr[(state_len - shift + s) * stride_istate_token] = val; + } + } else { + int zeros_prefix = state_len - seqlen; + for (int s = 0; s < zeros_prefix; s++) { + conv_states_ptr[s * stride_istate_token] = input_t(0); + } + for (int s = 0; s < seqlen; s++) { + int x_idx = seq_start + s; + input_t val = input_t(0); + if (x_idx < seq_end) { + val = x_ptr[x_idx * stride_x_token]; + } + conv_states_ptr[(zeros_prefix + s) * stride_istate_token] = val; + } + } + } + } + } else { + // chunk_idx > 0: read prior tokens from x + if constexpr (kWidth == 2) { + col0 = to_float(x_ptr[(seq_start + token_offset - 1) * stride_x_token]); + } else if constexpr (kWidth == 3) { + col1 = to_float(x_ptr[(seq_start + token_offset - 1) * stride_x_token]); + col0 = to_float(x_ptr[(seq_start + token_offset - 2) * stride_x_token]); + } else if constexpr (kWidth == 4) { + col2 = to_float(x_ptr[(seq_start + token_offset - 1) * stride_x_token]); + col1 = to_float(x_ptr[(seq_start + token_offset - 2) * stride_x_token]); + col0 = to_float(x_ptr[(seq_start + token_offset - 3) * stride_x_token]); + } else if constexpr (kWidth == 5) { + col3 = to_float(x_ptr[(seq_start + token_offset - 1) * stride_x_token]); + col2 = to_float(x_ptr[(seq_start + token_offset - 2) * stride_x_token]); + col1 = to_float(x_ptr[(seq_start + token_offset - 3) * stride_x_token]); + col0 = to_float(x_ptr[(seq_start + token_offset - 4) * stride_x_token]); + } + } + + // Load bias + float acc_bias = params.bias_ptr != nullptr ? to_float(reinterpret_cast(params.bias_ptr)[feat]) : 0.0f; + + // Load weights + float w0 = 0.0f, w1 = 0.0f, w2 = 0.0f, w3 = 0.0f, w4 = 0.0f; + if constexpr (kWidth >= 2) { + w0 = to_float(w_ptr[0 * stride_w_width]); + w1 = to_float(w_ptr[1 * stride_w_width]); + } + if constexpr (kWidth >= 3) { + w2 = to_float(w_ptr[2 * stride_w_width]); + } + if constexpr (kWidth >= 4) { + w3 = to_float(w_ptr[3 * stride_w_width]); + } + if constexpr (kWidth >= 5) { + w4 = to_float(w_ptr[4 * stride_w_width]); + } + + // Main convolution loop + for (int t = 0; t < segment_len; t++) { + float acc = acc_bias; + + // Load current token + float cur_x = 0.0f; + cur_x = to_float(x_ptr[(seq_start + token_offset + t) * stride_x_token]); + + if constexpr (kWidth == 2) { + acc += col0 * w0; + acc += cur_x * w1; + col0 = cur_x; + } else if constexpr (kWidth == 3) { + acc += col0 * w0; + acc += col1 * w1; + acc += cur_x * w2; + col0 = col1; + col1 = cur_x; + } else if constexpr (kWidth == 4) { + acc += col0 * w0; + acc += col1 * w1; + acc += col2 * w2; + acc += cur_x * w3; + col0 = col1; + col1 = col2; + col2 = cur_x; + } else if constexpr (kWidth == 5) { + acc += col0 * w0; + acc += col1 * w1; + acc += col2 * w2; + acc += col3 * w3; + acc += cur_x * w4; + col0 = col1; + col1 = col2; + col2 = col3; + col3 = cur_x; + } + + acc = params.silu_activation ? acc / (1.0f + expf(-acc)) : acc; + + out_ptr[t * stride_o_token] = from_float(acc); + } +} + +} // anonymous namespace + + +template +void causal_conv1d_varlen_fwd_launch(ConvParamsBase ¶ms, hipStream_t stream) { + BOOL_SWITCH(params.initial_states_ptr != nullptr, kHasCache, [&] { + BOOL_SWITCH(params.has_initial_state_ptr != nullptr, kHasInitialStates, [&] { + using Ktraits = Causal_conv1d_varlen_fwd_kernel_traits; + const int n_chunks_L = (params.max_seqlen + Ktraits::kBlockM - 1) / Ktraits::kBlockM; + const int n_chunks_C = (params.dim + Ktraits::kBlockN - 1) / Ktraits::kBlockN; + + dim3 grid(params.batch, n_chunks_L, n_chunks_C); + dim3 block(Ktraits::kNThreads); + + auto kernel = &causal_conv1d_fwd_kernel_hcu; + hipLaunchKernelGGL(( kernel), dim3(grid), dim3(Ktraits::kNThreads), 0, stream, params); + + C10_HIP_KERNEL_LAUNCH_CHECK(); + }); + }); +} + + +template +void causal_conv1d_varlen_fwd_cuda(ConvParamsBase ¶ms, hipStream_t stream) { + if (params.width == 2) { + causal_conv1d_varlen_fwd_launch<256, 2, input_t, weight_t>(params, stream); + } else if (params.width == 3) { + causal_conv1d_varlen_fwd_launch<256, 3, input_t, weight_t>(params, stream); + } else if (params.width == 4) { + causal_conv1d_varlen_fwd_launch<256, 4, input_t, weight_t>(params, stream); + } +} + +template +void causal_conv1d_varlen_channellast_fwd_cuda(ConvParamsBase ¶ms, hipStream_t stream) { + if (params.width == 2) { + causal_conv1d_varlen_fwd_launch<256, 2, input_t, weight_t>(params, stream); + } else if (params.width == 3) { + causal_conv1d_varlen_fwd_launch<256, 3, input_t, weight_t>(params, stream); + } else if (params.width == 4) { + causal_conv1d_varlen_fwd_launch<256, 4, input_t, weight_t>(params, stream); + } +} + +template void causal_conv1d_varlen_fwd_cuda(ConvParamsBase ¶ms, hipStream_t stream); +template void causal_conv1d_varlen_fwd_cuda(ConvParamsBase ¶ms, hipStream_t stream); +template void causal_conv1d_varlen_fwd_cuda(ConvParamsBase ¶ms, hipStream_t stream); +template void causal_conv1d_varlen_fwd_cuda(ConvParamsBase ¶ms, hipStream_t stream); +template void causal_conv1d_varlen_fwd_cuda(ConvParamsBase ¶ms, hipStream_t stream); +template void causal_conv1d_varlen_fwd_cuda(ConvParamsBase ¶ms, hipStream_t stream); +template void causal_conv1d_varlen_fwd_cuda(ConvParamsBase ¶ms, hipStream_t stream); +template void causal_conv1d_varlen_fwd_cuda(ConvParamsBase ¶ms, hipStream_t stream); +template void causal_conv1d_varlen_fwd_cuda(ConvParamsBase ¶ms, hipStream_t stream); + +template void causal_conv1d_varlen_channellast_fwd_cuda(ConvParamsBase ¶ms, hipStream_t stream); +template void causal_conv1d_varlen_channellast_fwd_cuda(ConvParamsBase ¶ms, hipStream_t stream); +template void causal_conv1d_varlen_channellast_fwd_cuda(ConvParamsBase ¶ms, hipStream_t stream); +template void causal_conv1d_varlen_channellast_fwd_cuda(ConvParamsBase ¶ms, hipStream_t stream); +template void causal_conv1d_varlen_channellast_fwd_cuda(ConvParamsBase ¶ms, hipStream_t stream); +template void causal_conv1d_varlen_channellast_fwd_cuda(ConvParamsBase ¶ms, hipStream_t stream); +template void causal_conv1d_varlen_channellast_fwd_cuda(ConvParamsBase ¶ms, hipStream_t stream); +template void causal_conv1d_varlen_channellast_fwd_cuda(ConvParamsBase ¶ms, hipStream_t stream); +template void causal_conv1d_varlen_channellast_fwd_cuda(ConvParamsBase ¶ms, hipStream_t stream); \ No newline at end of file diff --git a/xllm/core/kernels/dcu/causal_conv1d/causal_conv1d_update.hip b/xllm/core/kernels/dcu/causal_conv1d/causal_conv1d_update.hip new file mode 100644 index 0000000000..1bce937da6 --- /dev/null +++ b/xllm/core/kernels/dcu/causal_conv1d/causal_conv1d_update.hip @@ -0,0 +1,149 @@ +// !!! This is a file automatically generated by hipify!!! +#include +#include "hip/hip_runtime.h" +/****************************************************************************** + * Copyright (c) 2023, Tri Dao. + ******************************************************************************/ + +#include +#include +#include // For C10_HIP_CHECK and C10_HIP_KERNEL_LAUNCH_CHECK + +#include "causal_conv1d.h" +#include "causal_conv1d_common_hip.h" +#include "static_switch.h" + +template +struct Causal_conv1d_update_kernel_traits { + using input_t = input_t_; + using weight_t = weight_t_; + static constexpr int kNThreads = kNThreads_; + static constexpr int kWidth = kWidth_; + static constexpr int kNBytes = sizeof(input_t); + static_assert(kNBytes == 2 || kNBytes == 4); +}; + +template +__global__ __launch_bounds__(Ktraits::kNThreads) +void causal_conv1d_update_kernel(ConvParamsBase params) { + constexpr int kWidth = Ktraits::kWidth; + constexpr int kNThreads = Ktraits::kNThreads; + using input_t = typename Ktraits::input_t; + using weight_t = typename Ktraits::weight_t; + + const int tidx = threadIdx.x; + const int batch_id = blockIdx.x; + const int channel_id = blockIdx.y * kNThreads + tidx; + if (channel_id >= params.dim) return; + + input_t *x = reinterpret_cast(params.x_ptr) + batch_id * params.x_batch_stride + + channel_id * params.x_c_stride; + input_t *out = reinterpret_cast(params.out_ptr) + batch_id * params.out_batch_stride + + channel_id * params.out_c_stride; + + // If params.conv_state_batch_indices is set, then the conv state is gathered from the conv state tensor + // along the batch axis. Otherwise, the conv state coordinate is the same as the batch id. + const int conv_state_batch_coord = params.conv_state_indices_ptr == nullptr + ? batch_id + : params.conv_state_indices_ptr[batch_id]; + + // Skip padding tokens. + if (conv_state_batch_coord < 0) { + #pragma unroll 2 + for (int i = 0; i < params.seqlen; ++i) { + out[i * params.out_l_stride] = input_t(0.f); + } + return; + } + input_t *conv_state = reinterpret_cast(params.conv_state_ptr) + + conv_state_batch_coord * params.conv_state_batch_stride + + channel_id * params.conv_state_c_stride; + weight_t *weight = reinterpret_cast(params.weight_ptr) + channel_id * params.weight_c_stride; + float bias_val = params.bias_ptr == nullptr ? 0.f : float(reinterpret_cast(params.bias_ptr)[channel_id]); + + int state_len = params.conv_state_len; + int advance_len = params.seqlen; + int cache_seqlen = kIsCircularBuffer ? params.cache_seqlens[batch_id] % state_len : 0; + int update_idx = cache_seqlen - (kWidth - 1); + update_idx = update_idx < 0 ? update_idx + state_len : update_idx; + + float weight_vals[kWidth] = {0}; + #pragma unroll + for (int i = 0; i < kWidth; ++i) { weight_vals[i] = float(weight[i * params.weight_width_stride]); } + + float x_vals[kWidth] = {0}; + if constexpr (!kIsCircularBuffer) { + #pragma unroll 2 + for (int i = 0; i < state_len - advance_len - (kWidth - 1); ++i) { + conv_state[i * params.conv_state_l_stride] = conv_state[(i + advance_len) * params.conv_state_l_stride]; + } + #pragma unroll + for (int i = 0; i < kWidth - 1; ++i) { + input_t state_val = conv_state[(state_len - (kWidth - 1) + i) * params.conv_state_l_stride]; + if (i < advance_len + (kWidth - 1) && state_len - advance_len - (kWidth - 1) + i >= 0) { + conv_state[(state_len - advance_len - (kWidth - 1) + i) * params.conv_state_l_stride] = state_val; + } + x_vals[i] = float(state_val); + } + } else { + #pragma unroll + for (int i = 0; i < kWidth - 1; ++i, update_idx = update_idx + 1 >= state_len ? update_idx + 1 - state_len : update_idx + 1) { + input_t state_val = conv_state[update_idx * params.conv_state_l_stride]; + x_vals[i] = float(state_val); + } + } + #pragma unroll 2 + for (int i = 0; i < params.seqlen; ++i) { + input_t x_val = x[i * params.x_l_stride]; + if constexpr (!kIsCircularBuffer) { + if (i < advance_len && state_len - advance_len + i >= 0) { + conv_state[(state_len - advance_len + i) * params.conv_state_l_stride] = x_val; + } + } else { + conv_state[update_idx * params.conv_state_l_stride] = x_val; + ++update_idx; + update_idx = update_idx >= state_len ? update_idx - state_len : update_idx; + } + x_vals[kWidth - 1] = float(x_val); + float out_val = bias_val; + #pragma unroll + for (int j = 0; j < kWidth; ++j) { out_val += weight_vals[j] * x_vals[j]; } + if (params.silu_activation) { out_val = out_val / (1 + expf(-out_val)); } + out[i * params.out_l_stride] = input_t(out_val); + // Shift the input buffer by 1 + #pragma unroll + for (int i = 0; i < kWidth - 1; ++i) { x_vals[i] = x_vals[i + 1]; } + } +} + +template +void causal_conv1d_update_launch(ConvParamsBase ¶ms, hipStream_t stream) { + using Ktraits = Causal_conv1d_update_kernel_traits; + dim3 grid(params.batch, (params.dim + kNThreads - 1) / kNThreads); + auto kernel = params.cache_seqlens == nullptr + ? &causal_conv1d_update_kernel + : &causal_conv1d_update_kernel; + hipLaunchKernelGGL(( kernel), dim3(grid), dim3(Ktraits::kNThreads), 0, stream, params); + C10_HIP_KERNEL_LAUNCH_CHECK(); +} + +template +void causal_conv1d_update_cuda(ConvParamsBase ¶ms, hipStream_t stream) { + if (params.width == 2) { + causal_conv1d_update_launch<64, 2, input_t, weight_t>(params, stream); + } else if (params.width == 3) { + causal_conv1d_update_launch<64, 3, input_t, weight_t>(params, stream); + } else if (params.width == 4) { + causal_conv1d_update_launch<64, 4, input_t, weight_t>(params, stream); + } +} + +template void causal_conv1d_update_cuda(ConvParamsBase ¶ms, hipStream_t stream); +template void causal_conv1d_update_cuda(ConvParamsBase ¶ms, hipStream_t stream); +template void causal_conv1d_update_cuda(ConvParamsBase ¶ms, hipStream_t stream); +template void causal_conv1d_update_cuda(ConvParamsBase ¶ms, hipStream_t stream); +template void causal_conv1d_update_cuda(ConvParamsBase ¶ms, hipStream_t stream); +template void causal_conv1d_update_cuda(ConvParamsBase ¶ms, hipStream_t stream); +template void causal_conv1d_update_cuda(ConvParamsBase ¶ms, hipStream_t stream); +template void causal_conv1d_update_cuda(ConvParamsBase ¶ms, hipStream_t stream); +template void causal_conv1d_update_cuda(ConvParamsBase ¶ms, hipStream_t stream); diff --git a/xllm/core/kernels/dcu/causal_conv1d/static_switch.h b/xllm/core/kernels/dcu/causal_conv1d/static_switch.h new file mode 100644 index 0000000000..11c8768423 --- /dev/null +++ b/xllm/core/kernels/dcu/causal_conv1d/static_switch.h @@ -0,0 +1,26 @@ +// Inspired by +// https://github.com/NVIDIA/DALI/blob/main/include/dali/core/static_switch.h +// and https://github.com/pytorch/pytorch/blob/master/aten/src/ATen/Dispatch.h + +#pragma once + +/// @param COND - a boolean expression to switch by +/// @param CONST_NAME - a name given for the constexpr bool variable. +/// @param ... - code to execute for true and false +/// +/// Usage: +/// ``` +/// BOOL_SWITCH(flag, BoolConst, [&] { +/// some_function(...); +/// }); +/// ``` +#define BOOL_SWITCH(COND, CONST_NAME, ...) \ + [&] { \ + if (COND) { \ + static constexpr bool CONST_NAME = true; \ + return __VA_ARGS__(); \ + } else { \ + static constexpr bool CONST_NAME = false; \ + return __VA_ARGS__(); \ + } \ + }() diff --git a/xllm/core/kernels/dcu/causal_conv1d_adapter.cpp b/xllm/core/kernels/dcu/causal_conv1d_adapter.cpp new file mode 100644 index 0000000000..b47570fef7 --- /dev/null +++ b/xllm/core/kernels/dcu/causal_conv1d_adapter.cpp @@ -0,0 +1,253 @@ +/* Copyright 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. +==============================================================================*/ + +#include "causal_conv1d_adapter.h" + +#include +#include +#include +#include + +#include +#include + +#include "causal_conv1d/causal_conv1d.h" + +// Template kernel forward declarations from the copied .hip files. +template +void causal_conv1d_channellast_fwd_cuda(ConvParamsBase& params, + hipStream_t stream); +template +void causal_conv1d_varlen_channellast_fwd_cuda(ConvParamsBase& params, + hipStream_t stream); +template +void causal_conv1d_update_cuda(ConvParamsBase& params, hipStream_t stream); + +namespace xllm { +namespace kernel { +namespace dcu { + +namespace { + +// Fill common fields of ConvParamsBase for forward passes. +// `varlen=true` selects the packed-varlen stride layout for x/out. +void set_conv_params_fwd(ConvParamsBase& params, + int64_t batch, + int64_t dim, + int64_t seqlen, + int64_t width, + const torch::Tensor& x, + const torch::Tensor& weight, + const torch::Tensor& out, + void* bias_ptr, + bool silu_activation, + bool varlen) { + std::memset(¶ms, 0, sizeof(params)); + params.batch = static_cast(batch); + params.dim = static_cast(dim); + params.seqlen = static_cast(seqlen); + params.width = static_cast(width); + params.silu_activation = silu_activation; + params.pad_slot_id = -1; + + params.x_ptr = x.data_ptr(); + params.weight_ptr = weight.data_ptr(); + params.bias_ptr = bias_ptr; + params.out_ptr = out.data_ptr(); + + params.x_batch_stride = varlen ? 0 : x.stride(0); + params.x_c_stride = varlen ? x.stride(0) : x.stride(1); + params.x_l_stride = varlen ? x.stride(1) : x.stride(-1); + params.weight_c_stride = weight.stride(0); + params.weight_width_stride = weight.stride(1); + params.out_batch_stride = varlen ? 0 : out.stride(0); + params.out_c_stride = varlen ? out.stride(0) : out.stride(1); + params.out_l_stride = varlen ? out.stride(1) : out.stride(-1); +} + +#define DISPATCH_ITYPE(TYPE, NAME, ...) \ + do { \ + if ((TYPE) == at::ScalarType::Half) { \ + using input_t = at::Half; \ + __VA_ARGS__(); \ + } else if ((TYPE) == at::ScalarType::BFloat16) { \ + using input_t = at::BFloat16; \ + __VA_ARGS__(); \ + } else if ((TYPE) == at::ScalarType::Float) { \ + using input_t = float; \ + __VA_ARGS__(); \ + } else { \ + TORCH_CHECK(false, NAME " unsupported input dtype"); \ + } \ + } while (0) + +#define DISPATCH_WTYPE(TYPE, NAME, ...) \ + do { \ + if ((TYPE) == at::ScalarType::Half) { \ + using weight_t = at::Half; \ + __VA_ARGS__(); \ + } else if ((TYPE) == at::ScalarType::BFloat16) { \ + using weight_t = at::BFloat16; \ + __VA_ARGS__(); \ + } else if ((TYPE) == at::ScalarType::Float) { \ + using weight_t = float; \ + __VA_ARGS__(); \ + } else { \ + TORCH_CHECK(false, NAME " unsupported weight dtype"); \ + } \ + } while (0) + +} // namespace + +torch::Tensor causal_conv1d_varlen_fwd( + const torch::Tensor& x_tc, + const torch::Tensor& weight_ck, + const std::vector& cu_seqlens, + const std::vector& state_indices, + torch::Tensor& conv_cache, + bool activation) { + TORCH_CHECK(x_tc.dim() == 2, "x_tc must be [T, C]"); + TORCH_CHECK(weight_ck.dim() == 2, "weight_ck must be [C, K]"); + TORCH_CHECK(conv_cache.dim() == 3, "conv_cache must be [N, K-1, C]"); + TORCH_CHECK(cu_seqlens.size() >= 2, "cu_seqlens needs N+1 entries"); + const int64_t N = static_cast(cu_seqlens.size()) - 1; + TORCH_CHECK(static_cast(state_indices.size()) == N, + "state_indices size must equal N"); + TORCH_CHECK(x_tc.dtype() == weight_ck.dtype(), + "x and weight must share dtype"); + TORCH_CHECK(x_tc.dtype() == conv_cache.dtype(), + "x and conv_cache must share dtype"); + + const int64_t T_total = x_tc.size(0); + const int64_t C = x_tc.size(1); + const int64_t K = weight_ck.size(1); + TORCH_CHECK(weight_ck.size(0) == C, "weight_ck[0] must equal C"); + TORCH_CHECK(conv_cache.size(1) == K - 1 && conv_cache.size(2) == C, + "conv_cache shape must be [N, K-1, C]"); + TORCH_CHECK(K >= 2 && K <= 4, "kernel supports width in [2, 4]"); + TORCH_CHECK(C % 8 == 0, "channel dim must be divisible by 8"); + + auto device = x_tc.device(); + c10::DeviceGuard guard(device); + + // Reinterpret x/out as channel-last [C, T] views (stride(0)==1). + auto x_ct = x_tc.transpose(0, 1); + auto out_tc = torch::empty_like(x_tc); + auto out_ct = out_tc.transpose(0, 1); + // Reinterpret conv_cache to [N, C, K-1] view; stride(1)==1 (channel-first). + auto conv_state = conv_cache.transpose(1, 2); + + auto int32_opts = torch::TensorOptions().dtype(torch::kInt32).device(device); + std::vector cu32(cu_seqlens.begin(), cu_seqlens.end()); + auto query_start_loc = torch::tensor(cu32, int32_opts); + std::vector idx32(state_indices.begin(), state_indices.end()); + auto cache_indices = torch::tensor(idx32, int32_opts); + + // Determine max seqlen for grid sizing. + int64_t max_seqlen = 0; + for (int64_t i = 0; i < N; ++i) { + max_seqlen = std::max(max_seqlen, cu_seqlens[i + 1] - cu_seqlens[i]); + } + + ConvParamsBase params; + set_conv_params_fwd(params, + N, + C, + T_total, + K, + x_ct, + weight_ck, + out_ct, + /*bias_ptr=*/nullptr, + activation, + /*varlen=*/true); + params.max_seqlen = static_cast(max_seqlen); + params.query_start_loc_ptr = query_start_loc.data_ptr(); + params.conv_state_indices_ptr = cache_indices.data_ptr(); + params.conv_state_indices_stride = cache_indices.stride(0); + + params.initial_states_ptr = conv_state.data_ptr(); + params.initial_states_batch_stride = conv_state.stride(0); + params.initial_states_c_stride = conv_state.stride(1); + params.initial_states_l_stride = conv_state.stride(2); + params.conv_state_len = static_cast(K - 1); + + auto stream = at::hip::getCurrentHIPStreamMasqueradingAsCUDA().stream(); + DISPATCH_ITYPE(x_tc.scalar_type(), "causal_conv1d_varlen_fwd", [&] { + DISPATCH_WTYPE(weight_ck.scalar_type(), "causal_conv1d_varlen_fwd", [&] { + causal_conv1d_varlen_channellast_fwd_cuda(params, + stream); + }); + }); + return out_tc; +} + +torch::Tensor causal_conv1d_update(const torch::Tensor& x_bc, + const torch::Tensor& weight_ck, + const torch::Tensor& state_indices, + torch::Tensor& conv_cache, + bool activation) { + TORCH_CHECK(x_bc.dim() == 2, "x_bc must be [B, C]"); + TORCH_CHECK(weight_ck.dim() == 2, "weight_ck must be [C, K]"); + TORCH_CHECK(conv_cache.dim() == 3, "conv_cache must be [N, K-1, C]"); + TORCH_CHECK(state_indices.scalar_type() == torch::kInt32, + "state_indices must be int32"); + + const int64_t B = x_bc.size(0); + const int64_t C = x_bc.size(1); + const int64_t K = weight_ck.size(1); + TORCH_CHECK(weight_ck.size(0) == C, "weight_ck[0] must equal C"); + TORCH_CHECK(conv_cache.size(1) == K - 1 && conv_cache.size(2) == C, + "conv_cache shape must be [N, K-1, C]"); + TORCH_CHECK(K >= 2 && K <= 4, "kernel supports width in [2, 4]"); + + auto device = x_bc.device(); + c10::DeviceGuard guard(device); + + auto out_bc = torch::empty_like(x_bc); + // Present x/out as [B, C, seqlen=1] to match kernel expectation. + auto x_bcs = x_bc.unsqueeze(-1); + auto out_bcs = out_bc.unsqueeze(-1); + auto conv_state = conv_cache.transpose(1, 2); // [N, C, K-1] + auto state_idx_i32 = state_indices.contiguous(); + + ConvParamsBase params; + set_conv_params_fwd(params, + B, + C, + /*seqlen=*/1, + K, + x_bcs, + weight_ck, + out_bcs, + /*bias_ptr=*/nullptr, + activation, + /*varlen=*/false); + params.conv_state_ptr = conv_state.data_ptr(); + params.conv_state_len = static_cast(K - 1); + params.conv_state_batch_stride = conv_state.stride(0); + params.conv_state_c_stride = conv_state.stride(1); + params.conv_state_l_stride = conv_state.stride(2); + params.conv_state_indices_ptr = state_idx_i32.data_ptr(); + + auto stream = at::hip::getCurrentHIPStreamMasqueradingAsCUDA().stream(); + DISPATCH_ITYPE(x_bc.scalar_type(), "causal_conv1d_update", [&] { + DISPATCH_WTYPE(weight_ck.scalar_type(), "causal_conv1d_update", [&] { + causal_conv1d_update_cuda(params, stream); + }); + }); + return out_bc; +} + +} // namespace dcu +} // namespace kernel +} // namespace xllm diff --git a/xllm/core/kernels/dcu/causal_conv1d_adapter.h b/xllm/core/kernels/dcu/causal_conv1d_adapter.h new file mode 100644 index 0000000000..711062e981 --- /dev/null +++ b/xllm/core/kernels/dcu/causal_conv1d_adapter.h @@ -0,0 +1,58 @@ +/* Copyright 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 + +#include + +#include + +namespace xllm { +namespace kernel { +namespace dcu { + +// Varlen (packed) causal conv1d forward for prefill. +// +// x_tc : [T_total, C] bf16/fp16/fp32, contiguous +// weight_ck : [C, K] same dtype +// cu_seqlens : N+1 host vec, prefix sums (only >0 seqs, no padding rows) +// state_indices: N host vec of ssm-cache slot ids for each real sequence +// conv_cache : [num_slots, K-1, C] (xllm layout) — updated in place +// activation : true = silu +// +// Returns out_tc: [T_total, C] same dtype as x_tc. +torch::Tensor causal_conv1d_varlen_fwd( + const torch::Tensor& x_tc, + const torch::Tensor& weight_ck, + const std::vector& cu_seqlens, + const std::vector& state_indices, + torch::Tensor& conv_cache, + bool activation); + +// Decode single-step update. +// +// x_bc : [B, C] bf16/fp16/fp32 +// weight_ck : [C, K] same dtype +// state_indices : [B] int32 device tensor of ssm-cache slot ids +// conv_cache : [num_slots, K-1, C] updated in place +// activation : true = silu +// +// Returns out_bc: [B, C]. +torch::Tensor causal_conv1d_update(const torch::Tensor& x_bc, + const torch::Tensor& weight_ck, + const torch::Tensor& state_indices, + torch::Tensor& conv_cache, + bool activation); + +} // namespace dcu +} // namespace kernel +} // namespace xllm diff --git a/xllm/core/layers/dcu/CMakeLists.txt b/xllm/core/layers/dcu/CMakeLists.txt index 8dcb44cd13..43eac8d0e0 100644 --- a/xllm/core/layers/dcu/CMakeLists.txt +++ b/xllm/core/layers/dcu/CMakeLists.txt @@ -28,6 +28,8 @@ cc_library( deepseek_v2_decoder_layer_impl.h qwen3_5_decoder_layer.h qwen3_gated_delta_net_base.h + qwen3_gated_delta_net_helpers.h + qwen3_gated_delta_net_impls.h qwen3_5_gated_delta_net.h SRCS attention.cpp @@ -39,6 +41,8 @@ cc_library( deepseek_v2_decoder_layer_impl.cpp qwen3_5_decoder_layer.cpp qwen3_gated_delta_net_base.cpp + qwen3_gated_delta_net_torch_impl.cpp + qwen3_gated_delta_net_kernel_impl.cpp qwen3_5_gated_delta_net.cpp DEPS :common_layers diff --git a/xllm/core/layers/dcu/qwen3_gated_delta_net_base.cpp b/xllm/core/layers/dcu/qwen3_gated_delta_net_base.cpp index f9bb9eea21..8b656c1e22 100644 --- a/xllm/core/layers/dcu/qwen3_gated_delta_net_base.cpp +++ b/xllm/core/layers/dcu/qwen3_gated_delta_net_base.cpp @@ -15,270 +15,142 @@ limitations under the License. #include #include +#include +#include #include #include "qwen3_5_gated_delta_net.h" +#include "qwen3_gated_delta_net_impls.h" namespace xllm { namespace layer { namespace { -// ========================================================================= -// Pure-torch fallback implementations (no NPU/DCU kernels required) -// ========================================================================= - -torch::Tensor l2norm(const torch::Tensor& x, int64_t dim, double eps = 1e-6) { - auto norm = torch::sqrt(torch::sum(torch::square(x), dim, true) + eps); - return x / norm; +// Runtime backend selection. Both flags default off (use optimized kernel). +// Setting them to "1" swaps that operator to the pure-torch fallback for +// debugging / A-B testing. +inline bool use_torch_conv1d() { + static const bool v = [] { + const char* e = std::getenv("XLLM_DCU_USE_TORCH_CONV1D"); + return e && std::string_view(e) == "1"; + }(); + return v; } -torch::Tensor repeat_tensor_heads(const torch::Tensor& tensor, - int64_t target_heads, - int64_t head_dim) { - const int64_t current_heads = tensor.size(head_dim); - if (current_heads == target_heads) { - return tensor; - } - CHECK_GT(current_heads, 0) << "current heads must be positive"; - CHECK_EQ(target_heads % current_heads, 0) - << "target heads must be divisible by current heads"; - - const int64_t repeats = target_heads / current_heads; - std::vector view_shape = tensor.sizes().vec(); - view_shape.insert(view_shape.begin() + head_dim + 1, 1); - std::vector expand_shape = view_shape; - expand_shape[head_dim + 1] = repeats; - std::vector output_shape = tensor.sizes().vec(); - output_shape[head_dim] = target_heads; - return tensor.unsqueeze(head_dim + 1) - .expand(expand_shape) - .reshape(output_shape) - .contiguous(); +inline bool use_torch_recurrent_prefill() { + static const bool v = [] { + const char* e = std::getenv("XLLM_DCU_USE_TORCH_RECURRENT"); + return e && std::string_view(e) == "1"; + }(); + return v; } -// ----------------------------------------------------------------------- -// Torch causal conv1d (prefill) -// conv_weight_2d: [out_channels, kernel_size] from ColumnParallelLinear -// conv_cache: [num_sequences, kernel_size - 1, channels] -// ----------------------------------------------------------------------- -torch::Tensor torch_causal_conv1d(const torch::Tensor& flat_input, - const torch::Tensor& conv_weight_2d, - torch::Tensor& conv_cache, - const std::vector& cu_seqlens, - const std::vector& state_indices, - int32_t kernel_size, - bool activation) { - // flat_input: [total_tokens, channels] - // conv_weight_2d: [kernel_size, channels] (from conv1d_->weight()) - int64_t channels = flat_input.size(1); - int64_t batch_size = cu_seqlens.size() - 1; - auto options = flat_input.options(); - // Transpose to [channels, kernel_size] for per-channel dot product - auto weight = conv_weight_2d.transpose(0, 1).contiguous(); - // weight: [channels, kernel_size] - - std::vector outputs; - outputs.reserve(batch_size); - - for (int64_t b = 0; b < batch_size; ++b) { - int64_t start = cu_seqlens[b]; - int64_t end = cu_seqlens[b + 1]; - int64_t seq_len = end - start; - if (seq_len == 0) continue; - - // seq_input: [seq_len, channels] - auto seq_input = flat_input.slice(0, start, end); - - // Load / init conv state: [kernel_size - 1, channels] - int64_t state_idx = state_indices[b]; - torch::Tensor state; - if (conv_cache.defined() && conv_cache.numel() > 0) { - state = conv_cache[state_idx].clone(); - } else { - state = torch::zeros({kernel_size - 1, channels}, options); - } - - // Pad: [state(k-1,c); input(seq,c)] -> [seq + k - 1, channels] - auto padded = torch::cat({state, seq_input}, 0); - - // For each time step, depthwise conv: weight[c, ks] dot window[ks, c] - auto out = torch::zeros({seq_len, channels}, options); - for (int64_t t = 0; t < seq_len; ++t) { - auto window = padded.slice(0, t, t + kernel_size).to(torch::kFloat32); - // window: [kernel_size, channels] - // weight: [channels, kernel_size] - // out_t[c] = sum_i weight[c, i] * window[i, c] - auto out_t = torch::sum(weight * window.transpose(0, 1), {1}); - // out_t: [channels] - out[t] = out_t; - } - - // Store last state: last (kernel_size-1) tokens - if (conv_cache.defined() && conv_cache.numel() > 0) { - conv_cache[state_idx] = - padded.slice(0, seq_len, seq_len + kernel_size - 1).clone(); - } +// ------------------------------------------------------------------------- +// Backend dispatch — a single entry point per operator; picks torch_impl or +// kernel_impl based on the env flag. Callers stay free of if/else clutter. +// ------------------------------------------------------------------------- + +torch::Tensor causal_conv1d_prefill(const torch::Tensor& flat_input, + const torch::Tensor& conv_weight, + torch::Tensor& conv_cache, + const std::vector& cu_seqlens, + const std::vector& state_indices, + int32_t kernel_size) { + namespace gdn = qwen3_gdn; + return (use_torch_conv1d() ? gdn::torch_impl::causal_conv1d_prefill + : gdn::kernel_impl::causal_conv1d_prefill)( + flat_input.contiguous(), + conv_weight, + conv_cache, + cu_seqlens, + state_indices, + kernel_size, + /*activation=*/true); +} - if (activation) { - out = torch::silu(out); - } - outputs.push_back(out); - } - return torch::cat(outputs, 0).contiguous(); +torch::Tensor causal_conv1d_decode(const torch::Tensor& flat_input, + const torch::Tensor& conv_weight, + torch::Tensor& conv_cache, + const torch::Tensor& state_indices, + int32_t kernel_size) { + namespace gdn = qwen3_gdn; + return (use_torch_conv1d() + ? gdn::torch_impl::causal_conv1d_decode + : gdn::kernel_impl::causal_conv1d_decode)(flat_input, + conv_weight, + conv_cache, + state_indices, + kernel_size, + /*activation=*/true); } -// ----------------------------------------------------------------------- -// Torch causal conv1d update (decode): single-step -// conv_weight_2d: [out_channels, kernel_size] -// conv_cache: [num_sequences, kernel_size - 1, channels] -// ----------------------------------------------------------------------- -torch::Tensor torch_causal_conv1d_update(const torch::Tensor& flat_input, - const torch::Tensor& conv_weight_2d, - torch::Tensor& conv_cache, - const torch::Tensor& state_indices, - int32_t kernel_size, - bool activation) { - // flat_input: [batch, channels] - // conv_weight_2d: [kernel_size, channels] (from conv1d_->weight()) - int64_t batch = flat_input.size(0); - int64_t channels = flat_input.size(1); - auto options = flat_input.options(); - auto weight = conv_weight_2d.transpose(0, 1).contiguous(); - // weight: [channels, kernel_size] - auto weight_f32 = weight.to(torch::kFloat32); - auto outputs = torch::empty({batch, channels}, options); - - for (int64_t b = 0; b < batch; ++b) { - int64_t idx = state_indices[b].item(); - // cache format: [kernel_size - 1, channels] - auto state_t = conv_cache[idx].to(torch::kFloat32); - auto x_t = flat_input[b].to(torch::kFloat32); - // frame: [kernel_size, channels] - auto frame = torch::cat({state_t, x_t.unsqueeze(0)}, 0); - // depthwise conv: weight[c, ks] dot frame[ks, c] -> [c] - auto out_t = torch::sum(weight_f32 * frame.transpose(0, 1), {1}); - outputs[b] = out_t.to(options.dtype()); - // Update state: shift, keep last (kernel_size-1) tokens - conv_cache[idx] = frame.slice(0, 1, kernel_size).to(options.dtype()); +// Prefill selects between torch_recurrent (slow, exact) and aiter chunked +// (fast). Decode always uses torch_recurrent (T=1, no chunk speedup possible). +std::tuple gated_delta_rule_prefill( + const torch::Tensor& cat_q, + const torch::Tensor& cat_k, + const torch::Tensor& cat_v, + const torch::Tensor& cat_g, + const torch::Tensor& cat_beta, + const torch::Tensor& initial_state, + const std::vector& cu_seqlens) { + namespace gdn = qwen3_gdn; + if (use_torch_recurrent_prefill()) { + return gdn::torch_impl::gated_delta_rule( + cat_q, cat_k, cat_v, cat_g, cat_beta, initial_state); } + return gdn::kernel_impl::gated_delta_rule_prefill( + cat_q, + cat_k, + cat_v, + cat_g, + cat_beta, + std::optional(initial_state), + cu_seqlens, + /*output_final_state=*/true, + /*use_qk_l2norm_in_kernel=*/true); +} - if (activation) { - outputs = torch::silu(outputs); - } - return outputs; +std::tuple gated_delta_rule_decode( + const torch::Tensor& q, + const torch::Tensor& k, + const torch::Tensor& v, + const torch::Tensor& g, + const torch::Tensor& beta, + const torch::Tensor& init_state) { + return qwen3_gdn::torch_impl::gated_delta_rule(q, k, v, g, beta, init_state); } -// ----------------------------------------------------------------------- -// Torch fused_qkvzba_split_reshape_cat replacement -// ----------------------------------------------------------------------- +// ------------------------------------------------------------------------- +// Small helpers used only by the layer's forward orchestration. +// ------------------------------------------------------------------------- + +// Split the fused QKVZ/BA projections back into their four component tensors. std::tuple -torch_fused_qkvzba_split(const torch::Tensor& qkvz_flat, - const torch::Tensor& ba_flat, - int64_t num_k_heads_local, - int64_t num_v_heads_local, - int64_t head_k_dim, - int64_t head_v_dim) { - int64_t total_tokens = qkvz_flat.size(0); - int64_t k_size = num_k_heads_local * head_k_dim; - int64_t v_size = num_v_heads_local * head_v_dim; - - // qkvz_flat: [total_tokens, 2*k_size + 2*v_size] - // = [total_tokens, q(k_size) + k(k_size) + v(v_size) + z(v_size)] +split_qkvz_ba(const torch::Tensor& qkvz_flat, + const torch::Tensor& ba_flat, + int64_t num_k_heads_local, + int64_t num_v_heads_local, + int64_t head_k_dim, + int64_t head_v_dim) { + const int64_t k_size = num_k_heads_local * head_k_dim; + const int64_t v_size = num_v_heads_local * head_v_dim; + auto qkvz_split = torch::split(qkvz_flat, {k_size, k_size, v_size, v_size}, 1); - auto q = qkvz_split[0]; - auto k = qkvz_split[1]; - auto v = qkvz_split[2]; - auto z = qkvz_split[3]; - - // ba_flat: [total_tokens, 2 * num_v_heads_local] - // = [total_tokens, b(num_v_heads) + a(num_v_heads)] - int64_t num_v = num_v_heads_local; - auto ba_split = torch::split(ba_flat, {num_v, num_v}, 1); - auto b = ba_split[0]; - auto a = ba_split[1]; - - // mixed_qkv = cat([q, k, v], -1) → [total_tokens, 2*k_size + v_size] + const auto& q = qkvz_split[0]; + const auto& k = qkvz_split[1]; + const auto& v = qkvz_split[2]; + const auto& z = qkvz_split[3]; + auto ba_split = + torch::split(ba_flat, {num_v_heads_local, num_v_heads_local}, 1); + const auto& b = ba_split[0]; + const auto& a = ba_split[1]; auto mixed_qkv = torch::cat({q, k, v}, 1).contiguous(); - return std::make_tuple(mixed_qkv, z, b, a); } -// ========================================================================= -// Gated Delta Rule torch implementations (ported from NPU) -// ========================================================================= - -std::tuple torch_recurrent_gated_delta_rule( - torch::Tensor query, - torch::Tensor key, - torch::Tensor value, - torch::Tensor g, - torch::Tensor beta, - std::optional initial_state, - bool output_final_state = true, - bool use_qk_l2norm_in_kernel = true) { - auto initial_dtype = query.dtype(); - - if (use_qk_l2norm_in_kernel) { - query = l2norm(query, -1, 1e-6); - key = l2norm(key, -1, 1e-6); - } - - auto to_float32_and_transpose = [](torch::Tensor x) { - return x.transpose(1, 2).contiguous().to(torch::kFloat32); - }; - query = to_float32_and_transpose(query); - key = to_float32_and_transpose(key); - value = to_float32_and_transpose(value); - beta = to_float32_and_transpose(beta); - g = to_float32_and_transpose(g); - const int64_t value_num_heads = value.size(1); - query = repeat_tensor_heads(query, value_num_heads, 1); - key = repeat_tensor_heads(key, value_num_heads, 1); - - int64_t batch_size = key.size(0); - int64_t num_heads = key.size(1); - int64_t sequence_length = key.size(2); - int64_t k_head_dim = key.size(3); - int64_t v_head_dim = value.size(3); - - float scale_val = 1.0 / std::sqrt(static_cast(query.size(-1))); - torch::Tensor scale = torch::tensor(scale_val, query.options()); - query = query * scale; - torch::Tensor core_attn_out = torch::zeros( - {batch_size, num_heads, sequence_length, v_head_dim}, - torch::TensorOptions().dtype(torch::kFloat32).device(value.device())); - torch::Tensor last_recurrent_state; - if (!initial_state.has_value()) { - last_recurrent_state = torch::zeros( - {batch_size, num_heads, k_head_dim, v_head_dim}, - torch::TensorOptions().dtype(torch::kFloat32).device(value.device())); - } else { - last_recurrent_state = - initial_state.value().to(value.device(), torch::kFloat32); - } - - for (int64_t i = 0; i < sequence_length; ++i) { - torch::Tensor q_t = query.select(2, i); - torch::Tensor k_t = key.select(2, i); - torch::Tensor v_t = value.select(2, i); - torch::Tensor g_t = g.select(2, i).exp().unsqueeze(-1).unsqueeze(-1); - torch::Tensor beta_t = beta.select(2, i).unsqueeze(-1); - last_recurrent_state = last_recurrent_state * g_t; - torch::Tensor kv_mem = - torch::sum(last_recurrent_state * k_t.unsqueeze(-1), -2); - torch::Tensor delta = (v_t - kv_mem) * beta_t; - last_recurrent_state = - last_recurrent_state + k_t.unsqueeze(-1) * delta.unsqueeze(-2); - core_attn_out.select(2, i) = - torch::sum(last_recurrent_state * q_t.unsqueeze(-1), -2); - } - - core_attn_out = core_attn_out.transpose(1, 2).contiguous().to(initial_dtype); - return std::make_tuple(core_attn_out, last_recurrent_state); -} - int64_t get_checkpoint_stride(const torch::Tensor& conv_cache, const torch::Tensor& ssm_cache) { if (!conv_cache.defined() || !ssm_cache.defined() || @@ -300,6 +172,37 @@ torch::Tensor build_linear_state_base_indices( return logical_state_indices * checkpoint_stride; } +// cu_seqlens with only the positive-length sequences retained. Zero-length +// padding entries in q_seq_lens_vec must not become separate cu_seqlens rows +// because they don't consume linear state slots. +std::vector build_real_cu_seqlens( + const std::vector& q_seq_lens_vec) { + std::vector cu; + cu.reserve(q_seq_lens_vec.size() + 1); + cu.push_back(0); + for (const int32_t vlen : q_seq_lens_vec) { + if (vlen > 0) { + cu.push_back(cu.back() + vlen); + } + } + return cu; +} + +// Extract cu_seqlens from host-side lens (or fall back to device tensor). +std::vector extract_cu_seqlens( + const AttentionMetadata& attn_metadata) { + std::vector cu; + if (!attn_metadata.q_seq_lens_vec.empty()) { + cu.assign(attn_metadata.q_seq_lens_vec.begin(), + attn_metadata.q_seq_lens_vec.end()); + } else { + auto cpu_lens = attn_metadata.q_seq_lens.cpu(); + auto* ptr = cpu_lens.data_ptr(); + cu.assign(ptr, ptr + cpu_lens.numel()); + } + return cu; +} + } // namespace // ========================================================================= @@ -402,22 +305,20 @@ torch::Tensor Qwen3GatedDeltaNetBaseImpl::forward( const int64_t original_num_tokens = hidden_states.size(0); auto [qkvz_padded, ba_padded] = project_padded_inputs(hidden_states, attn_metadata); - int64_t batch_size = qkvz_padded.size(0); - int64_t seq_len = qkvz_padded.size(1); + const int64_t batch_size = qkvz_padded.size(0); + const int64_t seq_len = qkvz_padded.size(1); - // Flatten and split QKVZ/BA via torch fallback auto qkvz_flat = qkvz_padded.reshape({batch_size * seq_len, qkvz_padded.size(-1)}); auto ba_flat = ba_padded.reshape({batch_size * seq_len, ba_padded.size(-1)}); torch::Tensor mixed_qkv, z, b, a; - std::tie(mixed_qkv, z, b, a) = - torch_fused_qkvzba_split(qkvz_flat, - ba_flat, - num_k_heads_ / tp_size_, - num_v_heads_ / tp_size_, - head_k_dim_, - head_v_dim_); + std::tie(mixed_qkv, z, b, a) = split_qkvz_ba(qkvz_flat, + ba_flat, + num_k_heads_ / tp_size_, + num_v_heads_ / tp_size_, + head_k_dim_, + head_v_dim_); mixed_qkv = mixed_qkv.reshape({batch_size, seq_len, mixed_qkv.size(-1)}); z = z.reshape({batch_size, seq_len, num_v_heads_ / tp_size_, head_v_dim_}); @@ -426,126 +327,101 @@ torch::Tensor Qwen3GatedDeltaNetBaseImpl::forward( torch::Tensor conv_cache = kv_cache.get_conv_cache(); torch::Tensor ssm_cache = kv_cache.get_ssm_cache(); - torch::Device device = mixed_qkv.device(); torch::Tensor conv_weight = conv1d_->weight(); torch::Tensor logical_state_indices = - get_linear_state_indices(input_params, device); + get_linear_state_indices(input_params, mixed_qkv.device()); const int64_t checkpoint_stride = get_checkpoint_stride(conv_cache, ssm_cache); torch::Tensor linear_state_base_indices = build_linear_state_base_indices(logical_state_indices, checkpoint_stride); - const bool is_any_prefill = + const bool is_prefill = attn_metadata.is_prefill || attn_metadata.is_chunked_prefill; // ---- Causal Conv1d ---- - if (is_any_prefill) { - std::vector cu_seqlens_vec; - if (!attn_metadata.q_seq_lens_vec.empty()) { - cu_seqlens_vec.push_back(0); - for (size_t i = 0; i < attn_metadata.q_seq_lens_vec.size(); ++i) { - cu_seqlens_vec.push_back(cu_seqlens_vec.back() + - attn_metadata.q_seq_lens_vec[i]); - } - } else { - auto cpu_lens = attn_metadata.q_seq_lens.cpu(); - auto* ptr = cpu_lens.data_ptr(); - cu_seqlens_vec.assign(ptr, ptr + cpu_lens.numel()); - } - + if (is_prefill) { + auto cu_seqlens_vec = extract_cu_seqlens(attn_metadata); auto conv_input = reshape_qkvz_unpad(attn_metadata, mixed_qkv); - std::vector state_indices_vec( input_params.embedding.linear_state_ids.begin(), input_params.embedding.linear_state_ids.end()); - - mixed_qkv = torch_causal_conv1d(conv_input.contiguous(), - conv_weight, - conv_cache, - cu_seqlens_vec, - state_indices_vec, - conv_kernel_size_, - /*activation=*/true); - + mixed_qkv = causal_conv1d_prefill(conv_input, + conv_weight, + conv_cache, + cu_seqlens_vec, + state_indices_vec, + conv_kernel_size_); mixed_qkv = reshape_qkvz_with_pad(attn_metadata, mixed_qkv); mixed_qkv = mixed_qkv.transpose(1, 2); } else { - // Decode: single-step conv update auto flat_input = mixed_qkv.reshape({batch_size * seq_len, mixed_qkv.size(-1)}); - auto updated = torch_causal_conv1d_update(flat_input, - conv_weight, - conv_cache, - logical_state_indices, - conv_kernel_size_, - /*activation=*/true); + auto updated = causal_conv1d_decode(flat_input, + conv_weight, + conv_cache, + logical_state_indices, + conv_kernel_size_); mixed_qkv = updated.reshape({batch_size, seq_len, -1}).transpose(1, 2); } - // ---- Gating (torch path: no kernel required) ---- + // ---- Gating ---- + const auto beta = torch::sigmoid(b); torch::Tensor g; - torch::Tensor beta; { - beta = torch::sigmoid(b); auto A_log_exp = A_log_.exp(); - auto a_float = a.to(torch::kFloat32); - auto a_plus_dt = a_float + dt_bias_; auto softplus_out = torch::nn::functional::softplus( - a_plus_dt, + a.to(torch::kFloat32) + dt_bias_, torch::nn::functional::SoftplusFuncOptions().beta(1.0).threshold(20.0)); - g = -A_log_exp * softplus_out; - g = g.to(a.dtype()).contiguous(); + g = (-A_log_exp * softplus_out).to(a.dtype()).contiguous(); } auto [processed_q, processed_k, processed_v] = process_mixed_qkv(mixed_qkv); // ---- Gated Delta Rule ---- torch::Tensor core_attn_out; - if (is_any_prefill) { + if (is_prefill) { CHECK_GE(attn_metadata.q_seq_lens_vec.size(), static_cast(batch_size)) << "q_seq_lens_vec must be populated for Qwen3.5 prefill."; - // Pack valid tokens per batch, run chunked gated delta rule - std::vector packed_q, packed_k, packed_v, packed_g, - packed_beta; - packed_q.reserve(batch_size); - packed_k.reserve(batch_size); - packed_v.reserve(batch_size); - packed_g.reserve(batch_size); - packed_beta.reserve(batch_size); - + std::vector pq, pk, pv, pg, pbeta; + pq.reserve(batch_size); + pk.reserve(batch_size); + pv.reserve(batch_size); + pg.reserve(batch_size); + pbeta.reserve(batch_size); for (int64_t bidx = 0; bidx < batch_size; ++bidx) { - int64_t vlen = attn_metadata.q_seq_lens_vec[bidx]; - packed_q.push_back(processed_q[bidx].narrow(0, 0, vlen)); - packed_k.push_back(processed_k[bidx].narrow(0, 0, vlen)); - packed_v.push_back(processed_v[bidx].narrow(0, 0, vlen)); - packed_g.push_back(g[bidx].narrow(0, 0, vlen)); - packed_beta.push_back(beta[bidx].narrow(0, 0, vlen)); + const int64_t vlen = attn_metadata.q_seq_lens_vec[bidx]; + pq.push_back(processed_q[bidx].narrow(0, 0, vlen)); + pk.push_back(processed_k[bidx].narrow(0, 0, vlen)); + pv.push_back(processed_v[bidx].narrow(0, 0, vlen)); + pg.push_back(g[bidx].narrow(0, 0, vlen)); + pbeta.push_back(beta[bidx].narrow(0, 0, vlen)); } + auto cat_q = torch::cat(pq, 0).unsqueeze(0); + auto cat_k = torch::cat(pk, 0).unsqueeze(0); + auto cat_v = torch::cat(pv, 0).unsqueeze(0); + auto cat_g = torch::cat(pg, 0).unsqueeze(0); + auto cat_beta = torch::cat(pbeta, 0).unsqueeze(0); - auto cat_q = torch::cat(packed_q, 0).unsqueeze(0); - auto cat_k = torch::cat(packed_k, 0).unsqueeze(0); - auto cat_v = torch::cat(packed_v, 0).unsqueeze(0); - auto cat_g = torch::cat(packed_g, 0).unsqueeze(0); - auto cat_beta = torch::cat(packed_beta, 0).unsqueeze(0); - - // Get initial state from ssm_cache + // ssm_cache stores state as [N, H, V, K] (non-fla) or [N, H, K, V] (fla); + // gated_delta_rule expects [N, H, K, V] fla layout. auto initial_state = torch::index_select(ssm_cache, 0, linear_state_base_indices); if (!use_fla_ssm_state_layout()) { initial_state = initial_state.transpose(-1, -2).contiguous(); } + auto real_cu = build_real_cu_seqlens(attn_metadata.q_seq_lens_vec); torch::Tensor last_state; - std::tie(core_attn_out, last_state) = torch_recurrent_gated_delta_rule( - cat_q, cat_k, cat_v, cat_g, cat_beta, initial_state); + std::tie(core_attn_out, last_state) = gated_delta_rule_prefill( + cat_q, cat_k, cat_v, cat_g, cat_beta, initial_state, real_cu); - // Scatter back to per-batch output + // Scatter back to per-batch output. core_attn_out = core_attn_out.squeeze(0); auto final_out = torch::zeros_like(processed_v); int64_t offset = 0; for (int64_t bidx = 0; bidx < batch_size; ++bidx) { - int64_t vlen = attn_metadata.q_seq_lens_vec[bidx]; + const int64_t vlen = attn_metadata.q_seq_lens_vec[bidx]; final_out[bidx] .narrow(0, 0, vlen) .copy_(core_attn_out.narrow(0, offset, vlen)); @@ -553,13 +429,11 @@ torch::Tensor Qwen3GatedDeltaNetBaseImpl::forward( } core_attn_out = final_out; - // Store state auto state_to_store = use_fla_ssm_state_layout() ? last_state : last_state.transpose(-1, -2); ssm_cache.index_put_({linear_state_base_indices}, state_to_store.to(ssm_cache.dtype())); } else { - // Decode: recurrent step torch::Tensor init_state; if (ssm_cache.defined() && ssm_cache.numel() > 0) { init_state = torch::index_select(ssm_cache, 0, linear_state_base_indices); @@ -568,7 +442,7 @@ torch::Tensor Qwen3GatedDeltaNetBaseImpl::forward( } } torch::Tensor last_state; - std::tie(core_attn_out, last_state) = torch_recurrent_gated_delta_rule( + std::tie(core_attn_out, last_state) = gated_delta_rule_decode( processed_q, processed_k, processed_v, g, beta, init_state); auto state_to_store = @@ -582,14 +456,12 @@ torch::Tensor Qwen3GatedDeltaNetBaseImpl::forward( auto core_attn_out_reshaped = core_attn_out.reshape({-1, core_attn_out.size(-1)}); auto norm_out = norm_->forward(core_attn_out_reshaped, z_reshaped); - auto z_shape_og = z.sizes().vec(); - norm_out = norm_out.view(z_shape_og); + norm_out = norm_out.view(z.sizes().vec()); norm_out = norm_out.reshape({-1, norm_out.size(2), norm_out.size(3)}); auto rearranged_norm = norm_out.reshape({norm_out.size(0), norm_out.size(1) * norm_out.size(2)}); rearranged_norm = reshape_qkvz_unpad(attn_metadata, rearranged_norm); - if (rearranged_norm.size(0) > original_num_tokens) { rearranged_norm = rearranged_norm.slice(0, 0, original_num_tokens).contiguous(); @@ -605,20 +477,19 @@ torch::Tensor Qwen3GatedDeltaNetBaseImpl::reshape_qkvz_unpad( if (!has_padded_queries) { return padded_qkvz; } - std::vector valid_batches; const bool has_host_lens = !attn_metadata.q_seq_lens_vec.empty(); - int64_t bs = has_host_lens - ? static_cast(attn_metadata.q_seq_lens_vec.size()) - : attn_metadata.q_seq_lens.size(0); - valid_batches.reserve(bs); - int64_t max_len = attn_metadata.max_query_len; - const auto& ori_seq_lens = attn_metadata.q_seq_lens; + const int64_t bs = + has_host_lens ? static_cast(attn_metadata.q_seq_lens_vec.size()) + : attn_metadata.q_seq_lens.size(0); + const int64_t max_len = attn_metadata.max_query_len; auto reshaped_qkvz = padded_qkvz.reshape({bs, max_len, -1}); + std::vector valid_batches; + valid_batches.reserve(bs); for (int64_t b = 0; b < bs; ++b) { - int64_t ori_len = has_host_lens ? attn_metadata.q_seq_lens_vec[b] - : ori_seq_lens[b].template item(); - torch::Tensor valid_batch = reshaped_qkvz[b].slice(0, 0, ori_len); - valid_batches.emplace_back(valid_batch); + const int64_t ori_len = + has_host_lens ? attn_metadata.q_seq_lens_vec[b] + : attn_metadata.q_seq_lens[b].template item(); + valid_batches.emplace_back(reshaped_qkvz[b].slice(0, 0, ori_len)); } if (valid_batches.size() == 1) { return valid_batches[0].contiguous(); @@ -640,10 +511,10 @@ torch::Tensor Qwen3GatedDeltaNetBaseImpl::reshape_qkvz_with_pad( const AttentionMetadata& attn_metadata, const torch::Tensor& qkvz) const { const bool has_host_lens = !attn_metadata.q_seq_lens_vec.empty(); - int64_t bs = has_host_lens - ? static_cast(attn_metadata.q_seq_lens_vec.size()) - : attn_metadata.q_seq_lens.size(0); - int64_t max_len = attn_metadata.max_query_len; + const int64_t bs = + has_host_lens ? static_cast(attn_metadata.q_seq_lens_vec.size()) + : attn_metadata.q_seq_lens.size(0); + const int64_t max_len = attn_metadata.max_query_len; const bool need_padding = attn_metadata.is_prefill || attn_metadata.is_chunked_prefill; if (!need_padding) { @@ -653,9 +524,9 @@ torch::Tensor Qwen3GatedDeltaNetBaseImpl::reshape_qkvz_with_pad( batches.reserve(bs); int64_t idx = 0; for (int64_t b = 0; b < bs; ++b) { - int64_t cur_len = has_host_lens ? attn_metadata.q_seq_lens_vec[b] : 0; + const int64_t cur_len = has_host_lens ? attn_metadata.q_seq_lens_vec[b] : 0; torch::Tensor batch = qkvz.slice(0, idx, idx + cur_len).contiguous(); - idx = idx + cur_len; + idx += cur_len; if (batch.size(0) != max_len) { batch = batch.size(0) > max_len ? batch.slice(0, 0, max_len).contiguous() @@ -673,19 +544,16 @@ torch::Tensor Qwen3GatedDeltaNetBaseImpl::reshape_qkvz_with_pad( std::tuple Qwen3GatedDeltaNetBaseImpl::process_mixed_qkv(torch::Tensor& mixed_qkv) const { mixed_qkv = mixed_qkv.transpose(1, 2); - int64_t batch_size = mixed_qkv.size(0); - int64_t seq_len = mixed_qkv.size(1); - std::vector split_sizes = { + const int64_t batch_size = mixed_qkv.size(0); + const int64_t seq_len = mixed_qkv.size(1); + const std::vector split_sizes = { k_size_ / tp_size_, k_size_ / tp_size_, v_size_ / tp_size_}; auto processed_qkv = torch::split(mixed_qkv, split_sizes, 2); - auto processed_q = processed_qkv[0]; - auto processed_k = processed_qkv[1]; - auto processed_v = processed_qkv[2]; - processed_q = processed_q.reshape( + auto processed_q = processed_qkv[0].reshape( {batch_size, seq_len, num_k_heads_ / tp_size_, head_k_dim_}); - processed_k = processed_k.reshape( + auto processed_k = processed_qkv[1].reshape( {batch_size, seq_len, num_k_heads_ / tp_size_, head_k_dim_}); - processed_v = processed_v.reshape( + auto processed_v = processed_qkv[2].reshape( {batch_size, seq_len, num_v_heads_ / tp_size_, head_v_dim_}); return std::make_tuple(processed_q, processed_k, processed_v); } @@ -699,7 +567,6 @@ torch::Tensor Qwen3_5GatedDeltaNetImpl::forward( const AttentionMetadata& attn_metadata, KVCache& kv_cache, const ModelInputParams& input_params) { - // Direct projection (no merge/split), matching vLLM reference. auto mixed_qkv_flat = in_proj_qkv_->forward(hidden_states); auto z_flat = in_proj_z_->forward(hidden_states); auto b_flat = in_proj_b_->forward(hidden_states); @@ -731,7 +598,6 @@ torch::Tensor Qwen3_5GatedDeltaNetImpl::forward( a = a_flat.reshape({batch_size, seq_len, -1}); } - // z: reshape to [batch, seq, num_v_heads/tp, head_v_dim] z = z.reshape({batch_size, seq_len, num_v_heads_ / tp_size_, head_v_dim_}); // ---- Causal Conv1d ---- @@ -739,50 +605,37 @@ torch::Tensor Qwen3_5GatedDeltaNetImpl::forward( torch::Tensor conv_weight = conv1d_->weight(); if (is_prefill) { - std::vector cu_seqlens_vec; - if (!attn_metadata.q_seq_lens_vec.empty()) { - cu_seqlens_vec.push_back(0); - for (size_t i = 0; i < attn_metadata.q_seq_lens_vec.size(); ++i) - cu_seqlens_vec.push_back(cu_seqlens_vec.back() + - attn_metadata.q_seq_lens_vec[i]); - } else { - auto cpu_lens = attn_metadata.q_seq_lens.cpu(); - auto* ptr = cpu_lens.data_ptr(); - cu_seqlens_vec.assign(ptr, ptr + cpu_lens.numel()); - } + auto cu_seqlens_vec = extract_cu_seqlens(attn_metadata); auto conv_input = reshape_qkvz_unpad(attn_metadata, mixed_qkv); std::vector state_indices( input_params.embedding.linear_state_ids.begin(), input_params.embedding.linear_state_ids.end()); - mixed_qkv = torch_causal_conv1d(conv_input.contiguous(), - conv_weight, - conv_cache, - cu_seqlens_vec, - state_indices, - conv_kernel_size_, - /*activation=*/true); + mixed_qkv = causal_conv1d_prefill(conv_input, + conv_weight, + conv_cache, + cu_seqlens_vec, + state_indices, + conv_kernel_size_); mixed_qkv = reshape_qkvz_with_pad(attn_metadata, mixed_qkv); } else { auto flat_input = mixed_qkv.reshape({batch_size * seq_len, -1}); auto state_indices = get_linear_state_indices(input_params, mixed_qkv.device()); - auto updated = torch_causal_conv1d_update(flat_input, - conv_weight, - conv_cache, - state_indices, - conv_kernel_size_, - /*activation=*/true); - mixed_qkv = updated.reshape({batch_size, seq_len, -1}); + mixed_qkv = causal_conv1d_decode(flat_input, + conv_weight, + conv_cache, + state_indices, + conv_kernel_size_) + .reshape({batch_size, seq_len, -1}); } mixed_qkv = mixed_qkv.transpose(1, 2); auto [q, k, v] = process_mixed_qkv(mixed_qkv); // ---- Gating ---- - auto beta = torch::sigmoid(b); - auto a_f32 = a.to(torch::kFloat32); + const auto beta = torch::sigmoid(b); auto softplus_out = torch::nn::functional::softplus( - a_f32 + dt_bias_, + a.to(torch::kFloat32) + dt_bias_, torch::nn::functional::SoftplusFuncOptions().beta(1.0).threshold(20.0)); auto g = (-A_log_.exp() * softplus_out).to(b.dtype()).contiguous(); @@ -793,32 +646,33 @@ torch::Tensor Qwen3_5GatedDeltaNetImpl::forward( torch::Tensor core_attn_out; if (is_prefill) { - std::vector pq, pk, pv, pg, pb; + std::vector pq, pk, pv, pg, pbeta; for (int64_t bi = 0; bi < batch_size; ++bi) { - int64_t vlen = attn_metadata.q_seq_lens_vec[bi]; + const int64_t vlen = attn_metadata.q_seq_lens_vec[bi]; pq.push_back(q[bi].narrow(0, 0, vlen)); pk.push_back(k[bi].narrow(0, 0, vlen)); pv.push_back(v[bi].narrow(0, 0, vlen)); pg.push_back(g[bi].narrow(0, 0, vlen)); - pb.push_back(beta[bi].narrow(0, 0, vlen)); + pbeta.push_back(beta[bi].narrow(0, 0, vlen)); } auto cat_q = torch::cat(pq, 0).unsqueeze(0); auto cat_k = torch::cat(pk, 0).unsqueeze(0); auto cat_v = torch::cat(pv, 0).unsqueeze(0); auto cat_g = torch::cat(pg, 0).unsqueeze(0); - auto cat_beta = torch::cat(pb, 0).unsqueeze(0); + auto cat_beta = torch::cat(pbeta, 0).unsqueeze(0); auto init_state = torch::index_select(ssm_cache, 0, linear_state_base); + auto real_cu = build_real_cu_seqlens(attn_metadata.q_seq_lens_vec); torch::Tensor last_state; - std::tie(core_attn_out, last_state) = torch_recurrent_gated_delta_rule( - cat_q, cat_k, cat_v, cat_g, cat_beta, init_state); + std::tie(core_attn_out, last_state) = gated_delta_rule_prefill( + cat_q, cat_k, cat_v, cat_g, cat_beta, init_state, real_cu); core_attn_out = core_attn_out.squeeze(0); auto out3d = torch::zeros_like(v); int64_t off = 0; for (int64_t bi = 0; bi < batch_size; ++bi) { - int64_t vlen = attn_metadata.q_seq_lens_vec[bi]; + const int64_t vlen = attn_metadata.q_seq_lens_vec[bi]; out3d[bi].narrow(0, 0, vlen).copy_(core_attn_out.narrow(0, off, vlen)); off += vlen; } @@ -828,7 +682,7 @@ torch::Tensor Qwen3_5GatedDeltaNetImpl::forward( auto init_state = torch::index_select(ssm_cache, 0, linear_state_base); torch::Tensor last_state; std::tie(core_attn_out, last_state) = - torch_recurrent_gated_delta_rule(q, k, v, g, beta, init_state); + gated_delta_rule_decode(q, k, v, g, beta, init_state); ssm_cache.index_put_({linear_state_base}, last_state.to(ssm_cache.dtype())); } diff --git a/xllm/core/layers/dcu/qwen3_gated_delta_net_helpers.h b/xllm/core/layers/dcu/qwen3_gated_delta_net_helpers.h new file mode 100644 index 0000000000..af44f5e53b --- /dev/null +++ b/xllm/core/layers/dcu/qwen3_gated_delta_net_helpers.h @@ -0,0 +1,59 @@ +/* Copyright 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 + +#include + +namespace xllm { +namespace layer { +namespace qwen3_gdn { + +// Small tensor helpers shared by both torch and optimized-kernel backends. + +// L2-normalize along `dim` with an epsilon guard. +inline torch::Tensor l2norm(const torch::Tensor& x, + int64_t dim, + double eps = 1e-6) { + auto norm = torch::sqrt(torch::sum(torch::square(x), dim, true) + eps); + return x / norm; +} + +// Broadcast GQA: replicate `tensor` along `head_dim` from current head count to +// `target_heads` (target must be a multiple of current). +inline torch::Tensor repeat_tensor_heads(const torch::Tensor& tensor, + int64_t target_heads, + int64_t head_dim) { + const int64_t current_heads = tensor.size(head_dim); + if (current_heads == target_heads) { + return tensor; + } + CHECK_GT(current_heads, 0) << "current heads must be positive"; + CHECK_EQ(target_heads % current_heads, 0) + << "target heads must be divisible by current heads"; + + const int64_t repeats = target_heads / current_heads; + std::vector view_shape = tensor.sizes().vec(); + view_shape.insert(view_shape.begin() + head_dim + 1, 1); + std::vector expand_shape = view_shape; + expand_shape[head_dim + 1] = repeats; + std::vector output_shape = tensor.sizes().vec(); + output_shape[head_dim] = target_heads; + return tensor.unsqueeze(head_dim + 1) + .expand(expand_shape) + .reshape(output_shape) + .contiguous(); +} + +} // namespace qwen3_gdn +} // namespace layer +} // namespace xllm diff --git a/xllm/core/layers/dcu/qwen3_gated_delta_net_impls.h b/xllm/core/layers/dcu/qwen3_gated_delta_net_impls.h new file mode 100644 index 0000000000..88650b6921 --- /dev/null +++ b/xllm/core/layers/dcu/qwen3_gated_delta_net_impls.h @@ -0,0 +1,146 @@ +/* Copyright 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. +==============================================================================*/ + +// Common backend interface for Qwen3 gated-delta-net operators on DCU. +// +// Two implementations live behind identical function signatures: +// * torch_impl:: — pure-torch fallback (portable, no external kernels). +// * kernel_impl:: — optimized HIP kernels (aiter fla + causal-conv1d). +// +// The layer's forward() chooses between them via runtime env flags +// (XLLM_DCU_USE_TORCH_CONV1D=1, XLLM_DCU_USE_TORCH_RECURRENT=1). All shared +// tensor-shape conventions are documented once, here. + +#pragma once + +#include + +#include +#include +#include + +namespace xllm { +namespace layer { +namespace qwen3_gdn { + +// --------------------------------------------------------------------------- +// causal_conv1d +// --------------------------------------------------------------------------- +// Prefill (varlen, packed): +// flat_input : [T_total, C] contiguous, same dtype as conv_cache/weight. +// conv_weight_2d : [C, K] contiguous. +// conv_cache : [num_slots, K-1, C] in place — one row per real seq +// (matching state_indices below). +// cu_seqlens : N+1 int64 prefix sums (only real seqs; caller filters +// zero-length padding rows). +// state_indices : N int64 slot ids in conv_cache for the same N real seqs. +// kernel_size : K +// activation : true => silu. +// Returns [T_total, C] output with silu applied. +using CausalConv1dPrefillFn = + torch::Tensor (*)(const torch::Tensor& flat_input, + const torch::Tensor& conv_weight_2d, + torch::Tensor& conv_cache, + const std::vector& cu_seqlens, + const std::vector& state_indices, + int32_t kernel_size, + bool activation); + +// Decode (per-token): +// flat_input : [B, C] contiguous. +// conv_weight_2d: [C, K]. +// conv_cache : [num_slots, K-1, C] updated in place. +// state_indices : [B] int32/int64 device tensor of slot ids. +// Returns [B, C] output with silu applied. +using CausalConv1dDecodeFn = + torch::Tensor (*)(const torch::Tensor& flat_input, + const torch::Tensor& conv_weight_2d, + torch::Tensor& conv_cache, + const torch::Tensor& state_indices, + int32_t kernel_size, + bool activation); + +// --------------------------------------------------------------------------- +// gated_delta_rule +// --------------------------------------------------------------------------- +// Shared for both prefill (long T) and decode (T=1): +// q, k : [1, T, Hk, K] bf16, before optional l2norm. +// v : [1, T, Hv, V] bf16. +// g : [1, T, Hv] log-gate (bf16 or fp32 accepted). +// beta : [1, T, Hv]. +// initial_state (optional): [N, Hv, K, V] fp32 (fla layout). +// Returns: +// core_attn_out : [1, T, Hv, V] in q's dtype. +// last_state : [N, Hv, K, V] fp32 (fla layout) — final state per seq. + +namespace torch_impl { +torch::Tensor causal_conv1d_prefill(const torch::Tensor& flat_input, + const torch::Tensor& conv_weight_2d, + torch::Tensor& conv_cache, + const std::vector& cu_seqlens, + const std::vector& state_indices, + int32_t kernel_size, + bool activation); + +torch::Tensor causal_conv1d_decode(const torch::Tensor& flat_input, + const torch::Tensor& conv_weight_2d, + torch::Tensor& conv_cache, + const torch::Tensor& state_indices, + int32_t kernel_size, + bool activation); + +// Naive per-token recurrent; used for both prefill and decode. +std::tuple gated_delta_rule( + torch::Tensor query, + torch::Tensor key, + torch::Tensor value, + torch::Tensor g, + torch::Tensor beta, + std::optional initial_state, + bool output_final_state = true, + bool use_qk_l2norm_in_kernel = true); +} // namespace torch_impl + +namespace kernel_impl { +torch::Tensor causal_conv1d_prefill(const torch::Tensor& flat_input, + const torch::Tensor& conv_weight_2d, + torch::Tensor& conv_cache, + const std::vector& cu_seqlens, + const std::vector& state_indices, + int32_t kernel_size, + bool activation); + +torch::Tensor causal_conv1d_decode(const torch::Tensor& flat_input, + const torch::Tensor& conv_weight_2d, + torch::Tensor& conv_cache, + const torch::Tensor& state_indices, + int32_t kernel_size, + bool activation); + +// Chunked WY decomposition + aiter h-recurrence + chunk_fwd_o. Prefill only — +// decode always uses torch_impl::gated_delta_rule since its per-token cost is +// already the bottleneck-free case. +std::tuple gated_delta_rule_prefill( + torch::Tensor query, + torch::Tensor key, + torch::Tensor value, + torch::Tensor g, + torch::Tensor beta, + std::optional initial_state, + const std::vector& cu_seqlens, + bool output_final_state = true, + bool use_qk_l2norm_in_kernel = true); +} // namespace kernel_impl + +} // namespace qwen3_gdn +} // namespace layer +} // namespace xllm diff --git a/xllm/core/layers/dcu/qwen3_gated_delta_net_kernel_impl.cpp b/xllm/core/layers/dcu/qwen3_gated_delta_net_kernel_impl.cpp new file mode 100644 index 0000000000..84a5239539 --- /dev/null +++ b/xllm/core/layers/dcu/qwen3_gated_delta_net_kernel_impl.cpp @@ -0,0 +1,416 @@ +/* Copyright 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. +==============================================================================*/ + +// Optimized HIP-kernel backend for Qwen3 gated-delta-net. +// * causal_conv1d prefill/decode -> hipfied Tri Dao kernel via adapter. +// * gated_delta_rule prefill -> chunked WY decomposition + aiter FLA +// h-recurrence + chunk_fwd_o. + +#include +#include + +#include +#include + +#include "fla_api.h" +#include "kernels/dcu/causal_conv1d_adapter.h" +#include "qwen3_gated_delta_net_helpers.h" +#include "qwen3_gated_delta_net_impls.h" + +namespace xllm { +namespace layer { +namespace qwen3_gdn { +namespace kernel_impl { + +namespace { + +constexpr int64_t kChunkSize = 64; + +// Build [NT_total, 2] i64 tensor of (seq_id, local_chunk_id) rows. +torch::Tensor build_chunk_indices(const torch::Tensor& cu_seqlens_i64, + int64_t chunk_size) { + const auto opts_long = cu_seqlens_i64.options().dtype(torch::kInt64); + auto lengths = cu_seqlens_i64.narrow(0, 1, cu_seqlens_i64.size(0) - 1) - + cu_seqlens_i64.narrow(0, 0, cu_seqlens_i64.size(0) - 1); + auto num_chunks = ((lengths + chunk_size - 1) / chunk_size).to(torch::kInt64); + auto cumsum = num_chunks.cumsum(0); + int64_t total = cumsum[-1].item(); + auto arange_total = torch::arange(total, opts_long); + auto zeros = torch::zeros({1}, cumsum.options()); + auto prefix = torch::cat({zeros, cumsum.slice(0, 0, -1)}); + auto repeats_prefix = torch::repeat_interleave(prefix, num_chunks); + auto local_idx = arange_total - repeats_prefix; + auto is_start = (local_idx == 0).to(torch::kInt64); + auto seq_id = is_start.cumsum(0) - 1; + return torch::stack({seq_id, local_idx}, 1).contiguous(); +} + +// Cumulative sum of g within each chunk per sequence per head. Returns fp32. +torch::Tensor chunk_local_cumsum_g(const torch::Tensor& g_in, + const std::vector& cu_seqlens_vec, + int64_t chunk_size) { + auto g_f32 = g_in.to(torch::kFloat32).contiguous(); + auto out = torch::zeros_like(g_f32); + const int64_t N = static_cast(cu_seqlens_vec.size()) - 1; + for (int64_t s = 0; s < N; ++s) { + int64_t bos = cu_seqlens_vec[s]; + int64_t eos = cu_seqlens_vec[s + 1]; + int64_t len = eos - bos; + if (len == 0) continue; + auto seq = g_f32.select(0, 0).slice(0, bos, eos); + int64_t nchunks = (len + chunk_size - 1) / chunk_size; + int64_t full_end = nchunks * chunk_size; + if (full_end == len) { + auto view = seq.view({nchunks, chunk_size, seq.size(-1)}); + out.select(0, 0) + .slice(0, bos, eos) + .copy_(view.cumsum(1).reshape({len, seq.size(-1)})); + } else { + if (nchunks > 1) { + int64_t full_len = (nchunks - 1) * chunk_size; + auto full = seq.slice(0, 0, full_len) + .view({nchunks - 1, chunk_size, seq.size(-1)}); + out.select(0, 0) + .slice(0, bos, bos + full_len) + .copy_(full.cumsum(1).reshape({full_len, seq.size(-1)})); + } + int64_t tail_start = (nchunks - 1) * chunk_size; + auto tail = seq.slice(0, tail_start, len).cumsum(0); + out.select(0, 0).slice(0, bos + tail_start, eos).copy_(tail); + } + } + return out; +} + +// Per-chunk A[i, j] = beta_i * (k_i . k_j) * exp(g_i - g_j) for i > j. +torch::Tensor chunk_scaled_dot_kkt(const torch::Tensor& k, + const torch::Tensor& beta, + const torch::Tensor& g_cum, + const std::vector& cu_seqlens_vec, + int64_t chunk_size) { + const int64_t T = k.size(1); + const int64_t Hv = beta.size(-1); + auto opts_f32 = k.options().dtype(torch::kFloat32); + auto A = torch::zeros({1, T, Hv, chunk_size}, opts_f32); + auto k_v_full = repeat_tensor_heads(k, Hv, 2); + const int64_t N = static_cast(cu_seqlens_vec.size()) - 1; + // strict lower-triangular [BT, BT] mask shared across chunks. + auto tri_full = torch::tril(torch::ones({chunk_size, chunk_size}, opts_f32), + /*diagonal=*/-1); + for (int64_t s = 0; s < N; ++s) { + int64_t bos = cu_seqlens_vec[s]; + int64_t eos = cu_seqlens_vec[s + 1]; + int64_t len = eos - bos; + if (len == 0) continue; + for (int64_t t0 = 0; t0 < len; t0 += chunk_size) { + int64_t rows = std::min(chunk_size, len - t0); + auto k_c = k_v_full.select(0, 0).slice(0, bos + t0, bos + t0 + rows); + auto beta_c = beta.select(0, 0).slice(0, bos + t0, bos + t0 + rows); + auto g_c = g_cum.select(0, 0).slice(0, bos + t0, bos + t0 + rows); + auto k_th = k_c.permute({1, 0, 2}).to(torch::kFloat32); + auto kkT = torch::matmul(k_th, k_th.transpose(1, 2)) + .permute({1, 0, 2}) + .contiguous(); + auto g_diff = + (g_c.unsqueeze(1) - g_c.unsqueeze(0)).permute({0, 2, 1}).contiguous(); + auto mask = tri_full.slice(0, 0, rows).slice(1, 0, rows).unsqueeze(1); + // safe_exp guards masked-out upper triangle from fp32 overflow. + auto safe_g = torch::minimum(g_diff, torch::zeros_like(g_diff)); + auto a_chunk = beta_c.unsqueeze(-1) * kkT * torch::exp(safe_g) * mask; + A.select(0, 0) + .slice(0, bos + t0, bos + t0 + rows) + .slice(-1, 0, rows) + .copy_(a_chunk); + } + } + return A; +} + +// Per-chunk (I + A_strict_lower)^{-1} via triangular solve. +torch::Tensor solve_tril_batched(const torch::Tensor& A, + const std::vector& cu_seqlens_vec, + torch::ScalarType out_dtype) { + const int64_t T = A.size(1); + const int64_t Hv = A.size(2); + const int64_t BT = A.size(-1); + auto opts_f32 = A.options(); + auto out = torch::empty({1, T, Hv, BT}, opts_f32); + auto eye = torch::eye(BT, opts_f32); + const int64_t N = static_cast(cu_seqlens_vec.size()) - 1; + for (int64_t s = 0; s < N; ++s) { + int64_t bos = cu_seqlens_vec[s]; + int64_t eos = cu_seqlens_vec[s + 1]; + int64_t len = eos - bos; + if (len == 0) continue; + for (int64_t t0 = 0; t0 < len; t0 += BT) { + int64_t rows = std::min(BT, len - t0); + auto A_full = torch::zeros({BT, Hv, BT}, opts_f32); + A_full.slice(0, 0, rows) + .copy_(A.select(0, 0).slice(0, bos + t0, bos + t0 + rows)); + auto M = eye.unsqueeze(0) + A_full.permute({1, 0, 2}).contiguous(); + auto X = + at::linalg_solve_triangular(M, + eye.unsqueeze(0).expand({Hv, BT, BT}), + /*upper=*/false, + /*left=*/true, + /*unitriangular=*/true); + out.select(0, 0) + .slice(0, bos + t0, bos + t0 + rows) + .copy_(X.permute({1, 0, 2}).contiguous().slice(0, 0, rows)); + } + } + return out.to(out_dtype); +} + +// w = A_tril @ (beta * exp(g_cum) * k_repeated), u = A_tril @ (beta * v) +std::pair recompute_w_u( + const torch::Tensor& k, + const torch::Tensor& v, + const torch::Tensor& beta, + const torch::Tensor& A_tril, + const torch::Tensor& g_cum, + const std::vector& cu_seqlens_vec, + int64_t chunk_size) { + const int64_t T = k.size(1); + const int64_t Hv = v.size(-2); + const int64_t Kd = k.size(-1); + const int64_t Vd = v.size(-1); + auto w = torch::empty({1, T, Hv, Kd}, k.options()); + auto u = torch::empty({1, T, Hv, Vd}, v.options()); + auto k_v_full = repeat_tensor_heads(k, Hv, 2); + const int64_t N = static_cast(cu_seqlens_vec.size()) - 1; + for (int64_t s = 0; s < N; ++s) { + int64_t bos = cu_seqlens_vec[s]; + int64_t eos = cu_seqlens_vec[s + 1]; + int64_t len = eos - bos; + if (len == 0) continue; + for (int64_t t0 = 0; t0 < len; t0 += chunk_size) { + int64_t rows = std::min(chunk_size, len - t0); + auto A_c = A_tril.select(0, 0) + .slice(0, bos + t0, bos + t0 + rows) + .slice(-1, 0, rows); + auto beta_c = beta.select(0, 0).slice(0, bos + t0, bos + t0 + rows); + auto k_c = k_v_full.select(0, 0).slice(0, bos + t0, bos + t0 + rows); + auto v_c = v.select(0, 0).slice(0, bos + t0, bos + t0 + rows); + auto g_c = g_cum.select(0, 0).slice(0, bos + t0, bos + t0 + rows); + auto A_hrs = A_c.permute({1, 0, 2}).contiguous(); + auto bv = (beta_c.unsqueeze(-1) * v_c) + .to(torch::kFloat32) + .permute({1, 0, 2}) + .contiguous(); + auto u_hv = torch::matmul(A_hrs, bv); + auto weight = beta_c.to(torch::kFloat32) * torch::exp(g_c); + auto bk = (weight.unsqueeze(-1) * k_c.to(torch::kFloat32)) + .permute({1, 0, 2}) + .contiguous(); + auto w_hv = torch::matmul(A_hrs, bk); + w.select(0, 0) + .slice(0, bos + t0, bos + t0 + rows) + .copy_(w_hv.permute({1, 0, 2}).contiguous()); + u.select(0, 0) + .slice(0, bos + t0, bos + t0 + rows) + .copy_(u_hv.permute({1, 0, 2}).contiguous()); + } + } + return {w, u}; +} + +// Per-chunk cross-chunk + intra-chunk attention output. +torch::Tensor chunk_fwd_o(const torch::Tensor& q, + const torch::Tensor& k, + const torch::Tensor& v_new, + const torch::Tensor& h, + const torch::Tensor& g_cum, + double scale, + const std::vector& cu_seqlens_vec, + int64_t chunk_size) { + const int64_t T = q.size(1); + const int64_t Hv = v_new.size(-2); + const int64_t Vd = v_new.size(-1); + auto o = torch::empty({1, T, Hv, Vd}, v_new.options()); + auto q_v_full = repeat_tensor_heads(q, Hv, 2); + auto k_v_full = repeat_tensor_heads(k, Hv, 2); + auto opts_f32 = q.options().dtype(torch::kFloat32); + auto tri_full = torch::tril(torch::ones({chunk_size, chunk_size}, opts_f32), + /*diagonal=*/0); + const int64_t N = static_cast(cu_seqlens_vec.size()) - 1; + int64_t chunk_offset = 0; + for (int64_t s = 0; s < N; ++s) { + int64_t bos = cu_seqlens_vec[s]; + int64_t eos = cu_seqlens_vec[s + 1]; + int64_t len = eos - bos; + if (len == 0) continue; + int64_t local_chunk = 0; + for (int64_t t0 = 0; t0 < len; t0 += chunk_size, ++local_chunk) { + int64_t rows = std::min(chunk_size, len - t0); + auto q_c = q_v_full.select(0, 0) + .slice(0, bos + t0, bos + t0 + rows) + .to(torch::kFloat32); + auto k_c = k_v_full.select(0, 0) + .slice(0, bos + t0, bos + t0 + rows) + .to(torch::kFloat32); + auto v_c = v_new.select(0, 0) + .slice(0, bos + t0, bos + t0 + rows) + .to(torch::kFloat32); + auto g_c = g_cum.select(0, 0).slice(0, bos + t0, bos + t0 + rows); + // h_chunk: aiter stores [Hv, V, K] under transpose_state_layout=true; + // transpose to [Hv, K, V] for q @ h. + auto h_chunk = h.select(0, 0) + .select(0, chunk_offset + local_chunk) + .transpose(-1, -2) + .to(torch::kFloat32) + .contiguous(); + auto q_hv = q_c.permute({1, 0, 2}).contiguous(); + auto oh = torch::matmul(q_hv, h_chunk).permute({1, 0, 2}).contiguous(); + oh = oh * torch::exp(g_c).unsqueeze(-1); + auto k_hv = k_c.permute({1, 0, 2}).contiguous(); + auto A_qk = torch::matmul(q_hv, k_hv.transpose(1, 2)) + .permute({1, 0, 2}) + .contiguous(); + auto g_diff = + (g_c.unsqueeze(1) - g_c.unsqueeze(0)).permute({0, 2, 1}).contiguous(); + auto mask = tri_full.slice(0, 0, rows).slice(1, 0, rows).unsqueeze(1); + // safe_exp: clamp g_diff <= 0 to avoid overflow on masked upper triangle. + auto safe_g = torch::minimum(g_diff, torch::zeros_like(g_diff)); + A_qk = A_qk * torch::exp(safe_g) * mask; + auto A_hv = A_qk.permute({1, 0, 2}).contiguous(); + auto v_hv = v_c.permute({1, 0, 2}).contiguous(); + auto ov = torch::matmul(A_hv, v_hv).permute({1, 0, 2}).contiguous(); + auto out_c = ((oh + ov) * static_cast(scale)).to(v_new.dtype()); + o.select(0, 0).slice(0, bos + t0, bos + t0 + rows).copy_(out_c); + } + chunk_offset += (len + chunk_size - 1) / chunk_size; + } + return o; +} + +} // namespace + +torch::Tensor causal_conv1d_prefill(const torch::Tensor& flat_input, + const torch::Tensor& conv_weight_2d, + torch::Tensor& conv_cache, + const std::vector& cu_seqlens, + const std::vector& state_indices, + int32_t /*kernel_size*/, + bool activation) { + // Layer stores conv_weight as [K, C]; DCU kernel expects [C, K]. + auto weight_ck = conv_weight_2d.transpose(0, 1).contiguous(); + return xllm::kernel::dcu::causal_conv1d_varlen_fwd( + flat_input, weight_ck, cu_seqlens, state_indices, conv_cache, activation); +} + +torch::Tensor causal_conv1d_decode(const torch::Tensor& flat_input, + const torch::Tensor& conv_weight_2d, + torch::Tensor& conv_cache, + const torch::Tensor& state_indices, + int32_t /*kernel_size*/, + bool activation) { + auto weight_ck = conv_weight_2d.transpose(0, 1).contiguous(); + return xllm::kernel::dcu::causal_conv1d_update( + flat_input, + weight_ck, + state_indices.to(torch::kInt32), + conv_cache, + activation); +} + +std::tuple gated_delta_rule_prefill( + torch::Tensor q, + torch::Tensor k, + torch::Tensor v, + torch::Tensor g_in, + torch::Tensor beta_in, + std::optional initial_state, + const std::vector& cu_seqlens_vec, + bool output_final_state, + bool use_qk_l2norm_in_kernel) { + const auto initial_dtype = q.dtype(); + if (use_qk_l2norm_in_kernel) { + q = l2norm(q, -1, 1e-6); + k = l2norm(k, -1, 1e-6); + } + q = q.contiguous(); + k = k.contiguous(); + v = v.contiguous(); + beta_in = beta_in.contiguous(); + const double scale = 1.0 / std::sqrt(static_cast(k.size(-1))); + + auto cu_seqlens_i64 = torch::tensor( + cu_seqlens_vec, + torch::TensorOptions().dtype(torch::kInt64).device(q.device())); + auto chunk_indices = build_chunk_indices(cu_seqlens_i64, kChunkSize); + + auto g_cum = chunk_local_cumsum_g(g_in, cu_seqlens_vec, kChunkSize); + auto A = chunk_scaled_dot_kkt(k, beta_in, g_cum, cu_seqlens_vec, kChunkSize); + auto A_tril = solve_tril_batched(A, cu_seqlens_vec, torch::kFloat32); + auto [w, u] = + recompute_w_u(k, v, beta_in, A_tril, g_cum, cu_seqlens_vec, kChunkSize); + w = w.contiguous(); + u = u.contiguous(); + g_cum = g_cum.contiguous(); + + const int64_t N = static_cast(cu_seqlens_vec.size()) - 1; + std::optional state_idx_opt; + torch::Tensor sglang_state; + if (initial_state.has_value()) { + state_idx_opt = torch::arange( + N, torch::TensorOptions().dtype(torch::kInt32).device(q.device())); + // aiter sglang expects fp32 state in [N, H, V, K] layout when + // transpose_state_layout=true; caller passes fla layout [N, H, K, V]. + sglang_state = initial_state.value() + .to(torch::kFloat32) + .transpose(-1, -2) + .contiguous(); + } else { + sglang_state = torch::zeros( + {N, u.size(2), u.size(3), k.size(3)}, + torch::TensorOptions().dtype(torch::kFloat32).device(k.device())); + state_idx_opt = torch::arange( + N, torch::TensorOptions().dtype(torch::kInt32).device(q.device())); + } + + // The sglang variant is chosen because it updates `initial_state` in place; + // the vllm variant leaves untouched positions of a freshly torch::empty'd + // final_state uninitialized, poisoning ssm_cache later. + auto out = aiter::native::chunk_gated_delta_rule_fwd_sglang( + k, + w, + u, + /*g=*/std::optional(g_cum), + /*gk=*/std::nullopt, + std::optional(sglang_state), + state_idx_opt, + output_final_state, + /*chunk_size=*/static_cast(kChunkSize), + /*save_new_value=*/true, + /*cu_seqlens=*/std::optional(cu_seqlens_i64), + /*chunk_indices=*/std::nullopt, + /*use_exp2=*/false, + /*transpose_state_layout=*/true); + CHECK_EQ(out.size(), 2u) + << "aiter chunk_gated_delta_rule_fwd_sglang returned " + "unexpected number of tensors"; + auto& h = out[0]; + auto& v_new = out[1]; + // sglang wrote final state in-place into sglang_state; convert layout back + // to [N, H, K, V] to match the caller's fla convention. + auto final_state = sglang_state.transpose(-1, -2).contiguous(); + + auto o = + chunk_fwd_o(q, k, v_new, h, g_cum, scale, cu_seqlens_vec, kChunkSize); + return std::make_tuple(o.to(initial_dtype), final_state); +} + +} // namespace kernel_impl +} // namespace qwen3_gdn +} // namespace layer +} // namespace xllm diff --git a/xllm/core/layers/dcu/qwen3_gated_delta_net_torch_impl.cpp b/xllm/core/layers/dcu/qwen3_gated_delta_net_torch_impl.cpp new file mode 100644 index 0000000000..6eaa48ac44 --- /dev/null +++ b/xllm/core/layers/dcu/qwen3_gated_delta_net_torch_impl.cpp @@ -0,0 +1,184 @@ +/* Copyright 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. +==============================================================================*/ + +// Pure-torch fallback for Qwen3 gated-delta-net kernels. Used when the +// optimized HIP kernels are unavailable, or via env-flag A/B testing. + +#include + +#include +#include + +#include "qwen3_gated_delta_net_helpers.h" +#include "qwen3_gated_delta_net_impls.h" + +namespace xllm { +namespace layer { +namespace qwen3_gdn { +namespace torch_impl { + +torch::Tensor causal_conv1d_prefill(const torch::Tensor& flat_input, + const torch::Tensor& conv_weight_2d, + torch::Tensor& conv_cache, + const std::vector& cu_seqlens, + const std::vector& state_indices, + int32_t kernel_size, + bool activation) { + // flat_input: [total_tokens, channels] + // conv_weight_2d: [kernel_size, channels] + const int64_t channels = flat_input.size(1); + const int64_t batch_size = cu_seqlens.size() - 1; + auto options = flat_input.options(); + // Transpose to [channels, kernel_size] for per-channel dot product. + auto weight = conv_weight_2d.transpose(0, 1).contiguous(); + + std::vector outputs; + outputs.reserve(batch_size); + + for (int64_t b = 0; b < batch_size; ++b) { + const int64_t start = cu_seqlens[b]; + const int64_t end = cu_seqlens[b + 1]; + const int64_t seq_len = end - start; + if (seq_len == 0) continue; + + auto seq_input = flat_input.slice(0, start, end); + + const int64_t state_idx = state_indices[b]; + torch::Tensor state; + if (conv_cache.defined() && conv_cache.numel() > 0) { + state = conv_cache[state_idx].clone(); + } else { + state = torch::zeros({kernel_size - 1, channels}, options); + } + + // Pad: [state(k-1,c); input(seq,c)] -> [seq + k - 1, channels] + auto padded = torch::cat({state, seq_input}, 0); + + auto out = torch::zeros({seq_len, channels}, options); + for (int64_t t = 0; t < seq_len; ++t) { + auto window = padded.slice(0, t, t + kernel_size).to(torch::kFloat32); + // out_t[c] = sum_i weight[c, i] * window[i, c] + out[t] = torch::sum(weight * window.transpose(0, 1), {1}); + } + + if (conv_cache.defined() && conv_cache.numel() > 0) { + conv_cache[state_idx] = + padded.slice(0, seq_len, seq_len + kernel_size - 1).clone(); + } + + if (activation) { + out = torch::silu(out); + } + outputs.push_back(out); + } + return torch::cat(outputs, 0).contiguous(); +} + +torch::Tensor causal_conv1d_decode(const torch::Tensor& flat_input, + const torch::Tensor& conv_weight_2d, + torch::Tensor& conv_cache, + const torch::Tensor& state_indices, + int32_t kernel_size, + bool activation) { + const int64_t batch = flat_input.size(0); + const int64_t channels = flat_input.size(1); + auto options = flat_input.options(); + auto weight = conv_weight_2d.transpose(0, 1).contiguous(); + auto weight_f32 = weight.to(torch::kFloat32); + auto outputs = torch::empty({batch, channels}, options); + + for (int64_t b = 0; b < batch; ++b) { + const int64_t idx = state_indices[b].item(); + auto state_t = conv_cache[idx].to(torch::kFloat32); + auto x_t = flat_input[b].to(torch::kFloat32); + auto frame = torch::cat({state_t, x_t.unsqueeze(0)}, 0); + auto out_t = torch::sum(weight_f32 * frame.transpose(0, 1), {1}); + outputs[b] = out_t.to(options.dtype()); + conv_cache[idx] = frame.slice(0, 1, kernel_size).to(options.dtype()); + } + + if (activation) { + outputs = torch::silu(outputs); + } + return outputs; +} + +std::tuple gated_delta_rule( + torch::Tensor query, + torch::Tensor key, + torch::Tensor value, + torch::Tensor g, + torch::Tensor beta, + std::optional initial_state, + bool output_final_state, + bool use_qk_l2norm_in_kernel) { + (void)output_final_state; + const auto initial_dtype = query.dtype(); + + if (use_qk_l2norm_in_kernel) { + query = l2norm(query, -1, 1e-6); + key = l2norm(key, -1, 1e-6); + } + + auto to_f32_trans = [](torch::Tensor x) { + return x.transpose(1, 2).contiguous().to(torch::kFloat32); + }; + query = to_f32_trans(query); + key = to_f32_trans(key); + value = to_f32_trans(value); + beta = to_f32_trans(beta); + g = to_f32_trans(g); + const int64_t value_num_heads = value.size(1); + query = repeat_tensor_heads(query, value_num_heads, 1); + key = repeat_tensor_heads(key, value_num_heads, 1); + + const int64_t batch_size = key.size(0); + const int64_t num_heads = key.size(1); + const int64_t sequence_length = key.size(2); + const int64_t k_head_dim = key.size(3); + const int64_t v_head_dim = value.size(3); + + const float scale_val = 1.0f / std::sqrt(static_cast(query.size(-1))); + query = query * scale_val; + torch::Tensor core_attn_out = torch::zeros( + {batch_size, num_heads, sequence_length, v_head_dim}, + torch::TensorOptions().dtype(torch::kFloat32).device(value.device())); + torch::Tensor last_state; + if (!initial_state.has_value()) { + last_state = torch::zeros( + {batch_size, num_heads, k_head_dim, v_head_dim}, + torch::TensorOptions().dtype(torch::kFloat32).device(value.device())); + } else { + last_state = initial_state.value().to(value.device(), torch::kFloat32); + } + + for (int64_t i = 0; i < sequence_length; ++i) { + auto q_t = query.select(2, i); + auto k_t = key.select(2, i); + auto v_t = value.select(2, i); + auto g_t = g.select(2, i).exp().unsqueeze(-1).unsqueeze(-1); + auto beta_t = beta.select(2, i).unsqueeze(-1); + last_state = last_state * g_t; + auto kv_mem = torch::sum(last_state * k_t.unsqueeze(-1), -2); + auto delta = (v_t - kv_mem) * beta_t; + last_state = last_state + k_t.unsqueeze(-1) * delta.unsqueeze(-2); + core_attn_out.select(2, i) = torch::sum(last_state * q_t.unsqueeze(-1), -2); + } + + core_attn_out = core_attn_out.transpose(1, 2).contiguous().to(initial_dtype); + return std::make_tuple(core_attn_out, last_state); +} + +} // namespace torch_impl +} // namespace qwen3_gdn +} // namespace layer +} // namespace xllm