| 1 | // vLLM-native: DeepSeek-V4 sparse indexer Q-side RoPE + FP8 quant kernel. |
| 2 | // Ports the math of vLLM's own GPU/Triton reference |
| 3 | // (`_fused_indexer_q_rope_quant_kernel`, common/ops/fused_indexer_q.py) -- |
| 4 | // not a port of SGLang's own indexer-Q kernel, since that applies RoPE plus |
| 5 | // a 128-point Hadamard rotation the real DeepSeek-V4-Flash checkpoint |
| 6 | // doesn't use (verified: it applies RoPE only, no Hadamard step). |
| 7 | // |
| 8 | // Per (token, head): GPT-J interleaved RoPE on the last `rot_dim` dims; |
| 9 | // leading NoPE dims pass through unrotated. RoPE-rotated values are bf16- |
| 10 | // round-tripped before the absmax reduction (matching the K-side compressor |
| 11 | // kernel), NoPE values are not. One UE8M0 scale per (token, head) covers the |
| 12 | // full head_dim; since the FP8 Q tensor has no companion scale tensor, the |
| 13 | // scale is folded into `index_weights_out` instead: |
| 14 | // weights_out = weights * q_scale * softmax_scale * head_scale |
| 15 | // MXFP4 (`use_fp4=True`) has per-block scales that can't be folded this way |
| 16 | // and stays on triton-cpu -- out of scope here. |
| 17 | // |
| 18 | // AVX512 path: NoPE copy+absmax and the FP8 scale/clamp use plain 16-wide |
| 19 | // vector ops (final FP8 cast stays scalar, matching SGLang's own |
| 20 | // `act_quant_cpu_impl`). The RoPE rotation vectorizes the interleaved |
| 21 | // (x_even, x_odd) pairs via in-register permutes (`_mm512_permute_ps` / |
| 22 | // `_mm512_permutexvar_ps`), the same technique compressor.cpp's |
| 23 | // `apply_gptj_rope_inplace` and store_cache.cpp's |
| 24 | // `inverse_gptj_rope_o_proj_cpu` use, since gather/scatter is comparatively |
| 25 | // expensive on AVX512. |
| 26 | |
| 27 | #include <algorithm> |
| 28 | #include <cmath> |
| 29 | #include <vector> |
| 30 | |
| 31 | #include "common.h" |
| 32 | #include "vec.h" |
| 33 | |
| 34 | #if !defined(CPU_CAPABILITY_AVX512) |
| 35 | static_assert( |
| 36 | false, |
| 37 | "indexer.cpp is only ever compiled into the AVX512+AMX-flagged " |
| 38 | "_C target (see cmake/cpu_extension.cmake's VLLM_EXT_SRC_SGL) and " |
| 39 | "relies on that unconditionally -- CPU_CAPABILITY_AVX512 must be " |
| 40 | "defined here."); |
| 41 | #endif |
| 42 | |
| 43 | namespace { |
| 44 | |
| 45 | constexpr float kFp8Max = 448.0f; |
| 46 | |
| 47 | #if defined(CPU_CAPABILITY_AVX512) |
| 48 | |
| 49 | // fp32 -> bf16 -> fp32 round trip via the AVX512-BF16 conversion |
| 50 | // instruction, numerically matching `static_cast<float>(at::BFloat16(x))` |
| 51 | // (round-to-nearest-even). |
| 52 | inline __m512 bf16_round_trip_avx512(__m512 x) { |
| 53 | return CVT_BF16_TO_FP32((__m256i)(_mm512_cvtneps_pbh(x))); |
| 54 | } |
| 55 | |
| 56 | inline __m512 abs_ps(__m512 x) { |
| 57 | return _mm512_andnot_ps(_mm512_set1_ps(-0.0f), x); |
| 58 | } |
| 59 | |
| 60 | #endif // CPU_CAPABILITY_AVX512 |
| 61 | |
| 62 | void fused_indexer_q_rope_quant_impl( |
| 63 | const int64_t* __restrict__ positions, // [num_tokens] |
| 64 | const float* __restrict__ index_q, // [num_tokens, num_heads, head_dim] |
| 65 | int64_t index_q_stride0, int64_t index_q_stride1, |
| 66 | const float* __restrict__ cos_sin_cache, // [max_pos, rot_dim] |
| 67 | int64_t cos_sin_stride, int64_t half_rot_dim, |
| 68 | at::Float8_e4m3fn* __restrict__ index_q_fp8, // [num_tokens, num_heads, |
| 69 | // head_dim] |
| 70 | int64_t index_q_fp8_stride0, int64_t index_q_fp8_stride1, int64_t head_dim, |
| 71 | const float* __restrict__ index_weights, // [num_tokens, num_heads] |
| 72 | int64_t index_weights_stride, float softmax_scale, float head_scale, |
| 73 | float* __restrict__ index_weights_out, // [num_tokens, num_heads] |
| 74 | int64_t index_weights_out_stride, int64_t num_tokens, int64_t num_heads) { |
| 75 | const int64_t rot_dim = 2 * half_rot_dim; |
| 76 | const int64_t nope_dim = head_dim - rot_dim; |
| 77 | |
| 78 | at::parallel_for( |
| 79 | 0, num_tokens * num_heads, 0, [&](int64_t begin, int64_t end) { |
| 80 | std::vector<float> buf(head_dim); |
| 81 | for (int64_t idx = begin; idx < end; ++idx) { |
| 82 | const int64_t t = idx / num_heads; |
| 83 | const int64_t h = idx % num_heads; |
| 84 | |
| 85 | const int64_t pos = positions[t]; |
| 86 | const float* const cos = cos_sin_cache + pos * cos_sin_stride; |
| 87 | const float* const sin = cos + half_rot_dim; |
| 88 | |
| 89 | const float* const q = |
| 90 | index_q + t * index_q_stride0 + h * index_q_stride1; |
| 91 | |
| 92 | float amax = 1e-4f; |
| 93 | #if defined(CPU_CAPABILITY_AVX512) |
| 94 | { |
| 95 | __m512 vamax = _mm512_set1_ps(amax); |
| 96 | for (int64_t d = 0; d < nope_dim; d += 16) { |
| 97 | __m512 v = _mm512_loadu_ps(q + d); |
| 98 | _mm512_storeu_ps(&buf[d], v); |
| 99 | vamax = _mm512_max_ps(vamax, abs_ps(v)); |
| 100 | } |
| 101 | amax = _mm512_reduce_max_ps(vamax); |
| 102 | } |
| 103 | #endif |
| 104 | |
| 105 | const float* const q_rot = q + nope_dim; |
| 106 | #if defined(CPU_CAPABILITY_AVX512) |
| 107 | { |
| 108 | // Permute-based (no gather/scatter), same technique as |
| 109 | // compressor.cpp's apply_gptj_rope_inplace / store_cache.cpp's |
| 110 | // inverse_gptj_rope_o_proj_cpu: `_mm512_permute_ps` swaps each |
| 111 | // pair's (even, odd) lanes in one shuffle, and cos/sin -- loaded as |
| 112 | // 8 contiguous values -- are duplicated pairwise via |
| 113 | // `_mm512_permutexvar_ps` to align with that layout. Two batches of |
| 114 | // 8 pairs per iteration (16 total, matching the old gather step). |
| 115 | // Unlike the RoPE kernels above, the result here stays in its |
| 116 | // original interleaved layout, so no un-permute/scatter is needed |
| 117 | // before the bf16 round-trip and store. |
| 118 | const __m512i dup_idx = _mm512_set_epi32(7, 7, 6, 6, 5, 5, 4, 4, 3, |
| 119 | 3, 2, 2, 1, 1, 0, 0); |
| 120 | // Forward GPT-J rotation negates the sin broadcast on the even |
| 121 | // (real) lane of each pair; the odd (imag) lane keeps sin's sign. |
| 122 | const __m512 sign_mask = _mm512_castsi512_ps(_mm512_set_epi32( |
| 123 | 0, (int)0x80000000, 0, (int)0x80000000, 0, (int)0x80000000, 0, |
| 124 | (int)0x80000000, 0, (int)0x80000000, 0, (int)0x80000000, 0, |
| 125 | (int)0x80000000, 0, (int)0x80000000)); |
| 126 | __m512 vamax_rope = _mm512_setzero_ps(); |
| 127 | for (int64_t i = 0; i < half_rot_dim; i += 16) { |
| 128 | const float* const base = q_rot + 2 * i; |
| 129 | float* const dst = buf.data() + nope_dim + 2 * i; |
| 130 | |
| 131 | __m512 xv0 = _mm512_loadu_ps(base); |
| 132 | __m512 cos_b0 = _mm512_permutexvar_ps( |
| 133 | dup_idx, _mm512_castps256_ps512(_mm256_loadu_ps(cos + i))); |
| 134 | __m512 sin_b0 = _mm512_permutexvar_ps( |
| 135 | dup_idx, _mm512_castps256_ps512(_mm256_loadu_ps(sin + i))); |
| 136 | __m512 out0 = _mm512_fmadd_ps( |
| 137 | xv0, cos_b0, |
| 138 | _mm512_mul_ps(_mm512_permute_ps(xv0, 0xB1), |
| 139 | _mm512_xor_ps(sin_b0, sign_mask))); |
| 140 | // Match reference numerics: fp32 -> bf16 -> fp32 before the |
| 141 | // absmax. |
| 142 | out0 = bf16_round_trip_avx512(out0); |
| 143 | _mm512_storeu_ps(dst, out0); |
| 144 | vamax_rope = _mm512_max_ps(vamax_rope, abs_ps(out0)); |
| 145 | |
| 146 | __m512 xv1 = _mm512_loadu_ps(base + 16); |
| 147 | __m512 cos_b1 = _mm512_permutexvar_ps( |
| 148 | dup_idx, |
| 149 | _mm512_castps256_ps512(_mm256_loadu_ps(cos + i + 8))); |
| 150 | __m512 sin_b1 = _mm512_permutexvar_ps( |
| 151 | dup_idx, |
| 152 | _mm512_castps256_ps512(_mm256_loadu_ps(sin + i + 8))); |
| 153 | __m512 out1 = _mm512_fmadd_ps( |
| 154 | xv1, cos_b1, |
| 155 | _mm512_mul_ps(_mm512_permute_ps(xv1, 0xB1), |
| 156 | _mm512_xor_ps(sin_b1, sign_mask))); |
| 157 | out1 = bf16_round_trip_avx512(out1); |
| 158 | _mm512_storeu_ps(dst + 16, out1); |
| 159 | vamax_rope = _mm512_max_ps(vamax_rope, abs_ps(out1)); |
| 160 | } |
| 161 | amax = std::max(amax, _mm512_reduce_max_ps(vamax_rope)); |
| 162 | } |
| 163 | #endif |
| 164 | |
| 165 | const float exponent = std::ceil(std::log2(amax / kFp8Max)); |
| 166 | const float scale = std::exp2(exponent); |
| 167 | const float inv_scale = 1.0f / scale; |
| 168 | |
| 169 | at::Float8_e4m3fn* const out = |
| 170 | index_q_fp8 + t * index_q_fp8_stride0 + h * index_q_fp8_stride1; |
| 171 | #if defined(CPU_CAPABILITY_AVX512) |
| 172 | { |
| 173 | const __m512 vinv = _mm512_set1_ps(inv_scale); |
| 174 | const __m512 vmax = _mm512_set1_ps(kFp8Max); |
| 175 | const __m512 vmin = _mm512_set1_ps(-kFp8Max); |
| 176 | alignas(64) float scaled_buf[16]; |
| 177 | for (int64_t d2 = 0; d2 < head_dim; d2 += 16) { |
| 178 | __m512 v = _mm512_loadu_ps(&buf[d2]); |
| 179 | v = _mm512_min_ps(_mm512_max_ps(_mm512_mul_ps(v, vinv), vmin), |
| 180 | vmax); |
| 181 | _mm512_store_ps(scaled_buf, v); |
| 182 | for (int64_t j = 0; j < 16; ++j) |
| 183 | out[d2 + j] = at::Float8_e4m3fn(scaled_buf[j]); |
| 184 | } |
| 185 | } |
| 186 | #endif |
| 187 | |
| 188 | float w = index_weights[t * index_weights_stride + h]; |
| 189 | w *= scale; |
| 190 | w *= softmax_scale; |
| 191 | w *= head_scale; |
| 192 | index_weights_out[t * index_weights_out_stride + h] = w; |
| 193 | } |
| 194 | }); |
| 195 | } |
| 196 | |
| 197 | } // anonymous namespace |
| 198 | |
| 199 | void fused_indexer_q_rope_quant_cpu(at::Tensor& positions, at::Tensor& index_q, |
| 200 | at::Tensor& index_q_cos_sin_cache, |
| 201 | at::Tensor& index_q_fp8, |
| 202 | at::Tensor& index_weights, |
| 203 | double index_weights_softmax_scale, |
| 204 | double index_weights_head_scale, |
| 205 | at::Tensor& index_weights_out) { |
| 206 | TORCH_CHECK( |
| 207 | positions.dim() == 1 && positions.scalar_type() == at::kLong, |
| 208 | "fused_indexer_q_rope_quant_cpu: positions must be a 1D int64 tensor"); |
| 209 | TORCH_CHECK(positions.is_contiguous(), |
| 210 | "fused_indexer_q_rope_quant_cpu: positions must be contiguous"); |
| 211 | TORCH_CHECK(index_q.dim() == 3, |
| 212 | "fused_indexer_q_rope_quant_cpu: index_q must be 3D"); |
| 213 | TORCH_CHECK(index_q.scalar_type() == at::kFloat, |
| 214 | "fused_indexer_q_rope_quant_cpu: index_q must be float32 " |
| 215 | "(convert once in Python)"); |
| 216 | TORCH_CHECK(index_q.stride(2) == 1, |
| 217 | "fused_indexer_q_rope_quant_cpu: index_q must be contiguous in " |
| 218 | "the head_dim axis"); |
| 219 | TORCH_CHECK( |
| 220 | index_q_cos_sin_cache.dim() == 2 && index_q_cos_sin_cache.stride(1) == 1, |
| 221 | "fused_indexer_q_rope_quant_cpu: index_q_cos_sin_cache must be a " |
| 222 | "contiguous 2D tensor"); |
| 223 | TORCH_CHECK( |
| 224 | index_q_cos_sin_cache.scalar_type() == at::kFloat, |
| 225 | "fused_indexer_q_rope_quant_cpu: index_q_cos_sin_cache must be float32 " |
| 226 | "(convert once in Python)"); |
| 227 | TORCH_CHECK(index_q_fp8.sizes() == index_q.sizes() && |
| 228 | index_q_fp8.scalar_type() == at::kFloat8_e4m3fn && |
| 229 | index_q_fp8.stride(2) == 1, |
| 230 | "fused_indexer_q_rope_quant_cpu: index_q_fp8 must match " |
| 231 | "index_q's shape, be " |
| 232 | "float8_e4m3fn, and contiguous in the head_dim axis"); |
| 233 | TORCH_CHECK(index_weights.dim() == 2 && index_weights.stride(1) == 1, |
| 234 | "fused_indexer_q_rope_quant_cpu: index_weights must be a " |
| 235 | "contiguous-in-head 2D tensor"); |
| 236 | TORCH_CHECK(index_weights.scalar_type() == at::kFloat, |
| 237 | "fused_indexer_q_rope_quant_cpu: index_weights must be float32 " |
| 238 | "(convert once in Python)"); |
| 239 | TORCH_CHECK(index_weights_out.sizes() == index_weights.sizes() && |
| 240 | index_weights_out.scalar_type() == at::kFloat && |
| 241 | index_weights_out.stride(1) == 1, |
| 242 | "fused_indexer_q_rope_quant_cpu: index_weights_out must match " |
| 243 | "index_weights' shape, " |
| 244 | "be float32, and contiguous in the head axis"); |
| 245 | TORCH_CHECK(positions.device().is_cpu() && index_q.device().is_cpu() && |
| 246 | index_q_cos_sin_cache.device().is_cpu() && |
| 247 | index_q_fp8.device().is_cpu() && |
| 248 | index_weights.device().is_cpu() && |
| 249 | index_weights_out.device().is_cpu(), |
| 250 | "fused_indexer_q_rope_quant_cpu: all inputs must be CPU tensors"); |
| 251 | |
| 252 | const int64_t num_tokens = positions.size(0); |
| 253 | const int64_t num_heads = index_q.size(1); |
| 254 | const int64_t head_dim = index_q.size(2); |
| 255 | const int64_t half_rot_dim = index_q_cos_sin_cache.size(-1) / 2; |
| 256 | // Only ever called with head_dim=128, rot_dim=64 (DeepseekV4's sparse |
| 257 | // indexer Q side); the AVX512 loops below process 16 elements/pairs per |
| 258 | // iteration with no scalar remainder, so require even divisibility. |
| 259 | TORCH_CHECK(head_dim % 16 == 0 && (head_dim - 2 * half_rot_dim) % 16 == 0 && |
| 260 | half_rot_dim % 16 == 0, |
| 261 | "fused_indexer_q_rope_quant_cpu: head_dim, head_dim-rot_dim, and " |
| 262 | "rot_dim/2 " |
| 263 | "must be multiples of 16"); |
| 264 | |
| 265 | fused_indexer_q_rope_quant_impl( |
| 266 | positions.data_ptr<int64_t>(), index_q.data_ptr<float>(), |
| 267 | index_q.stride(0), index_q.stride(1), |
| 268 | index_q_cos_sin_cache.data_ptr<float>(), index_q_cos_sin_cache.stride(0), |
| 269 | half_rot_dim, index_q_fp8.data_ptr<at::Float8_e4m3fn>(), |
| 270 | index_q_fp8.stride(0), index_q_fp8.stride(1), head_dim, |
| 271 | index_weights.data_ptr<float>(), index_weights.stride(0), |
| 272 | static_cast<float>(index_weights_softmax_scale), |
| 273 | static_cast<float>(index_weights_head_scale), |
| 274 | index_weights_out.data_ptr<float>(), index_weights_out.stride(0), |
| 275 | num_tokens, num_heads); |
| 276 | } |