---
title: 'WaveTune: GPU Auto-Tuning Framework'
url: https://www.emergentmind.com/topics/wavetune
type: topic
---

# WaveTune: GPU Auto-Tuning Framework

WaveTune is a runtime auto-tuning framework for GPU kernels that targets the tile-based kernels dominating large language model inference, especially GEMM, grouped GEMM for mixture-of-experts, and FlashAttention. Its defining premise is that kernel latency is shaped by GPU wave scheduling rather than behaving as an arbitrary black-box function of input shape and configuration. On that basis, WaveTune combines a unified mapping from kernel-specific inputs into a common physical space, a decomposition of configuration space into macro and micro parameters, a wave-aware piecewise bilinear latency model, sparse wave-structured profiling, and a dual-table runtime retrieval mechanism. Across three representative kernels and five GPU architectures, it is reported to achieve near-optimal kernel performance while reducing runtime decision overhead by five orders of magnitude relative to exhaustive search, with up to \(1.83\times\) kernel-level speedup and up to \(1.33\times\) end-to-end TTFT reduction [2604.10187].

## 1. Problem formulation and design objective

WaveTune addresses runtime configuration of GPU kernels in LLM inference, where performance is dominated by tile-based kernels and remains sensitive to runtime parameters such as tile sizes, pipeline stages, number of warps, swizzling, and CTA structure [2604.10187]. The central difficulty is that these parameters interact with GPU resources including registers, shared memory, occupancy, and SM-level concurrency, producing a non-convex optimization landscape rather than a smooth or easily searchable response surface [2604.10187].

For tiled GEMM, the total number of blocks is
\[
G = \left\lceil \frac{M}{T_M} \right\rceil \times \left\lceil \frac{N}{T_N} \right\rceil,
\]
and the reduction-loop count is
\[
L = \left\lceil \frac{K}{T_K} \right\rceil.
\]
Changing \((T_M, T_N, T_K)\) therefore changes per-block work, grid size \(G\), loop count \(L\), the number of execution waves, and per-block resource usage simultaneously [2604.10187]. The same logic extends to grouped GEMM and FlashAttention through kernel-specific tiling variables.

The framework is motivated by the limitations of three established families of methods. Exhaustive or search-based auto-tuning can find excellent configurations but is too expensive for online use and cannot practically cover the full dynamic serving space; the paper notes that in DeepGEMM the joint space of input dimensions and configuration choices can reach on the order of \(10^{15}\) combinations [2604.10187]. Expert heuristics are runtime-cheap but brittle outside their calibration regime. Learned cost models, including XGBoost and decision-tree dispatch policies, can be accurate enough offline but still incur substantial online overhead because they require repeated evaluation across candidate configurations, with reported runtime decision latencies in the millisecond or high-microsecond range [2604.10187].

WaveTune is therefore constructed to eliminate the conventional trade-off between configuration quality and runtime overhead: its goal is to approach exhaustive-search performance while retaining microsecond-level decision cost [2604.10187].

## 2. Execution model: waves, physical coordinates, and configuration decomposition

The framework is explicitly wave-aware. A GPU kernel launches a grid of CTAs, but each SM can host only a limited number of blocks concurrently, constrained by hardware CTA slots, register pressure, and shared-memory usage. When the total grid exceeds concurrent capacity, execution proceeds in rounds, or waves [2604.10187]. The paper discusses the classical step approximation
\[
T_{\text{total}} \approx T_{\text{wave}} \cdot \left\lceil \frac{G}{N_{SM}} \right\rceil,
\]
with \(G\) the grid size and \(N_{SM}\) the number of SMs, as a useful intuition but not an adequate final model [2604.10187]. Crossing a wave boundary introduces latency discontinuities, especially through underfilled tail waves, yet empirical profiling shows that real kernels are not purely stepwise because block-runtime variability and greedy scheduling smooth the boundary behavior [2604.10187].

To make diverse kernels comparable, WaveTune maps them into a common physical space \((G,L)\), where \(G\) denotes spatial task decomposition and \(L\) denotes temporal work per task [2604.10187]. For dense GEMM with macro-config \(\langle T_M, T_N, T_K \rangle\),
\[
G = \left\lceil \frac{M}{T_M} \right\rceil \times \left\lceil \frac{N}{T_N} \right\rceil, \qquad
L = \left\lceil \frac{K}{T_K} \right\rceil.
\]
For grouped GEMM with experts \(i=1,\dots,E\),
\[
G = \sum_{i=1}^{E} \left( \left\lceil \frac{M_i}{T_M} \right\rceil \times \left\lceil \frac{N}{T_N} \right\rceil \right), \qquad
L = \left\lceil \frac{K}{T_K} \right\rceil.
\]
For FlashAttention with macro-config \(\langle T_Q, T_{KV} \rangle\),
\[
G = N_h \times \left\lceil \frac{S_q}{T_Q} \right\rceil, \qquad
L = \left\lceil \frac{S_{kv}}{T_{KV}} \right\rceil.
\]
This mapping reduces kernel-specific diversity to the physical variables most directly linked to wave behavior and per-task loop work [2604.10187].

