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 d5a55e7e132c..83509da8f3b9 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 37fc0269eca7..873fbafcdedc 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 \ @@ -1308,6 +1308,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 07fb708341f5..e95ef7708d77 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_engine.py +++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py @@ -32,6 +32,7 @@ from ..attention.backends.interface import (AttentionMetadata, AttentionRuntimeFeatures) from ..attention.backends.utils import get_attention_backend +from ..autotuner import AutoTuner, NvMMHConfig from ..compilation.backend import Backend from ..distributed import Distributed from ..distributed.communicator import init_pp_comm @@ -111,6 +112,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.""" @@ -167,6 +181,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( diff --git a/tensorrt_llm/llmapi/__init__.py b/tensorrt_llm/llmapi/__init__.py index 8d1dfd245852..2000ab7de133 100644 --- a/tensorrt_llm/llmapi/__init__.py +++ b/tensorrt_llm/llmapi/__init__.py @@ -9,8 +9,9 @@ from ..scheduling_params import SchedulingParams from .llm import LLM, RequestOutput # yapf: disable -from .llm_args import (AttentionDpConfig, AutoDecodingConfig, BatchingType, - BlockReuseConfig, CacheTransceiverConfig, CalibConfig, +from .llm_args import (AttentionDpConfig, AutoDecodingConfig, + AutoTunerNvMMHConfig, BatchingType, BlockReuseConfig, + CacheTransceiverConfig, CalibConfig, CapacitySchedulerPolicy, ColdPageQuantizationCompressionConfig, ContextChunkingPolicy, CudaGraphConfig, @@ -90,6 +91,7 @@ 'LlmArgs', 'TorchLlmArgs', 'AutoDecodingConfig', + 'AutoTunerNvMMHConfig', 'AttentionDpConfig', 'LoRARequest', 'SaveHiddenStatesDecodingConfig', diff --git a/tensorrt_llm/llmapi/llm_args.py b/tensorrt_llm/llmapi/llm_args.py index e02c42ae24f5..10965f3d3429 100644 --- a/tensorrt_llm/llmapi/llm_args.py +++ b/tensorrt_llm/llmapi/llm_args.py @@ -2006,6 +2006,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, @@ -6854,6 +6869,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 3db659f73b23..0d0ac5abef87 100644 --- a/tensorrt_llm/schemas/trtllm-serve-config.schema.json +++ b/tensorrt_llm/schemas/trtllm-serve-config.schema.json @@ -340,6 +340,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.", @@ -6427,6 +6442,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 371b27ff1e08..43716585ca94 100644 --- a/tensorrt_llm/schemas/trtllm-serve-disagg-config.schema.json +++ b/tensorrt_llm/schemas/trtllm-serve-disagg-config.schema.json @@ -340,6 +340,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.", @@ -1691,6 +1706,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", @@ -8003,6 +8031,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 f940f2a80937..8d95a449d17b 100644 --- a/tensorrt_llm/usage/llm_args_golden_manifest.json +++ b/tensorrt_llm/usage/llm_args_golden_manifest.json @@ -91,6 +91,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 83fd7d54f1ed..e03198aecb36 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