diff --git a/.buildkite/hardware_tests/cpu.yaml b/.buildkite/hardware_tests/cpu.yaml index f4cb2ad0a7..769e3d4adc 100644 --- a/.buildkite/hardware_tests/cpu.yaml +++ b/.buildkite/hardware_tests/cpu.yaml @@ -1,7 +1,7 @@ group: CPU depends_on: [] steps: -- label: CPU-Kernel Tests +- label: CPU-Kernel Tests Shard %N key: cpu-kernel-tests depends_on: [] device: intel_cpu @@ -16,6 +16,7 @@ steps: - tests/kernels/attention/test_amx_mla.py - vllm/v1/attention/backends/mla/amx_mla.py - vllm/model_executor/layers/attention/mla_attention.py + - tests/kernels/test_deepseek_v4_cpu_kernels.py - tests/kernels/moe/test_cpu_fused_moe.py - tests/kernels/moe/test_cpu_quant_fused_moe.py - tests/kernels/moe/test_zen_cpu_int8_moe.py @@ -33,20 +34,23 @@ steps: commands: - | bash .buildkite/scripts/hardware_ci/run-cpu-test.sh 30m " - pytest -x -v -s tests/kernels/attention/test_cpu_attn.py - pytest -x -v -s tests/v1/attention/test_group_head_counts.py - pytest -x -v -s tests/kernels/attention/test_amx_mla.py - pytest -x -v -s tests/kernels/moe/test_cpu_fused_moe.py - pytest -x -v -s tests/kernels/moe/test_cpu_quant_fused_moe.py - pytest -x -v -s tests/kernels/moe/test_zen_cpu_int8_moe.py - pytest -x -v -s tests/kernels/mamba/test_cpu_short_conv.py - pytest -x -v -s tests/kernels/test_onednn.py - pytest -x -v -s tests/kernels/test_awq_int4_to_int8.py - pytest -x -v -s tests/kernels/quantization/test_cpu_fp8_scaled_mm.py - pytest -x -v -s tests/kernels/mamba/cpu/test_cpu_gdn_ops.py - pytest -x -v -s tests/kernels/mamba/test_causal_conv1d.py - pytest -x -v -s tests/kernels/mamba/test_mamba_ssm.py - pytest -x -v -s tests/v1/sample/test_topk_topp_sampler.py::TestTritonTopkTopp" + pytest -x -v -s --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB \ + tests/kernels/attention/test_cpu_attn.py \ + tests/v1/attention/test_group_head_counts.py \ + tests/kernels/attention/test_amx_mla.py \ + tests/kernels/test_deepseek_v4_cpu_kernels.py \ + tests/kernels/moe/test_cpu_fused_moe.py \ + tests/kernels/moe/test_cpu_quant_fused_moe.py \ + tests/kernels/moe/test_zen_cpu_int8_moe.py \ + tests/kernels/mamba/test_cpu_short_conv.py \ + tests/kernels/test_onednn.py \ + tests/kernels/test_awq_int4_to_int8.py \ + tests/kernels/quantization/test_cpu_fp8_scaled_mm.py \ + tests/kernels/mamba/cpu/test_cpu_gdn_ops.py \ + tests/kernels/mamba/test_causal_conv1d.py \ + tests/kernels/mamba/test_mamba_ssm.py \ + tests/v1/sample/test_topk_topp_sampler.py::TestTritonTopkTopp" + parallelism: 2 # Note: SDE can't be downloaded from CI host because of AWS WAF # - label: CPU-Compatibility Tests diff --git a/benchmarks/kernels/cpu/benchmark_cpu_fused_moe.py b/benchmarks/kernels/cpu/benchmark_cpu_fused_moe.py index f5a5ed1dc5..9b1bd5a7a0 100644 --- a/benchmarks/kernels/cpu/benchmark_cpu_fused_moe.py +++ b/benchmarks/kernels/cpu/benchmark_cpu_fused_moe.py @@ -70,12 +70,14 @@ def main( packed_w13 = cpu_prepack_moe_weight(w13, isa) packed_w2 = cpu_prepack_moe_weight(w2, isa) + output = torch.empty_like(input_tensor) def run_benchmark(iters: int) -> list[float]: times = [] for _ in range(iters): start_time = time.perf_counter_ns() - _ = cpu_fused_moe( + cpu_fused_moe( + output, input_tensor, packed_w13, packed_w2, diff --git a/cmake/cpu_extension.cmake b/cmake/cpu_extension.cmake index 35f56f242d..ff2911dc11 100644 --- a/cmake/cpu_extension.cmake +++ b/cmake/cpu_extension.cmake @@ -548,7 +548,14 @@ if (ENABLE_X86_ISA) "csrc/cpu/sgl-kernels/bmm.cpp" "csrc/cpu/sgl-kernels/decode.cpp" "csrc/cpu/sgl-kernels/extend.cpp" - "csrc/cpu/sgl-kernels/mla_cache.cpp") + "csrc/cpu/sgl-kernels/mla_cache.cpp" + "csrc/cpu/sgl-kernels/mhc.cpp" + "csrc/cpu/sgl-kernels/store_cache.cpp" + "csrc/cpu/sgl-kernels/flash_mla.cpp" + "csrc/cpu/sgl-kernels/compressor.cpp" + "csrc/cpu/sgl-kernels/paged_mqa_logits.cpp" + "csrc/cpu/sgl-kernels/topk.cpp" + "csrc/cpu/sgl-kernels/indexer.cpp") set(VLLM_EXT_SRC_AVX512 "csrc/cpu/sgl-kernels/fla.cpp" diff --git a/csrc/cpu/sgl-kernels/common.h b/csrc/cpu/sgl-kernels/common.h index cbd7e691f7..6f27ba2d5c 100644 --- a/csrc/cpu/sgl-kernels/common.h +++ b/csrc/cpu/sgl-kernels/common.h @@ -343,6 +343,36 @@ inline void parallel_2d(int m, int n, const func_t& f) { #endif } +// Like parallel_2d but with explicit nth_m x nth_n decomposition. +// Caller is responsible for choosing nth_m and nth_n. +// +// Cherry-picked from a newer sglang commit than this file's last full sync +// (needed by mhc.cpp's phase-1 tiling); not part of a full common.h resync. +template +inline void parallel_2d_tiled(int m, int n, int nth_m, int nth_n, const func_t& f) { + const int nth = nth_m * nth_n; +#if defined(_OPENMP) +#pragma omp parallel num_threads(nth) + { + int ith = omp_get_thread_num(); + int ith_m = ith / nth_n; + int ith_n = ith % nth_n; + + int thread_block_m = div_up(m, nth_m); + int thread_block_n = div_up(n, nth_n); + + int begin_m = std::min(ith_m * thread_block_m, m); + int end_m = std::min(begin_m + thread_block_m, m); + int begin_n = std::min(ith_n * thread_block_n, n); + int end_n = std::min(begin_n + thread_block_n, n); + + f(begin_m, end_m, begin_n, end_n); + } +#else + f(0, m, 0, n); +#endif +} + // limit max cache blocks // when we need to do pre-unpack for weights, e.g. fp8 #define MAX_CACHE_BLOCK_SIZE 4 diff --git a/csrc/cpu/sgl-kernels/compressor.cpp b/csrc/cpu/sgl-kernels/compressor.cpp new file mode 100644 index 0000000000..08a7e3f0f7 --- /dev/null +++ b/csrc/cpu/sgl-kernels/compressor.cpp @@ -0,0 +1,747 @@ +// Adapted from +// https://github.com/sgl-project/sglang/blob/main/python/sglang/kernels/aot/csrc/cpu/compressor.cpp +// +// DeepSeek-V4 compressor: softmax-weighted pooling over a sliding window of +// raw KV/score state, RMSNorm, FP8/UE8M0 quant (NoPE) + GPT-J RoPE (RoPE), +// paged-insert into the fp8_ds_mla cache layout (see store_cache.cpp). +// Covers both the head_dim=512 attention-compressor path +// (`compress_norm_rope_store_cpu`) and the head_dim=128 indexer-compressor +// path (`compress_norm_rope_store_indexer_cpu`, distinct cache layout -- +// see its own comment below). +// +// Differences from upstream: +// - Upstream splits decode (`compress_decode_cpu`) from prefill/extend +// (eager `compress_extend_separate`); this port stays unified across +// decode/prefill/extend, one program per token, matching vLLM's own +// GPU/Triton reference (`_fused_kv_compress_norm_rope_insert_sparse_attn`, +// common/ops/fused_compress_quant_cache.py). +// - Window-gather indices arrive pre-resolved to flat state-cache row ids +// from the Python caller (mirroring flash_mla.cpp), not looked up via +// block_table in-kernel. +// - The state cache is pooled with the main KV cache's physical pages +// (CompressorStateCache in compressor.py), so its per-block byte stride +// is page-aligned, not `block_size * row_width` -- row ids are addressed +// as (block_idx, pos_in_block) against the tensor's real block stride, +// never by flattening to a 2D view. +// +// clang-format off + +#include +#include +#include +#include + +#include "common.h" +#include "vec.h" + +#if !defined(CPU_CAPABILITY_AVX512) +static_assert(false, + "compressor.cpp is only ever compiled into the AVX512+AMX-flagged " + "_C target (see cmake/cpu_extension.cmake's VLLM_EXT_SRC_SGL) and " + "relies on that unconditionally -- CPU_CAPABILITY_AVX512 must be " + "defined here."); +#endif + +namespace { + +using fVec = at::vec::Vectorized; + +constexpr int64_t kNopeDim = 448; +constexpr int64_t kRopeDim = 64; +constexpr int64_t kHeadDim = kNopeDim + kRopeDim; +constexpr int64_t kQuantBlock = 64; +constexpr int64_t kNumQuantBlocks = kNopeDim / kQuantBlock; +constexpr int64_t kScaleBytesPerToken = kNumQuantBlocks + 1; +constexpr int64_t kTokenDataBytes = kNopeDim + kRopeDim * 2; +constexpr float kFp8Max = 448.0f; +constexpr float kFp8Min = -448.0f; +constexpr float kQuantEps = 1e-4f; + +// In-place interleaved (GPT-J-style) RoPE: x[2i]/x[2i+1] are one (even, odd) +// pair, cos[i]/sin[i] its angle. Adapted from SGLang's `rope.cpp` +// (`apply_rotary_emb_row`), which avoids gather/scatter (slow on AVX512) by +// broadcasting cos/sin across each pair via in-register permutes instead: +// `_mm512_permute_ps(xv, 0xB1)` swaps the (even, odd) lanes within every +// 2-float pair in one shuffle, and cos/sin -- loaded as 8 contiguous values +// -- are duplicated pairwise into the same layout via `_mm512_permutexvar_ps`. +// SGLang's version reads cos/sin pre-interleaved with x (`[c0,s0,c1,s1,...]`) +// since its freqs tensor is laid out that way; ours are two separate flat +// arrays (`cos[i]`, `sin[i]`, matching store_cache.cpp's convention), hence +// the extra broadcast-duplicate permute here instead of a plain load. +// Both call sites in this file pass half_dim=32 (kRopeDim/2, kIndexerRopeDim/2 +// are both 64/2), an exact multiple of 8, so this file's build (always AVX512, +// see file header) never needs a scalar remainder. +inline void apply_gptj_rope_inplace(float* __restrict__ x, const float* __restrict__ cos, const float* __restrict__ sin, int64_t half_dim) { + // Duplicates each of 8 source lanes into 2 adjacent output lanes: + // [c0,c1,...,c7] -> [c0,c0,c1,c1,...,c7,c7]. + const __m512i dup_idx = _mm512_set_epi32(7, 7, 6, 6, 5, 5, 4, 4, 3, 3, 2, 2, 1, 1, 0, 0); + // Forward GPT-J rotation negates the sin broadcast on the even (real) + // lane of each pair; the odd (imag) lane keeps sin's sign. + const __m512 sign_mask = _mm512_castsi512_ps(_mm512_set_epi32( + 0, (int)0x80000000, 0, (int)0x80000000, 0, (int)0x80000000, 0, (int)0x80000000, + 0, (int)0x80000000, 0, (int)0x80000000, 0, (int)0x80000000, 0, (int)0x80000000)); + for (int64_t i = 0; i + 8 <= half_dim; i += 8) { + __m512 xv = _mm512_loadu_ps(x + 2 * i); + __m512 xv_swapped = _mm512_permute_ps(xv, 0xB1); + __m512 cos_b = _mm512_permutexvar_ps(dup_idx, _mm512_castps256_ps512(_mm256_loadu_ps(cos + i))); + __m512 sin_b = _mm512_permutexvar_ps(dup_idx, _mm512_castps256_ps512(_mm256_loadu_ps(sin + i))); + __m512 sin_signed = _mm512_xor_ps(sin_b, sign_mask); + __m512 out = _mm512_fmadd_ps(xv, cos_b, _mm512_mul_ps(xv_swapped, sin_signed)); + _mm512_storeu_ps(x + 2 * i, out); + } +} + +// AVX512 tile-quantization, adapted verbatim from SGLang's `store_cache.cpp` +// (`quant_to_nope_fp8_rope_bf16_pack_cpu`'s +// `cvt_fp32x16_to_fp8x16`/`quantize_tile_avx512`) -- SGLang itself keeps +// `compress_decode_cpu` fp32-output-only and calls this quant step as a +// separate op; vLLM's compressor kernel stays fused (per this file's header +// comment), so the same AVX512 tile-quant math is inlined here instead. +#if defined(CPU_CAPABILITY_AVX512) + +// Vectorized float32x16 -> fp8_e4m3fn x16 using AVX512 bit manipulation. +// Implements RNE (round-to-nearest-even). Flushes subnormal fp8 to zero. +// fp8_e4m3fn: sign(1), exp(4, bias=7), mant(3), no inf/nan special values. +// Normal: exp_field in [1..15], value = 2^(exp-7) * (1 + mant/8) +// Input assumed clamped to [-448, 448]. +inline __m128i cvt_fp32x16_to_fp8x16(__m512 input) { + __m512i bits = _mm512_castps_si512(input); + + __m512i signs = _mm512_srli_epi32(_mm512_and_si512(bits, _mm512_set1_epi32(0x80000000)), 24); + __m512i abs_bits = _mm512_and_si512(bits, _mm512_set1_epi32(0x7FFFFFFF)); + __m512i f32_exp = _mm512_srli_epi32(abs_bits, 23); + __m512i f32_mant = _mm512_and_si512(abs_bits, _mm512_set1_epi32(0x7FFFFF)); + __m512i fp8_exp = _mm512_sub_epi32(f32_exp, _mm512_set1_epi32(120)); + __m512i mant3 = _mm512_srli_epi32(f32_mant, 20); + + __m512i round_bit = _mm512_and_si512(_mm512_srli_epi32(f32_mant, 19), _mm512_set1_epi32(1)); + __m512i sticky_bits = _mm512_and_si512(f32_mant, _mm512_set1_epi32(0x7FFFF)); + __mmask16 has_sticky = _mm512_cmpneq_epi32_mask(sticky_bits, _mm512_setzero_si512()); + + __m512i lsb = _mm512_and_si512(mant3, _mm512_set1_epi32(1)); + __m512i sticky_or_lsb = _mm512_or_si512(_mm512_maskz_mov_epi32(has_sticky, _mm512_set1_epi32(1)), lsb); + __m512i do_round = _mm512_and_si512(round_bit, sticky_or_lsb); + mant3 = _mm512_add_epi32(mant3, do_round); + + __mmask16 mant_ovf = _mm512_cmpeq_epi32_mask(mant3, _mm512_set1_epi32(8)); + mant3 = _mm512_mask_mov_epi32(mant3, mant_ovf, _mm512_setzero_si512()); + fp8_exp = _mm512_mask_add_epi32(fp8_exp, mant_ovf, fp8_exp, _mm512_set1_epi32(1)); + + fp8_exp = _mm512_min_epi32(fp8_exp, _mm512_set1_epi32(15)); + + __m512i result = _mm512_or_si512(_mm512_slli_epi32(fp8_exp, 3), mant3); + + __mmask16 is_subnormal_or_zero = _mm512_cmplt_epi32_mask(f32_exp, _mm512_set1_epi32(121)); + result = _mm512_mask_mov_epi32(result, is_subnormal_or_zero, _mm512_setzero_si512()); + + result = _mm512_or_si512(result, signs); + + return _mm512_cvtepi32_epi8(result); +} + +// Process one 64-wide tile: find amax, compute a shared power-of-two scale, +// quantize to fp8. Returns the UE8M0 scale byte for the tile. +inline uint8_t quantize_tile_avx512(const at::BFloat16* __restrict__ src, uint8_t* __restrict__ dst) { + __m256i bf16_0 = _mm256_loadu_si256(reinterpret_cast(src)); + __m256i bf16_1 = _mm256_loadu_si256(reinterpret_cast(src + 16)); + __m256i bf16_2 = _mm256_loadu_si256(reinterpret_cast(src + 32)); + __m256i bf16_3 = _mm256_loadu_si256(reinterpret_cast(src + 48)); + + __m512 f0 = _mm512_castsi512_ps(_mm512_slli_epi32(_mm512_cvtepu16_epi32(bf16_0), 16)); + __m512 f1 = _mm512_castsi512_ps(_mm512_slli_epi32(_mm512_cvtepu16_epi32(bf16_1), 16)); + __m512 f2 = _mm512_castsi512_ps(_mm512_slli_epi32(_mm512_cvtepu16_epi32(bf16_2), 16)); + __m512 f3 = _mm512_castsi512_ps(_mm512_slli_epi32(_mm512_cvtepu16_epi32(bf16_3), 16)); + + const __m512i abs_mask_i = _mm512_set1_epi32(0x7FFFFFFF); + __m512 abs0 = _mm512_castsi512_ps(_mm512_and_si512(_mm512_castps_si512(f0), abs_mask_i)); + __m512 abs1 = _mm512_castsi512_ps(_mm512_and_si512(_mm512_castps_si512(f1), abs_mask_i)); + __m512 abs2 = _mm512_castsi512_ps(_mm512_and_si512(_mm512_castps_si512(f2), abs_mask_i)); + __m512 abs3 = _mm512_castsi512_ps(_mm512_and_si512(_mm512_castps_si512(f3), abs_mask_i)); + + __m512 max01 = _mm512_max_ps(abs0, abs1); + __m512 max23 = _mm512_max_ps(abs2, abs3); + __m512 max0123 = _mm512_max_ps(max01, max23); + float amax = _mm512_reduce_max_ps(max0123); + + float scale = std::max(amax / kFp8Max, kQuantEps); + float ceil_log2 = std::ceil(std::log2(scale)); + float scale_pow2 = std::exp2(ceil_log2); + float scale_inv = 1.0f / scale_pow2; + + int exponent = static_cast(ceil_log2); + uint8_t scale_uint8 = static_cast(exponent + 127); + + __m512 vinv = _mm512_set1_ps(scale_inv); + __m512 s0 = _mm512_mul_ps(f0, vinv); + __m512 s1 = _mm512_mul_ps(f1, vinv); + __m512 s2 = _mm512_mul_ps(f2, vinv); + __m512 s3 = _mm512_mul_ps(f3, vinv); + + __m512 vmax = _mm512_set1_ps(kFp8Max); + __m512 vmin = _mm512_set1_ps(kFp8Min); + s0 = _mm512_max_ps(_mm512_min_ps(s0, vmax), vmin); + s1 = _mm512_max_ps(_mm512_min_ps(s1, vmax), vmin); + s2 = _mm512_max_ps(_mm512_min_ps(s2, vmax), vmin); + s3 = _mm512_max_ps(_mm512_min_ps(s3, vmax), vmin); + + __m128i fp8_0 = cvt_fp32x16_to_fp8x16(s0); + __m128i fp8_1 = cvt_fp32x16_to_fp8x16(s1); + __m128i fp8_2 = cvt_fp32x16_to_fp8x16(s2); + __m128i fp8_3 = cvt_fp32x16_to_fp8x16(s3); + + _mm_storeu_si128(reinterpret_cast<__m128i*>(dst), fp8_0); + _mm_storeu_si128(reinterpret_cast<__m128i*>(dst + 16), fp8_1); + _mm_storeu_si128(reinterpret_cast<__m128i*>(dst + 32), fp8_2); + _mm_storeu_si128(reinterpret_cast<__m128i*>(dst + 48), fp8_3); + + // Patch subnormal-magnitude lanes: `cvt_fp32x16_to_fp8x16` flushes + // anything below the smallest fp8_e4m3fn normal (2^-6) to zero, but + // vLLM's dequant path relies on `at::Float8_e4m3fn`'s real subnormal + // encoding -- redo just those rare lanes with the exact scalar + // conversion to stay consistent with the read side. + alignas(64) float scaled[kQuantBlock]; + _mm512_store_ps(scaled, s0); + _mm512_store_ps(scaled + 16, s1); + _mm512_store_ps(scaled + 32, s2); + _mm512_store_ps(scaled + 48, s3); + for (int64_t j = 0; j < kQuantBlock; ++j) { + const float av = std::abs(scaled[j]); + if (av != 0.f && av < 0x1p-6f) { + dst[j] = at::Float8_e4m3fn(scaled[j]).x; + } + } + + return scale_uint8; +} + +#endif // CPU_CAPABILITY_AVX512 + +void save_partial_states_impl( + float* __restrict__ state_cache, // [num_blocks, block_stride] + int64_t block_stride, // element stride between blocks + int64_t block_size, // rows (slots) per block + const float* __restrict__ kv, // [num_tokens, state_width] + int64_t kv_row_stride, // element stride between kv rows + const float* __restrict__ score, // [num_tokens, state_width] + int64_t score_row_stride, // element stride between score rows + const float* __restrict__ ape, // [compress_ratio, state_width] + const int64_t* __restrict__ positions, // [num_tokens] + const int64_t* __restrict__ slot_mapping, // [num_tokens] + int64_t state_width, + int64_t compress_ratio, + int64_t num_tokens) { + const int64_t row_width = 2 * state_width; + at::parallel_for(0, num_tokens, 0, [&](int64_t begin, int64_t end) { + for (int64_t t = begin; t < end; ++t) { + const int64_t slot_id = slot_mapping[t]; + if (slot_id < 0) continue; + + const int64_t block_idx = slot_id / block_size; + const int64_t pos_in_block = slot_id % block_size; + float* const row = state_cache + block_idx * block_stride + pos_in_block * row_width; + const float* const kv_t = kv + t * kv_row_stride; + // state_width is always a multiple of 16 for this model (512/1024 for + // the main compressor, 128/256 for the indexer compressor). + for (int64_t i = 0; i + 16 <= state_width; i += 16) { + fVec::loadu(kv_t + i).store(row + i); + } + + const int64_t position = positions[t]; + const float* const ape_row = ape + (position % compress_ratio) * state_width; + const float* const score_t = score + t * score_row_stride; + float* const score_out = row + state_width; + for (int64_t i = 0; i + 16 <= state_width; i += 16) { + (fVec::loadu(score_t + i) + fVec::loadu(ape_row + i)).store(score_out + i); + } + } + }); +} + +void compress_norm_rope_store_impl( + const float* __restrict__ state_cache, // [num_blocks, state_block_stride] + int64_t state_block_stride, // element stride between state blocks + int64_t state_block_size, // rows (slots) per state block + const int64_t* __restrict__ gather_slots, // [num_tokens, window], -1 = invalid + const int64_t* __restrict__ positions, // [num_tokens] + const int64_t* __restrict__ kv_slot_mapping, // [num_tokens], -1 = skip + const float* __restrict__ rms_norm_weight, // [kHeadDim] + float rms_norm_eps, + const float* __restrict__ cos_sin_cache, // [max_pos, kRopeDim] + uint8_t* __restrict__ kv_cache, // [num_blocks, cache_row_stride] + int64_t cache_row_stride, + int64_t kv_cache_block_size, + int64_t compress_ratio, + int64_t window, + int64_t state_width, + int64_t num_tokens) { + const int64_t rope_half = kRopeDim / 2; + const int64_t state_row_width = 2 * state_width; + + at::parallel_for(0, num_tokens, 0, [&](int64_t begin, int64_t end) { + std::array max_v; + std::array denom; + std::array acc; + std::array normed; + std::array nope_bf16; + + for (int64_t t = begin; t < end; ++t) { + const int64_t position = positions[t]; + if ((position + 1) % compress_ratio != 0) continue; + const int64_t kv_slot = kv_slot_mapping[t]; + if (kv_slot < 0) continue; + + std::fill(max_v.begin(), max_v.end(), -std::numeric_limits::infinity()); + std::fill(denom.begin(), denom.end(), 0.f); + std::fill(acc.begin(), acc.end(), 0.f); + + // Per-channel (not per-position) softmax over the window: each of the + // kHeadDim columns is normalized independently across the window axis. + for (int64_t w = 0; w < window; ++w) { + const int64_t slot = gather_slots[t * window + w]; + if (slot < 0) continue; + const int64_t block_idx = slot / state_block_size; + const int64_t pos_in_block = slot % state_block_size; + const int64_t head_offset = (w >= compress_ratio) ? kHeadDim : 0; + const float* const score_row = state_cache + block_idx * state_block_stride + + pos_in_block * state_row_width + head_offset + state_width; + for (int64_t d = 0; d < kHeadDim; d += 16) { + __m512 mx = _mm512_loadu_ps(max_v.data() + d); + __m512 sv = _mm512_loadu_ps(score_row + d); + _mm512_storeu_ps(max_v.data() + d, _mm512_max_ps(mx, sv)); + } + } + for (int64_t w = 0; w < window; ++w) { + const int64_t slot = gather_slots[t * window + w]; + if (slot < 0) continue; + const int64_t block_idx = slot / state_block_size; + const int64_t pos_in_block = slot % state_block_size; + const int64_t head_offset = (w >= compress_ratio) ? kHeadDim : 0; + const float* const row = state_cache + block_idx * state_block_stride + + pos_in_block * state_row_width + head_offset; + const float* const score_row = row + state_width; + for (int64_t d = 0; d < kHeadDim; d += 16) { + fVec e_v = (fVec::loadu(score_row + d) - fVec::loadu(max_v.data() + d)).exp(); + (fVec::loadu(denom.data() + d) + e_v).store(denom.data() + d); + (fVec::loadu(acc.data() + d) + e_v * fVec::loadu(row + d)).store(acc.data() + d); + } + } + + float sq_sum = 0.f; + { + fVec sq_sum_v(0.f); + for (int64_t d = 0; d < kHeadDim; d += 16) { + fVec v = fVec::loadu(acc.data() + d) / fVec::loadu(denom.data() + d); + v.store(normed.data() + d); + sq_sum_v = sq_sum_v + v * v; + } + sq_sum = vec_reduce_sum(sq_sum_v); + } + const float inv_rms = 1.0f / std::sqrt(sq_sum / static_cast(kHeadDim) + rms_norm_eps); + { + fVec inv_rms_v(inv_rms); + for (int64_t d = 0; d < kHeadDim; d += 16) { + (fVec::loadu(normed.data() + d) * inv_rms_v * fVec::loadu(rms_norm_weight + d)).store(normed.data() + d); + } + } + + // bf16 round-trip feeds only the NoPE FP8 quantization (RoPE rotates + // the un-rounded `normed` below, then casts once to bf16 at store + // time for the separate bf16 rope region). + for (int64_t d = 0; d < kNopeDim; d += 16) + store_from_float_ext(nope_bf16.data() + d, fVec::loadu(normed.data() + d)); + + const int64_t compressed_pos = (position / compress_ratio) * compress_ratio; + const float* const cos = cos_sin_cache + compressed_pos * kRopeDim; + const float* const sin = cos + rope_half; + apply_gptj_rope_inplace(normed.data() + kNopeDim, cos, sin, rope_half); + + const int64_t block_idx = kv_slot / kv_cache_block_size; + const int64_t pos_in_block = kv_slot % kv_cache_block_size; + uint8_t* const block_base = kv_cache + block_idx * cache_row_stride; + uint8_t* const token_fp8_ptr = block_base + pos_in_block * kTokenDataBytes; + uint8_t* const token_bf16_ptr = token_fp8_ptr + kNopeDim; + uint8_t* const token_scale_ptr = + block_base + kv_cache_block_size * kTokenDataBytes + pos_in_block * kScaleBytesPerToken; + + // AVX512 tile-quantization (SGLang's `quantize_tile_avx512`), one + // 64-wide tile per iteration. + for (int64_t qb = 0; qb < kNumQuantBlocks; ++qb) { + const at::BFloat16* tile_src = nope_bf16.data() + qb * kQuantBlock; + uint8_t* tile_dst = token_fp8_ptr + qb * kQuantBlock; + token_scale_ptr[qb] = quantize_tile_avx512(tile_src, tile_dst); + } + token_scale_ptr[kNumQuantBlocks] = 0; + + auto* const out_rope = reinterpret_cast(token_bf16_ptr); + for (int64_t ri = 0; ri < kRopeDim; ri += 16) + store_from_float_ext(out_rope + ri, fVec::loadu(normed.data() + kNopeDim + ri)); + } + }); +} + +// Indexer compressor (head_dim=128): same softmax-pool + RMSNorm + GPT-J +// RoPE prologue as the head=512 path above, but a different cache layout: +// all 128 dims are FP8 in one combined quant block (no separate bf16 rope +// region, single raw fp32 scale per token) instead of 7 UE8M0 blocks over +// nope alone -- matching the paged indexer K-cache layout +// fp8_paged_mqa_logits_cpu/topk_transform_512_cpu (paged_mqa_logits.cpp/ +// topk.cpp) read, and the GPU/Triton reference +// (`_fused_kv_compress_norm_rope_insert_indexer_attn`, +// common/ops/fused_compress_quant_cache.py) exactly. +constexpr int64_t kIndexerHeadDim = 128; +constexpr int64_t kIndexerRopeDim = 64; +constexpr int64_t kIndexerNopeDim = kIndexerHeadDim - kIndexerRopeDim; +constexpr int64_t kIndexerTokenBytes = kIndexerHeadDim; +constexpr int64_t kIndexerScaleBytes = 4; + +void compress_norm_rope_store_indexer_impl( + const float* __restrict__ state_cache, // [num_blocks, state_block_stride] + int64_t state_block_stride, + int64_t state_block_size, + const int64_t* __restrict__ gather_slots, // [num_tokens, window], -1 = invalid + const int64_t* __restrict__ positions, // [num_tokens] + const int64_t* __restrict__ kv_slot_mapping, // [num_tokens], -1 = skip + const float* __restrict__ rms_norm_weight, // [kIndexerHeadDim] + float rms_norm_eps, + const float* __restrict__ cos_sin_cache, // [max_pos, kIndexerRopeDim] + uint8_t* __restrict__ kv_cache, // [num_blocks, cache_row_stride] + int64_t cache_row_stride, + int64_t kv_cache_block_size, + int64_t compress_ratio, + int64_t window, + int64_t state_width, + int64_t num_tokens) { + const int64_t rope_half = kIndexerRopeDim / 2; + const int64_t state_row_width = 2 * state_width; + + at::parallel_for(0, num_tokens, 0, [&](int64_t begin, int64_t end) { + std::array max_v; + std::array denom; + std::array acc; + std::array normed; + std::array quant_bf16; + + for (int64_t t = begin; t < end; ++t) { + const int64_t position = positions[t]; + if ((position + 1) % compress_ratio != 0) continue; + const int64_t kv_slot = kv_slot_mapping[t]; + if (kv_slot < 0) continue; + + std::fill(max_v.begin(), max_v.end(), -std::numeric_limits::infinity()); + std::fill(denom.begin(), denom.end(), 0.f); + std::fill(acc.begin(), acc.end(), 0.f); + + for (int64_t w = 0; w < window; ++w) { + const int64_t slot = gather_slots[t * window + w]; + if (slot < 0) continue; + const int64_t block_idx = slot / state_block_size; + const int64_t pos_in_block = slot % state_block_size; + const int64_t head_offset = (w >= compress_ratio) ? kIndexerHeadDim : 0; + const float* const score_row = state_cache + block_idx * state_block_stride + + pos_in_block * state_row_width + head_offset + state_width; + for (int64_t d = 0; d < kIndexerHeadDim; d += 16) { + __m512 mx = _mm512_loadu_ps(max_v.data() + d); + __m512 sv = _mm512_loadu_ps(score_row + d); + _mm512_storeu_ps(max_v.data() + d, _mm512_max_ps(mx, sv)); + } + } + for (int64_t w = 0; w < window; ++w) { + const int64_t slot = gather_slots[t * window + w]; + if (slot < 0) continue; + const int64_t block_idx = slot / state_block_size; + const int64_t pos_in_block = slot % state_block_size; + const int64_t head_offset = (w >= compress_ratio) ? kIndexerHeadDim : 0; + const float* const row = state_cache + block_idx * state_block_stride + + pos_in_block * state_row_width + head_offset; + const float* const score_row = row + state_width; + for (int64_t d = 0; d < kIndexerHeadDim; d += 16) { + fVec e_v = (fVec::loadu(score_row + d) - fVec::loadu(max_v.data() + d)).exp(); + (fVec::loadu(denom.data() + d) + e_v).store(denom.data() + d); + (fVec::loadu(acc.data() + d) + e_v * fVec::loadu(row + d)).store(acc.data() + d); + } + } + + float sq_sum = 0.f; + { + fVec sq_sum_v(0.f); + for (int64_t d = 0; d < kIndexerHeadDim; d += 16) { + fVec v = fVec::loadu(acc.data() + d) / fVec::loadu(denom.data() + d); + v.store(normed.data() + d); + sq_sum_v = sq_sum_v + v * v; + } + sq_sum = vec_reduce_sum(sq_sum_v); + } + const float inv_rms = 1.0f / std::sqrt(sq_sum / static_cast(kIndexerHeadDim) + rms_norm_eps); + { + fVec inv_rms_v(inv_rms); + for (int64_t d = 0; d < kIndexerHeadDim; d += 16) { + (fVec::loadu(normed.data() + d) * inv_rms_v * fVec::loadu(rms_norm_weight + d)).store(normed.data() + d); + } + } + + const int64_t compressed_pos = (position / compress_ratio) * compress_ratio; + const float* const cos = cos_sin_cache + compressed_pos * kIndexerRopeDim; + const float* const sin = cos + rope_half; + apply_gptj_rope_inplace(normed.data() + kIndexerNopeDim, cos, sin, rope_half); + + // Indexer path rounds the FULL post-RoPE vector through bf16 before + // quantizing (single combined block); the head=512 path above instead + // rounds only the NoPE region before RoPE (separate bf16 rope store). + for (int64_t d = 0; d < kIndexerHeadDim; d += 16) + store_from_float_ext(quant_bf16.data() + d, fVec::loadu(normed.data() + d)); + + const int64_t block_idx = kv_slot / kv_cache_block_size; + const int64_t pos_in_block = kv_slot % kv_cache_block_size; + uint8_t* const block_base = kv_cache + block_idx * cache_row_stride; + uint8_t* const token_fp8_ptr = block_base + pos_in_block * kIndexerTokenBytes; + uint8_t* const token_scale_ptr = + block_base + kv_cache_block_size * kIndexerTokenBytes + pos_in_block * kIndexerScaleBytes; + + float absmax = 1e-4f; + { + __m512 max_v_reg = _mm512_set1_ps(absmax); + for (int64_t di = 0; di < kIndexerHeadDim; di += 16) { + __m256i bf16_v = _mm256_loadu_si256(reinterpret_cast(quant_bf16.data() + di)); + __m512 f = _mm512_castsi512_ps(_mm512_slli_epi32(_mm512_cvtepu16_epi32(bf16_v), 16)); + __m512 av = _mm512_castsi512_ps(_mm512_and_si512(_mm512_castps_si512(f), _mm512_set1_epi32(0x7FFFFFFF))); + max_v_reg = _mm512_max_ps(max_v_reg, av); + } + absmax = _mm512_reduce_max_ps(max_v_reg); + } + const float exponent = std::ceil(std::log2(absmax / kFp8Max)); + const float scale = std::exp2(exponent); + const float inv_scale = 1.0f / scale; + + // AVX512 quantize+store, reusing SGLang's `cvt_fp32x16_to_fp8x16` bit- + // manipulation conversion (this single-shared-scale 128-wide block has + // no direct SGLang equivalent -- see file header -- so only the fp8 + // conversion primitive is ported, the amax/scale formula is unchanged + // from vLLM's own reference). + uint8_t* const out_fp8 = token_fp8_ptr; + { + const __m512 vinv = _mm512_set1_ps(inv_scale); + const __m512 vmax = _mm512_set1_ps(kFp8Max); + const __m512 vmin = _mm512_set1_ps(kFp8Min); + alignas(64) float scaled_buf[kIndexerHeadDim]; + for (int64_t d = 0; d < kIndexerHeadDim; d += 16) { + __m256i bf16_v = _mm256_loadu_si256(reinterpret_cast(quant_bf16.data() + d)); + __m512 f = _mm512_castsi512_ps(_mm512_slli_epi32(_mm512_cvtepu16_epi32(bf16_v), 16)); + __m512 s = _mm512_mul_ps(f, vinv); + s = _mm512_max_ps(_mm512_min_ps(s, vmax), vmin); + _mm512_store_ps(scaled_buf + d, s); + __m128i fp8_v = cvt_fp32x16_to_fp8x16(s); + _mm_storeu_si128(reinterpret_cast<__m128i*>(out_fp8 + d), fp8_v); + } + // Patch subnormal-magnitude lanes (see `quantize_tile_avx512`'s + // comment above for why the vectorized conversion's flush-to-zero + // must not be trusted here). + for (int64_t d = 0; d < kIndexerHeadDim; ++d) { + const float av = std::abs(scaled_buf[d]); + if (av != 0.f && av < 0x1p-6f) { + out_fp8[d] = at::Float8_e4m3fn(scaled_buf[d]).x; + } + } + } + // Raw fp32 scale (not UE8M0-encoded) -- single quant block, matches + // fp8_paged_mqa_logits_cpu's (paged_mqa_logits.cpp) read-side expectations. + *reinterpret_cast(token_scale_ptr) = scale; + } + }); +} + +} // anonymous namespace + +void save_partial_states_cpu( + at::Tensor& kv, + at::Tensor& score, + at::Tensor& ape, + at::Tensor& positions, + at::Tensor& state_cache, + at::Tensor& slot_mapping) { + TORCH_CHECK( + kv.dim() == 2 && kv.scalar_type() == at::kFloat, "save_partial_states_cpu: kv must be a 2D fp32 tensor"); + TORCH_CHECK(score.sizes() == kv.sizes() && score.scalar_type() == at::kFloat, + "save_partial_states_cpu: score must match kv's shape/dtype"); + TORCH_CHECK(ape.dim() == 2 && ape.size(1) == kv.size(1) && ape.scalar_type() == at::kFloat, + "save_partial_states_cpu: ape must be [compress_ratio, state_width] fp32"); + TORCH_CHECK(positions.scalar_type() == at::kLong && positions.numel() == kv.size(0), + "save_partial_states_cpu: positions must be int64 with num_tokens entries"); + TORCH_CHECK(slot_mapping.scalar_type() == at::kLong && slot_mapping.numel() == kv.size(0), + "save_partial_states_cpu: slot_mapping must be int64 with num_tokens entries"); + TORCH_CHECK(state_cache.dim() == 3 && state_cache.size(2) == 2 * kv.size(1) && + state_cache.scalar_type() == at::kFloat, + "save_partial_states_cpu: state_cache must be [num_blocks, block_size, " + "2*state_width] fp32"); + TORCH_CHECK(state_cache.stride(2) == 1 && state_cache.stride(1) == state_cache.size(2), + "save_partial_states_cpu: state_cache rows must be contiguous within a block"); + TORCH_CHECK( + kv.device().is_cpu() && score.device().is_cpu() && ape.device().is_cpu() && positions.device().is_cpu() && + state_cache.device().is_cpu() && slot_mapping.device().is_cpu(), + "save_partial_states_cpu: all inputs must be CPU tensors"); + + const int64_t num_tokens = kv.size(0); + const int64_t state_width = kv.size(1); + const int64_t compress_ratio = ape.size(0); + const int64_t block_size = state_cache.size(1); + + TORCH_CHECK( + kv.stride(1) == 1 && score.stride(1) == 1 && ape.is_contiguous() && positions.is_contiguous() && + slot_mapping.is_contiguous(), + "save_partial_states_cpu: kv/score must be contiguous in their last dim " + "(row stride may exceed state_width -- e.g. a column-split view of a " + "fused kv_score projection); ape/positions/slot_mapping must be fully " + "contiguous"); + + save_partial_states_impl( + state_cache.data_ptr(), + state_cache.stride(0), + block_size, + kv.data_ptr(), + kv.stride(0), + score.data_ptr(), + score.stride(0), + ape.data_ptr(), + positions.data_ptr(), + slot_mapping.data_ptr(), + state_width, + compress_ratio, + num_tokens); +} + +void compress_norm_rope_store_cpu( + at::Tensor& state_cache, + at::Tensor& gather_slots, + at::Tensor& positions, + at::Tensor& kv_slot_mapping, + at::Tensor& rms_norm_weight, + double rms_norm_eps, + at::Tensor& cos_sin_cache, + at::Tensor& kv_cache_2d, + int64_t kv_cache_block_size, + int64_t compress_ratio) { + TORCH_CHECK(state_cache.dim() == 3 && state_cache.scalar_type() == at::kFloat, + "compress_norm_rope_store_cpu: state_cache must be a 3D fp32 tensor"); + TORCH_CHECK(state_cache.stride(2) == 1 && state_cache.stride(1) == state_cache.size(2), + "compress_norm_rope_store_cpu: state_cache rows must be contiguous within a block"); + TORCH_CHECK(gather_slots.dim() == 2 && gather_slots.scalar_type() == at::kLong, + "compress_norm_rope_store_cpu: gather_slots must be a 2D int64 tensor"); + const int64_t num_tokens = gather_slots.size(0); + TORCH_CHECK(positions.scalar_type() == at::kLong && positions.numel() == num_tokens, + "compress_norm_rope_store_cpu: positions must be int64 with num_tokens entries"); + TORCH_CHECK(kv_slot_mapping.scalar_type() == at::kLong && kv_slot_mapping.numel() == num_tokens, + "compress_norm_rope_store_cpu: kv_slot_mapping must be int64 with num_tokens entries"); + TORCH_CHECK(rms_norm_weight.numel() == kHeadDim, + "compress_norm_rope_store_cpu: rms_norm_weight must have head_dim (512) entries"); + TORCH_CHECK(kv_cache_2d.dim() == 2 && kv_cache_2d.scalar_type() == at::kByte, + "compress_norm_rope_store_cpu: kv_cache_2d must be a 2D uint8 tensor"); + TORCH_CHECK(state_cache.size(2) % 2 == 0, "compress_norm_rope_store_cpu: state_cache width must be even"); + TORCH_CHECK( + rms_norm_weight.scalar_type() == at::kFloat && rms_norm_weight.is_contiguous() && + cos_sin_cache.scalar_type() == at::kFloat && cos_sin_cache.is_contiguous(), + "compress_norm_rope_store_cpu: rms_norm_weight/cos_sin_cache must be " + "float32 and contiguous (convert once in Python)"); + TORCH_CHECK( + gather_slots.is_contiguous() && positions.is_contiguous() && kv_slot_mapping.is_contiguous(), + "compress_norm_rope_store_cpu: gather_slots/positions/kv_slot_mapping must be contiguous"); + TORCH_CHECK( + state_cache.device().is_cpu() && gather_slots.device().is_cpu() && positions.device().is_cpu() && + kv_slot_mapping.device().is_cpu() && rms_norm_weight.device().is_cpu() && + cos_sin_cache.device().is_cpu() && kv_cache_2d.device().is_cpu(), + "compress_norm_rope_store_cpu: all inputs must be CPU tensors"); + + const int64_t window = gather_slots.size(1); + const int64_t state_width = state_cache.size(2) / 2; + const int64_t state_block_size = state_cache.size(1); + + // Row strides are threaded through explicitly rather than assumed, since + // both buffers may be pages of a larger pooled allocation (matches + // store_cache.cpp/flash_mla.cpp's precedent). + const int64_t state_block_stride = state_cache.stride(0); + const int64_t cache_row_stride = kv_cache_2d.stride(0); + + compress_norm_rope_store_impl( + state_cache.data_ptr(), + state_block_stride, + state_block_size, + gather_slots.data_ptr(), + positions.data_ptr(), + kv_slot_mapping.data_ptr(), + rms_norm_weight.data_ptr(), + static_cast(rms_norm_eps), + cos_sin_cache.data_ptr(), + kv_cache_2d.data_ptr(), + cache_row_stride, + kv_cache_block_size, + compress_ratio, + window, + state_width, + num_tokens); +} + +void compress_norm_rope_store_indexer_cpu( + at::Tensor& state_cache, + at::Tensor& gather_slots, + at::Tensor& positions, + at::Tensor& kv_slot_mapping, + at::Tensor& rms_norm_weight, + double rms_norm_eps, + at::Tensor& cos_sin_cache, + at::Tensor& kv_cache_2d, + int64_t kv_cache_block_size, + int64_t compress_ratio) { + TORCH_CHECK(state_cache.dim() == 3 && state_cache.scalar_type() == at::kFloat, + "compress_norm_rope_store_indexer_cpu: state_cache must be a 3D fp32 tensor"); + TORCH_CHECK(state_cache.stride(2) == 1 && state_cache.stride(1) == state_cache.size(2), + "compress_norm_rope_store_indexer_cpu: state_cache rows must be contiguous within a block"); + TORCH_CHECK(gather_slots.dim() == 2 && gather_slots.scalar_type() == at::kLong, + "compress_norm_rope_store_indexer_cpu: gather_slots must be a 2D int64 tensor"); + const int64_t num_tokens = gather_slots.size(0); + TORCH_CHECK(positions.scalar_type() == at::kLong && positions.numel() == num_tokens, + "compress_norm_rope_store_indexer_cpu: positions must be int64 with num_tokens entries"); + TORCH_CHECK(kv_slot_mapping.scalar_type() == at::kLong && kv_slot_mapping.numel() == num_tokens, + "compress_norm_rope_store_indexer_cpu: kv_slot_mapping must be int64 with num_tokens entries"); + TORCH_CHECK(rms_norm_weight.numel() == kIndexerHeadDim, + "compress_norm_rope_store_indexer_cpu: rms_norm_weight must have head_dim (128) entries"); + TORCH_CHECK(kv_cache_2d.dim() == 2 && kv_cache_2d.scalar_type() == at::kByte, + "compress_norm_rope_store_indexer_cpu: kv_cache_2d must be a 2D uint8 tensor"); + TORCH_CHECK(state_cache.size(2) % 2 == 0, + "compress_norm_rope_store_indexer_cpu: state_cache width must be even"); + TORCH_CHECK( + rms_norm_weight.scalar_type() == at::kFloat && rms_norm_weight.is_contiguous() && + cos_sin_cache.scalar_type() == at::kFloat && cos_sin_cache.is_contiguous(), + "compress_norm_rope_store_indexer_cpu: rms_norm_weight/cos_sin_cache must " + "be float32 and contiguous (convert once in Python)"); + TORCH_CHECK( + gather_slots.is_contiguous() && positions.is_contiguous() && kv_slot_mapping.is_contiguous(), + "compress_norm_rope_store_indexer_cpu: gather_slots/positions/kv_slot_mapping must be contiguous"); + TORCH_CHECK( + state_cache.device().is_cpu() && gather_slots.device().is_cpu() && positions.device().is_cpu() && + kv_slot_mapping.device().is_cpu() && rms_norm_weight.device().is_cpu() && + cos_sin_cache.device().is_cpu() && kv_cache_2d.device().is_cpu(), + "compress_norm_rope_store_indexer_cpu: all inputs must be CPU tensors"); + + const int64_t window = gather_slots.size(1); + const int64_t state_width = state_cache.size(2) / 2; + const int64_t state_block_size = state_cache.size(1); + + const int64_t state_block_stride = state_cache.stride(0); + const int64_t cache_row_stride = kv_cache_2d.stride(0); + + compress_norm_rope_store_indexer_impl( + state_cache.data_ptr(), + state_block_stride, + state_block_size, + gather_slots.data_ptr(), + positions.data_ptr(), + kv_slot_mapping.data_ptr(), + rms_norm_weight.data_ptr(), + static_cast(rms_norm_eps), + cos_sin_cache.data_ptr(), + kv_cache_2d.data_ptr(), + cache_row_stride, + kv_cache_block_size, + compress_ratio, + window, + state_width, + num_tokens); +} diff --git a/csrc/cpu/sgl-kernels/flash_mla.cpp b/csrc/cpu/sgl-kernels/flash_mla.cpp new file mode 100644 index 0000000000..8be8d9443b --- /dev/null +++ b/csrc/cpu/sgl-kernels/flash_mla.cpp @@ -0,0 +1,625 @@ +// Adapted from +// https://github.com/sgl-project/sglang/blob/main/python/sglang/kernels/aot/csrc/cpu/flash_mla.cpp +// +// DeepSeek-V4 sparse MQA attention kernel for the fp8_ds_mla cache layout. +// Near-verbatim port of `sparse_mla_decode_kernel_impl` (VNNI K/V packing + +// AMX brgemm QK^T/PV + online softmax); the `V4_FP8Sparse` byte layout +// matches vLLM's fp8_ds_mla layout exactly (see store_cache.cpp). +// +// Differences from upstream: +// - Outer shape: upstream takes a padded batch-major query `[B, S_q, H, D]` +// with per-batch topk sets; vLLM passes flat `[num_tokens, H, D]` with +// per-token resolved indices. The kernel's own `at::parallel_for` already +// flattens `(batch, s_q)` internally, so this is called with +// `batches = num_tokens, s_q = 1` -- no padding, no kernel-body changes. +// - Two physically distinct paged caches (SWA "window" + compressor +// "compressed") are addressed via block_table-resolved flat slot ids, +// matching upstream's `k_main`/`k_extra` two-index-set convention. +// - `attn_sink` uses upstream's own post-hoc LSE correction +// (`out *= 1 / (1 + exp(sink - lse))`). +// - bf16 query only (matching upstream's own restriction: its FP8->bf16 +// dequant intrinsics are bf16-specific); fp16 is not supported here. +// +// clang-format off + +#include +#include +#include +#include + +#include "common.h" +#include "gemm.h" +#include "vec.h" + +#if !defined(CPU_CAPABILITY_AVX512) +static_assert(false, + "flash_mla.cpp is only ever compiled into the AVX512+AMX-flagged " + "_C target (see cmake/cpu_extension.cmake's VLLM_EXT_SRC_SGL) and " + "relies on that unconditionally -- CPU_CAPABILITY_AVX512 must be " + "defined here."); +#endif + +namespace { + +constexpr int64_t kNopeDim = 448; +constexpr int64_t kRopeDim = 64; +constexpr int64_t kHeadDim = kNopeDim + kRopeDim; +constexpr int64_t kQuantBlock = 64; +constexpr int64_t kNumQuantBlocks = kNopeDim / kQuantBlock; +constexpr int64_t kScaleBytesPerToken = kNumQuantBlocks + 1; +constexpr int64_t kTokenDataBytes = kNopeDim + kRopeDim * 2; + +// --------------------------------------------------------------------------- +// V4_FP8Sparse layout (SGLang's `flash_mla.cpp`) -- byte-for-byte identical +// to vLLM's own fp8_ds_mla layout above (d=512, d_nope=448, d_rope=64, +// tile=64, num_tiles=7): +// per block: +// [block_size * (d_nope + 2*d_rope) bytes ; FP8 NoPE + bf16 RoPE per token] +// [block_size * 8 bytes ; 7 e8m0 scales + 1 pad byte/token] +// --------------------------------------------------------------------------- + +struct FP8LayoutMeta { + int64_t d; + int64_t d_nope; + int64_t d_rope; + int64_t tile_size; + int64_t num_tiles; +}; + +constexpr FP8LayoutMeta kV4Meta = {kHeadDim, kNopeDim, kRopeDim, kQuantBlock, kNumQuantBlocks}; + +// Convert one fp8_e8m0 byte (= unsigned 8-bit exponent) to float. +// e8m0 only stores an exponent (bias 127); value = 2^(e - 127), with 0xFF +// reserved for NaN. +inline float fp8_e8m0_to_float(uint8_t v) { + if (v == 0xFF) return std::numeric_limits::quiet_NaN(); + if (v == 0) return 0.f; + union { + uint32_t u; + float f; + } u; + u.u = static_cast(v) << 23; + return u.f; +} + +template +inline bool is_valid_sparse_index(index_t idx, int64_t pos, int64_t topk_limit, int64_t total_tokens) { + const int64_t v = static_cast(idx); + return pos < topk_limit && v >= 0 && v < total_tokens; +} + +#if defined(CPU_CAPABILITY_AVX512) +inline __attribute__((always_inline)) __m512i cvt_fp8_32_to_scaled_bf16(__m256i fp8, float scale) { + const __m512bh bf16_ext = CVT_FP8_TO_BF16_EXT(fp8); + const __m512 scale_vec = _mm512_mul_ps(_mm512_set1_ps(scale), _mm512_castsi512_ps(_mm512_set1_epi32(kFP8_BIAS))); + const __m512 f0 = _mm512_mul_ps(CVT_BF16_TO_FP32(_mm512_extracti32x8_epi32((__m512i)bf16_ext, 0)), scale_vec); + const __m512 f1 = _mm512_mul_ps(CVT_BF16_TO_FP32(_mm512_extracti32x8_epi32((__m512i)bf16_ext, 1)), scale_vec); + return (__m512i)_mm512_cvtne2ps_pbh(f1, f0); +} + +// Load 32 contiguous logical dims (NoPE fp8-dequant-to-bf16, or raw bf16 +// RoPE) starting at `dim_offset` from one token row of the V4_FP8Sparse +// layout. +inline __attribute__((always_inline)) __m512i load_fp8_kvcache_32_from_row( + const uint8_t* __restrict__ row_base, const uint8_t* __restrict__ scale_base, int dim_offset) { + if (dim_offset < kV4Meta.d_nope) { + const float scale = fp8_e8m0_to_float(scale_base[dim_offset / kV4Meta.tile_size]); + return cvt_fp8_32_to_scaled_bf16( + _mm256_loadu_si256(reinterpret_cast(row_base + dim_offset)), scale); + } + const auto* rope_ptr = reinterpret_cast(row_base + kV4Meta.d_nope); + return _mm512_loadu_si512(rope_ptr + dim_offset - kV4Meta.d_nope); +} + +inline __attribute__((always_inline)) void init_fp8_kvcache_tile_rows( + const uint8_t* __restrict__ fp8_storage, + const int64_t block_size, + const int64_t storage_block_stride_bytes, + const int64_t token_idx, + const uint8_t*& row_base, + const uint8_t*& scale_base) { + constexpr int64_t nope_rope_per_token = kV4Meta.d_nope + 2 * kV4Meta.d_rope; + constexpr int64_t scale_stride = 8; + const int64_t block_idx = token_idx / block_size; + const int64_t block_off = token_idx - block_idx * block_size; + const uint8_t* block_base = fp8_storage + block_idx * storage_block_stride_bytes; + row_base = block_base + block_off * nope_rope_per_token; + scale_base = block_base + block_size * nope_rope_per_token + block_off * scale_stride; +} + +template +inline __attribute__((always_inline)) void sparse_pack_vnni_Nx32( + scalar_t* __restrict__ dst0, + scalar_t* __restrict__ dst1, + const index_t* __restrict__ ind, + const bool* __restrict__ valid_mask, + int N, + int dim_offset, + int ld_dst0, + int ld_dst1, + const load_vec_t& load_vec) { + __m512i vinputs[16]; + int n = 0; + for (; n < N; ++n) { + if (!valid_mask[n]) { + vinputs[n] = _mm512_set1_epi32(0); + } else { + vinputs[n] = load_vec(n, static_cast(ind[n]), dim_offset); + } + } + for (; n < 16; ++n) { + vinputs[n] = _mm512_set1_epi32(0); + } + + if constexpr (convert_v) { + for (int nn = 0; nn < 16; nn += 2) { + __m512i d0, d1; + std::tie(d0, d1) = transpose_2x32_16bit(vinputs[nn], vinputs[nn + 1]); + _mm512_storeu_si512(dst1 + (nn >> 1) * ld_dst1 * 2, d0); + _mm512_storeu_si512(dst1 + (nn >> 1) * ld_dst1 * 2 + 32, d1); + } + } + + transpose_16x16_32bit(vinputs); + const __mmask16 vmask = (1 << N) - 1; + for (int k = 0; k < 16; ++k) { + _mm512_mask_storeu_epi32(dst0 + k * ld_dst0 * 2, vmask, vinputs[k]); + } +} +#endif // CPU_CAPABILITY_AVX512 + +// Gather+VNNI-pack `N` fp8-cache rows (dequantized to bf16 while packing) +// into `dst0` (K, vnni-packed for brgemm's B operand) and `dst1` (V, +// vnni-packed). +template +void sparse_pack_vnni_fp8( + scalar_t* __restrict__ dst0, + scalar_t* __restrict__ dst1, + const uint8_t* __restrict__ fp8_storage, + const index_t* __restrict__ ind, + const bool* __restrict__ valid_mask, + int N, + int K, + int Kv, + int64_t block_size, + int64_t storage_block_stride_bytes, + int ld_dst0, + int ld_dst1) { + const int NB = div_up(N, 16); + const int KB = K / 32; + const int KBv = std::min(Kv / 32, KB); + for (int nb = 0; nb < NB; ++nb) { + const uint8_t* row_base[16]; + const uint8_t* scale_base[16]; + const int nb_size = std::min(N - nb * 16, 16); + for (int n = 0; n < nb_size; ++n) { + if (!valid_mask[nb * 16 + n]) { + row_base[n] = nullptr; + scale_base[n] = nullptr; + continue; + } + init_fp8_kvcache_tile_rows( + fp8_storage, block_size, storage_block_stride_bytes, static_cast(ind[nb * 16 + n]), row_base[n], + scale_base[n]); + } + auto load_vec = [&row_base, &scale_base](int n, int64_t idx, int dim_offset) { + UNUSED(idx); + return load_fp8_kvcache_32_from_row(row_base[n], scale_base[n], dim_offset); + }; + for (int kb = 0; kb < KBv; ++kb) { + sparse_pack_vnni_Nx32( + dst0 + ((kb * 32) >> 1) * ld_dst0 * 2 + nb * 16 * 2, + dst1 + ((nb * 16) >> 1) * ld_dst1 * 2 + kb * 32 * 2, + ind + nb * 16, + valid_mask + nb * 16, + nb_size, + kb * 32, + ld_dst0, + ld_dst1, + load_vec); + } + for (int kb = KBv; kb < KB; ++kb) { + sparse_pack_vnni_Nx32( + dst0 + ((kb * 32) >> 1) * ld_dst0 * 2 + nb * 16 * 2, + dst1 + ((nb * 16) >> 1) * ld_dst1 * 2 + kb * 32 * 2, + ind + nb * 16, + valid_mask + nb * 16, + nb_size, + kb * 32, + ld_dst0, + ld_dst1, + load_vec); + } + } +} + +template +inline void fmla_fill_stub(scalar_t* __restrict__ out, float val, int64_t size) { + using Vec = at::vec::Vectorized; + constexpr int kVecSize = Vec::size(); + const Vec data_vec = Vec(static_cast(val)); + int64_t d = 0; +#pragma GCC unroll 4 + for (; d <= size - kVecSize; d += kVecSize) { + data_vec.store(out + d); + } + if (size - d > 0) { + data_vec.store(out + d, size - d); + } +} + +template +inline void fmla_copy_stub(scalar_t* __restrict__ out, const float* __restrict__ input) { + static_assert(BLOCK_N % 32 == 0); + using bVec = at::vec::Vectorized; + using fVec = at::vec::Vectorized; + constexpr int COLS = BLOCK_N / 16; + auto store = [&](auto i) { + constexpr int col = i % COLS; + if constexpr (col % 2 == 0) { + fVec a0 = fVec::loadu(input + col * 16); + fVec a1 = fVec::loadu(input + col * 16 + 16); + bVec out_bvec = convert_from_float_ext(a0, a1); + out_bvec.store(out + col * 16); + } + }; + Unroll{}(store); +} + +template +inline void fmla_finalize_out(scalar_t* __restrict__ out, const float* __restrict__ acc, float inv_s, int64_t size) { + using bVec = at::vec::Vectorized; + using fVec = at::vec::Vectorized; + constexpr int kVecSize = bVec::size(); + const fVec s_fvec = fVec(inv_s); + int64_t d = 0; +#pragma GCC unroll 4 + for (; d <= size - kVecSize; d += kVecSize) { + fVec a0 = fVec::loadu(acc + d) * s_fvec; + fVec a1 = fVec::loadu(acc + d + fVec::size()) * s_fvec; + bVec out_bvec = convert_from_float_ext(a0, a1); + out_bvec.store(out + d); + } + for (; d < size; ++d) { + out[d] = static_cast(acc[d] * inv_s); + } +} + +// --------------------------------------------------------------------------- +// query : [num_tokens, H_q, D_qk] bf16 (SGLang's `[B, S_q, H_q, D]` +// with S_q folded away -- vLLM's flat per-token convention) +// indices : [num_tokens, topk_main] int64 (flat, block-table-resolved +// slot ids, -1 = invalid; SGLang's "main"/window index set) +// extra_* : same, for the optional "extra"/compressed index set +// attn_sink: [H_q] float32 +// output : [num_tokens, H_q, D_v] bf16 +// --------------------------------------------------------------------------- + +template +void sparse_mla_decode_kernel_impl( + scalar_t* __restrict__ output, + float* __restrict__ lse_out, + const scalar_t* __restrict__ query, + const uint8_t* __restrict__ k_main_fp8, + const uint8_t* __restrict__ k_extra_fp8, + const index_t* __restrict__ indices, + const index_t* __restrict__ extra_indices, + const float* __restrict__ attn_sink, + scalar_t* __restrict__ buffer, + int64_t batches, + int64_t s_q, + int64_t num_heads, + int64_t head_size, + int64_t head_size_v, + int64_t topk_main, + int64_t topk_extra, + int64_t total_tokens_main, + int64_t total_tokens_extra, + int64_t q_strideB, + int64_t q_strideH, + int64_t k_main_block_size, + int64_t k_extra_block_size, + int64_t k_main_block_stride_bytes, + int64_t k_extra_block_stride_bytes, + int64_t idx_strideB, + int64_t extra_idx_strideB, + float scaling, + int64_t buffer_size_per_thread) { + using Vec = at::vec::Vectorized; + + // partition heads + constexpr int64_t kBLOCK_H_MAX = 16; + const int64_t BLOCK_H = (batches * s_q) >= 16 ? kBLOCK_H_MAX : 8; + const int64_t num_h_blocks = div_up(num_heads, BLOCK_H); + + // parallel on [B, S_q, head_block] -- with s_q == 1 (vLLM's flat call + // convention) this collapses to a plain [num_tokens, head_block] loop, + // one flat iteration per token with no padding to skip. + at::parallel_for(0, batches * s_q * num_h_blocks, 0, [&](int64_t begin, int64_t end) { + int64_t bs{0}, sq{0}, hb{0}; + data_index_init(begin, bs, batches, sq, s_q, hb, num_h_blocks); + + int tid = at::get_thread_num(); + scalar_t* __restrict__ Btmp0 = buffer + tid * buffer_size_per_thread; // K packed + scalar_t* __restrict__ Btmp1 = Btmp0 + BLOCK_N * head_size; // V packed + // f32 V accumulator follows the bf16 packing region (reinterpret cast). + float* __restrict__ v_acc_local = reinterpret_cast(Btmp1 + BLOCK_N * head_size_v); + fmla_fill_stub(Btmp1, 0.f, BLOCK_N * head_size_v); // initialize V padding + + alignas(64) float s_i[kBLOCK_H_MAX * BLOCK_N]; + float* __restrict__ s_delta = s_i; + alignas(64) scalar_t s_delta2[kBLOCK_H_MAX * BLOCK_N]; + + alignas(64) float s_prime[kBLOCK_H_MAX]; + alignas(64) float m_prime[kBLOCK_H_MAX]; + + for (int64_t i = begin; i < end; ++i) { + const int64_t h_start = hb * BLOCK_H; + const int64_t h_end = std::min(h_start + BLOCK_H, num_heads); + const int64_t h_size = h_end - h_start; + + const scalar_t* __restrict__ q_ptr = query + bs * q_strideB + h_start * q_strideH; + const index_t* __restrict__ idx_ptr = indices + bs * idx_strideB; + const index_t* __restrict__ extra_idx_ptr = extra_indices == nullptr ? nullptr : extra_indices + bs * extra_idx_strideB; + + fmla_fill_stub(s_prime, 0.f, BLOCK_H); + fmla_fill_stub(m_prime, -std::numeric_limits::infinity(), BLOCK_H); + for (int64_t h = 0; h < h_size; ++h) { + fmla_fill_stub(v_acc_local + h * head_size_v, 0.f, head_size_v); + } + + auto process_cache = [&](const uint8_t* __restrict__ fp8_ptr, + const index_t* __restrict__ cur_idx_ptr, + int64_t topk_count, + int64_t total_tokens, + int64_t fp8_block_size, + int64_t fp8_block_stride_bytes) { + if (fp8_ptr == nullptr || cur_idx_ptr == nullptr || topk_count == 0) { + return; + } + const int64_t topk_limit = topk_count; + + for (int64_t n = 0; n < topk_limit; n += BLOCK_N) { + int64_t n_size = std::min(BLOCK_N, topk_limit - n); + const int64_t padded_n_size = div_up(int(n_size), TILE_K) * TILE_K; + bool valid_mask[BLOCK_N]; + bool has_valid = false; + for (int64_t k = 0; k < n_size; ++k) { + const bool valid = is_valid_sparse_index(cur_idx_ptr[n + k], n + k, topk_limit, total_tokens); + valid_mask[k] = valid; + has_valid |= valid; + } + if (!has_valid) { + continue; + } + + // Pack K (BLOCK_N rows via gather) into Btmp0 (key, vnni) and Btmp1 + // (value, vnni), dequantizing FP8 cache entries while packing. + sparse_pack_vnni_fp8( + Btmp0, Btmp1, fp8_ptr, cur_idx_ptr + n, valid_mask, static_cast(n_size), static_cast(head_size), + static_cast(head_size_v), fp8_block_size, fp8_block_stride_bytes, static_cast(BLOCK_N), + static_cast(head_size_v)); + + // Q @ K + at::native::cpublas::brgemm( + /* M */ h_size, + /* N */ n_size, + /* K */ head_size, + /* lda */ q_strideH, + /* ldb */ BLOCK_N, + /* ldc */ BLOCK_N, + /* add_C */ false, + /* A */ q_ptr, + /* B */ Btmp0, + /* C */ s_i); + + const Vec scale_vec = Vec(scaling); + for (int64_t h = 0; h < h_size; ++h) { + float* row = s_i + h * BLOCK_N; + at::vec::map([scale_vec](Vec x) { return x * scale_vec; }, row, row, n_size); + + for (int64_t k = 0; k < n_size; ++k) { + if (!valid_mask[k]) { + row[k] = -std::numeric_limits::infinity(); + } + } + + // online softmax update + float m_i = at::vec::reduce_all([](Vec& x, Vec& y) { return at::vec::maximum(x, y); }, row, n_size); + m_i = std::max(m_i, m_prime[h]); + + if (!std::isfinite(m_i)) { + fmla_fill_stub(s_delta + h * BLOCK_N, 0.f, padded_n_size); + fmla_copy_stub(s_delta2 + h * BLOCK_N, s_delta + h * BLOCK_N); + continue; + } + + const float m_delta = std::exp(m_prime[h] - m_i); + at::vec::map([m_i](Vec x) { return (x - Vec(m_i)).exp_u20(); }, s_delta + h * BLOCK_N, row, n_size); + + s_prime[h] *= m_delta; + s_prime[h] += at::vec::reduce_all([](Vec& x, Vec& y) { return x + y; }, s_delta + h * BLOCK_N, n_size); + + m_prime[h] = m_i; + + float scale_m = m_delta; + at::vec::map( + [scale_m](Vec x) { return x * Vec(scale_m); }, v_acc_local + h * head_size_v, v_acc_local + h * head_size_v, + head_size_v); + + fmla_fill_stub(s_delta + h * BLOCK_N + n_size, 0.f, padded_n_size - n_size); + fmla_copy_stub(s_delta2 + h * BLOCK_N, s_delta + h * BLOCK_N); + } + + // V' <- s_delta @ V + V' (accumulate into v_acc_local at f32) + at::native::cpublas::brgemm( + /* M */ h_size, + /* N */ head_size_v, + /* K */ padded_n_size, + /* lda */ BLOCK_N, + /* ldb */ head_size_v, + /* ldc */ head_size_v, + /* add_C */ true, + /* A */ s_delta2, + /* B */ Btmp1, + /* C */ v_acc_local); + } + }; + + process_cache(k_main_fp8, idx_ptr, topk_main, total_tokens_main, k_main_block_size, k_main_block_stride_bytes); + process_cache( + k_extra_fp8, extra_idx_ptr, topk_extra, total_tokens_extra, k_extra_block_size, k_extra_block_stride_bytes); + + // Apply attention sink correction directly on the output and lse. + // out *= exp(lse_no_sink) / (exp(lse_no_sink) + exp(attn_sink)) + // = 1 / (1 + exp(attn_sink - lse_no_sink)) + // (Algebraically equivalent to treating attn_sink as a virtual, + // zero-value extra key inside the softmax.) + for (int64_t h = 0; h < h_size; ++h) { + const int64_t hh = h_start + h; + const bool lonely = !std::isfinite(m_prime[h]) || s_prime[h] == 0.f; + float lse_val = lonely ? std::numeric_limits::infinity() : (m_prime[h] + std::log(s_prime[h])); + float inv_s = lonely ? 0.f : (1.f / s_prime[h]); + + if (!lonely && attn_sink != nullptr) { + const float sink = attn_sink[hh]; + const float corr = 1.f / (1.f + std::exp(sink - lse_val)); + inv_s *= corr; + } + + scalar_t* out_row = output + bs * (s_q * num_heads * head_size_v) + sq * (num_heads * head_size_v) + hh * head_size_v; + if (lonely) { + fmla_fill_stub(out_row, 0.f, head_size_v); + } else { + fmla_finalize_out(out_row, v_acc_local + h * head_size_v, inv_s, head_size_v); + } + + lse_out[bs * num_heads * s_q + hh * s_q + sq] = lse_val; + } + + data_index_step(bs, batches, sq, s_q, hb, num_h_blocks); + } + at::native::cpublas::brgemm_release(); + }); +} + +} // anonymous namespace + +void flash_mla_with_kvcache_cpu( + at::Tensor& out, + at::Tensor& q, + at::Tensor& window_cache_2d, + at::Tensor& window_slots, + int64_t window_block_size, + at::Tensor& compressed_cache_2d, + at::Tensor& compressed_slots, + int64_t compressed_block_size, + at::Tensor& attn_sink, + double scale) { + TORCH_CHECK( + q.dim() == 3 && q.size(2) == kHeadDim, + "flash_mla_with_kvcache_cpu: q must be [num_tokens, num_heads, 512]"); + TORCH_CHECK( + out.dim() == 3 && out.size(0) == q.size(0) && out.size(1) == q.size(1) && out.size(2) == kHeadDim, + "flash_mla_with_kvcache_cpu: out must be [num_tokens, num_heads, 512]"); + TORCH_CHECK( + out.scalar_type() == at::kBFloat16 && out.is_contiguous(), + "flash_mla_with_kvcache_cpu: out must be bf16 and contiguous " + "(pre-allocate once in Python -- see DeepseekV4CPUAttention.forward_mqa)"); + TORCH_CHECK(out.device().is_cpu(), "flash_mla_with_kvcache_cpu: out must be a CPU tensor"); + TORCH_CHECK( + q.scalar_type() == at::kBFloat16, + "flash_mla_with_kvcache_cpu: q must be bf16 (the AMX/VNNI kernel's " + "FP8->bf16 dequant intrinsics are bf16-specific)"); + TORCH_CHECK( + window_slots.dim() == 2 && window_slots.size(0) == q.size(0), + "flash_mla_with_kvcache_cpu: window_slots must be [num_tokens, num_window]"); + TORCH_CHECK( + compressed_slots.dim() == 2 && compressed_slots.size(0) == q.size(0), + "flash_mla_with_kvcache_cpu: compressed_slots must be [num_tokens, num_compressed]"); + TORCH_CHECK( + window_slots.scalar_type() == at::kLong && compressed_slots.scalar_type() == at::kLong, + "flash_mla_with_kvcache_cpu: slot tensors must be int64"); + TORCH_CHECK( + attn_sink.numel() == q.size(1), "flash_mla_with_kvcache_cpu: attn_sink must have num_heads entries"); + TORCH_CHECK( + attn_sink.scalar_type() == at::kFloat && attn_sink.is_contiguous(), + "flash_mla_with_kvcache_cpu: attn_sink must be float32 and contiguous " + "(convert once in Python -- see DeepseekV4CPUAttention.forward_mqa)"); + TORCH_CHECK( + window_slots.is_contiguous() && compressed_slots.is_contiguous(), + "flash_mla_with_kvcache_cpu: window_slots/compressed_slots must be contiguous"); + TORCH_CHECK( + q.device().is_cpu() && window_cache_2d.device().is_cpu() && window_slots.device().is_cpu() && + compressed_cache_2d.device().is_cpu() && compressed_slots.device().is_cpu() && + attn_sink.device().is_cpu(), + "flash_mla_with_kvcache_cpu: all inputs must be CPU tensors"); + + const int64_t num_tokens = q.size(0); + const int64_t num_heads = q.size(1); + const int64_t num_window = window_slots.size(1); + const int64_t num_compressed = compressed_slots.size(1); + TORCH_CHECK(num_window > 0, "flash_mla_with_kvcache_cpu: window_slots must have at least one column"); + + // Row stride is threaded through explicitly rather than assumed to equal + // block_size * (token_data_bytes + scale_bytes_per_token): a paged/pooled + // cache buffer's physical row stride is not guaranteed to match the + // logical per-block byte count (matches store_cache.cpp's precedent). + const int64_t window_cache_row_stride = window_cache_2d.stride(0); + const int64_t compressed_cache_row_stride = num_compressed > 0 ? compressed_cache_2d.stride(0) : 0; + + // AMX/VNNI path: flatten (batch, s_q) -> batches = num_tokens, s_q = 1. + CHECK_LAST_DIM_CONTIGUOUS_INPUT(q); + constexpr int64_t BLOCK_N = 128; // multiple of 32 for AMX brgemm + constexpr int64_t kBLOCK_H_MAX = 16; + static_assert(kHeadDim % 32 == 0, "head_dim must be a multiple of 32"); + + const int num_threads = at::get_num_threads(); + // Layout per thread (in bf16 elements): + // [Btmp0 : BLOCK_N * D] [Btmp1 : BLOCK_N * D] [v_acc_local : kBLOCK_H_MAX * D floats] + // f32 takes 2 bf16 elements -> multiply by 2. + const int64_t buffer_size_per_thread = BLOCK_N * kHeadDim + BLOCK_N * kHeadDim + 2 * kBLOCK_H_MAX * kHeadDim; + auto buffer = at::empty({num_threads, buffer_size_per_thread}, q.options()); + // lse is an internal scratch (attn_sink correction needs it per-token/head) + // -- not surfaced to the Python call site, which only consumes `out`. + auto lse = at::empty({num_tokens, num_heads}, q.options().dtype(at::kFloat)); + + const int64_t q_strideB = q.stride(0); + const int64_t q_strideH = q.stride(1); + const int64_t idx_strideB = window_slots.stride(0); + const int64_t extra_idx_strideB = num_compressed > 0 ? compressed_slots.stride(0) : 0; + + const int64_t total_tokens_main = window_cache_2d.size(0) * window_block_size; + const int64_t total_tokens_extra = num_compressed > 0 ? compressed_cache_2d.size(0) * compressed_block_size : 0; + + sparse_mla_decode_kernel_impl( + out.data_ptr(), + lse.data_ptr(), + q.data_ptr(), + window_cache_2d.data_ptr(), + num_compressed > 0 ? compressed_cache_2d.data_ptr() : nullptr, + window_slots.data_ptr(), + num_compressed > 0 ? compressed_slots.data_ptr() : nullptr, + attn_sink.data_ptr(), + buffer.data_ptr(), + /* batches */ num_tokens, + /* s_q */ 1, + num_heads, + /* head_size */ kHeadDim, + /* head_size_v */ kHeadDim, + /* topk_main */ num_window, + /* topk_extra */ num_compressed, + total_tokens_main, + total_tokens_extra, + q_strideB, + q_strideH, + window_block_size, + num_compressed > 0 ? compressed_block_size : 1, + window_cache_row_stride, + compressed_cache_row_stride, + idx_strideB, + extra_idx_strideB, + static_cast(scale), + buffer_size_per_thread); +} + diff --git a/csrc/cpu/sgl-kernels/indexer.cpp b/csrc/cpu/sgl-kernels/indexer.cpp new file mode 100644 index 0000000000..038304c62d --- /dev/null +++ b/csrc/cpu/sgl-kernels/indexer.cpp @@ -0,0 +1,276 @@ +// vLLM-native: DeepSeek-V4 sparse indexer Q-side RoPE + FP8 quant kernel. +// Ports the math of vLLM's own GPU/Triton reference +// (`_fused_indexer_q_rope_quant_kernel`, common/ops/fused_indexer_q.py) -- +// not a port of SGLang's own indexer-Q kernel, since that applies RoPE plus +// a 128-point Hadamard rotation the real DeepSeek-V4-Flash checkpoint +// doesn't use (verified: it applies RoPE only, no Hadamard step). +// +// Per (token, head): GPT-J interleaved RoPE on the last `rot_dim` dims; +// leading NoPE dims pass through unrotated. RoPE-rotated values are bf16- +// round-tripped before the absmax reduction (matching the K-side compressor +// kernel), NoPE values are not. One UE8M0 scale per (token, head) covers the +// full head_dim; since the FP8 Q tensor has no companion scale tensor, the +// scale is folded into `index_weights_out` instead: +// weights_out = weights * q_scale * softmax_scale * head_scale +// MXFP4 (`use_fp4=True`) has per-block scales that can't be folded this way +// and stays on triton-cpu -- out of scope here. +// +// AVX512 path: NoPE copy+absmax and the FP8 scale/clamp use plain 16-wide +// vector ops (final FP8 cast stays scalar, matching SGLang's own +// `act_quant_cpu_impl`). The RoPE rotation vectorizes the interleaved +// (x_even, x_odd) pairs via in-register permutes (`_mm512_permute_ps` / +// `_mm512_permutexvar_ps`), the same technique compressor.cpp's +// `apply_gptj_rope_inplace` and store_cache.cpp's +// `inverse_gptj_rope_o_proj_cpu` use, since gather/scatter is comparatively +// expensive on AVX512. + +#include +#include +#include + +#include "common.h" +#include "vec.h" + +#if !defined(CPU_CAPABILITY_AVX512) +static_assert( + false, + "indexer.cpp is only ever compiled into the AVX512+AMX-flagged " + "_C target (see cmake/cpu_extension.cmake's VLLM_EXT_SRC_SGL) and " + "relies on that unconditionally -- CPU_CAPABILITY_AVX512 must be " + "defined here."); +#endif + +namespace { + +constexpr float kFp8Max = 448.0f; + +#if defined(CPU_CAPABILITY_AVX512) + +// fp32 -> bf16 -> fp32 round trip via the AVX512-BF16 conversion +// instruction, numerically matching `static_cast(at::BFloat16(x))` +// (round-to-nearest-even). +inline __m512 bf16_round_trip_avx512(__m512 x) { + return CVT_BF16_TO_FP32((__m256i)(_mm512_cvtneps_pbh(x))); +} + +inline __m512 abs_ps(__m512 x) { + return _mm512_andnot_ps(_mm512_set1_ps(-0.0f), x); +} + +#endif // CPU_CAPABILITY_AVX512 + +void fused_indexer_q_rope_quant_impl( + const int64_t* __restrict__ positions, // [num_tokens] + const float* __restrict__ index_q, // [num_tokens, num_heads, head_dim] + int64_t index_q_stride0, int64_t index_q_stride1, + const float* __restrict__ cos_sin_cache, // [max_pos, rot_dim] + int64_t cos_sin_stride, int64_t half_rot_dim, + at::Float8_e4m3fn* __restrict__ index_q_fp8, // [num_tokens, num_heads, + // head_dim] + int64_t index_q_fp8_stride0, int64_t index_q_fp8_stride1, int64_t head_dim, + const float* __restrict__ index_weights, // [num_tokens, num_heads] + int64_t index_weights_stride, float softmax_scale, float head_scale, + float* __restrict__ index_weights_out, // [num_tokens, num_heads] + int64_t index_weights_out_stride, int64_t num_tokens, int64_t num_heads) { + const int64_t rot_dim = 2 * half_rot_dim; + const int64_t nope_dim = head_dim - rot_dim; + + at::parallel_for( + 0, num_tokens * num_heads, 0, [&](int64_t begin, int64_t end) { + std::vector buf(head_dim); + for (int64_t idx = begin; idx < end; ++idx) { + const int64_t t = idx / num_heads; + const int64_t h = idx % num_heads; + + const int64_t pos = positions[t]; + const float* const cos = cos_sin_cache + pos * cos_sin_stride; + const float* const sin = cos + half_rot_dim; + + const float* const q = + index_q + t * index_q_stride0 + h * index_q_stride1; + + float amax = 1e-4f; +#if defined(CPU_CAPABILITY_AVX512) + { + __m512 vamax = _mm512_set1_ps(amax); + for (int64_t d = 0; d < nope_dim; d += 16) { + __m512 v = _mm512_loadu_ps(q + d); + _mm512_storeu_ps(&buf[d], v); + vamax = _mm512_max_ps(vamax, abs_ps(v)); + } + amax = _mm512_reduce_max_ps(vamax); + } +#endif + + const float* const q_rot = q + nope_dim; +#if defined(CPU_CAPABILITY_AVX512) + { + // Permute-based (no gather/scatter), same technique as + // compressor.cpp's apply_gptj_rope_inplace / store_cache.cpp's + // inverse_gptj_rope_o_proj_cpu: `_mm512_permute_ps` swaps each + // pair's (even, odd) lanes in one shuffle, and cos/sin -- loaded as + // 8 contiguous values -- are duplicated pairwise via + // `_mm512_permutexvar_ps` to align with that layout. Two batches of + // 8 pairs per iteration (16 total, matching the old gather step). + // Unlike the RoPE kernels above, the result here stays in its + // original interleaved layout, so no un-permute/scatter is needed + // before the bf16 round-trip and store. + const __m512i dup_idx = _mm512_set_epi32(7, 7, 6, 6, 5, 5, 4, 4, 3, + 3, 2, 2, 1, 1, 0, 0); + // Forward GPT-J rotation negates the sin broadcast on the even + // (real) lane of each pair; the odd (imag) lane keeps sin's sign. + const __m512 sign_mask = _mm512_castsi512_ps(_mm512_set_epi32( + 0, (int)0x80000000, 0, (int)0x80000000, 0, (int)0x80000000, 0, + (int)0x80000000, 0, (int)0x80000000, 0, (int)0x80000000, 0, + (int)0x80000000, 0, (int)0x80000000)); + __m512 vamax_rope = _mm512_setzero_ps(); + for (int64_t i = 0; i < half_rot_dim; i += 16) { + const float* const base = q_rot + 2 * i; + float* const dst = buf.data() + nope_dim + 2 * i; + + __m512 xv0 = _mm512_loadu_ps(base); + __m512 cos_b0 = _mm512_permutexvar_ps( + dup_idx, _mm512_castps256_ps512(_mm256_loadu_ps(cos + i))); + __m512 sin_b0 = _mm512_permutexvar_ps( + dup_idx, _mm512_castps256_ps512(_mm256_loadu_ps(sin + i))); + __m512 out0 = _mm512_fmadd_ps( + xv0, cos_b0, + _mm512_mul_ps(_mm512_permute_ps(xv0, 0xB1), + _mm512_xor_ps(sin_b0, sign_mask))); + // Match reference numerics: fp32 -> bf16 -> fp32 before the + // absmax. + out0 = bf16_round_trip_avx512(out0); + _mm512_storeu_ps(dst, out0); + vamax_rope = _mm512_max_ps(vamax_rope, abs_ps(out0)); + + __m512 xv1 = _mm512_loadu_ps(base + 16); + __m512 cos_b1 = _mm512_permutexvar_ps( + dup_idx, + _mm512_castps256_ps512(_mm256_loadu_ps(cos + i + 8))); + __m512 sin_b1 = _mm512_permutexvar_ps( + dup_idx, + _mm512_castps256_ps512(_mm256_loadu_ps(sin + i + 8))); + __m512 out1 = _mm512_fmadd_ps( + xv1, cos_b1, + _mm512_mul_ps(_mm512_permute_ps(xv1, 0xB1), + _mm512_xor_ps(sin_b1, sign_mask))); + out1 = bf16_round_trip_avx512(out1); + _mm512_storeu_ps(dst + 16, out1); + vamax_rope = _mm512_max_ps(vamax_rope, abs_ps(out1)); + } + amax = std::max(amax, _mm512_reduce_max_ps(vamax_rope)); + } +#endif + + const float exponent = std::ceil(std::log2(amax / kFp8Max)); + const float scale = std::exp2(exponent); + const float inv_scale = 1.0f / scale; + + at::Float8_e4m3fn* const out = + index_q_fp8 + t * index_q_fp8_stride0 + h * index_q_fp8_stride1; +#if defined(CPU_CAPABILITY_AVX512) + { + const __m512 vinv = _mm512_set1_ps(inv_scale); + const __m512 vmax = _mm512_set1_ps(kFp8Max); + const __m512 vmin = _mm512_set1_ps(-kFp8Max); + alignas(64) float scaled_buf[16]; + for (int64_t d2 = 0; d2 < head_dim; d2 += 16) { + __m512 v = _mm512_loadu_ps(&buf[d2]); + v = _mm512_min_ps(_mm512_max_ps(_mm512_mul_ps(v, vinv), vmin), + vmax); + _mm512_store_ps(scaled_buf, v); + for (int64_t j = 0; j < 16; ++j) + out[d2 + j] = at::Float8_e4m3fn(scaled_buf[j]); + } + } +#endif + + float w = index_weights[t * index_weights_stride + h]; + w *= scale; + w *= softmax_scale; + w *= head_scale; + index_weights_out[t * index_weights_out_stride + h] = w; + } + }); +} + +} // anonymous namespace + +void fused_indexer_q_rope_quant_cpu(at::Tensor& positions, at::Tensor& index_q, + at::Tensor& index_q_cos_sin_cache, + at::Tensor& index_q_fp8, + at::Tensor& index_weights, + double index_weights_softmax_scale, + double index_weights_head_scale, + at::Tensor& index_weights_out) { + TORCH_CHECK( + positions.dim() == 1 && positions.scalar_type() == at::kLong, + "fused_indexer_q_rope_quant_cpu: positions must be a 1D int64 tensor"); + TORCH_CHECK(positions.is_contiguous(), + "fused_indexer_q_rope_quant_cpu: positions must be contiguous"); + TORCH_CHECK(index_q.dim() == 3, + "fused_indexer_q_rope_quant_cpu: index_q must be 3D"); + TORCH_CHECK(index_q.scalar_type() == at::kFloat, + "fused_indexer_q_rope_quant_cpu: index_q must be float32 " + "(convert once in Python)"); + TORCH_CHECK(index_q.stride(2) == 1, + "fused_indexer_q_rope_quant_cpu: index_q must be contiguous in " + "the head_dim axis"); + TORCH_CHECK( + index_q_cos_sin_cache.dim() == 2 && index_q_cos_sin_cache.stride(1) == 1, + "fused_indexer_q_rope_quant_cpu: index_q_cos_sin_cache must be a " + "contiguous 2D tensor"); + TORCH_CHECK( + index_q_cos_sin_cache.scalar_type() == at::kFloat, + "fused_indexer_q_rope_quant_cpu: index_q_cos_sin_cache must be float32 " + "(convert once in Python)"); + TORCH_CHECK(index_q_fp8.sizes() == index_q.sizes() && + index_q_fp8.scalar_type() == at::kFloat8_e4m3fn && + index_q_fp8.stride(2) == 1, + "fused_indexer_q_rope_quant_cpu: index_q_fp8 must match " + "index_q's shape, be " + "float8_e4m3fn, and contiguous in the head_dim axis"); + TORCH_CHECK(index_weights.dim() == 2 && index_weights.stride(1) == 1, + "fused_indexer_q_rope_quant_cpu: index_weights must be a " + "contiguous-in-head 2D tensor"); + TORCH_CHECK(index_weights.scalar_type() == at::kFloat, + "fused_indexer_q_rope_quant_cpu: index_weights must be float32 " + "(convert once in Python)"); + TORCH_CHECK(index_weights_out.sizes() == index_weights.sizes() && + index_weights_out.scalar_type() == at::kFloat && + index_weights_out.stride(1) == 1, + "fused_indexer_q_rope_quant_cpu: index_weights_out must match " + "index_weights' shape, " + "be float32, and contiguous in the head axis"); + TORCH_CHECK(positions.device().is_cpu() && index_q.device().is_cpu() && + index_q_cos_sin_cache.device().is_cpu() && + index_q_fp8.device().is_cpu() && + index_weights.device().is_cpu() && + index_weights_out.device().is_cpu(), + "fused_indexer_q_rope_quant_cpu: all inputs must be CPU tensors"); + + const int64_t num_tokens = positions.size(0); + const int64_t num_heads = index_q.size(1); + const int64_t head_dim = index_q.size(2); + const int64_t half_rot_dim = index_q_cos_sin_cache.size(-1) / 2; + // Only ever called with head_dim=128, rot_dim=64 (DeepseekV4's sparse + // indexer Q side); the AVX512 loops below process 16 elements/pairs per + // iteration with no scalar remainder, so require even divisibility. + TORCH_CHECK(head_dim % 16 == 0 && (head_dim - 2 * half_rot_dim) % 16 == 0 && + half_rot_dim % 16 == 0, + "fused_indexer_q_rope_quant_cpu: head_dim, head_dim-rot_dim, and " + "rot_dim/2 " + "must be multiples of 16"); + + fused_indexer_q_rope_quant_impl( + positions.data_ptr(), index_q.data_ptr(), + index_q.stride(0), index_q.stride(1), + index_q_cos_sin_cache.data_ptr(), index_q_cos_sin_cache.stride(0), + half_rot_dim, index_q_fp8.data_ptr(), + index_q_fp8.stride(0), index_q_fp8.stride(1), head_dim, + index_weights.data_ptr(), index_weights.stride(0), + static_cast(index_weights_softmax_scale), + static_cast(index_weights_head_scale), + index_weights_out.data_ptr(), index_weights_out.stride(0), + num_tokens, num_heads); +} diff --git a/csrc/cpu/sgl-kernels/mhc.cpp b/csrc/cpu/sgl-kernels/mhc.cpp new file mode 100644 index 0000000000..632f395317 --- /dev/null +++ b/csrc/cpu/sgl-kernels/mhc.cpp @@ -0,0 +1,1046 @@ +// Adapted from +// https://github.com/sgl-project/sglang/blob/main/python/sglang/kernels/aot/csrc/cpu/mhc.cpp +// +// DeepSeek-V4 mHC (multi-residual-stream hyper-connection) gating kernels. +// +// Difference from upstream: vLLM's `mhc_pre_torch` reference (kernels/ +// mhc/torch.py) exposes separate `hc_pre_eps`/`hc_sinkhorn_eps`/ +// `hc_post_mult_value`/`n_splits` knobs this kernel doesn't implement -- +// only ever called with hc_pre_eps == hc_sinkhorn_eps, hc_post_mult_value +// == 2.0, n_splits == 1 (matches upstream), so the entry points take one +// merged `hc_eps` and hardcode the post-mix multiplier to 2.0. The Python +// wrapper (kernels/mhc/cpu.py) enforces the merged-eps assumption via +// TORCH_CHECK. +// +// clang-format off + +#include +#include +#include + +#include "common.h" +#include "vec.h" + +namespace { + +// Loop tiling heuristics. +constexpr int64_t kMhcTileCacheBudgetBytes = 36 * 1024; +constexpr int64_t kMhcDefaultTokenBlock = 32; + +inline int64_t round_to_nearest_multiple(int64_t value, int64_t multiple) { + return ((value + multiple / 2) / multiple) * multiple; +} + +inline int64_t choose_t_block(int64_t /*T*/) { + return kMhcDefaultTokenBlock; +} + +inline int64_t choose_k_block(int64_t weight_rows, int64_t vec_size, int64_t max_block) { + const int64_t bytes_per_k = (weight_rows + 1) * static_cast(sizeof(float)); + const int64_t raw = std::max(vec_size, kMhcTileCacheBudgetBytes / bytes_per_k); + const int64_t rounded = round_to_nearest_multiple(raw, vec_size); + return std::max(vec_size, std::min(max_block, rounded)); +} + +template +inline void parallel_mhc_phase1(int64_t n_t_blocks, int64_t n_k_blocks, const func_t& f) { + const int64_t nth = static_cast(at::get_num_threads()); + if (n_t_blocks >= nth) { + // Large prefill: enough token blocks exist to fill all workers. + parallel_2d_tiled(static_cast(n_t_blocks), static_cast(n_k_blocks), static_cast(nth), 1, f); + } else { + // Decode / small T: split both T and K to expose enough parallel tasks. + parallel_2d(static_cast(n_t_blocks), static_cast(n_k_blocks), f); + } +} + +// In-place Sinkhorn on a row-major [hc, hc] matrix. +inline void +sinkhorn_inplace(float* __restrict__ comb_matrix, int hc, int sinkhorn_iters, float eps, float* __restrict__ scratch) { + float* const row_sum = scratch; + float* const col_sum = scratch + hc; + + // First iteration. + for (int r = 0; r < hc; ++r) { + const float* row = comb_matrix + r * hc; + float row_acc = 0.f; +#pragma GCC unroll 4 + for (int c = 0; c < hc; ++c) + row_acc += row[c]; + row_sum[r] = row_acc; + } + for (int r = 0; r < hc; ++r) { + const float row_inv = 1.0f / row_sum[r]; + float* row = comb_matrix + r * hc; +#pragma GCC unroll 4 + for (int c = 0; c < hc; ++c) + row[c] = row[c] * row_inv + eps; + } + for (int c = 0; c < hc; ++c) { + float col_acc = 0.f; +#pragma GCC unroll 4 + for (int r = 0; r < hc; ++r) + col_acc += comb_matrix[r * hc + c]; + col_sum[c] = col_acc; + } + for (int r = 0; r < hc; ++r) +#pragma GCC unroll 4 + for (int c = 0; c < hc; ++c) + comb_matrix[r * hc + c] /= (col_sum[c] + eps); + + // Remaining iterations. + for (int iter = 1; iter < sinkhorn_iters; ++iter) { + for (int r = 0; r < hc; ++r) { + const float* row = comb_matrix + r * hc; + float row_acc = 0.f; +#pragma GCC unroll 4 + for (int c = 0; c < hc; ++c) + row_acc += row[c]; + row_sum[r] = row_acc; + } + for (int r = 0; r < hc; ++r) { + const float row_inv = 1.0f / (row_sum[r] + eps); + float* row = comb_matrix + r * hc; +#pragma GCC unroll 4 + for (int c = 0; c < hc; ++c) + row[c] *= row_inv; + } + for (int c = 0; c < hc; ++c) { + float col_acc = 0.f; +#pragma GCC unroll 4 + for (int r = 0; r < hc; ++r) + col_acc += comb_matrix[r * hc + c]; + col_sum[c] = col_acc; + } + for (int r = 0; r < hc; ++r) +#pragma GCC unroll 4 + for (int c = 0; c < hc; ++c) + comb_matrix[r * hc + c] /= (col_sum[c] + eps); + } +} + +// Parse mixes -> pre/post/comb and run sinkhorn on comb. +inline void parse_mixes_and_sinkhorn( + float* __restrict__ pre_gate, // [hc] output + float* __restrict__ post_gate, // [hc] output + float* __restrict__ comb_matrix, // [hc*hc] output (row-major) + const float* __restrict__ mixes, // [mix_hc] input + float s0, + float s1, + float s2, + const float* __restrict__ hc_base, + int sinkhorn_iters, + float eps, + int hc, + float* __restrict__ scratch) { +#pragma GCC unroll 4 + for (int h = 0; h < hc; ++h) { + const float v = mixes[h] * s0 + hc_base[h]; + pre_gate[h] = 1.f / (1.f + std::exp(-v)) + eps; + } +#pragma GCC unroll 4 + for (int h = 0; h < hc; ++h) { + const float v = mixes[hc + h] * s1 + hc_base[hc + h]; + post_gate[h] = 2.f / (1.f + std::exp(-v)); + } + float* const row_max = scratch; + for (int r = 0; r < hc; ++r) + row_max[r] = -std::numeric_limits::infinity(); + for (int r = 0; r < hc; ++r) { + float* row = comb_matrix + r * hc; + for (int c = 0; c < hc; ++c) { + const int idx = r * hc + c; + const float v = mixes[2 * hc + idx] * s2 + hc_base[2 * hc + idx]; + row[c] = v; + if (v > row_max[r]) row_max[r] = v; + } + } + for (int r = 0; r < hc; ++r) { + float* row = comb_matrix + r * hc; +#pragma GCC unroll 4 + for (int c = 0; c < hc; ++c) + row[c] = std::exp(row[c] - row_max[r]); + } + + sinkhorn_inplace(comb_matrix, hc, sinkhorn_iters, eps, scratch + hc); +} + +template +static void hc_post_fuse_impl( + scalar_t* __restrict__ out, // [T, hc, d] output + const scalar_t* __restrict__ x, // [T, d] input + const scalar_t* __restrict__ residual, // [T, hc, d] input + const float* __restrict__ post, // [T, hc] float32 + const float* __restrict__ comb, // [T, hc, hc] float32 row-major + int64_t T, + int64_t d, + int hc) { + using bVec = at::vec::Vectorized; + using fVec = at::vec::Vectorized; + constexpr int64_t kVecSize = bVec::size(); + + const int64_t nthreads = at::get_num_threads(); + const int64_t max_splits = std::max(int64_t(1), d / kVecSize); + const int64_t token_head_tasks = T * hc; + const int64_t k_splits = (token_head_tasks >= nthreads) ? int64_t(1) : std::min(nthreads, max_splits); + const int64_t k_chunk = (d + k_splits - 1) / k_splits; + + const int64_t fvec_floats = static_cast(sizeof(fVec) / sizeof(float)); + auto comb_fvec_storage = at::empty({nthreads * hc * fvec_floats}, at::kFloat); + fVec* const comb_fvec_base = reinterpret_cast(comb_fvec_storage.data_ptr()); + + at::parallel_for(0, T * hc * k_splits, 0, [&](int64_t begin, int64_t end) { + const int tid = static_cast(at::get_thread_num()); + fVec* const comb_fvec = comb_fvec_base + tid * hc; + + for (int64_t idx = begin; idx < end; ++idx) { + const int64_t task = idx / k_splits; + const int64_t split_id = idx % k_splits; + const int64_t t = task / hc; + const int h = static_cast(task % hc); + const int64_t k0 = split_id * k_chunk; + const int64_t k1 = std::min(k0 + k_chunk, d); + + const float post_val = post[t * hc + h]; + const float* comb_t = comb + t * hc * hc; + const scalar_t* x_t = x + t * d; + const scalar_t* res_t = residual + t * hc * d; + scalar_t* out_th = out + (t * hc + h) * d; + + const fVec post_fvec(post_val); +#pragma GCC unroll 4 + for (int i = 0; i < hc; ++i) + comb_fvec[i] = fVec(comb_t[i * hc + h]); + + int64_t k; + // bf16: convert and accumulate + for (k = k0; k <= k1 - kVecSize; k += kVecSize) { + fVec acc0, acc1; + std::tie(acc0, acc1) = at::vec::convert_to_float(bVec::loadu(x_t + k)); + acc0 = post_fvec * acc0; + acc1 = post_fvec * acc1; +#pragma GCC unroll 4 + for (int i = 0; i < hc; ++i) { + fVec r0, r1; + std::tie(r0, r1) = at::vec::convert_to_float(bVec::loadu(res_t + i * d + k)); + acc0 += comb_fvec[i] * r0; + acc1 += comb_fvec[i] * r1; + } + at::vec::convert_from_float(acc0, acc1).store(out_th + k); + } + if (k < k1) { + const int64_t rem = k1 - k; + fVec acc0, acc1; + std::tie(acc0, acc1) = at::vec::convert_to_float(bVec::loadu(x_t + k, rem)); + acc0 = post_fvec * acc0; + acc1 = post_fvec * acc1; +#pragma GCC unroll 4 + for (int i = 0; i < hc; ++i) { + fVec r0, r1; + std::tie(r0, r1) = at::vec::convert_to_float(bVec::loadu(res_t + i * d + k, rem)); + acc0 += comb_fvec[i] * r0; + acc1 += comb_fvec[i] * r1; + } + at::vec::convert_from_float(acc0, acc1).store(out_th + k, rem); + } + } + }); +} + +// Fused hc_pre path. +template +static void hc_pre_fuse_impl( + scalar_t* __restrict__ y, // [T, d] bf16 output + float* __restrict__ post_out, // [T, hc] float32 output + float* __restrict__ comb_out, // [T, hc*hc] float32 output + const scalar_t* __restrict__ x, // [T, hc, d] bf16 input + const float* __restrict__ hc_fn, // [mix_hc, hc_d] float32 row-major + const float* __restrict__ hc_scale, // [3] + const float* __restrict__ hc_base, // [mix_hc] + int64_t T, + int64_t d, + int hc, + int sinkhorn_iters, + float hc_eps, + float norm_eps) { + using bVec = at::vec::Vectorized; + using fVec = at::vec::Vectorized; + constexpr int64_t kVecSize = bVec::size(); // 32 bf16 + constexpr int64_t kFVecSize = fVec::size(); // 16 fp32 + constexpr int64_t kSmallTKBlockCap = 128; + constexpr int64_t kLargeTKBlockCap = 256; + constexpr int64_t kDVecsPerBlock = 16; + constexpr int64_t kCacheKBlockCap = kLargeTKBlockCap; + constexpr int64_t MAX_K_VECS = kCacheKBlockCap / kVecSize; + const int mix_hc = (2 + hc) * hc; + const int64_t hc_d = hc * d; + const float s0 = hc_scale[0], s1 = hc_scale[1], s2 = hc_scale[2]; + + const int64_t T_BLOCK = choose_t_block(T); + const int64_t num_threads = static_cast(at::get_num_threads()); + const int64_t K_BLOCK_CAP = (T >= num_threads) ? kLargeTKBlockCap : kSmallTKBlockCap; + const int64_t K_BLOCK = choose_k_block(mix_hc, kVecSize, K_BLOCK_CAP); + TORCH_CHECK( + K_BLOCK <= K_BLOCK_CAP && K_BLOCK_CAP <= kCacheKBlockCap && K_BLOCK % kVecSize == 0, + "hc_pre_fuse_impl: K_BLOCK must be within capacity and a multiple of vector size"); + const int64_t D_BLOCK = kDVecsPerBlock * kVecSize; + + const int64_t n_t_blocks = div_up(T, T_BLOCK); + const int64_t n_k_blocks = div_up(hc_d, K_BLOCK); + + // Partial buffer: [n_t_blocks, n_k_blocks, T_BLOCK, 1+mix_hc] + const int64_t partial_stride = 1 + mix_hc; + auto partials_tensor = at::empty({n_t_blocks * n_k_blocks * T_BLOCK * partial_stride}, at::kFloat); + float* const partials = partials_tensor.data_ptr(); + + // Per-thread scratch: [mix_buffer[mix_hc] | pre_gate[hc] | + // post_gate[hc] | comb_matrix[hc*hc] | sinkhorn_scratch[3*hc]] + const int64_t per_thread_scratch = mix_hc + hc + hc + hc * hc + 3 * hc; + auto scratch_tensor = at::empty({num_threads * per_thread_scratch}, at::kFloat); + float* const scratch_base = scratch_tensor.data_ptr(); + + // Phase 1. + parallel_mhc_phase1(n_t_blocks, n_k_blocks, [&](int64_t tb0, int64_t tb1, int64_t kb0, int64_t kb1) { + for (int64_t kb = kb0; kb < kb1; ++kb) { + const int64_t k0 = kb * K_BLOCK; + const int64_t klen = std::min(K_BLOCK, hc_d - k0); + + for (int64_t tb = tb0; tb < tb1; ++tb) { + const int64_t t0 = tb * T_BLOCK; + const int64_t tlen = std::min(T_BLOCK, T - t0); + float* partial_block = partials + (tb * n_k_blocks + kb) * T_BLOCK * partial_stride; + + const float* w_ptr[mix_hc]; + for (int h = 0; h < mix_hc; ++h) + w_ptr[h] = hc_fn + h * hc_d + k0; + + // Stack cache is sized for the largest selected K block, but only the first + // n_full entries are touched for the current K_BLOCK. + fVec x_cache_lo[MAX_K_VECS], x_cache_hi[MAX_K_VECS]; + + for (int64_t tl = 0; tl < tlen; ++tl) { + const scalar_t* x_t = x + (t0 + tl) * hc_d + k0; + float* partial_row = partial_block + tl * partial_stride; + + const int64_t n_full = klen / kVecSize; + const int64_t k_tail = n_full * kVecSize; + const int64_t rem = klen - k_tail; + const int64_t rem0 = std::min(rem, kFVecSize); + const int64_t rem1 = rem - rem0; + + // Prefetch next token. + if (tl + 1 < tlen) { + const scalar_t* xn = x + (t0 + tl + 1) * hc_d + k0; + __builtin_prefetch(xn, 0, 0); + if (klen > kVecSize * 4) __builtin_prefetch(xn + kVecSize * 4, 0, 0); + } + + // h=0..3: convert/cache x and accumulate sq + four dots. + if (tl == 0) { + for (int pw = 1; pw < mix_hc; ++pw) + __builtin_prefetch(w_ptr[pw], 0, 3); + } + { + const float* w0 = w_ptr[0]; + const float* w1 = w_ptr[1]; + const float* w2 = w_ptr[2]; + const float* w3 = w_ptr[3]; + fVec sq_lo(0.f), sq_hi(0.f); + fVec d0_lo(0.f), d0_hi(0.f); + fVec d1_lo(0.f), d1_hi(0.f); + fVec d2_lo(0.f), d2_hi(0.f); + fVec d3_lo(0.f), d3_hi(0.f); + for (int64_t ki = 0; ki < n_full; ++ki) { + const int64_t k = ki * kVecSize; + std::tie(x_cache_lo[ki], x_cache_hi[ki]) = at::vec::convert_to_float(bVec::loadu(x_t + k)); + sq_lo += x_cache_lo[ki] * x_cache_lo[ki]; + sq_hi += x_cache_hi[ki] * x_cache_hi[ki]; + d0_lo += x_cache_lo[ki] * fVec::loadu(w0 + k); + d0_hi += x_cache_hi[ki] * fVec::loadu(w0 + k + kFVecSize); + d1_lo += x_cache_lo[ki] * fVec::loadu(w1 + k); + d1_hi += x_cache_hi[ki] * fVec::loadu(w1 + k + kFVecSize); + d2_lo += x_cache_lo[ki] * fVec::loadu(w2 + k); + d2_hi += x_cache_hi[ki] * fVec::loadu(w2 + k + kFVecSize); + d3_lo += x_cache_lo[ki] * fVec::loadu(w3 + k); + d3_hi += x_cache_hi[ki] * fVec::loadu(w3 + k + kFVecSize); + } + + fVec x_tail_lo(0.f), x_tail_hi(0.f); + if (rem > 0) { + std::tie(x_tail_lo, x_tail_hi) = at::vec::convert_to_float(bVec::loadu(x_t + k_tail, rem)); + sq_lo += x_tail_lo * x_tail_lo; + if (rem1 > 0) sq_hi += x_tail_hi * x_tail_hi; + d0_lo += x_tail_lo * fVec::loadu(w0 + k_tail, rem0); + d1_lo += x_tail_lo * fVec::loadu(w1 + k_tail, rem0); + d2_lo += x_tail_lo * fVec::loadu(w2 + k_tail, rem0); + d3_lo += x_tail_lo * fVec::loadu(w3 + k_tail, rem0); + if (rem1 > 0) { + d0_hi += x_tail_hi * fVec::loadu(w0 + k_tail + kFVecSize, rem1); + d1_hi += x_tail_hi * fVec::loadu(w1 + k_tail + kFVecSize, rem1); + d2_hi += x_tail_hi * fVec::loadu(w2 + k_tail + kFVecSize, rem1); + d3_hi += x_tail_hi * fVec::loadu(w3 + k_tail + kFVecSize, rem1); + } + } + + partial_row[0] = vec_reduce_sum(sq_lo + sq_hi); + partial_row[1 + 0] = vec_reduce_sum(d0_lo + d0_hi); + partial_row[1 + 1] = vec_reduce_sum(d1_lo + d1_hi); + partial_row[1 + 2] = vec_reduce_sum(d2_lo + d2_hi); + partial_row[1 + 3] = vec_reduce_sum(d3_lo + d3_hi); + + // h=4..mix_hc-1: reuse cached x, four dots per pass. + int h = 4; + for (; h + 3 < mix_hc; h += 4) { + const float* wa = w_ptr[h]; + const float* wb = w_ptr[h + 1]; + const float* wc = w_ptr[h + 2]; + const float* wd = w_ptr[h + 3]; + fVec a0(0.f), a1(0.f), b0(0.f), b1(0.f); + fVec c0(0.f), c1(0.f), d0(0.f), d1(0.f); + for (int64_t ki = 0; ki < n_full; ++ki) { + const int64_t k = ki * kVecSize; + a0 += x_cache_lo[ki] * fVec::loadu(wa + k); + b0 += x_cache_lo[ki] * fVec::loadu(wb + k); + c0 += x_cache_lo[ki] * fVec::loadu(wc + k); + d0 += x_cache_lo[ki] * fVec::loadu(wd + k); + a1 += x_cache_hi[ki] * fVec::loadu(wa + k + kFVecSize); + b1 += x_cache_hi[ki] * fVec::loadu(wb + k + kFVecSize); + c1 += x_cache_hi[ki] * fVec::loadu(wc + k + kFVecSize); + d1 += x_cache_hi[ki] * fVec::loadu(wd + k + kFVecSize); + } + if (rem > 0) { + a0 += x_tail_lo * fVec::loadu(wa + k_tail, rem0); + b0 += x_tail_lo * fVec::loadu(wb + k_tail, rem0); + c0 += x_tail_lo * fVec::loadu(wc + k_tail, rem0); + d0 += x_tail_lo * fVec::loadu(wd + k_tail, rem0); + if (rem1 > 0) { + a1 += x_tail_hi * fVec::loadu(wa + k_tail + kFVecSize, rem1); + b1 += x_tail_hi * fVec::loadu(wb + k_tail + kFVecSize, rem1); + c1 += x_tail_hi * fVec::loadu(wc + k_tail + kFVecSize, rem1); + d1 += x_tail_hi * fVec::loadu(wd + k_tail + kFVecSize, rem1); + } + } + partial_row[1 + h] = vec_reduce_sum(a0 + a1); + partial_row[1 + h + 1] = vec_reduce_sum(b0 + b1); + partial_row[1 + h + 2] = vec_reduce_sum(c0 + c1); + partial_row[1 + h + 3] = vec_reduce_sum(d0 + d1); + } + for (; h + 1 < mix_hc; h += 2) { + const float* wa = w_ptr[h]; + const float* wb = w_ptr[h + 1]; + fVec a0(0.f), a1(0.f), b0(0.f), b1(0.f); + for (int64_t ki = 0; ki < n_full; ++ki) { + const int64_t k = ki * kVecSize; + a0 += x_cache_lo[ki] * fVec::loadu(wa + k); + b0 += x_cache_lo[ki] * fVec::loadu(wb + k); + a1 += x_cache_hi[ki] * fVec::loadu(wa + k + kFVecSize); + b1 += x_cache_hi[ki] * fVec::loadu(wb + k + kFVecSize); + } + if (rem > 0) { + a0 += x_tail_lo * fVec::loadu(wa + k_tail, rem0); + b0 += x_tail_lo * fVec::loadu(wb + k_tail, rem0); + if (rem1 > 0) { + a1 += x_tail_hi * fVec::loadu(wa + k_tail + kFVecSize, rem1); + b1 += x_tail_hi * fVec::loadu(wb + k_tail + kFVecSize, rem1); + } + } + partial_row[1 + h] = vec_reduce_sum(a0 + a1); + partial_row[1 + h + 1] = vec_reduce_sum(b0 + b1); + } + // odd leftover + if (h < mix_hc) { + const float* w = w_ptr[h]; + fVec hd0(0.f), hd1(0.f); + for (int64_t ki = 0; ki < n_full; ++ki) { + const int64_t k = ki * kVecSize; + hd0 += x_cache_lo[ki] * fVec::loadu(w + k); + hd1 += x_cache_hi[ki] * fVec::loadu(w + k + kFVecSize); + } + if (rem > 0) { + hd0 += x_tail_lo * fVec::loadu(w + k_tail, rem0); + if (rem1 > 0) hd1 += x_tail_hi * fVec::loadu(w + k_tail + kFVecSize, rem1); + } + partial_row[1 + h] = vec_reduce_sum(hd0 + hd1); + } + } + } // tl + + // Prefetch next t-block head. + if (tb + 1 < tb1) { + const scalar_t* x_next_tb = x + ((tb + 1) * T_BLOCK) * hc_d + k0; + __builtin_prefetch(x_next_tb, 0, 0); + if (klen > kVecSize * 4) __builtin_prefetch(x_next_tb + kVecSize * 4, 0, 0); + } + } // tb + } // kb + }); // Phase 1 + + // Phase 2+3: reduce partials, sinkhorn, and combine. + at::parallel_for(0, T, 1, [&](int64_t begin, int64_t end) { + const int64_t tid = at::get_thread_num(); + fVec v0, v1, pre_fvec; + float* const thread_scratch = scratch_base + tid * per_thread_scratch; + float* const mix_buffer = thread_scratch; + float* const pre_gate = thread_scratch + mix_hc; + float* const post_gate = pre_gate + hc; + float* const comb_matrix = post_gate + hc; + float* const sinkhorn_scratch = comb_matrix + hc * hc; + + for (int64_t t = begin; t < end; ++t) { + const int64_t tb = t / T_BLOCK; + const int64_t tl = t % T_BLOCK; + + std::fill(mix_buffer, mix_buffer + mix_hc, 0.f); + double sq_total = 0.0; + for (int64_t kb = 0; kb < n_k_blocks; ++kb) { + const float* partial_row = partials + (tb * n_k_blocks + kb) * T_BLOCK * partial_stride + tl * partial_stride; + sq_total += static_cast(partial_row[0]); +#pragma GCC unroll 4 + for (int i = 0; i < mix_hc; ++i) + mix_buffer[i] += partial_row[1 + i]; + } + const float inv_rms = + static_cast(1.0 / std::sqrt(sq_total / static_cast(hc_d) + static_cast(norm_eps))); +#pragma GCC unroll 4 + for (int i = 0; i < mix_hc; ++i) + mix_buffer[i] *= inv_rms; + + parse_mixes_and_sinkhorn( + pre_gate, + post_gate, + comb_matrix, + mix_buffer, + s0, + s1, + s2, + hc_base, + sinkhorn_iters, + hc_eps, + hc, + sinkhorn_scratch); + + std::copy_n(post_gate, hc, post_out + t * hc); + std::copy_n(comb_matrix, hc * hc, comb_out + t * hc * hc); + + // Phase 3: d-tiled combine. + const scalar_t* x_t = x + t * hc_d; + scalar_t* y_t = y + t * d; + for (int64_t j0 = 0; j0 < d; j0 += D_BLOCK) { + const int64_t jlen = std::min(D_BLOCK, d - j0); + + int64_t kk = 0; + for (; kk <= jlen - kVecSize; kk += kVecSize) { + pre_fvec = fVec(pre_gate[0]); + std::tie(v0, v1) = at::vec::convert_to_float(bVec::loadu(x_t + j0 + kk)); + fVec acc0 = pre_fvec * v0; + fVec acc1 = pre_fvec * v1; +#pragma GCC unroll 4 + for (int h = 1; h < hc; ++h) { + pre_fvec = fVec(pre_gate[h]); + const scalar_t* x_th = x_t + (int64_t)h * d + j0 + kk; + std::tie(v0, v1) = at::vec::convert_to_float(bVec::loadu(x_th)); + acc0 += pre_fvec * v0; + acc1 += pre_fvec * v1; + } + at::vec::convert_from_float(acc0, acc1).store(y_t + j0 + kk); + } + if (kk < jlen) { + const int64_t rem = jlen - kk, rem1 = rem - std::min(rem, kFVecSize); + pre_fvec = fVec(pre_gate[0]); + std::tie(v0, v1) = at::vec::convert_to_float(bVec::loadu(x_t + j0 + kk, rem)); + fVec acc0 = pre_fvec * v0; + fVec acc1 = pre_fvec * v1; +#pragma GCC unroll 4 + for (int h = 1; h < hc; ++h) { + pre_fvec = fVec(pre_gate[h]); + const scalar_t* x_th = x_t + (int64_t)h * d + j0 + kk; + std::tie(v0, v1) = at::vec::convert_to_float(bVec::loadu(x_th, rem)); + acc0 += pre_fvec * v0; + acc1 += pre_fvec * v1; + } + at::vec::convert_from_float(acc0, rem1 > 0 ? acc1 : fVec(0.f)).store(y_t + j0 + kk, rem); + } + } // d-tile + } // t + }); // Phase 2+3 +} + +// Fused hc_head path. +template +static void hc_head_fuse_impl( + scalar_t* __restrict__ y, + const scalar_t* __restrict__ x, + const float* __restrict__ hc_fn, // [hc, hc*d] row-major + float hc_scale_val, + const float* __restrict__ hc_base, + int64_t T, + int64_t d, + int hc, + float hc_eps, + float norm_eps) { + using bVec = at::vec::Vectorized; + using fVec = at::vec::Vectorized; + constexpr int64_t kVecSize = bVec::size(); // 32 bf16 + constexpr int64_t kFVecSize = fVec::size(); // 16 fp32 + constexpr int64_t kDVecsPerBlock = 16; + const int64_t hc_d = hc * d; + + constexpr int64_t K_BLOCK_CAP = 1024; + constexpr int64_t MAX_K_VECS = K_BLOCK_CAP / kVecSize; + + const int64_t T_BLOCK = choose_t_block(T); + const int64_t num_threads = static_cast(at::get_num_threads()); + const int64_t K_BLOCK = choose_k_block(hc, kVecSize, K_BLOCK_CAP); + TORCH_CHECK( + K_BLOCK <= K_BLOCK_CAP && K_BLOCK % kVecSize == 0, + "hc_head_fuse_impl: K_BLOCK must be <= K_BLOCK_CAP and a multiple of vector size"); + const int64_t D_BLOCK = kDVecsPerBlock * kVecSize; + + const int64_t n_t_blocks = div_up(T, T_BLOCK); + const int64_t n_k_blocks = div_up(hc_d, K_BLOCK); + + // Partial buffer: [n_t_blocks, n_k_blocks, T_BLOCK, 1+hc] + // p[0] = partial sq, p[1..hc] = partial dots + const int64_t partial_stride = 1 + hc; + auto partials_tensor = at::empty({n_t_blocks * n_k_blocks * T_BLOCK * partial_stride}, at::kFloat); + float* const partials = partials_tensor.data_ptr(); + + // Per-thread scratch: [pre_gate[hc]] + const int64_t per_thread_head = hc; + auto scratch_tensor = at::empty({num_threads * per_thread_head}, at::kFloat); + float* const scratch_base = scratch_tensor.data_ptr(); + + // Phase 1. + parallel_mhc_phase1(n_t_blocks, n_k_blocks, [&](int64_t tb0, int64_t tb1, int64_t kb0, int64_t kb1) { + for (int64_t kb = kb0; kb < kb1; ++kb) { + const int64_t k0 = kb * K_BLOCK; + const int64_t klen = std::min(K_BLOCK, hc_d - k0); + + for (int64_t tb = tb0; tb < tb1; ++tb) { + const int64_t t0 = tb * T_BLOCK; + const int64_t tlen = std::min(T_BLOCK, T - t0); + float* partial_block = partials + (tb * n_k_blocks + kb) * T_BLOCK * partial_stride; + + const float* w_ptr[hc]; + for (int h = 0; h < hc; ++h) + w_ptr[h] = hc_fn + h * hc_d + k0; + + // Stack cache is sized for the largest selected K block, but only the first + // n_full entries are touched for the current K_BLOCK. + fVec x_cache_lo[MAX_K_VECS], x_cache_hi[MAX_K_VECS]; + + for (int64_t tl = 0; tl < tlen; ++tl) { + const scalar_t* x_t = x + (t0 + tl) * hc_d + k0; + float* partial_row = partial_block + tl * partial_stride; + + const int64_t n_full = klen / kVecSize; + const int64_t k_tail = n_full * kVecSize; + const int64_t rem = klen - k_tail; + const int64_t rem0 = std::min(rem, kFVecSize); + const int64_t rem1 = rem - rem0; + + // Prefetch next token. + if (tl + 1 < tlen) { + const scalar_t* xn = x + (t0 + tl + 1) * hc_d + k0; + __builtin_prefetch(xn, 0, 0); + if (klen > kVecSize * 4) __builtin_prefetch(xn + kVecSize * 4, 0, 0); + } + + // Convert/cache x and accumulate sq + multiple dots. + if (tl == 0) { + for (int pw = 1; pw < hc; ++pw) + __builtin_prefetch(w_ptr[pw], 0, 3); + } + { + fVec x_tail_lo(0.f), x_tail_hi(0.f); + int h = 0; + + if (hc >= 4) { + const float* w0 = w_ptr[0]; + const float* w1 = w_ptr[1]; + const float* w2 = w_ptr[2]; + const float* w3 = w_ptr[3]; + fVec sq_lo(0.f), sq_hi(0.f); + fVec d0_lo(0.f), d0_hi(0.f); + fVec d1_lo(0.f), d1_hi(0.f); + fVec d2_lo(0.f), d2_hi(0.f); + fVec d3_lo(0.f), d3_hi(0.f); + for (int64_t ki = 0; ki < n_full; ++ki) { + const int64_t k = ki * kVecSize; + std::tie(x_cache_lo[ki], x_cache_hi[ki]) = at::vec::convert_to_float(bVec::loadu(x_t + k)); + sq_lo += x_cache_lo[ki] * x_cache_lo[ki]; + sq_hi += x_cache_hi[ki] * x_cache_hi[ki]; + d0_lo += x_cache_lo[ki] * fVec::loadu(w0 + k); + d0_hi += x_cache_hi[ki] * fVec::loadu(w0 + k + kFVecSize); + d1_lo += x_cache_lo[ki] * fVec::loadu(w1 + k); + d1_hi += x_cache_hi[ki] * fVec::loadu(w1 + k + kFVecSize); + d2_lo += x_cache_lo[ki] * fVec::loadu(w2 + k); + d2_hi += x_cache_hi[ki] * fVec::loadu(w2 + k + kFVecSize); + d3_lo += x_cache_lo[ki] * fVec::loadu(w3 + k); + d3_hi += x_cache_hi[ki] * fVec::loadu(w3 + k + kFVecSize); + } + if (rem > 0) { + std::tie(x_tail_lo, x_tail_hi) = at::vec::convert_to_float(bVec::loadu(x_t + k_tail, rem)); + sq_lo += x_tail_lo * x_tail_lo; + if (rem1 > 0) sq_hi += x_tail_hi * x_tail_hi; + d0_lo += x_tail_lo * fVec::loadu(w0 + k_tail, rem0); + d1_lo += x_tail_lo * fVec::loadu(w1 + k_tail, rem0); + d2_lo += x_tail_lo * fVec::loadu(w2 + k_tail, rem0); + d3_lo += x_tail_lo * fVec::loadu(w3 + k_tail, rem0); + if (rem1 > 0) { + d0_hi += x_tail_hi * fVec::loadu(w0 + k_tail + kFVecSize, rem1); + d1_hi += x_tail_hi * fVec::loadu(w1 + k_tail + kFVecSize, rem1); + d2_hi += x_tail_hi * fVec::loadu(w2 + k_tail + kFVecSize, rem1); + d3_hi += x_tail_hi * fVec::loadu(w3 + k_tail + kFVecSize, rem1); + } + } + partial_row[0] = vec_reduce_sum(sq_lo + sq_hi); + partial_row[1 + 0] = vec_reduce_sum(d0_lo + d0_hi); + partial_row[1 + 1] = vec_reduce_sum(d1_lo + d1_hi); + partial_row[1 + 2] = vec_reduce_sum(d2_lo + d2_hi); + partial_row[1 + 3] = vec_reduce_sum(d3_lo + d3_hi); + h = 4; + } else { + const float* w0 = w_ptr[0]; + fVec sq_lo(0.f), sq_hi(0.f); + fVec d0(0.f), d1(0.f); + for (int64_t ki = 0; ki < n_full; ++ki) { + const int64_t k = ki * kVecSize; + std::tie(x_cache_lo[ki], x_cache_hi[ki]) = at::vec::convert_to_float(bVec::loadu(x_t + k)); + sq_lo += x_cache_lo[ki] * x_cache_lo[ki]; + sq_hi += x_cache_hi[ki] * x_cache_hi[ki]; + d0 += x_cache_lo[ki] * fVec::loadu(w0 + k); + d1 += x_cache_hi[ki] * fVec::loadu(w0 + k + kFVecSize); + } + if (rem > 0) { + std::tie(x_tail_lo, x_tail_hi) = at::vec::convert_to_float(bVec::loadu(x_t + k_tail, rem)); + sq_lo += x_tail_lo * x_tail_lo; + if (rem1 > 0) sq_hi += x_tail_hi * x_tail_hi; + d0 += x_tail_lo * fVec::loadu(w0 + k_tail, rem0); + if (rem1 > 0) d1 += x_tail_hi * fVec::loadu(w0 + k_tail + kFVecSize, rem1); + } + partial_row[0] = vec_reduce_sum(sq_lo + sq_hi); + partial_row[1 + 0] = vec_reduce_sum(d0 + d1); + h = 1; + } + + // h: reuse cached x, four dots per pass. + for (; h + 3 < hc; h += 4) { + const float* wa = w_ptr[h]; + const float* wb = w_ptr[h + 1]; + const float* wc = w_ptr[h + 2]; + const float* wd = w_ptr[h + 3]; + fVec a0(0.f), a1(0.f), b0(0.f), b1(0.f); + fVec c0(0.f), c1(0.f), d0(0.f), d1(0.f); + for (int64_t ki = 0; ki < n_full; ++ki) { + const int64_t k = ki * kVecSize; + a0 += x_cache_lo[ki] * fVec::loadu(wa + k); + b0 += x_cache_lo[ki] * fVec::loadu(wb + k); + c0 += x_cache_lo[ki] * fVec::loadu(wc + k); + d0 += x_cache_lo[ki] * fVec::loadu(wd + k); + a1 += x_cache_hi[ki] * fVec::loadu(wa + k + kFVecSize); + b1 += x_cache_hi[ki] * fVec::loadu(wb + k + kFVecSize); + c1 += x_cache_hi[ki] * fVec::loadu(wc + k + kFVecSize); + d1 += x_cache_hi[ki] * fVec::loadu(wd + k + kFVecSize); + } + if (rem > 0) { + a0 += x_tail_lo * fVec::loadu(wa + k_tail, rem0); + b0 += x_tail_lo * fVec::loadu(wb + k_tail, rem0); + c0 += x_tail_lo * fVec::loadu(wc + k_tail, rem0); + d0 += x_tail_lo * fVec::loadu(wd + k_tail, rem0); + if (rem1 > 0) { + a1 += x_tail_hi * fVec::loadu(wa + k_tail + kFVecSize, rem1); + b1 += x_tail_hi * fVec::loadu(wb + k_tail + kFVecSize, rem1); + c1 += x_tail_hi * fVec::loadu(wc + k_tail + kFVecSize, rem1); + d1 += x_tail_hi * fVec::loadu(wd + k_tail + kFVecSize, rem1); + } + } + partial_row[1 + h] = vec_reduce_sum(a0 + a1); + partial_row[1 + h + 1] = vec_reduce_sum(b0 + b1); + partial_row[1 + h + 2] = vec_reduce_sum(c0 + c1); + partial_row[1 + h + 3] = vec_reduce_sum(d0 + d1); + } + for (; h + 1 < hc; h += 2) { + const float* wa = w_ptr[h]; + const float* wb = w_ptr[h + 1]; + fVec a0(0.f), a1(0.f), b0(0.f), b1(0.f); + for (int64_t ki = 0; ki < n_full; ++ki) { + const int64_t k = ki * kVecSize; + a0 += x_cache_lo[ki] * fVec::loadu(wa + k); + b0 += x_cache_lo[ki] * fVec::loadu(wb + k); + a1 += x_cache_hi[ki] * fVec::loadu(wa + k + kFVecSize); + b1 += x_cache_hi[ki] * fVec::loadu(wb + k + kFVecSize); + } + if (rem > 0) { + a0 += x_tail_lo * fVec::loadu(wa + k_tail, rem0); + b0 += x_tail_lo * fVec::loadu(wb + k_tail, rem0); + if (rem1 > 0) { + a1 += x_tail_hi * fVec::loadu(wa + k_tail + kFVecSize, rem1); + b1 += x_tail_hi * fVec::loadu(wb + k_tail + kFVecSize, rem1); + } + } + partial_row[1 + h] = vec_reduce_sum(a0 + a1); + partial_row[1 + h + 1] = vec_reduce_sum(b0 + b1); + } + // odd leftover + if (h < hc) { + const float* w = w_ptr[h]; + fVec hd0(0.f), hd1(0.f); + for (int64_t ki = 0; ki < n_full; ++ki) { + const int64_t k = ki * kVecSize; + hd0 += x_cache_lo[ki] * fVec::loadu(w + k); + hd1 += x_cache_hi[ki] * fVec::loadu(w + k + kFVecSize); + } + if (rem > 0) { + hd0 += x_tail_lo * fVec::loadu(w + k_tail, rem0); + if (rem1 > 0) hd1 += x_tail_hi * fVec::loadu(w + k_tail + kFVecSize, rem1); + } + partial_row[1 + h] = vec_reduce_sum(hd0 + hd1); + } + } + } // tl + + // Prefetch next t-block head. + if (tb + 1 < tb1) { + const scalar_t* x_next_tb = x + ((tb + 1) * T_BLOCK) * hc_d + k0; + __builtin_prefetch(x_next_tb, 0, 0); + if (klen > kVecSize * 2) __builtin_prefetch(x_next_tb + kVecSize * 2, 0, 0); + } + } // tb + } // kb + }); // Phase 1 + + // Phase 2+3. + at::parallel_for(0, T, 1, [&](int64_t begin, int64_t end) { + const int64_t tid = at::get_thread_num(); + float* const pre_gate = scratch_base + tid * per_thread_head; + fVec v0, v1, pre_fvec; + + for (int64_t t = begin; t < end; ++t) { + const int64_t tb = t / T_BLOCK; + const int64_t tl = t % T_BLOCK; + + // Phase 2: reduce partials + sigmoid gate. + double sq_total = 0.0; + std::fill(pre_gate, pre_gate + hc, 0.f); + for (int64_t kb = 0; kb < n_k_blocks; ++kb) { + const float* partial_row = partials + (tb * n_k_blocks + kb) * T_BLOCK * partial_stride + tl * partial_stride; + sq_total += static_cast(partial_row[0]); +#pragma GCC unroll 4 + for (int h = 0; h < hc; ++h) + pre_gate[h] += partial_row[1 + h]; + } + const float inv_rms = + static_cast(1.0 / std::sqrt(sq_total / static_cast(hc_d) + static_cast(norm_eps))); +#pragma GCC unroll 4 + for (int h = 0; h < hc; ++h) { + const float gate = pre_gate[h] * inv_rms * hc_scale_val + hc_base[h]; + pre_gate[h] = 1.f / (1.f + std::exp(-gate)) + hc_eps; + } + + // Phase 3: d-tiled combine. + const scalar_t* x_t = x + t * hc_d; + scalar_t* y_t = y + t * d; + for (int64_t j0 = 0; j0 < d; j0 += D_BLOCK) { + const int64_t jlen = std::min(D_BLOCK, d - j0); + + int64_t kk = 0; + for (; kk <= jlen - kVecSize; kk += kVecSize) { + pre_fvec = fVec(pre_gate[0]); + std::tie(v0, v1) = at::vec::convert_to_float(bVec::loadu(x_t + j0 + kk)); + fVec acc0 = pre_fvec * v0; + fVec acc1 = pre_fvec * v1; +#pragma GCC unroll 4 + for (int h = 1; h < hc; ++h) { + pre_fvec = fVec(pre_gate[h]); + const scalar_t* x_th = x_t + (int64_t)h * d + j0 + kk; + std::tie(v0, v1) = at::vec::convert_to_float(bVec::loadu(x_th)); + acc0 += pre_fvec * v0; + acc1 += pre_fvec * v1; + } + at::vec::convert_from_float(acc0, acc1).store(y_t + j0 + kk); + } + if (kk < jlen) { + const int64_t rem = jlen - kk, rem1 = rem - std::min(rem, kFVecSize); + pre_fvec = fVec(pre_gate[0]); + std::tie(v0, v1) = at::vec::convert_to_float(bVec::loadu(x_t + j0 + kk, rem)); + fVec acc0 = pre_fvec * v0; + fVec acc1 = pre_fvec * v1; +#pragma GCC unroll 4 + for (int h = 1; h < hc; ++h) { + pre_fvec = fVec(pre_gate[h]); + const scalar_t* x_th = x_t + (int64_t)h * d + j0 + kk; + std::tie(v0, v1) = at::vec::convert_to_float(bVec::loadu(x_th, rem)); + acc0 += pre_fvec * v0; + acc1 += pre_fvec * v1; + } + at::vec::convert_from_float(acc0, rem1 > 0 ? acc1 : fVec(0.f)).store(y_t + j0 + kk, rem); + } + } // d-tile + } // t + }); // Phase 2+3 +} + +} // anonymous namespace + +// Fused hc_pre entry. +std::tuple hc_pre_fused_cpu( + at::Tensor& x, + at::Tensor& hc_fn, + at::Tensor& hc_scale, + at::Tensor& hc_base, + int64_t hc_mult, + int64_t sinkhorn_iters, + double rms_eps, + double hc_eps) { + TORCH_CHECK(hc_mult >= 2, "hc_pre_fused_cpu: hc_mult must be >= 2"); + TORCH_CHECK( + x.dim() == 3 && (x.scalar_type() == at::kBFloat16 || x.scalar_type() == at::kHalf), + "hc_pre_fused_cpu: x must be bf16/fp16 [T, hc, d]"); + TORCH_CHECK(hc_fn.scalar_type() == at::kFloat, "hc_pre_fused_cpu: hc_fn must be float32"); + TORCH_CHECK( + hc_scale.scalar_type() == at::kFloat && hc_scale.numel() == 3, "hc_pre_fused_cpu: hc_scale must be float32 [3]"); + TORCH_CHECK(hc_base.scalar_type() == at::kFloat, "hc_pre_fused_cpu: hc_base must be float32"); + TORCH_CHECK( + x.device().is_cpu() && hc_fn.device().is_cpu() && hc_scale.device().is_cpu() && hc_base.device().is_cpu(), + "hc_pre_fused_cpu: all inputs must be CPU tensors"); + + const int64_t T = x.size(0); + const int64_t d = x.size(2); + const int64_t hc_d = hc_mult * d; + const int64_t mix_hc = (2 + hc_mult) * hc_mult; + + TORCH_CHECK(T > 0 && d > 0, "hc_pre_fused_cpu: T and d must be positive"); + TORCH_CHECK(x.size(1) == hc_mult, "hc_pre_fused_cpu: x shape must be [T, hc_mult, d]"); + TORCH_CHECK( + hc_fn.dim() == 2 && hc_fn.size(0) == mix_hc && hc_fn.size(1) == hc_d, + "hc_pre_fused_cpu: hc_fn must be [mix_hc, hc_mult*d]"); + TORCH_CHECK(hc_base.numel() == mix_hc, "hc_pre_fused_cpu: hc_base.numel() must equal mix_hc"); + + // Static per-layer weights (hc_fn/hc_scale/hc_base) are allocated fp32 and + // never reassigned after weight loading; x/residual come straight from a + // preceding op. All are expected contiguous already -- fail loudly instead + // of silently absorbing a copy if that contract is ever violated. + TORCH_CHECK( + x.is_contiguous() && hc_fn.is_contiguous() && hc_scale.is_contiguous() && hc_base.is_contiguous(), + "hc_pre_fused_cpu: x/hc_fn/hc_scale/hc_base must be contiguous"); + + auto f32_opts = at::TensorOptions().dtype(at::kFloat).device(x.device()); + auto post = at::empty({T, hc_mult}, f32_opts); + auto comb = at::empty({T, hc_mult, hc_mult}, f32_opts); + auto y = at::empty({T, d}, x.options()); + + AT_DISPATCH_REDUCED_FLOATING_TYPES(x.scalar_type(), "hc_pre_fused_cpu", [&] { + hc_pre_fuse_impl( + y.data_ptr(), + post.data_ptr(), + comb.data_ptr(), + x.data_ptr(), + hc_fn.data_ptr(), + hc_scale.data_ptr(), + hc_base.data_ptr(), + T, + d, + static_cast(hc_mult), + static_cast(sinkhorn_iters), + static_cast(hc_eps), + static_cast(rms_eps)); + }); + return {y, post, comb}; +} + +// Fused hc_post entry. +at::Tensor hc_post_fused_cpu(at::Tensor& x, at::Tensor& residual, at::Tensor& post, at::Tensor& comb) { + TORCH_CHECK( + x.dim() == 2 && (x.scalar_type() == at::kBFloat16 || x.scalar_type() == at::kHalf), + "hc_post_fused_cpu: x must be bf16/fp16 [T, d]"); + TORCH_CHECK( + residual.dim() == 3 && residual.scalar_type() == x.scalar_type(), + "hc_post_fused_cpu: residual must be same dtype as x [T, hc, d]"); + TORCH_CHECK(post.scalar_type() == at::kFloat, "hc_post_fused_cpu: post must be float32"); + TORCH_CHECK(comb.scalar_type() == at::kFloat, "hc_post_fused_cpu: comb must be float32"); + TORCH_CHECK( + x.device().is_cpu() && residual.device().is_cpu() && post.device().is_cpu() && comb.device().is_cpu(), + "hc_post_fused_cpu: all inputs must be CPU tensors"); + + const int64_t T = x.size(0); + const int64_t d = x.size(1); + const int64_t hc = residual.size(1); + TORCH_CHECK(hc >= 2, "hc_post_fused_cpu: hc_mult must be >= 2"); + TORCH_CHECK(T > 0 && d > 0, "hc_post_fused_cpu: T and d must be positive"); + TORCH_CHECK(residual.size(0) == T && residual.size(2) == d, "hc_post_fused_cpu: residual shape must be [T, hc, d]"); + TORCH_CHECK( + post.dim() == 2 && post.size(0) == T && post.size(1) == hc, "hc_post_fused_cpu: post shape must be [T, hc]"); + TORCH_CHECK( + comb.dim() == 3 && comb.size(0) == T && comb.size(1) == hc && comb.size(2) == hc, + "hc_post_fused_cpu: comb shape must be [T, hc, hc]"); + + TORCH_CHECK( + x.is_contiguous() && residual.is_contiguous() && post.is_contiguous() && comb.is_contiguous(), + "hc_post_fused_cpu: x/residual/post/comb must be contiguous"); + + auto out = at::empty({T, hc, d}, x.options()); // same dtype as x + + AT_DISPATCH_REDUCED_FLOATING_TYPES(x.scalar_type(), "hc_post_fused_cpu", [&] { + hc_post_fuse_impl( + out.data_ptr(), + x.data_ptr(), + residual.data_ptr(), + post.data_ptr(), + comb.data_ptr(), + T, + d, + static_cast(hc)); + }); + + return out; +} + +// Fused hc_head entry. +at::Tensor hc_head_fused_cpu( + at::Tensor& x, at::Tensor& hc_fn, at::Tensor& hc_scale, at::Tensor& hc_base, double hc_eps, double norm_eps) { + TORCH_CHECK( + x.dim() == 3 && (x.scalar_type() == at::kBFloat16 || x.scalar_type() == at::kHalf), + "hc_head_fused_cpu: x must be bf16/fp16 [T, hc, d]"); + const int64_t hc_mult = x.size(1); + TORCH_CHECK(hc_mult >= 2, "hc_head_fused_cpu: hc_mult must be >= 2"); + TORCH_CHECK(hc_fn.scalar_type() == at::kFloat, "hc_head_fused_cpu: hc_fn must be float32"); + TORCH_CHECK( + hc_scale.scalar_type() == at::kFloat && hc_scale.numel() == 1, + "hc_head_fused_cpu: hc_scale must be a scalar float32"); + TORCH_CHECK(hc_base.scalar_type() == at::kFloat, "hc_head_fused_cpu: hc_base must be float32"); + TORCH_CHECK(hc_base.numel() == hc_mult, "hc_head_fused_cpu: hc_base.numel() must equal hc_mult"); + TORCH_CHECK( + x.device().is_cpu() && hc_fn.device().is_cpu() && hc_scale.device().is_cpu() && hc_base.device().is_cpu(), + "hc_head_fused_cpu: all inputs must be CPU tensors"); + + const int64_t T = x.size(0); + const int64_t d = x.size(2); + TORCH_CHECK(T > 0 && d > 0, "hc_head_fused_cpu: T and d must be positive"); + TORCH_CHECK( + hc_fn.dim() == 2 && hc_fn.size(0) == hc_mult && hc_fn.size(1) == hc_mult * d, + "hc_head_fused_cpu: hc_fn must be [hc_mult, hc_mult*d]"); + + TORCH_CHECK( + x.is_contiguous() && hc_fn.is_contiguous() && hc_scale.is_contiguous() && hc_base.is_contiguous(), + "hc_head_fused_cpu: x/hc_fn/hc_scale/hc_base must be contiguous"); + const float hc_scale_val = hc_scale.data_ptr()[0]; + + auto y = at::empty({T, d}, x.options()); + + AT_DISPATCH_REDUCED_FLOATING_TYPES(x.scalar_type(), "hc_head_fused_cpu", [&] { + hc_head_fuse_impl( + y.data_ptr(), + x.data_ptr(), + hc_fn.data_ptr(), + hc_scale_val, + hc_base.data_ptr(), + T, + d, + static_cast(hc_mult), + static_cast(hc_eps), + static_cast(norm_eps)); + }); + + return y; +} diff --git a/csrc/cpu/sgl-kernels/mla_cache.cpp b/csrc/cpu/sgl-kernels/mla_cache.cpp index 26f52b9159..74caf6dd39 100644 --- a/csrc/cpu/sgl-kernels/mla_cache.cpp +++ b/csrc/cpu/sgl-kernels/mla_cache.cpp @@ -1,14 +1,14 @@ // vLLM-native CPU cache-write op for MLA's single-latent-buffer KV cache. +// `concat_and_cache_mla` (the generic op used by every GPU backend) is +// CUDA-only; this is the CPU counterpart, adapted in spirit from SGLang's +// `store_cache_cpu` (csrc/cpu/kvcache.cpp). // -// `concat_and_cache_mla` (the generic MLA cache-write op used by every GPU -// backend) is registered CUDA-only. This is the CPU counterpart, adapted in -// spirit from SGLang's `store_cache_cpu` (csrc/cpu/kvcache.cpp) but -// generalized to write two source tensors (`kv_c_normed`, `k_pe`) into two -// different column-offset ranges of the SAME destination row -- SGLang's -// version assumes k/v land in two independent, equal-row-width cache -// tensors, which doesn't hold here since MLA's cache is one 576-wide buffer -// and the two column ranges (512-wide, 64-wide) don't match the buffer's -// true per-token stride, so the write can't reuse `store_cache_cpu` as-is. +// Difference from upstream: writes two source tensors (`kv_c_normed`, +// `k_pe`) into two column-offset ranges of the SAME destination row. +// Upstream assumes k/v land in two independent, equal-row-width cache +// tensors -- MLA's cache is one 576-wide buffer whose two column ranges +// (512-wide, 64-wide) don't match that per-token stride, so +// `store_cache_cpu` can't be reused as-is. #include "common.h" #include "vec.h" diff --git a/csrc/cpu/sgl-kernels/moe.cpp b/csrc/cpu/sgl-kernels/moe.cpp index 5a810d8f38..3afdc1bd67 100644 --- a/csrc/cpu/sgl-kernels/moe.cpp +++ b/csrc/cpu/sgl-kernels/moe.cpp @@ -871,13 +871,13 @@ static inline void check_moe_scales( // topk_ids: [M, topk] (int32_t) // -at::Tensor fused_experts_cpu( +void fused_experts_cpu( + at::Tensor& out, at::Tensor& hidden_states, at::Tensor& w1, at::Tensor& w2, at::Tensor& topk_weights, at::Tensor& topk_ids, - bool inplace, int64_t moe_comp_method, const std::optional& w1_scale, const std::optional& w2_scale, @@ -943,7 +943,7 @@ at::Tensor fused_experts_cpu( // check scales check_moe_scales(moe_comp_method, w1_scale, w2_scale, block_size); - at::Tensor out_hidden_states = inplace ? hidden_states : at::empty_like(hidden_states); + at::Tensor& out_hidden_states = out; // NB: worst case is each expert holds a block with remainder of 1 // 1. sorted_ids : [M * topk + E * (BLOCK_M - 1)] @@ -1214,7 +1214,6 @@ at::Tensor fused_experts_cpu( with_bias); } }); - return out_hidden_states; } // shared expert kernel diff --git a/csrc/cpu/sgl-kernels/paged_mqa_logits.cpp b/csrc/cpu/sgl-kernels/paged_mqa_logits.cpp new file mode 100644 index 0000000000..3462ab45f3 --- /dev/null +++ b/csrc/cpu/sgl-kernels/paged_mqa_logits.cpp @@ -0,0 +1,217 @@ +// Adapted from +// https://github.com/sgl-project/sglang/tree/main/sgl-kernel/csrc/cpu: +// `fp8_paged_mqa_logits_cpu`/`fp8_paged_mqa_logits_cpu_impl` (DECODE path). +// +// Reads the paged indexer K-cache directly via `page_table` -- no separate +// eager gather step -- and processes the whole decode batch (one query token +// per request; the CPU indexer decode path doesn't support native MTP/spec +// decode yet, see `sparse_attn_indexer_cpu`) in a single `at::parallel_for` +// flattening `(batch, token)`, exactly like upstream. AVX512 `dot_fp8_128` +// dequant-dot-product and the split K-region-then-scale-region page byte +// layout are verbatim (that part of the upstream contract already matches +// vLLM's cache, written by `compress_norm_rope_store_indexer_cpu` in +// compressor.cpp). +// +// Differences from upstream: +// - `block_size` is a runtime argument instead of the hardcoded +// `kBlockSize = 64`. vLLM's indexer K-cache shares the model's single +// global `--block-size` KV-cache-group setting (confirmed empirically: +// the DeepSeek-V4-Flash TP1/TP4 baseline configs run with +// `block_size=256`, not 64 -- see `.agent-workspace/perf_baseline_tp*_ +// server.log`), so the upstream `kvcache_fp8.size(1) == 64` shape +// contract does not hold in general. +// - `q_fp8`/`weight` drop the redundant `S_q=1` dimension (`[batch, heads, +// 128]`/`[batch, heads]` instead of `[batch, 1, heads, 128]`) since S_q is +// always 1 here; `weight` is fp32-only (already folds q_scale/softmax_ +// scale/head_scale, matching this call site's actual dtype -- unlike +// upstream's bf16/fp16/fp32 dispatch, which vLLM's CPU indexer never +// produces). +// - `index_head_dim` (128) stays a compile-time constant, matching +// upstream: it's `config.index_head_dim`, a genuine DeepSeek-V4 model +// invariant, not a vLLM KV-cache config knob like `block_size`. + +#include "common.h" +#include "vec.h" + +#if !defined(CPU_CAPABILITY_AVX512) +static_assert(false, + "paged_mqa_logits.cpp is only ever compiled into the " + "AVX512+AMX-flagged _C target (see cmake/cpu_extension.cmake's " + "VLLM_EXT_SRC_SGL) and relies on that unconditionally -- " + "CPU_CAPABILITY_AVX512 must be defined here."); +#endif + +namespace { + +constexpr int64_t kIndexHeadDim = 128; +constexpr int64_t kIndexHeadDimWithScaleBytes = 132; + +inline float dot_fp8_128(const uint8_t* k, const uint8_t* q) { + __m512 acc = _mm512_setzero_ps(); + for (int64_t d = 0; d < kIndexHeadDim; d += 32) { + const __m256i k8 = + _mm256_loadu_si256(reinterpret_cast(k + d)); + const __m256i q8 = + _mm256_loadu_si256(reinterpret_cast(q + d)); + acc = _mm512_dpbf16_ps(acc, CVT_FP8_TO_BF16(k8), CVT_FP8_TO_BF16(q8)); + } + return _mm512_reduce_add_ps(acc); +} + +template +void fp8_paged_mqa_logits_cpu_impl(const at::Tensor& q_fp8, + const at::Tensor& kvcache_fp8, + const at::Tensor& weight, + const at::Tensor& seq_lens, + const at::Tensor& page_table, + at::Tensor& logits, int64_t block_size, + int64_t max_seq_len) { + const int64_t batch_size = q_fp8.size(0); + const int64_t num_heads = q_fp8.size(1); + const int64_t num_pages = kvcache_fp8.size(0); + // Row stride is threaded through explicitly rather than assumed equal to + // the page byte width: the K-cache may be a page of a larger pooled + // allocation shared across layers (matches store_cache.cpp/flash_mla.cpp's + // precedent). + const int64_t page_stride = kvcache_fp8.stride(0); + const int64_t pages_per_batch = page_table.size(1); + const int64_t scale_offset_bytes = block_size * kIndexHeadDim; + + const auto* q_ptr = reinterpret_cast(q_fp8.const_data_ptr()); + const auto* cache_ptr = + reinterpret_cast(kvcache_fp8.const_data_ptr()); + const auto* weight_ptr = weight.const_data_ptr(); + const auto* seq_ptr = seq_lens.const_data_ptr(); + const auto* page_ptr = page_table.const_data_ptr(); + auto* out_ptr = logits.data_ptr(); + + at::parallel_for( + 0, batch_size * max_seq_len, GRAIN_SIZE / kIndexHeadDim, + [&](int64_t begin, int64_t end) { + int64_t b{0}, token{0}; + data_index_init(begin, b, batch_size, token, max_seq_len); + for (int64_t i = begin; i < end; ++i) { + const int64_t seq_len = static_cast(seq_ptr[b]); + + if (token >= seq_len) { + data_index_step(b, batch_size, token, max_seq_len); + continue; + } + + const int64_t q_batch_offset = b * num_heads * kIndexHeadDim; + const int64_t weight_batch_offset = b * num_heads; + const int64_t page_batch_offset = b * pages_per_batch; + float* out_row = out_ptr + b * max_seq_len; + + const int64_t logical_page = token / block_size; + const int64_t token_in_page = token % block_size; + + const int64_t physical_page = + static_cast(page_ptr[page_batch_offset + logical_page]); + + const uint8_t* block = cache_ptr + physical_page * page_stride; + const uint8_t* k_token = block + token_in_page * kIndexHeadDim; + const float* scale_ptr = + reinterpret_cast(block + scale_offset_bytes); + const float k_scale = scale_ptr[token_in_page]; + + float score_sum = 0.0f; + for (int64_t h = 0; h < num_heads; ++h) { + const uint8_t* q_head = q_ptr + q_batch_offset + h * kIndexHeadDim; + float dot = dot_fp8_128(k_token, q_head); + dot = std::max(dot, 0.0f); + score_sum += dot * weight_ptr[weight_batch_offset + h]; + } + + out_row[token] = score_sum * k_scale; + data_index_step(b, batch_size, token, max_seq_len); + } + }); +} + +template +void dispatch_page_type(const at::Tensor& q_fp8, const at::Tensor& kvcache_fp8, + const at::Tensor& weight, const at::Tensor& seq_lens, + const at::Tensor& page_table, at::Tensor& logits, + int64_t block_size, int64_t max_seq_len) { + if (page_table.scalar_type() == at::kInt) { + fp8_paged_mqa_logits_cpu_impl(q_fp8, kvcache_fp8, weight, + seq_lens, page_table, logits, + block_size, max_seq_len); + } else if (page_table.scalar_type() == at::kLong) { + fp8_paged_mqa_logits_cpu_impl(q_fp8, kvcache_fp8, weight, + seq_lens, page_table, logits, + block_size, max_seq_len); + } else { + TORCH_CHECK(false, "page_table must be int32 or int64"); + } +} + +} // namespace + +at::Tensor fp8_paged_mqa_logits_cpu(at::Tensor& q_fp8, at::Tensor& kvcache_fp8, + at::Tensor& weight, at::Tensor& seq_lens, + at::Tensor& page_table, int64_t block_size, + int64_t max_seq_len) { + CHECK_INPUT(q_fp8); + CHECK_CPU(kvcache_fp8); + CHECK_INPUT(weight); + CHECK_INPUT(seq_lens); + CHECK_INPUT(page_table); + TORCH_CHECK( + q_fp8.dim() == 3 && q_fp8.scalar_type() == at::kFloat8_e4m3fn, + "fp8_paged_mqa_logits_cpu: q_fp8 must be a 3D float8_e4m3fn tensor " + "[batch, heads, 128]"); + TORCH_CHECK(q_fp8.size(2) == kIndexHeadDim, + "fp8_paged_mqa_logits_cpu: q_fp8 head_dim must be 128"); + TORCH_CHECK(kvcache_fp8.dim() == 2 && kvcache_fp8.scalar_type() == at::kByte, + "fp8_paged_mqa_logits_cpu: kvcache_fp8 must be a 2D uint8 tensor " + "[num_pages, buf_numel_per_page]"); + TORCH_CHECK(kvcache_fp8.stride(1) == 1 && + kvcache_fp8.stride(0) >= kvcache_fp8.size(1), + "fp8_paged_mqa_logits_cpu: kvcache_fp8 must be contiguous " + "within a page (may be a page of a larger pooled allocation " + "between pages)"); + TORCH_CHECK(block_size > 0, + "fp8_paged_mqa_logits_cpu: block_size must be positive"); + TORCH_CHECK( + kvcache_fp8.size(1) >= block_size * kIndexHeadDimWithScaleBytes, + "fp8_paged_mqa_logits_cpu: kvcache_fp8 page byte width too small for " + "block_size"); + + const int64_t batch_size = q_fp8.size(0); + const int64_t num_heads = q_fp8.size(1); + TORCH_CHECK( + weight.dim() == 2 && weight.scalar_type() == at::kFloat && + weight.size(0) == batch_size && weight.size(1) == num_heads, + "fp8_paged_mqa_logits_cpu: weight must be a 2D fp32 tensor [batch, " + "heads]"); + TORCH_CHECK(seq_lens.dim() == 1 && seq_lens.size(0) == batch_size, + "fp8_paged_mqa_logits_cpu: seq_lens must be a 1D tensor [batch]"); + TORCH_CHECK(page_table.dim() == 2 && page_table.size(0) == batch_size, + "fp8_paged_mqa_logits_cpu: page_table must be a 2D tensor " + "[batch, pages]"); + TORCH_CHECK(max_seq_len >= 0, + "fp8_paged_mqa_logits_cpu: max_seq_len must be non-negative"); + + // Positions with token >= seq_len are never written by the kernel below, + // and topk_transform_512_cpu (the only consumer) never reads them either + // -- left uninitialized, matching upstream. + auto logits = + at::empty({batch_size, max_seq_len}, q_fp8.options().dtype(at::kFloat)); + if (batch_size == 0 || max_seq_len == 0) { + return logits; + } + + if (seq_lens.scalar_type() == at::kInt) { + dispatch_page_type(q_fp8, kvcache_fp8, weight, seq_lens, + page_table, logits, block_size, max_seq_len); + } else if (seq_lens.scalar_type() == at::kLong) { + dispatch_page_type(q_fp8, kvcache_fp8, weight, seq_lens, + page_table, logits, block_size, max_seq_len); + } else { + TORCH_CHECK(false, "seq_lens must be int32 or int64"); + } + + return logits; +} diff --git a/csrc/cpu/sgl-kernels/store_cache.cpp b/csrc/cpu/sgl-kernels/store_cache.cpp new file mode 100644 index 0000000000..9da10ac111 --- /dev/null +++ b/csrc/cpu/sgl-kernels/store_cache.cpp @@ -0,0 +1,481 @@ +// Adapted from +// https://github.com/sgl-project/sglang/blob/main/python/sglang/kernels/aot/csrc/cpu/store_cache.cpp +// +// DeepSeek-V4 fused Q RMSNorm+RoPE / KV RoPE+UE8M0-quant+paged-cache-insert +// kernel for the fp8_ds_mla cache layout. Not a port of upstream's +// `set_k_and_s_cpu`/`quant_to_nope_fp8_rope_bf16_pack_cpu` -- their +// pool_kv/pool_score signature doesn't apply to vLLM's layout. This is the +// CPU counterpart to vLLM's own CUDA kernel +// (`fused_deepseek_v4_qnorm_rope_kv_insert_kernel.cu`); see +// `DeepseekV4Attention._fused_qnorm_rope_kv_insert_cpu` (attention.py) for +// the eager reference this replaces. +// +// clang-format off + +#include +#include +#include + +#include "common.h" +#include "vec.h" + +#if !defined(CPU_CAPABILITY_AVX512) +static_assert(false, + "store_cache.cpp is only ever compiled into the AVX512+AMX-flagged " + "_C target (see cmake/cpu_extension.cmake's VLLM_EXT_SRC_SGL) and " + "relies on that unconditionally -- CPU_CAPABILITY_AVX512 must be " + "defined here."); +#endif + +namespace { + +constexpr int64_t kRopeDim = 64; +constexpr int64_t kQuantBlock = 64; +constexpr float kFp8Max = 448.0f; +constexpr float kFp8Min = -448.0f; +constexpr float kQuantEps = 1e-4f; + +inline void apply_gptj_rope_inplace(float* __restrict__ x, const float* __restrict__ cos, const float* __restrict__ sin, int64_t half_dim) { + for (int64_t i = 0; i < half_dim; ++i) { + const float e = x[2 * i]; + const float o = x[2 * i + 1]; + x[2 * i] = e * cos[i] - o * sin[i]; + x[2 * i + 1] = e * sin[i] + o * cos[i]; + } +} + +// AVX512 tile-quantization, adapted verbatim from SGLang's +// `store_cache.cpp` (`quant_to_nope_fp8_rope_bf16_pack_cpu`'s +// `cvt_fp32x16_to_fp8x16`/`quantize_tile_avx512`): per-64-element tile, +// computes the UE8M0 power-of-two scale and quantizes to fp8_e4m3fn with +// RNE rounding via bit manipulation instead of a per-element scalar +// `at::Float8_e4m3fn` cast. Subnormal-magnitude fp8 lanes are flushed to +// zero (matches SGLang's own convention, `cvt_fp32x16_to_fp8x16`) -- +// vLLM's own AVX512 dequant of this cache (`flash_mla.cpp`'s +// `CVT_FP8_TO_BF16_EXT`) doesn't reconstruct exact subnormal magnitudes +// either way, so there is nothing to preserve by rounding them exactly on +// the write side. +#if defined(CPU_CAPABILITY_AVX512) + +// Vectorized float32x16 -> fp8_e4m3fn x16 using AVX512 bit manipulation. +// Implements RNE (round-to-nearest-even). Flushes subnormal fp8 to zero. +// fp8_e4m3fn: sign(1), exp(4, bias=7), mant(3), no inf/nan special values. +// Normal: exp_field in [1..15], value = 2^(exp-7) * (1 + mant/8) +// Input assumed clamped to [-448, 448]. +inline __m128i cvt_fp32x16_to_fp8x16(__m512 input) { + __m512i bits = _mm512_castps_si512(input); + + // Extract sign -> bit 7 of fp8 byte + __m512i signs = _mm512_srli_epi32(_mm512_and_si512(bits, _mm512_set1_epi32(0x80000000)), 24); + + // Absolute value bits + __m512i abs_bits = _mm512_and_si512(bits, _mm512_set1_epi32(0x7FFFFFFF)); + + // Float32 biased exponent + __m512i f32_exp = _mm512_srli_epi32(abs_bits, 23); + + // Float32 mantissa (23 bits) + __m512i f32_mant = _mm512_and_si512(abs_bits, _mm512_set1_epi32(0x7FFFFF)); + + // fp8_exp = f32_exp - 120 (rebias from 127 to 7) + __m512i fp8_exp = _mm512_sub_epi32(f32_exp, _mm512_set1_epi32(120)); + + // Top 3 mantissa bits + __m512i mant3 = _mm512_srli_epi32(f32_mant, 20); + + // RNE rounding: round_bit = bit 19, sticky = bits[18:0] + __m512i round_bit = _mm512_and_si512(_mm512_srli_epi32(f32_mant, 19), _mm512_set1_epi32(1)); + __m512i sticky_bits = _mm512_and_si512(f32_mant, _mm512_set1_epi32(0x7FFFF)); + __mmask16 has_sticky = _mm512_cmpneq_epi32_mask(sticky_bits, _mm512_setzero_si512()); + + // Round up if round_bit AND (sticky OR lsb_of_mant3) + __m512i lsb = _mm512_and_si512(mant3, _mm512_set1_epi32(1)); + __m512i sticky_or_lsb = _mm512_or_si512(_mm512_maskz_mov_epi32(has_sticky, _mm512_set1_epi32(1)), lsb); + __m512i do_round = _mm512_and_si512(round_bit, sticky_or_lsb); + mant3 = _mm512_add_epi32(mant3, do_round); + + // Mantissa overflow (mant3 == 8) -> increment exponent, zero mantissa + __mmask16 mant_ovf = _mm512_cmpeq_epi32_mask(mant3, _mm512_set1_epi32(8)); + mant3 = _mm512_mask_mov_epi32(mant3, mant_ovf, _mm512_setzero_si512()); + fp8_exp = _mm512_mask_add_epi32(fp8_exp, mant_ovf, fp8_exp, _mm512_set1_epi32(1)); + + // Clamp exponent: max 15 + fp8_exp = _mm512_min_epi32(fp8_exp, _mm512_set1_epi32(15)); + + // Result: (fp8_exp << 3) | mant3 + __m512i result = _mm512_or_si512(_mm512_slli_epi32(fp8_exp, 3), mant3); + + // Flush to zero: if f32_exp < 121 (would be subnormal in fp8), output 0 + __mmask16 is_subnormal_or_zero = _mm512_cmplt_epi32_mask(f32_exp, _mm512_set1_epi32(121)); + result = _mm512_mask_mov_epi32(result, is_subnormal_or_zero, _mm512_setzero_si512()); + + // Add sign + result = _mm512_or_si512(result, signs); + + // Pack 16 x i32 -> 16 x i8 + return _mm512_cvtepi32_epi8(result); +} + +// Process one tile (64 bf16 values): find amax, compute scale, quantize to fp8. +inline uint8_t quantize_tile_avx512(const at::BFloat16* __restrict__ src, uint8_t* __restrict__ dst) { + // Load 64 bf16 -> 4 x 16 floats + __m256i bf16_0 = _mm256_loadu_si256(reinterpret_cast(src)); + __m256i bf16_1 = _mm256_loadu_si256(reinterpret_cast(src + 16)); + __m256i bf16_2 = _mm256_loadu_si256(reinterpret_cast(src + 32)); + __m256i bf16_3 = _mm256_loadu_si256(reinterpret_cast(src + 48)); + + __m512 f0 = _mm512_castsi512_ps(_mm512_slli_epi32(_mm512_cvtepu16_epi32(bf16_0), 16)); + __m512 f1 = _mm512_castsi512_ps(_mm512_slli_epi32(_mm512_cvtepu16_epi32(bf16_1), 16)); + __m512 f2 = _mm512_castsi512_ps(_mm512_slli_epi32(_mm512_cvtepu16_epi32(bf16_2), 16)); + __m512 f3 = _mm512_castsi512_ps(_mm512_slli_epi32(_mm512_cvtepu16_epi32(bf16_3), 16)); + + // Absolute values + const __m512i abs_mask_i = _mm512_set1_epi32(0x7FFFFFFF); + __m512 abs0 = _mm512_castsi512_ps(_mm512_and_si512(_mm512_castps_si512(f0), abs_mask_i)); + __m512 abs1 = _mm512_castsi512_ps(_mm512_and_si512(_mm512_castps_si512(f1), abs_mask_i)); + __m512 abs2 = _mm512_castsi512_ps(_mm512_and_si512(_mm512_castps_si512(f2), abs_mask_i)); + __m512 abs3 = _mm512_castsi512_ps(_mm512_and_si512(_mm512_castps_si512(f3), abs_mask_i)); + + // Horizontal max + __m512 max01 = _mm512_max_ps(abs0, abs1); + __m512 max23 = _mm512_max_ps(abs2, abs3); + __m512 max0123 = _mm512_max_ps(max01, max23); + float amax = _mm512_reduce_max_ps(max0123); + + // Scale computation + float scale = std::max(amax / kFp8Max, kQuantEps); + float ceil_log2 = std::ceil(std::log2(scale)); + float scale_pow2 = std::exp2(ceil_log2); + float scale_inv = 1.0f / scale_pow2; + + int exponent = static_cast(ceil_log2); + uint8_t scale_uint8 = static_cast(exponent + 127); + + // Scale all values + __m512 vinv = _mm512_set1_ps(scale_inv); + __m512 s0 = _mm512_mul_ps(f0, vinv); + __m512 s1 = _mm512_mul_ps(f1, vinv); + __m512 s2 = _mm512_mul_ps(f2, vinv); + __m512 s3 = _mm512_mul_ps(f3, vinv); + + // Clamp + __m512 vmax = _mm512_set1_ps(kFp8Max); + __m512 vmin = _mm512_set1_ps(kFp8Min); + s0 = _mm512_max_ps(_mm512_min_ps(s0, vmax), vmin); + s1 = _mm512_max_ps(_mm512_min_ps(s1, vmax), vmin); + s2 = _mm512_max_ps(_mm512_min_ps(s2, vmax), vmin); + s3 = _mm512_max_ps(_mm512_min_ps(s3, vmax), vmin); + + // Vectorized float -> fp8 conversion (16 values at a time) + __m128i fp8_0 = cvt_fp32x16_to_fp8x16(s0); + __m128i fp8_1 = cvt_fp32x16_to_fp8x16(s1); + __m128i fp8_2 = cvt_fp32x16_to_fp8x16(s2); + __m128i fp8_3 = cvt_fp32x16_to_fp8x16(s3); + + _mm_storeu_si128(reinterpret_cast<__m128i*>(dst), fp8_0); + _mm_storeu_si128(reinterpret_cast<__m128i*>(dst + 16), fp8_1); + _mm_storeu_si128(reinterpret_cast<__m128i*>(dst + 32), fp8_2); + _mm_storeu_si128(reinterpret_cast<__m128i*>(dst + 48), fp8_3); + + return scale_uint8; +} + +#endif // CPU_CAPABILITY_AVX512 + +template +static void qnorm_rope_kv_insert_impl( + scalar_t* __restrict__ q_out, // [num_tokens_full, q_head_padded, head_dim] + uint8_t* __restrict__ cache, // [num_blocks, cache_row_stride] bytes + const scalar_t* __restrict__ q, // [num_tokens_full, num_heads_q, head_dim] + const scalar_t* __restrict__ kv, // [>=num_tokens_insert, head_dim] + const int64_t* __restrict__ positions, // [num_tokens_full] + const int64_t* __restrict__ slot_mapping, // [num_tokens_insert] + const float* __restrict__ cos_sin_cache, // [max_pos, rope_dim] (first half cos, second half sin) + int64_t num_tokens_full, + int64_t num_heads_q, + int64_t q_head_padded, + int64_t head_dim, + int64_t nope_dim, + int64_t num_quant_blocks, + int64_t token_data_bytes, + int64_t scale_bytes_per_token, + int64_t num_tokens_insert, + int64_t cache_block_size, + int64_t cache_row_stride, + float eps) { + const int64_t rope_half = kRopeDim / 2; + + // Q: per-head weight-free RMSNorm + GPT-J RoPE on the last rope_dim dims. + at::parallel_for(0, num_tokens_full * num_heads_q, 0, [&](int64_t begin, int64_t end) { + std::vector buf(head_dim); + for (int64_t idx = begin; idx < end; ++idx) { + const int64_t t = idx / num_heads_q; + const int64_t h = idx % num_heads_q; + const scalar_t* q_th = q + (t * num_heads_q + h) * head_dim; + + float sq_sum = 0.f; + for (int64_t i = 0; i < head_dim; ++i) { + const float v = static_cast(q_th[i]); + buf[i] = v; + sq_sum += v * v; + } + const float inv_rms = 1.0f / std::sqrt(sq_sum / static_cast(head_dim) + eps); + for (int64_t i = 0; i < head_dim; ++i) + buf[i] *= inv_rms; + + const int64_t pos = positions[t]; + const float* cos = cos_sin_cache + pos * kRopeDim; + const float* sin = cos + rope_half; + apply_gptj_rope_inplace(buf.data() + nope_dim, cos, sin, rope_half); + + scalar_t* out_th = q_out + (t * q_head_padded + h) * head_dim; + for (int64_t i = 0; i < head_dim; ++i) + out_th[i] = static_cast(buf[i]); + } + }); + + // KV: GPT-J RoPE (no RMSNorm) + bf16 round-trip + UE8M0 FP8 quant + paged insert. + at::parallel_for(0, num_tokens_insert, 0, [&](int64_t begin, int64_t end) { + std::vector buf(head_dim); + std::vector nope_bf16(nope_dim); + for (int64_t t = begin; t < end; ++t) { + const int64_t slot_id = slot_mapping[t]; + if (slot_id < 0) continue; + + const scalar_t* kv_t = kv + t * head_dim; + for (int64_t i = 0; i < head_dim; ++i) + buf[i] = static_cast(kv_t[i]); + + const int64_t pos = positions[t]; + const float* cos = cos_sin_cache + pos * kRopeDim; + const float* sin = cos + rope_half; + apply_gptj_rope_inplace(buf.data() + nope_dim, cos, sin, rope_half); + + // bf16 round-trip, matching the CUDA kernel's double convert. + for (int64_t i = 0; i < head_dim; ++i) + buf[i] = static_cast(at::BFloat16(buf[i])); + for (int64_t i = 0; i < nope_dim; ++i) + nope_bf16[i] = at::BFloat16(buf[i]); + + const int64_t block_idx = slot_id / cache_block_size; + const int64_t pos_in_block = slot_id % cache_block_size; + uint8_t* const block_base = cache + block_idx * cache_row_stride; + uint8_t* const token_fp8_ptr = block_base + pos_in_block * token_data_bytes; + uint8_t* const token_bf16_ptr = token_fp8_ptr + nope_dim; + uint8_t* const token_scale_ptr = + block_base + cache_block_size * token_data_bytes + pos_in_block * scale_bytes_per_token; + + // AVX512 tile-quantization (SGLang's `quantize_tile_avx512`), one + // 64-wide tile per iteration. + for (int64_t qb = 0; qb < num_quant_blocks; ++qb) { + const at::BFloat16* tile_src = nope_bf16.data() + qb * kQuantBlock; + uint8_t* tile_dst = token_fp8_ptr + qb * kQuantBlock; + token_scale_ptr[qb] = quantize_tile_avx512(tile_src, tile_dst); + } + token_scale_ptr[num_quant_blocks] = 0; + + auto* out_rope = reinterpret_cast(token_bf16_ptr); + for (int64_t i = 0; i < kRopeDim; ++i) + out_rope[i] = at::BFloat16(buf[nope_dim + i]); + } + }); +} + +} // anonymous namespace + +at::Tensor fused_qnorm_rope_kv_insert_cpu( + at::Tensor& q, + at::Tensor& kv, + at::Tensor& positions, + at::Tensor& swa_kv_cache_2d, + at::Tensor& slot_mapping, + at::Tensor& cos_sin_cache, + int64_t q_head_padded, + double eps, + int64_t cache_block_size) { + TORCH_CHECK( + q.dim() == 3 && (q.scalar_type() == at::kBFloat16 || q.scalar_type() == at::kHalf), + "fused_qnorm_rope_kv_insert_cpu: q must be bf16/fp16 [num_tokens, num_heads_q, head_dim]"); + TORCH_CHECK(kv.dim() == 2 && kv.scalar_type() == q.scalar_type(), "fused_qnorm_rope_kv_insert_cpu: kv dtype must match q"); + TORCH_CHECK(positions.scalar_type() == at::kLong, "fused_qnorm_rope_kv_insert_cpu: positions must be int64"); + TORCH_CHECK(slot_mapping.scalar_type() == at::kLong, "fused_qnorm_rope_kv_insert_cpu: slot_mapping must be int64"); + TORCH_CHECK( + swa_kv_cache_2d.dim() == 2 && swa_kv_cache_2d.scalar_type() == at::kByte, + "fused_qnorm_rope_kv_insert_cpu: swa_kv_cache_2d must be a 2D uint8 tensor"); + TORCH_CHECK( + q.device().is_cpu() && kv.device().is_cpu() && positions.device().is_cpu() && slot_mapping.device().is_cpu() && + swa_kv_cache_2d.device().is_cpu() && cos_sin_cache.device().is_cpu(), + "fused_qnorm_rope_kv_insert_cpu: all inputs must be CPU tensors"); + + const int64_t num_tokens_full = q.size(0); + const int64_t num_heads_q = q.size(1); + const int64_t head_dim = q.size(2); + const int64_t nope_dim = head_dim - kRopeDim; + TORCH_CHECK(nope_dim > 0 && nope_dim % kQuantBlock == 0, "fused_qnorm_rope_kv_insert_cpu: unexpected head_dim"); + const int64_t num_quant_blocks = nope_dim / kQuantBlock; + const int64_t scale_bytes_per_token = num_quant_blocks + 1; + const int64_t token_data_bytes = nope_dim + kRopeDim * 2; + + TORCH_CHECK(q_head_padded >= num_heads_q, "fused_qnorm_rope_kv_insert_cpu: q_head_padded must be >= num_heads_q"); + TORCH_CHECK(kv.size(1) == head_dim, "fused_qnorm_rope_kv_insert_cpu: kv head_dim must match q"); + TORCH_CHECK(positions.numel() == num_tokens_full, "fused_qnorm_rope_kv_insert_cpu: positions must have num_tokens_full entries"); + + const int64_t num_tokens_insert = slot_mapping.size(0); + TORCH_CHECK(kv.size(0) >= num_tokens_insert, "fused_qnorm_rope_kv_insert_cpu: kv must have >= num_tokens_insert rows"); + TORCH_CHECK(cache_block_size > 0, "fused_qnorm_rope_kv_insert_cpu: cache_block_size must be positive"); + TORCH_CHECK( + cos_sin_cache.scalar_type() == at::kFloat && cos_sin_cache.is_contiguous(), + "fused_qnorm_rope_kv_insert_cpu: cos_sin_cache must be float32 and " + "contiguous (convert once in Python)"); + TORCH_CHECK( + q.is_contiguous() && kv.is_contiguous() && positions.is_contiguous() && slot_mapping.is_contiguous(), + "fused_qnorm_rope_kv_insert_cpu: q/kv/positions/slot_mapping must be contiguous"); + + // swa_kv_cache_2d is mutated in place; index via its real row stride + // rather than assuming it equals cache_block_size * (token_data_bytes + + // scale_bytes_per_token) -- a paged/pooled buffer's row stride is not + // guaranteed to match the logical per-block byte count. + const int64_t cache_row_stride = swa_kv_cache_2d.stride(0); + + auto q_out = at::zeros({num_tokens_full, q_head_padded, head_dim}, q.options()); + + AT_DISPATCH_REDUCED_FLOATING_TYPES(q.scalar_type(), "fused_qnorm_rope_kv_insert_cpu", [&] { + qnorm_rope_kv_insert_impl( + q_out.data_ptr(), + swa_kv_cache_2d.data_ptr(), + q.data_ptr(), + kv.data_ptr(), + positions.data_ptr(), + slot_mapping.data_ptr(), + cos_sin_cache.data_ptr(), + num_tokens_full, + num_heads_q, + q_head_padded, + head_dim, + nope_dim, + num_quant_blocks, + token_data_bytes, + scale_bytes_per_token, + num_tokens_insert, + cache_block_size, + cache_row_stride, + static_cast(eps)); + }); + + return q_out; +} + +// vLLM-native: inverse GPT-J RoPE for `_o_proj`'s output de-rotation (undoes +// the query's own rotation on the attention output before the low-rank +// wo_a/wo_b projection -- see the math writeup in cpu_sparse.py's `_o_proj`). +// Upstream needs this same step (`fused_rope_inplace(..., inverse=True)` in +// `deepseek_v4.py`) but only has a CUDA-JIT implementation, no CPU kernel -- +// this fills that gap rather than porting one. Replaces the eager reference +// in `DeepseekV4CPUAttention._o_proj` (cpu_sparse.py). Uses the same +// half-half `[cos...|sin...]` cache layout as `apply_gptj_rope_inplace` +// above (not upstream's interleaved `[c0,s0,c1,s1,...]` layout). +// +// AVX512 path: NoPE dtype-convert copy uses `vec.h`'s `load_float_vec` +// (16-wide bf16/fp16 -> fp32). The interleaved (even,odd) rotation widens +// 32 raw elements to fp32 via `load_float_vec2`, round-trips them through +// an aligned scratch buffer, then gathers/scatters the even/odd lanes -- +// same technique as `indexer.cpp`'s Q-side RoPE, since `cos`/`sin` here +// are two contiguous halves rather than per-pair-interleaved. +at::Tensor inverse_gptj_rope_o_proj_cpu( + at::Tensor& o, // [num_tokens, num_heads, head_dim] + at::Tensor& positions, // [num_tokens] + at::Tensor& cos_sin_cache, // [max_pos, rope_dim] (first half cos, second half sin) + int64_t rope_dim) { + TORCH_CHECK(o.dim() == 3, "inverse_gptj_rope_o_proj_cpu: o must be 3D [num_tokens, num_heads, head_dim]"); + TORCH_CHECK(positions.scalar_type() == at::kLong, "inverse_gptj_rope_o_proj_cpu: positions must be int64"); + TORCH_CHECK( + o.device().is_cpu() && positions.device().is_cpu() && cos_sin_cache.device().is_cpu(), + "inverse_gptj_rope_o_proj_cpu: all inputs must be CPU tensors"); + + const int64_t num_tokens = o.size(0); + const int64_t num_heads = o.size(1); + const int64_t head_dim = o.size(2); + const int64_t nope_dim = head_dim - rope_dim; + const int64_t half_dim = rope_dim / 2; + TORCH_CHECK(nope_dim >= 0, "inverse_gptj_rope_o_proj_cpu: rope_dim must not exceed head_dim"); + // Only ever called with head_dim=512, rope_dim=64 (DeepseekV4CPUAttention's + // _o_proj); the AVX512 loops below process 16 elements/pairs per iteration + // with no scalar remainder, so require both to divide evenly. + TORCH_CHECK(nope_dim % 16 == 0 && half_dim % 16 == 0, + "inverse_gptj_rope_o_proj_cpu: head_dim-rope_dim and rope_dim/2 must be multiples of 16"); + TORCH_CHECK( + cos_sin_cache.scalar_type() == at::kFloat && cos_sin_cache.is_contiguous(), + "inverse_gptj_rope_o_proj_cpu: cos_sin_cache must be float32 and " + "contiguous (convert once in Python)"); + TORCH_CHECK(o.is_contiguous() && positions.is_contiguous(), + "inverse_gptj_rope_o_proj_cpu: o/positions must be contiguous"); + + auto out = at::empty({num_tokens, num_heads, head_dim}, o.options().dtype(at::kFloat)); + + AT_DISPATCH_REDUCED_FLOATING_TYPES(o.scalar_type(), "inverse_gptj_rope_o_proj_cpu", [&] { + const scalar_t* o_ptr = o.data_ptr(); + float* out_ptr = out.data_ptr(); + const int64_t* pos_ptr = positions.data_ptr(); + const float* cos_sin_ptr = cos_sin_cache.data_ptr(); + + at::parallel_for(0, num_tokens * num_heads, 0, [&](int64_t begin, int64_t end) { + for (int64_t idx = begin; idx < end; ++idx) { + const int64_t t = idx / num_heads; + const int64_t h = idx % num_heads; + const scalar_t* src = o_ptr + (t * num_heads + h) * head_dim; + float* dst = out_ptr + (t * num_heads + h) * head_dim; + + // head_dim=512, rope_dim=64 in practice (checked above), so + // nope_dim/half_dim are always multiples of 16 -- no scalar + // remainder needed. + for (int64_t i = 0; i < nope_dim; i += 16) { + load_float_vec(src + i).store(dst + i); + } + + const int64_t pos = pos_ptr[t]; + const float* cos = cos_sin_ptr + pos * rope_dim; + const float* sin = cos + half_dim; + + { + // Permute-based (no gather/scatter), adapted from SGLang's + // `rope.cpp` (`apply_rotary_emb_row`) the same way as + // compressor.cpp's `apply_gptj_rope_inplace`: `_mm512_permute_ps` + // swaps each pair's (even, odd) lanes in one shuffle, and cos/sin + // -- loaded as 8 contiguous values -- are duplicated pairwise via + // `_mm512_permutexvar_ps` to align with that layout, since our + // cos/sin are two separate flat arrays rather than SGLang's + // pre-interleaved freqs. Two batches of 8 pairs per iteration + // (matching `load_float_vec2`'s 16+16-lane split). + const __m512i dup_idx = _mm512_set_epi32(7, 7, 6, 6, 5, 5, 4, 4, 3, 3, 2, 2, 1, 1, 0, 0); + // Inverse GPT-J rotation negates the sin broadcast on the odd + // (imag) lane of each pair; the even (real) lane keeps sin's sign. + const __m512 sign_mask = _mm512_castsi512_ps(_mm512_set_epi32( + (int)0x80000000, 0, (int)0x80000000, 0, (int)0x80000000, 0, (int)0x80000000, 0, + (int)0x80000000, 0, (int)0x80000000, 0, (int)0x80000000, 0, (int)0x80000000, 0)); + for (int64_t j = 0; j < half_dim; j += 16) { + const scalar_t* rsrc = src + nope_dim + 2 * j; + at::vec::Vectorized x0, x1; + std::tie(x0, x1) = load_float_vec2(rsrc); + float* rdst = dst + nope_dim + 2 * j; + + __m512 xv0 = x0; + __m512 cos_b0 = _mm512_permutexvar_ps(dup_idx, _mm512_castps256_ps512(_mm256_loadu_ps(cos + j))); + __m512 sin_b0 = _mm512_permutexvar_ps(dup_idx, _mm512_castps256_ps512(_mm256_loadu_ps(sin + j))); + __m512 out0 = _mm512_fmadd_ps( + xv0, cos_b0, _mm512_mul_ps(_mm512_permute_ps(xv0, 0xB1), _mm512_xor_ps(sin_b0, sign_mask))); + _mm512_storeu_ps(rdst, out0); + + __m512 xv1 = x1; + __m512 cos_b1 = _mm512_permutexvar_ps(dup_idx, _mm512_castps256_ps512(_mm256_loadu_ps(cos + j + 8))); + __m512 sin_b1 = _mm512_permutexvar_ps(dup_idx, _mm512_castps256_ps512(_mm256_loadu_ps(sin + j + 8))); + __m512 out1 = _mm512_fmadd_ps( + xv1, cos_b1, _mm512_mul_ps(_mm512_permute_ps(xv1, 0xB1), _mm512_xor_ps(sin_b1, sign_mask))); + _mm512_storeu_ps(rdst + 16, out1); + } + } + } + }); + }); + + return out; +} diff --git a/csrc/cpu/sgl-kernels/topk.cpp b/csrc/cpu/sgl-kernels/topk.cpp new file mode 100644 index 0000000000..c41d680e9e --- /dev/null +++ b/csrc/cpu/sgl-kernels/topk.cpp @@ -0,0 +1,834 @@ +// Adapted from +// https://github.com/sgl-project/sglang/tree/main/sgl-kernel/csrc/cpu: +// `biased_topk_kernel_impl`/`biased_topk_cpu` (flat, non-grouped DeepSeek-V4 +// biased top-k) and `hash_topk_kernel_impl`/`hash_topk_cpu` (hash-routed +// layers, expert IDs from a precomputed `tid2eid` lookup table), plus their +// `sigmoid`/`softmax`/`apply_bias`/`sqrtsoftplus` scoring helpers. See +// `biased_topk_cpu`/`hash_topk_cpu` below for the one deliberate deviation +// from the upstream kernel (routed_scaling_factor handling when no shared +// expert is fused into routing, vLLM's CPU DeepSeek-V4 calling convention). +// +// clang-format off + +#include +#include +#include +#include +#include +#include + +#include "common.h" +#include "vec.h" + +namespace { + +template , int> = 0> +inline void sigmoid(float* __restrict__ out, const scalar_t* __restrict__ input) { + using bVec = at::vec::Vectorized; + using fVec = at::vec::Vectorized; + + const fVec one = fVec(1.f); + + constexpr int kVecSize = bVec::size(); + for (int d = 0; d < SIZE; d += kVecSize) { + bVec x_bvec = bVec::loadu(input + d); + fVec x_fvec0, x_fvec1; + std::tie(x_fvec0, x_fvec1) = at::vec::convert_to_float(x_bvec); + + x_fvec0 = one / (one + x_fvec0.neg().exp_u20()); + x_fvec1 = one / (one + x_fvec1.neg().exp_u20()); + + x_fvec0.store(out + d); + x_fvec1.store(out + d + fVec::size()); + } +} + +template , int> = 0> +inline void sigmoid(float* __restrict__ out, const float* __restrict__ input) { + using fVec = at::vec::Vectorized; + const fVec one = fVec(1.f); + constexpr int kVecSize = fVec::size(); + for (int d = 0; d < SIZE; d += kVecSize) { + fVec in_fvec = fVec::loadu(input + d); + in_fvec = one / (one + in_fvec.neg().exp_u20()); + in_fvec.store(out + d); + } +} + +template , int> = 0> +inline void softmax(float* __restrict__ out, const scalar_t* __restrict__ input) { + using bVec = at::vec::Vectorized; + using fVec = at::vec::Vectorized; + + constexpr int kVecSize = bVec::size(); + + // step 1: get max + fVec max_fvec = fVec(-std::numeric_limits::infinity()); + for (int d = 0; d < SIZE; d += kVecSize) { + bVec x_bvec = bVec::loadu(input + d); + fVec x_fvec0, x_fvec1; + std::tie(x_fvec0, x_fvec1) = at::vec::convert_to_float(x_bvec); + + max_fvec = at::vec::maximum(max_fvec, x_fvec0); + max_fvec = at::vec::maximum(max_fvec, x_fvec1); + x_fvec0.store(out + d); + x_fvec1.store(out + d + fVec::size()); + } + float max_val = vec_reduce_max(max_fvec); + max_fvec = fVec(max_val); + + // step 2: sum of (x - max).exp() + fVec sum_fvec = fVec(float(0)); + for (int d = 0; d < SIZE; d += fVec::size()) { + fVec x_fvec = (fVec::loadu(out + d) - max_fvec).exp_u20(); + sum_fvec += x_fvec; + x_fvec.store(out + d); + } + float sum_val = vec_reduce_sum(sum_fvec); + + // step 3: x * (1 / sum) + sum_fvec = fVec(1.f / sum_val); + for (int d = 0; d < SIZE; d += fVec::size()) { + fVec out_fvec = fVec::loadu(out + d) * sum_fvec; + out_fvec.store(out + d); + } +} + +template , int> = 0> +inline void softmax(float* __restrict__ out, const float* __restrict__ input) { + using fVec = at::vec::Vectorized; + + constexpr int kVecSize = fVec::size(); + + // step 1: get max + fVec max_fvec = fVec(-std::numeric_limits::infinity()); + for (int d = 0; d < SIZE; d += kVecSize) { + fVec x_fvec = fVec::loadu(input + d); + max_fvec = at::vec::maximum(max_fvec, x_fvec); + x_fvec.store(out + d); + } + float max_val = vec_reduce_max(max_fvec); + max_fvec = fVec(max_val); + + // step 2: sum of (x - max).exp() + fVec sum_fvec = fVec(float(0)); + for (int d = 0; d < SIZE; d += kVecSize) { + fVec x_fvec = (fVec::loadu(out + d) - max_fvec).exp_u20(); + sum_fvec += x_fvec; + x_fvec.store(out + d); + } + float sum_val = vec_reduce_sum(sum_fvec); + + // step 3: x * (1 / sum) + sum_fvec = fVec(1.f / sum_val); + for (int d = 0; d < SIZE; d += kVecSize) { + fVec out_fvec = fVec::loadu(out + d) * sum_fvec; + out_fvec.store(out + d); + } +} + +template +inline void +apply_bias(float* __restrict__ scores2, const float* __restrict__ scores, const param_t* __restrict__ bias) { + using fVec = at::vec::Vectorized; + auto vec_size = fVec::size() * 2; + int d = 0; + for (; d <= SIZE - vec_size; d += vec_size) { + fVec bias0, bias1, x0, x1; + std::tie(bias0, bias1) = load_float_vec2(bias + d); + std::tie(x0, x1) = load_float_vec2(scores + d); + x0 = x0 + bias0; + x1 = x1 + bias1; + x0.store(scores2 + d); + x1.store(scores2 + d + fVec::size()); + } + for (; d < SIZE; d++) { + scores2[d] = scores[d] + (float)bias[d]; + } +} + +// sqrtsoftplus: sqrt(softplus(x)) = sqrt(log(1 + exp(x))) +// For numerical stability: when x > threshold, softplus(x) ≈ x +// When x < -threshold, softplus(x) ≈ exp(x), so sqrt(softplus(x)) ≈ sqrt(exp(x)) = exp(x/2) +template , int> = 0> +inline void sqrtsoftplus(float* __restrict__ out, const scalar_t* __restrict__ input) { + using bVec = at::vec::Vectorized; + using fVec = at::vec::Vectorized; + + const fVec one = fVec(1.f); + const fVec half = fVec(0.5f); + const fVec threshold = fVec(20.f); + const fVec neg_threshold = fVec(-15.f); // below this, 1+exp(x) loses precision in float32 + + constexpr int kVecSize = bVec::size(); + for (int d = 0; d < SIZE; d += kVecSize) { + bVec x_bvec = bVec::loadu(input + d); + fVec x_fvec0, x_fvec1; + std::tie(x_fvec0, x_fvec1) = at::vec::convert_to_float(x_bvec); + + // default: softplus(x) = log(1 + exp(x)) + fVec sp0 = (one + x_fvec0.exp_u20()).log(); + fVec sp1 = (one + x_fvec1.exp_u20()).log(); + // x > 20: softplus(x) ≈ x + sp0 = fVec::blendv(sp0, x_fvec0, x_fvec0 > threshold); + sp1 = fVec::blendv(sp1, x_fvec1, x_fvec1 > threshold); + // x < -15: softplus(x) ≈ exp(x), sqrt(exp(x)) = exp(x/2) + fVec exp_half0 = (x_fvec0 * half).exp_u20(); + fVec exp_half1 = (x_fvec1 * half).exp_u20(); + sp0 = fVec::blendv(sp0.sqrt(), exp_half0, x_fvec0 < neg_threshold); + sp1 = fVec::blendv(sp1.sqrt(), exp_half1, x_fvec1 < neg_threshold); + + sp0.store(out + d); + sp1.store(out + d + fVec::size()); + } +} + +template , int> = 0> +inline void sqrtsoftplus(float* __restrict__ out, const float* __restrict__ input) { + using fVec = at::vec::Vectorized; + const fVec one = fVec(1.f); + const fVec half = fVec(0.5f); + const fVec threshold = fVec(20.f); + const fVec neg_threshold = fVec(-15.f); + constexpr int kVecSize = fVec::size(); + for (int d = 0; d < SIZE; d += kVecSize) { + fVec x = fVec::loadu(input + d); + fVec sp = (one + x.exp_u20()).log(); + sp = fVec::blendv(sp, x, x > threshold); + fVec exp_half = (x * half).exp_u20(); + sp = fVec::blendv(sp.sqrt(), exp_half, x < neg_threshold); + sp.store(out + d); + } +} + +// biased_topk: flat (non-grouped) biased topk for DeepSeek V4 +// scoring_func: 0 = sigmoid, 1 = sqrtsoftplus +template +void biased_topk_kernel_impl( + float* __restrict__ topk_weights, + int32_t* __restrict__ topk_ids, + const scalar_t* __restrict__ gating_output, + const param_t* __restrict__ bias, + int64_t num_tokens, + int64_t topk, + bool renormalize, + int scoring_func, + int64_t num_fused_shared_experts, + float routed_scaling_factor, + bool apply_routed_scaling_factor_on_output) { + at::parallel_for(0, num_tokens, 0, [&](int64_t begin, int64_t end) { + alignas(64) float scores[NUM_EXPERTS]; + alignas(64) float scores_biased[NUM_EXPERTS]; + + using elem_t = std::pair; + std::array queue; + + // simple RNG for fused shared expert random ID + uint64_t rng_state = begin * 6364136223846793005ULL + 1442695040888963407ULL; + + for (int64_t i = begin; i < end; ++i) { + // compute scores + if (scoring_func == 0) { + sigmoid(scores, gating_output + i * NUM_EXPERTS); + } else { + sqrtsoftplus(scores, gating_output + i * NUM_EXPERTS); + } + + // add bias for selection + apply_bias(scores_biased, scores, bias); + + // build queue and partial sort to find top-k + for (int64_t e = 0; e < NUM_EXPERTS; ++e) { + queue[e] = {scores_biased[e], static_cast(e)}; + } + + std::partial_sort( + queue.begin(), queue.begin() + topk, queue.end(), [](const elem_t& x, const elem_t& y) -> bool { + return x.first > y.first; + }); + + // gather original scores (without bias) as weights + for (int64_t j = 0; j < topk; ++j) { + int32_t idx = queue[j].second; + topk_ids[i * topk + j] = idx; + topk_weights[i * topk + j] = scores[idx]; + } + + // handle fused shared experts + if (num_fused_shared_experts > 0) { + // replace last slot with random shared expert ID + rng_state = rng_state * 6364136223846793005ULL + 1442695040888963407ULL; + int32_t shared_id = + NUM_EXPERTS + static_cast((rng_state >> 33) % static_cast(num_fused_shared_experts)); + topk_ids[i * topk + topk - 1] = shared_id; + + // shared expert weight = sum of routed weights / scaling_factor + if (routed_scaling_factor != 0.0f) { + float routed_sum = 0.f; + for (int64_t j = 0; j < topk - 1; ++j) { + routed_sum += topk_weights[i * topk + j]; + } + topk_weights[i * topk + topk - 1] = routed_sum / routed_scaling_factor; + } + } + + // renormalize + if (renormalize) { + float sum = 0.f; + int64_t norm_end = (num_fused_shared_experts == 0) ? topk : topk - 1; + for (int64_t j = 0; j < norm_end; ++j) { + sum += topk_weights[i * topk + j]; + } + float scale = 1.f / sum; + if (apply_routed_scaling_factor_on_output) { + scale *= routed_scaling_factor; + } + for (int64_t j = 0; j < norm_end; ++j) { + topk_weights[i * topk + j] *= scale; + } + // also scale the shared expert weight if present + if (num_fused_shared_experts > 0) { + topk_weights[i * topk + topk - 1] *= scale; + } + } + } + }); +} + +// hash_topk: expert IDs come from a precomputed lookup table tid2eid[input_ids] +// scoring_func: 0 = softmax, 1 = sigmoid, 2 = sqrtsoftplus +template +void hash_topk_kernel_impl( + float* __restrict__ topk_weights, + int32_t* __restrict__ topk_ids, + const scalar_t* __restrict__ gating_output, + const int32_t* __restrict__ tid2eid, // [num_tokens, routed_topk] + int64_t num_tokens, + int scoring_func, + int64_t num_fused_shared_experts, + int64_t num_experts, + float routed_scaling_factor, + int64_t topk) { + const int64_t routed_topk = topk - num_fused_shared_experts; + const bool need_renormalize = (scoring_func != 0); // renormalize for non-softmax + + at::parallel_for(0, num_tokens, 0, [&](int64_t begin, int64_t end) { + alignas(64) float scores[NUM_EXPERTS]; + + // simple RNG for fused shared expert random ID + uint64_t rng_state = begin * 6364136223846793005ULL + 1442695040888963407ULL; + + for (int64_t i = begin; i < end; ++i) { + // compute scores over all experts + if (scoring_func == 0) { + softmax(scores, gating_output + i * NUM_EXPERTS); + } else if (scoring_func == 1) { + sigmoid(scores, gating_output + i * NUM_EXPERTS); + } else { + sqrtsoftplus(scores, gating_output + i * NUM_EXPERTS); + } + + // gather expert IDs from lookup table + const int32_t* eid_row = tid2eid + i * routed_topk; + for (int64_t j = 0; j < routed_topk; ++j) { + int32_t eid = eid_row[j]; + topk_ids[i * topk + j] = eid; + topk_weights[i * topk + j] = scores[eid]; + } + + // renormalize routed weights (for non-softmax scoring). Unlike the + // upstream kernel, always fold routed_scaling_factor in here rather + // than only applying it to the fused-shared-expert slot below: vLLM's + // DeepSeek-V4 CPU hash-routed layers never fuse a shared expert into + // this routing table (num_fused_shared_experts is always 0 on this + // call path), so this is the only place the scaling factor can be + // applied. Matches the eager reference (`sqrtsoftplus_bias_topk` in + // cpu_moe.py) and the CUDA kernel (`dsv4HashTopkSoftplusSqrt`). + if (need_renormalize) { + float sum = 0.f; + for (int64_t j = 0; j < routed_topk; ++j) { + sum += topk_weights[i * topk + j]; + } + if (sum > 0.f) { + float scale = 1.f / sum; + if (num_fused_shared_experts == 0) { + scale *= routed_scaling_factor; + } + for (int64_t j = 0; j < routed_topk; ++j) { + topk_weights[i * topk + j] *= scale; + } + } + } + + // handle fused shared expert + if (num_fused_shared_experts > 0) { + // random shared expert ID in [num_experts, num_experts + num_fused_shared_experts) + rng_state = rng_state * 6364136223846793005ULL + 1442695040888963407ULL; + int32_t shared_id = + num_experts + static_cast((rng_state >> 33) % static_cast(num_fused_shared_experts)); + topk_ids[i * topk + topk - 1] = shared_id; + + // shared expert weight = sum of routed weights / scaling_factor + float routed_sum = 0.f; + for (int64_t j = 0; j < routed_topk; ++j) { + routed_sum += topk_weights[i * topk + j]; + } + topk_weights[i * topk + topk - 1] = routed_sum / routed_scaling_factor; + } + } + }); +} + +#define LAUNCH_HASH_TOPK_KERNEL(NE, NTOPK) \ + hash_topk_kernel_impl( \ + topk_weights.data_ptr(), \ + topk_ids.data_ptr(), \ + gating_output.data_ptr(), \ + tid2eid_flat.data_ptr(), \ + num_tokens, \ + scoring_func_id, \ + num_fused_shared_experts, \ + num_experts, \ + routed_scaling_factor_value, \ + topk); + +#define LAUNCH_BIASED_TOPK_KERNEL(NE, NTOPK) \ + biased_topk_kernel_impl( \ + topk_weights.data_ptr(), \ + topk_ids.data_ptr(), \ + gating_output.data_ptr(), \ + correction_bias.data_ptr(), \ + num_tokens, \ + topk, \ + renormalize, \ + scoring_func_id, \ + num_fused_shared_experts, \ + routed_scaling_factor_value, \ + apply_routed_scaling_factor_on_output); + +} // namespace + +// biased topk for DeepSeek V4 (flat, non-grouped) +std::tuple biased_topk_cpu( + at::Tensor& hidden_states, + at::Tensor& gating_output, + at::Tensor& correction_bias, + int64_t topk, + bool renormalize, + std::string scoring_func, + int64_t num_fused_shared_experts, + std::optional routed_scaling_factor, + bool apply_routed_scaling_factor_on_output) { + CHECK_INPUT(gating_output); + CHECK_INPUT(correction_bias); + + const auto st = gating_output.scalar_type(); + int64_t num_tokens = hidden_states.size(0); + int64_t num_experts = gating_output.size(1); + TORCH_CHECK(gating_output.size(0) == num_tokens, "Number of tokens mismatch"); + TORCH_CHECK(correction_bias.numel() == num_experts, "Bias shape mismatch"); + + int scoring_func_id = 0; + if (scoring_func == "sigmoid") { + scoring_func_id = 0; + } else if (scoring_func == "sqrtsoftplus") { + scoring_func_id = 1; + } else { + TORCH_CHECK(false, "Unsupported scoring_func: ", scoring_func); + } + + float routed_scaling_factor_value = routed_scaling_factor.has_value() ? routed_scaling_factor.value() : 0.0f; + + at::Tensor topk_weights = at::empty({num_tokens, topk}, hidden_states.options().dtype(at::kFloat)); + at::Tensor topk_ids = at::empty({num_tokens, topk}, hidden_states.options().dtype(at::kInt)); + + // The actual routed topk (excluding fused shared experts) + int64_t routed_topk = topk - num_fused_shared_experts; + + CPU_DISPATCH_FLOATING_TYPES_EXT(st, correction_bias.scalar_type(), "biased_topk_kernel", [&] { + // dispatch on num_experts and routed_topk. For DeepSeek V4: 256 experts, + // routed_topk=6 (num_fused_shared_experts is always 0 on this CPU call + // path -- see cpu_moe.py). Anything outside this table isn't a real + // model shape, so fail loudly instead of silently degrading to a slow + // generic path. + switch (num_experts) { + case 64: + switch (routed_topk) { + case 6: + LAUNCH_BIASED_TOPK_KERNEL(64, 6); + break; + case 8: + LAUNCH_BIASED_TOPK_KERNEL(64, 8); + break; + default: + TORCH_CHECK(false, "biased_topk_cpu: unsupported routed_topk ", routed_topk, + " for num_experts=64"); + } + break; + case 128: + switch (routed_topk) { + case 6: + LAUNCH_BIASED_TOPK_KERNEL(128, 6); + break; + case 8: + LAUNCH_BIASED_TOPK_KERNEL(128, 8); + break; + default: + TORCH_CHECK(false, "biased_topk_cpu: unsupported routed_topk ", routed_topk, + " for num_experts=128"); + } + break; + case 256: + switch (routed_topk) { + case 6: + LAUNCH_BIASED_TOPK_KERNEL(256, 6); + break; + case 8: + LAUNCH_BIASED_TOPK_KERNEL(256, 8); + break; + case 9: + LAUNCH_BIASED_TOPK_KERNEL(256, 9); + break; + default: + TORCH_CHECK(false, "biased_topk_cpu: unsupported routed_topk ", routed_topk, + " for num_experts=256"); + } + break; + case 384: + switch (routed_topk) { + case 6: + LAUNCH_BIASED_TOPK_KERNEL(384, 6); + break; + case 8: + LAUNCH_BIASED_TOPK_KERNEL(384, 8); + break; + default: + TORCH_CHECK(false, "biased_topk_cpu: unsupported routed_topk ", routed_topk, + " for num_experts=384"); + } + break; + default: + TORCH_CHECK(false, "biased_topk_cpu: unsupported num_experts ", num_experts); + } + }); + return std::make_tuple(topk_weights, topk_ids); +} + +// hash topk for DeepSeek V4 (expert IDs from precomputed lookup table) +std::tuple hash_topk_cpu( + at::Tensor& gating_output, + at::Tensor& tid2eid, + int64_t topk, + std::string scoring_func, + int64_t num_fused_shared_experts, + int64_t num_experts, + double routed_scaling_factor) { + CHECK_INPUT(gating_output); + CHECK_INPUT(tid2eid); + + const auto st = gating_output.scalar_type(); + int64_t num_tokens = gating_output.size(0); + int64_t num_experts_gating = gating_output.size(1); + int64_t routed_topk = topk - num_fused_shared_experts; + + TORCH_CHECK(tid2eid.size(0) == num_tokens, "tid2eid row count must match num_tokens"); + TORCH_CHECK(tid2eid.size(1) == routed_topk, "tid2eid column count must match routed_topk"); + TORCH_CHECK(tid2eid.scalar_type() == at::kInt, "tid2eid must be int32"); + TORCH_CHECK(num_experts_gating == num_experts, "num_experts mismatch"); + + int scoring_func_id = 0; + if (scoring_func == "softmax") { + scoring_func_id = 0; + } else if (scoring_func == "sigmoid") { + scoring_func_id = 1; + } else if (scoring_func == "sqrtsoftplus") { + scoring_func_id = 2; + } else { + TORCH_CHECK(false, "Unsupported scoring_func: ", scoring_func); + } + + float routed_scaling_factor_value = static_cast(routed_scaling_factor); + + at::Tensor topk_weights = at::empty({num_tokens, topk}, gating_output.options().dtype(at::kFloat)); + at::Tensor topk_ids = at::empty({num_tokens, topk}, gating_output.options().dtype(at::kInt)); + + // tid2eid is [num_tokens, routed_topk], already indexed by input_ids in + // Python; CHECK_INPUT above already guarantees it is contiguous. + const at::Tensor& tid2eid_flat = tid2eid; + + // Dispatch for bf16, fp16, and float32 + // Note: cannot use AT_DISPATCH_FLOATING_TYPES_AND2 since it includes double which lacks convert_to_float + auto dispatch_fn = [&]() { + // For DeepSeek V4: 256 experts, routed_topk in {6,7,8} depending on + // layer (num_fused_shared_experts is always 0 on this CPU call path -- + // see cpu_moe.py). Anything outside this table isn't a real model + // shape, so fail loudly instead of silently degrading to a slow + // generic path. + switch (num_experts) { + case 64: + switch (routed_topk) { + case 6: + LAUNCH_HASH_TOPK_KERNEL(64, 6); + break; + case 7: + LAUNCH_HASH_TOPK_KERNEL(64, 7); + break; + case 8: + LAUNCH_HASH_TOPK_KERNEL(64, 8); + break; + default: + TORCH_CHECK(false, "hash_topk_cpu: unsupported routed_topk ", routed_topk, + " for num_experts=64"); + } + break; + case 128: + switch (routed_topk) { + case 6: + LAUNCH_HASH_TOPK_KERNEL(128, 6); + break; + case 7: + LAUNCH_HASH_TOPK_KERNEL(128, 7); + break; + case 8: + LAUNCH_HASH_TOPK_KERNEL(128, 8); + break; + default: + TORCH_CHECK(false, "hash_topk_cpu: unsupported routed_topk ", routed_topk, + " for num_experts=128"); + } + break; + case 256: + switch (routed_topk) { + case 6: + LAUNCH_HASH_TOPK_KERNEL(256, 6); + break; + case 7: + LAUNCH_HASH_TOPK_KERNEL(256, 7); + break; + case 8: + LAUNCH_HASH_TOPK_KERNEL(256, 8); + break; + default: + TORCH_CHECK(false, "hash_topk_cpu: unsupported routed_topk ", routed_topk, + " for num_experts=256"); + } + break; + case 384: + switch (routed_topk) { + case 6: + LAUNCH_HASH_TOPK_KERNEL(384, 6); + break; + case 7: + LAUNCH_HASH_TOPK_KERNEL(384, 7); + break; + case 8: + LAUNCH_HASH_TOPK_KERNEL(384, 8); + break; + default: + TORCH_CHECK(false, "hash_topk_cpu: unsupported routed_topk ", routed_topk, + " for num_experts=384"); + } + break; + default: + TORCH_CHECK(false, "hash_topk_cpu: unsupported num_experts ", num_experts); + } + }; + + if (st == at::ScalarType::BFloat16) { + dispatch_fn.template operator()(); + } else if (st == at::ScalarType::Half) { + dispatch_fn.template operator()(); + } else if (st == at::ScalarType::Float) { + dispatch_fn.template operator()(); + } else { + TORCH_CHECK(false, "Unsupported dtype for hash_topk_cpu: ", st); + } + + return std::make_tuple(topk_weights, topk_ids); +} + +// Adapted from +// https://github.com/sgl-project/sglang/tree/main/sgl-kernel/csrc/cpu: +// `topk_transform_512_cpu`/`topk_transform_512_cpu_kernel_impl`. Near- +// verbatim port for the sparse indexer's DECODE path: hardcodes top-512 +// (the only `index_topk` DeepSeek-V4-Flash uses) and streams a min-heap over +// the tail of the row instead of sorting the whole row. +// +// Differences from upstream: +// - `out_page_indices` is still computed (required by the kernel contract) +// but vLLM's caller only consumes `out_raw_indices` (local/compressed- +// context positions -- `DeepseekV4CPUAttention.forward_mqa` resolves those +// to physical slots itself via `map_local_to_global_slots_cpu`). +// - `page_size` was already a runtime argument upstream, unlike +// `fp8_paged_mqa_logits_cpu` (paged_mqa_logits.cpp), which needed one. +namespace { + +constexpr int64_t kC4Topk = 512; + +template +inline int64_t load_index_value(const T* __restrict__ ptr, int64_t idx) { + return static_cast(ptr[idx]); +} + +struct TopKTransformElem { + float score; + int32_t index; +}; + +struct TopKTransformMinHeapCmp { + bool operator()(const TopKTransformElem& lhs, const TopKTransformElem& rhs) const { + if (lhs.score == rhs.score) { + return lhs.index > rhs.index; + } + return lhs.score > rhs.score; + } +}; + +template +void topk_transform_512_cpu_kernel_impl( + const float* __restrict__ scores, + const seq_t* __restrict__ seq_lens, + const page_t* __restrict__ page_tables, + int32_t* __restrict__ out_page_indices, + int32_t* __restrict__ out_raw_indices, + int64_t batch_size, + int64_t max_seq_len, + int64_t page_table_stride, + int64_t out_stride, + int64_t page_size) { + TORCH_CHECK(page_size > 0, "page_size must be positive"); + TORCH_CHECK((page_size & (page_size - 1)) == 0, "page_size must be a power of 2"); + const int page_bits = page_size > 1 ? static_cast(std::log2(static_cast(page_size))) : 0; + const int64_t page_mask = page_size - 1; + + at::parallel_for(0, batch_size, 0, [&](int64_t begin, int64_t end) { + std::array heap; + + for (int64_t b = begin; b < end; ++b) { + const float* __restrict__ scores_row = scores + b * max_seq_len; + const page_t* __restrict__ page_table_row = page_tables + b * page_table_stride; + int32_t* __restrict__ out_page_row = out_page_indices + b * out_stride; + int32_t* __restrict__ out_raw_row = out_raw_indices == nullptr ? nullptr : out_raw_indices + b * out_stride; + + int64_t seq_len = load_index_value(seq_lens, b); + seq_len = std::max(0, std::min(seq_len, max_seq_len)); + const int64_t valid_topk = std::min(seq_len, kC4Topk); + + auto store_slot = [&](int64_t slot, int32_t raw_index) { + if (raw_index < 0) { + out_page_row[slot] = -1; + if (out_raw_row != nullptr) { + out_raw_row[slot] = -1; + } + return; + } + + const int64_t page_idx = static_cast(raw_index) >> page_bits; + const int64_t offset_in_page = static_cast(raw_index) & page_mask; + const int64_t physical_page = load_index_value(page_table_row, page_idx); + out_page_row[slot] = static_cast((physical_page << page_bits) | offset_in_page); + if (out_raw_row != nullptr) { + out_raw_row[slot] = raw_index; + } + }; + + if (seq_len <= kC4Topk) { + for (int64_t i = 0; i < valid_topk; ++i) { + store_slot(i, static_cast(i)); + } + for (int64_t i = valid_topk; i < kC4Topk; ++i) { + store_slot(i, -1); + } + continue; + } + + for (int64_t i = 0; i < kC4Topk; ++i) { + heap[i] = {scores_row[i], static_cast(i)}; + } + std::make_heap(heap.begin(), heap.end(), TopKTransformMinHeapCmp()); + + for (int64_t i = kC4Topk; i < seq_len; ++i) { + const float score = scores_row[i]; + const TopKTransformElem& current_min = heap.front(); + if (score > current_min.score || + (score == current_min.score && static_cast(i) < current_min.index)) { + std::pop_heap(heap.begin(), heap.end(), TopKTransformMinHeapCmp()); + heap.back() = {score, static_cast(i)}; + std::push_heap(heap.begin(), heap.end(), TopKTransformMinHeapCmp()); + } + } + + for (int64_t i = 0; i < kC4Topk; ++i) { + store_slot(i, heap[i].index); + } + } + }); +} + +} // namespace + +void topk_transform_512_cpu( + at::Tensor& scores, + at::Tensor& seq_lens, + at::Tensor& page_tables, + at::Tensor& out_page_indices, + int64_t page_size, + const std::optional& out_raw_indices) { + CHECK_INPUT(scores); + CHECK_INPUT(seq_lens); + CHECK_INPUT(page_tables); + CHECK_INPUT(out_page_indices); + + TORCH_CHECK(scores.scalar_type() == at::kFloat, "scores must be float32"); + TORCH_CHECK(out_page_indices.scalar_type() == at::kInt, "out_page_indices must be int32"); + TORCH_CHECK(scores.dim() == 2, "scores must be a 2D tensor"); + TORCH_CHECK(seq_lens.dim() == 1, "seq_lens must be a 1D tensor"); + TORCH_CHECK(page_tables.dim() == 2, "page_tables must be a 2D tensor"); + TORCH_CHECK(out_page_indices.dim() == 2, "out_page_indices must be a 2D tensor"); + + const int64_t batch_size = scores.size(0); + const int64_t max_seq_len = scores.size(1); + TORCH_CHECK(seq_lens.size(0) == batch_size, "seq_lens row count must match scores"); + TORCH_CHECK(page_tables.size(0) == batch_size, "page_tables row count must match scores"); + TORCH_CHECK(out_page_indices.size(0) == batch_size, "out_page_indices row count must match scores"); + TORCH_CHECK(out_page_indices.size(1) >= kC4Topk, "out_page_indices must have at least 512 columns"); + + int32_t* raw_ptr = nullptr; + if (out_raw_indices.has_value()) { + at::Tensor raw = out_raw_indices.value(); + CHECK_INPUT(raw); + TORCH_CHECK(raw.scalar_type() == at::kInt, "out_raw_indices must be int32"); + TORCH_CHECK(raw.dim() == 2, "out_raw_indices must be a 2D tensor"); + TORCH_CHECK(raw.sizes() == out_page_indices.sizes(), "out_raw_indices shape must match out_page_indices"); + raw_ptr = raw.data_ptr(); + } + + if (batch_size == 0) { + return; + } + + auto launch_with_seq_type = [&](auto seq_zero) { + using seq_t = decltype(seq_zero); + if (page_tables.scalar_type() == at::kInt) { + topk_transform_512_cpu_kernel_impl( + scores.data_ptr(), seq_lens.data_ptr(), page_tables.data_ptr(), + out_page_indices.data_ptr(), raw_ptr, batch_size, max_seq_len, page_tables.stride(0), + out_page_indices.stride(0), page_size); + } else if (page_tables.scalar_type() == at::kLong) { + topk_transform_512_cpu_kernel_impl( + scores.data_ptr(), seq_lens.data_ptr(), page_tables.data_ptr(), + out_page_indices.data_ptr(), raw_ptr, batch_size, max_seq_len, page_tables.stride(0), + out_page_indices.stride(0), page_size); + } else { + TORCH_CHECK(false, "page_tables must be int32 or int64"); + } + }; + + if (seq_lens.scalar_type() == at::kInt) { + launch_with_seq_type(int32_t{0}); + } else if (seq_lens.scalar_type() == at::kLong) { + launch_with_seq_type(int64_t{0}); + } else { + TORCH_CHECK(false, "seq_lens must be int32 or int64"); + } +} diff --git a/csrc/cpu/torch_bindings.cpp b/csrc/cpu/torch_bindings.cpp index ea3469b30a..8f663a93c3 100644 --- a/csrc/cpu/torch_bindings.cpp +++ b/csrc/cpu/torch_bindings.cpp @@ -66,18 +66,18 @@ at::Tensor convert_weight_packed(at::Tensor& weight); at::Tensor convert_scale_packed(at::Tensor& scale); -at::Tensor fused_experts_cpu( - at::Tensor& hidden_states, at::Tensor& w1, at::Tensor& w2, - at::Tensor& topk_weights, at::Tensor& topk_ids, bool inplace, - int64_t moe_comp_method, const std::optional& w1_scale, - const std::optional& w2_scale, - const std::optional& w1_zero, - const std::optional& w2_zero, - const std::optional> block_size, - const std::optional& w1_bias, - const std::optional& w2_bias, - const std::optional& alpha, const std::optional& limit, - bool is_vnni); +void fused_experts_cpu(at::Tensor& out, at::Tensor& hidden_states, + at::Tensor& w1, at::Tensor& w2, at::Tensor& topk_weights, + at::Tensor& topk_ids, int64_t moe_comp_method, + const std::optional& w1_scale, + const std::optional& w2_scale, + const std::optional& w1_zero, + const std::optional& w2_zero, + const std::optional> block_size, + const std::optional& w1_bias, + const std::optional& w2_bias, + const std::optional& alpha, + const std::optional& limit, bool is_vnni); at::Tensor int8_scaled_mm_with_quant(at::Tensor& mat1, at::Tensor& mat2, at::Tensor& scales2, @@ -124,6 +124,99 @@ void concat_and_cache_mla_cpu(const at::Tensor& kv_c_normed, const at::Tensor& k_pe, at::Tensor& kv_cache, const at::Tensor& slot_mapping); +// Adapted from sglang: DeepSeek-V4 mHC gating kernels +std::tuple hc_pre_fused_cpu( + at::Tensor& x, at::Tensor& hc_fn, at::Tensor& hc_scale, at::Tensor& hc_base, + int64_t hc_mult, int64_t sinkhorn_iters, double rms_eps, double hc_eps); + +at::Tensor hc_post_fused_cpu(at::Tensor& x, at::Tensor& residual, + at::Tensor& post, at::Tensor& comb); + +at::Tensor hc_head_fused_cpu(at::Tensor& x, at::Tensor& hc_fn, + at::Tensor& hc_scale, at::Tensor& hc_base, + double hc_eps, double norm_eps); + +// Adapted from sglang: DeepSeek-V4 fp8_ds_mla cache-write kernel +at::Tensor fused_qnorm_rope_kv_insert_cpu(at::Tensor& q, at::Tensor& kv, + at::Tensor& positions, + at::Tensor& swa_kv_cache_2d, + at::Tensor& slot_mapping, + at::Tensor& cos_sin_cache, + int64_t q_head_padded, double eps, + int64_t cache_block_size); + +// vLLM-native: DeepSeek-V4 sparse MQA attention kernel for the fp8_ds_mla +// cache layout (SWA window + optional compressed top-k index sets). +void flash_mla_with_kvcache_cpu( + at::Tensor& out, at::Tensor& q, at::Tensor& window_cache_2d, + at::Tensor& window_slots, int64_t window_block_size, + at::Tensor& compressed_cache_2d, at::Tensor& compressed_slots, + int64_t compressed_block_size, at::Tensor& attn_sink, double scale); + +// Adapted from sglang: DeepSeek-V4 compressor state-cache write + fused +// compress+RMSNorm+quant+RoPE+store kernels. +void save_partial_states_cpu(at::Tensor& kv, at::Tensor& score, at::Tensor& ape, + at::Tensor& positions, at::Tensor& state_cache, + at::Tensor& slot_mapping); + +void compress_norm_rope_store_cpu( + at::Tensor& state_cache, at::Tensor& gather_slots, at::Tensor& positions, + at::Tensor& kv_slot_mapping, at::Tensor& rms_norm_weight, + double rms_norm_eps, at::Tensor& cos_sin_cache, at::Tensor& kv_cache_2d, + int64_t kv_cache_block_size, int64_t compress_ratio); + +// head_dim=128 indexer-compressor variant of the above (single fp8 quant +// block + raw fp32 scale -- matches the paged indexer K-cache layout +// fp8_paged_mqa_logits_cpu/topk_transform_512_cpu read). +void compress_norm_rope_store_indexer_cpu( + at::Tensor& state_cache, at::Tensor& gather_slots, at::Tensor& positions, + at::Tensor& kv_slot_mapping, at::Tensor& rms_norm_weight, + double rms_norm_eps, at::Tensor& cos_sin_cache, at::Tensor& kv_cache_2d, + int64_t kv_cache_block_size, int64_t compress_ratio); + +// Adapted from sglang: DeepSeek-V4 sparse indexer paged MQA-logits + +// top-512 transform kernels (DECODE path -- reads the paged K-cache +// directly via `page_table`, whole batch in one call). +at::Tensor fp8_paged_mqa_logits_cpu(at::Tensor& q_fp8, at::Tensor& kvcache_fp8, + at::Tensor& weight, at::Tensor& seq_lens, + at::Tensor& page_table, int64_t block_size, + int64_t max_seq_len); + +void topk_transform_512_cpu(at::Tensor& scores, at::Tensor& seq_lens, + at::Tensor& page_tables, + at::Tensor& out_page_indices, int64_t page_size, + const std::optional& out_raw_indices); + +// Adapted from sglang: DeepSeek-V4 MoE routing kernels (flat biased top-k and +// hash-routed-layer top-k). +std::tuple biased_topk_cpu( + at::Tensor& hidden_states, at::Tensor& gating_output, + at::Tensor& correction_bias, int64_t topk, bool renormalize, + std::string scoring_func, int64_t num_fused_shared_experts, + std::optional routed_scaling_factor, + bool apply_routed_scaling_factor_on_output); + +std::tuple hash_topk_cpu( + at::Tensor& gating_output, at::Tensor& tid2eid, int64_t topk, + std::string scoring_func, int64_t num_fused_shared_experts, + int64_t num_experts, double routed_scaling_factor); + +// vLLM-native: DeepSeek-V4 sparse indexer Q-side RoPE + FP8 quant kernel +// (fp8 path only; MXFP4 stays on triton-cpu). +void fused_indexer_q_rope_quant_cpu(at::Tensor& positions, at::Tensor& index_q, + at::Tensor& index_q_cos_sin_cache, + at::Tensor& index_q_fp8, + at::Tensor& index_weights, + double index_weights_softmax_scale, + double index_weights_head_scale, + at::Tensor& index_weights_out); + +// vLLM-native: inverse GPT-J RoPE for DeepSeek-V4 CPU attention's `_o_proj` +// output de-rotation. +at::Tensor inverse_gptj_rope_o_proj_cpu(at::Tensor& o, at::Tensor& positions, + at::Tensor& cos_sin_cache, + int64_t rope_dim); + // Adapted from sglang: INT4 W4A8 kernels std::tuple convert_weight_packed_scale_zp( at::Tensor qweight, // awq: (*, K, N / 8) || gptq: (*, K / 8, N) , int32 @@ -515,13 +608,13 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) { ops.def("convert_scale_packed(Tensor! scale) -> Tensor"); ops.impl("convert_scale_packed", torch::kCPU, &convert_scale_packed); ops.def( - "fused_experts_cpu(Tensor hidden_states, Tensor w1, Tensor w2, Tensor " - "topk_weights, Tensor topk_ids, bool " - "inplace, int moe_comp_method, Tensor? w1_scale, Tensor? w2_scale, " + "fused_experts_cpu(Tensor(a0!) out, Tensor hidden_states, Tensor w1, " + "Tensor w2, Tensor topk_weights, Tensor topk_ids, " + "int moe_comp_method, Tensor? w1_scale, Tensor? w2_scale, " "Tensor? w1_zero, Tensor? w2_zero, int[]? block_size, " "Tensor? w1_bias, Tensor? w2_bias, float? alpha, float? limit, " "bool is_vnni) -> " - "Tensor"); + "()"); ops.impl("fused_experts_cpu", torch::kCPU, &fused_experts_cpu); ops.def( "int8_scaled_mm_with_quant(Tensor mat1, Tensor mat2, Tensor scales2, " @@ -583,6 +676,110 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) { "concat_and_cache_mla_cpu(Tensor kv_c_normed, Tensor k_pe, " "Tensor(a!) kv_cache, Tensor slot_mapping) -> ()"); ops.impl("concat_and_cache_mla_cpu", torch::kCPU, &concat_and_cache_mla_cpu); + + // Adapted from sglang: DeepSeek-V4 mHC gating kernels + ops.def( + "hc_pre_fused_cpu(Tensor x, Tensor hc_fn, Tensor hc_scale, " + "Tensor hc_base, int hc_mult, int sinkhorn_iters, float rms_eps, " + "float hc_eps) -> (Tensor, Tensor, Tensor)"); + ops.impl("hc_pre_fused_cpu", torch::kCPU, &hc_pre_fused_cpu); + + ops.def( + "hc_post_fused_cpu(Tensor x, Tensor residual, Tensor post, " + "Tensor comb) -> Tensor"); + ops.impl("hc_post_fused_cpu", torch::kCPU, &hc_post_fused_cpu); + + ops.def( + "hc_head_fused_cpu(Tensor x, Tensor hc_fn, Tensor " + "hc_scale, Tensor hc_base, float hc_eps, float norm_eps) -> " + "Tensor"); + ops.impl("hc_head_fused_cpu", torch::kCPU, &hc_head_fused_cpu); + + // Adapted from sglang: DeepSeek-V4 fp8_ds_mla cache-write kernel + ops.def( + "fused_qnorm_rope_kv_insert_cpu(Tensor q, Tensor kv, Tensor positions, " + "Tensor(a!) swa_kv_cache_2d, Tensor slot_mapping, Tensor " + "cos_sin_cache, int q_head_padded, float eps, int cache_block_size) " + "-> Tensor"); + ops.impl("fused_qnorm_rope_kv_insert_cpu", torch::kCPU, + &fused_qnorm_rope_kv_insert_cpu); + + // vLLM-native: DeepSeek-V4 sparse MQA attention kernel + ops.def( + "flash_mla_with_kvcache_cpu(Tensor(a!) out, Tensor q, Tensor " + "window_cache_2d, Tensor window_slots, int window_block_size, Tensor " + "compressed_cache_2d, Tensor compressed_slots, int " + "compressed_block_size, Tensor attn_sink, float scale) -> ()"); + ops.impl("flash_mla_with_kvcache_cpu", torch::kCPU, + &flash_mla_with_kvcache_cpu); + + // Adapted from sglang: DeepSeek-V4 compressor kernels + ops.def( + "save_partial_states_cpu(Tensor kv, Tensor score, Tensor ape, " + "Tensor positions, Tensor(a!) state_cache, Tensor slot_mapping) " + "-> ()"); + ops.impl("save_partial_states_cpu", torch::kCPU, &save_partial_states_cpu); + + ops.def( + "compress_norm_rope_store_cpu(Tensor state_cache, Tensor " + "gather_slots, Tensor positions, Tensor kv_slot_mapping, Tensor " + "rms_norm_weight, float rms_norm_eps, Tensor cos_sin_cache, " + "Tensor(a!) kv_cache_2d, int kv_cache_block_size, int compress_ratio) " + "-> ()"); + ops.impl("compress_norm_rope_store_cpu", torch::kCPU, + &compress_norm_rope_store_cpu); + + ops.def( + "compress_norm_rope_store_indexer_cpu(Tensor state_cache, Tensor " + "gather_slots, Tensor positions, Tensor kv_slot_mapping, Tensor " + "rms_norm_weight, float rms_norm_eps, Tensor cos_sin_cache, " + "Tensor(a!) kv_cache_2d, int kv_cache_block_size, int compress_ratio) " + "-> ()"); + ops.impl("compress_norm_rope_store_indexer_cpu", torch::kCPU, + &compress_norm_rope_store_indexer_cpu); + + // Adapted from sglang: DeepSeek-V4 sparse indexer paged MQA-logits kernel + // (DECODE path) + ops.def( + "fp8_paged_mqa_logits_cpu(Tensor q_fp8, Tensor kvcache_fp8, Tensor " + "weight, Tensor seq_lens, Tensor page_table, int block_size, int " + "max_seq_len) -> Tensor"); + ops.impl("fp8_paged_mqa_logits_cpu", torch::kCPU, &fp8_paged_mqa_logits_cpu); + + ops.def( + "topk_transform_512_cpu(Tensor scores, Tensor seq_lens, Tensor " + "page_tables, Tensor(a!) out_page_indices, int page_size, Tensor(b!)? " + "out_raw_indices) -> ()"); + ops.impl("topk_transform_512_cpu", torch::kCPU, &topk_transform_512_cpu); + + ops.def( + "fused_indexer_q_rope_quant_cpu(Tensor positions, Tensor index_q, " + "Tensor index_q_cos_sin_cache, Tensor(a!) index_q_fp8, Tensor " + "index_weights, float index_weights_softmax_scale, float " + "index_weights_head_scale, Tensor(b!) index_weights_out) -> ()"); + ops.impl("fused_indexer_q_rope_quant_cpu", torch::kCPU, + &fused_indexer_q_rope_quant_cpu); + + // Adapted from sglang: DeepSeek-V4 MoE routing kernels + ops.def( + "biased_topk_cpu(Tensor hidden_states, Tensor gating_output, Tensor " + "correction_bias, int topk, bool renormalize, str scoring_func, int " + "num_fused_shared_experts, float? routed_scaling_factor, bool " + "apply_routed_scaling_factor_on_output) -> (Tensor, Tensor)"); + ops.impl("biased_topk_cpu", torch::kCPU, &biased_topk_cpu); + + ops.def( + "hash_topk_cpu(Tensor gating_output, Tensor tid2eid, int topk, str " + "scoring_func, int num_fused_shared_experts, int num_experts, float " + "routed_scaling_factor) -> (Tensor, Tensor)"); + ops.impl("hash_topk_cpu", torch::kCPU, &hash_topk_cpu); + + // vLLM-native: inverse GPT-J RoPE for _o_proj's output de-rotation + ops.def( + "inverse_gptj_rope_o_proj_cpu(Tensor o, Tensor positions, Tensor " + "cos_sin_cache, int rope_dim) -> Tensor"); + ops.impl("inverse_gptj_rope_o_proj_cpu", torch::kCPU, + &inverse_gptj_rope_o_proj_cpu); #endif #if (defined(__AVX512BF16__) && defined(__AVX512F__) && \ diff --git a/tests/kernels/moe/test_cpu_fused_moe.py b/tests/kernels/moe/test_cpu_fused_moe.py index cc5e30c03a..dccda70c0b 100644 --- a/tests/kernels/moe/test_cpu_fused_moe.py +++ b/tests/kernels/moe/test_cpu_fused_moe.py @@ -26,8 +26,8 @@ from vllm.model_executor.layers.fused_moe.experts.cpu_moe import ( CPUUnquantizedExperts, PowerCPUUnquantizedExperts, X86CPUUnquantizedExperts, - select_experts, ) +from vllm.model_executor.layers.fused_moe.router.cpu_router import _softmax_topk from vllm.platforms import CpuArchEnum, current_platform from vllm.utils.torch_utils import set_random_seed @@ -284,7 +284,9 @@ def test_cpu_fused_moe( packed_w13 = cpu_prepack_moe_weight(w13, isa) packed_w2 = cpu_prepack_moe_weight(w2, isa) - output = cpu_fused_moe( + output = torch.empty_like(input) + cpu_fused_moe( + output, input, packed_w13, packed_w2, @@ -368,7 +370,9 @@ def test_cpu_fused_moe_int8( ) packed_w13 = cpu_prepack_moe_weight_int8(w13, isa) packed_w2 = cpu_prepack_moe_weight_int8(w2, isa) - output = cpu_fused_moe_int8( + output = torch.empty_like(input) + cpu_fused_moe_int8( + output, input, packed_w13, packed_w2, @@ -521,15 +525,10 @@ def test_cpu_fused_moe_unaligned_intermediate_size( w2_bias = torch.randn((expert_num, hidden_size), dtype=dtype) / ( 0.5 * hidden_size**0.5 ) - # Route with the same helper apply() uses internally, so the reference - # only differs from the kernel in how the experts are evaluated. - topk_weight, topk_ids = select_experts( - hidden_states=input, - router_logits=router_logits, - top_k=topk_num, - use_grouped_topk=False, - renormalize=False, - ) + # Route with the same helper the router uses internally, so the + # reference only differs from the kernel in how the experts are + # evaluated. + topk_weight, topk_ids = _softmax_topk(router_logits, topk_num, False) ref_output = ref_fused_moe( input, w13, w2, w13_bias, w2_bias, topk_weight, topk_ids, act @@ -575,15 +574,22 @@ def test_cpu_fused_moe_unaligned_intermediate_size( assert experts.isa == isa experts.process_weights_after_loading(layer) - output = experts.apply( + output = torch.empty_like(input) + experts.apply( + output=output, hidden_states=input, w1=layer.w13_weight, w2=layer.w2_weight, - router_logits=router_logits, + topk_weights=topk_weight, + topk_ids=topk_ids, activation=act, global_num_experts=expert_num, expert_map=None, a1q_scale=None, + a2_scale=None, + workspace13=torch.empty(0), + workspace2=torch.empty(0), + expert_tokens_meta=None, apply_router_weight_on_input=False, ) diff --git a/tests/kernels/moe/test_cpu_quant_fused_moe.py b/tests/kernels/moe/test_cpu_quant_fused_moe.py index e039310b93..84ff1ae266 100644 --- a/tests/kernels/moe/test_cpu_quant_fused_moe.py +++ b/tests/kernels/moe/test_cpu_quant_fused_moe.py @@ -197,14 +197,15 @@ def test_w8a16_block_fp8_cpu_fused_moe(M, N, K, E, topk, seed): pw1, pw2 = _prepack_experts(w1), _prepack_experts(w2) - # Test inplace=False against reference - out = ops.fused_experts_cpu( + # Test writing into a fresh output buffer against reference + out = torch.empty_like(a) + ops.fused_experts_cpu( + out, a.clone(), pw1, pw2, topk_weight, topk_ids, - False, ops.CPUQuantMethod.FP8_W8A16, w1_s, w2_s, @@ -215,14 +216,16 @@ def test_w8a16_block_fp8_cpu_fused_moe(M, N, K, E, topk, seed): ) torch.testing.assert_close(ref_out.bfloat16(), out, atol=1e-2, rtol=1e-2) - # Test inplace=True produces identical output - out_inplace = ops.fused_experts_cpu( - a.clone(), + # Test writing into the input tensor itself (aliased in-place) produces + # an identical result + a_inplace = a.clone() + ops.fused_experts_cpu( + a_inplace, + a_inplace, pw1, pw2, topk_weight, topk_ids, - True, ops.CPUQuantMethod.FP8_W8A16, w1_s, w2_s, @@ -231,7 +234,7 @@ def test_w8a16_block_fp8_cpu_fused_moe(M, N, K, E, topk, seed): BLOCK_SIZE, is_vnni=True, ) - torch.testing.assert_close(out_inplace, out, atol=0, rtol=0) + torch.testing.assert_close(a_inplace, out, atol=0, rtol=0) 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(): a, w1, w2, w1_s, w2_s, topk_weight, topk_ids, BLOCK_SIZE ) pw1, pw2 = _prepack_experts(w1), _prepack_experts(w2) - out = ops.fused_experts_cpu( + out = torch.empty_like(a) + ops.fused_experts_cpu( + out, a, pw1, pw2, topk_weight, topk_ids, - False, ops.CPUQuantMethod.FP8_W8A16, w1_s, w2_s, @@ -469,13 +473,15 @@ def test_mxfp4_cpu_fused_moe(M, N, K, E, topk, seed): pw2, pw2s = _prepack_mxfp4_experts(w2q, w2s) # Kernel - out = ops.fused_experts_cpu( - a.clone(), + a_in = a.clone() + out = torch.empty_like(a_in) + ops.fused_experts_cpu( + out, + a_in, pw1, pw2, topk_weight, topk_ids, - False, # inplace ops.CPUQuantMethod.MXFP4, pw1s, # w1_scale pw2s, # w2_scale @@ -511,13 +517,14 @@ def test_mxfp4_cpu_fused_moe_small_expert_blocks(): pw1, pw1s = _prepack_mxfp4_experts(w1q, w1s) pw2, pw2s = _prepack_mxfp4_experts(w2q, w2s) - out = ops.fused_experts_cpu( + out = torch.empty_like(a) + ops.fused_experts_cpu( + out, a, pw1, pw2, topk_weight, topk_ids, - False, ops.CPUQuantMethod.MXFP4, pw1s, pw2s, @@ -569,13 +576,14 @@ def test_mxfp4_cpu_zero_codes_stay_zero(zero_byte, e8m0): pw1, pw1s = _prepack_mxfp4_experts(w1q, w1s) pw2, pw2s = _prepack_mxfp4_experts(w2q, w2s) - out = ops.fused_experts_cpu( + out = torch.empty_like(a) + ops.fused_experts_cpu( + out, a, pw1, pw2, topk_weight, topk_ids, - False, ops.CPUQuantMethod.MXFP4, pw1s, pw2s, @@ -630,13 +638,14 @@ def test_mxfp4_cpu_zero_codes_mixed_with_nonzero(e8m0): pw1, pw1s = _prepack_mxfp4_experts(w1q, w1s) pw2, pw2s = _prepack_mxfp4_experts(w2q, w2s) - out = ops.fused_experts_cpu( + out = torch.empty_like(a) + ops.fused_experts_cpu( + out, a, pw1, pw2, topk_weight, topk_ids, - False, ops.CPUQuantMethod.MXFP4, pw1s, pw2s, @@ -689,13 +698,15 @@ def test_mxfp4_cpu_fused_moe_bias_swiglu(M, N, K, E, topk, seed): pw2, pw2s = _prepack_mxfp4_experts(w2q, w2s) # Kernel - out = ops.fused_experts_cpu( - a.clone(), + a_in = a.clone() + out = torch.empty_like(a_in) + ops.fused_experts_cpu( + out, + a_in, pw1, pw2, topk_weight, topk_ids, - False, # inplace ops.CPUQuantMethod.MXFP4, pw1s, # w1_scale pw2s, # w2_scale @@ -944,13 +955,15 @@ def test_int4_w4a16_cpu_fused_moe(M, N, K, E, topk, group_size, quant_algo, seed ) ) - out = ops.fused_experts_cpu( - a.clone(), + a_in = a.clone() + out = torch.empty_like(a_in) + ops.fused_experts_cpu( + out, + a_in, blocked_w1, blocked_w2, topk_weight, topk_ids, - False, # inplace ops.CPUQuantMethod.INT4_W4A8, blocked_s1, blocked_s2, @@ -1051,7 +1064,11 @@ INT8_MOE_CONFIGS = [ @pytest.mark.parametrize("is_vnni", [False, True]) @pytest.mark.parametrize("inplace", [False, True]) def test_int8_w8a8_cpu_fused_moe(M, N, K, E, topk, seed, is_vnni, inplace): - """Test fused_experts_cpu INT8 W8A8 against torch reference.""" + """Test fused_experts_cpu INT8 W8A8 against torch reference. + + ``inplace`` exercises writing into a fresh output buffer vs. aliasing + the input tensor as the output buffer. + """ set_random_seed(seed) 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): w1 = _prepack_experts(w1_q) if is_vnni else w1_q w2 = _prepack_experts(w2_q) if is_vnni else w2_q - out = ops.fused_experts_cpu( - a.clone(), + a_in = a.clone() + out = a_in if inplace else torch.empty_like(a_in) + ops.fused_experts_cpu( + out, + a_in, w1, w2, topk_weight, topk_ids, - inplace, ops.CPUQuantMethod.INT8_W8A8, w1_s, w2_s, diff --git a/tests/kernels/moe/test_zen_cpu_fused_moe.py b/tests/kernels/moe/test_zen_cpu_fused_moe.py index 5ff0f9151f..297f80ecbb 100644 --- a/tests/kernels/moe/test_zen_cpu_fused_moe.py +++ b/tests/kernels/moe/test_zen_cpu_fused_moe.py @@ -2,8 +2,6 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """Tests for ZenTorch CPU fused MoE dispatch and forward.""" -from types import SimpleNamespace - import pytest import torch @@ -27,8 +25,8 @@ from vllm.model_executor.layers.fused_moe.config import ( ) from vllm.model_executor.layers.fused_moe.experts.cpu_moe import ( CPUUnquantizedExperts, - select_experts, ) +from vllm.model_executor.layers.fused_moe.router.cpu_router import _softmax_topk from vllm.platforms import current_platform from vllm.utils.torch_utils import set_random_seed @@ -95,18 +93,19 @@ def _make_layer_and_config( @pytest.mark.parametrize("act", ZENTORCH_ACT) -def test_zen_cpu_fused_moe_dispatches_to_zentorch(act: MoEActivation): - """When zentorch MoE is supported, experts select the zentorch path.""" - layer, moe_config = _make_layer_and_config(act) - w13_before = layer.w13_weight.detach().clone() - w2_before = layer.w2_weight.detach().clone() +def test_zen_cpu_fused_moe_config_supported(act: MoEActivation): + """CPUUnquantizedExperts is selected when zentorch MoE is supported.""" + _, moe_config = _make_layer_and_config(act) + assert is_zentorch_moe_config_supported(moe_config) - experts = CPUUnquantizedExperts(moe_config, FusedMoEQuantConfig.make()) - experts.process_weights_after_loading(layer) - - assert experts._use_zentorch is True - torch.testing.assert_close(layer.w13_weight, w13_before) - torch.testing.assert_close(layer.w2_weight, w2_before) + supported, reason = CPUUnquantizedExperts.is_supported_config( + CPUUnquantizedExperts, + moe_config, + None, + None, + mk.FusedMoEActivationFormat.Standard, + ) + assert supported, reason def test_zen_cpu_fused_moe_config_supported_unaligned_intermediate(): @@ -169,26 +168,34 @@ def test_zen_cpu_fused_moe_forward( experts.process_weights_after_loading(layer) assert experts._use_zentorch is True - topk_weight, topk_ids = select_experts( - hidden_states=input, - router_logits=router_logits, - top_k=topk_num, - use_grouped_topk=False, - renormalize=False, - ) + w13_before = layer.w13_weight.detach().clone() + w2_before = layer.w2_weight.detach().clone() - output = experts.apply( + topk_weight, topk_ids = _softmax_topk(router_logits, topk_num, False) + + output = torch.empty_like(input) + experts.apply( + output=output, hidden_states=input, w1=layer.w13_weight, w2=layer.w2_weight, - router_logits=router_logits, + topk_weights=topk_weight, + topk_ids=topk_ids, activation=act, global_num_experts=expert_num, expert_map=None, a1q_scale=None, + a2_scale=None, + workspace13=torch.empty(0), + workspace2=torch.empty(0), + expert_tokens_meta=None, apply_router_weight_on_input=False, ) + # zentorch needs no prepacking; weights are consumed as loaded. + torch.testing.assert_close(layer.w13_weight, w13_before) + torch.testing.assert_close(layer.w2_weight, w2_before) + ref_output = ref_fused_moe( input, layer.w13_weight, @@ -202,20 +209,3 @@ def test_zen_cpu_fused_moe_forward( atol, rtol = get_default_atol(output), get_default_rtol(output) torch.testing.assert_close(output, ref_output, atol=atol, rtol=rtol) - - -def test_zen_cpu_fused_moe_skips_without_activation(): - """Layers without an activation attribute fall back from zentorch.""" - layer = SimpleNamespace() - moe_config = _make_moe_config( - expert_num=8, - hidden_size=128, - intermediate_size=128, - topk_num=4, - dtype=torch.bfloat16, - act=MoEActivation.SILU, - ) - experts = CPUUnquantizedExperts(moe_config, FusedMoEQuantConfig.make()) - experts.process_weights_after_loading(layer) - - assert experts._use_zentorch is False diff --git a/tests/kernels/moe/test_zen_cpu_int8_moe.py b/tests/kernels/moe/test_zen_cpu_int8_moe.py index a01bc3d1bf..cdcfe3b165 100644 --- a/tests/kernels/moe/test_zen_cpu_int8_moe.py +++ b/tests/kernels/moe/test_zen_cpu_int8_moe.py @@ -18,14 +18,12 @@ from vllm.model_executor.layers.fused_moe.config import ( FusedMoEParallelConfig, int8_w8a8_moe_quant_config, ) -from vllm.model_executor.layers.fused_moe.experts.cpu_moe import ( - ZenCPUExpertsInt8, - select_experts, -) +from vllm.model_executor.layers.fused_moe.experts.cpu_moe import ZenCPUExpertsInt8 from vllm.model_executor.layers.fused_moe.oracle.int8 import ( Int8MoeBackend, backend_to_kernel_cls, ) +from vllm.model_executor.layers.fused_moe.router.cpu_router import _softmax_topk from vllm.model_executor.layers.quantization.utils.quant_utils import ( kInt8DynamicTokenSym, kInt8StaticChannelSym, @@ -199,19 +197,27 @@ def test_zen_int8_dispatch_contract( 0.5 * HIDDEN_SIZE**0.5 ) router_logits = torch.randn((BATCH_SIZE, EXPERT_NUM), dtype=torch.bfloat16) + topk_weights, topk_ids = _softmax_topk(router_logits, topk_num, False) experts = ZenCPUExpertsInt8(moe_config, quant_config) # zentorch_fused_moe takes no expert_map. assert not experts.supports_expert_map() - output = experts.apply( + output = torch.empty_like(input) + experts.apply( + output=output, hidden_states=input, w1=w13, w2=w2, - router_logits=router_logits, + topk_weights=topk_weights, + topk_ids=topk_ids, activation=act, global_num_experts=EXPERT_NUM, expert_map=None, a1q_scale=None, + a2_scale=None, + workspace13=torch.empty(0), + workspace2=torch.empty(0), + expert_tokens_meta=None, apply_router_weight_on_input=False, ) @@ -230,14 +236,6 @@ def test_zen_int8_dispatch_contract( else: assert call["w13_bias"] is None and call["w2_bias"] is None - topk_weights, topk_ids = select_experts( - hidden_states=input, - router_logits=router_logits, - use_grouped_topk=False, - top_k=topk_num, - renormalize=False, - scoring_func="softmax", - ) ref = ref_fused_moe_int8( input, w13, diff --git a/tests/kernels/test_deepseek_v4_cpu_kernels.py b/tests/kernels/test_deepseek_v4_cpu_kernels.py new file mode 100644 index 0000000000..81c03a0bca --- /dev/null +++ b/tests/kernels/test_deepseek_v4_cpu_kernels.py @@ -0,0 +1,1494 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""CPU-kernel-vs-eager-reference parity tests for the DeepSeek-V4 CPU port's +sgl-kernels ops (csrc/cpu/sgl-kernels/). Each test carries its own reference +reimplementation of the kernel's math rather than depending on production +model code. +""" + +from unittest import mock + +import pytest +import torch + +import vllm.model_executor.kernels.mhc as mhc_kernels +from tests.kernels.test_fused_deepseek_v4_qnorm_rope_kv_insert import ( + HEAD_DIM, + NOPE_DIM, + QUANT_BLOCK, + ROPE_DIM, + make_cos_sin_cache, +) +from tests.kernels.utils import bf16_ulp_distance, fp8_ulp_distance +from vllm._custom_ops import ( + compress_norm_rope_store_cpu, + compress_norm_rope_store_indexer_cpu, + flash_mla_with_kvcache_cpu, + fp8_paged_mqa_logits_cpu, + fused_indexer_q_rope_quant_cpu, + fused_qnorm_rope_kv_insert_cpu, + inverse_gptj_rope_o_proj_cpu, + save_partial_states_cpu, +) +from vllm.models.deepseek_v4.common.ops import fused_indexer_q_rope_quant +from vllm.platforms import current_platform +from vllm.utils.torch_utils import set_random_seed + +DEVICE = current_platform.device_type + +# Main attention cache (fp8_ds_mla): NoPE fp8 in QUANT_BLOCK-wide UE8M0 blocks, +# RoPE bf16, one scale byte per block (+1 pad byte) per token. +SCALE_BYTES_PER_TOKEN = NOPE_DIM // QUANT_BLOCK + 1 +TOKEN_DATA_BYTES = NOPE_DIM + ROPE_DIM * 2 +FP8_MAX = 448.0 + +# Indexer K-cache: all 128 dims in one fp8 block, one raw fp32 scale per token. +INDEXER_HEAD_DIM = 128 +INDEXER_NOPE_DIM = INDEXER_HEAD_DIM - ROPE_DIM +INDEXER_SCALE_BYTES = 4 + + +def _op_available(*names: str) -> bool: + return all(hasattr(torch.ops._C, name) for name in names) + + +pytestmark = pytest.mark.skipif( + not current_platform.is_cpu() + or not _op_available( + "save_partial_states_cpu", + "compress_norm_rope_store_cpu", + "compress_norm_rope_store_indexer_cpu", + "flash_mla_with_kvcache_cpu", + "fused_qnorm_rope_kv_insert_cpu", + "inverse_gptj_rope_o_proj_cpu", + "fp8_paged_mqa_logits_cpu", + "fused_indexer_q_rope_quant_cpu", + ), + reason="CPU not available or DeepSeek-V4 CPU kernels not built in", +) + + +def _apply_gptj_rope(x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor): + x_even = x[..., 0::2] + x_odd = x[..., 1::2] + out_even = x_even * cos - x_odd * sin + out_odd = x_even * sin + x_odd * cos + return torch.stack((out_even, out_odd), dim=-1).flatten(-2) + + +def _make_cache(num_blocks: int, block_size: int) -> torch.Tensor: + block_bytes = block_size * TOKEN_DATA_BYTES + block_size * SCALE_BYTES_PER_TOKEN + return torch.zeros(num_blocks, block_bytes, dtype=torch.uint8) + + +def _decode_token_regions( + cache: torch.Tensor, slot_mapping: torch.Tensor, block_size: int +): + """Slice out, per valid token, the raw NoPE-fp8 / RoPE-bf16 / scale byte + regions this layout defines.""" + valid = slot_mapping >= 0 + idx = valid.nonzero(as_tuple=True)[0] + slots = slot_mapping[idx] + block_idx = torch.div(slots, block_size, rounding_mode="floor") + pos_in_block = slots % block_size + + nope, rope, scale = [], [], [] + for b, p in zip(block_idx.tolist(), pos_in_block.tolist()): + token_off = p * TOKEN_DATA_BYTES + nope.append(cache[b, token_off : token_off + NOPE_DIM]) + rope.append(cache[b, token_off + NOPE_DIM : token_off + TOKEN_DATA_BYTES]) + scale_off = block_size * TOKEN_DATA_BYTES + p * SCALE_BYTES_PER_TOKEN + scale.append(cache[b, scale_off : scale_off + SCALE_BYTES_PER_TOKEN - 1]) + + if not nope: + empty = lambda n: torch.empty(0, n, dtype=torch.uint8) + return empty(NOPE_DIM), empty(ROPE_DIM * 2), empty(SCALE_BYTES_PER_TOKEN - 1) + return torch.stack(nope), torch.stack(rope), torch.stack(scale) + + +def _assert_cache_parity(cache_fused, cache_ref, slot_mapping, block_size): + """NoPE fp8 + the UE8M0 scale byte are a deterministic quant of an + un-rotated bf16 value, so they match tightly; the RoPE region is rotated + in fp32 before its bf16 store and eager/kernel can round a tie to + opposite sides, so allow <=1 ULP there.""" + nope_ref, rope_ref, scale_ref = _decode_token_regions( + cache_ref, slot_mapping, block_size + ) + nope_fused, rope_fused, scale_fused = _decode_token_regions( + cache_fused, slot_mapping, block_size + ) + if nope_ref.numel() == 0: + return + + max_ulp = int(fp8_ulp_distance(nope_fused, nope_ref).max().item()) + assert max_ulp <= 1, f"NoPE fp8 differs by {max_ulp} ULP (>1)" + + rope_ref_bf16 = rope_ref.contiguous().view(torch.bfloat16) + rope_fused_bf16 = rope_fused.contiguous().view(torch.bfloat16) + max_ulp_rope = int(bf16_ulp_distance(rope_fused_bf16, rope_ref_bf16).max().item()) + assert max_ulp_rope <= 1, f"RoPE bf16 differs by {max_ulp_rope} ULP (>1)" + + scale_diff = (scale_fused.int() - scale_ref.int()).abs().max().item() + assert scale_diff <= 1, f"scale byte differs by {scale_diff} (>1)" + + +def _make_indexer_cache(num_blocks: int, block_size: int) -> torch.Tensor: + block_bytes = block_size * INDEXER_HEAD_DIM + block_size * INDEXER_SCALE_BYTES + return torch.zeros(num_blocks, block_bytes, dtype=torch.uint8) + + +def _assert_indexer_cache_parity(cache_fused, cache_ref, slot_mapping, block_size): + valid = slot_mapping >= 0 + idx = valid.nonzero(as_tuple=True)[0] + if idx.numel() == 0: + return + slots = slot_mapping[idx] + block_idx = torch.div(slots, block_size, rounding_mode="floor") + pos_in_block = slots % block_size + + fp8_ref, scale_ref = [], [] + fp8_fused, scale_fused = [], [] + for b, p in zip(block_idx.tolist(), pos_in_block.tolist()): + token_off = p * INDEXER_HEAD_DIM + scale_off = block_size * INDEXER_HEAD_DIM + p * INDEXER_SCALE_BYTES + fp8_ref.append(cache_ref[b, token_off : token_off + INDEXER_HEAD_DIM]) + fp8_fused.append(cache_fused[b, token_off : token_off + INDEXER_HEAD_DIM]) + scale_ref.append(cache_ref[b, scale_off : scale_off + INDEXER_SCALE_BYTES]) + scale_fused.append(cache_fused[b, scale_off : scale_off + INDEXER_SCALE_BYTES]) + + fp8_ref_t = torch.stack(fp8_ref).view(torch.float8_e4m3fn) + fp8_fused_t = torch.stack(fp8_fused).view(torch.float8_e4m3fn) + torch.testing.assert_close(fp8_fused_t.float(), fp8_ref_t.float()) + + scale_ref_t = torch.stack(scale_ref).view(torch.float32) + scale_fused_t = torch.stack(scale_fused).view(torch.float32) + torch.testing.assert_close(scale_fused_t, scale_ref_t) + + +# --------------------------------------------------------------------------- +# save_partial_states_cpu / compress_norm_rope_store_cpu (compressor.cpp) +# --------------------------------------------------------------------------- +# +# The compressor's state cache is pooled with the main KV cache's physical +# pages and page-aligned, so it is genuinely 3D with a block-to-block stride +# that need not equal block_size * width -- built here via torch.as_strided +# over an over-sized flat buffer rather than a plain contiguous tensor, so a +# regression back to flattening via .view() would be caught. + + +def _make_padded_block_cache( + num_blocks: int, block_size: int, width: int, pad_elems: int +) -> torch.Tensor: + page_stride = block_size * width + pad_elems + flat = torch.randn(num_blocks, page_stride, dtype=torch.float32) + return torch.as_strided( + flat, (num_blocks, block_size, width), (page_stride, width, 1) + ) + + +def _dup_padded_block_cache( + num_blocks: int, block_size: int, width: int, pad_elems: int +) -> tuple[torch.Tensor, torch.Tensor]: + """Two independently-stored padded caches with identical values, for ops + that mutate the cache in place.""" + page_stride = block_size * width + pad_elems + template = torch.randn(num_blocks, page_stride, dtype=torch.float32) + + def _view(flat: torch.Tensor) -> torch.Tensor: + return torch.as_strided( + flat, (num_blocks, block_size, width), (page_stride, width, 1) + ) + + return _view(template.clone()), _view(template.clone()) + + +def _run_save_partial_states_eager( + state_cache, kv, score, ape, positions, slot_mapping, compress_ratio +): + """Reference: state_cache[block_idx, pos_in_block, :W] = kv; + state_cache[..., W:] = score + ape[position % compress_ratio]; skip + slot < 0. Mutates state_cache in place, matching the op's contract.""" + state_width = kv.shape[-1] + block_size = state_cache.shape[1] + valid = slot_mapping >= 0 + idx = valid.nonzero(as_tuple=True)[0] + slots = slot_mapping[idx] + block_idx = slots // block_size + pos_in_block = slots % block_size + + state_cache[block_idx, pos_in_block, :state_width] = kv[idx] + ape_row = ape[positions[idx] % compress_ratio] + state_cache[block_idx, pos_in_block, state_width:] = score[idx] + ape_row + + +@pytest.mark.parametrize("pad_elems", [0, 24]) +@pytest.mark.parametrize("num_tokens", [1, 17]) +def test_save_partial_states_cpu_matches_eager(num_tokens: int, pad_elems: int): + torch.manual_seed(0) + compress_ratio = 4 + block_size = 4 + state_width = 512 + width = 2 * state_width + num_blocks = (num_tokens + 4 + block_size - 1) // block_size + 1 + num_slots = num_blocks * block_size + + kv = torch.randn(num_tokens, state_width, dtype=torch.float32) + score = torch.randn(num_tokens, state_width, dtype=torch.float32) + ape = torch.randn(compress_ratio, state_width, dtype=torch.float32) + positions = torch.randint(0, 4096, (num_tokens,), dtype=torch.int64) + + ref_cache, fused_cache = _dup_padded_block_cache( + num_blocks, block_size, width, pad_elems + ) + + slot_mapping = torch.randperm(num_slots)[:num_tokens].to(torch.int64) + if num_tokens > 1: + slot_mapping[-1] = -1 # exercise the "skip invalid slot" branch + + _run_save_partial_states_eager( + ref_cache, kv, score, ape, positions, slot_mapping, compress_ratio + ) + save_partial_states_cpu(kv, score, ape, positions, fused_cache, slot_mapping) + + torch.testing.assert_close(fused_cache, ref_cache) + + +def _run_compress_eager( + state_cache, + gather_slots, + positions, + kv_slot_mapping, + rms_norm_weight, + rms_norm_eps, + cos_sin_cache, + kv_cache, + kv_cache_block_size, + compress_ratio, +): + """Reference for the per-channel-softmax-pooling + RMSNorm + NoPE-UE8M0 + quant + GPT-J RoPE math writing the main (head_dim=512) fp8_ds_mla + layout.""" + num_tokens, window = gather_slots.shape + state_width = state_cache.shape[-1] // 2 + state_block_size = state_cache.shape[1] + num_quant_blocks = NOPE_DIM // QUANT_BLOCK + + w_idx = torch.arange(window) + head_offset = torch.where(w_idx >= compress_ratio, HEAD_DIM, 0) + + for t in range(num_tokens): + position = int(positions[t]) + if (position + 1) % compress_ratio != 0: + continue + kv_slot = int(kv_slot_mapping[t]) + if kv_slot < 0: + continue + + slots = gather_slots[t] + valid_w = slots >= 0 + safe_slots = slots.clamp(min=0) + block_idx = safe_slots // state_block_size + pos_in_block = safe_slots % state_block_size + + rows = state_cache[block_idx, pos_in_block] # [window, 2*state_width] + col = head_offset.unsqueeze(-1) + torch.arange(HEAD_DIM) + kv_vals = torch.gather(rows, 1, col) + score_col = state_width + col + score_vals = torch.gather(rows, 1, score_col) + score_vals = torch.where( + valid_w.unsqueeze(-1), score_vals, torch.full_like(score_vals, -torch.inf) + ) + + weights = torch.softmax(score_vals, dim=0) + compressed = (kv_vals * weights).sum(dim=0) # [HEAD_DIM] + + variance = compressed.pow(2).mean() + normed = compressed * torch.rsqrt(variance + rms_norm_eps) * rms_norm_weight + + quant_input = normed.to(torch.bfloat16).to(torch.float32) + + compressed_pos = (position // compress_ratio) * compress_ratio + cos = cos_sin_cache[compressed_pos, : ROPE_DIM // 2] + sin = cos_sin_cache[compressed_pos, ROPE_DIM // 2 :] + rotated_rope = _apply_gptj_rope(normed[NOPE_DIM:].clone(), cos, sin) + + nope_vals = quant_input[:NOPE_DIM].view(num_quant_blocks, QUANT_BLOCK) + absmax = nope_vals.abs().amax(dim=-1).clamp(min=1e-4) + exponent = torch.ceil(torch.log2(absmax / FP8_MAX)) + inv_scale = torch.exp2(-exponent) + scaled = (nope_vals * inv_scale.unsqueeze(-1)).clamp(-FP8_MAX, FP8_MAX) + nope_fp8_bytes = scaled.to(torch.float8_e4m3fn).reshape(-1).view(torch.uint8) + scale_bytes = (exponent + 127.0).clamp(0, 255).to(torch.uint8) + rope_bf16_bytes = rotated_rope.to(torch.bfloat16).contiguous().view(torch.uint8) + + block_idx_kv = kv_slot // kv_cache_block_size + pos_in_block_kv = kv_slot % kv_cache_block_size + token_off = pos_in_block_kv * TOKEN_DATA_BYTES + kv_cache[block_idx_kv, token_off : token_off + NOPE_DIM] = nope_fp8_bytes + kv_cache[block_idx_kv, token_off + NOPE_DIM : token_off + TOKEN_DATA_BYTES] = ( + rope_bf16_bytes + ) + scale_off = ( + kv_cache_block_size * TOKEN_DATA_BYTES + + pos_in_block_kv * SCALE_BYTES_PER_TOKEN + ) + kv_cache[block_idx_kv, scale_off : scale_off + num_quant_blocks] = scale_bytes + kv_cache[block_idx_kv, scale_off + num_quant_blocks] = 0 + + +def _make_compress_inputs( + num_tokens: int, + compress_ratio: int, + overlap: bool, + pad_elems: int = 24, + max_pos: int = 4096, + invalid_window_frac: float = 0.1, + head_dim: int = HEAD_DIM, +): + torch.manual_seed(0) + coff = 1 + int(overlap) + state_width = coff * head_dim + window = coff * compress_ratio + state_block_size = 4 if compress_ratio == 4 else 8 + num_state_slots = num_tokens * window + 8 + num_state_blocks = (num_state_slots + state_block_size - 1) // state_block_size + 1 + + state_cache = _make_padded_block_cache( + num_state_blocks, state_block_size, 2 * state_width, pad_elems + ) + gather_slots = torch.randint( + 0, num_state_blocks * state_block_size, (num_tokens, window) + ).to(torch.int64) + invalid_mask = torch.rand(num_tokens, window) < invalid_window_frac + # Never invalidate every column of a row (would force max=-inf -> NaN). + invalid_mask[:, -1] = False + gather_slots = torch.where( + invalid_mask, torch.full_like(gather_slots, -1), gather_slots + ) + + base = (torch.arange(num_tokens) + 1) * compress_ratio + is_boundary = torch.arange(num_tokens) % 2 == 0 + positions = torch.where(is_boundary, base - 1, base + 1).to(torch.int64) + + kv_slot_mapping = torch.arange(num_tokens, dtype=torch.int64) + if num_tokens > 1: + kv_slot_mapping[1] = -1 # exercise the "skip invalid output slot" branch + + rms_norm_weight = torch.randn(head_dim, dtype=torch.float32) + cos_sin_cache = make_cos_sin_cache(max_pos, ROPE_DIM, torch.float32, "cpu") + + return ( + state_cache, + gather_slots, + positions, + kv_slot_mapping, + rms_norm_weight, + cos_sin_cache, + is_boundary, + ) + + +@pytest.mark.parametrize("compress_ratio,overlap", [(128, False), (4, True)]) +@pytest.mark.parametrize("num_tokens", [1, 17]) +def test_compress_norm_rope_store_cpu_matches_eager( + compress_ratio: int, overlap: bool, num_tokens: int +): + eps = 1e-6 + block_size = 16 + ( + state_cache, + gather_slots, + positions, + kv_slot_mapping, + rms_norm_weight, + cos_sin_cache, + is_boundary, + ) = _make_compress_inputs(num_tokens, compress_ratio, overlap) + + num_blocks = (num_tokens + block_size - 1) // block_size + 1 + kv_cache_ref = _make_cache(num_blocks, block_size) + kv_cache_fused = kv_cache_ref.clone() + + _run_compress_eager( + state_cache, + gather_slots, + positions, + kv_slot_mapping, + rms_norm_weight, + eps, + cos_sin_cache, + kv_cache_ref, + block_size, + compress_ratio, + ) + compress_norm_rope_store_cpu( + state_cache, + gather_slots, + positions, + kv_slot_mapping, + rms_norm_weight, + eps, + cos_sin_cache, + kv_cache_fused, + block_size, + compress_ratio, + ) + + # Only tokens that hit the compression boundary (and have a valid output + # slot) were written -- restrict parity checks to those. + written_slot_mapping = torch.where( + is_boundary, kv_slot_mapping, torch.full_like(kv_slot_mapping, -1) + ) + _assert_cache_parity(kv_cache_fused, kv_cache_ref, written_slot_mapping, block_size) + + +def test_compress_norm_rope_store_cpu_non_contiguous_kv_cache_row_stride(): + """The main (output) kv_cache is a separate physical buffer from the + state cache and can independently be a strided view.""" + eps = 1e-6 + compress_ratio, overlap = 128, False + num_tokens = 5 + block_size = 16 + cache_pad = 32 + + ( + state_cache, + gather_slots, + positions, + kv_slot_mapping, + rms_norm_weight, + cos_sin_cache, + is_boundary, + ) = _make_compress_inputs(num_tokens, compress_ratio, overlap) + + num_blocks = (num_tokens + block_size - 1) // block_size + 1 + kv_cache_ref = _make_cache(num_blocks, block_size) + block_bytes = kv_cache_ref.shape[1] + padded_cache_ref = torch.zeros( + num_blocks, block_bytes + cache_pad, dtype=torch.uint8 + ) + kv_cache_ref_view = padded_cache_ref[:, :block_bytes] + kv_cache_ref_view.copy_(kv_cache_ref) + padded_cache_fused = padded_cache_ref.clone() + kv_cache_fused_view = padded_cache_fused[:, :block_bytes] + + _run_compress_eager( + state_cache, + gather_slots, + positions, + kv_slot_mapping, + rms_norm_weight, + eps, + cos_sin_cache, + kv_cache_ref_view, + block_size, + compress_ratio, + ) + compress_norm_rope_store_cpu( + state_cache, + gather_slots, + positions, + kv_slot_mapping, + rms_norm_weight, + eps, + cos_sin_cache, + kv_cache_fused_view, + block_size, + compress_ratio, + ) + + written_slot_mapping = torch.where( + is_boundary, kv_slot_mapping, torch.full_like(kv_slot_mapping, -1) + ) + _assert_cache_parity( + kv_cache_fused_view, kv_cache_ref_view, written_slot_mapping, block_size + ) + + +def _run_compress_indexer_eager( + state_cache, + gather_slots, + positions, + kv_slot_mapping, + rms_norm_weight, + rms_norm_eps, + cos_sin_cache, + kv_cache, + kv_cache_block_size, + compress_ratio, +): + """Same softmax-pool + RMSNorm + GPT-J RoPE prologue as the head=512 + path, but the indexer's cache quantizes all 128 post-RoPE dims as one + FP8 block with a single raw fp32 scale per token (matching + _indexer_k_quant_and_cache_cpu's layout), not 7 separate UE8M0 blocks.""" + num_tokens, window = gather_slots.shape + state_width = state_cache.shape[-1] // 2 + state_block_size = state_cache.shape[1] + + w_idx = torch.arange(window) + head_offset = torch.where(w_idx >= compress_ratio, INDEXER_HEAD_DIM, 0) + + for t in range(num_tokens): + position = int(positions[t]) + if (position + 1) % compress_ratio != 0: + continue + kv_slot = int(kv_slot_mapping[t]) + if kv_slot < 0: + continue + + slots = gather_slots[t] + valid_w = slots >= 0 + safe_slots = slots.clamp(min=0) + block_idx = safe_slots // state_block_size + pos_in_block = safe_slots % state_block_size + + rows = state_cache[block_idx, pos_in_block] # [window, 2*state_width] + col = head_offset.unsqueeze(-1) + torch.arange(INDEXER_HEAD_DIM) + kv_vals = torch.gather(rows, 1, col) + score_col = state_width + col + score_vals = torch.gather(rows, 1, score_col) + score_vals = torch.where( + valid_w.unsqueeze(-1), score_vals, torch.full_like(score_vals, -torch.inf) + ) + + weights = torch.softmax(score_vals, dim=0) + compressed = (kv_vals * weights).sum(dim=0) # [INDEXER_HEAD_DIM] + + variance = compressed.pow(2).mean() + normed = compressed * torch.rsqrt(variance + rms_norm_eps) * rms_norm_weight + + compressed_pos = (position // compress_ratio) * compress_ratio + cos = cos_sin_cache[compressed_pos, : ROPE_DIM // 2] + sin = cos_sin_cache[compressed_pos, ROPE_DIM // 2 :] + rotated_rope = _apply_gptj_rope(normed[INDEXER_NOPE_DIM:].clone(), cos, sin) + full = torch.cat((normed[:INDEXER_NOPE_DIM], rotated_rope)) + + quant_input = full.to(torch.bfloat16).to(torch.float32) + absmax = quant_input.abs().amax().clamp(min=1e-4) + exponent = torch.ceil(torch.log2(absmax / FP8_MAX)) + scale = torch.exp2(exponent) + scaled = (quant_input / scale).clamp(-FP8_MAX, FP8_MAX) + fp8_bytes = scaled.to(torch.float8_e4m3fn).view(torch.uint8) + + block_idx_kv = kv_slot // kv_cache_block_size + pos_in_block_kv = kv_slot % kv_cache_block_size + token_off = pos_in_block_kv * INDEXER_HEAD_DIM + kv_cache[block_idx_kv, token_off : token_off + INDEXER_HEAD_DIM] = fp8_bytes + scale_off = ( + kv_cache_block_size * INDEXER_HEAD_DIM + + pos_in_block_kv * INDEXER_SCALE_BYTES + ) + kv_cache[block_idx_kv, scale_off : scale_off + INDEXER_SCALE_BYTES] = ( + scale.reshape(1).view(torch.uint8) + ) + + +@pytest.mark.parametrize("compress_ratio,overlap", [(128, False), (4, True)]) +@pytest.mark.parametrize("num_tokens", [1, 17]) +def test_compress_norm_rope_store_indexer_cpu_matches_eager( + compress_ratio: int, overlap: bool, num_tokens: int +): + eps = 1e-6 + block_size = 16 + ( + state_cache, + gather_slots, + positions, + kv_slot_mapping, + rms_norm_weight, + cos_sin_cache, + is_boundary, + ) = _make_compress_inputs( + num_tokens, compress_ratio, overlap, head_dim=INDEXER_HEAD_DIM + ) + + num_blocks = (num_tokens + block_size - 1) // block_size + 1 + kv_cache_ref = _make_indexer_cache(num_blocks, block_size) + kv_cache_fused = kv_cache_ref.clone() + + _run_compress_indexer_eager( + state_cache, + gather_slots, + positions, + kv_slot_mapping, + rms_norm_weight, + eps, + cos_sin_cache, + kv_cache_ref, + block_size, + compress_ratio, + ) + compress_norm_rope_store_indexer_cpu( + state_cache, + gather_slots, + positions, + kv_slot_mapping, + rms_norm_weight, + eps, + cos_sin_cache, + kv_cache_fused, + block_size, + compress_ratio, + ) + + written_slot_mapping = torch.where( + is_boundary, kv_slot_mapping, torch.full_like(kv_slot_mapping, -1) + ) + _assert_indexer_cache_parity( + kv_cache_fused, kv_cache_ref, written_slot_mapping, block_size + ) + + +# --------------------------------------------------------------------------- +# flash_mla_with_kvcache_cpu (flash_mla.cpp): fused sparse MQA decode attn +# --------------------------------------------------------------------------- + + +def _pack_flash_mla_cache(kv_f32: torch.Tensor, block_size: int): + """Quantize kv_f32 ([N, HEAD_DIM]) into the fp8_ds_mla byte layout and + return (cache, dequant_ref) -- the exact post-quantization value the + kernel's dequant should reproduce. Slot id == row index, packed + contiguously from block 0.""" + n = kv_f32.shape[0] + num_blocks = max(1, -(-n // block_size)) + cache = _make_cache(num_blocks, block_size) + if n == 0: + return cache, kv_f32.new_zeros(0, HEAD_DIM) + + num_quant_blocks = NOPE_DIM // QUANT_BLOCK + nope = kv_f32[:, :NOPE_DIM].reshape(n, num_quant_blocks, QUANT_BLOCK) + absmax = nope.abs().amax(dim=-1).clamp(min=1e-4) + exponent = torch.ceil(torch.log2(absmax / FP8_MAX)) + inv_scale = torch.exp2(-exponent) + scaled = (nope * inv_scale.unsqueeze(-1)).clamp(-FP8_MAX, FP8_MAX) + nope_fp8 = scaled.to(torch.float8_e4m3fn) + scale = torch.exp2(exponent) + nope_dequant = (nope_fp8.float() * scale.unsqueeze(-1)).reshape(n, NOPE_DIM) + + rope_bf16 = kv_f32[:, NOPE_DIM:].to(torch.bfloat16) + dequant_ref = torch.cat([nope_dequant, rope_bf16.float()], dim=-1) + + slot_ids = torch.arange(n) + block_idx = slot_ids // block_size + pos_in_block = slot_ids % block_size + + nope_fp8_bytes = nope_fp8.reshape(n, NOPE_DIM).contiguous().view(torch.uint8) + rope_bf16_bytes = rope_bf16.contiguous().view(torch.uint8) + scale_bytes = (exponent + 127.0).clamp(0, 255).to(torch.uint8) + + arange_nope = torch.arange(NOPE_DIM) + arange_rope = torch.arange(ROPE_DIM * 2) + arange_scale = torch.arange(num_quant_blocks) + token_base = pos_in_block.unsqueeze(-1) * TOKEN_DATA_BYTES + row_idx = block_idx.unsqueeze(-1) + + cache[row_idx, token_base + arange_nope] = nope_fp8_bytes + cache[row_idx, token_base + NOPE_DIM + arange_rope] = rope_bf16_bytes + scale_base = ( + block_size * TOKEN_DATA_BYTES + + pos_in_block.unsqueeze(-1) * SCALE_BYTES_PER_TOKEN + ) + cache[row_idx, scale_base + arange_scale] = scale_bytes + cache[block_idx, scale_base.squeeze(-1) + num_quant_blocks] = 0 + + return cache, dequant_ref + + +def _flash_mla_eager( + q: torch.Tensor, + window_dequant: torch.Tensor, + window_slots: torch.Tensor, + compressed_dequant: torch.Tensor, + compressed_slots: torch.Tensor, + attn_sink: torch.Tensor, + scale: float, +) -> torch.Tensor: + """MQA attention: gather the (already-dequantized) window+compressed KV + rows referenced by *_slots (-1 = invalid/masked), append a learned + per-head sink logit with zero value contribution, softmax, weighted + sum.""" + num_tokens, num_heads, _ = q.shape + q_f32 = q.float() + + def gather(dequant, slots): + valid = slots >= 0 + safe = slots.clamp(min=0) + kv = dequant[safe] # [T, K, D] + return kv, valid + + win_kv, win_valid = gather(window_dequant, window_slots) + if compressed_slots.shape[1] > 0: + comp_kv, comp_valid = gather(compressed_dequant, compressed_slots) + kv = torch.cat([win_kv, comp_kv], dim=1) + valid = torch.cat([win_valid, comp_valid], dim=1) + else: + kv = win_kv + valid = win_valid + + scores = torch.einsum("thd,tkd->thk", q_f32, kv) * scale + scores = scores.masked_fill(~valid.unsqueeze(1), float("-inf")) + sink = attn_sink.view(1, num_heads, 1).expand(num_tokens, num_heads, 1) + scores_full = torch.cat([scores, sink], dim=-1) + probs = torch.softmax(scores_full, dim=-1) + out = torch.einsum("thk,tkd->thd", probs[..., :-1], kv) + return out.to(q.dtype) + + +@pytest.mark.parametrize("num_heads", [1, 64]) +@pytest.mark.parametrize("num_tokens", [1, 17]) +@pytest.mark.parametrize("swa_only", [True, False]) +def test_flash_mla_cpu_kernel_matches_eager( + num_heads: int, num_tokens: int, swa_only: bool +): + torch.manual_seed(0) + dtype = torch.bfloat16 + window_block_size = 16 + compressed_block_size = 8 + num_window = 6 + num_compressed = 0 if swa_only else 5 + + q = torch.randn(num_tokens, num_heads, HEAD_DIM, dtype=dtype) + scale = 1.0 / (HEAD_DIM**0.5) + attn_sink = torch.randn(num_heads, dtype=torch.float32) + + num_window_slots = num_tokens * num_window + window_block_size + window_kv_f32 = torch.randn(num_window_slots, HEAD_DIM) + window_cache, window_dequant = _pack_flash_mla_cache( + window_kv_f32, window_block_size + ) + window_slots = torch.randint(0, num_window_slots, (num_tokens, num_window)) + window_slots[:, 0] = -1 # exercise invalid-slot masking + + if num_compressed > 0: + num_compressed_slots = num_tokens * num_compressed + compressed_block_size + compressed_kv_f32 = torch.randn(num_compressed_slots, HEAD_DIM) + compressed_cache, compressed_dequant = _pack_flash_mla_cache( + compressed_kv_f32, compressed_block_size + ) + compressed_slots = torch.randint( + 0, num_compressed_slots, (num_tokens, num_compressed) + ) + if num_compressed > 1: + compressed_slots[:, -1] = -1 + else: + compressed_cache = window_cache + compressed_dequant = window_dequant.new_zeros(0, HEAD_DIM) + compressed_slots = window_slots.new_full((num_tokens, 0), -1) + + out_ref = _flash_mla_eager( + q, + window_dequant, + window_slots, + compressed_dequant, + compressed_slots, + attn_sink, + scale, + ) + out_kernel = q.new_empty(num_tokens, num_heads, HEAD_DIM) + flash_mla_with_kvcache_cpu( + out_kernel, + q, + window_cache, + window_slots, + window_block_size, + compressed_cache, + compressed_slots, + compressed_block_size, + attn_sink, + scale, + ) + + torch.testing.assert_close( + out_kernel.float(), out_ref.float(), rtol=2e-2, atol=2e-2 + ) + + +def test_flash_mla_cpu_kernel_all_invalid_falls_back_to_sink(): + """When every KV slot is masked out, softmax degenerates to the sink + logit alone and the output must be all-zero (no valid value rows).""" + torch.manual_seed(2) + dtype = torch.bfloat16 + num_tokens, num_heads, num_window = 3, 4, 5 + window_block_size = 16 + + q = torch.randn(num_tokens, num_heads, HEAD_DIM, dtype=dtype) + scale = 1.0 / (HEAD_DIM**0.5) + attn_sink = torch.randn(num_heads, dtype=torch.float32) + + window_kv_f32 = torch.randn(window_block_size, HEAD_DIM) + window_cache, _ = _pack_flash_mla_cache(window_kv_f32, window_block_size) + window_slots = torch.full((num_tokens, num_window), -1, dtype=torch.int64) + compressed_slots = window_slots.new_full((num_tokens, 0), -1) + + out = q.new_empty(num_tokens, num_heads, HEAD_DIM) + flash_mla_with_kvcache_cpu( + out, + q, + window_cache, + window_slots, + window_block_size, + window_cache, + compressed_slots, + 1, + attn_sink, + scale, + ) + assert out.abs().max().item() == 0.0 + + +# --------------------------------------------------------------------------- +# fused_qnorm_rope_kv_insert_cpu (store_cache.cpp) +# --------------------------------------------------------------------------- + + +def _qnorm_kv_insert_eager( + q, kv, positions, cache, slot_mapping, cos_sin_cache, q_head_padded, eps, bs +): + """Reference: per-head weight-free RMSNorm + GPT-J RoPE on Q; GPT-J RoPE + + UE8M0 FP8 quant + paged insert on KV.""" + num_quant_blocks = NOPE_DIM // QUANT_BLOCK + + num_tokens_full, num_heads_q, _ = q.shape + out_dtype = q.dtype + + cos = cos_sin_cache[positions, : ROPE_DIM // 2] + sin = cos_sin_cache[positions, ROPE_DIM // 2 :] + + q_f32 = q.float() + var = q_f32.square().mean(dim=-1, keepdim=True) + q_f32 = q_f32 * torch.rsqrt(var + eps) + q_f32[..., NOPE_DIM:] = _apply_gptj_rope( + q_f32[..., NOPE_DIM:], cos.unsqueeze(1), sin.unsqueeze(1) + ) + q_out = q.new_zeros(num_tokens_full, q_head_padded, HEAD_DIM) + q_out[:, :num_heads_q, :] = q_f32.to(out_dtype) + + num_tokens_insert = slot_mapping.shape[0] + kv_f32 = kv[:num_tokens_insert].float() + kv_f32[..., NOPE_DIM:] = _apply_gptj_rope( + kv_f32[..., NOPE_DIM:], cos[:num_tokens_insert], sin[:num_tokens_insert] + ) + kv_bf16 = kv_f32.to(torch.bfloat16) + kv_f32 = kv_bf16.float() + + valid = slot_mapping >= 0 + if bool(valid.any()): + idx = valid.nonzero(as_tuple=True)[0] + slot_id = slot_mapping[idx] + block_idx = torch.div(slot_id, bs, rounding_mode="floor") + pos_in_block = slot_id % bs + + nope = kv_f32[idx, :NOPE_DIM].view(-1, num_quant_blocks, QUANT_BLOCK) + absmax = nope.abs().amax(dim=-1).clamp(min=1e-4) + exponent = torch.ceil(torch.log2(absmax / FP8_MAX)) + inv_scale = torch.exp2(-exponent) + scaled = (nope * inv_scale.unsqueeze(-1)).clamp(-FP8_MAX, FP8_MAX) + # Flush fp8-subnormal magnitudes to zero, matching the kernel's + # AVX512 cvt_fp32x16_to_fp8x16, which doesn't preserve exact + # subnormal encodings the way .to(torch.float8_e4m3fn) does. + fp8_min_normal = 2.0**-6 + scaled = torch.where(scaled.abs() < fp8_min_normal, 0.0, scaled) + nope_fp8_bytes = ( + scaled.to(torch.float8_e4m3fn).view(-1, NOPE_DIM).view(torch.uint8) + ) + scale_bytes = (exponent + 127.0).clamp(0, 255).to(torch.uint8) + rope_bf16_bytes = kv_bf16[idx, NOPE_DIM:].contiguous().view(torch.uint8) + + arange_nope = torch.arange(NOPE_DIM) + arange_rope = torch.arange(ROPE_DIM * 2) + arange_scale = torch.arange(num_quant_blocks) + row_idx = block_idx.unsqueeze(-1) + token_base = pos_in_block.unsqueeze(-1) * TOKEN_DATA_BYTES + + cache[row_idx, token_base + arange_nope] = nope_fp8_bytes + cache[row_idx, token_base + NOPE_DIM + arange_rope] = rope_bf16_bytes + scale_base = ( + bs * TOKEN_DATA_BYTES + pos_in_block.unsqueeze(-1) * SCALE_BYTES_PER_TOKEN + ) + cache[row_idx, scale_base + arange_scale] = scale_bytes + cache[block_idx, scale_base.squeeze(-1) + num_quant_blocks] = 0 + + return q_out + + +@pytest.mark.parametrize("num_tokens", [1, 64]) +@pytest.mark.parametrize("n_heads,padded_heads", [(1, 8), (8, 8), (8, 16), (64, 128)]) +def test_qnorm_rope_kv_insert_cpu_kernel_matches_eager( + num_tokens: int, n_heads: int, padded_heads: int +): + torch.manual_seed(0) + dtype = torch.bfloat16 + eps = 1e-6 + max_pos = 4096 + block_size = 16 + + q = torch.randn(num_tokens, n_heads, HEAD_DIM, dtype=dtype) + kv = torch.randn(num_tokens, HEAD_DIM, dtype=dtype) + positions = torch.randint(0, max_pos, (num_tokens,), dtype=torch.int64) + cos_sin_cache = make_cos_sin_cache(max_pos, ROPE_DIM, torch.float32, "cpu") + + num_blocks = (num_tokens + block_size - 1) // block_size + 1 + slot_mapping = torch.arange(num_tokens, dtype=torch.int64) + if num_tokens > 1: + slot_mapping[-1] = -1 # exercise the "skip invalid slot" branch + + cache_ref = _make_cache(num_blocks, block_size) + cache_fused = cache_ref.clone() + + q_ref = _qnorm_kv_insert_eager( + q, + kv, + positions, + cache_ref, + slot_mapping, + cos_sin_cache, + padded_heads, + eps, + block_size, + ) + q_fused = fused_qnorm_rope_kv_insert_cpu( + q, + kv, + positions, + cache_fused, + slot_mapping, + cos_sin_cache, + padded_heads, + eps, + block_size, + ) + + torch.testing.assert_close(q_fused.float(), q_ref.float(), rtol=1e-2, atol=1e-2) + if n_heads < padded_heads: + pad_region = q_fused[:, n_heads:padded_heads] + assert pad_region.abs().max().item() == 0.0, ( + "padded head slots must be exact zero" + ) + + _assert_cache_parity(cache_fused, cache_ref, slot_mapping, block_size) + + +@pytest.mark.parametrize("pad", [1, 5]) +def test_qnorm_rope_kv_insert_cpu_kernel_dp_padding(pad: int): + """slot_mapping shorter than q/kv rows: the KV branch must only touch + the first num_tokens rows while Q-norm+RoPE still runs on all rows.""" + torch.manual_seed(1) + dtype = torch.bfloat16 + eps = 1e-6 + max_pos = 4096 + block_size = 16 + num_tokens = 17 + total = num_tokens + pad + padded_heads = 8 + + q = torch.randn(total, padded_heads, HEAD_DIM, dtype=dtype) + kv = torch.randn(total, HEAD_DIM, dtype=dtype) + positions = torch.randint(0, max_pos, (total,), dtype=torch.int64) + cos_sin_cache = make_cos_sin_cache(max_pos, ROPE_DIM, torch.float32, "cpu") + + num_blocks = (num_tokens + block_size - 1) // block_size + 1 + slot_mapping = torch.arange(num_tokens, dtype=torch.int64) + + cache_ref = _make_cache(num_blocks, block_size) + cache_fused = cache_ref.clone() + + q_ref = _qnorm_kv_insert_eager( + q, + kv, + positions, + cache_ref, + slot_mapping, + cos_sin_cache, + padded_heads, + eps, + block_size, + ) + q_fused = fused_qnorm_rope_kv_insert_cpu( + q, + kv, + positions, + cache_fused, + slot_mapping, + cos_sin_cache, + padded_heads, + eps, + block_size, + ) + + torch.testing.assert_close(q_fused.float(), q_ref.float(), rtol=1e-2, atol=1e-2) + _assert_cache_parity(cache_fused, cache_ref, slot_mapping, block_size) + + +# --------------------------------------------------------------------------- +# inverse_gptj_rope_o_proj_cpu (store_cache.cpp): DeepseekV4CPUAttention's +# _o_proj de-rotation kernel +# --------------------------------------------------------------------------- + + +def _inverse_gptj_rope_o_proj_eager( + o: torch.Tensor, + positions: torch.Tensor, + cos_sin_cache: torch.Tensor, + rope_dim: int, +) -> torch.Tensor: + """NoPE dims pass through as fp32; the RoPE segment is de-rotated by + R(-pos) in GPT-J (interleaved even/odd) convention.""" + nope_dim = o.shape[-1] - rope_dim + half_dim = rope_dim // 2 + + o_f32 = o.float() + out = torch.empty_like(o_f32) + out[..., :nope_dim] = o_f32[..., :nope_dim] + + cos = cos_sin_cache[positions, :half_dim].unsqueeze(1) + sin = cos_sin_cache[positions, half_dim:].unsqueeze(1) + + e = o_f32[..., nope_dim::2] + od = o_f32[..., nope_dim + 1 :: 2] + out[..., nope_dim::2] = e * cos + od * sin + out[..., nope_dim + 1 :: 2] = od * cos - e * sin + return out + + +@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16]) +@pytest.mark.parametrize("num_tokens", [1, 17]) +@pytest.mark.parametrize( + "head_dim,rope_dim,num_heads", + [ + (512, 64, 8), # DeepseekV4CPUAttention._o_proj's actual shape + (64, 64, 1), # nope_dim=0: no NoPE region at all + ], +) +def test_inverse_gptj_rope_o_proj_cpu_kernel_matches_eager( + dtype: torch.dtype, + num_tokens: int, + head_dim: int, + rope_dim: int, + num_heads: int, +): + torch.manual_seed(0) + max_pos = 4096 + + o = torch.randn(num_tokens, num_heads, head_dim, dtype=dtype) + positions = torch.randint(0, max_pos, (num_tokens,), dtype=torch.int64) + cos_sin_cache = make_cos_sin_cache(max_pos, rope_dim, torch.float32, "cpu") + + out_ref = _inverse_gptj_rope_o_proj_eager(o, positions, cos_sin_cache, rope_dim) + out_fused = inverse_gptj_rope_o_proj_cpu(o, positions, cos_sin_cache, rope_dim) + + torch.testing.assert_close(out_fused, out_ref, rtol=1e-2, atol=1e-2) + + +# --------------------------------------------------------------------------- +# fp8_paged_mqa_logits_cpu (paged_mqa_logits.cpp): sparse indexer's +# decode-path MQA-logits kernel +# --------------------------------------------------------------------------- + + +def _rand_fp8(*shape: int) -> torch.Tensor: + return torch.randn(shape, dtype=torch.float32).clamp(-3, 3).to(torch.float8_e4m3fn) + + +def _pages_per_request(block_size: int, seq_lens: list[int]) -> list[int]: + return [-(-seq_len // block_size) for seq_len in seq_lens] + + +def _write_page( + page_row: torch.Tensor, block_size: int, k_fp8: torch.Tensor, scale: torch.Tensor +) -> None: + """Write k_fp8 ([block_size, INDEXER_HEAD_DIM]) and scale ([block_size] + fp32) into one page row: K-region then scale-region, matching + compress_norm_rope_store_indexer_cpu's layout.""" + k_region = page_row[: block_size * INDEXER_HEAD_DIM] + k_region.copy_(k_fp8.view(torch.uint8).flatten()) + scale_region = page_row[block_size * INDEXER_HEAD_DIM :].view(torch.float32) + scale_region[:].copy_(scale) + + +def _reference_paged_mqa_logits( + q_fp8: torch.Tensor, + weight: torch.Tensor, + ref_k: list[torch.Tensor], + ref_scale: list[torch.Tensor], + max_seq_len: int, +) -> torch.Tensor: + batch_size, _num_heads, _ = q_fp8.shape + out = torch.zeros((batch_size, max_seq_len), dtype=torch.float32) + for b in range(batch_size): + k = ref_k[b].float() # [seq_len, head_dim] + scale = ref_scale[b] # [seq_len] + q = q_fp8[b].float() # [heads, head_dim] + w = weight[b] # [heads] + dot = torch.einsum("jd,hd->jh", k, q).clamp(min=0.0) + score = (dot * w.unsqueeze(0)).sum(dim=1) + out[b, : k.shape[0]] = score * scale + return out + + +def _build_paged_mqa_logits_batch( + block_size: int, + seq_lens: list[int], + num_heads: int, + pages: torch.Tensor, + seed: int, +): + """Fill pages (a pre-sized [n_pages, buf_width] uint8 tensor, possibly + non-contiguous in dim 0) with deterministic-per-seed random K/scale + data.""" + torch.manual_seed(seed) + batch_size = len(seq_lens) + max_seq_len = max(seq_lens) + per_request = _pages_per_request(block_size, seq_lens) + pages_per_batch = max(per_request) + page_table = torch.full((batch_size, pages_per_batch), -1, dtype=torch.int32) + + ref_k: list[torch.Tensor] = [] + ref_scale: list[torch.Tensor] = [] + next_page = 0 + for b, seq_len in enumerate(seq_lens): + k_chunks = [] + scale_chunks = [] + for p in range(per_request[b]): + page_table[b, p] = next_page + k = _rand_fp8(block_size, INDEXER_HEAD_DIM) + scale = torch.rand(block_size, dtype=torch.float32) * 0.05 + 0.01 + _write_page(pages[next_page], block_size, k, scale) + tokens_here = min(block_size, seq_len - p * block_size) + k_chunks.append(k[:tokens_here]) + scale_chunks.append(scale[:tokens_here]) + next_page += 1 + ref_k.append(torch.cat(k_chunks, dim=0)) + ref_scale.append(torch.cat(scale_chunks, dim=0)) + + q_fp8 = _rand_fp8(batch_size, num_heads, INDEXER_HEAD_DIM) + weight = torch.rand(batch_size, num_heads, dtype=torch.float32) * 0.1 + seq_lens_t = torch.tensor(seq_lens, dtype=torch.int32) + return q_fp8, weight, seq_lens_t, page_table, ref_k, ref_scale, max_seq_len + + +def _paged_mqa_logits_page_layout( + block_size: int, seq_lens: list[int] +) -> tuple[int, int]: + n_pages = sum(_pages_per_request(block_size, seq_lens)) + buf_width = block_size * INDEXER_HEAD_DIM + block_size * INDEXER_SCALE_BYTES + return n_pages, buf_width + + +@pytest.mark.parametrize("block_size", [64, 256]) +def test_fp8_paged_mqa_logits_cpu_matches_eager(block_size: int): + num_heads = 4 + seq_lens = [1, block_size + 7, 3 * block_size] + n_pages, buf_width = _paged_mqa_logits_page_layout(block_size, seq_lens) + pages = torch.zeros((n_pages, buf_width), dtype=torch.uint8) + + q_fp8, weight, seq_lens_t, page_table, ref_k, ref_scale, max_seq_len = ( + _build_paged_mqa_logits_batch(block_size, seq_lens, num_heads, pages, seed=0) + ) + + logits = fp8_paged_mqa_logits_cpu( + q_fp8, pages, weight, seq_lens_t, page_table, block_size, max_seq_len + ) + ref = _reference_paged_mqa_logits(q_fp8, weight, ref_k, ref_scale, max_seq_len) + + for b, seq_len in enumerate(seq_lens): + torch.testing.assert_close( + logits[b, :seq_len], ref[b, :seq_len], atol=1e-2, rtol=1e-2 + ) + + +def test_fp8_paged_mqa_logits_cpu_non_contiguous_kv_cache_row_stride(): + """The paged K-cache in production is a per-layer view into a shared + multi-layer allocation, so its page stride generally exceeds the page + byte width -- check the kernel reads a padded/strided buffer identically + to an equivalent tightly-packed one.""" + block_size = 64 + num_heads = 4 + seq_lens = [1, block_size + 7, 3 * block_size] + n_pages, buf_width = _paged_mqa_logits_page_layout(block_size, seq_lens) + cache_pad = 48 + + tight_pages = torch.zeros((n_pages, buf_width), dtype=torch.uint8) + q_fp8, weight, seq_lens_t, page_table, _, _, max_seq_len = ( + _build_paged_mqa_logits_batch( + block_size, seq_lens, num_heads, tight_pages, seed=0 + ) + ) + + padded = torch.zeros((n_pages, buf_width + cache_pad), dtype=torch.uint8) + strided_pages = padded[:, :buf_width] + _build_paged_mqa_logits_batch( + block_size, seq_lens, num_heads, strided_pages, seed=0 + ) + assert strided_pages.stride(0) != strided_pages.size(1) + + logits_tight = fp8_paged_mqa_logits_cpu( + q_fp8, tight_pages, weight, seq_lens_t, page_table, block_size, max_seq_len + ) + logits_strided = fp8_paged_mqa_logits_cpu( + q_fp8, strided_pages, weight, seq_lens_t, page_table, block_size, max_seq_len + ) + # Only [:seq_len] is ever written (and ever read downstream by + # topk_transform_512_cpu) -- positions beyond that are uninitialized, so + # comparing the full row would be comparing unrelated garbage. + for b, seq_len in enumerate(seq_lens): + torch.testing.assert_close( + logits_tight[b, :seq_len], logits_strided[b, :seq_len] + ) + + +# --------------------------------------------------------------------------- +# fused_indexer_q_rope_quant_cpu (indexer.cpp): indexer Q-side RoPE + FP8 +# quant, and the dispatcher wiring that routes to it on CPU +# --------------------------------------------------------------------------- +# +# Independent eager reference: this kernel's absmax reduction is over an +# unusual mix (raw NoPE values plus a bf16-round-tripped RoPE half), which +# per_token_group_quant_fp8 does not replicate. Allow <=1 fp8 ULP: the RoPE +# halves round a value sitting exactly on a tie to either side. + +INDEXER_Q_N_HEAD = 64 +INDEXER_Q_MAX_POS = 4096 + + +def _indexer_q_rope_quant_reference( + positions: torch.Tensor, + q: torch.Tensor, + cos_sin_cache: torch.Tensor, + weights: torch.Tensor, + softmax_scale: float, + head_scale: float, +) -> tuple[torch.Tensor, torch.Tensor]: + _num_tokens, _num_heads, head_dim = q.shape + half_rot_dim = cos_sin_cache.shape[-1] // 2 + rot_dim = 2 * half_rot_dim + nope_dim = head_dim - rot_dim + + q_f32 = q.float() + cos_sin_f32 = cos_sin_cache.float() + pos_cos_sin = cos_sin_f32[positions] # [T, rot_dim] + cos = pos_cos_sin[:, :half_rot_dim].unsqueeze(1) + sin = pos_cos_sin[:, half_rot_dim:].unsqueeze(1) + + nope = q_f32[..., :nope_dim] + rot = q_f32[..., nope_dim:] + x_even = rot[..., 0::2] + x_odd = rot[..., 1::2] + r_even = (x_even * cos - x_odd * sin).to(torch.bfloat16).to(torch.float32) + r_odd = (x_odd * cos + x_even * sin).to(torch.bfloat16).to(torch.float32) + rope_out = torch.stack((r_even, r_odd), dim=-1).flatten(-2) + + combined = torch.cat((nope, rope_out), dim=-1) # [T, H, head_dim] + amax = combined.abs().amax(dim=-1).clamp(min=1e-4) # [T, H] + exponent = torch.ceil(torch.log2(amax / FP8_MAX)) + scale = torch.exp2(exponent) + inv_scale = scale.reciprocal().unsqueeze(-1) + q_fp8 = (combined * inv_scale).clamp(-FP8_MAX, FP8_MAX).to(torch.float8_e4m3fn) + + weights_out = weights.to(torch.float32) * scale * softmax_scale * head_scale + return q_fp8, weights_out + + +@pytest.mark.parametrize("num_tokens", [1, 257]) +def test_fused_indexer_q_rope_quant_cpu_matches_eager(num_tokens: int): + torch.manual_seed(0) + + q = torch.randn(num_tokens, INDEXER_Q_N_HEAD, INDEXER_HEAD_DIM, dtype=torch.float32) + positions = torch.randint(0, INDEXER_Q_MAX_POS, (num_tokens,), dtype=torch.int64) + cos_sin_cache = torch.randn(INDEXER_Q_MAX_POS, ROPE_DIM, dtype=torch.float32) + weights = torch.randn(num_tokens, INDEXER_Q_N_HEAD, dtype=torch.float32) + softmax_scale = INDEXER_HEAD_DIM**-0.5 + head_scale = INDEXER_Q_N_HEAD**-0.5 + + q_fp8_ref, weights_ref = _indexer_q_rope_quant_reference( + positions, q, cos_sin_cache, weights, softmax_scale, head_scale + ) + + q_fp8_fused = torch.empty_like(q, dtype=torch.float8_e4m3fn) + weights_fused = torch.empty_like(weights, dtype=torch.float32) + fused_indexer_q_rope_quant_cpu( + positions, + q, + cos_sin_cache, + q_fp8_fused, + weights, + softmax_scale, + head_scale, + weights_fused, + ) + + max_ulp = int(fp8_ulp_distance(q_fp8_fused, q_fp8_ref).max().item()) + assert max_ulp <= 1, f"q_fp8 differs by {max_ulp} ULP (>1)" + torch.testing.assert_close(weights_fused, weights_ref, rtol=1e-5, atol=1e-6) + + +def test_fused_indexer_q_rope_quant_dispatches_to_cpu_kernel(): + """End-to-end wiring check: the production dispatcher must route to the + CPU op on a CPU tensor, not silently fall through to the Triton kernel + (which would also "work" under triton-cpu, masking a dispatch + regression). The dispatcher does its CPU import inside the CPU branch, + so the patch target is the source module.""" + torch.manual_seed(0) + num_tokens = 1 + + q = torch.randn(num_tokens, INDEXER_Q_N_HEAD, INDEXER_HEAD_DIM, dtype=torch.float32) + positions = torch.randint(0, INDEXER_Q_MAX_POS, (num_tokens,), dtype=torch.int64) + cos_sin_cache = torch.randn(INDEXER_Q_MAX_POS, ROPE_DIM, dtype=torch.float32) + weights = torch.randn(num_tokens, INDEXER_Q_N_HEAD, dtype=torch.float32) + softmax_scale = INDEXER_HEAD_DIM**-0.5 + head_scale = INDEXER_Q_N_HEAD**-0.5 + + q_fp8_ref, weights_ref = _indexer_q_rope_quant_reference( + positions, q, cos_sin_cache, weights, softmax_scale, head_scale + ) + + with mock.patch("vllm._custom_ops.fused_indexer_q_rope_quant_cpu") as mocked: + fused_indexer_q_rope_quant( + positions, + q.clone(), + cos_sin_cache, + weights, + softmax_scale, + head_scale, + use_fp4=False, + ) + assert mocked.called, ( + "fused_indexer_q_rope_quant did not dispatch to the CPU kernel " + "on a CPU tensor" + ) + + q_fp8_fused, weights_fused = fused_indexer_q_rope_quant( + positions, + q.clone(), + cos_sin_cache, + weights, + softmax_scale, + head_scale, + use_fp4=False, + ) + max_ulp = int(fp8_ulp_distance(q_fp8_fused, q_fp8_ref).max().item()) + assert max_ulp <= 1, f"q_fp8 differs by {max_ulp} ULP (>1)" + torch.testing.assert_close(weights_fused, weights_ref, rtol=1e-5, atol=1e-6) + + +# --------------------------------------------------------------------------- +# mhc_pre_cpu / mhc_post_cpu / hc_head_fused_cpu (mhc.cpp): CPU ports of the +# mHC gating kernels +# --------------------------------------------------------------------------- + + +def _rmsnorm_nw_torch(x: torch.Tensor, eps: float) -> torch.Tensor: + """Weight-free RMSNorm over the last dim, matching hc_head_fused_cpu.""" + x_f32 = x.float() + var = x_f32.square().mean(dim=-1, keepdim=True) + return (x_f32 * torch.rsqrt(var + eps)).to(x.dtype) + + +def _hc_head_torch( + hidden_states: torch.Tensor, + hc_fn: torch.Tensor, + hc_scale: torch.Tensor, + hc_base: torch.Tensor, + rms_norm_eps: float, + hc_eps: float, +) -> torch.Tensor: + """Eager reference for hc_head_fused_cpu: round-trips the RMSNorm + output through bf16 before the linear, matching the kernel's bf16 + arithmetic.""" + hc_mult, hidden_size = hidden_states.shape[-2:] + outer_shape = hidden_states.shape[:-2] + hs_flat = hidden_states.reshape(-1, hc_mult, hidden_size) + + x_flat = hs_flat.flatten(-2) + x_normed = _rmsnorm_nw_torch(x_flat, rms_norm_eps) + mixes = torch.nn.functional.linear(x_normed.float(), hc_fn) + pre = torch.sigmoid(mixes * hc_scale + hc_base) + hc_eps + + out = torch.sum(pre.unsqueeze(-1) * hs_flat.float(), dim=1).to(hidden_states.dtype) + return out.view(*outer_shape, hidden_size) + + +@pytest.mark.parametrize("num_tokens", [1, 128]) +@pytest.mark.parametrize("hidden_size", [4096, 7168]) +def test_mhc_pre_cpu(num_tokens: int, hidden_size: int): + """mhc_pre_cpu vs. the mhc_pre_torch eager reference it replaces.""" + torch.set_default_device(DEVICE) + set_random_seed(0) + hc_mult = 4 + + hc_mult3 = 2 * hc_mult + hc_mult * hc_mult + residual = torch.randn((num_tokens, hc_mult, hidden_size), dtype=torch.bfloat16) + fn = torch.randn((hc_mult3, hc_mult * hidden_size), dtype=torch.float32) * 1e-4 + hc_scale = torch.randn((3,), dtype=torch.float32) * 0.1 + hc_base = torch.randn((hc_mult3,), dtype=torch.float32) * 0.1 + + rms_eps = hc_eps = 1e-6 + hc_post_mult_value = 2.0 + sinkhorn_repeat = 8 + + ref = mhc_kernels.mhc_pre_torch( + residual, + fn, + hc_scale, + hc_base, + rms_eps, + hc_eps, + hc_eps, + hc_post_mult_value, + sinkhorn_repeat, + ) + out = mhc_kernels.mhc_pre_cpu( + residual, + fn, + hc_scale, + hc_base, + rms_eps, + hc_eps, + hc_eps, + hc_post_mult_value, + sinkhorn_repeat, + ) + + for actual, expected in zip(out, ref, strict=True): + torch.testing.assert_close(actual, expected, atol=5e-2, rtol=1e-2) + + +@pytest.mark.parametrize("num_tokens", [1, 128]) +@pytest.mark.parametrize("hidden_size", [4096, 7168]) +def test_mhc_post_cpu(num_tokens: int, hidden_size: int): + """mhc_post_cpu vs. the mhc_post_torch eager reference it replaces.""" + torch.set_default_device(DEVICE) + set_random_seed(0) + hc_mult = 4 + + x = torch.randn((num_tokens, hidden_size), dtype=torch.bfloat16) + residual = torch.randn((num_tokens, hc_mult, hidden_size), dtype=torch.bfloat16) + post_layer_mix = torch.randn((num_tokens, hc_mult, 1), dtype=torch.float32) + comb_res_mix = torch.randn((num_tokens, hc_mult, hc_mult), dtype=torch.float32) + + ref = mhc_kernels.mhc_post_torch(x, residual, post_layer_mix, comb_res_mix) + out = mhc_kernels.mhc_post_cpu(x, residual, post_layer_mix, comb_res_mix) + + torch.testing.assert_close(out, ref, atol=5e-2, rtol=1e-2) + + +@pytest.mark.parametrize("num_tokens", [1, 128]) +@pytest.mark.parametrize("hidden_size", [4096, 7168]) +def test_hc_head_cpu(num_tokens: int, hidden_size: int): + """hc_head_fused_cpu vs. the _hc_head_torch eager reference it + replaces.""" + torch.set_default_device(DEVICE) + set_random_seed(0) + hc_mult = 4 + + residual = torch.randn((num_tokens, hc_mult, hidden_size), dtype=torch.bfloat16) + fn = torch.randn((hc_mult, hc_mult * hidden_size), dtype=torch.float32) * 1e-4 + hc_scale = torch.randn((1,), dtype=torch.float32) * 0.1 + hc_base = torch.randn((hc_mult,), dtype=torch.float32) * 0.1 + rms_eps = hc_eps = 1e-6 + + ref = _hc_head_torch(residual, fn, hc_scale, hc_base, rms_eps, hc_eps) + out = mhc_kernels.hc_head_fused_cpu( + residual, fn, hc_scale, hc_base, rms_eps, hc_eps + ) + + torch.testing.assert_close(out, ref, atol=5e-2, rtol=1e-2) diff --git a/tests/models/multimodal/processing/test_common.py b/tests/models/multimodal/processing/test_common.py index b84432cb95..bbc39ca171 100644 --- a/tests/models/multimodal/processing/test_common.py +++ b/tests/models/multimodal/processing/test_common.py @@ -92,6 +92,11 @@ _XPU_EXCLUDED_MODEL_IDS = { "thinkingmachines/Inkling-NVFP4", } +_CPU_EXCLUDED_MODEL_IDS = { + # DeepSeek-V4 vision variant is only supported on NVIDIA GPUs for now. + "deepseek-ai/DeepSeek-V4-Flash-Vision-Exp", +} + def _iter_model_ids_to_test(model_arch_list: AbstractSet[str]): for model_arch in model_arch_list: @@ -113,6 +118,11 @@ def _get_model_ids_to_test(model_arch_list: AbstractSet[str]): while excluded_model_id in model_ids: model_ids.remove(excluded_model_id) + if current_platform.is_cpu(): + for excluded_model_id in _CPU_EXCLUDED_MODEL_IDS: + while excluded_model_id in model_ids: + model_ids.remove(excluded_model_id) + return model_ids diff --git a/tests/v1/attention/test_sparse_mla_backends.py b/tests/v1/attention/test_sparse_mla_backends.py index b2490e391b..002aa5a4b7 100644 --- a/tests/v1/attention/test_sparse_mla_backends.py +++ b/tests/v1/attention/test_sparse_mla_backends.py @@ -56,7 +56,7 @@ from vllm.v1.attention.backends.mla.flashmla_sparse import ( FlashMLASparseMetadataBuilder, triton_convert_req_index_to_global_index, ) -from vllm.v1.attention.backends.mla.indexer import split_indexer_prefill_chunks +from vllm.v1.attention.backends.mla.indexer import DeepseekV32IndexerMetadataBuilder from vllm.v1.attention.backends.utils import ( split_decodes_and_prefills, split_prefill_chunks, @@ -1424,7 +1424,7 @@ def test_sparse_backend_prefill_correctness( def test_split_indexer_prefill_chunks( seq_lens, query_lens, workspace_size, max_logits_bytes, expected ): - out = split_indexer_prefill_chunks( + out = DeepseekV32IndexerMetadataBuilder._split_indexer_prefill_chunks( seq_lens, query_lens, workspace_size, @@ -1438,7 +1438,9 @@ def test_split_indexer_prefill_chunks_single_request_overflow(): seq_lens = torch.tensor([1000, 50]) query_lens = torch.tensor([100, 5]) - out = split_indexer_prefill_chunks(seq_lens, query_lens, 2000, 1000) + out = DeepseekV32IndexerMetadataBuilder._split_indexer_prefill_chunks( + seq_lens, query_lens, 2000, 1000 + ) # max_logits_elems = 250, N=1000 -> max_q = 1 -> 100 query sub-chunks expected = [(slice(0, 1), slice(i, i + 1)) for i in range(100)] # req1: M=5, N=50 -> 250 elems fits budget diff --git a/vllm/_custom_ops.py b/vllm/_custom_ops.py index 5856acc987..db20adb994 100644 --- a/vllm/_custom_ops.py +++ b/vllm/_custom_ops.py @@ -3357,31 +3357,6 @@ class CPUQuantMethod(IntEnum): MXFP4 = 4 -if hasattr(torch.ops._C, "fused_experts_cpu"): - - @register_fake("_C::fused_experts_cpu") - def fused_experts_cpu_fake( - hidden_states: torch.Tensor, - w1: torch.Tensor, - w2: torch.Tensor, - topk_weights: torch.Tensor, - topk_ids: torch.Tensor, - inplace: bool, - moe_comp_method: CPUQuantMethod, - w1_scale: torch.Tensor | None, - w2_scale: torch.Tensor | None, - w1_zero: torch.Tensor | None, - w2_zero: torch.Tensor | None, - block_size: list[int] | None, - w1_bias: torch.Tensor | None, - w2_bias: torch.Tensor | None, - alpha: float | None, - limit: float | None, - is_vnni: bool, - ) -> torch.Tensor: - return torch.empty_like(hidden_states) - - if hasattr(torch.ops._C, "dynamic_4bit_int_moe"): @register_fake("_C::dynamic_4bit_int_moe") @@ -3401,12 +3376,12 @@ if hasattr(torch.ops._C, "dynamic_4bit_int_moe"): def fused_experts_cpu( + out: torch.Tensor, hidden_states: torch.Tensor, w1: torch.Tensor, w2: torch.Tensor, topk_weights: torch.Tensor, topk_ids: torch.Tensor, - inplace: bool, moe_comp_method: CPUQuantMethod, w1_scale: torch.Tensor | None, w2_scale: torch.Tensor | None, @@ -3418,14 +3393,14 @@ def fused_experts_cpu( alpha: float | None = None, limit: float | None = None, is_vnni: bool = True, -) -> torch.Tensor: - return torch.ops._C.fused_experts_cpu( +) -> None: + torch.ops._C.fused_experts_cpu( + out, hidden_states, w1, w2, topk_weights, topk_ids, - inplace, moe_comp_method, w1_scale, w2_scale, @@ -4059,6 +4034,393 @@ def amx_mla_concat_and_cache( torch.ops._C.concat_and_cache_mla_cpu(kv_c_normed, k_pe, kv_cache, slot_mapping) +if hasattr(torch.ops._C, "hc_pre_fused_cpu"): + + @register_fake("_C::hc_pre_fused_cpu") + def hc_pre_fused_cpu_fake( + x: torch.Tensor, + hc_fn: torch.Tensor, + hc_scale: torch.Tensor, + hc_base: torch.Tensor, + hc_mult: int, + sinkhorn_iters: int, + rms_eps: float, + hc_eps: float, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + t, _, d = x.shape + y = x.new_empty((t, d)) + post = x.new_empty((t, hc_mult), dtype=torch.float32) + comb = x.new_empty((t, hc_mult, hc_mult), dtype=torch.float32) + return y, post, comb + + +def hc_pre_fused_cpu( + x: torch.Tensor, + hc_fn: torch.Tensor, + hc_scale: torch.Tensor, + hc_base: torch.Tensor, + hc_mult: int, + sinkhorn_iters: int, + rms_eps: float, + hc_eps: float, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + return torch.ops._C.hc_pre_fused_cpu( + x, hc_fn, hc_scale, hc_base, hc_mult, sinkhorn_iters, rms_eps, hc_eps + ) + + +if hasattr(torch.ops._C, "hc_post_fused_cpu"): + + @register_fake("_C::hc_post_fused_cpu") + def hc_post_fused_cpu_fake( + x: torch.Tensor, + residual: torch.Tensor, + post: torch.Tensor, + comb: torch.Tensor, + ) -> torch.Tensor: + t, d = x.shape + hc = residual.size(1) + return x.new_empty((t, hc, d)) + + +def hc_post_fused_cpu( + x: torch.Tensor, + residual: torch.Tensor, + post: torch.Tensor, + comb: torch.Tensor, +) -> torch.Tensor: + return torch.ops._C.hc_post_fused_cpu(x, residual, post, comb) + + +if hasattr(torch.ops._C, "hc_head_fused_cpu"): + + @register_fake("_C::hc_head_fused_cpu") + def hc_head_fused_cpu_fake( + x: torch.Tensor, + hc_fn: torch.Tensor, + hc_scale: torch.Tensor, + hc_base: torch.Tensor, + hc_eps: float, + norm_eps: float, + ) -> torch.Tensor: + t, _, d = x.shape + return x.new_empty((t, d)) + + +def hc_head_fused_cpu( + x: torch.Tensor, + hc_fn: torch.Tensor, + hc_scale: torch.Tensor, + hc_base: torch.Tensor, + hc_eps: float, + norm_eps: float, +) -> torch.Tensor: + return torch.ops._C.hc_head_fused_cpu(x, hc_fn, hc_scale, hc_base, hc_eps, norm_eps) + + +if hasattr(torch.ops._C, "fused_qnorm_rope_kv_insert_cpu"): + + @register_fake("_C::fused_qnorm_rope_kv_insert_cpu") + def fused_qnorm_rope_kv_insert_cpu_fake( + q: torch.Tensor, + kv: torch.Tensor, + positions: torch.Tensor, + swa_kv_cache_2d: torch.Tensor, + slot_mapping: torch.Tensor, + cos_sin_cache: torch.Tensor, + q_head_padded: int, + eps: float, + cache_block_size: int, + ) -> torch.Tensor: + num_tokens_full, _, head_dim = q.shape + return q.new_empty((num_tokens_full, q_head_padded, head_dim)) + + +def fused_qnorm_rope_kv_insert_cpu( + q: torch.Tensor, + kv: torch.Tensor, + positions: torch.Tensor, + swa_kv_cache_2d: torch.Tensor, + slot_mapping: torch.Tensor, + cos_sin_cache: torch.Tensor, + q_head_padded: int, + eps: float, + cache_block_size: int, +) -> torch.Tensor: + return torch.ops._C.fused_qnorm_rope_kv_insert_cpu( + q, + kv, + positions, + swa_kv_cache_2d, + slot_mapping, + cos_sin_cache, + q_head_padded, + eps, + cache_block_size, + ) + + +def flash_mla_with_kvcache_cpu( + out: torch.Tensor, + q: torch.Tensor, + window_cache_2d: torch.Tensor, + window_slots: torch.Tensor, + window_block_size: int, + compressed_cache_2d: torch.Tensor, + compressed_slots: torch.Tensor, + compressed_block_size: int, + attn_sink: torch.Tensor, + scale: float, +) -> None: + torch.ops._C.flash_mla_with_kvcache_cpu( + out, + q, + window_cache_2d, + window_slots, + window_block_size, + compressed_cache_2d, + compressed_slots, + compressed_block_size, + attn_sink, + scale, + ) + + +def save_partial_states_cpu( + kv: torch.Tensor, + score: torch.Tensor, + ape: torch.Tensor, + positions: torch.Tensor, + state_cache: torch.Tensor, + slot_mapping: torch.Tensor, +) -> None: + torch.ops._C.save_partial_states_cpu( + kv, score, ape, positions, state_cache, slot_mapping + ) + + +def compress_norm_rope_store_cpu( + state_cache: torch.Tensor, + gather_slots: torch.Tensor, + positions: torch.Tensor, + kv_slot_mapping: torch.Tensor, + rms_norm_weight: torch.Tensor, + rms_norm_eps: float, + cos_sin_cache: torch.Tensor, + kv_cache_2d: torch.Tensor, + kv_cache_block_size: int, + compress_ratio: int, +) -> None: + torch.ops._C.compress_norm_rope_store_cpu( + state_cache, + gather_slots, + positions, + kv_slot_mapping, + rms_norm_weight, + rms_norm_eps, + cos_sin_cache, + kv_cache_2d, + kv_cache_block_size, + compress_ratio, + ) + + +def compress_norm_rope_store_indexer_cpu( + state_cache: torch.Tensor, + gather_slots: torch.Tensor, + positions: torch.Tensor, + kv_slot_mapping: torch.Tensor, + rms_norm_weight: torch.Tensor, + rms_norm_eps: float, + cos_sin_cache: torch.Tensor, + kv_cache_2d: torch.Tensor, + kv_cache_block_size: int, + compress_ratio: int, +) -> None: + torch.ops._C.compress_norm_rope_store_indexer_cpu( + state_cache, + gather_slots, + positions, + kv_slot_mapping, + rms_norm_weight, + rms_norm_eps, + cos_sin_cache, + kv_cache_2d, + kv_cache_block_size, + compress_ratio, + ) + + +if hasattr(torch.ops._C, "fp8_paged_mqa_logits_cpu"): + + @register_fake("_C::fp8_paged_mqa_logits_cpu") + def fp8_paged_mqa_logits_cpu_fake( + q_fp8: torch.Tensor, + kvcache_fp8: torch.Tensor, + weight: torch.Tensor, + seq_lens: torch.Tensor, + page_table: torch.Tensor, + block_size: int, + max_seq_len: int, + ) -> torch.Tensor: + batch_size = q_fp8.size(0) + return q_fp8.new_empty((batch_size, max_seq_len), dtype=torch.float32) + + +def fp8_paged_mqa_logits_cpu( + q_fp8: torch.Tensor, + kvcache_fp8: torch.Tensor, + weight: torch.Tensor, + seq_lens: torch.Tensor, + page_table: torch.Tensor, + block_size: int, + max_seq_len: int, +) -> torch.Tensor: + return torch.ops._C.fp8_paged_mqa_logits_cpu( + q_fp8, kvcache_fp8, weight, seq_lens, page_table, block_size, max_seq_len + ) + + +def topk_transform_512_cpu( + scores: torch.Tensor, + seq_lens: torch.Tensor, + page_tables: torch.Tensor, + out_page_indices: torch.Tensor, + page_size: int, + out_raw_indices: torch.Tensor | None = None, +) -> None: + torch.ops._C.topk_transform_512_cpu( + scores, seq_lens, page_tables, out_page_indices, page_size, out_raw_indices + ) + + +def fused_indexer_q_rope_quant_cpu( + positions: torch.Tensor, + index_q: torch.Tensor, + index_q_cos_sin_cache: torch.Tensor, + index_q_fp8: torch.Tensor, + index_weights: torch.Tensor, + index_weights_softmax_scale: float, + index_weights_head_scale: float, + index_weights_out: torch.Tensor, +) -> None: + torch.ops._C.fused_indexer_q_rope_quant_cpu( + positions, + index_q, + index_q_cos_sin_cache, + index_q_fp8, + index_weights, + index_weights_softmax_scale, + index_weights_head_scale, + index_weights_out, + ) + + +if hasattr(torch.ops._C, "biased_topk_cpu"): + + @register_fake("_C::biased_topk_cpu") + def biased_topk_cpu_fake( + hidden_states: torch.Tensor, + gating_output: torch.Tensor, + correction_bias: torch.Tensor, + topk: int, + renormalize: bool, + scoring_func: str, + num_fused_shared_experts: int, + routed_scaling_factor: float | None, + apply_routed_scaling_factor_on_output: bool, + ) -> tuple[torch.Tensor, torch.Tensor]: + num_tokens = hidden_states.size(0) + topk_weights = hidden_states.new_empty((num_tokens, topk), dtype=torch.float32) + topk_ids = hidden_states.new_empty((num_tokens, topk), dtype=torch.int32) + return topk_weights, topk_ids + + +def biased_topk_cpu( + hidden_states: torch.Tensor, + gating_output: torch.Tensor, + correction_bias: torch.Tensor, + topk: int, + renormalize: bool, + scoring_func: str, + num_fused_shared_experts: int, + routed_scaling_factor: float | None, + apply_routed_scaling_factor_on_output: bool, +) -> tuple[torch.Tensor, torch.Tensor]: + return torch.ops._C.biased_topk_cpu( + hidden_states, + gating_output, + correction_bias, + topk, + renormalize, + scoring_func, + num_fused_shared_experts, + routed_scaling_factor, + apply_routed_scaling_factor_on_output, + ) + + +if hasattr(torch.ops._C, "hash_topk_cpu"): + + @register_fake("_C::hash_topk_cpu") + def hash_topk_cpu_fake( + gating_output: torch.Tensor, + tid2eid: torch.Tensor, + topk: int, + scoring_func: str, + num_fused_shared_experts: int, + num_experts: int, + routed_scaling_factor: float, + ) -> tuple[torch.Tensor, torch.Tensor]: + num_tokens = gating_output.size(0) + topk_weights = gating_output.new_empty((num_tokens, topk), dtype=torch.float32) + topk_ids = gating_output.new_empty((num_tokens, topk), dtype=torch.int32) + return topk_weights, topk_ids + + +def hash_topk_cpu( + gating_output: torch.Tensor, + tid2eid: torch.Tensor, + topk: int, + scoring_func: str, + num_fused_shared_experts: int, + num_experts: int, + routed_scaling_factor: float, +) -> tuple[torch.Tensor, torch.Tensor]: + return torch.ops._C.hash_topk_cpu( + gating_output, + tid2eid, + topk, + scoring_func, + num_fused_shared_experts, + num_experts, + routed_scaling_factor, + ) + + +if hasattr(torch.ops._C, "inverse_gptj_rope_o_proj_cpu"): + + @register_fake("_C::inverse_gptj_rope_o_proj_cpu") + def inverse_gptj_rope_o_proj_cpu_fake( + o: torch.Tensor, + positions: torch.Tensor, + cos_sin_cache: torch.Tensor, + rope_dim: int, + ) -> torch.Tensor: + return o.new_empty(o.shape, dtype=torch.float32) + + +def inverse_gptj_rope_o_proj_cpu( + o: torch.Tensor, + positions: torch.Tensor, + cos_sin_cache: torch.Tensor, + rope_dim: int, +) -> torch.Tensor: + return torch.ops._C.inverse_gptj_rope_o_proj_cpu( + o, positions, cos_sin_cache, rope_dim + ) + + def cpu_gemm_wna16( input: torch.Tensor, q_weight: torch.Tensor, @@ -4107,6 +4469,7 @@ def cpu_prepack_moe_weight_int8( def cpu_fused_moe( + out: torch.Tensor, input: torch.Tensor, w13: torch.Tensor, w2: torch.Tensor, @@ -4117,10 +4480,9 @@ def cpu_fused_moe( act: str, isa: str, skip_weighted: bool = False, -) -> torch.Tensor: - output = torch.empty_like(input) +) -> None: torch.ops._C.cpu_fused_moe( - output, + out, input, w13, w2, @@ -4132,10 +4494,10 @@ def cpu_fused_moe( act, isa, ) - return output def cpu_fused_moe_int8( + out: torch.Tensor, input: torch.Tensor, w13: torch.Tensor, w2: torch.Tensor, @@ -4148,10 +4510,9 @@ def cpu_fused_moe_int8( act: str, isa: str, skip_weighted: bool = False, -) -> torch.Tensor: - output = torch.empty_like(input) +) -> None: torch.ops._C.cpu_fused_moe_int8( - output, + out, input, w13, w2, @@ -4165,7 +4526,6 @@ def cpu_fused_moe_int8( act, isa, ) - return output if hasattr(torch.ops._qutlass_C, "matmul_mxf4_bf16_tn"): diff --git a/vllm/model_executor/kernels/linear/scaled_mm/cpu.py b/vllm/model_executor/kernels/linear/scaled_mm/cpu.py index cb0fe0325b..4ef1af0a70 100644 --- a/vllm/model_executor/kernels/linear/scaled_mm/cpu.py +++ b/vllm/model_executor/kernels/linear/scaled_mm/cpu.py @@ -6,6 +6,9 @@ import torch from vllm import _custom_ops as ops from vllm.model_executor.layers.quantization.utils import replace_parameter +from vllm.model_executor.layers.quantization.utils.fp8_utils import ( + _upcast_e8m0_to_fp32, +) from vllm.model_executor.layers.quantization.utils.w8a8_utils import ( convert_to_channelwise, ) @@ -281,10 +284,12 @@ class CPUFp8BlockScaledMMKernel(Fp8BlockScaledMMLinearKernel): else params.weight_scale ) assert weight_scale is not None + if weight_scale.dtype in (torch.float8_e8m0fnu, torch.uint8): + weight_scale = _upcast_e8m0_to_fp32(weight_scale.data) replace_parameter( layer, scale_attr, - torch.nn.Parameter(weight_scale.data, requires_grad=False), + torch.nn.Parameter(weight_scale.data.contiguous(), requires_grad=False), ) def apply_weights( diff --git a/vllm/model_executor/kernels/mhc/__init__.py b/vllm/model_executor/kernels/mhc/__init__.py index 3bde02e28a..324e3cec64 100644 --- a/vllm/model_executor/kernels/mhc/__init__.py +++ b/vllm/model_executor/kernels/mhc/__init__.py @@ -1,6 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project from .aiter import * +from .cpu import * from .tilelang import * from .torch import * from .triton import * @@ -14,6 +15,9 @@ __all__ = [ "mhc_post_aiter", "mhc_fused_post_pre_aiter", "hc_head_fused_aiter", + "mhc_pre_cpu", + "mhc_post_cpu", + "hc_head_fused_cpu", "mhc_pre_tilelang", "mhc_post_tilelang", "mhc_fused_post_pre_tilelang", diff --git a/vllm/model_executor/kernels/mhc/cpu.py b/vllm/model_executor/kernels/mhc/cpu.py new file mode 100644 index 0000000000..4659540698 --- /dev/null +++ b/vllm/model_executor/kernels/mhc/cpu.py @@ -0,0 +1,96 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import torch + +from vllm import _custom_ops as ops + + +def mhc_pre_cpu( + residual: torch.Tensor, + fn: torch.Tensor, + hc_scale: torch.Tensor, + hc_base: torch.Tensor, + rms_eps: float, + hc_pre_eps: float, + hc_sinkhorn_eps: float, + hc_post_mult_value: float, + sinkhorn_repeat: int, + n_splits: int = 1, + norm_weight: torch.Tensor | None = None, + norm_eps: float = 0.0, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """CPU-ported mHC pre block (see `mhc_pre_torch` for the eager reference). + + The ported kernel (`hc_pre_fused_cpu`) only exposes one merged `hc_eps` + (used for both `hc_pre_eps`/`hc_sinkhorn_eps`) and hardcodes the post-mix + multiplier to 2.0 -- true of every real call site in + `models/deepseek_v4/cpu/model.py`, so this is not a capability loss here. + """ + assert n_splits == 1, "mhc_pre_cpu does not support n_splits != 1" + assert hc_pre_eps == hc_sinkhorn_eps, ( + "mhc_pre_cpu requires hc_pre_eps == hc_sinkhorn_eps (single merged hc_eps)" + ) + assert hc_post_mult_value == 2.0, "mhc_pre_cpu hardcodes the post-mix multiplier" + + hc_mult, hidden_size = residual.shape[-2:] + outer_shape = residual.shape[:-2] + x_flat = residual.reshape(-1, hc_mult, hidden_size) + + layer_input, post, comb = ops.hc_pre_fused_cpu( + x_flat, + fn, + hc_scale, + hc_base, + hc_mult, + sinkhorn_repeat, + rms_eps, + hc_pre_eps, + ) + + post_mix = post.view(*outer_shape, hc_mult, 1) + comb_mix = comb.view(*outer_shape, hc_mult, hc_mult) + layer_input = layer_input.view(*outer_shape, hidden_size) + return post_mix, comb_mix, layer_input + + +def mhc_post_cpu( + x: torch.Tensor, + residual: torch.Tensor, + post_layer_mix: torch.Tensor, + comb_res_mix: torch.Tensor, +) -> torch.Tensor: + """CPU-ported mHC post block (see `mhc_post_torch` for the eager reference).""" + hc_mult, hidden_size = residual.shape[-2:] + outer_shape = residual.shape[:-2] + + x_flat = x.reshape(-1, hidden_size) + residual_flat = residual.reshape(-1, hc_mult, hidden_size) + post_flat = post_layer_mix.reshape(-1, hc_mult).float() + comb_flat = comb_res_mix.reshape(-1, hc_mult, hc_mult).float() + + out = ops.hc_post_fused_cpu(x_flat, residual_flat, post_flat, comb_flat) + return out.view(*outer_shape, hc_mult, hidden_size) + + +def hc_head_fused_cpu( + hidden_states: torch.Tensor, + hc_fn: torch.Tensor, + hc_scale: torch.Tensor, + hc_base: torch.Tensor, + rms_norm_eps: float, + hc_eps: float, +) -> torch.Tensor: + """CPU-ported HC head reduction (see `test_hc_head_cpu` in + tests/kernels/test_mhc_kernels.py for the eager reference this is tested + against). + + The ported kernel's C++ signature takes `(hc_eps, norm_eps)` -- the + opposite order from this wrapper's `(rms_norm_eps, hc_eps)`, which matches + the real call site in `models/deepseek_v4/cpu/model.py`. + """ + hc_mult, hidden_size = hidden_states.shape[-2:] + outer_shape = hidden_states.shape[:-2] + hs_flat = hidden_states.reshape(-1, hc_mult, hidden_size) + + out = ops.hc_head_fused_cpu(hs_flat, hc_fn, hc_scale, hc_base, hc_eps, rms_norm_eps) + return out.view(*outer_shape, hidden_size) diff --git a/vllm/model_executor/layers/fused_moe/experts/cpu_int4_moe.py b/vllm/model_executor/layers/fused_moe/experts/cpu_int4_moe.py index 5816bc5d3a..b1e32462d8 100644 --- a/vllm/model_executor/layers/fused_moe/experts/cpu_int4_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/cpu_int4_moe.py @@ -218,7 +218,7 @@ class CPUExpertsInt4(mk.FusedMoEExpertsMonolithic): Returns: Output tensor after MoE computation """ - from vllm.model_executor.layers.fused_moe.experts.cpu_moe import ( + from vllm.model_executor.layers.fused_moe.router.cpu_router import ( select_experts, ) diff --git a/vllm/model_executor/layers/fused_moe/experts/cpu_moe.py b/vllm/model_executor/layers/fused_moe/experts/cpu_moe.py index 00c3dcc17b..44a5bda21e 100644 --- a/vllm/model_executor/layers/fused_moe/experts/cpu_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/cpu_moe.py @@ -4,7 +4,6 @@ import math import sys -from collections.abc import Callable from typing import cast import torch @@ -34,6 +33,9 @@ from vllm.model_executor.layers.fused_moe.config import ( FusedMoEQuantConfig, RoutingMethodType, ) +from vllm.model_executor.layers.fused_moe.topk_weight_and_reduce import ( + TopKWeightAndReduceNoOP, +) from vllm.model_executor.layers.quantization.utils.quant_utils import ( QuantKey, kFp8Dynamic128Sym, @@ -48,133 +50,26 @@ from vllm.platforms import CpuArchEnum, current_platform from vllm.utils.math_utils import round_up logger = init_logger(__name__) -# =========================================================================== -# Routing -# =========================================================================== - - -def grouped_topk( - hidden_states: torch.Tensor, - gating_output: torch.Tensor, - topk: int, - renormalize: bool, - num_expert_group: int = 0, - topk_group: int = 0, - scoring_func: str = "softmax", - routed_scaling_factor: float = 1.0, - e_score_correction_bias: torch.Tensor | None = None, -) -> tuple[torch.Tensor, torch.Tensor]: - assert hidden_states.shape[0] == gating_output.shape[0], "Number of tokens mismatch" - - gating_output = gating_output.float() - if scoring_func == "softmax": - scores = torch.softmax(gating_output, dim=-1) - elif scoring_func == "sigmoid": - scores = gating_output.sigmoid() - else: - raise ValueError(f"Unsupported scoring function: {scoring_func}") - - num_token = scores.shape[0] - if e_score_correction_bias is not None: - original_scores = scores - scores = scores + e_score_correction_bias.unsqueeze(0) - group_scores = ( - scores.view(num_token, num_expert_group, -1).topk(2, dim=-1)[0].sum(dim=-1) - ) - else: - group_scores = ( - scores.view(num_token, num_expert_group, -1).max(dim=-1).values - ) # [n, n_group] - group_idx = torch.topk(group_scores, k=topk_group, dim=-1, sorted=False)[ - 1 - ] # [n, top_k_group] - group_mask = torch.zeros_like(group_scores) # [n, n_group] - group_mask.scatter_(1, group_idx, 1) # [n, n_group] - score_mask = ( - group_mask.unsqueeze(-1) - .expand(num_token, num_expert_group, scores.shape[-1] // num_expert_group) - .reshape(num_token, -1) - ) # [n, e] - tmp_scores = scores.masked_fill(~score_mask.bool(), float("-inf")) # [n, e] - - if e_score_correction_bias is not None: - topk_ids = torch.topk(tmp_scores, k=topk, dim=-1, sorted=False)[1] - topk_weights = original_scores.gather(1, topk_ids) - else: - topk_weights, topk_ids = torch.topk(tmp_scores, k=topk, dim=-1, sorted=False) - - if renormalize: - topk_weights = topk_weights / topk_weights.sum(dim=-1, keepdim=True) - - if routed_scaling_factor != 1.0: - topk_weights = topk_weights * routed_scaling_factor - return topk_weights, topk_ids.to(torch.int32) - - -def select_experts( - hidden_states: torch.Tensor, - router_logits: torch.Tensor, - top_k: int, - use_grouped_topk: bool, - renormalize: bool, - topk_group: int | None = None, - num_expert_group: int | None = None, - custom_routing_function: Callable | None = None, - scoring_func: str = "softmax", - routed_scaling_factor: float = 1.0, - e_score_correction_bias: torch.Tensor | None = None, -) -> tuple[torch.Tensor, torch.Tensor]: - if use_grouped_topk: - assert topk_group is not None - assert num_expert_group is not None - return grouped_topk( - hidden_states=hidden_states, - gating_output=router_logits, - topk=top_k, - renormalize=renormalize, - num_expert_group=num_expert_group, - topk_group=topk_group, - scoring_func=scoring_func, - routed_scaling_factor=routed_scaling_factor, - e_score_correction_bias=e_score_correction_bias, - ) - elif custom_routing_function is None: - assert scoring_func == "softmax" - topk_logit_vals, topk_idx = torch.topk( - router_logits, k=top_k, dim=-1, sorted=False - ) - if renormalize: - topk_vals = torch.softmax(topk_logit_vals, dim=-1) - else: - logZ = torch.logsumexp(router_logits, dim=-1, keepdim=True) - topk_vals = (topk_logit_vals - logZ).exp() - return topk_vals.to(torch.float32), topk_idx.to(torch.int32) - else: - topk_weights, topk_ids = custom_routing_function( - hidden_states=hidden_states, - gating_output=router_logits, - topk=top_k, - renormalize=renormalize, - ) - # cpu_fused_moe reads routing tensors as contiguous float32/int32 - # buffers and does not account for tensor strides. - topk_weights = topk_weights.to(torch.float32).contiguous() - topk_ids = topk_ids.to(torch.int32).contiguous() - return topk_weights, topk_ids - - # =========================================================================== # Unquantized (BF16/FP16/FP32) MoE # =========================================================================== -class CPUUnquantizedExperts(mk.FusedMoEExpertsMonolithic): +class CPUUnquantizedExperts(mk.FusedMoEExpertsModular): """Portable vector grouped-gemm unquantized MoE experts.""" isa = "vec" output_alignment = 32 reduction_alignment = 1 + def __init__( + self, + moe_config: FusedMoEConfig, + quant_config: FusedMoEQuantConfig, + ): + super().__init__(moe_config, quant_config) + self._use_zentorch = False + @classmethod def _intermediate_alignment(cls) -> int: return math.lcm(cls.output_alignment, cls.reduction_alignment) @@ -210,20 +105,6 @@ class CPUUnquantizedExperts(mk.FusedMoEExpertsMonolithic): ) return True, None - def __init__( - self, - moe_config: FusedMoEConfig, - quant_config: FusedMoEQuantConfig, - ): - super().__init__(moe_config, quant_config) - # Router configuration that the monolithic apply() signature cannot - # carry. Captured off the layer in process_weights_after_loading. - self.use_grouped_topk = False - self.renormalize = False - self.scoring_func = "softmax" - self.custom_routing_function: Callable | None = None - self._use_zentorch = False - @property def expects_unquantized_inputs(self) -> bool: return True @@ -286,8 +167,8 @@ class CPUUnquantizedExperts(mk.FusedMoEExpertsMonolithic): weight_key: QuantKey | None, activation_key: QuantKey | None, ) -> bool: - # Routing runs in select_experts(), which covers every routing method - # a layer can be configured with, including custom routing functions. + # Routing runs in the CPURouter, which covers every routing method a + # layer can be configured with, including custom routing functions. return True @staticmethod @@ -298,10 +179,6 @@ class CPUUnquantizedExperts(mk.FusedMoEExpertsMonolithic): return True def process_weights_after_loading(self, layer: torch.nn.Module) -> None: - self.use_grouped_topk = layer.use_grouped_topk - self.renormalize = layer.renormalize - self.scoring_func = layer.scoring_func - self.custom_routing_function = layer.custom_routing_function self._use_zentorch = is_zentorch_moe_supported(layer) if self._use_zentorch: return @@ -345,65 +222,44 @@ class CPUUnquantizedExperts(mk.FusedMoEExpertsMonolithic): # this very object, which is what feeds self.w1_bias in apply(). layer.w13_bias.data = new_bias - def _select_experts( + def workspace_shapes( self, - hidden_states: torch.Tensor, - router_logits: torch.Tensor, - num_expert_group: int | None, - topk_group: int | None, - e_score_correction_bias: torch.Tensor | None, - routed_scaling_factor: float | None, - ) -> tuple[torch.Tensor, torch.Tensor]: - return select_experts( - hidden_states=hidden_states, - router_logits=router_logits, - top_k=self.moe_config.experts_per_token, - use_grouped_topk=self.use_grouped_topk, - renormalize=self.renormalize, - topk_group=topk_group, - num_expert_group=num_expert_group, - custom_routing_function=self.custom_routing_function, - scoring_func=self.scoring_func, - routed_scaling_factor=( - routed_scaling_factor if routed_scaling_factor is not None else 1.0 - ), - e_score_correction_bias=e_score_correction_bias, - ) + M: int, + N: int, + K: int, + topk: int, + global_num_experts: int, + local_num_experts: int, + expert_tokens_meta: mk.ExpertTokensMetadata | None, + activation: MoEActivation, + ) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...]]: + # cpu_fused_moe manages its own scratch space. + return (0,), (0,), (M, K) + + def finalize_weight_and_reduce_impl(self) -> mk.TopKWeightAndReduce: + return TopKWeightAndReduceNoOP() def apply( self, + output: torch.Tensor, hidden_states: torch.Tensor, w1: torch.Tensor, w2: torch.Tensor, - router_logits: torch.Tensor, + topk_weights: torch.Tensor, + topk_ids: torch.Tensor, activation: MoEActivation, global_num_experts: int, expert_map: torch.Tensor | None, a1q_scale: torch.Tensor | None, + a2_scale: torch.Tensor | None, + workspace13: torch.Tensor, + workspace2: torch.Tensor, + expert_tokens_meta: mk.ExpertTokensMetadata | None, apply_router_weight_on_input: bool, - # grouped topk + fused topk bias parameters - num_expert_group: int | None = None, - e_score_correction_bias: torch.Tensor | None = None, - routed_scaling_factor: float | None = None, - topk_group: int | None = None, - ) -> torch.Tensor: - topk_weights, topk_ids = self._select_experts( - hidden_states, - router_logits, - num_expert_group, - topk_group, - e_score_correction_bias, - routed_scaling_factor, - ) - - if apply_router_weight_on_input: - assert topk_ids.size(1) == 1, ( - "apply_router_weight_on_input is only implemented for topk=1" - ) - hidden_states.mul_(topk_weights.to(hidden_states.dtype)) - + ) -> None: + # apply_router_weight_on_input (topk=1 only) is already applied to + # hidden_states by MoEPrepareAndFinalizeNoDPEPModular.prepare(). if self._use_zentorch: - output = torch.empty_like(hidden_states) torch.ops.zentorch.zentorch_fused_moe( output, hidden_states, @@ -416,9 +272,9 @@ class CPUUnquantizedExperts(mk.FusedMoEExpertsMonolithic): apply_router_weight_on_input, str(activation.value).lower(), ) - return output - - return cpu_fused_moe( + return + cpu_fused_moe( + output, hidden_states, w1, w2, @@ -460,8 +316,6 @@ class X86CPUUnquantizedExperts(CPUUnquantizedExperts): ) if not supported: return supported, reason - if is_zentorch_moe_config_supported(moe_config): - return True, None if moe_config.in_dtype != torch.bfloat16: return False, "kernel requires bfloat16 activations" cpu_cls = cast(type[CPUUnquantizedExperts], cls) @@ -496,8 +350,6 @@ class ArmCPUUnquantizedExperts(CPUUnquantizedExperts): ) if not supported: return supported, reason - if is_zentorch_moe_config_supported(moe_config): - return True, None if moe_config.in_dtype != torch.bfloat16: return False, "kernel requires bfloat16 activations" cpu_cls = cast(type[CPUUnquantizedExperts], cls) @@ -552,8 +404,8 @@ def prepare_fp8_moe_layer_for_cpu( return packed_w13, packed_w2 -class CPUExpertsFp8(mk.FusedMoEExpertsMonolithic): - """CPU FP8 W8A16 block-quantized monolithic MoE experts.""" +class CPUExpertsFp8(mk.FusedMoEExpertsModular): + """CPU FP8 W8A16 block-quantized modular MoE experts.""" def __init__( self, @@ -615,6 +467,7 @@ class CPUExpertsFp8(mk.FusedMoEExpertsMonolithic): RoutingMethodType.Default, RoutingMethodType.Renormalize, RoutingMethodType.RenormalizeNaive, + RoutingMethodType.DeepseekV4, ] @staticmethod @@ -624,42 +477,49 @@ class CPUExpertsFp8(mk.FusedMoEExpertsMonolithic): ) -> bool: return True + def workspace_shapes( + self, + M: int, + N: int, + K: int, + topk: int, + global_num_experts: int, + local_num_experts: int, + expert_tokens_meta: mk.ExpertTokensMetadata | None, + activation: MoEActivation, + ) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...]]: + # fused_experts_cpu manages its own scratch space. + return (0,), (0,), (M, K) + + def finalize_weight_and_reduce_impl(self) -> mk.TopKWeightAndReduce: + return TopKWeightAndReduceNoOP() + def apply( self, + output: torch.Tensor, hidden_states: torch.Tensor, w1: torch.Tensor, w2: torch.Tensor, - router_logits: torch.Tensor, + topk_weights: torch.Tensor, + topk_ids: torch.Tensor, activation: MoEActivation, global_num_experts: int, expert_map: torch.Tensor | None, a1q_scale: torch.Tensor | None, + a2_scale: torch.Tensor | None, + workspace13: torch.Tensor, + workspace2: torch.Tensor, + expert_tokens_meta: mk.ExpertTokensMetadata | None, apply_router_weight_on_input: bool, - # grouped topk + fused topk bias parameters - num_expert_group: int | None = None, - e_score_correction_bias: torch.Tensor | None = None, - routed_scaling_factor: float | None = None, - topk_group: int | None = None, - ) -> torch.Tensor: - topk_weights, topk_ids = select_experts( - hidden_states=hidden_states, - router_logits=router_logits, - use_grouped_topk=num_expert_group is not None, - top_k=self.moe_config.experts_per_token, - renormalize=self.moe_config.routing_method - in ( - RoutingMethodType.Renormalize, - RoutingMethodType.RenormalizeNaive, - ), - topk_group=topk_group, - num_expert_group=num_expert_group, - scoring_func="softmax", - routed_scaling_factor=( - routed_scaling_factor if routed_scaling_factor is not None else 1.0 - ), - e_score_correction_bias=e_score_correction_bias, - ) - + ) -> None: + if apply_router_weight_on_input: + # fused_experts_cpu always applies topk_weights internally on + # combine; MoEPrepareAndFinalizeNoDPEPModular.prepare() would + # also pre-apply it to hidden_states, double-weighting the + # output. Not needed by any CPU FP8 model today. + raise NotImplementedError( + "CPUExpertsFp8 does not support apply_router_weight_on_input=True." + ) block_shape = ( list(self.quant_config.block_shape) if self.quant_config.block_shape @@ -670,13 +530,13 @@ class CPUExpertsFp8(mk.FusedMoEExpertsMonolithic): ) ) - return fused_experts_cpu( + fused_experts_cpu( + output, hidden_states, w1, w2, topk_weights, topk_ids, - False, # inplace CPUQuantMethod.FP8_W8A16, # moe_comp_method self.w1_scale, # w1_scale self.w2_scale, # w2_scale @@ -710,8 +570,8 @@ def prepare_mxfp4_moe_layer_for_cpu( return packed_w13, packed_w2, packed_w13_scale, packed_w2_scale -class CPUExpertsMxfp4(mk.FusedMoEExpertsMonolithic): - """CPU MXFP4 W4A16 monolithic MoE experts.""" +class CPUExpertsMxfp4(mk.FusedMoEExpertsModular): + """CPU MXFP4 W4A16 modular MoE experts.""" def __init__( self, @@ -773,6 +633,7 @@ class CPUExpertsMxfp4(mk.FusedMoEExpertsMonolithic): RoutingMethodType.Default, RoutingMethodType.Renormalize, RoutingMethodType.RenormalizeNaive, + RoutingMethodType.DeepseekV4, ] @staticmethod @@ -782,55 +643,62 @@ class CPUExpertsMxfp4(mk.FusedMoEExpertsMonolithic): ) -> bool: return True + def workspace_shapes( + self, + M: int, + N: int, + K: int, + topk: int, + global_num_experts: int, + local_num_experts: int, + expert_tokens_meta: mk.ExpertTokensMetadata | None, + activation: MoEActivation, + ) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...]]: + # fused_experts_cpu manages its own scratch space. + return (0,), (0,), (M, K) + + def finalize_weight_and_reduce_impl(self) -> mk.TopKWeightAndReduce: + return TopKWeightAndReduceNoOP() + def apply( self, + output: torch.Tensor, hidden_states: torch.Tensor, w1: torch.Tensor, w2: torch.Tensor, - router_logits: torch.Tensor, + topk_weights: torch.Tensor, + topk_ids: torch.Tensor, activation: MoEActivation, global_num_experts: int, expert_map: torch.Tensor | None, a1q_scale: torch.Tensor | None, + a2_scale: torch.Tensor | None, + workspace13: torch.Tensor, + workspace2: torch.Tensor, + expert_tokens_meta: mk.ExpertTokensMetadata | None, apply_router_weight_on_input: bool, - # grouped topk + fused topk bias parameters - num_expert_group: int | None = None, - e_score_correction_bias: torch.Tensor | None = None, - routed_scaling_factor: float | None = None, - topk_group: int | None = None, - ) -> torch.Tensor: - topk_weights, topk_ids = select_experts( - hidden_states=hidden_states, - router_logits=router_logits, - use_grouped_topk=num_expert_group is not None, - top_k=self.moe_config.experts_per_token, - renormalize=self.moe_config.routing_method - in ( - RoutingMethodType.Renormalize, - RoutingMethodType.RenormalizeNaive, - ), - topk_group=topk_group, - num_expert_group=num_expert_group, - scoring_func="softmax", - routed_scaling_factor=( - routed_scaling_factor if routed_scaling_factor is not None else 1.0 - ), - e_score_correction_bias=e_score_correction_bias, - ) - + ) -> None: + if apply_router_weight_on_input: + # fused_experts_cpu always applies topk_weights internally on + # combine; MoEPrepareAndFinalizeNoDPEPModular.prepare() would + # also pre-apply it to hidden_states, double-weighting the + # output. Not needed by any CPU MXFP4 model today. + raise NotImplementedError( + "CPUExpertsMxfp4 does not support apply_router_weight_on_input=True." + ) # Get bias and swiglu params from quant config w1_bias = self.quant_config.w1_bias w2_bias = self.quant_config.w2_bias alpha = getattr(self.quant_config, "gemm1_alpha", None) limit = getattr(self.quant_config, "gemm1_clamp_limit", None) - return fused_experts_cpu( + fused_experts_cpu( + output, hidden_states, w1, w2, topk_weights, topk_ids, - False, # inplace CPUQuantMethod.MXFP4, # moe_comp_method self.w1_scale, # w1_scale self.w2_scale, # w2_scale @@ -911,8 +779,8 @@ def prepare_int4_moe_layer_for_cpu( return (blocked_w13, blocked_w2, blocked_s13, blocked_s2, blocked_z13, blocked_z2) -class CPUExpertsInt4(mk.FusedMoEExpertsMonolithic): - """CPU INT4 W4A16 group-quantized monolithic MoE experts. +class CPUExpertsInt4(mk.FusedMoEExpertsModular): + """CPU INT4 W4A16 group-quantized modular MoE experts. Weights are int4 (packed), activations are bf16/fp16. Internally uses int8 compute via fused_experts_cpu with INT4_W4A8. @@ -992,55 +860,75 @@ class CPUExpertsInt4(mk.FusedMoEExpertsMonolithic): def supports_expert_map(self) -> bool: return False + def moe_problem_size( + self, + a1: torch.Tensor, + w1: torch.Tensor, + w2: torch.Tensor, + topk_ids: torch.Tensor, + ) -> tuple[int, int, int, int, int]: + # convert_weight_packed_scale_zp blocks w1/w2 into AMX tiles + # (E, Nc, Kc, buffer_bytes), not the (E, N, K) layout the base + # implementation assumes -- N/K aren't recoverable from that shape, + # so read them from moe_config instead. + E = w1.shape[0] + K = a1.size(-1) + N = ( + self.moe_config.intermediate_size_per_partition + * self.moe_config.w13_num_shards + ) + M = a1.size(0) if a1.dim() == 2 else a1.size(1) + topk = topk_ids.size(1) + return E, M, N, K, topk + + def workspace_shapes( + self, + M: int, + N: int, + K: int, + topk: int, + global_num_experts: int, + local_num_experts: int, + expert_tokens_meta: mk.ExpertTokensMetadata | None, + activation: MoEActivation, + ) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...]]: + # fused_experts_cpu manages its own scratch space. + return (0,), (0,), (M, K) + + def finalize_weight_and_reduce_impl(self) -> mk.TopKWeightAndReduce: + return TopKWeightAndReduceNoOP() + def apply( self, + output: torch.Tensor, hidden_states: torch.Tensor, w1: torch.Tensor, w2: torch.Tensor, - router_logits: torch.Tensor, + topk_weights: torch.Tensor, + topk_ids: torch.Tensor, activation: MoEActivation, global_num_experts: int, expert_map: torch.Tensor | None, a1q_scale: torch.Tensor | None, + a2_scale: torch.Tensor | None, + workspace13: torch.Tensor, + workspace2: torch.Tensor, + expert_tokens_meta: mk.ExpertTokensMetadata | None, apply_router_weight_on_input: bool, - # grouped topk + fused topk bias parameters - num_expert_group: int | None = None, - e_score_correction_bias: torch.Tensor | None = None, - routed_scaling_factor: float | None = None, - topk_group: int | None = None, - ) -> torch.Tensor: + ) -> None: if apply_router_weight_on_input: raise NotImplementedError( "CPUExpertsInt4 (W4A16) does not support " "apply_router_weight_on_input=True. " ) - topk_weights, topk_ids = select_experts( - hidden_states=hidden_states, - router_logits=router_logits, - use_grouped_topk=num_expert_group is not None, - top_k=self.moe_config.experts_per_token, - renormalize=self.moe_config.routing_method - in ( - RoutingMethodType.Renormalize, - RoutingMethodType.RenormalizeNaive, - ), - topk_group=topk_group, - num_expert_group=num_expert_group, - scoring_func="softmax", - routed_scaling_factor=( - routed_scaling_factor if routed_scaling_factor is not None else 1.0 - ), - e_score_correction_bias=e_score_correction_bias, - ) - - return fused_experts_cpu( + fused_experts_cpu( + output, hidden_states, w1, w2, topk_weights, topk_ids, - False, # inplace CPUQuantMethod.INT4_W4A8, self.w1_scale, self.w2_scale, @@ -1060,9 +948,9 @@ class CPUExpertsInt4(mk.FusedMoEExpertsMonolithic): # =========================================================================== -class CPUExpertsInt8(mk.FusedMoEExpertsMonolithic): +class CPUExpertsInt8(mk.FusedMoEExpertsModular): """CPU INT8 W8A8 per-channel weight / dynamic per-token activation - monolithic MoE experts.""" + modular MoE experts.""" def __init__( self, @@ -1172,49 +1060,56 @@ class CPUExpertsInt8(mk.FusedMoEExpertsMonolithic): replace_parameter(layer, "w13_weight", w13) replace_parameter(layer, "w2_weight", w2) + def workspace_shapes( + self, + M: int, + N: int, + K: int, + topk: int, + global_num_experts: int, + local_num_experts: int, + expert_tokens_meta: mk.ExpertTokensMetadata | None, + activation: MoEActivation, + ) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...]]: + # fused_experts_cpu manages its own scratch space. + return (0,), (0,), (M, K) + + def finalize_weight_and_reduce_impl(self) -> mk.TopKWeightAndReduce: + return TopKWeightAndReduceNoOP() + def apply( self, + output: torch.Tensor, hidden_states: torch.Tensor, w1: torch.Tensor, w2: torch.Tensor, - router_logits: torch.Tensor, + topk_weights: torch.Tensor, + topk_ids: torch.Tensor, activation: MoEActivation, global_num_experts: int, expert_map: torch.Tensor | None, a1q_scale: torch.Tensor | None, + a2_scale: torch.Tensor | None, + workspace13: torch.Tensor, + workspace2: torch.Tensor, + expert_tokens_meta: mk.ExpertTokensMetadata | None, apply_router_weight_on_input: bool, - # grouped topk + fused topk bias parameters - num_expert_group: int | None = None, - e_score_correction_bias: torch.Tensor | None = None, - routed_scaling_factor: float | None = None, - topk_group: int | None = None, - ) -> torch.Tensor: - topk_weights, topk_ids = select_experts( - hidden_states=hidden_states, - router_logits=router_logits, - use_grouped_topk=num_expert_group is not None, - top_k=self.moe_config.experts_per_token, - renormalize=self.moe_config.routing_method - in ( - RoutingMethodType.Renormalize, - RoutingMethodType.RenormalizeNaive, - ), - topk_group=topk_group, - num_expert_group=num_expert_group, - scoring_func="softmax", - routed_scaling_factor=( - routed_scaling_factor if routed_scaling_factor is not None else 1.0 - ), - e_score_correction_bias=e_score_correction_bias, - ) - - return fused_experts_cpu( + ) -> None: + if apply_router_weight_on_input: + # fused_experts_cpu always applies topk_weights internally on + # combine; MoEPrepareAndFinalizeNoDPEPModular.prepare() would + # also pre-apply it to hidden_states, double-weighting the + # output. Not needed by any CPU INT8 W8A8 model today. + raise NotImplementedError( + "CPUExpertsInt8 does not support apply_router_weight_on_input=True." + ) + fused_experts_cpu( + output, hidden_states, w1, w2, topk_weights, topk_ids, - False, # inplace CPUQuantMethod.INT8_W8A8, self.w1_scale, self.w2_scale, @@ -1229,7 +1124,7 @@ class CPUExpertsInt8(mk.FusedMoEExpertsMonolithic): ) -class ArmCPUExpertsInt8(mk.FusedMoEExpertsMonolithic): +class ArmCPUExpertsInt8(mk.FusedMoEExpertsModular): """Arm INT8 MoE with per-token activation and channelwise weight quantization.""" @property @@ -1331,48 +1226,47 @@ class ArmCPUExpertsInt8(mk.FusedMoEExpertsMonolithic): replace_parameter(layer, "w13_weight", w13) replace_parameter(layer, "w2_weight", w2) + def workspace_shapes( + self, + M: int, + N: int, + K: int, + topk: int, + global_num_experts: int, + local_num_experts: int, + expert_tokens_meta: mk.ExpertTokensMetadata | None, + activation: MoEActivation, + ) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...]]: + # cpu_fused_moe_int8 manages its own scratch space. + return (0,), (0,), (M, K) + + def finalize_weight_and_reduce_impl(self) -> mk.TopKWeightAndReduce: + return TopKWeightAndReduceNoOP() + def apply( self, + output: torch.Tensor, hidden_states: torch.Tensor, w1: torch.Tensor, w2: torch.Tensor, - router_logits: torch.Tensor, + topk_weights: torch.Tensor, + topk_ids: torch.Tensor, activation: MoEActivation, global_num_experts: int, expert_map: torch.Tensor | None, a1q_scale: torch.Tensor | None, + a2_scale: torch.Tensor | None, + workspace13: torch.Tensor, + workspace2: torch.Tensor, + expert_tokens_meta: mk.ExpertTokensMetadata | None, apply_router_weight_on_input: bool, - num_expert_group: int | None = None, - e_score_correction_bias: torch.Tensor | None = None, - routed_scaling_factor: float | None = None, - topk_group: int | None = None, - ) -> torch.Tensor: - topk_weights, topk_ids = select_experts( - hidden_states=hidden_states, - router_logits=router_logits, - use_grouped_topk=num_expert_group is not None, - top_k=self.moe_config.experts_per_token, - renormalize=self.moe_config.routing_method - in ( - RoutingMethodType.Renormalize, - RoutingMethodType.RenormalizeNaive, - ), - topk_group=topk_group, - num_expert_group=num_expert_group, - scoring_func="softmax", - routed_scaling_factor=( - routed_scaling_factor if routed_scaling_factor is not None else 1.0 - ), - e_score_correction_bias=e_score_correction_bias, - ) - - if apply_router_weight_on_input: - assert topk_ids.size(1) == 1 - hidden_states.mul_(topk_weights.to(hidden_states.dtype)) - + ) -> None: + # apply_router_weight_on_input (topk=1 only) is already applied to + # hidden_states by MoEPrepareAndFinalizeNoDPEPModular.prepare(). assert self.w1_scale is not None assert self.w2_scale is not None - return cpu_fused_moe_int8( + cpu_fused_moe_int8( + output, hidden_states, w1, w2, @@ -1388,7 +1282,7 @@ class ArmCPUExpertsInt8(mk.FusedMoEExpertsMonolithic): ) -class ZenCPUExpertsInt8(mk.FusedMoEExpertsMonolithic): +class ZenCPUExpertsInt8(mk.FusedMoEExpertsModular): """AMD Zen INT8 MoE with per-token activation and channelwise weight quantization, dispatched through zentorch.""" @@ -1424,9 +1318,7 @@ class ZenCPUExpertsInt8(mk.FusedMoEExpertsMonolithic): if self.w2_bias is None else self.w2_bias.detach().to(torch.bfloat16).contiguous() ) - logger.info_once( - "[zen_cpu] Using zentorch_fused_moe for W8A8 INT8 MoE (monolithic experts)" - ) + logger.info_once("[zen_cpu] Using zentorch_fused_moe for W8A8 INT8 MoE") @property def expects_unquantized_inputs(self) -> bool: @@ -1487,48 +1379,47 @@ class ZenCPUExpertsInt8(mk.FusedMoEExpertsMonolithic): def supports_expert_map(self) -> bool: return False + def workspace_shapes( + self, + M: int, + N: int, + K: int, + topk: int, + global_num_experts: int, + local_num_experts: int, + expert_tokens_meta: mk.ExpertTokensMetadata | None, + activation: MoEActivation, + ) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...]]: + # zentorch_fused_moe manages its own scratch space. + return (0,), (0,), (M, K) + + def finalize_weight_and_reduce_impl(self) -> mk.TopKWeightAndReduce: + return TopKWeightAndReduceNoOP() + def apply( self, + output: torch.Tensor, hidden_states: torch.Tensor, w1: torch.Tensor, w2: torch.Tensor, - router_logits: torch.Tensor, + topk_weights: torch.Tensor, + topk_ids: torch.Tensor, activation: MoEActivation, global_num_experts: int, expert_map: torch.Tensor | None, a1q_scale: torch.Tensor | None, + a2_scale: torch.Tensor | None, + workspace13: torch.Tensor, + workspace2: torch.Tensor, + expert_tokens_meta: mk.ExpertTokensMetadata | None, apply_router_weight_on_input: bool, - num_expert_group: int | None = None, - e_score_correction_bias: torch.Tensor | None = None, - routed_scaling_factor: float | None = None, - topk_group: int | None = None, - ) -> torch.Tensor: - topk_weights, topk_ids = select_experts( - hidden_states=hidden_states, - router_logits=router_logits, - use_grouped_topk=num_expert_group is not None, - top_k=self.moe_config.experts_per_token, - renormalize=self.moe_config.routing_method - in ( - RoutingMethodType.Renormalize, - RoutingMethodType.RenormalizeNaive, - ), - topk_group=topk_group, - num_expert_group=num_expert_group, - scoring_func="softmax", - routed_scaling_factor=( - routed_scaling_factor if routed_scaling_factor is not None else 1.0 - ), - e_score_correction_bias=e_score_correction_bias, - ) - + ) -> None: if apply_router_weight_on_input: assert topk_ids.size(1) == 1, ( "apply_router_weight_on_input is only implemented for topk=1" ) hidden_states = hidden_states.mul(topk_weights.to(hidden_states.dtype)) - output = torch.empty_like(hidden_states) torch.ops.zentorch.zentorch_fused_moe( output, hidden_states, @@ -1543,4 +1434,3 @@ class ZenCPUExpertsInt8(mk.FusedMoEExpertsMonolithic): self._w1_scale_bf16, self._w2_scale_bf16, ) - return output diff --git a/vllm/model_executor/layers/fused_moe/modular_kernel.py b/vllm/model_executor/layers/fused_moe/modular_kernel.py index d7446a893a..778b95bf47 100644 --- a/vllm/model_executor/layers/fused_moe/modular_kernel.py +++ b/vllm/model_executor/layers/fused_moe/modular_kernel.py @@ -1175,6 +1175,24 @@ class FusedMoEKernelModularImpl: # time we need cache3, we're done with cache1. # Reuse workspace13 for the output since there is only one chunk. max_shape_size = max(prod(workspace13_shape), prod(fused_out_shape)) + + if current_platform.is_cpu(): + # Every CPU FusedMoEExpertsModular kernel reports zero-sized + # workspace13/workspace2 (it manages its own scratch space + # internally) and just needs the output buffer, so skip the + # workspace manager entirely here -- it's the one part of this + # call graph Dynamo can't trace (ContextVar-based lane lookup), + # and CPU never actually needs its cross-chunk memory reuse. + common_workspace = torch.empty( + (max_shape_size,), dtype=workspace_dtype, device=device + ) + workspace2 = torch.empty( + workspace2_shape, dtype=workspace_dtype, device=device + ) + workspace13 = _resize_cache(common_workspace, workspace13_shape) + fused_out = _resize_cache(common_workspace, fused_out_shape) + return workspace13, workspace2, fused_out + common_workspace, workspace2 = current_workspace_manager().get_simultaneous( ((max_shape_size,), workspace_dtype), (workspace2_shape, workspace_dtype), diff --git a/vllm/model_executor/layers/fused_moe/oracle/mxfp4.py b/vllm/model_executor/layers/fused_moe/oracle/mxfp4.py index 1fa2f63e5d..3b415acd81 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/mxfp4.py +++ b/vllm/model_executor/layers/fused_moe/oracle/mxfp4.py @@ -348,6 +348,7 @@ def _get_priority_backends_for_gpt_oss() -> list[Mxfp4MoeBackend]: Mxfp4MoeBackend.MARLIN, Mxfp4MoeBackend.BATCHED_MARLIN, Mxfp4MoeBackend.XPU, + Mxfp4MoeBackend.CPU, Mxfp4MoeBackend.EMULATION, ] return _AVAILABLE_BACKENDS @@ -366,6 +367,8 @@ def _get_priority_backends() -> list[Mxfp4MoeBackend]: ] if current_platform.is_xpu(): return [Mxfp4MoeBackend.XPU] + if current_platform.is_cpu(): + return [Mxfp4MoeBackend.CPU] _AVAILABLE_BACKENDS = [ Mxfp4MoeBackend.FLASHINFER_TRTLLM_MXFP4_MXFP8, Mxfp4MoeBackend.DEEPGEMM_MXFP4, @@ -627,17 +630,6 @@ def select_mxfp4_moe_backend( activation_format, ) - if current_platform.is_cpu(): - backend = Mxfp4MoeBackend.CPU - logger.info_once(_make_log_backend(backend)) - return _return_or_raise( - Mxfp4MoeBackend.CPU, - config, - kMxfp4Static, - None, - activation_format, - ) - unsupported_log = "; ".join( [ f"backend: {backend.value}, reason: {reason}" @@ -1813,11 +1805,36 @@ def convert_weight_to_mxfp4_moe_kernel_format( w2_bias=w2_bias, _cache_permute_indices=_cache_permute_indices, ) + elif mxfp4_backend == Mxfp4MoeBackend.CPU: + from vllm.model_executor.layers.fused_moe.experts.cpu_moe import ( + prepare_mxfp4_moe_layer_for_cpu, + ) + + packed_w13, packed_w2, packed_w13_scale, packed_w2_scale = ( + prepare_mxfp4_moe_layer_for_cpu( + w13_weight.data, + w2_weight.data, + w13_weight_scale.data, + w2_weight_scale.data, + ) + ) + if w13_bias is not None: + w13_bias = w13_bias.data.to(torch.float32) + if w2_bias is not None: + w2_bias = w2_bias.data.to(torch.float32) + return ( + packed_w13, + packed_w2, + packed_w13_scale, + packed_w2_scale, + w13_bias, + w2_bias, + ) else: raise ValueError( f"Unsupported mxfp4_backend for Mxfp4MoEMethod: {mxfp4_backend}. " - "Expected TRTLLM, FlashInfer CUTLASS, Triton, AITER, XPU, or " - "emulation backend." + "Expected TRTLLM, FlashInfer CUTLASS, Triton, AITER, XPU, " + "CPU, or emulation backend." ) diff --git a/vllm/model_executor/layers/fused_moe/router/cpu_router.py b/vllm/model_executor/layers/fused_moe/router/cpu_router.py new file mode 100644 index 0000000000..1f96e9107a --- /dev/null +++ b/vllm/model_executor/layers/fused_moe/router/cpu_router.py @@ -0,0 +1,294 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Router for CPU MoE experts. + +This is the routing logic that used to live inline in each CPU +FusedMoEExpertsMonolithic.apply() (see git history of +experts/cpu_moe.py); factoring it out into a proper FusedMoERouter is what +lets CPU MoE experts be FusedMoEExpertsModular like every other backend. +The math below is unchanged from that prior inline version -- this is a +relocation, not a rewrite. +""" + +from collections.abc import Callable + +import torch + +from vllm._custom_ops import biased_topk_cpu, hash_topk_cpu +from vllm.distributed.eplb.eplb_state import EplbLayerState +from vllm.model_executor.layers.fused_moe.config import ( + RoutingMethodType, + get_routing_method_type, +) +from vllm.model_executor.layers.fused_moe.router.base_router import BaseRouter + + +def _grouped_topk( + gating_output: torch.Tensor, + topk: int, + renormalize: bool, + num_expert_group: int, + topk_group: int, + scoring_func: str, + routed_scaling_factor: float, + e_score_correction_bias: torch.Tensor | None, +) -> tuple[torch.Tensor, torch.Tensor]: + gating_output = gating_output.float() + if scoring_func == "softmax": + scores = torch.softmax(gating_output, dim=-1) + elif scoring_func == "sigmoid": + scores = gating_output.sigmoid() + else: + raise ValueError(f"Unsupported scoring function: {scoring_func}") + + num_token = scores.shape[0] + if e_score_correction_bias is not None: + original_scores = scores + scores = scores + e_score_correction_bias.unsqueeze(0) + group_scores = ( + scores.view(num_token, num_expert_group, -1).topk(2, dim=-1)[0].sum(dim=-1) + ) + else: + group_scores = ( + scores.view(num_token, num_expert_group, -1).max(dim=-1).values + ) # [n, n_group] + group_idx = torch.topk(group_scores, k=topk_group, dim=-1, sorted=False)[ + 1 + ] # [n, top_k_group] + group_mask = torch.zeros_like(group_scores) # [n, n_group] + group_mask.scatter_(1, group_idx, 1) # [n, n_group] + score_mask = ( + group_mask.unsqueeze(-1) + .expand(num_token, num_expert_group, scores.shape[-1] // num_expert_group) + .reshape(num_token, -1) + ) # [n, e] + tmp_scores = scores.masked_fill(~score_mask.bool(), float("-inf")) # [n, e] + + if e_score_correction_bias is not None: + topk_ids = torch.topk(tmp_scores, k=topk, dim=-1, sorted=False)[1] + topk_weights = original_scores.gather(1, topk_ids) + else: + topk_weights, topk_ids = torch.topk(tmp_scores, k=topk, dim=-1, sorted=False) + + if renormalize: + topk_weights = topk_weights / topk_weights.sum(dim=-1, keepdim=True) + + if routed_scaling_factor != 1.0: + topk_weights = topk_weights * routed_scaling_factor + return topk_weights, topk_ids.to(torch.int32) + + +def _sqrtsoftplus_bias_topk( + router_logits: torch.Tensor, + top_k: int, + renormalize: bool, + e_score_correction_bias: torch.Tensor | None, + routed_scaling_factor: float, + input_ids: torch.Tensor | None, + hash_indices_table: torch.Tensor | None, +) -> tuple[torch.Tensor, torch.Tensor]: + """DeepSeek V4's routing: weight = sqrt(softplus(logit)); expert + selection uses a bias-corrected score (weight + correction_bias), but + the returned weight for each selected expert is the *unbiased* value + (the "noaux_tc" pattern -- bias steers selection only). Always + renormalizes selected weights to sum to ``routed_scaling_factor``, + matching the reference Triton kernel + (vllm/model_executor/layers/fused_moe/router/dsv4_topk.py). + + When ``hash_indices_table`` is given (hash-routed layers), expert + selection instead comes from ``hash_indices_table[input_ids]`` -- + the bias-corrected-score topk is skipped entirely, matching the CUDA + reference kernel (``dsv4HashTopkSoftplusSqrt``). + + Backed by the ported SGLang AVX512 kernels `biased_topk_cpu` (the + correction-bias path) and `hash_topk_cpu` (the hash-routed-layer path); + see csrc/cpu/sgl-kernels/topk.cpp. + """ + assert renormalize, "DeepseekV4 routing always renormalizes" + gating_output = router_logits.contiguous() + num_experts = gating_output.shape[1] + if hash_indices_table is not None: + assert input_ids is not None + # hash_indices_table is int32 by construction (see cpu/model.py) and + # advanced indexing always allocates a fresh contiguous tensor. + tid2eid = hash_indices_table[input_ids] + return hash_topk_cpu( + gating_output, + tid2eid, + top_k, + "sqrtsoftplus", + 0, # num_fused_shared_experts: never fused into routing on CPU + num_experts, + routed_scaling_factor, + ) + assert e_score_correction_bias is not None + return biased_topk_cpu( + gating_output, # only used for num_tokens/device (dtype forced fp32) + gating_output, + e_score_correction_bias.contiguous(), + top_k, + renormalize, + "sqrtsoftplus", + 0, # num_fused_shared_experts: never fused into routing on CPU + routed_scaling_factor, + True, # apply_routed_scaling_factor_on_output + ) + + +def _softmax_topk( + router_logits: torch.Tensor, + top_k: int, + renormalize: bool, +) -> tuple[torch.Tensor, torch.Tensor]: + topk_logit_vals, topk_idx = torch.topk(router_logits, k=top_k, dim=-1, sorted=False) + if renormalize: + topk_vals = torch.softmax(topk_logit_vals, dim=-1) + else: + logZ = torch.logsumexp(router_logits, dim=-1, keepdim=True) + topk_vals = (topk_logit_vals - logZ).exp() + return topk_vals.to(torch.float32), topk_idx.to(torch.int32) + + +class CPURouter(BaseRouter): + """Router covering every routing scheme CPU MoE experts support: plain + softmax top-k, grouped top-k (with an optional correction bias), an + arbitrary custom_routing_function, and DeepSeek V4's sqrtsoftplus scheme + (bias-corrected or hash-routed via the ported biased_topk_cpu/ + hash_topk_cpu kernels). All of the above except sqrtsoftplus are plain + torch ops so Dynamo traces straight through into them like it did when + this logic lived inline in FusedMoEExpertsMonolithic.apply() -- no new + compile boundary is introduced here. + """ + + def __init__( + self, + top_k: int, + global_num_experts: int, + use_grouped_topk: bool = False, + num_expert_group: int | None = None, + topk_group: int | None = None, + renormalize: bool = True, + scoring_func: str = "softmax", + routed_scaling_factor: float = 1.0, + e_score_correction_bias: torch.Tensor | None = None, + custom_routing_function: Callable | None = None, + hash_indices_table: torch.Tensor | None = None, + eplb_state: EplbLayerState | None = None, + ): + super().__init__( + top_k=top_k, + global_num_experts=global_num_experts, + eplb_state=eplb_state, + ) + self.use_grouped_topk = use_grouped_topk + self.num_expert_group = num_expert_group + self.topk_group = topk_group + self.renormalize = renormalize + self.scoring_func = scoring_func + self.routed_scaling_factor = routed_scaling_factor + self.e_score_correction_bias = e_score_correction_bias + self.custom_routing_function = custom_routing_function + self._hash_indices_table = hash_indices_table + + @property + def routing_method_type(self) -> RoutingMethodType: + return get_routing_method_type( + scoring_func=self.scoring_func, + top_k=self.top_k, + renormalize=self.renormalize, + num_expert_group=self.num_expert_group, + has_e_score_bias=self.e_score_correction_bias is not None, + routed_scaling_factor=self.routed_scaling_factor, + ) + + def _compute_routing( + self, + hidden_states: torch.Tensor, + router_logits: torch.Tensor, + indices_type: torch.dtype | None, + *, + input_ids: torch.Tensor | None = None, + ) -> tuple[torch.Tensor, torch.Tensor]: + if self.use_grouped_topk: + assert self.topk_group is not None + assert self.num_expert_group is not None + return _grouped_topk( + gating_output=router_logits, + topk=self.top_k, + renormalize=self.renormalize, + num_expert_group=self.num_expert_group, + topk_group=self.topk_group, + scoring_func=self.scoring_func, + routed_scaling_factor=self.routed_scaling_factor, + e_score_correction_bias=self.e_score_correction_bias, + ) + elif self.scoring_func == "sqrtsoftplus": + assert ( + self._hash_indices_table is not None + or self.e_score_correction_bias is not None + ) + return _sqrtsoftplus_bias_topk( + router_logits=router_logits, + top_k=self.top_k, + renormalize=self.renormalize, + e_score_correction_bias=self.e_score_correction_bias, + routed_scaling_factor=self.routed_scaling_factor, + input_ids=input_ids, + hash_indices_table=self._hash_indices_table, + ) + elif self.custom_routing_function is None: + assert self.scoring_func == "softmax" + return _softmax_topk(router_logits, self.top_k, self.renormalize) + else: + topk_weights, topk_ids = self.custom_routing_function( + hidden_states=hidden_states, + gating_output=router_logits, + topk=self.top_k, + renormalize=self.renormalize, + ) + # cpu_fused_moe reads routing tensors as contiguous float32/int32 + # buffers and does not account for tensor strides. + topk_weights = topk_weights.to(torch.float32).contiguous() + topk_ids = topk_ids.to(torch.int32).contiguous() + return topk_weights, topk_ids + + +def select_experts( + hidden_states: torch.Tensor, + router_logits: torch.Tensor, + top_k: int, + use_grouped_topk: bool, + renormalize: bool, + topk_group: int | None = None, + num_expert_group: int | None = None, + custom_routing_function: Callable | None = None, + scoring_func: str = "softmax", + routed_scaling_factor: float = 1.0, + e_score_correction_bias: torch.Tensor | None = None, + input_ids: torch.Tensor | None = None, + hash_indices_table: torch.Tensor | None = None, +) -> tuple[torch.Tensor, torch.Tensor]: + """Routing helper for the CPU MoE experts that still use the monolithic + contract and so must compute routing themselves from ``router_logits`` + (the zentorch-accelerated path and the ARM W4A8 dynamic-quant kernel, + kept monolithic since zentorch/Arm hardware isn't available to validate + a modular migration here). Modular CPU experts get routing from + ``CPURouter`` instead; this is a thin wrapper around the same logic for + the two classes that can't use it. + """ + router = CPURouter( + top_k=top_k, + global_num_experts=router_logits.shape[-1], + use_grouped_topk=use_grouped_topk, + num_expert_group=num_expert_group, + topk_group=topk_group, + renormalize=renormalize, + scoring_func=scoring_func, + routed_scaling_factor=routed_scaling_factor, + e_score_correction_bias=e_score_correction_bias, + custom_routing_function=custom_routing_function, + hash_indices_table=hash_indices_table, + ) + return router._compute_routing( + hidden_states, router_logits, None, input_ids=input_ids + ) diff --git a/vllm/model_executor/layers/fused_moe/router/router_factory.py b/vllm/model_executor/layers/fused_moe/router/router_factory.py index d80fa840f8..0f6ec15d45 100644 --- a/vllm/model_executor/layers/fused_moe/router/router_factory.py +++ b/vllm/model_executor/layers/fused_moe/router/router_factory.py @@ -14,6 +14,7 @@ from vllm.model_executor.layers.fused_moe.config import ( from vllm.model_executor.layers.fused_moe.router.aiter_shared_routed_fused_moe_router import ( # noqa: E501 AiterSharedRoutedFusedMoERouter, ) +from vllm.model_executor.layers.fused_moe.router.cpu_router import CPURouter from vllm.model_executor.layers.fused_moe.router.custom_routing_router import ( CustomRoutingRouter, ) @@ -35,6 +36,7 @@ from vllm.model_executor.layers.fused_moe.router.routing_simulator_router import from vllm.model_executor.layers.fused_moe.router.zero_expert_router import ( ZeroExpertRouter, ) +from vllm.platforms import current_platform def create_fused_moe_router( @@ -141,6 +143,26 @@ def create_fused_moe_router( routed_scaling_factor=routed_scaling_factor, ) + if current_platform.is_cpu(): + # CPURouter covers every routing scheme CPU MoE experts support + # (plain softmax, grouped topk, custom routing functions, and + # DeepSeek V4's sqrtsoftplus/hash routing) in a single class, so it + # takes priority over the scheme-specific routers below. + return CPURouter( + top_k=top_k, + global_num_experts=global_num_experts, + use_grouped_topk=use_grouped_topk, + num_expert_group=num_expert_group, + topk_group=topk_group, + renormalize=renormalize, + scoring_func=scoring_func, + routed_scaling_factor=routed_scaling_factor, + e_score_correction_bias=e_score_correction_bias, + custom_routing_function=custom_routing_function, + hash_indices_table=hash_indices_table, + eplb_state=eplb_state, + ) + if use_grouped_topk: assert custom_routing_function is None if num_expert_group is None or topk_group is None: diff --git a/vllm/model_executor/layers/fused_moe/runner/moe_runner.py b/vllm/model_executor/layers/fused_moe/runner/moe_runner.py index 2f072d14b4..daaff3b79b 100644 --- a/vllm/model_executor/layers/fused_moe/runner/moe_runner.py +++ b/vllm/model_executor/layers/fused_moe/runner/moe_runner.py @@ -317,10 +317,17 @@ class MoERunner(MoERunnerInterface): return self.routed_experts.load_weights(weights) def _select_forward(self) -> Callable: - if current_platform.is_tpu() or current_platform.is_cpu(): + if current_platform.is_tpu(): # TODO: Once the OOM issue for the TPU backend is resolved, we # will switch to using the moe_forward custom op. - # Note: CPU doesn't require wrapped _forward_impl. + return _moe_forward if self._shared_experts is None else _moe_forward_shared + + if current_platform.is_cpu(): + # CPU never touches the workspace manager (Monolithic experts + # skip it entirely; Modular experts' _allocate_buffers bypasses + # it too, see modular_kernel.py) -- the ContextVar-based lane + # lookup was the only part of this call graph Dynamo can't + # trace, so CPU can always call the fused-MoE op directly. return _moe_forward if self._shared_experts is None else _moe_forward_shared return ( diff --git a/vllm/model_executor/layers/mhc.py b/vllm/model_executor/layers/mhc.py index d4bc224af4..6ddf4ff066 100644 --- a/vllm/model_executor/layers/mhc.py +++ b/vllm/model_executor/layers/mhc.py @@ -261,6 +261,36 @@ class MHCPreOp(CustomOp): sinkhorn_repeat, ) + def forward_cpu( + self, + residual: torch.Tensor, + fn: torch.Tensor, + hc_scale: torch.Tensor, + hc_base: torch.Tensor, + rms_eps: float, + hc_pre_eps: float, + hc_sinkhorn_eps: float, + hc_post_mult_value: float, + sinkhorn_repeat: int, + n_splits: int = 1, + norm_weight: torch.Tensor | None = None, + norm_eps: float = 0.0, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + return mhc_kernels.mhc_pre_cpu( + residual, + fn, + hc_scale, + hc_base, + rms_eps, + hc_pre_eps, + hc_sinkhorn_eps, + hc_post_mult_value, + sinkhorn_repeat, + n_splits, + norm_weight, + norm_eps, + ) + # --8<-- [start:mhc_post] @CustomOp.register("mhc_post") @@ -337,6 +367,15 @@ class MHCPostOp(CustomOp): comb_res_mix, ) + def forward_cpu( + self, + x: torch.Tensor, + residual: torch.Tensor, + post_layer_mix: torch.Tensor, + comb_res_mix: torch.Tensor, + ) -> torch.Tensor: + return mhc_kernels.mhc_post_cpu(x, residual, post_layer_mix, comb_res_mix) + # --8<-- [start:hc_head] @CustomOp.register("hc_head") @@ -444,6 +483,19 @@ class HCHeadOp(CustomOp): ) return out.view(*outer_shape, hidden_size) + def forward_cpu( + self, + hidden_states: torch.Tensor, + hc_fn: torch.Tensor, + hc_scale: torch.Tensor, + hc_base: torch.Tensor, + rms_norm_eps: float, + hc_eps: float, + ) -> torch.Tensor: + return mhc_kernels.hc_head_fused_cpu( + hidden_states, hc_fn, hc_scale, hc_base, rms_norm_eps, hc_eps + ) + # --8<-- [start:mhc_fused_post_pre] @CustomOp.register("mhc_fused_post_pre") @@ -654,6 +706,42 @@ class MHCFusedPostPreOp(CustomOp): sinkhorn_repeat, ) + def forward_cpu( + self, + x: torch.Tensor, + residual: torch.Tensor, + post_layer_mix: torch.Tensor, + comb_res_mix: torch.Tensor, + fn: torch.Tensor, + hc_scale: torch.Tensor, + hc_base: torch.Tensor, + rms_eps: float, + hc_pre_eps: float, + hc_sinkhorn_eps: float, + hc_post_mult_value: float, + sinkhorn_repeat: int, + n_splits: int = 1, + tile_n: int = 1, + norm_weight: torch.Tensor | None = None, + norm_eps: float = 0.0, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + # Decompose into post + pre (no fused kernel available). + residual_cur = mhc_kernels.mhc_post_cpu( + x, residual, post_layer_mix, comb_res_mix + ) + post_mix_cur, comb_mix_cur, layer_input_cur = mhc_kernels.mhc_pre_cpu( + residual_cur, + fn, + hc_scale, + hc_base, + rms_eps, + hc_pre_eps, + hc_sinkhorn_eps, + hc_post_mult_value, + sinkhorn_repeat, + ) + return residual_cur, post_mix_cur, comb_mix_cur, layer_input_cur + def hc_expand(x: torch.Tensor, n: int) -> torch.Tensor: """[s, hidden_size] -> [s, n * hidden_size] by replication.""" diff --git a/vllm/model_executor/layers/sparse_attn_indexer.py b/vllm/model_executor/layers/sparse_attn_indexer.py index 488eb3899c..72a9a2d5a9 100644 --- a/vllm/model_executor/layers/sparse_attn_indexer.py +++ b/vllm/model_executor/layers/sparse_attn_indexer.py @@ -828,10 +828,12 @@ class SparseAttnIndexer(CustomOp): return self.forward_cuda(hidden_states, q_quant, k, weights) elif current_platform.is_rocm(): return self.forward_hip(hidden_states, q_quant, k, weights) + elif current_platform.is_cpu(): + return self.forward_cpu(hidden_states, q_quant, k, weights) else: raise NotImplementedError( "SparseAttnIndexer native forward is only implemented for " - "CUDA, ROCm and XPU platforms." + "CUDA, ROCm, XPU and CPU platforms." ) def forward_cuda( @@ -919,3 +921,196 @@ class SparseAttnIndexer(CustomOp): "Sparse attention indexer ROCm path is only supported on AITER. " "Please enable aiter with VLLM_ROCM_USE_AITER=1" ) + + def forward_cpu( + self, + hidden_states: torch.Tensor, + q_quant: torch.Tensor | tuple[torch.Tensor, torch.Tensor], + k: torch.Tensor | None, + weights: torch.Tensor, + ): + """CPU sparse attention indexer: cache write stays eager Python glue + (own K-cache layout, not shared with the main attention cache + write). PREFILL and DECODE both call the ported + ``fp8_paged_mqa_logits_cpu``/``topk_transform_512_cpu`` kernels, + which read the paged K-cache directly via ``page_table`` -- no + eager gather step, no per-request Python loop. + + ``prefill_metadata.chunks`` always has exactly one entry here: + ``DeepseekV4CPUIndexerMetadataBuilder`` overrides the base chunk + split to always return the whole step's prefill batch as one + chunk, since the base chunking only bounds CUDA/XPU's dense M*N + logits tensor and flat K-gather workspace, neither of which this + paged kernel allocates. + """ + assert not self.use_fp4_cache, ( + "CPU sparse indexer doesn't support fp4 cache yet" + ) + assert isinstance(q_quant, torch.Tensor), ( + "CPU sparse_attn_indexer expects a single FP8 q_quant tensor" + ) + assert self.dcp_world_size <= 1 and not self.use_pcp, ( + "CPU sparse indexer doesn't support decode/prefill context parallelism yet." + ) + + forward_context = get_forward_context() + attn_metadata = forward_context.attn_metadata + attn_metadata_narrowed: DeepseekV32IndexerMetadata | None = None + if isinstance(attn_metadata, dict): + metadata = attn_metadata[self.k_cache.prefix] + assert isinstance(metadata, DeepseekV32IndexerMetadata) + attn_metadata_narrowed = metadata + if attn_metadata_narrowed is None: + # Profiling/dummy run: no real metadata to act on. + return self.topk_indices_buffer + + kv_cache = self.k_cache.kv_cache + topk_tokens = self.topk_tokens + topk_indices_buffer = self.topk_indices_buffer + slot_mapping = attn_metadata_narrowed.slot_mapping + has_decode = attn_metadata_narrowed.num_decodes > 0 + has_prefill = attn_metadata_narrowed.num_prefills > 0 + num_decode_tokens = attn_metadata_narrowed.num_decode_tokens + + num_tokens = slot_mapping.shape[0] + if k is not None: + k = k[:num_tokens] + + if not self.skip_k_cache_insert: + # Only reachable via DeepseekV32Attention with + # prefill_context_parallel_size > 1 on CPU -- set_k_cpu/set_s_cpu + # (the kernels this used to call) have been removed as unused/ + # untested (csrc/cpu/sgl-kernels/store_cache.cpp). + raise NotImplementedError( + "SparseAttnIndexer.forward_cpu: skip_k_cache_insert=False " + "(prefill context parallel on CPU) is not supported." + ) + + topk_indices_buffer[: hidden_states.shape[0]] = -1 + + if has_prefill: + assert topk_tokens == 512, ( + "topk_transform_512_cpu only supports index_topk == 512." + ) + prefill_metadata = attn_metadata_narrowed.prefill + assert prefill_metadata is not None + assert len(prefill_metadata.chunks) == 1, ( + "forward_cpu expects the prefill metadata builder to always " + "produce a single chunk -- see " + "DeepseekV4CPUIndexerMetadataBuilder._split_indexer_prefill_chunks." + ) + chunk = prefill_metadata.chunks[0] + # kv_cache is a per-layer view into vLLM's shared multi-layer + # cache allocation, so its block stride generally exceeds + # block_size * page_width; fp8_paged_mqa_logits_cpu reads + # kv_view.stride(0) explicitly, so no copy is needed here. + kv_view = kv_cache.view(kv_cache.shape[0], -1) + block_size = kv_cache.shape[1] + q_slice = q_quant[chunk.token_start : chunk.token_end] + topk_indices = topk_indices_buffer[ + chunk.token_start : chunk.token_end, :topk_tokens + ] + if chunk.local_total_seq_lens == 0: + topk_indices.fill_(-1) + else: + assert chunk.local_cu_seq_lens is not None + # Each token's own (per-request, DCP-local) causal length. + local_seq_lens = chunk.cu_seqlen_ke - chunk.cu_seqlen_ks + # Recover each token's owning request from its row-start + # tag. Ties (from zero-length requests) are harmless: those + # tokens have local length 0 and never dereference + # page_table. + req_idx = ( + torch.searchsorted( + chunk.local_cu_seq_lens, chunk.cu_seqlen_ks, right=True + ) + - 1 + ) + page_table = chunk.block_table[req_idx] + # The true max over this chunk's own rows, NOT + # chunk.max_local_total_seq_lens (that field sums every + # request's length in the chunk, bounding the old + # flat-gather buffer this paged path no longer allocates). + max_seq_len = int(local_seq_lens.max().item()) + + logits = ops.fp8_paged_mqa_logits_cpu( + q_slice, + kv_view, + weights[chunk.token_start : chunk.token_end], + local_seq_lens, + page_table, + block_size, + max_seq_len, + ) + out_page_scratch = torch.empty( + (q_slice.shape[0], topk_tokens), + dtype=torch.int32, + device=kv_cache.device, + ) + ops.topk_transform_512_cpu( + logits, + local_seq_lens, + page_table, + out_page_scratch, + block_size, + topk_indices, + ) + + if has_decode: + decode_metadata = attn_metadata_narrowed.decode + assert decode_metadata is not None + assert not decode_metadata.requires_padding, ( + "CPU sparse indexer decode path does not support speculative " + "decoding (native MTP) yet." + ) + batch_size = decode_metadata.decode_lens.shape[0] + if batch_size > 0: + # No native MTP on CPU (asserted above) => exactly one + # query token per decode request, so the flat slice below + # is already the batch-major layout + # fp8_paged_mqa_logits_cpu wants. + assert num_decode_tokens == batch_size, ( + "CPU sparse indexer decode path expects exactly one query " + "token per decode request." + ) + assert topk_tokens == 512, ( + "topk_transform_512_cpu only supports index_topk == 512." + ) + seq_lens = decode_metadata.seq_lens + seq_lens = ( + seq_lens[:, -1].contiguous() if seq_lens.ndim == 2 else seq_lens + ) + block_table = decode_metadata.block_table[:batch_size] + block_size = kv_cache.shape[1] + kv_view = kv_cache.view(kv_cache.shape[0], -1) + + logits = ops.fp8_paged_mqa_logits_cpu( + q_quant[:num_decode_tokens], + kv_view, + weights[:num_decode_tokens], + seq_lens, + block_table, + block_size, + attn_metadata_narrowed.max_seq_len, + ) + + # out_page_indices is a required kernel output but unused: + # the indexer's topk output must stay local/compressed- + # context positions (resolved later by + # DeepseekV4CPUAttention.forward_mqa via + # map_local_to_global_slots_cpu). + out_page_scratch = torch.empty( + (batch_size, topk_tokens), + dtype=torch.int32, + device=kv_cache.device, + ) + ops.topk_transform_512_cpu( + logits, + seq_lens, + block_table, + out_page_scratch, + block_size, + topk_indices_buffer[:num_decode_tokens, :topk_tokens], + ) + + return topk_indices_buffer diff --git a/vllm/models/deepseek_v4/__init__.py b/vllm/models/deepseek_v4/__init__.py index d71bddae33..4170a39042 100644 --- a/vllm/models/deepseek_v4/__init__.py +++ b/vllm/models/deepseek_v4/__init__.py @@ -30,6 +30,13 @@ elif current_platform.is_xpu(): from .xpu.dspark import DSparkDeepseekV4ForCausalLM # type: ignore[assignment] from .xpu.model import DeepseekV4ForCausalLM # type: ignore[assignment] from .xpu.mtp import DeepSeekV4MTP # type: ignore[assignment] +elif current_platform.is_cpu(): + from .cpu.dspark import DSparkDeepseekV4ForCausalLM # type: ignore[assignment] + from .cpu.model import DeepseekV4ForCausalLM # type: ignore[assignment] + from .cpu.mtp import DeepSeekV4MTP # type: ignore[assignment] + from .vl_stub import ( # type: ignore[assignment] + DeepseekV4ForConditionalGeneration, + ) else: from .nvidia.dspark import ( # type: ignore[assignment] DSparkDeepseekV4ForCausalLM, diff --git a/vllm/models/deepseek_v4/common/ops/fused_indexer_q.py b/vllm/models/deepseek_v4/common/ops/fused_indexer_q.py index d5c7d669b2..ae014750bf 100644 --- a/vllm/models/deepseek_v4/common/ops/fused_indexer_q.py +++ b/vllm/models/deepseek_v4/common/ops/fused_indexer_q.py @@ -678,6 +678,19 @@ def fused_indexer_q_rope_quant( index_q_fp8, index_weights_out, ) + elif current_platform.is_cpu(): + from vllm._custom_ops import fused_indexer_q_rope_quant_cpu + + fused_indexer_q_rope_quant_cpu( + positions.contiguous(), + index_q.to(torch.float32), + index_q_cos_sin_cache.to(torch.float32).contiguous(), + index_q_fp8, + index_weights.to(torch.float32), + index_weights_softmax_scale, + index_weights_head_scale, + index_weights_out, + ) else: _FUSED_INDEXER_Q_ROPE_QUANT_TRITON_KERNEL( positions, diff --git a/vllm/models/deepseek_v4/cpu/__init__.py b/vllm/models/deepseek_v4/cpu/__init__.py new file mode 100644 index 0000000000..208f01a7cb --- /dev/null +++ b/vllm/models/deepseek_v4/cpu/__init__.py @@ -0,0 +1,2 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project diff --git a/vllm/models/deepseek_v4/cpu/cpu_compressor.py b/vllm/models/deepseek_v4/cpu/cpu_compressor.py new file mode 100644 index 0000000000..6f0b9e06d8 --- /dev/null +++ b/vllm/models/deepseek_v4/cpu/cpu_compressor.py @@ -0,0 +1,112 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""CPU DeepSeek-V4 compressor subclass and dispatch: partial-state-cache +write and the compress -> RMSNorm -> RoPE -> FP8 quant -> KV cache store +step, calling ``torch.ops._C.save_partial_states_cpu``/ +``compress_norm_rope_store_cpu``/``compress_norm_rope_store_indexer_cpu`` +(``csrc/cpu/sgl-kernels/compressor.cpp``) in place of triton-cpu. + +Covers both the head_dim=512 main-attention compressor path and the +head_dim=128 indexer compressor path (fp8 only -- MXFP4, ``use_fp4_cache``, +never occurs on CPU: the indexer's MXFP4 cache requires a Blackwell GPU). +""" + +from typing import Any, cast + +import torch + +from vllm._custom_ops import ( + compress_norm_rope_store_cpu, + compress_norm_rope_store_indexer_cpu, + save_partial_states_cpu, +) +from vllm.forward_context import get_forward_context +from vllm.models.deepseek_v4.compressor import CompressorMetadata, DeepseekCompressor +from vllm.models.deepseek_v4.cpu.cpu_utils import map_local_to_global_slots_cpu + + +class DeepseekV4CPUCompressor(DeepseekCompressor): + """CPU compressor: same state/weights as the shared base, but + ``forward``'s cache-write and compress->RMSNorm->RoPE->quant->store + steps always dispatch straight to the ported CPU kernels above instead + of through the shared method's platform-dispatch chain. This class + only ever runs on CPU, and ``use_fp4_cache`` is always ``False`` here + (the indexer's MXFP4 cache requires a Blackwell GPU), so the shared + method's MXFP4/two-stage/cutedsl/triton branches never apply. + """ + + _norm_weight_fp32: torch.Tensor + + def cache_norm_weight_fp32(self) -> None: + """Cache a fp32 contiguous copy of ``norm.weight`` once, right after + weight loading (see + ``DeepseekV4CPUAttention.process_weights_after_loading``) -- the CPU + kernel requires fp32/contiguous but the checkpoint loads this weight + in bf16, and it never changes after loading.""" + self._norm_weight_fp32 = self.norm.weight.to(torch.float32).contiguous() + + def forward( + self, + kv_score: torch.Tensor, + positions: torch.Tensor, + rotary_emb, + ) -> None: + kv, score = kv_score.split( + [self.coff * self.head_dim, self.coff * self.head_dim], dim=-1 + ) + + attn_metadata = get_forward_context().attn_metadata + if not isinstance(attn_metadata, dict): + return + + state_metadata = attn_metadata[self.state_cache.prefix] + assert isinstance(state_metadata, CompressorMetadata) + token_to_req_indices = state_metadata.token_to_req_indices + assert token_to_req_indices is not None + slot_mapping = state_metadata.slot_mapping + num_actual = slot_mapping.shape[0] + block_table = state_metadata.block_table + block_size = state_metadata.block_size + + state_cache = self.state_cache.kv_cache + + save_partial_states_cpu( + kv, score, self.ape, positions, state_cache, slot_mapping + ) + + cos_sin_cache = rotary_emb.cos_sin_cache + k_cache_metadata = cast(Any, attn_metadata[self.k_cache_prefix]) + k_cache_layer = self._static_forward_context[self.k_cache_prefix] + kv_cache = k_cache_layer.kv_cache + + assert self.head_dim in (512, 128), ( + f"unsupported compressor head_dim: {self.head_dim}" + ) + window = (1 + self.overlap) * self.compress_ratio + win_offsets = torch.arange(window, device=positions.device) + win_local = positions.unsqueeze(-1) - window + 1 + win_offsets + req_idx = token_to_req_indices[:num_actual] + gather_slots = map_local_to_global_slots_cpu( + win_local, req_idx, block_table, block_size + ) + + kv_cache_2d = kv_cache.view(kv_cache.shape[0], -1) + kv_cache_block_size = kv_cache.shape[1] + + store_fn = ( + compress_norm_rope_store_cpu + if self.head_dim == 512 + else compress_norm_rope_store_indexer_cpu + ) + store_fn( + state_cache, + gather_slots, + positions, + k_cache_metadata.slot_mapping, + self._norm_weight_fp32, + self.rms_norm_eps, + cos_sin_cache, + kv_cache_2d, + kv_cache_block_size, + self.compress_ratio, + ) diff --git a/vllm/models/deepseek_v4/cpu/cpu_mla.py b/vllm/models/deepseek_v4/cpu/cpu_mla.py new file mode 100644 index 0000000000..52d4f4c3b0 --- /dev/null +++ b/vllm/models/deepseek_v4/cpu/cpu_mla.py @@ -0,0 +1,126 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""CPU DeepSeek-V4 sparse-MLA backend descriptor. + +DeepSeek-V4 attention runs entirely through +``DeepseekV4Attention``/``DeepseekV4CPUAttention``, never through a generic +``AttentionImpl``. ``DeepseekV4CPUSparseBackend`` exists so the CPU attention +layer has its own backend name and its own metadata-builder class: CPU is not +allowed to run triton-cpu inside the model's forward path, so +``DeepseekV4CPUFlashMLAMetadataBuilder`` overrides the shared base's +``_build_c128a_metadata`` (triton-backed in +``DeepseekV4SparseMLAMetadataBuilder``) with a plain-eager-PyTorch +reimplementation of the same per-token block-table-resolution math. +""" + +import torch + +from vllm.models.deepseek_v4.sparse_mla import ( + DeepseekV4SparseMLABackend, + DeepseekV4SparseMLAMetadataBuilder, +) +from vllm.v1.attention.backend import CommonAttentionMetadata +from vllm.v1.attention.backends.mla.indexer import ( + DeepseekV4IndexerBackend, + DeepseekV32IndexerMetadataBuilder, +) +from vllm.v1.attention.backends.mla.sparse_swa import ( + _LAYER_TYPE_C4A, + _LAYER_TYPE_C128A, + _LAYER_TYPE_SWAONLY, + DeepseekSparseSWAMetadataBuilder, +) +from vllm.v1.attention.ops.flashmla import FlashMLASchedMeta + + +class DeepseekV4CPUFlashMLAMetadataBuilder(DeepseekV4SparseMLAMetadataBuilder): + """CPU sparse-MLA metadata builder: same fields as the shared base, but + skips the C128A dense-topk metadata entirely -- unlike CUDA/XPU/ROCm, + ``DeepseekV4CPUAttention.forward_mqa`` never reads + ``c128a_global_decode_topk_indices``/``c128a_decode_topk_lens``/ + ``c128a_prefill_topk_indices``; it recomputes the same local top-k + directly from ``positions`` instead, so building it here (whether + eagerly or via the shared base's triton kernel) would be wasted work. + """ + + def _build_c128a_metadata( + self, + cm: CommonAttentionMetadata, + req_id_per_token: torch.Tensor, + ) -> dict[str, torch.Tensor | None]: + return {} + + +class DeepseekV4CPUSparseBackend(DeepseekV4SparseMLABackend): + @staticmethod + def get_name() -> str: + return "CPU_V4_MLA_SPARSE" + + @staticmethod + def get_builder_cls() -> type[DeepseekV4CPUFlashMLAMetadataBuilder]: + return DeepseekV4CPUFlashMLAMetadataBuilder + + +class DeepseekV4CPUIndexerMetadataBuilder(DeepseekV32IndexerMetadataBuilder): + """CPU indexer metadata builder: same fields as the shared base, but + prefill requests are never split into multiple chunks. + + ``DeepseekV32IndexerMetadataBuilder._split_indexer_prefill_chunks`` (the + shared base's default) bounds two things the CUDA/XPU indexer kernels + need: the flat-gather workspace size and the dense M*N logits tensor + those kernels allocate for a chunk. The CPU indexer + (``sparse_attn_indexer_cpu``) allocates neither -- it reads the paged + K-cache directly via ``fp8_paged_mqa_logits_cpu``/ + ``topk_transform_512_cpu``, with no chunk-splitting of its own. + With chunked prefill now enabled for this model on CPU (see + ``CpuPlatform.check_and_update_config``'s ``amx_mla_or_dsv4_enabled``), + the total query-token count for one step is already bounded by + ``max_num_batched_tokens`` -- there is nothing left here to bound, so + this always returns a single chunk spanning the whole step's prefill + batch. + """ + + @staticmethod + def _split_indexer_prefill_chunks( + compressed_seq_lens_cpu: torch.Tensor, + prefill_query_lens_cpu: torch.Tensor, + workspace_size: int, + max_logits_bytes: int, + request_offset: int = 0, + ) -> list[tuple[slice, slice]]: + # workspace_size/max_logits_bytes are unused: no flat-gather + # workspace or dense M*N buffer is allocated. + num_reqs = compressed_seq_lens_cpu.shape[0] + total_query_len = int(prefill_query_lens_cpu.sum().item()) + return [ + ( + slice(request_offset, request_offset + num_reqs), + slice(0, total_query_len), + ) + ] + + +class DeepseekV4CPUIndexerBackend(DeepseekV4IndexerBackend): + @staticmethod + def get_name() -> str: + return "CPU_V4_INDEXER" + + @staticmethod + def get_builder_cls() -> type[DeepseekV4CPUIndexerMetadataBuilder]: + return DeepseekV4CPUIndexerMetadataBuilder + + +class DeepseekV4CPUSparseSWAMetadataBuilder(DeepseekSparseSWAMetadataBuilder): + """CPU SWA metadata builder: same fields as the shared base.""" + + def build_tile_scheduler( + self, num_decode_tokens: int + ) -> dict[str, FlashMLASchedMeta | None]: + """CPU never runs the FlashMLA tile-scheduler planner (that's a CUDA + C++ decode-path concern) -- always return the all-``None`` sentinel + the shared base's own CPU branch returns.""" + return { + _LAYER_TYPE_SWAONLY: None, + _LAYER_TYPE_C4A: None, + _LAYER_TYPE_C128A: None, + } diff --git a/vllm/models/deepseek_v4/cpu/cpu_sparse.py b/vllm/models/deepseek_v4/cpu/cpu_sparse.py new file mode 100644 index 0000000000..d51e589652 --- /dev/null +++ b/vllm/models/deepseek_v4/cpu/cpu_sparse.py @@ -0,0 +1,648 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""CPU DeepSeek-V4 attention subclass. + +``forward_mqa`` resolves SWA/compressed top-k indices to paged-cache slot +ids and calls the fused ``flash_mla_with_kvcache_cpu`` kernel. ``_o_proj`` +stays eager -- not on the attention hot path. +""" + +from typing import Any, cast + +import torch +import torch.nn as nn +import torch.nn.functional as F + +from vllm import _custom_ops as ops +from vllm.forward_context import get_forward_context +from vllm.models.deepseek_v4.attention import ( + DeepseekV4Attention, + DeepseekV4Indexer, + DeepseekV4IndexerCache, +) +from vllm.models.deepseek_v4.cpu.cpu_compressor import DeepseekV4CPUCompressor +from vllm.models.deepseek_v4.cpu.cpu_mla import ( + DeepseekV4CPUIndexerBackend, + DeepseekV4CPUSparseBackend, +) +from vllm.models.deepseek_v4.cpu.cpu_utils import map_local_to_global_slots_cpu +from vllm.models.deepseek_v4.sparse_mla import DeepseekV4FlashMLAMetadata +from vllm.platforms import current_platform +from vllm.utils.torch_utils import ( + LayerNameType, + _encode_layer_name, + _resolve_layer_name, + direct_register_custom_op, +) +from vllm.v1.attention.backend import AttentionBackend, AttentionMetadata +from vllm.v1.attention.backends.mla.sparse_swa import DeepseekSparseSWAMetadata + +_ROPE_DIM = 64 + + +def _deepseek_v4_cpu_prepare_and_attn( + hidden_states: torch.Tensor, + qr: torch.Tensor, + kv: torch.Tensor, + qr_scale: torch.Tensor | None, + kv_score: torch.Tensor | None, + indexer_kv_score: torch.Tensor | None, + indexer_weights: torch.Tensor | None, + positions: torch.Tensor, + o_padded: torch.Tensor, + layer_name: LayerNameType, +) -> torch.Tensor: + """Opaque custom-op boundary around attention-prep + sparse-indexer + + MLA attention + output projection -- CPU-only. + + Wraps cache writes with no explicit tensor args (SWA/compressor/indexer + state), which torch.compile can't otherwise track as mutations -- one + opaque call sidesteps that. ``_o_proj`` is folded in too because CPU's + ``DYNAMO_TRACE_ONCE`` mode traces exactly once, during warmup when + ``attn_metadata`` is ``None``; tracing ``_o_proj`` separately would bake + that warmup branch into the compiled graph permanently, dead-code- + eliminating the whole attention computation upstream. + """ + layer_name_str = _resolve_layer_name(layer_name) + self = get_forward_context().no_compile_layers[layer_name_str] + self._prepare_and_attn_fn( + hidden_states, + qr, + kv, + qr_scale, + kv_score, + indexer_kv_score, + indexer_weights, + positions, + o_padded, + ) + o = o_padded[:, : self.n_local_heads, :] + return self._o_proj(o, positions) + + +def _deepseek_v4_cpu_prepare_and_attn_fake( + hidden_states: torch.Tensor, + qr: torch.Tensor, + kv: torch.Tensor, + qr_scale: torch.Tensor | None, + kv_score: torch.Tensor | None, + indexer_kv_score: torch.Tensor | None, + indexer_weights: torch.Tensor | None, + positions: torch.Tensor, + o_padded: torch.Tensor, + layer_name: LayerNameType, +) -> torch.Tensor: + # hidden_states' last dim already equals _o_proj's output width, so we + # don't need to resolve layer_name (an opaque object) here. + return hidden_states.new_empty((hidden_states.shape[0], hidden_states.shape[-1])) + + +direct_register_custom_op( + op_name="deepseek_v4_cpu_prepare_and_attn", + op_func=_deepseek_v4_cpu_prepare_and_attn, + mutates_args=["o_padded"], + fake_impl=_deepseek_v4_cpu_prepare_and_attn_fake, +) + + +def _dequant_linear_weight(layer: torch.nn.Module) -> torch.Tensor: + """Dequantize a linear layer's weight, including FP8 block scales. + + ``_o_proj`` reads ``wo_a.weight`` directly (not via + ``quant_method.apply()``), so the per-block FP8 scale must be applied + here explicitly. + """ + weight = layer.weight + scale = getattr(layer, "weight_scale_inv", None) + if scale is None: + scale = getattr(layer, "weight_scale", None) + if scale is None: + return weight.float() + + if scale.dtype == torch.float8_e8m0fnu: + from vllm.model_executor.layers.quantization.utils.fp8_utils import ( + _upcast_e8m0_to_fp32, + ) + + scale = _upcast_e8m0_to_fp32(scale) + scale = scale.to(torch.float32) + + if scale.numel() == 1: + return weight.float() * scale + + block_size = getattr(layer, "weight_block_size", None) + if block_size is None: + # Per-channel scale: one row-scale per output row. + return weight.float() * scale.view(-1, 1) + + block_m, block_k = block_size + m, k = weight.shape + scale = scale.repeat_interleave(block_m, dim=0).repeat_interleave(block_k, dim=1) + return weight.float() * scale[:m, :k] + + +def _fused_indexer_q_rope_quant_cpu( + positions: torch.Tensor, + index_q: torch.Tensor, + index_q_cos_sin_cache: torch.Tensor, + index_weights: torch.Tensor, + index_weights_softmax_scale: float, + index_weights_head_scale: float, +) -> tuple[torch.Tensor, torch.Tensor]: + """CPU-only equivalent of ``fused_indexer_q_rope_quant``'s FP8 path + (CPU never reaches its MXFP4 arm -- Blackwell-only). Returns + ``(q_fp8, weights_out)``, folding the per-token q scale into + ``index_weights``. ``positions``/``index_q_cos_sin_cache`` are already + contiguous int64/fp32, so no conversion is needed here. + """ + assert positions.ndim == 1 + assert index_q.ndim == 3 + assert index_q_cos_sin_cache.ndim == 2 + + index_weights_out = torch.empty_like(index_weights, dtype=torch.float32) + index_q_fp8 = torch.empty_like(index_q, dtype=current_platform.fp8_dtype()) + ops.fused_indexer_q_rope_quant_cpu( + positions, + index_q.to(torch.float32), + index_q_cos_sin_cache, + index_q_fp8, + index_weights.to(torch.float32), + index_weights_softmax_scale, + index_weights_head_scale, + index_weights_out, + ) + return index_q_fp8, index_weights_out + + +class _NoOpEvent: + """Stand-in for ``torch.cuda.Event`` on CPU: since CPU never sets an aux + stream list, these events are constructed but never actually used.""" + + +class DeepseekV4CPUIndexerCache(DeepseekV4IndexerCache): + """CPU indexer K-cache descriptor: same fields as the shared base, just + pointing ``get_attn_backend()`` at ``DeepseekV4CPUIndexerBackend`` + instead of the shared, CUDA/XPU-oriented ``DeepseekV4IndexerBackend``.""" + + def get_attn_backend(self) -> type[AttentionBackend]: + return DeepseekV4CPUIndexerBackend + + +class DeepseekV4CPUIndexer(DeepseekV4Indexer): + """CPU indexer: the C4A short-context fallback runs as eager PyTorch + instead of Triton. Never constructed directly -- ``__class__`` is + swapped onto an existing ``DeepseekV4Indexer`` instance post-construction. + """ + + def forward( + self, + hidden_states: torch.Tensor, + qr: torch.Tensor, + compressed_kv_score: torch.Tensor, + indexer_weights: torch.Tensor, + positions: torch.Tensor, + rotary_emb: nn.Module, + qr_scale: torch.Tensor | None = None, + ) -> tuple[torch.Tensor | None, torch.Tensor | None, torch.Tensor | None]: + """CPU override: no aux streams, so wq_b_and_q_quant and the + compressor run straight-line instead of through + ``maybe_execute_in_parallel``; the short-context check drops the + CUDA-only ``is_current_stream_capturing()`` guard (always False + here).""" + compressor = self.compressor + + attn_metadata = get_forward_context().attn_metadata + if isinstance(attn_metadata, dict): + indexer_metadata = cast(Any, attn_metadata[self.k_cache.prefix]) + if indexer_metadata.max_seq_len // self.compress_ratio <= self.topk_tokens: + # Fewer candidates than topk: all are selected, but the K + # cache still needs to be built. + compressor(compressed_kv_score, positions, rotary_emb) + assert self.topk_indices_buffer is not None + num_tokens = ( + indexer_metadata.num_decode_tokens + + indexer_metadata.num_prefill_tokens + ) + if num_tokens > 0: + self._fill_short_context_topk(positions, num_tokens) + return None, None, None + + def wq_b_and_q_quant(): + q = self._wq_b_proj(qr, qr_scale) + q = q.view(-1, self.n_head, self.head_dim) + return _fused_indexer_q_rope_quant_cpu( + positions, + q, + rotary_emb.cos_sin_cache, + indexer_weights, + self.softmax_scale, + self.n_head**-0.5, + ) + + # Order matches the base class's parallel-join contract: compressor + # must complete before indexer_op runs (skip_k_cache_insert=True). + q, weights = wq_b_and_q_quant() + compressor(compressed_kv_score, positions, rotary_emb) + # No q_scale: FP8-only path (MXFP4 needs Blackwell). + return q, None, weights + + def _fill_short_context_topk( + self, positions: torch.Tensor, num_tokens: int + ) -> None: + assert self.topk_indices_buffer is not None + device = positions.device + arange = torch.arange(self.topk_tokens, device=device) + num_compressed = torch.div( + positions[:num_tokens] + 1, self.compress_ratio, rounding_mode="floor" + ) + valid = arange.unsqueeze(0) < num_compressed.unsqueeze(-1) + filled = torch.where( + valid, arange.unsqueeze(0), torch.full_like(arange.unsqueeze(0), -1) + ) + self.topk_indices_buffer[:num_tokens, : self.topk_tokens] = filled.to( + self.topk_indices_buffer.dtype + ) + + +class DeepseekV4CPUAttention(DeepseekV4Attention): + """CPU sparse MLA attention layer for DeepSeek V4.""" + + backend_cls = DeepseekV4CPUSparseBackend + + def __init__(self, *args, **kwargs) -> None: + super().__init__(*args, **kwargs) + + # Swap classes post-construction (no class-swap hook on the shared + # base) so forward()/get_attn_backend() dispatch to CPU kernels. + # SWA-only layers (compress_ratio <= 1) never build a + # compressor/indexer. + if self.compressor is not None: + self.compressor.__class__ = DeepseekV4CPUCompressor + if self.indexer is not None: + self.indexer.__class__ = DeepseekV4CPUIndexer + self.indexer.k_cache.__class__ = DeepseekV4CPUIndexerCache + self.indexer.compressor.__class__ = DeepseekV4CPUCompressor + + self._wo_a_packed: torch.Tensor | None = None + self._wrap_wo_a_process_weights_after_loading() + + def process_weights_after_loading(self, act_dtype: torch.dtype) -> None: + """Cache fp32-contiguous copies of tensors CPU kernels want that way + but that arrive bf16 (rotary cos/sin table, compressor RMSNorm + weights) -- cast once here instead of per forward call. + + Runs after every quantized layer's own + ``process_weights_after_loading`` (see ``is_deferred_attention_layer``); + ``wo_a``'s packing must happen *before* that phase instead, so it's a + separate monkeypatch in ``__init__``. + """ + self.rotary_emb.cos_sin_cache = self.rotary_emb.cos_sin_cache.to( + torch.float32 + ).contiguous() + if self.compressor is not None: + self.compressor.cache_norm_weight_fp32() + if self.indexer is not None: + self.indexer.compressor.cache_norm_weight_fp32() + + def _wrap_wo_a_process_weights_after_loading(self) -> None: + """Snapshot and pack ``wo_a``'s weight into a bf16 copy for + ``bmm_cpu`` before the FP8 kernel's own + ``process_weights_after_loading`` VNNI-repacks it in place for + row-major reads ``_o_proj`` never does.""" + wo_a = self.wo_a + quant_method = wo_a.quant_method + orig_pwal = quant_method.process_weights_after_loading + heads_per_group = self.n_local_heads // self.n_local_groups + k_dim = heads_per_group * self.head_dim + + def _capture_and_pack(layer: torch.nn.Module) -> None: + dequant = _dequant_linear_weight(layer).to(torch.bfloat16) + dequant = dequant.view(self.n_local_groups, self.o_lora_rank, k_dim) + self._wo_a_packed = torch.ops._C.convert_weight_packed(dequant.contiguous()) + orig_pwal(layer) + # orig_pwal VNNI-repacks layer.weight/scale for ordinary-linear + # use that _o_proj never needs -- free them, same pattern as + # dispatch_cpu_unquantized_gemm's remove_weight=True. + layer.weight = torch.nn.Parameter(torch.empty(0), requires_grad=False) + scale_attr = ( + "weight_scale_inv" + if hasattr(layer, "weight_scale_inv") + else "weight_scale" + ) + if hasattr(layer, scale_attr): + setattr( + layer, + scale_attr, + torch.nn.Parameter(torch.empty(0), requires_grad=False), + ) + + quant_method.process_weights_after_loading = ( # type: ignore[method-assign] + _capture_and_pack + ) + + def forward( + self, + positions: torch.Tensor, + hidden_states: torch.Tensor, + llama_4_scaling: torch.Tensor | None = None, + ) -> torch.Tensor: + """CPU override: wraps attention-prep + MLA + output projection in the + CPU-only opaque custom op (see ``_deepseek_v4_cpu_prepare_and_attn``). + """ + num_tokens = hidden_states.shape[0] + o_padded = torch.empty( + (num_tokens, self.padded_heads, self.head_dim), + dtype=hidden_states.dtype, + device=hidden_states.device, + ) + + qr_kv = self._fused_wqa_wkv_gemm(hidden_states) + + kv_score: torch.Tensor | None = None + if self.compressor is not None: + # Must go through the module's own forward (not a raw matmul on + # its weight): the packed-GEMM kernel frees the raw weight after + # packing, same as wo_a's. + kv_score = self.compressor.fused_wkv_wgate(hidden_states).to(torch.float32) + + indexer_weights: torch.Tensor | None = None + indexer_kv_score: torch.Tensor | None = None + if self.indexer is not None: + indexer_weights, _ = self.indexer.weights_proj(hidden_states) + indexer_kv_score = self.indexer.compressor.fused_wkv_wgate( + hidden_states + ).to(torch.float32) + + qr, qr_scale, kv = self._split_qkv_and_norm(qr_kv) + + return torch.ops.vllm.deepseek_v4_cpu_prepare_and_attn( + hidden_states, + qr, + kv, + qr_scale, + kv_score, + indexer_kv_score, + indexer_weights, + positions, + o_padded, + _encode_layer_name(self.prefix), + ) + + def _fused_qnorm_rope_kv_insert( + self, + q: torch.Tensor, + kv: torch.Tensor, + positions: torch.Tensor, + attn_metadata: ( + dict[str, AttentionMetadata] | list[dict[str, AttentionMetadata]] | None + ), + ) -> torch.Tensor: + """CPU override: only the fp8_ds_mla (uint8) SWA cache layout is + supported here, so the base method's bf16/per-tensor-fp8 branches + are dropped.""" + if not isinstance(attn_metadata, dict): + # Profile run: kernel doesn't fire; produce a padded tensor so + # downstream FlashMLA gets the right shape. + if self.n_local_heads < self.padded_heads: + return F.pad( + q, + (0, 0, 0, self.padded_heads - self.n_local_heads), + value=0.0, + ) + return q + + swa_metadata = cast( + "DeepseekSparseSWAMetadata | None", + attn_metadata.get(self.swa_cache_layer.prefix), + ) + assert swa_metadata is not None + + swa_kv_cache = self.swa_cache_layer.kv_cache + # positions is already int64 (the runner's buffer), no cast needed. + assert positions.dtype == torch.int64 + assert swa_kv_cache.dtype == torch.uint8, ( + "DeepseekV4CPUAttention only supports the fp8_ds_mla cache layout" + ) + # Fused: Q side does per-head weight-free RMSNorm + GPT-J RoPE + # (zero-filling padded head slots); KV side does GPT-J RoPE + UE8M0 + # FP8 quant + paged insert. + swa_kv_cache_2d = swa_kv_cache.view(swa_kv_cache.shape[0], -1) + return ops.fused_qnorm_rope_kv_insert_cpu( + q, + kv, + positions, + swa_kv_cache_2d, + swa_metadata.slot_mapping, + self.rotary_emb.cos_sin_cache, + self.padded_heads, + self.eps, + swa_metadata.block_size, + ) + + def _split_qkv_and_norm( + self, qr_kv: torch.Tensor + ) -> tuple[torch.Tensor, torch.Tensor | None, torch.Tensor]: + """CPU override: two ``RMSNorm`` calls instead of the shared + ``fused_q_kv_rmsnorm``, whose raw ``@triton.jit`` kernel fails to + link under triton-cpu (``undefined symbol: __truncdfbf2``).""" + qr, kv = qr_kv.split([self.q_lora_rank, self.head_dim], dim=-1) + return self.q_norm(qr), None, self.kv_norm(kv) + + def _prepare_and_attn( + self, + hidden_states: torch.Tensor, + qr: torch.Tensor, + kv: torch.Tensor, + qr_scale: torch.Tensor | None, + kv_score: torch.Tensor, + indexer_kv_score: torch.Tensor, + indexer_weights: torch.Tensor, + positions: torch.Tensor, + o_padded: torch.Tensor, + ) -> None: + """CPU override: no aux streams, so query-projection+KV-insert, the + indexer, and the compressor run straight-line instead of through + execute_in_parallel/maybe_execute_in_parallel.""" + attn_metadata = get_forward_context().attn_metadata + indexer = self.indexer + compressor = self.compressor + + q = self._wq_b_proj(qr, qr_scale).view(-1, self.n_local_heads, self.head_dim) + q = self._fused_qnorm_rope_kv_insert(q, kv, positions, attn_metadata) + + index_q: torch.Tensor | None = None + index_q_scale: torch.Tensor | None = None + index_weights_out: torch.Tensor | None = None + if indexer is not None: + assert compressor is not None + index_q, index_q_scale, index_weights_out = indexer( + hidden_states, + qr, + indexer_kv_score, + indexer_weights, + positions, + self.indexer_rotary_emb, + qr_scale, + ) + compressor(kv_score, positions, self.rotary_emb) + elif compressor is not None: + compressor(kv_score, positions, self.rotary_emb) + + self._sparse_indexer_and_attn( + hidden_states, + index_q, + index_q_scale, + index_weights_out, + q, + kv, + positions, + o_padded, + ) + + def _sparse_indexer_and_attn( + self, + hidden_states: torch.Tensor, + index_q: torch.Tensor | None, + index_q_scale: torch.Tensor | None, + index_weights: torch.Tensor | None, + q: torch.Tensor, + kv: torch.Tensor, + positions: torch.Tensor, + out: torch.Tensor, + ) -> None: + """CPU override: identical body, minus ``@eager_break_during_capture`` + (this platform never captures a CUDA graph, so it was a no-op).""" + if self.indexer is not None and index_q is not None: + assert index_weights is not None + # index_q_scale is always None on CPU (FP8-only path). + assert index_q_scale is None + self.indexer.indexer_op(hidden_states, index_q, None, index_weights) + + self.forward_mqa(q, kv, positions, out) + + @classmethod + def get_padded_num_q_heads(cls, num_heads: int) -> int: + # No head-count padding constraint for the (not yet ported) CPU kernel. + return num_heads + + def _o_proj(self, o: torch.Tensor, positions: torch.Tensor) -> torch.Tensor: + if get_forward_context().attn_metadata is None: + # Warmup/profile run: no real metadata, o is a zeroed placeholder. + return torch.zeros( + (o.shape[0], self.hidden_size), dtype=o.dtype, device=o.device + ) + + num_tokens, num_heads, _ = o.shape + cos_sin_cache = self.rotary_emb.cos_sin_cache + # Undo the query's own GPT-J RoPE rotation: the attention output is + # a weighted sum of KV rows at different positions, so R(-pos_q) + # recovers the translation-invariant quantity wo_a/wo_b expect. + o_derot = ops.inverse_gptj_rope_o_proj_cpu( + o, positions, cos_sin_cache, _ROPE_DIM + ) + + heads_per_group = num_heads // self.n_local_groups + o_grouped = o_derot.reshape( + num_tokens, self.n_local_groups, heads_per_group * self.head_dim + ).to(torch.bfloat16) + assert self._wo_a_packed is not None + z = o_grouped.new_empty((self.n_local_groups, num_tokens, self.o_lora_rank)) + ops.bmm_cpu( + z, + o_grouped.transpose(0, 1).contiguous(), + self._wo_a_packed, + True, + None, + ) + z = z.transpose(0, 1).to(o.dtype) + return self.wo_b(z.flatten(1)) + + def forward_mqa( + self, + q: torch.Tensor, + kv: torch.Tensor, + positions: torch.Tensor, + output: torch.Tensor, + ) -> None: + forward_context = get_forward_context() + attn_metadata = forward_context.attn_metadata + if attn_metadata is None: + # Warmup/profile run: reserve nothing extra, skip real kernels. + output.zero_() + return + + assert isinstance(attn_metadata, dict) + swa_metadata = cast( + DeepseekSparseSWAMetadata, + attn_metadata.get(self.swa_cache_layer.prefix), + ) + assert swa_metadata is not None + assert swa_metadata.token_to_req_indices is not None + req_idx = swa_metadata.token_to_req_indices + + swa_only = self.compress_ratio <= 1 + flashmla_metadata = cast( + DeepseekV4FlashMLAMetadata | None, attn_metadata.get(self.prefix) + ) + + num_tokens = q.shape[0] + device = q.device + + # ---- SWA window: causal sliding window over the raw per-token cache + window = self.window_size + win_offsets = torch.arange(window, device=device) + start_pos = (positions - window + 1).clamp(min=0) + local_pos = start_pos.unsqueeze(-1) + win_offsets + valid_win = local_pos <= positions.unsqueeze(-1) + win_local = torch.where(valid_win, local_pos, torch.full_like(local_pos, -1)) + window_slots = map_local_to_global_slots_cpu( + win_local, req_idx, swa_metadata.block_table, swa_metadata.block_size + ) + window_cache_2d = self.swa_cache_layer.kv_cache.view( + self.swa_cache_layer.kv_cache.shape[0], -1 + ) + + # ---- Compressed KV: C4A's learned top-k, or C128A's dense top-k ---- + if swa_only: + compressed_slots = window_slots.new_full((num_tokens, 0), -1) + # Unused when num_compressed == 0 (only dtype/2D-ness matter); + # reuse the window cache to avoid an allocation. + compressed_cache_2d = window_cache_2d + compressed_block_size = 1 + else: + assert flashmla_metadata is not None + main_block_size = flashmla_metadata.block_size // self.compress_ratio + if self.compress_ratio == 4: + assert self.topk_indices_buffer is not None + local_topk = self.topk_indices_buffer[:num_tokens] + else: + width = -(-self.max_model_len // self.compress_ratio) + num_compressed = torch.div( + positions + 1, self.compress_ratio, rounding_mode="floor" + ).clamp(max=width) + arange_w = torch.arange(width, device=device) + local_topk = torch.where( + arange_w.unsqueeze(0) < num_compressed.unsqueeze(-1), + arange_w.unsqueeze(0).expand(num_tokens, width), + torch.full( + (num_tokens, width), -1, device=device, dtype=torch.int64 + ), + ) + compressed_slots = map_local_to_global_slots_cpu( + local_topk, req_idx, flashmla_metadata.block_table, main_block_size + ) + compressed_cache_2d = self.kv_cache.view(self.kv_cache.shape[0], -1) + compressed_block_size = main_block_size + + torch.ops._C.flash_mla_with_kvcache_cpu( + output, + q, + window_cache_2d, + window_slots, + swa_metadata.block_size, + compressed_cache_2d, + compressed_slots, + compressed_block_size, + self.attn_sink, + self.scale, + ) diff --git a/vllm/models/deepseek_v4/cpu/cpu_utils.py b/vllm/models/deepseek_v4/cpu/cpu_utils.py new file mode 100644 index 0000000000..f03bf60b58 --- /dev/null +++ b/vllm/models/deepseek_v4/cpu/cpu_utils.py @@ -0,0 +1,34 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Shared helpers for DeepSeek-V4's CPU-ported kernels.""" + +import torch + + +def map_local_to_global_slots_cpu( + local_indices: torch.Tensor, + req_idx: torch.Tensor, + block_table: torch.Tensor, + block_size: int, +) -> torch.Tensor: + """Map per-request-local (token or compressed-token) indices to global + paged-cache slot ids via a block table, mirroring + ``compute_global_topk_indices_and_lens``. + + Args: + local_indices: [N, K], -1 sentinel for invalid entries. + req_idx: [N], row into ``block_table`` for each of the N rows. + block_table: [num_reqs, max_blocks_per_seq]. + block_size: tokens (or compressed tokens) per physical block. + + Returns: + [N, K] int64 global slot ids, -1 where ``local_indices`` was invalid. + """ + valid = local_indices >= 0 + safe_local = local_indices.clamp(min=0).to(torch.int64) + block_pos = torch.div(safe_local, block_size, rounding_mode="floor") + offset = safe_local % block_size + req_expand = req_idx.to(torch.int64).unsqueeze(-1).expand_as(local_indices) + block_num = block_table[req_expand, block_pos].to(torch.int64) + slot = block_num * block_size + offset + return torch.where(valid, slot, torch.full_like(slot, -1)) diff --git a/vllm/models/deepseek_v4/cpu/dspark.py b/vllm/models/deepseek_v4/cpu/dspark.py new file mode 100644 index 0000000000..36ab1bb50c --- /dev/null +++ b/vllm/models/deepseek_v4/cpu/dspark.py @@ -0,0 +1,18 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""CPU DeepSeek-V4 DSpark (speculative decoding) — deferred, not implemented. + +``DSparkDeepseekV4ForCausalLM`` is only instantiated when the speculative- +decoding registry resolves it by name for a DSpark speculative config; a +plain (non-spec) DeepSeek-V4 CPU model never constructs this class. This +stub exists solely so ``vllm.models.deepseek_v4.cpu`` is import-clean. +""" + +from vllm.config import VllmConfig + + +class DSparkDeepseekV4ForCausalLM: + def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None: + raise NotImplementedError( + "DeepSeek-V4 speculative decoding (DSpark) is not implemented on CPU." + ) diff --git a/vllm/models/deepseek_v4/cpu/model.py b/vllm/models/deepseek_v4/cpu/model.py new file mode 100644 index 0000000000..6e88d21c27 --- /dev/null +++ b/vllm/models/deepseek_v4/cpu/model.py @@ -0,0 +1,922 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import typing +from collections.abc import Callable, Iterable +from itertools import islice + +import regex as re +import torch +import torch.nn as nn + +from vllm.compilation.decorators import support_torch_compile +from vllm.config import VllmConfig +from vllm.distributed import ( + get_pp_group, + get_tensor_model_parallel_rank, + get_tensor_model_parallel_world_size, +) +from vllm.model_executor.layers.activation import SiluAndMul, SiluAndMulWithClamp +from vllm.model_executor.layers.fused_moe import ( + FusedMoEFactory, + GateLinear, + fused_moe_make_expert_params_mapping, +) +from vllm.model_executor.layers.fused_moe.utils import ( + is_model_fused_shared_expert_compatible, +) +from vllm.model_executor.layers.layernorm import RMSNorm +from vllm.model_executor.layers.linear import ( + MergedColumnParallelLinear, + RowParallelLinear, +) +from vllm.model_executor.layers.logits_processor import LogitsProcessor +from vllm.model_executor.layers.mhc import ( + HAS_AITER_MHC, + HAS_TILELANG_MHC, + HCHeadOp, + MHCFusedPostPreOp, + MHCPostOp, + MHCPreOp, +) +from vllm.model_executor.layers.quantization import QuantizationConfig +from vllm.model_executor.layers.vocab_parallel_embedding import ( + ParallelLMHead, + VocabParallelEmbedding, +) +from vllm.model_executor.model_loader.weight_utils import default_weight_loader +from vllm.model_executor.models.interfaces import ( + EagleModelMixin, + SupportsEagle3, + SupportsPP, +) +from vllm.model_executor.models.utils import ( + AutoWeightsLoader, + PPMissingLayer, + WeightsMapper, + extract_layer_index, + is_pp_missing_parameter, + make_layers, + maybe_prefix, +) +from vllm.models.deepseek_v4.cpu.cpu_sparse import DeepseekV4CPUAttention +from vllm.platforms import current_platform +from vllm.sequence import IntermediateTensors + + +class DeepseekV4MLP(nn.Module): + def __init__( + self, + hidden_size: int, + intermediate_size: int, + hidden_act: str, + swiglu_limit: float | None = None, + quant_config: QuantizationConfig | None = None, + reduce_results: bool = True, + is_sequence_parallel: bool = False, + prefix: str = "", + ) -> None: + super().__init__() + + # is_sequence_parallel: weights are replicated, no collective ops + # needed. Otherwise: standard TP with an allreduce at the end. + self.gate_up_proj = MergedColumnParallelLinear( + hidden_size, + [intermediate_size] * 2, + bias=False, + quant_config=quant_config, + disable_tp=is_sequence_parallel, + prefix=f"{prefix}.gate_up_proj", + ) + self.down_proj = RowParallelLinear( + intermediate_size, + hidden_size, + bias=False, + quant_config=quant_config, + reduce_results=reduce_results, + disable_tp=is_sequence_parallel, + prefix=f"{prefix}.down_proj", + ) + if hidden_act != "silu": + raise ValueError( + f"Unsupported activation: {hidden_act}. Only silu is supported for now." + ) + if swiglu_limit is not None: + self.act_fn = SiluAndMulWithClamp(swiglu_limit) + else: + self.act_fn = SiluAndMul() + + def forward(self, x): + gate_up, _ = self.gate_up_proj(x) + x = self.act_fn(gate_up) + x, _ = self.down_proj(x) + return x + + +class DeepseekV4MoE(nn.Module): + def __init__( + self, + vllm_config: VllmConfig, + prefix: str = "", + ): + super().__init__() + + self.tp_size = get_tensor_model_parallel_world_size() + config = vllm_config.model_config.hf_config + quant_config = vllm_config.quant_config + self.prefix = prefix + + self.routed_scaling_factor = getattr(config, "routed_scaling_factor", 1.0) + self.hidden_size = config.hidden_size + + self.n_routed_experts = config.n_routed_experts + self.n_activated_experts = config.num_experts_per_tok + self.moe_intermediate_size = config.moe_intermediate_size + self.swiglu_limit = config.swiglu_limit + self.renormalize = config.norm_topk_prob + self.scoring_func = getattr(config, "scoring_func", "sqrtsoftplus") + + self.gate = GateLinear( + input_size=config.hidden_size, + output_size=config.n_routed_experts, + bias=False, + out_dtype=torch.float32, + prefix=f"{prefix}.gate", + ) + + self.gate.e_score_correction_bias = None + self.gate.tid2eid = None + is_hash_moe = extract_layer_index(prefix) < config.num_hash_layers + self.hash_indices_dtype = torch.int32 + if is_hash_moe: + # hash MoE doesn't use e_score_correction_bias + # Use randint instead of empty to avoid garbage values causing + # invalid memory access in dummy mode (--load-format="dummy") + self.gate.tid2eid = nn.Parameter( + torch.randint( + 0, + config.n_routed_experts, + (config.vocab_size, config.num_experts_per_tok), + dtype=self.hash_indices_dtype, + ), + requires_grad=False, + ) + elif getattr(config, "topk_method", None) == "noaux_tc": + self.gate.e_score_correction_bias = nn.Parameter( + torch.empty(config.n_routed_experts, dtype=torch.float32), + requires_grad=False, + ) + + self.n_shared_experts = config.n_shared_experts + + # No CPU fused-shared-expert grouped-GEMM path yet. + self.is_fused_shared_expert_enabled = False + + if config.n_shared_experts is None or self.is_fused_shared_expert_enabled: + self.shared_experts = None + else: + intermediate_size = config.moe_intermediate_size * config.n_shared_experts + + self.shared_experts = DeepseekV4MLP( + hidden_size=config.hidden_size, + intermediate_size=intermediate_size, + hidden_act=config.hidden_act, + swiglu_limit=self.swiglu_limit, + quant_config=quant_config, + reduce_results=False, + prefix=f"{prefix}.shared_experts", + ) + + self.tp_rank = get_tensor_model_parallel_rank() + assert config.n_routed_experts % self.tp_size == 0 + + self.n_local_experts = config.n_routed_experts // self.tp_size + self.experts_start_idx = self.tp_rank * self.n_local_experts + self.experts_end_idx = self.experts_start_idx + self.n_local_experts + + self.experts = FusedMoEFactory( + shared_experts=self.shared_experts, + n_shared_experts=( + config.n_shared_experts if self.is_fused_shared_expert_enabled else None + ), + fuse_shared_experts=self.is_fused_shared_expert_enabled, + gate=self.gate, + num_experts=config.n_routed_experts, + top_k=config.num_experts_per_tok, + hidden_size=config.hidden_size, + intermediate_size=config.moe_intermediate_size, + renormalize=config.norm_topk_prob, + quant_config=quant_config, + prefix=f"{prefix}.experts", + scoring_func=self.scoring_func, + routed_scaling_factor=self.routed_scaling_factor, + e_score_correction_bias=self.gate.e_score_correction_bias, + hash_indices_table=self.gate.tid2eid, + swiglu_limit=self.swiglu_limit, + router_logits_dtype=torch.float32, + ) + if is_hash_moe: + # apply_monolithic() never sees hash_indices_table via the + # router path (modular-only) -- stash it on RoutedExperts + # directly. object.__setattr__ avoids nn.Module re-registering + # (and duplicating in the state_dict) this alias. + object.__setattr__( + self.experts.routed_experts, "hash_indices_table", self.gate.tid2eid + ) + + def forward( + self, hidden_states: torch.Tensor, input_ids: torch.Tensor | None = None + ) -> torch.Tensor: + if self.gate.tid2eid is not None and input_ids is None: + raise ValueError("DeepSeek V4 hash MoE routing requires input_ids.") + + org_shape = hidden_states.shape + final_hidden_states = self.experts( + hidden_states=hidden_states, + router_logits=hidden_states, + input_ids=input_ids, + ) + + return final_hidden_states.view(org_shape) + + +class DeepseekV4DecoderLayer(nn.Module): + def __init__( + self, + vllm_config, + prefix, + topk_indices_buffer: torch.Tensor | None = None, + ): + super().__init__() + + # Lazy import to avoid top-level tilelang dependency. + # Registers both torch.ops.vllm.mhc_pre and mhc_post + import vllm.model_executor.layers.mhc # noqa: F401 + + config = vllm_config.model_config.hf_config + self.hidden_size = config.hidden_size + + self.rms_norm_eps = config.rms_norm_eps + self.attn = DeepseekV4CPUAttention( + vllm_config, + prefix=f"{prefix}.attn", + topk_indices_buffer=topk_indices_buffer, + ) + self.ffn = DeepseekV4MoE(vllm_config, prefix=f"{prefix}.ffn") + + self.attn_norm = RMSNorm(self.hidden_size, self.rms_norm_eps) + self.ffn_norm = RMSNorm(self.hidden_size, self.rms_norm_eps) + self.hc_mult = config.hc_mult + self.hc_sinkhorn_iters = config.hc_sinkhorn_iters + self.hc_eps = config.hc_eps + self.hc_post_alpha = 2.0 + mix_hc = (2 + self.hc_mult) * self.hc_mult + hc_dim = self.hc_mult * self.hidden_size + self.hc_attn_fn = nn.Parameter( + torch.empty( + (mix_hc, hc_dim), + dtype=torch.float32, + ), + requires_grad=False, + ) + self.hc_ffn_fn = nn.Parameter( + torch.empty( + (mix_hc, hc_dim), + dtype=torch.float32, + ), + requires_grad=False, + ) + self.hc_attn_base = nn.Parameter( + torch.empty( + mix_hc, + dtype=torch.float32, + ), + requires_grad=False, + ) + self.hc_ffn_base = nn.Parameter( + torch.empty( + mix_hc, + dtype=torch.float32, + ), + requires_grad=False, + ) + self.hc_attn_scale = nn.Parameter( + torch.empty( + 3, + dtype=torch.float32, + ), + requires_grad=False, + ) + self.hc_ffn_scale = nn.Parameter( + torch.empty( + 3, + dtype=torch.float32, + ), + requires_grad=False, + ) + self.mhc_pre = MHCPreOp() + self.mhc_post = MHCPostOp() + self.mhc_fused_post_pre = MHCFusedPostPreOp() + self.use_fused_mhc = HAS_TILELANG_MHC and not ( + HAS_AITER_MHC and self.hidden_size % 256 == 0 + ) + + def hc_pre( + self, + x: torch.Tensor, + hc_fn: torch.Tensor, + hc_scale: torch.Tensor, + hc_base: torch.Tensor, + ): + post_mix, res_mix, layer_input = self.mhc_pre( + residual=x, + fn=hc_fn, + hc_scale=hc_scale, + hc_base=hc_base, + rms_eps=self.rms_norm_eps, + hc_pre_eps=self.hc_eps, + hc_sinkhorn_eps=self.hc_eps, + hc_post_mult_value=self.hc_post_alpha, + sinkhorn_repeat=self.hc_sinkhorn_iters, + ) + return layer_input, post_mix, res_mix + + def hc_post( + self, + x: torch.Tensor, + residual: torch.Tensor, + post: torch.Tensor, + comb: torch.Tensor, + ): + return self.mhc_post(x, residual, post, comb) + + def _forward_fused_post_pre( + self, + x: torch.Tensor, + positions: torch.Tensor, + input_ids: torch.Tensor | None, + post_mix: torch.Tensor | None = None, + res_mix: torch.Tensor | None = None, + residual: torch.Tensor | None = None, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + if residual is None: + # Run standalone hc_pre on first layer + residual = x + x, post_mix, res_mix = self.hc_pre( + x, self.hc_attn_fn, self.hc_attn_scale, self.hc_attn_base + ) + else: + residual, post_mix, res_mix, x = self.mhc_fused_post_pre( + x, + residual, + post_mix, + res_mix, + self.hc_attn_fn, + self.hc_attn_scale, + self.hc_attn_base, + self.rms_norm_eps, + self.hc_eps, + self.hc_eps, + self.hc_post_alpha, + self.hc_sinkhorn_iters, + ) + + x = self.attn_norm(x) + x = self.attn(positions, x, None) + + residual, post_mix, res_mix, x = self.mhc_fused_post_pre( + x, + residual, + post_mix, + res_mix, + self.hc_ffn_fn, + self.hc_ffn_scale, + self.hc_ffn_base, + self.rms_norm_eps, + self.hc_eps, + self.hc_eps, + self.hc_post_alpha, + self.hc_sinkhorn_iters, + ) + x = self.ffn_norm(x) + x = self.ffn(x, input_ids) + return x, residual, post_mix, res_mix + + def _forward_unfused_post_pre( + self, + x: torch.Tensor, + positions: torch.Tensor, + input_ids: torch.Tensor | None, + post_mix: torch.Tensor | None = None, + res_mix: torch.Tensor | None = None, + residual: torch.Tensor | None = None, + ) -> tuple[ + torch.Tensor, torch.Tensor | None, torch.Tensor | None, torch.Tensor | None + ]: + residual = x + x, post, comb = self.hc_pre( + x, self.hc_attn_fn, self.hc_attn_scale, self.hc_attn_base + ) + x = self.attn_norm(x) + x = self.attn(positions, x, None) + x = self.hc_post(x, residual, post, comb) + + residual = x + x, post, comb = self.hc_pre( + x, self.hc_ffn_fn, self.hc_ffn_scale, self.hc_ffn_base + ) + x = self.ffn_norm(x) + x = self.ffn(x, input_ids) + x = self.hc_post(x, residual, post, comb) + return x, None, None, None + + def forward( + self, + x: torch.Tensor, + positions: torch.Tensor, + input_ids: torch.Tensor | None, + post_mix: torch.Tensor | None = None, + res_mix: torch.Tensor | None = None, + residual: torch.Tensor | None = None, + ) -> tuple[ + torch.Tensor, torch.Tensor | None, torch.Tensor | None, torch.Tensor | None + ]: + if not self.use_fused_mhc: + return self._forward_unfused_post_pre( + x, positions, input_ids, post_mix, res_mix, residual + ) + return self._forward_fused_post_pre( + x, positions, input_ids, post_mix, res_mix, residual + ) + + +@support_torch_compile( + dynamic_arg_dims={ + "input_ids": 0, + "positions": -1, + "intermediate_tensors": 0, + "inputs_embeds": 0, + } +) +class DeepseekV4Model(nn.Module, EagleModelMixin): + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): + super().__init__() + + config = vllm_config.model_config.hf_config + quant_config = vllm_config.quant_config + self.config = config + self.vocab_size = config.vocab_size + self.hc_eps = config.hc_eps + self.hc_mult = config.hc_mult + self.hc_dim = self.hc_mult * config.hidden_size + self.rms_norm_eps = config.rms_norm_eps + + self.device = current_platform.device_type + # Reserved topk indices buffer for all Indexer layers to reuse. + self.topk_indices_buffer = torch.empty( + vllm_config.scheduler_config.max_num_batched_tokens, + config.index_topk, + dtype=torch.int32, + device=self.device, + ) + + if get_pp_group().is_first_rank: + self.embed_tokens = VocabParallelEmbedding( + config.vocab_size, + config.hidden_size, + quant_config=quant_config, + prefix=f"{prefix}.embed_tokens", + ) + else: + self.embed_tokens = PPMissingLayer() + + self.start_layer, self.end_layer, self.layers = make_layers( + config.num_hidden_layers, + lambda prefix: DeepseekV4DecoderLayer( + vllm_config, + prefix=prefix, + topk_indices_buffer=self.topk_indices_buffer, + ), + prefix=f"{prefix}.layers", + ) + self.is_fused_shared_expert_enabled = is_model_fused_shared_expert_compatible( + self.layers, + DeepseekV4MoE, + "ffn", + ) + + if get_pp_group().is_last_rank: + self.norm = RMSNorm(config.hidden_size, self.rms_norm_eps) + else: + self.norm = PPMissingLayer() + + self.hc_head_fn = nn.Parameter( + torch.empty( + self.hc_mult, + self.hc_dim, + dtype=torch.float32, + ), + requires_grad=False, + ) + self.hc_head_base = nn.Parameter( + torch.empty( + self.hc_mult, + dtype=torch.float32, + ), + requires_grad=False, + ) + self.hc_head_scale = nn.Parameter( + torch.empty(1, dtype=torch.float32), + requires_grad=False, + ) + self.hc_head_op = HCHeadOp() + spec_config = vllm_config.speculative_config + needs_mtp_hidden_states = spec_config is not None and ( + spec_config.use_eagle() or spec_config.uses_draft_model() + ) + if get_pp_group().is_last_rank and needs_mtp_hidden_states: + self._mtp_hidden_buffer = torch.empty( + vllm_config.scheduler_config.max_num_batched_tokens, + self.hc_dim, + dtype=vllm_config.model_config.dtype, + device=self.device, + ) + else: + self._mtp_hidden_buffer = None + + def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor: + return self.embed_tokens(input_ids) + + def make_empty_intermediate_tensors( + self, + batch_size: int, + dtype: torch.dtype, + device: torch.device, + ) -> IntermediateTensors: + # PP intermediate tensors keep the (num_tokens, hc_mult, hidden_size) + # multi-stream shape until hc_head() collapses it. + return IntermediateTensors( + { + "hidden_states": torch.zeros( + (batch_size, self.hc_mult, self.config.hidden_size), + dtype=dtype, + device=device, + ), + } + ) + + def forward( + self, + input_ids: torch.Tensor, + positions: torch.Tensor, + intermediate_tensors: IntermediateTensors | None, + inputs_embeds: torch.Tensor | None = None, + ) -> torch.Tensor | IntermediateTensors: + if get_pp_group().is_first_rank: + if inputs_embeds is not None: + hidden_states = inputs_embeds + else: + hidden_states = self.embed_input_ids(input_ids) + hidden_states = hidden_states.unsqueeze(-2).repeat(1, self.hc_mult, 1) + else: + assert intermediate_tensors is not None + hidden_states = intermediate_tensors["hidden_states"] + + residual, post_mix, res_mix = None, None, None + # EAGLE3/DSpark/DFlash aux hidden states: post-mhc hidden state at + # each target layer, averaged over hc_mult streams to [T, + # hidden_size]. Empty unless a draft model set aux_hidden_state_layers. + aux_hidden_states: list[torch.Tensor] = [] + # Reused below to avoid computing hc_post twice if the last layer is + # also an aux layer. + final_aux_recon: torch.Tensor | None = None + for idx, layer in enumerate( + islice(self.layers, self.start_layer, self.end_layer), + start=self.start_layer, + ): + hidden_states, residual, post_mix, res_mix = layer( + hidden_states, + positions, + input_ids, + post_mix, + res_mix, + residual, + ) + if (idx + 1) in self.aux_hidden_state_layers: + # Unfused path already applied hc_post; fused path needs it + # applied here before averaging. + if layer.use_fused_mhc: + aux_recon = layer.hc_post( + hidden_states, residual, post_mix, res_mix + ) + final_aux_recon = aux_recon + else: + aux_recon = hidden_states + aux_hidden_states.append(aux_recon.mean(dim=1)) + if layer is not None and layer.use_fused_mhc: + # Reuse the last layer's hc_post output if already computed above. + if ( + final_aux_recon is not None + and self.end_layer in self.aux_hidden_state_layers + ): + hidden_states = final_aux_recon + else: + hidden_states = layer.hc_post( + hidden_states, residual, post_mix, res_mix + ) + + if not get_pp_group().is_last_rank: + return IntermediateTensors({"hidden_states": hidden_states}) + + if self._mtp_hidden_buffer is not None: + num_tokens = hidden_states.shape[0] + self._mtp_hidden_buffer[:num_tokens].copy_(hidden_states.flatten(1)) + + hidden_states = self.hc_head_op( + hidden_states, + self.hc_head_fn, + self.hc_head_scale, + self.hc_head_base, + self.rms_norm_eps, + self.hc_eps, + ) + hidden_states = self.norm(hidden_states) + if len(aux_hidden_states) > 0: + return hidden_states, aux_hidden_states + return hidden_states + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: + stacked_params_mapping = [ + # (param_name, shard_name, shard_id) + ("gate_up_proj", "w1", 0), + ("gate_up_proj", "w3", 1), + ("attn.fused_wqa_wkv", "attn.wq_a", 0), + ("attn.fused_wqa_wkv", "attn.wkv", 1), + ("compressor.fused_wkv_wgate", "compressor.wkv", 0), + ("compressor.fused_wkv_wgate", "compressor.wgate", 1), + ] + params_dict = dict(self.named_parameters()) + loaded_params: set[str] = set() + + def _resolve_param_name(name: str) -> str: + inv_name = f"{name}_inv" + if name not in params_dict and inv_name in params_dict: + return inv_name + return name + + # TP for attention + tp_size = get_tensor_model_parallel_world_size() + tp_rank = get_tensor_model_parallel_rank() + n_head = self.config.num_attention_heads + n_local_head = n_head // tp_size + head_rank_start = n_local_head * tp_rank + head_rank_end = n_local_head * (tp_rank + 1) + + # Pre-compute expert mapping ONCE. + expert_mapping = self.get_expert_mapping() + + # Each MoE's own per-layer fusion decision is the source of truth, + # so the redirect below can't diverge from how it was built. + fuse_by_layer = { + extract_layer_index(mod_name): mod.is_fused_shared_expert_enabled + for mod_name, mod in self.named_modules() + if isinstance(mod, DeepseekV4MoE) + } + n_routed = self.config.n_routed_experts + # Maps the single shared-expert tensor group to one appended slot; + # multiple shared experts would need per-expert slicing (see + # deepseek_v2.py). DeepSeek-V4 has n_shared_experts == 1. + if any(fuse_by_layer.values()) and self.config.n_shared_experts != 1: + raise NotImplementedError( + "deepseek-v4 fused shared-expert loading supports only " + f"n_shared_experts == 1, got {self.config.n_shared_experts}" + ) + + for name, loaded_weight in weights: + # Shared-expert fusion: redirect shared_experts.w{1,2,3} into + # the appended routed-expert slot so it loads through the + # routed expert (grouped GEMM) loader. + if ".ffn.shared_experts.w" in name and fuse_by_layer.get( + extract_layer_index(name), False + ): + name = name.replace( + ".ffn.shared_experts.w", + f".ffn.experts.{n_routed}.w", + ) + + for param_name, weight_name, shard_id in stacked_params_mapping: + # Skip non-stacked layers and experts (experts handled below). + if ".experts." in name: + continue + if weight_name not in name: + continue + name = name.replace(weight_name, param_name) + + if is_pp_missing_parameter(name, self): + break + param_name = _resolve_param_name(name) + param = params_dict[param_name] + weight_loader = param.weight_loader + weight_loader(param, loaded_weight, shard_id) + loaded_params.add(param_name) + break + else: + if ".experts." in name: + # E8M0 scales are float8_e8m0fnu in checkpoints but the + # MoE param is uint8; copy_() would numerically convert + # (e.g. 2^-7 -> 0) instead of preserving raw bytes. + if ( + "weight_scale" in name + and loaded_weight.dtype == torch.float8_e8m0fnu + ): + loaded_weight = loaded_weight.view(torch.uint8) + for mapping in expert_mapping: + param_name, weight_name, expert_id, expert_shard_id = mapping + if weight_name not in name: + continue + name_mapped = name.replace(weight_name, param_name) + if is_pp_missing_parameter(name_mapped, self): + continue + param = params_dict[name_mapped] + # return_success lets us skip experts that have + # another replica available. + weight_loader = typing.cast( + Callable[..., bool], param.weight_loader + ) + success = weight_loader( + param, + loaded_weight, + name_mapped, + shard_id=expert_shard_id, + expert_id=expert_id, + return_success=True, + ) + if success: + name = name_mapped + break + loaded_params.add(name_mapped) + continue + elif "attn_sink" in name: + if is_pp_missing_parameter(name, self): + continue + narrow_weight = loaded_weight[head_rank_start:head_rank_end] + n = narrow_weight.shape[0] + params_dict[name][:n].copy_(narrow_weight) + loaded_params.add(name) + continue + else: + if is_pp_missing_parameter(name, self): + continue + param_name = _resolve_param_name(name) + param = params_dict[param_name] + weight_loader = getattr( + param, "weight_loader", default_weight_loader + ) + weight_loader(param, loaded_weight) + loaded_params.add(param_name) + continue + + return loaded_params + + def get_expert_mapping(self) -> list[tuple[str, str, int, str]]: + # (param_name, weight_name, expert_id, shard_id). When fusing shared + # experts, include the appended slots so redirected shared-expert + # weights route through the expert loader. + n_shared = getattr(self.config, "n_shared_experts", 0) or 0 + num_experts = self.config.n_routed_experts + ( + n_shared if self.is_fused_shared_expert_enabled else 0 + ) + return fused_moe_make_expert_params_mapping( + self, + ckpt_gate_proj_name="w1", + ckpt_down_proj_name="w2", + ckpt_up_proj_name="w3", + num_experts=num_experts, + ) + + +def _make_deepseek_v4_weights_mapper( + expert_dtype: str, fuse_shared_experts: bool = False +) -> WeightsMapper: + if expert_dtype == "fp4": + # MXFP4 experts register scales as ``w{1,2,3}_weight_scale`` (no + # _inv); FP8 linear/shared-expert scales register as + # ``weight_scale_inv``. Route DeepSeek-native ``.scale`` and + # AMD-Quark ``.weight_scale`` keys accordingly, leaving expert/ + # shared-expert ``w{1,2,3}.weight_scale`` untouched for the MXFP4 + # loader. + scale_regex = { + re.compile(r"(\.experts\.\d+\.w[123])\.scale$"): r"\1.weight_scale", + re.compile(r"\.scale$"): ".weight_scale_inv", + re.compile(r"(? torch.Tensor: + return self.model.embed_input_ids(input_ids) + + def compute_logits( + self, + hidden_states: torch.Tensor, + ) -> torch.Tensor | None: + logits = self.logits_processor(self.lm_head, hidden_states) + return logits + + def forward( + self, + input_ids: torch.Tensor, + positions: torch.Tensor, + intermediate_tensors: IntermediateTensors | None = None, + inputs_embeds: torch.Tensor | None = None, + ) -> torch.Tensor | IntermediateTensors: + hidden_states = self.model( + input_ids, positions, intermediate_tensors, inputs_embeds + ) + return hidden_states + + def get_mtp_target_hidden_states(self) -> torch.Tensor | None: + """Pre-hc_head residual stream buffer (max_num_batched_tokens, + hc_mult * hidden_size) for the MTP draft model. Populated by + forward(); valid after each target step.""" + return getattr(self.model, "_mtp_hidden_buffer", None) + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: + loader = AutoWeightsLoader(self) + return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper) + + def process_weights_after_loading(self) -> None: + # No CPU AMX weight-packing hook yet. + pass + + def get_expert_mapping(self) -> list[tuple[str, str, int, str]]: + return self.model.get_expert_mapping() diff --git a/vllm/models/deepseek_v4/cpu/mtp.py b/vllm/models/deepseek_v4/cpu/mtp.py new file mode 100644 index 0000000000..c723047a89 --- /dev/null +++ b/vllm/models/deepseek_v4/cpu/mtp.py @@ -0,0 +1,19 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""CPU DeepSeek-V4 MTP (speculative decoding) — deferred, not implemented. + +``DeepSeekV4MTP`` is only instantiated when +``vllm/model_executor/models/registry.py``'s speculative-decoding registry +resolves it by name for a speculative-decoding config; a plain (non-spec) +DeepSeek-V4 CPU model never constructs this class. This stub exists solely +so ``vllm.models.deepseek_v4.cpu`` is import-clean. +""" + +from vllm.config import VllmConfig + + +class DeepSeekV4MTP: + def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None: + raise NotImplementedError( + "DeepSeek-V4 speculative decoding (MTP) is not implemented on CPU." + ) diff --git a/vllm/platforms/cpu.py b/vllm/platforms/cpu.py index b0b5228829..ecd2dc5fe6 100644 --- a/vllm/platforms/cpu.py +++ b/vllm/platforms/cpu.py @@ -203,14 +203,26 @@ class CpuPlatform(Platform): cache_config = vllm_config.cache_config + is_deepseek_v4 = ( + model_config is not None + and getattr(model_config.hf_config, "model_type", None) == "deepseek_v4" + ) + # The CPU MLA decode kernel only compiles with block_size=16 today # (see csrc/cpu/mla_decode.cpp). If the model uses MLA we override # the default block size regardless of user preference to avoid a # runtime kernel dispatch failure. AMX MLA has no such constraint # (same AMX-available condition as get_attn_backend_cls), so it's - # excluded from this override. - cpu_mla_enabled = model_config is not None and getattr( - model_config, "use_mla", False + # excluded from this override. DeepSeek-V4 is also excluded here and + # handled in its own branch below: its sparse-MLA and indexer + # backends declare block_size=256 as their only supported kernel + # block size (DeepseekV4SparseMLABackend/DeepseekV4IndexerBackend), + # so it needs the same override-regardless-of-preference treatment + # as CPU MLA, just with a different value. + cpu_mla_enabled = ( + not is_deepseek_v4 + and model_config is not None + and getattr(model_config, "use_mla", False) ) amx_mla_enabled = ( cpu_mla_enabled @@ -219,6 +231,21 @@ class CpuPlatform(Platform): and vllm_config.attention_config.backend != AttentionBackendEnum.CPU_MLA ) reference_cpu_mla_enabled = cpu_mla_enabled and not amx_mla_enabled + # DeepSeek-V4's CPU attention/indexer kernels + # (csrc/cpu/sgl-kernels/{flash_mla,store_cache,compressor, + # paged_mqa_logits,topk}.cpp) are AMX-kernel-backed and built on the + # same paged/position-indexed conventions as the GPU/XPU backends + # (block_table/slot_mapping addressing throughout, chunk metadata + # that already carries per-token causal offsets for partial/extend + # continuation) -- so they support chunked prefill and prefix + # caching the same way. `amx_mla_enabled` above deliberately + # excludes DeepSeek-V4 (it has its own block-size requirements, + # unrelated to chunked-prefill support), so it can't be reused here. + amx_mla_or_dsv4_enabled = amx_mla_enabled or ( + is_deepseek_v4 + and cls.get_cpu_architecture() == CpuArchEnum.X86 + and torch.cpu._is_amx_tile_supported() + ) if reference_cpu_mla_enabled: if cache_config.user_specified_block_size and cache_config.block_size != 16: logger.warning( @@ -227,6 +254,17 @@ class CpuPlatform(Platform): cache_config.block_size, ) cache_config.block_size = 16 + elif is_deepseek_v4: + if ( + cache_config.user_specified_block_size + and cache_config.block_size != 256 + ): + logger.warning( + "DeepSeek-V4 CPU backend requires block_size=256, " + "overriding user-specified block_size=%s.", + cache_config.block_size, + ) + cache_config.block_size = 256 elif not cache_config.user_specified_block_size: cache_config.block_size = 128 @@ -439,7 +477,11 @@ class CpuPlatform(Platform): vllm_config.parallel_config.tensor_parallel_size ) - if model_config is not None and model_config.use_mla and not amx_mla_enabled: + if ( + model_config is not None + and model_config.use_mla + and not amx_mla_or_dsv4_enabled + ): logger.info_once( "MLA is enabled on a non-GPU platform; forcing chunked " "prefill and prefix caching to be disabled." diff --git a/vllm/v1/attention/backends/mla/indexer.py b/vllm/v1/attention/backends/mla/indexer.py index 6198c14eb7..5ec07ad211 100644 --- a/vllm/v1/attention/backends/mla/indexer.py +++ b/vllm/v1/attention/backends/mla/indexer.py @@ -177,55 +177,6 @@ class PrepareUniformDecodeKernel( ) -def split_indexer_prefill_chunks( - seq_lens_cpu: torch.Tensor, - query_lens_cpu: torch.Tensor, - workspace_size: int, - max_logits_bytes: int, - request_offset: int = 0, -) -> list[tuple[slice, slice]]: - """ - Split prefill requests into chunks for the sparse indexer, respecting: - - N constraint: total_seq_lens <= workspace_size (existing O(N) workspace) - - Logits constraint: M * N * 4 <= max_logits_bytes - - When a single request-level chunk still exceeds the logits budget, - sub-chunks on the query dimension (M) to bound peak memory. - - Returns list of (req_slice, query_slice) tuples. - """ - chunks: list[tuple[slice, slice]] = [] - n = len(seq_lens_cpu) - max_logits_elems = max_logits_bytes // 4 - end = 0 - - while end < n: - start, chunk_m, chunk_n = end, 0, 0 - - while end < n: - q, s = query_lens_cpu[end].item(), seq_lens_cpu[end].item() - new_m, new_n = chunk_m + q, chunk_n + s - if new_n <= workspace_size and new_m * new_n <= max_logits_elems: - chunk_m, chunk_n = new_m, new_n - end += 1 - else: - break - - # A single request can exceed the budget, requiring sub-chunking - # on the query dimension. - if end == start: - chunk_m, chunk_n = query_lens_cpu[end].item(), seq_lens_cpu[end].item() - end += 1 - - req_slice = slice(start + request_offset, end + request_offset) - max_q = max(1, max_logits_elems // chunk_n) if chunk_n > 0 else max(1, chunk_m) - for q_off in range(0, chunk_m, max_q): - sub_m = min(max_q, chunk_m - q_off) - chunks.append((req_slice, slice(q_off, q_off + sub_m))) - - return chunks - - class DeepseekV32IndexerBackend(AttentionBackend): @classmethod def supports_pcp(cls) -> bool: @@ -1020,6 +971,63 @@ class DeepseekV32IndexerMetadataBuilder(AttentionMetadataBuilder): ) return indices + @staticmethod + def _split_indexer_prefill_chunks( + compressed_seq_lens_cpu: torch.Tensor, + prefill_query_lens_cpu: torch.Tensor, + workspace_size: int, + max_logits_bytes: int, + request_offset: int = 0, + ) -> list[tuple[slice, slice]]: + """Split this step's prefill requests into chunks, respecting: + - N constraint: total_seq_lens <= workspace_size (existing O(N) + workspace) + - Logits constraint: M * N * 4 <= max_logits_bytes + + When a single request-level chunk still exceeds the logits budget, + sub-chunks on the query dimension (M) to bound peak memory. + + Returns list of (req_slice, query_slice) tuples. + """ + chunks: list[tuple[slice, slice]] = [] + n = len(compressed_seq_lens_cpu) + max_logits_elems = max_logits_bytes // 4 + end = 0 + + while end < n: + start, chunk_m, chunk_n = end, 0, 0 + + while end < n: + q, s = ( + prefill_query_lens_cpu[end].item(), + compressed_seq_lens_cpu[end].item(), + ) + new_m, new_n = chunk_m + q, chunk_n + s + if new_n <= workspace_size and new_m * new_n <= max_logits_elems: + chunk_m, chunk_n = new_m, new_n + end += 1 + else: + break + + # A single request can exceed the budget, requiring sub-chunking + # on the query dimension. + if end == start: + chunk_m, chunk_n = ( + prefill_query_lens_cpu[end].item(), + compressed_seq_lens_cpu[end].item(), + ) + end += 1 + + req_slice = slice(start + request_offset, end + request_offset) + max_q = ( + max(1, max_logits_elems // chunk_n) if chunk_n > 0 else max(1, chunk_m) + ) + for q_off in range(0, chunk_m, max_q): + sub_m = min(max_q, chunk_m - q_off) + chunks.append((req_slice, slice(q_off, q_off + sub_m))) + + return chunks + def build( self, common_prefix_len: int, @@ -1097,7 +1105,7 @@ class DeepseekV32IndexerMetadataBuilder(AttentionMetadataBuilder): # slice below). assert common_attn_metadata.seq_lens_cpu_upper_bound is not None seq_lens_cpu = common_attn_metadata.seq_lens_cpu_upper_bound - chunk_specs = split_indexer_prefill_chunks( + chunk_specs = self._split_indexer_prefill_chunks( compressed_seq_lens_cpu[num_decodes:], prefill_query_lens_cpu, self.max_prefill_buffer_size, diff --git a/vllm/v1/attention/backends/mla/sparse_swa.py b/vllm/v1/attention/backends/mla/sparse_swa.py index ccdf47f52e..0e4c360e42 100644 --- a/vllm/v1/attention/backends/mla/sparse_swa.py +++ b/vllm/v1/attention/backends/mla/sparse_swa.py @@ -148,6 +148,12 @@ class DeepseekSparseSWABackend(AttentionBackend): ) return DeepseekV4ROCMAiterSparseSWAMetadataBuilder + if current_platform.is_cpu(): + from vllm.models.deepseek_v4.cpu.cpu_mla import ( + DeepseekV4CPUSparseSWAMetadataBuilder, + ) + + return DeepseekV4CPUSparseSWAMetadataBuilder return DeepseekSparseSWAMetadataBuilder