@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
cpp 217 lines 9.92 KB
Raw
1 // Adapted from
2 // https://github.com/sgl-project/sglang/tree/main/sgl-kernel/csrc/cpu:
3 // `fp8_paged_mqa_logits_cpu`/`fp8_paged_mqa_logits_cpu_impl` (DECODE path).
4 //
5 // Reads the paged indexer K-cache directly via `page_table` -- no separate
6 // eager gather step -- and processes the whole decode batch (one query token
7 // per request; the CPU indexer decode path doesn't support native MTP/spec
8 // decode yet, see `sparse_attn_indexer_cpu`) in a single `at::parallel_for`
9 // flattening `(batch, token)`, exactly like upstream. AVX512 `dot_fp8_128`
10 // dequant-dot-product and the split K-region-then-scale-region page byte
11 // layout are verbatim (that part of the upstream contract already matches
12 // vLLM's cache, written by `compress_norm_rope_store_indexer_cpu` in
13 // compressor.cpp).
14 //
15 // Differences from upstream:
16 // - `block_size` is a runtime argument instead of the hardcoded
17 // `kBlockSize = 64`. vLLM's indexer K-cache shares the model's single
18 // global `--block-size` KV-cache-group setting (confirmed empirically:
19 // the DeepSeek-V4-Flash TP1/TP4 baseline configs run with
20 // `block_size=256`, not 64 -- see `.agent-workspace/perf_baseline_tp*_
21 // server.log`), so the upstream `kvcache_fp8.size(1) == 64` shape
22 // contract does not hold in general.
23 // - `q_fp8`/`weight` drop the redundant `S_q=1` dimension (`[batch, heads,
24 // 128]`/`[batch, heads]` instead of `[batch, 1, heads, 128]`) since S_q is
25 // always 1 here; `weight` is fp32-only (already folds q_scale/softmax_
26 // scale/head_scale, matching this call site's actual dtype -- unlike
27 // upstream's bf16/fp16/fp32 dispatch, which vLLM's CPU indexer never
28 // produces).
29 // - `index_head_dim` (128) stays a compile-time constant, matching
30 // upstream: it's `config.index_head_dim`, a genuine DeepSeek-V4 model
31 // invariant, not a vLLM KV-cache config knob like `block_size`.
32
33 #include "common.h"
34 #include "vec.h"
35
36 #if !defined(CPU_CAPABILITY_AVX512)
37 static_assert(false,
38 "paged_mqa_logits.cpp is only ever compiled into the "
39 "AVX512+AMX-flagged _C target (see cmake/cpu_extension.cmake's "
40 "VLLM_EXT_SRC_SGL) and relies on that unconditionally -- "
41 "CPU_CAPABILITY_AVX512 must be defined here.");
42 #endif
43
44 namespace {
45
46 constexpr int64_t kIndexHeadDim = 128;
47 constexpr int64_t kIndexHeadDimWithScaleBytes = 132;
48
49 inline float dot_fp8_128(const uint8_t* k, const uint8_t* q) {
50 __m512 acc = _mm512_setzero_ps();
51 for (int64_t d = 0; d < kIndexHeadDim; d += 32) {
52 const __m256i k8 =
53 _mm256_loadu_si256(reinterpret_cast<const __m256i*>(k + d));
54 const __m256i q8 =
55 _mm256_loadu_si256(reinterpret_cast<const __m256i*>(q + d));
56 acc = _mm512_dpbf16_ps(acc, CVT_FP8_TO_BF16(k8), CVT_FP8_TO_BF16(q8));
57 }
58 return _mm512_reduce_add_ps(acc);
59 }
60
61 template <typename seq_t, typename page_t>
62 void fp8_paged_mqa_logits_cpu_impl(const at::Tensor& q_fp8,
63 const at::Tensor& kvcache_fp8,
64 const at::Tensor& weight,
65 const at::Tensor& seq_lens,
66 const at::Tensor& page_table,
67 at::Tensor& logits, int64_t block_size,
68 int64_t max_seq_len) {
69 const int64_t batch_size = q_fp8.size(0);
70 const int64_t num_heads = q_fp8.size(1);
71 const int64_t num_pages = kvcache_fp8.size(0);
72 // Row stride is threaded through explicitly rather than assumed equal to
73 // the page byte width: the K-cache may be a page of a larger pooled
74 // allocation shared across layers (matches store_cache.cpp/flash_mla.cpp's
75 // precedent).
76 const int64_t page_stride = kvcache_fp8.stride(0);
77 const int64_t pages_per_batch = page_table.size(1);
78 const int64_t scale_offset_bytes = block_size * kIndexHeadDim;
79
80 const auto* q_ptr = reinterpret_cast<const uint8_t*>(q_fp8.const_data_ptr());
81 const auto* cache_ptr =
82 reinterpret_cast<const uint8_t*>(kvcache_fp8.const_data_ptr());
83 const auto* weight_ptr = weight.const_data_ptr<float>();
84 const auto* seq_ptr = seq_lens.const_data_ptr<seq_t>();
85 const auto* page_ptr = page_table.const_data_ptr<page_t>();
86 auto* out_ptr = logits.data_ptr<float>();
87
88 at::parallel_for(
89 0, batch_size * max_seq_len, GRAIN_SIZE / kIndexHeadDim,
90 [&](int64_t begin, int64_t end) {
91 int64_t b{0}, token{0};
92 data_index_init(begin, b, batch_size, token, max_seq_len);
93 for (int64_t i = begin; i < end; ++i) {
94 const int64_t seq_len = static_cast<int64_t>(seq_ptr[b]);
95
96 if (token >= seq_len) {
97 data_index_step(b, batch_size, token, max_seq_len);
98 continue;
99 }
100
101 const int64_t q_batch_offset = b * num_heads * kIndexHeadDim;
102 const int64_t weight_batch_offset = b * num_heads;
103 const int64_t page_batch_offset = b * pages_per_batch;
104 float* out_row = out_ptr + b * max_seq_len;
105
106 const int64_t logical_page = token / block_size;
107 const int64_t token_in_page = token % block_size;
108
109 const int64_t physical_page =
110 static_cast<int64_t>(page_ptr[page_batch_offset + logical_page]);
111
112 const uint8_t* block = cache_ptr + physical_page * page_stride;
113 const uint8_t* k_token = block + token_in_page * kIndexHeadDim;
114 const float* scale_ptr =
115 reinterpret_cast<const float*>(block + scale_offset_bytes);
116 const float k_scale = scale_ptr[token_in_page];
117
118 float score_sum = 0.0f;
119 for (int64_t h = 0; h < num_heads; ++h) {
120 const uint8_t* q_head = q_ptr + q_batch_offset + h * kIndexHeadDim;
121 float dot = dot_fp8_128(k_token, q_head);
122 dot = std::max(dot, 0.0f);
123 score_sum += dot * weight_ptr[weight_batch_offset + h];
124 }
125
126 out_row[token] = score_sum * k_scale;
127 data_index_step(b, batch_size, token, max_seq_len);
128 }
129 });
130 }
131
132 template <typename seq_t>
133 void dispatch_page_type(const at::Tensor& q_fp8, const at::Tensor& kvcache_fp8,
134 const at::Tensor& weight, const at::Tensor& seq_lens,
135 const at::Tensor& page_table, at::Tensor& logits,
136 int64_t block_size, int64_t max_seq_len) {
137 if (page_table.scalar_type() == at::kInt) {
138 fp8_paged_mqa_logits_cpu_impl<seq_t, int32_t>(q_fp8, kvcache_fp8, weight,
139 seq_lens, page_table, logits,
140 block_size, max_seq_len);
141 } else if (page_table.scalar_type() == at::kLong) {
142 fp8_paged_mqa_logits_cpu_impl<seq_t, int64_t>(q_fp8, kvcache_fp8, weight,
143 seq_lens, page_table, logits,
144 block_size, max_seq_len);
145 } else {
146 TORCH_CHECK(false, "page_table must be int32 or int64");
147 }
148 }
149
150 } // namespace
151
152 at::Tensor fp8_paged_mqa_logits_cpu(at::Tensor& q_fp8, at::Tensor& kvcache_fp8,
153 at::Tensor& weight, at::Tensor& seq_lens,
154 at::Tensor& page_table, int64_t block_size,
155 int64_t max_seq_len) {
156 CHECK_INPUT(q_fp8);
157 CHECK_CPU(kvcache_fp8);
158 CHECK_INPUT(weight);
159 CHECK_INPUT(seq_lens);
160 CHECK_INPUT(page_table);
161 TORCH_CHECK(
162 q_fp8.dim() == 3 && q_fp8.scalar_type() == at::kFloat8_e4m3fn,
163 "fp8_paged_mqa_logits_cpu: q_fp8 must be a 3D float8_e4m3fn tensor "
164 "[batch, heads, 128]");
165 TORCH_CHECK(q_fp8.size(2) == kIndexHeadDim,
166 "fp8_paged_mqa_logits_cpu: q_fp8 head_dim must be 128");
167 TORCH_CHECK(kvcache_fp8.dim() == 2 && kvcache_fp8.scalar_type() == at::kByte,
168 "fp8_paged_mqa_logits_cpu: kvcache_fp8 must be a 2D uint8 tensor "
169 "[num_pages, buf_numel_per_page]");
170 TORCH_CHECK(kvcache_fp8.stride(1) == 1 &&
171 kvcache_fp8.stride(0) >= kvcache_fp8.size(1),
172 "fp8_paged_mqa_logits_cpu: kvcache_fp8 must be contiguous "
173 "within a page (may be a page of a larger pooled allocation "
174 "between pages)");
175 TORCH_CHECK(block_size > 0,
176 "fp8_paged_mqa_logits_cpu: block_size must be positive");
177 TORCH_CHECK(
178 kvcache_fp8.size(1) >= block_size * kIndexHeadDimWithScaleBytes,
179 "fp8_paged_mqa_logits_cpu: kvcache_fp8 page byte width too small for "
180 "block_size");
181
182 const int64_t batch_size = q_fp8.size(0);
183 const int64_t num_heads = q_fp8.size(1);
184 TORCH_CHECK(
185 weight.dim() == 2 && weight.scalar_type() == at::kFloat &&
186 weight.size(0) == batch_size && weight.size(1) == num_heads,
187 "fp8_paged_mqa_logits_cpu: weight must be a 2D fp32 tensor [batch, "
188 "heads]");
189 TORCH_CHECK(seq_lens.dim() == 1 && seq_lens.size(0) == batch_size,
190 "fp8_paged_mqa_logits_cpu: seq_lens must be a 1D tensor [batch]");
191 TORCH_CHECK(page_table.dim() == 2 && page_table.size(0) == batch_size,
192 "fp8_paged_mqa_logits_cpu: page_table must be a 2D tensor "
193 "[batch, pages]");
194 TORCH_CHECK(max_seq_len >= 0,
195 "fp8_paged_mqa_logits_cpu: max_seq_len must be non-negative");
196
197 // Positions with token >= seq_len are never written by the kernel below,
198 // and topk_transform_512_cpu (the only consumer) never reads them either
199 // -- left uninitialized, matching upstream.
200 auto logits =
201 at::empty({batch_size, max_seq_len}, q_fp8.options().dtype(at::kFloat));
202 if (batch_size == 0 || max_seq_len == 0) {
203 return logits;
204 }
205
206 if (seq_lens.scalar_type() == at::kInt) {
207 dispatch_page_type<int32_t>(q_fp8, kvcache_fp8, weight, seq_lens,
208 page_table, logits, block_size, max_seq_len);
209 } else if (seq_lens.scalar_type() == at::kLong) {
210 dispatch_page_type<int64_t>(q_fp8, kvcache_fp8, weight, seq_lens,
211 page_table, logits, block_size, max_seq_len);
212 } else {
213 TORCH_CHECK(false, "seq_lens must be int32 or int64");
214 }
215
216 return logits;
217 }