change block tile and tests ok

example/ck_tile/02_layernorm2d/instances/layernorm2d_fwd_api.cpp

@@ -33,7 +33,7 @@ float layernorm2d_fwd_b16_(layernorm2d_fwd_traits /*t*/,
     // clang-format off
     //                                            rm  rn  tm   tn  vn  pd     mv     2p
     if(a.n <= 64) {
-            r = layernorm2d_fwd_<trait_<data_type, 1,  1,  4,  64, 1,  true,  false, false>>(s, a);
+            r = layernorm2d_fwd_<trait_<data_type, 1,  1,  4,  16, 4,  true,  false, false>>(s, a);
     }
     else if(a.n <= 128) {
         if (a.n % 2 == 0)

example/ck_tile/02_layernorm2d/instances/layernorm2d_fwd_bf16_n64_n128_instance.cpp

@@ -7,6 +7,7 @@
 // clang-format off
 //                                                       rm  rn  tm  tn  vn  pd     mv     2p
 template float layernorm2d_fwd_<trait_<ck_tile::bf16_t, 1,  1,  4,  64, 1,  true , false, false>>(const S&, A);
+template float layernorm2d_fwd_<trait_<ck_tile::bf16_t, 1,  1,  4,  16, 4,  true , false, false>>(const S&, A);
 template float layernorm2d_fwd_<trait_<ck_tile::bf16_t, 1,  1,  4,  64, 2,  true , false, false>>(const S&, A);
 template float layernorm2d_fwd_<trait_<ck_tile::bf16_t, 1,  2,  4,  64, 1,  true , false, false>>(const S&, A);
 // clang-format on

example/ck_tile/02_layernorm2d/instances/layernorm2d_fwd_fp16_n64_n128_instance.cpp

@@ -7,6 +7,7 @@
 // clang-format off
 //                                                       rm  rn  tm  tn  vn  pd     mv     2p
 template float layernorm2d_fwd_<trait_<ck_tile::fp16_t, 1,  1,  4,  64, 1,  true , false, false>>(const S&, A);
+template float layernorm2d_fwd_<trait_<ck_tile::fp16_t, 1,  1,  4,  16, 4,  true , false, false>>(const S&, A);
 template float layernorm2d_fwd_<trait_<ck_tile::fp16_t, 1,  1,  4,  64, 2,  true , false, false>>(const S&, A);
 template float layernorm2d_fwd_<trait_<ck_tile::fp16_t, 1,  2,  4,  64, 1,  true , false, false>>(const S&, A);
 // clang-format on

example/ck_tile/02_layernorm2d/layernorm2d_fwd.hpp

@@ -101,7 +101,8 @@ struct layernorm2d_fwd_traits_
     using WarpTile   = ck_tile::sequence<Warp_M, Warp_N>;
     using Vector     = ck_tile::sequence<1, Vector_N_>;
 
-    using Shape = ck_tile::Layernorm2dShape<BlockTile, BlockWarps, WarpTile, Vector>;
+    using Shape = ck_tile::Layernorm2dShape<BlockTile, BlockWarps, 
+            WarpTile, Vector, ThreadPerBlock_M_ * ThreadPerBlock_N_>;
 
     static constexpr bool kPadN           = kPadN_;
     static constexpr bool kSaveMeanInvStd = kSaveMeanInvStd_;