From e4d61d0d222ad21437a1507bb34ca7e916aac880 Mon Sep 17 00:00:00 2001 From: "Li, Jiang" Date: Wed, 19 Aug 2026 10:28:52 +0800 Subject: [PATCH] [CPU] Add AMX-only high-performance MLA backend for DeepSeek V2/V3/R1 (#52616) Signed-off-by: jiang1.li --- .buildkite/hardware_tests/cpu.yaml | 4 + cmake/cpu_extension.cmake | 6 +- csrc/cpu/sgl-kernels/bmm.cpp | 208 ++ csrc/cpu/sgl-kernels/decode.cpp | 1821 +++++++++++++++++ csrc/cpu/sgl-kernels/extend.cpp | 568 +++++ csrc/cpu/sgl-kernels/flash_attn.h | 250 +++ csrc/cpu/sgl-kernels/mla_cache.cpp | 111 + csrc/cpu/torch_bindings.cpp | 62 + tests/kernels/attention/test_amx_mla.py | 471 +++++ tests/v1/attention/test_mla_backends.py | 2 + .../v1/attention/test_mla_prefill_selector.py | 1 + vllm/_custom_ops.py | 103 + .../layers/attention/mla_attention.py | 41 + vllm/model_executor/layers/linear.py | 6 + vllm/platforms/cpu.py | 36 +- vllm/v1/attention/backends/mla/amx_mla.py | 405 ++++ .../backends/mla/prefill/cpu_native.py | 61 + .../backends/mla/prefill/registry.py | 3 + .../backends/mla/prefill/selector.py | 2 + vllm/v1/attention/backends/registry.py | 1 + 20 files changed, 4156 insertions(+), 6 deletions(-) create mode 100644 csrc/cpu/sgl-kernels/bmm.cpp create mode 100644 csrc/cpu/sgl-kernels/decode.cpp create mode 100644 csrc/cpu/sgl-kernels/extend.cpp create mode 100644 csrc/cpu/sgl-kernels/flash_attn.h create mode 100644 csrc/cpu/sgl-kernels/mla_cache.cpp create mode 100644 tests/kernels/attention/test_amx_mla.py create mode 100644 vllm/v1/attention/backends/mla/amx_mla.py create mode 100644 vllm/v1/attention/backends/mla/prefill/cpu_native.py diff --git a/.buildkite/hardware_tests/cpu.yaml b/.buildkite/hardware_tests/cpu.yaml index efd3e30f6f..77ebfb99be 100644 --- a/.buildkite/hardware_tests/cpu.yaml +++ b/.buildkite/hardware_tests/cpu.yaml @@ -12,6 +12,9 @@ steps: - vllm/_custom_ops.py - tests/kernels/attention/test_cpu_attn.py - tests/v1/attention/test_group_head_counts.py + - tests/kernels/attention/test_amx_mla.py + - vllm/v1/attention/backends/mla/amx_mla.py + - vllm/model_executor/layers/attention/mla_attention.py - tests/kernels/moe/test_cpu_fused_moe.py - tests/kernels/moe/test_cpu_quant_fused_moe.py - tests/kernels/test_onednn.py @@ -29,6 +32,7 @@ steps: bash .buildkite/scripts/hardware_ci/run-cpu-test.sh 30m " pytest -x -v -s tests/kernels/attention/test_cpu_attn.py pytest -x -v -s tests/v1/attention/test_group_head_counts.py + pytest -x -v -s tests/kernels/attention/test_amx_mla.py pytest -x -v -s tests/kernels/moe/test_cpu_fused_moe.py pytest -x -v -s tests/kernels/moe/test_cpu_quant_fused_moe.py pytest -x -v -s tests/kernels/mamba/test_cpu_short_conv.py diff --git a/cmake/cpu_extension.cmake b/cmake/cpu_extension.cmake index 46ed96c600..92c3daae64 100644 --- a/cmake/cpu_extension.cmake +++ b/cmake/cpu_extension.cmake @@ -519,7 +519,11 @@ if (ENABLE_X86_ISA) "csrc/cpu/sgl-kernels/moe.cpp" "csrc/cpu/sgl-kernels/moe_int8.cpp" "csrc/cpu/sgl-kernels/moe_int4.cpp" - "csrc/cpu/sgl-kernels/moe_fp8.cpp") + "csrc/cpu/sgl-kernels/moe_fp8.cpp" + "csrc/cpu/sgl-kernels/bmm.cpp" + "csrc/cpu/sgl-kernels/decode.cpp" + "csrc/cpu/sgl-kernels/extend.cpp" + "csrc/cpu/sgl-kernels/mla_cache.cpp") set(VLLM_EXT_SRC_AVX512 "csrc/cpu/sgl-kernels/fla.cpp" diff --git a/csrc/cpu/sgl-kernels/bmm.cpp b/csrc/cpu/sgl-kernels/bmm.cpp new file mode 100644 index 0000000000..7f1ef9fb8b --- /dev/null +++ b/csrc/cpu/sgl-kernels/bmm.cpp @@ -0,0 +1,208 @@ +// Adapted from +// https://github.com/sgl-project/sglang/tree/main/sgl-kernel/csrc/cpu + +// clang-format off + +#include "common.h" +#include "gemm.h" +#include "vec.h" + +namespace { + +template +void bmm_kernel_impl( + scalar_t* __restrict__ out, + const scalar_t* __restrict__ mat1, + const packed_t* __restrict__ mat2, + int64_t B, + int64_t M, + int64_t N, + int64_t K, + int64_t mat1_strideB, + int64_t mat1_strideM, + int64_t out_strideB, + int64_t out_strideM, + float scale = 0.f) { + constexpr int64_t BLOCK_M = block_size_m(); + constexpr int64_t BLOCK_N = block_size_n(); + const int64_t MB = div_up(M, BLOCK_M); + const int64_t NB = div_up(N, BLOCK_N); + + // mat2 contiguous in [B, N, K] + int64_t mat2_strideB = N * K; + int64_t mat2_strideN = K; + + const bool use_brgemm = can_use_brgemm(M); + + // parallel on [B, MB, NB] + at::parallel_for(0, B * MB * NB, 0, [&](int64_t begin, int64_t end) { + int64_t bs{0}, mb{0}, nb{0}; + data_index_init(begin, bs, B, mb, MB, nb, NB); + + // for brgemm, use float32 for accumulate + alignas(64) float Ctmp[BLOCK_M * BLOCK_N]; + + for (int i = begin; i < end; ++i) { + UNUSED(i); + int mb_start = mb * BLOCK_M; + int mb_size = std::min(M - mb_start, BLOCK_M); + int nb_start = nb * BLOCK_N; + int nb_size = std::min(N - nb_start, BLOCK_N); + + tinygemm_kernel( + /* A */ mat1 + bs * mat1_strideB + mb_start * mat1_strideM, + /* B */ mat2 + bs * mat2_strideB + nb_start * mat2_strideN /* nb * BLOCK_N * K */, + /* C */ out + bs * out_strideB + mb_start * out_strideM + nb_start, + /* Ctmp*/ Ctmp, + /* M */ mb_size, + /* N */ nb_size, + /* K */ K, + /* lda */ mat1_strideM, + /* ldb */ nb_size, + /* ldc */ out_strideM, + /* brg */ use_brgemm); + + // move to the next index + data_index_step(bs, B, mb, MB, nb, NB); + } + + if (use_brgemm) { + at::native::cpublas::brgemm_release(); + } + }); +} + +template <> +void bmm_kernel_impl( + at::BFloat16* __restrict__ out, + const at::BFloat16* __restrict__ mat1, + const at::Float8_e4m3fn* __restrict__ mat2, + int64_t B, + int64_t M, + int64_t N, + int64_t K, + int64_t mat1_strideB, + int64_t mat1_strideM, + int64_t out_strideB, + int64_t out_strideM, + float scale) { + constexpr int64_t BLOCK_M = block_size_m(); + constexpr int64_t BLOCK_N = block_size_n(); + const int64_t MB = div_up(M, BLOCK_M); + const int64_t NB = div_up(N, BLOCK_N); + + // mat2 contiguous in [B, N, K] + int64_t mat2_strideB = N * K; + int64_t mat2_strideN = K; + + const bool use_brgemm = can_use_brgemm(M); + + // parallel on [B, MB, NB] + parallel_2d(B * MB, NB, [&](int64_t mb0, int64_t mb1, int64_t nb0, int64_t nb1) { + // for brgemm, use float32 for accumulate + alignas(64) float Ctmp[BLOCK_M * BLOCK_N]; + // for brgemm when mat2 is float8_e4m3 + alignas(64) at::BFloat16 Btmp[BLOCK_N * BLOCK_K]; + + loop_2d(mb0, mb1, nb0, nb1, BLOCK_N * K, [&](int64_t mb, int64_t nb, int64_t nb_offset) { + int64_t bs = mb / MB; + int64_t mb_start = (mb % MB) * BLOCK_M; + int64_t mb_size = std::min(M - mb_start, BLOCK_M); + int64_t nb_start = nb * BLOCK_N; + int64_t nb_size = std::min(N - nb_start, BLOCK_N); + + tinygemm_kernel( + /* A */ mat1 + bs * mat1_strideB + mb_start * mat1_strideM, + /* B */ mat2 + bs * mat2_strideB + nb_start * mat2_strideN /* nb * BLOCK_N * K */, + /* C */ out + bs * out_strideB + mb_start * out_strideM + nb_start, + /* Btmp*/ Btmp, + /* Ctmp*/ Ctmp, + /*scale*/ scale, + /* M */ mb_size, + /* N */ nb_size, + /* K */ K, + /* lda */ mat1_strideM, + /* ldb */ nb_size, + /* ldc */ out_strideM, + /* brg */ use_brgemm); + }); + + if (use_brgemm) { + at::native::cpublas::brgemm_release(); + } + }); +} + +} // anonymous namespace + +// mat1 : [B, M, K] +// mat2 : [B, N, K] or [B, OC, IC] +// out : [B, M, N] +// scale: [] 0-dim tensor for per tensor quant +// +void bmm_cpu( + at::Tensor& out, at::Tensor& mat1, at::Tensor& mat2, bool is_vnni, const std::optional& scale) { + auto packed_w = is_vnni ? mat2 : convert_weight_packed(mat2); + + // input and out could be non-contiguous + // weight needs to be contiguous in [OC, IC] order + CHECK_LAST_DIM_CONTIGUOUS_INPUT(mat1); + CHECK_LAST_DIM_CONTIGUOUS_INPUT(out); + CHECK_INPUT(mat2); + CHECK_DIM(3, out); + CHECK_DIM(3, mat1); + CHECK_DIM(3, mat2); + + int64_t B = mat1.size(0); + int64_t M = mat1.size(1); + int64_t N = mat2.size(1); + int64_t K = mat1.size(2); + + const bool use_fp8_w8a16 = scale.has_value(); + TORCH_CHECK(N % 32 == 0, "tinygemm requires N to be 32x."); + + int64_t mat1_strideB = mat1.stride(0); + int64_t mat1_strideM = mat1.stride(1); + int64_t out_strideB = out.stride(0); + int64_t out_strideM = out.stride(1); + + // check shapes + TORCH_CHECK(mat2.size(0) == B && mat2.size(2) == K, "bmm: mat2 shape mismatch!"); + TORCH_CHECK(out.size(0) == B && out.size(1) == M, "bmm: out shape mismatch!"); + if (!use_fp8_w8a16) { + AT_DISPATCH_REDUCED_FLOATING_TYPES(mat1.scalar_type(), "bmm_kernel_impl", [&] { + bmm_kernel_impl( + out.data_ptr(), + mat1.data_ptr(), + packed_w.data_ptr(), + B, + M, + N, + K, + mat1_strideB, + mat1_strideM, + out_strideB, + out_strideM); + }); + } else { // fp8 bmm + float scale_val = 0.f; + + auto scale_tensor = scale.value(); + TORCH_CHECK(scale_tensor.ndimension() == 0, "bmm: expect scale to be 0-dim tensor."); + scale_val = scale_tensor.item(); + + bmm_kernel_impl( + out.data_ptr(), + mat1.data_ptr(), + packed_w.data_ptr(), + B, + M, + N, + K, + mat1_strideB, + mat1_strideM, + out_strideB, + out_strideM, + scale_val); + } +} diff --git a/csrc/cpu/sgl-kernels/decode.cpp b/csrc/cpu/sgl-kernels/decode.cpp new file mode 100644 index 0000000000..edc6e3818c --- /dev/null +++ b/csrc/cpu/sgl-kernels/decode.cpp @@ -0,0 +1,1821 @@ +// Adapted from +// https://github.com/sgl-project/sglang/tree/main/sgl-kernel/csrc/cpu + +// clang-format off + +#include "common.h" +#include "gemm.h" +#include "vec.h" + +namespace { + +// [NOTE] TODO list for this kernel: +// 1. tune the value for BLOCK_N +// 2. planning for {batches, num_heads, num_kv_splits} +// and use actual num_kv_splits for small seq length +// 3. try fast impl of `.tanh()` +// 4. provide amx kernel for index_gemm_kernel_nn when M = 16 +// + +#if defined(CPU_CAPABILITY_AVX512) +// key: from [N, 32] to [32/2, N, 2] +// val: from [N, 32] to [N/2, 32, 2] +template +inline void pack_vnni_Nx32( + scalar_t* __restrict__ dst0, + scalar_t* __restrict__ dst1, + const scalar_t* __restrict__ src, + const index_t* __restrict__ ind, + int N, + int ld_src, + int ld_dst0, + int ld_dst1, + bool convert_v) { + __m512i vinputs[16]; + int n = 0; + for (; n < N; ++n) { + vinputs[n] = _mm512_loadu_si512(src + ind[n] * ld_src); + } + // padding with zero to avoid uninitialized vectors + for (; n < 16; ++n) { + vinputs[n] = _mm512_set1_epi32(0); + } + + // pack value, skip 64 elems for deepseek + // handle 2 vectors at a time from [2, 32] to [32, 2] + if (convert_v) { + for (int n = 0; n < 16; n += 2) { + __m512i d0, d1; + std::tie(d0, d1) = transpose_2x32_16bit(vinputs[n], vinputs[n + 1]); + _mm512_storeu_si512(dst1 + (n >> 1) * ld_dst1 * 2, d0); + _mm512_storeu_si512(dst1 + (n >> 1) * ld_dst1 * 2 + 32, d1); + } + } + + // pack key + transpose_16x16_32bit(vinputs); + + const __mmask16 vmask = (1 << N) - 1; + for (int k = 0; k < 16; ++k) { + _mm512_mask_storeu_epi32(dst0 + k * ld_dst0 * 2, vmask, vinputs[k]); + } +} +#endif + +// [NOTE]: MLA vnni format conversion +// +// here we apply same strategy as `FlashMLA`: +// each kv_cache is loaded once and packed twice (L2 cache hit) +// +// * for key: from [N, K/2, 2] to [K/2, N, 2] +// * for value: from [N/2, 2, Kv] to [N/2, Kv, 2] +// +template +void pack_vnni( + scalar_t* __restrict__ dst0, + scalar_t* __restrict__ dst1, + const scalar_t* __restrict__ src, + const index_t* __restrict__ ind, + int N, + int K, + int Kv, + int ld_src, + int ld_dst0, + int ld_dst1) { +#if defined(CPU_CAPABILITY_AVX512) + const int NB = div_up(N, 16); + const int KB = K / 32; // no remainder + const int KBv = Kv / 32; // no remainder + + for (int nb = 0; nb < NB; ++nb) { + for (int kb = 0; kb < KB; ++kb) { + // handle 16x512bits each block + int nb_size = std::min(N - nb * 16, 16); + pack_vnni_Nx32( + /* dst0 */ dst0 + ((kb * 32) >> 1) * ld_dst0 * 2 + nb * 16 * 2, + /* dst1 */ dst1 + ((nb * 16) >> 1) * ld_dst1 * 2 + kb * 32 * 2, + /* src */ src + kb * 32, + /* ind */ ind + nb * 16, + /* N */ nb_size, + /* ld_src */ ld_src, + /* ld_dst0 */ ld_dst0, + /* ld_dst1 */ ld_dst1, + /* cvt_v */ kb < KBv); + } + } +#else + for (int n = 0; n < N; ++n) { + index_t index = ind[n]; + for (int k = 0; k < K / 2; ++k) { + for (int d = 0; d < 2; ++d) { + dst0[k * ld_dst0 * 2 + n * 2 + d] = src[index * ld_src + k * 2 + d]; + } + } + } + // from [N/2, 2, K] to [N/2, K, 2] + for (int n = 0; n < (N >> 1) * 2; n += 2) { + index_t index0 = ind[n + 0]; + index_t index1 = ind[n + 1]; + for (int k = 0; k < Kv; ++k) { + dst1[(n >> 1) * ld_dst1 * 2 + k * 2 + 0] = src[index0 * ld_src + k]; + dst1[(n >> 1) * ld_dst1 * 2 + k * 2 + 1] = src[index1 * ld_src + k]; + } + } + if (N % 2 != 0) { + index_t index = ind[N - 1]; + for (int k = 0; k < Kv; ++k) { + dst1[(N >> 1) * ld_dst1 * 2 + k * 2 + 0] = src[index * ld_src + k]; + dst1[(N >> 1) * ld_dst1 * 2 + k * 2 + 1] = 0; + } + } +#endif +} + +template +inline void fill_stub(scalar_t* __restrict__ out, float val, int64_t size) { + using Vec = at::vec::Vectorized; + constexpr int kVecSize = Vec::size(); + const Vec data_vec = Vec(static_cast(val)); + int64_t d = 0; +#pragma GCC unroll 4 + for (; d <= size - kVecSize; d += kVecSize) { + data_vec.store(out + d); + } + if (size - d > 0) { + data_vec.store(out + d, size - d); + } +} + +template +inline void copy_stub(scalar_t* __restrict__ out, const float* __restrict__ acc, float s, int64_t size) { + using bVec = at::vec::Vectorized; + using fVec = at::vec::Vectorized; + constexpr int kVecSize = bVec::size(); + const fVec s_fvec = fVec(s); + int64_t d = 0; +#pragma GCC unroll 4 + for (; d <= size - kVecSize; d += kVecSize) { + auto [a_fvec0, a_fvec1] = load_float_vec2(acc + d); + a_fvec0 = a_fvec0 * s_fvec; + a_fvec1 = a_fvec1 * s_fvec; + bVec out_bvec = convert_from_float_ext(a_fvec0, a_fvec1); + out_bvec.store(out + d); + } + for (; d < size; ++d) { + out[d] = static_cast(acc[d] * s); + } +} + +template +inline void copy_stub(scalar_t* __restrict__ out, const scalar_t* __restrict__ src, int64_t size) { + using bVec = at::vec::Vectorized; + constexpr int kVecSize = bVec::size(); + int64_t d = 0; +#pragma GCC unroll 4 + for (; d <= size - kVecSize; d += kVecSize) { + bVec out_bvec = bVec::loadu(src + d); + out_bvec.store(out + d); + } + for (; d < size; ++d) { + out[d] = src[d]; + } +} + +template +inline void copy_stub(scalar_t* __restrict__ out, const float* __restrict__ input) { + static_assert(BLOCK_N % 32 == 0); + using bVec = at::vec::Vectorized; + using fVec = at::vec::Vectorized; + + constexpr int COLS = BLOCK_N / 16; + auto store = [&](auto i) { + constexpr int col = i % COLS; + // for COLS = 2, 4 use 512bit store + if constexpr (col % 2 == 0) { + auto [a_fvec0, a_fvec1] = load_float_vec2(input + col * 16); + bVec out_bvec = convert_from_float_ext(a_fvec0, a_fvec1); + out_bvec.store(out + col * 16); + } + }; + Unroll{}(store); +} + +// GEMM handles query @ key (indexed) x scale +// A : [M, K] +// B : [N, K] indexed +// C : [M, N] +// +template +struct tinygemm_kernel_nt { + static inline void apply( + const scalar_t* __restrict__ A, + const scalar_t* __restrict__ B, + float* __restrict__ C, + const index_t* __restrict__ indices, + float scale, + int64_t lda, + int64_t ldb, + int64_t ldc, + int64_t K, + int64_t max_tokens) { + for (int64_t m = 0; m < BLOCK_M; ++m) { + for (int64_t n = 0; n < BLOCK_N; ++n) { + float sum = 0.f; + int64_t b_idx = indices[n]; + TORCH_CHECK(b_idx < max_tokens, "token index out of scope!"); + for (int64_t k = 0; k < K; ++k) { + sum += scale * static_cast(A[m * lda + k]) * static_cast(B[b_idx * ldb + k]); + } + C[m * ldc + n] = sum; + } + } + } +}; + +#if defined(CPU_CAPABILITY_AVX512) +template +struct tinygemm_kernel_nt { + static inline void apply( + const at::BFloat16* __restrict__ A, + const at::BFloat16* __restrict__ B, + float* __restrict__ C, + const index_t* __restrict__ indices, + float scale, + int64_t lda, + int64_t ldb, + int64_t ldc, + int64_t K, + int64_t max_tokens) { + constexpr int ROWS = BLOCK_M; + constexpr int COLS = BLOCK_N; + + __m512bh va; + __m512bh vb[COLS]; + __m512 vc[ROWS * COLS]; + __m512 vscale = _mm512_set1_ps(scale); + + auto loadc = [&](auto i) { vc[i] = _mm512_setzero_ps(); }; + Unroll{}(loadc); + + // for main loop + auto compute = [&](auto i, int64_t k) { + constexpr int row = i / COLS; + constexpr int col = i % COLS; + + if constexpr (col == 0) { + va = (__m512bh)(_mm512_loadu_si512(A + row * lda + k)); + } + if constexpr (row == 0) { + if constexpr (col + 1 < COLS) { + int64_t b_idx_prefetch = indices[col + 1]; + _mm_prefetch(B + b_idx_prefetch * ldb + k, _MM_HINT_T0); + } + int64_t b_idx = indices[col]; + TORCH_CHECK(b_idx < max_tokens, "token index out of scope!"); + vb[col] = (__m512bh)(_mm512_loadu_si512(B + b_idx * ldb + k)); + } + vc[i] = _mm512_dpbf16_ps(vc[i], va, vb[col]); + }; + + // for remainder + auto compute2 = [&](auto i, int64_t k, __mmask32 mask) { + constexpr int row = i / COLS; + constexpr int col = i % COLS; + + if constexpr (col == 0) { + va = (__m512bh)(_mm512_maskz_loadu_epi16(mask, A + row * lda + k)); + } + if constexpr (row == 0) { + int64_t b_idx = indices[col]; + TORCH_CHECK(b_idx < max_tokens, "token index out of scope!"); + vb[col] = (__m512bh)(_mm512_maskz_loadu_epi16(mask, B + b_idx * ldb + k)); + } + vc[i] = _mm512_dpbf16_ps(vc[i], va, vb[col]); + }; + + int64_t k = 0; + for (; k <= K - 32; k += 32) { + Unroll{}(compute, k); + } + int64_t count = K - k; + if (count > 0) { + __mmask32 mask = (1ULL << count) - 1; + Unroll{}(compute2, k, mask); + } + + auto storec = [&](auto i) { + constexpr int row = i / COLS; + constexpr int col = i % COLS; + C[row * ldc + col] = _mm512_reduce_add_ps(_mm512_mul_ps(vc[i], vscale)); + }; + Unroll{}(storec); + } +}; +#endif + +#if defined(CPU_CAPABILITY_AVX512) +template +struct tinygemm_kernel_nt { + static inline void apply( + const at::Half* __restrict__ A, + const at::Half* __restrict__ B, + float* __restrict__ C, + const index_t* __restrict__ indices, + float scale, + int64_t lda, + int64_t ldb, + int64_t ldc, + int64_t K, + int64_t max_tokens) { + constexpr int ROWS = BLOCK_M; + constexpr int COLS = BLOCK_N; + + __m512 va0, va1; + __m512 vb0[COLS], vb1[COLS]; + __m512 vc[ROWS * COLS]; + __m512 vscale = _mm512_set1_ps(scale); + + auto loadc = [&](auto i) { vc[i] = _mm512_setzero_ps(); }; + Unroll{}(loadc); + + auto compute = [&](auto i, int64_t k) { + constexpr int row = i / COLS; + constexpr int col = i % COLS; + + if constexpr (col == 0) { + __m512i a16 = _mm512_loadu_si512((__m512i const*)(A + row * lda + k)); + va0 = CVT_FP16_TO_FP32(_mm512_extracti32x8_epi32(a16, 0)); + va1 = CVT_FP16_TO_FP32(_mm512_extracti32x8_epi32(a16, 1)); + } + + if constexpr (row == 0) { + int64_t b_idx = indices[col]; + TORCH_CHECK(b_idx < max_tokens, "token index out of scope!"); + __m512i b16 = _mm512_loadu_si512((__m512i const*)(B + b_idx * ldb + k)); + vb0[col] = CVT_FP16_TO_FP32(_mm512_extracti32x8_epi32(b16, 0)); + vb1[col] = CVT_FP16_TO_FP32(_mm512_extracti32x8_epi32(b16, 1)); + } + + vc[i] = _mm512_fmadd_ps(va0, vb0[col], _mm512_fmadd_ps(va1, vb1[col], vc[i])); + }; + + auto compute2 = [&](auto i, int64_t k, __mmask32 mask) { + constexpr int row = i / COLS; + constexpr int col = i % COLS; + + if constexpr (col == 0) { + __m512i a16 = _mm512_maskz_loadu_epi16(mask, (const void*)(A + row * lda + k)); + va0 = CVT_FP16_TO_FP32(_mm512_extracti32x8_epi32(a16, 0)); + va1 = CVT_FP16_TO_FP32(_mm512_extracti32x8_epi32(a16, 1)); + } + + if constexpr (row == 0) { + int64_t b_idx = indices[col]; + TORCH_CHECK(b_idx < max_tokens, "token index out of scope!"); + __m512i b16 = _mm512_maskz_loadu_epi16(mask, (const void*)(B + b_idx * ldb + k)); + vb0[col] = CVT_FP16_TO_FP32(_mm512_extracti32x8_epi32(b16, 0)); + vb1[col] = CVT_FP16_TO_FP32(_mm512_extracti32x8_epi32(b16, 1)); + } + + vc[i] = _mm512_fmadd_ps(va0, vb0[col], _mm512_fmadd_ps(va1, vb1[col], vc[i])); + }; + + int64_t k = 0; + for (; k <= K - 32; k += 32) { + Unroll{}(compute, k); + } + int64_t count = K - k; + if (count > 0) { + __mmask32 mask = (1ULL << count) - 1; + Unroll{}(compute2, k, mask); + } + + auto storec = [&](auto i) { + constexpr int row = i / COLS; + constexpr int col = i % COLS; + C[row * ldc + col] = _mm512_reduce_add_ps(_mm512_mul_ps(vc[i], vscale)); + }; + Unroll{}(storec); + } +}; +#endif + +#define LAUNCH_TINYGEMM_KERNEL_NT(MB_SIZE, NB_SIZE) \ + tinygemm_kernel_nt::apply( \ + A + mb_start * lda, B, C + mb_start * ldc + nb_start, indices + nb_start, scale, lda, ldb, ldc, K, max_tokens); + +// this is used when N isn't multiple of 16, +// N corresponds to `head_size_v` which should be 16x +template +inline void tinygemm_kernel_nn_scalar( + const float* __restrict__ A, + const scalar_t* __restrict__ B, + float* __restrict__ C, + const index_t* __restrict__ indices, + const float* __restrict__ scale, + int64_t M, + int64_t N, + int64_t K, + int64_t lda, + int64_t ldb, + int64_t ldc, + int64_t max_tokens) { + for (int64_t m = 0; m < M; ++m) { + for (int64_t n = 0; n < N; ++n) { + C[m * ldc + n] *= scale[m]; + for (int64_t k = 0; k < K; ++k) { + int64_t b_idx = indices[k]; + TORCH_CHECK(b_idx < max_tokens, "token index out of scope!"); + C[m * ldc + n] += A[m * lda + k] * static_cast(B[b_idx * ldb + n]); + } + } + } +} + +// GEMM handles v' * scale + attn @ value (indexed) +// A : [M, K] +// B : [K, N] indexed +// C ：[M, N] +// +template +struct tinygemm_kernel_nn { + static inline void apply( + const float* __restrict__ A, + const scalar_t* __restrict__ B, + float* __restrict__ C, + const index_t* __restrict__ indices, + const float* __restrict__ scale, + int64_t lda, + int64_t ldb, + int64_t ldc, + int64_t K, + int64_t max_tokens) { + tinygemm_kernel_nn_scalar(A, B, C, indices, scale, BLOCK_M, BLOCK_N, K, lda, ldb, ldc, max_tokens); + } +}; + +#if defined(CPU_CAPABILITY_AVX512) +template +struct tinygemm_kernel_nn { + static inline void apply( + const float* __restrict__ A, + const at::BFloat16* __restrict__ B, + float* __restrict__ C, + const index_t* __restrict__ indices, + const float* __restrict__ scale, + int64_t lda, + int64_t ldb, + int64_t ldc, + int64_t K, + int64_t max_tokens) { + constexpr int ROWS = BLOCK_M; + constexpr int COLS = BLOCK_N / 16; + + __m512 va; + __m512 vb[COLS]; + __m512 vc[ROWS * COLS]; + __m512 vscale; + + auto loadc = [&](auto i) { + constexpr int row = i / COLS; + constexpr int col = i % COLS; +#pragma GCC diagnostic push +#pragma GCC diagnostic ignored "-Warray-bounds" + if constexpr (col == 0) { + vscale = _mm512_set1_ps(scale[row]); + } +#pragma GCC diagnostic pop + vc[i] = _mm512_loadu_ps(C + row * ldc + col * 16); + vc[i] = _mm512_mul_ps(vc[i], vscale); + }; + Unroll{}(loadc); + + auto compute = [&](auto i, int64_t k) { + constexpr int row = i / COLS; + constexpr int col = i % COLS; + + if constexpr (col == 0) { + va = _mm512_set1_ps(A[row * lda + k]); + } + if constexpr (row == 0) { + if (k + 1 < K) { + int64_t b_idx_prefetch = indices[k + 1]; + _mm_prefetch(B + b_idx_prefetch * ldb + col * 16, _MM_HINT_T0); + } + int64_t b_idx = indices[k]; + TORCH_CHECK(b_idx < max_tokens, "token index out of scope!"); + + // for COLS = 2, 4, 6, 8 use 512 bit load + // for COLS = 1, 3, 5, 7 use 256 bit load + if constexpr (COLS % 2 == 0) { + if constexpr (col % 2 == 0) { + __m512i b16 = _mm512_loadu_si512(reinterpret_cast(B + b_idx * ldb + col * 16)); + vb[col + 0] = CVT_BF16_TO_FP32(_mm512_extracti32x8_epi32(b16, 0)); + vb[col + 1] = CVT_BF16_TO_FP32(_mm512_extracti32x8_epi32(b16, 1)); + } + } else { + __m256i b16 = _mm256_loadu_si256(reinterpret_cast(B + b_idx * ldb + col * 16)); + vb[col] = CVT_BF16_TO_FP32(b16); + } + } + vc[i] = _mm512_fmadd_ps(va, vb[col], vc[i]); + }; + + for (int64_t k = 0; k < K; ++k) { + Unroll{}(compute, k); + } + + auto storec = [&](auto i) { + constexpr int row = i / COLS; + constexpr int col = i % COLS; + _mm512_storeu_ps(C + row * ldc + col * 16, vc[i]); + }; + Unroll{}(storec); + } +}; +#endif + +#if defined(CPU_CAPABILITY_AVX512) +template +struct tinygemm_kernel_nn { + static inline void apply( + const float* __restrict__ A, + const at::Half* __restrict__ B, + float* __restrict__ C, + const index_t* __restrict__ indices, + const float* __restrict__ scale, + int64_t lda, + int64_t ldb, + int64_t ldc, + int64_t K, + int64_t max_tokens) { + constexpr int ROWS = BLOCK_M; + constexpr int COLS = BLOCK_N / 16; + + __m512 va; + __m512 vb[COLS]; + __m512 vc[ROWS * COLS]; + __m512 vscale; + + auto loadc = [&](auto i) { + constexpr int row = i / COLS; + constexpr int col = i % COLS; +#pragma GCC diagnostic push +#pragma GCC diagnostic ignored "-Warray-bounds" + if constexpr (col == 0) { + vscale = _mm512_set1_ps(scale[row]); + } +#pragma GCC diagnostic pop + vc[i] = _mm512_loadu_ps(C + row * ldc + col * 16); + vc[i] = _mm512_mul_ps(vc[i], vscale); + }; + Unroll{}(loadc); + + auto compute = [&](auto i, int64_t k) { + constexpr int row = i / COLS; + constexpr int col = i % COLS; + + if constexpr (col == 0) { + va = _mm512_set1_ps(A[row * lda + k]); + } + if constexpr (row == 0) { + if (k + 1 < K) { + int64_t b_idx_prefetch = indices[k + 1]; + _mm_prefetch(B + b_idx_prefetch * ldb + col * 16, _MM_HINT_T0); + } + int64_t b_idx = indices[k]; + TORCH_CHECK(b_idx < max_tokens, "token index out of scope!"); + + // for COLS = 2, 4, 6, 8 use 512 bit load + // for COLS = 1, 3, 5, 7 use 256 bit load + if constexpr (COLS % 2 == 0) { + if constexpr (col % 2 == 0) { + __m512i b16 = _mm512_loadu_si512(reinterpret_cast(B + b_idx * ldb + col * 16)); + vb[col + 0] = CVT_FP16_TO_FP32(_mm512_extracti32x8_epi32(b16, 0)); + vb[col + 1] = CVT_FP16_TO_FP32(_mm512_extracti32x8_epi32(b16, 1)); + } + } else { + __m256i b16 = _mm256_loadu_si256(reinterpret_cast(B + b_idx * ldb + col * 16)); + vb[col] = CVT_FP16_TO_FP32(b16); + } + } + vc[i] = _mm512_fmadd_ps(va, vb[col], vc[i]); + }; + + for (int64_t k = 0; k < K; ++k) { + Unroll{}(compute, k); + } + + auto storec = [&](auto i) { + constexpr int row = i / COLS; + constexpr int col = i % COLS; + _mm512_storeu_ps(C + row * ldc + col * 16, vc[i]); + }; + Unroll{}(storec); + } +}; +#endif + +#define LAUNCH_TINYGEMM_KERNEL_NN(MB_SIZE, NB_SIZE) \ + tinygemm_kernel_nn::apply( \ + A + mb_start * lda, \ + B + nb_start, \ + C + mb_start * ldc + nb_start, \ + indices, \ + scale + mb_start, \ + lda, \ + ldb, \ + ldc, \ + K, \ + max_tokens); + +template +void index_gemm_kernel_nt( + const scalar_t* __restrict__ A, + const scalar_t* __restrict__ B, + float* __restrict__ C, + const index_t* __restrict__ indices, + float scale, + int64_t M, + int64_t N, + int64_t K, + int64_t lda, + int64_t ldb, + int64_t ldc, + int64_t max_tokens) { + // pattern: 1-8-8 + if (M == 1) { + constexpr int64_t BLOCK_N = 8; + const int64_t NB = div_up(N, BLOCK_N); + int64_t mb_start = 0, lda = 1, ldc = 1; + + for (int64_t nb = 0; nb < NB; ++nb) { + int64_t nb_start = nb * BLOCK_N; + int64_t nb_size = std::min(BLOCK_N, N - nb_start); + + switch (nb_size) { + case 1: + LAUNCH_TINYGEMM_KERNEL_NT(1, 1); + break; + case 2: + LAUNCH_TINYGEMM_KERNEL_NT(1, 2); + break; + case 3: + LAUNCH_TINYGEMM_KERNEL_NT(1, 3); + break; + case 4: + LAUNCH_TINYGEMM_KERNEL_NT(1, 4); + break; + case 5: + LAUNCH_TINYGEMM_KERNEL_NT(1, 5); + break; + case 6: + LAUNCH_TINYGEMM_KERNEL_NT(1, 6); + break; + case 7: + LAUNCH_TINYGEMM_KERNEL_NT(1, 7); + break; + case 8: + LAUNCH_TINYGEMM_KERNEL_NT(1, 8); + break; + default: + TORCH_CHECK(false, "Unexpected block size, 1x", "nb_size"); + } + } + return; + } + + // default pattern: 1-6-24 + // FP16 pattern: 2-8-16 + constexpr int64_t BLOCK_M = 4; + constexpr int64_t BLOCK_N = std::is_same_v ? 4 : 6; + const int64_t MB = div_up(M, BLOCK_M); + const int64_t NB = div_up(N, BLOCK_N); + + for (int64_t mb = 0; mb < MB; ++mb) { + int64_t mb_start = mb * BLOCK_M; + int64_t mb_size = std::min(BLOCK_M, M - mb_start); + for (int64_t nb = 0; nb < NB; ++nb) { + int64_t nb_start = nb * BLOCK_N; + int64_t nb_size = std::min(BLOCK_N, N - nb_start); + + switch (mb_size << 4 | nb_size) { + // mb_size = 1 + case 0x11: + LAUNCH_TINYGEMM_KERNEL_NT(1, 1); + break; + case 0x12: + LAUNCH_TINYGEMM_KERNEL_NT(1, 2); + break; + case 0x13: + LAUNCH_TINYGEMM_KERNEL_NT(1, 3); + break; + case 0x14: + LAUNCH_TINYGEMM_KERNEL_NT(1, 4); + break; + case 0x15: + LAUNCH_TINYGEMM_KERNEL_NT(1, 5); + break; + case 0x16: + LAUNCH_TINYGEMM_KERNEL_NT(1, 6); + break; + // mb_size = 2 + case 0x21: + LAUNCH_TINYGEMM_KERNEL_NT(2, 1); + break; + case 0x22: + LAUNCH_TINYGEMM_KERNEL_NT(2, 2); + break; + case 0x23: + LAUNCH_TINYGEMM_KERNEL_NT(2, 3); + break; + case 0x24: + LAUNCH_TINYGEMM_KERNEL_NT(2, 4); + break; + case 0x25: + LAUNCH_TINYGEMM_KERNEL_NT(2, 5); + break; + case 0x26: + LAUNCH_TINYGEMM_KERNEL_NT(2, 6); + break; + // mb_size = 3 + case 0x31: + LAUNCH_TINYGEMM_KERNEL_NT(3, 1); + break; + case 0x32: + LAUNCH_TINYGEMM_KERNEL_NT(3, 2); + break; + case 0x33: + LAUNCH_TINYGEMM_KERNEL_NT(3, 3); + break; + case 0x34: + LAUNCH_TINYGEMM_KERNEL_NT(3, 4); + break; + case 0x35: + LAUNCH_TINYGEMM_KERNEL_NT(3, 5); + break; + case 0x36: + LAUNCH_TINYGEMM_KERNEL_NT(3, 6); + break; + // mb_size = 4 + case 0x41: + LAUNCH_TINYGEMM_KERNEL_NT(4, 1); + break; + case 0x42: + LAUNCH_TINYGEMM_KERNEL_NT(4, 2); + break; + case 0x43: + LAUNCH_TINYGEMM_KERNEL_NT(4, 3); + break; + case 0x44: + LAUNCH_TINYGEMM_KERNEL_NT(4, 4); + break; + case 0x45: + LAUNCH_TINYGEMM_KERNEL_NT(4, 5); + break; + case 0x46: + LAUNCH_TINYGEMM_KERNEL_NT(4, 6); + break; + default: + TORCH_CHECK(false, "Unexpected block size, ", mb_size, "x", "nb_size"); + } + } + } +} + +template +void index_gemm_kernel_nn( + const float* __restrict__ A, + const scalar_t* __restrict__ B, + float* __restrict__ C, + const index_t* __restrict__ indices, + float* __restrict__ scale, + int64_t M, + int64_t N, + int64_t K, + int64_t lda, + int64_t ldb, + int64_t ldc, + int64_t max_tokens) { + constexpr int kVecSize = 16; + if ((N & (kVecSize - 1)) != 0) { + tinygemm_kernel_nn_scalar(A, B, C, indices, scale, M, N, K, lda, ldb, ldc, max_tokens); + return; + } + + // pattern: 1-8-8 + if (M == 1) { + constexpr int64_t BLOCK_N = 8 * kVecSize; + const int64_t NB = div_up(N, BLOCK_N); + int64_t mb_start = 0, lda = 1, ldc = 1; + + for (int64_t nb = 0; nb < NB; ++nb) { + int64_t nb_start = nb * BLOCK_N; + int64_t nb_size = std::min(BLOCK_N, N - nb_start); + + switch (nb_size >> 4) { + case 1: + LAUNCH_TINYGEMM_KERNEL_NN(1, 16); + break; + case 2: + LAUNCH_TINYGEMM_KERNEL_NN(1, 32); + break; + case 3: + LAUNCH_TINYGEMM_KERNEL_NN(1, 48); + break; + case 4: + LAUNCH_TINYGEMM_KERNEL_NN(1, 64); + break; + case 5: + LAUNCH_TINYGEMM_KERNEL_NN(1, 80); + break; + case 6: + LAUNCH_TINYGEMM_KERNEL_NN(1, 96); + break; + case 7: + LAUNCH_TINYGEMM_KERNEL_NN(1, 112); + break; + case 8: + LAUNCH_TINYGEMM_KERNEL_NN(1, 128); + break; + default: + TORCH_CHECK(false, "Unexpected block size, 1x", "nb_size"); + } + } + return; + } + + constexpr int64_t BLOCK_M = 4; + constexpr int64_t BLOCK_N = 6 * kVecSize; + const int64_t MB = div_up(M, BLOCK_M); + const int64_t NB = div_up(N, BLOCK_N); + + for (int64_t mb = 0; mb < MB; ++mb) { + int64_t mb_start = mb * BLOCK_M; + int64_t mb_size = std::min(BLOCK_M, M - mb_start); + for (int64_t nb = 0; nb < NB; ++nb) { + int64_t nb_start = nb * BLOCK_N; + int64_t nb_size = std::min(BLOCK_N, N - nb_start); + + switch (mb_size << 4 | nb_size >> 4) { + // mb_size = 1 + case 0x11: + LAUNCH_TINYGEMM_KERNEL_NN(1, 16); + break; + case 0x12: + LAUNCH_TINYGEMM_KERNEL_NN(1, 32); + break; + case 0x13: + LAUNCH_TINYGEMM_KERNEL_NN(1, 48); + break; + case 0x14: + LAUNCH_TINYGEMM_KERNEL_NN(1, 64); + break; + case 0x15: + LAUNCH_TINYGEMM_KERNEL_NN(1, 80); + break; + case 0x16: + LAUNCH_TINYGEMM_KERNEL_NN(1, 96); + break; + // mb_size = 2 + case 0x21: + LAUNCH_TINYGEMM_KERNEL_NN(2, 16); + break; + case 0x22: + LAUNCH_TINYGEMM_KERNEL_NN(2, 32); + break; + case 0x23: + LAUNCH_TINYGEMM_KERNEL_NN(2, 48); + break; + case 0x24: + LAUNCH_TINYGEMM_KERNEL_NN(2, 64); + break; + case 0x25: + LAUNCH_TINYGEMM_KERNEL_NN(2, 80); + break; + case 0x26: + LAUNCH_TINYGEMM_KERNEL_NN(2, 96); + break; + // mb_size = 3 + case 0x31: + LAUNCH_TINYGEMM_KERNEL_NN(3, 16); + break; + case 0x32: + LAUNCH_TINYGEMM_KERNEL_NN(3, 32); + break; + case 0x33: + LAUNCH_TINYGEMM_KERNEL_NN(3, 48); + break; + case 0x34: + LAUNCH_TINYGEMM_KERNEL_NN(3, 64); + break; + case 0x35: + LAUNCH_TINYGEMM_KERNEL_NN(3, 80); + break; + case 0x36: + LAUNCH_TINYGEMM_KERNEL_NN(3, 96); + break; + // mb_size = 4 + case 0x41: + LAUNCH_TINYGEMM_KERNEL_NN(4, 16); + break; + case 0x42: + LAUNCH_TINYGEMM_KERNEL_NN(4, 32); + break; + case 0x43: + LAUNCH_TINYGEMM_KERNEL_NN(4, 48); + break; + case 0x44: + LAUNCH_TINYGEMM_KERNEL_NN(4, 64); + break; + case 0x45: + LAUNCH_TINYGEMM_KERNEL_NN(4, 80); + break; + case 0x46: + LAUNCH_TINYGEMM_KERNEL_NN(4, 96); + break; + default: + TORCH_CHECK(false, "Unexpected block size, ", mb_size, "x", "nb_size"); + } + } + } +} + +template +void decode_set_kv_buffer( + scalar_t* __restrict__ k_buffer, + scalar_t* __restrict__ v_buffer, + const scalar_t* __restrict__ key, + const scalar_t* __restrict__ value, + const int64_t* __restrict__ loc, + int64_t batches, + int64_t num_heads_kv, + int64_t head_size, + int64_t head_size_v, + int64_t k_strideN, + int64_t k_strideH, + int64_t v_strideN, + int64_t v_strideH, + int64_t nk_strideN, + int64_t nk_strideH, + int64_t nv_strideN, + int64_t nv_strideH, + bool is_mla) { + at::parallel_for(0, batches * num_heads_kv, 0, [&](int64_t begin, int64_t end) { + int64_t bs{0}, head_kv_id{0}; + data_index_init(begin, bs, batches, head_kv_id, num_heads_kv); + + for (int64_t i = begin; i < end; i++) { + int64_t loc_val = loc[bs]; + scalar_t* k_buffer_ptr = k_buffer + loc_val * k_strideN + head_kv_id * k_strideH; + const scalar_t* new_key_ptr = key + bs * nk_strideN + head_kv_id * nk_strideH; + copy_stub(k_buffer_ptr, new_key_ptr, head_size); + if (!is_mla) { + scalar_t* v_buffer_ptr = v_buffer + loc_val * v_strideN + head_kv_id * v_strideH; + const scalar_t* new_value_ptr = value + bs * nv_strideN + head_kv_id * nv_strideH; + copy_stub(v_buffer_ptr, new_value_ptr, head_size_v); + } + + // move to the next index + data_index_step(bs, batches, head_kv_id, num_heads_kv); + } + }); +} + +template +void decode_accumulate_kv_splits( + scalar_t* __restrict__ output, + float* __restrict__ attn_logits, + const scalar_t* __restrict__ sinks_ptr, + int64_t batches, + int64_t num_heads, + int64_t head_size_v, + int64_t num_kv_splits, + int64_t l_stride1, + int64_t l_stride2, + bool has_sink) { + using Vec = at::vec::Vectorized; + + // parallel on [batches, num_heads] + at::parallel_for(0, batches * num_heads, 0, [&](int64_t begin, int64_t end) { + int64_t bi{0}, ni{0}; + data_index_init(begin, bi, batches, ni, num_heads); + // NB: here we use logits[b][h][0] as acc, since + // for the first kv split (kv_id == 0): + // m_delta = std::exp(-inf) = 0 + // e_logic = std::exp(0) = 1 + // acc = acc * m_delta + tv * e_logic = tv + for (int64_t i = begin; i < end; ++i) { + float* __restrict__ acc = attn_logits + i * l_stride1; + + float s_prime = 0.f; + float m_prime = -std::numeric_limits::infinity(); + + // update acc with from each kv_split + for (int64_t kv_id = 0; kv_id < num_kv_splits; ++kv_id) { + float* __restrict__ tv = acc + kv_id * l_stride2; + const float tlogic = (acc + kv_id * l_stride2)[head_size_v]; + + float m_i = std::max(tlogic, m_prime); + float m_delta = std::exp(m_prime - m_i); + float e_logic = std::exp(tlogic - m_i); + if (kv_id != 0) { + at::vec::map2( + [m_delta, e_logic](Vec x, Vec y) { return x * Vec(m_delta) + y * Vec(e_logic); }, + acc, + acc, + tv, + head_size_v); + } + + s_prime = s_prime * m_delta + e_logic; + m_prime = m_i; + } + if (has_sink) { + s_prime += std::exp(sinks_ptr[ni] - m_prime); + } + copy_stub(output + i * head_size_v, acc, 1 / s_prime, head_size_v); + // move to the next index + data_index_step(bi, batches, ni, num_heads); + } + }); +} + +template +void decode_attention_kernel_impl( + scalar_t* __restrict__ output, + float* __restrict__ attn_logits, + const scalar_t* __restrict__ query, + const scalar_t* __restrict__ k_buffer, + const scalar_t* __restrict__ v_buffer, + const index_t* __restrict__ req_to_token, + const int64_t* __restrict__ req_pool_indices, + const int64_t* __restrict__ seq_lens, + const int64_t* __restrict__ encoder_lens, + const scalar_t* __restrict__ sinks, + int64_t batches, + int64_t num_heads, + int64_t head_size, + int64_t head_size_v, + int64_t num_kv_splits, + int64_t q_strideM, + int64_t q_strideH, + int64_t k_strideN, + int64_t k_strideH, + int64_t v_strideN, + int64_t v_strideH, + float sm_scale, + float logit_cap, + int64_t max_num_reqs, + int64_t max_context_len, + int64_t max_total_num_tokens, + int64_t sliding_window_size, + bool is_cross_attn, + bool has_encoder_lens, + bool has_sink) { + using Vec = at::vec::Vectorized; + + // strides + const int64_t l_stride1 = num_kv_splits * (head_size_v + 1); + const int64_t l_stride2 = head_size_v + 1; + + const bool has_logit_cap = logit_cap > 0; + float rlogit_cap = has_logit_cap ? 1 / logit_cap : 0.f; + + // parallel on [batches, num_heads, num_kv_splits] + at::parallel_for(0, batches * num_heads * num_kv_splits, 0, [&](int64_t begin, int64_t end) { + int64_t bs{0}, head_id{0}, kv_id{0}; + data_index_init(begin, bs, batches, head_id, num_heads, kv_id, num_kv_splits); + + // s_prime and s_delta + alignas(64) float s_i[BLOCK_N]; + float* __restrict__ s_delta = s_i; + + for (int64_t i = begin; i < end; ++i) { + // get query + const scalar_t* __restrict__ q_ptr = query + bs * q_strideM + head_id * q_strideH; + + // get key/value + int64_t seq_len_kv = is_cross_attn ? encoder_lens[bs] : seq_lens[bs]; + int64_t req_pool_id = req_pool_indices[bs]; + int64_t kv_offset = (has_encoder_lens && (!is_cross_attn)) ? encoder_lens[bs] : 0; + if (sliding_window_size > 0 && seq_len_kv > sliding_window_size) { + kv_offset = seq_len_kv - sliding_window_size; + seq_len_kv = sliding_window_size; + } + TORCH_CHECK(seq_len_kv <= max_context_len, "seq_len_kv out of scope!"); + TORCH_CHECK(req_pool_id < max_num_reqs, "req_pool_id out of scope!"); + + const int64_t SPLIT_SIZE = div_up(seq_len_kv, num_kv_splits); + const int64_t kv_start = kv_id * SPLIT_SIZE; + const int64_t kv_end = std::min(kv_start + SPLIT_SIZE, seq_len_kv); + + float m_prime = -std::numeric_limits::infinity(); + float s_prime = 0.f; + + // get v_prime, and init to zero + float* __restrict__ v_prime = attn_logits + i * (head_size_v + 1); + fill_stub(v_prime, 0.f, head_size_v); + + // loop over K and V sequence with BLOCK_N + for (int64_t n = kv_start; n < kv_end; n += BLOCK_N) { + int64_t n_size = std::min(BLOCK_N, kv_end - n); + + // calculate s_i <- scale * Q @ K + index_gemm_kernel_nt( + /* A */ q_ptr, + /* B */ k_buffer + head_id * k_strideH, + /* C */ s_i, + /* ind */ req_to_token + req_pool_id * max_context_len + n + kv_offset, + /* scl */ sm_scale, + /* M */ 1, + /* N */ n_size, + /* K */ head_size, + /* lda */ 1, + /* ldb */ k_strideN, + /* ldc */ 1, + /* mtt */ max_total_num_tokens); + + // TODO: `tanh` from torch uses sleef u10, going to be slow + if (has_logit_cap) { + at::vec::map( + [logit_cap, rlogit_cap](Vec x) { return Vec(logit_cap) * (x * Vec(rlogit_cap)).tanh(); }, + s_i, + s_i, + n_size); + } + + // m_i: max value per row + float m_i = at::vec::reduce_all([](Vec& x, Vec& y) { return at::vec::maximum(x, y); }, s_i, n_size); + m_i = std::max(m_i, m_prime); + + // m_delta <- exp(m' - m_i) + float m_delta = std::exp(m_prime - m_i); + + // s_delta <- exp(s_i - m_i) + at::vec::map([m_i](Vec x) { return (x - Vec(m_i)).exp_u20(); }, s_delta, s_i, n_size); + + // s' <- s' * m_delta + sum(s_delta) + s_prime *= m_delta; + s_prime += at::vec::reduce_all([](Vec& x, Vec& y) { return x + y; }, s_delta, n_size); + + m_prime = m_i; + + // calculate V' <- s_delta @ V + V' * m_delta + index_gemm_kernel_nn( + /* A */ s_delta, + /* B */ v_buffer + head_id * v_strideH, + /* C */ v_prime, + /* ind */ req_to_token + req_pool_id * max_context_len + n + kv_offset, + /* scl */ &m_delta, + /* M */ 1, + /* N */ head_size_v, + /* K */ n_size, + /* lda */ 1, + /* ldb */ v_strideN, + /* ldc */ 1, + /* mtt */ max_total_num_tokens); + } // loop with KV blocks + + // only update v' when kv_split_size > 0 + if (kv_end > kv_start) { + float s = 1 / s_prime; + at::vec::map([s](Vec out) { return out * Vec(s); }, v_prime, v_prime, head_size_v); + + v_prime[head_size_v] = m_prime + std::log(s_prime); + } else { + v_prime[head_size_v] = -std::numeric_limits::infinity(); + } + + // move to the next index + data_index_step(bs, batches, head_id, num_heads, kv_id, num_kv_splits); + } + }); + + decode_accumulate_kv_splits( + output, attn_logits, sinks, batches, num_heads, head_size_v, num_kv_splits, l_stride1, l_stride2, has_sink); +} // MHA + +template +void decode_attention_mla_kernel_impl( + scalar_t* __restrict__ output, + float* __restrict__ attn_logits, + const scalar_t* __restrict__ query, + const scalar_t* __restrict__ k_buffer, + const scalar_t* __restrict__ v_buffer, + const index_t* __restrict__ req_to_token, + const int64_t* __restrict__ req_pool_indices, + const int64_t* __restrict__ seq_lens, + scalar_t* __restrict__ buffer, + const scalar_t* __restrict__ sinks, + int64_t batches, + int64_t num_heads, + int64_t head_size, + int64_t head_size_v, + int64_t num_kv_splits, + int64_t q_strideM, + int64_t q_strideH, + int64_t k_strideN, + int64_t k_strideH, + int64_t v_strideN, + int64_t v_strideH, + float sm_scale, + float logit_cap, + int64_t max_num_reqs, + int64_t max_context_len, + int64_t max_total_num_tokens, + int64_t buffer_size_per_thread, + bool has_sink) { + using Vec = at::vec::Vectorized; + + // block length for heads + const int64_t BLOCK_H = batches == 1 ? 6 : (batches > 16 ? 22 : 11); + + // strides + const int64_t l_stride0 = num_heads * num_kv_splits * (head_size_v + 1); + const int64_t l_stride1 = num_kv_splits * (head_size_v + 1); + const int64_t l_stride2 = head_size_v + 1; + + TORCH_CHECK(logit_cap == 0.f, "decode MLA: expect no logit_cap."); + + // partition the heads into blocks for parallel + const int64_t num_blocks = div_up(num_heads, BLOCK_H); + + // parallel on [batches, num_blocks, num_kv_splits] + at::parallel_for(0, batches * num_blocks * num_kv_splits, 0, [&](int64_t begin, int64_t end) { + int64_t bs{0}, block_id{0}, kv_id{0}; + data_index_init(begin, bs, batches, block_id, num_blocks, kv_id, num_kv_splits); + + int tid = at::get_thread_num(); + scalar_t* __restrict__ Btmp0 = buffer + tid * buffer_size_per_thread; + scalar_t* __restrict__ Btmp1 = Btmp0 + BLOCK_N * head_size; + + // init Btmp1 just once for each thread to prevent NaN + // Btmp0 is not needed as it computes full K every single time + fill_stub(Btmp1, 0.f, BLOCK_N * head_size_v); + + alignas(64) float s_i[BLOCK_H * BLOCK_N]; + float* __restrict__ s_delta = s_i; + alignas(64) scalar_t s_delta2[BLOCK_H * BLOCK_N]; + + alignas(64) float s_prime[BLOCK_H]; + alignas(64) float m_prime[BLOCK_H]; + alignas(64) float m_delta[BLOCK_H]; + + for (int64_t i = begin; i < end; ++i) { + const int64_t h_start = block_id * BLOCK_H; + const int64_t h_end = std::min(block_id * BLOCK_H + BLOCK_H, num_heads); + const int64_t h_size = h_end - h_start; + + // get query + const scalar_t* __restrict__ q_ptr = query + bs * q_strideM + h_start * q_strideH; + + int64_t seq_len_kv = seq_lens[bs]; + int64_t req_pool_id = req_pool_indices[bs]; + TORCH_CHECK(seq_len_kv <= max_context_len, "seq_len_kv out of scope!"); + TORCH_CHECK(req_pool_id < max_num_reqs, "req_pool_id out of scope!"); + + const int64_t SPLIT_SIZE = div_up(seq_len_kv, num_kv_splits); + const int64_t kv_start = kv_id * SPLIT_SIZE; + const int64_t kv_end = std::min(kv_start + SPLIT_SIZE, seq_len_kv); + + fill_stub(s_prime, 0.f, BLOCK_H); + fill_stub(m_prime, -std::numeric_limits::infinity(), BLOCK_H); + + // get v_prime, and init to zero + float* __restrict__ v_prime = attn_logits + bs * l_stride0 + h_start * l_stride1 + kv_id * l_stride2; + for (int64_t h = 0; h < h_size; ++h) { + fill_stub(v_prime + h * l_stride1, 0.f, head_size_v); + } + + // loop over K and V sequence with BLOCK_N + for (int64_t n = kv_start; n < kv_end; n += BLOCK_N) { + int64_t n_size = std::min(BLOCK_N, kv_end - n); + const int64_t padded_n_size = div_up(int(n_size), TILE_K) * TILE_K; + + // get key and pack + pack_vnni( + /* dst0 */ Btmp0, + /* dst1 */ Btmp1, + /* src */ k_buffer + /* head_kv_id */ 0 * k_strideH, + /* ind */ req_to_token + req_pool_id * max_context_len + n, + /* N */ n_size, + /* K */ head_size, + /* Kv */ head_size_v, + /* ld_src */ k_strideN, + /* ld_dst0 */ BLOCK_N, + /* ld_dst1 */ head_size_v); + + // calculate s_i <- Q @ K + at::native::cpublas::brgemm( + /* M */ h_size, + /* N */ n_size, + /* K */ head_size, + /* lda */ q_strideH, + /* ldb */ BLOCK_N, + /* ldc */ BLOCK_N, + /* add_C */ false, + /* A */ q_ptr, + /* B */ Btmp0, + /* C */ s_i); + + const Vec scale_vec = Vec(sm_scale); + for (int64_t h = 0; h < h_size; ++h) { + // s_i <- s_i * scale + at::vec::map( + [scale_vec](Vec x) { return x * scale_vec; }, s_i + h * BLOCK_N, s_i + h * BLOCK_N, n_size); + + // m_i: max value per row + float m_i = at::vec::reduce_all( + [](Vec& x, Vec& y) { return at::vec::maximum(x, y); }, s_i + h * BLOCK_N, n_size); + m_i = std::max(m_i, m_prime[h]); + + // m_delta <- exp(m' - m_i) + m_delta[h] = std::exp(m_prime[h] - m_i); + + // s_delta <- exp(s_i - m_i) + at::vec::map( + [m_i](Vec x) { return (x - Vec(m_i)).exp_u20(); }, s_delta + h * BLOCK_N, s_i + h * BLOCK_N, n_size); + + // s' <- s' * m_delta + sum(s_delta) + s_prime[h] *= m_delta[h]; + s_prime[h] += at::vec::reduce_all([](Vec& x, Vec& y) { return x + y; }, s_delta + h * BLOCK_N, n_size); + + m_prime[h] = m_i; + + // v' <- v' * m_delta + float scale_m = m_delta[h]; + at::vec::map( + [scale_m](Vec x) { return x * Vec(scale_m); }, + v_prime + h * l_stride1, + v_prime + h * l_stride1, + head_size_v); + + // pad s_delta with 0 first and then convert to scalar_t + fill_stub(s_delta + h * BLOCK_N + n_size, 0.f, padded_n_size - n_size); + copy_stub(s_delta2 + h * BLOCK_N, s_delta + h * BLOCK_N); + } + + // calculate V' <- s_delta @ V + V' + at::native::cpublas::brgemm( + /* M */ h_size, + /* N */ head_size_v, + /* K */ padded_n_size, // n_size + /* lda */ BLOCK_N, + /* ldb */ head_size_v, + /* ldc */ l_stride1, + /* add_C */ true, + /* A */ s_delta2, + /* B */ Btmp1, + /* C */ v_prime); + } // loop with KV blocks + + // only update v' when kv_split_size > 0 + if (kv_end > kv_start) { + for (int64_t h = 0; h < h_size; ++h) { + float s = 1 / s_prime[h]; + at::vec::map( + [s](Vec out) { return out * Vec(s); }, v_prime + h * l_stride1, v_prime + h * l_stride1, head_size_v); + (v_prime + h * l_stride1)[head_size_v] = m_prime[h] + std::log(s_prime[h]); + } + } else { + for (int64_t h = 0; h < h_size; ++h) { + (v_prime + h * l_stride1)[head_size_v] = -std::numeric_limits::infinity(); + } + } + + // move to the next index + data_index_step(bs, batches, block_id, num_blocks, kv_id, num_kv_splits); + } + at::native::cpublas::brgemm_release(); + }); + + decode_accumulate_kv_splits( + output, attn_logits, sinks, batches, num_heads, head_size_v, num_kv_splits, l_stride1, l_stride2, has_sink); +} // MLA + +template +void decode_attention_grouped_kernel_impl( + scalar_t* __restrict__ output, + float* __restrict__ attn_logits, + const scalar_t* __restrict__ query, + const scalar_t* __restrict__ k_buffer, + const scalar_t* __restrict__ v_buffer, + const index_t* __restrict__ req_to_token, + const int64_t* __restrict__ req_pool_indices, + const int64_t* __restrict__ seq_lens, + const int64_t* __restrict__ encoder_lens, + const scalar_t* __restrict__ sinks, + int64_t batches, + int64_t num_heads, + int64_t num_heads_kv, + int64_t head_size, + int64_t head_size_v, + int64_t num_kv_splits, + int64_t q_strideM, + int64_t q_strideH, + int64_t k_strideN, + int64_t k_strideH, + int64_t v_strideN, + int64_t v_strideH, + float sm_scale, + float logit_cap, + int64_t max_num_reqs, + int64_t max_context_len, + int64_t max_total_num_tokens, + int64_t sliding_window_size, + bool is_cross_attn, + bool has_encoder_lens, + bool has_sink) { + using Vec = at::vec::Vectorized; + + // block length for heads + // we parallel on [batches, divup(num_heads, BLOCK_H), num_kv_splits] + // use smaller BLOCK_H when batches is small to utilize all cores + constexpr int64_t kBLOCK_H = 16; + const int64_t BLOCK_H = std::min(4 * batches, kBLOCK_H); + + // strides + const int64_t l_stride0 = num_heads * num_kv_splits * (head_size_v + 1); + const int64_t l_stride1 = num_kv_splits * (head_size_v + 1); + const int64_t l_stride2 = head_size_v + 1; + + const bool has_logit_cap = logit_cap > 0; + float rlogit_cap = has_logit_cap ? 1 / logit_cap : 0.f; + + // partition the heads into blocks for parallel + const int64_t num_groups = num_heads / num_heads_kv; + const int64_t num_blocks = div_up(num_groups, BLOCK_H); + + // parallel on [batches, num_heads_kv, num_blocks, num_kv_splits] + at::parallel_for(0, batches * num_heads_kv * num_blocks * num_kv_splits, 0, [&](int64_t begin, int64_t end) { + int64_t bs{0}, head_kv_id{0}, block_id{0}, kv_id{0}; + data_index_init(begin, bs, batches, head_kv_id, num_heads_kv, block_id, num_blocks, kv_id, num_kv_splits); + + alignas(64) float s_i[BLOCK_H * BLOCK_N]; + float* __restrict__ s_delta = s_i; + + alignas(64) float s_prime[BLOCK_H]; + alignas(64) float m_prime[BLOCK_H]; + alignas(64) float m_delta[BLOCK_H]; + + for (int64_t i = begin; i < end; ++i) { + const int64_t h_start = head_kv_id * num_groups + block_id * BLOCK_H; + const int64_t h_end = head_kv_id * num_groups + std::min(block_id * BLOCK_H + BLOCK_H, num_groups); + const int64_t h_size = h_end - h_start; + + // get query + const scalar_t* __restrict__ q_ptr = query + bs * q_strideM + h_start * q_strideH; + + int64_t seq_len_kv = is_cross_attn ? encoder_lens[bs] : seq_lens[bs]; + int64_t req_pool_id = req_pool_indices[bs]; + int64_t kv_offset = (has_encoder_lens && (!is_cross_attn)) ? encoder_lens[bs] : 0; + TORCH_CHECK(seq_len_kv <= max_context_len, "seq_len_kv out of scope!"); + TORCH_CHECK(req_pool_id < max_num_reqs, "req_pool_id out of scope!"); + if (sliding_window_size > 0 && seq_len_kv > sliding_window_size) { + kv_offset = seq_len_kv - sliding_window_size; + seq_len_kv = sliding_window_size; + } + const int64_t SPLIT_SIZE = div_up(seq_len_kv, num_kv_splits); + const int64_t kv_start = kv_id * SPLIT_SIZE; + const int64_t kv_end = std::min(kv_start + SPLIT_SIZE, seq_len_kv); + + fill_stub(s_prime, 0.f, BLOCK_H); + fill_stub(m_prime, -std::numeric_limits::infinity(), BLOCK_H); + + // get v_prime, and init to zero + float* __restrict__ v_prime = attn_logits + bs * l_stride0 + h_start * l_stride1 + kv_id * l_stride2; + for (int64_t h = 0; h < h_size; ++h) { + fill_stub(v_prime + h * l_stride1, 0.f, head_size_v); + } + + // loop over K and V sequence with BLOCK_N + for (int64_t n = kv_start; n < kv_end; n += BLOCK_N) { + int64_t n_size = std::min(BLOCK_N, kv_end - n); + + // calculate Q @ K + index_gemm_kernel_nt( + /* A */ q_ptr, + /* B */ k_buffer + head_kv_id * k_strideH, + /* C */ s_i, + /* ind */ req_to_token + req_pool_id * max_context_len + n + kv_offset, + /* scl */ sm_scale, + /* M */ h_size, + /* N */ n_size, + /* K */ head_size, + /* lda */ q_strideH, + /* ldb */ k_strideN, + /* ldc */ BLOCK_N, + /* mtt */ max_total_num_tokens); + + if (has_logit_cap) { + at::vec::map( + [logit_cap, rlogit_cap](Vec x) { return Vec(logit_cap) * (x * Vec(rlogit_cap)).tanh(); }, + s_i, + s_i, + BLOCK_H * BLOCK_N); + } + + // update the sm_scale coefficients + for (int64_t h = 0; h < h_size; ++h) { + // m_i: max value per row + float m_i = at::vec::reduce_all( + [](Vec& x, Vec& y) { return at::vec::maximum(x, y); }, s_i + h * BLOCK_N, n_size); + m_i = std::max(m_i, m_prime[h]); + + // m_delta <- exp(m' - m_i) + m_delta[h] = std::exp(m_prime[h] - m_i); + + // s_delta <- exp(s_i - m_i) + at::vec::map( + [m_i](Vec x) { return (x - Vec(m_i)).exp_u20(); }, s_delta + h * BLOCK_N, s_i + h * BLOCK_N, n_size); + + // s' <- s' * m_delta + sum(s_delta) + s_prime[h] *= m_delta[h]; + s_prime[h] += at::vec::reduce_all([](Vec& x, Vec& y) { return x + y; }, s_delta + h * BLOCK_N, n_size); + + m_prime[h] = m_i; + } + + // calculate V' <- s_delta @ V + V' * m_delta + index_gemm_kernel_nn( + /* A */ s_delta, + /* B */ v_buffer + head_kv_id * v_strideH, + /* C */ v_prime, + /* ind */ req_to_token + req_pool_id * max_context_len + n + kv_offset, + /* scl */ m_delta, + /* M */ h_size, + /* N */ head_size_v, + /* K */ n_size, + /* lda */ BLOCK_N, + /* ldb */ v_strideN, + /* ldc */ l_stride1, + /* mtt */ max_total_num_tokens); + } // loop with KV blocks + + // only update v' when kv_split_size > 0 + if (kv_end > kv_start) { + for (int64_t h = 0; h < h_size; ++h) { + float s = 1 / s_prime[h]; + at::vec::map( + [s](Vec out) { return out * Vec(s); }, v_prime + h * l_stride1, v_prime + h * l_stride1, head_size_v); + (v_prime + h * l_stride1)[head_size_v] = m_prime[h] + std::log(s_prime[h]); + } + } else { + for (int64_t h = 0; h < h_size; ++h) { + (v_prime + h * l_stride1)[head_size_v] = -std::numeric_limits::infinity(); + } + } + + // move to the next index + data_index_step(bs, batches, head_kv_id, num_heads_kv, block_id, num_blocks, kv_id, num_kv_splits); + } + }); + + decode_accumulate_kv_splits( + output, attn_logits, sinks, batches, num_heads, head_size_v, num_kv_splits, l_stride1, l_stride2, has_sink); +} // GQA/MQA + +} // anonymous namespace + +// query: [num_tokens, num_heads, head_size] +// output: [num_tokens, num_heads, head_size] +// k_buffer: [max_total_num_tokens, num_heads, head_size] +// v_buffer: [max_total_num_tokens, num_heads, head_size_v] +// attn_logits: [num_seqs, num_heads, num_kv_splits, head_size_v + 1] +// req_to_token: [max_num_reqs, max_context_len] int32 or int64 +// req_pool_indices: [num_seqs] int64 +// seq_lens: [num_seqs] int64 +// encoder_lens: [num_seqs] int64 or None +// sinks: [num_heads] or None +void decode_attention_cpu( + at::Tensor& query, + at::Tensor& k_buffer, + at::Tensor& v_buffer, + at::Tensor& output, + const std::optional& key, + const std::optional& value, + const std::optional& loc, + at::Tensor& attn_logits, + at::Tensor& req_to_token, + at::Tensor& req_pool_indices, + at::Tensor& seq_lens, + double sm_scale, + double logit_cap, + bool is_cross_attn, + int64_t sliding_window_size, + std::optional encoder_lens, + std::optional sinks) { + CHECK_LAST_DIM_CONTIGUOUS_INPUT(query); + CHECK_LAST_DIM_CONTIGUOUS_INPUT(k_buffer); + CHECK_LAST_DIM_CONTIGUOUS_INPUT(v_buffer); + CHECK_DIM(3, query); + CHECK_DIM(3, k_buffer); + CHECK_DIM(3, v_buffer); + + int64_t num_seqs = seq_lens.size(0); + int64_t max_num_reqs = req_to_token.size(0); + int64_t max_context_len = req_to_token.size(1); + int64_t max_total_num_tokens = k_buffer.size(0); + + int64_t num_heads = query.size(1); + int64_t num_heads_kv = k_buffer.size(1); + int64_t head_size = query.size(2); + int64_t head_size_v = v_buffer.size(2); + + int64_t num_kv_splits = attn_logits.size(2); + + CHECK_EQ(attn_logits.size(0), num_seqs); + CHECK_EQ(attn_logits.size(1), num_heads); + CHECK_EQ(attn_logits.size(3), head_size_v + 1); + CHECK_EQ(attn_logits.scalar_type(), at::kFloat); + + // strides for query + int64_t q_strideM = query.stride(0); + int64_t q_strideH = query.stride(1); + + // strides for k_buffer and v_buffer + int64_t k_strideN = k_buffer.stride(0); + int64_t k_strideH = k_buffer.stride(1); + int64_t v_strideN = v_buffer.stride(0); + int64_t v_strideH = v_buffer.stride(1); + + // check index data types + const auto index_dtype = req_to_token.scalar_type(); + TORCH_CHECK( + index_dtype == at::kInt || index_dtype == at::kLong, + "decode: expect req_to_token to be int32 or int64, got ", + index_dtype); + TORCH_CHECK(seq_lens.scalar_type() == at::kLong, "decode: expect req_lens to be int64, got ", seq_lens.scalar_type()); + TORCH_CHECK( + req_pool_indices.scalar_type() == at::kLong, + "decode: expect req_pool_indices to be int64, got ", + req_pool_indices.scalar_type()); + + // check if we have MLA here + void* k_buffer_data = k_buffer.data_ptr(); + void* v_buffer_data = v_buffer.data_ptr(); + const bool is_mla = (k_buffer_data == v_buffer_data) && (num_heads_kv == 1) && (head_size == head_size_v + 64); + + // block length for k_buffer and v_buffer + constexpr int BLOCK_N = 256; + + // buffer for packing k_cache and v_cache + int num_threads = at::get_num_threads(); + int64_t size_per_thread = is_mla ? BLOCK_N * head_size + BLOCK_N * head_size_v : 0; + auto buffer = at::empty({num_threads, size_per_thread}, k_buffer.options()); + bool has_encoder_lens = encoder_lens.has_value(); + // Since encoder_lens is not used when it is None, encoder_lens_t can be initialized as any tensor of int64_t dtype. + at::Tensor encoder_lens_t = seq_lens; + if (has_encoder_lens) { + encoder_lens_t = encoder_lens.value(); + CHECK_EQ(encoder_lens_t.size(0), num_seqs); + } + bool has_sink = sinks.has_value(); + at::Tensor sinks_tensor = has_sink ? sinks.value() : at::empty({num_heads}, query.options()); + CHECK_DIM(1, sinks_tensor); + CHECK_EQ(sinks_tensor.size(0), num_heads); + AT_DISPATCH_REDUCED_FLOATING_TYPES(query.scalar_type(), "decode_attention_kernel", [&] { + AT_DISPATCH_INDEX_TYPES(index_dtype, "decode_attention_indices", [&] { + if (key.has_value()) { + TORCH_CHECK(value.has_value(), "key and value should have values at the same time") + TORCH_CHECK(loc.has_value(), "loc must be given when key/value are given") + auto loc_tensor = loc.value(); + CHECK_DIM(1, loc_tensor); + CHECK_EQ(loc_tensor.numel(), num_seqs); + auto key_tensor = key.value(); + auto value_tensor = value.value(); + // for MLA, key and value shares the same storage and value could be non-contiguous + CHECK_LAST_DIM_CONTIGUOUS_INPUT(key_tensor); + CHECK_LAST_DIM_CONTIGUOUS_INPUT(value_tensor); + CHECK_DIM(3, key_tensor); + CHECK_DIM(3, value_tensor); + // strides for new key and value + int64_t nk_strideN = key_tensor.stride(0); + int64_t nk_strideH = key_tensor.stride(1); + int64_t nv_strideN = value_tensor.stride(0); + int64_t nv_strideH = value_tensor.stride(1); + // update the kv buffer + decode_set_kv_buffer( + (scalar_t*)k_buffer_data, + (scalar_t*)v_buffer_data, + key_tensor.data_ptr(), + value_tensor.data_ptr(), + loc_tensor.data_ptr(), + num_seqs, + num_heads_kv, + head_size, + head_size_v, + k_strideN, + k_strideH, + v_strideN, + v_strideH, + nk_strideN, + nk_strideH, + nv_strideN, + nv_strideH, + is_mla); + } + + if (num_heads == num_heads_kv) { + // MHA + decode_attention_kernel_impl( + output.data_ptr(), + attn_logits.data_ptr(), + query.data_ptr(), + (const scalar_t*)k_buffer_data, + (const scalar_t*)v_buffer_data, + req_to_token.data_ptr(), + req_pool_indices.data_ptr(), + seq_lens.data_ptr(), + encoder_lens_t.data_ptr(), + sinks_tensor.data_ptr(), + num_seqs, + num_heads, + head_size, + head_size_v, + num_kv_splits, + q_strideM, + q_strideH, + k_strideN, + k_strideH, + v_strideN, + v_strideH, + sm_scale, + logit_cap, + max_num_reqs, + max_context_len, + max_total_num_tokens, + sliding_window_size, + is_cross_attn, + has_encoder_lens, + has_sink); + } else if (is_mla) { + // MLA + decode_attention_mla_kernel_impl( + output.data_ptr(), + attn_logits.data_ptr(), + query.data_ptr(), + (const scalar_t*)k_buffer_data, + (const scalar_t*)v_buffer_data, + req_to_token.data_ptr(), + req_pool_indices.data_ptr(), + seq_lens.data_ptr(), + buffer.data_ptr(), + sinks_tensor.data_ptr(), + num_seqs, + num_heads, + head_size, + head_size_v, + num_kv_splits, + q_strideM, + q_strideH, + k_strideN, + k_strideH, + v_strideN, + v_strideH, + sm_scale, + logit_cap, + max_num_reqs, + max_context_len, + max_total_num_tokens, + size_per_thread, + has_sink); + } else { + // GQA/MQA + decode_attention_grouped_kernel_impl( + output.data_ptr(), + attn_logits.data_ptr(), + query.data_ptr(), + (const scalar_t*)k_buffer_data, + (const scalar_t*)v_buffer_data, + req_to_token.data_ptr(), + req_pool_indices.data_ptr(), + seq_lens.data_ptr(), + encoder_lens_t.data_ptr(), + sinks_tensor.data_ptr(), + num_seqs, + num_heads, + num_heads_kv, + head_size, + head_size_v, + num_kv_splits, + q_strideM, + q_strideH, + k_strideN, + k_strideH, + v_strideN, + v_strideH, + sm_scale, + logit_cap, + max_num_reqs, + max_context_len, + max_total_num_tokens, + sliding_window_size, + is_cross_attn, + has_encoder_lens, + has_sink); + } + }); + }); +} diff --git a/csrc/cpu/sgl-kernels/extend.cpp b/csrc/cpu/sgl-kernels/extend.cpp new file mode 100644 index 0000000000..3c0a9ebe54 --- /dev/null +++ b/csrc/cpu/sgl-kernels/extend.cpp @@ -0,0 +1,568 @@ +// Adapted from +// https://github.com/sgl-project/sglang/tree/main/sgl-kernel/csrc/cpu + +// clang-format off + +#include "common.h" +#include "flash_attn.h" +#include "gemm.h" + +namespace { + +// [NOTE]: extend attention for CPU +// 1. BLOCK_M and BLOCK_N tuned for various seq lengths +// 2. can handle non-contiguous k_extend and v_extend +// 3. computes attention for prefix and extend separately +// 4. TODO: apply head dimension blocking to optimize GQA +// 5. optional tree mask for speculative decoding TARGET_VERIFY (EAGLE topk > 1): +// `tree_mask` is a flat [batches * qlen * qlen] bool tensor in +// TreeMaskMode::QLEN_ONLY layout, where qlen == extend_seq_lens[bs] == +// max_len_extend (uniform across the batch, equal to draft_token_num). +// Row i = query draft token, column j = key draft token; true means query i +// may attend key j (each row marks self + ancestors + root). The committed +// prefix (stage 1) is implicitly fully visible to every draft token, which +// is why the mask only covers the qlen x qlen new-token block; the GPU +// FULL_MASK layout carries the prefix columns explicitly but they are +// all-true for EAGLE. When tree_mask is absent, stage 2 falls back to the +// plain causal mask (correct for non-spec extend and topk == 1 chains). +// + +template +void extend_attention_kernel_impl( + scalar_t* __restrict__ o_extend, + const scalar_t* __restrict__ q_extend, + const scalar_t* __restrict__ k_extend, + const scalar_t* __restrict__ v_extend, + const scalar_t* __restrict__ k_buffer, + const scalar_t* __restrict__ v_buffer, + const index_t* __restrict__ req_to_token, + const int64_t* __restrict__ req_pool_indices, + const int64_t* __restrict__ seq_lens, + const int64_t* __restrict__ encoder_lens, + const index_t* __restrict__ extend_seq_lens, + const index_t* __restrict__ extend_start_loc, + const void* __restrict__ buffer, + const scalar_t* __restrict__ sinks, + const bool* __restrict__ tree_mask, + int batches, + int num_heads, + int num_heads_kv, + int head_size, + int head_size_v, + int q_strideM, + int q_strideH, + int ke_strideN, + int ke_strideH, + int ve_strideN, + int ve_strideH, + int k_strideN, + int k_strideH, + int v_strideN, + int v_strideH, + float sm_scale, + int max_num_reqs, + int max_context_len, + int max_total_num_tokens, + int max_len_extend, + int buffer_size_per_thread, + int64_t sliding_window_size, + bool is_prefix_skipped, + bool is_cross_attn, + bool has_encoder_lens, + bool has_sink) { + // strides + const int o_strideM = num_heads * head_size_v; + const int o_strideH = head_size_v; + + // we use same buffer for packed key and value + const int ldb_tmp = std::max(head_size, head_size_v); + + const int num_groups = num_heads / num_heads_kv; + TORCH_CHECK(num_groups * num_heads_kv == num_heads); + + // number of blocks along M + int MB = div_up(max_len_extend, BLOCK_M); + + // parallel on [batches, num_heads, BM] + at::parallel_for(0, batches * num_heads * MB, 0, [&](int begin, int end) { + int bs{0}, head_id{0}, mb{0}; + data_index_init(begin, bs, batches, head_id, num_heads, mb, MB); + + int tid = at::get_thread_num(); + // s_i: [BLOCK_M, BLOCK_N] + float* __restrict__ s_i = reinterpret_cast((char*)(buffer) + tid * buffer_size_per_thread); + + // v_prime: [BLOCK_M, head_size_v] + float* __restrict__ v_prime = s_i + BLOCK_M * BLOCK_N; + + // s_delta: [BLOCK_M, BLOCK_N] + scalar_t* __restrict__ s_delta = reinterpret_cast(v_prime + BLOCK_M * head_size_v); + + // Btmp: [BLOCK_N, max(head_size, head_size_v)] + scalar_t* __restrict__ Btmp = reinterpret_cast(s_delta + BLOCK_M * BLOCK_N); + + // init Btmp just once for each thread to prevent NaN + fill_stub(Btmp, 0.f, BLOCK_N * ldb_tmp); + fill_stub(s_delta, 0.f, BLOCK_M * BLOCK_N); + + alignas(64) float s_prime[BLOCK_M]; + alignas(64) float m_prime[BLOCK_M]; + + for (int i = begin; i < end; ++i) { + // seq_len = prefix + extend + int head_kv_id = head_id / num_groups; + int seq_len = seq_lens[bs]; + int seq_len_extend = extend_seq_lens[bs]; + int seq_len_prefix = seq_len - seq_len_extend; + int seq_extend_start_loc = extend_start_loc[bs]; + + int req_pool_id = req_pool_indices[bs]; + int kv_offset = (has_encoder_lens && (!is_cross_attn)) ? encoder_lens[bs] : 0; + TORCH_CHECK(seq_len_prefix >= 0, "prefix len < 0!"); + TORCH_CHECK(seq_len <= max_context_len, "seq_len out of scope!"); + TORCH_CHECK(req_pool_id < max_num_reqs, "req_pool_id out of scope!"); + + if (is_prefix_skipped) { + TORCH_CHECK(seq_len_prefix == 0, "extend attention: expect seq_len_prefix to be 0, got ", seq_len_prefix); + } + + if (tree_mask != nullptr) { + // QLEN_ONLY layout assumes a uniform qlen across the batch (TARGET_VERIFY) + TORCH_CHECK( + seq_len_extend == max_len_extend, + "extend attention: tree_mask requires uniform extend_seq_lens, got ", + seq_len_extend, + " vs ", + max_len_extend); + } + + // offset and size in MB + int m = mb * BLOCK_M; + int m_size = std::min(BLOCK_M, seq_len_extend - m); + + if (m_size <= 0) { + data_index_step(bs, batches, head_id, num_heads, mb, MB); + continue; + } + + // get query + const scalar_t* __restrict__ q_ptr = q_extend + (seq_extend_start_loc + m) * q_strideM + head_id * q_strideH; + + // init v', s' and m' + fill_stub(v_prime, 0.f, m_size * head_size_v); + fill_stub(s_prime, 0.f, m_size); + fill_stub(m_prime, -std::numeric_limits::infinity(), m_size); + // stage 1: compute scores with prefix + int kv_start = 0; + int kv_end = is_cross_attn ? encoder_lens[bs] : seq_len_prefix; + for (int n = kv_start; n < kv_end; n += BLOCK_N) { + int n_size = std::min(BLOCK_N, kv_end - n); + + // `n_size` is K in 2nd gemm, pad to TILE_K; + const int padded_n_size = div_up(n_size, TILE_K) * TILE_K; + + // get key and pack + pack_vnni( + /* dst */ Btmp, + /* src */ k_buffer + head_kv_id * k_strideH, + /* ind */ req_to_token + req_pool_id * max_context_len + n + kv_offset, + /* N */ n_size, + /* K */ head_size, + /* ld_src */ k_strideN, + /* ld_dst */ BLOCK_N); + + // calculate s_i <- Q @ K + at::native::cpublas::brgemm( + /* M */ m_size, + /* N */ n_size, + /* K */ head_size, + /* lda */ q_strideM, + /* ldb */ BLOCK_N, + /* ldc */ BLOCK_N, + /* add_C */ false, + /* A */ q_ptr, + /* B */ Btmp, + /* C */ s_i); + + for (int row = 0; row < m_size; ++row) { + if (sliding_window_size > 0) { + int last_col = seq_len_prefix + row + m - sliding_window_size + 1; + if (last_col >= n + n_size) { + continue; + } + fill_stub(s_i + row * BLOCK_N, -std::numeric_limits::infinity(), last_col - n); + } + flash_attn_softmax::apply( + s_i, s_delta, v_prime, s_prime, m_prime, m_size, n_size, padded_n_size, head_size_v, sm_scale, row); + } + + // get value and pack + pack_vnni2( + /* dst */ Btmp, + /* src */ v_buffer + head_kv_id * v_strideH, + /* ind */ req_to_token + req_pool_id * max_context_len + n + kv_offset, + /* K */ n_size, + /* N */ head_size_v, + /* ld_src */ v_strideN, + /* ld_dst */ head_size_v); + + // calculate V' <- s_delta @ V + V' + at::native::cpublas::brgemm( + /* M */ m_size, + /* N */ head_size_v, + /* K */ padded_n_size, // n_size + /* lda */ BLOCK_N, + /* ldb */ head_size_v, + /* ldc */ head_size_v, + /* add_C */ true, + /* A */ s_delta, + /* B */ Btmp, + /* C */ v_prime); + } // loop with seq_len_prefix + if (!is_cross_attn) { + // stage 2: compute the triangle part + int num_keys = std::min(seq_len_extend, m + BLOCK_M); + for (int n = 0; n < num_keys; n += BLOCK_N) { + int n_size = std::min(BLOCK_N, num_keys - n); + + // `n_size` is K in 2nd gemm, pad to TILE_K; + const int padded_n_size = div_up(n_size, TILE_K) * TILE_K; + + // get key and pack + pack_vnni( + /* dst */ Btmp, + /* src */ k_extend + (seq_extend_start_loc + n) * ke_strideN + head_kv_id * ke_strideH, + /* N */ n_size, + /* K */ head_size, + /* ld_src */ ke_strideN, + /* ld_dst */ BLOCK_N); + + // calculate s_i <- Q @ K + at::native::cpublas::brgemm( + /* M */ m_size, + /* N */ n_size, + /* K */ head_size, + /* lda */ q_strideM, + /* ldb */ BLOCK_N, + /* ldc */ BLOCK_N, + /* add_C */ false, + /* A */ q_ptr, + /* B */ Btmp, + /* C */ s_i); + + // apply tree mask (speculative TARGET_VERIFY) or causal mask + if (tree_mask != nullptr) { + // [Note] tree mask for EAGLE topk > 1 (TreeMaskMode::QLEN_ONLY). + // mask[bs][m + row][n + col] == false -> query draft token (m + row) + // may not attend key draft token (n + col); set the score to -inf + // before softmax. The tree mask subsumes the causal constraint: + // ancestors always precede descendants in the draft token ordering, + // so permitted keys satisfy j <= i and the causal `num_keys` bound + // above remains valid. + const bool* __restrict__ mask_base = + tree_mask + (static_cast(bs) * seq_len_extend + m) * seq_len_extend + n; + for (int row = 0; row < m_size; ++row) { + float* __restrict__ row_ptr = s_i + row * BLOCK_N; + const bool* __restrict__ mask_ptr = mask_base + static_cast(row) * seq_len_extend; + for (int col = 0; col < n_size; ++col) { + if (!mask_ptr[col]) { + row_ptr[col] = -std::numeric_limits::infinity(); + } + } + } + } else if (n + n_size - 1 > m) { + // apply causal mask + // [Note] condition to apply causal mask. + // Mask any block whose last key (n + n_size - 1) is strictly after the first query position (m), i.e. n + + // n_size - 1 > m. The original condition was `num_keys - n <= BLOCK_N` (last n-block only). That was + // correct when BLOCK_M <= BLOCK_N/2 because earlier n-blocks were guaranteed to contain only past keys. + // With BLOCK_M=512, BLOCK_N=768: + // BLOCK_M > BLOCK_N/2, so the first n-block can contain future keys. + // Example: m=512 (mb=1), num_keys=1024, first n-block covers keys [0, 768). + // Query row=0 is at position 512, so keys 513..767 are future and must be + // masked — but `num_keys - 0 = 1024 > BLOCK_N` skips masking entirely, + // producing wrong (non-causal) attention for rows 0..254 of this m-block. + for (int row = 0; row < m_size; ++row) { + int last_col = m + row - n; + // [Note] mask the entire row if last_col < 0. + // Clamp to -1: when n > m + row every key in this block is a future + // key, so the entire row should be masked. Without this clamp, + // last_col+1 <= 0 and fill_stub would write before row_ptr. + last_col = std::max(last_col, -1); + // fill [last_col + 1, n_size) to -inf + float* row_ptr = s_i + row * BLOCK_N; + fill_stub(row_ptr + last_col + 1, -std::numeric_limits::infinity(), n_size - last_col - 1); + } + } + + for (int row = 0; row < m_size; ++row) { + if (sliding_window_size > 0 && row + m + 1 >= n + sliding_window_size - 1 && + row + m + 1 < n + sliding_window_size + n_size) { + fill_stub( + s_i + row * BLOCK_N, -std::numeric_limits::infinity(), row + m - n - sliding_window_size + 1); + } else if (sliding_window_size > 0 && row + m + 1 >= n + sliding_window_size) { + continue; + } + flash_attn_softmax::apply( + s_i, s_delta, v_prime, s_prime, m_prime, m_size, n_size, padded_n_size, head_size_v, sm_scale, row); + } + + // get value and pack + pack_vnni2( + /* dst */ Btmp, + /* src */ v_extend + (seq_extend_start_loc + n) * ve_strideN + head_kv_id * ve_strideH, + /* K */ n_size, + /* N */ head_size_v, + /* ld_src */ ve_strideN, + /* ld_dst */ head_size_v); + + // calculate V' <- s_delta @ V + V' + at::native::cpublas::brgemm( + /* M */ m_size, + /* N */ head_size_v, + /* K */ padded_n_size, // n_size + /* lda */ BLOCK_N, + /* ldb */ head_size_v, + /* ldc */ head_size_v, + /* add_C */ true, + /* A */ s_delta, + /* B */ Btmp, + /* C */ v_prime); + } // loop with seq_len_extend + } + scalar_t* __restrict__ out_ptr = o_extend + (seq_extend_start_loc + m) * o_strideM + head_id * o_strideH; + for (int row = 0; row < m_size; ++row) { + if (has_sink) { + s_prime[row] += std::exp(sinks[head_id] - m_prime[row]); + } + float s = 1 / s_prime[row]; + copy_stub(out_ptr + row * o_strideM, v_prime + row * head_size_v, s, head_size_v); + } + + // move to the next index + data_index_step(bs, batches, head_id, num_heads, mb, MB); + } + at::native::cpublas::brgemm_release(); + }); +} + +} // anonymous namespace + +template +inline int resize_buffer(at::Tensor& buffer, int num_threads, int head_size, int head_size_v) { + static_assert(BLOCK_M <= BLOCK_N, "Make sure BLOCK_M <= BLOCK_N to prevent buffer overflows during causal masking"); + const int size_per_thread = + /* s_i */ BLOCK_M * BLOCK_N * sizeof(float) + + /* v_prime */ BLOCK_M * head_size_v * sizeof(float) + + /* s_delta */ BLOCK_M * BLOCK_N * sizeof(uint16_t) + + /* Btmp */ BLOCK_N * std::max(head_size, head_size_v) * sizeof(uint16_t); + + buffer.resize_({num_threads, size_per_thread}); + return size_per_thread; +} + +#define LAUNCH_EXTEND_ATTENTION_KERNEL(BLOCK_M, BLOCK_N) \ + do { \ + int sz = resize_buffer(buffer, num_threads, head_size, head_size_v); \ + \ + extend_attention_kernel_impl( \ + o_extend.data_ptr(), \ + q_extend.data_ptr(), \ + k_extend.data_ptr(), \ + v_extend.data_ptr(), \ + k_buffer.data_ptr(), \ + v_buffer.data_ptr(), \ + req_to_token.data_ptr(), \ + req_pool_indices.data_ptr(), \ + seq_lens.data_ptr(), \ + encoder_lens_t.data_ptr(), \ + extend_seq_lens.data_ptr(), \ + extend_start_loc.data_ptr(), \ + buffer.data_ptr(), \ + sinks_tensor.data_ptr(), \ + tree_mask_ptr, \ + num_seqs, \ + num_heads, \ + num_heads_kv, \ + head_size, \ + head_size_v, \ + q_strideM, \ + q_strideH, \ + ke_strideN, \ + ke_strideH, \ + ve_strideN, \ + ve_strideH, \ + k_strideN, \ + k_strideH, \ + v_strideN, \ + v_strideH, \ + sm_scale, \ + max_num_reqs, \ + max_context_len, \ + max_total_num_tokens, \ + max_len_extend, \ + sz, \ + sliding_window_size, \ + is_prefix_skipped, \ + is_cross_attn, \ + has_encoder_lens, \ + has_sink); \ + } while (0) + +// q_extend, k_extend, v_extend, o_extend: contiguous tensors +// k_buffer, v_buffer: (prefix + extend) tensors in mem_manager +// +// q_extend: [num_tokens, num_heads, head_size] +// k_extend: [num_extend_tokens, num_heads, head_size] +// v_extend: [num_extend_tokens, num_heads, head_size] +// o_extend: [num_tokens, num_heads, head_size] +// k_buffer: [max_total_num_tokens, num_heads, head_size] +// v_buffer: [max_total_num_tokens, num_heads, head_size] +// req_to_token: [max_num_reqs, max_context_len] int32 or int64 +// req_pool_indices: [num_seqs] int64 +// seq_lens: [num_seqs] int64 +// extend_seq_lens: [num_seqs] +// extend_start_loc: [num_seqs] +// encoder_lens: [num_seqs] int64 or None +// sinks: [num_heads] or None +// tree_mask: [num_seqs * max_len_extend * max_len_extend] bool or None +// TreeMaskMode::QLEN_ONLY tree mask for speculative TARGET_VERIFY; see [NOTE] 5 above. +void extend_attention_cpu( + at::Tensor& q_extend, + const std::optional& k_extend_opt, + const std::optional& v_extend_opt, + at::Tensor& o_extend, + at::Tensor& k_buffer, + at::Tensor& v_buffer, + at::Tensor& req_to_token, + at::Tensor& req_pool_indices, + at::Tensor& seq_lens, + at::Tensor& extend_seq_lens, + at::Tensor& extend_start_loc, + int64_t max_len_extend, + double sm_scale, + double logit_cap, + bool is_cross_attn, + int64_t sliding_window_size, + std::optional encoder_lens, + std::optional sinks, + std::optional tree_mask) { + if (!is_cross_attn) { + TORCH_CHECK( + k_extend_opt.has_value() && v_extend_opt.has_value(), + "k_extend and v_extend are required for non-cross attention"); + } + // Since k_extend and v_extend are not used for cross attention, they can be initialized as k_buffer and v_buffer + // here. + auto k_extend = k_extend_opt.has_value() ? k_extend_opt.value() : k_buffer; + auto v_extend = v_extend_opt.has_value() ? v_extend_opt.value() : v_buffer; + + CHECK_LAST_DIM_CONTIGUOUS_INPUT(q_extend); + CHECK_INPUT(o_extend); + CHECK_LAST_DIM_CONTIGUOUS_INPUT(k_extend); + CHECK_LAST_DIM_CONTIGUOUS_INPUT(v_extend); + CHECK_LAST_DIM_CONTIGUOUS_INPUT(k_buffer); + CHECK_LAST_DIM_CONTIGUOUS_INPUT(v_buffer); + + int num_seqs = seq_lens.size(0); + int max_num_reqs = req_to_token.size(0); + int max_context_len = req_to_token.size(1); + int max_total_num_tokens = k_buffer.size(0); + + int num_heads = q_extend.size(1); + int num_heads_kv = k_extend.size(1); + int head_size = q_extend.size(2); + int head_size_v = v_extend.size(2); + + // strides for q_extend, k_extend and v_extend + int q_strideM = q_extend.stride(0); + int q_strideH = q_extend.stride(1); + int ke_strideN = k_extend.stride(0); + int ke_strideH = k_extend.stride(1); + int ve_strideN = v_extend.stride(0); + int ve_strideH = v_extend.stride(1); + + // strides for k_buffer and v_buffer + int k_strideN = k_buffer.stride(0); + int k_strideH = k_buffer.stride(1); + int v_strideN = v_buffer.stride(0); + int v_strideH = v_buffer.stride(1); + + // check sizes + CHECK_EQ(req_pool_indices.size(0), num_seqs); + CHECK_EQ(extend_seq_lens.size(0), num_seqs); + CHECK_EQ(extend_start_loc.size(0), num_seqs); + CHECK_EQ(v_extend.size(1), num_heads_kv); + CHECK_EQ(k_buffer.size(1), v_buffer.size(1)); + + // MLA will skip prefix part + const bool is_prefix_skipped = k_buffer.size(1) != num_heads_kv; + + // check index data types + const auto index_dtype = req_to_token.scalar_type(); + TORCH_CHECK( + index_dtype == at::kInt || index_dtype == at::kLong, + "extend: expect req_to_token to be int32 or int64, got ", + index_dtype); + TORCH_CHECK(seq_lens.scalar_type() == at::kLong, "extend: expect req_lens to be int64, got ", seq_lens.scalar_type()); + TORCH_CHECK( + req_pool_indices.scalar_type() == at::kLong, + "extend: expect req_pool_indices to be int64, got ", + req_pool_indices.scalar_type()); + TORCH_CHECK( + extend_seq_lens.scalar_type() == index_dtype && extend_start_loc.scalar_type() == index_dtype, + "extend: expect extend_seq_lens and extend_start_loc to have same dtype as req_to_token."); + + // D and DV need to be 32x as we transpose by 512-bit + TORCH_CHECK(head_size % 32 == 0, "invalid head_size ", head_size); + TORCH_CHECK(head_size_v % 32 == 0, "invalid head_size_v ", head_size_v); + + int num_threads = at::get_num_threads(); + auto buffer = at::empty({}, q_extend.options().dtype(at::kChar)); + + bool has_encoder_lens = encoder_lens.has_value(); + // Since encoder_lens is not used when it is None, encoder_lens_t can be initialized as any tensor of int64_t dtype. + at::Tensor encoder_lens_t = seq_lens; + if (has_encoder_lens) { + encoder_lens_t = encoder_lens.value(); + CHECK_EQ(encoder_lens_t.size(0), num_seqs); + } + bool has_sink = sinks.has_value(); + at::Tensor sinks_tensor = has_sink ? sinks.value() : at::empty({num_heads}, q_extend.options()); + CHECK_DIM(1, sinks_tensor); + CHECK_EQ(sinks_tensor.size(0), num_heads); + + const bool* tree_mask_ptr = nullptr; + if (tree_mask.has_value()) { + const at::Tensor& tree_mask_t = tree_mask.value(); + CHECK_INPUT(tree_mask_t); + TORCH_CHECK( + tree_mask_t.scalar_type() == at::kBool, "extend: expect tree_mask to be bool, got ", tree_mask_t.scalar_type()); + TORCH_CHECK( + tree_mask_t.numel() == static_cast(num_seqs) * max_len_extend * max_len_extend, + "extend: expect tree_mask numel to be num_seqs * max_len_extend^2 = ", + static_cast(num_seqs) * max_len_extend * max_len_extend, + ", got ", + tree_mask_t.numel()); + TORCH_CHECK(!is_cross_attn, "extend: tree_mask is not supported for cross attention"); + // The window mask derives query positions from the row index + // (seq_len_prefix + m + row), but tree-mask rows sit at their tree depth, + // which is <= the row index; combining the two would over-mask the prefix. + TORCH_CHECK(sliding_window_size <= 0, "extend: tree_mask is not supported with sliding window attention"); + tree_mask_ptr = tree_mask_t.data_ptr(); + } + + AT_DISPATCH_REDUCED_FLOATING_TYPES(q_extend.scalar_type(), "extend_attention_kernel", [&] { + AT_DISPATCH_INDEX_TYPES(index_dtype, "extend_attention_indices", [&] { + if (max_len_extend <= 256) { + LAUNCH_EXTEND_ATTENTION_KERNEL(32, 64); + } else if (max_len_extend <= 1024) { + LAUNCH_EXTEND_ATTENTION_KERNEL(128, 256); + } else if (max_len_extend <= 4096) { + LAUNCH_EXTEND_ATTENTION_KERNEL(256, 768); + } else { // max_len_extend > 4096 + LAUNCH_EXTEND_ATTENTION_KERNEL(512, 768); + } + }); + }); +} diff --git a/csrc/cpu/sgl-kernels/flash_attn.h b/csrc/cpu/sgl-kernels/flash_attn.h new file mode 100644 index 0000000000..24ceec64b5 --- /dev/null +++ b/csrc/cpu/sgl-kernels/flash_attn.h @@ -0,0 +1,250 @@ +// Adapted from +// https://github.com/sgl-project/sglang/tree/main/sgl-kernel/csrc/cpu + +// clang-format off + +#pragma once +#include "common.h" +#include "vec.h" +#include "vec_pack.h" + +template +inline void fill_stub(scalar_t* __restrict__ out, float val, int size) { + using Vec = at::vec::Vectorized; + constexpr int kVecSize = Vec::size(); + const Vec data_vec = Vec(static_cast(val)); + int d = 0; +#pragma GCC unroll 4 + for (; d <= size - kVecSize; d += kVecSize) { + data_vec.store(out + d); + } + if (size - d > 0) { + data_vec.store(out + d, size - d); + } +} + +template +inline void copy_stub(scalar_t* __restrict__ out, const float* __restrict__ input) { + static_assert(BLOCK_N % 32 == 0); + using bVec = at::vec::Vectorized; + using fVec = at::vec::Vectorized; + + constexpr int COLS = BLOCK_N / 16; + auto store = [&](auto i) { + constexpr int col = i % COLS; + // for COLS = 2, 4 use 512bit store + if constexpr (col % 2 == 0) { + auto [a_fvec0, a_fvec1] = load_float_vec2(input + col * 16); + bVec out_bvec = convert_from_float_ext(a_fvec0, a_fvec1); + out_bvec.store(out + col * 16); + } + }; + Unroll{}(store); +} + +template +inline void copy_stub(scalar_t* __restrict__ out, const float* __restrict__ acc, float s, int size) { + using bVec = at::vec::Vectorized; + using fVec = at::vec::Vectorized; + constexpr int kVecSize = bVec::size(); + const fVec s_fvec = fVec(s); + int d = 0; +#pragma GCC unroll 4 + for (; d <= size - kVecSize; d += kVecSize) { + auto [a_fvec0, a_fvec1] = load_float_vec2(acc + d); + a_fvec0 = a_fvec0 * s_fvec; + a_fvec1 = a_fvec1 * s_fvec; + bVec out_bvec = convert_from_float_ext(a_fvec0, a_fvec1); + out_bvec.store(out + d); + } + for (; d < size; ++d) { + out[d] = static_cast(acc[d] * s); + } +} + +#if defined(CPU_CAPABILITY_AVX512) +template <> +inline void copy_stub(at::BFloat16* __restrict__ out, const float* __restrict__ acc, float s, int size) { + const __m512 vscale = _mm512_set1_ps(s); + int d = 0; +#pragma GCC unroll 4 + for (; d <= size - 32; d += 32) { + __m512 va0 = _mm512_mul_ps(_mm512_loadu_ps(acc + d), vscale); + __m512 va1 = _mm512_mul_ps(_mm512_loadu_ps(acc + d + 16), vscale); + __m512i vb = (__m512i)(_mm512_cvtne2ps_pbh(va1, va0)); + _mm512_storeu_si512(out + d, vb); + } + int remainder = size - d; + if (remainder > 0) { + if (remainder <= 16) { + const __mmask16 vmask = (1ULL << remainder) - 1; + __m512 va = _mm512_mul_ps(_mm512_maskz_loadu_ps(vmask, acc + d), vscale); + __m256i vb = (__m256i)(_mm512_cvtneps_pbh(va)); + _mm256_mask_storeu_epi16(reinterpret_cast<__m256i*>(out + d), vmask, vb); + } else { // remainder > 16 + const __mmask16 vmask = (1ULL << (remainder - 16)) - 1; + __m512 va0 = _mm512_mul_ps(_mm512_loadu_ps(acc + d), vscale); + __m512 va1 = _mm512_mul_ps(_mm512_maskz_loadu_ps(vmask, acc + d + 16), vscale); + __m512i vb = (__m512i)(_mm512_cvtne2ps_pbh(va1, va0)); + const __mmask32 vmask2 = (1ULL << remainder) - 1; + _mm512_mask_storeu_epi16(reinterpret_cast<__m512i*>(out + d), vmask2, vb); + } + } +} +#endif + +template +struct flash_attn_softmax { + static inline void apply( + float* __restrict__ s_i, + scalar_t* __restrict__ s_delta2, + float* __restrict__ v_prime, + float* __restrict__ s_prime, + float* __restrict__ m_prime, + int m_size, + int n_size, + int padded_n_size, + int head_size_v, + const float sm_scale, + int row) { + using Vec = at::vec::Vectorized; + const Vec scale_vec = Vec(sm_scale); + float* s_delta = s_i; + // s_i <- s_i * scale + at::vec::map([scale_vec](Vec x) { return x * scale_vec; }, s_i + row * BLOCK_N, s_i + row * BLOCK_N, n_size); + + // m_i: max value per row + float m_i = + at::vec::reduce_all([](Vec& x, Vec& y) { return at::vec::maximum(x, y); }, s_i + row * BLOCK_N, n_size); + m_i = std::max(m_i, m_prime[row]); + + // m_delta <- exp(m' - m_i) + float m_delta = std::exp(m_prime[row] - m_i); + + // s_delta <- exp(s_i - m_i) + at::vec::map( + [m_i](Vec x) { return (x - Vec(m_i)).fexp_u20(); }, s_delta + row * BLOCK_N, s_i + row * BLOCK_N, n_size); + + // s' <- s' * m_delta + sum(s_delta) + s_prime[row] *= m_delta; + s_prime[row] += at::vec::reduce_all([](Vec& x, Vec& y) { return x + y; }, s_delta + row * BLOCK_N, n_size); + + m_prime[row] = m_i; + + // v' <- v' * m_delta + at::vec::map( + [m_delta](Vec x) { return x * Vec(m_delta); }, + v_prime + row * head_size_v, + v_prime + row * head_size_v, + head_size_v); + + // Keep s_delta row-major for the following brgemm(P @ V), and only + // convert the columns that brgemm will consume. + fill_stub(s_delta + row * BLOCK_N + n_size, 0.f, padded_n_size - n_size); + copy_stub(s_delta2 + row * BLOCK_N, s_delta + row * BLOCK_N, 1.f, padded_n_size); + } +}; + +#if defined(CPU_CAPABILITY_AVX512) +template +struct flash_attn_softmax { + static inline void apply( + float* __restrict__ s_i, + at::BFloat16* __restrict__ s_delta2, + float* __restrict__ v_prime, + float* __restrict__ s_prime, + float* __restrict__ m_prime, + int m_size, + int n_size, + int padded_n_size, + int head_size_v, + const float sm_scale, + int row) { + float* s_delta = s_i; + const __m512 vscale = _mm512_set1_ps(sm_scale); + + int n_remainder = n_size & 15; // 0xF + const __mmask16 vmask = (1ULL << n_remainder) - 1; + + int v_remainder = head_size_v & 15; // 0xF + const __mmask16 vmask1 = (1ULL << v_remainder) - 1; + + constexpr float NEG_INF = -std::numeric_limits::infinity(); + + __m512 va; + __m256i vb; + __m512 vmax; + __m512 vsum; + __m512 vmdelta; + + const __m512 vneg_inf = _mm512_set1_ps(NEG_INF); + + int m = row; + vmax = vneg_inf; + + // s_i <- s_i * scale + int n = 0; + for (; n <= n_size - 16; n += 16) { + va = _mm512_mul_ps(_mm512_loadu_ps(s_i + m * BLOCK_N + n), vscale); + vmax = _mm512_max_ps(va, vmax); + } + if (n_remainder > 0) { + va = _mm512_mul_ps(_mm512_mask_loadu_ps(vneg_inf, vmask, s_i + m * BLOCK_N + n), vscale); + vmax = _mm512_max_ps(va, vmax); + } + + // m_i: max value per row + float m_i = _mm512_reduce_max_ps(vmax); + m_i = std::max(m_i, m_prime[m]); + vmax = _mm512_set1_ps(m_i); + + // m_delta <- exp(m' - m_i) + float m_delta = std::exp(m_prime[m] - m_i); + + // s_delta <- exp(s_i - m_i) + vsum = _mm512_setzero_ps(); + for (n = 0; n <= n_size - 16; n += 16) { + va = _mm512_mul_ps(_mm512_loadu_ps(s_i + m * BLOCK_N + n), vscale); + va = _mm512_fexp_u20_ps(_mm512_sub_ps(va, vmax)); + vsum = _mm512_add_ps(vsum, va); + + vb = (__m256i)(_mm512_cvtneps_pbh(va)); + _mm256_storeu_si256(reinterpret_cast<__m256i*>(s_delta2 + m * BLOCK_N + n), vb); + } + if (n_remainder > 0) { + va = _mm512_mul_ps(_mm512_mask_loadu_ps(vneg_inf, vmask, s_i + m * BLOCK_N + n), vscale); + va = _mm512_fexp_u20_ps(_mm512_sub_ps(va, vmax)); + vsum = _mm512_add_ps(vsum, va); + + vb = (__m256i)(_mm512_cvtneps_pbh(va)); + _mm256_mask_storeu_epi16(reinterpret_cast<__m256i*>(s_delta2 + m * BLOCK_N + n), vmask, vb); + } + + // s' <- s' * m_delta + sum(s_delta) + s_prime[m] *= m_delta; + s_prime[m] += _mm512_reduce_add_ps(vsum); + + m_prime[m] = m_i; + + // pad s_delta with 0, pad_size range from [0, 32) + int pad_size = padded_n_size - n_size; + if (pad_size > 0) { + const __m512i vzero = _mm512_setzero_si512(); + __mmask32 vmask2 = (1ULL << pad_size) - 1; + _mm512_mask_storeu_epi16(reinterpret_cast<__m512i*>(s_delta2 + m * BLOCK_N + n_size), vmask2, vzero); + } + + // v' <- v' * m_delta + vmdelta = _mm512_set1_ps(m_delta); + int k = 0; + for (; k <= head_size_v - 16; k += 16) { + va = _mm512_mul_ps(_mm512_loadu_ps(v_prime + m * head_size_v + k), vmdelta); + _mm512_storeu_ps(reinterpret_cast<__m512*>(v_prime + m * head_size_v + k), va); + } + if (v_remainder > 0) { + va = _mm512_mul_ps(_mm512_maskz_loadu_ps(vmask1, v_prime + m * head_size_v + k), vmdelta); + _mm512_mask_storeu_ps(reinterpret_cast<__m512*>(v_prime + m * head_size_v + k), vmask1, va); + } + } +}; +#endif diff --git a/csrc/cpu/sgl-kernels/mla_cache.cpp b/csrc/cpu/sgl-kernels/mla_cache.cpp new file mode 100644 index 0000000000..e4629f84a2 --- /dev/null +++ b/csrc/cpu/sgl-kernels/mla_cache.cpp @@ -0,0 +1,111 @@ +// vLLM-native CPU cache-write op for MLA's single-latent-buffer KV cache. +// +// `concat_and_cache_mla` (the generic MLA cache-write op used by every GPU +// backend) is registered CUDA-only. This is the CPU counterpart, adapted in +// spirit from SGLang's `store_cache_cpu` (csrc/cpu/kvcache.cpp) but +// generalized to write two source tensors (`kv_c_normed`, `k_pe`) into two +// different column-offset ranges of the SAME destination row -- SGLang's +// version assumes k/v land in two independent, equal-row-width cache +// tensors, which doesn't hold here since MLA's cache is one 576-wide buffer +// and the two column ranges (512-wide, 64-wide) don't match the buffer's +// true per-token stride, so the write can't reuse `store_cache_cpu` as-is. + +#include "common.h" +#include "vec.h" + +namespace { + +template +inline void copy_stub(scalar_t* __restrict__ dst, + const scalar_t* __restrict__ src, int64_t size) { + int64_t d = 0; +#if defined(CPU_CAPABILITY_AVX512) + using Vec = at::vec::Vectorized; + constexpr int64_t kVecSize = Vec::size(); + for (; d <= size - kVecSize; d += kVecSize) { + Vec data = Vec::loadu(src + d); + data.store(dst + d); + } +#endif + for (; d < size; ++d) { + dst[d] = src[d]; + } +} + +template +void concat_and_cache_mla_kernel_impl( + const scalar_t* __restrict__ kv_c_normed, // [num_tokens, kv_lora_rank] + const scalar_t* __restrict__ k_pe, // [num_tokens, qk_rope_head_dim] + scalar_t* __restrict__ kv_cache, // [.., kv_lora_rank + qk_rope_head_dim] + const index_t* __restrict__ slot_mapping, // [num_tokens] + int64_t num_tokens, int64_t kv_lora_rank, int64_t qk_rope_head_dim, + int64_t kv_c_stride, int64_t k_pe_stride, int64_t cache_stride) { + at::parallel_for(0, num_tokens, 0, [&](int64_t begin, int64_t end) { + for (int64_t i = begin; i < end; ++i) { + int64_t slot = static_cast(slot_mapping[i]); + if (slot < 0) { + // padded/invalid token, matches `concat_and_cache_mla`'s semantics. + continue; + } + scalar_t* __restrict__ cache_row = kv_cache + slot * cache_stride; + copy_stub(cache_row, kv_c_normed + i * kv_c_stride, kv_lora_rank); + copy_stub(cache_row + kv_lora_rank, k_pe + i * k_pe_stride, + qk_rope_head_dim); + } + }); +} + +} // namespace + +// kv_c_normed : [num_tokens, kv_lora_rank] +// k_pe : [num_tokens, qk_rope_head_dim] or [num_tokens, 1, +// qk_rope_head_dim] kv_cache : [num_blocks, block_size, kv_lora_rank + +// qk_rope_head_dim] slot_mapping: [num_tokens] int32/int64, absolute physical +// row index +// (block_id * block_size + block_offset); negative entries are +// skipped (padded tokens), matching `concat_and_cache_mla`. +void concat_and_cache_mla_cpu(const at::Tensor& kv_c_normed, + const at::Tensor& k_pe, at::Tensor& kv_cache, + const at::Tensor& slot_mapping) { + TORCH_CHECK(kv_c_normed.dim() == 2, + "kv_c_normed must be 2D [num_tokens, kv_lora_rank]"); + TORCH_CHECK(k_pe.dim() == 2 || k_pe.dim() == 3, "k_pe must be 2D or 3D"); + TORCH_CHECK(kv_cache.dim() == 3, + "kv_cache must be 3D [num_blocks, block_size, head_size]"); + TORCH_CHECK(kv_c_normed.stride(-1) == 1, + "kv_c_normed innermost dim must be contiguous"); + TORCH_CHECK(k_pe.stride(-1) == 1, "k_pe innermost dim must be contiguous"); + TORCH_CHECK(kv_cache.stride(-1) == 1, + "kv_cache innermost dim must be contiguous"); + + int64_t num_tokens = kv_c_normed.size(0); + int64_t kv_lora_rank = kv_c_normed.size(1); + int64_t qk_rope_head_dim = k_pe.size(-1); + int64_t head_size = kv_cache.size(-1); + TORCH_CHECK(head_size == kv_lora_rank + qk_rope_head_dim, + "kv_cache head_size must equal kv_lora_rank + qk_rope_head_dim"); + TORCH_CHECK(slot_mapping.size(0) == num_tokens, "slot_mapping size mismatch"); + + // Real physical per-token stride of the cache, read from the tensor's own + // strides rather than assumed to equal head_size (paged/pooled buffers are + // not guaranteed to have zero inter-row padding in general, even though in + // practice a freshly-allocated MLA cache is fully contiguous). + int64_t cache_stride = kv_cache.stride(1); + + const auto dtype = kv_cache.scalar_type(); + TORCH_CHECK(dtype == kv_c_normed.scalar_type() && dtype == k_pe.scalar_type(), + "concat_and_cache_mla_cpu: dtype mismatch"); + const auto index_dtype = slot_mapping.scalar_type(); + TORCH_CHECK(index_dtype == at::kLong || index_dtype == at::kInt, + "slot_mapping must be int32 or int64"); + + AT_DISPATCH_REDUCED_FLOATING_TYPES(dtype, "concat_and_cache_mla_cpu", [&] { + AT_DISPATCH_INDEX_TYPES(index_dtype, "concat_and_cache_mla_cpu_index", [&] { + concat_and_cache_mla_kernel_impl( + kv_c_normed.data_ptr(), k_pe.data_ptr(), + kv_cache.data_ptr(), slot_mapping.data_ptr(), + num_tokens, kv_lora_rank, qk_rope_head_dim, kv_c_normed.stride(0), + k_pe.stride(0), cache_stride); + }); + }); +} diff --git a/csrc/cpu/torch_bindings.cpp b/csrc/cpu/torch_bindings.cpp index 4bb0edf2b2..c393000617 100644 --- a/csrc/cpu/torch_bindings.cpp +++ b/csrc/cpu/torch_bindings.cpp @@ -91,6 +91,39 @@ at::Tensor fp8_scaled_mm_cpu(at::Tensor& mat1, at::Tensor& mat2, const std::optional& bias, at::ScalarType out_dtype, bool is_vnni); +// Adapted from sglang: MLA CPU kernels (AMX-only) +void decode_attention_cpu(at::Tensor& query, at::Tensor& k_buffer, + at::Tensor& v_buffer, at::Tensor& output, + const std::optional& key, + const std::optional& value, + const std::optional& loc, + at::Tensor& attn_logits, at::Tensor& req_to_token, + at::Tensor& req_pool_indices, at::Tensor& seq_lens, + double sm_scale, double logit_cap, bool is_cross_attn, + int64_t sliding_window_size, + std::optional encoder_lens, + std::optional sinks); + +void extend_attention_cpu( + at::Tensor& q_extend, const std::optional& k_extend, + const std::optional& v_extend, at::Tensor& o_extend, + at::Tensor& k_buffer, at::Tensor& v_buffer, at::Tensor& req_to_token, + at::Tensor& req_pool_indices, at::Tensor& seq_lens, + at::Tensor& extend_seq_lens, at::Tensor& extend_start_loc, + int64_t max_len_extend, double sm_scale, double logit_cap, + bool is_cross_attn, int64_t sliding_window_size, + std::optional encoder_lens, std::optional sinks, + std::optional tree_mask); + +void bmm_cpu(at::Tensor& out, at::Tensor& mat1, at::Tensor& mat2, bool is_vnni, + const std::optional& scale); + +// vLLM-native: CPU cache-write op for MLA's single-latent-buffer KV cache +// (the CUDA-only `concat_and_cache_mla` has no CPU dispatch). +void concat_and_cache_mla_cpu(const at::Tensor& kv_c_normed, + const at::Tensor& k_pe, at::Tensor& kv_cache, + const at::Tensor& slot_mapping); + // Adapted from sglang: INT4 W4A8 kernels std::tuple convert_weight_packed_scale_zp( at::Tensor qweight, // awq: (*, K, N / 8) || gptq: (*, K / 8, N) , int32 @@ -521,6 +554,35 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) { "pad_slot_id, " "bool is_vnni) -> Tensor"); ops.impl("causal_conv1d_update_cpu", torch::kCPU, &causal_conv1d_update_cpu); + + // Adapted from sglang: MLA CPU kernels (AMX-only, DeepSeek V2/V3/R1) + ops.def( + "decode_attention_cpu(Tensor query, Tensor k_buffer, Tensor v_buffer, " + "Tensor(a!) output, Tensor? key, Tensor? value, Tensor? loc, Tensor " + "attn_logits, Tensor req_to_token, Tensor req_pool_indices, Tensor " + "seq_lens, float sm_scale, float logit_cap, bool is_cross_attn, int " + "sliding_window_size, Tensor? encoder_lens, Tensor? sinks) -> ()"); + ops.impl("decode_attention_cpu", torch::kCPU, &decode_attention_cpu); + + ops.def( + "extend_attention_cpu(Tensor q_extend, Tensor? k_extend, Tensor? " + "v_extend, Tensor(a!) o_extend, Tensor k_buffer, Tensor v_buffer, " + "Tensor req_to_token, Tensor req_pool_indices, Tensor seq_lens, Tensor " + "extend_seq_lens, Tensor extend_start_loc, int max_len_extend, float " + "sm_scale, float logit_cap, bool is_cross_attn, int " + "sliding_window_size, Tensor? encoder_lens, Tensor? sinks, Tensor? " + "tree_mask=None) -> ()"); + ops.impl("extend_attention_cpu", torch::kCPU, &extend_attention_cpu); + + ops.def( + "bmm_cpu(Tensor(a!) out, Tensor mat1, Tensor mat2, bool is_vnni, " + "Tensor? scale) -> ()"); + ops.impl("bmm_cpu", torch::kCPU, &bmm_cpu); + + ops.def( + "concat_and_cache_mla_cpu(Tensor kv_c_normed, Tensor k_pe, " + "Tensor(a!) kv_cache, Tensor slot_mapping) -> ()"); + ops.impl("concat_and_cache_mla_cpu", torch::kCPU, &concat_and_cache_mla_cpu); #endif #if (defined(__AVX512BF16__) && defined(__AVX512F__) && \ diff --git a/tests/kernels/attention/test_amx_mla.py b/tests/kernels/attention/test_amx_mla.py new file mode 100644 index 0000000000..f5fa2fad16 --- /dev/null +++ b/tests/kernels/attention/test_amx_mla.py @@ -0,0 +1,471 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Correctness tests for the AMX-only CPU MLA backend: the vendored +decode/extend/bmm kernels, the KV cache write, and the ``AMXMLAImpl`` +backend built on top of them. +""" + +import pytest +import torch + +from vllm.platforms import current_platform + +if not current_platform.is_cpu(): + pytest.skip("skipping CPU-only tests", allow_module_level=True) + +if not torch.cpu._is_amx_tile_supported(): + pytest.skip("AMX MLA requires an AMX-capable host", allow_module_level=True) + +torch.cpu._init_amx() + +from vllm import _custom_ops as ops # noqa: E402 +from vllm.utils.torch_utils import set_random_seed # noqa: E402 +from vllm.v1.attention.backends.mla.amx_mla import ( # noqa: E402 + AMXMLAImpl, + _compute_num_kv_splits, + _expand_block_table, +) + +KV_LORA_RANK = 64 +QK_ROPE_HEAD_DIM = 32 +HEAD_SIZE = KV_LORA_RANK + QK_ROPE_HEAD_DIM # 96, kv-cache row width +# The extend/decode kernels transpose by 512-bit lanes: both the cache row +# width and the (kv_lora_rank-wide) value width must be 32-element aligned. +DTYPE = torch.bfloat16 +ATOL = 2e-2 +RTOL = 2e-2 +# Backend-level tests chain absorb -> attend -> de-absorb through bf16, +# accumulating more rounding error than a single kernel op. +_IMPL_ATOL = 1.5e-1 +_IMPL_RTOL = 1.5e-1 + + +def _flatten_cache(kv_cache: torch.Tensor) -> torch.Tensor: + """(num_blocks, block_size, head_size) -> (num_slots, head_size).""" + return kv_cache.view(-1, kv_cache.size(-1)) + + +def _random_paged_cache(num_slots: int, head_size: int) -> torch.Tensor: + return torch.randn(num_slots, head_size, dtype=DTYPE) + + +def _ref_latent_attn( + q: torch.Tensor, + keys: torch.Tensor, + scale: float, + kv_lora_rank: int, +) -> torch.Tensor: + """Plain-PyTorch causal-free latent-space MQA attention for one request. + + q: (num_heads, head_size); keys: (ctx_len, head_size). Values are the + first ``kv_lora_rank`` columns of the same latent rows (MLA's K/V + aliasing), matching what decode_attention_cpu/extend_attention_cpu + compute internally for the MLA-shaped case. + """ + values = keys[:, :kv_lora_rank] + scores = (q.float() @ keys.float().T) * scale + probs = torch.softmax(scores, dim=-1) + return (probs @ values.float()).to(q.dtype) + + +def test_bmm_cpu_matches_torch_bmm(): + set_random_seed(0) + # bmm_cpu requires the output's last dim (mat2's out-features) to be a + # multiple of 32 (tinygemm's tile width); the contraction dim is free. + n, b, p, l = 8, 32, 32, 32 # noqa: E741 + mat1 = torch.randn(n, b, p, dtype=DTYPE) + mat2 = torch.randn(n, l, p, dtype=DTYPE) # (N, OUT, IN) Linear convention + + ref = torch.bmm(mat1.float(), mat2.float().transpose(1, 2)).to(DTYPE) + + out = torch.empty(n, b, l, dtype=DTYPE) + ops.bmm_cpu(out, mat1, mat2, False, None) + torch.testing.assert_close(out, ref, atol=ATOL, rtol=RTOL) + + packed = torch.ops._C.convert_weight_packed(mat2.contiguous()) + out_vnni = torch.empty(n, b, l, dtype=DTYPE) + ops.bmm_cpu(out_vnni, mat1, packed, True, None) + torch.testing.assert_close(out_vnni, ref, atol=ATOL, rtol=RTOL) + + +def test_concat_and_cache_mla_round_trip(): + set_random_seed(0) + num_tokens = 37 + num_blocks, block_size = 8, 16 + kv_cache = torch.zeros(num_blocks, block_size, HEAD_SIZE, dtype=DTYPE) + + kv_c_normed = torch.randn(num_tokens, KV_LORA_RANK, dtype=DTYPE) + k_pe = torch.randn(num_tokens, QK_ROPE_HEAD_DIM, dtype=DTYPE) + slot_mapping = torch.randperm(num_blocks * block_size)[:num_tokens].to(torch.int64) + + ops.amx_mla_concat_and_cache(kv_c_normed, k_pe, kv_cache, slot_mapping) + + flat = _flatten_cache(kv_cache) + read_back = flat[slot_mapping] + expected = torch.cat([kv_c_normed, k_pe], dim=-1) + torch.testing.assert_close(read_back, expected) + + untouched = torch.ones(num_blocks * block_size, dtype=torch.bool) + untouched[slot_mapping] = False + assert (flat[untouched] == 0).all() + + +@pytest.mark.parametrize("seq_lens", [[5], [1, 300, 33], [513, 1, 129, 7]]) +def test_decode_attention_cpu_matches_reference(seq_lens): + set_random_seed(0) + num_seqs = len(seq_lens) + num_heads = 4 + block_size = 32 + max_len = max(seq_lens) + max_blocks = (max_len + block_size - 1) // block_size + num_blocks = num_seqs * max_blocks + + kv_cache = _random_paged_cache(num_blocks * block_size, HEAD_SIZE) + block_table = torch.randperm(num_blocks).view(num_seqs, max_blocks).to(torch.int64) + req_to_token = _expand_block_table(block_table, block_size) + req_pool_indices = torch.arange(num_seqs, dtype=torch.int64) + seq_lens_t = torch.tensor(seq_lens, dtype=torch.int64) + scale = HEAD_SIZE**-0.5 + + q = torch.randn(num_seqs, num_heads, HEAD_SIZE, dtype=DTYPE) + kv_cache_flat = kv_cache.view(-1, 1, HEAD_SIZE) + v_buffer = kv_cache_flat[..., :KV_LORA_RANK] + + o = torch.zeros(num_seqs, num_heads, KV_LORA_RANK, dtype=DTYPE) + attn_logits = torch.zeros( + num_seqs, num_heads, 4, KV_LORA_RANK + 1, dtype=torch.float32 + ) + ops.cpu_mla_decode( + q, + kv_cache_flat, + v_buffer, + o, + None, + None, + None, + attn_logits, + req_to_token, + req_pool_indices, + seq_lens_t, + scale, + 0.0, + False, + 0, + None, + None, + ) + + flat_cache = _flatten_cache(kv_cache) + for i, seq_len in enumerate(seq_lens): + token_ids = req_to_token[i, :seq_len] + keys = flat_cache[token_ids] + ref = _ref_latent_attn(q[i], keys, scale, KV_LORA_RANK) + torch.testing.assert_close(o[i], ref, atol=ATOL, rtol=RTOL) + + +@pytest.mark.parametrize( + "prefix_lens,extend_lens", + [ + ([0], [9]), # fresh prefill, no cached prefix + ([64], [17]), # continuation of a cached prefix + ([0, 40, 128], [23, 5, 31]), # mixed batch + ], +) +def test_extend_attention_cpu_matches_reference(prefix_lens, extend_lens): + set_random_seed(0) + num_seqs = len(prefix_lens) + num_heads = 4 + block_size = 32 + seq_lens = [p + e for p, e in zip(prefix_lens, extend_lens)] + max_len = max(seq_lens) + max_blocks = (max_len + block_size - 1) // block_size + num_blocks = num_seqs * max_blocks + + kv_cache = _random_paged_cache(num_blocks * block_size, HEAD_SIZE) + block_table = torch.randperm(num_blocks).view(num_seqs, max_blocks).to(torch.int64) + req_to_token = _expand_block_table(block_table, block_size) + req_pool_indices = torch.arange(num_seqs, dtype=torch.int64) + seq_lens_t = torch.tensor(seq_lens, dtype=torch.int64) + extend_seq_lens_t = torch.tensor(extend_lens, dtype=torch.int64) + extend_start_loc = torch.cumsum( + torch.tensor([0, *extend_lens[:-1]], dtype=torch.int64), dim=0 + ) + scale = HEAD_SIZE**-0.5 + total_new_tokens = sum(extend_lens) + + # The new tokens' latent K/V, already written into the cache (mirrors + # do_kv_cache_update running before forward_mha in real usage). + k_extend_flat = torch.randn(total_new_tokens, HEAD_SIZE, dtype=DTYPE) + flat_cache = _flatten_cache(kv_cache) + for i in range(num_seqs): + prefix_len = prefix_lens[i] + ext_len = extend_lens[i] + new_token_ids = req_to_token[i, prefix_len : prefix_len + ext_len] + start = extend_start_loc[i].item() + flat_cache[new_token_ids] = k_extend_flat[start : start + ext_len] + + q_extend = torch.randn(total_new_tokens, num_heads, HEAD_SIZE, dtype=DTYPE) + k_extend = k_extend_flat.unsqueeze(1) + v_extend = k_extend[..., :KV_LORA_RANK] + kv_cache_flat = kv_cache.view(-1, 1, HEAD_SIZE) + v_buffer = kv_cache_flat[..., :KV_LORA_RANK] + + o_extend = torch.empty(total_new_tokens, num_heads, KV_LORA_RANK, dtype=DTYPE) + ops.cpu_mla_extend( + q_extend, + k_extend, + v_extend, + o_extend, + kv_cache_flat, + v_buffer, + req_to_token, + req_pool_indices, + seq_lens_t, + extend_seq_lens_t, + extend_start_loc, + max(extend_lens), + scale, + 0.0, + False, + 0, + None, + None, + None, + ) + + for i in range(num_seqs): + prefix_len = prefix_lens[i] + ext_len = extend_lens[i] + start = extend_start_loc[i].item() + for t in range(ext_len): + visible = prefix_len + t + 1 + token_ids = req_to_token[i, :visible] + keys = flat_cache[token_ids] + ref = _ref_latent_attn(q_extend[start + t], keys, scale, KV_LORA_RANK) + torch.testing.assert_close(o_extend[start + t], ref, atol=ATOL, rtol=RTOL) + + +class _FakeLinear: + """Minimal duck-typed stand-in for kv_b_proj: only .weight and + .quant_method are read by get_and_maybe_dequant_weights().""" + + def __init__(self, weight: torch.Tensor): + self.weight = weight + self.quant_method = None + + +def _make_amx_mla_impl(num_heads, qk_nope_head_dim, v_head_dim, kv_lora_rank): + weight = torch.randn( + num_heads * (qk_nope_head_dim + v_head_dim), kv_lora_rank, dtype=DTYPE + ) + kv_b_proj = _FakeLinear(weight) + impl = AMXMLAImpl( + num_heads=num_heads, + head_size=kv_lora_rank + QK_ROPE_HEAD_DIM, + scale=(qk_nope_head_dim + QK_ROPE_HEAD_DIM) ** -0.5, + num_kv_heads=1, + alibi_slopes=None, + sliding_window=None, + kv_cache_dtype="auto", + logits_soft_cap=None, + attn_type="decoder", + kv_sharing_target_layer_name=None, + q_lora_rank=None, + kv_lora_rank=kv_lora_rank, + qk_nope_head_dim=qk_nope_head_dim, + qk_rope_head_dim=QK_ROPE_HEAD_DIM, + qk_head_dim=qk_nope_head_dim + QK_ROPE_HEAD_DIM, + v_head_dim=v_head_dim, + kv_b_proj=kv_b_proj, + ) + impl.process_weights_after_loading(DTYPE) + return impl, weight + + +class _FakePrefillMetadata: + def __init__( + self, block_table, cpu_seq_lens, query_start_loc, req_to_token, req_pool_indices + ): + self.block_table = block_table + self.cpu_seq_lens = cpu_seq_lens + query_start_loc_i64 = query_start_loc.to(torch.int64) + extend_seq_lens = query_start_loc_i64[1:] - query_start_loc_i64[:-1] + self.extend_seq_lens = extend_seq_lens + self.extend_start_loc = query_start_loc_i64[:-1] + self.max_len_extend = int(extend_seq_lens.max().item()) + self.req_to_token = req_to_token + self.req_pool_indices = req_pool_indices + + +class _FakeDecodeMetadata: + def __init__(self, block_table, seq_lens, req_to_token, req_pool_indices): + self.block_table = block_table + self.seq_lens_i64 = seq_lens.to(torch.int64) + self.req_to_token = req_to_token + self.req_pool_indices = req_pool_indices + self.num_kv_splits = _compute_num_kv_splits( + int(seq_lens.max().item()), current_platform.num_compute_units() + ) + + +class _FakeAttnMetadata: + def __init__(self, decode=None, prefill=None, max_seq_len=0): + self.decode = decode + self.prefill = prefill + self.max_seq_len = max_seq_len + + +def test_amx_mla_impl_forward_mqa_matches_reference(default_vllm_config): + """forward_mqa receives already-absorbed Q (as MLAAttention.forward_impl + would produce via layer.W_UK_T) and must reproduce plain per-head MLA + decode attention through W_UK/W_UV.""" + set_random_seed(1) + num_heads, qk_nope_head_dim, v_head_dim, kv_lora_rank = 4, 32, 32, 64 + impl, kv_b_weight = _make_amx_mla_impl( + num_heads, qk_nope_head_dim, v_head_dim, kv_lora_rank + ) + w_uk, w_uv = kv_b_weight.T.view( + kv_lora_rank, num_heads, qk_nope_head_dim + v_head_dim + ).split([qk_nope_head_dim, v_head_dim], dim=-1) + + seq_lens = [1, 40, 5] + num_seqs = len(seq_lens) + block_size = 32 + max_blocks = (max(seq_lens) + block_size - 1) // block_size + num_blocks = num_seqs * max_blocks + head_size = kv_lora_rank + QK_ROPE_HEAD_DIM + + kv_cache = torch.randn(num_blocks, block_size, head_size, dtype=DTYPE) + block_table = torch.randperm(num_blocks).view(num_seqs, max_blocks).to(torch.int64) + req_to_token = _expand_block_table(block_table, block_size) + + q_nope = torch.randn(num_seqs, num_heads, qk_nope_head_dim, dtype=DTYPE) + q_pe = torch.randn(num_seqs, num_heads, QK_ROPE_HEAD_DIM, dtype=DTYPE) + # Pre-absorb, mirroring MLAAttention.forward_impl's own bmm against W_UK_T. + ql_nope = torch.einsum("bnp,lnp->bnl", q_nope.float(), w_uk.float()).to(DTYPE) + + attn_metadata = _FakeAttnMetadata( + decode=_FakeDecodeMetadata( + block_table=block_table, + seq_lens=torch.tensor(seq_lens, dtype=torch.int64), + req_to_token=req_to_token, + req_pool_indices=torch.arange(num_seqs, dtype=torch.int64), + ), + max_seq_len=max(seq_lens), + ) + o, lse = impl.forward_mqa((ql_nope, q_pe), kv_cache, attn_metadata, layer=None) + assert lse is None + # De-absorb, mirroring MLAAttention.forward_impl's own bmm against + # layer._v_up_proj, to get back into real per-head V space for + # comparison against a from-scratch reference. + out_real = torch.einsum("bnl,lnv->bnv", o.float(), w_uv.float()) + + flat_cache = _flatten_cache(kv_cache) + scale = impl.scale + for i, seq_len in enumerate(seq_lens): + token_ids = req_to_token[i, :seq_len] + keys_latent = flat_cache[token_ids] # (L, head_size) + k_nope_latent, k_pe = keys_latent.split( + [kv_lora_rank, QK_ROPE_HEAD_DIM], dim=-1 + ) + k_real = torch.einsum("lp,pnd->nld", k_nope_latent.float(), w_uk.float()) + k_pe_b = k_pe.float().unsqueeze(0).expand(num_heads, -1, -1) + k_full = torch.cat([k_real, k_pe_b], dim=-1) # (num_heads, L, qk_head_dim) + q_full = torch.cat( + [q_nope[i], q_pe[i]], dim=-1 + ).float() # (num_heads, qk_head_dim) + scores = torch.einsum("nd,nld->nl", q_full, k_full) * scale + probs = torch.softmax(scores, dim=-1) + v_real = torch.einsum("lp,pnv->nlv", k_nope_latent.float(), w_uv.float()) + ref = torch.einsum("nl,nlv->nv", probs, v_real) + torch.testing.assert_close(out_real[i], ref, atol=_IMPL_ATOL, rtol=_IMPL_RTOL) + + +def test_amx_mla_impl_forward_mha_matches_reference(default_vllm_config): + """forward_mha receives raw unabsorbed Q and must attend correctly for a + prefill batch with a mix of fresh and continued (cached-prefix) + sequences.""" + set_random_seed(2) + num_heads, qk_nope_head_dim, v_head_dim, kv_lora_rank = 4, 32, 32, 64 + impl, kv_b_weight = _make_amx_mla_impl( + num_heads, qk_nope_head_dim, v_head_dim, kv_lora_rank + ) + w_uk, w_uv = kv_b_weight.T.view( + kv_lora_rank, num_heads, qk_nope_head_dim + v_head_dim + ).split([qk_nope_head_dim, v_head_dim], dim=-1) + + prefix_lens = [0, 20] + extend_lens = [13, 7] + seq_lens = [p + e for p, e in zip(prefix_lens, extend_lens)] + num_seqs = len(prefix_lens) + block_size = 32 + max_blocks = (max(seq_lens) + block_size - 1) // block_size + num_blocks = num_seqs * max_blocks + head_size = kv_lora_rank + QK_ROPE_HEAD_DIM + + kv_cache = torch.randn(num_blocks, block_size, head_size, dtype=DTYPE) + block_table = torch.randperm(num_blocks).view(num_seqs, max_blocks).to(torch.int64) + req_to_token = _expand_block_table(block_table, block_size) + flat_cache = _flatten_cache(kv_cache) + + total_new_tokens = sum(extend_lens) + query_start_loc = torch.tensor( + [0, *torch.cumsum(torch.tensor(extend_lens), dim=0).tolist()], + dtype=torch.int64, + ) + + q_nope = torch.randn(total_new_tokens, num_heads, qk_nope_head_dim, dtype=DTYPE) + q_pe = torch.randn(total_new_tokens, num_heads, QK_ROPE_HEAD_DIM, dtype=DTYPE) + q = torch.cat([q_nope, q_pe], dim=-1) + + kv_c_normed = torch.randn(total_new_tokens, kv_lora_rank, dtype=DTYPE) + k_pe = torch.randn(total_new_tokens, 1, QK_ROPE_HEAD_DIM, dtype=DTYPE) + + # Write the new tokens' latent K/V into the cache ahead of time, mirroring + # do_kv_cache_update running before forward_mha in real usage. + for i in range(num_seqs): + prefix_len = prefix_lens[i] + ext_len = extend_lens[i] + new_token_ids = req_to_token[i, prefix_len : prefix_len + ext_len] + start = query_start_loc[i].item() + flat_cache[new_token_ids, :kv_lora_rank] = kv_c_normed[start : start + ext_len] + flat_cache[new_token_ids, kv_lora_rank:] = k_pe[start : start + ext_len, 0] + + attn_metadata = _FakeAttnMetadata( + prefill=_FakePrefillMetadata( + block_table=block_table, + cpu_seq_lens=torch.tensor(seq_lens, dtype=torch.int64), + query_start_loc=query_start_loc, + req_to_token=req_to_token, + req_pool_indices=torch.arange(num_seqs, dtype=torch.int64), + ) + ) + output = torch.empty(total_new_tokens, num_heads * v_head_dim, dtype=DTYPE) + impl.forward_mha( + q, kv_c_normed, k_pe, kv_cache, attn_metadata, k_scale=None, output=output + ) + out_view = output.view(total_new_tokens, num_heads, v_head_dim).float() + + scale = impl.scale + for i in range(num_seqs): + prefix_len = prefix_lens[i] + ext_len = extend_lens[i] + start = query_start_loc[i].item() + for t in range(ext_len): + visible = prefix_len + t + 1 + token_ids = req_to_token[i, :visible] + keys_latent = flat_cache[token_ids] + k_nope_latent, k_pe_ref = keys_latent.split( + [kv_lora_rank, QK_ROPE_HEAD_DIM], dim=-1 + ) + k_real = torch.einsum("lp,pnd->nld", k_nope_latent.float(), w_uk.float()) + k_pe_b = k_pe_ref.float().unsqueeze(0).expand(num_heads, -1, -1) + k_full = torch.cat([k_real, k_pe_b], dim=-1) + q_full = torch.cat([q_nope[start + t], q_pe[start + t]], dim=-1).float() + scores = torch.einsum("nd,nld->nl", q_full, k_full) * scale + probs = torch.softmax(scores, dim=-1) + v_real = torch.einsum("lp,pnv->nlv", k_nope_latent.float(), w_uv.float()) + ref = torch.einsum("nl,nlv->nv", probs, v_real) + torch.testing.assert_close( + out_view[start + t], ref, atol=_IMPL_ATOL, rtol=_IMPL_RTOL + ) diff --git a/tests/v1/attention/test_mla_backends.py b/tests/v1/attention/test_mla_backends.py index cd164250a2..950de76410 100644 --- a/tests/v1/attention/test_mla_backends.py +++ b/tests/v1/attention/test_mla_backends.py @@ -172,9 +172,11 @@ def test_mla_post_load_preserves_runtime_weight_addresses(monkeypatch): layer.kv_b_proj.quant_method = None layer.is_aiter_triton_fp4_bmm_enabled = False layer.is_aiter_triton_fp8_bmm_enabled = False + layer.is_amx_bmm_enabled = False layer.dcp_q_replicate = False layer.quant_config = None layer.layer_name = "test" + layer.impl = SimpleNamespace(process_weights_after_loading=lambda act_dtype: None) monkeypatch.setattr( mla_attention_module, "set_default_quant_scales", lambda *_, **__: None diff --git a/tests/v1/attention/test_mla_prefill_selector.py b/tests/v1/attention/test_mla_prefill_selector.py index e6d9f939ea..bb03aa496a 100644 --- a/tests/v1/attention/test_mla_prefill_selector.py +++ b/tests/v1/attention/test_mla_prefill_selector.py @@ -63,6 +63,7 @@ class TestGetMLAPrefillBackend: with patch("vllm.platforms.current_platform") as mock_platform: mock_platform.get_device_capability.return_value = None + mock_platform.is_cpu.return_value = False backend = get_mla_prefill_backend(vllm_config) assert backend.get_name() == "FLASH_ATTN" diff --git a/vllm/_custom_ops.py b/vllm/_custom_ops.py index 5d63ffa0de..550bf94c43 100644 --- a/vllm/_custom_ops.py +++ b/vllm/_custom_ops.py @@ -3937,6 +3937,109 @@ def cpu_attention_with_kv_cache( ) +def cpu_mla_decode( + query: torch.Tensor, + k_buffer: torch.Tensor, + v_buffer: torch.Tensor, + output: torch.Tensor, + key: torch.Tensor | None, + value: torch.Tensor | None, + loc: torch.Tensor | None, + attn_logits: torch.Tensor, + req_to_token: torch.Tensor, + req_pool_indices: torch.Tensor, + seq_lens: torch.Tensor, + sm_scale: float, + logit_cap: float, + is_cross_attn: bool, + sliding_window_size: int, + encoder_lens: torch.Tensor | None, + sinks: torch.Tensor | None, +) -> None: + torch.ops._C.decode_attention_cpu( + query, + k_buffer, + v_buffer, + output, + key, + value, + loc, + attn_logits, + req_to_token, + req_pool_indices, + seq_lens, + sm_scale, + logit_cap, + is_cross_attn, + sliding_window_size, + encoder_lens, + sinks, + ) + + +def cpu_mla_extend( + q_extend: torch.Tensor, + k_extend: torch.Tensor | None, + v_extend: torch.Tensor | None, + o_extend: torch.Tensor, + k_buffer: torch.Tensor, + v_buffer: torch.Tensor, + req_to_token: torch.Tensor, + req_pool_indices: torch.Tensor, + seq_lens: torch.Tensor, + extend_seq_lens: torch.Tensor, + extend_start_loc: torch.Tensor, + max_len_extend: int, + sm_scale: float, + logit_cap: float, + is_cross_attn: bool, + sliding_window_size: int, + encoder_lens: torch.Tensor | None, + sinks: torch.Tensor | None, + tree_mask: torch.Tensor | None = None, +) -> None: + torch.ops._C.extend_attention_cpu( + q_extend, + k_extend, + v_extend, + o_extend, + k_buffer, + v_buffer, + req_to_token, + req_pool_indices, + seq_lens, + extend_seq_lens, + extend_start_loc, + max_len_extend, + sm_scale, + logit_cap, + is_cross_attn, + sliding_window_size, + encoder_lens, + sinks, + tree_mask, + ) + + +def bmm_cpu( + out: torch.Tensor, + mat1: torch.Tensor, + mat2: torch.Tensor, + is_vnni: bool, + scale: torch.Tensor | None = None, +) -> None: + torch.ops._C.bmm_cpu(out, mat1, mat2, is_vnni, scale) + + +def amx_mla_concat_and_cache( + kv_c_normed: torch.Tensor, + k_pe: torch.Tensor, + kv_cache: torch.Tensor, + slot_mapping: torch.Tensor, +) -> None: + torch.ops._C.concat_and_cache_mla_cpu(kv_c_normed, k_pe, kv_cache, slot_mapping) + + def cpu_gemm_wna16( input: torch.Tensor, q_weight: torch.Tensor, diff --git a/vllm/model_executor/layers/attention/mla_attention.py b/vllm/model_executor/layers/attention/mla_attention.py index 27ccd8ca0c..37dc2c95ff 100644 --- a/vllm/model_executor/layers/attention/mla_attention.py +++ b/vllm/model_executor/layers/attention/mla_attention.py @@ -540,6 +540,11 @@ class MLAAttention(nn.Module, AttentionLayerBase): **extra_impl_args, ) self.q_pad_num_heads = getattr(self.impl, "q_pad_num_heads", None) + self.is_amx_bmm_enabled = getattr(self.impl, "uses_amx_bmm", False) + # AMX reads kv_b_proj's weight directly and never calls it live; the + # reference CPU MLA backend calls it but isn't perf-critical. Skip + # the packed-kernel dispatch either way. + kv_b_proj._cpu_skip_gemm_dispatch = True self.use_direct_call = not current_platform.opaque_attention_op() vllm_config = get_current_vllm_config() @@ -908,6 +913,20 @@ class MLAAttention(nn.Module, AttentionLayerBase): group_size=128, transpose_bm=True, ) + elif self.is_amx_bmm_enabled: + # bmm_cpu computes out[n] = mat1[n] @ mat2[n]^T against + # AMXMLAImpl's own (N, L, P) packed W_UK -- same as prefill. + N, B, P = mqa_q_nope.shape + L = self.kv_lora_rank + mqa_ql_nope = mqa_q_nope.new_empty((N, B, L)) + ops.bmm_cpu( + mqa_ql_nope, + mqa_q_nope, + self.impl._w_uk_packed, # type: ignore[attr-defined] + True, + None, + ) + mqa_ql_nope = mqa_ql_nope.transpose(0, 1) else: # Pads the head_dim if necessary (for the underlying kernel) N, B, P = mqa_q_nope.shape @@ -1026,6 +1045,18 @@ class MLAAttention(nn.Module, AttentionLayerBase): return output_padded def process_weights_after_loading(self, act_dtype: torch.dtype): + # Let per-backend impls do their own weight packing first (no-op + # unless overridden), mirroring Attention.process_weights_after_loading. + self.impl.process_weights_after_loading(act_dtype) + + if self.is_amx_bmm_enabled: + # AMXMLAImpl already packed its own W_UK/W_UV above, for both + # prefill and decode. Release the now-unused raw weight. + self.kv_b_proj.weight = torch.nn.Parameter( + torch.empty(0), requires_grad=False + ) + return + # we currently do not have quantized bmm's which are needed for # `W_UV` and `W_UK_T`, we just store fp16/bf16 copies and perform # the bmm's in 16-bit, the extra memory overhead of this is fairly low @@ -1192,6 +1223,16 @@ class MLAAttention(nn.Module, AttentionLayerBase): x = rocm_aiter_ops.triton_fp8_bmm( x, self.W_V, self.W_V_scale, group_size=128, transpose_bm=True, YQ=out ) + elif self.is_amx_bmm_enabled: + # bmm_cpu computes out[n] = mat1[n] @ mat2[n]^T against + # AMXMLAImpl's own (N, V, L) packed W_UV -- same as prefill. + ops.bmm_cpu( + out.transpose(0, 1), + x, + self.impl._w_uv_packed, # type: ignore[attr-defined] + True, + None, + ) else: # Multiply + Transpose (N, B, L) x (N, L, V)->(N, B, V)->(B, N, V) torch.bmm(x, self.W_UV, out=out.transpose(0, 1)) diff --git a/vllm/model_executor/layers/linear.py b/vllm/model_executor/layers/linear.py index 050e7a0acd..1e31a64d68 100644 --- a/vllm/model_executor/layers/linear.py +++ b/vllm/model_executor/layers/linear.py @@ -197,6 +197,12 @@ class UnquantizedLinearMethod(LinearMethodBase): def process_weights_after_loading(self, layer: torch.nn.Module) -> None: if current_platform.is_cpu(): + # MLA's kv_b_proj (see `_cpu_skip_gemm_dispatch`): not + # perf-critical, so skip packing and use a plain fallback. + if getattr(layer, "_cpu_skip_gemm_dispatch", False): + layer.cpu_linear = torch.nn.functional.linear + return + from vllm.model_executor.layers.utils import dispatch_cpu_unquantized_gemm dispatch_cpu_unquantized_gemm(layer, remove_weight=True) diff --git a/vllm/platforms/cpu.py b/vllm/platforms/cpu.py index dd18363e0d..02ca99ff33 100644 --- a/vllm/platforms/cpu.py +++ b/vllm/platforms/cpu.py @@ -89,13 +89,30 @@ class CpuPlatform(Platform): if attn_selector_config.use_sparse: raise NotImplementedError("Sparse Attention is not supported on CPU.") if attn_selector_config.use_mla: + amx_available = ( + cls.get_cpu_architecture() == CpuArchEnum.X86 + and torch.cpu._is_amx_tile_supported() + ) + if amx_available and selected_backend != AttentionBackendEnum.CPU_MLA: + # Prefer AMX when available, unless CPU_MLA was explicitly requested. + if ( + selected_backend + and selected_backend != AttentionBackendEnum.AMX_MLA + ): + logger.info("Cannot use %s backend on CPU.", selected_backend) + logger.info_once("Using %s backend.", AttentionBackendEnum.AMX_MLA.name) + return AttentionBackendEnum.AMX_MLA.get_path() # Reference MLA implementation on CPU. Performance is not the # goal here; the backend simply wires the CPU decode kernel # (`mla_decode_kvcache`) and an SDPA-based prefill together with # the shared MLA scaffolding so that DeepSeek-style models can # execute on CPU. - if selected_backend and selected_backend != AttentionBackendEnum.CPU_MLA: + if selected_backend and selected_backend not in ( + AttentionBackendEnum.CPU_MLA, + AttentionBackendEnum.AMX_MLA, + ): logger.info("Cannot use %s backend on CPU.", selected_backend) + logger.info_once("Using %s backend.", AttentionBackendEnum.CPU_MLA.name) return AttentionBackendEnum.CPU_MLA.get_path() if selected_backend and selected_backend != AttentionBackendEnum.CPU_ATTN: logger.info("Cannot use %s backend on CPU.", selected_backend) @@ -134,11 +151,20 @@ class CpuPlatform(Platform): # The CPU MLA decode kernel only compiles with block_size=16 today # (see csrc/cpu/mla_decode.cpp). If the model uses MLA we override # the default block size regardless of user preference to avoid a - # runtime kernel dispatch failure. + # runtime kernel dispatch failure. AMX MLA has no such constraint + # (same AMX-available condition as get_attn_backend_cls), so it's + # excluded from this override. cpu_mla_enabled = model_config is not None and getattr( model_config, "use_mla", False ) - if cpu_mla_enabled: + amx_mla_enabled = ( + cpu_mla_enabled + and cls.get_cpu_architecture() == CpuArchEnum.X86 + and torch.cpu._is_amx_tile_supported() + and vllm_config.attention_config.backend != AttentionBackendEnum.CPU_MLA + ) + reference_cpu_mla_enabled = cpu_mla_enabled and not amx_mla_enabled + if reference_cpu_mla_enabled: if cache_config.user_specified_block_size and cache_config.block_size != 16: logger.warning( "CPU MLA backend requires block_size=16, overriding " @@ -149,7 +175,7 @@ class CpuPlatform(Platform): elif not cache_config.user_specified_block_size: cache_config.block_size = 128 - if not cpu_mla_enabled and cache_config.block_size % 32 != 0: + if not reference_cpu_mla_enabled and cache_config.block_size % 32 != 0: logger.warning( "CPU backend prefers block_size is multiples of 32, " "otherwise the performance is not optimized." @@ -342,7 +368,7 @@ class CpuPlatform(Platform): vllm_config.parallel_config.tensor_parallel_size ) - if model_config is not None and model_config.use_mla: + if model_config is not None and model_config.use_mla and not amx_mla_enabled: logger.info_once( "MLA is enabled on a non-GPU platform; forcing chunked " "prefill and prefix caching to be disabled." diff --git a/vllm/v1/attention/backends/mla/amx_mla.py b/vllm/v1/attention/backends/mla/amx_mla.py new file mode 100644 index 0000000000..5f9ac7f17a --- /dev/null +++ b/vllm/v1/attention/backends/mla/amx_mla.py @@ -0,0 +1,405 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""AMX-only, high-performance MLA backend for DeepSeek V2/V3/R1 on CPU. + +Built on the AMX decode/extend/bmm kernels vendored under +``csrc/cpu/sgl-kernels/``, plugged into vLLM's ``MLACommonBackend``/ +``MLACommonImpl`` abstraction the same way every other concrete MLA backend +(``TritonMLAImpl``, etc.) does. + +This is a separate backend from the reference ``CPUMLABackend`` +(``vllm/v1/attention/backends/mla/cpu_mla.py``): that one targets every CPU +(any dtype, block_size=16, ``mla_decode_kvcache`` decode kernel + inherited +SDPA-style prefill) as a functional/CI reference, not performance. This +backend instead requires AMX (bf16 only, block_size a multiple of 32) and is +selected by the platform layer in preference to the reference backend +whenever the host supports it -- see ``CpuPlatform.get_attn_backend_cls``. + +Two points where this backend differs structurally from the GPU backends, +both explained in the CPU MLA design plan: + +- ``forward_mha`` is fully overridden (not inherited from + ``MLACommonBaseImpl``): the GPU "compute-friendly" prefill path depends on + a pluggable ``MLAPrefillBackend`` and CUDA-only chunked-context gather ops, + neither of which exist on CPU. Instead, this attends directly in + latent-MQA space via the ``extend_attention_cpu`` kernel, which handles + cached-prefix continuation and fresh prefill in one causal pass (the KV + cache already contains the new tokens by the time this runs, since + ``do_kv_cache_update`` executes first). +- Weight absorption for the prefill path is done with this impl's own + VNNI-packed copies of ``W_UK``/``W_UV`` (computed in + ``process_weights_after_loading``, using ``bmm_cpu``), rather than reading + the generic ``layer.W_UK_T``/``layer.W_UV``/``layer._v_up_proj`` the way + GPU backends do, since ``forward_mha``'s abstract signature has no + ``layer`` parameter (only ``forward_mqa`` does) and that signature is left + untouched. The decode path needs no such packing: ``MLAAttention.forward_impl`` + already absorbs/de-absorbs Q around the ``forward_mqa`` call using the + generic (unpacked) ``layer.W_UK_T``/``layer._v_up_proj``, so + ``forward_mqa`` here only has to invoke the decode kernel. +""" + +from __future__ import annotations + +from typing import ClassVar + +import torch + +from vllm import _custom_ops as ops +from vllm.model_executor.layers.attention.mla_attention import ( + MLACommonBackend, + MLACommonImpl, + MLACommonMetadata, + MLACommonMetadataBuilder, +) +from vllm.model_executor.layers.quantization.utils.quant_utils import ( + get_and_maybe_dequant_weights, +) +from vllm.platforms import current_platform +from vllm.v1.attention.backend import AttentionLayer, AttentionType, MultipleOf + +_MIN_WORK_PER_SPLIT = 512 +_SPLIT_OCCUPANCY_MULTIPLIER = 2 + + +def _compute_num_kv_splits(max_seq_len: int, num_threads: int) -> int: + """Mirrors TritonMLAImpl's _compute_num_kv_splits, using the CPU thread + count in place of SM count.""" + ideal_splits = 1 + while ideal_splits < max(1, max_seq_len // _MIN_WORK_PER_SPLIT): + ideal_splits *= 2 + max_splits = num_threads * _SPLIT_OCCUPANCY_MULTIPLIER + return min(ideal_splits, max_splits) + + +def _expand_block_table(block_table: torch.Tensor, block_size: int) -> torch.Tensor: + """Adapter: vLLM's block-table paging -> the flat per-(request, position) + physical row index the decode/extend kernels expect. Called once per step + from ``AMXMLAMetadataBuilder.build``, not per layer. + """ + offsets = torch.arange( + block_size, device=block_table.device, dtype=block_table.dtype + ) + flat = block_table.unsqueeze(-1) * block_size + offsets + return flat.reshape(block_table.size(0), -1).contiguous() + + +class AMXMLABackend(MLACommonBackend): + supported_dtypes: ClassVar[list[torch.dtype]] = [torch.bfloat16] + + @classmethod + def get_supported_head_sizes(cls) -> list[int]: + # kv_lora_rank(512) + qk_rope_head_dim(64), DeepSeek V2/V3/R1/V3.2. + return [576] + + @staticmethod + def get_supported_kernel_block_sizes() -> list[int | MultipleOf]: + return [MultipleOf(32)] + + @classmethod + def supports_block_size(cls, block_size: int | None) -> bool: + if block_size is None: + return True + return block_size % 32 == 0 + + @staticmethod + def get_name() -> str: + return "AMX_MLA" + + @staticmethod + def get_impl_cls() -> type[AMXMLAImpl]: + return AMXMLAImpl + + @staticmethod + def get_builder_cls() -> type[AMXMLAMetadataBuilder]: + return AMXMLAMetadataBuilder + + +class AMXMLAMetadataBuilder(MLACommonMetadataBuilder[MLACommonMetadata]): + def build(self, common_prefix_len, common_attn_metadata, fast_build: bool = False): + attn_metadata = super().build( + common_prefix_len, common_attn_metadata, fast_build + ) + block_size = self.kv_cache_spec.block_size + if attn_metadata.decode is not None: + # Built once per step and reused by every layer, instead of + # recomputing the same tensor arithmetic in every forward_mqa call. + decode = attn_metadata.decode + decode.req_to_token = _expand_block_table( # type: ignore[attr-defined] + decode.block_table, block_size + ) + decode.seq_lens_i64 = decode.seq_lens.to(torch.int64) # type: ignore[attr-defined] + decode.req_pool_indices = torch.arange( # type: ignore[attr-defined] + decode.block_table.size(0), + dtype=torch.int64, + device=decode.block_table.device, + ) + decode.num_kv_splits = _compute_num_kv_splits( # type: ignore[attr-defined] + attn_metadata.max_seq_len, current_platform.num_compute_units() + ) + if attn_metadata.prefill is not None: + # cpu_seq_lens: total per-request context length (prefix + new + # tokens), the one thing the extend kernel needs that isn't + # already on the shared MLACommonPrefillMetadata. + prefill = attn_metadata.prefill + num_decodes = attn_metadata.num_decodes + num_prefills = attn_metadata.num_prefills + prefill.cpu_seq_lens = common_attn_metadata.seq_lens[ # type: ignore[attr-defined] + num_decodes : num_decodes + num_prefills + ].to(torch.int64) + prefill.req_to_token = _expand_block_table( # type: ignore[attr-defined] + prefill.block_table, block_size + ).to(torch.int64) + prefill.req_pool_indices = torch.arange( # type: ignore[attr-defined] + prefill.block_table.size(0), + dtype=torch.int64, + device=prefill.block_table.device, + ) + query_start_loc_i64 = prefill.query_start_loc.to(torch.int64) + extend_seq_lens = query_start_loc_i64[1:] - query_start_loc_i64[:-1] + prefill.extend_seq_lens = extend_seq_lens # type: ignore[attr-defined] + prefill.extend_start_loc = query_start_loc_i64[:-1] # type: ignore[attr-defined] + prefill.max_len_extend = int(extend_seq_lens.max().item()) # type: ignore[attr-defined] + return attn_metadata + + +class AMXMLAImpl(MLACommonImpl[MLACommonMetadata]): + # Tells MLAAttention (mla_attention.py) to use this impl's own packed + # W_UK/W_UV + bmm_cpu for the shared decode absorb/de-absorb, mirroring + # is_aiter_triton_fp8_bmm_enabled/is_aiter_triton_fp4_bmm_enabled. + uses_amx_bmm = True + + def __init__( + self, + num_heads: int, + head_size: int, + scale: float, + num_kv_heads: int, + alibi_slopes: list[float] | None, + sliding_window: int | None, + kv_cache_dtype: str, + logits_soft_cap: float | None, + attn_type: str, + kv_sharing_target_layer_name: str | None, + **mla_args, + ) -> None: + super().__init__( + num_heads, + head_size, + scale, + num_kv_heads, + alibi_slopes, + sliding_window, + kv_cache_dtype, + logits_soft_cap, + attn_type, + kv_sharing_target_layer_name, + **mla_args, + ) + unsupported_features = [alibi_slopes, sliding_window, logits_soft_cap] + if any(unsupported_features): + raise NotImplementedError( + "AMXMLAImpl does not support one of the following: " + "alibi_slopes, sliding_window, logits_soft_cap" + ) + if attn_type != AttentionType.DECODER: + raise NotImplementedError( + "Encoder self-attention and encoder/decoder cross-attention " + "are not implemented for AMXMLAImpl" + ) + self._w_uk_packed: torch.Tensor | None = None + self._w_uv_packed: torch.Tensor | None = None + + def process_weights_after_loading(self, act_dtype: torch.dtype) -> None: + # forward_mha (unlike forward_mqa) receives no `layer` argument, so it + # cannot read the generic layer.W_UK_T/W_UV/._v_up_proj that + # MLAAttention.process_weights_after_loading already computes. This + # impl derives its own copies from self.kv_b_proj (already available + # via the constructor) and VNNI-packs them once, ahead of time, for + # use with bmm_cpu in forward_mha. + kv_b_proj_weight = get_and_maybe_dequant_weights( + self.kv_b_proj, out_dtype=act_dtype + ).T + kv_b_proj_weight = kv_b_proj_weight.view( + self.kv_lora_rank, self.num_heads, self.qk_nope_head_dim + self.v_head_dim + ) + w_uk, w_uv = kv_b_proj_weight.split( + [self.qk_nope_head_dim, self.v_head_dim], dim=-1 + ) + # bmm_cpu computes out[b] = mat1[b] @ mat2[b]^T (Linear-weight + # convention: mat2[b] is (OUT, IN)). + # absorb: ql_nope(N,B,L) = q_nope(N,B,P) @ w_uk_for_bmm(N,L,P)^T + # de-absorb: out(N,B,V) = attn_out(N,B,L) @ w_uv_for_bmm(N,V,L)^T + w_uk_for_bmm = w_uk.permute(1, 0, 2).contiguous() # (N, L, P) + w_uv_for_bmm = w_uv.permute(1, 2, 0).contiguous() # (N, V, L) + self._w_uk_packed = torch.ops._C.convert_weight_packed(w_uk_for_bmm) + self._w_uv_packed = torch.ops._C.convert_weight_packed(w_uv_for_bmm) + + def do_kv_cache_update( + self, + kv_c_normed: torch.Tensor, + k_pe: torch.Tensor, + kv_cache: torch.Tensor, + slot_mapping: torch.Tensor, + kv_cache_dtype: str, + k_scale: torch.Tensor, + ) -> None: + # Overrides the default (CUDA-only concat_and_cache_mla). + if kv_cache.numel() == 0: + return + assert kv_cache_dtype == "auto", ( + "AMXMLAImpl only supports an unquantized (bf16) KV cache; fp8 " + "KV cache for CPU MLA is not yet implemented." + ) + ops.amx_mla_concat_and_cache( + kv_c_normed, k_pe.squeeze(1), kv_cache, slot_mapping.flatten() + ) + + def forward_mqa( + self, + q: torch.Tensor | tuple[torch.Tensor, torch.Tensor], + kv_c_and_k_pe_cache: torch.Tensor, + attn_metadata: MLACommonMetadata, + layer: AttentionLayer, + ) -> tuple[torch.Tensor, torch.Tensor | None]: + # q arrives already absorbed (tensor or (ql_nope, q_pe) tuple) -- + # MLAAttention.forward_impl does the Q-absorption bmm generically + # before calling forward_mqa, using layer.W_UK_T. De-absorption is + # likewise done generically afterwards via layer._v_up_proj. Nothing + # left for this impl to do beyond invoking the decode kernel. + if isinstance(q, tuple): + q = torch.cat(q, dim=-1) + assert attn_metadata.decode is not None + + num_tokens = q.shape[0] + kv_cache_flat = kv_c_and_k_pe_cache.view(-1, 1, self.head_size) + v_buffer = kv_cache_flat[..., : self.kv_lora_rank] + + seq_lens = attn_metadata.decode.seq_lens_i64 # type: ignore[attr-defined] + req_to_token = attn_metadata.decode.req_to_token # type: ignore[attr-defined] + req_pool_indices = attn_metadata.decode.req_pool_indices # type: ignore[attr-defined] + + num_kv_splits = attn_metadata.decode.num_kv_splits # type: ignore[attr-defined] + o = torch.zeros( + num_tokens, + self.num_heads, + self.kv_lora_rank, + dtype=q.dtype, + device=q.device, + ) + attn_logits = torch.zeros( + num_tokens, + self.num_heads, + num_kv_splits, + self.kv_lora_rank + 1, + dtype=torch.float32, + device=q.device, + ) + + ops.cpu_mla_decode( + q, + kv_cache_flat, + v_buffer, + o, + None, + None, + None, # loc: only meaningful when key/value are given; we + # already wrote the new K/V via do_kv_cache_update before this. + attn_logits, + req_to_token, + req_pool_indices, + seq_lens, + self.scale, + 0.0, + False, + 0, + None, + None, + ) + return o, None + + def forward_mha( + self, + q: torch.Tensor, + kv_c_normed: torch.Tensor, + k_pe: torch.Tensor, + kv_c_and_k_pe_cache: torch.Tensor, + attn_metadata: MLACommonMetadata, + k_scale: torch.Tensor, + output: torch.Tensor, + output_scale: torch.Tensor | None = None, + ) -> None: + assert output_scale is None, ( + "AMXMLAImpl.forward_mha does not support fused output quantization" + ) + prefill = attn_metadata.prefill + assert prefill is not None + assert self._w_uk_packed is not None and self._w_uv_packed is not None + + num_tokens = q.shape[0] + num_heads = self.num_heads + + q_nope, q_pe = q.split([self.qk_nope_head_dim, self.qk_rope_head_dim], dim=-1) + + # Absorb: (N, B, P) x (N, L, P)^T -> (N, B, L) + # bmm_cpu tolerates non-contiguous mat1/out as long as their last dim + # is contiguous, so q_nope's transpose and ql_nope's transposed view + # can be passed directly -- no copy needed for either. + q_nope_t = q_nope.transpose(0, 1) + ql_nope = torch.empty( + num_tokens, num_heads, self.kv_lora_rank, dtype=q.dtype, device=q.device + ) + ops.bmm_cpu(ql_nope.transpose(0, 1), q_nope_t, self._w_uk_packed, True, None) + mqa_q = torch.cat([ql_nope, q_pe], dim=-1) + + kv_cache_flat = kv_c_and_k_pe_cache.view(-1, 1, self.head_size) + v_buffer = kv_cache_flat[..., : self.kv_lora_rank] + + req_to_token = prefill.req_to_token # type: ignore[attr-defined] + req_pool_indices = prefill.req_pool_indices # type: ignore[attr-defined] + + seq_lens = prefill.cpu_seq_lens # type: ignore[attr-defined] + extend_seq_lens = prefill.extend_seq_lens # type: ignore[attr-defined] + extend_start_loc = prefill.extend_start_loc # type: ignore[attr-defined] + max_len_extend = prefill.max_len_extend # type: ignore[attr-defined] + + # k_extend/v_extend: the new tokens' own latent K/V, aliased in one + # 576-wide buffer (v_extend is a view of k_extend's first + # kv_lora_rank columns) so the kernel can gather both from one read, + # mirroring the cache's own layout. + k_extend = torch.cat( + [kv_c_normed, k_pe.reshape(num_tokens, self.qk_rope_head_dim)], dim=-1 + ).unsqueeze(1) + v_extend = k_extend[..., : self.kv_lora_rank] + + attn_out = torch.empty( + num_tokens, num_heads, self.kv_lora_rank, dtype=q.dtype, device=q.device + ) + ops.cpu_mla_extend( + mqa_q, + k_extend, + v_extend, + attn_out, + kv_cache_flat, + v_buffer, + req_to_token, + req_pool_indices, + seq_lens, + extend_seq_lens, + extend_start_loc, + max_len_extend, + self.scale, + 0.0, + False, + 0, + None, + None, + None, + ) + + # De-absorb directly into `output` (view, no extra copy on either side): + # (N, B, L) x (N, V, L)^T -> (N, B, V) + attn_out_t = attn_out.transpose(0, 1) + output_view = output.view(num_tokens, num_heads, self.v_head_dim).transpose( + 0, 1 + ) + ops.bmm_cpu(output_view, attn_out_t, self._w_uv_packed, True, None) diff --git a/vllm/v1/attention/backends/mla/prefill/cpu_native.py b/vllm/v1/attention/backends/mla/prefill/cpu_native.py new file mode 100644 index 0000000000..52ba8b7e24 --- /dev/null +++ b/vllm/v1/attention/backends/mla/prefill/cpu_native.py @@ -0,0 +1,61 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Inert MLA prefill backend for the CPU platform. + +CPU's MLA attention impl (``CPUMLAImpl``) fully overrides ``forward_mha`` with a +kernel that attends directly against the paged latent KV cache (covering both +fresh prefill and cached-prefix continuation in one pass), so it never calls +into a pluggable ``MLAPrefillBackend``'s ``run_prefill_new_tokens``/ +``run_prefill_context_chunk``. This class exists only to satisfy +``MLACommonMetadataBuilder``'s structural requirement that +``MLAAttention.prefill_backend`` be a real, cloneable object. +""" + +from typing import TYPE_CHECKING + +import torch + +from vllm.v1.attention.backends.mla.prefill.base import MLAPrefillBackend + +if TYPE_CHECKING: + from vllm.model_executor.layers.attention.mla_attention import ( + MLACommonPrefillMetadata, + ) + + +class CPUNativeMLAPrefillBackend(MLAPrefillBackend): + """Placeholder MLA prefill backend for CPU; never actually invoked.""" + + @staticmethod + def get_name() -> str: + return "CPU_NATIVE" + + def run_prefill_new_tokens( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + return_softmax_lse: bool, + out: torch.Tensor | None = None, + output_scale: torch.Tensor | None = None, + ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: + raise AssertionError( + "CPUNativeMLAPrefillBackend.run_prefill_new_tokens is unreachable: " + "CPUMLAImpl.forward_mha fully overrides the dense-MHA prefill path " + "and never calls the pluggable prefill backend." + ) + + def run_prefill_context_chunk( + self, + chunk: "MLACommonPrefillMetadata.ContextChunk", + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + out: torch.Tensor | None = None, + ) -> tuple[torch.Tensor, torch.Tensor]: + raise AssertionError( + "CPUNativeMLAPrefillBackend.run_prefill_context_chunk is " + "unreachable: CPUMLAImpl.forward_mha fully overrides the " + "dense-MHA prefill path and never calls the pluggable prefill " + "backend." + ) diff --git a/vllm/v1/attention/backends/mla/prefill/registry.py b/vllm/v1/attention/backends/mla/prefill/registry.py index 0d818a084b..b521f4ff1e 100644 --- a/vllm/v1/attention/backends/mla/prefill/registry.py +++ b/vllm/v1/attention/backends/mla/prefill/registry.py @@ -52,6 +52,9 @@ class MLAPrefillBackendEnum(Enum, metaclass=_MLAPrefillBackendEnumMeta): "vllm.v1.attention.backends.mla.prefill.aiter_flash_attn." "AiterFlashAttnPrefillBackend" ) + CPU_NATIVE = ( + "vllm.v1.attention.backends.mla.prefill.cpu_native.CPUNativeMLAPrefillBackend" + ) # Placeholder for third-party/custom backends - must be registered before use # set to None to avoid alias with other backend, whose value is an empty string CUSTOM = None diff --git a/vllm/v1/attention/backends/mla/prefill/selector.py b/vllm/v1/attention/backends/mla/prefill/selector.py index c313c90177..5cdfe3c6ca 100644 --- a/vllm/v1/attention/backends/mla/prefill/selector.py +++ b/vllm/v1/attention/backends/mla/prefill/selector.py @@ -109,6 +109,8 @@ def get_mla_prefill_backend( device_capability = current_platform.get_device_capability() if device_capability is None: + if current_platform.is_cpu(): + return MLAPrefillBackendEnum.CPU_NATIVE.get_class() logger.info_once( "Device capability not available, using FlashAttention MLA prefill backend." ) diff --git a/vllm/v1/attention/backends/registry.py b/vllm/v1/attention/backends/registry.py index 646227cc4f..6fc9dd1e1f 100644 --- a/vllm/v1/attention/backends/registry.py +++ b/vllm/v1/attention/backends/registry.py @@ -123,6 +123,7 @@ class AttentionBackendEnum(Enum, metaclass=_AttentionBackendEnumMeta): ) CPU_ATTN = "vllm.v1.attention.backends.cpu_attn.CPUAttentionBackend" CPU_MLA = "vllm.v1.attention.backends.mla.cpu_mla.CPUMLABackend" + AMX_MLA = "vllm.v1.attention.backends.mla.amx_mla.AMXMLABackend" TURBOQUANT = "vllm.v1.attention.backends.turboquant_attn.TurboQuantAttentionBackend" # Placeholder for third-party/custom backends - must be registered before use # set to None to avoid alias with other backend, whose value is an empty string