Remove head size 16 examples

Author: petercad

applications/flash_attention_v2/collective/xe_fmha_fwd_epilogue.hpp

@@ -285,6 +285,7 @@ class FMHAFwdEpilogue {
           }
         }
       }
+
       return std::make_tuple(rA, rA_sum, active);
     }
   }

applications/flash_attention_v2/collective/xe_fmha_fwd_mainloop.hpp

@@ -293,9 +293,11 @@ struct FMHAFwdMainloop<XeDefault<Stages>, CausalMask_,
       if (check_remainder_k && K == blk_k1 - 1) {
         FragSRow k_rem_mask;
         int k = get<0>(tKgK(0,0,0,K,0)) + get_sub_group().get_local_id()[0];
+        CUTLASS_PRAGMA_UNROLL
         for (int i = 0; i < k_rem_mask.size(); i++, k += intel::sg_size) {
           k_rem_mask(i) = (k < shape<0>(K_2D)) ? ElementS(sycl::nan(0u)) : ElementS(-INFINITY);
         }
+        CUTLASS_PRAGMA_UNROLL
         for (int i = 0; i < tSrS.size(); i++) {
           tSrS(i) = sycl::fmin(tSrS(i), broadcast<1>(k_rem_mask, tSrS, i));
         }
@@ -312,6 +314,11 @@ struct FMHAFwdMainloop<XeDefault<Stages>, CausalMask_,
       for (int i = 0; i < tArP.size(); i++)
         tArP(i) = static_cast<typename TiledMMAPV::ValTypeA>(tSrS(i));
 #endif
+#define XSHOW(x) if (cute::thread0()) { print(#x ": "); print(x); print("\n"); }
+      XSHOW(tSrS.layout())
+      XSHOW(tSrS.tv_layout())
+      XSHOW(tArP.layout())
+      XSHOW(tArP.tv_layout())
 
       /* GEMM 2: A += P * V, split in v dimension */
       CUTLASS_PRAGMA_UNROLL

examples/06_bmg_flash_attention/06_xe_fmha_fwd.cpp

@@ -133,6 +133,7 @@ int main(int argc, const char **argv) {
     using SubgroupLayoutQK = Layout<Shape<_1, _8, _1>>;
 
 #elif HEAD_DIM == 192
+    /* FIXME: does not compile yet due to 1x24 output tile */
     using ShapeQK = Shape<_1, _512, _64>;
     using ShapePV = Shape<_1, _32, _512>;
     using ShapeOut = Shape<_1, _192>;

examples/06_bmg_flash_attention/CMakeLists.txt

@@ -32,7 +32,7 @@ set(CUTLASS_APPLICATIONS_DIR ${CMAKE_SOURCE_DIR}/applications)
 set(TEST_NO_PAGED "")
 set(TEST_PAGED "--use_paged_kv")
 
-foreach(HEAD_DIM 16 64 96 128 192)
+foreach(HEAD_DIM 64 96 128 192)
 
   cutlass_example_add_executable(
     06_xe_fmha_fwd_prefill_hdim${HEAD_DIM}