diff --git a/python/sglang/srt/layers/attention/intel_amx_backend.py b/python/sglang/srt/layers/attention/intel_amx_backend.py
index e64bb45c0..46b657d64 100644
--- a/python/sglang/srt/layers/attention/intel_amx_backend.py
+++ b/python/sglang/srt/layers/attention/intel_amx_backend.py
@@ -99,14 +99,15 @@ class IntelAMXAttnBackend(AttentionBackend):
             o = q.new_empty((q.shape[0], layer.tp_q_head_num * layer.v_head_dim))
         else:
             o = torch.empty_like(q)
-
-        if save_kv_cache:
-            forward_batch.token_to_kv_pool.set_kv_buffer(
-                layer, forward_batch.out_cache_loc, k, v
-            )
+        cache_loc = (
+            forward_batch.out_cache_loc
+            if not layer.is_cross_attention
+            else forward_batch.encoder_out_cache_loc
+        )
+        if save_kv_cache and k is not None and v is not None:
+            forward_batch.token_to_kv_pool.set_kv_buffer(layer, cache_loc, k, v)
 
         _, max_extend_len = self.forward_metadata
-
         self.extend_attention_fwd(
             q.view(-1, layer.tp_q_head_num, layer.qk_head_dim),
             k,
@@ -122,6 +123,8 @@ class IntelAMXAttnBackend(AttentionBackend):
             max_extend_len,
             layer.scaling,
             layer.logit_cap,
+            layer.is_cross_attention,
+            forward_batch.encoder_lens,
         )
         return o
 
@@ -142,7 +145,11 @@ class IntelAMXAttnBackend(AttentionBackend):
             o = q.new_empty((q.shape[0], layer.tp_q_head_num * layer.v_head_dim))
         else:
             o = torch.empty_like(q)
-
+        cache_loc = (
+            forward_batch.out_cache_loc
+            if not layer.is_cross_attention
+            else forward_batch.encoder_out_cache_loc
+        )
         self.decode_attention_fwd(
             q.view(-1, layer.tp_q_head_num, layer.qk_head_dim),
             forward_batch.token_to_kv_pool.get_key_buffer(layer.layer_id),
@@ -150,15 +157,16 @@ class IntelAMXAttnBackend(AttentionBackend):
             o.view(-1, layer.tp_q_head_num, layer.v_head_dim),
             k,
             v,
-            forward_batch.out_cache_loc,
+            cache_loc,
             attn_logits,
             forward_batch.req_to_token_pool.req_to_token,
             forward_batch.req_pool_indices,
             forward_batch.seq_lens,
             layer.scaling,
             layer.logit_cap,
+            layer.is_cross_attention,
+            forward_batch.encoder_lens,
         )
-
         return o
 
     def support_triton(self):
diff --git a/python/sglang/srt/layers/attention/torch_native_backend.py b/python/sglang/srt/layers/attention/torch_native_backend.py
index 85f0cb42f..00d424c44 100644
--- a/python/sglang/srt/layers/attention/torch_native_backend.py
+++ b/python/sglang/srt/layers/attention/torch_native_backend.py
@@ -1,6 +1,6 @@
 from __future__ import annotations
 
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, Optional
 
 import torch
 from torch.nn.functional import scaled_dot_product_attention
@@ -35,9 +35,11 @@ class TorchNativeAttnBackend(AttentionBackend):
         seq_lens: torch.Tensor,
         extend_prefix_lens: torch.Tensor,
         extend_seq_lens: torch.Tensor,
+        encoder_lens: Optional[torch.Tensor] = None,
         scaling=None,
         enable_gqa=False,
         causal=False,
+        is_cross_attn=False,
     ):
         """Run the extend forward by using torch native sdpa op.
 
@@ -48,12 +50,14 @@ class TorchNativeAttnBackend(AttentionBackend):
             v_cache: [max_total_num_tokens, num_heads, head_size]
             req_to_token: [max_num_reqs, max_context_len]
             req_pool_indices: [num_seqs]
+            encoder_lens: [num_seqs] or None
             seq_lens: [num_seqs]
             extend_prefix_lens: [num_seqs]
             extend_seq_lens: [num_seqs]
             scaling: float or None
             enable_gqa: bool
             causal: bool
+            is_cross_attn: bool
 
         Returns:
             output: [num_tokens, num_heads, head_size]
@@ -75,8 +79,16 @@ class TorchNativeAttnBackend(AttentionBackend):
 
             seq_len_kv = seq_lens[seq_idx]
             end_q = start_q + extend_seq_len_q
-            end_kv = start_kv + seq_len_kv
-
+            if encoder_lens is not None:
+                if is_cross_attn:
+                    start_kv = 0
+                    end_kv = encoder_lens[seq_idx]
+                else:
+                    start_kv = encoder_lens[seq_idx]
+                    end_kv = start_kv + seq_len_kv
+            else:
+                start_kv = 0
+                end_kv = start_kv + seq_len_kv
             per_req_query = query[:, start_q:end_q, :]
             per_req_query_redudant = torch.empty(
                 (per_req_query.shape[0], seq_len_kv, per_req_query.shape[2]),
@@ -89,7 +101,7 @@ class TorchNativeAttnBackend(AttentionBackend):
             # get key and value from cache. per_req_tokens contains the kv cache
             # index for each token in the sequence.
             req_pool_idx = req_pool_indices[seq_idx]
-            per_req_tokens = req_to_token[req_pool_idx, :seq_len_kv]
+            per_req_tokens = req_to_token[req_pool_idx, start_kv:end_kv]
             per_req_key = k_cache[per_req_tokens].movedim(0, query.dim() - 2)
             per_req_value = v_cache[per_req_tokens].movedim(0, query.dim() - 2)
 
@@ -123,9 +135,11 @@ class TorchNativeAttnBackend(AttentionBackend):
         req_to_token: torch.Tensor,
         req_pool_indices: torch.Tensor,
         seq_lens: torch.Tensor,
+        encoder_lens: Optional[torch.Tensor] = None,
         scaling=None,
         enable_gqa=False,
         causal=False,
+        is_cross_attn=False,
     ):
         """Run the decode forward by using torch native sdpa op.
 
@@ -137,9 +151,11 @@ class TorchNativeAttnBackend(AttentionBackend):
             req_to_token: [max_num_reqs, max_context_len]
             req_pool_indices: [num_seqs]
             seq_lens: [num_seqs]
+            encoder_lens: [num_seqs] or None
             scaling: float or None
             enable_gqa: bool
             causal: bool
+            is_cross_attn: bool
 
         Returns:
             output: [num_tokens, num_heads, head_size]
@@ -156,14 +172,24 @@ class TorchNativeAttnBackend(AttentionBackend):
             seq_len_q = 1
             seq_len_kv = seq_lens[seq_idx]
             end_q = start_q + seq_len_q
-            end_kv = start_kv + seq_len_kv
+            if encoder_lens is not None:
+                if is_cross_attn:
+                    start_kv = 0
+                    end_kv = encoder_lens[seq_idx]
+                else:
+                    start_kv = encoder_lens[seq_idx]
+                    end_kv = start_kv + seq_len_kv
+            else:
+                start_kv = 0
+                end_kv = start_kv + seq_len_kv
 
             per_req_query = query[:, start_q:end_q, :]
 
             # get key and value from cache. per_req_tokens contains the kv cache
             # index for each token in the sequence.
+
             req_pool_idx = req_pool_indices[seq_idx]
-            per_req_tokens = req_to_token[req_pool_idx, :seq_len_kv]
+            per_req_tokens = req_to_token[req_pool_idx, start_kv:end_kv]
             per_req_key = k_cache[per_req_tokens].movedim(0, query.dim() - 2)
             per_req_value = v_cache[per_req_tokens].movedim(0, query.dim() - 2)
 
@@ -208,7 +234,7 @@ class TorchNativeAttnBackend(AttentionBackend):
         else:
             cache_loc = forward_batch.out_cache_loc
 
-        if save_kv_cache:
+        if save_kv_cache and k is not None and v is not None:
             forward_batch.token_to_kv_pool.set_kv_buffer(layer, cache_loc, k, v)
 
         use_gqa = layer.tp_q_head_num != layer.tp_k_head_num
@@ -230,9 +256,11 @@ class TorchNativeAttnBackend(AttentionBackend):
             forward_batch.seq_lens,
             forward_batch.extend_prefix_lens,
             forward_batch.extend_seq_lens,
+            forward_batch.encoder_lens,
             scaling=layer.scaling,
             enable_gqa=use_gqa,
             causal=causal,
+            is_cross_attn=layer.is_cross_attention,
         )
         return o
 
@@ -253,6 +281,11 @@ class TorchNativeAttnBackend(AttentionBackend):
             o = q.new_empty((q.shape[0], layer.tp_q_head_num * layer.v_head_dim))
         else:
             o = torch.empty_like(q)
+        cache_loc = (
+            forward_batch.out_cache_loc
+            if not layer.is_cross_attention
+            else forward_batch.encoder_out_cache_loc
+        )
 
         if layer.is_cross_attention:
             cache_loc = forward_batch.encoder_out_cache_loc
@@ -260,7 +293,8 @@ class TorchNativeAttnBackend(AttentionBackend):
             cache_loc = forward_batch.out_cache_loc
 
         if save_kv_cache:
-            forward_batch.token_to_kv_pool.set_kv_buffer(layer, cache_loc, k, v)
+            if k is not None and v is not None:
+                forward_batch.token_to_kv_pool.set_kv_buffer(layer, cache_loc, k, v)
 
         use_gqa = layer.tp_q_head_num != layer.tp_k_head_num
 
