Fused Matmul-Epilogue Sampling Techniques

• The paper demonstrates that fusing sampling operations into the matmul kernel eliminates intermediate data storage, reducing latency and memory traffic.
• This method employs advanced scheduling, tiling, and register-level dataflow optimizations to seamlessly integrate nonlinear post-processing steps into computation.
• Empirical benchmarks on GPUs reveal significant speedups in language model decoding, underscoring practical improvements in efficiency and performance.

Fused matmul-epilogue sampling refers to a class of GPU and tensor-accelerator techniques that integrate random sampling or other nonlinear pointwise transformations directly into the matrix multiplication (matmul) kernel, eliminating the need to materialize intermediate result matrices (such as logits) in device memory. By fusing complex epilogues—particularly sampling operations required in LLM decoding—into the matmul operation itself, these methods overcome the bandwidth and launch overheads associated with standard multi-step pipelines. This approach is enabled by scheduling and compiler infrastructures capable of precise control over computation tiling, register-level dataflow, and kernel epilogue customization (Zhang et al., 2 Dec 2025, Ruiz et al., 16 Mar 2026).

1. Foundations of Fused Matmul-Epilogue Pipelines

Traditional accelerator workflows separate the dense matrix multiplication and subsequent post-processing steps (e.g., activation, bias addition, softmax, sampling). The fused matmul-epilogue paradigm collapses these stages by scheduling post-matmul operations—such as nonlinear activations, interpolation, or random sampling—directly within the output tile computation, before any data is spilled from on-chip registers back to high-bandwidth memory (HBM). The canonical example in this context is the integration of categorical sampling, specifically the Gumbel-Max trick for exact sampling, into the transformer LLM head matmul, thereby bypassing intermediate storage of dense logits (Ruiz et al., 16 Mar 2026).

This approach is generalizable to any pointwise epilogue, encompassing nonlinearities (e.g., ReLU, GELU), interpolation, or per-element random number generation, as long as the epilogue can be computed at the register or tile granularity (Zhang et al., 2 Dec 2025).

2. Algorithmic Structure and Epilogue Fusion Strategies

A typical fused matmul-epilogue kernel is structured around on-chip tiling and tile-level accumulation. The matmul is realized as a sequence of block-level and warp-level tiling operations, with per-tile accumulators held in specialized register fragments (e.g., WMMAAccumulator registers). The matmul pipeline for a row-major output $C(i,j)$ can be represented as:

$$C(i, j) = \text{sample}\left(\sum_k A(i, k) B(k, j) + \text{bias}(j)\right)$$

where $\text{sample}(\cdot)$ is any feasible elementwise function, including stochastic samplers. In the specific case of categorical sampling via Gumbel-Max, the fused operation becomes:

$$s_{b,i} = \tilde y_{b,i} + g_{b,i}, \qquad g_{b,i} \sim \text{Gumbel}(0,1)$$

and one retains the $\arg\max_i s_{b,i}$ as the sample for each row $b$ (Ruiz et al., 16 Mar 2026).

Kernel scheduling primitives (as illustrated in Halide) tile output loops, map tiles to GPU blocks/warps, and select the compute granularity for MMA (matrix multiply-accumulate) instructions. Per-tile epilogue code is inserted immediately after MMA accumulation, ensuring all nonlinearity, transformation, or sampling occurs while accumulator fragments are still resident in fast registers (Zhang et al., 2 Dec 2025).

3. Instruction Selection, Compiler Techniques, and Hardware Mapping

The transformation of fused pipelines into efficient hardware code is facilitated by equality saturation-based instruction selectors (e.g., the "HardBoiled" selector). Following vectorization and tiling, the intermediate representation is analyzed to detect patterns corresponding to hardware-specific MMA instruction usage. These selectors apply rule classes including normalization of data layout (ramp/broadcast normalization), memory layout detection (e.g., VNNI), and lowering to intrinsic calls for wmma.load.{a,b}.sync and wmma.mma.sync.

Register allocation and producer-consumer fusion are managed to avoid global memory round-trips. Tensor instruction selectors ensure that register-resident accumulators (e.g., WMMAAccumulator) directly feed into the fused epilogue before storage, maximizing on-chip data residency and minimizing memory traffic (Zhang et al., 2 Dec 2025).