WaveTune also decomposes the full configuration space as
\[
\mathcal{C} = \mathcal{C}_{\text{macro}} \times \mathcal{C}_{\text{micro}},
\]
with a configuration written as
\[
\mathbf{c} = (\mathbf{c}_{\text{macro}}, \mathbf{c}_{\text{micro}}).
\]
Macro-configs determine workload geometry, wave structure, grid size, and loop count; for GEMM these are the tile dimensions, and for attention they are the query and key-value tile sizes [2604.10187]. Micro-configs refine local execution efficiency once the macro geometry is fixed and include the number of software pipeline stages \(N_{stg}\), the number of warps \(N_{wrp}\), and rasterization or swizzle-related choices [2604.10187]. This hierarchical split is essential because it allows WaveTune to model structural latency at the macro level while caching micro-config decisions through anchors rather than searching the full Cartesian product at runtime.

## 3. Analytical wave-aware bilinear latency model

The core latency model rests on the observation that, within a fixed wave regime, kernel latency exhibits strong linearity in both \(G\) and \(L\), but with coupled slopes. The global empirical form introduced in the paper is
\[
T(G, L) \approx \alpha G L + \beta G + \gamma L + \delta,
\]
where the \(\alpha G L\) term represents the dominant workload-volume contribution, \(\beta G\) captures the marginal spatial cost, \(\gamma L\) captures the marginal temporal or loop cost, and \(\delta\) is a fixed overhead [2604.10187]. The model is therefore bilinear rather than purely linear.

Wave awareness is incorporated by fitting this bilinear form separately for each pair \(\langle \mathbf{c}_{\text{macro}}, w \rangle\), where \(w\) is the wave count induced by the runtime workload and hardware [2604.10187]. For each such bucket, WaveTune uses
\[
\hat{T}(G,L \mid \boldsymbol{\theta}_{\mathbf{c}_{\text{macro}}, w}) =
\alpha GL + \beta G + \gamma L + \delta,
\]
with
\[
\boldsymbol{\theta}_{\mathbf{c}_{\text{macro}}, w} =
\langle \alpha, \beta, \gamma, \delta \rangle.
\]
The model is thus piecewise bilinear: piecewise across wave-count regions and bilinear within each region [2604.10187].

This structure is the basis for the framework’s claim to outperform both a pure linear approximation and a pure step model. A global linear model cannot represent wave-boundary discontinuities, while a step model ignores the substantial within-wave variation that remains after fixing wave count [2604.10187]. In ablations on FlashAttention, the full WaveTune model reaches \(1.83\times\) speedup, compared with \(1.60\times\) for a linear-only variant and \(1.39\times\) for a step-only variant [2604.10187]. This suggests that both components—inter-wave partitioning and intra-wave bilinear modeling—are necessary.

For wave counts beyond the profiled maximum \(W\), WaveTune uses an extrapolation model:
\[
\hat{T}_{\mathbf{c}_{\text{macro}}}(G,L,w)=
\begin{cases}
\hat{T}(G,L \mid \boldsymbol{\theta}_{\mathbf{c}_{\text{macro}}, w}), & w \le W,\\[1ex]
\hat{T}(G,L \mid \boldsymbol{\theta}^{\text{ext}}_{\mathbf{c}_{\text{macro}}}), & w > W.
\end{cases}
\]
The extrapolation coefficients are fitted from the last \(p\) observed waves, with the paper giving \(p=10\) as an example [2604.10187]. This suggests that larger-wave behavior is sufficiently regular for stable asymptotic continuation.

## 4. Sparse profiling and wave-structured sampling

WaveTune does not fit its model from dense profiling. Instead, it constructs a sparse sampling grid aligned with wave regions [2604.10187]. Let \(N_{SM}\) denote the number of SMs, \(W\) the maximum profiled wave count, and \(I\) the number of sub-intervals per wave. For each wave count \(w\),
\[
\mathcal{R}_w = [(w-1)N_{SM}+1,\; wN_{SM}], \quad w=1,\dots,W.
\]
Each wave region is partitioned into \(I\) sub-intervals, and one representative grid size \(g_{w,i}\) is selected from each, giving
\[
\mathcal{G} = \{ g_{w,i} \mid w=1,\dots,W,\; i=1,\dots,I \}.
\]
This preserves both the wave boundaries and the variation inside each wave [2604.10187].

