[CPU] Fix wrongly causal-masked bidirectional attention (#35434)

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 is_causal,
     bool kv_from_cache,
     bool has_encoder_lens,
     bool has_sink) {
@@ -236,8 +237,8 @@ void extend_attention_kernel_impl(
             /* C     */ v_prime);
       }  // loop with seq_len_prefix
       if (!is_cross_attn && !kv_from_cache) {
-        // stage 2: compute the triangle part
-        int num_keys = std::min(seq_len_extend, m + BLOCK_M);
+        // stage 2: compute the triangle part (or, when !is_causal, the full square)
+        int num_keys = is_causal ? std::min(seq_len_extend, m + BLOCK_M) : seq_len_extend;
         for (int n = 0; n < num_keys; n += BLOCK_N) {
           int n_size = std::min(BLOCK_N, num_keys - n);
 
@@ -286,7 +287,7 @@ void extend_attention_kernel_impl(
                 }
               }
             }
-          } else if (n + n_size - 1 > m) {
+          } else if (is_causal && n + n_size - 1 > m) {
             // 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 +
@@ -421,6 +422,7 @@ inline int resize_buffer(at::Tensor& buffer, int num_threads, int head_size, int
         sliding_window_size,                                                               \
         is_prefix_skipped,                                                                 \
         is_cross_attn,                                                                     \
+        is_causal,                                                                         \
         kv_from_cache,                                                                     \
         has_encoder_lens,                                                                  \
         has_sink);                                                                         \
@@ -463,7 +465,8 @@ void extend_attention_cpu(
     int64_t sliding_window_size,
     std::optional<at::Tensor> encoder_lens,
     std::optional<at::Tensor> sinks,
-    std::optional<at::Tensor> tree_mask) {
+    std::optional<at::Tensor> tree_mask,
+    bool is_causal = true) {
   TORCH_CHECK(k_extend_opt.has_value() == v_extend_opt.has_value(), "k_extend and v_extend must be given together");
   // A KV-shared layer (Gemma 4) passes no extend K/V - the layer it shares with
   // already wrote them to the cache, so this kernel masks causally itself. Cross

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

@@ -245,7 +245,8 @@ void extend_attention_cpu(
     int64_t sliding_window_size,
     std::optional<at::Tensor> encoder_lens,
     std::optional<at::Tensor> sinks,
-    std::optional<at::Tensor> tree_mask);
+    std::optional<at::Tensor> tree_mask,
+    bool is_causal = true);
 
 // flash attention
 at::Tensor flash_attn_varlen_func(
@@ -719,7 +720,7 @@ TORCH_LIBRARY_FRAGMENT(sgl_kernel, m) {
       "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, bool is_cross_attn, int "
       "sliding_window_size, Tensor? "
-      "encoder_lens, Tensor? sinks, Tensor? tree_mask=None) -> ()");
+      "encoder_lens, Tensor? sinks, Tensor? tree_mask=None, bool is_causal=True) -> ()");
   m.impl("extend_attention_cpu", torch::kCPU, &extend_attention_cpu);
 
   // flash attn

python/sglang/srt/layers/attention/intel_amx_backend.py

@@ -5,6 +5,7 @@ from typing import TYPE_CHECKING
 import torch
 
 from sglang.srt.layers.attention.base_attn_backend import AttentionBackend
+from sglang.srt.layers.radix_attention import AttentionType
 from sglang.srt.mem_cache.memory_pool import KVWriteLoc
 from sglang.srt.mem_cache.swa_memory_pool import SWAKVPool
 from sglang.srt.model_executor.forward_batch_info import ForwardBatch
@@ -218,6 +219,10 @@ class IntelAMXAttnBackend(AttentionBackend):
         if seq_lens.dtype != torch.int64:
             seq_lens = seq_lens.to(torch.int64)
 
+        is_causal = True
+        if layer.is_cross_attention or layer.attn_type == AttentionType.ENCODER_ONLY:
+            is_causal = False
+
         # 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(
@@ -240,6 +245,7 @@ class IntelAMXAttnBackend(AttentionBackend):
             forward_batch.encoder_lens,
             sinks,
             tree_mask,
+            is_causal,
         )
         return o.view(-1, layer.tp_q_head_num * layer.v_head_dim)
 

test/registered/cpu/test_extend.py

@@ -198,6 +198,7 @@ class TestExtendAttention(CustomTestCase):
         b_seq_len_prefix=None,
         b_seq_len_extend=None,
         kv_from_cache=False,
+        is_causal=True,
     ):
         dtype = torch.bfloat16
 
@@ -316,7 +317,7 @@ class TestExtendAttention(CustomTestCase):
                 b_seq_len_extend,
                 scaling=sm_scale,
                 enable_gqa=enable_gqa,
-                causal=not is_cross_attn,
+                causal=(not is_cross_attn) and is_causal,
                 is_cross_attn=is_cross_attn,
                 encoder_lens=encoder_lens,
             )
@@ -341,6 +342,8 @@ class TestExtendAttention(CustomTestCase):
             sliding_window if sliding_window is not None else 0,
             encoder_lens,
             sinks if has_sink else None,
+            None,  # tree_mask
+            is_causal,
         )
 
         torch.testing.assert_close(o_ref, o_extend, atol=1e-2, rtol=1e-2)
@@ -421,6 +424,20 @@ class TestExtendAttention(CustomTestCase):
             b_seq_len_extend=[37],
         )
 
+    def test_extend_attention_bidirectional(self):
+        # Test for is_causal=False: encoder-only self-attention (e.g. bge-reranker)
+        self._test_extend_attention_once(
+            B=4,
+            N_CTX=123,
+            H_Q=16,
+            H_KV=4,
+            D=128,
+            DV=96,
+            b_seq_len_prefix=[0, 0, 0, 0],
+            b_seq_len_extend=[41, 90, 123, 5],
+            is_causal=False,
+        )
+
 
 if __name__ == "__main__":
     unittest.main()