@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
hpp 327 lines 12.1 KB
Raw
1 #ifndef CPU_ATTN_VEC_HPP
2 #define CPU_ATTN_VEC_HPP
3
4 #include "cpu_attn_fp8.hpp"
5 #include "cpu_attn_impl.hpp"
6
7 namespace cpu_attention {
8
9 namespace {
10
11 // Load 32 kv_cache_t elements starting at ptr and return them as two FP32Vec16s
12 // covering the lower 16 and upper 16 positions.
13 // For FP8: both halves come from a single BF16Vec32 dequant of 32 bytes.
14 // For BF16/FP16/FP32: two separate vector loads at ptr and ptr+16.
15 template <typename kv_cache_t>
16 FORCE_INLINE std::pair<vec_op::FP32Vec16, vec_op::FP32Vec16> load_b_pair_vec(
17 const kv_cache_t* ptr) {
18 if constexpr (std::is_same_v<kv_cache_t, c10::Float8_e4m3fn>) {
19 // BF16 container, but values are in the FP16 exponent range (bias 15 not
20 // 127).
21 vec_op::BF16Vec32 bf16_b_reg(reinterpret_cast<const uint8_t*>(ptr),
22 vec_op::fp8_e4m3_tag{});
23 return {vec_op::FP32Vec16(bf16_b_reg, 0), vec_op::FP32Vec16(bf16_b_reg, 1)};
24 } else if constexpr (std::is_same_v<kv_cache_t, c10::Float8_e5m2>) {
25 vec_op::BF16Vec32 bf16_b_reg(reinterpret_cast<const uint8_t*>(ptr),
26 vec_op::fp8_e5m2_tag{});
27 return {vec_op::FP32Vec16(bf16_b_reg, 0), vec_op::FP32Vec16(bf16_b_reg, 1)};
28 } else {
29 using load_vec_t = typename VecTypeTrait<kv_cache_t>::vec_t;
30 return std::make_pair(vec_op::FP32Vec16(load_vec_t(ptr)),
31 vec_op::FP32Vec16(load_vec_t(ptr + 16)));
32 }
33 }
34
35 // 8-2-16 pattern, 8 regs for A, 2 regs for B, 16 regs for C, [8, K] @ [k, 32]
36 template <typename kv_cache_t>
37 class TileGemm82 {
38 public:
39 template <AttentionGemmPhase phase, int32_t k_size>
40 FORCE_INLINE static void gemm(const int32_t m_size,
41 float* __restrict__ a_tile,
42 kv_cache_t* __restrict__ b_tile,
43 float* __restrict__ c_tile, const int64_t lda,
44 const int64_t ldb, const int64_t ldc,
45 const int32_t block_size,
46 const int32_t dynamic_k_size,
47 const bool accum_c) {
48 switch (m_size) {
49 case 1:
50 gemm_micro<1>(a_tile, b_tile, c_tile, lda, ldb, ldc, block_size,
51 dynamic_k_size, accum_c);
52 break;
53 case 2:
54 gemm_micro<2>(a_tile, b_tile, c_tile, lda, ldb, ldc, block_size,
55 dynamic_k_size, accum_c);
56 break;
57 case 3:
58 case 4:
59 gemm_micro<4>(a_tile, b_tile, c_tile, lda, ldb, ldc, block_size,
60 dynamic_k_size, accum_c);
61 break;
62 case 5:
63 case 6:
64 gemm_micro<6>(a_tile, b_tile, c_tile, lda, ldb, ldc, block_size,
65 dynamic_k_size, accum_c);
66 break;
67 case 7:
68 case 8:
69 gemm_micro<8>(a_tile, b_tile, c_tile, lda, ldb, ldc, block_size,
70 dynamic_k_size, accum_c);
71 break;
72 }
73 }
74
75 template <int32_t M>
76 static void gemm_micro(float* __restrict__ a_tile,
77 kv_cache_t* __restrict__ b_tile,
78 float* __restrict__ c_tile, const int64_t lda,
79 const int64_t ldb, const int64_t ldc,
80 const int32_t block_size, const int32_t dynamic_k_size,
81 const bool accum_c) {
82 static_assert(0 < M && M <= 8);
83
84 float* __restrict__ curr_c_0 = c_tile;
85 float* __restrict__ curr_c_1 = c_tile + 16;
86
87 vec_op::FP32Vec16 c_regs[M * 2];
88 if (accum_c) {
89 float* __restrict__ curr_m_c_0 = curr_c_0;
90 float* __restrict__ curr_m_c_1 = curr_c_1;
91 vec_op::unroll_loop<int32_t, M>([&](int32_t i) {
92 c_regs[i * 2] = vec_op::FP32Vec16(curr_m_c_0);
93 c_regs[i * 2 + 1] = vec_op::FP32Vec16(curr_m_c_1);
94
95 // update
96 curr_m_c_0 += ldc;
97 curr_m_c_1 += ldc;
98 });
99 }
100
101 float* __restrict__ curr_a = a_tile;
102 kv_cache_t* __restrict__ curr_b = b_tile;
103
104 for (int32_t k = 0; k < dynamic_k_size; ++k) {
105 auto fp32_b_regs = load_b_pair_vec(curr_b);
106 auto fp32_b_0_reg = fp32_b_regs.first;
107 auto fp32_b_1_reg = fp32_b_regs.second;
108
109 float* __restrict__ curr_m_a = curr_a;
110 vec_op::unroll_loop<int32_t, M>([&](int32_t i) {
111 vec_op::FP32Vec16 a_reg(*curr_m_a);
112 c_regs[i * 2] = c_regs[i * 2] + a_reg * fp32_b_0_reg;
113 c_regs[i * 2 + 1] = c_regs[i * 2 + 1] + a_reg * fp32_b_1_reg;
114
115 // update
116 curr_m_a += lda;
117 });
118
119 // update
120 curr_a += 1;
121 curr_b += ldb;
122 }
123
124 vec_op::unroll_loop<int32_t, M>([&](int32_t i) {
125 c_regs[i * 2].save(curr_c_0);
126 c_regs[i * 2 + 1].save(curr_c_1);
127
128 // update
129 curr_c_0 += ldc;
130 curr_c_1 += ldc;
131 });
132 }
133 };
134
135 } // namespace
136
137 // This is a general but naive implementation based on vector instructions
138 template <typename scalar_t, int64_t head_dim, typename kv_cache_scalar_t>
139 class AttentionImpl<ISA::VEC, scalar_t, head_dim, kv_cache_scalar_t> {
140 static constexpr bool fp8_kv =
141 std::is_same_v<kv_cache_scalar_t, c10::Float8_e4m3fn> ||
142 std::is_same_v<kv_cache_scalar_t, c10::Float8_e5m2>;
143
144 public:
145 using query_t = scalar_t;
146 using q_buffer_t = float;
147 using kv_cache_t = kv_cache_scalar_t;
148 using logits_buffer_t = float;
149 using partial_output_buffer_t = float;
150 using prob_buffer_t = float;
151
152 constexpr static int64_t BlockSizeAlignment =
153 32; // KV token num unit of QK and PV phases
154 constexpr static int64_t HeadDimAlignment =
155 32; // headdim num unit of PV phase
156 constexpr static int64_t MaxQHeadNumPerIteration = 8;
157 constexpr static int64_t HeadDim = head_dim;
158 constexpr static ISA ISAType = ISA::VEC;
159 constexpr static bool scale_on_logits = fp8_kv;
160
161 float k_scale = 1.0f;
162 float v_scale = 1.0f;
163
164 public:
165 void init_from_input(const AttentionInput* input) {
166 if constexpr (fp8_kv) {
167 k_scale = input->k_scale_fp8;
168 v_scale = input->v_scale_fp8;
169 }
170 }
171
172 float get_output_v_scale() const noexcept {
173 if constexpr (fp8_kv) {
174 // VEC dequant unpacks FP8 into a pseudo-FP16 layout (exponent bias 15).
175 // E4M3 (bias=7) needs correction 2^(15-7) = 2^8; E5M2 bias matches FP16
176 // so no correction.
177 if constexpr (std::is_same_v<kv_cache_t, c10::Float8_e5m2>) {
178 return v_scale;
179 } else {
180 return v_scale * 0x1p8f;
181 }
182 }
183 return 1.0f;
184 }
185
186 template <template <typename tile_gemm_t> typename attention>
187 FORCE_INLINE void execute_attention(DEFINE_CPU_ATTENTION_PARAMS) {
188 if constexpr (fp8_kv) {
189 // Same bias correction as get_output_v_scale: VEC FP8→pseudo-FP16 dequant
190 // uses bias 15; E4M3 (bias=7) needs ×2^8, E5M2 (bias=15) needs no
191 // correction.
192 if constexpr (std::is_same_v<kv_cache_t, c10::Float8_e5m2>) {
193 scale *= k_scale;
194 } else {
195 scale *= k_scale * 0x1p8f;
196 }
197 }
198 attention<TileGemm82<kv_cache_t>> attention_iteration;
199 attention_iteration(CPU_ATTENTION_PARAMS);
200 }
201
202 // k_cache_token_group_stride: stride of K cache when move to next
203 // BlockSizeAlignment tokens in a block
204 constexpr static int64_t k_cache_token_group_stride(
205 const int32_t block_size) {
206 return BlockSizeAlignment; // layout of k_cache block is [head_dim,
207 // block_size], row-major
208 }
209
210 // v_cache_token_group_stride: stride of V cache when move to next
211 // BlockSizeAlignment tokens in a block
212 constexpr static int64_t v_cache_token_group_stride(
213 const int32_t block_size) {
214 return head_dim * BlockSizeAlignment; // layout of v_cache is [block_size,
215 // head_dim], row-major
216 }
217
218 // v_cache_head_group_stride: stride of V cache when move to next
219 // HeadDimAlignment head dims in a block
220 constexpr static int64_t v_cache_head_group_stride(const int32_t block_size) {
221 return HeadDimAlignment; // layout of v_cache is [block_size, head_dim],
222 // row-major
223 }
224
225 // Copy q to q_buffer and cast it to fp32.
226 // FP8: QK scale is folded into execute_attention; copy Q unscaled here.
227 void copy_q_heads_tile(scalar_t* __restrict__ src,
228 float* __restrict__ q_buffer, const int32_t q_num,
229 const int32_t q_heads_per_kv,
230 const int64_t q_num_stride,
231 const int64_t q_head_stride, float scale) {
232 static_assert(head_dim % 16 == 0);
233 constexpr int32_t unroll_size = head_dim / 16;
234 using load_vec_t = typename VecTypeTrait<scalar_t>::vec_t;
235
236 const float effective_scale = fp8_kv ? 1.0f : scale;
237 vec_op::FP32Vec16 scale_vec(effective_scale);
238 for (int32_t q_num_idx = 0; q_num_idx < q_num; ++q_num_idx) {
239 for (int32_t q_head_idx = 0; q_head_idx < q_heads_per_kv; ++q_head_idx) {
240 scalar_t* __restrict__ curr_q =
241 src + q_num_idx * q_num_stride + q_head_idx * q_head_stride;
242 float* __restrict__ curr_q_buffer =
243 q_buffer + q_num_idx * q_heads_per_kv * head_dim +
244 q_head_idx * head_dim;
245
246 vec_op::unroll_loop<int32_t, unroll_size>([&](int32_t i) {
247 load_vec_t vec(curr_q);
248 vec_op::FP32Vec16 fp32_vec(vec);
249 fp32_vec = fp32_vec * scale_vec;
250 fp32_vec.save(curr_q_buffer);
251
252 curr_q += 16;
253 curr_q_buffer += 16;
254 });
255 }
256 }
257 }
258
259 // reshape K as column-major and V as row-major
260 static void reshape_and_cache(
261 const scalar_t* __restrict__ key, const scalar_t* __restrict__ value,
262 kv_cache_t* __restrict__ key_cache, kv_cache_t* __restrict__ value_cache,
263 const int64_t* __restrict__ slot_mapping, const int64_t token_num,
264 const int64_t key_token_num_stride, const int64_t value_token_num_stride,
265 const int64_t head_num, const int64_t key_head_num_stride,
266 const int64_t value_head_num_stride, const int64_t num_blocks,
267 const int64_t num_blocks_stride, const int64_t cache_head_num_stride,
268 const int64_t block_size, const int64_t block_size_stride,
269 const float k_inv = 0.0f, const float v_inv = 0.0f) {
270 if constexpr (fp8_kv) {
271 constexpr auto qfn = select_fp8_quant_fn<kv_cache_t>();
272 reshape_and_cache_fp8_vec_impl<scalar_t, qfn>(
273 key, value, reinterpret_cast<uint8_t*>(key_cache),
274 reinterpret_cast<uint8_t*>(value_cache), slot_mapping, token_num,
275 head_num, head_dim, block_size, key_token_num_stride,
276 key_head_num_stride, value_token_num_stride, value_head_num_stride,
277 num_blocks_stride, cache_head_num_stride, num_blocks_stride,
278 cache_head_num_stride, k_inv, v_inv);
279 return;
280 }
281
282 #pragma omp parallel for collapse(2)
283 for (int64_t token_idx = 0; token_idx < token_num; ++token_idx) {
284 for (int64_t head_idx = 0; head_idx < head_num; ++head_idx) {
285 const int64_t pos = slot_mapping[token_idx];
286 if (pos < 0) {
287 // skip
288 continue;
289 }
290
291 const int64_t block_idx = pos / block_size;
292 const int64_t block_offset = pos % block_size;
293 {
294 // Write Key as column-major
295 const scalar_t* key_start_ptr = key +
296 token_idx * key_token_num_stride +
297 head_idx * key_head_num_stride;
298 scalar_t* key_cache_start_ptr =
299 reinterpret_cast<scalar_t*>(key_cache) +
300 block_idx * num_blocks_stride + head_idx * cache_head_num_stride +
301 block_offset;
302
303 #pragma GCC unroll 8
304 for (int64_t i = 0, j = 0; i < head_dim; ++i, j += block_size) {
305 key_cache_start_ptr[j] = key_start_ptr[i];
306 }
307 }
308 {
309 // Write Value as row-major
310 const scalar_t* value_start_ptr = value +
311 token_idx * value_token_num_stride +
312 head_idx * value_head_num_stride;
313 scalar_t* value_cache_start_ptr =
314 reinterpret_cast<scalar_t*>(value_cache) +
315 block_idx * num_blocks_stride + head_idx * cache_head_num_stride +
316 block_offset * head_dim;
317 std::memcpy(value_cache_start_ptr, value_start_ptr,
318 sizeof(scalar_t) * head_dim);
319 }
320 }
321 }
322 }
323 };
324
325 } // namespace cpu_attention
326
327 #endif