For GEMM-like kernels, the representative is chosen as the largest value in the sub-interval that can be factorized into a valid 2D block grid
\[
g_{w,i} = m_G n_G, \qquad m_G \le n_G,
\]
subject to aspect-ratio constraint
\[
\rho(g_{w,i}) = \frac{n_G}{m_G} \le \tau,
\]
with \(\tau\) set for example to \(1.1\) [2604.10187]. This excludes pathological skinny matrices that would distort memory behavior. For attention kernels, representatives are aligned to the head count:
\[
g_{w,i} = \left\lfloor \frac{b_{w,i}}{N_h} \right\rfloor \cdot N_h, \qquad g_{w,i} \ge a_{w,i}.
\]
Loop counts are sampled separately through a set \(\mathcal{L}\), and the profiling anchors are all pairs \(\langle G,L \rangle\) with \(G \in \mathcal{G}\) and \(L \in \mathcal{L}\) [2604.10187].

These anchors are then mapped back to executable problem sizes. For GEMM-like kernels with macro-config \(\langle T_M,T_N,T_K\rangle\),
\[
M = m_G T_M,\qquad N = n_G T_N,\qquad K = L T_K.
\]
For attention with macro-config \(\langle T_Q, T_{KV}\rangle\),
\[
S_q = \frac{G}{N_h} T_Q,\qquad S_{kv} = L T_{KV}.
\]
For each anchor and macro-config, all feasible micro-configs are benchmarked using 3 warm-up iterations and 5 measured iterations through PyTorch Profiler [2604.10187].

The total number of profiled points is
\[
N_{\text{profile}} =
|\mathcal{G}| \cdot |\mathcal{L}| \cdot
|\mathcal{C}_{\text{macro}}| \cdot
|\mathcal{C}_{\text{micro}}|.
\]
This is still substantial but far smaller than exhaustive search over the full online space [2604.10187]. The paper reports approximate profiling volumes per hardware-kernel pair of 74K configurations for Dense GEMM, 73K for Grouped GEMM, and 6.6K for FlashAttention, corresponding to about 0.5 to 2 GPU-hours per \(\langle\)hardware, kernel\(\rangle\) pair [2604.10187]. This suggests an offline cost compatible with practical deployment pipelines.

## 5. Dual-table runtime selection

WaveTune’s low online cost derives from a dual-table retrieval design [2604.10187]. The first table is a coefficient table indexed by \(\langle \mathbf{c}_{\text{macro}}, w \rangle\), storing the bilinear coefficients
\[
\boldsymbol{\theta}_{\mathbf{c}_{\text{macro}}, w}
= \langle \alpha,\beta,\gamma,\delta \rangle.
\]
The second is an anchor-based micro-config table indexed by \(\langle \mathbf{c}_{\text{macro}}, w, L \rangle\), storing a shared micro-config selected for that bucket [2604.10187]. Importantly, the stored micro-config is not pointwise-optimal for each sampled \(G\); instead, within each bucket WaveTune selects a shared \(\mathbf{c}^{*}_{\text{micro}}(L)\) minimizing average latency across the sampled grid sizes [2604.10187]. This makes the retrieved choice locally robust rather than overfit to individual anchors.

At runtime, WaveTune solves
\[
\mathbf{c}^*(\mathbf{x}) = \operatorname*{argmin}_{\mathbf{c} \in \mathcal{C}} T(\mathcal{K}, \mathbf{x}, \mathbf{c}),
\]
but implements it in two stages [2604.10187]. First, for each candidate macro-config it maps the input to \((G,L)\), computes the wave count \(w\), retrieves the corresponding coefficients, evaluates the bilinear latency estimate, and selects
\[
\mathbf{c}_{\text{macro}}^{*} =
\arg\min_{\mathbf{c}_{\text{macro}} \in \mathcal{C}_{\text{macro}}}
\hat{T}(G, L \mid \boldsymbol{\theta}_{\mathbf{c}_{\text{macro}}, w}).
\]
Second, given \(\mathbf{c}_{\text{macro}}^{*}\), it retrieves the nearest stored loop anchor
\[
\tilde{L} =
\arg\min_{L_i \in \mathcal{L}_{\mathbf{c}_{\text{macro}}^{*}, w^{*}}}
|L^{*} - L_i|,
\]
and returns the associated cached micro-config [2604.10187].

This changes online complexity from multiplicative to additive:
\[
O(|\mathcal{C}_{\text{macro}}| \cdot |\mathcal{C}_{\text{micro}}|)
\quad \text{to} \quad
O(|\mathcal{C}_{\text{macro}}| + \log |\mathcal{L}|).
\]
Because the online path consists only of coordinate mapping, coefficient lookup, a few arithmetic evaluations, and nearest-anchor retrieval, the reported decision overhead is only 5–6 \(\mu s\) [2604.10187].

