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()