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)