[TRTLLM-14024][fix] Address nvMMH integration review findings

Keep split-K admission filtering when restoring unmatched swap orientations.
Pin the process-wide NVMMH configuration while engines remain active, reject
conflicting updates atomically, and release ownership during cleanup or GC.
Equivalent policies reuse the installed snapshot.

Cover engine configuration wiring and lifecycle, unmatched-swap admission,
and deterministic complete swizzle annotations with regression tests. Add
docstrings for changed helpers and correct the documented Blackwell BF16
field support.

Validation: 43 autotuner tests and 2 Rubin MXFP8 scheduler tests passed on
Hecate SM107. The four policy/admission cases passed again after the final
identity-preservation adjustment. The LLM args golden manifest is unchanged.

Signed-off-by: peaceh <103117813+peaceh-nv@users.noreply.github.com>

Author: peaceh-nv

docs/source/torch/adding_custom_kernels.md

@@ -268,8 +268,8 @@ This is wired up for the following dense GEMM runners. Each builds its validated
 |---|---|---|
 | SM100/SM103 | NVFP4 | `tile`, `cluster` (prune); `swizzle`, `cta_order` (annotate only) |
 | SM100/SM103 | blockwise FP8 | `tile`, `cluster` |
-| SM100/SM103 | BF16 | all five |
-| SM107 (Rubin) | block-scaled NVFP4, incl. the uGPU in-place variant | all five |
+| SM100/SM103 | BF16 | `tile`, `cluster`, `swizzle`, `cta_order` |
+| SM107 (Rubin) | block-scaled NVFP4 | all five |
 | SM107 (Rubin) | block-scaled MXFP8 | all five |
 | SM107 (Rubin) | per-tensor FP8 | `tile`, `cluster`, `cta_order` |
 | SM107 (Rubin) | BF16 | all five |
@@ -319,6 +319,12 @@ autotuner_nvmmh_config:
   max_tactics: 5  # Top-K signatures per modeled problem/orientation.
 ```
 
+All active model engines in one process must use the same effective policy,
+including the disabled default. Creating an engine with a conflicting policy
+or changing that policy through `configure_nvmmh()` raises `ValueError` without
+replacing the current configuration. Identical policies can be reused. After
+the last engine is cleaned up or garbage-collected, a new policy can be installed.
+
 Compatible validated runner-specific tactics that NVMMH cannot represent are
 retained automatically and empirically profiled. The BF16 and SM107
 block-scaled runners supplement their validated preferred/fallback-cluster

tensorrt_llm/_torch/autotuner.py

@@ -11,7 +11,9 @@
 import json
 import math
 import os
+import threading
 import time
+import weakref
 from abc import ABC, abstractmethod
 from dataclasses import dataclass, field
 from functools import lru_cache
@@ -172,6 +174,7 @@ class NvMMHConfig:
     max_tactics: int = 5
 
     def __post_init__(self) -> None:
+        """Validate policy values and normalize aliases and coupled search fields."""
         if not isinstance(self.enabled, bool):
             raise ValueError("NvMMHConfig.enabled must be a bool")
         if isinstance(self.max_tactics,
@@ -360,9 +363,11 @@ class CuteDSLTunableRunner(TunableRunner):
     """Marker base for CuTe DSL runners that prefer CUPTI timing."""
 
     def use_torch_profiler(self, tactic: Any) -> bool:
+        """Prefer CUDA activity timing for every CuTe DSL tactic."""
         return True
 
     def profiling_timer_cache_key(self) -> Optional[str]:
+        """Separate CUPTI-ranked tactics from legacy timer cache entries."""
         return _TORCH_PROFILER_TIMER_KEY
 
 
@@ -659,6 +664,7 @@ def get_cache_key(
         apply_map_to_tuning_buckets: bool = True,
         timer_key_override: Optional[str] = None,
     ) -> Tuple:
+        """Combine operation, runner, shapes, search policy, and timing backend."""
         cache_key = (
             custom_op,
             runner.__class__.__name__,
@@ -1061,12 +1067,15 @@ class AutoTuner:
     _CUDA_GRAPH_DELAY_MICRO_SECS = 100
     _instance = None
     _nvmmh_config = NvMMHConfig()
+    _nvmmh_policy_lock = threading.RLock()
+    _nvmmh_policy_owners: weakref.WeakSet = weakref.WeakSet()
     # CUPTI availability is process-wide. Keep the failure sticky across
     # AutoTuner instances so every tactic does not repeat the same failed
     # profiler initialization before falling back to CUDA events.
     _torch_profiler_unavailable = False
 
     def __init__(self, warmup=2, repeat=10, stream_delay_micro_secs=1000):
+        """Initialize profiling settings, distributed state, statistics, and caches."""
         # Increase log level for AutoTuner associated logger`
         self._log_level_to_info = os.getenv(
             "TLLM_AUTOTUNER_LOG_LEVEL_DEBUG_TO_INFO", '0') == '1'
