From 2c3f007a656b8c9414d15320f9a49f15420950b9 Mon Sep 17 00:00:00 2001 From: Ma Mingfei Date: Thu, 25 Jun 2026 09:04:34 +0800 Subject: [PATCH] [CPU] optimize GDN prefill performance (#29117) --- sgl-kernel/csrc/cpu/mamba/fla.cpp | 2129 +++++++++++++++++------------ sgl-kernel/csrc/cpu/vec.h | 71 + test/registered/cpu/test_mamba.py | 70 +- 3 files changed, 1350 insertions(+), 920 deletions(-) diff --git a/sgl-kernel/csrc/cpu/mamba/fla.cpp b/sgl-kernel/csrc/cpu/mamba/fla.cpp index 2ba702789..a05fb1a15 100644 --- a/sgl-kernel/csrc/cpu/mamba/fla.cpp +++ b/sgl-kernel/csrc/cpu/mamba/fla.cpp @@ -4,793 +4,1053 @@ #include "vec_pack.h" namespace { -// For this cpu kernel, we have some innovations aside from the existing gpu kernels: -// 1) Use less parallel loops, i.e. 4 including l2_norm. -// 2) Fuse part of l2_norm with the rest of the computation. -#define THREAD_BUFFER_ALLOC(dst, base_ptr, offset, type, size) \ - type* dst = reinterpret_cast((base_ptr) + (offset)); \ - offset += (size); +// [NOTE] GDN Optimizations on AMX CPU +// * intra loop: fuse `kkt_solve` and `recompute_w_u` so as to avoid materialize `A`. +// * inter loop: fuse `recompute_w_u` and `update_v` so as to avoid materialize `h` and `v_new`. +// * intra loop parallel on H instead of Hv, remove duplicated key @ key.T +// * fuse format pack with elemwise OP as much as possible. +// * update state (FP32) with amx-bf16 where C(FP32) += A(BF16) * B(BF16) +// * compile time mask out upper triangular part in decay mask and tril solve, reduce fma needed. -template -inline void fill_stub(scalar_t* __restrict__ out, float val, int size) { +// * convert to vnni format， expect contiguous input and output +// from [K/2, 2, N] FP32 to [K/2, N, 2] BF16 +// * update src = src * exp(g_last) +template +void pack_vnni2(scalar_t* __restrict__ dst, float* __restrict__ src, const float g_last, int ld_src, int ld_dst) { + static_assert(K % 32 == 0); + static_assert(N % 32 == 0); + + const float scale = std::exp(g_last); +#if defined(CPU_CAPABILITY_AVX512) + constexpr int KB = K / 2; + constexpr int NB = N / 32; + + __m512i s0, s1, d0, d1; + __m512 vd = _mm512_set1_ps(scale); + + const auto trans = [&](auto i) { + constexpr int kb = i / NB; + constexpr int nb = i % NB; + + // [K/2, 2, N/32, 32] -> [K/2, N/32, 32, 2] + constexpr int k0 = kb * 2 + 0; + constexpr int k1 = kb * 2 + 1; + __m512 v00 = _mm512_loadu_ps(src + k0 * ld_src + nb * 32); + __m512 v01 = _mm512_loadu_ps(src + k0 * ld_src + nb * 32 + 16); + __m512 v10 = _mm512_loadu_ps(src + k1 * ld_src + nb * 32); + __m512 v11 = _mm512_loadu_ps(src + k1 * ld_src + nb * 32 + 16); + s0 = (__m512i)_mm512_cvtne2ps_pbh(v01, v00); + s1 = (__m512i)_mm512_cvtne2ps_pbh(v11, v10); + + std::tie(d0, d1) = transpose_2x32_16bit(s0, s1); + _mm512_storeu_si512(dst + kb * ld_dst * 2 + nb * 32 * 2, d0); + _mm512_storeu_si512(dst + kb * ld_dst * 2 + nb * 32 * 2 + 32, d1); + + // update src = src * exp(g_last) + _mm512_storeu_ps(src + k0 * ld_src + nb * 32, _mm512_mul_ps(v00, vd)); + _mm512_storeu_ps(src + k0 * ld_src + nb * 32 + 16, _mm512_mul_ps(v01, vd)); + _mm512_storeu_ps(src + k1 * ld_src + nb * 32, _mm512_mul_ps(v10, vd)); + _mm512_storeu_ps(src + k1 * ld_src + nb * 32 + 16, _mm512_mul_ps(v11, vd)); + }; + Unroll{}(trans); +#else + // [K/2, 2, N] -> [K/2, N, 2] + for (int k = 0; k < K; k += 2) { + for (int n = 0; n < N; ++n) { + const float v0 = src[(k + 0) * ld_src + n]; + const float v1 = src[(k + 1) * ld_src + n]; + dst[(k >> 1) * ld_dst * 2 + n * 2 + 0] = static_cast(v0); + dst[(k >> 1) * ld_dst * 2 + n * 2 + 1] = static_cast(v1); + src[(k + 0) * ld_src + n] = v0 * scale; + src[(k + 1) * ld_src + n] = v1 * scale; + } + } +#endif +} + +template +inline void fill_stub(scalar_t* __restrict__ out, float val) { using Vec = at::vec::Vectorized; constexpr int kVecSize = Vec::size(); + static_assert(SIZE % kVecSize == 0); const Vec data_vec = Vec(static_cast(val)); - int d = 0; -#pragma GCC unroll 4 - for (; d <= size - kVecSize; d += kVecSize) { +#pragma GCC unroll 8 + for (int d = 0; d < SIZE; d += kVecSize) { data_vec.store(out + d); } - if (size - d > 0) { - data_vec.store(out + d, size - d); - } } -template -void chunk_gated_delta_rule_kernel_impl( - scalar_t* __restrict__ out, // [B, T, HV, EV] - float* __restrict__ final_state_data, // [N, HV, EK, EV] - const scalar_t* __restrict__ q_orig, // [B, T, HK, EK] - const scalar_t* __restrict__ k_orig, // [B, T, HK, EK] - const scalar_t* __restrict__ v_orig, // [B, T, HV, EV] - const float* __restrict__ g_orig, // [B, T, HV] FP32 - const scalar_t* __restrict__ b_orig, // [B, T, HV] - const int32_t* __restrict__ cu_seqlens_ptr, // [N + 1] INT32 - float* __restrict__ buff, - scalar_t* __restrict__ reduced_buff, - scalar_t* __restrict__ thread_buff, - const int32_t* __restrict__ chunk_offsets_ptr, - const int32_t* __restrict__ chunk_indices_ptr, - bool use_qk_l2norm_in_kernel, - const int64_t& batch_size, - const int64_t& global_seq_len, - const int64_t& qk_num_head, - const int64_t& v_num_head, - const int64_t& qk_head_size, - const int64_t& v_head_size, - const int64_t& qStrideH, - const int64_t& qStrideT, - const int64_t& kStrideH, - const int64_t& kStrideT, - const int64_t& vStrideH, - const int64_t& vStrideT, - const int64_t& oStrideH, - const int64_t& oStrideT, - const int64_t& global_total_seq_length, - const int64_t& global_num_chunk, - const int64_t& buff_size_16bit_per_thread, - double eps = 1e-5) { - int64_t gStrideH = 1; - int64_t gStrideT = v_num_head; - int64_t bStrideH = 1; - int64_t bStrideT = v_num_head; - int64_t final_state_StrideN = v_num_head * qk_head_size * v_head_size; - int64_t final_state_StrideH = qk_head_size * v_head_size; - int64_t final_state_StrideE = v_head_size; - int64_t head_group = v_num_head / qk_num_head; - float scale = 1.0 / std::sqrt(qk_head_size); - using bVec = at::vec::Vectorized; - using fVec = at::vec::Vectorized; - constexpr int64_t VecSize = bVec::size(); - constexpr int64_t fVecSize = fVec::size(); - - // Data pointers - float* g_pad = buff; - float* core_attn_out = g_pad + v_num_head * global_total_seq_length; - float* decay_mask = core_attn_out + batch_size * v_num_head * global_total_seq_length * v_head_size; - float* v_beta_attn = decay_mask + v_num_head * global_total_seq_length * chunk_size; - - scalar_t* q_pad = reduced_buff; - scalar_t* k_pad = q_pad + qk_num_head * global_total_seq_length * qk_head_size; - scalar_t* v_pad = k_pad + qk_num_head * global_total_seq_length * qk_head_size; - scalar_t* k_beta = v_pad + v_num_head * global_total_seq_length * v_head_size; - scalar_t* v_beta = k_beta + v_num_head * global_total_seq_length * qk_head_size; - scalar_t* k_cumdecay_reduced = v_beta + v_num_head * global_total_seq_length * v_head_size; - scalar_t* q_norm_sum = k_cumdecay_reduced + v_num_head * global_total_seq_length * qk_head_size; - scalar_t* k_norm_sum = q_norm_sum + qk_num_head * global_seq_len; - - if (use_qk_l2norm_in_kernel) { - at::parallel_for(0, qk_num_head * global_seq_len, 0, [&](int64_t begin, int64_t end) { - int64_t h_qk = 0, l = 0; - data_index_init(begin, h_qk, qk_num_head, l, global_seq_len); - for (int64_t i = begin; i < end; ++i) { - auto q_norm_sum_ptr = q_norm_sum + h_qk * global_seq_len + l; - auto k_norm_sum_ptr = k_norm_sum + h_qk * global_seq_len + l; - float sum_q = float(0); - float sum_k = float(0); - fVec sum_q_fvec = fVec(float(0)); - fVec sum_k_fvec = fVec(float(0)); - int64_t q_offset = l * qStrideT + h_qk * qStrideH; - int64_t k_offset = l * qStrideT + h_qk * qStrideH; - int64_t d; - for (d = 0; d <= qk_head_size - VecSize; d += VecSize) { - bVec q_bvec = bVec::loadu(q_orig + q_offset + d); - fVec q_fvec0, q_fvec1; - std::tie(q_fvec0, q_fvec1) = at::vec::convert_to_float(q_bvec); - sum_q_fvec += q_fvec0 * q_fvec0; - sum_q_fvec += q_fvec1 * q_fvec1; - bVec k_bvec = bVec::loadu(k_orig + k_offset + d); - fVec k_fvec0, k_fvec1; - std::tie(k_fvec0, k_fvec1) = at::vec::convert_to_float(k_bvec); - sum_k_fvec += k_fvec0 * k_fvec0; - sum_k_fvec += k_fvec1 * k_fvec1; - } - sum_q += vec_reduce_sum(sum_q_fvec); - sum_k += vec_reduce_sum(sum_k_fvec); - q_norm_sum_ptr[0] = static_cast(float(1) / std::sqrt(sum_q + eps)); - k_norm_sum_ptr[0] = static_cast(float(1) / std::sqrt(sum_k + eps)); - data_index_step(h_qk, qk_num_head, l, global_seq_len); - } - }); +template +struct l2norm_kernel { + static inline void apply(scalar_t* __restrict__ out, const scalar_t* __restrict__ input, float eps) { + TORCH_CHECK(false, "l2norm_kernel: scalar path not implemented!"); } +}; - // query = query * scale - // k_beta = key * beta.unsqueeze(-1) - // v_beta = value * beta.unsqueeze(-1) - // Padding for q/k/v/beta - at::parallel_for(0, qk_num_head * global_num_chunk, 1, [&](int64_t begin, int64_t end) { - int ompIdx = at::get_thread_num(); - int64_t h_qk = 0, c = 0; - data_index_init(begin, h_qk, qk_num_head, c, global_num_chunk); - for ([[maybe_unused]] auto z : c10::irange(begin, end)) { - int64_t ib = chunk_indices_ptr[c * 2]; // idx_batch - int64_t ic = chunk_indices_ptr[c * 2 + 1]; // idx_chunk - int64_t l_orig = cu_seqlens_ptr[ib] + ic * chunk_size; - int64_t l = c * chunk_size; - bool is_tail = (c + 1 == chunk_offsets_ptr[ib + 1]); - int64_t seq_len = cu_seqlens_ptr[ib + 1] - cu_seqlens_ptr[ib]; - int64_t real_chunk_size = is_tail ? seq_len - ic * chunk_size : chunk_size; - auto q_orig_ptr = q_orig + h_qk * qStrideH + l_orig * qStrideT; - auto k_orig_ptr = k_orig + h_qk * kStrideH + l_orig * kStrideT; - auto v_orig_ptr = v_orig + l_orig * vStrideT; - auto b_orig_ptr = b_orig + l_orig * bStrideT; - auto q_pad_ptr = q_pad + h_qk * global_total_seq_length * qk_head_size + l * qk_head_size; - auto k_pad_ptr = k_pad + h_qk * global_total_seq_length * qk_head_size + l * qk_head_size; - auto v_pad_ptr = v_pad + l * v_head_size; - auto k_beta_ptr = k_beta + l * qk_head_size; - auto v_beta_ptr = v_beta + l * v_head_size; +#if defined(CPU_CAPABILITY_AVX512) +template +struct l2norm_kernel { + static inline void apply(at::BFloat16* __restrict__ out, const at::BFloat16* __restrict__ input, float eps) { + static_assert(D % 32 == 0); + constexpr int COLS = D / 32; - for (int64_t j = 0; j < real_chunk_size; j++) { - auto curr_q_orig = q_orig_ptr + j * qStrideT; - auto curr_k_orig = k_orig_ptr + j * kStrideT; - auto curr_q_pad = q_pad_ptr + j * qk_head_size; - auto curr_k_pad = k_pad_ptr + j * qk_head_size; - auto q_scale = - use_qk_l2norm_in_kernel ? *(q_norm_sum + h_qk * global_seq_len + l_orig + j) : static_cast(1); - auto k_scale = - use_qk_l2norm_in_kernel ? *(k_norm_sum + h_qk * global_seq_len + l_orig + j) : static_cast(1); - auto q_scale_vec = bVec(q_scale); - auto k_scale_vec = bVec(k_scale); - int64_t i = 0; - scalar_t scale_reduced = static_cast(scale); - auto vec_scale_reduced = bVec(scale_reduced); - for (; i < fVecSize * (qk_head_size / fVecSize); i += fVecSize) { - auto tmp0 = bVec::loadu(curr_q_orig + i, fVecSize); - auto tmp1 = tmp0 * q_scale_vec * vec_scale_reduced; - tmp1.store(curr_q_pad + i, fVecSize); - auto tmp3 = bVec::loadu(curr_k_orig + i, fVecSize); - auto tmp4 = tmp3 * k_scale_vec; - tmp4.store(curr_k_pad + i, fVecSize); - } + __m512bh va[COLS]; + __m512 vrscale; - for (auto hi = 0; hi < head_group; hi++) { - int64_t h = h_qk * head_group + hi; - auto curr_v_orig = v_orig_ptr + h * vStrideH + j * vStrideT; - auto curr_b_orig = b_orig_ptr + h * bStrideH + j * bStrideT; - scalar_t b_orig_val_reduced = *(curr_b_orig); - auto curr_v_pad = v_pad_ptr + h * global_total_seq_length * v_head_size + j * v_head_size; - auto curr_k_beta = k_beta_ptr + h * global_total_seq_length * qk_head_size + j * qk_head_size; - auto curr_v_beta = v_beta_ptr + h * global_total_seq_length * v_head_size + j * v_head_size; + constexpr float scale = 1.f / std::sqrt(D); + __m512 vscale = _mm512_set1_ps(scale); - // query = query * scale - // k_beta = key * beta.unsqueeze(-1) - int64_t i = 0; - auto vec_b_reduced = bVec(b_orig_val_reduced); - for (; i < fVecSize * (qk_head_size / fVecSize); i += fVecSize) { - auto tmp0 = bVec::loadu(curr_k_orig + i, fVecSize); - auto tmp2 = tmp0 * k_scale_vec * vec_b_reduced; - tmp2.store(curr_k_beta + i, fVecSize); - } - // v_beta = value * beta.unsqueeze(-1) - i = 0; - for (; i < VecSize * (v_head_size / VecSize); i += VecSize) { - auto tmp3 = bVec::loadu(curr_v_orig + i); - tmp3.store(curr_v_pad + i); - auto tmp5 = tmp3 * vec_b_reduced; - tmp5.store(curr_v_beta + i); - } - } + // step 1: load input and do reduce with avx512-bf16 + __m512 vsum = _mm512_set1_ps(0.f); + auto reduce = [&](auto col) { + va[col] = (__m512bh)(_mm512_loadu_si512(input + col * 32)); + vsum = _mm512_dpbf16_ps(vsum, va[col], va[col]); + }; + Unroll{}(reduce); + + float sqsum = _mm512_reduce_add_ps(vsum); + float rscale = 1.f / std::sqrt(sqsum + eps); + vrscale = _mm512_set1_ps(rscale); + + // step 2: apply scale to output + auto map = [&](auto col) { + __m512i a16 = (__m512i)va[col]; + __m512 va0 = CVT_BF16_TO_FP32(_mm512_extracti32x8_epi32(a16, 0)); + __m512 va1 = CVT_BF16_TO_FP32(_mm512_extracti32x8_epi32(a16, 1)); + va0 = _mm512_mul_ps(va0, vrscale); + va1 = _mm512_mul_ps(va1, vrscale); + // keep the mul order same as torch code: + // query = l2norm(query) * scale + if constexpr (has_scale) { + va0 = _mm512_mul_ps(va0, vscale); + va1 = _mm512_mul_ps(va1, vscale); + } + _mm512_storeu_si512(out + col * 32, (__m512i)(_mm512_cvtne2ps_pbh(va1, va0))); + }; + Unroll{}(map); + } +}; +#endif + +template +struct cumsum_kernel { + static inline void + apply(scalar_t* __restrict__ out, const scalar_t* __restrict__ input, int size, int ld_src, int ld_dst) { + TORCH_CHECK(false, "cumsum_kernel: scalar path not implemented!"); + } +}; + +#if defined(CPU_CAPABILITY_AVX512) +template +struct cumsum_kernel { + static inline void apply(float* __restrict__ out, const float* __restrict__ input, int size, int ld_src, int ld_dst) { + // vector length of fp32 for avx512 + static_assert(BLOCK_H == 16); + + __m512i va[16]; + __m512 vsum = _mm512_set1_ps(0.f); + + for (int i = 0; i < CHUNK_SIZE; i += 16) { + // load input data + Unroll<16>{}([&](auto j) { + __m512 v = (i + j < size) ? _mm512_loadu_ps(input + (i + j) * ld_src) : _mm512_setzero_ps(); + vsum = _mm512_add_ps(vsum, v); + va[j] = _mm512_castps_si512(vsum); + }); + // transpose + transpose_16x16_32bit(va); + // store output data + Unroll<16>{}([&](auto j) { _mm512_storeu_si512(out + j * ld_dst + i, va[j]); }); + } + } +}; +#endif + +template +struct decay_mask_kernel { + static inline void apply(scalar_t* __restrict__ out, const scalar_t* __restrict__ input) { + TORCH_CHECK(false, "decay_mask_kernel: scalar path not implemented!"); + } +}; + +#if defined(CPU_CAPABILITY_AVX512) +template +struct decay_mask_kernel { + static inline void apply(float* __restrict__ out, const float* __restrict__ input) { + static_assert(CHUNK_SIZE % 16 == 0); + + constexpr int ROWS = CHUNK_SIZE; + constexpr int COLS = CHUNK_SIZE / 16; + + __m512 va; + __m512 vb[COLS]; + + // step 1: load g[j] + auto loadb = [&](auto i) { vb[i] = _mm512_loadu_ps(input + i * 16); }; + Unroll{}(loadb); + + // step2: exp(g[i] - g[j]) + auto compute = [&](auto i) { + constexpr int row = i / COLS; + constexpr int col = i % COLS; + + if constexpr (col == 0) { + va = _mm512_set1_ps(input[row]); } - for (int64_t j = real_chunk_size; j < chunk_size; j++) { - auto curr_q_pad = q_pad_ptr + j * qk_head_size; - auto curr_k_pad = k_pad_ptr + j * qk_head_size; - int64_t i = 0; - auto vec_zero = bVec(0.0); - for (; i < VecSize * (qk_head_size / VecSize); i += VecSize) { - vec_zero.store(curr_q_pad + i); - vec_zero.store(curr_k_pad + i); - } - for (auto hi = 0; hi < head_group; hi++) { - int64_t h = h_qk * head_group + hi; - auto curr_v_pad = v_pad_ptr + h * global_total_seq_length * v_head_size + j * v_head_size; - auto curr_k_beta = k_beta_ptr + h * global_total_seq_length * qk_head_size + j * qk_head_size; - auto curr_v_beta = v_beta_ptr + h * global_total_seq_length * v_head_size + j * v_head_size; - int64_t i = 0; - for (; i < VecSize * (qk_head_size / VecSize); i += VecSize) { - vec_zero.store(curr_k_beta + i); - } - i = 0; - for (; i < VecSize * (v_head_size / VecSize); i += VecSize) { - vec_zero.store(curr_v_pad + i); - vec_zero.store(curr_v_beta + i); - } - } + // mask vb[col] (already loaded in step 1) for the lower-triangular region + constexpr int len = std::max(0, std::min(row + 1 - col * 16, 16)); + + __m512 vc; + if constexpr (len == 16) { + vc = _mm512_fexp_u20_ps(va - vb[col]); + } else if constexpr (len == 0) { + vc = _mm512_setzero_ps(); + } else { + vc = _mm512_fexp_u20_ps(va - vb[col]); + // do mask for vc + constexpr __mmask16 vmask = (1 << len) - 1; + vc = _mm512_mask_blend_ps(vmask, _mm512_setzero_ps(), vc); } - // Move to the next query - data_index_step(h_qk, qk_num_head, c, global_num_chunk); + _mm512_storeu_ps(out + row * CHUNK_SIZE + col * 16, vc); + }; + Unroll{}(compute); + } +}; +#endif + +template +struct apply_mask_kernel { + static inline void apply( + scalar_t* __restrict__ attn2, + const float* __restrict__ attn, + const scalar_t* __restrict__ beta, + const float* __restrict__ d, + int size, + int b_stride = 0) { + TORCH_CHECK(false, "apply_mask_kernel: scalar path not implemented!"); + } +}; + +#if defined(CPU_CAPABILITY_AVX512) +template +struct apply_mask_kernel { + static inline void apply( + at::BFloat16* __restrict__ attn2, + const float* __restrict__ attn, + const at::BFloat16* __restrict__ beta, + const float* __restrict__ d, + int size, + int b_stride = 0) { + static_assert(CHUNK_SIZE % 16 == 0); + + constexpr int ROWS = CHUNK_SIZE; + constexpr int COLS = CHUNK_SIZE / 16; + + __m512 vbeta; + + // has_beta: attn2 = -attn * beta * d (strict lower) + // !has_beta: attn2 = attn * d (lower incl. diagonal) + auto compute = [&](auto i) { + constexpr int row = i / COLS; + constexpr int col = i % COLS; + + constexpr int len = + has_beta ? std::max(0, std::min(row - col * 16, 16)) : std::max(0, std::min(row + 1 - col * 16, 16)); + if (row < size) { + if constexpr (has_beta) { + if constexpr (col == 0) { + vbeta = _mm512_set1_ps(-static_cast(beta[row * b_stride])); + } + } + + __m512 vc; + if constexpr (len == 0) { + vc = _mm512_setzero_ps(); + } else { + constexpr __mmask16 vmask = (1 << len) - 1; + __m512 va = _mm512_maskz_loadu_ps(vmask, attn + row * CHUNK_SIZE + col * 16); + __m512 vd = _mm512_maskz_loadu_ps(vmask, d + row * CHUNK_SIZE + col * 16); + if constexpr (has_beta) { + vc = _mm512_mul_ps(_mm512_mul_ps(va, vbeta), vd); + } else { + vc = _mm512_mul_ps(va, vd); + } + } + _mm256_storeu_si256( + reinterpret_cast<__m256i*>(attn2 + row * CHUNK_SIZE + col * 16), (__m256i)(_mm512_cvtneps_pbh(vc))); + } + }; + Unroll{}(compute); + } +}; +#endif + +template +struct solve_tril_kernel { + static inline void apply(scalar_t* __restrict__ attn2, int size) { + TORCH_CHECK(false, "solve_tril_kernel: scalar path not implemented!"); + } +}; + +#if defined(CPU_CAPABILITY_AVX512) +template +struct solve_tril_kernel { + static inline void apply(at::BFloat16* __restrict__ attn2, int size) { + static_assert(CHUNK_SIZE % 16 == 0); + + constexpr int COLS = CHUNK_SIZE / 16; + + __m512 va; + __m512 vb[COLS]; + __m512 vsum[COLS]; + + // for len == 0 and row < size, we don't have to write back zero again + // as in `apply_mask_kernel`, we already set zero for the upper-triangular region + for (int i = 1; i < size; ++i) { + // load row attn[..., i, :i] + at::BFloat16* __restrict__ row_ptr = attn2 + i * CHUNK_SIZE; + Unroll{}([&](auto col) { + int len = std::min(i - col * 16, 16); + if (len > 0) { + const __mmask16 vmask = (1 << len) - 1; + vsum[col] = CVT_BF16_TO_FP32(_mm256_maskz_loadu_epi16(vmask, row_ptr + col * 16)); + } + }); + + // row = attn[..., i, :i].clone() + // sub = attn[..., :i, :i].clone() + // vsum = row + (row.unsqueeze(-1) * sub).sum(-2) + for (int k = 0; k < i; ++k) { + va = _mm512_set1_ps(static_cast(row_ptr[k])); + + const at::BFloat16* __restrict__ row_k_ptr = attn2 + k * CHUNK_SIZE; + Unroll{}([&](auto col) { + int len = std::min(k - col * 16, 16); + if (len > 0) { + const __mmask16 vmask = (1 << len) - 1; + vb[col] = CVT_BF16_TO_FP32(_mm256_maskz_loadu_epi16(vmask, row_k_ptr + col * 16)); + vsum[col] = _mm512_fmadd_ps(va, vb[col], vsum[col]); + } + }); + } + + // attn[..., i, :i] = vsum + Unroll{}([&](auto col) { + int len = std::min(i - col * 16, 16); + if (len > 0) { + const __mmask16 vmask = (1 << len) - 1; + _mm256_mask_storeu_epi16(row_ptr + col * 16, vmask, (__m256i)(_mm512_cvtneps_pbh(vsum[col]))); + } + }); + } + + // attn = attn + torch.eye(chunk_size) + for (int i = 0; i < size; ++i) { + attn2[i * CHUNK_SIZE + i] += 1.f; + } + } +}; +#endif + +template +struct apply_beta_kernel { + static inline void apply( + scalar_t* __restrict__ out, + const scalar_t* __restrict__ input, + const scalar_t* __restrict__ beta, + const float* __restrict__ g, + int size, + int ld_src, + int ld_dst, + int b_stride) { + TORCH_CHECK(false, "apply_beta_kernel: scalar path not implemented!"); + } +}; + +#if defined(CPU_CAPABILITY_AVX512) +template +struct apply_beta_kernel { + static inline void apply( + at::BFloat16* __restrict__ out, + const at::BFloat16* __restrict__ input, + const at::BFloat16* __restrict__ beta, + const float* __restrict__ g, + int size, + int ld_src, + int ld_dst, + int b_stride) { + static_assert(D % 32 == 0); + constexpr int COLS = D / 16; + + // get g.exp() and g is padded to CHUNK_SIZE + alignas(64) float g_arr[CHUNK_SIZE]; + if constexpr (has_g) { + Unroll{}([&](auto col) { + __m512 vg = _mm512_loadu_ps(g + col * 16); + __m512 vg_exp = _mm512_fexp_u20_ps(vg); + _mm512_storeu_ps(g_arr + col * 16, vg_exp); + }); + } + + for (int i = 0; i < size; ++i) { + __m512 vbeta; + if constexpr (has_beta) { + vbeta = _mm512_set1_ps(static_cast(beta[i * b_stride])); + } + __m512 vg; + if constexpr (has_g) { + vg = _mm512_set1_ps(g_arr[i]); + } + + Unroll{}([&](auto col) { + // load for 0, 2, 4, 6 + if constexpr (col % 2 == 0) { + __m512i a16 = _mm512_loadu_si512(input + i * ld_src + col * 16); + __m512 va0 = CVT_BF16_TO_FP32(_mm512_extracti32x8_epi32(a16, 0)); + __m512 va1 = CVT_BF16_TO_FP32(_mm512_extracti32x8_epi32(a16, 1)); + if constexpr (has_beta) { + va0 = _mm512_mul_ps(va0, vbeta); + va1 = _mm512_mul_ps(va1, vbeta); + } + if constexpr (has_g) { + va0 = _mm512_mul_ps(va0, vg); + va1 = _mm512_mul_ps(va1, vg); + } + _mm512_storeu_si512(out + i * ld_dst + col * 16, (__m512i)(_mm512_cvtne2ps_pbh(va1, va0))); + } + }); + } + } +}; +#endif + +template +struct update_kernel { + static inline void + apply(scalar_t* __restrict__ out, const float* __restrict__ input, int size, int ld_src, int ld_dst) { + TORCH_CHECK(false, "update_kernel: scalar path not implemented!"); + } +}; + +#if defined(CPU_CAPABILITY_AVX512) +template +struct update_kernel { + static inline void + apply(at::BFloat16* __restrict__ out, const float* __restrict__ input, int size, int ld_src, int ld_dst) { + static_assert(D % 32 == 0); + constexpr int COLS = D / 16; + + for (int i = 0; i < size; ++i) { + Unroll{}([&](auto col) { + if constexpr (col % 2 == 0) { + __m512 va0 = _mm512_loadu_ps(input + i * ld_src + (col + 0) * 16); + __m512 va1 = _mm512_loadu_ps(input + i * ld_src + (col + 1) * 16); + __m512i a16 = (__m512i)(_mm512_cvtne2ps_pbh(va1, va0)); + _mm512_storeu_si512(out + i * ld_dst + col * 16, a16); + } + }); + } + } +}; +#endif + +template +struct update_value_kernel { + static inline void apply( + scalar_t* __restrict__ v_prime2, + const scalar_t* __restrict__ v, + const float* __restrict__ v_prime, + int size, + int padded_size, + int v_strideT) { + TORCH_CHECK(false, "update_kernel: scalar path not implemented!"); + } +}; + +#if defined(CPU_CAPABILITY_AVX512) +template +struct update_value_kernel { + static inline void apply( + at::BFloat16* __restrict__ v_prime2, + const at::BFloat16* __restrict__ v, + const float* __restrict__ v_prime, + int size, + int padded_size, + int v_strideT) { + static_assert(D % 32 == 0); + constexpr int COLS = D / 16; + + // v2' = v - v' + for (int i = 0; i < size; ++i) { + Unroll{}([&](auto col) { + // load for 0, 2, 4, 6 + if constexpr (col % 2 == 0) { + __m512i v16 = _mm512_loadu_si512(v + i * v_strideT + col * 16); + __m512 va0 = CVT_BF16_TO_FP32(_mm512_extracti32x8_epi32(v16, 0)); + __m512 va1 = CVT_BF16_TO_FP32(_mm512_extracti32x8_epi32(v16, 1)); + + __m512 v_prime0 = _mm512_loadu_ps(v_prime + i * D + col * 16); + __m512 v_prime1 = _mm512_loadu_ps(v_prime + i * D + col * 16 + 16); + va0 = _mm512_sub_ps(va0, v_prime0); + va1 = _mm512_sub_ps(va1, v_prime1); + __m512i o16 = (__m512i)(_mm512_cvtne2ps_pbh(va1, va0)); + _mm512_storeu_si512(v_prime2 + i * D + col * 16, o16); + } + }); + } + + // pad the last chunk + for (int i = size; i < padded_size; ++i) { + Unroll{}([&](auto col) { + if constexpr (col % 2 == 0) { + __m512i v16 = _mm512_setzero_si512(); + _mm512_storeu_si512(v_prime2 + i * D + col * 16, v16); + } + }); + } + } +}; +#endif + +template +struct update_key_kernel { + static inline void apply( + scalar_t* __restrict__ k_updated, + const scalar_t* __restrict__ k, + const float* __restrict__ g, + int size, + int k_strideT) { + TORCH_CHECK(false, "update_key_kernel: scalar path not implemented!"); + } +}; + +#if defined(CPU_CAPABILITY_AVX512) +template +struct update_key_kernel { + static inline void apply( + at::BFloat16* __restrict__ k_updated, + const at::BFloat16* __restrict__ k, + const float* __restrict__ g, + int size, + int k_strideT) { + static_assert(D % 32 == 0); + const int MB = div_up(size, 16); + const int KB = D / 16; + + const float g_last = g[size - 1]; + const __m512 vg_last = _mm512_set1_ps(g_last); + + float scale_arr[16]; + __m256i va[16]; + + // from [C, D](MB, KB) to [D, C](KB, MB) + // pad size to 16 in this kernel so that transpose can be done in one loop + for (int mb = 0; mb < MB; ++mb) { + const int mb_size = std::min(size - mb * 16, 16); + // prepare exp(g_last - g) + __m512 vg = _mm512_loadu_ps(g + mb * 16); + _mm512_storeu_ps(scale_arr, _mm512_fexp_u20_ps(_mm512_sub_ps(vg_last, vg))); + for (int kb = 0; kb < KB; ++kb) { + const at::BFloat16* __restrict__ k_ptr = k + mb * 16 * k_strideT + kb * 16; + at::BFloat16* __restrict__ k_updated_ptr = k_updated + kb * 16 * CHUNK_SIZE + mb * 16; + // load 16 regs + Unroll<16>{}([&](auto m) { + if (m < mb_size) { + __m256i v16 = _mm256_loadu_si256(reinterpret_cast(k_ptr + m * k_strideT)); + __m512 v32 = _mm512_mul_ps(CVT_BF16_TO_FP32(v16), _mm512_set1_ps(scale_arr[m])); + va[m] = (__m256i)_mm512_cvtneps_pbh(v32); + } else { + va[m] = _mm256_setzero_si256(); + } + }); + // transpose 16x16 + transpose_16x16_16bit(va); + // store 16 regs + Unroll<16>{}( + [&](auto k) { _mm256_storeu_si256(reinterpret_cast<__m256i*>(k_updated_ptr + k * CHUNK_SIZE), va[k]); }); + } + } + } +}; +#endif + +// template head_dim here to reduce extra read +// * normal approach: read inputs 2 times: +// - reduce: 1R +// - scale: 1R + 1W +// * keep input data in register: +// - reduce: 1R +// - scale: 1W +template +void l2norm_fwd_kernel_impl( + scalar_t* __restrict__ query_norm, + scalar_t* __restrict__ key_norm, + const scalar_t* __restrict__ query, + const scalar_t* __restrict__ key, + float eps, + int64_t T, + int64_t H, + int64_t q_strideT, + int64_t q_strideH, + int64_t k_strideT, + int64_t k_strideH) { + // expected to be contuguous + int64_t qn_strideH = D; + int64_t kn_strideH = D; + + // parallel on [B, T, H] + at::parallel_for(0, T * H, 0, [&](int64_t begin, int64_t end) { + int64_t t{0}, h{0}; + data_index_init(begin, t, T, h, H); + + for (int64_t i = begin; i < end; ++i) { + const scalar_t* __restrict__ q_ptr = query + t * q_strideT + h * q_strideH; + const scalar_t* __restrict__ k_ptr = key + t * k_strideT + h * k_strideH; + scalar_t* __restrict__ qn_ptr = query_norm + i * qn_strideH; + scalar_t* __restrict__ kn_ptr = key_norm + i * kn_strideH; + + l2norm_kernel::apply(qn_ptr, q_ptr, eps); + l2norm_kernel::apply(kn_ptr, k_ptr, eps); + + // move to the next index + data_index_step(t, T, h, H); } }); +} - at::parallel_for(0, v_num_head * global_num_chunk, 1, [&](int64_t begin, int64_t end) { - int64_t h = 0, c = 0; - data_index_init(begin, h, v_num_head, c, global_num_chunk); - int ompIdx = at::get_thread_num(); - int64_t offset = 0; - scalar_t* thread_buff_ptr = thread_buff + ompIdx * buff_size_16bit_per_thread; - THREAD_BUFFER_ALLOC(k_transpose, thread_buff_ptr, offset, scalar_t, qk_head_size * chunk_size); - THREAD_BUFFER_ALLOC(v_pack, thread_buff_ptr, offset, scalar_t, chunk_size * v_head_size); - THREAD_BUFFER_ALLOC(k_beta_g, thread_buff_ptr, offset, scalar_t, chunk_size * qk_head_size); - THREAD_BUFFER_ALLOC(k_beta_g_pack, thread_buff_ptr, offset, scalar_t, chunk_size * qk_head_size); - THREAD_BUFFER_ALLOC(curr_attn, thread_buff_ptr, offset, float, chunk_size* chunk_size * 2); - THREAD_BUFFER_ALLOC(curr_attn_reduced, thread_buff_ptr, offset, scalar_t, chunk_size * chunk_size); - THREAD_BUFFER_ALLOC(k_cumdecay, thread_buff_ptr, offset, float, chunk_size* qk_head_size * 2); - THREAD_BUFFER_ALLOC(row, thread_buff_ptr, offset, float, chunk_size * 2); - THREAD_BUFFER_ALLOC(updated, thread_buff_ptr, offset, float, chunk_size * 2); - for ([[maybe_unused]] auto z : c10::irange(begin, end)) { - int64_t ib = chunk_indices_ptr[c * 2]; // idx_batch - int64_t ic = chunk_indices_ptr[c * 2 + 1]; // idx_chunk - int64_t l_orig = cu_seqlens_ptr[ib] + ic * chunk_size; - int64_t seq_len = cu_seqlens_ptr[ib + 1] - cu_seqlens_ptr[ib]; - int64_t h_qk = h / head_group; - auto curr_g_orig = g_orig + h * gStrideH + l_orig * gStrideT; - auto curr_g_pad = g_pad + h * global_total_seq_length + c * chunk_size; - auto curr_decay_mask = decay_mask + h * global_total_seq_length * chunk_size + c * chunk_size * chunk_size; - auto curr_k_pad = k_pad + h_qk * global_total_seq_length * qk_head_size + c * chunk_size * qk_head_size; - auto curr_k_beta = k_beta + h * global_total_seq_length * qk_head_size + c * chunk_size * qk_head_size; - auto curr_k_cumdecay_reduced = - k_cumdecay_reduced + h * global_total_seq_length * qk_head_size + c * chunk_size * qk_head_size; - auto curr_v_beta = v_beta + h * global_total_seq_length * v_head_size + c * chunk_size * v_head_size; - auto curr_value = v_beta_attn + h * global_total_seq_length * v_head_size + c * chunk_size * v_head_size; +// g : [B, T, Hv] +// g_ : [B, Hv, NT, C] -> [B, NT, HB, BLOCK_H, C] +// cu_seqlens : [num_seqs + 1] +// chunk_indices : [NT * 2] +template +void chunk_local_cumsum_kernel_impl( + scalar_t* __restrict__ g_, + const scalar_t* __restrict__ g, + const int32_t* __restrict__ cu_seqlens, + const int32_t* __restrict__ chunk_indices, + int64_t Hv, + int64_t NT) { + constexpr int BLOCK_H = 16; + // TODO: now we only support qwen3.5 configs (H/Hv == 16/32) + TORCH_CHECK(Hv % BLOCK_H == 0); + int64_t HB = Hv / BLOCK_H; - float acc_val = 0; - for (int64_t i = 0; i < chunk_size; i++) { - // Padding for g - // g = g.cumsum(dim=-1) - // g: [B, HV, num_chunk, chunk_size] - if (ic * chunk_size + i < seq_len) { - acc_val += curr_g_orig[i * gStrideT]; - } - curr_g_pad[i] = acc_val; - // decay_mask = ((g.unsqueeze(-1) - g.unsqueeze(-2)).tril().exp().float()).tril() - // decay_mask: [B, HV, num_chunk, chunk_size, chunk_size] - float curr_g_pad_i = static_cast(curr_g_pad[i]); - auto vec_curr_g_pad_i = fVec(curr_g_pad_i); - int64_t j = 0; - int64_t len = i + 1; - for (; j < fVecSize * (len / fVecSize); j += fVecSize) { - auto tmp0 = fVec::loadu(curr_g_pad + j); - auto tmp1 = vec_curr_g_pad_i - tmp0; - auto tmp2 = tmp1.exp_u20(); - tmp2.store(curr_decay_mask + i * chunk_size + j); - } - if (j < len) { - auto tmp0 = fVec::loadu(curr_g_pad + j, len - j); - auto tmp1 = vec_curr_g_pad_i - tmp0; - auto tmp2 = tmp1.exp_u20(); - tmp2.store(curr_decay_mask + i * chunk_size + j, len - j); - } + // parallel on [NT * HB] to increase parallelism + at::parallel_for(0, NT * HB, 0, [&](int64_t begin, int64_t end) { + int64_t nt{0}, hb{0}; + data_index_init(begin, nt, NT, hb, HB); + + for (int64_t i = begin; i < end; ++i) { + int32_t bs = chunk_indices[nt * 2 + 0]; + int32_t batch_offset = cu_seqlens[bs]; + int32_t seqlen = cu_seqlens[bs + 1] - cu_seqlens[bs]; + int64_t mb_start = chunk_indices[nt * 2 + 1] * CHUNK_SIZE; + int64_t mb_size = std::min(seqlen - mb_start, int64_t(CHUNK_SIZE)); + + const scalar_t* __restrict__ g_ptr = g + (batch_offset + mb_start) * Hv + hb * BLOCK_H; + scalar_t* __restrict__ gsum_ptr = g_ + nt * (Hv * CHUNK_SIZE) + hb * (BLOCK_H * CHUNK_SIZE); + cumsum_kernel::apply(gsum_ptr, g_ptr, mb_size, Hv, CHUNK_SIZE); + + // move to the next index + data_index_step(nt, NT, hb, HB); + } + }); +} + +#define DECL_BUF(type, name, size_expr) alignas(64) type name[(size_expr)] +#define DECL_ZERO_BUF(type, name, size_expr) \ + DECL_BUF(type, name, size_expr); \ + fill_stub(name, 0.f) + +// w : [B, T, Hv, D] +// u : [B, T, Hv, Dv] +// d : [B, NT, Hv, C, C] +// k : [B, T, H, D] +// v : [B, T, Hv, Dv] +// g : [B, NT, Hv, C] +// beta : [B, T, Hv] +// cu_seqlens : [num_seqs + 1] +// chunk_indices : [NT * 2] +template +void chunk_gated_delta_rule_fwd_intra_kernel_impl( + scalar_t* __restrict__ w, + scalar_t* __restrict__ u, + float* __restrict__ d, + const scalar_t* __restrict__ k, + const scalar_t* __restrict__ v, + const float* __restrict__ g, + const scalar_t* __restrict__ beta, + const int32_t* __restrict__ cu_seqlens, + const int32_t* __restrict__ chunk_indices, + int64_t H, + int64_t Hv, + int64_t NT, + int64_t k_strideT, + int64_t k_strideH, + int64_t v_strideT, + int64_t v_strideH) { + // head group, expect to be 1，2，4 for qwen3.5 + const int64_t HG = Hv / H; + + // strides + const int64_t w_strideT = Hv * D; + const int64_t w_strideH = D; + const int64_t u_strideT = Hv * D; + const int64_t u_strideH = D; + + // [NB]: parallel on [NT, H] + // * parallel on num_heads and go sequential on num_heads_v, + // * avoid instantialize k_beta (beta * k) + // * compute key @ key^T * beta instead of k_beta @ key^T, same as triton impl + // * compute key @ key^T once for each k head index and reuse for v head index + at::parallel_for(0, NT * H, 0, [&](int64_t begin, int64_t end) { + int64_t nt{0}, h{0}; + data_index_init(begin, nt, NT, h, H); + + // thread local temp buffer + DECL_ZERO_BUF(scalar_t, tmp, CHUNK_SIZE * D); + DECL_ZERO_BUF(scalar_t, tmp2, CHUNK_SIZE * D); + DECL_ZERO_BUF(float, attn, CHUNK_SIZE* CHUNK_SIZE); + DECL_ZERO_BUF(scalar_t, attn2, CHUNK_SIZE * CHUNK_SIZE); + DECL_ZERO_BUF(float, tmp3, CHUNK_SIZE* D); + + // alias + scalar_t* __restrict__ k_packed = tmp; + scalar_t* __restrict__ k_beta = tmp; + scalar_t* __restrict__ v_beta = tmp; + scalar_t* __restrict__ k_beta_packed = tmp2; + scalar_t* __restrict__ v_beta_packed = tmp2; + float* __restrict__ k_updated = tmp3; + float* __restrict__ v_updated = tmp3; + + for (int64_t i = begin; i < end; ++i) { + int32_t bs = chunk_indices[nt * 2 + 0]; + int32_t batch_offset = cu_seqlens[bs]; + int32_t seqlen = cu_seqlens[bs + 1] - cu_seqlens[bs]; + int64_t mb_start = chunk_indices[nt * 2 + 1] * CHUNK_SIZE; + int64_t mb_size = std::min(seqlen - mb_start, int64_t(CHUNK_SIZE)); + + // mb_size` is K in 5.c, 5.g, pad to TILE_K; + const int64_t padded_mb_size = div_up((int)mb_size, TILE_K) * TILE_K; + + // step 1: decay_mask = ((g.unsqueeze(-1) - g.unsqueeze(-2)).tril().exp().float()).tril() + for (int64_t hv = h * HG; hv < h * HG + HG; ++hv) { + const float* __restrict__ g_ptr = g + nt * (Hv * CHUNK_SIZE) + hv * CHUNK_SIZE; + float* __restrict__ d_ptr = d + nt * (Hv * CHUNK_SIZE * CHUNK_SIZE) + hv * (CHUNK_SIZE * CHUNK_SIZE); + decay_mask_kernel::apply(d_ptr, g_ptr); } - // attn = k_beta @ key.transpose(-1, -2) - // attn: [B, HV, num_chunk, chunk_size, chunk_size] - // transpose and pack for key + // step 2: attn = key @ key^T + const scalar_t* __restrict__ k_ptr = k + (batch_offset + mb_start) * k_strideT + h * k_strideH; pack_vnni( - /* dst */ k_transpose, - /* src */ curr_k_pad, - /* N */ chunk_size, - /* K */ qk_head_size, - /* ld_src */ qk_head_size, - /* ld_dst */ chunk_size); - // k_beta @ key.transpose(-1, -2) + /* dst */ k_packed, + /* src */ k_ptr, + /* N */ mb_size, + /* K */ D, + /* ld_src */ k_strideT, + /* ld_dst */ CHUNK_SIZE); + at::native::cpublas::brgemm( - /* M */ chunk_size, - /* N */ chunk_size, - /* K */ qk_head_size, - /* lda */ qk_head_size, - /* ldb */ chunk_size, - /* ldc */ chunk_size, + /* M */ mb_size, + /* N */ mb_size, + /* K */ D, + /* lda */ k_strideT, + /* ldb */ CHUNK_SIZE, + /* ldc */ CHUNK_SIZE, /* add_C */ false, - /* A */ curr_k_beta, - /* B */ k_transpose, - /* C */ curr_attn); - // attn = attn * decay_mask - for (int64_t m = 0; m < chunk_size; m++) { - at::vec::map2( - [](fVec x, fVec y) { return fVec(0) - x * y; }, - curr_attn + m * chunk_size, - curr_attn + m * chunk_size, - curr_decay_mask + m * chunk_size, - chunk_size); - } + /* A */ k_ptr, + /* B */ k_packed, + /* C */ attn); - // chunk decay - // attn: [B, HV, num_chunk, chunk_size, chunk_size] - // mask = torch.triu(torch.ones(chunk_size, chunk_size, dtype=torch.bool, device=query.device), diagonal=0) - // attn = -attn.masked_fill(mask, 0) - // attn[..., i, :i] = row + (row.unsqueeze(-1) * sub).sum(-2) [B, HV, num_chunk, i] - // attn = attn + torch.eye(chunk_size, dtype=attn.dtype, device=attn.device) - // attn = -attn.masked_fill(mask, 0) - for (int i = 0; i < chunk_size; i++) { - const auto vec_zero = fVec(0); - int64_t len = chunk_size - i; - int64_t front = len % fVecSize; - int64_t j = i; - // first masked vec for alignment - if (front > 0) { - vec_zero.store(curr_attn + i * chunk_size + j, front); - j += front; - } - for (; j < fVecSize * (chunk_size / fVecSize); j += fVecSize) { - vec_zero.store(curr_attn + i * chunk_size + j); - } - } - for (int i = 1; i < chunk_size; i++) { - // row = attn[..., i, :i] [B, HK, num_chunk, i] - int64_t j = 0; - int64_t len = i; - for (; j < fVecSize * (len / fVecSize); j += fVecSize) { - auto tmp0 = fVec::loadu(curr_attn + i * chunk_size + j); - tmp0.store(row + j); - } - if (j < len) { - auto tmp0 = fVec::loadu(curr_attn + i * chunk_size + j, len - j); - tmp0.store(row + j, len - j); - } - // (row.unsqueeze(-1) * sub).sum(-2) - fill_stub(updated, 0, i); - for (int k = 0; k < i; k++) { - float row_k = row[k]; - auto vec_row_k = fVec(row_k); - int64_t j = 0; - int64_t len = i; - for (; j < fVecSize * (len / fVecSize); j += fVecSize) { - auto tmp0 = fVec::loadu(curr_attn + k * chunk_size + j); - auto tmp1 = vec_row_k * tmp0; - auto tmp2 = fVec::loadu(updated + j); - auto tmp3 = tmp1 + tmp2; - tmp3.store(updated + j); - } - if (j < len) { - auto tmp0 = fVec::loadu(curr_attn + k * chunk_size + j, len - j); - auto tmp1 = vec_row_k * tmp0; - auto tmp2 = fVec::loadu(updated + j); - auto tmp3 = tmp1 + tmp2; - tmp3.store(updated + j, len - j); - } - } - // attn[..., i, :i] = row + sum(...) - j = 0; - len = i; - for (; j < fVecSize * (len / fVecSize); j += fVecSize) { - auto tmp0 = fVec::loadu(row + j); - auto tmp1 = fVec::loadu(updated + j); - auto tmp2 = tmp0 + tmp1; - tmp2.store(curr_attn + i * chunk_size + j); - } - if (j < len) { - auto tmp0 = fVec::loadu(row + j, len - j); - auto tmp1 = fVec::loadu(updated + j, len - j); - auto tmp2 = tmp0 + tmp1; - tmp2.store(curr_attn + i * chunk_size + j, len - j); - } - } - for (int i = 0; i < chunk_size; i++) { - curr_attn[i * chunk_size + i] += 1.0f; - at::vec::map( - [](fVec x) { return x; }, curr_attn_reduced + i * chunk_size, curr_attn + i * chunk_size, chunk_size); - } + for (int64_t hv = h * HG; hv < h * HG + HG; ++hv) { + // step 3: attn2 = -attn * beta * d + const scalar_t* __restrict__ beta_ptr = beta + (batch_offset + mb_start) * Hv + hv; + const float* __restrict__ d_ptr = d + nt * (Hv * CHUNK_SIZE * CHUNK_SIZE) + hv * (CHUNK_SIZE * CHUNK_SIZE); + apply_mask_kernel::apply(attn2, attn, beta_ptr, d_ptr, mb_size, Hv); - // v_beta_attn = attn @ v_beta - // k_cumdecay = attn @ (k_beta * g.exp().unsqueeze(-1)) - // v_beta_attn: [B, HV, num_chunk, chunk_size, EV] - // k_beta_g = k_beta * g: [B, HV, num_chunk, chunk_size, EK] - // k_cumdecay: [B, HV, num_chunk, chunk_size, EK] - // pack for value - pack_vnni2( - /* dst */ v_pack, - /* src */ curr_v_beta, - /* N */ chunk_size, - /* K */ v_head_size, - /* ld_src */ v_head_size, - /* ld_dst */ v_head_size); - // value = attn @ v_beta - at::native::cpublas::brgemm( - /* M */ chunk_size, - /* N */ v_head_size, - /* K */ chunk_size, - /* lda */ chunk_size, - /* ldb */ v_head_size, - /* ldc */ v_head_size, - /* add_C */ false, - /* A */ curr_attn_reduced, - /* B */ v_pack, - /* C */ curr_value); - // k_beta_g = k_beta * g.exp().unsqueeze(-1) - for (int64_t j = 0; j < chunk_size; j++) { - int64_t i = 0; - float g_exp = std::exp(curr_g_pad[j]); - scalar_t g_exp_reduced = static_cast(g_exp); - auto vec_g_exp_reduced = bVec(g_exp_reduced); - for (; i < VecSize * (qk_head_size / VecSize); i += VecSize) { - auto tmp0 = bVec::loadu(curr_k_beta + j * qk_head_size + i); - auto tmp1 = tmp0 * vec_g_exp_reduced; - tmp1.store(k_beta_g + j * qk_head_size + i); - } - } - // pack for k_beta_g - pack_vnni2( - /* dst */ k_beta_g_pack, - /* src */ k_beta_g, - /* N */ chunk_size, - /* K */ qk_head_size, - /* ld_src */ qk_head_size, - /* ld_dst */ qk_head_size); - // k_cumdecay = attn @ k_beta_g - at::native::cpublas::brgemm( - /* M */ chunk_size, - /* N */ qk_head_size, - /* K */ chunk_size, - /* lda */ chunk_size, - /* ldb */ qk_head_size, - /* ldc */ qk_head_size, - /* add_C */ false, - /* A */ curr_attn_reduced, - /* B */ k_beta_g_pack, - /* C */ k_cumdecay); - for (int i = 0; i < chunk_size; i++) { - at::vec::map( - [](fVec x) { return x; }, - curr_k_cumdecay_reduced + i * qk_head_size, - k_cumdecay + i * qk_head_size, - qk_head_size); - } + // step 4: solve_tril(attn2) -> (I + L)^{-1}, L = strict-lower from step 3 + // for i in 1..C-1: attn2[i, :i] += (attn2[i, :i] * attn2[:i, :i]).sum(-1) + // attn2 += eye(C) + solve_tril_kernel::apply(attn2, mb_size); - // Move to the next query - data_index_step(h, v_num_head, c, global_num_chunk); - } - }); + // step 5: recompute_w_u + // w = attn2 @ (k_beta * g.exp().unsqueeze(-1)) + // u = attn2 @ value * beta.unsqueeze(-1) + const float* __restrict__ g_ptr = g + nt * (Hv * CHUNK_SIZE) + hv * CHUNK_SIZE; + const scalar_t* __restrict__ v_ptr = v + (batch_offset + mb_start) * v_strideT + hv * v_strideH; - // for each chunk - at::parallel_for(0, batch_size * v_num_head, 1, [&](int64_t begin, int64_t end) { - int64_t b = 0, h = 0; - data_index_init(begin, b, batch_size, h, v_num_head); - int ompIdx = at::get_thread_num(); - int64_t offset = - /* k_transpose */ qk_head_size * chunk_size + - /* v_pack */ chunk_size * v_head_size + - /* k_beta_g */ chunk_size * qk_head_size + - /* k_beta_g_pack */ chunk_size * qk_head_size + - /* attn */ chunk_size * chunk_size * 2 + - /* attn_reduced */ chunk_size * chunk_size + - /* k_cumdecay */ chunk_size * qk_head_size * 2 + - /* row */ chunk_size * 2 + - /* updated */ chunk_size * 2; - scalar_t* thread_buff_ptr = thread_buff + ompIdx * buff_size_16bit_per_thread; - THREAD_BUFFER_ALLOC( - curr_last_recurrent_state_reduced, thread_buff_ptr, offset, scalar_t, qk_head_size * v_head_size); - THREAD_BUFFER_ALLOC( - curr_last_recurrent_state_pack_reduced, thread_buff_ptr, offset, scalar_t, qk_head_size * v_head_size); - THREAD_BUFFER_ALLOC(k_transpose_i, thread_buff_ptr, offset, scalar_t, qk_head_size * chunk_size); - THREAD_BUFFER_ALLOC(attn_i, thread_buff_ptr, offset, float, chunk_size* chunk_size * 2); - THREAD_BUFFER_ALLOC(attn_i_reduced, thread_buff_ptr, offset, scalar_t, chunk_size * chunk_size); - THREAD_BUFFER_ALLOC(v_prime, thread_buff_ptr, offset, float, chunk_size* v_head_size * 2); - THREAD_BUFFER_ALLOC(v_prime_reduced, thread_buff_ptr, offset, scalar_t, chunk_size * v_head_size); - THREAD_BUFFER_ALLOC(v_prime_pack_reduced, thread_buff_ptr, offset, scalar_t, chunk_size * v_head_size); - THREAD_BUFFER_ALLOC(qg, thread_buff_ptr, offset, scalar_t, chunk_size * qk_head_size); - THREAD_BUFFER_ALLOC(attn_inter, thread_buff_ptr, offset, float, chunk_size* v_head_size * 2); - THREAD_BUFFER_ALLOC(kg, thread_buff_ptr, offset, scalar_t, chunk_size * qk_head_size); - THREAD_BUFFER_ALLOC(kg_transpose, thread_buff_ptr, offset, scalar_t, qk_head_size * chunk_size); - THREAD_BUFFER_ALLOC(kgv, thread_buff_ptr, offset, float, qk_head_size* v_head_size * 2); + // 5.a key = key * beta * g.exp + apply_beta_kernel::apply( + k_beta, k_ptr, beta_ptr, g_ptr, mb_size, k_strideT, D, Hv); - for ([[maybe_unused]] auto z : c10::irange(begin, end)) { - int64_t start_q = cu_seqlens_ptr[b]; - int64_t seq_len = cu_seqlens_ptr[b + 1] - start_q; - int64_t num_chunk = chunk_offsets_ptr[b + 1] - chunk_offsets_ptr[b]; - int64_t chunk_offset = chunk_offsets_ptr[b]; - int64_t len_offset = chunk_offset * chunk_size; - - int64_t h_qk = h / head_group; - auto out_ptr = out + start_q * oStrideT; - auto curr_q = q_pad + len_offset * qk_head_size + - h_qk * global_total_seq_length * qk_head_size; // [num_chunk, chunk_size, EK] - auto curr_k = k_pad + len_offset * qk_head_size + - h_qk * global_total_seq_length * qk_head_size; // [num_chunk, chunk_size, EK] - auto curr_v = v_beta_attn + h * global_total_seq_length * v_head_size; // [num_chunk, chunk_size, EV] - auto curr_decay_mask = - decay_mask + h * global_total_seq_length * chunk_size; // [num_chunk, chunk_size, chunk_size] - auto curr_k_cumdecay_reduced = - k_cumdecay_reduced + h * global_total_seq_length * qk_head_size; // [num_chunk, chunk_size, EK] - auto curr_last_recurrent_state = - final_state_data + b * final_state_StrideN + h * final_state_StrideH; // [EK, EV] - auto curr_g_pad = g_pad + len_offset + h * global_total_seq_length; // [num_chunk, chunk_size] - auto curr_core_attn_out = core_attn_out + len_offset * v_head_size + - h * global_total_seq_length * v_head_size; // [num_chunk, chunk_size, EV] - for (int64_t c = 0; c < num_chunk; c++) { - for (int i = 0; i < qk_head_size; i++) { - at::vec::map( - [](fVec x) { return x; }, - curr_last_recurrent_state_reduced + i * v_head_size, - curr_last_recurrent_state + i * v_head_size, - v_head_size); - } - auto q_i = curr_q + c * chunk_size * qk_head_size; // [chunk_size, EK] - auto k_i = curr_k + c * chunk_size * qk_head_size; // [chunk_size, EK] - auto v_i = curr_v + (chunk_offset + c) * chunk_size * v_head_size; // [chunk_size, EV] - auto decay_mask_i = curr_decay_mask + (chunk_offset + c) * chunk_size * chunk_size; // [chunk_size, chunk_size] - auto k_cumdecay_i_reduced = - curr_k_cumdecay_reduced + (chunk_offset + c) * chunk_size * qk_head_size; // [chunk_size, EK] - auto g_pad_i = curr_g_pad + c * chunk_size; // [chunk_size] - auto core_attn_out_i = curr_core_attn_out + c * chunk_size * v_head_size; // [chunk_size, EV] - - // attn_i = (q_i @ k_i.transpose(-1, -2) * decay_mask[:, :, i]).masked_fill_(mask, 0) - // k_transpose_i = k_i.transpose(-1, -2) - pack_vnni( - /* dst */ k_transpose_i, - /* src */ k_i, - /* N */ chunk_size, - /* K */ qk_head_size, - /* ld_src */ qk_head_size, - /* ld_dst */ chunk_size); - // attn_i = q_i @ k_transpose_i - at::native::cpublas::brgemm( - /* M */ chunk_size, - /* N */ chunk_size, - /* K */ qk_head_size, - /* lda */ qk_head_size, - /* ldb */ chunk_size, - /* ldc */ chunk_size, - /* add_C */ false, - /* A */ q_i, - /* B */ k_transpose_i, - /* C */ attn_i); - // attn_i = attn_i * decay_mask_i - for (int64_t m = 0; m < chunk_size; m++) { - auto attn_i_m = attn_i + m * chunk_size; - auto attn_i_reduced_m = attn_i_reduced + m * chunk_size; - auto decay_mask_i_m = decay_mask_i + m * chunk_size; - int64_t n = 0; - for (; n < fVecSize * (chunk_size / fVecSize); n += fVecSize) { - auto tmp0 = fVec::loadu(attn_i_m + n); - auto tmp1 = fVec::loadu(decay_mask_i_m + n); - auto tmp2 = tmp0 * tmp1; - auto tmp3 = at::vec::convert(tmp2); - tmp3.store(attn_i_reduced_m + n, fVecSize); - } - if (n < chunk_size) { - auto tmp0 = fVec::loadu(attn_i_m + n, chunk_size - n); - auto tmp1 = fVec::loadu(decay_mask_i_m + n, chunk_size - n); - auto tmp2 = tmp0 * tmp1; - auto tmp3 = at::vec::convert(tmp2); - tmp3.store(attn_i_reduced_m + n, chunk_size - n); - } - } - // mask = torch.triu(torch.ones(chunk_size, chunk_size, dtype=torch.bool, device=query.device), diagonal=1) - // attn_i = attn_i.masked_fill_(mask, 0) - for (int i = 0; i < chunk_size - 1; i++) { - const auto vec_zero = bVec(0); - int64_t len = chunk_size - i - 1; - int64_t front = len % VecSize; - int64_t j = i + 1; - // first masked vec for alignment - if (front > 0) { - vec_zero.store(attn_i_reduced + i * chunk_size + j, front); - j += front; - } - for (; j < VecSize * (chunk_size / VecSize); j += VecSize) { - vec_zero.store(attn_i_reduced + i * chunk_size + j); - } - } - - // pack for curr_last_recurrent_state + // 5.b pack key pack_vnni2( - /* dst */ curr_last_recurrent_state_pack_reduced, - /* src */ curr_last_recurrent_state_reduced, - /* N */ qk_head_size, - /* K */ v_head_size, - /* ld_src */ v_head_size, - /* ld_dst */ v_head_size); + /* dst */ k_beta_packed, + /* src */ k_beta, + /* K */ mb_size, + /* N */ D, + /* ld_src */ D, + /* ld_dst */ D); - // v_prime = k_cumdecay_i @ curr_last_recurrent_state: [chunk_size, EV] - // k_cumdecay_i: [chunk_size, EK] - // curr_last_recurrent_state: [EK, EV] + // 5.c w = attn2 @ k_beta at::native::cpublas::brgemm( - /* M */ chunk_size, - /* N */ v_head_size, - /* K */ qk_head_size, - /* lda */ qk_head_size, - /* ldb */ v_head_size, - /* ldc */ v_head_size, + /* M */ mb_size, + /* N */ D, + /* K */ padded_mb_size, // mb_size + /* lda */ CHUNK_SIZE, + /* ldb */ D, + /* ldc */ D, /* add_C */ false, - /* A */ k_cumdecay_i_reduced, - /* B */ curr_last_recurrent_state_pack_reduced, + /* A */ attn2, + /* B */ k_beta_packed, + /* C */ k_updated); + + // 5.d k_updated -> w + scalar_t* __restrict__ w_ptr = w + (batch_offset + mb_start) * w_strideT + hv * w_strideH; + update_kernel::apply(w_ptr, k_updated, mb_size, D, w_strideT); + + // 5.e value = value * beta + apply_beta_kernel::apply( + v_beta, v_ptr, beta_ptr, nullptr, mb_size, v_strideT, D, Hv); + + // 5.f pack value + pack_vnni2( + /* dst */ v_beta_packed, + /* src */ v_beta, + /* K */ mb_size, + /* N */ D, + /* ld_src */ D, + /* ld_dst */ D); + + // 5.g u = attn2 @ v_beta + at::native::cpublas::brgemm( + /* M */ mb_size, + /* N */ D, + /* K */ padded_mb_size, // mb_size + /* lda */ CHUNK_SIZE, + /* ldb */ D, + /* ldc */ D, + /* add_C */ false, + /* A */ attn2, + /* B */ v_beta_packed, + /* C */ v_updated); + + // 5.h v_updated -> u + scalar_t* __restrict__ u_ptr = u + (batch_offset + mb_start) * u_strideT + hv * u_strideH; + update_kernel::apply(u_ptr, v_updated, mb_size, D, u_strideT); + } + + // move to the next index + data_index_step(nt, NT, h, H); + } + at::native::cpublas::brgemm_release(); + }); +} + +// +// out : [B, T, Hv, Dv] +// state : [num_seqs, Hv, Dv, D] +// q : [B, T, H, D] +// k : [B, T, H, D] +// w : [B, T, Hv, D] +// u : [B, T, Hv, Dv] +// g : [B, NT, Hv, C] +// d : [B, NT, Hv, C, C] +// cu_seqlens : [num_seqs + 1] +// chunk_offsets : [num_seqs + 1] +template +void chunk_gated_delta_rule_fwd_inter_kernel_impl( + scalar_t* __restrict__ out, + float* __restrict__ state, + const scalar_t* __restrict__ q, + const scalar_t* __restrict__ k, + const scalar_t* __restrict__ w, + const scalar_t* __restrict__ u, + const float* __restrict__ g, + const float* __restrict__ d, + const int32_t* __restrict__ cu_seqlens, + const int32_t* __restrict__ chunk_offsets, + int64_t H, + int64_t Hv, + int64_t num_seqs, + int64_t q_strideT, + int64_t q_strideH, + int64_t k_strideT, + int64_t k_strideH) { + // head group, expect to be 1，2，4 for qwen3.5 + const int64_t HG = Hv / H; + + // strides + const int64_t w_strideT = Hv * D; + const int64_t w_strideH = D; + const int64_t u_strideT = Hv * D; + const int64_t u_strideH = D; + const int64_t o_strideT = Hv * D; + const int64_t o_strideH = D; + + // [NB]: parallel on [num_seqs, Hv] + // * choose to parallel on Hv instead of H, though this means q @ kT has duplicated compute + // * H might be 16 which is not enough to use 32C when num_seqs is small + at::parallel_for(0, num_seqs * Hv, 0, [&](int64_t begin, int64_t end) { + int64_t bs{0}, hv{0}; + data_index_init(begin, bs, num_seqs, hv, Hv); + + // thread local temp buffer + DECL_ZERO_BUF(scalar_t, tmp, CHUNK_SIZE * D); + DECL_ZERO_BUF(scalar_t, tmp2, D * D); + DECL_ZERO_BUF(float, tmp3, CHUNK_SIZE* D); + DECL_ZERO_BUF(scalar_t, tmp4, CHUNK_SIZE * D); + DECL_ZERO_BUF(float, attn, CHUNK_SIZE* CHUNK_SIZE); + DECL_ZERO_BUF(scalar_t, attn2, CHUNK_SIZE * CHUNK_SIZE); + + // alias + scalar_t* __restrict__ k_packed = tmp; + scalar_t* __restrict__ s_packed = tmp2; + float* __restrict__ v_prime = tmp3; + scalar_t* __restrict__ v_prime2 = tmp; + float* __restrict__ attn_inter = tmp3; + scalar_t* __restrict__ qg_exp = tmp4; + scalar_t* __restrict__ v_packed = tmp4; + scalar_t* __restrict__ k_updated = tmp; + + for (int64_t i = begin; i < end; ++i) { + int64_t h = hv / HG; + int32_t batch_offset = cu_seqlens[bs]; + int32_t seqlen = cu_seqlens[bs + 1] - cu_seqlens[bs]; + int64_t nt = chunk_offsets[bs]; + + for (int64_t mb_start = 0; mb_start < seqlen; mb_start += CHUNK_SIZE, ++nt) { + int64_t mb_size = std::min(seqlen - mb_start, int64_t(CHUNK_SIZE)); + + // mb_size` is K in 4.a, pad to TILE_K; + const int64_t padded_mb_size = div_up((int)mb_size, TILE_K) * TILE_K; + + // step 1.a: attn = query @ key^T + // attn_i = (q_i @ k_i.transpose(-1, -2) * decay_mask[:, :, i]).masked_fill_(mask, 0) + const scalar_t* __restrict__ q_ptr = q + (batch_offset + mb_start) * q_strideT + h * q_strideH; + const scalar_t* __restrict__ k_ptr = k + (batch_offset + mb_start) * k_strideT + h * k_strideH; + pack_vnni( + /* dst */ k_packed, + /* src */ k_ptr, + /* N */ mb_size, + /* K */ D, + /* ld_src */ k_strideT, + /* ld_dst */ CHUNK_SIZE); + + at::native::cpublas::brgemm( + /* M */ mb_size, + /* N */ mb_size, + /* K */ D, + /* lda */ q_strideT, + /* ldb */ CHUNK_SIZE, + /* ldc */ CHUNK_SIZE, + /* add_C */ false, + /* A */ q_ptr, + /* B */ k_packed, + /* C */ attn); + + // step 1.b: attn = attn * decay_mask.masked_fill_(mask, 0) + const float* __restrict__ d_ptr = d + nt * (Hv * CHUNK_SIZE * CHUNK_SIZE) + hv * (CHUNK_SIZE * CHUNK_SIZE); + apply_mask_kernel::apply(attn2, attn, nullptr, d_ptr, mb_size); + + // step 2.a: v' = w @ state (fuse state *= exp(g_last) with packing) + float* __restrict__ s_ptr = state + bs * (Hv * D * D) + hv * (D * D); + const float* __restrict__ g_ptr = g + nt * (Hv * CHUNK_SIZE) + hv * (CHUNK_SIZE); + float g_last = g_ptr[mb_size - 1]; + pack_vnni2( + /* dst */ s_packed, + /* src */ s_ptr, + /* g_last */ g_last, + /* ld_src */ D, + /* ld_dst */ D); + + const scalar_t* __restrict__ w_ptr = w + (batch_offset + mb_start) * w_strideT + hv * w_strideH; + at::native::cpublas::brgemm( + /* M */ mb_size, + /* N */ D, + /* K */ D, + /* lda */ w_strideT, + /* ldb */ D, + /* ldc */ D, + /* add_C */ false, + /* A */ w_ptr, + /* B */ s_packed, /* C */ v_prime); - // v_new = v_prime = v_i - v_prime - // v_i: [chunk_size, EV] - for (int64_t m = 0; m < chunk_size; m++) { - int64_t i = 0; - for (; i < fVecSize * (v_head_size / fVecSize); i += fVecSize) { - auto tmp0 = fVec::loadu(v_i + m * v_head_size + i); - auto tmp1 = fVec::loadu(v_prime + m * v_head_size + i); - auto tmp2 = tmp0 - tmp1; - auto tmp3 = at::vec::convert(tmp2); - tmp3.store(v_prime_reduced + m * v_head_size + i, fVecSize); - } - } + // step 2.b: v2' = u - v' + const scalar_t* __restrict__ u_ptr = u + (batch_offset + mb_start) * u_strideT + hv * u_strideH; + update_value_kernel::apply(v_prime2, u_ptr, v_prime, mb_size, padded_mb_size, u_strideT); - // attn_inter = (q_i * g[:, :, i, :, None].exp()) @ last_recurrent_state - // qg = q_i * g[:, :, i, :, None].exp(): [chunk_size, EK] - // q_i: [chunk_size, EK] - // g[:, :, i, :, None]: [chunk_size, 1] - for (int64_t m = 0; m < chunk_size; m++) { - auto g_pad_i_m = g_pad_i + m; - auto g_exp = std::exp(*g_pad_i_m); - int64_t i = 0; - scalar_t g_exp_reduced = static_cast(g_exp); - auto vec_g_exp_reduced = bVec(g_exp_reduced); - for (; i < VecSize * (qk_head_size / VecSize); i += VecSize) { - auto tmp0 = bVec::loadu(q_i + m * qk_head_size + i); - auto tmp2 = tmp0 * vec_g_exp_reduced; - tmp2.store(qg + m * qk_head_size + i); - } - } - // attn_inter = qg @ curr_last_recurrent_state: [chunk_size, EV] - // curr_last_recurrent_state: [EK, EV] + // step 3.a: qg_exp = q * exp(g) + apply_beta_kernel::apply( + qg_exp, q_ptr, nullptr, g_ptr, mb_size, q_strideT, D, /*b_stride*/ 0); + + // step 3.b: attn_inter = qg_exp @ state at::native::cpublas::brgemm( - /* M */ chunk_size, - /* N */ v_head_size, - /* K */ qk_head_size, - /* lda */ qk_head_size, - /* ldb */ v_head_size, - /* ldc */ v_head_size, + /* M */ mb_size, + /* N */ D, + /* K */ D, + /* lda */ D, + /* ldb */ D, + /* ldc */ D, /* add_C */ false, - /* A */ qg, - /* B */ curr_last_recurrent_state_pack_reduced, - /* C */ attn_inter); + /* A */ qg_exp, + /* B */ s_packed, + /* C */ attn_inter); - // core_attn_out[:, :, i] = attn_inter + attn_i @ v_new - // pack for v_prime + // step 4.a: attn_inter += attn2 @ v2' pack_vnni2( - /* dst */ v_prime_pack_reduced, - /* src */ v_prime_reduced, - /* N */ chunk_size, - /* K */ v_head_size, - /* ld_src */ v_head_size, - /* ld_dst */ v_head_size); - // attn_inter = attn_inter + attn_i @ v_new: [chunk_size, EV] - // attn_i: [chunk_size, chunk_size] - // v_new: [chunk_size, EV] + /* dst */ v_packed, + /* src */ v_prime2, + /* K */ padded_mb_size, + /* N */ D, + /* ld_src */ D, + /* ld_dst */ D); + at::native::cpublas::brgemm( - /* M */ chunk_size, - /* N */ v_head_size, - /* K */ chunk_size, - /* lda */ chunk_size, - /* ldb */ v_head_size, - /* ldc */ v_head_size, + /* M */ mb_size, + /* N */ D, + /* K */ padded_mb_size, + /* lda */ CHUNK_SIZE, + /* ldb */ D, + /* ldc */ D, /* add_C */ true, - /* A */ attn_i_reduced, - /* B */ v_prime_pack_reduced, - /* C */ attn_inter); + /* A */ attn2, + /* B */ v_packed, + /* C */ attn_inter); - // core_attn_out[:, :, i] = attn_inter - for (int64_t m = 0; m < chunk_size; m++) { - at::vec::map( - [](fVec x) { return x; }, core_attn_out_i + m * v_head_size, attn_inter + m * v_head_size, v_head_size); - } + // step 4.b: write attn_inter -> out + scalar_t* __restrict__ o_ptr = out + (batch_offset + mb_start) * o_strideT + hv * o_strideH; + update_kernel::apply(o_ptr, attn_inter, mb_size, D, o_strideT); - // last_recurrent_state = ( - // last_recurrent_state * g[:, :, i, -1, None, None].exp() - // + (k_i * (g[:, :, i, -1, None] - g[:, :, i]).exp()[..., None]).transpose(-1, -2) @ v_new - // ) - // 1) last_recurrent_state * g[:, :, i, -1, None, None].exp() - // curr_last_recurrent_state: [EK, EV] - // g[:, :, i, -1, None, None]: [1, 1] - // last_recurrent_state * g[:, :, i, -1, None, None].exp(): [EK, EV] - auto g_pad_i_last = g_pad_i + chunk_size - 1; - auto g_exp_last = std::exp(g_pad_i_last[0]); - for (int64_t m = 0; m < qk_head_size; m++) { - int64_t i = 0; - auto vec_g_exp_last = fVec(g_exp_last); - for (; i < fVecSize * (v_head_size / fVecSize); i += fVecSize) { - auto tmp0 = bVec::loadu(curr_last_recurrent_state_reduced + m * v_head_size + i); - auto tmp1 = at::vec::convert(tmp0); - auto tmp2 = tmp1 * vec_g_exp_last; - tmp2.store(curr_last_recurrent_state + m * v_head_size + i); - } - if (i < v_head_size) { - auto tmp0 = bVec::loadu(curr_last_recurrent_state_reduced + m * v_head_size + i, v_head_size - i); - auto tmp1 = at::vec::convert(tmp0); - auto tmp2 = tmp1 * vec_g_exp_last; - tmp2.store(curr_last_recurrent_state + m * v_head_size + i, v_head_size - i); - } - } - // 2) (k_i * (g[:, :, i, -1, None] - g[:, :, i]).exp()[..., None]).transpose(-1, -2) @ v_new - // k_i: [chunk_size, EK] - // g[:, :, i, -1, None]: [1] - // g[:, :, i]: [chunk_size] - // (g[:, :, i, -1, None] - g[:, :, i]).exp()[..., None]: [chunk_size, 1] - // kg = k_i * (g[:, :, i, -1, None] - g[:, :, i]).exp()[..., None]: [chunk_size, EK] - // (k_i * (g[:, :, i, -1, None] - g[:, :, i]).exp()[..., None]).transpose(-1, -2): [EK, chunk_size] - // v_new: [chunk_size, EV] - // (k_i * (g[:, :, i, -1, None] - g[:, :, i]).exp()[..., None]).transpose(-1, -2) @ v_new: [EK, EV] - // kg = k_i * (g[:, :, i, -1, None] - g[:, :, i]).exp()[..., None] - for (int64_t m = 0; m < chunk_size; m++) { - auto g_exp = std::exp((g_pad_i_last[0] - g_pad_i[m])); - int64_t i = 0; - scalar_t g_exp_reduced = static_cast(g_exp); - auto vec_g_exp_reduced = bVec(g_exp_reduced); - for (; i < VecSize * (qk_head_size / VecSize); i += VecSize) { - auto tmp0 = bVec::loadu(k_i + m * qk_head_size + i); - auto tmp2 = tmp0 * vec_g_exp_reduced; - tmp2.store(kg + m * qk_head_size + i); - } - } - // kg.transpose(-1, -2): [EK, chunk_size] - at::native::utils::transpose( - /* M */ chunk_size, - /* N */ qk_head_size, - /* src */ kg, - /* ld_src */ qk_head_size, - /* dst */ kg_transpose, - /* ld_dst */ chunk_size); - // kgv = kg.transpose(-1, -2) @ v_new - // v_new: [chunk_size, EV] + // step 5: update state + // state_new = state * exp(g_last) + (k * exp(g_last - g)).T @ v2' + + // step 5.1 state *= exp(g_last) fused with step 2.a + + // step 5.2 k' = k * exp(g_last - g).T; TODO: fuse this with 1.a + update_key_kernel::apply(k_updated, k_ptr, g_ptr, mb_size, k_strideT); + + // step 5.3 state += k' @ v2' at::native::cpublas::brgemm( - /* M */ qk_head_size, - /* N */ v_head_size, - /* K */ chunk_size, - /* lda */ chunk_size, - /* ldb */ v_head_size, - /* ldc */ v_head_size, - /* add_C */ false, - /* A */ kg_transpose, - /* B */ v_prime_pack_reduced, - /* C */ kgv); - // last_recurrent_state = 1) + 2) - for (int64_t m = 0; m < qk_head_size; m++) { - at::vec::map2( - [](fVec x, fVec y) { return x + y; }, - curr_last_recurrent_state + m * v_head_size, - curr_last_recurrent_state + m * v_head_size, - kgv + m * v_head_size, - v_head_size); - } + /* M */ D, + /* N */ D, + /* K */ padded_mb_size, // mb_size + /* lda */ CHUNK_SIZE, + /* ldb */ D, + /* ldc */ D, + /* add_C */ true, + /* A */ k_updated, + /* B */ v_packed, + /* C */ s_ptr); } - // core_attn_out -> output - // output: [B, T, HV, EV] - // core_attn_out: [B, HV, padded_T, EV] - auto curr_out = out_ptr + h * oStrideH; - for (int64_t m = 0; m < seq_len; m++) { - at::vec::map( - [](fVec x) { return x; }, curr_out + m * oStrideT, curr_core_attn_out + m * v_head_size, v_head_size); - } - - // Move to the next query - data_index_step(b, batch_size, h, v_num_head); + // move to the next index + data_index_step(bs, num_seqs, hv, Hv); } + at::native::cpublas::brgemm_release(); }); } @@ -1072,29 +1332,205 @@ void fused_gdn_gating_kernel_impl( } // anonymous namespace -template -inline void -CHECK_INPUT_SHAPE_DTYPE(const at::Tensor& tensor, const int64_t& dim, const at::IntArrayRef& sizes, at::ScalarType st) { - TORCH_CHECK(tensor.sizes() == sizes, "Input tensor shape mismatch: expected ", sizes, ", got ", tensor.sizes()); - TORCH_CHECK(tensor.dtype() == st, "Input tensor dtype mismatch"); - CHECK_DIM(dim, tensor); - if (is_last_dim_contiguous) { - CHECK_LAST_DIM_CONTIGUOUS_INPUT(tensor); - } else { - CHECK_CONTIGUOUS(tensor); +template +std::tuple prepare_chunk_indices(const at::Tensor& cu_seqlens) { + int64_t num_seqs = cu_seqlens.size(0) - 1; + at::Tensor chunk_offsets = at::empty({num_seqs + 1}, cu_seqlens.options()); + // get number of chunks and chunk offsets + const int32_t* offsets_data = cu_seqlens.data_ptr(); + int32_t num_chunks = 0; + chunk_offsets[0] = 0; + for (int64_t row = 0; row < num_seqs; ++row) { + num_chunks += div_up(offsets_data[row + 1] - offsets_data[row], CHUNK_SIZE); + chunk_offsets[row + 1] = num_chunks; } + // get chunk indices + at::Tensor chunk_indices = at::empty({num_chunks, 2}, cu_seqlens.options()); + int32_t* indices_data = chunk_indices.data_ptr(); + + int64_t idx = 0; + for (int32_t row = 0; row < num_seqs; ++row) { + int32_t num_chunks = div_up(offsets_data[row + 1] - offsets_data[row], CHUNK_SIZE); + + for (int32_t col = 0; col < num_chunks; ++col) { + indices_data[idx * 2 + 0] = row; + indices_data[idx * 2 + 1] = col; + idx++; + } + } + return std::make_tuple(chunk_indices, chunk_offsets); } -// query: [B, T, HK, EK] -// key: [B, T, HK, EK] -// value: [B, T, HV, EV] -// g: [B, T, HV] FP32 -// beta: [B, T, HV] -// initial_state: [N, HV, EK, EV] FP32 -// output_final_state: bool -// cu_seqlens: [N + 1] INT32 -// head_first: bool -// use_qk_l2norm_in_kernel: bool +#define DISPATCH_HEAD_DIM_CASE(launch_macro, hd) \ + case hd: { \ + launch_macro(hd); \ + break; \ + } + +// [NB]: add new head_dim support here +#define DISPATCH_HEAD_DIM(dim, launch_macro) \ + switch (dim) { \ + DISPATCH_HEAD_DIM_CASE(launch_macro, 64) \ + DISPATCH_HEAD_DIM_CASE(launch_macro, 128) \ + default: \ + TORCH_CHECK(false, "Unexpected head dim size, ", dim); \ + } + +#define LAUNCH_L2NORM_KERNEL(HD) \ + l2norm_fwd_kernel_impl( \ + query_norm.data_ptr(), \ + key_norm.data_ptr(), \ + query.data_ptr(), \ + key.data_ptr(), \ + eps, \ + T, \ + H, \ + query.stride(1), \ + query.stride(2), \ + key.stride(1), \ + key.stride(2)); + +std::tuple l2norm_fwd(const at::Tensor& query, const at::Tensor& key, double eps) { + int64_t B = query.size(0); + int64_t T = query.size(1); + int64_t H = query.size(2); + int64_t D = query.size(3); + + at::Tensor query_norm = at::empty_like(query); + at::Tensor key_norm = at::empty_like(key); + + AT_DISPATCH_REDUCED_FLOATING_TYPES( + query.scalar_type(), "l2norm_fwd", [&] { DISPATCH_HEAD_DIM(D, LAUNCH_L2NORM_KERNEL); }); + + return std::make_tuple(query_norm, key_norm); +} + +// [NB]: instantiate decay_mask to avoid heavy recomputation in the kernel with exp +template +at::Tensor chunk_local_cumsum(const at::Tensor& g, const at::Tensor& cu_seqlens, const at::Tensor& chunk_indices) { + int64_t B = g.size(0); + // int64_t T = g.size(1); + int64_t Hv = g.size(2); + int64_t NT = chunk_indices.size(0); + + at::Tensor g_ = at::empty({B, NT, Hv, CHUNK_SIZE}, g.options()); + AT_DISPATCH_FLOATING_TYPES(g.scalar_type(), "chunk_local_cumsum", [&] { + chunk_local_cumsum_kernel_impl( + g_.data_ptr(), + g.data_ptr(), + cu_seqlens.data_ptr(), + chunk_indices.data_ptr(), + Hv, + NT); + }); + return g_; +} + +#define LAUNCH_CHUNK_GATED_DELTA_RULE_FWD_INTRA_KERNEL(HD) \ + chunk_gated_delta_rule_fwd_intra_kernel_impl( \ + w.data_ptr(), \ + u.data_ptr(), \ + decay_mask.data_ptr(), \ + k.data_ptr(), \ + v.data_ptr(), \ + g.data_ptr(), \ + beta.data_ptr(), \ + cu_seqlens.data_ptr(), \ + chunk_indices.data_ptr(), \ + H, \ + Hv, \ + NT, \ + k.stride(1), \ + k.stride(2), \ + v.stride(1), \ + v.stride(2)); + +template +std::tuple chunk_gated_delta_rule_fwd_intra( + const at::Tensor& k, + const at::Tensor& v, + const at::Tensor& g, + const at::Tensor& beta, + const at::Tensor& cu_seqlens, + const at::Tensor& chunk_indices) { + int64_t B = k.size(0); + int64_t T = k.size(1); + int64_t H = k.size(2); + int64_t D = k.size(3); + int64_t Hv = v.size(2); + int64_t Dv = v.size(3); + int64_t NT = chunk_indices.size(0); + + at::Tensor w = at::empty({B, T, Hv, D}, k.options()); // BFloat16 + at::Tensor u = at::empty({B, T, Hv, Dv}, k.options()); // BFloat16 + at::Tensor decay_mask = at::empty({B, NT, Hv, CHUNK_SIZE, CHUNK_SIZE}, g.options()); // Float + AT_DISPATCH_REDUCED_FLOATING_TYPES(k.scalar_type(), "chunk_gated_delta_rule_fwd_intra", [&] { + DISPATCH_HEAD_DIM(D, LAUNCH_CHUNK_GATED_DELTA_RULE_FWD_INTRA_KERNEL); + }); + + return std::make_tuple(w, u, decay_mask); +} + +#define LAUNCH_CHUNK_GATED_DELTA_RULE_FWD_INTER_KERNEL(HD) \ + chunk_gated_delta_rule_fwd_inter_kernel_impl( \ + o.data_ptr(), \ + initial_state.data_ptr(), \ + q.data_ptr(), \ + k.data_ptr(), \ + w.data_ptr(), \ + u.data_ptr(), \ + g.data_ptr(), \ + decay_mask.data_ptr(), \ + cu_seqlens.data_ptr(), \ + chunk_offsets.data_ptr(), \ + H, \ + Hv, \ + num_seqs, \ + q.stride(1), \ + q.stride(2), \ + k.stride(1), \ + k.stride(2)); + +template +std::tuple chunk_gated_delta_rule_fwd_inter( + const at::Tensor& q, + const at::Tensor& k, + const at::Tensor& w, + const at::Tensor& u, + const at::Tensor& g, + const at::Tensor& decay_mask, + const at::Tensor& initial_state, + bool output_final_state, + const at::Tensor& cu_seqlens, + const at::Tensor& chunk_offsets) { + const int64_t B = q.size(0); + const int64_t T = q.size(1); + const int64_t H = q.size(2); + const int64_t D = q.size(3); + const int64_t Hv = w.size(2); + const int64_t Dv = u.size(3); + const int64_t num_seqs = initial_state.size(0); + + at::Tensor o = at::empty({B, T, Hv, Dv}, q.options()); + AT_DISPATCH_REDUCED_FLOATING_TYPES(q.scalar_type(), "chunk_gated_delta_rule_fwd_inter", [&] { + DISPATCH_HEAD_DIM(D, LAUNCH_CHUNK_GATED_DELTA_RULE_FWD_INTER_KERNEL); + }); + + return std::make_tuple(o, initial_state); +} + +// [NB]: Support only varlen inputs +// B: packed batch dim of q/k/v (== 1) +// num_seqs: number of variable-length sequences +// +// query: [B, T, H, D] +// key: [B, T, H, D] +// value: [B, T, Hv, Dv] +// g: [B, T, Hv] FP32 +// beta: [B, T, Hv] +// initial_state: [num_seqs, Hv, Dv, D] FP32 +// cu_seqlens: [num_seqs + 1] INT32 +// std::tuple chunk_gated_delta_rule_cpu( const at::Tensor& query, const at::Tensor& key, @@ -1106,154 +1542,49 @@ std::tuple chunk_gated_delta_rule_cpu( const at::Tensor& cu_seqlens, bool head_first, bool use_qk_l2norm_in_kernel, - double eps = 1e-5) { - TORCH_CHECK(head_first == false, "chunk_gated_delta_rule_cpu does not support head first"); + double eps = 1e-6) { + TORCH_CHECK(!head_first, "chunk_gated_delta_rule_cpu: does not support head first"); + int64_t B = query.size(0); - int64_t global_seq_len = query.size(1); - int64_t qk_num_head = query.size(2); - int64_t qk_head_size = query.size(3); - int64_t v_num_head = value.size(2); - int64_t v_head_size = value.size(3); - int64_t batch_size = initial_state.size(0); - CHECK_EQ(B, 1); - TORCH_CHECK(v_num_head % qk_num_head == 0, "expect v_num_head multiple of qk_num_head."); - TORCH_CHECK(qk_head_size % 32 == 0, "expect qk_head_size to be multiples of 32."); - TORCH_CHECK(v_head_size % 32 == 0, "expect v_head_size to be multiples of 32."); - CHECK_INPUT_SHAPE_DTYPE(query, 4, {B, global_seq_len, qk_num_head, qk_head_size}, at::kBFloat16); - CHECK_INPUT_SHAPE_DTYPE(key, 4, {B, global_seq_len, qk_num_head, qk_head_size}, at::kBFloat16); - CHECK_INPUT_SHAPE_DTYPE(value, 4, {B, global_seq_len, v_num_head, v_head_size}, at::kBFloat16); - CHECK_INPUT_SHAPE_DTYPE(g, 3, {B, global_seq_len, v_num_head}, at::kFloat); - CHECK_INPUT_SHAPE_DTYPE(beta, 3, {B, global_seq_len, v_num_head}, at::kBFloat16); - CHECK_INPUT_SHAPE_DTYPE(cu_seqlens, 1, {batch_size + 1}, at::kInt); - CHECK_INPUT_SHAPE_DTYPE(initial_state, 4, {batch_size, v_num_head, qk_head_size, v_head_size}, at::kFloat); + int64_t T = query.size(1); + int64_t H = query.size(2); + int64_t D = query.size(3); + int64_t Hv = value.size(2); + int64_t Dv = value.size(3); + int64_t num_seqs = initial_state.size(0); - at::Tensor output = at::empty_like(value, value.options()); // [B, T, HV, EV] - at::Tensor final_state = initial_state.to(at::kFloat); // [N, HV, EK, EV] + TORCH_CHECK(B == 1, __func__, ": expect batch size to be 1"); + TORCH_CHECK(Hv % H == 0, __func__, ": expect num_heads_kv multiple of num_heads."); + TORCH_CHECK(D % 32 == 0, __func__, ": expect head_dim to be multiples of 32."); + TORCH_CHECK(Dv % 32 == 0, __func__, ": expect head_dim_v to be multiples of 32."); + TORCH_CHECK(D == Dv, __func__, ": expect head_dim to be equal to head_dim_v."); + CHECK_INPUT_SHAPE_DTYPE(query, {B, T, H, D}, at::kBFloat16); + CHECK_INPUT_SHAPE_DTYPE(key, {B, T, H, D}, at::kBFloat16); + CHECK_INPUT_SHAPE_DTYPE(value, {B, T, Hv, Dv}, at::kBFloat16); + CHECK_INPUT_SHAPE_DTYPE(g, {B, T, Hv}, at::kFloat); + CHECK_INPUT_SHAPE_DTYPE(beta, {B, T, Hv}, at::kBFloat16); + CHECK_INPUT_SHAPE_DTYPE(cu_seqlens, {num_seqs + 1}, at::kInt); + CHECK_INPUT_SHAPE_DTYPE(initial_state, {num_seqs, Hv, Dv, D}, at::kFloat); - // Strides - int64_t qStrideH = query.stride(2); - int64_t qStrideT = query.stride(1); - int64_t kStrideH = key.stride(2); - int64_t kStrideT = key.stride(1); - int64_t vStrideH = value.stride(2); - int64_t vStrideT = value.stride(1); - int64_t oStrideH = output.stride(2); - int64_t oStrideT = output.stride(1); + constexpr int CHUNK_SIZE = 64; - constexpr int64_t chunk_size = 64; - // Deduce the global chunks - // e.g. cu_seqlens: [0, 5, 13, 16], chunk_size = 4 - // chunk_offsets: [0, 2, 4, 5] - // chunk_indices (batch_id, local_chunk_id): [[0, 0], [0, 1], [1, 0], [1, 1], [2, 0]] - at::Tensor chunk_offsets = at::empty(batch_size + 1, cu_seqlens.options()); - auto chunk_offsets_ptr = chunk_offsets.data_ptr(); - chunk_offsets_ptr[0] = 0; - int32_t* cu_seqlens_ptr = cu_seqlens.data_ptr(); - int64_t s = 0; - int64_t e = 0; - int64_t s_pad = 0; - int64_t e_pad = 0; - for (int64_t b = 0; b < batch_size; b++) { - e = cu_seqlens_ptr[b + 1]; - int64_t seq_len = e - s; - int64_t pad_size = (chunk_size - seq_len % chunk_size) % chunk_size; - int64_t total_seq_length = seq_len + pad_size; - e_pad = s_pad + total_seq_length; - chunk_offsets[b + 1] = e_pad / chunk_size; - s = e; - s_pad = e_pad; - } - int64_t global_total_seq_length = e_pad; - int64_t global_num_chunk = chunk_offsets_ptr[batch_size]; - at::Tensor chunk_indices = at::empty(global_num_chunk * 2, cu_seqlens.options()); - auto chunk_indices_ptr = chunk_indices.data_ptr(); - int64_t curr_c = 0; - for (int64_t b = 0; b < batch_size; b++) { - int64_t batch_chunk_num = chunk_offsets_ptr[b + 1] - chunk_offsets_ptr[b]; - for (int64_t c = 0; c < batch_chunk_num; c++) { - chunk_indices_ptr[curr_c * 2] = b; - chunk_indices_ptr[curr_c * 2 + 1] = c; - curr_c += 1; - } - } + // prepare chunk indices + auto [chunk_indices, chunk_offsets] = prepare_chunk_indices(cu_seqlens); - // Allocate buffer - int64_t buff_size = v_num_head * global_total_seq_length // g_pad_data - + batch_size * v_num_head * global_total_seq_length * v_head_size // core_attn - + v_num_head * global_total_seq_length * chunk_size // decay_mask - + v_num_head * global_total_seq_length * v_head_size; // v_beta_attn - at::Tensor buff_data = at::empty({buff_size}, query.options().dtype(at::kFloat)); - int64_t reduced_buff_size = qk_num_head * global_total_seq_length * qk_head_size // q_pad_data - + qk_num_head * global_total_seq_length * qk_head_size // k_pad_data - + v_num_head * global_total_seq_length * v_head_size // v_pad_data - + v_num_head * global_total_seq_length * qk_head_size // k_beta_data - + v_num_head * global_total_seq_length * v_head_size // v_beta_data - + v_num_head * global_total_seq_length * qk_head_size // k_cumdecay_reduced - + qk_num_head * global_seq_len // q_norm_sum - + qk_num_head * global_seq_len; // k_norm_sum - at::Tensor reduced_buff_data = at::empty({reduced_buff_size}, query.options()); - int64_t num_thread = at::get_num_threads(); - int64_t buff_size_16bit_per_thread = - /* k_transpose */ qk_head_size * chunk_size + - /* v_pack */ chunk_size * v_head_size + - /* k_beta_g */ chunk_size * qk_head_size + - /* k_beta_g_pack */ chunk_size * qk_head_size + - /* attn */ chunk_size * chunk_size * 2 + - /* attn_reduced */ chunk_size * chunk_size + - /* k_cumdecay */ chunk_size * qk_head_size * 2 + - /* row */ chunk_size * 2 + - /* updated */ chunk_size * 2 + - /* curr_last_recurrent_state_reduced */ qk_head_size * v_head_size + - /* curr_last_recurrent_state_pack_reduced */ qk_head_size * v_head_size + - /* k_transpose_i */ qk_head_size * chunk_size + - /* attn_i */ chunk_size * chunk_size * 2 + - /* attn_i_reduced */ chunk_size * chunk_size + - /* v_prime */ chunk_size * v_head_size * 2 + - /* v_prime_reduced */ chunk_size * v_head_size + - /* v_prime_pack_reduced */ chunk_size * v_head_size + - /* qg */ chunk_size * qk_head_size + - /* attn_inter */ chunk_size * v_head_size * 2 + - /* kg */ chunk_size * qk_head_size + - /* kg_transpose */ qk_head_size * chunk_size + - /* kgv */ qk_head_size * v_head_size * 2; - at::Tensor thread_buff_data = at::empty({num_thread, buff_size_16bit_per_thread}, query.options()); + float scale = 1.0 / std::sqrt(D); + auto [query_, key_] = use_qk_l2norm_in_kernel ? l2norm_fwd(query, key, eps) : std::make_tuple(query.mul(scale), key); - AT_DISPATCH_REDUCED_FLOATING_TYPES(query.scalar_type(), "chunk_gated_delta_rule_kernel", [&] { - chunk_gated_delta_rule_kernel_impl( - output.data_ptr(), - final_state.data_ptr(), - query.data_ptr(), - key.data_ptr(), - value.data_ptr(), - g.data_ptr(), - beta.data_ptr(), - cu_seqlens_ptr, - buff_data.data_ptr(), - reduced_buff_data.data_ptr(), - thread_buff_data.data_ptr(), - chunk_offsets_ptr, - chunk_indices_ptr, - use_qk_l2norm_in_kernel, - batch_size, - global_seq_len, - qk_num_head, - v_num_head, - qk_head_size, - v_head_size, - qStrideH, - qStrideT, - kStrideH, - kStrideT, - vStrideH, - vStrideT, - oStrideH, - oStrideT, - global_total_seq_length, - global_num_chunk, - buff_size_16bit_per_thread, - eps); - }); - return std::make_tuple(std::move(output), std::move(final_state)); + auto g_ = chunk_local_cumsum(g, cu_seqlens, chunk_indices); + + // fused kkt + solve_tril + recompute_w_u + auto [w, u, decay_mask] = + chunk_gated_delta_rule_fwd_intra(key_, value, g_, beta, cu_seqlens, chunk_indices); + + // fused `chunk_gated_delta_rule_fwd_h` + `chunk_fwd_o` + auto [output, final_state] = chunk_gated_delta_rule_fwd_inter( + query_, key_, w, u, g_, decay_mask, initial_state, output_final_state, cu_seqlens, chunk_offsets); + + return std::make_tuple(output, final_state); } // A_log: [v_num_heads] diff --git a/sgl-kernel/csrc/cpu/vec.h b/sgl-kernel/csrc/cpu/vec.h index f1de4f9ca..a235cb4a7 100644 --- a/sgl-kernel/csrc/cpu/vec.h +++ b/sgl-kernel/csrc/cpu/vec.h @@ -287,6 +287,77 @@ inline void quantize_row_int8( // transpose utils // taken from my PR in ggml: https://github.com/ggml-org/llama.cpp/pull/8998 #if defined(CPU_CAPABILITY_AVX512) +inline void transpose_16x16_16bit(__m256i* v) { + __m256i v1[16]; + v1[0] = _mm256_unpacklo_epi16(v[0], v[1]); + v1[1] = _mm256_unpackhi_epi16(v[0], v[1]); + v1[2] = _mm256_unpacklo_epi16(v[2], v[3]); + v1[3] = _mm256_unpackhi_epi16(v[2], v[3]); + v1[4] = _mm256_unpacklo_epi16(v[4], v[5]); + v1[5] = _mm256_unpackhi_epi16(v[4], v[5]); + v1[6] = _mm256_unpacklo_epi16(v[6], v[7]); + v1[7] = _mm256_unpackhi_epi16(v[6], v[7]); + v1[8] = _mm256_unpacklo_epi16(v[8], v[9]); + v1[9] = _mm256_unpackhi_epi16(v[8], v[9]); + v1[10] = _mm256_unpacklo_epi16(v[10], v[11]); + v1[11] = _mm256_unpackhi_epi16(v[10], v[11]); + v1[12] = _mm256_unpacklo_epi16(v[12], v[13]); + v1[13] = _mm256_unpackhi_epi16(v[12], v[13]); + v1[14] = _mm256_unpacklo_epi16(v[14], v[15]); + v1[15] = _mm256_unpackhi_epi16(v[14], v[15]); + + v[0] = _mm256_unpacklo_epi32(v1[0], v1[2]); + v[1] = _mm256_unpackhi_epi32(v1[0], v1[2]); + v[2] = _mm256_unpacklo_epi32(v1[1], v1[3]); + v[3] = _mm256_unpackhi_epi32(v1[1], v1[3]); + v[4] = _mm256_unpacklo_epi32(v1[4], v1[6]); + v[5] = _mm256_unpackhi_epi32(v1[4], v1[6]); + v[6] = _mm256_unpacklo_epi32(v1[5], v1[7]); + v[7] = _mm256_unpackhi_epi32(v1[5], v1[7]); + v[8] = _mm256_unpacklo_epi32(v1[8], v1[10]); + v[9] = _mm256_unpackhi_epi32(v1[8], v1[10]); + v[10] = _mm256_unpacklo_epi32(v1[9], v1[11]); + v[11] = _mm256_unpackhi_epi32(v1[9], v1[11]); + v[12] = _mm256_unpacklo_epi32(v1[12], v1[14]); + v[13] = _mm256_unpackhi_epi32(v1[12], v1[14]); + v[14] = _mm256_unpacklo_epi32(v1[13], v1[15]); + v[15] = _mm256_unpackhi_epi32(v1[13], v1[15]); + + v1[0] = _mm256_unpacklo_epi64(v[0], v[4]); + v1[1] = _mm256_unpackhi_epi64(v[0], v[4]); + v1[2] = _mm256_unpacklo_epi64(v[1], v[5]); + v1[3] = _mm256_unpackhi_epi64(v[1], v[5]); + v1[4] = _mm256_unpacklo_epi64(v[2], v[6]); + v1[5] = _mm256_unpackhi_epi64(v[2], v[6]); + v1[6] = _mm256_unpacklo_epi64(v[3], v[7]); + v1[7] = _mm256_unpackhi_epi64(v[3], v[7]); + v1[8] = _mm256_unpacklo_epi64(v[8], v[12]); + v1[9] = _mm256_unpackhi_epi64(v[8], v[12]); + v1[10] = _mm256_unpacklo_epi64(v[9], v[13]); + v1[11] = _mm256_unpackhi_epi64(v[9], v[13]); + v1[12] = _mm256_unpacklo_epi64(v[10], v[14]); + v1[13] = _mm256_unpackhi_epi64(v[10], v[14]); + v1[14] = _mm256_unpacklo_epi64(v[11], v[15]); + v1[15] = _mm256_unpackhi_epi64(v[11], v[15]); + + v[0] = _mm256_permute2x128_si256(v1[0], v1[8], 0x20); + v[1] = _mm256_permute2x128_si256(v1[1], v1[9], 0x20); + v[2] = _mm256_permute2x128_si256(v1[2], v1[10], 0x20); + v[3] = _mm256_permute2x128_si256(v1[3], v1[11], 0x20); + v[4] = _mm256_permute2x128_si256(v1[4], v1[12], 0x20); + v[5] = _mm256_permute2x128_si256(v1[5], v1[13], 0x20); + v[6] = _mm256_permute2x128_si256(v1[6], v1[14], 0x20); + v[7] = _mm256_permute2x128_si256(v1[7], v1[15], 0x20); + v[8] = _mm256_permute2x128_si256(v1[0], v1[8], 0x31); + v[9] = _mm256_permute2x128_si256(v1[1], v1[9], 0x31); + v[10] = _mm256_permute2x128_si256(v1[2], v1[10], 0x31); + v[11] = _mm256_permute2x128_si256(v1[3], v1[11], 0x31); + v[12] = _mm256_permute2x128_si256(v1[4], v1[12], 0x31); + v[13] = _mm256_permute2x128_si256(v1[5], v1[13], 0x31); + v[14] = _mm256_permute2x128_si256(v1[6], v1[14], 0x31); + v[15] = _mm256_permute2x128_si256(v1[7], v1[15], 0x31); +} + inline void transpose_16x16_32bit(__m512i* v) { __m512i v1[16]; v1[0] = _mm512_unpacklo_epi32(v[0], v[1]); diff --git a/test/registered/cpu/test_mamba.py b/test/registered/cpu/test_mamba.py index 053339c76..8e30afd64 100644 --- a/test/registered/cpu/test_mamba.py +++ b/test/registered/cpu/test_mamba.py @@ -12,6 +12,12 @@ register_cpu_ci(est_time=10, suite="base-b-test-cpu") torch.manual_seed(1234) +# [NB]: State-layout convention for this test file: +# - CPU kernel path in fla.cpp uses VK state layout, same as triton impl. +# - Torch naive reference follows KV semantics from: +# https://github.com/fla-org/flash-linear-attention/blob/main/fla/ops/gated_delta_rule/naive.py +# - Transposes in these tests only bridge VK (kernel-facing) and KV (ref-facing) views. + def l2norm(x: torch.Tensor, dim: int = -1, eps: float = 1e-6): """This function is intended to align with the l2norm implementation in the FLA library.""" @@ -119,12 +125,13 @@ def chunk_gated_delta_rule_update( g, # [B, T, HV] beta, # [B, T, HV] cu_seqlens, # [N+1] - initial_state, # [N, HV, K, V] + initial_state, # [N, HV, V, K] use_qk_l2norm_in_kernel, # True ): num_heads = query.shape[2] num_value_heads = value.shape[2] batch_size = initial_state.shape[0] + initial_state_kv = initial_state.transpose(-1, -2).contiguous() if num_value_heads // num_heads > 1: query = query.repeat_interleave(num_value_heads // num_heads, dim=2) key = key.repeat_interleave(num_value_heads // num_heads, dim=2) @@ -139,12 +146,12 @@ def chunk_gated_delta_rule_update( value=value[:, start_q:end_q, :, :], g=g[:, start_q:end_q, :], beta=beta[:, start_q:end_q, :], - initial_state=initial_state[i], + initial_state=initial_state_kv[i], output_final_state=True, use_qk_l2norm_in_kernel=use_qk_l2norm_in_kernel, ) output[:, start_q:end_q, :, :] = core_attn_outi - final_state[i] = last_recurrent_state + final_state[i] = last_recurrent_state.transpose(-1, -2).contiguous() start_q = end_q return output, final_state @@ -217,16 +224,24 @@ def sigmoid_gating_delta_rule_update( ): beta = b.sigmoid() g = -A_log.float().exp() * softplus(a.float() + dt_bias) - return torch_recurrent_gated_delta_rule( + initial_state_kv = ( + initial_state.transpose(-1, -2).contiguous() + if initial_state is not None + else None + ) + core_attn_out, last_recurrent_state = torch_recurrent_gated_delta_rule( query, key, value, g.unsqueeze(1), beta.unsqueeze(1), - initial_state, + initial_state_kv, output_final_state, use_qk_l2norm_in_kernel=use_qk_l2norm_in_kernel, ) + if last_recurrent_state is not None: + last_recurrent_state = last_recurrent_state.transpose(-1, -2).contiguous() + return core_attn_out, last_recurrent_state def torch_gdn_gating(A_log, a, b, dt_bias): @@ -237,19 +252,27 @@ def torch_gdn_gating(A_log, a, b, dt_bias): class TestMambaAttention(CustomTestCase): def test_chunk_gated_delta_rule(self): - B, L, HK, HV, EK, EV, N = 1, 100, 3, 6, 64, 64, 4 - seqlens = torch.randint(1, L, (N + 1,)) - seqlens[0] = 0 - cu_seqlens_ = torch.cumsum(seqlens, dim=0).to(torch.int32) + B, T_PER_SEQ, HK, HV, K, V, N = 1, 128, 16, 32, 128, 128, 4 + seq_lens = torch.tensor( + [T_PER_SEQ - 7, T_PER_SEQ + 11, T_PER_SEQ - 13, T_PER_SEQ + 9], + dtype=torch.int32, + ) + cu_seqlens_ = torch.cat( + [ + torch.zeros(1, dtype=torch.int32), + seq_lens.cumsum(dim=0, dtype=torch.int32), + ] + ) T = cu_seqlens_[-1].item() - query_ = torch.rand((B, T, HK, EK), dtype=torch.bfloat16) * 0.05 - key_ = torch.rand((B, T, HK, EK), dtype=torch.bfloat16) * 0.05 - value_ = torch.rand((B, T, HV, EV), dtype=torch.bfloat16) * 0.05 - g_ = torch.rand((B, T, HV), dtype=torch.float32) * 0.05 - beta_ = torch.rand((B, T, HV), dtype=torch.bfloat16) * 0.05 - initial_state_ = torch.rand((N, HV, EK, EV), dtype=torch.float32) * 0.05 + query_ = torch.randn((B, T, HK, K), dtype=torch.bfloat16) + key_ = torch.randn((B, T, HK, K), dtype=torch.bfloat16) + value_ = torch.randn((B, T, HV, V), dtype=torch.bfloat16) + g_ = F.logsigmoid(torch.randn((B, T, HV), dtype=torch.float32)) + beta_ = torch.sigmoid(torch.randn((B, T, HV), dtype=torch.bfloat16)) + initial_state_ = torch.randn((N, HV, V, K), dtype=torch.float32) * 0.1 - for use_qk_l2norm_in_kernel in [True, False]: + # skip `use_qk_l2norm_in_kernel=False` case since it's not numerically stable in bfloat16 + for use_qk_l2norm_in_kernel in [True]: core_attn_out_ref, last_recurrent_state_ref = chunk_gated_delta_rule_update( query=query_, key=key_, @@ -267,7 +290,7 @@ class TestMambaAttention(CustomTestCase): g = g_.clone() beta = beta_.clone() cu_seqlens = cu_seqlens_.clone() - initial_state = initial_state_.clone() + initial_state = initial_state_.clone().transpose(-1, -2).contiguous() core_attn_out, last_recurrent_state = ( torch.ops.sgl_kernel.chunk_gated_delta_rule_cpu( @@ -283,6 +306,7 @@ class TestMambaAttention(CustomTestCase): use_qk_l2norm_in_kernel=use_qk_l2norm_in_kernel, ) ) + last_recurrent_state = last_recurrent_state.transpose(-1, -2).contiguous() atol = rtol = precision[core_attn_out.dtype] torch.testing.assert_close( core_attn_out, core_attn_out_ref, atol=atol, rtol=rtol @@ -350,7 +374,7 @@ class TestMambaAttention(CustomTestCase): a = torch.rand(batch_size, num_value_heads, dtype=torch.bfloat16) b = torch.rand(batch_size, num_value_heads, dtype=torch.bfloat16) dt_bias = torch.rand(num_value_heads, dtype=torch.bfloat16) - ssm_states = torch.rand( + ssm_states_kv = torch.rand( 513, num_value_heads, head_k_dim, head_v_dim, dtype=torch.float32 ) cache_indices = torch.randint(0, 513, (batch_size,), dtype=torch.int32) @@ -372,7 +396,9 @@ class TestMambaAttention(CustomTestCase): a, dt_bias, b, - initial_state=ssm_states[cache_indices], + initial_state=ssm_states_kv[cache_indices] + .transpose(-1, -2) + .contiguous(), output_final_state=True, use_qk_l2norm_in_kernel=use_qk_l2norm_in_kernel, ) @@ -386,7 +412,7 @@ class TestMambaAttention(CustomTestCase): v=value, a=a, b=b, - initial_state_source=ssm_states, + initial_state_source=ssm_states_kv, initial_state_indices=cache_indices, cu_seqlens=query_start_loc, use_qk_l2norm_in_kernel=use_qk_l2norm_in_kernel, @@ -394,7 +420,9 @@ class TestMambaAttention(CustomTestCase): softplus_threshold=20.0, ) ) - last_recurrent_state = ssm_states[cache_indices] + last_recurrent_state = ( + ssm_states_kv[cache_indices].transpose(-1, -2).contiguous() + ) atol = rtol = precision[core_attn_out.dtype] torch.testing.assert_close( core_attn_out, core_attn_out_ref, atol=atol, rtol=rtol