fix index

Author: wuxun-zhang

applications/flash_attention_v2/kernel/xe_fhma_fwd_kernel.hpp

@@ -562,8 +562,8 @@ class XeFMHAFwdDynamicSplitKernel {
         }
         CUTLASS_PRAGMA_UNROLL
         for (int i = 0; i < size(FragARow{}.shape()); ++i) {
-          merged_res(i + size(FragA{}.shape())) = tA_max(i);
-          merged_res(i + 1 + size(FragA{}.shape())) = tA_sum(i);
+          merged_res(2 * i + size(FragA{}.shape())) = tA_max(i);
+          merged_res(2 * i + 1 + size(FragA{}.shape())) = tA_sum(i);
         }
         copy(merged_res, tPartial);
 
@@ -608,8 +608,8 @@ class XeFMHAFwdDynamicSplitKernel {
 
             CUTLASS_PRAGMA_UNROLL
             for (int i = 0; i < size(FragARow{}.shape()); ++i) {
-              tA_max(i) = merged_res(i + size(FragA{}.shape()));
-              tA_sum(i) = merged_res(i + 1 + size(FragA{}.shape()));
+              tA_max(i) = merged_res(2 * i + size(FragA{}.shape()));
+              tA_sum(i) = merged_res(2 * i + 1 + size(FragA{}.shape()));
             }
 
             continue;

examples/06_bmg_flash_attention/06_xe_fmha_fwd.cpp

@@ -112,7 +112,7 @@ int main(int argc, const char **argv) {
 #define NUM_SG _16
 #define KV_TILE_SIZE _256
 #else
-#define NUM_SG _16
+#define NUM_SG _8
 #define KV_TILE_SIZE _512
 #endif