diff --git a/docs/source/torch/adding_custom_kernels.md b/docs/source/torch/adding_custom_kernels.md index 97aaa3ffac94..1aaef8a06cc9 100644 --- a/docs/source/torch/adding_custom_kernels.md +++ b/docs/source/torch/adding_custom_kernels.md @@ -260,27 +260,111 @@ 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 | `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 | + +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). +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 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 +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..f6539c6b239e 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: false +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..29b3eab0b4e4 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: false +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..231e9a12d264 100644 --- a/tensorrt_llm/_torch/autotuner.py +++ b/tensorrt_llm/_torch/autotuner.py @@ -20,13 +20,17 @@ import inspect import itertools 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 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 +45,16 @@ # 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 _TorchProfilerEmptyTraceError(RuntimeError): + """Raised when one completed profile contains no CUDA activity.""" class DistributedTuningStrategy(enum.Enum): @@ -152,6 +166,72 @@ 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: + """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, + 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 +351,42 @@ 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 with optional CUPTI timing.""" + + def use_torch_profiler(self, tactic: Any) -> bool: + """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 opt-in CUPTI-ranked tactics from legacy cache entries.""" + return _TORCH_PROFILER_TIMER_KEY if self.use_torch_profiler( + None) else None + @contextlib.contextmanager def autotune(tune_mode: bool = True, @@ -503,6 +619,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 +636,37 @@ 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 + 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 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=alternate_timer_key, + ) + if alternate_key in self.cache: + cached_runner_id, tactic, min_time = self.cache[ + alternate_key] + return True, idx, tactic, min_time return False, *self.fallback_entry() @@ -535,8 +677,10 @@ 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 ( + """Combine operation, runner, shapes, search policy, and timing backend.""" + cache_key = ( custom_op, runner.__class__.__name__, str(runner.unique_id()), @@ -548,6 +692,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,8 +1081,16 @@ 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' @@ -939,13 +1101,17 @@ 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._force_cuda_event_for_sweep = False 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 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": @@ -976,6 +1142,78 @@ 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. Active model engines pin + this policy: a differing update raises without changing any state. + """ + 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) + + 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. @@ -1026,7 +1264,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 +1280,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 +1402,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 +1491,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 +1575,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 +1588,10 @@ 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 + 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: + 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,14 +1731,61 @@ 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 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}." @@ -1521,6 +1821,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 _TorchProfilerEmptyTraceError( + "torch.profiler reported no CUDA activity for this 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: + """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 " + "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 +1885,11 @@ 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) + """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): if use_cuda_graph: @@ -1601,6 +1907,104 @@ 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(): + """Replay the captured batch or launch the same batch eagerly.""" + if use_cuda_graph: + graph.replay() + else: + for r in range(repeat): + runner( + input_tensor_batches[r % + len(input_tensor_batches)], + tactic=tactic, + **kwargs, + ) + + 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: + 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 isinstance(error, _TorchProfilerEmptyTraceError): + raise + 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: + start_ts = torch.empty(1, dtype=torch.int64, device="cuda") + 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 + else: + start_evt = torch.cuda.Event(enable_timing=True) + 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. if use_cuda_graph: delay_kernel(self._CUDA_GRAPH_DELAY_MICRO_SECS, stream) @@ -1608,17 +2012,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() @@ -1628,7 +2022,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" @@ -1640,12 +2041,18 @@ 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 + 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 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..3dfbe5792a3d 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, nvmmh_cta_order_to_cute_raster, + nvmmh_library_version, rank_configs) try: from cuda.bindings import driver as cuda @@ -108,6 +112,32 @@ def _kimi_k3_mxfp8_tuning_bucket(num_tokens: int) -> int: return next_positive_power_of_2(num_tokens) +# 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. + + 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", + _CUTEDSL_NVMMH_TACTIC_CACHE_SCHEMA, + nvmmh_library_version(), + 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 +740,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 +776,10 @@ def unique_id(self): self.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 __hash__(self): return hash(self.unique_id()) @@ -761,6 +795,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. " @@ -884,9 +919,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 +1041,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 +1099,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 +1490,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 +2040,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 +2931,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 +3230,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 +3757,7 @@ def _( device=input.device) class Sm100BlockScaledContiguousGroupedGemmSwigluFusionRunner( - TunableRunner): + CuteDSLTunableRunner): kernel_class = Sm100BlockScaledContiguousGroupedGemmSwigluFusionKernel kernel_cache = dict() tuning_config_cache = dict() @@ -4047,7 +4081,7 @@ def _( return output, output_scale class Sm100BlockScaledContiguousGatherGroupedGemmActFusionRunner( - TunableRunner): + CuteDSLTunableRunner): kernel_class = BlockScaledContiguousGatherGroupedGemmKernel kernel_cache = dict() tuning_config_cache = dict() @@ -4643,7 +4677,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 +4966,7 @@ def _( input.new_empty((m, n // 128), dtype=torch.int32), ) - class CuteDSLFp8BlackwellRunner(TunableRunner): + class CuteDSLFp8BlackwellRunner(CuteDSLTunableRunner): kernel_class = Sm100BlockwiseGemmKernel kernel_cache = dict() @@ -4954,12 +4988,17 @@ def __init__(self, self.output_dtype = output_dtype 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( self, inputs: List[torch.Tensor], 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. " @@ -4991,7 +5030,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 +5051,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 +5340,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 +5665,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 +6051,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 +9643,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 +9733,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 +9792,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 +9819,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 +10063,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 +10074,11 @@ class CuteDSLBf16BlackwellGemmRunner(TunableRunner): """ kernel_class = PersistentDenseGemmKernel kernel_cache = dict() + 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( 0, 0, get_last_power_of_2_num_tokens_buckets, @@ -9985,7 +10094,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. " @@ -10019,7 +10128,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 +10150,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 +10173,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 +10234,7 @@ def forward( use_2cta_instrs, mma_tiler_mn, cluster_shape_mn, + *kernel_options, self.use_tvm_ffi, c_dtype_cutlass, ) @@ -10142,6 +10262,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 +10315,92 @@ def forward( if c_needs_copy: c_tensor.copy_(c_buf) + 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 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 + 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" or not allow_split_k: + # Direct reduction requires Rubin and a matching output shape. + 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) + 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, + k, + precision, + query_count, + layout_name=BF16_LAYOUT, + split_k_kind=split_k_kind, + ) + if not ranked_configs: + return fallback_tactics + + candidate_tactics = 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(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 +10587,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 +10622,11 @@ 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): + """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( @@ -10426,6 +10644,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): @@ -10442,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 @@ -10467,18 +10692,45 @@ 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, 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) + 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 +10793,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 +10814,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 +11186,12 @@ 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 the active NVMMH search policy discriminator for profiling caches.""" + return _cutedsl_nvmmh_tactic_search_cache_key() # See CuteDSLBf16RubinBmmRunner.tuning_config. tuning_config = TuningConfig( @@ -10956,6 +11219,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. " @@ -10981,17 +11245,30 @@ 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, + allow_split_k=inputs[2].shape[0] == m) def forward( self, 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) + 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 +11338,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 +11360,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 +13252,7 @@ def _( (4, 4), ] - class CuteDSLFp8RubinGemmRunner(TunableRunner): + class CuteDSLFp8RubinGemmRunner(CuteDSLTunableRunner): """FP8 1x128/128x128 blockwise GEMM runner for SM107.""" kernel_class = SM107BlockwiseGemmKernel @@ -13231,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, @@ -13255,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, @@ -13650,7 +13974,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 +13994,17 @@ 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 + # 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, ...]] = ("static", @@ -13697,6 +14032,12 @@ def unique_id(self): return (self.output_dtype, self.to_userbuffers, 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() + def __hash__(self): return hash( (self.output_dtype, self.to_userbuffers, self.use_tvm_ffi)) @@ -13713,6 +14054,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. " @@ -13797,7 +14139,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 +14146,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 +14207,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 +14239,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 +14267,451 @@ 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 _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. + + 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 + 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 + model_fields = fields + 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 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}) + # 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: + 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) + if swap_ab in tiles_by_swap: + query_kwargs["tiles"] = tiles_by_swap[swap_ab] + 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 + + 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, + self.nvmmh_sweep_cluster_shape, + ) + + valid_keys = { + model_key(signature) + for signature in candidate_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 + 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: + 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 +14742,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 +14755,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 +14781,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 +14807,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 +14841,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 +14860,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 +15030,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 +15086,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 +15097,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 +15200,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 +15232,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. 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 @@ -14940,13 +15766,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,13 +15788,17 @@ def __init__(self, self.output_dtype = output_dtype 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( self, inputs: List[torch.Tensor], 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. " @@ -15047,7 +15877,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 +16201,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 +17008,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 +17604,7 @@ def _( Sm107BlockScaledContiguousGroupedGemmFinalizeFusionKernel class Sm107BlockScaledContiguousGroupedGemmFinalizeFusionRunner( - TunableRunner): + CuteDSLTunableRunner): """Rubin (SM107) runner for grouped GEMM + finalize fusion (FC2). This is the Rubin counterpart to @@ -17400,7 +18305,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..921f2601a48e 100644 --- a/tensorrt_llm/_torch/custom_ops/cutedsl_matmul_heuristics.py +++ b/tensorrt_llm/_torch/custom_ops/cutedsl_matmul_heuristics.py @@ -14,10 +14,11 @@ ranking so callers fall back to the full tactic list and never regress. """ -import os +import ctypes +import importlib.metadata import re from functools import lru_cache -from typing import List, NamedTuple, Optional, Tuple +from typing import Callable, Iterable, List, NamedTuple, Optional, Tuple from tensorrt_llm.logger import logger @@ -25,10 +26,24 @@ import nvMatmulHeuristics as _nvmmh IS_NVMMH_AVAILABLE = True -except ImportError: +except (ImportError, OSError): _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 @@ -40,84 +55,100 @@ 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" + +# 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). + 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. """ - 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 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 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 +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'") @lru_cache(maxsize=None) @@ -182,11 +213,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 +278,76 @@ 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) + split_k = _extract_split_k(config, kernel) + 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. -def rank_configs(m: int, n: int, k: int, precision: str, count: int) -> List[HeuristicConfig]: + 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, + k: int, + precision: str, + 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. - 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 [] - layout = _get_layout(NVFP4_LAYOUT) + layout = _get_layout(layout_name) if layout is None: return [] try: @@ -244,11 +356,98 @@ 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) + + 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}") + + 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), + ) + 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: + # Older wheels take M, N, and K as separate arguments. + 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: + 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}), 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", ) @@ -261,7 +460,13 @@ 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. + # 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) @@ -275,3 +480,347 @@ 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, + 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: + 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 + + 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): + """Ignore the locally swept cluster-N dimension in the matching key.""" + 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 or widen(signature) + ): + selected.append(tactic) + else: + kept = set(kept_signatures) + 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]: + """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``. 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. + """ + recommended = {int(config.split_k) for config in ranked_configs if int(config.split_k) > 1} + 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) + # 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: + 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"}: + 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: + # 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 _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])) + 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, + 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 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 + 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: + 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) + 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..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) @@ -58,8 +59,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 +1509,22 @@ def forward( else: raise ValueError(f"Invalid tactic: {tactic}") + def use_torch_profiler(self, tactic) -> bool: + """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]: + """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() + @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/_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/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py index 221f9145b17c..503ec255ce2f 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,24 @@ def _configure_deep_gemm_pdl() -> None: _DEEP_GEMM_PDL_CONFIGURED = True +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() + else: + policy = NvMMHConfig( + enabled=True, + fields=config.fields, + max_tactics=config.max_tactics, + ) + AutoTuner.get()._acquire_nvmmh_policy(owner, policy) + + class PyTorchModelEngine(ModelEngine): def __init__( @@ -177,6 +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, self) _configure_deep_gemm_pdl() self._metrics: dict[str, float] = defaultdict(float) @@ -1142,8 +1166,13 @@ def cleanup(self) -> None: # Release model weights. release_gc() + 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..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,9 +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'`. - # 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. + 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/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..869abf9fb51a 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 selected-field 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/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", 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/kernels/parallel/test_fp8_block_scale_gemm.py b/tests/unittest/_torch/kernels/parallel/test_fp8_block_scale_gemm.py index 9567164ca7d3..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 @@ -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) @@ -867,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 @@ -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,75 @@ 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) + 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) - assert static_tactic in valid_tactics - assert dynamic_tactic in valid_tactics - assert split_k_tactic 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 + # 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 + 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) 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..e96685b12f70 100644 --- a/tests/unittest/_torch/misc/test_autotuner.py +++ b/tests/unittest/_torch/misc/test_autotuner.py @@ -55,6 +55,17 @@ pytestmark = pytest.mark.threadleak(enabled=False) +@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: + yield tuner + finally: + tuner.configure_nvmmh(previous) + + def test_multi_dynamic_dims(): tuner = autotuner.AutoTuner() x = torch.rand([5, 1024]) @@ -1513,48 +1524,473 @@ 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) + 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 + 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): + """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": + 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 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) + 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): + """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]) + 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(): + """Execute the cuBLASLt reference used for BF16 timing comparisons.""" + 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 + 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)") + 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): + """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 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 + }.issubset({_nvmmh_tactic_family(t) + 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]) + 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(): + """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( + 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 + 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)") + 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, + 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 +2002,79 @@ 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): + """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() - 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): + """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]) + 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,33 +2085,103 @@ 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} " 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 cuBLAS 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( + "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") + + monkeypatch.setenv("TLLM_PROFILING_TIMER", "torch_profiler") + _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"]) 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