| 1 | // Adapted from |
| 2 | // https://github.com/sgl-project/sglang/tree/main/sgl-kernel/csrc/cpu |
| 3 | |
| 4 | // clang-format off |
| 5 | |
| 6 | #include "common.h" |
| 7 | #include "gemm.h" |
| 8 | #include "vec.h" |
| 9 | #include "vec_pack.h" |
| 10 | |
| 11 | namespace { |
| 12 | |
| 13 | // [NOTE] GDN Optimizations on AMX CPU |
| 14 | // * intra loop: fuse `kkt_solve` and `recompute_w_u` so as to avoid materialize `A`. |
| 15 | // * inter loop: fuse `recompute_w_u` and `update_v` so as to avoid materialize `h` and `v_new`. |
| 16 | // * intra loop parallel on H instead of Hv, remove duplicated key @ key.T |
| 17 | // * fuse format pack with elemwise OP as much as possible. |
| 18 | // * update state (FP32) with amx-bf16 where C(FP32) += A(BF16) * B(BF16) |
| 19 | // * compile time mask out upper triangular part in decay mask and tril solve, reduce fma needed. |
| 20 | |
| 21 | // * convert to vnni format， expect contiguous input and output |
| 22 | // from [K/2, 2, N] FP32 to [K/2, N, 2] BF16 |
| 23 | // * update src = src * exp(g_last) |
| 24 | template <typename scalar_t, int K, int N> |
| 25 | void pack_vnni2(scalar_t* __restrict__ dst, float* __restrict__ src, const float g_last, int ld_src, int ld_dst) { |
| 26 | static_assert(K % 32 == 0); |
| 27 | static_assert(N % 32 == 0); |
| 28 | |
| 29 | const float scale = std::exp(g_last); |
| 30 | #if defined(CPU_CAPABILITY_AVX512) |
| 31 | constexpr int KB = K / 2; |
| 32 | constexpr int NB = N / 32; |
| 33 | |
| 34 | __m512i s0, s1, d0, d1; |
| 35 | __m512 vd = _mm512_set1_ps(scale); |
| 36 | |
| 37 | const auto trans = [&](auto i) { |
| 38 | constexpr int kb = i / NB; |
| 39 | constexpr int nb = i % NB; |
| 40 | |
| 41 | // [K/2, 2, N/32, 32] -> [K/2, N/32, 32, 2] |
| 42 | constexpr int k0 = kb * 2 + 0; |
| 43 | constexpr int k1 = kb * 2 + 1; |
| 44 | __m512 v00 = _mm512_loadu_ps(src + k0 * ld_src + nb * 32); |
| 45 | __m512 v01 = _mm512_loadu_ps(src + k0 * ld_src + nb * 32 + 16); |
| 46 | __m512 v10 = _mm512_loadu_ps(src + k1 * ld_src + nb * 32); |
| 47 | __m512 v11 = _mm512_loadu_ps(src + k1 * ld_src + nb * 32 + 16); |
| 48 | s0 = (__m512i)_mm512_cvtne2ps_pbh(v01, v00); |
| 49 | s1 = (__m512i)_mm512_cvtne2ps_pbh(v11, v10); |
| 50 | |
| 51 | std::tie(d0, d1) = transpose_2x32_16bit(s0, s1); |
| 52 | _mm512_storeu_si512(dst + kb * ld_dst * 2 + nb * 32 * 2, d0); |
| 53 | _mm512_storeu_si512(dst + kb * ld_dst * 2 + nb * 32 * 2 + 32, d1); |
| 54 | |
| 55 | // update src = src * exp(g_last) |
| 56 | _mm512_storeu_ps(src + k0 * ld_src + nb * 32, _mm512_mul_ps(v00, vd)); |
| 57 | _mm512_storeu_ps(src + k0 * ld_src + nb * 32 + 16, _mm512_mul_ps(v01, vd)); |
| 58 | _mm512_storeu_ps(src + k1 * ld_src + nb * 32, _mm512_mul_ps(v10, vd)); |
| 59 | _mm512_storeu_ps(src + k1 * ld_src + nb * 32 + 16, _mm512_mul_ps(v11, vd)); |
| 60 | }; |
| 61 | Unroll<KB * NB>{}(trans); |
| 62 | #else |
| 63 | // [K/2, 2, N] -> [K/2, N, 2] |
| 64 | for (int k = 0; k < K; k += 2) { |
| 65 | for (int n = 0; n < N; ++n) { |
| 66 | const float v0 = src[(k + 0) * ld_src + n]; |
| 67 | const float v1 = src[(k + 1) * ld_src + n]; |
| 68 | dst[(k >> 1) * ld_dst * 2 + n * 2 + 0] = static_cast<scalar_t>(v0); |
| 69 | dst[(k >> 1) * ld_dst * 2 + n * 2 + 1] = static_cast<scalar_t>(v1); |
| 70 | src[(k + 0) * ld_src + n] = v0 * scale; |
| 71 | src[(k + 1) * ld_src + n] = v1 * scale; |
| 72 | } |
| 73 | } |
| 74 | #endif |
| 75 | } |
| 76 | |
| 77 | template <typename state_t> |
| 78 | inline void store_state_vec2( |
| 79 | state_t* __restrict__ dst, |
| 80 | const at::vec::Vectorized<float>& src0, |
| 81 | const at::vec::Vectorized<float>& src1) { |
| 82 | if constexpr (std::is_same_v<state_t, float>) { |
| 83 | src0.store(dst); |
| 84 | src1.store(dst + at::vec::Vectorized<float>::size()); |
| 85 | } else { |
| 86 | auto state_vec = convert_from_float_ext<state_t>(src0, src1); |
| 87 | state_vec.store(dst); |
| 88 | } |
| 89 | } |
| 90 | |
| 91 | template <typename scalar_t, int SIZE> |
| 92 | inline void fill_stub(scalar_t* __restrict__ out, float val) { |
| 93 | using Vec = at::vec::Vectorized<scalar_t>; |
| 94 | constexpr int kVecSize = Vec::size(); |
| 95 | static_assert(SIZE % kVecSize == 0); |
| 96 | const Vec data_vec = Vec(static_cast<scalar_t>(val)); |
| 97 | #pragma GCC unroll 8 |
| 98 | for (int d = 0; d < SIZE; d += kVecSize) { |
| 99 | data_vec.store(out + d); |
| 100 | } |
| 101 | } |
| 102 | |
| 103 | // Portable fallback for non-AVX512 builds (ARM/NEON, old x86 without |
| 104 | // AVX512BF16), vectorized via at::vec::Vectorized<T> (portable across |
| 105 | // AVX2/NEON/generic) mirroring the idioms used elsewhere in this file (see |
| 106 | // l2norm_fwd_kernel_impl's predecessor and fused_gdn_gating_kernel_impl): |
| 107 | // bVec/fVec pairs with convert_to_float/convert_from_float for bf16<->float, |
| 108 | // plain scalar tails for remainders. Two kernels (cumsum_kernel, |
| 109 | // update_key_kernel) write a transposed layout relative to their vectorized |
| 110 | // read axis; those vectorize the load/compute and unpack lanes for the |
| 111 | // (unavoidably strided) store. |
| 112 | template <typename scalar_t, int D, bool has_scale> |
| 113 | struct l2norm_kernel { |
| 114 | static inline void apply(scalar_t* __restrict__ out, const scalar_t* __restrict__ input, float eps) { |
| 115 | using bVec = at::vec::Vectorized<scalar_t>; |
| 116 | using fVec = at::vec::Vectorized<float>; |
| 117 | constexpr int bVecSize = bVec::size(); |
| 118 | const float scale = 1.f / std::sqrt(static_cast<float>(D)); |
| 119 | |
| 120 | fVec sum_fvec0(0.f), sum_fvec1(0.f); |
| 121 | int d = 0; |
| 122 | for (; d <= D - bVecSize; d += bVecSize) { |
| 123 | bVec in_bvec = bVec::loadu(input + d); |
| 124 | fVec in0, in1; |
| 125 | std::tie(in0, in1) = at::vec::convert_to_float(in_bvec); |
| 126 | sum_fvec0 = sum_fvec0 + in0 * in0; |
| 127 | sum_fvec1 = sum_fvec1 + in1 * in1; |
| 128 | } |
| 129 | float sqsum = vec_reduce_sum(sum_fvec0 + sum_fvec1); |
| 130 | for (; d < D; ++d) { |
| 131 | float v = static_cast<float>(input[d]); |
| 132 | sqsum += v * v; |
| 133 | } |
| 134 | |
| 135 | float rscale = 1.f / std::sqrt(sqsum + eps); |
| 136 | fVec rscale_fvec(rscale); |
| 137 | fVec scale_fvec(scale); |
| 138 | d = 0; |
| 139 | for (; d <= D - bVecSize; d += bVecSize) { |
| 140 | bVec in_bvec = bVec::loadu(input + d); |
| 141 | fVec in0, in1; |
| 142 | std::tie(in0, in1) = at::vec::convert_to_float(in_bvec); |
| 143 | in0 = in0 * rscale_fvec; |
| 144 | in1 = in1 * rscale_fvec; |
| 145 | if constexpr (has_scale) { |
| 146 | in0 = in0 * scale_fvec; |
| 147 | in1 = in1 * scale_fvec; |
| 148 | } |
| 149 | bVec out_bvec = at::vec::convert_from_float<scalar_t>(in0, in1); |
| 150 | out_bvec.store(out + d); |
| 151 | } |
| 152 | for (; d < D; ++d) { |
| 153 | float v = static_cast<float>(input[d]) * rscale; |
| 154 | if constexpr (has_scale) { |
| 155 | v *= scale; |
| 156 | } |
| 157 | out[d] = static_cast<scalar_t>(v); |
| 158 | } |
| 159 | } |
| 160 | }; |
| 161 | |
| 162 | #if defined(CPU_CAPABILITY_AVX512) |
| 163 | template <int D, bool has_scale> |
| 164 | struct l2norm_kernel<at::BFloat16, D, has_scale> { |
| 165 | static inline void apply(at::BFloat16* __restrict__ out, const at::BFloat16* __restrict__ input, float eps) { |
| 166 | static_assert(D % 32 == 0); |
| 167 | constexpr int COLS = D / 32; |
| 168 | |
| 169 | __m512bh va[COLS]; |
| 170 | __m512 vrscale; |
| 171 | |
| 172 | const float scale = 1.f / std::sqrt(D); |
| 173 | __m512 vscale = _mm512_set1_ps(scale); |
| 174 | |
| 175 | // step 1: load input and do reduce with avx512-bf16 |
| 176 | __m512 vsum = _mm512_set1_ps(0.f); |
| 177 | auto reduce = [&](auto col) { |
| 178 | va[col] = (__m512bh)(_mm512_loadu_si512(input + col * 32)); |
| 179 | vsum = _mm512_dpbf16_ps(vsum, va[col], va[col]); |
| 180 | }; |
| 181 | Unroll<COLS>{}(reduce); |
| 182 | |
| 183 | float sqsum = _mm512_reduce_add_ps(vsum); |
| 184 | float rscale = 1.f / std::sqrt(sqsum + eps); |
| 185 | vrscale = _mm512_set1_ps(rscale); |
| 186 | |
| 187 | // step 2: apply scale to output |
| 188 | auto map = [&](auto col) { |
| 189 | __m512i a16 = (__m512i)va[col]; |
| 190 | __m512 va0 = CVT_BF16_TO_FP32(_mm512_extracti32x8_epi32(a16, 0)); |
| 191 | __m512 va1 = CVT_BF16_TO_FP32(_mm512_extracti32x8_epi32(a16, 1)); |
| 192 | va0 = _mm512_mul_ps(va0, vrscale); |
| 193 | va1 = _mm512_mul_ps(va1, vrscale); |
| 194 | // keep the mul order same as torch code: |
| 195 | // query = l2norm(query) * scale |
| 196 | if constexpr (has_scale) { |
| 197 | va0 = _mm512_mul_ps(va0, vscale); |
| 198 | va1 = _mm512_mul_ps(va1, vscale); |
| 199 | } |
| 200 | _mm512_storeu_si512(out + col * 32, (__m512i)(_mm512_cvtne2ps_pbh(va1, va0))); |
| 201 | }; |
| 202 | Unroll<COLS>{}(map); |
| 203 | } |
| 204 | }; |
| 205 | #endif |
| 206 | |
| 207 | template <typename scalar_t, int CHUNK_SIZE, int BLOCK_H> |
| 208 | struct cumsum_kernel { |
| 209 | static inline void apply( |
| 210 | scalar_t* __restrict__ out, |
| 211 | const scalar_t* __restrict__ input, |
| 212 | int mb_size, |
| 213 | int hb_size, |
| 214 | int ld_src, |
| 215 | int ld_dst) { |
| 216 | // out: [hb_size valid rows, CHUNK_SIZE] within a [BLOCK_H, ld_dst] buffer |
| 217 | // input: [mb_size valid rows, hb_size] within a [CHUNK_SIZE(padded), ld_src] buffer |
| 218 | // input is contiguous along j (vectorize the load/accumulate); out is |
| 219 | // contiguous along i instead (transposed), so the store is per-lane. |
| 220 | using Vec = at::vec::Vectorized<scalar_t>; |
| 221 | constexpr int VecSize = Vec::size(); |
| 222 | alignas(64) scalar_t lane_buf[VecSize]; |
| 223 | int j = 0; |
| 224 | for (; j <= hb_size - VecSize; j += VecSize) { |
| 225 | Vec running(static_cast<scalar_t>(0)); |
| 226 | for (int i = 0; i < CHUNK_SIZE; ++i) { |
| 227 | if (i < mb_size) { |
| 228 | running = running + Vec::loadu(input + i * ld_src + j); |
| 229 | } |
| 230 | running.store(lane_buf); |
| 231 | for (int lane = 0; lane < VecSize; ++lane) { |
| 232 | out[(j + lane) * ld_dst + i] = lane_buf[lane]; |
| 233 | } |
| 234 | } |
| 235 | } |
| 236 | for (; j < hb_size; ++j) { |
| 237 | float running = 0.f; |
| 238 | for (int i = 0; i < CHUNK_SIZE; ++i) { |
| 239 | if (i < mb_size) { |
| 240 | running += static_cast<float>(input[i * ld_src + j]); |
| 241 | } |
| 242 | out[j * ld_dst + i] = static_cast<scalar_t>(running); |
| 243 | } |
| 244 | } |
| 245 | } |
| 246 | }; |
| 247 | |
| 248 | #if defined(CPU_CAPABILITY_AVX512) |
| 249 | template <int CHUNK_SIZE, int BLOCK_H> |
| 250 | struct cumsum_kernel<float, CHUNK_SIZE, BLOCK_H> { |
| 251 | static inline void |
| 252 | apply(float* __restrict__ out, const float* __restrict__ input, int mb_size, int hb_size, int ld_src, int ld_dst) { |
| 253 | // vector length of fp32 for avx512 |
| 254 | static_assert(BLOCK_H == 16); |
| 255 | TORCH_CHECK(hb_size > 0 && hb_size <= BLOCK_H); |
| 256 | const __mmask16 vmask = static_cast<__mmask16>((1u << hb_size) - 1u); |
| 257 | |
| 258 | __m512i va[16]; |
| 259 | __m512 vsum = _mm512_set1_ps(0.f); |
| 260 | |
| 261 | for (int i = 0; i < CHUNK_SIZE; i += 16) { |
| 262 | // load input data |
| 263 | Unroll<16>{}([&](auto j) { |
| 264 | __m512 v; |
| 265 | if (i + j < mb_size) { |
| 266 | v = _mm512_maskz_loadu_ps(vmask, input + (i + j) * ld_src); |
| 267 | } else { |
| 268 | v = _mm512_setzero_ps(); |
| 269 | } |
| 270 | vsum = _mm512_add_ps(vsum, v); |
| 271 | va[j] = _mm512_castps_si512(vsum); |
| 272 | }); |
| 273 | // transpose |
| 274 | transpose_16x16_32bit(va); |
| 275 | // store output data |
| 276 | Unroll<16>{}([&](auto j) { |
| 277 | if (j < hb_size) { |
| 278 | _mm512_storeu_si512(out + j * ld_dst + i, va[j]); |
| 279 | } |
| 280 | }); |
| 281 | } |
| 282 | } |
| 283 | }; |
| 284 | #endif |
| 285 | |
| 286 | template <typename scalar_t, int CHUNK_SIZE> |
| 287 | struct decay_mask_kernel { |
| 288 | // decay_mask = ((g.unsqueeze(-1) - g.unsqueeze(-2)).tril().exp().float()).tril() |
| 289 | static inline void apply(scalar_t* __restrict__ out, const scalar_t* __restrict__ input) { |
| 290 | using Vec = at::vec::Vectorized<scalar_t>; |
| 291 | constexpr int VecSize = Vec::size(); |
| 292 | const Vec zero(static_cast<scalar_t>(0)); |
| 293 | for (int row = 0; row < CHUNK_SIZE; ++row) { |
| 294 | Vec g_row(input[row]); |
| 295 | Vec limit_vec(static_cast<scalar_t>(row)); |
| 296 | int col = 0; |
| 297 | for (; col <= CHUNK_SIZE - VecSize; col += VecSize) { |
| 298 | Vec g_col = Vec::loadu(input + col); |
| 299 | Vec vc = (g_row - g_col).exp_u20(); |
| 300 | Vec idx = Vec::arange(static_cast<scalar_t>(col), static_cast<scalar_t>(1)); |
| 301 | Vec result = Vec::blendv(zero, vc, idx <= limit_vec); |
| 302 | result.store(out + row * CHUNK_SIZE + col); |
| 303 | } |
| 304 | for (; col < CHUNK_SIZE; ++col) { |
| 305 | out[row * CHUNK_SIZE + col] = |
| 306 | col <= row |
| 307 | ? static_cast<scalar_t>(std::exp(static_cast<float>(input[row]) - static_cast<float>(input[col]))) |
| 308 | : static_cast<scalar_t>(0); |
| 309 | } |
| 310 | } |
| 311 | } |
| 312 | }; |
| 313 | |
| 314 | #if defined(CPU_CAPABILITY_AVX512) |
| 315 | template <int CHUNK_SIZE> |
| 316 | struct decay_mask_kernel<float, CHUNK_SIZE> { |
| 317 | static inline void apply(float* __restrict__ out, const float* __restrict__ input) { |
| 318 | static_assert(CHUNK_SIZE % 16 == 0); |
| 319 | |
| 320 | constexpr int ROWS = CHUNK_SIZE; |
| 321 | constexpr int COLS = CHUNK_SIZE / 16; |
| 322 | |
| 323 | __m512 va; |
| 324 | __m512 vb[COLS]; |
| 325 | |
| 326 | // step 1: load g[j] |
| 327 | auto loadb = [&](auto i) { vb[i] = _mm512_loadu_ps(input + i * 16); }; |
| 328 | Unroll<COLS>{}(loadb); |
| 329 | |
| 330 | // step2: exp(g[i] - g[j]) |
| 331 | auto compute = [&](auto i) { |
| 332 | constexpr int row = i / COLS; |
| 333 | constexpr int col = i % COLS; |
| 334 | |
| 335 | if constexpr (col == 0) { |
| 336 | va = _mm512_set1_ps(input[row]); |
| 337 | } |
| 338 | |
| 339 | // mask vb[col] (already loaded in step 1) for the lower-triangular region |
| 340 | constexpr int len = std::max(0, std::min(row + 1 - col * 16, 16)); |
| 341 | |
| 342 | __m512 vc; |
| 343 | if constexpr (len == 16) { |
| 344 | vc = _mm512_fexp_u20_ps(va - vb[col]); |
| 345 | } else if constexpr (len == 0) { |
| 346 | vc = _mm512_setzero_ps(); |
| 347 | } else { |
| 348 | vc = _mm512_fexp_u20_ps(va - vb[col]); |
| 349 | // do mask for vc |
| 350 | constexpr __mmask16 vmask = (1 << len) - 1; |
| 351 | vc = _mm512_mask_blend_ps(vmask, _mm512_setzero_ps(), vc); |
| 352 | } |
| 353 | _mm512_storeu_ps(out + row * CHUNK_SIZE + col * 16, vc); |
| 354 | }; |
| 355 | Unroll<ROWS * COLS>{}(compute); |
| 356 | } |
| 357 | }; |
| 358 | #endif |
| 359 | |
| 360 | template <typename scalar_t, int CHUNK_SIZE, bool has_beta> |
| 361 | struct apply_mask_kernel { |
| 362 | // has_beta: attn2 = -attn * beta * d (strict lower, col < row) |
| 363 | // !has_beta: attn2 = attn * d (lower incl. diagonal, col <= row) |
| 364 | static inline void apply( |
| 365 | scalar_t* __restrict__ attn2, |
| 366 | const float* __restrict__ attn, |
| 367 | const scalar_t* __restrict__ beta, |
| 368 | const float* __restrict__ d, |
| 369 | int size, |
| 370 | int b_stride = 0) { |
| 371 | using bVec = at::vec::Vectorized<scalar_t>; |
| 372 | using fVec = at::vec::Vectorized<float>; |
| 373 | constexpr int fVecSize = fVec::size(); |
| 374 | constexpr int bVecSize = bVec::size(); |
| 375 | const fVec zero(0.f); |
| 376 | for (int row = 0; row < size; ++row) { |
| 377 | int col_limit = has_beta ? row : row + 1; |
| 378 | float beta_val = 1.f; |
| 379 | if constexpr (has_beta) { |
| 380 | beta_val = -static_cast<float>(beta[row * b_stride]); |
| 381 | } |
| 382 | fVec beta_fvec(beta_val); |
| 383 | fVec limit_fvec(static_cast<float>(col_limit)); |
| 384 | int col = 0; |
| 385 | for (; col <= CHUNK_SIZE - bVecSize; col += bVecSize) { |
| 386 | fVec a0 = fVec::loadu(attn + row * CHUNK_SIZE + col); |
| 387 | fVec a1 = fVec::loadu(attn + row * CHUNK_SIZE + col + fVecSize); |
| 388 | fVec d0 = fVec::loadu(d + row * CHUNK_SIZE + col); |
| 389 | fVec d1 = fVec::loadu(d + row * CHUNK_SIZE + col + fVecSize); |
| 390 | fVec v0 = a0 * beta_fvec * d0; |
| 391 | fVec v1 = a1 * beta_fvec * d1; |
| 392 | fVec idx0 = fVec::arange(static_cast<float>(col), 1.f); |
| 393 | fVec idx1 = fVec::arange(static_cast<float>(col + fVecSize), 1.f); |
| 394 | v0 = fVec::blendv(zero, v0, idx0 < limit_fvec); |
| 395 | v1 = fVec::blendv(zero, v1, idx1 < limit_fvec); |
| 396 | bVec out_bvec = at::vec::convert_from_float<scalar_t>(v0, v1); |
| 397 | out_bvec.store(attn2 + row * CHUNK_SIZE + col); |
| 398 | } |
| 399 | for (; col < CHUNK_SIZE; ++col) { |
| 400 | float v = 0.f; |
| 401 | if (col < col_limit) { |
| 402 | v = attn[row * CHUNK_SIZE + col] * beta_val * d[row * CHUNK_SIZE + col]; |
| 403 | } |
| 404 | attn2[row * CHUNK_SIZE + col] = static_cast<scalar_t>(v); |
| 405 | } |
| 406 | } |
| 407 | } |
| 408 | }; |
| 409 | |
| 410 | #if defined(CPU_CAPABILITY_AVX512) |
| 411 | template <int CHUNK_SIZE, bool has_beta> |
| 412 | struct apply_mask_kernel<at::BFloat16, CHUNK_SIZE, has_beta> { |
| 413 | static inline void apply( |
| 414 | at::BFloat16* __restrict__ attn2, |
| 415 | const float* __restrict__ attn, |
| 416 | const at::BFloat16* __restrict__ beta, |
| 417 | const float* __restrict__ d, |
| 418 | int size, |
| 419 | int b_stride = 0) { |
| 420 | static_assert(CHUNK_SIZE % 16 == 0); |
| 421 | |
| 422 | constexpr int ROWS = CHUNK_SIZE; |
| 423 | constexpr int COLS = CHUNK_SIZE / 16; |
| 424 | |
| 425 | __m512 vbeta; |
| 426 | |
| 427 | // has_beta: attn2 = -attn * beta * d (strict lower) |
| 428 | // !has_beta: attn2 = attn * d (lower incl. diagonal) |
| 429 | auto compute = [&](auto i) { |
| 430 | constexpr int row = i / COLS; |
| 431 | constexpr int col = i % COLS; |
| 432 | |
| 433 | constexpr int len = |
| 434 | has_beta ? std::max(0, std::min(row - col * 16, 16)) : std::max(0, std::min(row + 1 - col * 16, 16)); |
| 435 | if (row < size) { |
| 436 | if constexpr (has_beta) { |
| 437 | if constexpr (col == 0) { |
| 438 | vbeta = _mm512_set1_ps(-static_cast<float>(beta[row * b_stride])); |
| 439 | } |
| 440 | } |
| 441 | |
| 442 | __m512 vc; |
| 443 | if constexpr (len == 0) { |
| 444 | vc = _mm512_setzero_ps(); |
| 445 | } else { |
| 446 | constexpr __mmask16 vmask = (1 << len) - 1; |
| 447 | __m512 va = _mm512_maskz_loadu_ps(vmask, attn + row * CHUNK_SIZE + col * 16); |
| 448 | __m512 vd = _mm512_maskz_loadu_ps(vmask, d + row * CHUNK_SIZE + col * 16); |
| 449 | if constexpr (has_beta) { |
| 450 | vc = _mm512_mul_ps(_mm512_mul_ps(va, vbeta), vd); |
| 451 | } else { |
| 452 | vc = _mm512_mul_ps(va, vd); |
| 453 | } |
| 454 | } |
| 455 | _mm256_storeu_si256( |
| 456 | reinterpret_cast<__m256i*>(attn2 + row * CHUNK_SIZE + col * 16), (__m256i)(_mm512_cvtneps_pbh(vc))); |
| 457 | } |
| 458 | }; |
| 459 | Unroll<ROWS * COLS>{}(compute); |
| 460 | } |
| 461 | }; |
| 462 | #endif |
| 463 | |
| 464 | template <typename scalar_t, int CHUNK_SIZE> |
| 465 | struct solve_tril_kernel { |
| 466 | // (I + L)^{-1} via forward substitution, L = strict-lower part of attn2. |
| 467 | static inline void apply(scalar_t* __restrict__ attn2, int size) { |
| 468 | using bVec = at::vec::Vectorized<scalar_t>; |
| 469 | using fVec = at::vec::Vectorized<float>; |
| 470 | constexpr int fVecSize = fVec::size(); |
| 471 | constexpr int bVecSize = bVec::size(); |
| 472 | |
| 473 | // for len == 0 and row < size, we don't have to write back zero again |
| 474 | // as in `apply_mask_kernel`, we already set zero for the upper-triangular region |
| 475 | for (int i = 1; i < size; ++i) { |
| 476 | scalar_t* __restrict__ row_ptr = attn2 + i * CHUNK_SIZE; |
| 477 | float vsum[CHUNK_SIZE]; |
| 478 | int j = 0; |
| 479 | for (; j <= i - bVecSize; j += bVecSize) { |
| 480 | bVec row_bvec = bVec::loadu(row_ptr + j); |
| 481 | fVec f0, f1; |
| 482 | std::tie(f0, f1) = at::vec::convert_to_float(row_bvec); |
| 483 | f0.store(vsum + j); |
| 484 | f1.store(vsum + j + fVecSize); |
| 485 | } |
| 486 | for (; j < i; ++j) { |
| 487 | vsum[j] = static_cast<float>(row_ptr[j]); |
| 488 | } |
| 489 | |
| 490 | // row = attn[..., i, :i].clone() |
| 491 | // sub = attn[..., :i, :i].clone() |
| 492 | // vsum = row + (row.unsqueeze(-1) * sub).sum(-2) |
| 493 | for (int k = 0; k < i; ++k) { |
| 494 | // read BEFORE row_ptr is written back below (row k was finalized in |
| 495 | // an earlier outer iteration; row i itself is untouched until the |
| 496 | // final write-back after this loop) |
| 497 | float va = static_cast<float>(row_ptr[k]); |
| 498 | fVec va_vec(va); |
| 499 | const scalar_t* __restrict__ row_k_ptr = attn2 + k * CHUNK_SIZE; |
| 500 | int jj = 0; |
| 501 | for (; jj <= k - bVecSize; jj += bVecSize) { |
| 502 | bVec rk_bvec = bVec::loadu(row_k_ptr + jj); |
| 503 | fVec rk0, rk1; |
| 504 | std::tie(rk0, rk1) = at::vec::convert_to_float(rk_bvec); |
| 505 | fVec vsum0 = fVec::loadu(vsum + jj); |
| 506 | fVec vsum1 = fVec::loadu(vsum + jj + fVecSize); |
| 507 | vsum0 = vsum0 + va_vec * rk0; |
| 508 | vsum1 = vsum1 + va_vec * rk1; |
| 509 | vsum0.store(vsum + jj); |
| 510 | vsum1.store(vsum + jj + fVecSize); |
| 511 | } |
| 512 | for (; jj < k; ++jj) { |
| 513 | vsum[jj] += va * static_cast<float>(row_k_ptr[jj]); |
| 514 | } |
| 515 | } |
| 516 | |
| 517 | j = 0; |
| 518 | for (; j <= i - bVecSize; j += bVecSize) { |
| 519 | fVec f0 = fVec::loadu(vsum + j); |
| 520 | fVec f1 = fVec::loadu(vsum + j + fVecSize); |
| 521 | bVec out_bvec = at::vec::convert_from_float<scalar_t>(f0, f1); |
| 522 | out_bvec.store(row_ptr + j); |
| 523 | } |
| 524 | for (; j < i; ++j) { |
| 525 | row_ptr[j] = static_cast<scalar_t>(vsum[j]); |
| 526 | } |
| 527 | } |
| 528 | |
| 529 | // attn = attn + torch.eye(chunk_size) |
| 530 | for (int i = 0; i < size; ++i) { |
| 531 | attn2[i * CHUNK_SIZE + i] = static_cast<scalar_t>(static_cast<float>(attn2[i * CHUNK_SIZE + i]) + 1.f); |
| 532 | } |
| 533 | } |
| 534 | }; |
| 535 | |
| 536 | #if defined(CPU_CAPABILITY_AVX512) |
| 537 | template <int CHUNK_SIZE> |
| 538 | struct solve_tril_kernel<at::BFloat16, CHUNK_SIZE> { |
| 539 | static inline void apply(at::BFloat16* __restrict__ attn2, int size) { |
| 540 | static_assert(CHUNK_SIZE % 16 == 0); |
| 541 | |
| 542 | constexpr int COLS = CHUNK_SIZE / 16; |
| 543 | |
| 544 | __m512 va; |
| 545 | __m512 vb[COLS]; |
| 546 | __m512 vsum[COLS]; |
| 547 | |
| 548 | // for len == 0 and row < size, we don't have to write back zero again |
| 549 | // as in `apply_mask_kernel`, we already set zero for the upper-triangular region |
| 550 | for (int i = 1; i < size; ++i) { |
| 551 | // load row attn[..., i, :i] |
| 552 | at::BFloat16* __restrict__ row_ptr = attn2 + i * CHUNK_SIZE; |
| 553 | Unroll<COLS>{}([&](auto col) { |
| 554 | int len = std::min(i - col * 16, 16); |
| 555 | if (len > 0) { |
| 556 | const __mmask16 vmask = (1 << len) - 1; |
| 557 | vsum[col] = CVT_BF16_TO_FP32(_mm256_maskz_loadu_epi16(vmask, row_ptr + col * 16)); |
| 558 | } |
| 559 | }); |
| 560 | |
| 561 | // row = attn[..., i, :i].clone() |
| 562 | // sub = attn[..., :i, :i].clone() |
| 563 | // vsum = row + (row.unsqueeze(-1) * sub).sum(-2) |
| 564 | for (int k = 0; k < i; ++k) { |
| 565 | va = _mm512_set1_ps(static_cast<float>(row_ptr[k])); |
| 566 | |
| 567 | const at::BFloat16* __restrict__ row_k_ptr = attn2 + k * CHUNK_SIZE; |
| 568 | Unroll<COLS>{}([&](auto col) { |
| 569 | int len = std::min(k - col * 16, 16); |
| 570 | if (len > 0) { |
| 571 | const __mmask16 vmask = (1 << len) - 1; |
| 572 | vb[col] = CVT_BF16_TO_FP32(_mm256_maskz_loadu_epi16(vmask, row_k_ptr + col * 16)); |
| 573 | vsum[col] = _mm512_fmadd_ps(va, vb[col], vsum[col]); |
| 574 | } |
| 575 | }); |
| 576 | } |
| 577 | |
| 578 | // attn[..., i, :i] = vsum |
| 579 | Unroll<COLS>{}([&](auto col) { |
| 580 | int len = std::min(i - col * 16, 16); |
| 581 | if (len > 0) { |
| 582 | const __mmask16 vmask = (1 << len) - 1; |
| 583 | _mm256_mask_storeu_epi16(row_ptr + col * 16, vmask, (__m256i)(_mm512_cvtneps_pbh(vsum[col]))); |
| 584 | } |
| 585 | }); |
| 586 | } |
| 587 | |
| 588 | // attn = attn + torch.eye(chunk_size) |
| 589 | for (int i = 0; i < size; ++i) { |
| 590 | attn2[i * CHUNK_SIZE + i] += 1.f; |
| 591 | } |
| 592 | } |
| 593 | }; |
| 594 | #endif |
| 595 | |
| 596 | template <typename scalar_t, int CHUNK_SIZE, int D, bool has_beta, bool has_g> |
| 597 | struct apply_beta_kernel { |
| 598 | static inline void apply( |
| 599 | scalar_t* __restrict__ out, |
| 600 | const scalar_t* __restrict__ input, |
| 601 | const scalar_t* __restrict__ beta, |
| 602 | const float* __restrict__ g, |
| 603 | int size, |
| 604 | int ld_src, |
| 605 | int ld_dst, |
| 606 | int b_stride) { |
| 607 | using bVec = at::vec::Vectorized<scalar_t>; |
| 608 | using fVec = at::vec::Vectorized<float>; |
| 609 | constexpr int fVecSize = fVec::size(); |
| 610 | constexpr int bVecSize = bVec::size(); |
| 611 | |
| 612 | for (int i = 0; i < size; ++i) { |
| 613 | float scale = 1.f; |
| 614 | if constexpr (has_beta) { |
| 615 | scale *= static_cast<float>(beta[i * b_stride]); |
| 616 | } |
| 617 | if constexpr (has_g) { |
| 618 | scale *= std::exp(g[i]); |
| 619 | } |
| 620 | fVec scale_fvec(scale); |
| 621 | int d = 0; |
| 622 | for (; d <= D - bVecSize; d += bVecSize) { |
| 623 | bVec in_bvec = bVec::loadu(input + i * ld_src + d); |
| 624 | fVec in0, in1; |
| 625 | std::tie(in0, in1) = at::vec::convert_to_float(in_bvec); |
| 626 | in0 = in0 * scale_fvec; |
| 627 | in1 = in1 * scale_fvec; |
| 628 | bVec out_bvec = at::vec::convert_from_float<scalar_t>(in0, in1); |
| 629 | out_bvec.store(out + i * ld_dst + d); |
| 630 | } |
| 631 | for (; d < D; ++d) { |
| 632 | out[i * ld_dst + d] = static_cast<scalar_t>(static_cast<float>(input[i * ld_src + d]) * scale); |
| 633 | } |
| 634 | } |
| 635 | } |
| 636 | }; |
| 637 | |
| 638 | #if defined(CPU_CAPABILITY_AVX512) |
| 639 | template <int CHUNK_SIZE, int D, bool has_beta, bool has_g> |
| 640 | struct apply_beta_kernel<at::BFloat16, CHUNK_SIZE, D, has_beta, has_g> { |
| 641 | static inline void apply( |
| 642 | at::BFloat16* __restrict__ out, |
| 643 | const at::BFloat16* __restrict__ input, |
| 644 | const at::BFloat16* __restrict__ beta, |
| 645 | const float* __restrict__ g, |
| 646 | int size, |
| 647 | int ld_src, |
| 648 | int ld_dst, |
| 649 | int b_stride) { |
| 650 | static_assert(D % 32 == 0); |
| 651 | constexpr int COLS = D / 16; |
| 652 | |
| 653 | // get g.exp() and g is padded to CHUNK_SIZE |
| 654 | alignas(64) float g_arr[CHUNK_SIZE]; |
| 655 | if constexpr (has_g) { |
| 656 | Unroll<CHUNK_SIZE / 16>{}([&](auto col) { |
| 657 | __m512 vg = _mm512_loadu_ps(g + col * 16); |
| 658 | __m512 vg_exp = _mm512_fexp_u20_ps(vg); |
| 659 | _mm512_storeu_ps(g_arr + col * 16, vg_exp); |
| 660 | }); |
| 661 | } |
| 662 | |
| 663 | for (int i = 0; i < size; ++i) { |
| 664 | __m512 vbeta; |
| 665 | if constexpr (has_beta) { |
| 666 | vbeta = _mm512_set1_ps(static_cast<float>(beta[i * b_stride])); |
| 667 | } |
| 668 | __m512 vg; |
| 669 | if constexpr (has_g) { |
| 670 | vg = _mm512_set1_ps(g_arr[i]); |
| 671 | } |
| 672 | |
| 673 | Unroll<COLS>{}([&](auto col) { |
| 674 | // load for 0, 2, 4, 6 |
| 675 | if constexpr (col % 2 == 0) { |
| 676 | __m512i a16 = _mm512_loadu_si512(input + i * ld_src + col * 16); |
| 677 | __m512 va0 = CVT_BF16_TO_FP32(_mm512_extracti32x8_epi32(a16, 0)); |
| 678 | __m512 va1 = CVT_BF16_TO_FP32(_mm512_extracti32x8_epi32(a16, 1)); |
| 679 | if constexpr (has_beta) { |
| 680 | va0 = _mm512_mul_ps(va0, vbeta); |
| 681 | va1 = _mm512_mul_ps(va1, vbeta); |
| 682 | } |
| 683 | if constexpr (has_g) { |
| 684 | va0 = _mm512_mul_ps(va0, vg); |
| 685 | va1 = _mm512_mul_ps(va1, vg); |
| 686 | } |
| 687 | _mm512_storeu_si512(out + i * ld_dst + col * 16, (__m512i)(_mm512_cvtne2ps_pbh(va1, va0))); |
| 688 | } |
| 689 | }); |
| 690 | } |
| 691 | } |
| 692 | }; |
| 693 | #endif |
| 694 | |
| 695 | template <typename scalar_t, int D> |
| 696 | struct update_kernel { |
| 697 | static inline void |
| 698 | apply(scalar_t* __restrict__ out, const float* __restrict__ input, int size, int ld_src, int ld_dst) { |
| 699 | using bVec = at::vec::Vectorized<scalar_t>; |
| 700 | using fVec = at::vec::Vectorized<float>; |
| 701 | constexpr int fVecSize = fVec::size(); |
| 702 | constexpr int bVecSize = bVec::size(); |
| 703 | |
| 704 | for (int i = 0; i < size; ++i) { |
| 705 | int d = 0; |
| 706 | for (; d <= D - bVecSize; d += bVecSize) { |
| 707 | fVec f0 = fVec::loadu(input + i * ld_src + d); |
| 708 | fVec f1 = fVec::loadu(input + i * ld_src + d + fVecSize); |
| 709 | bVec out_bvec = at::vec::convert_from_float<scalar_t>(f0, f1); |
| 710 | out_bvec.store(out + i * ld_dst + d); |
| 711 | } |
| 712 | for (; d < D; ++d) { |
| 713 | out[i * ld_dst + d] = static_cast<scalar_t>(input[i * ld_src + d]); |
| 714 | } |
| 715 | } |
| 716 | } |
| 717 | }; |
| 718 | |
| 719 | #if defined(CPU_CAPABILITY_AVX512) |
| 720 | template <int D> |
| 721 | struct update_kernel<at::BFloat16, D> { |
| 722 | static inline void |
| 723 | apply(at::BFloat16* __restrict__ out, const float* __restrict__ input, int size, int ld_src, int ld_dst) { |
| 724 | static_assert(D % 32 == 0); |
| 725 | constexpr int COLS = D / 16; |
| 726 | |
| 727 | for (int i = 0; i < size; ++i) { |
| 728 | Unroll<COLS>{}([&](auto col) { |
| 729 | if constexpr (col % 2 == 0) { |
| 730 | __m512 va0 = _mm512_loadu_ps(input + i * ld_src + (col + 0) * 16); |
| 731 | __m512 va1 = _mm512_loadu_ps(input + i * ld_src + (col + 1) * 16); |
| 732 | __m512i a16 = (__m512i)(_mm512_cvtne2ps_pbh(va1, va0)); |
| 733 | _mm512_storeu_si512(out + i * ld_dst + col * 16, a16); |
| 734 | } |
| 735 | }); |
| 736 | } |
| 737 | } |
| 738 | }; |
| 739 | #endif |
| 740 | |
| 741 | template <typename scalar_t, int D> |
| 742 | struct update_value_kernel { |
| 743 | static inline void apply( |
| 744 | scalar_t* __restrict__ v_prime2, |
| 745 | const scalar_t* __restrict__ v, |
| 746 | const float* __restrict__ v_prime, |
| 747 | int size, |
| 748 | int padded_size, |
| 749 | int v_strideT) { |
| 750 | using bVec = at::vec::Vectorized<scalar_t>; |
| 751 | using fVec = at::vec::Vectorized<float>; |
| 752 | constexpr int fVecSize = fVec::size(); |
| 753 | constexpr int bVecSize = bVec::size(); |
| 754 | |
| 755 | // v2' = v - v' |
| 756 | for (int i = 0; i < size; ++i) { |
| 757 | int d = 0; |
| 758 | for (; d <= D - bVecSize; d += bVecSize) { |
| 759 | bVec v_bvec = bVec::loadu(v + i * v_strideT + d); |
| 760 | fVec v0, v1; |
| 761 | std::tie(v0, v1) = at::vec::convert_to_float(v_bvec); |
| 762 | fVec vp0 = fVec::loadu(v_prime + i * D + d); |
| 763 | fVec vp1 = fVec::loadu(v_prime + i * D + d + fVecSize); |
| 764 | v0 = v0 - vp0; |
| 765 | v1 = v1 - vp1; |
| 766 | bVec out_bvec = at::vec::convert_from_float<scalar_t>(v0, v1); |
| 767 | out_bvec.store(v_prime2 + i * D + d); |
| 768 | } |
| 769 | for (; d < D; ++d) { |
| 770 | float val = static_cast<float>(v[i * v_strideT + d]) - v_prime[i * D + d]; |
| 771 | v_prime2[i * D + d] = static_cast<scalar_t>(val); |
| 772 | } |
| 773 | } |
| 774 | // pad the last chunk |
| 775 | const bVec zero_bvec(static_cast<scalar_t>(0)); |
| 776 | for (int i = size; i < padded_size; ++i) { |
| 777 | int d = 0; |
| 778 | for (; d <= D - bVecSize; d += bVecSize) { |
| 779 | zero_bvec.store(v_prime2 + i * D + d); |
| 780 | } |
| 781 | for (; d < D; ++d) { |
| 782 | v_prime2[i * D + d] = static_cast<scalar_t>(0); |
| 783 | } |
| 784 | } |
| 785 | } |
| 786 | }; |
| 787 | |
| 788 | #if defined(CPU_CAPABILITY_AVX512) |
| 789 | template <int D> |
| 790 | struct update_value_kernel<at::BFloat16, D> { |
| 791 | static inline void apply( |
| 792 | at::BFloat16* __restrict__ v_prime2, |
| 793 | const at::BFloat16* __restrict__ v, |
| 794 | const float* __restrict__ v_prime, |
| 795 | int size, |
| 796 | int padded_size, |
| 797 | int v_strideT) { |
| 798 | static_assert(D % 32 == 0); |
| 799 | constexpr int COLS = D / 16; |
| 800 | |
| 801 | // v2' = v - v' |
| 802 | for (int i = 0; i < size; ++i) { |
| 803 | Unroll<COLS>{}([&](auto col) { |
| 804 | // load for 0, 2, 4, 6 |
| 805 | if constexpr (col % 2 == 0) { |
| 806 | __m512i v16 = _mm512_loadu_si512(v + i * v_strideT + col * 16); |
| 807 | __m512 va0 = CVT_BF16_TO_FP32(_mm512_extracti32x8_epi32(v16, 0)); |
| 808 | __m512 va1 = CVT_BF16_TO_FP32(_mm512_extracti32x8_epi32(v16, 1)); |
| 809 | |
| 810 | __m512 v_prime0 = _mm512_loadu_ps(v_prime + i * D + col * 16); |
| 811 | __m512 v_prime1 = _mm512_loadu_ps(v_prime + i * D + col * 16 + 16); |
| 812 | va0 = _mm512_sub_ps(va0, v_prime0); |
| 813 | va1 = _mm512_sub_ps(va1, v_prime1); |
| 814 | __m512i o16 = (__m512i)(_mm512_cvtne2ps_pbh(va1, va0)); |
| 815 | _mm512_storeu_si512(v_prime2 + i * D + col * 16, o16); |
| 816 | } |
| 817 | }); |
| 818 | } |
| 819 | |
| 820 | // pad the last chunk |
| 821 | for (int i = size; i < padded_size; ++i) { |
| 822 | Unroll<COLS>{}([&](auto col) { |
| 823 | if constexpr (col % 2 == 0) { |
| 824 | __m512i v16 = _mm512_setzero_si512(); |
| 825 | _mm512_storeu_si512(v_prime2 + i * D + col * 16, v16); |
| 826 | } |
| 827 | }); |
| 828 | } |
| 829 | } |
| 830 | }; |
| 831 | #endif |
| 832 | |
| 833 | template <typename scalar_t, int CHUNK_SIZE, int D> |
| 834 | struct update_key_kernel { |
| 835 | static inline void apply( |
| 836 | scalar_t* __restrict__ k_updated, |
| 837 | const scalar_t* __restrict__ k, |
| 838 | const float* __restrict__ g, |
| 839 | int size, |
| 840 | int k_strideT) { |
| 841 | // k_updated is transposed: [D, CHUNK_SIZE], k_updated[d, t] = k[t, d] * exp(g_last - g[t]). |
| 842 | // k's D dim is contiguous (vectorize load/compute); k_updated's D dim is |
| 843 | // strided (CHUNK_SIZE apart), so the store is unpacked per lane. |
| 844 | using bVec = at::vec::Vectorized<scalar_t>; |
| 845 | using fVec = at::vec::Vectorized<float>; |
| 846 | constexpr int fVecSize = fVec::size(); |
| 847 | constexpr int bVecSize = bVec::size(); |
| 848 | alignas(64) scalar_t lane_buf[bVecSize]; |
| 849 | |
| 850 | const float g_last = g[size - 1]; |
| 851 | for (int t = 0; t < size; ++t) { |
| 852 | float scale = std::exp(g_last - g[t]); |
| 853 | fVec scale_fvec(scale); |
| 854 | int d = 0; |
| 855 | for (; d <= D - bVecSize; d += bVecSize) { |
| 856 | bVec k_bvec = bVec::loadu(k + t * k_strideT + d); |
| 857 | fVec k0, k1; |
| 858 | std::tie(k0, k1) = at::vec::convert_to_float(k_bvec); |
| 859 | k0 = k0 * scale_fvec; |
| 860 | k1 = k1 * scale_fvec; |
| 861 | bVec out_bvec = at::vec::convert_from_float<scalar_t>(k0, k1); |
| 862 | out_bvec.store(lane_buf); |
| 863 | for (int lane = 0; lane < bVecSize; ++lane) { |
| 864 | k_updated[(d + lane) * CHUNK_SIZE + t] = lane_buf[lane]; |
| 865 | } |
| 866 | } |
| 867 | for (; d < D; ++d) { |
| 868 | k_updated[d * CHUNK_SIZE + t] = static_cast<scalar_t>(static_cast<float>(k[t * k_strideT + d]) * scale); |
| 869 | } |
| 870 | } |
| 871 | const bVec zero_bvec(static_cast<scalar_t>(0)); |
| 872 | for (int t = size; t < CHUNK_SIZE; ++t) { |
| 873 | int d = 0; |
| 874 | for (; d <= D - bVecSize; d += bVecSize) { |
| 875 | zero_bvec.store(lane_buf); |
| 876 | for (int lane = 0; lane < bVecSize; ++lane) { |
| 877 | k_updated[(d + lane) * CHUNK_SIZE + t] = lane_buf[lane]; |
| 878 | } |
| 879 | } |
| 880 | for (; d < D; ++d) { |
| 881 | k_updated[d * CHUNK_SIZE + t] = static_cast<scalar_t>(0); |
| 882 | } |
| 883 | } |
| 884 | } |
| 885 | }; |
| 886 | |
| 887 | #if defined(CPU_CAPABILITY_AVX512) |
| 888 | template <int CHUNK_SIZE, int D> |
| 889 | struct update_key_kernel<at::BFloat16, CHUNK_SIZE, D> { |
| 890 | static inline void apply( |
| 891 | at::BFloat16* __restrict__ k_updated, |
| 892 | const at::BFloat16* __restrict__ k, |
| 893 | const float* __restrict__ g, |
| 894 | int size, |
| 895 | int k_strideT) { |
| 896 | static_assert(D % 32 == 0); |
| 897 | const int MB = div_up(size, 16); |
| 898 | const int KB = D / 16; |
| 899 | |
| 900 | const float g_last = g[size - 1]; |
| 901 | const __m512 vg_last = _mm512_set1_ps(g_last); |
| 902 | |
| 903 | float scale_arr[16]; |
| 904 | __m256i va[16]; |
| 905 | |
| 906 | // from [C, D](MB, KB) to [D, C](KB, MB) |
| 907 | // pad size to 16 in this kernel so that transpose can be done in one loop |
| 908 | for (int mb = 0; mb < MB; ++mb) { |
| 909 | const int mb_size = std::min(size - mb * 16, 16); |
| 910 | // prepare exp(g_last - g) |
| 911 | __m512 vg = _mm512_loadu_ps(g + mb * 16); |
| 912 | _mm512_storeu_ps(scale_arr, _mm512_fexp_u20_ps(_mm512_sub_ps(vg_last, vg))); |
| 913 | for (int kb = 0; kb < KB; ++kb) { |
| 914 | const at::BFloat16* __restrict__ k_ptr = k + mb * 16 * k_strideT + kb * 16; |
| 915 | at::BFloat16* __restrict__ k_updated_ptr = k_updated + kb * 16 * CHUNK_SIZE + mb * 16; |
| 916 | // load 16 regs |
| 917 | Unroll<16>{}([&](auto m) { |
| 918 | if (m < mb_size) { |
| 919 | __m256i v16 = _mm256_loadu_si256(reinterpret_cast<const __m256i*>(k_ptr + m * k_strideT)); |
| 920 | __m512 v32 = _mm512_mul_ps(CVT_BF16_TO_FP32(v16), _mm512_set1_ps(scale_arr[m])); |
| 921 | va[m] = (__m256i)_mm512_cvtneps_pbh(v32); |
| 922 | } else { |
| 923 | va[m] = _mm256_setzero_si256(); |
| 924 | } |
| 925 | }); |
| 926 | // transpose 16x16 |
| 927 | transpose_16x16_16bit(va); |
| 928 | // store 16 regs |
| 929 | Unroll<16>{}( |
| 930 | [&](auto k) { _mm256_storeu_si256(reinterpret_cast<__m256i*>(k_updated_ptr + k * CHUNK_SIZE), va[k]); }); |
| 931 | } |
| 932 | } |
| 933 | } |
| 934 | }; |
| 935 | #endif |
| 936 | |
| 937 | |
| 938 | // template head_dim here to reduce extra read |
| 939 | // * normal approach: read inputs 2 times: |
| 940 | // - reduce: 1R |
| 941 | // - scale: 1R + 1W |
| 942 | // * keep input data in register: |
| 943 | // - reduce: 1R |
| 944 | // - scale: 1W |
| 945 | template <typename scalar_t, int D> |
| 946 | void l2norm_fwd_kernel_impl( |
| 947 | scalar_t* __restrict__ query_norm, |
| 948 | scalar_t* __restrict__ key_norm, |
| 949 | const scalar_t* __restrict__ query, |
| 950 | const scalar_t* __restrict__ key, |
| 951 | float eps, |
| 952 | int64_t T, |
| 953 | int64_t H, |
| 954 | int64_t q_strideT, |
| 955 | int64_t q_strideH, |
| 956 | int64_t k_strideT, |
| 957 | int64_t k_strideH) { |
| 958 | // expected to be contuguous |
| 959 | int64_t qn_strideH = D; |
| 960 | int64_t kn_strideH = D; |
| 961 | |
| 962 | // parallel on [B, T, H] |
| 963 | at::parallel_for(0, T * H, 0, [&](int64_t begin, int64_t end) { |
| 964 | int64_t t{0}, h{0}; |
| 965 | data_index_init(begin, t, T, h, H); |
| 966 | |
| 967 | for (int64_t i = begin; i < end; ++i) { |
| 968 | const scalar_t* __restrict__ q_ptr = query + t * q_strideT + h * q_strideH; |
| 969 | const scalar_t* __restrict__ k_ptr = key + t * k_strideT + h * k_strideH; |
| 970 | scalar_t* __restrict__ qn_ptr = query_norm + i * qn_strideH; |
| 971 | scalar_t* __restrict__ kn_ptr = key_norm + i * kn_strideH; |
| 972 | |
| 973 | l2norm_kernel<scalar_t, D, true>::apply(qn_ptr, q_ptr, eps); |
| 974 | l2norm_kernel<scalar_t, D, false>::apply(kn_ptr, k_ptr, eps); |
| 975 | |
| 976 | // move to the next index |
| 977 | data_index_step(t, T, h, H); |
| 978 | } |
| 979 | }); |
| 980 | } |
| 981 | |
| 982 | // g : [B, T, Hv] |
| 983 | // g_ : [B, Hv, NT, C] -> [B, NT, HB, BLOCK_H, C] |
| 984 | // cu_seqlens : [num_seqs + 1] |
| 985 | // chunk_indices : [NT * 2] |
| 986 | template <typename scalar_t, int CHUNK_SIZE> |
| 987 | void chunk_local_cumsum_kernel_impl( |
| 988 | scalar_t* __restrict__ g_, |
| 989 | const scalar_t* __restrict__ g, |
| 990 | const int32_t* __restrict__ cu_seqlens, |
| 991 | const int32_t* __restrict__ chunk_indices, |
| 992 | int64_t Hv, |
| 993 | int64_t NT) { |
| 994 | constexpr int BLOCK_H = 16; |
| 995 | int64_t HB = div_up(Hv, int64_t(BLOCK_H)); |
| 996 | |
| 997 | // parallel on [NT * HB] to increase parallelism |
| 998 | at::parallel_for(0, NT * HB, 0, [&](int64_t begin, int64_t end) { |
| 999 | int64_t nt{0}, hb{0}; |
| 1000 | data_index_init(begin, nt, NT, hb, HB); |
| 1001 | |
| 1002 | for (int64_t i = begin; i < end; ++i) { |
| 1003 | int32_t bs = chunk_indices[nt * 2 + 0]; |
| 1004 | int32_t batch_offset = cu_seqlens[bs]; |
| 1005 | int32_t seqlen = cu_seqlens[bs + 1] - cu_seqlens[bs]; |
| 1006 | int64_t mb_start = chunk_indices[nt * 2 + 1] * CHUNK_SIZE; |
| 1007 | int64_t mb_size = std::min(seqlen - mb_start, int64_t(CHUNK_SIZE)); |
| 1008 | int64_t hb_size = std::min(Hv - hb * BLOCK_H, int64_t(BLOCK_H)); |
| 1009 | |
| 1010 | const scalar_t* __restrict__ g_ptr = g + (batch_offset + mb_start) * Hv + hb * BLOCK_H; |
| 1011 | scalar_t* __restrict__ gsum_ptr = g_ + nt * (Hv * CHUNK_SIZE) + hb * (BLOCK_H * CHUNK_SIZE); |
| 1012 | cumsum_kernel<scalar_t, CHUNK_SIZE, BLOCK_H>::apply(gsum_ptr, g_ptr, mb_size, hb_size, Hv, CHUNK_SIZE); |
| 1013 | |
| 1014 | // move to the next index |
| 1015 | data_index_step(nt, NT, hb, HB); |
| 1016 | } |
| 1017 | }); |
| 1018 | } |
| 1019 | |
| 1020 | #define DECL_BUF(type, name, size_expr) alignas(64) type name[(size_expr)] |
| 1021 | #define DECL_ZERO_BUF(type, name, size_expr) \ |
| 1022 | DECL_BUF(type, name, size_expr); \ |
| 1023 | fill_stub<type, (size_expr)>(name, 0.f) |
| 1024 | |
| 1025 | // w : [B, T, Hv, D] |
| 1026 | // u : [B, T, Hv, Dv] |
| 1027 | // d : [B, NT, Hv, C, C] |
| 1028 | // k : [B, T, H, D] |
| 1029 | // v : [B, T, Hv, Dv] |
| 1030 | // g : [B, NT, Hv, C] |
| 1031 | // beta : [B, T, Hv] |
| 1032 | // cu_seqlens : [num_seqs + 1] |
| 1033 | // chunk_indices : [NT * 2] |
| 1034 | template <typename scalar_t, int D, int CHUNK_SIZE> |
| 1035 | void chunk_gated_delta_rule_fwd_intra_kernel_impl( |
| 1036 | scalar_t* __restrict__ w, |
| 1037 | scalar_t* __restrict__ u, |
| 1038 | float* __restrict__ d, |
| 1039 | const scalar_t* __restrict__ k, |
| 1040 | const scalar_t* __restrict__ v, |
| 1041 | const float* __restrict__ g, |
| 1042 | const scalar_t* __restrict__ beta, |
| 1043 | const int32_t* __restrict__ cu_seqlens, |
| 1044 | const int32_t* __restrict__ chunk_indices, |
| 1045 | int64_t H, |
| 1046 | int64_t Hv, |
| 1047 | int64_t NT, |
| 1048 | int64_t k_strideT, |
| 1049 | int64_t k_strideH, |
| 1050 | int64_t v_strideT, |
| 1051 | int64_t v_strideH) { |
| 1052 | // head group, expect to be 1，2，4 for qwen3.5 |
| 1053 | const int64_t HG = Hv / H; |
| 1054 | |
| 1055 | // strides |
| 1056 | const int64_t w_strideT = Hv * D; |
| 1057 | const int64_t w_strideH = D; |
| 1058 | const int64_t u_strideT = Hv * D; |
| 1059 | const int64_t u_strideH = D; |
| 1060 | |
| 1061 | // [NB]: parallel on [NT, H] |
| 1062 | // * parallel on num_heads and go sequential on num_heads_v, |
| 1063 | // * avoid instantialize k_beta (beta * k) |
| 1064 | // * compute key @ key^T * beta instead of k_beta @ key^T, same as triton impl |
| 1065 | // * compute key @ key^T once for each k head index and reuse for v head index |
| 1066 | at::parallel_for(0, NT * H, 0, [&](int64_t begin, int64_t end) { |
| 1067 | int64_t nt{0}, h{0}; |
| 1068 | data_index_init(begin, nt, NT, h, H); |
| 1069 | |
| 1070 | // thread local temp buffer |
| 1071 | DECL_ZERO_BUF(scalar_t, tmp, CHUNK_SIZE * D); |
| 1072 | DECL_ZERO_BUF(scalar_t, tmp2, CHUNK_SIZE * D); |
| 1073 | DECL_ZERO_BUF(float, attn, CHUNK_SIZE* CHUNK_SIZE); |
| 1074 | DECL_ZERO_BUF(scalar_t, attn2, CHUNK_SIZE * CHUNK_SIZE); |
| 1075 | DECL_ZERO_BUF(float, tmp3, CHUNK_SIZE* D); |
| 1076 | |
| 1077 | // alias |
| 1078 | scalar_t* __restrict__ k_packed = tmp; |
| 1079 | scalar_t* __restrict__ k_beta = tmp; |
| 1080 | scalar_t* __restrict__ v_beta = tmp; |
| 1081 | scalar_t* __restrict__ k_beta_packed = tmp2; |
| 1082 | scalar_t* __restrict__ v_beta_packed = tmp2; |
| 1083 | float* __restrict__ k_updated = tmp3; |
| 1084 | float* __restrict__ v_updated = tmp3; |
| 1085 | |
| 1086 | for (int64_t i = begin; i < end; ++i) { |
| 1087 | int32_t bs = chunk_indices[nt * 2 + 0]; |
| 1088 | int32_t batch_offset = cu_seqlens[bs]; |
| 1089 | int32_t seqlen = cu_seqlens[bs + 1] - cu_seqlens[bs]; |
| 1090 | int64_t mb_start = chunk_indices[nt * 2 + 1] * CHUNK_SIZE; |
| 1091 | int64_t mb_size = std::min(seqlen - mb_start, int64_t(CHUNK_SIZE)); |
| 1092 | |
| 1093 | // mb_size` is K in 5.c, 5.g, pad to TILE_K; |
| 1094 | const int64_t padded_mb_size = div_up((int)mb_size, TILE_K) * TILE_K; |
| 1095 | |
| 1096 | // step 1: decay_mask = ((g.unsqueeze(-1) - g.unsqueeze(-2)).tril().exp().float()).tril() |
| 1097 | for (int64_t hv = h * HG; hv < h * HG + HG; ++hv) { |
| 1098 | const float* __restrict__ g_ptr = g + nt * (Hv * CHUNK_SIZE) + hv * CHUNK_SIZE; |
| 1099 | float* __restrict__ d_ptr = d + nt * (Hv * CHUNK_SIZE * CHUNK_SIZE) + hv * (CHUNK_SIZE * CHUNK_SIZE); |
| 1100 | decay_mask_kernel<float, CHUNK_SIZE>::apply(d_ptr, g_ptr); |
| 1101 | } |
| 1102 | |
| 1103 | // step 2: attn = key @ key^T |
| 1104 | const scalar_t* __restrict__ k_ptr = k + (batch_offset + mb_start) * k_strideT + h * k_strideH; |
| 1105 | if constexpr (brgemm_supported()) { |
| 1106 | pack_vnni<scalar_t>( |
| 1107 | /* dst */ k_packed, |
| 1108 | /* src */ k_ptr, |
| 1109 | /* N */ mb_size, |
| 1110 | /* K */ D, |
| 1111 | /* ld_src */ k_strideT, |
| 1112 | /* ld_dst */ CHUNK_SIZE); |
| 1113 | |
| 1114 | at::native::cpublas::brgemm( |
| 1115 | /* M */ mb_size, |
| 1116 | /* N */ mb_size, |
| 1117 | /* K */ D, |
| 1118 | /* lda */ k_strideT, |
| 1119 | /* ldb */ CHUNK_SIZE, |
| 1120 | /* ldc */ CHUNK_SIZE, |
| 1121 | /* add_C */ false, |
| 1122 | /* A */ k_ptr, |
| 1123 | /* B */ k_packed, |
| 1124 | /* C */ attn); |
| 1125 | } else { |
| 1126 | blas_gemm( |
| 1127 | at::native::TransposeType::Transpose, |
| 1128 | at::native::TransposeType::NoTranspose, |
| 1129 | mb_size, |
| 1130 | mb_size, |
| 1131 | D, |
| 1132 | 1.0f, |
| 1133 | k_ptr, |
| 1134 | k_strideT, |
| 1135 | k_ptr, |
| 1136 | k_strideT, |
| 1137 | 0.0f, |
| 1138 | attn, |
| 1139 | CHUNK_SIZE); |
| 1140 | } |
| 1141 | |
| 1142 | for (int64_t hv = h * HG; hv < h * HG + HG; ++hv) { |
| 1143 | // step 3: attn2 = -attn * beta * d |
| 1144 | const scalar_t* __restrict__ beta_ptr = beta + (batch_offset + mb_start) * Hv + hv; |
| 1145 | const float* __restrict__ d_ptr = d + nt * (Hv * CHUNK_SIZE * CHUNK_SIZE) + hv * (CHUNK_SIZE * CHUNK_SIZE); |
| 1146 | apply_mask_kernel<scalar_t, CHUNK_SIZE, true>::apply(attn2, attn, beta_ptr, d_ptr, mb_size, Hv); |
| 1147 | |
| 1148 | // step 4: solve_tril(attn2) -> (I + L)^{-1}, L = strict-lower from step 3 |
| 1149 | // for i in 1..C-1: attn2[i, :i] += (attn2[i, :i] * attn2[:i, :i]).sum(-1) |
| 1150 | // attn2 += eye(C) |
| 1151 | solve_tril_kernel<scalar_t, CHUNK_SIZE>::apply(attn2, mb_size); |
| 1152 | |
| 1153 | // step 5: recompute_w_u |
| 1154 | // w = attn2 @ (k_beta * g.exp().unsqueeze(-1)) |
| 1155 | // u = attn2 @ value * beta.unsqueeze(-1) |
| 1156 | const float* __restrict__ g_ptr = g + nt * (Hv * CHUNK_SIZE) + hv * CHUNK_SIZE; |
| 1157 | const scalar_t* __restrict__ v_ptr = v + (batch_offset + mb_start) * v_strideT + hv * v_strideH; |
| 1158 | |
| 1159 | // 5.a key = key * beta * g.exp |
| 1160 | apply_beta_kernel<scalar_t, CHUNK_SIZE, D, true, true>::apply( |
| 1161 | k_beta, k_ptr, beta_ptr, g_ptr, mb_size, k_strideT, D, Hv); |
| 1162 | |
| 1163 | // 5.b pack key |
| 1164 | if constexpr (brgemm_supported()) { |
| 1165 | pack_vnni2<scalar_t>( |
| 1166 | /* dst */ k_beta_packed, |
| 1167 | /* src */ k_beta, |
| 1168 | /* K */ mb_size, |
| 1169 | /* N */ D, |
| 1170 | /* ld_src */ D, |
| 1171 | /* ld_dst */ D); |
| 1172 | |
| 1173 | // 5.c w = attn2 @ k_beta |
| 1174 | at::native::cpublas::brgemm( |
| 1175 | /* M */ mb_size, |
| 1176 | /* N */ D, |
| 1177 | /* K */ padded_mb_size, // mb_size |
| 1178 | /* lda */ CHUNK_SIZE, |
| 1179 | /* ldb */ D, |
| 1180 | /* ldc */ D, |
| 1181 | /* add_C */ false, |
| 1182 | /* A */ attn2, |
| 1183 | /* B */ k_beta_packed, |
| 1184 | /* C */ k_updated); |
| 1185 | } else { |
| 1186 | blas_gemm( |
| 1187 | at::native::TransposeType::NoTranspose, |
| 1188 | at::native::TransposeType::NoTranspose, |
| 1189 | D, |
| 1190 | mb_size, |
| 1191 | padded_mb_size, |
| 1192 | 1.0f, |
| 1193 | k_beta, |
| 1194 | D, |
| 1195 | attn2, |
| 1196 | CHUNK_SIZE, |
| 1197 | 0.0f, |
| 1198 | k_updated, |
| 1199 | D); |
| 1200 | } |
| 1201 | |
| 1202 | // 5.d k_updated -> w |
| 1203 | scalar_t* __restrict__ w_ptr = w + (batch_offset + mb_start) * w_strideT + hv * w_strideH; |
| 1204 | update_kernel<scalar_t, D>::apply(w_ptr, k_updated, mb_size, D, w_strideT); |
| 1205 | |
| 1206 | // 5.e value = value * beta |
| 1207 | apply_beta_kernel<scalar_t, CHUNK_SIZE, D, true, false>::apply( |
| 1208 | v_beta, v_ptr, beta_ptr, nullptr, mb_size, v_strideT, D, Hv); |
| 1209 | |
| 1210 | // 5.f pack value |
| 1211 | if constexpr (brgemm_supported()) { |
| 1212 | pack_vnni2<scalar_t>( |
| 1213 | /* dst */ v_beta_packed, |
| 1214 | /* src */ v_beta, |
| 1215 | /* K */ mb_size, |
| 1216 | /* N */ D, |
| 1217 | /* ld_src */ D, |
| 1218 | /* ld_dst */ D); |
| 1219 | |
| 1220 | // 5.g u = attn2 @ v_beta |
| 1221 | at::native::cpublas::brgemm( |
| 1222 | /* M */ mb_size, |
| 1223 | /* N */ D, |
| 1224 | /* K */ padded_mb_size, // mb_size |
| 1225 | /* lda */ CHUNK_SIZE, |
| 1226 | /* ldb */ D, |
| 1227 | /* ldc */ D, |
| 1228 | /* add_C */ false, |
| 1229 | /* A */ attn2, |
| 1230 | /* B */ v_beta_packed, |
| 1231 | /* C */ v_updated); |
| 1232 | } else { |
| 1233 | blas_gemm( |
| 1234 | at::native::TransposeType::NoTranspose, |
| 1235 | at::native::TransposeType::NoTranspose, |
| 1236 | D, |
| 1237 | mb_size, |
| 1238 | padded_mb_size, |
| 1239 | 1.0f, |
| 1240 | v_beta, |
| 1241 | D, |
| 1242 | attn2, |
| 1243 | CHUNK_SIZE, |
| 1244 | 0.0f, |
| 1245 | v_updated, |
| 1246 | D); |
| 1247 | } |
| 1248 | |
| 1249 | // 5.h v_updated -> u |
| 1250 | scalar_t* __restrict__ u_ptr = u + (batch_offset + mb_start) * u_strideT + hv * u_strideH; |
| 1251 | update_kernel<scalar_t, D>::apply(u_ptr, v_updated, mb_size, D, u_strideT); |
| 1252 | } |
| 1253 | |
| 1254 | // move to the next index |
| 1255 | data_index_step(nt, NT, h, H); |
| 1256 | } |
| 1257 | at::native::cpublas::brgemm_release(); |
| 1258 | }); |
| 1259 | } |
| 1260 | |
| 1261 | // |
| 1262 | // out : [B, T, Hv, Dv] |
| 1263 | // state : [num_seqs, Hv, Dv, D] |
| 1264 | // q : [B, T, H, D] |
| 1265 | // k : [B, T, H, D] |
| 1266 | // w : [B, T, Hv, D] |
| 1267 | // u : [B, T, Hv, Dv] |
| 1268 | // g : [B, NT, Hv, C] |
| 1269 | // d : [B, NT, Hv, C, C] |
| 1270 | // cu_seqlens : [num_seqs + 1] |
| 1271 | // chunk_offsets : [num_seqs + 1] |
| 1272 | template <typename scalar_t, typename state_t, int D, int CHUNK_SIZE> |
| 1273 | void chunk_gated_delta_rule_fwd_inter_kernel_impl( |
| 1274 | scalar_t* __restrict__ out, |
| 1275 | state_t* __restrict__ state, |
| 1276 | const int32_t* __restrict__ indices, |
| 1277 | const scalar_t* __restrict__ q, |
| 1278 | const scalar_t* __restrict__ k, |
| 1279 | const scalar_t* __restrict__ w, |
| 1280 | const scalar_t* __restrict__ u, |
| 1281 | const float* __restrict__ g, |
| 1282 | const float* __restrict__ d, |
| 1283 | const int32_t* __restrict__ cu_seqlens, |
| 1284 | const int32_t* __restrict__ chunk_offsets, |
| 1285 | int64_t H, |
| 1286 | int64_t Hv, |
| 1287 | int64_t num_seqs, |
| 1288 | int64_t q_strideT, |
| 1289 | int64_t q_strideH, |
| 1290 | int64_t k_strideT, |
| 1291 | int64_t k_strideH, |
| 1292 | int64_t state_strideS) { |
| 1293 | // head group, expect to be 1，2，4 for qwen3.5 |
| 1294 | const int64_t HG = Hv / H; |
| 1295 | |
| 1296 | // strides |
| 1297 | const int64_t w_strideT = Hv * D; |
| 1298 | const int64_t w_strideH = D; |
| 1299 | const int64_t u_strideT = Hv * D; |
| 1300 | const int64_t u_strideH = D; |
| 1301 | const int64_t o_strideT = Hv * D; |
| 1302 | const int64_t o_strideH = D; |
| 1303 | |
| 1304 | // [NB]: parallel on [num_seqs, Hv] |
| 1305 | // * choose to parallel on Hv instead of H, though this means q @ kT has duplicated compute |
| 1306 | // * H might be 16 which is not enough to use 32C when num_seqs is small |
| 1307 | at::parallel_for(0, num_seqs * Hv, 0, [&](int64_t begin, int64_t end) { |
| 1308 | int64_t bs{0}, hv{0}; |
| 1309 | data_index_init(begin, bs, num_seqs, hv, Hv); |
| 1310 | |
| 1311 | // thread local temp buffer |
| 1312 | DECL_ZERO_BUF(scalar_t, tmp, CHUNK_SIZE * D); |
| 1313 | DECL_ZERO_BUF(scalar_t, tmp2, D * D); |
| 1314 | DECL_ZERO_BUF(float, state_tmp, D * D); |
| 1315 | DECL_ZERO_BUF(float, tmp3, CHUNK_SIZE* D); |
| 1316 | DECL_ZERO_BUF(scalar_t, tmp4, CHUNK_SIZE * D); |
| 1317 | DECL_ZERO_BUF(float, attn, CHUNK_SIZE* CHUNK_SIZE); |
| 1318 | DECL_ZERO_BUF(scalar_t, attn2, CHUNK_SIZE * CHUNK_SIZE); |
| 1319 | |
| 1320 | // alias |
| 1321 | scalar_t* __restrict__ k_packed = tmp; |
| 1322 | scalar_t* __restrict__ s_packed = tmp2; |
| 1323 | float* __restrict__ v_prime = tmp3; |
| 1324 | scalar_t* __restrict__ v_prime2 = tmp; |
| 1325 | float* __restrict__ attn_inter = tmp3; |
| 1326 | scalar_t* __restrict__ qg_exp = tmp4; |
| 1327 | scalar_t* __restrict__ v_packed = tmp4; |
| 1328 | scalar_t* __restrict__ k_updated = tmp; |
| 1329 | |
| 1330 | for (int64_t i = begin; i < end; ++i) { |
| 1331 | int64_t h = hv / HG; |
| 1332 | int32_t batch_offset = cu_seqlens[bs]; |
| 1333 | int32_t seqlen = cu_seqlens[bs + 1] - cu_seqlens[bs]; |
| 1334 | int64_t nt = chunk_offsets[bs]; |
| 1335 | |
| 1336 | for (int64_t mb_start = 0; mb_start < seqlen; mb_start += CHUNK_SIZE, ++nt) { |
| 1337 | int64_t mb_size = std::min(seqlen - mb_start, int64_t(CHUNK_SIZE)); |
| 1338 | |
| 1339 | // mb_size` is K in 4.a, pad to TILE_K; |
| 1340 | const int64_t padded_mb_size = div_up((int)mb_size, TILE_K) * TILE_K; |
| 1341 | |
| 1342 | // step 1.a: attn = query @ key^T |
| 1343 | // attn_i = (q_i @ k_i.transpose(-1, -2) * decay_mask[:, :, i]).masked_fill_(mask, 0) |
| 1344 | const scalar_t* __restrict__ q_ptr = q + (batch_offset + mb_start) * q_strideT + h * q_strideH; |
| 1345 | const scalar_t* __restrict__ k_ptr = k + (batch_offset + mb_start) * k_strideT + h * k_strideH; |
| 1346 | if constexpr (brgemm_supported()) { |
| 1347 | pack_vnni<scalar_t>( |
| 1348 | /* dst */ k_packed, |
| 1349 | /* src */ k_ptr, |
| 1350 | /* N */ mb_size, |
| 1351 | /* K */ D, |
| 1352 | /* ld_src */ k_strideT, |
| 1353 | /* ld_dst */ CHUNK_SIZE); |
| 1354 | |
| 1355 | at::native::cpublas::brgemm( |
| 1356 | /* M */ mb_size, |
| 1357 | /* N */ mb_size, |
| 1358 | /* K */ D, |
| 1359 | /* lda */ q_strideT, |
| 1360 | /* ldb */ CHUNK_SIZE, |
| 1361 | /* ldc */ CHUNK_SIZE, |
| 1362 | /* add_C */ false, |
| 1363 | /* A */ q_ptr, |
| 1364 | /* B */ k_packed, |
| 1365 | /* C */ attn); |
| 1366 | } else { |
| 1367 | blas_gemm( |
| 1368 | at::native::TransposeType::Transpose, |
| 1369 | at::native::TransposeType::NoTranspose, |
| 1370 | mb_size, |
| 1371 | mb_size, |
| 1372 | D, |
| 1373 | 1.0f, |
| 1374 | k_ptr, |
| 1375 | k_strideT, |
| 1376 | q_ptr, |
| 1377 | q_strideT, |
| 1378 | 0.0f, |
| 1379 | attn, |
| 1380 | CHUNK_SIZE); |
| 1381 | } |
| 1382 | |
| 1383 | // step 1.b: attn = attn * decay_mask.masked_fill_(mask, 0) |
| 1384 | const float* __restrict__ d_ptr = d + nt * (Hv * CHUNK_SIZE * CHUNK_SIZE) + hv * (CHUNK_SIZE * CHUNK_SIZE); |
| 1385 | apply_mask_kernel<scalar_t, CHUNK_SIZE, false>::apply(attn2, attn, nullptr, d_ptr, mb_size); |
| 1386 | |
| 1387 | // step 2.a: v' = w @ state (fuse state *= exp(g_last) with packing) |
| 1388 | state_t* __restrict__ state_slot = |
| 1389 | state + indices[bs] * state_strideS + hv * (D * D); |
| 1390 | float* s_ptr; |
| 1391 | if constexpr (std::is_same_v<state_t, float>) { |
| 1392 | s_ptr = state_slot; |
| 1393 | } else { |
| 1394 | using bVec = at::vec::Vectorized<state_t>; |
| 1395 | using fVec = at::vec::Vectorized<float>; |
| 1396 | constexpr int kVecSize = bVec::size(); |
| 1397 | constexpr int fVecSize = fVec::size(); |
| 1398 | int64_t d0 = 0; |
| 1399 | for (; d0 + kVecSize <= D * D; d0 += kVecSize) { |
| 1400 | auto [state_vec0, state_vec1] = load_float_vec2(state_slot + d0); |
| 1401 | state_vec0.store(state_tmp + d0); |
| 1402 | state_vec1.store(state_tmp + d0 + fVecSize); |
| 1403 | } |
| 1404 | for (; d0 < D * D; ++d0) { |
| 1405 | state_tmp[d0] = static_cast<float>(state_slot[d0]); |
| 1406 | } |
| 1407 | s_ptr = state_tmp; |
| 1408 | } |
| 1409 | const float* __restrict__ g_ptr = g + nt * (Hv * CHUNK_SIZE) + hv * (CHUNK_SIZE); |
| 1410 | float g_last = g_ptr[mb_size - 1]; |
| 1411 | const scalar_t* __restrict__ w_ptr = w + (batch_offset + mb_start) * w_strideT + hv * w_strideH; |
| 1412 | if constexpr (brgemm_supported()) { |
| 1413 | pack_vnni2<scalar_t, D, D>( |
| 1414 | /* dst */ s_packed, |
| 1415 | /* src */ s_ptr, |
| 1416 | /* g_last */ g_last, |
| 1417 | /* ld_src */ D, |
| 1418 | /* ld_dst */ D); |
| 1419 | |
| 1420 | at::native::cpublas::brgemm( |
| 1421 | /* M */ mb_size, |
| 1422 | /* N */ D, |
| 1423 | /* K */ D, |
| 1424 | /* lda */ w_strideT, |
| 1425 | /* ldb */ D, |
| 1426 | /* ldc */ D, |
| 1427 | /* add_C */ false, |
| 1428 | /* A */ w_ptr, |
| 1429 | /* B */ s_packed, |
| 1430 | /* C */ v_prime); |
| 1431 | } else { |
| 1432 | // brgemm_supported()==false path: pack_vnni2 above packs the |
| 1433 | // *unscaled* state into its dst (for this GEMM's B operand) while |
| 1434 | // separately scaling src in place by exp(g_last) (consumed later, |
| 1435 | // at step 5.3's state accumulation). Replicate both halves in the |
| 1436 | // same order: snapshot s_ptr into s_packed BEFORE scaling s_ptr, |
| 1437 | // not after, since the GEMM below needs the pre-scale state. |
| 1438 | float g_last_scale = std::exp(g_last); |
| 1439 | for (int64_t d0 = 0; d0 < D * D; ++d0) { |
| 1440 | s_packed[d0] = static_cast<scalar_t>(s_ptr[d0]); |
| 1441 | s_ptr[d0] *= g_last_scale; |
| 1442 | } |
| 1443 | blas_gemm( |
| 1444 | at::native::TransposeType::NoTranspose, |
| 1445 | at::native::TransposeType::NoTranspose, |
| 1446 | D, |
| 1447 | mb_size, |
| 1448 | D, |
| 1449 | 1.0f, |
| 1450 | s_packed, |
| 1451 | D, |
| 1452 | w_ptr, |
| 1453 | w_strideT, |
| 1454 | 0.0f, |
| 1455 | v_prime, |
| 1456 | D); |
| 1457 | } |
| 1458 | |
| 1459 | // step 2.b: v2' = u - v' |
| 1460 | const scalar_t* __restrict__ u_ptr = u + (batch_offset + mb_start) * u_strideT + hv * u_strideH; |
| 1461 | update_value_kernel<scalar_t, D>::apply(v_prime2, u_ptr, v_prime, mb_size, padded_mb_size, u_strideT); |
| 1462 | |
| 1463 | // step 3.a: qg_exp = q * exp(g) |
| 1464 | apply_beta_kernel<scalar_t, CHUNK_SIZE, D, false, true>::apply( |
| 1465 | qg_exp, q_ptr, nullptr, g_ptr, mb_size, q_strideT, D, /*b_stride*/ 0); |
| 1466 | |
| 1467 | // step 3.b: attn_inter = qg_exp @ state |
| 1468 | if constexpr (brgemm_supported()) { |
| 1469 | at::native::cpublas::brgemm( |
| 1470 | /* M */ mb_size, |
| 1471 | /* N */ D, |
| 1472 | /* K */ D, |
| 1473 | /* lda */ D, |
| 1474 | /* ldb */ D, |
| 1475 | /* ldc */ D, |
| 1476 | /* add_C */ false, |
| 1477 | /* A */ qg_exp, |
| 1478 | /* B */ s_packed, |
| 1479 | /* C */ attn_inter); |
| 1480 | } else { |
| 1481 | blas_gemm( |
| 1482 | at::native::TransposeType::NoTranspose, |
| 1483 | at::native::TransposeType::NoTranspose, |
| 1484 | D, |
| 1485 | mb_size, |
| 1486 | D, |
| 1487 | 1.0f, |
| 1488 | s_packed, |
| 1489 | D, |
| 1490 | qg_exp, |
| 1491 | D, |
| 1492 | 0.0f, |
| 1493 | attn_inter, |
| 1494 | D); |
| 1495 | } |
| 1496 | |
| 1497 | // step 4.a: attn_inter += attn2 @ v2' |
| 1498 | if constexpr (brgemm_supported()) { |
| 1499 | pack_vnni2<scalar_t>( |
| 1500 | /* dst */ v_packed, |
| 1501 | /* src */ v_prime2, |
| 1502 | /* K */ padded_mb_size, |
| 1503 | /* N */ D, |
| 1504 | /* ld_src */ D, |
| 1505 | /* ld_dst */ D); |
| 1506 | |
| 1507 | at::native::cpublas::brgemm( |
| 1508 | /* M */ mb_size, |
| 1509 | /* N */ D, |
| 1510 | /* K */ padded_mb_size, |
| 1511 | /* lda */ CHUNK_SIZE, |
| 1512 | /* ldb */ D, |
| 1513 | /* ldc */ D, |
| 1514 | /* add_C */ true, |
| 1515 | /* A */ attn2, |
| 1516 | /* B */ v_packed, |
| 1517 | /* C */ attn_inter); |
| 1518 | } else { |
| 1519 | blas_gemm( |
| 1520 | at::native::TransposeType::NoTranspose, |
| 1521 | at::native::TransposeType::NoTranspose, |
| 1522 | D, |
| 1523 | mb_size, |
| 1524 | padded_mb_size, |
| 1525 | 1.0f, |
| 1526 | v_prime2, |
| 1527 | D, |
| 1528 | attn2, |
| 1529 | CHUNK_SIZE, |
| 1530 | 1.0f, |
| 1531 | attn_inter, |
| 1532 | D); |
| 1533 | } |
| 1534 | |
| 1535 | // step 4.b: write attn_inter -> out |
| 1536 | scalar_t* __restrict__ o_ptr = out + (batch_offset + mb_start) * o_strideT + hv * o_strideH; |
| 1537 | update_kernel<scalar_t, D>::apply(o_ptr, attn_inter, mb_size, D, o_strideT); |
| 1538 | |
| 1539 | // brgemm_supported()==false path: step 5.2 below overwrites k_updated, |
| 1540 | // which aliases the same buffer as v_prime2 (both are `tmp`). Snapshot |
| 1541 | // v_prime2 into v_packed (otherwise brgemm-only, and free by now since |
| 1542 | // qg_exp was consumed at step 3.b) before that happens, so step 5.3 |
| 1543 | // doesn't read k_updated's data under the v_prime2 name. |
| 1544 | if constexpr (!brgemm_supported()) { |
| 1545 | std::copy(v_prime2, v_prime2 + padded_mb_size * D, v_packed); |
| 1546 | } |
| 1547 | |
| 1548 | // step 5: update state |
| 1549 | // state_new = state * exp(g_last) + (k * exp(g_last - g)).T @ v2' |
| 1550 | |
| 1551 | // step 5.1 state *= exp(g_last) fused with step 2.a |
| 1552 | |
| 1553 | // step 5.2 k' = k * exp(g_last - g).T; TODO: fuse this with 1.a |
| 1554 | update_key_kernel<scalar_t, CHUNK_SIZE, D>::apply(k_updated, k_ptr, g_ptr, mb_size, k_strideT); |
| 1555 | |
| 1556 | // step 5.3 state += k' @ v2' |
| 1557 | if constexpr (brgemm_supported()) { |
| 1558 | at::native::cpublas::brgemm( |
| 1559 | /* M */ D, |
| 1560 | /* N */ D, |
| 1561 | /* K */ padded_mb_size, // mb_size |
| 1562 | /* lda */ CHUNK_SIZE, |
| 1563 | /* ldb */ D, |
| 1564 | /* ldc */ D, |
| 1565 | /* add_C */ true, |
| 1566 | /* A */ k_updated, |
| 1567 | /* B */ v_packed, |
| 1568 | /* C */ s_ptr); |
| 1569 | } else { |
| 1570 | blas_gemm( |
| 1571 | at::native::TransposeType::NoTranspose, |
| 1572 | at::native::TransposeType::NoTranspose, |
| 1573 | D, |
| 1574 | D, |
| 1575 | padded_mb_size, |
| 1576 | 1.0f, |
| 1577 | v_packed, |
| 1578 | D, |
| 1579 | k_updated, |
| 1580 | CHUNK_SIZE, |
| 1581 | 1.0f, |
| 1582 | s_ptr, |
| 1583 | D); |
| 1584 | } |
| 1585 | |
| 1586 | if constexpr (!std::is_same_v<state_t, float>) { |
| 1587 | using bVec = at::vec::Vectorized<state_t>; |
| 1588 | using fVec = at::vec::Vectorized<float>; |
| 1589 | constexpr int kVecSize = bVec::size(); |
| 1590 | constexpr int fVecSize = fVec::size(); |
| 1591 | int64_t d0 = 0; |
| 1592 | for (; d0 + kVecSize <= D * D; d0 += kVecSize) { |
| 1593 | fVec state_vec0 = fVec::loadu(s_ptr + d0); |
| 1594 | fVec state_vec1 = fVec::loadu(s_ptr + d0 + fVecSize); |
| 1595 | store_state_vec2<state_t>(state_slot + d0, state_vec0, state_vec1); |
| 1596 | } |
| 1597 | for (; d0 < D * D; ++d0) { |
| 1598 | state_slot[d0] = static_cast<state_t>(s_ptr[d0]); |
| 1599 | } |
| 1600 | } |
| 1601 | } |
| 1602 | |
| 1603 | // move to the next index |
| 1604 | data_index_step(bs, num_seqs, hv, Hv); |
| 1605 | } |
| 1606 | at::native::cpublas::brgemm_release(); |
| 1607 | }); |
| 1608 | } |
| 1609 | |
| 1610 | inline float softplus(float x, double threshold = 20.0) { |
| 1611 | if (x > threshold) |
| 1612 | return x; |
| 1613 | else if (x < -threshold) |
| 1614 | return std::exp(x); |
| 1615 | else |
| 1616 | return std::log1p(std::exp(x)); |
| 1617 | } |
| 1618 | |
| 1619 | inline at::vec::Vectorized<float> softplus(const at::vec::Vectorized<float>& x, double threshold = 20.0) { |
| 1620 | using Vec = at::vec::Vectorized<float>; |
| 1621 | Vec mask_hi = x > Vec(threshold); |
| 1622 | Vec mask_lo = x < Vec(-threshold); |
| 1623 | |
| 1624 | Vec expx = x.exp_u20(); |
| 1625 | Vec log1pex = (expx + Vec(1.0f)).log(); |
| 1626 | |
| 1627 | return Vec::blendv(Vec::blendv(log1pex, expx, mask_lo), x, mask_hi); |
| 1628 | } |
| 1629 | |
| 1630 | template <typename scalar_t, typename param_t, typename state_t> |
| 1631 | void fused_sigmoid_gating_delta_rule_update_kernel_impl( |
| 1632 | const scalar_t* __restrict__ q_ptr, |
| 1633 | const scalar_t* __restrict__ k_ptr, |
| 1634 | const scalar_t* __restrict__ v_ptr, |
| 1635 | const param_t* __restrict__ A_log_ptr, |
| 1636 | const scalar_t* __restrict__ a_ptr, |
| 1637 | const scalar_t* __restrict__ dt_bias_ptr, |
| 1638 | const scalar_t* __restrict__ b_ptr, |
| 1639 | const int32_t* __restrict__ indices_ptr, |
| 1640 | state_t* __restrict__ state_ptr, |
| 1641 | scalar_t* __restrict__ o_ptr, |
| 1642 | float* __restrict__ qk_scale_buf, |
| 1643 | int64_t seq_len, |
| 1644 | int64_t batch_size, |
| 1645 | int64_t num_heads, |
| 1646 | int64_t head_dim, |
| 1647 | int64_t v_num_heads, |
| 1648 | int64_t v_head_dim, |
| 1649 | int64_t q_strideB, |
| 1650 | int64_t q_strideS, |
| 1651 | int64_t q_strideH, |
| 1652 | int64_t k_strideB, |
| 1653 | int64_t k_strideS, |
| 1654 | int64_t k_strideH, |
| 1655 | int64_t v_strideB, |
| 1656 | int64_t v_strideS, |
| 1657 | int64_t v_strideH, |
| 1658 | int64_t state_slot_stride, |
| 1659 | bool use_qk_l2norm_in_kernel, |
| 1660 | double softplus_threshold) { |
| 1661 | using bVec = at::vec::Vectorized<scalar_t>; |
| 1662 | using fVec = at::vec::Vectorized<float>; |
| 1663 | |
| 1664 | constexpr int64_t VecSize = bVec::size(); |
| 1665 | constexpr int64_t fVecSize = fVec::size(); |
| 1666 | int64_t group_size = v_num_heads / num_heads; |
| 1667 | double scale = 1 / std::sqrt(head_dim); |
| 1668 | fVec scale_vec = fVec(scale); |
| 1669 | if (use_qk_l2norm_in_kernel) { |
| 1670 | float eps = 1e-5; |
| 1671 | at::parallel_for(0, batch_size * seq_len * num_heads, 0, [&](int64_t begin, int64_t end) { |
| 1672 | int64_t bi{0}, si{0}, ni{0}; |
| 1673 | data_index_init(begin, bi, batch_size, si, seq_len, ni, num_heads); |
| 1674 | for (int64_t i = begin; i < end; ++i) { |
| 1675 | float sum_q = float(0); |
| 1676 | float sum_k = float(0); |
| 1677 | fVec sum_q_fvec = fVec(float(0)); |
| 1678 | fVec sum_k_fvec = fVec(float(0)); |
| 1679 | int64_t q_offset = bi * q_strideB + si * q_strideS + ni * q_strideH; |
| 1680 | int64_t k_offset = bi * k_strideB + si * k_strideS + ni * k_strideH; |
| 1681 | int64_t q_scale_offset = bi * seq_len * num_heads + si * num_heads + ni; |
| 1682 | int64_t k_scale_offset = q_scale_offset + batch_size * seq_len * num_heads; |
| 1683 | int64_t d; |
| 1684 | #pragma GCC unroll 4 |
| 1685 | for (d = 0; d <= head_dim - VecSize; d += VecSize) { |
| 1686 | auto [q_fvec0, q_fvec1] = load_float_vec2(q_ptr + q_offset + d); |
| 1687 | sum_q_fvec += q_fvec0 * q_fvec0; |
| 1688 | sum_q_fvec += q_fvec1 * q_fvec1; |
| 1689 | auto [k_fvec0, k_fvec1] = load_float_vec2(k_ptr + k_offset + d); |
| 1690 | sum_k_fvec += k_fvec0 * k_fvec0; |
| 1691 | sum_k_fvec += k_fvec1 * k_fvec1; |
| 1692 | } |
| 1693 | #pragma GCC unroll 4 |
| 1694 | for (; d < head_dim; ++d) { |
| 1695 | float q_val = static_cast<float>(q_ptr[q_offset + d]); |
| 1696 | sum_q += q_val * q_val; |
| 1697 | float k_val = static_cast<float>(k_ptr[k_offset + d]); |
| 1698 | sum_k += k_val * k_val; |
| 1699 | } |
| 1700 | |
| 1701 | sum_q += vec_reduce_sum(sum_q_fvec); |
| 1702 | sum_k += vec_reduce_sum(sum_k_fvec); |
| 1703 | qk_scale_buf[q_scale_offset] = float(1) / std::sqrt(sum_q + eps); |
| 1704 | qk_scale_buf[k_scale_offset] = float(1) / std::sqrt(sum_k + eps); |
| 1705 | |
| 1706 | data_index_step(bi, batch_size, si, seq_len, ni, num_heads); |
| 1707 | } |
| 1708 | }); |
| 1709 | } |
| 1710 | at::parallel_for(0, batch_size * seq_len * v_num_heads, 0, [&](int64_t begin, int64_t end) { |
| 1711 | int64_t bi{0}, si{0}, ni{0}; |
| 1712 | data_index_init(begin, bi, batch_size, si, seq_len, ni, v_num_heads); |
| 1713 | for (int64_t i = begin; i < end; ++i) { |
| 1714 | int64_t cache_index = indices_ptr[bi]; |
| 1715 | int64_t state_offset = cache_index * state_slot_stride + ni * head_dim * v_head_dim; |
| 1716 | float g_val = -std::exp(float(A_log_ptr[ni])) * |
| 1717 | softplus(float(a_ptr[bi * v_num_heads + ni]) + float(dt_bias_ptr[ni]), softplus_threshold); |
| 1718 | float g_val_exp = std::exp(g_val); |
| 1719 | fVec g_val_exp_vec = fVec(g_val_exp); |
| 1720 | int64_t q_offset = si * q_strideS + bi * q_strideB + (ni / group_size) * q_strideH; |
| 1721 | int64_t k_offset = si * k_strideS + bi * k_strideB + (ni / group_size) * k_strideH; |
| 1722 | int64_t q_scale_offset = bi * seq_len * num_heads + si * num_heads + (ni / group_size); |
| 1723 | int64_t k_scale_offset = q_scale_offset + batch_size * seq_len * num_heads; |
| 1724 | float q_scale = use_qk_l2norm_in_kernel ? qk_scale_buf[q_scale_offset] : 1.0f; |
| 1725 | float k_scale = use_qk_l2norm_in_kernel ? qk_scale_buf[k_scale_offset] : 1.0f; |
| 1726 | int64_t v_offset = si * v_strideS + bi * v_strideB + ni * v_strideH; |
| 1727 | int64_t o_offset = ((bi * seq_len + si) * v_num_heads + ni) * v_head_dim; |
| 1728 | // See: https://github.com/sgl-project/sglang/pull/26634 |
| 1729 | float beta_val = 1 / (1 + std::exp(-b_ptr[bi * v_num_heads + ni])); |
| 1730 | fVec beta_vec = fVec(beta_val); |
| 1731 | int64_t dvi = 0; |
| 1732 | for (; dvi <= v_head_dim - VecSize; dvi += VecSize) { |
| 1733 | fVec kv_mem_vec0 = fVec(float(0)); |
| 1734 | fVec kv_mem_vec1 = fVec(float(0)); |
| 1735 | for (int di = 0; di < head_dim; ++di) { |
| 1736 | fVec k_val_vec = fVec(k_ptr[k_offset + di] * k_scale); |
| 1737 | auto [state_vec0, state_vec1] = load_float_vec2(state_ptr + state_offset + di * v_head_dim + dvi); |
| 1738 | kv_mem_vec0 = kv_mem_vec0 + state_vec0 * g_val_exp_vec * k_val_vec; |
| 1739 | kv_mem_vec1 = kv_mem_vec1 + state_vec1 * g_val_exp_vec * k_val_vec; |
| 1740 | } |
| 1741 | auto [v_vec0, v_vec1] = load_float_vec2(v_ptr + v_offset + dvi); |
| 1742 | fVec dt_vec0 = (v_vec0 - kv_mem_vec0) * beta_vec; |
| 1743 | fVec dt_vec1 = (v_vec1 - kv_mem_vec1) * beta_vec; |
| 1744 | fVec o_vec0 = fVec(float(0)); |
| 1745 | fVec o_vec1 = fVec(float(0)); |
| 1746 | for (int di = 0; di < head_dim; ++di) { |
| 1747 | fVec q_vec = fVec(q_ptr[q_offset + di] * q_scale); |
| 1748 | fVec k_vec = fVec(k_ptr[k_offset + di] * k_scale); |
| 1749 | auto [state_vec0, state_vec1] = load_float_vec2(state_ptr + state_offset + di * v_head_dim + dvi); |
| 1750 | state_vec0 = state_vec0 * g_val_exp_vec + k_vec * dt_vec0; |
| 1751 | state_vec1 = state_vec1 * g_val_exp_vec + k_vec * dt_vec1; |
| 1752 | o_vec0 = o_vec0 + state_vec0 * q_vec * scale_vec; |
| 1753 | o_vec1 = o_vec1 + state_vec1 * q_vec * scale_vec; |
| 1754 | store_state_vec2<state_t>( |
| 1755 | state_ptr + state_offset + di * v_head_dim + dvi, |
| 1756 | state_vec0, |
| 1757 | state_vec1); |
| 1758 | } |
| 1759 | bVec o_vec = at::vec::convert_from_float<scalar_t>(o_vec0, o_vec1); |
| 1760 | o_vec.store(o_ptr + o_offset + dvi); |
| 1761 | } |
| 1762 | for (; dvi < v_head_dim; ++dvi) { |
| 1763 | float kv_mem_val = 0; |
| 1764 | for (int di = 0; di < head_dim; ++di) { |
| 1765 | float k_val = k_ptr[k_offset + di] * k_scale; |
| 1766 | float state_val = |
| 1767 | static_cast<float>(state_ptr[state_offset + di * v_head_dim + dvi]) * |
| 1768 | g_val_exp; |
| 1769 | kv_mem_val += state_val * k_val; |
| 1770 | } |
| 1771 | float v_val = v_ptr[v_offset + dvi]; |
| 1772 | float dt_val = (v_val - kv_mem_val) * beta_val; |
| 1773 | float o_val = 0; |
| 1774 | for (int di = 0; di < head_dim; ++di) { |
| 1775 | float q_val = q_ptr[q_offset + di] * q_scale; |
| 1776 | float k_val = k_ptr[k_offset + di] * k_scale; |
| 1777 | float state_val = |
| 1778 | static_cast<float>(state_ptr[state_offset + di * v_head_dim + dvi]) * |
| 1779 | g_val_exp; |
| 1780 | state_val += k_val * dt_val; |
| 1781 | state_ptr[state_offset + di * v_head_dim + dvi] = |
| 1782 | static_cast<state_t>(state_val); |
| 1783 | o_val += state_val * q_val * scale; |
| 1784 | } |
| 1785 | o_ptr[o_offset + dvi] = o_val; |
| 1786 | } |
| 1787 | data_index_step(bi, batch_size, si, seq_len, ni, v_num_heads); |
| 1788 | } |
| 1789 | }); |
| 1790 | } |
| 1791 | |
| 1792 | // Speculative-decode variant: processes a varlen batch where each sequence has |
| 1793 | // ``q_len`` draft tokens, runs the recurrence sequentially over those tokens |
| 1794 | // (inside the kernel, so one dispatch handles the whole draft block), reads the |
| 1795 | // initial state from cache slot ``num_accepted-1`` and stores the state *after* |
| 1796 | // token ``t`` into cache slot ``t`` (multi-slot rollback, matching the GPU |
| 1797 | // kernel). Parallelized over (sequence, v_head); the per-sequence token loop is |
| 1798 | // sequential as required by the recurrence. |
| 1799 | template <typename scalar_t, typename param_t, typename state_t> |
| 1800 | void fused_sigmoid_gating_delta_rule_update_spec_kernel_impl( |
| 1801 | const scalar_t* __restrict__ q_ptr, // [T, HK, EK] |
| 1802 | const scalar_t* __restrict__ k_ptr, // [T, HK, EK] |
| 1803 | const scalar_t* __restrict__ v_ptr, // [T, HV, EV] |
| 1804 | const param_t* __restrict__ A_log_ptr, |
| 1805 | const scalar_t* __restrict__ a_ptr, // [T, HV] |
| 1806 | const scalar_t* __restrict__ dt_bias_ptr, |
| 1807 | const scalar_t* __restrict__ b_ptr, // [T, HV] |
| 1808 | const int32_t* __restrict__ spec_indices_ptr, // [N, S] |
| 1809 | const int32_t* __restrict__ num_accepted_ptr, // [N] |
| 1810 | const int32_t* __restrict__ cu_seqlens_ptr, // [N + 1] |
| 1811 | state_t* __restrict__ state_ptr, |
| 1812 | scalar_t* __restrict__ o_ptr, // [T, HV, EV] |
| 1813 | float* __restrict__ qk_scale_buf, // [2, T, HK] |
| 1814 | int64_t total_tokens, |
| 1815 | int64_t batch_size, |
| 1816 | int64_t spec_stride, |
| 1817 | int64_t num_heads, |
| 1818 | int64_t head_dim, |
| 1819 | int64_t v_num_heads, |
| 1820 | int64_t v_head_dim, |
| 1821 | int64_t q_strideT, |
| 1822 | int64_t q_strideH, |
| 1823 | int64_t k_strideT, |
| 1824 | int64_t k_strideH, |
| 1825 | int64_t v_strideT, |
| 1826 | int64_t v_strideH, |
| 1827 | int64_t state_slot_stride, |
| 1828 | bool use_qk_l2norm_in_kernel, |
| 1829 | double softplus_threshold) { |
| 1830 | using bVec = at::vec::Vectorized<scalar_t>; |
| 1831 | using fVec = at::vec::Vectorized<float>; |
| 1832 | constexpr int64_t VecSize = bVec::size(); |
| 1833 | int64_t group_size = v_num_heads / num_heads; |
| 1834 | double scale = 1 / std::sqrt((double)head_dim); |
| 1835 | fVec scale_vec = fVec((float)scale); |
| 1836 | |
| 1837 | if (use_qk_l2norm_in_kernel) { |
| 1838 | float eps = 1e-5f; |
| 1839 | at::parallel_for(0, total_tokens * num_heads, 0, [&](int64_t begin, int64_t end) { |
| 1840 | for (int64_t i = begin; i < end; ++i) { |
| 1841 | int64_t ti = i / num_heads; |
| 1842 | int64_t ni = i % num_heads; |
| 1843 | const scalar_t* qp = q_ptr + ti * q_strideT + ni * q_strideH; |
| 1844 | const scalar_t* kp = k_ptr + ti * k_strideT + ni * k_strideH; |
| 1845 | float sq = 0.f, sk = 0.f; |
| 1846 | for (int64_t d = 0; d < head_dim; ++d) { |
| 1847 | float qv = (float)qp[d]; |
| 1848 | sq += qv * qv; |
| 1849 | float kv = (float)kp[d]; |
| 1850 | sk += kv * kv; |
| 1851 | } |
| 1852 | qk_scale_buf[ti * num_heads + ni] = 1.f / std::sqrt(sq + eps); |
| 1853 | qk_scale_buf[total_tokens * num_heads + ti * num_heads + ni] = 1.f / std::sqrt(sk + eps); |
| 1854 | } |
| 1855 | }); |
| 1856 | } |
| 1857 | |
| 1858 | at::parallel_for(0, batch_size * v_num_heads, 0, [&](int64_t begin, int64_t end) { |
| 1859 | for (int64_t idx = begin; idx < end; ++idx) { |
| 1860 | int64_t bi = idx / v_num_heads; |
| 1861 | int64_t ni = idx % v_num_heads; |
| 1862 | int64_t kh = ni / group_size; |
| 1863 | int64_t q_start = cu_seqlens_ptr[bi]; |
| 1864 | int64_t q_len = cu_seqlens_ptr[bi + 1] - q_start; |
| 1865 | if (q_len <= 0) { |
| 1866 | continue; |
| 1867 | } |
| 1868 | int64_t acc = (int64_t)num_accepted_ptr[bi]; |
| 1869 | // Clamp acc-1 to >=0: when num_accepted is 0 the unclamped index reads |
| 1870 | // out of bounds and yields an arbitrary prev_slot used to index the SSM |
| 1871 | // state. Mirrors the GPU guard tl.maximum(num_accepted - 1, 0). |
| 1872 | int64_t prev_slot = |
| 1873 | (int64_t)spec_indices_ptr[bi * spec_stride + (acc > 0 ? acc - 1 : 0)]; |
| 1874 | for (int64_t t = 0; t < q_len; ++t) { |
| 1875 | int64_t cur_slot = (int64_t)spec_indices_ptr[bi * spec_stride + t]; |
| 1876 | int64_t token = q_start + t; |
| 1877 | const state_t* src = |
| 1878 | state_ptr + prev_slot * state_slot_stride + ni * head_dim * v_head_dim; |
| 1879 | state_t* dst = |
| 1880 | state_ptr + cur_slot * state_slot_stride + ni * head_dim * v_head_dim; |
| 1881 | float g_val = -std::exp((float)A_log_ptr[ni]) * |
| 1882 | softplus((float)a_ptr[token * v_num_heads + ni] + (float)dt_bias_ptr[ni], softplus_threshold); |
| 1883 | float g_val_exp = std::exp(g_val); |
| 1884 | fVec g_val_exp_vec = fVec(g_val_exp); |
| 1885 | float beta_val = 1.f / (1.f + std::exp(-(float)b_ptr[token * v_num_heads + ni])); |
| 1886 | fVec beta_vec = fVec(beta_val); |
| 1887 | int64_t q_offset = token * q_strideT + kh * q_strideH; |
| 1888 | int64_t k_offset = token * k_strideT + kh * k_strideH; |
| 1889 | float q_scale = use_qk_l2norm_in_kernel ? qk_scale_buf[token * num_heads + kh] : 1.f; |
| 1890 | float k_scale = |
| 1891 | use_qk_l2norm_in_kernel ? qk_scale_buf[total_tokens * num_heads + token * num_heads + kh] : 1.f; |
| 1892 | int64_t v_offset = token * v_strideT + ni * v_strideH; |
| 1893 | int64_t o_offset = (token * v_num_heads + ni) * v_head_dim; |
| 1894 | int64_t dvi = 0; |
| 1895 | for (; dvi <= v_head_dim - VecSize; dvi += VecSize) { |
| 1896 | fVec kv_mem_vec0 = fVec(0.f); |
| 1897 | fVec kv_mem_vec1 = fVec(0.f); |
| 1898 | for (int di = 0; di < head_dim; ++di) { |
| 1899 | fVec k_val_vec = fVec((float)k_ptr[k_offset + di] * k_scale); |
| 1900 | auto [sv0, sv1] = |
| 1901 | load_float_vec2(src + di * v_head_dim + dvi); |
| 1902 | kv_mem_vec0 = kv_mem_vec0 + sv0 * g_val_exp_vec * k_val_vec; |
| 1903 | kv_mem_vec1 = kv_mem_vec1 + sv1 * g_val_exp_vec * k_val_vec; |
| 1904 | } |
| 1905 | bVec v_bvec = bVec::loadu(v_ptr + v_offset + dvi); |
| 1906 | fVec v_vec0, v_vec1; |
| 1907 | std::tie(v_vec0, v_vec1) = at::vec::convert_to_float(v_bvec); |
| 1908 | fVec dt_vec0 = (v_vec0 - kv_mem_vec0) * beta_vec; |
| 1909 | fVec dt_vec1 = (v_vec1 - kv_mem_vec1) * beta_vec; |
| 1910 | fVec o_vec0 = fVec(0.f); |
| 1911 | fVec o_vec1 = fVec(0.f); |
| 1912 | for (int di = 0; di < head_dim; ++di) { |
| 1913 | fVec q_vec = fVec((float)q_ptr[q_offset + di] * q_scale); |
| 1914 | fVec k_vec = fVec((float)k_ptr[k_offset + di] * k_scale); |
| 1915 | auto [sv0, sv1] = |
| 1916 | load_float_vec2(src + di * v_head_dim + dvi); |
| 1917 | sv0 = sv0 * g_val_exp_vec + k_vec * dt_vec0; |
| 1918 | sv1 = sv1 * g_val_exp_vec + k_vec * dt_vec1; |
| 1919 | o_vec0 = o_vec0 + sv0 * q_vec * scale_vec; |
| 1920 | o_vec1 = o_vec1 + sv1 * q_vec * scale_vec; |
| 1921 | store_state_vec2<state_t>( |
| 1922 | dst + di * v_head_dim + dvi, sv0, sv1); |
| 1923 | } |
| 1924 | bVec o_vec = at::vec::convert_from_float<scalar_t>(o_vec0, o_vec1); |
| 1925 | o_vec.store(o_ptr + o_offset + dvi); |
| 1926 | } |
| 1927 | for (; dvi < v_head_dim; ++dvi) { |
| 1928 | float kv_mem_val = 0.f; |
| 1929 | for (int di = 0; di < head_dim; ++di) { |
| 1930 | float k_val = (float)k_ptr[k_offset + di] * k_scale; |
| 1931 | kv_mem_val += static_cast<float>(src[di * v_head_dim + dvi]) * |
| 1932 | g_val_exp * k_val; |
| 1933 | } |
| 1934 | float v_val = (float)v_ptr[v_offset + dvi]; |
| 1935 | float dt_val = (v_val - kv_mem_val) * beta_val; |
| 1936 | float o_val = 0.f; |
| 1937 | for (int di = 0; di < head_dim; ++di) { |
| 1938 | float q_val = (float)q_ptr[q_offset + di] * q_scale; |
| 1939 | float k_val = (float)k_ptr[k_offset + di] * k_scale; |
| 1940 | float ns = static_cast<float>(src[di * v_head_dim + dvi]) * |
| 1941 | g_val_exp + |
| 1942 | k_val * dt_val; |
| 1943 | dst[di * v_head_dim + dvi] = static_cast<state_t>(ns); |
| 1944 | o_val += ns * q_val * scale; |
| 1945 | } |
| 1946 | o_ptr[o_offset + dvi] = (scalar_t)o_val; |
| 1947 | } |
| 1948 | prev_slot = cur_slot; |
| 1949 | } |
| 1950 | } |
| 1951 | }); |
| 1952 | } |
| 1953 | |
| 1954 | template <typename scalar_t> |
| 1955 | void fused_gdn_gating_kernel_impl( |
| 1956 | float* __restrict__ A_log, |
| 1957 | const scalar_t* __restrict__ a, |
| 1958 | const scalar_t* __restrict__ b, |
| 1959 | const scalar_t* __restrict__ dt_bias, |
| 1960 | float* __restrict__ out, |
| 1961 | scalar_t* __restrict__ beta, |
| 1962 | int64_t batch, |
| 1963 | int64_t num_heads) { |
| 1964 | using bVec = at::vec::Vectorized<scalar_t>; |
| 1965 | using fVec = at::vec::Vectorized<float>; |
| 1966 | constexpr int vec_size = bVec::size(); |
| 1967 | constexpr int fvec_size = fVec::size(); |
| 1968 | const fVec neg_one(-1.0f); |
| 1969 | at::parallel_for(0, batch, 0, [&](int64_t begin, int64_t end) { |
| 1970 | for (int64_t i = begin; i < end; ++i) { |
| 1971 | int64_t j = 0; |
| 1972 | for (; j < num_heads - (num_heads % vec_size); j += vec_size) { |
| 1973 | auto [A_log_vec0, A_log_vec1] = load_float_vec2(A_log + j); |
| 1974 | auto [dt_bias_vec0, dt_bias_vec1] = load_float_vec2(dt_bias + j); |
| 1975 | auto [a0, a1] = load_float_vec2(a + i * num_heads + j); |
| 1976 | auto [b0, b1] = load_float_vec2(b + i * num_heads + j); |
| 1977 | |
| 1978 | fVec g0 = neg_one * A_log_vec0.exp_u20() * softplus(a0 + dt_bias_vec0); |
| 1979 | fVec g1 = neg_one * A_log_vec1.exp_u20() * softplus(a1 + dt_bias_vec1); |
| 1980 | fVec beta0 = fast_sigmoid(b0); |
| 1981 | fVec beta1 = fast_sigmoid(b1); |
| 1982 | |
| 1983 | g0.store(out + i * num_heads + j); |
| 1984 | g1.store(out + i * num_heads + j + fvec_size); |
| 1985 | bVec beta_vec = at::vec::convert_from_float<scalar_t>(beta0, beta1); |
| 1986 | beta_vec.store(beta + i * num_heads + j); |
| 1987 | } |
| 1988 | for (; j < num_heads; ++j) { |
| 1989 | out[i * num_heads + j] = -std::exp(A_log[j]) * softplus(float(a[i * num_heads + j]) + float(dt_bias[j])); |
| 1990 | beta[i * num_heads + j] = 1 / (1 + std::exp(-b[i * num_heads + j])); |
| 1991 | } |
| 1992 | } |
| 1993 | }); |
| 1994 | } |
| 1995 | |
| 1996 | template <typename scalar_t> |
| 1997 | void fused_gdn_gating_kernel_impl( |
| 1998 | scalar_t* __restrict__ A_log, |
| 1999 | const scalar_t* __restrict__ a, |
| 2000 | const scalar_t* __restrict__ b, |
| 2001 | const scalar_t* __restrict__ dt_bias, |
| 2002 | float* __restrict__ out, |
| 2003 | scalar_t* __restrict__ beta, |
| 2004 | int64_t batch, |
| 2005 | int64_t num_heads) { |
| 2006 | using bVec = at::vec::Vectorized<scalar_t>; |
| 2007 | using fVec = at::vec::Vectorized<float>; |
| 2008 | constexpr int vec_size = bVec::size(); |
| 2009 | constexpr int fvec_size = fVec::size(); |
| 2010 | const fVec neg_one(-1.0f); |
| 2011 | at::parallel_for(0, batch, 0, [&](int64_t begin, int64_t end) { |
| 2012 | for (int64_t i = begin; i < end; ++i) { |
| 2013 | int64_t j = 0; |
| 2014 | for (; j < num_heads - (num_heads % vec_size); j += vec_size) { |
| 2015 | auto [A_log_vec0, A_log_vec1] = load_float_vec2(A_log + j); |
| 2016 | auto [dt_bias_vec0, dt_bias_vec1] = load_float_vec2(dt_bias + j); |
| 2017 | auto [a0, a1] = load_float_vec2(a + i * num_heads + j); |
| 2018 | auto [b0, b1] = load_float_vec2(b + i * num_heads + j); |
| 2019 | |
| 2020 | fVec g0 = neg_one * A_log_vec0.exp_u20() * softplus(a0 + dt_bias_vec0); |
| 2021 | fVec g1 = neg_one * A_log_vec1.exp_u20() * softplus(a1 + dt_bias_vec1); |
| 2022 | fVec beta0 = fast_sigmoid(b0); |
| 2023 | fVec beta1 = fast_sigmoid(b1); |
| 2024 | |
| 2025 | g0.store(out + i * num_heads + j); |
| 2026 | g1.store(out + i * num_heads + j + fvec_size); |
| 2027 | bVec beta_vec = at::vec::convert_from_float<scalar_t>(beta0, beta1); |
| 2028 | beta_vec.store(beta + i * num_heads + j); |
| 2029 | } |
| 2030 | for (; j < num_heads; ++j) { |
| 2031 | out[i * num_heads + j] = -std::exp(float(A_log[j])) * softplus(float(a[i * num_heads + j]) + float(dt_bias[j])); |
| 2032 | beta[i * num_heads + j] = 1 / (1 + std::exp(-b[i * num_heads + j])); |
| 2033 | } |
| 2034 | } |
| 2035 | }); |
| 2036 | } |
| 2037 | |
| 2038 | } // anonymous namespace |
| 2039 | |
| 2040 | template <int CHUNK_SIZE> |
| 2041 | std::tuple<at::Tensor, at::Tensor> prepare_chunk_indices(const at::Tensor& cu_seqlens) { |
| 2042 | int64_t num_seqs = cu_seqlens.size(0) - 1; |
| 2043 | at::Tensor chunk_offsets = at::empty({num_seqs + 1}, cu_seqlens.options()); |
| 2044 | // get number of chunks and chunk offsets |
| 2045 | const int32_t* offsets_data = cu_seqlens.data_ptr<int32_t>(); |
| 2046 | int32_t num_chunks = 0; |
| 2047 | chunk_offsets[0] = 0; |
| 2048 | for (int64_t row = 0; row < num_seqs; ++row) { |
| 2049 | num_chunks += div_up(offsets_data[row + 1] - offsets_data[row], CHUNK_SIZE); |
| 2050 | chunk_offsets[row + 1] = num_chunks; |
| 2051 | } |
| 2052 | // get chunk indices |
| 2053 | at::Tensor chunk_indices = at::empty({num_chunks, 2}, cu_seqlens.options()); |
| 2054 | int32_t* indices_data = chunk_indices.data_ptr<int32_t>(); |
| 2055 | |
| 2056 | int64_t idx = 0; |
| 2057 | for (int32_t row = 0; row < num_seqs; ++row) { |
| 2058 | int32_t num_chunks = div_up(offsets_data[row + 1] - offsets_data[row], CHUNK_SIZE); |
| 2059 | |
| 2060 | for (int32_t col = 0; col < num_chunks; ++col) { |
| 2061 | indices_data[idx * 2 + 0] = row; |
| 2062 | indices_data[idx * 2 + 1] = col; |
| 2063 | idx++; |
| 2064 | } |
| 2065 | } |
| 2066 | return std::make_tuple(chunk_indices, chunk_offsets); |
| 2067 | } |
| 2068 | |
| 2069 | #define DISPATCH_HEAD_DIM_CASE(launch_macro, hd) \ |
| 2070 | case hd: { \ |
| 2071 | launch_macro(hd); \ |
| 2072 | break; \ |
| 2073 | } |
| 2074 | |
| 2075 | // [NB]: add new head_dim support here |
| 2076 | #define DISPATCH_HEAD_DIM(dim, launch_macro) \ |
| 2077 | switch (dim) { \ |
| 2078 | DISPATCH_HEAD_DIM_CASE(launch_macro, 64) \ |
| 2079 | DISPATCH_HEAD_DIM_CASE(launch_macro, 128) \ |
| 2080 | default: \ |
| 2081 | TORCH_CHECK(false, "Unexpected head dim size, ", dim); \ |
| 2082 | } |
| 2083 | |
| 2084 | #define LAUNCH_L2NORM_KERNEL(HD) \ |
| 2085 | l2norm_fwd_kernel_impl<scalar_t, HD>( \ |
| 2086 | query_norm.data_ptr<scalar_t>(), \ |
| 2087 | key_norm.data_ptr<scalar_t>(), \ |
| 2088 | query.data_ptr<scalar_t>(), \ |
| 2089 | key.data_ptr<scalar_t>(), \ |
| 2090 | eps, \ |
| 2091 | T, \ |
| 2092 | H, \ |
| 2093 | query.stride(1), \ |
| 2094 | query.stride(2), \ |
| 2095 | key.stride(1), \ |
| 2096 | key.stride(2)); |
| 2097 | |
| 2098 | std::tuple<at::Tensor, at::Tensor> l2norm_fwd(const at::Tensor& query, const at::Tensor& key, double eps) { |
| 2099 | int64_t B = query.size(0); |
| 2100 | int64_t T = query.size(1); |
| 2101 | int64_t H = query.size(2); |
| 2102 | int64_t D = query.size(3); |
| 2103 | |
| 2104 | at::Tensor query_norm = at::empty_like(query); |
| 2105 | at::Tensor key_norm = at::empty_like(key); |
| 2106 | |
| 2107 | AT_DISPATCH_REDUCED_FLOATING_TYPES( |
| 2108 | query.scalar_type(), "l2norm_fwd", [&] { DISPATCH_HEAD_DIM(D, LAUNCH_L2NORM_KERNEL); }); |
| 2109 | |
| 2110 | return std::make_tuple(query_norm, key_norm); |
| 2111 | } |
| 2112 | |
| 2113 | // [NB]: instantiate decay_mask to avoid heavy recomputation in the kernel with exp |
| 2114 | template <int CHUNK_SIZE> |
| 2115 | at::Tensor chunk_local_cumsum(const at::Tensor& g, const at::Tensor& cu_seqlens, const at::Tensor& chunk_indices) { |
| 2116 | int64_t B = g.size(0); |
| 2117 | // int64_t T = g.size(1); |
| 2118 | int64_t Hv = g.size(2); |
| 2119 | int64_t NT = chunk_indices.size(0); |
| 2120 | |
| 2121 | at::Tensor g_ = at::empty({B, NT, Hv, CHUNK_SIZE}, g.options()); |
| 2122 | AT_DISPATCH_FLOATING_TYPES(g.scalar_type(), "chunk_local_cumsum", [&] { |
| 2123 | chunk_local_cumsum_kernel_impl<scalar_t, CHUNK_SIZE>( |
| 2124 | g_.data_ptr<scalar_t>(), |
| 2125 | g.data_ptr<scalar_t>(), |
| 2126 | cu_seqlens.data_ptr<int32_t>(), |
| 2127 | chunk_indices.data_ptr<int32_t>(), |
| 2128 | Hv, |
| 2129 | NT); |
| 2130 | }); |
| 2131 | return g_; |
| 2132 | } |
| 2133 | |
| 2134 | #define LAUNCH_CHUNK_GATED_DELTA_RULE_FWD_INTRA_KERNEL(HD) \ |
| 2135 | chunk_gated_delta_rule_fwd_intra_kernel_impl<scalar_t, HD, CHUNK_SIZE>( \ |
| 2136 | w.data_ptr<scalar_t>(), \ |
| 2137 | u.data_ptr<scalar_t>(), \ |
| 2138 | decay_mask.data_ptr<float>(), \ |
| 2139 | k.data_ptr<scalar_t>(), \ |
| 2140 | v.data_ptr<scalar_t>(), \ |
| 2141 | g.data_ptr<float>(), \ |
| 2142 | beta.data_ptr<scalar_t>(), \ |
| 2143 | cu_seqlens.data_ptr<int32_t>(), \ |
| 2144 | chunk_indices.data_ptr<int32_t>(), \ |
| 2145 | H, \ |
| 2146 | Hv, \ |
| 2147 | NT, \ |
| 2148 | k.stride(1), \ |
| 2149 | k.stride(2), \ |
| 2150 | v.stride(1), \ |
| 2151 | v.stride(2)); |
| 2152 | |
| 2153 | template <int CHUNK_SIZE> |
| 2154 | std::tuple<at::Tensor, at::Tensor, at::Tensor> chunk_gated_delta_rule_fwd_intra( |
| 2155 | const at::Tensor& k, |
| 2156 | const at::Tensor& v, |
| 2157 | const at::Tensor& g, |
| 2158 | const at::Tensor& beta, |
| 2159 | const at::Tensor& cu_seqlens, |
| 2160 | const at::Tensor& chunk_indices) { |
| 2161 | int64_t B = k.size(0); |
| 2162 | int64_t T = k.size(1); |
| 2163 | int64_t H = k.size(2); |
| 2164 | int64_t D = k.size(3); |
| 2165 | int64_t Hv = v.size(2); |
| 2166 | int64_t Dv = v.size(3); |
| 2167 | int64_t NT = chunk_indices.size(0); |
| 2168 | |
| 2169 | at::Tensor w = at::empty({B, T, Hv, D}, k.options()); // BFloat16 |
| 2170 | at::Tensor u = at::empty({B, T, Hv, Dv}, k.options()); // BFloat16 |
| 2171 | at::Tensor decay_mask = at::empty({B, NT, Hv, CHUNK_SIZE, CHUNK_SIZE}, g.options()); // Float |
| 2172 | AT_DISPATCH_REDUCED_FLOATING_TYPES(k.scalar_type(), "chunk_gated_delta_rule_fwd_intra", [&] { |
| 2173 | DISPATCH_HEAD_DIM(D, LAUNCH_CHUNK_GATED_DELTA_RULE_FWD_INTRA_KERNEL); |
| 2174 | }); |
| 2175 | |
| 2176 | return std::make_tuple(w, u, decay_mask); |
| 2177 | } |
| 2178 | |
| 2179 | #define LAUNCH_CHUNK_GATED_DELTA_RULE_FWD_INTER_KERNEL(HD) \ |
| 2180 | chunk_gated_delta_rule_fwd_inter_kernel_impl<scalar_t, state_t, HD, CHUNK_SIZE>( \ |
| 2181 | o.data_ptr<scalar_t>(), \ |
| 2182 | initial_state.data_ptr<state_t>(), \ |
| 2183 | initial_state_indices.data_ptr<int32_t>(), \ |
| 2184 | q.data_ptr<scalar_t>(), \ |
| 2185 | k.data_ptr<scalar_t>(), \ |
| 2186 | w.data_ptr<scalar_t>(), \ |
| 2187 | u.data_ptr<scalar_t>(), \ |
| 2188 | g.data_ptr<float>(), \ |
| 2189 | decay_mask.data_ptr<float>(), \ |
| 2190 | cu_seqlens.data_ptr<int32_t>(), \ |
| 2191 | chunk_offsets.data_ptr<int32_t>(), \ |
| 2192 | H, \ |
| 2193 | Hv, \ |
| 2194 | num_seqs, \ |
| 2195 | q.stride(1), \ |
| 2196 | q.stride(2), \ |
| 2197 | k.stride(1), \ |
| 2198 | k.stride(2), \ |
| 2199 | initial_state.stride(0)); |
| 2200 | |
| 2201 | template <typename state_t, int CHUNK_SIZE> |
| 2202 | std::tuple<at::Tensor, at::Tensor> chunk_gated_delta_rule_fwd_inter( |
| 2203 | const at::Tensor& q, |
| 2204 | const at::Tensor& k, |
| 2205 | const at::Tensor& w, |
| 2206 | const at::Tensor& u, |
| 2207 | const at::Tensor& g, |
| 2208 | const at::Tensor& decay_mask, |
| 2209 | const at::Tensor& initial_state, |
| 2210 | bool output_final_state, |
| 2211 | const at::Tensor& cu_seqlens, |
| 2212 | const at::Tensor& chunk_offsets, |
| 2213 | const at::Tensor& initial_state_indices) { |
| 2214 | const int64_t B = q.size(0); |
| 2215 | const int64_t T = q.size(1); |
| 2216 | const int64_t H = q.size(2); |
| 2217 | const int64_t D = q.size(3); |
| 2218 | const int64_t Hv = w.size(2); |
| 2219 | const int64_t Dv = u.size(3); |
| 2220 | const int64_t num_seqs = initial_state_indices.size(0); |
| 2221 | |
| 2222 | at::Tensor o = at::empty({B, T, Hv, Dv}, q.options()); |
| 2223 | AT_DISPATCH_REDUCED_FLOATING_TYPES(q.scalar_type(), "chunk_gated_delta_rule_fwd_inter", [&] { |
| 2224 | DISPATCH_HEAD_DIM(D, LAUNCH_CHUNK_GATED_DELTA_RULE_FWD_INTER_KERNEL); |
| 2225 | }); |
| 2226 | |
| 2227 | return std::make_tuple(o, initial_state); |
| 2228 | } |
| 2229 | |
| 2230 | // [NB]: Support only varlen inputs |
| 2231 | // B: packed batch dim of q/k/v (== 1) |
| 2232 | // num_seqs: number of variable-length sequences |
| 2233 | // |
| 2234 | // query: [B, T, H, D] |
| 2235 | // key: [B, T, H, D] |
| 2236 | // value: [B, T, Hv, Dv] |
| 2237 | // g: [B, T, Hv] FP32 |
| 2238 | // beta: [B, T, Hv] |
| 2239 | // initial_state: [num_seqs, Hv, Dv, D] FP16, BF16, or FP32 persisted state |
| 2240 | // cu_seqlens: [num_seqs + 1] INT32 |
| 2241 | // |
| 2242 | std::tuple<at::Tensor, at::Tensor> chunk_gated_delta_rule_cpu( |
| 2243 | const at::Tensor& query, |
| 2244 | const at::Tensor& key, |
| 2245 | const at::Tensor& value, |
| 2246 | const at::Tensor& g, |
| 2247 | const at::Tensor& beta, |
| 2248 | const at::Tensor& initial_state, |
| 2249 | bool output_final_state, |
| 2250 | const at::Tensor& cu_seqlens, |
| 2251 | bool head_first, |
| 2252 | bool use_qk_l2norm_in_kernel, |
| 2253 | const at::Tensor& initial_state_indices, |
| 2254 | double eps = 1e-6) { |
| 2255 | TORCH_CHECK(!head_first, "chunk_gated_delta_rule_cpu: does not support head first"); |
| 2256 | |
| 2257 | int64_t B = query.size(0); |
| 2258 | int64_t T = query.size(1); |
| 2259 | int64_t H = query.size(2); |
| 2260 | int64_t D = query.size(3); |
| 2261 | int64_t Hv = value.size(2); |
| 2262 | int64_t Dv = value.size(3); |
| 2263 | int64_t num_seqs = initial_state_indices.size(0); |
| 2264 | |
| 2265 | TORCH_CHECK(B == 1, __func__, ": expect batch size to be 1"); |
| 2266 | TORCH_CHECK(Hv % H == 0, __func__, ": expect num_heads_kv multiple of num_heads."); |
| 2267 | TORCH_CHECK(D % 32 == 0, __func__, ": expect head_dim to be multiples of 32."); |
| 2268 | TORCH_CHECK(Dv % 32 == 0, __func__, ": expect head_dim_v to be multiples of 32."); |
| 2269 | TORCH_CHECK(D == Dv, __func__, ": expect head_dim to be equal to head_dim_v."); |
| 2270 | CHECK_INPUT_SHAPE_DTYPE<true>(query, {B, T, H, D}, at::kBFloat16); |
| 2271 | CHECK_INPUT_SHAPE_DTYPE<true>(key, {B, T, H, D}, at::kBFloat16); |
| 2272 | CHECK_INPUT_SHAPE_DTYPE<true>(value, {B, T, Hv, Dv}, at::kBFloat16); |
| 2273 | CHECK_INPUT_SHAPE_DTYPE<false>(g, {B, T, Hv}, at::kFloat); |
| 2274 | CHECK_INPUT_SHAPE_DTYPE<false>(beta, {B, T, Hv}, at::kBFloat16); |
| 2275 | CHECK_INPUT_SHAPE_DTYPE<false>(cu_seqlens, {num_seqs + 1}, at::kInt); |
| 2276 | TORCH_CHECK(initial_state.sizes() == at::IntArrayRef({initial_state.size(0), Hv, Dv, D}), |
| 2277 | "chunk_gated_delta_rule_cpu: initial_state shape mismatch, got ", initial_state.sizes()); |
| 2278 | TORCH_CHECK( |
| 2279 | initial_state.scalar_type() == at::kFloat || |
| 2280 | initial_state.scalar_type() == at::kBFloat16 || |
| 2281 | initial_state.scalar_type() == at::kHalf, |
| 2282 | "chunk_gated_delta_rule_cpu: initial_state dtype must be float16, " |
| 2283 | "bfloat16, or float32"); |
| 2284 | CHECK_CPU(initial_state); |
| 2285 | // initial_state may be a pooled/paged buffer with padding between slots |
| 2286 | // (e.g. mamba cache-align mode), so only the per-slot (Hv, Dv, D) layout |
| 2287 | // needs to be densely packed; dim 0's stride is read at runtime instead of |
| 2288 | // assumed, mirroring conv.cpp's conv_state_slot_stride handling. |
| 2289 | TORCH_CHECK(initial_state.stride(-1) == 1 && initial_state.stride(-2) == D && |
| 2290 | initial_state.stride(-3) == Dv * D, |
| 2291 | "chunk_gated_delta_rule_cpu: expect initial_state to be contiguous per pool slot."); |
| 2292 | CHECK_INPUT_SHAPE_DTYPE<false>(initial_state_indices, {num_seqs}, at::kInt); |
| 2293 | |
| 2294 | constexpr int CHUNK_SIZE = 64; |
| 2295 | |
| 2296 | // prepare chunk indices |
| 2297 | auto [chunk_indices, chunk_offsets] = prepare_chunk_indices<CHUNK_SIZE>(cu_seqlens); |
| 2298 | |
| 2299 | float scale = 1.0 / std::sqrt(D); |
| 2300 | auto [query_, key_] = use_qk_l2norm_in_kernel ? l2norm_fwd(query, key, eps) : std::make_tuple(query.mul(scale), key); |
| 2301 | |
| 2302 | auto g_ = chunk_local_cumsum<CHUNK_SIZE>(g, cu_seqlens, chunk_indices); |
| 2303 | |
| 2304 | // fused kkt + solve_tril + recompute_w_u |
| 2305 | auto [w, u, decay_mask] = |
| 2306 | chunk_gated_delta_rule_fwd_intra<CHUNK_SIZE>(key_, value, g_, beta, cu_seqlens, chunk_indices); |
| 2307 | |
| 2308 | // fused `chunk_gated_delta_rule_fwd_h` + `chunk_fwd_o` |
| 2309 | auto launch_inter = [local_query = query_, |
| 2310 | local_key = key_, |
| 2311 | local_w = w, |
| 2312 | local_u = u, |
| 2313 | local_decay_mask = decay_mask, |
| 2314 | local_chunk_offsets = chunk_offsets, |
| 2315 | g_, |
| 2316 | initial_state, |
| 2317 | output_final_state, |
| 2318 | cu_seqlens, |
| 2319 | initial_state_indices](auto state_tag) { |
| 2320 | using state_t = decltype(state_tag); |
| 2321 | return chunk_gated_delta_rule_fwd_inter<state_t, CHUNK_SIZE>( |
| 2322 | local_query, |
| 2323 | local_key, |
| 2324 | local_w, |
| 2325 | local_u, |
| 2326 | g_, |
| 2327 | local_decay_mask, |
| 2328 | initial_state, |
| 2329 | output_final_state, |
| 2330 | cu_seqlens, |
| 2331 | local_chunk_offsets, |
| 2332 | initial_state_indices); |
| 2333 | }; |
| 2334 | |
| 2335 | |
| 2336 | |
| 2337 | |
| 2338 | auto output_and_state = |
| 2339 | initial_state.scalar_type() == at::kFloat |
| 2340 | ? launch_inter(float{}) |
| 2341 | : initial_state.scalar_type() == at::kBFloat16 |
| 2342 | ? launch_inter(at::BFloat16{}) |
| 2343 | : launch_inter(at::Half{}); |
| 2344 | auto [output, final_state] = output_and_state; |
| 2345 | |
| 2346 | return std::make_tuple(output, final_state); |
| 2347 | } |
| 2348 | |
| 2349 | // A_log: [v_num_heads] |
| 2350 | // dt_bias: [v_num_heads] |
| 2351 | // query: [seq_len, batch_size, num_heads, head_dim] |
| 2352 | // key: [seq_len, batch_size, num_heads, head_dim] |
| 2353 | // value: [seq_len, batch_size, v_num_heads, v_head_dim] |
| 2354 | // a: [batch_size, v_num_heads] |
| 2355 | // b: [batch_size, v_num_heads] |
| 2356 | // initial_state_source:[num_tokens, v_num_heads, head_dim, v_head_dim] |
| 2357 | // FP16, BF16, or FP32 persisted state; recurrent arithmetic remains FP32. |
| 2358 | // initial_state_indices: [batch_size] |
| 2359 | // cu_seqlens: [batch_size + 1] |
| 2360 | at::Tensor fused_sigmoid_gating_delta_rule_update_cpu( |
| 2361 | const at::Tensor& A_log, |
| 2362 | const at::Tensor& dt_bias, |
| 2363 | const at::Tensor& q, |
| 2364 | const at::Tensor& k, |
| 2365 | const at::Tensor& v, |
| 2366 | const at::Tensor& a, |
| 2367 | const at::Tensor& b, |
| 2368 | at::Tensor& initial_state_source, |
| 2369 | const at::Tensor& initial_state_indices, |
| 2370 | const at::Tensor& cu_seqlens, |
| 2371 | bool use_qk_l2norm_in_kernel, |
| 2372 | double softplus_beta = 1.0, |
| 2373 | double softplus_threshold = 20.0) { |
| 2374 | CHECK_DIM(4, q); |
| 2375 | CHECK_DIM(4, v); |
| 2376 | CHECK_LAST_DIM_CONTIGUOUS_INPUT(q); |
| 2377 | int64_t seq_len = q.size(0); |
| 2378 | int64_t batch_size = q.size(1); |
| 2379 | int64_t num_heads = q.size(2); |
| 2380 | int64_t head_dim = q.size(3); |
| 2381 | int64_t v_num_heads = v.size(2); |
| 2382 | int64_t v_head_dim = v.size(3); |
| 2383 | CHECK_INPUT_SHAPE_DTYPE<true>(k, {seq_len, batch_size, num_heads, head_dim}, q.scalar_type()); |
| 2384 | CHECK_INPUT_SHAPE_DTYPE<true>(v, {seq_len, batch_size, v_num_heads, v_head_dim}, q.scalar_type()); |
| 2385 | CHECK_INPUT_SHAPE_DTYPE<true>(a, {batch_size, v_num_heads}, q.scalar_type()); |
| 2386 | CHECK_INPUT_SHAPE_DTYPE<true>(dt_bias, {v_num_heads}, q.scalar_type()); |
| 2387 | CHECK_INPUT_SHAPE_DTYPE<true>(b, {batch_size, v_num_heads}, q.scalar_type()); |
| 2388 | CHECK_INPUT_SHAPE_DTYPE<true>(initial_state_indices, {batch_size}, at::kInt); |
| 2389 | CHECK_INPUT_SHAPE_DTYPE<true>(cu_seqlens, {batch_size + 1}, at::kInt); |
| 2390 | TORCH_CHECK( |
| 2391 | initial_state_source.scalar_type() == at::kFloat || |
| 2392 | initial_state_source.scalar_type() == at::kBFloat16 || |
| 2393 | initial_state_source.scalar_type() == at::kHalf, |
| 2394 | "fused_sigmoid_gating_delta_rule_update_cpu: initial_state dtype must " |
| 2395 | "be float16, bfloat16, or float32"); |
| 2396 | TORCH_CHECK( |
| 2397 | initial_state_source.sizes() == |
| 2398 | at::IntArrayRef( |
| 2399 | {initial_state_source.size(0), v_num_heads, head_dim, v_head_dim}), |
| 2400 | "Input tensor shape mismatch for initial_state_source"); |
| 2401 | CHECK_CPU(initial_state_source); |
| 2402 | TORCH_CHECK( |
| 2403 | initial_state_source.stride(-1) == 1 && |
| 2404 | initial_state_source.stride(-2) == v_head_dim && |
| 2405 | initial_state_source.stride(-3) == head_dim * v_head_dim, |
| 2406 | "fused_sigmoid_gating_delta_rule_update_cpu: expect initial_state_source " |
| 2407 | "to be contiguous per pool slot."); |
| 2408 | CHECK(initial_state_source.size(0) >= batch_size); |
| 2409 | CHECK_EQ(v_num_heads % num_heads, 0); |
| 2410 | TORCH_CHECK( |
| 2411 | A_log.sizes() == at::IntArrayRef({v_num_heads}), |
| 2412 | "Input tensor shape mismatch: expected ", |
| 2413 | at::IntArrayRef({v_num_heads}), |
| 2414 | ", got ", |
| 2415 | A_log.sizes()); |
| 2416 | |
| 2417 | int64_t q_strideB = q.stride(1); |
| 2418 | int64_t q_strideS = q.stride(0); |
| 2419 | int64_t q_strideH = q.stride(2); |
| 2420 | int64_t k_strideB = k.stride(1); |
| 2421 | int64_t k_strideS = k.stride(0); |
| 2422 | int64_t k_strideH = k.stride(2); |
| 2423 | int64_t v_strideB = v.stride(1); |
| 2424 | int64_t v_strideS = v.stride(0); |
| 2425 | int64_t v_strideH = v.stride(2); |
| 2426 | // IMPORTANT: To make the kernal compatible with vLLM KV cache layout |
| 2427 | int64_t state_slot_stride = initial_state_source.stride(0); |
| 2428 | at::Tensor core_attn_out = at::empty({batch_size, seq_len, v_num_heads, v_head_dim}, q.options()); |
| 2429 | at::Tensor qk_scale_buf = at::empty({2 * batch_size, seq_len, num_heads}, at::kFloat); |
| 2430 | |
| 2431 | auto launch = [&](auto state_tag) { |
| 2432 | using state_t = decltype(state_tag); |
| 2433 | CPU_DISPATCH_REDUCED_FLOATING_TYPES_EXT( |
| 2434 | q.scalar_type(), A_log.scalar_type(), |
| 2435 | "fused_sigmoid_gating_delta_rule_update_kernel_impl", [&] { |
| 2436 | fused_sigmoid_gating_delta_rule_update_kernel_impl<scalar_t, param_t, state_t>( |
| 2437 | q.data_ptr<scalar_t>(), |
| 2438 | k.data_ptr<scalar_t>(), |
| 2439 | v.data_ptr<scalar_t>(), |
| 2440 | A_log.data_ptr<param_t>(), |
| 2441 | a.data_ptr<scalar_t>(), |
| 2442 | dt_bias.data_ptr<scalar_t>(), |
| 2443 | b.data_ptr<scalar_t>(), |
| 2444 | initial_state_indices.data_ptr<int32_t>(), |
| 2445 | initial_state_source.data_ptr<state_t>(), |
| 2446 | core_attn_out.data_ptr<scalar_t>(), |
| 2447 | qk_scale_buf.data_ptr<float>(), |
| 2448 | seq_len, |
| 2449 | batch_size, |
| 2450 | num_heads, |
| 2451 | head_dim, |
| 2452 | v_num_heads, |
| 2453 | v_head_dim, |
| 2454 | q_strideB, |
| 2455 | q_strideS, |
| 2456 | q_strideH, |
| 2457 | k_strideB, |
| 2458 | k_strideS, |
| 2459 | k_strideH, |
| 2460 | v_strideB, |
| 2461 | v_strideS, |
| 2462 | v_strideH, |
| 2463 | state_slot_stride, |
| 2464 | use_qk_l2norm_in_kernel, |
| 2465 | softplus_threshold); |
| 2466 | }); |
| 2467 | }; |
| 2468 | if (initial_state_source.scalar_type() == at::kFloat) { |
| 2469 | launch(float()); |
| 2470 | } else if (initial_state_source.scalar_type() == at::kBFloat16) { |
| 2471 | launch(at::BFloat16()); |
| 2472 | } else if (initial_state_source.scalar_type() == at::kHalf) { |
| 2473 | launch(at::Half()); |
| 2474 | } else { |
| 2475 | TORCH_CHECK( |
| 2476 | false, |
| 2477 | "fused_sigmoid_gating_delta_rule_update_cpu: initial_state dtype must " |
| 2478 | "be float16, bfloat16, or float32"); |
| 2479 | } |
| 2480 | return core_attn_out; |
| 2481 | } |
| 2482 | |
| 2483 | // Speculative-decode update (multi-token, multi-slot rollback). |
| 2484 | // q: [T, HK, EK] k: [T, HK, EK] v: [T, HV, EV] |
| 2485 | // a: [T, HV] b: [T, HV] |
| 2486 | // initial_state_source: [N_slots, HV, EK, EV] FP16, BF16, or FP32 persisted state |
| 2487 | // (updated in place; recurrent arithmetic remains FP32) |
| 2488 | // spec_state_indices: [batch, S] INT32 (S = num_spec + 1) |
| 2489 | // num_accepted_tokens: [batch] INT32 |
| 2490 | // cu_seqlens: [batch + 1] INT32 |
| 2491 | // Returns output: [T, HV, EV] |
| 2492 | at::Tensor fused_sigmoid_gating_delta_rule_update_spec_cpu( |
| 2493 | const at::Tensor& A_log, |
| 2494 | const at::Tensor& dt_bias, |
| 2495 | const at::Tensor& q, |
| 2496 | const at::Tensor& k, |
| 2497 | const at::Tensor& v, |
| 2498 | const at::Tensor& a, |
| 2499 | const at::Tensor& b, |
| 2500 | at::Tensor& initial_state_source, |
| 2501 | const at::Tensor& spec_state_indices, |
| 2502 | const at::Tensor& num_accepted_tokens, |
| 2503 | const at::Tensor& cu_seqlens, |
| 2504 | bool use_qk_l2norm_in_kernel, |
| 2505 | double softplus_beta = 1.0, |
| 2506 | double softplus_threshold = 20.0) { |
| 2507 | CHECK_DIM(3, q); |
| 2508 | CHECK_DIM(3, v); |
| 2509 | CHECK_LAST_DIM_CONTIGUOUS_INPUT(q); |
| 2510 | int64_t total_tokens = q.size(0); |
| 2511 | int64_t num_heads = q.size(1); |
| 2512 | int64_t head_dim = q.size(2); |
| 2513 | int64_t v_num_heads = v.size(1); |
| 2514 | int64_t v_head_dim = v.size(2); |
| 2515 | CHECK_INPUT_SHAPE_DTYPE<true>(k, {total_tokens, num_heads, head_dim}, q.scalar_type()); |
| 2516 | CHECK_INPUT_SHAPE_DTYPE<true>(v, {total_tokens, v_num_heads, v_head_dim}, q.scalar_type()); |
| 2517 | CHECK_INPUT_SHAPE_DTYPE<true>(a, {total_tokens, v_num_heads}, q.scalar_type()); |
| 2518 | CHECK_INPUT_SHAPE_DTYPE<true>(b, {total_tokens, v_num_heads}, q.scalar_type()); |
| 2519 | CHECK_INPUT_SHAPE_DTYPE<true>(dt_bias, {v_num_heads}, q.scalar_type()); |
| 2520 | CHECK_DIM(1, cu_seqlens); |
| 2521 | int64_t batch_size = cu_seqlens.size(0) - 1; |
| 2522 | CHECK_INPUT_SHAPE_DTYPE<true>(num_accepted_tokens, {batch_size}, at::kInt); |
| 2523 | CHECK_INPUT_SHAPE_DTYPE<true>(cu_seqlens, {batch_size + 1}, at::kInt); |
| 2524 | CHECK_DIM(2, spec_state_indices); |
| 2525 | CHECK_CPU(spec_state_indices); |
| 2526 | TORCH_CHECK( |
| 2527 | spec_state_indices.scalar_type() == at::kInt, |
| 2528 | "spec_state_indices must be an int32 CPU tensor"); |
| 2529 | TORCH_CHECK( |
| 2530 | spec_state_indices.size(0) == batch_size, |
| 2531 | "spec_state_indices first dimension must equal batch_size"); |
| 2532 | CHECK_LAST_DIM_CONTIGUOUS_INPUT(spec_state_indices); |
| 2533 | const int32_t* cu_seqlens_ptr = cu_seqlens.data_ptr<int32_t>(); |
| 2534 | int64_t max_seq_len = 0; |
| 2535 | for (int64_t i = 0; i < batch_size; ++i) { |
| 2536 | int64_t span = static_cast<int64_t>(cu_seqlens_ptr[i + 1]) - |
| 2537 | static_cast<int64_t>(cu_seqlens_ptr[i]); |
| 2538 | TORCH_CHECK( |
| 2539 | span >= 0, |
| 2540 | "cu_seqlens must be non-decreasing at index ", |
| 2541 | i); |
| 2542 | if (span > max_seq_len) { |
| 2543 | max_seq_len = span; |
| 2544 | } |
| 2545 | } |
| 2546 | TORCH_CHECK( |
| 2547 | spec_state_indices.size(1) >= max_seq_len, |
| 2548 | "spec_state_indices second dimension must be at least the longest " |
| 2549 | "cu_seqlens span; got ", |
| 2550 | spec_state_indices.size(1), |
| 2551 | " but required ", |
| 2552 | max_seq_len); |
| 2553 | int64_t spec_stride = spec_state_indices.stride(0); |
| 2554 | CHECK_EQ(v_num_heads % num_heads, 0); |
| 2555 | TORCH_CHECK(A_log.sizes() == at::IntArrayRef({v_num_heads})); |
| 2556 | TORCH_CHECK( |
| 2557 | initial_state_source.scalar_type() == at::kFloat || |
| 2558 | initial_state_source.scalar_type() == at::kBFloat16 || |
| 2559 | initial_state_source.scalar_type() == at::kHalf, |
| 2560 | "fused_sigmoid_gating_delta_rule_update_spec_cpu: initial_state dtype " |
| 2561 | "must be float16, bfloat16, or float32"); |
| 2562 | TORCH_CHECK( |
| 2563 | initial_state_source.sizes() == |
| 2564 | at::IntArrayRef( |
| 2565 | {initial_state_source.size(0), v_num_heads, head_dim, v_head_dim}), |
| 2566 | "Input tensor shape mismatch for initial_state_source"); |
| 2567 | CHECK_CPU(initial_state_source); |
| 2568 | TORCH_CHECK( |
| 2569 | initial_state_source.stride(-1) == 1 && |
| 2570 | initial_state_source.stride(-2) == v_head_dim && |
| 2571 | initial_state_source.stride(-3) == head_dim * v_head_dim, |
| 2572 | "fused_sigmoid_gating_delta_rule_update_spec_cpu: expect " |
| 2573 | "initial_state_source to be contiguous per pool slot."); |
| 2574 | TORCH_CHECK(initial_state_source.size(0) >= batch_size, |
| 2575 | "initial_state_source capacity too small: size(0)=", |
| 2576 | initial_state_source.size(0), ", batch_size=", batch_size); |
| 2577 | |
| 2578 | int64_t q_strideT = q.stride(0); |
| 2579 | int64_t q_strideH = q.stride(1); |
| 2580 | int64_t k_strideT = k.stride(0); |
| 2581 | int64_t k_strideH = k.stride(1); |
| 2582 | int64_t v_strideT = v.stride(0); |
| 2583 | int64_t v_strideH = v.stride(1); |
| 2584 | int64_t state_slot_stride = initial_state_source.stride(0); |
| 2585 | |
| 2586 | at::Tensor o = at::empty({total_tokens, v_num_heads, v_head_dim}, q.options()); |
| 2587 | at::Tensor qk_scale_buf = at::empty({2, total_tokens, num_heads}, at::kFloat); |
| 2588 | |
| 2589 | auto launch = [&](auto state_tag) { |
| 2590 | using state_t = decltype(state_tag); |
| 2591 | CPU_DISPATCH_REDUCED_FLOATING_TYPES_EXT( |
| 2592 | q.scalar_type(), A_log.scalar_type(), |
| 2593 | "fused_sigmoid_gating_delta_rule_update_spec_kernel_impl", [&] { |
| 2594 | fused_sigmoid_gating_delta_rule_update_spec_kernel_impl< |
| 2595 | scalar_t, param_t, state_t>( |
| 2596 | q.data_ptr<scalar_t>(), |
| 2597 | k.data_ptr<scalar_t>(), |
| 2598 | v.data_ptr<scalar_t>(), |
| 2599 | A_log.data_ptr<param_t>(), |
| 2600 | a.data_ptr<scalar_t>(), |
| 2601 | dt_bias.data_ptr<scalar_t>(), |
| 2602 | b.data_ptr<scalar_t>(), |
| 2603 | spec_state_indices.data_ptr<int32_t>(), |
| 2604 | num_accepted_tokens.data_ptr<int32_t>(), |
| 2605 | cu_seqlens.data_ptr<int32_t>(), |
| 2606 | initial_state_source.data_ptr<state_t>(), |
| 2607 | o.data_ptr<scalar_t>(), |
| 2608 | qk_scale_buf.data_ptr<float>(), |
| 2609 | total_tokens, |
| 2610 | batch_size, |
| 2611 | spec_stride, |
| 2612 | num_heads, |
| 2613 | head_dim, |
| 2614 | v_num_heads, |
| 2615 | v_head_dim, |
| 2616 | q_strideT, |
| 2617 | q_strideH, |
| 2618 | k_strideT, |
| 2619 | k_strideH, |
| 2620 | v_strideT, |
| 2621 | v_strideH, |
| 2622 | state_slot_stride, |
| 2623 | use_qk_l2norm_in_kernel, |
| 2624 | softplus_threshold); |
| 2625 | }); |
| 2626 | }; |
| 2627 | if (initial_state_source.scalar_type() == at::kFloat) { |
| 2628 | launch(float()); |
| 2629 | } else if (initial_state_source.scalar_type() == at::kBFloat16) { |
| 2630 | launch(at::BFloat16()); |
| 2631 | } else if (initial_state_source.scalar_type() == at::kHalf) { |
| 2632 | launch(at::Half()); |
| 2633 | } else { |
| 2634 | TORCH_CHECK( |
| 2635 | false, |
| 2636 | "fused_sigmoid_gating_delta_rule_update_spec_cpu: initial_state " |
| 2637 | "dtype must be float16, bfloat16, or float32"); |
| 2638 | } |
| 2639 | return o; |
| 2640 | } |
| 2641 | |
| 2642 | // A_log: [num_v_heads] |
| 2643 | // a: [batch, num_v_heads] |
| 2644 | // b: [batch, num_v_heads] |
| 2645 | // dt_bias: [num_v_heads] |
| 2646 | // -A_log.float().exp() * F.softplus(a.float() + dt_bias) |
| 2647 | std::tuple<at::Tensor, at::Tensor> |
| 2648 | fused_gdn_gating_cpu(const at::Tensor& A_log, const at::Tensor& a, const at::Tensor& b, const at::Tensor& dt_bias) { |
| 2649 | CHECK_DIM(1, A_log); |
| 2650 | CHECK_DIM(2, a); |
| 2651 | CHECK_DIM(2, b); |
| 2652 | CHECK_DIM(1, dt_bias); |
| 2653 | CHECK_CONTIGUOUS(a); |
| 2654 | CHECK_EQ(A_log.size(0), a.size(1)); |
| 2655 | CHECK_EQ(A_log.size(0), dt_bias.size(0)); |
| 2656 | int batch = a.size(0); |
| 2657 | int num_heads = a.size(1); |
| 2658 | CHECK_EQ(b.size(0), batch); |
| 2659 | CHECK_EQ(b.size(1), num_heads); |
| 2660 | at::Tensor out = at::empty({1, batch, num_heads}, a.options().dtype(at::kFloat)); |
| 2661 | at::Tensor beta = at::empty({1, batch, num_heads}, b.options()); |
| 2662 | CPU_DISPATCH_REDUCED_FLOATING_TYPES_EXT(a.scalar_type(), A_log.scalar_type(), "fused_gdn_gating_kernel", [&] { |
| 2663 | fused_gdn_gating_kernel_impl<scalar_t>( |
| 2664 | A_log.data_ptr<param_t>(), |
| 2665 | a.data_ptr<scalar_t>(), |
| 2666 | b.data_ptr<scalar_t>(), |
| 2667 | dt_bias.data_ptr<scalar_t>(), |
| 2668 | out.data_ptr<float>(), |
| 2669 | beta.data_ptr<scalar_t>(), |
| 2670 | batch, |
| 2671 | num_heads); |
| 2672 | }); |
| 2673 | return std::make_tuple(out, beta); |
| 2674 | } |