dbg

stack-info: PR: #1911, branch: shunting314/stack/30

Author: shunting314

examples/distributed/fp8_matmul_reduce_scatter.py

@@ -33,13 +33,16 @@
     max_mismatch_pct=1e-3,
 )
 
+config = helion.Config(
+    block_sizes=[64, 64, 32],  # M, N, K
+    num_warps=8,
+    num_stages=3,
+)
+
+# config = helion.Config(block_sizes=[64, 128, 128], indexing=['pointer', 'pointer', 'pointer', 'tensor_descriptor', 'pointer', 'pointer', 'pointer', 'pointer', 'tensor_descriptor', 'pointer', 'pointer', 'pointer', 'tensor_descriptor', 'pointer'], l2_groupings=[1], load_eviction_policies=['last', '', '', '', '', 'first', '', '', '', '', '', 'last'], loop_orders=[[0, 1]], num_sm_multiplier=2, num_stages=1, num_warps=8, pid_type='persistent_blocked', range_flattens=[None, None], range_multi_buffers=[True, None], range_unroll_factors=[3, 0], range_warp_specializes=[])
 
 @helion.kernel(
-    config=helion.Config(
-        block_sizes=[64, 64, 32],  # M, N, K
-        num_warps=8,
-        num_stages=3,
-    ),
+    config=config,
     static_shapes=True,
     ignore_warnings=[helion.exc.TensorOperationInWrapper],
     autotune_baseline_accuracy_check_fn=functools.partial(

helion/autotuner/base_search.py

@@ -635,8 +635,9 @@ def _validate_against_baseline(
             custom_check = self.settings.autotune_baseline_accuracy_check_fn
             if custom_check is not None:
                 custom_check(output, self._baseline_output)
-                if len(self._mutated_arg_indices) > 0:
-                    custom_check(args, self._baseline_post_args)
+                if os.getenv("CHECK_INPUT_ACCURACY", "1") == "1":
+                    if len(self._mutated_arg_indices) > 0:
+                        custom_check(args, self._baseline_post_args)
             else:
                 _assert_close(
                     output,