From 96f75b508139d787087939482ed8fab58d6af421 Mon Sep 17 00:00:00 2001 From: Barry Kang <43644113+Barry-Delaney@users.noreply.github.com> Date: Wed, 28 Oct 2026 04:42:24 -0700 Subject: [PATCH] [None][fix] CuteDSL MoE: keep inplace FC1 outputs out of the Rubin autotune key Sm107BlockScaledContiguousGatherGroupedGemmActFusionRunner takes twelve inputs; 10 and 11 are the inplace output and output scale-factor buffers the locality-domain half-GEMM writes into. Its TuningConfig buckets input 7 (permuted_idx_to_expanded_idx) and constrains inputs 0/2/5/6 to it, but leaves 10/11 unconstrained, so the autotuner copies their raw row counts into the cache key. Those buffers are sized from the actual permuted-token count, which defeats the bucketing on input 7: entries are stored per distinct buffer size, and later lookups at another token count in the same bucket miss and fall back to the default tactic. Constrain both buffers to input 7's bucket. The scale-factor buffer size matches what the locality-domain caller allocates and what the kernel asserts (m * n // scaling_vector_size with n the per-partition weight width), which keeps the profiling tensors consistent with _resize_locality_domain_outputs. When the outputs are None (the single-tensor path) the autotuner skips these specs, so those cache keys are unchanged. Signed-off-by: Barry Kang <43644113+Barry-Delaney@users.noreply.github.com> --- .../_torch/custom_ops/cute_dsl_custom_ops.py | 24 +++++++++ tensorrt_llm/_torch/cute_dsl_utils.py | 6 +-- .../moe/kernels/parallel/test_cute_dsl_moe.py | 52 +++++++++++++++++++ 3 files changed, 78 insertions(+), 4 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 54792dab7866..02b4f13b208d 100644 --- a/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py +++ b/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py @@ -560,6 +560,20 @@ class GatherGroupedGemmInputsHelper(GroupedGemmInputsHelper): IDX_PERMUTED_IDX_TO_EXPANDED_IDX = 7 IDX_SHAPE_INFER = IDX_PERMUTED_IDX_TO_EXPANDED_IDX + def infer_shape_locality_domain_output_sf_numel( + self, + input_shapes: List[torch.Size], + scaling_vector_size: int = 16) -> int: + """Element count of the shared FC1 output scale-factor buffer. + + Both locality-domain partitions write ``m * full_interm // + scaling_vector_size`` scales; each partition's weight width ``n`` + (input 1) equals ``full_interm``. + """ + n = input_shapes[1][1] + return (self.infer_shape_max_num_permuted_tokens(input_shapes) * n // + scaling_vector_size) + @staticmethod def _resize_locality_domain_outputs( m: int, @@ -15947,6 +15961,9 @@ def get_tuning_config(self) -> TuningConfig: helper = GatherGroupedGemmInputsHelper( self.num_experts, self.top_k, self.num_local_experts, self.local_expert_offset, self.tile_size) + infer_shape_output_sf_numel = functools.partial( + helper.infer_shape_locality_domain_output_sf_numel, + scaling_vector_size=self.scaling_vector_size) self.__class__.tuning_config_cache[key] = TuningConfig( # Use permuted_idx_to_expanded_idx (IDX_SHAPE_INFER) for tuning dynamic_tensor_specs=(DynamicTensorSpec( @@ -15960,6 +15977,13 @@ def get_tuning_config(self) -> TuningConfig: helper.infer_shape_max_num_tiles), ConstraintSpec(6, 0, helper.infer_shape_max_num_tiles), + # Inplace outputs of the locality-domain half-GEMM + # (10: output, 11: output_sf) follow input 7's bucket; + # the autotuner skips these specs when they are None. + ConstraintSpec( + 10, 0, + helper.infer_shape_max_num_permuted_tokens), + ConstraintSpec(11, 0, infer_shape_output_sf_numel), ), inputs_pre_hook=helper.inputs_pre_hook, ) diff --git a/tensorrt_llm/_torch/cute_dsl_utils.py b/tensorrt_llm/_torch/cute_dsl_utils.py index b99362b2cffb..4af6f1429a53 100644 --- a/tensorrt_llm/_torch/cute_dsl_utils.py +++ b/tensorrt_llm/_torch/cute_dsl_utils.py @@ -38,10 +38,8 @@ def _skip_legacy_cutlass_mlir_helpers() -> None: IS_CUTLASS_DSL_AVAILABLE = False -# Whether the public CuTe DSL package provides the SM107/Rubin helper module. -# Rubin kernels stay disabled when it is absent and callers retain their -# existing fallback paths. -# TODO: flips to True once a Rubin-capable CuTe DSL package ships in the image. +# Whether the installed CuTe DSL provides ``cutlass.utils.rubin_helpers``; set +# by the import probe below. Rubin kernels stay disabled when it is absent. IS_CUTLASS_DSL_RUBIN_AVAILABLE = False if platform.system() != "Windows": diff --git a/tests/unittest/_torch/moe/kernels/parallel/test_cute_dsl_moe.py b/tests/unittest/_torch/moe/kernels/parallel/test_cute_dsl_moe.py index b13d3b6747aa..10d7a6832631 100644 --- a/tests/unittest/_torch/moe/kernels/parallel/test_cute_dsl_moe.py +++ b/tests/unittest/_torch/moe/kernels/parallel/test_cute_dsl_moe.py @@ -3250,6 +3250,58 @@ def test_sm107_nvfp4_finalize_accepts_two_cta_n_tiling(): assert runner_cls._is_n_tiling_compatible(n=256, mma_n=128, cluster_n=2) +@pytest.mark.cpu_only +@pytest.mark.skipif( + not IS_CUTLASS_DSL_RUBIN_AVAILABLE, + reason="This test requires the public Rubin CuTe DSL package", +) +def test_sm107_nvfp4_gather_act_fusion_cache_key_ignores_inplace_output_rows(monkeypatch): + monkeypatch.setattr(cute_dsl_custom_ops, "get_sm_version", lambda: 107) + runner = cute_dsl_custom_ops.Sm107BlockScaledContiguousGatherGroupedGemmActFusionRunner( + num_experts=1, top_k=1, num_local_experts=1, local_expert_offset=0, tile_size=128 + ) + tuning_config = runner.get_tuning_config() + shard_n, packed_k, tile_size = 512, 128, 128 + + def input_shapes(m: int) -> tuple[torch.Size, ...]: + num_tiles, num_tokens, scale_k = m // tile_size, m // 2, packed_k * 2 // 16 + return ( + torch.Size((num_tokens, packed_k)), + torch.Size((1, shard_n, packed_k)), + torch.Size((num_tokens, scale_k)), + torch.Size((1, shard_n, scale_k)), + torch.Size((1,)), + torch.Size((num_tiles,)), + torch.Size((num_tiles,)), + torch.Size((m,)), + torch.Size((1,)), + torch.Size((1,)), + # Inplace output and output scale factors, sized from the actual m. + torch.Size((m, shard_n // 2)), + torch.Size((m * shard_n // 16,)), + ) + + def cache_key(m: int) -> tuple: + return AutoTuner._find_nearest_profile( + input_shapes(m), + tuning_config.dynamic_tensor_specs, + tuning_config.constraint_specs, + tuning_config.tune_max_num_tokens, + True, + ) + + # Two permuted-token counts in one tuning bucket must share a cache entry. + bucket_map = tuning_config.dynamic_tensor_specs[0].map_to_tuning_buckets + assert bucket_map(1024) == bucket_map(1152) + assert cache_key(1024) == cache_key(1152) + + # Profiling sizes the scale-factor buffer as the kernel asserts it. + inputs = [torch.empty(shape, dtype=torch.uint8, device="meta") for shape in input_shapes(1152)] + for profile in AutoTuner.get()._optimization_profiles(tuning_config, inputs): + shapes = profile.get_opt_shapes() + assert shapes[11] == (shapes[7][0] * shard_n // 16,) + + def _get_registered_rubin_moe_op(name: str): try: return getattr(torch.ops.trtllm, name)