Autotuning kernel launch parameters

[AntonOresten]

There are often significant performance gains from e.g. increasing the tile size when a smaller type is used, or sometimes simply setting the entry hint occupancy=2 will have a big impact. It would be nice to be able to automatically find a good combination of parameters specific to the current array sizes, hardware, and arch.

See _autotuner.py with example usage in AttentionFMHA.py

[AntonOresten]

I've been looking into this. Regardless of whether it ends up here or in an external package, I think I would need access to CUDAExt.emit_function, and to be able to safely call it in parallel, such that compilation, warmup, and timing become distinct stages.

I’ve been prototyping an autotuning interface (similar to the Python one), and I’m trying to separate:

1. compile candidate configs
2. warm up candidates
3. benchmark candidates

Right now this seems to require extension internals. Before I go further: is this a reasonable direction?

From my local runs, tuning latency appears to be dominated by the tileiras subprocess step (in CUDAExt.emit_binary), so parallel compile would be useful. I’m less sure about the safety/perf tradeoffs of doing module/function load (CuModule/CuFunction) in parallel.

[maleadt]

Should be fine, but we probably need a better API just like @cuda launch=false. For now, especially if the tuner lives here, calling emit_function is fine.

[AntonOresten]

Makes sense, thanks. I’ll use emit_function for now while my prototype lives in CUDAExt.

A small launch-construction API (similar in spirit to @cuda launch=false) could make autotuning cleaner.
One idea is a macro where $ marks config-dependent expressions, e.g.:

• cld(M, $tm) expanding to cfg -> cld(M, cfg.tm) for grid construction.
• same with kernel args/hints, so we can write one concise launch template and evaluate it with the selected config.

I prefer this because juggling 5 different functions gets awkward. For larger expressions one could define the function explicitly or there could be a separate @withconfig macro that uses the same trick. This is similar to Makie's @lift.

[AntonOresten]

There are a lot of arguments to accommodate (⚙️ = depends on current config):

• kernel function
• config search space (+ sampling strategy, warmup, runs, etc)
• ⚙️ grid
• ⚙️ kernel arguments (often many, requiring multiple lines)
• ⚙️ entry hints (see bottom note)
• key for caching best config (e.g. tuple of head dim, data types, etc. depending on what warrants config re-tuning)

Currently I have a syntax resembling:

ct.autotune_launch

# candidate config is an element of the cartesian product (with field names).
# could have different types to embed search space and sampling strategies
search_space = (tm=(32, 64, 128, 256), tn=(32, 64, 128, 256), occupancy=(1, 2, 4, 8))
ct.autotune_launch(
    mha_fwd_kernel,
    search_space,
    cfg -> (cld(SeqLen_Q, cfg.tm), Heads * Batch),
    cfg -> (
        Q, K, V,
        qk_scale, input_pos, Heads,
        ct.Constant(Dk_pow2), ct.Constant(Dv_pow2),
        ct.Constant(cfg.tm), ct.Constant(cfg.tn),
        ct.Constant(query_group_size),
        ct.Constant(causal),
        ct.Constant(iszero(SeqLen_K % cfg.tn))
    );
    key = (T, Dk_pow2, Dv_pow2)
)

which is slightly more verbose than:

ct.launch

ct.launch(
    mha_fwd_kernel,
    (cld(SeqLen_Q, TILE_M), Heads * Batch),
    Q, K, V,
    qk_scale, input_pos, Heads,
    ct.Constant(Dk_pow2), ct.Constant(Dv_pow2),
    ct.Constant(TILE_M), ct.Constant(TILE_N),
    ct.Constant(query_group_size),
    ct.Constant(causal),
    ct.Constant(iszero(SeqLen_K % TILE_N))
)

A single line macro might not be able to accommodate all the options as the line quickly becomes congested:

autotune macro

search_space = (tm=(32, 64, 128, 256), tn=(32, 64, 128, 256), occupancy=(1, 2, 4, 8))
grid = cfg -> (cld(SeqLen_Q, cfg.tm), Heads * Batch)
key = (T, Dk_pow2, Dv_pow2)
@autotune search_space grid key mha_fwd_kernel(
    Q, K, V,
    qk_scale, input_pos, Heads,
    ct.Constant(Dk_pow2), ct.Constant(Dv_pow2),
    ct.Constant($tm), ct.Constant($tn), # $tm becomes cfg.tm
    ct.Constant(query_group_size),
    ct.Constant(causal),
    ct.Constant(iszero(SeqLen_K % $tn))
)

and with block syntax, maybe:

autotune macro in block

@autotune begin
    search_space = (tm=(32, 64, 128, 256), tn=(32, 64, 128, 256), occupancy=(1, 2, 4, 8))
    grid = (cld(SeqLen_Q, $tm), Heads * Batch)
    key = (T, Dk_pow2, Dv_pow2)

    mha_fwd_kernel(
        Q, K, V,
        qk_scale, input_pos, Heads,
        ct.Constant(Dk_pow2), ct.Constant(Dv_pow2),
        ct.Constant($tm), ct.Constant($tn),
        ct.Constant(query_group_size),
        ct.Constant(causal),
        ct.Constant(iszero(SeqLen_K % $tn))
    )
end

EDIT: I made the realization that the entry hints function hints_fn in the AttentionFMHA.py example is just an identity mapping of occupancy and num_ctas from the candidate config, so doesn't need to be explicitly passed. Autotuning calls became slightly more concise after removing it. Any hints_fn behavior could instead be done with a custom search space maybe.

[maleadt]

For reference, there's also https://github.com/KernelTuner/kernel_tuner, which IIRC a Julia back-end is being made for (cc @fjwillemsen)

[fjwillemsen]

@maleadt thanks for CCing me, that's right - we're working on Kernel Tuner as a Julia package, so not only will it have a Julia backend, you'll be able to interact with it entirely from Julia!

[AntonOresten]

Nice, looking forward to seeing the Julia package! The autotuner I'm prototyping does tuning at runtime, with parallel compilation (because tileiras is most often the bottleneck) and caching of the tuned config across calls.

I'd be curious to hear more about the scope of the Julia package, and in particular whether the search strategies could be used as a standalone component. Right now my prototype does random sampling with a refinement pass, so better search would be good.

[fjwillemsen]

@AntonOresten Awesome, that is great to hear! You might want to take a look at a nice runtime autotuner: Kernel Tuning Toolkit and related publications.

There are of course substantial distinctions in considerations between runtime and compile time (deployment vs development stage) auto-tuning; Kernel Tuner falls in the latter category. Regarding the search strategies, they can make a massive difference in the effectivity of the auto-tuning process (topic of my PhD in case you're interested). Kernel Tuner is quite modular; we've actually recently discussed making the SearchSpace and strategies standalone packages, as we've already done previously with the constraint resolution solver (which sees quite wide adoption), so if you have a good use case for that it could definitely weigh in.