[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);