[CPU][sgl-kernel] extend_attention_cpu and flash_attn_varlen_func: fix nan for large seq (#22434)

Co-authored-by: Ma Mingfei <mingfei.ma@intel.com>

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<float>::infinity(), n_size - last_col - 1);

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<float>::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<float>::infinity(), n_size - last_col - 1);

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

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