@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
hpp 171 lines 5.66 KB
Raw
1 #ifndef CPU_ATTN_VEC16_HPP
2 #define CPU_ATTN_VEC16_HPP
3
4 #include "cpu_attn_vec.hpp"
5
6 namespace cpu_attention {
7
8 namespace {
9 // 16-1-16 pattern, 16 regs for A, 1 regs for B, 16 regs for C, [16, K] @ [k,
10 // 16]
11 template <typename kv_cache_t>
12 class TileGemm161 {
13 public:
14 template <AttentionGemmPhase phase, int32_t k_size>
15 FORCE_INLINE static void gemm(const int32_t m_size,
16 float* __restrict__ a_tile,
17 kv_cache_t* __restrict__ b_tile,
18 float* __restrict__ c_tile, const int64_t lda,
19 const int64_t ldb, const int64_t ldc,
20 const int32_t block_size,
21 const int32_t dynamic_k_size,
22 const bool accum_c) {
23 switch (m_size) {
24 case 1:
25 gemm_micro<1>(a_tile, b_tile, c_tile, lda, ldb, ldc, block_size,
26 dynamic_k_size, accum_c);
27 break;
28 case 2:
29 gemm_micro<2>(a_tile, b_tile, c_tile, lda, ldb, ldc, block_size,
30 dynamic_k_size, accum_c);
31 break;
32 case 3:
33 case 4:
34 gemm_micro<4>(a_tile, b_tile, c_tile, lda, ldb, ldc, block_size,
35 dynamic_k_size, accum_c);
36 break;
37 case 5:
38 case 6:
39 gemm_micro<6>(a_tile, b_tile, c_tile, lda, ldb, ldc, block_size,
40 dynamic_k_size, accum_c);
41 break;
42 case 7:
43 case 8:
44 gemm_micro<8>(a_tile, b_tile, c_tile, lda, ldb, ldc, block_size,
45 dynamic_k_size, accum_c);
46 break;
47 case 9:
48 case 10:
49 case 11:
50 case 12:
51 gemm_micro<12>(a_tile, b_tile, c_tile, lda, ldb, ldc, block_size,
52 dynamic_k_size, accum_c);
53 break;
54 case 13:
55 case 14:
56 case 15:
57 case 16:
58 gemm_micro<16>(a_tile, b_tile, c_tile, lda, ldb, ldc, block_size,
59 dynamic_k_size, accum_c);
60 break;
61 }
62 }
63
64 template <int32_t M>
65 static void gemm_micro(float* __restrict__ a_tile,
66 kv_cache_t* __restrict__ b_tile,
67 float* __restrict__ c_tile, const int64_t lda,
68 const int64_t ldb, const int64_t ldc,
69 const int32_t block_size, const int32_t dynamic_k_size,
70 const bool accum_c) {
71 static_assert(0 < M && M <= 16);
72 using load_vec_t = typename VecTypeTrait<kv_cache_t>::vec_t;
73
74 kv_cache_t* __restrict__ curr_b_0 = b_tile;
75 float* __restrict__ curr_c_0 = c_tile;
76
77 vec_op::FP32Vec16 c_regs[M];
78 if (accum_c) {
79 float* __restrict__ curr_m_c_0 = curr_c_0;
80 vec_op::unroll_loop<int32_t, M>([&](int32_t i) {
81 c_regs[i] = vec_op::FP32Vec16(curr_m_c_0);
82
83 // update
84 curr_m_c_0 += ldc;
85 });
86 }
87
88 float* __restrict__ curr_a = a_tile;
89 for (int32_t k = 0; k < dynamic_k_size; ++k) {
90 load_vec_t b_0_reg(curr_b_0);
91 vec_op::FP32Vec16 fp32_b_0_reg(b_0_reg);
92
93 float* __restrict__ curr_m_a = curr_a;
94 vec_op::unroll_loop<int32_t, M>([&](int32_t i) {
95 float v = *curr_m_a;
96 vec_op::FP32Vec16 a_reg(v);
97 c_regs[i] = c_regs[i] + a_reg * fp32_b_0_reg;
98
99 // update
100 curr_m_a += lda;
101 });
102
103 // update
104 curr_a += 1;
105 curr_b_0 += ldb;
106 }
107
108 vec_op::unroll_loop<int32_t, M>([&](int32_t i) {
109 c_regs[i].save(curr_c_0);
110
111 // update
112 curr_c_0 += ldc;
113 });
114 }
115 };
116 } // namespace
117
118 // This is a general but naive implementation based on vector instructions
119 template <typename scalar_t, int64_t head_dim, typename kv_cache_scalar_t>
120 class AttentionImpl<ISA::VEC16, scalar_t, head_dim, kv_cache_scalar_t>
121 : public AttentionImpl<ISA::VEC, scalar_t, head_dim, kv_cache_scalar_t> {
122 public:
123 using query_t = scalar_t;
124 using q_buffer_t = float;
125 using kv_cache_t = scalar_t;
126 using logits_buffer_t = float;
127 using partial_output_buffer_t = float;
128 using prob_buffer_t = float;
129
130 constexpr static int64_t BlockSizeAlignment =
131 16; // KV token num unit of QK and PV phases
132 constexpr static int64_t HeadDimAlignment =
133 16; // headdim num unit of PV phase
134 constexpr static int64_t MaxQHeadNumPerIteration = 16;
135 constexpr static int64_t HeadDim = head_dim;
136 constexpr static ISA ISAType = ISA::VEC16;
137 constexpr static bool scale_on_logits = false; // apply scale on q_buffer
138
139 public:
140 template <template <typename tile_gemm_t> typename attention>
141 FORCE_INLINE void execute_attention(DEFINE_CPU_ATTENTION_PARAMS) {
142 attention<TileGemm161<kv_cache_t>> attention_iteration;
143 attention_iteration(CPU_ATTENTION_PARAMS);
144 }
145
146 // k_cache_token_group_stride: stride of K cache when move to next
147 // BlockSizeAlignment tokens in a block
148 constexpr static int64_t k_cache_token_group_stride(
149 const int32_t block_size) {
150 return BlockSizeAlignment; // layout of k_cache block is [head_dim,
151 // block_size], row-major
152 }
153
154 // v_cache_token_group_stride: stride of V cache when move to next
155 // BlockSizeAlignment tokens in a block
156 constexpr static int64_t v_cache_token_group_stride(
157 const int32_t block_size) {
158 return head_dim * BlockSizeAlignment; // layout of v_cache is [block_size,
159 // head_dim], row-major
160 }
161
162 // v_cache_head_group_stride: stride of V cache when move to next
163 // HeadDimAlignment head dims in a block
164 constexpr static int64_t v_cache_head_group_stride(const int32_t block_size) {
165 return HeadDimAlignment; // layout of v_cache is [block_size, head_dim],
166 // row-major
167 }
168 };
169 } // namespace cpu_attention
170
171 #endif