From da47886e34bfc3bd2038a1fdff5dd889dc74af08 Mon Sep 17 00:00:00 2001
From: botbw <wang1570@e.ntu.edu.sg>
Date: Thu, 31 Jul 2025 10:12:33 +0800
Subject: [PATCH] Fix example bug (#2351)

---
 examples/cute/tutorial/sgemm_sm80.cu | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/examples/cute/tutorial/sgemm_sm80.cu b/examples/cute/tutorial/sgemm_sm80.cu
index 50914548..29fb51e5 100644
--- a/examples/cute/tutorial/sgemm_sm80.cu
+++ b/examples/cute/tutorial/sgemm_sm80.cu
@@ -232,6 +232,7 @@ gemm_device(ProblemShape shape_MNK, CtaTiler cta_tiler,
 
   // Size of the register pipeline
   auto K_BLOCK_MAX = size<2>(tCrA);
+  CUTE_STATIC_ASSERT_V(K_BLOCK_MAX == size<2>(tXrA));
 
   // PREFETCH register pipeline
   if (K_BLOCK_MAX > 1) {
@@ -371,7 +372,7 @@ gemm_tn(int m, int n, int k,
                                     Layout<Shape<_16,_8>,Stride<_8,_1>>{},  // Thr layout 16x8 k-major
                                     Layout<Shape< _1,_8>>{});               // Val layout  1x8 n-major
 
-  TiledMMA mmaC = make_tiled_mma(SM80_16x8x8_F16F16F16F16_TN{},
+  TiledMMA mmaC = make_tiled_mma(SM80_16x8x16_F16F16F16F16_TN{},
                                  Layout<Shape<_2,_2>>{},    // 2x2x1 MMA Atoms
                                  Tile<_32,_32,_16>{});      // 32x32x16 Tiled MMA for LDSM