[Inductor][FP8] Validate exhaustive autotuning for FP8 Inductor templates (pytorch#161442)

Summary:
X-link: meta-pytorch/tritonbench#355


Validate exhaustive autotuning for FP8 Inductor templates: scaled MM templates require `block_k >= 32`. Before, exhaustive autotuning defaulted to a limited set of autotuning configs, as limitations for exhaustively autotuning on FP8 shapes had not been tested.

Test Plan:
```
CUDA_VISIBLE_DEVICES=0 TRITON_PRINT_AUTOTUNING=1 TRITON_ALWAYS_COMPILE=1 TORCH_LOGS=+inductor TORCHINDUCTOR_FORCE_DISABLE_CACHES=1 ENABLE_PERSISTENT_TMA_MATMUL=1 TORCHINDUCTOR_MAX_AUTOTUNE_GEMM=1 TORCHINDUCTOR_MAX_AUTOTUNE_GEMM_SEARCH_SPACE=DEFAULT buck2 run mode/{opt,inplace} pytorch/t
ritonbench:run -- --op fp8_gemm --only torch_fp8_gemm,pt2_fp8_gemm --metrics tflops,accuracy --input-loader=/home/jananisriram/personal/exhaustive_autotune_rowwise_persistent_tma/json_fi
les/rowwise_ptma_0.json --output="/home/jananisriram/personal/exhaustive_autotune_rowwise_persistent_tma/autotune/gpu0_bench.csv" --atol=1e-2 --rtol=0.5 2>&1 | tee ~/personal/exhaustive_
autotune_rowwise_persistent_tma/autotune/gpu0.log
```
autotunes on the maximum configs available, rather than the defaults, and skips configs not compatible with TMA.

Rollback Plan:

Reviewed By: coconutruben

Differential Revision: D80958642

Author: jananisriram

torch/_inductor/template_heuristics.py

@@ -1293,9 +1293,9 @@ def get_template_configs(
         Convert config lists to template kwargs.
         This replaces the logic from choices.get_mm_configs and inlines mm_options.
         """
-        assert isinstance(kernel_inputs, MMKernelInputs), (
-            f"{self.__class__.__name__} requires MMKernelInputs"
-        )
+        assert isinstance(
+            kernel_inputs, MMKernelInputs
+        ), f"{self.__class__.__name__} requires MMKernelInputs"
         input_nodes = kernel_inputs.nodes()
         if len(input_nodes) < 2:
             raise ValueError(f"Need at least 2 input tensors, got {len(input_nodes)}")
@@ -1467,9 +1467,9 @@ def get_template_configs(
         input_nodes = kernel_inputs.nodes()
 
         # Initial assertion from mm_common.scaled_mm_options
-        assert len(input_nodes) >= 4, (
-            f"scaled_mm requires at least 4 inputs, got {len(input_nodes)}"
-        )
+        assert (
+            len(input_nodes) >= 4
+        ), f"scaled_mm requires at least 4 inputs, got {len(input_nodes)}"
 
         # Extract scale tensors (typically scale_a and scale_b are input_nodes[2] and input_nodes[3])
         scale_a = input_nodes[2]
@@ -1522,9 +1522,11 @@ class ScaledTMAConfigMixin(ScaledMMConfigMixin):
 
     def _filter_configs(self, configs: list[BaseConfig]) -> list[BaseConfig]:
         """
-        TMA specific filtering, as num_warps=2 not safe for TMA
+        TMA specific filtering:
+        - num_warps=2 not safe for TMA
+        - block_k >= 32 required for TMA (requires inner-most dimension >= 32)
         """
-        configs = [c for c in configs if c.num_warps != 2]
+        configs = [c for c in configs if c.num_warps != 2 and c.block_k >= 32]
         return super()._filter_configs(configs)
 
     def get_template_configs(
@@ -1596,11 +1598,10 @@ def __init__(self) -> None:
         super().__init__()
         # Override mm_configs to use scaled_mm_configs
         self.mm_configs = self.scaled_mm_configs
-        # NOTE: overriding exhaustive configs here to be the same as mm_configs
-        # as we haven't validated exhaustive support here yet
-        # TODO(coconutruben): remove this once we have validated exhaustive support
-        # for scaled_mm
-        self.exhaustive_configs = self.scaled_mm_configs
+
+    def _filter_configs(self, configs: list[BaseConfig]) -> list[BaseConfig]:
+        configs = [c for c in configs if c.block_k >= 32]
+        return super()._filter_configs(configs)
 
 
 # TODO(coconutruben): replace with template.name once templates are importable
@@ -1614,11 +1615,6 @@ def __init__(self) -> None:
         super().__init__()
         # Override mm_configs to use scaled_persistent_mm_configs for TMA
         self.mm_configs = self.scaled_persistent_mm_configs
-        # NOTE: overriding exhaustive configs here to be the same as mm_configs
-        # as we haven't validated exhaustive support here yet
-        # TODO(coconutruben): remove this once we have validated exhaustive support
-        # for scaled_mm
-        self.exhaustive_configs = self.scaled_persistent_mm_configs
 
 
 # TODO(coconutruben): replace with template.name once templates are importable