WaveTune: GPU Auto-Tuning Framework

• WaveTune is a runtime auto-tuning framework for GPU kernels that optimizes tile-based operations in LLM inference by leveraging wave scheduling.
• It employs a unified mapping to a common physical space, decomposing configuration into macro and micro parameters with a wave-aware bilinear latency model.
• Empirical results demonstrate up to 1.83× kernel speedup and significant TTFT reduction across multiple GPUs while maintaining microsecond-level decision overhead.

WaveTune is a runtime auto-tuning framework for GPU kernels that targets the tile-based kernels dominating LLM 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 (Zhang et al., 11 Apr 2026).

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 (Zhang et al., 11 Apr 2026). 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 (Zhang et al., 11 Apr 2026).

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 (Zhang et al., 11 Apr 2026). 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 (Zhang et al., 11 Apr 2026). 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 (Zhang et al., 11 Apr 2026).

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 (Zhang et al., 11 Apr 2026).

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 (Zhang et al., 11 Apr 2026). 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 $1.33\times$0 the number of SMs, as a useful intuition but not an adequate final model (Zhang et al., 11 Apr 2026). 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 (Zhang et al., 11 Apr 2026).

To make diverse kernels comparable, WaveTune maps them into a common physical space $1.33\times$1, where $1.33\times$2 denotes spatial task decomposition and $1.33\times$3 denotes temporal work per task (Zhang et al., 11 Apr 2026). For dense GEMM with macro-config $1.33\times$4,

$1.33\times$5

For grouped GEMM with experts $1.33\times$6,

$1.33\times$7

For FlashAttention with macro-config $1.33\times$8,

$1.33\times$9

This mapping reduces kernel-specific diversity to the physical variables most directly linked to wave behavior and per-task loop work (Zhang et al., 11 Apr 2026).

WaveTune also decomposes the full configuration space as

$$G = \left\lceil \frac{M}{T_M} \right\rceil \times \left\lceil \frac{N}{T_N} \right\rceil,$$0

with a configuration written as

$$G = \left\lceil \frac{M}{T_M} \right\rceil \times \left\lceil \frac{N}{T_N} \right\rceil,$$1

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 (Zhang et al., 11 Apr 2026). Micro-configs refine local execution efficiency once the macro geometry is fixed and include the number of software pipeline stages $$G = \left\lceil \frac{M}{T_M} \right\rceil \times \left\lceil \frac{N}{T_N} \right\rceil,$$2, the number of warps $$G = \left\lceil \frac{M}{T_M} \right\rceil \times \left\lceil \frac{N}{T_N} \right\rceil,$$3, and rasterization or swizzle-related choices (Zhang et al., 11 Apr 2026). 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 = \left\lceil \frac{M}{T_M} \right\rceil \times \left\lceil \frac{N}{T_N} \right\rceil,$$4 and $$G = \left\lceil \frac{M}{T_M} \right\rceil \times \left\lceil \frac{N}{T_N} \right\rceil,$$5, but with coupled slopes. The global empirical form introduced in the paper is

$$G = \left\lceil \frac{M}{T_M} \right\rceil \times \left\lceil \frac{N}{T_N} \right\rceil,$$6

where the $$G = \left\lceil \frac{M}{T_M} \right\rceil \times \left\lceil \frac{N}{T_N} \right\rceil,$$7 term represents the dominant workload-volume contribution, $$G = \left\lceil \frac{M}{T_M} \right\rceil \times \left\lceil \frac{N}{T_N} \right\rceil,$$8 captures the marginal spatial cost, $$G = \left\lceil \frac{M}{T_M} \right\rceil \times \left\lceil \frac{N}{T_N} \right\rceil,$$9 captures the marginal temporal or loop cost, and $L = \left\lceil \frac{K}{T_K} \right\rceil.$0 is a fixed overhead (Zhang et al., 11 Apr 2026). The model is therefore bilinear rather than purely linear.

