| 1 | // Adapted from |
| 2 | // https://github.com/sgl-project/sglang/tree/main/sgl-kernel/csrc/cpu |
| 3 | |
| 4 | // clang-format off |
| 5 | |
| 6 | #include "common.h" |
| 7 | #include "flash_attn.h" |
| 8 | #include "gemm.h" |
| 9 | |
| 10 | namespace { |
| 11 | |
| 12 | // [NOTE]: extend attention for CPU |
| 13 | // 1. BLOCK_M and BLOCK_N tuned for various seq lengths |
| 14 | // 2. can handle non-contiguous k_extend and v_extend |
| 15 | // 3. computes attention for prefix and extend separately |
| 16 | // 4. TODO: apply head dimension blocking to optimize GQA |
| 17 | // 5. optional tree mask for speculative decoding TARGET_VERIFY (EAGLE topk > 1): |
| 18 | // `tree_mask` is a flat [batches * qlen * qlen] bool tensor in |
| 19 | // TreeMaskMode::QLEN_ONLY layout, where qlen == extend_seq_lens[bs] == |
| 20 | // max_len_extend (uniform across the batch, equal to draft_token_num). |
| 21 | // Row i = query draft token, column j = key draft token; true means query i |
| 22 | // may attend key j (each row marks self + ancestors + root). The committed |
| 23 | // prefix (stage 1) is implicitly fully visible to every draft token, which |
| 24 | // is why the mask only covers the qlen x qlen new-token block; the GPU |
| 25 | // FULL_MASK layout carries the prefix columns explicitly but they are |
| 26 | // all-true for EAGLE. When tree_mask is absent, stage 2 falls back to the |
| 27 | // plain causal mask (correct for non-spec extend and topk == 1 chains). |
| 28 | // |
| 29 | |
| 30 | template <typename scalar_t, typename index_t, int BLOCK_M, int BLOCK_N> |
| 31 | void extend_attention_kernel_impl( |
| 32 | scalar_t* __restrict__ o_extend, |
| 33 | const scalar_t* __restrict__ q_extend, |
| 34 | const scalar_t* __restrict__ k_extend, |
| 35 | const scalar_t* __restrict__ v_extend, |
| 36 | const scalar_t* __restrict__ k_buffer, |
| 37 | const scalar_t* __restrict__ v_buffer, |
| 38 | const index_t* __restrict__ req_to_token, |
| 39 | const int64_t* __restrict__ req_pool_indices, |
| 40 | const int64_t* __restrict__ seq_lens, |
| 41 | const int64_t* __restrict__ encoder_lens, |
| 42 | const index_t* __restrict__ extend_seq_lens, |
| 43 | const index_t* __restrict__ extend_start_loc, |
| 44 | const void* __restrict__ buffer, |
| 45 | const scalar_t* __restrict__ sinks, |
| 46 | const bool* __restrict__ tree_mask, |
| 47 | int batches, |
| 48 | int num_heads, |
| 49 | int num_heads_kv, |
| 50 | int head_size, |
| 51 | int head_size_v, |
| 52 | int q_strideM, |
| 53 | int q_strideH, |
| 54 | int ke_strideN, |
| 55 | int ke_strideH, |
| 56 | int ve_strideN, |
| 57 | int ve_strideH, |
| 58 | int k_strideN, |
| 59 | int k_strideH, |
| 60 | int v_strideN, |
| 61 | int v_strideH, |
| 62 | float sm_scale, |
| 63 | int max_num_reqs, |
| 64 | int max_context_len, |
| 65 | int max_total_num_tokens, |
| 66 | int max_len_extend, |
| 67 | int buffer_size_per_thread, |
| 68 | int64_t sliding_window_size, |
| 69 | bool is_prefix_skipped, |
| 70 | bool is_cross_attn, |
| 71 | bool has_encoder_lens, |
| 72 | bool has_sink) { |
| 73 | // strides |
| 74 | const int o_strideM = num_heads * head_size_v; |
| 75 | const int o_strideH = head_size_v; |
| 76 | |
| 77 | // we use same buffer for packed key and value |
| 78 | const int ldb_tmp = std::max(head_size, head_size_v); |
| 79 | |
| 80 | const int num_groups = num_heads / num_heads_kv; |
| 81 | TORCH_CHECK(num_groups * num_heads_kv == num_heads); |
| 82 | |
| 83 | // number of blocks along M |
| 84 | int MB = div_up(max_len_extend, BLOCK_M); |
| 85 | |
| 86 | // parallel on [batches, num_heads, BM] |
| 87 | at::parallel_for(0, batches * num_heads * MB, 0, [&](int begin, int end) { |
| 88 | int bs{0}, head_id{0}, mb{0}; |
| 89 | data_index_init(begin, bs, batches, head_id, num_heads, mb, MB); |
| 90 | |
| 91 | int tid = at::get_thread_num(); |
| 92 | // s_i: [BLOCK_M, BLOCK_N] |
| 93 | float* __restrict__ s_i = reinterpret_cast<float*>((char*)(buffer) + tid * buffer_size_per_thread); |
| 94 | |
| 95 | // v_prime: [BLOCK_M, head_size_v] |
| 96 | float* __restrict__ v_prime = s_i + BLOCK_M * BLOCK_N; |
| 97 | |
| 98 | // s_delta: [BLOCK_M, BLOCK_N] |
| 99 | scalar_t* __restrict__ s_delta = reinterpret_cast<scalar_t*>(v_prime + BLOCK_M * head_size_v); |
| 100 | |
| 101 | // Btmp: [BLOCK_N, max(head_size, head_size_v)] |
| 102 | scalar_t* __restrict__ Btmp = reinterpret_cast<scalar_t*>(s_delta + BLOCK_M * BLOCK_N); |
| 103 | |
| 104 | // init Btmp just once for each thread to prevent NaN |
| 105 | fill_stub(Btmp, 0.f, BLOCK_N * ldb_tmp); |
| 106 | fill_stub(s_delta, 0.f, BLOCK_M * BLOCK_N); |
| 107 | |
| 108 | alignas(64) float s_prime[BLOCK_M]; |
| 109 | alignas(64) float m_prime[BLOCK_M]; |
| 110 | |
| 111 | for (int i = begin; i < end; ++i) { |
| 112 | // seq_len = prefix + extend |
| 113 | int head_kv_id = head_id / num_groups; |
| 114 | int seq_len = seq_lens[bs]; |
| 115 | int seq_len_extend = extend_seq_lens[bs]; |
| 116 | int seq_len_prefix = seq_len - seq_len_extend; |
| 117 | int seq_extend_start_loc = extend_start_loc[bs]; |
| 118 | |
| 119 | int req_pool_id = req_pool_indices[bs]; |
| 120 | int kv_offset = (has_encoder_lens && (!is_cross_attn)) ? encoder_lens[bs] : 0; |
| 121 | TORCH_CHECK(seq_len_prefix >= 0, "prefix len < 0!"); |
| 122 | TORCH_CHECK(seq_len <= max_context_len, "seq_len out of scope!"); |
| 123 | TORCH_CHECK(req_pool_id < max_num_reqs, "req_pool_id out of scope!"); |
| 124 | |
| 125 | if (is_prefix_skipped) { |
| 126 | TORCH_CHECK(seq_len_prefix == 0, "extend attention: expect seq_len_prefix to be 0, got ", seq_len_prefix); |
| 127 | } |
| 128 | |
| 129 | if (tree_mask != nullptr) { |
| 130 | // QLEN_ONLY layout assumes a uniform qlen across the batch (TARGET_VERIFY) |
| 131 | TORCH_CHECK( |
| 132 | seq_len_extend == max_len_extend, |
| 133 | "extend attention: tree_mask requires uniform extend_seq_lens, got ", |
| 134 | seq_len_extend, |
| 135 | " vs ", |
| 136 | max_len_extend); |
| 137 | } |
| 138 | |
| 139 | // offset and size in MB |
| 140 | int m = mb * BLOCK_M; |
| 141 | int m_size = std::min(BLOCK_M, seq_len_extend - m); |
| 142 | |
| 143 | if (m_size <= 0) { |
| 144 | data_index_step(bs, batches, head_id, num_heads, mb, MB); |
| 145 | continue; |
| 146 | } |
| 147 | |
| 148 | // get query |
| 149 | const scalar_t* __restrict__ q_ptr = q_extend + (seq_extend_start_loc + m) * q_strideM + head_id * q_strideH; |
| 150 | |
| 151 | // init v', s' and m' |
| 152 | fill_stub(v_prime, 0.f, m_size * head_size_v); |
| 153 | fill_stub(s_prime, 0.f, m_size); |
| 154 | fill_stub(m_prime, -std::numeric_limits<scalar_t>::infinity(), m_size); |
| 155 | // stage 1: compute scores with prefix |
| 156 | int kv_start = 0; |
| 157 | int kv_end = is_cross_attn ? encoder_lens[bs] : seq_len_prefix; |
| 158 | for (int n = kv_start; n < kv_end; n += BLOCK_N) { |
| 159 | int n_size = std::min(BLOCK_N, kv_end - n); |
| 160 | |
| 161 | // `n_size` is K in 2nd gemm, pad to TILE_K; |
| 162 | const int padded_n_size = div_up(n_size, TILE_K) * TILE_K; |
| 163 | |
| 164 | // get key and pack |
| 165 | pack_vnni<scalar_t, index_t>( |
| 166 | /* dst */ Btmp, |
| 167 | /* src */ k_buffer + head_kv_id * k_strideH, |
| 168 | /* ind */ req_to_token + req_pool_id * max_context_len + n + kv_offset, |
| 169 | /* N */ n_size, |
| 170 | /* K */ head_size, |
| 171 | /* ld_src */ k_strideN, |
| 172 | /* ld_dst */ BLOCK_N); |
| 173 | |
| 174 | // calculate s_i <- Q @ K |
| 175 | at::native::cpublas::brgemm( |
| 176 | /* M */ m_size, |
| 177 | /* N */ n_size, |
| 178 | /* K */ head_size, |
| 179 | /* lda */ q_strideM, |
| 180 | /* ldb */ BLOCK_N, |
| 181 | /* ldc */ BLOCK_N, |
| 182 | /* add_C */ false, |
| 183 | /* A */ q_ptr, |
| 184 | /* B */ Btmp, |
| 185 | /* C */ s_i); |
| 186 | |
| 187 | for (int row = 0; row < m_size; ++row) { |
| 188 | if (sliding_window_size > 0) { |
| 189 | int last_col = seq_len_prefix + row + m - sliding_window_size + 1; |
| 190 | if (last_col >= n + n_size) { |
| 191 | continue; |
| 192 | } |
| 193 | fill_stub(s_i + row * BLOCK_N, -std::numeric_limits<float>::infinity(), last_col - n); |
| 194 | } |
| 195 | flash_attn_softmax<scalar_t, BLOCK_M, BLOCK_N>::apply( |
| 196 | s_i, s_delta, v_prime, s_prime, m_prime, m_size, n_size, padded_n_size, head_size_v, sm_scale, row); |
| 197 | } |
| 198 | |
| 199 | // get value and pack |
| 200 | pack_vnni2<scalar_t>( |
| 201 | /* dst */ Btmp, |
| 202 | /* src */ v_buffer + head_kv_id * v_strideH, |
| 203 | /* ind */ req_to_token + req_pool_id * max_context_len + n + kv_offset, |
| 204 | /* K */ n_size, |
| 205 | /* N */ head_size_v, |
| 206 | /* ld_src */ v_strideN, |
| 207 | /* ld_dst */ head_size_v); |
| 208 | |
| 209 | // calculate V' <- s_delta @ V + V' |
| 210 | at::native::cpublas::brgemm( |
| 211 | /* M */ m_size, |
| 212 | /* N */ head_size_v, |
| 213 | /* K */ padded_n_size, // n_size |
| 214 | /* lda */ BLOCK_N, |
| 215 | /* ldb */ head_size_v, |
| 216 | /* ldc */ head_size_v, |
| 217 | /* add_C */ true, |
| 218 | /* A */ s_delta, |
| 219 | /* B */ Btmp, |
| 220 | /* C */ v_prime); |
| 221 | } // loop with seq_len_prefix |
| 222 | if (!is_cross_attn) { |
| 223 | // stage 2: compute the triangle part |
| 224 | int num_keys = std::min(seq_len_extend, m + BLOCK_M); |
| 225 | for (int n = 0; n < num_keys; n += BLOCK_N) { |
| 226 | int n_size = std::min(BLOCK_N, num_keys - n); |
| 227 | |
| 228 | // `n_size` is K in 2nd gemm, pad to TILE_K; |
| 229 | const int padded_n_size = div_up(n_size, TILE_K) * TILE_K; |
| 230 | |
| 231 | // get key and pack |
| 232 | pack_vnni<scalar_t>( |
| 233 | /* dst */ Btmp, |
| 234 | /* src */ k_extend + (seq_extend_start_loc + n) * ke_strideN + head_kv_id * ke_strideH, |
| 235 | /* N */ n_size, |
| 236 | /* K */ head_size, |
| 237 | /* ld_src */ ke_strideN, |
| 238 | /* ld_dst */ BLOCK_N); |
| 239 | |
| 240 | // calculate s_i <- Q @ K |
| 241 | at::native::cpublas::brgemm( |
| 242 | /* M */ m_size, |
| 243 | /* N */ n_size, |
| 244 | /* K */ head_size, |
| 245 | /* lda */ q_strideM, |
| 246 | /* ldb */ BLOCK_N, |
| 247 | /* ldc */ BLOCK_N, |
| 248 | /* add_C */ false, |
| 249 | /* A */ q_ptr, |
| 250 | /* B */ Btmp, |
| 251 | /* C */ s_i); |
| 252 | |
| 253 | // apply tree mask (speculative TARGET_VERIFY) or causal mask |
| 254 | if (tree_mask != nullptr) { |
| 255 | // [Note] tree mask for EAGLE topk > 1 (TreeMaskMode::QLEN_ONLY). |
| 256 | // mask[bs][m + row][n + col] == false -> query draft token (m + row) |
| 257 | // may not attend key draft token (n + col); set the score to -inf |
| 258 | // before softmax. The tree mask subsumes the causal constraint: |
| 259 | // ancestors always precede descendants in the draft token ordering, |
| 260 | // so permitted keys satisfy j <= i and the causal `num_keys` bound |
| 261 | // above remains valid. |
| 262 | const bool* __restrict__ mask_base = |
| 263 | tree_mask + (static_cast<int64_t>(bs) * seq_len_extend + m) * seq_len_extend + n; |
| 264 | for (int row = 0; row < m_size; ++row) { |
| 265 | float* __restrict__ row_ptr = s_i + row * BLOCK_N; |
| 266 | const bool* __restrict__ mask_ptr = mask_base + static_cast<int64_t>(row) * seq_len_extend; |
| 267 | for (int col = 0; col < n_size; ++col) { |
| 268 | if (!mask_ptr[col]) { |
| 269 | row_ptr[col] = -std::numeric_limits<float>::infinity(); |
| 270 | } |
| 271 | } |
| 272 | } |
| 273 | } else if (n + n_size - 1 > m) { |
| 274 | // apply causal mask |
| 275 | // [Note] condition to apply causal mask. |
| 276 | // Mask any block whose last key (n + n_size - 1) is strictly after the first query position (m), i.e. n + |
| 277 | // n_size - 1 > m. The original condition was `num_keys - n <= BLOCK_N` (last n-block only). That was |
| 278 | // correct when BLOCK_M <= BLOCK_N/2 because earlier n-blocks were guaranteed to contain only past keys. |
| 279 | // With BLOCK_M=512, BLOCK_N=768: |
| 280 | // BLOCK_M > BLOCK_N/2, so the first n-block can contain future keys. |
| 281 | // Example: m=512 (mb=1), num_keys=1024, first n-block covers keys [0, 768). |
| 282 | // Query row=0 is at position 512, so keys 513..767 are future and must be |
| 283 | // masked — but `num_keys - 0 = 1024 > BLOCK_N` skips masking entirely, |
| 284 | // producing wrong (non-causal) attention for rows 0..254 of this m-block. |
| 285 | for (int row = 0; row < m_size; ++row) { |
| 286 | int last_col = m + row - n; |
| 287 | // [Note] mask the entire row if last_col < 0. |
| 288 | // Clamp to -1: when n > m + row every key in this block is a future |
| 289 | // key, so the entire row should be masked. Without this clamp, |
| 290 | // last_col+1 <= 0 and fill_stub would write before row_ptr. |
| 291 | last_col = std::max(last_col, -1); |
| 292 | // fill [last_col + 1, n_size) to -inf |
| 293 | float* row_ptr = s_i + row * BLOCK_N; |
| 294 | fill_stub(row_ptr + last_col + 1, -std::numeric_limits<float>::infinity(), n_size - last_col - 1); |
| 295 | } |
| 296 | } |
| 297 | |
| 298 | for (int row = 0; row < m_size; ++row) { |
| 299 | if (sliding_window_size > 0 && row + m + 1 >= n + sliding_window_size - 1 && |
| 300 | row + m + 1 < n + sliding_window_size + n_size) { |
| 301 | fill_stub( |
| 302 | s_i + row * BLOCK_N, -std::numeric_limits<float>::infinity(), row + m - n - sliding_window_size + 1); |
| 303 | } else if (sliding_window_size > 0 && row + m + 1 >= n + sliding_window_size) { |
| 304 | continue; |
| 305 | } |
| 306 | flash_attn_softmax<scalar_t, BLOCK_M, BLOCK_N>::apply( |
| 307 | s_i, s_delta, v_prime, s_prime, m_prime, m_size, n_size, padded_n_size, head_size_v, sm_scale, row); |
| 308 | } |
| 309 | |
| 310 | // get value and pack |
| 311 | pack_vnni2<scalar_t>( |
| 312 | /* dst */ Btmp, |
| 313 | /* src */ v_extend + (seq_extend_start_loc + n) * ve_strideN + head_kv_id * ve_strideH, |
| 314 | /* K */ n_size, |
| 315 | /* N */ head_size_v, |
| 316 | /* ld_src */ ve_strideN, |
| 317 | /* ld_dst */ head_size_v); |
| 318 | |
| 319 | // calculate V' <- s_delta @ V + V' |
| 320 | at::native::cpublas::brgemm( |
| 321 | /* M */ m_size, |
| 322 | /* N */ head_size_v, |
| 323 | /* K */ padded_n_size, // n_size |
| 324 | /* lda */ BLOCK_N, |
| 325 | /* ldb */ head_size_v, |
| 326 | /* ldc */ head_size_v, |
| 327 | /* add_C */ true, |
| 328 | /* A */ s_delta, |
| 329 | /* B */ Btmp, |
| 330 | /* C */ v_prime); |
| 331 | } // loop with seq_len_extend |
| 332 | } |
| 333 | scalar_t* __restrict__ out_ptr = o_extend + (seq_extend_start_loc + m) * o_strideM + head_id * o_strideH; |
| 334 | for (int row = 0; row < m_size; ++row) { |
| 335 | if (has_sink) { |
| 336 | s_prime[row] += std::exp(sinks[head_id] - m_prime[row]); |
| 337 | } |
| 338 | float s = 1 / s_prime[row]; |
| 339 | copy_stub<scalar_t>(out_ptr + row * o_strideM, v_prime + row * head_size_v, s, head_size_v); |
| 340 | } |
| 341 | |
| 342 | // move to the next index |
| 343 | data_index_step(bs, batches, head_id, num_heads, mb, MB); |
| 344 | } |
| 345 | at::native::cpublas::brgemm_release(); |
| 346 | }); |
| 347 | } |
| 348 | |
| 349 | } // anonymous namespace |
| 350 | |
| 351 | template <int BLOCK_M, int BLOCK_N> |
| 352 | inline int resize_buffer(at::Tensor& buffer, int num_threads, int head_size, int head_size_v) { |
| 353 | static_assert(BLOCK_M <= BLOCK_N, "Make sure BLOCK_M <= BLOCK_N to prevent buffer overflows during causal masking"); |
| 354 | const int size_per_thread = |
| 355 | /* s_i */ BLOCK_M * BLOCK_N * sizeof(float) + |
| 356 | /* v_prime */ BLOCK_M * head_size_v * sizeof(float) + |
| 357 | /* s_delta */ BLOCK_M * BLOCK_N * sizeof(uint16_t) + |
| 358 | /* Btmp */ BLOCK_N * std::max(head_size, head_size_v) * sizeof(uint16_t); |
| 359 | |
| 360 | buffer.resize_({num_threads, size_per_thread}); |
| 361 | return size_per_thread; |
| 362 | } |
| 363 | |
| 364 | #define LAUNCH_EXTEND_ATTENTION_KERNEL(BLOCK_M, BLOCK_N) \ |
| 365 | do { \ |
| 366 | int sz = resize_buffer<BLOCK_M, BLOCK_N>(buffer, num_threads, head_size, head_size_v); \ |
| 367 | \ |
| 368 | extend_attention_kernel_impl<scalar_t, index_t, BLOCK_M, BLOCK_N>( \ |
| 369 | o_extend.data_ptr<scalar_t>(), \ |
| 370 | q_extend.data_ptr<scalar_t>(), \ |
| 371 | k_extend.data_ptr<scalar_t>(), \ |
| 372 | v_extend.data_ptr<scalar_t>(), \ |
| 373 | k_buffer.data_ptr<scalar_t>(), \ |
| 374 | v_buffer.data_ptr<scalar_t>(), \ |
| 375 | req_to_token.data_ptr<index_t>(), \ |
| 376 | req_pool_indices.data_ptr<int64_t>(), \ |
| 377 | seq_lens.data_ptr<int64_t>(), \ |
| 378 | encoder_lens_t.data_ptr<int64_t>(), \ |
| 379 | extend_seq_lens.data_ptr<index_t>(), \ |
| 380 | extend_start_loc.data_ptr<index_t>(), \ |
| 381 | buffer.data_ptr(), \ |
| 382 | sinks_tensor.data_ptr<scalar_t>(), \ |
| 383 | tree_mask_ptr, \ |
| 384 | num_seqs, \ |
| 385 | num_heads, \ |
| 386 | num_heads_kv, \ |
| 387 | head_size, \ |
| 388 | head_size_v, \ |
| 389 | q_strideM, \ |
| 390 | q_strideH, \ |
| 391 | ke_strideN, \ |
| 392 | ke_strideH, \ |
| 393 | ve_strideN, \ |
| 394 | ve_strideH, \ |
| 395 | k_strideN, \ |
| 396 | k_strideH, \ |
| 397 | v_strideN, \ |
| 398 | v_strideH, \ |
| 399 | sm_scale, \ |
| 400 | max_num_reqs, \ |
| 401 | max_context_len, \ |
| 402 | max_total_num_tokens, \ |
| 403 | max_len_extend, \ |
| 404 | sz, \ |
| 405 | sliding_window_size, \ |
| 406 | is_prefix_skipped, \ |
| 407 | is_cross_attn, \ |
| 408 | has_encoder_lens, \ |
| 409 | has_sink); \ |
| 410 | } while (0) |
| 411 | |
| 412 | // q_extend, k_extend, v_extend, o_extend: contiguous tensors |
| 413 | // k_buffer, v_buffer: (prefix + extend) tensors in mem_manager |
| 414 | // |
| 415 | // q_extend: [num_tokens, num_heads, head_size] |
| 416 | // k_extend: [num_extend_tokens, num_heads, head_size] |
| 417 | // v_extend: [num_extend_tokens, num_heads, head_size] |
| 418 | // o_extend: [num_tokens, num_heads, head_size] |
| 419 | // k_buffer: [max_total_num_tokens, num_heads, head_size] |
| 420 | // v_buffer: [max_total_num_tokens, num_heads, head_size] |
| 421 | // req_to_token: [max_num_reqs, max_context_len] int32 or int64 |
| 422 | // req_pool_indices: [num_seqs] int64 |
| 423 | // seq_lens: [num_seqs] int64 |
| 424 | // extend_seq_lens: [num_seqs] |
| 425 | // extend_start_loc: [num_seqs] |
| 426 | // encoder_lens: [num_seqs] int64 or None |
| 427 | // sinks: [num_heads] or None |
| 428 | // tree_mask: [num_seqs * max_len_extend * max_len_extend] bool or None |
| 429 | // TreeMaskMode::QLEN_ONLY tree mask for speculative TARGET_VERIFY; see [NOTE] 5 above. |
| 430 | void extend_attention_cpu( |
| 431 | at::Tensor& q_extend, |
| 432 | const std::optional<at::Tensor>& k_extend_opt, |
| 433 | const std::optional<at::Tensor>& v_extend_opt, |
| 434 | at::Tensor& o_extend, |
| 435 | at::Tensor& k_buffer, |
| 436 | at::Tensor& v_buffer, |
| 437 | at::Tensor& req_to_token, |
| 438 | at::Tensor& req_pool_indices, |
| 439 | at::Tensor& seq_lens, |
| 440 | at::Tensor& extend_seq_lens, |
| 441 | at::Tensor& extend_start_loc, |
| 442 | int64_t max_len_extend, |
| 443 | double sm_scale, |
| 444 | double logit_cap, |
| 445 | bool is_cross_attn, |
| 446 | int64_t sliding_window_size, |
| 447 | std::optional<at::Tensor> encoder_lens, |
| 448 | std::optional<at::Tensor> sinks, |
| 449 | std::optional<at::Tensor> tree_mask) { |
| 450 | if (!is_cross_attn) { |
| 451 | TORCH_CHECK( |
| 452 | k_extend_opt.has_value() && v_extend_opt.has_value(), |
| 453 | "k_extend and v_extend are required for non-cross attention"); |
| 454 | } |
| 455 | // Since k_extend and v_extend are not used for cross attention, they can be initialized as k_buffer and v_buffer |
| 456 | // here. |
| 457 | auto k_extend = k_extend_opt.has_value() ? k_extend_opt.value() : k_buffer; |
| 458 | auto v_extend = v_extend_opt.has_value() ? v_extend_opt.value() : v_buffer; |
| 459 | |
| 460 | CHECK_LAST_DIM_CONTIGUOUS_INPUT(q_extend); |
| 461 | CHECK_INPUT(o_extend); |
| 462 | CHECK_LAST_DIM_CONTIGUOUS_INPUT(k_extend); |
| 463 | CHECK_LAST_DIM_CONTIGUOUS_INPUT(v_extend); |
| 464 | CHECK_LAST_DIM_CONTIGUOUS_INPUT(k_buffer); |
| 465 | CHECK_LAST_DIM_CONTIGUOUS_INPUT(v_buffer); |
| 466 | |
| 467 | int num_seqs = seq_lens.size(0); |
| 468 | int max_num_reqs = req_to_token.size(0); |
| 469 | int max_context_len = req_to_token.size(1); |
| 470 | int max_total_num_tokens = k_buffer.size(0); |
| 471 | |
| 472 | int num_heads = q_extend.size(1); |
| 473 | int num_heads_kv = k_extend.size(1); |
| 474 | int head_size = q_extend.size(2); |
| 475 | int head_size_v = v_extend.size(2); |
| 476 | |
| 477 | // strides for q_extend, k_extend and v_extend |
| 478 | int q_strideM = q_extend.stride(0); |
| 479 | int q_strideH = q_extend.stride(1); |
| 480 | int ke_strideN = k_extend.stride(0); |
| 481 | int ke_strideH = k_extend.stride(1); |
| 482 | int ve_strideN = v_extend.stride(0); |
| 483 | int ve_strideH = v_extend.stride(1); |
| 484 | |
| 485 | // strides for k_buffer and v_buffer |
| 486 | int k_strideN = k_buffer.stride(0); |
| 487 | int k_strideH = k_buffer.stride(1); |
| 488 | int v_strideN = v_buffer.stride(0); |
| 489 | int v_strideH = v_buffer.stride(1); |
| 490 | |
| 491 | // check sizes |
| 492 | CHECK_EQ(req_pool_indices.size(0), num_seqs); |
| 493 | CHECK_EQ(extend_seq_lens.size(0), num_seqs); |
| 494 | CHECK_EQ(extend_start_loc.size(0), num_seqs); |
| 495 | CHECK_EQ(v_extend.size(1), num_heads_kv); |
| 496 | CHECK_EQ(k_buffer.size(1), v_buffer.size(1)); |
| 497 | |
| 498 | // MLA will skip prefix part |
| 499 | const bool is_prefix_skipped = k_buffer.size(1) != num_heads_kv; |
| 500 | |
| 501 | // check index data types |
| 502 | const auto index_dtype = req_to_token.scalar_type(); |
| 503 | TORCH_CHECK( |
| 504 | index_dtype == at::kInt || index_dtype == at::kLong, |
| 505 | "extend: expect req_to_token to be int32 or int64, got ", |
| 506 | index_dtype); |
| 507 | TORCH_CHECK(seq_lens.scalar_type() == at::kLong, "extend: expect req_lens to be int64, got ", seq_lens.scalar_type()); |
| 508 | TORCH_CHECK( |
| 509 | req_pool_indices.scalar_type() == at::kLong, |
| 510 | "extend: expect req_pool_indices to be int64, got ", |
| 511 | req_pool_indices.scalar_type()); |
| 512 | TORCH_CHECK( |
| 513 | extend_seq_lens.scalar_type() == index_dtype && extend_start_loc.scalar_type() == index_dtype, |
| 514 | "extend: expect extend_seq_lens and extend_start_loc to have same dtype as req_to_token."); |
| 515 | |
| 516 | // D and DV need to be 32x as we transpose by 512-bit |
| 517 | TORCH_CHECK(head_size % 32 == 0, "invalid head_size ", head_size); |
| 518 | TORCH_CHECK(head_size_v % 32 == 0, "invalid head_size_v ", head_size_v); |
| 519 | |
| 520 | int num_threads = at::get_num_threads(); |
| 521 | auto buffer = at::empty({}, q_extend.options().dtype(at::kChar)); |
| 522 | |
| 523 | bool has_encoder_lens = encoder_lens.has_value(); |
| 524 | // Since encoder_lens is not used when it is None, encoder_lens_t can be initialized as any tensor of int64_t dtype. |
| 525 | at::Tensor encoder_lens_t = seq_lens; |
| 526 | if (has_encoder_lens) { |
| 527 | encoder_lens_t = encoder_lens.value(); |
| 528 | CHECK_EQ(encoder_lens_t.size(0), num_seqs); |
| 529 | } |
| 530 | bool has_sink = sinks.has_value(); |
| 531 | at::Tensor sinks_tensor = has_sink ? sinks.value() : at::empty({num_heads}, q_extend.options()); |
| 532 | CHECK_DIM(1, sinks_tensor); |
| 533 | CHECK_EQ(sinks_tensor.size(0), num_heads); |
| 534 | |
| 535 | const bool* tree_mask_ptr = nullptr; |
| 536 | if (tree_mask.has_value()) { |
| 537 | const at::Tensor& tree_mask_t = tree_mask.value(); |
| 538 | CHECK_INPUT(tree_mask_t); |
| 539 | TORCH_CHECK( |
| 540 | tree_mask_t.scalar_type() == at::kBool, "extend: expect tree_mask to be bool, got ", tree_mask_t.scalar_type()); |
| 541 | TORCH_CHECK( |
| 542 | tree_mask_t.numel() == static_cast<int64_t>(num_seqs) * max_len_extend * max_len_extend, |
| 543 | "extend: expect tree_mask numel to be num_seqs * max_len_extend^2 = ", |
| 544 | static_cast<int64_t>(num_seqs) * max_len_extend * max_len_extend, |
| 545 | ", got ", |
| 546 | tree_mask_t.numel()); |
| 547 | TORCH_CHECK(!is_cross_attn, "extend: tree_mask is not supported for cross attention"); |
| 548 | // The window mask derives query positions from the row index |
| 549 | // (seq_len_prefix + m + row), but tree-mask rows sit at their tree depth, |
| 550 | // which is <= the row index; combining the two would over-mask the prefix. |
| 551 | TORCH_CHECK(sliding_window_size <= 0, "extend: tree_mask is not supported with sliding window attention"); |
| 552 | tree_mask_ptr = tree_mask_t.data_ptr<bool>(); |
| 553 | } |
| 554 | |
| 555 | AT_DISPATCH_REDUCED_FLOATING_TYPES(q_extend.scalar_type(), "extend_attention_kernel", [&] { |
| 556 | AT_DISPATCH_INDEX_TYPES(index_dtype, "extend_attention_indices", [&] { |
| 557 | if (max_len_extend <= 256) { |
| 558 | LAUNCH_EXTEND_ATTENTION_KERNEL(32, 64); |
| 559 | } else if (max_len_extend <= 1024) { |
| 560 | LAUNCH_EXTEND_ATTENTION_KERNEL(128, 256); |
| 561 | } else if (max_len_extend <= 4096) { |
| 562 | LAUNCH_EXTEND_ATTENTION_KERNEL(256, 768); |
| 563 | } else { // max_len_extend > 4096 |
| 564 | LAUNCH_EXTEND_ATTENTION_KERNEL(512, 768); |
| 565 | } |
| 566 | }); |
| 567 | }); |
| 568 | } |