How to use the new autotune intoruced in https://github.com/pytorch/torchdynamo/pull/1338

ayoub_louati: <p>Hello,<br>
Please how can we use the new decorator related to the caching_autotune introduced in:</p><aside class="onebox githubpullrequest" data-onebox-src="https://github.com/pytorch/torchdynamo/pull/1338">
  <header class="source">

      <a href="https://github.com/pytorch/torchdynamo/pull/1338" target="_blank" rel="noopener nofollow ugc">github.com/pytorch/torchdynamo</a>
  </header>

  <article class="onebox-body">
    <div class="github-row">



    <div class="github-icon-container" title="Pull Request">
      <svg width="60" height="60" class="github-icon" viewbox="0 0 12 16" aria-hidden="true"><path fill-rule="evenodd" d="M11 11.28V5c-.03-.78-.34-1.47-.94-2.06C9.46 2.35 8.78 2.03 8 2H7V0L4 3l3 3V4h1c.27.02.48.11.69.31.21.2.3.42.31.69v6.28A1.993 1.993 0 0 0 10 15a1.993 1.993 0 0 0 1-3.72zm-1 2.92c-.66 0-1.2-.55-1.2-1.2 0-.65.55-1.2 1.2-1.2.65 0 1.2.55 1.2 1.2 0 .65-.55 1.2-1.2 1.2zM4 3c0-1.11-.89-2-2-2a1.993 1.993 0 0 0-1 3.72v6.56A1.993 1.993 0 0 0 2 15a1.993 1.993 0 0 0 1-3.72V4.72c.59-.34 1-.98 1-1.72zm-.8 10c0 .66-.55 1.2-1.2 1.2-.65 0-1.2-.55-1.2-1.2 0-.65.55-1.2 1.2-1.2.65 0 1.2.55 1.2 1.2zM2 4.2C1.34 4.2.8 3.65.8 3c0-.65.55-1.2 1.2-1.2.65 0 1.2.55 1.2 1.2 0 .65-.55 1.2-1.2 1.2z"></path></svg>
    </div>

  <div class="github-info-container">



      <h4>
        <a href="https://github.com/pytorch/torchdynamo/pull/1338" target="_blank" rel="noopener nofollow ugc">Use new Triton runtime</a>
      </h4>

    <div class="branches">
      <code>pytorch:main</code> ← <code>jansel:newruntime202209</code>
    </div>

      <div class="github-info">
        <div class="date">
          opened <span class="discourse-local-date" data-format="ll" data-date="2022-09-24" data-time="23:59:54" data-timezone="UTC">11:59PM - 24 Sep 22 UTC</span>
        </div>

        <div class="user">
          <a href="https://github.com/jansel" target="_blank" rel="noopener nofollow ugc">
            <img alt="jansel" src="https://avatars.githubusercontent.com/u/533820?v=4" class="onebox-avatar-inline" width="20" height="20">
            jansel
          </a>
        </div>

        <div class="lines" title="26 commits changed 18 files with 453 additions and 176 deletions">
          <a href="https://github.com/pytorch/torchdynamo/pull/1338/files" target="_blank" rel="noopener nofollow ugc">
            <span class="added">+453</span>
            <span class="removed">-176</span>
          </a>
        </div>
      </div>
  </div>
</div>

  <div class="github-row">
    <p class="github-body-container">@ptillet recently rewrote the Triton runtime in https://github.com/openai/triton<span class="show-more-container"><a href="https://github.com/pytorch/torchdynamo/pull/1338" target="_blank" rel="noopener nofollow ugc" class="show-more">…</a></span><span class="excerpt hidden">/pull/644

This should dramatically reduce CPU overheads when cudagraphs is disabled.

This updates TorchInductor to use that new runtime.  We now call `triton.compile()`, and no longer use `triton.jit()`.  

There is also some early support for parallel compiles, but still need optimize that part.

Note this PR breaks support for Triton versions prior to 998fd5f9afe166247f441999c605dfe624ca9331.</span></p>
  </div>

  </article>

  <div class="onebox-metadata">
    
    
  </div>

  <div style="clear: both"></div>
</aside>
<p>
with a new defined kernel. Here is an example of the kernel’s signature:</p>
<pre><code class="lang-auto">def kernel_fma(
    C,  # Pointers to matrices
    ACT_INPUTS,
    A,
    B,
    bias,
    # Matrix dimensions
    M,
    N,
    K,
    CACHE_KEY_M,
    CACHE_KEY_N,
    CACHE_KEY_K,
    # The stride variables represent how much to increase the ptr by when moving by 1
    # element in a particular dimension. E.g. stride_am is how much to increase a_ptr
    # by to get the element one row down (A has M rows)
    stride_om,
    stride_on,
    stride_im,
    stride_ik,
    stride_wn,
    stride_wk,
    # Meta-parameters
    BLOCK_M: tl.constexpr,
    GROUP_M: tl.constexpr,
    BLOCK_N: tl.constexpr,
    BLOCK_K: tl.constexpr,
    # split k not used, not performant with activation, kept because early_config_prune is expecting it
    SPLIT_K: tl.constexpr,
    EVEN_K: tl.constexpr,
    BIAS: tl.constexpr,
    SAVE_ACT_INPUTS: tl.constexpr,
    ACTIVATION: tl.constexpr,
)
</code></pre>
<p>I introduced this decorator <img src="https://discuss.pytorch.org/images/emoji/apple/slight_smile.png?v=12" title=":slight_smile:" class="emoji" alt=":slight_smile:" loading="lazy" width="20" height="20"></p>
<pre><code class="lang-auto">def autotune(configs, meta, save_cache_hook=False):
    def decorator(fn):
        return CachingAutotuner(
            # force autotune by setting save_cache_hook to False
            fn,
            meta=meta,
            configs=configs,
            save_cache_hook=save_cache_hook,
        )

    return decorator

</code></pre>
<p>based on this example of test: <a href="https://github.com/pytorch/pytorch/blob/fae821c2f166fccab6a3c34e293c7268f61e82ba/test/inductor/test_torchinductor.py#L5821" class="inline-onebox" rel="noopener nofollow ugc">pytorch/test_torchinductor.py at fae821c2f166fccab6a3c34e293c7268f61e82ba · pytorch/pytorch · GitHub</a></p>
<p>But i thought it might be a better way to use the caching_autotune.</p>
<p>Thanks in advance,</p>