perf: cache the bf16 x fp4 runner factories (cute-dsl + cudnn)

flashinfer/gemm/gemm_bf16_fp4_cudnn.py

@@ -355,11 +355,9 @@ def execute_cudnn_bf16_fp4_graph_override_shape(
 )
 
 
-def _cudnn_bf16_fp4_runner(tuning_config):
-    """Build a ``CudnnBf16Fp4Runner`` bound to the active tuning config."""
-    m_bucket_mapper = AutoTuner.get().get_effective_map_to_tuning_buckets(
-        tuning_config, spec_idx=0
-    )
+@functools.lru_cache(maxsize=1024)
+def _cudnn_bf16_fp4_runner_for_mapper(m_bucket_mapper):
+    """Build a ``CudnnBf16Fp4Runner`` bound to the given M-bucket mapper."""
 
     class CudnnBf16Fp4Runner(TunableRunner):
         def __init__(self):
@@ -487,6 +485,21 @@ def forward(
     return CudnnBf16Fp4Runner()
 
 
+def _cudnn_bf16_fp4_runner(tuning_config):
+    """Return the runner for the currently effective M-bucket mapper.
+
+    Cached per effective mapper rather than per ``tuning_config``: the mapper
+    reflects any thread-local ``autotune(tuning_buckets=..., round_up=...)``
+    override, and a runner built under one bucket scheme must not be reused
+    under another (its override-shape graphs would be keyed on stale
+    ``cache_m`` buckets).
+    """
+    m_bucket_mapper = AutoTuner.get().get_effective_map_to_tuning_buckets(
+        tuning_config, spec_idx=0
+    )
+    return _cudnn_bf16_fp4_runner_for_mapper(m_bucket_mapper)
+
+
 def _prepare_cudnn(
     b: torch.Tensor,
     b_descale: torch.Tensor,

flashinfer/gemm/gemm_bf16_fp4_cute_dsl.py

@@ -2,6 +2,7 @@
 # SPDX-License-Identifier: Apache-2.0
 """CuTe-DSL backend for the bf16 x fp4 GEMM (weight repack / kernel launch)."""
 
+import functools
 from typing import List, Optional, Tuple, cast
 
 import torch
@@ -365,6 +366,7 @@ def add(tile_m, atom, pdepth, fp16, tile_n=64, tk=None, swz=1):
 )
 
 
+@functools.lru_cache(maxsize=1024)
 def _cute_dsl_bf16_fp4_runner(enable_pdl: bool = True) -> TunableRunner:
     """Build a ``CuteDslBf16Fp4Runner`` for the cute-DSL bf16 x fp4 GEMM."""