@@ -275,9 +309,11 @@ class TorchNativeAttnBackend(AttentionBackend):
             forward_batch.req_to_token_pool.req_to_token,
             forward_batch.req_pool_indices,
             forward_batch.seq_lens,
+            forward_batch.encoder_lens,
             scaling=layer.scaling,
             enable_gqa=use_gqa,
             causal=False,
+            is_cross_attn=layer.is_cross_attention,
         )
 
         return o
diff --git a/python/sglang/srt/models/mllama.py b/python/sglang/srt/models/mllama.py
index c7ebba787..8f05d9432 100644
--- a/python/sglang/srt/models/mllama.py
+++ b/python/sglang/srt/models/mllama.py
@@ -565,6 +565,7 @@ class MllamaTextCrossAttention(nn.Module):
         )
 
         output = self.attn(q, k, v, forward_batch)
+        output = output.view(-1, self.num_local_heads * self.head_dim)
         out, _ = self.o_proj(output)
         return out
 
@@ -865,9 +866,7 @@ class MllamaForConditionalGeneration(nn.Module):
                 self.image_size,
                 dtype=torch.float32,
             )
-            batched_ar_ids = torch.ones(
-                bs, max_num_images, dtype=torch.int64, device="cuda"
-            )
+            batched_ar_ids = torch.ones(bs, max_num_images, dtype=torch.int64)
             batched_ar_mask = torch.zeros(
                 bs, max_num_images, max_num_tiles, dtype=torch.int64
             )
@@ -886,11 +885,13 @@ class MllamaForConditionalGeneration(nn.Module):
                     img = pixel_values[0, j]
                     num_tiles = img.shape[0]
                     batched_images[i, j, :num_tiles] = img
-                    batched_ar_ids[i, j] = mm_input.mm_items[0].aspect_ratio_ids[0, j]
+                    batched_ar_ids[i, j] = mm_input.mm_items[0].model_specific_data[
+                        "aspect_ratio_ids"
+                    ][0, j]
 
                     batched_ar_mask[i, j, :num_tiles] = mm_input.mm_items[
                         0
-                    ].aspect_ratio_mask[0, j]
+                    ].model_specific_data["aspect_ratio_mask"][0, j]
                 i += 1
 
         return batched_images, batched_ar_ids, batched_ar_mask, encoder_lens_need
