From 0e6fb4eb03a9050f96e9ac3a2a510d8719c2611b Mon Sep 17 00:00:00 2001 From: peihengh <259410613+peihu-nv@users.noreply.github.com> Date: Tue, 28 Jul 2026 08:25:52 -0700 Subject: [PATCH] [None][perf] Add CUTEDSL FC2 N-tile tuning override Signed-off-by: peihengh <259410613+peihu-nv@users.noreply.github.com> --- .../_torch/custom_ops/cute_dsl_custom_ops.py | 29 ++++++- .../custom_ops/test_cutedsl_fc2_tuning.py | 78 +++++++++++++++++++ 2 files changed, 106 insertions(+), 1 deletion(-) create mode 100644 tests/unittest/_torch/custom_ops/test_cutedsl_fc2_tuning.py 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 e53d58693a4b..de92e5950d5c 100644 --- a/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py +++ b/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py @@ -4,6 +4,7 @@ import functools import itertools import math +import os from typing import List, Optional, Tuple import torch @@ -34,12 +35,27 @@ # Torch schema parsing rejects ``inf`` as a default value. SWIGLU_LIMIT_SCALAR_DISABLED = -1.0 +_CUTEDSL_FC2_N_TILE_SIZE_ENV = "TRTLLM_CUTEDSL_FC2_N_TILE_SIZE" +_CUTEDSL_FC2_N_TILE_SIZES = (128, 256) def _canonicalize_swiglu_limit_scalar(swiglu_limit_scalar: float) -> float: return float("inf") if swiglu_limit_scalar < 0 else swiglu_limit_scalar +def _get_cutedsl_fc2_n_tile_size_override() -> Optional[int]: + value = os.environ.get(_CUTEDSL_FC2_N_TILE_SIZE_ENV) + if value is None: + return None + valid_values = tuple( + str(tile_size) for tile_size in _CUTEDSL_FC2_N_TILE_SIZES) + if value not in valid_values: + raise ValueError( + f"{_CUTEDSL_FC2_N_TILE_SIZE_ENV} must be unset or one of " + f"{', '.join(valid_values)}; got {value!r}") + return int(value) + + class GroupedGemmInputsHelper: """Base helper class for grouped GEMM input preparation and tuning. @@ -2508,6 +2524,8 @@ def __init__(self, assert output_dtype == torch.bfloat16 self.output_dtype = output_dtype self.scaling_vector_size = scaling_vector_size + self.fc2_n_tile_size_override = ( + _get_cutedsl_fc2_n_tile_size_override()) if (sm_version := get_sm_version()) not in (100, 103): raise ValueError( @@ -2520,7 +2538,7 @@ def __init__(self, ) def unique_id(self): - return ( + runner_id = ( self.num_experts, self.top_k, self.num_local_experts, @@ -2529,6 +2547,12 @@ def unique_id(self): self.output_dtype, self.scaling_vector_size, ) + if self.fc2_n_tile_size_override is not None: + runner_id += ( + "fc2_n_tile_size_override", + self.fc2_n_tile_size_override, + ) + return runner_id def get_valid_tactics( self, @@ -2542,6 +2566,9 @@ def get_valid_tactics( mma_tiler_mn_candidates = [(self.tile_size, 128), (self.tile_size, 256)] + if self.fc2_n_tile_size_override is not None: + mma_tiler_mn_candidates = [(self.tile_size, + self.fc2_n_tile_size_override)] cluster_shape_mn_candidates = [(self.tile_size // 128, 1), (self.tile_size // 128, 2)] # raster_along_m=False should be theoretically more performant than raster_along_m=True. diff --git a/tests/unittest/_torch/custom_ops/test_cutedsl_fc2_tuning.py b/tests/unittest/_torch/custom_ops/test_cutedsl_fc2_tuning.py new file mode 100644 index 000000000000..7c10cdb11f85 --- /dev/null +++ b/tests/unittest/_torch/custom_ops/test_cutedsl_fc2_tuning.py @@ -0,0 +1,78 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from types import SimpleNamespace + +import pytest +import torch + +import tensorrt_llm._torch.custom_ops.cute_dsl_custom_ops as cute_dsl_custom_ops +from tensorrt_llm._torch.autotuner import OptimizationProfile +from tensorrt_llm._torch.cute_dsl_utils import IS_CUTLASS_DSL_AVAILABLE + +if IS_CUTLASS_DSL_AVAILABLE: + from tensorrt_llm._torch.custom_ops.cute_dsl_custom_ops import ( + Sm100BlockScaledContiguousGroupedGemmFinalizeFusionRunner, + ) + + +_CUTEDSL_FC2_N_TILE_SIZE_ENV = "TRTLLM_CUTEDSL_FC2_N_TILE_SIZE" + +pytestmark = pytest.mark.skipif( + not IS_CUTLASS_DSL_AVAILABLE, + reason="Requires CUTLASS DSL", +) + + +def _make_runner(monkeypatch) -> "Sm100BlockScaledContiguousGroupedGemmFinalizeFusionRunner": + monkeypatch.setattr(cute_dsl_custom_ops, "get_sm_version", lambda: 100) + monkeypatch.setattr( + Sm100BlockScaledContiguousGroupedGemmFinalizeFusionRunner, + "kernel_class", + SimpleNamespace(can_implement=lambda **kwargs: True), + ) + return Sm100BlockScaledContiguousGroupedGemmFinalizeFusionRunner( + num_experts=256, + top_k=8, + num_local_experts=8, + local_expert_offset=0, + tile_size=128, + output_dtype=torch.bfloat16, + ) + + +@pytest.mark.parametrize( + ("env_value", "expected_n_tiles"), + [ + (None, {128, 256}), + ("128", {128}), + ("256", {256}), + ], +) +def test_fc2_n_tile_size_override(monkeypatch, env_value, expected_n_tiles) -> None: + if env_value is None: + monkeypatch.delenv(_CUTEDSL_FC2_N_TILE_SIZE_ENV, raising=False) + else: + monkeypatch.setenv(_CUTEDSL_FC2_N_TILE_SIZE_ENV, env_value) + + runner = _make_runner(monkeypatch) + inputs = [torch.empty(256, 64), torch.empty(8, 256)] + + tactics = runner.get_valid_tactics(inputs, OptimizationProfile()) + + assert {tactic[0][1] for tactic in tactics} == expected_n_tiles + if env_value is None: + assert "fc2_n_tile_size_override" not in runner.unique_id() + else: + assert runner.unique_id()[-2:] == ( + "fc2_n_tile_size_override", + int(env_value), + ) + + +@pytest.mark.parametrize("env_value", ["", "64", "128,256", " 128", "invalid"]) +def test_fc2_n_tile_size_override_rejects_invalid_values(monkeypatch, env_value) -> None: + monkeypatch.setenv(_CUTEDSL_FC2_N_TILE_SIZE_ENV, env_value) + + with pytest.raises(ValueError, match=f"{_CUTEDSL_FC2_N_TILE_SIZE_ENV} must be unset"): + _make_runner(monkeypatch)