From 6c8921458475d00170b5d3020bd9162539c94cab Mon Sep 17 00:00:00 2001 From: Chunyuan WU Date: Fri, 17 Apr 2026 13:01:01 +0800 Subject: [PATCH] [CPU][sgl-kernel] `extend_attention_cpu` and `flash_attn_varlen_func`: fix `nan` for large seq (#22434) Co-authored-by: Ma Mingfei --- sgl-kernel/csrc/cpu/extend.cpp | 21 ++++++++++++++- sgl-kernel/csrc/cpu/flash_attn.cpp | 10 +++++-- test/srt/cpu/test_extend.py | 42 ++++++++++++++++++++++++++---- test/srt/cpu/test_flash_attn.py | 34 ++++++++++++++++++++++++ 4 files changed, 99 insertions(+), 8 deletions(-) diff --git a/sgl-kernel/csrc/cpu/extend.cpp b/sgl-kernel/csrc/cpu/extend.cpp index 63da654e4..1d0e2ecae 100644 --- a/sgl-kernel/csrc/cpu/extend.cpp +++ b/sgl-kernel/csrc/cpu/extend.cpp @@ -204,9 +204,28 @@ void extend_attention_kernel_impl( /* C */ s_i); // apply causal mask - if (num_keys - n <= BLOCK_N) { + // [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::infinity(), n_size - last_col - 1); diff --git a/sgl-kernel/csrc/cpu/flash_attn.cpp b/sgl-kernel/csrc/cpu/flash_attn.cpp index 58852671a..8f202de86 100644 --- a/sgl-kernel/csrc/cpu/flash_attn.cpp +++ b/sgl-kernel/csrc/cpu/flash_attn.cpp @@ -149,9 +149,12 @@ void flash_attn_kernel_impl( /* C */ s_i); // apply causal mask - if (causal && num_keys - n <= BLOCK_N) { + // See [Note] condition to apply causal mask. + if (causal && n + n_size - 1 > m) { for (int row = 0; row < m_size; ++row) { int last_col = m + row - n; + // See [Note] mask the entire row if last_col < 0. + 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::infinity(), n_size - last_col - 1); @@ -329,9 +332,12 @@ void flash_attn_varlen_kernel_impl( /* C */ s_i); // apply causal mask - if (causal && num_keys - n <= BLOCK_N) { + // See [Note] condition to apply causal mask. + if (causal && n + n_size - 1 > m) { for (int row = 0; row < m_size; ++row) { int last_col = m + row - n; + // See [Note] mask the entire row if last_col < 0. + 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::infinity(), n_size - last_col - 1); diff --git a/test/srt/cpu/test_extend.py b/test/srt/cpu/test_extend.py index 5e0585933..ce1a888be 100644 --- a/test/srt/cpu/test_extend.py +++ b/test/srt/cpu/test_extend.py @@ -74,13 +74,33 @@ class TestExtendAttention(CustomTestCase): start_q, start_kv = end_q, end_kv return output - def _test_extend_attention_once(self, B, N_CTX, H_Q, H_KV, D, DV, mla=False): + def _test_extend_attention_once( + self, + B, + N_CTX, + H_Q, + H_KV, + D, + DV, + mla=False, + *, + b_seq_len_prefix=None, + b_seq_len_extend=None, + ): dtype = torch.bfloat16 - b_seq_len_prefix = torch.randint(1, N_CTX // 2, (B,), dtype=torch.int32) - if mla: - b_seq_len_prefix.zero_() - b_seq_len_extend = torch.randint(1, N_CTX // 2, (B,), dtype=torch.int32) + 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) + + 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() @@ -185,6 +205,18 @@ class TestExtendAttention(CustomTestCase): 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) + def test_extend_attention_large_seq_causal_mask(self): + self._test_extend_attention_once( + B=1, + N_CTX=5001, + H_Q=8, + H_KV=2, + D=64, + DV=64, + b_seq_len_prefix=[0], + b_seq_len_extend=[5000], + ) + if __name__ == "__main__": unittest.main() diff --git a/test/srt/cpu/test_flash_attn.py b/test/srt/cpu/test_flash_attn.py index e964636d1..bd012d187 100644 --- a/test/srt/cpu/test_flash_attn.py +++ b/test/srt/cpu/test_flash_attn.py @@ -203,6 +203,40 @@ class TestFlashAttn(CustomTestCase): atol = rtol = precision[dtype] torch.testing.assert_close(out_ref, out, atol=atol, rtol=rtol) + def _test_flash_attn_large_seq_causal_mask_once(self, seqlens): + dtype = torch.bfloat16 + num_heads = 8 + num_heads_kv = 2 + head_dim = 64 + + seqlens_t = torch.tensor(seqlens, dtype=torch.int32) + cu_seqlens = torch.zeros(len(seqlens) + 1, dtype=torch.int32) + cu_seqlens[1:] = torch.cumsum(seqlens_t, 0) + total = cu_seqlens[-1].item() + max_seqlen = seqlens_t.max().item() + + q = torch.randn(total, num_heads, head_dim, dtype=dtype) + k = torch.randn(total, num_heads_kv, head_dim, dtype=dtype) + v = torch.randn(total, num_heads_kv, head_dim, dtype=dtype) + + out_ref = flash_attn_varlen_ref( + q, k, v, cu_seqlens, cu_seqlens, is_causal=True, enable_gqa=True + ) + out = flash_attn_varlen_func( + q, k, v, cu_seqlens, cu_seqlens, max_seqlen, max_seqlen, True + ) + + atol = rtol = precision[dtype] + torch.testing.assert_close(out_ref, out, atol=atol, rtol=rtol) + + def test_flash_attn_large_seq_causal_mask(self): + # Non-varlen path: single sequence, has_varlen_sequences returns False + # → dispatches to flash_attn_kernel_impl. + self._test_flash_attn_large_seq_causal_mask_once([5000]) + # Varlen path: sequences with different lengths, has_varlen_sequences + # returns True → dispatches to flash_attn_varlen_kernel_impl + self._test_flash_attn_large_seq_causal_mask_once([5000, 4999]) + if __name__ == "__main__": unittest.main()