[CPU] Bump up CPU kernels to latest version (#50387)
Signed-off-by: jiang1.li <jiang1.li@intel.com> Co-authored-by: Claude <noreply@anthropic.com>
Li, Jiang committed
Jul 30, 2026 at 19:18 UTC
072a4727820732fb9e1480dd3deffb9070520b41
15 files changed
+2183
-1292
csrc/cpu/sgl-kernels/common.h
+13
@@ -1,11 +1,16 @@
1
// Adapted from
2
// https://github.com/sgl-project/sglang/tree/main/sgl-kernel/csrc/cpu
3
+//
4
+// Synced from
5
+// https://github.com/sgl-project/sglang/tree/7c248dde7fe1f3b5100966f8143f97a9932c22a4/sgl-kernel/csrc/cpu
6
+// Sync date: 2026-07-29
7
8
// clang-format off
9
10
#pragma once
11
12
#include <ATen/ATen.h>
13
+#include <ATen/Dispatch.h>
14
#include <ATen/Parallel.h>
15
16
#if defined(_OPENMP)
@@ -49,6 +54,14 @@ namespace {
54
} \
55
}()
56
57
+// Half + BFloat16, plus one extra scalar type
58
+#define AT_DISPATCH_CASE_REDUCED_FLOATING_TYPES_AND(SCALARTYPE, ...) \
59
+ AT_DISPATCH_CASE_REDUCED_FLOATING_TYPES(__VA_ARGS__) \
60
+ AT_DISPATCH_CASE(SCALARTYPE, __VA_ARGS__)
61
+
62
+#define AT_DISPATCH_REDUCED_FLOATING_TYPES_AND(SCALARTYPE, TYPE, NAME, ...) \
63
+ AT_DISPATCH_SWITCH(TYPE, NAME, AT_DISPATCH_CASE_REDUCED_FLOATING_TYPES_AND(SCALARTYPE, __VA_ARGS__))
64
+
65
// dispatch: bfloat16, float16, int8_t, fp8_e4m3, uint8_t(mxfp4/int4)
66
#define CPU_DISPATCH_PACKED_TYPES(TYPE, ...) \
67
[&] { \
csrc/cpu/sgl-kernels/conv.cpp
+2
-3
@@ -182,14 +182,13 @@ struct tinygemm_kernel<at::BFloat16, K, BLOCK_N, has_bias, has_silu> {
182
183
using fVec = at::vec::Vectorized<float>;
184
using bVec = at::vec::Vectorized<at::BFloat16>;
185
- const fVec one = fVec(1.f);
185
auto storec = [&](auto i, int64_t m) {
186
constexpr int col = i;
187
fVec x0 = fVec(vc[col * 2 + 0]);
188
fVec x1 = fVec(vc[col * 2 + 1]);
189
if constexpr (has_silu) {
191
- x0 = x0 / (one + x0.neg().exp_u20());
192
- x1 = x1 / (one + x1.neg().exp_u20());
190
+ x0 = fast_silu(x0);
191
+ x1 = fast_silu(x1);
192
}
193
bVec out_vec = convert_from_float_ext<at::BFloat16>(x0, x1);
194
out_vec.store(C + m * lda + col * 32);
csrc/cpu/sgl-kernels/fla.cpp
+1701
-1005
@@ -9,929 +9,1550 @@
9
#include "vec_pack.h"
10
11
namespace {
12
-// For this cpu kernel, we have some innovations aside from the existing gpu kernels:
13
-// 1) Use less parallel loops, i.e. 4 including l2_norm.
14
-// 2) Fuse part of l2_norm with the rest of the computation.
12
16
-#define THREAD_BUFFER_ALLOC(dst, base_ptr, offset, type, size) \
17
- type* dst = reinterpret_cast<type*>((base_ptr) + (offset)); \
18
- offset += (size);
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
20
-template <typename scalar_t>
21
-inline void fill_stub(scalar_t* __restrict__ out, float val, int size) {
77
+template <typename scalar_t, int SIZE>
78
+inline void fill_stub(scalar_t* __restrict__ out, float val) {
79
using Vec = at::vec::Vectorized<scalar_t>;
80
constexpr int kVecSize = Vec::size();
81
+ static_assert(SIZE % kVecSize == 0);
82
const Vec data_vec = Vec(static_cast<scalar_t>(val));
25
- int d = 0;
26
-#pragma GCC unroll 4
27
- for (; d <= size - kVecSize; d += kVecSize) {
83
+#pragma GCC unroll 8
84
+ for (int d = 0; d < SIZE; d += kVecSize) {
85
data_vec.store(out + d);
86
}
30
- if (size - d > 0) {
31
- data_vec.store(out + d, size - d);
32
- }
87
}
88
35
-template <typename scalar_t, int64_t chunk_size = 64>
36
-void chunk_gated_delta_rule_kernel_impl(
37
- scalar_t* __restrict__ out, // [B, T, HV, EV]
38
- float* __restrict__ final_state_data, // [N, HV, EK, EV]
39
- const scalar_t* __restrict__ q_orig, // [B, T, HK, EK]
40
- const scalar_t* __restrict__ k_orig, // [B, T, HK, EK]
41
- const scalar_t* __restrict__ v_orig, // [B, T, HV, EV]
42
- const float* __restrict__ g_orig, // [B, T, HV] FP32
43
- const scalar_t* __restrict__ b_orig, // [B, T, HV]
44
- const int32_t* __restrict__ cu_seqlens_ptr, // [N + 1] INT32
45
- float* __restrict__ buff,
46
- scalar_t* __restrict__ reduced_buff,
47
- scalar_t* __restrict__ thread_buff,
48
- const int32_t* __restrict__ chunk_offsets_ptr,
49
- const int32_t* __restrict__ chunk_indices_ptr,
50
- bool use_qk_l2norm_in_kernel,
51
- const int64_t& batch_size,
52
- const int64_t& global_seq_len,
53
- const int64_t& qk_num_head,
54
- const int64_t& v_num_head,
55
- const int64_t& qk_head_size,
56
- const int64_t& v_head_size,
57
- const int64_t& qStrideH,
58
- const int64_t& qStrideT,
59
- const int64_t& kStrideH,
60
- const int64_t& kStrideT,
61
- const int64_t& vStrideH,
62
- const int64_t& vStrideT,
63
- const int64_t& oStrideH,
64
- const int64_t& oStrideT,
65
- const int64_t& global_total_seq_length,
66
- const int64_t& global_num_chunk,
67
- const int64_t& buff_size_16bit_per_thread,
68
- double eps = 1e-5) {
69
- int64_t gStrideH = 1;
70
- int64_t gStrideT = v_num_head;
71
- int64_t bStrideH = 1;
72
- int64_t bStrideT = v_num_head;
73
- int64_t final_state_StrideN = v_num_head * qk_head_size * v_head_size;
74
- int64_t final_state_StrideH = qk_head_size * v_head_size;
75
- int64_t final_state_StrideE = v_head_size;
76
- int64_t head_group = v_num_head / qk_num_head;
77
- float scale = 1.0 / std::sqrt(qk_head_size);
78
- using bVec = at::vec::Vectorized<scalar_t>;
79
- using fVec = at::vec::Vectorized<float>;
80
- constexpr int64_t VecSize = bVec::size();
81
- constexpr int64_t fVecSize = fVec::size();
82
-
83
- // Data pointers
84
- float* g_pad = buff;
85
- float* core_attn_out = g_pad + v_num_head * global_total_seq_length;
86
- float* decay_mask = core_attn_out + batch_size * v_num_head * global_total_seq_length * v_head_size;
87
- float* v_beta_attn = decay_mask + v_num_head * global_total_seq_length * chunk_size;
88
-
89
- scalar_t* q_pad = reduced_buff;
90
- scalar_t* k_pad = q_pad + qk_num_head * global_total_seq_length * qk_head_size;
91
- scalar_t* v_pad = k_pad + qk_num_head * global_total_seq_length * qk_head_size;
92
- scalar_t* k_beta = v_pad + v_num_head * global_total_seq_length * v_head_size;
93
- scalar_t* v_beta = k_beta + v_num_head * global_total_seq_length * qk_head_size;
94
- scalar_t* k_cumdecay_reduced = v_beta + v_num_head * global_total_seq_length * v_head_size;
95
- scalar_t* q_norm_sum = k_cumdecay_reduced + v_num_head * global_total_seq_length * qk_head_size;
96
- scalar_t* k_norm_sum = q_norm_sum + qk_num_head * global_seq_len;
89
+// Portable fallback for non-AVX512 builds (ARM/NEON, old x86 without
90
+// AVX512BF16), vectorized via at::vec::Vectorized<T> (portable across
91
+// AVX2/NEON/generic) mirroring the idioms used elsewhere in this file (see
92
+// l2norm_fwd_kernel_impl's predecessor and fused_gdn_gating_kernel_impl):
93
+// bVec/fVec pairs with convert_to_float/convert_from_float for bf16<->float,
94
+// plain scalar tails for remainders. Two kernels (cumsum_kernel,
95
+// update_key_kernel) write a transposed layout relative to their vectorized
96
+// read axis; those vectorize the load/compute and unpack lanes for the
97
+// (unavoidably strided) store.
98
+template <typename scalar_t, int D, bool has_scale>
99
+struct l2norm_kernel {
100
+ static inline void apply(scalar_t* __restrict__ out, const scalar_t* __restrict__ input, float eps) {
101
+ using bVec = at::vec::Vectorized<scalar_t>;
102
+ using fVec = at::vec::Vectorized<float>;
103
+ constexpr int bVecSize = bVec::size();
104
+ constexpr float scale = 1.f / std::sqrt(static_cast<float>(D));
105
+
106
+ fVec sum_fvec0(0.f), sum_fvec1(0.f);
107
+ int d = 0;
108
+ for (; d <= D - bVecSize; d += bVecSize) {
109
+ bVec in_bvec = bVec::loadu(input + d);
110
+ fVec in0, in1;
111
+ std::tie(in0, in1) = at::vec::convert_to_float(in_bvec);
112
+ sum_fvec0 = sum_fvec0 + in0 * in0;
113
+ sum_fvec1 = sum_fvec1 + in1 * in1;
114
+ }
115
+ float sqsum = vec_reduce_sum(sum_fvec0 + sum_fvec1);
116
+ for (; d < D; ++d) {
117
+ float v = static_cast<float>(input[d]);
118
+ sqsum += v * v;
119
+ }
120
98
- if (use_qk_l2norm_in_kernel) {
99
- at::parallel_for(0, qk_num_head * global_seq_len, 0, [&](int64_t begin, int64_t end) {
100
- int64_t h_qk = 0, l = 0;
101
- data_index_init(begin, h_qk, qk_num_head, l, global_seq_len);
102
- for (int64_t i = begin; i < end; ++i) {
103
- auto q_norm_sum_ptr = q_norm_sum + h_qk * global_seq_len + l;
104
- auto k_norm_sum_ptr = k_norm_sum + h_qk * global_seq_len + l;
105
- float sum_q = float(0);
106
- float sum_k = float(0);
107
- fVec sum_q_fvec = fVec(float(0));
108
- fVec sum_k_fvec = fVec(float(0));
109
- int64_t q_offset = l * qStrideT + h_qk * qStrideH;
110
- int64_t k_offset = l * qStrideT + h_qk * qStrideH;
111
- int64_t d;
112
- for (d = 0; d <= qk_head_size - VecSize; d += VecSize) {
113
- bVec q_bvec = bVec::loadu(q_orig + q_offset + d);
114
- fVec q_fvec0, q_fvec1;
115
- std::tie(q_fvec0, q_fvec1) = at::vec::convert_to_float(q_bvec);
116
- sum_q_fvec += q_fvec0 * q_fvec0;
117
- sum_q_fvec += q_fvec1 * q_fvec1;
118
- bVec k_bvec = bVec::loadu(k_orig + k_offset + d);
119
- fVec k_fvec0, k_fvec1;
120
- std::tie(k_fvec0, k_fvec1) = at::vec::convert_to_float(k_bvec);
121
- sum_k_fvec += k_fvec0 * k_fvec0;
122
- sum_k_fvec += k_fvec1 * k_fvec1;
121
+ float rscale = 1.f / std::sqrt(sqsum + eps);
122
+ fVec rscale_fvec(rscale);
123
+ fVec scale_fvec(scale);
124
+ d = 0;
125
+ for (; d <= D - bVecSize; d += bVecSize) {
126
+ bVec in_bvec = bVec::loadu(input + d);
127
+ fVec in0, in1;
128
+ std::tie(in0, in1) = at::vec::convert_to_float(in_bvec);
129
+ in0 = in0 * rscale_fvec;
130
+ in1 = in1 * rscale_fvec;
131
+ if constexpr (has_scale) {
132
+ in0 = in0 * scale_fvec;
133
+ in1 = in1 * scale_fvec;
134
+ }
135
+ bVec out_bvec = at::vec::convert_from_float<scalar_t>(in0, in1);
136
+ out_bvec.store(out + d);
137
+ }
138
+ for (; d < D; ++d) {
139
+ float v = static_cast<float>(input[d]) * rscale;
140
+ if constexpr (has_scale) {
141
+ v *= scale;
142
+ }
143
+ out[d] = static_cast<scalar_t>(v);
144
+ }
145
+ }
146
+};
147
+
148
+#if defined(CPU_CAPABILITY_AVX512)
149
+template <int D, bool has_scale>
150
+struct l2norm_kernel<at::BFloat16, D, has_scale> {
151
+ static inline void apply(at::BFloat16* __restrict__ out, const at::BFloat16* __restrict__ input, float eps) {
152
+ static_assert(D % 32 == 0);
153
+ constexpr int COLS = D / 32;
154
+
155
+ __m512bh va[COLS];
156
+ __m512 vrscale;
157
+
158
+ constexpr float scale = 1.f / std::sqrt(D);
159
+ __m512 vscale = _mm512_set1_ps(scale);
160
+
161
+ // step 1: load input and do reduce with avx512-bf16
162
+ __m512 vsum = _mm512_set1_ps(0.f);
163
+ auto reduce = [&](auto col) {
164
+ va[col] = (__m512bh)(_mm512_loadu_si512(input + col * 32));
165
+ vsum = _mm512_dpbf16_ps(vsum, va[col], va[col]);
166
+ };
167
+ Unroll<COLS>{}(reduce);
168
+
169
+ float sqsum = _mm512_reduce_add_ps(vsum);
170
+ float rscale = 1.f / std::sqrt(sqsum + eps);
171
+ vrscale = _mm512_set1_ps(rscale);
172
+
173
+ // step 2: apply scale to output
174
+ auto map = [&](auto col) {
175
+ __m512i a16 = (__m512i)va[col];
176
+ __m512 va0 = CVT_BF16_TO_FP32(_mm512_extracti32x8_epi32(a16, 0));
177
+ __m512 va1 = CVT_BF16_TO_FP32(_mm512_extracti32x8_epi32(a16, 1));
178
+ va0 = _mm512_mul_ps(va0, vrscale);
179
+ va1 = _mm512_mul_ps(va1, vrscale);
180
+ // keep the mul order same as torch code:
181
+ // query = l2norm(query) * scale
182
+ if constexpr (has_scale) {
183
+ va0 = _mm512_mul_ps(va0, vscale);
184
+ va1 = _mm512_mul_ps(va1, vscale);
185
+ }
186
+ _mm512_storeu_si512(out + col * 32, (__m512i)(_mm512_cvtne2ps_pbh(va1, va0)));
187
+ };
188
+ Unroll<COLS>{}(map);
189
+ }
190
+};
191
+#endif
192
+
193
+template <typename scalar_t, int CHUNK_SIZE, int BLOCK_H>
194
+struct cumsum_kernel {
195
+ static inline void apply(
196
+ scalar_t* __restrict__ out,
197
+ const scalar_t* __restrict__ input,
198
+ int mb_size,
199
+ int hb_size,
200
+ int ld_src,
201
+ int ld_dst) {
202
+ // out: [hb_size valid rows, CHUNK_SIZE] within a [BLOCK_H, ld_dst] buffer
203
+ // input: [mb_size valid rows, hb_size] within a [CHUNK_SIZE(padded), ld_src] buffer
204
+ // input is contiguous along j (vectorize the load/accumulate); out is
205
+ // contiguous along i instead (transposed), so the store is per-lane.
206
+ using Vec = at::vec::Vectorized<scalar_t>;
207
+ constexpr int VecSize = Vec::size();
208
+ alignas(64) scalar_t lane_buf[VecSize];
209
+ int j = 0;
210
+ for (; j <= hb_size - VecSize; j += VecSize) {
211
+ Vec running(static_cast<scalar_t>(0));
212
+ for (int i = 0; i < CHUNK_SIZE; ++i) {
213
+ if (i < mb_size) {
214
+ running = running + Vec::loadu(input + i * ld_src + j);
215
+ }
216
+ running.store(lane_buf);
217
+ for (int lane = 0; lane < VecSize; ++lane) {
218
+ out[(j + lane) * ld_dst + i] = lane_buf[lane];
219
}
124
- sum_q += vec_reduce_sum(sum_q_fvec);
125
- sum_k += vec_reduce_sum(sum_k_fvec);
126
- q_norm_sum_ptr[0] = static_cast<scalar_t>(float(1) / std::sqrt(sum_q + eps));
127
- k_norm_sum_ptr[0] = static_cast<scalar_t>(float(1) / std::sqrt(sum_k + eps));
128
- data_index_step(h_qk, qk_num_head, l, global_seq_len);
220
}
130
- });
221
+ }
222
+ for (; j < hb_size; ++j) {
223
+ float running = 0.f;
224
+ for (int i = 0; i < CHUNK_SIZE; ++i) {
225
+ if (i < mb_size) {
226
+ running += static_cast<float>(input[i * ld_src + j]);
227
+ }
228
+ out[j * ld_dst + i] = static_cast<scalar_t>(running);
229
+ }
230
+ }
231
}
132
-
133
- // query = query * scale
134
- // k_beta = key * beta.unsqueeze(-1)
135
- // v_beta = value * beta.unsqueeze(-1)
136
- // Padding for q/k/v/beta
137
- at::parallel_for(0, qk_num_head * global_num_chunk, 1, [&](int64_t begin, int64_t end) {
138
- int ompIdx = at::get_thread_num();
139
- int64_t h_qk = 0, c = 0;
140
- data_index_init(begin, h_qk, qk_num_head, c, global_num_chunk);
141
- for ([[maybe_unused]] auto z : c10::irange(begin, end)) {
142
- int64_t ib = chunk_indices_ptr[c * 2]; // idx_batch
143
- int64_t ic = chunk_indices_ptr[c * 2 + 1]; // idx_chunk
144
- int64_t l_orig = cu_seqlens_ptr[ib] + ic * chunk_size;
145
- int64_t l = c * chunk_size;
146
- bool is_tail = (c + 1 == chunk_offsets_ptr[ib + 1]);
147
- int64_t seq_len = cu_seqlens_ptr[ib + 1] - cu_seqlens_ptr[ib];
148
- int64_t real_chunk_size = is_tail ? seq_len - ic * chunk_size : chunk_size;
149
- auto q_orig_ptr = q_orig + h_qk * qStrideH + l_orig * qStrideT;
150
- auto k_orig_ptr = k_orig + h_qk * kStrideH + l_orig * kStrideT;
151
- auto v_orig_ptr = v_orig + l_orig * vStrideT;
152
- auto b_orig_ptr = b_orig + l_orig * bStrideT;
153
- auto q_pad_ptr = q_pad + h_qk * global_total_seq_length * qk_head_size + l * qk_head_size;
154
- auto k_pad_ptr = k_pad + h_qk * global_total_seq_length * qk_head_size + l * qk_head_size;
155
- auto v_pad_ptr = v_pad + l * v_head_size;
156
- auto k_beta_ptr = k_beta + l * qk_head_size;
157
- auto v_beta_ptr = v_beta + l * v_head_size;
158
-
159
- for (int64_t j = 0; j < real_chunk_size; j++) {
160
- auto curr_q_orig = q_orig_ptr + j * qStrideT;
161
- auto curr_k_orig = k_orig_ptr + j * kStrideT;
162
- auto curr_q_pad = q_pad_ptr + j * qk_head_size;
163
- auto curr_k_pad = k_pad_ptr + j * qk_head_size;
164
- auto q_scale =
165
- use_qk_l2norm_in_kernel ? *(q_norm_sum + h_qk * global_seq_len + l_orig + j) : static_cast<scalar_t>(1);
166
- auto k_scale =
167
- use_qk_l2norm_in_kernel ? *(k_norm_sum + h_qk * global_seq_len + l_orig + j) : static_cast<scalar_t>(1);
168
- auto q_scale_vec = bVec(q_scale);
169
- auto k_scale_vec = bVec(k_scale);
170
- int64_t i = 0;
171
- scalar_t scale_reduced = static_cast<scalar_t>(scale);
172
- auto vec_scale_reduced = bVec(scale_reduced);
173
- for (; i < fVecSize * (qk_head_size / fVecSize); i += fVecSize) {
174
- auto tmp0 = bVec::loadu(curr_q_orig + i, fVecSize);
175
- auto tmp1 = tmp0 * q_scale_vec * vec_scale_reduced;
176
- tmp1.store(curr_q_pad + i, fVecSize);
177
- auto tmp3 = bVec::loadu(curr_k_orig + i, fVecSize);
178
- auto tmp4 = tmp3 * k_scale_vec;
179
- tmp4.store(curr_k_pad + i, fVecSize);
232
+};
233
+
234
+#if defined(CPU_CAPABILITY_AVX512)
235
+template <int CHUNK_SIZE, int BLOCK_H>
236
+struct cumsum_kernel<float, CHUNK_SIZE, BLOCK_H> {
237
+ static inline void
238
+ apply(float* __restrict__ out, const float* __restrict__ input, int mb_size, int hb_size, int ld_src, int ld_dst) {
239
+ // vector length of fp32 for avx512
240
+ static_assert(BLOCK_H == 16);
241
+ TORCH_CHECK(hb_size > 0 && hb_size <= BLOCK_H);
242
+ const __mmask16 vmask = static_cast<__mmask16>((1u << hb_size) - 1u);
243
+
244
+ __m512i va[16];
245
+ __m512 vsum = _mm512_set1_ps(0.f);
246
+
247
+ for (int i = 0; i < CHUNK_SIZE; i += 16) {
248
+ // load input data
249
+ Unroll<16>{}([&](auto j) {
250
+ __m512 v;
251
+ if (i + j < mb_size) {
252
+ v = _mm512_maskz_loadu_ps(vmask, input + (i + j) * ld_src);
253
+ } else {
254
+ v = _mm512_setzero_ps();
255
}
181
-
182
- for (auto hi = 0; hi < head_group; hi++) {
183
- int64_t h = h_qk * head_group + hi;
184
- auto curr_v_orig = v_orig_ptr + h * vStrideH + j * vStrideT;
185
- auto curr_b_orig = b_orig_ptr + h * bStrideH + j * bStrideT;
186
- scalar_t b_orig_val_reduced = *(curr_b_orig);
187
- auto curr_v_pad = v_pad_ptr + h * global_total_seq_length * v_head_size + j * v_head_size;
188
- auto curr_k_beta = k_beta_ptr + h * global_total_seq_length * qk_head_size + j * qk_head_size;
189
- auto curr_v_beta = v_beta_ptr + h * global_total_seq_length * v_head_size + j * v_head_size;
190
-
191
- // query = query * scale
192
- // k_beta = key * beta.unsqueeze(-1)
193
- int64_t i = 0;
194
- auto vec_b_reduced = bVec(b_orig_val_reduced);
195
- for (; i < fVecSize * (qk_head_size / fVecSize); i += fVecSize) {
196
- auto tmp0 = bVec::loadu(curr_k_orig + i, fVecSize);
197
- auto tmp2 = tmp0 * k_scale_vec * vec_b_reduced;
198
- tmp2.store(curr_k_beta + i, fVecSize);
199
- }
200
- // v_beta = value * beta.unsqueeze(-1)
201
- i = 0;
202
- for (; i < VecSize * (v_head_size / VecSize); i += VecSize) {
203
- auto tmp3 = bVec::loadu(curr_v_orig + i);
204
- tmp3.store(curr_v_pad + i);
205
- auto tmp5 = tmp3 * vec_b_reduced;
206
- tmp5.store(curr_v_beta + i);
207
- }
256
+ vsum = _mm512_add_ps(vsum, v);
257
+ va[j] = _mm512_castps_si512(vsum);
258
+ });
259
+ // transpose
260
+ transpose_16x16_32bit(va);
261
+ // store output data
262
+ Unroll<16>{}([&](auto j) {
263
+ if (j < hb_size) {
264
+ _mm512_storeu_si512(out + j * ld_dst + i, va[j]);
265
}
266
+ });
267
+ }
268
+ }
269
+};
270
+#endif
271
+
272
+template <typename scalar_t, int CHUNK_SIZE>
273
+struct decay_mask_kernel {
274
+ // decay_mask = ((g.unsqueeze(-1) - g.unsqueeze(-2)).tril().exp().float()).tril()
275
+ static inline void apply(scalar_t* __restrict__ out, const scalar_t* __restrict__ input) {
276
+ using Vec = at::vec::Vectorized<scalar_t>;
277
+ constexpr int VecSize = Vec::size();
278
+ const Vec zero(static_cast<scalar_t>(0));
279
+ for (int row = 0; row < CHUNK_SIZE; ++row) {
280
+ Vec g_row(input[row]);
281
+ Vec limit_vec(static_cast<scalar_t>(row));
282
+ int col = 0;
283
+ for (; col <= CHUNK_SIZE - VecSize; col += VecSize) {
284
+ Vec g_col = Vec::loadu(input + col);
285
+ Vec vc = (g_row - g_col).exp_u20();
286
+ Vec idx = Vec::arange(static_cast<scalar_t>(col), static_cast<scalar_t>(1));
287
+ Vec result = Vec::blendv(zero, vc, idx <= limit_vec);
288
+ result.store(out + row * CHUNK_SIZE + col);
289
+ }
290
+ for (; col < CHUNK_SIZE; ++col) {
291
+ out[row * CHUNK_SIZE + col] =
292
+ col <= row
293
+ ? static_cast<scalar_t>(std::exp(static_cast<float>(input[row]) - static_cast<float>(input[col])))
294
+ : static_cast<scalar_t>(0);
295
+ }
296
+ }
297
+ }
298
+};
299
+
300
+#if defined(CPU_CAPABILITY_AVX512)
301
+template <int CHUNK_SIZE>
302
+struct decay_mask_kernel<float, CHUNK_SIZE> {
303
+ static inline void apply(float* __restrict__ out, const float* __restrict__ input) {
304
+ static_assert(CHUNK_SIZE % 16 == 0);
305
+
306
+ constexpr int ROWS = CHUNK_SIZE;
307
+ constexpr int COLS = CHUNK_SIZE / 16;
308
+
309
+ __m512 va;
310
+ __m512 vb[COLS];
311
+
312
+ // step 1: load g[j]
313
+ auto loadb = [&](auto i) { vb[i] = _mm512_loadu_ps(input + i * 16); };
314
+ Unroll<COLS>{}(loadb);
315
+
316
+ // step2: exp(g[i] - g[j])
317
+ auto compute = [&](auto i) {
318
+ constexpr int row = i / COLS;
319
+ constexpr int col = i % COLS;
320
+
321
+ if constexpr (col == 0) {
322
+ va = _mm512_set1_ps(input[row]);
323
}
324
211
- for (int64_t j = real_chunk_size; j < chunk_size; j++) {
212
- auto curr_q_pad = q_pad_ptr + j * qk_head_size;
213
- auto curr_k_pad = k_pad_ptr + j * qk_head_size;
214
- int64_t i = 0;
215
- auto vec_zero = bVec(0.0);
216
- for (; i < VecSize * (qk_head_size / VecSize); i += VecSize) {
217
- vec_zero.store(curr_q_pad + i);
218
- vec_zero.store(curr_k_pad + i);
325
+ // mask vb[col] (already loaded in step 1) for the lower-triangular region
326
+ constexpr int len = std::max(0, std::min(row + 1 - col * 16, 16));
327
+
328
+ __m512 vc;
329
+ if constexpr (len == 16) {
330
+ vc = _mm512_fexp_u20_ps(va - vb[col]);
331
+ } else if constexpr (len == 0) {
332
+ vc = _mm512_setzero_ps();
333
+ } else {
334
+ vc = _mm512_fexp_u20_ps(va - vb[col]);
335
+ // do mask for vc
336
+ constexpr __mmask16 vmask = (1 << len) - 1;
337
+ vc = _mm512_mask_blend_ps(vmask, _mm512_setzero_ps(), vc);
338
+ }
339
+ _mm512_storeu_ps(out + row * CHUNK_SIZE + col * 16, vc);
340
+ };
341
+ Unroll<ROWS * COLS>{}(compute);
342
+ }
343
+};
344
+#endif
345
+
346
+template <typename scalar_t, int CHUNK_SIZE, bool has_beta>
347
+struct apply_mask_kernel {
348
+ // has_beta: attn2 = -attn * beta * d (strict lower, col < row)
349
+ // !has_beta: attn2 = attn * d (lower incl. diagonal, col <= row)
350
+ static inline void apply(
351
+ scalar_t* __restrict__ attn2,
352
+ const float* __restrict__ attn,
353
+ const scalar_t* __restrict__ beta,
354
+ const float* __restrict__ d,
355
+ int size,
356
+ int b_stride = 0) {
357
+ using bVec = at::vec::Vectorized<scalar_t>;
358
+ using fVec = at::vec::Vectorized<float>;
359
+ constexpr int fVecSize = fVec::size();
360
+ constexpr int bVecSize = bVec::size();
361
+ const fVec zero(0.f);
362
+ for (int row = 0; row < size; ++row) {
363
+ int col_limit = has_beta ? row : row + 1;
364
+ float beta_val = 1.f;
365
+ if constexpr (has_beta) {
366
+ beta_val = -static_cast<float>(beta[row * b_stride]);
367
+ }
368
+ fVec beta_fvec(beta_val);
369
+ fVec limit_fvec(static_cast<float>(col_limit));
370
+ int col = 0;
371
+ for (; col <= CHUNK_SIZE - bVecSize; col += bVecSize) {
372
+ fVec a0 = fVec::loadu(attn + row * CHUNK_SIZE + col);
373
+ fVec a1 = fVec::loadu(attn + row * CHUNK_SIZE + col + fVecSize);
374
+ fVec d0 = fVec::loadu(d + row * CHUNK_SIZE + col);
375
+ fVec d1 = fVec::loadu(d + row * CHUNK_SIZE + col + fVecSize);
376
+ fVec v0 = a0 * beta_fvec * d0;
377
+ fVec v1 = a1 * beta_fvec * d1;
378
+ fVec idx0 = fVec::arange(static_cast<float>(col), 1.f);
379
+ fVec idx1 = fVec::arange(static_cast<float>(col + fVecSize), 1.f);
380
+ v0 = fVec::blendv(zero, v0, idx0 < limit_fvec);
381
+ v1 = fVec::blendv(zero, v1, idx1 < limit_fvec);
382
+ bVec out_bvec = at::vec::convert_from_float<scalar_t>(v0, v1);
383
+ out_bvec.store(attn2 + row * CHUNK_SIZE + col);
384
+ }
385
+ for (; col < CHUNK_SIZE; ++col) {
386
+ float v = 0.f;
387
+ if (col < col_limit) {
388
+ v = attn[row * CHUNK_SIZE + col] * beta_val * d[row * CHUNK_SIZE + col];
389
}
220
- for (auto hi = 0; hi < head_group; hi++) {
221
- int64_t h = h_qk * head_group + hi;
222
- auto curr_v_pad = v_pad_ptr + h * global_total_seq_length * v_head_size + j * v_head_size;
223
- auto curr_k_beta = k_beta_ptr + h * global_total_seq_length * qk_head_size + j * qk_head_size;
224
- auto curr_v_beta = v_beta_ptr + h * global_total_seq_length * v_head_size + j * v_head_size;
225
- int64_t i = 0;
226
- for (; i < VecSize * (qk_head_size / VecSize); i += VecSize) {
227
- vec_zero.store(curr_k_beta + i);
390
+ attn2[row * CHUNK_SIZE + col] = static_cast<scalar_t>(v);
391
+ }
392
+ }
393
+ }
394
+};
395
+
396
+#if defined(CPU_CAPABILITY_AVX512)
397
+template <int CHUNK_SIZE, bool has_beta>
398
+struct apply_mask_kernel<at::BFloat16, CHUNK_SIZE, has_beta> {
399
+ static inline void apply(
400
+ at::BFloat16* __restrict__ attn2,
401
+ const float* __restrict__ attn,
402
+ const at::BFloat16* __restrict__ beta,
403
+ const float* __restrict__ d,
404
+ int size,
405
+ int b_stride = 0) {
406
+ static_assert(CHUNK_SIZE % 16 == 0);
407
+
408
+ constexpr int ROWS = CHUNK_SIZE;
409
+ constexpr int COLS = CHUNK_SIZE / 16;
410
+
411
+ __m512 vbeta;
412
+
413
+ // has_beta: attn2 = -attn * beta * d (strict lower)
414
+ // !has_beta: attn2 = attn * d (lower incl. diagonal)
415
+ auto compute = [&](auto i) {
416
+ constexpr int row = i / COLS;
417
+ constexpr int col = i % COLS;
418
+
419
+ constexpr int len =
420
+ has_beta ? std::max(0, std::min(row - col * 16, 16)) : std::max(0, std::min(row + 1 - col * 16, 16));
421
+ if (row < size) {
422
+ if constexpr (has_beta) {
423
+ if constexpr (col == 0) {
424
+ vbeta = _mm512_set1_ps(-static_cast<float>(beta[row * b_stride]));
425
}
229
- i = 0;
230
- for (; i < VecSize * (v_head_size / VecSize); i += VecSize) {
231
- vec_zero.store(curr_v_pad + i);
232
- vec_zero.store(curr_v_beta + i);
426
+ }
427
+
428
+ __m512 vc;
429
+ if constexpr (len == 0) {
430
+ vc = _mm512_setzero_ps();
431
+ } else {
432
+ constexpr __mmask16 vmask = (1 << len) - 1;
433
+ __m512 va = _mm512_maskz_loadu_ps(vmask, attn + row * CHUNK_SIZE + col * 16);
434
+ __m512 vd = _mm512_maskz_loadu_ps(vmask, d + row * CHUNK_SIZE + col * 16);
435
+ if constexpr (has_beta) {
436
+ vc = _mm512_mul_ps(_mm512_mul_ps(va, vbeta), vd);
437
+ } else {
438
+ vc = _mm512_mul_ps(va, vd);
439
}
440
}
441
+ _mm256_storeu_si256(
442
+ reinterpret_cast<__m256i*>(attn2 + row * CHUNK_SIZE + col * 16), (__m256i)(_mm512_cvtneps_pbh(vc)));
443
+ }
444
+ };
445
+ Unroll<ROWS * COLS>{}(compute);
446
+ }
447
+};
448
+#endif
449
+
450
+template <typename scalar_t, int CHUNK_SIZE>
451
+struct solve_tril_kernel {
452
+ // (I + L)^{-1} via forward substitution, L = strict-lower part of attn2.
453
+ static inline void apply(scalar_t* __restrict__ attn2, int size) {
454
+ using bVec = at::vec::Vectorized<scalar_t>;
455
+ using fVec = at::vec::Vectorized<float>;
456
+ constexpr int fVecSize = fVec::size();
457
+ constexpr int bVecSize = bVec::size();
458
+
459
+ // for len == 0 and row < size, we don't have to write back zero again
460
+ // as in `apply_mask_kernel`, we already set zero for the upper-triangular region
461
+ for (int i = 1; i < size; ++i) {
462
+ scalar_t* __restrict__ row_ptr = attn2 + i * CHUNK_SIZE;
463
+ float vsum[CHUNK_SIZE];
464
+ int j = 0;
465
+ for (; j <= i - bVecSize; j += bVecSize) {
466
+ bVec row_bvec = bVec::loadu(row_ptr + j);
467
+ fVec f0, f1;
468
+ std::tie(f0, f1) = at::vec::convert_to_float(row_bvec);
469
+ f0.store(vsum + j);
470
+ f1.store(vsum + j + fVecSize);
471
+ }
472
+ for (; j < i; ++j) {
473
+ vsum[j] = static_cast<float>(row_ptr[j]);
474
}
236
- // Move to the next query
237
- data_index_step(h_qk, qk_num_head, c, global_num_chunk);
238
- }
239
- });
475
241
- at::parallel_for(0, v_num_head * global_num_chunk, 1, [&](int64_t begin, int64_t end) {
242
- int64_t h = 0, c = 0;
243
- data_index_init(begin, h, v_num_head, c, global_num_chunk);
244
- int ompIdx = at::get_thread_num();
245
- int64_t offset = 0;
246
- scalar_t* thread_buff_ptr = thread_buff + ompIdx * buff_size_16bit_per_thread;
247
- THREAD_BUFFER_ALLOC(k_transpose, thread_buff_ptr, offset, scalar_t, qk_head_size * chunk_size);
248
- THREAD_BUFFER_ALLOC(v_pack, thread_buff_ptr, offset, scalar_t, chunk_size * v_head_size);
249
- THREAD_BUFFER_ALLOC(k_beta_g, thread_buff_ptr, offset, scalar_t, chunk_size * qk_head_size);
250
- THREAD_BUFFER_ALLOC(k_beta_g_pack, thread_buff_ptr, offset, scalar_t, chunk_size * qk_head_size);
251
- THREAD_BUFFER_ALLOC(curr_attn, thread_buff_ptr, offset, float, chunk_size* chunk_size * 2);
252
- THREAD_BUFFER_ALLOC(curr_attn_reduced, thread_buff_ptr, offset, scalar_t, chunk_size * chunk_size);
253
- THREAD_BUFFER_ALLOC(k_cumdecay, thread_buff_ptr, offset, float, chunk_size* qk_head_size * 2);
254
- THREAD_BUFFER_ALLOC(row, thread_buff_ptr, offset, float, chunk_size * 2);
255
- THREAD_BUFFER_ALLOC(updated, thread_buff_ptr, offset, float, chunk_size * 2);
256
- for ([[maybe_unused]] auto z : c10::irange(begin, end)) {
257
- int64_t ib = chunk_indices_ptr[c * 2]; // idx_batch
258
- int64_t ic = chunk_indices_ptr[c * 2 + 1]; // idx_chunk
259
- int64_t l_orig = cu_seqlens_ptr[ib] + ic * chunk_size;
260
- int64_t seq_len = cu_seqlens_ptr[ib + 1] - cu_seqlens_ptr[ib];
261
- int64_t h_qk = h / head_group;
262
- auto curr_g_orig = g_orig + h * gStrideH + l_orig * gStrideT;
263
- auto curr_g_pad = g_pad + h * global_total_seq_length + c * chunk_size;
264
- auto curr_decay_mask = decay_mask + h * global_total_seq_length * chunk_size + c * chunk_size * chunk_size;
265
- auto curr_k_pad = k_pad + h_qk * global_total_seq_length * qk_head_size + c * chunk_size * qk_head_size;
266
- auto curr_k_beta = k_beta + h * global_total_seq_length * qk_head_size + c * chunk_size * qk_head_size;
267
- auto curr_k_cumdecay_reduced =
268
- k_cumdecay_reduced + h * global_total_seq_length * qk_head_size + c * chunk_size * qk_head_size;
269
- auto curr_v_beta = v_beta + h * global_total_seq_length * v_head_size + c * chunk_size * v_head_size;
270
- auto curr_value = v_beta_attn + h * global_total_seq_length * v_head_size + c * chunk_size * v_head_size;
271
-
272
- float acc_val = 0;
273
- for (int64_t i = 0; i < chunk_size; i++) {
274
- // Padding for g
275
- // g = g.cumsum(dim=-1)
276
- // g: [B, HV, num_chunk, chunk_size]
277
- if (ic * chunk_size + i < seq_len) {
278
- acc_val += curr_g_orig[i * gStrideT];
476
+ // row = attn[..., i, :i].clone()
477
+ // sub = attn[..., :i, :i].clone()
478
+ // vsum = row + (row.unsqueeze(-1) * sub).sum(-2)
479
+ for (int k = 0; k < i; ++k) {
480
+ // read BEFORE row_ptr is written back below (row k was finalized in
481
+ // an earlier outer iteration; row i itself is untouched until the
482
+ // final write-back after this loop)
483
+ float va = static_cast<float>(row_ptr[k]);
484
+ fVec va_vec(va);
485
+ const scalar_t* __restrict__ row_k_ptr = attn2 + k * CHUNK_SIZE;
486
+ int jj = 0;
487
+ for (; jj <= k - bVecSize; jj += bVecSize) {
488
+ bVec rk_bvec = bVec::loadu(row_k_ptr + jj);
489
+ fVec rk0, rk1;
490
+ std::tie(rk0, rk1) = at::vec::convert_to_float(rk_bvec);
491
+ fVec vsum0 = fVec::loadu(vsum + jj);
492
+ fVec vsum1 = fVec::loadu(vsum + jj + fVecSize);
493
+ vsum0 = vsum0 + va_vec * rk0;
494
+ vsum1 = vsum1 + va_vec * rk1;
495
+ vsum0.store(vsum + jj);
496
+ vsum1.store(vsum + jj + fVecSize);
497
}
280
- curr_g_pad[i] = acc_val;
281
- // decay_mask = ((g.unsqueeze(-1) - g.unsqueeze(-2)).tril().exp().float()).tril()
282
- // decay_mask: [B, HV, num_chunk, chunk_size, chunk_size]
283
- float curr_g_pad_i = static_cast<float>(curr_g_pad[i]);
284
- auto vec_curr_g_pad_i = fVec(curr_g_pad_i);
285
- int64_t j = 0;
286
- int64_t len = i + 1;
287
- for (; j < fVecSize * (len / fVecSize); j += fVecSize) {
288
- auto tmp0 = fVec::loadu(curr_g_pad + j);
289
- auto tmp1 = vec_curr_g_pad_i - tmp0;
290
- auto tmp2 = tmp1.exp_u20();
291
- tmp2.store(curr_decay_mask + i * chunk_size + j);
292
- }
293
- if (j < len) {
294
- auto tmp0 = fVec::loadu(curr_g_pad + j, len - j);
295
- auto tmp1 = vec_curr_g_pad_i - tmp0;
296
- auto tmp2 = tmp1.exp_u20();
297
- tmp2.store(curr_decay_mask + i * chunk_size + j, len - j);
498
+ for (; jj < k; ++jj) {
499
+ vsum[jj] += va * static_cast<float>(row_k_ptr[jj]);
500
}
501
}
502
301
- // attn = k_beta @ key.transpose(-1, -2)
302
- // attn: [B, HV, num_chunk, chunk_size, chunk_size]
303
- // transpose and pack for key
304
- if constexpr (brgemm_supported()) {
305
- pack_vnni<scalar_t>(
306
- /* dst */ k_transpose,
307
- /* src */ curr_k_pad,
308
- /* N */ chunk_size,
309
- /* K */ qk_head_size,
310
- /* ld_src */ qk_head_size,
311
- /* ld_dst */ chunk_size);
312
- // k_beta @ key.transpose(-1, -2)
313
- at::native::cpublas::brgemm(
314
- /* M */ chunk_size,
315
- /* N */ chunk_size,
316
- /* K */ qk_head_size,
317
- /* lda */ qk_head_size,
318
- /* ldb */ chunk_size,
319
- /* ldc */ chunk_size,
320
- /* add_C */ false,
321
- /* A */ curr_k_beta,
322
- /* B */ k_transpose,
323
- /* C */ curr_attn);
324
- } else {
325
- blas_gemm(
326
- at::native::TransposeType::Transpose,
327
- at::native::TransposeType::NoTranspose,
328
- chunk_size,
329
- chunk_size,
330
- qk_head_size,
331
- 1.0f,
332
- curr_k_pad,
333
- qk_head_size,
334
- curr_k_beta,
335
- qk_head_size,
336
- 0.0f,
337
- curr_attn,
338
- chunk_size);
503
+ j = 0;
504
+ for (; j <= i - bVecSize; j += bVecSize) {
505
+ fVec f0 = fVec::loadu(vsum + j);
506
+ fVec f1 = fVec::loadu(vsum + j + fVecSize);
507
+ bVec out_bvec = at::vec::convert_from_float<scalar_t>(f0, f1);
508
+ out_bvec.store(row_ptr + j);
509
}
340
- // attn = attn * decay_mask
341
- for (int64_t m = 0; m < chunk_size; m++) {
342
- at::vec::map2<float>(
343
- [](fVec x, fVec y) { return fVec(0) - x * y; },
344
- curr_attn + m * chunk_size,
345
- curr_attn + m * chunk_size,
346
- curr_decay_mask + m * chunk_size,
347
- chunk_size);
510
+ for (; j < i; ++j) {
511
+ row_ptr[j] = static_cast<scalar_t>(vsum[j]);
512
}
513
+ }
514
350
- // chunk decay
351
- // attn: [B, HV, num_chunk, chunk_size, chunk_size]
352
- // mask = torch.triu(torch.ones(chunk_size, chunk_size, dtype=torch.bool, device=query.device), diagonal=0)
353
- // attn = -attn.masked_fill(mask, 0)
354
- // attn[..., i, :i] = row + (row.unsqueeze(-1) * sub).sum(-2) [B, HV, num_chunk, i]
355
- // attn = attn + torch.eye(chunk_size, dtype=attn.dtype, device=attn.device)
356
- // attn = -attn.masked_fill(mask, 0)
357
- for (int i = 0; i < chunk_size; i++) {
358
- const auto vec_zero = fVec(0);
359
- int64_t len = chunk_size - i;
360
- int64_t front = len % fVecSize;
361
- int64_t j = i;
362
- // first masked vec for alignment
363
- if (front > 0) {
364
- vec_zero.store(curr_attn + i * chunk_size + j, front);
365
- j += front;
366
- }
367
- for (; j < fVecSize * (chunk_size / fVecSize); j += fVecSize) {
368
- vec_zero.store(curr_attn + i * chunk_size + j);
515
+ // attn = attn + torch.eye(chunk_size)
516
+ for (int i = 0; i < size; ++i) {
517
+ attn2[i * CHUNK_SIZE + i] = static_cast<scalar_t>(static_cast<float>(attn2[i * CHUNK_SIZE + i]) + 1.f);
518
+ }
519
+ }
520
+};
521
+
522
+#if defined(CPU_CAPABILITY_AVX512)
523
+template <int CHUNK_SIZE>
524
+struct solve_tril_kernel<at::BFloat16, CHUNK_SIZE> {
525
+ static inline void apply(at::BFloat16* __restrict__ attn2, int size) {
526
+ static_assert(CHUNK_SIZE % 16 == 0);
527
+
528
+ constexpr int COLS = CHUNK_SIZE / 16;
529
+
530
+ __m512 va;
531
+ __m512 vb[COLS];
532
+ __m512 vsum[COLS];
533
+
534
+ // for len == 0 and row < size, we don't have to write back zero again
535
+ // as in `apply_mask_kernel`, we already set zero for the upper-triangular region
536
+ for (int i = 1; i < size; ++i) {
537
+ // load row attn[..., i, :i]
538
+ at::BFloat16* __restrict__ row_ptr = attn2 + i * CHUNK_SIZE;
539
+ Unroll<COLS>{}([&](auto col) {
540
+ int len = std::min(i - col * 16, 16);
541
+ if (len > 0) {
542
+ const __mmask16 vmask = (1 << len) - 1;
543
+ vsum[col] = CVT_BF16_TO_FP32(_mm256_maskz_loadu_epi16(vmask, row_ptr + col * 16));
544
}
545
+ });
546
+
547
+ // row = attn[..., i, :i].clone()
548
+ // sub = attn[..., :i, :i].clone()
549
+ // vsum = row + (row.unsqueeze(-1) * sub).sum(-2)
550
+ for (int k = 0; k < i; ++k) {
551
+ va = _mm512_set1_ps(static_cast<float>(row_ptr[k]));
552
+
553
+ const at::BFloat16* __restrict__ row_k_ptr = attn2 + k * CHUNK_SIZE;
554
+ Unroll<COLS>{}([&](auto col) {
555
+ int len = std::min(k - col * 16, 16);
556
+ if (len > 0) {
557
+ const __mmask16 vmask = (1 << len) - 1;
558
+ vb[col] = CVT_BF16_TO_FP32(_mm256_maskz_loadu_epi16(vmask, row_k_ptr + col * 16));
559
+ vsum[col] = _mm512_fmadd_ps(va, vb[col], vsum[col]);
560
+ }
561
+ });
562
}
371
- for (int i = 1; i < chunk_size; i++) {
372
- // row = attn[..., i, :i] [B, HK, num_chunk, i]
373
- int64_t j = 0;
374
- int64_t len = i;
375
- for (; j < fVecSize * (len / fVecSize); j += fVecSize) {
376
- auto tmp0 = fVec::loadu(curr_attn + i * chunk_size + j);
377
- tmp0.store(row + j);
378
- }
379
- if (j < len) {
380
- auto tmp0 = fVec::loadu(curr_attn + i * chunk_size + j, len - j);
381
- tmp0.store(row + j, len - j);
563
+
564
+ // attn[..., i, :i] = vsum
565
+ Unroll<COLS>{}([&](auto col) {
566
+ int len = std::min(i - col * 16, 16);
567
+ if (len > 0) {
568
+ const __mmask16 vmask = (1 << len) - 1;
569
+ _mm256_mask_storeu_epi16(row_ptr + col * 16, vmask, (__m256i)(_mm512_cvtneps_pbh(vsum[col])));
570
}
383
- // (row.unsqueeze(-1) * sub).sum(-2)
384
- fill_stub(updated, 0, i);
385
- for (int k = 0; k < i; k++) {
386
- float row_k = row[k];
387
- auto vec_row_k = fVec(row_k);
388
- int64_t j = 0;
389
- int64_t len = i;
390
- for (; j < fVecSize * (len / fVecSize); j += fVecSize) {
391
- auto tmp0 = fVec::loadu(curr_attn + k * chunk_size + j);
392
- auto tmp1 = vec_row_k * tmp0;
393
- auto tmp2 = fVec::loadu(updated + j);
394
- auto tmp3 = tmp1 + tmp2;
395
- tmp3.store(updated + j);
571
+ });
572
+ }
573
+
574
+ // attn = attn + torch.eye(chunk_size)
575
+ for (int i = 0; i < size; ++i) {
576
+ attn2[i * CHUNK_SIZE + i] += 1.f;
577
+ }
578
+ }
579
+};
580
+#endif
581
+
582
+template <typename scalar_t, int CHUNK_SIZE, int D, bool has_beta, bool has_g>
583
+struct apply_beta_kernel {
584
+ static inline void apply(
585
+ scalar_t* __restrict__ out,
586
+ const scalar_t* __restrict__ input,
587
+ const scalar_t* __restrict__ beta,
588
+ const float* __restrict__ g,
589
+ int size,
590
+ int ld_src,
591
+ int ld_dst,
592
+ int b_stride) {
593
+ using bVec = at::vec::Vectorized<scalar_t>;
594
+ using fVec = at::vec::Vectorized<float>;
595
+ constexpr int fVecSize = fVec::size();
596
+ constexpr int bVecSize = bVec::size();
597
+
598
+ for (int i = 0; i < size; ++i) {
599
+ float scale = 1.f;
600
+ if constexpr (has_beta) {
601
+ scale *= static_cast<float>(beta[i * b_stride]);
602
+ }
603
+ if constexpr (has_g) {
604
+ scale *= std::exp(g[i]);
605
+ }
606
+ fVec scale_fvec(scale);
607
+ int d = 0;
608
+ for (; d <= D - bVecSize; d += bVecSize) {
609
+ bVec in_bvec = bVec::loadu(input + i * ld_src + d);
610
+ fVec in0, in1;
611
+ std::tie(in0, in1) = at::vec::convert_to_float(in_bvec);
612
+ in0 = in0 * scale_fvec;
613
+ in1 = in1 * scale_fvec;
614
+ bVec out_bvec = at::vec::convert_from_float<scalar_t>(in0, in1);
615
+ out_bvec.store(out + i * ld_dst + d);
616
+ }
617
+ for (; d < D; ++d) {
618
+ out[i * ld_dst + d] = static_cast<scalar_t>(static_cast<float>(input[i * ld_src + d]) * scale);
619
+ }
620
+ }
621
+ }
622
+};
623
+
624
+#if defined(CPU_CAPABILITY_AVX512)
625
+template <int CHUNK_SIZE, int D, bool has_beta, bool has_g>
626
+struct apply_beta_kernel<at::BFloat16, CHUNK_SIZE, D, has_beta, has_g> {
627
+ static inline void apply(
628
+ at::BFloat16* __restrict__ out,
629
+ const at::BFloat16* __restrict__ input,
630
+ const at::BFloat16* __restrict__ beta,
631
+ const float* __restrict__ g,
632
+ int size,
633
+ int ld_src,
634
+ int ld_dst,
635
+ int b_stride) {
636
+ static_assert(D % 32 == 0);
637
+ constexpr int COLS = D / 16;
638
+
639
+ // get g.exp() and g is padded to CHUNK_SIZE
640
+ alignas(64) float g_arr[CHUNK_SIZE];
641
+ if constexpr (has_g) {
642
+ Unroll<CHUNK_SIZE / 16>{}([&](auto col) {
643
+ __m512 vg = _mm512_loadu_ps(g + col * 16);
644
+ __m512 vg_exp = _mm512_fexp_u20_ps(vg);
645
+ _mm512_storeu_ps(g_arr + col * 16, vg_exp);
646
+ });
647
+ }
648
+
649
+ for (int i = 0; i < size; ++i) {
650
+ __m512 vbeta;
651
+ if constexpr (has_beta) {
652
+ vbeta = _mm512_set1_ps(static_cast<float>(beta[i * b_stride]));
653
+ }
654
+ __m512 vg;
655
+ if constexpr (has_g) {
656
+ vg = _mm512_set1_ps(g_arr[i]);
657
+ }
658
+
659
+ Unroll<COLS>{}([&](auto col) {
660
+ // load for 0, 2, 4, 6
661
+ if constexpr (col % 2 == 0) {
662
+ __m512i a16 = _mm512_loadu_si512(input + i * ld_src + col * 16);
663
+ __m512 va0 = CVT_BF16_TO_FP32(_mm512_extracti32x8_epi32(a16, 0));
664
+ __m512 va1 = CVT_BF16_TO_FP32(_mm512_extracti32x8_epi32(a16, 1));
665
+ if constexpr (has_beta) {
666
+ va0 = _mm512_mul_ps(va0, vbeta);
667
+ va1 = _mm512_mul_ps(va1, vbeta);
668
}
397
- if (j < len) {
398
- auto tmp0 = fVec::loadu(curr_attn + k * chunk_size + j, len - j);
399
- auto tmp1 = vec_row_k * tmp0;
400
- auto tmp2 = fVec::loadu(updated + j);
401
- auto tmp3 = tmp1 + tmp2;
402
- tmp3.store(updated + j, len - j);
669
+ if constexpr (has_g) {
670
+ va0 = _mm512_mul_ps(va0, vg);
671
+ va1 = _mm512_mul_ps(va1, vg);
672
}
673
+ _mm512_storeu_si512(out + i * ld_dst + col * 16, (__m512i)(_mm512_cvtne2ps_pbh(va1, va0)));
674
}
405
- // attn[..., i, :i] = row + sum(...)
406
- j = 0;
407
- len = i;
408
- for (; j < fVecSize * (len / fVecSize); j += fVecSize) {
409
- auto tmp0 = fVec::loadu(row + j);
410
- auto tmp1 = fVec::loadu(updated + j);
411
- auto tmp2 = tmp0 + tmp1;
412
- tmp2.store(curr_attn + i * chunk_size + j);
413
- }
414
- if (j < len) {
415
- auto tmp0 = fVec::loadu(row + j, len - j);
416
- auto tmp1 = fVec::loadu(updated + j, len - j);
417
- auto tmp2 = tmp0 + tmp1;
418
- tmp2.store(curr_attn + i * chunk_size + j, len - j);
675
+ });
676
+ }
677
+ }
678
+};
679
+#endif
680
+
681
+template <typename scalar_t, int D>
682
+struct update_kernel {
683
+ static inline void
684
+ apply(scalar_t* __restrict__ out, const float* __restrict__ input, int size, int ld_src, int ld_dst) {
685
+ using bVec = at::vec::Vectorized<scalar_t>;
686
+ using fVec = at::vec::Vectorized<float>;
687
+ constexpr int fVecSize = fVec::size();
688
+ constexpr int bVecSize = bVec::size();
689
+
690
+ for (int i = 0; i < size; ++i) {
691
+ int d = 0;
692
+ for (; d <= D - bVecSize; d += bVecSize) {
693
+ fVec f0 = fVec::loadu(input + i * ld_src + d);
694
+ fVec f1 = fVec::loadu(input + i * ld_src + d + fVecSize);
695
+ bVec out_bvec = at::vec::convert_from_float<scalar_t>(f0, f1);
696
+ out_bvec.store(out + i * ld_dst + d);
697
+ }
698
+ for (; d < D; ++d) {
699
+ out[i * ld_dst + d] = static_cast<scalar_t>(input[i * ld_src + d]);
700
+ }
701
+ }
702
+ }
703
+};
704
+
705
+#if defined(CPU_CAPABILITY_AVX512)
706
+template <int D>
707
+struct update_kernel<at::BFloat16, D> {
708
+ static inline void
709
+ apply(at::BFloat16* __restrict__ out, const float* __restrict__ input, int size, int ld_src, int ld_dst) {
710
+ static_assert(D % 32 == 0);
711
+ constexpr int COLS = D / 16;
712
+
713
+ for (int i = 0; i < size; ++i) {
714
+ Unroll<COLS>{}([&](auto col) {
715
+ if constexpr (col % 2 == 0) {
716
+ __m512 va0 = _mm512_loadu_ps(input + i * ld_src + (col + 0) * 16);
717
+ __m512 va1 = _mm512_loadu_ps(input + i * ld_src + (col + 1) * 16);
718
+ __m512i a16 = (__m512i)(_mm512_cvtne2ps_pbh(va1, va0));
719
+ _mm512_storeu_si512(out + i * ld_dst + col * 16, a16);
720
}
721
+ });
722
+ }
723
+ }
724
+};
725
+#endif
726
+
727
+template <typename scalar_t, int D>
728
+struct update_value_kernel {
729
+ static inline void apply(
730
+ scalar_t* __restrict__ v_prime2,
731
+ const scalar_t* __restrict__ v,
732
+ const float* __restrict__ v_prime,
733
+ int size,
734
+ int padded_size,
735
+ int v_strideT) {
736
+ using bVec = at::vec::Vectorized<scalar_t>;
737
+ using fVec = at::vec::Vectorized<float>;
738
+ constexpr int fVecSize = fVec::size();
739
+ constexpr int bVecSize = bVec::size();
740
+
741
+ // v2' = v - v'
742
+ for (int i = 0; i < size; ++i) {
743
+ int d = 0;
744
+ for (; d <= D - bVecSize; d += bVecSize) {
745
+ bVec v_bvec = bVec::loadu(v + i * v_strideT + d);
746
+ fVec v0, v1;
747
+ std::tie(v0, v1) = at::vec::convert_to_float(v_bvec);
748
+ fVec vp0 = fVec::loadu(v_prime + i * D + d);
749
+ fVec vp1 = fVec::loadu(v_prime + i * D + d + fVecSize);
750
+ v0 = v0 - vp0;
751
+ v1 = v1 - vp1;
752
+ bVec out_bvec = at::vec::convert_from_float<scalar_t>(v0, v1);
753
+ out_bvec.store(v_prime2 + i * D + d);
754
}
421
- for (int i = 0; i < chunk_size; i++) {
422
- curr_attn[i * chunk_size + i] += 1.0f;
423
- at::vec::map<scalar_t>(
424
- [](fVec x) { return x; }, curr_attn_reduced + i * chunk_size, curr_attn + i * chunk_size, chunk_size);
755
+ for (; d < D; ++d) {
756
+ float val = static_cast<float>(v[i * v_strideT + d]) - v_prime[i * D + d];
757
+ v_prime2[i * D + d] = static_cast<scalar_t>(val);
758
}
759
+ }
760
+ // pad the last chunk
761
+ const bVec zero_bvec(static_cast<scalar_t>(0));
762
+ for (int i = size; i < padded_size; ++i) {
763
+ int d = 0;
764
+ for (; d <= D - bVecSize; d += bVecSize) {
765
+ zero_bvec.store(v_prime2 + i * D + d);
766
+ }
767
+ for (; d < D; ++d) {
768
+ v_prime2[i * D + d] = static_cast<scalar_t>(0);
769
+ }
770
+ }
771
+ }
772
+};
773
+
774
+#if defined(CPU_CAPABILITY_AVX512)
775
+template <int D>
776
+struct update_value_kernel<at::BFloat16, D> {
777
+ static inline void apply(
778
+ at::BFloat16* __restrict__ v_prime2,
779
+ const at::BFloat16* __restrict__ v,
780
+ const float* __restrict__ v_prime,
781
+ int size,
782
+ int padded_size,
783
+ int v_strideT) {
784
+ static_assert(D % 32 == 0);
785
+ constexpr int COLS = D / 16;
786
+
787
+ // v2' = v - v'
788
+ for (int i = 0; i < size; ++i) {
789
+ Unroll<COLS>{}([&](auto col) {
790
+ // load for 0, 2, 4, 6
791
+ if constexpr (col % 2 == 0) {
792
+ __m512i v16 = _mm512_loadu_si512(v + i * v_strideT + col * 16);
793
+ __m512 va0 = CVT_BF16_TO_FP32(_mm512_extracti32x8_epi32(v16, 0));
794
+ __m512 va1 = CVT_BF16_TO_FP32(_mm512_extracti32x8_epi32(v16, 1));
795
+
796
+ __m512 v_prime0 = _mm512_loadu_ps(v_prime + i * D + col * 16);
797
+ __m512 v_prime1 = _mm512_loadu_ps(v_prime + i * D + col * 16 + 16);
798
+ va0 = _mm512_sub_ps(va0, v_prime0);
799
+ va1 = _mm512_sub_ps(va1, v_prime1);
800
+ __m512i o16 = (__m512i)(_mm512_cvtne2ps_pbh(va1, va0));
801
+ _mm512_storeu_si512(v_prime2 + i * D + col * 16, o16);
802
+ }
803
+ });
804
+ }
805
427
- // v_beta_attn = attn @ v_beta
428
- // k_cumdecay = attn @ (k_beta * g.exp().unsqueeze(-1))
429
- // v_beta_attn: [B, HV, num_chunk, chunk_size, EV]
430
- // k_beta_g = k_beta * g: [B, HV, num_chunk, chunk_size, EK]
431
- // k_cumdecay: [B, HV, num_chunk, chunk_size, EK]
432
- // pack for value
433
- if constexpr (brgemm_supported()) {
434
- pack_vnni2<scalar_t>(
435
- /* dst */ v_pack,
436
- /* src */ curr_v_beta,
437
- /* N */ chunk_size,
438
- /* K */ v_head_size,
439
- /* ld_src */ v_head_size,
440
- /* ld_dst */ v_head_size);
441
- // value = attn @ v_beta
442
- at::native::cpublas::brgemm(
443
- /* M */ chunk_size,
444
- /* N */ v_head_size,
445
- /* K */ chunk_size,
446
- /* lda */ chunk_size,
447
- /* ldb */ v_head_size,
448
- /* ldc */ v_head_size,
449
- /* add_C */ false,
450
- /* A */ curr_attn_reduced,
451
- /* B */ v_pack,
452
- /* C */ curr_value);
453
- } else {
454
- blas_gemm(
455
- at::native::TransposeType::NoTranspose,
456
- at::native::TransposeType::NoTranspose,
457
- v_head_size,
458
- chunk_size,
459
- chunk_size,
460
- 1.0f,
461
- curr_v_beta,
462
- v_head_size,
463
- curr_attn_reduced,
464
- chunk_size,
465
- 0.0f,
466
- curr_value,
467
- v_head_size);
806
+ // pad the last chunk
807
+ for (int i = size; i < padded_size; ++i) {
808
+ Unroll<COLS>{}([&](auto col) {
809
+ if constexpr (col % 2 == 0) {
810
+ __m512i v16 = _mm512_setzero_si512();
811
+ _mm512_storeu_si512(v_prime2 + i * D + col * 16, v16);
812
+ }
813
+ });
814
+ }
815
+ }
816
+};
817
+#endif
818
+
819
+template <typename scalar_t, int CHUNK_SIZE, int D>
820
+struct update_key_kernel {
821
+ static inline void apply(
822
+ scalar_t* __restrict__ k_updated,
823
+ const scalar_t* __restrict__ k,
824
+ const float* __restrict__ g,
825
+ int size,
826
+ int k_strideT) {
827
+ // k_updated is transposed: [D, CHUNK_SIZE], k_updated[d, t] = k[t, d] * exp(g_last - g[t]).
828
+ // k's D dim is contiguous (vectorize load/compute); k_updated's D dim is
829
+ // strided (CHUNK_SIZE apart), so the store is unpacked per lane.
830
+ using bVec = at::vec::Vectorized<scalar_t>;
831
+ using fVec = at::vec::Vectorized<float>;
832
+ constexpr int fVecSize = fVec::size();
833
+ constexpr int bVecSize = bVec::size();
834
+ alignas(64) scalar_t lane_buf[bVecSize];
835
+
836
+ const float g_last = g[size - 1];
837
+ for (int t = 0; t < size; ++t) {
838
+ float scale = std::exp(g_last - g[t]);
839
+ fVec scale_fvec(scale);
840
+ int d = 0;
841
+ for (; d <= D - bVecSize; d += bVecSize) {
842
+ bVec k_bvec = bVec::loadu(k + t * k_strideT + d);
843
+ fVec k0, k1;
844
+ std::tie(k0, k1) = at::vec::convert_to_float(k_bvec);
845
+ k0 = k0 * scale_fvec;
846
+ k1 = k1 * scale_fvec;
847
+ bVec out_bvec = at::vec::convert_from_float<scalar_t>(k0, k1);
848
+ out_bvec.store(lane_buf);
849
+ for (int lane = 0; lane < bVecSize; ++lane) {
850
+ k_updated[(d + lane) * CHUNK_SIZE + t] = lane_buf[lane];
851
+ }
852
+ }
853
+ for (; d < D; ++d) {
854
+ k_updated[d * CHUNK_SIZE + t] = static_cast<scalar_t>(static_cast<float>(k[t * k_strideT + d]) * scale);
855
}
469
- // k_beta_g = k_beta * g.exp().unsqueeze(-1)
470
- for (int64_t j = 0; j < chunk_size; j++) {
471
- int64_t i = 0;
472
- float g_exp = std::exp(curr_g_pad[j]);
473
- scalar_t g_exp_reduced = static_cast<scalar_t>(g_exp);
474
- auto vec_g_exp_reduced = bVec(g_exp_reduced);
475
- for (; i < VecSize * (qk_head_size / VecSize); i += VecSize) {
476
- auto tmp0 = bVec::loadu(curr_k_beta + j * qk_head_size + i);
477
- auto tmp1 = tmp0 * vec_g_exp_reduced;
478
- tmp1.store(k_beta_g + j * qk_head_size + i);
856
+ }
857
+ const bVec zero_bvec(static_cast<scalar_t>(0));
858
+ for (int t = size; t < CHUNK_SIZE; ++t) {
859
+ int d = 0;
860
+ for (; d <= D - bVecSize; d += bVecSize) {
861
+ zero_bvec.store(lane_buf);
862
+ for (int lane = 0; lane < bVecSize; ++lane) {
863
+ k_updated[(d + lane) * CHUNK_SIZE + t] = lane_buf[lane];
864
}
865
}
481
- // pack for k_beta_g
866
+ for (; d < D; ++d) {
867
+ k_updated[d * CHUNK_SIZE + t] = static_cast<scalar_t>(0);
868
+ }
869
+ }
870
+ }
871
+};
872
+
873
+#if defined(CPU_CAPABILITY_AVX512)
874
+template <int CHUNK_SIZE, int D>
875
+struct update_key_kernel<at::BFloat16, CHUNK_SIZE, D> {
876
+ static inline void apply(
877
+ at::BFloat16* __restrict__ k_updated,
878
+ const at::BFloat16* __restrict__ k,
879
+ const float* __restrict__ g,
880
+ int size,
881
+ int k_strideT) {
882
+ static_assert(D % 32 == 0);
883
+ const int MB = div_up(size, 16);
884
+ const int KB = D / 16;
885
+
886
+ const float g_last = g[size - 1];
887
+ const __m512 vg_last = _mm512_set1_ps(g_last);
888
+
889
+ float scale_arr[16];
890
+ __m256i va[16];
891
+
892
+ // from [C, D](MB, KB) to [D, C](KB, MB)
893
+ // pad size to 16 in this kernel so that transpose can be done in one loop
894
+ for (int mb = 0; mb < MB; ++mb) {
895
+ const int mb_size = std::min(size - mb * 16, 16);
896
+ // prepare exp(g_last - g)
897
+ __m512 vg = _mm512_loadu_ps(g + mb * 16);
898
+ _mm512_storeu_ps(scale_arr, _mm512_fexp_u20_ps(_mm512_sub_ps(vg_last, vg)));
899
+ for (int kb = 0; kb < KB; ++kb) {
900
+ const at::BFloat16* __restrict__ k_ptr = k + mb * 16 * k_strideT + kb * 16;
901
+ at::BFloat16* __restrict__ k_updated_ptr = k_updated + kb * 16 * CHUNK_SIZE + mb * 16;
902
+ // load 16 regs
903
+ Unroll<16>{}([&](auto m) {
904
+ if (m < mb_size) {
905
+ __m256i v16 = _mm256_loadu_si256(reinterpret_cast<const __m256i*>(k_ptr + m * k_strideT));
906
+ __m512 v32 = _mm512_mul_ps(CVT_BF16_TO_FP32(v16), _mm512_set1_ps(scale_arr[m]));
907
+ va[m] = (__m256i)_mm512_cvtneps_pbh(v32);
908
+ } else {
909
+ va[m] = _mm256_setzero_si256();
910
+ }
911
+ });
912
+ // transpose 16x16
913
+ transpose_16x16_16bit(va);
914
+ // store 16 regs
915
+ Unroll<16>{}(
916
+ [&](auto k) { _mm256_storeu_si256(reinterpret_cast<__m256i*>(k_updated_ptr + k * CHUNK_SIZE), va[k]); });
917
+ }
918
+ }
919
+ }
920
+};
921
+#endif
922
+
923
+
924
+// template head_dim here to reduce extra read
925
+// * normal approach: read inputs 2 times:
926
+// - reduce: 1R
927
+// - scale: 1R + 1W
928
+// * keep input data in register:
929
+// - reduce: 1R
930
+// - scale: 1W
931
+template <typename scalar_t, int D>
932
+void l2norm_fwd_kernel_impl(
933
+ scalar_t* __restrict__ query_norm,
934
+ scalar_t* __restrict__ key_norm,
935
+ const scalar_t* __restrict__ query,
936
+ const scalar_t* __restrict__ key,
937
+ float eps,
938
+ int64_t T,
939
+ int64_t H,
940
+ int64_t q_strideT,
941
+ int64_t q_strideH,
942
+ int64_t k_strideT,
943
+ int64_t k_strideH) {
944
+ // expected to be contuguous
945
+ int64_t qn_strideH = D;
946
+ int64_t kn_strideH = D;
947
+
948
+ // parallel on [B, T, H]
949
+ at::parallel_for(0, T * H, 0, [&](int64_t begin, int64_t end) {
950
+ int64_t t{0}, h{0};
951
+ data_index_init(begin, t, T, h, H);
952
+
953
+ for (int64_t i = begin; i < end; ++i) {
954
+ const scalar_t* __restrict__ q_ptr = query + t * q_strideT + h * q_strideH;
955
+ const scalar_t* __restrict__ k_ptr = key + t * k_strideT + h * k_strideH;
956
+ scalar_t* __restrict__ qn_ptr = query_norm + i * qn_strideH;
957
+ scalar_t* __restrict__ kn_ptr = key_norm + i * kn_strideH;
958
+
959
+ l2norm_kernel<scalar_t, D, true>::apply(qn_ptr, q_ptr, eps);
960
+ l2norm_kernel<scalar_t, D, false>::apply(kn_ptr, k_ptr, eps);
961
+
962
+ // move to the next index
963
+ data_index_step(t, T, h, H);
964
+ }
965
+ });
966
+}
967
+
968
+// g : [B, T, Hv]
969
+// g_ : [B, Hv, NT, C] -> [B, NT, HB, BLOCK_H, C]
970
+// cu_seqlens : [num_seqs + 1]
971
+// chunk_indices : [NT * 2]
972
+template <typename scalar_t, int CHUNK_SIZE>
973
+void chunk_local_cumsum_kernel_impl(
974
+ scalar_t* __restrict__ g_,
975
+ const scalar_t* __restrict__ g,
976
+ const int32_t* __restrict__ cu_seqlens,
977
+ const int32_t* __restrict__ chunk_indices,
978
+ int64_t Hv,
979
+ int64_t NT) {
980
+ constexpr int BLOCK_H = 16;
981
+ int64_t HB = div_up(Hv, int64_t(BLOCK_H));
982
+
983
+ // parallel on [NT * HB] to increase parallelism
984
+ at::parallel_for(0, NT * HB, 0, [&](int64_t begin, int64_t end) {
985
+ int64_t nt{0}, hb{0};
986
+ data_index_init(begin, nt, NT, hb, HB);
987
+
988
+ for (int64_t i = begin; i < end; ++i) {
989
+ int32_t bs = chunk_indices[nt * 2 + 0];
990
+ int32_t batch_offset = cu_seqlens[bs];
991
+ int32_t seqlen = cu_seqlens[bs + 1] - cu_seqlens[bs];
992
+ int64_t mb_start = chunk_indices[nt * 2 + 1] * CHUNK_SIZE;
993
+ int64_t mb_size = std::min(seqlen - mb_start, int64_t(CHUNK_SIZE));
994
+ int64_t hb_size = std::min(Hv - hb * BLOCK_H, int64_t(BLOCK_H));
995
+
996
+ const scalar_t* __restrict__ g_ptr = g + (batch_offset + mb_start) * Hv + hb * BLOCK_H;
997
+ scalar_t* __restrict__ gsum_ptr = g_ + nt * (Hv * CHUNK_SIZE) + hb * (BLOCK_H * CHUNK_SIZE);
998
+ cumsum_kernel<scalar_t, CHUNK_SIZE, BLOCK_H>::apply(gsum_ptr, g_ptr, mb_size, hb_size, Hv, CHUNK_SIZE);
999
+
1000
+ // move to the next index
1001
+ data_index_step(nt, NT, hb, HB);
1002
+ }
1003
+ });
1004
+}
1005
+
1006
+#define DECL_BUF(type, name, size_expr) alignas(64) type name[(size_expr)]
1007
+#define DECL_ZERO_BUF(type, name, size_expr) \
1008
+ DECL_BUF(type, name, size_expr); \
1009
+ fill_stub<type, (size_expr)>(name, 0.f)
1010
+
1011
+// w : [B, T, Hv, D]
1012
+// u : [B, T, Hv, Dv]
1013
+// d : [B, NT, Hv, C, C]
1014
+// k : [B, T, H, D]
1015
+// v : [B, T, Hv, Dv]
1016
+// g : [B, NT, Hv, C]
1017
+// beta : [B, T, Hv]
1018
+// cu_seqlens : [num_seqs + 1]
1019
+// chunk_indices : [NT * 2]
1020
+template <typename scalar_t, int D, int CHUNK_SIZE>
1021
+void chunk_gated_delta_rule_fwd_intra_kernel_impl(
1022
+ scalar_t* __restrict__ w,
1023
+ scalar_t* __restrict__ u,
1024
+ float* __restrict__ d,
1025
+ const scalar_t* __restrict__ k,
1026
+ const scalar_t* __restrict__ v,
1027
+ const float* __restrict__ g,
1028
+ const scalar_t* __restrict__ beta,
1029
+ const int32_t* __restrict__ cu_seqlens,
1030
+ const int32_t* __restrict__ chunk_indices,
1031
+ int64_t H,
1032
+ int64_t Hv,
1033
+ int64_t NT,
1034
+ int64_t k_strideT,
1035
+ int64_t k_strideH,
1036
+ int64_t v_strideT,
1037
+ int64_t v_strideH) {
1038
+ // head group, expect to be 1，2，4 for qwen3.5
1039
+ const int64_t HG = Hv / H;
1040
+
1041
+ // strides
1042
+ const int64_t w_strideT = Hv * D;
1043
+ const int64_t w_strideH = D;
1044
+ const int64_t u_strideT = Hv * D;
1045
+ const int64_t u_strideH = D;
1046
+
1047
+ // [NB]: parallel on [NT, H]
1048
+ // * parallel on num_heads and go sequential on num_heads_v,
1049
+ // * avoid instantialize k_beta (beta * k)
1050
+ // * compute key @ key^T * beta instead of k_beta @ key^T, same as triton impl
1051
+ // * compute key @ key^T once for each k head index and reuse for v head index
1052
+ at::parallel_for(0, NT * H, 0, [&](int64_t begin, int64_t end) {
1053
+ int64_t nt{0}, h{0};
1054
+ data_index_init(begin, nt, NT, h, H);
1055
+
1056
+ // thread local temp buffer
1057
+ DECL_ZERO_BUF(scalar_t, tmp, CHUNK_SIZE * D);
1058
+ DECL_ZERO_BUF(scalar_t, tmp2, CHUNK_SIZE * D);
1059
+ DECL_ZERO_BUF(float, attn, CHUNK_SIZE* CHUNK_SIZE);
1060
+ DECL_ZERO_BUF(scalar_t, attn2, CHUNK_SIZE * CHUNK_SIZE);
1061
+ DECL_ZERO_BUF(float, tmp3, CHUNK_SIZE* D);
1062
+
1063
+ // alias
1064
+ scalar_t* __restrict__ k_packed = tmp;
1065
+ scalar_t* __restrict__ k_beta = tmp;
1066
+ scalar_t* __restrict__ v_beta = tmp;
1067
+ scalar_t* __restrict__ k_beta_packed = tmp2;
1068
+ scalar_t* __restrict__ v_beta_packed = tmp2;
1069
+ float* __restrict__ k_updated = tmp3;
1070
+ float* __restrict__ v_updated = tmp3;
1071
+
1072
+ for (int64_t i = begin; i < end; ++i) {
1073
+ int32_t bs = chunk_indices[nt * 2 + 0];
1074
+ int32_t batch_offset = cu_seqlens[bs];
1075
+ int32_t seqlen = cu_seqlens[bs + 1] - cu_seqlens[bs];
1076
+ int64_t mb_start = chunk_indices[nt * 2 + 1] * CHUNK_SIZE;
1077
+ int64_t mb_size = std::min(seqlen - mb_start, int64_t(CHUNK_SIZE));
1078
+
1079
+ // mb_size` is K in 5.c, 5.g, pad to TILE_K;
1080
+ const int64_t padded_mb_size = div_up((int)mb_size, TILE_K) * TILE_K;
1081
+
1082
+ // step 1: decay_mask = ((g.unsqueeze(-1) - g.unsqueeze(-2)).tril().exp().float()).tril()
1083
+ for (int64_t hv = h * HG; hv < h * HG + HG; ++hv) {
1084
+ const float* __restrict__ g_ptr = g + nt * (Hv * CHUNK_SIZE) + hv * CHUNK_SIZE;
1085
+ float* __restrict__ d_ptr = d + nt * (Hv * CHUNK_SIZE * CHUNK_SIZE) + hv * (CHUNK_SIZE * CHUNK_SIZE);
1086
+ decay_mask_kernel<float, CHUNK_SIZE>::apply(d_ptr, g_ptr);
1087
+ }
1088
+
1089
+ // step 2: attn = key @ key^T
1090
+ const scalar_t* __restrict__ k_ptr = k + (batch_offset + mb_start) * k_strideT + h * k_strideH;
1091
if constexpr (brgemm_supported()) {
483
- pack_vnni2<scalar_t>(
484
- /* dst */ k_beta_g_pack,
485
- /* src */ k_beta_g,
486
- /* N */ chunk_size,
487
- /* K */ qk_head_size,
488
- /* ld_src */ qk_head_size,
489
- /* ld_dst */ qk_head_size);
490
- // k_cumdecay = attn @ k_beta_g
1092
+ pack_vnni<scalar_t>(
1093
+ /* dst */ k_packed,
1094
+ /* src */ k_ptr,
1095
+ /* N */ mb_size,
1096
+ /* K */ D,
1097
+ /* ld_src */ k_strideT,
1098
+ /* ld_dst */ CHUNK_SIZE);
1099
+
1100
at::native::cpublas::brgemm(
492
- /* M */ chunk_size,
493
- /* N */ qk_head_size,
494
- /* K */ chunk_size,
495
- /* lda */ chunk_size,
496
- /* ldb */ qk_head_size,
497
- /* ldc */ qk_head_size,
1101
+ /* M */ mb_size,
1102
+ /* N */ mb_size,
1103
+ /* K */ D,
1104
+ /* lda */ k_strideT,
1105
+ /* ldb */ CHUNK_SIZE,
1106
+ /* ldc */ CHUNK_SIZE,
1107
/* add_C */ false,
499
- /* A */ curr_attn_reduced,
500
- /* B */ k_beta_g_pack,
501
- /* C */ k_cumdecay);
1108
+ /* A */ k_ptr,
1109
+ /* B */ k_packed,
1110
+ /* C */ attn);
1111
} else {
1112
blas_gemm(
1113
+ at::native::TransposeType::Transpose,
1114
at::native::TransposeType::NoTranspose,
505
- at::native::TransposeType::NoTranspose,
506
- qk_head_size,
507
- chunk_size,
508
- chunk_size,
1115
+ mb_size,
1116
+ mb_size,
1117
+ D,
1118
1.0f,
510
- k_beta_g,
511
- qk_head_size,
512
- curr_attn_reduced,
513
- chunk_size,
1119
+ k_ptr,
1120
+ k_strideT,
1121
+ k_ptr,
1122
+ k_strideT,
1123
0.0f,
515
- k_cumdecay,
516
- qk_head_size);
1124
+ attn,
1125
+ CHUNK_SIZE);
1126
}
518
- for (int i = 0; i < chunk_size; i++) {
519
- at::vec::map<scalar_t>(
520
- [](fVec x) { return x; },
521
- curr_k_cumdecay_reduced + i * qk_head_size,
522
- k_cumdecay + i * qk_head_size,
523
- qk_head_size);
1127
+
1128
+ for (int64_t hv = h * HG; hv < h * HG + HG; ++hv) {
1129
+ // step 3: attn2 = -attn * beta * d
1130
+ const scalar_t* __restrict__ beta_ptr = beta + (batch_offset + mb_start) * Hv + hv;
1131
+ const float* __restrict__ d_ptr = d + nt * (Hv * CHUNK_SIZE * CHUNK_SIZE) + hv * (CHUNK_SIZE * CHUNK_SIZE);
1132
+ apply_mask_kernel<scalar_t, CHUNK_SIZE, true>::apply(attn2, attn, beta_ptr, d_ptr, mb_size, Hv);
1133
+
1134
+ // step 4: solve_tril(attn2) -> (I + L)^{-1}, L = strict-lower from step 3
1135
+ // for i in 1..C-1: attn2[i, :i] += (attn2[i, :i] * attn2[:i, :i]).sum(-1)
1136
+ // attn2 += eye(C)
1137
+ solve_tril_kernel<scalar_t, CHUNK_SIZE>::apply(attn2, mb_size);
1138
+
1139
+ // step 5: recompute_w_u
1140
+ // w = attn2 @ (k_beta * g.exp().unsqueeze(-1))
1141
+ // u = attn2 @ value * beta.unsqueeze(-1)
1142
+ const float* __restrict__ g_ptr = g + nt * (Hv * CHUNK_SIZE) + hv * CHUNK_SIZE;
1143
+ const scalar_t* __restrict__ v_ptr = v + (batch_offset + mb_start) * v_strideT + hv * v_strideH;
1144
+
1145
+ // 5.a key = key * beta * g.exp
1146
+ apply_beta_kernel<scalar_t, CHUNK_SIZE, D, true, true>::apply(
1147
+ k_beta, k_ptr, beta_ptr, g_ptr, mb_size, k_strideT, D, Hv);
1148
+
1149
+ // 5.b pack key
1150
+ if constexpr (brgemm_supported()) {
1151
+ pack_vnni2<scalar_t>(
1152
+ /* dst */ k_beta_packed,
1153
+ /* src */ k_beta,
1154
+ /* K */ mb_size,
1155
+ /* N */ D,
1156
+ /* ld_src */ D,
1157
+ /* ld_dst */ D);
1158
+
1159
+ // 5.c w = attn2 @ k_beta
1160
+ at::native::cpublas::brgemm(
1161
+ /* M */ mb_size,
1162
+ /* N */ D,
1163
+ /* K */ padded_mb_size, // mb_size
1164
+ /* lda */ CHUNK_SIZE,
1165
+ /* ldb */ D,
1166
+ /* ldc */ D,
1167
+ /* add_C */ false,
1168
+ /* A */ attn2,
1169
+ /* B */ k_beta_packed,
1170
+ /* C */ k_updated);
1171
+ } else {
1172
+ blas_gemm(
1173
+ at::native::TransposeType::NoTranspose,
1174
+ at::native::TransposeType::NoTranspose,
1175
+ D,
1176
+ mb_size,
1177
+ padded_mb_size,
1178
+ 1.0f,
1179
+ k_beta,
1180
+ D,
1181
+ attn2,
1182
+ CHUNK_SIZE,
1183
+ 0.0f,
1184
+ k_updated,
1185
+ D);
1186
+ }
1187
+
1188
+ // 5.d k_updated -> w
1189
+ scalar_t* __restrict__ w_ptr = w + (batch_offset + mb_start) * w_strideT + hv * w_strideH;
1190
+ update_kernel<scalar_t, D>::apply(w_ptr, k_updated, mb_size, D, w_strideT);
1191
+
1192
+ // 5.e value = value * beta
1193
+ apply_beta_kernel<scalar_t, CHUNK_SIZE, D, true, false>::apply(
1194
+ v_beta, v_ptr, beta_ptr, nullptr, mb_size, v_strideT, D, Hv);
1195
+
1196
+ // 5.f pack value
1197
+ if constexpr (brgemm_supported()) {
1198
+ pack_vnni2<scalar_t>(
1199
+ /* dst */ v_beta_packed,
1200
+ /* src */ v_beta,
1201
+ /* K */ mb_size,
1202
+ /* N */ D,
1203
+ /* ld_src */ D,
1204
+ /* ld_dst */ D);
1205
+
1206
+ // 5.g u = attn2 @ v_beta
1207
+ at::native::cpublas::brgemm(
1208
+ /* M */ mb_size,
1209
+ /* N */ D,
1210
+ /* K */ padded_mb_size, // mb_size
1211
+ /* lda */ CHUNK_SIZE,
1212
+ /* ldb */ D,
1213
+ /* ldc */ D,
1214
+ /* add_C */ false,
1215
+ /* A */ attn2,
1216
+ /* B */ v_beta_packed,
1217
+ /* C */ v_updated);
1218
+ } else {
1219
+ blas_gemm(
1220
+ at::native::TransposeType::NoTranspose,
1221
+ at::native::TransposeType::NoTranspose,
1222
+ D,
1223
+ mb_size,
1224
+ padded_mb_size,
1225
+ 1.0f,
1226
+ v_beta,
1227
+ D,
1228
+ attn2,
1229
+ CHUNK_SIZE,
1230
+ 0.0f,
1231
+ v_updated,
1232
+ D);
1233
+ }
1234
+
1235
+ // 5.h v_updated -> u
1236
+ scalar_t* __restrict__ u_ptr = u + (batch_offset + mb_start) * u_strideT + hv * u_strideH;
1237
+ update_kernel<scalar_t, D>::apply(u_ptr, v_updated, mb_size, D, u_strideT);
1238
}
1239
526
- // Move to the next query
527
- data_index_step(h, v_num_head, c, global_num_chunk);
1240
+ // move to the next index
1241
+ data_index_step(nt, NT, h, H);
1242
}
1243
+ at::native::cpublas::brgemm_release();
1244
});
1245
+}
1246
531
- // for each chunk
532
- at::parallel_for(0, batch_size * v_num_head, 1, [&](int64_t begin, int64_t end) {
533
- int64_t b = 0, h = 0;
534
- data_index_init(begin, b, batch_size, h, v_num_head);
535
- int ompIdx = at::get_thread_num();
536
- int64_t offset =
537
- /* k_transpose */ qk_head_size * chunk_size +
538
- /* v_pack */ chunk_size * v_head_size +
539
- /* k_beta_g */ chunk_size * qk_head_size +
540
- /* k_beta_g_pack */ chunk_size * qk_head_size +
541
- /* attn */ chunk_size * chunk_size * 2 +
542
- /* attn_reduced */ chunk_size * chunk_size +
543
- /* k_cumdecay */ chunk_size * qk_head_size * 2 +
544
- /* row */ chunk_size * 2 +
545
- /* updated */ chunk_size * 2;
546
- scalar_t* thread_buff_ptr = thread_buff + ompIdx * buff_size_16bit_per_thread;
547
- THREAD_BUFFER_ALLOC(
548
- curr_last_recurrent_state_reduced, thread_buff_ptr, offset, scalar_t, qk_head_size * v_head_size);
549
- THREAD_BUFFER_ALLOC(
550
- curr_last_recurrent_state_pack_reduced, thread_buff_ptr, offset, scalar_t, qk_head_size * v_head_size);
551
- THREAD_BUFFER_ALLOC(k_transpose_i, thread_buff_ptr, offset, scalar_t, qk_head_size * chunk_size);
552
- THREAD_BUFFER_ALLOC(attn_i, thread_buff_ptr, offset, float, chunk_size* chunk_size * 2);
553
- THREAD_BUFFER_ALLOC(attn_i_reduced, thread_buff_ptr, offset, scalar_t, chunk_size * chunk_size);
554
- THREAD_BUFFER_ALLOC(v_prime, thread_buff_ptr, offset, float, chunk_size* v_head_size * 2);
555
- THREAD_BUFFER_ALLOC(v_prime_reduced, thread_buff_ptr, offset, scalar_t, chunk_size * v_head_size);
556
- THREAD_BUFFER_ALLOC(v_prime_pack_reduced, thread_buff_ptr, offset, scalar_t, chunk_size * v_head_size);
557
- THREAD_BUFFER_ALLOC(qg, thread_buff_ptr, offset, scalar_t, chunk_size * qk_head_size);
558
- THREAD_BUFFER_ALLOC(attn_inter, thread_buff_ptr, offset, float, chunk_size* v_head_size * 2);
559
- THREAD_BUFFER_ALLOC(kg, thread_buff_ptr, offset, scalar_t, chunk_size * qk_head_size);
560
- THREAD_BUFFER_ALLOC(kg_transpose, thread_buff_ptr, offset, scalar_t, qk_head_size * chunk_size);
561
- THREAD_BUFFER_ALLOC(kgv, thread_buff_ptr, offset, float, qk_head_size* v_head_size * 2);
562
-
563
- for ([[maybe_unused]] auto z : c10::irange(begin, end)) {
564
- int64_t start_q = cu_seqlens_ptr[b];
565
- int64_t seq_len = cu_seqlens_ptr[b + 1] - start_q;
566
- int64_t num_chunk = chunk_offsets_ptr[b + 1] - chunk_offsets_ptr[b];
567
- int64_t chunk_offset = chunk_offsets_ptr[b];
568
- int64_t len_offset = chunk_offset * chunk_size;
569
-
570
- int64_t h_qk = h / head_group;
571
- auto out_ptr = out + start_q * oStrideT;
572
- auto curr_q = q_pad + len_offset * qk_head_size +
573
- h_qk * global_total_seq_length * qk_head_size; // [num_chunk, chunk_size, EK]
574
- auto curr_k = k_pad + len_offset * qk_head_size +
575
- h_qk * global_total_seq_length * qk_head_size; // [num_chunk, chunk_size, EK]
576
- auto curr_v = v_beta_attn + h * global_total_seq_length * v_head_size; // [num_chunk, chunk_size, EV]
577
- auto curr_decay_mask =
578
- decay_mask + h * global_total_seq_length * chunk_size; // [num_chunk, chunk_size, chunk_size]
579
- auto curr_k_cumdecay_reduced =
580
- k_cumdecay_reduced + h * global_total_seq_length * qk_head_size; // [num_chunk, chunk_size, EK]
581
- auto curr_last_recurrent_state =
582
- final_state_data + b * final_state_StrideN + h * final_state_StrideH; // [EK, EV]
583
- auto curr_g_pad = g_pad + len_offset + h * global_total_seq_length; // [num_chunk, chunk_size]
584
- auto curr_core_attn_out = core_attn_out + len_offset * v_head_size +
585
- h * global_total_seq_length * v_head_size; // [num_chunk, chunk_size, EV]
586
- for (int64_t c = 0; c < num_chunk; c++) {
587
- for (int i = 0; i < qk_head_size; i++) {
588
- at::vec::map<scalar_t>(
589
- [](fVec x) { return x; },
590
- curr_last_recurrent_state_reduced + i * v_head_size,
591
- curr_last_recurrent_state + i * v_head_size,
592
- v_head_size);
593
- }
594
- auto q_i = curr_q + c * chunk_size * qk_head_size; // [chunk_size, EK]
595
- auto k_i = curr_k + c * chunk_size * qk_head_size; // [chunk_size, EK]
596
- auto v_i = curr_v + (chunk_offset + c) * chunk_size * v_head_size; // [chunk_size, EV]
597
- auto decay_mask_i = curr_decay_mask + (chunk_offset + c) * chunk_size * chunk_size; // [chunk_size, chunk_size]
598
- auto k_cumdecay_i_reduced =
599
- curr_k_cumdecay_reduced + (chunk_offset + c) * chunk_size * qk_head_size; // [chunk_size, EK]
600
- auto g_pad_i = curr_g_pad + c * chunk_size; // [chunk_size]
601
- auto core_attn_out_i = curr_core_attn_out + c * chunk_size * v_head_size; // [chunk_size, EV]
1247
+//
1248
+// out : [B, T, Hv, Dv]
1249
+// state : [num_seqs, Hv, Dv, D]
1250
+// q : [B, T, H, D]
1251
+// k : [B, T, H, D]
1252
+// w : [B, T, Hv, D]
1253
+// u : [B, T, Hv, Dv]
1254
+// g : [B, NT, Hv, C]
1255
+// d : [B, NT, Hv, C, C]
1256
+// cu_seqlens : [num_seqs + 1]
1257
+// chunk_offsets : [num_seqs + 1]
1258
+template <typename scalar_t, int D, int CHUNK_SIZE>
1259
+void chunk_gated_delta_rule_fwd_inter_kernel_impl(
1260
+ scalar_t* __restrict__ out,
1261
+ float* __restrict__ state,
1262
+ const int32_t* __restrict__ indices,
1263
+ const scalar_t* __restrict__ q,
1264
+ const scalar_t* __restrict__ k,
1265
+ const scalar_t* __restrict__ w,
1266
+ const scalar_t* __restrict__ u,
1267
+ const float* __restrict__ g,
1268
+ const float* __restrict__ d,
1269
+ const int32_t* __restrict__ cu_seqlens,
1270
+ const int32_t* __restrict__ chunk_offsets,
1271
+ int64_t H,
1272
+ int64_t Hv,
1273
+ int64_t num_seqs,
1274
+ int64_t q_strideT,
1275
+ int64_t q_strideH,
1276
+ int64_t k_strideT,
1277
+ int64_t k_strideH,
1278
+ int64_t state_strideS) {
1279
+ // head group, expect to be 1，2，4 for qwen3.5
1280
+ const int64_t HG = Hv / H;
1281
+
1282
+ // strides
1283
+ const int64_t w_strideT = Hv * D;
1284
+ const int64_t w_strideH = D;
1285
+ const int64_t u_strideT = Hv * D;
1286
+ const int64_t u_strideH = D;
1287
+ const int64_t o_strideT = Hv * D;
1288
+ const int64_t o_strideH = D;
1289
+
1290
+ // [NB]: parallel on [num_seqs, Hv]
1291
+ // * choose to parallel on Hv instead of H, though this means q @ kT has duplicated compute
1292
+ // * H might be 16 which is not enough to use 32C when num_seqs is small
1293
+ at::parallel_for(0, num_seqs * Hv, 0, [&](int64_t begin, int64_t end) {
1294
+ int64_t bs{0}, hv{0};
1295
+ data_index_init(begin, bs, num_seqs, hv, Hv);
1296
+
1297
+ // thread local temp buffer
1298
+ DECL_ZERO_BUF(scalar_t, tmp, CHUNK_SIZE * D);
1299
+ DECL_ZERO_BUF(scalar_t, tmp2, D * D);
1300
+ DECL_ZERO_BUF(float, tmp3, CHUNK_SIZE* D);
1301
+ DECL_ZERO_BUF(scalar_t, tmp4, CHUNK_SIZE * D);
1302
+ DECL_ZERO_BUF(float, attn, CHUNK_SIZE* CHUNK_SIZE);
1303
+ DECL_ZERO_BUF(scalar_t, attn2, CHUNK_SIZE * CHUNK_SIZE);
1304
+
1305
+ // alias
1306
+ scalar_t* __restrict__ k_packed = tmp;
1307
+ scalar_t* __restrict__ s_packed = tmp2;
1308
+ float* __restrict__ v_prime = tmp3;
1309
+ scalar_t* __restrict__ v_prime2 = tmp;
1310
+ float* __restrict__ attn_inter = tmp3;
1311
+ scalar_t* __restrict__ qg_exp = tmp4;
1312
+ scalar_t* __restrict__ v_packed = tmp4;
1313
+ scalar_t* __restrict__ k_updated = tmp;
1314
+
1315
+ for (int64_t i = begin; i < end; ++i) {
1316
+ int64_t h = hv / HG;
1317
+ int32_t batch_offset = cu_seqlens[bs];
1318
+ int32_t seqlen = cu_seqlens[bs + 1] - cu_seqlens[bs];
1319
+ int64_t nt = chunk_offsets[bs];
1320
+
1321
+ for (int64_t mb_start = 0; mb_start < seqlen; mb_start += CHUNK_SIZE, ++nt) {
1322
+ int64_t mb_size = std::min(seqlen - mb_start, int64_t(CHUNK_SIZE));
1323
1324
+ // mb_size` is K in 4.a, pad to TILE_K;
1325
+ const int64_t padded_mb_size = div_up((int)mb_size, TILE_K) * TILE_K;
1326
+
1327
+ // step 1.a: attn = query @ key^T
1328
// attn_i = (q_i @ k_i.transpose(-1, -2) * decay_mask[:, :, i]).masked_fill_(mask, 0)
604
- // k_transpose_i = k_i.transpose(-1, -2)
1329
+ const scalar_t* __restrict__ q_ptr = q + (batch_offset + mb_start) * q_strideT + h * q_strideH;
1330
+ const scalar_t* __restrict__ k_ptr = k + (batch_offset + mb_start) * k_strideT + h * k_strideH;
1331
if constexpr (brgemm_supported()) {
1332
pack_vnni<scalar_t>(
607
- /* dst */ k_transpose_i,
608
- /* src */ k_i,
609
- /* N */ chunk_size,
610
- /* K */ qk_head_size,
611
- /* ld_src */ qk_head_size,
612
- /* ld_dst */ chunk_size);
613
- // attn_i = q_i @ k_transpose_i
1333
+ /* dst */ k_packed,
1334
+ /* src */ k_ptr,
1335
+ /* N */ mb_size,
1336
+ /* K */ D,
1337
+ /* ld_src */ k_strideT,
1338
+ /* ld_dst */ CHUNK_SIZE);
1339
+
1340
at::native::cpublas::brgemm(
615
- /* M */ chunk_size,
616
- /* N */ chunk_size,
617
- /* K */ qk_head_size,
618
- /* lda */ qk_head_size,
619
- /* ldb */ chunk_size,
620
- /* ldc */ chunk_size,
1341
+ /* M */ mb_size,
1342
+ /* N */ mb_size,
1343
+ /* K */ D,
1344
+ /* lda */ q_strideT,
1345
+ /* ldb */ CHUNK_SIZE,
1346
+ /* ldc */ CHUNK_SIZE,
1347
/* add_C */ false,
622
- /* A */ q_i,
623
- /* B */ k_transpose_i,
624
- /* C */ attn_i);
1348
+ /* A */ q_ptr,
1349
+ /* B */ k_packed,
1350
+ /* C */ attn);
1351
} else {
1352
blas_gemm(
1353
at::native::TransposeType::Transpose,
1354
at::native::TransposeType::NoTranspose,
629
- chunk_size,
630
- chunk_size,
631
- qk_head_size,
1355
+ mb_size,
1356
+ mb_size,
1357
+ D,
1358
1.0f,
633
- k_i,
634
- qk_head_size,
635
- q_i,
636
- qk_head_size,
1359
+ k_ptr,
1360
+ k_strideT,
1361
+ q_ptr,
1362
+ q_strideT,
1363
0.0f,
638
- attn_i,
639
- chunk_size);
640
- }
641
- // attn_i = attn_i * decay_mask_i
642
- for (int64_t m = 0; m < chunk_size; m++) {
643
- auto attn_i_m = attn_i + m * chunk_size;
644
- auto attn_i_reduced_m = attn_i_reduced + m * chunk_size;
645
- auto decay_mask_i_m = decay_mask_i + m * chunk_size;
646
- int64_t n = 0;
647
- for (; n < fVecSize * (chunk_size / fVecSize); n += fVecSize) {
648
- auto tmp0 = fVec::loadu(attn_i_m + n);
649
- auto tmp1 = fVec::loadu(decay_mask_i_m + n);
650
- auto tmp2 = tmp0 * tmp1;
651
- auto tmp3 = at::vec::convert<scalar_t>(tmp2);
652
- tmp3.store(attn_i_reduced_m + n, fVecSize);
653
- }
654
- if (n < chunk_size) {
655
- auto tmp0 = fVec::loadu(attn_i_m + n, chunk_size - n);
656
- auto tmp1 = fVec::loadu(decay_mask_i_m + n, chunk_size - n);
657
- auto tmp2 = tmp0 * tmp1;
658
- auto tmp3 = at::vec::convert<scalar_t>(tmp2);
659
- tmp3.store(attn_i_reduced_m + n, chunk_size - n);
660
- }
661
- }
662
- // mask = torch.triu(torch.ones(chunk_size, chunk_size, dtype=torch.bool, device=query.device), diagonal=1)
663
- // attn_i = attn_i.masked_fill_(mask, 0)
664
- for (int i = 0; i < chunk_size - 1; i++) {
665
- const auto vec_zero = bVec(0);
666
- int64_t len = chunk_size - i - 1;
667
- int64_t front = len % VecSize;
668
- int64_t j = i + 1;
669
- // first masked vec for alignment
670
- if (front > 0) {
671
- vec_zero.store(attn_i_reduced + i * chunk_size + j, front);
672
- j += front;
673
- }
674
- for (; j < VecSize * (chunk_size / VecSize); j += VecSize) {
675
- vec_zero.store(attn_i_reduced + i * chunk_size + j);
676
- }
1364
+ attn,
1365
+ CHUNK_SIZE);
1366
}
1367
679
- // pack for curr_last_recurrent_state
1368
+ // step 1.b: attn = attn * decay_mask.masked_fill_(mask, 0)
1369
+ const float* __restrict__ d_ptr = d + nt * (Hv * CHUNK_SIZE * CHUNK_SIZE) + hv * (CHUNK_SIZE * CHUNK_SIZE);
1370
+ apply_mask_kernel<scalar_t, CHUNK_SIZE, false>::apply(attn2, attn, nullptr, d_ptr, mb_size);
1371
+
1372
+ // step 2.a: v' = w @ state (fuse state *= exp(g_last) with packing)
1373
+ float* __restrict__ s_ptr = state + indices[bs] * state_strideS + hv * (D * D);
1374
+ const float* __restrict__ g_ptr = g + nt * (Hv * CHUNK_SIZE) + hv * (CHUNK_SIZE);
1375
+ float g_last = g_ptr[mb_size - 1];
1376
+ const scalar_t* __restrict__ w_ptr = w + (batch_offset + mb_start) * w_strideT + hv * w_strideH;
1377
if constexpr (brgemm_supported()) {
681
- pack_vnni2<scalar_t>(
682
- /* dst */ curr_last_recurrent_state_pack_reduced,
683
- /* src */ curr_last_recurrent_state_reduced,
684
- /* N */ qk_head_size,
685
- /* K */ v_head_size,
686
- /* ld_src */ v_head_size,
687
- /* ld_dst */ v_head_size);
688
-
689
- // v_prime = k_cumdecay_i @ curr_last_recurrent_state: [chunk_size, EV]
690
- // k_cumdecay_i: [chunk_size, EK]
691
- // curr_last_recurrent_state: [EK, EV]
1378
+ pack_vnni2<scalar_t, D, D>(
1379
+ /* dst */ s_packed,
1380
+ /* src */ s_ptr,
1381
+ /* g_last */ g_last,
1382
+ /* ld_src */ D,
1383
+ /* ld_dst */ D);
1384
+
1385
at::native::cpublas::brgemm(
693
- /* M */ chunk_size,
694
- /* N */ v_head_size,
695
- /* K */ qk_head_size,
696
- /* lda */ qk_head_size,
697
- /* ldb */ v_head_size,
698
- /* ldc */ v_head_size,
1386
+ /* M */ mb_size,
1387
+ /* N */ D,
1388
+ /* K */ D,
1389
+ /* lda */ w_strideT,
1390
+ /* ldb */ D,
1391
+ /* ldc */ D,
1392
/* add_C */ false,
700
- /* A */ k_cumdecay_i_reduced,
701
- /* B */ curr_last_recurrent_state_pack_reduced,
1393
+ /* A */ w_ptr,
1394
+ /* B */ s_packed,
1395
/* C */ v_prime);
1396
} else {
1397
+ // brgemm_supported()==false path: pack_vnni2 above packs the
1398
+ // *unscaled* state into its dst (for this GEMM's B operand) while
1399
+ // separately scaling src in place by exp(g_last) (consumed later,
1400
+ // at step 5.3's state accumulation). Replicate both halves in the
1401
+ // same order: snapshot s_ptr into s_packed BEFORE scaling s_ptr,
1402
+ // not after, since the GEMM below needs the pre-scale state.
1403
+ float g_last_scale = std::exp(g_last);
1404
+ for (int64_t d0 = 0; d0 < D * D; ++d0) {
1405
+ s_packed[d0] = static_cast<scalar_t>(s_ptr[d0]);
1406
+ s_ptr[d0] *= g_last_scale;
1407
+ }
1408
blas_gemm(
1409
at::native::TransposeType::NoTranspose,
1410
at::native::TransposeType::NoTranspose,
707
- v_head_size,
708
- chunk_size,
709
- qk_head_size,
1411
+ D,
1412
+ mb_size,
1413
+ D,
1414
1.0f,
711
- curr_last_recurrent_state_reduced,
712
- v_head_size,
713
- k_cumdecay_i_reduced,
714
- qk_head_size,
1415
+ s_packed,
1416
+ D,
1417
+ w_ptr,
1418
+ w_strideT,
1419
0.0f,
1420
v_prime,
717
- v_head_size);
1421
+ D);
1422
}
1423
720
- // v_new = v_prime = v_i - v_prime
721
- // v_i: [chunk_size, EV]
722
- for (int64_t m = 0; m < chunk_size; m++) {
723
- int64_t i = 0;
724
- for (; i < fVecSize * (v_head_size / fVecSize); i += fVecSize) {
725
- auto tmp0 = fVec::loadu(v_i + m * v_head_size + i);
726
- auto tmp1 = fVec::loadu(v_prime + m * v_head_size + i);
727
- auto tmp2 = tmp0 - tmp1;
728
- auto tmp3 = at::vec::convert<scalar_t>(tmp2);
729
- tmp3.store(v_prime_reduced + m * v_head_size + i, fVecSize);
730
- }
731
- }
1424
+ // step 2.b: v2' = u - v'
1425
+ const scalar_t* __restrict__ u_ptr = u + (batch_offset + mb_start) * u_strideT + hv * u_strideH;
1426
+ update_value_kernel<scalar_t, D>::apply(v_prime2, u_ptr, v_prime, mb_size, padded_mb_size, u_strideT);
1427
733
- // attn_inter = (q_i * g[:, :, i, :, None].exp()) @ last_recurrent_state
734
- // qg = q_i * g[:, :, i, :, None].exp(): [chunk_size, EK]
735
- // q_i: [chunk_size, EK]
736
- // g[:, :, i, :, None]: [chunk_size, 1]
737
- for (int64_t m = 0; m < chunk_size; m++) {
738
- auto g_pad_i_m = g_pad_i + m;
739
- auto g_exp = std::exp(*g_pad_i_m);
740
- int64_t i = 0;
741
- scalar_t g_exp_reduced = static_cast<scalar_t>(g_exp);
742
- auto vec_g_exp_reduced = bVec(g_exp_reduced);
743
- for (; i < VecSize * (qk_head_size / VecSize); i += VecSize) {
744
- auto tmp0 = bVec::loadu(q_i + m * qk_head_size + i);
745
- auto tmp2 = tmp0 * vec_g_exp_reduced;
746
- tmp2.store(qg + m * qk_head_size + i);
747
- }
748
- }
749
- // attn_inter = qg @ curr_last_recurrent_state: [chunk_size, EV]
750
- // curr_last_recurrent_state: [EK, EV]
1428
+ // step 3.a: qg_exp = q * exp(g)
1429
+ apply_beta_kernel<scalar_t, CHUNK_SIZE, D, false, true>::apply(
1430
+ qg_exp, q_ptr, nullptr, g_ptr, mb_size, q_strideT, D, /*b_stride*/ 0);
1431
+
1432
+ // step 3.b: attn_inter = qg_exp @ state
1433
if constexpr (brgemm_supported()) {
1434
at::native::cpublas::brgemm(
753
- /* M */ chunk_size,
754
- /* N */ v_head_size,
755
- /* K */ qk_head_size,
756
- /* lda */ qk_head_size,
757
- /* ldb */ v_head_size,
758
- /* ldc */ v_head_size,
1435
+ /* M */ mb_size,
1436
+ /* N */ D,
1437
+ /* K */ D,
1438
+ /* lda */ D,
1439
+ /* ldb */ D,
1440
+ /* ldc */ D,
1441
/* add_C */ false,
760
- /* A */ qg,
761
- /* B */ curr_last_recurrent_state_pack_reduced,
762
- /* C */ attn_inter);
1442
+ /* A */ qg_exp,
1443
+ /* B */ s_packed,
1444
+ /* C */ attn_inter);
1445
} else {
1446
blas_gemm(
1447
at::native::TransposeType::NoTranspose,
1448
at::native::TransposeType::NoTranspose,
767
- v_head_size,
768
- chunk_size,
769
- qk_head_size,
1449
+ D,
1450
+ mb_size,
1451
+ D,
1452
1.0f,
771
- curr_last_recurrent_state_reduced,
772
- v_head_size,
773
- qg,
774
- qk_head_size,
1453
+ s_packed,
1454
+ D,
1455
+ qg_exp,
1456
+ D,
1457
0.0f,
1458
attn_inter,
777
- v_head_size);
1459
+ D);
1460
}
1461
780
- // core_attn_out[:, :, i] = attn_inter + attn_i @ v_new
781
- // pack for v_prime
1462
+ // step 4.a: attn_inter += attn2 @ v2'
1463
if constexpr (brgemm_supported()) {
1464
pack_vnni2<scalar_t>(
784
- /* dst */ v_prime_pack_reduced,
785
- /* src */ v_prime_reduced,
786
- /* N */ chunk_size,
787
- /* K */ v_head_size,
788
- /* ld_src */ v_head_size,
789
- /* ld_dst */ v_head_size);
790
- // attn_inter = attn_inter + attn_i @ v_new: [chunk_size, EV]
791
- // attn_i: [chunk_size, chunk_size]
792
- // v_new: [chunk_size, EV]
1465
+ /* dst */ v_packed,
1466
+ /* src */ v_prime2,
1467
+ /* K */ padded_mb_size,
1468
+ /* N */ D,
1469
+ /* ld_src */ D,
1470
+ /* ld_dst */ D);
1471
+
1472
at::native::cpublas::brgemm(
794
- /* M */ chunk_size,
795
- /* N */ v_head_size,
796
- /* K */ chunk_size,
797
- /* lda */ chunk_size,
798
- /* ldb */ v_head_size,
799
- /* ldc */ v_head_size,
1473
+ /* M */ mb_size,
1474
+ /* N */ D,
1475
+ /* K */ padded_mb_size,
1476
+ /* lda */ CHUNK_SIZE,
1477
+ /* ldb */ D,
1478
+ /* ldc */ D,
1479
/* add_C */ true,
801
- /* A */ attn_i_reduced,
802
- /* B */ v_prime_pack_reduced,
803
- /* C */ attn_inter);
1480
+ /* A */ attn2,
1481
+ /* B */ v_packed,
1482
+ /* C */ attn_inter);
1483
} else {
1484
blas_gemm(
1485
at::native::TransposeType::NoTranspose,
1486
at::native::TransposeType::NoTranspose,
808
- v_head_size,
809
- chunk_size,
810
- chunk_size,
1487
+ D,
1488
+ mb_size,
1489
+ padded_mb_size,
1490
1.0f,
812
- v_prime_reduced,
813
- v_head_size,
814
- attn_i_reduced,
815
- chunk_size,
1491
+ v_prime2,
1492
+ D,
1493
+ attn2,
1494
+ CHUNK_SIZE,
1495
1.0f,
1496
attn_inter,
818
- v_head_size);
1497
+ D);
1498
}
1499
821
- // core_attn_out[:, :, i] = attn_inter
822
- for (int64_t m = 0; m < chunk_size; m++) {
823
- at::vec::map<float>(
824
- [](fVec x) { return x; }, core_attn_out_i + m * v_head_size, attn_inter + m * v_head_size, v_head_size);
1500
+ // step 4.b: write attn_inter -> out
1501
+ scalar_t* __restrict__ o_ptr = out + (batch_offset + mb_start) * o_strideT + hv * o_strideH;
1502
+ update_kernel<scalar_t, D>::apply(o_ptr, attn_inter, mb_size, D, o_strideT);
1503
+
1504
+ // brgemm_supported()==false path: step 5.2 below overwrites k_updated,
1505
+ // which aliases the same buffer as v_prime2 (both are `tmp`). Snapshot
1506
+ // v_prime2 into v_packed (otherwise brgemm-only, and free by now since
1507
+ // qg_exp was consumed at step 3.b) before that happens, so step 5.3
1508
+ // doesn't read k_updated's data under the v_prime2 name.
1509
+ if constexpr (!brgemm_supported()) {
1510
+ std::copy(v_prime2, v_prime2 + padded_mb_size * D, v_packed);
1511
}
1512
827
- // last_recurrent_state = (
828
- // last_recurrent_state * g[:, :, i, -1, None, None].exp()
829
- // + (k_i * (g[:, :, i, -1, None] - g[:, :, i]).exp()[..., None]).transpose(-1, -2) @ v_new
830
- // )
831
- // 1) last_recurrent_state * g[:, :, i, -1, None, None].exp()
832
- // curr_last_recurrent_state: [EK, EV]
833
- // g[:, :, i, -1, None, None]: [1, 1]
834
- // last_recurrent_state * g[:, :, i, -1, None, None].exp(): [EK, EV]
835
- auto g_pad_i_last = g_pad_i + chunk_size - 1;
836
- auto g_exp_last = std::exp(g_pad_i_last[0]);
837
- for (int64_t m = 0; m < qk_head_size; m++) {
838
- int64_t i = 0;
839
- auto vec_g_exp_last = fVec(g_exp_last);
840
- for (; i < fVecSize * (v_head_size / fVecSize); i += fVecSize) {
841
- auto tmp0 = bVec::loadu(curr_last_recurrent_state_reduced + m * v_head_size + i);
842
- auto tmp1 = at::vec::convert<float>(tmp0);
843
- auto tmp2 = tmp1 * vec_g_exp_last;
844
- tmp2.store(curr_last_recurrent_state + m * v_head_size + i);
845
- }
846
- if (i < v_head_size) {
847
- auto tmp0 = bVec::loadu(curr_last_recurrent_state_reduced + m * v_head_size + i, v_head_size - i);
848
- auto tmp1 = at::vec::convert<float>(tmp0);
849
- auto tmp2 = tmp1 * vec_g_exp_last;
850
- tmp2.store(curr_last_recurrent_state + m * v_head_size + i, v_head_size - i);
851
- }
852
- }
853
- // 2) (k_i * (g[:, :, i, -1, None] - g[:, :, i]).exp()[..., None]).transpose(-1, -2) @ v_new
854
- // k_i: [chunk_size, EK]
855
- // g[:, :, i, -1, None]: [1]
856
- // g[:, :, i]: [chunk_size]
857
- // (g[:, :, i, -1, None] - g[:, :, i]).exp()[..., None]: [chunk_size, 1]
858
- // kg = k_i * (g[:, :, i, -1, None] - g[:, :, i]).exp()[..., None]: [chunk_size, EK]
859
- // (k_i * (g[:, :, i, -1, None] - g[:, :, i]).exp()[..., None]).transpose(-1, -2): [EK, chunk_size]
860
- // v_new: [chunk_size, EV]
861
- // (k_i * (g[:, :, i, -1, None] - g[:, :, i]).exp()[..., None]).transpose(-1, -2) @ v_new: [EK, EV]
862
- // kg = k_i * (g[:, :, i, -1, None] - g[:, :, i]).exp()[..., None]
863
- for (int64_t m = 0; m < chunk_size; m++) {
864
- auto g_exp = std::exp((g_pad_i_last[0] - g_pad_i[m]));
865
- int64_t i = 0;
866
- scalar_t g_exp_reduced = static_cast<scalar_t>(g_exp);
867
- auto vec_g_exp_reduced = bVec(g_exp_reduced);
868
- for (; i < VecSize * (qk_head_size / VecSize); i += VecSize) {
869
- auto tmp0 = bVec::loadu(k_i + m * qk_head_size + i);
870
- auto tmp2 = tmp0 * vec_g_exp_reduced;
871
- tmp2.store(kg + m * qk_head_size + i);
872
- }
873
- }
874
- // kg.transpose(-1, -2): [EK, chunk_size]
875
- at::native::utils::transpose<scalar_t>(
876
- /* M */ chunk_size,
877
- /* N */ qk_head_size,
878
- /* src */ kg,
879
- /* ld_src */ qk_head_size,
880
- /* dst */ kg_transpose,
881
- /* ld_dst */ chunk_size);
882
- // kgv = kg.transpose(-1, -2) @ v_new
883
- // v_new: [chunk_size, EV]
1513
+ // step 5: update state
1514
+ // state_new = state * exp(g_last) + (k * exp(g_last - g)).T @ v2'
1515
+
1516
+ // step 5.1 state *= exp(g_last) fused with step 2.a
1517
+
1518
+ // step 5.2 k' = k * exp(g_last - g).T; TODO: fuse this with 1.a
1519
+ update_key_kernel<scalar_t, CHUNK_SIZE, D>::apply(k_updated, k_ptr, g_ptr, mb_size, k_strideT);
1520
+
1521
+ // step 5.3 state += k' @ v2'
1522
if constexpr (brgemm_supported()) {
1523
at::native::cpublas::brgemm(
886
- /* M */ qk_head_size,
887
- /* N */ v_head_size,
888
- /* K */ chunk_size,
889
- /* lda */ chunk_size,
890
- /* ldb */ v_head_size,
891
- /* ldc */ v_head_size,
892
- /* add_C */ false,
893
- /* A */ kg_transpose,
894
- /* B */ v_prime_pack_reduced,
895
- /* C */ kgv);
1524
+ /* M */ D,
1525
+ /* N */ D,
1526
+ /* K */ padded_mb_size, // mb_size
1527
+ /* lda */ CHUNK_SIZE,
1528
+ /* ldb */ D,
1529
+ /* ldc */ D,
1530
+ /* add_C */ true,
1531
+ /* A */ k_updated,
1532
+ /* B */ v_packed,
1533
+ /* C */ s_ptr);
1534
} else {
1535
blas_gemm(
1536
at::native::TransposeType::NoTranspose,
1537
at::native::TransposeType::NoTranspose,
900
- v_head_size,
901
- qk_head_size,
902
- chunk_size,
1538
+ D,
1539
+ D,
1540
+ padded_mb_size,
1541
1.0f,
904
- v_prime_reduced,
905
- v_head_size,
906
- kg_transpose,
907
- chunk_size,
908
- 0.0f,
909
- kgv,
910
- v_head_size);
911
- }
912
- // last_recurrent_state = 1) + 2)
913
- for (int64_t m = 0; m < qk_head_size; m++) {
914
- at::vec::map2<float>(
915
- [](fVec x, fVec y) { return x + y; },
916
- curr_last_recurrent_state + m * v_head_size,
917
- curr_last_recurrent_state + m * v_head_size,
918
- kgv + m * v_head_size,
919
- v_head_size);
1542
+ v_packed,
1543
+ D,
1544
+ k_updated,
1545
+ CHUNK_SIZE,
1546
+ 1.0f,
1547
+ s_ptr,
1548
+ D);
1549
}
1550
}
1551
923
- // core_attn_out -> output
924
- // output: [B, T, HV, EV]
925
- // core_attn_out: [B, HV, padded_T, EV]
926
- auto curr_out = out_ptr + h * oStrideH;
927
- for (int64_t m = 0; m < seq_len; m++) {
928
- at::vec::map<scalar_t>(
929
- [](fVec x) { return x; }, curr_out + m * oStrideT, curr_core_attn_out + m * v_head_size, v_head_size);
930
- }
931
-
932
- // Move to the next query
933
- data_index_step(b, batch_size, h, v_num_head);
1552
+ // move to the next index
1553
+ data_index_step(bs, num_seqs, hv, Hv);
1554
}
1555
+ at::native::cpublas::brgemm_release();
1556
});
1557
}
1558
@@ -1011,14 +1632,10 @@ void fused_sigmoid_gating_delta_rule_update_kernel_impl(
1632
int64_t d;
1633
#pragma GCC unroll 4
1634
for (d = 0; d <= head_dim - VecSize; d += VecSize) {
1014
- bVec q_bvec = bVec::loadu(q_ptr + q_offset + d);
1015
- fVec q_fvec0, q_fvec1;
1016
- std::tie(q_fvec0, q_fvec1) = at::vec::convert_to_float(q_bvec);
1635
+ auto [q_fvec0, q_fvec1] = load_float_vec2(q_ptr + q_offset + d);
1636
sum_q_fvec += q_fvec0 * q_fvec0;
1637
sum_q_fvec += q_fvec1 * q_fvec1;
1019
- bVec k_bvec = bVec::loadu(k_ptr + k_offset + d);
1020
- fVec k_fvec0, k_fvec1;
1021
- std::tie(k_fvec0, k_fvec1) = at::vec::convert_to_float(k_bvec);
1638
+ auto [k_fvec0, k_fvec1] = load_float_vec2(k_ptr + k_offset + d);
1639
sum_k_fvec += k_fvec0 * k_fvec0;
1640
sum_k_fvec += k_fvec1 * k_fvec1;
1641
}
@@ -1066,14 +1683,11 @@ void fused_sigmoid_gating_delta_rule_update_kernel_impl(
1683
fVec kv_mem_vec1 = fVec(float(0));
1684
for (int di = 0; di < head_dim; ++di) {
1685
fVec k_val_vec = fVec(k_ptr[k_offset + di] * k_scale);
1069
- fVec state_vec0 = fVec::loadu(state_ptr + state_offset + di * v_head_dim + dvi);
1070
- fVec state_vec1 = fVec::loadu(state_ptr + state_offset + di * v_head_dim + dvi + fVecSize);
1686
+ auto [state_vec0, state_vec1] = load_float_vec2(state_ptr + state_offset + di * v_head_dim + dvi);
1687
kv_mem_vec0 = kv_mem_vec0 + state_vec0 * g_val_exp_vec * k_val_vec;
1688
kv_mem_vec1 = kv_mem_vec1 + state_vec1 * g_val_exp_vec * k_val_vec;
1689
}
1074
- bVec v_bvec = bVec::loadu(v_ptr + v_offset + dvi);
1075
- fVec v_vec0, v_vec1;
1076
- std::tie(v_vec0, v_vec1) = at::vec::convert_to_float(v_bvec);
1690
+ auto [v_vec0, v_vec1] = load_float_vec2(v_ptr + v_offset + dvi);
1691
fVec dt_vec0 = (v_vec0 - kv_mem_vec0) * beta_vec;
1692
fVec dt_vec1 = (v_vec1 - kv_mem_vec1) * beta_vec;
1693
fVec o_vec0 = fVec(float(0));
@@ -1081,8 +1695,7 @@ void fused_sigmoid_gating_delta_rule_update_kernel_impl(
1695
for (int di = 0; di < head_dim; ++di) {
1696
fVec q_vec = fVec(q_ptr[q_offset + di] * q_scale);
1697
fVec k_vec = fVec(k_ptr[k_offset + di] * k_scale);
1084
- fVec state_vec0 = fVec::loadu(state_ptr + state_offset + di * v_head_dim + dvi);
1085
- fVec state_vec1 = fVec::loadu(state_ptr + state_offset + di * v_head_dim + dvi + fVecSize);
1698
+ auto [state_vec0, state_vec1] = load_float_vec2(state_ptr + state_offset + di * v_head_dim + dvi);
1699
state_vec0 = state_vec0 * g_val_exp_vec + k_vec * dt_vec0;
1700
state_vec1 = state_vec1 * g_val_exp_vec + k_vec * dt_vec1;
1701
o_vec0 = o_vec0 + state_vec0 * q_vec * scale_vec;
@@ -1289,25 +1902,19 @@ void fused_gdn_gating_kernel_impl(
1902
constexpr int vec_size = bVec::size();
1903
constexpr int fvec_size = fVec::size();
1904
const fVec neg_one(-1.0f);
1292
- const fVec one(1.0f);
1905
at::parallel_for(0, batch, 0, [&](int64_t begin, int64_t end) {
1906
for (int64_t i = begin; i < end; ++i) {
1907
int64_t j = 0;
1908
for (; j < num_heads - (num_heads % vec_size); j += vec_size) {
1297
- fVec A_log_vec0 = fVec::loadu(A_log + j);
1298
- fVec A_log_vec1 = fVec::loadu(A_log + j + fvec_size);
1299
- bVec dt_bias_vec = bVec::loadu(dt_bias + j);
1300
- bVec a_bvec = bVec::loadu(a + i * num_heads + j);
1301
- bVec b_bvec = bVec::loadu(b + i * num_heads + j);
1302
- fVec a0, a1, dt_bias_vec0, dt_bias_vec1, b0, b1;
1303
- std::tie(a0, a1) = at::vec::convert_to_float(a_bvec);
1304
- std::tie(b0, b1) = at::vec::convert_to_float(b_bvec);
1305
- std::tie(dt_bias_vec0, dt_bias_vec1) = at::vec::convert_to_float(dt_bias_vec);
1909
+ auto [A_log_vec0, A_log_vec1] = load_float_vec2(A_log + j);
1910
+ auto [dt_bias_vec0, dt_bias_vec1] = load_float_vec2(dt_bias + j);
1911
+ auto [a0, a1] = load_float_vec2(a + i * num_heads + j);
1912
+ auto [b0, b1] = load_float_vec2(b + i * num_heads + j);
1913
1914
fVec g0 = neg_one * A_log_vec0.exp_u20() * softplus(a0 + dt_bias_vec0);
1915
fVec g1 = neg_one * A_log_vec1.exp_u20() * softplus(a1 + dt_bias_vec1);
1309
- fVec beta0 = one / (one + (neg_one * b0).exp_u20());
1310
- fVec beta1 = one / (one + (neg_one * b1).exp_u20());
1916
+ fVec beta0 = fast_sigmoid(b0);
1917
+ fVec beta1 = fast_sigmoid(b1);
1918
1919
g0.store(out + i * num_heads + j);
1920
g1.store(out + i * num_heads + j + fvec_size);
@@ -1337,26 +1944,19 @@ void fused_gdn_gating_kernel_impl(
1944
constexpr int vec_size = bVec::size();
1945
constexpr int fvec_size = fVec::size();
1946
const fVec neg_one(-1.0f);
1340
- const fVec one(1.0f);
1947
at::parallel_for(0, batch, 0, [&](int64_t begin, int64_t end) {
1948
for (int64_t i = begin; i < end; ++i) {
1949
int64_t j = 0;
1950
for (; j < num_heads - (num_heads % vec_size); j += vec_size) {
1345
- bVec A_log_bvec = bVec::loadu(A_log + j);
1346
- fVec A_log_vec0, A_log_vec1;
1347
- std::tie(A_log_vec0, A_log_vec1) = at::vec::convert_to_float(A_log_bvec);
1348
- bVec dt_bias_vec = bVec::loadu(dt_bias + j);
1349
- bVec a_bvec = bVec::loadu(a + i * num_heads + j);
1350
- bVec b_bvec = bVec::loadu(b + i * num_heads + j);
1351
- fVec a0, a1, dt_bias_vec0, dt_bias_vec1, b0, b1;
1352
- std::tie(a0, a1) = at::vec::convert_to_float(a_bvec);
1353
- std::tie(b0, b1) = at::vec::convert_to_float(b_bvec);
1354
- std::tie(dt_bias_vec0, dt_bias_vec1) = at::vec::convert_to_float(dt_bias_vec);
1951
+ auto [A_log_vec0, A_log_vec1] = load_float_vec2(A_log + j);
1952
+ auto [dt_bias_vec0, dt_bias_vec1] = load_float_vec2(dt_bias + j);
1953
+ auto [a0, a1] = load_float_vec2(a + i * num_heads + j);
1954
+ auto [b0, b1] = load_float_vec2(b + i * num_heads + j);
1955
1956
fVec g0 = neg_one * A_log_vec0.exp_u20() * softplus(a0 + dt_bias_vec0);
1957
fVec g1 = neg_one * A_log_vec1.exp_u20() * softplus(a1 + dt_bias_vec1);
1358
- fVec beta0 = one / (one + (neg_one * b0).exp_u20());
1359
- fVec beta1 = one / (one + (neg_one * b1).exp_u20());
1958
+ fVec beta0 = fast_sigmoid(b0);
1959
+ fVec beta1 = fast_sigmoid(b1);
1960
1961
g0.store(out + i * num_heads + j);
1962
g1.store(out + i * num_heads + j + fvec_size);
@@ -1373,29 +1973,208 @@ void fused_gdn_gating_kernel_impl(
1973
1974
} // anonymous namespace
1975
1376
-template <bool is_last_dim_contiguous>
1377
-inline void
1378
-CHECK_INPUT_SHAPE_DTYPE(const at::Tensor& tensor, const int64_t& dim, const at::IntArrayRef& sizes, at::ScalarType st) {
1379
- TORCH_CHECK(tensor.sizes() == sizes, "Input tensor shape mismatch: expected ", sizes, ", got ", tensor.sizes());
1380
- TORCH_CHECK(tensor.dtype() == st, "Input tensor dtype mismatch");
1381
- CHECK_DIM(dim, tensor);
1382
- if (is_last_dim_contiguous) {
1383
- CHECK_LAST_DIM_CONTIGUOUS_INPUT(tensor);
1384
- } else {
1385
- CHECK_CONTIGUOUS(tensor);
1976
+template <int CHUNK_SIZE>
1977
+std::tuple<at::Tensor, at::Tensor> prepare_chunk_indices(const at::Tensor& cu_seqlens) {
1978
+ int64_t num_seqs = cu_seqlens.size(0) - 1;
1979
+ at::Tensor chunk_offsets = at::empty({num_seqs + 1}, cu_seqlens.options());
1980
+ // get number of chunks and chunk offsets
1981
+ const int32_t* offsets_data = cu_seqlens.data_ptr<int32_t>();
1982
+ int32_t num_chunks = 0;
1983
+ chunk_offsets[0] = 0;
1984
+ for (int64_t row = 0; row < num_seqs; ++row) {
1985
+ num_chunks += div_up(offsets_data[row + 1] - offsets_data[row], CHUNK_SIZE);
1986
+ chunk_offsets[row + 1] = num_chunks;
1987
}
1988
+ // get chunk indices
1989
+ at::Tensor chunk_indices = at::empty({num_chunks, 2}, cu_seqlens.options());
1990
+ int32_t* indices_data = chunk_indices.data_ptr<int32_t>();
1991
+
1992
+ int64_t idx = 0;
1993
+ for (int32_t row = 0; row < num_seqs; ++row) {
1994
+ int32_t num_chunks = div_up(offsets_data[row + 1] - offsets_data[row], CHUNK_SIZE);
1995
+
1996
+ for (int32_t col = 0; col < num_chunks; ++col) {
1997
+ indices_data[idx * 2 + 0] = row;
1998
+ indices_data[idx * 2 + 1] = col;
1999
+ idx++;
2000
+ }
2001
+ }
2002
+ return std::make_tuple(chunk_indices, chunk_offsets);
2003
}
2004
1389
-// query: [B, T, HK, EK]
1390
-// key: [B, T, HK, EK]
1391
-// value: [B, T, HV, EV]
1392
-// g: [B, T, HV] FP32
1393
-// beta: [B, T, HV]
1394
-// initial_state: [N, HV, EK, EV] FP32
1395
-// output_final_state: bool
1396
-// cu_seqlens: [N + 1] INT32
1397
-// head_first: bool
1398
-// use_qk_l2norm_in_kernel: bool
2005
+#define DISPATCH_HEAD_DIM_CASE(launch_macro, hd) \
2006
+ case hd: { \
2007
+ launch_macro(hd); \
2008
+ break; \
2009
+ }
2010
+
2011
+// [NB]: add new head_dim support here
2012
+#define DISPATCH_HEAD_DIM(dim, launch_macro) \
2013
+ switch (dim) { \
2014
+ DISPATCH_HEAD_DIM_CASE(launch_macro, 64) \
2015
+ DISPATCH_HEAD_DIM_CASE(launch_macro, 128) \
2016
+ default: \
2017
+ TORCH_CHECK(false, "Unexpected head dim size, ", dim); \
2018
+ }
2019
+
2020
+#define LAUNCH_L2NORM_KERNEL(HD) \
2021
+ l2norm_fwd_kernel_impl<scalar_t, HD>( \
2022
+ query_norm.data_ptr<scalar_t>(), \
2023
+ key_norm.data_ptr<scalar_t>(), \
2024
+ query.data_ptr<scalar_t>(), \
2025
+ key.data_ptr<scalar_t>(), \
2026
+ eps, \
2027
+ T, \
2028
+ H, \
2029
+ query.stride(1), \
2030
+ query.stride(2), \
2031
+ key.stride(1), \
2032
+ key.stride(2));
2033
+
2034
+std::tuple<at::Tensor, at::Tensor> l2norm_fwd(const at::Tensor& query, const at::Tensor& key, double eps) {
2035
+ int64_t B = query.size(0);
2036
+ int64_t T = query.size(1);
2037
+ int64_t H = query.size(2);
2038
+ int64_t D = query.size(3);
2039
+
2040
+ at::Tensor query_norm = at::empty_like(query);
2041
+ at::Tensor key_norm = at::empty_like(key);
2042
+
2043
+ AT_DISPATCH_REDUCED_FLOATING_TYPES(
2044
+ query.scalar_type(), "l2norm_fwd", [&] { DISPATCH_HEAD_DIM(D, LAUNCH_L2NORM_KERNEL); });
2045
+
2046
+ return std::make_tuple(query_norm, key_norm);
2047
+}
2048
+
2049
+// [NB]: instantiate decay_mask to avoid heavy recomputation in the kernel with exp
2050
+template <int CHUNK_SIZE>
2051
+at::Tensor chunk_local_cumsum(const at::Tensor& g, const at::Tensor& cu_seqlens, const at::Tensor& chunk_indices) {
2052
+ int64_t B = g.size(0);
2053
+ // int64_t T = g.size(1);
2054
+ int64_t Hv = g.size(2);
2055
+ int64_t NT = chunk_indices.size(0);
2056
+
2057
+ at::Tensor g_ = at::empty({B, NT, Hv, CHUNK_SIZE}, g.options());
2058
+ AT_DISPATCH_FLOATING_TYPES(g.scalar_type(), "chunk_local_cumsum", [&] {
2059
+ chunk_local_cumsum_kernel_impl<scalar_t, CHUNK_SIZE>(
2060
+ g_.data_ptr<scalar_t>(),
2061
+ g.data_ptr<scalar_t>(),
2062
+ cu_seqlens.data_ptr<int32_t>(),
2063
+ chunk_indices.data_ptr<int32_t>(),
2064
+ Hv,
2065
+ NT);
2066
+ });
2067
+ return g_;
2068
+}
2069
+
2070
+#define LAUNCH_CHUNK_GATED_DELTA_RULE_FWD_INTRA_KERNEL(HD) \
2071
+ chunk_gated_delta_rule_fwd_intra_kernel_impl<scalar_t, HD, CHUNK_SIZE>( \
2072
+ w.data_ptr<scalar_t>(), \
2073
+ u.data_ptr<scalar_t>(), \
2074
+ decay_mask.data_ptr<float>(), \
2075
+ k.data_ptr<scalar_t>(), \
2076
+ v.data_ptr<scalar_t>(), \
2077
+ g.data_ptr<float>(), \
2078
+ beta.data_ptr<scalar_t>(), \
2079
+ cu_seqlens.data_ptr<int32_t>(), \
2080
+ chunk_indices.data_ptr<int32_t>(), \
2081
+ H, \
2082
+ Hv, \
2083
+ NT, \
2084
+ k.stride(1), \
2085
+ k.stride(2), \
2086
+ v.stride(1), \
2087
+ v.stride(2));
2088
+
2089
+template <int CHUNK_SIZE>
2090
+std::tuple<at::Tensor, at::Tensor, at::Tensor> chunk_gated_delta_rule_fwd_intra(
2091
+ const at::Tensor& k,
2092
+ const at::Tensor& v,
2093
+ const at::Tensor& g,
2094
+ const at::Tensor& beta,
2095
+ const at::Tensor& cu_seqlens,
2096
+ const at::Tensor& chunk_indices) {
2097
+ int64_t B = k.size(0);
2098
+ int64_t T = k.size(1);
2099
+ int64_t H = k.size(2);
2100
+ int64_t D = k.size(3);
2101
+ int64_t Hv = v.size(2);
2102
+ int64_t Dv = v.size(3);
2103
+ int64_t NT = chunk_indices.size(0);
2104
+
2105
+ at::Tensor w = at::empty({B, T, Hv, D}, k.options()); // BFloat16
2106
+ at::Tensor u = at::empty({B, T, Hv, Dv}, k.options()); // BFloat16
2107
+ at::Tensor decay_mask = at::empty({B, NT, Hv, CHUNK_SIZE, CHUNK_SIZE}, g.options()); // Float
2108
+ AT_DISPATCH_REDUCED_FLOATING_TYPES(k.scalar_type(), "chunk_gated_delta_rule_fwd_intra", [&] {
2109
+ DISPATCH_HEAD_DIM(D, LAUNCH_CHUNK_GATED_DELTA_RULE_FWD_INTRA_KERNEL);
2110
+ });
2111
+
2112
+ return std::make_tuple(w, u, decay_mask);
2113
+}
2114
+
2115
+#define LAUNCH_CHUNK_GATED_DELTA_RULE_FWD_INTER_KERNEL(HD) \
2116
+ chunk_gated_delta_rule_fwd_inter_kernel_impl<scalar_t, HD, CHUNK_SIZE>( \
2117
+ o.data_ptr<scalar_t>(), \
2118
+ initial_state.data_ptr<float>(), \
2119
+ initial_state_indices.data_ptr<int32_t>(), \
2120
+ q.data_ptr<scalar_t>(), \
2121
+ k.data_ptr<scalar_t>(), \
2122
+ w.data_ptr<scalar_t>(), \
2123
+ u.data_ptr<scalar_t>(), \
2124
+ g.data_ptr<float>(), \
2125
+ decay_mask.data_ptr<float>(), \
2126
+ cu_seqlens.data_ptr<int32_t>(), \
2127
+ chunk_offsets.data_ptr<int32_t>(), \
2128
+ H, \
2129
+ Hv, \
2130
+ num_seqs, \
2131
+ q.stride(1), \
2132
+ q.stride(2), \
2133
+ k.stride(1), \
2134
+ k.stride(2), \
2135
+ initial_state.stride(0));
2136
+
2137
+template <int CHUNK_SIZE>
2138
+std::tuple<at::Tensor, at::Tensor> chunk_gated_delta_rule_fwd_inter(
2139
+ const at::Tensor& q,
2140
+ const at::Tensor& k,
2141
+ const at::Tensor& w,
2142
+ const at::Tensor& u,
2143
+ const at::Tensor& g,
2144
+ const at::Tensor& decay_mask,
2145
+ const at::Tensor& initial_state,
2146
+ bool output_final_state,
2147
+ const at::Tensor& cu_seqlens,
2148
+ const at::Tensor& chunk_offsets,
2149
+ const at::Tensor& initial_state_indices) {
2150
+ const int64_t B = q.size(0);
2151
+ const int64_t T = q.size(1);
2152
+ const int64_t H = q.size(2);
2153
+ const int64_t D = q.size(3);
2154
+ const int64_t Hv = w.size(2);
2155
+ const int64_t Dv = u.size(3);
2156
+ const int64_t num_seqs = initial_state_indices.size(0);
2157
+
2158
+ at::Tensor o = at::empty({B, T, Hv, Dv}, q.options());
2159
+ AT_DISPATCH_REDUCED_FLOATING_TYPES(q.scalar_type(), "chunk_gated_delta_rule_fwd_inter", [&] {
2160
+ DISPATCH_HEAD_DIM(D, LAUNCH_CHUNK_GATED_DELTA_RULE_FWD_INTER_KERNEL);
2161
+ });
2162
+
2163
+ return std::make_tuple(o, initial_state);
2164
+}
2165
+
2166
+// [NB]: Support only varlen inputs
2167
+// B: packed batch dim of q/k/v (== 1)
2168
+// num_seqs: number of variable-length sequences
2169
+//
2170
+// query: [B, T, H, D]
2171
+// key: [B, T, H, D]
2172
+// value: [B, T, Hv, Dv]
2173
+// g: [B, T, Hv] FP32
2174
+// beta: [B, T, Hv]
2175
+// initial_state: [num_seqs, Hv, Dv, D] FP32
2176
+// cu_seqlens: [num_seqs + 1] INT32
2177
+//
2178
std::tuple<at::Tensor, at::Tensor> chunk_gated_delta_rule_cpu(
2179
const at::Tensor& query,
2180
const at::Tensor& key,
@@ -1407,154 +2186,71 @@ std::tuple<at::Tensor, at::Tensor> chunk_gated_delta_rule_cpu(
2186
const at::Tensor& cu_seqlens,
2187
bool head_first,
2188
bool use_qk_l2norm_in_kernel,
1410
- double eps = 1e-5) {
1411
- TORCH_CHECK(head_first == false, "chunk_gated_delta_rule_cpu does not support head first");
1412
- int64_t B = query.size(0);
1413
- int64_t global_seq_len = query.size(1);
1414
- int64_t qk_num_head = query.size(2);
1415
- int64_t qk_head_size = query.size(3);
1416
- int64_t v_num_head = value.size(2);
1417
- int64_t v_head_size = value.size(3);
1418
- int64_t batch_size = initial_state.size(0);
1419
- CHECK_EQ(B, 1);
1420
- TORCH_CHECK(v_num_head % qk_num_head == 0, "expect v_num_head multiple of qk_num_head.");
1421
- TORCH_CHECK(qk_head_size % 32 == 0, "expect qk_head_size to be multiples of 32.");
1422
- TORCH_CHECK(v_head_size % 32 == 0, "expect v_head_size to be multiples of 32.");
1423
- CHECK_INPUT_SHAPE_DTYPE<true>(query, 4, {B, global_seq_len, qk_num_head, qk_head_size}, at::kBFloat16);
1424
- CHECK_INPUT_SHAPE_DTYPE<true>(key, 4, {B, global_seq_len, qk_num_head, qk_head_size}, at::kBFloat16);
1425
- CHECK_INPUT_SHAPE_DTYPE<true>(value, 4, {B, global_seq_len, v_num_head, v_head_size}, at::kBFloat16);
1426
- CHECK_INPUT_SHAPE_DTYPE<false>(g, 3, {B, global_seq_len, v_num_head}, at::kFloat);
1427
- CHECK_INPUT_SHAPE_DTYPE<false>(beta, 3, {B, global_seq_len, v_num_head}, at::kBFloat16);
1428
- CHECK_INPUT_SHAPE_DTYPE<false>(cu_seqlens, 1, {batch_size + 1}, at::kInt);
1429
- CHECK_INPUT_SHAPE_DTYPE<false>(initial_state, 4, {batch_size, v_num_head, qk_head_size, v_head_size}, at::kFloat);
1430
-
1431
- at::Tensor output = at::empty_like(value, value.options()); // [B, T, HV, EV]
1432
- at::Tensor final_state = initial_state.to(at::kFloat); // [N, HV, EK, EV]
1433
-
1434
- // Strides
1435
- int64_t qStrideH = query.stride(2);
1436
- int64_t qStrideT = query.stride(1);
1437
- int64_t kStrideH = key.stride(2);
1438
- int64_t kStrideT = key.stride(1);
1439
- int64_t vStrideH = value.stride(2);
1440
- int64_t vStrideT = value.stride(1);
1441
- int64_t oStrideH = output.stride(2);
1442
- int64_t oStrideT = output.stride(1);
1443
-
1444
- constexpr int64_t chunk_size = 64;
1445
- // Deduce the global chunks
1446
- // e.g. cu_seqlens: [0, 5, 13, 16], chunk_size = 4
1447
- // chunk_offsets: [0, 2, 4, 5]
1448
- // chunk_indices (batch_id, local_chunk_id): [[0, 0], [0, 1], [1, 0], [1, 1], [2, 0]]
1449
- at::Tensor chunk_offsets = at::empty(batch_size + 1, cu_seqlens.options());
1450
- auto chunk_offsets_ptr = chunk_offsets.data_ptr<int32_t>();
1451
- chunk_offsets_ptr[0] = 0;
1452
- int32_t* cu_seqlens_ptr = cu_seqlens.data_ptr<int32_t>();
1453
- int64_t s = 0;
1454
- int64_t e = 0;
1455
- int64_t s_pad = 0;
1456
- int64_t e_pad = 0;
1457
- for (int64_t b = 0; b < batch_size; b++) {
1458
- e = cu_seqlens_ptr[b + 1];
1459
- int64_t seq_len = e - s;
1460
- int64_t pad_size = (chunk_size - seq_len % chunk_size) % chunk_size;
1461
- int64_t total_seq_length = seq_len + pad_size;
1462
- e_pad = s_pad + total_seq_length;
1463
- chunk_offsets[b + 1] = e_pad / chunk_size;
1464
- s = e;
1465
- s_pad = e_pad;
1466
- }
1467
- int64_t global_total_seq_length = e_pad;
1468
- int64_t global_num_chunk = chunk_offsets_ptr[batch_size];
1469
- at::Tensor chunk_indices = at::empty(global_num_chunk * 2, cu_seqlens.options());
1470
- auto chunk_indices_ptr = chunk_indices.data_ptr<int32_t>();
1471
- int64_t curr_c = 0;
1472
- for (int64_t b = 0; b < batch_size; b++) {
1473
- int64_t batch_chunk_num = chunk_offsets_ptr[b + 1] - chunk_offsets_ptr[b];
1474
- for (int64_t c = 0; c < batch_chunk_num; c++) {
1475
- chunk_indices_ptr[curr_c * 2] = b;
1476
- chunk_indices_ptr[curr_c * 2 + 1] = c;
1477
- curr_c += 1;
1478
- }
1479
- }
2189
+ const at::Tensor& initial_state_indices,
2190
+ double eps = 1e-6) {
2191
+ TORCH_CHECK(!head_first, "chunk_gated_delta_rule_cpu: does not support head first");
2192
1481
- // Allocate buffer
1482
- int64_t buff_size = v_num_head * global_total_seq_length // g_pad_data
1483
- + batch_size * v_num_head * global_total_seq_length * v_head_size // core_attn
1484
- + v_num_head * global_total_seq_length * chunk_size // decay_mask
1485
- + v_num_head * global_total_seq_length * v_head_size; // v_beta_attn
1486
- at::Tensor buff_data = at::empty({buff_size}, query.options().dtype(at::kFloat));
1487
- int64_t reduced_buff_size = qk_num_head * global_total_seq_length * qk_head_size // q_pad_data
1488
- + qk_num_head * global_total_seq_length * qk_head_size // k_pad_data
1489
- + v_num_head * global_total_seq_length * v_head_size // v_pad_data
1490
- + v_num_head * global_total_seq_length * qk_head_size // k_beta_data
1491
- + v_num_head * global_total_seq_length * v_head_size // v_beta_data
1492
- + v_num_head * global_total_seq_length * qk_head_size // k_cumdecay_reduced
1493
- + qk_num_head * global_seq_len // q_norm_sum
1494
- + qk_num_head * global_seq_len; // k_norm_sum
1495
- at::Tensor reduced_buff_data = at::empty({reduced_buff_size}, query.options());
1496
- int64_t num_thread = at::get_num_threads();
1497
- int64_t buff_size_16bit_per_thread =
1498
- /* k_transpose */ qk_head_size * chunk_size +
1499
- /* v_pack */ chunk_size * v_head_size +
1500
- /* k_beta_g */ chunk_size * qk_head_size +
1501
- /* k_beta_g_pack */ chunk_size * qk_head_size +
1502
- /* attn */ chunk_size * chunk_size * 2 +
1503
- /* attn_reduced */ chunk_size * chunk_size +
1504
- /* k_cumdecay */ chunk_size * qk_head_size * 2 +
1505
- /* row */ chunk_size * 2 +
1506
- /* updated */ chunk_size * 2 +
1507
- /* curr_last_recurrent_state_reduced */ qk_head_size * v_head_size +
1508
- /* curr_last_recurrent_state_pack_reduced */ qk_head_size * v_head_size +
1509
- /* k_transpose_i */ qk_head_size * chunk_size +
1510
- /* attn_i */ chunk_size * chunk_size * 2 +
1511
- /* attn_i_reduced */ chunk_size * chunk_size +
1512
- /* v_prime */ chunk_size * v_head_size * 2 +
1513
- /* v_prime_reduced */ chunk_size * v_head_size +
1514
- /* v_prime_pack_reduced */ chunk_size * v_head_size +
1515
- /* qg */ chunk_size * qk_head_size +
1516
- /* attn_inter */ chunk_size * v_head_size * 2 +
1517
- /* kg */ chunk_size * qk_head_size +
1518
- /* kg_transpose */ qk_head_size * chunk_size +
1519
- /* kgv */ qk_head_size * v_head_size * 2;
1520
- at::Tensor thread_buff_data = at::empty({num_thread, buff_size_16bit_per_thread}, query.options());
1521
-
1522
- AT_DISPATCH_REDUCED_FLOATING_TYPES(query.scalar_type(), "chunk_gated_delta_rule_kernel", [&] {
1523
- chunk_gated_delta_rule_kernel_impl<scalar_t, chunk_size>(
1524
- output.data_ptr<scalar_t>(),
1525
- final_state.data_ptr<float>(),
1526
- query.data_ptr<scalar_t>(),
1527
- key.data_ptr<scalar_t>(),
1528
- value.data_ptr<scalar_t>(),
1529
- g.data_ptr<float>(),
1530
- beta.data_ptr<scalar_t>(),
1531
- cu_seqlens_ptr,
1532
- buff_data.data_ptr<float>(),
1533
- reduced_buff_data.data_ptr<scalar_t>(),
1534
- thread_buff_data.data_ptr<scalar_t>(),
1535
- chunk_offsets_ptr,
1536
- chunk_indices_ptr,
1537
- use_qk_l2norm_in_kernel,
1538
- batch_size,
1539
- global_seq_len,
1540
- qk_num_head,
1541
- v_num_head,
1542
- qk_head_size,
1543
- v_head_size,
1544
- qStrideH,
1545
- qStrideT,
1546
- kStrideH,
1547
- kStrideT,
1548
- vStrideH,
1549
- vStrideT,
1550
- oStrideH,
1551
- oStrideT,
1552
- global_total_seq_length,
1553
- global_num_chunk,
1554
- buff_size_16bit_per_thread,
1555
- eps);
1556
- });
1557
- return std::make_tuple(std::move(output), std::move(final_state));
2193
+ int64_t B = query.size(0);
2194
+ int64_t T = query.size(1);
2195
+ int64_t H = query.size(2);
2196
+ int64_t D = query.size(3);
2197
+ int64_t Hv = value.size(2);
2198
+ int64_t Dv = value.size(3);
2199
+ int64_t num_seqs = initial_state_indices.size(0);
2200
+
2201
+ TORCH_CHECK(B == 1, __func__, ": expect batch size to be 1");
2202
+ TORCH_CHECK(Hv % H == 0, __func__, ": expect num_heads_kv multiple of num_heads.");
2203
+ TORCH_CHECK(D % 32 == 0, __func__, ": expect head_dim to be multiples of 32.");
2204
+ TORCH_CHECK(Dv % 32 == 0, __func__, ": expect head_dim_v to be multiples of 32.");
2205
+ TORCH_CHECK(D == Dv, __func__, ": expect head_dim to be equal to head_dim_v.");
2206
+ CHECK_INPUT_SHAPE_DTYPE<true>(query, {B, T, H, D}, at::kBFloat16);
2207
+ CHECK_INPUT_SHAPE_DTYPE<true>(key, {B, T, H, D}, at::kBFloat16);
2208
+ CHECK_INPUT_SHAPE_DTYPE<true>(value, {B, T, Hv, Dv}, at::kBFloat16);
2209
+ CHECK_INPUT_SHAPE_DTYPE<false>(g, {B, T, Hv}, at::kFloat);
2210
+ CHECK_INPUT_SHAPE_DTYPE<false>(beta, {B, T, Hv}, at::kBFloat16);
2211
+ CHECK_INPUT_SHAPE_DTYPE<false>(cu_seqlens, {num_seqs + 1}, at::kInt);
2212
+ TORCH_CHECK(initial_state.sizes() == at::IntArrayRef({initial_state.size(0), Hv, Dv, D}),
2213
+ "chunk_gated_delta_rule_cpu: initial_state shape mismatch, got ", initial_state.sizes());
2214
+ TORCH_CHECK(initial_state.scalar_type() == at::kFloat, "chunk_gated_delta_rule_cpu: initial_state dtype mismatch");
2215
+ CHECK_CPU(initial_state);
2216
+ // initial_state may be a pooled/paged buffer with padding between slots
2217
+ // (e.g. mamba cache-align mode), so only the per-slot (Hv, Dv, D) layout
2218
+ // needs to be densely packed; dim 0's stride is read at runtime instead of
2219
+ // assumed, mirroring conv.cpp's conv_state_slot_stride handling.
2220
+ TORCH_CHECK(initial_state.stride(-1) == 1 && initial_state.stride(-2) == D &&
2221
+ initial_state.stride(-3) == Dv * D,
2222
+ "chunk_gated_delta_rule_cpu: expect initial_state to be contiguous per pool slot.");
2223
+ CHECK_INPUT_SHAPE_DTYPE<false>(initial_state_indices, {num_seqs}, at::kInt);
2224
+
2225
+ constexpr int CHUNK_SIZE = 64;
2226
+
2227
+ // prepare chunk indices
2228
+ auto [chunk_indices, chunk_offsets] = prepare_chunk_indices<CHUNK_SIZE>(cu_seqlens);
2229
+
2230
+ float scale = 1.0 / std::sqrt(D);
2231
+ auto [query_, key_] = use_qk_l2norm_in_kernel ? l2norm_fwd(query, key, eps) : std::make_tuple(query.mul(scale), key);
2232
+
2233
+ auto g_ = chunk_local_cumsum<CHUNK_SIZE>(g, cu_seqlens, chunk_indices);
2234
+
2235
+ // fused kkt + solve_tril + recompute_w_u
2236
+ auto [w, u, decay_mask] =
2237
+ chunk_gated_delta_rule_fwd_intra<CHUNK_SIZE>(key_, value, g_, beta, cu_seqlens, chunk_indices);
2238
+
2239
+ // fused `chunk_gated_delta_rule_fwd_h` + `chunk_fwd_o`
2240
+ auto [output, final_state] = chunk_gated_delta_rule_fwd_inter<CHUNK_SIZE>(
2241
+ query_,
2242
+ key_,
2243
+ w,
2244
+ u,
2245
+ g_,
2246
+ decay_mask,
2247
+ initial_state,
2248
+ output_final_state,
2249
+ cu_seqlens,
2250
+ chunk_offsets,
2251
+ initial_state_indices);
2252
+
2253
+ return std::make_tuple(output, final_state);
2254
}
2255
2256
// A_log: [v_num_heads]
csrc/cpu/sgl-kernels/gemm.cpp
+22
-28
@@ -116,8 +116,7 @@ inline void copy_stub(scalar_t* __restrict__ out, const float* __restrict__ inpu
116
int64_t d;
117
#pragma GCC unroll 4
118
for (d = 0; d <= size - kVecSize; d += kVecSize) {
119
- fVec data0 = fVec::loadu(input + d);
120
- fVec data1 = fVec::loadu(input + d + fVec::size());
119
+ auto [data0, data1] = load_float_vec2(input + d);
120
bVec out_vec = convert_from_float_ext<scalar_t>(data0, data1);
121
out_vec.store(out + d);
122
}
@@ -135,9 +134,7 @@ inline void copy_stub(float* __restrict__ out, const scalar_t* __restrict__ inpu
134
int64_t d;
135
#pragma GCC unroll 4
136
for (d = 0; d <= size - kVecSize; d += kVecSize) {
138
- fVec data0, data1;
139
- bVec b_vec = bVec::loadu(input + d);
140
- std::tie(data0, data1) = at::vec::convert_to_float(b_vec);
137
+ auto [data0, data1] = load_float_vec2(input + d);
138
data0.store(out + d);
139
data1.store(out + d + fVec::size());
140
}
@@ -156,9 +153,9 @@ inline void copy_add_stub(
153
int64_t d;
154
#pragma GCC unroll 4
155
for (d = 0; d <= size - kVecSize; d += kVecSize) {
159
- fVec data0 = fVec::loadu(input + d) + fVec::loadu(bias + d);
160
- fVec data1 = fVec::loadu(input + d + fVec::size()) + fVec::loadu(bias + d + fVec::size());
161
- bVec out_vec = convert_from_float_ext<scalar_t>(data0, data1);
156
+ auto [data0, data1] = load_float_vec2(input + d);
157
+ auto [bias0, bias1] = load_float_vec2(bias + d);
158
+ bVec out_vec = convert_from_float_ext<scalar_t>(data0 + bias0, data1 + bias1);
159
out_vec.store(out + d);
160
}
161
for (; d < size; ++d) {
@@ -176,7 +173,6 @@ inline void scalar_sigmoid_and_mul(
173
using bVec = at::vec::Vectorized<scalar_t>;
174
using fVec = at::vec::Vectorized<float>;
175
// scalar sigmoid
179
- const fVec one = fVec(1.f);
176
fVec X;
177
if constexpr (has_bias) {
178
assert(bias != nullptr);
@@ -184,18 +180,13 @@ inline void scalar_sigmoid_and_mul(
180
} else {
181
X = fVec(input[0]);
182
}
187
- X = one / (one + X.neg().exp_u20());
183
+ X = fast_sigmoid(X);
184
185
// vec mul
186
constexpr int kVecSize = bVec::size();
187
for (int d = 0; d < SIZE; d += kVecSize) {
192
- bVec m_bvec = bVec::loadu(mul + d);
193
- fVec m_fvec0, m_fvec1;
194
- std::tie(m_fvec0, m_fvec1) = at::vec::convert_to_float(m_bvec);
195
- m_fvec0 = m_fvec0 * X;
196
- m_fvec1 = m_fvec1 * X;
197
-
198
- bVec out_vec = convert_from_float_ext<scalar_t>(m_fvec0, m_fvec1);
188
+ auto [m_fvec0, m_fvec1] = load_float_vec2(mul + d);
189
+ bVec out_vec = convert_from_float_ext<scalar_t>(m_fvec0 * X, m_fvec1 * X);
190
out_vec.store(out + d);
191
}
192
}
@@ -727,10 +718,10 @@ at::Tensor convert_scale_packed(at::Tensor& scale) {
718
return packed_scale;
719
}
720
730
-// mat1 : [M, K]
721
+// mat1 : [*, K]
722
// mat2 : [N, K] ([K, N] if use_fma_gemm)
723
// bias : [N]
733
-// out : [M, N]
724
+// out : [*, N]
725
//
726
at::Tensor
727
weight_packed_linear(at::Tensor& mat1, at::Tensor& mat2, const std::optional<at::Tensor>& bias, bool is_vnni) {
@@ -740,23 +731,25 @@ weight_packed_linear(at::Tensor& mat1, at::Tensor& mat2, const std::optional<at:
731
use_fma_gemm = true;
732
}
733
743
- int64_t M = mat1.size(0);
744
- int64_t K = mat1.size(1);
745
- int64_t N = use_fma_gemm ? mat2.size(1) : mat2.size(0);
746
-
734
CHECK_LAST_DIM_CONTIGUOUS_INPUT(mat1);
735
CHECK_INPUT(mat2);
749
- CHECK_DIM(2, mat1);
736
+ const int64_t ndim = mat1.ndimension();
737
+ auto input_sizes = mat1.sizes().vec();
738
+ int64_t N = use_fma_gemm ? mat2.size(1) : mat2.size(0);
739
+ int64_t K = use_fma_gemm ? mat1.size(1) : mat2.size(1);
740
+ int64_t M = use_fma_gemm ? mat1.size(0) : mat1.numel() / K;
741
CHECK_DIM(2, mat2);
751
- if (!use_fma_gemm) {
752
- CHECK_EQ(mat1.size(1), K);
742
+ if (use_fma_gemm) {
743
+ CHECK_DIM(2, mat1);
744
+ } else {
745
+ CHECK_EQ(mat1.size(ndim - 1), K);
746
}
747
748
auto dispatch_type = mat1.scalar_type();
749
auto out = at::empty({M, N}, mat1.options());
750
// strides
751
int64_t out_strideM = out.stride(0);
759
- int64_t mat1_strideM = mat1.stride(0);
752
+ int64_t mat1_strideM = mat1.stride(-2);
753
754
const bool has_bias = bias.has_value();
755
const float* bias_data = nullptr;
@@ -792,7 +785,8 @@ weight_packed_linear(at::Tensor& mat1, at::Tensor& mat2, const std::optional<at:
785
}
786
});
787
795
- return out;
788
+ input_sizes[ndim - 1] = N;
789
+ return out.view(input_sizes);
790
}
791
792
// mat1 : [M, K]
csrc/cpu/sgl-kernels/gemm.h
+17
-10
@@ -83,15 +83,11 @@ inline int64_t get_row_size(int64_t K, bool use_int8_w8a8) {
83
return use_int8_w8a8 ? K + sizeof(int32_t) : K;
84
}
85
86
-enum class CPUAcTMethod : int { silu_and_mul = 0, swiglu = 1 };
87
-
88
-constexpr bool operator==(CPUAcTMethod a, int b) {
89
- return static_cast<int>(a) == b;
90
-}
91
-
92
-constexpr bool operator==(int a, CPUAcTMethod b) {
93
- return a == static_cast<int>(b);
94
-}
86
+enum class CPUActMethod : int {
87
+ silu_and_mul = 0,
88
+ swiglu = 1,
89
+ gelu_and_mul = 2,
90
+};
91
92
enum class CPUQuantMethod : int64_t { BF16 = 0, INT8_W8A8 = 1, FP8_W8A16 = 2, INT4_W4A8 = 3, MXFP4 = 4 };
93
@@ -113,6 +109,17 @@ constexpr bool operator==(int64_t a, CPUQuantAlgo b) {
109
return a == static_cast<int64_t>(b);
110
}
111
112
+inline int64_t get_row_size(CPUQuantMethod quant, int64_t K) {
113
+ switch (quant) {
114
+ case CPUQuantMethod::INT8_W8A8:
115
+ return K + sizeof(int32_t);
116
+ case CPUQuantMethod::MXFP4:
117
+ return K >> 1;
118
+ default:
119
+ return K;
120
+ }
121
+}
122
+
123
inline int64_t get_4bit_block_k_size(int64_t group_size) {
124
return group_size > 128 ? 128 : group_size;
125
}
@@ -184,7 +191,7 @@ void fused_experts_fp_kernel_impl(
191
int64_t num_tokens_post_pad,
192
float alpha,
193
float limit,
187
- CPUAcTMethod act_func,
194
+ CPUActMethod act_func,
195
bool with_bias);
196
197
// shared expert implementation for int8 w8a8
csrc/cpu/sgl-kernels/gemm_fp8.cpp
+6
-8
@@ -18,8 +18,7 @@ inline void copy_stub(scalar_t* __restrict__ out, const float* __restrict__ inpu
18
int64_t d;
19
#pragma GCC unroll 4
20
for (d = 0; d <= size - kVecSize; d += kVecSize) {
21
- fVec data0 = fVec::loadu(input + d);
22
- fVec data1 = fVec::loadu(input + d + fVec::size());
21
+ auto [data0, data1] = load_float_vec2(input + d);
22
bVec out_vec = convert_from_float_ext<scalar_t>(data0, data1);
23
out_vec.store(out + d);
24
}
@@ -38,9 +37,9 @@ inline void copy_add_stub(
37
int64_t d;
38
#pragma GCC unroll 4
39
for (d = 0; d <= size - kVecSize; d += kVecSize) {
41
- fVec data0 = fVec::loadu(input + d) + fVec::loadu(bias + d);
42
- fVec data1 = fVec::loadu(input + d + fVec::size()) + fVec::loadu(bias + d + fVec::size());
43
- bVec out_vec = convert_from_float_ext<scalar_t>(data0, data1);
40
+ auto [data0, data1] = load_float_vec2(input + d);
41
+ auto [bias0, bias1] = load_float_vec2(bias + d);
42
+ bVec out_vec = convert_from_float_ext<scalar_t>(data0 + bias0, data1 + bias1);
43
out_vec.store(out + d);
44
}
45
for (; d < size; ++d) {
@@ -57,9 +56,8 @@ inline void copy_mul_stub(scalar_t* __restrict__ out, const float* __restrict__
56
int d;
57
#pragma GCC unroll 4
58
for (d = 0; d <= size - kVecSize; d += kVecSize) {
60
- fVec data0 = fVec::loadu(input + d) * vscale;
61
- fVec data1 = fVec::loadu(input + d + fVec::size()) * vscale;
62
- bVec out_vec = convert_from_float_ext<scalar_t>(data0, data1);
59
+ auto [data0, data1] = load_float_vec2(input + d);
60
+ bVec out_vec = convert_from_float_ext<scalar_t>(data0 * vscale, data1 * vscale);
61
out_vec.store(out + d);
62
}
63
for (; d < size; ++d) {
csrc/cpu/sgl-kernels/moe.cpp
+124
-100
@@ -123,46 +123,6 @@ int moe_align_block_size(
123
return num_tokens_post_pad;
124
}
125
126
-// silu : shape leading dimension
127
-// input0 [m_size, BLOCK_N] BLOCK_N
128
-// input1 [m_size, BLOCK_N] BLOCK_N
129
-// output [M * topk, N] N
130
-template <typename scalar_t, int BLOCK_N>
131
-inline void silu_and_mul(
132
- scalar_t* __restrict__ output,
133
- const float* __restrict__ input0, // x: x0, x1
134
- const float* __restrict__ input1, // y: y0, y1
135
- int64_t m_size,
136
- int64_t N) {
137
- using bVec = at::vec::Vectorized<scalar_t>;
138
- using fVec = at::vec::Vectorized<float>;
139
-
140
- const fVec one = fVec(1.f);
141
-
142
- // no remainder
143
- for (int64_t m = 0; m < m_size; ++m) {
144
- scalar_t* __restrict__ out = output + m * N;
145
- const float* __restrict__ x = input0 + m * BLOCK_N;
146
- const float* __restrict__ y = input1 + m * BLOCK_N;
147
-
148
- for (int64_t d = 0; d < BLOCK_N; d += bVec::size()) {
149
- fVec x0 = fVec::loadu(x + d);
150
- fVec x1 = fVec::loadu(x + d + fVec::size());
151
- fVec y0 = fVec::loadu(y + d);
152
- fVec y1 = fVec::loadu(y + d + fVec::size());
153
- // silu
154
- x0 = x0 / (one + x0.neg().exp_u20());
155
- x1 = x1 / (one + x1.neg().exp_u20());
156
- // mul
157
- x0 = x0 * y0;
158
- x1 = x1 * y1;
159
- // convert
160
- bVec out_vec = convert_from_float_ext<scalar_t>(x0, x1);
161
- out_vec.store(out + d);
162
- }
163
- }
164
-}
165
-
126
template <typename scalar_t, int BLOCK_M, int BLOCK_N>
127
struct tinygemm_kernel_nn2 {
128
static inline void apply(
@@ -239,23 +199,17 @@ struct tinygemm_kernel_nn2<at::BFloat16, BLOCK_M, BLOCK_N> {
199
Unroll<ROWS * COLS>{}(compute, k);
200
}
201
242
- using Vec = at::vec::Vectorized<float>;
243
- const Vec one = Vec(1.f);
202
auto storec = [&](auto i) {
203
constexpr int row = i / COLS;
204
constexpr int col = i % COLS;
205
// for COLS = 2, 4 use 512bit store
206
if constexpr (col % 2 == 0) {
249
- Vec x0 = vc0[row * COLS + col + 0];
250
- Vec x1 = vc0[row * COLS + col + 1];
251
- Vec y0 = vc1[row * COLS + col + 0];
252
- Vec y1 = vc1[row * COLS + col + 1];
253
- // silu
254
- x0 = x0 / (one + x0.neg().exp_u20());
255
- x1 = x1 / (one + x1.neg().exp_u20());
256
- // mul
257
- x0 = x0 * y0;
258
- x1 = x1 * y1;
207
+ __m512 x0 = vc0[row * COLS + col + 0];
208
+ __m512 x1 = vc0[row * COLS + col + 1];
209
+ __m512 y0 = vc1[row * COLS + col + 0];
210
+ __m512 y1 = vc1[row * COLS + col + 1];
211
+ x0 = _mm512_mul_ps(_mm512_rcp14_silu_ps(x0), y0);
212
+ x1 = _mm512_mul_ps(_mm512_rcp14_silu_ps(x1), y1);
213
214
_mm512_storeu_si512(
215
reinterpret_cast<__m512i*>((C + row * ldc + col * 16)),
@@ -455,6 +409,8 @@ void fused_experts_kernel_impl(
409
const scalar_t* __restrict__ input,
410
const scalar_t* __restrict__ packed_w1,
411
const scalar_t* __restrict__ packed_w2,
412
+ const float* __restrict__ w1_bias,
413
+ const float* __restrict__ w2_bias,
414
const float* __restrict__ topk_weights,
415
const int32_t* __restrict__ sorted_ids,
416
const int32_t* __restrict__ expert_ids,
@@ -464,7 +420,11 @@ void fused_experts_kernel_impl(
420
int64_t K,
421
int64_t E,
422
int64_t topk,
467
- int64_t num_tokens_post_pad) {
423
+ int64_t num_tokens_post_pad,
424
+ float alpha,
425
+ float limit,
426
+ CPUActMethod act_func,
427
+ bool with_bias) {
428
// handle 2 tiles per block
429
constexpr int64_t BLOCK_M = block_size_m();
430
constexpr int64_t BLOCK_N = block_size_n();
@@ -499,6 +459,8 @@ void fused_experts_kernel_impl(
459
int32_t expert_id = expert_ids[mb];
460
const scalar_t* __restrict__ B0 = packed_w1 + expert_id * stride_e + nb_upper * BLOCK_N * stride_n;
461
const scalar_t* __restrict__ B1 = packed_w1 + expert_id * stride_e + nb_lower * BLOCK_N * stride_n;
462
+ const float* __restrict__ B0_bias = w1_bias + expert_id * 2 * N + nb_upper * BLOCK_N;
463
+ const float* __restrict__ B1_bias = w1_bias + expert_id * 2 * N + nb_lower * BLOCK_N;
464
465
int64_t m_size = offsets[mb + 1] - offsets[mb];
466
@@ -538,23 +500,62 @@ void fused_experts_kernel_impl(
500
/* B */ B1,
501
/* C */ C1);
502
541
- // 1.d silu and mul
542
- const int64_t offset = offsets[mb];
543
- silu_and_mul<scalar_t, BLOCK_N>(ic1 + offset * N + nb * BLOCK_N, C0, C1, m_size, N);
503
} else {
545
- // fused 1.bcd: silu_and_mul(A @ B0, A @ B1)
504
const int64_t offset = offsets[mb];
547
- tinygemm_kernel(
548
- /* A */ A,
549
- /* B0 */ B0,
550
- /* B1 */ B1,
551
- /* C */ ic1 + offset * N + nb * BLOCK_N,
552
- /* M */ m_size,
553
- /* N */ n_size,
554
- /* K */ K,
555
- /* lda */ K,
556
- /* ldb */ n_size,
557
- /* ldc */ N);
505
+ if (act_func == CPUActMethod::swiglu) {
506
+ tinygemm_kernel(
507
+ /* A */ A,
508
+ /* B */ B0,
509
+ /* C */ C0,
510
+ /* M */ m_size,
511
+ /* N */ n_size,
512
+ /* K */ K,
513
+ /* lda */ K,
514
+ /* ldb */ n_size,
515
+ /* ldc */ BLOCK_N);
516
+ tinygemm_kernel(
517
+ /* A */ A,
518
+ /* B */ B1,
519
+ /* C */ C1,
520
+ /* M */ m_size,
521
+ /* N */ n_size,
522
+ /* K */ K,
523
+ /* lda */ K,
524
+ /* ldb */ n_size,
525
+ /* ldc */ BLOCK_N);
526
+ } else {
527
+ // fused 1.bcd: silu_and_mul(A @ B0, A @ B1)
528
+ tinygemm_kernel(
529
+ /* A */ A,
530
+ /* B0 */ B0,
531
+ /* B1 */ B1,
532
+ /* C */ ic1 + offset * N + nb * BLOCK_N,
533
+ /* M */ m_size,
534
+ /* N */ n_size,
535
+ /* K */ K,
536
+ /* lda */ K,
537
+ /* ldb */ n_size,
538
+ /* ldc */ N);
539
+ }
540
+ }
541
+ if (with_bias) {
542
+ for (int64_t m = 0; m < m_size; ++m) {
543
+ add_bias_stub(C0 + m * BLOCK_N, B0_bias, n_size);
544
+ add_bias_stub(C1 + m * BLOCK_N, B1_bias, n_size);
545
+ }
546
+ }
547
+ // 1.d silu and mul
548
+ const int64_t offset = offsets[mb];
549
+ if (act_func == CPUActMethod::silu_and_mul && use_brgemm) {
550
+ for (int64_t m = 0; m < m_size; ++m) {
551
+ silu_and_mul_stub(ic1 + (offset + m) * N + nb * BLOCK_N, C0 + m * BLOCK_N, C1 + m * BLOCK_N, BLOCK_N);
552
+ }
553
+ } else if (act_func == CPUActMethod::swiglu) {
554
+ for (int64_t m = 0; m < m_size; ++m) {
555
+ scalar_t* __restrict__ ic1_row = ic1 + (offset + m) * N;
556
+ clamp_sigmoid_and_mul_stub(ic1_row + nb * BLOCK_N / 2, C0 + m * BLOCK_N, BLOCK_N / 2, alpha, limit);
557
+ clamp_sigmoid_and_mul_stub(ic1_row + N / 2 + nb * BLOCK_N / 2, C1 + m * BLOCK_N, BLOCK_N / 2, alpha, limit);
558
+ }
559
}
560
});
561
@@ -591,6 +592,7 @@ void fused_experts_kernel_impl(
592
// B shape [IC, n_size] in vnni format
593
int32_t expert_id = expert_ids[mb];
594
const scalar_t* __restrict__ B = packed_w2 + expert_id * stride_e2 + nb * BLOCK_N * stride_oc;
595
+ const float* __restrict__ B_bias = w2_bias + expert_id * OC + nb * BLOCK_N;
596
597
// 2.a gemm: C = A @ B
598
if (use_brgemm) {
@@ -618,6 +620,11 @@ void fused_experts_kernel_impl(
620
/* ldc */ BLOCK_N);
621
}
622
623
+ if (with_bias) {
624
+ for (int64_t m = 0; m < m_size; ++m) {
625
+ add_bias_stub(C + m * BLOCK_N, B_bias, n_size);
626
+ }
627
+ }
628
// 2.b copy from C to ic2 in original order
629
// and also mul topk_weights in float32
630
for (int64_t m = 0; m < m_size; ++m) {
@@ -717,7 +724,9 @@ void shared_expert_kernel_impl(
724
/* C */ C1);
725
726
// 1.d silu and mul
720
- silu_and_mul<scalar_t, BLOCK_N>(ic1 + mb * BLOCK_M * N + nb * BLOCK_N, C0, C1, m_size, N);
727
+ for (int64_t m = 0; m < m_size; ++m) {
728
+ silu_and_mul_stub(ic1 + (mb * BLOCK_M + m) * N + nb * BLOCK_N, C0 + m * BLOCK_N, C1 + m * BLOCK_N, BLOCK_N);
729
+ }
730
} else {
731
// fused 1.bcd: silu_and_mul(A @ B0, A @ B1)
732
tinygemm_kernel(
@@ -809,24 +818,20 @@ void shared_expert_kernel_impl(
818
} // anonymous namespace
819
820
// common checks
821
+template <CPUQuantMethod quant>
822
static inline void check_moe_scales(
813
- bool use_int8_w8a8,
814
- bool use_fp8_w8a16,
815
- bool use_mxfp4,
823
const std::optional<at::Tensor>& w1_scale,
824
const std::optional<at::Tensor>& w2_scale,
825
const std::optional<std::vector<int64_t>> block_size) {
819
- if (use_int8_w8a8) {
826
+ if constexpr (quant == CPUQuantMethod::INT8_W8A8) {
827
TORCH_CHECK(w1_scale.has_value(), "missing w1_scale for int8 w8a8.");
828
TORCH_CHECK(w2_scale.has_value(), "missing w2_scale for int8 w8a8.");
822
- }
823
- if (use_fp8_w8a16) {
829
+ } else if constexpr (quant == CPUQuantMethod::FP8_W8A16) {
830
TORCH_CHECK(w1_scale.has_value(), "missing w1_scale for fp8 w8a16.");
831
TORCH_CHECK(w2_scale.has_value(), "missing w2_scale for fp8 w8a16.");
832
TORCH_CHECK(block_size.has_value(), "missing block_size for fp8 w8a16.");
833
TORCH_CHECK(block_size.value().size() == 2, "expect block_size.size() to be 2.");
828
- }
829
- if (use_mxfp4) {
834
+ } else if constexpr (quant == CPUQuantMethod::MXFP4) {
835
TORCH_CHECK(w1_scale.has_value(), "missing w1_scale for mxfp4.");
836
TORCH_CHECK(w2_scale.has_value(), "missing w2_scale for mxfp4.");
837
TORCH_CHECK(w1_scale.value().scalar_type() == at::kByte, "expect w1_scale to be uint8.");
@@ -834,6 +839,20 @@ static inline void check_moe_scales(
839
}
840
}
841
842
+static inline void check_moe_scales(
843
+ int64_t moe_comp_method,
844
+ const std::optional<at::Tensor>& w1_scale,
845
+ const std::optional<at::Tensor>& w2_scale,
846
+ const std::optional<std::vector<int64_t>> block_size) {
847
+ if (moe_comp_method == CPUQuantMethod::INT8_W8A8) {
848
+ check_moe_scales<CPUQuantMethod::INT8_W8A8>(w1_scale, w2_scale, block_size);
849
+ } else if (moe_comp_method == CPUQuantMethod::FP8_W8A16) {
850
+ check_moe_scales<CPUQuantMethod::FP8_W8A16>(w1_scale, w2_scale, block_size);
851
+ } else if (moe_comp_method == CPUQuantMethod::MXFP4) {
852
+ check_moe_scales<CPUQuantMethod::MXFP4>(w1_scale, w2_scale, block_size);
853
+ }
854
+}
855
+
856
#define CHECK_MOE_SCALES_FP8(DIM0, DIM1) \
857
auto w1s = w1_scale.value(); \
858
auto w2s = w2_scale.value(); \
@@ -877,10 +896,14 @@ at::Tensor fused_experts_cpu(
896
constexpr int64_t BLOCK_N = block_size_n();
897
898
const auto st = hidden_states.scalar_type();
899
+ // TODO: fused_topk_torch_native (CPU fallback for models like MiniMax)
900
+ // returns int64 topk_ids; fused_experts_cpu requires int32. Remove the typecast after topk kernel is provided
901
+ auto topk_ids_ = topk_ids.scalar_type() == at::kInt ? topk_ids : topk_ids.to(at::kInt);
902
+
903
CHECK_INPUT(hidden_states);
904
CHECK_INPUT(w1);
905
CHECK_INPUT(w2);
883
- CHECK_EQ(topk_weights.sizes(), topk_ids.sizes());
906
+ CHECK_EQ(topk_weights.sizes(), topk_ids_.sizes());
907
CHECK_DIM(2, hidden_states);
908
if (moe_comp_method == CPUQuantMethod::INT4_W4A8 && is_vnni) {
909
CHECK_DIM(4, w1);
@@ -890,9 +913,8 @@ at::Tensor fused_experts_cpu(
913
CHECK_DIM(3, w2);
914
}
915
CHECK_DIM(2, topk_weights);
893
- CHECK_DIM(2, topk_ids);
894
-
895
- CHECK_EQ(topk_ids.scalar_type(), at::kInt);
916
+ CHECK_DIM(2, topk_ids_);
917
+ CHECK_EQ(topk_ids_.scalar_type(), at::kInt);
918
919
// TODO: support topk_weights to be bf16 or fp16 in the kernel.
920
// The topk_weights of llama4 is computed via Llama4MoE:custom_routing_function and is bf16/fp16
@@ -908,12 +930,8 @@ at::Tensor fused_experts_cpu(
930
int64_t topk = topk_weights_.size(1);
931
932
// we use int32_t compensation for int8 w8a8
911
- int64_t packed_K = moe_comp_method == CPUQuantMethod::MXFP4
912
- ? get_row_size<uint8_t>(K)
913
- : get_row_size(K, moe_comp_method == CPUQuantMethod::INT8_W8A8);
914
- int64_t packed_N = moe_comp_method == CPUQuantMethod::MXFP4
915
- ? get_row_size<uint8_t>(N)
916
- : get_row_size(N, moe_comp_method == CPUQuantMethod::INT8_W8A8);
933
+ int64_t packed_K = get_row_size(static_cast<CPUQuantMethod>(moe_comp_method), K);
934
+ int64_t packed_N = get_row_size(static_cast<CPUQuantMethod>(moe_comp_method), N);
935
936
// check weight shapes
937
CHECK_EQ(w2.size(0), E);
@@ -923,13 +941,7 @@ at::Tensor fused_experts_cpu(
941
CHECK_EQ(packed_w2.size(2), packed_N / (moe_comp_method == CPUQuantMethod::INT4_W4A8 ? 2 : 1));
942
}
943
// check scales
926
- check_moe_scales(
927
- moe_comp_method == CPUQuantMethod::INT8_W8A8,
928
- moe_comp_method == CPUQuantMethod::FP8_W8A16,
929
- moe_comp_method == CPUQuantMethod::MXFP4,
930
- w1_scale,
931
- w2_scale,
932
- block_size);
944
+ check_moe_scales(moe_comp_method, w1_scale, w2_scale, block_size);
945
946
at::Tensor out_hidden_states = inplace ? hidden_states : at::empty_like(hidden_states);
947
@@ -945,7 +957,7 @@ at::Tensor fused_experts_cpu(
957
int64_t max_num_blocks = div_up(max_num_tokens_padded, BLOCK_M);
958
auto buffer = at::empty(
959
{max_num_tokens_padded + max_num_blocks + (num_threads + 1) * E + (E + 1) + (max_num_blocks + 1)},
948
- topk_ids.options());
960
+ topk_ids_.options());
961
962
int32_t* __restrict__ sorted_ids = buffer.data_ptr<int32_t>();
963
int32_t* __restrict__ expert_ids = sorted_ids + max_num_tokens_padded;
@@ -969,7 +981,7 @@ at::Tensor fused_experts_cpu(
981
982
// align experts index
983
int64_t num_tokens_post_pad = moe_align_block_size<BLOCK_M>(
972
- sorted_ids, expert_ids, topk_ids.data_ptr<int32_t>(), total_cnts, cumsums, offsets, E, numel, num_threads);
984
+ sorted_ids, expert_ids, topk_ids_.data_ptr<int32_t>(), total_cnts, cumsums, offsets, E, numel, num_threads);
985
986
// unlike triton kernel, we fuse silu with gemm1 so only need 2 intermediate_caches:
987
// 1. intermediate_cache1 : [M * topk, N]
@@ -1048,7 +1060,7 @@ at::Tensor fused_experts_cpu(
1060
scalar_t* __restrict__ intermediate_cache0 = (scalar_t*)((void*)(C_tmp + num_threads * 2 * BLOCK_M * BLOCK_N));
1061
scalar_t* __restrict__ B_tmp = (scalar_t*)((void*)(intermediate_cache0 + M * topk * 2 * N));
1062
bool with_bias = w1_bias.has_value();
1051
- auto act_func = alpha.has_value() && limit.has_value() ? CPUAcTMethod::swiglu : CPUAcTMethod::silu_and_mul;
1063
+ auto act_func = alpha.has_value() && limit.has_value() ? CPUActMethod::swiglu : CPUActMethod::silu_and_mul;
1064
1065
CHECK_MOE_SCALES_FP8(1, 2);
1066
fused_experts_fp_kernel_impl<scalar_t, at::Float8_e4m3fn, float, false>(
@@ -1088,7 +1100,7 @@ at::Tensor fused_experts_cpu(
1100
scalar_t* __restrict__ intermediate_cache0 = (scalar_t*)((void*)(C_tmp + num_threads * 2 * BLOCK_M * BLOCK_N));
1101
scalar_t* __restrict__ B_tmp = (scalar_t*)((void*)(intermediate_cache0 + M * topk * 2 * N));
1102
bool with_bias = w1_bias.has_value();
1091
- auto act_func = alpha.has_value() && limit.has_value() ? CPUAcTMethod::swiglu : CPUAcTMethod::silu_and_mul;
1103
+ auto act_func = alpha.has_value() && limit.has_value() ? CPUActMethod::swiglu : CPUActMethod::silu_and_mul;
1104
1105
// mxfp4 supports only group size of 32 (2^5)
1106
constexpr int64_t group_size = 32;
@@ -1172,6 +1184,8 @@ at::Tensor fused_experts_cpu(
1184
} else {
1185
scalar_t* __restrict__ A_tmp = intermediate_cache2 + M * topk * K;
1186
float* __restrict__ C_tmp = (float*)((void*)(A_tmp + num_threads * BLOCK_M * K));
1187
+ bool with_bias = w1_bias.has_value();
1188
+ auto act_func = alpha.has_value() && limit.has_value() ? CPUActMethod::swiglu : CPUActMethod::silu_and_mul;
1189
1190
fused_experts_kernel_impl<scalar_t>(
1191
out_hidden_states.data_ptr<scalar_t>(),
@@ -1182,6 +1196,8 @@ at::Tensor fused_experts_cpu(
1196
hidden_states.data_ptr<scalar_t>(),
1197
packed_w1.data_ptr<scalar_t>(),
1198
packed_w2.data_ptr<scalar_t>(),
1199
+ with_bias ? w1_bias.value().data_ptr<float>() : nullptr,
1200
+ with_bias ? w2_bias.value().data_ptr<float>() : nullptr,
1201
topk_weights_.data_ptr<float>(),
1202
sorted_ids,
1203
expert_ids,
@@ -1191,7 +1207,11 @@ at::Tensor fused_experts_cpu(
1207
K,
1208
E,
1209
topk,
1194
- num_tokens_post_pad);
1210
+ num_tokens_post_pad,
1211
+ alpha.has_value() ? float(alpha.value()) : 0,
1212
+ limit.has_value() ? float(limit.value()) : 0,
1213
+ act_func,
1214
+ with_bias);
1215
}
1216
});
1217
return out_hidden_states;
@@ -1254,7 +1274,11 @@ at::Tensor shared_expert_cpu(
1274
CHECK_EQ(packed_w2.size(1), packed_N);
1275
1276
// check scales
1257
- check_moe_scales(use_int8_w8a8, use_fp8_w8a16, false, w1_scale, w2_scale, block_size);
1277
+ if (use_int8_w8a8) {
1278
+ check_moe_scales<CPUQuantMethod::INT8_W8A8>(w1_scale, w2_scale, block_size);
1279
+ } else if (use_fp8_w8a16) {
1280
+ check_moe_scales<CPUQuantMethod::FP8_W8A16>(w1_scale, w2_scale, block_size);
1281
+ }
1282
1283
at::Tensor out_hidden_states = inplace ? hidden_states : at::empty_like(hidden_states);
1284
csrc/cpu/sgl-kernels/moe.h
+41
-82
@@ -64,9 +64,7 @@ inline void copy_mul_stub(scalar_t* __restrict__ out, const input_t* __restrict_
64
#pragma GCC unroll 4
65
for (d = 0; d <= size - kVecSize; d += kVecSize) {
66
auto [x0, x1] = load_float_vec2(input + d);
67
- x0 = x0 * weight_vec;
68
- x1 = x1 * weight_vec;
69
- bVec out_vec = convert_from_float_ext<scalar_t>(x0, x1);
67
+ bVec out_vec = convert_from_float_ext<scalar_t>(x0 * weight_vec, x1 * weight_vec);
68
out_vec.store(out + d);
69
}
70
for (; d < size; ++d) {
@@ -91,10 +89,7 @@ inline void sum_stub(scalar_t* __restrict__ out, const scalar_t* __restrict__ in
89
fVec sum_fvec0 = fVec(0.f);
90
fVec sum_fvec1 = fVec(0.f);
91
for (int t = 0; t < topk; ++t) {
94
- bVec x_bvec = bVec::loadu(input + t * K + d);
95
- fVec x_fvec0, x_fvec1;
96
- std::tie(x_fvec0, x_fvec1) = at::vec::convert_to_float(x_bvec);
97
-
92
+ auto [x_fvec0, x_fvec1] = load_float_vec2(input + t * K + d);
93
sum_fvec0 += x_fvec0;
94
sum_fvec1 += x_fvec1;
95
}
@@ -137,11 +132,7 @@ inline void add_mul_stub(
132
#pragma GCC unroll 4
133
for (d = 0; d <= size - kVecSize; d += kVecSize) {
134
auto [x0, x1] = load_float_vec2(input + d);
140
-
141
- bVec y_bvec = bVec::loadu(input2 + d);
142
- fVec y0, y1;
143
- std::tie(y0, y1) = at::vec::convert_to_float(y_bvec);
144
-
135
+ auto [y0, y1] = load_float_vec2(input2 + d);
136
x0 = x0 + y0 * s_vec;
137
x1 = x1 + y1 * s_vec;
138
bVec out_vec = convert_from_float_ext<scalar_t>(x0, x1);
@@ -152,31 +143,52 @@ inline void add_mul_stub(
143
}
144
}
145
155
-template <typename scalar_t>
146
+template <typename scalar_t, typename input_t>
147
inline void silu_and_mul_stub(
157
- scalar_t* __restrict__ out, const scalar_t* __restrict__ input, const scalar_t* __restrict__ input2, int64_t size) {
148
+ scalar_t* __restrict__ out, const input_t* __restrict__ input, const input_t* __restrict__ input2, int64_t size) {
149
+ static_assert(
150
+ std::is_same_v<input_t, float> || std::is_same_v<input_t, scalar_t>,
151
+ "silu_and_mul_stub only supports input_t == float or input_t == scalar_t");
152
using bVec = at::vec::Vectorized<scalar_t>;
153
using fVec = at::vec::Vectorized<float>;
160
- const fVec one = fVec(1.f);
154
155
// no remainder
156
#pragma GCC unroll 4
157
for (int64_t d = 0; d < size; d += bVec::size()) {
165
- bVec x = bVec::loadu(input + d);
166
- fVec x0, x1;
167
- std::tie(x0, x1) = at::vec::convert_to_float(x);
168
- bVec y = bVec::loadu(input2 + d);
169
- fVec y0, y1;
170
- std::tie(y0, y1) = at::vec::convert_to_float(y);
171
- x0 = x0 / (one + x0.neg().exp_u20());
172
- x1 = x1 / (one + x1.neg().exp_u20());
173
- x0 = x0 * y0;
174
- x1 = x1 * y1;
158
+ auto [x0, x1] = load_float_vec2(input + d);
159
+ auto [y0, y1] = load_float_vec2(input2 + d);
160
+ x0 = fast_silu(x0) * y0;
161
+ x1 = fast_silu(x1) * y1;
162
bVec out_vec = convert_from_float_ext<scalar_t>(x0, x1);
163
out_vec.store(out + d);
164
}
165
}
166
167
+template <typename scalar_t, typename input_t>
168
+inline void clamp_sigmoid_and_mul_stub(
169
+ scalar_t* __restrict__ out, const input_t* __restrict__ input, int64_t size, const float alpha, const float limit) {
170
+ static_assert(
171
+ std::is_same_v<input_t, float> || std::is_same_v<input_t, scalar_t>,
172
+ "clamp_sigmoid_and_mul_stub only supports input_t == float or input_t == scalar_t");
173
+ using bVec = at::vec::Vectorized<scalar_t>;
174
+ using fVec = at::vec::Vectorized<float>;
175
+ const fVec one = fVec(1.f);
176
+ const fVec limit_v = fVec(limit);
177
+ const fVec nlimit_v = fVec(-limit);
178
+ const fVec alpha_v = fVec(alpha);
179
+
180
+#pragma GCC unroll 4
181
+ for (int64_t d = 0; d < 2 * size; d += bVec::size()) {
182
+ auto [x0_, y0_] = load_float_vec2(input + d);
183
+ auto [x0, y0] = at::vec::deinterleave2<float>(x0_, y0_);
184
+
185
+ x0 = at::vec::minimum(x0, limit_v);
186
+ y0 = at::vec::minimum(limit_v, at::vec::maximum(nlimit_v, y0));
187
+ x0 = fast_sigmoid_glu(x0, alpha_v) * (y0 + one);
188
+ store_from_float_ext(out + d / 2, x0);
189
+ }
190
+}
191
+
192
template <typename scalar_t>
193
inline void copy_mul_stub(scalar_t* __restrict__ out, const float* __restrict__ input, float weight, int64_t size) {
194
using bVec = at::vec::Vectorized<scalar_t>;
@@ -186,9 +198,8 @@ inline void copy_mul_stub(scalar_t* __restrict__ out, const float* __restrict__
198
int64_t d;
199
#pragma GCC unroll 4
200
for (d = 0; d <= size - kVecSize; d += kVecSize) {
189
- fVec data0 = fVec::loadu(input + d) * weight_vec;
190
- fVec data1 = fVec::loadu(input + d + fVec::size()) * weight_vec;
191
- bVec out_vec = convert_from_float_ext<scalar_t>(data0, data1);
201
+ auto [x0, x1] = load_float_vec2(input + d);
202
+ bVec out_vec = convert_from_float_ext<scalar_t>(x0 * weight_vec, x1 * weight_vec);
203
out_vec.store(out + d);
204
}
205
for (; d < size; ++d) {
@@ -222,63 +233,11 @@ inline void copy_mul_stub(scalar_t* __restrict__ out, const scalar_t* __restrict
233
int64_t d;
234
#pragma GCC unroll 4
235
for (d = 0; d <= size - kVecSize; d += kVecSize) {
225
- bVec x = bVec::loadu(input + d);
226
- fVec x0, x1;
227
- std::tie(x0, x1) = at::vec::convert_to_float(x);
228
- x0 = x0 * weight_vec;
229
- x1 = x1 * weight_vec;
230
- bVec out_vec = convert_from_float_ext<scalar_t>(x0, x1);
236
+ auto [x0, x1] = load_float_vec2(input + d);
237
+ bVec out_vec = convert_from_float_ext<scalar_t>(x0 * weight_vec, x1 * weight_vec);
238
out_vec.store(out + d);
239
}
240
for (; d < size; ++d) {
241
out[d] = static_cast<scalar_t>(input[d] * weight);
242
}
243
}
237
-
238
-template <typename scalar_t>
239
-inline void clamp_sigmoid_and_mul_stub(
240
- scalar_t* __restrict__ out,
241
- const scalar_t* __restrict__ input,
242
- int64_t size,
243
- const float alpha,
244
- const float limit) {
245
- using bVec = at::vec::Vectorized<scalar_t>;
246
- using fVec = at::vec::Vectorized<float>;
247
- const fVec one = fVec(1.f);
248
- const fVec zero = fVec(0.f);
249
- const fVec limit_v = fVec(limit);
250
- const fVec nlimit_v = fVec(-limit);
251
- const fVec alpha_v = fVec(alpha);
252
-
253
- // no remainder
254
-#pragma GCC unroll 4
255
- for (int64_t d = 0; d < size; d += bVec::size()) {
256
- bVec x = bVec::loadu(input + d);
257
- fVec x0_, y0_;
258
- std::tie(x0_, y0_) = at::vec::convert_to_float(x);
259
- float tmp_buffer[fVec::size() * 2]; // 32
260
- float tmp_glu[fVec::size()]; // 16
261
- float tmp_linear[fVec::size()]; // 16
262
- x0_.store(tmp_buffer);
263
- y0_.store(tmp_buffer + fVec::size());
264
- // interleaved: x[2i] = glu, x[2i+1] = linear
265
- for (int j = 0; j < fVec::size(); ++j) {
266
- // x0 [0,2,..30]
267
- tmp_glu[j] = tmp_buffer[j * 2];
268
- // y0 [1,3,...31]
269
- tmp_linear[j] = tmp_buffer[j * 2 + 1];
270
- }
271
- fVec x0 = fVec::loadu(tmp_glu);
272
- fVec y0 = fVec::loadu(tmp_linear);
273
-
274
- // clamp
275
- x0 = at::vec::minimum(x0, limit_v);
276
- y0 = at::vec::minimum(limit_v, at::vec::maximum(nlimit_v, y0));
277
- // x * sigmoid(x * alpha)
278
- x0 = x0 / (one + (x0 * alpha_v).neg().exp_u20());
279
- // (y + 1) * x
280
- y0 = y0 + one;
281
- x0 = x0 * y0;
282
- convert_from_float_and_store<scalar_t>(out + d / 2, x0);
283
- }
284
-}
csrc/cpu/sgl-kernels/moe_fp8.cpp
+6
-6
@@ -37,7 +37,7 @@ void fused_experts_fp_kernel_impl(
37
int64_t num_tokens_post_pad,
38
float alpha,
39
float limit,
40
- CPUAcTMethod act_func,
40
+ CPUActMethod act_func,
41
bool with_bias) {
42
constexpr int64_t BLOCK_M = block_size_m();
43
constexpr int64_t BLOCK_N = block_size_n();
@@ -122,17 +122,17 @@ void fused_experts_fp_kernel_impl(
122
});
123
124
// stage 1.5: intermediate_cache1 = silu(intermediate_cache0)
125
- if (act_func == CPUAcTMethod::silu_and_mul) {
125
+ if (act_func == CPUActMethod::silu_and_mul) {
126
at::parallel_for(0, M * topk, 0, [&](int64_t begin, int64_t end) {
127
for (int64_t m = begin; m < end; ++m) {
128
silu_and_mul_stub(ic1 + m * N, ic0 + m * 2 * N, ic0 + m * 2 * N + N, N);
129
}
130
});
131
- } else if (act_func == CPUAcTMethod::swiglu) {
131
+ } else if (act_func == CPUActMethod::swiglu) {
132
at::parallel_for(0, M * topk, 0, [&](int64_t begin, int64_t end) {
133
for (int64_t m = begin; m < end; ++m) {
134
- clamp_sigmoid_and_mul_stub(ic1 + m * N, ic0 + m * 2 * N, N, alpha, limit);
135
- clamp_sigmoid_and_mul_stub(ic1 + m * N + N / 2, ic0 + m * 2 * N + N, N, alpha, limit);
134
+ clamp_sigmoid_and_mul_stub(ic1 + m * N, ic0 + m * 2 * N, N / 2, alpha, limit);
135
+ clamp_sigmoid_and_mul_stub(ic1 + m * N + N / 2, ic0 + m * 2 * N + N, N / 2, alpha, limit);
136
}
137
});
138
}
@@ -243,7 +243,7 @@ void fused_experts_fp_kernel_impl(
243
int64_t num_tokens_post_pad, \
244
float alpha, \
245
float limit, \
246
- CPUAcTMethod act_func, \
246
+ CPUActMethod act_func, \
247
bool with_bias)
248
249
INSTANTIATE_MOE_FP_TEMPLATE(at::BFloat16, at::Float8_e4m3fn, float, false);
csrc/cpu/sgl-kernels/moe_int8.cpp
+11
-19
@@ -53,16 +53,14 @@ inline void silu_and_mul(
53
vc1[col] = _mm512_mul_ps(_mm512_mul_ps(vc1[col], vas), vbs1[col]);
54
};
55
56
- using bVec = at::vec::Vectorized<scalar_t>;
57
- using fVec = at::vec::Vectorized<float>;
58
- const fVec one = fVec(1.f);
56
auto silu_and_mul = [&](auto col) {
60
- fVec x = fVec(vc0[col]);
61
- fVec y = fVec(vc1[col]);
62
- x = x / (one + x.neg().exp_u20());
63
- vc0[col] = x * y;
57
+ __m512 x = vc0[col];
58
+ __m512 y = vc1[col];
59
+ vc0[col] = _mm512_mul_ps(_mm512_rcp14_silu_ps(x), y);
60
};
61
62
+ using bVec = at::vec::Vectorized<scalar_t>;
63
+ using fVec = at::vec::Vectorized<float>;
64
auto storec = [&](auto col, int64_t m) {
65
if constexpr (col % 2 == 0) {
66
fVec x0 = fVec(vc0[col + 0]);
@@ -229,23 +227,17 @@ struct tinygemm_kernel_vnni<at::BFloat16, BLOCK_M, BLOCK_N> {
227
};
228
Unroll<ROWS * COLS>{}(scalec);
229
232
- using Vec = at::vec::Vectorized<float>;
233
- const Vec one = Vec(1.f);
230
auto storec = [&](auto i) {
231
constexpr int row = i / COLS;
232
constexpr int col = i % COLS;
233
// for COLS = 2, 4 use 512bit store
234
if constexpr (col % 2 == 0) {
239
- Vec x0 = _mm512_castsi512_ps(vc0[row * COLS + col + 0]);
240
- Vec x1 = _mm512_castsi512_ps(vc0[row * COLS + col + 1]);
241
- Vec y0 = _mm512_castsi512_ps(vc1[row * COLS + col + 0]);
242
- Vec y1 = _mm512_castsi512_ps(vc1[row * COLS + col + 1]);
243
- // silu
244
- x0 = x0 / (one + x0.neg().exp_u20());
245
- x1 = x1 / (one + x1.neg().exp_u20());
246
- // mul
247
- x0 = x0 * y0;
248
- x1 = x1 * y1;
235
+ __m512 x0 = _mm512_castsi512_ps(vc0[row * COLS + col + 0]);
236
+ __m512 x1 = _mm512_castsi512_ps(vc0[row * COLS + col + 1]);
237
+ __m512 y0 = _mm512_castsi512_ps(vc1[row * COLS + col + 0]);
238
+ __m512 y1 = _mm512_castsi512_ps(vc1[row * COLS + col + 1]);
239
+ x0 = _mm512_mul_ps(_mm512_rcp14_silu_ps(x0), y0);
240
+ x1 = _mm512_mul_ps(_mm512_rcp14_silu_ps(x1), y1);
241
242
_mm512_storeu_si512(
243
reinterpret_cast<__m512i*>((C + row * ldc + col * 16)),
csrc/cpu/sgl-kernels/vec.h
+195
-6
@@ -32,11 +32,11 @@ inline Vectorized<scalar_t> convert_from_float_ext(const Vectorized<float>& a, c
32
}
33
34
template <typename scalar_t>
35
-inline void convert_from_float_and_store(scalar_t* out, const Vectorized<float>& a) {
36
- float out_buffer[at::vec::Vectorized<float>::size()];
35
+inline void store_from_float_ext(scalar_t* out, const Vectorized<float>& a) {
36
+ float out_buffer[Vectorized<float>::size()];
37
a.store(out_buffer);
38
- for (int i = 0; i < 16; i++) {
39
- out[i] = (scalar_t)out_buffer[i];
38
+ for (int i = 0; i < Vectorized<float>::size(); ++i) {
39
+ out[i] = static_cast<scalar_t>(out_buffer[i]);
40
}
41
}
42
@@ -59,6 +59,17 @@ inline std::tuple<Vectorized<float>, Vectorized<float>> load_float_vec2(const fl
59
return std::make_tuple(x0, x1);
60
}
61
62
+template <typename scalar_t, typename std::enable_if_t<is_reduced_floating_point_v<scalar_t>, int> = 1>
63
+inline at::vec::Vectorized<float> load_float_vec(const scalar_t* __restrict__ data) {
64
+ at::vec::Vectorized<float> out;
65
+ if constexpr (std::is_same_v<scalar_t, at::BFloat16>) {
66
+ at::vec::load_fp32_from_bf16(data, out);
67
+ } else {
68
+ at::vec::load_fp32_from_fp16(data, out);
69
+ }
70
+ return out;
71
+}
72
+
73
#if defined(CPU_CAPABILITY_AVX512)
74
75
// `at::vec::convert_from_float<>` from PyTorch doesn't have avx512-bf16 intrinsics
@@ -70,8 +81,14 @@ convert_from_float_ext<at::BFloat16>(const Vectorized<float>& a, const Vectorize
81
}
82
83
template <>
73
-inline void convert_from_float_and_store<at::BFloat16>(at::BFloat16* out, const Vectorized<float>& a) {
74
- _mm256_storeu_si256((__m256i*)out, (__m256i)(_mm512_cvtneps_pbh(__m512(a))));
84
+inline void store_from_float_ext<at::BFloat16>(at::BFloat16* out, const Vectorized<float>& a) {
85
+ _mm256_storeu_si256(reinterpret_cast<__m256i*>(out), (__m256i)(_mm512_cvtneps_pbh(__m512(a))));
86
+}
87
+
88
+template <>
89
+inline void store_from_float_ext<at::Half>(at::Half* out, const Vectorized<float>& a) {
90
+ _mm256_storeu_si256(
91
+ reinterpret_cast<__m256i*>(out), _mm512_cvtps_ph(__m512(a), _MM_FROUND_TO_NEAREST_INT | _MM_FROUND_NO_EXC));
92
}
93
94
#define CVT_BF16_TO_FP32(a) _mm512_castsi512_ps(_mm512_slli_epi32(_mm512_cvtepu16_epi32(a), 16))
@@ -301,6 +318,77 @@ inline void quantize_row_int8<at::BFloat16>(
318
// transpose utils
319
// taken from my PR in ggml: https://github.com/ggml-org/llama.cpp/pull/8998
320
#if defined(CPU_CAPABILITY_AVX512)
321
+inline void transpose_16x16_16bit(__m256i* v) {
322
+ __m256i v1[16];
323
+ v1[0] = _mm256_unpacklo_epi16(v[0], v[1]);
324
+ v1[1] = _mm256_unpackhi_epi16(v[0], v[1]);
325
+ v1[2] = _mm256_unpacklo_epi16(v[2], v[3]);
326
+ v1[3] = _mm256_unpackhi_epi16(v[2], v[3]);
327
+ v1[4] = _mm256_unpacklo_epi16(v[4], v[5]);
328
+ v1[5] = _mm256_unpackhi_epi16(v[4], v[5]);
329
+ v1[6] = _mm256_unpacklo_epi16(v[6], v[7]);
330
+ v1[7] = _mm256_unpackhi_epi16(v[6], v[7]);
331
+ v1[8] = _mm256_unpacklo_epi16(v[8], v[9]);
332
+ v1[9] = _mm256_unpackhi_epi16(v[8], v[9]);
333
+ v1[10] = _mm256_unpacklo_epi16(v[10], v[11]);
334
+ v1[11] = _mm256_unpackhi_epi16(v[10], v[11]);
335
+ v1[12] = _mm256_unpacklo_epi16(v[12], v[13]);
336
+ v1[13] = _mm256_unpackhi_epi16(v[12], v[13]);
337
+ v1[14] = _mm256_unpacklo_epi16(v[14], v[15]);
338
+ v1[15] = _mm256_unpackhi_epi16(v[14], v[15]);
339
+
340
+ v[0] = _mm256_unpacklo_epi32(v1[0], v1[2]);
341
+ v[1] = _mm256_unpackhi_epi32(v1[0], v1[2]);
342
+ v[2] = _mm256_unpacklo_epi32(v1[1], v1[3]);
343
+ v[3] = _mm256_unpackhi_epi32(v1[1], v1[3]);
344
+ v[4] = _mm256_unpacklo_epi32(v1[4], v1[6]);
345
+ v[5] = _mm256_unpackhi_epi32(v1[4], v1[6]);
346
+ v[6] = _mm256_unpacklo_epi32(v1[5], v1[7]);
347
+ v[7] = _mm256_unpackhi_epi32(v1[5], v1[7]);
348
+ v[8] = _mm256_unpacklo_epi32(v1[8], v1[10]);
349
+ v[9] = _mm256_unpackhi_epi32(v1[8], v1[10]);
350
+ v[10] = _mm256_unpacklo_epi32(v1[9], v1[11]);
351
+ v[11] = _mm256_unpackhi_epi32(v1[9], v1[11]);
352
+ v[12] = _mm256_unpacklo_epi32(v1[12], v1[14]);
353
+ v[13] = _mm256_unpackhi_epi32(v1[12], v1[14]);
354
+ v[14] = _mm256_unpacklo_epi32(v1[13], v1[15]);
355
+ v[15] = _mm256_unpackhi_epi32(v1[13], v1[15]);
356
+
357
+ v1[0] = _mm256_unpacklo_epi64(v[0], v[4]);
358
+ v1[1] = _mm256_unpackhi_epi64(v[0], v[4]);
359
+ v1[2] = _mm256_unpacklo_epi64(v[1], v[5]);
360
+ v1[3] = _mm256_unpackhi_epi64(v[1], v[5]);
361
+ v1[4] = _mm256_unpacklo_epi64(v[2], v[6]);
362
+ v1[5] = _mm256_unpackhi_epi64(v[2], v[6]);
363
+ v1[6] = _mm256_unpacklo_epi64(v[3], v[7]);
364
+ v1[7] = _mm256_unpackhi_epi64(v[3], v[7]);
365
+ v1[8] = _mm256_unpacklo_epi64(v[8], v[12]);
366
+ v1[9] = _mm256_unpackhi_epi64(v[8], v[12]);
367
+ v1[10] = _mm256_unpacklo_epi64(v[9], v[13]);
368
+ v1[11] = _mm256_unpackhi_epi64(v[9], v[13]);
369
+ v1[12] = _mm256_unpacklo_epi64(v[10], v[14]);
370
+ v1[13] = _mm256_unpackhi_epi64(v[10], v[14]);
371
+ v1[14] = _mm256_unpacklo_epi64(v[11], v[15]);
372
+ v1[15] = _mm256_unpackhi_epi64(v[11], v[15]);
373
+
374
+ v[0] = _mm256_permute2x128_si256(v1[0], v1[8], 0x20);
375
+ v[1] = _mm256_permute2x128_si256(v1[1], v1[9], 0x20);
376
+ v[2] = _mm256_permute2x128_si256(v1[2], v1[10], 0x20);
377
+ v[3] = _mm256_permute2x128_si256(v1[3], v1[11], 0x20);
378
+ v[4] = _mm256_permute2x128_si256(v1[4], v1[12], 0x20);
379
+ v[5] = _mm256_permute2x128_si256(v1[5], v1[13], 0x20);
380
+ v[6] = _mm256_permute2x128_si256(v1[6], v1[14], 0x20);
381
+ v[7] = _mm256_permute2x128_si256(v1[7], v1[15], 0x20);
382
+ v[8] = _mm256_permute2x128_si256(v1[0], v1[8], 0x31);
383
+ v[9] = _mm256_permute2x128_si256(v1[1], v1[9], 0x31);
384
+ v[10] = _mm256_permute2x128_si256(v1[2], v1[10], 0x31);
385
+ v[11] = _mm256_permute2x128_si256(v1[3], v1[11], 0x31);
386
+ v[12] = _mm256_permute2x128_si256(v1[4], v1[12], 0x31);
387
+ v[13] = _mm256_permute2x128_si256(v1[5], v1[13], 0x31);
388
+ v[14] = _mm256_permute2x128_si256(v1[6], v1[14], 0x31);
389
+ v[15] = _mm256_permute2x128_si256(v1[7], v1[15], 0x31);
390
+}
391
+
392
inline void transpose_16x16_32bit(__m512i* v) {
393
__m512i v1[16];
394
v1[0] = _mm512_unpacklo_epi32(v[0], v[1]);
@@ -394,6 +482,62 @@ inline std::tuple<__m512i, __m512i> transpose_2x32_16bit(__m512i r0, __m512i r1)
482
}
483
#pragma GCC diagnostic pop
484
485
+// Note: mapped from aten exp_u20
486
+inline __attribute__((always_inline)) __m512 _mm512_exp_u20_ps(const __m512 values) {
487
+ const __m512 vec_factorial_1 = _mm512_set1_ps(0.999999701f);
488
+ const __m512 vec_factorial_2 = _mm512_set1_ps(0.499991506f);
489
+ const __m512 vec_factorial_3 = _mm512_set1_ps(0.166676521f);
490
+ const __m512 vec_factorial_4 = _mm512_set1_ps(0.0418978221f);
491
+ const __m512 vec_factorial_5 = _mm512_set1_ps(0.00828929059f);
492
+ const __m512 vec_exp_log2ef = _mm512_castsi512_ps(_mm512_set1_epi32(0x3fb8aa3b)); // log2(e)
493
+ const __m512 vec_half = _mm512_set1_ps(0.5f);
494
+ const __m512 vec_one = _mm512_set1_ps(1.f);
495
+ const __m512 vec_zero = _mm512_set1_ps(0.f);
496
+ const __m512 vec_two = _mm512_set1_ps(2.f);
497
+ const __m512 vec_ln2f = _mm512_castsi512_ps(_mm512_set1_epi32(0x3f317218));
498
+ const __m512 vec_ln_flt_min = _mm512_castsi512_ps(_mm512_set1_epi32(0xc2aeac50));
499
+ const __m512 vec_ln_flt_max = _mm512_castsi512_ps(_mm512_set1_epi32(0x42b17218));
500
+ const __m512i vec_127 = _mm512_set1_epi32(0x0000007f);
501
+ const int n_mantissa_bits = 23;
502
+
503
+ // exp(x) =
504
+ // = exp(n * ln(2) + r) // divide x by ln(2) and get quot and rem
505
+ // = 2^n * exp(r) // simplify the exp(n*ln(2)) expression
506
+
507
+ auto less_ln_flt_min_mask = _mm512_cmp_ps_mask(values, vec_ln_flt_min, 1 /*_CMP_LT_OS*/);
508
+ auto vec_src = _mm512_min_ps(values, vec_ln_flt_max);
509
+ vec_src = _mm512_max_ps(vec_src, vec_ln_flt_min);
510
+
511
+ // fx = floorf(x * log2ef + 0.5)
512
+ auto vec_fx = _mm512_fmadd_ps(vec_src, vec_exp_log2ef, vec_half);
513
+ auto vec_fx_i = _mm512_cvt_roundps_epi32(vec_fx, _MM_FROUND_TO_NEG_INF | _MM_FROUND_NO_EXC);
514
+ vec_fx = _mm512_cvtepi32_ps(vec_fx_i);
515
+
516
+ // x = x - fx * ln2
517
+ auto vec_exp_poly = _mm512_fnmadd_ps(vec_fx, vec_ln2f, vec_src);
518
+
519
+ // compute polynomial
520
+ auto vec_res = _mm512_fmadd_ps(vec_exp_poly, vec_factorial_5, vec_factorial_4);
521
+ vec_res = _mm512_fmadd_ps(vec_exp_poly, vec_res, vec_factorial_3);
522
+ vec_res = _mm512_fmadd_ps(vec_exp_poly, vec_res, vec_factorial_2);
523
+ vec_res = _mm512_fmadd_ps(vec_exp_poly, vec_res, vec_factorial_1);
524
+ vec_res = _mm512_fmadd_ps(vec_exp_poly, vec_res, vec_one);
525
+
526
+ // compute 2^(n-1)
527
+ auto vec_exp_number = _mm512_sub_ps(vec_fx, vec_one);
528
+ auto vec_exp_number_i = _mm512_cvtps_epi32(vec_exp_number);
529
+ auto vec_two_pow_n_i = _mm512_add_epi32(vec_exp_number_i, vec_127);
530
+ vec_two_pow_n_i = _mm512_slli_epi32(vec_two_pow_n_i, n_mantissa_bits);
531
+ auto vec_two_pow_n = _mm512_castsi512_ps(vec_two_pow_n_i);
532
+ vec_two_pow_n = _mm512_mask_blend_ps(less_ln_flt_min_mask, vec_two_pow_n, vec_zero);
533
+
534
+ // y = y * 2^n
535
+ vec_res = _mm512_mul_ps(vec_res, vec_two_pow_n);
536
+ vec_res = _mm512_mul_ps(vec_res, vec_two);
537
+ return vec_res;
538
+}
539
+
540
+// Note: mapped from aten fexp_u20
541
inline __attribute__((always_inline)) __m512 _mm512_fexp_u20_ps(const __m512 values) {
542
const __m512 vec_c0 = _mm512_set1_ps(0.00010703434948458272f);
543
const __m512 vec_c1 = _mm512_set1_ps(0.30354260500649682f);
@@ -444,6 +588,51 @@ inline __attribute__((always_inline)) __m512 _mm512_fexp_u20_ps(const __m512 val
588
// final interpretation to float
589
return _mm512_castsi512_ps(casted_integer);
590
}
591
+
592
+// sigmoid(x) = 1 / (1 + exp(-x)); avoid vdivps via rcp14
593
+inline __attribute__((always_inline)) __m512 _mm512_rcp14_sigmoid_ps(__m512 x) {
594
+ __m512 minus_x = _mm512_xor_ps(_mm512_set1_ps(-0.f), x);
595
+ __m512 denom = _mm512_add_ps(_mm512_exp_u20_ps(minus_x), _mm512_set1_ps(1.f));
596
+ return _mm512_rcp14_ps(denom);
597
+}
598
+
599
+// SiLU(x) = x * sigmoid(x)
600
+inline __attribute__((always_inline)) __m512 _mm512_rcp14_silu_ps(__m512 x) {
601
+ return _mm512_mul_ps(x, _mm512_rcp14_sigmoid_ps(x));
602
+}
603
+
604
+// x * sigmoid(x * alpha) for clamped SwiGLU
605
+inline __attribute__((always_inline)) __m512 _mm512_rcp14_sigmoid_glu_ps(__m512 x, __m512 alpha) {
606
+ __m512 xa = _mm512_mul_ps(x, alpha);
607
+ return _mm512_mul_ps(x, _mm512_rcp14_sigmoid_ps(xa));
608
+}
609
+
610
+#endif
611
+
612
+inline at::vec::Vectorized<float> fast_sigmoid(const at::vec::Vectorized<float>& x) {
613
+#if defined(CPU_CAPABILITY_AVX512)
614
+ return at::vec::Vectorized<float>(_mm512_rcp14_sigmoid_ps(x));
615
+#else
616
+ const auto one = at::vec::Vectorized<float>(1.f);
617
+ return one / (one + x.neg().exp_u20());
618
+#endif
619
+}
620
+
621
+inline at::vec::Vectorized<float> fast_silu(const at::vec::Vectorized<float>& x) {
622
+#if defined(CPU_CAPABILITY_AVX512)
623
+ return at::vec::Vectorized<float>(_mm512_rcp14_silu_ps(x));
624
+#else
625
+ return x * fast_sigmoid(x);
626
#endif
627
+}
628
+
629
+inline at::vec::Vectorized<float>
630
+fast_sigmoid_glu(const at::vec::Vectorized<float>& x, const at::vec::Vectorized<float>& alpha) {
631
+#if defined(CPU_CAPABILITY_AVX512)
632
+ return at::vec::Vectorized<float>(_mm512_rcp14_sigmoid_glu_ps(x, alpha));
633
+#else
634
+ return x * fast_sigmoid(x * alpha);
635
+#endif
636
+}
637
638
} // anonymous namespace
csrc/cpu/torch_bindings.cpp
+3
-2
@@ -110,7 +110,7 @@ std::tuple<at::Tensor, at::Tensor> chunk_gated_delta_rule_cpu(
110
const at::Tensor& g, const at::Tensor& beta,
111
const at::Tensor& initial_state, bool output_final_state,
112
const at::Tensor& cu_seqlens, bool head_first, bool use_qk_l2norm_in_kernel,
113
- double eps = 1e-5);
113
+ const at::Tensor& initial_state_indices, double eps = 1e-5);
114
115
at::Tensor fused_sigmoid_gating_delta_rule_update_cpu(
116
const at::Tensor& A_log, const at::Tensor& dt_bias, const at::Tensor& q,
@@ -546,7 +546,8 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
546
"Tensor g, Tensor beta, "
547
"Tensor initial_state, bool output_final_state, Tensor cu_seqlens, bool "
548
"head_first, "
549
- "bool use_qk_l2norm_in_kernel, float eps=1e-5) -> (Tensor, Tensor)");
549
+ "bool use_qk_l2norm_in_kernel, Tensor initial_state_indices, float "
550
+ "eps=1e-5) -> (Tensor, Tensor)");
551
ops.impl("chunk_gated_delta_rule_cpu", torch::kCPU,
552
&chunk_gated_delta_rule_cpu);
553
ops.def(
tests/kernels/mamba/cpu/test_cpu_gdn_ops.py
+13
-2
@@ -26,6 +26,13 @@ HEAD_DIMS = [
26
(32, 32),
27
(64, 32),
28
]
29
+# chunk_gated_delta_rule_cpu (the chunked-prefill kernel) only supports
30
+# head_dim == head_dim_v in {64, 128}; the decode-path update kernel above
31
+# has no such restriction and keeps using the wider HEAD_DIMS list.
32
+CHUNK_HEAD_DIMS = [
33
+ (64, 64),
34
+ (128, 128),
35
+]
36
CHUNK_SIZE = 64
37
CONV_DIM = 128
38
CONV_KERNEL = 4
@@ -258,7 +265,7 @@ def test_fused_sigmoid_gating_delta_rule_update_cpu(
265
# prefill path
266
@pytest.mark.parametrize("seq_lens", PREFILL_SEQ_LENS)
267
@pytest.mark.parametrize("num_heads", NUM_HEADS)
261
-@pytest.mark.parametrize("head_dims", HEAD_DIMS)
268
+@pytest.mark.parametrize("head_dims", CHUNK_HEAD_DIMS)
269
@torch.inference_mode()
270
def test_chunk_gated_delta_rule_cpu(
271
seq_lens: list[int],
@@ -307,6 +314,7 @@ def test_chunk_gated_delta_rule_cpu(
314
cu_seqlens=cu_seqlens,
315
head_first=False,
316
use_qk_l2norm_in_kernel=True,
317
+ initial_state_indices=torch.arange(len(seq_lens), dtype=torch.int32),
318
)
319
320
torch.testing.assert_close(out, out_ref, atol=1e-2, rtol=1e-2)
@@ -331,7 +339,7 @@ TWO_CALL_SPLITS = [
339
340
@pytest.mark.parametrize("total_tokens, split", TWO_CALL_SPLITS)
341
@pytest.mark.parametrize("num_heads", NUM_HEADS)
334
-@pytest.mark.parametrize("head_dims", HEAD_DIMS)
342
+@pytest.mark.parametrize("head_dims", CHUNK_HEAD_DIMS)
343
@torch.inference_mode()
344
def test_chunk_gated_delta_rule_cpu_two_call_split(
345
total_tokens: int,
@@ -370,6 +378,7 @@ def test_chunk_gated_delta_rule_cpu_two_call_split(
378
cu_seqlens=torch.tensor([0, total_tokens], dtype=torch.int32),
379
head_first=False,
380
use_qk_l2norm_in_kernel=True,
381
+ initial_state_indices=torch.zeros(1, dtype=torch.int32),
382
)
383
384
# Call 1: tokens [0:split], no initial state, capture final state.
@@ -384,6 +393,7 @@ def test_chunk_gated_delta_rule_cpu_two_call_split(
393
cu_seqlens=torch.tensor([0, split], dtype=torch.int32),
394
head_first=False,
395
use_qk_l2norm_in_kernel=True,
396
+ initial_state_indices=torch.zeros(1, dtype=torch.int32),
397
)
398
# Call 2: tokens [split:T] seeded with call 1's final state and a cu_seqlens
399
# rebased to start at 0, as cpu_gdn_attention_core continues a prefill chunk.
@@ -399,6 +409,7 @@ def test_chunk_gated_delta_rule_cpu_two_call_split(
409
cu_seqlens=torch.tensor([0, tail], dtype=torch.int32),
410
head_first=False,
411
use_qk_l2norm_in_kernel=True,
412
+ initial_state_indices=torch.zeros(1, dtype=torch.int32),
413
)
414
415
out_split = torch.cat([out1, out2], dim=1)
vllm/_custom_ops.py
+2
@@ -3491,6 +3491,7 @@ def chunk_gated_delta_rule_cpu(
3491
cu_seqlens: torch.Tensor,
3492
head_first: bool,
3493
use_qk_l2norm_in_kernel: bool,
3494
+ initial_state_indices: torch.Tensor,
3495
eps: float = 1e-5,
3496
) -> tuple[torch.Tensor, torch.Tensor]:
3497
return torch.ops._C.chunk_gated_delta_rule_cpu(
@@ -3504,6 +3505,7 @@ def chunk_gated_delta_rule_cpu(
3505
cu_seqlens,
3506
head_first,
3507
use_qk_l2norm_in_kernel,
3508
+ initial_state_indices,
3509
eps,
3510
)
3511
vllm/model_executor/layers/mamba/ops/cpu/gdn_attention.py
+27
-21
@@ -236,22 +236,24 @@ def _cpu_gdn_attention_nonspec(
236
A_log=layer.A_log, a=prefill_a, b=prefill_b, dt_bias=layer.dt_bias
237
)
238
239
- initial_state = ssm_state[prefill_state_indices]
240
- initial_state[~prefill_has_initial_state, ...] = 0
241
- attn_out, last_recurrent_state = ops.chunk_gated_delta_rule_cpu(
239
+ # zero pool slots for sequences without a prior state; the kernel
240
+ # gathers/mutates ssm_state in place via prefill_state_indices, so no
241
+ # separate scatter-back is needed after the call.
242
+ no_initial_state = prefill_state_indices[~prefill_has_initial_state]
243
+ if no_initial_state.numel() > 0:
244
+ ssm_state[no_initial_state] = 0
245
+ attn_out, _ = ops.chunk_gated_delta_rule_cpu(
246
query=query,
247
key=key,
248
value=value,
249
g=g,
250
beta=beta,
247
- initial_state=initial_state,
251
+ initial_state=ssm_state,
252
output_final_state=True,
253
cu_seqlens=prefill_query_start_loc,
254
head_first=False,
255
use_qk_l2norm_in_kernel=True,
252
- )
253
- ssm_state[prefill_state_indices] = last_recurrent_state.to(
254
- ssm_state.dtype, copy=False
256
+ initial_state_indices=prefill_state_indices,
257
)
258
core_attn_out[prefill_token_start:prefill_token_end] = attn_out.squeeze(0)
259
@@ -616,22 +618,24 @@ def _spec_aware_nonspec(
618
g, beta = ops.fused_gdn_gating_cpu(
619
A_log=layer.A_log, a=prefill_a, b=prefill_b, dt_bias=layer.dt_bias
620
)
619
- initial_state = ssm_state[prefill_state_indices]
620
- initial_state[~prefill_has_initial_state, ...] = 0
621
- attn_out, last_recurrent_state = ops.chunk_gated_delta_rule_cpu(
621
+ # zero pool slots for sequences without a prior state; the kernel
622
+ # gathers/mutates ssm_state in place via prefill_state_indices, so no
623
+ # separate scatter-back is needed after the call.
624
+ no_initial_state = prefill_state_indices[~prefill_has_initial_state]
625
+ if no_initial_state.numel() > 0:
626
+ ssm_state[no_initial_state] = 0
627
+ attn_out, _ = ops.chunk_gated_delta_rule_cpu(
628
query=query,
629
key=key,
630
value=value,
631
g=g,
632
beta=beta,
627
- initial_state=initial_state,
633
+ initial_state=ssm_state,
634
output_final_state=True,
635
cu_seqlens=prefill_query_start_loc,
636
head_first=False,
637
use_qk_l2norm_in_kernel=True,
632
- )
633
- ssm_state[prefill_state_indices] = last_recurrent_state.to(
634
- ssm_state.dtype, copy=False
638
+ initial_state_indices=prefill_state_indices,
639
)
640
core_attn_out[prefill_token_start:prefill_token_end] = attn_out.squeeze(0)
641
@@ -690,22 +694,24 @@ def _spec_aware_nonspec_subset(
694
g, beta = ops.fused_gdn_gating_cpu(
695
A_log=layer.A_log, a=a, b=b, dt_bias=layer.dt_bias
696
)
693
- initial_state = ssm_state[prefill_state_indices]
694
- initial_state[~has_initial_state, ...] = 0
695
- attn_out, last_recurrent_state = ops.chunk_gated_delta_rule_cpu(
697
+ # zero pool slots for sequences without a prior state; the kernel
698
+ # gathers/mutates ssm_state in place via prefill_state_indices, so no
699
+ # separate scatter-back is needed after the call.
700
+ no_initial_state = prefill_state_indices[~has_initial_state]
701
+ if no_initial_state.numel() > 0:
702
+ ssm_state[no_initial_state] = 0
703
+ attn_out, _ = ops.chunk_gated_delta_rule_cpu(
704
query=query,
705
key=key,
706
value=value,
707
g=g,
708
beta=beta,
701
- initial_state=initial_state,
709
+ initial_state=ssm_state,
710
output_final_state=True,
711
cu_seqlens=prefill_qsl,
712
head_first=False,
713
use_qk_l2norm_in_kernel=True,
706
- )
707
- ssm_state[prefill_state_indices] = last_recurrent_state.to(
708
- ssm_state.dtype, copy=False
714
+ initial_state_indices=prefill_state_indices,
715
)
716
return attn_out.squeeze(0)
717