| 1 | // Adapted from |
| 2 | // https://github.com/sgl-project/sglang/blob/main/python/sglang/kernels/aot/csrc/cpu/flash_mla.cpp |
| 3 | // |
| 4 | // DeepSeek-V4 sparse MQA attention kernel for the fp8_ds_mla cache layout. |
| 5 | // Near-verbatim port of `sparse_mla_decode_kernel_impl` (VNNI K/V packing + |
| 6 | // AMX brgemm QK^T/PV + online softmax); the `V4_FP8Sparse` byte layout |
| 7 | // matches vLLM's fp8_ds_mla layout exactly (see store_cache.cpp). |
| 8 | // |
| 9 | // Differences from upstream: |
| 10 | // - Outer shape: upstream takes a padded batch-major query `[B, S_q, H, D]` |
| 11 | // with per-batch topk sets; vLLM passes flat `[num_tokens, H, D]` with |
| 12 | // per-token resolved indices. The kernel's own `at::parallel_for` already |
| 13 | // flattens `(batch, s_q)` internally, so this is called with |
| 14 | // `batches = num_tokens, s_q = 1` -- no padding, no kernel-body changes. |
| 15 | // - Two physically distinct paged caches (SWA "window" + compressor |
| 16 | // "compressed") are addressed via block_table-resolved flat slot ids, |
| 17 | // matching upstream's `k_main`/`k_extra` two-index-set convention. |
| 18 | // - `attn_sink` uses upstream's own post-hoc LSE correction |
| 19 | // (`out *= 1 / (1 + exp(sink - lse))`). |
| 20 | // - bf16 query only (matching upstream's own restriction: its FP8->bf16 |
| 21 | // dequant intrinsics are bf16-specific); fp16 is not supported here. |
| 22 | // |
| 23 | // clang-format off |
| 24 | |
| 25 | #include <algorithm> |
| 26 | #include <cmath> |
| 27 | #include <limits> |
| 28 | #include <vector> |
| 29 | |
| 30 | #include "common.h" |
| 31 | #include "gemm.h" |
| 32 | #include "vec.h" |
| 33 | |
| 34 | #if !defined(CPU_CAPABILITY_AVX512) |
| 35 | static_assert(false, |
| 36 | "flash_mla.cpp is only ever compiled into the AVX512+AMX-flagged " |
| 37 | "_C target (see cmake/cpu_extension.cmake's VLLM_EXT_SRC_SGL) and " |
| 38 | "relies on that unconditionally -- CPU_CAPABILITY_AVX512 must be " |
| 39 | "defined here."); |
| 40 | #endif |
| 41 | |
| 42 | namespace { |
| 43 | |
| 44 | constexpr int64_t kNopeDim = 448; |
| 45 | constexpr int64_t kRopeDim = 64; |
| 46 | constexpr int64_t kHeadDim = kNopeDim + kRopeDim; |
| 47 | constexpr int64_t kQuantBlock = 64; |
| 48 | constexpr int64_t kNumQuantBlocks = kNopeDim / kQuantBlock; |
| 49 | constexpr int64_t kScaleBytesPerToken = kNumQuantBlocks + 1; |
| 50 | constexpr int64_t kTokenDataBytes = kNopeDim + kRopeDim * 2; |
| 51 | |
| 52 | // --------------------------------------------------------------------------- |
| 53 | // V4_FP8Sparse layout (SGLang's `flash_mla.cpp`) -- byte-for-byte identical |
| 54 | // to vLLM's own fp8_ds_mla layout above (d=512, d_nope=448, d_rope=64, |
| 55 | // tile=64, num_tiles=7): |
| 56 | // per block: |
| 57 | // [block_size * (d_nope + 2*d_rope) bytes ; FP8 NoPE + bf16 RoPE per token] |
| 58 | // [block_size * 8 bytes ; 7 e8m0 scales + 1 pad byte/token] |
| 59 | // --------------------------------------------------------------------------- |
| 60 | |
| 61 | struct FP8LayoutMeta { |
| 62 | int64_t d; |
| 63 | int64_t d_nope; |
| 64 | int64_t d_rope; |
| 65 | int64_t tile_size; |
| 66 | int64_t num_tiles; |
| 67 | }; |
| 68 | |
| 69 | constexpr FP8LayoutMeta kV4Meta = {kHeadDim, kNopeDim, kRopeDim, kQuantBlock, kNumQuantBlocks}; |
| 70 | |
| 71 | // Convert one fp8_e8m0 byte (= unsigned 8-bit exponent) to float. |
| 72 | // e8m0 only stores an exponent (bias 127); value = 2^(e - 127), with 0xFF |
| 73 | // reserved for NaN. |
| 74 | inline float fp8_e8m0_to_float(uint8_t v) { |
| 75 | if (v == 0xFF) return std::numeric_limits<float>::quiet_NaN(); |
| 76 | if (v == 0) return 0.f; |
| 77 | union { |
| 78 | uint32_t u; |
| 79 | float f; |
| 80 | } u; |
| 81 | u.u = static_cast<uint32_t>(v) << 23; |
| 82 | return u.f; |
| 83 | } |
| 84 | |
| 85 | template <typename index_t> |
| 86 | inline bool is_valid_sparse_index(index_t idx, int64_t pos, int64_t topk_limit, int64_t total_tokens) { |
| 87 | const int64_t v = static_cast<int64_t>(idx); |
| 88 | return pos < topk_limit && v >= 0 && v < total_tokens; |
| 89 | } |
| 90 | |
| 91 | #if defined(CPU_CAPABILITY_AVX512) |
| 92 | inline __attribute__((always_inline)) __m512i cvt_fp8_32_to_scaled_bf16(__m256i fp8, float scale) { |
| 93 | const __m512bh bf16_ext = CVT_FP8_TO_BF16_EXT(fp8); |
| 94 | const __m512 scale_vec = _mm512_mul_ps(_mm512_set1_ps(scale), _mm512_castsi512_ps(_mm512_set1_epi32(kFP8_BIAS))); |
| 95 | const __m512 f0 = _mm512_mul_ps(CVT_BF16_TO_FP32(_mm512_extracti32x8_epi32((__m512i)bf16_ext, 0)), scale_vec); |
| 96 | const __m512 f1 = _mm512_mul_ps(CVT_BF16_TO_FP32(_mm512_extracti32x8_epi32((__m512i)bf16_ext, 1)), scale_vec); |
| 97 | return (__m512i)_mm512_cvtne2ps_pbh(f1, f0); |
| 98 | } |
| 99 | |
| 100 | // Load 32 contiguous logical dims (NoPE fp8-dequant-to-bf16, or raw bf16 |
| 101 | // RoPE) starting at `dim_offset` from one token row of the V4_FP8Sparse |
| 102 | // layout. |
| 103 | inline __attribute__((always_inline)) __m512i load_fp8_kvcache_32_from_row( |
| 104 | const uint8_t* __restrict__ row_base, const uint8_t* __restrict__ scale_base, int dim_offset) { |
| 105 | if (dim_offset < kV4Meta.d_nope) { |
| 106 | const float scale = fp8_e8m0_to_float(scale_base[dim_offset / kV4Meta.tile_size]); |
| 107 | return cvt_fp8_32_to_scaled_bf16( |
| 108 | _mm256_loadu_si256(reinterpret_cast<const __m256i*>(row_base + dim_offset)), scale); |
| 109 | } |
| 110 | const auto* rope_ptr = reinterpret_cast<const at::BFloat16*>(row_base + kV4Meta.d_nope); |
| 111 | return _mm512_loadu_si512(rope_ptr + dim_offset - kV4Meta.d_nope); |
| 112 | } |
| 113 | |
| 114 | inline __attribute__((always_inline)) void init_fp8_kvcache_tile_rows( |
| 115 | const uint8_t* __restrict__ fp8_storage, |
| 116 | const int64_t block_size, |
| 117 | const int64_t storage_block_stride_bytes, |
| 118 | const int64_t token_idx, |
| 119 | const uint8_t*& row_base, |
| 120 | const uint8_t*& scale_base) { |
| 121 | constexpr int64_t nope_rope_per_token = kV4Meta.d_nope + 2 * kV4Meta.d_rope; |
| 122 | constexpr int64_t scale_stride = 8; |
| 123 | const int64_t block_idx = token_idx / block_size; |
| 124 | const int64_t block_off = token_idx - block_idx * block_size; |
| 125 | const uint8_t* block_base = fp8_storage + block_idx * storage_block_stride_bytes; |
| 126 | row_base = block_base + block_off * nope_rope_per_token; |
| 127 | scale_base = block_base + block_size * nope_rope_per_token + block_off * scale_stride; |
| 128 | } |
| 129 | |
| 130 | template <bool convert_v, typename scalar_t, typename index_t, typename load_vec_t> |
| 131 | inline __attribute__((always_inline)) void sparse_pack_vnni_Nx32( |
| 132 | scalar_t* __restrict__ dst0, |
| 133 | scalar_t* __restrict__ dst1, |
| 134 | const index_t* __restrict__ ind, |
| 135 | const bool* __restrict__ valid_mask, |
| 136 | int N, |
| 137 | int dim_offset, |
| 138 | int ld_dst0, |
| 139 | int ld_dst1, |
| 140 | const load_vec_t& load_vec) { |
| 141 | __m512i vinputs[16]; |
| 142 | int n = 0; |
| 143 | for (; n < N; ++n) { |
| 144 | if (!valid_mask[n]) { |
| 145 | vinputs[n] = _mm512_set1_epi32(0); |
| 146 | } else { |
| 147 | vinputs[n] = load_vec(n, static_cast<int64_t>(ind[n]), dim_offset); |
| 148 | } |
| 149 | } |
| 150 | for (; n < 16; ++n) { |
| 151 | vinputs[n] = _mm512_set1_epi32(0); |
| 152 | } |
| 153 | |
| 154 | if constexpr (convert_v) { |
| 155 | for (int nn = 0; nn < 16; nn += 2) { |
| 156 | __m512i d0, d1; |
| 157 | std::tie(d0, d1) = transpose_2x32_16bit(vinputs[nn], vinputs[nn + 1]); |
| 158 | _mm512_storeu_si512(dst1 + (nn >> 1) * ld_dst1 * 2, d0); |
| 159 | _mm512_storeu_si512(dst1 + (nn >> 1) * ld_dst1 * 2 + 32, d1); |
| 160 | } |
| 161 | } |
| 162 | |
| 163 | transpose_16x16_32bit(vinputs); |
| 164 | const __mmask16 vmask = (1 << N) - 1; |
| 165 | for (int k = 0; k < 16; ++k) { |
| 166 | _mm512_mask_storeu_epi32(dst0 + k * ld_dst0 * 2, vmask, vinputs[k]); |
| 167 | } |
| 168 | } |
| 169 | #endif // CPU_CAPABILITY_AVX512 |
| 170 | |
| 171 | // Gather+VNNI-pack `N` fp8-cache rows (dequantized to bf16 while packing) |
| 172 | // into `dst0` (K, vnni-packed for brgemm's B operand) and `dst1` (V, |
| 173 | // vnni-packed). |
| 174 | template <typename scalar_t, typename index_t> |
| 175 | void sparse_pack_vnni_fp8( |
| 176 | scalar_t* __restrict__ dst0, |
| 177 | scalar_t* __restrict__ dst1, |
| 178 | const uint8_t* __restrict__ fp8_storage, |
| 179 | const index_t* __restrict__ ind, |
| 180 | const bool* __restrict__ valid_mask, |
| 181 | int N, |
| 182 | int K, |
| 183 | int Kv, |
| 184 | int64_t block_size, |
| 185 | int64_t storage_block_stride_bytes, |
| 186 | int ld_dst0, |
| 187 | int ld_dst1) { |
| 188 | const int NB = div_up(N, 16); |
| 189 | const int KB = K / 32; |
| 190 | const int KBv = std::min(Kv / 32, KB); |
| 191 | for (int nb = 0; nb < NB; ++nb) { |
| 192 | const uint8_t* row_base[16]; |
| 193 | const uint8_t* scale_base[16]; |
| 194 | const int nb_size = std::min(N - nb * 16, 16); |
| 195 | for (int n = 0; n < nb_size; ++n) { |
| 196 | if (!valid_mask[nb * 16 + n]) { |
| 197 | row_base[n] = nullptr; |
| 198 | scale_base[n] = nullptr; |
| 199 | continue; |
| 200 | } |
| 201 | init_fp8_kvcache_tile_rows( |
| 202 | fp8_storage, block_size, storage_block_stride_bytes, static_cast<int64_t>(ind[nb * 16 + n]), row_base[n], |
| 203 | scale_base[n]); |
| 204 | } |
| 205 | auto load_vec = [&row_base, &scale_base](int n, int64_t idx, int dim_offset) { |
| 206 | UNUSED(idx); |
| 207 | return load_fp8_kvcache_32_from_row(row_base[n], scale_base[n], dim_offset); |
| 208 | }; |
| 209 | for (int kb = 0; kb < KBv; ++kb) { |
| 210 | sparse_pack_vnni_Nx32<true, scalar_t, index_t>( |
| 211 | dst0 + ((kb * 32) >> 1) * ld_dst0 * 2 + nb * 16 * 2, |
| 212 | dst1 + ((nb * 16) >> 1) * ld_dst1 * 2 + kb * 32 * 2, |
| 213 | ind + nb * 16, |
| 214 | valid_mask + nb * 16, |
| 215 | nb_size, |
| 216 | kb * 32, |
| 217 | ld_dst0, |
| 218 | ld_dst1, |
| 219 | load_vec); |
| 220 | } |
| 221 | for (int kb = KBv; kb < KB; ++kb) { |
| 222 | sparse_pack_vnni_Nx32<false, scalar_t, index_t>( |
| 223 | dst0 + ((kb * 32) >> 1) * ld_dst0 * 2 + nb * 16 * 2, |
| 224 | dst1 + ((nb * 16) >> 1) * ld_dst1 * 2 + kb * 32 * 2, |
| 225 | ind + nb * 16, |
| 226 | valid_mask + nb * 16, |
| 227 | nb_size, |
| 228 | kb * 32, |
| 229 | ld_dst0, |
| 230 | ld_dst1, |
| 231 | load_vec); |
| 232 | } |
| 233 | } |
| 234 | } |
| 235 | |
| 236 | template <typename scalar_t> |
| 237 | inline void fmla_fill_stub(scalar_t* __restrict__ out, float val, int64_t size) { |
| 238 | using Vec = at::vec::Vectorized<scalar_t>; |
| 239 | constexpr int kVecSize = Vec::size(); |
| 240 | const Vec data_vec = Vec(static_cast<scalar_t>(val)); |
| 241 | int64_t d = 0; |
| 242 | #pragma GCC unroll 4 |
| 243 | for (; d <= size - kVecSize; d += kVecSize) { |
| 244 | data_vec.store(out + d); |
| 245 | } |
| 246 | if (size - d > 0) { |
| 247 | data_vec.store(out + d, size - d); |
| 248 | } |
| 249 | } |
| 250 | |
| 251 | template <typename scalar_t, int BLOCK_N> |
| 252 | inline void fmla_copy_stub(scalar_t* __restrict__ out, const float* __restrict__ input) { |
| 253 | static_assert(BLOCK_N % 32 == 0); |
| 254 | using bVec = at::vec::Vectorized<scalar_t>; |
| 255 | using fVec = at::vec::Vectorized<float>; |
| 256 | constexpr int COLS = BLOCK_N / 16; |
| 257 | auto store = [&](auto i) { |
| 258 | constexpr int col = i % COLS; |
| 259 | if constexpr (col % 2 == 0) { |
| 260 | fVec a0 = fVec::loadu(input + col * 16); |
| 261 | fVec a1 = fVec::loadu(input + col * 16 + 16); |
| 262 | bVec out_bvec = convert_from_float_ext<scalar_t>(a0, a1); |
| 263 | out_bvec.store(out + col * 16); |
| 264 | } |
| 265 | }; |
| 266 | Unroll<COLS>{}(store); |
| 267 | } |
| 268 | |
| 269 | template <typename scalar_t> |
| 270 | inline void fmla_finalize_out(scalar_t* __restrict__ out, const float* __restrict__ acc, float inv_s, int64_t size) { |
| 271 | using bVec = at::vec::Vectorized<scalar_t>; |
| 272 | using fVec = at::vec::Vectorized<float>; |
| 273 | constexpr int kVecSize = bVec::size(); |
| 274 | const fVec s_fvec = fVec(inv_s); |
| 275 | int64_t d = 0; |
| 276 | #pragma GCC unroll 4 |
| 277 | for (; d <= size - kVecSize; d += kVecSize) { |
| 278 | fVec a0 = fVec::loadu(acc + d) * s_fvec; |
| 279 | fVec a1 = fVec::loadu(acc + d + fVec::size()) * s_fvec; |
| 280 | bVec out_bvec = convert_from_float_ext<scalar_t>(a0, a1); |
| 281 | out_bvec.store(out + d); |
| 282 | } |
| 283 | for (; d < size; ++d) { |
| 284 | out[d] = static_cast<scalar_t>(acc[d] * inv_s); |
| 285 | } |
| 286 | } |
| 287 | |
| 288 | // --------------------------------------------------------------------------- |
| 289 | // query : [num_tokens, H_q, D_qk] bf16 (SGLang's `[B, S_q, H_q, D]` |
| 290 | // with S_q folded away -- vLLM's flat per-token convention) |
| 291 | // indices : [num_tokens, topk_main] int64 (flat, block-table-resolved |
| 292 | // slot ids, -1 = invalid; SGLang's "main"/window index set) |
| 293 | // extra_* : same, for the optional "extra"/compressed index set |
| 294 | // attn_sink: [H_q] float32 |
| 295 | // output : [num_tokens, H_q, D_v] bf16 |
| 296 | // --------------------------------------------------------------------------- |
| 297 | |
| 298 | template <typename scalar_t, typename index_t, int64_t BLOCK_N> |
| 299 | void sparse_mla_decode_kernel_impl( |
| 300 | scalar_t* __restrict__ output, |
| 301 | float* __restrict__ lse_out, |
| 302 | const scalar_t* __restrict__ query, |
| 303 | const uint8_t* __restrict__ k_main_fp8, |
| 304 | const uint8_t* __restrict__ k_extra_fp8, |
| 305 | const index_t* __restrict__ indices, |
| 306 | const index_t* __restrict__ extra_indices, |
| 307 | const float* __restrict__ attn_sink, |
| 308 | scalar_t* __restrict__ buffer, |
| 309 | int64_t batches, |
| 310 | int64_t s_q, |
| 311 | int64_t num_heads, |
| 312 | int64_t head_size, |
| 313 | int64_t head_size_v, |
| 314 | int64_t topk_main, |
| 315 | int64_t topk_extra, |
| 316 | int64_t total_tokens_main, |
| 317 | int64_t total_tokens_extra, |
| 318 | int64_t q_strideB, |
| 319 | int64_t q_strideH, |
| 320 | int64_t k_main_block_size, |
| 321 | int64_t k_extra_block_size, |
| 322 | int64_t k_main_block_stride_bytes, |
| 323 | int64_t k_extra_block_stride_bytes, |
| 324 | int64_t idx_strideB, |
| 325 | int64_t extra_idx_strideB, |
| 326 | float scaling, |
| 327 | int64_t buffer_size_per_thread) { |
| 328 | using Vec = at::vec::Vectorized<float>; |
| 329 | |
| 330 | // partition heads |
| 331 | constexpr int64_t kBLOCK_H_MAX = 16; |
| 332 | const int64_t BLOCK_H = (batches * s_q) >= 16 ? kBLOCK_H_MAX : 8; |
| 333 | const int64_t num_h_blocks = div_up(num_heads, BLOCK_H); |
| 334 | |
| 335 | // parallel on [B, S_q, head_block] -- with s_q == 1 (vLLM's flat call |
| 336 | // convention) this collapses to a plain [num_tokens, head_block] loop, |
| 337 | // one flat iteration per token with no padding to skip. |
| 338 | at::parallel_for(0, batches * s_q * num_h_blocks, 0, [&](int64_t begin, int64_t end) { |
| 339 | int64_t bs{0}, sq{0}, hb{0}; |
| 340 | data_index_init(begin, bs, batches, sq, s_q, hb, num_h_blocks); |
| 341 | |
| 342 | int tid = at::get_thread_num(); |
| 343 | scalar_t* __restrict__ Btmp0 = buffer + tid * buffer_size_per_thread; // K packed |
| 344 | scalar_t* __restrict__ Btmp1 = Btmp0 + BLOCK_N * head_size; // V packed |
| 345 | // f32 V accumulator follows the bf16 packing region (reinterpret cast). |
| 346 | float* __restrict__ v_acc_local = reinterpret_cast<float*>(Btmp1 + BLOCK_N * head_size_v); |
| 347 | fmla_fill_stub(Btmp1, 0.f, BLOCK_N * head_size_v); // initialize V padding |
| 348 | |
| 349 | alignas(64) float s_i[kBLOCK_H_MAX * BLOCK_N]; |
| 350 | float* __restrict__ s_delta = s_i; |
| 351 | alignas(64) scalar_t s_delta2[kBLOCK_H_MAX * BLOCK_N]; |
| 352 | |
| 353 | alignas(64) float s_prime[kBLOCK_H_MAX]; |
| 354 | alignas(64) float m_prime[kBLOCK_H_MAX]; |
| 355 | |
| 356 | for (int64_t i = begin; i < end; ++i) { |
| 357 | const int64_t h_start = hb * BLOCK_H; |
| 358 | const int64_t h_end = std::min(h_start + BLOCK_H, num_heads); |
| 359 | const int64_t h_size = h_end - h_start; |
| 360 | |
| 361 | const scalar_t* __restrict__ q_ptr = query + bs * q_strideB + h_start * q_strideH; |
| 362 | const index_t* __restrict__ idx_ptr = indices + bs * idx_strideB; |
| 363 | const index_t* __restrict__ extra_idx_ptr = extra_indices == nullptr ? nullptr : extra_indices + bs * extra_idx_strideB; |
| 364 | |
| 365 | fmla_fill_stub(s_prime, 0.f, BLOCK_H); |
| 366 | fmla_fill_stub(m_prime, -std::numeric_limits<float>::infinity(), BLOCK_H); |
| 367 | for (int64_t h = 0; h < h_size; ++h) { |
| 368 | fmla_fill_stub(v_acc_local + h * head_size_v, 0.f, head_size_v); |
| 369 | } |
| 370 | |
| 371 | auto process_cache = [&](const uint8_t* __restrict__ fp8_ptr, |
| 372 | const index_t* __restrict__ cur_idx_ptr, |
| 373 | int64_t topk_count, |
| 374 | int64_t total_tokens, |
| 375 | int64_t fp8_block_size, |
| 376 | int64_t fp8_block_stride_bytes) { |
| 377 | if (fp8_ptr == nullptr || cur_idx_ptr == nullptr || topk_count == 0) { |
| 378 | return; |
| 379 | } |
| 380 | const int64_t topk_limit = topk_count; |
| 381 | |
| 382 | for (int64_t n = 0; n < topk_limit; n += BLOCK_N) { |
| 383 | int64_t n_size = std::min<int64_t>(BLOCK_N, topk_limit - n); |
| 384 | const int64_t padded_n_size = div_up(int(n_size), TILE_K) * TILE_K; |
| 385 | bool valid_mask[BLOCK_N]; |
| 386 | bool has_valid = false; |
| 387 | for (int64_t k = 0; k < n_size; ++k) { |
| 388 | const bool valid = is_valid_sparse_index(cur_idx_ptr[n + k], n + k, topk_limit, total_tokens); |
| 389 | valid_mask[k] = valid; |
| 390 | has_valid |= valid; |
| 391 | } |
| 392 | if (!has_valid) { |
| 393 | continue; |
| 394 | } |
| 395 | |
| 396 | // Pack K (BLOCK_N rows via gather) into Btmp0 (key, vnni) and Btmp1 |
| 397 | // (value, vnni), dequantizing FP8 cache entries while packing. |
| 398 | sparse_pack_vnni_fp8<scalar_t, index_t>( |
| 399 | Btmp0, Btmp1, fp8_ptr, cur_idx_ptr + n, valid_mask, static_cast<int>(n_size), static_cast<int>(head_size), |
| 400 | static_cast<int>(head_size_v), fp8_block_size, fp8_block_stride_bytes, static_cast<int>(BLOCK_N), |
| 401 | static_cast<int>(head_size_v)); |
| 402 | |
| 403 | // Q @ K |
| 404 | at::native::cpublas::brgemm( |
| 405 | /* M */ h_size, |
| 406 | /* N */ n_size, |
| 407 | /* K */ head_size, |
| 408 | /* lda */ q_strideH, |
| 409 | /* ldb */ BLOCK_N, |
| 410 | /* ldc */ BLOCK_N, |
| 411 | /* add_C */ false, |
| 412 | /* A */ q_ptr, |
| 413 | /* B */ Btmp0, |
| 414 | /* C */ s_i); |
| 415 | |
| 416 | const Vec scale_vec = Vec(scaling); |
| 417 | for (int64_t h = 0; h < h_size; ++h) { |
| 418 | float* row = s_i + h * BLOCK_N; |
| 419 | at::vec::map<float>([scale_vec](Vec x) { return x * scale_vec; }, row, row, n_size); |
| 420 | |
| 421 | for (int64_t k = 0; k < n_size; ++k) { |
| 422 | if (!valid_mask[k]) { |
| 423 | row[k] = -std::numeric_limits<float>::infinity(); |
| 424 | } |
| 425 | } |
| 426 | |
| 427 | // online softmax update |
| 428 | float m_i = at::vec::reduce_all<float>([](Vec& x, Vec& y) { return at::vec::maximum(x, y); }, row, n_size); |
| 429 | m_i = std::max(m_i, m_prime[h]); |
| 430 | |
| 431 | if (!std::isfinite(m_i)) { |
| 432 | fmla_fill_stub(s_delta + h * BLOCK_N, 0.f, padded_n_size); |
| 433 | fmla_copy_stub<scalar_t, BLOCK_N>(s_delta2 + h * BLOCK_N, s_delta + h * BLOCK_N); |
| 434 | continue; |
| 435 | } |
| 436 | |
| 437 | const float m_delta = std::exp(m_prime[h] - m_i); |
| 438 | at::vec::map<float>([m_i](Vec x) { return (x - Vec(m_i)).exp_u20(); }, s_delta + h * BLOCK_N, row, n_size); |
| 439 | |
| 440 | s_prime[h] *= m_delta; |
| 441 | s_prime[h] += at::vec::reduce_all<float>([](Vec& x, Vec& y) { return x + y; }, s_delta + h * BLOCK_N, n_size); |
| 442 | |
| 443 | m_prime[h] = m_i; |
| 444 | |
| 445 | float scale_m = m_delta; |
| 446 | at::vec::map<float>( |
| 447 | [scale_m](Vec x) { return x * Vec(scale_m); }, v_acc_local + h * head_size_v, v_acc_local + h * head_size_v, |
| 448 | head_size_v); |
| 449 | |
| 450 | fmla_fill_stub(s_delta + h * BLOCK_N + n_size, 0.f, padded_n_size - n_size); |
| 451 | fmla_copy_stub<scalar_t, BLOCK_N>(s_delta2 + h * BLOCK_N, s_delta + h * BLOCK_N); |
| 452 | } |
| 453 | |
| 454 | // V' <- s_delta @ V + V' (accumulate into v_acc_local at f32) |
| 455 | at::native::cpublas::brgemm( |
| 456 | /* M */ h_size, |
| 457 | /* N */ head_size_v, |
| 458 | /* K */ padded_n_size, |
| 459 | /* lda */ BLOCK_N, |
| 460 | /* ldb */ head_size_v, |
| 461 | /* ldc */ head_size_v, |
| 462 | /* add_C */ true, |
| 463 | /* A */ s_delta2, |
| 464 | /* B */ Btmp1, |
| 465 | /* C */ v_acc_local); |
| 466 | } |
| 467 | }; |
| 468 | |
| 469 | process_cache(k_main_fp8, idx_ptr, topk_main, total_tokens_main, k_main_block_size, k_main_block_stride_bytes); |
| 470 | process_cache( |
| 471 | k_extra_fp8, extra_idx_ptr, topk_extra, total_tokens_extra, k_extra_block_size, k_extra_block_stride_bytes); |
| 472 | |
| 473 | // Apply attention sink correction directly on the output and lse. |
| 474 | // out *= exp(lse_no_sink) / (exp(lse_no_sink) + exp(attn_sink)) |
| 475 | // = 1 / (1 + exp(attn_sink - lse_no_sink)) |
| 476 | // (Algebraically equivalent to treating attn_sink as a virtual, |
| 477 | // zero-value extra key inside the softmax.) |
| 478 | for (int64_t h = 0; h < h_size; ++h) { |
| 479 | const int64_t hh = h_start + h; |
| 480 | const bool lonely = !std::isfinite(m_prime[h]) || s_prime[h] == 0.f; |
| 481 | float lse_val = lonely ? std::numeric_limits<float>::infinity() : (m_prime[h] + std::log(s_prime[h])); |
| 482 | float inv_s = lonely ? 0.f : (1.f / s_prime[h]); |
| 483 | |
| 484 | if (!lonely && attn_sink != nullptr) { |
| 485 | const float sink = attn_sink[hh]; |
| 486 | const float corr = 1.f / (1.f + std::exp(sink - lse_val)); |
| 487 | inv_s *= corr; |
| 488 | } |
| 489 | |
| 490 | scalar_t* out_row = output + bs * (s_q * num_heads * head_size_v) + sq * (num_heads * head_size_v) + hh * head_size_v; |
| 491 | if (lonely) { |
| 492 | fmla_fill_stub(out_row, 0.f, head_size_v); |
| 493 | } else { |
| 494 | fmla_finalize_out<scalar_t>(out_row, v_acc_local + h * head_size_v, inv_s, head_size_v); |
| 495 | } |
| 496 | |
| 497 | lse_out[bs * num_heads * s_q + hh * s_q + sq] = lse_val; |
| 498 | } |
| 499 | |
| 500 | data_index_step(bs, batches, sq, s_q, hb, num_h_blocks); |
| 501 | } |
| 502 | at::native::cpublas::brgemm_release(); |
| 503 | }); |
| 504 | } |
| 505 | |
| 506 | } // anonymous namespace |
| 507 | |
| 508 | void flash_mla_with_kvcache_cpu( |
| 509 | at::Tensor& out, |
| 510 | at::Tensor& q, |
| 511 | at::Tensor& window_cache_2d, |
| 512 | at::Tensor& window_slots, |
| 513 | int64_t window_block_size, |
| 514 | at::Tensor& compressed_cache_2d, |
| 515 | at::Tensor& compressed_slots, |
| 516 | int64_t compressed_block_size, |
| 517 | at::Tensor& attn_sink, |
| 518 | double scale) { |
| 519 | TORCH_CHECK( |
| 520 | q.dim() == 3 && q.size(2) == kHeadDim, |
| 521 | "flash_mla_with_kvcache_cpu: q must be [num_tokens, num_heads, 512]"); |
| 522 | TORCH_CHECK( |
| 523 | out.dim() == 3 && out.size(0) == q.size(0) && out.size(1) == q.size(1) && out.size(2) == kHeadDim, |
| 524 | "flash_mla_with_kvcache_cpu: out must be [num_tokens, num_heads, 512]"); |
| 525 | TORCH_CHECK( |
| 526 | out.scalar_type() == at::kBFloat16 && out.is_contiguous(), |
| 527 | "flash_mla_with_kvcache_cpu: out must be bf16 and contiguous " |
| 528 | "(pre-allocate once in Python -- see DeepseekV4CPUAttention.forward_mqa)"); |
| 529 | TORCH_CHECK(out.device().is_cpu(), "flash_mla_with_kvcache_cpu: out must be a CPU tensor"); |
| 530 | TORCH_CHECK( |
| 531 | q.scalar_type() == at::kBFloat16, |
| 532 | "flash_mla_with_kvcache_cpu: q must be bf16 (the AMX/VNNI kernel's " |
| 533 | "FP8->bf16 dequant intrinsics are bf16-specific)"); |
| 534 | TORCH_CHECK( |
| 535 | window_slots.dim() == 2 && window_slots.size(0) == q.size(0), |
| 536 | "flash_mla_with_kvcache_cpu: window_slots must be [num_tokens, num_window]"); |
| 537 | TORCH_CHECK( |
| 538 | compressed_slots.dim() == 2 && compressed_slots.size(0) == q.size(0), |
| 539 | "flash_mla_with_kvcache_cpu: compressed_slots must be [num_tokens, num_compressed]"); |
| 540 | TORCH_CHECK( |
| 541 | window_slots.scalar_type() == at::kLong && compressed_slots.scalar_type() == at::kLong, |
| 542 | "flash_mla_with_kvcache_cpu: slot tensors must be int64"); |
| 543 | TORCH_CHECK( |
| 544 | attn_sink.numel() == q.size(1), "flash_mla_with_kvcache_cpu: attn_sink must have num_heads entries"); |
| 545 | TORCH_CHECK( |
| 546 | attn_sink.scalar_type() == at::kFloat && attn_sink.is_contiguous(), |
| 547 | "flash_mla_with_kvcache_cpu: attn_sink must be float32 and contiguous " |
| 548 | "(convert once in Python -- see DeepseekV4CPUAttention.forward_mqa)"); |
| 549 | TORCH_CHECK( |
| 550 | window_slots.is_contiguous() && compressed_slots.is_contiguous(), |
| 551 | "flash_mla_with_kvcache_cpu: window_slots/compressed_slots must be contiguous"); |
| 552 | TORCH_CHECK( |
| 553 | q.device().is_cpu() && window_cache_2d.device().is_cpu() && window_slots.device().is_cpu() && |
| 554 | compressed_cache_2d.device().is_cpu() && compressed_slots.device().is_cpu() && |
| 555 | attn_sink.device().is_cpu(), |
| 556 | "flash_mla_with_kvcache_cpu: all inputs must be CPU tensors"); |
| 557 | |
| 558 | const int64_t num_tokens = q.size(0); |
| 559 | const int64_t num_heads = q.size(1); |
| 560 | const int64_t num_window = window_slots.size(1); |
| 561 | const int64_t num_compressed = compressed_slots.size(1); |
| 562 | TORCH_CHECK(num_window > 0, "flash_mla_with_kvcache_cpu: window_slots must have at least one column"); |
| 563 | |
| 564 | // Row stride is threaded through explicitly rather than assumed to equal |
| 565 | // block_size * (token_data_bytes + scale_bytes_per_token): a paged/pooled |
| 566 | // cache buffer's physical row stride is not guaranteed to match the |
| 567 | // logical per-block byte count (matches store_cache.cpp's precedent). |
| 568 | const int64_t window_cache_row_stride = window_cache_2d.stride(0); |
| 569 | const int64_t compressed_cache_row_stride = num_compressed > 0 ? compressed_cache_2d.stride(0) : 0; |
| 570 | |
| 571 | // AMX/VNNI path: flatten (batch, s_q) -> batches = num_tokens, s_q = 1. |
| 572 | CHECK_LAST_DIM_CONTIGUOUS_INPUT(q); |
| 573 | constexpr int64_t BLOCK_N = 128; // multiple of 32 for AMX brgemm |
| 574 | constexpr int64_t kBLOCK_H_MAX = 16; |
| 575 | static_assert(kHeadDim % 32 == 0, "head_dim must be a multiple of 32"); |
| 576 | |
| 577 | const int num_threads = at::get_num_threads(); |
| 578 | // Layout per thread (in bf16 elements): |
| 579 | // [Btmp0 : BLOCK_N * D] [Btmp1 : BLOCK_N * D] [v_acc_local : kBLOCK_H_MAX * D floats] |
| 580 | // f32 takes 2 bf16 elements -> multiply by 2. |
| 581 | const int64_t buffer_size_per_thread = BLOCK_N * kHeadDim + BLOCK_N * kHeadDim + 2 * kBLOCK_H_MAX * kHeadDim; |
| 582 | auto buffer = at::empty({num_threads, buffer_size_per_thread}, q.options()); |
| 583 | // lse is an internal scratch (attn_sink correction needs it per-token/head) |
| 584 | // -- not surfaced to the Python call site, which only consumes `out`. |
| 585 | auto lse = at::empty({num_tokens, num_heads}, q.options().dtype(at::kFloat)); |
| 586 | |
| 587 | const int64_t q_strideB = q.stride(0); |
| 588 | const int64_t q_strideH = q.stride(1); |
| 589 | const int64_t idx_strideB = window_slots.stride(0); |
| 590 | const int64_t extra_idx_strideB = num_compressed > 0 ? compressed_slots.stride(0) : 0; |
| 591 | |
| 592 | const int64_t total_tokens_main = window_cache_2d.size(0) * window_block_size; |
| 593 | const int64_t total_tokens_extra = num_compressed > 0 ? compressed_cache_2d.size(0) * compressed_block_size : 0; |
| 594 | |
| 595 | sparse_mla_decode_kernel_impl<at::BFloat16, int64_t, BLOCK_N>( |
| 596 | out.data_ptr<at::BFloat16>(), |
| 597 | lse.data_ptr<float>(), |
| 598 | q.data_ptr<at::BFloat16>(), |
| 599 | window_cache_2d.data_ptr<uint8_t>(), |
| 600 | num_compressed > 0 ? compressed_cache_2d.data_ptr<uint8_t>() : nullptr, |
| 601 | window_slots.data_ptr<int64_t>(), |
| 602 | num_compressed > 0 ? compressed_slots.data_ptr<int64_t>() : nullptr, |
| 603 | attn_sink.data_ptr<float>(), |
| 604 | buffer.data_ptr<at::BFloat16>(), |
| 605 | /* batches */ num_tokens, |
| 606 | /* s_q */ 1, |
| 607 | num_heads, |
| 608 | /* head_size */ kHeadDim, |
| 609 | /* head_size_v */ kHeadDim, |
| 610 | /* topk_main */ num_window, |
| 611 | /* topk_extra */ num_compressed, |
| 612 | total_tokens_main, |
| 613 | total_tokens_extra, |
| 614 | q_strideB, |
| 615 | q_strideH, |
| 616 | window_block_size, |
| 617 | num_compressed > 0 ? compressed_block_size : 1, |
| 618 | window_cache_row_stride, |
| 619 | compressed_cache_row_stride, |
| 620 | idx_strideB, |
| 621 | extra_idx_strideB, |
| 622 | static_cast<float>(scale), |
| 623 | buffer_size_per_thread); |
| 624 | } |
| 625 |