From 2497bdacb6bfad6fb04ddb7e22b82269cff5f5ce Mon Sep 17 00:00:00 2001 From: peaceh <103117813+peaceh-nv@users.noreply.github.com> Date: Wed, 16 Sep 2026 22:26:34 -0700 Subject: [PATCH 1/8] [None][feat] Integrate nvMMH guidance for CuTe DSL GEMM autotuning Configure optional nvMMH tactic pruning and scheduler guidance for CuTe DSL BF16, FP8, NVFP4 and MXFP8 runners. Separate search policies and profiling timers in cache identity, prefer CUPTI timing and fall back to CUDA events. Expose the policy through LLM args and layerwise YAML configuration. Use main's separate Rubin BF16 runners, direct split-K and public DSL gates. Preserve current dispatch, telemetry policies and validated baseline tactic families. Use in-process profiling. Validation after decoupling: 61 checks passed on Hecate SM107 (39 autotuner, 12 focused NVMMH, 10 Rubin kernel tests). Golden manifest regenerated and touched-file pre-commit passed. CUPTI is unavailable in this setup; CUDA-event fallback is verified. Reused recent native binaries, not an exact-tip rebuild. Signed-off-by: peaceh <103117813+peaceh-nv@users.noreply.github.com> --- docs/source/torch/adding_custom_kernels.md | 95 +- examples/layer_wise_benchmarks/README.md | 57 ++ .../layer_wise_benchmarks/config_ctx.yaml | 4 + .../layer_wise_benchmarks/config_gen.yaml | 4 + examples/layer_wise_benchmarks/run.py | 18 +- tensorrt_llm/_torch/autotuner.py | 427 ++++++-- .../_torch/custom_ops/cute_dsl_custom_ops.py | 933 ++++++++++++++++-- .../custom_ops/cutedsl_matmul_heuristics.py | 568 +++++++++-- .../_torch/custom_ops/torch_custom_ops.py | 16 +- .../rubin/dense_bf16_gemm_persistent.py | 6 + .../_torch/locality_domain/autotune.py | 4 + .../moe/fused_moe/fused_moe_cute_dsl.py | 9 +- .../_torch/pyexecutor/model_engine.py | 19 + tensorrt_llm/llmapi/__init__.py | 6 +- tensorrt_llm/llmapi/llm_args.py | 34 + .../usage/llm_args_golden_manifest.json | 17 + .../parallel/test_fp8_block_scale_gemm.py | 57 +- tests/unittest/_torch/misc/test_autotuner.py | 724 ++++++++++++-- .../api_stability/references/llm.yaml | 4 + 19 files changed, 2679 insertions(+), 323 deletions(-) diff --git a/docs/source/torch/adding_custom_kernels.md b/docs/source/torch/adding_custom_kernels.md index 97aaa3ffac94..a64dde92bcfc 100644 --- a/docs/source/torch/adding_custom_kernels.md +++ b/docs/source/torch/adding_custom_kernels.md @@ -260,27 +260,96 @@ The same pattern (`@torch.library.custom_op("trtllm::...")`, `register_fake`, co ### 4.5 Pruning 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 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 or empty match, so a heuristics failure never *reduces* the tactic set. Successful pruning deliberately does drop valid candidates — that is the point — so it can exclude the full-sweep winner and change the selected tactic and steady-state performance. -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. +This is wired up for the following dense GEMM runners. Each builds its validated sweep and then hands it to a private `_rank_prune_tactics` helper, which honors only the intersection of the configured `fields` with the runner's own modeled fields: -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): +| Arch | Precision | Modeled fields | +|---|---|---| +| 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 | +| SM107 (Rubin) | block-scaled MXFP8 | all five | +| SM107 (Rubin) | per-tensor FP8 | `tile`, `cluster`, `cta_order` | +| SM107 (Rubin) | BF16 | all five | + +Batched (BMM) runners, grouped/MoE runners, and the SM107 blockwise FP8 GEMM are not pruned. A runner whose modeled fields do not intersect the configured `fields` skips pruning entirely and runs its full sweep. Do not confuse the unsupported SM107 blockwise FP8 runner with the block-scaled MXFP8 and NVFP4 dense runners in the table above. + +The SM107 MXFP8 adapter uses NVMMH's closest per-tensor E4M3 compute model because NVMMH does not expose the UE8M0 block-scale format. CuTe DSL validity checks still enforce the real MXFP8 scale layout before any modeled tactic is retained. -| Environment variable | Default | Meaning | +The pruned set preserves the runner's validated tile/cluster templates and falls back to the full sweep when the heuristic cannot produce a usable match. On a 4096³ NVFP4 GEMM with the default `fields=("tile", "cluster")` and `max_tactics=5`, this cuts a cold autotune from ~120 tactics to ~5 ranked tile/cluster keys — roughly 10 tactics once the unmodeled prefetch variants are added back — while keeping the selected kernel within a few percent of the full-sweep best. Scheduler-only fields behave differently: on SM100/SM103 NVFP4 every validated tile/cluster is retained and merely annotated, so the tactic count is unchanged. + +The behavior is configured atomically on the process-wide `AutoTuner`. The +general library default is disabled; when requested but +`nvidia-matmul-heuristics` is unavailable, tactic generation falls back to the +full sweep. + +| `NvMMHConfig` field | 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. | +| `enabled` | `False` | Opt in to heuristic tactic pruning. | +| `fields` | `("tile", "cluster")` | Knobs selected for nvMatmulHeuristics-guided pruning: `tile`, `cluster`, `swizzle`, `cta_order`, and `split_k`. Knobs not listed remain swept. `tile` and `cluster` are coupled. Direct Python calls also accept `cta_tile` / `mma_tiler` as aliases for `tile`; strict LLM API YAML accepts canonical names only. Unknown fields raise `ValueError`. Scheduler-only fields preserve every valid tile/cluster on the runners that model a scheduler knob. `split_k` is actionable on the SM107 BF16 GEMM runner and the SM107 block-scaled NVFP4/MXFP8 base kernels. On SM107 it enables the bounded local split sweep described below; mixed-CGA block-scaled tactics do not implement split-K and remain at 1. | +| `max_tactics` | `5` | Top-K distinct model-visible signatures per modeled problem/orientation over the selected `fields` — tile/cluster when those are selected, otherwise the scheduler knobs. A block-scaled runner can query native and swapped M/N orientations separately. Runner-local variants that the model cannot represent (prefetch, MMA instruction shape, cluster-N, scheduler family) ride along with each retained signature, so the final tactic count is usually larger. Must be at least 1. | + +Every `configure_nvmmh` keyword defaults to `None`, which *preserves* the currently +installed value rather than resetting it; the defaults above are the +`NvMMHConfig` dataclass defaults. Pass a complete `NvMMHConfig` to set all +fields at once. + +Example (enable pruning and keep the top 8 tile/cluster candidates): + +```python +from tensorrt_llm._torch.autotuner import AutoTuner + +AutoTuner.get().configure_nvmmh( + enabled=True, + fields=("tile", "cluster"), + max_tactics=8, +) +``` -Example (enable pruning, keep the top 8 tile/cluster candidates): +For an LLM API process, configure the same immutable policy through +`extra_llm_api_options.yaml`; each executor worker installs it before model +loading and tactic enumeration: -```bash -export TRTLLM_CUTEDSL_NVMMH_ENABLE=1 -export TRTLLM_CUTEDSL_NVMMH_FIELDS=tile,cluster -export TRTLLM_CUTEDSL_NVMMH_MAX_TACTICS=8 +```yaml +enable_autotuner: true +autotuner_nvmmh_config: + fields: [swizzle, cta_order, split_k] + max_tactics: 5 # Top-K signatures per modeled problem/orientation. ``` -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). +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 +family because NVMMH has only one cluster field and cannot encode both shapes. +Selected `swizzle` and `cta_order` guidance still applies to the block-scaled +supplement. The current public NVMMH API does not expose its native +`minTilesKLowerBound`, so selecting `split_k` on SM107 uses a shared TRT-LLM +admission rule: split 1 is always retained, and larger factors must leave +strictly more than 6 CTA-K tiles per split. The AutoTuner profiles every +admitted, runner-validated factor instead of pinning the set to NVMMH's native +rank-1 split. The admitted set remains a subset of the runner's normal full sweep. +The CTA-K is 256 for NVFP4, 128 for MXFP8, and 64 for BF16. This SM107 lower +bound is code-defined and is not user-configurable. + +Pass the file with `trtllm-serve --backend pytorch +--extra_llm_api_options /path/to/extra_llm_api_options.yaml` (or the equivalent +`--config` alias). Omitting the nested block leaves NVMMH disabled and retains +the full tactic sweep; setting it to `null` does the same. Any mapping, including +`autotuner_nvmmh_config: {}`, enables NVMMH. YAML `fields` must use only `tile`, +`cluster`, `swizzle`, `cta_order`, and `split_k`; aliases and alternate split-K +spellings are rejected so configuration typos fail early. + +The immutable configuration remains installed across `autotune()` contexts +and `clear_cache()` calls. Profiling-cache keys include the complete nvMMH +search policy, so full sweep, scheduler-only, and all-fields results cannot +alias. `AutoTuner.get().configure_nvmmh(...)` is the recommended public entry +point; the installed policy is class-wide, so every `AutoTuner` instance and +CuTe DSL runner observes the same snapshot. The adapter lives 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). --- diff --git a/examples/layer_wise_benchmarks/README.md b/examples/layer_wise_benchmarks/README.md index 9fecec99a595..f4e45ee93671 100644 --- a/examples/layer_wise_benchmarks/README.md +++ b/examples/layer_wise_benchmarks/README.md @@ -136,6 +136,63 @@ NP=4 ./mpi_launch.sh ./run.sh config_ctx.yaml --balance-method ImbalancedExperts NP=4 ./mpi_launch.sh ./run.sh config_gen.yaml --balance-method ImbalancedExperts --balance-ratio 0.5 ``` +### CuTe DSL nvMMH defaults + +The shipped `config_ctx.yaml` and `config_gen.yaml` enable nvMMH in +scheduler-only mode. nvMMH selects `swizzle`, `cta_order`, and `split_k`; the +autotuner continues to profile the full tile and cluster search space. The +layer-wise nvMMH policy is YAML-only so the effective search policy stays with +the benchmark configuration. + +```yaml +# config_ctx.yaml / config_gen.yaml: scheduler-only nvMMH +enable_autotuner: true +autotuner_nvmmh: true +autotuner_nvmmh_fields: [swizzle, cta_order, split_k] +autotuner_nvmmh_max_tactics: 5 + +# To use the full autotuner sweep instead: +# autotuner_nvmmh: false + +# To let nvMMH select every modeled field instead: +# autotuner_nvmmh_fields: [tile, cluster, swizzle, cta_order, split_k] +``` + +The YAML keys are `autotuner_nvmmh`, `autotuner_nvmmh_fields`, and +`autotuner_nvmmh_max_tactics`. The corresponding `--autotuner-nvmmh*` +command-line options are intentionally not accepted. A custom layer-wise YAML +that omits the keys preserves the general library default, where nvMMH is +disabled. + +Validated runner-specific tactics that nvMMH cannot represent are retained +automatically. BF16 and Rubin block-scaled runners preserve their validated +preferred/fallback-cluster family. On SM107, split-capable runners admit a +bounded local set of `split_k` values using TRT-LLM's shared K-tile lower bound, +then let CUPTI select among them; this avoids depending on nvMMH's fixed native +`minTilesKLowerBound`. `autotuner_nvmmh_max_tactics` caps the top-K signatures +for the selected fields per modeled problem/orientation; it is not limited to +tile/cluster. + +These flat keys belong to the layer-wise benchmark's own `config_ctx.yaml` / +`config_gen.yaml` parser. `run.py` has no `--extra_llm_api_options` flag; the +typed nested form below applies to PyTorch LLM API files consumed by +`trtllm-serve` / `trtllm-bench` via `--extra_llm_api_options` or `--config`: + +```yaml +enable_autotuner: true +autotuner_nvmmh_config: + fields: [swizzle, cta_order, split_k] + max_tactics: 5 # Top-K signatures per modeled problem/orientation. +``` + +For this nested LLM API form, any mapping enables nvMMH; omission or `null` +disables it. This differs from the layer-wise benchmark's flat +`autotuner_nvmmh` switch above. Both interfaces preserve the library's disabled +default when their nvMMH keys are omitted; the shipped layer-wise YAML files +explicitly opt in to scheduler-only mode. Moving that policy into +`extra_llm_api_options.yaml` therefore requires adding the nested block +explicitly, or nvMMH is off. + ### Run with Slurm > Tips: diff --git a/examples/layer_wise_benchmarks/config_ctx.yaml b/examples/layer_wise_benchmarks/config_ctx.yaml index 21d36b95234f..f2ce0274b97f 100644 --- a/examples/layer_wise_benchmarks/config_ctx.yaml +++ b/examples/layer_wise_benchmarks/config_ctx.yaml @@ -9,6 +9,10 @@ enable_attention_dp: true # Model init args moe_backend: CUTLASS +enable_autotuner: true +autotuner_nvmmh: true +autotuner_nvmmh_fields: [swizzle, cta_order, split_k] +autotuner_nvmmh_max_tactics: 5 use_cuda_graph: false # Per iteration args diff --git a/examples/layer_wise_benchmarks/config_gen.yaml b/examples/layer_wise_benchmarks/config_gen.yaml index 6384cd0ea222..2b52c8e5887e 100644 --- a/examples/layer_wise_benchmarks/config_gen.yaml +++ b/examples/layer_wise_benchmarks/config_gen.yaml @@ -9,6 +9,10 @@ enable_attention_dp: true # Model init args moe_backend: CUTLASS +enable_autotuner: true +autotuner_nvmmh: true +autotuner_nvmmh_fields: [swizzle, cta_order, split_k] +autotuner_nvmmh_max_tactics: 5 use_cuda_graph: true # Per iteration args diff --git a/examples/layer_wise_benchmarks/run.py b/examples/layer_wise_benchmarks/run.py index eb712bd66164..9d71b76ec55f 100644 --- a/examples/layer_wise_benchmarks/run.py +++ b/examples/layer_wise_benchmarks/run.py @@ -29,6 +29,10 @@ def comma_separated_floats(s): return [float(x) for x in s.split(",")] +def comma_separated_strings(s): + return [value.strip() for value in s.split(",") if value.strip()] + + # Parse cmdline parser = argparse.ArgumentParser() parser.add_argument("config_path", type=str) @@ -93,6 +97,10 @@ def comma_separated_floats(s): group.add_argument("--enable-autotuner", action="store_true", dest="enable_autotuner") group.add_argument("--no-enable-autotuner", action="store_false", dest="enable_autotuner") parser.set_defaults(enable_autotuner=None) +# NVMMH policy is configured through YAML. +parser.set_defaults( + autotuner_nvmmh=None, autotuner_nvmmh_fields=None, autotuner_nvmmh_max_tactics=None +) group = parser.add_mutually_exclusive_group() group.add_argument("--use-cuda-graph", action="store_true", dest="use_cuda_graph") group.add_argument("--no-use-cuda-graph", action="store_false", dest="use_cuda_graph") @@ -169,6 +177,8 @@ def comma_separated_floats(s): args.use_low_precision_moe_combine = False if args.enable_autotuner is None: args.enable_autotuner = True +if isinstance(args.autotuner_nvmmh_fields, str): + args.autotuner_nvmmh_fields = comma_separated_strings(args.autotuner_nvmmh_fields) if args.use_cuda_graph is None: args.use_cuda_graph = False if (args.replay_start_iter is None) != (args.replay_stop_iter is None): @@ -224,7 +234,13 @@ def comma_separated_floats(s): logger.info("Layer-wise benchmarks: Create KV cache manager ... Done") # Create other global objects -AutoTuner.get().clear_cache() +tuner = AutoTuner.get() +tuner.configure_nvmmh( + enabled=args.autotuner_nvmmh, + fields=args.autotuner_nvmmh_fields, + max_tactics=args.autotuner_nvmmh_max_tactics, +) +tuner.clear_cache() capture_stream = torch.cuda.Stream() mark_ranges() diff --git a/tensorrt_llm/_torch/autotuner.py b/tensorrt_llm/_torch/autotuner.py index b475caa0e368..87da00ea2c8f 100644 --- a/tensorrt_llm/_torch/autotuner.py +++ b/tensorrt_llm/_torch/autotuner.py @@ -20,13 +20,15 @@ import inspect import itertools import json +import math import os import time from abc import ABC, abstractmethod from dataclasses import dataclass, field from functools import lru_cache from pathlib import Path -from typing import Any, Callable, Dict, List, Optional, Set, Tuple, Union +from typing import (Any, Callable, Dict, FrozenSet, Iterable, List, Optional, + Set, Tuple, Union) import torch from cuda.bindings import driver @@ -41,6 +43,12 @@ # Unique tag to avoid collisions with other comms PP_COMM_TAG_AUTOTUNING = 30000 +_TORCH_PROFILER_TIMER_KEY = "torch_profiler" +_CUDA_EVENT_FALLBACK_TIMER_KEY = "cuda_event_fallback" + + +class _TorchProfilerUnavailableError(RuntimeError): + """Raised when Kineto/CUPTI cannot provide CUDA timing activities.""" class DistributedTuningStrategy(enum.Enum): @@ -152,6 +160,71 @@ class TuningConfig: exclude_from_cache: bool = False +_NVMMH_CANONICAL_FIELDS = frozenset( + {"tile", "cluster", "swizzle", "cta_order", "split_k"}) +_NVMMH_FIELD_ALIASES = { + "cta_tile": "tile", + "mma_tiler": "tile", +} + + +@dataclass(frozen=True, slots=True) +class NvMMHConfig: + """Immutable nvMatmulHeuristics tactic-search policy. + + ``fields`` accepts either an iterable of field names or a comma-separated + string. Field names are normalized once at construction so every consumer + observes one complete, validated policy snapshot. ``tile`` and ``cluster`` + are coupled because nvMatmulHeuristics models their joint configuration. + """ + + enabled: bool = False + fields: FrozenSet[str] = frozenset({"tile", "cluster"}) + max_tactics: int = 5 + + def __post_init__(self) -> None: + if not isinstance(self.enabled, bool): + raise ValueError("NvMMHConfig.enabled must be a bool") + if isinstance(self.max_tactics, + bool) or not isinstance(self.max_tactics, int): + raise ValueError("NvMMHConfig.max_tactics must be an integer") + if self.max_tactics < 1: + raise ValueError("NvMMHConfig.max_tactics must be at least 1") + + fields = self.fields + if isinstance(fields, str): + raw_fields: Iterable[str] = fields.split(",") + else: + try: + raw_fields = iter(fields) + except TypeError as error: + raise ValueError( + "NvMMHConfig.fields must be an iterable of strings or a " + "comma-separated string") from error + + normalized_fields = set() + for field_name in raw_fields: + if not isinstance(field_name, str): + raise ValueError( + "NvMMHConfig.fields entries must all be strings") + field_name = field_name.strip().lower() + if not field_name: + continue + normalized_fields.add( + _NVMMH_FIELD_ALIASES.get(field_name, field_name)) + + unknown_fields = normalized_fields - _NVMMH_CANONICAL_FIELDS + if unknown_fields: + raise ValueError("Unknown NvMMHConfig fields: " + f"{sorted(unknown_fields)}; expected a subset of " + f"{sorted(_NVMMH_CANONICAL_FIELDS)}") + + # The 2-SM encoding models the joint CTA tile and cluster shape. + if normalized_fields & {"tile", "cluster"}: + normalized_fields.update({"tile", "cluster"}) + object.__setattr__(self, "fields", frozenset(normalized_fields)) + + @dataclass(unsafe_hash=True) class StaticDim: val: int @@ -271,6 +344,38 @@ def unique_id(self): """ return tuple(self.__dict__.values()) + def use_torch_profiler(self, tactic: Any) -> bool: + """Whether this tactic should prefer Kineto/CUPTI timing. + + The default keeps the existing globaltimer/CUDA-event path. Runners + that prefer CUPTI timing override this per-tactic hook; AutoTuner + falls back to CUDA events when CUPTI is unavailable. + """ + return False + + def profiling_timer_cache_key(self) -> Optional[str]: + """Return a cache discriminator when timing changes tactic ranking.""" + return None + + def tactic_search_cache_key(self) -> Optional[Tuple[Any, ...]]: + """Return a cache discriminator when tactic generation changes. + + Runners whose candidate set depends on process-wide search policy + should override this hook with a stable, versioned tuple. The search + discriminator is composed before the profiling-timer discriminator. + """ + return None + + +class CuteDSLTunableRunner(TunableRunner): + """Marker base for CuTe DSL runners that prefer CUPTI timing.""" + + def use_torch_profiler(self, tactic: Any) -> bool: + return True + + def profiling_timer_cache_key(self) -> Optional[str]: + return _TORCH_PROFILER_TIMER_KEY + @contextlib.contextmanager def autotune(tune_mode: bool = True, @@ -503,6 +608,7 @@ def search_cache( input_shapes: Tuple[torch.Size], tuning_config: TuningConfig, apply_map_to_tuning_buckets: bool = True, + allow_timer_fallback: bool = False, ) -> Tuple[bool, int, int, Dict[str, Any], OptimizationProfile]: """Search for cached profiling results matching the current configuration. @@ -519,12 +625,39 @@ def search_cache( runner_id is the index in the current runners list """ for idx, r in enumerate(runners): - if (cache_key := self.get_cache_key( - custom_op, r, input_shapes, tuning_config, - apply_map_to_tuning_buckets)) in self.cache: - # Return the current index in runners list, not the cached runner_id - cached_runner_id, tactic, min_time = self.cache[cache_key] - return True, idx, tactic, min_time + cache_keys = [ + self.get_cache_key( + custom_op, + r, + input_shapes, + tuning_config, + apply_map_to_tuning_buckets, + ) + ] + # A fresh inference process has not probed CUPTI yet. It may reuse + # an explicitly tagged CUDA-event fallback entry when no preferred + # CUPTI entry exists. Tuning mode leaves this disabled so a + # CUPTI-capable process re-profiles instead of reusing event-ranked + # tactics. + if (allow_timer_fallback + and not AutoTuner._torch_profiler_unavailable and + r.profiling_timer_cache_key() == _TORCH_PROFILER_TIMER_KEY): + cache_keys.append( + self.get_cache_key( + custom_op, + r, + input_shapes, + tuning_config, + apply_map_to_tuning_buckets, + timer_key_override=_CUDA_EVENT_FALLBACK_TIMER_KEY, + )) + + for cache_key in cache_keys: + if cache_key in self.cache: + # Return the current index in runners list, not the cached + # runner_id. + cached_runner_id, tactic, min_time = self.cache[cache_key] + return True, idx, tactic, min_time return False, *self.fallback_entry() @@ -535,8 +668,9 @@ def get_cache_key( input_shapes: Tuple[torch.Size], tuning_config: TuningConfig, apply_map_to_tuning_buckets: bool = True, + timer_key_override: Optional[str] = None, ) -> Tuple: - return ( + cache_key = ( custom_op, runner.__class__.__name__, str(runner.unique_id()), @@ -548,6 +682,16 @@ def get_cache_key( apply_map_to_tuning_buckets, ), ) + tactic_search_key = runner.tactic_search_cache_key() + if tactic_search_key is not None: + cache_key += (tactic_search_key, ) + timer_key = (runner.profiling_timer_cache_key() + if timer_key_override is None else timer_key_override) + if (timer_key_override is None + and timer_key == _TORCH_PROFILER_TIMER_KEY + and AutoTuner._torch_profiler_unavailable): + timer_key = _CUDA_EVENT_FALLBACK_TIMER_KEY + return cache_key if timer_key is None else cache_key + (timer_key, ) def merge_cache_data(self, cache_data: Dict[Tuple, Tuple]): self.cache.update(cache_data) @@ -927,6 +1071,11 @@ class AutoTuner: """ _CUDA_GRAPH_DELAY_MICRO_SECS = 100 _instance = None + _nvmmh_config = NvMMHConfig() + # 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): # Increase log level for AutoTuner associated logger` @@ -939,10 +1088,13 @@ def __init__(self, warmup=2, repeat=10, stream_delay_micro_secs=1000): self.stream_delay_micro_secs = stream_delay_micro_secs self.profiling_cache = AutoTunerProfilingCache() self._primed_cached_tactics: set[tuple[str, str, str, str]] = set() + self._last_profile_timer: Optional[str] = None self.is_tuning_mode = False self.skip_dynamic_tuning_buckets = False - # Timing backend: globaltimer kernel vs cuda events. + # Default timing backend: globaltimer kernel vs CUDA events. CuTe DSL + # runners prefer torch.profiler/CUPTI and fall back to CUDA events when + # CUPTI is unavailable in the current process. # TLLM_PROFILING_TIMER env var overrides auto-detection: # "globaltimer" -> force globaltimer # "cuda_event" -> force cuda events @@ -976,6 +1128,52 @@ def get(cls): cls._instance = AutoTuner() return cls._instance + @property + def nvmmh_config(self) -> NvMMHConfig: + """Return the current immutable NVMMH policy snapshot.""" + return AutoTuner._nvmmh_config + + def configure_nvmmh( + self, + config: Optional[NvMMHConfig] = None, + *, + enabled: Optional[bool] = None, + fields: Optional[Union[str, Iterable[str]]] = None, + max_tactics: Optional[int] = None, + ) -> NvMMHConfig: + """Atomically install a validated NVMMH tactic-search policy. + + Pass either a complete ``NvMMHConfig`` or keyword updates to the + current policy. Validation happens before the single config reference + 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. + """ + 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), + ) + + AutoTuner._nvmmh_config = new_config + return new_config + class TacticsCapture: """Object returned by capture() that can be iterated to get all tactic combinations. @@ -1026,7 +1224,7 @@ def __iter__(self): yield tuple(runner_tactic_pairs) def _generate_context_tactics_lists(self): - """Generate all valid tactic combinations.""" + """Generate all runner-valid tactic combinations.""" if not self._captured_contexts: raise RuntimeError( "No context available for testing.\n" @@ -1042,11 +1240,14 @@ def _generate_context_tactics_lists(self): inputs = context['inputs'] kwargs = context.get('kwargs', {}) - # Collect all valid (runner, tactic) combinations for this context + profile = OptimizationProfile() + + # Collect the same runner-valid tactic combinations considered + # by normal autotuning for this context. tactics_lists = [] for runner_idx, runner in enumerate(runners): valid_tactics = runner.get_valid_tactics( - inputs, OptimizationProfile(), **kwargs) + inputs, profile, **kwargs) for tactic in valid_tactics: tactics_lists.append((runner_idx, tactic)) context_tactics_lists.append(tactics_lists) @@ -1161,7 +1362,9 @@ def choose_one( runners, input_shapes, tuning_config, - apply_map_to_tuning_buckets=True)) + apply_map_to_tuning_buckets=True, + allow_timer_fallback=not self.is_tuning_mode, + )) # Early return if it's not tuning, use cache found one or fallback one if not self.is_tuning_mode: @@ -1248,7 +1451,14 @@ def choose_one( # Get the best runner and tactic from cache # If no valid tactic is found, the fallback runner and tactic will be used _, runner_id, tactic, _ = self.profiling_cache.search_cache( - custom_op, runners, input_shapes, tuning_config) + custom_op, + runners, + input_shapes, + tuning_config, + # Non-profiling distributed ranks receive the selected cache entry + # but do not inherit the profiling process's CUPTI failure latch. + allow_timer_fallback=True, + ) tuning_end_time = time.perf_counter() self.stats.tuned_op_time_cost[ @@ -1325,6 +1535,7 @@ def _profile_runners( input_tensors: List[torch.Tensor], profile: OptimizationProfile, tuning_config: TuningConfig, + _inputs_prepared: bool = False, **kwargs, ) -> float: """Profile runners and select the best tactic. @@ -1337,8 +1548,9 @@ def _profile_runners( min_time = float('inf') has_tuning_failure_occurred = False best_runner_id, best_tactic = None, None + torch_profiler_was_available = not AutoTuner._torch_profiler_unavailable # If the inputs_pre_hook is provided, it will be called before profiling. - if tuning_config.inputs_pre_hook is not None: + if tuning_config.inputs_pre_hook is not None and not _inputs_prepared: input_tensors = tuning_config.inputs_pre_hook(input_tensors) # Pre-walk the runners so we know the total candidate (runner, tactic) @@ -1478,6 +1690,32 @@ def _profile_runners( min_time = time_measured best_runner_id, best_tactic = runner_id, tac + if (torch_profiler_was_available + and AutoTuner._torch_profiler_unavailable): + # At least one timing used the CUDA-event fallback after this sweep + # began. Discard every measurement from the mixed-backend sweep and + # re-profile the complete candidate set. The process-wide latch + # forces every CUPTI-preferring candidate onto CUDA events during + # the retry. + self._debug_logger( + f"[Autotuner] Re-profiling all tactics for custom_op={custom_op} " + "with CUDA events after CUPTI became unavailable.") + retry_result = self._profile_runners( + custom_op, + runners, + input_tensors, + profile, + tuning_config, + _inputs_prepared=True, + **kwargs, + ) + return ( + retry_result[0], + retry_result[1], + retry_result[2], + retry_result[3] or has_tuning_failure_occurred, + ) + if best_runner_id is not None: # At least one valid (runner, tactic) pair is found cache_key = self.profiling_cache.get_cache_key( @@ -1521,6 +1759,37 @@ def _get_input_sizes(self, inputs: List[torch.Tensor]) -> List[torch.Size]: return sizes + @staticmethod + def _torch_profiler_elapsed_time_ms(prof, repeat: int) -> float: + """Return summed CUDA activity time per launch in milliseconds.""" + total_cuda_us = sum( + float(event.device_time_total or 0.0) for event in prof.events() + if event.device_type == torch.autograd.DeviceType.CUDA and ( + event.device_time_total or 0.0) > 0) + if not math.isfinite(total_cuda_us) or total_cuda_us <= 0: + raise _TorchProfilerUnavailableError( + "torch.profiler/CUPTI reported no CUDA activity while " + "profiling a CuTe DSL tactic") + return total_cuda_us / max(1, int(repeat)) / 1000.0 + + @staticmethod + def _is_torch_profiler_unavailable_error(error: Exception) -> bool: + """Whether an exception reports a Kineto/CUPTI backend failure.""" + if isinstance(error, _TorchProfilerUnavailableError): + return True + message = str(error).lower() + return "cupti" in message or "kineto" in message + + @classmethod + def _disable_torch_profiler(cls, error: Exception) -> None: + cls._torch_profiler_unavailable = True + logger.warning_once( + "[Autotuner] torch.profiler/CUPTI is unavailable; falling back " + "to CUDA-event timing for CuTe DSL tactics in this process. " + f"Error: {error}", + key=("autotuner", "warning_torch_profiler_cuda_event_fallback"), + ) + def _profile_single_kernel( self, runner: TunableRunner, @@ -1554,36 +1823,10 @@ def _profile_single_kernel( short_profile_threshold_ms = 1 avg_time = float('inf') + use_torch_profiler = runner.use_torch_profiler(tactic) def pure_profile(stream: torch.cuda.Stream, repeat: int): - graph = torch.cuda.CUDAGraph() - - if self._use_global_timer: - start_ts = torch.empty(1, dtype=torch.int64, device='cuda') - end_ts = torch.empty(1, dtype=torch.int64, device='cuda') - - def record_start(): - record_global_timer(start_ts.data_ptr(), stream) - - def record_end(): - record_global_timer(end_ts.data_ptr(), stream) - - def elapsed_time(): - # 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 - else: - start_evt = torch.cuda.Event(enable_timing=True) - end_evt = torch.cuda.Event(enable_timing=True) - - def record_start(): - start_evt.record() - - def record_end(): - end_evt.record() - - def elapsed_time(): - return start_evt.elapsed_time(end_evt) + graph = torch.cuda.CUDAGraph() if use_cuda_graph else None with torch.cuda.stream(stream): if use_cuda_graph: @@ -1601,6 +1844,83 @@ def elapsed_time(): # Currently only AllReduce will use this strategy, and only MPI parallel will enable tuning. self._dist.tp_barrier() + def launch_profiled_work(): + if use_cuda_graph: + graph.replay() + else: + for r in range(repeat): + runner( + input_tensor_batches[r % + len(input_tensor_batches)], + tactic=tactic, + **kwargs, + ) + + if use_torch_profiler and not AutoTuner._torch_profiler_unavailable: + profiled_work_started = False + profiled_work_completed = False + try: + with torch.profiler.profile( + activities=[ + torch.profiler.ProfilerActivity.CUDA + ], + acc_events=True, + record_shapes=False, + profile_memory=False, + with_stack=False, + ) as prof: + profiled_work_started = True + launch_profiled_work() + # CuTe DSL fused kernels may launch auxiliary + # streams; wait for every device stream before + # closing CUPTI. + torch.cuda.synchronize() + profiled_work_completed = True + return self._torch_profiler_elapsed_time_ms( + prof, repeat) + except Exception as error: + # A runner or CUDA synchronization failure is a real + # tactic failure. Do not hide it behind the timing + # fallback, even if it occurred inside the profiler + # context. + if profiled_work_started and not profiled_work_completed: + raise + if not self._is_torch_profiler_unavailable_error(error): + raise + self._disable_torch_profiler(error) + + # CUPTI was unavailable either earlier in this process or in + # the attempt above. Re-run this measurement with the legacy + # CUDA-event timer, irrespective of _use_global_timer. + force_cuda_event = use_torch_profiler + + if not force_cuda_event and self._use_global_timer: + start_ts = torch.empty(1, dtype=torch.int64, device="cuda") + end_ts = torch.empty(1, dtype=torch.int64, device="cuda") + + def record_start(): + record_global_timer(start_ts.data_ptr(), stream) + + def record_end(): + record_global_timer(end_ts.data_ptr(), stream) + + def elapsed_time(): + # 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 + else: + start_evt = torch.cuda.Event(enable_timing=True) + end_evt = torch.cuda.Event(enable_timing=True) + + def record_start(): + start_evt.record() + + def record_end(): + end_evt.record() + + def elapsed_time(): + return start_evt.elapsed_time(end_evt) + # Delay the profiled kernel launch to eliminate affects of host time overhead in profiling. if use_cuda_graph: delay_kernel(self._CUDA_GRAPH_DELAY_MICRO_SECS, stream) @@ -1608,17 +1928,7 @@ def elapsed_time(): delay_kernel(self.stream_delay_micro_secs, stream) record_start() - - if use_cuda_graph: - graph.replay() - else: - for r in range(repeat): - runner( - input_tensor_batches[r % len(input_tensor_batches)], - tactic=tactic, - **kwargs, - ) - + launch_profiled_work() record_end() stream.synchronize() @@ -1643,9 +1953,14 @@ def elapsed_time(): avg_time = pure_profile(stream, self.repeat) shapes = self._get_input_sizes(inputs) + timer_name = (_CUDA_EVENT_FALLBACK_TIMER_KEY if use_torch_profiler + and AutoTuner._torch_profiler_unavailable else + _TORCH_PROFILER_TIMER_KEY if use_torch_profiler else + "globaltimer" if self._use_global_timer else "cuda_event") + self._last_profile_timer = timer_name self._debug_logger( - f"[Autotuner] Profiled runner={runner}, tactic={tactic}, shapes={shapes}: {avg_time:.6f}ms." - ) + f"[Autotuner] Profiled runner={runner}, tactic={tactic}, " + f"shapes={shapes}, timer={timer_name}: {avg_time:.6f}ms.") return avg_time diff --git a/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py b/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py index 0d7fe4907ebd..60c1a09e12d8 100644 --- a/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py +++ b/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py @@ -16,9 +16,9 @@ from ..._utils import get_sm_version, is_sm_100f from ...math_utils import ceil_div, pad_up -from ..autotuner import (AutoTuner, ConstraintSpec, DistributedTuningStrategy, - DynamicTensorSpec, OptimizationProfile, TunableRunner, - TuningConfig) +from ..autotuner import (AutoTuner, ConstraintSpec, CuteDSLTunableRunner, + DistributedTuningStrategy, DynamicTensorSpec, + OptimizationProfile, TunableRunner, TuningConfig) from ..cute_dsl_utils import (IS_CUTLASS_DSL_AVAILABLE, IS_CUTLASS_DSL_RUBIN_AVAILABLE) from ..locality_domain.autotune import tune_locality_domain_concurrent @@ -31,9 +31,13 @@ infer_output_m_shape, is_gated_activation, last_positive_power_of_2, mxfp8_scale_infer_shape, next_positive_power_of_2) -from .cutedsl_matmul_heuristics import (NVFP4_PRECISION, - nvmmh_enabled_for_nvfp4, nvmmh_fields, - nvmmh_max_tactics, rank_configs) +from .cutedsl_matmul_heuristics import ( + BF16_FP32_OUTPUT_PRECISION, BF16_LAYOUT, BF16_PRECISION, + FP8_BF16_OUTPUT_PRECISION, FP8_FP16_OUTPUT_PRECISION, FP8_LAYOUT, + IN_PLACE_SPLIT_K_KIND, IS_NVMMH_AVAILABLE, NVFP4_LAYOUT, NVFP4_PRECISION, + cute_raster_to_nvmmh_cta_order, expand_bf16_tactics_with_heuristic_splits, + filter_bf16_tactics, filter_fp8_tactics, is_sm107_nvmmh_split_k_eligible, + nvmmh_cta_order_to_cute_raster, rank_configs) try: from cuda.bindings import driver as cuda @@ -107,6 +111,24 @@ def _kimi_k3_mxfp8_tuning_bucket(num_tokens: int) -> int: return next(m for m in _KIMI_K3_MXFP8_TUNING_BUCKETS if num_tokens <= m) return next_positive_power_of_2(num_tokens) +_CUTEDSL_NVMMH_TACTIC_SEARCH_CACHE_VERSION = "cutedsl-nvmmh-v1" + + +def _cutedsl_nvmmh_tactic_search_cache_key(): + """Return the cache discriminator for NVMMH-aware tactic generation.""" + config = AutoTuner.get().nvmmh_config + if not config.enabled or not config.fields or not IS_NVMMH_AVAILABLE: + # These paths preserve the pre-NVMMH candidate list, so policy details + # cannot affect tactic search or fragment its cache identity. + return None + return ( + _CUTEDSL_NVMMH_TACTIC_SEARCH_CACHE_VERSION, + bool(config.enabled), + bool(IS_NVMMH_AVAILABLE), + tuple(sorted(config.fields)), + int(config.max_tactics), + ) + def _with_input_cuda_device(function): """Run a custom-op implementation under its input tensor's CUDA device.""" @@ -710,7 +732,7 @@ def get_max_activate_clusters(cluster_size): 1, max_active // 2) return max_active - class CuteDSLNVFP4BlackwellRunner(TunableRunner): + class CuteDSLNVFP4BlackwellRunner(CuteDSLTunableRunner): kernel_class = Sm100BlockScaledPersistentDenseGemmKernel kernel_cache = dict() tuning_config = TuningConfig( @@ -746,6 +768,9 @@ def unique_id(self): self.use_tvm_ffi, ) + def tactic_search_cache_key(self): + return _cutedsl_nvmmh_tactic_search_cache_key() + def __hash__(self): return hash(self.unique_id()) @@ -884,9 +909,9 @@ def get_valid_tactics( logger.debug( f"CuteDSL: Found {len(valid_tactics)} valid tactics for M={m}, N={n}, K={real_k}" ) - # Optionally rank/prune the sweep with nvMatmulHeuristics (opt-in via - # TRTLLM_CUTEDSL_NVMMH_ENABLE). Returns the full sweep unchanged when - # disabled, unconfigured, or on any failure. + # Optionally rank/prune the sweep with nvMatmulHeuristics. Returns + # the full sweep unchanged when disabled, unconfigured, unavailable, + # or on any failure. return self._rank_prune_tactics(valid_tactics, m, n, real_k) def _swap_ab_candidates(self, m, n): @@ -1006,28 +1031,29 @@ def _rank_prune_tactics(self, tactics, m, n, real_k): Called at the end of get_valid_tactics. A strict, re-validated subset of ``tactics`` -- it never introduces a (tile, cluster, swap) the - kernel validator rejected. Gated by TRTLLM_CUTEDSL_NVMMH_ENABLE; - TRTLLM_CUTEDSL_NVMMH_FIELDS selects the model-driven knobs. When - "tile"/"cluster" is selected we keep only the top - TRTLLM_CUTEDSL_NVMMH_MAX_TACTICS ranked (mma_tiler, cluster, swap) - keys (all matching prefetch variants retained); when only scheduler - knobs (swizzle / cta_order) are selected we sweep every tile/cluster - and just annotate it with the model's swizzle / raster (a safe - 6-tuple extension -- they do not affect can_implement). Deterministic - swap_ab selection (small M swaps) is applied here, in the opt-in path - only. + kernel validator rejected. The AutoTuner NVMMH policy selects the + model-driven knobs and maximum number of ranked signatures. When + "tile"/"cluster" is selected we retain the top ranked + (mma_tiler, cluster, swap) keys (all matching prefetch variants); + when only scheduler knobs (swizzle / cta_order) are selected we + sweep every tile/cluster and annotate it with the model's swizzle / + raster (a safe 6-tuple extension -- they do not affect + can_implement). Deterministic swap_ab selection (small M swaps) is + applied here, in the opt-in path only. Purely additive: on any failure, an empty match, or when heuristics are disabled / unconfigured, returns ``tactics`` unchanged so profiling never loses a valid candidate. """ - if not nvmmh_enabled_for_nvfp4() or not tactics: + nvmmh_config = AutoTuner.get().nvmmh_config + if (not nvmmh_config.enabled or not IS_NVMMH_AVAILABLE + or not tactics): return tactics - fields = nvmmh_fields() + fields = set(nvmmh_config.fields) if not fields: return tactics try: - max_tactics = nvmmh_max_tactics() + max_tactics = nvmmh_config.max_tactics # Only prune the tile/cluster set when the model is asked to # drive it; scheduler-only fields keep the full sweep. prune_tile_cluster = bool(fields & {"tile", "cluster"}) @@ -1063,10 +1089,8 @@ def _rank_prune_tactics(self, tactics, m, n, real_k): continue pref_rank[key] = len(pref_rank) swizzle_size = cfg.swizzle_factor if use_swizzle else 1 - # nvMatmulHeuristics cta_order==0 is row-major (N-major) - # raster -> raster_along_m=False; !=0 -> M-major -> True. - raster_along_m = ((cfg.cta_order != 0) - if use_cta_order else True) + raster_along_m = (nvmmh_cta_order_to_cute_raster( + cfg.cta_order) == "m" if use_cta_order else True) pref_sched[key] = (max(1, int(swizzle_size)), bool(raster_along_m)) @@ -1456,7 +1480,7 @@ def _( ret = mat_a.new_empty(shape, dtype=torch.bfloat16) return ret - class CuteDSLNVFP4SwigluBlackwellRunner(TunableRunner): + class CuteDSLNVFP4SwigluBlackwellRunner(CuteDSLTunableRunner): """Runner for dense GEMM + SwiGLU fusion on Blackwell GPUs using CuteDSL. Fuses the FC1 (gate_up projection) GEMM and SwiGLU activation into a @@ -2006,7 +2030,7 @@ def _( ret = mat_a.new_empty(shape, dtype=torch.bfloat16) return ret - class CuteDSLNVFP4SwigluFP4OutBlackwellRunner(TunableRunner): + class CuteDSLNVFP4SwigluFP4OutBlackwellRunner(CuteDSLTunableRunner): """Runner for dense GEMM + SwiGLU fusion with FP4 output on Blackwell. Same as CuteDSLNVFP4SwigluBlackwellRunner but produces Float4E2M1FN @@ -2897,7 +2921,7 @@ def _( raw_sf = mat_a.new_empty(raw_shape, dtype=torch.float32) return fp4_output, raw_sf - class Sm100BlockScaledContiguousGroupedGemmRunner(TunableRunner): + class Sm100BlockScaledContiguousGroupedGemmRunner(CuteDSLTunableRunner): kernel_class = Sm100BlockScaledContiguousGroupedGemmKernel kernel_cache = dict() tuning_config_cache = dict() @@ -3196,7 +3220,7 @@ def _( return torch.empty(m, n, dtype=output_dtype, device=input.device) class Sm100BlockScaledContiguousGroupedGemmFinalizeFusionRunner( - TunableRunner): + CuteDSLTunableRunner): kernel_class = Sm100BlockScaledContiguousGroupedGemmFinalizeFusionKernel kernel_cache = dict() tuning_config_cache = dict() @@ -3723,7 +3747,7 @@ def _( device=input.device) class Sm100BlockScaledContiguousGroupedGemmSwigluFusionRunner( - TunableRunner): + CuteDSLTunableRunner): kernel_class = Sm100BlockScaledContiguousGroupedGemmSwigluFusionKernel kernel_cache = dict() tuning_config_cache = dict() @@ -4047,7 +4071,7 @@ def _( return output, output_scale class Sm100BlockScaledContiguousGatherGroupedGemmActFusionRunner( - TunableRunner): + CuteDSLTunableRunner): kernel_class = BlockScaledContiguousGatherGroupedGemmKernel kernel_cache = dict() tuning_config_cache = dict() @@ -4643,7 +4667,7 @@ def _prepare_cutedsl_indexer_q_tuning_inputs( position_ids) return inputs - class CuteDSLIndexerQBlackwellRunner(TunableRunner): + class CuteDSLIndexerQBlackwellRunner(CuteDSLTunableRunner): """Native MXF8 GEMM with fused DSv4 indexer-Q RoPE/MXFP4 output.""" kernel_class = Sm100BlockScaledPersistentDenseGemmActFusionKernel @@ -4932,7 +4956,7 @@ def _( input.new_empty((m, n // 128), dtype=torch.int32), ) - class CuteDSLFp8BlackwellRunner(TunableRunner): + class CuteDSLFp8BlackwellRunner(CuteDSLTunableRunner): kernel_class = Sm100BlockwiseGemmKernel kernel_cache = dict() @@ -4954,6 +4978,9 @@ def __init__(self, self.output_dtype = output_dtype self.use_tvm_ffi = use_tvm_ffi + def tactic_search_cache_key(self): + return _cutedsl_nvmmh_tactic_search_cache_key() + def get_valid_tactics( self, inputs: List[torch.Tensor], @@ -4991,7 +5018,7 @@ def get_valid_tactics( (4, 2), (4, 4), ] - return [ + valid_tactics = [ (use_2cta_instrs, mma_tiler_mn, cluster_shape_mn) for use_2cta_instrs in use_2cta_instrs_candi for mma_tiler_mn in mma_tiler_mn_candi @@ -5012,6 +5039,64 @@ def get_valid_tactics( c_major, ) ] + return self._rank_prune_tactics(valid_tactics, m, n, k) + + @staticmethod + def _nvmmh_tactic_signature(tactic, fields): + """Return fields nvMMH can model for blockwise FP8 GEMM.""" + if not isinstance(tactic, tuple) or len(tactic) != 3: + return None + use_2cta, mma_tiler_mn, cluster_shape_mn = tactic + signature = [] + if fields & {"tile", "cluster"}: + cta_group = 2 if use_2cta else 1 + cta = (int(mma_tiler_mn[0]) // cta_group, int(mma_tiler_mn[1])) + signature.extend((cta, tuple(cluster_shape_mn))) + return tuple(signature) + + def _rank_prune_tactics(self, tactics, m, n, k): + """Rank/prune validated SM100/SM103 FP8 tactics with nvMMH. + + Called at the end of get_valid_tactics. Any nvMMH failure or empty + match returns the validated full sweep unchanged. + """ + nvmmh_config = AutoTuner.get().nvmmh_config + if (not nvmmh_config.enabled or not IS_NVMMH_AVAILABLE + or not tactics): + return tactics + fields = set(nvmmh_config.fields) & {"tile", "cluster"} + if not fields: + return tactics + + try: + m, n, k = int(m), int(n), int(k) + max_tactics = nvmmh_config.max_tactics + ranked_configs = rank_configs( + m, + n, + k, + FP8_BF16_OUTPUT_PRECISION, + max(max_tactics * 16, 64), + layout_name=FP8_LAYOUT, + ) + selected = filter_fp8_tactics( + tactics, + ranked_configs, + fields, + max_tactics, + self._nvmmh_tactic_signature, + ) + logger.debug( + f"CuteDSL FP8 nvMatmulHeuristics: {len(tactics)} valid " + f"-> {len(selected)} tactics for M={m}, N={n}, K={k}") + return selected + except Exception as e: # noqa: BLE001 - tuning must fall back + logger.warning_once( + f"[nvMatmulHeuristics] FP8 tactic filtering failed: {e}. " + f"Falling back to the full tactic list.", + key="nvmmh_fp8_filter_failure", + ) + return tactics def forward( self, @@ -5243,7 +5328,7 @@ def _( ret = mat_a.new_empty(shape, dtype=torch.bfloat16) return ret - class CuteDSLFp8BlackwellBmmRunner(TunableRunner): + class CuteDSLFp8BlackwellBmmRunner(CuteDSLTunableRunner): kernel_class = Sm100BlockwiseGemmKernel kernel_cache = dict() @@ -5568,7 +5653,7 @@ def _( # Dense GEMM with SwiGLU Fusion (FC1 Kernel for MoE as Dense GEMM) # ============================================================================= - class CuteDSLNVFP4DenseGemmSwigluRunner(TunableRunner): + class CuteDSLNVFP4DenseGemmSwigluRunner(CuteDSLTunableRunner): """Runner for Dense GEMM with SwiGLU fusion (MoE FC1 layer as dense GEMM). This kernel performs: C = SwiGLU(alpha * (SFA * A) @ (SFB * B)) @@ -5954,7 +6039,7 @@ def _( from ..moe.kernels.blackwell.moe_as_dense_gemm.fc2 import \ Sm100BlockScaledPersistentDenseGemmKernel as DenseGemmFC2Kernel - class CuteDSLNVFP4DenseGemmFC2Runner(TunableRunner): + class CuteDSLNVFP4DenseGemmFC2Runner(CuteDSLTunableRunner): """Runner for Dense GEMM FC2 layer (MoE second projection). This kernel performs: C = (A * SFA) @ (B * SFB) * alpha_scale @@ -9546,7 +9631,7 @@ def _bf16_cluster_m_fits( cta_tile_m = mma_tiler_mn[0] // (2 if use_2cta_instrs else 1) return ceil_div(m, cta_tile_m) >= cluster_shape_mn[0] - class CuteDSLBf16BlackwellBmmRunner(TunableRunner): + class CuteDSLBf16BlackwellBmmRunner(CuteDSLTunableRunner): kernel_class = PersistentDenseGemmKernel kernel_cache = dict() # Output element type; subclasses override (see the FP8-out runner). @@ -9636,8 +9721,11 @@ def forward( tactic: Tiling and cluster strategy, typically a tuple (use_2cta_instrs, mma_tiler_mn, cluster_shape_mn). """ + kernel_options = ("m", 1, True) if isinstance(tactic, tuple): - use_2cta_instrs, mma_tiler_mn, cluster_shape_mn = tactic + use_2cta_instrs, mma_tiler_mn, cluster_shape_mn = tactic[:3] + if len(tactic) == 6: + kernel_options = tactic[3:] else: use_2cta_instrs, mma_tiler_mn, cluster_shape_mn = [ False, @@ -9692,6 +9780,7 @@ def forward( use_2cta_instrs, mma_tiler_mn, cluster_shape_mn, + *kernel_options, self.use_tvm_ffi, ) if cache_key not in self.__class__.kernel_cache: @@ -9718,6 +9807,9 @@ def forward( use_2cta_instrs=use_2cta_instrs, mma_tiler_mn=mma_tiler_mn, cluster_shape_mn=cluster_shape_mn, + raster_along=kernel_options[0], + swizzle_size=kernel_options[1], + use_tma_store=kernel_options[2], ) hardware_info = cutlass.utils.HardwareInfo() max_active_clusters = hardware_info.get_max_active_clusters( @@ -9959,7 +10051,7 @@ def _( # BF16 Dense Persistent GEMM (CuTe DSL) for Blackwell - Linear layers # ====================================================================== - class CuteDSLBf16BlackwellGemmRunner(TunableRunner): + class CuteDSLBf16BlackwellGemmRunner(CuteDSLTunableRunner): """ CuTe DSL BF16 GEMM runner for Linear layers. @@ -9970,6 +10062,10 @@ class CuteDSLBf16BlackwellGemmRunner(TunableRunner): """ kernel_class = PersistentDenseGemmKernel kernel_cache = dict() + target_sm = "blackwell" + + def tactic_search_cache_key(self): + return _cutedsl_nvmmh_tactic_search_cache_key() tuning_config = TuningConfig(dynamic_tensor_specs=(DynamicTensorSpec( 0, 0, get_last_power_of_2_num_tokens_buckets, @@ -10019,7 +10115,7 @@ def get_valid_tactics( (4, 2), (4, 4), ] - return [ + tactics = [ (use_2cta_instrs, mma_tiler_mn, cluster_shape_mn) for use_2cta_instrs in use_2cta_instrs_candi for mma_tiler_mn in mma_tiler_mn_candi @@ -10041,6 +10137,13 @@ def get_valid_tactics( ) ] + canonical = [("base", *tactic, 0, 1) for tactic in tactics] + selected = self._rank_prune_tactics(canonical, m, n, k, + inputs[2].dtype) + return [ + tuple(tactic[1:4]) + tuple(tactic[6:]) for tactic in selected + ] + def forward( self, inputs: List[torch.Tensor], @@ -10057,8 +10160,11 @@ def forward( tactic: Tiling and cluster strategy, typically a tuple (use_2cta_instrs, mma_tiler_mn, cluster_shape_mn). """ + kernel_options = ("m", 1, True) if isinstance(tactic, tuple): - use_2cta_instrs, mma_tiler_mn, cluster_shape_mn = tactic + use_2cta_instrs, mma_tiler_mn, cluster_shape_mn = tactic[:3] + if len(tactic) == 6: + kernel_options = tactic[3:] else: use_2cta_instrs, mma_tiler_mn, cluster_shape_mn = [ False, @@ -10115,6 +10221,7 @@ def forward( use_2cta_instrs, mma_tiler_mn, cluster_shape_mn, + *kernel_options, self.use_tvm_ffi, c_dtype_cutlass, ) @@ -10142,6 +10249,9 @@ def forward( use_2cta_instrs=use_2cta_instrs, mma_tiler_mn=mma_tiler_mn, cluster_shape_mn=cluster_shape_mn, + raster_along=kernel_options[0], + swizzle_size=kernel_options[1], + use_tma_store=kernel_options[2], ) hardware_info = cutlass.utils.HardwareInfo() max_active_clusters = hardware_info.get_max_active_clusters( @@ -10192,6 +10302,145 @@ def forward( if c_needs_copy: c_tensor.copy_(c_buf) + def _rank_prune_tactics(self, tactics, m, n, k, output_dtype): + """Rank/prune BF16 GEMM tactics with nvMatmulHeuristics. + + Called at the end of get_valid_tactics. The model controls only + fields selected by the AutoTuner NVMMH policy. Emitted tuples are + from the runner's validated tactic list. On SM107, split-K itself + remains a bounded local CUPTI sweep because the installed NVMMH + cannot expose its native K-tile lower bound. + Any query or matching failure keeps the validated tactic families. + """ + nvmmh_config = AutoTuner.get().nvmmh_config + if (not nvmmh_config.enabled or not IS_NVMMH_AVAILABLE + or not tactics): + return tactics + + fields = set(nvmmh_config.fields) + if self.target_sm != "rubin": + # Main supports direct split-K only in the Rubin runner. + fields.discard("split_k") + if not fields: + return tactics + fallback_tactics = tactics + try: + m, n, k = int(m), int(n), int(k) + if output_dtype == torch.bfloat16: + precision = BF16_PRECISION + elif output_dtype == torch.float32: + precision = BF16_FP32_OUTPUT_PRECISION + else: + return fallback_tactics + + max_tactics = nvmmh_config.max_tactics + # The public API returns multiple warp/stage variants for the + # same runner-visible key. Over-query so those duplicates do not + # crowd out valid CTA/cluster/split matches. + query_count = max(max_tactics * 16, 64) + ranked_configs = rank_configs( + m, + n, + k, + precision, + query_count, + layout_name=BF16_LAYOUT, + ) + if not ranked_configs: + return fallback_tactics + + candidate_tactics = tactics + model_fields = fields + if self.target_sm == "rubin" and "split_k" in fields: + # NVMMH's native minTilesKLowerBound is not configurable in + # this wheel. Admit Rubin splits with TRT-LLM's shared + # lower bound, then let CUPTI select among the runner's + # normal full-sweep candidates. + + def _split_k(tactic): + if (isinstance(tactic, tuple) and len(tactic) >= 6 + and tactic[0] == "base"): + return max(1, int(tactic[5])) + return 1 + + candidate_tactics = [ + tactic for tactic in candidate_tactics + if is_sm107_nvmmh_split_k_eligible( + k, self.nvmmh_split_k_cta_k, _split_k(tactic)) + ] + model_fields = fields - {"split_k"} + non_split_tactics = [ + tactic for tactic in candidate_tactics + if _split_k(tactic) == 1 + ] + split_tactics = [ + tactic for tactic in candidate_tactics + if _split_k(tactic) > 1 + ] + + if model_fields: + selected_non_split = filter_bf16_tactics( + non_split_tactics, + ranked_configs, + model_fields, + max_tactics, + include_known_good=not bool( + model_fields & {"swizzle", "cta_order"}), + fallback_tactics=[], + ) + if not selected_non_split: + return fallback_tactics + else: + selected_non_split = non_split_tactics + + # Rubin's split-K reduction path has fixed scheduler + # defaults, so swizzle/CTA order cannot participate in its + # model match. Tile/cluster still can when requested. + split_model_fields = model_fields - {"swizzle", "cta_order"} + if split_tactics and split_model_fields: + selected_split = filter_bf16_tactics( + split_tactics, + ranked_configs, + split_model_fields, + max_tactics, + include_known_good=False, + fallback_tactics=[], + ) + if not selected_split: + return fallback_tactics + else: + selected_split = split_tactics + selected = list( + dict.fromkeys(selected_non_split + selected_split)) + else: + if "split_k" in fields: + candidate_tactics = ( + expand_bf16_tactics_with_heuristic_splits( + tactics, ranked_configs)) + selected = filter_bf16_tactics( + candidate_tactics, + ranked_configs, + fields, + max_tactics, + include_known_good=not bool(fields + & {"swizzle", "cta_order"}), + fallback_tactics=fallback_tactics, + ) + logger.debug( + f"CuTe DSL BF16 nvMatmulHeuristics: {len(tactics)} baseline " + f"/ {len(candidate_tactics)} candidates " + f"-> {len(selected)} tactics for M={m}, N={n}, K={k}, " + f"precision={precision}, fields={sorted(fields)}, " + f"model_fields={sorted(model_fields)}") + return selected + except Exception as e: # noqa: BLE001 - tuning must keep working + logger.warning_once( + f"[nvMatmulHeuristics] BF16 tactic filtering failed: {e}. " + f"Falling back to the validated tactic list.", + key="nvmmh_bf16_filter_failure", + ) + return fallback_tactics + # input: [M, K], weight: [N, K], output: [M, N] @torch.library.custom_op("trtllm::cute_dsl_bf16_gemm_blackwell", mutates_args=("output", ), @@ -10378,6 +10627,11 @@ def _parse_sm107_bf16_tactic(tactic): """Decode a tactic into (kernel_variant, use_2cta_instrs, mma_tiler_mn, preferred_cluster_shape_mn, cluster_shape_mn, max_num_ab_stage, split_k_slices).""" + if isinstance(tactic, tuple) and tactic and tactic[0] == "base": + if len(tactic) == 9: + tactic = tactic[:6] + elif len(tactic) == 8: + tactic = tactic[:5] if (isinstance(tactic, tuple) and len(tactic) > 0 and isinstance(tactic[0], str)): kernel_variant = tactic[0] @@ -10408,7 +10662,10 @@ def _sm107_bf16_gemm_kernel(kernel_variant: str, preferred_cluster_shape_mn, cluster_shape_mn, max_num_ab_stage: int, - split_k_slices: int = 1): + split_k_slices: int = 1, + raster_along: str = "m", + swizzle_size: int = 1, + use_tma_store: bool = True): kernel_class = _sm107_bf16_kernel_class(kernel_variant) if kernel_variant == "preferred_cluster": return kernel_class( @@ -10426,6 +10683,9 @@ def _sm107_bf16_gemm_kernel(kernel_variant: str, cluster_shape_mn=cluster_shape_mn, max_num_ab_stage=max_num_ab_stage, split_k_slices=split_k_slices, + raster_along=raster_along, + swizzle_size=swizzle_size, + use_tma_store=use_tma_store, ) def _sm107_bf16_occupancy(cluster_shape_mn, preferred_cluster_shape_mn): @@ -10479,6 +10739,10 @@ def forward( (kernel_variant, use_2cta_instrs, mma_tiler_mn, preferred_cluster_shape_mn, cluster_shape_mn, max_num_ab_stage, _) = _parse_sm107_bf16_tactic(tactic) + kernel_options = ("m", 1, True) + if isinstance(tactic, tuple) and tactic and tactic[0] == "base": + if len(tactic) in (8, 9): + kernel_options = tactic[-3:] a_tensor, b_tensor, c_tensor = inputs batch_size, m, k = a_tensor.shape @@ -10541,6 +10805,7 @@ def forward( preferred_cluster_shape_mn, cluster_shape_mn, max_num_ab_stage, + *kernel_options, self.use_tvm_ffi, c_layout_key, max_active_clusters, @@ -10561,11 +10826,15 @@ def forward( stream = cute.runtime.make_fake_stream( use_tvm_ffi_env_stream=True) - gemm = _sm107_bf16_gemm_kernel(kernel_variant, use_2cta_instrs, + gemm = _sm107_bf16_gemm_kernel(kernel_variant, + use_2cta_instrs, mma_tiler_mn, preferred_cluster_shape_mn, cluster_shape_mn, - max_num_ab_stage) + max_num_ab_stage, + raster_along=kernel_options[0], + swizzle_size=kernel_options[1], + use_tma_store=kernel_options[2]) compile_args = [ m, n, k, batch_size, a_ptr, b_ptr, c_cute_tensor, a_stride_m, a_stride_batch, b_stride_n, b_stride_batch @@ -10929,6 +11198,11 @@ class CuteDSLBf16RubinGemmRunner(CuteDSLBf16BlackwellGemmRunner): # Split-K artifacts are shape-independent (tactic plus occupancy) and # cached separately from the split=1 path. split_k_gemm_cache = dict() + target_sm = "rubin" + nvmmh_split_k_cta_k = 64 + + def tactic_search_cache_key(self): + return _cutedsl_nvmmh_tactic_search_cache_key() # See CuteDSLBf16RubinBmmRunner.tuning_config. tuning_config = TuningConfig( @@ -10981,8 +11255,10 @@ def get_valid_tactics( if k >= 4096 and n <= 512 and inputs[2].shape[0] == m: split_k_candi = [1, 2, 4, 8] - return _sm107_bf16_valid_tactics(m, n, k, 1, c_dtype_cutlass, - mma_tiler_mn_candi, split_k_candi) + tactics = _sm107_bf16_valid_tactics(m, n, k, 1, c_dtype_cutlass, + mma_tiler_mn_candi, + split_k_candi) + return self._rank_prune_tactics(tactics, m, n, k, inputs[2].dtype) def forward( self, @@ -10992,6 +11268,10 @@ def forward( (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) + kernel_options = ("m", 1, True) + if isinstance(tactic, tuple) and tactic and tactic[0] == "base": + if len(tactic) in (8, 9): + kernel_options = tactic[-3:] a_tensor, b_tensor, c_tensor = inputs m, k = a_tensor.shape @@ -11061,6 +11341,7 @@ def forward( preferred_cluster_shape_mn, cluster_shape_mn, max_num_ab_stage, + *kernel_options, self.use_tvm_ffi, c_dtype_cutlass, c_layout_key, @@ -11082,11 +11363,15 @@ def forward( stream = cute.runtime.make_fake_stream( use_tvm_ffi_env_stream=True) - gemm = _sm107_bf16_gemm_kernel(kernel_variant, use_2cta_instrs, + gemm = _sm107_bf16_gemm_kernel(kernel_variant, + use_2cta_instrs, mma_tiler_mn, preferred_cluster_shape_mn, cluster_shape_mn, - max_num_ab_stage) + max_num_ab_stage, + raster_along=kernel_options[0], + swizzle_size=kernel_options[1], + use_tma_store=kernel_options[2]) compile_args = [ m, n, k, batch_size, a_ptr, b_ptr, c_cute_tensor ] @@ -12970,7 +13255,7 @@ def _( (4, 4), ] - class CuteDSLFp8RubinGemmRunner(TunableRunner): + class CuteDSLFp8RubinGemmRunner(CuteDSLTunableRunner): """FP8 1x128/128x128 blockwise GEMM runner for SM107.""" kernel_class = SM107BlockwiseGemmKernel @@ -13650,7 +13935,7 @@ def _( assert packed_scale.shape == (expected_scale_numel, ) assert packed_scale.is_contiguous() - class CuteDSLBlockScaledRubinLinear(TunableRunner): + class CuteDSLBlockScaledRubinLinear(CuteDSLTunableRunner): """Datatype-agnostic block-scaled Linear runner for SM107. Concrete datatypes provide their element types, scale-factor @@ -13670,6 +13955,15 @@ class CuteDSLBlockScaledRubinLinear(TunableRunner): mma_tiler_k: ClassVar[int] mma_inst_shape_k: ClassVar[int] allow_unswapped_small_m_fallback: ClassVar[bool] + # nvMatmulHeuristics models the GEMM element types and layout, but + # not the block-scale format. Concrete runners opt in with the + # closest matching compute precision. A None precision keeps the + # exact full-sweep behavior for unsupported datatypes. + nvmmh_precision: ClassVar[Optional[str]] = None + nvmmh_layout: ClassVar[Optional[str]] = None + nvmmh_supported_fields: ClassVar[frozenset] = frozenset() + nvmmh_sweep_cluster_shape: ClassVar[bool] = False + nvmmh_supplement_mixed_clusters: ClassVar[bool] = False raster_order_candidates: ClassVar[Tuple[str, ...]] = ("m", ) scheduler_mode_candidates: ClassVar[Tuple[str, ...]] = ("static", @@ -13697,6 +13991,11 @@ def unique_id(self): return (self.output_dtype, self.to_userbuffers, self.use_tvm_ffi) + def tactic_search_cache_key(self): + if (self.nvmmh_precision is None or self.nvmmh_layout is None): + return None + return _cutedsl_nvmmh_tactic_search_cache_key() + def __hash__(self): return hash( (self.output_dtype, self.to_userbuffers, self.use_tvm_ffi)) @@ -13797,7 +14096,6 @@ def get_valid_tactics( (2, 2), (4, 2), ] - raster_order_candidates = self.raster_order_candidates use_prefetch_candidates = [True, False] valid_tactics = [] @@ -13805,8 +14103,9 @@ def get_valid_tactics( mma_tiler_m, mma_inst_m ), mma_n, cluster_shape_mn, raster_order, swap_ab, use_prefetch in itertools.product( mma_m_candidates, mma_n_candidates, - cluster_shape_mn_candidates, raster_order_candidates, - base_swap_ab_candidates, use_prefetch_candidates): + cluster_shape_mn_candidates, + self.raster_order_candidates, base_swap_ab_candidates, + use_prefetch_candidates): # Check 2CTA cluster constraint: cluster_m must be divisible by 2 if mma_inst_m == 256 and cluster_shape_mn[0] % 2 != 0: @@ -13865,7 +14164,7 @@ def get_valid_tactics( mma_tiler_m, mma_inst_m ), mma_n, raster_order, swap_ab, use_prefetch in itertools.product( mma_m_candidates, mma_n_candidates, - raster_order_candidates, mixed_swap_ab_candidates, + self.raster_order_candidates, mixed_swap_ab_candidates, use_prefetch_candidates): mma_tiler_mnk = (mma_tiler_m, mma_n, mma_tiler_k) mma_inst_shape = (mma_inst_m, mma_n, mma_inst_shape_k) @@ -13897,7 +14196,7 @@ def get_valid_tactics( if self.mixed_scheduler_mode_candidates == ("static", ): # Preserve the existing raster-only tactic schema # for runners that do not profile mixed-cluster - # schedulers (currently NVFP4). + # scheduler fields. valid_tactics.append( ("mixed_clusters", mma_tiler_mnk, mma_inst_shape, preferred_cluster_shape_mn, @@ -13925,7 +14224,368 @@ def get_valid_tactics( logger.debug( f"CuteDSL SM107: Found {len(valid_tactics)} valid tactics for M={m}, N={n}, K={real_k}" ) - return valid_tactics + return self._rank_prune_tactics(valid_tactics, m, n, real_k) + + @staticmethod + def _nvmmh_tactic_signature(tactic, fields): + """Return the selected NVMMH fields plus the swap orientation. + + Rubin encodes two-CTA MMA in ``mma_inst_shape[0] == 256``; + nvMatmulHeuristics reports the corresponding per-CTA M tile. + The model has one cluster field, so mixed-cluster tactics use + their fallback cluster. Instruction mode, B reuse, prefetch, + and persistent scheduler mode remain local sweep choices. + """ + if not isinstance(tactic, tuple) or not tactic: + return None + raster_order = "m" + split_k = 1 + swizzle_size = 1 + if tactic[0] == "base" and len(tactic) in (6, 7, 8, 9, 10): + mma_tiler_mnk, mma_inst_shape = tactic[1], tactic[2] + cluster_shape_mn, swap_ab = tactic[3], tactic[4] + if len(tactic) >= 8: + raster_order = tactic[7] + if len(tactic) >= 9: + split_k = tactic[8] + if len(tactic) == 10: + swizzle_size = tactic[9] + elif (tactic[0] == "mixed_clusters" + and len(tactic) in (7, 8, 9, 10)): + mma_tiler_mnk, mma_inst_shape = tactic[1], tactic[2] + cluster_shape_mn, swap_ab = tactic[4], tactic[5] + # Length eight is the legacy raster-only schema, but still + # older cache entries used the same slot for a fixed + # "static" scheduler. Length nine is either the new + # raster/swizzle schema or the old scheduler/raster schema. + if len(tactic) == 8 and tactic[7] in {"m", "n"}: + raster_order = tactic[7] + elif len(tactic) == 9: + if tactic[7] in {"m", "n"}: + raster_order = tactic[7] + swizzle_size = tactic[8] + elif tactic[7] in {"static", "clc_dynamic"}: + raster_order = tactic[8] + else: + return None + elif len(tactic) == 10: + if tactic[7] not in {"static", "clc_dynamic"}: + return None + raster_order = tactic[8] + swizzle_size = tactic[9] + else: + return None + + if raster_order not in {"m", "n"}: + return None + if (isinstance(swizzle_size, bool) + or not isinstance(swizzle_size, int) + or swizzle_size < 1): + return None + if (isinstance(split_k, bool) or not isinstance(split_k, int) + or split_k < 1): + return None + + signature = [] + if fields & {"tile", "cluster"}: + cta_group = (2 if int(mma_inst_shape[0]) == 256 else 1) + cta = (int(mma_tiler_mnk[0]) // cta_group, + int(mma_tiler_mnk[1])) + signature.extend((cta, tuple(cluster_shape_mn))) + if "split_k" in fields: + signature.append(int(split_k)) + if "swizzle" in fields: + signature.append(int(swizzle_size)) + if "cta_order" in fields: + signature.append( + cute_raster_to_nvmmh_cta_order(raster_order)) + # swap_ab is not an NVMMH field. Keep it in the key because + # each orientation is ranked as a separate M/N problem below. + signature.append(bool(swap_ab)) + return tuple(signature) + + @staticmethod + def _nvmmh_config_signature(config, fields, swap_ab): + """Map one model result to the block-scaled runner key.""" + signature = [] + if fields & {"tile", "cluster"}: + signature.extend((tuple(config.cta), tuple(config.cluster))) + if "split_k" in fields: + signature.append(max(1, int(config.split_k))) + if "swizzle" in fields: + signature.append(max(1, int(config.swizzle_factor))) + if "cta_order" in fields: + signature.append(int(config.cta_order)) + signature.append(bool(swap_ab)) + return tuple(signature) + + @staticmethod + def _nvmmh_cluster_family(signature, fields, sweep_shape): + """Drop runner-local cluster fields from a model-match key.""" + if not (fields & {"tile", "cluster"}): + return signature + cta, cluster, *remaining = signature + if sweep_shape: + return (cta, *remaining) + return signature + + def _apply_nvmmh_scheduler(self, tactic, config, fields): + """Copy selected scheduler fields directly from one config. + + Raw validated tactics use the runner's native scheduler + settings. This helper materializes no local scheduler domain: + it only annotates values returned by NVMMH after structural + tactic matching. + """ + scheduler_fields = fields & {"swizzle", "cta_order"} + if not scheduler_fields: + return tactic + raster_order = (nvmmh_cta_order_to_cute_raster(config.cta_order) + if "cta_order" in scheduler_fields else None) + swizzle_size = (max(1, int(config.swizzle_factor)) + if "swizzle" in scheduler_fields else 1) + + if tactic[0] == "base": + if len(tactic) not in {9, 10}: + return None + annotated = list(tactic[:9]) + if raster_order is not None: + annotated[7] = raster_order + if "swizzle" in scheduler_fields: + if annotated[6] == "clc_dynamic" and swizzle_size != 1: + return None + annotated.append(swizzle_size) + return tuple(annotated) + + if tactic[0] != "mixed_clusters" or len(tactic) not in { + 8, 9, 10 + }: + return None + if len(tactic) == 8 and tactic[7] in {"m", "n"}: + scheduler_mode = "static" + current_raster = tactic[7] + annotated = list(tactic[:7]) + annotated.append(raster_order or current_raster) + elif len(tactic) >= 9 and tactic[7] in { + "static", "clc_dynamic" + }: + scheduler_mode = tactic[7] + current_raster = tactic[8] + annotated = list(tactic[:8]) + annotated.append(raster_order or current_raster) + elif len(tactic) == 9 and tactic[7] in {"m", "n"}: + # Accept the historical raster/swizzle schema, but replace + # its scheduler value rather than carrying it forward. + scheduler_mode = "static" + current_raster = tactic[7] + annotated = list(tactic[:7]) + annotated.append(raster_order or current_raster) + else: + return None + effective_raster = annotated[-1] + if (scheduler_mode == "clc_dynamic" + and (effective_raster != "m" or swizzle_size != 1)): + return None + if "swizzle" in scheduler_fields: + annotated.append(swizzle_size) + return tuple(annotated) + + def _rank_prune_tactics(self, tactics, m, n, real_k): + """Rank/prune Rubin block-scaled tactics with NVMMH. + + Split-K uses TRT-LLM's SM107 K-tile admission policy and remains + a local CUPTI sweep field. This is necessary because the native + scheduler's hard-coded lower bound can prevent useful factors + from being modeled at all. NVMMH still controls every other + selected field. Scheduler values are copied directly from the + matched model config instead of being locally swept. + + MXFP8 uses the generic E4M3 FP8 compute model because NVMMH + does not encode scale-factor formats. Runner-native fields + absent from the model key remain swept. + """ + if (self.nvmmh_precision is None or self.nvmmh_layout is None + or not tactics): + return tactics + nvmmh_config = AutoTuner.get().nvmmh_config + if not nvmmh_config.enabled or not IS_NVMMH_AVAILABLE: + return tactics + fields = (set(nvmmh_config.fields) + & set(self.nvmmh_supported_fields)) + if not fields: + return tactics + fallback_tactics = tactics + # The installed NVMMH has no public property for overriding + # minTilesKLowerBound. Treat split-K as an admitted local sweep + # dimension instead of matching its native rank-1 value. + model_fields = fields - {"split_k"} + scheduler_fields = model_fields & {"swizzle", "cta_order"} + match_fields = model_fields - scheduler_fields + + try: + m, n, real_k = int(m), int(n), int(real_k) + max_tactics = nvmmh_config.max_tactics + candidate_tactics = tactics + if "split_k" in fields: + + def _split_k(tactic): + if (isinstance(tactic, tuple) and tactic + and tactic[0] == "base" + and len(tactic) >= 9): + return int(tactic[8]) + return 1 + + candidate_tactics = [ + tactic for tactic in candidate_tactics + if is_sm107_nvmmh_split_k_eligible( + real_k, self.mma_tiler_k, _split_k(tactic)) + ] + if not model_fields: + return candidate_tactics or fallback_tactics + + baseline_signatures = { + signature + for tactic in candidate_tactics + if (signature := self._nvmmh_tactic_signature( + tactic, match_fields)) is not None + } + if not baseline_signatures: + return fallback_tactics + + valid_swaps = sorted( + {signature[-1] + for signature in baseline_signatures}) + query_count = max(max_tactics * 16, 64) + ranked_configs_by_swap = {} + for swap_ab in valid_swaps: + kernel_m, kernel_n = (n, m) if swap_ab else (m, n) + query_kwargs = {"layout_name": self.nvmmh_layout} + if "split_k" in fields: + query_kwargs["split_k_kind"] = ( + IN_PLACE_SPLIT_K_KIND) + ranked_configs = rank_configs( + kernel_m, + kernel_n, + real_k, + self.nvmmh_precision, + query_count, + **query_kwargs, + ) + ranked_configs_by_swap[swap_ab] = ranked_configs + valid_signatures = { + signature + for tactic in candidate_tactics + if (signature := self._nvmmh_tactic_signature( + tactic, match_fields)) is not None + } + if not valid_signatures: + return fallback_tactics + + def model_key(signature): + return self._nvmmh_cluster_family( + signature, + match_fields, + self.nvmmh_sweep_cluster_shape, + ) + + valid_keys = { + model_key(signature) + for signature in valid_signatures + } + config_by_key = {} + rank1_config_by_swap = {} + matched_swaps = set() + for swap_ab in valid_swaps: + matches_for_swap = 0 + for config in ranked_configs_by_swap[swap_ab]: + rank1_config_by_swap.setdefault(swap_ab, config) + signature = self._nvmmh_config_signature( + config, match_fields, swap_ab) + key = model_key(signature) + if key not in valid_keys or key in config_by_key: + continue + config_by_key[key] = config + matches_for_swap += 1 + if matches_for_swap >= max_tactics: + break + if matches_for_swap: + matched_swaps.add(swap_ab) + + if not matched_swaps: + return fallback_tactics + + # Each swap orientation is a different model problem. If + # only one orientation has a representable model result, + # fail open for the unmatched orientation instead of + # discarding its validated full-sweep candidates. + unmatched_swaps = set(valid_swaps) - matched_swaps + + selected = [] + for tactic in candidate_tactics: + signature = self._nvmmh_tactic_signature( + tactic, match_fields) + if signature is None: + continue + config = config_by_key.get(model_key(signature)) + if config is None: + continue + annotated = self._apply_nvmmh_scheduler( + tactic, config, scheduler_fields) + if annotated is not None: + selected.append(annotated) + selected.extend( + tactic for tactic in fallback_tactics + if ((signature := self._nvmmh_tactic_signature( + tactic, match_fields)) is not None + and signature[-1] in unmatched_swaps)) + + # NVMMH cannot encode the preferred/fallback cluster pair. + # Keep the canonical validated mixed-cluster family for + # runners that have established it as a known-good tactic. + if self.nvmmh_supplement_mixed_clusters: + selected_mixed_swaps = { + bool(tactic[5]) + for tactic in selected + if tactic[0] == "mixed_clusters" + } + for tactic in tactics: + if (tactic[0] != "mixed_clusters" + or bool(tactic[5]) in selected_mixed_swaps + or tuple(tactic[1]) != (256, 256, + self.mma_tiler_k) + or tuple(tactic[2]) + != (256, 256, self.mma_inst_shape_k) + or tuple(tactic[3]) != (4, 2) + or tuple(tactic[4]) != (2, 1)): + continue + swap_ab = bool(tactic[5]) + config = rank1_config_by_swap.get(swap_ab) + if config is None: + continue + annotated = self._apply_nvmmh_scheduler( + tactic, config, scheduler_fields) + if annotated is not None: + selected.append(annotated) + + selected = list(dict.fromkeys(selected)) + logger.debug( + f"CuteDSL Rubin {self.gemm_name} " + f"nvMatmulHeuristics: {len(tactics)} valid -> " + f"{len(selected)} tactics ({len(config_by_key)} model keys) " + f"for M={m}, N={n}, K={real_k}, " + f"precision={self.nvmmh_precision}, " + f"fields={sorted(fields)}, " + f"model_fields={sorted(model_fields)}") + return selected if selected else fallback_tactics + except Exception as e: # noqa: BLE001 - tuning must fall back + logger.warning_once( + f"[nvMatmulHeuristics] Rubin {self.gemm_name} tactic " + f"filtering failed: {e}. Falling back to the " + f"validated tactic list.", + key=("nvmmh_rubin_" + f"{self.gemm_name.lower()}_filter_failure"), + ) + return fallback_tactics def make_cute_dsl_global_pointer(self, tensor: torch.Tensor, dtype, assumed_align: int): @@ -13956,11 +14616,11 @@ def forward( tactic: Tiling and cluster strategy: base: ("base", mma_tiler_mnk, mma_inst_shape, cluster_shape, swap_ab, use_prefetch, scheduler_mode, - raster_order, split_k) + raster_order, split_k[, swizzle_size]) mixed: ("mixed_clusters", mma_tiler_mnk, mma_inst_shape, preferred_cluster_shape, fallback_cluster_shape, swap_ab, use_prefetch, - scheduler_mode, raster_order) + scheduler_mode, raster_order[, swizzle_size]) bias: Optional per-N bias [N]. Added post-GEMM for the non-inplace path. @@ -13969,16 +14629,17 @@ def forward( """ sf_vec_size = self.sf_vec_size split_k = 1 + swizzle_size = 1 if (isinstance(tactic, tuple) and len(tactic) > 0 and isinstance(tactic[0], str)): kernel_variant = tactic[0] if kernel_variant == "mixed_clusters": - if len(tactic) not in {7, 8, 9}: + if len(tactic) not in {7, 8, 9, 10}: raise ValueError( "CuteDSL SM107 mixed-cluster tactics must use " "Boolean use_prefetch with optional scheduler " - "mode and raster order") + "mode, raster order, and swizzle") (_, mma_tiler_mnk, mma_inst_shape, preferred_cluster_shape_mn, cluster_shape_mn, swap_ab, use_prefetch, *tuning_options) = tactic @@ -13994,14 +14655,23 @@ def forward( raster_order = "m" else: raster_order = tuning_options[0] + elif len(tuning_options) == 2: + if tuning_options[ + 0] in self.mixed_scheduler_mode_candidates: + scheduler_mode, raster_order = tuning_options + else: + # Accept the raster/swizzle schema emitted + # before mixed-cluster scheduler profiling. + raster_order, swizzle_size = tuning_options else: - scheduler_mode, raster_order = tuning_options + scheduler_mode, raster_order, swizzle_size = tuning_options else: - if len(tactic) not in {6, 7, 8, 9}: + if len(tactic) not in {6, 7, 8, 9, 10}: raise ValueError( "CuteDSL SM107 base tactics must use Boolean " "use_prefetch with an optional scheduler mode " - "and raster order, followed by split_k") + "and raster order, followed by split_k and " + "optional swizzle_size") (_, mma_tiler_mnk, mma_inst_shape, cluster_shape_mn, swap_ab, use_prefetch, *tuning_options) = tactic preferred_cluster_shape_mn = None @@ -14011,8 +14681,10 @@ def forward( scheduler_mode = tuning_options[0] if len(tuning_options) >= 2: raster_order = tuning_options[1] - if len(tuning_options) == 3: + if len(tuning_options) >= 3: split_k = tuning_options[2] + if len(tuning_options) == 4: + swizzle_size = tuning_options[3] elif isinstance(tactic, tuple) and len(tactic) == 5: kernel_variant = "base" (mma_tiler_mnk, mma_inst_shape, cluster_shape_mn, swap_ab, @@ -14043,6 +14715,11 @@ def forward( if raster_order not in {"m", "n"}: raise ValueError("CuteDSL SM107 raster_order must be 'm' " f"or 'n', got {raster_order!r}") + if (not isinstance(swizzle_size, int) + or isinstance(swizzle_size, bool) or swizzle_size < 1): + raise ValueError( + "CuteDSL Rubin swizzle_size must be a positive integer, " + f"got {swizzle_size!r}") if (kernel_variant == "mixed_clusters" and scheduler_mode not in self.mixed_scheduler_mode_candidates): raise ValueError( @@ -14057,6 +14734,10 @@ def forward( if scheduler_mode not in {"static", "clc_dynamic"}: raise ValueError("Unsupported CuteDSL SM107 scheduler mode " f"{scheduler_mode!r}") + if scheduler_mode == "clc_dynamic" and swizzle_size != 1: + raise ValueError( + "CuteDSL Rubin clc_dynamic scheduler only supports " + "swizzle_size=1 with the current CuTe DSL release") assert len(inputs) == 5 or len(inputs) == 6 if len(inputs) == 5: @@ -14223,9 +14904,10 @@ def forward( cache_key = (kernel_variant, sf_vec_size, mma_tiler_mnk, mma_inst_shape, preferred_cluster_shape_mn, cluster_shape_mn, swap_ab, use_prefetch, - raster_order, scheduler_mode, split_k, - locality_domain_half_gemm, self.use_tvm_ffi, - max_active_preferred_clusters, max_active_clusters) + raster_order, swizzle_size, scheduler_mode, + split_k, locality_domain_half_gemm, + self.use_tvm_ffi, max_active_preferred_clusters, + max_active_clusters) alpha_cute_tensor = cute.runtime.from_dlpack(alpha_tensor) # get stream @@ -14278,6 +14960,7 @@ def forward( preferred_cluster_shape_mn, cluster_shape_mn, None if use_prefetch else 0, + swizzle_size=swizzle_size, raster_order=raster_order, scheduler_type=scheduler_type, ) @@ -14288,6 +14971,7 @@ def forward( mma_tiler_mnk, cluster_shape_mn, None if use_prefetch else 0, + swizzle_size=swizzle_size, raster_order=raster_order, scheduler_type=scheduler_type, split_k=split_k, @@ -14390,6 +15074,11 @@ class CuteDSLNVFP4RubinLinear(CuteDSLBlockScaledRubinLinear): kernel_cache = dict() gemm_name = "NVFP4" + nvmmh_precision = NVFP4_PRECISION + nvmmh_layout = NVFP4_LAYOUT + nvmmh_supported_fields = frozenset( + {"tile", "cluster", "swizzle", "cta_order", "split_k"}) + nvmmh_supplement_mixed_clusters = True a_dtype = cutlass.Float4E2M1FN b_dtype = cutlass.Float4E2M1FN sf_dtype = cutlass.Float8E4M3FN @@ -14417,6 +15106,17 @@ class CuteDSLMXFP8RubinLinear(CuteDSLBlockScaledRubinLinear): kernel_cache = dict() gemm_name = "MXFP8" + # NVMMH has no block-scale descriptor. QQTST is the matching E4M3 + # A/B, FP32-accumulate, BF16-output compute model; the CuTe DSL + # validity checks continue to enforce MXFP8's UE8M0 scale layout. + nvmmh_precision = FP8_BF16_OUTPUT_PRECISION + nvmmh_layout = FP8_LAYOUT + # NVMMH can prune tile/cluster choices or fix the rank-1 scheduler + # tuple. Split-K uses the shared SM107 admission rule and remains + # locally profiled. + nvmmh_supported_fields = frozenset( + {"tile", "cluster", "swizzle", "cta_order", "split_k"}) + nvmmh_supplement_mixed_clusters = True a_dtype = cutlass.Float8E4M3FN b_dtype = cutlass.Float8E4M3FN sf_dtype = cutlass.Float8E8M0FNU @@ -14940,13 +15640,13 @@ def _( ) -> None: return None - class CuteDSLFp8RubinPerTensorGemmRunner(TunableRunner): + class CuteDSLFp8RubinPerTensorGemmRunner(CuteDSLTunableRunner): kernel_class = SM107PersistentDenseGemmKernel kernel_cache = dict() tuning_config = TuningConfig( dynamic_tensor_specs=(DynamicTensorSpec( - 0, 1, get_last_power_of_2_num_tokens_buckets, + 0, 0, get_last_power_of_2_num_tokens_buckets, last_positive_power_of_2), ), use_cold_l2_cache=True, ) @@ -14962,6 +15662,9 @@ def __init__(self, self.output_dtype = output_dtype self.use_tvm_ffi = use_tvm_ffi + def tactic_search_cache_key(self): + return _cutedsl_nvmmh_tactic_search_cache_key() + def get_valid_tactics( self, inputs: List[torch.Tensor], @@ -15047,7 +15750,81 @@ def get_valid_tactics( valid_tactics.append( (use_2cta_instrs, mma_tiler, mma_inst_shape, cluster_shape_mn, use_tma_store, raster_along)) - return valid_tactics + return self._rank_prune_tactics(valid_tactics, m, n, k) + + @staticmethod + def _nvmmh_tactic_signature(tactic, fields): + """Return the nvMMH-visible FP8 tile, cluster, and CTA order.""" + if not isinstance(tactic, tuple) or len(tactic) != 6: + return None + use_2cta, mma_tiler_mnk, _, cluster_shape_mn, _, raster_along = tactic + signature = [] + if fields & {"tile", "cluster"}: + cta_group = 2 if use_2cta else 1 + cta = ( + int(mma_tiler_mnk[0]) // cta_group, + int(mma_tiler_mnk[1]), + ) + signature.extend((cta, tuple(cluster_shape_mn))) + if "cta_order" in fields: + signature.append( + cute_raster_to_nvmmh_cta_order(raster_along)) + return tuple(signature) + + def _rank_prune_tactics(self, tactics, m, n, k): + """Rank/prune validated Rubin per-tensor FP8 tactics with nvMMH. + + Called at the end of get_valid_tactics. MMA instruction shape, + B-reuse, and TMA-store mode remain local choices. Raster order + remains swept unless ``cta_order`` is explicitly + model-controlled. Cluster-N is also swept within each + nvMMH-selected CTA/cluster-M family: current SM107 model rankings + consistently select cluster-N=1, while validated cluster-N + alternatives are required to avoid a performance regression. + """ + nvmmh_config = AutoTuner.get().nvmmh_config + if (not nvmmh_config.enabled or not IS_NVMMH_AVAILABLE + or not tactics): + return tactics + fields = set( + nvmmh_config.fields) & {"tile", "cluster", "cta_order"} + if not fields: + return tactics + + try: + m, n, k = int(m), int(n), int(k) + max_tactics = nvmmh_config.max_tactics + precision = (FP8_BF16_OUTPUT_PRECISION + if self.output_dtype == torch.bfloat16 else + FP8_FP16_OUTPUT_PRECISION) + ranked_configs = rank_configs( + m, + n, + k, + precision, + max(max_tactics * 16, 64), + layout_name=FP8_LAYOUT, + ) + selected = filter_fp8_tactics( + tactics, + ranked_configs, + fields, + max_tactics, + self._nvmmh_tactic_signature, + sweep_cluster_n=True, + ) + logger.debug( + f"CuteDSL Rubin FP8 nvMatmulHeuristics: {len(tactics)} " + f"valid -> {len(selected)} tactics for M={m}, N={n}, K={k}, " + f"precision={precision}, fields={sorted(fields)}") + return selected + except Exception as e: # noqa: BLE001 - tuning must fall back + logger.warning_once( + f"[nvMatmulHeuristics] Rubin FP8 tactic filtering failed: " + f"{e}. Falling back to the full tactic list.", + key="nvmmh_rubin_fp8_filter_failure", + ) + return tactics def forward( self, @@ -15297,7 +16074,7 @@ def _( validate_activation_type as validate_rubin_activation_type class Sm107BlockScaledContiguousGatherGroupedGemmActFusionRunner( - TunableRunner): + CuteDSLTunableRunner): """Rubin runner for gather + grouped GEMM + activation fusion. SM107 counterpart to Blackwell's @@ -16104,7 +16881,8 @@ def _( from ..moe.kernels.rubin.rubin_contiguous_gather_grouped_gemm_swiglu_fusion import \ Sm107ContiguousGatherGroupedGemmSwigluFusionKernel - class Sm107ContiguousGatherGroupedGemmSwigluFusionRunner(TunableRunner): + class Sm107ContiguousGatherGroupedGemmSwigluFusionRunner( + CuteDSLTunableRunner): """Rubin (SM107) runner for BF16/FP16 gather + grouped GEMM + SwiGLU fusion (FC1). Similar to the blockscaled runner but without scale factors. @@ -16699,7 +17477,7 @@ def _( Sm107BlockScaledContiguousGroupedGemmFinalizeFusionKernel class Sm107BlockScaledContiguousGroupedGemmFinalizeFusionRunner( - TunableRunner): + CuteDSLTunableRunner): """Rubin (SM107) runner for grouped GEMM + finalize fusion (FC2). This is the Rubin counterpart to @@ -17400,7 +18178,8 @@ def _( from ..moe.kernels.rubin.rubin_contiguous_grouped_gemm_finalize_fusion import \ Sm107ContiguousGroupedGemmFinalizeFusionKernel - class Sm107ContiguousGroupedGemmFinalizeFusionRunner(TunableRunner): + class Sm107ContiguousGroupedGemmFinalizeFusionRunner( + CuteDSLTunableRunner): """Rubin (SM107) runner for BF16/FP16 grouped GEMM + finalize fusion (FC2). Similar to the blockscaled finalize runner but without scale factors. diff --git a/tensorrt_llm/_torch/custom_ops/cutedsl_matmul_heuristics.py b/tensorrt_llm/_torch/custom_ops/cutedsl_matmul_heuristics.py index 0113a6c3a515..50841a84bf98 100644 --- a/tensorrt_llm/_torch/custom_ops/cutedsl_matmul_heuristics.py +++ b/tensorrt_llm/_torch/custom_ops/cutedsl_matmul_heuristics.py @@ -14,10 +14,10 @@ ranking so callers fall back to the full tactic list and never regress. """ -import os +import ctypes import re from functools import lru_cache -from typing import List, NamedTuple, Optional, Tuple +from typing import Callable, List, NamedTuple, Optional, Tuple from tensorrt_llm.logger import logger @@ -25,7 +25,7 @@ import nvMatmulHeuristics as _nvmmh IS_NVMMH_AVAILABLE = True -except ImportError: +except (ImportError, OSError): _nvmmh = None IS_NVMMH_AVAILABLE = False @@ -40,84 +40,119 @@ NVFP4_PRECISION = "F4F4T" NVFP4_LAYOUT = "TN_ROW_MAJOR" -CtaCluster = Tuple[Tuple[int, int], Tuple[int, int]] - -# The knobs a caller may let nvMatmulHeuristics drive (instead of sweeping the -# runner's candidate list). A field named here is supplied by the model; a field -# not named is swept. "tile" and "cluster" are a coupled pair (the 2-SM encoding -# maps the joint (cta, cluster)), so selecting either drives both jointly. -# "swizzle" (swizzle_factor) and "cta_order" (rasterization order) are per-knob. -NVMMH_OPTIONAL_FIELDS = frozenset({"tile", "cluster", "swizzle", "cta_order"}) +# Dense BF16 GEMM uses BF16 A/B with FP32 accumulation. The output may be BF16 +# or FP32, which is the third character in nvMatmulHeuristics' three-character +# precision form. The runner stores A as [M, K] and weight as [N, K] while +# consuming weight^T, matching the TN row-major layout used by the NVFP4 path. +BF16_PRECISION = "TST" +BF16_FP32_OUTPUT_PRECISION = "TSS" +BF16_LAYOUT = "TN_ROW_MAJOR" + +# Dense FP8 GEMM uses E4M3 A/B with FP32 accumulation and BF16 or FP16 output. +# nvMatmulHeuristics uses its five-character FP8 form: A, B, C, compute, D. +# The installed Rubin wheel accepts both tokens below. A/B are K-major, matching +# the physical [M, K] activation and [N, K] weight tensors. +FP8_BF16_OUTPUT_PRECISION = "QQTST" +FP8_FP16_OUTPUT_PRECISION = "QQHSH" +FP8_LAYOUT = "TN_ROW_MAJOR" + +# nvMatmulHeuristics often recommends non-power-of-two split counts. The BF16 +# filter accepts model-recommended values through this conservative cap after a +# successful query; get_valid_tactics itself remains unchanged. +BF16_MAX_HEURISTIC_SPLIT_K = 8 + +# Rubin's native split-K scheduler currently requires more than 20 K tiles per +# split. That lower bound hides useful split factors for the relatively small-K +# CuTe DSL kernels used by TRT-LLM. Keep the Rubin-specific override here so +# every split-capable CuTe DSL runner applies the same policy. This is an +# admission rule only: CUPTI still chooses the fastest admitted factor. +SM107_NVMMH_MIN_K_TILES_PER_SPLIT = 6 + + +def is_sm107_nvmmh_split_k_eligible(k: int, cta_k: int, split_k: int) -> bool: + """Whether an SM107 split-K candidate clears TRT-LLM's K-tile bound. + + This mirrors nvMatmulHeuristics' strict ``tiles_per_split > lower_bound`` + gate while replacing its hard-coded lower bound for SM107. Split-K 1 is + always retained so the heuristic path cannot remove the unsplit fallback. + """ + k, cta_k, split_k = int(k), int(cta_k), int(split_k) + if k <= 0 or cta_k <= 0 or split_k <= 0: + return False + if split_k == 1: + return True + tiles_per_split = (k + cta_k * split_k - 1) // (cta_k * split_k) + return tiles_per_split > SM107_NVMMH_MIN_K_TILES_PER_SPLIT + + +# Canonical query mode used by kernels that reduce split-K partials directly +# into the destination tensor. Kept as an adapter-owned string so callers do +# not need to import nvMatmulHeuristics enums (or even have the wheel installed). +IN_PLACE_SPLIT_K_KIND = "in_place" + +_BF16_KNOWN_GOOD_PREFERRED_TACTIC = ( + "preferred_cluster", + True, + (256, 256), + (4, 2), + (2, 1), + 0, +) -# Accepted spellings that normalize onto a canonical field name. -_NVMMH_FIELD_ALIASES = {"cta_tile": "tile", "mma_tiler": "tile"} +CtaCluster = Tuple[Tuple[int, int], Tuple[int, int]] class HeuristicConfig(NamedTuple): - """One nvMatmulHeuristics candidate, reduced to the knobs we can map to a - CuTe DSL NVFP4 tactic. ``swizzle_factor``/``cta_order`` default to the - kernel's neutral values (no swizzle, M-major raster) when absent.""" + """One model candidate reduced to knobs understood by CuTe DSL runners. + + Optional fields use neutral defaults when absent. ``split_k`` is appended + to preserve positional compatibility with the original four-field tuple. + """ cta: Tuple[int, int] cluster: Tuple[int, int] swizzle_factor: int = 1 cta_order: int = 0 + split_k: int = 1 -def nvmmh_fields() -> set: - """Which knobs nvMatmulHeuristics drives (vs. the runner sweeping them). +def nvmmh_cta_order_to_cute_raster(cta_order: int) -> str: + """Translate CUTLASS3 nvMMH CTA order to CuTe's raster name. - Read from TRTLLM_CUTEDSL_NVMMH_FIELDS (comma-separated). Recognized tokens - are ``NVMMH_OPTIONAL_FIELDS`` (plus the ``cta_tile``/``mma_tiler`` aliases - for "tile"); unknown tokens are warned about once and ignored. Because the - 2-SM encoding maps the joint (cta, cluster), "tile" and "cluster" are - coupled -- naming either drives both. Default: "tile,cluster" (today's - behavior: heuristic tile+cluster, swept swizzle/cta_order). + nvMMH's CUTLASS swizzler and DKG integration use order 0 for an + M-fast raster and order 1 for an N-fast raster. Keep this conversion at + the adapter boundary because the public nvMMH header's row/column wording + does not match the CUTLASS3 implementation used to score the configs. """ - raw = os.environ.get("TRTLLM_CUTEDSL_NVMMH_FIELDS", "tile,cluster") - fields = set() - for tok in (f.strip().lower() for f in raw.split(",")): - if not tok: - continue - canonical = _NVMMH_FIELD_ALIASES.get(tok, tok) - if canonical not in NVMMH_OPTIONAL_FIELDS: - logger.warning_once( - f"[nvMatmulHeuristics] Ignoring unknown field '{tok}' in " - f"TRTLLM_CUTEDSL_NVMMH_FIELDS. Recognized: " - f"{sorted(NVMMH_OPTIONAL_FIELDS)} (aliases: " - f"{sorted(_NVMMH_FIELD_ALIASES)}).", - key=f"nvmmh_unknown_field_{tok}", - ) - continue - fields.add(canonical) - # tile/cluster are a coupled pair: naming either drives both. - if fields & {"tile", "cluster"}: - fields.update({"tile", "cluster"}) - return fields - + if cta_order == 0: + return "m" + if cta_order == 1: + return "n" + raise ValueError(f"Unsupported nvMMH CTA order {cta_order!r}; expected 0 or 1") -def nvmmh_enabled() -> bool: - """Master opt-in switch for heuristic tactic pruning.""" - return IS_NVMMH_AVAILABLE and os.environ.get("TRTLLM_CUTEDSL_NVMMH_ENABLE", "0") == "1" +def cute_raster_to_nvmmh_cta_order(raster_along: str) -> int: + """Translate CuTe's M/N-fast raster name to CUTLASS3 nvMMH CTA order.""" + if raster_along == "m": + return 0 + if raster_along == "n": + return 1 + raise ValueError(f"Unsupported CuTe raster {raster_along!r}; expected 'm' or 'n'") -def nvmmh_enabled_for_nvfp4() -> bool: - """Whether heuristic pruning applies to the NVFP4 dense GEMM. Gated solely - by the master switch (TRTLLM_CUTEDSL_NVMMH_ENABLE).""" - return nvmmh_enabled() +def _pin_split_k_to_rank1( + ranked_configs: List[HeuristicConfig], +) -> List[HeuristicConfig]: + """Use rank-1 split-K for every precision while preserving other fields.""" + if not ranked_configs: + return ranked_configs -def nvmmh_max_tactics() -> int: - """Number of heuristic-ranked (tile, cluster) configs to keep per problem. - - Mirrors TorchInductor's nvgemm_max_profiling_configs. Both use_prefetch - variants of a kept config are profiled on top of this (prefetch is not - modeled by the heuristic), so the final tactic count is up to ~2x this. - """ - try: - return max(1, int(os.environ.get("TRTLLM_CUTEDSL_NVMMH_MAX_TACTICS", "5"))) - except ValueError: - return 5 + rank1_split_k = max(1, int(ranked_configs[0].split_k)) + pinned = [ + config if int(config.split_k) == rank1_split_k else config._replace(split_k=rank1_split_k) + for config in ranked_configs + ] + return list(dict.fromkeys(pinned)) @lru_cache(maxsize=None) @@ -182,11 +217,41 @@ def _extract_int(config, kernel, dict_key: str, attr: str, str_re: str, default: return default +def _extract_split_k(config, kernel) -> int: + """Read split-K across the public, raw-ctypes, and printable API forms.""" + candidates = [] + if isinstance(config, dict): + candidates.extend(config.get(key) for key in ("split_k", "splitK", "split_k_slices")) + for attr in ("split_k", "splitK", "split_k_slices"): + if hasattr(kernel, attr): + candidates.append(getattr(kernel, attr)) + if isinstance(config, dict): + raw = config.get("nvmmhKernelConfiguration") + if raw is not None and hasattr(raw, "splitK"): + candidates.append(raw.splitK) + for value in candidates: + try: + value = int(value) + except (TypeError, ValueError): + continue + if value >= 1: + return value + + text = kernel if isinstance(kernel, str) else str(kernel) + match = re.search(r"\bsplitK\(\s*(\d+)\s*\)", text) + if match: + try: + return max(1, int(match.group(1))) + except ValueError: + pass + return 1 + + def _extract_config(config) -> Optional[HeuristicConfig]: - """Pull (cta, cluster, swizzle_factor, cta_order) out of one heuristic - config. Tolerant of the several shapes the config may take across versions: - a dict with flat fields, a dict whose ``kernel`` is an object with - attributes, or a dict whose ``kernel`` is a printable string. + """Extract runner-visible fields from one heuristic configuration. + + Tolerates flat dictionaries, public ``GemmConfig`` objects, raw ctypes + wrappers, and printable kernel strings used across package versions. """ kernel = config.get("kernel") if isinstance(config, dict) else config cta = cluster = None @@ -217,25 +282,35 @@ def _extract_config(config) -> Optional[HeuristicConfig]: return None # swizzle_factor -> "swizz(N)"; cta_order -> "ctaOrder(N)". Default to the - # kernel's neutral values (no swizzle, M-major raster) when the field is + # kernel's neutral values (no swizzle, along-M raster) when the field is # absent so an unselected knob never changes behavior. swizzle = _extract_int( config, kernel, "swizzle_factor", "swizzle_factor", r"swizz\(\s*(\d+)\)", 1 ) cta_order = _extract_int(config, kernel, "cta_order", "cta_order", r"ctaOrder\(\s*(\d+)\)", 0) - return HeuristicConfig(cta, cluster, max(1, swizzle), cta_order) - - -def rank_configs(m: int, n: int, k: int, precision: str, count: int) -> List[HeuristicConfig]: + split_k = _extract_split_k(config, kernel) + return HeuristicConfig(cta, cluster, max(1, swizzle), cta_order, max(1, split_k)) + + +def rank_configs( + m: int, + n: int, + k: int, + precision: str, + count: int, + layout_name: str = NVFP4_LAYOUT, + split_k_kind: Optional[str] = None, +) -> List[HeuristicConfig]: """Return up to ``count`` HeuristicConfig entries ranked best-first. - Returns an empty list if heuristics are unavailable or the query fails, so + Every returned precision uses the first parsed rank's split-K value. Returns + an empty list if heuristics are unavailable or the query fails, so callers can fall back to their full tactic list. """ interface = _get_interface(precision) if interface is None: return [] - layout = _get_layout(NVFP4_LAYOUT) + layout = _get_layout(layout_name) if layout is None: return [] try: @@ -244,11 +319,64 @@ def rank_configs(m: int, n: int, k: int, precision: str, count: int) -> List[Heu interface.loadInternalDiscoverySet(layout, None) except Exception as e: # noqa: BLE001 - discovery data is optional logger.debug(f"[nvMatmulHeuristics] loadInternalDiscoverySet skipped: {e}") - configs = interface.get_with_mnk(int(m), int(n), int(k), layout, int(count), None) + if split_k_kind is None: + configs = interface.get_with_mnk(int(m), int(n), int(k), layout, int(count), None) + else: + if split_k_kind != IN_PLACE_SPLIT_K_KIND: + raise ValueError(f"Unsupported nvMatmulHeuristics split-K kind {split_k_kind!r}") + + def _method(*names): + for name in names: + method = getattr(interface, name, None) + if method is not None: + return method + raise AttributeError(f"nvMatmulHeuristics interface has none of {names!r}") + + backend = None + try: + create_backend = _method("createBackend", "create_backend") + backend = create_backend(_nvmmh.NvMatmulHeuristicsTarget.CUTLASS3) + property_name = _nvmmh.NvMatmulHeuristicsBackendProperty.SPLIT_K_KIND + property_value = int(_nvmmh.NvMatmulHeuristicsSplitKKind.IN_PLACE) + set_backend_property = getattr( + interface, + "setBackendPropertyValue", + getattr(interface, "set_backend_property_value", None), + ) + if set_backend_property is not None: + # Newer wrappers accept the scalar value directly. + set_backend_property(backend, property_name, property_value) + else: + # The wheel bundled in the Rubin container exposes the + # original raw-buffer InterfaceEx spelling. + set_backend_property = _method( + "setBackendValueProperty", "set_backend_value_property" + ) + property_buffer = ctypes.c_int32(property_value) + set_backend_property( + backend, + property_name, + ctypes.byref(property_buffer), + ctypes.sizeof(property_buffer), + ) + make_problem = _method( + "makeNvMatmulHeuristicsProblem", "make_nv_matmul_heuristics_problem" + ) + try: + problem = make_problem((int(m), int(n), int(k)), layout, batch_size=1) + except TypeError: + # Older wheels take M, N, and K as separate arguments. + problem = make_problem(int(m), int(n), int(k), layout, batch_size=1) + get_ex = _method("getEx", "get_ex") + configs = get_ex(problem, int(count), backend) + finally: + if backend is not None: + _method("destroyBackend", "destroy_backend")(backend) except Exception as e: # noqa: BLE001 - any failure must degrade gracefully logger.warning_once( f"[nvMatmulHeuristics] Query failed for " - f"precision={precision}, mnk=({m},{n},{k}): {e}. " + f"precision={precision}, mnk=({m},{n},{k}), " + f"split_k_kind={split_k_kind}: {e}. " f"Falling back to full tactic list.", key="nvmmh_query_failure", ) @@ -261,13 +389,17 @@ def rank_configs(m: int, n: int, k: int, precision: str, count: int) -> List[Heu def _runtime(c): return c.get("runtime", float("inf")) if isinstance(c, dict) else float("inf") - ordered = sorted(configs, key=_runtime) + # getEx already returns a custom backend's ranking. Its convenience + # ``runtime`` value may be estimated with the interface's default backend, + # so sorting it again could reorder the IN_PLACE split-K recommendation + # using OUT_OF_PLACE costs. + ordered = list(configs) if split_k_kind is not None else sorted(configs, key=_runtime) ranked: List[HeuristicConfig] = [] for c in ordered: cfg = _extract_config(c) if cfg is not None and cfg not in ranked: ranked.append(cfg) - return ranked + return _pin_split_k_to_rank1(ranked) def rank_tile_cluster_configs( @@ -275,3 +407,283 @@ def rank_tile_cluster_configs( ) -> List[CtaCluster]: """Back-compat helper: (cta, cluster) pairs only, ranked best-first.""" return [(c.cta, c.cluster) for c in rank_configs(m, n, k, precision, count)] + + +def _fp8_config_signature(config: HeuristicConfig, fields: set) -> tuple: + """Map a model config to fields actionable by dense FP8 runners.""" + signature = [] + if fields & {"tile", "cluster"}: + signature.extend((tuple(config.cta), tuple(config.cluster))) + if "cta_order" in fields: + signature.append(int(config.cta_order)) + return tuple(signature) + + +def filter_fp8_tactics( + tactics: List[tuple], + ranked_configs: List[HeuristicConfig], + fields: set, + max_tactics: int, + tactic_signature: Callable[[tuple, set], Optional[tuple]], + sweep_cluster_n: bool = False, +) -> List[tuple]: + """Keep FP8 tactics matching the top model-visible field signatures. + + Instruction shape, B-reuse, TMA-store mode, and other runner-local fields + stay swept because ``tactic_signature`` deliberately omits them. Rubin may + additionally sweep cluster-N within each selected (CTA, cluster-M) family. + An empty query or intersection returns the original validated tactic list. + """ + if not tactics or not ranked_configs or not fields: + return tactics + + valid_signatures = { + signature + for tactic in tactics + if (signature := tactic_signature(tactic, fields)) is not None + } + kept_signatures = [] + seen = set() + for config in ranked_configs: + signature = _fp8_config_signature(config, fields) + if signature in valid_signatures and signature not in seen: + kept_signatures.append(signature) + seen.add(signature) + if len(kept_signatures) >= max(1, int(max_tactics)): + break + if not kept_signatures: + return tactics + + if sweep_cluster_n and fields & {"tile", "cluster"}: + + def _cluster_n_family(signature): + cta, cluster, *remaining = signature + return (cta, int(cluster[0]), *remaining) + + kept_families = {_cluster_n_family(signature) for signature in kept_signatures} + selected = [] + for tactic in tactics: + signature = tactic_signature(tactic, fields) + if signature is not None and _cluster_n_family(signature) in kept_families: + selected.append(tactic) + else: + kept = set(kept_signatures) + selected = [tactic for tactic in tactics if tactic_signature(tactic, fields) in kept] + return selected if selected else tactics + + +def expand_bf16_tactics_with_heuristic_splits( + tactics: List[tuple], ranked_configs: List[HeuristicConfig] +) -> List[tuple]: + """Add only the rank-1 model split count after a successful BF16 query. + + The runner has already validated each base tactic's tile and cluster. Split + count does not participate in ``can_implement``, so a six-field split tactic + can be derived safely from those validated templates. Baseline tactics are + returned byte-for-byte when the model has no supported split recommendation. + """ + # ``get_valid_tactics`` only emits a split greater than one when the + # runner's split-K shape guard passes. Preserve that eligibility decision + # before adding non-power-of-two model recommendations. + split_k_eligible = any( + isinstance(tactic, tuple) + and len(tactic) == 6 + and tactic[0] == "base" + and int(tactic[5]) > 1 + for tactic in tactics + ) + if not split_k_eligible: + return tactics + + ranked_configs = _pin_split_k_to_rank1(ranked_configs) + recommended = { + int(config.split_k) + for config in ranked_configs + if 1 < int(config.split_k) <= BF16_MAX_HEURISTIC_SPLIT_K + } + if not recommended: + return tactics + + expanded = list(tactics) + seen = set(expanded) + templates = [] + seen_templates = set() + for tactic in tactics: + if not isinstance(tactic, tuple) or len(tactic) < 6 or tactic[0] != "base": + continue + template = (tactic[1], tuple(tactic[2]), tuple(tactic[3]), tactic[4]) + if template not in seen_templates: + templates.append(template) + seen_templates.add(template) + for use_2cta, mma_tiler_mn, cluster_shape_mn, max_num_ab_stage in templates: + for split_k in sorted(recommended): + tactic = ( + "base", + use_2cta, + mma_tiler_mn, + cluster_shape_mn, + max_num_ab_stage, + split_k, + ) + if tactic not in seen: + expanded.append(tactic) + seen.add(tactic) + return expanded + + +def _bf16_actionable_fields(fields: set, split_k: int) -> set: + """Return model fields the selected BF16 kernel path can control. + + Split-K>1 uses the reduction path's legacy six-field tactic and fixed + scheduler defaults, so swizzle and CTA order are not actionable there. + Split-K itself, when selected, remains model-controlled; tile/cluster + fields also remain part of the match when requested. + """ + if split_k > 1: + return fields - {"swizzle", "cta_order"} + return fields + + +def _bf16_config_signature(config: HeuristicConfig, fields: set) -> tuple: + """Map a model config into the BF16 runner's selected-field key.""" + # The model's split recommendation only narrows the actionable fields when + # split-K was selected. Honoring it otherwise strips swizzle/cta_order from + # the key and matches the split-K tactics instead of the scheduler sweep the + # caller asked the model to rank. + config_split_k = max(1, int(config.split_k)) if "split_k" in fields else 1 + fields = _bf16_actionable_fields(fields, config_split_k) + signature = [] + if fields & {"tile", "cluster"}: + cta_m, cta_n = (int(v) for v in config.cta) + cluster_m, cluster_n = (int(v) for v in config.cluster) + # libheuristics encodes its two-CTA MMA with cluster_m == 2 while the + # CuTe DSL runner encodes it by doubling mma_tiler_mn[0]. + n_align = 32 if cta_n > 256 else 16 + use_2cta = cluster_m == 2 and cta_n % n_align == 0 + mma_tiler_mn = ((2 * cta_m if use_2cta else cta_m), cta_n) + signature.extend((use_2cta, mma_tiler_mn, (cluster_m, cluster_n))) + if "split_k" in fields: + signature.append(max(1, int(config.split_k))) + if "swizzle" in fields: + signature.append(max(1, int(config.swizzle_factor))) + if "cta_order" in fields: + signature.append(int(config.cta_order)) + return tuple(signature) + + +def _bf16_tactic_signature(tactic: tuple, fields: set) -> Optional[tuple]: + """Map a base BF16 tactic into the same selected-field key as the model.""" + if not isinstance(tactic, tuple) or len(tactic) < 6 or tactic[0] != "base": + return None + split_k = max(1, int(tactic[5])) + fields = _bf16_actionable_fields(fields, split_k) + signature = [] + if fields & {"tile", "cluster"}: + signature.extend((bool(tactic[1]), tuple(tactic[2]), tuple(tactic[3]))) + if "split_k" in fields: + signature.append(split_k) + if "swizzle" in fields: + # Split-K tactics use the kernel scheduler defaults and therefore keep + # their legacy six-field form. + signature.append(max(1, int(tactic[7])) if len(tactic) > 7 else 1) + if "cta_order" in fields: + raster_along = tactic[6] if len(tactic) > 6 else "m" + signature.append(cute_raster_to_nvmmh_cta_order(raster_along)) + # An empty key carries no information, so it would match every config the + # model returns. Report the tactic as unmappable instead: a split-K tactic + # has no actionable scheduler knob, and callers already sweep unmatched + # tactics via the fallback path. + return tuple(signature) if signature else None + + +def _bf16_structural_tactic_signature(tactic: tuple, fields: set) -> Optional[tuple]: + """Return the model key after removing directly-applied scheduler knobs.""" + if not isinstance(tactic, tuple) or len(tactic) < 6 or tactic[0] != "base": + return None + split_k = max(1, int(tactic[5])) + actionable_fields = _bf16_actionable_fields(fields, split_k) + structural_fields = actionable_fields - {"swizzle", "cta_order"} + if not structural_fields: + # Scheduler-only guidance applies to the non-split base path. Split-K + # kernels have fixed scheduler defaults and must stay unannotated. + return () if split_k == 1 and actionable_fields else None + return _bf16_tactic_signature(tactic, structural_fields) + + +def _bf16_structural_config_signature(config: HeuristicConfig, fields: set) -> tuple: + """Return the config key used before scheduler values are materialized.""" + config_split_k = max(1, int(config.split_k)) if "split_k" in fields else 1 + actionable_fields = _bf16_actionable_fields(fields, config_split_k) + structural_fields = actionable_fields - {"swizzle", "cta_order"} + return _bf16_config_signature(config, structural_fields) + + +def _apply_bf16_scheduler_config(tactic: tuple, config: HeuristicConfig, fields: set) -> tuple: + """Copy selected scheduler values from one NVMMH config to a base tactic.""" + split_k = max(1, int(tactic[5])) + actionable_fields = _bf16_actionable_fields(fields, split_k) + scheduler_fields = actionable_fields & {"swizzle", "cta_order"} + if not scheduler_fields: + return tactic + raster_order = ( + nvmmh_cta_order_to_cute_raster(config.cta_order) if "cta_order" in scheduler_fields else "m" + ) + swizzle_size = max(1, int(config.swizzle_factor)) if "swizzle" in scheduler_fields else 1 + return tactic[:6] + (raster_order, swizzle_size, True) + + +def filter_bf16_tactics( + tactics: List[tuple], + ranked_configs: List[HeuristicConfig], + fields: set, + max_tactics: int, + include_known_good: bool = True, + fallback_tactics: Optional[List[tuple]] = None, +) -> List[tuple]: + """Return top model-matched BF16 tactics with direct scheduler annotation. + + Only base tactics participate in model matching. Valid preferred/fallback + cluster schedulers cannot be represented directly by the model's one cluster + field, so only the validated q_b known-good tactic is supplemented when its + implicit split-K=1 agrees with rank 1. When split-K is selected, every lower + rank keeps its other fields but uses the rank-1 split count. An empty + intersection returns ``fallback_tactics`` (normally the exact list produced + by ``get_valid_tactics``). + """ + fallback = tactics if fallback_tactics is None else fallback_tactics + if not tactics or not ranked_configs or not fields: + return fallback + allow_known_good_supplement = include_known_good + if "split_k" in fields: + ranked_configs = _pin_split_k_to_rank1(ranked_configs) + allow_known_good_supplement = include_known_good and int(ranked_configs[0].split_k) == 1 + + valid_signatures = { + signature + for tactic in tactics + if (signature := _bf16_structural_tactic_signature(tactic, fields)) is not None + } + config_by_signature = {} + for config in ranked_configs: + signature = _bf16_structural_config_signature(config, fields) + if signature in valid_signatures and signature not in config_by_signature: + config_by_signature[signature] = config + if len(config_by_signature) >= max(1, int(max_tactics)): + break + if not config_by_signature: + return fallback + + selected = [] + for tactic in tactics: + signature = _bf16_structural_tactic_signature(tactic, fields) + config = config_by_signature.get(signature) + if config is None: + continue + selected.append(_apply_bf16_scheduler_config(tactic, config, fields)) + if ( + allow_known_good_supplement + and _BF16_KNOWN_GOOD_PREFERRED_TACTIC in tactics + and _BF16_KNOWN_GOOD_PREFERRED_TACTIC not in selected + ): + selected.append(_BF16_KNOWN_GOOD_PREFERRED_TACTIC) + return selected if selected else fallback diff --git a/tensorrt_llm/_torch/custom_ops/torch_custom_ops.py b/tensorrt_llm/_torch/custom_ops/torch_custom_ops.py index d7a1c4714137..f258d9c307dc 100644 --- a/tensorrt_llm/_torch/custom_ops/torch_custom_ops.py +++ b/tensorrt_llm/_torch/custom_ops/torch_custom_ops.py @@ -58,8 +58,8 @@ next_positive_power_of_2) if IS_CUTLASS_DSL_AVAILABLE: - from tensorrt_llm._torch.custom_ops.cute_dsl_custom_ops import \ - CuteDSLNVFP4BlackwellRunner + from tensorrt_llm._torch.custom_ops.cute_dsl_custom_ops import ( + CuteDSLNVFP4BlackwellRunner, _cutedsl_nvmmh_tactic_search_cache_key) if IS_CUTLASS_DSL_RUBIN_AVAILABLE: from tensorrt_llm._torch.custom_ops.cute_dsl_custom_ops import \ @@ -1508,6 +1508,18 @@ def forward( else: raise ValueError(f"Invalid tactic: {tactic}") + def use_torch_profiler(self, tactic) -> bool: + # Keep one timing backend for every candidate in a mixed-backend sweep. + return "cutedsl" in self.allowed_backends + + def profiling_timer_cache_key(self) -> Optional[str]: + return "torch_profiler" if "cutedsl" in self.allowed_backends else None + + def tactic_search_cache_key(self): + if "cutedsl" not in self.allowed_backends or not IS_CUTLASS_DSL_AVAILABLE: + return None + return _cutedsl_nvmmh_tactic_search_cache_key() + @fast_custom_op("trtllm::nvfp4_gemm", mutates_args=()) def nvfp4_gemm( diff --git a/tensorrt_llm/_torch/kernels/rubin/dense_bf16_gemm_persistent.py b/tensorrt_llm/_torch/kernels/rubin/dense_bf16_gemm_persistent.py index 930f8538bb72..76efbfdf1952 100644 --- a/tensorrt_llm/_torch/kernels/rubin/dense_bf16_gemm_persistent.py +++ b/tensorrt_llm/_torch/kernels/rubin/dense_bf16_gemm_persistent.py @@ -122,6 +122,9 @@ def _kernel_tactic_name(self) -> str: f"_stage{self._max_num_ab_stage}" f"_split{self.split_k_slices}" f"_direct{int(self.use_direct_split_k_reduce)}" + f"_raster{self.raster_along}" + f"_swizzle{self.swizzle_size}" + f"_tmastore{int(self.use_tma_store)}" ) def __repr__(self) -> str: @@ -298,6 +301,9 @@ def _kernel_tactic_name(self) -> str: f"_preferred{self._format_mn_pair(self.preferred_cluster_shape_mn)}" f"_fallback{self._format_mn_pair(self.fallback_cluster_shape_mn)}" f"_stage{self._max_num_ab_stage}" + f"_raster{self.raster_along}" + f"_swizzle{self.swizzle_size}" + f"_tmastore{int(self.use_tma_store)}" ) def _setup_attributes(self): diff --git a/tensorrt_llm/_torch/locality_domain/autotune.py b/tensorrt_llm/_torch/locality_domain/autotune.py index 0f446e02dbb0..a215f404b38c 100644 --- a/tensorrt_llm/_torch/locality_domain/autotune.py +++ b/tensorrt_llm/_torch/locality_domain/autotune.py @@ -75,6 +75,10 @@ def get_valid_tactics( ) -> list[Any]: return self._op_runner.get_valid_tactics(inputs, profile, **kwargs) + def tactic_search_cache_key(self) -> tuple[Any, ...] | None: + """Use the wrapped runner's candidate-search discriminator.""" + return self._op_runner.tactic_search_cache_key() + def should_profile_tactic_in_subprocess( self, custom_op: str, diff --git a/tensorrt_llm/_torch/moe/fused_moe/fused_moe_cute_dsl.py b/tensorrt_llm/_torch/moe/fused_moe/fused_moe_cute_dsl.py index 91879d73d94d..b0a60d6ae2c9 100644 --- a/tensorrt_llm/_torch/moe/fused_moe/fused_moe_cute_dsl.py +++ b/tensorrt_llm/_torch/moe/fused_moe/fused_moe_cute_dsl.py @@ -24,8 +24,9 @@ from tensorrt_llm._utils import get_sm_version, is_sm_100f from tensorrt_llm.models.modeling_utils import QuantAlgo -from ...autotuner import (AutoTuner, ConstraintSpec, DynamicTensorSpec, - OptimizationProfile, TunableRunner, TuningConfig) +from ...autotuner import (AutoTuner, ConstraintSpec, CuteDSLTunableRunner, + DynamicTensorSpec, OptimizationProfile, TunableRunner, + TuningConfig) from ...custom_ops.cute_dsl_custom_ops import GroupedGemmInputsHelper from ...cute_dsl_utils import (IS_CUTLASS_DSL_AVAILABLE, IS_CUTLASS_DSL_RUBIN_AVAILABLE) @@ -331,7 +332,7 @@ def inputs_pre_hook(self, inputs: List[torch.Tensor]) -> List[torch.Tensor]: return x, new_token_selected_experts, *others -class CuteDslFusedMoENvfp4Runner(TunableRunner): +class CuteDslFusedMoENvfp4Runner(CuteDSLTunableRunner): tuning_config_cache = dict() def __init__(self, @@ -555,7 +556,7 @@ def inputs_pre_hook(self, inputs: List[torch.Tensor]) -> List[torch.Tensor]: return x, new_token_selected_experts, *others -class CuteDslFusedMoEBF16Runner(TunableRunner): +class CuteDslFusedMoEBF16Runner(CuteDSLTunableRunner): """Autotuner runner for BF16/FP16 MoE on Rubin (SM107). Selects tile_size from {64, 128, 256} and delegates to run_moe_bf16_impl. diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py index 221f9145b17c..43ee5dd59a01 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_engine.py +++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py @@ -32,6 +32,7 @@ from ..attention.backends.interface import (AttentionMetadata, AttentionRuntimeFeatures) from ..attention.backends.utils import get_attention_backend +from ..autotuner import AutoTuner, NvMMHConfig from ..compilation.backend import Backend from ..distributed import Distributed from ..distributed.communicator import init_pp_comm @@ -161,6 +162,20 @@ def _configure_deep_gemm_pdl() -> None: _DEEP_GEMM_PDL_CONFIGURED = True +def _configure_autotuner_nvmmh(llm_args: TorchLlmArgs) -> None: + """Install the per-worker NVMMH policy before tactic enumeration.""" + config = llm_args.autotuner_nvmmh_config + if config is None: + policy = NvMMHConfig() + else: + policy = NvMMHConfig( + enabled=True, + fields=config.fields, + max_tactics=config.max_tactics, + ) + AutoTuner.get().configure_nvmmh(policy) + + class PyTorchModelEngine(ModelEngine): def __init__( @@ -177,6 +192,10 @@ def __init__( model_weights_memory_tag: Optional[str] = None, model_weights_restore_mode=None, ): + # Every executor worker receives its own TorchLlmArgs copy. Install the + # policy before model loading because CuTe DSL runners may enumerate + # tactics while modules are being constructed. + _configure_autotuner_nvmmh(llm_args) _configure_deep_gemm_pdl() self._metrics: dict[str, float] = defaultdict(float) diff --git a/tensorrt_llm/llmapi/__init__.py b/tensorrt_llm/llmapi/__init__.py index 8d1dfd245852..2000ab7de133 100644 --- a/tensorrt_llm/llmapi/__init__.py +++ b/tensorrt_llm/llmapi/__init__.py @@ -9,8 +9,9 @@ from ..scheduling_params import SchedulingParams from .llm import LLM, RequestOutput # yapf: disable -from .llm_args import (AttentionDpConfig, AutoDecodingConfig, BatchingType, - BlockReuseConfig, CacheTransceiverConfig, CalibConfig, +from .llm_args import (AttentionDpConfig, AutoDecodingConfig, + AutoTunerNvMMHConfig, BatchingType, BlockReuseConfig, + CacheTransceiverConfig, CalibConfig, CapacitySchedulerPolicy, ColdPageQuantizationCompressionConfig, ContextChunkingPolicy, CudaGraphConfig, @@ -90,6 +91,7 @@ 'LlmArgs', 'TorchLlmArgs', 'AutoDecodingConfig', + 'AutoTunerNvMMHConfig', 'AttentionDpConfig', 'LoRARequest', 'SaveHiddenStatesDecodingConfig', diff --git a/tensorrt_llm/llmapi/llm_args.py b/tensorrt_llm/llmapi/llm_args.py index 55a23ac1a87b..be040cefab47 100644 --- a/tensorrt_llm/llmapi/llm_args.py +++ b/tensorrt_llm/llmapi/llm_args.py @@ -6346,6 +6346,34 @@ def _check_removed_sampler_type(value: Any) -> None: TORCH_LLMARGS_REMOVED_ARGS = frozenset(_TORCH_LLMARGS_REMOVED_ARG_CHECKS) +class AutoTunerNvMMHConfig(StrictBaseModel): + """Enabled nvMatmulHeuristics policy for CuTe DSL search reduction.""" + + model_config = ConfigDict(extra="forbid", frozen=True) + + fields: Tuple[Literal[ + "tile", "cluster", "swizzle", "cta_order", "split_k"], ...] = Field( + default=("tile", "cluster"), + description= + "Tactic fields selected by NVMMH. Tile and cluster are coupled.") + + max_tactics: PositiveInt = Field( + default=5, + strict=True, + description= + "Maximum number of NVMMH-ranked tile/cluster signatures to retain.") + + @field_validator('fields') + @classmethod + def normalize_fields(cls, fields: Tuple[str, ...]) -> Tuple[str, ...]: + selected = set(fields) + if selected & {"tile", "cluster"}: + selected.update({"tile", "cluster"}) + return tuple(field for field in ("tile", "cluster", "swizzle", + "cta_order", "split_k") + if field in selected) + + class TorchLlmArgs(BaseLlmArgs): # PyTorch backend specific configurations generation_config: Literal["auto", "trtllm"] = Field( @@ -6700,6 +6728,12 @@ def validate_mla_skip_correction_config(self) -> 'TorchLlmArgs': "Enable autotuner for all tunable ops. This flag is for debugging purposes only, and the performance may significantly degrade if set to false.", status="prototype") + autotuner_nvmmh_config: Optional[AutoTunerNvMMHConfig] = Field( + default=None, + description= + "nvMatmulHeuristics policy used to reduce the CuTe DSL autotuner search space. A mapping enables NVMMH; null disables it.", + status="prototype") + use_fine_grained_sync: bool = Field( default=False, description= diff --git a/tensorrt_llm/usage/llm_args_golden_manifest.json b/tensorrt_llm/usage/llm_args_golden_manifest.json index e58f42598f2d..3a8eb75557f8 100644 --- a/tensorrt_llm/usage/llm_args_golden_manifest.json +++ b/tensorrt_llm/usage/llm_args_golden_manifest.json @@ -91,6 +91,23 @@ "kind": "categorical", "path": "attn_backend" }, + { + "allowed_values": [ + "tile", + "cluster", + "swizzle", + "cta_order", + "split_k" + ], + "capture_policy": "tuple[literal]", + "kind": "categorical", + "path": "autotuner_nvmmh_config.fields" + }, + { + "capture_policy": "int", + "kind": "value", + "path": "autotuner_nvmmh_config.max_tactics" + }, { "allowed_values": [ "pytorch" diff --git a/tests/unittest/_torch/kernels/parallel/test_fp8_block_scale_gemm.py b/tests/unittest/_torch/kernels/parallel/test_fp8_block_scale_gemm.py index 9567164ca7d3..af11313f1751 100644 --- a/tests/unittest/_torch/kernels/parallel/test_fp8_block_scale_gemm.py +++ b/tests/unittest/_torch/kernels/parallel/test_fp8_block_scale_gemm.py @@ -696,6 +696,7 @@ def test_cute_dsl_mxfp8_gemm_rubin_mixed_clusters_multi_wave( "static", "m", 1, + 4, ) mixed_tactic = ( "mixed_clusters", @@ -706,13 +707,14 @@ def test_cute_dsl_mxfp8_gemm_rubin_mixed_clusters_multi_wave( False, use_prefetch, "static", - "m", + "n", + 2, ) runner = CuteDSLMXFP8RubinLinear(output_dtype=torch.bfloat16, use_tvm_ffi=True) tactics = runner.get_valid_tactics(inputs, None) - assert base_tactic in tactics - assert mixed_tactic in tactics + assert base_tactic[:-1] in tactics + assert mixed_tactic[:-1] in tactics base_output = runner(inputs, tactic=base_tactic) mixed_output = runner(inputs, tactic=mixed_tactic) @@ -758,17 +760,19 @@ def test_cute_dsl_mxfp8_gemm_rubin_mixed_clusters_clc_dynamic_prefetch_multi_wav True, "static", "m", + 1, ) - dynamic_tactic = (*static_tactic[:-2], "clc_dynamic", "m") + dynamic_tactic = (*static_tactic[:-3], "clc_dynamic", "m", 1) runner = CuteDSLMXFP8RubinLinear(output_dtype=torch.bfloat16, use_tvm_ffi=True) valid_tactics = runner.get_valid_tactics(inputs, None) + # Model-selected swizzles extend the validated baseline tactic tuple. for raster_order in ("m", "n"): - tactic = (*static_tactic[:-2], "static", raster_order) - assert tactic in valid_tactics - assert dynamic_tactic in valid_tactics - unsafe_dynamic_tactic = (*static_tactic[:-2], "clc_dynamic", "n") - assert unsafe_dynamic_tactic not in valid_tactics + tactic = (*static_tactic[:-3], "static", raster_order, 1) + assert tactic[:-1] in valid_tactics + assert dynamic_tactic[:-1] in valid_tactics + unsafe_dynamic_tactic = (*static_tactic[:-3], "clc_dynamic", "n", 1) + assert unsafe_dynamic_tactic[:-1] not in valid_tactics static_output = runner(inputs, tactic=static_tactic) dynamic_output = runner(inputs, tactic=dynamic_tactic) @@ -895,6 +899,7 @@ def test_cute_dsl_mxfp8_gemm_rubin_clc_dynamic_prefetch_multi_wave(): "static", "m", 1, + 4, ) dynamic_tactic = ( "base", @@ -906,31 +911,53 @@ def test_cute_dsl_mxfp8_gemm_rubin_clc_dynamic_prefetch_multi_wave(): "clc_dynamic", "n", 1, + 1, ) - split_k_tactic = (*dynamic_tactic[:-1], 4) + split_k_tactic = (*dynamic_tactic[:-2], 4, 1) + non_power_split_k_tactic = (*dynamic_tactic[:-2], 3, 1) runner = CuteDSLMXFP8RubinLinear(output_dtype=torch.bfloat16, use_tvm_ffi=True) valid_tactics = runner.get_valid_tactics(inputs, None) - assert static_tactic in valid_tactics - assert dynamic_tactic in valid_tactics - assert split_k_tactic in valid_tactics + # Model-selected swizzles extend the validated baseline tactic tuple. + assert static_tactic[:-1] in valid_tactics + assert dynamic_tactic[:-1] in valid_tactics + assert split_k_tactic[:-1] in valid_tactics for scheduler_mode in ("static", "clc_dynamic"): for raster_order in ("m", "n"): - tactic = (*static_tactic[:-3], scheduler_mode, raster_order, 1) - assert tactic in valid_tactics + swizzle_sizes = ((1, 2, 4, 8) if scheduler_mode == "static" else + (1, )) + for swizzle_size in swizzle_sizes: + tactic = ( + *static_tactic[:-4], + scheduler_mode, + raster_order, + 1, + swizzle_size, + ) + assert tactic[:-1] in valid_tactics static_output = runner(inputs, tactic=static_tactic) dynamic_output = runner(inputs, tactic=dynamic_tactic) split_k_output = runner(inputs, tactic=split_k_tactic) + non_power_split_k_output = runner(inputs, tactic=non_power_split_k_tactic) + non_power_split_k_replay = runner(inputs, tactic=non_power_split_k_tactic) torch.cuda.synchronize() expected = a @ b.t() assert torch.isfinite(dynamic_output).all() assert torch.isfinite(split_k_output).all() + assert torch.isfinite(non_power_split_k_output).all() + assert torch.isfinite(non_power_split_k_replay).all() assert calc_diff(static_output, expected) < 1e-3 assert calc_diff(dynamic_output, expected) < 1e-3 assert calc_diff(split_k_output, expected) < 1e-3 + assert calc_diff(non_power_split_k_output, expected) < 1e-3 + assert calc_diff(non_power_split_k_replay, expected) < 1e-3 torch.testing.assert_close(dynamic_output, static_output, rtol=0, atol=0) + torch.testing.assert_close(non_power_split_k_replay, + non_power_split_k_output, + rtol=0, + atol=0) @pytest.mark.skipif( diff --git a/tests/unittest/_torch/misc/test_autotuner.py b/tests/unittest/_torch/misc/test_autotuner.py index 7b7d44422ff6..05910e78f479 100644 --- a/tests/unittest/_torch/misc/test_autotuner.py +++ b/tests/unittest/_torch/misc/test_autotuner.py @@ -55,6 +55,16 @@ pytestmark = pytest.mark.threadleak(enabled=False) +@pytest.fixture +def nvmmh_config_guard(): + tuner = AutoTuner.get() + previous = tuner.nvmmh_config + try: + yield tuner + finally: + tuner.configure_nvmmh(previous) + + def test_multi_dynamic_dims(): tuner = autotuner.AutoTuner() x = torch.rand([5, 1024]) @@ -890,6 +900,45 @@ def test_kernel_testing_single_context(): f"Expected 3 tactics to be tested, got {len(tested_tactics)}" +def test_kernel_testing_uses_runner_valid_tactics(): + + class CaptureRunner(TunableRunner): + + def get_valid_tactics(self, inputs: List[FakeTensor], + profile: OptimizationProfile, + **kwargs) -> List[int]: + return [1, 2] + + def forward(self, + /, + inputs: List[torch.Tensor], + *, + tactic: int = -1, + **kwargs) -> torch.Tensor: + assert tactic in [-1, 0, 1, 2] + return inputs[0] + + x = torch.randn(4, 4) + runners = [CaptureRunner()] + tuning_config = TuningConfig() + tuner = AutoTuner.get() + tuner.clear_cache() + + with tuner.capture() as all_tactics: + tuner.choose_one("test_runner_valid_capture", runners, tuning_config, + [x]) + + captured = list(all_tactics) + assert [tactic for ((_, tactic), ) in captured] == [1, 2] + + for ((runner, tactic), ) in captured: + with tuner.replay(((runner, tactic), )): + replay_runner, replay_tactic = tuner.choose_one( + "test_runner_valid_capture", runners, tuning_config, [x]) + assert replay_runner is runner + assert replay_tactic == tactic + + class MultiContextRunner(TunableRunner): def get_valid_tactics(self, inputs: List[FakeTensor], @@ -1513,48 +1562,471 @@ def forward(self, ), (f"a failed single-pair candidate must be logged as inf; got {lines}") -def test_cutedsl_nvfp4_heuristic_matches_full_sweep(monkeypatch): - """End-to-end guard for the nvMatmulHeuristics tactic pruning. +_CUPTI_PREFLIGHT_STATE = {} +_CUTE_DSL_NVMMH_TEST_MNK = (16, 256, 7168) - For one representative problem size, the tactic the AutoTuner selects when - nvMatmulHeuristics prunes the tile/cluster candidates must be no slower (up - to a small tolerance) than the tactic it selects from the full CuteDSL - NVFP4 tactic sweep. This validates that pruning does not cost performance. - It additionally compares the heuristic-chosen CuteDSL kernel against the - cuBLASLt NVFP4 GEMM on the same fp4 inputs: the CuteDSL kernel-only device - time must be within a small tolerance of cuBLAS. Both kernel symbol names - are logged for manual inspection -- cuBLAS exposes no API for its selected - kernel's CTA tile / cluster shape, so that is not asserted. +def _skip_if_cupti_unavailable(error, context): + """Skip CUPTI-only perf tests without hiding kernel/runtime failures.""" + if not AutoTuner._is_torch_profiler_unavailable_error(error): + raise error + reason = f"{context} requires working torch.profiler/CUPTI: {error}" + _CUPTI_PREFLIGHT_STATE[torch.cuda.current_device()] = reason + pytest.skip(reason) - Note: the sweep and heuristic paths do NOT profile identical candidate sets - -- the heuristic path is a strict validated subset of the sweep -- so the - exact winning tactic can differ. Only the achieved runtime is a meaningful - invariant, hence the tolerance comparisons. - Blackwell only (SM100/SM103) and requires the nvMatmulHeuristics library; - skipped otherwise. - """ - if not torch.cuda.is_available(): - pytest.skip("requires a CUDA device") +def _require_cupti_cuda_activity(): + """Require torch.profiler to report positive CUDA activity on this GPU.""" + device_index = torch.cuda.current_device() + if AutoTuner._torch_profiler_unavailable: + pytest.skip( + "CUPTI-only performance comparison cannot use the process-wide " + "CUDA-event fallback") + + cached_state = _CUPTI_PREFLIGHT_STATE.get(device_index) + if cached_state is True: + return + if isinstance(cached_state, str): + pytest.skip(cached_state) + + # Warm CUDA before starting the profiler so the probe tests CUPTI activity, + # not first-use CUDA initialization. A CUDA/kernel failure here is real and + # intentionally remains a test failure. + probe = torch.ones(1, device="cuda") + probe.add_(1) + torch.cuda.synchronize() try: - from tensorrt_llm._utils import get_sm_version - sm_version = get_sm_version() - except Exception: - sm_version = None - if sm_version not in (100, 103): - pytest.skip("CuteDSL NVFP4 requires SM100 (B200) / SM103 (B300)") + with torch.profiler.profile( + activities=[torch.profiler.ProfilerActivity.CUDA], + acc_events=True, + record_shapes=False, + profile_memory=False, + with_stack=False, + ) as prof: + probe.add_(1) + torch.cuda.synchronize() + AutoTuner._torch_profiler_elapsed_time_ms(prof, repeat=1) + except Exception as error: + _skip_if_cupti_unavailable(error, "CuTe DSL performance test") + + _CUPTI_PREFLIGHT_STATE[device_index] = True + + +def _require_cupti_autotuner_entry(tuner, name, expected_candidates): + """Return a measured cache entry only when it was timed with CUPTI.""" + assert len(expected_candidates) > 1, ( + "CUPTI cache-time comparison requires multiple candidates; the " + "AutoTuner single-pair shortcut records an unmeasured 0.0 entry") + cache_entries = tuner.profiling_cache.get_specific_custom_op(name) + assert len(cache_entries) == 1 + cache_key, cache_value = next(iter(cache_entries.items())) + _, cached_tactic, profile_ms = cache_value + timer_key = cache_key[-1] + if (AutoTuner._torch_profiler_unavailable + or timer_key == "cuda_event_fallback"): + pytest.skip( + f"CUPTI became unavailable while profiling {name}; refusing to " + "compare CUDA-event fallback timings") + assert timer_key == "torch_profiler", ( + f"Expected CUPTI timing for {name}, got timer key {timer_key!r}") + assert math.isfinite(profile_ms) and profile_ms > 0.0, ( + f"Expected a measured positive CUPTI time for {name}, got {profile_ms}") + return cached_tactic, profile_ms * 1000.0 + + +def _choose_cupti_autotuner_tactic(tuner, name, runner, tuning_config, inputs, + expected_candidates): + """Choose one tactic and return its total CUPTI time in microseconds.""" + assert len(expected_candidates) > 1, ( + "performance comparison requires a measured tactic sweep") + tuner.clear_cache() + with autotune(skip_dynamic_tuning_buckets=True): + _, tactic = tuner.choose_one(name, [runner], tuning_config, inputs) + cached_tactic, profile_us = _require_cupti_autotuner_entry( + tuner, name, expected_candidates) + assert cached_tactic == tactic + return tactic, profile_us - from tensorrt_llm._torch.custom_ops import \ - cutedsl_matmul_heuristics as nvmmh - if not nvmmh.IS_NVMMH_AVAILABLE: - pytest.skip("nvMatmulHeuristics library not installed") - from tensorrt_llm._torch.custom_ops.cute_dsl_custom_ops import \ - CuteDSLNVFP4BlackwellRunner - # One representative square problem. fp4 packing / scale-factor layout - # follows shmoo_nvfp4_cutedsl_heuristics.py::_quantize_inputs. - m = n = k = 4096 +def _profile_cupti_total_us(fn, context, warmup=20, iterations=200, trials=5): + """Return median per-call device time, including every CUDA activity.""" + for _ in range(warmup): + fn() + torch.cuda.synchronize() + + samples = [] + for _ in range(trials): + try: + with torch.profiler.profile( + activities=[torch.profiler.ProfilerActivity.CUDA], + acc_events=True, + record_shapes=False, + profile_memory=False, + with_stack=False, + ) as prof: + for _ in range(iterations): + fn() + torch.cuda.synchronize() + elapsed_ms = AutoTuner._torch_profiler_elapsed_time_ms( + prof, iterations) + except Exception as error: + _skip_if_cupti_unavailable(error, context) + samples.append(elapsed_ms * 1000.0) + return statistics.median(samples) + + +def _nvmmh_tactic_family(tactic): + # Scheduler-only guidance annotates validated families with model-selected + # raster/swizzle values; the resulting tuples need not occur in the sweep. + if tactic[0] == "mixed_clusters": + return tactic[:7] + if tactic[0] == "base": + return tactic[:5] if isinstance(tactic[1], bool) else tactic[:6] + return tactic[:5] + + +def _run_cute_dsl_bf16_heuristic_comparison(monkeypatch, tuner, nvmmh): + """Compare BF16 full sweep, scheduler-only NVMMH, and cuBLASLt.""" + from tensorrt_llm._torch.custom_ops import cute_dsl_custom_ops + + m, n, k = _CUTE_DSL_NVMMH_TEST_MNK + torch.manual_seed(2028) + act = torch.randn(m, k, dtype=torch.bfloat16, device="cuda") + weight = torch.randn(n, k, dtype=torch.bfloat16, device="cuda") + output = torch.empty(m, n, dtype=torch.bfloat16, device="cuda") + inputs = [act, weight, output] + + runner = cute_dsl_custom_ops.CuteDSLBf16RubinGemmRunner( + use_tvm_ffi=True, output_dtype=torch.bfloat16) + runner.__class__.kernel_cache.clear() + runner.__class__.split_k_gemm_cache.clear() + tuning_config = runner.__class__.tuning_config + monkeypatch.setattr(tuning_config, "use_cuda_graph", False) + + scheduler_fields = ("swizzle", "cta_order", "split_k") + tuner.configure_nvmmh(enabled=False, fields=scheduler_fields, max_tactics=5) + baseline_tactics = runner.get_valid_tactics(inputs, None) + assert baseline_tactics + sweep_tactic, sweep_us = _choose_cupti_autotuner_tactic( + tuner, + "test::cute_dsl_bf16_exact_full_sweep", + runner, + tuning_config, + inputs, + baseline_tactics, + ) + + tuner.configure_nvmmh(enabled=True, fields=scheduler_fields, max_tactics=5) + heuristic_candidates = runner.get_valid_tactics(inputs, None) + assert 0 < len(heuristic_candidates) <= len(baseline_tactics), ( + "nvMMH did not prune SM107 BF16 tactics: " + f"{len(heuristic_candidates)} vs {len(baseline_tactics)}") + assert {_nvmmh_tactic_family(t) + for t in heuristic_candidates + }.issubset({_nvmmh_tactic_family(t) + for t in baseline_tactics}) + + model_configs = nvmmh.rank_configs( + m, + n, + k, + nvmmh.BF16_PRECISION, + max(tuner.nvmmh_config.max_tactics * 16, 64), + layout_name=nvmmh.BF16_LAYOUT, + ) + assert model_configs + native_rank1_split_k = int(model_configs[0].split_k) + + def _split_k(tactic): + if (isinstance(tactic, tuple) and tactic and tactic[0] == "base" + and len(tactic) >= 6): + return int(tactic[5]) + return 1 + + candidate_splits = {_split_k(tactic) for tactic in heuristic_candidates} + expected_local_splits = { + split_k + for split_k in (1, 2, 4, 8) if nvmmh.is_sm107_nvmmh_split_k_eligible( + k, runner.nvmmh_split_k_cta_k, split_k) + } + assert expected_local_splits.issubset(candidate_splits), ( + "SM107 BF16 scheduler-only filtering removed locally admitted " + f"split-K candidates: expected={expected_local_splits}, " + f"actual={candidate_splits}") + assert all( + nvmmh.is_sm107_nvmmh_split_k_eligible(k, runner.nvmmh_split_k_cta_k, + split_k) + for split_k in candidate_splits), ( + f"SM107 BF16 retained an ineligible split: {candidate_splits}") + baseline_base_families = { + tactic[:5] + for tactic in baseline_tactics + if isinstance(tactic, tuple) and tactic and tactic[0] == "base" + } + heuristic_base_families = { + tactic[:5] + for tactic in heuristic_candidates + if isinstance(tactic, tuple) and tactic and tactic[0] == "base" + } + assert heuristic_base_families == baseline_base_families, ( + "scheduler-only BF16 NVMMH must keep every base tile/cluster family") + heuristic_tactic, heuristic_us = _choose_cupti_autotuner_tactic( + tuner, + "test::cute_dsl_bf16_scheduler_only_heuristic", + runner, + tuning_config, + inputs, + heuristic_candidates, + ) + + def _cublaslt_call(): + return torch.ops.trtllm.cublas_mm( + act, + weight.t(), + bias=None, + out_dtype=None, + ) + + output.fill_(torch.nan) + runner(inputs, tactic=sweep_tactic) + sweep_output = output.clone() + output.fill_(torch.nan) + runner(inputs, tactic=heuristic_tactic) + heuristic_output = output.clone() + cublaslt_output = _cublaslt_call() + torch.cuda.synchronize() + reference = act.float() @ weight.t().float() + torch.testing.assert_close(sweep_output.float(), + reference, + rtol=1e-2, + atol=1.0) + torch.testing.assert_close(heuristic_output.float(), + reference, + rtol=1e-2, + atol=1.0) + torch.testing.assert_close(cublaslt_output.float(), + reference, + rtol=1e-2, + atol=1.0) + + # Use one CUPTI whole-call timer for the three selected implementations. + # This counts the reduction kernel when a CuTe DSL tactic uses split-K. + sweep_cupti_us = _profile_cupti_total_us( + lambda: runner(inputs, tactic=sweep_tactic), + "BF16 full-sweep tactic comparison", + ) + heuristic_cupti_us = _profile_cupti_total_us( + lambda: runner(inputs, tactic=heuristic_tactic), + "BF16 scheduler-only nvMMH tactic comparison", + ) + cublaslt_cupti_us = _profile_cupti_total_us( + _cublaslt_call, + "BF16 cuBLASLt comparison", + ) + + print(f"\nSM107 BF16 [{m},{k}] x [{k},{n}] -> [{m},{n}]") + print(f"full sweep: candidates={len(baseline_tactics)}, " + f"tactic={sweep_tactic}, autotuner_cupti_time={sweep_us:.3f} us, " + f"comparison_cupti_time={sweep_cupti_us:.3f} us") + print(f"scheduler-only heuristic: " + f"candidates={len(heuristic_candidates)}, " + f"native_rank1_split_k={native_rank1_split_k}, " + f"candidate_splits={sorted(candidate_splits)}, " + f"tactic={heuristic_tactic}, " + f"autotuner_cupti_time={heuristic_us:.3f} us, " + f"comparison_cupti_time={heuristic_cupti_us:.3f} us") + print(f"cuBLASLt: comparison_cupti_time={cublaslt_cupti_us:.3f} us") + + tolerance = 1.1 + cublas_tolerance = 1.1 + assert heuristic_us <= sweep_us * tolerance, ( + f"BF16 heuristic tactic {heuristic_tactic} ({heuristic_us:.3f} us) " + f"is >{tolerance:.2f}x slower than full-sweep tactic " + f"{sweep_tactic} ({sweep_us:.3f} us)") + assert heuristic_cupti_us <= sweep_cupti_us * tolerance, ( + f"BF16 heuristic tactic {heuristic_tactic} " + f"({heuristic_cupti_us:.3f} us CUPTI) is >{tolerance:.2f}x slower " + f"than full-sweep tactic {sweep_tactic} " + f"({sweep_cupti_us:.3f} us CUPTI)") + assert heuristic_cupti_us <= cublaslt_cupti_us * cublas_tolerance, ( + f"BF16 CuTe DSL heuristic tactic {heuristic_tactic} " + f"({heuristic_cupti_us:.3f} us CUPTI) is " + f">{cublas_tolerance:.2f}x slower " + f"than cuBLASLt ({cublaslt_cupti_us:.3f} us CUPTI)") + + +def _run_cute_dsl_mxfp8_heuristic_comparison(monkeypatch, tuner, nvmmh): + """Compare MXFP8 full sweep with scheduler-only NVMMH on SM107.""" + from _torch.helpers import calc_diff, per_block_cast_to_fp8_e8m0 + + import tensorrt_llm.quantization.utils.fp8_utils as fp8_utils + from tensorrt_llm._torch.custom_ops import cute_dsl_custom_ops + + m, n, k = _CUTE_DSL_NVMMH_TEST_MNK + torch.manual_seed(2029) + act = torch.randn(m, k, dtype=torch.bfloat16, device="cuda") / math.sqrt(k) + weight = (torch.randn(n, k, dtype=torch.bfloat16, device="cuda") / + math.sqrt(k)) + act_fp8, act_sf = \ + torch.ops.trtllm.fp8_quantize_1x128_packed_ue8m0(act) + weight_fp8, weight_sf_k128 = per_block_cast_to_fp8_e8m0(weight) + weight_sf = fp8_utils.transform_k128_scales_to_cutedsl_mxfp8_layout( + weight_sf_k128, mn=n, k=k) + alpha = torch.ones((), dtype=torch.float32, device="cuda") + inputs = [act_fp8, weight_fp8, act_sf, weight_sf, alpha] + + runner = cute_dsl_custom_ops.CuteDSLMXFP8RubinLinear( + output_dtype=torch.bfloat16, use_tvm_ffi=True) + runner.__class__.kernel_cache.clear() + tuning_config = runner.__class__.tuning_config + monkeypatch.setattr(tuning_config, "use_cuda_graph", False) + + scheduler_fields = ("swizzle", "cta_order", "split_k") + tuner.configure_nvmmh(enabled=False, fields=scheduler_fields, max_tactics=5) + baseline_tactics = runner.get_valid_tactics(inputs, None) + assert baseline_tactics + sweep_tactic, sweep_us = _choose_cupti_autotuner_tactic( + tuner, + "test::cute_dsl_mxfp8_exact_full_sweep", + runner, + tuning_config, + inputs, + baseline_tactics, + ) + + tuner.configure_nvmmh(enabled=True, fields=scheduler_fields, max_tactics=5) + heuristic_candidates = runner.get_valid_tactics(inputs, None) + assert 0 < len(heuristic_candidates) <= len(baseline_tactics), ( + "nvMMH did not prune SM107 MXFP8 tactics: " + f"{len(heuristic_candidates)} vs {len(baseline_tactics)}") + + assert {_nvmmh_tactic_family(t) + for t in heuristic_candidates + }.issubset({_nvmmh_tactic_family(t) + for t in baseline_tactics}) + + def _split_k(tactic): + if (isinstance(tactic, tuple) and tactic and tactic[0] == "base" + and len(tactic) >= 9): + return int(tactic[8]) + return 1 + + candidate_splits = {_split_k(tactic) for tactic in heuristic_candidates} + expected_local_splits = { + split_k + for split_k in runner.split_k_candidates + if nvmmh.is_sm107_nvmmh_split_k_eligible(k, runner.mma_tiler_k, split_k) + } + assert expected_local_splits.issubset(candidate_splits), ( + "SM107 MXFP8 scheduler-only filtering removed locally admitted " + f"split-K candidates: expected={expected_local_splits}, " + f"actual={candidate_splits}") + assert all( + nvmmh.is_sm107_nvmmh_split_k_eligible(k, runner.mma_tiler_k, split_k) + for split_k in candidate_splits), ( + f"SM107 MXFP8 retained an ineligible split: {candidate_splits}") + assert 8 in candidate_splits + baseline_base_families = { + tactic[:7] + for tactic in baseline_tactics + if isinstance(tactic, tuple) and tactic and tactic[0] == "base" + } + heuristic_base_families = { + tactic[:7] + for tactic in heuristic_candidates + if isinstance(tactic, tuple) and tactic and tactic[0] == "base" + } + assert heuristic_base_families == baseline_base_families, ( + "scheduler-only MXFP8 NVMMH must keep every base tile/cluster family") + + heuristic_tactic, heuristic_us = _choose_cupti_autotuner_tactic( + tuner, + "test::cute_dsl_mxfp8_scheduler_only_heuristic", + runner, + tuning_config, + inputs, + heuristic_candidates, + ) + + def _cublaslt_call(): + # PyTorch's Rubin MXFP8 reference dispatches this block-scaled call to + # cuBLASLt. The CuTe DSL inputs already use its flat R128c4 SF layout. + return torch._scaled_mm( + act_fp8, + weight_fp8.t(), + scale_a=act_sf.view(torch.float8_e8m0fnu), + scale_b=weight_sf.view(torch.float8_e8m0fnu), + out_dtype=torch.bfloat16, + ) + + sweep_output = runner(inputs, tactic=sweep_tactic) + heuristic_output = runner(inputs, tactic=heuristic_tactic) + cublaslt_output = _cublaslt_call() + torch.cuda.synchronize() + reference = act @ weight.t() + assert calc_diff(sweep_output, reference) < 1e-3 + assert calc_diff(heuristic_output, reference) < 1e-3 + assert calc_diff(cublaslt_output, reference) < 1e-3 + + sweep_cupti_us = _profile_cupti_total_us( + lambda: runner(inputs, tactic=sweep_tactic), + "MXFP8 full-sweep tactic comparison", + ) + heuristic_cupti_us = _profile_cupti_total_us( + lambda: runner(inputs, tactic=heuristic_tactic), + "MXFP8 scheduler-only nvMMH tactic comparison", + ) + cublaslt_cupti_us = _profile_cupti_total_us( + _cublaslt_call, + "MXFP8 cuBLASLt comparison", + ) + + print(f"\nSM107 MXFP8 [{m},{k}] x [{k},{n}] -> [{m},{n}]") + print(f"full sweep: candidates={len(baseline_tactics)}, " + f"tactic={sweep_tactic}, autotuner_cupti_time={sweep_us:.3f} us, " + f"comparison_cupti_time={sweep_cupti_us:.3f} us") + print(f"scheduler-only heuristic: " + f"candidates={len(heuristic_candidates)}, " + f"candidate_splits={sorted(candidate_splits)}, " + f"tactic={heuristic_tactic}, " + f"autotuner_cupti_time={heuristic_us:.3f} us, " + f"comparison_cupti_time={heuristic_cupti_us:.3f} us") + print(f"cuBLASLt: comparison_cupti_time={cublaslt_cupti_us:.3f} us") + + tolerance = 1.1 + cublas_tolerance = 1.1 + assert heuristic_us <= sweep_us * tolerance, ( + f"MXFP8 heuristic tactic {heuristic_tactic} ({heuristic_us:.3f} us) " + f"is >{tolerance:.2f}x slower than full-sweep tactic " + f"{sweep_tactic} ({sweep_us:.3f} us)") + assert heuristic_cupti_us <= sweep_cupti_us * tolerance, ( + f"MXFP8 heuristic tactic {heuristic_tactic} " + f"({heuristic_cupti_us:.3f} us CUPTI) is >{tolerance:.2f}x slower " + f"than full-sweep tactic {sweep_tactic} " + f"({sweep_cupti_us:.3f} us CUPTI)") + assert heuristic_cupti_us <= cublaslt_cupti_us * cublas_tolerance, ( + f"MXFP8 CuTe DSL heuristic tactic {heuristic_tactic} " + f"({heuristic_cupti_us:.3f} us CUPTI) is " + f">{cublas_tolerance:.2f}x slower " + f"than cuBLASLt ({cublaslt_cupti_us:.3f} us CUPTI)") + + +def _run_cute_dsl_nvfp4_heuristic_comparison(monkeypatch, tuner, nvmmh, + sm_version): + """Compare NVFP4 full sweep, NVMMH pruning, and cuBLASLt.""" + from _torch.helpers import calc_diff + + if sm_version == 107: + from tensorrt_llm._torch.custom_ops.cute_dsl_custom_ops import \ + CuteDSLNVFP4RubinLinear as NVFP4Runner + else: + from tensorrt_llm._torch.custom_ops.cute_dsl_custom_ops import \ + CuteDSLNVFP4BlackwellRunner as NVFP4Runner + + # Keep the NVFP4 operands aligned with the BF16 and MXFP8 subcases. FP4 + # packing / scale-factor layout follows + # shmoo_nvfp4_cutedsl_heuristics.py::_quantize_inputs. + m, n, k = _CUTE_DSL_NVMMH_TEST_MNK dtype = torch.bfloat16 sf_vec_size = 16 torch.manual_seed(0) @@ -1566,52 +2038,77 @@ def test_cutedsl_nvfp4_heuristic_matches_full_sweep(monkeypatch): False) w_fp4, w_sf = torch.ops.trtllm.fp4_quantize(w, w_sf_global, sf_vec_size, False) - alpha = torch.tensor([1.0], device="cuda") + alpha = (1.0 / (x_sf_global * w_sf_global)).reshape(1) inputs = [x_fp4, w_fp4, x_sf, w_sf, alpha] - runner = CuteDSLNVFP4BlackwellRunner(output_dtype=dtype) + runner = NVFP4Runner(output_dtype=dtype) tuning_config = runner.__class__.tuning_config + monkeypatch.setattr(tuning_config, "use_cuda_graph", False) + nvfp4_fields = (("tile", "cluster", "split_k") if sm_version == 107 else + ("tile", "cluster")) + tuner.configure_nvmmh( + enabled=False, + fields=nvfp4_fields, + max_tactics=5, + ) - def _best_tactic(): - tuner = AutoTuner.get() + def _best_tactic(name, expected_candidates): + assert len(expected_candidates) > 1, ( + "NVFP4 performance comparison requires a measured tactic sweep") tuner.clear_cache() - with autotune(): + with autotune(skip_dynamic_tuning_buckets=True): _, tactic = tuner.choose_one( - "test::cutedsl_nvfp4_heuristic_match", + name, [runner], tuning_config, inputs, ) + cached_tactic, _ = _require_cupti_autotuner_entry( + tuner, name, expected_candidates) + assert cached_tactic == tactic return tactic - def _dominant_kernel(fn, iters=20): - """(name, per-call device-us) of the longest-running CUDA kernel that - fn launches, isolating kernel time from host/op overhead.""" - from torch.profiler import ProfilerActivity, profile - fn() - torch.cuda.synchronize() - with profile(activities=[ProfilerActivity.CUDA]) as prof: - for _ in range(iters): - fn() - torch.cuda.synchronize() - best_name, best_total, best_count = "", -1.0, 1 - for e in prof.key_averages(): - t = (getattr(e, "self_device_time_total", 0) - or getattr(e, "self_cuda_time_total", 0)) - if t and t > best_total: - best_name, best_total, best_count = e.key, t, max(1, e.count) - return best_name, best_total / best_count + baseline_tactics = runner.get_valid_tactics(inputs, None) + assert baseline_tactics # Full sweep: heuristics disabled. - monkeypatch.delenv("TRTLLM_CUTEDSL_NVMMH_ENABLE", raising=False) - sweep_tactic = _best_tactic() + sweep_tactic = _best_tactic("test::cute_dsl_nvfp4_exact_full_sweep", + baseline_tactics) # Pruned: nvMatmulHeuristics drives the (coupled) tile+cluster candidates. - # Pin MAX_TACTICS so the tolerance below is not affected by an env override. - monkeypatch.setenv("TRTLLM_CUTEDSL_NVMMH_ENABLE", "1") - monkeypatch.setenv("TRTLLM_CUTEDSL_NVMMH_FIELDS", "tile,cluster") - monkeypatch.setenv("TRTLLM_CUTEDSL_NVMMH_MAX_TACTICS", "5") - heuristic_tactic = _best_tactic() + tuner.configure_nvmmh(enabled=True) + heuristic_candidates = runner.get_valid_tactics(inputs, None) + assert 0 < len(heuristic_candidates) <= len(baseline_tactics), ( + f"nvMMH did not prune SM{sm_version} NVFP4 tactics: " + f"{len(heuristic_candidates)} vs {len(baseline_tactics)}") + assert {_nvmmh_tactic_family(t) + for t in heuristic_candidates + }.issubset({_nvmmh_tactic_family(t) + for t in baseline_tactics}) + if sm_version == 107: + + def _nvfp4_split_k(tactic): + if (isinstance(tactic, tuple) and tactic and tactic[0] == "base" + and len(tactic) >= 9): + return int(tactic[8]) + return 1 + + candidate_splits = { + _nvfp4_split_k(tactic) + for tactic in heuristic_candidates + } + expected_local_splits = { + split_k + for split_k in runner.split_k_candidates + if nvmmh.is_sm107_nvmmh_split_k_eligible(k, runner.mma_tiler_k, + split_k) + } + assert expected_local_splits.issubset(candidate_splits) + assert all( + nvmmh.is_sm107_nvmmh_split_k_eligible( + k, runner.mma_tiler_k, split_k) for split_k in candidate_splits) + heuristic_tactic = _best_tactic("test::cute_dsl_nvfp4_exact_heuristic", + heuristic_candidates) # cuBLASLt runs its own heuristic auto-tuning; warm it under autotune(). def _cublas_call(): @@ -1622,18 +2119,39 @@ def _cublas_call(): _cublas_call() torch.cuda.synchronize() - # All comparisons use kernel-only device time (isolates the GEMM kernel from - # host/op dispatch overhead), measured via the CUDA profiler. - _, sweep_us = _dominant_kernel(lambda: runner(inputs, tactic=sweep_tactic)) - _, heuristic_us = _dominant_kernel( - lambda: runner(inputs, tactic=heuristic_tactic)) - _, cublas_us = _dominant_kernel(_cublas_call) + sweep_output = runner(inputs, tactic=sweep_tactic) + heuristic_output = runner(inputs, tactic=heuristic_tactic) + cublas_output = _cublas_call() + torch.cuda.synchronize() + assert calc_diff(sweep_output, cublas_output) < 1e-3 + assert calc_diff(heuristic_output, cublas_output) < 1e-3 + + # Use whole-call CUPTI time so in-place split-K's required output zeroing is + # included alongside the GEMM, matching the AutoTuner's timing semantics. + sweep_us = _profile_cupti_total_us( + lambda: runner(inputs, tactic=sweep_tactic), + "NVFP4 full-sweep tactic comparison", + ) + heuristic_us = _profile_cupti_total_us( + lambda: runner(inputs, tactic=heuristic_tactic), + "NVFP4 heuristic tactic comparison", + ) + cublas_us = _profile_cupti_total_us( + _cublas_call, + "NVFP4 cuBLASLt comparison", + ) + print(f"\nSM{sm_version} NVFP4 [{m},{k}] x [{k},{n}] -> [{m},{n}]") + print(f"full sweep: candidates={len(baseline_tactics)}, " + f"tactic={sweep_tactic}, time={sweep_us:.3f} us") + print(f"heuristic: candidates={len(heuristic_candidates)}, " + f"tactic={heuristic_tactic}, time={heuristic_us:.3f} us") + print(f"cuBLASLt: time={cublas_us:.3f} us") # Pruning must not degrade the achieved kernel runtime beyond this tolerance. # With the default MAX_TACTICS=5 the heuristic set includes the empirical - # best tile (its cluster ranking can still be slightly off, ~2-3% on square - # 4096), so a tight bound catches gross regressions while allowing that. - tolerance = 1.05 + # best tile, but its cluster ranking can be slightly off across Rubin nodes, + # so a 10% bound catches gross regressions while allowing that variance. + tolerance = 1.1 assert heuristic_us <= sweep_us * tolerance, ( f"heuristic-pruned tactic {heuristic_tactic} ({heuristic_us:.2f} us) is " f">{tolerance:.2f}x slower than full-sweep tactic {sweep_tactic} " @@ -1647,10 +2165,66 @@ def _cublas_call(): cublas_tolerance = 1.10 assert heuristic_us <= cublas_us * cublas_tolerance, ( f"CuteDSL heuristic kernel ({heuristic_us:.2f} us) is " - f">{cublas_tolerance:.2f}x slower than cuBLAS NVFP4 " + f">{cublas_tolerance:.2f}x slower than cuBLASLt NVFP4 " f"({cublas_us:.2f} us) for M={m}, N={n}, K={k}") +@pytest.mark.parametrize( + "precision,supported_sms", + [ + pytest.param("nvfp4", (100, 103, 107), id="nvfp4"), + pytest.param("bf16", (107, ), id="bf16"), + pytest.param("mxfp8", (107, ), id="mxfp8"), + ], +) +def test_cute_dsl_nvmmh_matches_full_sweep(precision, supported_sms, + monkeypatch, nvmmh_config_guard): + """Compare CuTe DSL full-sweep and NVMMH-selected performance. + + Each precision uses the same representative M=16, N=256, K=7168 problem + and compares the NVMMH-selected tactic with both the full CuTe DSL sweep + and cuBLASLt using torch.profiler/CUPTI. NVFP4 runs on SM100/SM103/SM107; + BF16 and MXFP8 run on SM107. + + The heuristic path is a strict validated subset of the sweep, so the exact + winning tactics may differ. The achieved runtime is the test invariant. + """ + if not torch.cuda.is_available(): + pytest.skip("requires a CUDA device") + + try: + from tensorrt_llm._utils import get_sm_version + sm_version = get_sm_version() + except Exception: + sm_version = None + if sm_version not in supported_sms: + supported_sm_names = ", ".join(f"SM{sm}" for sm in supported_sms) + pytest.skip(f"CuTe DSL {precision} comparison requires " + f"{supported_sm_names}") + + from tensorrt_llm._torch.custom_ops import \ + cutedsl_matmul_heuristics as nvmmh + if not nvmmh.IS_NVMMH_AVAILABLE: + pytest.skip("nvMatmulHeuristics library not installed") + if sm_version == 107: + from tensorrt_llm._torch.cute_dsl_utils import \ + IS_CUTLASS_DSL_RUBIN_AVAILABLE + if not IS_CUTLASS_DSL_RUBIN_AVAILABLE: + pytest.skip("CuTe DSL Rubin support is not available") + + _require_cupti_cuda_activity() + + comparison, args = { + "nvfp4": (_run_cute_dsl_nvfp4_heuristic_comparison, + (monkeypatch, nvmmh_config_guard, nvmmh, sm_version)), + "bf16": (_run_cute_dsl_bf16_heuristic_comparison, + (monkeypatch, nvmmh_config_guard, nvmmh)), + "mxfp8": (_run_cute_dsl_mxfp8_heuristic_comparison, + (monkeypatch, nvmmh_config_guard, nvmmh)), + }[precision] + comparison(*args) + + @pytest.mark.parametrize("distribution", ["random", "balanced"]) def test_trtllm_gen_moe_dummy_topk_local_experts_less_than_topk( distribution, monkeypatch): diff --git a/tests/unittest/api_stability/references/llm.yaml b/tests/unittest/api_stability/references/llm.yaml index d2d1739473ff..b8e00f41ff4b 100644 --- a/tests/unittest/api_stability/references/llm.yaml +++ b/tests/unittest/api_stability/references/llm.yaml @@ -286,6 +286,10 @@ methods: annotation: bool default: True status: prototype + autotuner_nvmmh_config: + annotation: Optional[tensorrt_llm.llmapi.llm_args.AutoTunerNvMMHConfig] + default: null + status: prototype use_fine_grained_sync: annotation: bool default: False From b706c24c610e118c88aa5b16cf118fddf0bfc171 Mon Sep 17 00:00:00 2001 From: peaceh <103117813+peaceh-nv@users.noreply.github.com> Date: Thu, 17 Sep 2026 00:31:01 -0700 Subject: [PATCH 2/8] [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> --- docs/source/torch/adding_custom_kernels.md | 10 +- tensorrt_llm/_torch/autotuner.py | 88 +++++++++---- .../_torch/custom_ops/cute_dsl_custom_ops.py | 22 +++- .../custom_ops/cutedsl_matmul_heuristics.py | 2 + .../_torch/pyexecutor/model_engine.py | 14 +- .../parallel/test_fp8_block_scale_gemm.py | 50 ++++++-- tests/unittest/_torch/misc/test_autotuner.py | 120 ++++++++++++++++++ 7 files changed, 261 insertions(+), 45 deletions(-) diff --git a/docs/source/torch/adding_custom_kernels.md b/docs/source/torch/adding_custom_kernels.md index a64dde92bcfc..edfa32813dce 100644 --- a/docs/source/torch/adding_custom_kernels.md +++ b/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 diff --git a/tensorrt_llm/_torch/autotuner.py b/tensorrt_llm/_torch/autotuner.py index 87da00ea2c8f..dfcf035593c3 100644 --- a/tensorrt_llm/_torch/autotuner.py +++ b/tensorrt_llm/_torch/autotuner.py @@ -22,7 +22,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 @@ -183,6 +185,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, @@ -371,9 +374,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 @@ -670,6 +675,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__, @@ -1072,12 +1078,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' @@ -1148,31 +1157,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. @@ -1782,6 +1817,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 " @@ -1826,6 +1862,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): @@ -1845,6 +1882,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: @@ -1899,12 +1937,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 @@ -1913,12 +1954,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. diff --git a/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py b/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py index 60c1a09e12d8..ff86bb085743 100644 --- a/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py +++ b/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py @@ -769,6 +769,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): @@ -786,6 +787,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. " @@ -4979,6 +4981,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( @@ -4987,6 +4990,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. " @@ -10065,6 +10069,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( @@ -10081,7 +10086,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. " @@ -10358,6 +10363,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])) @@ -10666,6 +10672,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( @@ -10736,6 +10743,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) @@ -11202,6 +11210,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. @@ -11230,6 +11239,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. " @@ -11265,6 +11275,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) @@ -13992,6 +14003,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() @@ -14012,6 +14024,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. " @@ -14429,6 +14442,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): @@ -14482,6 +14496,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, @@ -14534,7 +14549,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)) @@ -15663,6 +15678,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( @@ -15671,7 +15687,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. " diff --git a/tensorrt_llm/_torch/custom_ops/cutedsl_matmul_heuristics.py b/tensorrt_llm/_torch/custom_ops/cutedsl_matmul_heuristics.py index 50841a84bf98..48df5aece788 100644 --- a/tensorrt_llm/_torch/custom_ops/cutedsl_matmul_heuristics.py +++ b/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) diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py index 43ee5dd59a01..14182772793a 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_engine.py +++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py @@ -162,8 +162,12 @@ def _configure_deep_gemm_pdl() -> None: _DEEP_GEMM_PDL_CONFIGURED = True -def _configure_autotuner_nvmmh(llm_args: TorchLlmArgs) -> None: - """Install the per-worker NVMMH policy before tactic enumeration.""" +def _configure_autotuner_nvmmh(llm_args: TorchLlmArgs, owner: object) -> None: + """Install and pin an engine's NVMMH policy before tactic enumeration. + + Engines in one worker share the autotuner, so simultaneous owners must + agree on the effective policy, including the disabled default. + """ config = llm_args.autotuner_nvmmh_config if config is None: policy = NvMMHConfig() @@ -173,7 +177,7 @@ def _configure_autotuner_nvmmh(llm_args: TorchLlmArgs) -> None: fields=config.fields, max_tactics=config.max_tactics, ) - AutoTuner.get().configure_nvmmh(policy) + AutoTuner.get()._acquire_nvmmh_policy(owner, policy) class PyTorchModelEngine(ModelEngine): @@ -192,10 +196,11 @@ def __init__( model_weights_memory_tag: Optional[str] = None, model_weights_restore_mode=None, ): + """Pin the worker policy and initialize model loading and execution resources.""" # Every executor worker receives its own TorchLlmArgs copy. Install the # policy before model loading because CuTe DSL runners may enumerate # tactics while modules are being constructed. - _configure_autotuner_nvmmh(llm_args) + _configure_autotuner_nvmmh(llm_args, self) _configure_deep_gemm_pdl() self._metrics: dict[str, float] = defaultdict(float) @@ -1161,6 +1166,7 @@ def cleanup(self) -> None: # Release model weights. release_gc() + AutoTuner.get()._release_nvmmh_policy(self) self._cleanup_done = True def __del__(self) -> None: diff --git a/tests/unittest/_torch/kernels/parallel/test_fp8_block_scale_gemm.py b/tests/unittest/_torch/kernels/parallel/test_fp8_block_scale_gemm.py index af11313f1751..35cf1348bda1 100644 --- a/tests/unittest/_torch/kernels/parallel/test_fp8_block_scale_gemm.py +++ b/tests/unittest/_torch/kernels/parallel/test_fp8_block_scale_gemm.py @@ -871,7 +871,7 @@ def test_cute_dsl_mxfp8_gemm_rubin_rejects_clc_grid_z_overflow(): getSMVersion() != 107 or not IS_CUTLASS_DSL_RUBIN_AVAILABLE, reason="The test requires SM107 and SM107 CuTe DSL support.", ) -def test_cute_dsl_mxfp8_gemm_rubin_clc_dynamic_prefetch_multi_wave(): +def test_cute_dsl_mxfp8_gemm_rubin_clc_dynamic_prefetch_multi_wave(monkeypatch): """CLC dynamic scheduling supports Boolean-enabled prefetch.""" from tensorrt_llm._torch.custom_ops.cute_dsl_custom_ops import \ CuteDSLMXFP8RubinLinear @@ -917,24 +917,46 @@ def test_cute_dsl_mxfp8_gemm_rubin_clc_dynamic_prefetch_multi_wave(): non_power_split_k_tactic = (*dynamic_tactic[:-2], 3, 1) runner = CuteDSLMXFP8RubinLinear(output_dtype=torch.bfloat16, use_tvm_ffi=True) + from tensorrt_llm._torch import autotuner + from tensorrt_llm._torch.custom_ops import cute_dsl_custom_ops as ops + from tensorrt_llm._torch.custom_ops import \ + cutedsl_matmul_heuristics as heuristics + + # Isolate this test from the process-wide policy of preceding tests. + monkeypatch.setattr(autotuner.AutoTuner, "_nvmmh_config", + autotuner.NvMMHConfig()) valid_tactics = runner.get_valid_tactics(inputs, None) # Model-selected swizzles extend the validated baseline tactic tuple. assert static_tactic[:-1] in valid_tactics assert dynamic_tactic[:-1] in valid_tactics assert split_k_tactic[:-1] in valid_tactics - for scheduler_mode in ("static", "clc_dynamic"): - for raster_order in ("m", "n"): - swizzle_sizes = ((1, 2, 4, 8) if scheduler_mode == "static" else - (1, )) - for swizzle_size in swizzle_sizes: - tactic = ( - *static_tactic[:-4], - scheduler_mode, - raster_order, - 1, - swizzle_size, - ) - assert tactic[:-1] in valid_tactics + monkeypatch.setattr(ops, "IS_NVMMH_AVAILABLE", True) + monkeypatch.setattr( + autotuner.AutoTuner, "_nvmmh_config", + autotuner.NvMMHConfig(enabled=True, fields=("swizzle", ))) + for swizzle_size in (1, 2, 4, 8): + config = heuristics.HeuristicConfig((128, 256), (4, 2), swizzle_size, 0) + monkeypatch.setattr(ops, + "rank_configs", + lambda *args, _config=config, **kwargs: [_config]) + annotated_tactics = runner.get_valid_tactics(inputs, None) + base_tactics = [t for t in annotated_tactics if t[0] == "base"] + assert base_tactics and all(len(t) == 10 for t in base_tactics) + for scheduler_mode in ("static", "clc_dynamic"): + for raster_order in ("m", "n"): + baseline = (*static_tactic[:-4], scheduler_mode, raster_order, + 1) + assert baseline in valid_tactics + expected = (*baseline, swizzle_size) + annotated = runner._apply_nvmmh_scheduler( + baseline, config, {"swizzle"}) + if scheduler_mode == "clc_dynamic" and swizzle_size != 1: + assert annotated is None + assert expected not in annotated_tactics + else: + assert annotated == expected + assert expected in annotated_tactics + assert all(t[-1] == 1 for t in base_tactics if t[6] == "clc_dynamic") static_output = runner(inputs, tactic=static_tactic) dynamic_output = runner(inputs, tactic=dynamic_tactic) diff --git a/tests/unittest/_torch/misc/test_autotuner.py b/tests/unittest/_torch/misc/test_autotuner.py index 05910e78f479..60de2455f9fc 100644 --- a/tests/unittest/_torch/misc/test_autotuner.py +++ b/tests/unittest/_torch/misc/test_autotuner.py @@ -57,6 +57,7 @@ @pytest.fixture def nvmmh_config_guard(): + """Restore the process-wide NVMMH policy after each test that changes it.""" tuner = AutoTuner.get() previous = tuner.nvmmh_config try: @@ -901,12 +902,14 @@ def test_kernel_testing_single_context(): def test_kernel_testing_uses_runner_valid_tactics(): + """Verify captured tactic combinations use the runner-provided candidate list.""" class CaptureRunner(TunableRunner): def get_valid_tactics(self, inputs: List[FakeTensor], profile: OptimizationProfile, **kwargs) -> List[int]: + """Provide a fixed candidate list for capture and replay assertions.""" return [1, 2] def forward(self, @@ -915,6 +918,7 @@ def forward(self, *, tactic: int = -1, **kwargs) -> torch.Tensor: + """Return the input unchanged while the test observes tactic selection.""" assert tactic in [-1, 0, 1, 2] return inputs[0] @@ -1676,6 +1680,7 @@ def _profile_cupti_total_us(fn, context, warmup=20, iterations=200, trials=5): def _nvmmh_tactic_family(tactic): + """Strip model-selected scheduler fields before comparing structural families.""" # Scheduler-only guidance annotates validated families with model-selected # raster/swizzle values; the resulting tuples need not occur in the sweep. if tactic[0] == "mixed_clusters": @@ -1738,6 +1743,7 @@ def _run_cute_dsl_bf16_heuristic_comparison(monkeypatch, tuner, nvmmh): native_rank1_split_k = int(model_configs[0].split_k) 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) >= 6): return int(tactic[5]) @@ -1780,6 +1786,7 @@ def _split_k(tactic): ) def _cublaslt_call(): + """Execute the cuBLASLt reference used for BF16 timing comparisons.""" return torch.ops.trtllm.cublas_mm( act, weight.t(), @@ -1906,6 +1913,7 @@ def _run_cute_dsl_mxfp8_heuristic_comparison(monkeypatch, tuner, nvmmh): for t in baseline_tactics}) 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): return int(tactic[8]) @@ -1949,6 +1957,7 @@ def _split_k(tactic): ) def _cublaslt_call(): + """Execute the cuBLASLt reference for the MXFP8 comparison inputs.""" # PyTorch's Rubin MXFP8 reference dispatches this block-scaled call to # cuBLASLt. The CuTe DSL inputs already use its flat R128c4 SF layout. return torch._scaled_mm( @@ -2053,6 +2062,7 @@ def _run_cute_dsl_nvfp4_heuristic_comparison(monkeypatch, tuner, nvmmh, ) def _best_tactic(name, expected_candidates): + """Tune the supplied candidates and require a CUPTI-ranked cache entry.""" assert len(expected_candidates) > 1, ( "NVFP4 performance comparison requires a measured tactic sweep") tuner.clear_cache() @@ -2088,6 +2098,7 @@ def _best_tactic(name, expected_candidates): if sm_version == 107: def _nvfp4_split_k(tactic): + """Read a base NVFP4 split factor, treating other families as unsplit.""" if (isinstance(tactic, tuple) and tactic and tactic[0] == "base" and len(tactic) >= 9): return int(tactic[8]) @@ -2514,3 +2525,112 @@ def test_deep_gemm_jit_warmup_buckets(): assert max(deep_gemm_jit_warmup_buckets(32)) < 128 assert deep_gemm_jit_warmup_buckets(16384) == deep_gemm_jit_warmup_buckets( 8192) + +@pytest.mark.parametrize("conflicting_config", + [None, { + "fields": ("tile", ), + "max_tactics": 2 + }]) +def test_nvmmh_engine_policy_lifetime(nvmmh_config_guard, monkeypatch, + conflicting_config): + """Engine wiring pins non-default policy until the last owner is cleaned up.""" + from tensorrt_llm._torch.pyexecutor import model_engine + from tensorrt_llm.llmapi.llm_args import AutoTunerNvMMHConfig, TorchLlmArgs + + class EngineOwner: + """Minimal engine resources for exercising the real cleanup path.""" + _cleanup_done = False + model_loader = None + + def _release_cuda_graphs(self): + """No graph resources are allocated by this wiring test.""" + + monkeypatch.setattr(model_engine, "release_gc", lambda: None) + args = TorchLlmArgs.model_construct( + autotuner_nvmmh_config=AutoTunerNvMMHConfig( + fields=("swizzle", "cta_order"), max_tactics=3)) + first, second, rejected = EngineOwner(), EngineOwner(), EngineOwner() + tuner = nvmmh_config_guard + try: + model_engine._configure_autotuner_nvmmh(args, first) + expected = autotuner.NvMMHConfig(enabled=True, + fields=("swizzle", "cta_order"), + max_tactics=3) + assert tuner.nvmmh_config == expected + installed = tuner.nvmmh_config + model_engine._configure_autotuner_nvmmh(args, second) + assert tuner.nvmmh_config is installed + other_args = TorchLlmArgs.model_construct(autotuner_nvmmh_config=( + None if conflicting_config is None else AutoTunerNvMMHConfig( + **conflicting_config))) + with pytest.raises(ValueError, match="active model engines"): + model_engine._configure_autotuner_nvmmh(other_args, rejected) + assert tuner.nvmmh_config == expected + with pytest.raises(ValueError, match="active model engines"): + tuner.configure_nvmmh(enabled=False) + assert tuner.nvmmh_config == expected + + model_engine.PyTorchModelEngine.cleanup(first) + model_engine.PyTorchModelEngine.cleanup(first) # Idempotent release. + with pytest.raises(ValueError, match="active model engines"): + model_engine._configure_autotuner_nvmmh(other_args, rejected) + model_engine.PyTorchModelEngine.cleanup(second) + model_engine._configure_autotuner_nvmmh(other_args, rejected) + assert tuner.nvmmh_config != expected + model_engine.PyTorchModelEngine.cleanup(rejected) + finally: + tuner._release_nvmmh_policy(first) + tuner._release_nvmmh_policy(second) + tuner._release_nvmmh_policy(rejected) + + +def test_nvmmh_policy_owner_collection(nvmmh_config_guard): + """An abandoned or partially initialized engine cannot pin policy forever.""" + import gc + import weakref + + class Owner: + """Weak-referenceable stand-in for an engine that failed to initialize.""" + + owner = Owner() + tuner = nvmmh_config_guard + tuner._acquire_nvmmh_policy(owner, autotuner.NvMMHConfig(enabled=True)) + ref = weakref.ref(owner) + del owner + gc.collect() + assert ref() is None + tuner.configure_nvmmh(enabled=False) + assert not tuner.nvmmh_config.enabled + + +def test_nvmmh_unmatched_swap_preserves_split_k_admission( + nvmmh_config_guard, monkeypatch): + """A model miss for one swap must not resurrect rejected split-K factors.""" + from tensorrt_llm._torch.custom_ops import cute_dsl_custom_ops as ops + from tensorrt_llm._torch.custom_ops import \ + cutedsl_matmul_heuristics as heuristics + + if not hasattr(ops, "CuteDSLMXFP8RubinLinear"): + pytest.skip("Rubin CuTe DSL is unavailable") + monkeypatch.setattr(ops, "IS_NVMMH_AVAILABLE", True) + runner = ops.CuteDSLMXFP8RubinLinear(output_dtype=torch.bfloat16) + nvmmh_config_guard.configure_nvmmh(enabled=True, + fields=("tile", "cluster", "split_k")) + tactics = [("base", (128, 128, 128), (128, 128, 64), (1, 1), swap, False, + "static", "m", split) for swap in (False, True) + for split in (1, 2, 4, 8)] + + def rank_one_orientation(m, n, *args, **kwargs): + """Only the unswapped model problem has a representable result.""" + return ([heuristics.HeuristicConfig((128, 128), + (1, 1), 1, 0)] if m == 128 else []) + + monkeypatch.setattr(ops, "rank_configs", rank_one_orientation) + selected = runner._rank_prune_tactics(tactics, 128, 256, 1024) + admitted = [ + t for t in tactics if heuristics.is_sm107_nvmmh_split_k_eligible( + 1024, runner.mma_tiler_k, t[8]) + ] + assert len(admitted) < len(tactics) + assert {t[4] for t in selected} == {False, True} + assert set(selected) == set(admitted) From c93799f55e3ac176afe64dd5bdf5ee2b2643ad5b Mon Sep 17 00:00:00 2001 From: peaceh <103117813+peaceh-nv@users.noreply.github.com> Date: Mon, 21 Sep 2026 01:17:03 -0700 Subject: [PATCH 3/8] [TRTLLM-14024][chore] Refresh generated metadata after main rebase Regenerate the LLM args manifest for main's existing uses_spcompress field and apply the formatter's required spacing at the nvMMH integration boundary. Validation on Hecate SM107: 47 autotuner regressions, 12 focused nvMMH checks, and 10 Rubin kernel tests passed. Golden manifest and all applicable pre-commit checks passed. Signed-off-by: peaceh <103117813+peaceh-nv@users.noreply.github.com> --- tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py | 1 + 1 file changed, 1 insertion(+) diff --git a/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py b/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py index ff86bb085743..25306c70c591 100644 --- a/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py +++ b/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py @@ -111,6 +111,7 @@ def _kimi_k3_mxfp8_tuning_bucket(num_tokens: int) -> int: return next(m for m in _KIMI_K3_MXFP8_TUNING_BUCKETS if num_tokens <= m) return next_positive_power_of_2(num_tokens) + _CUTEDSL_NVMMH_TACTIC_SEARCH_CACHE_VERSION = "cutedsl-nvmmh-v1" From 88999c753bdf07b7d229d42b2ea17194d219207d Mon Sep 17 00:00:00 2001 From: peaceh <103117813+peaceh-nv@users.noreply.github.com> Date: Tue, 29 Sep 2026 06:24:05 -0700 Subject: [PATCH 4/8] [TRTLLM-14024][fix] Address remaining nvMMH review feedback Signed-off-by: peaceh <103117813+peaceh-nv@users.noreply.github.com> --- docs/source/torch/adding_custom_kernels.md | 11 +- .../layer_wise_benchmarks/config_ctx.yaml | 2 +- .../layer_wise_benchmarks/config_gen.yaml | 2 +- tensorrt_llm/_torch/autotuner.py | 139 +++++++++----- .../custom_ops/cutedsl_matmul_heuristics.py | 24 +-- .../_torch/custom_ops/torch_custom_ops.py | 17 +- .../_torch/pyexecutor/model_engine.py | 6 +- tensorrt_llm/_torch/pyexecutor/py_executor.py | 9 +- tensorrt_llm/llmapi/llm_args.py | 2 +- tests/unittest/_torch/misc/test_autotuner.py | 169 +----------------- 10 files changed, 148 insertions(+), 233 deletions(-) diff --git a/docs/source/torch/adding_custom_kernels.md b/docs/source/torch/adding_custom_kernels.md index edfa32813dce..1aaef8a06cc9 100644 --- a/docs/source/torch/adding_custom_kernels.md +++ b/docs/source/torch/adding_custom_kernels.md @@ -323,7 +323,16 @@ 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. +the last executor shuts down (or its engine is cleaned up or collected), a new +policy can be installed. + +CuTe DSL runners use the existing globaltimer or CUDA-event timing by default. +Set `TLLM_PROFILING_TIMER=torch_profiler` to opt in to Kineto/CUPTI timing for +these runners. Explicit `globaltimer` and `cuda_event` values keep their +existing meanings. CUPTI results and CUDA-event fallbacks have separate cache +keys; enabling CUPTI or NVMMH can require re-tuning entries in an existing +`TLLM_AUTOTUNER_CACHE_PATH` file. The layer-wise benchmark YAML examples leave +NVMMH disabled; set `autotuner_nvmmh: true` to benchmark pruning. Compatible validated runner-specific tactics that NVMMH cannot represent are retained automatically and empirically profiled. The BF16 and SM107 diff --git a/examples/layer_wise_benchmarks/config_ctx.yaml b/examples/layer_wise_benchmarks/config_ctx.yaml index f2ce0274b97f..f6539c6b239e 100644 --- a/examples/layer_wise_benchmarks/config_ctx.yaml +++ b/examples/layer_wise_benchmarks/config_ctx.yaml @@ -10,7 +10,7 @@ enable_attention_dp: true # Model init args moe_backend: CUTLASS enable_autotuner: true -autotuner_nvmmh: true +autotuner_nvmmh: false autotuner_nvmmh_fields: [swizzle, cta_order, split_k] autotuner_nvmmh_max_tactics: 5 use_cuda_graph: false diff --git a/examples/layer_wise_benchmarks/config_gen.yaml b/examples/layer_wise_benchmarks/config_gen.yaml index 2b52c8e5887e..29b3eab0b4e4 100644 --- a/examples/layer_wise_benchmarks/config_gen.yaml +++ b/examples/layer_wise_benchmarks/config_gen.yaml @@ -10,7 +10,7 @@ enable_attention_dp: true # Model init args moe_backend: CUTLASS enable_autotuner: true -autotuner_nvmmh: true +autotuner_nvmmh: false autotuner_nvmmh_fields: [swizzle, cta_order, split_k] autotuner_nvmmh_max_tactics: 5 use_cuda_graph: true diff --git a/tensorrt_llm/_torch/autotuner.py b/tensorrt_llm/_torch/autotuner.py index dfcf035593c3..c54899ad99c3 100644 --- a/tensorrt_llm/_torch/autotuner.py +++ b/tensorrt_llm/_torch/autotuner.py @@ -53,6 +53,10 @@ class _TorchProfilerUnavailableError(RuntimeError): """Raised when Kineto/CUPTI cannot provide CUDA timing activities.""" +class _TorchProfilerEmptyTraceError(RuntimeError): + """Raised when one completed profile contains no CUDA activity.""" + + class DistributedTuningStrategy(enum.Enum): """ Strategy for distributed tuning. @@ -371,15 +375,17 @@ def tactic_search_cache_key(self) -> Optional[Tuple[Any, ...]]: class CuteDSLTunableRunner(TunableRunner): - """Marker base for CuTe DSL runners that prefer CUPTI timing.""" + """Marker base for CuTe DSL runners with optional CUPTI timing.""" def use_torch_profiler(self, tactic: Any) -> bool: - """Prefer CUDA activity timing for every CuTe DSL tactic.""" - return True + """Use CUDA activity timing only when explicitly requested.""" + return os.getenv("TLLM_PROFILING_TIMER", + "").lower() == _TORCH_PROFILER_TIMER_KEY def profiling_timer_cache_key(self) -> Optional[str]: - """Separate CUPTI-ranked tactics from legacy timer cache entries.""" - return _TORCH_PROFILER_TIMER_KEY + """Separate opt-in CUPTI-ranked tactics from legacy cache entries.""" + return _TORCH_PROFILER_TIMER_KEY if self.use_torch_profiler( + None) else None @contextlib.contextmanager @@ -630,15 +636,16 @@ def search_cache( runner_id is the index in the current runners list """ for idx, r in enumerate(runners): - cache_keys = [ - self.get_cache_key( - custom_op, - r, - input_shapes, - tuning_config, - apply_map_to_tuning_buckets, - ) - ] + cache_key = self.get_cache_key( + custom_op, + r, + input_shapes, + tuning_config, + apply_map_to_tuning_buckets, + ) + if cache_key in self.cache: + cached_runner_id, tactic, min_time = self.cache[cache_key] + return True, idx, tactic, min_time # A fresh inference process has not probed CUPTI yet. It may reuse # an explicitly tagged CUDA-event fallback entry when no preferred # CUPTI entry exists. Tuning mode leaves this disabled so a @@ -647,21 +654,17 @@ def search_cache( if (allow_timer_fallback and not AutoTuner._torch_profiler_unavailable and r.profiling_timer_cache_key() == _TORCH_PROFILER_TIMER_KEY): - cache_keys.append( - self.get_cache_key( - custom_op, - r, - input_shapes, - tuning_config, - apply_map_to_tuning_buckets, - timer_key_override=_CUDA_EVENT_FALLBACK_TIMER_KEY, - )) - - for cache_key in cache_keys: - if cache_key in self.cache: - # Return the current index in runners list, not the cached - # runner_id. - cached_runner_id, tactic, min_time = self.cache[cache_key] + fallback_key = self.get_cache_key( + custom_op, + r, + input_shapes, + tuning_config, + apply_map_to_tuning_buckets, + timer_key_override=_CUDA_EVENT_FALLBACK_TIMER_KEY, + ) + if fallback_key in self.cache: + cached_runner_id, tactic, min_time = self.cache[ + fallback_key] return True, idx, tactic, min_time return False, *self.fallback_entry() @@ -1098,15 +1101,16 @@ def __init__(self, warmup=2, repeat=10, stream_delay_micro_secs=1000): self.profiling_cache = AutoTunerProfilingCache() self._primed_cached_tactics: set[tuple[str, str, str, str]] = set() self._last_profile_timer: Optional[str] = None + self._force_cuda_event_for_sweep = False self.is_tuning_mode = False self.skip_dynamic_tuning_buckets = False # Default timing backend: globaltimer kernel vs CUDA events. CuTe DSL - # runners prefer torch.profiler/CUPTI and fall back to CUDA events when - # CUPTI is unavailable in the current process. + # runners use torch.profiler/CUPTI only when explicitly requested. # TLLM_PROFILING_TIMER env var overrides auto-detection: # "globaltimer" -> force globaltimer # "cuda_event" -> force cuda events + # "torch_profiler" -> use CUPTI for CuTe DSL runners # unset/default -> auto-detect via confidential_compute_enabled() timer_env = os.getenv("TLLM_PROFILING_TIMER", "").lower() if timer_env == "globaltimer": @@ -1584,6 +1588,7 @@ def _profile_runners( has_tuning_failure_occurred = False best_runner_id, best_tactic = None, None torch_profiler_was_available = not AutoTuner._torch_profiler_unavailable + local_timer_fallback_started = self._force_cuda_event_for_sweep # If the inputs_pre_hook is provided, it will be called before profiling. if tuning_config.inputs_pre_hook is not None and not _inputs_prepared: input_tensors = tuning_config.inputs_pre_hook(input_tensors) @@ -1751,14 +1756,35 @@ def _profile_runners( retry_result[3] or has_tuning_failure_occurred, ) + if not local_timer_fallback_started and self._force_cuda_event_for_sweep: + # An outer profiler (or another profiler entry failure) forced one + # measurement onto CUDA events. Re-rank the complete sweep with + # events so no candidate is compared across different timers. + try: + return self._profile_runners(custom_op, + runners, + input_tensors, + profile, + tuning_config, + _inputs_prepared=True, + **kwargs) + finally: + self._force_cuda_event_for_sweep = False + if best_runner_id is not None: # At least one valid (runner, tactic) pair is found + timer_key_override = ( + _CUDA_EVENT_FALLBACK_TIMER_KEY + if self._force_cuda_event_for_sweep + and runners[best_runner_id].profiling_timer_cache_key() + == _TORCH_PROFILER_TIMER_KEY else None) cache_key = self.profiling_cache.get_cache_key( custom_op, runners[best_runner_id], profile.get_opt_shapes(), tuning_config, - apply_map_to_tuning_buckets=False) + apply_map_to_tuning_buckets=False, + timer_key_override=timer_key_override) self._debug_logger( f"[Autotuner] Profiling runner={runners[best_runner_id]}, tactic={best_tactic} for cache_key={cache_key}." @@ -1802,9 +1828,8 @@ def _torch_profiler_elapsed_time_ms(prof, repeat: int) -> float: if event.device_type == torch.autograd.DeviceType.CUDA and ( event.device_time_total or 0.0) > 0) if not math.isfinite(total_cuda_us) or total_cuda_us <= 0: - raise _TorchProfilerUnavailableError( - "torch.profiler/CUPTI reported no CUDA activity while " - "profiling a CuTe DSL tactic") + raise _TorchProfilerEmptyTraceError( + "torch.profiler reported no CUDA activity for this tactic") return total_cuda_us / max(1, int(repeat)) / 1000.0 @staticmethod @@ -1894,7 +1919,15 @@ def launch_profiled_work(): **kwargs, ) - if use_torch_profiler and not AutoTuner._torch_profiler_unavailable: + profiler_enabled = getattr(torch.autograd, "_profiler_enabled", + None) + if (use_torch_profiler and profiler_enabled is not None + and profiler_enabled()): + self._force_cuda_event_for_sweep = True + + if (use_torch_profiler + and not AutoTuner._torch_profiler_unavailable + and not self._force_cuda_event_for_sweep): profiled_work_started = False profiled_work_completed = False try: @@ -1923,13 +1956,19 @@ def launch_profiled_work(): # context. if profiled_work_started and not profiled_work_completed: raise - if not self._is_torch_profiler_unavailable_error(error): + if isinstance(error, _TorchProfilerEmptyTraceError): raise - self._disable_torch_profiler(error) - - # CUPTI was unavailable either earlier in this process or in - # the attempt above. Re-run this measurement with the legacy - # CUDA-event timer, irrespective of _use_global_timer. + if self._is_torch_profiler_unavailable_error(error): + self._disable_torch_profiler(error) + else: + # Entering the profiler can fail while an outer + # profiler is active. An exit failure after work + # completes is also a timer failure. Use events + # for this sweep without disabling CUPTI globally. + self._force_cuda_event_for_sweep = True + + # CUPTI was unavailable or profiler entry/exit failed. Re-run + # this measurement with CUDA events regardless of globaltimer. force_cuda_event = use_torch_profiler if not force_cuda_event and self._use_global_timer: @@ -1982,7 +2021,14 @@ def elapsed_time(): for _ in range(self.warmup): runner(input_tensor_batches[-1], tactic=tactic, **kwargs) - fewer_repeat_avg_time = pure_profile(stream, profile_fewer_repeat) + def profile_with_empty_retry(repeat: int) -> float: + """Retry one empty CUPTI trace without changing the process timer.""" + try: + return pure_profile(stream, repeat) + except _TorchProfilerEmptyTraceError: + return pure_profile(stream, repeat) + + fewer_repeat_avg_time = profile_with_empty_retry(profile_fewer_repeat) disable_short_profile = os.environ.get( "TLLM_AUTOTUNER_DISABLE_SHORT_PROFILE", "0") == "1" @@ -1994,11 +2040,12 @@ def elapsed_time(): avg_time = fewer_repeat_avg_time else: # profile the kernel with the full repeat to get precise time - avg_time = pure_profile(stream, self.repeat) + avg_time = profile_with_empty_retry(self.repeat) shapes = self._get_input_sizes(inputs) - timer_name = (_CUDA_EVENT_FALLBACK_TIMER_KEY if use_torch_profiler - and AutoTuner._torch_profiler_unavailable else + timer_name = (_CUDA_EVENT_FALLBACK_TIMER_KEY if use_torch_profiler and + (AutoTuner._torch_profiler_unavailable + or self._force_cuda_event_for_sweep) else _TORCH_PROFILER_TIMER_KEY if use_torch_profiler else "globaltimer" if self._use_global_timer else "cuda_event") self._last_profile_timer = timer_name diff --git a/tensorrt_llm/_torch/custom_ops/cutedsl_matmul_heuristics.py b/tensorrt_llm/_torch/custom_ops/cutedsl_matmul_heuristics.py index 48df5aece788..b4c1b2f87e83 100644 --- a/tensorrt_llm/_torch/custom_ops/cutedsl_matmul_heuristics.py +++ b/tensorrt_llm/_torch/custom_ops/cutedsl_matmul_heuristics.py @@ -119,10 +119,11 @@ class HeuristicConfig(NamedTuple): def nvmmh_cta_order_to_cute_raster(cta_order: int) -> str: """Translate CUTLASS3 nvMMH CTA order to CuTe's raster name. - nvMMH's CUTLASS swizzler and DKG integration use order 0 for an - M-fast raster and order 1 for an N-fast raster. Keep this conversion at - the adapter boundary because the public nvMMH header's row/column wording - does not match the CUTLASS3 implementation used to score the configs. + CUTLASS's nvMMH adapter maps order 0 to ``along_m`` and order 1 to + ``along_n`` (python/cutlass_library/heuristics_provider.py in CUTLASS, + see https://github.com/NVIDIA/cutlass/blob/main/python/cutlass_library/heuristics_provider.py#L222). + Keep this conversion at the adapter boundary because the public nvMMH + header's row/column wording differs from this CUTLASS3 mapping. """ if cta_order == 0: return "m" @@ -558,12 +559,9 @@ def _bf16_config_signature(config: HeuristicConfig, fields: set) -> tuple: if fields & {"tile", "cluster"}: cta_m, cta_n = (int(v) for v in config.cta) cluster_m, cluster_n = (int(v) for v in config.cluster) - # libheuristics encodes its two-CTA MMA with cluster_m == 2 while the - # CuTe DSL runner encodes it by doubling mma_tiler_mn[0]. - n_align = 32 if cta_n > 256 else 16 - use_2cta = cluster_m == 2 and cta_n % n_align == 0 - mma_tiler_mn = ((2 * cta_m if use_2cta else cta_m), cta_n) - signature.extend((use_2cta, mma_tiler_mn, (cluster_m, cluster_n))) + # The model reports a per-CTA tile. Both MMA instruction modes remain + # local choices for this tile and cluster. + signature.extend(((cta_m, cta_n), (cluster_m, cluster_n))) if "split_k" in fields: signature.append(max(1, int(config.split_k))) if "swizzle" in fields: @@ -581,7 +579,11 @@ def _bf16_tactic_signature(tactic: tuple, fields: set) -> Optional[tuple]: fields = _bf16_actionable_fields(fields, split_k) signature = [] if fields & {"tile", "cluster"}: - signature.extend((bool(tactic[1]), tuple(tactic[2]), tuple(tactic[3]))) + cta_group = 2 if tactic[1] else 1 + mma_tiler_m, mma_tiler_n = tactic[2] + if mma_tiler_m % cta_group: + return None + signature.extend(((mma_tiler_m // cta_group, mma_tiler_n), tuple(tactic[3]))) if "split_k" in fields: signature.append(split_k) if "swizzle" in fields: diff --git a/tensorrt_llm/_torch/custom_ops/torch_custom_ops.py b/tensorrt_llm/_torch/custom_ops/torch_custom_ops.py index f258d9c307dc..c32305a95a6c 100644 --- a/tensorrt_llm/_torch/custom_ops/torch_custom_ops.py +++ b/tensorrt_llm/_torch/custom_ops/torch_custom_ops.py @@ -33,9 +33,10 @@ from tensorrt_llm.logger import logger from tensorrt_llm.quantization.utils import fp8_quantize -from ..autotuner import (AutoTuner, ConstraintSpec, DistributedTuningStrategy, - DynamicTensorSpec, OptimizationProfile, TunableRunner, - TuningConfig, autotune) +from ..autotuner import (_TORCH_PROFILER_TIMER_KEY, AutoTuner, ConstraintSpec, + DistributedTuningStrategy, DynamicTensorSpec, + OptimizationProfile, TunableRunner, TuningConfig, + autotune) from ..cublaslt_utils import IS_CUBLASLT_AVAILABLE from ..cute_dsl_utils import (IS_CUTLASS_DSL_AVAILABLE, IS_CUTLASS_DSL_RUBIN_AVAILABLE) @@ -1509,13 +1510,17 @@ def forward( raise ValueError(f"Invalid tactic: {tactic}") def use_torch_profiler(self, tactic) -> bool: - # Keep one timing backend for every candidate in a mixed-backend sweep. - return "cutedsl" in self.allowed_backends + """Use CUPTI for a mixed-backend sweep only when explicitly requested.""" + return ("cutedsl" in self.allowed_backends and os.getenv( + "TLLM_PROFILING_TIMER", "").lower() == _TORCH_PROFILER_TIMER_KEY) def profiling_timer_cache_key(self) -> Optional[str]: - return "torch_profiler" if "cutedsl" in self.allowed_backends else None + """Tag mixed-backend results ranked with the opt-in CUPTI timer.""" + return _TORCH_PROFILER_TIMER_KEY if self.use_torch_profiler( + None) else None def tactic_search_cache_key(self): + """Distinguish candidate sets affected by the active NVMMH policy.""" if "cutedsl" not in self.allowed_backends or not IS_CUTLASS_DSL_AVAILABLE: return None return _cutedsl_nvmmh_tactic_search_cache_key() diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py index 14182772793a..503ec255ce2f 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_engine.py +++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py @@ -1166,9 +1166,13 @@ def cleanup(self) -> None: # Release model weights. release_gc() - AutoTuner.get()._release_nvmmh_policy(self) + self._release_autotuner_nvmmh_policy() self._cleanup_done = True + def _release_autotuner_nvmmh_policy(self) -> None: + """Release the policy pin after the executor stops using this engine.""" + AutoTuner.get()._release_nvmmh_policy(self) + def __del__(self) -> None: """Best-effort cleanup during garbage collection. diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py index 4bcbcfe65cbd..aac3f867c5ca 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py @@ -2035,9 +2035,12 @@ def shutdown(self): # to compute kv_cache_max_memory. cleanup() would set # model_engine.model = None, breaking that read with # `'NoneType' object has no attribute 'model_config'`. - # The engine's __del__ still calls cleanup() at terminal teardown - # (when the executor's reference is dropped), which is sufficient for - # the GMS daemon registry eviction the cleanup hook was added for. + # Release the NVMMH policy pin explicitly; cleanup() still runs when + # the engine is collected and releases its remaining resources. + for engine in (self.model_engine, self.draft_model_engine): + if engine is not None and hasattr( + engine, '_release_autotuner_nvmmh_policy'): + engine._release_autotuner_nvmmh_policy() del self.model_engine if self.draft_model_engine is not None: del self.draft_model_engine diff --git a/tensorrt_llm/llmapi/llm_args.py b/tensorrt_llm/llmapi/llm_args.py index be040cefab47..869abf9fb51a 100644 --- a/tensorrt_llm/llmapi/llm_args.py +++ b/tensorrt_llm/llmapi/llm_args.py @@ -6361,7 +6361,7 @@ class AutoTunerNvMMHConfig(StrictBaseModel): default=5, strict=True, description= - "Maximum number of NVMMH-ranked tile/cluster signatures to retain.") + "Maximum number of NVMMH-ranked selected-field signatures to retain.") @field_validator('fields') @classmethod diff --git a/tests/unittest/_torch/misc/test_autotuner.py b/tests/unittest/_torch/misc/test_autotuner.py index 60de2455f9fc..21dbd7285b72 100644 --- a/tests/unittest/_torch/misc/test_autotuner.py +++ b/tests/unittest/_torch/misc/test_autotuner.py @@ -901,47 +901,6 @@ def test_kernel_testing_single_context(): f"Expected 3 tactics to be tested, got {len(tested_tactics)}" -def test_kernel_testing_uses_runner_valid_tactics(): - """Verify captured tactic combinations use the runner-provided candidate list.""" - - class CaptureRunner(TunableRunner): - - def get_valid_tactics(self, inputs: List[FakeTensor], - profile: OptimizationProfile, - **kwargs) -> List[int]: - """Provide a fixed candidate list for capture and replay assertions.""" - return [1, 2] - - def forward(self, - /, - inputs: List[torch.Tensor], - *, - tactic: int = -1, - **kwargs) -> torch.Tensor: - """Return the input unchanged while the test observes tactic selection.""" - assert tactic in [-1, 0, 1, 2] - return inputs[0] - - x = torch.randn(4, 4) - runners = [CaptureRunner()] - tuning_config = TuningConfig() - tuner = AutoTuner.get() - tuner.clear_cache() - - with tuner.capture() as all_tactics: - tuner.choose_one("test_runner_valid_capture", runners, tuning_config, - [x]) - - captured = list(all_tactics) - assert [tactic for ((_, tactic), ) in captured] == [1, 2] - - for ((runner, tactic), ) in captured: - with tuner.replay(((runner, tactic), )): - replay_runner, replay_tactic = tuner.choose_one( - "test_runner_valid_capture", runners, tuning_config, [x]) - assert replay_runner is runner - assert replay_tactic == tactic - class MultiContextRunner(TunableRunner): @@ -1572,7 +1531,8 @@ def forward(self, def _skip_if_cupti_unavailable(error, context): """Skip CUPTI-only perf tests without hiding kernel/runtime failures.""" - if not AutoTuner._is_torch_profiler_unavailable_error(error): + if (not AutoTuner._is_torch_profiler_unavailable_error(error) + and not isinstance(error, autotuner._TorchProfilerEmptyTraceError)): raise error reason = f"{context} requires working torch.profiler/CUPTI: {error}" _CUPTI_PREFLIGHT_STATE[torch.cuda.current_device()] = reason @@ -1845,7 +1805,6 @@ def _cublaslt_call(): print(f"cuBLASLt: comparison_cupti_time={cublaslt_cupti_us:.3f} us") tolerance = 1.1 - cublas_tolerance = 1.1 assert heuristic_us <= sweep_us * tolerance, ( f"BF16 heuristic tactic {heuristic_tactic} ({heuristic_us:.3f} us) " f"is >{tolerance:.2f}x slower than full-sweep tactic " @@ -1855,11 +1814,8 @@ def _cublaslt_call(): f"({heuristic_cupti_us:.3f} us CUPTI) is >{tolerance:.2f}x slower " f"than full-sweep tactic {sweep_tactic} " f"({sweep_cupti_us:.3f} us CUPTI)") - assert heuristic_cupti_us <= cublaslt_cupti_us * cublas_tolerance, ( - f"BF16 CuTe DSL heuristic tactic {heuristic_tactic} " - f"({heuristic_cupti_us:.3f} us CUPTI) is " - f">{cublas_tolerance:.2f}x slower " - f"than cuBLASLt ({cublaslt_cupti_us:.3f} us CUPTI)") + print(f"BF16 CuTe DSL / cuBLASLt CUPTI ratio: " + f"{heuristic_cupti_us / cublaslt_cupti_us:.3f}x") def _run_cute_dsl_mxfp8_heuristic_comparison(monkeypatch, tuner, nvmmh): @@ -2003,7 +1959,6 @@ def _cublaslt_call(): print(f"cuBLASLt: comparison_cupti_time={cublaslt_cupti_us:.3f} us") tolerance = 1.1 - cublas_tolerance = 1.1 assert heuristic_us <= sweep_us * tolerance, ( f"MXFP8 heuristic tactic {heuristic_tactic} ({heuristic_us:.3f} us) " f"is >{tolerance:.2f}x slower than full-sweep tactic " @@ -2013,11 +1968,8 @@ def _cublaslt_call(): f"({heuristic_cupti_us:.3f} us CUPTI) is >{tolerance:.2f}x slower " f"than full-sweep tactic {sweep_tactic} " f"({sweep_cupti_us:.3f} us CUPTI)") - assert heuristic_cupti_us <= cublaslt_cupti_us * cublas_tolerance, ( - f"MXFP8 CuTe DSL heuristic tactic {heuristic_tactic} " - f"({heuristic_cupti_us:.3f} us CUPTI) is " - f">{cublas_tolerance:.2f}x slower " - f"than cuBLASLt ({cublaslt_cupti_us:.3f} us CUPTI)") + print(f"MXFP8 CuTe DSL / cuBLASLt CUPTI ratio: " + f"{heuristic_cupti_us / cublaslt_cupti_us:.3f}x") def _run_cute_dsl_nvfp4_heuristic_comparison(monkeypatch, tuner, nvmmh, @@ -2223,6 +2175,7 @@ def test_cute_dsl_nvmmh_matches_full_sweep(precision, supported_sms, if not IS_CUTLASS_DSL_RUBIN_AVAILABLE: pytest.skip("CuTe DSL Rubin support is not available") + monkeypatch.setenv("TLLM_PROFILING_TIMER", "torch_profiler") _require_cupti_cuda_activity() comparison, args = { @@ -2526,111 +2479,3 @@ def test_deep_gemm_jit_warmup_buckets(): assert deep_gemm_jit_warmup_buckets(16384) == deep_gemm_jit_warmup_buckets( 8192) -@pytest.mark.parametrize("conflicting_config", - [None, { - "fields": ("tile", ), - "max_tactics": 2 - }]) -def test_nvmmh_engine_policy_lifetime(nvmmh_config_guard, monkeypatch, - conflicting_config): - """Engine wiring pins non-default policy until the last owner is cleaned up.""" - from tensorrt_llm._torch.pyexecutor import model_engine - from tensorrt_llm.llmapi.llm_args import AutoTunerNvMMHConfig, TorchLlmArgs - - class EngineOwner: - """Minimal engine resources for exercising the real cleanup path.""" - _cleanup_done = False - model_loader = None - - def _release_cuda_graphs(self): - """No graph resources are allocated by this wiring test.""" - - monkeypatch.setattr(model_engine, "release_gc", lambda: None) - args = TorchLlmArgs.model_construct( - autotuner_nvmmh_config=AutoTunerNvMMHConfig( - fields=("swizzle", "cta_order"), max_tactics=3)) - first, second, rejected = EngineOwner(), EngineOwner(), EngineOwner() - tuner = nvmmh_config_guard - try: - model_engine._configure_autotuner_nvmmh(args, first) - expected = autotuner.NvMMHConfig(enabled=True, - fields=("swizzle", "cta_order"), - max_tactics=3) - assert tuner.nvmmh_config == expected - installed = tuner.nvmmh_config - model_engine._configure_autotuner_nvmmh(args, second) - assert tuner.nvmmh_config is installed - other_args = TorchLlmArgs.model_construct(autotuner_nvmmh_config=( - None if conflicting_config is None else AutoTunerNvMMHConfig( - **conflicting_config))) - with pytest.raises(ValueError, match="active model engines"): - model_engine._configure_autotuner_nvmmh(other_args, rejected) - assert tuner.nvmmh_config == expected - with pytest.raises(ValueError, match="active model engines"): - tuner.configure_nvmmh(enabled=False) - assert tuner.nvmmh_config == expected - - model_engine.PyTorchModelEngine.cleanup(first) - model_engine.PyTorchModelEngine.cleanup(first) # Idempotent release. - with pytest.raises(ValueError, match="active model engines"): - model_engine._configure_autotuner_nvmmh(other_args, rejected) - model_engine.PyTorchModelEngine.cleanup(second) - model_engine._configure_autotuner_nvmmh(other_args, rejected) - assert tuner.nvmmh_config != expected - model_engine.PyTorchModelEngine.cleanup(rejected) - finally: - tuner._release_nvmmh_policy(first) - tuner._release_nvmmh_policy(second) - tuner._release_nvmmh_policy(rejected) - - -def test_nvmmh_policy_owner_collection(nvmmh_config_guard): - """An abandoned or partially initialized engine cannot pin policy forever.""" - import gc - import weakref - - class Owner: - """Weak-referenceable stand-in for an engine that failed to initialize.""" - - owner = Owner() - tuner = nvmmh_config_guard - tuner._acquire_nvmmh_policy(owner, autotuner.NvMMHConfig(enabled=True)) - ref = weakref.ref(owner) - del owner - gc.collect() - assert ref() is None - tuner.configure_nvmmh(enabled=False) - assert not tuner.nvmmh_config.enabled - - -def test_nvmmh_unmatched_swap_preserves_split_k_admission( - nvmmh_config_guard, monkeypatch): - """A model miss for one swap must not resurrect rejected split-K factors.""" - from tensorrt_llm._torch.custom_ops import cute_dsl_custom_ops as ops - from tensorrt_llm._torch.custom_ops import \ - cutedsl_matmul_heuristics as heuristics - - if not hasattr(ops, "CuteDSLMXFP8RubinLinear"): - pytest.skip("Rubin CuTe DSL is unavailable") - monkeypatch.setattr(ops, "IS_NVMMH_AVAILABLE", True) - runner = ops.CuteDSLMXFP8RubinLinear(output_dtype=torch.bfloat16) - nvmmh_config_guard.configure_nvmmh(enabled=True, - fields=("tile", "cluster", "split_k")) - tactics = [("base", (128, 128, 128), (128, 128, 64), (1, 1), swap, False, - "static", "m", split) for swap in (False, True) - for split in (1, 2, 4, 8)] - - def rank_one_orientation(m, n, *args, **kwargs): - """Only the unswapped model problem has a representable result.""" - return ([heuristics.HeuristicConfig((128, 128), - (1, 1), 1, 0)] if m == 128 else []) - - monkeypatch.setattr(ops, "rank_configs", rank_one_orientation) - selected = runner._rank_prune_tactics(tactics, 128, 256, 1024) - admitted = [ - t for t in tactics if heuristics.is_sm107_nvmmh_split_k_eligible( - 1024, runner.mma_tiler_k, t[8]) - ] - assert len(admitted) < len(tactics) - assert {t[4] for t in selected} == {False, True} - assert set(selected) == set(admitted) From 701f543952cfae3a868b2bc6721826d6a63b8db6 Mon Sep 17 00:00:00 2001 From: peaceh <103117813+peaceh-nv@users.noreply.github.com> Date: Wed, 30 Sep 2026 08:08:42 -0700 Subject: [PATCH 5/8] [TRTLLM-14024][fix] Address follow-up nvMMH review findings Signed-off-by: peaceh <103117813+peaceh-nv@users.noreply.github.com> --- tensorrt_llm/_torch/autotuner.py | 26 ++++++++++--------- tensorrt_llm/_torch/pyexecutor/_util.py | 2 +- tensorrt_llm/_torch/pyexecutor/py_executor.py | 20 +++++++------- tests/unittest/_torch/misc/test_autotuner.py | 22 +++++++--------- 4 files changed, 35 insertions(+), 35 deletions(-) diff --git a/tensorrt_llm/_torch/autotuner.py b/tensorrt_llm/_torch/autotuner.py index c54899ad99c3..6eb4053a5580 100644 --- a/tensorrt_llm/_torch/autotuner.py +++ b/tensorrt_llm/_torch/autotuner.py @@ -646,25 +646,27 @@ def search_cache( if cache_key in self.cache: cached_runner_id, tactic, min_time = self.cache[cache_key] return True, idx, tactic, min_time - # A fresh inference process has not probed CUPTI yet. It may reuse - # an explicitly tagged CUDA-event fallback entry when no preferred - # CUPTI entry exists. Tuning mode leaves this disabled so a - # CUPTI-capable process re-profiles instead of reusing event-ranked - # tactics. - if (allow_timer_fallback - and not AutoTuner._torch_profiler_unavailable and - r.profiling_timer_cache_key() == _TORCH_PROFILER_TIMER_KEY): - fallback_key = self.get_cache_key( + # A CUPTI-capable inference process may reuse an event-ranked + # entry. Once CUPTI becomes unavailable, an earlier CUPTI-ranked + # entry remains valid and must remain reachable as well. + if (r.profiling_timer_cache_key() == _TORCH_PROFILER_TIMER_KEY + and (allow_timer_fallback + or AutoTuner._torch_profiler_unavailable)): + alternate_timer_key = ( + _TORCH_PROFILER_TIMER_KEY + if AutoTuner._torch_profiler_unavailable else + _CUDA_EVENT_FALLBACK_TIMER_KEY) + alternate_key = self.get_cache_key( custom_op, r, input_shapes, tuning_config, apply_map_to_tuning_buckets, - timer_key_override=_CUDA_EVENT_FALLBACK_TIMER_KEY, + timer_key_override=alternate_timer_key, ) - if fallback_key in self.cache: + if alternate_key in self.cache: cached_runner_id, tactic, min_time = self.cache[ - fallback_key] + alternate_key] return True, idx, tactic, min_time return False, *self.fallback_entry() diff --git a/tensorrt_llm/_torch/pyexecutor/_util.py b/tensorrt_llm/_torch/pyexecutor/_util.py index de9ab1ca3acc..374d925d0e24 100644 --- a/tensorrt_llm/_torch/pyexecutor/_util.py +++ b/tensorrt_llm/_torch/pyexecutor/_util.py @@ -1841,7 +1841,7 @@ def configure_kv_cache_capacity(self, kv_stats_draft.allocated_bytes if kv_stats_draft is not None else 0) py_executor.is_warmup = False - py_executor.shutdown() + py_executor.shutdown(release_nvmmh_policy=False) py_executor.enable_iter_perf_stats = origin_iter_stats py_executor.set_gather_responses(False) diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py index aac3f867c5ca..588662db38f7 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py @@ -1958,9 +1958,12 @@ def _stop_mpi_progress_thread(self) -> None: self._mpi_progress_thread = None self._mpi_progress_thread_exit_hook = None - def shutdown(self): - """ - Signals the server to shutdown. + def shutdown(self, *, release_nvmmh_policy: bool = True): + """Signal the server to shut down. + + Args: + release_nvmmh_policy: Keep false when the model engine will be + reused by the final executor after KV-cache estimation. """ node_server = getattr(self, "_openengine_node_server", None) if node_server is not None: @@ -2035,12 +2038,11 @@ def shutdown(self): # to compute kv_cache_max_memory. cleanup() would set # model_engine.model = None, breaking that read with # `'NoneType' object has no attribute 'model_config'`. - # Release the NVMMH policy pin explicitly; cleanup() still runs when - # the engine is collected and releases its remaining resources. - for engine in (self.model_engine, self.draft_model_engine): - if engine is not None and hasattr( - engine, '_release_autotuner_nvmmh_policy'): - engine._release_autotuner_nvmmh_policy() + if release_nvmmh_policy: + for engine in (self.model_engine, self.draft_model_engine): + if engine is not None and hasattr( + engine, '_release_autotuner_nvmmh_policy'): + engine._release_autotuner_nvmmh_policy() del self.model_engine if self.draft_model_engine is not None: del self.draft_model_engine diff --git a/tests/unittest/_torch/misc/test_autotuner.py b/tests/unittest/_torch/misc/test_autotuner.py index 21dbd7285b72..f9d4ac9a3dfb 100644 --- a/tests/unittest/_torch/misc/test_autotuner.py +++ b/tests/unittest/_torch/misc/test_autotuner.py @@ -1684,8 +1684,10 @@ def _run_cute_dsl_bf16_heuristic_comparison(monkeypatch, tuner, nvmmh): tuner.configure_nvmmh(enabled=True, fields=scheduler_fields, max_tactics=5) heuristic_candidates = runner.get_valid_tactics(inputs, None) assert 0 < len(heuristic_candidates) <= len(baseline_tactics), ( - "nvMMH did not prune SM107 BF16 tactics: " + "nvMMH returned an invalid SM107 BF16 candidate count: " f"{len(heuristic_candidates)} vs {len(baseline_tactics)}") + assert set(heuristic_candidates) != set(baseline_tactics), ( + "nvMMH scheduler-only guidance fell back to the full BF16 sweep") assert {_nvmmh_tactic_family(t) for t in heuristic_candidates }.issubset({_nvmmh_tactic_family(t) @@ -1860,8 +1862,10 @@ def _run_cute_dsl_mxfp8_heuristic_comparison(monkeypatch, tuner, nvmmh): tuner.configure_nvmmh(enabled=True, fields=scheduler_fields, max_tactics=5) heuristic_candidates = runner.get_valid_tactics(inputs, None) assert 0 < len(heuristic_candidates) <= len(baseline_tactics), ( - "nvMMH did not prune SM107 MXFP8 tactics: " + "nvMMH returned an invalid SM107 MXFP8 candidate count: " f"{len(heuristic_candidates)} vs {len(baseline_tactics)}") + assert set(heuristic_candidates) != set(baseline_tactics), ( + "nvMMH scheduler-only guidance fell back to the full MXFP8 sweep") assert {_nvmmh_tactic_family(t) for t in heuristic_candidates @@ -2040,7 +2044,7 @@ def _best_tactic(name, expected_candidates): # Pruned: nvMatmulHeuristics drives the (coupled) tile+cluster candidates. tuner.configure_nvmmh(enabled=True) heuristic_candidates = runner.get_valid_tactics(inputs, None) - assert 0 < len(heuristic_candidates) <= len(baseline_tactics), ( + assert 0 < len(heuristic_candidates) < len(baseline_tactics), ( f"nvMMH did not prune SM{sm_version} NVFP4 tactics: " f"{len(heuristic_candidates)} vs {len(baseline_tactics)}") assert {_nvmmh_tactic_family(t) @@ -2120,16 +2124,8 @@ def _cublas_call(): f">{tolerance:.2f}x slower than full-sweep tactic {sweep_tactic} " f"({sweep_us:.2f} us) for M={m}, N={n}, K={k}") - # The heuristic CuteDSL kernel should beat cuBLAS or be within tolerance. - # This is a cross-library comparison (CuteDSL vs cuBLASLt) at ~28us per call - # profiled over only 20 iterations; run-to-run jitter easily reaches a few - # percent from DVFS, L2 residency, and interleaving with cuBLAS autotune - # warmup, so keep this bound looser than the intra-CuteDSL one above. - cublas_tolerance = 1.10 - assert heuristic_us <= cublas_us * cublas_tolerance, ( - f"CuteDSL heuristic kernel ({heuristic_us:.2f} us) is " - f">{cublas_tolerance:.2f}x slower than cuBLASLt NVFP4 " - f"({cublas_us:.2f} us) for M={m}, N={n}, K={k}") + print(f"NVFP4 CuTe DSL / cuBLASLt CUPTI ratio: " + f"{heuristic_us / cublas_us:.3f}x") @pytest.mark.parametrize( From cdda22f647e2f7398a583146a5bb162acc85b9b5 Mon Sep 17 00:00:00 2001 From: peaceh <103117813+peaceh-nv@users.noreply.github.com> Date: Mon, 5 Oct 2026 08:47:52 -0700 Subject: [PATCH 6/8] [TRTLLM-14024][fix] Complete nvMMH review and pre-commit cleanup Signed-off-by: peaceh <103117813+peaceh-nv@users.noreply.github.com> --- tensorrt_llm/_torch/autotuner.py | 7 +++---- tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py | 10 +--------- .../parallel/test_cute_dsl_quant_dense_rubin.py | 6 ++++++ tests/unittest/_torch/misc/test_autotuner.py | 2 -- 4 files changed, 10 insertions(+), 15 deletions(-) diff --git a/tensorrt_llm/_torch/autotuner.py b/tensorrt_llm/_torch/autotuner.py index 6eb4053a5580..231e9a12d264 100644 --- a/tensorrt_llm/_torch/autotuner.py +++ b/tensorrt_llm/_torch/autotuner.py @@ -652,10 +652,9 @@ def search_cache( if (r.profiling_timer_cache_key() == _TORCH_PROFILER_TIMER_KEY and (allow_timer_fallback or AutoTuner._torch_profiler_unavailable)): - alternate_timer_key = ( - _TORCH_PROFILER_TIMER_KEY - if AutoTuner._torch_profiler_unavailable else - _CUDA_EVENT_FALLBACK_TIMER_KEY) + alternate_timer_key = (_TORCH_PROFILER_TIMER_KEY + if AutoTuner._torch_profiler_unavailable + else _CUDA_EVENT_FALLBACK_TIMER_KEY) alternate_key = self.get_cache_key( custom_op, r, diff --git a/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py b/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py index 25306c70c591..0c54c48682f6 100644 --- a/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py +++ b/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py @@ -14487,14 +14487,6 @@ def _split_k(tactic): **query_kwargs, ) ranked_configs_by_swap[swap_ab] = ranked_configs - valid_signatures = { - signature - for tactic in candidate_tactics - if (signature := self._nvmmh_tactic_signature( - tactic, match_fields)) is not None - } - if not valid_signatures: - return fallback_tactics def model_key(signature): """Drop locally swept cluster fields from a model-matching signature.""" @@ -14506,7 +14498,7 @@ def model_key(signature): valid_keys = { model_key(signature) - for signature in valid_signatures + for signature in baseline_signatures } config_by_key = {} rank1_config_by_swap = {} diff --git a/tests/unittest/_torch/kernels/parallel/test_cute_dsl_quant_dense_rubin.py b/tests/unittest/_torch/kernels/parallel/test_cute_dsl_quant_dense_rubin.py index efc54bf36225..e3a5348a3322 100644 --- a/tests/unittest/_torch/kernels/parallel/test_cute_dsl_quant_dense_rubin.py +++ b/tests/unittest/_torch/kernels/parallel/test_cute_dsl_quant_dense_rubin.py @@ -144,6 +144,12 @@ def _op_calls(): def test_sm107_quant_ops_registered_only_with_sm107_cute_dsl(): for name in SM107_QUANT_OPS: assert hasattr(torch.ops.trtllm, name) == IS_CUTLASS_DSL_RUBIN_AVAILABLE, name + if IS_CUTLASS_DSL_RUBIN_AVAILABLE: + from tensorrt_llm._torch.custom_ops import cute_dsl_custom_ops + + runner = cute_dsl_custom_ops.CuteDSLFp8RubinPerTensorGemmRunner + (dynamic_spec,) = runner.tuning_config.dynamic_tensor_specs + assert (dynamic_spec.input_idx, dynamic_spec.dim_idx) == (0, 0) @skip_no_sm107_cute_dsl diff --git a/tests/unittest/_torch/misc/test_autotuner.py b/tests/unittest/_torch/misc/test_autotuner.py index f9d4ac9a3dfb..e96685b12f70 100644 --- a/tests/unittest/_torch/misc/test_autotuner.py +++ b/tests/unittest/_torch/misc/test_autotuner.py @@ -901,7 +901,6 @@ def test_kernel_testing_single_context(): f"Expected 3 tactics to be tested, got {len(tested_tactics)}" - class MultiContextRunner(TunableRunner): def get_valid_tactics(self, inputs: List[FakeTensor], @@ -2474,4 +2473,3 @@ def test_deep_gemm_jit_warmup_buckets(): assert max(deep_gemm_jit_warmup_buckets(32)) < 128 assert deep_gemm_jit_warmup_buckets(16384) == deep_gemm_jit_warmup_buckets( 8192) - From 0b79cf6dc0f87ae4047b654a001b957f60581858 Mon Sep 17 00:00:00 2001 From: peaceh <103117813+peaceh-nv@users.noreply.github.com> Date: Thu, 8 Oct 2026 01:24:32 -0700 Subject: [PATCH 7/8] [TRTLLM-14024][fix] Isolate nvMMH setup in PCG fallback tests Mock autotuner policy configuration alongside the other unrelated engine initializers in the PCG routing test. Its partial SimpleNamespace args omit autotuner_nvmmh_config, so the new constructor hook fails before the routing assertions and unnecessarily touches the process-wide autotuner. Validation: reproduced all 15 parameterized failures on Hecate ARM before the fix; all 47 CPU warmup tests and two subtests pass afterward. The LLM args golden manifest check also passes after rebasing onto latest main. Signed-off-by: peaceh <103117813+peaceh-nv@users.noreply.github.com> --- .../unittest/_torch/executor/test_pytorch_model_engine_warmup.py | 1 + 1 file changed, 1 insertion(+) diff --git a/tests/unittest/_torch/executor/test_pytorch_model_engine_warmup.py b/tests/unittest/_torch/executor/test_pytorch_model_engine_warmup.py index bce4e0eb9b15..8176245a92cf 100644 --- a/tests/unittest/_torch/executor/test_pytorch_model_engine_warmup.py +++ b/tests/unittest/_torch/executor/test_pytorch_model_engine_warmup.py @@ -138,6 +138,7 @@ def test_pcg_fx_fallback_policy_is_model_specific( monkeypatch.setattr(_util, name, Mock(return_value=False)) monkeypatch.setattr(_util, "compute_max_num_sequences", Mock(return_value=4)) for name in ( + "_configure_autotuner_nvmmh", "_configure_deep_gemm_pdl", "create_input_processor", "setup_mm_encoder_attn_metadata", From ab4141ef85018e7c8e0129598a91e72797202908 Mon Sep 17 00:00:00 2001 From: peaceh <103117813+peaceh-nv@users.noreply.github.com> Date: Mon, 19 Oct 2026 03:03:29 -0700 Subject: [PATCH 8/8] [TRTLLM-14024][perf] Fix Rubin NVMMH tactic pruning gaps - BF16: keep each model split-K, drop the split cap of 8. - Block-scaled: materialize model clusters outside the default sweep. - Rank only runner-implementable CTA tiles via the NVMMH validity callback. - BF16: keep model tiles whose recommended cluster is illegal. - Sweep clusters locally around model tiles (MXFP8 rank-1 cluster-M; BF16 and FP8 BMM rank-1 plus <=2-CTA clusters). - Key the tactic-search cache on the NVMMH wheel version and a tactic schema number instead of a hand-bumped version string. - Support the nvMatmulHeuristics 0.2.0 Python API (getWithMNK, and the 0.1.0.27 callback setter) so plain BF16 queries still prune. Rubin, 29 document GEMM/BMM cases, same-GPU CUPTI vs the full sweep: none >5% slower (was +24.6% on UKV), every case prunes (31382 -> 5240 tactics), geomean vs the previous wheel 1.0005x -> 1.0204x. Same on the 0.2.0 wheel (max +2.28%). Signed-off-by: peaceh <103117813+peaceh-nv@users.noreply.github.com> --- .../_torch/custom_ops/cute_dsl_custom_ops.py | 412 +++++++++++------- .../custom_ops/cutedsl_matmul_heuristics.py | 293 +++++++++---- 2 files changed, 478 insertions(+), 227 deletions(-) diff --git a/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py b/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py index 0c54c48682f6..3dfbe5792a3d 100644 --- a/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py +++ b/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py @@ -36,8 +36,8 @@ FP8_BF16_OUTPUT_PRECISION, FP8_FP16_OUTPUT_PRECISION, FP8_LAYOUT, IN_PLACE_SPLIT_K_KIND, IS_NVMMH_AVAILABLE, NVFP4_LAYOUT, NVFP4_PRECISION, cute_raster_to_nvmmh_cta_order, expand_bf16_tactics_with_heuristic_splits, - filter_bf16_tactics, filter_fp8_tactics, is_sm107_nvmmh_split_k_eligible, - nvmmh_cta_order_to_cute_raster, rank_configs) + filter_bf16_tactics, filter_fp8_tactics, nvmmh_cta_order_to_cute_raster, + nvmmh_library_version, rank_configs) try: from cuda.bindings import driver as cuda @@ -112,20 +112,27 @@ def _kimi_k3_mxfp8_tuning_bucket(num_tokens: int) -> int: return next_positive_power_of_2(num_tokens) -_CUTEDSL_NVMMH_TACTIC_SEARCH_CACHE_VERSION = "cutedsl-nvmmh-v1" +# Bump only when NVMMH-emitted tactic tuples change shape, so that cached +# entries could no longer be parsed by the runners. Selection changes do not +# need a bump: cached winners stay valid tactics. +_CUTEDSL_NVMMH_TACTIC_CACHE_SCHEMA = 1 def _cutedsl_nvmmh_tactic_search_cache_key(): - """Return the cache discriminator for NVMMH-aware tactic generation.""" + """Return the cache discriminator for NVMMH-aware tactic generation. + + The key covers the policy (fields, tactic budget) and the installed NVMMH + wheel, whose ranking decides the candidate set. + """ config = AutoTuner.get().nvmmh_config if not config.enabled or not config.fields or not IS_NVMMH_AVAILABLE: # These paths preserve the pre-NVMMH candidate list, so policy details # cannot affect tactic search or fragment its cache identity. return None return ( - _CUTEDSL_NVMMH_TACTIC_SEARCH_CACHE_VERSION, - bool(config.enabled), - bool(IS_NVMMH_AVAILABLE), + "cutedsl-nvmmh", + _CUTEDSL_NVMMH_TACTIC_CACHE_SCHEMA, + nvmmh_library_version(), tuple(sorted(config.fields)), int(config.max_tactics), ) @@ -10308,14 +10315,20 @@ def forward( if c_needs_copy: c_tensor.copy_(c_buf) - def _rank_prune_tactics(self, tactics, m, n, k, output_dtype): + def _rank_prune_tactics(self, + tactics, + m, + n, + k, + output_dtype, + *, + allow_split_k: bool = True): """Rank/prune BF16 GEMM tactics with nvMatmulHeuristics. Called at the end of get_valid_tactics. The model controls only fields selected by the AutoTuner NVMMH policy. Emitted tuples are - from the runner's validated tactic list. On SM107, split-K itself - remains a bounded local CUPTI sweep because the installed NVMMH - cannot expose its native K-tile lower bound. + from the runner's validated tactic families. When split-K is + selected, each model configuration retains its own split count. Any query or matching failure keeps the validated tactic families. """ nvmmh_config = AutoTuner.get().nvmmh_config @@ -10324,8 +10337,8 @@ def _rank_prune_tactics(self, tactics, m, n, k, output_dtype): return tactics fields = set(nvmmh_config.fields) - if self.target_sm != "rubin": - # Main supports direct split-K only in the Rubin runner. + if self.target_sm != "rubin" or not allow_split_k: + # Direct reduction requires Rubin and a matching output shape. fields.discard("split_k") if not fields: return tactics @@ -10344,6 +10357,9 @@ def _rank_prune_tactics(self, tactics, m, n, k, output_dtype): # same runner-visible key. Over-query so those duplicates do not # crowd out valid CTA/cluster/split matches. query_count = max(max_tactics * 16, 64) + split_k_kind = (IN_PLACE_SPLIT_K_KIND + if self.target_sm == "rubin" + and "split_k" in fields else None) ranked_configs = rank_configs( m, n, @@ -10351,94 +10367,31 @@ def _rank_prune_tactics(self, tactics, m, n, k, output_dtype): precision, query_count, layout_name=BF16_LAYOUT, + split_k_kind=split_k_kind, ) if not ranked_configs: return fallback_tactics candidate_tactics = tactics - model_fields = fields - if self.target_sm == "rubin" and "split_k" in fields: - # NVMMH's native minTilesKLowerBound is not configurable in - # this wheel. Admit Rubin splits with TRT-LLM's shared - # lower bound, then let CUPTI select among the runner's - # 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])) - return 1 - - candidate_tactics = [ - tactic for tactic in candidate_tactics - if is_sm107_nvmmh_split_k_eligible( - k, self.nvmmh_split_k_cta_k, _split_k(tactic)) - ] - model_fields = fields - {"split_k"} - non_split_tactics = [ - tactic for tactic in candidate_tactics - if _split_k(tactic) == 1 - ] - split_tactics = [ - tactic for tactic in candidate_tactics - if _split_k(tactic) > 1 - ] - - if model_fields: - selected_non_split = filter_bf16_tactics( - non_split_tactics, - ranked_configs, - model_fields, - max_tactics, - include_known_good=not bool( - model_fields & {"swizzle", "cta_order"}), - fallback_tactics=[], - ) - if not selected_non_split: - return fallback_tactics - else: - selected_non_split = non_split_tactics - - # Rubin's split-K reduction path has fixed scheduler - # defaults, so swizzle/CTA order cannot participate in its - # model match. Tile/cluster still can when requested. - split_model_fields = model_fields - {"swizzle", "cta_order"} - if split_tactics and split_model_fields: - selected_split = filter_bf16_tactics( - split_tactics, - ranked_configs, - split_model_fields, - max_tactics, - include_known_good=False, - fallback_tactics=[], - ) - if not selected_split: - return fallback_tactics - else: - selected_split = split_tactics - selected = list( - dict.fromkeys(selected_non_split + selected_split)) - else: - if "split_k" in fields: - candidate_tactics = ( - expand_bf16_tactics_with_heuristic_splits( - tactics, ranked_configs)) - selected = filter_bf16_tactics( - candidate_tactics, - ranked_configs, - fields, - max_tactics, - include_known_good=not bool(fields - & {"swizzle", "cta_order"}), - fallback_tactics=fallback_tactics, - ) + if "split_k" in fields: + candidate_tactics = expand_bf16_tactics_with_heuristic_splits( + tactics, ranked_configs) + selected = filter_bf16_tactics( + candidate_tactics, + ranked_configs, + fields, + max_tactics, + include_known_good=not bool(fields + & {"swizzle", "cta_order"}), + fallback_tactics=fallback_tactics, + local_cluster_sweep=self.target_sm == "rubin", + ) logger.debug( f"CuTe DSL BF16 nvMatmulHeuristics: {len(tactics)} baseline " f"/ {len(candidate_tactics)} candidates " f"-> {len(selected)} tactics for M={m}, N={n}, K={k}, " f"precision={precision}, fields={sorted(fields)}, " - f"model_fields={sorted(model_fields)}") + f"model_fields={sorted(fields)}") return selected except Exception as e: # noqa: BLE001 - tuning must keep working logger.warning_once( @@ -10710,6 +10663,10 @@ class CuteDSLBf16RubinBmmRunner(CuteDSLBf16BlackwellBmmRunner): """SM107 BF16 batched GEMM: [B, M, K] @ [B, N, K] -> [B, M, N].""" kernel_cache = dict() + def tactic_search_cache_key(self): + """Include the batch-aware NVMMH policy in the profiling key.""" + return _cutedsl_nvmmh_tactic_search_cache_key() + # The output's M dim (inputs[2] dim 1) must track input0's bucketed # M: without the constraint, profiling-cache keys embed the raw # tuning-time M and every other runtime M misses the cache and falls @@ -10735,9 +10692,31 @@ def get_valid_tactics( return [] batch_size, m, k = inputs[0].shape n = inputs[1].shape[1] - return _sm107_bf16_valid_tactics( + tactics = _sm107_bf16_valid_tactics( m, n, k, batch_size, cutlass.BFloat16, _SM107_BF16_MMA_TILER_MN_CANDIDATES, None) + config = AutoTuner.get().nvmmh_config + fields = set(config.fields) & {"tile", "cluster"} + if not config.enabled or not IS_NVMMH_AVAILABLE or not fields: + return tactics + ranked_configs = rank_configs(m, + n, + k, + BF16_PRECISION, + max(config.max_tactics * 16, 64), + layout_name=BF16_LAYOUT, + batch_size=batch_size) + # The shared BF16 matcher includes an explicit unsplit field. + normalized = { + tactic + (1, ) if tactic[0] == "base" else tactic: tactic + for tactic in tactics + } + selected = filter_bf16_tactics(list(normalized), + ranked_configs, + fields, + config.max_tactics, + local_cluster_sweep=True) + return [normalized[tactic] for tactic in selected] def forward( self, @@ -11269,7 +11248,13 @@ def get_valid_tactics( tactics = _sm107_bf16_valid_tactics(m, n, k, 1, c_dtype_cutlass, mma_tiler_mn_candi, split_k_candi) - return self._rank_prune_tactics(tactics, m, n, k, inputs[2].dtype) + return self._rank_prune_tactics( + tactics, + m, + n, + k, + inputs[2].dtype, + allow_split_k=inputs[2].shape[0] == m) def forward( self, @@ -13528,6 +13513,18 @@ class CuteDSLFp8RubinBmmRunner(CuteDSLFp8BlackwellBmmRunner): kernel_class = SM107BlockwiseGemmKernel kernel_cache = dict() + def tactic_search_cache_key(self): + """Include the batch-aware NVMMH policy in the profiling key.""" + return _cutedsl_nvmmh_tactic_search_cache_key() + + @staticmethod + def _nvmmh_tactic_signature(tactic, fields): + """Map the blockwise BMM instruction tile to a per-CTA tile.""" + use_2cta, mma_tiler, _, cluster_shape = tactic + cta_group = 2 if use_2cta else 1 + return ((int(mma_tiler[0]) // cta_group, int(mma_tiler[1])), + tuple(cluster_shape)) + def __init__( self, output_dtype: torch.dtype = torch.bfloat16, @@ -13552,31 +13549,61 @@ def get_valid_tactics( batch_size, m, k = inputs[0].shape n = inputs[1].shape[1] - return [(use_2cta_instrs, mma_tiler, mma_inst_shape, - cluster_shape_mn) for use_2cta_instrs, mma_tiler, - mma_inst_shape in _SM107_FP8_BLOCKWISE_MMA_CANDIDATES - for cluster_shape_mn in - _SM107_FP8_BLOCKWISE_CLUSTER_CANDIDATES - if (not self.fp8_quantize_1x128 or mma_tiler[0] // - (2 if use_2cta_instrs else 1) == 128) - and self.__class__.kernel_class.can_implement( - cutlass.Float8E4M3FN, - cutlass.Float8E4M3FN, - cutlass.Float32, - (cutlass.Float8E4M3FN if self. - fp8_quantize_1x128 else cutlass.BFloat16), - use_2cta_instrs, - mma_tiler, - mma_inst_shape, - cluster_shape_mn, - m, - n, - k, - batch_size, - "k", - "k", - "n", - )] + tactics = [(use_2cta_instrs, mma_tiler, mma_inst_shape, + cluster_shape_mn) for use_2cta_instrs, mma_tiler, + mma_inst_shape in _SM107_FP8_BLOCKWISE_MMA_CANDIDATES + for cluster_shape_mn in + _SM107_FP8_BLOCKWISE_CLUSTER_CANDIDATES + if (not self.fp8_quantize_1x128 or mma_tiler[0] // + (2 if use_2cta_instrs else 1) == 128) + and self.__class__.kernel_class.can_implement( + cutlass.Float8E4M3FN, + cutlass.Float8E4M3FN, + cutlass.Float32, + (cutlass.Float8E4M3FN if self. + fp8_quantize_1x128 else cutlass.BFloat16), + use_2cta_instrs, + mma_tiler, + mma_inst_shape, + cluster_shape_mn, + m, + n, + k, + batch_size, + "k", + "k", + "n", + )] + + config = AutoTuner.get().nvmmh_config + fields = set(config.fields) & {"tile", "cluster"} + if (not config.enabled or not IS_NVMMH_AVAILABLE or not fields + or self.fp8_quantize_1x128): + return tactics + # Only rank tiles this blockwise kernel implements: for skinny + # M the unrestricted model returns CTA_M=64 tiles that match no + # tactic, so the full sweep was kept. Cluster-N stays swept + # within each matched (tile, cluster-M) family, as for the + # dense Rubin FP8 runner. + tiles = { + self._nvmmh_tactic_signature(tactic, fields)[0] + for tactic in tactics + } + ranked_configs = rank_configs(m, + n, + k, + FP8_BF16_OUTPUT_PRECISION, + max(config.max_tactics * 16, 64), + layout_name=FP8_LAYOUT, + batch_size=batch_size, + tiles=tiles) + return filter_fp8_tactics(tactics, + ranked_configs, + fields, + config.max_tactics, + self._nvmmh_tactic_signature, + sweep_cluster_n=True, + local_cluster_sweep=True) def forward( self, @@ -13975,6 +14002,8 @@ class CuteDSLBlockScaledRubinLinear(CuteDSLTunableRunner): nvmmh_layout: ClassVar[Optional[str]] = None nvmmh_supported_fields: ClassVar[frozenset] = frozenset() nvmmh_sweep_cluster_shape: ClassVar[bool] = False + # Sweep cluster-M locally inside the model's rank-1 tile family. + nvmmh_sweep_cluster_m: ClassVar[bool] = False nvmmh_supplement_mixed_clusters: ClassVar[bool] = False raster_order_candidates: ClassVar[Tuple[str, ...]] = ("m", ) scheduler_mode_candidates: ClassVar[Tuple[str, @@ -14404,15 +14433,80 @@ def _apply_nvmmh_scheduler(self, tactic, config, fields): annotated.append(swizzle_size) return tuple(annotated) + def _expand_nvmmh_clusters( + self, + tactics: List[Tuple], + ranked_configs_by_swap: dict, + fields: set, + m: int, + n: int, + real_k: int, + ) -> List[Tuple]: + """Materialize model clusters beyond the default sweep domain. + + Retain the runner's validated instruction/B-reuse templates. + A new cluster must pass kernel legality checks for the actual + problem and split-K. Mixed-cluster pairs remain runner-local. + """ + if (not fields & {"tile", "cluster"} + or self.nvmmh_sweep_cluster_shape): + return tactics + + templates_by_key = {} + for tactic in tactics: + if tactic[0] != "base": + continue + signature = self._nvmmh_tactic_signature(tactic, fields) + if signature is None: + continue + key = self._nvmmh_cluster_family(signature, fields, True) + templates_by_key.setdefault(key, []).append(tactic) + + expanded = dict.fromkeys(tactics) + validity = {} + for swap_ab, configs in ranked_configs_by_swap.items(): + kernel_m, kernel_n = (n, m) if swap_ab else (m, n) + for config in configs: + signature = self._nvmmh_config_signature( + config, fields, swap_ab) + key = self._nvmmh_cluster_family( + signature, fields, True) + cluster = tuple(config.cluster) + for tactic in templates_by_key.get(key, ()): + candidate = (*tactic[:3], cluster, *tactic[4:]) + if candidate in expanded: + continue + split_k = tactic[8] if len(tactic) >= 9 else 1 + check_key = (tactic[1], tactic[2], cluster, swap_ab, + split_k) + if check_key not in validity: + validity[ + check_key] = self.kernel_class.can_implement( + (kernel_m, kernel_n, real_k, 1), + self.a_dtype, + self.b_dtype, + self.sf_dtype, + cutlass.BFloat16, + "k", + "k", + "m" if swap_ab else "n", + self.sf_vec_size, + tactic[1], + tactic[2], + cluster, + split_k, + ) + if validity[check_key]: + expanded[candidate] = None + return list(expanded) + def _rank_prune_tactics(self, tactics, m, n, real_k): """Rank/prune Rubin block-scaled tactics with NVMMH. - Split-K uses TRT-LLM's SM107 K-tile admission policy and remains - a local CUPTI sweep field. This is necessary because the native - scheduler's hard-coded lower bound can prevent useful factors - from being modeled at all. NVMMH still controls every other - selected field. Scheduler values are copied directly from the - matched model config instead of being locally swept. + Model cluster shapes are materialized and validated even when + absent from the default sweep. Selected split-K values still + participate in structural matching. Scheduler values are copied + from the matched config. Unsupported results retain fallback. MXFP8 uses the generic E4M3 FP8 compute model because NVMMH does not encode scale-factor formats. Runner-native fields @@ -14429,10 +14523,7 @@ def _rank_prune_tactics(self, tactics, m, n, real_k): if not fields: return tactics fallback_tactics = tactics - # The installed NVMMH has no public property for overriding - # minTilesKLowerBound. Treat split-K as an admitted local sweep - # dimension instead of matching its native rank-1 value. - model_fields = fields - {"split_k"} + model_fields = fields scheduler_fields = model_fields & {"swizzle", "cta_order"} match_fields = model_fields - scheduler_fields @@ -14440,21 +14531,6 @@ def _rank_prune_tactics(self, tactics, m, n, real_k): m, n, real_k = int(m), int(n), int(real_k) max_tactics = nvmmh_config.max_tactics candidate_tactics = tactics - 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): - return int(tactic[8]) - return 1 - - candidate_tactics = [ - tactic for tactic in candidate_tactics - if is_sm107_nvmmh_split_k_eligible( - real_k, self.mma_tiler_k, _split_k(tactic)) - ] if not model_fields: return candidate_tactics or fallback_tactics @@ -14470,6 +14546,16 @@ def _split_k(tactic): valid_swaps = sorted( {signature[-1] for signature in baseline_signatures}) + # Per-CTA tiles each orientation can launch. The model + # otherwise ranks tiles this kernel never implements + # (CTA_M=64, CTA_N=32/96) for skinny shapes, and nothing + # matches. Clusters stay open: model clusters outside the + # sweep are materialized and validated below. + tiles_by_swap = {} + if match_fields & {"tile", "cluster"}: + for signature in baseline_signatures: + tiles_by_swap.setdefault(signature[-1], + set()).add(signature[0]) query_count = max(max_tactics * 16, 64) ranked_configs_by_swap = {} for swap_ab in valid_swaps: @@ -14478,6 +14564,8 @@ def _split_k(tactic): if "split_k" in fields: query_kwargs["split_k_kind"] = ( IN_PLACE_SPLIT_K_KIND) + if swap_ab in tiles_by_swap: + query_kwargs["tiles"] = tiles_by_swap[swap_ab] ranked_configs = rank_configs( kernel_m, kernel_n, @@ -14488,8 +14576,36 @@ def _split_k(tactic): ) ranked_configs_by_swap[swap_ab] = ranked_configs + candidate_tactics = self._expand_nvmmh_clusters( + tactics, ranked_configs_by_swap, match_fields, m, n, + real_k) + candidate_signatures = { + signature + for tactic in candidate_tactics + if (signature := self._nvmmh_tactic_signature( + tactic, match_fields)) is not None + } + + # The model's top tile is right far more often than its + # cluster: small-M MXFP8 ranks only cluster 4x1 although + # 1x1/2x1 often win. Sweep cluster-M inside the rank-1 + # tile family of each orientation; other tiles keep the + # model cluster so the pool stays small. + sweep_tile_by_swap = {} + if (self.nvmmh_sweep_cluster_m + and match_fields & {"tile", "cluster"} + and not self.nvmmh_sweep_cluster_shape): + for swap_ab, configs in ranked_configs_by_swap.items(): + if configs: + sweep_tile_by_swap[swap_ab] = tuple( + configs[0].cta) + def model_key(signature): """Drop locally swept cluster fields from a model-matching signature.""" + if sweep_tile_by_swap.get( + signature[-1]) == signature[0]: + cta, cluster, *remaining = signature + return (cta, int(cluster[1]), *remaining) return self._nvmmh_cluster_family( signature, match_fields, @@ -14498,7 +14614,7 @@ def model_key(signature): valid_keys = { model_key(signature) - for signature in baseline_signatures + for signature in candidate_signatures } config_by_key = {} rank1_config_by_swap = {} @@ -14542,7 +14658,7 @@ def model_key(signature): if annotated is not None: selected.append(annotated) selected.extend( - tactic for tactic in candidate_tactics + tactic for tactic in fallback_tactics if ((signature := self._nvmmh_tactic_signature( tactic, match_fields)) is not None and signature[-1] in unmatched_swaps)) @@ -14570,6 +14686,8 @@ def model_key(signature): config = rank1_config_by_swap.get(swap_ab) if config is None: continue + if "split_k" in fields and int(config.split_k) != 1: + continue annotated = self._apply_nvmmh_scheduler( tactic, config, scheduler_fields) if annotated is not None: @@ -15120,11 +15238,11 @@ class CuteDSLMXFP8RubinLinear(CuteDSLBlockScaledRubinLinear): nvmmh_precision = FP8_BF16_OUTPUT_PRECISION nvmmh_layout = FP8_LAYOUT # NVMMH can prune tile/cluster choices or fix the rank-1 scheduler - # tuple. Split-K uses the shared SM107 admission rule and remains - # locally profiled. + # tuple. Selected split-K values follow the model recommendation. nvmmh_supported_fields = frozenset( {"tile", "cluster", "swizzle", "cta_order", "split_k"}) nvmmh_supplement_mixed_clusters = True + nvmmh_sweep_cluster_m = True a_dtype = cutlass.Float8E4M3FN b_dtype = cutlass.Float8E4M3FN sf_dtype = cutlass.Float8E8M0FNU diff --git a/tensorrt_llm/_torch/custom_ops/cutedsl_matmul_heuristics.py b/tensorrt_llm/_torch/custom_ops/cutedsl_matmul_heuristics.py index b4c1b2f87e83..921f2601a48e 100644 --- a/tensorrt_llm/_torch/custom_ops/cutedsl_matmul_heuristics.py +++ b/tensorrt_llm/_torch/custom_ops/cutedsl_matmul_heuristics.py @@ -15,9 +15,10 @@ """ import ctypes +import importlib.metadata import re from functools import lru_cache -from typing import Callable, List, NamedTuple, Optional, Tuple +from typing import Callable, Iterable, List, NamedTuple, Optional, Tuple from tensorrt_llm.logger import logger @@ -29,6 +30,20 @@ _nvmmh = None IS_NVMMH_AVAILABLE = False + +@lru_cache(maxsize=1) +def nvmmh_library_version() -> str: + """Return the installed nvMatmulHeuristics distribution version. + + Rankings change between wheels, so callers include this in persistent + tactic-search cache keys. + """ + try: + return importlib.metadata.version("nvidia-matmul-heuristics") + except importlib.metadata.PackageNotFoundError: + return getattr(_nvmmh, "__version__", "unknown") + + # NVFP4 dense GEMM precision/layout for nvMatmulHeuristics, fixed to the # CuteDSLNVFP4BlackwellRunner kernel: A/B are Float4E2M1FN (fp4), the output is # BFloat16, and A/B are K-major (see the can_implement call in @@ -56,11 +71,6 @@ FP8_FP16_OUTPUT_PRECISION = "QQHSH" FP8_LAYOUT = "TN_ROW_MAJOR" -# nvMatmulHeuristics often recommends non-power-of-two split counts. The BF16 -# filter accepts model-recommended values through this conservative cap after a -# successful query; get_valid_tactics itself remains unchanged. -BF16_MAX_HEURISTIC_SPLIT_K = 8 - # Rubin's native split-K scheduler currently requires more than 20 K tiles per # split. That lower bound hides useful split factors for the relatively small-K # CuTe DSL kernels used by TRT-LLM. Keep the Rubin-specific override here so @@ -119,11 +129,11 @@ class HeuristicConfig(NamedTuple): def nvmmh_cta_order_to_cute_raster(cta_order: int) -> str: """Translate CUTLASS3 nvMMH CTA order to CuTe's raster name. - CUTLASS's nvMMH adapter maps order 0 to ``along_m`` and order 1 to - ``along_n`` (python/cutlass_library/heuristics_provider.py in CUTLASS, - see https://github.com/NVIDIA/cutlass/blob/main/python/cutlass_library/heuristics_provider.py#L222). - Keep this conversion at the adapter boundary because the public nvMMH - header's row/column wording differs from this CUTLASS3 mapping. + Order 0 is ``AlongM`` and order 1 is ``AlongN``, as in CUTLASS's nvMMH + provider (https://github.com/NVIDIA/cutlass/blob/78e46228a1bfc5c972e2f9a634223af784cd6949/python/cutlass_library/heuristics_provider.py#L222) + and nvMMH's own CUTLASS swizzler model. Keep this conversion at the + adapter boundary because the public nvMMH header's row/column wording + differs from this CUTLASS3 mapping. """ if cta_order == 0: return "m" @@ -141,21 +151,6 @@ def cute_raster_to_nvmmh_cta_order(raster_along: str) -> int: raise ValueError(f"Unsupported CuTe raster {raster_along!r}; expected 'm' or 'n'") -def _pin_split_k_to_rank1( - ranked_configs: List[HeuristicConfig], -) -> List[HeuristicConfig]: - """Use rank-1 split-K for every precision while preserving other fields.""" - if not ranked_configs: - return ranked_configs - - rank1_split_k = max(1, int(ranked_configs[0].split_k)) - pinned = [ - config if int(config.split_k) == rank1_split_k else config._replace(split_k=rank1_split_k) - for config in ranked_configs - ] - return list(dict.fromkeys(pinned)) - - @lru_cache(maxsize=None) def _get_interface(precision: str): """Build and cache an interface for a given precision (None on failure).""" @@ -293,6 +288,38 @@ def _extract_config(config) -> Optional[HeuristicConfig]: return HeuristicConfig(cta, cluster, max(1, swizzle), cta_order, max(1, split_k)) +def _set_tile_validity_check(interface, backend, tiles: frozenset) -> None: + """Restrict a backend's search to the per-CTA tiles a runner can launch. + + NVMMH calls the check for every candidate it would otherwise rank. Without + it the model can spend its whole ranking on tiles the CuTe DSL kernel does + not implement (e.g. CTA_M=64 on block-scaled Rubin GEMMs), and the caller + then falls back to the full sweep. Clusters stay unrestricted: runners + either materialize model clusters or substitute legal ones per tile. + """ + + def admit(config) -> int: + return int((int(config.cta[0]), int(config.cta[1])) in tiles) + + def admit_pointer(config_ptr, _problem_ptr) -> int: + return admit(config_ptr.contents) + + set_callback = getattr(interface, "backendSetCallbackProperty", None) + callback = admit + if set_callback is None: + # nvMatmulHeuristics 0.1.0.27 spells it setBackendCallbackProperty and + # passes a configuration pointer plus an unused problem pointer. + set_callback = getattr(interface, "setBackendCallbackProperty", None) + callback = admit_pointer + if set_callback is None: + raise AttributeError("nvMatmulHeuristics interface cannot set a validity callback") + set_callback( + backend, + _nvmmh.NvMatmulHeuristicsBackendPropertyCallbackKind.KERNEL_ADDITIONAL_VALIDITY_CHECK, + callback, + ) + + def rank_configs( m: int, n: int, @@ -301,13 +328,22 @@ def rank_configs( count: int, layout_name: str = NVFP4_LAYOUT, split_k_kind: Optional[str] = None, + batch_size: int = 1, + tiles: Optional[Iterable[Tuple[int, int]]] = None, ) -> List[HeuristicConfig]: """Return up to ``count`` HeuristicConfig entries ranked best-first. - Every returned precision uses the first parsed rank's split-K value. Returns - an empty list if heuristics are unavailable or the query fails, so - callers can fall back to their full tactic list. + Each entry retains its own model-recommended split-K value. Returns an + empty list if heuristics are unavailable or the query fails, so callers + can fall back to their full tactic list. In-place queries preserve the + backend's native split-K search range. + + ``tiles`` lists the per-CTA (M, N) tiles the caller can launch. When + given, the model only ranks configurations with one of those tiles. """ + tiles = None if tiles is None else frozenset((int(m_), int(n_)) for m_, n_ in tiles) + if tiles is not None and not tiles: + return [] interface = _get_interface(precision) if interface is None: return [] @@ -320,20 +356,46 @@ def rank_configs( interface.loadInternalDiscoverySet(layout, None) except Exception as e: # noqa: BLE001 - discovery data is optional logger.debug(f"[nvMatmulHeuristics] loadInternalDiscoverySet skipped: {e}") - if split_k_kind is None: - configs = interface.get_with_mnk(int(m), int(n), int(k), layout, int(count), None) + + 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: + return method + raise AttributeError(f"nvMatmulHeuristics interface has none of {names!r}") + + if split_k_kind is None and batch_size == 1 and tiles is None: + get_with_mnk = getattr(interface, "get_with_mnk", None) + if get_with_mnk is not None: + configs = get_with_mnk(int(m), int(n), int(k), layout, int(count), None) + else: + configs = _method("getWithMNK")((int(m), int(n), int(k)), layout, int(count)) + elif split_k_kind is None: + backend = _method("createBackend", "create_backend")( + _nvmmh.NvMatmulHeuristicsTarget.CUTLASS3 + ) + try: + if tiles is not None: + _set_tile_validity_check(interface, backend, tiles) + make_problem = _method( + "makeNvMatmulHeuristicsProblem", "make_nv_matmul_heuristics_problem" + ) + try: + problem = make_problem( + (int(m), int(n), int(k)), layout, batch_size=int(batch_size) + ) + except TypeError: + problem = make_problem( + int(m), int(n), int(k), layout, batch_size=int(batch_size) + ) + configs = _method("getEx", "get_ex")(problem, int(count), backend) + finally: + _method("destroyBackend", "destroy_backend")(backend) else: if split_k_kind != IN_PLACE_SPLIT_K_KIND: 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: - return method - raise AttributeError(f"nvMatmulHeuristics interface has none of {names!r}") - backend = None try: create_backend = _method("createBackend", "create_backend") @@ -361,14 +423,20 @@ def _method(*names): ctypes.byref(property_buffer), ctypes.sizeof(property_buffer), ) + if tiles is not None: + _set_tile_validity_check(interface, backend, tiles) make_problem = _method( "makeNvMatmulHeuristicsProblem", "make_nv_matmul_heuristics_problem" ) try: - problem = make_problem((int(m), int(n), int(k)), layout, batch_size=1) + problem = make_problem( + (int(m), int(n), int(k)), layout, batch_size=int(batch_size) + ) except TypeError: # Older wheels take M, N, and K as separate arguments. - problem = make_problem(int(m), int(n), int(k), layout, batch_size=1) + problem = make_problem( + int(m), int(n), int(k), layout, batch_size=int(batch_size) + ) get_ex = _method("getEx", "get_ex") configs = get_ex(problem, int(count), backend) finally: @@ -377,7 +445,8 @@ def _method(*names): except Exception as e: # noqa: BLE001 - any failure must degrade gracefully logger.warning_once( f"[nvMatmulHeuristics] Query failed for " - f"precision={precision}, mnk=({m},{n},{k}), " + f"precision={precision}, mnk=({m},{n},{k}), batch_size={batch_size}, " + f"tiles={None if tiles is None else sorted(tiles)}, " f"split_k_kind={split_k_kind}: {e}. " f"Falling back to full tactic list.", key="nvmmh_query_failure", @@ -395,13 +464,15 @@ def _runtime(c): # ``runtime`` value may be estimated with the interface's default backend, # so sorting it again could reorder the IN_PLACE split-K recommendation # using OUT_OF_PLACE costs. + # A tile restriction only filters candidates, so non-split queries keep the + # same runtime ordering with or without it. ordered = list(configs) if split_k_kind is not None else sorted(configs, key=_runtime) ranked: List[HeuristicConfig] = [] for c in ordered: cfg = _extract_config(c) if cfg is not None and cfg not in ranked: ranked.append(cfg) - return _pin_split_k_to_rank1(ranked) + return ranked def rank_tile_cluster_configs( @@ -428,12 +499,16 @@ def filter_fp8_tactics( max_tactics: int, tactic_signature: Callable[[tuple, set], Optional[tuple]], sweep_cluster_n: bool = False, + local_cluster_sweep: bool = False, ) -> List[tuple]: """Keep FP8 tactics matching the top model-visible field signatures. Instruction shape, B-reuse, TMA-store mode, and other runner-local fields stay swept because ``tactic_signature`` deliberately omits them. Rubin may additionally sweep cluster-N within each selected (CTA, cluster-M) family. + ``local_cluster_sweep`` also keeps every cluster of the rank-1 tile and the + small (<= 2 CTA) clusters of the other selected tiles; see + ``_widened_cluster``. An empty query or intersection returns the original validated tactic list. """ if not tactics or not ranked_configs or not fields: @@ -456,6 +531,11 @@ def filter_fp8_tactics( if not kept_signatures: return tactics + widen = ( + _widened_cluster(kept_signatures) + if local_cluster_sweep and fields & {"tile", "cluster"} + else (lambda signature: False) + ) if sweep_cluster_n and fields & {"tile", "cluster"}: def _cluster_n_family(signature): @@ -467,43 +547,55 @@ def _cluster_n_family(signature): selected = [] for tactic in tactics: signature = tactic_signature(tactic, fields) - if signature is not None and _cluster_n_family(signature) in kept_families: + if signature is not None and ( + _cluster_n_family(signature) in kept_families or widen(signature) + ): selected.append(tactic) else: kept = set(kept_signatures) - selected = [tactic for tactic in tactics if tactic_signature(tactic, fields) in kept] + selected = [] + for tactic in tactics: + signature = tactic_signature(tactic, fields) + if signature is not None and (signature in kept or widen(signature)): + selected.append(tactic) return selected if selected else tactics +def _widened_cluster(ranked_signatures: List[tuple]) -> Callable[[tuple], bool]: + """Admit locally swept clusters around model-ranked (tile, cluster) keys. + + On Rubin the model ranks tiles well but over-values large clusters (4x1, + 4x2, 2x4), while the measured winner often uses the same tile with a + 1x1/1x2/2x1 cluster. Keep every cluster of the rank-1 tile and the + small clusters of the other ranked tiles. Keys are (tile, cluster, *rest); + ``rest`` (e.g. a model CTA order) must match a ranked key of that tile. + """ + if not ranked_signatures: + return lambda signature: False + rank1_tile = ranked_signatures[0][0] + ranked_rest = {(signature[0], tuple(signature[2:])) for signature in ranked_signatures} + + def admit(signature: tuple) -> bool: + cta, cluster, *rest = signature + if (cta, tuple(rest)) not in ranked_rest: + return False + return cta == rank1_tile or int(cluster[0]) * int(cluster[1]) <= 2 + + return admit + + def expand_bf16_tactics_with_heuristic_splits( tactics: List[tuple], ranked_configs: List[HeuristicConfig] ) -> List[tuple]: - """Add only the rank-1 model split count after a successful BF16 query. + """Make model-recommended split counts available for per-config matching. The runner has already validated each base tactic's tile and cluster. Split - count does not participate in ``can_implement``, so a six-field split tactic - can be derived safely from those validated templates. Baseline tactics are - returned byte-for-byte when the model has no supported split recommendation. + count does not participate in ``can_implement``. The caller must ensure + its output shape supports direct reduction before expanding the validated + templates. Every positive model-recommended split is retained, including + values absent from the baseline sweep. """ - # ``get_valid_tactics`` only emits a split greater than one when the - # runner's split-K shape guard passes. Preserve that eligibility decision - # before adding non-power-of-two model recommendations. - split_k_eligible = any( - isinstance(tactic, tuple) - and len(tactic) == 6 - and tactic[0] == "base" - and int(tactic[5]) > 1 - for tactic in tactics - ) - if not split_k_eligible: - return tactics - - ranked_configs = _pin_split_k_to_rank1(ranked_configs) - recommended = { - int(config.split_k) - for config in ranked_configs - if 1 < int(config.split_k) <= BF16_MAX_HEURISTIC_SPLIT_K - } + recommended = {int(config.split_k) for config in ranked_configs if int(config.split_k) > 1} if not recommended: return tactics @@ -622,6 +714,11 @@ def _bf16_structural_config_signature(config: HeuristicConfig, fields: set) -> t return _bf16_config_signature(config, structural_fields) +def _bf16_tile_family(signature: tuple) -> tuple: + """Drop the cluster from a tile/cluster-led BF16 key; keep the rest.""" + return (signature[0],) + tuple(signature[2:]) + + def _apply_bf16_scheduler_config(tactic: tuple, config: HeuristicConfig, fields: set) -> tuple: """Copy selected scheduler values from one NVMMH config to a base tactic.""" split_k = max(1, int(tactic[5])) @@ -643,40 +740,76 @@ def filter_bf16_tactics( max_tactics: int, include_known_good: bool = True, fallback_tactics: Optional[List[tuple]] = None, + local_cluster_sweep: bool = False, ) -> List[tuple]: """Return top model-matched BF16 tactics with direct scheduler annotation. + ``local_cluster_sweep`` (Rubin) also keeps every legal cluster of the + rank-1 matched tile and the small (<= 2 CTA) clusters of the other + matched tiles; see ``_widened_cluster``. + Only base tactics participate in model matching. Valid preferred/fallback cluster schedulers cannot be represented directly by the model's one cluster field, so only the validated q_b known-good tactic is supplemented when its - implicit split-K=1 agrees with rank 1. When split-K is selected, every lower - rank keeps its other fields but uses the rank-1 split count. An empty + implicit split-K=1 agrees with a matched model configuration. Each model + configuration keeps its own split count during matching. An empty intersection returns ``fallback_tactics`` (normally the exact list produced by ``get_valid_tactics``). """ fallback = tactics if fallback_tactics is None else fallback_tactics if not tactics or not ranked_configs or not fields: return fallback - allow_known_good_supplement = include_known_good - if "split_k" in fields: - ranked_configs = _pin_split_k_to_rank1(ranked_configs) - allow_known_good_supplement = include_known_good and int(ranked_configs[0].split_k) == 1 - valid_signatures = { signature for tactic in tactics if (signature := _bf16_structural_tactic_signature(tactic, fields)) is not None } + # The model sometimes ranks a supported tile with a cluster the kernel + # cannot launch (SM107 rejects CTA_N=256 with cluster_n=2). Dropping that + # config discards the tile entirely, even when it is the fastest one, so + # keep the tile and let the autotuner pick among its legal clusters. The + # substituted family counts as one model key. + substitute_cluster = bool(fields & {"tile", "cluster"}) + signatures_by_tile = {} + if substitute_cluster: + for signature in valid_signatures: + signatures_by_tile.setdefault(_bf16_tile_family(signature), set()).add(signature) config_by_signature = {} + model_keys = 0 for config in ranked_configs: signature = _bf16_structural_config_signature(config, fields) - if signature in valid_signatures and signature not in config_by_signature: - config_by_signature[signature] = config - if len(config_by_signature) >= max(1, int(max_tactics)): - break + if signature in valid_signatures: + matched = [signature] + elif substitute_cluster: + matched = sorted(signatures_by_tile.get(_bf16_tile_family(signature), ())) + else: + matched = [] + matched = [s for s in matched if s not in config_by_signature] + if not matched: + continue + for s in matched: + config_by_signature[s] = config + model_keys += 1 + if model_keys >= max(1, int(max_tactics)): + break if not config_by_signature: return fallback - + if local_cluster_sweep and substitute_cluster: + # Insertion order follows the model ranking, so the first key is the + # rank-1 matched tile family. + ranked_signatures = list(config_by_signature) + config_by_family = {} + for signature, config in config_by_signature.items(): + config_by_family.setdefault(_bf16_tile_family(signature), config) + widen = _widened_cluster(ranked_signatures) + for signature in sorted(valid_signatures): + if signature not in config_by_signature and widen(signature): + config_by_signature[signature] = config_by_family[_bf16_tile_family(signature)] + + allow_known_good_supplement = include_known_good and ( + "split_k" not in fields + or any(int(config.split_k) == 1 for config in config_by_signature.values()) + ) selected = [] for tactic in tactics: signature = _bf16_structural_tactic_signature(tactic, fields)