diff --git a/sgl-kernel/csrc/cpu/decode.cpp b/sgl-kernel/csrc/cpu/decode.cpp
index 21334238e..b3ea06f10 100644
--- a/sgl-kernel/csrc/cpu/decode.cpp
+++ b/sgl-kernel/csrc/cpu/decode.cpp
@@ -1038,6 +1038,7 @@ void decode_attention_kernel_impl(
     const index_t* __restrict__ req_to_token,
     const int64_t* __restrict__ req_pool_indices,
     const int64_t* __restrict__ seq_lens,
+    const int64_t* __restrict__ encoder_lens,
     int64_t batches,
     int64_t num_heads,
     int64_t head_size,
@@ -1053,7 +1054,9 @@ void decode_attention_kernel_impl(
     float logit_cap,
     int64_t max_num_reqs,
     int64_t max_context_len,
-    int64_t max_total_num_tokens) {
+    int64_t max_total_num_tokens,
+    bool is_cross_attn,
+    bool has_encoder_lens) {
   using Vec = at::vec::Vectorized<float>;
 
   // strides
@@ -1077,8 +1080,9 @@ void decode_attention_kernel_impl(
       const scalar_t* __restrict__ q_ptr = query + bs * q_strideM + head_id * q_strideH;
 
       // get key/value
-      int64_t seq_len_kv = seq_lens[bs];
+      int64_t seq_len_kv = is_cross_attn ? encoder_lens[bs] : seq_lens[bs];
       int64_t req_pool_id = req_pool_indices[bs];
+      int64_t kv_offset = (has_encoder_lens && (!is_cross_attn)) ? encoder_lens[bs] : 0;
       TORCH_CHECK(seq_len_kv <= max_context_len, "seq_len_kv out of scope!");
       TORCH_CHECK(req_pool_id < max_num_reqs, "req_pool_id out of scope!");
 
@@ -1102,7 +1106,7 @@ void decode_attention_kernel_impl(
             /* A   */ q_ptr,
             /* B   */ k_buffer + head_id * k_strideH,
             /* C   */ s_i,
-            /* ind */ req_to_token + req_pool_id * max_context_len + n,
+            /* ind */ req_to_token + req_pool_id * max_context_len + n + kv_offset,
             /* scl */ sm_scale,
             /* M   */ 1,
             /* N   */ n_size,
@@ -1142,7 +1146,7 @@ void decode_attention_kernel_impl(
             /* A   */ s_delta,
             /* B   */ v_buffer + head_id * v_strideH,
             /* C   */ v_prime,
-            /* ind */ req_to_token + req_pool_id * max_context_len + n,
+            /* ind */ req_to_token + req_pool_id * max_context_len + n + kv_offset,
             /* scl */ &m_delta,
             /* M   */ 1,
             /* N   */ head_size_v,
@@ -1159,6 +1163,8 @@ void decode_attention_kernel_impl(
         at::vec::map<float>([s](Vec out) { return out * Vec(s); }, v_prime, v_prime, head_size_v);
 
         v_prime[head_size_v] = m_prime + std::log(s_prime);
+      } else {
+        v_prime[head_size_v] = -std::numeric_limits<float>::infinity();
       }
 
       // move to the next index
@@ -1350,6 +1356,10 @@ void decode_attention_mla_kernel_impl(
               [s](Vec out) { return out * Vec(s); }, v_prime + h * l_stride1, v_prime + h * l_stride1, head_size_v);
           (v_prime + h * l_stride1)[head_size_v] = m_prime[h] + std::log(s_prime[h]);
         }
+      } else {
+        for (int64_t h = 0; h < h_size; ++h) {
+          (v_prime + h * l_stride1)[head_size_v] = -std::numeric_limits<float>::infinity();
+        }
       }
 
       // move to the next index
@@ -1372,6 +1382,7 @@ void decode_attention_grouped_kernel_impl(
     const index_t* __restrict__ req_to_token,
     const int64_t* __restrict__ req_pool_indices,
     const int64_t* __restrict__ seq_lens,
+    const int64_t* __restrict__ encoder_lens,
     int64_t batches,
     int64_t num_heads,
     int64_t num_heads_kv,
@@ -1388,7 +1399,9 @@ void decode_attention_grouped_kernel_impl(
     float logit_cap,
     int64_t max_num_reqs,
     int64_t max_context_len,
-    int64_t max_total_num_tokens) {
+    int64_t max_total_num_tokens,
+    bool is_cross_attn,
+    bool has_encoder_lens) {
   using Vec = at::vec::Vectorized<float>;
 
   // block length for heads
@@ -1429,8 +1442,9 @@ void decode_attention_grouped_kernel_impl(
       // get query
       const scalar_t* __restrict__ q_ptr = query + bs * q_strideM + h_start * q_strideH;
 
-      int64_t seq_len_kv = seq_lens[bs];
+      int64_t seq_len_kv = is_cross_attn ? encoder_lens[bs] : seq_lens[bs];
       int64_t req_pool_id = req_pool_indices[bs];
+      int64_t kv_offset = (has_encoder_lens && (!is_cross_attn)) ? encoder_lens[bs] : 0;
       TORCH_CHECK(seq_len_kv <= max_context_len, "seq_len_kv out of scope!");
       TORCH_CHECK(req_pool_id < max_num_reqs, "req_pool_id out of scope!");
 
@@ -1456,7 +1470,7 @@ void decode_attention_grouped_kernel_impl(
             /* A   */ q_ptr,
             /* B   */ k_buffer + head_kv_id * k_strideH,
             /* C   */ s_i,
-            /* ind */ req_to_token + req_pool_id * max_context_len + n,
+            /* ind */ req_to_token + req_pool_id * max_context_len + n + kv_offset,
             /* scl */ sm_scale,
             /* M   */ h_size,
             /* N   */ n_size,
@@ -1500,7 +1514,7 @@ void decode_attention_grouped_kernel_impl(
             /* A   */ s_delta,
             /* B   */ v_buffer + head_kv_id * v_strideH,
             /* C   */ v_prime,
-            /* ind */ req_to_token + req_pool_id * max_context_len + n,
+            /* ind */ req_to_token + req_pool_id * max_context_len + n + kv_offset,
             /* scl */ m_delta,
             /* M   */ h_size,
             /* N   */ head_size_v,
@@ -1519,6 +1533,10 @@ void decode_attention_grouped_kernel_impl(
               [s](Vec out) { return out * Vec(s); }, v_prime + h * l_stride1, v_prime + h * l_stride1, head_size_v);
           (v_prime + h * l_stride1)[head_size_v] = m_prime[h] + std::log(s_prime[h]);
         }
+      } else {
+        for (int64_t h = 0; h < h_size; ++h) {
+          (v_prime + h * l_stride1)[head_size_v] = -std::numeric_limits<float>::infinity();
+        }
       }
 
       // move to the next index
@@ -1540,32 +1558,30 @@ void decode_attention_grouped_kernel_impl(
 // req_to_token:     [max_num_reqs, max_context_len] int32 or int64
 // req_pool_indices: [num_seqs] int64
 // seq_lens:         [num_seqs] int64
+// encoder_lens:     [num_seqs] int64 or None
 //
 void decode_attention_cpu(
     at::Tensor& query,
     at::Tensor& k_buffer,
     at::Tensor& v_buffer,
     at::Tensor& output,
-    at::Tensor& key,
-    at::Tensor& value,
+    const std::optional<at::Tensor>& key,
+    const std::optional<at::Tensor>& value,
     at::Tensor& loc,
     at::Tensor& attn_logits,
     at::Tensor& req_to_token,
     at::Tensor& req_pool_indices,
     at::Tensor& seq_lens,
     double sm_scale,
-    double logit_cap) {
+    double logit_cap,
+    bool is_cross_attn,
+    std::optional<at::Tensor> encoder_lens) {
   CHECK_LAST_DIM_CONTIGUOUS_INPUT(query);
   CHECK_LAST_DIM_CONTIGUOUS_INPUT(k_buffer);
   CHECK_LAST_DIM_CONTIGUOUS_INPUT(v_buffer);
-  // for MLA, key and value shares the same storage and value could be non-contiguous
-  CHECK_LAST_DIM_CONTIGUOUS_INPUT(key);
-  CHECK_LAST_DIM_CONTIGUOUS_INPUT(value);
   CHECK_DIM(3, query);
   CHECK_DIM(3, k_buffer);
   CHECK_DIM(3, v_buffer);
-  CHECK_DIM(3, key);
-  CHECK_DIM(3, value);
   CHECK_DIM(1, loc);
 
   int64_t num_seqs = seq_lens.size(0);
@@ -1580,7 +1596,6 @@ void decode_attention_cpu(
 
   int64_t num_kv_splits = attn_logits.size(2);
 
-  CHECK_EQ(loc.numel(), num_seqs);
   CHECK_EQ(attn_logits.size(0), num_seqs);
   CHECK_EQ(attn_logits.size(1), num_heads);
   CHECK_EQ(attn_logits.size(3), head_size_v + 1);
@@ -1595,11 +1610,6 @@ void decode_attention_cpu(
   int64_t k_strideH = k_buffer.stride(1);
   int64_t v_strideN = v_buffer.stride(0);
   int64_t v_strideH = v_buffer.stride(1);
-  // strides for new key and value
-  int64_t nk_strideN = key.stride(0);
-  int64_t nk_strideH = key.stride(1);
-  int64_t nv_strideN = value.stride(0);
-  int64_t nv_strideH = value.stride(1);
 
   // check index data types
   const auto index_dtype = req_to_token.scalar_type();
@@ -1625,29 +1635,51 @@ void decode_attention_cpu(
   int num_threads = at::get_num_threads();
   int64_t size_per_thread = is_mla ? BLOCK_N * head_size + BLOCK_N * head_size_v : 0;
   auto buffer = at::empty({num_threads, size_per_thread}, k_buffer.options());
-
+  bool has_encoder_lens = encoder_lens.has_value();
+  // Since encoder_lens is not used when it is None, encoder_lens_t can be initialized as any tensor of int64_t dtype.
+  at::Tensor encoder_lens_t = seq_lens;
+  if (has_encoder_lens) {
+    encoder_lens_t = encoder_lens.value();
+    CHECK_EQ(encoder_lens_t.size(0), num_seqs);
+  }
   AT_DISPATCH_REDUCED_FLOATING_TYPES(query.scalar_type(), "decode_attention_kernel", [&] {
     AT_DISPATCH_INDEX_TYPES(index_dtype, "decode_attention_indices", [&] {
-      // update the kv buffer
-      decode_set_kv_buffer(
-          (scalar_t*)k_buffer_data,
-          (scalar_t*)v_buffer_data,
-          key.data_ptr<scalar_t>(),
-          value.data_ptr<scalar_t>(),
-          loc.data_ptr<int64_t>(),
-          num_seqs,
-          num_heads_kv,
-          head_size,
-          head_size_v,
-          k_strideN,
-          k_strideH,
-          v_strideN,
-          v_strideH,
-          nk_strideN,
-          nk_strideH,
-          nv_strideN,
-          nv_strideH,
-          is_mla);
+      if (key.has_value()) {
+        TORCH_CHECK(value.has_value(), "key and value should have values at the same time")
+        CHECK_EQ(loc.numel(), num_seqs);
+        auto key_tensor = key.value();
+        auto value_tensor = value.value();
+        // for MLA, key and value shares the same storage and value could be non-contiguous
+        CHECK_LAST_DIM_CONTIGUOUS_INPUT(key_tensor);
+        CHECK_LAST_DIM_CONTIGUOUS_INPUT(value_tensor);
+        CHECK_DIM(3, key_tensor);
+        CHECK_DIM(3, value_tensor);
+        // strides for new key and value
+        int64_t nk_strideN = key_tensor.stride(0);
+        int64_t nk_strideH = key_tensor.stride(1);
+        int64_t nv_strideN = value_tensor.stride(0);
+        int64_t nv_strideH = value_tensor.stride(1);
+        // update the kv buffer
+        decode_set_kv_buffer(
+            (scalar_t*)k_buffer_data,
+            (scalar_t*)v_buffer_data,
+            key_tensor.data_ptr<scalar_t>(),
+            value_tensor.data_ptr<scalar_t>(),
+            loc.data_ptr<int64_t>(),
+            num_seqs,
+            num_heads_kv,
+            head_size,
+            head_size_v,
+            k_strideN,
+            k_strideH,
+            v_strideN,
+            v_strideH,
+            nk_strideN,
+            nk_strideH,
+            nv_strideN,
+            nv_strideH,
+            is_mla);
+      }
 
       if (num_heads == num_heads_kv) {
         // MHA
@@ -1660,6 +1692,7 @@ void decode_attention_cpu(
             req_to_token.data_ptr<index_t>(),
             req_pool_indices.data_ptr<int64_t>(),
             seq_lens.data_ptr<int64_t>(),
+            encoder_lens_t.data_ptr<int64_t>(),
             num_seqs,
             num_heads,
             head_size,
@@ -1675,7 +1708,9 @@ void decode_attention_cpu(
             logit_cap,
             max_num_reqs,
             max_context_len,
-            max_total_num_tokens);
+            max_total_num_tokens,
+            is_cross_attn,
+            has_encoder_lens);
       } else if (is_mla) {
         // MLA
         decode_attention_mla_kernel_impl<scalar_t, index_t, BLOCK_N>(
@@ -1716,6 +1751,7 @@ void decode_attention_cpu(
             req_to_token.data_ptr<index_t>(),
             req_pool_indices.data_ptr<int64_t>(),
             seq_lens.data_ptr<int64_t>(),
+            encoder_lens_t.data_ptr<int64_t>(),
             num_seqs,
             num_heads,
             num_heads_kv,
@@ -1732,7 +1768,9 @@ void decode_attention_cpu(
             logit_cap,
             max_num_reqs,
             max_context_len,
-            max_total_num_tokens);
+            max_total_num_tokens,
+            is_cross_attn,
+            has_encoder_lens);
       }
     });
   });
diff --git a/sgl-kernel/csrc/cpu/extend.cpp b/sgl-kernel/csrc/cpu/extend.cpp
index b74f35a71..9cc378602 100644
--- a/sgl-kernel/csrc/cpu/extend.cpp
+++ b/sgl-kernel/csrc/cpu/extend.cpp
@@ -22,6 +22,7 @@ void extend_attention_kernel_impl(
     const index_t* __restrict__ req_to_token,
     const int64_t* __restrict__ req_pool_indices,
     const int64_t* __restrict__ seq_lens,
+    const int64_t* __restrict__ encoder_lens,
     const index_t* __restrict__ extend_seq_lens,
     const index_t* __restrict__ extend_start_loc,
     const void* __restrict__ buffer,
@@ -46,7 +47,9 @@ void extend_attention_kernel_impl(
     int max_total_num_tokens,
     int max_len_extend,
     int buffer_size_per_thread,
-    bool is_prefix_skipped) {
+    bool is_prefix_skipped,
+    bool is_cross_attn,
+    bool has_encoder_lens) {
   // strides
   const int o_strideM = num_heads * head_size_v;
   const int o_strideH = head_size_v;
@@ -91,6 +94,7 @@ void extend_attention_kernel_impl(
       int seq_extend_start_loc = extend_start_loc[bs];
 
       int req_pool_id = req_pool_indices[bs];
+      int kv_offset = (has_encoder_lens && (!is_cross_attn)) ? encoder_lens[bs] : 0;
       TORCH_CHECK(seq_len_prefix >= 0, "prefix len < 0!");
       TORCH_CHECK(seq_len <= max_context_len, "seq_len out of scope!");
       TORCH_CHECK(req_pool_id < max_num_reqs, "req_pool_id out of scope!");
@@ -115,10 +119,11 @@ void extend_attention_kernel_impl(
       fill_stub(v_prime, 0.f, m_size * head_size_v);
       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
-      for (int n = 0; n < seq_len_prefix; n += BLOCK_N) {
-        int n_size = std::min(BLOCK_N, seq_len_prefix - n);
+      int kv_start = 0;
+      int kv_end = is_cross_attn ? encoder_lens[bs] : seq_len_prefix;
+      for (int n = kv_start; n < kv_end; n += BLOCK_N) {
+        int n_size = std::min(BLOCK_N, kv_end - n);
 
         // `n_size` is K in 2nd gemm, pad to TILE_K;
         const int padded_n_size = div_up(n_size, TILE_K) * TILE_K;
@@ -127,7 +132,7 @@ void extend_attention_kernel_impl(
         pack_vnni<scalar_t, index_t>(
             /*    dst */ Btmp,
             /*    src */ k_buffer + head_kv_id * k_strideH,
-            /*    ind */ req_to_token + req_pool_id * max_context_len + n,
+            /*    ind */ req_to_token + req_pool_id * max_context_len + n + kv_offset,
             /*     N  */ n_size,
             /*     K  */ head_size,
             /* ld_src */ k_strideN,
@@ -153,7 +158,7 @@ void extend_attention_kernel_impl(
         pack_vnni2<scalar_t, index_t>(
             /*    dst */ Btmp,
             /*    src */ v_buffer + head_kv_id * v_strideH,
-            /*    ind */ req_to_token + req_pool_id * max_context_len + n,
+            /*    ind */ req_to_token + req_pool_id * max_context_len + n + kv_offset,
             /*     K  */ n_size,
             /*     N  */ head_size_v,
             /* ld_src */ v_strideN,
@@ -172,92 +177,88 @@ void extend_attention_kernel_impl(
             /* B     */ Btmp,
             /* C     */ v_prime);
       }  // loop with seq_len_prefix
+      if (!is_cross_attn) {
+        // 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) {
+          int n_size = std::min(BLOCK_N, num_keys - n);
 
-      // 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) {
-        int n_size = std::min(BLOCK_N, num_keys - n);
+          // `n_size` is K in 2nd gemm, pad to TILE_K;
+          const int padded_n_size = div_up(n_size, TILE_K) * TILE_K;
 
-        // `n_size` is K in 2nd gemm, pad to TILE_K;
-        const int padded_n_size = div_up(n_size, TILE_K) * TILE_K;
+          // get key and pack
+          pack_vnni<scalar_t>(
+              /*    dst */ Btmp,
+              /*    src */ k_extend + (seq_extend_start_loc + n) * ke_strideN + head_kv_id * ke_strideH,
+              /*     N  */ n_size,
+              /*     K  */ head_size,
+              /* ld_src */ ke_strideN,
+              /* ld_dst */ BLOCK_N);
 
-        // get key and pack
-        pack_vnni<scalar_t>(
-            /*    dst */ Btmp,
-            /*    src */ k_extend + (seq_extend_start_loc + n) * ke_strideN + head_kv_id * ke_strideH,
-            /*     N  */ n_size,
-            /*     K  */ head_size,
-            /* ld_src */ ke_strideN,
-            /* ld_dst */ BLOCK_N);
+          // calculate s_i <- Q @ K
+          at::native::cpublas::brgemm(
+              /* M     */ m_size,
+              /* N     */ n_size,
+              /* K     */ head_size,
+              /* lda   */ q_strideM,
+              /* ldb   */ BLOCK_N,
+              /* ldc   */ BLOCK_N,
+              /* add_C */ false,
+              /* A     */ q_ptr,
+              /* B     */ Btmp,
+              /* C     */ s_i);
 
-        // calculate s_i <- Q @ K
-        at::native::cpublas::brgemm(
-            /* M     */ m_size,
-            /* N     */ n_size,
-            /* K     */ head_size,
-            /* lda   */ q_strideM,
-            /* ldb   */ BLOCK_N,
-            /* ldc   */ BLOCK_N,
-            /* add_C */ false,
-            /* A     */ q_ptr,
-            /* B     */ Btmp,
-            /* C     */ s_i);
-
-        // apply causal mask
-        // [Note] condition to apply causal mask.
-        // Mask any block whose last key (n + n_size - 1) is strictly after the first query position (m), i.e. n +
-        // n_size - 1 > m. The original condition was `num_keys - n <= BLOCK_N` (last n-block only). That was correct
-        // when BLOCK_M <= BLOCK_N/2 because earlier n-blocks were guaranteed to contain only past keys.  With
-        // BLOCK_M=512, BLOCK_N=768:
-        //   BLOCK_M > BLOCK_N/2, so the first n-block can contain future keys.
-        //   Example: m=512 (mb=1), num_keys=1024, first n-block covers keys [0, 768).
-        //   Query row=0 is at position 512, so keys 513..767 are future and must be
-        //   masked — but `num_keys - 0 = 1024 > BLOCK_N` skips masking entirely,
-        //   producing wrong (non-causal) attention for rows 0..254 of this m-block.
-        if (n + n_size - 1 > m) {
-          for (int row = 0; row < m_size; ++row) {
-            int last_col = m + row - n;
-            // [Note] mask the entire row if last_col < 0.
-            // Clamp to -1: when n > m + row every key in this block is a future
-            // key, so the entire row should be masked.  Without this clamp,
-            // last_col+1 <= 0 and fill_stub would write before row_ptr.
-            // Example:
-            //  For max_len_extend > 4096 → selects BLOCK_M=512, BLOCK_N=768
-            //  m + BLOCK_M = 512 + 512 = 1024 > BLOCK_N = 768, this means we can have a a second n-block at n=768.
-            //  For m = 512, row = 0, n = 768, last_col = 512 + 0 - 768 = -256 → out of bounds write in fill_stub
-            last_col = std::max(last_col, -1);
-            // fill [last_col + 1, n_size) to -inf
-            float* row_ptr = s_i + row * BLOCK_N;
-            fill_stub(row_ptr + last_col + 1, -std::numeric_limits<float>::infinity(), n_size - last_col - 1);
+          // apply causal mask
+          // [Note] condition to apply causal mask.
+          // Mask any block whose last key (n + n_size - 1) is strictly after the first query position (m), i.e. n +
+          // n_size - 1 > m. The original condition was `num_keys - n <= BLOCK_N` (last n-block only). That was correct
+          // when BLOCK_M <= BLOCK_N/2 because earlier n-blocks were guaranteed to contain only past keys.  With
+          // BLOCK_M=512, BLOCK_N=768:
+          //   BLOCK_M > BLOCK_N/2, so the first n-block can contain future keys.
+          //   Example: m=512 (mb=1), num_keys=1024, first n-block covers keys [0, 768).
+          //   Query row=0 is at position 512, so keys 513..767 are future and must be
+          //   masked — but `num_keys - 0 = 1024 > BLOCK_N` skips masking entirely,
+          //   producing wrong (non-causal) attention for rows 0..254 of this m-block.
+          if (n + n_size - 1 > m) {
+            for (int row = 0; row < m_size; ++row) {
+              int last_col = m + row - n;
+              // [Note] mask the entire row if last_col < 0.
+              // Clamp to -1: when n > m + row every key in this block is a future
+              // key, so the entire row should be masked.  Without this clamp,
+              // last_col+1 <= 0 and fill_stub would write before row_ptr.
+              last_col = std::max(last_col, -1);
+              // fill [last_col + 1, n_size) to -inf
+              float* row_ptr = s_i + row * BLOCK_N;
+              fill_stub(row_ptr + last_col + 1, -std::numeric_limits<float>::infinity(), n_size - last_col - 1);
+            }
           }
-        }
 
-        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);
+          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);
 
-        // get value and pack
-        pack_vnni2<scalar_t>(
-            /*    dst */ Btmp,
-            /*    src */ v_extend + (seq_extend_start_loc + n) * ve_strideN + head_kv_id * ve_strideH,
-            /*     K  */ n_size,
-            /*     N  */ head_size_v,
-            /* ld_src */ ve_strideN,
-            /* ld_dst */ head_size_v);
-
-        // calculate V' <- s_delta @ V + V'
-        at::native::cpublas::brgemm(
-            /* M     */ m_size,
-            /* N     */ head_size_v,
-            /* K     */ padded_n_size,  // n_size
-            /* lda   */ BLOCK_N,
-            /* ldb   */ head_size_v,
-            /* ldc   */ head_size_v,
-            /* add_C */ true,
-            /* A     */ s_delta,
-            /* B     */ Btmp,
-            /* C     */ v_prime);
-      }  // loop with seq_len_extend
+          // get value and pack
+          pack_vnni2<scalar_t>(
+              /*    dst */ Btmp,
+              /*    src */ v_extend + (seq_extend_start_loc + n) * ve_strideN + head_kv_id * ve_strideH,
+              /*     K  */ n_size,
+              /*     N  */ head_size_v,
+              /* ld_src */ ve_strideN,
+              /* ld_dst */ head_size_v);
 
+          // calculate V' <- s_delta @ V + V'
+          at::native::cpublas::brgemm(
+              /* M     */ m_size,
+              /* N     */ head_size_v,
+              /* K     */ padded_n_size,  // n_size
+              /* lda   */ BLOCK_N,
+              /* ldb   */ head_size_v,
+              /* ldc   */ head_size_v,
+              /* add_C */ true,
+              /* A     */ s_delta,
+              /* B     */ Btmp,
+              /* C     */ v_prime);
+        }  // loop with seq_len_extend
+      }
       scalar_t* __restrict__ out_ptr = o_extend + (seq_extend_start_loc + m) * o_strideM + head_id * o_strideH;
       for (int row = 0; row < m_size; ++row) {
         float s = 1 / s_prime[row];
@@ -299,6 +300,7 @@ inline int resize_buffer(at::Tensor& buffer, int num_threads, int head_size, int
         req_to_token.data_ptr<index_t>(),                                                  \
         req_pool_indices.data_ptr<int64_t>(),                                              \
         seq_lens.data_ptr<int64_t>(),                                                      \
+        encoder_lens_t.data_ptr<int64_t>(),                                                \
         extend_seq_lens.data_ptr<index_t>(),                                               \
         extend_start_loc.data_ptr<index_t>(),                                              \
         buffer.data_ptr(),                                                                 \
@@ -323,7 +325,9 @@ inline int resize_buffer(at::Tensor& buffer, int num_threads, int head_size, int
         max_total_num_tokens,                                                              \
         max_len_extend,                                                                    \
         sz,                                                                                \
-        is_prefix_skipped);                                                                \
+        is_prefix_skipped,                                                                 \
+        is_cross_attn,                                                                     \
+        has_encoder_lens);                                                                 \
   } while (0)
 
 // q_extend, k_extend, v_extend, o_extend: contiguous tensors
@@ -340,11 +344,12 @@ inline int resize_buffer(at::Tensor& buffer, int num_threads, int head_size, int
 // seq_lens: [num_seqs] int64
 // extend_seq_lens: [num_seqs]
 // extend_start_loc: [num_seqs]
+// encoder_lens: [num_seqs] int64
 //
 void extend_attention_cpu(
     at::Tensor& q_extend,
-    at::Tensor& k_extend,
-    at::Tensor& v_extend,
+    const std::optional<at::Tensor>& k_extend_opt,
+    const std::optional<at::Tensor>& v_extend_opt,
     at::Tensor& o_extend,
     at::Tensor& k_buffer,
     at::Tensor& v_buffer,
@@ -355,7 +360,19 @@ void extend_attention_cpu(
     at::Tensor& extend_start_loc,
     int64_t max_len_extend,
     double sm_scale,
-    double logit_cap) {
+    double logit_cap,
+    bool is_cross_attn,
+    std::optional<at::Tensor> encoder_lens) {
+  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.
+  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;
+
   CHECK_LAST_DIM_CONTIGUOUS_INPUT(q_extend);
   CHECK_INPUT(o_extend);
   CHECK_LAST_DIM_CONTIGUOUS_INPUT(k_extend);
@@ -419,6 +436,13 @@ void extend_attention_cpu(
   int num_threads = at::get_num_threads();
   auto buffer = at::empty({}, q_extend.options().dtype(at::kChar));
 
+  bool has_encoder_lens = encoder_lens.has_value();
+  // Since encoder_lens is not used when it is None, encoder_lens_t can be initialized as any tensor of int64_t dtype.
+  at::Tensor encoder_lens_t = seq_lens;
+  if (has_encoder_lens) {
+    encoder_lens_t = encoder_lens.value();
+    CHECK_EQ(encoder_lens_t.size(0), num_seqs);
+  }
   AT_DISPATCH_REDUCED_FLOATING_TYPES(q_extend.scalar_type(), "extend_attention_kernel", [&] {
     AT_DISPATCH_INDEX_TYPES(index_dtype, "extend_attention_indices", [&] {
       if (max_len_extend <= 256) {
diff --git a/sgl-kernel/csrc/cpu/gemm.cpp b/sgl-kernel/csrc/cpu/gemm.cpp
index 09299887f..e66edacd4 100644
--- a/sgl-kernel/csrc/cpu/gemm.cpp
+++ b/sgl-kernel/csrc/cpu/gemm.cpp
@@ -722,10 +722,10 @@ at::Tensor convert_scale_packed(at::Tensor& scale) {
   return packed_scale;
 }
 
-// mat1 : [M, K]
+// mat1 : [*, K]
 // mat2 : [N, K] ([K, N] if use_fma_gemm)
 // bias : [N]
-// out  : [M, N]
+// out  : [*, N]
 //
 at::Tensor
 weight_packed_linear(at::Tensor& mat1, at::Tensor& mat2, const std::optional<at::Tensor>& bias, bool is_vnni) {
@@ -735,23 +735,25 @@ weight_packed_linear(at::Tensor& mat1, at::Tensor& mat2, const std::optional<at:
     use_fma_gemm = true;
   }
 
-  int64_t M = mat1.size(0);
-  int64_t K = mat1.size(1);
-  int64_t N = use_fma_gemm ? mat2.size(1) : mat2.size(0);
-
   CHECK_LAST_DIM_CONTIGUOUS_INPUT(mat1);
   CHECK_INPUT(mat2);
-  CHECK_DIM(2, mat1);
+  const int64_t ndim = mat1.ndimension();
+  auto input_sizes = mat1.sizes().vec();
+  int64_t N = use_fma_gemm ? mat2.size(1) : mat2.size(0);
+  int64_t K = use_fma_gemm ? mat1.size(1) : mat2.size(1);
+  int64_t M = use_fma_gemm ? mat1.size(0) : mat1.numel() / K;
   CHECK_DIM(2, mat2);
-  if (!use_fma_gemm) {
-    CHECK_EQ(mat1.size(1), K);
+  if (use_fma_gemm) {
+    CHECK_DIM(2, mat1);
+  } else {
+    CHECK_EQ(mat1.size(ndim - 1), K);
   }
 
   auto dispatch_type = mat1.scalar_type();
   auto out = at::empty({M, N}, mat1.options());
   // strides
   int64_t out_strideM = out.stride(0);
-  int64_t mat1_strideM = mat1.stride(0);
+  int64_t mat1_strideM = mat1.stride(-2);
 
   const bool has_bias = bias.has_value();
   const float* bias_data = nullptr;
@@ -787,7 +789,8 @@ weight_packed_linear(at::Tensor& mat1, at::Tensor& mat2, const std::optional<at:
     }
   });
 
-  return out;
+  input_sizes[ndim - 1] = N;
+  return out.view(input_sizes);
 }
 
 // mat1         : [M, K]
diff --git a/sgl-kernel/csrc/cpu/torch_extension_cpu.cpp b/sgl-kernel/csrc/cpu/torch_extension_cpu.cpp
index 16370d146..e86d249da 100644
--- a/sgl-kernel/csrc/cpu/torch_extension_cpu.cpp
+++ b/sgl-kernel/csrc/cpu/torch_extension_cpu.cpp
@@ -90,20 +90,22 @@ void decode_attention_cpu(
     at::Tensor& k_cache,
     at::Tensor& v_cache,
     at::Tensor& output,
-    at::Tensor& key,
-    at::Tensor& value,
+    const std::optional<at::Tensor>& key,
+    const std::optional<at::Tensor>& value,
     at::Tensor& loc,
     at::Tensor& attn_logits,
     at::Tensor& req_to_token,
     at::Tensor& req_pool_indices,
     at::Tensor& seq_lens,
     double sm_scale,
-    double logit_cap);
+    double logit_cap,
+    bool is_cross_attn,
+    std::optional<at::Tensor> encoder_lens);
 
 void extend_attention_cpu(
     at::Tensor& q_extend,
-    at::Tensor& k_extend,
-    at::Tensor& v_extend,
+    const std::optional<at::Tensor>& k_extend,
+    const std::optional<at::Tensor>& v_extend,
     at::Tensor& o_extend,
     at::Tensor& k_buffer,
     at::Tensor& v_buffer,
@@ -114,7 +116,9 @@ void extend_attention_cpu(
     at::Tensor& extend_start_loc,
     int64_t max_len_extend,
     double sm_scale,
-    double logit_cap);
+    double logit_cap,
+    bool is_cross_attn,
+    std::optional<at::Tensor> encoder_lens);
 
 // flash attention
 at::Tensor flash_attn_varlen_func(
@@ -462,16 +466,18 @@ TORCH_LIBRARY_FRAGMENT(sgl_kernel, m) {
 
   // decode
   m.def(
-      "decode_attention_cpu(Tensor query, Tensor k_cache, Tensor v_cahce, Tensor(a!) output, Tensor key, Tensor value, "
+      "decode_attention_cpu(Tensor query, Tensor k_cache, Tensor v_cahce, Tensor(a!) output, Tensor? key, Tensor? "
+      "value, "
       "Tensor loc, Tensor attn_logits, Tensor req_to_token, Tensor req_pool_indices, Tensor seq_lens, float sm_scale, "
-      "float logit_cap) -> ()");
+      "float logit_cap, bool is_cross_attn, Tensor? encoder_lens) -> ()");
   m.impl("decode_attention_cpu", torch::kCPU, &decode_attention_cpu);
 
   // extend
   m.def(
-      "extend_attention_cpu(Tensor q_extend, Tensor k_extend, Tensor v_extend, Tensor(a!) o_extend, Tensor k_buffer, "
+      "extend_attention_cpu(Tensor q_extend, Tensor? k_extend, Tensor? v_extend, Tensor(a!) o_extend, Tensor k_buffer, "
       "Tensor v_buffer, Tensor req_to_token, Tensor req_pool_indices, Tensor seq_lens, Tensor extend_seq_lens, Tensor "
-      "extend_start_loc, int max_len_extend, float sm_scale, float logit_cap) -> ()");
+      "extend_start_loc, int max_len_extend, float sm_scale, float logit_cap, bool is_cross_attn, Tensor? "
+      "encoder_lens) -> ()");
   m.impl("extend_attention_cpu", torch::kCPU, &extend_attention_cpu);
 
   // flash attn
diff --git a/test/registered/cpu/test_decode.py b/test/registered/cpu/test_decode.py
index 81385d3b2..a29d9b76c 100644
--- a/test/registered/cpu/test_decode.py
+++ b/test/registered/cpu/test_decode.py
@@ -21,9 +21,11 @@ class TestDecodeAttention(CustomTestCase):
         req_to_token: torch.Tensor,
         req_pool_indices: torch.Tensor,
         seq_lens: torch.Tensor,
+        encoder_lens=None,
         scaling=None,
         enable_gqa=False,
         causal=False,
+        is_cross_attn=False,
     ):
         # [num_tokens, num_heads, head_size] -> [num_heads, num_tokens, head_size]
         query = query.movedim(0, query.dim() - 2)
@@ -33,14 +35,21 @@ class TestDecodeAttention(CustomTestCase):
             seq_len_q = 1
             seq_len_kv = seq_lens[seq_idx]
             end_q = start_q + seq_len_q
-            end_kv = start_kv + seq_len_kv
+            if encoder_lens is not None:
+                start_kv = 0 if is_cross_attn else encoder_lens[seq_idx]
+                end_kv = (
+                    encoder_lens[seq_idx] if is_cross_attn else start_kv + seq_len_kv
+                )
+            else:
+                start_kv = 0
+                end_kv = start_kv + seq_len_kv
 
             per_req_query = query[:, start_q:end_q, :]
 
             # get key and value from cache. per_req_tokens contains the kv cache
             # index for each token in the sequence.
             req_pool_idx = req_pool_indices[seq_idx]
-            per_req_tokens = req_to_token[req_pool_idx, :seq_len_kv]
+            per_req_tokens = req_to_token[req_pool_idx, start_kv:end_kv]
             per_req_key = k_cache[per_req_tokens].movedim(0, query.dim() - 2)
             per_req_value = v_cache[per_req_tokens].movedim(0, query.dim() - 2)
 
@@ -61,10 +70,13 @@ class TestDecodeAttention(CustomTestCase):
 
         return output
 
-    def _test_grouped_decode_attention_once(self, B, H_Q, H_KV, D, D_V, dtype, device):
+    def _test_grouped_decode_attention_once(
+        self, B, H_Q, H_KV, D, D_V, is_cross_attn, dtype, device
+    ):
         # This represents the number of tokens already in the sequence
         seq_len = 1024
-        total_tokens = B * seq_len
+        encoder_len = 10
+        total_tokens = B * (seq_len + encoder_len)
         sm_scale = 1.0 / (D**0.5)
         logit_cap = 0.0
         num_kv_splits = 8
@@ -91,11 +103,12 @@ class TestDecodeAttention(CustomTestCase):
 
         req_to_token = (
             torch.arange(total_tokens, device=device)
-            .reshape(B, seq_len)
+            .reshape(B, seq_len + encoder_len)
             .to(torch.int32)
         )
         b_req_idx = torch.arange(B, device=device).to(torch.int64)
         b_seq_len = torch.full((B,), seq_len, device=device).to(torch.int64)
+        encoder_lens = torch.full((B,), encoder_len, device=device).to(torch.int64)
 
         attn_logits = torch.empty(
             (B, H_Q, num_kv_splits, D_V + 1),
@@ -114,8 +127,8 @@ class TestDecodeAttention(CustomTestCase):
             k_buffer,
             v_buffer,
             o,
-            key,
-            value,
+            key if not is_cross_attn else None,
+            value if not is_cross_attn else None,
             loc,
             attn_logits,
             req_to_token,
@@ -123,6 +136,8 @@ class TestDecodeAttention(CustomTestCase):
             b_seq_len,
             sm_scale,
             logit_cap,
+            is_cross_attn,
+            encoder_lens,
         )
 
         self._run_sdpa_forward_decode(
@@ -135,15 +150,16 @@ class TestDecodeAttention(CustomTestCase):
             b_seq_len,
             scaling=sm_scale,
             enable_gqa=enable_gqa,
+            encoder_lens=encoder_lens,
+            is_cross_attn=is_cross_attn,
         )
-
         cos_sim = torch.nn.functional.cosine_similarity(
             o.flatten(), o_grouped.flatten(), dim=0
         )
         self.assertGreater(cos_sim.item(), 0.99)
         torch.testing.assert_close(o, o_grouped, atol=3e-2, rtol=1e-6)
 
-    def _test_grouped_decode_attention(self, device="cpu"):
+    def _test_grouped_decode_attention(self, device="cuda"):
         configs = [
             (2, 16, 16, 64, 64),
             (2, 16, 1, 16, 16),
@@ -161,7 +177,10 @@ class TestDecodeAttention(CustomTestCase):
         for B, H_Q, H_KV, D, D_V in configs:
             for dtype in [torch.bfloat16, torch.float16]:
                 self._test_grouped_decode_attention_once(
-                    B, H_Q, H_KV, D, D_V, dtype=dtype, device=device
+                    B, H_Q, H_KV, D, D_V, False, dtype=dtype, device=device
+                )
+                self._test_grouped_decode_attention_once(
+                    B, H_Q, H_KV, D, D_V, True, dtype=dtype, device=device
                 )
 
     def test_grouped_decode_attention(self):
diff --git a/test/registered/cpu/test_extend.py b/test/registered/cpu/test_extend.py
index 2ea6e3fba..ab55fe383 100644
--- a/test/registered/cpu/test_extend.py
+++ b/test/registered/cpu/test_extend.py
@@ -24,9 +24,11 @@ class TestExtendAttention(CustomTestCase):
         seq_lens: torch.Tensor,
         extend_prefix_lens: torch.Tensor,
         extend_seq_lens: torch.Tensor,
+        encoder_lens=None,
         scaling=None,
         enable_gqa=False,
         causal=False,
+        is_cross_attn=False,
     ):
 
         assert seq_lens.shape[0] == extend_prefix_lens.shape[0]
@@ -43,7 +45,14 @@ class TestExtendAttention(CustomTestCase):
 
             seq_len_kv = seq_lens[seq_idx]
             end_q = start_q + extend_seq_len_q
-            end_kv = start_kv + seq_len_kv
+            if encoder_lens is not None:
+                start_kv = 0 if is_cross_attn else encoder_lens[seq_idx]
+                end_kv = (
+                    encoder_lens[seq_idx] if is_cross_attn else start_kv + seq_len_kv
+                )
+            else:
+                start_kv = 0
+                end_kv = start_kv + seq_len_kv
 
             per_req_query = query[:, start_q:end_q, :]
             per_req_query_redudant = torch.empty(
@@ -57,7 +66,7 @@ class TestExtendAttention(CustomTestCase):
             # get key and value from cache. per_req_tokens contains the kv cache
             # index for each token in the sequence.
             req_pool_idx = req_pool_indices[seq_idx]
-            per_req_tokens = req_to_token[req_pool_idx, :seq_len_kv]
+            per_req_tokens = req_to_token[req_pool_idx, start_kv:end_kv]
             per_req_key = k_cache[per_req_tokens].movedim(0, query.dim() - 2)
             per_req_value = v_cache[per_req_tokens].movedim(0, query.dim() - 2)
 
@@ -86,6 +95,7 @@ class TestExtendAttention(CustomTestCase):
         D,
         DV,
         mla=False,
+        is_cross_attn=False,
         *,
         b_seq_len_prefix=None,
         b_seq_len_extend=None,
@@ -94,32 +104,36 @@ class TestExtendAttention(CustomTestCase):
 
         if b_seq_len_prefix is None:
             b_seq_len_prefix = torch.randint(1, N_CTX // 2, (B,), dtype=torch.int32)
-            if mla:
-                b_seq_len_prefix.zero_()
         else:
             b_seq_len_prefix = torch.as_tensor(b_seq_len_prefix, dtype=torch.int32)
 
+        encoder_lens = torch.randint(1, N_CTX // 2, (B,), dtype=torch.int64)
+        if mla:
+            b_seq_len_prefix.zero_()
+            encoder_lens.zero_()
+
         if b_seq_len_extend is None:
             b_seq_len_extend = torch.randint(1, N_CTX // 2, (B,), dtype=torch.int32)
         else:
             b_seq_len_extend = torch.as_tensor(b_seq_len_extend, dtype=torch.int32)
-
         b_seq_len = b_seq_len_prefix + b_seq_len_extend
-        max_len_in_batch = torch.max(b_seq_len, 0)[0].item()
+        max_len_in_batch = (
+            torch.max(b_seq_len, 0)[0].item() + torch.max(encoder_lens, 0)[0].item()
+        )
 
         b_req_idx = torch.arange(B, dtype=torch.int32)
         req_to_tokens = torch.empty((B, max_len_in_batch), dtype=torch.int32)
         b_start_loc = torch.zeros((B,), dtype=torch.int32)
-        b_start_loc[1:] = torch.cumsum(b_seq_len[:-1], 0)
+        b_start_loc[1:] = torch.cumsum(b_seq_len[:-1] + encoder_lens[:-1], 0)
         b_start_loc_extend = torch.zeros((B,), dtype=torch.int32)
         b_start_loc_extend[1:] = torch.cumsum(b_seq_len_extend[:-1], 0)
 
         for i in range(B):
-            req_to_tokens[i, : b_seq_len[i]] = torch.arange(
-                b_start_loc[i], b_start_loc[i] + b_seq_len[i]
+            req_to_tokens[i, : b_seq_len[i] + encoder_lens[i]] = torch.arange(
+                b_start_loc[i], b_start_loc[i] + b_seq_len[i] + encoder_lens[i]
             )
 
-        total_token_num = torch.sum(b_seq_len).item()
+        total_token_num = torch.sum(b_seq_len).item() + torch.sum(encoder_lens).item()
         extend_token_num = torch.sum(b_seq_len_extend).item()
 
         H_BUF = 1 if mla else H_KV
@@ -131,8 +145,10 @@ class TestExtendAttention(CustomTestCase):
         q_extend = torch.empty((extend_token_num, H_Q, D), dtype=dtype)
 
         for i in range(B):
-            extend_start_in_buffer = b_start_loc[i] + b_seq_len_prefix[i]
-            extend_end_in_buffer = b_start_loc[i] + b_seq_len[i]
+            extend_start_in_buffer = (
+                b_start_loc[i] + b_seq_len_prefix[i] + encoder_lens[i]
+            )
+            extend_end_in_buffer = b_start_loc[i] + b_seq_len[i] + encoder_lens[i]
             extend_start = b_start_loc_extend[i]
             extend_end = b_start_loc_extend[i] + b_seq_len_extend[i]
             k_extend[extend_start:extend_end] = k_buffer[
@@ -178,7 +194,9 @@ class TestExtendAttention(CustomTestCase):
             b_seq_len_extend,
             scaling=sm_scale,
             enable_gqa=enable_gqa,
-            causal=True,
+            causal=not is_cross_attn,
+            is_cross_attn=is_cross_attn,
+            encoder_lens=encoder_lens,
         )
 
         o_extend = torch.empty((extend_token_num, H_Q, DV), dtype=dtype)
@@ -197,16 +215,29 @@ class TestExtendAttention(CustomTestCase):
             max_len_extend,
             sm_scale,
             logit_cap,
+            is_cross_attn,
+            encoder_lens,
         )
 
         torch.testing.assert_close(o_ref, o_extend, atol=1e-2, rtol=1e-2)
 
     def test_extend_attention(self):
         for is_mla in [True, False]:
-            self._test_extend_attention_once(1, 123, 1, 1, 128, 96, is_mla)
-            self._test_extend_attention_once(1, 123, 16, 1, 128, 96, is_mla)
-            self._test_extend_attention_once(4, 1230, 16, 4, 128, 96, is_mla)
-            self._test_extend_attention_once(1, 9000, 16, 1, 32, 32, is_mla)
+            for is_cross_attn in [True, False]:
+                if is_mla and is_cross_attn:
+                    continue
+                self._test_extend_attention_once(
+                    1, 123, 1, 1, 128, 96, is_mla, is_cross_attn
+                )
+                self._test_extend_attention_once(
+                    1, 123, 16, 1, 128, 96, is_mla, is_cross_attn
+                )
+                self._test_extend_attention_once(
+                    4, 1230, 16, 4, 128, 96, is_mla, is_cross_attn
+                )
+                self._test_extend_attention_once(
+                    1, 9000, 16, 1, 32, 32, is_mla, is_cross_attn
+                )
 
     def test_extend_attention_large_seq_causal_mask(self):
         self._test_extend_attention_once(
@@ -220,6 +251,18 @@ class TestExtendAttention(CustomTestCase):
             b_seq_len_extend=[5000],
         )
 
+    def test_extend_attention_gqa_partial_extend_with_prefix(self):
+        self._test_extend_attention_once(
+            B=1,
+            N_CTX=256,
+            H_Q=16,
+            H_KV=4,
+            D=128,
+            DV=96,
+            b_seq_len_prefix=[97],
+            b_seq_len_extend=[37],
+        )
+
 
 if __name__ == "__main__":
     unittest.main()
diff --git a/test/registered/cpu/test_gemm.py b/test/registered/cpu/test_gemm.py
index 909e0ab27..3bd0bb961 100644
--- a/test/registered/cpu/test_gemm.py
+++ b/test/registered/cpu/test_gemm.py
@@ -35,6 +35,7 @@ class TestGemm(CustomTestCase):
     N = [16, 32 * 13]
     K = [32 * 16]
     has_bias = [False, True]
+    dim = [2, 3, 4, 5]
 
     M_int8 = [2, 128]
     N_int8 = [32 * 12]
@@ -52,10 +53,16 @@ class TestGemm(CustomTestCase):
     N_gptq = [4096]
     K_gptq = [4096]
 
-    def _bf16_gemm(self, M, N, K, has_bias):
+    def _bf16_gemm(self, M, N, K, has_bias, dim):
 
         mat1 = torch.randn(M, K, dtype=torch.bfloat16)
         mat2 = torch.randn(N, K, dtype=torch.bfloat16)
+        if dim == 3:
+            mat1 = mat1.unsqueeze(0).repeat(2, 1, 1)
+        if dim == 4:
+            mat1 = mat1.unsqueeze(0).unsqueeze(0).repeat(2, 2, 1, 1)
+        if dim == 5:
+            mat1 = mat1.unsqueeze(0).unsqueeze(0).unsqueeze(0).repeat(2, 2, 2, 1, 1)
 
         ref = torch.matmul(mat1.float(), mat2.float().t())
         if has_bias:
@@ -83,12 +90,14 @@ class TestGemm(CustomTestCase):
             self.N,
             self.K,
             self.has_bias,
+            self.dim,
         ):
             with self.subTest(
                 M=params[0],
                 N=params[1],
                 K=params[2],
                 has_bias=params[3],
+                dim=params[4],
             ):
                 self._bf16_gemm(*params)
 
diff --git a/test/registered/cpu/test_mla.py b/test/registered/cpu/test_mla.py
index 93867ea79..da11a281f 100644
--- a/test/registered/cpu/test_mla.py
+++ b/test/registered/cpu/test_mla.py
@@ -118,6 +118,8 @@ class TestMLA(CustomTestCase):
             b_seq_len,
             sm_scale,
             logit_cap,
+            False,
+            None,
         )
 
         self._run_sdpa_forward_decode(
diff --git a/test/srt/cpu/test_decode.py b/test/srt/cpu/test_decode.py
index aaf5f5d47..b04f3d51c 100644
--- a/test/srt/cpu/test_decode.py
+++ b/test/srt/cpu/test_decode.py
@@ -18,9 +18,11 @@ class TestDecodeAttention(CustomTestCase):
         req_to_token: torch.Tensor,
         req_pool_indices: torch.Tensor,
         seq_lens: torch.Tensor,
+        encoder_lens=None,
         scaling=None,
         enable_gqa=False,
         causal=False,
+        is_cross_attn=False,
     ):
         # [num_tokens, num_heads, head_size] -> [num_heads, num_tokens, head_size]
         query = query.movedim(0, query.dim() - 2)
@@ -30,14 +32,21 @@ class TestDecodeAttention(CustomTestCase):
             seq_len_q = 1
             seq_len_kv = seq_lens[seq_idx]
             end_q = start_q + seq_len_q
-            end_kv = start_kv + seq_len_kv
+            if encoder_lens is not None:
+                start_kv = 0 if is_cross_attn else encoder_lens[seq_idx]
+                end_kv = (
+                    encoder_lens[seq_idx] if is_cross_attn else start_kv + seq_len_kv
+                )
+            else:
+                start_kv = 0
+                end_kv = start_kv + seq_len_kv
 
             per_req_query = query[:, start_q:end_q, :]
 
             # get key and value from cache. per_req_tokens contains the kv cache
             # index for each token in the sequence.
             req_pool_idx = req_pool_indices[seq_idx]
-            per_req_tokens = req_to_token[req_pool_idx, :seq_len_kv]
+            per_req_tokens = req_to_token[req_pool_idx, start_kv:end_kv]
             per_req_key = k_cache[per_req_tokens].movedim(0, query.dim() - 2)
             per_req_value = v_cache[per_req_tokens].movedim(0, query.dim() - 2)
 
@@ -58,10 +67,13 @@ class TestDecodeAttention(CustomTestCase):
 
         return output
 
-    def _test_grouped_decode_attention_once(self, B, H_Q, H_KV, D, D_V, dtype, device):
+    def _test_grouped_decode_attention_once(
+        self, B, H_Q, H_KV, D, D_V, is_cross_attn, dtype, device
+    ):
         # This represents the number of tokens already in the sequence
         seq_len = 1024
-        total_tokens = B * seq_len
+        encoder_len = 10
+        total_tokens = B * (seq_len + encoder_len)
         sm_scale = 1.0 / (D**0.5)
         logit_cap = 0.0
         num_kv_splits = 8
@@ -88,11 +100,12 @@ class TestDecodeAttention(CustomTestCase):
 
         req_to_token = (
             torch.arange(total_tokens, device=device)
-            .reshape(B, seq_len)
+            .reshape(B, seq_len + encoder_len)
             .to(torch.int32)
         )
         b_req_idx = torch.arange(B, device=device).to(torch.int64)
         b_seq_len = torch.full((B,), seq_len, device=device).to(torch.int64)
+        encoder_lens = torch.full((B,), encoder_len, device=device).to(torch.int64)
 
         attn_logits = torch.empty(
             (B, H_Q, num_kv_splits, D_V + 1),
@@ -111,8 +124,8 @@ class TestDecodeAttention(CustomTestCase):
             k_buffer,
             v_buffer,
             o,
-            key,
-            value,
+            key if not is_cross_attn else None,
+            value if not is_cross_attn else None,
             loc,
             attn_logits,
             req_to_token,
@@ -120,6 +133,8 @@ class TestDecodeAttention(CustomTestCase):
             b_seq_len,
             sm_scale,
             logit_cap,
+            is_cross_attn,
+            encoder_lens,
         )
 
         self._run_sdpa_forward_decode(
@@ -132,8 +147,9 @@ class TestDecodeAttention(CustomTestCase):
             b_seq_len,
             scaling=sm_scale,
             enable_gqa=enable_gqa,
+            encoder_lens=encoder_lens,
+            is_cross_attn=is_cross_attn,
         )
-
         cos_sim = torch.nn.functional.cosine_similarity(
             o.flatten(), o_grouped.flatten(), dim=0
         )
@@ -158,7 +174,10 @@ class TestDecodeAttention(CustomTestCase):
         for B, H_Q, H_KV, D, D_V in configs:
             for dtype in [torch.bfloat16, torch.float16]:
                 self._test_grouped_decode_attention_once(
-                    B, H_Q, H_KV, D, D_V, dtype=dtype, device=device
+                    B, H_Q, H_KV, D, D_V, False, dtype=dtype, device=device
+                )
+                self._test_grouped_decode_attention_once(
+                    B, H_Q, H_KV, D, D_V, True, dtype=dtype, device=device
                 )
 
     def test_grouped_decode_attention(self):
diff --git a/test/srt/cpu/test_extend.py b/test/srt/cpu/test_extend.py
index 215241862..58a22c37b 100644
--- a/test/srt/cpu/test_extend.py
+++ b/test/srt/cpu/test_extend.py
@@ -21,9 +21,11 @@ class TestExtendAttention(CustomTestCase):
         seq_lens: torch.Tensor,
         extend_prefix_lens: torch.Tensor,
         extend_seq_lens: torch.Tensor,
+        encoder_lens=None,
         scaling=None,
         enable_gqa=False,
         causal=False,
+        is_cross_attn=False,
     ):
 
         assert seq_lens.shape[0] == extend_prefix_lens.shape[0]
@@ -40,7 +42,14 @@ class TestExtendAttention(CustomTestCase):
 
             seq_len_kv = seq_lens[seq_idx]
             end_q = start_q + extend_seq_len_q
-            end_kv = start_kv + seq_len_kv
+            if encoder_lens is not None:
+                start_kv = 0 if is_cross_attn else encoder_lens[seq_idx]
+                end_kv = (
+                    encoder_lens[seq_idx] if is_cross_attn else start_kv + seq_len_kv
+                )
+            else:
+                start_kv = 0
+                end_kv = start_kv + seq_len_kv
 
             per_req_query = query[:, start_q:end_q, :]
             per_req_query_redudant = torch.empty(
@@ -54,7 +63,7 @@ class TestExtendAttention(CustomTestCase):
             # get key and value from cache. per_req_tokens contains the kv cache
             # index for each token in the sequence.
             req_pool_idx = req_pool_indices[seq_idx]
-            per_req_tokens = req_to_token[req_pool_idx, :seq_len_kv]
+            per_req_tokens = req_to_token[req_pool_idx, start_kv:end_kv]
             per_req_key = k_cache[per_req_tokens].movedim(0, query.dim() - 2)
             per_req_value = v_cache[per_req_tokens].movedim(0, query.dim() - 2)
 
@@ -83,6 +92,7 @@ class TestExtendAttention(CustomTestCase):
         D,
         DV,
         mla=False,
+        is_cross_attn=False,
         *,
         b_seq_len_prefix=None,
         b_seq_len_extend=None,
@@ -91,32 +101,36 @@ class TestExtendAttention(CustomTestCase):
 
         if b_seq_len_prefix is None:
             b_seq_len_prefix = torch.randint(1, N_CTX // 2, (B,), dtype=torch.int32)
-            if mla:
-                b_seq_len_prefix.zero_()
         else:
             b_seq_len_prefix = torch.as_tensor(b_seq_len_prefix, dtype=torch.int32)
 
+        encoder_lens = torch.randint(1, N_CTX // 2, (B,), dtype=torch.int64)
+        if mla:
+            b_seq_len_prefix.zero_()
+            encoder_lens.zero_()
+
         if b_seq_len_extend is None:
             b_seq_len_extend = torch.randint(1, N_CTX // 2, (B,), dtype=torch.int32)
         else:
             b_seq_len_extend = torch.as_tensor(b_seq_len_extend, dtype=torch.int32)
-
         b_seq_len = b_seq_len_prefix + b_seq_len_extend
-        max_len_in_batch = torch.max(b_seq_len, 0)[0].item()
+        max_len_in_batch = (
+            torch.max(b_seq_len, 0)[0].item() + torch.max(encoder_lens, 0)[0].item()
+        )
 
         b_req_idx = torch.arange(B, dtype=torch.int32)
         req_to_tokens = torch.empty((B, max_len_in_batch), dtype=torch.int32)
         b_start_loc = torch.zeros((B,), dtype=torch.int32)
-        b_start_loc[1:] = torch.cumsum(b_seq_len[:-1], 0)
+        b_start_loc[1:] = torch.cumsum(b_seq_len[:-1] + encoder_lens[:-1], 0)
         b_start_loc_extend = torch.zeros((B,), dtype=torch.int32)
         b_start_loc_extend[1:] = torch.cumsum(b_seq_len_extend[:-1], 0)
 
         for i in range(B):
-            req_to_tokens[i, : b_seq_len[i]] = torch.arange(
-                b_start_loc[i], b_start_loc[i] + b_seq_len[i]
+            req_to_tokens[i, : b_seq_len[i] + encoder_lens[i]] = torch.arange(
+                b_start_loc[i], b_start_loc[i] + b_seq_len[i] + encoder_lens[i]
             )
 
-        total_token_num = torch.sum(b_seq_len).item()
+        total_token_num = torch.sum(b_seq_len).item() + torch.sum(encoder_lens).item()
         extend_token_num = torch.sum(b_seq_len_extend).item()
 
         H_BUF = 1 if mla else H_KV
@@ -128,8 +142,10 @@ class TestExtendAttention(CustomTestCase):
         q_extend = torch.empty((extend_token_num, H_Q, D), dtype=dtype)
 
         for i in range(B):
-            extend_start_in_buffer = b_start_loc[i] + b_seq_len_prefix[i]
-            extend_end_in_buffer = b_start_loc[i] + b_seq_len[i]
+            extend_start_in_buffer = (
+                b_start_loc[i] + b_seq_len_prefix[i] + encoder_lens[i]
+            )
+            extend_end_in_buffer = b_start_loc[i] + b_seq_len[i] + encoder_lens[i]
             extend_start = b_start_loc_extend[i]
             extend_end = b_start_loc_extend[i] + b_seq_len_extend[i]
             k_extend[extend_start:extend_end] = k_buffer[
@@ -175,7 +191,9 @@ class TestExtendAttention(CustomTestCase):
             b_seq_len_extend,
             scaling=sm_scale,
             enable_gqa=enable_gqa,
-            causal=True,
+            causal=not is_cross_attn,
+            is_cross_attn=is_cross_attn,
+            encoder_lens=encoder_lens,
         )
 
         o_extend = torch.empty((extend_token_num, H_Q, DV), dtype=dtype)
@@ -194,16 +212,29 @@ class TestExtendAttention(CustomTestCase):
             max_len_extend,
             sm_scale,
             logit_cap,
+            is_cross_attn,
+            encoder_lens,
         )
 
         torch.testing.assert_close(o_ref, o_extend, atol=1e-2, rtol=1e-2)
 
     def test_extend_attention(self):
         for is_mla in [True, False]:
-            self._test_extend_attention_once(1, 123, 1, 1, 128, 96, is_mla)
-            self._test_extend_attention_once(1, 123, 16, 1, 128, 96, is_mla)
-            self._test_extend_attention_once(4, 1230, 16, 4, 128, 96, is_mla)
-            self._test_extend_attention_once(1, 9000, 16, 1, 32, 32, is_mla)
+            for is_cross_attn in [True, False]:
+                if is_mla and is_cross_attn:
+                    continue
+                self._test_extend_attention_once(
+                    1, 123, 1, 1, 128, 96, is_mla, is_cross_attn
+                )
+                self._test_extend_attention_once(
+                    1, 123, 16, 1, 128, 96, is_mla, is_cross_attn
+                )
+                self._test_extend_attention_once(
+                    4, 1230, 16, 4, 128, 96, is_mla, is_cross_attn
+                )
+                self._test_extend_attention_once(
+                    1, 9000, 16, 1, 32, 32, is_mla, is_cross_attn
+                )
 
     def test_extend_attention_large_seq_causal_mask(self):
         self._test_extend_attention_once(
diff --git a/test/srt/cpu/test_gemm.py b/test/srt/cpu/test_gemm.py
index cde23a371..8395ac298 100644
--- a/test/srt/cpu/test_gemm.py
+++ b/test/srt/cpu/test_gemm.py
@@ -32,6 +32,7 @@ class TestGemm(CustomTestCase):
     N = [16, 32 * 13]
     K = [32 * 16]
     has_bias = [False, True]
+    dim = [2, 3, 4, 5]
 
     M_int8 = [2, 128]
     N_int8 = [32 * 12]
@@ -49,10 +50,16 @@ class TestGemm(CustomTestCase):
     N_gptq = [4096]
     K_gptq = [4096]
 
-    def _bf16_gemm(self, M, N, K, has_bias):
+    def _bf16_gemm(self, M, N, K, has_bias, dim):
 
         mat1 = torch.randn(M, K, dtype=torch.bfloat16)
         mat2 = torch.randn(N, K, dtype=torch.bfloat16)
+        if dim == 3:
+            mat1 = mat1.unsqueeze(0).repeat(2, 1, 1)
+        if dim == 4:
+            mat1 = mat1.unsqueeze(0).unsqueeze(0).repeat(2, 2, 1, 1)
+        if dim == 5:
+            mat1 = mat1.unsqueeze(0).unsqueeze(0).unsqueeze(0).repeat(2, 2, 2, 1, 1)
 
         ref = torch.matmul(mat1.float(), mat2.float().t())
         if has_bias:
@@ -80,12 +87,14 @@ class TestGemm(CustomTestCase):
             self.N,
             self.K,
             self.has_bias,
+            self.dim,
         ):
             with self.subTest(
                 M=params[0],
                 N=params[1],
                 K=params[2],
                 has_bias=params[3],
+                dim=params[4],
             ):
                 self._bf16_gemm(*params)
 
diff --git a/test/srt/cpu/test_mla.py b/test/srt/cpu/test_mla.py
index 3caa109d5..3db205581 100644
--- a/test/srt/cpu/test_mla.py
+++ b/test/srt/cpu/test_mla.py
@@ -115,6 +115,8 @@ class TestMLA(CustomTestCase):
             b_seq_len,
             sm_scale,
             logit_cap,
+            False,
+            None,
         )
 
         self._run_sdpa_forward_decode(