@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
cpp 111 lines 4.99 KB
Raw
1 // vLLM-native CPU cache-write op for MLA's single-latent-buffer KV cache.
2 // `concat_and_cache_mla` (the generic op used by every GPU backend) is
3 // CUDA-only; this is the CPU counterpart, adapted in spirit from SGLang's
4 // `store_cache_cpu` (csrc/cpu/kvcache.cpp).
5 //
6 // Difference from upstream: writes two source tensors (`kv_c_normed`,
7 // `k_pe`) into two column-offset ranges of the SAME destination row.
8 // Upstream assumes k/v land in two independent, equal-row-width cache
9 // tensors -- MLA's cache is one 576-wide buffer whose two column ranges
10 // (512-wide, 64-wide) don't match that per-token stride, so
11 // `store_cache_cpu` can't be reused as-is.
12
13 #include "common.h"
14 #include "vec.h"
15
16 namespace {
17
18 template <typename scalar_t>
19 inline void copy_stub(scalar_t* __restrict__ dst,
20 const scalar_t* __restrict__ src, int64_t size) {
21 int64_t d = 0;
22 #if defined(CPU_CAPABILITY_AVX512)
23 using Vec = at::vec::Vectorized<scalar_t>;
24 constexpr int64_t kVecSize = Vec::size();
25 for (; d <= size - kVecSize; d += kVecSize) {
26 Vec data = Vec::loadu(src + d);
27 data.store(dst + d);
28 }
29 #endif
30 for (; d < size; ++d) {
31 dst[d] = src[d];
32 }
33 }
34
35 template <typename scalar_t, typename index_t>
36 void concat_and_cache_mla_kernel_impl(
37 const scalar_t* __restrict__ kv_c_normed, // [num_tokens, kv_lora_rank]
38 const scalar_t* __restrict__ k_pe, // [num_tokens, qk_rope_head_dim]
39 scalar_t* __restrict__ kv_cache, // [.., kv_lora_rank + qk_rope_head_dim]
40 const index_t* __restrict__ slot_mapping, // [num_tokens]
41 int64_t num_tokens, int64_t kv_lora_rank, int64_t qk_rope_head_dim,
42 int64_t kv_c_stride, int64_t k_pe_stride, int64_t cache_stride) {
43 at::parallel_for(0, num_tokens, 0, [&](int64_t begin, int64_t end) {
44 for (int64_t i = begin; i < end; ++i) {
45 int64_t slot = static_cast<int64_t>(slot_mapping[i]);
46 if (slot < 0) {
47 // padded/invalid token, matches `concat_and_cache_mla`'s semantics.
48 continue;
49 }
50 scalar_t* __restrict__ cache_row = kv_cache + slot * cache_stride;
51 copy_stub(cache_row, kv_c_normed + i * kv_c_stride, kv_lora_rank);
52 copy_stub(cache_row + kv_lora_rank, k_pe + i * k_pe_stride,
53 qk_rope_head_dim);
54 }
55 });
56 }
57
58 } // namespace
59
60 // kv_c_normed : [num_tokens, kv_lora_rank]
61 // k_pe : [num_tokens, qk_rope_head_dim] or [num_tokens, 1,
62 // qk_rope_head_dim]
63 // kv_cache : [num_blocks, block_size, kv_lora_rank + qk_rope_head_dim]
64 // slot_mapping: [num_tokens] int32/int64, absolute physical row index
65 // (block_id * block_size + block_offset); negative entries are
66 // skipped (padded tokens), matching `concat_and_cache_mla`.
67 void concat_and_cache_mla_cpu(const at::Tensor& kv_c_normed,
68 const at::Tensor& k_pe, at::Tensor& kv_cache,
69 const at::Tensor& slot_mapping) {
70 TORCH_CHECK(kv_c_normed.dim() == 2,
71 "kv_c_normed must be 2D [num_tokens, kv_lora_rank]");
72 TORCH_CHECK(k_pe.dim() == 2 || k_pe.dim() == 3, "k_pe must be 2D or 3D");
73 TORCH_CHECK(kv_cache.dim() == 3,
74 "kv_cache must be 3D [num_blocks, block_size, head_size]");
75 TORCH_CHECK(kv_c_normed.stride(-1) == 1,
76 "kv_c_normed innermost dim must be contiguous");
77 TORCH_CHECK(k_pe.stride(-1) == 1, "k_pe innermost dim must be contiguous");
78 TORCH_CHECK(kv_cache.stride(-1) == 1,
79 "kv_cache innermost dim must be contiguous");
80
81 int64_t num_tokens = kv_c_normed.size(0);
82 int64_t kv_lora_rank = kv_c_normed.size(1);
83 int64_t qk_rope_head_dim = k_pe.size(-1);
84 int64_t head_size = kv_cache.size(-1);
85 TORCH_CHECK(head_size == kv_lora_rank + qk_rope_head_dim,
86 "kv_cache head_size must equal kv_lora_rank + qk_rope_head_dim");
87 TORCH_CHECK(slot_mapping.size(0) == num_tokens, "slot_mapping size mismatch");
88
89 // Real physical per-token stride of the cache, read from the tensor's own
90 // strides rather than assumed to equal head_size (paged/pooled buffers are
91 // not guaranteed to have zero inter-row padding in general, even though in
92 // practice a freshly-allocated MLA cache is fully contiguous).
93 int64_t cache_stride = kv_cache.stride(1);
94
95 const auto dtype = kv_cache.scalar_type();
96 TORCH_CHECK(dtype == kv_c_normed.scalar_type() && dtype == k_pe.scalar_type(),
97 "concat_and_cache_mla_cpu: dtype mismatch");
98 const auto index_dtype = slot_mapping.scalar_type();
99 TORCH_CHECK(index_dtype == at::kLong || index_dtype == at::kInt,
100 "slot_mapping must be int32 or int64");
101
102 AT_DISPATCH_REDUCED_FLOATING_TYPES(dtype, "concat_and_cache_mla_cpu", [&] {
103 AT_DISPATCH_INDEX_TYPES(index_dtype, "concat_and_cache_mla_cpu_index", [&] {
104 concat_and_cache_mla_kernel_impl<scalar_t, index_t>(
105 kv_c_normed.data_ptr<scalar_t>(), k_pe.data_ptr<scalar_t>(),
106 kv_cache.data_ptr<scalar_t>(), slot_mapping.data_ptr<index_t>(),
107 num_tokens, kv_lora_rank, qk_rope_head_dim, kv_c_normed.stride(0),
108 k_pe.stride(0), cache_stride);
109 });
110 });
111 }