On platforms that expose custom epilogue capabilities (e.g., via Triton or user-schedulable Halide backends), developers express both matmul and epilogue seamlessly for compilation into single fused device kernels.

4. Complexity, Memory Traffic, and Computational Intensity

Fused matmul-epilogue sampling significantly improves computational efficiency by minimizing external memory transfers. In the unfused variant, the GEMM writes and reads the full $[B,V]$ logits matrix, incurring $2BV$ elementwise transfers in BF16 ($\approx4BV$ bytes). The fused approach writes only per-tile or per-row-maximizer candidates (score/index pairs), a drastically smaller volume.

Arithmetic intensity for the fused operator is:

$$I_{\rm fused} \approx \frac{2BVD}{2(VD+BD)} = \frac{BV}{V+B} \quad [\text{FLOP/byte}]$$

which outperforms the non-fused baseline, particularly in small-batch, large-vocabulary (decode) regimes. The main runtime improvement arises from eliminating multiple kernel launches, device-host synchronization points, and excessive global memory bandwidth utilization (Ruiz et al., 16 Mar 2026).

5. Empirical Performance and Application Benchmarks

Empirical evaluations demonstrate considerable performance improvements across key machine learning pipelines:

• On NVIDIA A100 GPUs, a fused GEMM (1024×1024) + bias + activation kernel achieves $3.4\times$ speedup over unfused variants (66 μs vs. 223 μs), closely approaching cuBLASLt performance (Zhang et al., 2 Dec 2025).
• Fused convolutional layers (4k×64×64×16, with bias and ReLU) and downsampling routines realize $3.7\times$–$6.1\times$ speedups relative to PyTorch and CUDA baselines (Zhang et al., 2 Dec 2025).
• In the language modeling context, FlashSampling achieves up to a $19\%$ reduction in time-per-output-token for small- to medium-size models (e.g., Qwen3-1.7B), with kernel-level speedups up to $1.84\times$ against PyTorch multinomial-softmax and up to $2.52\times$ against FlashInfer top-$k$/$p$ (Ruiz et al., 16 Mar 2026).
• Performance benefits decrease for larger batch sizes ($B \geq 128$) or models where the head matmul is not the bottleneck.

These results validate the effectiveness of full fusion not only in theory but also for practical, on-device inference scenarios.

6. Extensions, Epilogue Generality, and Integration Considerations

Fused matmul-epilogue sampling supports any pointwise or tile-local epilogue, including:

• Activation functions (ReLU, GELU, sigmoid),
• Random sampling (e.g., generated via Philox/Xorshift on per-lane basis),
• Mixed-precision epilogues, with accumulation in FP32 and conversion/epilogue in lower precision as allowed,
• Gumbel-Top-$k$ and Top-$p$ sampling, extendable via postprocessing over per-row candidates.

For random sampling, counter-based RNG ensures reproducibility and correct parallelization for Gumbel noise injection. Numerical stability is maintained by FP32 accumulation for both GEMM and RNG stages.

A crucial integration constraint is the requirement for a GEMM backend that allows custom epilogues; not all vendor libraries (e.g., stock cuBLAS) expose such functionality, though Triton and many modern inference stack implementations do (Ruiz et al., 16 Mar 2026). Pipelines must be adapted to permit custom kernel injection at the head matmul.

7. Significance and Limitations

Fused matmul-epilogue sampling transforms a traditionally memory-bound postprocessing workload into a lightweight on-chip operation. By eliminating intermediate tensor materializations and kernel launches, it leads to reduced latency and more scalable hardware utilization, particularly in memory-bandwidth-limited settings such as LLM decoding. However, implementation is nontrivial in platforms that lack programmable epilogues, and the benefits are maximal in regimes where output tile sizes and activation/sampling costs are balanced with compute throughput. For expensive epilogues, vectorized or staged register-based evaluation may be required to avoid register pressure or stall cycles (Zhang et al., 2 Dec 2025, Ruiz et al., 16 Mar 2026).

These techniques represent a convergence of compiler scheduling, hardware tensorization, and stochastic algorithmic insight, demonstrating the potential for tensor accelerators to efficiently address a wide spectrum of postprocessing workloads beyond classical matrix multiplication.