diff --git a/kernel/include/gemmini_mmio.h b/kernel/include/gemmini_mmio.h
index 45030b94..cb86ffca 100644
--- a/kernel/include/gemmini_mmio.h
+++ b/kernel/include/gemmini_mmio.h
@@ -10,9 +10,9 @@
 // 64KB
 // #define SMEM_SIZE 0x10000
 // 128KB (FP16 GEMM)
-// #define SMEM_SIZE 0x20000
+#define SMEM_SIZE 0x20000
 // 256KB (FlashAttention)
-#define SMEM_SIZE 0x40000
+// #define SMEM_SIZE 0x40000
 
 #define SMEM_MASK (SMEM_SIZE - 1)
 #define SMEM_ADDR_END (SMEM_BASE + SMEM_SIZE)