[Gemm,Sm100] Fix num_sf_tmem_cols when not blockscaled

Author: tridao

quack/gemm_sm100.py

@@ -480,15 +480,18 @@ def _setup_attributes(self, epilogue_args: EpilogueArguments, varlen_args: Varle
         # to release acc_pipeline early.
         # The two approaches perform about the same.
         self.overlap_accum_sf = self.blockscaled and self.num_acc_stage == 1
-        num_sf_tmem_cols = (
-            (
-                cute.ceil_div(self.cta_tile_shape_mnk[0], 128)
-                + cute.ceil_div(self.cta_tile_shape_mnk[1], 128)
+        if const_expr(self.overlap_accum_sf):
+            num_sf_tmem_cols = (
+                (
+                    cute.ceil_div(self.cta_tile_shape_mnk[0], 128)
+                    + cute.ceil_div(self.cta_tile_shape_mnk[1], 128)
+                )
+                * 4  # 4 cols per stage
+                * (self.mma_inst_shape_mnk[2] // self.sf_vec_size)
             )
-            * 4  # 4 cols per stage
-            * (self.mma_inst_shape_mnk[2] // self.sf_vec_size)
-        )
-        self.iter_acc_early_release = num_sf_tmem_cols // cute.size(self.epi_tile[1])
+            self.iter_acc_early_release = num_sf_tmem_cols // cute.size(self.epi_tile[1])
+        else:
+            self.iter_acc_early_release = -1
 
     @cute.jit
     def __call__(

tools/dump_sass.py

@@ -1,8 +1,9 @@
 #!/usr/bin/env python3
 """Dump PTX and SASS of cute-dsl kernels from any script.
 
-Sets CUTE_DSL_KEEP_CUBIN=1 and CUTE_DSL_KEEP_PTX=1, runs the target script,
-then disassembles all generated .cubin files with nvdisasm.
+Disables the QuACK persistent kernel cache, sets CUTE_DSL_KEEP_CUBIN=1 and
+CUTE_DSL_KEEP_PTX=1, runs the target script, then disassembles all generated
+.cubin files with nvdisasm.
 
 Usage::
 
@@ -45,6 +46,11 @@ def main():
     parser.add_argument("script", help="Python script to run")
     parser.add_argument("-o", "--output-dir", default="dump_sass_out", help="Output directory")
     parser.add_argument("--ptx-only", action="store_true", help="Skip SASS disassembly")
+    parser.add_argument(
+        "--use-cache",
+        action="store_true",
+        help="Allow QuACK to use its persistent .o cache instead of forcing recompilation",
+    )
     args = parser.parse_args(our_argv)
 
     script = Path(args.script)
@@ -59,13 +65,17 @@ def main():
             f.unlink()
 
     env = os.environ.copy()
+    if not args.use_cache:
+        env["QUACK_CACHE_ENABLED"] = "0"
     env["CUTE_DSL_KEEP_PTX"] = "1"
     env["CUTE_DSL_KEEP_CUBIN"] = "1"
     env["CUTE_DSL_DUMP_DIR"] = str(out_dir.resolve())
 
     cmd = [sys.executable, str(script)] + script_args
     print(f"Running: {' '.join(cmd)}")
     print(f"Dump dir: {out_dir.resolve()}\n")
+    if not args.use_cache:
+        print("QuACK cache: disabled via QUACK_CACHE_ENABLED=0\n")
     subprocess.run(cmd, env=env)
 
     ptx_files = sorted(out_dir.glob("*.ptx"))