diff --git a/lib/include/gemmini_mmio.h b/lib/include/gemmini_mmio.h
index ef783c8a..833ac9f0 100644
--- a/lib/include/gemmini_mmio.h
+++ b/lib/include/gemmini_mmio.h
@@ -12,9 +12,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)