[None][perf] Add CUTEDSL FC2 N-tile tuning override

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.

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)