@@ -1136,31 +1145,57 @@ def configure_nvmmh(
         is replaced, so an invalid update leaves the previous policy intact.
         Every ``AutoTuner`` instance shares this process-wide policy, which is
         intentionally persistent across ``autotune`` contexts, instance
-        construction, and profiling-cache clears.
+        construction, and profiling-cache clears. Active model engines pin
+        this policy: a differing update raises without changing any state.
         """
-        keyword_values = (
-            enabled,
-            fields,
-            max_tactics,
-        )
-        if config is not None:
-            if any(value is not None for value in keyword_values):
-                raise ValueError(
-                    "Pass either config or NVMMH keyword settings, not both")
-            if not isinstance(config, NvMMHConfig):
-                raise TypeError("config must be an NvMMHConfig")
-            new_config = config
-        else:
-            current = AutoTuner._nvmmh_config
-            new_config = NvMMHConfig(
-                enabled=current.enabled if enabled is None else enabled,
-                fields=current.fields if fields is None else fields,
-                max_tactics=(current.max_tactics
-                             if max_tactics is None else max_tactics),
+        with AutoTuner._nvmmh_policy_lock:
+            keyword_values = (
+                enabled,
+                fields,
+                max_tactics,
             )
+            if config is not None:
+                if any(value is not None for value in keyword_values):
+                    raise ValueError(
+                        "Pass either config or NVMMH keyword settings, not both"
+                    )
+                if not isinstance(config, NvMMHConfig):
+                    raise TypeError("config must be an NvMMHConfig")
+                new_config = config
+            else:
+                current = AutoTuner._nvmmh_config
+                new_config = NvMMHConfig(
+                    enabled=current.enabled if enabled is None else enabled,
+                    fields=current.fields if fields is None else fields,
+                    max_tactics=(current.max_tactics
+                                 if max_tactics is None else max_tactics),
+                )
+
+            if new_config == AutoTuner._nvmmh_config:
+                return AutoTuner._nvmmh_config
+            if AutoTuner._nvmmh_policy_owners:
+                raise ValueError(
+                    "All active model engines in one process must use the same "
+                    "autotuner_nvmmh_config; clean up existing engines before "
+                    "changing the policy.")
+            AutoTuner._nvmmh_config = new_config
+            return new_config
+
+    def _acquire_nvmmh_policy(self, owner: object, config: NvMMHConfig) -> None:
+        """Pin one policy until the owning engine is cleaned up or collected.
+
+        Weak ownership also releases failed, partially initialized engines.
+        Registering an owner and validating its policy share the configuration
+        lock, so concurrent engine construction cannot replace a pinned policy.
+        """
+        with AutoTuner._nvmmh_policy_lock:
+            self.configure_nvmmh(config)
+            AutoTuner._nvmmh_policy_owners.add(owner)
 
-        AutoTuner._nvmmh_config = new_config
-        return new_config
+    def _release_nvmmh_policy(self, owner: object) -> None:
+        """Release an engine's policy pin without resetting the current policy."""
+        with AutoTuner._nvmmh_policy_lock:
+            AutoTuner._nvmmh_policy_owners.discard(owner)
 
     class TacticsCapture:
         """Object returned by capture() that can be iterated to get all tactic combinations.
@@ -1688,6 +1723,7 @@ def _is_torch_profiler_unavailable_error(error: Exception) -> bool:
 
     @classmethod
     def _disable_torch_profiler(cls, error: Exception) -> None:
+        """Latch a process-wide CUPTI failure and report the CUDA-event fallback."""
         cls._torch_profiler_unavailable = True
         logger.warning_once(
             "[Autotuner] torch.profiler/CUPTI is unavailable; falling back "
@@ -1732,6 +1768,7 @@ def _profile_single_kernel(
         use_torch_profiler = runner.use_torch_profiler(tactic)
 
         def pure_profile(stream: torch.cuda.Stream, repeat: int):
+            """Measure repeated launches on one stream using the selected timer."""
             graph = torch.cuda.CUDAGraph() if use_cuda_graph else None
 
             with torch.cuda.stream(stream):
@@ -1751,6 +1788,7 @@ def pure_profile(stream: torch.cuda.Stream, repeat: int):
                     self._dist.tp_barrier()
 
                 def launch_profiled_work():
+                    """Replay the captured batch or launch the same batch eagerly."""
                     if use_cuda_graph:
                         graph.replay()
                     else:
@@ -1805,12 +1843,15 @@ def launch_profiled_work():
                     end_ts = torch.empty(1, dtype=torch.int64, device="cuda")
 
                     def record_start():
+                        """Record the start of the measured interval on the profiling stream."""
                         record_global_timer(start_ts.data_ptr(), stream)
 
                     def record_end():
+                        """Record the end of the measured interval on the profiling stream."""
                         record_global_timer(end_ts.data_ptr(), stream)
 
                     def elapsed_time():
+                        """Return the recorded profiling interval in milliseconds."""
                         # GPU %globaltimer counts in ns; convert to ms to match
                         # the units of torch.cuda.Event.elapsed_time().
                         return (end_ts.item() - start_ts.item()) / 1e6
@@ -1819,12 +1860,15 @@ def elapsed_time():
                     end_evt = torch.cuda.Event(enable_timing=True)
 
                     def record_start():
+                        """Record the start of the measured interval on the profiling stream."""
                         start_evt.record()
 
                     def record_end():
+                        """Record the end of the measured interval on the profiling stream."""
                         end_evt.record()
 
                     def elapsed_time():
+                        """Return the recorded profiling interval in milliseconds."""
                         return start_evt.elapsed_time(end_evt)
 
                 # Delay the profiled kernel launch to eliminate affects of host time overhead in profiling.

tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py

@@ -586,6 +586,7 @@ def unique_id(self):
             )
 
         def tactic_search_cache_key(self):
+            """Return the active NVMMH search policy discriminator for profiling caches."""
             return _cutedsl_nvmmh_tactic_search_cache_key()
 
         def __hash__(self):
@@ -603,6 +604,7 @@ def get_valid_tactics(
             **kwargs,
         ) -> List[Tuple[int, int]]:
             # Early exit: Check SM version - CuteDSL NVFP4 only supports SM 100 and SM 103
+            """Enumerate supported Blackwell NVFP4 tactics and apply optional NVMMH guidance."""
             if (sm_version := get_sm_version()) not in (100, 103):
                 logger.debug(
                     f"CuteDSL: SM version {sm_version} is not supported. "
@@ -4391,6 +4393,7 @@ def __init__(self,
             self.use_tvm_ffi = use_tvm_ffi
 
         def tactic_search_cache_key(self):
+            """Return the active NVMMH search policy discriminator for profiling caches."""
             return _cutedsl_nvmmh_tactic_search_cache_key()
 
         def get_valid_tactics(
@@ -4399,6 +4402,7 @@ def get_valid_tactics(
             profile: OptimizationProfile,
             **kwargs,
         ) -> List[int]:
+            """Enumerate blockwise FP8 tactics and retain NVMMH-ranked structural families."""
             if not is_sm_100f():
                 logger.debug(
                     f"CuteDSL: SM version {get_sm_version()} is not supported. "
@@ -9338,6 +9342,7 @@ class CuteDSLBf16BlackwellGemmRunner(CuteDSLTunableRunner):
         target_sm = "blackwell"
 
         def tactic_search_cache_key(self):
+            """Return the active NVMMH search policy discriminator for profiling caches."""
             return _cutedsl_nvmmh_tactic_search_cache_key()
 
         tuning_config = TuningConfig(dynamic_tensor_specs=(DynamicTensorSpec(
@@ -9354,7 +9359,7 @@ def get_valid_tactics(
             profile: OptimizationProfile,
             **kwargs,
         ) -> List[int]:
-
+            """Generate Blackwell BF16 tactics with optional tile and scheduler guidance."""
             if not is_sm_100f():
                 logger.debug(
                     f"CuteDSL: SM version {get_sm_version()} is not supported. "
@@ -9631,6 +9636,7 @@ def _rank_prune_tactics(self, tactics, m, n, k, output_dtype):
                     # normal full-sweep candidates.
 
                     def _split_k(tactic):
+                        """Read the tactic split factor, defaulting unsplit variants to one."""
                         if (isinstance(tactic, tuple) and len(tactic) >= 6
                                 and tactic[0] == "base"):
                             return max(1, int(tactic[5]))
@@ -9939,6 +9945,7 @@ def _sm107_bf16_gemm_kernel(kernel_variant: str,
                                 raster_along: str = "m",
                                 swizzle_size: int = 1,
                                 use_tma_store: bool = True):
+        """Compile and cache the Rubin BF16 kernel for its shape and scheduler options."""
         kernel_class = _sm107_bf16_kernel_class(kernel_variant)
         if kernel_variant == "preferred_cluster":
             return kernel_class(
@@ -10009,6 +10016,7 @@ def forward(
             inputs: List[torch.Tensor],
             tactic,
         ) -> None:
+            """Launch a Rubin batched GEMM into the supplied output tensor."""
             (kernel_variant, use_2cta_instrs, mma_tiler_mn,
              preferred_cluster_shape_mn, cluster_shape_mn, max_num_ab_stage,
              _) = _parse_sm107_bf16_tactic(tactic)
@@ -10211,6 +10219,7 @@ class CuteDSLBf16RubinGemmRunner(CuteDSLBf16BlackwellGemmRunner):
         nvmmh_split_k_cta_k = 64
 
         def tactic_search_cache_key(self):
+            """Return the active NVMMH search policy discriminator for profiling caches."""
             return _cutedsl_nvmmh_tactic_search_cache_key()
 
         # See CuteDSLBf16RubinBmmRunner.tuning_config.
@@ -10239,6 +10248,7 @@ def get_valid_tactics(
             profile: OptimizationProfile,
             **kwargs,
         ) -> List[Tuple]:
+            """Generate Rubin BF16 tactics and apply configured NVMMH search fields."""
             if not _is_sm107_cute_dsl_available():
                 logger.debug(
                     f"CuteDSL: SM version {get_sm_version()} is not supported. "
@@ -10274,6 +10284,7 @@ def forward(
             inputs: List[torch.Tensor],
             tactic,
         ) -> None:
+            """Launch the selected Rubin BF16 GEMM with its scheduler and split-K settings."""
             (kernel_variant, use_2cta_instrs, mma_tiler_mn,
              preferred_cluster_shape_mn, cluster_shape_mn, max_num_ab_stage,
              split_k_slices) = _parse_sm107_bf16_tactic(tactic)
@@ -12834,6 +12845,7 @@ def unique_id(self):
                         self.use_tvm_ffi)
 
             def tactic_search_cache_key(self):
+                """Return the active NVMMH search policy discriminator for profiling caches."""
                 if (self.nvmmh_precision is None or self.nvmmh_layout is None):
                     return None
                 return _cutedsl_nvmmh_tactic_search_cache_key()
@@ -12854,6 +12866,7 @@ def get_valid_tactics(
                 **kwargs,
             ) -> List[Tuple[int, int]]:
                 # SM107 dense block-scaled GEMM only.
+                """Enumerate valid Rubin block-scaled families before optional heuristic pruning."""
                 if (sm_version := get_sm_version()) != 107:
                     logger.debug(
                         f"CuteDSL: SM version {sm_version} is not supported. "
@@ -13258,6 +13271,7 @@ def _rank_prune_tactics(self, tactics, m, n, real_k):
                     if "split_k" in fields:
 
                         def _split_k(tactic):
+                            """Read the tactic split factor, defaulting unsplit variants to one."""
                             if (isinstance(tactic, tuple) and tactic
                                     and tactic[0] == "base"
                                     and len(tactic) >= 9):
@@ -13311,6 +13325,7 @@ def _split_k(tactic):
                         return fallback_tactics
 
                     def model_key(signature):
+                        """Drop locally swept cluster fields from a model-matching signature."""
                         return self._nvmmh_cluster_family(
                             signature,
                             match_fields,
@@ -13363,7 +13378,7 @@ def model_key(signature):
                         if annotated is not None:
                             selected.append(annotated)
                     selected.extend(
-                        tactic for tactic in fallback_tactics
+                        tactic for tactic in candidate_tactics
                         if ((signature := self._nvmmh_tactic_signature(
                             tactic, match_fields)) is not None
                             and signature[-1] in unmatched_swaps))
@@ -14453,6 +14468,7 @@ def __init__(self,
                 self.use_tvm_ffi = use_tvm_ffi
 
             def tactic_search_cache_key(self):
+                """Return the active NVMMH search policy discriminator for profiling caches."""
                 return _cutedsl_nvmmh_tactic_search_cache_key()
 
             def get_valid_tactics(
@@ -14461,7 +14477,7 @@ def get_valid_tactics(
                 profile: OptimizationProfile,
                 **kwargs,
             ) -> List[int]:
-
+                """Generate per-tensor Rubin FP8 tactics and apply optional NVMMH filtering."""
                 if (sm_version := get_sm_version()) != 107:
                     logger.debug(
                         f"CuteDSL: SM version {sm_version} is not supported. "

tensorrt_llm/_torch/custom_ops/cutedsl_matmul_heuristics.py

@@ -326,6 +326,7 @@ def rank_configs(
                 raise ValueError(f"Unsupported nvMatmulHeuristics split-K kind {split_k_kind!r}")
 
             def _method(*names):
+                """Resolve a supported method name across NVMMH Python API versions."""
                 for name in names:
                     method = getattr(interface, name, None)
                     if method is not None:
@@ -457,6 +458,7 @@ def filter_fp8_tactics(
     if sweep_cluster_n and fields & {"tile", "cluster"}:
 
         def _cluster_n_family(signature):
+            """Ignore the locally swept cluster-N dimension in the matching key."""
             cta, cluster, *remaining = signature
             return (cta, int(cluster[0]), *remaining)