From b6d7602914d1d54253e79b1cc5b8f9017152b6e0 Mon Sep 17 00:00:00 2001 From: blzheng Date: Mon, 17 Aug 2026 10:52:26 +0800 Subject: [PATCH] [CPU] Add support for Gemma4 on Xeon (#22498) Co-authored-by: Copilot <198982749+Copilot@users.noreply.github.com> Co-authored-by: jianan-gu Co-authored-by: Haotong Zou --- .../kernels/aot/csrc/cpu/aarch64/moe.cpp | 7 +- .../kernels/aot/csrc/cpu/activation.cpp | 2 +- python/sglang/kernels/aot/csrc/cpu/extend.cpp | 47 +++++++--- python/sglang/kernels/aot/csrc/cpu/gemm.h | 20 +++++ python/sglang/kernels/aot/csrc/cpu/moe.cpp | 25 ++++-- python/sglang/kernels/aot/csrc/cpu/moe.h | 20 +++++ .../sglang/kernels/aot/csrc/cpu/moe_fp8.cpp | 8 +- python/sglang/kernels/aot/csrc/cpu/rope.cpp | 44 +++++++++- .../aot/csrc/cpu/torch_extension_cpu.cpp | 18 ++-- python/sglang/kernels/aot/csrc/cpu/vec.h | 8 ++ python/sglang/srt/configs/update_config.py | 42 ++++++++- .../srt/layers/attention/intel_amx_backend.py | 36 +++++++- python/sglang/srt/layers/layernorm.py | 4 + python/sglang/srt/layers/quantization/fp8.py | 1 + .../sglang/srt/layers/quantization/mxfp4.py | 1 + .../sglang/srt/layers/quantization/unquant.py | 5 +- .../srt/layers/quantization/w8a8_int8.py | 1 + .../srt/layers/rotary_embedding/mrope.py | 3 +- .../srt/model_executor/cpu_graph_runner.py | 22 ++--- python/sglang/srt/models/gemma4_causal.py | 36 ++++++-- python/sglang/srt/models/gemma4_mm.py | 40 ++++++--- python/sglang/srt/models/gemma4_vision.py | 31 +++++-- python/sglang/srt/server_args.py | 12 ++- python/sglang/test/cpu_test_utils.py | 25 ++++-- test/registered/cpu/test_extend.py | 26 +++++- test/registered/cpu/test_moe.py | 49 +++++++++-- test/registered/cpu/test_rope.py | 87 +++++++++++++++++-- 27 files changed, 514 insertions(+), 106 deletions(-) diff --git a/python/sglang/kernels/aot/csrc/cpu/aarch64/moe.cpp b/python/sglang/kernels/aot/csrc/cpu/aarch64/moe.cpp index 171481fd3..1348ecda2 100644 --- a/python/sglang/kernels/aot/csrc/cpu/aarch64/moe.cpp +++ b/python/sglang/kernels/aot/csrc/cpu/aarch64/moe.cpp @@ -242,7 +242,8 @@ at::Tensor fused_experts_cpu( const std::optional& /*w2_bias*/, const std::optional& /*alpha*/, const std::optional& /*limit*/, - bool /*is_vnni*/) { + bool /*is_vnni*/, + const std::optional& activation) { const auto st = hidden_states.scalar_type(); CHECK_INPUT(hidden_states); CHECK_INPUT(w13); @@ -256,6 +257,10 @@ at::Tensor fused_experts_cpu( CHECK_EQ(topk_ids.scalar_type(), at::kInt); + if (activation.has_value() && activation.value() != "silu") { + TORCH_CHECK(false, "fused_experts_cpu on ARM64 only supports activation='silu', got: ", activation.value()); + } + // TODO: support topk_weights to be bf16 or fp16 in the kernel auto topk_weights_ = topk_weights.to(at::kFloat); diff --git a/python/sglang/kernels/aot/csrc/cpu/activation.cpp b/python/sglang/kernels/aot/csrc/cpu/activation.cpp index 8263e4ebe..658577bc1 100644 --- a/python/sglang/kernels/aot/csrc/cpu/activation.cpp +++ b/python/sglang/kernels/aot/csrc/cpu/activation.cpp @@ -167,7 +167,7 @@ at::Tensor gelu_and_mul_cpu(const at::Tensor& input) { num_tokens, d, [inv_sqrt2](float x) { return 0.5f * x * (1.f + std::erf(x * inv_sqrt2)); }, - [inv_sqrt2](Vec x) { return Vec(0.5f) * x * (Vec(1.f) + (x * Vec(inv_sqrt2)).erf()); }); + [](Vec x) { return fast_gelu(x); }); }); return out; diff --git a/python/sglang/kernels/aot/csrc/cpu/extend.cpp b/python/sglang/kernels/aot/csrc/cpu/extend.cpp index 4592b1425..db2d46857 100644 --- a/python/sglang/kernels/aot/csrc/cpu/extend.cpp +++ b/python/sglang/kernels/aot/csrc/cpu/extend.cpp @@ -63,6 +63,7 @@ void extend_attention_kernel_impl( int64_t sliding_window_size, bool is_prefix_skipped, bool is_cross_attn, + bool kv_from_cache, bool has_encoder_lens, bool has_sink) { // strides @@ -148,8 +149,9 @@ void extend_attention_kernel_impl( fill_stub(s_prime, 0.f, m_size); fill_stub(m_prime, -std::numeric_limits::infinity(), m_size); // stage 1: compute scores with prefix + // kv_from_cache has no stage 2, so cover the extend range and stop at the diagonal int kv_start = 0; - int kv_end = is_cross_attn ? encoder_lens[bs] : seq_len_prefix; + int kv_end = is_cross_attn ? encoder_lens[bs] : (kv_from_cache ? seq_len_prefix + m + m_size : seq_len_prefix); for (int n = kv_start; n < kv_end; n += BLOCK_N) { int n_size = std::min(BLOCK_N, kv_end - n); @@ -180,12 +182,31 @@ void extend_attention_kernel_impl( /* C */ s_i); for (int row = 0; row < m_size; ++row) { - if (sliding_window_size > 0) { + bool row_is_empty = false; + if (kv_from_cache) { + int first_future_col = seq_len_prefix + m + row + 1; + if (n >= first_future_col) { + row_is_empty = true; + } else if (first_future_col < n + n_size) { + fill_stub( + s_i + row * BLOCK_N + (first_future_col - n), + -std::numeric_limits::infinity(), + n + n_size - first_future_col); + } + } + if (!row_is_empty && sliding_window_size > 0) { int last_col = seq_len_prefix + row + m - sliding_window_size + 1; if (last_col >= n + n_size) { - continue; + row_is_empty = true; + } else { + fill_stub(s_i + row * BLOCK_N, -std::numeric_limits::infinity(), last_col - n); } - fill_stub(s_i + row * BLOCK_N, -std::numeric_limits::infinity(), last_col - n); + } + if (row_is_empty) { + // s_delta is reused across blocks - zero an empty row rather than + // skip it, or P @ V below applies the previous block's weights here + fill_stub(s_delta + row * BLOCK_N, 0.f, padded_n_size); + continue; } flash_attn_softmax::apply( s_i, s_delta, v_prime, s_prime, m_prime, m_size, n_size, padded_n_size, head_size_v, sm_scale, row); @@ -214,7 +235,7 @@ void extend_attention_kernel_impl( /* B */ Btmp, /* C */ v_prime); } // loop with seq_len_prefix - if (!is_cross_attn) { + if (!is_cross_attn && !kv_from_cache) { // stage 2: compute the triangle part int num_keys = std::min(seq_len_extend, m + BLOCK_M); for (int n = 0; n < num_keys; n += BLOCK_N) { @@ -400,6 +421,7 @@ inline int resize_buffer(at::Tensor& buffer, int num_threads, int head_size, int sliding_window_size, \ is_prefix_skipped, \ is_cross_attn, \ + kv_from_cache, \ has_encoder_lens, \ has_sink); \ } while (0) @@ -442,13 +464,13 @@ void extend_attention_cpu( std::optional encoder_lens, std::optional sinks, std::optional tree_mask) { - if (!is_cross_attn) { - TORCH_CHECK( - k_extend_opt.has_value() && v_extend_opt.has_value(), - "k_extend and v_extend are required for non-cross attention"); - } - // Since k_extend and v_extend are not used for cross attention, they can be initialized as k_buffer and v_buffer - // here. + TORCH_CHECK(k_extend_opt.has_value() == v_extend_opt.has_value(), "k_extend and v_extend must be given together"); + // A KV-shared layer (Gemma 4) passes no extend K/V - the layer it shares with + // already wrote them to the cache, so this kernel masks causally itself. Cross + // attention can also arrive without K/V but reads an encoder sequence that + // carries no causal order, so it keeps its own path. + const bool kv_from_cache = !is_cross_attn && !k_extend_opt.has_value(); + // unused when the range comes from the cache - bind them to the buffers auto k_extend = k_extend_opt.has_value() ? k_extend_opt.value() : k_buffer; auto v_extend = v_extend_opt.has_value() ? v_extend_opt.value() : v_buffer; @@ -540,6 +562,7 @@ void extend_attention_cpu( ", got ", tree_mask_t.numel()); TORCH_CHECK(!is_cross_attn, "extend: tree_mask is not supported for cross attention"); + TORCH_CHECK(!kv_from_cache, "extend: tree_mask is not supported for KV-shared layers"); // The window mask derives query positions from the row index // (seq_len_prefix + m + row), but tree-mask rows sit at their tree depth, // which is <= the row index; combining the two would over-mask the prefix. diff --git a/python/sglang/kernels/aot/csrc/cpu/gemm.h b/python/sglang/kernels/aot/csrc/cpu/gemm.h index 102779a7a..b8a583170 100644 --- a/python/sglang/kernels/aot/csrc/cpu/gemm.h +++ b/python/sglang/kernels/aot/csrc/cpu/gemm.h @@ -73,6 +73,26 @@ enum class CPUActMethod : int { gelu_and_mul = 2, }; +// swiglu is selected by its (alpha, limit) parameters rather than by name and +// takes precedence; the name is still validated, so it cannot pass silently. +inline CPUActMethod act_method_from_string(const std::optional& activation, bool has_swiglu_params) { + const bool unnamed_or_silu = !activation.has_value() || activation.value() == "silu"; + if (has_swiglu_params) { + TORCH_CHECK( + unnamed_or_silu || activation.value() == "swiglu", + "Unsupported activation with clamped swiglu parameters: ", + activation.value()); + return CPUActMethod::swiglu; + } + if (unnamed_or_silu) { + return CPUActMethod::silu_and_mul; + } + if (activation.value() == "gelu") { + return CPUActMethod::gelu_and_mul; + } + TORCH_CHECK(false, "Unsupported activation: ", activation.value(), ". Supported: silu, gelu"); +} + enum class CPUQuantMethod : int64_t { BF16 = 0, INT8_W8A8 = 1, FP8_W8A16 = 2, INT4_W4A8 = 3, MXFP4 = 4 }; constexpr bool operator==(CPUQuantMethod a, int64_t b) { diff --git a/python/sglang/kernels/aot/csrc/cpu/moe.cpp b/python/sglang/kernels/aot/csrc/cpu/moe.cpp index 9b2ed699c..1a71bab5e 100644 --- a/python/sglang/kernels/aot/csrc/cpu/moe.cpp +++ b/python/sglang/kernels/aot/csrc/cpu/moe.cpp @@ -497,7 +497,8 @@ void fused_experts_kernel_impl( } else { const int64_t offset = offsets[mb]; - if (act_func == CPUActMethod::swiglu) { + // the fused path below hardcodes SiLU in its store step + if (act_func != CPUActMethod::silu_and_mul) { tinygemm_kernel( /* A */ A, /* B */ B0, @@ -539,12 +540,16 @@ void fused_experts_kernel_impl( add_bias_stub(C1 + m * BLOCK_N, B1_bias, n_size); } } - // 1.d silu and mul + // 1.d activation and mul const int64_t offset = offsets[mb]; if (act_func == CPUActMethod::silu_and_mul && use_brgemm) { for (int64_t m = 0; m < m_size; ++m) { silu_and_mul_stub(ic1 + (offset + m) * N + nb * BLOCK_N, C0 + m * BLOCK_N, C1 + m * BLOCK_N, BLOCK_N); } + } else if (act_func == CPUActMethod::gelu_and_mul) { + for (int64_t m = 0; m < m_size; ++m) { + gelu_and_mul_stub(ic1 + (offset + m) * N + nb * BLOCK_N, C0 + m * BLOCK_N, C1 + m * BLOCK_N, BLOCK_N); + } } else if (act_func == CPUActMethod::swiglu) { for (int64_t m = 0; m < m_size; ++m) { scalar_t* __restrict__ ic1_row = ic1 + (offset + m) * N; @@ -883,7 +888,17 @@ at::Tensor fused_experts_cpu( const std::optional& w2_bias, const std::optional& alpha, const std::optional& limit, - bool is_vnni) { + bool is_vnni, + const std::optional& activation) { + const CPUActMethod act_func = act_method_from_string(activation, alpha.has_value() && limit.has_value()); + // the int8 and int4 kernels hardcode silu in their fused store step + const bool is_int_quant = + moe_comp_method == CPUQuantMethod::INT8_W8A8 || moe_comp_method == CPUQuantMethod::INT4_W4A8; + TORCH_CHECK( + !is_int_quant || act_func == CPUActMethod::silu_and_mul, + "fused_experts_cpu: INT8_W8A8 and INT4_W4A8 support activation='silu' only, got: ", + activation.value_or("silu")); + auto packed_w1 = is_vnni ? w1 : convert_weight_packed(w1); auto packed_w2 = is_vnni ? w2 : convert_weight_packed(w2); @@ -909,6 +924,7 @@ at::Tensor fused_experts_cpu( } CHECK_DIM(2, topk_weights); CHECK_DIM(2, topk_ids_); + CHECK_EQ(topk_ids_.scalar_type(), at::kInt); // TODO: support topk_weights to be bf16 or fp16 in the kernel. @@ -1055,7 +1071,6 @@ at::Tensor fused_experts_cpu( scalar_t* __restrict__ intermediate_cache0 = (scalar_t*)((void*)(C_tmp + num_threads * 2 * BLOCK_M * BLOCK_N)); scalar_t* __restrict__ B_tmp = (scalar_t*)((void*)(intermediate_cache0 + M * topk * 2 * N)); bool with_bias = w1_bias.has_value(); - auto act_func = alpha.has_value() && limit.has_value() ? CPUActMethod::swiglu : CPUActMethod::silu_and_mul; CHECK_MOE_SCALES_FP8(1, 2); fused_experts_fp_kernel_impl( @@ -1095,7 +1110,6 @@ at::Tensor fused_experts_cpu( scalar_t* __restrict__ intermediate_cache0 = (scalar_t*)((void*)(C_tmp + num_threads * 2 * BLOCK_M * BLOCK_N)); scalar_t* __restrict__ B_tmp = (scalar_t*)((void*)(intermediate_cache0 + M * topk * 2 * N)); bool with_bias = w1_bias.has_value(); - auto act_func = alpha.has_value() && limit.has_value() ? CPUActMethod::swiglu : CPUActMethod::silu_and_mul; // mxfp4 supports only group size of 32 (2^5) constexpr int64_t group_size = 32; @@ -1180,7 +1194,6 @@ at::Tensor fused_experts_cpu( scalar_t* __restrict__ A_tmp = intermediate_cache2 + M * topk * K; float* __restrict__ C_tmp = (float*)((void*)(A_tmp + num_threads * BLOCK_M * K)); bool with_bias = w1_bias.has_value(); - auto act_func = alpha.has_value() && limit.has_value() ? CPUActMethod::swiglu : CPUActMethod::silu_and_mul; fused_experts_kernel_impl( out_hidden_states.data_ptr(), diff --git a/python/sglang/kernels/aot/csrc/cpu/moe.h b/python/sglang/kernels/aot/csrc/cpu/moe.h index 275a06b39..287551ca5 100644 --- a/python/sglang/kernels/aot/csrc/cpu/moe.h +++ b/python/sglang/kernels/aot/csrc/cpu/moe.h @@ -159,6 +159,26 @@ inline void silu_and_mul_stub( } } +template +inline void gelu_and_mul_stub( + scalar_t* __restrict__ out, const input_t* __restrict__ input, const input_t* __restrict__ input2, int64_t size) { + static_assert( + std::is_same_v || std::is_same_v, + "gelu_and_mul_stub only supports input_t == float or input_t == scalar_t"); + using bVec = at::vec::Vectorized; + + // no remainder +#pragma GCC unroll 4 + for (int64_t d = 0; d < size; d += bVec::size()) { + auto [x0, x1] = load_float_vec2(input + d); + auto [y0, y1] = load_float_vec2(input2 + d); + x0 = fast_gelu(x0) * y0; + x1 = fast_gelu(x1) * y1; + bVec out_vec = convert_from_float_ext(x0, x1); + out_vec.store(out + d); + } +} + template inline void clamp_sigmoid_and_mul_stub( scalar_t* __restrict__ out, const input_t* __restrict__ input, int64_t size, const float alpha, const float limit) { diff --git a/python/sglang/kernels/aot/csrc/cpu/moe_fp8.cpp b/python/sglang/kernels/aot/csrc/cpu/moe_fp8.cpp index 9dd2bb0ee..4ab29c354 100644 --- a/python/sglang/kernels/aot/csrc/cpu/moe_fp8.cpp +++ b/python/sglang/kernels/aot/csrc/cpu/moe_fp8.cpp @@ -116,13 +116,19 @@ void fused_experts_fp_kernel_impl( } }); - // stage 1.5: intermediate_cache1 = silu(intermediate_cache0) + // stage 1.5: intermediate_cache1 = activation(intermediate_cache0) if (act_func == CPUActMethod::silu_and_mul) { at::parallel_for(0, M * topk, 0, [&](int64_t begin, int64_t end) { for (int64_t m = begin; m < end; ++m) { silu_and_mul_stub(ic1 + m * N, ic0 + m * 2 * N, ic0 + m * 2 * N + N, N); } }); + } else if (act_func == CPUActMethod::gelu_and_mul) { + at::parallel_for(0, M * topk, 0, [&](int64_t begin, int64_t end) { + for (int64_t m = begin; m < end; ++m) { + gelu_and_mul_stub(ic1 + m * N, ic0 + m * 2 * N, ic0 + m * 2 * N + N, N); + } + }); } else if (act_func == CPUActMethod::swiglu) { at::parallel_for(0, M * topk, 0, [&](int64_t begin, int64_t end) { for (int64_t m = begin; m < end; ++m) { diff --git a/python/sglang/kernels/aot/csrc/cpu/rope.cpp b/python/sglang/kernels/aot/csrc/cpu/rope.cpp index 80fca4702..83c5b1ad0 100644 --- a/python/sglang/kernels/aot/csrc/cpu/rope.cpp +++ b/python/sglang/kernels/aot/csrc/cpu/rope.cpp @@ -265,6 +265,47 @@ void rotary_embedding_kernel_impl( } // namespace +// query: [num_tokens, num_heads, head_dim] +// key: [num_tokens, num_heads, head_dim] +// cos: [num_tokens, head_dim] +// sin: [num_tokens, head_dim] +// Gemma 4's vision tower rotates ndim = 2 head_dim chunks independently, so +// this is apply_rotary_pos_emb_cpu run once per chunk, in place. +void apply_multidimensional_rope_cpu(at::Tensor& query, at::Tensor& key, at::Tensor& cos, at::Tensor& sin) { + CHECK_DIM(3, query); + const auto input_dtype = query.scalar_type(); + int64_t num_tokens = query.size(0); + int64_t num_heads = query.size(1); + int64_t head_size = query.size(2); + + CHECK_LAST_DIM_CONTIGUOUS_INPUT(query); + CHECK_INPUT_SHAPE_DTYPE(key, {num_tokens, num_heads, head_size}, input_dtype); + CHECK_INPUT_SHAPE_DTYPE(cos, {num_tokens, head_size}, cos.scalar_type()); + CHECK_INPUT_SHAPE_DTYPE(sin, {num_tokens, head_size}, sin.scalar_type()); + CHECK_EQ(cos.scalar_type(), sin.scalar_type()); + + constexpr int64_t ndim = 2; + TORCH_CHECK( + head_size % (2 * ndim) == 0, "head_size must be divisible by ", 2 * ndim, " for ndim = ", ndim, " rotary chunks"); + const int64_t chunk_size = head_size / ndim; + + const RopeParams p{query, key, head_size, chunk_size}; + CPU_DISPATCH_REDUCED_FLOATING_TYPES_EXT(input_dtype, cos.scalar_type(), [&] { + scalar_t* q_ptr = query.data_ptr(); + scalar_t* k_ptr = key.data_ptr(); + const param_t* cos_ptr = cos.data_ptr(); + const param_t* sin_ptr = sin.data_ptr(); + for (int64_t d = 0; d < ndim; ++d) { + const int64_t offset = d * chunk_size; + auto cache_pos = [cos_ptr, sin_ptr, head_size, offset](int64_t token) -> SplitCosSinRow { + return {cos_ptr + token * head_size + offset, sin_ptr + token * head_size + offset}; + }; + rotary_embedding_kernel_impl( + q_ptr + offset, k_ptr + offset, q_ptr + offset, k_ptr + offset, p, cache_pos); + } + }); +} + // 2D : [num_tokens, num_heads*head_size] inplace // 3D : [num_tokens, num_heads, head_size] outplace // 4D : [batch_size, seq_len, num_heads, head_size] inplace @@ -376,7 +417,7 @@ apply_rotary_pos_emb_cpu(at::Tensor& query, at::Tensor& key, at::Tensor& cos, at // key: [num_tokens, num_kv_heads * head_size] // cos_sin_cache: [max_position_embeddings, rotary_dim] // mrope_section: [t, h, w] -std::tuple multimodal_rotary_embedding_cpu( +void multimodal_rotary_embedding_cpu( at::Tensor& positions, at::Tensor& query, at::Tensor& key, @@ -449,5 +490,4 @@ std::tuple multimodal_rotary_embedding_cpu( } } }); - return std::make_tuple(query, key); } diff --git a/python/sglang/kernels/aot/csrc/cpu/torch_extension_cpu.cpp b/python/sglang/kernels/aot/csrc/cpu/torch_extension_cpu.cpp index 51b83beb6..e9aca568b 100644 --- a/python/sglang/kernels/aot/csrc/cpu/torch_extension_cpu.cpp +++ b/python/sglang/kernels/aot/csrc/cpu/torch_extension_cpu.cpp @@ -361,7 +361,8 @@ at::Tensor fused_experts_cpu( const std::optional& w2_bias, const std::optional& alpha, const std::optional& limit, - bool is_vnni); + bool is_vnni, + const std::optional& activation); #if !defined(SGLANG_CPU_ARM64_SKIP_X86_ONLY_OPS) at::Tensor shared_expert_cpu( @@ -478,8 +479,11 @@ std::tuple rotary_embedding_cpu( std::tuple apply_rotary_pos_emb_cpu(at::Tensor& query, at::Tensor& key, at::Tensor& cos, at::Tensor& sin); +// multidimensional rope +void apply_multidimensional_rope_cpu(at::Tensor& query, at::Tensor& key, at::Tensor& cos, at::Tensor& sin); + // mrope -std::tuple multimodal_rotary_embedding_cpu( +void multimodal_rotary_embedding_cpu( at::Tensor& positions, at::Tensor& query, at::Tensor& key, @@ -795,7 +799,7 @@ TORCH_LIBRARY_FRAGMENT(sgl_kernel, m) { "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, " "Tensor? w1_zero, Tensor? w2_zero, int[]? block_size, Tensor? w1_bias, Tensor? w2_bias, float? alpha, float? " - "limit, bool is_vnni) -> Tensor"); + "limit, bool is_vnni, str? activation=None) -> Tensor"); m.impl("fused_experts_cpu", torch::kCPU, &fused_experts_cpu); #if !defined(SGLANG_CPU_ARM64_SKIP_X86_ONLY_OPS) @@ -864,10 +868,14 @@ TORCH_LIBRARY_FRAGMENT(sgl_kernel, m) { m.def("apply_rotary_pos_emb_cpu(Tensor query, Tensor key, Tensor cos, Tensor sin) -> (Tensor, Tensor)"); m.impl("apply_rotary_pos_emb_cpu", torch::kCPU, &apply_rotary_pos_emb_cpu); + // multidimensional rope + m.def("apply_multidimensional_rope_cpu(Tensor(a!) query, Tensor(b!) key, Tensor cos, Tensor sin) -> ()"); + m.impl("apply_multidimensional_rope_cpu", torch::kCPU, &apply_multidimensional_rope_cpu); + // multimodal rope m.def( - "multimodal_rotary_embedding_cpu(Tensor positions, Tensor query, Tensor key, int head_size, Tensor " - "cos_sin_cache, int[]? mrope_section, bool mrope_interleaved, bool is_neox) -> (Tensor, Tensor)"); + "multimodal_rotary_embedding_cpu(Tensor positions, Tensor(a!) query, Tensor(b!) key, int head_size, Tensor " + "cos_sin_cache, int[]? mrope_section, bool mrope_interleaved, bool is_neox) -> ()"); m.impl("multimodal_rotary_embedding_cpu", torch::kCPU, &multimodal_rotary_embedding_cpu); // CPU and memory binding diff --git a/python/sglang/kernels/aot/csrc/cpu/vec.h b/python/sglang/kernels/aot/csrc/cpu/vec.h index e0fa0c364..9f7334152 100644 --- a/python/sglang/kernels/aot/csrc/cpu/vec.h +++ b/python/sglang/kernels/aot/csrc/cpu/vec.h @@ -612,6 +612,14 @@ inline at::vec::Vectorized fast_silu(const at::vec::Vectorized& x) #endif } +// exact (erf) gelu, matching torch.nn.functional.gelu(approximate="none") +inline at::vec::Vectorized fast_gelu(const at::vec::Vectorized& x) { + const auto half = at::vec::Vectorized(0.5f); + const auto one = at::vec::Vectorized(1.f); + const auto inv_sqrt2 = at::vec::Vectorized(0.70710678118654752440f); + return half * x * (one + (x * inv_sqrt2).erf()); +} + inline at::vec::Vectorized fast_sigmoid_glu(const at::vec::Vectorized& x, const at::vec::Vectorized& alpha) { #if defined(CPU_CAPABILITY_AVX512) diff --git a/python/sglang/srt/configs/update_config.py b/python/sglang/srt/configs/update_config.py index 4a1968692..e0f0de2b0 100644 --- a/python/sglang/srt/configs/update_config.py +++ b/python/sglang/srt/configs/update_config.py @@ -137,6 +137,32 @@ def adjust_tp_num_heads_if_necessary(model_config, tp_size, is_post_update): ) +def adjust_swa_num_heads_if_necessary(model_config, tp_size, weight_block_size): + # Sliding-window layers carry their own head counts, so the padded + # full-attention num_attention_heads does not describe them + from sglang.srt.layers.vocab_parallel_embedding import pad_vocab_size + + text_config = model_config.hf_text_config + if not hasattr(text_config, "swa_num_key_value_heads"): + return + + swa_num_key_value_heads = text_config.swa_num_key_value_heads + swa_num_attention_heads = getattr( + text_config, "swa_num_attention_heads", model_config.num_attention_heads + ) + # ModelConfig always materializes swa_head_dim, defaulting it to head_dim. + swa_pad_size = get_num_heads_padding_size( + tp_size, weight_block_size, text_config.swa_head_dim + ) + padded_num_key_value_heads = pad_vocab_size(swa_num_key_value_heads, swa_pad_size) + padded_num_attention_heads = padded_num_key_value_heads * ( + swa_num_attention_heads // swa_num_key_value_heads + ) + + update_config(text_config, "swa_num_key_value_heads", padded_num_key_value_heads) + update_config(text_config, "swa_num_attention_heads", padded_num_attention_heads) + + def update_intermediate_size(model_config, attr_name, intermediate_padding_size): attr_value = intermediate_padding_size if ( @@ -152,7 +178,7 @@ def update_intermediate_size(model_config, attr_name, intermediate_padding_size) elif hasattr(model_config, attr_name): attr_value = getattr(model_config, attr_name) - if attr_value % intermediate_padding_size != 0: + if attr_value is not None and attr_value % intermediate_padding_size != 0: from sglang.srt.layers.vocab_parallel_embedding import pad_vocab_size origin_value = attr_value @@ -223,6 +249,10 @@ def adjust_config_with_unaligned_cpu_tp( + model_config.hf_config.qk_rope_head_dim, ) + # gate is on full-attention counts; Gemma 4 has 2 full-attention KV + # heads, so only TP 1 and 2 clear it and both align the 8 sliding heads + adjust_swa_num_heads_if_necessary(model_config, tp_size, weight_block_size) + query_heads_per_kv = ( model_config.num_attention_heads // model_config.get_total_num_kv_heads() ) @@ -306,9 +336,13 @@ def adjust_config_with_unaligned_cpu_tp( ) for m_config, config_name, model_type, num_head_str in multimodal_config: - if hasattr(m_config, config_name) and ( - m_config.model_type == model_type - or getattr(m_config, config_name).model_type == model_type + if ( + hasattr(m_config, config_name) + and getattr(m_config, config_name) is not None + and ( + m_config.model_type == model_type + or getattr(m_config, config_name).model_type == model_type + ) ): num_heads = getattr(getattr(m_config, config_name), num_head_str) diff --git a/python/sglang/srt/layers/attention/intel_amx_backend.py b/python/sglang/srt/layers/attention/intel_amx_backend.py index c54dd395a..7a3f138a1 100644 --- a/python/sglang/srt/layers/attention/intel_amx_backend.py +++ b/python/sglang/srt/layers/attention/intel_amx_backend.py @@ -59,6 +59,8 @@ class IntelAMXAttnBackend(AttentionBackend): # sized [bs, num_head, num_kv_splits, v_head_dim + 1] to match. self.num_kv_splits = 8 + self._attn_logits_buffers: dict[tuple[int, int], torch.Tensor] = {} + # speculative decoding params self.num_draft_tokens = get_spec().speculative_num_draft_tokens @@ -216,6 +218,9 @@ class IntelAMXAttnBackend(AttentionBackend): seq_lens = forward_batch.seq_lens if seq_lens.dtype != torch.int64: seq_lens = seq_lens.to(torch.int64) + + # Gemma4's KV-shared layers pass k=v=None - the layer they share with + # already wrote their extend K/V to the cache self.extend_attention_fwd( q.view(-1, layer.tp_q_head_num, layer.qk_head_dim), k, @@ -237,7 +242,7 @@ class IntelAMXAttnBackend(AttentionBackend): sinks, tree_mask, ) - return o + return o.view(-1, layer.tp_q_head_num * layer.v_head_dim) def forward_decode( self, @@ -249,8 +254,6 @@ class IntelAMXAttnBackend(AttentionBackend): save_kv_cache=True, sinks=None, ): - attn_logits, _ = self.forward_metadata - if self.draft_decode_metadata is not None: req_to_token, seq_lens, req_pool_indices = self.draft_decode_metadata else: @@ -263,6 +266,15 @@ class IntelAMXAttnBackend(AttentionBackend): if seq_lens.dtype != torch.int64: seq_lens = seq_lens.to(torch.int64) + if layer.v_head_dim == self.v_head_dim and layer.tp_q_head_num == self.num_head: + attn_logits, _ = self.forward_metadata + else: + # This layer's shape differs from the model-wide metadata buffer - + # size from the same seq_lens the kernel derives num_seqs from + attn_logits = self._get_attn_logits_buffer( + seq_lens.shape[0], layer.tp_q_head_num, layer.v_head_dim + ) + if layer.qk_head_dim != layer.v_head_dim: o = q.new_empty((q.shape[0], layer.tp_q_head_num * layer.v_head_dim)) else: @@ -291,7 +303,23 @@ class IntelAMXAttnBackend(AttentionBackend): forward_batch.encoder_lens, sinks, ) - return o + return o.view(-1, layer.tp_q_head_num * layer.v_head_dim) + + def _get_attn_logits_buffer( + self, num_seqs: int, num_heads: int, v_head_dim: int + ) -> torch.Tensor: + key = (num_heads, v_head_dim) + buffer = self._attn_logits_buffers.get(key) + if buffer is None or buffer.shape[0] < num_seqs: + # decode_attention_cpu writes every element it later reads, so the + # buffer needs no initialization and can be reused; it only grows. + buffer = torch.empty( + (num_seqs, num_heads, self.num_kv_splits, v_head_dim + 1), + dtype=torch.float32, + device=self.device, + ) + self._attn_logits_buffers[key] = buffer + return buffer[:num_seqs] def support_triton(self): return False diff --git a/python/sglang/srt/layers/layernorm.py b/python/sglang/srt/layers/layernorm.py index a3ee03580..c2a0ba5f7 100644 --- a/python/sglang/srt/layers/layernorm.py +++ b/python/sglang/srt/layers/layernorm.py @@ -1310,6 +1310,10 @@ class Gemma4RMSNorm(BaseFusedOp): def forward_cpu(self, x: torch.Tensor) -> torch.Tensor: if _is_cpu_amx_available: + # the kernel needs a last-dim-contiguous input; the audio conformer + # normalizes its depthwise conv output, which arrives permuted + if x.stride(-1) != 1: + x = x.contiguous() return torch.ops.sgl_kernel.gemma4_rmsnorm_cpu( x, self.weight.data, self.eps, self.scale_shift, self.with_scale ) diff --git a/python/sglang/srt/layers/quantization/fp8.py b/python/sglang/srt/layers/quantization/fp8.py index a77e1fbe4..3dee0a97d 100644 --- a/python/sglang/srt/layers/quantization/fp8.py +++ b/python/sglang/srt/layers/quantization/fp8.py @@ -2406,6 +2406,7 @@ class Fp8MoEMethod(FusedMoEMethodBase): None, # alpha None, # limit True, # is_vnni + moe_runner_config.activation, # activation ) return StandardCombineInput(hidden_states=output) diff --git a/python/sglang/srt/layers/quantization/mxfp4.py b/python/sglang/srt/layers/quantization/mxfp4.py index 8da877441..12ed0cf27 100644 --- a/python/sglang/srt/layers/quantization/mxfp4.py +++ b/python/sglang/srt/layers/quantization/mxfp4.py @@ -1417,6 +1417,7 @@ class Mxfp4MoEMethod(FusedMoEMethodBase): layer.moe_runner_config.gemm1_alpha, layer.moe_runner_config.gemm1_clamp_limit, True, # is_vnni + layer.moe_runner_config.activation, # activation ) else: from sglang.srt.layers.moe.fused_moe_native import moe_forward_native diff --git a/python/sglang/srt/layers/quantization/unquant.py b/python/sglang/srt/layers/quantization/unquant.py index ed89a2be1..eb05ade8b 100644 --- a/python/sglang/srt/layers/quantization/unquant.py +++ b/python/sglang/srt/layers/quantization/unquant.py @@ -849,10 +849,6 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, BaseFusedOp): moe_runner_config = self.moe_runner_config - assert ( - moe_runner_config.activation == "silu" - ), f"activation = {moe_runner_config.activation} is not supported." - if use_intel_amx_backend(layer): from sglang.srt.layers.moe.topk import apply_topk_weights_cpu @@ -878,6 +874,7 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, BaseFusedOp): layer.moe_runner_config.gemm1_alpha, layer.moe_runner_config.gemm1_clamp_limit, True, # is_vnni + moe_runner_config.activation, # activation ) return StandardCombineInput(hidden_states=output) else: diff --git a/python/sglang/srt/layers/quantization/w8a8_int8.py b/python/sglang/srt/layers/quantization/w8a8_int8.py index 6bcc3549f..8dc8dd2c5 100644 --- a/python/sglang/srt/layers/quantization/w8a8_int8.py +++ b/python/sglang/srt/layers/quantization/w8a8_int8.py @@ -380,6 +380,7 @@ class W8A8Int8MoEMethod(FusedMoEMethodBase): None, # alpha None, # limit True, # is_vnni + self.moe_runner_config.activation, # activation ) return StandardCombineInput(hidden_states=output) diff --git a/python/sglang/srt/layers/rotary_embedding/mrope.py b/python/sglang/srt/layers/rotary_embedding/mrope.py index dd1d5df95..b8b57bf97 100644 --- a/python/sglang/srt/layers/rotary_embedding/mrope.py +++ b/python/sglang/srt/layers/rotary_embedding/mrope.py @@ -224,7 +224,7 @@ class MRotaryEmbedding(RotaryEmbedding): fused_set_kv_buffer_arg=None, ) -> Tuple[torch.Tensor, torch.Tensor]: if _is_cpu_amx_available: - return torch.ops.sgl_kernel.multimodal_rotary_embedding_cpu( + torch.ops.sgl_kernel.multimodal_rotary_embedding_cpu( positions, query, key, @@ -234,6 +234,7 @@ class MRotaryEmbedding(RotaryEmbedding): self.mrope_interleaved, self.is_neox_style, ) + return query, key return self.forward_native(positions, query, key, fused_set_kv_buffer_arg) def forward_cuda( diff --git a/python/sglang/srt/model_executor/cpu_graph_runner.py b/python/sglang/srt/model_executor/cpu_graph_runner.py index 30719aa26..ab14d8f0b 100644 --- a/python/sglang/srt/model_executor/cpu_graph_runner.py +++ b/python/sglang/srt/model_executor/cpu_graph_runner.py @@ -176,6 +176,8 @@ def register_fake_ops(tp_size: int): "gemma_fused_add_rmsnorm_cpu", "layernorm_cpu", "fused_add_layernorm_cpu", + "multimodal_rotary_embedding_cpu", + "apply_multidimensional_rope_cpu", ] for op in none_return_ops: @@ -259,28 +261,14 @@ def register_fake_ops(tp_size: int): @register_cpu_compile_fake("rotary_embedding_cpu") def _(positions, query, key, head_size, cos_sin_cache, is_neox): - if query.ndim == 2: - return query, key - else: - return torch.empty_like(query), torch.empty_like(key) + # TODO: the kernel aliases query/key for 2D and 4D but allocates for 3D, + # which no schema expresses; an accurate fake needs it to pick one + return torch.empty_like(query), torch.empty_like(key) @register_cpu_compile_fake("apply_rotary_pos_emb_cpu") def _(query, key, cos, sin): return query, key - @register_cpu_compile_fake("multimodal_rotary_embedding_cpu") - def _( - positions, - query, - key, - head_size, - cos_sin_cache, - mrope_section, - mrope_interleaved, - is_neox, - ): - return query, key - @register_cpu_compile_fake("qkv_proj_with_rope_fused_weight") def _( hidden_states, diff --git a/python/sglang/srt/models/gemma4_causal.py b/python/sglang/srt/models/gemma4_causal.py index b37cc7dd4..4ebe925c5 100644 --- a/python/sglang/srt/models/gemma4_causal.py +++ b/python/sglang/srt/models/gemma4_causal.py @@ -74,6 +74,21 @@ Gemma4MLP = Gemma3MLP Gemma4TextScaledWordEmbedding = Gemma3TextScaledWordEmbedding +def load_tied_lm_head( + loaded_weight, *, params_dict, loaded_params, head_param_name="lm_head.weight" +): + """Load a tied embedding into an lm_head the runtime could not alias. + + No-op when this rank holds no lm_head. + """ + head_param = params_dict.get(head_param_name) + if head_param is None: + return + wl = getattr(head_param, "weight_loader", default_weight_loader) + wl(head_param, loaded_weight) + loaded_params.add(head_param_name) + + def pp_filter_load_weight( name, loaded_weight, @@ -109,11 +124,12 @@ def pp_filter_load_weight( return True if tie_word_embeddings and pp_group.is_last_rank and name == embed_weight_name: - head_param = params_dict.get(head_param_name) - if head_param is not None: - wl = getattr(head_param, "weight_loader", default_weight_loader) - wl(head_param, loaded_weight) - loaded_params.add(head_param_name) + load_tied_lm_head( + loaded_weight, + params_dict=params_dict, + loaded_params=loaded_params, + head_param_name=head_param_name, + ) return True if not pp_group.is_first_rank and any(p in name for p in first_rank_only_patterns): @@ -297,16 +313,18 @@ class Gemma4Attention(nn.Module): else -1 ) - self.total_num_heads = config.num_attention_heads - assert self.total_num_heads % tp_size == 0 - self.num_heads = self.total_num_heads // tp_size - if layer_type == "sliding_attention": + self.total_num_heads = getattr( + config, "swa_num_attention_heads", config.num_attention_heads + ) self.total_num_kv_heads = getattr( config, "swa_num_key_value_heads", config.num_key_value_heads ) else: + self.total_num_heads = config.num_attention_heads self.total_num_kv_heads = config.num_key_value_heads + assert self.total_num_heads % tp_size == 0 + self.num_heads = self.total_num_heads // tp_size self.num_kv_heads = max(1, self.total_num_kv_heads // tp_size) diff --git a/python/sglang/srt/models/gemma4_mm.py b/python/sglang/srt/models/gemma4_mm.py index 56d033ff4..afa9a0c09 100644 --- a/python/sglang/srt/models/gemma4_mm.py +++ b/python/sglang/srt/models/gemma4_mm.py @@ -59,13 +59,20 @@ from sglang.srt.model_loader.weight_utils import ( maybe_remap_kv_scale_name, ) from sglang.srt.models.gemma4_audio import Gemma4AudioEncoder -from sglang.srt.models.gemma4_causal import Gemma4TextModel, pp_filter_load_weight +from sglang.srt.models.gemma4_causal import ( + Gemma4TextModel, + load_tied_lm_head, + pp_filter_load_weight, +) from sglang.srt.models.gemma4_vision import Gemma4VisionEncoder -from sglang.srt.utils import add_prefix +from sglang.srt.utils import add_prefix, cpu_has_amx_support, is_cpu from sglang.srt.utils.hf_transformers_utils import get_processor logger = logging.getLogger(__name__) +_is_cpu_amx_available = cpu_has_amx_support() +_is_cpu = is_cpu() + cached_get_processor = lru_cache(get_processor) @@ -241,9 +248,15 @@ class Gemma4ForConditionalGeneration(PreTrainedModel): # while logits run on the last rank, so we can't reuse the embedding # module directly. For PP=1 keep the original tying; for PP>1 # materialize a real ParallelLMHead on the last rank and route the - # checkpoint embedding into it during load_weights. + # checkpoint embedding into it during load_weights. CPU with AMX does + # the same: the packed head weights cannot alias the embedding table. text_tie = getattr(text_config, "tie_word_embeddings", True) - if self.pp_group.world_size == 1 and text_tie: + self.lm_head_is_tied = ( + self.pp_group.world_size == 1 + and text_tie + and not (_is_cpu and _is_cpu_amx_available) + ) + if self.lm_head_is_tied: self.lm_head = self.language_model.embed_tokens elif self.pp_group.is_last_rank: self.lm_head = ParallelLMHead( @@ -668,13 +681,8 @@ class Gemma4ForConditionalGeneration(PreTrainedModel): if self.capture_aux_hidden_states: hidden_states, aux_hidden_states = hidden_states - # PP=1 keeps the original tied-weight behavior of using embed_tokens - # directly; under PP we route through the dedicated lm_head module. head = ( - self.language_model.embed_tokens - if self.pp_group.world_size == 1 - and getattr(self.config.text_config, "tie_word_embeddings", True) - else self.lm_head + self.language_model.embed_tokens if self.lm_head_is_tied else self.lm_head ) return self.logits_processor( input_ids, @@ -1016,6 +1024,16 @@ class Gemma4ForConditionalGeneration(PreTrainedModel): param, "weight_loader", default_weight_loader ) weight_loader(param, loaded_weight) + if ( + text_tie + and not self.lm_head_is_tied + and name == "language_model.embed_tokens.weight" + ): + load_tied_lm_head( + loaded_weight, + params_dict=params_dict, + loaded_params=loaded_params, + ) loaded_params.add(name) unloaded_params = params_dict.keys() - loaded_params if unloaded_params: @@ -1095,7 +1113,7 @@ class Gemma4ForConditionalGeneration(PreTrainedModel): "--pp-size 1 if you need this API." ) embed = self.language_model.embed_tokens.weight - # Gemma4 ties word embeddings, so embed_tokens serves as lm_head + # a materialized lm_head is loaded from this very tensor, so it is exact return embed, embed def set_eagle3_layers_to_capture(self, layer_ids: Optional[List[int]] = None): diff --git a/python/sglang/srt/models/gemma4_vision.py b/python/sglang/srt/models/gemma4_vision.py index 63fa2d064..bf8a990dd 100644 --- a/python/sglang/srt/models/gemma4_vision.py +++ b/python/sglang/srt/models/gemma4_vision.py @@ -30,7 +30,17 @@ from sglang.srt.layers.clippable_linear import ( from sglang.srt.layers.layernorm import Gemma4RMSNorm from sglang.srt.layers.quantization.base_config import QuantizationConfig from sglang.srt.runtime_context import get_mm, get_parallel -from sglang.srt.utils import add_prefix, get_device_capability, is_cuda, is_hip +from sglang.srt.utils import ( + add_prefix, + cpu_has_amx_support, + get_device_capability, + is_cpu, + is_cuda, + is_hip, +) + +_is_cpu = is_cpu() +_is_cpu_amx_available = _is_cpu and cpu_has_amx_support() # --------------------------------------------------------------------------- # 2-D Multidimensional RoPE (matches HF Gemma4RotaryEmbedding for vision) @@ -173,7 +183,9 @@ class Gemma4VisionAttention(nn.Module): num_heads=self.num_heads_per_partition, num_kv_heads=self.num_kv_heads_per_partition, dropout=0.0, - flatten_batch=True, + # sdpa asserts bsz == 1 under flatten_batch, which batched video + # frames violate; Gemma 4 passes its own 4-D mask regardless + flatten_batch=backend != "sdpa", softmax_in_single_precision=False, softmax_scale=1.0, ) @@ -198,6 +210,8 @@ class Gemma4VisionAttention(nn.Module): # ROCm: use triton_attn to avoid SDPA flatten_batch issues # with multi-image/video inputs return "triton_attn" + # not amx_attn: VisionAMXAttention swallows softmax_scale and the mask + # in **kwargs, and the CPU flash_attn hardcodes sm_scale return "sdpa" def forward( @@ -219,10 +233,15 @@ class Gemma4VisionAttention(nn.Module): k = self.k_norm(k.reshape(-1, self.head_dim)).reshape(k.shape) v = self.v_norm(v.reshape(-1, self.head_dim)).reshape(v.shape) - cos_flat = cos.reshape(bsz * seq_len, 1, self.head_dim) - sin_flat = sin.reshape(bsz * seq_len, 1, self.head_dim) - q = _apply_multidimensional_rope(q, cos_flat, sin_flat) - k = _apply_multidimensional_rope(k, cos_flat, sin_flat) + if _is_cpu_amx_available: + cos = cos.reshape(bsz * seq_len, self.head_dim) + sin = sin.reshape(bsz * seq_len, self.head_dim) + torch.ops.sgl_kernel.apply_multidimensional_rope_cpu(q, k, cos, sin) + else: + cos_flat = cos.reshape(bsz * seq_len, 1, self.head_dim) + sin_flat = sin.reshape(bsz * seq_len, 1, self.head_dim) + q = _apply_multidimensional_rope(q, cos_flat, sin_flat) + k = _apply_multidimensional_rope(k, cos_flat, sin_flat) if attention_mask is not None: attn_mask_4d = ( diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index ed46b25d2..f48d478d5 100644 --- a/python/sglang/srt/server_args.py +++ b/python/sglang/srt/server_args.py @@ -5635,13 +5635,19 @@ class ServerArgs: # Default attention backend selection moved to the override registry # (arg_groups/overrides.py: _gemma4_overrides). prefill_backend, decode_backend = self._resolved_attention_backends() - accepted_backends = ("trtllm_mha", "triton", "ascend", "intel_xpu") + accepted_backends = ( + "trtllm_mha", + "triton", + "ascend", + "intel_xpu", + "intel_amx", + ) assert ( prefill_backend in accepted_backends and decode_backend in accepted_backends ), ( - "Gemma4 only supports trtllm_mha, triton, or intel_xpu attention backend, " - f"got prefill={prefill_backend}, decode={decode_backend}" + "Gemma4 only supports trtllm_mha, triton, ascend, intel_xpu, or intel_amx " + f"attention backend, got prefill={prefill_backend}, decode={decode_backend}" ) # The quantization/moe_runner_backend resolution moved to the override diff --git a/python/sglang/test/cpu_test_utils.py b/python/sglang/test/cpu_test_utils.py index 83dba491d..f75279687 100644 --- a/python/sglang/test/cpu_test_utils.py +++ b/python/sglang/test/cpu_test_utils.py @@ -1,3 +1,4 @@ +import functools import itertools import math @@ -198,7 +199,17 @@ def scaled_weight(weight, scales): return weight_scaled -def torch_naive_fused_moe(a, w1, w2, score, topk, renormalize): +def _activation_fn(activation): + """Reference gate-and-multiply for the activation fused_experts_cpu applies.""" + if activation == "silu": + return SiluAndMul + if activation == "gelu": + # matches the erf gelu the kernel computes + return functools.partial(GeluAndMul, approximate="none") + raise ValueError(f"Unsupported activation: {activation}") + + +def torch_naive_fused_moe(a, w1, w2, score, topk, renormalize, activation="silu"): B, D = a.shape a = a.view(B, -1, D).repeat(1, topk, 1).reshape(-1, D) out = torch.zeros(B * topk, w2.shape[1], dtype=a.dtype, device=a.device) @@ -208,14 +219,14 @@ def torch_naive_fused_moe(a, w1, w2, score, topk, renormalize): if renormalize: topk_weight = topk_weight / topk_weight.sum(dim=-1, keepdim=True) + act_fn = _activation_fn(activation) + topk_weight = topk_weight.view(-1) topk_ids = topk_ids.view(-1) for i in range(w1.shape[0]): mask = topk_ids == i if mask.sum(): - out[mask] = SiluAndMul(a[mask] @ w1[i].transpose(0, 1)) @ w2[i].transpose( - 0, 1 - ) + out[mask] = act_fn(a[mask] @ w1[i].transpose(0, 1)) @ w2[i].transpose(0, 1) return ( out.view(B, -1, w2.shape[1]) * topk_weight.view(B, -1, 1).to(out.dtype) ).sum(dim=1) @@ -372,11 +383,13 @@ def torch_w8a8_per_column_fused_moe(a, w1, w2, w1_s, w2_s, topk_weight, topk_ids ) -def native_fp8_fused_moe(a, w1, w2, topk_weight, topk_ids, topk): +def native_fp8_fused_moe(a, w1, w2, topk_weight, topk_ids, topk, activation="silu"): B, D = a.shape a = a.view(B, -1, D).repeat(1, topk, 1).reshape(-1, D).float() out = torch.zeros(B * topk, w2.shape[1], dtype=torch.float32, device=a.device) + act_fn = _activation_fn(activation) + # Calculate routing topk_weight = topk_weight.view(-1) topk_ids = topk_ids.view(-1) @@ -385,7 +398,7 @@ def native_fp8_fused_moe(a, w1, w2, topk_weight, topk_ids, topk): mask = topk_ids == i if mask.sum(): ic0 = torch.matmul(a[mask], w1[i].transpose(0, 1)) - ic1 = SiluAndMul(ic0) + ic1 = act_fn(ic0) out[mask] = torch.matmul(ic1, w2[i].transpose(0, 1)) return ( diff --git a/test/registered/cpu/test_extend.py b/test/registered/cpu/test_extend.py index 994ead071..4400d45ed 100644 --- a/test/registered/cpu/test_extend.py +++ b/test/registered/cpu/test_extend.py @@ -196,6 +196,7 @@ class TestExtendAttention(CustomTestCase): *, b_seq_len_prefix=None, b_seq_len_extend=None, + kv_from_cache=False, ): dtype = torch.bfloat16 @@ -322,8 +323,8 @@ class TestExtendAttention(CustomTestCase): o_extend = torch.empty((extend_token_num, H_Q, DV), dtype=dtype) torch.ops.sgl_kernel.extend_attention_cpu( q_extend, - k_extend, - v_extend, + None if kv_from_cache else k_extend, + None if kv_from_cache else v_extend, o_extend, k_buffer, v_buffer, @@ -374,6 +375,27 @@ class TestExtendAttention(CustomTestCase): 1, 20, 1, 1, 64, 64, sliding_window, has_sink, False, False ) + def test_extend_attention_kv_from_cache(self): + # KV-shared layers pass no extend K/V, so the kernel masks the extend + # range causally itself; sizes straddle several BLOCK_N. + # Window only tested with a sink - _run_sdpa_forward_extend models none, + # the same restriction test_extend_attention applies. + for sliding_window, has_sink in [(None, False), (128, True)]: + for prefix, extend in [([0], [343]), ([100], [343]), ([0], [1500])]: + self._test_extend_attention_once( + B=1, + N_CTX=4096, + H_Q=16, + H_KV=4, + D=64, + DV=64, + sliding_window=sliding_window, + has_sink=has_sink, + b_seq_len_prefix=prefix, + b_seq_len_extend=extend, + kv_from_cache=True, + ) + def test_extend_attention_large_seq_causal_mask(self): self._test_extend_attention_once( B=1, diff --git a/test/registered/cpu/test_moe.py b/test/registered/cpu/test_moe.py index 8f540ee6e..d6a3d1d26 100644 --- a/test/registered/cpu/test_moe.py +++ b/test/registered/cpu/test_moe.py @@ -54,6 +54,7 @@ def run_fused_experts( alpha=None, limit=None, is_vnni=True, + activation=None, inplace=False, ): return kernel.fused_experts_cpu( @@ -74,6 +75,7 @@ def run_fused_experts( alpha, limit, is_vnni, + activation, ) @@ -138,13 +140,35 @@ def make_mxfp4_weights(e, out_dim, in_dim, dtype, with_bias=False): class TestFusedExperts: + def test_unsupported_activation_is_rejected(self): + m, n, k, e, topk = 2, 32, 32, 4, 2 + a = torch.randn((m, k), dtype=dtype) / 10 + w1 = make_bf16_weights(e, 2 * n, k) + w2 = make_bf16_weights(e, k, n) + topk_weights, topk_ids = make_routing(m, e, topk, dtype=dtype) + packed_w1 = kernel.convert_weight_packed(w1) if prepack else w1 + packed_w2 = kernel.convert_weight_packed(w2) if prepack else w2 + + with pytest.raises(RuntimeError, match="Unsupported activation"): + run_fused_experts( + a, + packed_w1, + packed_w2, + topk_weights, + topk_ids, + quant=CPUQuantMethod.UNQUANT, + is_vnni=prepack, + activation="relu", + ) + @pytest.mark.parametrize("m", [2, 114]) @pytest.mark.parametrize("n", [32]) @pytest.mark.parametrize("k", [32]) @pytest.mark.parametrize("e", [4]) @pytest.mark.parametrize("topk", [2]) @pytest.mark.parametrize("renormalize", [False, True]) - def test_bf16_moe(self, m, n, k, e, topk, renormalize): + @pytest.mark.parametrize("activation", ["silu", "gelu"]) + def test_bf16_moe(self, m, n, k, e, topk, renormalize, activation): a = torch.randn((m, k), dtype=dtype) / 10 w1 = make_bf16_weights(e, 2 * n, k) w2 = make_bf16_weights(e, k, n) @@ -156,7 +180,9 @@ class TestFusedExperts: renormalize=renormalize, return_score=True, ) - torch_output = torch_naive_fused_moe(a, w1, w2, score, topk, renormalize) + torch_output = torch_naive_fused_moe( + a, w1, w2, score, topk, renormalize, activation=activation + ) packed_w1 = kernel.convert_weight_packed(w1) if prepack else w1 packed_w2 = kernel.convert_weight_packed(w2) if prepack else w2 @@ -168,6 +194,7 @@ class TestFusedExperts: topk_ids, quant=CPUQuantMethod.UNQUANT, is_vnni=prepack, + activation=activation, inplace=True, ) @@ -276,7 +303,8 @@ class TestFusedExperts: @pytest.mark.parametrize("K", [256, 320]) @pytest.mark.parametrize("E", [8]) @pytest.mark.parametrize("topk", [4]) - def test_fp8_moe(self, M, N, K, E, topk): + @pytest.mark.parametrize("activation", ["silu", "gelu"]) + def test_fp8_moe(self, M, N, K, E, topk, activation): a = torch.randn(M, K, dtype=dtype) / math.sqrt(K) w1, w1s, w1_scaled = make_fp8_weights(E, 2 * N, K) @@ -288,7 +316,7 @@ class TestFusedExperts: w2 = kernel.convert_weight_packed(w2) ref_out = native_fp8_fused_moe( - a, w1_scaled, w2_scaled, topk_weight, topk_ids, topk + a, w1_scaled, w2_scaled, topk_weight, topk_ids, topk, activation=activation ) out = run_fused_experts( a, @@ -301,6 +329,7 @@ class TestFusedExperts: w2_scale=w2s, block_size=[BLOCK_N, BLOCK_K], is_vnni=True, + activation=activation, inplace=False, ) @@ -372,7 +401,8 @@ class TestFusedExperts: @pytest.mark.parametrize("K", [256, 320]) @pytest.mark.parametrize("E", [8]) @pytest.mark.parametrize("topk", [4]) - def test_mxfp4_moe(self, M, N, K, E, topk): + @pytest.mark.parametrize("activation", ["silu", "gelu"]) + def test_mxfp4_moe(self, M, N, K, E, topk, activation): a = torch.randn(M, K, dtype=dtype) / 10 w1dq, w1_packed, w1s_packed = make_mxfp4_weights(E, 2 * N, K, dtype=dtype) @@ -381,7 +411,13 @@ class TestFusedExperts: topk_weight, topk_ids = make_routing(M, E, topk, dtype=dtype) ref_out = native_fp8_fused_moe( - a, w1dq.float(), w2dq.float(), topk_weight, topk_ids, topk + a, + w1dq.float(), + w2dq.float(), + topk_weight, + topk_ids, + topk, + activation=activation, ) out = run_fused_experts( a, @@ -393,6 +429,7 @@ class TestFusedExperts: w1_scale=w1s_packed, w2_scale=w2s_packed, is_vnni=True, + activation=activation, inplace=False, ) diff --git a/test/registered/cpu/test_rope.py b/test/registered/cpu/test_rope.py index d790e1f83..4a04011dc 100644 --- a/test/registered/cpu/test_rope.py +++ b/test/registered/cpu/test_rope.py @@ -70,9 +70,9 @@ class TestROPE(CustomTestCase): with torch.no_grad(), torch.amp.autocast("cpu", enabled=enable_autocast): q = torch.randn(seq_len, num_heads * head_size, dtype=dtype) - q_clone = q.clone() + q_sgl = q.clone() k = torch.randn(seq_len, num_kv_heads * head_size, dtype=dtype) - k_clone = k.clone() + k_sgl = k.clone() # ref kernel q_ref, k_ref = rope.forward_native( @@ -81,10 +81,10 @@ class TestROPE(CustomTestCase): positions=positions, ) # fused rope kernel - q_sgl, k_sgl = torch.ops.sgl_kernel.multimodal_rotary_embedding_cpu( + torch.ops.sgl_kernel.multimodal_rotary_embedding_cpu( positions, - q_clone, - k_clone, + q_sgl, + k_sgl, rope.head_size, rope.cos_sin_cache, rope.mrope_section, @@ -286,6 +286,83 @@ class TestROPE(CustomTestCase): torch.testing.assert_close(q_out_ref, q_out_sgl, atol=1e-2, rtol=1e-2) torch.testing.assert_close(k_out_ref, k_out_sgl, atol=1e-2, rtol=1e-2) + def test_apply_multidimensional_rope(self): + """Test apply_multidimensional_rope_cpu against the native Python reference.""" + + def _rotate_half(x): + x1 = x[..., : x.shape[-1] // 2] + x2 = x[..., x.shape[-1] // 2 :] + return torch.cat((-x2, x1), dim=-1) + + def _apply_rotary(x, cos, sin): + return (x * cos) + (_rotate_half(x) * sin) + + def _apply_multidimensional_rope_ref(x, cos, sin): + ndim = 2 + chunk_size = x.shape[-1] // ndim + cos_3d = cos.unsqueeze(1) + sin_3d = sin.unsqueeze(1) + x_parts = x.split(chunk_size, dim=-1) + cos_parts = cos_3d.split(chunk_size, dim=-1) + sin_parts = sin_3d.split(chunk_size, dim=-1) + y_parts = [ + _apply_rotary(x_parts[k], cos_parts[k], sin_parts[k]) + for k in range(ndim) + ] + return torch.cat(y_parts, dim=-1) + + test_configs = [ + # (num_tokens, num_heads, head_dim, dtype, sincos_dtype) + (4, 8, 64, torch.bfloat16, torch.bfloat16), + (32, 16, 128, torch.bfloat16, torch.bfloat16), + (128, 4, 256, torch.bfloat16, torch.bfloat16), + (1, 1, 32, torch.bfloat16, torch.float32), + (32, 16, 128, torch.bfloat16, torch.float32), + (2520, 12, 64, torch.bfloat16, torch.bfloat16), + (2520, 12, 64, torch.bfloat16, torch.float32), + # head_dim 160 -> 40 elements per rotary half, so the 32-wide + # vector loop runs once and leaves an 8-element scalar tail + (17, 3, 160, torch.bfloat16, torch.bfloat16), + (17, 3, 160, torch.float16, torch.float32), + (32, 16, 128, torch.float16, torch.float16), + ] + + for num_tokens, num_heads, head_dim, dtype, sincos_dtype in test_configs: + with self.subTest( + num_tokens=num_tokens, + num_heads=num_heads, + head_dim=head_dim, + dtype=dtype, + sincos_dtype=sincos_dtype, + ): + torch.manual_seed(42) + query = torch.randn( + num_tokens, num_heads, head_dim, dtype=dtype, device="cpu" + ) + key = torch.randn( + num_tokens, num_heads, head_dim, dtype=dtype, device="cpu" + ) + cos = torch.randn( + num_tokens, head_dim, dtype=sincos_dtype, device="cpu" + ) + sin = torch.randn( + num_tokens, head_dim, dtype=sincos_dtype, device="cpu" + ) + + q_expected = _apply_multidimensional_rope_ref( + query.float(), cos.float(), sin.float() + ).to(dtype) + k_expected = _apply_multidimensional_rope_ref( + key.float(), cos.float(), sin.float() + ).to(dtype) + + torch.ops.sgl_kernel.apply_multidimensional_rope_cpu( + query, key, cos, sin + ) + atol = rtol = precision[dtype] + torch.testing.assert_close(query, q_expected, atol=atol, rtol=rtol) + torch.testing.assert_close(key, k_expected, atol=atol, rtol=rtol) + if __name__ == "__main__": unittest.main()