From 6d101d9b8e1bb66d8c41be5342c4fa33b9a3cbbe Mon Sep 17 00:00:00 2001 From: chenfeiz0326 <203214996+chenfeiz0326@users.noreply.github.com> Date: Wed, 14 Oct 2026 09:06:44 -0700 Subject: [PATCH] [https://nvbugs/6919819][perf] Autotune the CuTeDSL FC2 N tile again when the override is unset The FC2 finalize-fusion runner's unset default pinned the N tile to 128, which removed N256 from the autotuner's candidate set for every shape. Prefill-sized FC2 calls and attention-TP layouts profile faster at N256 and lost it. Leave TRTLLM_CUTEDSL_FC2_N_TILE_SIZE unset to autotune over both tiles, as before the override existed. Setting it to 128 or 256 still pins that tile, and 0 still autotunes. Signed-off-by: chenfeiz0326 <203214996+chenfeiz0326@users.noreply.github.com> --- tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py | 11 ++++++----- .../_torch/moe/fused_moe/test_cutedsl_fc2_tuning.py | 3 ++- 2 files changed, 8 insertions(+), 6 deletions(-) diff --git a/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py b/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py index f01da1dc55ed..da8b2175bdcf 100644 --- a/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py +++ b/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py @@ -44,7 +44,6 @@ 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) -_CUTEDSL_FC2_DEFAULT_N_TILE_SIZE = 128 # Legacy Rubin ops use a float sentinel for absent SiTU soft-caps. # The Blackwell act-fusion op accepts Optional[float] directly. @@ -192,10 +191,12 @@ def _get_sm107_nvfp4_default_mma_config( @functools.lru_cache(maxsize=1) def _get_cutedsl_fc2_n_tile_size_override() -> Optional[int]: - value = os.environ.get( - _CUTEDSL_FC2_N_TILE_SIZE_ENV, - str(_CUTEDSL_FC2_DEFAULT_N_TILE_SIZE), - ) + """Return the FC2 N tile pinned by the user, or None to autotune over all of them. + + Unset and ``0`` both autotune. ``128`` or ``256`` pins that tile for shapes + where it is known to win end to end. + """ + value = os.environ.get(_CUTEDSL_FC2_N_TILE_SIZE_ENV, "0") if value == "0": return None valid_values = tuple( diff --git a/tests/unittest/_torch/moe/fused_moe/test_cutedsl_fc2_tuning.py b/tests/unittest/_torch/moe/fused_moe/test_cutedsl_fc2_tuning.py index 0c7e4cf01a54..bddf78405490 100644 --- a/tests/unittest/_torch/moe/fused_moe/test_cutedsl_fc2_tuning.py +++ b/tests/unittest/_torch/moe/fused_moe/test_cutedsl_fc2_tuning.py @@ -61,6 +61,7 @@ def _make_runner( [ ((128, 128), (1, 1), False), ((128, 128), (1, 2), False), + ((128, 256), (1, 1), False), ], ), ( @@ -97,7 +98,7 @@ def test_fc2_n_tile_size_override( tactics = runner.get_valid_tactics(inputs, OptimizationProfile()) assert tactics == expected_tactics - expected_override = 128 if env_value is None else (None if env_value == "0" else int(env_value)) + expected_override = None if env_value in (None, "0") else int(env_value) assert runner._get_default_tactic()[0][1] == (expected_override or 128) if expected_override is None: assert "fc2_n_tile_size_override" not in runner.unique_id()