[Model] Add DeepSeek-V4 CPU backend (#55355)
Signed-off-by: jiang1.li <jiang1.li@intel.com>
This commit is contained in:
Li, Jiang
@@ -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
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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"
|
||||
|
||||
@@ -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 <typename func_t>
|
||||
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
|
||||
|
||||
@@ -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 <algorithm>
|
||||
#include <array>
|
||||
#include <cmath>
|
||||
#include <limits>
|
||||
|
||||
#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<float>;
|
||||
|
||||
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<const __m256i*>(src));
|
||||
__m256i bf16_1 = _mm256_loadu_si256(reinterpret_cast<const __m256i*>(src + 16));
|
||||
__m256i bf16_2 = _mm256_loadu_si256(reinterpret_cast<const __m256i*>(src + 32));
|
||||
__m256i bf16_3 = _mm256_loadu_si256(reinterpret_cast<const __m256i*>(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<int>(ceil_log2);
|
||||
uint8_t scale_uint8 = static_cast<uint8_t>(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<float, kHeadDim> max_v;
|
||||
std::array<float, kHeadDim> denom;
|
||||
std::array<float, kHeadDim> acc;
|
||||
std::array<float, kHeadDim> normed;
|
||||
std::array<at::BFloat16, kNopeDim> 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<float>::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<float>(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<at::BFloat16>(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<at::BFloat16*>(token_bf16_ptr);
|
||||
for (int64_t ri = 0; ri < kRopeDim; ri += 16)
|
||||
store_from_float_ext<at::BFloat16>(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<float, kIndexerHeadDim> max_v;
|
||||
std::array<float, kIndexerHeadDim> denom;
|
||||
std::array<float, kIndexerHeadDim> acc;
|
||||
std::array<float, kIndexerHeadDim> normed;
|
||||
std::array<at::BFloat16, kIndexerHeadDim> 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<float>::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<float>(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<at::BFloat16>(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<const __m256i*>(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<const __m256i*>(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<float*>(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<float>(),
|
||||
state_cache.stride(0),
|
||||
block_size,
|
||||
kv.data_ptr<float>(),
|
||||
kv.stride(0),
|
||||
score.data_ptr<float>(),
|
||||
score.stride(0),
|
||||
ape.data_ptr<float>(),
|
||||
positions.data_ptr<int64_t>(),
|
||||
slot_mapping.data_ptr<int64_t>(),
|
||||
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<float>(),
|
||||
state_block_stride,
|
||||
state_block_size,
|
||||
gather_slots.data_ptr<int64_t>(),
|
||||
positions.data_ptr<int64_t>(),
|
||||
kv_slot_mapping.data_ptr<int64_t>(),
|
||||
rms_norm_weight.data_ptr<float>(),
|
||||
static_cast<float>(rms_norm_eps),
|
||||
cos_sin_cache.data_ptr<float>(),
|
||||
kv_cache_2d.data_ptr<uint8_t>(),
|
||||
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<float>(),
|
||||
state_block_stride,
|
||||
state_block_size,
|
||||
gather_slots.data_ptr<int64_t>(),
|
||||
positions.data_ptr<int64_t>(),
|
||||
kv_slot_mapping.data_ptr<int64_t>(),
|
||||
rms_norm_weight.data_ptr<float>(),
|
||||
static_cast<float>(rms_norm_eps),
|
||||
cos_sin_cache.data_ptr<float>(),
|
||||
kv_cache_2d.data_ptr<uint8_t>(),
|
||||
cache_row_stride,
|
||||
kv_cache_block_size,
|
||||
compress_ratio,
|
||||
window,
|
||||
state_width,
|
||||
num_tokens);
|
||||
}
|
||||
@@ -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 <algorithm>
|
||||
#include <cmath>
|
||||
#include <limits>
|
||||
#include <vector>
|
||||
|
||||
#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<float>::quiet_NaN();
|
||||
if (v == 0) return 0.f;
|
||||
union {
|
||||
uint32_t u;
|
||||
float f;
|
||||
} u;
|
||||
u.u = static_cast<uint32_t>(v) << 23;
|
||||
return u.f;
|
||||
}
|
||||
|
||||
template <typename index_t>
|
||||
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<int64_t>(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<const __m256i*>(row_base + dim_offset)), scale);
|
||||
}
|
||||
const auto* rope_ptr = reinterpret_cast<const at::BFloat16*>(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 <bool convert_v, typename scalar_t, typename index_t, typename load_vec_t>
|
||||
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<int64_t>(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 <typename scalar_t, typename index_t>
|
||||
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<int64_t>(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<true, scalar_t, index_t>(
|
||||
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<false, scalar_t, index_t>(
|
||||
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 <typename scalar_t>
|
||||
inline void fmla_fill_stub(scalar_t* __restrict__ out, float val, int64_t size) {
|
||||
using Vec = at::vec::Vectorized<scalar_t>;
|
||||
constexpr int kVecSize = Vec::size();
|
||||
const Vec data_vec = Vec(static_cast<scalar_t>(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 <typename scalar_t, int BLOCK_N>
|
||||
inline void fmla_copy_stub(scalar_t* __restrict__ out, const float* __restrict__ input) {
|
||||
static_assert(BLOCK_N % 32 == 0);
|
||||
using bVec = at::vec::Vectorized<scalar_t>;
|
||||
using fVec = at::vec::Vectorized<float>;
|
||||
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<scalar_t>(a0, a1);
|
||||
out_bvec.store(out + col * 16);
|
||||
}
|
||||
};
|
||||
Unroll<COLS>{}(store);
|
||||
}
|
||||
|
||||
template <typename scalar_t>
|
||||
inline void fmla_finalize_out(scalar_t* __restrict__ out, const float* __restrict__ acc, float inv_s, int64_t size) {
|
||||
using bVec = at::vec::Vectorized<scalar_t>;
|
||||
using fVec = at::vec::Vectorized<float>;
|
||||
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<scalar_t>(a0, a1);
|
||||
out_bvec.store(out + d);
|
||||
}
|
||||
for (; d < size; ++d) {
|
||||
out[d] = static_cast<scalar_t>(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 <typename scalar_t, typename index_t, int64_t BLOCK_N>
|
||||
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<float>;
|
||||
|
||||
// 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<float*>(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<float>::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<int64_t>(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<scalar_t, index_t>(
|
||||
Btmp0, Btmp1, fp8_ptr, cur_idx_ptr + n, valid_mask, static_cast<int>(n_size), static_cast<int>(head_size),
|
||||
static_cast<int>(head_size_v), fp8_block_size, fp8_block_stride_bytes, static_cast<int>(BLOCK_N),
|
||||
static_cast<int>(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<float>([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<float>::infinity();
|
||||
}
|
||||
}
|
||||
|
||||
// online softmax update
|
||||
float m_i = at::vec::reduce_all<float>([](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<scalar_t, BLOCK_N>(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<float>([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<float>([](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<float>(
|
||||
[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<scalar_t, BLOCK_N>(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<float>::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<scalar_t>(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<at::BFloat16, int64_t, BLOCK_N>(
|
||||
out.data_ptr<at::BFloat16>(),
|
||||
lse.data_ptr<float>(),
|
||||
q.data_ptr<at::BFloat16>(),
|
||||
window_cache_2d.data_ptr<uint8_t>(),
|
||||
num_compressed > 0 ? compressed_cache_2d.data_ptr<uint8_t>() : nullptr,
|
||||
window_slots.data_ptr<int64_t>(),
|
||||
num_compressed > 0 ? compressed_slots.data_ptr<int64_t>() : nullptr,
|
||||
attn_sink.data_ptr<float>(),
|
||||
buffer.data_ptr<at::BFloat16>(),
|
||||
/* 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<float>(scale),
|
||||
buffer_size_per_thread);
|
||||
}
|
||||
|
||||
@@ -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 <algorithm>
|
||||
#include <cmath>
|
||||
#include <vector>
|
||||
|
||||
#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<float>(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<float> 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<int64_t>(), index_q.data_ptr<float>(),
|
||||
index_q.stride(0), index_q.stride(1),
|
||||
index_q_cos_sin_cache.data_ptr<float>(), index_q_cos_sin_cache.stride(0),
|
||||
half_rot_dim, index_q_fp8.data_ptr<at::Float8_e4m3fn>(),
|
||||
index_q_fp8.stride(0), index_q_fp8.stride(1), head_dim,
|
||||
index_weights.data_ptr<float>(), index_weights.stride(0),
|
||||
static_cast<float>(index_weights_softmax_scale),
|
||||
static_cast<float>(index_weights_head_scale),
|
||||
index_weights_out.data_ptr<float>(), index_weights_out.stride(0),
|
||||
num_tokens, num_heads);
|
||||
}
|
||||
File diff suppressed because it is too large
Load Diff
@@ -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"
|
||||
|
||||
@@ -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<at::Tensor>& w1_scale,
|
||||
const std::optional<at::Tensor>& 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
|
||||
|
||||
@@ -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<const __m256i*>(k + d));
|
||||
const __m256i q8 =
|
||||
_mm256_loadu_si256(reinterpret_cast<const __m256i*>(q + d));
|
||||
acc = _mm512_dpbf16_ps(acc, CVT_FP8_TO_BF16(k8), CVT_FP8_TO_BF16(q8));
|
||||
}
|
||||
return _mm512_reduce_add_ps(acc);
|
||||
}
|
||||
|
||||
template <typename seq_t, typename page_t>
|
||||
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<const uint8_t*>(q_fp8.const_data_ptr());
|
||||
const auto* cache_ptr =
|
||||
reinterpret_cast<const uint8_t*>(kvcache_fp8.const_data_ptr());
|
||||
const auto* weight_ptr = weight.const_data_ptr<float>();
|
||||
const auto* seq_ptr = seq_lens.const_data_ptr<seq_t>();
|
||||
const auto* page_ptr = page_table.const_data_ptr<page_t>();
|
||||
auto* out_ptr = logits.data_ptr<float>();
|
||||
|
||||
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<int64_t>(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<int64_t>(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<const float*>(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 <typename seq_t>
|
||||
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<seq_t, int32_t>(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<seq_t, int64_t>(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<int32_t>(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<int64_t>(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;
|
||||
}
|
||||
@@ -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 <algorithm>
|
||||
#include <cmath>
|
||||
#include <vector>
|
||||
|
||||
#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<const __m256i*>(src));
|
||||
__m256i bf16_1 = _mm256_loadu_si256(reinterpret_cast<const __m256i*>(src + 16));
|
||||
__m256i bf16_2 = _mm256_loadu_si256(reinterpret_cast<const __m256i*>(src + 32));
|
||||
__m256i bf16_3 = _mm256_loadu_si256(reinterpret_cast<const __m256i*>(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<int>(ceil_log2);
|
||||
uint8_t scale_uint8 = static_cast<uint8_t>(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 <typename scalar_t>
|
||||
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<float> 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<float>(q_th[i]);
|
||||
buf[i] = v;
|
||||
sq_sum += v * v;
|
||||
}
|
||||
const float inv_rms = 1.0f / std::sqrt(sq_sum / static_cast<float>(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<scalar_t>(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<float> buf(head_dim);
|
||||
std::vector<at::BFloat16> 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<float>(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<float>(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<at::BFloat16*>(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<scalar_t>(
|
||||
q_out.data_ptr<scalar_t>(),
|
||||
swa_kv_cache_2d.data_ptr<uint8_t>(),
|
||||
q.data_ptr<scalar_t>(),
|
||||
kv.data_ptr<scalar_t>(),
|
||||
positions.data_ptr<int64_t>(),
|
||||
slot_mapping.data_ptr<int64_t>(),
|
||||
cos_sin_cache.data_ptr<float>(),
|
||||
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<float>(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<scalar_t>();
|
||||
float* out_ptr = out.data_ptr<float>();
|
||||
const int64_t* pos_ptr = positions.data_ptr<int64_t>();
|
||||
const float* cos_sin_ptr = cos_sin_cache.data_ptr<float>();
|
||||
|
||||
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<scalar_t>(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<float> 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;
|
||||
}
|
||||
@@ -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 <algorithm>
|
||||
#include <array>
|
||||
#include <cmath>
|
||||
#include <limits>
|
||||
#include <type_traits>
|
||||
#include <utility>
|
||||
|
||||
#include "common.h"
|
||||
#include "vec.h"
|
||||
|
||||
namespace {
|
||||
|
||||
template <typename scalar_t, int SIZE, std::enable_if_t<!std::is_same_v<scalar_t, float>, int> = 0>
|
||||
inline void sigmoid(float* __restrict__ out, const scalar_t* __restrict__ input) {
|
||||
using bVec = at::vec::Vectorized<scalar_t>;
|
||||
using fVec = at::vec::Vectorized<float>;
|
||||
|
||||
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 <typename scalar_t, int SIZE, std::enable_if_t<std::is_same_v<scalar_t, float>, int> = 0>
|
||||
inline void sigmoid(float* __restrict__ out, const float* __restrict__ input) {
|
||||
using fVec = at::vec::Vectorized<float>;
|
||||
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 <typename scalar_t, int SIZE, std::enable_if_t<!std::is_same_v<scalar_t, float>, int> = 0>
|
||||
inline void softmax(float* __restrict__ out, const scalar_t* __restrict__ input) {
|
||||
using bVec = at::vec::Vectorized<scalar_t>;
|
||||
using fVec = at::vec::Vectorized<float>;
|
||||
|
||||
constexpr int kVecSize = bVec::size();
|
||||
|
||||
// step 1: get max
|
||||
fVec max_fvec = fVec(-std::numeric_limits<float>::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 <typename scalar_t, int SIZE, std::enable_if_t<std::is_same_v<scalar_t, float>, int> = 0>
|
||||
inline void softmax(float* __restrict__ out, const float* __restrict__ input) {
|
||||
using fVec = at::vec::Vectorized<float>;
|
||||
|
||||
constexpr int kVecSize = fVec::size();
|
||||
|
||||
// step 1: get max
|
||||
fVec max_fvec = fVec(-std::numeric_limits<float>::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 <typename param_t, int SIZE>
|
||||
inline void
|
||||
apply_bias(float* __restrict__ scores2, const float* __restrict__ scores, const param_t* __restrict__ bias) {
|
||||
using fVec = at::vec::Vectorized<float>;
|
||||
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 <typename scalar_t, int SIZE, std::enable_if_t<!std::is_same_v<scalar_t, float>, int> = 0>
|
||||
inline void sqrtsoftplus(float* __restrict__ out, const scalar_t* __restrict__ input) {
|
||||
using bVec = at::vec::Vectorized<scalar_t>;
|
||||
using fVec = at::vec::Vectorized<float>;
|
||||
|
||||
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 <typename scalar_t, int SIZE, std::enable_if_t<std::is_same_v<scalar_t, float>, int> = 0>
|
||||
inline void sqrtsoftplus(float* __restrict__ out, const float* __restrict__ input) {
|
||||
using fVec = at::vec::Vectorized<float>;
|
||||
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 <typename scalar_t, typename param_t, int NUM_EXPERTS, int TOPK>
|
||||
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<float, int32_t>;
|
||||
std::array<elem_t, NUM_EXPERTS> 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<scalar_t, NUM_EXPERTS>(scores, gating_output + i * NUM_EXPERTS);
|
||||
} else {
|
||||
sqrtsoftplus<scalar_t, NUM_EXPERTS>(scores, gating_output + i * NUM_EXPERTS);
|
||||
}
|
||||
|
||||
// add bias for selection
|
||||
apply_bias<param_t, NUM_EXPERTS>(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<int32_t>(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<int32_t>((rng_state >> 33) % static_cast<uint64_t>(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 <typename scalar_t, int NUM_EXPERTS, int TOPK>
|
||||
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<scalar_t, NUM_EXPERTS>(scores, gating_output + i * NUM_EXPERTS);
|
||||
} else if (scoring_func == 1) {
|
||||
sigmoid<scalar_t, NUM_EXPERTS>(scores, gating_output + i * NUM_EXPERTS);
|
||||
} else {
|
||||
sqrtsoftplus<scalar_t, NUM_EXPERTS>(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<int32_t>((rng_state >> 33) % static_cast<uint64_t>(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<scalar_t, NE, NTOPK>( \
|
||||
topk_weights.data_ptr<float>(), \
|
||||
topk_ids.data_ptr<int32_t>(), \
|
||||
gating_output.data_ptr<scalar_t>(), \
|
||||
tid2eid_flat.data_ptr<int32_t>(), \
|
||||
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<scalar_t, param_t, NE, NTOPK>( \
|
||||
topk_weights.data_ptr<float>(), \
|
||||
topk_ids.data_ptr<int32_t>(), \
|
||||
gating_output.data_ptr<scalar_t>(), \
|
||||
correction_bias.data_ptr<param_t>(), \
|
||||
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<at::Tensor, at::Tensor> 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<double> 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<at::Tensor, at::Tensor> 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<float>(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 = [&]<typename scalar_t>() {
|
||||
// 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()<at::BFloat16>();
|
||||
} else if (st == at::ScalarType::Half) {
|
||||
dispatch_fn.template operator()<at::Half>();
|
||||
} else if (st == at::ScalarType::Float) {
|
||||
dispatch_fn.template operator()<float>();
|
||||
} 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 <typename T>
|
||||
inline int64_t load_index_value(const T* __restrict__ ptr, int64_t idx) {
|
||||
return static_cast<int64_t>(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 <typename seq_t, typename page_t>
|
||||
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<int>(std::log2(static_cast<double>(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<TopKTransformElem, kC4Topk> 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<int64_t>(0, std::min<int64_t>(seq_len, max_seq_len));
|
||||
const int64_t valid_topk = std::min<int64_t>(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<int64_t>(raw_index) >> page_bits;
|
||||
const int64_t offset_in_page = static_cast<int64_t>(raw_index) & page_mask;
|
||||
const int64_t physical_page = load_index_value(page_table_row, page_idx);
|
||||
out_page_row[slot] = static_cast<int32_t>((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<int32_t>(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<int32_t>(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<int32_t>(i) < current_min.index)) {
|
||||
std::pop_heap(heap.begin(), heap.end(), TopKTransformMinHeapCmp());
|
||||
heap.back() = {score, static_cast<int32_t>(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<at::Tensor>& 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<int32_t>();
|
||||
}
|
||||
|
||||
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<seq_t, int32_t>(
|
||||
scores.data_ptr<float>(), seq_lens.data_ptr<seq_t>(), page_tables.data_ptr<int32_t>(),
|
||||
out_page_indices.data_ptr<int32_t>(), 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<seq_t, int64_t>(
|
||||
scores.data_ptr<float>(), seq_lens.data_ptr<seq_t>(), page_tables.data_ptr<int64_t>(),
|
||||
out_page_indices.data_ptr<int32_t>(), 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");
|
||||
}
|
||||
}
|
||||
+213
-16
@@ -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<at::Tensor>& w1_scale,
|
||||
const std::optional<at::Tensor>& w2_scale,
|
||||
const std::optional<at::Tensor>& w1_zero,
|
||||
const std::optional<at::Tensor>& w2_zero,
|
||||
const std::optional<std::vector<int64_t>> block_size,
|
||||
const std::optional<at::Tensor>& w1_bias,
|
||||
const std::optional<at::Tensor>& w2_bias,
|
||||
const std::optional<double>& alpha, const std::optional<double>& 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<at::Tensor>& w1_scale,
|
||||
const std::optional<at::Tensor>& w2_scale,
|
||||
const std::optional<at::Tensor>& w1_zero,
|
||||
const std::optional<at::Tensor>& w2_zero,
|
||||
const std::optional<std::vector<int64_t>> block_size,
|
||||
const std::optional<at::Tensor>& w1_bias,
|
||||
const std::optional<at::Tensor>& w2_bias,
|
||||
const std::optional<double>& alpha,
|
||||
const std::optional<double>& 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<at::Tensor, at::Tensor, at::Tensor> 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<at::Tensor>& out_raw_indices);
|
||||
|
||||
// Adapted from sglang: DeepSeek-V4 MoE routing kernels (flat biased top-k and
|
||||
// hash-routed-layer top-k).
|
||||
std::tuple<at::Tensor, at::Tensor> 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<double> routed_scaling_factor,
|
||||
bool apply_routed_scaling_factor_on_output);
|
||||
|
||||
std::tuple<at::Tensor, at::Tensor> 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<at::Tensor, at::Tensor, at::Tensor> 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__) && \
|
||||
|
||||
@@ -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,
|
||||
)
|
||||
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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,
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
@@ -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
|
||||
|
||||
|
||||
|
||||
@@ -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
|
||||
|
||||
+397
-37
@@ -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"):
|
||||
|
||||
@@ -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(
|
||||
|
||||
@@ -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",
|
||||
|
||||
@@ -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)
|
||||
@@ -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,
|
||||
)
|
||||
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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),
|
||||
|
||||
@@ -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."
|
||||
)
|
||||
|
||||
|
||||
|
||||
@@ -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
|
||||
)
|
||||
@@ -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:
|
||||
|
||||
@@ -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 (
|
||||
|
||||
@@ -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."""
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -0,0 +1,2 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
@@ -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,
|
||||
)
|
||||
@@ -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,
|
||||
}
|
||||
@@ -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,
|
||||
)
|
||||
@@ -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))
|
||||
@@ -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."
|
||||
)
|
||||
@@ -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"(?<!\.w[123])\.weight_scale$"): ".weight_scale_inv",
|
||||
}
|
||||
else:
|
||||
# FP8 experts use Fp8MoEMethod (block_quant=True), which registers
|
||||
# scales as ``w{13,2}_weight_scale_inv``. Map all ``.scale`` keys
|
||||
# there.
|
||||
scale_regex = {
|
||||
re.compile(r"\.scale$"): ".weight_scale_inv",
|
||||
}
|
||||
# Fused shared experts are redirected to routed expert slots elsewhere;
|
||||
# leave their names untouched here.
|
||||
substr_map: dict[str, str | None] = (
|
||||
{}
|
||||
if fuse_shared_experts
|
||||
else {".shared_experts.w2": ".shared_experts.down_proj"}
|
||||
)
|
||||
# No MTP support on CPU yet: drop mtp.* weights (no model.mtp submodule).
|
||||
substr_map["mtp."] = None
|
||||
return WeightsMapper(
|
||||
orig_to_new_prefix={
|
||||
"layers.": "model.layers.",
|
||||
"embed.": "model.embed.",
|
||||
"norm.": "model.norm.",
|
||||
"hc_head": "model.hc_head",
|
||||
"mtp.": "model.mtp.",
|
||||
},
|
||||
orig_to_new_regex=scale_regex,
|
||||
orig_to_new_suffix={
|
||||
"head.weight": "lm_head.weight",
|
||||
"embed.weight": "embed_tokens.weight",
|
||||
".ffn.gate.bias": ".ffn.gate.e_score_correction_bias",
|
||||
".input_scale": ".input_scale_2",
|
||||
},
|
||||
orig_to_new_substr=substr_map,
|
||||
)
|
||||
|
||||
|
||||
class DeepseekV4ForCausalLM(nn.Module, SupportsPP, SupportsEagle3):
|
||||
model_cls = DeepseekV4Model
|
||||
|
||||
# Default mapper assumes the original FP4-expert checkpoint layout.
|
||||
# Overridden per-instance in __init__ when expert_dtype != "fp4".
|
||||
hf_to_vllm_mapper = _make_deepseek_v4_weights_mapper("fp4")
|
||||
packed_modules_mapping = {
|
||||
"gate_up_proj": ["w1", "w3"],
|
||||
"fused_wqa_wkv": ["wq_a", "wkv"],
|
||||
}
|
||||
|
||||
def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
|
||||
super().__init__()
|
||||
|
||||
config = vllm_config.model_config.hf_config
|
||||
self.config = config
|
||||
expert_dtype = getattr(config, "expert_dtype", "fp4")
|
||||
self.model = self.model_cls(
|
||||
vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model")
|
||||
)
|
||||
if expert_dtype != "fp4" or self.model.is_fused_shared_expert_enabled:
|
||||
self.hf_to_vllm_mapper = _make_deepseek_v4_weights_mapper(
|
||||
expert_dtype,
|
||||
fuse_shared_experts=self.model.is_fused_shared_expert_enabled,
|
||||
)
|
||||
if get_pp_group().is_last_rank:
|
||||
self.lm_head = ParallelLMHead(
|
||||
config.vocab_size,
|
||||
config.hidden_size,
|
||||
prefix=maybe_prefix(prefix, "lm_head"),
|
||||
)
|
||||
else:
|
||||
self.lm_head = PPMissingLayer()
|
||||
self.logits_processor = LogitsProcessor(config.vocab_size)
|
||||
self.make_empty_intermediate_tensors = ( # type: ignore[method-assign]
|
||||
self.model.make_empty_intermediate_tensors
|
||||
)
|
||||
|
||||
def embed_input_ids(self, input_ids: torch.Tensor) -> 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()
|
||||
@@ -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."
|
||||
)
|
||||
+46
-4
@@ -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."
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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
|
||||
|
||||
|
||||
|
||||
Reference in New Issue
Block a user