From 92b0d447b0fc37909ba24a4e751d688026627081 Mon Sep 17 00:00:00 2001 From: Chunyuan WU Date: Sat, 29 Aug 2026 10:49:43 +0800 Subject: [PATCH] [CPU] Fix wrongly causal-masked bidirectional attention (#35434) --- python/sglang/kernels/aot/csrc/cpu/extend.cpp | 11 +++++++---- .../aot/csrc/cpu/torch_extension_cpu.cpp | 5 +++-- .../srt/layers/attention/intel_amx_backend.py | 6 ++++++ test/registered/cpu/test_extend.py | 19 ++++++++++++++++++- 4 files changed, 34 insertions(+), 7 deletions(-) diff --git a/python/sglang/kernels/aot/csrc/cpu/extend.cpp b/python/sglang/kernels/aot/csrc/cpu/extend.cpp index db2d46857..3a3280e69 100644 --- a/python/sglang/kernels/aot/csrc/cpu/extend.cpp +++ b/python/sglang/kernels/aot/csrc/cpu/extend.cpp @@ -63,6 +63,7 @@ void extend_attention_kernel_impl( int64_t sliding_window_size, bool is_prefix_skipped, bool is_cross_attn, + bool 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 encoder_lens, std::optional sinks, - std::optional tree_mask) { + std::optional 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 diff --git a/python/sglang/kernels/aot/csrc/cpu/torch_extension_cpu.cpp b/python/sglang/kernels/aot/csrc/cpu/torch_extension_cpu.cpp index 573674ab2..dedb4985f 100644 --- a/python/sglang/kernels/aot/csrc/cpu/torch_extension_cpu.cpp +++ b/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 encoder_lens, std::optional sinks, - std::optional tree_mask); + std::optional 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 diff --git a/python/sglang/srt/layers/attention/intel_amx_backend.py b/python/sglang/srt/layers/attention/intel_amx_backend.py index 8dc211d1b..476aeb835 100644 --- a/python/sglang/srt/layers/attention/intel_amx_backend.py +++ b/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) diff --git a/test/registered/cpu/test_extend.py b/test/registered/cpu/test_extend.py index 99d40e446..af76a08f5 100644 --- a/test/registered/cpu/test_extend.py +++ b/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()