[Model] Add DeepSeek-V4 CPU backend (#55355)
Signed-off-by: jiang1.li <jiang1.li@intel.com>
Li, Jiang committed
Sep 10, 2026 at 14:14 UTC
c3ccc0e957fbcb1f4b55fb10ff935cf0a114e880
47 files changed
+9564
-627
.buildkite/hardware_tests/cpu.yaml
+19
-15
@@ -1,7 +1,7 @@
1
group: CPU
2
depends_on: []
3
steps:
4
-- label: CPU-Kernel Tests
4
+- label: CPU-Kernel Tests Shard %N
5
key: cpu-kernel-tests
6
depends_on: []
7
device: intel_cpu
@@ -16,6 +16,7 @@ steps:
16
- tests/kernels/attention/test_amx_mla.py
17
- vllm/v1/attention/backends/mla/amx_mla.py
18
- vllm/model_executor/layers/attention/mla_attention.py
19
+ - tests/kernels/test_deepseek_v4_cpu_kernels.py
20
- tests/kernels/moe/test_cpu_fused_moe.py
21
- tests/kernels/moe/test_cpu_quant_fused_moe.py
22
- tests/kernels/moe/test_zen_cpu_int8_moe.py
@@ -33,20 +34,23 @@ steps:
34
commands:
35
- |
36
bash .buildkite/scripts/hardware_ci/run-cpu-test.sh 30m "
36
- pytest -x -v -s tests/kernels/attention/test_cpu_attn.py
37
- pytest -x -v -s tests/v1/attention/test_group_head_counts.py
38
- pytest -x -v -s tests/kernels/attention/test_amx_mla.py
39
- pytest -x -v -s tests/kernels/moe/test_cpu_fused_moe.py
40
- pytest -x -v -s tests/kernels/moe/test_cpu_quant_fused_moe.py
41
- pytest -x -v -s tests/kernels/moe/test_zen_cpu_int8_moe.py
42
- pytest -x -v -s tests/kernels/mamba/test_cpu_short_conv.py
43
- pytest -x -v -s tests/kernels/test_onednn.py
44
- pytest -x -v -s tests/kernels/test_awq_int4_to_int8.py
45
- pytest -x -v -s tests/kernels/quantization/test_cpu_fp8_scaled_mm.py
46
- pytest -x -v -s tests/kernels/mamba/cpu/test_cpu_gdn_ops.py
47
- pytest -x -v -s tests/kernels/mamba/test_causal_conv1d.py
48
- pytest -x -v -s tests/kernels/mamba/test_mamba_ssm.py
49
- pytest -x -v -s tests/v1/sample/test_topk_topp_sampler.py::TestTritonTopkTopp"
37
+ pytest -x -v -s --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB \
38
+ tests/kernels/attention/test_cpu_attn.py \
39
+ tests/v1/attention/test_group_head_counts.py \
40
+ tests/kernels/attention/test_amx_mla.py \
41
+ tests/kernels/test_deepseek_v4_cpu_kernels.py \
42
+ tests/kernels/moe/test_cpu_fused_moe.py \
43
+ tests/kernels/moe/test_cpu_quant_fused_moe.py \
44
+ tests/kernels/moe/test_zen_cpu_int8_moe.py \
45
+ tests/kernels/mamba/test_cpu_short_conv.py \
46
+ tests/kernels/test_onednn.py \
47
+ tests/kernels/test_awq_int4_to_int8.py \
48
+ tests/kernels/quantization/test_cpu_fp8_scaled_mm.py \
49
+ tests/kernels/mamba/cpu/test_cpu_gdn_ops.py \
50
+ tests/kernels/mamba/test_causal_conv1d.py \
51
+ tests/kernels/mamba/test_mamba_ssm.py \
52
+ tests/v1/sample/test_topk_topp_sampler.py::TestTritonTopkTopp"
53
+ parallelism: 2
54
55
# Note: SDE can't be downloaded from CI host because of AWS WAF
56
# - label: CPU-Compatibility Tests
benchmarks/kernels/cpu/benchmark_cpu_fused_moe.py
+3
-1
@@ -70,12 +70,14 @@ def main(
70
71
packed_w13 = cpu_prepack_moe_weight(w13, isa)
72
packed_w2 = cpu_prepack_moe_weight(w2, isa)
73
+ output = torch.empty_like(input_tensor)
74
75
def run_benchmark(iters: int) -> list[float]:
76
times = []
77
for _ in range(iters):
78
start_time = time.perf_counter_ns()
78
- _ = cpu_fused_moe(
79
+ cpu_fused_moe(
80
+ output,
81
input_tensor,
82
packed_w13,
83
packed_w2,
cmake/cpu_extension.cmake
+8
-1
@@ -548,7 +548,14 @@ if (ENABLE_X86_ISA)
548
"csrc/cpu/sgl-kernels/bmm.cpp"
549
"csrc/cpu/sgl-kernels/decode.cpp"
550
"csrc/cpu/sgl-kernels/extend.cpp"
551
- "csrc/cpu/sgl-kernels/mla_cache.cpp")
551
+ "csrc/cpu/sgl-kernels/mla_cache.cpp"
552
+ "csrc/cpu/sgl-kernels/mhc.cpp"
553
+ "csrc/cpu/sgl-kernels/store_cache.cpp"
554
+ "csrc/cpu/sgl-kernels/flash_mla.cpp"
555
+ "csrc/cpu/sgl-kernels/compressor.cpp"
556
+ "csrc/cpu/sgl-kernels/paged_mqa_logits.cpp"
557
+ "csrc/cpu/sgl-kernels/topk.cpp"
558
+ "csrc/cpu/sgl-kernels/indexer.cpp")
559
560
set(VLLM_EXT_SRC_AVX512
561
"csrc/cpu/sgl-kernels/fla.cpp"
csrc/cpu/sgl-kernels/common.h
+30
@@ -343,6 +343,36 @@ inline void parallel_2d(int m, int n, const func_t& f) {
343
#endif
344
}
345
346
+// Like parallel_2d but with explicit nth_m x nth_n decomposition.
347
+// Caller is responsible for choosing nth_m and nth_n.
348
+//
349
+// Cherry-picked from a newer sglang commit than this file's last full sync
350
+// (needed by mhc.cpp's phase-1 tiling); not part of a full common.h resync.
351
+template <typename func_t>
352
+inline void parallel_2d_tiled(int m, int n, int nth_m, int nth_n, const func_t& f) {
353
+ const int nth = nth_m * nth_n;
354
+#if defined(_OPENMP)
355
+#pragma omp parallel num_threads(nth)
356
+ {
357
+ int ith = omp_get_thread_num();
358
+ int ith_m = ith / nth_n;
359
+ int ith_n = ith % nth_n;
360
+
361
+ int thread_block_m = div_up(m, nth_m);
362
+ int thread_block_n = div_up(n, nth_n);
363
+
364
+ int begin_m = std::min(ith_m * thread_block_m, m);
365
+ int end_m = std::min(begin_m + thread_block_m, m);
366
+ int begin_n = std::min(ith_n * thread_block_n, n);
367
+ int end_n = std::min(begin_n + thread_block_n, n);
368
+
369
+ f(begin_m, end_m, begin_n, end_n);
370
+ }
371
+#else
372
+ f(0, m, 0, n);
373
+#endif
374
+}
375
+
376
// limit max cache blocks
377
// when we need to do pre-unpack for weights, e.g. fp8
378
#define MAX_CACHE_BLOCK_SIZE 4
csrc/cpu/sgl-kernels/compressor.cpp
new
+747
@@ -0,0 +1,747 @@
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
+}
csrc/cpu/sgl-kernels/flash_mla.cpp
new
+625
@@ -0,0 +1,625 @@
1
+// Adapted from
2
+// https://github.com/sgl-project/sglang/blob/main/python/sglang/kernels/aot/csrc/cpu/flash_mla.cpp
3
+//
4
+// DeepSeek-V4 sparse MQA attention kernel for the fp8_ds_mla cache layout.
5
+// Near-verbatim port of `sparse_mla_decode_kernel_impl` (VNNI K/V packing +
6
+// AMX brgemm QK^T/PV + online softmax); the `V4_FP8Sparse` byte layout
7
+// matches vLLM's fp8_ds_mla layout exactly (see store_cache.cpp).
8
+//
9
+// Differences from upstream:
10
+// - Outer shape: upstream takes a padded batch-major query `[B, S_q, H, D]`
11
+// with per-batch topk sets; vLLM passes flat `[num_tokens, H, D]` with
12
+// per-token resolved indices. The kernel's own `at::parallel_for` already
13
+// flattens `(batch, s_q)` internally, so this is called with
14
+// `batches = num_tokens, s_q = 1` -- no padding, no kernel-body changes.
15
+// - Two physically distinct paged caches (SWA "window" + compressor
16
+// "compressed") are addressed via block_table-resolved flat slot ids,
17
+// matching upstream's `k_main`/`k_extra` two-index-set convention.
18
+// - `attn_sink` uses upstream's own post-hoc LSE correction
19
+// (`out *= 1 / (1 + exp(sink - lse))`).
20
+// - bf16 query only (matching upstream's own restriction: its FP8->bf16
21
+// dequant intrinsics are bf16-specific); fp16 is not supported here.
22
+//
23
+// clang-format off
24
+
25
+#include <algorithm>
26
+#include <cmath>
27
+#include <limits>
28
+#include <vector>
29
+
30
+#include "common.h"
31
+#include "gemm.h"
32
+#include "vec.h"
33
+
34
+#if !defined(CPU_CAPABILITY_AVX512)
35
+static_assert(false,
36
+ "flash_mla.cpp is only ever compiled into the AVX512+AMX-flagged "
37
+ "_C target (see cmake/cpu_extension.cmake's VLLM_EXT_SRC_SGL) and "
38
+ "relies on that unconditionally -- CPU_CAPABILITY_AVX512 must be "
39
+ "defined here.");
40
+#endif
41
+
42
+namespace {
43
+
44
+constexpr int64_t kNopeDim = 448;
45
+constexpr int64_t kRopeDim = 64;
46
+constexpr int64_t kHeadDim = kNopeDim + kRopeDim;
47
+constexpr int64_t kQuantBlock = 64;
48
+constexpr int64_t kNumQuantBlocks = kNopeDim / kQuantBlock;
49
+constexpr int64_t kScaleBytesPerToken = kNumQuantBlocks + 1;
50
+constexpr int64_t kTokenDataBytes = kNopeDim + kRopeDim * 2;
51
+
52
+// ---------------------------------------------------------------------------
53
+// V4_FP8Sparse layout (SGLang's `flash_mla.cpp`) -- byte-for-byte identical
54
+// to vLLM's own fp8_ds_mla layout above (d=512, d_nope=448, d_rope=64,
55
+// tile=64, num_tiles=7):
56
+// per block:
57
+// [block_size * (d_nope + 2*d_rope) bytes ; FP8 NoPE + bf16 RoPE per token]
58
+// [block_size * 8 bytes ; 7 e8m0 scales + 1 pad byte/token]
59
+// ---------------------------------------------------------------------------
60
+
61
+struct FP8LayoutMeta {
62
+ int64_t d;
63
+ int64_t d_nope;
64
+ int64_t d_rope;
65
+ int64_t tile_size;
66
+ int64_t num_tiles;
67
+};
68
+
69
+constexpr FP8LayoutMeta kV4Meta = {kHeadDim, kNopeDim, kRopeDim, kQuantBlock, kNumQuantBlocks};
70
+
71
+// Convert one fp8_e8m0 byte (= unsigned 8-bit exponent) to float.
72
+// e8m0 only stores an exponent (bias 127); value = 2^(e - 127), with 0xFF
73
+// reserved for NaN.
74
+inline float fp8_e8m0_to_float(uint8_t v) {
75
+ if (v == 0xFF) return std::numeric_limits<float>::quiet_NaN();
76
+ if (v == 0) return 0.f;
77
+ union {
78
+ uint32_t u;
79
+ float f;
80
+ } u;
81
+ u.u = static_cast<uint32_t>(v) << 23;
82
+ return u.f;
83
+}
84
+
85
+template <typename index_t>
86
+inline bool is_valid_sparse_index(index_t idx, int64_t pos, int64_t topk_limit, int64_t total_tokens) {
87
+ const int64_t v = static_cast<int64_t>(idx);
88
+ return pos < topk_limit && v >= 0 && v < total_tokens;
89
+}
90
+
91
+#if defined(CPU_CAPABILITY_AVX512)
92
+inline __attribute__((always_inline)) __m512i cvt_fp8_32_to_scaled_bf16(__m256i fp8, float scale) {
93
+ const __m512bh bf16_ext = CVT_FP8_TO_BF16_EXT(fp8);
94
+ const __m512 scale_vec = _mm512_mul_ps(_mm512_set1_ps(scale), _mm512_castsi512_ps(_mm512_set1_epi32(kFP8_BIAS)));
95
+ const __m512 f0 = _mm512_mul_ps(CVT_BF16_TO_FP32(_mm512_extracti32x8_epi32((__m512i)bf16_ext, 0)), scale_vec);
96
+ const __m512 f1 = _mm512_mul_ps(CVT_BF16_TO_FP32(_mm512_extracti32x8_epi32((__m512i)bf16_ext, 1)), scale_vec);
97
+ return (__m512i)_mm512_cvtne2ps_pbh(f1, f0);
98
+}
99
+
100
+// Load 32 contiguous logical dims (NoPE fp8-dequant-to-bf16, or raw bf16
101
+// RoPE) starting at `dim_offset` from one token row of the V4_FP8Sparse
102
+// layout.
103
+inline __attribute__((always_inline)) __m512i load_fp8_kvcache_32_from_row(
104
+ const uint8_t* __restrict__ row_base, const uint8_t* __restrict__ scale_base, int dim_offset) {
105
+ if (dim_offset < kV4Meta.d_nope) {
106
+ const float scale = fp8_e8m0_to_float(scale_base[dim_offset / kV4Meta.tile_size]);
107
+ return cvt_fp8_32_to_scaled_bf16(
108
+ _mm256_loadu_si256(reinterpret_cast<const __m256i*>(row_base + dim_offset)), scale);
109
+ }
110
+ const auto* rope_ptr = reinterpret_cast<const at::BFloat16*>(row_base + kV4Meta.d_nope);
111
+ return _mm512_loadu_si512(rope_ptr + dim_offset - kV4Meta.d_nope);
112
+}
113
+
114
+inline __attribute__((always_inline)) void init_fp8_kvcache_tile_rows(
115
+ const uint8_t* __restrict__ fp8_storage,
116
+ const int64_t block_size,
117
+ const int64_t storage_block_stride_bytes,
118
+ const int64_t token_idx,
119
+ const uint8_t*& row_base,
120
+ const uint8_t*& scale_base) {
121
+ constexpr int64_t nope_rope_per_token = kV4Meta.d_nope + 2 * kV4Meta.d_rope;
122
+ constexpr int64_t scale_stride = 8;
123
+ const int64_t block_idx = token_idx / block_size;
124
+ const int64_t block_off = token_idx - block_idx * block_size;
125
+ const uint8_t* block_base = fp8_storage + block_idx * storage_block_stride_bytes;
126
+ row_base = block_base + block_off * nope_rope_per_token;
127
+ scale_base = block_base + block_size * nope_rope_per_token + block_off * scale_stride;
128
+}
129
+
130
+template <bool convert_v, typename scalar_t, typename index_t, typename load_vec_t>
131
+inline __attribute__((always_inline)) void sparse_pack_vnni_Nx32(
132
+ scalar_t* __restrict__ dst0,
133
+ scalar_t* __restrict__ dst1,
134
+ const index_t* __restrict__ ind,
135
+ const bool* __restrict__ valid_mask,
136
+ int N,
137
+ int dim_offset,
138
+ int ld_dst0,
139
+ int ld_dst1,
140
+ const load_vec_t& load_vec) {
141
+ __m512i vinputs[16];
142
+ int n = 0;
143
+ for (; n < N; ++n) {
144
+ if (!valid_mask[n]) {
145
+ vinputs[n] = _mm512_set1_epi32(0);
146
+ } else {
147
+ vinputs[n] = load_vec(n, static_cast<int64_t>(ind[n]), dim_offset);
148
+ }
149
+ }
150
+ for (; n < 16; ++n) {
151
+ vinputs[n] = _mm512_set1_epi32(0);
152
+ }
153
+
154
+ if constexpr (convert_v) {
155
+ for (int nn = 0; nn < 16; nn += 2) {
156
+ __m512i d0, d1;
157
+ std::tie(d0, d1) = transpose_2x32_16bit(vinputs[nn], vinputs[nn + 1]);
158
+ _mm512_storeu_si512(dst1 + (nn >> 1) * ld_dst1 * 2, d0);
159
+ _mm512_storeu_si512(dst1 + (nn >> 1) * ld_dst1 * 2 + 32, d1);
160
+ }
161
+ }
162
+
163
+ transpose_16x16_32bit(vinputs);
164
+ const __mmask16 vmask = (1 << N) - 1;
165
+ for (int k = 0; k < 16; ++k) {
166
+ _mm512_mask_storeu_epi32(dst0 + k * ld_dst0 * 2, vmask, vinputs[k]);
167
+ }
168
+}
169
+#endif // CPU_CAPABILITY_AVX512
170
+
171
+// Gather+VNNI-pack `N` fp8-cache rows (dequantized to bf16 while packing)
172
+// into `dst0` (K, vnni-packed for brgemm's B operand) and `dst1` (V,
173
+// vnni-packed).
174
+template <typename scalar_t, typename index_t>
175
+void sparse_pack_vnni_fp8(
176
+ scalar_t* __restrict__ dst0,
177
+ scalar_t* __restrict__ dst1,
178
+ const uint8_t* __restrict__ fp8_storage,
179
+ const index_t* __restrict__ ind,
180
+ const bool* __restrict__ valid_mask,
181
+ int N,
182
+ int K,
183
+ int Kv,
184
+ int64_t block_size,
185
+ int64_t storage_block_stride_bytes,
186
+ int ld_dst0,
187
+ int ld_dst1) {
188
+ const int NB = div_up(N, 16);
189
+ const int KB = K / 32;
190
+ const int KBv = std::min(Kv / 32, KB);
191
+ for (int nb = 0; nb < NB; ++nb) {
192
+ const uint8_t* row_base[16];
193
+ const uint8_t* scale_base[16];
194
+ const int nb_size = std::min(N - nb * 16, 16);
195
+ for (int n = 0; n < nb_size; ++n) {
196
+ if (!valid_mask[nb * 16 + n]) {
197
+ row_base[n] = nullptr;
198
+ scale_base[n] = nullptr;
199
+ continue;
200
+ }
201
+ init_fp8_kvcache_tile_rows(
202
+ fp8_storage, block_size, storage_block_stride_bytes, static_cast<int64_t>(ind[nb * 16 + n]), row_base[n],
203
+ scale_base[n]);
204
+ }
205
+ auto load_vec = [&row_base, &scale_base](int n, int64_t idx, int dim_offset) {
206
+ UNUSED(idx);
207
+ return load_fp8_kvcache_32_from_row(row_base[n], scale_base[n], dim_offset);
208
+ };
209
+ for (int kb = 0; kb < KBv; ++kb) {
210
+ sparse_pack_vnni_Nx32<true, scalar_t, index_t>(
211
+ dst0 + ((kb * 32) >> 1) * ld_dst0 * 2 + nb * 16 * 2,
212
+ dst1 + ((nb * 16) >> 1) * ld_dst1 * 2 + kb * 32 * 2,
213
+ ind + nb * 16,
214
+ valid_mask + nb * 16,
215
+ nb_size,
216
+ kb * 32,
217
+ ld_dst0,
218
+ ld_dst1,
219
+ load_vec);
220
+ }
221
+ for (int kb = KBv; kb < KB; ++kb) {
222
+ sparse_pack_vnni_Nx32<false, scalar_t, index_t>(
223
+ dst0 + ((kb * 32) >> 1) * ld_dst0 * 2 + nb * 16 * 2,
224
+ dst1 + ((nb * 16) >> 1) * ld_dst1 * 2 + kb * 32 * 2,
225
+ ind + nb * 16,
226
+ valid_mask + nb * 16,
227
+ nb_size,
228
+ kb * 32,
229
+ ld_dst0,
230
+ ld_dst1,
231
+ load_vec);
232
+ }
233
+ }
234
+}
235
+
236
+template <typename scalar_t>
237
+inline void fmla_fill_stub(scalar_t* __restrict__ out, float val, int64_t size) {
238
+ using Vec = at::vec::Vectorized<scalar_t>;
239
+ constexpr int kVecSize = Vec::size();
240
+ const Vec data_vec = Vec(static_cast<scalar_t>(val));
241
+ int64_t d = 0;
242
+#pragma GCC unroll 4
243
+ for (; d <= size - kVecSize; d += kVecSize) {
244
+ data_vec.store(out + d);
245
+ }
246
+ if (size - d > 0) {
247
+ data_vec.store(out + d, size - d);
248
+ }
249
+}
250
+
251
+template <typename scalar_t, int BLOCK_N>
252
+inline void fmla_copy_stub(scalar_t* __restrict__ out, const float* __restrict__ input) {
253
+ static_assert(BLOCK_N % 32 == 0);
254
+ using bVec = at::vec::Vectorized<scalar_t>;
255
+ using fVec = at::vec::Vectorized<float>;
256
+ constexpr int COLS = BLOCK_N / 16;
257
+ auto store = [&](auto i) {
258
+ constexpr int col = i % COLS;
259
+ if constexpr (col % 2 == 0) {
260
+ fVec a0 = fVec::loadu(input + col * 16);
261
+ fVec a1 = fVec::loadu(input + col * 16 + 16);
262
+ bVec out_bvec = convert_from_float_ext<scalar_t>(a0, a1);
263
+ out_bvec.store(out + col * 16);
264
+ }
265
+ };
266
+ Unroll<COLS>{}(store);
267
+}
268
+
269
+template <typename scalar_t>
270
+inline void fmla_finalize_out(scalar_t* __restrict__ out, const float* __restrict__ acc, float inv_s, int64_t size) {
271
+ using bVec = at::vec::Vectorized<scalar_t>;
272
+ using fVec = at::vec::Vectorized<float>;
273
+ constexpr int kVecSize = bVec::size();
274
+ const fVec s_fvec = fVec(inv_s);
275
+ int64_t d = 0;
276
+#pragma GCC unroll 4
277
+ for (; d <= size - kVecSize; d += kVecSize) {
278
+ fVec a0 = fVec::loadu(acc + d) * s_fvec;
279
+ fVec a1 = fVec::loadu(acc + d + fVec::size()) * s_fvec;
280
+ bVec out_bvec = convert_from_float_ext<scalar_t>(a0, a1);
281
+ out_bvec.store(out + d);
282
+ }
283
+ for (; d < size; ++d) {
284
+ out[d] = static_cast<scalar_t>(acc[d] * inv_s);
285
+ }
286
+}
287
+
288
+// ---------------------------------------------------------------------------
289
+// query : [num_tokens, H_q, D_qk] bf16 (SGLang's `[B, S_q, H_q, D]`
290
+// with S_q folded away -- vLLM's flat per-token convention)
291
+// indices : [num_tokens, topk_main] int64 (flat, block-table-resolved
292
+// slot ids, -1 = invalid; SGLang's "main"/window index set)
293
+// extra_* : same, for the optional "extra"/compressed index set
294
+// attn_sink: [H_q] float32
295
+// output : [num_tokens, H_q, D_v] bf16
296
+// ---------------------------------------------------------------------------
297
+
298
+template <typename scalar_t, typename index_t, int64_t BLOCK_N>
299
+void sparse_mla_decode_kernel_impl(
300
+ scalar_t* __restrict__ output,
301
+ float* __restrict__ lse_out,
302
+ const scalar_t* __restrict__ query,
303
+ const uint8_t* __restrict__ k_main_fp8,
304
+ const uint8_t* __restrict__ k_extra_fp8,
305
+ const index_t* __restrict__ indices,
306
+ const index_t* __restrict__ extra_indices,
307
+ const float* __restrict__ attn_sink,
308
+ scalar_t* __restrict__ buffer,
309
+ int64_t batches,
310
+ int64_t s_q,
311
+ int64_t num_heads,
312
+ int64_t head_size,
313
+ int64_t head_size_v,
314
+ int64_t topk_main,
315
+ int64_t topk_extra,
316
+ int64_t total_tokens_main,
317
+ int64_t total_tokens_extra,
318
+ int64_t q_strideB,
319
+ int64_t q_strideH,
320
+ int64_t k_main_block_size,
321
+ int64_t k_extra_block_size,
322
+ int64_t k_main_block_stride_bytes,
323
+ int64_t k_extra_block_stride_bytes,
324
+ int64_t idx_strideB,
325
+ int64_t extra_idx_strideB,
326
+ float scaling,
327
+ int64_t buffer_size_per_thread) {
328
+ using Vec = at::vec::Vectorized<float>;
329
+
330
+ // partition heads
331
+ constexpr int64_t kBLOCK_H_MAX = 16;
332
+ const int64_t BLOCK_H = (batches * s_q) >= 16 ? kBLOCK_H_MAX : 8;
333
+ const int64_t num_h_blocks = div_up(num_heads, BLOCK_H);
334
+
335
+ // parallel on [B, S_q, head_block] -- with s_q == 1 (vLLM's flat call
336
+ // convention) this collapses to a plain [num_tokens, head_block] loop,
337
+ // one flat iteration per token with no padding to skip.
338
+ at::parallel_for(0, batches * s_q * num_h_blocks, 0, [&](int64_t begin, int64_t end) {
339
+ int64_t bs{0}, sq{0}, hb{0};
340
+ data_index_init(begin, bs, batches, sq, s_q, hb, num_h_blocks);
341
+
342
+ int tid = at::get_thread_num();
343
+ scalar_t* __restrict__ Btmp0 = buffer + tid * buffer_size_per_thread; // K packed
344
+ scalar_t* __restrict__ Btmp1 = Btmp0 + BLOCK_N * head_size; // V packed
345
+ // f32 V accumulator follows the bf16 packing region (reinterpret cast).
346
+ float* __restrict__ v_acc_local = reinterpret_cast<float*>(Btmp1 + BLOCK_N * head_size_v);
347
+ fmla_fill_stub(Btmp1, 0.f, BLOCK_N * head_size_v); // initialize V padding
348
+
349
+ alignas(64) float s_i[kBLOCK_H_MAX * BLOCK_N];
350
+ float* __restrict__ s_delta = s_i;
351
+ alignas(64) scalar_t s_delta2[kBLOCK_H_MAX * BLOCK_N];
352
+
353
+ alignas(64) float s_prime[kBLOCK_H_MAX];
354
+ alignas(64) float m_prime[kBLOCK_H_MAX];
355
+
356
+ for (int64_t i = begin; i < end; ++i) {
357
+ const int64_t h_start = hb * BLOCK_H;
358
+ const int64_t h_end = std::min(h_start + BLOCK_H, num_heads);
359
+ const int64_t h_size = h_end - h_start;
360
+
361
+ const scalar_t* __restrict__ q_ptr = query + bs * q_strideB + h_start * q_strideH;
362
+ const index_t* __restrict__ idx_ptr = indices + bs * idx_strideB;
363
+ const index_t* __restrict__ extra_idx_ptr = extra_indices == nullptr ? nullptr : extra_indices + bs * extra_idx_strideB;
364
+
365
+ fmla_fill_stub(s_prime, 0.f, BLOCK_H);
366
+ fmla_fill_stub(m_prime, -std::numeric_limits<float>::infinity(), BLOCK_H);
367
+ for (int64_t h = 0; h < h_size; ++h) {
368
+ fmla_fill_stub(v_acc_local + h * head_size_v, 0.f, head_size_v);
369
+ }
370
+
371
+ auto process_cache = [&](const uint8_t* __restrict__ fp8_ptr,
372
+ const index_t* __restrict__ cur_idx_ptr,
373
+ int64_t topk_count,
374
+ int64_t total_tokens,
375
+ int64_t fp8_block_size,
376
+ int64_t fp8_block_stride_bytes) {
377
+ if (fp8_ptr == nullptr || cur_idx_ptr == nullptr || topk_count == 0) {
378
+ return;
379
+ }
380
+ const int64_t topk_limit = topk_count;
381
+
382
+ for (int64_t n = 0; n < topk_limit; n += BLOCK_N) {
383
+ int64_t n_size = std::min<int64_t>(BLOCK_N, topk_limit - n);
384
+ const int64_t padded_n_size = div_up(int(n_size), TILE_K) * TILE_K;
385
+ bool valid_mask[BLOCK_N];
386
+ bool has_valid = false;
387
+ for (int64_t k = 0; k < n_size; ++k) {
388
+ const bool valid = is_valid_sparse_index(cur_idx_ptr[n + k], n + k, topk_limit, total_tokens);
389
+ valid_mask[k] = valid;
390
+ has_valid |= valid;
391
+ }
392
+ if (!has_valid) {
393
+ continue;
394
+ }
395
+
396
+ // Pack K (BLOCK_N rows via gather) into Btmp0 (key, vnni) and Btmp1
397
+ // (value, vnni), dequantizing FP8 cache entries while packing.
398
+ sparse_pack_vnni_fp8<scalar_t, index_t>(
399
+ Btmp0, Btmp1, fp8_ptr, cur_idx_ptr + n, valid_mask, static_cast<int>(n_size), static_cast<int>(head_size),
400
+ static_cast<int>(head_size_v), fp8_block_size, fp8_block_stride_bytes, static_cast<int>(BLOCK_N),
401
+ static_cast<int>(head_size_v));
402
+
403
+ // Q @ K
404
+ at::native::cpublas::brgemm(
405
+ /* M */ h_size,
406
+ /* N */ n_size,
407
+ /* K */ head_size,
408
+ /* lda */ q_strideH,
409
+ /* ldb */ BLOCK_N,
410
+ /* ldc */ BLOCK_N,
411
+ /* add_C */ false,
412
+ /* A */ q_ptr,
413
+ /* B */ Btmp0,
414
+ /* C */ s_i);
415
+
416
+ const Vec scale_vec = Vec(scaling);
417
+ for (int64_t h = 0; h < h_size; ++h) {
418
+ float* row = s_i + h * BLOCK_N;
419
+ at::vec::map<float>([scale_vec](Vec x) { return x * scale_vec; }, row, row, n_size);
420
+
421
+ for (int64_t k = 0; k < n_size; ++k) {
422
+ if (!valid_mask[k]) {
423
+ row[k] = -std::numeric_limits<float>::infinity();
424
+ }
425
+ }
426
+
427
+ // online softmax update
428
+ float m_i = at::vec::reduce_all<float>([](Vec& x, Vec& y) { return at::vec::maximum(x, y); }, row, n_size);
429
+ m_i = std::max(m_i, m_prime[h]);
430
+
431
+ if (!std::isfinite(m_i)) {
432
+ fmla_fill_stub(s_delta + h * BLOCK_N, 0.f, padded_n_size);
433
+ fmla_copy_stub<scalar_t, BLOCK_N>(s_delta2 + h * BLOCK_N, s_delta + h * BLOCK_N);
434
+ continue;
435
+ }
436
+
437
+ const float m_delta = std::exp(m_prime[h] - m_i);
438
+ at::vec::map<float>([m_i](Vec x) { return (x - Vec(m_i)).exp_u20(); }, s_delta + h * BLOCK_N, row, n_size);
439
+
440
+ s_prime[h] *= m_delta;
441
+ s_prime[h] += at::vec::reduce_all<float>([](Vec& x, Vec& y) { return x + y; }, s_delta + h * BLOCK_N, n_size);
442
+
443
+ m_prime[h] = m_i;
444
+
445
+ float scale_m = m_delta;
446
+ at::vec::map<float>(
447
+ [scale_m](Vec x) { return x * Vec(scale_m); }, v_acc_local + h * head_size_v, v_acc_local + h * head_size_v,
448
+ head_size_v);
449
+
450
+ fmla_fill_stub(s_delta + h * BLOCK_N + n_size, 0.f, padded_n_size - n_size);
451
+ fmla_copy_stub<scalar_t, BLOCK_N>(s_delta2 + h * BLOCK_N, s_delta + h * BLOCK_N);
452
+ }
453
+
454
+ // V' <- s_delta @ V + V' (accumulate into v_acc_local at f32)
455
+ at::native::cpublas::brgemm(
456
+ /* M */ h_size,
457
+ /* N */ head_size_v,
458
+ /* K */ padded_n_size,
459
+ /* lda */ BLOCK_N,
460
+ /* ldb */ head_size_v,
461
+ /* ldc */ head_size_v,
462
+ /* add_C */ true,
463
+ /* A */ s_delta2,
464
+ /* B */ Btmp1,
465
+ /* C */ v_acc_local);
466
+ }
467
+ };
468
+
469
+ process_cache(k_main_fp8, idx_ptr, topk_main, total_tokens_main, k_main_block_size, k_main_block_stride_bytes);
470
+ process_cache(
471
+ k_extra_fp8, extra_idx_ptr, topk_extra, total_tokens_extra, k_extra_block_size, k_extra_block_stride_bytes);
472
+
473
+ // Apply attention sink correction directly on the output and lse.
474
+ // out *= exp(lse_no_sink) / (exp(lse_no_sink) + exp(attn_sink))
475
+ // = 1 / (1 + exp(attn_sink - lse_no_sink))
476
+ // (Algebraically equivalent to treating attn_sink as a virtual,
477
+ // zero-value extra key inside the softmax.)
478
+ for (int64_t h = 0; h < h_size; ++h) {
479
+ const int64_t hh = h_start + h;
480
+ const bool lonely = !std::isfinite(m_prime[h]) || s_prime[h] == 0.f;
481
+ float lse_val = lonely ? std::numeric_limits<float>::infinity() : (m_prime[h] + std::log(s_prime[h]));
482
+ float inv_s = lonely ? 0.f : (1.f / s_prime[h]);
483
+
484
+ if (!lonely && attn_sink != nullptr) {
485
+ const float sink = attn_sink[hh];
486
+ const float corr = 1.f / (1.f + std::exp(sink - lse_val));
487
+ inv_s *= corr;
488
+ }
489
+
490
+ scalar_t* out_row = output + bs * (s_q * num_heads * head_size_v) + sq * (num_heads * head_size_v) + hh * head_size_v;
491
+ if (lonely) {
492
+ fmla_fill_stub(out_row, 0.f, head_size_v);
493
+ } else {
494
+ fmla_finalize_out<scalar_t>(out_row, v_acc_local + h * head_size_v, inv_s, head_size_v);
495
+ }
496
+
497
+ lse_out[bs * num_heads * s_q + hh * s_q + sq] = lse_val;
498
+ }
499
+
500
+ data_index_step(bs, batches, sq, s_q, hb, num_h_blocks);
501
+ }
502
+ at::native::cpublas::brgemm_release();
503
+ });
504
+}
505
+
506
+} // anonymous namespace
507
+
508
+void flash_mla_with_kvcache_cpu(
509
+ at::Tensor& out,
510
+ at::Tensor& q,
511
+ at::Tensor& window_cache_2d,
512
+ at::Tensor& window_slots,
513
+ int64_t window_block_size,
514
+ at::Tensor& compressed_cache_2d,
515
+ at::Tensor& compressed_slots,
516
+ int64_t compressed_block_size,
517
+ at::Tensor& attn_sink,
518
+ double scale) {
519
+ TORCH_CHECK(
520
+ q.dim() == 3 && q.size(2) == kHeadDim,
521
+ "flash_mla_with_kvcache_cpu: q must be [num_tokens, num_heads, 512]");
522
+ TORCH_CHECK(
523
+ out.dim() == 3 && out.size(0) == q.size(0) && out.size(1) == q.size(1) && out.size(2) == kHeadDim,
524
+ "flash_mla_with_kvcache_cpu: out must be [num_tokens, num_heads, 512]");
525
+ TORCH_CHECK(
526
+ out.scalar_type() == at::kBFloat16 && out.is_contiguous(),
527
+ "flash_mla_with_kvcache_cpu: out must be bf16 and contiguous "
528
+ "(pre-allocate once in Python -- see DeepseekV4CPUAttention.forward_mqa)");
529
+ TORCH_CHECK(out.device().is_cpu(), "flash_mla_with_kvcache_cpu: out must be a CPU tensor");
530
+ TORCH_CHECK(
531
+ q.scalar_type() == at::kBFloat16,
532
+ "flash_mla_with_kvcache_cpu: q must be bf16 (the AMX/VNNI kernel's "
533
+ "FP8->bf16 dequant intrinsics are bf16-specific)");
534
+ TORCH_CHECK(
535
+ window_slots.dim() == 2 && window_slots.size(0) == q.size(0),
536
+ "flash_mla_with_kvcache_cpu: window_slots must be [num_tokens, num_window]");
537
+ TORCH_CHECK(
538
+ compressed_slots.dim() == 2 && compressed_slots.size(0) == q.size(0),
539
+ "flash_mla_with_kvcache_cpu: compressed_slots must be [num_tokens, num_compressed]");
540
+ TORCH_CHECK(
541
+ window_slots.scalar_type() == at::kLong && compressed_slots.scalar_type() == at::kLong,
542
+ "flash_mla_with_kvcache_cpu: slot tensors must be int64");
543
+ TORCH_CHECK(
544
+ attn_sink.numel() == q.size(1), "flash_mla_with_kvcache_cpu: attn_sink must have num_heads entries");
545
+ TORCH_CHECK(
546
+ attn_sink.scalar_type() == at::kFloat && attn_sink.is_contiguous(),
547
+ "flash_mla_with_kvcache_cpu: attn_sink must be float32 and contiguous "
548
+ "(convert once in Python -- see DeepseekV4CPUAttention.forward_mqa)");
549
+ TORCH_CHECK(
550
+ window_slots.is_contiguous() && compressed_slots.is_contiguous(),
551
+ "flash_mla_with_kvcache_cpu: window_slots/compressed_slots must be contiguous");
552
+ TORCH_CHECK(
553
+ q.device().is_cpu() && window_cache_2d.device().is_cpu() && window_slots.device().is_cpu() &&
554
+ compressed_cache_2d.device().is_cpu() && compressed_slots.device().is_cpu() &&
555
+ attn_sink.device().is_cpu(),
556
+ "flash_mla_with_kvcache_cpu: all inputs must be CPU tensors");
557
+
558
+ const int64_t num_tokens = q.size(0);
559
+ const int64_t num_heads = q.size(1);
560
+ const int64_t num_window = window_slots.size(1);
561
+ const int64_t num_compressed = compressed_slots.size(1);
562
+ TORCH_CHECK(num_window > 0, "flash_mla_with_kvcache_cpu: window_slots must have at least one column");
563
+
564
+ // Row stride is threaded through explicitly rather than assumed to equal
565
+ // block_size * (token_data_bytes + scale_bytes_per_token): a paged/pooled
566
+ // cache buffer's physical row stride is not guaranteed to match the
567
+ // logical per-block byte count (matches store_cache.cpp's precedent).
568
+ const int64_t window_cache_row_stride = window_cache_2d.stride(0);
569
+ const int64_t compressed_cache_row_stride = num_compressed > 0 ? compressed_cache_2d.stride(0) : 0;
570
+
571
+ // AMX/VNNI path: flatten (batch, s_q) -> batches = num_tokens, s_q = 1.
572
+ CHECK_LAST_DIM_CONTIGUOUS_INPUT(q);
573
+ constexpr int64_t BLOCK_N = 128; // multiple of 32 for AMX brgemm
574
+ constexpr int64_t kBLOCK_H_MAX = 16;
575
+ static_assert(kHeadDim % 32 == 0, "head_dim must be a multiple of 32");
576
+
577
+ const int num_threads = at::get_num_threads();
578
+ // Layout per thread (in bf16 elements):
579
+ // [Btmp0 : BLOCK_N * D] [Btmp1 : BLOCK_N * D] [v_acc_local : kBLOCK_H_MAX * D floats]
580
+ // f32 takes 2 bf16 elements -> multiply by 2.
581
+ const int64_t buffer_size_per_thread = BLOCK_N * kHeadDim + BLOCK_N * kHeadDim + 2 * kBLOCK_H_MAX * kHeadDim;
582
+ auto buffer = at::empty({num_threads, buffer_size_per_thread}, q.options());
583
+ // lse is an internal scratch (attn_sink correction needs it per-token/head)
584
+ // -- not surfaced to the Python call site, which only consumes `out`.
585
+ auto lse = at::empty({num_tokens, num_heads}, q.options().dtype(at::kFloat));
586
+
587
+ const int64_t q_strideB = q.stride(0);
588
+ const int64_t q_strideH = q.stride(1);
589
+ const int64_t idx_strideB = window_slots.stride(0);
590
+ const int64_t extra_idx_strideB = num_compressed > 0 ? compressed_slots.stride(0) : 0;
591
+
592
+ const int64_t total_tokens_main = window_cache_2d.size(0) * window_block_size;
593
+ const int64_t total_tokens_extra = num_compressed > 0 ? compressed_cache_2d.size(0) * compressed_block_size : 0;
594
+
595
+ sparse_mla_decode_kernel_impl<at::BFloat16, int64_t, BLOCK_N>(
596
+ out.data_ptr<at::BFloat16>(),
597
+ lse.data_ptr<float>(),
598
+ q.data_ptr<at::BFloat16>(),
599
+ window_cache_2d.data_ptr<uint8_t>(),
600
+ num_compressed > 0 ? compressed_cache_2d.data_ptr<uint8_t>() : nullptr,
601
+ window_slots.data_ptr<int64_t>(),
602
+ num_compressed > 0 ? compressed_slots.data_ptr<int64_t>() : nullptr,
603
+ attn_sink.data_ptr<float>(),
604
+ buffer.data_ptr<at::BFloat16>(),
605
+ /* batches */ num_tokens,
606
+ /* s_q */ 1,
607
+ num_heads,
608
+ /* head_size */ kHeadDim,
609
+ /* head_size_v */ kHeadDim,
610
+ /* topk_main */ num_window,
611
+ /* topk_extra */ num_compressed,
612
+ total_tokens_main,
613
+ total_tokens_extra,
614
+ q_strideB,
615
+ q_strideH,
616
+ window_block_size,
617
+ num_compressed > 0 ? compressed_block_size : 1,
618
+ window_cache_row_stride,
619
+ compressed_cache_row_stride,
620
+ idx_strideB,
621
+ extra_idx_strideB,
622
+ static_cast<float>(scale),
623
+ buffer_size_per_thread);
624
+}
625
+
csrc/cpu/sgl-kernels/indexer.cpp
new
+276
@@ -0,0 +1,276 @@
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
+}
csrc/cpu/sgl-kernels/mhc.cpp
new
+1046
@@ -0,0 +1,1046 @@
1
+// Adapted from
2
+// https://github.com/sgl-project/sglang/blob/main/python/sglang/kernels/aot/csrc/cpu/mhc.cpp
3
+//
4
+// DeepSeek-V4 mHC (multi-residual-stream hyper-connection) gating kernels.
5
+//
6
+// Difference from upstream: vLLM's `mhc_pre_torch` reference (kernels/
7
+// mhc/torch.py) exposes separate `hc_pre_eps`/`hc_sinkhorn_eps`/
8
+// `hc_post_mult_value`/`n_splits` knobs this kernel doesn't implement --
9
+// only ever called with hc_pre_eps == hc_sinkhorn_eps, hc_post_mult_value
10
+// == 2.0, n_splits == 1 (matches upstream), so the entry points take one
11
+// merged `hc_eps` and hardcode the post-mix multiplier to 2.0. The Python
12
+// wrapper (kernels/mhc/cpu.py) enforces the merged-eps assumption via
13
+// TORCH_CHECK.
14
+//
15
+// clang-format off
16
+
17
+#include <algorithm>
18
+#include <cmath>
19
+#include <limits>
20
+
21
+#include "common.h"
22
+#include "vec.h"
23
+
24
+namespace {
25
+
26
+// Loop tiling heuristics.
27
+constexpr int64_t kMhcTileCacheBudgetBytes = 36 * 1024;
28
+constexpr int64_t kMhcDefaultTokenBlock = 32;
29
+
30
+inline int64_t round_to_nearest_multiple(int64_t value, int64_t multiple) {
31
+ return ((value + multiple / 2) / multiple) * multiple;
32
+}
33
+
34
+inline int64_t choose_t_block(int64_t /*T*/) {
35
+ return kMhcDefaultTokenBlock;
36
+}
37
+
38
+inline int64_t choose_k_block(int64_t weight_rows, int64_t vec_size, int64_t max_block) {
39
+ const int64_t bytes_per_k = (weight_rows + 1) * static_cast<int64_t>(sizeof(float));
40
+ const int64_t raw = std::max(vec_size, kMhcTileCacheBudgetBytes / bytes_per_k);
41
+ const int64_t rounded = round_to_nearest_multiple(raw, vec_size);
42
+ return std::max(vec_size, std::min(max_block, rounded));
43
+}
44
+
45
+template <typename func_t>
46
+inline void parallel_mhc_phase1(int64_t n_t_blocks, int64_t n_k_blocks, const func_t& f) {
47
+ const int64_t nth = static_cast<int64_t>(at::get_num_threads());
48
+ if (n_t_blocks >= nth) {
49
+ // Large prefill: enough token blocks exist to fill all workers.
50
+ parallel_2d_tiled(static_cast<int>(n_t_blocks), static_cast<int>(n_k_blocks), static_cast<int>(nth), 1, f);
51
+ } else {
52
+ // Decode / small T: split both T and K to expose enough parallel tasks.
53
+ parallel_2d(static_cast<int>(n_t_blocks), static_cast<int>(n_k_blocks), f);
54
+ }
55
+}
56
+
57
+// In-place Sinkhorn on a row-major [hc, hc] matrix.
58
+inline void
59
+sinkhorn_inplace(float* __restrict__ comb_matrix, int hc, int sinkhorn_iters, float eps, float* __restrict__ scratch) {
60
+ float* const row_sum = scratch;
61
+ float* const col_sum = scratch + hc;
62
+
63
+ // First iteration.
64
+ for (int r = 0; r < hc; ++r) {
65
+ const float* row = comb_matrix + r * hc;
66
+ float row_acc = 0.f;
67
+#pragma GCC unroll 4
68
+ for (int c = 0; c < hc; ++c)
69
+ row_acc += row[c];
70
+ row_sum[r] = row_acc;
71
+ }
72
+ for (int r = 0; r < hc; ++r) {
73
+ const float row_inv = 1.0f / row_sum[r];
74
+ float* row = comb_matrix + r * hc;
75
+#pragma GCC unroll 4
76
+ for (int c = 0; c < hc; ++c)
77
+ row[c] = row[c] * row_inv + eps;
78
+ }
79
+ for (int c = 0; c < hc; ++c) {
80
+ float col_acc = 0.f;
81
+#pragma GCC unroll 4
82
+ for (int r = 0; r < hc; ++r)
83
+ col_acc += comb_matrix[r * hc + c];
84
+ col_sum[c] = col_acc;
85
+ }
86
+ for (int r = 0; r < hc; ++r)
87
+#pragma GCC unroll 4
88
+ for (int c = 0; c < hc; ++c)
89
+ comb_matrix[r * hc + c] /= (col_sum[c] + eps);
90
+
91
+ // Remaining iterations.
92
+ for (int iter = 1; iter < sinkhorn_iters; ++iter) {
93
+ for (int r = 0; r < hc; ++r) {
94
+ const float* row = comb_matrix + r * hc;
95
+ float row_acc = 0.f;
96
+#pragma GCC unroll 4
97
+ for (int c = 0; c < hc; ++c)
98
+ row_acc += row[c];
99
+ row_sum[r] = row_acc;
100
+ }
101
+ for (int r = 0; r < hc; ++r) {
102
+ const float row_inv = 1.0f / (row_sum[r] + eps);
103
+ float* row = comb_matrix + r * hc;
104
+#pragma GCC unroll 4
105
+ for (int c = 0; c < hc; ++c)
106
+ row[c] *= row_inv;
107
+ }
108
+ for (int c = 0; c < hc; ++c) {
109
+ float col_acc = 0.f;
110
+#pragma GCC unroll 4
111
+ for (int r = 0; r < hc; ++r)
112
+ col_acc += comb_matrix[r * hc + c];
113
+ col_sum[c] = col_acc;
114
+ }
115
+ for (int r = 0; r < hc; ++r)
116
+#pragma GCC unroll 4
117
+ for (int c = 0; c < hc; ++c)
118
+ comb_matrix[r * hc + c] /= (col_sum[c] + eps);
119
+ }
120
+}
121
+
122
+// Parse mixes -> pre/post/comb and run sinkhorn on comb.
123
+inline void parse_mixes_and_sinkhorn(
124
+ float* __restrict__ pre_gate, // [hc] output
125
+ float* __restrict__ post_gate, // [hc] output
126
+ float* __restrict__ comb_matrix, // [hc*hc] output (row-major)
127
+ const float* __restrict__ mixes, // [mix_hc] input
128
+ float s0,
129
+ float s1,
130
+ float s2,
131
+ const float* __restrict__ hc_base,
132
+ int sinkhorn_iters,
133
+ float eps,
134
+ int hc,
135
+ float* __restrict__ scratch) {
136
+#pragma GCC unroll 4
137
+ for (int h = 0; h < hc; ++h) {
138
+ const float v = mixes[h] * s0 + hc_base[h];
139
+ pre_gate[h] = 1.f / (1.f + std::exp(-v)) + eps;
140
+ }
141
+#pragma GCC unroll 4
142
+ for (int h = 0; h < hc; ++h) {
143
+ const float v = mixes[hc + h] * s1 + hc_base[hc + h];
144
+ post_gate[h] = 2.f / (1.f + std::exp(-v));
145
+ }
146
+ float* const row_max = scratch;
147
+ for (int r = 0; r < hc; ++r)
148
+ row_max[r] = -std::numeric_limits<float>::infinity();
149
+ for (int r = 0; r < hc; ++r) {
150
+ float* row = comb_matrix + r * hc;
151
+ for (int c = 0; c < hc; ++c) {
152
+ const int idx = r * hc + c;
153
+ const float v = mixes[2 * hc + idx] * s2 + hc_base[2 * hc + idx];
154
+ row[c] = v;
155
+ if (v > row_max[r]) row_max[r] = v;
156
+ }
157
+ }
158
+ for (int r = 0; r < hc; ++r) {
159
+ float* row = comb_matrix + r * hc;
160
+#pragma GCC unroll 4
161
+ for (int c = 0; c < hc; ++c)
162
+ row[c] = std::exp(row[c] - row_max[r]);
163
+ }
164
+
165
+ sinkhorn_inplace(comb_matrix, hc, sinkhorn_iters, eps, scratch + hc);
166
+}
167
+
168
+template <typename scalar_t>
169
+static void hc_post_fuse_impl(
170
+ scalar_t* __restrict__ out, // [T, hc, d] output
171
+ const scalar_t* __restrict__ x, // [T, d] input
172
+ const scalar_t* __restrict__ residual, // [T, hc, d] input
173
+ const float* __restrict__ post, // [T, hc] float32
174
+ const float* __restrict__ comb, // [T, hc, hc] float32 row-major
175
+ int64_t T,
176
+ int64_t d,
177
+ int hc) {
178
+ using bVec = at::vec::Vectorized<scalar_t>;
179
+ using fVec = at::vec::Vectorized<float>;
180
+ constexpr int64_t kVecSize = bVec::size();
181
+
182
+ const int64_t nthreads = at::get_num_threads();
183
+ const int64_t max_splits = std::max(int64_t(1), d / kVecSize);
184
+ const int64_t token_head_tasks = T * hc;
185
+ const int64_t k_splits = (token_head_tasks >= nthreads) ? int64_t(1) : std::min(nthreads, max_splits);
186
+ const int64_t k_chunk = (d + k_splits - 1) / k_splits;
187
+
188
+ const int64_t fvec_floats = static_cast<int64_t>(sizeof(fVec) / sizeof(float));
189
+ auto comb_fvec_storage = at::empty({nthreads * hc * fvec_floats}, at::kFloat);
190
+ fVec* const comb_fvec_base = reinterpret_cast<fVec*>(comb_fvec_storage.data_ptr<float>());
191
+
192
+ at::parallel_for(0, T * hc * k_splits, 0, [&](int64_t begin, int64_t end) {
193
+ const int tid = static_cast<int>(at::get_thread_num());
194
+ fVec* const comb_fvec = comb_fvec_base + tid * hc;
195
+
196
+ for (int64_t idx = begin; idx < end; ++idx) {
197
+ const int64_t task = idx / k_splits;
198
+ const int64_t split_id = idx % k_splits;
199
+ const int64_t t = task / hc;
200
+ const int h = static_cast<int>(task % hc);
201
+ const int64_t k0 = split_id * k_chunk;
202
+ const int64_t k1 = std::min(k0 + k_chunk, d);
203
+
204
+ const float post_val = post[t * hc + h];
205
+ const float* comb_t = comb + t * hc * hc;
206
+ const scalar_t* x_t = x + t * d;
207
+ const scalar_t* res_t = residual + t * hc * d;
208
+ scalar_t* out_th = out + (t * hc + h) * d;
209
+
210
+ const fVec post_fvec(post_val);
211
+#pragma GCC unroll 4
212
+ for (int i = 0; i < hc; ++i)
213
+ comb_fvec[i] = fVec(comb_t[i * hc + h]);
214
+
215
+ int64_t k;
216
+ // bf16: convert and accumulate
217
+ for (k = k0; k <= k1 - kVecSize; k += kVecSize) {
218
+ fVec acc0, acc1;
219
+ std::tie(acc0, acc1) = at::vec::convert_to_float(bVec::loadu(x_t + k));
220
+ acc0 = post_fvec * acc0;
221
+ acc1 = post_fvec * acc1;
222
+#pragma GCC unroll 4
223
+ for (int i = 0; i < hc; ++i) {
224
+ fVec r0, r1;
225
+ std::tie(r0, r1) = at::vec::convert_to_float(bVec::loadu(res_t + i * d + k));
226
+ acc0 += comb_fvec[i] * r0;
227
+ acc1 += comb_fvec[i] * r1;
228
+ }
229
+ at::vec::convert_from_float<scalar_t>(acc0, acc1).store(out_th + k);
230
+ }
231
+ if (k < k1) {
232
+ const int64_t rem = k1 - k;
233
+ fVec acc0, acc1;
234
+ std::tie(acc0, acc1) = at::vec::convert_to_float(bVec::loadu(x_t + k, rem));
235
+ acc0 = post_fvec * acc0;
236
+ acc1 = post_fvec * acc1;
237
+#pragma GCC unroll 4
238
+ for (int i = 0; i < hc; ++i) {
239
+ fVec r0, r1;
240
+ std::tie(r0, r1) = at::vec::convert_to_float(bVec::loadu(res_t + i * d + k, rem));
241
+ acc0 += comb_fvec[i] * r0;
242
+ acc1 += comb_fvec[i] * r1;
243
+ }
244
+ at::vec::convert_from_float<scalar_t>(acc0, acc1).store(out_th + k, rem);
245
+ }
246
+ }
247
+ });
248
+}
249
+
250
+// Fused hc_pre path.
251
+template <typename scalar_t>
252
+static void hc_pre_fuse_impl(
253
+ scalar_t* __restrict__ y, // [T, d] bf16 output
254
+ float* __restrict__ post_out, // [T, hc] float32 output
255
+ float* __restrict__ comb_out, // [T, hc*hc] float32 output
256
+ const scalar_t* __restrict__ x, // [T, hc, d] bf16 input
257
+ const float* __restrict__ hc_fn, // [mix_hc, hc_d] float32 row-major
258
+ const float* __restrict__ hc_scale, // [3]
259
+ const float* __restrict__ hc_base, // [mix_hc]
260
+ int64_t T,
261
+ int64_t d,
262
+ int hc,
263
+ int sinkhorn_iters,
264
+ float hc_eps,
265
+ float norm_eps) {
266
+ using bVec = at::vec::Vectorized<scalar_t>;
267
+ using fVec = at::vec::Vectorized<float>;
268
+ constexpr int64_t kVecSize = bVec::size(); // 32 bf16
269
+ constexpr int64_t kFVecSize = fVec::size(); // 16 fp32
270
+ constexpr int64_t kSmallTKBlockCap = 128;
271
+ constexpr int64_t kLargeTKBlockCap = 256;
272
+ constexpr int64_t kDVecsPerBlock = 16;
273
+ constexpr int64_t kCacheKBlockCap = kLargeTKBlockCap;
274
+ constexpr int64_t MAX_K_VECS = kCacheKBlockCap / kVecSize;
275
+ const int mix_hc = (2 + hc) * hc;
276
+ const int64_t hc_d = hc * d;
277
+ const float s0 = hc_scale[0], s1 = hc_scale[1], s2 = hc_scale[2];
278
+
279
+ const int64_t T_BLOCK = choose_t_block(T);
280
+ const int64_t num_threads = static_cast<int64_t>(at::get_num_threads());
281
+ const int64_t K_BLOCK_CAP = (T >= num_threads) ? kLargeTKBlockCap : kSmallTKBlockCap;
282
+ const int64_t K_BLOCK = choose_k_block(mix_hc, kVecSize, K_BLOCK_CAP);
283
+ TORCH_CHECK(
284
+ K_BLOCK <= K_BLOCK_CAP && K_BLOCK_CAP <= kCacheKBlockCap && K_BLOCK % kVecSize == 0,
285
+ "hc_pre_fuse_impl: K_BLOCK must be within capacity and a multiple of vector size");
286
+ const int64_t D_BLOCK = kDVecsPerBlock * kVecSize;
287
+
288
+ const int64_t n_t_blocks = div_up(T, T_BLOCK);
289
+ const int64_t n_k_blocks = div_up(hc_d, K_BLOCK);
290
+
291
+ // Partial buffer: [n_t_blocks, n_k_blocks, T_BLOCK, 1+mix_hc]
292
+ const int64_t partial_stride = 1 + mix_hc;
293
+ auto partials_tensor = at::empty({n_t_blocks * n_k_blocks * T_BLOCK * partial_stride}, at::kFloat);
294
+ float* const partials = partials_tensor.data_ptr<float>();
295
+
296
+ // Per-thread scratch: [mix_buffer[mix_hc] | pre_gate[hc] |
297
+ // post_gate[hc] | comb_matrix[hc*hc] | sinkhorn_scratch[3*hc]]
298
+ const int64_t per_thread_scratch = mix_hc + hc + hc + hc * hc + 3 * hc;
299
+ auto scratch_tensor = at::empty({num_threads * per_thread_scratch}, at::kFloat);
300
+ float* const scratch_base = scratch_tensor.data_ptr<float>();
301
+
302
+ // Phase 1.
303
+ parallel_mhc_phase1(n_t_blocks, n_k_blocks, [&](int64_t tb0, int64_t tb1, int64_t kb0, int64_t kb1) {
304
+ for (int64_t kb = kb0; kb < kb1; ++kb) {
305
+ const int64_t k0 = kb * K_BLOCK;
306
+ const int64_t klen = std::min(K_BLOCK, hc_d - k0);
307
+
308
+ for (int64_t tb = tb0; tb < tb1; ++tb) {
309
+ const int64_t t0 = tb * T_BLOCK;
310
+ const int64_t tlen = std::min(T_BLOCK, T - t0);
311
+ float* partial_block = partials + (tb * n_k_blocks + kb) * T_BLOCK * partial_stride;
312
+
313
+ const float* w_ptr[mix_hc];
314
+ for (int h = 0; h < mix_hc; ++h)
315
+ w_ptr[h] = hc_fn + h * hc_d + k0;
316
+
317
+ // Stack cache is sized for the largest selected K block, but only the first
318
+ // n_full entries are touched for the current K_BLOCK.
319
+ fVec x_cache_lo[MAX_K_VECS], x_cache_hi[MAX_K_VECS];
320
+
321
+ for (int64_t tl = 0; tl < tlen; ++tl) {
322
+ const scalar_t* x_t = x + (t0 + tl) * hc_d + k0;
323
+ float* partial_row = partial_block + tl * partial_stride;
324
+
325
+ const int64_t n_full = klen / kVecSize;
326
+ const int64_t k_tail = n_full * kVecSize;
327
+ const int64_t rem = klen - k_tail;
328
+ const int64_t rem0 = std::min(rem, kFVecSize);
329
+ const int64_t rem1 = rem - rem0;
330
+
331
+ // Prefetch next token.
332
+ if (tl + 1 < tlen) {
333
+ const scalar_t* xn = x + (t0 + tl + 1) * hc_d + k0;
334
+ __builtin_prefetch(xn, 0, 0);
335
+ if (klen > kVecSize * 4) __builtin_prefetch(xn + kVecSize * 4, 0, 0);
336
+ }
337
+
338
+ // h=0..3: convert/cache x and accumulate sq + four dots.
339
+ if (tl == 0) {
340
+ for (int pw = 1; pw < mix_hc; ++pw)
341
+ __builtin_prefetch(w_ptr[pw], 0, 3);
342
+ }
343
+ {
344
+ const float* w0 = w_ptr[0];
345
+ const float* w1 = w_ptr[1];
346
+ const float* w2 = w_ptr[2];
347
+ const float* w3 = w_ptr[3];
348
+ fVec sq_lo(0.f), sq_hi(0.f);
349
+ fVec d0_lo(0.f), d0_hi(0.f);
350
+ fVec d1_lo(0.f), d1_hi(0.f);
351
+ fVec d2_lo(0.f), d2_hi(0.f);
352
+ fVec d3_lo(0.f), d3_hi(0.f);
353
+ for (int64_t ki = 0; ki < n_full; ++ki) {
354
+ const int64_t k = ki * kVecSize;
355
+ std::tie(x_cache_lo[ki], x_cache_hi[ki]) = at::vec::convert_to_float(bVec::loadu(x_t + k));
356
+ sq_lo += x_cache_lo[ki] * x_cache_lo[ki];
357
+ sq_hi += x_cache_hi[ki] * x_cache_hi[ki];
358
+ d0_lo += x_cache_lo[ki] * fVec::loadu(w0 + k);
359
+ d0_hi += x_cache_hi[ki] * fVec::loadu(w0 + k + kFVecSize);
360
+ d1_lo += x_cache_lo[ki] * fVec::loadu(w1 + k);
361
+ d1_hi += x_cache_hi[ki] * fVec::loadu(w1 + k + kFVecSize);
362
+ d2_lo += x_cache_lo[ki] * fVec::loadu(w2 + k);
363
+ d2_hi += x_cache_hi[ki] * fVec::loadu(w2 + k + kFVecSize);
364
+ d3_lo += x_cache_lo[ki] * fVec::loadu(w3 + k);
365
+ d3_hi += x_cache_hi[ki] * fVec::loadu(w3 + k + kFVecSize);
366
+ }
367
+
368
+ fVec x_tail_lo(0.f), x_tail_hi(0.f);
369
+ if (rem > 0) {
370
+ std::tie(x_tail_lo, x_tail_hi) = at::vec::convert_to_float(bVec::loadu(x_t + k_tail, rem));
371
+ sq_lo += x_tail_lo * x_tail_lo;
372
+ if (rem1 > 0) sq_hi += x_tail_hi * x_tail_hi;
373
+ d0_lo += x_tail_lo * fVec::loadu(w0 + k_tail, rem0);
374
+ d1_lo += x_tail_lo * fVec::loadu(w1 + k_tail, rem0);
375
+ d2_lo += x_tail_lo * fVec::loadu(w2 + k_tail, rem0);
376
+ d3_lo += x_tail_lo * fVec::loadu(w3 + k_tail, rem0);
377
+ if (rem1 > 0) {
378
+ d0_hi += x_tail_hi * fVec::loadu(w0 + k_tail + kFVecSize, rem1);
379
+ d1_hi += x_tail_hi * fVec::loadu(w1 + k_tail + kFVecSize, rem1);
380
+ d2_hi += x_tail_hi * fVec::loadu(w2 + k_tail + kFVecSize, rem1);
381
+ d3_hi += x_tail_hi * fVec::loadu(w3 + k_tail + kFVecSize, rem1);
382
+ }
383
+ }
384
+
385
+ partial_row[0] = vec_reduce_sum(sq_lo + sq_hi);
386
+ partial_row[1 + 0] = vec_reduce_sum(d0_lo + d0_hi);
387
+ partial_row[1 + 1] = vec_reduce_sum(d1_lo + d1_hi);
388
+ partial_row[1 + 2] = vec_reduce_sum(d2_lo + d2_hi);
389
+ partial_row[1 + 3] = vec_reduce_sum(d3_lo + d3_hi);
390
+
391
+ // h=4..mix_hc-1: reuse cached x, four dots per pass.
392
+ int h = 4;
393
+ for (; h + 3 < mix_hc; h += 4) {
394
+ const float* wa = w_ptr[h];
395
+ const float* wb = w_ptr[h + 1];
396
+ const float* wc = w_ptr[h + 2];
397
+ const float* wd = w_ptr[h + 3];
398
+ fVec a0(0.f), a1(0.f), b0(0.f), b1(0.f);
399
+ fVec c0(0.f), c1(0.f), d0(0.f), d1(0.f);
400
+ for (int64_t ki = 0; ki < n_full; ++ki) {
401
+ const int64_t k = ki * kVecSize;
402
+ a0 += x_cache_lo[ki] * fVec::loadu(wa + k);
403
+ b0 += x_cache_lo[ki] * fVec::loadu(wb + k);
404
+ c0 += x_cache_lo[ki] * fVec::loadu(wc + k);
405
+ d0 += x_cache_lo[ki] * fVec::loadu(wd + k);
406
+ a1 += x_cache_hi[ki] * fVec::loadu(wa + k + kFVecSize);
407
+ b1 += x_cache_hi[ki] * fVec::loadu(wb + k + kFVecSize);
408
+ c1 += x_cache_hi[ki] * fVec::loadu(wc + k + kFVecSize);
409
+ d1 += x_cache_hi[ki] * fVec::loadu(wd + k + kFVecSize);
410
+ }
411
+ if (rem > 0) {
412
+ a0 += x_tail_lo * fVec::loadu(wa + k_tail, rem0);
413
+ b0 += x_tail_lo * fVec::loadu(wb + k_tail, rem0);
414
+ c0 += x_tail_lo * fVec::loadu(wc + k_tail, rem0);
415
+ d0 += x_tail_lo * fVec::loadu(wd + k_tail, rem0);
416
+ if (rem1 > 0) {
417
+ a1 += x_tail_hi * fVec::loadu(wa + k_tail + kFVecSize, rem1);
418
+ b1 += x_tail_hi * fVec::loadu(wb + k_tail + kFVecSize, rem1);
419
+ c1 += x_tail_hi * fVec::loadu(wc + k_tail + kFVecSize, rem1);
420
+ d1 += x_tail_hi * fVec::loadu(wd + k_tail + kFVecSize, rem1);
421
+ }
422
+ }
423
+ partial_row[1 + h] = vec_reduce_sum(a0 + a1);
424
+ partial_row[1 + h + 1] = vec_reduce_sum(b0 + b1);
425
+ partial_row[1 + h + 2] = vec_reduce_sum(c0 + c1);
426
+ partial_row[1 + h + 3] = vec_reduce_sum(d0 + d1);
427
+ }
428
+ for (; h + 1 < mix_hc; h += 2) {
429
+ const float* wa = w_ptr[h];
430
+ const float* wb = w_ptr[h + 1];
431
+ fVec a0(0.f), a1(0.f), b0(0.f), b1(0.f);
432
+ for (int64_t ki = 0; ki < n_full; ++ki) {
433
+ const int64_t k = ki * kVecSize;
434
+ a0 += x_cache_lo[ki] * fVec::loadu(wa + k);
435
+ b0 += x_cache_lo[ki] * fVec::loadu(wb + k);
436
+ a1 += x_cache_hi[ki] * fVec::loadu(wa + k + kFVecSize);
437
+ b1 += x_cache_hi[ki] * fVec::loadu(wb + k + kFVecSize);
438
+ }
439
+ if (rem > 0) {
440
+ a0 += x_tail_lo * fVec::loadu(wa + k_tail, rem0);
441
+ b0 += x_tail_lo * fVec::loadu(wb + k_tail, rem0);
442
+ if (rem1 > 0) {
443
+ a1 += x_tail_hi * fVec::loadu(wa + k_tail + kFVecSize, rem1);
444
+ b1 += x_tail_hi * fVec::loadu(wb + k_tail + kFVecSize, rem1);
445
+ }
446
+ }
447
+ partial_row[1 + h] = vec_reduce_sum(a0 + a1);
448
+ partial_row[1 + h + 1] = vec_reduce_sum(b0 + b1);
449
+ }
450
+ // odd leftover
451
+ if (h < mix_hc) {
452
+ const float* w = w_ptr[h];
453
+ fVec hd0(0.f), hd1(0.f);
454
+ for (int64_t ki = 0; ki < n_full; ++ki) {
455
+ const int64_t k = ki * kVecSize;
456
+ hd0 += x_cache_lo[ki] * fVec::loadu(w + k);
457
+ hd1 += x_cache_hi[ki] * fVec::loadu(w + k + kFVecSize);
458
+ }
459
+ if (rem > 0) {
460
+ hd0 += x_tail_lo * fVec::loadu(w + k_tail, rem0);
461
+ if (rem1 > 0) hd1 += x_tail_hi * fVec::loadu(w + k_tail + kFVecSize, rem1);
462
+ }
463
+ partial_row[1 + h] = vec_reduce_sum(hd0 + hd1);
464
+ }
465
+ }
466
+ } // tl
467
+
468
+ // Prefetch next t-block head.
469
+ if (tb + 1 < tb1) {
470
+ const scalar_t* x_next_tb = x + ((tb + 1) * T_BLOCK) * hc_d + k0;
471
+ __builtin_prefetch(x_next_tb, 0, 0);
472
+ if (klen > kVecSize * 4) __builtin_prefetch(x_next_tb + kVecSize * 4, 0, 0);
473
+ }
474
+ } // tb
475
+ } // kb
476
+ }); // Phase 1
477
+
478
+ // Phase 2+3: reduce partials, sinkhorn, and combine.
479
+ at::parallel_for(0, T, 1, [&](int64_t begin, int64_t end) {
480
+ const int64_t tid = at::get_thread_num();
481
+ fVec v0, v1, pre_fvec;
482
+ float* const thread_scratch = scratch_base + tid * per_thread_scratch;
483
+ float* const mix_buffer = thread_scratch;
484
+ float* const pre_gate = thread_scratch + mix_hc;
485
+ float* const post_gate = pre_gate + hc;
486
+ float* const comb_matrix = post_gate + hc;
487
+ float* const sinkhorn_scratch = comb_matrix + hc * hc;
488
+
489
+ for (int64_t t = begin; t < end; ++t) {
490
+ const int64_t tb = t / T_BLOCK;
491
+ const int64_t tl = t % T_BLOCK;
492
+
493
+ std::fill(mix_buffer, mix_buffer + mix_hc, 0.f);
494
+ double sq_total = 0.0;
495
+ for (int64_t kb = 0; kb < n_k_blocks; ++kb) {
496
+ const float* partial_row = partials + (tb * n_k_blocks + kb) * T_BLOCK * partial_stride + tl * partial_stride;
497
+ sq_total += static_cast<double>(partial_row[0]);
498
+#pragma GCC unroll 4
499
+ for (int i = 0; i < mix_hc; ++i)
500
+ mix_buffer[i] += partial_row[1 + i];
501
+ }
502
+ const float inv_rms =
503
+ static_cast<float>(1.0 / std::sqrt(sq_total / static_cast<double>(hc_d) + static_cast<double>(norm_eps)));
504
+#pragma GCC unroll 4
505
+ for (int i = 0; i < mix_hc; ++i)
506
+ mix_buffer[i] *= inv_rms;
507
+
508
+ parse_mixes_and_sinkhorn(
509
+ pre_gate,
510
+ post_gate,
511
+ comb_matrix,
512
+ mix_buffer,
513
+ s0,
514
+ s1,
515
+ s2,
516
+ hc_base,
517
+ sinkhorn_iters,
518
+ hc_eps,
519
+ hc,
520
+ sinkhorn_scratch);
521
+
522
+ std::copy_n(post_gate, hc, post_out + t * hc);
523
+ std::copy_n(comb_matrix, hc * hc, comb_out + t * hc * hc);
524
+
525
+ // Phase 3: d-tiled combine.
526
+ const scalar_t* x_t = x + t * hc_d;
527
+ scalar_t* y_t = y + t * d;
528
+ for (int64_t j0 = 0; j0 < d; j0 += D_BLOCK) {
529
+ const int64_t jlen = std::min(D_BLOCK, d - j0);
530
+
531
+ int64_t kk = 0;
532
+ for (; kk <= jlen - kVecSize; kk += kVecSize) {
533
+ pre_fvec = fVec(pre_gate[0]);
534
+ std::tie(v0, v1) = at::vec::convert_to_float(bVec::loadu(x_t + j0 + kk));
535
+ fVec acc0 = pre_fvec * v0;
536
+ fVec acc1 = pre_fvec * v1;
537
+#pragma GCC unroll 4
538
+ for (int h = 1; h < hc; ++h) {
539
+ pre_fvec = fVec(pre_gate[h]);
540
+ const scalar_t* x_th = x_t + (int64_t)h * d + j0 + kk;
541
+ std::tie(v0, v1) = at::vec::convert_to_float(bVec::loadu(x_th));
542
+ acc0 += pre_fvec * v0;
543
+ acc1 += pre_fvec * v1;
544
+ }
545
+ at::vec::convert_from_float<scalar_t>(acc0, acc1).store(y_t + j0 + kk);
546
+ }
547
+ if (kk < jlen) {
548
+ const int64_t rem = jlen - kk, rem1 = rem - std::min(rem, kFVecSize);
549
+ pre_fvec = fVec(pre_gate[0]);
550
+ std::tie(v0, v1) = at::vec::convert_to_float(bVec::loadu(x_t + j0 + kk, rem));
551
+ fVec acc0 = pre_fvec * v0;
552
+ fVec acc1 = pre_fvec * v1;
553
+#pragma GCC unroll 4
554
+ for (int h = 1; h < hc; ++h) {
555
+ pre_fvec = fVec(pre_gate[h]);
556
+ const scalar_t* x_th = x_t + (int64_t)h * d + j0 + kk;
557
+ std::tie(v0, v1) = at::vec::convert_to_float(bVec::loadu(x_th, rem));
558
+ acc0 += pre_fvec * v0;
559
+ acc1 += pre_fvec * v1;
560
+ }
561
+ at::vec::convert_from_float<scalar_t>(acc0, rem1 > 0 ? acc1 : fVec(0.f)).store(y_t + j0 + kk, rem);
562
+ }
563
+ } // d-tile
564
+ } // t
565
+ }); // Phase 2+3
566
+}
567
+
568
+// Fused hc_head path.
569
+template <typename scalar_t>
570
+static void hc_head_fuse_impl(
571
+ scalar_t* __restrict__ y,
572
+ const scalar_t* __restrict__ x,
573
+ const float* __restrict__ hc_fn, // [hc, hc*d] row-major
574
+ float hc_scale_val,
575
+ const float* __restrict__ hc_base,
576
+ int64_t T,
577
+ int64_t d,
578
+ int hc,
579
+ float hc_eps,
580
+ float norm_eps) {
581
+ using bVec = at::vec::Vectorized<scalar_t>;
582
+ using fVec = at::vec::Vectorized<float>;
583
+ constexpr int64_t kVecSize = bVec::size(); // 32 bf16
584
+ constexpr int64_t kFVecSize = fVec::size(); // 16 fp32
585
+ constexpr int64_t kDVecsPerBlock = 16;
586
+ const int64_t hc_d = hc * d;
587
+
588
+ constexpr int64_t K_BLOCK_CAP = 1024;
589
+ constexpr int64_t MAX_K_VECS = K_BLOCK_CAP / kVecSize;
590
+
591
+ const int64_t T_BLOCK = choose_t_block(T);
592
+ const int64_t num_threads = static_cast<int64_t>(at::get_num_threads());
593
+ const int64_t K_BLOCK = choose_k_block(hc, kVecSize, K_BLOCK_CAP);
594
+ TORCH_CHECK(
595
+ K_BLOCK <= K_BLOCK_CAP && K_BLOCK % kVecSize == 0,
596
+ "hc_head_fuse_impl: K_BLOCK must be <= K_BLOCK_CAP and a multiple of vector size");
597
+ const int64_t D_BLOCK = kDVecsPerBlock * kVecSize;
598
+
599
+ const int64_t n_t_blocks = div_up(T, T_BLOCK);
600
+ const int64_t n_k_blocks = div_up(hc_d, K_BLOCK);
601
+
602
+ // Partial buffer: [n_t_blocks, n_k_blocks, T_BLOCK, 1+hc]
603
+ // p[0] = partial sq, p[1..hc] = partial dots
604
+ const int64_t partial_stride = 1 + hc;
605
+ auto partials_tensor = at::empty({n_t_blocks * n_k_blocks * T_BLOCK * partial_stride}, at::kFloat);
606
+ float* const partials = partials_tensor.data_ptr<float>();
607
+
608
+ // Per-thread scratch: [pre_gate[hc]]
609
+ const int64_t per_thread_head = hc;
610
+ auto scratch_tensor = at::empty({num_threads * per_thread_head}, at::kFloat);
611
+ float* const scratch_base = scratch_tensor.data_ptr<float>();
612
+
613
+ // Phase 1.
614
+ parallel_mhc_phase1(n_t_blocks, n_k_blocks, [&](int64_t tb0, int64_t tb1, int64_t kb0, int64_t kb1) {
615
+ for (int64_t kb = kb0; kb < kb1; ++kb) {
616
+ const int64_t k0 = kb * K_BLOCK;
617
+ const int64_t klen = std::min(K_BLOCK, hc_d - k0);
618
+
619
+ for (int64_t tb = tb0; tb < tb1; ++tb) {
620
+ const int64_t t0 = tb * T_BLOCK;
621
+ const int64_t tlen = std::min(T_BLOCK, T - t0);
622
+ float* partial_block = partials + (tb * n_k_blocks + kb) * T_BLOCK * partial_stride;
623
+
624
+ const float* w_ptr[hc];
625
+ for (int h = 0; h < hc; ++h)
626
+ w_ptr[h] = hc_fn + h * hc_d + k0;
627
+
628
+ // Stack cache is sized for the largest selected K block, but only the first
629
+ // n_full entries are touched for the current K_BLOCK.
630
+ fVec x_cache_lo[MAX_K_VECS], x_cache_hi[MAX_K_VECS];
631
+
632
+ for (int64_t tl = 0; tl < tlen; ++tl) {
633
+ const scalar_t* x_t = x + (t0 + tl) * hc_d + k0;
634
+ float* partial_row = partial_block + tl * partial_stride;
635
+
636
+ const int64_t n_full = klen / kVecSize;
637
+ const int64_t k_tail = n_full * kVecSize;
638
+ const int64_t rem = klen - k_tail;
639
+ const int64_t rem0 = std::min(rem, kFVecSize);
640
+ const int64_t rem1 = rem - rem0;
641
+
642
+ // Prefetch next token.
643
+ if (tl + 1 < tlen) {
644
+ const scalar_t* xn = x + (t0 + tl + 1) * hc_d + k0;
645
+ __builtin_prefetch(xn, 0, 0);
646
+ if (klen > kVecSize * 4) __builtin_prefetch(xn + kVecSize * 4, 0, 0);
647
+ }
648
+
649
+ // Convert/cache x and accumulate sq + multiple dots.
650
+ if (tl == 0) {
651
+ for (int pw = 1; pw < hc; ++pw)
652
+ __builtin_prefetch(w_ptr[pw], 0, 3);
653
+ }
654
+ {
655
+ fVec x_tail_lo(0.f), x_tail_hi(0.f);
656
+ int h = 0;
657
+
658
+ if (hc >= 4) {
659
+ const float* w0 = w_ptr[0];
660
+ const float* w1 = w_ptr[1];
661
+ const float* w2 = w_ptr[2];
662
+ const float* w3 = w_ptr[3];
663
+ fVec sq_lo(0.f), sq_hi(0.f);
664
+ fVec d0_lo(0.f), d0_hi(0.f);
665
+ fVec d1_lo(0.f), d1_hi(0.f);
666
+ fVec d2_lo(0.f), d2_hi(0.f);
667
+ fVec d3_lo(0.f), d3_hi(0.f);
668
+ for (int64_t ki = 0; ki < n_full; ++ki) {
669
+ const int64_t k = ki * kVecSize;
670
+ std::tie(x_cache_lo[ki], x_cache_hi[ki]) = at::vec::convert_to_float(bVec::loadu(x_t + k));
671
+ sq_lo += x_cache_lo[ki] * x_cache_lo[ki];
672
+ sq_hi += x_cache_hi[ki] * x_cache_hi[ki];
673
+ d0_lo += x_cache_lo[ki] * fVec::loadu(w0 + k);
674
+ d0_hi += x_cache_hi[ki] * fVec::loadu(w0 + k + kFVecSize);
675
+ d1_lo += x_cache_lo[ki] * fVec::loadu(w1 + k);
676
+ d1_hi += x_cache_hi[ki] * fVec::loadu(w1 + k + kFVecSize);
677
+ d2_lo += x_cache_lo[ki] * fVec::loadu(w2 + k);
678
+ d2_hi += x_cache_hi[ki] * fVec::loadu(w2 + k + kFVecSize);
679
+ d3_lo += x_cache_lo[ki] * fVec::loadu(w3 + k);
680
+ d3_hi += x_cache_hi[ki] * fVec::loadu(w3 + k + kFVecSize);
681
+ }
682
+ if (rem > 0) {
683
+ std::tie(x_tail_lo, x_tail_hi) = at::vec::convert_to_float(bVec::loadu(x_t + k_tail, rem));
684
+ sq_lo += x_tail_lo * x_tail_lo;
685
+ if (rem1 > 0) sq_hi += x_tail_hi * x_tail_hi;
686
+ d0_lo += x_tail_lo * fVec::loadu(w0 + k_tail, rem0);
687
+ d1_lo += x_tail_lo * fVec::loadu(w1 + k_tail, rem0);
688
+ d2_lo += x_tail_lo * fVec::loadu(w2 + k_tail, rem0);
689
+ d3_lo += x_tail_lo * fVec::loadu(w3 + k_tail, rem0);
690
+ if (rem1 > 0) {
691
+ d0_hi += x_tail_hi * fVec::loadu(w0 + k_tail + kFVecSize, rem1);
692
+ d1_hi += x_tail_hi * fVec::loadu(w1 + k_tail + kFVecSize, rem1);
693
+ d2_hi += x_tail_hi * fVec::loadu(w2 + k_tail + kFVecSize, rem1);
694
+ d3_hi += x_tail_hi * fVec::loadu(w3 + k_tail + kFVecSize, rem1);
695
+ }
696
+ }
697
+ partial_row[0] = vec_reduce_sum(sq_lo + sq_hi);
698
+ partial_row[1 + 0] = vec_reduce_sum(d0_lo + d0_hi);
699
+ partial_row[1 + 1] = vec_reduce_sum(d1_lo + d1_hi);
700
+ partial_row[1 + 2] = vec_reduce_sum(d2_lo + d2_hi);
701
+ partial_row[1 + 3] = vec_reduce_sum(d3_lo + d3_hi);
702
+ h = 4;
703
+ } else {
704
+ const float* w0 = w_ptr[0];
705
+ fVec sq_lo(0.f), sq_hi(0.f);
706
+ fVec d0(0.f), d1(0.f);
707
+ for (int64_t ki = 0; ki < n_full; ++ki) {
708
+ const int64_t k = ki * kVecSize;
709
+ std::tie(x_cache_lo[ki], x_cache_hi[ki]) = at::vec::convert_to_float(bVec::loadu(x_t + k));
710
+ sq_lo += x_cache_lo[ki] * x_cache_lo[ki];
711
+ sq_hi += x_cache_hi[ki] * x_cache_hi[ki];
712
+ d0 += x_cache_lo[ki] * fVec::loadu(w0 + k);
713
+ d1 += x_cache_hi[ki] * fVec::loadu(w0 + k + kFVecSize);
714
+ }
715
+ if (rem > 0) {
716
+ std::tie(x_tail_lo, x_tail_hi) = at::vec::convert_to_float(bVec::loadu(x_t + k_tail, rem));
717
+ sq_lo += x_tail_lo * x_tail_lo;
718
+ if (rem1 > 0) sq_hi += x_tail_hi * x_tail_hi;
719
+ d0 += x_tail_lo * fVec::loadu(w0 + k_tail, rem0);
720
+ if (rem1 > 0) d1 += x_tail_hi * fVec::loadu(w0 + k_tail + kFVecSize, rem1);
721
+ }
722
+ partial_row[0] = vec_reduce_sum(sq_lo + sq_hi);
723
+ partial_row[1 + 0] = vec_reduce_sum(d0 + d1);
724
+ h = 1;
725
+ }
726
+
727
+ // h: reuse cached x, four dots per pass.
728
+ for (; h + 3 < hc; h += 4) {
729
+ const float* wa = w_ptr[h];
730
+ const float* wb = w_ptr[h + 1];
731
+ const float* wc = w_ptr[h + 2];
732
+ const float* wd = w_ptr[h + 3];
733
+ fVec a0(0.f), a1(0.f), b0(0.f), b1(0.f);
734
+ fVec c0(0.f), c1(0.f), d0(0.f), d1(0.f);
735
+ for (int64_t ki = 0; ki < n_full; ++ki) {
736
+ const int64_t k = ki * kVecSize;
737
+ a0 += x_cache_lo[ki] * fVec::loadu(wa + k);
738
+ b0 += x_cache_lo[ki] * fVec::loadu(wb + k);
739
+ c0 += x_cache_lo[ki] * fVec::loadu(wc + k);
740
+ d0 += x_cache_lo[ki] * fVec::loadu(wd + k);
741
+ a1 += x_cache_hi[ki] * fVec::loadu(wa + k + kFVecSize);
742
+ b1 += x_cache_hi[ki] * fVec::loadu(wb + k + kFVecSize);
743
+ c1 += x_cache_hi[ki] * fVec::loadu(wc + k + kFVecSize);
744
+ d1 += x_cache_hi[ki] * fVec::loadu(wd + k + kFVecSize);
745
+ }
746
+ if (rem > 0) {
747
+ a0 += x_tail_lo * fVec::loadu(wa + k_tail, rem0);
748
+ b0 += x_tail_lo * fVec::loadu(wb + k_tail, rem0);
749
+ c0 += x_tail_lo * fVec::loadu(wc + k_tail, rem0);
750
+ d0 += x_tail_lo * fVec::loadu(wd + k_tail, rem0);
751
+ if (rem1 > 0) {
752
+ a1 += x_tail_hi * fVec::loadu(wa + k_tail + kFVecSize, rem1);
753
+ b1 += x_tail_hi * fVec::loadu(wb + k_tail + kFVecSize, rem1);
754
+ c1 += x_tail_hi * fVec::loadu(wc + k_tail + kFVecSize, rem1);
755
+ d1 += x_tail_hi * fVec::loadu(wd + k_tail + kFVecSize, rem1);
756
+ }
757
+ }
758
+ partial_row[1 + h] = vec_reduce_sum(a0 + a1);
759
+ partial_row[1 + h + 1] = vec_reduce_sum(b0 + b1);
760
+ partial_row[1 + h + 2] = vec_reduce_sum(c0 + c1);
761
+ partial_row[1 + h + 3] = vec_reduce_sum(d0 + d1);
762
+ }
763
+ for (; h + 1 < hc; h += 2) {
764
+ const float* wa = w_ptr[h];
765
+ const float* wb = w_ptr[h + 1];
766
+ fVec a0(0.f), a1(0.f), b0(0.f), b1(0.f);
767
+ for (int64_t ki = 0; ki < n_full; ++ki) {
768
+ const int64_t k = ki * kVecSize;
769
+ a0 += x_cache_lo[ki] * fVec::loadu(wa + k);
770
+ b0 += x_cache_lo[ki] * fVec::loadu(wb + k);
771
+ a1 += x_cache_hi[ki] * fVec::loadu(wa + k + kFVecSize);
772
+ b1 += x_cache_hi[ki] * fVec::loadu(wb + k + kFVecSize);
773
+ }
774
+ if (rem > 0) {
775
+ a0 += x_tail_lo * fVec::loadu(wa + k_tail, rem0);
776
+ b0 += x_tail_lo * fVec::loadu(wb + k_tail, rem0);
777
+ if (rem1 > 0) {
778
+ a1 += x_tail_hi * fVec::loadu(wa + k_tail + kFVecSize, rem1);
779
+ b1 += x_tail_hi * fVec::loadu(wb + k_tail + kFVecSize, rem1);
780
+ }
781
+ }
782
+ partial_row[1 + h] = vec_reduce_sum(a0 + a1);
783
+ partial_row[1 + h + 1] = vec_reduce_sum(b0 + b1);
784
+ }
785
+ // odd leftover
786
+ if (h < hc) {
787
+ const float* w = w_ptr[h];
788
+ fVec hd0(0.f), hd1(0.f);
789
+ for (int64_t ki = 0; ki < n_full; ++ki) {
790
+ const int64_t k = ki * kVecSize;
791
+ hd0 += x_cache_lo[ki] * fVec::loadu(w + k);
792
+ hd1 += x_cache_hi[ki] * fVec::loadu(w + k + kFVecSize);
793
+ }
794
+ if (rem > 0) {
795
+ hd0 += x_tail_lo * fVec::loadu(w + k_tail, rem0);
796
+ if (rem1 > 0) hd1 += x_tail_hi * fVec::loadu(w + k_tail + kFVecSize, rem1);
797
+ }
798
+ partial_row[1 + h] = vec_reduce_sum(hd0 + hd1);
799
+ }
800
+ }
801
+ } // tl
802
+
803
+ // Prefetch next t-block head.
804
+ if (tb + 1 < tb1) {
805
+ const scalar_t* x_next_tb = x + ((tb + 1) * T_BLOCK) * hc_d + k0;
806
+ __builtin_prefetch(x_next_tb, 0, 0);
807
+ if (klen > kVecSize * 2) __builtin_prefetch(x_next_tb + kVecSize * 2, 0, 0);
808
+ }
809
+ } // tb
810
+ } // kb
811
+ }); // Phase 1
812
+
813
+ // Phase 2+3.
814
+ at::parallel_for(0, T, 1, [&](int64_t begin, int64_t end) {
815
+ const int64_t tid = at::get_thread_num();
816
+ float* const pre_gate = scratch_base + tid * per_thread_head;
817
+ fVec v0, v1, pre_fvec;
818
+
819
+ for (int64_t t = begin; t < end; ++t) {
820
+ const int64_t tb = t / T_BLOCK;
821
+ const int64_t tl = t % T_BLOCK;
822
+
823
+ // Phase 2: reduce partials + sigmoid gate.
824
+ double sq_total = 0.0;
825
+ std::fill(pre_gate, pre_gate + hc, 0.f);
826
+ for (int64_t kb = 0; kb < n_k_blocks; ++kb) {
827
+ const float* partial_row = partials + (tb * n_k_blocks + kb) * T_BLOCK * partial_stride + tl * partial_stride;
828
+ sq_total += static_cast<double>(partial_row[0]);
829
+#pragma GCC unroll 4
830
+ for (int h = 0; h < hc; ++h)
831
+ pre_gate[h] += partial_row[1 + h];
832
+ }
833
+ const float inv_rms =
834
+ static_cast<float>(1.0 / std::sqrt(sq_total / static_cast<double>(hc_d) + static_cast<double>(norm_eps)));
835
+#pragma GCC unroll 4
836
+ for (int h = 0; h < hc; ++h) {
837
+ const float gate = pre_gate[h] * inv_rms * hc_scale_val + hc_base[h];
838
+ pre_gate[h] = 1.f / (1.f + std::exp(-gate)) + hc_eps;
839
+ }
840
+
841
+ // Phase 3: d-tiled combine.
842
+ const scalar_t* x_t = x + t * hc_d;
843
+ scalar_t* y_t = y + t * d;
844
+ for (int64_t j0 = 0; j0 < d; j0 += D_BLOCK) {
845
+ const int64_t jlen = std::min(D_BLOCK, d - j0);
846
+
847
+ int64_t kk = 0;
848
+ for (; kk <= jlen - kVecSize; kk += kVecSize) {
849
+ pre_fvec = fVec(pre_gate[0]);
850
+ std::tie(v0, v1) = at::vec::convert_to_float(bVec::loadu(x_t + j0 + kk));
851
+ fVec acc0 = pre_fvec * v0;
852
+ fVec acc1 = pre_fvec * v1;
853
+#pragma GCC unroll 4
854
+ for (int h = 1; h < hc; ++h) {
855
+ pre_fvec = fVec(pre_gate[h]);
856
+ const scalar_t* x_th = x_t + (int64_t)h * d + j0 + kk;
857
+ std::tie(v0, v1) = at::vec::convert_to_float(bVec::loadu(x_th));
858
+ acc0 += pre_fvec * v0;
859
+ acc1 += pre_fvec * v1;
860
+ }
861
+ at::vec::convert_from_float<scalar_t>(acc0, acc1).store(y_t + j0 + kk);
862
+ }
863
+ if (kk < jlen) {
864
+ const int64_t rem = jlen - kk, rem1 = rem - std::min(rem, kFVecSize);
865
+ pre_fvec = fVec(pre_gate[0]);
866
+ std::tie(v0, v1) = at::vec::convert_to_float(bVec::loadu(x_t + j0 + kk, rem));
867
+ fVec acc0 = pre_fvec * v0;
868
+ fVec acc1 = pre_fvec * v1;
869
+#pragma GCC unroll 4
870
+ for (int h = 1; h < hc; ++h) {
871
+ pre_fvec = fVec(pre_gate[h]);
872
+ const scalar_t* x_th = x_t + (int64_t)h * d + j0 + kk;
873
+ std::tie(v0, v1) = at::vec::convert_to_float(bVec::loadu(x_th, rem));
874
+ acc0 += pre_fvec * v0;
875
+ acc1 += pre_fvec * v1;
876
+ }
877
+ at::vec::convert_from_float<scalar_t>(acc0, rem1 > 0 ? acc1 : fVec(0.f)).store(y_t + j0 + kk, rem);
878
+ }
879
+ } // d-tile
880
+ } // t
881
+ }); // Phase 2+3
882
+}
883
+
884
+} // anonymous namespace
885
+
886
+// Fused hc_pre entry.
887
+std::tuple<at::Tensor, at::Tensor, at::Tensor> hc_pre_fused_cpu(
888
+ at::Tensor& x,
889
+ at::Tensor& hc_fn,
890
+ at::Tensor& hc_scale,
891
+ at::Tensor& hc_base,
892
+ int64_t hc_mult,
893
+ int64_t sinkhorn_iters,
894
+ double rms_eps,
895
+ double hc_eps) {
896
+ TORCH_CHECK(hc_mult >= 2, "hc_pre_fused_cpu: hc_mult must be >= 2");
897
+ TORCH_CHECK(
898
+ x.dim() == 3 && (x.scalar_type() == at::kBFloat16 || x.scalar_type() == at::kHalf),
899
+ "hc_pre_fused_cpu: x must be bf16/fp16 [T, hc, d]");
900
+ TORCH_CHECK(hc_fn.scalar_type() == at::kFloat, "hc_pre_fused_cpu: hc_fn must be float32");
901
+ TORCH_CHECK(
902
+ hc_scale.scalar_type() == at::kFloat && hc_scale.numel() == 3, "hc_pre_fused_cpu: hc_scale must be float32 [3]");
903
+ TORCH_CHECK(hc_base.scalar_type() == at::kFloat, "hc_pre_fused_cpu: hc_base must be float32");
904
+ TORCH_CHECK(
905
+ x.device().is_cpu() && hc_fn.device().is_cpu() && hc_scale.device().is_cpu() && hc_base.device().is_cpu(),
906
+ "hc_pre_fused_cpu: all inputs must be CPU tensors");
907
+
908
+ const int64_t T = x.size(0);
909
+ const int64_t d = x.size(2);
910
+ const int64_t hc_d = hc_mult * d;
911
+ const int64_t mix_hc = (2 + hc_mult) * hc_mult;
912
+
913
+ TORCH_CHECK(T > 0 && d > 0, "hc_pre_fused_cpu: T and d must be positive");
914
+ TORCH_CHECK(x.size(1) == hc_mult, "hc_pre_fused_cpu: x shape must be [T, hc_mult, d]");
915
+ TORCH_CHECK(
916
+ hc_fn.dim() == 2 && hc_fn.size(0) == mix_hc && hc_fn.size(1) == hc_d,
917
+ "hc_pre_fused_cpu: hc_fn must be [mix_hc, hc_mult*d]");
918
+ TORCH_CHECK(hc_base.numel() == mix_hc, "hc_pre_fused_cpu: hc_base.numel() must equal mix_hc");
919
+
920
+ // Static per-layer weights (hc_fn/hc_scale/hc_base) are allocated fp32 and
921
+ // never reassigned after weight loading; x/residual come straight from a
922
+ // preceding op. All are expected contiguous already -- fail loudly instead
923
+ // of silently absorbing a copy if that contract is ever violated.
924
+ TORCH_CHECK(
925
+ x.is_contiguous() && hc_fn.is_contiguous() && hc_scale.is_contiguous() && hc_base.is_contiguous(),
926
+ "hc_pre_fused_cpu: x/hc_fn/hc_scale/hc_base must be contiguous");
927
+
928
+ auto f32_opts = at::TensorOptions().dtype(at::kFloat).device(x.device());
929
+ auto post = at::empty({T, hc_mult}, f32_opts);
930
+ auto comb = at::empty({T, hc_mult, hc_mult}, f32_opts);
931
+ auto y = at::empty({T, d}, x.options());
932
+
933
+ AT_DISPATCH_REDUCED_FLOATING_TYPES(x.scalar_type(), "hc_pre_fused_cpu", [&] {
934
+ hc_pre_fuse_impl<scalar_t>(
935
+ y.data_ptr<scalar_t>(),
936
+ post.data_ptr<float>(),
937
+ comb.data_ptr<float>(),
938
+ x.data_ptr<scalar_t>(),
939
+ hc_fn.data_ptr<float>(),
940
+ hc_scale.data_ptr<float>(),
941
+ hc_base.data_ptr<float>(),
942
+ T,
943
+ d,
944
+ static_cast<int>(hc_mult),
945
+ static_cast<int>(sinkhorn_iters),
946
+ static_cast<float>(hc_eps),
947
+ static_cast<float>(rms_eps));
948
+ });
949
+ return {y, post, comb};
950
+}
951
+
952
+// Fused hc_post entry.
953
+at::Tensor hc_post_fused_cpu(at::Tensor& x, at::Tensor& residual, at::Tensor& post, at::Tensor& comb) {
954
+ TORCH_CHECK(
955
+ x.dim() == 2 && (x.scalar_type() == at::kBFloat16 || x.scalar_type() == at::kHalf),
956
+ "hc_post_fused_cpu: x must be bf16/fp16 [T, d]");
957
+ TORCH_CHECK(
958
+ residual.dim() == 3 && residual.scalar_type() == x.scalar_type(),
959
+ "hc_post_fused_cpu: residual must be same dtype as x [T, hc, d]");
960
+ TORCH_CHECK(post.scalar_type() == at::kFloat, "hc_post_fused_cpu: post must be float32");
961
+ TORCH_CHECK(comb.scalar_type() == at::kFloat, "hc_post_fused_cpu: comb must be float32");
962
+ TORCH_CHECK(
963
+ x.device().is_cpu() && residual.device().is_cpu() && post.device().is_cpu() && comb.device().is_cpu(),
964
+ "hc_post_fused_cpu: all inputs must be CPU tensors");
965
+
966
+ const int64_t T = x.size(0);
967
+ const int64_t d = x.size(1);
968
+ const int64_t hc = residual.size(1);
969
+ TORCH_CHECK(hc >= 2, "hc_post_fused_cpu: hc_mult must be >= 2");
970
+ TORCH_CHECK(T > 0 && d > 0, "hc_post_fused_cpu: T and d must be positive");
971
+ TORCH_CHECK(residual.size(0) == T && residual.size(2) == d, "hc_post_fused_cpu: residual shape must be [T, hc, d]");
972
+ TORCH_CHECK(
973
+ post.dim() == 2 && post.size(0) == T && post.size(1) == hc, "hc_post_fused_cpu: post shape must be [T, hc]");
974
+ TORCH_CHECK(
975
+ comb.dim() == 3 && comb.size(0) == T && comb.size(1) == hc && comb.size(2) == hc,
976
+ "hc_post_fused_cpu: comb shape must be [T, hc, hc]");
977
+
978
+ TORCH_CHECK(
979
+ x.is_contiguous() && residual.is_contiguous() && post.is_contiguous() && comb.is_contiguous(),
980
+ "hc_post_fused_cpu: x/residual/post/comb must be contiguous");
981
+
982
+ auto out = at::empty({T, hc, d}, x.options()); // same dtype as x
983
+
984
+ AT_DISPATCH_REDUCED_FLOATING_TYPES(x.scalar_type(), "hc_post_fused_cpu", [&] {
985
+ hc_post_fuse_impl<scalar_t>(
986
+ out.data_ptr<scalar_t>(),
987
+ x.data_ptr<scalar_t>(),
988
+ residual.data_ptr<scalar_t>(),
989
+ post.data_ptr<float>(),
990
+ comb.data_ptr<float>(),
991
+ T,
992
+ d,
993
+ static_cast<int>(hc));
994
+ });
995
+
996
+ return out;
997
+}
998
+
999
+// Fused hc_head entry.
1000
+at::Tensor hc_head_fused_cpu(
1001
+ at::Tensor& x, at::Tensor& hc_fn, at::Tensor& hc_scale, at::Tensor& hc_base, double hc_eps, double norm_eps) {
1002
+ TORCH_CHECK(
1003
+ x.dim() == 3 && (x.scalar_type() == at::kBFloat16 || x.scalar_type() == at::kHalf),
1004
+ "hc_head_fused_cpu: x must be bf16/fp16 [T, hc, d]");
1005
+ const int64_t hc_mult = x.size(1);
1006
+ TORCH_CHECK(hc_mult >= 2, "hc_head_fused_cpu: hc_mult must be >= 2");
1007
+ TORCH_CHECK(hc_fn.scalar_type() == at::kFloat, "hc_head_fused_cpu: hc_fn must be float32");
1008
+ TORCH_CHECK(
1009
+ hc_scale.scalar_type() == at::kFloat && hc_scale.numel() == 1,
1010
+ "hc_head_fused_cpu: hc_scale must be a scalar float32");
1011
+ TORCH_CHECK(hc_base.scalar_type() == at::kFloat, "hc_head_fused_cpu: hc_base must be float32");
1012
+ TORCH_CHECK(hc_base.numel() == hc_mult, "hc_head_fused_cpu: hc_base.numel() must equal hc_mult");
1013
+ TORCH_CHECK(
1014
+ x.device().is_cpu() && hc_fn.device().is_cpu() && hc_scale.device().is_cpu() && hc_base.device().is_cpu(),
1015
+ "hc_head_fused_cpu: all inputs must be CPU tensors");
1016
+
1017
+ const int64_t T = x.size(0);
1018
+ const int64_t d = x.size(2);
1019
+ TORCH_CHECK(T > 0 && d > 0, "hc_head_fused_cpu: T and d must be positive");
1020
+ TORCH_CHECK(
1021
+ hc_fn.dim() == 2 && hc_fn.size(0) == hc_mult && hc_fn.size(1) == hc_mult * d,
1022
+ "hc_head_fused_cpu: hc_fn must be [hc_mult, hc_mult*d]");
1023
+
1024
+ TORCH_CHECK(
1025
+ x.is_contiguous() && hc_fn.is_contiguous() && hc_scale.is_contiguous() && hc_base.is_contiguous(),
1026
+ "hc_head_fused_cpu: x/hc_fn/hc_scale/hc_base must be contiguous");
1027
+ const float hc_scale_val = hc_scale.data_ptr<float>()[0];
1028
+
1029
+ auto y = at::empty({T, d}, x.options());
1030
+
1031
+ AT_DISPATCH_REDUCED_FLOATING_TYPES(x.scalar_type(), "hc_head_fused_cpu", [&] {
1032
+ hc_head_fuse_impl<scalar_t>(
1033
+ y.data_ptr<scalar_t>(),
1034
+ x.data_ptr<scalar_t>(),
1035
+ hc_fn.data_ptr<float>(),
1036
+ hc_scale_val,
1037
+ hc_base.data_ptr<float>(),
1038
+ T,
1039
+ d,
1040
+ static_cast<int>(hc_mult),
1041
+ static_cast<float>(hc_eps),
1042
+ static_cast<float>(norm_eps));
1043
+ });
1044
+
1045
+ return y;
1046
+}
csrc/cpu/sgl-kernels/mla_cache.cpp
+9
-9
@@ -1,14 +1,14 @@
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
//
3
-// `concat_and_cache_mla` (the generic MLA cache-write op used by every GPU
4
-// backend) is registered CUDA-only. This is the CPU counterpart, adapted in
5
-// spirit from SGLang's `store_cache_cpu` (csrc/cpu/kvcache.cpp) but
6
-// generalized to write two source tensors (`kv_c_normed`, `k_pe`) into two
7
-// different column-offset ranges of the SAME destination row -- SGLang's
8
-// version assumes k/v land in two independent, equal-row-width cache
9
-// tensors, which doesn't hold here since MLA's cache is one 576-wide buffer
10
-// and the two column ranges (512-wide, 64-wide) don't match the buffer's
11
-// true per-token stride, so the write can't reuse `store_cache_cpu` as-is.
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"
csrc/cpu/sgl-kernels/moe.cpp
+3
-4
@@ -871,13 +871,13 @@ static inline void check_moe_scales(
871
// topk_ids: [M, topk] (int32_t)
872
//
873
874
-at::Tensor fused_experts_cpu(
874
+void fused_experts_cpu(
875
+ at::Tensor& out,
876
at::Tensor& hidden_states,
877
at::Tensor& w1,
878
at::Tensor& w2,
879
at::Tensor& topk_weights,
880
at::Tensor& topk_ids,
880
- bool inplace,
881
int64_t moe_comp_method,
882
const std::optional<at::Tensor>& w1_scale,
883
const std::optional<at::Tensor>& w2_scale,
@@ -943,7 +943,7 @@ at::Tensor fused_experts_cpu(
943
// check scales
944
check_moe_scales(moe_comp_method, w1_scale, w2_scale, block_size);
945
946
- at::Tensor out_hidden_states = inplace ? hidden_states : at::empty_like(hidden_states);
946
+ at::Tensor& out_hidden_states = out;
947
948
// NB: worst case is each expert holds a block with remainder of 1
949
// 1. sorted_ids : [M * topk + E * (BLOCK_M - 1)]
@@ -1214,7 +1214,6 @@ at::Tensor fused_experts_cpu(
1214
with_bias);
1215
}
1216
});
1217
- return out_hidden_states;
1217
}
1218
1219
// shared expert kernel
csrc/cpu/sgl-kernels/paged_mqa_logits.cpp
new
+217
@@ -0,0 +1,217 @@
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
+}
csrc/cpu/sgl-kernels/store_cache.cpp
new
+481
@@ -0,0 +1,481 @@
1
+// Adapted from
2
+// https://github.com/sgl-project/sglang/blob/main/python/sglang/kernels/aot/csrc/cpu/store_cache.cpp
3
+//
4
+// DeepSeek-V4 fused Q RMSNorm+RoPE / KV RoPE+UE8M0-quant+paged-cache-insert
5
+// kernel for the fp8_ds_mla cache layout. Not a port of upstream's
6
+// `set_k_and_s_cpu`/`quant_to_nope_fp8_rope_bf16_pack_cpu` -- their
7
+// pool_kv/pool_score signature doesn't apply to vLLM's layout. This is the
8
+// CPU counterpart to vLLM's own CUDA kernel
9
+// (`fused_deepseek_v4_qnorm_rope_kv_insert_kernel.cu`); see
10
+// `DeepseekV4Attention._fused_qnorm_rope_kv_insert_cpu` (attention.py) for
11
+// the eager reference this replaces.
12
+//
13
+// clang-format off
14
+
15
+#include <algorithm>
16
+#include <cmath>
17
+#include <vector>
18
+
19
+#include "common.h"
20
+#include "vec.h"
21
+
22
+#if !defined(CPU_CAPABILITY_AVX512)
23
+static_assert(false,
24
+ "store_cache.cpp is only ever compiled into the AVX512+AMX-flagged "
25
+ "_C target (see cmake/cpu_extension.cmake's VLLM_EXT_SRC_SGL) and "
26
+ "relies on that unconditionally -- CPU_CAPABILITY_AVX512 must be "
27
+ "defined here.");
28
+#endif
29
+
30
+namespace {
31
+
32
+constexpr int64_t kRopeDim = 64;
33
+constexpr int64_t kQuantBlock = 64;
34
+constexpr float kFp8Max = 448.0f;
35
+constexpr float kFp8Min = -448.0f;
36
+constexpr float kQuantEps = 1e-4f;
37
+
38
+inline void apply_gptj_rope_inplace(float* __restrict__ x, const float* __restrict__ cos, const float* __restrict__ sin, int64_t half_dim) {
39
+ for (int64_t i = 0; i < half_dim; ++i) {
40
+ const float e = x[2 * i];
41
+ const float o = x[2 * i + 1];
42
+ x[2 * i] = e * cos[i] - o * sin[i];
43
+ x[2 * i + 1] = e * sin[i] + o * cos[i];
44
+ }
45
+}
46
+
47
+// AVX512 tile-quantization, adapted verbatim from SGLang's
48
+// `store_cache.cpp` (`quant_to_nope_fp8_rope_bf16_pack_cpu`'s
49
+// `cvt_fp32x16_to_fp8x16`/`quantize_tile_avx512`): per-64-element tile,
50
+// computes the UE8M0 power-of-two scale and quantizes to fp8_e4m3fn with
51
+// RNE rounding via bit manipulation instead of a per-element scalar
52
+// `at::Float8_e4m3fn` cast. Subnormal-magnitude fp8 lanes are flushed to
53
+// zero (matches SGLang's own convention, `cvt_fp32x16_to_fp8x16`) --
54
+// vLLM's own AVX512 dequant of this cache (`flash_mla.cpp`'s
55
+// `CVT_FP8_TO_BF16_EXT`) doesn't reconstruct exact subnormal magnitudes
56
+// either way, so there is nothing to preserve by rounding them exactly on
57
+// the write side.
58
+#if defined(CPU_CAPABILITY_AVX512)
59
+
60
+// Vectorized float32x16 -> fp8_e4m3fn x16 using AVX512 bit manipulation.
61
+// Implements RNE (round-to-nearest-even). Flushes subnormal fp8 to zero.
62
+// fp8_e4m3fn: sign(1), exp(4, bias=7), mant(3), no inf/nan special values.
63
+// Normal: exp_field in [1..15], value = 2^(exp-7) * (1 + mant/8)
64
+// Input assumed clamped to [-448, 448].
65
+inline __m128i cvt_fp32x16_to_fp8x16(__m512 input) {
66
+ __m512i bits = _mm512_castps_si512(input);
67
+
68
+ // Extract sign -> bit 7 of fp8 byte
69
+ __m512i signs = _mm512_srli_epi32(_mm512_and_si512(bits, _mm512_set1_epi32(0x80000000)), 24);
70
+
71
+ // Absolute value bits
72
+ __m512i abs_bits = _mm512_and_si512(bits, _mm512_set1_epi32(0x7FFFFFFF));
73
+
74
+ // Float32 biased exponent
75
+ __m512i f32_exp = _mm512_srli_epi32(abs_bits, 23);
76
+
77
+ // Float32 mantissa (23 bits)
78
+ __m512i f32_mant = _mm512_and_si512(abs_bits, _mm512_set1_epi32(0x7FFFFF));
79
+
80
+ // fp8_exp = f32_exp - 120 (rebias from 127 to 7)
81
+ __m512i fp8_exp = _mm512_sub_epi32(f32_exp, _mm512_set1_epi32(120));
82
+
83
+ // Top 3 mantissa bits
84
+ __m512i mant3 = _mm512_srli_epi32(f32_mant, 20);
85
+
86
+ // RNE rounding: round_bit = bit 19, sticky = bits[18:0]
87
+ __m512i round_bit = _mm512_and_si512(_mm512_srli_epi32(f32_mant, 19), _mm512_set1_epi32(1));
88
+ __m512i sticky_bits = _mm512_and_si512(f32_mant, _mm512_set1_epi32(0x7FFFF));
89
+ __mmask16 has_sticky = _mm512_cmpneq_epi32_mask(sticky_bits, _mm512_setzero_si512());
90
+
91
+ // Round up if round_bit AND (sticky OR lsb_of_mant3)
92
+ __m512i lsb = _mm512_and_si512(mant3, _mm512_set1_epi32(1));
93
+ __m512i sticky_or_lsb = _mm512_or_si512(_mm512_maskz_mov_epi32(has_sticky, _mm512_set1_epi32(1)), lsb);
94
+ __m512i do_round = _mm512_and_si512(round_bit, sticky_or_lsb);
95
+ mant3 = _mm512_add_epi32(mant3, do_round);
96
+
97
+ // Mantissa overflow (mant3 == 8) -> increment exponent, zero mantissa
98
+ __mmask16 mant_ovf = _mm512_cmpeq_epi32_mask(mant3, _mm512_set1_epi32(8));
99
+ mant3 = _mm512_mask_mov_epi32(mant3, mant_ovf, _mm512_setzero_si512());
100
+ fp8_exp = _mm512_mask_add_epi32(fp8_exp, mant_ovf, fp8_exp, _mm512_set1_epi32(1));
101
+
102
+ // Clamp exponent: max 15
103
+ fp8_exp = _mm512_min_epi32(fp8_exp, _mm512_set1_epi32(15));
104
+
105
+ // Result: (fp8_exp << 3) | mant3
106
+ __m512i result = _mm512_or_si512(_mm512_slli_epi32(fp8_exp, 3), mant3);
107
+
108
+ // Flush to zero: if f32_exp < 121 (would be subnormal in fp8), output 0
109
+ __mmask16 is_subnormal_or_zero = _mm512_cmplt_epi32_mask(f32_exp, _mm512_set1_epi32(121));
110
+ result = _mm512_mask_mov_epi32(result, is_subnormal_or_zero, _mm512_setzero_si512());
111
+
112
+ // Add sign
113
+ result = _mm512_or_si512(result, signs);
114
+
115
+ // Pack 16 x i32 -> 16 x i8
116
+ return _mm512_cvtepi32_epi8(result);
117
+}
118
+
119
+// Process one tile (64 bf16 values): find amax, compute scale, quantize to fp8.
120
+inline uint8_t quantize_tile_avx512(const at::BFloat16* __restrict__ src, uint8_t* __restrict__ dst) {
121
+ // Load 64 bf16 -> 4 x 16 floats
122
+ __m256i bf16_0 = _mm256_loadu_si256(reinterpret_cast<const __m256i*>(src));
123
+ __m256i bf16_1 = _mm256_loadu_si256(reinterpret_cast<const __m256i*>(src + 16));
124
+ __m256i bf16_2 = _mm256_loadu_si256(reinterpret_cast<const __m256i*>(src + 32));
125
+ __m256i bf16_3 = _mm256_loadu_si256(reinterpret_cast<const __m256i*>(src + 48));
126
+
127
+ __m512 f0 = _mm512_castsi512_ps(_mm512_slli_epi32(_mm512_cvtepu16_epi32(bf16_0), 16));
128
+ __m512 f1 = _mm512_castsi512_ps(_mm512_slli_epi32(_mm512_cvtepu16_epi32(bf16_1), 16));
129
+ __m512 f2 = _mm512_castsi512_ps(_mm512_slli_epi32(_mm512_cvtepu16_epi32(bf16_2), 16));
130
+ __m512 f3 = _mm512_castsi512_ps(_mm512_slli_epi32(_mm512_cvtepu16_epi32(bf16_3), 16));
131
+
132
+ // Absolute values
133
+ const __m512i abs_mask_i = _mm512_set1_epi32(0x7FFFFFFF);
134
+ __m512 abs0 = _mm512_castsi512_ps(_mm512_and_si512(_mm512_castps_si512(f0), abs_mask_i));
135
+ __m512 abs1 = _mm512_castsi512_ps(_mm512_and_si512(_mm512_castps_si512(f1), abs_mask_i));
136
+ __m512 abs2 = _mm512_castsi512_ps(_mm512_and_si512(_mm512_castps_si512(f2), abs_mask_i));
137
+ __m512 abs3 = _mm512_castsi512_ps(_mm512_and_si512(_mm512_castps_si512(f3), abs_mask_i));
138
+
139
+ // Horizontal max
140
+ __m512 max01 = _mm512_max_ps(abs0, abs1);
141
+ __m512 max23 = _mm512_max_ps(abs2, abs3);
142
+ __m512 max0123 = _mm512_max_ps(max01, max23);
143
+ float amax = _mm512_reduce_max_ps(max0123);
144
+
145
+ // Scale computation
146
+ float scale = std::max(amax / kFp8Max, kQuantEps);
147
+ float ceil_log2 = std::ceil(std::log2(scale));
148
+ float scale_pow2 = std::exp2(ceil_log2);
149
+ float scale_inv = 1.0f / scale_pow2;
150
+
151
+ int exponent = static_cast<int>(ceil_log2);
152
+ uint8_t scale_uint8 = static_cast<uint8_t>(exponent + 127);
153
+
154
+ // Scale all values
155
+ __m512 vinv = _mm512_set1_ps(scale_inv);
156
+ __m512 s0 = _mm512_mul_ps(f0, vinv);
157
+ __m512 s1 = _mm512_mul_ps(f1, vinv);
158
+ __m512 s2 = _mm512_mul_ps(f2, vinv);
159
+ __m512 s3 = _mm512_mul_ps(f3, vinv);
160
+
161
+ // Clamp
162
+ __m512 vmax = _mm512_set1_ps(kFp8Max);
163
+ __m512 vmin = _mm512_set1_ps(kFp8Min);
164
+ s0 = _mm512_max_ps(_mm512_min_ps(s0, vmax), vmin);
165
+ s1 = _mm512_max_ps(_mm512_min_ps(s1, vmax), vmin);
166
+ s2 = _mm512_max_ps(_mm512_min_ps(s2, vmax), vmin);
167
+ s3 = _mm512_max_ps(_mm512_min_ps(s3, vmax), vmin);
168
+
169
+ // Vectorized float -> fp8 conversion (16 values at a time)
170
+ __m128i fp8_0 = cvt_fp32x16_to_fp8x16(s0);
171
+ __m128i fp8_1 = cvt_fp32x16_to_fp8x16(s1);
172
+ __m128i fp8_2 = cvt_fp32x16_to_fp8x16(s2);
173
+ __m128i fp8_3 = cvt_fp32x16_to_fp8x16(s3);
174
+
175
+ _mm_storeu_si128(reinterpret_cast<__m128i*>(dst), fp8_0);
176
+ _mm_storeu_si128(reinterpret_cast<__m128i*>(dst + 16), fp8_1);
177
+ _mm_storeu_si128(reinterpret_cast<__m128i*>(dst + 32), fp8_2);
178
+ _mm_storeu_si128(reinterpret_cast<__m128i*>(dst + 48), fp8_3);
179
+
180
+ return scale_uint8;
181
+}
182
+
183
+#endif // CPU_CAPABILITY_AVX512
184
+
185
+template <typename scalar_t>
186
+static void qnorm_rope_kv_insert_impl(
187
+ scalar_t* __restrict__ q_out, // [num_tokens_full, q_head_padded, head_dim]
188
+ uint8_t* __restrict__ cache, // [num_blocks, cache_row_stride] bytes
189
+ const scalar_t* __restrict__ q, // [num_tokens_full, num_heads_q, head_dim]
190
+ const scalar_t* __restrict__ kv, // [>=num_tokens_insert, head_dim]
191
+ const int64_t* __restrict__ positions, // [num_tokens_full]
192
+ const int64_t* __restrict__ slot_mapping, // [num_tokens_insert]
193
+ const float* __restrict__ cos_sin_cache, // [max_pos, rope_dim] (first half cos, second half sin)
194
+ int64_t num_tokens_full,
195
+ int64_t num_heads_q,
196
+ int64_t q_head_padded,
197
+ int64_t head_dim,
198
+ int64_t nope_dim,
199
+ int64_t num_quant_blocks,
200
+ int64_t token_data_bytes,
201
+ int64_t scale_bytes_per_token,
202
+ int64_t num_tokens_insert,
203
+ int64_t cache_block_size,
204
+ int64_t cache_row_stride,
205
+ float eps) {
206
+ const int64_t rope_half = kRopeDim / 2;
207
+
208
+ // Q: per-head weight-free RMSNorm + GPT-J RoPE on the last rope_dim dims.
209
+ at::parallel_for(0, num_tokens_full * num_heads_q, 0, [&](int64_t begin, int64_t end) {
210
+ std::vector<float> buf(head_dim);
211
+ for (int64_t idx = begin; idx < end; ++idx) {
212
+ const int64_t t = idx / num_heads_q;
213
+ const int64_t h = idx % num_heads_q;
214
+ const scalar_t* q_th = q + (t * num_heads_q + h) * head_dim;
215
+
216
+ float sq_sum = 0.f;
217
+ for (int64_t i = 0; i < head_dim; ++i) {
218
+ const float v = static_cast<float>(q_th[i]);
219
+ buf[i] = v;
220
+ sq_sum += v * v;
221
+ }
222
+ const float inv_rms = 1.0f / std::sqrt(sq_sum / static_cast<float>(head_dim) + eps);
223
+ for (int64_t i = 0; i < head_dim; ++i)
224
+ buf[i] *= inv_rms;
225
+
226
+ const int64_t pos = positions[t];
227
+ const float* cos = cos_sin_cache + pos * kRopeDim;
228
+ const float* sin = cos + rope_half;
229
+ apply_gptj_rope_inplace(buf.data() + nope_dim, cos, sin, rope_half);
230
+
231
+ scalar_t* out_th = q_out + (t * q_head_padded + h) * head_dim;
232
+ for (int64_t i = 0; i < head_dim; ++i)
233
+ out_th[i] = static_cast<scalar_t>(buf[i]);
234
+ }
235
+ });
236
+
237
+ // KV: GPT-J RoPE (no RMSNorm) + bf16 round-trip + UE8M0 FP8 quant + paged insert.
238
+ at::parallel_for(0, num_tokens_insert, 0, [&](int64_t begin, int64_t end) {
239
+ std::vector<float> buf(head_dim);
240
+ std::vector<at::BFloat16> nope_bf16(nope_dim);
241
+ for (int64_t t = begin; t < end; ++t) {
242
+ const int64_t slot_id = slot_mapping[t];
243
+ if (slot_id < 0) continue;
244
+
245
+ const scalar_t* kv_t = kv + t * head_dim;
246
+ for (int64_t i = 0; i < head_dim; ++i)
247
+ buf[i] = static_cast<float>(kv_t[i]);
248
+
249
+ const int64_t pos = positions[t];
250
+ const float* cos = cos_sin_cache + pos * kRopeDim;
251
+ const float* sin = cos + rope_half;
252
+ apply_gptj_rope_inplace(buf.data() + nope_dim, cos, sin, rope_half);
253
+
254
+ // bf16 round-trip, matching the CUDA kernel's double convert.
255
+ for (int64_t i = 0; i < head_dim; ++i)
256
+ buf[i] = static_cast<float>(at::BFloat16(buf[i]));
257
+ for (int64_t i = 0; i < nope_dim; ++i)
258
+ nope_bf16[i] = at::BFloat16(buf[i]);
259
+
260
+ const int64_t block_idx = slot_id / cache_block_size;
261
+ const int64_t pos_in_block = slot_id % cache_block_size;
262
+ uint8_t* const block_base = cache + block_idx * cache_row_stride;
263
+ uint8_t* const token_fp8_ptr = block_base + pos_in_block * token_data_bytes;
264
+ uint8_t* const token_bf16_ptr = token_fp8_ptr + nope_dim;
265
+ uint8_t* const token_scale_ptr =
266
+ block_base + cache_block_size * token_data_bytes + pos_in_block * scale_bytes_per_token;
267
+
268
+ // AVX512 tile-quantization (SGLang's `quantize_tile_avx512`), one
269
+ // 64-wide tile per iteration.
270
+ for (int64_t qb = 0; qb < num_quant_blocks; ++qb) {
271
+ const at::BFloat16* tile_src = nope_bf16.data() + qb * kQuantBlock;
272
+ uint8_t* tile_dst = token_fp8_ptr + qb * kQuantBlock;
273
+ token_scale_ptr[qb] = quantize_tile_avx512(tile_src, tile_dst);
274
+ }
275
+ token_scale_ptr[num_quant_blocks] = 0;
276
+
277
+ auto* out_rope = reinterpret_cast<at::BFloat16*>(token_bf16_ptr);
278
+ for (int64_t i = 0; i < kRopeDim; ++i)
279
+ out_rope[i] = at::BFloat16(buf[nope_dim + i]);
280
+ }
281
+ });
282
+}
283
+
284
+} // anonymous namespace
285
+
286
+at::Tensor fused_qnorm_rope_kv_insert_cpu(
287
+ at::Tensor& q,
288
+ at::Tensor& kv,
289
+ at::Tensor& positions,
290
+ at::Tensor& swa_kv_cache_2d,
291
+ at::Tensor& slot_mapping,
292
+ at::Tensor& cos_sin_cache,
293
+ int64_t q_head_padded,
294
+ double eps,
295
+ int64_t cache_block_size) {
296
+ TORCH_CHECK(
297
+ q.dim() == 3 && (q.scalar_type() == at::kBFloat16 || q.scalar_type() == at::kHalf),
298
+ "fused_qnorm_rope_kv_insert_cpu: q must be bf16/fp16 [num_tokens, num_heads_q, head_dim]");
299
+ TORCH_CHECK(kv.dim() == 2 && kv.scalar_type() == q.scalar_type(), "fused_qnorm_rope_kv_insert_cpu: kv dtype must match q");
300
+ TORCH_CHECK(positions.scalar_type() == at::kLong, "fused_qnorm_rope_kv_insert_cpu: positions must be int64");
301
+ TORCH_CHECK(slot_mapping.scalar_type() == at::kLong, "fused_qnorm_rope_kv_insert_cpu: slot_mapping must be int64");
302
+ TORCH_CHECK(
303
+ swa_kv_cache_2d.dim() == 2 && swa_kv_cache_2d.scalar_type() == at::kByte,
304
+ "fused_qnorm_rope_kv_insert_cpu: swa_kv_cache_2d must be a 2D uint8 tensor");
305
+ TORCH_CHECK(
306
+ q.device().is_cpu() && kv.device().is_cpu() && positions.device().is_cpu() && slot_mapping.device().is_cpu() &&
307
+ swa_kv_cache_2d.device().is_cpu() && cos_sin_cache.device().is_cpu(),
308
+ "fused_qnorm_rope_kv_insert_cpu: all inputs must be CPU tensors");
309
+
310
+ const int64_t num_tokens_full = q.size(0);
311
+ const int64_t num_heads_q = q.size(1);
312
+ const int64_t head_dim = q.size(2);
313
+ const int64_t nope_dim = head_dim - kRopeDim;
314
+ TORCH_CHECK(nope_dim > 0 && nope_dim % kQuantBlock == 0, "fused_qnorm_rope_kv_insert_cpu: unexpected head_dim");
315
+ const int64_t num_quant_blocks = nope_dim / kQuantBlock;
316
+ const int64_t scale_bytes_per_token = num_quant_blocks + 1;
317
+ const int64_t token_data_bytes = nope_dim + kRopeDim * 2;
318
+
319
+ TORCH_CHECK(q_head_padded >= num_heads_q, "fused_qnorm_rope_kv_insert_cpu: q_head_padded must be >= num_heads_q");
320
+ TORCH_CHECK(kv.size(1) == head_dim, "fused_qnorm_rope_kv_insert_cpu: kv head_dim must match q");
321
+ TORCH_CHECK(positions.numel() == num_tokens_full, "fused_qnorm_rope_kv_insert_cpu: positions must have num_tokens_full entries");
322
+
323
+ const int64_t num_tokens_insert = slot_mapping.size(0);
324
+ TORCH_CHECK(kv.size(0) >= num_tokens_insert, "fused_qnorm_rope_kv_insert_cpu: kv must have >= num_tokens_insert rows");
325
+ TORCH_CHECK(cache_block_size > 0, "fused_qnorm_rope_kv_insert_cpu: cache_block_size must be positive");
326
+ TORCH_CHECK(
327
+ cos_sin_cache.scalar_type() == at::kFloat && cos_sin_cache.is_contiguous(),
328
+ "fused_qnorm_rope_kv_insert_cpu: cos_sin_cache must be float32 and "
329
+ "contiguous (convert once in Python)");
330
+ TORCH_CHECK(
331
+ q.is_contiguous() && kv.is_contiguous() && positions.is_contiguous() && slot_mapping.is_contiguous(),
332
+ "fused_qnorm_rope_kv_insert_cpu: q/kv/positions/slot_mapping must be contiguous");
333
+
334
+ // swa_kv_cache_2d is mutated in place; index via its real row stride
335
+ // rather than assuming it equals cache_block_size * (token_data_bytes +
336
+ // scale_bytes_per_token) -- a paged/pooled buffer's row stride is not
337
+ // guaranteed to match the logical per-block byte count.
338
+ const int64_t cache_row_stride = swa_kv_cache_2d.stride(0);
339
+
340
+ auto q_out = at::zeros({num_tokens_full, q_head_padded, head_dim}, q.options());
341
+
342
+ AT_DISPATCH_REDUCED_FLOATING_TYPES(q.scalar_type(), "fused_qnorm_rope_kv_insert_cpu", [&] {
343
+ qnorm_rope_kv_insert_impl<scalar_t>(
344
+ q_out.data_ptr<scalar_t>(),
345
+ swa_kv_cache_2d.data_ptr<uint8_t>(),
346
+ q.data_ptr<scalar_t>(),
347
+ kv.data_ptr<scalar_t>(),
348
+ positions.data_ptr<int64_t>(),
349
+ slot_mapping.data_ptr<int64_t>(),
350
+ cos_sin_cache.data_ptr<float>(),
351
+ num_tokens_full,
352
+ num_heads_q,
353
+ q_head_padded,
354
+ head_dim,
355
+ nope_dim,
356
+ num_quant_blocks,
357
+ token_data_bytes,
358
+ scale_bytes_per_token,
359
+ num_tokens_insert,
360
+ cache_block_size,
361
+ cache_row_stride,
362
+ static_cast<float>(eps));
363
+ });
364
+
365
+ return q_out;
366
+}
367
+
368
+// vLLM-native: inverse GPT-J RoPE for `_o_proj`'s output de-rotation (undoes
369
+// the query's own rotation on the attention output before the low-rank
370
+// wo_a/wo_b projection -- see the math writeup in cpu_sparse.py's `_o_proj`).
371
+// Upstream needs this same step (`fused_rope_inplace(..., inverse=True)` in
372
+// `deepseek_v4.py`) but only has a CUDA-JIT implementation, no CPU kernel --
373
+// this fills that gap rather than porting one. Replaces the eager reference
374
+// in `DeepseekV4CPUAttention._o_proj` (cpu_sparse.py). Uses the same
375
+// half-half `[cos...|sin...]` cache layout as `apply_gptj_rope_inplace`
376
+// above (not upstream's interleaved `[c0,s0,c1,s1,...]` layout).
377
+//
378
+// AVX512 path: NoPE dtype-convert copy uses `vec.h`'s `load_float_vec`
379
+// (16-wide bf16/fp16 -> fp32). The interleaved (even,odd) rotation widens
380
+// 32 raw elements to fp32 via `load_float_vec2`, round-trips them through
381
+// an aligned scratch buffer, then gathers/scatters the even/odd lanes --
382
+// same technique as `indexer.cpp`'s Q-side RoPE, since `cos`/`sin` here
383
+// are two contiguous halves rather than per-pair-interleaved.
384
+at::Tensor inverse_gptj_rope_o_proj_cpu(
385
+ at::Tensor& o, // [num_tokens, num_heads, head_dim]
386
+ at::Tensor& positions, // [num_tokens]
387
+ at::Tensor& cos_sin_cache, // [max_pos, rope_dim] (first half cos, second half sin)
388
+ int64_t rope_dim) {
389
+ TORCH_CHECK(o.dim() == 3, "inverse_gptj_rope_o_proj_cpu: o must be 3D [num_tokens, num_heads, head_dim]");
390
+ TORCH_CHECK(positions.scalar_type() == at::kLong, "inverse_gptj_rope_o_proj_cpu: positions must be int64");
391
+ TORCH_CHECK(
392
+ o.device().is_cpu() && positions.device().is_cpu() && cos_sin_cache.device().is_cpu(),
393
+ "inverse_gptj_rope_o_proj_cpu: all inputs must be CPU tensors");
394
+
395
+ const int64_t num_tokens = o.size(0);
396
+ const int64_t num_heads = o.size(1);
397
+ const int64_t head_dim = o.size(2);
398
+ const int64_t nope_dim = head_dim - rope_dim;
399
+ const int64_t half_dim = rope_dim / 2;
400
+ TORCH_CHECK(nope_dim >= 0, "inverse_gptj_rope_o_proj_cpu: rope_dim must not exceed head_dim");
401
+ // Only ever called with head_dim=512, rope_dim=64 (DeepseekV4CPUAttention's
402
+ // _o_proj); the AVX512 loops below process 16 elements/pairs per iteration
403
+ // with no scalar remainder, so require both to divide evenly.
404
+ TORCH_CHECK(nope_dim % 16 == 0 && half_dim % 16 == 0,
405
+ "inverse_gptj_rope_o_proj_cpu: head_dim-rope_dim and rope_dim/2 must be multiples of 16");
406
+ TORCH_CHECK(
407
+ cos_sin_cache.scalar_type() == at::kFloat && cos_sin_cache.is_contiguous(),
408
+ "inverse_gptj_rope_o_proj_cpu: cos_sin_cache must be float32 and "
409
+ "contiguous (convert once in Python)");
410
+ TORCH_CHECK(o.is_contiguous() && positions.is_contiguous(),
411
+ "inverse_gptj_rope_o_proj_cpu: o/positions must be contiguous");
412
+
413
+ auto out = at::empty({num_tokens, num_heads, head_dim}, o.options().dtype(at::kFloat));
414
+
415
+ AT_DISPATCH_REDUCED_FLOATING_TYPES(o.scalar_type(), "inverse_gptj_rope_o_proj_cpu", [&] {
416
+ const scalar_t* o_ptr = o.data_ptr<scalar_t>();
417
+ float* out_ptr = out.data_ptr<float>();
418
+ const int64_t* pos_ptr = positions.data_ptr<int64_t>();
419
+ const float* cos_sin_ptr = cos_sin_cache.data_ptr<float>();
420
+
421
+ at::parallel_for(0, num_tokens * num_heads, 0, [&](int64_t begin, int64_t end) {
422
+ for (int64_t idx = begin; idx < end; ++idx) {
423
+ const int64_t t = idx / num_heads;
424
+ const int64_t h = idx % num_heads;
425
+ const scalar_t* src = o_ptr + (t * num_heads + h) * head_dim;
426
+ float* dst = out_ptr + (t * num_heads + h) * head_dim;
427
+
428
+ // head_dim=512, rope_dim=64 in practice (checked above), so
429
+ // nope_dim/half_dim are always multiples of 16 -- no scalar
430
+ // remainder needed.
431
+ for (int64_t i = 0; i < nope_dim; i += 16) {
432
+ load_float_vec<scalar_t>(src + i).store(dst + i);
433
+ }
434
+
435
+ const int64_t pos = pos_ptr[t];
436
+ const float* cos = cos_sin_ptr + pos * rope_dim;
437
+ const float* sin = cos + half_dim;
438
+
439
+ {
440
+ // Permute-based (no gather/scatter), adapted from SGLang's
441
+ // `rope.cpp` (`apply_rotary_emb_row`) the same way as
442
+ // compressor.cpp's `apply_gptj_rope_inplace`: `_mm512_permute_ps`
443
+ // swaps each pair's (even, odd) lanes in one shuffle, and cos/sin
444
+ // -- loaded as 8 contiguous values -- are duplicated pairwise via
445
+ // `_mm512_permutexvar_ps` to align with that layout, since our
446
+ // cos/sin are two separate flat arrays rather than SGLang's
447
+ // pre-interleaved freqs. Two batches of 8 pairs per iteration
448
+ // (matching `load_float_vec2`'s 16+16-lane split).
449
+ const __m512i dup_idx = _mm512_set_epi32(7, 7, 6, 6, 5, 5, 4, 4, 3, 3, 2, 2, 1, 1, 0, 0);
450
+ // Inverse GPT-J rotation negates the sin broadcast on the odd
451
+ // (imag) lane of each pair; the even (real) lane keeps sin's sign.
452
+ const __m512 sign_mask = _mm512_castsi512_ps(_mm512_set_epi32(
453
+ (int)0x80000000, 0, (int)0x80000000, 0, (int)0x80000000, 0, (int)0x80000000, 0,
454
+ (int)0x80000000, 0, (int)0x80000000, 0, (int)0x80000000, 0, (int)0x80000000, 0));
455
+ for (int64_t j = 0; j < half_dim; j += 16) {
456
+ const scalar_t* rsrc = src + nope_dim + 2 * j;
457
+ at::vec::Vectorized<float> x0, x1;
458
+ std::tie(x0, x1) = load_float_vec2(rsrc);
459
+ float* rdst = dst + nope_dim + 2 * j;
460
+
461
+ __m512 xv0 = x0;
462
+ __m512 cos_b0 = _mm512_permutexvar_ps(dup_idx, _mm512_castps256_ps512(_mm256_loadu_ps(cos + j)));
463
+ __m512 sin_b0 = _mm512_permutexvar_ps(dup_idx, _mm512_castps256_ps512(_mm256_loadu_ps(sin + j)));
464
+ __m512 out0 = _mm512_fmadd_ps(
465
+ xv0, cos_b0, _mm512_mul_ps(_mm512_permute_ps(xv0, 0xB1), _mm512_xor_ps(sin_b0, sign_mask)));
466
+ _mm512_storeu_ps(rdst, out0);
467
+
468
+ __m512 xv1 = x1;
469
+ __m512 cos_b1 = _mm512_permutexvar_ps(dup_idx, _mm512_castps256_ps512(_mm256_loadu_ps(cos + j + 8)));
470
+ __m512 sin_b1 = _mm512_permutexvar_ps(dup_idx, _mm512_castps256_ps512(_mm256_loadu_ps(sin + j + 8)));
471
+ __m512 out1 = _mm512_fmadd_ps(
472
+ xv1, cos_b1, _mm512_mul_ps(_mm512_permute_ps(xv1, 0xB1), _mm512_xor_ps(sin_b1, sign_mask)));
473
+ _mm512_storeu_ps(rdst + 16, out1);
474
+ }
475
+ }
476
+ }
477
+ });
478
+ });
479
+
480
+ return out;
481
+}
csrc/cpu/sgl-kernels/topk.cpp
new
+834
@@ -0,0 +1,834 @@
1
+// Adapted from
2
+// https://github.com/sgl-project/sglang/tree/main/sgl-kernel/csrc/cpu:
3
+// `biased_topk_kernel_impl`/`biased_topk_cpu` (flat, non-grouped DeepSeek-V4
4
+// biased top-k) and `hash_topk_kernel_impl`/`hash_topk_cpu` (hash-routed
5
+// layers, expert IDs from a precomputed `tid2eid` lookup table), plus their
6
+// `sigmoid`/`softmax`/`apply_bias`/`sqrtsoftplus` scoring helpers. See
7
+// `biased_topk_cpu`/`hash_topk_cpu` below for the one deliberate deviation
8
+// from the upstream kernel (routed_scaling_factor handling when no shared
9
+// expert is fused into routing, vLLM's CPU DeepSeek-V4 calling convention).
10
+//
11
+// clang-format off
12
+
13
+#include <algorithm>
14
+#include <array>
15
+#include <cmath>
16
+#include <limits>
17
+#include <type_traits>
18
+#include <utility>
19
+
20
+#include "common.h"
21
+#include "vec.h"
22
+
23
+namespace {
24
+
25
+template <typename scalar_t, int SIZE, std::enable_if_t<!std::is_same_v<scalar_t, float>, int> = 0>
26
+inline void sigmoid(float* __restrict__ out, const scalar_t* __restrict__ input) {
27
+ using bVec = at::vec::Vectorized<scalar_t>;
28
+ using fVec = at::vec::Vectorized<float>;
29
+
30
+ const fVec one = fVec(1.f);
31
+
32
+ constexpr int kVecSize = bVec::size();
33
+ for (int d = 0; d < SIZE; d += kVecSize) {
34
+ bVec x_bvec = bVec::loadu(input + d);
35
+ fVec x_fvec0, x_fvec1;
36
+ std::tie(x_fvec0, x_fvec1) = at::vec::convert_to_float(x_bvec);
37
+
38
+ x_fvec0 = one / (one + x_fvec0.neg().exp_u20());
39
+ x_fvec1 = one / (one + x_fvec1.neg().exp_u20());
40
+
41
+ x_fvec0.store(out + d);
42
+ x_fvec1.store(out + d + fVec::size());
43
+ }
44
+}
45
+
46
+template <typename scalar_t, int SIZE, std::enable_if_t<std::is_same_v<scalar_t, float>, int> = 0>
47
+inline void sigmoid(float* __restrict__ out, const float* __restrict__ input) {
48
+ using fVec = at::vec::Vectorized<float>;
49
+ const fVec one = fVec(1.f);
50
+ constexpr int kVecSize = fVec::size();
51
+ for (int d = 0; d < SIZE; d += kVecSize) {
52
+ fVec in_fvec = fVec::loadu(input + d);
53
+ in_fvec = one / (one + in_fvec.neg().exp_u20());
54
+ in_fvec.store(out + d);
55
+ }
56
+}
57
+
58
+template <typename scalar_t, int SIZE, std::enable_if_t<!std::is_same_v<scalar_t, float>, int> = 0>
59
+inline void softmax(float* __restrict__ out, const scalar_t* __restrict__ input) {
60
+ using bVec = at::vec::Vectorized<scalar_t>;
61
+ using fVec = at::vec::Vectorized<float>;
62
+
63
+ constexpr int kVecSize = bVec::size();
64
+
65
+ // step 1: get max
66
+ fVec max_fvec = fVec(-std::numeric_limits<float>::infinity());
67
+ for (int d = 0; d < SIZE; d += kVecSize) {
68
+ bVec x_bvec = bVec::loadu(input + d);
69
+ fVec x_fvec0, x_fvec1;
70
+ std::tie(x_fvec0, x_fvec1) = at::vec::convert_to_float(x_bvec);
71
+
72
+ max_fvec = at::vec::maximum(max_fvec, x_fvec0);
73
+ max_fvec = at::vec::maximum(max_fvec, x_fvec1);
74
+ x_fvec0.store(out + d);
75
+ x_fvec1.store(out + d + fVec::size());
76
+ }
77
+ float max_val = vec_reduce_max(max_fvec);
78
+ max_fvec = fVec(max_val);
79
+
80
+ // step 2: sum of (x - max).exp()
81
+ fVec sum_fvec = fVec(float(0));
82
+ for (int d = 0; d < SIZE; d += fVec::size()) {
83
+ fVec x_fvec = (fVec::loadu(out + d) - max_fvec).exp_u20();
84
+ sum_fvec += x_fvec;
85
+ x_fvec.store(out + d);
86
+ }
87
+ float sum_val = vec_reduce_sum(sum_fvec);
88
+
89
+ // step 3: x * (1 / sum)
90
+ sum_fvec = fVec(1.f / sum_val);
91
+ for (int d = 0; d < SIZE; d += fVec::size()) {
92
+ fVec out_fvec = fVec::loadu(out + d) * sum_fvec;
93
+ out_fvec.store(out + d);
94
+ }
95
+}
96
+
97
+template <typename scalar_t, int SIZE, std::enable_if_t<std::is_same_v<scalar_t, float>, int> = 0>
98
+inline void softmax(float* __restrict__ out, const float* __restrict__ input) {
99
+ using fVec = at::vec::Vectorized<float>;
100
+
101
+ constexpr int kVecSize = fVec::size();
102
+
103
+ // step 1: get max
104
+ fVec max_fvec = fVec(-std::numeric_limits<float>::infinity());
105
+ for (int d = 0; d < SIZE; d += kVecSize) {
106
+ fVec x_fvec = fVec::loadu(input + d);
107
+ max_fvec = at::vec::maximum(max_fvec, x_fvec);
108
+ x_fvec.store(out + d);
109
+ }
110
+ float max_val = vec_reduce_max(max_fvec);
111
+ max_fvec = fVec(max_val);
112
+
113
+ // step 2: sum of (x - max).exp()
114
+ fVec sum_fvec = fVec(float(0));
115
+ for (int d = 0; d < SIZE; d += kVecSize) {
116
+ fVec x_fvec = (fVec::loadu(out + d) - max_fvec).exp_u20();
117
+ sum_fvec += x_fvec;
118
+ x_fvec.store(out + d);
119
+ }
120
+ float sum_val = vec_reduce_sum(sum_fvec);
121
+
122
+ // step 3: x * (1 / sum)
123
+ sum_fvec = fVec(1.f / sum_val);
124
+ for (int d = 0; d < SIZE; d += kVecSize) {
125
+ fVec out_fvec = fVec::loadu(out + d) * sum_fvec;
126
+ out_fvec.store(out + d);
127
+ }
128
+}
129
+
130
+template <typename param_t, int SIZE>
131
+inline void
132
+apply_bias(float* __restrict__ scores2, const float* __restrict__ scores, const param_t* __restrict__ bias) {
133
+ using fVec = at::vec::Vectorized<float>;
134
+ auto vec_size = fVec::size() * 2;
135
+ int d = 0;
136
+ for (; d <= SIZE - vec_size; d += vec_size) {
137
+ fVec bias0, bias1, x0, x1;
138
+ std::tie(bias0, bias1) = load_float_vec2(bias + d);
139
+ std::tie(x0, x1) = load_float_vec2(scores + d);
140
+ x0 = x0 + bias0;
141
+ x1 = x1 + bias1;
142
+ x0.store(scores2 + d);
143
+ x1.store(scores2 + d + fVec::size());
144
+ }
145
+ for (; d < SIZE; d++) {
146
+ scores2[d] = scores[d] + (float)bias[d];
147
+ }
148
+}
149
+
150
+// sqrtsoftplus: sqrt(softplus(x)) = sqrt(log(1 + exp(x)))
151
+// For numerical stability: when x > threshold, softplus(x) ≈ x
152
+// When x < -threshold, softplus(x) ≈ exp(x), so sqrt(softplus(x)) ≈ sqrt(exp(x)) = exp(x/2)
153
+template <typename scalar_t, int SIZE, std::enable_if_t<!std::is_same_v<scalar_t, float>, int> = 0>
154
+inline void sqrtsoftplus(float* __restrict__ out, const scalar_t* __restrict__ input) {
155
+ using bVec = at::vec::Vectorized<scalar_t>;
156
+ using fVec = at::vec::Vectorized<float>;
157
+
158
+ const fVec one = fVec(1.f);
159
+ const fVec half = fVec(0.5f);
160
+ const fVec threshold = fVec(20.f);
161
+ const fVec neg_threshold = fVec(-15.f); // below this, 1+exp(x) loses precision in float32
162
+
163
+ constexpr int kVecSize = bVec::size();
164
+ for (int d = 0; d < SIZE; d += kVecSize) {
165
+ bVec x_bvec = bVec::loadu(input + d);
166
+ fVec x_fvec0, x_fvec1;
167
+ std::tie(x_fvec0, x_fvec1) = at::vec::convert_to_float(x_bvec);
168
+
169
+ // default: softplus(x) = log(1 + exp(x))
170
+ fVec sp0 = (one + x_fvec0.exp_u20()).log();
171
+ fVec sp1 = (one + x_fvec1.exp_u20()).log();
172
+ // x > 20: softplus(x) ≈ x
173
+ sp0 = fVec::blendv(sp0, x_fvec0, x_fvec0 > threshold);
174
+ sp1 = fVec::blendv(sp1, x_fvec1, x_fvec1 > threshold);
175
+ // x < -15: softplus(x) ≈ exp(x), sqrt(exp(x)) = exp(x/2)
176
+ fVec exp_half0 = (x_fvec0 * half).exp_u20();
177
+ fVec exp_half1 = (x_fvec1 * half).exp_u20();
178
+ sp0 = fVec::blendv(sp0.sqrt(), exp_half0, x_fvec0 < neg_threshold);
179
+ sp1 = fVec::blendv(sp1.sqrt(), exp_half1, x_fvec1 < neg_threshold);
180
+
181
+ sp0.store(out + d);
182
+ sp1.store(out + d + fVec::size());
183
+ }
184
+}
185
+
186
+template <typename scalar_t, int SIZE, std::enable_if_t<std::is_same_v<scalar_t, float>, int> = 0>
187
+inline void sqrtsoftplus(float* __restrict__ out, const float* __restrict__ input) {
188
+ using fVec = at::vec::Vectorized<float>;
189
+ const fVec one = fVec(1.f);
190
+ const fVec half = fVec(0.5f);
191
+ const fVec threshold = fVec(20.f);
192
+ const fVec neg_threshold = fVec(-15.f);
193
+ constexpr int kVecSize = fVec::size();
194
+ for (int d = 0; d < SIZE; d += kVecSize) {
195
+ fVec x = fVec::loadu(input + d);
196
+ fVec sp = (one + x.exp_u20()).log();
197
+ sp = fVec::blendv(sp, x, x > threshold);
198
+ fVec exp_half = (x * half).exp_u20();
199
+ sp = fVec::blendv(sp.sqrt(), exp_half, x < neg_threshold);
200
+ sp.store(out + d);
201
+ }
202
+}
203
+
204
+// biased_topk: flat (non-grouped) biased topk for DeepSeek V4
205
+// scoring_func: 0 = sigmoid, 1 = sqrtsoftplus
206
+template <typename scalar_t, typename param_t, int NUM_EXPERTS, int TOPK>
207
+void biased_topk_kernel_impl(
208
+ float* __restrict__ topk_weights,
209
+ int32_t* __restrict__ topk_ids,
210
+ const scalar_t* __restrict__ gating_output,
211
+ const param_t* __restrict__ bias,
212
+ int64_t num_tokens,
213
+ int64_t topk,
214
+ bool renormalize,
215
+ int scoring_func,
216
+ int64_t num_fused_shared_experts,
217
+ float routed_scaling_factor,
218
+ bool apply_routed_scaling_factor_on_output) {
219
+ at::parallel_for(0, num_tokens, 0, [&](int64_t begin, int64_t end) {
220
+ alignas(64) float scores[NUM_EXPERTS];
221
+ alignas(64) float scores_biased[NUM_EXPERTS];
222
+
223
+ using elem_t = std::pair<float, int32_t>;
224
+ std::array<elem_t, NUM_EXPERTS> queue;
225
+
226
+ // simple RNG for fused shared expert random ID
227
+ uint64_t rng_state = begin * 6364136223846793005ULL + 1442695040888963407ULL;
228
+
229
+ for (int64_t i = begin; i < end; ++i) {
230
+ // compute scores
231
+ if (scoring_func == 0) {
232
+ sigmoid<scalar_t, NUM_EXPERTS>(scores, gating_output + i * NUM_EXPERTS);
233
+ } else {
234
+ sqrtsoftplus<scalar_t, NUM_EXPERTS>(scores, gating_output + i * NUM_EXPERTS);
235
+ }
236
+
237
+ // add bias for selection
238
+ apply_bias<param_t, NUM_EXPERTS>(scores_biased, scores, bias);
239
+
240
+ // build queue and partial sort to find top-k
241
+ for (int64_t e = 0; e < NUM_EXPERTS; ++e) {
242
+ queue[e] = {scores_biased[e], static_cast<int32_t>(e)};
243
+ }
244
+
245
+ std::partial_sort(
246
+ queue.begin(), queue.begin() + topk, queue.end(), [](const elem_t& x, const elem_t& y) -> bool {
247
+ return x.first > y.first;
248
+ });
249
+
250
+ // gather original scores (without bias) as weights
251
+ for (int64_t j = 0; j < topk; ++j) {
252
+ int32_t idx = queue[j].second;
253
+ topk_ids[i * topk + j] = idx;
254
+ topk_weights[i * topk + j] = scores[idx];
255
+ }
256
+
257
+ // handle fused shared experts
258
+ if (num_fused_shared_experts > 0) {
259
+ // replace last slot with random shared expert ID
260
+ rng_state = rng_state * 6364136223846793005ULL + 1442695040888963407ULL;
261
+ int32_t shared_id =
262
+ NUM_EXPERTS + static_cast<int32_t>((rng_state >> 33) % static_cast<uint64_t>(num_fused_shared_experts));
263
+ topk_ids[i * topk + topk - 1] = shared_id;
264
+
265
+ // shared expert weight = sum of routed weights / scaling_factor
266
+ if (routed_scaling_factor != 0.0f) {
267
+ float routed_sum = 0.f;
268
+ for (int64_t j = 0; j < topk - 1; ++j) {
269
+ routed_sum += topk_weights[i * topk + j];
270
+ }
271
+ topk_weights[i * topk + topk - 1] = routed_sum / routed_scaling_factor;
272
+ }
273
+ }
274
+
275
+ // renormalize
276
+ if (renormalize) {
277
+ float sum = 0.f;
278
+ int64_t norm_end = (num_fused_shared_experts == 0) ? topk : topk - 1;
279
+ for (int64_t j = 0; j < norm_end; ++j) {
280
+ sum += topk_weights[i * topk + j];
281
+ }
282
+ float scale = 1.f / sum;
283
+ if (apply_routed_scaling_factor_on_output) {
284
+ scale *= routed_scaling_factor;
285
+ }
286
+ for (int64_t j = 0; j < norm_end; ++j) {
287
+ topk_weights[i * topk + j] *= scale;
288
+ }
289
+ // also scale the shared expert weight if present
290
+ if (num_fused_shared_experts > 0) {
291
+ topk_weights[i * topk + topk - 1] *= scale;
292
+ }
293
+ }
294
+ }
295
+ });
296
+}
297
+
298
+// hash_topk: expert IDs come from a precomputed lookup table tid2eid[input_ids]
299
+// scoring_func: 0 = softmax, 1 = sigmoid, 2 = sqrtsoftplus
300
+template <typename scalar_t, int NUM_EXPERTS, int TOPK>
301
+void hash_topk_kernel_impl(
302
+ float* __restrict__ topk_weights,
303
+ int32_t* __restrict__ topk_ids,
304
+ const scalar_t* __restrict__ gating_output,
305
+ const int32_t* __restrict__ tid2eid, // [num_tokens, routed_topk]
306
+ int64_t num_tokens,
307
+ int scoring_func,
308
+ int64_t num_fused_shared_experts,
309
+ int64_t num_experts,
310
+ float routed_scaling_factor,
311
+ int64_t topk) {
312
+ const int64_t routed_topk = topk - num_fused_shared_experts;
313
+ const bool need_renormalize = (scoring_func != 0); // renormalize for non-softmax
314
+
315
+ at::parallel_for(0, num_tokens, 0, [&](int64_t begin, int64_t end) {
316
+ alignas(64) float scores[NUM_EXPERTS];
317
+
318
+ // simple RNG for fused shared expert random ID
319
+ uint64_t rng_state = begin * 6364136223846793005ULL + 1442695040888963407ULL;
320
+
321
+ for (int64_t i = begin; i < end; ++i) {
322
+ // compute scores over all experts
323
+ if (scoring_func == 0) {
324
+ softmax<scalar_t, NUM_EXPERTS>(scores, gating_output + i * NUM_EXPERTS);
325
+ } else if (scoring_func == 1) {
326
+ sigmoid<scalar_t, NUM_EXPERTS>(scores, gating_output + i * NUM_EXPERTS);
327
+ } else {
328
+ sqrtsoftplus<scalar_t, NUM_EXPERTS>(scores, gating_output + i * NUM_EXPERTS);
329
+ }
330
+
331
+ // gather expert IDs from lookup table
332
+ const int32_t* eid_row = tid2eid + i * routed_topk;
333
+ for (int64_t j = 0; j < routed_topk; ++j) {
334
+ int32_t eid = eid_row[j];
335
+ topk_ids[i * topk + j] = eid;
336
+ topk_weights[i * topk + j] = scores[eid];
337
+ }
338
+
339
+ // renormalize routed weights (for non-softmax scoring). Unlike the
340
+ // upstream kernel, always fold routed_scaling_factor in here rather
341
+ // than only applying it to the fused-shared-expert slot below: vLLM's
342
+ // DeepSeek-V4 CPU hash-routed layers never fuse a shared expert into
343
+ // this routing table (num_fused_shared_experts is always 0 on this
344
+ // call path), so this is the only place the scaling factor can be
345
+ // applied. Matches the eager reference (`sqrtsoftplus_bias_topk` in
346
+ // cpu_moe.py) and the CUDA kernel (`dsv4HashTopkSoftplusSqrt`).
347
+ if (need_renormalize) {
348
+ float sum = 0.f;
349
+ for (int64_t j = 0; j < routed_topk; ++j) {
350
+ sum += topk_weights[i * topk + j];
351
+ }
352
+ if (sum > 0.f) {
353
+ float scale = 1.f / sum;
354
+ if (num_fused_shared_experts == 0) {
355
+ scale *= routed_scaling_factor;
356
+ }
357
+ for (int64_t j = 0; j < routed_topk; ++j) {
358
+ topk_weights[i * topk + j] *= scale;
359
+ }
360
+ }
361
+ }
362
+
363
+ // handle fused shared expert
364
+ if (num_fused_shared_experts > 0) {
365
+ // random shared expert ID in [num_experts, num_experts + num_fused_shared_experts)
366
+ rng_state = rng_state * 6364136223846793005ULL + 1442695040888963407ULL;
367
+ int32_t shared_id =
368
+ num_experts + static_cast<int32_t>((rng_state >> 33) % static_cast<uint64_t>(num_fused_shared_experts));
369
+ topk_ids[i * topk + topk - 1] = shared_id;
370
+
371
+ // shared expert weight = sum of routed weights / scaling_factor
372
+ float routed_sum = 0.f;
373
+ for (int64_t j = 0; j < routed_topk; ++j) {
374
+ routed_sum += topk_weights[i * topk + j];
375
+ }
376
+ topk_weights[i * topk + topk - 1] = routed_sum / routed_scaling_factor;
377
+ }
378
+ }
379
+ });
380
+}
381
+
382
+#define LAUNCH_HASH_TOPK_KERNEL(NE, NTOPK) \
383
+ hash_topk_kernel_impl<scalar_t, NE, NTOPK>( \
384
+ topk_weights.data_ptr<float>(), \
385
+ topk_ids.data_ptr<int32_t>(), \
386
+ gating_output.data_ptr<scalar_t>(), \
387
+ tid2eid_flat.data_ptr<int32_t>(), \
388
+ num_tokens, \
389
+ scoring_func_id, \
390
+ num_fused_shared_experts, \
391
+ num_experts, \
392
+ routed_scaling_factor_value, \
393
+ topk);
394
+
395
+#define LAUNCH_BIASED_TOPK_KERNEL(NE, NTOPK) \
396
+ biased_topk_kernel_impl<scalar_t, param_t, NE, NTOPK>( \
397
+ topk_weights.data_ptr<float>(), \
398
+ topk_ids.data_ptr<int32_t>(), \
399
+ gating_output.data_ptr<scalar_t>(), \
400
+ correction_bias.data_ptr<param_t>(), \
401
+ num_tokens, \
402
+ topk, \
403
+ renormalize, \
404
+ scoring_func_id, \
405
+ num_fused_shared_experts, \
406
+ routed_scaling_factor_value, \
407
+ apply_routed_scaling_factor_on_output);
408
+
409
+} // namespace
410
+
411
+// biased topk for DeepSeek V4 (flat, non-grouped)
412
+std::tuple<at::Tensor, at::Tensor> biased_topk_cpu(
413
+ at::Tensor& hidden_states,
414
+ at::Tensor& gating_output,
415
+ at::Tensor& correction_bias,
416
+ int64_t topk,
417
+ bool renormalize,
418
+ std::string scoring_func,
419
+ int64_t num_fused_shared_experts,
420
+ std::optional<double> routed_scaling_factor,
421
+ bool apply_routed_scaling_factor_on_output) {
422
+ CHECK_INPUT(gating_output);
423
+ CHECK_INPUT(correction_bias);
424
+
425
+ const auto st = gating_output.scalar_type();
426
+ int64_t num_tokens = hidden_states.size(0);
427
+ int64_t num_experts = gating_output.size(1);
428
+ TORCH_CHECK(gating_output.size(0) == num_tokens, "Number of tokens mismatch");
429
+ TORCH_CHECK(correction_bias.numel() == num_experts, "Bias shape mismatch");
430
+
431
+ int scoring_func_id = 0;
432
+ if (scoring_func == "sigmoid") {
433
+ scoring_func_id = 0;
434
+ } else if (scoring_func == "sqrtsoftplus") {
435
+ scoring_func_id = 1;
436
+ } else {
437
+ TORCH_CHECK(false, "Unsupported scoring_func: ", scoring_func);
438
+ }
439
+
440
+ float routed_scaling_factor_value = routed_scaling_factor.has_value() ? routed_scaling_factor.value() : 0.0f;
441
+
442
+ at::Tensor topk_weights = at::empty({num_tokens, topk}, hidden_states.options().dtype(at::kFloat));
443
+ at::Tensor topk_ids = at::empty({num_tokens, topk}, hidden_states.options().dtype(at::kInt));
444
+
445
+ // The actual routed topk (excluding fused shared experts)
446
+ int64_t routed_topk = topk - num_fused_shared_experts;
447
+
448
+ CPU_DISPATCH_FLOATING_TYPES_EXT(st, correction_bias.scalar_type(), "biased_topk_kernel", [&] {
449
+ // dispatch on num_experts and routed_topk. For DeepSeek V4: 256 experts,
450
+ // routed_topk=6 (num_fused_shared_experts is always 0 on this CPU call
451
+ // path -- see cpu_moe.py). Anything outside this table isn't a real
452
+ // model shape, so fail loudly instead of silently degrading to a slow
453
+ // generic path.
454
+ switch (num_experts) {
455
+ case 64:
456
+ switch (routed_topk) {
457
+ case 6:
458
+ LAUNCH_BIASED_TOPK_KERNEL(64, 6);
459
+ break;
460
+ case 8:
461
+ LAUNCH_BIASED_TOPK_KERNEL(64, 8);
462
+ break;
463
+ default:
464
+ TORCH_CHECK(false, "biased_topk_cpu: unsupported routed_topk ", routed_topk,
465
+ " for num_experts=64");
466
+ }
467
+ break;
468
+ case 128:
469
+ switch (routed_topk) {
470
+ case 6:
471
+ LAUNCH_BIASED_TOPK_KERNEL(128, 6);
472
+ break;
473
+ case 8:
474
+ LAUNCH_BIASED_TOPK_KERNEL(128, 8);
475
+ break;
476
+ default:
477
+ TORCH_CHECK(false, "biased_topk_cpu: unsupported routed_topk ", routed_topk,
478
+ " for num_experts=128");
479
+ }
480
+ break;
481
+ case 256:
482
+ switch (routed_topk) {
483
+ case 6:
484
+ LAUNCH_BIASED_TOPK_KERNEL(256, 6);
485
+ break;
486
+ case 8:
487
+ LAUNCH_BIASED_TOPK_KERNEL(256, 8);
488
+ break;
489
+ case 9:
490
+ LAUNCH_BIASED_TOPK_KERNEL(256, 9);
491
+ break;
492
+ default:
493
+ TORCH_CHECK(false, "biased_topk_cpu: unsupported routed_topk ", routed_topk,
494
+ " for num_experts=256");
495
+ }
496
+ break;
497
+ case 384:
498
+ switch (routed_topk) {
499
+ case 6:
500
+ LAUNCH_BIASED_TOPK_KERNEL(384, 6);
501
+ break;
502
+ case 8:
503
+ LAUNCH_BIASED_TOPK_KERNEL(384, 8);
504
+ break;
505
+ default:
506
+ TORCH_CHECK(false, "biased_topk_cpu: unsupported routed_topk ", routed_topk,
507
+ " for num_experts=384");
508
+ }
509
+ break;
510
+ default:
511
+ TORCH_CHECK(false, "biased_topk_cpu: unsupported num_experts ", num_experts);
512
+ }
513
+ });
514
+ return std::make_tuple(topk_weights, topk_ids);
515
+}
516
+
517
+// hash topk for DeepSeek V4 (expert IDs from precomputed lookup table)
518
+std::tuple<at::Tensor, at::Tensor> hash_topk_cpu(
519
+ at::Tensor& gating_output,
520
+ at::Tensor& tid2eid,
521
+ int64_t topk,
522
+ std::string scoring_func,
523
+ int64_t num_fused_shared_experts,
524
+ int64_t num_experts,
525
+ double routed_scaling_factor) {
526
+ CHECK_INPUT(gating_output);
527
+ CHECK_INPUT(tid2eid);
528
+
529
+ const auto st = gating_output.scalar_type();
530
+ int64_t num_tokens = gating_output.size(0);
531
+ int64_t num_experts_gating = gating_output.size(1);
532
+ int64_t routed_topk = topk - num_fused_shared_experts;
533
+
534
+ TORCH_CHECK(tid2eid.size(0) == num_tokens, "tid2eid row count must match num_tokens");
535
+ TORCH_CHECK(tid2eid.size(1) == routed_topk, "tid2eid column count must match routed_topk");
536
+ TORCH_CHECK(tid2eid.scalar_type() == at::kInt, "tid2eid must be int32");
537
+ TORCH_CHECK(num_experts_gating == num_experts, "num_experts mismatch");
538
+
539
+ int scoring_func_id = 0;
540
+ if (scoring_func == "softmax") {
541
+ scoring_func_id = 0;
542
+ } else if (scoring_func == "sigmoid") {
543
+ scoring_func_id = 1;
544
+ } else if (scoring_func == "sqrtsoftplus") {
545
+ scoring_func_id = 2;
546
+ } else {
547
+ TORCH_CHECK(false, "Unsupported scoring_func: ", scoring_func);
548
+ }
549
+
550
+ float routed_scaling_factor_value = static_cast<float>(routed_scaling_factor);
551
+
552
+ at::Tensor topk_weights = at::empty({num_tokens, topk}, gating_output.options().dtype(at::kFloat));
553
+ at::Tensor topk_ids = at::empty({num_tokens, topk}, gating_output.options().dtype(at::kInt));
554
+
555
+ // tid2eid is [num_tokens, routed_topk], already indexed by input_ids in
556
+ // Python; CHECK_INPUT above already guarantees it is contiguous.
557
+ const at::Tensor& tid2eid_flat = tid2eid;
558
+
559
+ // Dispatch for bf16, fp16, and float32
560
+ // Note: cannot use AT_DISPATCH_FLOATING_TYPES_AND2 since it includes double which lacks convert_to_float
561
+ auto dispatch_fn = [&]<typename scalar_t>() {
562
+ // For DeepSeek V4: 256 experts, routed_topk in {6,7,8} depending on
563
+ // layer (num_fused_shared_experts is always 0 on this CPU call path --
564
+ // see cpu_moe.py). Anything outside this table isn't a real model
565
+ // shape, so fail loudly instead of silently degrading to a slow
566
+ // generic path.
567
+ switch (num_experts) {
568
+ case 64:
569
+ switch (routed_topk) {
570
+ case 6:
571
+ LAUNCH_HASH_TOPK_KERNEL(64, 6);
572
+ break;
573
+ case 7:
574
+ LAUNCH_HASH_TOPK_KERNEL(64, 7);
575
+ break;
576
+ case 8:
577
+ LAUNCH_HASH_TOPK_KERNEL(64, 8);
578
+ break;
579
+ default:
580
+ TORCH_CHECK(false, "hash_topk_cpu: unsupported routed_topk ", routed_topk,
581
+ " for num_experts=64");
582
+ }
583
+ break;
584
+ case 128:
585
+ switch (routed_topk) {
586
+ case 6:
587
+ LAUNCH_HASH_TOPK_KERNEL(128, 6);
588
+ break;
589
+ case 7:
590
+ LAUNCH_HASH_TOPK_KERNEL(128, 7);
591
+ break;
592
+ case 8:
593
+ LAUNCH_HASH_TOPK_KERNEL(128, 8);
594
+ break;
595
+ default:
596
+ TORCH_CHECK(false, "hash_topk_cpu: unsupported routed_topk ", routed_topk,
597
+ " for num_experts=128");
598
+ }
599
+ break;
600
+ case 256:
601
+ switch (routed_topk) {
602
+ case 6:
603
+ LAUNCH_HASH_TOPK_KERNEL(256, 6);
604
+ break;
605
+ case 7:
606
+ LAUNCH_HASH_TOPK_KERNEL(256, 7);
607
+ break;
608
+ case 8:
609
+ LAUNCH_HASH_TOPK_KERNEL(256, 8);
610
+ break;
611
+ default:
612
+ TORCH_CHECK(false, "hash_topk_cpu: unsupported routed_topk ", routed_topk,
613
+ " for num_experts=256");
614
+ }
615
+ break;
616
+ case 384:
617
+ switch (routed_topk) {
618
+ case 6:
619
+ LAUNCH_HASH_TOPK_KERNEL(384, 6);
620
+ break;
621
+ case 7:
622
+ LAUNCH_HASH_TOPK_KERNEL(384, 7);
623
+ break;
624
+ case 8:
625
+ LAUNCH_HASH_TOPK_KERNEL(384, 8);
626
+ break;
627
+ default:
628
+ TORCH_CHECK(false, "hash_topk_cpu: unsupported routed_topk ", routed_topk,
629
+ " for num_experts=384");
630
+ }
631
+ break;
632
+ default:
633
+ TORCH_CHECK(false, "hash_topk_cpu: unsupported num_experts ", num_experts);
634
+ }
635
+ };
636
+
637
+ if (st == at::ScalarType::BFloat16) {
638
+ dispatch_fn.template operator()<at::BFloat16>();
639
+ } else if (st == at::ScalarType::Half) {
640
+ dispatch_fn.template operator()<at::Half>();
641
+ } else if (st == at::ScalarType::Float) {
642
+ dispatch_fn.template operator()<float>();
643
+ } else {
644
+ TORCH_CHECK(false, "Unsupported dtype for hash_topk_cpu: ", st);
645
+ }
646
+
647
+ return std::make_tuple(topk_weights, topk_ids);
648
+}
649
+
650
+// Adapted from
651
+// https://github.com/sgl-project/sglang/tree/main/sgl-kernel/csrc/cpu:
652
+// `topk_transform_512_cpu`/`topk_transform_512_cpu_kernel_impl`. Near-
653
+// verbatim port for the sparse indexer's DECODE path: hardcodes top-512
654
+// (the only `index_topk` DeepSeek-V4-Flash uses) and streams a min-heap over
655
+// the tail of the row instead of sorting the whole row.
656
+//
657
+// Differences from upstream:
658
+// - `out_page_indices` is still computed (required by the kernel contract)
659
+// but vLLM's caller only consumes `out_raw_indices` (local/compressed-
660
+// context positions -- `DeepseekV4CPUAttention.forward_mqa` resolves those
661
+// to physical slots itself via `map_local_to_global_slots_cpu`).
662
+// - `page_size` was already a runtime argument upstream, unlike
663
+// `fp8_paged_mqa_logits_cpu` (paged_mqa_logits.cpp), which needed one.
664
+namespace {
665
+
666
+constexpr int64_t kC4Topk = 512;
667
+
668
+template <typename T>
669
+inline int64_t load_index_value(const T* __restrict__ ptr, int64_t idx) {
670
+ return static_cast<int64_t>(ptr[idx]);
671
+}
672
+
673
+struct TopKTransformElem {
674
+ float score;
675
+ int32_t index;
676
+};
677
+
678
+struct TopKTransformMinHeapCmp {
679
+ bool operator()(const TopKTransformElem& lhs, const TopKTransformElem& rhs) const {
680
+ if (lhs.score == rhs.score) {
681
+ return lhs.index > rhs.index;
682
+ }
683
+ return lhs.score > rhs.score;
684
+ }
685
+};
686
+
687
+template <typename seq_t, typename page_t>
688
+void topk_transform_512_cpu_kernel_impl(
689
+ const float* __restrict__ scores,
690
+ const seq_t* __restrict__ seq_lens,
691
+ const page_t* __restrict__ page_tables,
692
+ int32_t* __restrict__ out_page_indices,
693
+ int32_t* __restrict__ out_raw_indices,
694
+ int64_t batch_size,
695
+ int64_t max_seq_len,
696
+ int64_t page_table_stride,
697
+ int64_t out_stride,
698
+ int64_t page_size) {
699
+ TORCH_CHECK(page_size > 0, "page_size must be positive");
700
+ TORCH_CHECK((page_size & (page_size - 1)) == 0, "page_size must be a power of 2");
701
+ const int page_bits = page_size > 1 ? static_cast<int>(std::log2(static_cast<double>(page_size))) : 0;
702
+ const int64_t page_mask = page_size - 1;
703
+
704
+ at::parallel_for(0, batch_size, 0, [&](int64_t begin, int64_t end) {
705
+ std::array<TopKTransformElem, kC4Topk> heap;
706
+
707
+ for (int64_t b = begin; b < end; ++b) {
708
+ const float* __restrict__ scores_row = scores + b * max_seq_len;
709
+ const page_t* __restrict__ page_table_row = page_tables + b * page_table_stride;
710
+ int32_t* __restrict__ out_page_row = out_page_indices + b * out_stride;
711
+ int32_t* __restrict__ out_raw_row = out_raw_indices == nullptr ? nullptr : out_raw_indices + b * out_stride;
712
+
713
+ int64_t seq_len = load_index_value(seq_lens, b);
714
+ seq_len = std::max<int64_t>(0, std::min<int64_t>(seq_len, max_seq_len));
715
+ const int64_t valid_topk = std::min<int64_t>(seq_len, kC4Topk);
716
+
717
+ auto store_slot = [&](int64_t slot, int32_t raw_index) {
718
+ if (raw_index < 0) {
719
+ out_page_row[slot] = -1;
720
+ if (out_raw_row != nullptr) {
721
+ out_raw_row[slot] = -1;
722
+ }
723
+ return;
724
+ }
725
+
726
+ const int64_t page_idx = static_cast<int64_t>(raw_index) >> page_bits;
727
+ const int64_t offset_in_page = static_cast<int64_t>(raw_index) & page_mask;
728
+ const int64_t physical_page = load_index_value(page_table_row, page_idx);
729
+ out_page_row[slot] = static_cast<int32_t>((physical_page << page_bits) | offset_in_page);
730
+ if (out_raw_row != nullptr) {
731
+ out_raw_row[slot] = raw_index;
732
+ }
733
+ };
734
+
735
+ if (seq_len <= kC4Topk) {
736
+ for (int64_t i = 0; i < valid_topk; ++i) {
737
+ store_slot(i, static_cast<int32_t>(i));
738
+ }
739
+ for (int64_t i = valid_topk; i < kC4Topk; ++i) {
740
+ store_slot(i, -1);
741
+ }
742
+ continue;
743
+ }
744
+
745
+ for (int64_t i = 0; i < kC4Topk; ++i) {
746
+ heap[i] = {scores_row[i], static_cast<int32_t>(i)};
747
+ }
748
+ std::make_heap(heap.begin(), heap.end(), TopKTransformMinHeapCmp());
749
+
750
+ for (int64_t i = kC4Topk; i < seq_len; ++i) {
751
+ const float score = scores_row[i];
752
+ const TopKTransformElem& current_min = heap.front();
753
+ if (score > current_min.score ||
754
+ (score == current_min.score && static_cast<int32_t>(i) < current_min.index)) {
755
+ std::pop_heap(heap.begin(), heap.end(), TopKTransformMinHeapCmp());
756
+ heap.back() = {score, static_cast<int32_t>(i)};
757
+ std::push_heap(heap.begin(), heap.end(), TopKTransformMinHeapCmp());
758
+ }
759
+ }
760
+
761
+ for (int64_t i = 0; i < kC4Topk; ++i) {
762
+ store_slot(i, heap[i].index);
763
+ }
764
+ }
765
+ });
766
+}
767
+
768
+} // namespace
769
+
770
+void topk_transform_512_cpu(
771
+ at::Tensor& scores,
772
+ at::Tensor& seq_lens,
773
+ at::Tensor& page_tables,
774
+ at::Tensor& out_page_indices,
775
+ int64_t page_size,
776
+ const std::optional<at::Tensor>& out_raw_indices) {
777
+ CHECK_INPUT(scores);
778
+ CHECK_INPUT(seq_lens);
779
+ CHECK_INPUT(page_tables);
780
+ CHECK_INPUT(out_page_indices);
781
+
782
+ TORCH_CHECK(scores.scalar_type() == at::kFloat, "scores must be float32");
783
+ TORCH_CHECK(out_page_indices.scalar_type() == at::kInt, "out_page_indices must be int32");
784
+ TORCH_CHECK(scores.dim() == 2, "scores must be a 2D tensor");
785
+ TORCH_CHECK(seq_lens.dim() == 1, "seq_lens must be a 1D tensor");
786
+ TORCH_CHECK(page_tables.dim() == 2, "page_tables must be a 2D tensor");
787
+ TORCH_CHECK(out_page_indices.dim() == 2, "out_page_indices must be a 2D tensor");
788
+
789
+ const int64_t batch_size = scores.size(0);
790
+ const int64_t max_seq_len = scores.size(1);
791
+ TORCH_CHECK(seq_lens.size(0) == batch_size, "seq_lens row count must match scores");
792
+ TORCH_CHECK(page_tables.size(0) == batch_size, "page_tables row count must match scores");
793
+ TORCH_CHECK(out_page_indices.size(0) == batch_size, "out_page_indices row count must match scores");
794
+ TORCH_CHECK(out_page_indices.size(1) >= kC4Topk, "out_page_indices must have at least 512 columns");
795
+
796
+ int32_t* raw_ptr = nullptr;
797
+ if (out_raw_indices.has_value()) {
798
+ at::Tensor raw = out_raw_indices.value();
799
+ CHECK_INPUT(raw);
800
+ TORCH_CHECK(raw.scalar_type() == at::kInt, "out_raw_indices must be int32");
801
+ TORCH_CHECK(raw.dim() == 2, "out_raw_indices must be a 2D tensor");
802
+ TORCH_CHECK(raw.sizes() == out_page_indices.sizes(), "out_raw_indices shape must match out_page_indices");
803
+ raw_ptr = raw.data_ptr<int32_t>();
804
+ }
805
+
806
+ if (batch_size == 0) {
807
+ return;
808
+ }
809
+
810
+ auto launch_with_seq_type = [&](auto seq_zero) {
811
+ using seq_t = decltype(seq_zero);
812
+ if (page_tables.scalar_type() == at::kInt) {
813
+ topk_transform_512_cpu_kernel_impl<seq_t, int32_t>(
814
+ scores.data_ptr<float>(), seq_lens.data_ptr<seq_t>(), page_tables.data_ptr<int32_t>(),
815
+ out_page_indices.data_ptr<int32_t>(), raw_ptr, batch_size, max_seq_len, page_tables.stride(0),
816
+ out_page_indices.stride(0), page_size);
817
+ } else if (page_tables.scalar_type() == at::kLong) {
818
+ topk_transform_512_cpu_kernel_impl<seq_t, int64_t>(
819
+ scores.data_ptr<float>(), seq_lens.data_ptr<seq_t>(), page_tables.data_ptr<int64_t>(),
820
+ out_page_indices.data_ptr<int32_t>(), raw_ptr, batch_size, max_seq_len, page_tables.stride(0),
821
+ out_page_indices.stride(0), page_size);
822
+ } else {
823
+ TORCH_CHECK(false, "page_tables must be int32 or int64");
824
+ }
825
+ };
826
+
827
+ if (seq_lens.scalar_type() == at::kInt) {
828
+ launch_with_seq_type(int32_t{0});
829
+ } else if (seq_lens.scalar_type() == at::kLong) {
830
+ launch_with_seq_type(int64_t{0});
831
+ } else {
832
+ TORCH_CHECK(false, "seq_lens must be int32 or int64");
833
+ }
834
+}
csrc/cpu/torch_bindings.cpp
+213
-16
@@ -66,18 +66,18 @@ at::Tensor convert_weight_packed(at::Tensor& weight);
66
67
at::Tensor convert_scale_packed(at::Tensor& scale);
68
69
-at::Tensor fused_experts_cpu(
70
- at::Tensor& hidden_states, at::Tensor& w1, at::Tensor& w2,
71
- at::Tensor& topk_weights, at::Tensor& topk_ids, bool inplace,
72
- int64_t moe_comp_method, const std::optional<at::Tensor>& w1_scale,
73
- const std::optional<at::Tensor>& w2_scale,
74
- const std::optional<at::Tensor>& w1_zero,
75
- const std::optional<at::Tensor>& w2_zero,
76
- const std::optional<std::vector<int64_t>> block_size,
77
- const std::optional<at::Tensor>& w1_bias,
78
- const std::optional<at::Tensor>& w2_bias,
79
- const std::optional<double>& alpha, const std::optional<double>& limit,
80
- bool is_vnni);
69
+void fused_experts_cpu(at::Tensor& out, at::Tensor& hidden_states,
70
+ at::Tensor& w1, at::Tensor& w2, at::Tensor& topk_weights,
71
+ at::Tensor& topk_ids, int64_t moe_comp_method,
72
+ const std::optional<at::Tensor>& w1_scale,
73
+ const std::optional<at::Tensor>& w2_scale,
74
+ const std::optional<at::Tensor>& w1_zero,
75
+ const std::optional<at::Tensor>& w2_zero,
76
+ const std::optional<std::vector<int64_t>> block_size,
77
+ const std::optional<at::Tensor>& w1_bias,
78
+ const std::optional<at::Tensor>& w2_bias,
79
+ const std::optional<double>& alpha,
80
+ const std::optional<double>& limit, bool is_vnni);
81
82
at::Tensor int8_scaled_mm_with_quant(at::Tensor& mat1, at::Tensor& mat2,
83
at::Tensor& scales2,
@@ -124,6 +124,99 @@ void concat_and_cache_mla_cpu(const at::Tensor& kv_c_normed,
124
const at::Tensor& k_pe, at::Tensor& kv_cache,
125
const at::Tensor& slot_mapping);
126
127
+// Adapted from sglang: DeepSeek-V4 mHC gating kernels
128
+std::tuple<at::Tensor, at::Tensor, at::Tensor> hc_pre_fused_cpu(
129
+ at::Tensor& x, at::Tensor& hc_fn, at::Tensor& hc_scale, at::Tensor& hc_base,
130
+ int64_t hc_mult, int64_t sinkhorn_iters, double rms_eps, double hc_eps);
131
+
132
+at::Tensor hc_post_fused_cpu(at::Tensor& x, at::Tensor& residual,
133
+ at::Tensor& post, at::Tensor& comb);
134
+
135
+at::Tensor hc_head_fused_cpu(at::Tensor& x, at::Tensor& hc_fn,
136
+ at::Tensor& hc_scale, at::Tensor& hc_base,
137
+ double hc_eps, double norm_eps);
138
+
139
+// Adapted from sglang: DeepSeek-V4 fp8_ds_mla cache-write kernel
140
+at::Tensor fused_qnorm_rope_kv_insert_cpu(at::Tensor& q, at::Tensor& kv,
141
+ at::Tensor& positions,
142
+ at::Tensor& swa_kv_cache_2d,
143
+ at::Tensor& slot_mapping,
144
+ at::Tensor& cos_sin_cache,
145
+ int64_t q_head_padded, double eps,
146
+ int64_t cache_block_size);
147
+
148
+// vLLM-native: DeepSeek-V4 sparse MQA attention kernel for the fp8_ds_mla
149
+// cache layout (SWA window + optional compressed top-k index sets).
150
+void flash_mla_with_kvcache_cpu(
151
+ at::Tensor& out, at::Tensor& q, at::Tensor& window_cache_2d,
152
+ at::Tensor& window_slots, int64_t window_block_size,
153
+ at::Tensor& compressed_cache_2d, at::Tensor& compressed_slots,
154
+ int64_t compressed_block_size, at::Tensor& attn_sink, double scale);
155
+
156
+// Adapted from sglang: DeepSeek-V4 compressor state-cache write + fused
157
+// compress+RMSNorm+quant+RoPE+store kernels.
158
+void save_partial_states_cpu(at::Tensor& kv, at::Tensor& score, at::Tensor& ape,
159
+ at::Tensor& positions, at::Tensor& state_cache,
160
+ at::Tensor& slot_mapping);
161
+
162
+void compress_norm_rope_store_cpu(
163
+ at::Tensor& state_cache, at::Tensor& gather_slots, at::Tensor& positions,
164
+ at::Tensor& kv_slot_mapping, at::Tensor& rms_norm_weight,
165
+ double rms_norm_eps, at::Tensor& cos_sin_cache, at::Tensor& kv_cache_2d,
166
+ int64_t kv_cache_block_size, int64_t compress_ratio);
167
+
168
+// head_dim=128 indexer-compressor variant of the above (single fp8 quant
169
+// block + raw fp32 scale -- matches the paged indexer K-cache layout
170
+// fp8_paged_mqa_logits_cpu/topk_transform_512_cpu read).
171
+void compress_norm_rope_store_indexer_cpu(
172
+ at::Tensor& state_cache, at::Tensor& gather_slots, at::Tensor& positions,
173
+ at::Tensor& kv_slot_mapping, at::Tensor& rms_norm_weight,
174
+ double rms_norm_eps, at::Tensor& cos_sin_cache, at::Tensor& kv_cache_2d,
175
+ int64_t kv_cache_block_size, int64_t compress_ratio);
176
+
177
+// Adapted from sglang: DeepSeek-V4 sparse indexer paged MQA-logits +
178
+// top-512 transform kernels (DECODE path -- reads the paged K-cache
179
+// directly via `page_table`, whole batch in one call).
180
+at::Tensor fp8_paged_mqa_logits_cpu(at::Tensor& q_fp8, at::Tensor& kvcache_fp8,
181
+ at::Tensor& weight, at::Tensor& seq_lens,
182
+ at::Tensor& page_table, int64_t block_size,
183
+ int64_t max_seq_len);
184
+
185
+void topk_transform_512_cpu(at::Tensor& scores, at::Tensor& seq_lens,
186
+ at::Tensor& page_tables,
187
+ at::Tensor& out_page_indices, int64_t page_size,
188
+ const std::optional<at::Tensor>& out_raw_indices);
189
+
190
+// Adapted from sglang: DeepSeek-V4 MoE routing kernels (flat biased top-k and
191
+// hash-routed-layer top-k).
192
+std::tuple<at::Tensor, at::Tensor> biased_topk_cpu(
193
+ at::Tensor& hidden_states, at::Tensor& gating_output,
194
+ at::Tensor& correction_bias, int64_t topk, bool renormalize,
195
+ std::string scoring_func, int64_t num_fused_shared_experts,
196
+ std::optional<double> routed_scaling_factor,
197
+ bool apply_routed_scaling_factor_on_output);
198
+
199
+std::tuple<at::Tensor, at::Tensor> hash_topk_cpu(
200
+ at::Tensor& gating_output, at::Tensor& tid2eid, int64_t topk,
201
+ std::string scoring_func, int64_t num_fused_shared_experts,
202
+ int64_t num_experts, double routed_scaling_factor);
203
+
204
+// vLLM-native: DeepSeek-V4 sparse indexer Q-side RoPE + FP8 quant kernel
205
+// (fp8 path only; MXFP4 stays on triton-cpu).
206
+void fused_indexer_q_rope_quant_cpu(at::Tensor& positions, at::Tensor& index_q,
207
+ at::Tensor& index_q_cos_sin_cache,
208
+ at::Tensor& index_q_fp8,
209
+ at::Tensor& index_weights,
210
+ double index_weights_softmax_scale,
211
+ double index_weights_head_scale,
212
+ at::Tensor& index_weights_out);
213
+
214
+// vLLM-native: inverse GPT-J RoPE for DeepSeek-V4 CPU attention's `_o_proj`
215
+// output de-rotation.
216
+at::Tensor inverse_gptj_rope_o_proj_cpu(at::Tensor& o, at::Tensor& positions,
217
+ at::Tensor& cos_sin_cache,
218
+ int64_t rope_dim);
219
+
220
// Adapted from sglang: INT4 W4A8 kernels
221
std::tuple<at::Tensor, at::Tensor, at::Tensor> convert_weight_packed_scale_zp(
222
at::Tensor qweight, // awq: (*, K, N / 8) || gptq: (*, K / 8, N) , int32
@@ -515,13 +608,13 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
608
ops.def("convert_scale_packed(Tensor! scale) -> Tensor");
609
ops.impl("convert_scale_packed", torch::kCPU, &convert_scale_packed);
610
ops.def(
518
- "fused_experts_cpu(Tensor hidden_states, Tensor w1, Tensor w2, Tensor "
519
- "topk_weights, Tensor topk_ids, bool "
520
- "inplace, int moe_comp_method, Tensor? w1_scale, Tensor? w2_scale, "
611
+ "fused_experts_cpu(Tensor(a0!) out, Tensor hidden_states, Tensor w1, "
612
+ "Tensor w2, Tensor topk_weights, Tensor topk_ids, "
613
+ "int moe_comp_method, Tensor? w1_scale, Tensor? w2_scale, "
614
"Tensor? w1_zero, Tensor? w2_zero, int[]? block_size, "
615
"Tensor? w1_bias, Tensor? w2_bias, float? alpha, float? limit, "
616
"bool is_vnni) -> "
524
- "Tensor");
617
+ "()");
618
ops.impl("fused_experts_cpu", torch::kCPU, &fused_experts_cpu);
619
ops.def(
620
"int8_scaled_mm_with_quant(Tensor mat1, Tensor mat2, Tensor scales2, "
@@ -583,6 +676,110 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
676
"concat_and_cache_mla_cpu(Tensor kv_c_normed, Tensor k_pe, "
677
"Tensor(a!) kv_cache, Tensor slot_mapping) -> ()");
678
ops.impl("concat_and_cache_mla_cpu", torch::kCPU, &concat_and_cache_mla_cpu);
679
+
680
+ // Adapted from sglang: DeepSeek-V4 mHC gating kernels
681
+ ops.def(
682
+ "hc_pre_fused_cpu(Tensor x, Tensor hc_fn, Tensor hc_scale, "
683
+ "Tensor hc_base, int hc_mult, int sinkhorn_iters, float rms_eps, "
684
+ "float hc_eps) -> (Tensor, Tensor, Tensor)");
685
+ ops.impl("hc_pre_fused_cpu", torch::kCPU, &hc_pre_fused_cpu);
686
+
687
+ ops.def(
688
+ "hc_post_fused_cpu(Tensor x, Tensor residual, Tensor post, "
689
+ "Tensor comb) -> Tensor");
690
+ ops.impl("hc_post_fused_cpu", torch::kCPU, &hc_post_fused_cpu);
691
+
692
+ ops.def(
693
+ "hc_head_fused_cpu(Tensor x, Tensor hc_fn, Tensor "
694
+ "hc_scale, Tensor hc_base, float hc_eps, float norm_eps) -> "
695
+ "Tensor");
696
+ ops.impl("hc_head_fused_cpu", torch::kCPU, &hc_head_fused_cpu);
697
+
698
+ // Adapted from sglang: DeepSeek-V4 fp8_ds_mla cache-write kernel
699
+ ops.def(
700
+ "fused_qnorm_rope_kv_insert_cpu(Tensor q, Tensor kv, Tensor positions, "
701
+ "Tensor(a!) swa_kv_cache_2d, Tensor slot_mapping, Tensor "
702
+ "cos_sin_cache, int q_head_padded, float eps, int cache_block_size) "
703
+ "-> Tensor");
704
+ ops.impl("fused_qnorm_rope_kv_insert_cpu", torch::kCPU,
705
+ &fused_qnorm_rope_kv_insert_cpu);
706
+
707
+ // vLLM-native: DeepSeek-V4 sparse MQA attention kernel
708
+ ops.def(
709
+ "flash_mla_with_kvcache_cpu(Tensor(a!) out, Tensor q, Tensor "
710
+ "window_cache_2d, Tensor window_slots, int window_block_size, Tensor "
711
+ "compressed_cache_2d, Tensor compressed_slots, int "
712
+ "compressed_block_size, Tensor attn_sink, float scale) -> ()");
713
+ ops.impl("flash_mla_with_kvcache_cpu", torch::kCPU,
714
+ &flash_mla_with_kvcache_cpu);
715
+
716
+ // Adapted from sglang: DeepSeek-V4 compressor kernels
717
+ ops.def(
718
+ "save_partial_states_cpu(Tensor kv, Tensor score, Tensor ape, "
719
+ "Tensor positions, Tensor(a!) state_cache, Tensor slot_mapping) "
720
+ "-> ()");
721
+ ops.impl("save_partial_states_cpu", torch::kCPU, &save_partial_states_cpu);
722
+
723
+ ops.def(
724
+ "compress_norm_rope_store_cpu(Tensor state_cache, Tensor "
725
+ "gather_slots, Tensor positions, Tensor kv_slot_mapping, Tensor "
726
+ "rms_norm_weight, float rms_norm_eps, Tensor cos_sin_cache, "
727
+ "Tensor(a!) kv_cache_2d, int kv_cache_block_size, int compress_ratio) "
728
+ "-> ()");
729
+ ops.impl("compress_norm_rope_store_cpu", torch::kCPU,
730
+ &compress_norm_rope_store_cpu);
731
+
732
+ ops.def(
733
+ "compress_norm_rope_store_indexer_cpu(Tensor state_cache, Tensor "
734
+ "gather_slots, Tensor positions, Tensor kv_slot_mapping, Tensor "
735
+ "rms_norm_weight, float rms_norm_eps, Tensor cos_sin_cache, "
736
+ "Tensor(a!) kv_cache_2d, int kv_cache_block_size, int compress_ratio) "
737
+ "-> ()");
738
+ ops.impl("compress_norm_rope_store_indexer_cpu", torch::kCPU,
739
+ &compress_norm_rope_store_indexer_cpu);
740
+
741
+ // Adapted from sglang: DeepSeek-V4 sparse indexer paged MQA-logits kernel
742
+ // (DECODE path)
743
+ ops.def(
744
+ "fp8_paged_mqa_logits_cpu(Tensor q_fp8, Tensor kvcache_fp8, Tensor "
745
+ "weight, Tensor seq_lens, Tensor page_table, int block_size, int "
746
+ "max_seq_len) -> Tensor");
747
+ ops.impl("fp8_paged_mqa_logits_cpu", torch::kCPU, &fp8_paged_mqa_logits_cpu);
748
+
749
+ ops.def(
750
+ "topk_transform_512_cpu(Tensor scores, Tensor seq_lens, Tensor "
751
+ "page_tables, Tensor(a!) out_page_indices, int page_size, Tensor(b!)? "
752
+ "out_raw_indices) -> ()");
753
+ ops.impl("topk_transform_512_cpu", torch::kCPU, &topk_transform_512_cpu);
754
+
755
+ ops.def(
756
+ "fused_indexer_q_rope_quant_cpu(Tensor positions, Tensor index_q, "
757
+ "Tensor index_q_cos_sin_cache, Tensor(a!) index_q_fp8, Tensor "
758
+ "index_weights, float index_weights_softmax_scale, float "
759
+ "index_weights_head_scale, Tensor(b!) index_weights_out) -> ()");
760
+ ops.impl("fused_indexer_q_rope_quant_cpu", torch::kCPU,
761
+ &fused_indexer_q_rope_quant_cpu);
762
+
763
+ // Adapted from sglang: DeepSeek-V4 MoE routing kernels
764
+ ops.def(
765
+ "biased_topk_cpu(Tensor hidden_states, Tensor gating_output, Tensor "
766
+ "correction_bias, int topk, bool renormalize, str scoring_func, int "
767
+ "num_fused_shared_experts, float? routed_scaling_factor, bool "
768
+ "apply_routed_scaling_factor_on_output) -> (Tensor, Tensor)");
769
+ ops.impl("biased_topk_cpu", torch::kCPU, &biased_topk_cpu);
770
+
771
+ ops.def(
772
+ "hash_topk_cpu(Tensor gating_output, Tensor tid2eid, int topk, str "
773
+ "scoring_func, int num_fused_shared_experts, int num_experts, float "
774
+ "routed_scaling_factor) -> (Tensor, Tensor)");
775
+ ops.impl("hash_topk_cpu", torch::kCPU, &hash_topk_cpu);
776
+
777
+ // vLLM-native: inverse GPT-J RoPE for _o_proj's output de-rotation
778
+ ops.def(
779
+ "inverse_gptj_rope_o_proj_cpu(Tensor o, Tensor positions, Tensor "
780
+ "cos_sin_cache, int rope_dim) -> Tensor");
781
+ ops.impl("inverse_gptj_rope_o_proj_cpu", torch::kCPU,
782
+ &inverse_gptj_rope_o_proj_cpu);
783
#endif
784
785
#if (defined(__AVX512BF16__) && defined(__AVX512F__) && \
tests/kernels/moe/test_cpu_fused_moe.py
+20
-14
@@ -26,8 +26,8 @@ from vllm.model_executor.layers.fused_moe.experts.cpu_moe import (
26
CPUUnquantizedExperts,
27
PowerCPUUnquantizedExperts,
28
X86CPUUnquantizedExperts,
29
- select_experts,
29
)
30
+from vllm.model_executor.layers.fused_moe.router.cpu_router import _softmax_topk
31
from vllm.platforms import CpuArchEnum, current_platform
32
from vllm.utils.torch_utils import set_random_seed
33
@@ -284,7 +284,9 @@ def test_cpu_fused_moe(
284
285
packed_w13 = cpu_prepack_moe_weight(w13, isa)
286
packed_w2 = cpu_prepack_moe_weight(w2, isa)
287
- output = cpu_fused_moe(
287
+ output = torch.empty_like(input)
288
+ cpu_fused_moe(
289
+ output,
290
input,
291
packed_w13,
292
packed_w2,
@@ -368,7 +370,9 @@ def test_cpu_fused_moe_int8(
370
)
371
packed_w13 = cpu_prepack_moe_weight_int8(w13, isa)
372
packed_w2 = cpu_prepack_moe_weight_int8(w2, isa)
371
- output = cpu_fused_moe_int8(
373
+ output = torch.empty_like(input)
374
+ cpu_fused_moe_int8(
375
+ output,
376
input,
377
packed_w13,
378
packed_w2,
@@ -521,15 +525,10 @@ def test_cpu_fused_moe_unaligned_intermediate_size(
525
w2_bias = torch.randn((expert_num, hidden_size), dtype=dtype) / (
526
0.5 * hidden_size**0.5
527
)
524
- # Route with the same helper apply() uses internally, so the reference
525
- # only differs from the kernel in how the experts are evaluated.
526
- topk_weight, topk_ids = select_experts(
527
- hidden_states=input,
528
- router_logits=router_logits,
529
- top_k=topk_num,
530
- use_grouped_topk=False,
531
- renormalize=False,
532
- )
528
+ # Route with the same helper the router uses internally, so the
529
+ # reference only differs from the kernel in how the experts are
530
+ # evaluated.
531
+ topk_weight, topk_ids = _softmax_topk(router_logits, topk_num, False)
532
533
ref_output = ref_fused_moe(
534
input, w13, w2, w13_bias, w2_bias, topk_weight, topk_ids, act
@@ -575,15 +574,22 @@ def test_cpu_fused_moe_unaligned_intermediate_size(
574
assert experts.isa == isa
575
experts.process_weights_after_loading(layer)
576
578
- output = experts.apply(
577
+ output = torch.empty_like(input)
578
+ experts.apply(
579
+ output=output,
580
hidden_states=input,
581
w1=layer.w13_weight,
582
w2=layer.w2_weight,
582
- router_logits=router_logits,
583
+ topk_weights=topk_weight,
584
+ topk_ids=topk_ids,
585
activation=act,
586
global_num_experts=expert_num,
587
expert_map=None,
588
a1q_scale=None,
589
+ a2_scale=None,
590
+ workspace13=torch.empty(0),
591
+ workspace2=torch.empty(0),
592
+ expert_tokens_meta=None,
593
apply_router_weight_on_input=False,
594
)
595
tests/kernels/moe/test_cpu_quant_fused_moe.py
+48
-29
@@ -197,14 +197,15 @@ def test_w8a16_block_fp8_cpu_fused_moe(M, N, K, E, topk, seed):
197
198
pw1, pw2 = _prepack_experts(w1), _prepack_experts(w2)
199
200
- # Test inplace=False against reference
201
- out = ops.fused_experts_cpu(
200
+ # Test writing into a fresh output buffer against reference
201
+ out = torch.empty_like(a)
202
+ ops.fused_experts_cpu(
203
+ out,
204
a.clone(),
205
pw1,
206
pw2,
207
topk_weight,
208
topk_ids,
207
- False,
209
ops.CPUQuantMethod.FP8_W8A16,
210
w1_s,
211
w2_s,
@@ -215,14 +216,16 @@ def test_w8a16_block_fp8_cpu_fused_moe(M, N, K, E, topk, seed):
216
)
217
torch.testing.assert_close(ref_out.bfloat16(), out, atol=1e-2, rtol=1e-2)
218
218
- # Test inplace=True produces identical output
219
- out_inplace = ops.fused_experts_cpu(
220
- a.clone(),
219
+ # Test writing into the input tensor itself (aliased in-place) produces
220
+ # an identical result
221
+ a_inplace = a.clone()
222
+ ops.fused_experts_cpu(
223
+ a_inplace,
224
+ a_inplace,
225
pw1,
226
pw2,
227
topk_weight,
228
topk_ids,
225
- True,
229
ops.CPUQuantMethod.FP8_W8A16,
230
w1_s,
231
w2_s,
@@ -231,7 +234,7 @@ def test_w8a16_block_fp8_cpu_fused_moe(M, N, K, E, topk, seed):
234
BLOCK_SIZE,
235
is_vnni=True,
236
)
234
- torch.testing.assert_close(out_inplace, out, atol=0, rtol=0)
237
+ torch.testing.assert_close(a_inplace, out, atol=0, rtol=0)
238
239
240
def test_w8a16_block_fp8_cpu_fused_moe_small_expert_blocks():
@@ -249,13 +252,14 @@ def test_w8a16_block_fp8_cpu_fused_moe_small_expert_blocks():
252
a, w1, w2, w1_s, w2_s, topk_weight, topk_ids, BLOCK_SIZE
253
)
254
pw1, pw2 = _prepack_experts(w1), _prepack_experts(w2)
252
- out = ops.fused_experts_cpu(
255
+ out = torch.empty_like(a)
256
+ ops.fused_experts_cpu(
257
+ out,
258
a,
259
pw1,
260
pw2,
261
topk_weight,
262
topk_ids,
258
- False,
263
ops.CPUQuantMethod.FP8_W8A16,
264
w1_s,
265
w2_s,
@@ -469,13 +473,15 @@ def test_mxfp4_cpu_fused_moe(M, N, K, E, topk, seed):
473
pw2, pw2s = _prepack_mxfp4_experts(w2q, w2s)
474
475
# Kernel
472
- out = ops.fused_experts_cpu(
473
- a.clone(),
476
+ a_in = a.clone()
477
+ out = torch.empty_like(a_in)
478
+ ops.fused_experts_cpu(
479
+ out,
480
+ a_in,
481
pw1,
482
pw2,
483
topk_weight,
484
topk_ids,
478
- False, # inplace
485
ops.CPUQuantMethod.MXFP4,
486
pw1s, # w1_scale
487
pw2s, # w2_scale
@@ -511,13 +517,14 @@ def test_mxfp4_cpu_fused_moe_small_expert_blocks():
517
518
pw1, pw1s = _prepack_mxfp4_experts(w1q, w1s)
519
pw2, pw2s = _prepack_mxfp4_experts(w2q, w2s)
514
- out = ops.fused_experts_cpu(
520
+ out = torch.empty_like(a)
521
+ ops.fused_experts_cpu(
522
+ out,
523
a,
524
pw1,
525
pw2,
526
topk_weight,
527
topk_ids,
520
- False,
528
ops.CPUQuantMethod.MXFP4,
529
pw1s,
530
pw2s,
@@ -569,13 +576,14 @@ def test_mxfp4_cpu_zero_codes_stay_zero(zero_byte, e8m0):
576
577
pw1, pw1s = _prepack_mxfp4_experts(w1q, w1s)
578
pw2, pw2s = _prepack_mxfp4_experts(w2q, w2s)
572
- out = ops.fused_experts_cpu(
579
+ out = torch.empty_like(a)
580
+ ops.fused_experts_cpu(
581
+ out,
582
a,
583
pw1,
584
pw2,
585
topk_weight,
586
topk_ids,
578
- False,
587
ops.CPUQuantMethod.MXFP4,
588
pw1s,
589
pw2s,
@@ -630,13 +638,14 @@ def test_mxfp4_cpu_zero_codes_mixed_with_nonzero(e8m0):
638
639
pw1, pw1s = _prepack_mxfp4_experts(w1q, w1s)
640
pw2, pw2s = _prepack_mxfp4_experts(w2q, w2s)
633
- out = ops.fused_experts_cpu(
641
+ out = torch.empty_like(a)
642
+ ops.fused_experts_cpu(
643
+ out,
644
a,
645
pw1,
646
pw2,
647
topk_weight,
648
topk_ids,
639
- False,
649
ops.CPUQuantMethod.MXFP4,
650
pw1s,
651
pw2s,
@@ -689,13 +698,15 @@ def test_mxfp4_cpu_fused_moe_bias_swiglu(M, N, K, E, topk, seed):
698
pw2, pw2s = _prepack_mxfp4_experts(w2q, w2s)
699
700
# Kernel
692
- out = ops.fused_experts_cpu(
693
- a.clone(),
701
+ a_in = a.clone()
702
+ out = torch.empty_like(a_in)
703
+ ops.fused_experts_cpu(
704
+ out,
705
+ a_in,
706
pw1,
707
pw2,
708
topk_weight,
709
topk_ids,
698
- False, # inplace
710
ops.CPUQuantMethod.MXFP4,
711
pw1s, # w1_scale
712
pw2s, # w2_scale
@@ -944,13 +955,15 @@ def test_int4_w4a16_cpu_fused_moe(M, N, K, E, topk, group_size, quant_algo, seed
955
)
956
)
957
947
- out = ops.fused_experts_cpu(
948
- a.clone(),
958
+ a_in = a.clone()
959
+ out = torch.empty_like(a_in)
960
+ ops.fused_experts_cpu(
961
+ out,
962
+ a_in,
963
blocked_w1,
964
blocked_w2,
965
topk_weight,
966
topk_ids,
953
- False, # inplace
967
ops.CPUQuantMethod.INT4_W4A8,
968
blocked_s1,
969
blocked_s2,
@@ -1051,7 +1064,11 @@ INT8_MOE_CONFIGS = [
1064
@pytest.mark.parametrize("is_vnni", [False, True])
1065
@pytest.mark.parametrize("inplace", [False, True])
1066
def test_int8_w8a8_cpu_fused_moe(M, N, K, E, topk, seed, is_vnni, inplace):
1054
- """Test fused_experts_cpu INT8 W8A8 against torch reference."""
1067
+ """Test fused_experts_cpu INT8 W8A8 against torch reference.
1068
+
1069
+ ``inplace`` exercises writing into a fresh output buffer vs. aliasing
1070
+ the input tensor as the output buffer.
1071
+ """
1072
set_random_seed(seed)
1073
1074
a = torch.randn(M, K, dtype=torch.bfloat16) / (0.5 * K**0.5)
@@ -1067,13 +1084,15 @@ def test_int8_w8a8_cpu_fused_moe(M, N, K, E, topk, seed, is_vnni, inplace):
1084
w1 = _prepack_experts(w1_q) if is_vnni else w1_q
1085
w2 = _prepack_experts(w2_q) if is_vnni else w2_q
1086
1070
- out = ops.fused_experts_cpu(
1071
- a.clone(),
1087
+ a_in = a.clone()
1088
+ out = a_in if inplace else torch.empty_like(a_in)
1089
+ ops.fused_experts_cpu(
1090
+ out,
1091
+ a_in,
1092
w1,
1093
w2,
1094
topk_weight,
1095
topk_ids,
1076
- inplace,
1096
ops.CPUQuantMethod.INT8_W8A8,
1097
w1_s,
1098
w2_s,
tests/kernels/moe/test_zen_cpu_fused_moe.py
+30
-40
@@ -2,8 +2,6 @@
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
"""Tests for ZenTorch CPU fused MoE dispatch and forward."""
4
5
-from types import SimpleNamespace
6
-
5
import pytest
6
import torch
7
@@ -27,8 +25,8 @@ from vllm.model_executor.layers.fused_moe.config import (
25
)
26
from vllm.model_executor.layers.fused_moe.experts.cpu_moe import (
27
CPUUnquantizedExperts,
30
- select_experts,
28
)
29
+from vllm.model_executor.layers.fused_moe.router.cpu_router import _softmax_topk
30
from vllm.platforms import current_platform
31
from vllm.utils.torch_utils import set_random_seed
32
@@ -95,18 +93,19 @@ def _make_layer_and_config(
93
94
95
@pytest.mark.parametrize("act", ZENTORCH_ACT)
98
-def test_zen_cpu_fused_moe_dispatches_to_zentorch(act: MoEActivation):
99
- """When zentorch MoE is supported, experts select the zentorch path."""
100
- layer, moe_config = _make_layer_and_config(act)
101
- w13_before = layer.w13_weight.detach().clone()
102
- w2_before = layer.w2_weight.detach().clone()
103
-
104
- experts = CPUUnquantizedExperts(moe_config, FusedMoEQuantConfig.make())
105
- experts.process_weights_after_loading(layer)
96
+def test_zen_cpu_fused_moe_config_supported(act: MoEActivation):
97
+ """CPUUnquantizedExperts is selected when zentorch MoE is supported."""
98
+ _, moe_config = _make_layer_and_config(act)
99
+ assert is_zentorch_moe_config_supported(moe_config)
100
107
- assert experts._use_zentorch is True
108
- torch.testing.assert_close(layer.w13_weight, w13_before)
109
- torch.testing.assert_close(layer.w2_weight, w2_before)
101
+ supported, reason = CPUUnquantizedExperts.is_supported_config(
102
+ CPUUnquantizedExperts,
103
+ moe_config,
104
+ None,
105
+ None,
106
+ mk.FusedMoEActivationFormat.Standard,
107
+ )
108
+ assert supported, reason
109
110
111
def test_zen_cpu_fused_moe_config_supported_unaligned_intermediate():
@@ -169,26 +168,34 @@ def test_zen_cpu_fused_moe_forward(
168
experts.process_weights_after_loading(layer)
169
assert experts._use_zentorch is True
170
172
- topk_weight, topk_ids = select_experts(
173
- hidden_states=input,
174
- router_logits=router_logits,
175
- top_k=topk_num,
176
- use_grouped_topk=False,
177
- renormalize=False,
178
- )
171
+ w13_before = layer.w13_weight.detach().clone()
172
+ w2_before = layer.w2_weight.detach().clone()
173
180
- output = experts.apply(
174
+ topk_weight, topk_ids = _softmax_topk(router_logits, topk_num, False)
175
+
176
+ output = torch.empty_like(input)
177
+ experts.apply(
178
+ output=output,
179
hidden_states=input,
180
w1=layer.w13_weight,
181
w2=layer.w2_weight,
184
- router_logits=router_logits,
182
+ topk_weights=topk_weight,
183
+ topk_ids=topk_ids,
184
activation=act,
185
global_num_experts=expert_num,
186
expert_map=None,
187
a1q_scale=None,
188
+ a2_scale=None,
189
+ workspace13=torch.empty(0),
190
+ workspace2=torch.empty(0),
191
+ expert_tokens_meta=None,
192
apply_router_weight_on_input=False,
193
)
194
195
+ # zentorch needs no prepacking; weights are consumed as loaded.
196
+ torch.testing.assert_close(layer.w13_weight, w13_before)
197
+ torch.testing.assert_close(layer.w2_weight, w2_before)
198
+
199
ref_output = ref_fused_moe(
200
input,
201
layer.w13_weight,
@@ -202,20 +209,3 @@ def test_zen_cpu_fused_moe_forward(
209
210
atol, rtol = get_default_atol(output), get_default_rtol(output)
211
torch.testing.assert_close(output, ref_output, atol=atol, rtol=rtol)
205
-
206
-
207
-def test_zen_cpu_fused_moe_skips_without_activation():
208
- """Layers without an activation attribute fall back from zentorch."""
209
- layer = SimpleNamespace()
210
- moe_config = _make_moe_config(
211
- expert_num=8,
212
- hidden_size=128,
213
- intermediate_size=128,
214
- topk_num=4,
215
- dtype=torch.bfloat16,
216
- act=MoEActivation.SILU,
217
- )
218
- experts = CPUUnquantizedExperts(moe_config, FusedMoEQuantConfig.make())
219
- experts.process_weights_after_loading(layer)
220
-
221
- assert experts._use_zentorch is False
tests/kernels/moe/test_zen_cpu_int8_moe.py
+12
-14
@@ -18,14 +18,12 @@ from vllm.model_executor.layers.fused_moe.config import (
18
FusedMoEParallelConfig,
19
int8_w8a8_moe_quant_config,
20
)
21
-from vllm.model_executor.layers.fused_moe.experts.cpu_moe import (
22
- ZenCPUExpertsInt8,
23
- select_experts,
24
-)
21
+from vllm.model_executor.layers.fused_moe.experts.cpu_moe import ZenCPUExpertsInt8
22
from vllm.model_executor.layers.fused_moe.oracle.int8 import (
23
Int8MoeBackend,
24
backend_to_kernel_cls,
25
)
26
+from vllm.model_executor.layers.fused_moe.router.cpu_router import _softmax_topk
27
from vllm.model_executor.layers.quantization.utils.quant_utils import (
28
kInt8DynamicTokenSym,
29
kInt8StaticChannelSym,
@@ -199,19 +197,27 @@ def test_zen_int8_dispatch_contract(
197
0.5 * HIDDEN_SIZE**0.5
198
)
199
router_logits = torch.randn((BATCH_SIZE, EXPERT_NUM), dtype=torch.bfloat16)
200
+ topk_weights, topk_ids = _softmax_topk(router_logits, topk_num, False)
201
202
experts = ZenCPUExpertsInt8(moe_config, quant_config)
203
# zentorch_fused_moe takes no expert_map.
204
assert not experts.supports_expert_map()
206
- output = experts.apply(
205
+ output = torch.empty_like(input)
206
+ experts.apply(
207
+ output=output,
208
hidden_states=input,
209
w1=w13,
210
w2=w2,
210
- router_logits=router_logits,
211
+ topk_weights=topk_weights,
212
+ topk_ids=topk_ids,
213
activation=act,
214
global_num_experts=EXPERT_NUM,
215
expert_map=None,
216
a1q_scale=None,
217
+ a2_scale=None,
218
+ workspace13=torch.empty(0),
219
+ workspace2=torch.empty(0),
220
+ expert_tokens_meta=None,
221
apply_router_weight_on_input=False,
222
)
223
@@ -230,14 +236,6 @@ def test_zen_int8_dispatch_contract(
236
else:
237
assert call["w13_bias"] is None and call["w2_bias"] is None
238
233
- topk_weights, topk_ids = select_experts(
234
- hidden_states=input,
235
- router_logits=router_logits,
236
- use_grouped_topk=False,
237
- top_k=topk_num,
238
- renormalize=False,
239
- scoring_func="softmax",
240
- )
239
ref = ref_fused_moe_int8(
240
input,
241
w13,
tests/kernels/test_deepseek_v4_cpu_kernels.py
new
+1494
@@ -0,0 +1,1494 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""CPU-kernel-vs-eager-reference parity tests for the DeepSeek-V4 CPU port's
4
+sgl-kernels ops (csrc/cpu/sgl-kernels/). Each test carries its own reference
5
+reimplementation of the kernel's math rather than depending on production
6
+model code.
7
+"""
8
+
9
+from unittest import mock
10
+
11
+import pytest
12
+import torch
13
+
14
+import vllm.model_executor.kernels.mhc as mhc_kernels
15
+from tests.kernels.test_fused_deepseek_v4_qnorm_rope_kv_insert import (
16
+ HEAD_DIM,
17
+ NOPE_DIM,
18
+ QUANT_BLOCK,
19
+ ROPE_DIM,
20
+ make_cos_sin_cache,
21
+)
22
+from tests.kernels.utils import bf16_ulp_distance, fp8_ulp_distance
23
+from vllm._custom_ops import (
24
+ compress_norm_rope_store_cpu,
25
+ compress_norm_rope_store_indexer_cpu,
26
+ flash_mla_with_kvcache_cpu,
27
+ fp8_paged_mqa_logits_cpu,
28
+ fused_indexer_q_rope_quant_cpu,
29
+ fused_qnorm_rope_kv_insert_cpu,
30
+ inverse_gptj_rope_o_proj_cpu,
31
+ save_partial_states_cpu,
32
+)
33
+from vllm.models.deepseek_v4.common.ops import fused_indexer_q_rope_quant
34
+from vllm.platforms import current_platform
35
+from vllm.utils.torch_utils import set_random_seed
36
+
37
+DEVICE = current_platform.device_type
38
+
39
+# Main attention cache (fp8_ds_mla): NoPE fp8 in QUANT_BLOCK-wide UE8M0 blocks,
40
+# RoPE bf16, one scale byte per block (+1 pad byte) per token.
41
+SCALE_BYTES_PER_TOKEN = NOPE_DIM // QUANT_BLOCK + 1
42
+TOKEN_DATA_BYTES = NOPE_DIM + ROPE_DIM * 2
43
+FP8_MAX = 448.0
44
+
45
+# Indexer K-cache: all 128 dims in one fp8 block, one raw fp32 scale per token.
46
+INDEXER_HEAD_DIM = 128
47
+INDEXER_NOPE_DIM = INDEXER_HEAD_DIM - ROPE_DIM
48
+INDEXER_SCALE_BYTES = 4
49
+
50
+
51
+def _op_available(*names: str) -> bool:
52
+ return all(hasattr(torch.ops._C, name) for name in names)
53
+
54
+
55
+pytestmark = pytest.mark.skipif(
56
+ not current_platform.is_cpu()
57
+ or not _op_available(
58
+ "save_partial_states_cpu",
59
+ "compress_norm_rope_store_cpu",
60
+ "compress_norm_rope_store_indexer_cpu",
61
+ "flash_mla_with_kvcache_cpu",
62
+ "fused_qnorm_rope_kv_insert_cpu",
63
+ "inverse_gptj_rope_o_proj_cpu",
64
+ "fp8_paged_mqa_logits_cpu",
65
+ "fused_indexer_q_rope_quant_cpu",
66
+ ),
67
+ reason="CPU not available or DeepSeek-V4 CPU kernels not built in",
68
+)
69
+
70
+
71
+def _apply_gptj_rope(x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor):
72
+ x_even = x[..., 0::2]
73
+ x_odd = x[..., 1::2]
74
+ out_even = x_even * cos - x_odd * sin
75
+ out_odd = x_even * sin + x_odd * cos
76
+ return torch.stack((out_even, out_odd), dim=-1).flatten(-2)
77
+
78
+
79
+def _make_cache(num_blocks: int, block_size: int) -> torch.Tensor:
80
+ block_bytes = block_size * TOKEN_DATA_BYTES + block_size * SCALE_BYTES_PER_TOKEN
81
+ return torch.zeros(num_blocks, block_bytes, dtype=torch.uint8)
82
+
83
+
84
+def _decode_token_regions(
85
+ cache: torch.Tensor, slot_mapping: torch.Tensor, block_size: int
86
+):
87
+ """Slice out, per valid token, the raw NoPE-fp8 / RoPE-bf16 / scale byte
88
+ regions this layout defines."""
89
+ valid = slot_mapping >= 0
90
+ idx = valid.nonzero(as_tuple=True)[0]
91
+ slots = slot_mapping[idx]
92
+ block_idx = torch.div(slots, block_size, rounding_mode="floor")
93
+ pos_in_block = slots % block_size
94
+
95
+ nope, rope, scale = [], [], []
96
+ for b, p in zip(block_idx.tolist(), pos_in_block.tolist()):
97
+ token_off = p * TOKEN_DATA_BYTES
98
+ nope.append(cache[b, token_off : token_off + NOPE_DIM])
99
+ rope.append(cache[b, token_off + NOPE_DIM : token_off + TOKEN_DATA_BYTES])
100
+ scale_off = block_size * TOKEN_DATA_BYTES + p * SCALE_BYTES_PER_TOKEN
101
+ scale.append(cache[b, scale_off : scale_off + SCALE_BYTES_PER_TOKEN - 1])
102
+
103
+ if not nope:
104
+ empty = lambda n: torch.empty(0, n, dtype=torch.uint8)
105
+ return empty(NOPE_DIM), empty(ROPE_DIM * 2), empty(SCALE_BYTES_PER_TOKEN - 1)
106
+ return torch.stack(nope), torch.stack(rope), torch.stack(scale)
107
+
108
+
109
+def _assert_cache_parity(cache_fused, cache_ref, slot_mapping, block_size):
110
+ """NoPE fp8 + the UE8M0 scale byte are a deterministic quant of an
111
+ un-rotated bf16 value, so they match tightly; the RoPE region is rotated
112
+ in fp32 before its bf16 store and eager/kernel can round a tie to
113
+ opposite sides, so allow <=1 ULP there."""
114
+ nope_ref, rope_ref, scale_ref = _decode_token_regions(
115
+ cache_ref, slot_mapping, block_size
116
+ )
117
+ nope_fused, rope_fused, scale_fused = _decode_token_regions(
118
+ cache_fused, slot_mapping, block_size
119
+ )
120
+ if nope_ref.numel() == 0:
121
+ return
122
+
123
+ max_ulp = int(fp8_ulp_distance(nope_fused, nope_ref).max().item())
124
+ assert max_ulp <= 1, f"NoPE fp8 differs by {max_ulp} ULP (>1)"
125
+
126
+ rope_ref_bf16 = rope_ref.contiguous().view(torch.bfloat16)
127
+ rope_fused_bf16 = rope_fused.contiguous().view(torch.bfloat16)
128
+ max_ulp_rope = int(bf16_ulp_distance(rope_fused_bf16, rope_ref_bf16).max().item())
129
+ assert max_ulp_rope <= 1, f"RoPE bf16 differs by {max_ulp_rope} ULP (>1)"
130
+
131
+ scale_diff = (scale_fused.int() - scale_ref.int()).abs().max().item()
132
+ assert scale_diff <= 1, f"scale byte differs by {scale_diff} (>1)"
133
+
134
+
135
+def _make_indexer_cache(num_blocks: int, block_size: int) -> torch.Tensor:
136
+ block_bytes = block_size * INDEXER_HEAD_DIM + block_size * INDEXER_SCALE_BYTES
137
+ return torch.zeros(num_blocks, block_bytes, dtype=torch.uint8)
138
+
139
+
140
+def _assert_indexer_cache_parity(cache_fused, cache_ref, slot_mapping, block_size):
141
+ valid = slot_mapping >= 0
142
+ idx = valid.nonzero(as_tuple=True)[0]
143
+ if idx.numel() == 0:
144
+ return
145
+ slots = slot_mapping[idx]
146
+ block_idx = torch.div(slots, block_size, rounding_mode="floor")
147
+ pos_in_block = slots % block_size
148
+
149
+ fp8_ref, scale_ref = [], []
150
+ fp8_fused, scale_fused = [], []
151
+ for b, p in zip(block_idx.tolist(), pos_in_block.tolist()):
152
+ token_off = p * INDEXER_HEAD_DIM
153
+ scale_off = block_size * INDEXER_HEAD_DIM + p * INDEXER_SCALE_BYTES
154
+ fp8_ref.append(cache_ref[b, token_off : token_off + INDEXER_HEAD_DIM])
155
+ fp8_fused.append(cache_fused[b, token_off : token_off + INDEXER_HEAD_DIM])
156
+ scale_ref.append(cache_ref[b, scale_off : scale_off + INDEXER_SCALE_BYTES])
157
+ scale_fused.append(cache_fused[b, scale_off : scale_off + INDEXER_SCALE_BYTES])
158
+
159
+ fp8_ref_t = torch.stack(fp8_ref).view(torch.float8_e4m3fn)
160
+ fp8_fused_t = torch.stack(fp8_fused).view(torch.float8_e4m3fn)
161
+ torch.testing.assert_close(fp8_fused_t.float(), fp8_ref_t.float())
162
+
163
+ scale_ref_t = torch.stack(scale_ref).view(torch.float32)
164
+ scale_fused_t = torch.stack(scale_fused).view(torch.float32)
165
+ torch.testing.assert_close(scale_fused_t, scale_ref_t)
166
+
167
+
168
+# ---------------------------------------------------------------------------
169
+# save_partial_states_cpu / compress_norm_rope_store_cpu (compressor.cpp)
170
+# ---------------------------------------------------------------------------
171
+#
172
+# The compressor's state cache is pooled with the main KV cache's physical
173
+# pages and page-aligned, so it is genuinely 3D with a block-to-block stride
174
+# that need not equal block_size * width -- built here via torch.as_strided
175
+# over an over-sized flat buffer rather than a plain contiguous tensor, so a
176
+# regression back to flattening via .view() would be caught.
177
+
178
+
179
+def _make_padded_block_cache(
180
+ num_blocks: int, block_size: int, width: int, pad_elems: int
181
+) -> torch.Tensor:
182
+ page_stride = block_size * width + pad_elems
183
+ flat = torch.randn(num_blocks, page_stride, dtype=torch.float32)
184
+ return torch.as_strided(
185
+ flat, (num_blocks, block_size, width), (page_stride, width, 1)
186
+ )
187
+
188
+
189
+def _dup_padded_block_cache(
190
+ num_blocks: int, block_size: int, width: int, pad_elems: int
191
+) -> tuple[torch.Tensor, torch.Tensor]:
192
+ """Two independently-stored padded caches with identical values, for ops
193
+ that mutate the cache in place."""
194
+ page_stride = block_size * width + pad_elems
195
+ template = torch.randn(num_blocks, page_stride, dtype=torch.float32)
196
+
197
+ def _view(flat: torch.Tensor) -> torch.Tensor:
198
+ return torch.as_strided(
199
+ flat, (num_blocks, block_size, width), (page_stride, width, 1)
200
+ )
201
+
202
+ return _view(template.clone()), _view(template.clone())
203
+
204
+
205
+def _run_save_partial_states_eager(
206
+ state_cache, kv, score, ape, positions, slot_mapping, compress_ratio
207
+):
208
+ """Reference: state_cache[block_idx, pos_in_block, :W] = kv;
209
+ state_cache[..., W:] = score + ape[position % compress_ratio]; skip
210
+ slot < 0. Mutates state_cache in place, matching the op's contract."""
211
+ state_width = kv.shape[-1]
212
+ block_size = state_cache.shape[1]
213
+ valid = slot_mapping >= 0
214
+ idx = valid.nonzero(as_tuple=True)[0]
215
+ slots = slot_mapping[idx]
216
+ block_idx = slots // block_size
217
+ pos_in_block = slots % block_size
218
+
219
+ state_cache[block_idx, pos_in_block, :state_width] = kv[idx]
220
+ ape_row = ape[positions[idx] % compress_ratio]
221
+ state_cache[block_idx, pos_in_block, state_width:] = score[idx] + ape_row
222
+
223
+
224
+@pytest.mark.parametrize("pad_elems", [0, 24])
225
+@pytest.mark.parametrize("num_tokens", [1, 17])
226
+def test_save_partial_states_cpu_matches_eager(num_tokens: int, pad_elems: int):
227
+ torch.manual_seed(0)
228
+ compress_ratio = 4
229
+ block_size = 4
230
+ state_width = 512
231
+ width = 2 * state_width
232
+ num_blocks = (num_tokens + 4 + block_size - 1) // block_size + 1
233
+ num_slots = num_blocks * block_size
234
+
235
+ kv = torch.randn(num_tokens, state_width, dtype=torch.float32)
236
+ score = torch.randn(num_tokens, state_width, dtype=torch.float32)
237
+ ape = torch.randn(compress_ratio, state_width, dtype=torch.float32)
238
+ positions = torch.randint(0, 4096, (num_tokens,), dtype=torch.int64)
239
+
240
+ ref_cache, fused_cache = _dup_padded_block_cache(
241
+ num_blocks, block_size, width, pad_elems
242
+ )
243
+
244
+ slot_mapping = torch.randperm(num_slots)[:num_tokens].to(torch.int64)
245
+ if num_tokens > 1:
246
+ slot_mapping[-1] = -1 # exercise the "skip invalid slot" branch
247
+
248
+ _run_save_partial_states_eager(
249
+ ref_cache, kv, score, ape, positions, slot_mapping, compress_ratio
250
+ )
251
+ save_partial_states_cpu(kv, score, ape, positions, fused_cache, slot_mapping)
252
+
253
+ torch.testing.assert_close(fused_cache, ref_cache)
254
+
255
+
256
+def _run_compress_eager(
257
+ state_cache,
258
+ gather_slots,
259
+ positions,
260
+ kv_slot_mapping,
261
+ rms_norm_weight,
262
+ rms_norm_eps,
263
+ cos_sin_cache,
264
+ kv_cache,
265
+ kv_cache_block_size,
266
+ compress_ratio,
267
+):
268
+ """Reference for the per-channel-softmax-pooling + RMSNorm + NoPE-UE8M0
269
+ quant + GPT-J RoPE math writing the main (head_dim=512) fp8_ds_mla
270
+ layout."""
271
+ num_tokens, window = gather_slots.shape
272
+ state_width = state_cache.shape[-1] // 2
273
+ state_block_size = state_cache.shape[1]
274
+ num_quant_blocks = NOPE_DIM // QUANT_BLOCK
275
+
276
+ w_idx = torch.arange(window)
277
+ head_offset = torch.where(w_idx >= compress_ratio, HEAD_DIM, 0)
278
+
279
+ for t in range(num_tokens):
280
+ position = int(positions[t])
281
+ if (position + 1) % compress_ratio != 0:
282
+ continue
283
+ kv_slot = int(kv_slot_mapping[t])
284
+ if kv_slot < 0:
285
+ continue
286
+
287
+ slots = gather_slots[t]
288
+ valid_w = slots >= 0
289
+ safe_slots = slots.clamp(min=0)
290
+ block_idx = safe_slots // state_block_size
291
+ pos_in_block = safe_slots % state_block_size
292
+
293
+ rows = state_cache[block_idx, pos_in_block] # [window, 2*state_width]
294
+ col = head_offset.unsqueeze(-1) + torch.arange(HEAD_DIM)
295
+ kv_vals = torch.gather(rows, 1, col)
296
+ score_col = state_width + col
297
+ score_vals = torch.gather(rows, 1, score_col)
298
+ score_vals = torch.where(
299
+ valid_w.unsqueeze(-1), score_vals, torch.full_like(score_vals, -torch.inf)
300
+ )
301
+
302
+ weights = torch.softmax(score_vals, dim=0)
303
+ compressed = (kv_vals * weights).sum(dim=0) # [HEAD_DIM]
304
+
305
+ variance = compressed.pow(2).mean()
306
+ normed = compressed * torch.rsqrt(variance + rms_norm_eps) * rms_norm_weight
307
+
308
+ quant_input = normed.to(torch.bfloat16).to(torch.float32)
309
+
310
+ compressed_pos = (position // compress_ratio) * compress_ratio
311
+ cos = cos_sin_cache[compressed_pos, : ROPE_DIM // 2]
312
+ sin = cos_sin_cache[compressed_pos, ROPE_DIM // 2 :]
313
+ rotated_rope = _apply_gptj_rope(normed[NOPE_DIM:].clone(), cos, sin)
314
+
315
+ nope_vals = quant_input[:NOPE_DIM].view(num_quant_blocks, QUANT_BLOCK)
316
+ absmax = nope_vals.abs().amax(dim=-1).clamp(min=1e-4)
317
+ exponent = torch.ceil(torch.log2(absmax / FP8_MAX))
318
+ inv_scale = torch.exp2(-exponent)
319
+ scaled = (nope_vals * inv_scale.unsqueeze(-1)).clamp(-FP8_MAX, FP8_MAX)
320
+ nope_fp8_bytes = scaled.to(torch.float8_e4m3fn).reshape(-1).view(torch.uint8)
321
+ scale_bytes = (exponent + 127.0).clamp(0, 255).to(torch.uint8)
322
+ rope_bf16_bytes = rotated_rope.to(torch.bfloat16).contiguous().view(torch.uint8)
323
+
324
+ block_idx_kv = kv_slot // kv_cache_block_size
325
+ pos_in_block_kv = kv_slot % kv_cache_block_size
326
+ token_off = pos_in_block_kv * TOKEN_DATA_BYTES
327
+ kv_cache[block_idx_kv, token_off : token_off + NOPE_DIM] = nope_fp8_bytes
328
+ kv_cache[block_idx_kv, token_off + NOPE_DIM : token_off + TOKEN_DATA_BYTES] = (
329
+ rope_bf16_bytes
330
+ )
331
+ scale_off = (
332
+ kv_cache_block_size * TOKEN_DATA_BYTES
333
+ + pos_in_block_kv * SCALE_BYTES_PER_TOKEN
334
+ )
335
+ kv_cache[block_idx_kv, scale_off : scale_off + num_quant_blocks] = scale_bytes
336
+ kv_cache[block_idx_kv, scale_off + num_quant_blocks] = 0
337
+
338
+
339
+def _make_compress_inputs(
340
+ num_tokens: int,
341
+ compress_ratio: int,
342
+ overlap: bool,
343
+ pad_elems: int = 24,
344
+ max_pos: int = 4096,
345
+ invalid_window_frac: float = 0.1,
346
+ head_dim: int = HEAD_DIM,
347
+):
348
+ torch.manual_seed(0)
349
+ coff = 1 + int(overlap)
350
+ state_width = coff * head_dim
351
+ window = coff * compress_ratio
352
+ state_block_size = 4 if compress_ratio == 4 else 8
353
+ num_state_slots = num_tokens * window + 8
354
+ num_state_blocks = (num_state_slots + state_block_size - 1) // state_block_size + 1
355
+
356
+ state_cache = _make_padded_block_cache(
357
+ num_state_blocks, state_block_size, 2 * state_width, pad_elems
358
+ )
359
+ gather_slots = torch.randint(
360
+ 0, num_state_blocks * state_block_size, (num_tokens, window)
361
+ ).to(torch.int64)
362
+ invalid_mask = torch.rand(num_tokens, window) < invalid_window_frac
363
+ # Never invalidate every column of a row (would force max=-inf -> NaN).
364
+ invalid_mask[:, -1] = False
365
+ gather_slots = torch.where(
366
+ invalid_mask, torch.full_like(gather_slots, -1), gather_slots
367
+ )
368
+
369
+ base = (torch.arange(num_tokens) + 1) * compress_ratio
370
+ is_boundary = torch.arange(num_tokens) % 2 == 0
371
+ positions = torch.where(is_boundary, base - 1, base + 1).to(torch.int64)
372
+
373
+ kv_slot_mapping = torch.arange(num_tokens, dtype=torch.int64)
374
+ if num_tokens > 1:
375
+ kv_slot_mapping[1] = -1 # exercise the "skip invalid output slot" branch
376
+
377
+ rms_norm_weight = torch.randn(head_dim, dtype=torch.float32)
378
+ cos_sin_cache = make_cos_sin_cache(max_pos, ROPE_DIM, torch.float32, "cpu")
379
+
380
+ return (
381
+ state_cache,
382
+ gather_slots,
383
+ positions,
384
+ kv_slot_mapping,
385
+ rms_norm_weight,
386
+ cos_sin_cache,
387
+ is_boundary,
388
+ )
389
+
390
+
391
+@pytest.mark.parametrize("compress_ratio,overlap", [(128, False), (4, True)])
392
+@pytest.mark.parametrize("num_tokens", [1, 17])
393
+def test_compress_norm_rope_store_cpu_matches_eager(
394
+ compress_ratio: int, overlap: bool, num_tokens: int
395
+):
396
+ eps = 1e-6
397
+ block_size = 16
398
+ (
399
+ state_cache,
400
+ gather_slots,
401
+ positions,
402
+ kv_slot_mapping,
403
+ rms_norm_weight,
404
+ cos_sin_cache,
405
+ is_boundary,
406
+ ) = _make_compress_inputs(num_tokens, compress_ratio, overlap)
407
+
408
+ num_blocks = (num_tokens + block_size - 1) // block_size + 1
409
+ kv_cache_ref = _make_cache(num_blocks, block_size)
410
+ kv_cache_fused = kv_cache_ref.clone()
411
+
412
+ _run_compress_eager(
413
+ state_cache,
414
+ gather_slots,
415
+ positions,
416
+ kv_slot_mapping,
417
+ rms_norm_weight,
418
+ eps,
419
+ cos_sin_cache,
420
+ kv_cache_ref,
421
+ block_size,
422
+ compress_ratio,
423
+ )
424
+ compress_norm_rope_store_cpu(
425
+ state_cache,
426
+ gather_slots,
427
+ positions,
428
+ kv_slot_mapping,
429
+ rms_norm_weight,
430
+ eps,
431
+ cos_sin_cache,
432
+ kv_cache_fused,
433
+ block_size,
434
+ compress_ratio,
435
+ )
436
+
437
+ # Only tokens that hit the compression boundary (and have a valid output
438
+ # slot) were written -- restrict parity checks to those.
439
+ written_slot_mapping = torch.where(
440
+ is_boundary, kv_slot_mapping, torch.full_like(kv_slot_mapping, -1)
441
+ )
442
+ _assert_cache_parity(kv_cache_fused, kv_cache_ref, written_slot_mapping, block_size)
443
+
444
+
445
+def test_compress_norm_rope_store_cpu_non_contiguous_kv_cache_row_stride():
446
+ """The main (output) kv_cache is a separate physical buffer from the
447
+ state cache and can independently be a strided view."""
448
+ eps = 1e-6
449
+ compress_ratio, overlap = 128, False
450
+ num_tokens = 5
451
+ block_size = 16
452
+ cache_pad = 32
453
+
454
+ (
455
+ state_cache,
456
+ gather_slots,
457
+ positions,
458
+ kv_slot_mapping,
459
+ rms_norm_weight,
460
+ cos_sin_cache,
461
+ is_boundary,
462
+ ) = _make_compress_inputs(num_tokens, compress_ratio, overlap)
463
+
464
+ num_blocks = (num_tokens + block_size - 1) // block_size + 1
465
+ kv_cache_ref = _make_cache(num_blocks, block_size)
466
+ block_bytes = kv_cache_ref.shape[1]
467
+ padded_cache_ref = torch.zeros(
468
+ num_blocks, block_bytes + cache_pad, dtype=torch.uint8
469
+ )
470
+ kv_cache_ref_view = padded_cache_ref[:, :block_bytes]
471
+ kv_cache_ref_view.copy_(kv_cache_ref)
472
+ padded_cache_fused = padded_cache_ref.clone()
473
+ kv_cache_fused_view = padded_cache_fused[:, :block_bytes]
474
+
475
+ _run_compress_eager(
476
+ state_cache,
477
+ gather_slots,
478
+ positions,
479
+ kv_slot_mapping,
480
+ rms_norm_weight,
481
+ eps,
482
+ cos_sin_cache,
483
+ kv_cache_ref_view,
484
+ block_size,
485
+ compress_ratio,
486
+ )
487
+ compress_norm_rope_store_cpu(
488
+ state_cache,
489
+ gather_slots,
490
+ positions,
491
+ kv_slot_mapping,
492
+ rms_norm_weight,
493
+ eps,
494
+ cos_sin_cache,
495
+ kv_cache_fused_view,
496
+ block_size,
497
+ compress_ratio,
498
+ )
499
+
500
+ written_slot_mapping = torch.where(
501
+ is_boundary, kv_slot_mapping, torch.full_like(kv_slot_mapping, -1)
502
+ )
503
+ _assert_cache_parity(
504
+ kv_cache_fused_view, kv_cache_ref_view, written_slot_mapping, block_size
505
+ )
506
+
507
+
508
+def _run_compress_indexer_eager(
509
+ state_cache,
510
+ gather_slots,
511
+ positions,
512
+ kv_slot_mapping,
513
+ rms_norm_weight,
514
+ rms_norm_eps,
515
+ cos_sin_cache,
516
+ kv_cache,
517
+ kv_cache_block_size,
518
+ compress_ratio,
519
+):
520
+ """Same softmax-pool + RMSNorm + GPT-J RoPE prologue as the head=512
521
+ path, but the indexer's cache quantizes all 128 post-RoPE dims as one
522
+ FP8 block with a single raw fp32 scale per token (matching
523
+ _indexer_k_quant_and_cache_cpu's layout), not 7 separate UE8M0 blocks."""
524
+ num_tokens, window = gather_slots.shape
525
+ state_width = state_cache.shape[-1] // 2
526
+ state_block_size = state_cache.shape[1]
527
+
528
+ w_idx = torch.arange(window)
529
+ head_offset = torch.where(w_idx >= compress_ratio, INDEXER_HEAD_DIM, 0)
530
+
531
+ for t in range(num_tokens):
532
+ position = int(positions[t])
533
+ if (position + 1) % compress_ratio != 0:
534
+ continue
535
+ kv_slot = int(kv_slot_mapping[t])
536
+ if kv_slot < 0:
537
+ continue
538
+
539
+ slots = gather_slots[t]
540
+ valid_w = slots >= 0
541
+ safe_slots = slots.clamp(min=0)
542
+ block_idx = safe_slots // state_block_size
543
+ pos_in_block = safe_slots % state_block_size
544
+
545
+ rows = state_cache[block_idx, pos_in_block] # [window, 2*state_width]
546
+ col = head_offset.unsqueeze(-1) + torch.arange(INDEXER_HEAD_DIM)
547
+ kv_vals = torch.gather(rows, 1, col)
548
+ score_col = state_width + col
549
+ score_vals = torch.gather(rows, 1, score_col)
550
+ score_vals = torch.where(
551
+ valid_w.unsqueeze(-1), score_vals, torch.full_like(score_vals, -torch.inf)
552
+ )
553
+
554
+ weights = torch.softmax(score_vals, dim=0)
555
+ compressed = (kv_vals * weights).sum(dim=0) # [INDEXER_HEAD_DIM]
556
+
557
+ variance = compressed.pow(2).mean()
558
+ normed = compressed * torch.rsqrt(variance + rms_norm_eps) * rms_norm_weight
559
+
560
+ compressed_pos = (position // compress_ratio) * compress_ratio
561
+ cos = cos_sin_cache[compressed_pos, : ROPE_DIM // 2]
562
+ sin = cos_sin_cache[compressed_pos, ROPE_DIM // 2 :]
563
+ rotated_rope = _apply_gptj_rope(normed[INDEXER_NOPE_DIM:].clone(), cos, sin)
564
+ full = torch.cat((normed[:INDEXER_NOPE_DIM], rotated_rope))
565
+
566
+ quant_input = full.to(torch.bfloat16).to(torch.float32)
567
+ absmax = quant_input.abs().amax().clamp(min=1e-4)
568
+ exponent = torch.ceil(torch.log2(absmax / FP8_MAX))
569
+ scale = torch.exp2(exponent)
570
+ scaled = (quant_input / scale).clamp(-FP8_MAX, FP8_MAX)
571
+ fp8_bytes = scaled.to(torch.float8_e4m3fn).view(torch.uint8)
572
+
573
+ block_idx_kv = kv_slot // kv_cache_block_size
574
+ pos_in_block_kv = kv_slot % kv_cache_block_size
575
+ token_off = pos_in_block_kv * INDEXER_HEAD_DIM
576
+ kv_cache[block_idx_kv, token_off : token_off + INDEXER_HEAD_DIM] = fp8_bytes
577
+ scale_off = (
578
+ kv_cache_block_size * INDEXER_HEAD_DIM
579
+ + pos_in_block_kv * INDEXER_SCALE_BYTES
580
+ )
581
+ kv_cache[block_idx_kv, scale_off : scale_off + INDEXER_SCALE_BYTES] = (
582
+ scale.reshape(1).view(torch.uint8)
583
+ )
584
+
585
+
586
+@pytest.mark.parametrize("compress_ratio,overlap", [(128, False), (4, True)])
587
+@pytest.mark.parametrize("num_tokens", [1, 17])
588
+def test_compress_norm_rope_store_indexer_cpu_matches_eager(
589
+ compress_ratio: int, overlap: bool, num_tokens: int
590
+):
591
+ eps = 1e-6
592
+ block_size = 16
593
+ (
594
+ state_cache,
595
+ gather_slots,
596
+ positions,
597
+ kv_slot_mapping,
598
+ rms_norm_weight,
599
+ cos_sin_cache,
600
+ is_boundary,
601
+ ) = _make_compress_inputs(
602
+ num_tokens, compress_ratio, overlap, head_dim=INDEXER_HEAD_DIM
603
+ )
604
+
605
+ num_blocks = (num_tokens + block_size - 1) // block_size + 1
606
+ kv_cache_ref = _make_indexer_cache(num_blocks, block_size)
607
+ kv_cache_fused = kv_cache_ref.clone()
608
+
609
+ _run_compress_indexer_eager(
610
+ state_cache,
611
+ gather_slots,
612
+ positions,
613
+ kv_slot_mapping,
614
+ rms_norm_weight,
615
+ eps,
616
+ cos_sin_cache,
617
+ kv_cache_ref,
618
+ block_size,
619
+ compress_ratio,
620
+ )
621
+ compress_norm_rope_store_indexer_cpu(
622
+ state_cache,
623
+ gather_slots,
624
+ positions,
625
+ kv_slot_mapping,
626
+ rms_norm_weight,
627
+ eps,
628
+ cos_sin_cache,
629
+ kv_cache_fused,
630
+ block_size,
631
+ compress_ratio,
632
+ )
633
+
634
+ written_slot_mapping = torch.where(
635
+ is_boundary, kv_slot_mapping, torch.full_like(kv_slot_mapping, -1)
636
+ )
637
+ _assert_indexer_cache_parity(
638
+ kv_cache_fused, kv_cache_ref, written_slot_mapping, block_size
639
+ )
640
+
641
+
642
+# ---------------------------------------------------------------------------
643
+# flash_mla_with_kvcache_cpu (flash_mla.cpp): fused sparse MQA decode attn
644
+# ---------------------------------------------------------------------------
645
+
646
+
647
+def _pack_flash_mla_cache(kv_f32: torch.Tensor, block_size: int):
648
+ """Quantize kv_f32 ([N, HEAD_DIM]) into the fp8_ds_mla byte layout and
649
+ return (cache, dequant_ref) -- the exact post-quantization value the
650
+ kernel's dequant should reproduce. Slot id == row index, packed
651
+ contiguously from block 0."""
652
+ n = kv_f32.shape[0]
653
+ num_blocks = max(1, -(-n // block_size))
654
+ cache = _make_cache(num_blocks, block_size)
655
+ if n == 0:
656
+ return cache, kv_f32.new_zeros(0, HEAD_DIM)
657
+
658
+ num_quant_blocks = NOPE_DIM // QUANT_BLOCK
659
+ nope = kv_f32[:, :NOPE_DIM].reshape(n, num_quant_blocks, QUANT_BLOCK)
660
+ absmax = nope.abs().amax(dim=-1).clamp(min=1e-4)
661
+ exponent = torch.ceil(torch.log2(absmax / FP8_MAX))
662
+ inv_scale = torch.exp2(-exponent)
663
+ scaled = (nope * inv_scale.unsqueeze(-1)).clamp(-FP8_MAX, FP8_MAX)
664
+ nope_fp8 = scaled.to(torch.float8_e4m3fn)
665
+ scale = torch.exp2(exponent)
666
+ nope_dequant = (nope_fp8.float() * scale.unsqueeze(-1)).reshape(n, NOPE_DIM)
667
+
668
+ rope_bf16 = kv_f32[:, NOPE_DIM:].to(torch.bfloat16)
669
+ dequant_ref = torch.cat([nope_dequant, rope_bf16.float()], dim=-1)
670
+
671
+ slot_ids = torch.arange(n)
672
+ block_idx = slot_ids // block_size
673
+ pos_in_block = slot_ids % block_size
674
+
675
+ nope_fp8_bytes = nope_fp8.reshape(n, NOPE_DIM).contiguous().view(torch.uint8)
676
+ rope_bf16_bytes = rope_bf16.contiguous().view(torch.uint8)
677
+ scale_bytes = (exponent + 127.0).clamp(0, 255).to(torch.uint8)
678
+
679
+ arange_nope = torch.arange(NOPE_DIM)
680
+ arange_rope = torch.arange(ROPE_DIM * 2)
681
+ arange_scale = torch.arange(num_quant_blocks)
682
+ token_base = pos_in_block.unsqueeze(-1) * TOKEN_DATA_BYTES
683
+ row_idx = block_idx.unsqueeze(-1)
684
+
685
+ cache[row_idx, token_base + arange_nope] = nope_fp8_bytes
686
+ cache[row_idx, token_base + NOPE_DIM + arange_rope] = rope_bf16_bytes
687
+ scale_base = (
688
+ block_size * TOKEN_DATA_BYTES
689
+ + pos_in_block.unsqueeze(-1) * SCALE_BYTES_PER_TOKEN
690
+ )
691
+ cache[row_idx, scale_base + arange_scale] = scale_bytes
692
+ cache[block_idx, scale_base.squeeze(-1) + num_quant_blocks] = 0
693
+
694
+ return cache, dequant_ref
695
+
696
+
697
+def _flash_mla_eager(
698
+ q: torch.Tensor,
699
+ window_dequant: torch.Tensor,
700
+ window_slots: torch.Tensor,
701
+ compressed_dequant: torch.Tensor,
702
+ compressed_slots: torch.Tensor,
703
+ attn_sink: torch.Tensor,
704
+ scale: float,
705
+) -> torch.Tensor:
706
+ """MQA attention: gather the (already-dequantized) window+compressed KV
707
+ rows referenced by *_slots (-1 = invalid/masked), append a learned
708
+ per-head sink logit with zero value contribution, softmax, weighted
709
+ sum."""
710
+ num_tokens, num_heads, _ = q.shape
711
+ q_f32 = q.float()
712
+
713
+ def gather(dequant, slots):
714
+ valid = slots >= 0
715
+ safe = slots.clamp(min=0)
716
+ kv = dequant[safe] # [T, K, D]
717
+ return kv, valid
718
+
719
+ win_kv, win_valid = gather(window_dequant, window_slots)
720
+ if compressed_slots.shape[1] > 0:
721
+ comp_kv, comp_valid = gather(compressed_dequant, compressed_slots)
722
+ kv = torch.cat([win_kv, comp_kv], dim=1)
723
+ valid = torch.cat([win_valid, comp_valid], dim=1)
724
+ else:
725
+ kv = win_kv
726
+ valid = win_valid
727
+
728
+ scores = torch.einsum("thd,tkd->thk", q_f32, kv) * scale
729
+ scores = scores.masked_fill(~valid.unsqueeze(1), float("-inf"))
730
+ sink = attn_sink.view(1, num_heads, 1).expand(num_tokens, num_heads, 1)
731
+ scores_full = torch.cat([scores, sink], dim=-1)
732
+ probs = torch.softmax(scores_full, dim=-1)
733
+ out = torch.einsum("thk,tkd->thd", probs[..., :-1], kv)
734
+ return out.to(q.dtype)
735
+
736
+
737
+@pytest.mark.parametrize("num_heads", [1, 64])
738
+@pytest.mark.parametrize("num_tokens", [1, 17])
739
+@pytest.mark.parametrize("swa_only", [True, False])
740
+def test_flash_mla_cpu_kernel_matches_eager(
741
+ num_heads: int, num_tokens: int, swa_only: bool
742
+):
743
+ torch.manual_seed(0)
744
+ dtype = torch.bfloat16
745
+ window_block_size = 16
746
+ compressed_block_size = 8
747
+ num_window = 6
748
+ num_compressed = 0 if swa_only else 5
749
+
750
+ q = torch.randn(num_tokens, num_heads, HEAD_DIM, dtype=dtype)
751
+ scale = 1.0 / (HEAD_DIM**0.5)
752
+ attn_sink = torch.randn(num_heads, dtype=torch.float32)
753
+
754
+ num_window_slots = num_tokens * num_window + window_block_size
755
+ window_kv_f32 = torch.randn(num_window_slots, HEAD_DIM)
756
+ window_cache, window_dequant = _pack_flash_mla_cache(
757
+ window_kv_f32, window_block_size
758
+ )
759
+ window_slots = torch.randint(0, num_window_slots, (num_tokens, num_window))
760
+ window_slots[:, 0] = -1 # exercise invalid-slot masking
761
+
762
+ if num_compressed > 0:
763
+ num_compressed_slots = num_tokens * num_compressed + compressed_block_size
764
+ compressed_kv_f32 = torch.randn(num_compressed_slots, HEAD_DIM)
765
+ compressed_cache, compressed_dequant = _pack_flash_mla_cache(
766
+ compressed_kv_f32, compressed_block_size
767
+ )
768
+ compressed_slots = torch.randint(
769
+ 0, num_compressed_slots, (num_tokens, num_compressed)
770
+ )
771
+ if num_compressed > 1:
772
+ compressed_slots[:, -1] = -1
773
+ else:
774
+ compressed_cache = window_cache
775
+ compressed_dequant = window_dequant.new_zeros(0, HEAD_DIM)
776
+ compressed_slots = window_slots.new_full((num_tokens, 0), -1)
777
+
778
+ out_ref = _flash_mla_eager(
779
+ q,
780
+ window_dequant,
781
+ window_slots,
782
+ compressed_dequant,
783
+ compressed_slots,
784
+ attn_sink,
785
+ scale,
786
+ )
787
+ out_kernel = q.new_empty(num_tokens, num_heads, HEAD_DIM)
788
+ flash_mla_with_kvcache_cpu(
789
+ out_kernel,
790
+ q,
791
+ window_cache,
792
+ window_slots,
793
+ window_block_size,
794
+ compressed_cache,
795
+ compressed_slots,
796
+ compressed_block_size,
797
+ attn_sink,
798
+ scale,
799
+ )
800
+
801
+ torch.testing.assert_close(
802
+ out_kernel.float(), out_ref.float(), rtol=2e-2, atol=2e-2
803
+ )
804
+
805
+
806
+def test_flash_mla_cpu_kernel_all_invalid_falls_back_to_sink():
807
+ """When every KV slot is masked out, softmax degenerates to the sink
808
+ logit alone and the output must be all-zero (no valid value rows)."""
809
+ torch.manual_seed(2)
810
+ dtype = torch.bfloat16
811
+ num_tokens, num_heads, num_window = 3, 4, 5
812
+ window_block_size = 16
813
+
814
+ q = torch.randn(num_tokens, num_heads, HEAD_DIM, dtype=dtype)
815
+ scale = 1.0 / (HEAD_DIM**0.5)
816
+ attn_sink = torch.randn(num_heads, dtype=torch.float32)
817
+
818
+ window_kv_f32 = torch.randn(window_block_size, HEAD_DIM)
819
+ window_cache, _ = _pack_flash_mla_cache(window_kv_f32, window_block_size)
820
+ window_slots = torch.full((num_tokens, num_window), -1, dtype=torch.int64)
821
+ compressed_slots = window_slots.new_full((num_tokens, 0), -1)
822
+
823
+ out = q.new_empty(num_tokens, num_heads, HEAD_DIM)
824
+ flash_mla_with_kvcache_cpu(
825
+ out,
826
+ q,
827
+ window_cache,
828
+ window_slots,
829
+ window_block_size,
830
+ window_cache,
831
+ compressed_slots,
832
+ 1,
833
+ attn_sink,
834
+ scale,
835
+ )
836
+ assert out.abs().max().item() == 0.0
837
+
838
+
839
+# ---------------------------------------------------------------------------
840
+# fused_qnorm_rope_kv_insert_cpu (store_cache.cpp)
841
+# ---------------------------------------------------------------------------
842
+
843
+
844
+def _qnorm_kv_insert_eager(
845
+ q, kv, positions, cache, slot_mapping, cos_sin_cache, q_head_padded, eps, bs
846
+):
847
+ """Reference: per-head weight-free RMSNorm + GPT-J RoPE on Q; GPT-J RoPE
848
+ + UE8M0 FP8 quant + paged insert on KV."""
849
+ num_quant_blocks = NOPE_DIM // QUANT_BLOCK
850
+
851
+ num_tokens_full, num_heads_q, _ = q.shape
852
+ out_dtype = q.dtype
853
+
854
+ cos = cos_sin_cache[positions, : ROPE_DIM // 2]
855
+ sin = cos_sin_cache[positions, ROPE_DIM // 2 :]
856
+
857
+ q_f32 = q.float()
858
+ var = q_f32.square().mean(dim=-1, keepdim=True)
859
+ q_f32 = q_f32 * torch.rsqrt(var + eps)
860
+ q_f32[..., NOPE_DIM:] = _apply_gptj_rope(
861
+ q_f32[..., NOPE_DIM:], cos.unsqueeze(1), sin.unsqueeze(1)
862
+ )
863
+ q_out = q.new_zeros(num_tokens_full, q_head_padded, HEAD_DIM)
864
+ q_out[:, :num_heads_q, :] = q_f32.to(out_dtype)
865
+
866
+ num_tokens_insert = slot_mapping.shape[0]
867
+ kv_f32 = kv[:num_tokens_insert].float()
868
+ kv_f32[..., NOPE_DIM:] = _apply_gptj_rope(
869
+ kv_f32[..., NOPE_DIM:], cos[:num_tokens_insert], sin[:num_tokens_insert]
870
+ )
871
+ kv_bf16 = kv_f32.to(torch.bfloat16)
872
+ kv_f32 = kv_bf16.float()
873
+
874
+ valid = slot_mapping >= 0
875
+ if bool(valid.any()):
876
+ idx = valid.nonzero(as_tuple=True)[0]
877
+ slot_id = slot_mapping[idx]
878
+ block_idx = torch.div(slot_id, bs, rounding_mode="floor")
879
+ pos_in_block = slot_id % bs
880
+
881
+ nope = kv_f32[idx, :NOPE_DIM].view(-1, num_quant_blocks, QUANT_BLOCK)
882
+ absmax = nope.abs().amax(dim=-1).clamp(min=1e-4)
883
+ exponent = torch.ceil(torch.log2(absmax / FP8_MAX))
884
+ inv_scale = torch.exp2(-exponent)
885
+ scaled = (nope * inv_scale.unsqueeze(-1)).clamp(-FP8_MAX, FP8_MAX)
886
+ # Flush fp8-subnormal magnitudes to zero, matching the kernel's
887
+ # AVX512 cvt_fp32x16_to_fp8x16, which doesn't preserve exact
888
+ # subnormal encodings the way .to(torch.float8_e4m3fn) does.
889
+ fp8_min_normal = 2.0**-6
890
+ scaled = torch.where(scaled.abs() < fp8_min_normal, 0.0, scaled)
891
+ nope_fp8_bytes = (
892
+ scaled.to(torch.float8_e4m3fn).view(-1, NOPE_DIM).view(torch.uint8)
893
+ )
894
+ scale_bytes = (exponent + 127.0).clamp(0, 255).to(torch.uint8)
895
+ rope_bf16_bytes = kv_bf16[idx, NOPE_DIM:].contiguous().view(torch.uint8)
896
+
897
+ arange_nope = torch.arange(NOPE_DIM)
898
+ arange_rope = torch.arange(ROPE_DIM * 2)
899
+ arange_scale = torch.arange(num_quant_blocks)
900
+ row_idx = block_idx.unsqueeze(-1)
901
+ token_base = pos_in_block.unsqueeze(-1) * TOKEN_DATA_BYTES
902
+
903
+ cache[row_idx, token_base + arange_nope] = nope_fp8_bytes
904
+ cache[row_idx, token_base + NOPE_DIM + arange_rope] = rope_bf16_bytes
905
+ scale_base = (
906
+ bs * TOKEN_DATA_BYTES + pos_in_block.unsqueeze(-1) * SCALE_BYTES_PER_TOKEN
907
+ )
908
+ cache[row_idx, scale_base + arange_scale] = scale_bytes
909
+ cache[block_idx, scale_base.squeeze(-1) + num_quant_blocks] = 0
910
+
911
+ return q_out
912
+
913
+
914
+@pytest.mark.parametrize("num_tokens", [1, 64])
915
+@pytest.mark.parametrize("n_heads,padded_heads", [(1, 8), (8, 8), (8, 16), (64, 128)])
916
+def test_qnorm_rope_kv_insert_cpu_kernel_matches_eager(
917
+ num_tokens: int, n_heads: int, padded_heads: int
918
+):
919
+ torch.manual_seed(0)
920
+ dtype = torch.bfloat16
921
+ eps = 1e-6
922
+ max_pos = 4096
923
+ block_size = 16
924
+
925
+ q = torch.randn(num_tokens, n_heads, HEAD_DIM, dtype=dtype)
926
+ kv = torch.randn(num_tokens, HEAD_DIM, dtype=dtype)
927
+ positions = torch.randint(0, max_pos, (num_tokens,), dtype=torch.int64)
928
+ cos_sin_cache = make_cos_sin_cache(max_pos, ROPE_DIM, torch.float32, "cpu")
929
+
930
+ num_blocks = (num_tokens + block_size - 1) // block_size + 1
931
+ slot_mapping = torch.arange(num_tokens, dtype=torch.int64)
932
+ if num_tokens > 1:
933
+ slot_mapping[-1] = -1 # exercise the "skip invalid slot" branch
934
+
935
+ cache_ref = _make_cache(num_blocks, block_size)
936
+ cache_fused = cache_ref.clone()
937
+
938
+ q_ref = _qnorm_kv_insert_eager(
939
+ q,
940
+ kv,
941
+ positions,
942
+ cache_ref,
943
+ slot_mapping,
944
+ cos_sin_cache,
945
+ padded_heads,
946
+ eps,
947
+ block_size,
948
+ )
949
+ q_fused = fused_qnorm_rope_kv_insert_cpu(
950
+ q,
951
+ kv,
952
+ positions,
953
+ cache_fused,
954
+ slot_mapping,
955
+ cos_sin_cache,
956
+ padded_heads,
957
+ eps,
958
+ block_size,
959
+ )
960
+
961
+ torch.testing.assert_close(q_fused.float(), q_ref.float(), rtol=1e-2, atol=1e-2)
962
+ if n_heads < padded_heads:
963
+ pad_region = q_fused[:, n_heads:padded_heads]
964
+ assert pad_region.abs().max().item() == 0.0, (
965
+ "padded head slots must be exact zero"
966
+ )
967
+
968
+ _assert_cache_parity(cache_fused, cache_ref, slot_mapping, block_size)
969
+
970
+
971
+@pytest.mark.parametrize("pad", [1, 5])
972
+def test_qnorm_rope_kv_insert_cpu_kernel_dp_padding(pad: int):
973
+ """slot_mapping shorter than q/kv rows: the KV branch must only touch
974
+ the first num_tokens rows while Q-norm+RoPE still runs on all rows."""
975
+ torch.manual_seed(1)
976
+ dtype = torch.bfloat16
977
+ eps = 1e-6
978
+ max_pos = 4096
979
+ block_size = 16
980
+ num_tokens = 17
981
+ total = num_tokens + pad
982
+ padded_heads = 8
983
+
984
+ q = torch.randn(total, padded_heads, HEAD_DIM, dtype=dtype)
985
+ kv = torch.randn(total, HEAD_DIM, dtype=dtype)
986
+ positions = torch.randint(0, max_pos, (total,), dtype=torch.int64)
987
+ cos_sin_cache = make_cos_sin_cache(max_pos, ROPE_DIM, torch.float32, "cpu")
988
+
989
+ num_blocks = (num_tokens + block_size - 1) // block_size + 1
990
+ slot_mapping = torch.arange(num_tokens, dtype=torch.int64)
991
+
992
+ cache_ref = _make_cache(num_blocks, block_size)
993
+ cache_fused = cache_ref.clone()
994
+
995
+ q_ref = _qnorm_kv_insert_eager(
996
+ q,
997
+ kv,
998
+ positions,
999
+ cache_ref,
1000
+ slot_mapping,
1001
+ cos_sin_cache,
1002
+ padded_heads,
1003
+ eps,
1004
+ block_size,
1005
+ )
1006
+ q_fused = fused_qnorm_rope_kv_insert_cpu(
1007
+ q,
1008
+ kv,
1009
+ positions,
1010
+ cache_fused,
1011
+ slot_mapping,
1012
+ cos_sin_cache,
1013
+ padded_heads,
1014
+ eps,
1015
+ block_size,
1016
+ )
1017
+
1018
+ torch.testing.assert_close(q_fused.float(), q_ref.float(), rtol=1e-2, atol=1e-2)
1019
+ _assert_cache_parity(cache_fused, cache_ref, slot_mapping, block_size)
1020
+
1021
+
1022
+# ---------------------------------------------------------------------------
1023
+# inverse_gptj_rope_o_proj_cpu (store_cache.cpp): DeepseekV4CPUAttention's
1024
+# _o_proj de-rotation kernel
1025
+# ---------------------------------------------------------------------------
1026
+
1027
+
1028
+def _inverse_gptj_rope_o_proj_eager(
1029
+ o: torch.Tensor,
1030
+ positions: torch.Tensor,
1031
+ cos_sin_cache: torch.Tensor,
1032
+ rope_dim: int,
1033
+) -> torch.Tensor:
1034
+ """NoPE dims pass through as fp32; the RoPE segment is de-rotated by
1035
+ R(-pos) in GPT-J (interleaved even/odd) convention."""
1036
+ nope_dim = o.shape[-1] - rope_dim
1037
+ half_dim = rope_dim // 2
1038
+
1039
+ o_f32 = o.float()
1040
+ out = torch.empty_like(o_f32)
1041
+ out[..., :nope_dim] = o_f32[..., :nope_dim]
1042
+
1043
+ cos = cos_sin_cache[positions, :half_dim].unsqueeze(1)
1044
+ sin = cos_sin_cache[positions, half_dim:].unsqueeze(1)
1045
+
1046
+ e = o_f32[..., nope_dim::2]
1047
+ od = o_f32[..., nope_dim + 1 :: 2]
1048
+ out[..., nope_dim::2] = e * cos + od * sin
1049
+ out[..., nope_dim + 1 :: 2] = od * cos - e * sin
1050
+ return out
1051
+
1052
+
1053
+@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16])
1054
+@pytest.mark.parametrize("num_tokens", [1, 17])
1055
+@pytest.mark.parametrize(
1056
+ "head_dim,rope_dim,num_heads",
1057
+ [
1058
+ (512, 64, 8), # DeepseekV4CPUAttention._o_proj's actual shape
1059
+ (64, 64, 1), # nope_dim=0: no NoPE region at all
1060
+ ],
1061
+)
1062
+def test_inverse_gptj_rope_o_proj_cpu_kernel_matches_eager(
1063
+ dtype: torch.dtype,
1064
+ num_tokens: int,
1065
+ head_dim: int,
1066
+ rope_dim: int,
1067
+ num_heads: int,
1068
+):
1069
+ torch.manual_seed(0)
1070
+ max_pos = 4096
1071
+
1072
+ o = torch.randn(num_tokens, num_heads, head_dim, dtype=dtype)
1073
+ positions = torch.randint(0, max_pos, (num_tokens,), dtype=torch.int64)
1074
+ cos_sin_cache = make_cos_sin_cache(max_pos, rope_dim, torch.float32, "cpu")
1075
+
1076
+ out_ref = _inverse_gptj_rope_o_proj_eager(o, positions, cos_sin_cache, rope_dim)
1077
+ out_fused = inverse_gptj_rope_o_proj_cpu(o, positions, cos_sin_cache, rope_dim)
1078
+
1079
+ torch.testing.assert_close(out_fused, out_ref, rtol=1e-2, atol=1e-2)
1080
+
1081
+
1082
+# ---------------------------------------------------------------------------
1083
+# fp8_paged_mqa_logits_cpu (paged_mqa_logits.cpp): sparse indexer's
1084
+# decode-path MQA-logits kernel
1085
+# ---------------------------------------------------------------------------
1086
+
1087
+
1088
+def _rand_fp8(*shape: int) -> torch.Tensor:
1089
+ return torch.randn(shape, dtype=torch.float32).clamp(-3, 3).to(torch.float8_e4m3fn)
1090
+
1091
+
1092
+def _pages_per_request(block_size: int, seq_lens: list[int]) -> list[int]:
1093
+ return [-(-seq_len // block_size) for seq_len in seq_lens]
1094
+
1095
+
1096
+def _write_page(
1097
+ page_row: torch.Tensor, block_size: int, k_fp8: torch.Tensor, scale: torch.Tensor
1098
+) -> None:
1099
+ """Write k_fp8 ([block_size, INDEXER_HEAD_DIM]) and scale ([block_size]
1100
+ fp32) into one page row: K-region then scale-region, matching
1101
+ compress_norm_rope_store_indexer_cpu's layout."""
1102
+ k_region = page_row[: block_size * INDEXER_HEAD_DIM]
1103
+ k_region.copy_(k_fp8.view(torch.uint8).flatten())
1104
+ scale_region = page_row[block_size * INDEXER_HEAD_DIM :].view(torch.float32)
1105
+ scale_region[:].copy_(scale)
1106
+
1107
+
1108
+def _reference_paged_mqa_logits(
1109
+ q_fp8: torch.Tensor,
1110
+ weight: torch.Tensor,
1111
+ ref_k: list[torch.Tensor],
1112
+ ref_scale: list[torch.Tensor],
1113
+ max_seq_len: int,
1114
+) -> torch.Tensor:
1115
+ batch_size, _num_heads, _ = q_fp8.shape
1116
+ out = torch.zeros((batch_size, max_seq_len), dtype=torch.float32)
1117
+ for b in range(batch_size):
1118
+ k = ref_k[b].float() # [seq_len, head_dim]
1119
+ scale = ref_scale[b] # [seq_len]
1120
+ q = q_fp8[b].float() # [heads, head_dim]
1121
+ w = weight[b] # [heads]
1122
+ dot = torch.einsum("jd,hd->jh", k, q).clamp(min=0.0)
1123
+ score = (dot * w.unsqueeze(0)).sum(dim=1)
1124
+ out[b, : k.shape[0]] = score * scale
1125
+ return out
1126
+
1127
+
1128
+def _build_paged_mqa_logits_batch(
1129
+ block_size: int,
1130
+ seq_lens: list[int],
1131
+ num_heads: int,
1132
+ pages: torch.Tensor,
1133
+ seed: int,
1134
+):
1135
+ """Fill pages (a pre-sized [n_pages, buf_width] uint8 tensor, possibly
1136
+ non-contiguous in dim 0) with deterministic-per-seed random K/scale
1137
+ data."""
1138
+ torch.manual_seed(seed)
1139
+ batch_size = len(seq_lens)
1140
+ max_seq_len = max(seq_lens)
1141
+ per_request = _pages_per_request(block_size, seq_lens)
1142
+ pages_per_batch = max(per_request)
1143
+ page_table = torch.full((batch_size, pages_per_batch), -1, dtype=torch.int32)
1144
+
1145
+ ref_k: list[torch.Tensor] = []
1146
+ ref_scale: list[torch.Tensor] = []
1147
+ next_page = 0
1148
+ for b, seq_len in enumerate(seq_lens):
1149
+ k_chunks = []
1150
+ scale_chunks = []
1151
+ for p in range(per_request[b]):
1152
+ page_table[b, p] = next_page
1153
+ k = _rand_fp8(block_size, INDEXER_HEAD_DIM)
1154
+ scale = torch.rand(block_size, dtype=torch.float32) * 0.05 + 0.01
1155
+ _write_page(pages[next_page], block_size, k, scale)
1156
+ tokens_here = min(block_size, seq_len - p * block_size)
1157
+ k_chunks.append(k[:tokens_here])
1158
+ scale_chunks.append(scale[:tokens_here])
1159
+ next_page += 1
1160
+ ref_k.append(torch.cat(k_chunks, dim=0))
1161
+ ref_scale.append(torch.cat(scale_chunks, dim=0))
1162
+
1163
+ q_fp8 = _rand_fp8(batch_size, num_heads, INDEXER_HEAD_DIM)
1164
+ weight = torch.rand(batch_size, num_heads, dtype=torch.float32) * 0.1
1165
+ seq_lens_t = torch.tensor(seq_lens, dtype=torch.int32)
1166
+ return q_fp8, weight, seq_lens_t, page_table, ref_k, ref_scale, max_seq_len
1167
+
1168
+
1169
+def _paged_mqa_logits_page_layout(
1170
+ block_size: int, seq_lens: list[int]
1171
+) -> tuple[int, int]:
1172
+ n_pages = sum(_pages_per_request(block_size, seq_lens))
1173
+ buf_width = block_size * INDEXER_HEAD_DIM + block_size * INDEXER_SCALE_BYTES
1174
+ return n_pages, buf_width
1175
+
1176
+
1177
+@pytest.mark.parametrize("block_size", [64, 256])
1178
+def test_fp8_paged_mqa_logits_cpu_matches_eager(block_size: int):
1179
+ num_heads = 4
1180
+ seq_lens = [1, block_size + 7, 3 * block_size]
1181
+ n_pages, buf_width = _paged_mqa_logits_page_layout(block_size, seq_lens)
1182
+ pages = torch.zeros((n_pages, buf_width), dtype=torch.uint8)
1183
+
1184
+ q_fp8, weight, seq_lens_t, page_table, ref_k, ref_scale, max_seq_len = (
1185
+ _build_paged_mqa_logits_batch(block_size, seq_lens, num_heads, pages, seed=0)
1186
+ )
1187
+
1188
+ logits = fp8_paged_mqa_logits_cpu(
1189
+ q_fp8, pages, weight, seq_lens_t, page_table, block_size, max_seq_len
1190
+ )
1191
+ ref = _reference_paged_mqa_logits(q_fp8, weight, ref_k, ref_scale, max_seq_len)
1192
+
1193
+ for b, seq_len in enumerate(seq_lens):
1194
+ torch.testing.assert_close(
1195
+ logits[b, :seq_len], ref[b, :seq_len], atol=1e-2, rtol=1e-2
1196
+ )
1197
+
1198
+
1199
+def test_fp8_paged_mqa_logits_cpu_non_contiguous_kv_cache_row_stride():
1200
+ """The paged K-cache in production is a per-layer view into a shared
1201
+ multi-layer allocation, so its page stride generally exceeds the page
1202
+ byte width -- check the kernel reads a padded/strided buffer identically
1203
+ to an equivalent tightly-packed one."""
1204
+ block_size = 64
1205
+ num_heads = 4
1206
+ seq_lens = [1, block_size + 7, 3 * block_size]
1207
+ n_pages, buf_width = _paged_mqa_logits_page_layout(block_size, seq_lens)
1208
+ cache_pad = 48
1209
+
1210
+ tight_pages = torch.zeros((n_pages, buf_width), dtype=torch.uint8)
1211
+ q_fp8, weight, seq_lens_t, page_table, _, _, max_seq_len = (
1212
+ _build_paged_mqa_logits_batch(
1213
+ block_size, seq_lens, num_heads, tight_pages, seed=0
1214
+ )
1215
+ )
1216
+
1217
+ padded = torch.zeros((n_pages, buf_width + cache_pad), dtype=torch.uint8)
1218
+ strided_pages = padded[:, :buf_width]
1219
+ _build_paged_mqa_logits_batch(
1220
+ block_size, seq_lens, num_heads, strided_pages, seed=0
1221
+ )
1222
+ assert strided_pages.stride(0) != strided_pages.size(1)
1223
+
1224
+ logits_tight = fp8_paged_mqa_logits_cpu(
1225
+ q_fp8, tight_pages, weight, seq_lens_t, page_table, block_size, max_seq_len
1226
+ )
1227
+ logits_strided = fp8_paged_mqa_logits_cpu(
1228
+ q_fp8, strided_pages, weight, seq_lens_t, page_table, block_size, max_seq_len
1229
+ )
1230
+ # Only [:seq_len] is ever written (and ever read downstream by
1231
+ # topk_transform_512_cpu) -- positions beyond that are uninitialized, so
1232
+ # comparing the full row would be comparing unrelated garbage.
1233
+ for b, seq_len in enumerate(seq_lens):
1234
+ torch.testing.assert_close(
1235
+ logits_tight[b, :seq_len], logits_strided[b, :seq_len]
1236
+ )
1237
+
1238
+
1239
+# ---------------------------------------------------------------------------
1240
+# fused_indexer_q_rope_quant_cpu (indexer.cpp): indexer Q-side RoPE + FP8
1241
+# quant, and the dispatcher wiring that routes to it on CPU
1242
+# ---------------------------------------------------------------------------
1243
+#
1244
+# Independent eager reference: this kernel's absmax reduction is over an
1245
+# unusual mix (raw NoPE values plus a bf16-round-tripped RoPE half), which
1246
+# per_token_group_quant_fp8 does not replicate. Allow <=1 fp8 ULP: the RoPE
1247
+# halves round a value sitting exactly on a tie to either side.
1248
+
1249
+INDEXER_Q_N_HEAD = 64
1250
+INDEXER_Q_MAX_POS = 4096
1251
+
1252
+
1253
+def _indexer_q_rope_quant_reference(
1254
+ positions: torch.Tensor,
1255
+ q: torch.Tensor,
1256
+ cos_sin_cache: torch.Tensor,
1257
+ weights: torch.Tensor,
1258
+ softmax_scale: float,
1259
+ head_scale: float,
1260
+) -> tuple[torch.Tensor, torch.Tensor]:
1261
+ _num_tokens, _num_heads, head_dim = q.shape
1262
+ half_rot_dim = cos_sin_cache.shape[-1] // 2
1263
+ rot_dim = 2 * half_rot_dim
1264
+ nope_dim = head_dim - rot_dim
1265
+
1266
+ q_f32 = q.float()
1267
+ cos_sin_f32 = cos_sin_cache.float()
1268
+ pos_cos_sin = cos_sin_f32[positions] # [T, rot_dim]
1269
+ cos = pos_cos_sin[:, :half_rot_dim].unsqueeze(1)
1270
+ sin = pos_cos_sin[:, half_rot_dim:].unsqueeze(1)
1271
+
1272
+ nope = q_f32[..., :nope_dim]
1273
+ rot = q_f32[..., nope_dim:]
1274
+ x_even = rot[..., 0::2]
1275
+ x_odd = rot[..., 1::2]
1276
+ r_even = (x_even * cos - x_odd * sin).to(torch.bfloat16).to(torch.float32)
1277
+ r_odd = (x_odd * cos + x_even * sin).to(torch.bfloat16).to(torch.float32)
1278
+ rope_out = torch.stack((r_even, r_odd), dim=-1).flatten(-2)
1279
+
1280
+ combined = torch.cat((nope, rope_out), dim=-1) # [T, H, head_dim]
1281
+ amax = combined.abs().amax(dim=-1).clamp(min=1e-4) # [T, H]
1282
+ exponent = torch.ceil(torch.log2(amax / FP8_MAX))
1283
+ scale = torch.exp2(exponent)
1284
+ inv_scale = scale.reciprocal().unsqueeze(-1)
1285
+ q_fp8 = (combined * inv_scale).clamp(-FP8_MAX, FP8_MAX).to(torch.float8_e4m3fn)
1286
+
1287
+ weights_out = weights.to(torch.float32) * scale * softmax_scale * head_scale
1288
+ return q_fp8, weights_out
1289
+
1290
+
1291
+@pytest.mark.parametrize("num_tokens", [1, 257])
1292
+def test_fused_indexer_q_rope_quant_cpu_matches_eager(num_tokens: int):
1293
+ torch.manual_seed(0)
1294
+
1295
+ q = torch.randn(num_tokens, INDEXER_Q_N_HEAD, INDEXER_HEAD_DIM, dtype=torch.float32)
1296
+ positions = torch.randint(0, INDEXER_Q_MAX_POS, (num_tokens,), dtype=torch.int64)
1297
+ cos_sin_cache = torch.randn(INDEXER_Q_MAX_POS, ROPE_DIM, dtype=torch.float32)
1298
+ weights = torch.randn(num_tokens, INDEXER_Q_N_HEAD, dtype=torch.float32)
1299
+ softmax_scale = INDEXER_HEAD_DIM**-0.5
1300
+ head_scale = INDEXER_Q_N_HEAD**-0.5
1301
+
1302
+ q_fp8_ref, weights_ref = _indexer_q_rope_quant_reference(
1303
+ positions, q, cos_sin_cache, weights, softmax_scale, head_scale
1304
+ )
1305
+
1306
+ q_fp8_fused = torch.empty_like(q, dtype=torch.float8_e4m3fn)
1307
+ weights_fused = torch.empty_like(weights, dtype=torch.float32)
1308
+ fused_indexer_q_rope_quant_cpu(
1309
+ positions,
1310
+ q,
1311
+ cos_sin_cache,
1312
+ q_fp8_fused,
1313
+ weights,
1314
+ softmax_scale,
1315
+ head_scale,
1316
+ weights_fused,
1317
+ )
1318
+
1319
+ max_ulp = int(fp8_ulp_distance(q_fp8_fused, q_fp8_ref).max().item())
1320
+ assert max_ulp <= 1, f"q_fp8 differs by {max_ulp} ULP (>1)"
1321
+ torch.testing.assert_close(weights_fused, weights_ref, rtol=1e-5, atol=1e-6)
1322
+
1323
+
1324
+def test_fused_indexer_q_rope_quant_dispatches_to_cpu_kernel():
1325
+ """End-to-end wiring check: the production dispatcher must route to the
1326
+ CPU op on a CPU tensor, not silently fall through to the Triton kernel
1327
+ (which would also "work" under triton-cpu, masking a dispatch
1328
+ regression). The dispatcher does its CPU import inside the CPU branch,
1329
+ so the patch target is the source module."""
1330
+ torch.manual_seed(0)
1331
+ num_tokens = 1
1332
+
1333
+ q = torch.randn(num_tokens, INDEXER_Q_N_HEAD, INDEXER_HEAD_DIM, dtype=torch.float32)
1334
+ positions = torch.randint(0, INDEXER_Q_MAX_POS, (num_tokens,), dtype=torch.int64)
1335
+ cos_sin_cache = torch.randn(INDEXER_Q_MAX_POS, ROPE_DIM, dtype=torch.float32)
1336
+ weights = torch.randn(num_tokens, INDEXER_Q_N_HEAD, dtype=torch.float32)
1337
+ softmax_scale = INDEXER_HEAD_DIM**-0.5
1338
+ head_scale = INDEXER_Q_N_HEAD**-0.5
1339
+
1340
+ q_fp8_ref, weights_ref = _indexer_q_rope_quant_reference(
1341
+ positions, q, cos_sin_cache, weights, softmax_scale, head_scale
1342
+ )
1343
+
1344
+ with mock.patch("vllm._custom_ops.fused_indexer_q_rope_quant_cpu") as mocked:
1345
+ fused_indexer_q_rope_quant(
1346
+ positions,
1347
+ q.clone(),
1348
+ cos_sin_cache,
1349
+ weights,
1350
+ softmax_scale,
1351
+ head_scale,
1352
+ use_fp4=False,
1353
+ )
1354
+ assert mocked.called, (
1355
+ "fused_indexer_q_rope_quant did not dispatch to the CPU kernel "
1356
+ "on a CPU tensor"
1357
+ )
1358
+
1359
+ q_fp8_fused, weights_fused = fused_indexer_q_rope_quant(
1360
+ positions,
1361
+ q.clone(),
1362
+ cos_sin_cache,
1363
+ weights,
1364
+ softmax_scale,
1365
+ head_scale,
1366
+ use_fp4=False,
1367
+ )
1368
+ max_ulp = int(fp8_ulp_distance(q_fp8_fused, q_fp8_ref).max().item())
1369
+ assert max_ulp <= 1, f"q_fp8 differs by {max_ulp} ULP (>1)"
1370
+ torch.testing.assert_close(weights_fused, weights_ref, rtol=1e-5, atol=1e-6)
1371
+
1372
+
1373
+# ---------------------------------------------------------------------------
1374
+# mhc_pre_cpu / mhc_post_cpu / hc_head_fused_cpu (mhc.cpp): CPU ports of the
1375
+# mHC gating kernels
1376
+# ---------------------------------------------------------------------------
1377
+
1378
+
1379
+def _rmsnorm_nw_torch(x: torch.Tensor, eps: float) -> torch.Tensor:
1380
+ """Weight-free RMSNorm over the last dim, matching hc_head_fused_cpu."""
1381
+ x_f32 = x.float()
1382
+ var = x_f32.square().mean(dim=-1, keepdim=True)
1383
+ return (x_f32 * torch.rsqrt(var + eps)).to(x.dtype)
1384
+
1385
+
1386
+def _hc_head_torch(
1387
+ hidden_states: torch.Tensor,
1388
+ hc_fn: torch.Tensor,
1389
+ hc_scale: torch.Tensor,
1390
+ hc_base: torch.Tensor,
1391
+ rms_norm_eps: float,
1392
+ hc_eps: float,
1393
+) -> torch.Tensor:
1394
+ """Eager reference for hc_head_fused_cpu: round-trips the RMSNorm
1395
+ output through bf16 before the linear, matching the kernel's bf16
1396
+ arithmetic."""
1397
+ hc_mult, hidden_size = hidden_states.shape[-2:]
1398
+ outer_shape = hidden_states.shape[:-2]
1399
+ hs_flat = hidden_states.reshape(-1, hc_mult, hidden_size)
1400
+
1401
+ x_flat = hs_flat.flatten(-2)
1402
+ x_normed = _rmsnorm_nw_torch(x_flat, rms_norm_eps)
1403
+ mixes = torch.nn.functional.linear(x_normed.float(), hc_fn)
1404
+ pre = torch.sigmoid(mixes * hc_scale + hc_base) + hc_eps
1405
+
1406
+ out = torch.sum(pre.unsqueeze(-1) * hs_flat.float(), dim=1).to(hidden_states.dtype)
1407
+ return out.view(*outer_shape, hidden_size)
1408
+
1409
+
1410
+@pytest.mark.parametrize("num_tokens", [1, 128])
1411
+@pytest.mark.parametrize("hidden_size", [4096, 7168])
1412
+def test_mhc_pre_cpu(num_tokens: int, hidden_size: int):
1413
+ """mhc_pre_cpu vs. the mhc_pre_torch eager reference it replaces."""
1414
+ torch.set_default_device(DEVICE)
1415
+ set_random_seed(0)
1416
+ hc_mult = 4
1417
+
1418
+ hc_mult3 = 2 * hc_mult + hc_mult * hc_mult
1419
+ residual = torch.randn((num_tokens, hc_mult, hidden_size), dtype=torch.bfloat16)
1420
+ fn = torch.randn((hc_mult3, hc_mult * hidden_size), dtype=torch.float32) * 1e-4
1421
+ hc_scale = torch.randn((3,), dtype=torch.float32) * 0.1
1422
+ hc_base = torch.randn((hc_mult3,), dtype=torch.float32) * 0.1
1423
+
1424
+ rms_eps = hc_eps = 1e-6
1425
+ hc_post_mult_value = 2.0
1426
+ sinkhorn_repeat = 8
1427
+
1428
+ ref = mhc_kernels.mhc_pre_torch(
1429
+ residual,
1430
+ fn,
1431
+ hc_scale,
1432
+ hc_base,
1433
+ rms_eps,
1434
+ hc_eps,
1435
+ hc_eps,
1436
+ hc_post_mult_value,
1437
+ sinkhorn_repeat,
1438
+ )
1439
+ out = mhc_kernels.mhc_pre_cpu(
1440
+ residual,
1441
+ fn,
1442
+ hc_scale,
1443
+ hc_base,
1444
+ rms_eps,
1445
+ hc_eps,
1446
+ hc_eps,
1447
+ hc_post_mult_value,
1448
+ sinkhorn_repeat,
1449
+ )
1450
+
1451
+ for actual, expected in zip(out, ref, strict=True):
1452
+ torch.testing.assert_close(actual, expected, atol=5e-2, rtol=1e-2)
1453
+
1454
+
1455
+@pytest.mark.parametrize("num_tokens", [1, 128])
1456
+@pytest.mark.parametrize("hidden_size", [4096, 7168])
1457
+def test_mhc_post_cpu(num_tokens: int, hidden_size: int):
1458
+ """mhc_post_cpu vs. the mhc_post_torch eager reference it replaces."""
1459
+ torch.set_default_device(DEVICE)
1460
+ set_random_seed(0)
1461
+ hc_mult = 4
1462
+
1463
+ x = torch.randn((num_tokens, hidden_size), dtype=torch.bfloat16)
1464
+ residual = torch.randn((num_tokens, hc_mult, hidden_size), dtype=torch.bfloat16)
1465
+ post_layer_mix = torch.randn((num_tokens, hc_mult, 1), dtype=torch.float32)
1466
+ comb_res_mix = torch.randn((num_tokens, hc_mult, hc_mult), dtype=torch.float32)
1467
+
1468
+ ref = mhc_kernels.mhc_post_torch(x, residual, post_layer_mix, comb_res_mix)
1469
+ out = mhc_kernels.mhc_post_cpu(x, residual, post_layer_mix, comb_res_mix)
1470
+
1471
+ torch.testing.assert_close(out, ref, atol=5e-2, rtol=1e-2)
1472
+
1473
+
1474
+@pytest.mark.parametrize("num_tokens", [1, 128])
1475
+@pytest.mark.parametrize("hidden_size", [4096, 7168])
1476
+def test_hc_head_cpu(num_tokens: int, hidden_size: int):
1477
+ """hc_head_fused_cpu vs. the _hc_head_torch eager reference it
1478
+ replaces."""
1479
+ torch.set_default_device(DEVICE)
1480
+ set_random_seed(0)
1481
+ hc_mult = 4
1482
+
1483
+ residual = torch.randn((num_tokens, hc_mult, hidden_size), dtype=torch.bfloat16)
1484
+ fn = torch.randn((hc_mult, hc_mult * hidden_size), dtype=torch.float32) * 1e-4
1485
+ hc_scale = torch.randn((1,), dtype=torch.float32) * 0.1
1486
+ hc_base = torch.randn((hc_mult,), dtype=torch.float32) * 0.1
1487
+ rms_eps = hc_eps = 1e-6
1488
+
1489
+ ref = _hc_head_torch(residual, fn, hc_scale, hc_base, rms_eps, hc_eps)
1490
+ out = mhc_kernels.hc_head_fused_cpu(
1491
+ residual, fn, hc_scale, hc_base, rms_eps, hc_eps
1492
+ )
1493
+
1494
+ torch.testing.assert_close(out, ref, atol=5e-2, rtol=1e-2)
tests/models/multimodal/processing/test_common.py
+10
@@ -92,6 +92,11 @@ _XPU_EXCLUDED_MODEL_IDS = {
92
"thinkingmachines/Inkling-NVFP4",
93
}
94
95
+_CPU_EXCLUDED_MODEL_IDS = {
96
+ # DeepSeek-V4 vision variant is only supported on NVIDIA GPUs for now.
97
+ "deepseek-ai/DeepSeek-V4-Flash-Vision-Exp",
98
+}
99
+
100
101
def _iter_model_ids_to_test(model_arch_list: AbstractSet[str]):
102
for model_arch in model_arch_list:
@@ -113,6 +118,11 @@ def _get_model_ids_to_test(model_arch_list: AbstractSet[str]):
118
while excluded_model_id in model_ids:
119
model_ids.remove(excluded_model_id)
120
121
+ if current_platform.is_cpu():
122
+ for excluded_model_id in _CPU_EXCLUDED_MODEL_IDS:
123
+ while excluded_model_id in model_ids:
124
+ model_ids.remove(excluded_model_id)
125
+
126
return model_ids
127
128
tests/v1/attention/test_sparse_mla_backends.py
+5
-3
@@ -56,7 +56,7 @@ from vllm.v1.attention.backends.mla.flashmla_sparse import (
56
FlashMLASparseMetadataBuilder,
57
triton_convert_req_index_to_global_index,
58
)
59
-from vllm.v1.attention.backends.mla.indexer import split_indexer_prefill_chunks
59
+from vllm.v1.attention.backends.mla.indexer import DeepseekV32IndexerMetadataBuilder
60
from vllm.v1.attention.backends.utils import (
61
split_decodes_and_prefills,
62
split_prefill_chunks,
@@ -1424,7 +1424,7 @@ def test_sparse_backend_prefill_correctness(
1424
def test_split_indexer_prefill_chunks(
1425
seq_lens, query_lens, workspace_size, max_logits_bytes, expected
1426
):
1427
- out = split_indexer_prefill_chunks(
1427
+ out = DeepseekV32IndexerMetadataBuilder._split_indexer_prefill_chunks(
1428
seq_lens,
1429
query_lens,
1430
workspace_size,
@@ -1438,7 +1438,9 @@ def test_split_indexer_prefill_chunks_single_request_overflow():
1438
seq_lens = torch.tensor([1000, 50])
1439
query_lens = torch.tensor([100, 5])
1440
1441
- out = split_indexer_prefill_chunks(seq_lens, query_lens, 2000, 1000)
1441
+ out = DeepseekV32IndexerMetadataBuilder._split_indexer_prefill_chunks(
1442
+ seq_lens, query_lens, 2000, 1000
1443
+ )
1444
# max_logits_elems = 250, N=1000 -> max_q = 1 -> 100 query sub-chunks
1445
expected = [(slice(0, 1), slice(i, i + 1)) for i in range(100)]
1446
# req1: M=5, N=50 -> 250 elems fits budget
vllm/_custom_ops.py
+397
-37
@@ -3357,31 +3357,6 @@ class CPUQuantMethod(IntEnum):
3357
MXFP4 = 4
3358
3359
3360
-if hasattr(torch.ops._C, "fused_experts_cpu"):
3361
-
3362
- @register_fake("_C::fused_experts_cpu")
3363
- def fused_experts_cpu_fake(
3364
- hidden_states: torch.Tensor,
3365
- w1: torch.Tensor,
3366
- w2: torch.Tensor,
3367
- topk_weights: torch.Tensor,
3368
- topk_ids: torch.Tensor,
3369
- inplace: bool,
3370
- moe_comp_method: CPUQuantMethod,
3371
- w1_scale: torch.Tensor | None,
3372
- w2_scale: torch.Tensor | None,
3373
- w1_zero: torch.Tensor | None,
3374
- w2_zero: torch.Tensor | None,
3375
- block_size: list[int] | None,
3376
- w1_bias: torch.Tensor | None,
3377
- w2_bias: torch.Tensor | None,
3378
- alpha: float | None,
3379
- limit: float | None,
3380
- is_vnni: bool,
3381
- ) -> torch.Tensor:
3382
- return torch.empty_like(hidden_states)
3383
-
3384
-
3360
if hasattr(torch.ops._C, "dynamic_4bit_int_moe"):
3361
3362
@register_fake("_C::dynamic_4bit_int_moe")
@@ -3401,12 +3376,12 @@ if hasattr(torch.ops._C, "dynamic_4bit_int_moe"):
3376
3377
3378
def fused_experts_cpu(
3379
+ out: torch.Tensor,
3380
hidden_states: torch.Tensor,
3381
w1: torch.Tensor,
3382
w2: torch.Tensor,
3383
topk_weights: torch.Tensor,
3384
topk_ids: torch.Tensor,
3409
- inplace: bool,
3385
moe_comp_method: CPUQuantMethod,
3386
w1_scale: torch.Tensor | None,
3387
w2_scale: torch.Tensor | None,
@@ -3418,14 +3393,14 @@ def fused_experts_cpu(
3393
alpha: float | None = None,
3394
limit: float | None = None,
3395
is_vnni: bool = True,
3421
-) -> torch.Tensor:
3422
- return torch.ops._C.fused_experts_cpu(
3396
+) -> None:
3397
+ torch.ops._C.fused_experts_cpu(
3398
+ out,
3399
hidden_states,
3400
w1,
3401
w2,
3402
topk_weights,
3403
topk_ids,
3428
- inplace,
3404
moe_comp_method,
3405
w1_scale,
3406
w2_scale,
@@ -4059,6 +4034,393 @@ def amx_mla_concat_and_cache(
4034
torch.ops._C.concat_and_cache_mla_cpu(kv_c_normed, k_pe, kv_cache, slot_mapping)
4035
4036
4037
+if hasattr(torch.ops._C, "hc_pre_fused_cpu"):
4038
+
4039
+ @register_fake("_C::hc_pre_fused_cpu")
4040
+ def hc_pre_fused_cpu_fake(
4041
+ x: torch.Tensor,
4042
+ hc_fn: torch.Tensor,
4043
+ hc_scale: torch.Tensor,
4044
+ hc_base: torch.Tensor,
4045
+ hc_mult: int,
4046
+ sinkhorn_iters: int,
4047
+ rms_eps: float,
4048
+ hc_eps: float,
4049
+ ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
4050
+ t, _, d = x.shape
4051
+ y = x.new_empty((t, d))
4052
+ post = x.new_empty((t, hc_mult), dtype=torch.float32)
4053
+ comb = x.new_empty((t, hc_mult, hc_mult), dtype=torch.float32)
4054
+ return y, post, comb
4055
+
4056
+
4057
+def hc_pre_fused_cpu(
4058
+ x: torch.Tensor,
4059
+ hc_fn: torch.Tensor,
4060
+ hc_scale: torch.Tensor,
4061
+ hc_base: torch.Tensor,
4062
+ hc_mult: int,
4063
+ sinkhorn_iters: int,
4064
+ rms_eps: float,
4065
+ hc_eps: float,
4066
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
4067
+ return torch.ops._C.hc_pre_fused_cpu(
4068
+ x, hc_fn, hc_scale, hc_base, hc_mult, sinkhorn_iters, rms_eps, hc_eps
4069
+ )
4070
+
4071
+
4072
+if hasattr(torch.ops._C, "hc_post_fused_cpu"):
4073
+
4074
+ @register_fake("_C::hc_post_fused_cpu")
4075
+ def hc_post_fused_cpu_fake(
4076
+ x: torch.Tensor,
4077
+ residual: torch.Tensor,
4078
+ post: torch.Tensor,
4079
+ comb: torch.Tensor,
4080
+ ) -> torch.Tensor:
4081
+ t, d = x.shape
4082
+ hc = residual.size(1)
4083
+ return x.new_empty((t, hc, d))
4084
+
4085
+
4086
+def hc_post_fused_cpu(
4087
+ x: torch.Tensor,
4088
+ residual: torch.Tensor,
4089
+ post: torch.Tensor,
4090
+ comb: torch.Tensor,
4091
+) -> torch.Tensor:
4092
+ return torch.ops._C.hc_post_fused_cpu(x, residual, post, comb)
4093
+
4094
+
4095
+if hasattr(torch.ops._C, "hc_head_fused_cpu"):
4096
+
4097
+ @register_fake("_C::hc_head_fused_cpu")
4098
+ def hc_head_fused_cpu_fake(
4099
+ x: torch.Tensor,
4100
+ hc_fn: torch.Tensor,
4101
+ hc_scale: torch.Tensor,
4102
+ hc_base: torch.Tensor,
4103
+ hc_eps: float,
4104
+ norm_eps: float,
4105
+ ) -> torch.Tensor:
4106
+ t, _, d = x.shape
4107
+ return x.new_empty((t, d))
4108
+
4109
+
4110
+def hc_head_fused_cpu(
4111
+ x: torch.Tensor,
4112
+ hc_fn: torch.Tensor,
4113
+ hc_scale: torch.Tensor,
4114
+ hc_base: torch.Tensor,
4115
+ hc_eps: float,
4116
+ norm_eps: float,
4117
+) -> torch.Tensor:
4118
+ return torch.ops._C.hc_head_fused_cpu(x, hc_fn, hc_scale, hc_base, hc_eps, norm_eps)
4119
+
4120
+
4121
+if hasattr(torch.ops._C, "fused_qnorm_rope_kv_insert_cpu"):
4122
+
4123
+ @register_fake("_C::fused_qnorm_rope_kv_insert_cpu")
4124
+ def fused_qnorm_rope_kv_insert_cpu_fake(
4125
+ q: torch.Tensor,
4126
+ kv: torch.Tensor,
4127
+ positions: torch.Tensor,
4128
+ swa_kv_cache_2d: torch.Tensor,
4129
+ slot_mapping: torch.Tensor,
4130
+ cos_sin_cache: torch.Tensor,
4131
+ q_head_padded: int,
4132
+ eps: float,
4133
+ cache_block_size: int,
4134
+ ) -> torch.Tensor:
4135
+ num_tokens_full, _, head_dim = q.shape
4136
+ return q.new_empty((num_tokens_full, q_head_padded, head_dim))
4137
+
4138
+
4139
+def fused_qnorm_rope_kv_insert_cpu(
4140
+ q: torch.Tensor,
4141
+ kv: torch.Tensor,
4142
+ positions: torch.Tensor,
4143
+ swa_kv_cache_2d: torch.Tensor,
4144
+ slot_mapping: torch.Tensor,
4145
+ cos_sin_cache: torch.Tensor,
4146
+ q_head_padded: int,
4147
+ eps: float,
4148
+ cache_block_size: int,
4149
+) -> torch.Tensor:
4150
+ return torch.ops._C.fused_qnorm_rope_kv_insert_cpu(
4151
+ q,
4152
+ kv,
4153
+ positions,
4154
+ swa_kv_cache_2d,
4155
+ slot_mapping,
4156
+ cos_sin_cache,
4157
+ q_head_padded,
4158
+ eps,
4159
+ cache_block_size,
4160
+ )
4161
+
4162
+
4163
+def flash_mla_with_kvcache_cpu(
4164
+ out: torch.Tensor,
4165
+ q: torch.Tensor,
4166
+ window_cache_2d: torch.Tensor,
4167
+ window_slots: torch.Tensor,
4168
+ window_block_size: int,
4169
+ compressed_cache_2d: torch.Tensor,
4170
+ compressed_slots: torch.Tensor,
4171
+ compressed_block_size: int,
4172
+ attn_sink: torch.Tensor,
4173
+ scale: float,
4174
+) -> None:
4175
+ torch.ops._C.flash_mla_with_kvcache_cpu(
4176
+ out,
4177
+ q,
4178
+ window_cache_2d,
4179
+ window_slots,
4180
+ window_block_size,
4181
+ compressed_cache_2d,
4182
+ compressed_slots,
4183
+ compressed_block_size,
4184
+ attn_sink,
4185
+ scale,
4186
+ )
4187
+
4188
+
4189
+def save_partial_states_cpu(
4190
+ kv: torch.Tensor,
4191
+ score: torch.Tensor,
4192
+ ape: torch.Tensor,
4193
+ positions: torch.Tensor,
4194
+ state_cache: torch.Tensor,
4195
+ slot_mapping: torch.Tensor,
4196
+) -> None:
4197
+ torch.ops._C.save_partial_states_cpu(
4198
+ kv, score, ape, positions, state_cache, slot_mapping
4199
+ )
4200
+
4201
+
4202
+def compress_norm_rope_store_cpu(
4203
+ state_cache: torch.Tensor,
4204
+ gather_slots: torch.Tensor,
4205
+ positions: torch.Tensor,
4206
+ kv_slot_mapping: torch.Tensor,
4207
+ rms_norm_weight: torch.Tensor,
4208
+ rms_norm_eps: float,
4209
+ cos_sin_cache: torch.Tensor,
4210
+ kv_cache_2d: torch.Tensor,
4211
+ kv_cache_block_size: int,
4212
+ compress_ratio: int,
4213
+) -> None:
4214
+ torch.ops._C.compress_norm_rope_store_cpu(
4215
+ state_cache,
4216
+ gather_slots,
4217
+ positions,
4218
+ kv_slot_mapping,
4219
+ rms_norm_weight,
4220
+ rms_norm_eps,
4221
+ cos_sin_cache,
4222
+ kv_cache_2d,
4223
+ kv_cache_block_size,
4224
+ compress_ratio,
4225
+ )
4226
+
4227
+
4228
+def compress_norm_rope_store_indexer_cpu(
4229
+ state_cache: torch.Tensor,
4230
+ gather_slots: torch.Tensor,
4231
+ positions: torch.Tensor,
4232
+ kv_slot_mapping: torch.Tensor,
4233
+ rms_norm_weight: torch.Tensor,
4234
+ rms_norm_eps: float,
4235
+ cos_sin_cache: torch.Tensor,
4236
+ kv_cache_2d: torch.Tensor,
4237
+ kv_cache_block_size: int,
4238
+ compress_ratio: int,
4239
+) -> None:
4240
+ torch.ops._C.compress_norm_rope_store_indexer_cpu(
4241
+ state_cache,
4242
+ gather_slots,
4243
+ positions,
4244
+ kv_slot_mapping,
4245
+ rms_norm_weight,
4246
+ rms_norm_eps,
4247
+ cos_sin_cache,
4248
+ kv_cache_2d,
4249
+ kv_cache_block_size,
4250
+ compress_ratio,
4251
+ )
4252
+
4253
+
4254
+if hasattr(torch.ops._C, "fp8_paged_mqa_logits_cpu"):
4255
+
4256
+ @register_fake("_C::fp8_paged_mqa_logits_cpu")
4257
+ def fp8_paged_mqa_logits_cpu_fake(
4258
+ q_fp8: torch.Tensor,
4259
+ kvcache_fp8: torch.Tensor,
4260
+ weight: torch.Tensor,
4261
+ seq_lens: torch.Tensor,
4262
+ page_table: torch.Tensor,
4263
+ block_size: int,
4264
+ max_seq_len: int,
4265
+ ) -> torch.Tensor:
4266
+ batch_size = q_fp8.size(0)
4267
+ return q_fp8.new_empty((batch_size, max_seq_len), dtype=torch.float32)
4268
+
4269
+
4270
+def fp8_paged_mqa_logits_cpu(
4271
+ q_fp8: torch.Tensor,
4272
+ kvcache_fp8: torch.Tensor,
4273
+ weight: torch.Tensor,
4274
+ seq_lens: torch.Tensor,
4275
+ page_table: torch.Tensor,
4276
+ block_size: int,
4277
+ max_seq_len: int,
4278
+) -> torch.Tensor:
4279
+ return torch.ops._C.fp8_paged_mqa_logits_cpu(
4280
+ q_fp8, kvcache_fp8, weight, seq_lens, page_table, block_size, max_seq_len
4281
+ )
4282
+
4283
+
4284
+def topk_transform_512_cpu(
4285
+ scores: torch.Tensor,
4286
+ seq_lens: torch.Tensor,
4287
+ page_tables: torch.Tensor,
4288
+ out_page_indices: torch.Tensor,
4289
+ page_size: int,
4290
+ out_raw_indices: torch.Tensor | None = None,
4291
+) -> None:
4292
+ torch.ops._C.topk_transform_512_cpu(
4293
+ scores, seq_lens, page_tables, out_page_indices, page_size, out_raw_indices
4294
+ )
4295
+
4296
+
4297
+def fused_indexer_q_rope_quant_cpu(
4298
+ positions: torch.Tensor,
4299
+ index_q: torch.Tensor,
4300
+ index_q_cos_sin_cache: torch.Tensor,
4301
+ index_q_fp8: torch.Tensor,
4302
+ index_weights: torch.Tensor,
4303
+ index_weights_softmax_scale: float,
4304
+ index_weights_head_scale: float,
4305
+ index_weights_out: torch.Tensor,
4306
+) -> None:
4307
+ torch.ops._C.fused_indexer_q_rope_quant_cpu(
4308
+ positions,
4309
+ index_q,
4310
+ index_q_cos_sin_cache,
4311
+ index_q_fp8,
4312
+ index_weights,
4313
+ index_weights_softmax_scale,
4314
+ index_weights_head_scale,
4315
+ index_weights_out,
4316
+ )
4317
+
4318
+
4319
+if hasattr(torch.ops._C, "biased_topk_cpu"):
4320
+
4321
+ @register_fake("_C::biased_topk_cpu")
4322
+ def biased_topk_cpu_fake(
4323
+ hidden_states: torch.Tensor,
4324
+ gating_output: torch.Tensor,
4325
+ correction_bias: torch.Tensor,
4326
+ topk: int,
4327
+ renormalize: bool,
4328
+ scoring_func: str,
4329
+ num_fused_shared_experts: int,
4330
+ routed_scaling_factor: float | None,
4331
+ apply_routed_scaling_factor_on_output: bool,
4332
+ ) -> tuple[torch.Tensor, torch.Tensor]:
4333
+ num_tokens = hidden_states.size(0)
4334
+ topk_weights = hidden_states.new_empty((num_tokens, topk), dtype=torch.float32)
4335
+ topk_ids = hidden_states.new_empty((num_tokens, topk), dtype=torch.int32)
4336
+ return topk_weights, topk_ids
4337
+
4338
+
4339
+def biased_topk_cpu(
4340
+ hidden_states: torch.Tensor,
4341
+ gating_output: torch.Tensor,
4342
+ correction_bias: torch.Tensor,
4343
+ topk: int,
4344
+ renormalize: bool,
4345
+ scoring_func: str,
4346
+ num_fused_shared_experts: int,
4347
+ routed_scaling_factor: float | None,
4348
+ apply_routed_scaling_factor_on_output: bool,
4349
+) -> tuple[torch.Tensor, torch.Tensor]:
4350
+ return torch.ops._C.biased_topk_cpu(
4351
+ hidden_states,
4352
+ gating_output,
4353
+ correction_bias,
4354
+ topk,
4355
+ renormalize,
4356
+ scoring_func,
4357
+ num_fused_shared_experts,
4358
+ routed_scaling_factor,
4359
+ apply_routed_scaling_factor_on_output,
4360
+ )
4361
+
4362
+
4363
+if hasattr(torch.ops._C, "hash_topk_cpu"):
4364
+
4365
+ @register_fake("_C::hash_topk_cpu")
4366
+ def hash_topk_cpu_fake(
4367
+ gating_output: torch.Tensor,
4368
+ tid2eid: torch.Tensor,
4369
+ topk: int,
4370
+ scoring_func: str,
4371
+ num_fused_shared_experts: int,
4372
+ num_experts: int,
4373
+ routed_scaling_factor: float,
4374
+ ) -> tuple[torch.Tensor, torch.Tensor]:
4375
+ num_tokens = gating_output.size(0)
4376
+ topk_weights = gating_output.new_empty((num_tokens, topk), dtype=torch.float32)
4377
+ topk_ids = gating_output.new_empty((num_tokens, topk), dtype=torch.int32)
4378
+ return topk_weights, topk_ids
4379
+
4380
+
4381
+def hash_topk_cpu(
4382
+ gating_output: torch.Tensor,
4383
+ tid2eid: torch.Tensor,
4384
+ topk: int,
4385
+ scoring_func: str,
4386
+ num_fused_shared_experts: int,
4387
+ num_experts: int,
4388
+ routed_scaling_factor: float,
4389
+) -> tuple[torch.Tensor, torch.Tensor]:
4390
+ return torch.ops._C.hash_topk_cpu(
4391
+ gating_output,
4392
+ tid2eid,
4393
+ topk,
4394
+ scoring_func,
4395
+ num_fused_shared_experts,
4396
+ num_experts,
4397
+ routed_scaling_factor,
4398
+ )
4399
+
4400
+
4401
+if hasattr(torch.ops._C, "inverse_gptj_rope_o_proj_cpu"):
4402
+
4403
+ @register_fake("_C::inverse_gptj_rope_o_proj_cpu")
4404
+ def inverse_gptj_rope_o_proj_cpu_fake(
4405
+ o: torch.Tensor,
4406
+ positions: torch.Tensor,
4407
+ cos_sin_cache: torch.Tensor,
4408
+ rope_dim: int,
4409
+ ) -> torch.Tensor:
4410
+ return o.new_empty(o.shape, dtype=torch.float32)
4411
+
4412
+
4413
+def inverse_gptj_rope_o_proj_cpu(
4414
+ o: torch.Tensor,
4415
+ positions: torch.Tensor,
4416
+ cos_sin_cache: torch.Tensor,
4417
+ rope_dim: int,
4418
+) -> torch.Tensor:
4419
+ return torch.ops._C.inverse_gptj_rope_o_proj_cpu(
4420
+ o, positions, cos_sin_cache, rope_dim
4421
+ )
4422
+
4423
+
4424
def cpu_gemm_wna16(
4425
input: torch.Tensor,
4426
q_weight: torch.Tensor,
@@ -4107,6 +4469,7 @@ def cpu_prepack_moe_weight_int8(
4469
4470
4471
def cpu_fused_moe(
4472
+ out: torch.Tensor,
4473
input: torch.Tensor,
4474
w13: torch.Tensor,
4475
w2: torch.Tensor,
@@ -4117,10 +4480,9 @@ def cpu_fused_moe(
4480
act: str,
4481
isa: str,
4482
skip_weighted: bool = False,
4120
-) -> torch.Tensor:
4121
- output = torch.empty_like(input)
4483
+) -> None:
4484
torch.ops._C.cpu_fused_moe(
4123
- output,
4485
+ out,
4486
input,
4487
w13,
4488
w2,
@@ -4132,10 +4494,10 @@ def cpu_fused_moe(
4494
act,
4495
isa,
4496
)
4135
- return output
4497
4498
4499
def cpu_fused_moe_int8(
4500
+ out: torch.Tensor,
4501
input: torch.Tensor,
4502
w13: torch.Tensor,
4503
w2: torch.Tensor,
@@ -4148,10 +4510,9 @@ def cpu_fused_moe_int8(
4510
act: str,
4511
isa: str,
4512
skip_weighted: bool = False,
4151
-) -> torch.Tensor:
4152
- output = torch.empty_like(input)
4513
+) -> None:
4514
torch.ops._C.cpu_fused_moe_int8(
4154
- output,
4515
+ out,
4516
input,
4517
w13,
4518
w2,
@@ -4165,7 +4526,6 @@ def cpu_fused_moe_int8(
4526
act,
4527
isa,
4528
)
4168
- return output
4529
4530
4531
if hasattr(torch.ops._qutlass_C, "matmul_mxf4_bf16_tn"):
vllm/model_executor/kernels/linear/scaled_mm/cpu.py
+6
-1
@@ -6,6 +6,9 @@ import torch
6
7
from vllm import _custom_ops as ops
8
from vllm.model_executor.layers.quantization.utils import replace_parameter
9
+from vllm.model_executor.layers.quantization.utils.fp8_utils import (
10
+ _upcast_e8m0_to_fp32,
11
+)
12
from vllm.model_executor.layers.quantization.utils.w8a8_utils import (
13
convert_to_channelwise,
14
)
@@ -281,10 +284,12 @@ class CPUFp8BlockScaledMMKernel(Fp8BlockScaledMMLinearKernel):
284
else params.weight_scale
285
)
286
assert weight_scale is not None
287
+ if weight_scale.dtype in (torch.float8_e8m0fnu, torch.uint8):
288
+ weight_scale = _upcast_e8m0_to_fp32(weight_scale.data)
289
replace_parameter(
290
layer,
291
scale_attr,
287
- torch.nn.Parameter(weight_scale.data, requires_grad=False),
292
+ torch.nn.Parameter(weight_scale.data.contiguous(), requires_grad=False),
293
)
294
295
def apply_weights(
vllm/model_executor/kernels/mhc/__init__.py
+4
@@ -1,6 +1,7 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
from .aiter import *
4
+from .cpu import *
5
from .tilelang import *
6
from .torch import *
7
from .triton import *
@@ -14,6 +15,9 @@ __all__ = [
15
"mhc_post_aiter",
16
"mhc_fused_post_pre_aiter",
17
"hc_head_fused_aiter",
18
+ "mhc_pre_cpu",
19
+ "mhc_post_cpu",
20
+ "hc_head_fused_cpu",
21
"mhc_pre_tilelang",
22
"mhc_post_tilelang",
23
"mhc_fused_post_pre_tilelang",
vllm/model_executor/kernels/mhc/cpu.py
new
+96
@@ -0,0 +1,96 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+import torch
4
+
5
+from vllm import _custom_ops as ops
6
+
7
+
8
+def mhc_pre_cpu(
9
+ residual: torch.Tensor,
10
+ fn: torch.Tensor,
11
+ hc_scale: torch.Tensor,
12
+ hc_base: torch.Tensor,
13
+ rms_eps: float,
14
+ hc_pre_eps: float,
15
+ hc_sinkhorn_eps: float,
16
+ hc_post_mult_value: float,
17
+ sinkhorn_repeat: int,
18
+ n_splits: int = 1,
19
+ norm_weight: torch.Tensor | None = None,
20
+ norm_eps: float = 0.0,
21
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
22
+ """CPU-ported mHC pre block (see `mhc_pre_torch` for the eager reference).
23
+
24
+ The ported kernel (`hc_pre_fused_cpu`) only exposes one merged `hc_eps`
25
+ (used for both `hc_pre_eps`/`hc_sinkhorn_eps`) and hardcodes the post-mix
26
+ multiplier to 2.0 -- true of every real call site in
27
+ `models/deepseek_v4/cpu/model.py`, so this is not a capability loss here.
28
+ """
29
+ assert n_splits == 1, "mhc_pre_cpu does not support n_splits != 1"
30
+ assert hc_pre_eps == hc_sinkhorn_eps, (
31
+ "mhc_pre_cpu requires hc_pre_eps == hc_sinkhorn_eps (single merged hc_eps)"
32
+ )
33
+ assert hc_post_mult_value == 2.0, "mhc_pre_cpu hardcodes the post-mix multiplier"
34
+
35
+ hc_mult, hidden_size = residual.shape[-2:]
36
+ outer_shape = residual.shape[:-2]
37
+ x_flat = residual.reshape(-1, hc_mult, hidden_size)
38
+
39
+ layer_input, post, comb = ops.hc_pre_fused_cpu(
40
+ x_flat,
41
+ fn,
42
+ hc_scale,
43
+ hc_base,
44
+ hc_mult,
45
+ sinkhorn_repeat,
46
+ rms_eps,
47
+ hc_pre_eps,
48
+ )
49
+
50
+ post_mix = post.view(*outer_shape, hc_mult, 1)
51
+ comb_mix = comb.view(*outer_shape, hc_mult, hc_mult)
52
+ layer_input = layer_input.view(*outer_shape, hidden_size)
53
+ return post_mix, comb_mix, layer_input
54
+
55
+
56
+def mhc_post_cpu(
57
+ x: torch.Tensor,
58
+ residual: torch.Tensor,
59
+ post_layer_mix: torch.Tensor,
60
+ comb_res_mix: torch.Tensor,
61
+) -> torch.Tensor:
62
+ """CPU-ported mHC post block (see `mhc_post_torch` for the eager reference)."""
63
+ hc_mult, hidden_size = residual.shape[-2:]
64
+ outer_shape = residual.shape[:-2]
65
+
66
+ x_flat = x.reshape(-1, hidden_size)
67
+ residual_flat = residual.reshape(-1, hc_mult, hidden_size)
68
+ post_flat = post_layer_mix.reshape(-1, hc_mult).float()
69
+ comb_flat = comb_res_mix.reshape(-1, hc_mult, hc_mult).float()
70
+
71
+ out = ops.hc_post_fused_cpu(x_flat, residual_flat, post_flat, comb_flat)
72
+ return out.view(*outer_shape, hc_mult, hidden_size)
73
+
74
+
75
+def hc_head_fused_cpu(
76
+ hidden_states: torch.Tensor,
77
+ hc_fn: torch.Tensor,
78
+ hc_scale: torch.Tensor,
79
+ hc_base: torch.Tensor,
80
+ rms_norm_eps: float,
81
+ hc_eps: float,
82
+) -> torch.Tensor:
83
+ """CPU-ported HC head reduction (see `test_hc_head_cpu` in
84
+ tests/kernels/test_mhc_kernels.py for the eager reference this is tested
85
+ against).
86
+
87
+ The ported kernel's C++ signature takes `(hc_eps, norm_eps)` -- the
88
+ opposite order from this wrapper's `(rms_norm_eps, hc_eps)`, which matches
89
+ the real call site in `models/deepseek_v4/cpu/model.py`.
90
+ """
91
+ hc_mult, hidden_size = hidden_states.shape[-2:]
92
+ outer_shape = hidden_states.shape[:-2]
93
+ hs_flat = hidden_states.reshape(-1, hc_mult, hidden_size)
94
+
95
+ out = ops.hc_head_fused_cpu(hs_flat, hc_fn, hc_scale, hc_base, hc_eps, rms_norm_eps)
96
+ return out.view(*outer_shape, hidden_size)
vllm/model_executor/layers/fused_moe/experts/cpu_int4_moe.py
+1
-1
@@ -218,7 +218,7 @@ class CPUExpertsInt4(mk.FusedMoEExpertsMonolithic):
218
Returns:
219
Output tensor after MoE computation
220
"""
221
- from vllm.model_executor.layers.fused_moe.experts.cpu_moe import (
221
+ from vllm.model_executor.layers.fused_moe.router.cpu_router import (
222
select_experts,
223
)
224
vllm/model_executor/layers/fused_moe/experts/cpu_moe.py
+262
-372
@@ -4,7 +4,6 @@
4
5
import math
6
import sys
7
-from collections.abc import Callable
7
from typing import cast
8
9
import torch
@@ -34,6 +33,9 @@ from vllm.model_executor.layers.fused_moe.config import (
33
FusedMoEQuantConfig,
34
RoutingMethodType,
35
)
36
+from vllm.model_executor.layers.fused_moe.topk_weight_and_reduce import (
37
+ TopKWeightAndReduceNoOP,
38
+)
39
from vllm.model_executor.layers.quantization.utils.quant_utils import (
40
QuantKey,
41
kFp8Dynamic128Sym,
@@ -48,133 +50,26 @@ from vllm.platforms import CpuArchEnum, current_platform
50
from vllm.utils.math_utils import round_up
51
52
logger = init_logger(__name__)
51
-# ===========================================================================
52
-# Routing
53
-# ===========================================================================
54
-
55
-
56
-def grouped_topk(
57
- hidden_states: torch.Tensor,
58
- gating_output: torch.Tensor,
59
- topk: int,
60
- renormalize: bool,
61
- num_expert_group: int = 0,
62
- topk_group: int = 0,
63
- scoring_func: str = "softmax",
64
- routed_scaling_factor: float = 1.0,
65
- e_score_correction_bias: torch.Tensor | None = None,
66
-) -> tuple[torch.Tensor, torch.Tensor]:
67
- assert hidden_states.shape[0] == gating_output.shape[0], "Number of tokens mismatch"
68
-
69
- gating_output = gating_output.float()
70
- if scoring_func == "softmax":
71
- scores = torch.softmax(gating_output, dim=-1)
72
- elif scoring_func == "sigmoid":
73
- scores = gating_output.sigmoid()
74
- else:
75
- raise ValueError(f"Unsupported scoring function: {scoring_func}")
76
-
77
- num_token = scores.shape[0]
78
- if e_score_correction_bias is not None:
79
- original_scores = scores
80
- scores = scores + e_score_correction_bias.unsqueeze(0)
81
- group_scores = (
82
- scores.view(num_token, num_expert_group, -1).topk(2, dim=-1)[0].sum(dim=-1)
83
- )
84
- else:
85
- group_scores = (
86
- scores.view(num_token, num_expert_group, -1).max(dim=-1).values
87
- ) # [n, n_group]
88
- group_idx = torch.topk(group_scores, k=topk_group, dim=-1, sorted=False)[
89
- 1
90
- ] # [n, top_k_group]
91
- group_mask = torch.zeros_like(group_scores) # [n, n_group]
92
- group_mask.scatter_(1, group_idx, 1) # [n, n_group]
93
- score_mask = (
94
- group_mask.unsqueeze(-1)
95
- .expand(num_token, num_expert_group, scores.shape[-1] // num_expert_group)
96
- .reshape(num_token, -1)
97
- ) # [n, e]
98
- tmp_scores = scores.masked_fill(~score_mask.bool(), float("-inf")) # [n, e]
99
-
100
- if e_score_correction_bias is not None:
101
- topk_ids = torch.topk(tmp_scores, k=topk, dim=-1, sorted=False)[1]
102
- topk_weights = original_scores.gather(1, topk_ids)
103
- else:
104
- topk_weights, topk_ids = torch.topk(tmp_scores, k=topk, dim=-1, sorted=False)
105
-
106
- if renormalize:
107
- topk_weights = topk_weights / topk_weights.sum(dim=-1, keepdim=True)
108
-
109
- if routed_scaling_factor != 1.0:
110
- topk_weights = topk_weights * routed_scaling_factor
111
- return topk_weights, topk_ids.to(torch.int32)
112
-
113
-
114
-def select_experts(
115
- hidden_states: torch.Tensor,
116
- router_logits: torch.Tensor,
117
- top_k: int,
118
- use_grouped_topk: bool,
119
- renormalize: bool,
120
- topk_group: int | None = None,
121
- num_expert_group: int | None = None,
122
- custom_routing_function: Callable | None = None,
123
- scoring_func: str = "softmax",
124
- routed_scaling_factor: float = 1.0,
125
- e_score_correction_bias: torch.Tensor | None = None,
126
-) -> tuple[torch.Tensor, torch.Tensor]:
127
- if use_grouped_topk:
128
- assert topk_group is not None
129
- assert num_expert_group is not None
130
- return grouped_topk(
131
- hidden_states=hidden_states,
132
- gating_output=router_logits,
133
- topk=top_k,
134
- renormalize=renormalize,
135
- num_expert_group=num_expert_group,
136
- topk_group=topk_group,
137
- scoring_func=scoring_func,
138
- routed_scaling_factor=routed_scaling_factor,
139
- e_score_correction_bias=e_score_correction_bias,
140
- )
141
- elif custom_routing_function is None:
142
- assert scoring_func == "softmax"
143
- topk_logit_vals, topk_idx = torch.topk(
144
- router_logits, k=top_k, dim=-1, sorted=False
145
- )
146
- if renormalize:
147
- topk_vals = torch.softmax(topk_logit_vals, dim=-1)
148
- else:
149
- logZ = torch.logsumexp(router_logits, dim=-1, keepdim=True)
150
- topk_vals = (topk_logit_vals - logZ).exp()
151
- return topk_vals.to(torch.float32), topk_idx.to(torch.int32)
152
- else:
153
- topk_weights, topk_ids = custom_routing_function(
154
- hidden_states=hidden_states,
155
- gating_output=router_logits,
156
- topk=top_k,
157
- renormalize=renormalize,
158
- )
159
- # cpu_fused_moe reads routing tensors as contiguous float32/int32
160
- # buffers and does not account for tensor strides.
161
- topk_weights = topk_weights.to(torch.float32).contiguous()
162
- topk_ids = topk_ids.to(torch.int32).contiguous()
163
- return topk_weights, topk_ids
164
-
165
-
53
# ===========================================================================
54
# Unquantized (BF16/FP16/FP32) MoE
55
# ===========================================================================
56
57
171
-class CPUUnquantizedExperts(mk.FusedMoEExpertsMonolithic):
58
+class CPUUnquantizedExperts(mk.FusedMoEExpertsModular):
59
"""Portable vector grouped-gemm unquantized MoE experts."""
60
61
isa = "vec"
62
output_alignment = 32
63
reduction_alignment = 1
64
65
+ def __init__(
66
+ self,
67
+ moe_config: FusedMoEConfig,
68
+ quant_config: FusedMoEQuantConfig,
69
+ ):
70
+ super().__init__(moe_config, quant_config)
71
+ self._use_zentorch = False
72
+
73
@classmethod
74
def _intermediate_alignment(cls) -> int:
75
return math.lcm(cls.output_alignment, cls.reduction_alignment)
@@ -210,20 +105,6 @@ class CPUUnquantizedExperts(mk.FusedMoEExpertsMonolithic):
105
)
106
return True, None
107
213
- def __init__(
214
- self,
215
- moe_config: FusedMoEConfig,
216
- quant_config: FusedMoEQuantConfig,
217
- ):
218
- super().__init__(moe_config, quant_config)
219
- # Router configuration that the monolithic apply() signature cannot
220
- # carry. Captured off the layer in process_weights_after_loading.
221
- self.use_grouped_topk = False
222
- self.renormalize = False
223
- self.scoring_func = "softmax"
224
- self.custom_routing_function: Callable | None = None
225
- self._use_zentorch = False
226
-
108
@property
109
def expects_unquantized_inputs(self) -> bool:
110
return True
@@ -286,8 +167,8 @@ class CPUUnquantizedExperts(mk.FusedMoEExpertsMonolithic):
167
weight_key: QuantKey | None,
168
activation_key: QuantKey | None,
169
) -> bool:
289
- # Routing runs in select_experts(), which covers every routing method
290
- # a layer can be configured with, including custom routing functions.
170
+ # Routing runs in the CPURouter, which covers every routing method a
171
+ # layer can be configured with, including custom routing functions.
172
return True
173
174
@staticmethod
@@ -298,10 +179,6 @@ class CPUUnquantizedExperts(mk.FusedMoEExpertsMonolithic):
179
return True
180
181
def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
301
- self.use_grouped_topk = layer.use_grouped_topk
302
- self.renormalize = layer.renormalize
303
- self.scoring_func = layer.scoring_func
304
- self.custom_routing_function = layer.custom_routing_function
182
self._use_zentorch = is_zentorch_moe_supported(layer)
183
if self._use_zentorch:
184
return
@@ -345,65 +222,44 @@ class CPUUnquantizedExperts(mk.FusedMoEExpertsMonolithic):
222
# this very object, which is what feeds self.w1_bias in apply().
223
layer.w13_bias.data = new_bias
224
348
- def _select_experts(
225
+ def workspace_shapes(
226
self,
350
- hidden_states: torch.Tensor,
351
- router_logits: torch.Tensor,
352
- num_expert_group: int | None,
353
- topk_group: int | None,
354
- e_score_correction_bias: torch.Tensor | None,
355
- routed_scaling_factor: float | None,
356
- ) -> tuple[torch.Tensor, torch.Tensor]:
357
- return select_experts(
358
- hidden_states=hidden_states,
359
- router_logits=router_logits,
360
- top_k=self.moe_config.experts_per_token,
361
- use_grouped_topk=self.use_grouped_topk,
362
- renormalize=self.renormalize,
363
- topk_group=topk_group,
364
- num_expert_group=num_expert_group,
365
- custom_routing_function=self.custom_routing_function,
366
- scoring_func=self.scoring_func,
367
- routed_scaling_factor=(
368
- routed_scaling_factor if routed_scaling_factor is not None else 1.0
369
- ),
370
- e_score_correction_bias=e_score_correction_bias,
371
- )
227
+ M: int,
228
+ N: int,
229
+ K: int,
230
+ topk: int,
231
+ global_num_experts: int,
232
+ local_num_experts: int,
233
+ expert_tokens_meta: mk.ExpertTokensMetadata | None,
234
+ activation: MoEActivation,
235
+ ) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...]]:
236
+ # cpu_fused_moe manages its own scratch space.
237
+ return (0,), (0,), (M, K)
238
+
239
+ def finalize_weight_and_reduce_impl(self) -> mk.TopKWeightAndReduce:
240
+ return TopKWeightAndReduceNoOP()
241
242
def apply(
243
self,
244
+ output: torch.Tensor,
245
hidden_states: torch.Tensor,
246
w1: torch.Tensor,
247
w2: torch.Tensor,
378
- router_logits: torch.Tensor,
248
+ topk_weights: torch.Tensor,
249
+ topk_ids: torch.Tensor,
250
activation: MoEActivation,
251
global_num_experts: int,
252
expert_map: torch.Tensor | None,
253
a1q_scale: torch.Tensor | None,
254
+ a2_scale: torch.Tensor | None,
255
+ workspace13: torch.Tensor,
256
+ workspace2: torch.Tensor,
257
+ expert_tokens_meta: mk.ExpertTokensMetadata | None,
258
apply_router_weight_on_input: bool,
384
- # grouped topk + fused topk bias parameters
385
- num_expert_group: int | None = None,
386
- e_score_correction_bias: torch.Tensor | None = None,
387
- routed_scaling_factor: float | None = None,
388
- topk_group: int | None = None,
389
- ) -> torch.Tensor:
390
- topk_weights, topk_ids = self._select_experts(
391
- hidden_states,
392
- router_logits,
393
- num_expert_group,
394
- topk_group,
395
- e_score_correction_bias,
396
- routed_scaling_factor,
397
- )
398
-
399
- if apply_router_weight_on_input:
400
- assert topk_ids.size(1) == 1, (
401
- "apply_router_weight_on_input is only implemented for topk=1"
402
- )
403
- hidden_states.mul_(topk_weights.to(hidden_states.dtype))
404
-
259
+ ) -> None:
260
+ # apply_router_weight_on_input (topk=1 only) is already applied to
261
+ # hidden_states by MoEPrepareAndFinalizeNoDPEPModular.prepare().
262
if self._use_zentorch:
406
- output = torch.empty_like(hidden_states)
263
torch.ops.zentorch.zentorch_fused_moe(
264
output,
265
hidden_states,
@@ -416,9 +272,9 @@ class CPUUnquantizedExperts(mk.FusedMoEExpertsMonolithic):
272
apply_router_weight_on_input,
273
str(activation.value).lower(),
274
)
419
- return output
420
-
421
- return cpu_fused_moe(
275
+ return
276
+ cpu_fused_moe(
277
+ output,
278
hidden_states,
279
w1,
280
w2,
@@ -460,8 +316,6 @@ class X86CPUUnquantizedExperts(CPUUnquantizedExperts):
316
)
317
if not supported:
318
return supported, reason
463
- if is_zentorch_moe_config_supported(moe_config):
464
- return True, None
319
if moe_config.in_dtype != torch.bfloat16:
320
return False, "kernel requires bfloat16 activations"
321
cpu_cls = cast(type[CPUUnquantizedExperts], cls)
@@ -496,8 +350,6 @@ class ArmCPUUnquantizedExperts(CPUUnquantizedExperts):
350
)
351
if not supported:
352
return supported, reason
499
- if is_zentorch_moe_config_supported(moe_config):
500
- return True, None
353
if moe_config.in_dtype != torch.bfloat16:
354
return False, "kernel requires bfloat16 activations"
355
cpu_cls = cast(type[CPUUnquantizedExperts], cls)
@@ -552,8 +404,8 @@ def prepare_fp8_moe_layer_for_cpu(
404
return packed_w13, packed_w2
405
406
555
-class CPUExpertsFp8(mk.FusedMoEExpertsMonolithic):
556
- """CPU FP8 W8A16 block-quantized monolithic MoE experts."""
407
+class CPUExpertsFp8(mk.FusedMoEExpertsModular):
408
+ """CPU FP8 W8A16 block-quantized modular MoE experts."""
409
410
def __init__(
411
self,
@@ -615,6 +467,7 @@ class CPUExpertsFp8(mk.FusedMoEExpertsMonolithic):
467
RoutingMethodType.Default,
468
RoutingMethodType.Renormalize,
469
RoutingMethodType.RenormalizeNaive,
470
+ RoutingMethodType.DeepseekV4,
471
]
472
473
@staticmethod
@@ -624,42 +477,49 @@ class CPUExpertsFp8(mk.FusedMoEExpertsMonolithic):
477
) -> bool:
478
return True
479
480
+ def workspace_shapes(
481
+ self,
482
+ M: int,
483
+ N: int,
484
+ K: int,
485
+ topk: int,
486
+ global_num_experts: int,
487
+ local_num_experts: int,
488
+ expert_tokens_meta: mk.ExpertTokensMetadata | None,
489
+ activation: MoEActivation,
490
+ ) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...]]:
491
+ # fused_experts_cpu manages its own scratch space.
492
+ return (0,), (0,), (M, K)
493
+
494
+ def finalize_weight_and_reduce_impl(self) -> mk.TopKWeightAndReduce:
495
+ return TopKWeightAndReduceNoOP()
496
+
497
def apply(
498
self,
499
+ output: torch.Tensor,
500
hidden_states: torch.Tensor,
501
w1: torch.Tensor,
502
w2: torch.Tensor,
632
- router_logits: torch.Tensor,
503
+ topk_weights: torch.Tensor,
504
+ topk_ids: torch.Tensor,
505
activation: MoEActivation,
506
global_num_experts: int,
507
expert_map: torch.Tensor | None,
508
a1q_scale: torch.Tensor | None,
509
+ a2_scale: torch.Tensor | None,
510
+ workspace13: torch.Tensor,
511
+ workspace2: torch.Tensor,
512
+ expert_tokens_meta: mk.ExpertTokensMetadata | None,
513
apply_router_weight_on_input: bool,
638
- # grouped topk + fused topk bias parameters
639
- num_expert_group: int | None = None,
640
- e_score_correction_bias: torch.Tensor | None = None,
641
- routed_scaling_factor: float | None = None,
642
- topk_group: int | None = None,
643
- ) -> torch.Tensor:
644
- topk_weights, topk_ids = select_experts(
645
- hidden_states=hidden_states,
646
- router_logits=router_logits,
647
- use_grouped_topk=num_expert_group is not None,
648
- top_k=self.moe_config.experts_per_token,
649
- renormalize=self.moe_config.routing_method
650
- in (
651
- RoutingMethodType.Renormalize,
652
- RoutingMethodType.RenormalizeNaive,
653
- ),
654
- topk_group=topk_group,
655
- num_expert_group=num_expert_group,
656
- scoring_func="softmax",
657
- routed_scaling_factor=(
658
- routed_scaling_factor if routed_scaling_factor is not None else 1.0
659
- ),
660
- e_score_correction_bias=e_score_correction_bias,
661
- )
662
-
514
+ ) -> None:
515
+ if apply_router_weight_on_input:
516
+ # fused_experts_cpu always applies topk_weights internally on
517
+ # combine; MoEPrepareAndFinalizeNoDPEPModular.prepare() would
518
+ # also pre-apply it to hidden_states, double-weighting the
519
+ # output. Not needed by any CPU FP8 model today.
520
+ raise NotImplementedError(
521
+ "CPUExpertsFp8 does not support apply_router_weight_on_input=True."
522
+ )
523
block_shape = (
524
list(self.quant_config.block_shape)
525
if self.quant_config.block_shape
@@ -670,13 +530,13 @@ class CPUExpertsFp8(mk.FusedMoEExpertsMonolithic):
530
)
531
)
532
673
- return fused_experts_cpu(
533
+ fused_experts_cpu(
534
+ output,
535
hidden_states,
536
w1,
537
w2,
538
topk_weights,
539
topk_ids,
679
- False, # inplace
540
CPUQuantMethod.FP8_W8A16, # moe_comp_method
541
self.w1_scale, # w1_scale
542
self.w2_scale, # w2_scale
@@ -710,8 +570,8 @@ def prepare_mxfp4_moe_layer_for_cpu(
570
return packed_w13, packed_w2, packed_w13_scale, packed_w2_scale
571
572
713
-class CPUExpertsMxfp4(mk.FusedMoEExpertsMonolithic):
714
- """CPU MXFP4 W4A16 monolithic MoE experts."""
573
+class CPUExpertsMxfp4(mk.FusedMoEExpertsModular):
574
+ """CPU MXFP4 W4A16 modular MoE experts."""
575
576
def __init__(
577
self,
@@ -773,6 +633,7 @@ class CPUExpertsMxfp4(mk.FusedMoEExpertsMonolithic):
633
RoutingMethodType.Default,
634
RoutingMethodType.Renormalize,
635
RoutingMethodType.RenormalizeNaive,
636
+ RoutingMethodType.DeepseekV4,
637
]
638
639
@staticmethod
@@ -782,55 +643,62 @@ class CPUExpertsMxfp4(mk.FusedMoEExpertsMonolithic):
643
) -> bool:
644
return True
645
646
+ def workspace_shapes(
647
+ self,
648
+ M: int,
649
+ N: int,
650
+ K: int,
651
+ topk: int,
652
+ global_num_experts: int,
653
+ local_num_experts: int,
654
+ expert_tokens_meta: mk.ExpertTokensMetadata | None,
655
+ activation: MoEActivation,
656
+ ) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...]]:
657
+ # fused_experts_cpu manages its own scratch space.
658
+ return (0,), (0,), (M, K)
659
+
660
+ def finalize_weight_and_reduce_impl(self) -> mk.TopKWeightAndReduce:
661
+ return TopKWeightAndReduceNoOP()
662
+
663
def apply(
664
self,
665
+ output: torch.Tensor,
666
hidden_states: torch.Tensor,
667
w1: torch.Tensor,
668
w2: torch.Tensor,
790
- router_logits: torch.Tensor,
669
+ topk_weights: torch.Tensor,
670
+ topk_ids: torch.Tensor,
671
activation: MoEActivation,
672
global_num_experts: int,
673
expert_map: torch.Tensor | None,
674
a1q_scale: torch.Tensor | None,
675
+ a2_scale: torch.Tensor | None,
676
+ workspace13: torch.Tensor,
677
+ workspace2: torch.Tensor,
678
+ expert_tokens_meta: mk.ExpertTokensMetadata | None,
679
apply_router_weight_on_input: bool,
796
- # grouped topk + fused topk bias parameters
797
- num_expert_group: int | None = None,
798
- e_score_correction_bias: torch.Tensor | None = None,
799
- routed_scaling_factor: float | None = None,
800
- topk_group: int | None = None,
801
- ) -> torch.Tensor:
802
- topk_weights, topk_ids = select_experts(
803
- hidden_states=hidden_states,
804
- router_logits=router_logits,
805
- use_grouped_topk=num_expert_group is not None,
806
- top_k=self.moe_config.experts_per_token,
807
- renormalize=self.moe_config.routing_method
808
- in (
809
- RoutingMethodType.Renormalize,
810
- RoutingMethodType.RenormalizeNaive,
811
- ),
812
- topk_group=topk_group,
813
- num_expert_group=num_expert_group,
814
- scoring_func="softmax",
815
- routed_scaling_factor=(
816
- routed_scaling_factor if routed_scaling_factor is not None else 1.0
817
- ),
818
- e_score_correction_bias=e_score_correction_bias,
819
- )
820
-
680
+ ) -> None:
681
+ if apply_router_weight_on_input:
682
+ # fused_experts_cpu always applies topk_weights internally on
683
+ # combine; MoEPrepareAndFinalizeNoDPEPModular.prepare() would
684
+ # also pre-apply it to hidden_states, double-weighting the
685
+ # output. Not needed by any CPU MXFP4 model today.
686
+ raise NotImplementedError(
687
+ "CPUExpertsMxfp4 does not support apply_router_weight_on_input=True."
688
+ )
689
# Get bias and swiglu params from quant config
690
w1_bias = self.quant_config.w1_bias
691
w2_bias = self.quant_config.w2_bias
692
alpha = getattr(self.quant_config, "gemm1_alpha", None)
693
limit = getattr(self.quant_config, "gemm1_clamp_limit", None)
694
827
- return fused_experts_cpu(
695
+ fused_experts_cpu(
696
+ output,
697
hidden_states,
698
w1,
699
w2,
700
topk_weights,
701
topk_ids,
833
- False, # inplace
702
CPUQuantMethod.MXFP4, # moe_comp_method
703
self.w1_scale, # w1_scale
704
self.w2_scale, # w2_scale
@@ -911,8 +779,8 @@ def prepare_int4_moe_layer_for_cpu(
779
return (blocked_w13, blocked_w2, blocked_s13, blocked_s2, blocked_z13, blocked_z2)
780
781
914
-class CPUExpertsInt4(mk.FusedMoEExpertsMonolithic):
915
- """CPU INT4 W4A16 group-quantized monolithic MoE experts.
782
+class CPUExpertsInt4(mk.FusedMoEExpertsModular):
783
+ """CPU INT4 W4A16 group-quantized modular MoE experts.
784
785
Weights are int4 (packed), activations are bf16/fp16.
786
Internally uses int8 compute via fused_experts_cpu with INT4_W4A8.
@@ -992,55 +860,75 @@ class CPUExpertsInt4(mk.FusedMoEExpertsMonolithic):
860
def supports_expert_map(self) -> bool:
861
return False
862
863
+ def moe_problem_size(
864
+ self,
865
+ a1: torch.Tensor,
866
+ w1: torch.Tensor,
867
+ w2: torch.Tensor,
868
+ topk_ids: torch.Tensor,
869
+ ) -> tuple[int, int, int, int, int]:
870
+ # convert_weight_packed_scale_zp blocks w1/w2 into AMX tiles
871
+ # (E, Nc, Kc, buffer_bytes), not the (E, N, K) layout the base
872
+ # implementation assumes -- N/K aren't recoverable from that shape,
873
+ # so read them from moe_config instead.
874
+ E = w1.shape[0]
875
+ K = a1.size(-1)
876
+ N = (
877
+ self.moe_config.intermediate_size_per_partition
878
+ * self.moe_config.w13_num_shards
879
+ )
880
+ M = a1.size(0) if a1.dim() == 2 else a1.size(1)
881
+ topk = topk_ids.size(1)
882
+ return E, M, N, K, topk
883
+
884
+ def workspace_shapes(
885
+ self,
886
+ M: int,
887
+ N: int,
888
+ K: int,
889
+ topk: int,
890
+ global_num_experts: int,
891
+ local_num_experts: int,
892
+ expert_tokens_meta: mk.ExpertTokensMetadata | None,
893
+ activation: MoEActivation,
894
+ ) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...]]:
895
+ # fused_experts_cpu manages its own scratch space.
896
+ return (0,), (0,), (M, K)
897
+
898
+ def finalize_weight_and_reduce_impl(self) -> mk.TopKWeightAndReduce:
899
+ return TopKWeightAndReduceNoOP()
900
+
901
def apply(
902
self,
903
+ output: torch.Tensor,
904
hidden_states: torch.Tensor,
905
w1: torch.Tensor,
906
w2: torch.Tensor,
1000
- router_logits: torch.Tensor,
907
+ topk_weights: torch.Tensor,
908
+ topk_ids: torch.Tensor,
909
activation: MoEActivation,
910
global_num_experts: int,
911
expert_map: torch.Tensor | None,
912
a1q_scale: torch.Tensor | None,
913
+ a2_scale: torch.Tensor | None,
914
+ workspace13: torch.Tensor,
915
+ workspace2: torch.Tensor,
916
+ expert_tokens_meta: mk.ExpertTokensMetadata | None,
917
apply_router_weight_on_input: bool,
1006
- # grouped topk + fused topk bias parameters
1007
- num_expert_group: int | None = None,
1008
- e_score_correction_bias: torch.Tensor | None = None,
1009
- routed_scaling_factor: float | None = None,
1010
- topk_group: int | None = None,
1011
- ) -> torch.Tensor:
918
+ ) -> None:
919
if apply_router_weight_on_input:
920
raise NotImplementedError(
921
"CPUExpertsInt4 (W4A16) does not support "
922
"apply_router_weight_on_input=True. "
923
)
924
1018
- topk_weights, topk_ids = select_experts(
1019
- hidden_states=hidden_states,
1020
- router_logits=router_logits,
1021
- use_grouped_topk=num_expert_group is not None,
1022
- top_k=self.moe_config.experts_per_token,
1023
- renormalize=self.moe_config.routing_method
1024
- in (
1025
- RoutingMethodType.Renormalize,
1026
- RoutingMethodType.RenormalizeNaive,
1027
- ),
1028
- topk_group=topk_group,
1029
- num_expert_group=num_expert_group,
1030
- scoring_func="softmax",
1031
- routed_scaling_factor=(
1032
- routed_scaling_factor if routed_scaling_factor is not None else 1.0
1033
- ),
1034
- e_score_correction_bias=e_score_correction_bias,
1035
- )
1036
-
1037
- return fused_experts_cpu(
925
+ fused_experts_cpu(
926
+ output,
927
hidden_states,
928
w1,
929
w2,
930
topk_weights,
931
topk_ids,
1043
- False, # inplace
932
CPUQuantMethod.INT4_W4A8,
933
self.w1_scale,
934
self.w2_scale,
@@ -1060,9 +948,9 @@ class CPUExpertsInt4(mk.FusedMoEExpertsMonolithic):
948
# ===========================================================================
949
950
1063
-class CPUExpertsInt8(mk.FusedMoEExpertsMonolithic):
951
+class CPUExpertsInt8(mk.FusedMoEExpertsModular):
952
"""CPU INT8 W8A8 per-channel weight / dynamic per-token activation
1065
- monolithic MoE experts."""
953
+ modular MoE experts."""
954
955
def __init__(
956
self,
@@ -1172,49 +1060,56 @@ class CPUExpertsInt8(mk.FusedMoEExpertsMonolithic):
1060
replace_parameter(layer, "w13_weight", w13)
1061
replace_parameter(layer, "w2_weight", w2)
1062
1063
+ def workspace_shapes(
1064
+ self,
1065
+ M: int,
1066
+ N: int,
1067
+ K: int,
1068
+ topk: int,
1069
+ global_num_experts: int,
1070
+ local_num_experts: int,
1071
+ expert_tokens_meta: mk.ExpertTokensMetadata | None,
1072
+ activation: MoEActivation,
1073
+ ) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...]]:
1074
+ # fused_experts_cpu manages its own scratch space.
1075
+ return (0,), (0,), (M, K)
1076
+
1077
+ def finalize_weight_and_reduce_impl(self) -> mk.TopKWeightAndReduce:
1078
+ return TopKWeightAndReduceNoOP()
1079
+
1080
def apply(
1081
self,
1082
+ output: torch.Tensor,
1083
hidden_states: torch.Tensor,
1084
w1: torch.Tensor,
1085
w2: torch.Tensor,
1180
- router_logits: torch.Tensor,
1086
+ topk_weights: torch.Tensor,
1087
+ topk_ids: torch.Tensor,
1088
activation: MoEActivation,
1089
global_num_experts: int,
1090
expert_map: torch.Tensor | None,
1091
a1q_scale: torch.Tensor | None,
1092
+ a2_scale: torch.Tensor | None,
1093
+ workspace13: torch.Tensor,
1094
+ workspace2: torch.Tensor,
1095
+ expert_tokens_meta: mk.ExpertTokensMetadata | None,
1096
apply_router_weight_on_input: bool,
1186
- # grouped topk + fused topk bias parameters
1187
- num_expert_group: int | None = None,
1188
- e_score_correction_bias: torch.Tensor | None = None,
1189
- routed_scaling_factor: float | None = None,
1190
- topk_group: int | None = None,
1191
- ) -> torch.Tensor:
1192
- topk_weights, topk_ids = select_experts(
1193
- hidden_states=hidden_states,
1194
- router_logits=router_logits,
1195
- use_grouped_topk=num_expert_group is not None,
1196
- top_k=self.moe_config.experts_per_token,
1197
- renormalize=self.moe_config.routing_method
1198
- in (
1199
- RoutingMethodType.Renormalize,
1200
- RoutingMethodType.RenormalizeNaive,
1201
- ),
1202
- topk_group=topk_group,
1203
- num_expert_group=num_expert_group,
1204
- scoring_func="softmax",
1205
- routed_scaling_factor=(
1206
- routed_scaling_factor if routed_scaling_factor is not None else 1.0
1207
- ),
1208
- e_score_correction_bias=e_score_correction_bias,
1209
- )
1210
-
1211
- return fused_experts_cpu(
1097
+ ) -> None:
1098
+ if apply_router_weight_on_input:
1099
+ # fused_experts_cpu always applies topk_weights internally on
1100
+ # combine; MoEPrepareAndFinalizeNoDPEPModular.prepare() would
1101
+ # also pre-apply it to hidden_states, double-weighting the
1102
+ # output. Not needed by any CPU INT8 W8A8 model today.
1103
+ raise NotImplementedError(
1104
+ "CPUExpertsInt8 does not support apply_router_weight_on_input=True."
1105
+ )
1106
+ fused_experts_cpu(
1107
+ output,
1108
hidden_states,
1109
w1,
1110
w2,
1111
topk_weights,
1112
topk_ids,
1217
- False, # inplace
1113
CPUQuantMethod.INT8_W8A8,
1114
self.w1_scale,
1115
self.w2_scale,
@@ -1229,7 +1124,7 @@ class CPUExpertsInt8(mk.FusedMoEExpertsMonolithic):
1124
)
1125
1126
1232
-class ArmCPUExpertsInt8(mk.FusedMoEExpertsMonolithic):
1127
+class ArmCPUExpertsInt8(mk.FusedMoEExpertsModular):
1128
"""Arm INT8 MoE with per-token activation and channelwise weight quantization."""
1129
1130
@property
@@ -1331,48 +1226,47 @@ class ArmCPUExpertsInt8(mk.FusedMoEExpertsMonolithic):
1226
replace_parameter(layer, "w13_weight", w13)
1227
replace_parameter(layer, "w2_weight", w2)
1228
1229
+ def workspace_shapes(
1230
+ self,
1231
+ M: int,
1232
+ N: int,
1233
+ K: int,
1234
+ topk: int,
1235
+ global_num_experts: int,
1236
+ local_num_experts: int,
1237
+ expert_tokens_meta: mk.ExpertTokensMetadata | None,
1238
+ activation: MoEActivation,
1239
+ ) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...]]:
1240
+ # cpu_fused_moe_int8 manages its own scratch space.
1241
+ return (0,), (0,), (M, K)
1242
+
1243
+ def finalize_weight_and_reduce_impl(self) -> mk.TopKWeightAndReduce:
1244
+ return TopKWeightAndReduceNoOP()
1245
+
1246
def apply(
1247
self,
1248
+ output: torch.Tensor,
1249
hidden_states: torch.Tensor,
1250
w1: torch.Tensor,
1251
w2: torch.Tensor,
1339
- router_logits: torch.Tensor,
1252
+ topk_weights: torch.Tensor,
1253
+ topk_ids: torch.Tensor,
1254
activation: MoEActivation,
1255
global_num_experts: int,
1256
expert_map: torch.Tensor | None,
1257
a1q_scale: torch.Tensor | None,
1258
+ a2_scale: torch.Tensor | None,
1259
+ workspace13: torch.Tensor,
1260
+ workspace2: torch.Tensor,
1261
+ expert_tokens_meta: mk.ExpertTokensMetadata | None,
1262
apply_router_weight_on_input: bool,
1345
- num_expert_group: int | None = None,
1346
- e_score_correction_bias: torch.Tensor | None = None,
1347
- routed_scaling_factor: float | None = None,
1348
- topk_group: int | None = None,
1349
- ) -> torch.Tensor:
1350
- topk_weights, topk_ids = select_experts(
1351
- hidden_states=hidden_states,
1352
- router_logits=router_logits,
1353
- use_grouped_topk=num_expert_group is not None,
1354
- top_k=self.moe_config.experts_per_token,
1355
- renormalize=self.moe_config.routing_method
1356
- in (
1357
- RoutingMethodType.Renormalize,
1358
- RoutingMethodType.RenormalizeNaive,
1359
- ),
1360
- topk_group=topk_group,
1361
- num_expert_group=num_expert_group,
1362
- scoring_func="softmax",
1363
- routed_scaling_factor=(
1364
- routed_scaling_factor if routed_scaling_factor is not None else 1.0
1365
- ),
1366
- e_score_correction_bias=e_score_correction_bias,
1367
- )
1368
-
1369
- if apply_router_weight_on_input:
1370
- assert topk_ids.size(1) == 1
1371
- hidden_states.mul_(topk_weights.to(hidden_states.dtype))
1372
-
1263
+ ) -> None:
1264
+ # apply_router_weight_on_input (topk=1 only) is already applied to
1265
+ # hidden_states by MoEPrepareAndFinalizeNoDPEPModular.prepare().
1266
assert self.w1_scale is not None
1267
assert self.w2_scale is not None
1375
- return cpu_fused_moe_int8(
1268
+ cpu_fused_moe_int8(
1269
+ output,
1270
hidden_states,
1271
w1,
1272
w2,
@@ -1388,7 +1282,7 @@ class ArmCPUExpertsInt8(mk.FusedMoEExpertsMonolithic):
1282
)
1283
1284
1391
-class ZenCPUExpertsInt8(mk.FusedMoEExpertsMonolithic):
1285
+class ZenCPUExpertsInt8(mk.FusedMoEExpertsModular):
1286
"""AMD Zen INT8 MoE with per-token activation and channelwise weight
1287
quantization, dispatched through zentorch."""
1288
@@ -1424,9 +1318,7 @@ class ZenCPUExpertsInt8(mk.FusedMoEExpertsMonolithic):
1318
if self.w2_bias is None
1319
else self.w2_bias.detach().to(torch.bfloat16).contiguous()
1320
)
1427
- logger.info_once(
1428
- "[zen_cpu] Using zentorch_fused_moe for W8A8 INT8 MoE (monolithic experts)"
1429
- )
1321
+ logger.info_once("[zen_cpu] Using zentorch_fused_moe for W8A8 INT8 MoE")
1322
1323
@property
1324
def expects_unquantized_inputs(self) -> bool:
@@ -1487,48 +1379,47 @@ class ZenCPUExpertsInt8(mk.FusedMoEExpertsMonolithic):
1379
def supports_expert_map(self) -> bool:
1380
return False
1381
1382
+ def workspace_shapes(
1383
+ self,
1384
+ M: int,
1385
+ N: int,
1386
+ K: int,
1387
+ topk: int,
1388
+ global_num_experts: int,
1389
+ local_num_experts: int,
1390
+ expert_tokens_meta: mk.ExpertTokensMetadata | None,
1391
+ activation: MoEActivation,
1392
+ ) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...]]:
1393
+ # zentorch_fused_moe manages its own scratch space.
1394
+ return (0,), (0,), (M, K)
1395
+
1396
+ def finalize_weight_and_reduce_impl(self) -> mk.TopKWeightAndReduce:
1397
+ return TopKWeightAndReduceNoOP()
1398
+
1399
def apply(
1400
self,
1401
+ output: torch.Tensor,
1402
hidden_states: torch.Tensor,
1403
w1: torch.Tensor,
1404
w2: torch.Tensor,
1495
- router_logits: torch.Tensor,
1405
+ topk_weights: torch.Tensor,
1406
+ topk_ids: torch.Tensor,
1407
activation: MoEActivation,
1408
global_num_experts: int,
1409
expert_map: torch.Tensor | None,
1410
a1q_scale: torch.Tensor | None,
1411
+ a2_scale: torch.Tensor | None,
1412
+ workspace13: torch.Tensor,
1413
+ workspace2: torch.Tensor,
1414
+ expert_tokens_meta: mk.ExpertTokensMetadata | None,
1415
apply_router_weight_on_input: bool,
1501
- num_expert_group: int | None = None,
1502
- e_score_correction_bias: torch.Tensor | None = None,
1503
- routed_scaling_factor: float | None = None,
1504
- topk_group: int | None = None,
1505
- ) -> torch.Tensor:
1506
- topk_weights, topk_ids = select_experts(
1507
- hidden_states=hidden_states,
1508
- router_logits=router_logits,
1509
- use_grouped_topk=num_expert_group is not None,
1510
- top_k=self.moe_config.experts_per_token,
1511
- renormalize=self.moe_config.routing_method
1512
- in (
1513
- RoutingMethodType.Renormalize,
1514
- RoutingMethodType.RenormalizeNaive,
1515
- ),
1516
- topk_group=topk_group,
1517
- num_expert_group=num_expert_group,
1518
- scoring_func="softmax",
1519
- routed_scaling_factor=(
1520
- routed_scaling_factor if routed_scaling_factor is not None else 1.0
1521
- ),
1522
- e_score_correction_bias=e_score_correction_bias,
1523
- )
1524
-
1416
+ ) -> None:
1417
if apply_router_weight_on_input:
1418
assert topk_ids.size(1) == 1, (
1419
"apply_router_weight_on_input is only implemented for topk=1"
1420
)
1421
hidden_states = hidden_states.mul(topk_weights.to(hidden_states.dtype))
1422
1531
- output = torch.empty_like(hidden_states)
1423
torch.ops.zentorch.zentorch_fused_moe(
1424
output,
1425
hidden_states,
@@ -1543,4 +1434,3 @@ class ZenCPUExpertsInt8(mk.FusedMoEExpertsMonolithic):
1434
self._w1_scale_bf16,
1435
self._w2_scale_bf16,
1436
)
1546
- return output
vllm/model_executor/layers/fused_moe/modular_kernel.py
+18
@@ -1175,6 +1175,24 @@ class FusedMoEKernelModularImpl:
1175
# time we need cache3, we're done with cache1.
1176
# Reuse workspace13 for the output since there is only one chunk.
1177
max_shape_size = max(prod(workspace13_shape), prod(fused_out_shape))
1178
+
1179
+ if current_platform.is_cpu():
1180
+ # Every CPU FusedMoEExpertsModular kernel reports zero-sized
1181
+ # workspace13/workspace2 (it manages its own scratch space
1182
+ # internally) and just needs the output buffer, so skip the
1183
+ # workspace manager entirely here -- it's the one part of this
1184
+ # call graph Dynamo can't trace (ContextVar-based lane lookup),
1185
+ # and CPU never actually needs its cross-chunk memory reuse.
1186
+ common_workspace = torch.empty(
1187
+ (max_shape_size,), dtype=workspace_dtype, device=device
1188
+ )
1189
+ workspace2 = torch.empty(
1190
+ workspace2_shape, dtype=workspace_dtype, device=device
1191
+ )
1192
+ workspace13 = _resize_cache(common_workspace, workspace13_shape)
1193
+ fused_out = _resize_cache(common_workspace, fused_out_shape)
1194
+ return workspace13, workspace2, fused_out
1195
+
1196
common_workspace, workspace2 = current_workspace_manager().get_simultaneous(
1197
((max_shape_size,), workspace_dtype),
1198
(workspace2_shape, workspace_dtype),
vllm/model_executor/layers/fused_moe/oracle/mxfp4.py
+30
-13
@@ -348,6 +348,7 @@ def _get_priority_backends_for_gpt_oss() -> list[Mxfp4MoeBackend]:
348
Mxfp4MoeBackend.MARLIN,
349
Mxfp4MoeBackend.BATCHED_MARLIN,
350
Mxfp4MoeBackend.XPU,
351
+ Mxfp4MoeBackend.CPU,
352
Mxfp4MoeBackend.EMULATION,
353
]
354
return _AVAILABLE_BACKENDS
@@ -366,6 +367,8 @@ def _get_priority_backends() -> list[Mxfp4MoeBackend]:
367
]
368
if current_platform.is_xpu():
369
return [Mxfp4MoeBackend.XPU]
370
+ if current_platform.is_cpu():
371
+ return [Mxfp4MoeBackend.CPU]
372
_AVAILABLE_BACKENDS = [
373
Mxfp4MoeBackend.FLASHINFER_TRTLLM_MXFP4_MXFP8,
374
Mxfp4MoeBackend.DEEPGEMM_MXFP4,
@@ -627,17 +630,6 @@ def select_mxfp4_moe_backend(
630
activation_format,
631
)
632
630
- if current_platform.is_cpu():
631
- backend = Mxfp4MoeBackend.CPU
632
- logger.info_once(_make_log_backend(backend))
633
- return _return_or_raise(
634
- Mxfp4MoeBackend.CPU,
635
- config,
636
- kMxfp4Static,
637
- None,
638
- activation_format,
639
- )
640
-
633
unsupported_log = "; ".join(
634
[
635
f"backend: {backend.value}, reason: {reason}"
@@ -1813,11 +1805,36 @@ def convert_weight_to_mxfp4_moe_kernel_format(
1805
w2_bias=w2_bias,
1806
_cache_permute_indices=_cache_permute_indices,
1807
)
1808
+ elif mxfp4_backend == Mxfp4MoeBackend.CPU:
1809
+ from vllm.model_executor.layers.fused_moe.experts.cpu_moe import (
1810
+ prepare_mxfp4_moe_layer_for_cpu,
1811
+ )
1812
+
1813
+ packed_w13, packed_w2, packed_w13_scale, packed_w2_scale = (
1814
+ prepare_mxfp4_moe_layer_for_cpu(
1815
+ w13_weight.data,
1816
+ w2_weight.data,
1817
+ w13_weight_scale.data,
1818
+ w2_weight_scale.data,
1819
+ )
1820
+ )
1821
+ if w13_bias is not None:
1822
+ w13_bias = w13_bias.data.to(torch.float32)
1823
+ if w2_bias is not None:
1824
+ w2_bias = w2_bias.data.to(torch.float32)
1825
+ return (
1826
+ packed_w13,
1827
+ packed_w2,
1828
+ packed_w13_scale,
1829
+ packed_w2_scale,
1830
+ w13_bias,
1831
+ w2_bias,
1832
+ )
1833
else:
1834
raise ValueError(
1835
f"Unsupported mxfp4_backend for Mxfp4MoEMethod: {mxfp4_backend}. "
1819
- "Expected TRTLLM, FlashInfer CUTLASS, Triton, AITER, XPU, or "
1820
- "emulation backend."
1836
+ "Expected TRTLLM, FlashInfer CUTLASS, Triton, AITER, XPU, "
1837
+ "CPU, or emulation backend."
1838
)
1839
1840
vllm/model_executor/layers/fused_moe/router/cpu_router.py
new
+294
@@ -0,0 +1,294 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""Router for CPU MoE experts.
4
+
5
+This is the routing logic that used to live inline in each CPU
6
+FusedMoEExpertsMonolithic.apply() (see git history of
7
+experts/cpu_moe.py); factoring it out into a proper FusedMoERouter is what
8
+lets CPU MoE experts be FusedMoEExpertsModular like every other backend.
9
+The math below is unchanged from that prior inline version -- this is a
10
+relocation, not a rewrite.
11
+"""
12
+
13
+from collections.abc import Callable
14
+
15
+import torch
16
+
17
+from vllm._custom_ops import biased_topk_cpu, hash_topk_cpu
18
+from vllm.distributed.eplb.eplb_state import EplbLayerState
19
+from vllm.model_executor.layers.fused_moe.config import (
20
+ RoutingMethodType,
21
+ get_routing_method_type,
22
+)
23
+from vllm.model_executor.layers.fused_moe.router.base_router import BaseRouter
24
+
25
+
26
+def _grouped_topk(
27
+ gating_output: torch.Tensor,
28
+ topk: int,
29
+ renormalize: bool,
30
+ num_expert_group: int,
31
+ topk_group: int,
32
+ scoring_func: str,
33
+ routed_scaling_factor: float,
34
+ e_score_correction_bias: torch.Tensor | None,
35
+) -> tuple[torch.Tensor, torch.Tensor]:
36
+ gating_output = gating_output.float()
37
+ if scoring_func == "softmax":
38
+ scores = torch.softmax(gating_output, dim=-1)
39
+ elif scoring_func == "sigmoid":
40
+ scores = gating_output.sigmoid()
41
+ else:
42
+ raise ValueError(f"Unsupported scoring function: {scoring_func}")
43
+
44
+ num_token = scores.shape[0]
45
+ if e_score_correction_bias is not None:
46
+ original_scores = scores
47
+ scores = scores + e_score_correction_bias.unsqueeze(0)
48
+ group_scores = (
49
+ scores.view(num_token, num_expert_group, -1).topk(2, dim=-1)[0].sum(dim=-1)
50
+ )
51
+ else:
52
+ group_scores = (
53
+ scores.view(num_token, num_expert_group, -1).max(dim=-1).values
54
+ ) # [n, n_group]
55
+ group_idx = torch.topk(group_scores, k=topk_group, dim=-1, sorted=False)[
56
+ 1
57
+ ] # [n, top_k_group]
58
+ group_mask = torch.zeros_like(group_scores) # [n, n_group]
59
+ group_mask.scatter_(1, group_idx, 1) # [n, n_group]
60
+ score_mask = (
61
+ group_mask.unsqueeze(-1)
62
+ .expand(num_token, num_expert_group, scores.shape[-1] // num_expert_group)
63
+ .reshape(num_token, -1)
64
+ ) # [n, e]
65
+ tmp_scores = scores.masked_fill(~score_mask.bool(), float("-inf")) # [n, e]
66
+
67
+ if e_score_correction_bias is not None:
68
+ topk_ids = torch.topk(tmp_scores, k=topk, dim=-1, sorted=False)[1]
69
+ topk_weights = original_scores.gather(1, topk_ids)
70
+ else:
71
+ topk_weights, topk_ids = torch.topk(tmp_scores, k=topk, dim=-1, sorted=False)
72
+
73
+ if renormalize:
74
+ topk_weights = topk_weights / topk_weights.sum(dim=-1, keepdim=True)
75
+
76
+ if routed_scaling_factor != 1.0:
77
+ topk_weights = topk_weights * routed_scaling_factor
78
+ return topk_weights, topk_ids.to(torch.int32)
79
+
80
+
81
+def _sqrtsoftplus_bias_topk(
82
+ router_logits: torch.Tensor,
83
+ top_k: int,
84
+ renormalize: bool,
85
+ e_score_correction_bias: torch.Tensor | None,
86
+ routed_scaling_factor: float,
87
+ input_ids: torch.Tensor | None,
88
+ hash_indices_table: torch.Tensor | None,
89
+) -> tuple[torch.Tensor, torch.Tensor]:
90
+ """DeepSeek V4's routing: weight = sqrt(softplus(logit)); expert
91
+ selection uses a bias-corrected score (weight + correction_bias), but
92
+ the returned weight for each selected expert is the *unbiased* value
93
+ (the "noaux_tc" pattern -- bias steers selection only). Always
94
+ renormalizes selected weights to sum to ``routed_scaling_factor``,
95
+ matching the reference Triton kernel
96
+ (vllm/model_executor/layers/fused_moe/router/dsv4_topk.py).
97
+
98
+ When ``hash_indices_table`` is given (hash-routed layers), expert
99
+ selection instead comes from ``hash_indices_table[input_ids]`` --
100
+ the bias-corrected-score topk is skipped entirely, matching the CUDA
101
+ reference kernel (``dsv4HashTopkSoftplusSqrt``).
102
+
103
+ Backed by the ported SGLang AVX512 kernels `biased_topk_cpu` (the
104
+ correction-bias path) and `hash_topk_cpu` (the hash-routed-layer path);
105
+ see csrc/cpu/sgl-kernels/topk.cpp.
106
+ """
107
+ assert renormalize, "DeepseekV4 routing always renormalizes"
108
+ gating_output = router_logits.contiguous()
109
+ num_experts = gating_output.shape[1]
110
+ if hash_indices_table is not None:
111
+ assert input_ids is not None
112
+ # hash_indices_table is int32 by construction (see cpu/model.py) and
113
+ # advanced indexing always allocates a fresh contiguous tensor.
114
+ tid2eid = hash_indices_table[input_ids]
115
+ return hash_topk_cpu(
116
+ gating_output,
117
+ tid2eid,
118
+ top_k,
119
+ "sqrtsoftplus",
120
+ 0, # num_fused_shared_experts: never fused into routing on CPU
121
+ num_experts,
122
+ routed_scaling_factor,
123
+ )
124
+ assert e_score_correction_bias is not None
125
+ return biased_topk_cpu(
126
+ gating_output, # only used for num_tokens/device (dtype forced fp32)
127
+ gating_output,
128
+ e_score_correction_bias.contiguous(),
129
+ top_k,
130
+ renormalize,
131
+ "sqrtsoftplus",
132
+ 0, # num_fused_shared_experts: never fused into routing on CPU
133
+ routed_scaling_factor,
134
+ True, # apply_routed_scaling_factor_on_output
135
+ )
136
+
137
+
138
+def _softmax_topk(
139
+ router_logits: torch.Tensor,
140
+ top_k: int,
141
+ renormalize: bool,
142
+) -> tuple[torch.Tensor, torch.Tensor]:
143
+ topk_logit_vals, topk_idx = torch.topk(router_logits, k=top_k, dim=-1, sorted=False)
144
+ if renormalize:
145
+ topk_vals = torch.softmax(topk_logit_vals, dim=-1)
146
+ else:
147
+ logZ = torch.logsumexp(router_logits, dim=-1, keepdim=True)
148
+ topk_vals = (topk_logit_vals - logZ).exp()
149
+ return topk_vals.to(torch.float32), topk_idx.to(torch.int32)
150
+
151
+
152
+class CPURouter(BaseRouter):
153
+ """Router covering every routing scheme CPU MoE experts support: plain
154
+ softmax top-k, grouped top-k (with an optional correction bias), an
155
+ arbitrary custom_routing_function, and DeepSeek V4's sqrtsoftplus scheme
156
+ (bias-corrected or hash-routed via the ported biased_topk_cpu/
157
+ hash_topk_cpu kernels). All of the above except sqrtsoftplus are plain
158
+ torch ops so Dynamo traces straight through into them like it did when
159
+ this logic lived inline in FusedMoEExpertsMonolithic.apply() -- no new
160
+ compile boundary is introduced here.
161
+ """
162
+
163
+ def __init__(
164
+ self,
165
+ top_k: int,
166
+ global_num_experts: int,
167
+ use_grouped_topk: bool = False,
168
+ num_expert_group: int | None = None,
169
+ topk_group: int | None = None,
170
+ renormalize: bool = True,
171
+ scoring_func: str = "softmax",
172
+ routed_scaling_factor: float = 1.0,
173
+ e_score_correction_bias: torch.Tensor | None = None,
174
+ custom_routing_function: Callable | None = None,
175
+ hash_indices_table: torch.Tensor | None = None,
176
+ eplb_state: EplbLayerState | None = None,
177
+ ):
178
+ super().__init__(
179
+ top_k=top_k,
180
+ global_num_experts=global_num_experts,
181
+ eplb_state=eplb_state,
182
+ )
183
+ self.use_grouped_topk = use_grouped_topk
184
+ self.num_expert_group = num_expert_group
185
+ self.topk_group = topk_group
186
+ self.renormalize = renormalize
187
+ self.scoring_func = scoring_func
188
+ self.routed_scaling_factor = routed_scaling_factor
189
+ self.e_score_correction_bias = e_score_correction_bias
190
+ self.custom_routing_function = custom_routing_function
191
+ self._hash_indices_table = hash_indices_table
192
+
193
+ @property
194
+ def routing_method_type(self) -> RoutingMethodType:
195
+ return get_routing_method_type(
196
+ scoring_func=self.scoring_func,
197
+ top_k=self.top_k,
198
+ renormalize=self.renormalize,
199
+ num_expert_group=self.num_expert_group,
200
+ has_e_score_bias=self.e_score_correction_bias is not None,
201
+ routed_scaling_factor=self.routed_scaling_factor,
202
+ )
203
+
204
+ def _compute_routing(
205
+ self,
206
+ hidden_states: torch.Tensor,
207
+ router_logits: torch.Tensor,
208
+ indices_type: torch.dtype | None,
209
+ *,
210
+ input_ids: torch.Tensor | None = None,
211
+ ) -> tuple[torch.Tensor, torch.Tensor]:
212
+ if self.use_grouped_topk:
213
+ assert self.topk_group is not None
214
+ assert self.num_expert_group is not None
215
+ return _grouped_topk(
216
+ gating_output=router_logits,
217
+ topk=self.top_k,
218
+ renormalize=self.renormalize,
219
+ num_expert_group=self.num_expert_group,
220
+ topk_group=self.topk_group,
221
+ scoring_func=self.scoring_func,
222
+ routed_scaling_factor=self.routed_scaling_factor,
223
+ e_score_correction_bias=self.e_score_correction_bias,
224
+ )
225
+ elif self.scoring_func == "sqrtsoftplus":
226
+ assert (
227
+ self._hash_indices_table is not None
228
+ or self.e_score_correction_bias is not None
229
+ )
230
+ return _sqrtsoftplus_bias_topk(
231
+ router_logits=router_logits,
232
+ top_k=self.top_k,
233
+ renormalize=self.renormalize,
234
+ e_score_correction_bias=self.e_score_correction_bias,
235
+ routed_scaling_factor=self.routed_scaling_factor,
236
+ input_ids=input_ids,
237
+ hash_indices_table=self._hash_indices_table,
238
+ )
239
+ elif self.custom_routing_function is None:
240
+ assert self.scoring_func == "softmax"
241
+ return _softmax_topk(router_logits, self.top_k, self.renormalize)
242
+ else:
243
+ topk_weights, topk_ids = self.custom_routing_function(
244
+ hidden_states=hidden_states,
245
+ gating_output=router_logits,
246
+ topk=self.top_k,
247
+ renormalize=self.renormalize,
248
+ )
249
+ # cpu_fused_moe reads routing tensors as contiguous float32/int32
250
+ # buffers and does not account for tensor strides.
251
+ topk_weights = topk_weights.to(torch.float32).contiguous()
252
+ topk_ids = topk_ids.to(torch.int32).contiguous()
253
+ return topk_weights, topk_ids
254
+
255
+
256
+def select_experts(
257
+ hidden_states: torch.Tensor,
258
+ router_logits: torch.Tensor,
259
+ top_k: int,
260
+ use_grouped_topk: bool,
261
+ renormalize: bool,
262
+ topk_group: int | None = None,
263
+ num_expert_group: int | None = None,
264
+ custom_routing_function: Callable | None = None,
265
+ scoring_func: str = "softmax",
266
+ routed_scaling_factor: float = 1.0,
267
+ e_score_correction_bias: torch.Tensor | None = None,
268
+ input_ids: torch.Tensor | None = None,
269
+ hash_indices_table: torch.Tensor | None = None,
270
+) -> tuple[torch.Tensor, torch.Tensor]:
271
+ """Routing helper for the CPU MoE experts that still use the monolithic
272
+ contract and so must compute routing themselves from ``router_logits``
273
+ (the zentorch-accelerated path and the ARM W4A8 dynamic-quant kernel,
274
+ kept monolithic since zentorch/Arm hardware isn't available to validate
275
+ a modular migration here). Modular CPU experts get routing from
276
+ ``CPURouter`` instead; this is a thin wrapper around the same logic for
277
+ the two classes that can't use it.
278
+ """
279
+ router = CPURouter(
280
+ top_k=top_k,
281
+ global_num_experts=router_logits.shape[-1],
282
+ use_grouped_topk=use_grouped_topk,
283
+ num_expert_group=num_expert_group,
284
+ topk_group=topk_group,
285
+ renormalize=renormalize,
286
+ scoring_func=scoring_func,
287
+ routed_scaling_factor=routed_scaling_factor,
288
+ e_score_correction_bias=e_score_correction_bias,
289
+ custom_routing_function=custom_routing_function,
290
+ hash_indices_table=hash_indices_table,
291
+ )
292
+ return router._compute_routing(
293
+ hidden_states, router_logits, None, input_ids=input_ids
294
+ )
vllm/model_executor/layers/fused_moe/router/router_factory.py
+22
@@ -14,6 +14,7 @@ from vllm.model_executor.layers.fused_moe.config import (
14
from vllm.model_executor.layers.fused_moe.router.aiter_shared_routed_fused_moe_router import ( # noqa: E501
15
AiterSharedRoutedFusedMoERouter,
16
)
17
+from vllm.model_executor.layers.fused_moe.router.cpu_router import CPURouter
18
from vllm.model_executor.layers.fused_moe.router.custom_routing_router import (
19
CustomRoutingRouter,
20
)
@@ -35,6 +36,7 @@ from vllm.model_executor.layers.fused_moe.router.routing_simulator_router import
36
from vllm.model_executor.layers.fused_moe.router.zero_expert_router import (
37
ZeroExpertRouter,
38
)
39
+from vllm.platforms import current_platform
40
41
42
def create_fused_moe_router(
@@ -141,6 +143,26 @@ def create_fused_moe_router(
143
routed_scaling_factor=routed_scaling_factor,
144
)
145
146
+ if current_platform.is_cpu():
147
+ # CPURouter covers every routing scheme CPU MoE experts support
148
+ # (plain softmax, grouped topk, custom routing functions, and
149
+ # DeepSeek V4's sqrtsoftplus/hash routing) in a single class, so it
150
+ # takes priority over the scheme-specific routers below.
151
+ return CPURouter(
152
+ top_k=top_k,
153
+ global_num_experts=global_num_experts,
154
+ use_grouped_topk=use_grouped_topk,
155
+ num_expert_group=num_expert_group,
156
+ topk_group=topk_group,
157
+ renormalize=renormalize,
158
+ scoring_func=scoring_func,
159
+ routed_scaling_factor=routed_scaling_factor,
160
+ e_score_correction_bias=e_score_correction_bias,
161
+ custom_routing_function=custom_routing_function,
162
+ hash_indices_table=hash_indices_table,
163
+ eplb_state=eplb_state,
164
+ )
165
+
166
if use_grouped_topk:
167
assert custom_routing_function is None
168
if num_expert_group is None or topk_group is None:
vllm/model_executor/layers/fused_moe/runner/moe_runner.py
+9
-2
@@ -317,10 +317,17 @@ class MoERunner(MoERunnerInterface):
317
return self.routed_experts.load_weights(weights)
318
319
def _select_forward(self) -> Callable:
320
- if current_platform.is_tpu() or current_platform.is_cpu():
320
+ if current_platform.is_tpu():
321
# TODO: Once the OOM issue for the TPU backend is resolved, we
322
# will switch to using the moe_forward custom op.
323
- # Note: CPU doesn't require wrapped _forward_impl.
323
+ return _moe_forward if self._shared_experts is None else _moe_forward_shared
324
+
325
+ if current_platform.is_cpu():
326
+ # CPU never touches the workspace manager (Monolithic experts
327
+ # skip it entirely; Modular experts' _allocate_buffers bypasses
328
+ # it too, see modular_kernel.py) -- the ContextVar-based lane
329
+ # lookup was the only part of this call graph Dynamo can't
330
+ # trace, so CPU can always call the fused-MoE op directly.
331
return _moe_forward if self._shared_experts is None else _moe_forward_shared
332
333
return (
vllm/model_executor/layers/mhc.py
+88
@@ -261,6 +261,36 @@ class MHCPreOp(CustomOp):
261
sinkhorn_repeat,
262
)
263
264
+ def forward_cpu(
265
+ self,
266
+ residual: torch.Tensor,
267
+ fn: torch.Tensor,
268
+ hc_scale: torch.Tensor,
269
+ hc_base: torch.Tensor,
270
+ rms_eps: float,
271
+ hc_pre_eps: float,
272
+ hc_sinkhorn_eps: float,
273
+ hc_post_mult_value: float,
274
+ sinkhorn_repeat: int,
275
+ n_splits: int = 1,
276
+ norm_weight: torch.Tensor | None = None,
277
+ norm_eps: float = 0.0,
278
+ ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
279
+ return mhc_kernels.mhc_pre_cpu(
280
+ residual,
281
+ fn,
282
+ hc_scale,
283
+ hc_base,
284
+ rms_eps,
285
+ hc_pre_eps,
286
+ hc_sinkhorn_eps,
287
+ hc_post_mult_value,
288
+ sinkhorn_repeat,
289
+ n_splits,
290
+ norm_weight,
291
+ norm_eps,
292
+ )
293
+
294
295
# --8<-- [start:mhc_post]
296
@CustomOp.register("mhc_post")
@@ -337,6 +367,15 @@ class MHCPostOp(CustomOp):
367
comb_res_mix,
368
)
369
370
+ def forward_cpu(
371
+ self,
372
+ x: torch.Tensor,
373
+ residual: torch.Tensor,
374
+ post_layer_mix: torch.Tensor,
375
+ comb_res_mix: torch.Tensor,
376
+ ) -> torch.Tensor:
377
+ return mhc_kernels.mhc_post_cpu(x, residual, post_layer_mix, comb_res_mix)
378
+
379
380
# --8<-- [start:hc_head]
381
@CustomOp.register("hc_head")
@@ -444,6 +483,19 @@ class HCHeadOp(CustomOp):
483
)
484
return out.view(*outer_shape, hidden_size)
485
486
+ def forward_cpu(
487
+ self,
488
+ hidden_states: torch.Tensor,
489
+ hc_fn: torch.Tensor,
490
+ hc_scale: torch.Tensor,
491
+ hc_base: torch.Tensor,
492
+ rms_norm_eps: float,
493
+ hc_eps: float,
494
+ ) -> torch.Tensor:
495
+ return mhc_kernels.hc_head_fused_cpu(
496
+ hidden_states, hc_fn, hc_scale, hc_base, rms_norm_eps, hc_eps
497
+ )
498
+
499
500
# --8<-- [start:mhc_fused_post_pre]
501
@CustomOp.register("mhc_fused_post_pre")
@@ -654,6 +706,42 @@ class MHCFusedPostPreOp(CustomOp):
706
sinkhorn_repeat,
707
)
708
709
+ def forward_cpu(
710
+ self,
711
+ x: torch.Tensor,
712
+ residual: torch.Tensor,
713
+ post_layer_mix: torch.Tensor,
714
+ comb_res_mix: torch.Tensor,
715
+ fn: torch.Tensor,
716
+ hc_scale: torch.Tensor,
717
+ hc_base: torch.Tensor,
718
+ rms_eps: float,
719
+ hc_pre_eps: float,
720
+ hc_sinkhorn_eps: float,
721
+ hc_post_mult_value: float,
722
+ sinkhorn_repeat: int,
723
+ n_splits: int = 1,
724
+ tile_n: int = 1,
725
+ norm_weight: torch.Tensor | None = None,
726
+ norm_eps: float = 0.0,
727
+ ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
728
+ # Decompose into post + pre (no fused kernel available).
729
+ residual_cur = mhc_kernels.mhc_post_cpu(
730
+ x, residual, post_layer_mix, comb_res_mix
731
+ )
732
+ post_mix_cur, comb_mix_cur, layer_input_cur = mhc_kernels.mhc_pre_cpu(
733
+ residual_cur,
734
+ fn,
735
+ hc_scale,
736
+ hc_base,
737
+ rms_eps,
738
+ hc_pre_eps,
739
+ hc_sinkhorn_eps,
740
+ hc_post_mult_value,
741
+ sinkhorn_repeat,
742
+ )
743
+ return residual_cur, post_mix_cur, comb_mix_cur, layer_input_cur
744
+
745
746
def hc_expand(x: torch.Tensor, n: int) -> torch.Tensor:
747
"""[s, hidden_size] -> [s, n * hidden_size] by replication."""
vllm/model_executor/layers/sparse_attn_indexer.py
+196
-1
@@ -828,10 +828,12 @@ class SparseAttnIndexer(CustomOp):
828
return self.forward_cuda(hidden_states, q_quant, k, weights)
829
elif current_platform.is_rocm():
830
return self.forward_hip(hidden_states, q_quant, k, weights)
831
+ elif current_platform.is_cpu():
832
+ return self.forward_cpu(hidden_states, q_quant, k, weights)
833
else:
834
raise NotImplementedError(
835
"SparseAttnIndexer native forward is only implemented for "
834
- "CUDA, ROCm and XPU platforms."
836
+ "CUDA, ROCm, XPU and CPU platforms."
837
)
838
839
def forward_cuda(
@@ -919,3 +921,196 @@ class SparseAttnIndexer(CustomOp):
921
"Sparse attention indexer ROCm path is only supported on AITER. "
922
"Please enable aiter with VLLM_ROCM_USE_AITER=1"
923
)
924
+
925
+ def forward_cpu(
926
+ self,
927
+ hidden_states: torch.Tensor,
928
+ q_quant: torch.Tensor | tuple[torch.Tensor, torch.Tensor],
929
+ k: torch.Tensor | None,
930
+ weights: torch.Tensor,
931
+ ):
932
+ """CPU sparse attention indexer: cache write stays eager Python glue
933
+ (own K-cache layout, not shared with the main attention cache
934
+ write). PREFILL and DECODE both call the ported
935
+ ``fp8_paged_mqa_logits_cpu``/``topk_transform_512_cpu`` kernels,
936
+ which read the paged K-cache directly via ``page_table`` -- no
937
+ eager gather step, no per-request Python loop.
938
+
939
+ ``prefill_metadata.chunks`` always has exactly one entry here:
940
+ ``DeepseekV4CPUIndexerMetadataBuilder`` overrides the base chunk
941
+ split to always return the whole step's prefill batch as one
942
+ chunk, since the base chunking only bounds CUDA/XPU's dense M*N
943
+ logits tensor and flat K-gather workspace, neither of which this
944
+ paged kernel allocates.
945
+ """
946
+ assert not self.use_fp4_cache, (
947
+ "CPU sparse indexer doesn't support fp4 cache yet"
948
+ )
949
+ assert isinstance(q_quant, torch.Tensor), (
950
+ "CPU sparse_attn_indexer expects a single FP8 q_quant tensor"
951
+ )
952
+ assert self.dcp_world_size <= 1 and not self.use_pcp, (
953
+ "CPU sparse indexer doesn't support decode/prefill context parallelism yet."
954
+ )
955
+
956
+ forward_context = get_forward_context()
957
+ attn_metadata = forward_context.attn_metadata
958
+ attn_metadata_narrowed: DeepseekV32IndexerMetadata | None = None
959
+ if isinstance(attn_metadata, dict):
960
+ metadata = attn_metadata[self.k_cache.prefix]
961
+ assert isinstance(metadata, DeepseekV32IndexerMetadata)
962
+ attn_metadata_narrowed = metadata
963
+ if attn_metadata_narrowed is None:
964
+ # Profiling/dummy run: no real metadata to act on.
965
+ return self.topk_indices_buffer
966
+
967
+ kv_cache = self.k_cache.kv_cache
968
+ topk_tokens = self.topk_tokens
969
+ topk_indices_buffer = self.topk_indices_buffer
970
+ slot_mapping = attn_metadata_narrowed.slot_mapping
971
+ has_decode = attn_metadata_narrowed.num_decodes > 0
972
+ has_prefill = attn_metadata_narrowed.num_prefills > 0
973
+ num_decode_tokens = attn_metadata_narrowed.num_decode_tokens
974
+
975
+ num_tokens = slot_mapping.shape[0]
976
+ if k is not None:
977
+ k = k[:num_tokens]
978
+
979
+ if not self.skip_k_cache_insert:
980
+ # Only reachable via DeepseekV32Attention with
981
+ # prefill_context_parallel_size > 1 on CPU -- set_k_cpu/set_s_cpu
982
+ # (the kernels this used to call) have been removed as unused/
983
+ # untested (csrc/cpu/sgl-kernels/store_cache.cpp).
984
+ raise NotImplementedError(
985
+ "SparseAttnIndexer.forward_cpu: skip_k_cache_insert=False "
986
+ "(prefill context parallel on CPU) is not supported."
987
+ )
988
+
989
+ topk_indices_buffer[: hidden_states.shape[0]] = -1
990
+
991
+ if has_prefill:
992
+ assert topk_tokens == 512, (
993
+ "topk_transform_512_cpu only supports index_topk == 512."
994
+ )
995
+ prefill_metadata = attn_metadata_narrowed.prefill
996
+ assert prefill_metadata is not None
997
+ assert len(prefill_metadata.chunks) == 1, (
998
+ "forward_cpu expects the prefill metadata builder to always "
999
+ "produce a single chunk -- see "
1000
+ "DeepseekV4CPUIndexerMetadataBuilder._split_indexer_prefill_chunks."
1001
+ )
1002
+ chunk = prefill_metadata.chunks[0]
1003
+ # kv_cache is a per-layer view into vLLM's shared multi-layer
1004
+ # cache allocation, so its block stride generally exceeds
1005
+ # block_size * page_width; fp8_paged_mqa_logits_cpu reads
1006
+ # kv_view.stride(0) explicitly, so no copy is needed here.
1007
+ kv_view = kv_cache.view(kv_cache.shape[0], -1)
1008
+ block_size = kv_cache.shape[1]
1009
+ q_slice = q_quant[chunk.token_start : chunk.token_end]
1010
+ topk_indices = topk_indices_buffer[
1011
+ chunk.token_start : chunk.token_end, :topk_tokens
1012
+ ]
1013
+ if chunk.local_total_seq_lens == 0:
1014
+ topk_indices.fill_(-1)
1015
+ else:
1016
+ assert chunk.local_cu_seq_lens is not None
1017
+ # Each token's own (per-request, DCP-local) causal length.
1018
+ local_seq_lens = chunk.cu_seqlen_ke - chunk.cu_seqlen_ks
1019
+ # Recover each token's owning request from its row-start
1020
+ # tag. Ties (from zero-length requests) are harmless: those
1021
+ # tokens have local length 0 and never dereference
1022
+ # page_table.
1023
+ req_idx = (
1024
+ torch.searchsorted(
1025
+ chunk.local_cu_seq_lens, chunk.cu_seqlen_ks, right=True
1026
+ )
1027
+ - 1
1028
+ )
1029
+ page_table = chunk.block_table[req_idx]
1030
+ # The true max over this chunk's own rows, NOT
1031
+ # chunk.max_local_total_seq_lens (that field sums every
1032
+ # request's length in the chunk, bounding the old
1033
+ # flat-gather buffer this paged path no longer allocates).
1034
+ max_seq_len = int(local_seq_lens.max().item())
1035
+
1036
+ logits = ops.fp8_paged_mqa_logits_cpu(
1037
+ q_slice,
1038
+ kv_view,
1039
+ weights[chunk.token_start : chunk.token_end],
1040
+ local_seq_lens,
1041
+ page_table,
1042
+ block_size,
1043
+ max_seq_len,
1044
+ )
1045
+ out_page_scratch = torch.empty(
1046
+ (q_slice.shape[0], topk_tokens),
1047
+ dtype=torch.int32,
1048
+ device=kv_cache.device,
1049
+ )
1050
+ ops.topk_transform_512_cpu(
1051
+ logits,
1052
+ local_seq_lens,
1053
+ page_table,
1054
+ out_page_scratch,
1055
+ block_size,
1056
+ topk_indices,
1057
+ )
1058
+
1059
+ if has_decode:
1060
+ decode_metadata = attn_metadata_narrowed.decode
1061
+ assert decode_metadata is not None
1062
+ assert not decode_metadata.requires_padding, (
1063
+ "CPU sparse indexer decode path does not support speculative "
1064
+ "decoding (native MTP) yet."
1065
+ )
1066
+ batch_size = decode_metadata.decode_lens.shape[0]
1067
+ if batch_size > 0:
1068
+ # No native MTP on CPU (asserted above) => exactly one
1069
+ # query token per decode request, so the flat slice below
1070
+ # is already the batch-major layout
1071
+ # fp8_paged_mqa_logits_cpu wants.
1072
+ assert num_decode_tokens == batch_size, (
1073
+ "CPU sparse indexer decode path expects exactly one query "
1074
+ "token per decode request."
1075
+ )
1076
+ assert topk_tokens == 512, (
1077
+ "topk_transform_512_cpu only supports index_topk == 512."
1078
+ )
1079
+ seq_lens = decode_metadata.seq_lens
1080
+ seq_lens = (
1081
+ seq_lens[:, -1].contiguous() if seq_lens.ndim == 2 else seq_lens
1082
+ )
1083
+ block_table = decode_metadata.block_table[:batch_size]
1084
+ block_size = kv_cache.shape[1]
1085
+ kv_view = kv_cache.view(kv_cache.shape[0], -1)
1086
+
1087
+ logits = ops.fp8_paged_mqa_logits_cpu(
1088
+ q_quant[:num_decode_tokens],
1089
+ kv_view,
1090
+ weights[:num_decode_tokens],
1091
+ seq_lens,
1092
+ block_table,
1093
+ block_size,
1094
+ attn_metadata_narrowed.max_seq_len,
1095
+ )
1096
+
1097
+ # out_page_indices is a required kernel output but unused:
1098
+ # the indexer's topk output must stay local/compressed-
1099
+ # context positions (resolved later by
1100
+ # DeepseekV4CPUAttention.forward_mqa via
1101
+ # map_local_to_global_slots_cpu).
1102
+ out_page_scratch = torch.empty(
1103
+ (batch_size, topk_tokens),
1104
+ dtype=torch.int32,
1105
+ device=kv_cache.device,
1106
+ )
1107
+ ops.topk_transform_512_cpu(
1108
+ logits,
1109
+ seq_lens,
1110
+ block_table,
1111
+ out_page_scratch,
1112
+ block_size,
1113
+ topk_indices_buffer[:num_decode_tokens, :topk_tokens],
1114
+ )
1115
+
1116
+ return topk_indices_buffer
vllm/models/deepseek_v4/__init__.py
+7
@@ -30,6 +30,13 @@ elif current_platform.is_xpu():
30
from .xpu.dspark import DSparkDeepseekV4ForCausalLM # type: ignore[assignment]
31
from .xpu.model import DeepseekV4ForCausalLM # type: ignore[assignment]
32
from .xpu.mtp import DeepSeekV4MTP # type: ignore[assignment]
33
+elif current_platform.is_cpu():
34
+ from .cpu.dspark import DSparkDeepseekV4ForCausalLM # type: ignore[assignment]
35
+ from .cpu.model import DeepseekV4ForCausalLM # type: ignore[assignment]
36
+ from .cpu.mtp import DeepSeekV4MTP # type: ignore[assignment]
37
+ from .vl_stub import ( # type: ignore[assignment]
38
+ DeepseekV4ForConditionalGeneration,
39
+ )
40
else:
41
from .nvidia.dspark import ( # type: ignore[assignment]
42
DSparkDeepseekV4ForCausalLM,
vllm/models/deepseek_v4/common/ops/fused_indexer_q.py
+13
@@ -678,6 +678,19 @@ def fused_indexer_q_rope_quant(
678
index_q_fp8,
679
index_weights_out,
680
)
681
+ elif current_platform.is_cpu():
682
+ from vllm._custom_ops import fused_indexer_q_rope_quant_cpu
683
+
684
+ fused_indexer_q_rope_quant_cpu(
685
+ positions.contiguous(),
686
+ index_q.to(torch.float32),
687
+ index_q_cos_sin_cache.to(torch.float32).contiguous(),
688
+ index_q_fp8,
689
+ index_weights.to(torch.float32),
690
+ index_weights_softmax_scale,
691
+ index_weights_head_scale,
692
+ index_weights_out,
693
+ )
694
else:
695
_FUSED_INDEXER_Q_ROPE_QUANT_TRITON_KERNEL(
696
positions,
vllm/models/deepseek_v4/cpu/__init__.py
new
+2
@@ -0,0 +1,2 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
vllm/models/deepseek_v4/cpu/cpu_compressor.py
new
+112
@@ -0,0 +1,112 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""CPU DeepSeek-V4 compressor subclass and dispatch: partial-state-cache
4
+write and the compress -> RMSNorm -> RoPE -> FP8 quant -> KV cache store
5
+step, calling ``torch.ops._C.save_partial_states_cpu``/
6
+``compress_norm_rope_store_cpu``/``compress_norm_rope_store_indexer_cpu``
7
+(``csrc/cpu/sgl-kernels/compressor.cpp``) in place of triton-cpu.
8
+
9
+Covers both the head_dim=512 main-attention compressor path and the
10
+head_dim=128 indexer compressor path (fp8 only -- MXFP4, ``use_fp4_cache``,
11
+never occurs on CPU: the indexer's MXFP4 cache requires a Blackwell GPU).
12
+"""
13
+
14
+from typing import Any, cast
15
+
16
+import torch
17
+
18
+from vllm._custom_ops import (
19
+ compress_norm_rope_store_cpu,
20
+ compress_norm_rope_store_indexer_cpu,
21
+ save_partial_states_cpu,
22
+)
23
+from vllm.forward_context import get_forward_context
24
+from vllm.models.deepseek_v4.compressor import CompressorMetadata, DeepseekCompressor
25
+from vllm.models.deepseek_v4.cpu.cpu_utils import map_local_to_global_slots_cpu
26
+
27
+
28
+class DeepseekV4CPUCompressor(DeepseekCompressor):
29
+ """CPU compressor: same state/weights as the shared base, but
30
+ ``forward``'s cache-write and compress->RMSNorm->RoPE->quant->store
31
+ steps always dispatch straight to the ported CPU kernels above instead
32
+ of through the shared method's platform-dispatch chain. This class
33
+ only ever runs on CPU, and ``use_fp4_cache`` is always ``False`` here
34
+ (the indexer's MXFP4 cache requires a Blackwell GPU), so the shared
35
+ method's MXFP4/two-stage/cutedsl/triton branches never apply.
36
+ """
37
+
38
+ _norm_weight_fp32: torch.Tensor
39
+
40
+ def cache_norm_weight_fp32(self) -> None:
41
+ """Cache a fp32 contiguous copy of ``norm.weight`` once, right after
42
+ weight loading (see
43
+ ``DeepseekV4CPUAttention.process_weights_after_loading``) -- the CPU
44
+ kernel requires fp32/contiguous but the checkpoint loads this weight
45
+ in bf16, and it never changes after loading."""
46
+ self._norm_weight_fp32 = self.norm.weight.to(torch.float32).contiguous()
47
+
48
+ def forward(
49
+ self,
50
+ kv_score: torch.Tensor,
51
+ positions: torch.Tensor,
52
+ rotary_emb,
53
+ ) -> None:
54
+ kv, score = kv_score.split(
55
+ [self.coff * self.head_dim, self.coff * self.head_dim], dim=-1
56
+ )
57
+
58
+ attn_metadata = get_forward_context().attn_metadata
59
+ if not isinstance(attn_metadata, dict):
60
+ return
61
+
62
+ state_metadata = attn_metadata[self.state_cache.prefix]
63
+ assert isinstance(state_metadata, CompressorMetadata)
64
+ token_to_req_indices = state_metadata.token_to_req_indices
65
+ assert token_to_req_indices is not None
66
+ slot_mapping = state_metadata.slot_mapping
67
+ num_actual = slot_mapping.shape[0]
68
+ block_table = state_metadata.block_table
69
+ block_size = state_metadata.block_size
70
+
71
+ state_cache = self.state_cache.kv_cache
72
+
73
+ save_partial_states_cpu(
74
+ kv, score, self.ape, positions, state_cache, slot_mapping
75
+ )
76
+
77
+ cos_sin_cache = rotary_emb.cos_sin_cache
78
+ k_cache_metadata = cast(Any, attn_metadata[self.k_cache_prefix])
79
+ k_cache_layer = self._static_forward_context[self.k_cache_prefix]
80
+ kv_cache = k_cache_layer.kv_cache
81
+
82
+ assert self.head_dim in (512, 128), (
83
+ f"unsupported compressor head_dim: {self.head_dim}"
84
+ )
85
+ window = (1 + self.overlap) * self.compress_ratio
86
+ win_offsets = torch.arange(window, device=positions.device)
87
+ win_local = positions.unsqueeze(-1) - window + 1 + win_offsets
88
+ req_idx = token_to_req_indices[:num_actual]
89
+ gather_slots = map_local_to_global_slots_cpu(
90
+ win_local, req_idx, block_table, block_size
91
+ )
92
+
93
+ kv_cache_2d = kv_cache.view(kv_cache.shape[0], -1)
94
+ kv_cache_block_size = kv_cache.shape[1]
95
+
96
+ store_fn = (
97
+ compress_norm_rope_store_cpu
98
+ if self.head_dim == 512
99
+ else compress_norm_rope_store_indexer_cpu
100
+ )
101
+ store_fn(
102
+ state_cache,
103
+ gather_slots,
104
+ positions,
105
+ k_cache_metadata.slot_mapping,
106
+ self._norm_weight_fp32,
107
+ self.rms_norm_eps,
108
+ cos_sin_cache,
109
+ kv_cache_2d,
110
+ kv_cache_block_size,
111
+ self.compress_ratio,
112
+ )
vllm/models/deepseek_v4/cpu/cpu_mla.py
new
+126
@@ -0,0 +1,126 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""CPU DeepSeek-V4 sparse-MLA backend descriptor.
4
+
5
+DeepSeek-V4 attention runs entirely through
6
+``DeepseekV4Attention``/``DeepseekV4CPUAttention``, never through a generic
7
+``AttentionImpl``. ``DeepseekV4CPUSparseBackend`` exists so the CPU attention
8
+layer has its own backend name and its own metadata-builder class: CPU is not
9
+allowed to run triton-cpu inside the model's forward path, so
10
+``DeepseekV4CPUFlashMLAMetadataBuilder`` overrides the shared base's
11
+``_build_c128a_metadata`` (triton-backed in
12
+``DeepseekV4SparseMLAMetadataBuilder``) with a plain-eager-PyTorch
13
+reimplementation of the same per-token block-table-resolution math.
14
+"""
15
+
16
+import torch
17
+
18
+from vllm.models.deepseek_v4.sparse_mla import (
19
+ DeepseekV4SparseMLABackend,
20
+ DeepseekV4SparseMLAMetadataBuilder,
21
+)
22
+from vllm.v1.attention.backend import CommonAttentionMetadata
23
+from vllm.v1.attention.backends.mla.indexer import (
24
+ DeepseekV4IndexerBackend,
25
+ DeepseekV32IndexerMetadataBuilder,
26
+)
27
+from vllm.v1.attention.backends.mla.sparse_swa import (
28
+ _LAYER_TYPE_C4A,
29
+ _LAYER_TYPE_C128A,
30
+ _LAYER_TYPE_SWAONLY,
31
+ DeepseekSparseSWAMetadataBuilder,
32
+)
33
+from vllm.v1.attention.ops.flashmla import FlashMLASchedMeta
34
+
35
+
36
+class DeepseekV4CPUFlashMLAMetadataBuilder(DeepseekV4SparseMLAMetadataBuilder):
37
+ """CPU sparse-MLA metadata builder: same fields as the shared base, but
38
+ skips the C128A dense-topk metadata entirely -- unlike CUDA/XPU/ROCm,
39
+ ``DeepseekV4CPUAttention.forward_mqa`` never reads
40
+ ``c128a_global_decode_topk_indices``/``c128a_decode_topk_lens``/
41
+ ``c128a_prefill_topk_indices``; it recomputes the same local top-k
42
+ directly from ``positions`` instead, so building it here (whether
43
+ eagerly or via the shared base's triton kernel) would be wasted work.
44
+ """
45
+
46
+ def _build_c128a_metadata(
47
+ self,
48
+ cm: CommonAttentionMetadata,
49
+ req_id_per_token: torch.Tensor,
50
+ ) -> dict[str, torch.Tensor | None]:
51
+ return {}
52
+
53
+
54
+class DeepseekV4CPUSparseBackend(DeepseekV4SparseMLABackend):
55
+ @staticmethod
56
+ def get_name() -> str:
57
+ return "CPU_V4_MLA_SPARSE"
58
+
59
+ @staticmethod
60
+ def get_builder_cls() -> type[DeepseekV4CPUFlashMLAMetadataBuilder]:
61
+ return DeepseekV4CPUFlashMLAMetadataBuilder
62
+
63
+
64
+class DeepseekV4CPUIndexerMetadataBuilder(DeepseekV32IndexerMetadataBuilder):
65
+ """CPU indexer metadata builder: same fields as the shared base, but
66
+ prefill requests are never split into multiple chunks.
67
+
68
+ ``DeepseekV32IndexerMetadataBuilder._split_indexer_prefill_chunks`` (the
69
+ shared base's default) bounds two things the CUDA/XPU indexer kernels
70
+ need: the flat-gather workspace size and the dense M*N logits tensor
71
+ those kernels allocate for a chunk. The CPU indexer
72
+ (``sparse_attn_indexer_cpu``) allocates neither -- it reads the paged
73
+ K-cache directly via ``fp8_paged_mqa_logits_cpu``/
74
+ ``topk_transform_512_cpu``, with no chunk-splitting of its own.
75
+ With chunked prefill now enabled for this model on CPU (see
76
+ ``CpuPlatform.check_and_update_config``'s ``amx_mla_or_dsv4_enabled``),
77
+ the total query-token count for one step is already bounded by
78
+ ``max_num_batched_tokens`` -- there is nothing left here to bound, so
79
+ this always returns a single chunk spanning the whole step's prefill
80
+ batch.
81
+ """
82
+
83
+ @staticmethod
84
+ def _split_indexer_prefill_chunks(
85
+ compressed_seq_lens_cpu: torch.Tensor,
86
+ prefill_query_lens_cpu: torch.Tensor,
87
+ workspace_size: int,
88
+ max_logits_bytes: int,
89
+ request_offset: int = 0,
90
+ ) -> list[tuple[slice, slice]]:
91
+ # workspace_size/max_logits_bytes are unused: no flat-gather
92
+ # workspace or dense M*N buffer is allocated.
93
+ num_reqs = compressed_seq_lens_cpu.shape[0]
94
+ total_query_len = int(prefill_query_lens_cpu.sum().item())
95
+ return [
96
+ (
97
+ slice(request_offset, request_offset + num_reqs),
98
+ slice(0, total_query_len),
99
+ )
100
+ ]
101
+
102
+
103
+class DeepseekV4CPUIndexerBackend(DeepseekV4IndexerBackend):
104
+ @staticmethod
105
+ def get_name() -> str:
106
+ return "CPU_V4_INDEXER"
107
+
108
+ @staticmethod
109
+ def get_builder_cls() -> type[DeepseekV4CPUIndexerMetadataBuilder]:
110
+ return DeepseekV4CPUIndexerMetadataBuilder
111
+
112
+
113
+class DeepseekV4CPUSparseSWAMetadataBuilder(DeepseekSparseSWAMetadataBuilder):
114
+ """CPU SWA metadata builder: same fields as the shared base."""
115
+
116
+ def build_tile_scheduler(
117
+ self, num_decode_tokens: int
118
+ ) -> dict[str, FlashMLASchedMeta | None]:
119
+ """CPU never runs the FlashMLA tile-scheduler planner (that's a CUDA
120
+ C++ decode-path concern) -- always return the all-``None`` sentinel
121
+ the shared base's own CPU branch returns."""
122
+ return {
123
+ _LAYER_TYPE_SWAONLY: None,
124
+ _LAYER_TYPE_C4A: None,
125
+ _LAYER_TYPE_C128A: None,
126
+ }
vllm/models/deepseek_v4/cpu/cpu_sparse.py
new
+648
@@ -0,0 +1,648 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""CPU DeepSeek-V4 attention subclass.
4
+
5
+``forward_mqa`` resolves SWA/compressed top-k indices to paged-cache slot
6
+ids and calls the fused ``flash_mla_with_kvcache_cpu`` kernel. ``_o_proj``
7
+stays eager -- not on the attention hot path.
8
+"""
9
+
10
+from typing import Any, cast
11
+
12
+import torch
13
+import torch.nn as nn
14
+import torch.nn.functional as F
15
+
16
+from vllm import _custom_ops as ops
17
+from vllm.forward_context import get_forward_context
18
+from vllm.models.deepseek_v4.attention import (
19
+ DeepseekV4Attention,
20
+ DeepseekV4Indexer,
21
+ DeepseekV4IndexerCache,
22
+)
23
+from vllm.models.deepseek_v4.cpu.cpu_compressor import DeepseekV4CPUCompressor
24
+from vllm.models.deepseek_v4.cpu.cpu_mla import (
25
+ DeepseekV4CPUIndexerBackend,
26
+ DeepseekV4CPUSparseBackend,
27
+)
28
+from vllm.models.deepseek_v4.cpu.cpu_utils import map_local_to_global_slots_cpu
29
+from vllm.models.deepseek_v4.sparse_mla import DeepseekV4FlashMLAMetadata
30
+from vllm.platforms import current_platform
31
+from vllm.utils.torch_utils import (
32
+ LayerNameType,
33
+ _encode_layer_name,
34
+ _resolve_layer_name,
35
+ direct_register_custom_op,
36
+)
37
+from vllm.v1.attention.backend import AttentionBackend, AttentionMetadata
38
+from vllm.v1.attention.backends.mla.sparse_swa import DeepseekSparseSWAMetadata
39
+
40
+_ROPE_DIM = 64
41
+
42
+
43
+def _deepseek_v4_cpu_prepare_and_attn(
44
+ hidden_states: torch.Tensor,
45
+ qr: torch.Tensor,
46
+ kv: torch.Tensor,
47
+ qr_scale: torch.Tensor | None,
48
+ kv_score: torch.Tensor | None,
49
+ indexer_kv_score: torch.Tensor | None,
50
+ indexer_weights: torch.Tensor | None,
51
+ positions: torch.Tensor,
52
+ o_padded: torch.Tensor,
53
+ layer_name: LayerNameType,
54
+) -> torch.Tensor:
55
+ """Opaque custom-op boundary around attention-prep + sparse-indexer +
56
+ MLA attention + output projection -- CPU-only.
57
+
58
+ Wraps cache writes with no explicit tensor args (SWA/compressor/indexer
59
+ state), which torch.compile can't otherwise track as mutations -- one
60
+ opaque call sidesteps that. ``_o_proj`` is folded in too because CPU's
61
+ ``DYNAMO_TRACE_ONCE`` mode traces exactly once, during warmup when
62
+ ``attn_metadata`` is ``None``; tracing ``_o_proj`` separately would bake
63
+ that warmup branch into the compiled graph permanently, dead-code-
64
+ eliminating the whole attention computation upstream.
65
+ """
66
+ layer_name_str = _resolve_layer_name(layer_name)
67
+ self = get_forward_context().no_compile_layers[layer_name_str]
68
+ self._prepare_and_attn_fn(
69
+ hidden_states,
70
+ qr,
71
+ kv,
72
+ qr_scale,
73
+ kv_score,
74
+ indexer_kv_score,
75
+ indexer_weights,
76
+ positions,
77
+ o_padded,
78
+ )
79
+ o = o_padded[:, : self.n_local_heads, :]
80
+ return self._o_proj(o, positions)
81
+
82
+
83
+def _deepseek_v4_cpu_prepare_and_attn_fake(
84
+ hidden_states: torch.Tensor,
85
+ qr: torch.Tensor,
86
+ kv: torch.Tensor,
87
+ qr_scale: torch.Tensor | None,
88
+ kv_score: torch.Tensor | None,
89
+ indexer_kv_score: torch.Tensor | None,
90
+ indexer_weights: torch.Tensor | None,
91
+ positions: torch.Tensor,
92
+ o_padded: torch.Tensor,
93
+ layer_name: LayerNameType,
94
+) -> torch.Tensor:
95
+ # hidden_states' last dim already equals _o_proj's output width, so we
96
+ # don't need to resolve layer_name (an opaque object) here.
97
+ return hidden_states.new_empty((hidden_states.shape[0], hidden_states.shape[-1]))
98
+
99
+
100
+direct_register_custom_op(
101
+ op_name="deepseek_v4_cpu_prepare_and_attn",
102
+ op_func=_deepseek_v4_cpu_prepare_and_attn,
103
+ mutates_args=["o_padded"],
104
+ fake_impl=_deepseek_v4_cpu_prepare_and_attn_fake,
105
+)
106
+
107
+
108
+def _dequant_linear_weight(layer: torch.nn.Module) -> torch.Tensor:
109
+ """Dequantize a linear layer's weight, including FP8 block scales.
110
+
111
+ ``_o_proj`` reads ``wo_a.weight`` directly (not via
112
+ ``quant_method.apply()``), so the per-block FP8 scale must be applied
113
+ here explicitly.
114
+ """
115
+ weight = layer.weight
116
+ scale = getattr(layer, "weight_scale_inv", None)
117
+ if scale is None:
118
+ scale = getattr(layer, "weight_scale", None)
119
+ if scale is None:
120
+ return weight.float()
121
+
122
+ if scale.dtype == torch.float8_e8m0fnu:
123
+ from vllm.model_executor.layers.quantization.utils.fp8_utils import (
124
+ _upcast_e8m0_to_fp32,
125
+ )
126
+
127
+ scale = _upcast_e8m0_to_fp32(scale)
128
+ scale = scale.to(torch.float32)
129
+
130
+ if scale.numel() == 1:
131
+ return weight.float() * scale
132
+
133
+ block_size = getattr(layer, "weight_block_size", None)
134
+ if block_size is None:
135
+ # Per-channel scale: one row-scale per output row.
136
+ return weight.float() * scale.view(-1, 1)
137
+
138
+ block_m, block_k = block_size
139
+ m, k = weight.shape
140
+ scale = scale.repeat_interleave(block_m, dim=0).repeat_interleave(block_k, dim=1)
141
+ return weight.float() * scale[:m, :k]
142
+
143
+
144
+def _fused_indexer_q_rope_quant_cpu(
145
+ positions: torch.Tensor,
146
+ index_q: torch.Tensor,
147
+ index_q_cos_sin_cache: torch.Tensor,
148
+ index_weights: torch.Tensor,
149
+ index_weights_softmax_scale: float,
150
+ index_weights_head_scale: float,
151
+) -> tuple[torch.Tensor, torch.Tensor]:
152
+ """CPU-only equivalent of ``fused_indexer_q_rope_quant``'s FP8 path
153
+ (CPU never reaches its MXFP4 arm -- Blackwell-only). Returns
154
+ ``(q_fp8, weights_out)``, folding the per-token q scale into
155
+ ``index_weights``. ``positions``/``index_q_cos_sin_cache`` are already
156
+ contiguous int64/fp32, so no conversion is needed here.
157
+ """
158
+ assert positions.ndim == 1
159
+ assert index_q.ndim == 3
160
+ assert index_q_cos_sin_cache.ndim == 2
161
+
162
+ index_weights_out = torch.empty_like(index_weights, dtype=torch.float32)
163
+ index_q_fp8 = torch.empty_like(index_q, dtype=current_platform.fp8_dtype())
164
+ ops.fused_indexer_q_rope_quant_cpu(
165
+ positions,
166
+ index_q.to(torch.float32),
167
+ index_q_cos_sin_cache,
168
+ index_q_fp8,
169
+ index_weights.to(torch.float32),
170
+ index_weights_softmax_scale,
171
+ index_weights_head_scale,
172
+ index_weights_out,
173
+ )
174
+ return index_q_fp8, index_weights_out
175
+
176
+
177
+class _NoOpEvent:
178
+ """Stand-in for ``torch.cuda.Event`` on CPU: since CPU never sets an aux
179
+ stream list, these events are constructed but never actually used."""
180
+
181
+
182
+class DeepseekV4CPUIndexerCache(DeepseekV4IndexerCache):
183
+ """CPU indexer K-cache descriptor: same fields as the shared base, just
184
+ pointing ``get_attn_backend()`` at ``DeepseekV4CPUIndexerBackend``
185
+ instead of the shared, CUDA/XPU-oriented ``DeepseekV4IndexerBackend``."""
186
+
187
+ def get_attn_backend(self) -> type[AttentionBackend]:
188
+ return DeepseekV4CPUIndexerBackend
189
+
190
+
191
+class DeepseekV4CPUIndexer(DeepseekV4Indexer):
192
+ """CPU indexer: the C4A short-context fallback runs as eager PyTorch
193
+ instead of Triton. Never constructed directly -- ``__class__`` is
194
+ swapped onto an existing ``DeepseekV4Indexer`` instance post-construction.
195
+ """
196
+
197
+ def forward(
198
+ self,
199
+ hidden_states: torch.Tensor,
200
+ qr: torch.Tensor,
201
+ compressed_kv_score: torch.Tensor,
202
+ indexer_weights: torch.Tensor,
203
+ positions: torch.Tensor,
204
+ rotary_emb: nn.Module,
205
+ qr_scale: torch.Tensor | None = None,
206
+ ) -> tuple[torch.Tensor | None, torch.Tensor | None, torch.Tensor | None]:
207
+ """CPU override: no aux streams, so wq_b_and_q_quant and the
208
+ compressor run straight-line instead of through
209
+ ``maybe_execute_in_parallel``; the short-context check drops the
210
+ CUDA-only ``is_current_stream_capturing()`` guard (always False
211
+ here)."""
212
+ compressor = self.compressor
213
+
214
+ attn_metadata = get_forward_context().attn_metadata
215
+ if isinstance(attn_metadata, dict):
216
+ indexer_metadata = cast(Any, attn_metadata[self.k_cache.prefix])
217
+ if indexer_metadata.max_seq_len // self.compress_ratio <= self.topk_tokens:
218
+ # Fewer candidates than topk: all are selected, but the K
219
+ # cache still needs to be built.
220
+ compressor(compressed_kv_score, positions, rotary_emb)
221
+ assert self.topk_indices_buffer is not None
222
+ num_tokens = (
223
+ indexer_metadata.num_decode_tokens
224
+ + indexer_metadata.num_prefill_tokens
225
+ )
226
+ if num_tokens > 0:
227
+ self._fill_short_context_topk(positions, num_tokens)
228
+ return None, None, None
229
+
230
+ def wq_b_and_q_quant():
231
+ q = self._wq_b_proj(qr, qr_scale)
232
+ q = q.view(-1, self.n_head, self.head_dim)
233
+ return _fused_indexer_q_rope_quant_cpu(
234
+ positions,
235
+ q,
236
+ rotary_emb.cos_sin_cache,
237
+ indexer_weights,
238
+ self.softmax_scale,
239
+ self.n_head**-0.5,
240
+ )
241
+
242
+ # Order matches the base class's parallel-join contract: compressor
243
+ # must complete before indexer_op runs (skip_k_cache_insert=True).
244
+ q, weights = wq_b_and_q_quant()
245
+ compressor(compressed_kv_score, positions, rotary_emb)
246
+ # No q_scale: FP8-only path (MXFP4 needs Blackwell).
247
+ return q, None, weights
248
+
249
+ def _fill_short_context_topk(
250
+ self, positions: torch.Tensor, num_tokens: int
251
+ ) -> None:
252
+ assert self.topk_indices_buffer is not None
253
+ device = positions.device
254
+ arange = torch.arange(self.topk_tokens, device=device)
255
+ num_compressed = torch.div(
256
+ positions[:num_tokens] + 1, self.compress_ratio, rounding_mode="floor"
257
+ )
258
+ valid = arange.unsqueeze(0) < num_compressed.unsqueeze(-1)
259
+ filled = torch.where(
260
+ valid, arange.unsqueeze(0), torch.full_like(arange.unsqueeze(0), -1)
261
+ )
262
+ self.topk_indices_buffer[:num_tokens, : self.topk_tokens] = filled.to(
263
+ self.topk_indices_buffer.dtype
264
+ )
265
+
266
+
267
+class DeepseekV4CPUAttention(DeepseekV4Attention):
268
+ """CPU sparse MLA attention layer for DeepSeek V4."""
269
+
270
+ backend_cls = DeepseekV4CPUSparseBackend
271
+
272
+ def __init__(self, *args, **kwargs) -> None:
273
+ super().__init__(*args, **kwargs)
274
+
275
+ # Swap classes post-construction (no class-swap hook on the shared
276
+ # base) so forward()/get_attn_backend() dispatch to CPU kernels.
277
+ # SWA-only layers (compress_ratio <= 1) never build a
278
+ # compressor/indexer.
279
+ if self.compressor is not None:
280
+ self.compressor.__class__ = DeepseekV4CPUCompressor
281
+ if self.indexer is not None:
282
+ self.indexer.__class__ = DeepseekV4CPUIndexer
283
+ self.indexer.k_cache.__class__ = DeepseekV4CPUIndexerCache
284
+ self.indexer.compressor.__class__ = DeepseekV4CPUCompressor
285
+
286
+ self._wo_a_packed: torch.Tensor | None = None
287
+ self._wrap_wo_a_process_weights_after_loading()
288
+
289
+ def process_weights_after_loading(self, act_dtype: torch.dtype) -> None:
290
+ """Cache fp32-contiguous copies of tensors CPU kernels want that way
291
+ but that arrive bf16 (rotary cos/sin table, compressor RMSNorm
292
+ weights) -- cast once here instead of per forward call.
293
+
294
+ Runs after every quantized layer's own
295
+ ``process_weights_after_loading`` (see ``is_deferred_attention_layer``);
296
+ ``wo_a``'s packing must happen *before* that phase instead, so it's a
297
+ separate monkeypatch in ``__init__``.
298
+ """
299
+ self.rotary_emb.cos_sin_cache = self.rotary_emb.cos_sin_cache.to(
300
+ torch.float32
301
+ ).contiguous()
302
+ if self.compressor is not None:
303
+ self.compressor.cache_norm_weight_fp32()
304
+ if self.indexer is not None:
305
+ self.indexer.compressor.cache_norm_weight_fp32()
306
+
307
+ def _wrap_wo_a_process_weights_after_loading(self) -> None:
308
+ """Snapshot and pack ``wo_a``'s weight into a bf16 copy for
309
+ ``bmm_cpu`` before the FP8 kernel's own
310
+ ``process_weights_after_loading`` VNNI-repacks it in place for
311
+ row-major reads ``_o_proj`` never does."""
312
+ wo_a = self.wo_a
313
+ quant_method = wo_a.quant_method
314
+ orig_pwal = quant_method.process_weights_after_loading
315
+ heads_per_group = self.n_local_heads // self.n_local_groups
316
+ k_dim = heads_per_group * self.head_dim
317
+
318
+ def _capture_and_pack(layer: torch.nn.Module) -> None:
319
+ dequant = _dequant_linear_weight(layer).to(torch.bfloat16)
320
+ dequant = dequant.view(self.n_local_groups, self.o_lora_rank, k_dim)
321
+ self._wo_a_packed = torch.ops._C.convert_weight_packed(dequant.contiguous())
322
+ orig_pwal(layer)
323
+ # orig_pwal VNNI-repacks layer.weight/scale for ordinary-linear
324
+ # use that _o_proj never needs -- free them, same pattern as
325
+ # dispatch_cpu_unquantized_gemm's remove_weight=True.
326
+ layer.weight = torch.nn.Parameter(torch.empty(0), requires_grad=False)
327
+ scale_attr = (
328
+ "weight_scale_inv"
329
+ if hasattr(layer, "weight_scale_inv")
330
+ else "weight_scale"
331
+ )
332
+ if hasattr(layer, scale_attr):
333
+ setattr(
334
+ layer,
335
+ scale_attr,
336
+ torch.nn.Parameter(torch.empty(0), requires_grad=False),
337
+ )
338
+
339
+ quant_method.process_weights_after_loading = ( # type: ignore[method-assign]
340
+ _capture_and_pack
341
+ )
342
+
343
+ def forward(
344
+ self,
345
+ positions: torch.Tensor,
346
+ hidden_states: torch.Tensor,
347
+ llama_4_scaling: torch.Tensor | None = None,
348
+ ) -> torch.Tensor:
349
+ """CPU override: wraps attention-prep + MLA + output projection in the
350
+ CPU-only opaque custom op (see ``_deepseek_v4_cpu_prepare_and_attn``).
351
+ """
352
+ num_tokens = hidden_states.shape[0]
353
+ o_padded = torch.empty(
354
+ (num_tokens, self.padded_heads, self.head_dim),
355
+ dtype=hidden_states.dtype,
356
+ device=hidden_states.device,
357
+ )
358
+
359
+ qr_kv = self._fused_wqa_wkv_gemm(hidden_states)
360
+
361
+ kv_score: torch.Tensor | None = None
362
+ if self.compressor is not None:
363
+ # Must go through the module's own forward (not a raw matmul on
364
+ # its weight): the packed-GEMM kernel frees the raw weight after
365
+ # packing, same as wo_a's.
366
+ kv_score = self.compressor.fused_wkv_wgate(hidden_states).to(torch.float32)
367
+
368
+ indexer_weights: torch.Tensor | None = None
369
+ indexer_kv_score: torch.Tensor | None = None
370
+ if self.indexer is not None:
371
+ indexer_weights, _ = self.indexer.weights_proj(hidden_states)
372
+ indexer_kv_score = self.indexer.compressor.fused_wkv_wgate(
373
+ hidden_states
374
+ ).to(torch.float32)
375
+
376
+ qr, qr_scale, kv = self._split_qkv_and_norm(qr_kv)
377
+
378
+ return torch.ops.vllm.deepseek_v4_cpu_prepare_and_attn(
379
+ hidden_states,
380
+ qr,
381
+ kv,
382
+ qr_scale,
383
+ kv_score,
384
+ indexer_kv_score,
385
+ indexer_weights,
386
+ positions,
387
+ o_padded,
388
+ _encode_layer_name(self.prefix),
389
+ )
390
+
391
+ def _fused_qnorm_rope_kv_insert(
392
+ self,
393
+ q: torch.Tensor,
394
+ kv: torch.Tensor,
395
+ positions: torch.Tensor,
396
+ attn_metadata: (
397
+ dict[str, AttentionMetadata] | list[dict[str, AttentionMetadata]] | None
398
+ ),
399
+ ) -> torch.Tensor:
400
+ """CPU override: only the fp8_ds_mla (uint8) SWA cache layout is
401
+ supported here, so the base method's bf16/per-tensor-fp8 branches
402
+ are dropped."""
403
+ if not isinstance(attn_metadata, dict):
404
+ # Profile run: kernel doesn't fire; produce a padded tensor so
405
+ # downstream FlashMLA gets the right shape.
406
+ if self.n_local_heads < self.padded_heads:
407
+ return F.pad(
408
+ q,
409
+ (0, 0, 0, self.padded_heads - self.n_local_heads),
410
+ value=0.0,
411
+ )
412
+ return q
413
+
414
+ swa_metadata = cast(
415
+ "DeepseekSparseSWAMetadata | None",
416
+ attn_metadata.get(self.swa_cache_layer.prefix),
417
+ )
418
+ assert swa_metadata is not None
419
+
420
+ swa_kv_cache = self.swa_cache_layer.kv_cache
421
+ # positions is already int64 (the runner's buffer), no cast needed.
422
+ assert positions.dtype == torch.int64
423
+ assert swa_kv_cache.dtype == torch.uint8, (
424
+ "DeepseekV4CPUAttention only supports the fp8_ds_mla cache layout"
425
+ )
426
+ # Fused: Q side does per-head weight-free RMSNorm + GPT-J RoPE
427
+ # (zero-filling padded head slots); KV side does GPT-J RoPE + UE8M0
428
+ # FP8 quant + paged insert.
429
+ swa_kv_cache_2d = swa_kv_cache.view(swa_kv_cache.shape[0], -1)
430
+ return ops.fused_qnorm_rope_kv_insert_cpu(
431
+ q,
432
+ kv,
433
+ positions,
434
+ swa_kv_cache_2d,
435
+ swa_metadata.slot_mapping,
436
+ self.rotary_emb.cos_sin_cache,
437
+ self.padded_heads,
438
+ self.eps,
439
+ swa_metadata.block_size,
440
+ )
441
+
442
+ def _split_qkv_and_norm(
443
+ self, qr_kv: torch.Tensor
444
+ ) -> tuple[torch.Tensor, torch.Tensor | None, torch.Tensor]:
445
+ """CPU override: two ``RMSNorm`` calls instead of the shared
446
+ ``fused_q_kv_rmsnorm``, whose raw ``@triton.jit`` kernel fails to
447
+ link under triton-cpu (``undefined symbol: __truncdfbf2``)."""
448
+ qr, kv = qr_kv.split([self.q_lora_rank, self.head_dim], dim=-1)
449
+ return self.q_norm(qr), None, self.kv_norm(kv)
450
+
451
+ def _prepare_and_attn(
452
+ self,
453
+ hidden_states: torch.Tensor,
454
+ qr: torch.Tensor,
455
+ kv: torch.Tensor,
456
+ qr_scale: torch.Tensor | None,
457
+ kv_score: torch.Tensor,
458
+ indexer_kv_score: torch.Tensor,
459
+ indexer_weights: torch.Tensor,
460
+ positions: torch.Tensor,
461
+ o_padded: torch.Tensor,
462
+ ) -> None:
463
+ """CPU override: no aux streams, so query-projection+KV-insert, the
464
+ indexer, and the compressor run straight-line instead of through
465
+ execute_in_parallel/maybe_execute_in_parallel."""
466
+ attn_metadata = get_forward_context().attn_metadata
467
+ indexer = self.indexer
468
+ compressor = self.compressor
469
+
470
+ q = self._wq_b_proj(qr, qr_scale).view(-1, self.n_local_heads, self.head_dim)
471
+ q = self._fused_qnorm_rope_kv_insert(q, kv, positions, attn_metadata)
472
+
473
+ index_q: torch.Tensor | None = None
474
+ index_q_scale: torch.Tensor | None = None
475
+ index_weights_out: torch.Tensor | None = None
476
+ if indexer is not None:
477
+ assert compressor is not None
478
+ index_q, index_q_scale, index_weights_out = indexer(
479
+ hidden_states,
480
+ qr,
481
+ indexer_kv_score,
482
+ indexer_weights,
483
+ positions,
484
+ self.indexer_rotary_emb,
485
+ qr_scale,
486
+ )
487
+ compressor(kv_score, positions, self.rotary_emb)
488
+ elif compressor is not None:
489
+ compressor(kv_score, positions, self.rotary_emb)
490
+
491
+ self._sparse_indexer_and_attn(
492
+ hidden_states,
493
+ index_q,
494
+ index_q_scale,
495
+ index_weights_out,
496
+ q,
497
+ kv,
498
+ positions,
499
+ o_padded,
500
+ )
501
+
502
+ def _sparse_indexer_and_attn(
503
+ self,
504
+ hidden_states: torch.Tensor,
505
+ index_q: torch.Tensor | None,
506
+ index_q_scale: torch.Tensor | None,
507
+ index_weights: torch.Tensor | None,
508
+ q: torch.Tensor,
509
+ kv: torch.Tensor,
510
+ positions: torch.Tensor,
511
+ out: torch.Tensor,
512
+ ) -> None:
513
+ """CPU override: identical body, minus ``@eager_break_during_capture``
514
+ (this platform never captures a CUDA graph, so it was a no-op)."""
515
+ if self.indexer is not None and index_q is not None:
516
+ assert index_weights is not None
517
+ # index_q_scale is always None on CPU (FP8-only path).
518
+ assert index_q_scale is None
519
+ self.indexer.indexer_op(hidden_states, index_q, None, index_weights)
520
+
521
+ self.forward_mqa(q, kv, positions, out)
522
+
523
+ @classmethod
524
+ def get_padded_num_q_heads(cls, num_heads: int) -> int:
525
+ # No head-count padding constraint for the (not yet ported) CPU kernel.
526
+ return num_heads
527
+
528
+ def _o_proj(self, o: torch.Tensor, positions: torch.Tensor) -> torch.Tensor:
529
+ if get_forward_context().attn_metadata is None:
530
+ # Warmup/profile run: no real metadata, o is a zeroed placeholder.
531
+ return torch.zeros(
532
+ (o.shape[0], self.hidden_size), dtype=o.dtype, device=o.device
533
+ )
534
+
535
+ num_tokens, num_heads, _ = o.shape
536
+ cos_sin_cache = self.rotary_emb.cos_sin_cache
537
+ # Undo the query's own GPT-J RoPE rotation: the attention output is
538
+ # a weighted sum of KV rows at different positions, so R(-pos_q)
539
+ # recovers the translation-invariant quantity wo_a/wo_b expect.
540
+ o_derot = ops.inverse_gptj_rope_o_proj_cpu(
541
+ o, positions, cos_sin_cache, _ROPE_DIM
542
+ )
543
+
544
+ heads_per_group = num_heads // self.n_local_groups
545
+ o_grouped = o_derot.reshape(
546
+ num_tokens, self.n_local_groups, heads_per_group * self.head_dim
547
+ ).to(torch.bfloat16)
548
+ assert self._wo_a_packed is not None
549
+ z = o_grouped.new_empty((self.n_local_groups, num_tokens, self.o_lora_rank))
550
+ ops.bmm_cpu(
551
+ z,
552
+ o_grouped.transpose(0, 1).contiguous(),
553
+ self._wo_a_packed,
554
+ True,
555
+ None,
556
+ )
557
+ z = z.transpose(0, 1).to(o.dtype)
558
+ return self.wo_b(z.flatten(1))
559
+
560
+ def forward_mqa(
561
+ self,
562
+ q: torch.Tensor,
563
+ kv: torch.Tensor,
564
+ positions: torch.Tensor,
565
+ output: torch.Tensor,
566
+ ) -> None:
567
+ forward_context = get_forward_context()
568
+ attn_metadata = forward_context.attn_metadata
569
+ if attn_metadata is None:
570
+ # Warmup/profile run: reserve nothing extra, skip real kernels.
571
+ output.zero_()
572
+ return
573
+
574
+ assert isinstance(attn_metadata, dict)
575
+ swa_metadata = cast(
576
+ DeepseekSparseSWAMetadata,
577
+ attn_metadata.get(self.swa_cache_layer.prefix),
578
+ )
579
+ assert swa_metadata is not None
580
+ assert swa_metadata.token_to_req_indices is not None
581
+ req_idx = swa_metadata.token_to_req_indices
582
+
583
+ swa_only = self.compress_ratio <= 1
584
+ flashmla_metadata = cast(
585
+ DeepseekV4FlashMLAMetadata | None, attn_metadata.get(self.prefix)
586
+ )
587
+
588
+ num_tokens = q.shape[0]
589
+ device = q.device
590
+
591
+ # ---- SWA window: causal sliding window over the raw per-token cache
592
+ window = self.window_size
593
+ win_offsets = torch.arange(window, device=device)
594
+ start_pos = (positions - window + 1).clamp(min=0)
595
+ local_pos = start_pos.unsqueeze(-1) + win_offsets
596
+ valid_win = local_pos <= positions.unsqueeze(-1)
597
+ win_local = torch.where(valid_win, local_pos, torch.full_like(local_pos, -1))
598
+ window_slots = map_local_to_global_slots_cpu(
599
+ win_local, req_idx, swa_metadata.block_table, swa_metadata.block_size
600
+ )
601
+ window_cache_2d = self.swa_cache_layer.kv_cache.view(
602
+ self.swa_cache_layer.kv_cache.shape[0], -1
603
+ )
604
+
605
+ # ---- Compressed KV: C4A's learned top-k, or C128A's dense top-k ----
606
+ if swa_only:
607
+ compressed_slots = window_slots.new_full((num_tokens, 0), -1)
608
+ # Unused when num_compressed == 0 (only dtype/2D-ness matter);
609
+ # reuse the window cache to avoid an allocation.
610
+ compressed_cache_2d = window_cache_2d
611
+ compressed_block_size = 1
612
+ else:
613
+ assert flashmla_metadata is not None
614
+ main_block_size = flashmla_metadata.block_size // self.compress_ratio
615
+ if self.compress_ratio == 4:
616
+ assert self.topk_indices_buffer is not None
617
+ local_topk = self.topk_indices_buffer[:num_tokens]
618
+ else:
619
+ width = -(-self.max_model_len // self.compress_ratio)
620
+ num_compressed = torch.div(
621
+ positions + 1, self.compress_ratio, rounding_mode="floor"
622
+ ).clamp(max=width)
623
+ arange_w = torch.arange(width, device=device)
624
+ local_topk = torch.where(
625
+ arange_w.unsqueeze(0) < num_compressed.unsqueeze(-1),
626
+ arange_w.unsqueeze(0).expand(num_tokens, width),
627
+ torch.full(
628
+ (num_tokens, width), -1, device=device, dtype=torch.int64
629
+ ),
630
+ )
631
+ compressed_slots = map_local_to_global_slots_cpu(
632
+ local_topk, req_idx, flashmla_metadata.block_table, main_block_size
633
+ )
634
+ compressed_cache_2d = self.kv_cache.view(self.kv_cache.shape[0], -1)
635
+ compressed_block_size = main_block_size
636
+
637
+ torch.ops._C.flash_mla_with_kvcache_cpu(
638
+ output,
639
+ q,
640
+ window_cache_2d,
641
+ window_slots,
642
+ swa_metadata.block_size,
643
+ compressed_cache_2d,
644
+ compressed_slots,
645
+ compressed_block_size,
646
+ self.attn_sink,
647
+ self.scale,
648
+ )
vllm/models/deepseek_v4/cpu/cpu_utils.py
new
+34
@@ -0,0 +1,34 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""Shared helpers for DeepSeek-V4's CPU-ported kernels."""
4
+
5
+import torch
6
+
7
+
8
+def map_local_to_global_slots_cpu(
9
+ local_indices: torch.Tensor,
10
+ req_idx: torch.Tensor,
11
+ block_table: torch.Tensor,
12
+ block_size: int,
13
+) -> torch.Tensor:
14
+ """Map per-request-local (token or compressed-token) indices to global
15
+ paged-cache slot ids via a block table, mirroring
16
+ ``compute_global_topk_indices_and_lens``.
17
+
18
+ Args:
19
+ local_indices: [N, K], -1 sentinel for invalid entries.
20
+ req_idx: [N], row into ``block_table`` for each of the N rows.
21
+ block_table: [num_reqs, max_blocks_per_seq].
22
+ block_size: tokens (or compressed tokens) per physical block.
23
+
24
+ Returns:
25
+ [N, K] int64 global slot ids, -1 where ``local_indices`` was invalid.
26
+ """
27
+ valid = local_indices >= 0
28
+ safe_local = local_indices.clamp(min=0).to(torch.int64)
29
+ block_pos = torch.div(safe_local, block_size, rounding_mode="floor")
30
+ offset = safe_local % block_size
31
+ req_expand = req_idx.to(torch.int64).unsqueeze(-1).expand_as(local_indices)
32
+ block_num = block_table[req_expand, block_pos].to(torch.int64)
33
+ slot = block_num * block_size + offset
34
+ return torch.where(valid, slot, torch.full_like(slot, -1))
vllm/models/deepseek_v4/cpu/dspark.py
new
+18
@@ -0,0 +1,18 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""CPU DeepSeek-V4 DSpark (speculative decoding) — deferred, not implemented.
4
+
5
+``DSparkDeepseekV4ForCausalLM`` is only instantiated when the speculative-
6
+decoding registry resolves it by name for a DSpark speculative config; a
7
+plain (non-spec) DeepSeek-V4 CPU model never constructs this class. This
8
+stub exists solely so ``vllm.models.deepseek_v4.cpu`` is import-clean.
9
+"""
10
+
11
+from vllm.config import VllmConfig
12
+
13
+
14
+class DSparkDeepseekV4ForCausalLM:
15
+ def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None:
16
+ raise NotImplementedError(
17
+ "DeepSeek-V4 speculative decoding (DSpark) is not implemented on CPU."
18
+ )
vllm/models/deepseek_v4/cpu/model.py
new
+922
@@ -0,0 +1,922 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+import typing
4
+from collections.abc import Callable, Iterable
5
+from itertools import islice
6
+
7
+import regex as re
8
+import torch
9
+import torch.nn as nn
10
+
11
+from vllm.compilation.decorators import support_torch_compile
12
+from vllm.config import VllmConfig
13
+from vllm.distributed import (
14
+ get_pp_group,
15
+ get_tensor_model_parallel_rank,
16
+ get_tensor_model_parallel_world_size,
17
+)
18
+from vllm.model_executor.layers.activation import SiluAndMul, SiluAndMulWithClamp
19
+from vllm.model_executor.layers.fused_moe import (
20
+ FusedMoEFactory,
21
+ GateLinear,
22
+ fused_moe_make_expert_params_mapping,
23
+)
24
+from vllm.model_executor.layers.fused_moe.utils import (
25
+ is_model_fused_shared_expert_compatible,
26
+)
27
+from vllm.model_executor.layers.layernorm import RMSNorm
28
+from vllm.model_executor.layers.linear import (
29
+ MergedColumnParallelLinear,
30
+ RowParallelLinear,
31
+)
32
+from vllm.model_executor.layers.logits_processor import LogitsProcessor
33
+from vllm.model_executor.layers.mhc import (
34
+ HAS_AITER_MHC,
35
+ HAS_TILELANG_MHC,
36
+ HCHeadOp,
37
+ MHCFusedPostPreOp,
38
+ MHCPostOp,
39
+ MHCPreOp,
40
+)
41
+from vllm.model_executor.layers.quantization import QuantizationConfig
42
+from vllm.model_executor.layers.vocab_parallel_embedding import (
43
+ ParallelLMHead,
44
+ VocabParallelEmbedding,
45
+)
46
+from vllm.model_executor.model_loader.weight_utils import default_weight_loader
47
+from vllm.model_executor.models.interfaces import (
48
+ EagleModelMixin,
49
+ SupportsEagle3,
50
+ SupportsPP,
51
+)
52
+from vllm.model_executor.models.utils import (
53
+ AutoWeightsLoader,
54
+ PPMissingLayer,
55
+ WeightsMapper,
56
+ extract_layer_index,
57
+ is_pp_missing_parameter,
58
+ make_layers,
59
+ maybe_prefix,
60
+)
61
+from vllm.models.deepseek_v4.cpu.cpu_sparse import DeepseekV4CPUAttention
62
+from vllm.platforms import current_platform
63
+from vllm.sequence import IntermediateTensors
64
+
65
+
66
+class DeepseekV4MLP(nn.Module):
67
+ def __init__(
68
+ self,
69
+ hidden_size: int,
70
+ intermediate_size: int,
71
+ hidden_act: str,
72
+ swiglu_limit: float | None = None,
73
+ quant_config: QuantizationConfig | None = None,
74
+ reduce_results: bool = True,
75
+ is_sequence_parallel: bool = False,
76
+ prefix: str = "",
77
+ ) -> None:
78
+ super().__init__()
79
+
80
+ # is_sequence_parallel: weights are replicated, no collective ops
81
+ # needed. Otherwise: standard TP with an allreduce at the end.
82
+ self.gate_up_proj = MergedColumnParallelLinear(
83
+ hidden_size,
84
+ [intermediate_size] * 2,
85
+ bias=False,
86
+ quant_config=quant_config,
87
+ disable_tp=is_sequence_parallel,
88
+ prefix=f"{prefix}.gate_up_proj",
89
+ )
90
+ self.down_proj = RowParallelLinear(
91
+ intermediate_size,
92
+ hidden_size,
93
+ bias=False,
94
+ quant_config=quant_config,
95
+ reduce_results=reduce_results,
96
+ disable_tp=is_sequence_parallel,
97
+ prefix=f"{prefix}.down_proj",
98
+ )
99
+ if hidden_act != "silu":
100
+ raise ValueError(
101
+ f"Unsupported activation: {hidden_act}. Only silu is supported for now."
102
+ )
103
+ if swiglu_limit is not None:
104
+ self.act_fn = SiluAndMulWithClamp(swiglu_limit)
105
+ else:
106
+ self.act_fn = SiluAndMul()
107
+
108
+ def forward(self, x):
109
+ gate_up, _ = self.gate_up_proj(x)
110
+ x = self.act_fn(gate_up)
111
+ x, _ = self.down_proj(x)
112
+ return x
113
+
114
+
115
+class DeepseekV4MoE(nn.Module):
116
+ def __init__(
117
+ self,
118
+ vllm_config: VllmConfig,
119
+ prefix: str = "",
120
+ ):
121
+ super().__init__()
122
+
123
+ self.tp_size = get_tensor_model_parallel_world_size()
124
+ config = vllm_config.model_config.hf_config
125
+ quant_config = vllm_config.quant_config
126
+ self.prefix = prefix
127
+
128
+ self.routed_scaling_factor = getattr(config, "routed_scaling_factor", 1.0)
129
+ self.hidden_size = config.hidden_size
130
+
131
+ self.n_routed_experts = config.n_routed_experts
132
+ self.n_activated_experts = config.num_experts_per_tok
133
+ self.moe_intermediate_size = config.moe_intermediate_size
134
+ self.swiglu_limit = config.swiglu_limit
135
+ self.renormalize = config.norm_topk_prob
136
+ self.scoring_func = getattr(config, "scoring_func", "sqrtsoftplus")
137
+
138
+ self.gate = GateLinear(
139
+ input_size=config.hidden_size,
140
+ output_size=config.n_routed_experts,
141
+ bias=False,
142
+ out_dtype=torch.float32,
143
+ prefix=f"{prefix}.gate",
144
+ )
145
+
146
+ self.gate.e_score_correction_bias = None
147
+ self.gate.tid2eid = None
148
+ is_hash_moe = extract_layer_index(prefix) < config.num_hash_layers
149
+ self.hash_indices_dtype = torch.int32
150
+ if is_hash_moe:
151
+ # hash MoE doesn't use e_score_correction_bias
152
+ # Use randint instead of empty to avoid garbage values causing
153
+ # invalid memory access in dummy mode (--load-format="dummy")
154
+ self.gate.tid2eid = nn.Parameter(
155
+ torch.randint(
156
+ 0,
157
+ config.n_routed_experts,
158
+ (config.vocab_size, config.num_experts_per_tok),
159
+ dtype=self.hash_indices_dtype,
160
+ ),
161
+ requires_grad=False,
162
+ )
163
+ elif getattr(config, "topk_method", None) == "noaux_tc":
164
+ self.gate.e_score_correction_bias = nn.Parameter(
165
+ torch.empty(config.n_routed_experts, dtype=torch.float32),
166
+ requires_grad=False,
167
+ )
168
+
169
+ self.n_shared_experts = config.n_shared_experts
170
+
171
+ # No CPU fused-shared-expert grouped-GEMM path yet.
172
+ self.is_fused_shared_expert_enabled = False
173
+
174
+ if config.n_shared_experts is None or self.is_fused_shared_expert_enabled:
175
+ self.shared_experts = None
176
+ else:
177
+ intermediate_size = config.moe_intermediate_size * config.n_shared_experts
178
+
179
+ self.shared_experts = DeepseekV4MLP(
180
+ hidden_size=config.hidden_size,
181
+ intermediate_size=intermediate_size,
182
+ hidden_act=config.hidden_act,
183
+ swiglu_limit=self.swiglu_limit,
184
+ quant_config=quant_config,
185
+ reduce_results=False,
186
+ prefix=f"{prefix}.shared_experts",
187
+ )
188
+
189
+ self.tp_rank = get_tensor_model_parallel_rank()
190
+ assert config.n_routed_experts % self.tp_size == 0
191
+
192
+ self.n_local_experts = config.n_routed_experts // self.tp_size
193
+ self.experts_start_idx = self.tp_rank * self.n_local_experts
194
+ self.experts_end_idx = self.experts_start_idx + self.n_local_experts
195
+
196
+ self.experts = FusedMoEFactory(
197
+ shared_experts=self.shared_experts,
198
+ n_shared_experts=(
199
+ config.n_shared_experts if self.is_fused_shared_expert_enabled else None
200
+ ),
201
+ fuse_shared_experts=self.is_fused_shared_expert_enabled,
202
+ gate=self.gate,
203
+ num_experts=config.n_routed_experts,
204
+ top_k=config.num_experts_per_tok,
205
+ hidden_size=config.hidden_size,
206
+ intermediate_size=config.moe_intermediate_size,
207
+ renormalize=config.norm_topk_prob,
208
+ quant_config=quant_config,
209
+ prefix=f"{prefix}.experts",
210
+ scoring_func=self.scoring_func,
211
+ routed_scaling_factor=self.routed_scaling_factor,
212
+ e_score_correction_bias=self.gate.e_score_correction_bias,
213
+ hash_indices_table=self.gate.tid2eid,
214
+ swiglu_limit=self.swiglu_limit,
215
+ router_logits_dtype=torch.float32,
216
+ )
217
+ if is_hash_moe:
218
+ # apply_monolithic() never sees hash_indices_table via the
219
+ # router path (modular-only) -- stash it on RoutedExperts
220
+ # directly. object.__setattr__ avoids nn.Module re-registering
221
+ # (and duplicating in the state_dict) this alias.
222
+ object.__setattr__(
223
+ self.experts.routed_experts, "hash_indices_table", self.gate.tid2eid
224
+ )
225
+
226
+ def forward(
227
+ self, hidden_states: torch.Tensor, input_ids: torch.Tensor | None = None
228
+ ) -> torch.Tensor:
229
+ if self.gate.tid2eid is not None and input_ids is None:
230
+ raise ValueError("DeepSeek V4 hash MoE routing requires input_ids.")
231
+
232
+ org_shape = hidden_states.shape
233
+ final_hidden_states = self.experts(
234
+ hidden_states=hidden_states,
235
+ router_logits=hidden_states,
236
+ input_ids=input_ids,
237
+ )
238
+
239
+ return final_hidden_states.view(org_shape)
240
+
241
+
242
+class DeepseekV4DecoderLayer(nn.Module):
243
+ def __init__(
244
+ self,
245
+ vllm_config,
246
+ prefix,
247
+ topk_indices_buffer: torch.Tensor | None = None,
248
+ ):
249
+ super().__init__()
250
+
251
+ # Lazy import to avoid top-level tilelang dependency.
252
+ # Registers both torch.ops.vllm.mhc_pre and mhc_post
253
+ import vllm.model_executor.layers.mhc # noqa: F401
254
+
255
+ config = vllm_config.model_config.hf_config
256
+ self.hidden_size = config.hidden_size
257
+
258
+ self.rms_norm_eps = config.rms_norm_eps
259
+ self.attn = DeepseekV4CPUAttention(
260
+ vllm_config,
261
+ prefix=f"{prefix}.attn",
262
+ topk_indices_buffer=topk_indices_buffer,
263
+ )
264
+ self.ffn = DeepseekV4MoE(vllm_config, prefix=f"{prefix}.ffn")
265
+
266
+ self.attn_norm = RMSNorm(self.hidden_size, self.rms_norm_eps)
267
+ self.ffn_norm = RMSNorm(self.hidden_size, self.rms_norm_eps)
268
+ self.hc_mult = config.hc_mult
269
+ self.hc_sinkhorn_iters = config.hc_sinkhorn_iters
270
+ self.hc_eps = config.hc_eps
271
+ self.hc_post_alpha = 2.0
272
+ mix_hc = (2 + self.hc_mult) * self.hc_mult
273
+ hc_dim = self.hc_mult * self.hidden_size
274
+ self.hc_attn_fn = nn.Parameter(
275
+ torch.empty(
276
+ (mix_hc, hc_dim),
277
+ dtype=torch.float32,
278
+ ),
279
+ requires_grad=False,
280
+ )
281
+ self.hc_ffn_fn = nn.Parameter(
282
+ torch.empty(
283
+ (mix_hc, hc_dim),
284
+ dtype=torch.float32,
285
+ ),
286
+ requires_grad=False,
287
+ )
288
+ self.hc_attn_base = nn.Parameter(
289
+ torch.empty(
290
+ mix_hc,
291
+ dtype=torch.float32,
292
+ ),
293
+ requires_grad=False,
294
+ )
295
+ self.hc_ffn_base = nn.Parameter(
296
+ torch.empty(
297
+ mix_hc,
298
+ dtype=torch.float32,
299
+ ),
300
+ requires_grad=False,
301
+ )
302
+ self.hc_attn_scale = nn.Parameter(
303
+ torch.empty(
304
+ 3,
305
+ dtype=torch.float32,
306
+ ),
307
+ requires_grad=False,
308
+ )
309
+ self.hc_ffn_scale = nn.Parameter(
310
+ torch.empty(
311
+ 3,
312
+ dtype=torch.float32,
313
+ ),
314
+ requires_grad=False,
315
+ )
316
+ self.mhc_pre = MHCPreOp()
317
+ self.mhc_post = MHCPostOp()
318
+ self.mhc_fused_post_pre = MHCFusedPostPreOp()
319
+ self.use_fused_mhc = HAS_TILELANG_MHC and not (
320
+ HAS_AITER_MHC and self.hidden_size % 256 == 0
321
+ )
322
+
323
+ def hc_pre(
324
+ self,
325
+ x: torch.Tensor,
326
+ hc_fn: torch.Tensor,
327
+ hc_scale: torch.Tensor,
328
+ hc_base: torch.Tensor,
329
+ ):
330
+ post_mix, res_mix, layer_input = self.mhc_pre(
331
+ residual=x,
332
+ fn=hc_fn,
333
+ hc_scale=hc_scale,
334
+ hc_base=hc_base,
335
+ rms_eps=self.rms_norm_eps,
336
+ hc_pre_eps=self.hc_eps,
337
+ hc_sinkhorn_eps=self.hc_eps,
338
+ hc_post_mult_value=self.hc_post_alpha,
339
+ sinkhorn_repeat=self.hc_sinkhorn_iters,
340
+ )
341
+ return layer_input, post_mix, res_mix
342
+
343
+ def hc_post(
344
+ self,
345
+ x: torch.Tensor,
346
+ residual: torch.Tensor,
347
+ post: torch.Tensor,
348
+ comb: torch.Tensor,
349
+ ):
350
+ return self.mhc_post(x, residual, post, comb)
351
+
352
+ def _forward_fused_post_pre(
353
+ self,
354
+ x: torch.Tensor,
355
+ positions: torch.Tensor,
356
+ input_ids: torch.Tensor | None,
357
+ post_mix: torch.Tensor | None = None,
358
+ res_mix: torch.Tensor | None = None,
359
+ residual: torch.Tensor | None = None,
360
+ ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
361
+ if residual is None:
362
+ # Run standalone hc_pre on first layer
363
+ residual = x
364
+ x, post_mix, res_mix = self.hc_pre(
365
+ x, self.hc_attn_fn, self.hc_attn_scale, self.hc_attn_base
366
+ )
367
+ else:
368
+ residual, post_mix, res_mix, x = self.mhc_fused_post_pre(
369
+ x,
370
+ residual,
371
+ post_mix,
372
+ res_mix,
373
+ self.hc_attn_fn,
374
+ self.hc_attn_scale,
375
+ self.hc_attn_base,
376
+ self.rms_norm_eps,
377
+ self.hc_eps,
378
+ self.hc_eps,
379
+ self.hc_post_alpha,
380
+ self.hc_sinkhorn_iters,
381
+ )
382
+
383
+ x = self.attn_norm(x)
384
+ x = self.attn(positions, x, None)
385
+
386
+ residual, post_mix, res_mix, x = self.mhc_fused_post_pre(
387
+ x,
388
+ residual,
389
+ post_mix,
390
+ res_mix,
391
+ self.hc_ffn_fn,
392
+ self.hc_ffn_scale,
393
+ self.hc_ffn_base,
394
+ self.rms_norm_eps,
395
+ self.hc_eps,
396
+ self.hc_eps,
397
+ self.hc_post_alpha,
398
+ self.hc_sinkhorn_iters,
399
+ )
400
+ x = self.ffn_norm(x)
401
+ x = self.ffn(x, input_ids)
402
+ return x, residual, post_mix, res_mix
403
+
404
+ def _forward_unfused_post_pre(
405
+ self,
406
+ x: torch.Tensor,
407
+ positions: torch.Tensor,
408
+ input_ids: torch.Tensor | None,
409
+ post_mix: torch.Tensor | None = None,
410
+ res_mix: torch.Tensor | None = None,
411
+ residual: torch.Tensor | None = None,
412
+ ) -> tuple[
413
+ torch.Tensor, torch.Tensor | None, torch.Tensor | None, torch.Tensor | None
414
+ ]:
415
+ residual = x
416
+ x, post, comb = self.hc_pre(
417
+ x, self.hc_attn_fn, self.hc_attn_scale, self.hc_attn_base
418
+ )
419
+ x = self.attn_norm(x)
420
+ x = self.attn(positions, x, None)
421
+ x = self.hc_post(x, residual, post, comb)
422
+
423
+ residual = x
424
+ x, post, comb = self.hc_pre(
425
+ x, self.hc_ffn_fn, self.hc_ffn_scale, self.hc_ffn_base
426
+ )
427
+ x = self.ffn_norm(x)
428
+ x = self.ffn(x, input_ids)
429
+ x = self.hc_post(x, residual, post, comb)
430
+ return x, None, None, None
431
+
432
+ def forward(
433
+ self,
434
+ x: torch.Tensor,
435
+ positions: torch.Tensor,
436
+ input_ids: torch.Tensor | None,
437
+ post_mix: torch.Tensor | None = None,
438
+ res_mix: torch.Tensor | None = None,
439
+ residual: torch.Tensor | None = None,
440
+ ) -> tuple[
441
+ torch.Tensor, torch.Tensor | None, torch.Tensor | None, torch.Tensor | None
442
+ ]:
443
+ if not self.use_fused_mhc:
444
+ return self._forward_unfused_post_pre(
445
+ x, positions, input_ids, post_mix, res_mix, residual
446
+ )
447
+ return self._forward_fused_post_pre(
448
+ x, positions, input_ids, post_mix, res_mix, residual
449
+ )
450
+
451
+
452
+@support_torch_compile(
453
+ dynamic_arg_dims={
454
+ "input_ids": 0,
455
+ "positions": -1,
456
+ "intermediate_tensors": 0,
457
+ "inputs_embeds": 0,
458
+ }
459
+)
460
+class DeepseekV4Model(nn.Module, EagleModelMixin):
461
+ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
462
+ super().__init__()
463
+
464
+ config = vllm_config.model_config.hf_config
465
+ quant_config = vllm_config.quant_config
466
+ self.config = config
467
+ self.vocab_size = config.vocab_size
468
+ self.hc_eps = config.hc_eps
469
+ self.hc_mult = config.hc_mult
470
+ self.hc_dim = self.hc_mult * config.hidden_size
471
+ self.rms_norm_eps = config.rms_norm_eps
472
+
473
+ self.device = current_platform.device_type
474
+ # Reserved topk indices buffer for all Indexer layers to reuse.
475
+ self.topk_indices_buffer = torch.empty(
476
+ vllm_config.scheduler_config.max_num_batched_tokens,
477
+ config.index_topk,
478
+ dtype=torch.int32,
479
+ device=self.device,
480
+ )
481
+
482
+ if get_pp_group().is_first_rank:
483
+ self.embed_tokens = VocabParallelEmbedding(
484
+ config.vocab_size,
485
+ config.hidden_size,
486
+ quant_config=quant_config,
487
+ prefix=f"{prefix}.embed_tokens",
488
+ )
489
+ else:
490
+ self.embed_tokens = PPMissingLayer()
491
+
492
+ self.start_layer, self.end_layer, self.layers = make_layers(
493
+ config.num_hidden_layers,
494
+ lambda prefix: DeepseekV4DecoderLayer(
495
+ vllm_config,
496
+ prefix=prefix,
497
+ topk_indices_buffer=self.topk_indices_buffer,
498
+ ),
499
+ prefix=f"{prefix}.layers",
500
+ )
501
+ self.is_fused_shared_expert_enabled = is_model_fused_shared_expert_compatible(
502
+ self.layers,
503
+ DeepseekV4MoE,
504
+ "ffn",
505
+ )
506
+
507
+ if get_pp_group().is_last_rank:
508
+ self.norm = RMSNorm(config.hidden_size, self.rms_norm_eps)
509
+ else:
510
+ self.norm = PPMissingLayer()
511
+
512
+ self.hc_head_fn = nn.Parameter(
513
+ torch.empty(
514
+ self.hc_mult,
515
+ self.hc_dim,
516
+ dtype=torch.float32,
517
+ ),
518
+ requires_grad=False,
519
+ )
520
+ self.hc_head_base = nn.Parameter(
521
+ torch.empty(
522
+ self.hc_mult,
523
+ dtype=torch.float32,
524
+ ),
525
+ requires_grad=False,
526
+ )
527
+ self.hc_head_scale = nn.Parameter(
528
+ torch.empty(1, dtype=torch.float32),
529
+ requires_grad=False,
530
+ )
531
+ self.hc_head_op = HCHeadOp()
532
+ spec_config = vllm_config.speculative_config
533
+ needs_mtp_hidden_states = spec_config is not None and (
534
+ spec_config.use_eagle() or spec_config.uses_draft_model()
535
+ )
536
+ if get_pp_group().is_last_rank and needs_mtp_hidden_states:
537
+ self._mtp_hidden_buffer = torch.empty(
538
+ vllm_config.scheduler_config.max_num_batched_tokens,
539
+ self.hc_dim,
540
+ dtype=vllm_config.model_config.dtype,
541
+ device=self.device,
542
+ )
543
+ else:
544
+ self._mtp_hidden_buffer = None
545
+
546
+ def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
547
+ return self.embed_tokens(input_ids)
548
+
549
+ def make_empty_intermediate_tensors(
550
+ self,
551
+ batch_size: int,
552
+ dtype: torch.dtype,
553
+ device: torch.device,
554
+ ) -> IntermediateTensors:
555
+ # PP intermediate tensors keep the (num_tokens, hc_mult, hidden_size)
556
+ # multi-stream shape until hc_head() collapses it.
557
+ return IntermediateTensors(
558
+ {
559
+ "hidden_states": torch.zeros(
560
+ (batch_size, self.hc_mult, self.config.hidden_size),
561
+ dtype=dtype,
562
+ device=device,
563
+ ),
564
+ }
565
+ )
566
+
567
+ def forward(
568
+ self,
569
+ input_ids: torch.Tensor,
570
+ positions: torch.Tensor,
571
+ intermediate_tensors: IntermediateTensors | None,
572
+ inputs_embeds: torch.Tensor | None = None,
573
+ ) -> torch.Tensor | IntermediateTensors:
574
+ if get_pp_group().is_first_rank:
575
+ if inputs_embeds is not None:
576
+ hidden_states = inputs_embeds
577
+ else:
578
+ hidden_states = self.embed_input_ids(input_ids)
579
+ hidden_states = hidden_states.unsqueeze(-2).repeat(1, self.hc_mult, 1)
580
+ else:
581
+ assert intermediate_tensors is not None
582
+ hidden_states = intermediate_tensors["hidden_states"]
583
+
584
+ residual, post_mix, res_mix = None, None, None
585
+ # EAGLE3/DSpark/DFlash aux hidden states: post-mhc hidden state at
586
+ # each target layer, averaged over hc_mult streams to [T,
587
+ # hidden_size]. Empty unless a draft model set aux_hidden_state_layers.
588
+ aux_hidden_states: list[torch.Tensor] = []
589
+ # Reused below to avoid computing hc_post twice if the last layer is
590
+ # also an aux layer.
591
+ final_aux_recon: torch.Tensor | None = None
592
+ for idx, layer in enumerate(
593
+ islice(self.layers, self.start_layer, self.end_layer),
594
+ start=self.start_layer,
595
+ ):
596
+ hidden_states, residual, post_mix, res_mix = layer(
597
+ hidden_states,
598
+ positions,
599
+ input_ids,
600
+ post_mix,
601
+ res_mix,
602
+ residual,
603
+ )
604
+ if (idx + 1) in self.aux_hidden_state_layers:
605
+ # Unfused path already applied hc_post; fused path needs it
606
+ # applied here before averaging.
607
+ if layer.use_fused_mhc:
608
+ aux_recon = layer.hc_post(
609
+ hidden_states, residual, post_mix, res_mix
610
+ )
611
+ final_aux_recon = aux_recon
612
+ else:
613
+ aux_recon = hidden_states
614
+ aux_hidden_states.append(aux_recon.mean(dim=1))
615
+ if layer is not None and layer.use_fused_mhc:
616
+ # Reuse the last layer's hc_post output if already computed above.
617
+ if (
618
+ final_aux_recon is not None
619
+ and self.end_layer in self.aux_hidden_state_layers
620
+ ):
621
+ hidden_states = final_aux_recon
622
+ else:
623
+ hidden_states = layer.hc_post(
624
+ hidden_states, residual, post_mix, res_mix
625
+ )
626
+
627
+ if not get_pp_group().is_last_rank:
628
+ return IntermediateTensors({"hidden_states": hidden_states})
629
+
630
+ if self._mtp_hidden_buffer is not None:
631
+ num_tokens = hidden_states.shape[0]
632
+ self._mtp_hidden_buffer[:num_tokens].copy_(hidden_states.flatten(1))
633
+
634
+ hidden_states = self.hc_head_op(
635
+ hidden_states,
636
+ self.hc_head_fn,
637
+ self.hc_head_scale,
638
+ self.hc_head_base,
639
+ self.rms_norm_eps,
640
+ self.hc_eps,
641
+ )
642
+ hidden_states = self.norm(hidden_states)
643
+ if len(aux_hidden_states) > 0:
644
+ return hidden_states, aux_hidden_states
645
+ return hidden_states
646
+
647
+ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
648
+ stacked_params_mapping = [
649
+ # (param_name, shard_name, shard_id)
650
+ ("gate_up_proj", "w1", 0),
651
+ ("gate_up_proj", "w3", 1),
652
+ ("attn.fused_wqa_wkv", "attn.wq_a", 0),
653
+ ("attn.fused_wqa_wkv", "attn.wkv", 1),
654
+ ("compressor.fused_wkv_wgate", "compressor.wkv", 0),
655
+ ("compressor.fused_wkv_wgate", "compressor.wgate", 1),
656
+ ]
657
+ params_dict = dict(self.named_parameters())
658
+ loaded_params: set[str] = set()
659
+
660
+ def _resolve_param_name(name: str) -> str:
661
+ inv_name = f"{name}_inv"
662
+ if name not in params_dict and inv_name in params_dict:
663
+ return inv_name
664
+ return name
665
+
666
+ # TP for attention
667
+ tp_size = get_tensor_model_parallel_world_size()
668
+ tp_rank = get_tensor_model_parallel_rank()
669
+ n_head = self.config.num_attention_heads
670
+ n_local_head = n_head // tp_size
671
+ head_rank_start = n_local_head * tp_rank
672
+ head_rank_end = n_local_head * (tp_rank + 1)
673
+
674
+ # Pre-compute expert mapping ONCE.
675
+ expert_mapping = self.get_expert_mapping()
676
+
677
+ # Each MoE's own per-layer fusion decision is the source of truth,
678
+ # so the redirect below can't diverge from how it was built.
679
+ fuse_by_layer = {
680
+ extract_layer_index(mod_name): mod.is_fused_shared_expert_enabled
681
+ for mod_name, mod in self.named_modules()
682
+ if isinstance(mod, DeepseekV4MoE)
683
+ }
684
+ n_routed = self.config.n_routed_experts
685
+ # Maps the single shared-expert tensor group to one appended slot;
686
+ # multiple shared experts would need per-expert slicing (see
687
+ # deepseek_v2.py). DeepSeek-V4 has n_shared_experts == 1.
688
+ if any(fuse_by_layer.values()) and self.config.n_shared_experts != 1:
689
+ raise NotImplementedError(
690
+ "deepseek-v4 fused shared-expert loading supports only "
691
+ f"n_shared_experts == 1, got {self.config.n_shared_experts}"
692
+ )
693
+
694
+ for name, loaded_weight in weights:
695
+ # Shared-expert fusion: redirect shared_experts.w{1,2,3} into
696
+ # the appended routed-expert slot so it loads through the
697
+ # routed expert (grouped GEMM) loader.
698
+ if ".ffn.shared_experts.w" in name and fuse_by_layer.get(
699
+ extract_layer_index(name), False
700
+ ):
701
+ name = name.replace(
702
+ ".ffn.shared_experts.w",
703
+ f".ffn.experts.{n_routed}.w",
704
+ )
705
+
706
+ for param_name, weight_name, shard_id in stacked_params_mapping:
707
+ # Skip non-stacked layers and experts (experts handled below).
708
+ if ".experts." in name:
709
+ continue
710
+ if weight_name not in name:
711
+ continue
712
+ name = name.replace(weight_name, param_name)
713
+
714
+ if is_pp_missing_parameter(name, self):
715
+ break
716
+ param_name = _resolve_param_name(name)
717
+ param = params_dict[param_name]
718
+ weight_loader = param.weight_loader
719
+ weight_loader(param, loaded_weight, shard_id)
720
+ loaded_params.add(param_name)
721
+ break
722
+ else:
723
+ if ".experts." in name:
724
+ # E8M0 scales are float8_e8m0fnu in checkpoints but the
725
+ # MoE param is uint8; copy_() would numerically convert
726
+ # (e.g. 2^-7 -> 0) instead of preserving raw bytes.
727
+ if (
728
+ "weight_scale" in name
729
+ and loaded_weight.dtype == torch.float8_e8m0fnu
730
+ ):
731
+ loaded_weight = loaded_weight.view(torch.uint8)
732
+ for mapping in expert_mapping:
733
+ param_name, weight_name, expert_id, expert_shard_id = mapping
734
+ if weight_name not in name:
735
+ continue
736
+ name_mapped = name.replace(weight_name, param_name)
737
+ if is_pp_missing_parameter(name_mapped, self):
738
+ continue
739
+ param = params_dict[name_mapped]
740
+ # return_success lets us skip experts that have
741
+ # another replica available.
742
+ weight_loader = typing.cast(
743
+ Callable[..., bool], param.weight_loader
744
+ )
745
+ success = weight_loader(
746
+ param,
747
+ loaded_weight,
748
+ name_mapped,
749
+ shard_id=expert_shard_id,
750
+ expert_id=expert_id,
751
+ return_success=True,
752
+ )
753
+ if success:
754
+ name = name_mapped
755
+ break
756
+ loaded_params.add(name_mapped)
757
+ continue
758
+ elif "attn_sink" in name:
759
+ if is_pp_missing_parameter(name, self):
760
+ continue
761
+ narrow_weight = loaded_weight[head_rank_start:head_rank_end]
762
+ n = narrow_weight.shape[0]
763
+ params_dict[name][:n].copy_(narrow_weight)
764
+ loaded_params.add(name)
765
+ continue
766
+ else:
767
+ if is_pp_missing_parameter(name, self):
768
+ continue
769
+ param_name = _resolve_param_name(name)
770
+ param = params_dict[param_name]
771
+ weight_loader = getattr(
772
+ param, "weight_loader", default_weight_loader
773
+ )
774
+ weight_loader(param, loaded_weight)
775
+ loaded_params.add(param_name)
776
+ continue
777
+
778
+ return loaded_params
779
+
780
+ def get_expert_mapping(self) -> list[tuple[str, str, int, str]]:
781
+ # (param_name, weight_name, expert_id, shard_id). When fusing shared
782
+ # experts, include the appended slots so redirected shared-expert
783
+ # weights route through the expert loader.
784
+ n_shared = getattr(self.config, "n_shared_experts", 0) or 0
785
+ num_experts = self.config.n_routed_experts + (
786
+ n_shared if self.is_fused_shared_expert_enabled else 0
787
+ )
788
+ return fused_moe_make_expert_params_mapping(
789
+ self,
790
+ ckpt_gate_proj_name="w1",
791
+ ckpt_down_proj_name="w2",
792
+ ckpt_up_proj_name="w3",
793
+ num_experts=num_experts,
794
+ )
795
+
796
+
797
+def _make_deepseek_v4_weights_mapper(
798
+ expert_dtype: str, fuse_shared_experts: bool = False
799
+) -> WeightsMapper:
800
+ if expert_dtype == "fp4":
801
+ # MXFP4 experts register scales as ``w{1,2,3}_weight_scale`` (no
802
+ # _inv); FP8 linear/shared-expert scales register as
803
+ # ``weight_scale_inv``. Route DeepSeek-native ``.scale`` and
804
+ # AMD-Quark ``.weight_scale`` keys accordingly, leaving expert/
805
+ # shared-expert ``w{1,2,3}.weight_scale`` untouched for the MXFP4
806
+ # loader.
807
+ scale_regex = {
808
+ re.compile(r"(\.experts\.\d+\.w[123])\.scale$"): r"\1.weight_scale",
809
+ re.compile(r"\.scale$"): ".weight_scale_inv",
810
+ re.compile(r"(?<!\.w[123])\.weight_scale$"): ".weight_scale_inv",
811
+ }
812
+ else:
813
+ # FP8 experts use Fp8MoEMethod (block_quant=True), which registers
814
+ # scales as ``w{13,2}_weight_scale_inv``. Map all ``.scale`` keys
815
+ # there.
816
+ scale_regex = {
817
+ re.compile(r"\.scale$"): ".weight_scale_inv",
818
+ }
819
+ # Fused shared experts are redirected to routed expert slots elsewhere;
820
+ # leave their names untouched here.
821
+ substr_map: dict[str, str | None] = (
822
+ {}
823
+ if fuse_shared_experts
824
+ else {".shared_experts.w2": ".shared_experts.down_proj"}
825
+ )
826
+ # No MTP support on CPU yet: drop mtp.* weights (no model.mtp submodule).
827
+ substr_map["mtp."] = None
828
+ return WeightsMapper(
829
+ orig_to_new_prefix={
830
+ "layers.": "model.layers.",
831
+ "embed.": "model.embed.",
832
+ "norm.": "model.norm.",
833
+ "hc_head": "model.hc_head",
834
+ "mtp.": "model.mtp.",
835
+ },
836
+ orig_to_new_regex=scale_regex,
837
+ orig_to_new_suffix={
838
+ "head.weight": "lm_head.weight",
839
+ "embed.weight": "embed_tokens.weight",
840
+ ".ffn.gate.bias": ".ffn.gate.e_score_correction_bias",
841
+ ".input_scale": ".input_scale_2",
842
+ },
843
+ orig_to_new_substr=substr_map,
844
+ )
845
+
846
+
847
+class DeepseekV4ForCausalLM(nn.Module, SupportsPP, SupportsEagle3):
848
+ model_cls = DeepseekV4Model
849
+
850
+ # Default mapper assumes the original FP4-expert checkpoint layout.
851
+ # Overridden per-instance in __init__ when expert_dtype != "fp4".
852
+ hf_to_vllm_mapper = _make_deepseek_v4_weights_mapper("fp4")
853
+ packed_modules_mapping = {
854
+ "gate_up_proj": ["w1", "w3"],
855
+ "fused_wqa_wkv": ["wq_a", "wkv"],
856
+ }
857
+
858
+ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
859
+ super().__init__()
860
+
861
+ config = vllm_config.model_config.hf_config
862
+ self.config = config
863
+ expert_dtype = getattr(config, "expert_dtype", "fp4")
864
+ self.model = self.model_cls(
865
+ vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model")
866
+ )
867
+ if expert_dtype != "fp4" or self.model.is_fused_shared_expert_enabled:
868
+ self.hf_to_vllm_mapper = _make_deepseek_v4_weights_mapper(
869
+ expert_dtype,
870
+ fuse_shared_experts=self.model.is_fused_shared_expert_enabled,
871
+ )
872
+ if get_pp_group().is_last_rank:
873
+ self.lm_head = ParallelLMHead(
874
+ config.vocab_size,
875
+ config.hidden_size,
876
+ prefix=maybe_prefix(prefix, "lm_head"),
877
+ )
878
+ else:
879
+ self.lm_head = PPMissingLayer()
880
+ self.logits_processor = LogitsProcessor(config.vocab_size)
881
+ self.make_empty_intermediate_tensors = ( # type: ignore[method-assign]
882
+ self.model.make_empty_intermediate_tensors
883
+ )
884
+
885
+ def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
886
+ return self.model.embed_input_ids(input_ids)
887
+
888
+ def compute_logits(
889
+ self,
890
+ hidden_states: torch.Tensor,
891
+ ) -> torch.Tensor | None:
892
+ logits = self.logits_processor(self.lm_head, hidden_states)
893
+ return logits
894
+
895
+ def forward(
896
+ self,
897
+ input_ids: torch.Tensor,
898
+ positions: torch.Tensor,
899
+ intermediate_tensors: IntermediateTensors | None = None,
900
+ inputs_embeds: torch.Tensor | None = None,
901
+ ) -> torch.Tensor | IntermediateTensors:
902
+ hidden_states = self.model(
903
+ input_ids, positions, intermediate_tensors, inputs_embeds
904
+ )
905
+ return hidden_states
906
+
907
+ def get_mtp_target_hidden_states(self) -> torch.Tensor | None:
908
+ """Pre-hc_head residual stream buffer (max_num_batched_tokens,
909
+ hc_mult * hidden_size) for the MTP draft model. Populated by
910
+ forward(); valid after each target step."""
911
+ return getattr(self.model, "_mtp_hidden_buffer", None)
912
+
913
+ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
914
+ loader = AutoWeightsLoader(self)
915
+ return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)
916
+
917
+ def process_weights_after_loading(self) -> None:
918
+ # No CPU AMX weight-packing hook yet.
919
+ pass
920
+
921
+ def get_expert_mapping(self) -> list[tuple[str, str, int, str]]:
922
+ return self.model.get_expert_mapping()
vllm/models/deepseek_v4/cpu/mtp.py
new
+19
@@ -0,0 +1,19 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""CPU DeepSeek-V4 MTP (speculative decoding) — deferred, not implemented.
4
+
5
+``DeepSeekV4MTP`` is only instantiated when
6
+``vllm/model_executor/models/registry.py``'s speculative-decoding registry
7
+resolves it by name for a speculative-decoding config; a plain (non-spec)
8
+DeepSeek-V4 CPU model never constructs this class. This stub exists solely
9
+so ``vllm.models.deepseek_v4.cpu`` is import-clean.
10
+"""
11
+
12
+from vllm.config import VllmConfig
13
+
14
+
15
+class DeepSeekV4MTP:
16
+ def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None:
17
+ raise NotImplementedError(
18
+ "DeepSeek-V4 speculative decoding (MTP) is not implemented on CPU."
19
+ )
vllm/platforms/cpu.py
+46
-4
@@ -203,14 +203,26 @@ class CpuPlatform(Platform):
203
204
cache_config = vllm_config.cache_config
205
206
+ is_deepseek_v4 = (
207
+ model_config is not None
208
+ and getattr(model_config.hf_config, "model_type", None) == "deepseek_v4"
209
+ )
210
+
211
# The CPU MLA decode kernel only compiles with block_size=16 today
212
# (see csrc/cpu/mla_decode.cpp). If the model uses MLA we override
213
# the default block size regardless of user preference to avoid a
214
# runtime kernel dispatch failure. AMX MLA has no such constraint
215
# (same AMX-available condition as get_attn_backend_cls), so it's
211
- # excluded from this override.
212
- cpu_mla_enabled = model_config is not None and getattr(
213
- model_config, "use_mla", False
216
+ # excluded from this override. DeepSeek-V4 is also excluded here and
217
+ # handled in its own branch below: its sparse-MLA and indexer
218
+ # backends declare block_size=256 as their only supported kernel
219
+ # block size (DeepseekV4SparseMLABackend/DeepseekV4IndexerBackend),
220
+ # so it needs the same override-regardless-of-preference treatment
221
+ # as CPU MLA, just with a different value.
222
+ cpu_mla_enabled = (
223
+ not is_deepseek_v4
224
+ and model_config is not None
225
+ and getattr(model_config, "use_mla", False)
226
)
227
amx_mla_enabled = (
228
cpu_mla_enabled
@@ -219,6 +231,21 @@ class CpuPlatform(Platform):
231
and vllm_config.attention_config.backend != AttentionBackendEnum.CPU_MLA
232
)
233
reference_cpu_mla_enabled = cpu_mla_enabled and not amx_mla_enabled
234
+ # DeepSeek-V4's CPU attention/indexer kernels
235
+ # (csrc/cpu/sgl-kernels/{flash_mla,store_cache,compressor,
236
+ # paged_mqa_logits,topk}.cpp) are AMX-kernel-backed and built on the
237
+ # same paged/position-indexed conventions as the GPU/XPU backends
238
+ # (block_table/slot_mapping addressing throughout, chunk metadata
239
+ # that already carries per-token causal offsets for partial/extend
240
+ # continuation) -- so they support chunked prefill and prefix
241
+ # caching the same way. `amx_mla_enabled` above deliberately
242
+ # excludes DeepSeek-V4 (it has its own block-size requirements,
243
+ # unrelated to chunked-prefill support), so it can't be reused here.
244
+ amx_mla_or_dsv4_enabled = amx_mla_enabled or (
245
+ is_deepseek_v4
246
+ and cls.get_cpu_architecture() == CpuArchEnum.X86
247
+ and torch.cpu._is_amx_tile_supported()
248
+ )
249
if reference_cpu_mla_enabled:
250
if cache_config.user_specified_block_size and cache_config.block_size != 16:
251
logger.warning(
@@ -227,6 +254,17 @@ class CpuPlatform(Platform):
254
cache_config.block_size,
255
)
256
cache_config.block_size = 16
257
+ elif is_deepseek_v4:
258
+ if (
259
+ cache_config.user_specified_block_size
260
+ and cache_config.block_size != 256
261
+ ):
262
+ logger.warning(
263
+ "DeepSeek-V4 CPU backend requires block_size=256, "
264
+ "overriding user-specified block_size=%s.",
265
+ cache_config.block_size,
266
+ )
267
+ cache_config.block_size = 256
268
elif not cache_config.user_specified_block_size:
269
cache_config.block_size = 128
270
@@ -439,7 +477,11 @@ class CpuPlatform(Platform):
477
vllm_config.parallel_config.tensor_parallel_size
478
)
479
442
- if model_config is not None and model_config.use_mla and not amx_mla_enabled:
480
+ if (
481
+ model_config is not None
482
+ and model_config.use_mla
483
+ and not amx_mla_or_dsv4_enabled
484
+ ):
485
logger.info_once(
486
"MLA is enabled on a non-GPU platform; forcing chunked "
487
"prefill and prefix caching to be disabled."
vllm/v1/attention/backends/mla/indexer.py
+58
-50
@@ -177,55 +177,6 @@ class PrepareUniformDecodeKernel(
177
)
178
179
180
-def split_indexer_prefill_chunks(
181
- seq_lens_cpu: torch.Tensor,
182
- query_lens_cpu: torch.Tensor,
183
- workspace_size: int,
184
- max_logits_bytes: int,
185
- request_offset: int = 0,
186
-) -> list[tuple[slice, slice]]:
187
- """
188
- Split prefill requests into chunks for the sparse indexer, respecting:
189
- - N constraint: total_seq_lens <= workspace_size (existing O(N) workspace)
190
- - Logits constraint: M * N * 4 <= max_logits_bytes
191
-
192
- When a single request-level chunk still exceeds the logits budget,
193
- sub-chunks on the query dimension (M) to bound peak memory.
194
-
195
- Returns list of (req_slice, query_slice) tuples.
196
- """
197
- chunks: list[tuple[slice, slice]] = []
198
- n = len(seq_lens_cpu)
199
- max_logits_elems = max_logits_bytes // 4
200
- end = 0
201
-
202
- while end < n:
203
- start, chunk_m, chunk_n = end, 0, 0
204
-
205
- while end < n:
206
- q, s = query_lens_cpu[end].item(), seq_lens_cpu[end].item()
207
- new_m, new_n = chunk_m + q, chunk_n + s
208
- if new_n <= workspace_size and new_m * new_n <= max_logits_elems:
209
- chunk_m, chunk_n = new_m, new_n
210
- end += 1
211
- else:
212
- break
213
-
214
- # A single request can exceed the budget, requiring sub-chunking
215
- # on the query dimension.
216
- if end == start:
217
- chunk_m, chunk_n = query_lens_cpu[end].item(), seq_lens_cpu[end].item()
218
- end += 1
219
-
220
- req_slice = slice(start + request_offset, end + request_offset)
221
- max_q = max(1, max_logits_elems // chunk_n) if chunk_n > 0 else max(1, chunk_m)
222
- for q_off in range(0, chunk_m, max_q):
223
- sub_m = min(max_q, chunk_m - q_off)
224
- chunks.append((req_slice, slice(q_off, q_off + sub_m)))
225
-
226
- return chunks
227
-
228
-
180
class DeepseekV32IndexerBackend(AttentionBackend):
181
@classmethod
182
def supports_pcp(cls) -> bool:
@@ -1020,6 +971,63 @@ class DeepseekV32IndexerMetadataBuilder(AttentionMetadataBuilder):
971
)
972
return indices
973
974
+ @staticmethod
975
+ def _split_indexer_prefill_chunks(
976
+ compressed_seq_lens_cpu: torch.Tensor,
977
+ prefill_query_lens_cpu: torch.Tensor,
978
+ workspace_size: int,
979
+ max_logits_bytes: int,
980
+ request_offset: int = 0,
981
+ ) -> list[tuple[slice, slice]]:
982
+ """Split this step's prefill requests into chunks, respecting:
983
+ - N constraint: total_seq_lens <= workspace_size (existing O(N)
984
+ workspace)
985
+ - Logits constraint: M * N * 4 <= max_logits_bytes
986
+
987
+ When a single request-level chunk still exceeds the logits budget,
988
+ sub-chunks on the query dimension (M) to bound peak memory.
989
+
990
+ Returns list of (req_slice, query_slice) tuples.
991
+ """
992
+ chunks: list[tuple[slice, slice]] = []
993
+ n = len(compressed_seq_lens_cpu)
994
+ max_logits_elems = max_logits_bytes // 4
995
+ end = 0
996
+
997
+ while end < n:
998
+ start, chunk_m, chunk_n = end, 0, 0
999
+
1000
+ while end < n:
1001
+ q, s = (
1002
+ prefill_query_lens_cpu[end].item(),
1003
+ compressed_seq_lens_cpu[end].item(),
1004
+ )
1005
+ new_m, new_n = chunk_m + q, chunk_n + s
1006
+ if new_n <= workspace_size and new_m * new_n <= max_logits_elems:
1007
+ chunk_m, chunk_n = new_m, new_n
1008
+ end += 1
1009
+ else:
1010
+ break
1011
+
1012
+ # A single request can exceed the budget, requiring sub-chunking
1013
+ # on the query dimension.
1014
+ if end == start:
1015
+ chunk_m, chunk_n = (
1016
+ prefill_query_lens_cpu[end].item(),
1017
+ compressed_seq_lens_cpu[end].item(),
1018
+ )
1019
+ end += 1
1020
+
1021
+ req_slice = slice(start + request_offset, end + request_offset)
1022
+ max_q = (
1023
+ max(1, max_logits_elems // chunk_n) if chunk_n > 0 else max(1, chunk_m)
1024
+ )
1025
+ for q_off in range(0, chunk_m, max_q):
1026
+ sub_m = min(max_q, chunk_m - q_off)
1027
+ chunks.append((req_slice, slice(q_off, q_off + sub_m)))
1028
+
1029
+ return chunks
1030
+
1031
def build(
1032
self,
1033
common_prefix_len: int,
@@ -1097,7 +1105,7 @@ class DeepseekV32IndexerMetadataBuilder(AttentionMetadataBuilder):
1105
# slice below).
1106
assert common_attn_metadata.seq_lens_cpu_upper_bound is not None
1107
seq_lens_cpu = common_attn_metadata.seq_lens_cpu_upper_bound
1100
- chunk_specs = split_indexer_prefill_chunks(
1108
+ chunk_specs = self._split_indexer_prefill_chunks(
1109
compressed_seq_lens_cpu[num_decodes:],
1110
prefill_query_lens_cpu,
1111
self.max_prefill_buffer_size,
vllm/v1/attention/backends/mla/sparse_swa.py
+6
@@ -148,6 +148,12 @@ class DeepseekSparseSWABackend(AttentionBackend):
148
)
149
150
return DeepseekV4ROCMAiterSparseSWAMetadataBuilder
151
+ if current_platform.is_cpu():
152
+ from vllm.models.deepseek_v4.cpu.cpu_mla import (
153
+ DeepseekV4CPUSparseSWAMetadataBuilder,
154
+ )
155
+
156
+ return DeepseekV4CPUSparseSWAMetadataBuilder
157
return DeepseekSparseSWAMetadataBuilder
158
159