Wave awareness is incorporated by fitting this bilinear form separately for each pair $L = \left\lceil \frac{K}{T_K} \right\rceil.$1, where $L = \left\lceil \frac{K}{T_K} \right\rceil.$2 is the wave count induced by the runtime workload and hardware (Zhang et al., 11 Apr 2026). For each such bucket, WaveTune uses

$L = \left\lceil \frac{K}{T_K} \right\rceil.$3

with

$L = \left\lceil \frac{K}{T_K} \right\rceil.$4

The model is thus piecewise bilinear: piecewise across wave-count regions and bilinear within each region (Zhang et al., 11 Apr 2026).

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 (Zhang et al., 11 Apr 2026). In ablations on FlashAttention, the full WaveTune model reaches $L = \left\lceil \frac{K}{T_K} \right\rceil.$5 speedup, compared with $L = \left\lceil \frac{K}{T_K} \right\rceil.$6 for a linear-only variant and $L = \left\lceil \frac{K}{T_K} \right\rceil.$7 for a step-only variant (Zhang et al., 11 Apr 2026). This suggests that both components—inter-wave partitioning and intra-wave bilinear modeling—are necessary.

For wave counts beyond the profiled maximum $L = \left\lceil \frac{K}{T_K} \right\rceil.$8, WaveTune uses an extrapolation model: $L = \left\lceil \frac{K}{T_K} \right\rceil.$9 The extrapolation coefficients are fitted from the last $(T_M, T_N, T_K)$0 observed waves, with the paper giving $(T_M, T_N, T_K)$1 as an example (Zhang et al., 11 Apr 2026). 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 (Zhang et al., 11 Apr 2026). Let $(T_M, T_N, T_K)$2 denote the number of SMs, $(T_M, T_N, T_K)$3 the maximum profiled wave count, and $(T_M, T_N, T_K)$4 the number of sub-intervals per wave. For each wave count $(T_M, T_N, T_K)$5,

$(T_M, T_N, T_K)$6

Each wave region is partitioned into $(T_M, T_N, T_K)$7 sub-intervals, and one representative grid size $(T_M, T_N, T_K)$8 is selected from each, giving

$(T_M, T_N, T_K)$9

This preserves both the wave boundaries and the variation inside each wave (Zhang et al., 11 Apr 2026).

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$0

subject to aspect-ratio constraint

$G$1

with $G$2 set for example to $G$3 (Zhang et al., 11 Apr 2026). This excludes pathological skinny matrices that would distort memory behavior. For attention kernels, representatives are aligned to the head count: $G$4 Loop counts are sampled separately through a set $G$5, and the profiling anchors are all pairs $G$6 with $G$7 and $G$8 (Zhang et al., 11 Apr 2026).

These anchors are then mapped back to executable problem sizes. For GEMM-like kernels with macro-config $G$9,

$L$0

For attention with macro-config $L$1,

$L$2

For each anchor and macro-config, all feasible micro-configs are benchmarked using 3 warm-up iterations and 5 measured iterations through PyTorch Profiler (Zhang et al., 11 Apr 2026).

The total number of profiled points is

$L$3

This is still substantial but far smaller than exhaustive search over the full online space (Zhang et al., 11 Apr 2026). 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 $L$4hardware, kernel$L$5 pair (Zhang et al., 11 Apr 2026). 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 (Zhang et al., 11 Apr 2026). The first table is a coefficient table indexed by $L$6, storing the bilinear coefficients

$L$7

The second is an anchor-based micro-config table indexed by $L$8, storing a shared micro-config selected for that bucket (Zhang et al., 11 Apr 2026). Importantly, the stored micro-config is not pointwise-optimal for each sampled $L$9; instead, within each bucket WaveTune selects a shared $10^{15}$0 minimizing average latency across the sampled grid sizes (Zhang et al., 11 Apr 2026). This makes the retrieved choice locally robust rather than overfit to individual anchors.

At runtime, WaveTune solves

