[TRTLLM-14024][feat] Prune CuTe DSL GEMM autotuner tactics with nvMatmulHeuristics

docs/source/torch/adding_custom_kernels.md

@@ -258,29 +258,24 @@ The same pattern (`@torch.library.custom_op("trtllm::...")`, `register_fake`, co
 - **DLPack/CUDA Graphs.** When exporting tensors via DLPack, use the stream override that mimics the `CUDAGraphCompatibleWrapper` in `argmax.py` so the capture replays cleanly.  
 - **JIT compile cache.** Always cache compiled kernels by the keys that affect codegen (dtype, last-dim size, hardware variant). `argmax.py` and the `cute_dsl_custom_ops.py` runners are good references.
 
-### 4.5 Pruning autotuner tactics with nvMatmulHeuristics
+### 4.5 Guiding autotuner tactics with nvMatmulHeuristics
 
-A CuTe DSL runner that exposes many tactics (tile / cluster / swap / prefetch combinations) can be expensive to autotune, because the AutoTuner JIT-compiles and benchmarks every candidate returned by `get_valid_tactics`. A runner can therefore rank and prune its own candidate list with [nvMatmulHeuristics](https://pypi.org/project/nvidia-matmul-heuristics/) *inside* `get_valid_tactics`, before the tactics ever reach the compile+benchmark loop — all shapes and inputs are already visible at that point, so no extra AutoTuner hook is needed. The pruning is a no-op by default and must degrade gracefully to the full list on any failure, so tuning never loses a valid candidate.
+A CuTe DSL GEMM/BMM runner can expose hundreds of tactics (tile / cluster / split-K / rasterization / swap / prefetch combinations), and the AutoTuner JIT-compiles and benchmarks every candidate returned by `get_valid_tactics`. With `autotuner_nvmmh_config` set, nvMMH-aware runners instead build their tactic list from the top kernel configurations that [nvMatmulHeuristics](https://pypi.org/project/nvidia-matmul-heuristics/) (nvMMH) ranks for the problem:
 
-Today this is wired up for the SM100/SM103 (Blackwell B200/B300) NVFP4 dense GEMM: `CuteDSLNVFP4BlackwellRunner.get_valid_tactics` builds the full sweep and then hands it to the private `_rank_prune_tactics` helper. The pruned set is a strict, re-validated subset of the full sweep — it never introduces a tactic the kernel validator rejected — so correctness is unchanged and only autotuner wall-time is affected. On a 4096³ NVFP4 GEMM this cuts a cold autotune from ~120 tactics to ~5 (roughly 7× less wall-time) while keeping the selected kernel within a few percent of the full-sweep best.
-
-The behavior is controlled entirely by environment variables (all read at tune time; the package must be installed, otherwise the path silently falls back to the full sweep):
-
-| Environment variable | Default | Meaning |
-|---|---|---|
-| `TRTLLM_CUTEDSL_NVMMH_ENABLE` | `0` | Master opt-in switch. Set to `1` to enable heuristic tactic pruning. When unset/`0`, or when `nvidia-matmul-heuristics` is not installed, the full tactic sweep runs unchanged. |
-| `TRTLLM_CUTEDSL_NVMMH_FIELDS` | `tile,cluster` | Comma-separated list of the knobs nvMatmulHeuristics drives; any knob not listed is swept by the runner. Recognized tokens: `tile`, `cluster`, `swizzle`, `cta_order`. `tile` and `cluster` are a coupled pair (the 2-SM encoding maps the joint `(cta, cluster)`), so naming either drives both. Accepted aliases: `cta_tile` / `mma_tiler` → `tile`. Unknown tokens are warned once and ignored. **Top-K tile/cluster pruning happens only when `tile`/`cluster` is selected; with scheduler-only fields (`swizzle` and/or `cta_order`) every valid tile/cluster is swept and merely annotated with the model's scheduler knobs.** |
-| `TRTLLM_CUTEDSL_NVMMH_MAX_TACTICS` | `5` | Top-K distinct `(tile, cluster, swap)` keys to keep from the model's ranking (applies only when `tile`/`cluster` is in `FIELDS`). Both `use_prefetch` variants of a kept key are profiled on top (prefetch is not modeled), so the final tactic count can be up to ~2× this. Larger K = closer to the full-sweep best but slower autotune; smaller K = faster autotune with a larger potential perf gap. |
+```yaml
+# extra_llm_api_options.yaml
+autotuner_nvmmh_config:
+  max_tactics: 5   # nvMMH configurations to profile per GEMM shape
+```
 
-Example (enable pruning, keep the top 8 tile/cluster candidates):
+or `LLM(..., autotuner_nvmmh_config=AutoTunerNvMMHConfig(max_tactics=5))`. The option is off by default; when it is unset, when `nvidia-matmul-heuristics` is not installed, or when a query fails, runners fall back to their full tactic list.
 
-```bash
-export TRTLLM_CUTEDSL_NVMMH_ENABLE=1
-export TRTLLM_CUTEDSL_NVMMH_FIELDS=tile,cluster
-export TRTLLM_CUTEDSL_NVMMH_MAX_TACTICS=8
-```
+- **nvMMH chooses, the runner realizes.** Every field nvMMH models (CTA tile, cluster shape, split-K, rasterization order, swizzle) is taken as returned. Only knobs nvMMH does not model (prefetch, scheduler mode, the preferred/fallback cluster launch) are swept inside each configuration. If a chosen configuration performs poorly, the fix belongs in nvMMH's model, not in TensorRT-LLM.
+- **Kernels describe their capabilities.** Each runner passes nvMMH a validity check (built on the kernel's `can_implement`) and backend properties (split-K kind, column-major rasterization, fallback clusters), so nvMMH ranks implementable configurations only. Runners that support both A/B orientations query both problems and merge the rankings by nvMMH's ranking score.
+- **Cache keys.** Runners return `("nvmmh", <library version>, max_tactics)` from `TunableRunner.tactic_space_key()`, so profiling-cache entries tuned with and without the guidance are never mixed.
+- **Library versions.** Both the `nvidia-matmul-heuristics` 0.1.0.x wheels and newer libheuristics builds are supported; the adapter in [`cutedsl_matmul_heuristics.py`](https://github.com/NVIDIA/TensorRT-LLM/blob/main/tensorrt_llm/_torch/custom_ops/cutedsl_matmul_heuristics.py) hides their API differences.
 
-The adapter and env helpers live in [`tensorrt_llm/_torch/custom_ops/cutedsl_matmul_heuristics.py`](https://github.com/NVIDIA/TensorRT-LLM/blob/main/tensorrt_llm/_torch/custom_ops/cutedsl_matmul_heuristics.py); the `get_valid_tactics` / `_rank_prune_tactics` pruning is in [`cute_dsl_custom_ops.py`](https://github.com/NVIDIA/TensorRT-LLM/blob/main/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py).
+To guide a new runner, return `_nvmmh_tactic_space_key()` from `tactic_space_key`, and in `get_valid_tactics` (when `AutoTuner.get().nvmmh_config` is set) call `cutedsl_matmul_heuristics.query(...)` with a validity check that maps a configuration onto the kernel, expand each returned configuration into its tactics, and keep the first `max_tactics` configurations with `select_tactics`. The Rubin block-scaled runner (`CuteDSLBlockScaledRubinLinear`) in [`cute_dsl_custom_ops.py`](https://github.com/NVIDIA/TensorRT-LLM/blob/main/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py) is a complete example.
 
 ---
 

tensorrt_llm/_torch/autotuner.py

@@ -157,6 +157,27 @@ class TuningConfig:
     exclude_from_cache: bool = False
 
 
+@dataclass(frozen=True)
+class NvMMHConfig:
+    """nvMatmulHeuristics (nvMMH) guidance for CuTe DSL GEMM/BMM runners.
+
+    When installed on the AutoTuner, nvMMH-aware runners build their tactic
+    list from the top ``max_tactics`` kernel configurations nvMMH ranks for
+    the problem instead of enumerating their full tactic space. Every field
+    nvMMH models (CTA tile, cluster shape, split-K, rasterization order,
+    swizzle) is taken as returned; only knobs nvMMH does not model are swept.
+    """
+    max_tactics: int = 5
+
+    def __post_init__(self):
+        if (isinstance(self.max_tactics, bool)
+                or not isinstance(self.max_tactics, int)
+                or self.max_tactics < 1):
+            raise ValueError(
+                f"NvMMHConfig.max_tactics must be a positive integer, got "
+                f"{self.max_tactics!r}")
+
+
 @dataclass(unsafe_hash=True)
 class StaticDim:
     val: int
@@ -276,6 +297,16 @@ def unique_id(self):
         """
         return tuple(self.__dict__.values())
 
+    def tactic_space_key(self) -> Optional[Tuple]:
+        """Discriminates cache entries profiled over a non-default tactic space.
+
+        Runners whose candidate tactics depend on a process-wide search policy
+        (e.g. nvMatmulHeuristics guidance) return a literal tuple describing
+        that policy, so entries tuned under one policy are not reused under
+        another. None keeps the default cache key.
+        """
+        return None
+
 
 @contextlib.contextmanager
 def autotune(tune_mode: bool = True,
@@ -541,7 +572,7 @@ def get_cache_key(
         tuning_config: TuningConfig,
         apply_map_to_tuning_buckets: bool = True,
     ) -> Tuple:
-        return (
+        cache_key = (
             custom_op,
             runner.__class__.__name__,
             str(runner.unique_id()),
@@ -553,6 +584,10 @@ def get_cache_key(
                 apply_map_to_tuning_buckets,
             ),
         )
+        tactic_space_key = runner.tactic_space_key()
+        if tactic_space_key is not None:
+            cache_key += (tactic_space_key, )
+        return cache_key
 
     def merge_cache_data(self, cache_data: Dict[Tuple, Tuple]):
         self.cache.update(cache_data)
@@ -946,6 +981,9 @@ def __init__(self, warmup=2, repeat=10, stream_delay_micro_secs=1000):
         self._primed_cached_tactics: set[tuple[str, str, str, str]] = set()
         self.is_tuning_mode = False
         self.skip_dynamic_tuning_buckets = False
+        # nvMatmulHeuristics guidance for CuTe DSL GEMM/BMM runners; None
+        # profiles their full tactic space.
+        self.nvmmh_config: Optional[NvMMHConfig] = None
 
         # Timing backend: globaltimer kernel vs cuda events.
         # TLLM_PROFILING_TIMER env var overrides auto-detection: