From 0a39cdc877fe1b49fda4f54dae5ac32264a5f52f Mon Sep 17 00:00:00 2001 From: peaceh <103117813+peaceh-nv@users.noreply.github.com> Date: Tue, 27 Oct 2026 10:26:52 +0800 Subject: [PATCH 1/3] [TRTLLM-14024][feat] Prune CuTe DSL GEMM autotuner tactics with nvMatmulHeuristics Signed-off-by: peaceh <103117813+peaceh-nv@users.noreply.github.com> --- docs/source/torch/adding_custom_kernels.md | 31 +- tensorrt_llm/_torch/autotuner.py | 40 +- .../_torch/custom_ops/cute_dsl_custom_ops.py | 1164 +++++++++++++---- .../custom_ops/cutedsl_matmul_heuristics.py | 530 +++++--- .../_torch/custom_ops/torch_custom_ops.py | 10 +- .../blockwise_gemm/blockwise_gemm.py | 192 ++- .../dense_blockscaled_gemm_persistent.py | 97 +- .../rubin/blockwise_gemm/blockwise_gemm.py | 1122 ++++++++++------ .../dense_blockscaled_gemm_persistent.py | 45 +- .../_torch/locality_domain/autotune.py | 3 + .../_torch/pyexecutor/model_engine.py | 15 + tensorrt_llm/llmapi/__init__.py | 4 +- tensorrt_llm/llmapi/llm_args.py | 21 + .../schemas/trtllm-serve-config.schema.json | 28 + .../trtllm-serve-disagg-config.schema.json | 41 + .../usage/llm_args_golden_manifest.json | 5 + ...e_rubin.py => test_cute_dsl_bf16_dense.py} | 61 +- .../parallel/test_fp8_block_scale_gemm.py | 58 + tests/unittest/_torch/misc/test_autotuner.py | 220 +++- .../api_stability/references/llm.yaml | 4 + 20 files changed, 2680 insertions(+), 1011 deletions(-) rename tests/unittest/_torch/kernels/parallel/{test_cute_dsl_bf16_dense_rubin.py => test_cute_dsl_bf16_dense.py} (80%) diff --git a/docs/source/torch/adding_custom_kernels.md b/docs/source/torch/adding_custom_kernels.md index 97aaa3ffac94..c91172e8769b 100644 --- a/docs/source/torch/adding_custom_kernels.md +++ b/docs/source/torch/adding_custom_kernels.md @@ -258,29 +258,24 @@ The same pattern (`@torch.library.custom_op("trtllm::...")`, `register_fake`, co - **DLPack/CUDA Graphs.** When exporting tensors via DLPack, use the stream override that mimics the `CUDAGraphCompatibleWrapper` in `argmax.py` so the capture replays cleanly. - **JIT compile cache.** Always cache compiled kernels by the keys that affect codegen (dtype, last-dim size, hardware variant). `argmax.py` and the `cute_dsl_custom_ops.py` runners are good references. -### 4.5 Pruning autotuner tactics with nvMatmulHeuristics +### 4.5 Guiding autotuner tactics with nvMatmulHeuristics -A CuTe DSL runner that exposes many tactics (tile / cluster / swap / prefetch combinations) can be expensive to autotune, because the AutoTuner JIT-compiles and benchmarks every candidate returned by `get_valid_tactics`. A runner can therefore rank and prune its own candidate list with [nvMatmulHeuristics](https://pypi.org/project/nvidia-matmul-heuristics/) *inside* `get_valid_tactics`, before the tactics ever reach the compile+benchmark loop — all shapes and inputs are already visible at that point, so no extra AutoTuner hook is needed. The pruning is a no-op by default and must degrade gracefully to the full list on any failure, so tuning never loses a valid candidate. +A CuTe DSL GEMM/BMM runner can expose hundreds of tactics (tile / cluster / split-K / rasterization / swap / prefetch combinations), and the AutoTuner JIT-compiles and benchmarks every candidate returned by `get_valid_tactics`. With `autotuner_nvmmh_config` set, nvMMH-aware runners instead build their tactic list from the top kernel configurations that [nvMatmulHeuristics](https://pypi.org/project/nvidia-matmul-heuristics/) (nvMMH) ranks for the problem: -Today this is wired up for the SM100/SM103 (Blackwell B200/B300) NVFP4 dense GEMM: `CuteDSLNVFP4BlackwellRunner.get_valid_tactics` builds the full sweep and then hands it to the private `_rank_prune_tactics` helper. The pruned set is a strict, re-validated subset of the full sweep — it never introduces a tactic the kernel validator rejected — so correctness is unchanged and only autotuner wall-time is affected. On a 4096³ NVFP4 GEMM this cuts a cold autotune from ~120 tactics to ~5 (roughly 7× less wall-time) while keeping the selected kernel within a few percent of the full-sweep best. - -The behavior is controlled entirely by environment variables (all read at tune time; the package must be installed, otherwise the path silently falls back to the full sweep): - -| Environment variable | Default | Meaning | -|---|---|---| -| `TRTLLM_CUTEDSL_NVMMH_ENABLE` | `0` | Master opt-in switch. Set to `1` to enable heuristic tactic pruning. When unset/`0`, or when `nvidia-matmul-heuristics` is not installed, the full tactic sweep runs unchanged. | -| `TRTLLM_CUTEDSL_NVMMH_FIELDS` | `tile,cluster` | Comma-separated list of the knobs nvMatmulHeuristics drives; any knob not listed is swept by the runner. Recognized tokens: `tile`, `cluster`, `swizzle`, `cta_order`. `tile` and `cluster` are a coupled pair (the 2-SM encoding maps the joint `(cta, cluster)`), so naming either drives both. Accepted aliases: `cta_tile` / `mma_tiler` → `tile`. Unknown tokens are warned once and ignored. **Top-K tile/cluster pruning happens only when `tile`/`cluster` is selected; with scheduler-only fields (`swizzle` and/or `cta_order`) every valid tile/cluster is swept and merely annotated with the model's scheduler knobs.** | -| `TRTLLM_CUTEDSL_NVMMH_MAX_TACTICS` | `5` | Top-K distinct `(tile, cluster, swap)` keys to keep from the model's ranking (applies only when `tile`/`cluster` is in `FIELDS`). Both `use_prefetch` variants of a kept key are profiled on top (prefetch is not modeled), so the final tactic count can be up to ~2× this. Larger K = closer to the full-sweep best but slower autotune; smaller K = faster autotune with a larger potential perf gap. | +```yaml +# extra_llm_api_options.yaml +autotuner_nvmmh_config: + max_tactics: 5 # nvMMH configurations to profile per GEMM shape +``` -Example (enable pruning, keep the top 8 tile/cluster candidates): +or `LLM(..., autotuner_nvmmh_config=AutoTunerNvMMHConfig(max_tactics=5))`. The option is off by default; when it is unset, when `nvidia-matmul-heuristics` is not installed, or when a query fails, runners fall back to their full tactic list. -```bash -export TRTLLM_CUTEDSL_NVMMH_ENABLE=1 -export TRTLLM_CUTEDSL_NVMMH_FIELDS=tile,cluster -export TRTLLM_CUTEDSL_NVMMH_MAX_TACTICS=8 -``` +- **nvMMH chooses, the runner realizes.** Every field nvMMH models (CTA tile, cluster shape, split-K, rasterization order, swizzle) is taken as returned. Only knobs nvMMH does not model (prefetch, scheduler mode, the preferred/fallback cluster launch) are swept inside each configuration. If a chosen configuration performs poorly, the fix belongs in nvMMH's model, not in TensorRT-LLM. +- **Kernels describe their capabilities.** Each runner passes nvMMH a validity check (built on the kernel's `can_implement`) and backend properties (split-K kind, column-major rasterization, fallback clusters), so nvMMH ranks implementable configurations only. Runners that support both A/B orientations query both problems and merge the rankings by nvMMH's ranking score. +- **Cache keys.** Runners return `("nvmmh", , max_tactics)` from `TunableRunner.tactic_space_key()`, so profiling-cache entries tuned with and without the guidance are never mixed. +- **Library versions.** Both the `nvidia-matmul-heuristics` 0.1.0.x wheels and newer libheuristics builds are supported; the adapter in [`cutedsl_matmul_heuristics.py`](https://github.com/NVIDIA/TensorRT-LLM/blob/main/tensorrt_llm/_torch/custom_ops/cutedsl_matmul_heuristics.py) hides their API differences. -The adapter and env helpers live in [`tensorrt_llm/_torch/custom_ops/cutedsl_matmul_heuristics.py`](https://github.com/NVIDIA/TensorRT-LLM/blob/main/tensorrt_llm/_torch/custom_ops/cutedsl_matmul_heuristics.py); the `get_valid_tactics` / `_rank_prune_tactics` pruning is in [`cute_dsl_custom_ops.py`](https://github.com/NVIDIA/TensorRT-LLM/blob/main/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py). +To guide a new runner, return `_nvmmh_tactic_space_key()` from `tactic_space_key`, and in `get_valid_tactics` (when `AutoTuner.get().nvmmh_config` is set) call `cutedsl_matmul_heuristics.query(...)` with a validity check that maps a configuration onto the kernel, expand each returned configuration into its tactics, and keep the first `max_tactics` configurations with `select_tactics`. The Rubin block-scaled runner (`CuteDSLBlockScaledRubinLinear`) in [`cute_dsl_custom_ops.py`](https://github.com/NVIDIA/TensorRT-LLM/blob/main/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py) is a complete example. --- diff --git a/tensorrt_llm/_torch/autotuner.py b/tensorrt_llm/_torch/autotuner.py index f55006af4213..cf509856dfe9 100644 --- a/tensorrt_llm/_torch/autotuner.py +++ b/tensorrt_llm/_torch/autotuner.py @@ -157,6 +157,27 @@ class TuningConfig: exclude_from_cache: bool = False +@dataclass(frozen=True) +class NvMMHConfig: + """nvMatmulHeuristics (nvMMH) guidance for CuTe DSL GEMM/BMM runners. + + When installed on the AutoTuner, nvMMH-aware runners build their tactic + list from the top ``max_tactics`` kernel configurations nvMMH ranks for + the problem instead of enumerating their full tactic space. Every field + nvMMH models (CTA tile, cluster shape, split-K, rasterization order, + swizzle) is taken as returned; only knobs nvMMH does not model are swept. + """ + max_tactics: int = 5 + + def __post_init__(self): + if (isinstance(self.max_tactics, bool) + or not isinstance(self.max_tactics, int) + or self.max_tactics < 1): + raise ValueError( + f"NvMMHConfig.max_tactics must be a positive integer, got " + f"{self.max_tactics!r}") + + @dataclass(unsafe_hash=True) class StaticDim: val: int @@ -276,6 +297,16 @@ def unique_id(self): """ return tuple(self.__dict__.values()) + def tactic_space_key(self) -> Optional[Tuple]: + """Discriminates cache entries profiled over a non-default tactic space. + + Runners whose candidate tactics depend on a process-wide search policy + (e.g. nvMatmulHeuristics guidance) return a literal tuple describing + that policy, so entries tuned under one policy are not reused under + another. None keeps the default cache key. + """ + return None + @contextlib.contextmanager def autotune(tune_mode: bool = True, @@ -541,7 +572,7 @@ def get_cache_key( tuning_config: TuningConfig, apply_map_to_tuning_buckets: bool = True, ) -> Tuple: - return ( + cache_key = ( custom_op, runner.__class__.__name__, str(runner.unique_id()), @@ -553,6 +584,10 @@ def get_cache_key( apply_map_to_tuning_buckets, ), ) + tactic_space_key = runner.tactic_space_key() + if tactic_space_key is not None: + cache_key += (tactic_space_key, ) + return cache_key def merge_cache_data(self, cache_data: Dict[Tuple, Tuple]): self.cache.update(cache_data) @@ -946,6 +981,9 @@ def __init__(self, warmup=2, repeat=10, stream_delay_micro_secs=1000): self._primed_cached_tactics: set[tuple[str, str, str, str]] = set() self.is_tuning_mode = False self.skip_dynamic_tuning_buckets = False + # nvMatmulHeuristics guidance for CuTe DSL GEMM/BMM runners; None + # profiles their full tactic space. + self.nvmmh_config: Optional[NvMMHConfig] = None # Timing backend: globaltimer kernel vs cuda events. # TLLM_PROFILING_TIMER env var overrides auto-detection: 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 54792dab7866..5d95864a034d 100644 --- a/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py +++ b/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py @@ -32,9 +32,7 @@ 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 . import cutedsl_matmul_heuristics as cutedsl_nvmmh try: from cuda.bindings import driver as cuda @@ -694,6 +692,49 @@ def _launchable_raster_order( return raster_order +def _nvmmh_tactic_space_key() -> Optional[Tuple]: + """Profiling-cache discriminator of nvMMH-guided tactic lists.""" + config = AutoTuner.get().nvmmh_config + if config is None or not cutedsl_nvmmh.IS_NVMMH_AVAILABLE: + return None + return ("nvmmh", cutedsl_nvmmh.library_version(), config.max_tactics) + + +def _merge_nvmmh_rankings(rankings: List[List[Tuple]]) -> List[Tuple]: + """Merge per-orientation nvMMH rankings of (config, tactics) pairs. + + Each ranking keeps nvMMH's own order; across rankings (plain and A/B + swapped problems) the next pick is the head with the lowest nvMMH ranking + score. + """ + heads = [0] * len(rankings) + merged = [] + while True: + best = None + for i, ranking in enumerate(rankings): + if heads[i] < len(ranking) and ( + best is None or ranking[heads[i]][0].score + < rankings[best][heads[best]][0].score): + best = i + if best is None: + return merged + merged.append(rankings[best][heads[best]]) + heads[best] += 1 + + +# The Blackwell block-scaled GEMM reduces split-K partials with TMA reduce-add +# into the zeroed output and supports both rasters. +_SM100_BLOCKSCALED_NVMMH_BACKEND = cutedsl_nvmmh.Backend( + split_k_kind=cutedsl_nvmmh.SPLIT_K_IN_PLACE, supports_col_major_raster=True) + +# Rubin block-scaled GEMMs reduce split-K partials with TMA reduce-add into the +# zeroed output (no workspace, no reduction launch), support both rasters and +# launch 2-CTA clusters above two CTAs with a 2x1 fallback cluster. +_SM107_BLOCKSCALED_NVMMH_BACKEND = cutedsl_nvmmh.Backend( + split_k_kind=cutedsl_nvmmh.SPLIT_K_IN_PLACE, + supports_col_major_raster=True, + supports_fallback_cluster=True) + if IS_CUTLASS_DSL_AVAILABLE: import cutlass @@ -732,6 +773,117 @@ def _launchable_raster_order( from ..moe.kernels.blackwell.moe_as_dense_gemm.fc1 import \ Sm100BlockScaledPersistentDenseGemmKernel as DenseGemmSwigluKernel + def _parse_sm100_tactic(tactic): + """Decode an SM100 BF16 / FP8 blockwise tactic into (use_2cta_instrs, + mma_tiler_mn, cluster_shape_mn, raster_along, swizzle_size). + + The 5-field form appends the scheduler's raster ("m" or "n") and + swizzle width to (use_2cta_instrs, mma_tiler_mn, cluster_shape_mn); + the 3-field form rasterizes along M without swizzle. The 6-field form + adds the split-K factor (see _sm100_tactic_split_k). + """ + if not isinstance(tactic, tuple): + return False, (128, 128), (1, 1), "m", 1 + if len(tactic) >= 5: + return tactic[:5] + return (*tactic, "m", 1) + + def _sm100_tactic_split_k(tactic) -> int: + """Split-K factor of an SM100 BF16 / FP8 blockwise tactic (1 unless + 6 fields).""" + if isinstance(tactic, tuple) and len(tactic) == 6: + return int(tactic[5]) + return 1 + + def _sm100_fp8_blockwise_mma_tiler_candidates( + m: int) -> List[Tuple[int, int]]: + """Full-sweep MMA tiles of the SM100 FP8 blockwise GEMM / BMM. + + Tiles narrower than a 128-column scale block of B (N = 32, 64) put more + CTAs on the weight stream of decode-sized M; they are swept only there. + """ + tiles = [(64, 128), (128, 128), (256, 128)] + if m <= 256: + tiles += [(tile_m, tile_n) for tile_n in (32, 64) + for tile_m in (64, 128, 256)] + return tiles + + # SM100 BF16 and FP8 blockwise GEMMs support both rasters. The BF16 GEMM + # and the FP8 blockwise GEMM/BMM split K: every slice reduce-adds its + # partial into the zeroed output with TMA (no workspace, no reduction + # launch). + _SM100_NVMMH_BACKEND = cutedsl_nvmmh.Backend(split_k_kind="NONE", + supports_col_major_raster=True) + _SM100_SPLIT_K_NVMMH_BACKEND = cutedsl_nvmmh.Backend( + split_k_kind=cutedsl_nvmmh.SPLIT_K_IN_PLACE, + supports_col_major_raster=True) + + def _sm100_nvmmh_tactics(kernel_class, + ab_dtype, + c_dtype, + precision: str, + m: int, + n: int, + k: int, + batch_size: int, + mma_tiler_n: Optional[int], + cta_tile_k: int, + max_tactics: int, + allow_split_k: bool = False) -> List[Tuple]: + """SM100 BF16 / FP8 blockwise tactics of the top ``max_tactics`` + nvMMH configs (tile, cluster, raster, swizzle and split-K from nvMMH). + + Blackwell MMAs are at most 128 rows per CTA; an even nvMMH cluster M + means a 2-CTA MMA, whose MMA tile covers both CTAs. ``mma_tiler_n`` + pins the kernel's tile N (None: any N the kernel implements) and + ``cta_tile_k`` is the kernel's fixed K tile. ``allow_split_k`` lets + nvMMH split K (in place) up to one K tile per slice. + """ + if precision == cutedsl_nvmmh.BF16_PRECISION: + if c_dtype == cutlass.Float32: + precision = cutedsl_nvmmh.BF16_FP32_OUTPUT_PRECISION + elif c_dtype == cutlass.Float8E4M3FN: + precision = cutedsl_nvmmh.BF16_FP8_OUTPUT_PRECISION + + def shape(config): + cta_m, cta_n, cta_k = config.cta + if (cta_k != cta_tile_k or cta_m not in (64, 128) + or (mma_tiler_n is not None and cta_n != mma_tiler_n) + or cta_n % 32 != 0 or not 32 <= cta_n <= 256 + or (config.split_k != 1 and + (not allow_split_k + or config.split_k > ceil_div(k, cta_tile_k))) + or any(c not in (1, 2, 4) for c in config.cluster)): + return None + use_2cta_instrs = cutedsl_nvmmh.pairs_ctas(config) + mma_tiler_mn = (cta_m * (2 if use_2cta_instrs else 1), cta_n) + if not kernel_class.can_implement( + ab_dtype, cutlass.Float32, c_dtype, use_2cta_instrs, + mma_tiler_mn, config.cluster, m, n, k, batch_size, "k", "k", + "n"): + return None + return use_2cta_instrs, mma_tiler_mn + + configs = cutedsl_nvmmh.query( + cutedsl_nvmmh.Problem(m, n, k, batch_size, precision), + cutedsl_nvmmh.query_count(max_tactics), + lambda config: shape(config) is not None, + backend=(_SM100_SPLIT_K_NVMMH_BACKEND + if allow_split_k else _SM100_NVMMH_BACKEND), + cache_key=(kernel_class.__name__, mma_tiler_n, str(c_dtype))) + ranked = [] + for config in configs: + use_2cta_instrs, mma_tiler_mn = shape(config) + raster_along = cutedsl_nvmmh.cta_order_to_raster(config.cta_order) + scheduler = (() if raster_along == "m" and config.swizzle == 1 + and config.split_k == 1 else + (raster_along, config.swizzle)) + split_k = (config.split_k, ) if config.split_k > 1 else () + ranked.append( + (config, [(use_2cta_instrs, mma_tiler_mn, config.cluster) + + scheduler + split_k])) + return cutedsl_nvmmh.select_tactics(ranked, max_tactics) + @functools.cache def _get_full_device_max_active_clusters(device_id: int, cluster_size: int, reserved_sms: int) -> int: @@ -855,6 +1007,14 @@ def get_valid_tactics( f"CuteDSL: M={m}, N={n}(aligned={output_aligned}), K={real_k}(aligned=True). " f"Using swap_ab={swap_ab_candidates}") + nvmmh_config = AutoTuner.get().nvmmh_config + if nvmmh_config is not None: + nvmmh_tactics = self._nvmmh_tactics(m, n, real_k, + swap_ab_candidates, + nvmmh_config.max_tactics) + if nvmmh_tactics: + return nvmmh_tactics + # full shamoo mma_tiler_mn_candidates = [ (128, 64), @@ -927,88 +1087,105 @@ 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. - return self._rank_prune_tactics(valid_tactics, m, n, real_k) - - def _swap_ab_candidates(self, m, n): - """Deterministic swap_ab choice (not swept), constrained by C-layout - alignment. - - swap_ab=True maps the kernel M to the GEMM N (and vice versa), which - is preferred when M is small (<=128) so the kernel works on the - larger dimension; large M does not swap. The chosen value must still - satisfy the output (C) 16-byte alignment: swap_ab=False needs - N%8==0, swap_ab=True needs M%8==0. Falls back to the feasible value - if the preferred one violates alignment; returns [] if neither works. - """ - m_aligned = m % 8 == 0 - n_aligned = n % 8 == 0 - prefer_swap = m <= 128 - if prefer_swap and m_aligned: - return [True] - if not prefer_swap and n_aligned: - return [False] - if n_aligned: - return [False] - if m_aligned: - return [True] - return [] + return valid_tactics - @staticmethod - def _heuristic_to_tactic_tile(cta, cluster): - """Translate a nvMatmulHeuristics (cta, cluster) config to this - kernel's (mma_tiler_mn, cluster_shape_mn). - - nvMatmulHeuristics caps the per-CTA tile M at 128 on Blackwell and - encodes the 2-SM (2-CTA) MMA as cluster_m == 2 (in the queried - kernel frame), so its effective M tile is cluster_m * cta_m. This - kernel instead encodes the same 2-SM op as mma_tiler_m == 256 - (use_2cta_instrs). So a 2-CTA config (cluster_m == 2) doubles the M - tile while keeping the cluster; everything else passes through. - - The 2-SM op additionally requires the N tile to be aligned (the - ``mma_n_align_requirement_2cta`` in libheuristics; 16 for NVFP4/bf16 - output, 32 when the N tile exceeds the 256 UTCMMA max). If N is not - aligned it is not a valid 2-SM config, so leave it single-CTA. + def tactic_space_key(self): + return _nvmmh_tactic_space_key() + + # K extent of one MMA tile (FP4 A/B) + mma_tiler_k = 256 + + def _nvmmh_tactic(self, config, m: int, n: int, real_k: int, + swap_ab: bool, use_prefetch: bool): + """Map an nvMMH config onto this kernel, or None if it cannot run it. + + Blackwell block-scaled MMAs are 128 rows per CTA; an even nvMMH + cluster M means a 2-CTA MMA, which this kernel encodes as + mma_tiler_m == 256. The K tile is fixed; split-K slices take + balanced K-tile ranges, so any split up to the K-tile count runs. + """ + cta_m, cta_n, cta_k = config.cta + # Cheap rejections first: can_implement prints every refusal. + if (cta_m != 128 or cta_k != self.mma_tiler_k + or config.split_k > real_k // self.mma_tiler_k + or cta_n not in (64, 128, 192, 256) + or any(c not in (1, 2, 4) for c in config.cluster)): + return None + mma_tiler_mn = (256 if cutedsl_nvmmh.pairs_ctas(config) else 128, + cta_n) + if not self._tactic_is_supported(mma_tiler_mn, config.cluster, + swap_ab, use_prefetch, m, n, + real_k): + return None + raster_along_m = cutedsl_nvmmh.cta_order_to_raster( + config.cta_order) == "m" + if config.split_k > 1: + return (mma_tiler_mn, config.cluster, swap_ab, use_prefetch, + config.swizzle, raster_along_m, config.split_k) + if config.swizzle == 1 and raster_along_m: + return (mma_tiler_mn, config.cluster, swap_ab, use_prefetch) + return (mma_tiler_mn, config.cluster, swap_ab, use_prefetch, + config.swizzle, raster_along_m) + + def _nvmmh_tactics(self, m: int, n: int, real_k: int, + swap_ab_candidates: List[bool], + max_tactics: int) -> List[Tuple]: + """Tactics of the top ``max_tactics`` nvMMH configs. + + Tile, cluster, raster and swizzle come from nvMMH; prefetch is not + modeled and keeps the full sweep's pruning rule. """ - cta_m, cta_n = int(cta[0]), int(cta[1]) - cluster_m, cluster_n = int(cluster[0]), int(cluster[1]) - n_align = 32 if cta_n > 256 else 16 - if cluster_m == 2 and cta_n % n_align == 0: - # 2-SM MMA along the kernel M dimension. - return (2 * cta_m, cta_n), (cluster_m, cluster_n) - return (cta_m, cta_n), (cluster_m, cluster_n) + rankings = [] + for swap_ab in swap_ab_candidates: + kernel_m, kernel_n = (n, m) if swap_ab else (m, n) + problem = cutedsl_nvmmh.Problem( + kernel_m, kernel_n, real_k, 1, cutedsl_nvmmh.FP4_PRECISION, + cutedsl_nvmmh.LAYOUT_COL_MAJOR_C + if swap_ab else cutedsl_nvmmh.LAYOUT_ROW_MAJOR_C) + configs = cutedsl_nvmmh.query( + problem, + cutedsl_nvmmh.query_count(max_tactics), + lambda config: self._nvmmh_tactic( + config, m, n, real_k, swap_ab, False) is not None, + backend=_SM100_BLOCKSCALED_NVMMH_BACKEND, + cache_key=(self.__class__.__name__, swap_ab)) + rankings.append([(config, [ + tactic for use_prefetch in (False, True) + if (tactic := self._nvmmh_tactic(config, m, n, real_k, + swap_ab, use_prefetch) + ) is not None + ]) for config in configs]) + return cutedsl_nvmmh.select_tactics(_merge_nvmmh_rankings(rankings), + max_tactics) @staticmethod def _unpack_tactic(tactic): """Unpack a tactic into the full knob set with back-compat defaults. The base tactic is (mma_tiler_mn, cluster_shape_mn, swap_ab, - use_prefetch). The heuristic path may append two tile-scheduler - knobs (swizzle_size, raster_along_m); older 4-tuples (and the full - sweep) default to the kernel's neutral values (no swizzle, M-major - raster). Non-tuple tactics use the default kernel tactic. + use_prefetch). Two tile-scheduler knobs (swizzle_size, + raster_along_m) and the split-K factor may follow; shorter tuples + default to the kernel's neutral values (no swizzle, M-major + raster, no split-K). Non-tuple tactics use the default tactic. """ if not isinstance(tactic, (tuple, list)): - return (128, 128), (1, 1), False, False, 1, True + return (128, 128), (1, 1), False, False, 1, True, 1 mma_tiler_mn = tactic[0] cluster_shape_mn = tactic[1] swap_ab = tactic[2] use_prefetch = tactic[3] swizzle_size = int(tactic[4]) if len(tactic) > 4 else 1 raster_along_m = bool(tactic[5]) if len(tactic) > 5 else True + split_k = int(tactic[6]) if len(tactic) > 6 else 1 return (mma_tiler_mn, cluster_shape_mn, swap_ab, use_prefetch, - swizzle_size, raster_along_m) + swizzle_size, raster_along_m, split_k) def _tactic_is_supported(self, mma_tiler_mn, cluster_shape_mn, swap_ab, use_prefetch, m, n, real_k) -> bool: """Whether the kernel can run this (tile, cluster, swap, prefetch). - Validity gate used by get_valid_tactics() (the full enumeration). - The nvMatmulHeuristics path does NOT call this -- it trusts the - model's mapped tiles and emits them directly. + Validity gate of the nvMatmulHeuristics path; it applies the same + prefetch pruning rule as the full enumeration. """ sf_vec_size = 16 if swap_ab: @@ -1044,126 +1221,6 @@ def _tactic_is_supported(self, mma_tiler_mn, cluster_shape_mn, swap_ab, return False return True - def _rank_prune_tactics(self, tactics, m, n, real_k): - """Rank/prune the full-sweep tactics with nvMatmulHeuristics (opt-in). - - 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. - - 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: - return tactics - fields = nvmmh_fields() - if not fields: - return tactics - try: - max_tactics = nvmmh_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"}) - use_swizzle = "swizzle" in fields - use_cta_order = "cta_order" in fields - emit_extended = use_swizzle or use_cta_order - - # Deterministic swap orientation (opt-in only), intersected with - # the orientations actually present in the valid tactics. - valid_swaps = {t[2] for t in tactics} - swap_pref = [ - s - for s in self._swap_ab_candidates(m, n) if s in valid_swaps - ] or list(valid_swaps) - - # Build the model's ranked preference over (mma_tiler, cluster, - # swap) keys, plus the per-key scheduler knobs. Query enough - # configs to find matches within the valid candidate list. - query_count = max(max_tactics * 4, 16) - pref_rank = {} - pref_sched = {} - for swap_ab in swap_pref: - kernel_m, kernel_n = (n, m) if swap_ab else (m, n) - for cfg in rank_configs(kernel_m, kernel_n, real_k, - NVFP4_PRECISION, query_count): - # Map per-CTA tile + cluster to this kernel's mma_tiler / - # cluster (cluster_m==2 encodes an mma_tiler_m==256 2-SM - # op). Off-grid maps simply won't match the valid list. - mma_tiler_mn, cluster_shape_mn = \ - self._heuristic_to_tactic_tile(cfg.cta, cfg.cluster) - key = (mma_tiler_mn, cluster_shape_mn, swap_ab) - if key in pref_rank: - 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) - pref_sched[key] = (max(1, int(swizzle_size)), - bool(raster_along_m)) - - # Restrict to the top-K ranked tile/cluster keys ONLY when the - # model drives tile/cluster. For scheduler-only fields keep every - # valid tile/cluster key and just annotate it below. - valid_keys = {(t[0], t[1], t[2]) for t in tactics} - if prune_tile_cluster: - matched_keys = valid_keys & pref_rank.keys() - if not matched_keys: - return tactics - kept_keys = set( - sorted(matched_keys, - key=lambda kk: pref_rank[kk])[:max_tactics]) - else: - kept_keys = valid_keys - - # Emit every supplied (already-valid) tactic whose key is kept, - # re-validating as a safety net and optionally annotating the - # model's swizzle / raster (neutral defaults for unranked keys). - selected = [] - for t in tactics: - key = (t[0], t[1], t[2]) - if key not in kept_keys: - continue - mma_tiler_mn, cluster_shape_mn, swap_ab, use_prefetch = t[: - 4] - if not self._tactic_is_supported( - mma_tiler_mn, cluster_shape_mn, swap_ab, - use_prefetch, m, n, real_k): - continue - if emit_extended: - swizzle_size, raster_along_m = pref_sched.get( - key, (1, True)) - selected.append( - (mma_tiler_mn, cluster_shape_mn, swap_ab, - use_prefetch, swizzle_size, raster_along_m)) - else: - selected.append((mma_tiler_mn, cluster_shape_mn, - swap_ab, use_prefetch)) - - logger.debug( - f"CuteDSL nvMatmulHeuristics: {len(tactics)} valid -> " - f"{len(selected)} tactics ({len(kept_keys)} keys) for " - f"M={m}, N={n}, K={real_k}; fields={sorted(fields)}") - return selected if selected else tactics - except Exception as e: # noqa: BLE001 - must never break tuning - logger.warning_once( - f"[nvMatmulHeuristics] NVFP4 tactic filtering failed: {e}. " - f"Falling back to full tactic list.", - key="nvmmh_nvfp4_filter_failure", - ) - return tactics - def should_profile_tactic_in_subprocess( self, custom_op: str, @@ -1172,9 +1229,10 @@ def should_profile_tactic_in_subprocess( tuning_config: TuningConfig, **kwargs, ) -> bool: - # get_valid_tactics emits 4 fields: - # (mma_tiler_mn, cluster_shape_mn, swap_ab, use_prefetch). - return isinstance(tactic, tuple) and len(tactic) == 4 + # get_valid_tactics emits (mma_tiler_mn, cluster_shape_mn, swap_ab, + # use_prefetch), plus (swizzle_size, raster_along_m) for a + # non-default scheduler and the split-K factor. + return isinstance(tactic, tuple) and len(tactic) in (4, 6, 7) def make_cute_dsl_global_pointer(self, tensor: torch.Tensor, dtype, assumed_align: int): @@ -1212,7 +1270,8 @@ def forward( sf_vec_size = 16 (mma_tiler_mn, cluster_shape_mn, swap_ab, use_prefetch, - swizzle_size, raster_along_m) = self._unpack_tactic(tactic) + swizzle_size, raster_along_m, + split_k) = self._unpack_tactic(tactic) a_tensor, b_tensor, a_sf_tensor, b_sf_tensor, alpha_tensor = inputs m, k, n = a_tensor.shape[0], a_tensor.shape[1], b_tensor.shape[0] @@ -1275,7 +1334,7 @@ def forward( stream = cuda.CUstream(torch_stream.cuda_stream) cache_key = (sf_vec_size, mma_tiler_mn, cluster_shape_mn, swap_ab, - use_prefetch, swizzle_size, raster_along_m, + use_prefetch, swizzle_size, raster_along_m, split_k, self.use_tvm_ffi) if swap_ab: kernel_m = n @@ -1345,6 +1404,7 @@ def forward( use_prefetch, swizzle_size, raster_along_m, + split_k=split_k, ) # Compute max active clusters on current device hardware_info = cutlass.utils.HardwareInfo() @@ -1379,6 +1439,11 @@ def forward( else: compiled_gemm = self.__class__.kernel_cache[cache_key] + # Split-K slices reduce-add into C through TMA; zero it on the + # launch stream right before the GEMM so CUDA graph replay is safe. + if split_k > 1: + c_tensor.zero_() + # launch gemm kernel if self.use_tvm_ffi: # call with torch pointer types and no need to pass stream. @@ -4979,11 +5044,14 @@ class CuteDSLFp8BlackwellRunner(TunableRunner): kernel_class = Sm100BlockwiseGemmKernel kernel_cache = dict() + # Profile with a cold L2 like the other CuTe DSL GEMMs: a decode-sized + # weight fits in L2 and would otherwise rank tactics by L2 bandwidth. tuning_config = TuningConfig( dynamic_tensor_specs=(DynamicTensorSpec( 0, 0, get_last_power_of_2_num_tokens_buckets, last_positive_power_of_2), ), constraint_specs=(ConstraintSpec(2, 1, fp8_scale_infer_shape), ), + use_cold_l2_cache=True, ) def __init__(self, @@ -4997,6 +5065,9 @@ def __init__(self, self.output_dtype = output_dtype self.use_tvm_ffi = use_tvm_ffi + def tactic_space_key(self): + return _nvmmh_tactic_space_key() + def get_valid_tactics( self, inputs: List[torch.Tensor], @@ -5021,8 +5092,28 @@ def get_valid_tactics( # m, n c_major = "n" + nvmmh_config = AutoTuner.get().nvmmh_config + if nvmmh_config is not None: + # forward allocates the whole [M, N] output, so nvMMH may + # split K. + nvmmh_tactics = _sm100_nvmmh_tactics( + self.__class__.kernel_class, + cutlass.Float8E4M3FN, + cutlass.BFloat16, + cutedsl_nvmmh.FP8_PRECISION, + m, + n, + k, + batch_size, + None, + 128, + nvmmh_config.max_tactics, + allow_split_k=True) + if nvmmh_tactics: + return nvmmh_tactics + use_2cta_instrs_candi = [False, True] - mma_tiler_mn_candi = [(64, 128), (128, 128), (256, 128)] + mma_tiler_mn_candi = _sm100_fp8_blockwise_mma_tiler_candidates(m) cluster_shape_mn_candi = [ (1, 1), (1, 2), @@ -5075,15 +5166,9 @@ def forward( Returns: torch.Tensor: Output tensor of shape (m, n), dtype: bf16. """ - if isinstance(tactic, tuple): - use_2cta_instrs, mma_tiler_mn, cluster_shape_mn = tactic - else: - # fallback to default tactic - use_2cta_instrs, mma_tiler_mn, cluster_shape_mn = [ - False, - (128, 128), - (1, 1), - ] + (use_2cta_instrs, mma_tiler_mn, cluster_shape_mn, raster_along, + swizzle_size) = _parse_sm100_tactic(tactic) + split_k = _sm100_tactic_split_k(tactic) a_tensor, b_tensor, a_sf_tensor, b_sf_tensor = inputs m, n, k = a_tensor.shape[0], b_tensor.shape[0], b_tensor.shape[1] sf_m = m @@ -5130,6 +5215,9 @@ def forward( use_2cta_instrs, mma_tiler_mn, cluster_shape_mn, + raster_along, + swizzle_size, + split_k, self.use_tvm_ffi, ) if cache_key not in self.__class__.kernel_cache: @@ -5169,6 +5257,9 @@ def forward( use_2cta_instrs=use_2cta_instrs, mma_tiler_mn=mma_tiler_mn, cluster_shape_mn=cluster_shape_mn, + swizzle_size=swizzle_size, + raster_along_m=raster_along == "m", + split_k=split_k, ) # Compute max active clusters on current device hardware_info = cutlass.utils.HardwareInfo() @@ -5198,6 +5289,11 @@ def forward( else: compiled_gemm = self.__class__.kernel_cache[cache_key] + # Split-K slices reduce-add into C through TMA; zero it on the + # launch stream right before the GEMM. + if split_k > 1: + c_tensor.zero_() + # launch gemm kernel if self.use_tvm_ffi: # call with torch pointer types and no need to pass stream. @@ -5301,6 +5397,7 @@ class CuteDSLFp8BlackwellBmmRunner(TunableRunner): ConstraintSpec( 4, 1, lambda input_shapes: input_shapes[0][1])), + use_cold_l2_cache=True, ) def __init__(self, @@ -5314,6 +5411,9 @@ def __init__(self, self.output_dtype = output_dtype self.use_tvm_ffi = use_tvm_ffi + def tactic_space_key(self): + return _nvmmh_tactic_space_key() + def get_valid_tactics( self, inputs: List[torch.Tensor], @@ -5339,8 +5439,27 @@ def get_valid_tactics( # m, n c_major = "n" + nvmmh_config = AutoTuner.get().nvmmh_config + if nvmmh_config is not None: + # Split-K reduce-adds into the whole [B, M, N] output. + nvmmh_tactics = _sm100_nvmmh_tactics( + self.__class__.kernel_class, + cutlass.Float8E4M3FN, + cutlass.BFloat16, + cutedsl_nvmmh.FP8_PRECISION, + m, + n, + k, + batch_size, + None, + 128, + nvmmh_config.max_tactics, + allow_split_k=inputs[4].shape[1] == m) + if nvmmh_tactics: + return nvmmh_tactics + use_2cta_instrs_candi = [False, True] - mma_tiler_mn_candi = [(64, 128), (128, 128), (256, 128)] + mma_tiler_mn_candi = _sm100_fp8_blockwise_mma_tiler_candidates(m) cluster_shape_mn_candi = [ (1, 1), (1, 2), @@ -5393,15 +5512,9 @@ def forward( Returns: torch.Tensor: Output tensor of shape (batch_size, m, n), dtype: bf16. """ - if isinstance(tactic, tuple): - use_2cta_instrs, mma_tiler_mn, cluster_shape_mn = tactic - else: - # fallback to default tactic - use_2cta_instrs, mma_tiler_mn, cluster_shape_mn = [ - False, - (128, 128), - (1, 1), - ] + (use_2cta_instrs, mma_tiler_mn, cluster_shape_mn, raster_along, + swizzle_size) = _parse_sm100_tactic(tactic) + split_k = _sm100_tactic_split_k(tactic) a_tensor, b_tensor, a_sf_tensor, b_sf_tensor, c_tensor = inputs c_tmp = c_tensor.permute(1, 2, 0) @@ -5449,6 +5562,9 @@ def forward( use_2cta_instrs, mma_tiler_mn, cluster_shape_mn, + raster_along, + swizzle_size, + split_k, self.use_tvm_ffi, ) if cache_key not in self.__class__.kernel_cache: @@ -5489,6 +5605,9 @@ def forward( use_2cta_instrs=use_2cta_instrs, mma_tiler_mn=mma_tiler_mn, cluster_shape_mn=cluster_shape_mn, + swizzle_size=swizzle_size, + raster_along_m=raster_along == "m", + split_k=split_k, ) # Compute max active clusters on current device hardware_info = cutlass.utils.HardwareInfo() @@ -5518,6 +5637,11 @@ def forward( else: compiled_gemm = self.__class__.kernel_cache[cache_key] + # Split-K slices reduce-add into C through TMA; zero it on the + # launch stream right before the GEMM. + if split_k > 1: + c_tensor.zero_() + # launch gemm kernel if self.use_tvm_ffi: # call with torch pointer types and no need to pass stream. @@ -9664,14 +9788,25 @@ class CuteDSLBf16BlackwellBmmRunner(TunableRunner): # Output element type; subclasses override (see the FP8-out runner). c_dtype = cutlass.BFloat16 - tuning_config = TuningConfig(dynamic_tensor_specs=(DynamicTensorSpec( - 0, 1, get_last_power_of_2_num_tokens_buckets, - last_positive_power_of_2), ), ) + # The output's M (inputs[2] dim 1) tracks input0's bucketed M, so the + # profiling-cache key does not embed the raw tuning-time M (which every + # other runtime M would miss). + tuning_config = TuningConfig( + dynamic_tensor_specs=(DynamicTensorSpec( + 0, 1, get_last_power_of_2_num_tokens_buckets, + last_positive_power_of_2), ), + constraint_specs=(ConstraintSpec( + 2, 1, lambda input_shapes: input_shapes[0][1]), ), + use_cold_l2_cache=True, + ) def __init__(self, use_tvm_ffi: bool = True): super().__init__() self.use_tvm_ffi = use_tvm_ffi + def tactic_space_key(self): + return _nvmmh_tactic_space_key() + def get_valid_tactics( self, inputs: List[torch.Tensor], @@ -9697,6 +9832,15 @@ def get_valid_tactics( # m, n c_major = "n" + nvmmh_config = AutoTuner.get().nvmmh_config + if nvmmh_config is not None: + nvmmh_tactics = _sm100_nvmmh_tactics( + self.__class__.kernel_class, cutlass.BFloat16, + self.__class__.c_dtype, cutedsl_nvmmh.BF16_PRECISION, m, n, + k, batch_size, None, 64, nvmmh_config.max_tactics) + if nvmmh_tactics: + return nvmmh_tactics + use_2cta_instrs_candi = [False, True] mma_tiler_mn_candi = [(64, 128), (128, 128), (256, 128)] cluster_shape_mn_candi = [ @@ -9748,14 +9892,8 @@ def forward( tactic: Tiling and cluster strategy, typically a tuple (use_2cta_instrs, mma_tiler_mn, cluster_shape_mn). """ - if isinstance(tactic, tuple): - use_2cta_instrs, mma_tiler_mn, cluster_shape_mn = tactic - else: - use_2cta_instrs, mma_tiler_mn, cluster_shape_mn = [ - False, - (128, 128), - (1, 1), - ] + (use_2cta_instrs, mma_tiler_mn, cluster_shape_mn, raster_along, + swizzle_size) = _parse_sm100_tactic(tactic) a_tensor, b_tensor, c_tensor = inputs @@ -9804,6 +9942,8 @@ def forward( use_2cta_instrs, mma_tiler_mn, cluster_shape_mn, + raster_along, + swizzle_size, self.use_tvm_ffi, ) if cache_key not in self.__class__.kernel_cache: @@ -9830,6 +9970,8 @@ def forward( use_2cta_instrs=use_2cta_instrs, mma_tiler_mn=mma_tiler_mn, cluster_shape_mn=cluster_shape_mn, + swizzle_size=swizzle_size, + raster_along=raster_along, ) hardware_info = cutlass.utils.HardwareInfo() max_active_clusters = hardware_info.get_max_active_clusters( @@ -9957,14 +10099,8 @@ def forward( inputs: List[torch.Tensor], tactic, ) -> None: - if isinstance(tactic, tuple): - use_2cta_instrs, mma_tiler_mn, cluster_shape_mn = tactic - else: - use_2cta_instrs, mma_tiler_mn, cluster_shape_mn = [ - False, - (128, 128), - (1, 1), - ] + (use_2cta_instrs, mma_tiler_mn, cluster_shape_mn, raster_along, + swizzle_size) = _parse_sm100_tactic(tactic) a_tensor, b_tensor, c_tensor = inputs batch_size, m, k = a_tensor.shape[0], a_tensor.shape[ 1], a_tensor.shape[2] @@ -9988,13 +10124,16 @@ def forward( assumed_align=16) stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream) - cache_key = (use_2cta_instrs, mma_tiler_mn, cluster_shape_mn) + cache_key = (use_2cta_instrs, mma_tiler_mn, cluster_shape_mn, + raster_along, swizzle_size) if cache_key not in self.__class__.kernel_cache: gemm = self.__class__.kernel_class( cutlass.Float32, # acc_dtype use_2cta_instrs=use_2cta_instrs, mma_tiler_mn=mma_tiler_mn, cluster_shape_mn=cluster_shape_mn, + swizzle_size=swizzle_size, + raster_along=raster_along, ) hardware_info = cutlass.utils.HardwareInfo() max_active_clusters = hardware_info.get_max_active_clusters( @@ -10083,13 +10222,31 @@ class CuteDSLBf16BlackwellGemmRunner(TunableRunner): kernel_class = PersistentDenseGemmKernel kernel_cache = dict() - tuning_config = TuningConfig(dynamic_tensor_specs=(DynamicTensorSpec( - 0, 0, get_last_power_of_2_num_tokens_buckets, - last_positive_power_of_2), ), ) + # See CuteDSLBf16BlackwellBmmRunner.tuning_config. + tuning_config = TuningConfig( + dynamic_tensor_specs=(DynamicTensorSpec( + 0, 0, get_last_power_of_2_num_tokens_buckets, + last_positive_power_of_2), ), + constraint_specs=(ConstraintSpec( + 2, 0, lambda input_shapes: input_shapes[0][0]), ), + use_cold_l2_cache=True, + ) - def __init__(self, use_tvm_ffi: bool = True): + def __init__(self, + use_tvm_ffi: bool = True, + output_dtype: Optional[torch.dtype] = None): super().__init__() self.use_tvm_ffi = use_tvm_ffi + # The tactic space depends on the output dtype (the full sweep + # splits K only for FP32 output, and nvMMH models the output + # type), so unique_id() keys the profiling cache on it. + self.output_dtype = output_dtype + + def unique_id(self): + return (self.use_tvm_ffi, self.output_dtype) + + def tactic_space_key(self): + return _nvmmh_tactic_space_key() def get_valid_tactics( self, @@ -10118,6 +10275,20 @@ def get_valid_tactics( b_major = "k" c_major = "n" + # Split-K reduce-adds into the whole [M, N] output, so the output + # must hold exactly the M rows of A. + allow_split_k = inputs[2].shape[0] == m + + nvmmh_config = AutoTuner.get().nvmmh_config + if nvmmh_config is not None: + nvmmh_tactics = _sm100_nvmmh_tactics( + self.__class__.kernel_class, cutlass.BFloat16, + c_dtype_cutlass, cutedsl_nvmmh.BF16_PRECISION, m, n, k, + batch_size, None, 64, nvmmh_config.max_tactics, + allow_split_k) + if nvmmh_tactics: + return nvmmh_tactics + use_2cta_instrs_candi = [False, True] mma_tiler_mn_candi = [(64, 128), (128, 128), (256, 128)] cluster_shape_mn_candi = [ @@ -10131,8 +10302,16 @@ def get_valid_tactics( (4, 2), (4, 4), ] + # Split-K pays off for long K with few output tiles. The full + # sweep offers it only for FP32 output, where the TMA reduce-add + # of the slices does not round partial sums to BF16. + split_k_candi = [1] + if (allow_split_k and c_dtype_cutlass == cutlass.Float32 + and k >= 4096 and n <= 512): + split_k_candi = [1, 2, 4, 8] return [ - (use_2cta_instrs, mma_tiler_mn, cluster_shape_mn) + (use_2cta_instrs, mma_tiler_mn, cluster_shape_mn) + + (("m", 1, split_k) if split_k > 1 else ()) for use_2cta_instrs in use_2cta_instrs_candi for mma_tiler_mn in mma_tiler_mn_candi for cluster_shape_mn in cluster_shape_mn_candi @@ -10150,7 +10329,7 @@ def get_valid_tactics( a_major, b_major, c_major, - ) + ) for split_k in split_k_candi ] def forward( @@ -10169,14 +10348,9 @@ def forward( tactic: Tiling and cluster strategy, typically a tuple (use_2cta_instrs, mma_tiler_mn, cluster_shape_mn). """ - if isinstance(tactic, tuple): - use_2cta_instrs, mma_tiler_mn, cluster_shape_mn = tactic - else: - use_2cta_instrs, mma_tiler_mn, cluster_shape_mn = [ - False, - (128, 128), - (1, 1), - ] + (use_2cta_instrs, mma_tiler_mn, cluster_shape_mn, raster_along, + swizzle_size) = _parse_sm100_tactic(tactic) + split_k_slices = _sm100_tactic_split_k(tactic) a_tensor, b_tensor, c_tensor = inputs @@ -10184,6 +10358,10 @@ def forward( m, k = a_tensor.shape[0], a_tensor.shape[1] n = b_tensor.shape[0] batch_size = 1 + if split_k_slices > 1 and tuple(c_tensor.shape) != (m, n): + raise RuntimeError( + "BF16 split-K GEMM requires an [M, N] output, got " + f"output.shape={tuple(c_tensor.shape)}, expected={(m, n)}.") # Ensure inputs are contiguous a_tensor = a_tensor.contiguous() @@ -10227,6 +10405,9 @@ def forward( use_2cta_instrs, mma_tiler_mn, cluster_shape_mn, + raster_along, + swizzle_size, + split_k_slices, self.use_tvm_ffi, c_dtype_cutlass, ) @@ -10254,6 +10435,10 @@ def forward( use_2cta_instrs=use_2cta_instrs, mma_tiler_mn=mma_tiler_mn, cluster_shape_mn=cluster_shape_mn, + swizzle_size=swizzle_size, + raster_along=raster_along, + split_k_slices=split_k_slices, + use_direct_split_k_reduce=split_k_slices > 1, ) hardware_info = cutlass.utils.HardwareInfo() max_active_clusters = hardware_info.get_max_active_clusters( @@ -10277,6 +10462,11 @@ def forward( else: compiled_gemm = self.__class__.kernel_cache[cache_key] + # TMA ADD accumulates the split-K slices into C; zero it on the + # launch stream right before the GEMM so CUDA graph replay is safe. + if split_k_slices > 1: + c_buf.zero_() + # launch gemm kernel if self.use_tvm_ffi: compiled_gemm( @@ -10329,7 +10519,8 @@ def cute_dsl_bf16_gemm_blackwell( tuner = AutoTuner.get() - runner = CuteDSLBf16BlackwellGemmRunner(use_tvm_ffi=use_tvm_ffi) + runner = CuteDSLBf16BlackwellGemmRunner(use_tvm_ffi=use_tvm_ffi, + output_dtype=output.dtype) inputs = [input, weight, output] @@ -10489,22 +10680,34 @@ def _sm107_bf16_valid_tactics( 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).""" + split_k_slices, raster_along, swizzle_size). + + The 8-field forms append (raster_along, swizzle_size) to the base + ("base", 2cta, tile, cluster, max_num_ab_stage, split_k) and preferred + ("preferred_cluster", 2cta, tile, preferred, fallback, + max_num_ab_stage) tactics; shorter forms use raster "m", swizzle 1. + """ if (isinstance(tactic, tuple) and len(tactic) > 0 and isinstance(tactic[0], str)): kernel_variant = tactic[0] if kernel_variant == "preferred_cluster": - _, use_2cta_instrs, mma_tiler_mn, preferred_cluster_shape_mn, cluster_shape_mn, max_num_ab_stage = tactic + (_, use_2cta_instrs, mma_tiler_mn, preferred_cluster_shape_mn, + cluster_shape_mn, max_num_ab_stage, *scheduler) = tactic + raster_along, swizzle_size = scheduler or ("m", 1) return (kernel_variant, use_2cta_instrs, mma_tiler_mn, preferred_cluster_shape_mn, cluster_shape_mn, - max_num_ab_stage, 1) - if len(tactic) == 6: + max_num_ab_stage, 1, raster_along, swizzle_size) + raster_along, swizzle_size = "m", 1 + if len(tactic) == 8: + _, use_2cta_instrs, mma_tiler_mn, cluster_shape_mn, max_num_ab_stage, split_k_slices, raster_along, swizzle_size = tactic + elif len(tactic) == 6: _, use_2cta_instrs, mma_tiler_mn, cluster_shape_mn, max_num_ab_stage, split_k_slices = tactic else: _, use_2cta_instrs, mma_tiler_mn, cluster_shape_mn, max_num_ab_stage = tactic split_k_slices = 1 return (kernel_variant, use_2cta_instrs, mma_tiler_mn, None, - cluster_shape_mn, max_num_ab_stage, split_k_slices) + cluster_shape_mn, max_num_ab_stage, split_k_slices, + raster_along, swizzle_size) if isinstance(tactic, tuple): use_2cta_instrs, mma_tiler_mn, cluster_shape_mn = tactic else: @@ -10512,7 +10715,94 @@ def _parse_sm107_bf16_tactic(tactic): mma_tiler_mn = (128, 128) cluster_shape_mn = (1, 1) return ("base", use_2cta_instrs, mma_tiler_mn, None, cluster_shape_mn, - 0, 1) + 0, 1, "m", 1) + + # SM107 BF16 GEMMs reduce split-K partials with TMA reduce-add into the + # zeroed output (no workspace, no reduction launch) and launch 2-CTA + # clusters above two CTAs with a 2x1 fallback cluster. + _SM107_BF16_NVMMH_BACKEND = cutedsl_nvmmh.Backend( + split_k_kind=cutedsl_nvmmh.SPLIT_K_IN_PLACE, + supports_col_major_raster=True, + supports_fallback_cluster=True) + + def _sm107_bf16_nvmmh_shape(config, m: int, n: int, k: int, batch_size: int, + c_dtype_cutlass, allow_split_k: bool): + """Map an nvMMH config onto the SM107 BF16 kernel, or None. + + Returns (use_2cta_instrs, mma_tiler_mn). nvMMH's per-CTA tile is + 64 or 128 rows; an even cluster M means 2-CTA MMA, whose MMA tile + covers both CTAs. The kernel's K tile is 64. + """ + cta_m, cta_n, cta_k = config.cta + cluster_m, cluster_n = config.cluster + if (cta_k != 64 or cta_m not in (64, 128) or cta_n % 32 != 0 + or not 32 <= cta_n <= 256 or cluster_m not in (1, 2, 4) + or cluster_n not in (1, 2, 4) + # CTA_N=256 with cluster_n=2 is an illegal memory access on SM107. + or (cta_n == 256 and cluster_n == 2) or + (config.split_k > 1 and + (not allow_split_k or config.split_k > ceil_div(k, 64)))): + return None + use_2cta_instrs = cutedsl_nvmmh.pairs_ctas(config) + mma_tiler_mn = (cta_m * (2 if use_2cta_instrs else 1), cta_n) + if not _bf16_cluster_m_fits(m, use_2cta_instrs, mma_tiler_mn, + config.cluster): + return None + if not _sm107_bf16_kernel_class("base").can_implement( + cutlass.BFloat16, cutlass.Float32, c_dtype_cutlass, + use_2cta_instrs, mma_tiler_mn, config.cluster, m, n, k, + batch_size, "k", "k", "n"): + return None + return use_2cta_instrs, mma_tiler_mn + + def _sm107_bf16_nvmmh_tactics(m: int, n: int, k: int, batch_size: int, + c_dtype_cutlass, allow_split_k: bool, + max_tactics: int) -> List[Tuple]: + """SM107 BF16 tactics of the top ``max_tactics`` nvMMH configs. + + Tile, cluster, split-K, raster, swizzle and the A/B stage count come + from nvMMH. A preferred cluster above two 2-CTA CTAs is also launched + with a 2x1 fallback cluster, which nvMMH does not model. + """ + precision = (cutedsl_nvmmh.BF16_FP32_OUTPUT_PRECISION if c_dtype_cutlass + == cutlass.Float32 else cutedsl_nvmmh.BF16_PRECISION) + configs = cutedsl_nvmmh.query( + cutedsl_nvmmh.Problem(m, n, k, batch_size, precision), + cutedsl_nvmmh.query_count(max_tactics), + lambda config: _sm107_bf16_nvmmh_shape( + config, m, n, k, batch_size, c_dtype_cutlass, allow_split_k) is + not None, + backend=_SM107_BF16_NVMMH_BACKEND, + cache_key=("sm107_bf16", allow_split_k)) + fallback_cluster_shape_mn = _SM107_BF16_FALLBACK_CLUSTER_SHAPE_MN + ranked = [] + for config in configs: + use_2cta_instrs, mma_tiler_mn = _sm107_bf16_nvmmh_shape( + config, m, n, k, batch_size, c_dtype_cutlass, allow_split_k) + max_num_ab_stage = config.stages + raster_along = cutedsl_nvmmh.cta_order_to_raster(config.cta_order) + scheduler = (() if raster_along == "m" and config.swizzle == 1 else + (raster_along, config.swizzle)) + split_k = ((config.split_k, ) if allow_split_k or scheduler else ()) + tactics = [("base", use_2cta_instrs, mma_tiler_mn, config.cluster, + max_num_ab_stage) + split_k + scheduler] + preferred_class = _sm107_bf16_kernel_class("preferred_cluster") + if (config.split_k == 1 and use_2cta_instrs + and math.prod(config.cluster) > 2 + and _bf16_preferred_cluster_has_launchable_grid( + m, n, batch_size, use_2cta_instrs, mma_tiler_mn, + config.cluster, fallback_cluster_shape_mn) + and preferred_class.can_implement( + cutlass.BFloat16, cutlass.Float32, c_dtype_cutlass, + use_2cta_instrs, mma_tiler_mn, + fallback_cluster_shape_mn, m, n, k, batch_size, "k", + "k", "n")): + tactics.append(("preferred_cluster", use_2cta_instrs, + mma_tiler_mn, config.cluster, + fallback_cluster_shape_mn, max_num_ab_stage) + + scheduler) + ranked.append((config, tactics)) + return cutedsl_nvmmh.select_tactics(ranked, max_tactics) def _sm107_bf16_gemm_kernel(kernel_variant: str, use_2cta_instrs: bool, @@ -10520,7 +10810,9 @@ 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): kernel_class = _sm107_bf16_kernel_class(kernel_variant) if kernel_variant == "preferred_cluster": return kernel_class( @@ -10529,6 +10821,8 @@ def _sm107_bf16_gemm_kernel(kernel_variant: str, mma_tiler_mn=mma_tiler_mn, preferred_cluster_shape_mn=preferred_cluster_shape_mn, fallback_cluster_shape_mn=cluster_shape_mn, + swizzle_size=swizzle_size, + raster_along=raster_along, max_num_ab_stage=max_num_ab_stage, ) return kernel_class( @@ -10536,6 +10830,8 @@ def _sm107_bf16_gemm_kernel(kernel_variant: str, use_2cta_instrs=use_2cta_instrs, mma_tiler_mn=mma_tiler_mn, cluster_shape_mn=cluster_shape_mn, + swizzle_size=swizzle_size, + raster_along=raster_along, max_num_ab_stage=max_num_ab_stage, split_k_slices=split_k_slices, ) @@ -10579,18 +10875,28 @@ def get_valid_tactics( return [] batch_size, m, k = inputs[0].shape n = inputs[1].shape[1] + nvmmh_config = AutoTuner.get().nvmmh_config + if nvmmh_config is not None: + nvmmh_tactics = _sm107_bf16_nvmmh_tactics( + m, n, k, batch_size, cutlass.BFloat16, False, + nvmmh_config.max_tactics) + if nvmmh_tactics: + return nvmmh_tactics return _sm107_bf16_valid_tactics( m, n, k, batch_size, cutlass.BFloat16, _SM107_BF16_MMA_TILER_MN_CANDIDATES, None) + def tactic_space_key(self): + return _nvmmh_tactic_space_key() + def forward( self, inputs: List[torch.Tensor], tactic, ) -> None: (kernel_variant, use_2cta_instrs, mma_tiler_mn, - preferred_cluster_shape_mn, cluster_shape_mn, max_num_ab_stage, - _) = _parse_sm107_bf16_tactic(tactic) + preferred_cluster_shape_mn, cluster_shape_mn, max_num_ab_stage, _, + raster_along, swizzle_size) = _parse_sm107_bf16_tactic(tactic) a_tensor, b_tensor, c_tensor = inputs batch_size, m, k = a_tensor.shape @@ -10653,6 +10959,8 @@ def forward( preferred_cluster_shape_mn, cluster_shape_mn, max_num_ab_stage, + raster_along, + swizzle_size, self.use_tvm_ffi, c_layout_key, max_active_clusters, @@ -10673,11 +10981,14 @@ 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=raster_along, + swizzle_size=swizzle_size) 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 @@ -11062,6 +11373,9 @@ def __init__( def unique_id(self): return (self.use_tvm_ffi, self.output_dtype) + def tactic_space_key(self): + return _nvmmh_tactic_space_key() + def get_valid_tactics( self, inputs: List[torch.Tensor], @@ -11078,6 +11392,16 @@ def get_valid_tactics( n = inputs[1].shape[0] c_dtype_cutlass = _TORCH_TO_CUTLASS_DTYPE[inputs[2].dtype] + nvmmh_config = AutoTuner.get().nvmmh_config + if nvmmh_config is not None: + # Split-K writes the full [M, N] output; dynamic tuning may + # bucket input M while leaving C at the full M. + nvmmh_tactics = _sm107_bf16_nvmmh_tactics( + m, n, k, 1, c_dtype_cutlass, inputs[2].shape[0] == m, + nvmmh_config.max_tactics) + if nvmmh_tactics: + return nvmmh_tactics + mma_tiler_mn_candi = list(_SM107_BF16_MMA_TILER_MN_CANDIDATES) # Skinny-M / small-N shapes under-use a 128x128 MMA tile (half the # M rows wasted, coarse N tiling, few CTAs); offer smaller tiles so @@ -11103,7 +11427,8 @@ def forward( ) -> None: (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) + split_k_slices, raster_along, + swizzle_size) = _parse_sm107_bf16_tactic(tactic) a_tensor, b_tensor, c_tensor = inputs m, k = a_tensor.shape @@ -11146,6 +11471,8 @@ def forward( max_num_ab_stage, split_k_slices, c_dtype_cutlass, + raster_along, + swizzle_size, ) if c_needs_copy: c_tensor.copy_(c_buf) @@ -11173,6 +11500,8 @@ def forward( preferred_cluster_shape_mn, cluster_shape_mn, max_num_ab_stage, + raster_along, + swizzle_size, self.use_tvm_ffi, c_dtype_cutlass, c_layout_key, @@ -11194,11 +11523,14 @@ 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=raster_along, + swizzle_size=swizzle_size) compile_args = [ m, n, k, batch_size, a_ptr, b_ptr, c_cute_tensor ] @@ -11257,6 +11589,8 @@ def _forward_split_k( max_num_ab_stage: int, split_k_slices: int, c_dtype_cutlass, + raster_along: str = "m", + swizzle_size: int = 1, ) -> None: """Direct split-K: each K-slice reduce-adds into the zeroed C.""" batch_size = 1 @@ -11270,6 +11604,8 @@ def _forward_split_k( cluster_shape_mn, max_num_ab_stage, split_k_slices, + raster_along, + swizzle_size, self.use_tvm_ffi, c_dtype_cutlass, c_layout_key, @@ -11302,7 +11638,9 @@ def _c_cute_view(): None, cluster_shape_mn, max_num_ab_stage, - split_k_slices=split_k_slices) + split_k_slices=split_k_slices, + raster_along=raster_along, + swizzle_size=swizzle_size) compiled_gemm = cute.compile( gemm.wrapper, m, @@ -13113,6 +13451,75 @@ def _( (4, 4), ] + def _parse_sm107_fp8_blockwise_tactic(tactic): + """Decode (use_2cta_instrs, mma_tiler, mma_inst_shape, + cluster_shape_mn[, raster_along_m, swizzle_size]); the scheduler + fields default to raster along M without swizzle.""" + if not isinstance(tactic, tuple): + return (False, (128, 128, 128), (128, 128, 64), (1, 1), True, 1) + if len(tactic) == 6: + return tactic + return (*tactic, True, 1) + + def _sm107_fp8_blockwise_nvmmh_tactics( + kernel_class, + m: int, + n: int, + k: int, + batch_size: int, + c_dtype, + max_tactics: int, + fp8_quantize_1x128: bool = False) -> List[Tuple]: + """SM107 FP8 blockwise tactics of the top ``max_tactics`` nvMMH + configs (tile, cluster, raster and swizzle from nvMMH). + + The kernel's CTA tile is 128x128 or 256x128 (two 128-row MMAs + sharing B) with K=128, and an even cluster M means 2-CTA MMA. It + has no split-K. The fused FP8 1x128 output needs 128x128 tiles. + """ + + def shape(config): + cta_m, cta_n, cta_k = config.cta + cluster_m, cluster_n = config.cluster + if (cta_k != 128 or cta_n != 128 or cta_m not in (128, 256) + or config.split_k != 1 or cluster_m not in (1, 2, 4) + or cluster_n not in (1, 2, 4)): + return None + use_2cta_instrs = cutedsl_nvmmh.pairs_ctas(config) + mma_tiler = (cta_m * (2 if use_2cta_instrs else 1), 128, 128) + mma_inst_shape = (256 if use_2cta_instrs else 128, 128, 64) + if not kernel_class.can_implement( + cutlass.Float8E4M3FN, cutlass.Float8E4M3FN, + cutlass.Float32, c_dtype, use_2cta_instrs, mma_tiler, + mma_inst_shape, config.cluster, m, n, k, batch_size, + "k", "k", "n", fp8_quantize_1x128): + return None + return use_2cta_instrs, mma_tiler, mma_inst_shape + + configs = cutedsl_nvmmh.query( + cutedsl_nvmmh.Problem(m, n, k, batch_size, + cutedsl_nvmmh.FP8_PRECISION), + cutedsl_nvmmh.query_count(max_tactics), + lambda config: shape(config) is not None, + backend=_SM107_FP8_BLOCKWISE_NVMMH_BACKEND, + cache_key=("sm107_fp8_blockwise", c_dtype, fp8_quantize_1x128)) + ranked = [] + for config in configs: + use_2cta_instrs, mma_tiler, mma_inst_shape = shape(config) + raster_along_m = cutedsl_nvmmh.cta_order_to_raster( + config.cta_order) == "m" + scheduler = (() if raster_along_m and config.swizzle == 1 else + (raster_along_m, config.swizzle)) + ranked.append( + (config, [(use_2cta_instrs, mma_tiler, mma_inst_shape, + config.cluster) + scheduler])) + return cutedsl_nvmmh.select_tactics(ranked, max_tactics) + + # The SM107 FP8 blockwise kernel has no split-K and supports both + # rasters. + _SM107_FP8_BLOCKWISE_NVMMH_BACKEND = cutedsl_nvmmh.Backend( + split_k_kind="NONE", supports_col_major_raster=True) + class CuteDSLFp8RubinGemmRunner(TunableRunner): """FP8 1x128/128x128 blockwise GEMM runner for SM107.""" @@ -13138,6 +13545,9 @@ def __init__(self, self.output_dtype = output_dtype self.use_tvm_ffi = use_tvm_ffi + def tactic_space_key(self): + return _nvmmh_tactic_space_key() + def get_valid_tactics( self, inputs: List[torch.Tensor], @@ -13158,6 +13568,14 @@ def get_valid_tactics( b_major = "k" c_major = "n" + nvmmh_config = AutoTuner.get().nvmmh_config + if nvmmh_config is not None: + nvmmh_tactics = _sm107_fp8_blockwise_nvmmh_tactics( + self.__class__.kernel_class, m, n, k, batch_size, + cutlass.BFloat16, nvmmh_config.max_tactics) + if nvmmh_tactics: + return nvmmh_tactics + 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 @@ -13186,15 +13604,9 @@ def forward( inputs: List[torch.Tensor], tactic, ) -> torch.Tensor: - if isinstance(tactic, tuple): - use_2cta_instrs, mma_tiler, mma_inst_shape, cluster_shape_mn = tactic - else: - use_2cta_instrs, mma_tiler, mma_inst_shape, cluster_shape_mn = [ - False, - (128, 128, 128), - (128, 128, 64), - (1, 1), - ] + (use_2cta_instrs, mma_tiler, mma_inst_shape, cluster_shape_mn, + raster_along_m, + swizzle_size) = _parse_sm107_fp8_blockwise_tactic(tactic) a_tensor, b_tensor, a_sf_tensor, b_sf_tensor = inputs m, k = a_tensor.shape n = b_tensor.shape[0] @@ -13233,6 +13645,8 @@ def forward( mma_tiler, mma_inst_shape, cluster_shape_mn, + raster_along_m, + swizzle_size, self.use_tvm_ffi, ) if cache_key not in self.__class__.kernel_cache: @@ -13264,6 +13678,8 @@ def forward( mma_tiler=mma_tiler, mma_inst_shape=mma_inst_shape, cluster_shape_mn=cluster_shape_mn, + swizzle_size=swizzle_size, + raster_along_m=raster_along_m, ) hardware_info = cutlass.utils.HardwareInfo() max_active_clusters = hardware_info.get_max_active_clusters( @@ -13384,6 +13800,9 @@ def __init__( use_tvm_ffi=use_tvm_ffi) self.fp8_quantize_1x128 = fp8_quantize_1x128 + def tactic_space_key(self): + return _nvmmh_tactic_space_key() + def get_valid_tactics( self, inputs: List[torch.Tensor], @@ -13398,6 +13817,16 @@ def get_valid_tactics( batch_size, m, k = inputs[0].shape n = inputs[1].shape[1] + nvmmh_config = AutoTuner.get().nvmmh_config + if nvmmh_config is not None: + nvmmh_tactics = _sm107_fp8_blockwise_nvmmh_tactics( + self.__class__.kernel_class, m, n, k, batch_size, + (cutlass.Float8E4M3FN + if self.fp8_quantize_1x128 else cutlass.BFloat16), + nvmmh_config.max_tactics, self.fp8_quantize_1x128) + if nvmmh_tactics: + return nvmmh_tactics + 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 @@ -13430,14 +13859,9 @@ def forward( tactic, ) -> None: """Run FP8 block-scaled BMM with TRT-LLM tensor layouts.""" - if isinstance(tactic, tuple): - (use_2cta_instrs, mma_tiler, mma_inst_shape, - cluster_shape_mn) = tactic - else: - use_2cta_instrs = False - mma_tiler = (128, 128, 128) - mma_inst_shape = (128, 128, 64) - cluster_shape_mn = (1, 1) + (use_2cta_instrs, mma_tiler, mma_inst_shape, cluster_shape_mn, + raster_along_m, + swizzle_size) = _parse_sm107_fp8_blockwise_tactic(tactic) if self.fp8_quantize_1x128: ( @@ -13509,6 +13933,8 @@ def forward( mma_tiler, mma_inst_shape, cluster_shape_mn, + raster_along_m, + swizzle_size, self.use_tvm_ffi, self.fp8_quantize_1x128, ) @@ -13550,6 +13976,8 @@ def forward( mma_inst_shape=mma_inst_shape, cluster_shape_mn=cluster_shape_mn, fp8_quantize_1x128=self.fp8_quantize_1x128, + swizzle_size=swizzle_size, + raster_along_m=raster_along_m, ) hardware_info = cutlass.utils.HardwareInfo() max_active_clusters = hardware_info.get_max_active_clusters( @@ -13802,6 +14230,7 @@ class CuteDSLBlockScaledRubinLinear(TunableRunner): kernel_class = Sm107BlockScaledPersistentDenseGemmKernel kernel_cache: ClassVar[dict] gemm_name: ClassVar[str] + nvmmh_precision: ClassVar[str] a_dtype: ClassVar[type] b_dtype: ClassVar[type] sf_dtype: ClassVar[type] @@ -13859,6 +14288,150 @@ def __eq__(self, other): return False return self.unique_id() == other.unique_id() + def tactic_space_key(self): + return _nvmmh_tactic_space_key() + + def _nvmmh_kernel_shapes(self, config, kernel_m: int, kernel_n: int, + real_k: int, swap_ab: bool): + """Map an nvMMH config onto this kernel, or None if it cannot run it. + + nvMMH's per-CTA tile M is 128, or 256 for B-reuse (two + 128-row MMAs per CTA); an even cluster M means 2-CTA MMA. + """ + cta_m, cta_n, cta_k = config.cta + cluster_m, cluster_n = config.cluster + # Cheap rejections first: can_implement prints every refusal. + if (cta_k != self.mma_tiler_k or cta_m not in (128, 256) + or cta_n not in (64, 128, 192, 256) + or cluster_m not in (1, 2, 4) + or cluster_n not in (1, 2, 4) + or real_k % self.mma_tiler_k != 0 + or config.split_k > real_k // self.mma_tiler_k + or -(-kernel_m // cta_m) < cluster_m): + return None + use_2cta = cutedsl_nvmmh.pairs_ctas(config) + mma_tiler_mnk = (cta_m * (2 if use_2cta else 1), cta_n, + self.mma_tiler_k) + mma_inst_shape = (256 if use_2cta else 128, cta_n, + self.mma_inst_shape_k) + if not self.__class__.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, + mma_tiler_mnk, + mma_inst_shape, + config.cluster, + config.split_k, + ): + return None + return mma_tiler_mnk, mma_inst_shape + + def _nvmmh_realizations(self, config, kernel_m: int, kernel_n: int, + real_k: int, swap_ab: bool) -> List[Tuple]: + """Tactics realizing one nvMMH config. + + Tile, cluster, split-K, raster and swizzle come from nvMMH. + Prefetch, the scheduler and the preferred/fallback cluster + launch (a preferred cluster above two CTAs with a 2x1 + fallback) are not modeled by nvMMH and are swept. + """ + shapes = self._nvmmh_kernel_shapes(config, kernel_m, kernel_n, + real_k, swap_ab) + if shapes is None: + return [] + mma_tiler_mnk, mma_inst_shape = shapes + use_2cta = mma_inst_shape[0] == 256 + cluster_shape_mn = config.cluster + split_k = config.split_k + swizzle_size = config.swizzle + raster_order = cutedsl_nvmmh.cta_order_to_raster( + config.cta_order) + swizzle_suffix = (swizzle_size, ) if swizzle_size > 1 else () + tactics = [] + for use_prefetch, scheduler_mode in itertools.product( + (False, True), + _dense_gemm_scheduler_modes( + self.scheduler_mode_candidates)): + if (scheduler_mode == "clc_dynamic" and raster_order == "n" + and not _clc_raster_n_has_launchable_grid( + kernel_m, kernel_n, split_k, mma_tiler_mnk[:2], + use_2cta, cluster_shape_mn)): + continue + tactics.append(("base", mma_tiler_mnk, mma_inst_shape, + cluster_shape_mn, swap_ab, use_prefetch, + scheduler_mode, raster_order, split_k) + + swizzle_suffix) + fallback_cluster_shape_mn = (2, 1) + if (split_k == 1 and use_2cta + and math.prod(cluster_shape_mn) > 2 and + Sm107BlockScaledPersistentDenseGemmMixedClustersKernel. + 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, + mma_tiler_mnk, + mma_inst_shape, + cluster_shape_mn, + fallback_cluster_shape_mn, + )): + for use_prefetch, scheduler_mode in itertools.product( + (False, True), + _dense_gemm_scheduler_modes( + self.mixed_scheduler_mode_candidates)): + # CLC mixed clusters only launch with raster M. + if scheduler_mode == "clc_dynamic" and raster_order != "m": + continue + scheduler = ((scheduler_mode, ) + if self.mixed_scheduler_mode_candidates + != ("static", ) or swizzle_suffix else ()) + tactics.append(("mixed_clusters", mma_tiler_mnk, + mma_inst_shape, cluster_shape_mn, + fallback_cluster_shape_mn, swap_ab, + use_prefetch) + scheduler + + (raster_order, ) + swizzle_suffix) + return tactics + + def _nvmmh_tactics(self, m: int, n: int, real_k: int, + swap_ab_candidates: List[bool], + max_tactics: int) -> List[Tuple]: + """Tactics of the top ``max_tactics`` nvMMH configs.""" + rankings = [] + for swap_ab in swap_ab_candidates: + kernel_m, kernel_n = (n, m) if swap_ab else (m, n) + problem = cutedsl_nvmmh.Problem( + kernel_m, kernel_n, real_k, 1, self.nvmmh_precision, + cutedsl_nvmmh.LAYOUT_COL_MAJOR_C + if swap_ab else cutedsl_nvmmh.LAYOUT_ROW_MAJOR_C) + configs = cutedsl_nvmmh.query( + problem, + cutedsl_nvmmh.query_count(max_tactics), + lambda config: self._nvmmh_kernel_shapes( + config, kernel_m, kernel_n, real_k, swap_ab + ) is not None, + backend=_SM107_BLOCKSCALED_NVMMH_BACKEND, + cache_key=(self.__class__.kernel_class.__name__, + self.mma_tiler_k, swap_ab)) + rankings.append([ + (config, + self._nvmmh_realizations(config, kernel_m, kernel_n, + real_k, swap_ab)) + for config in configs + ]) + return cutedsl_nvmmh.select_tactics( + _merge_nvmmh_rankings(rankings), max_tactics) + def get_valid_tactics( self, inputs: List[torch.Tensor], @@ -13928,6 +14501,15 @@ def get_valid_tactics( f"Using base swap_ab={base_swap_ab_candidates}, " f"mixed-cluster swap_ab={mixed_swap_ab_candidates}") + nvmmh_config = AutoTuner.get().nvmmh_config + if nvmmh_config is not None: + # nvMMH ranks both A/B orientations; it decides which one + # each kernel config runs in. + nvmmh_tactics = self._nvmmh_tactics( + m, n, real_k, [False, True], nvmmh_config.max_tactics) + if nvmmh_tactics: + return nvmmh_tactics + # SM107 dense block-scaled tile sizes. # b-reuse is enabled when mma_tiler_m = 2 * mma_inst_shape_m mma_tiler_k = self.mma_tiler_k @@ -14129,16 +14711,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 size") (_, mma_tiler_mnk, mma_inst_shape, preferred_cluster_shape_mn, cluster_shape_mn, swap_ab, use_prefetch, *tuning_options) = tactic @@ -14154,14 +14737,17 @@ def forward( raster_order = "m" else: raster_order = tuning_options[0] - else: + elif len(tuning_options) == 2: scheduler_mode, raster_order = tuning_options + else: + 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 " + "swizzle size") (_, mma_tiler_mnk, mma_inst_shape, cluster_shape_mn, swap_ab, use_prefetch, *tuning_options) = tactic preferred_cluster_shape_mn = None @@ -14171,8 +14757,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, @@ -14200,6 +14788,11 @@ def forward( raise ValueError( f"CuteDSL SM107 split_k must be a positive integer, got {split_k!r}" ) + if (not isinstance(swizzle_size, int) + or isinstance(swizzle_size, bool) or swizzle_size < 1): + raise ValueError( + "CuteDSL SM107 swizzle_size must be a positive " + f"integer, got {swizzle_size!r}") if raster_order not in {"m", "n"}: raise ValueError("CuteDSL SM107 raster_order must be 'm' " f"or 'n', got {raster_order!r}") @@ -14402,8 +14995,9 @@ def forward( 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) + swizzle_size, 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 @@ -14456,6 +15050,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, ) @@ -14466,6 +15061,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, @@ -14568,6 +15164,7 @@ class CuteDSLNVFP4RubinLinear(CuteDSLBlockScaledRubinLinear): kernel_cache = dict() gemm_name = "NVFP4" + nvmmh_precision = cutedsl_nvmmh.FP4_PRECISION a_dtype = cutlass.Float4E2M1FN b_dtype = cutlass.Float4E2M1FN sf_dtype = cutlass.Float8E4M3FN @@ -14595,6 +15192,7 @@ class CuteDSLMXFP8RubinLinear(CuteDSLBlockScaledRubinLinear): kernel_cache = dict() gemm_name = "MXFP8" + nvmmh_precision = cutedsl_nvmmh.FP8_PRECISION a_dtype = cutlass.Float8E4M3FN b_dtype = cutlass.Float8E4M3FN sf_dtype = cutlass.Float8E8M0FNU diff --git a/tensorrt_llm/_torch/custom_ops/cutedsl_matmul_heuristics.py b/tensorrt_llm/_torch/custom_ops/cutedsl_matmul_heuristics.py index 0113a6c3a515..3a4b5a0c314e 100644 --- a/tensorrt_llm/_torch/custom_ops/cutedsl_matmul_heuristics.py +++ b/tensorrt_llm/_torch/custom_ops/cutedsl_matmul_heuristics.py @@ -1,23 +1,30 @@ # SPDX-FileCopyrightText: Copyright (c) 2022-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""Adapter around NVIDIA Matmul Heuristics (nvMatmulHeuristics) for the CuTe DSL -GEMM autotuner. - -The autotuner enumerates every valid CuTe DSL tactic and then compiles + -benchmarks all of them, which dominates autotuning wall-time. This module uses -nvMatmulHeuristics' analytical performance model to rank the candidate -(CTA-tile, cluster) configurations so the runner can keep only the top-K before -profiling. - -Everything here is best-effort and opt-in: if the library is missing, the -problem is unsupported, or the query fails, the public helpers return an empty -ranking so callers fall back to the full tactic list and never regress. +"""nvMatmulHeuristics (nvMMH) guidance for the CuTe DSL GEMM/BMM autotuner. + +The autotuner normally compiles and benchmarks every tactic a CuTe DSL GEMM or +BMM runner enumerates. When ``autotuner_nvmmh_config`` installs an +:class:`~tensorrt_llm._torch.autotuner.NvMMHConfig` on the AutoTuner, runners +instead build their tactics from the kernel configurations nvMMH's analytical +model ranks highest for the problem. Every field nvMMH models (CTA tile, +cluster shape, split-K, rasterization order, swizzle) is taken as returned; +runners only sweep knobs nvMMH does not model. + +Each runner describes what its kernel can execute through a validity check, so +nvMMH ranks implementable configurations only. Both the 0.1.0.x wheels (the +``nvidia-matmul-heuristics`` requirement) and newer libheuristics builds are +supported. If the library is missing or a query fails, :func:`query` returns +an empty list and the runner falls back to its full tactic space. """ -import os +import ctypes +import importlib.metadata +import inspect import re +import threading +from dataclasses import dataclass from functools import lru_cache -from typing import List, NamedTuple, Optional, Tuple +from typing import Callable, Dict, Hashable, List, NamedTuple, Optional, Tuple from tensorrt_llm.logger import logger @@ -25,253 +32,340 @@ import nvMatmulHeuristics as _nvmmh IS_NVMMH_AVAILABLE = True -except ImportError: +except (ImportError, OSError): _nvmmh = None IS_NVMMH_AVAILABLE = False -# 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 -# cute_dsl_custom_ops.py). "F4F4T" is the cuBLAS-style token (fp4 A, fp4 B, -# bf16 output; bf16->T) and TN_ROW_MAJOR is the K-major A/B layout enum name on -# NvMatmulHeuristicsMatmulLayout. The whole path is best-effort: if the token is -# ever rejected by the installed wheel, rank_configs degrades to an empty -# ranking and the caller falls back to the full tactic sweep. -NVFP4_PRECISION = "F4F4T" -NVFP4_LAYOUT = "TN_ROW_MAJOR" +# nvMMH precision strings name one data type per character (T = BF16, +# S = FP32, H = FP16, Q = FP8 E4M3, O = FP4 E2M1). Three characters are +# (A/B, compute, C/D); five are (A, B, C, compute, D). +BF16_PRECISION = "TST" +BF16_FP32_OUTPUT_PRECISION = "TSS" +BF16_FP8_OUTPUT_PRECISION = "TTQSQ" +FP8_PRECISION = "QQTST" +FP8_FP16_OUTPUT_PRECISION = "QQHSH" +FP4_PRECISION = "OOTST" + +# Layouts of the CuTe DSL GEMMs: A [M, K] and B [N, K] are K-major ("TN") and +# C [M, N] is row-major. Swapping A and B computes C^T, a column-major [N, M]. +LAYOUT_ROW_MAJOR_C = "TN_ROW_MAJOR" +LAYOUT_COL_MAJOR_C = "TN_COL_MAJOR" + +# nvMMH backend split-K kinds, by name so callers need not import nvMMH. +SPLIT_K_IN_PLACE = "IN_PLACE" +SPLIT_K_OUT_OF_PLACE = "OUT_OF_PLACE" + +# Configurations requested from nvMMH per kept tactic: several returned +# configurations differ only in fields a CuTe DSL kernel does not expose +# (e.g. the warp tile) and collapse onto one tactic. +QUERY_OVERSAMPLING = 8 +MIN_QUERY_COUNT = 32 + + +class KernelConfig(NamedTuple): + """One nvMMH kernel configuration (a CUTLASS3 ``nvmmhKernelConfiguration``). + + ``cta`` is the per-CTA (M, N, K) tile. On Blackwell-class GPUs nvMMH runs + 2-CTA MMA instructions whenever ``cluster[0]`` is even (see + :func:`pairs_ctas`). ``cta_order`` 0 rasterizes along M and 1 along N + (see :func:`cta_order_to_raster`). + ``score`` is the value nvMMH ranks configurations by (lower is better), + comparable across problems that compute the same product (see + :func:`query`). + """ -CtaCluster = Tuple[Tuple[int, int], Tuple[int, int]] + cta: Tuple[int, int, int] + cluster: Tuple[int, int] + split_k: int + swizzle: int + cta_order: int + stages: int + warp: Tuple[int, int, int] + instr: Tuple[int, int, int] + score: float -# 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"}) -# Accepted spellings that normalize onto a canonical field name. -_NVMMH_FIELD_ALIASES = {"cta_tile": "tile", "mma_tiler": "tile"} +@dataclass(frozen=True) +class Problem: + """A GEMM/BMM problem in nvMMH terms (after any A/B swap).""" + m: int + n: int + k: int + batch: int + precision: str + layout: str = LAYOUT_ROW_MAJOR_C -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.""" - cta: Tuple[int, int] - cluster: Tuple[int, int] - swizzle_factor: int = 1 - cta_order: int = 0 +@dataclass(frozen=True) +class Backend: + """CUTLASS3 backend properties describing a CuTe DSL kernel family. + Properties left as None keep nvMMH's CUTLASS3 defaults. + """ + + split_k_kind: Optional[str] = None + supports_col_major_raster: Optional[bool] = None + supports_odd_cluster_n: Optional[bool] = None + # The kernel launches a 2-CTA cluster above two CTAs with a 2x1 fallback + # cluster (preferred-cluster launch). + supports_fallback_cluster: Optional[bool] = None + cta_tile_m_div: Optional[int] = None + cta_tile_n_div: Optional[int] = None -def nvmmh_fields() -> set: - """Which knobs nvMatmulHeuristics drives (vs. the runner sweeping them). - 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). +@lru_cache(maxsize=1) +def library_version() -> str: + """Installed nvMMH distribution version (part of cache keys).""" + try: + return importlib.metadata.version("nvidia-matmul-heuristics") + except importlib.metadata.PackageNotFoundError: + return str(getattr(_nvmmh, "__version__", "unknown")) + + +def cta_order_to_raster(cta_order: int) -> str: + """nvMMH CUTLASS3 CTA order -> CuTe DSL raster ("m": M-fastest). + + 0 rasterizes along M and 1 along N, as in CUTLASS's own nvMMH provider for + the same CUTLASS3 target (``raster_order``): + https://github.com/NVIDIA/cutlass/blob/d4bbf728ca52ce23a549a299e229666e734ffb7c/python/cutlass_library/heuristics_provider.py#L222 """ - 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 + return "n" if int(cta_order) else "m" + +@lru_cache(maxsize=1) +def _pairs_only_cluster_m_2() -> bool: + """0.1.0.x PyPI releases model 2-CTA MMA for cluster M == 2 only.""" + match = re.fullmatch(r"0\.1\.0\.(\d+)", library_version()) + return match is not None and int(match.group(1)) < 10000 -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 pairs_ctas(config: KernelConfig) -> bool: + """Whether nvMMH models this config with 2-CTA (CTA pair) MMA. -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() + nvMMH's CUTLASS3 target pairs CTAs along an even cluster M extent on + Blackwell-class GPUs (only cluster M == 2 before libheuristics 0.2), so a + CuTe DSL kernel realizes the config with ``use_2cta_instrs`` and an MMA + tile of twice the per-CTA M. + """ + cluster_m = int(config.cluster[0]) + if _pairs_only_cluster_m_2(): + return cluster_m == 2 + return cluster_m % 2 == 0 -def nvmmh_max_tactics() -> int: - """Number of heuristic-ranked (tile, cluster) configs to keep per problem. +def _interface_class(): + """The interface class that manages discovery sets implicitly. - 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. + 0.1.0.x names it NvMatmulHeuristicsInterfaceEx; newer wheels moved the + implicit discovery handling to NvMatmulHeuristicsInterface. """ - try: - return max(1, int(os.environ.get("TRTLLM_CUTEDSL_NVMMH_MAX_TACTICS", "5"))) - except ValueError: - return 5 + interface = _nvmmh.NvMatmulHeuristicsInterface + if hasattr(interface, "_ensure_discovery_loaded"): + return interface + return _nvmmh.NvMatmulHeuristicsInterfaceEx + + +_interfaces: Dict[Tuple[int, str], object] = {} +_interfaces_lock = threading.Lock() -@lru_cache(maxsize=None) def _get_interface(precision: str): - """Build and cache an interface for a given precision (None on failure).""" - if not IS_NVMMH_AVAILABLE: - return None - try: - return _nvmmh.NvMatmulHeuristicsInterface( - _nvmmh.NvMatmulHeuristicsTarget.CUTLASS3, - precision=precision, - flags=_nvmmh.NvMatmulHeuristicsFlags.PERF_MODEL_BASED_AUTO_TUNING, - ) - except Exception as e: # noqa: BLE001 - any failure must degrade gracefully - logger.warning_once( - f"[nvMatmulHeuristics] Failed to build interface for " - f"precision={precision}: {e}. Falling back to full tactic list.", - key="nvmmh_interface_init_failure", + """One interface per (thread, precision); nvMMH handles are not shared.""" + key = (threading.get_ident(), precision) + with _interfaces_lock: + interface = _interfaces.get(key) + if interface is not None: + return interface + cls = _interface_class() + kwargs = dict(flags=_nvmmh.NvMatmulHeuristicsFlags.PERF_MODEL_BASED_AUTO_TUNING) + if "precision" in inspect.signature(cls.__init__).parameters: + kwargs["precision"] = precision + interface = cls(_nvmmh.NvMatmulHeuristicsTarget.CUTLASS3, **kwargs) + with _interfaces_lock: + _interfaces[key] = interface + return interface + + +def _make_problem(interface, problem: Problem): + layout = getattr(_nvmmh.NvMatmulHeuristicsMatmulLayout, problem.layout) + make = interface.makeNvMatmulHeuristicsProblem + if "mnk" in inspect.signature(make).parameters: + return make((problem.m, problem.n, problem.k), layout, batch_size=problem.batch) + return make(problem.m, problem.n, problem.k, layout, batch_size=problem.batch) + + +def _set_property(interface, backend, name: str, value: int) -> None: + prop = getattr(_nvmmh.NvMatmulHeuristicsBackendProperty, name) + setter = getattr(interface, "setBackendPropertyValue", None) + if setter is not None: + setter(backend, prop, int(value)) + return + # 0.1.0.x takes a raw value buffer. + buffer = ctypes.c_int32(int(value)) + interface.setBackendValueProperty(backend, prop, ctypes.byref(buffer), ctypes.sizeof(buffer)) + + +def _set_validity_check(interface, backend, accept: Callable[[object], bool]) -> None: + """Register ``accept(raw_config) -> bool`` as the additional validity check.""" + kind = _nvmmh.NvMatmulHeuristicsBackendPropertyCallbackKind.KERNEL_ADDITIONAL_VALIDITY_CHECK + setter = getattr(interface, "backendSetCallbackProperty", None) + if setter is not None: + setter(backend, kind, lambda config: int(bool(accept(config)))) + return + # 0.1.0.x passes a configuration pointer and an unused problem pointer. + interface.setBackendCallbackProperty( + backend, kind, lambda config_ptr, _problem_ptr: int(bool(accept(config_ptr.contents))) + ) + + +def _configure_backend(interface, backend, spec: Backend) -> None: + if spec.split_k_kind is not None: + _set_property( + interface, + backend, + "SPLIT_K_KIND", + int(getattr(_nvmmh.NvMatmulHeuristicsSplitKKind, spec.split_k_kind)), ) - return None + for name, value in ( + ("HAS_COL_MAJOR_RASTER", spec.supports_col_major_raster), + ("SUPPORTS_ODD_CLUSTER_N", spec.supports_odd_cluster_n), + ("SUPPORTS_FALLBACK_CLUSTER", spec.supports_fallback_cluster), + ("CTA_TILE_M_DIV_REQUIREMENT", spec.cta_tile_m_div), + ("CTA_TILE_N_DIV_REQUIREMENT", spec.cta_tile_n_div), + ): + if value is not None: + _set_property(interface, backend, name, int(value)) + + +def _to_kernel_config(raw, score: float) -> KernelConfig: + return KernelConfig( + cta=tuple(int(v) for v in raw.cta), + cluster=(max(1, int(raw.cluster[0])), max(1, int(raw.cluster[1]))), + split_k=max(1, int(raw.splitK)), + swizzle=max(1, int(raw.gridSwizzle)), + cta_order=int(raw.ctaOrder), + stages=int(raw.loadStages), + warp=tuple(int(v) for v in raw.warp), + instr=tuple(int(v) for v in raw.instr), + score=float(score), + ) + +def _ranking_score(interface, backend, nv_problem, raw, precision: str, fallback: float) -> float: + """The score nvMMH ranks ``raw`` by, with this backend's properties. -def _get_layout(name: str): + Libraries with ``RANKING_SCORE`` report it directly; older libheuristics + builds approximate it with ``RUNTIME_RELATIVE_FAST_S`` (their ranking + metric without the Rubin L2 simulation). 0.1.0.x has no per-backend + estimate: ``fallback`` is the wrapper's ``runtime`` field, a default-backend + absolute runtime. + """ + estimate = getattr(interface, "estimateSiliconMetricEx", None) + if estimate is None or not hasattr(interface, "setBackendPropertyValue"): + return fallback + metrics = _nvmmh.NvMatmulHeuristicsSiliconMetric + metric = getattr(metrics, "RANKING_SCORE", metrics.RUNTIME_RELATIVE_FAST_S) try: - return getattr(_nvmmh.NvMatmulHeuristicsMatmulLayout, name) - except AttributeError: - logger.warning_once( - f"[nvMatmulHeuristics] Unknown layout '{name}'. Falling back to full tactic list.", - key="nvmmh_unknown_layout", - ) - return None + value = estimate(precision, backend, nv_problem, raw, int(metric)) + except Exception: # noqa: BLE001 - only orders merged rankings + return fallback + return float(value) if value and value > 0 else fallback -def _as_int_pair(values) -> Optional[Tuple[int, int]]: +def _query_uncached( + problem: Problem, count: int, backend_spec: Backend, is_valid: Callable[[KernelConfig], bool] +) -> List[KernelConfig]: + interface = _get_interface(problem.precision) + backend = interface.createBackend(_nvmmh.NvMatmulHeuristicsTarget.CUTLASS3) try: - nums = [int(v) for v in values] - except (TypeError, ValueError): - return None - if len(nums) < 2: - return None - return (nums[0], nums[1]) - - -def _extract_int(config, kernel, dict_key: str, attr: str, str_re: str, default: int) -> int: - """Read one scalar field, tolerant of dict / object / printable-string.""" - if isinstance(config, dict) and dict_key in config: - try: - return int(config[dict_key]) - except (TypeError, ValueError): - pass - if hasattr(kernel, attr): - try: - return int(getattr(kernel, attr)) - except (TypeError, ValueError): - pass - m = re.search(str_re, kernel if isinstance(kernel, str) else str(kernel)) - if m: - try: - return int(m.group(1)) - except (TypeError, ValueError): - pass - return default - - -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. - """ - kernel = config.get("kernel") if isinstance(config, dict) else config - cta = cluster = None - - # Form 1: flat dict fields (DKG getEx-style wrapper). - if isinstance(config, dict) and "cta_tile_m" in config: - cta = _as_int_pair((config["cta_tile_m"], config["cta_tile_n"])) - cluster = _as_int_pair((config.get("cluster_m", 1), config.get("cluster_n", 1))) - - # Form 2: kernel object with attributes (public wheel's GemmConfig uses - # cluster_m/cluster_n; the DKG internal wrapper used cga_m/cga_n). - if (cta is None or cluster is None) and hasattr(kernel, "cta_tile_m"): - cta = _as_int_pair((kernel.cta_tile_m, kernel.cta_tile_n)) - cluster_m = getattr(kernel, "cluster_m", getattr(kernel, "cga_m", 1)) - cluster_n = getattr(kernel, "cluster_n", getattr(kernel, "cga_n", 1)) - cluster = _as_int_pair((cluster_m, cluster_n)) - - # Form 3: printable string like "... cta(128 16 128) ... cluster(2 1) ...". - if cta is None or cluster is None: - text = kernel if isinstance(kernel, str) else str(kernel) - cta_m = re.search(r"cta\(\s*([\d\s]+?)\)", text) - cluster_m = re.search(r"cluster\(\s*([\d\s]+?)\)", text) - if cta_m: - cta = _as_int_pair(cta_m.group(1).split()) - cluster = _as_int_pair(cluster_m.group(1).split()) if cluster_m else (1, 1) - - if cta is None or cluster is None: - 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 - # 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) + _configure_backend(interface, backend, backend_spec) + _set_validity_check(interface, backend, lambda raw: is_valid(_to_kernel_config(raw, 0.0))) + nv_problem = _make_problem(interface, problem) + configs = interface.getEx(nv_problem, int(count), backend, problem.precision) + ranked = [] + for entry in configs: + raw = entry["nvmmhKernelConfiguration"] + score = _ranking_score( + interface, backend, nv_problem, raw, problem.precision, entry.get("runtime", 0.0) + ) + config = _to_kernel_config(raw, score) + # Wrappers that drop the callback's verdict still return only + # nvMMH-ranked configs; re-check them here. + if is_valid(config): + ranked.append(config) + finally: + interface.destroyBackend(backend) + return ranked + + +_query_cache: Dict[Hashable, List[KernelConfig]] = {} +_query_cache_lock = threading.Lock() -def rank_configs(m: int, n: int, k: int, precision: str, count: int) -> List[HeuristicConfig]: - """Return up to ``count`` HeuristicConfig entries ranked best-first. +def query( + problem: Problem, + count: int, + is_valid: Callable[[KernelConfig], bool], + backend: Backend = Backend(), + cache_key: Optional[Hashable] = None, +) -> List[KernelConfig]: + """Return up to ``count`` valid nvMMH configs in nvMMH's ranking order. - Returns an empty list if heuristics are unavailable or the query fails, so - callers can fall back to their full tactic list. + nvMMH's own order is kept. Each config carries its ranking ``score``, so + rankings of several problems computing the same product (e.g. both A/B + orientations) can be merged consistently with that order. + ``cache_key`` identifies the validity check; results are memoized per + (problem, count, backend, cache_key) when it is given. Failures are + logged once and yield an empty list. """ - interface = _get_interface(precision) - if interface is None: - return [] - layout = _get_layout(NVFP4_LAYOUT) - if layout is None: + if not IS_NVMMH_AVAILABLE: return [] + key = None + if cache_key is not None: + key = (problem, int(count), backend, cache_key) + with _query_cache_lock: + cached = _query_cache.get(key) + if cached is not None: + return list(cached) try: - # hw=None targets the current GPU. - try: - 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) - except Exception as e: # noqa: BLE001 - any failure must degrade gracefully + ranked = _query_uncached(problem, count, backend, is_valid) + except Exception as e: # noqa: BLE001 - guidance must never break tuning logger.warning_once( - f"[nvMatmulHeuristics] Query failed for " - f"precision={precision}, mnk=({m},{n},{k}): {e}. " - f"Falling back to full tactic list.", + f"[nvMatmulHeuristics] Query failed for {problem} with nvMMH " + f"{library_version()}: {e}. Falling back to the full tactic list.", key="nvmmh_query_failure", ) return [] + if key is not None: + with _query_cache_lock: + _query_cache[key] = list(ranked) + return ranked - if not configs: - return [] - # Sort by estimated runtime when present; the API may already be ranked. - def _runtime(c): - return c.get("runtime", float("inf")) if isinstance(c, dict) else float("inf") +def query_count(max_tactics: int) -> int: + """How many configs to request so ``max_tactics`` distinct tactics remain.""" + return max(MIN_QUERY_COUNT, QUERY_OVERSAMPLING * int(max_tactics)) - ordered = sorted(configs, key=_runtime) - ranked: List[HeuristicConfig] = [] - for c in ordered: - cfg = _extract_config(c) - if cfg is not None and cfg not in ranked: - ranked.append(cfg) - return ranked +def select_tactics(ranked: List[Tuple[KernelConfig, List[tuple]]], max_tactics: int) -> List[tuple]: + """Expand the top ``max_tactics`` nvMMH configs into runner tactics. -def rank_tile_cluster_configs( - m: int, n: int, k: int, precision: str, count: int -) -> 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)] + ``ranked`` pairs each config (nvMMH order) with the tactics realizing it: + one per combination of the knobs nvMMH does not model. Configs that map + onto an already kept realization (e.g. differing only in warp tile) do not + count against ``max_tactics``. + """ + selected: List[tuple] = [] + seen = set() + kept = 0 + for _, tactics in ranked: + fresh = [t for t in tactics if t not in seen] + if not fresh: + continue + selected.extend(fresh) + seen.update(fresh) + kept += 1 + if kept >= max_tactics: + break + return selected diff --git a/tensorrt_llm/_torch/custom_ops/torch_custom_ops.py b/tensorrt_llm/_torch/custom_ops/torch_custom_ops.py index 24ac66fe074c..80463966872a 100644 --- a/tensorrt_llm/_torch/custom_ops/torch_custom_ops.py +++ b/tensorrt_llm/_torch/custom_ops/torch_custom_ops.py @@ -59,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, _nvmmh_tactic_space_key) if IS_CUTLASS_DSL_RUBIN_AVAILABLE: from tensorrt_llm._torch.custom_ops.cute_dsl_custom_ops import \ @@ -1311,6 +1311,12 @@ def unique_id(self): allowed_tuple = tuple(self.allowed_backends) return (self.output_buffer_kind, self.output_dtype, allowed_tuple) + def tactic_space_key(self) -> Optional[Tuple]: + """The CuTe DSL backend's tactics are pruned when nvMMH is enabled.""" + if IS_CUTLASS_DSL_AVAILABLE and self._is_backend_allowed("cutedsl"): + return _nvmmh_tactic_space_key() + return None + def _is_backend_allowed(self, backend_name: str) -> bool: """Check if a backend is allowed based on allowed_backends list.""" return backend_name in self.allowed_backends diff --git a/tensorrt_llm/_torch/kernels/blackwell/blockwise_gemm/blockwise_gemm.py b/tensorrt_llm/_torch/kernels/blackwell/blockwise_gemm/blockwise_gemm.py index b9338d4464c9..cbdb2be9d7f6 100644 --- a/tensorrt_llm/_torch/kernels/blackwell/blockwise_gemm/blockwise_gemm.py +++ b/tensorrt_llm/_torch/kernels/blackwell/blockwise_gemm/blockwise_gemm.py @@ -113,7 +113,7 @@ see detailed valid dtype combinations in below Sm100BlockwiseGemmKernel class documentation * A/B tensor must have the same data type * Mma tiler M must be 64/128/256 -* Mma tiler N must be 128, align with the scaleB requirement +* Mma tiler N must be 32, 64 or 128, so that a tile never spans two scaleB blocks * Cluster shape M/N must be positive and power of 2, total cluster size <= 16 * Cluster shape M must be multiple of 2 * The contiguous dimension of A/B/C tensors must be at least 16 bytes aligned @@ -145,7 +145,7 @@ class Sm100BlockwiseGemmKernel: :note: Constraints: - MMA tiler M must be 64/128/256 - - MMA tiler N must be 128 + - MMA tiler N must be 32, 64 or 128 - Cluster shape M must be multiple of 2 - Cluster shape M/N must be positive and power of 2, total cluster size <= 16 @@ -165,6 +165,9 @@ def __init__( use_2cta_instrs: bool, mma_tiler_mn: Tuple[int, int], cluster_shape_mn: Tuple[int, int], + swizzle_size: int = 1, + raster_along_m: bool = True, + split_k: int = 1, ): """Initializes the configuration for a Blackwell blockwise dense GEMM kernel. @@ -187,11 +190,24 @@ def __init__( :type use_2cta_instrs: bool :param cluster_shape_mn: Tuple (ClusterM, ClusterN) shape of the cluster. :type cluster_shape_mn: Tuple[int, int] + :param swizzle_size: Cluster swizzle width of the persistent tile scheduler. + :type swizzle_size: int + :param raster_along_m: Whether consecutive clusters advance along M. + :type raster_along_m: bool + :param split_k: Number of K slices per output tile. Every slice + reduce-adds its partial into C through TMA, so callers must zero C + before every launch. + :type split_k: int """ self.acc_dtype: Type[cutlass.Numeric] = acc_dtype self.use_2cta_instrs = use_2cta_instrs self.cluster_shape_mn = cluster_shape_mn + self.swizzle_size = swizzle_size + self.raster_along_m = raster_along_m + if split_k < 1: + raise ValueError(f"split_k must be >= 1, got {split_k}") + self.split_k = split_k # K dimension is deferred in _setup_attributes self.mma_tiler = (*mma_tiler_mn, 1) @@ -248,6 +264,20 @@ def __init__( # TMEM offset for final accumulator self.tmem_final_offset = 384 + def _split_k_range(self, k_tile_cnt, l_coord): + """Return ``(k_tile_begin, k_tiles)`` of the split-K slice at scheduler L ``l_coord``. + + Slices take balanced ranges, ``[floor(k_tile_cnt * s / split_k), + floor(k_tile_cnt * (s + 1) / split_k))`` for slice ``s = l_coord % split_k``, + so any ``split_k <= k_tile_cnt`` gives every slice at least one K tile. + """ + if cutlass.const_expr(self.split_k == 1): + return 0, k_tile_cnt + split_idx = l_coord % self.split_k + k_tile_begin = (k_tile_cnt * split_idx) // self.split_k + k_tile_end = (k_tile_cnt * (split_idx + 1)) // self.split_k + return k_tile_begin, k_tile_end - k_tile_begin + def _setup_attributes(self): """Set up configurations that are dependent on GEMM inputs @@ -295,8 +325,11 @@ def _setup_attributes(self): self.scale_granularity_m = 1 self.scale_granularity_n = 128 self.scale_granularity_k = 128 + # A CTA tile narrower than a scale block (N = 32 or 64) lies inside one + # 128-column block of B and applies that block's scale to all columns. + self.sfb_tile_granularity_n = min(self.scale_granularity_n, self.cta_tile_shape_mnk[1]) self.scale_m_per_tile = self.cta_tile_shape_mnk[0] // self.scale_granularity_m - self.scale_n_per_tile = self.cta_tile_shape_mnk[1] // self.scale_granularity_n + self.scale_n_per_tile = self.cta_tile_shape_mnk[1] // self.sfb_tile_granularity_n self.scale_k_per_tile = self.cta_tile_shape_mnk[2] // self.scale_granularity_k if self.scale_k_per_tile != 1: @@ -376,7 +409,7 @@ def _setup_attributes(self): ) self.sfb_smem_layout_staged = cute.make_layout( ( - (self.scale_granularity_n, self.scale_n_per_tile), + (self.sfb_tile_granularity_n, self.scale_n_per_tile), (self.scale_granularity_k, self.scale_k_per_tile), self.num_scale_stage, ), @@ -489,8 +522,11 @@ def __call__( tma_tensor_c = None c_cta_v_layout = cute.composition(cute.make_identity_layout(c.shape), self.epi_tile) epi_smem_layout = cute.slice_(self.c_smem_layout_staged, (None, None, 0)) + c_store_op = cpasync.CopyBulkTensorTileS2GOp() + if cutlass.const_expr(self.split_k > 1): + c_store_op = cpasync.CopyReduceBulkTensorTileS2GOp() tma_atom_c, tma_tensor_c = cpasync.make_tiled_tma_atom( - cpasync.CopyBulkTensorTileS2GOp(), + c_store_op, c, epi_smem_layout, c_cta_v_layout, @@ -529,7 +565,13 @@ def __call__( # Compute grid size self.tile_sched_params, grid = self._compute_grid( - c, self.cta_tile_shape_mnk, self.cluster_shape_mn, max_active_clusters + c, + self.cta_tile_shape_mnk, + self.cluster_shape_mn, + max_active_clusters, + self.swizzle_size, + self.raster_along_m, + self.split_k, ) self.buffer_align_bytes = 1024 @@ -865,7 +907,7 @@ def kernel( sSFB_view_as_C_layout = cute.make_layout( ( self.cta_tile_shape_mnk[0], - (self.scale_granularity_n, self.scale_n_per_tile), + (self.sfb_tile_granularity_n, self.scale_n_per_tile), self.num_scale_stage, ), stride=(0, (0, 1), self.scale_n_per_tile), @@ -1033,6 +1075,16 @@ def kernel( tile_info[1], tile_info[2], ) + # The scheduler's L enumerates (batch, split-K slice) pairs. + k_tile_begin, tile_k_tile_cnt = self._split_k_range( + k_tile_cnt, mma_tile_coord_mnl[2] + ) + if cutlass.const_expr(self.split_k > 1): + mma_tile_coord_mnl = ( + mma_tile_coord_mnl[0], + mma_tile_coord_mnl[1], + mma_tile_coord_mnl[2] // self.split_k, + ) # # Slice to per mma tile index # @@ -1044,14 +1096,17 @@ def kernel( # Peek (try_wait) AB buffer empty for k_tile = prefetch_k_tile_cnt ab_producer_state.reset_count() peek_ab_empty_status = cutlass.Boolean(1) - if ab_producer_state.count < k_tile_cnt: + if ab_producer_state.count < tile_k_tile_cnt: peek_ab_empty_status = ab_pipeline.producer_try_acquire(ab_producer_state) # # Tma load loop # - for k_tile in cutlass.range(0, k_tile_cnt, 1, unroll=1): - tAgA_k = tAgA_slice[(None, ab_producer_state.count)] - tBgB_k = tBgB_slice[(None, ab_producer_state.count)] + for k_tile in cutlass.range(0, tile_k_tile_cnt, 1, unroll=1): + k_coord = ab_producer_state.count + if cutlass.const_expr(self.split_k > 1): + k_coord = k_tile_begin + ab_producer_state.count + tAgA_k = tAgA_slice[(None, k_coord)] + tBgB_k = tBgB_slice[(None, k_coord)] tAsA_pipe = tAsA[(None, ab_producer_state.index)] tBsB_pipe = tBsB[(None, ab_producer_state.index)] @@ -1063,7 +1118,7 @@ def kernel( # Peek (try_wait) AB buffer empty for k_tile = prefetch_k_tile_cnt + k_tile + 1 ab_producer_state.advance() peek_ab_empty_status = cutlass.Boolean(1) - if ab_producer_state.count < k_tile_cnt: + if ab_producer_state.count < tile_k_tile_cnt: peek_ab_empty_status = ab_pipeline.producer_try_acquire(ab_producer_state) # TMA load A/B @@ -1137,6 +1192,10 @@ def kernel( is_valid_tile = tile_info[3] == 1 while is_valid_tile: + k_tile_begin, tile_k_tile_cnt = self._split_k_range(k_tile_cnt, tile_info[2]) + l_coord = tile_info[2] + if cutlass.const_expr(self.split_k > 1): + l_coord = tile_info[2] // self.split_k # # Prepare the mask for scaleA/scaleB # @@ -1156,7 +1215,7 @@ def kernel( # Peek (try_wait) SCALE buffer empty scale_producer_state.reset_count() peek_scale_empty_status = cutlass.Boolean(1) - if scale_producer_state.count < k_tile_cnt: + if scale_producer_state.count < tile_k_tile_cnt: peek_scale_empty_status = scale_pipeline.producer_try_acquire( scale_producer_state ) @@ -1164,7 +1223,10 @@ def kernel( # # load loop # - for k_tile in cutlass.range(0, k_tile_cnt, 1, unroll=1): + for k_tile in cutlass.range(0, tile_k_tile_cnt, 1, unroll=1): + k_coord = scale_producer_state.count + if cutlass.const_expr(self.split_k > 1): + k_coord = k_tile_begin + scale_producer_state.count # # Slice to per mma tile index # @@ -1181,8 +1243,8 @@ def kernel( None, None, tile_info[0], - scale_producer_state.count, - tile_info[2], + k_coord, + l_coord, ) ] ) @@ -1193,8 +1255,8 @@ def kernel( None, None, tile_info[1], - scale_producer_state.count, - tile_info[2], + k_coord, + l_coord, ) ] ) @@ -1207,8 +1269,8 @@ def kernel( None, None, tile_info[0], - scale_producer_state.count, - tile_info[2], + k_coord, + l_coord, ), ) ) @@ -1220,8 +1282,8 @@ def kernel( None, None, tile_info[1], - scale_producer_state.count, - tile_info[2], + k_coord, + l_coord, ), ) ) @@ -1246,7 +1308,7 @@ def kernel( # Peek (try_wait) Scale buffer empty scale_producer_state.advance() peek_scale_empty_status = cutlass.Boolean(1) - if scale_producer_state.count < k_tile_cnt: + if scale_producer_state.count < tile_k_tile_cnt: peek_scale_empty_status = scale_pipeline.producer_try_acquire( scale_producer_state ) @@ -1320,22 +1382,23 @@ def kernel( is_valid_tile = tile_info[3] == 1 while is_valid_tile: + _, tile_k_tile_cnt = self._split_k_range(k_tile_cnt, tile_info[2]) # Peek (try_wait) AB buffer full for k_tile = 0 ab_consumer_state.reset_count() peek_ab_full_status = cutlass.Boolean(1) - if ab_consumer_state.count < k_tile_cnt and is_leader_cta: + if ab_consumer_state.count < tile_k_tile_cnt and is_leader_cta: peek_ab_full_status = ab_pipeline.consumer_try_wait(ab_consumer_state) # Peek (try_wait) Acc buffer empty for k_tile = 0 acc_producer_state.reset_count() peek_acc_empty_status = cutlass.Boolean(1) - if ab_consumer_state.count < k_tile_cnt and is_leader_cta: + if ab_consumer_state.count < tile_k_tile_cnt and is_leader_cta: peek_acc_empty_status = acc_pipeline.producer_try_acquire(acc_producer_state) # # Mma mainloop # - for k_tile in cutlass.range(0, k_tile_cnt, 1, unroll=1): + for k_tile in cutlass.range(0, tile_k_tile_cnt, 1, unroll=1): # Set tensor memory buffer for current tile # (MMA, MMA_M, MMA_N) tCtAcc = tCtAcc_base[(None, None, None, acc_producer_state.index)] @@ -1381,7 +1444,7 @@ def kernel( # Peek (try_wait) AB buffer full for k_tile = k_tile + 1 ab_consumer_state.advance() peek_ab_full_status = cutlass.Boolean(1) - if ab_consumer_state.count < k_tile_cnt: + if ab_consumer_state.count < tile_k_tile_cnt: if is_leader_cta: peek_ab_full_status = ab_pipeline.consumer_try_wait(ab_consumer_state) @@ -1393,7 +1456,7 @@ def kernel( # Peek (try_wait) Acc buffer empty for k_tile = k_tile + 1 acc_producer_state.advance() - if acc_producer_state.count < k_tile_cnt: + if acc_producer_state.count < tile_k_tile_cnt: if is_leader_cta: peek_acc_empty_status = acc_pipeline.producer_try_acquire( acc_producer_state @@ -1498,6 +1561,7 @@ def kernel( is_valid_tile = tile_info[3] == 1 while is_valid_tile: + _, tile_k_tile_cnt = self._split_k_range(k_tile_cnt, tile_info[2]) # initialize the final accumulator tTR_rAcc_final.fill(0.0) @@ -1512,15 +1576,15 @@ def kernel( scale_consumer_state.reset_count() peek_scale_full_status = cutlass.Boolean(1) - if scale_consumer_state.count < k_tile_cnt: + if scale_consumer_state.count < tile_k_tile_cnt: peek_scale_full_status = scale_pipeline.consumer_try_wait(scale_consumer_state) acc_consumer_state.reset_count() peek_acc_full_status = cutlass.Boolean(1) - if acc_consumer_state.count < k_tile_cnt: + if acc_consumer_state.count < tile_k_tile_cnt: peek_acc_full_status = acc_pipeline.consumer_try_wait(acc_consumer_state) - for k_tile in cutlass.range(0, k_tile_cnt, 1, unroll=1): + for k_tile in cutlass.range(0, tile_k_tile_cnt, 1, unroll=1): # Set tensor memory buffer for current tile # (T2R, T2R_M, T2R_N, EPI_M, EPI_M) tTR_tAcc = tTR_tAcc_base[ @@ -1591,7 +1655,7 @@ def kernel( scale_consumer_state.advance() peek_scale_full_status = cutlass.Boolean(1) - if scale_consumer_state.count < k_tile_cnt: + if scale_consumer_state.count < tile_k_tile_cnt: peek_scale_full_status = scale_pipeline.consumer_try_wait( scale_consumer_state ) @@ -1603,7 +1667,7 @@ def kernel( acc_consumer_state.advance() peek_acc_full_status = cutlass.Boolean(1) - if acc_consumer_state.count < k_tile_cnt: + if acc_consumer_state.count < tile_k_tile_cnt: peek_acc_full_status = acc_pipeline.consumer_try_wait(acc_consumer_state) tRT_tAcc = tRT_tAcc_base[(None, None, None, None, None, 0)] @@ -1738,6 +1802,12 @@ def kernel( tile_info[1], tile_info[2], ) + if cutlass.const_expr(self.split_k > 1): + mma_tile_coord_mnl = ( + mma_tile_coord_mnl[0], + mma_tile_coord_mnl[1], + mma_tile_coord_mnl[2] // self.split_k, + ) # # Slice to per mma tile index # @@ -1901,27 +1971,38 @@ def acc_update_tmem_copy_and_partition( :rtype: Tuple[cute.TiledCopy, cute.TiledCopy, cute.Tensor, cute.Tensor, cute.Tensor, cute.Tensor, cute.Tensor, cute.Tensor, cute.Tensor] """ - # Make tiledCopy for tensor memory load + # Make tiledCopy for tensor memory load/store. The warps split an + # epilogue subtile as in sm100_utils.get_tmem_load_op; an atom spans at + # most a warp's share of it (16 datapaths x 64 or 32 x 32 FP32 + # columns), so narrow tiles (N = 32) use fewer repetitions. + tmem_warp_shape_mn = ( + (2, 2) if (self.cta_tile_shape_mnk[0] == 64 and self.use_2cta_instrs) else (4, 1) + ) + epi_warp_tile_mn = cute.shape_div( + cute.product_each(cute.shape(epi_tile)), tmem_warp_shape_mn + ) + num_warp_cols = cute.size(epi_warp_tile_mn[1]) tmem_load_atom = None tmem_store_atom = None if cutlass.const_expr(self.mma_tiler[0] == 64): + # 256 bits = 8 FP32 columns per repetition + num_rep = min(8, num_warp_cols * self.acc_dtype.width // 256) tmem_load_atom = cute.make_copy_atom( - tcgen05.copy.Ld16x256bOp(tcgen05.copy.Repetition(8)), - self.acc_dtype, - ) - else: - tmem_load_atom = cute.make_copy_atom( - tcgen05.copy.Ld32x32bOp(tcgen05.copy.Repetition(32)), + tcgen05.copy.Ld16x256bOp(tcgen05.copy.Repetition(num_rep)), self.acc_dtype, ) - if cutlass.const_expr(self.mma_tiler[0] == 64): tmem_store_atom = cute.make_copy_atom( - tcgen05.copy.St16x256bOp(tcgen05.copy.Repetition(8)), + tcgen05.copy.St16x256bOp(tcgen05.copy.Repetition(num_rep)), self.acc_dtype, ) else: + num_rep = min(32, num_warp_cols) + tmem_load_atom = cute.make_copy_atom( + tcgen05.copy.Ld32x32bOp(tcgen05.copy.Repetition(num_rep)), + self.acc_dtype, + ) tmem_store_atom = cute.make_copy_atom( - tcgen05.copy.St32x32bOp(tcgen05.copy.Repetition(32)), + tcgen05.copy.St32x32bOp(tcgen05.copy.Repetition(num_rep)), self.acc_dtype, ) @@ -2230,6 +2311,9 @@ def _compute_grid( cta_tile_shape_mnk: Tuple[int, int, int], cluster_shape_mn: Tuple[int, int], max_active_clusters: cutlass.Constexpr, + swizzle_size: int = 1, + raster_along_m: bool = True, + split_k: int = 1, ) -> Tuple[utils.PersistentTileSchedulerParams, Tuple[int, int, int]]: """Use persistent tile scheduler to compute the grid size for the output tensor C. @@ -2241,6 +2325,13 @@ def _compute_grid( :type cluster_shape_mn: tuple[int, int] :param max_active_clusters: Maximum number of active clusters. :type max_active_clusters: cutlass.Constexpr + :param swizzle_size: Cluster swizzle width of the persistent tile scheduler. + :type swizzle_size: int + :param raster_along_m: Whether consecutive clusters advance along M. + :type raster_along_m: bool + :param split_k: K slices per output tile; the scheduler's L enumerates + (batch, slice) pairs. + :type split_k: int :return: A tuple containing: - tile_sched_params: Parameters for the persistent tile scheduler. @@ -2250,9 +2341,17 @@ def _compute_grid( c_shape = cute.slice_(cta_tile_shape_mnk, (None, None, 0)) gc = cute.zipped_divide(c, tiler=c_shape) num_ctas_mnl = gc[(0, (None, None, None))].shape + if split_k > 1: + num_ctas_mnl = ( + num_ctas_mnl[0], + num_ctas_mnl[1], + num_ctas_mnl[2] * split_k, + ) cluster_shape_mnl = (*cluster_shape_mn, 1) - tile_sched_params = utils.PersistentTileSchedulerParams(num_ctas_mnl, cluster_shape_mnl) + tile_sched_params = utils.PersistentTileSchedulerParams( + num_ctas_mnl, cluster_shape_mnl, swizzle_size, raster_along_m + ) grid = utils.StaticPersistentTileScheduler.get_grid_shape( tile_sched_params, max_active_clusters ) @@ -2344,8 +2443,9 @@ def is_valid_mma_tiler_and_cluster_shape( or (use_2cta_instrs and mma_tiler_mn[0] in [128, 256]) ): is_valid = False - # Skip invalid mma tile n - if mma_tiler_mn[1] not in (128,): + # Skip invalid mma tile n: a tile covers one 128-column scale block of B + # or lies inside one. + if mma_tiler_mn[1] not in (32, 64, 128): is_valid = False # Skip illegal cluster shape if cluster_shape_mn[0] % (2 if use_2cta_instrs else 1) != 0: diff --git a/tensorrt_llm/_torch/kernels/blackwell/dense_blockscaled_gemm_persistent.py b/tensorrt_llm/_torch/kernels/blackwell/dense_blockscaled_gemm_persistent.py index 793bab1c2727..bc67b5304344 100644 --- a/tensorrt_llm/_torch/kernels/blackwell/dense_blockscaled_gemm_persistent.py +++ b/tensorrt_llm/_torch/kernels/blackwell/dense_blockscaled_gemm_persistent.py @@ -137,6 +137,7 @@ def __init__( raster_along_m: bool = True, indexer_q_fusion: bool = False, indexer_transform_warps: int = 4, + split_k: int = 1, ): """Initializes the configuration for a Blackwell dense GEMM kernel. @@ -153,6 +154,7 @@ def __init__( 3. Persistent tile scheduler: - swizzle_size: cluster-swizzle size for L2-locality (1 = no swizzle). - raster_along_m: cluster rasterization order (M-major when True). + - split_k: number of K slices per output tile. Args: sf_vec_size (int): Scalefactor vector size. @@ -160,6 +162,9 @@ def __init__( cluster_shape_mn (Tuple[int, int]): Cluster dimensions (M, N) for parallel processing. swizzle_size (int): Tile-scheduler swizzle size in clusters (default 1). raster_along_m (bool): Tile-scheduler rasterization order (default M-major). + split_k (int): Number of K slices per output tile (default 1). Each + slice reduce-adds its partial into C through TMA, so callers + must zero C before every launch. """ self.acc_dtype = cutlass.Float32 @@ -173,6 +178,11 @@ def __init__( self.raster_along_m = raster_along_m self.indexer_q_fusion = indexer_q_fusion self.indexer_transform_warps = indexer_transform_warps + if split_k < 1: + raise ValueError(f"split_k must be >= 1, got {split_k}") + if indexer_q_fusion and split_k > 1: + raise ValueError("The indexer-Q epilogue does not support split-K") + self.split_k = split_k if self.indexer_q_fusion and (sf_vec_size != 32 or mma_tiler_mn[0] != 128 or mma_tiler_mn[1] != 16 or cluster_shape_mn != (1, 1) @@ -218,6 +228,20 @@ def __init__( SM100_TMEM_CAPACITY_COLUMNS = 512 self.num_tmem_alloc_cols = SM100_TMEM_CAPACITY_COLUMNS + def _split_k_range(self, k_block_cnt, l_coord): + """Return ``(k_block_begin, k_blocks)`` of the split-K slice at scheduler L ``l_coord``. + + Slices take balanced ranges, ``[floor(k_block_cnt * s / split_k), + floor(k_block_cnt * (s + 1) / split_k))`` for slice ``s = l_coord % split_k``, + so any ``split_k <= k_block_cnt`` gives every slice at least one K block. + """ + if cutlass.const_expr(self.split_k == 1): + return 0, k_block_cnt + split_idx = l_coord % self.split_k + k_block_begin = (k_block_cnt * split_idx) // self.split_k + k_block_end = (k_block_cnt * (split_idx + 1)) // self.split_k + return k_block_begin, k_block_end - k_block_begin + def _setup_attributes(self): """Set up configurations that are dependent on GEMM inputs @@ -578,8 +602,11 @@ def __call__( # Setup TMA store for C epi_smem_layout = cute.slice_(self.c_smem_layout_staged, (None, None, 0)) + c_store_op = cpasync.CopyBulkTensorTileS2GOp() + if cutlass.const_expr(self.split_k > 1): + c_store_op = cpasync.CopyReduceBulkTensorTileS2GOp() tma_atom_c, tma_tensor_c = cpasync.make_tiled_tma_atom( - cpasync.CopyBulkTensorTileS2GOp(), + c_store_op, c_tensor, epi_smem_layout, self.epi_tile, @@ -593,6 +620,7 @@ def __call__( max_active_clusters, self.swizzle_size, self.raster_along_m, + self.split_k, ) self.buffer_align_bytes = 1024 @@ -1004,6 +1032,15 @@ def kernel( cur_tile_coord[1], cur_tile_coord[2], ) + # The scheduler's L enumerates (batch, split-K slice) pairs. + k_block_begin, tile_k_block_cnt = self._split_k_range( + k_block_cnt, mma_tile_coord_mnl[2]) + if cutlass.const_expr(self.split_k > 1): + mma_tile_coord_mnl = ( + mma_tile_coord_mnl[0], + mma_tile_coord_mnl[1], + mma_tile_coord_mnl[2] // self.split_k, + ) # # Slice to per mma tile index @@ -1029,43 +1066,50 @@ def kernel( # Prefetch both A and B (default behavior) for pf_k_block in cutlass.range(0, min(self.prefetch_dist, - k_block_cnt), + tile_k_block_cnt), unroll=1): + pf_k_coord = pf_k_block + if cutlass.const_expr(self.split_k > 1): + pf_k_coord = k_block_begin + pf_k_block cute.prefetch( tma_atom_a, - tAgA_slice[(None, pf_k_block)], + tAgA_slice[(None, pf_k_coord)], ) cute.prefetch( tma_atom_b, - tBgB_slice[(None, pf_k_block)], + tBgB_slice[(None, pf_k_coord)], ) cute.prefetch( tma_atom_sfa, - tAgSFA_slice[(None, pf_k_block)], + tAgSFA_slice[(None, pf_k_coord)], ) cute.prefetch( tma_atom_sfb, - tBgSFB_slice[(None, pf_k_block)], + tBgSFB_slice[(None, pf_k_coord)], ) # Peek (try_wait) AB buffer empty for k_block = prefetch_k_block_cnt ab_producer_state.reset_count() peek_ab_empty_status = cutlass.Boolean(1) - if ab_producer_state.count < k_block_cnt: + if ab_producer_state.count < tile_k_block_cnt: peek_ab_empty_status = ab_pipeline.producer_try_acquire( ab_producer_state) # # Tma load loop # - for k_block in cutlass.range(0, k_block_cnt, 1, unroll=1): + for k_block in cutlass.range(0, tile_k_block_cnt, 1, unroll=1): # Conditionally wait for AB buffer empty ab_pipeline.producer_acquire(ab_producer_state, peek_ab_empty_status) + k_coord = ab_producer_state.count + if cutlass.const_expr(self.split_k > 1): + k_coord = k_block_begin + ab_producer_state.count + # TMA load A/B/SFA/SFB cute.copy( tma_atom_a, - tAgA_slice[(None, ab_producer_state.count)], + tAgA_slice[(None, k_coord)], tAsA[(None, ab_producer_state.index)], tma_bar_ptr=ab_pipeline.producer_get_barrier( ab_producer_state), @@ -1073,7 +1117,7 @@ def kernel( ) cute.copy( tma_atom_b, - tBgB_slice[(None, ab_producer_state.count)], + tBgB_slice[(None, k_coord)], tBsB[(None, ab_producer_state.index)], tma_bar_ptr=ab_pipeline.producer_get_barrier( ab_producer_state), @@ -1081,7 +1125,7 @@ def kernel( ) cute.copy( tma_atom_sfa, - tAgSFA_slice[(None, ab_producer_state.count)], + tAgSFA_slice[(None, k_coord)], tAsSFA[(None, ab_producer_state.index)], tma_bar_ptr=ab_pipeline.producer_get_barrier( ab_producer_state), @@ -1089,7 +1133,7 @@ def kernel( ) cute.copy( tma_atom_sfb, - tBgSFB_slice[(None, ab_producer_state.count)], + tBgSFB_slice[(None, k_coord)], tBsSFB[(None, ab_producer_state.index)], tma_bar_ptr=ab_pipeline.producer_get_barrier( ab_producer_state), @@ -1098,8 +1142,8 @@ def kernel( # Prefetch: Rolling prefetch for next tiles if cutlass.const_expr(self.use_prefetch): - if k_block < k_block_cnt - self.prefetch_dist: - future_k_block = ab_producer_state.count + self.prefetch_dist + if k_block < tile_k_block_cnt - self.prefetch_dist: + future_k_block = k_coord + self.prefetch_dist # Prefetch both A and B (default behavior) cute.prefetch( tma_atom_a, @@ -1121,7 +1165,7 @@ def kernel( # Peek (try_wait) AB buffer empty for k_block = prefetch_k_block_cnt + k_block + 1 ab_producer_state.advance() peek_ab_empty_status = cutlass.Boolean(1) - if ab_producer_state.count < k_block_cnt: + if ab_producer_state.count < tile_k_block_cnt: peek_ab_empty_status = ab_pipeline.producer_try_acquire( ab_producer_state) @@ -1218,6 +1262,8 @@ def kernel( cur_tile_coord[1], cur_tile_coord[2], ) + _, tile_k_block_cnt = self._split_k_range( + k_block_cnt, cur_tile_coord[2]) if cutlass.const_expr(self.overlapping_accum): acc_stage_index = acc_producer_state.phase ^ 1 @@ -1231,7 +1277,7 @@ def kernel( # Peek (try_wait) AB buffer full for k_block = 0 ab_consumer_state.reset_count() peek_ab_full_status = cutlass.Boolean(1) - if ab_consumer_state.count < k_block_cnt and is_leader_cta: + if ab_consumer_state.count < tile_k_block_cnt and is_leader_cta: peek_ab_full_status = ab_pipeline.consumer_try_wait( ab_consumer_state) @@ -1271,7 +1317,7 @@ def kernel( # # Mma mainloop # - for k_block in range(k_block_cnt): + for k_block in range(tile_k_block_cnt): if is_leader_cta: # Conditionally wait for AB buffer full ab_pipeline.consumer_wait(ab_consumer_state, @@ -1339,7 +1385,7 @@ def kernel( # Peek (try_wait) AB buffer full for k_block = k_block + 1 ab_consumer_state.advance() peek_ab_full_status = cutlass.Boolean(1) - if ab_consumer_state.count < k_block_cnt: + if ab_consumer_state.count < tile_k_block_cnt: if is_leader_cta: peek_ab_full_status = ab_pipeline.consumer_try_wait( ab_consumer_state) @@ -1440,6 +1486,12 @@ def kernel( cur_tile_coord[1], cur_tile_coord[2], ) + if cutlass.const_expr(self.split_k > 1): + mma_tile_coord_mnl = ( + mma_tile_coord_mnl[0], + mma_tile_coord_mnl[1], + mma_tile_coord_mnl[2] // self.split_k, + ) # # Slice to per mma tile index @@ -2069,6 +2121,7 @@ def _compute_grid( max_active_clusters: cutlass.Constexpr, swizzle_size: cutlass.Constexpr = 1, raster_along_m: cutlass.Constexpr = True, + split_k: cutlass.Constexpr = 1, ) -> Tuple[utils.PersistentTileSchedulerParams, Tuple[int, int, int]]: """Computes the grid size for the kernel launch. @@ -2086,6 +2139,8 @@ def _compute_grid( clusters (1 = no swizzle). raster_along_m (cutlass.Constexpr): Tile-scheduler rasterization order (M-major when True). + split_k (cutlass.Constexpr): K slices per output tile; the + scheduler's L enumerates (batch, slice) pairs. Returns: A tuple containing the tile scheduler parameters and the computed @@ -2094,6 +2149,12 @@ def _compute_grid( c_shape = cute.slice_(cta_tile_shape_mnk, (None, None, 0)) gc = cute.zipped_divide(c, tiler=c_shape) num_ctas_mnl = gc[(0, (None, None, None))].shape + if split_k > 1: + num_ctas_mnl = ( + num_ctas_mnl[0], + num_ctas_mnl[1], + num_ctas_mnl[2] * split_k, + ) cluster_shape_mnl = (*cluster_shape_mn, 1) tile_sched_params = utils.PersistentTileSchedulerParams( diff --git a/tensorrt_llm/_torch/kernels/rubin/blockwise_gemm/blockwise_gemm.py b/tensorrt_llm/_torch/kernels/rubin/blockwise_gemm/blockwise_gemm.py index 638cd7e023f4..200d69d7eb51 100644 --- a/tensorrt_llm/_torch/kernels/rubin/blockwise_gemm/blockwise_gemm.py +++ b/tensorrt_llm/_torch/kernels/rubin/blockwise_gemm/blockwise_gemm.py @@ -278,6 +278,8 @@ def __init__( mma_inst_shape: Tuple[int, int, int], cluster_shape_mn: Tuple[int, int], fp8_quantize_1x128: bool = False, + swizzle_size: int = 1, + raster_along_m: bool = True, ): """Initialize the Rubin persistent blockwise dense GEMM kernel. @@ -291,6 +293,10 @@ def __init__( :type mma_inst_shape: Tuple[int, int, int] :param cluster_shape_mn: Tuple (ClusterM, ClusterN) shape of the cluster. :type cluster_shape_mn: Tuple[int, int] + :param swizzle_size: Cluster swizzle width of the persistent tile scheduler. + :type swizzle_size: int + :param raster_along_m: Whether consecutive clusters advance along M. + :type raster_along_m: bool """ # Call parent constructor with mma_tiler_mn (M, N only) super().__init__( @@ -298,12 +304,26 @@ def __init__( use_2cta_instrs, (mma_tiler[0], mma_tiler[1]), cluster_shape_mn, + swizzle_size, + raster_along_m, ) # Match the register split used by the Rubin source kernel's Blackwell # base. TRT-LLM's vendored Blackwell kernel uses a balanced 216/216 # split, while Rubin dedicates more registers to accumulator updates. self.num_regs_epilogue_warps = 168 self.num_regs_acc_update_warps = 256 + # 256-row CTA tiles split the accumulator update over warps 0-7 (see + # _split_acc_update). Both warp groups grow from the 168 launch + # registers into what warps 8-11 release: 2 * (216 - 168) <= 168 - 64. + self.num_regs_split_acc_update_warps = 216 + # Named barriers of the split update: 4 + h syncs the stores of MMA_M + # half h, 6 joins both halves before the TMEM is freed. + self.split_store_barrier_id = 4 + self.split_acc_update_done_barrier = pipeline.NamedBarrier( + barrier_id=6, + num_threads=self.threads_per_warp + * len((*self.acc_update_warp_id, *self.epilog_warp_id)), + ) self.arch = "sm_107" self.num_smem_capacity = utils.get_smem_capacity_in_bytes(self.arch) self.mma_tiler = mma_tiler @@ -410,6 +430,13 @@ def _setup_attributes(self): ) if self.fp8_quantize_1x128 and self.cta_tile_shape_mnk[:2] != (128, 128): raise ValueError("Fused FP8 quantization requires a 128x128 CTA output tile") + # Four accumulator-update warps would keep two rows (256 FP32 values per + # thread) of a 256-row tile's final accumulator in registers and spill; + # such tiles split the update by MMA_M half over warps 0-7 instead. + self.split_acc_update = self.cta_tile_shape_mnk[0] == 256 + self.num_acc_update_warps = len(self.acc_update_warp_id) * ( + 2 if self.split_acc_update else 1 + ) # Compute cluster layout self.cluster_layout_vmnk = cute.tiled_divide( @@ -446,6 +473,12 @@ def _setup_attributes(self): self.c_dtype, ) + # The split update stores each MMA_M half through C stages of its own, + # so plan the C stages in pairs of epilogue subtiles. + stage_epi_tile = self.epi_tile + if self.split_acc_update: + stage_epi_tile = (cute.make_layout(self.cta_tile_shape_mnk[0]), self.epi_tile[1]) + # Setup A/B/C/Scale stage count in shared memory and ACC stage count in tensor memory ( self.num_acc_stage, @@ -458,7 +491,7 @@ def _setup_attributes(self): self.mma_tiler, self.a_dtype, self.b_dtype, - self.epi_tile, + stage_epi_tile, self.c_dtype, self.c_layout, self.sfa_dtype, @@ -468,16 +501,19 @@ def _setup_attributes(self): self.num_smem_capacity, self.occupancy, ) + if self.split_acc_update: + self.num_c_stage *= 2 if self.fp8_quantize_1x128: # The direct fused epilogue reuses this allocation as one # row-major 128x128 FP8 scratch tile (16 KiB). Two 128x32 BF16 # epilogue stages provide exactly that capacity. self.num_c_stage = 2 - # SM107 TMEM (576 cols) can not fit multiple staging stages when - # MMA_M=2 (cta_tile_M==256), so only keep one staging + one final. - if self.cta_tile_shape_mnk[0] == 256: - self.num_acc_stage = 1 + # SM107 TMEM (576 cols) fits two 256-column stages of a 256-row + # (MMA_M=2) accumulator; the split update keeps no final accumulator + # in TMEM. + if self.split_acc_update: + self.num_acc_stage = 2 elif self.cta_tile_shape_mnk[0] == 128: self.num_acc_stage = 3 else: @@ -674,7 +710,12 @@ def __call__( # Compute grid size self.tile_sched_params, grid = self._compute_grid( - c, self.cta_tile_shape_mnk, self.cluster_shape_mn, max_active_clusters + c, + self.cta_tile_shape_mnk, + self.cluster_shape_mn, + max_active_clusters, + self.swizzle_size, + self.raster_along_m, ) self.buffer_align_bytes = 1024 @@ -855,7 +896,7 @@ def kernel( ) scale_pipeline_consumer_group = pipeline.CooperativeGroup( pipeline.Agent.Thread, - self.threads_per_warp * len(self.epilog_warp_id), + self.threads_per_warp * self.num_acc_update_warps, ) scale_pipeline = pipeline.PipelineCpAsync.create( barrier_storage=storage.scale_mbar_ptr.data_ptr(), @@ -867,7 +908,8 @@ def kernel( # Initialize acc_pipeline (barrier) and states acc_pipeline_producer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread) - num_acc_consumer_threads = len(self.epilog_warp_id) * (2 if use_2cta_instrs else 1) + # One elected thread per accumulator-update warp (of both CTAs for 2CTA) + num_acc_consumer_threads = self.num_acc_update_warps * (2 if use_2cta_instrs else 1) acc_pipeline_consumer_group = pipeline.CooperativeGroup( pipeline.Agent.Thread, num_acc_consumer_threads ) @@ -1622,468 +1664,778 @@ def kernel( # acc_pipeline.producer_tail(acc_producer_state) - # - # Specialized acc update warps - # - if warp_idx <= self.acc_update_warp_id[-1]: - cute.arch.setmaxregister_increase(self.num_regs_acc_update_warps) + if cutlass.const_expr(self.split_acc_update): # - # Bar sync for retrieve tensor memory ptr from shared memory - # - tmem.wait_for_alloc() - + # Specialized split acc update warps (256-row CTA tiles) # - # Retrieving tensor memory ptr and make accumulator tensor - # - tmem_ptr = tmem.retrieve_ptr(self.acc_dtype) - # (MMA, MMA_M, MMA_N, STAGE) - tCtAcc_base = cute.make_tensor(tmem_ptr, tCtAcc_fake.layout) - # Use dynamic cosize of tCtAcc_base to place the final accumulator - # right after the staging region, so it does not overlap when MMA_M>1. - tCtAcc_final = cute.make_tensor( - tCtAcc_base.iterator + tcgen05.find_tmem_tensor_col_offset(tCtAcc_base), - tCtAcc_base.layout, - ) - - # - # Partition for epilogue - # - epi_tidx = tidx % 128 - ( - tiled_copy_t2r, - tiled_copy_r2t, - tTR_tAcc_base, - tTR_rAcc, - tTR_rAcc_final, - tTR_sSFA, - tTR_sSFB, - tRT_rAcc, - tRT_tAcc_base, - ) = self.acc_update_tmem_copy_and_partition( - epi_tidx, - tCtAcc_base, - tCtAcc_final, - tCgC, - sSFA_view_as_C, - sSFB_view_as_C, - epi_tile, - ) - + if warp_idx <= self.epilog_warp_id[-1]: + self._split_acc_update( + warp_idx, + tidx, + tiled_mma, + tmem, + tCtAcc_fake, + tCgC, + sC, + sSFA_view_as_C, + sSFB_view_as_C, + sInfo, + tma_atom_c, + acc_pipeline, + scale_pipeline, + tile_info_pipeline, + epi_tile, + tile_sched_params, + k_tile_cnt, + epilogue_op, + ) + else: # - # Persistent tile scheduling loop + # Specialized acc update warps # - tile_sched = utils.StaticPersistentTileScheduler.create( - tile_sched_params, cute.arch.block_idx(), cute.arch.grid_dim() - ) - work_tile = tile_sched.initial_work_tile_info() - - acc_consumer_state = pipeline.make_pipeline_state( - pipeline.PipelineUserType.Consumer, self.num_acc_stage - ) - - scale_consumer_state = pipeline.make_pipeline_state( - pipeline.PipelineUserType.Consumer, self.num_scale_stage - ) + if warp_idx <= self.acc_update_warp_id[-1]: + cute.arch.setmaxregister_increase(self.num_regs_acc_update_warps) + # + # Bar sync for retrieve tensor memory ptr from shared memory + # + tmem.wait_for_alloc() - epi_producer_state = pipeline.make_pipeline_state(pipeline.PipelineUserType.Producer, 1) + # + # Retrieving tensor memory ptr and make accumulator tensor + # + tmem_ptr = tmem.retrieve_ptr(self.acc_dtype) + # (MMA, MMA_M, MMA_N, STAGE) + tCtAcc_base = cute.make_tensor(tmem_ptr, tCtAcc_fake.layout) + # Use dynamic cosize of tCtAcc_base to place the final accumulator + # right after the staging region, so it does not overlap when MMA_M>1. + tCtAcc_final = cute.make_tensor( + tCtAcc_base.iterator + tcgen05.find_tmem_tensor_col_offset(tCtAcc_base), + tCtAcc_base.layout, + ) - tile_info_consumer_state = pipeline.make_pipeline_state( - pipeline.PipelineUserType.Consumer, self.num_tile_stage - ) + # + # Partition for epilogue + # + epi_tidx = tidx % 128 + ( + tiled_copy_t2r, + tiled_copy_r2t, + tTR_tAcc_base, + tTR_rAcc, + tTR_rAcc_final, + tTR_sSFA, + tTR_sSFB, + tRT_rAcc, + tRT_tAcc_base, + ) = self.acc_update_tmem_copy_and_partition( + epi_tidx, + tCtAcc_base, + tCtAcc_final, + tCgC, + sSFA_view_as_C, + sSFB_view_as_C, + epi_tile, + ) - # get the first tile info - tile_info = cute.make_rmem_tensor((4,), cutlass.Int32) - # Get tile coord from tile scheduler - cur_tile_coord = work_tile.tile_idx - # initialize the tile info - tile_info[0] = cur_tile_coord[0] - tile_info[1] = cur_tile_coord[1] - tile_info[2] = cur_tile_coord[2] - tile_info[3] = work_tile.is_valid_tile + # + # Persistent tile scheduling loop + # + tile_sched = utils.StaticPersistentTileScheduler.create( + tile_sched_params, cute.arch.block_idx(), cute.arch.grid_dim() + ) + work_tile = tile_sched.initial_work_tile_info() - is_valid_tile = cutlass.Boolean(1) - is_valid_tile = tile_info[3] == 1 + acc_consumer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Consumer, self.num_acc_stage + ) - while is_valid_tile: - # initialize the final accumulator - tTR_rAcc_final.fill(0.0) - - # Keep both EPI_M and EPI_N so per-subtile_idx scale is correct - # when cta_tile_M > EPI_TILE_M (i.e. EPI_M > 1). - tTR_rSFA_ = cute.make_rmem_tensor( - cute.slice_(tTR_sSFA, (None, None, None, None, None, 0)).shape, - self.acc_dtype, + scale_consumer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Consumer, self.num_scale_stage ) - tTR_rSFB_ = cute.make_rmem_tensor( - cute.slice_(tTR_sSFB, (None, None, None, None, None, 0)).shape, - self.acc_dtype, + + epi_producer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Producer, 1 ) - # Group (EPI_M, EPI_N) into one mode so tTR_rSFA[.., subtile_idx] - # matches the grouped tTR_tAcc indexing below. - tTR_rSFA = cute.group_modes(tTR_rSFA_, 3, cute.rank(tTR_rSFA_)) - tTR_rSFB = cute.group_modes(tTR_rSFB_, 3, cute.rank(tTR_rSFB_)) - scale_consumer_state.reset_count() - peek_scale_full_status = cutlass.Boolean(1) - if scale_consumer_state.count < k_tile_cnt: - peek_scale_full_status = scale_pipeline.consumer_try_wait(scale_consumer_state) + tile_info_consumer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Consumer, self.num_tile_stage + ) - acc_consumer_state.reset_count() - peek_acc_full_status = cutlass.Boolean(1) - if acc_consumer_state.count < k_tile_cnt: - peek_acc_full_status = acc_pipeline.consumer_try_wait(acc_consumer_state) + # get the first tile info + tile_info = cute.make_rmem_tensor((4,), cutlass.Int32) + # Get tile coord from tile scheduler + cur_tile_coord = work_tile.tile_idx + # initialize the tile info + tile_info[0] = cur_tile_coord[0] + tile_info[1] = cur_tile_coord[1] + tile_info[2] = cur_tile_coord[2] + tile_info[3] = work_tile.is_valid_tile - for k_tile in cutlass.range(0, k_tile_cnt, 1, unroll=1): - # Set tensor memory buffer for current tile - # (T2R, T2R_M, T2R_N, EPI_M, EPI_M) - tTR_tAcc = tTR_tAcc_base[ - (None, None, None, None, None, acc_consumer_state.index) - ] + is_valid_tile = cutlass.Boolean(1) + is_valid_tile = tile_info[3] == 1 - # - # Wait for scale buffer full - # - scale_pipeline.consumer_wait(scale_consumer_state, peek_scale_full_status) + while is_valid_tile: + # initialize the final accumulator + tTR_rAcc_final.fill(0.0) - tTR_sSFA_slice = cute.slice_( - tTR_sSFA, - (None, None, None, None, None, scale_consumer_state.index), - ) - tTR_sSFB_slice = cute.slice_( - tTR_sSFB, - (None, None, None, None, None, scale_consumer_state.index), + # Keep both EPI_M and EPI_N so per-subtile_idx scale is correct + # when cta_tile_M > EPI_TILE_M (i.e. EPI_M > 1). + tTR_rSFA_ = cute.make_rmem_tensor( + cute.slice_(tTR_sSFA, (None, None, None, None, None, 0)).shape, + self.acc_dtype, ) - - scale_atom_copy = cute.make_copy_atom( - cute.nvgpu.CopyUniversalOp(), + tTR_rSFB_ = cute.make_rmem_tensor( + cute.slice_(tTR_sSFB, (None, None, None, None, None, 0)).shape, self.acc_dtype, - num_bits_per_copy=self.acc_dtype.width, ) + # Group (EPI_M, EPI_N) into one mode so tTR_rSFA[.., subtile_idx] + # matches the grouped tTR_tAcc indexing below. + tTR_rSFA = cute.group_modes(tTR_rSFA_, 3, cute.rank(tTR_rSFA_)) + tTR_rSFB = cute.group_modes(tTR_rSFB_, 3, cute.rank(tTR_rSFB_)) + + scale_consumer_state.reset_count() + peek_scale_full_status = cutlass.Boolean(1) + if scale_consumer_state.count < k_tile_cnt: + peek_scale_full_status = scale_pipeline.consumer_try_wait( + scale_consumer_state + ) - # Copy into the un-grouped rmem tensors (shape matches slice); - # tTR_rSFA / tTR_rSFB share storage via group_modes aliasing. - cute.copy(scale_atom_copy, tTR_sSFA_slice, tTR_rSFA_) - cute.copy(scale_atom_copy, tTR_sSFB_slice, tTR_rSFB_) + acc_consumer_state.reset_count() + peek_acc_full_status = cutlass.Boolean(1) + if acc_consumer_state.count < k_tile_cnt: + peek_acc_full_status = acc_pipeline.consumer_try_wait(acc_consumer_state) - # - # Wait for accumulator buffer full - # + for k_tile in cutlass.range(0, k_tile_cnt, 1, unroll=1): + # Set tensor memory buffer for current tile + # (T2R, T2R_M, T2R_N, EPI_M, EPI_M) + tTR_tAcc = tTR_tAcc_base[ + (None, None, None, None, None, acc_consumer_state.index) + ] - acc_pipeline.consumer_wait(acc_consumer_state, peek_acc_full_status) + # + # Wait for scale buffer full + # + scale_pipeline.consumer_wait(scale_consumer_state, peek_scale_full_status) - tTR_tAcc = cute.group_modes(tTR_tAcc, 3, cute.rank(tTR_tAcc)) + tTR_sSFA_slice = cute.slice_( + tTR_sSFA, + (None, None, None, None, None, scale_consumer_state.index), + ) + tTR_sSFB_slice = cute.slice_( + tTR_sSFB, + (None, None, None, None, None, scale_consumer_state.index), + ) + + scale_atom_copy = cute.make_copy_atom( + cute.nvgpu.CopyUniversalOp(), + self.acc_dtype, + num_bits_per_copy=self.acc_dtype.width, + ) + + # Copy into the un-grouped rmem tensors (shape matches slice); + # tTR_rSFA / tTR_rSFB share storage via group_modes aliasing. + cute.copy(scale_atom_copy, tTR_sSFA_slice, tTR_rSFA_) + cute.copy(scale_atom_copy, tTR_sSFB_slice, tTR_rSFB_) - # - # Update accumulator by scale factor in subtiles - # - subtile_cnt = cute.size(tTR_tAcc.shape, mode=[3]) - for subtile_idx in cutlass.range(subtile_cnt): # - # Load accumulator from tensor memory buffer to register + # Wait for accumulator buffer full # - tTR_tAcc_mn = tTR_tAcc[(None, None, None, subtile_idx)] - cute.copy(tiled_copy_t2r, tTR_tAcc_mn, tTR_rAcc) + + acc_pipeline.consumer_wait(acc_consumer_state, peek_acc_full_status) + + tTR_tAcc = cute.group_modes(tTR_tAcc, 3, cute.rank(tTR_tAcc)) # - # Update accumulator by scale factor + # Update accumulator by scale factor in subtiles # - tTR_rAcc_subtile = tTR_rAcc_final[(None, None, None, subtile_idx)] - tTR_rSFA_subtile = tTR_rSFA[(None, None, None, subtile_idx)] - tTR_rSFB_subtile = tTR_rSFB[(None, None, None, subtile_idx)] - - acc_vec = tTR_rAcc.load() - final_vec = tTR_rAcc_subtile.load() - scale_a = tTR_rSFA_subtile.load() - scale_b = tTR_rSFB_subtile.load() - scale = scale_a * scale_b - final_vec = acc_vec * scale + final_vec - tTR_rAcc_subtile.store(final_vec.to(self.acc_dtype)) + subtile_cnt = cute.size(tTR_tAcc.shape, mode=[3]) + for subtile_idx in cutlass.range(subtile_cnt): + # + # Load accumulator from tensor memory buffer to register + # + tTR_tAcc_mn = tTR_tAcc[(None, None, None, subtile_idx)] + cute.copy(tiled_copy_t2r, tTR_tAcc_mn, tTR_rAcc) + + # + # Update accumulator by scale factor + # + tTR_rAcc_subtile = tTR_rAcc_final[(None, None, None, subtile_idx)] + tTR_rSFA_subtile = tTR_rSFA[(None, None, None, subtile_idx)] + tTR_rSFB_subtile = tTR_rSFB[(None, None, None, subtile_idx)] + + acc_vec = tTR_rAcc.load() + final_vec = tTR_rAcc_subtile.load() + scale_a = tTR_rSFA_subtile.load() + scale_b = tTR_rSFB_subtile.load() + scale = scale_a * scale_b + final_vec = acc_vec * scale + final_vec + tTR_rAcc_subtile.store(final_vec.to(self.acc_dtype)) + + # + # Async arrive scale buffer empty + # + scale_pipeline.consumer_release(scale_consumer_state) + scale_consumer_state.advance() + + peek_scale_full_status = cutlass.Boolean(1) + if scale_consumer_state.count < k_tile_cnt: + peek_scale_full_status = scale_pipeline.consumer_try_wait( + scale_consumer_state + ) + # + # Async arrive accumulator buffer empty + # + with cute.arch.elect_one(): + acc_pipeline.consumer_release(acc_consumer_state) + acc_consumer_state.advance() + + peek_acc_full_status = cutlass.Boolean(1) + if acc_consumer_state.count < k_tile_cnt: + peek_acc_full_status = acc_pipeline.consumer_try_wait( + acc_consumer_state + ) + + tRT_tAcc = tRT_tAcc_base[(None, None, None, None, None, 0)] + tRT_tAcc = cute.group_modes(tRT_tAcc, 3, cute.rank(tRT_tAcc)) # - # Async arrive scale buffer empty + # Wait for epilogue buffer empty # - scale_pipeline.consumer_release(scale_consumer_state) - scale_consumer_state.advance() + epi_pipeline.producer_acquire(epi_producer_state) + + # copy the accumulator to tensor memory buffer + cute.copy(tiled_copy_r2t, tTR_rAcc_final, tRT_tAcc) + cute.arch.fence_view_async_tmem_store() - peek_scale_full_status = cutlass.Boolean(1) - if scale_consumer_state.count < k_tile_cnt: - peek_scale_full_status = scale_pipeline.consumer_try_wait( - scale_consumer_state - ) # - # Async arrive accumulator buffer empty + # Async arrive epilogue buffer full # - with cute.arch.elect_one(): - acc_pipeline.consumer_release(acc_consumer_state) - acc_consumer_state.advance() + epi_pipeline.producer_commit(epi_producer_state) + epi_producer_state.advance() - peek_acc_full_status = cutlass.Boolean(1) - if acc_consumer_state.count < k_tile_cnt: - peek_acc_full_status = acc_pipeline.consumer_try_wait(acc_consumer_state) + # + # Advance to next tile + # + tile_info_pipeline.consumer_wait(tile_info_consumer_state) + for idx in cutlass.range(4, unroll_full=True): + tile_info[idx] = sInfo[(idx, tile_info_consumer_state.index)] + is_valid_tile = tile_info[3] == 1 + cute.arch.fence_proxy( + "async.shared", + space="cta", + ) + tile_info_pipeline.consumer_release(tile_info_consumer_state) + tile_info_consumer_state.advance() - tRT_tAcc = tRT_tAcc_base[(None, None, None, None, None, 0)] - tRT_tAcc = cute.group_modes(tRT_tAcc, 3, cute.rank(tRT_tAcc)) + # + # Specialized epilogue warps + # + if warp_idx <= self.epilog_warp_id[-1] and warp_idx >= self.epilog_warp_id[0]: + cute.arch.setmaxregister_increase(self.num_regs_epilogue_warps) + # + # Alloc tensor memory buffer + # + tmem.allocate(self.num_tmem_alloc_cols) # - # Wait for epilogue buffer empty + # Bar sync for retrieve tensor memory ptr from shared memory # - epi_pipeline.producer_acquire(epi_producer_state) + tmem.wait_for_alloc() - # copy the accumulator to tensor memory buffer - cute.copy(tiled_copy_r2t, tTR_rAcc_final, tRT_tAcc) - cute.arch.fence_view_async_tmem_store() + # + # Retrieving tensor memory ptr and make accumulator tensor + # + tmem_ptr = tmem.retrieve_ptr(self.acc_dtype) + # (MMA, MMA_M, MMA_N, STAGE) + tCtAcc_base_ = cute.make_tensor(tmem_ptr, tCtAcc_fake.layout) + tCtAcc_final = cute.make_tensor( + tCtAcc_base_.iterator + tcgen05.find_tmem_tensor_col_offset(tCtAcc_base_), + tCtAcc_base_.layout, + ) # - # Async arrive epilogue buffer full + # Partition for epilogue # - epi_pipeline.producer_commit(epi_producer_state) - epi_producer_state.advance() + epi_tidx = tidx % 128 + ( + tiled_copy_t2r, + tTR_tAcc_base, + tTR_rAcc, + ) = self.epilog_tmem_copy_and_partition( + epi_tidx, tCtAcc_final, tCgC, epi_tile, use_2cta_instrs + ) + tTR_cC = None + if cutlass.const_expr(not self.fp8_quantize_1x128): + cC_t = transform_partitioned_tensor_layout(tCcC) + cC_epi = cute.flat_divide(cC_t, epi_tile) + tTR_cC = tiled_copy_t2r.get_slice(epi_tidx).partition_D(cC_epi) + + tTR_rC = None + tiled_copy_r2s = None + tRS_rC = None + tRS_sC = None + bSG_sC = None + bSG_gC_partitioned = None + if cutlass.const_expr(not self.fp8_quantize_1x128): + tTR_rC = cute.make_rmem_tensor(tTR_rAcc.shape, self.c_dtype) + tiled_copy_r2s, tRS_rC, tRS_sC = self.epilog_smem_copy_and_partition( + tiled_copy_t2r, tTR_rC, epi_tidx, sC + ) + ( + tma_atom_c, + bSG_sC, + bSG_gC_partitioned, + ) = self.epilog_gmem_copy_and_partition( + epi_tidx, tma_atom_c, tCgC, epi_tile, sC + ) # - # Advance to next tile + # Persistent tile scheduling loop # - tile_info_pipeline.consumer_wait(tile_info_consumer_state) - for idx in cutlass.range(4, unroll_full=True): - tile_info[idx] = sInfo[(idx, tile_info_consumer_state.index)] - is_valid_tile = tile_info[3] == 1 - cute.arch.fence_proxy( - "async.shared", - space="cta", + tile_sched = utils.StaticPersistentTileScheduler.create( + tile_sched_params, cute.arch.block_idx(), cute.arch.grid_dim() ) - tile_info_pipeline.consumer_release(tile_info_consumer_state) - tile_info_consumer_state.advance() + work_tile = tile_sched.initial_work_tile_info() + + epi_consumer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Consumer, 1 + ) + + c_pipeline = None + if cutlass.const_expr(not self.fp8_quantize_1x128): + # Threads/warps participating in tma store pipeline + c_producer_group = pipeline.CooperativeGroup( + pipeline.Agent.Thread, + 32 * len(self.epilog_warp_id), + ) + c_pipeline = pipeline.PipelineTmaStore.create( + num_stages=self.num_c_stage, + producer_group=c_producer_group, + ) + + tile_info_consumer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Consumer, self.num_tile_stage + ) + + # get the first tile info + tile_info = cute.make_rmem_tensor((4,), cutlass.Int32) + # Get tile coord from tile scheduler + cur_tile_coord = work_tile.tile_idx + # initialize the tile info + tile_info[0] = cur_tile_coord[0] + tile_info[1] = cur_tile_coord[1] + tile_info[2] = cur_tile_coord[2] + tile_info[3] = work_tile.is_valid_tile + + is_valid_tile = cutlass.Boolean(1) + is_valid_tile = tile_info[3] == 1 + + num_prev_subtiles = cutlass.Int32(0) + + while is_valid_tile: + mma_tile_coord_mnl = ( + tile_info[0] // cute.size(tiled_mma.thr_id.shape), + tile_info[1], + tile_info[2], + ) + # + # Slice to per mma tile index + # + bSG_gC = None + if cutlass.const_expr(not self.fp8_quantize_1x128): + # ((ATOM_V, REST_V), EPI_M, EPI_N) + bSG_gC = bSG_gC_partitioned[ + ( + None, + None, + None, + mma_tile_coord_mnl[0], + mma_tile_coord_mnl[1], + mma_tile_coord_mnl[2], + ) + ] + + # Set tensor memory buffer for current tile + # (T2R, T2R_M, T2R_N, EPI_M, EPI_M) + tTR_tAcc = tTR_tAcc_base[ + (None, None, None, None, None, epi_consumer_state.index) + ] + + # + # Wait for accumulator buffer full + # + epi_pipeline.consumer_wait(epi_consumer_state) + + tTR_tAcc = cute.group_modes(tTR_tAcc, 3, cute.rank(tTR_tAcc)) + tTR_cC_tile = None + if cutlass.const_expr(not self.fp8_quantize_1x128): + bSG_gC = cute.group_modes(bSG_gC, 1, cute.rank(bSG_gC)) + tTR_cC_tile = tTR_cC[(None, None, None, None, None, *mma_tile_coord_mnl)] + tTR_cC_tile = cute.group_modes(tTR_cC_tile, 3, cute.rank(tTR_cC_tile)) + + # + # Store accumulator to global memory in subtiles + # + if cutlass.const_expr(self.fp8_quantize_1x128): + assert fp8_scale is not None + self._epilog_fp8_register_scratch( + tiled_copy_t2r, + tTR_tAcc, + tTR_rAcc, + sC, + epi_consumer_state, + epi_pipeline, + epi_tidx // 32, + lane_idx, + mma_tile_coord_mnl, + mma_tile_coord_v, + epilogue_op, + mC_mnl, + fp8_scale, + ) + else: + subtile_cnt = cute.size(tTR_tAcc.shape, mode=[3]) + for subtile_idx in cutlass.range(subtile_cnt): + # Load accumulator from tensor memory buffer to register. + tTR_tAcc_mn = tTR_tAcc[(None, None, None, subtile_idx)] + cute.copy(tiled_copy_t2r, tTR_tAcc_mn, tTR_rAcc) + + # Convert to C type. + acc_vec = tiled_copy_r2s.retile(tTR_rAcc).load() + acc_vec = epilogue_op(acc_vec.to(self.c_dtype)) + tRS_rC.store(acc_vec) + + # Store C to shared memory. + num_prev_subtiles = num_prev_subtiles + 1 + c_buffer = num_prev_subtiles % self.num_c_stage + cute.copy( + tiled_copy_r2s, + tRS_rC, + tRS_sC[(None, None, None, c_buffer)], + ) + # Fence and barrier to make the shared-memory store + # visible to the TMA store. + cute.arch.fence_proxy( + "async.shared", + space="cta", + ) + self.epilog_sync_barrier.arrive_and_wait() + + # TMA store C to global memory. + if warp_idx == self.epilog_warp_id[0]: + cute.copy( + tma_atom_c, + bSG_sC[(None, c_buffer)], + bSG_gC[(None, subtile_idx)], + ) + c_pipeline.producer_commit() + c_pipeline.producer_acquire() + self.epilog_sync_barrier.arrive_and_wait() + + # + # Async arrive accumulator buffer empty + # + if cutlass.const_expr(not self.fp8_quantize_1x128): + epi_pipeline.consumer_release(epi_consumer_state) + epi_consumer_state.advance() + + # + # Advance to next tile + # + tile_info_pipeline.consumer_wait(tile_info_consumer_state) + for idx in cutlass.range(4, unroll_full=True): + tile_info[idx] = sInfo[(idx, tile_info_consumer_state.index)] + is_valid_tile = tile_info[3] == 1 + cute.arch.fence_proxy( + "async.shared", + space="cta", + ) + tile_info_pipeline.consumer_release(tile_info_consumer_state) + tile_info_consumer_state.advance() + + # + # Dealloc the tensor memory buffer + # + tmem.relinquish_alloc_permit() + self.epilog_sync_barrier.arrive_and_wait() + tmem.free(tmem_ptr) + # + # Wait for C store complete + # + if cutlass.const_expr(not self.fp8_quantize_1x128): + c_pipeline.producer_tail() + + @cute.jit + def _split_acc_update( + self, + warp_idx: cutlass.Int32, + tidx: cutlass.Int32, + tiled_mma: cute.TiledMma, + tmem: utils.TmemAllocator, + tCtAcc_fake: cute.Tensor, + tCgC: cute.Tensor, + sC: cute.Tensor, + sSFA_view_as_C: cute.Tensor, + sSFB_view_as_C: cute.Tensor, + sInfo: cute.Tensor, + tma_atom_c: cute.CopyAtom, + acc_pipeline: pipeline.PipelineUmmaAsync, + scale_pipeline: pipeline.PipelineCpAsync, + tile_info_pipeline: pipeline.PipelineAsync, + epi_tile: cute.Tile, + tile_sched_params: utils.PersistentTileSchedulerParams, + k_tile_cnt: cutlass.Int32, + epilogue_op: cutlass.Constexpr, + ): + """Accumulator update and store of 256-row CTA tiles on warps 0-7. + + Warp group 0 (warps 0-3) and warp group 1 (warps 4-7) each own one + 128-row MMA_M half: every thread keeps one row (128 FP32 values) of + the final accumulator in registers and stores its half through its + own C stages, so the final accumulator never goes back to TMEM. + """ + cute.arch.setmaxregister_increase(self.num_regs_split_acc_update_warps) + half_idx = warp_idx // 4 + epi_tidx = tidx % 128 + + # The allocator warp (epilog_warp_id[0]) belongs to warp group 1. + tmem.allocate(self.num_tmem_alloc_cols) + tmem.wait_for_alloc() + tmem_ptr = tmem.retrieve_ptr(self.acc_dtype) + # (MMA, MMA_M, MMA_N, STAGE) + tCtAcc_base = cute.make_tensor(tmem_ptr, tCtAcc_fake.layout) # - # Specialized epilogue warps + # Partition for the accumulator update: one thread reads one TMEM lane # - if warp_idx <= self.epilog_warp_id[-1] and warp_idx >= self.epilog_warp_id[0]: - cute.arch.setmaxregister_increase(self.num_regs_epilogue_warps) - # - # Alloc tensor memory buffer - # - tmem.allocate(self.num_tmem_alloc_cols) + tmem_load_atom = cute.make_copy_atom( + tcgen05.copy.Ld32x32bOp(tcgen05.copy.Repetition(32)), + self.acc_dtype, + ) + # (EPI_TILE_M, EPI_TILE_N, EPI_M, EPI_N, STAGE) + tAcc_epi = cute.flat_divide(transform_partitioned_tensor_layout(tCtAcc_base), epi_tile) + tiled_copy_t2r = tcgen05.make_tmem_copy(tmem_load_atom, tAcc_epi[(None, None, 0, 0, 0)]) + thr_copy_t2r = tiled_copy_t2r.get_slice(epi_tidx) + # (T2R, T2R_M, T2R_N, EPI_M, EPI_N, STAGE) + tTR_tAcc_base = thr_copy_t2r.partition_S(tAcc_epi) + tTR_sSFA = thr_copy_t2r.partition_D(cute.flat_divide(sSFA_view_as_C, epi_tile)) + tTR_sSFB = thr_copy_t2r.partition_D(cute.flat_divide(sSFB_view_as_C, epi_tile)) + # (T2R, T2R_M, T2R_N, EPI_M, EPI_N, loopM, loopN, loopL) + tTR_gC = thr_copy_t2r.partition_D( + cute.flat_divide(transform_partitioned_tensor_layout(tCgC), epi_tile) + ) + # (T2R, T2R_M, T2R_N) + tTR_rAcc = cute.make_rmem_tensor( + tTR_gC[(None, None, None, 0, 0, 0, 0, 0)].shape, self.acc_dtype + ) + # (T2R, T2R_M, T2R_N, EPI_N): this thread's row of its MMA_M half + tTR_rAcc_final = cute.make_rmem_tensor( + tTR_gC[(None, None, None, 0, None, 0, 0, 0)].shape, self.acc_dtype + ) + subtile_cnt = cute.size(tTR_rAcc_final, mode=[3]) - # - # Bar sync for retrieve tensor memory ptr from shared memory - # - tmem.wait_for_alloc() + # + # Partition for the store: registers -> sC -> TMA + # + tTR_rC = cute.make_rmem_tensor(tTR_rAcc.shape, self.c_dtype) + tiled_copy_r2s, tRS_rC, tRS_sC = self.epilog_smem_copy_and_partition( + tiled_copy_t2r, tTR_rC, epi_tidx, sC + ) + # ((ATOM_V, REST_V), STAGE) + # ((ATOM_V, REST_V), EPI_M, EPI_N, loopM, loopN, loopL) + tma_atom_c, bSG_sC, bSG_gC_partitioned = self.epilog_gmem_copy_and_partition( + epi_tidx, tma_atom_c, tCgC, epi_tile, sC + ) + # Half h stores through C stages [h * num_c_stage_half, (h + 1) * num_c_stage_half). + num_c_stage_half = self.num_c_stage // 2 + c_pipeline = pipeline.PipelineTmaStore.create( + num_stages=num_c_stage_half, + producer_group=pipeline.CooperativeGroup( + pipeline.Agent.Thread, + self.threads_per_warp * len(self.acc_update_warp_id), + ), + ) + store_sync_barrier = pipeline.NamedBarrier( + barrier_id=self.split_store_barrier_id + half_idx, + num_threads=self.threads_per_warp * len(self.acc_update_warp_id), + ) - # - # Retrieving tensor memory ptr and make accumulator tensor - # - tmem_ptr = tmem.retrieve_ptr(self.acc_dtype) - # (MMA, MMA_M, MMA_N, STAGE) - tCtAcc_base_ = cute.make_tensor(tmem_ptr, tCtAcc_fake.layout) - tCtAcc_final = cute.make_tensor( - tCtAcc_base_.iterator + tcgen05.find_tmem_tensor_col_offset(tCtAcc_base_), - tCtAcc_base_.layout, - ) + # + # Persistent tile scheduling loop + # + tile_sched = utils.StaticPersistentTileScheduler.create( + tile_sched_params, cute.arch.block_idx(), cute.arch.grid_dim() + ) + work_tile = tile_sched.initial_work_tile_info() - # - # Partition for epilogue - # - epi_tidx = tidx % 128 - ( - tiled_copy_t2r, - tTR_tAcc_base, - tTR_rAcc, - ) = self.epilog_tmem_copy_and_partition( - epi_tidx, tCtAcc_final, tCgC, epi_tile, use_2cta_instrs - ) - tTR_cC = None - if cutlass.const_expr(not self.fp8_quantize_1x128): - cC_t = transform_partitioned_tensor_layout(tCcC) - cC_epi = cute.flat_divide(cC_t, epi_tile) - tTR_cC = tiled_copy_t2r.get_slice(epi_tidx).partition_D(cC_epi) - - tTR_rC = None - tiled_copy_r2s = None - tRS_rC = None - tRS_sC = None - bSG_sC = None - bSG_gC_partitioned = None - if cutlass.const_expr(not self.fp8_quantize_1x128): - tTR_rC = cute.make_rmem_tensor(tTR_rAcc.shape, self.c_dtype) - tiled_copy_r2s, tRS_rC, tRS_sC = self.epilog_smem_copy_and_partition( - tiled_copy_t2r, tTR_rC, epi_tidx, sC - ) - ( - tma_atom_c, - bSG_sC, - bSG_gC_partitioned, - ) = self.epilog_gmem_copy_and_partition(epi_tidx, tma_atom_c, tCgC, epi_tile, sC) + acc_consumer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Consumer, self.num_acc_stage + ) - # - # Persistent tile scheduling loop - # - tile_sched = utils.StaticPersistentTileScheduler.create( - tile_sched_params, cute.arch.block_idx(), cute.arch.grid_dim() - ) - work_tile = tile_sched.initial_work_tile_info() + scale_consumer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Consumer, self.num_scale_stage + ) - epi_consumer_state = pipeline.make_pipeline_state(pipeline.PipelineUserType.Consumer, 1) + tile_info_consumer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Consumer, self.num_tile_stage + ) - c_pipeline = None - if cutlass.const_expr(not self.fp8_quantize_1x128): - # Threads/warps participating in tma store pipeline - c_producer_group = pipeline.CooperativeGroup( - pipeline.Agent.Thread, - 32 * len(self.epilog_warp_id), - ) - c_pipeline = pipeline.PipelineTmaStore.create( - num_stages=self.num_c_stage, - producer_group=c_producer_group, - ) + # get the first tile info + tile_info = cute.make_rmem_tensor((4,), cutlass.Int32) + # Get tile coord from tile scheduler + cur_tile_coord = work_tile.tile_idx + # initialize the tile info + tile_info[0] = cur_tile_coord[0] + tile_info[1] = cur_tile_coord[1] + tile_info[2] = cur_tile_coord[2] + tile_info[3] = work_tile.is_valid_tile - tile_info_consumer_state = pipeline.make_pipeline_state( - pipeline.PipelineUserType.Consumer, self.num_tile_stage - ) + is_valid_tile = cutlass.Boolean(1) + is_valid_tile = tile_info[3] == 1 - # get the first tile info - tile_info = cute.make_rmem_tensor((4,), cutlass.Int32) - # Get tile coord from tile scheduler - cur_tile_coord = work_tile.tile_idx - # initialize the tile info - tile_info[0] = cur_tile_coord[0] - tile_info[1] = cur_tile_coord[1] - tile_info[2] = cur_tile_coord[2] - tile_info[3] = work_tile.is_valid_tile + num_prev_subtiles = cutlass.Int32(0) - is_valid_tile = cutlass.Boolean(1) - is_valid_tile = tile_info[3] == 1 + while is_valid_tile: + mma_tile_coord_mnl = ( + tile_info[0] // cute.size(tiled_mma.thr_id.shape), + tile_info[1], + tile_info[2], + ) + # initialize the final accumulator + tTR_rAcc_final.fill(0.0) - num_prev_subtiles = cutlass.Int32(0) + scale_consumer_state.reset_count() + peek_scale_full_status = cutlass.Boolean(1) + if scale_consumer_state.count < k_tile_cnt: + peek_scale_full_status = scale_pipeline.consumer_try_wait(scale_consumer_state) - while is_valid_tile: - mma_tile_coord_mnl = ( - tile_info[0] // cute.size(tiled_mma.thr_id.shape), - tile_info[1], - tile_info[2], - ) + acc_consumer_state.reset_count() + peek_acc_full_status = cutlass.Boolean(1) + if acc_consumer_state.count < k_tile_cnt: + peek_acc_full_status = acc_pipeline.consumer_try_wait(acc_consumer_state) + + for k_tile in cutlass.range(0, k_tile_cnt, 1, unroll=1): # - # Slice to per mma tile index + # Wait for scale buffer full # - bSG_gC = None - if cutlass.const_expr(not self.fp8_quantize_1x128): - # ((ATOM_V, REST_V), EPI_M, EPI_N) - bSG_gC = bSG_gC_partitioned[ - ( - None, - None, - None, - mma_tile_coord_mnl[0], - mma_tile_coord_mnl[1], - mma_tile_coord_mnl[2], - ) - ] + scale_pipeline.consumer_wait(scale_consumer_state, peek_scale_full_status) - # Set tensor memory buffer for current tile - # (T2R, T2R_M, T2R_N, EPI_M, EPI_M) - tTR_tAcc = tTR_tAcc_base[(None, None, None, None, None, epi_consumer_state.index)] + # A thread's values share one row (scale A) and one 128-column + # block (scale B). + scale = ( + tTR_sSFA[(0, 0, 0, half_idx, 0, scale_consumer_state.index)] + * tTR_sSFB[(0, 0, 0, half_idx, 0, scale_consumer_state.index)] + ) # # Wait for accumulator buffer full # - epi_pipeline.consumer_wait(epi_consumer_state) - - tTR_tAcc = cute.group_modes(tTR_tAcc, 3, cute.rank(tTR_tAcc)) - tTR_cC_tile = None - if cutlass.const_expr(not self.fp8_quantize_1x128): - bSG_gC = cute.group_modes(bSG_gC, 1, cute.rank(bSG_gC)) - tTR_cC_tile = tTR_cC[(None, None, None, None, None, *mma_tile_coord_mnl)] - tTR_cC_tile = cute.group_modes(tTR_cC_tile, 3, cute.rank(tTR_cC_tile)) + acc_pipeline.consumer_wait(acc_consumer_state, peek_acc_full_status) # - # Store accumulator to global memory in subtiles + # Update this half's accumulator by the scale factor in subtiles # - if cutlass.const_expr(self.fp8_quantize_1x128): - assert fp8_scale is not None - self._epilog_fp8_register_scratch( + for subtile_idx in cutlass.range_constexpr(subtile_cnt): + cute.copy( tiled_copy_t2r, - tTR_tAcc, + tTR_tAcc_base[ + ( + None, + None, + None, + half_idx, + subtile_idx, + acc_consumer_state.index, + ) + ], tTR_rAcc, - sC, - epi_consumer_state, - epi_pipeline, - epi_tidx // 32, - lane_idx, - mma_tile_coord_mnl, - mma_tile_coord_v, - epilogue_op, - mC_mnl, - fp8_scale, ) - else: - subtile_cnt = cute.size(tTR_tAcc.shape, mode=[3]) - for subtile_idx in cutlass.range(subtile_cnt): - # Load accumulator from tensor memory buffer to register. - tTR_tAcc_mn = tTR_tAcc[(None, None, None, subtile_idx)] - cute.copy(tiled_copy_t2r, tTR_tAcc_mn, tTR_rAcc) - - # Convert to C type. - acc_vec = tiled_copy_r2s.retile(tTR_rAcc).load() - acc_vec = epilogue_op(acc_vec.to(self.c_dtype)) - tRS_rC.store(acc_vec) - - # Store C to shared memory. - num_prev_subtiles = num_prev_subtiles + 1 - c_buffer = num_prev_subtiles % self.num_c_stage - cute.copy( - tiled_copy_r2s, - tRS_rC, - tRS_sC[(None, None, None, c_buffer)], - ) - # Fence and barrier to make the shared-memory store - # visible to the TMA store. - cute.arch.fence_proxy( - "async.shared", - space="cta", - ) - self.epilog_sync_barrier.arrive_and_wait() - - # TMA store C to global memory. - if warp_idx == self.epilog_warp_id[0]: - cute.copy( - tma_atom_c, - bSG_sC[(None, c_buffer)], - bSG_gC[(None, subtile_idx)], - ) - c_pipeline.producer_commit() - c_pipeline.producer_acquire() - self.epilog_sync_barrier.arrive_and_wait() + tTR_rAcc_subtile = tTR_rAcc_final[(None, None, None, subtile_idx)] + tTR_rAcc_subtile.store(tTR_rAcc.load() * scale + tTR_rAcc_subtile.load()) # - # Async arrive accumulator buffer empty + # Async arrive scale buffer empty # - if cutlass.const_expr(not self.fp8_quantize_1x128): - epi_pipeline.consumer_release(epi_consumer_state) - epi_consumer_state.advance() + scale_pipeline.consumer_release(scale_consumer_state) + scale_consumer_state.advance() + peek_scale_full_status = cutlass.Boolean(1) + if scale_consumer_state.count < k_tile_cnt: + peek_scale_full_status = scale_pipeline.consumer_try_wait(scale_consumer_state) # - # Advance to next tile + # Async arrive accumulator buffer empty # - tile_info_pipeline.consumer_wait(tile_info_consumer_state) - for idx in cutlass.range(4, unroll_full=True): - tile_info[idx] = sInfo[(idx, tile_info_consumer_state.index)] - is_valid_tile = tile_info[3] == 1 + with cute.arch.elect_one(): + acc_pipeline.consumer_release(acc_consumer_state) + acc_consumer_state.advance() + + peek_acc_full_status = cutlass.Boolean(1) + if acc_consumer_state.count < k_tile_cnt: + peek_acc_full_status = acc_pipeline.consumer_try_wait(acc_consumer_state) + + # + # Store this half in subtiles + # + # ((ATOM_V, REST_V), EPI_M, EPI_N) + bSG_gC = bSG_gC_partitioned[(None, None, None, *mma_tile_coord_mnl)] + for subtile_idx in cutlass.range_constexpr(subtile_cnt): + # Convert to C type. + acc_vec = tiled_copy_r2s.retile( + tTR_rAcc_final[(None, None, None, subtile_idx)] + ).load() + tRS_rC.store(epilogue_op(acc_vec.to(self.c_dtype))) + + # Store C to shared memory. + num_prev_subtiles = num_prev_subtiles + 1 + c_buffer = half_idx * num_c_stage_half + num_prev_subtiles % num_c_stage_half + cute.copy( + tiled_copy_r2s, + tRS_rC, + tRS_sC[(None, None, None, c_buffer)], + ) + # Fence and barrier to make the shared-memory store visible + # to the TMA store. cute.arch.fence_proxy( "async.shared", space="cta", ) - tile_info_pipeline.consumer_release(tile_info_consumer_state) - tile_info_consumer_state.advance() + store_sync_barrier.arrive_and_wait() + + # TMA store C to global memory from the half's first warp. + if warp_idx % 4 == 0: + cute.copy( + tma_atom_c, + bSG_sC[(None, c_buffer)], + bSG_gC[(None, half_idx, subtile_idx)], + ) + c_pipeline.producer_commit() + c_pipeline.producer_acquire() + store_sync_barrier.arrive_and_wait() # - # Dealloc the tensor memory buffer - # - tmem.relinquish_alloc_permit() - self.epilog_sync_barrier.arrive_and_wait() - tmem.free(tmem_ptr) - # - # Wait for C store complete + # Advance to next tile # - if cutlass.const_expr(not self.fp8_quantize_1x128): - c_pipeline.producer_tail() + tile_info_pipeline.consumer_wait(tile_info_consumer_state) + for idx in cutlass.range(4, unroll_full=True): + tile_info[idx] = sInfo[(idx, tile_info_consumer_state.index)] + is_valid_tile = tile_info[3] == 1 + cute.arch.fence_proxy( + "async.shared", + space="cta", + ) + tile_info_pipeline.consumer_release(tile_info_consumer_state) + tile_info_consumer_state.advance() + + # + # Dealloc the tensor memory buffer once both halves stopped reading it + # + tmem.relinquish_alloc_permit() + self.split_acc_update_done_barrier.arrive_and_wait() + tmem.free(tmem_ptr) + # + # Wait for C store complete + # + c_pipeline.producer_tail() @cute.jit def _epilog_fp8_register_scratch( @@ -2500,9 +2852,11 @@ def can_implement( Check if the gemm can be implemented on SM107. """ can_impl = True - # The fused packed-FP8 epilogue writes full 128-column tiles and only + # The fused packed-FP8 epilogue writes full 128x128 CTA tiles and only # predicates rows. - if fuse_fp8_quantize_1x128 and (n % 128 != 0 or c_major != "n"): + if fuse_fp8_quantize_1x128 and ( + n % 128 != 0 or c_major != "n" or mma_tiler[0] // (2 if use_2cta_instrs else 1) != 128 + ): can_impl = False # Check data types diff --git a/tensorrt_llm/_torch/kernels/rubin/dense_blockscaled_gemm_persistent.py b/tensorrt_llm/_torch/kernels/rubin/dense_blockscaled_gemm_persistent.py index dddbb27ee0bf..f8bf0d78a5b0 100644 --- a/tensorrt_llm/_torch/kernels/rubin/dense_blockscaled_gemm_persistent.py +++ b/tensorrt_llm/_torch/kernels/rubin/dense_blockscaled_gemm_persistent.py @@ -238,6 +238,20 @@ def __init__( raise ValueError(f"split_k must be >= 1, got {split_k}") self.split_k = split_k + def _split_k_range(self, k_tile_cnt, l_coord): + """Return ``(k_tile_begin, k_tiles)`` of the split-K slice at scheduler L ``l_coord``. + + Slices take balanced ranges, ``[floor(k_tile_cnt * s / split_k), + floor(k_tile_cnt * (s + 1) / split_k))`` for slice ``s = l_coord % split_k``, + so any ``split_k <= k_tile_cnt`` gives every slice at least one K tile. + """ + if cutlass.const_expr(self.split_k == 1): + return 0, k_tile_cnt + split_idx = l_coord % self.split_k + k_tile_begin = (k_tile_cnt * split_idx) // self.split_k + k_tile_end = (k_tile_cnt * (split_idx + 1)) // self.split_k + return k_tile_begin, k_tile_end - k_tile_begin + def _compute_grid( self, c: cute.Tensor, @@ -1381,9 +1395,8 @@ def kernel_impl( gC_mnl = cute.local_tile( mC_mnl, cute.slice_(self.mma_tiler, (None, None, 0)), (None, None, None) ) + # Total K tiles; each split-K slice takes a balanced range of them. k_tile_cnt = cutlass.Int32(cute.size(gA_mkl, mode=[3])) - if cutlass.const_expr(self.split_k > 1): - k_tile_cnt = k_tile_cnt // self.split_k # # Partition global tensor for TiledMMA_A/B/C @@ -1510,11 +1523,8 @@ def kernel_impl( cur_tile_coord[2], ) - k_tile_start = 0 + k_tile_start, tile_k_cnt = self._split_k_range(k_tile_cnt, mma_tile_coord_mnl[2]) if cutlass.const_expr(self.split_k > 1): - k_tile_start = cutlass.Int32( - (mma_tile_coord_mnl[2] % self.split_k) * k_tile_cnt - ) mma_tile_coord_mnl = ( mma_tile_coord_mnl[0], mma_tile_coord_mnl[1], @@ -1545,7 +1555,7 @@ def kernel_impl( # if self.prefetch_enabled: for pf_k_tile in cutlass.range( - 0, min(self.prefetch_dist, k_tile_cnt), unroll=1 + 0, min(self.prefetch_dist, tile_k_cnt), unroll=1 ): prefetch_k = k_tile_start + pf_k_tile cute.prefetch(tma_atom_a, tAgA_slice[(None, prefetch_k)]) @@ -1556,12 +1566,12 @@ def kernel_impl( # Peek (try_wait) AB buffer empty for k_tile = prefetch_k_tile_cnt ab_producer_state.reset_count() peek_ab_empty_status = cutlass.Boolean(1) - if ab_producer_state.count < k_tile_cnt: + if ab_producer_state.count < tile_k_cnt: peek_ab_empty_status = ab_pipeline.producer_try_acquire(ab_producer_state) # # Tma load loop # - for k_tile in cutlass.range(0, k_tile_cnt, 1, unroll=1): + for k_tile in cutlass.range(0, tile_k_cnt, 1, unroll=1): # Conditionally wait for AB buffer empty ab_pipeline.producer_acquire(ab_producer_state, peek_ab_empty_status) @@ -1598,7 +1608,7 @@ def kernel_impl( # Prefetch: Rolling prefetch for next tiles if self.prefetch_enabled: - if k_tile < k_tile_cnt - self.prefetch_dist: + if k_tile < tile_k_cnt - self.prefetch_dist: future_k_tile = load_k + self.prefetch_dist cute.prefetch( tma_atom_a, @@ -1620,7 +1630,7 @@ def kernel_impl( # Peek (try_wait) AB buffer empty for k_tile = prefetch_k_tile_cnt + k_tile + 1 ab_producer_state.advance() peek_ab_empty_status = cutlass.Boolean(1) - if ab_producer_state.count < k_tile_cnt: + if ab_producer_state.count < tile_k_cnt: peek_ab_empty_status = ab_pipeline.producer_try_acquire(ab_producer_state) # @@ -1706,6 +1716,7 @@ def kernel_impl( cur_tile_coord[1], cur_tile_coord[2], ) + _, tile_k_cnt = self._split_k_range(k_tile_cnt, cur_tile_coord[2]) if cutlass.const_expr(self.split_k > 1): mma_tile_coord_mnl = ( mma_tile_coord_mnl[0], @@ -1723,7 +1734,7 @@ def kernel_impl( # Peek (try_wait) AB buffer full for k_tile = 0 ab_consumer_state.reset_count() peek_ab_full_status = cutlass.Boolean(1) - if ab_consumer_state.count < k_tile_cnt and is_leader_cta: + if ab_consumer_state.count < tile_k_cnt and is_leader_cta: peek_ab_full_status = ab_pipeline.consumer_try_wait(ab_consumer_state) # @@ -1750,7 +1761,7 @@ def kernel_impl( # # Mma mainloop # - for k_tile in range(k_tile_cnt): + for k_tile in range(tile_k_cnt): if is_leader_cta: # Conditionally wait for AB buffer full ab_pipeline.consumer_wait(ab_consumer_state, peek_ab_full_status) @@ -1853,7 +1864,7 @@ def kernel_impl( # Peek (try_wait) AB buffer full for k_tile = k_tile + 1 ab_consumer_state.advance() peek_ab_full_status = cutlass.Boolean(1) - if ab_consumer_state.count < k_tile_cnt: + if ab_consumer_state.count < tile_k_cnt: if is_leader_cta: peek_ab_full_status = ab_pipeline.consumer_try_wait(ab_consumer_state) @@ -2376,10 +2387,10 @@ def can_implement( raise testing.CantImplementError(f"split_k must be >= 1, got {split_k}") if split_k > 1: num_k_tiles = k // mma_tiler[2] - if l != 1 or k % mma_tiler[2] != 0 or num_k_tiles % split_k != 0: + if l != 1 or k % mma_tiler[2] != 0 or split_k > num_k_tiles: raise testing.CantImplementError( - f"split_k={split_k} requires L==1 and K tiles " - f"({num_k_tiles}) divisible by split_k" + f"split_k={split_k} requires L==1 and at least split_k " + f"full K tiles (got {num_k_tiles})" ) # The grid uses a half-sized per-CTA M tile for 2-CTA MMA. # Requiring enough real M tiles avoids phantom cluster peers, which diff --git a/tensorrt_llm/_torch/locality_domain/autotune.py b/tensorrt_llm/_torch/locality_domain/autotune.py index 0f446e02dbb0..1624d2bc59a5 100644 --- a/tensorrt_llm/_torch/locality_domain/autotune.py +++ b/tensorrt_llm/_torch/locality_domain/autotune.py @@ -67,6 +67,9 @@ def unique_id(self) -> tuple[Any, int, tuple[tuple[int, int], ...]]: self._runtime.topology_identity(), ) + def tactic_space_key(self) -> tuple[Any, ...] | None: + return self._op_runner.tactic_space_key() + def get_valid_tactics( self, inputs: list[torch.Tensor], diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py index 0fdb49950ec4..e3cea3538ecd 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_engine.py +++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py @@ -33,6 +33,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 @@ -121,6 +122,19 @@ def _configure_deep_gemm_pdl() -> None: _DEEP_GEMM_PDL_CONFIGURED = True +def _configure_autotuner_nvmmh(llm_args: TorchLlmArgs) -> None: + """Install the engine's nvMatmulHeuristics policy on the process AutoTuner. + + Runs before model construction because CuTe DSL runners may enumerate + tactics while modules are built. Engines of one process (e.g. target and + draft) share the LLM args and therefore the policy. + """ + # Partial test doubles of TorchLlmArgs may omit the field. + config = getattr(llm_args, "autotuner_nvmmh_config", None) + AutoTuner.get().nvmmh_config = (None if config is None else NvMMHConfig( + max_tactics=config.max_tactics)) + + class ModelEngine: """Own model setup and coordinate runner execution and resource lifetimes.""" @@ -177,6 +191,7 @@ def _init_from_args( services derived from the arguments; capacity may still be unset. Returns the encoder token budget before any multimodal increase.""" _configure_deep_gemm_pdl() + _configure_autotuner_nvmmh(llm_args) if llm_args.encode_only and llm_args.mm_encoder_only: raise ValueError("encode_only and mm_encoder_only are mutually exclusive.") diff --git a/tensorrt_llm/llmapi/__init__.py b/tensorrt_llm/llmapi/__init__.py index 7075805f47d0..a1425e9cfee2 100644 --- a/tensorrt_llm/llmapi/__init__.py +++ b/tensorrt_llm/llmapi/__init__.py @@ -9,7 +9,8 @@ from ..scheduling_params import SchedulingParams from .llm import LLM, RequestOutput # yapf: disable -from .llm_args import (AttentionDpConfig, AutoDecodingConfig, BlockReuseConfig, +from .llm_args import (AttentionDpConfig, AutoDecodingConfig, + AutoTunerNvMMHConfig, BlockReuseConfig, CacheTransceiverConfig, CalibConfig, CapacitySchedulerPolicy, ColdPageQuantizationCompressionConfig, @@ -87,6 +88,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 5c423261ed9d..6d4f3d6ae958 100644 --- a/tensorrt_llm/llmapi/llm_args.py +++ b/tensorrt_llm/llmapi/llm_args.py @@ -2004,6 +2004,21 @@ class Nvfp4GemmConfig(StrictBaseModel): description="List of backends to consider for auto-selection.") +class AutoTunerNvMMHConfig(StrictBaseModel): + """nvMatmulHeuristics (nvMMH) guidance for the CuTe DSL GEMM/BMM autotuner. + + CuTe DSL GEMM and BMM runners profile only the kernel configurations that + nvMatmulHeuristics ranks highest for each problem shape, taking every + field it models (CTA tile, cluster shape, split-K, rasterization order and + swizzle) as returned, instead of profiling their full tactic space. + """ + max_tactics: PositiveInt = Field( + default=5, + description= + "Number of top-ranked nvMatmulHeuristics kernel configurations profiled per GEMM/BMM shape." + ) + + class AttentionDpConfig(StrictBaseModel): """Configuration for attention DP.""" enable_balance: bool = Field(default=False, @@ -6830,6 +6845,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= + "Use nvMatmulHeuristics to shrink the CuTe DSL GEMM/BMM autotuning search space to its top-ranked kernel configurations. None profiles every tactic.", + status="prototype") + use_fine_grained_sync: bool = Field( default=False, description= diff --git a/tensorrt_llm/schemas/trtllm-serve-config.schema.json b/tensorrt_llm/schemas/trtllm-serve-config.schema.json index b0e84576f7bf..cbcdb3582e15 100644 --- a/tensorrt_llm/schemas/trtllm-serve-config.schema.json +++ b/tensorrt_llm/schemas/trtllm-serve-config.schema.json @@ -347,6 +347,21 @@ "title": "AutoDecodingConfig", "type": "object" }, + "AutoTunerNvMMHConfig": { + "additionalProperties": false, + "description": "nvMatmulHeuristics (nvMMH) guidance for the CuTe DSL GEMM/BMM autotuner.\n\nCuTe DSL GEMM and BMM runners profile only the kernel configurations that\nnvMatmulHeuristics ranks highest for each problem shape, taking every\nfield it models (CTA tile, cluster shape, split-K, rasterization order and\nswizzle) as returned, instead of profiling their full tactic space.", + "properties": { + "max_tactics": { + "default": 5, + "description": "Number of top-ranked nvMatmulHeuristics kernel configurations profiled per GEMM/BMM shape.", + "exclusiveMinimum": 0, + "title": "Max Tactics", + "type": "integer" + } + }, + "title": "AutoTunerNvMMHConfig", + "type": "object" + }, "BlockReuseConfig": { "additionalProperties": false, "description": "Configuration for KV cache block reuse policies.", @@ -6415,6 +6430,19 @@ "title": "Attn Backend", "type": "string" }, + "autotuner_nvmmh_config": { + "anyOf": [ + { + "$ref": "#/$defs/AutoTunerNvMMHConfig" + }, + { + "type": "null" + } + ], + "default": null, + "description": "Use nvMatmulHeuristics to shrink the CuTe DSL GEMM/BMM autotuning search space to its top-ranked kernel configurations. None profiles every tactic.", + "status": "prototype" + }, "backend": { "const": "pytorch", "default": "pytorch", diff --git a/tensorrt_llm/schemas/trtllm-serve-disagg-config.schema.json b/tensorrt_llm/schemas/trtllm-serve-disagg-config.schema.json index abca43a18d20..aa6ad717d05e 100644 --- a/tensorrt_llm/schemas/trtllm-serve-disagg-config.schema.json +++ b/tensorrt_llm/schemas/trtllm-serve-disagg-config.schema.json @@ -347,6 +347,21 @@ "title": "AutoDecodingConfig", "type": "object" }, + "AutoTunerNvMMHConfig": { + "additionalProperties": false, + "description": "nvMatmulHeuristics (nvMMH) guidance for the CuTe DSL GEMM/BMM autotuner.\n\nCuTe DSL GEMM and BMM runners profile only the kernel configurations that\nnvMatmulHeuristics ranks highest for each problem shape, taking every\nfield it models (CTA tile, cluster shape, split-K, rasterization order and\nswizzle) as returned, instead of profiling their full tactic space.", + "properties": { + "max_tactics": { + "default": 5, + "description": "Number of top-ranked nvMatmulHeuristics kernel configurations profiled per GEMM/BMM shape.", + "exclusiveMinimum": 0, + "title": "Max Tactics", + "type": "integer" + } + }, + "title": "AutoTunerNvMMHConfig", + "type": "object" + }, "BlockReuseConfig": { "additionalProperties": false, "description": "Configuration for KV cache block reuse policies.", @@ -1698,6 +1713,19 @@ "title": "Attn Backend", "type": "string" }, + "autotuner_nvmmh_config": { + "anyOf": [ + { + "$ref": "#/$defs/AutoTunerNvMMHConfig" + }, + { + "type": "null" + } + ], + "default": null, + "description": "Use nvMatmulHeuristics to shrink the CuTe DSL GEMM/BMM autotuning search space to its top-ranked kernel configurations. None profiles every tactic.", + "status": "prototype" + }, "backend": { "const": "pytorch", "default": "pytorch", @@ -7991,6 +8019,19 @@ "title": "Attn Backend", "type": "string" }, + "autotuner_nvmmh_config": { + "anyOf": [ + { + "$ref": "#/$defs/AutoTunerNvMMHConfig" + }, + { + "type": "null" + } + ], + "default": null, + "description": "Use nvMatmulHeuristics to shrink the CuTe DSL GEMM/BMM autotuning search space to its top-ranked kernel configurations. None profiles every tactic.", + "status": "prototype" + }, "backend": { "const": "pytorch", "default": "pytorch", diff --git a/tensorrt_llm/usage/llm_args_golden_manifest.json b/tensorrt_llm/usage/llm_args_golden_manifest.json index 093c1f0ee281..6adf8354ec2e 100644 --- a/tensorrt_llm/usage/llm_args_golden_manifest.json +++ b/tensorrt_llm/usage/llm_args_golden_manifest.json @@ -96,6 +96,11 @@ "kind": "categorical", "path": "attn_backend" }, + { + "capture_policy": "int", + "kind": "value", + "path": "autotuner_nvmmh_config.max_tactics" + }, { "allowed_values": [ "pytorch" diff --git a/tests/unittest/_torch/kernels/parallel/test_cute_dsl_bf16_dense_rubin.py b/tests/unittest/_torch/kernels/parallel/test_cute_dsl_bf16_dense.py similarity index 80% rename from tests/unittest/_torch/kernels/parallel/test_cute_dsl_bf16_dense_rubin.py rename to tests/unittest/_torch/kernels/parallel/test_cute_dsl_bf16_dense.py index 2cc3b4380406..798f70d1b573 100644 --- a/tests/unittest/_torch/kernels/parallel/test_cute_dsl_bf16_dense_rubin.py +++ b/tests/unittest/_torch/kernels/parallel/test_cute_dsl_bf16_dense.py @@ -12,11 +12,13 @@ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. -"""Tests for the SM107 CuTe DSL BF16 dense GEMM/BMM custom ops. +"""Tests for the CuTe DSL BF16 dense GEMM/BMM custom ops. The dispatch tests run on every architecture with CuTe DSL installed and check that ``trtllm::cute_dsl_bf16_{gemm,bmm}_rubin`` are unreachable off SM107. The -correctness tests need SM107 and a CuTe DSL package with the SM107 helpers. +SM107 correctness tests need SM107 and a CuTe DSL package with the SM107 +helpers. The split-K test of ``trtllm::cute_dsl_bf16_gemm_blackwell`` needs +SM100/SM103. """ import pytest @@ -39,6 +41,10 @@ reason="requires SM107 and a CuTe DSL package with SM107 support", ) skip_on_sm107 = pytest.mark.skipif(getSMVersion() == 107, reason="checks the non-SM107 path") +skip_not_sm100f = pytest.mark.skipif( + not IS_CUTLASS_DSL_AVAILABLE or getSMVersion() not in (100, 103), + reason="requires SM100/SM103 and cutlass-dsl", +) def _runner_class(name): @@ -247,3 +253,54 @@ def test_cute_dsl_bf16_bmm_rubin_strided_views(): bad_weight = torch.randn(b, k, n, dtype=torch.bfloat16, device="cuda").transpose(1, 2) with pytest.raises(ValueError, match="K innermost"): runner([act, bad_weight, output], tactic=_select_tactic(tactics, "base")) + + +# -------------------------------------------------------------------------- +# SM100/SM103 split-K. +# -------------------------------------------------------------------------- + + +@skip_not_sm100f +@pytest.mark.parametrize("split_k_slices", [2, 4, 8]) +@pytest.mark.parametrize("c_dtype", [torch.bfloat16, torch.float32]) +def test_cute_dsl_bf16_split_k_gemm_blackwell(split_k_slices, c_dtype): + """Split-K matches the dense reference for large-K, small-N shapes.""" + from tensorrt_llm._torch.custom_ops.cute_dsl_custom_ops import CuteDSLBf16BlackwellGemmRunner + + torch.manual_seed(2026) + runner = CuteDSLBf16BlackwellGemmRunner() + act, weight, output = _gemm_operands(64, 256, 7168, c_dtype) + base = (False, (128, 128), (1, 1)) + tactic = base + ("m", 1, split_k_slices) + + # Direct split-K rounds each partial to the output dtype before the + # atomic TMA ADD, so BF16 output needs tolerance for both rounding and + # arrival-order changes. + rtol, atol = (2e-2, 2.5) if c_dtype == torch.bfloat16 else (1e-2, 1.0) + + expected = act.float() @ weight.t().float() + runner([act, weight, output], tactic=tactic) + torch.cuda.synchronize() + torch.testing.assert_close(output.float(), expected, rtol=rtol, atol=atol) + + # A second launch must not accumulate on the previous output; poisoning C + # also catches a missing zero inside CUDA-graph replay. + output.fill_(float("nan")) + runner([act, weight, output], tactic=tactic) + torch.cuda.synchronize() + torch.testing.assert_close(output.float(), expected, rtol=rtol, atol=atol) + + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + runner([act, weight, output], tactic=tactic) + output.fill_(float("nan")) + graph.replay() + torch.cuda.synchronize() + torch.testing.assert_close(output.float(), expected, rtol=rtol, atol=atol) + + if c_dtype == torch.float32: + split1_output = torch.empty_like(output) + runner([act, weight, split1_output], tactic=base) + torch.cuda.synchronize() + # Both write FP32; direct split-K only changes the reduction order. + torch.testing.assert_close(output, split1_output, rtol=1e-3, atol=1e-2) 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 7b2d08ae5742..b23987dbff6a 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 @@ -1069,6 +1069,64 @@ def test_cute_dsl_fp8_block_scale_bmm_rubin(m, use_tvm_ffi): torch.testing.assert_close(output, output_expected, atol=1e-3, rtol=1e-3) +@pytest.mark.skipif( + getSMVersion() != 107 or not IS_CUTLASS_DSL_RUBIN_AVAILABLE, + reason="The test requires SM107 and SM107 CuTe DSL support.", +) +@pytest.mark.parametrize("m", [7, 300]) +def test_cute_dsl_fp8_block_scale_bmm_rubin_tile_shapes(m, monkeypatch): + """Every MMA tile of the SM107 blockwise BMM, including the 256-row tiles + whose accumulator update is split over two warp groups.""" + from tensorrt_llm._torch.custom_ops.cute_dsl_custom_ops import \ + CuteDSLFp8RubinBmmRunner + + # 8 K tiles wrap the A/B pipeline; 3 N tiles and 3 groups give every CTA + # several persistent tiles. + k, n, num_groups = 1024, 384, 3 + torch.random.manual_seed(0) + a = torch.randn((m, num_groups, k), device="cuda", dtype=torch.bfloat16) / k + a_fp8, a_scales = torch.ops.trtllm.fp8_batched_quantize_1x128_permute102(a) + b = torch.randn((num_groups, n, k), device="cuda", dtype=torch.bfloat16) / k + b_fp8 = torch.empty_like(b, dtype=torch.float8_e4m3fn) + b_scales = torch.empty((num_groups, n // 128, k // 128), + device="cuda", + dtype=torch.float32) + for group_idx in range(num_groups): + b_fp8[group_idx], b_scales[group_idx] = per_block_cast_to_fp8( + b[group_idx]) + output_expected = torch.einsum("mgk,gnk->gmn", a, b) + output = torch.empty((num_groups, m, n), + device="cuda", + dtype=torch.bfloat16) + + monkeypatch.setattr(AutoTuner.get(), "nvmmh_config", None) + runner = CuteDSLFp8RubinBmmRunner() + inputs = [a_fp8, b_fp8, a_scales, b_scales, output] + clusters = {False: ((1, 1), (2, 2)), True: ((2, 1), (2, 2))} + tactics = [ + tactic for tactic in runner.get_valid_tactics(inputs, None) + if tactic[3] in clusters[tactic[0]] + ] + assert len({tactic[:3] for tactic in tactics}) == 4 + + outputs_128_rows = {} + for tactic in tactics: + use_2cta_instrs, mma_tiler, mma_inst_shape = tactic[:3] + output.fill_(float("nan")) + runner(inputs, tactic=tactic) + torch.cuda.synchronize() + assert calc_diff(output, output_expected) < 1e-3, tactic + # A 256-row tile runs the same MMAs and FP32 promotion per element as + # the 128-row tile of its instruction. + if mma_tiler[0] // (2 if use_2cta_instrs else 1) == 128: + outputs_128_rows.setdefault(mma_inst_shape, output.clone()) + else: + torch.testing.assert_close(output, + outputs_128_rows[mma_inst_shape], + rtol=0, + atol=0) + + def _allocate_fp8_bmm_quant_outputs(a_fp8, b_fp8): num_groups, m, _ = a_fp8.shape n = b_fp8.shape[1] diff --git a/tests/unittest/_torch/misc/test_autotuner.py b/tests/unittest/_torch/misc/test_autotuner.py index 7b7d44422ff6..815ad4f39d0c 100644 --- a/tests/unittest/_torch/misc/test_autotuner.py +++ b/tests/unittest/_torch/misc/test_autotuner.py @@ -33,9 +33,10 @@ import tensorrt_llm._torch.autotuner as autotuner from tensorrt_llm._torch.autotuner import (AutoTuner, DistributedTuningStrategy, DynamicDim, DynamicTensorSpec, - FakeTensor, OptimizationProfile, - StaticDim, TunableRunner, - TuningConfig, autotune) + FakeTensor, NvMMHConfig, + OptimizationProfile, StaticDim, + TunableRunner, TuningConfig, + autotune) from tensorrt_llm._torch.distributed import Distributed from tensorrt_llm._torch.utils import (deep_gemm_jit_warmup_buckets, get_power_of_2_num_tokens_buckets, @@ -1514,12 +1515,13 @@ def forward(self, def test_cutedsl_nvfp4_heuristic_matches_full_sweep(monkeypatch): - """End-to-end guard for the nvMatmulHeuristics tactic pruning. + """End-to-end guard for the nvMatmulHeuristics tactic generation. 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. + ``AutoTuner.nvmmh_config`` builds the candidates from nvMatmulHeuristics' + top configurations must be no slower (up to a small tolerance) than the + tactic it selects from the full CuteDSL NVFP4 tactic sweep. This validates + that the nvMMH guidance 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 @@ -1528,9 +1530,9 @@ def test_cutedsl_nvfp4_heuristic_matches_full_sweep(monkeypatch): kernel's CTA tile / cluster shape, so that is not asserted. 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. + -- nvMMH may also pick rasterization orders or swizzles the sweep does not + enumerate -- 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. @@ -1602,15 +1604,13 @@ def _dominant_kernel(fn, iters=20): best_name, best_total, best_count = e.key, t, max(1, e.count) return best_name, best_total / best_count - # Full sweep: heuristics disabled. - monkeypatch.delenv("TRTLLM_CUTEDSL_NVMMH_ENABLE", raising=False) + # Full sweep: no nvMMH guidance. + tuner = AutoTuner.get() + monkeypatch.setattr(tuner, "nvmmh_config", None) sweep_tactic = _best_tactic() - # 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") + # nvMatmulHeuristics builds the candidates from its top 5 configurations. + monkeypatch.setattr(tuner, "nvmmh_config", NvMMHConfig(max_tactics=5)) heuristic_tactic = _best_tactic() # cuBLASLt runs its own heuristic auto-tuning; warm it under autotune(). @@ -1629,10 +1629,10 @@ def _cublas_call(): lambda: runner(inputs, tactic=heuristic_tactic)) _, cublas_us = _dominant_kernel(_cublas_call) - # 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. + # The guidance must not degrade the achieved kernel runtime beyond this + # tolerance. With the default max_tactics=5 the nvMMH 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. tolerance = 1.05 assert heuristic_us <= sweep_us * tolerance, ( f"heuristic-pruned tactic {heuristic_tactic} ({heuristic_us:.2f} us) is " @@ -1651,6 +1651,184 @@ def _cublas_call(): f"({cublas_us:.2f} us) for M={m}, N={n}, K={k}") +def _nvmmh_problem(case, sm_version): + """(runner, inputs, run, reference) of one nvMMH-aware CuTe DSL runner of + this GPU, or None if it has none for the case; run(tactic) returns that + tactic's output.""" + from _torch.helpers import per_block_cast_to_fp8, 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 as ops + + torch.manual_seed(0) + bf16 = torch.bfloat16 + rubin = sm_version == 107 + + def randn(*shape): + return torch.randn(*shape, device="cuda", dtype=bf16) + + def into(runner, inputs): + + def run(tactic): + inputs[-1].fill_(float("nan")) + runner(inputs, tactic=tactic) + return inputs[-1] + + return run + + def returned(runner, inputs): + return lambda tactic: runner(inputs, tactic=tactic) + + if case in ("bf16_gemm", "bf16_gemm_fp32_out"): + # nvMMH splits K of the FP32-output router GEMM and tiles the + # BF16-output decode GEMM narrowly. + fp32_out = case == "bf16_gemm_fp32_out" + m, n, k = (64, 256, 7168) if fp32_out else (64, 2112, 7168) + c_dtype = torch.float32 if fp32_out else bf16 + a, b = randn(m, k), randn(n, k) + runner_class = (ops.CuteDSLBf16RubinGemmRunner + if rubin else ops.CuteDSLBf16BlackwellGemmRunner) + runner = runner_class(output_dtype=c_dtype) + inputs = [a, b, torch.empty(m, n, device="cuda", dtype=c_dtype)] + return runner, inputs, into(runner, inputs), a.float() @ b.float().t() + if case in ("bf16_bmm", "bf16_bmm_fp8_out"): + fp8_out = case == "bf16_bmm_fp8_out" + if fp8_out and rubin: + return None + batch, m, n, k = 8, 300, 256, 512 + a, b = randn(batch, m, k), randn(batch, n, k) + if rubin: + runner = ops.CuteDSLBf16RubinBmmRunner() + elif fp8_out: + runner = ops.CuteDSLBf16BlackwellBmmFp8OutRunner() + else: + runner = ops.CuteDSLBf16BlackwellBmmRunner() + c_dtype = torch.float8_e4m3fn if fp8_out else bf16 + inputs = [a, b, torch.empty(batch, m, n, device="cuda", dtype=c_dtype)] + reference = torch.bmm(a.float(), b.float().transpose(1, 2)) + if fp8_out: + reference = reference.to(c_dtype).float() + return runner, inputs, into(runner, inputs), reference + if case == "fp8_gemm": + m, n, k = 64, 2048, 7168 + a, b = randn(m, k) / k, randn(n, k) / k + a_fp8, a_sf = torch.ops.trtllm.fp8_quantize_1x128(a) + b_fp8, b_sf = per_block_cast_to_fp8(b) + runner = (ops.CuteDSLFp8RubinGemmRunner() + if rubin else ops.CuteDSLFp8BlackwellRunner()) + inputs = [a_fp8, b_fp8, a_sf, b_sf] + reference = a.float() @ b.float().t() + return runner, inputs, returned(runner, inputs), reference + if case == "fp8_bmm": + # On SM100 nvMMH splits the 64 K tiles in 3 (unequal slices). + batch, m, n, k = 4, 64, 512, 8192 + a, b = randn(m, batch, k) / k, randn(batch, n, k) / k + a_fp8, a_sf = torch.ops.trtllm.fp8_batched_quantize_1x128_permute102(a) + b_fp8 = torch.empty_like(b, dtype=torch.float8_e4m3fn) + b_sf = torch.empty(batch, + n // 128, + k // 128, + device="cuda", + dtype=torch.float32) + for group in range(batch): + b_fp8[group], b_sf[group] = per_block_cast_to_fp8(b[group]) + runner = (ops.CuteDSLFp8RubinBmmRunner() + if rubin else ops.CuteDSLFp8BlackwellBmmRunner()) + inputs = [ + a_fp8, b_fp8, a_sf, b_sf, + torch.empty(batch, m, n, device="cuda", dtype=bf16) + ] + reference = torch.einsum("mgk,gnk->gmn", a.float(), b.float()) + return runner, inputs, into(runner, inputs), reference + if case == "mxfp8_gemm": + if not rubin: + return None + # A DeepSeek-V4 decode projection; nvMMH splits its 56 K tiles. + m, n, k = 32, 2048, 7168 + a, b = randn(m, k) / k, randn(n, k) / k + a_fp8, a_sf = torch.ops.trtllm.fp8_quantize_1x128_packed_ue8m0(a) + b_fp8, b_sf_k128 = per_block_cast_to_fp8_e8m0(b) + b_sf = fp8_utils.transform_k128_scales_to_cutedsl_mxfp8_layout( + b_sf_k128, mn=n, k=k) + alpha = ops._get_mxfp8_gemm_alpha(a.device) + runner = ops.CuteDSLMXFP8RubinLinear(output_dtype=bf16) + inputs = [a_fp8, b_fp8, a_sf, b_sf, alpha] + reference = a.float() @ b.float().t() + return runner, inputs, returned(runner, inputs), reference + assert case == "nvfp4_gemm" + # nvMMH splits the 64 K blocks in 2 and 3 on SM100. + m, n, k = 64, 2048, 16384 + x, w = randn(m, k), randn(n, k) + x_scale = 2688 / x.abs().max().float() + w_scale = 2688 / w.abs().max().float() + x_fp4, x_sf = torch.ops.trtllm.fp4_quantize(x, x_scale, 16, False) + w_fp4, w_sf = torch.ops.trtllm.fp4_quantize(w, w_scale, 16, False) + alpha = (1 / (x_scale * w_scale)).reshape(1) + inputs = [x_fp4, w_fp4, x_sf, w_sf, alpha] + runner = (ops.CuteDSLNVFP4RubinLinear(bf16) + if rubin else ops.CuteDSLNVFP4BlackwellRunner(bf16)) + reference = torch.ops.trtllm.nvfp4_gemm_cutlass(*inputs, bf16).float() + return runner, inputs, returned(runner, inputs), reference + + +@pytest.mark.parametrize( + "case", + [ + "bf16_gemm", + "bf16_gemm_fp32_out", + "bf16_bmm", + "bf16_bmm_fp8_out", + "fp8_gemm", + "fp8_bmm", + "mxfp8_gemm", + "nvfp4_gemm", + ], +) +def test_nvmmh_tactics_compute_the_product(case, monkeypatch): + """Every nvMMH-guided tactic of the SM100/SM103 or SM107 CuTe DSL + GEMM/BMM runners is correct, including tiles, rasters and split-K + factors their full sweeps never enumerate.""" + from tensorrt_llm._torch.custom_ops import \ + cutedsl_matmul_heuristics as nvmmh + from tensorrt_llm._torch.cute_dsl_utils import ( + IS_CUTLASS_DSL_AVAILABLE, IS_CUTLASS_DSL_RUBIN_AVAILABLE) + from tensorrt_llm._utils import get_sm_version + + sm_version = get_sm_version() + if not IS_CUTLASS_DSL_AVAILABLE or sm_version not in (100, 103, 107): + pytest.skip("The test requires SM100/SM103/SM107 and CuTe DSL.") + if sm_version == 107 and not IS_CUTLASS_DSL_RUBIN_AVAILABLE: + pytest.skip("The test requires the SM107 CuTe DSL helpers.") + if not nvmmh.IS_NVMMH_AVAILABLE: + pytest.skip("nvMatmulHeuristics is not installed") + from _torch.helpers import calc_diff + + problem = _nvmmh_problem(case, sm_version) + if problem is None: + pytest.skip(f"No nvMMH-aware {case} runner on SM{sm_version}.") + runner, inputs, run, reference = problem + selected = [] + select_tactics = nvmmh.select_tactics + + def recording_select_tactics(ranked, max_tactics): + selected.append(select_tactics(ranked, max_tactics)) + return selected[-1] + + monkeypatch.setattr(nvmmh, "select_tactics", recording_select_tactics) + monkeypatch.setattr(AutoTuner.get(), "nvmmh_config", + NvMMHConfig(max_tactics=3)) + tactics = runner.get_valid_tactics(inputs, None) + # The tactics come from nvMMH, not from the full-sweep fallback. + assert tactics and tactics == selected[-1] + # FP8 inputs are quantized from the BF16 reference operands. + max_diff = 1e-3 if case.startswith(("fp8", "mxfp8")) else 1e-4 + for tactic in tactics: + output = run(tactic).float() + torch.cuda.synchronize() + assert torch.isfinite(output).all(), tactic + assert calc_diff(output, reference) < max_diff, tactic + + @pytest.mark.parametrize("distribution", ["random", "balanced"]) def test_trtllm_gen_moe_dummy_topk_local_experts_less_than_topk( distribution, monkeypatch): diff --git a/tests/unittest/api_stability/references/llm.yaml b/tests/unittest/api_stability/references/llm.yaml index e3e0e82ddcaa..c765f42c1491 100644 --- a/tests/unittest/api_stability/references/llm.yaml +++ b/tests/unittest/api_stability/references/llm.yaml @@ -290,6 +290,10 @@ methods: annotation: bool default: True status: prototype + autotuner_nvmmh_config: + annotation: Optional[tensorrt_llm.llmapi.llm_args.AutoTunerNvMMHConfig] + default: null + status: prototype use_fine_grained_sync: annotation: bool default: False From b8c22140f09f620781ba8b3a6798bdd86011fbc5 Mon Sep 17 00:00:00 2001 From: peaceh <103117813+peaceh-nv@users.noreply.github.com> Date: Thu, 29 Oct 2026 20:06:33 +0800 Subject: [PATCH 2/3] [TRTLLM-14024][feat] Split K in the SM100 CuTe DSL GEMM full sweeps The full sweeps of the SM100 BF16 GEMM, NVFP4 GEMM and FP8 blockwise GEMM/BMM offer in-place split-K 2/4/8 for K >= 4096 and N <= 512, also with BF16 output, as nvMMH and the SM107 BF16 GEMM already do. The BF16 GEMM and the FP8 blockwise BMM split only while the output holds every row of A. The NVFP4 and FP8 blockwise sweeps split K only with the smallest cluster along M of each tile, and NVFP4 only with a single-column cluster and no prefetch. Signed-off-by: peaceh <103117813+peaceh-nv@users.noreply.github.com> --- .../_torch/custom_ops/cute_dsl_custom_ops.py | 61 ++++++++---- .../parallel/test_cute_dsl_bf16_dense.py | 4 +- .../kernels/parallel/test_fp4_linear.py | 39 ++++++++ .../parallel/test_fp8_block_scale_gemm.py | 94 +++++++++++++++++++ tests/unittest/_torch/misc/test_autotuner.py | 2 +- 5 files changed, 180 insertions(+), 20 deletions(-) diff --git a/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py b/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py index 5d95864a034d..a5e778afce2e 100644 --- a/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py +++ b/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py @@ -808,6 +808,14 @@ def _sm100_fp8_blockwise_mma_tiler_candidates( for tile_m in (64, 128, 256)] return tiles + def _sm100_split_k_candidates(n: int, k: int) -> List[int]: + """Full-sweep split-K factors of the SM100 dense GEMMs / BMMs whose + K slices reduce-add into the zeroed output. + + Split-K pays off for long K with few output tiles. + """ + return [1, 2, 4, 8] if k >= 4096 and n <= 512 else [1] + # SM100 BF16 and FP8 blockwise GEMMs support both rasters. The BF16 GEMM # and the FP8 blockwise GEMM/BMM split K: every slice reduce-adds its # partial into the zeroed output with TMA (no workspace, no reduction @@ -1039,6 +1047,8 @@ def get_valid_tactics( ] # prune: prefetch is beneficial only when K is large enough use_prefetch_candidates = [True, False] + # forward allocates the whole [M, N] output, so K may be split. + split_k_candidates = _sm100_split_k_candidates(n, real_k) valid_tactics = [] for mma_tiler_mn, cluster_shape_mn, swap_ab, use_prefetch in itertools.product( @@ -1081,8 +1091,16 @@ def get_valid_tactics( )): continue - valid_tactics.append( - (mma_tiler_mn, cluster_shape_mn, swap_ab, use_prefetch)) + # Split-K serves grids of few output tiles, so only the + # smallest cluster without prefetch splits K. + splits_k = (cluster_shape_mn == (mma_tiler_mn[0] // 128, 1) + and not use_prefetch) + valid_tactics.extend( + (mma_tiler_mn, cluster_shape_mn, swap_ab, + use_prefetch) + + ((1, True, split_k) if split_k > 1 else ()) + for split_k in ( + split_k_candidates if splits_k else [1])) logger.debug( f"CuteDSL: Found {len(valid_tactics)} valid tactics for M={m}, N={n}, K={real_k}" @@ -5094,8 +5112,8 @@ def get_valid_tactics( nvmmh_config = AutoTuner.get().nvmmh_config if nvmmh_config is not None: - # forward allocates the whole [M, N] output, so nvMMH may - # split K. + # forward allocates the whole [M, N] output, so K may be + # split. nvmmh_tactics = _sm100_nvmmh_tactics( self.__class__.kernel_class, cutlass.Float8E4M3FN, @@ -5125,8 +5143,10 @@ def get_valid_tactics( (4, 2), (4, 4), ] + split_k_candi = _sm100_split_k_candidates(n, k) return [ - (use_2cta_instrs, mma_tiler_mn, cluster_shape_mn) + (use_2cta_instrs, mma_tiler_mn, cluster_shape_mn) + + (("m", 1, split_k) if split_k > 1 else ()) for use_2cta_instrs in use_2cta_instrs_candi for mma_tiler_mn in mma_tiler_mn_candi for cluster_shape_mn in cluster_shape_mn_candi @@ -5144,7 +5164,11 @@ def get_valid_tactics( a_major, b_major, c_major, - ) + ) for split_k in split_k_candi + # Split-K serves grids of few output tiles, so only the + # smallest cluster M splits K. + if split_k == 1 or cluster_shape_mn[0] == ( + 2 if use_2cta_instrs else 1) ] def forward( @@ -5471,8 +5495,11 @@ def get_valid_tactics( (4, 2), (4, 4), ] + split_k_candi = (_sm100_split_k_candidates(n, k) + if inputs[4].shape[1] == m else [1]) return [ - (use_2cta_instrs, mma_tiler_mn, cluster_shape_mn) + (use_2cta_instrs, mma_tiler_mn, cluster_shape_mn) + + (("m", 1, split_k) if split_k > 1 else ()) for use_2cta_instrs in use_2cta_instrs_candi for mma_tiler_mn in mma_tiler_mn_candi for cluster_shape_mn in cluster_shape_mn_candi @@ -5490,7 +5517,11 @@ def get_valid_tactics( a_major, b_major, c_major, - ) + ) for split_k in split_k_candi + # Split-K serves grids of few output tiles, so only the + # smallest cluster M splits K. + if split_k == 1 or cluster_shape_mn[0] == ( + 2 if use_2cta_instrs else 1) ] def forward( @@ -10237,9 +10268,8 @@ def __init__(self, output_dtype: Optional[torch.dtype] = None): super().__init__() self.use_tvm_ffi = use_tvm_ffi - # The tactic space depends on the output dtype (the full sweep - # splits K only for FP32 output, and nvMMH models the output - # type), so unique_id() keys the profiling cache on it. + # The best tactic and the nvMMH ranking depend on the output + # dtype, so unique_id() keys the profiling cache on it. self.output_dtype = output_dtype def unique_id(self): @@ -10302,13 +10332,8 @@ def get_valid_tactics( (4, 2), (4, 4), ] - # Split-K pays off for long K with few output tiles. The full - # sweep offers it only for FP32 output, where the TMA reduce-add - # of the slices does not round partial sums to BF16. - split_k_candi = [1] - if (allow_split_k and c_dtype_cutlass == cutlass.Float32 - and k >= 4096 and n <= 512): - split_k_candi = [1, 2, 4, 8] + split_k_candi = (_sm100_split_k_candidates(n, k) + if allow_split_k else [1]) return [ (use_2cta_instrs, mma_tiler_mn, cluster_shape_mn) + (("m", 1, split_k) if split_k > 1 else ()) diff --git a/tests/unittest/_torch/kernels/parallel/test_cute_dsl_bf16_dense.py b/tests/unittest/_torch/kernels/parallel/test_cute_dsl_bf16_dense.py index 798f70d1b573..d8dc5c7b5b46 100644 --- a/tests/unittest/_torch/kernels/parallel/test_cute_dsl_bf16_dense.py +++ b/tests/unittest/_torch/kernels/parallel/test_cute_dsl_bf16_dense.py @@ -268,10 +268,12 @@ def test_cute_dsl_bf16_split_k_gemm_blackwell(split_k_slices, c_dtype): from tensorrt_llm._torch.custom_ops.cute_dsl_custom_ops import CuteDSLBf16BlackwellGemmRunner torch.manual_seed(2026) - runner = CuteDSLBf16BlackwellGemmRunner() + runner = CuteDSLBf16BlackwellGemmRunner(output_dtype=c_dtype) act, weight, output = _gemm_operands(64, 256, 7168, c_dtype) base = (False, (128, 128), (1, 1)) tactic = base + ("m", 1, split_k_slices) + # Large K and small N, so the full sweep offers split-K for either output dtype. + assert tactic in runner.get_valid_tactics([act, weight, output], None) # Direct split-K rounds each partial to the output dtype before the # atomic TMA ADD, so BF16 output needs tolerance for both rounding and diff --git a/tests/unittest/_torch/kernels/parallel/test_fp4_linear.py b/tests/unittest/_torch/kernels/parallel/test_fp4_linear.py index 3838273e04be..7da9f53823ff 100644 --- a/tests/unittest/_torch/kernels/parallel/test_fp4_linear.py +++ b/tests/unittest/_torch/kernels/parallel/test_fp4_linear.py @@ -17,6 +17,7 @@ import pytest import torch +from _torch.helpers import calc_diff from utils.util import skip_pre_blackwell import tensorrt_llm.quantization.utils.fp4_utils as fp4_utils @@ -1020,6 +1021,44 @@ def test_nvfp4_gemm_cutedsl(m, with_bias, monkeypatch, tmp_path): tuner.clear_cache() +@pytest.mark.skipif(not IS_CUTLASS_DSL_AVAILABLE + or get_sm_version() not in (100, 103), + reason="Requires SM100/SM103 CuTe DSL") +@pytest.mark.parametrize("m", [64, 256]) +@pytest.mark.parametrize("split_k", [2, 4, 8]) +def test_nvfp4_gemm_cutedsl_split_k(m, split_k): + """The full sweep splits long K for small N with and without swapping A + and B, and every split launch computes the product, also on graph replay.""" + from tensorrt_llm._torch.custom_ops import cute_dsl_custom_ops + + torch.manual_seed(0) + dtype = torch.bfloat16 + n, k = 512, 7168 + x = torch.randn(m, k, dtype=dtype, device="cuda") + w = torch.randn(n, k, dtype=dtype, device="cuda") + x_scale = 2688 / x.abs().max().float() + w_scale = 2688 / w.abs().max().float() + x_fp4, x_sf = torch.ops.trtllm.fp4_quantize(x, x_scale, 16, False) + w_fp4, w_sf = torch.ops.trtllm.fp4_quantize(w, w_scale, 16, False) + alpha = (1 / (x_scale * w_scale)).reshape(1) + inputs = [x_fp4, w_fp4, x_sf, w_sf, alpha] + reference = torch.ops.trtllm.nvfp4_gemm_cutlass(*inputs, dtype) + + runner = cute_dsl_custom_ops.CuteDSLNVFP4BlackwellRunner(dtype) + swap_ab = m <= 128 + tactic = ((128, 128), (1, 1), swap_ab, False, 1, True, split_k) + assert tactic in runner.get_valid_tactics(inputs, None) + with torch.inference_mode(): + assert calc_diff(runner(inputs, tactic=tactic), reference) < 1e-4 + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + output = runner(inputs, tactic=tactic) + output.fill_(float("nan")) + graph.replay() + torch.cuda.synchronize() + assert calc_diff(output, reference) < 1e-4 + + def _skip_if_no_locality_domain(): if not IS_CUTLASS_DSL_RUBIN_AVAILABLE: pytest.skip("CuTe DSL package with Rubin support is not available") 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 b23987dbff6a..6c112128a4d4 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 @@ -404,6 +404,100 @@ def cache_profile(m: int) -> tuple[tuple[int, ...], ...]: assert cache_profile(5) == cache_profile(7) +def _k_block_magnitudes(k: int) -> torch.Tensor: + """Per-128-column magnitudes, so a K slice that applied the block scales + of another slice would not compute the product.""" + exponents = torch.randint(-3, 4, (k // 128, ), device='cuda') + return torch.exp2(exponents.float()).repeat_interleave(128) + + +@pytest.mark.skipif( + getSMVersion() not in (100, 103), + reason="The test is for SM100 and SM103 only. Current SM is %d." % + getSMVersion(), +) +@pytest.mark.parametrize("split_k", [2, 4, 8]) +def test_cute_dsl_fp8_block_scale_gemm_split_k(split_k): + """The full sweep splits long K for small N, and every split launch + (also inside a CUDA graph) computes the product.""" + torch.random.manual_seed(0) + m, n, k = 64, 512, 7168 + a = (torch.randn( + (m, k), device='cuda') * _k_block_magnitudes(k) / k).to(torch.bfloat16) + b = (torch.randn( + (n, k), device='cuda') * _k_block_magnitudes(k) / k).to(torch.bfloat16) + a_fp8, a_sf = torch.ops.trtllm.fp8_quantize_1x128(a) + b_fp8, b_sf = per_block_cast_to_fp8(b) + inputs = [a_fp8, b_fp8, a_sf, b_sf] + expected = a.float() @ b.float().t() + + runner = cute_dsl_custom_ops.CuteDSLFp8BlackwellRunner() + tactic = (False, (128, 128), (1, 1), "m", 1, split_k) + assert tactic in runner.get_valid_tactics(inputs, None) + + assert calc_diff(runner(inputs, tactic=tactic), expected) < 1e-3 + + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + output = runner(inputs, tactic=tactic) + output.fill_(float("nan")) + graph.replay() + torch.cuda.synchronize() + assert calc_diff(output, expected) < 1e-3 + + +@pytest.mark.skipif( + getSMVersion() not in (100, 103), + reason="The test is for SM100 and SM103 only. Current SM is %d." % + getSMVersion(), +) +@pytest.mark.parametrize("split_k", [2, 4, 8]) +def test_cute_dsl_fp8_block_scale_bmm_split_k(split_k): + """The full sweep splits long K for small N when the output holds exactly + M rows, and every split launch zeroes the output first.""" + torch.random.manual_seed(0) + num_groups, m, n, k = 4, 64, 512, 7168 + a = (torch.randn((m, num_groups, k), device='cuda') * + _k_block_magnitudes(k) / k).to(torch.bfloat16) + b = (torch.randn((num_groups, n, k), device='cuda') * + _k_block_magnitudes(k) / k).to(torch.bfloat16) + a_fp8, a_scales = torch.ops.trtllm.fp8_batched_quantize_1x128_permute102(a) + b_fp8 = torch.empty_like(b, dtype=torch.float8_e4m3fn) + b_scales = torch.empty((num_groups, n // 128, k // 128), + device='cuda', + dtype=torch.float) + for i in range(num_groups): + b_fp8[i], b_scales[i] = per_block_cast_to_fp8(b[i]) + output = torch.empty((num_groups, m, n), + device='cuda', + dtype=torch.bfloat16) + inputs = [a_fp8, b_fp8, a_scales, b_scales, output] + expected = torch.einsum('mgk,gnk->gmn', a.float(), b.float()) + + runner = cute_dsl_custom_ops.CuteDSLFp8BlackwellBmmRunner() + tactic = (False, (128, 128), (1, 1), "m", 1, split_k) + assert tactic in runner.get_valid_tactics(inputs, None) + taller_output = torch.empty((num_groups, m + 8, n), + device='cuda', + dtype=torch.bfloat16) + assert all( + len(t) < 6 + for t in runner.get_valid_tactics(inputs[:4] + [taller_output], None)) + + output.fill_(float("nan")) + runner(inputs, tactic=tactic) + torch.cuda.synchronize() + assert calc_diff(output, expected) < 1e-3 + + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + runner(inputs, tactic=tactic) + output.fill_(float("nan")) + graph.replay() + torch.cuda.synchronize() + assert calc_diff(output, expected) < 1e-3 + + @pytest.mark.skipif( getSMVersion() != 107 or not IS_CUTLASS_DSL_RUBIN_AVAILABLE, reason="The test requires SM107 and SM107 CuTe DSL support.", diff --git a/tests/unittest/_torch/misc/test_autotuner.py b/tests/unittest/_torch/misc/test_autotuner.py index 815ad4f39d0c..efbba2697f5d 100644 --- a/tests/unittest/_torch/misc/test_autotuner.py +++ b/tests/unittest/_torch/misc/test_autotuner.py @@ -1787,7 +1787,7 @@ def returned(runner, inputs): def test_nvmmh_tactics_compute_the_product(case, monkeypatch): """Every nvMMH-guided tactic of the SM100/SM103 or SM107 CuTe DSL GEMM/BMM runners is correct, including tiles, rasters and split-K - factors their full sweeps never enumerate.""" + factors outside their full sweeps.""" from tensorrt_llm._torch.custom_ops import \ cutedsl_matmul_heuristics as nvmmh from tensorrt_llm._torch.cute_dsl_utils import ( From 5e0622b44dc953a4206278a74735314497183485 Mon Sep 17 00:00:00 2001 From: peaceh <103117813+peaceh-nv@users.noreply.github.com> Date: Fri, 30 Oct 2026 07:47:53 +0800 Subject: [PATCH 3/3] [TRTLLM-14024][fix] Report nvMMH fallbacks, trim BF16 full-sweep split-K and extend the CuTe DSL GEMM tests Tactics and checks: - The full sweep of the SM100 BF16 GEMM offers split-K only with the smallest cluster along M of each tile, like the FP8 blockwise GEMM/BMM sweeps. This cuts the sweep from 120 to 66 tactics for split-K shapes. - Log the shapes nvMMH returns no tactic for, which run the full sweep: once per runner at info level, each shape at debug level. - Warn once when autotuner_nvmmh_config is set but nvMatmulHeuristics cannot be imported. - Treat an nvMMH installation without version metadata as a 0.1.0.x release only if it lacks the per-backend silicon-metric API. - Reject split_k > 1 in the SM100 block-scaled GEMM unless the output is FP32, FP16 or BF16. - Describe max_tactics as a count of kernel configurations, and regenerate the serving schemas. Tests: - Pure-Python AutoTuner tests: tactic_space_key() extends the profiling cache key, survives a cache save and load, and follows the nvMMH policy; NVFP4GemmUnifiedRunner keys only its CuTe DSL tactic space and LocalityDomainConcurrentTunableRunner forwards the key; the engine installs its nvMMH policy; rankings merge by score; select_tactics counts configs that add new tactics; and AutoTunerNvMMHConfig rejects non-positive or non-integer max_tactics. - The SM100 FP8 blockwise GEMM and BMM run every MMA tile of their full sweeps, including the N = 32/64 tiles inside a 128-row scale block of B, against a dequantized reference whose B scale blocks differ in magnitude along N. - Every nvMMH-guided tactic of the SM107 FP8 BMM with fused 1x128 quantization matches the BF16 BMM of the same tactic quantized separately. - CuTe DSL split-K tactics round each K slice's partial sum to BF16 before reducing it into the output. Where the CuTe DSL backend can split K (SM100/SM103 with K >= 4096 and N <= 512, every shape on SM107), test_nvfp4_gemm_unified_all_tactics compares with atol = 0.05 x the RMS of the reference instead of 0.15. - Run test_nvmmh_tactics_compute_the_product pre-merge on B200 and GB300, and list the SM107 CuTe DSL GEMM/BMM tests in l0_gr100. Signed-off-by: peaceh <103117813+peaceh-nv@users.noreply.github.com> --- .../_torch/custom_ops/cute_dsl_custom_ops.py | 27 +++ .../custom_ops/cutedsl_matmul_heuristics.py | 22 +- .../dense_blockscaled_gemm_persistent.py | 4 + .../_torch/pyexecutor/model_engine.py | 20 +- tensorrt_llm/llmapi/llm_args.py | 7 +- .../schemas/trtllm-serve-config.schema.json | 4 +- .../trtllm-serve-disagg-config.schema.json | 4 +- .../test_lists/test-db/l0_b200.yml | 1 + .../test_lists/test-db/l0_gb300.yml | 1 + .../test_lists/test-db/l0_gr100.yml | 9 + .../kernels/parallel/test_fp4_linear.py | 14 +- .../parallel/test_fp8_block_scale_gemm.py | 188 +++++++++++++++++- tests/unittest/_torch/misc/test_autotuner.py | 169 ++++++++++++++++ 13 files changed, 455 insertions(+), 15 deletions(-) diff --git a/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py b/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py index a5e778afce2e..7d4b1981855e 100644 --- a/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py +++ b/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py @@ -722,6 +722,17 @@ def _merge_nvmmh_rankings(rankings: List[List[Tuple]]) -> List[Tuple]: heads[best] += 1 +def _log_nvmmh_fallback(runner: str, m: int, n: int, k: int, + batch_size: int) -> None: + """Log a shape nvMMH gave no tactic for, so it runs the full sweep.""" + logger.debug(f"CuteDSL: nvMMH returned no tactic for {runner} M={m}, " + f"N={n}, K={k}, batch={batch_size}; using the full sweep.") + logger.info_once( + f"CuteDSL: nvMMH returned no tactic for some {runner} shapes, which " + "use the full tactic sweep; debug logging lists the shapes.", + key=f"nvmmh_fallback_{runner}") + + # The Blackwell block-scaled GEMM reduces split-K partials with TMA reduce-add # into the zeroed output and supports both rasters. _SM100_BLOCKSCALED_NVMMH_BACKEND = cutedsl_nvmmh.Backend( @@ -1022,6 +1033,7 @@ def get_valid_tactics( nvmmh_config.max_tactics) if nvmmh_tactics: return nvmmh_tactics + _log_nvmmh_fallback(type(self).__name__, m, n, real_k, 1) # full shamoo mma_tiler_mn_candidates = [ @@ -5129,6 +5141,7 @@ def get_valid_tactics( allow_split_k=True) if nvmmh_tactics: return nvmmh_tactics + _log_nvmmh_fallback(type(self).__name__, m, n, k, batch_size) use_2cta_instrs_candi = [False, True] mma_tiler_mn_candi = _sm100_fp8_blockwise_mma_tiler_candidates(m) @@ -5481,6 +5494,7 @@ def get_valid_tactics( allow_split_k=inputs[4].shape[1] == m) if nvmmh_tactics: return nvmmh_tactics + _log_nvmmh_fallback(type(self).__name__, m, n, k, batch_size) use_2cta_instrs_candi = [False, True] mma_tiler_mn_candi = _sm100_fp8_blockwise_mma_tiler_candidates(m) @@ -9871,6 +9885,7 @@ def get_valid_tactics( k, batch_size, None, 64, nvmmh_config.max_tactics) if nvmmh_tactics: return nvmmh_tactics + _log_nvmmh_fallback(type(self).__name__, m, n, k, batch_size) use_2cta_instrs_candi = [False, True] mma_tiler_mn_candi = [(64, 128), (128, 128), (256, 128)] @@ -10318,6 +10333,7 @@ def get_valid_tactics( allow_split_k) if nvmmh_tactics: return nvmmh_tactics + _log_nvmmh_fallback(type(self).__name__, m, n, k, batch_size) use_2cta_instrs_candi = [False, True] mma_tiler_mn_candi = [(64, 128), (128, 128), (256, 128)] @@ -10355,6 +10371,10 @@ def get_valid_tactics( b_major, c_major, ) for split_k in split_k_candi + # Split-K serves grids of few output tiles, so only the + # smallest cluster M splits K. + if split_k == 1 or cluster_shape_mn[0] == ( + 2 if use_2cta_instrs else 1) ] def forward( @@ -10907,6 +10927,7 @@ def get_valid_tactics( nvmmh_config.max_tactics) if nvmmh_tactics: return nvmmh_tactics + _log_nvmmh_fallback(type(self).__name__, m, n, k, batch_size) return _sm107_bf16_valid_tactics( m, n, k, batch_size, cutlass.BFloat16, _SM107_BF16_MMA_TILER_MN_CANDIDATES, None) @@ -11426,6 +11447,7 @@ def get_valid_tactics( nvmmh_config.max_tactics) if nvmmh_tactics: return nvmmh_tactics + _log_nvmmh_fallback(type(self).__name__, m, n, k, 1) mma_tiler_mn_candi = list(_SM107_BF16_MMA_TILER_MN_CANDIDATES) # Skinny-M / small-N shapes under-use a 128x128 MMA tile (half the @@ -13600,6 +13622,8 @@ def get_valid_tactics( cutlass.BFloat16, nvmmh_config.max_tactics) if nvmmh_tactics: return nvmmh_tactics + _log_nvmmh_fallback( + type(self).__name__, m, n, k, batch_size) return [(use_2cta_instrs, mma_tiler, mma_inst_shape, cluster_shape_mn) for use_2cta_instrs, mma_tiler, @@ -13851,6 +13875,8 @@ def get_valid_tactics( nvmmh_config.max_tactics, self.fp8_quantize_1x128) if nvmmh_tactics: return nvmmh_tactics + _log_nvmmh_fallback( + type(self).__name__, m, n, k, batch_size) return [(use_2cta_instrs, mma_tiler, mma_inst_shape, cluster_shape_mn) for use_2cta_instrs, mma_tiler, @@ -14534,6 +14560,7 @@ def get_valid_tactics( m, n, real_k, [False, True], nvmmh_config.max_tactics) if nvmmh_tactics: return nvmmh_tactics + _log_nvmmh_fallback(type(self).__name__, m, n, real_k, 1) # SM107 dense block-scaled tile sizes. # b-reuse is enabled when mma_tiler_m = 2 * mma_inst_shape_m diff --git a/tensorrt_llm/_torch/custom_ops/cutedsl_matmul_heuristics.py b/tensorrt_llm/_torch/custom_ops/cutedsl_matmul_heuristics.py index 3a4b5a0c314e..ec9c69fa5a42 100644 --- a/tensorrt_llm/_torch/custom_ops/cutedsl_matmul_heuristics.py +++ b/tensorrt_llm/_torch/custom_ops/cutedsl_matmul_heuristics.py @@ -135,8 +135,26 @@ def cta_order_to_raster(cta_order: int) -> str: @lru_cache(maxsize=1) def _pairs_only_cluster_m_2() -> bool: - """0.1.0.x PyPI releases model 2-CTA MMA for cluster M == 2 only.""" - match = re.fullmatch(r"0\.1\.0\.(\d+)", library_version()) + """0.1.0.x PyPI releases model 2-CTA MMA for cluster M == 2 only. + + An installation without version metadata is a 0.1.0.x release if its + interface lacks the per-backend silicon-metric API newer builds add. + """ + version = library_version() + if version == "unknown": + interface = _interface_class() + old = not ( + hasattr(interface, "estimateSiliconMetricEx") + and hasattr(interface, "setBackendPropertyValue") + ) + logger.warning_once( + "nvMatmulHeuristics reports no version; assuming " + f"{'a 0.1.0.x release' if old else 'a libheuristics build'} " + "from its API.", + key="nvmmh_unknown_version", + ) + return old + match = re.fullmatch(r"0\.1\.0\.(\d+)", version) return match is not None and int(match.group(1)) < 10000 diff --git a/tensorrt_llm/_torch/kernels/blackwell/dense_blockscaled_gemm_persistent.py b/tensorrt_llm/_torch/kernels/blackwell/dense_blockscaled_gemm_persistent.py index bc67b5304344..d75c8470ee38 100644 --- a/tensorrt_llm/_torch/kernels/blackwell/dense_blockscaled_gemm_persistent.py +++ b/tensorrt_llm/_torch/kernels/blackwell/dense_blockscaled_gemm_persistent.py @@ -479,6 +479,10 @@ def __call__( "The swapped indexer-Q epilogue requires a column-major BF16 " "GEMM boundary plus packed output, scale, position, and RoPE tensors" ) + if cutlass.const_expr( + self.split_k > 1 and self.c_dtype + not in (cutlass.Float32, cutlass.Float16, cutlass.BFloat16)): + raise ValueError("split_k > 1 needs an FP32, FP16 or BF16 output") # Check if input data types are compatible with MMA instruction if cutlass.const_expr(self.a_dtype != self.b_dtype): diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py index e3cea3538ecd..d022b127ae3a 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_engine.py +++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py @@ -126,13 +126,25 @@ def _configure_autotuner_nvmmh(llm_args: TorchLlmArgs) -> None: """Install the engine's nvMatmulHeuristics policy on the process AutoTuner. Runs before model construction because CuTe DSL runners may enumerate - tactics while modules are built. Engines of one process (e.g. target and - draft) share the LLM args and therefore the policy. + tactics while modules are built. The policy is process-wide, so engines + that share a process should use the same config; target and draft engines + share the LLM args and therefore the policy. """ # Partial test doubles of TorchLlmArgs may omit the field. config = getattr(llm_args, "autotuner_nvmmh_config", None) - AutoTuner.get().nvmmh_config = (None if config is None else NvMMHConfig( - max_tactics=config.max_tactics)) + if config is not None: + from ..custom_ops import cutedsl_matmul_heuristics + + if not cutedsl_matmul_heuristics.IS_NVMMH_AVAILABLE: + logger.warning_once( + "autotuner_nvmmh_config is set but nvMatmulHeuristics cannot " + "be imported; CuTe DSL GEMM/BMM runners tune their full " + "tactic space.", + key="autotuner_nvmmh_unavailable", + ) + AutoTuner.get().nvmmh_config = ( + None if config is None else NvMMHConfig(max_tactics=config.max_tactics) + ) class ModelEngine: diff --git a/tensorrt_llm/llmapi/llm_args.py b/tensorrt_llm/llmapi/llm_args.py index 6d4f3d6ae958..7a1c953dff8a 100644 --- a/tensorrt_llm/llmapi/llm_args.py +++ b/tensorrt_llm/llmapi/llm_args.py @@ -2011,12 +2011,15 @@ class AutoTunerNvMMHConfig(StrictBaseModel): nvMatmulHeuristics ranks highest for each problem shape, taking every field it models (CTA tile, cluster shape, split-K, rasterization order and swizzle) as returned, instead of profiling their full tactic space. + ``max_tactics`` counts kernel configurations, not profiled tactics. """ max_tactics: PositiveInt = Field( default=5, description= - "Number of top-ranked nvMatmulHeuristics kernel configurations profiled per GEMM/BMM shape." - ) + "Number of top-ranked nvMatmulHeuristics kernel configurations kept per GEMM/BMM shape. " + "Each configuration expands into one tactic per combination of the knobs nvMatmulHeuristics " + "does not model (e.g. SM107 prefetch, scheduler and mixed-cluster variants), so the number " + "of profiled tactics can be several times max_tactics.") class AttentionDpConfig(StrictBaseModel): diff --git a/tensorrt_llm/schemas/trtllm-serve-config.schema.json b/tensorrt_llm/schemas/trtllm-serve-config.schema.json index cbcdb3582e15..2a5d21b50270 100644 --- a/tensorrt_llm/schemas/trtllm-serve-config.schema.json +++ b/tensorrt_llm/schemas/trtllm-serve-config.schema.json @@ -349,11 +349,11 @@ }, "AutoTunerNvMMHConfig": { "additionalProperties": false, - "description": "nvMatmulHeuristics (nvMMH) guidance for the CuTe DSL GEMM/BMM autotuner.\n\nCuTe DSL GEMM and BMM runners profile only the kernel configurations that\nnvMatmulHeuristics ranks highest for each problem shape, taking every\nfield it models (CTA tile, cluster shape, split-K, rasterization order and\nswizzle) as returned, instead of profiling their full tactic space.", + "description": "nvMatmulHeuristics (nvMMH) guidance for the CuTe DSL GEMM/BMM autotuner.\n\nCuTe DSL GEMM and BMM runners profile only the kernel configurations that\nnvMatmulHeuristics ranks highest for each problem shape, taking every\nfield it models (CTA tile, cluster shape, split-K, rasterization order and\nswizzle) as returned, instead of profiling their full tactic space.\n``max_tactics`` counts kernel configurations, not profiled tactics.", "properties": { "max_tactics": { "default": 5, - "description": "Number of top-ranked nvMatmulHeuristics kernel configurations profiled per GEMM/BMM shape.", + "description": "Number of top-ranked nvMatmulHeuristics kernel configurations kept per GEMM/BMM shape. Each configuration expands into one tactic per combination of the knobs nvMatmulHeuristics does not model (e.g. SM107 prefetch, scheduler and mixed-cluster variants), so the number of profiled tactics can be several times max_tactics.", "exclusiveMinimum": 0, "title": "Max Tactics", "type": "integer" diff --git a/tensorrt_llm/schemas/trtllm-serve-disagg-config.schema.json b/tensorrt_llm/schemas/trtllm-serve-disagg-config.schema.json index aa6ad717d05e..3917b8cbd743 100644 --- a/tensorrt_llm/schemas/trtllm-serve-disagg-config.schema.json +++ b/tensorrt_llm/schemas/trtllm-serve-disagg-config.schema.json @@ -349,11 +349,11 @@ }, "AutoTunerNvMMHConfig": { "additionalProperties": false, - "description": "nvMatmulHeuristics (nvMMH) guidance for the CuTe DSL GEMM/BMM autotuner.\n\nCuTe DSL GEMM and BMM runners profile only the kernel configurations that\nnvMatmulHeuristics ranks highest for each problem shape, taking every\nfield it models (CTA tile, cluster shape, split-K, rasterization order and\nswizzle) as returned, instead of profiling their full tactic space.", + "description": "nvMatmulHeuristics (nvMMH) guidance for the CuTe DSL GEMM/BMM autotuner.\n\nCuTe DSL GEMM and BMM runners profile only the kernel configurations that\nnvMatmulHeuristics ranks highest for each problem shape, taking every\nfield it models (CTA tile, cluster shape, split-K, rasterization order and\nswizzle) as returned, instead of profiling their full tactic space.\n``max_tactics`` counts kernel configurations, not profiled tactics.", "properties": { "max_tactics": { "default": 5, - "description": "Number of top-ranked nvMatmulHeuristics kernel configurations profiled per GEMM/BMM shape.", + "description": "Number of top-ranked nvMatmulHeuristics kernel configurations kept per GEMM/BMM shape. Each configuration expands into one tactic per combination of the knobs nvMatmulHeuristics does not model (e.g. SM107 prefetch, scheduler and mixed-cluster variants), so the number of profiled tactics can be several times max_tactics.", "exclusiveMinimum": 0, "title": "Max Tactics", "type": "integer" diff --git a/tests/integration/test_lists/test-db/l0_b200.yml b/tests/integration/test_lists/test-db/l0_b200.yml index e8e1f903dea9..3d8079ccf945 100644 --- a/tests/integration/test_lists/test-db/l0_b200.yml +++ b/tests/integration/test_lists/test-db/l0_b200.yml @@ -361,6 +361,7 @@ l0_b200: - unittest/_torch/moe/kernels/parallel - unittest/_torch/kernels/parallel TIMEOUT (90) - unittest/_torch/visual_gen/kernels/parallel + - unittest/_torch/misc/test_autotuner.py::test_nvmmh_tactics_compute_the_product - unittest/_torch/moe/kernels/serial - unittest/_torch/kernels/serial - unittest/_torch/modeling -k "modeling_llama" diff --git a/tests/integration/test_lists/test-db/l0_gb300.yml b/tests/integration/test_lists/test-db/l0_gb300.yml index ce21a01c9ddc..639f5d316d0f 100644 --- a/tests/integration/test_lists/test-db/l0_gb300.yml +++ b/tests/integration/test_lists/test-db/l0_gb300.yml @@ -45,6 +45,7 @@ l0_gb300: - accuracy/test_disaggregated_serving.py::TestQwen3_8_Flash_Next::test_fp8_nixl_python[prefix_cache] - unittest/_torch/moe/kernels/parallel - unittest/_torch/kernels/parallel TIMEOUT (90) + - unittest/_torch/misc/test_autotuner.py::test_nvmmh_tactics_compute_the_product # Kimi K3 KDA unit tests (single GPU, random weights, no checkpoint). - unittest/_torch/modules/kimi_kda/test_kda_flashinfer_prefill.py - unittest/_torch/modules/kimi_kda/test_kda_fused_prefill_conv.py diff --git a/tests/integration/test_lists/test-db/l0_gr100.yml b/tests/integration/test_lists/test-db/l0_gr100.yml index a9cbbc03f821..9370fb161f8e 100644 --- a/tests/integration/test_lists/test-db/l0_gr100.yml +++ b/tests/integration/test_lists/test-db/l0_gr100.yml @@ -17,6 +17,15 @@ l0_gr100: - unittest/_torch/attention/test_flashinfer_trtllm_gen_sm107_p128.py - unittest/_torch/modeling/test_kimi_k3_mla.py::test_k3_mla_nvfp4_kv_with_fp8_kv_b_matches_huggingface - unittest/_torch/modules/kimi_kda/test_kda_flashinfer_prefill.py::test_auto_prefill_stays_finite_on_repeated_token_with_mixed_decay + # CuTe DSL GEMM/BMM kernels and their nvMatmulHeuristics-guided tactics. + - unittest/_torch/kernels/parallel/test_cute_dsl_bf16_dense.py + - unittest/_torch/kernels/parallel/test_fp4_linear.py::test_nvfp4_gemm_cutedsl + - unittest/_torch/kernels/parallel/test_fp8_block_scale_gemm.py::test_cute_dsl_fp8_block_scale_gemm_rubin + - unittest/_torch/kernels/parallel/test_fp8_block_scale_gemm.py::test_cute_dsl_fp8_block_scale_bmm_rubin + - unittest/_torch/kernels/parallel/test_fp8_block_scale_gemm.py::test_cute_dsl_fp8_block_scale_bmm_rubin_tile_shapes + - unittest/_torch/kernels/parallel/test_fp8_block_scale_gemm.py::test_cute_dsl_fp8_bmm_quantize_rubin_matches_separate + - unittest/_torch/kernels/parallel/test_fp8_block_scale_gemm.py::test_cute_dsl_fp8_bmm_quantize_rubin_nvmmh_tactics + - unittest/_torch/misc/test_autotuner.py::test_nvmmh_tactics_compute_the_product # The GLM-5.2 target's unit tests, one entry per op directory. Its comm tests # run 4 ranks, and CI has no 4-GPU GR100 stage yet. - unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/attention diff --git a/tests/unittest/_torch/kernels/parallel/test_fp4_linear.py b/tests/unittest/_torch/kernels/parallel/test_fp4_linear.py index 7da9f53823ff..87b43eee18ad 100644 --- a/tests/unittest/_torch/kernels/parallel/test_fp4_linear.py +++ b/tests/unittest/_torch/kernels/parallel/test_fp4_linear.py @@ -508,9 +508,19 @@ def test_nvfp4_gemm_unified_all_tactics(dtype, mnk): AutoTuner.get().print_profiling_cache() + # CuTe DSL split-K tactics round each K slice's partial sum to BF16 before + # reducing it into the output, so their error is a few BF16 ulps of the + # partial sums, which scale with the output RMS rather than per element. + sm_version = get_sm_version() + if sm_version == 107 or (sm_version in (100, 103) and HIDDEN_SIZE >= 4096 + and OUTPUT_SIZE <= 512): + atol = 0.05 * output_ref.float().pow(2).mean().sqrt().item() + else: + atol = 0.15 + # Verify auto mode result matches reference torch.cuda.synchronize() - torch.testing.assert_close(output_auto, output_ref, rtol=1e-2, atol=0.15) + torch.testing.assert_close(output_auto, output_ref, rtol=1e-2, atol=atol) # Test all combinations of outer layer (backend selection) and inner layer (backend tactics) # Outer layer: nvfp4_gemm selects backend @@ -611,7 +621,7 @@ def test_nvfp4_gemm_unified_all_tactics(dtype, mnk): torch.testing.assert_close(output, output_ref, rtol=1e-2, - atol=0.15) + atol=atol) total_tactics_tested += 1 tactics_by_backend[runner_name].append(total_tactics_tested) 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 6c112128a4d4..530d78784076 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 @@ -24,7 +24,7 @@ from utils.util import getSMVersion, isSM100Family, isSM120Family import tensorrt_llm.quantization.utils.fp8_utils as fp8_utils -from tensorrt_llm._torch.autotuner import AutoTuner, autotune +from tensorrt_llm._torch.autotuner import AutoTuner, NvMMHConfig, autotune from tensorrt_llm._torch.custom_ops import cute_dsl_custom_ops from tensorrt_llm._torch.cute_dsl_utils import (IS_CUTLASS_DSL_AVAILABLE, IS_CUTLASS_DSL_RUBIN_AVAILABLE) @@ -498,6 +498,120 @@ def test_cute_dsl_fp8_block_scale_bmm_split_k(split_k): assert calc_diff(output, expected) < 1e-3 +def _n_block_magnitudes(n: int, k: int) -> torch.Tensor: + """Per-128x128-block magnitudes of B that differ between the N blocks of + every K block, so a tile that applied the scale of another N block would + not compute the product.""" + n_blocks = torch.arange(n // 128, device='cuda')[:, None] + k_blocks = torch.arange(k // 128, device='cuda')[None, :] + magnitudes = torch.exp2(((n_blocks + 2 * k_blocks) % 5).float()) + return magnitudes.repeat_interleave(128, 0).repeat_interleave(128, 1) + + +def _dequantize_1x128(x_fp8: torch.Tensor, + scales: torch.Tensor) -> torch.Tensor: + """FP64 values of rows quantized per 1x128 block, where + scales[..., j, i] scales K block j of row i.""" + m = x_fp8.shape[-2] + return x_fp8.double() * scales[..., :m].transpose( + -1, -2).double().repeat_interleave(128, -1) + + +def _dequantize_128x128(x_fp8: torch.Tensor, + scales: torch.Tensor) -> torch.Tensor: + return x_fp8.double() * scales.double().repeat_interleave( + 128, -2).repeat_interleave(128, -1) + + +def _sm100_fp8_tile_shape_tactics(runner, inputs): + """Full-sweep 1-CTA and 2-CTA tactics, with and without a cluster along + N, of every MMA tile.""" + clusters = {False: ((1, 1), (1, 2)), True: ((2, 1), (2, 2))} + tactics = [ + tactic for tactic in runner.get_valid_tactics(inputs, None) + if tactic[2] in clusters[tactic[0]] + ] + assert {tactic[1][1] for tactic in tactics} == {32, 64, 128} + return tactics + + +def _assert_matches_dequantized(output, expected, tactic): + assert torch.isfinite(output).all(), tactic + torch.testing.assert_close(output.double(), + expected, + rtol=1e-2, + atol=1e-3 * expected.abs().amax().item(), + msg=lambda msg: f"{tactic}: {msg}") + assert calc_diff(output, expected) < 1e-5, tactic + + +@pytest.mark.skipif( + getSMVersion() not in (100, 103), + reason="The test is for SM100 and SM103 only. Current SM is %d." % + getSMVersion(), +) +@pytest.mark.parametrize("m", [7, 200]) +def test_cute_dsl_fp8_block_scale_gemm_tile_shapes(m, monkeypatch): + """Every MMA tile of the SM100 blockwise GEMM computes the product, + including the N = 32 and 64 tiles inside a 128-row scale block of B.""" + k, n = 1024, 384 + torch.random.manual_seed(0) + a = torch.randn((m, k), device='cuda', dtype=torch.bfloat16) + b = (torch.randn( + (n, k), device='cuda') * _n_block_magnitudes(n, k)).to(torch.bfloat16) + a_fp8, a_sf = torch.ops.trtllm.fp8_quantize_1x128(a) + b_fp8, b_sf = per_block_cast_to_fp8(b) + expected = _dequantize_1x128(a_fp8, a_sf) @ _dequantize_128x128( + b_fp8, b_sf).t() + + monkeypatch.setattr(AutoTuner.get(), "nvmmh_config", None) + runner = cute_dsl_custom_ops.CuteDSLFp8BlackwellRunner() + inputs = [a_fp8, b_fp8, a_sf, b_sf] + for tactic in _sm100_fp8_tile_shape_tactics(runner, inputs): + output = runner(inputs, tactic=tactic) + torch.cuda.synchronize() + _assert_matches_dequantized(output, expected, tactic) + # The next tactic's output is likely allocated in this memory. + output.fill_(float("nan")) + + +@pytest.mark.skipif( + getSMVersion() not in (100, 103), + reason="The test is for SM100 and SM103 only. Current SM is %d." % + getSMVersion(), +) +@pytest.mark.parametrize("m", [7, 200]) +def test_cute_dsl_fp8_block_scale_bmm_tile_shapes(m, monkeypatch): + """Every MMA tile of the SM100 blockwise BMM computes the product, + including the N = 32 and 64 tiles inside a 128-row scale block of B.""" + k, n, num_groups = 1024, 384, 3 + torch.random.manual_seed(0) + a = torch.randn((m, num_groups, k), device='cuda', dtype=torch.bfloat16) + a_fp8, a_scales = torch.ops.trtllm.fp8_batched_quantize_1x128_permute102(a) + b = (torch.randn((num_groups, n, k), device='cuda') * + _n_block_magnitudes(n, k)).to(torch.bfloat16) + b_fp8 = torch.empty_like(b, dtype=torch.float8_e4m3fn) + b_scales = torch.empty((num_groups, n // 128, k // 128), + device='cuda', + dtype=torch.float) + for i in range(num_groups): + b_fp8[i], b_scales[i] = per_block_cast_to_fp8(b[i]) + expected = _dequantize_1x128(a_fp8, a_scales) @ _dequantize_128x128( + b_fp8, b_scales).transpose(1, 2) + output = torch.empty((num_groups, m, n), + device='cuda', + dtype=torch.bfloat16) + + monkeypatch.setattr(AutoTuner.get(), "nvmmh_config", None) + runner = cute_dsl_custom_ops.CuteDSLFp8BlackwellBmmRunner() + inputs = [a_fp8, b_fp8, a_scales, b_scales, output] + for tactic in _sm100_fp8_tile_shape_tactics(runner, inputs): + output.fill_(float("nan")) + runner(inputs, tactic=tactic) + torch.cuda.synchronize() + _assert_matches_dequantized(output, expected, tactic) + + @pytest.mark.skipif( getSMVersion() != 107 or not IS_CUTLASS_DSL_RUBIN_AVAILABLE, reason="The test requires SM107 and SM107 CuTe DSL support.", @@ -1310,6 +1424,78 @@ def test_cute_dsl_fp8_bmm_quantize_rubin_matches_separate(m): atol=1.0 / 16.0) +@pytest.mark.skipif( + getSMVersion() != 107 or not IS_CUTLASS_DSL_RUBIN_AVAILABLE, + reason="The test requires SM107 and SM107 CuTe DSL support.", +) +@pytest.mark.parametrize("m", [16, 256]) +def test_cute_dsl_fp8_bmm_quantize_rubin_nvmmh_tactics(m, monkeypatch): + """Every nvMMH-guided tactic of the fused FP8 BMM and 1x128 quantization + matches the BF16 BMM of the same tactic quantized separately.""" + from tensorrt_llm._torch.custom_ops import \ + cutedsl_matmul_heuristics as nvmmh + from tensorrt_llm._torch.custom_ops.cute_dsl_custom_ops import ( + CuteDSLFp8RubinBmmQuantRunner, CuteDSLFp8RubinBmmRunner) + + if not nvmmh.IS_NVMMH_AVAILABLE: + pytest.skip("nvMatmulHeuristics is not installed") + k, n, num_groups = 4096, 1024, 8 + flat_n = num_groups * n + torch.manual_seed(314) + a = torch.randn((m, num_groups, k), device="cuda", dtype=torch.bfloat16) / k + a_fp8, a_scales = torch.ops.trtllm.fp8_batched_quantize_1x128_permute102(a) + b = torch.randn((num_groups, n, k), device="cuda", dtype=torch.bfloat16) / k + b_fp8 = torch.empty_like(b, dtype=torch.float8_e4m3fn) + b_scales = torch.empty((num_groups, n // 128, k // 128), + device="cuda", + dtype=torch.float32) + for group_idx in range(num_groups): + b_fp8[group_idx], b_scales[group_idx] = per_block_cast_to_fp8( + b[group_idx]) + bf16_output = torch.empty((num_groups, m, n), + device="cuda", + dtype=torch.bfloat16) + fused_fp8, fused_batched, fused_scale = _allocate_fp8_bmm_quant_outputs( + a_fp8, b_fp8) + + selected = [] + select_tactics = nvmmh.select_tactics + + def recording_select_tactics(ranked, max_tactics): + selected.append(select_tactics(ranked, max_tactics)) + return selected[-1] + + monkeypatch.setattr(nvmmh, "select_tactics", recording_select_tactics) + monkeypatch.setattr(AutoTuner.get(), "nvmmh_config", + NvMMHConfig(max_tactics=3)) + runner = CuteDSLFp8RubinBmmQuantRunner() + inputs = [a_fp8, b_fp8, a_scales, b_scales, fused_batched, fused_scale] + tactics = runner.get_valid_tactics(inputs, None) + # The tactics come from nvMMH, not from the full-sweep fallback. + assert tactics and tactics == selected[-1] + + bf16_runner = CuteDSLFp8RubinBmmRunner() + for tactic in tactics: + bf16_output.fill_(float("nan")) + bf16_runner([a_fp8, b_fp8, a_scales, b_scales, bf16_output], + tactic=tactic) + ref_fp8, ref_scale = torch.ops.trtllm.fp8_quantize_1x128_packed_ue8m0( + bf16_output.permute(1, 0, 2).reshape(m, flat_n)) + # 0xFF is NaN in both E4M3 and UE8M0. + fused_fp8.view(torch.uint8).fill_(0xFF) + fused_scale.fill_(0xFF) + runner(inputs, tactic=tactic) + torch.cuda.synchronize() + + _assert_valid_fp8_bmm_quant_scales_equal(fused_scale, ref_scale, m, + flat_n) + torch.testing.assert_close(fused_fp8.float(), + ref_fp8.float(), + rtol=0.0, + atol=1.0 / 16.0, + msg=lambda msg: f"{tactic}: {msg}") + + @pytest.mark.skipif( getSMVersion() != 107 or not IS_CUTLASS_DSL_RUBIN_AVAILABLE, reason="The test requires SM107 and SM107 CuTe DSL support.", diff --git a/tests/unittest/_torch/misc/test_autotuner.py b/tests/unittest/_torch/misc/test_autotuner.py index efbba2697f5d..867aa1091835 100644 --- a/tests/unittest/_torch/misc/test_autotuner.py +++ b/tests/unittest/_torch/misc/test_autotuner.py @@ -22,6 +22,7 @@ import statistics import sys import tempfile +from types import SimpleNamespace from typing import Any, List import cloudpickle @@ -1829,6 +1830,174 @@ def recording_select_tactics(ranked, max_tactics): assert calc_diff(output, reference) < max_diff, tactic +class _TacticSpaceRunner(TunableRunner): + """Runner whose tactic space key is set by the test.""" + + def __init__(self, space_key=None): + self.space_key = space_key + + def unique_id(self): + return () + + def tactic_space_key(self): + return self.space_key + + def get_valid_tactics(self, inputs: List[FakeTensor], + profile: OptimizationProfile, **kwargs) -> List[int]: + return [0] + + def forward(self, + /, + inputs: List[torch.Tensor], + *, + tactic: int = -1, + **kwargs) -> torch.Tensor: + return inputs[0] + + +_TACTIC_SPACE_SHAPES = (torch.Size([16, 64]), ) + + +def _tactic_space_cache_key(cache, runner): + return cache.get_cache_key("test_tactic_space", runner, + _TACTIC_SPACE_SHAPES, TuningConfig()) + + +def test_tactic_space_key_extends_the_cache_key(): + cache = autotuner.AutoTuner().profiling_cache + default = _tactic_space_cache_key(cache, _TacticSpaceRunner()) + assert default == ("test_tactic_space", "_TacticSpaceRunner", "()", + ((16, 64), )) + guided = _tactic_space_cache_key(cache, _TacticSpaceRunner( + ("nvmmh", "x", 3))) + assert guided == default + (("nvmmh", "x", 3), ) + assert _tactic_space_cache_key(cache, _TacticSpaceRunner( + ("nvmmh", "x", 5))) != guided + + +def test_search_cache_misses_across_tactic_spaces(): + cache = autotuner.AutoTuner().profiling_cache + guided = _TacticSpaceRunner(("nvmmh", "x", 3)) + cache.cache[_tactic_space_cache_key(cache, guided)] = (0, 7, 1.0) + + def search(runner): + return cache.search_cache("test_tactic_space", [runner], + _TACTIC_SPACE_SHAPES, TuningConfig()) + + assert search(guided) == (True, 0, 7, 1.0) + assert not search(_TacticSpaceRunner())[0] + assert not search(_TacticSpaceRunner(("nvmmh", "x", 5)))[0] + + +def test_tactic_space_key_survives_cache_save_and_load(tmp_path): + src = autotuner.AutoTuner().profiling_cache + key = _tactic_space_cache_key(src, + _TacticSpaceRunner(("nvmmh", "0.1.0.27", 3))) + src.cache[key] = (0, 7, 1.0) + path = str(tmp_path / "cache.json") + src.save_cache(path, rank=0) + + dst = autotuner.AutoTuner().profiling_cache + dst.load_cache(path, rank=0) + assert dst.cache[key][1] == 7 + + +def test_nvmmh_tactic_space_key_follows_the_policy(monkeypatch): + from tensorrt_llm._torch.custom_ops import cute_dsl_custom_ops as ops + + tuner = AutoTuner.get() + monkeypatch.setattr(ops.cutedsl_nvmmh, "IS_NVMMH_AVAILABLE", True) + monkeypatch.setattr(ops.cutedsl_nvmmh, "library_version", lambda: "x") + monkeypatch.setattr(tuner, "nvmmh_config", None) + assert ops._nvmmh_tactic_space_key() is None + monkeypatch.setattr(tuner, "nvmmh_config", NvMMHConfig(max_tactics=3)) + assert ops._nvmmh_tactic_space_key() == ("nvmmh", "x", 3) + monkeypatch.setattr(ops.cutedsl_nvmmh, "IS_NVMMH_AVAILABLE", False) + assert ops._nvmmh_tactic_space_key() is None + + +def test_nvfp4_unified_runner_keys_only_the_cutedsl_tactic_space(monkeypatch): + from tensorrt_llm._torch.custom_ops import torch_custom_ops + + monkeypatch.setattr(torch_custom_ops, "IS_CUTLASS_DSL_AVAILABLE", True) + monkeypatch.setattr(torch_custom_ops, + "_nvmmh_tactic_space_key", + lambda: ("nvmmh", "x", 3), + raising=False) + runner = torch_custom_ops.NVFP4GemmUnifiedRunner + assert runner(0, torch.bfloat16, + ["cutlass", "cutedsl"]).tactic_space_key() == ("nvmmh", "x", + 3) + assert runner(0, torch.bfloat16, + ["cutlass", "cublaslt"]).tactic_space_key() is None + + +def test_locality_domain_runner_forwards_the_tactic_space_key(): + from tensorrt_llm._torch.locality_domain.autotune import \ + LocalityDomainConcurrentTunableRunner + + inner = _TacticSpaceRunner() + wrapper = LocalityDomainConcurrentTunableRunner( + inner, SimpleNamespace(num_partitions=2), 2, lambda *args: None) + assert wrapper.tactic_space_key() is None + inner.space_key = ("nvmmh", "x", 3) + assert wrapper.tactic_space_key() == ("nvmmh", "x", 3) + + +def test_configure_autotuner_nvmmh_installs_the_engine_policy(monkeypatch): + from tensorrt_llm._torch.pyexecutor.model_engine import \ + _configure_autotuner_nvmmh + from tensorrt_llm.llmapi.llm_args import AutoTunerNvMMHConfig + + tuner = AutoTuner.get() + monkeypatch.setattr(tuner, "nvmmh_config", None) + _configure_autotuner_nvmmh( + SimpleNamespace(autotuner_nvmmh_config=AutoTunerNvMMHConfig( + max_tactics=7))) + assert tuner.nvmmh_config == NvMMHConfig(max_tactics=7) + _configure_autotuner_nvmmh(SimpleNamespace(autotuner_nvmmh_config=None)) + assert tuner.nvmmh_config is None + + +def test_merge_nvmmh_rankings_picks_the_lowest_scoring_head(): + from tensorrt_llm._torch.custom_ops.cute_dsl_custom_ops import \ + _merge_nvmmh_rankings + + def ranking(name, *scores): + return [(SimpleNamespace(score=score), f"{name}{i}") + for i, score in enumerate(scores)] + + def merge(*rankings): + return [name for _, name in _merge_nvmmh_rankings(list(rankings))] + + # Each ranking keeps its own order, even where it is not sorted by score. + assert merge(ranking("a", 1, 4, 2), + ranking("b", 3, 5)) == ["a0", "b0", "a1", "a2", "b1"] + # Ties go to the earlier ranking. + assert merge(ranking("a", 1, 1), ranking("b", 1)) == ["a0", "a1", "b0"] + assert merge(ranking("a"), ranking("b", 2)) == ["b0"] + assert merge() == [] + + +def test_select_tactics_counts_configs_with_new_tactics(): + from tensorrt_llm._torch.custom_ops.cutedsl_matmul_heuristics import \ + select_tactics + + ranked = [(None, ["a", "b"]), (None, ["b"]), (None, ["b", "c"]), + (None, ["d"])] + # The second config only repeats "b", so it is not counted. + assert select_tactics(ranked, 1) == ["a", "b"] + assert select_tactics(ranked, 2) == ["a", "b", "c"] + assert select_tactics(ranked, 10) == ["a", "b", "c", "d"] + assert select_tactics([], 3) == [] + + +@pytest.mark.parametrize("max_tactics", [True, 0, -1, "5", 2.0]) +def test_nvmmh_config_rejects_non_positive_integers(max_tactics): + with pytest.raises(ValueError, match="positive integer"): + NvMMHConfig(max_tactics=max_tactics) + + @pytest.mark.parametrize("distribution", ["random", "balanced"]) def test_trtllm_gen_moe_dummy_topk_local_experts_less_than_topk( distribution, monkeypatch):