$10^{15}$1

but implements it in two stages (Zhang et al., 11 Apr 2026). First, for each candidate macro-config it maps the input to $10^{15}$2, computes the wave count $10^{15}$3, retrieves the corresponding coefficients, evaluates the bilinear latency estimate, and selects

$10^{15}$4

Second, given $10^{15}$5, it retrieves the nearest stored loop anchor

$10^{15}$6

and returns the associated cached micro-config (Zhang et al., 11 Apr 2026).

This changes online complexity from multiplicative to additive: $10^{15}$7 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 $10^{15}$8 (Zhang et al., 11 Apr 2026).

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 (Zhang et al., 11 Apr 2026). 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 (Zhang et al., 11 Apr 2026).

The principal quantitative results are summarized below.

Setting	Reported result	Source
Kernel-level speedup	Up to $10^{15}$9	(Zhang et al., 11 Apr 2026)
End-to-end TTFT reduction	Up to $$T_{\text{total}} \approx T_{\text{wave}} \cdot \left\lceil \frac{G}{N_{SM}} \right\rceil,$$0	(Zhang et al., 11 Apr 2026)
Decision overhead	5–6 $$T_{\text{total}} \approx T_{\text{wave}} \cdot \left\lceil \frac{G}{N_{SM}} \right\rceil,$$1	(Zhang et al., 11 Apr 2026)
Storage per hardware-kernel pair	0.09 MB average	(Zhang et al., 11 Apr 2026)

Against default heuristics, the reported geometric-mean speedups are $$T_{\text{total}} \approx T_{\text{wave}} \cdot \left\lceil \frac{G}{N_{SM}} \right\rceil,$$2 for FlashAttention, $$T_{\text{total}} \approx T_{\text{wave}} \cdot \left\lceil \frac{G}{N_{SM}} \right\rceil,$$3 for Grouped GEMM, and $$T_{\text{total}} \approx T_{\text{wave}} \cdot \left\lceil \frac{G}{N_{SM}} \right\rceil,$$4 for Dense GEMM, with corresponding oracle upper bounds of $$T_{\text{total}} \approx T_{\text{wave}} \cdot \left\lceil \frac{G}{N_{SM}} \right\rceil,$$5, $$T_{\text{total}} \approx T_{\text{wave}} \cdot \left\lceil \frac{G}{N_{SM}} \right\rceil,$$6, and $$T_{\text{total}} \approx T_{\text{wave}} \cdot \left\lceil \frac{G}{N_{SM}} \right\rceil,$$7 respectively (Zhang et al., 11 Apr 2026). 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 $$T_{\text{total}} \approx T_{\text{wave}} \cdot \left\lceil \frac{G}{N_{SM}} \right\rceil,$$8 TTFT speedup on MI355X and up to $$T_{\text{total}} \approx T_{\text{wave}} \cdot \left\lceil \frac{G}{N_{SM}} \right\rceil,$$9 on B200 (Zhang et al., 11 Apr 2026). 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 (Zhang et al., 11 Apr 2026). 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 $G$0 decision time; decision-tree dispatch is around 60 $G$1; and even the native DeepGEMM heuristic requires 76 $G$2 in one reported case, compared with 6 $G$3 for WaveTune (Zhang et al., 11 Apr 2026). 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 (Zhang et al., 11 Apr 2026).

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 (Zhang et al., 11 Apr 2026). Models are trained per device, so portability is methodological rather than parameter-free: each $G$4hardware, kernel$G$5 pair requires some offline profiling (Zhang et al., 11 Apr 2026). 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 (Zhang et al., 11 Apr 2026).

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 (Zhang et al., 11 Apr 2026). 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 (Zhang et al., 11 Apr 2026).

In that sense, WaveTune is less a generic ML cost model than a hybrid analytical systems model. Its unified $G$6 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 (Zhang et al., 11 Apr 2026). 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 (Zhang et al., 11 Apr 2026).