[CPU] Add support for Gemma4 on Xeon (#22498)

Co-authored-by: Copilot <198982749+Copilot@users.noreply.github.com>
Co-authored-by: jianan-gu <jianan.gu@intel.com>
Co-authored-by: Haotong Zou <haotong.zou@intel.com>

python/sglang/kernels/aot/csrc/cpu/aarch64/moe.cpp

@@ -242,7 +242,8 @@ at::Tensor fused_experts_cpu(
     const std::optional<at::Tensor>& /*w2_bias*/,
     const std::optional<double>& /*alpha*/,
     const std::optional<double>& /*limit*/,
-    bool /*is_vnni*/) {
+    bool /*is_vnni*/,
+    const std::optional<std::string>& 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);
 

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;

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<scalar_t>::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,13 +182,32 @@ 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<float>::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<float>::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<scalar_t, BLOCK_M, BLOCK_N>::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<at::Tensor> encoder_lens,
     std::optional<at::Tensor> sinks,
     std::optional<at::Tensor> tree_mask) {
-  if (!is_cross_attn) {
+  TORCH_CHECK(k_extend_opt.has_value() == v_extend_opt.has_value(), "k_extend and v_extend must be given together");
-    TORCH_CHECK(
+  // A KV-shared layer (Gemma 4) passes no extend K/V - the layer it shares with
-        k_extend_opt.has_value() && v_extend_opt.has_value(),
+  // already wrote them to the cache, so this kernel masks causally itself. Cross
-        "k_extend and v_extend are required for non-cross attention");
+  // attention can also arrive without K/V but reads an encoder sequence that
-  }
+  // carries no causal order, so it keeps its own path.
-  // Since k_extend and v_extend are not used for cross attention, they can be initialized as k_buffer and v_buffer
+  const bool kv_from_cache = !is_cross_attn && !k_extend_opt.has_value();
-  // here.
+  // 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.

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<std::string>& 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) {

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<at::Tensor>& w2_bias,
     const std::optional<double>& alpha,
     const std::optional<double>& limit,
-    bool is_vnni) {
+    bool is_vnni,
+    const std::optional<std::string>& 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<scalar_t, at::Float8_e4m3fn, float, false>(
@@ -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<scalar_t>(
           out_hidden_states.data_ptr<scalar_t>(),

python/sglang/kernels/aot/csrc/cpu/moe.h

@@ -159,6 +159,26 @@ inline void silu_and_mul_stub(
   }
 }
 
+template <typename scalar_t, typename input_t>
+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<input_t, float> || std::is_same_v<input_t, scalar_t>,
+      "gelu_and_mul_stub only supports input_t == float or input_t == scalar_t");
+  using bVec = at::vec::Vectorized<scalar_t>;
+
+  // 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<scalar_t>(x0, x1);
+    out_vec.store(out + d);
+  }
+}
+
 template <typename scalar_t, typename input_t>
 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) {

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) {

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<true>(key, {num_tokens, num_heads, head_size}, input_dtype);
+  CHECK_INPUT_SHAPE_DTYPE<false>(cos, {num_tokens, head_size}, cos.scalar_type());
+  CHECK_INPUT_SHAPE_DTYPE<false>(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>();
+    scalar_t* k_ptr = key.data_ptr<scalar_t>();
+    const param_t* cos_ptr = cos.data_ptr<param_t>();
+    const param_t* sin_ptr = sin.data_ptr<param_t>();
+    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<param_t> {
+        return {cos_ptr + token * head_size + offset, sin_ptr + token * head_size + offset};
+      };
+      rotary_embedding_kernel_impl<scalar_t, RotaryMode::NeoxFull, true>(
+          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<at::Tensor, at::Tensor> multimodal_rotary_embedding_cpu(
+void multimodal_rotary_embedding_cpu(
     at::Tensor& positions,
     at::Tensor& query,
     at::Tensor& key,
@@ -449,5 +490,4 @@ std::tuple<at::Tensor, at::Tensor> multimodal_rotary_embedding_cpu(
       }
     }
   });
-  return std::make_tuple(query, key);
 }

python/sglang/kernels/aot/csrc/cpu/torch_extension_cpu.cpp

@@ -361,7 +361,8 @@ at::Tensor fused_experts_cpu(
     const std::optional<at::Tensor>& w2_bias,
     const std::optional<double>& alpha,
     const std::optional<double>& limit,
-    bool is_vnni);
+    bool is_vnni,
+    const std::optional<std::string>& activation);
 
 #if !defined(SGLANG_CPU_ARM64_SKIP_X86_ONLY_OPS)
 at::Tensor shared_expert_cpu(
@@ -478,8 +479,11 @@ std::tuple<at::Tensor, at::Tensor> rotary_embedding_cpu(
 std::tuple<at::Tensor, at::Tensor>
 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<at::Tensor, at::Tensor> 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 "
+      "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) -> (Tensor, 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

python/sglang/kernels/aot/csrc/cpu/vec.h

@@ -612,6 +612,14 @@ inline at::vec::Vectorized<float> fast_silu(const at::vec::Vectorized<float>& x)
 #endif
 }
 
+// exact (erf) gelu, matching torch.nn.functional.gelu(approximate="none")
+inline at::vec::Vectorized<float> fast_gelu(const at::vec::Vectorized<float>& x) {
+  const auto half = at::vec::Vectorized<float>(0.5f);
+  const auto one = at::vec::Vectorized<float>(1.f);
+  const auto inv_sqrt2 = at::vec::Vectorized<float>(0.70710678118654752440f);
+  return half * x * (one + (x * inv_sqrt2).erf());
+}
+
 inline at::vec::Vectorized<float>
 fast_sigmoid_glu(const at::vec::Vectorized<float>& x, const at::vec::Vectorized<float>& alpha) {
 #if defined(CPU_CAPABILITY_AVX512)

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 (
+        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)
 

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

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
             )

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)
 

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

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:

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)
 

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(

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:
+        # TODO: the kernel aliases query/key for 2D and 4D but allocates for 3D,
-            return query, key
+        # which no schema expresses; an accurate fake needs it to pick one
-        else:
         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,

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)
+        load_tied_lm_head(
-        if head_param is not None:
+            loaded_weight,
-            wl = getattr(head_param, "weight_loader", default_weight_loader)
+            params_dict=params_dict,
-            wl(head_param, loaded_weight)
+            loaded_params=loaded_params,
-            loaded_params.add(head_param_name)
+            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)
 

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
+            self.language_model.embed_tokens if self.lm_head_is_tied else self.lm_head
-            if self.pp_group.world_size == 1
-            and getattr(self.config.text_config, "tie_word_embeddings", True)
-            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):

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,6 +233,11 @@ 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)
 
+        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)

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, "
+                "Gemma4 only supports trtllm_mha, triton, ascend, intel_xpu, or intel_amx "
-                f"got prefill={prefill_backend}, decode={decode_backend}"
+                f"attention backend, got prefill={prefill_backend}, decode={decode_backend}"
             )
 
             # The quantization/moe_runner_backend resolution moved to the override

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(
+            out[mask] = act_fn(a[mask] @ w1[i].transpose(0, 1)) @ w2[i].transpose(0, 1)
-                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 (

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,
+            None if kv_from_cache else k_extend,
-            v_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,

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,
         )
 

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,
+                    q_sgl,
-                    k_clone,
+                    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()