## 6. Empirical results, scope, and limitations

WaveTune is evaluated on three kernel classes—Dense GEMM, Grouped GEMM, and FlashAttention—across five GPUs: NVIDIA A100, NVIDIA H20, NVIDIA B200, AMD MI308X, and AMD MI355X [2604.10187]. Training data include 26K Dense GEMM test samples, 7.1K Grouped GEMM test samples, and 2.8K FlashAttention test samples, with brute-force oracle generation costing about 10 to 600 GPU-hours depending on the case [2604.10187].

The principal quantitative results are summarized below.

| Setting | Reported result | Source |
|---|---:|---|
| Kernel-level speedup | Up to \(1.83\times\) | [2604.10187] |
| End-to-end TTFT reduction | Up to \(1.33\times\) | [2604.10187] |
| Decision overhead | 5–6 \(\mu s\) | [2604.10187] |
| Storage per hardware-kernel pair | 0.09 MB average | [2604.10187] |

Against default heuristics, the reported geometric-mean speedups are \(1.83\times\) for FlashAttention, \(1.27\times\) for Grouped GEMM, and \(1.04\times\) for Dense GEMM, with corresponding oracle upper bounds of \(2.00\times\), \(1.46\times\), and \(1.07\times\) respectively [2604.10187]. Dense GEMM therefore shows relatively limited remaining headroom, but WaveTune still approaches oracle without regression.

At the serving level, when integrated into SGLang for Qwen3-30B-A3B mixture-of-experts inference, WaveTune yields up to \(1.33\times\) TTFT speedup on MI355X and up to \(1.28\times\) on B200 [2604.10187]. The paper attributes these gains mainly to prefill, where kernels are large and configuration-sensitive, rather than decode, where kernels are smaller and offer less tuning opportunity [2604.10187]. This suggests that WaveTune is most valuable in latency-critical, shape-diverse online inference where prefill dominates responsiveness.

Relative to learned baselines, WaveTune is reported to match or exceed their configuration quality while greatly reducing online cost. XGBoost-based models incur 1,822–2,965 \(\mu s\) decision time; decision-tree dispatch is around 60 \(\mu s\); and even the native DeepGEMM heuristic requires 76 \(\mu s\) in one reported case, compared with 6 \(\mu s\) for WaveTune [2604.10187]. The storage footprint is correspondingly smaller: 0.09 MB average per hardware-kernel pair, compared with 15.25 MB for a decision tree and 38.78 MB for XGBoost [2604.10187].

The framework’s limitations are also explicit. Grouped GEMM requires an approximation because exact per-expert token distributions are not queried at runtime; instead, a uniform approximation across experts is used to avoid synchronization that would break CPU–GPU asynchrony [2604.10187]. Models are trained per device, so portability is methodological rather than parameter-free: each \(\langle\)hardware, kernel\(\rangle\) pair requires some offline profiling [2604.10187]. Finally, the method is targeted at tile-based kernels with meaningful macro/micro decomposition and is not presented as a universal solution for arbitrary lightweight kernels [2604.10187].

A plausible implication is that WaveTune’s strongest domain is not generic kernel optimization but structured, high-impact kernels whose latency is jointly governed by discrete wave effects and continuous workload scaling. Within that scope, the paper presents it as a practical runtime system rather than merely a predictive model.

## 7. Position within auto-tuning research

WaveTune’s contribution is best understood as a shift from feature-heavy black-box prediction toward analytical modeling grounded in GPU execution structure [2604.10187]. Search-based tuning remains the oracle baseline for optimality, but WaveTune shows that much of the benefit can be recovered if the latency surface is factored according to wave count and workload geometry. Heuristic rules remain attractive for simplicity, but WaveTune indicates that the relevant structure is richer than fixed hand-coded rules can easily capture. Learned models remain flexible, yet WaveTune argues that in critical serving paths their online evaluation cost is itself part of the systems problem [2604.10187].

In that sense, WaveTune is less a generic ML cost model than a hybrid analytical systems model. Its unified \((G,L)\) mapping, wave-conditioned bilinear fitting, sparse sampling, and dual-table retrieval jointly define a runtime tuning architecture in which execution physics, not just statistical regression, determines the representation [2604.10187]. This suggests a broader methodological direction: when scheduling artifacts induce piecewise structure, explicitly modeling those artifacts may outperform both end-to-end search and generic learned predictors.

WaveTune therefore occupies a distinct position in GPU auto-tuning: it is a runtime framework that treats wave scheduling as the primary organizing principle for latency prediction and configuration selection. Its empirical results indicate that this principle is sufficient to recover near-optimal behavior for major LLM inference kernels while keeping online overhead at microsecond scale [2604.10187].

Source: https://www.emergentmind.com/topics/wavetune