Skip to content

[CuTeDSL] Make autotuning more user friendly & new docs - #3408

Open
kainzhong wants to merge 7 commits into
NVIDIA:mainfrom
kainzhong:improve_autotune
Open

[CuTeDSL] Make autotuning more user friendly & new docs#3408
kainzhong wants to merge 7 commits into
NVIDIA:mainfrom
kainzhong:improve_autotune

Conversation

@kainzhong

@kainzhong kainzhong commented Jul 24, 2026

Copy link
Copy Markdown
Contributor

This PR adds a lot of things to the overall CuTeDSL autotuning process to make it easier to use and more convenient:

  • Add derived_params to autotuning API so it computes derived arguments that must be decided during compile time but rely on autotuned parameters (whose values are different per configuration tuned)
  • Make testing.autotune_jit able to decorate a class. This is for people who prefer to implement kernels as classes that fix the compile time params in __init__ method and pass runtime params to __call__, so now they can do this
@testing.autotune_jit(
    params_dict={"COPY_BITS": [32, 64, 128], "THREADS_M": [4, 8]},
    update_on_change=["M", "N"],
    warmup_iterations=10,
    iterations=50,
)
class ElementwiseAddKernel:
    THREADS_M = 4

    def __init__(self, COPY_BITS=128):
        self.COPY_BITS = COPY_BITS

    @cute.jit
    def __call__(self, mA, mB, mC, M, N):
        elementwise_add(
            mA, mB, mC, M, N,
            copy_bits=self.COPY_BITS,
            threads_m=self.THREADS_M,
        )
  • Support pruning undesired configurations (either invalid or known to be suboptimal) -- this aligns with triton's prune_configs_by in their autotuning API
  • Add autotune_suite API so now user can provide a list of representative shapes and find the best overall kernel candidates with the ranking rule: sort according to how many times a configuration is accepted across all inputs, where "accepted" mean for a input case the configuration is in the top accept_percentile fastest ones
  • Better autotuning docs. Previously the doc only explains how to autotune manually which is quite painful at least for me

Signed-off-by: Kaining Zhong <kainingz@nvidia.com>
Signed-off-by: Kaining Zhong <kainingz@nvidia.com>
Signed-off-by: Kaining Zhong <kainingz@nvidia.com>
Signed-off-by: Kaining Zhong <kainingz@nvidia.com>
Signed-off-by: Kaining Zhong <kainingz@nvidia.com>
Signed-off-by: Kaining Zhong <kainingz@nvidia.com>
Signed-off-by: Kaining Zhong <kainingz@nvidia.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant