Fix AMX GQA extend attention (#25180)

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

sgl-kernel/csrc/cpu/flash_attn.h

@@ -134,9 +134,10 @@ struct flash_attn_softmax {
           v_prime + row * head_size_v,
           head_size_v);
 
-      // pad s_delta with 0 first and then convert to scalar_t
+      // Keep s_delta row-major for the following brgemm(P @ V), and only
+      // convert the columns that brgemm will consume.
       fill_stub(s_delta + row * BLOCK_N + n_size, 0.f, padded_n_size - n_size);
-      copy_stub<scalar_t, BLOCK_N>(s_delta2 + row * BLOCK_N, s_delta + row * BLOCK_N);
+      copy_stub<scalar_t>(s_delta2 + row * BLOCK_N, s_delta + row * BLOCK_N, 1.f, padded_n_size);
     }
   }
 };

test/srt/cpu/test_extend.py

@@ -217,6 +217,18 @@ class TestExtendAttention(CustomTestCase):
             b_seq_len_extend=[5000],
         )
 
+    def test_extend_attention_gqa_partial_extend_with_prefix(self):
+        self._test_extend_attention_once(
+            B=1,
+            N_CTX=256,
+            H_Q=16,
+            H_KV=4,
+            D=128,
+            DV=96,
+            b_seq_len_prefix=[97],
+            b_seq_len_extend=[37],
+        )
+
 
 if __name__ == "__main__":
     unittest.main()