| 1 | // Adapted from |
| 2 | // https://github.com/sgl-project/sglang/blob/main/python/sglang/kernels/aot/csrc/cpu/compressor.cpp |
| 3 | // |
| 4 | // DeepSeek-V4 compressor: softmax-weighted pooling over a sliding window of |
| 5 | // raw KV/score state, RMSNorm, FP8/UE8M0 quant (NoPE) + GPT-J RoPE (RoPE), |
| 6 | // paged-insert into the fp8_ds_mla cache layout (see store_cache.cpp). |
| 7 | // Covers both the head_dim=512 attention-compressor path |
| 8 | // (`compress_norm_rope_store_cpu`) and the head_dim=128 indexer-compressor |
| 9 | // path (`compress_norm_rope_store_indexer_cpu`, distinct cache layout -- |
| 10 | // see its own comment below). |
| 11 | // |
| 12 | // Differences from upstream: |
| 13 | // - Upstream splits decode (`compress_decode_cpu`) from prefill/extend |
| 14 | // (eager `compress_extend_separate`); this port stays unified across |
| 15 | // decode/prefill/extend, one program per token, matching vLLM's own |
| 16 | // GPU/Triton reference (`_fused_kv_compress_norm_rope_insert_sparse_attn`, |
| 17 | // common/ops/fused_compress_quant_cache.py). |
| 18 | // - Window-gather indices arrive pre-resolved to flat state-cache row ids |
| 19 | // from the Python caller (mirroring flash_mla.cpp), not looked up via |
| 20 | // block_table in-kernel. |
| 21 | // - The state cache is pooled with the main KV cache's physical pages |
| 22 | // (CompressorStateCache in compressor.py), so its per-block byte stride |
| 23 | // is page-aligned, not `block_size * row_width` -- row ids are addressed |
| 24 | // as (block_idx, pos_in_block) against the tensor's real block stride, |
| 25 | // never by flattening to a 2D view. |
| 26 | // |
| 27 | // clang-format off |
| 28 | |
| 29 | #include <algorithm> |
| 30 | #include <array> |
| 31 | #include <cmath> |
| 32 | #include <limits> |
| 33 | |
| 34 | #include "common.h" |
| 35 | #include "vec.h" |
| 36 | |
| 37 | #if !defined(CPU_CAPABILITY_AVX512) |
| 38 | static_assert(false, |
| 39 | "compressor.cpp is only ever compiled into the AVX512+AMX-flagged " |
| 40 | "_C target (see cmake/cpu_extension.cmake's VLLM_EXT_SRC_SGL) and " |
| 41 | "relies on that unconditionally -- CPU_CAPABILITY_AVX512 must be " |
| 42 | "defined here."); |
| 43 | #endif |
| 44 | |
| 45 | namespace { |
| 46 | |
| 47 | using fVec = at::vec::Vectorized<float>; |
| 48 | |
| 49 | constexpr int64_t kNopeDim = 448; |
| 50 | constexpr int64_t kRopeDim = 64; |
| 51 | constexpr int64_t kHeadDim = kNopeDim + kRopeDim; |
| 52 | constexpr int64_t kQuantBlock = 64; |
| 53 | constexpr int64_t kNumQuantBlocks = kNopeDim / kQuantBlock; |
| 54 | constexpr int64_t kScaleBytesPerToken = kNumQuantBlocks + 1; |
| 55 | constexpr int64_t kTokenDataBytes = kNopeDim + kRopeDim * 2; |
| 56 | constexpr float kFp8Max = 448.0f; |
| 57 | constexpr float kFp8Min = -448.0f; |
| 58 | constexpr float kQuantEps = 1e-4f; |
| 59 | |
| 60 | // In-place interleaved (GPT-J-style) RoPE: x[2i]/x[2i+1] are one (even, odd) |
| 61 | // pair, cos[i]/sin[i] its angle. Adapted from SGLang's `rope.cpp` |
| 62 | // (`apply_rotary_emb_row`), which avoids gather/scatter (slow on AVX512) by |
| 63 | // broadcasting cos/sin across each pair via in-register permutes instead: |
| 64 | // `_mm512_permute_ps(xv, 0xB1)` swaps the (even, odd) lanes within every |
| 65 | // 2-float pair in one shuffle, and cos/sin -- loaded as 8 contiguous values |
| 66 | // -- are duplicated pairwise into the same layout via `_mm512_permutexvar_ps`. |
| 67 | // SGLang's version reads cos/sin pre-interleaved with x (`[c0,s0,c1,s1,...]`) |
| 68 | // since its freqs tensor is laid out that way; ours are two separate flat |
| 69 | // arrays (`cos[i]`, `sin[i]`, matching store_cache.cpp's convention), hence |
| 70 | // the extra broadcast-duplicate permute here instead of a plain load. |
| 71 | // Both call sites in this file pass half_dim=32 (kRopeDim/2, kIndexerRopeDim/2 |
| 72 | // are both 64/2), an exact multiple of 8, so this file's build (always AVX512, |
| 73 | // see file header) never needs a scalar remainder. |
| 74 | inline void apply_gptj_rope_inplace(float* __restrict__ x, const float* __restrict__ cos, const float* __restrict__ sin, int64_t half_dim) { |
| 75 | // Duplicates each of 8 source lanes into 2 adjacent output lanes: |
| 76 | // [c0,c1,...,c7] -> [c0,c0,c1,c1,...,c7,c7]. |
| 77 | const __m512i dup_idx = _mm512_set_epi32(7, 7, 6, 6, 5, 5, 4, 4, 3, 3, 2, 2, 1, 1, 0, 0); |
| 78 | // Forward GPT-J rotation negates the sin broadcast on the even (real) |
| 79 | // lane of each pair; the odd (imag) lane keeps sin's sign. |
| 80 | const __m512 sign_mask = _mm512_castsi512_ps(_mm512_set_epi32( |
| 81 | 0, (int)0x80000000, 0, (int)0x80000000, 0, (int)0x80000000, 0, (int)0x80000000, |
| 82 | 0, (int)0x80000000, 0, (int)0x80000000, 0, (int)0x80000000, 0, (int)0x80000000)); |
| 83 | for (int64_t i = 0; i + 8 <= half_dim; i += 8) { |
| 84 | __m512 xv = _mm512_loadu_ps(x + 2 * i); |
| 85 | __m512 xv_swapped = _mm512_permute_ps(xv, 0xB1); |
| 86 | __m512 cos_b = _mm512_permutexvar_ps(dup_idx, _mm512_castps256_ps512(_mm256_loadu_ps(cos + i))); |
| 87 | __m512 sin_b = _mm512_permutexvar_ps(dup_idx, _mm512_castps256_ps512(_mm256_loadu_ps(sin + i))); |
| 88 | __m512 sin_signed = _mm512_xor_ps(sin_b, sign_mask); |
| 89 | __m512 out = _mm512_fmadd_ps(xv, cos_b, _mm512_mul_ps(xv_swapped, sin_signed)); |
| 90 | _mm512_storeu_ps(x + 2 * i, out); |
| 91 | } |
| 92 | } |
| 93 | |
| 94 | // AVX512 tile-quantization, adapted verbatim from SGLang's `store_cache.cpp` |
| 95 | // (`quant_to_nope_fp8_rope_bf16_pack_cpu`'s |
| 96 | // `cvt_fp32x16_to_fp8x16`/`quantize_tile_avx512`) -- SGLang itself keeps |
| 97 | // `compress_decode_cpu` fp32-output-only and calls this quant step as a |
| 98 | // separate op; vLLM's compressor kernel stays fused (per this file's header |
| 99 | // comment), so the same AVX512 tile-quant math is inlined here instead. |
| 100 | #if defined(CPU_CAPABILITY_AVX512) |
| 101 | |
| 102 | // Vectorized float32x16 -> fp8_e4m3fn x16 using AVX512 bit manipulation. |
| 103 | // Implements RNE (round-to-nearest-even). Flushes subnormal fp8 to zero. |
| 104 | // fp8_e4m3fn: sign(1), exp(4, bias=7), mant(3), no inf/nan special values. |
| 105 | // Normal: exp_field in [1..15], value = 2^(exp-7) * (1 + mant/8) |
| 106 | // Input assumed clamped to [-448, 448]. |
| 107 | inline __m128i cvt_fp32x16_to_fp8x16(__m512 input) { |
| 108 | __m512i bits = _mm512_castps_si512(input); |
| 109 | |
| 110 | __m512i signs = _mm512_srli_epi32(_mm512_and_si512(bits, _mm512_set1_epi32(0x80000000)), 24); |
| 111 | __m512i abs_bits = _mm512_and_si512(bits, _mm512_set1_epi32(0x7FFFFFFF)); |
| 112 | __m512i f32_exp = _mm512_srli_epi32(abs_bits, 23); |
| 113 | __m512i f32_mant = _mm512_and_si512(abs_bits, _mm512_set1_epi32(0x7FFFFF)); |
| 114 | __m512i fp8_exp = _mm512_sub_epi32(f32_exp, _mm512_set1_epi32(120)); |
| 115 | __m512i mant3 = _mm512_srli_epi32(f32_mant, 20); |
| 116 | |
| 117 | __m512i round_bit = _mm512_and_si512(_mm512_srli_epi32(f32_mant, 19), _mm512_set1_epi32(1)); |
| 118 | __m512i sticky_bits = _mm512_and_si512(f32_mant, _mm512_set1_epi32(0x7FFFF)); |
| 119 | __mmask16 has_sticky = _mm512_cmpneq_epi32_mask(sticky_bits, _mm512_setzero_si512()); |
| 120 | |
| 121 | __m512i lsb = _mm512_and_si512(mant3, _mm512_set1_epi32(1)); |
| 122 | __m512i sticky_or_lsb = _mm512_or_si512(_mm512_maskz_mov_epi32(has_sticky, _mm512_set1_epi32(1)), lsb); |
| 123 | __m512i do_round = _mm512_and_si512(round_bit, sticky_or_lsb); |
| 124 | mant3 = _mm512_add_epi32(mant3, do_round); |
| 125 | |
| 126 | __mmask16 mant_ovf = _mm512_cmpeq_epi32_mask(mant3, _mm512_set1_epi32(8)); |
| 127 | mant3 = _mm512_mask_mov_epi32(mant3, mant_ovf, _mm512_setzero_si512()); |
| 128 | fp8_exp = _mm512_mask_add_epi32(fp8_exp, mant_ovf, fp8_exp, _mm512_set1_epi32(1)); |
| 129 | |
| 130 | fp8_exp = _mm512_min_epi32(fp8_exp, _mm512_set1_epi32(15)); |
| 131 | |
| 132 | __m512i result = _mm512_or_si512(_mm512_slli_epi32(fp8_exp, 3), mant3); |
| 133 | |
| 134 | __mmask16 is_subnormal_or_zero = _mm512_cmplt_epi32_mask(f32_exp, _mm512_set1_epi32(121)); |
| 135 | result = _mm512_mask_mov_epi32(result, is_subnormal_or_zero, _mm512_setzero_si512()); |
| 136 | |
| 137 | result = _mm512_or_si512(result, signs); |
| 138 | |
| 139 | return _mm512_cvtepi32_epi8(result); |
| 140 | } |
| 141 | |
| 142 | // Process one 64-wide tile: find amax, compute a shared power-of-two scale, |
| 143 | // quantize to fp8. Returns the UE8M0 scale byte for the tile. |
| 144 | inline uint8_t quantize_tile_avx512(const at::BFloat16* __restrict__ src, uint8_t* __restrict__ dst) { |
| 145 | __m256i bf16_0 = _mm256_loadu_si256(reinterpret_cast<const __m256i*>(src)); |
| 146 | __m256i bf16_1 = _mm256_loadu_si256(reinterpret_cast<const __m256i*>(src + 16)); |
| 147 | __m256i bf16_2 = _mm256_loadu_si256(reinterpret_cast<const __m256i*>(src + 32)); |
| 148 | __m256i bf16_3 = _mm256_loadu_si256(reinterpret_cast<const __m256i*>(src + 48)); |
| 149 | |
| 150 | __m512 f0 = _mm512_castsi512_ps(_mm512_slli_epi32(_mm512_cvtepu16_epi32(bf16_0), 16)); |
| 151 | __m512 f1 = _mm512_castsi512_ps(_mm512_slli_epi32(_mm512_cvtepu16_epi32(bf16_1), 16)); |
| 152 | __m512 f2 = _mm512_castsi512_ps(_mm512_slli_epi32(_mm512_cvtepu16_epi32(bf16_2), 16)); |
| 153 | __m512 f3 = _mm512_castsi512_ps(_mm512_slli_epi32(_mm512_cvtepu16_epi32(bf16_3), 16)); |
| 154 | |
| 155 | const __m512i abs_mask_i = _mm512_set1_epi32(0x7FFFFFFF); |
| 156 | __m512 abs0 = _mm512_castsi512_ps(_mm512_and_si512(_mm512_castps_si512(f0), abs_mask_i)); |
| 157 | __m512 abs1 = _mm512_castsi512_ps(_mm512_and_si512(_mm512_castps_si512(f1), abs_mask_i)); |
| 158 | __m512 abs2 = _mm512_castsi512_ps(_mm512_and_si512(_mm512_castps_si512(f2), abs_mask_i)); |
| 159 | __m512 abs3 = _mm512_castsi512_ps(_mm512_and_si512(_mm512_castps_si512(f3), abs_mask_i)); |
| 160 | |
| 161 | __m512 max01 = _mm512_max_ps(abs0, abs1); |
| 162 | __m512 max23 = _mm512_max_ps(abs2, abs3); |
| 163 | __m512 max0123 = _mm512_max_ps(max01, max23); |
| 164 | float amax = _mm512_reduce_max_ps(max0123); |
| 165 | |
| 166 | float scale = std::max(amax / kFp8Max, kQuantEps); |
| 167 | float ceil_log2 = std::ceil(std::log2(scale)); |
| 168 | float scale_pow2 = std::exp2(ceil_log2); |
| 169 | float scale_inv = 1.0f / scale_pow2; |
| 170 | |
| 171 | int exponent = static_cast<int>(ceil_log2); |
| 172 | uint8_t scale_uint8 = static_cast<uint8_t>(exponent + 127); |
| 173 | |
| 174 | __m512 vinv = _mm512_set1_ps(scale_inv); |
| 175 | __m512 s0 = _mm512_mul_ps(f0, vinv); |
| 176 | __m512 s1 = _mm512_mul_ps(f1, vinv); |
| 177 | __m512 s2 = _mm512_mul_ps(f2, vinv); |
| 178 | __m512 s3 = _mm512_mul_ps(f3, vinv); |
| 179 | |
| 180 | __m512 vmax = _mm512_set1_ps(kFp8Max); |
| 181 | __m512 vmin = _mm512_set1_ps(kFp8Min); |
| 182 | s0 = _mm512_max_ps(_mm512_min_ps(s0, vmax), vmin); |
| 183 | s1 = _mm512_max_ps(_mm512_min_ps(s1, vmax), vmin); |
| 184 | s2 = _mm512_max_ps(_mm512_min_ps(s2, vmax), vmin); |
| 185 | s3 = _mm512_max_ps(_mm512_min_ps(s3, vmax), vmin); |
| 186 | |
| 187 | __m128i fp8_0 = cvt_fp32x16_to_fp8x16(s0); |
| 188 | __m128i fp8_1 = cvt_fp32x16_to_fp8x16(s1); |
| 189 | __m128i fp8_2 = cvt_fp32x16_to_fp8x16(s2); |
| 190 | __m128i fp8_3 = cvt_fp32x16_to_fp8x16(s3); |
| 191 | |
| 192 | _mm_storeu_si128(reinterpret_cast<__m128i*>(dst), fp8_0); |
| 193 | _mm_storeu_si128(reinterpret_cast<__m128i*>(dst + 16), fp8_1); |
| 194 | _mm_storeu_si128(reinterpret_cast<__m128i*>(dst + 32), fp8_2); |
| 195 | _mm_storeu_si128(reinterpret_cast<__m128i*>(dst + 48), fp8_3); |
| 196 | |
| 197 | // Patch subnormal-magnitude lanes: `cvt_fp32x16_to_fp8x16` flushes |
| 198 | // anything below the smallest fp8_e4m3fn normal (2^-6) to zero, but |
| 199 | // vLLM's dequant path relies on `at::Float8_e4m3fn`'s real subnormal |
| 200 | // encoding -- redo just those rare lanes with the exact scalar |
| 201 | // conversion to stay consistent with the read side. |
| 202 | alignas(64) float scaled[kQuantBlock]; |
| 203 | _mm512_store_ps(scaled, s0); |
| 204 | _mm512_store_ps(scaled + 16, s1); |
| 205 | _mm512_store_ps(scaled + 32, s2); |
| 206 | _mm512_store_ps(scaled + 48, s3); |
| 207 | for (int64_t j = 0; j < kQuantBlock; ++j) { |
| 208 | const float av = std::abs(scaled[j]); |
| 209 | if (av != 0.f && av < 0x1p-6f) { |
| 210 | dst[j] = at::Float8_e4m3fn(scaled[j]).x; |
| 211 | } |
| 212 | } |
| 213 | |
| 214 | return scale_uint8; |
| 215 | } |
| 216 | |
| 217 | #endif // CPU_CAPABILITY_AVX512 |
| 218 | |
| 219 | void save_partial_states_impl( |
| 220 | float* __restrict__ state_cache, // [num_blocks, block_stride] |
| 221 | int64_t block_stride, // element stride between blocks |
| 222 | int64_t block_size, // rows (slots) per block |
| 223 | const float* __restrict__ kv, // [num_tokens, state_width] |
| 224 | int64_t kv_row_stride, // element stride between kv rows |
| 225 | const float* __restrict__ score, // [num_tokens, state_width] |
| 226 | int64_t score_row_stride, // element stride between score rows |
| 227 | const float* __restrict__ ape, // [compress_ratio, state_width] |
| 228 | const int64_t* __restrict__ positions, // [num_tokens] |
| 229 | const int64_t* __restrict__ slot_mapping, // [num_tokens] |
| 230 | int64_t state_width, |
| 231 | int64_t compress_ratio, |
| 232 | int64_t num_tokens) { |
| 233 | const int64_t row_width = 2 * state_width; |
| 234 | at::parallel_for(0, num_tokens, 0, [&](int64_t begin, int64_t end) { |
| 235 | for (int64_t t = begin; t < end; ++t) { |
| 236 | const int64_t slot_id = slot_mapping[t]; |
| 237 | if (slot_id < 0) continue; |
| 238 | |
| 239 | const int64_t block_idx = slot_id / block_size; |
| 240 | const int64_t pos_in_block = slot_id % block_size; |
| 241 | float* const row = state_cache + block_idx * block_stride + pos_in_block * row_width; |
| 242 | const float* const kv_t = kv + t * kv_row_stride; |
| 243 | // state_width is always a multiple of 16 for this model (512/1024 for |
| 244 | // the main compressor, 128/256 for the indexer compressor). |
| 245 | for (int64_t i = 0; i + 16 <= state_width; i += 16) { |
| 246 | fVec::loadu(kv_t + i).store(row + i); |
| 247 | } |
| 248 | |
| 249 | const int64_t position = positions[t]; |
| 250 | const float* const ape_row = ape + (position % compress_ratio) * state_width; |
| 251 | const float* const score_t = score + t * score_row_stride; |
| 252 | float* const score_out = row + state_width; |
| 253 | for (int64_t i = 0; i + 16 <= state_width; i += 16) { |
| 254 | (fVec::loadu(score_t + i) + fVec::loadu(ape_row + i)).store(score_out + i); |
| 255 | } |
| 256 | } |
| 257 | }); |
| 258 | } |
| 259 | |
| 260 | void compress_norm_rope_store_impl( |
| 261 | const float* __restrict__ state_cache, // [num_blocks, state_block_stride] |
| 262 | int64_t state_block_stride, // element stride between state blocks |
| 263 | int64_t state_block_size, // rows (slots) per state block |
| 264 | const int64_t* __restrict__ gather_slots, // [num_tokens, window], -1 = invalid |
| 265 | const int64_t* __restrict__ positions, // [num_tokens] |
| 266 | const int64_t* __restrict__ kv_slot_mapping, // [num_tokens], -1 = skip |
| 267 | const float* __restrict__ rms_norm_weight, // [kHeadDim] |
| 268 | float rms_norm_eps, |
| 269 | const float* __restrict__ cos_sin_cache, // [max_pos, kRopeDim] |
| 270 | uint8_t* __restrict__ kv_cache, // [num_blocks, cache_row_stride] |
| 271 | int64_t cache_row_stride, |
| 272 | int64_t kv_cache_block_size, |
| 273 | int64_t compress_ratio, |
| 274 | int64_t window, |
| 275 | int64_t state_width, |
| 276 | int64_t num_tokens) { |
| 277 | const int64_t rope_half = kRopeDim / 2; |
| 278 | const int64_t state_row_width = 2 * state_width; |
| 279 | |
| 280 | at::parallel_for(0, num_tokens, 0, [&](int64_t begin, int64_t end) { |
| 281 | std::array<float, kHeadDim> max_v; |
| 282 | std::array<float, kHeadDim> denom; |
| 283 | std::array<float, kHeadDim> acc; |
| 284 | std::array<float, kHeadDim> normed; |
| 285 | std::array<at::BFloat16, kNopeDim> nope_bf16; |
| 286 | |
| 287 | for (int64_t t = begin; t < end; ++t) { |
| 288 | const int64_t position = positions[t]; |
| 289 | if ((position + 1) % compress_ratio != 0) continue; |
| 290 | const int64_t kv_slot = kv_slot_mapping[t]; |
| 291 | if (kv_slot < 0) continue; |
| 292 | |
| 293 | std::fill(max_v.begin(), max_v.end(), -std::numeric_limits<float>::infinity()); |
| 294 | std::fill(denom.begin(), denom.end(), 0.f); |
| 295 | std::fill(acc.begin(), acc.end(), 0.f); |
| 296 | |
| 297 | // Per-channel (not per-position) softmax over the window: each of the |
| 298 | // kHeadDim columns is normalized independently across the window axis. |
| 299 | for (int64_t w = 0; w < window; ++w) { |
| 300 | const int64_t slot = gather_slots[t * window + w]; |
| 301 | if (slot < 0) continue; |
| 302 | const int64_t block_idx = slot / state_block_size; |
| 303 | const int64_t pos_in_block = slot % state_block_size; |
| 304 | const int64_t head_offset = (w >= compress_ratio) ? kHeadDim : 0; |
| 305 | const float* const score_row = state_cache + block_idx * state_block_stride + |
| 306 | pos_in_block * state_row_width + head_offset + state_width; |
| 307 | for (int64_t d = 0; d < kHeadDim; d += 16) { |
| 308 | __m512 mx = _mm512_loadu_ps(max_v.data() + d); |
| 309 | __m512 sv = _mm512_loadu_ps(score_row + d); |
| 310 | _mm512_storeu_ps(max_v.data() + d, _mm512_max_ps(mx, sv)); |
| 311 | } |
| 312 | } |
| 313 | for (int64_t w = 0; w < window; ++w) { |
| 314 | const int64_t slot = gather_slots[t * window + w]; |
| 315 | if (slot < 0) continue; |
| 316 | const int64_t block_idx = slot / state_block_size; |
| 317 | const int64_t pos_in_block = slot % state_block_size; |
| 318 | const int64_t head_offset = (w >= compress_ratio) ? kHeadDim : 0; |
| 319 | const float* const row = state_cache + block_idx * state_block_stride + |
| 320 | pos_in_block * state_row_width + head_offset; |
| 321 | const float* const score_row = row + state_width; |
| 322 | for (int64_t d = 0; d < kHeadDim; d += 16) { |
| 323 | fVec e_v = (fVec::loadu(score_row + d) - fVec::loadu(max_v.data() + d)).exp(); |
| 324 | (fVec::loadu(denom.data() + d) + e_v).store(denom.data() + d); |
| 325 | (fVec::loadu(acc.data() + d) + e_v * fVec::loadu(row + d)).store(acc.data() + d); |
| 326 | } |
| 327 | } |
| 328 | |
| 329 | float sq_sum = 0.f; |
| 330 | { |
| 331 | fVec sq_sum_v(0.f); |
| 332 | for (int64_t d = 0; d < kHeadDim; d += 16) { |
| 333 | fVec v = fVec::loadu(acc.data() + d) / fVec::loadu(denom.data() + d); |
| 334 | v.store(normed.data() + d); |
| 335 | sq_sum_v = sq_sum_v + v * v; |
| 336 | } |
| 337 | sq_sum = vec_reduce_sum(sq_sum_v); |
| 338 | } |
| 339 | const float inv_rms = 1.0f / std::sqrt(sq_sum / static_cast<float>(kHeadDim) + rms_norm_eps); |
| 340 | { |
| 341 | fVec inv_rms_v(inv_rms); |
| 342 | for (int64_t d = 0; d < kHeadDim; d += 16) { |
| 343 | (fVec::loadu(normed.data() + d) * inv_rms_v * fVec::loadu(rms_norm_weight + d)).store(normed.data() + d); |
| 344 | } |
| 345 | } |
| 346 | |
| 347 | // bf16 round-trip feeds only the NoPE FP8 quantization (RoPE rotates |
| 348 | // the un-rounded `normed` below, then casts once to bf16 at store |
| 349 | // time for the separate bf16 rope region). |
| 350 | for (int64_t d = 0; d < kNopeDim; d += 16) |
| 351 | store_from_float_ext<at::BFloat16>(nope_bf16.data() + d, fVec::loadu(normed.data() + d)); |
| 352 | |
| 353 | const int64_t compressed_pos = (position / compress_ratio) * compress_ratio; |
| 354 | const float* const cos = cos_sin_cache + compressed_pos * kRopeDim; |
| 355 | const float* const sin = cos + rope_half; |
| 356 | apply_gptj_rope_inplace(normed.data() + kNopeDim, cos, sin, rope_half); |
| 357 | |
| 358 | const int64_t block_idx = kv_slot / kv_cache_block_size; |
| 359 | const int64_t pos_in_block = kv_slot % kv_cache_block_size; |
| 360 | uint8_t* const block_base = kv_cache + block_idx * cache_row_stride; |
| 361 | uint8_t* const token_fp8_ptr = block_base + pos_in_block * kTokenDataBytes; |
| 362 | uint8_t* const token_bf16_ptr = token_fp8_ptr + kNopeDim; |
| 363 | uint8_t* const token_scale_ptr = |
| 364 | block_base + kv_cache_block_size * kTokenDataBytes + pos_in_block * kScaleBytesPerToken; |
| 365 | |
| 366 | // AVX512 tile-quantization (SGLang's `quantize_tile_avx512`), one |
| 367 | // 64-wide tile per iteration. |
| 368 | for (int64_t qb = 0; qb < kNumQuantBlocks; ++qb) { |
| 369 | const at::BFloat16* tile_src = nope_bf16.data() + qb * kQuantBlock; |
| 370 | uint8_t* tile_dst = token_fp8_ptr + qb * kQuantBlock; |
| 371 | token_scale_ptr[qb] = quantize_tile_avx512(tile_src, tile_dst); |
| 372 | } |
| 373 | token_scale_ptr[kNumQuantBlocks] = 0; |
| 374 | |
| 375 | auto* const out_rope = reinterpret_cast<at::BFloat16*>(token_bf16_ptr); |
| 376 | for (int64_t ri = 0; ri < kRopeDim; ri += 16) |
| 377 | store_from_float_ext<at::BFloat16>(out_rope + ri, fVec::loadu(normed.data() + kNopeDim + ri)); |
| 378 | } |
| 379 | }); |
| 380 | } |
| 381 | |
| 382 | // Indexer compressor (head_dim=128): same softmax-pool + RMSNorm + GPT-J |
| 383 | // RoPE prologue as the head=512 path above, but a different cache layout: |
| 384 | // all 128 dims are FP8 in one combined quant block (no separate bf16 rope |
| 385 | // region, single raw fp32 scale per token) instead of 7 UE8M0 blocks over |
| 386 | // nope alone -- matching the paged indexer K-cache layout |
| 387 | // fp8_paged_mqa_logits_cpu/topk_transform_512_cpu (paged_mqa_logits.cpp/ |
| 388 | // topk.cpp) read, and the GPU/Triton reference |
| 389 | // (`_fused_kv_compress_norm_rope_insert_indexer_attn`, |
| 390 | // common/ops/fused_compress_quant_cache.py) exactly. |
| 391 | constexpr int64_t kIndexerHeadDim = 128; |
| 392 | constexpr int64_t kIndexerRopeDim = 64; |
| 393 | constexpr int64_t kIndexerNopeDim = kIndexerHeadDim - kIndexerRopeDim; |
| 394 | constexpr int64_t kIndexerTokenBytes = kIndexerHeadDim; |
| 395 | constexpr int64_t kIndexerScaleBytes = 4; |
| 396 | |
| 397 | void compress_norm_rope_store_indexer_impl( |
| 398 | const float* __restrict__ state_cache, // [num_blocks, state_block_stride] |
| 399 | int64_t state_block_stride, |
| 400 | int64_t state_block_size, |
| 401 | const int64_t* __restrict__ gather_slots, // [num_tokens, window], -1 = invalid |
| 402 | const int64_t* __restrict__ positions, // [num_tokens] |
| 403 | const int64_t* __restrict__ kv_slot_mapping, // [num_tokens], -1 = skip |
| 404 | const float* __restrict__ rms_norm_weight, // [kIndexerHeadDim] |
| 405 | float rms_norm_eps, |
| 406 | const float* __restrict__ cos_sin_cache, // [max_pos, kIndexerRopeDim] |
| 407 | uint8_t* __restrict__ kv_cache, // [num_blocks, cache_row_stride] |
| 408 | int64_t cache_row_stride, |
| 409 | int64_t kv_cache_block_size, |
| 410 | int64_t compress_ratio, |
| 411 | int64_t window, |
| 412 | int64_t state_width, |
| 413 | int64_t num_tokens) { |
| 414 | const int64_t rope_half = kIndexerRopeDim / 2; |
| 415 | const int64_t state_row_width = 2 * state_width; |
| 416 | |
| 417 | at::parallel_for(0, num_tokens, 0, [&](int64_t begin, int64_t end) { |
| 418 | std::array<float, kIndexerHeadDim> max_v; |
| 419 | std::array<float, kIndexerHeadDim> denom; |
| 420 | std::array<float, kIndexerHeadDim> acc; |
| 421 | std::array<float, kIndexerHeadDim> normed; |
| 422 | std::array<at::BFloat16, kIndexerHeadDim> quant_bf16; |
| 423 | |
| 424 | for (int64_t t = begin; t < end; ++t) { |
| 425 | const int64_t position = positions[t]; |
| 426 | if ((position + 1) % compress_ratio != 0) continue; |
| 427 | const int64_t kv_slot = kv_slot_mapping[t]; |
| 428 | if (kv_slot < 0) continue; |
| 429 | |
| 430 | std::fill(max_v.begin(), max_v.end(), -std::numeric_limits<float>::infinity()); |
| 431 | std::fill(denom.begin(), denom.end(), 0.f); |
| 432 | std::fill(acc.begin(), acc.end(), 0.f); |
| 433 | |
| 434 | for (int64_t w = 0; w < window; ++w) { |
| 435 | const int64_t slot = gather_slots[t * window + w]; |
| 436 | if (slot < 0) continue; |
| 437 | const int64_t block_idx = slot / state_block_size; |
| 438 | const int64_t pos_in_block = slot % state_block_size; |
| 439 | const int64_t head_offset = (w >= compress_ratio) ? kIndexerHeadDim : 0; |
| 440 | const float* const score_row = state_cache + block_idx * state_block_stride + |
| 441 | pos_in_block * state_row_width + head_offset + state_width; |
| 442 | for (int64_t d = 0; d < kIndexerHeadDim; d += 16) { |
| 443 | __m512 mx = _mm512_loadu_ps(max_v.data() + d); |
| 444 | __m512 sv = _mm512_loadu_ps(score_row + d); |
| 445 | _mm512_storeu_ps(max_v.data() + d, _mm512_max_ps(mx, sv)); |
| 446 | } |
| 447 | } |
| 448 | for (int64_t w = 0; w < window; ++w) { |
| 449 | const int64_t slot = gather_slots[t * window + w]; |
| 450 | if (slot < 0) continue; |
| 451 | const int64_t block_idx = slot / state_block_size; |
| 452 | const int64_t pos_in_block = slot % state_block_size; |
| 453 | const int64_t head_offset = (w >= compress_ratio) ? kIndexerHeadDim : 0; |
| 454 | const float* const row = state_cache + block_idx * state_block_stride + |
| 455 | pos_in_block * state_row_width + head_offset; |
| 456 | const float* const score_row = row + state_width; |
| 457 | for (int64_t d = 0; d < kIndexerHeadDim; d += 16) { |
| 458 | fVec e_v = (fVec::loadu(score_row + d) - fVec::loadu(max_v.data() + d)).exp(); |
| 459 | (fVec::loadu(denom.data() + d) + e_v).store(denom.data() + d); |
| 460 | (fVec::loadu(acc.data() + d) + e_v * fVec::loadu(row + d)).store(acc.data() + d); |
| 461 | } |
| 462 | } |
| 463 | |
| 464 | float sq_sum = 0.f; |
| 465 | { |
| 466 | fVec sq_sum_v(0.f); |
| 467 | for (int64_t d = 0; d < kIndexerHeadDim; d += 16) { |
| 468 | fVec v = fVec::loadu(acc.data() + d) / fVec::loadu(denom.data() + d); |
| 469 | v.store(normed.data() + d); |
| 470 | sq_sum_v = sq_sum_v + v * v; |
| 471 | } |
| 472 | sq_sum = vec_reduce_sum(sq_sum_v); |
| 473 | } |
| 474 | const float inv_rms = 1.0f / std::sqrt(sq_sum / static_cast<float>(kIndexerHeadDim) + rms_norm_eps); |
| 475 | { |
| 476 | fVec inv_rms_v(inv_rms); |
| 477 | for (int64_t d = 0; d < kIndexerHeadDim; d += 16) { |
| 478 | (fVec::loadu(normed.data() + d) * inv_rms_v * fVec::loadu(rms_norm_weight + d)).store(normed.data() + d); |
| 479 | } |
| 480 | } |
| 481 | |
| 482 | const int64_t compressed_pos = (position / compress_ratio) * compress_ratio; |
| 483 | const float* const cos = cos_sin_cache + compressed_pos * kIndexerRopeDim; |
| 484 | const float* const sin = cos + rope_half; |
| 485 | apply_gptj_rope_inplace(normed.data() + kIndexerNopeDim, cos, sin, rope_half); |
| 486 | |
| 487 | // Indexer path rounds the FULL post-RoPE vector through bf16 before |
| 488 | // quantizing (single combined block); the head=512 path above instead |
| 489 | // rounds only the NoPE region before RoPE (separate bf16 rope store). |
| 490 | for (int64_t d = 0; d < kIndexerHeadDim; d += 16) |
| 491 | store_from_float_ext<at::BFloat16>(quant_bf16.data() + d, fVec::loadu(normed.data() + d)); |
| 492 | |
| 493 | const int64_t block_idx = kv_slot / kv_cache_block_size; |
| 494 | const int64_t pos_in_block = kv_slot % kv_cache_block_size; |
| 495 | uint8_t* const block_base = kv_cache + block_idx * cache_row_stride; |
| 496 | uint8_t* const token_fp8_ptr = block_base + pos_in_block * kIndexerTokenBytes; |
| 497 | uint8_t* const token_scale_ptr = |
| 498 | block_base + kv_cache_block_size * kIndexerTokenBytes + pos_in_block * kIndexerScaleBytes; |
| 499 | |
| 500 | float absmax = 1e-4f; |
| 501 | { |
| 502 | __m512 max_v_reg = _mm512_set1_ps(absmax); |
| 503 | for (int64_t di = 0; di < kIndexerHeadDim; di += 16) { |
| 504 | __m256i bf16_v = _mm256_loadu_si256(reinterpret_cast<const __m256i*>(quant_bf16.data() + di)); |
| 505 | __m512 f = _mm512_castsi512_ps(_mm512_slli_epi32(_mm512_cvtepu16_epi32(bf16_v), 16)); |
| 506 | __m512 av = _mm512_castsi512_ps(_mm512_and_si512(_mm512_castps_si512(f), _mm512_set1_epi32(0x7FFFFFFF))); |
| 507 | max_v_reg = _mm512_max_ps(max_v_reg, av); |
| 508 | } |
| 509 | absmax = _mm512_reduce_max_ps(max_v_reg); |
| 510 | } |
| 511 | const float exponent = std::ceil(std::log2(absmax / kFp8Max)); |
| 512 | const float scale = std::exp2(exponent); |
| 513 | const float inv_scale = 1.0f / scale; |
| 514 | |
| 515 | // AVX512 quantize+store, reusing SGLang's `cvt_fp32x16_to_fp8x16` bit- |
| 516 | // manipulation conversion (this single-shared-scale 128-wide block has |
| 517 | // no direct SGLang equivalent -- see file header -- so only the fp8 |
| 518 | // conversion primitive is ported, the amax/scale formula is unchanged |
| 519 | // from vLLM's own reference). |
| 520 | uint8_t* const out_fp8 = token_fp8_ptr; |
| 521 | { |
| 522 | const __m512 vinv = _mm512_set1_ps(inv_scale); |
| 523 | const __m512 vmax = _mm512_set1_ps(kFp8Max); |
| 524 | const __m512 vmin = _mm512_set1_ps(kFp8Min); |
| 525 | alignas(64) float scaled_buf[kIndexerHeadDim]; |
| 526 | for (int64_t d = 0; d < kIndexerHeadDim; d += 16) { |
| 527 | __m256i bf16_v = _mm256_loadu_si256(reinterpret_cast<const __m256i*>(quant_bf16.data() + d)); |
| 528 | __m512 f = _mm512_castsi512_ps(_mm512_slli_epi32(_mm512_cvtepu16_epi32(bf16_v), 16)); |
| 529 | __m512 s = _mm512_mul_ps(f, vinv); |
| 530 | s = _mm512_max_ps(_mm512_min_ps(s, vmax), vmin); |
| 531 | _mm512_store_ps(scaled_buf + d, s); |
| 532 | __m128i fp8_v = cvt_fp32x16_to_fp8x16(s); |
| 533 | _mm_storeu_si128(reinterpret_cast<__m128i*>(out_fp8 + d), fp8_v); |
| 534 | } |
| 535 | // Patch subnormal-magnitude lanes (see `quantize_tile_avx512`'s |
| 536 | // comment above for why the vectorized conversion's flush-to-zero |
| 537 | // must not be trusted here). |
| 538 | for (int64_t d = 0; d < kIndexerHeadDim; ++d) { |
| 539 | const float av = std::abs(scaled_buf[d]); |
| 540 | if (av != 0.f && av < 0x1p-6f) { |
| 541 | out_fp8[d] = at::Float8_e4m3fn(scaled_buf[d]).x; |
| 542 | } |
| 543 | } |
| 544 | } |
| 545 | // Raw fp32 scale (not UE8M0-encoded) -- single quant block, matches |
| 546 | // fp8_paged_mqa_logits_cpu's (paged_mqa_logits.cpp) read-side expectations. |
| 547 | *reinterpret_cast<float*>(token_scale_ptr) = scale; |
| 548 | } |
| 549 | }); |
| 550 | } |
| 551 | |
| 552 | } // anonymous namespace |
| 553 | |
| 554 | void save_partial_states_cpu( |
| 555 | at::Tensor& kv, |
| 556 | at::Tensor& score, |
| 557 | at::Tensor& ape, |
| 558 | at::Tensor& positions, |
| 559 | at::Tensor& state_cache, |
| 560 | at::Tensor& slot_mapping) { |
| 561 | TORCH_CHECK( |
| 562 | kv.dim() == 2 && kv.scalar_type() == at::kFloat, "save_partial_states_cpu: kv must be a 2D fp32 tensor"); |
| 563 | TORCH_CHECK(score.sizes() == kv.sizes() && score.scalar_type() == at::kFloat, |
| 564 | "save_partial_states_cpu: score must match kv's shape/dtype"); |
| 565 | TORCH_CHECK(ape.dim() == 2 && ape.size(1) == kv.size(1) && ape.scalar_type() == at::kFloat, |
| 566 | "save_partial_states_cpu: ape must be [compress_ratio, state_width] fp32"); |
| 567 | TORCH_CHECK(positions.scalar_type() == at::kLong && positions.numel() == kv.size(0), |
| 568 | "save_partial_states_cpu: positions must be int64 with num_tokens entries"); |
| 569 | TORCH_CHECK(slot_mapping.scalar_type() == at::kLong && slot_mapping.numel() == kv.size(0), |
| 570 | "save_partial_states_cpu: slot_mapping must be int64 with num_tokens entries"); |
| 571 | TORCH_CHECK(state_cache.dim() == 3 && state_cache.size(2) == 2 * kv.size(1) && |
| 572 | state_cache.scalar_type() == at::kFloat, |
| 573 | "save_partial_states_cpu: state_cache must be [num_blocks, block_size, " |
| 574 | "2*state_width] fp32"); |
| 575 | TORCH_CHECK(state_cache.stride(2) == 1 && state_cache.stride(1) == state_cache.size(2), |
| 576 | "save_partial_states_cpu: state_cache rows must be contiguous within a block"); |
| 577 | TORCH_CHECK( |
| 578 | kv.device().is_cpu() && score.device().is_cpu() && ape.device().is_cpu() && positions.device().is_cpu() && |
| 579 | state_cache.device().is_cpu() && slot_mapping.device().is_cpu(), |
| 580 | "save_partial_states_cpu: all inputs must be CPU tensors"); |
| 581 | |
| 582 | const int64_t num_tokens = kv.size(0); |
| 583 | const int64_t state_width = kv.size(1); |
| 584 | const int64_t compress_ratio = ape.size(0); |
| 585 | const int64_t block_size = state_cache.size(1); |
| 586 | |
| 587 | TORCH_CHECK( |
| 588 | kv.stride(1) == 1 && score.stride(1) == 1 && ape.is_contiguous() && positions.is_contiguous() && |
| 589 | slot_mapping.is_contiguous(), |
| 590 | "save_partial_states_cpu: kv/score must be contiguous in their last dim " |
| 591 | "(row stride may exceed state_width -- e.g. a column-split view of a " |
| 592 | "fused kv_score projection); ape/positions/slot_mapping must be fully " |
| 593 | "contiguous"); |
| 594 | |
| 595 | save_partial_states_impl( |
| 596 | state_cache.data_ptr<float>(), |
| 597 | state_cache.stride(0), |
| 598 | block_size, |
| 599 | kv.data_ptr<float>(), |
| 600 | kv.stride(0), |
| 601 | score.data_ptr<float>(), |
| 602 | score.stride(0), |
| 603 | ape.data_ptr<float>(), |
| 604 | positions.data_ptr<int64_t>(), |
| 605 | slot_mapping.data_ptr<int64_t>(), |
| 606 | state_width, |
| 607 | compress_ratio, |
| 608 | num_tokens); |
| 609 | } |
| 610 | |
| 611 | void compress_norm_rope_store_cpu( |
| 612 | at::Tensor& state_cache, |
| 613 | at::Tensor& gather_slots, |
| 614 | at::Tensor& positions, |
| 615 | at::Tensor& kv_slot_mapping, |
| 616 | at::Tensor& rms_norm_weight, |
| 617 | double rms_norm_eps, |
| 618 | at::Tensor& cos_sin_cache, |
| 619 | at::Tensor& kv_cache_2d, |
| 620 | int64_t kv_cache_block_size, |
| 621 | int64_t compress_ratio) { |
| 622 | TORCH_CHECK(state_cache.dim() == 3 && state_cache.scalar_type() == at::kFloat, |
| 623 | "compress_norm_rope_store_cpu: state_cache must be a 3D fp32 tensor"); |
| 624 | TORCH_CHECK(state_cache.stride(2) == 1 && state_cache.stride(1) == state_cache.size(2), |
| 625 | "compress_norm_rope_store_cpu: state_cache rows must be contiguous within a block"); |
| 626 | TORCH_CHECK(gather_slots.dim() == 2 && gather_slots.scalar_type() == at::kLong, |
| 627 | "compress_norm_rope_store_cpu: gather_slots must be a 2D int64 tensor"); |
| 628 | const int64_t num_tokens = gather_slots.size(0); |
| 629 | TORCH_CHECK(positions.scalar_type() == at::kLong && positions.numel() == num_tokens, |
| 630 | "compress_norm_rope_store_cpu: positions must be int64 with num_tokens entries"); |
| 631 | TORCH_CHECK(kv_slot_mapping.scalar_type() == at::kLong && kv_slot_mapping.numel() == num_tokens, |
| 632 | "compress_norm_rope_store_cpu: kv_slot_mapping must be int64 with num_tokens entries"); |
| 633 | TORCH_CHECK(rms_norm_weight.numel() == kHeadDim, |
| 634 | "compress_norm_rope_store_cpu: rms_norm_weight must have head_dim (512) entries"); |
| 635 | TORCH_CHECK(kv_cache_2d.dim() == 2 && kv_cache_2d.scalar_type() == at::kByte, |
| 636 | "compress_norm_rope_store_cpu: kv_cache_2d must be a 2D uint8 tensor"); |
| 637 | TORCH_CHECK(state_cache.size(2) % 2 == 0, "compress_norm_rope_store_cpu: state_cache width must be even"); |
| 638 | TORCH_CHECK( |
| 639 | rms_norm_weight.scalar_type() == at::kFloat && rms_norm_weight.is_contiguous() && |
| 640 | cos_sin_cache.scalar_type() == at::kFloat && cos_sin_cache.is_contiguous(), |
| 641 | "compress_norm_rope_store_cpu: rms_norm_weight/cos_sin_cache must be " |
| 642 | "float32 and contiguous (convert once in Python)"); |
| 643 | TORCH_CHECK( |
| 644 | gather_slots.is_contiguous() && positions.is_contiguous() && kv_slot_mapping.is_contiguous(), |
| 645 | "compress_norm_rope_store_cpu: gather_slots/positions/kv_slot_mapping must be contiguous"); |
| 646 | TORCH_CHECK( |
| 647 | state_cache.device().is_cpu() && gather_slots.device().is_cpu() && positions.device().is_cpu() && |
| 648 | kv_slot_mapping.device().is_cpu() && rms_norm_weight.device().is_cpu() && |
| 649 | cos_sin_cache.device().is_cpu() && kv_cache_2d.device().is_cpu(), |
| 650 | "compress_norm_rope_store_cpu: all inputs must be CPU tensors"); |
| 651 | |
| 652 | const int64_t window = gather_slots.size(1); |
| 653 | const int64_t state_width = state_cache.size(2) / 2; |
| 654 | const int64_t state_block_size = state_cache.size(1); |
| 655 | |
| 656 | // Row strides are threaded through explicitly rather than assumed, since |
| 657 | // both buffers may be pages of a larger pooled allocation (matches |
| 658 | // store_cache.cpp/flash_mla.cpp's precedent). |
| 659 | const int64_t state_block_stride = state_cache.stride(0); |
| 660 | const int64_t cache_row_stride = kv_cache_2d.stride(0); |
| 661 | |
| 662 | compress_norm_rope_store_impl( |
| 663 | state_cache.data_ptr<float>(), |
| 664 | state_block_stride, |
| 665 | state_block_size, |
| 666 | gather_slots.data_ptr<int64_t>(), |
| 667 | positions.data_ptr<int64_t>(), |
| 668 | kv_slot_mapping.data_ptr<int64_t>(), |
| 669 | rms_norm_weight.data_ptr<float>(), |
| 670 | static_cast<float>(rms_norm_eps), |
| 671 | cos_sin_cache.data_ptr<float>(), |
| 672 | kv_cache_2d.data_ptr<uint8_t>(), |
| 673 | cache_row_stride, |
| 674 | kv_cache_block_size, |
| 675 | compress_ratio, |
| 676 | window, |
| 677 | state_width, |
| 678 | num_tokens); |
| 679 | } |
| 680 | |
| 681 | void compress_norm_rope_store_indexer_cpu( |
| 682 | at::Tensor& state_cache, |
| 683 | at::Tensor& gather_slots, |
| 684 | at::Tensor& positions, |
| 685 | at::Tensor& kv_slot_mapping, |
| 686 | at::Tensor& rms_norm_weight, |
| 687 | double rms_norm_eps, |
| 688 | at::Tensor& cos_sin_cache, |
| 689 | at::Tensor& kv_cache_2d, |
| 690 | int64_t kv_cache_block_size, |
| 691 | int64_t compress_ratio) { |
| 692 | TORCH_CHECK(state_cache.dim() == 3 && state_cache.scalar_type() == at::kFloat, |
| 693 | "compress_norm_rope_store_indexer_cpu: state_cache must be a 3D fp32 tensor"); |
| 694 | TORCH_CHECK(state_cache.stride(2) == 1 && state_cache.stride(1) == state_cache.size(2), |
| 695 | "compress_norm_rope_store_indexer_cpu: state_cache rows must be contiguous within a block"); |
| 696 | TORCH_CHECK(gather_slots.dim() == 2 && gather_slots.scalar_type() == at::kLong, |
| 697 | "compress_norm_rope_store_indexer_cpu: gather_slots must be a 2D int64 tensor"); |
| 698 | const int64_t num_tokens = gather_slots.size(0); |
| 699 | TORCH_CHECK(positions.scalar_type() == at::kLong && positions.numel() == num_tokens, |
| 700 | "compress_norm_rope_store_indexer_cpu: positions must be int64 with num_tokens entries"); |
| 701 | TORCH_CHECK(kv_slot_mapping.scalar_type() == at::kLong && kv_slot_mapping.numel() == num_tokens, |
| 702 | "compress_norm_rope_store_indexer_cpu: kv_slot_mapping must be int64 with num_tokens entries"); |
| 703 | TORCH_CHECK(rms_norm_weight.numel() == kIndexerHeadDim, |
| 704 | "compress_norm_rope_store_indexer_cpu: rms_norm_weight must have head_dim (128) entries"); |
| 705 | TORCH_CHECK(kv_cache_2d.dim() == 2 && kv_cache_2d.scalar_type() == at::kByte, |
| 706 | "compress_norm_rope_store_indexer_cpu: kv_cache_2d must be a 2D uint8 tensor"); |
| 707 | TORCH_CHECK(state_cache.size(2) % 2 == 0, |
| 708 | "compress_norm_rope_store_indexer_cpu: state_cache width must be even"); |
| 709 | TORCH_CHECK( |
| 710 | rms_norm_weight.scalar_type() == at::kFloat && rms_norm_weight.is_contiguous() && |
| 711 | cos_sin_cache.scalar_type() == at::kFloat && cos_sin_cache.is_contiguous(), |
| 712 | "compress_norm_rope_store_indexer_cpu: rms_norm_weight/cos_sin_cache must " |
| 713 | "be float32 and contiguous (convert once in Python)"); |
| 714 | TORCH_CHECK( |
| 715 | gather_slots.is_contiguous() && positions.is_contiguous() && kv_slot_mapping.is_contiguous(), |
| 716 | "compress_norm_rope_store_indexer_cpu: gather_slots/positions/kv_slot_mapping must be contiguous"); |
| 717 | TORCH_CHECK( |
| 718 | state_cache.device().is_cpu() && gather_slots.device().is_cpu() && positions.device().is_cpu() && |
| 719 | kv_slot_mapping.device().is_cpu() && rms_norm_weight.device().is_cpu() && |
| 720 | cos_sin_cache.device().is_cpu() && kv_cache_2d.device().is_cpu(), |
| 721 | "compress_norm_rope_store_indexer_cpu: all inputs must be CPU tensors"); |
| 722 | |
| 723 | const int64_t window = gather_slots.size(1); |
| 724 | const int64_t state_width = state_cache.size(2) / 2; |
| 725 | const int64_t state_block_size = state_cache.size(1); |
| 726 | |
| 727 | const int64_t state_block_stride = state_cache.stride(0); |
| 728 | const int64_t cache_row_stride = kv_cache_2d.stride(0); |
| 729 | |
| 730 | compress_norm_rope_store_indexer_impl( |
| 731 | state_cache.data_ptr<float>(), |
| 732 | state_block_stride, |
| 733 | state_block_size, |
| 734 | gather_slots.data_ptr<int64_t>(), |
| 735 | positions.data_ptr<int64_t>(), |
| 736 | kv_slot_mapping.data_ptr<int64_t>(), |
| 737 | rms_norm_weight.data_ptr<float>(), |
| 738 | static_cast<float>(rms_norm_eps), |
| 739 | cos_sin_cache.data_ptr<float>(), |
| 740 | kv_cache_2d.data_ptr<uint8_t>(), |
| 741 | cache_row_stride, |
| 742 | kv_cache_block_size, |
| 743 | compress_ratio, |
| 744 | window, |
| 745 | state_width, |
| 746 | num_tokens); |
| 747 | } |