GEMM-Oriented Mode Reordering Techniques

• GEMM-oriented mode reordering is a method that reorganizes tensor modes and operator sequences to map complex computations directly to optimized matrix multiplications.
• It employs deterministic passes and algebraic reparameterizations to ensure contiguous memory layouts and GEMM-friendly dimensions across tensor contractions, transformer systems, and LLM pruning.
• By aligning computation with GEMM paradigms, this approach enhances GPU throughput, minimizes data movement, and improves energy-delay performance through reduced runtime transposes.

GEMM-oriented mode reordering is a family of techniques that reorganize tensor modes, operator sequences, loop traversals, or pruned computational axes so that the resulting computation maps cleanly to general matrix multiplication (GEMM), preserves or increases arithmetic intensity, and reduces data movement. In tensor-network contraction, it is a global, path-wide reordering that makes every pairwise contraction look like a matrix multiplication by placing retained modes first and reduced modes last in each operand (Pan et al., 1 Jun 2026). In transformer systems, it is the practice of algebraically reparameterizing transformer block computations so that memory-bound operators execute inside GEMM epilogues while the GEMM output tile is resident on chip (Guo et al., 19 May 2026). In LLM pruning, it is the practice of reorganizing tokens, channels, heads, and layers so that the remaining work maps cleanly to the logical $M$, $N$, and $K$ dimensions of GEMM (Hu et al., 8 Jun 2026). Related formulations also reorder GEMM traversal and mapping itself, including generalized space filling curves for communication-avoiding matrix multiplication (Georganas et al., 22 Jan 2026), GEMM-like tensor contractions without global transposes (Springer et al., 2016), and analytical optimization of $M/N/K$ walking axes and residency on spatial accelerators (Yang et al., 9 Mar 2026).

1. Conceptual basis and formal structure

Across these formulations, the central operation is the reduction of a higher-level computation to GEMM-compatible dimensions. A generic tensor contraction can be written with output modes that appear only in the left operand, output modes that appear only in the right operand, and contracted modes shared between both operands. With extents $\{d_x\}$, one defines $m = \prod_{x\in M} d_x$, $n = \prod_{x\in N} d_x$, and $k = \prod_{x\in K} d_x$, then permutes the operands so that $A' = [M \mid K]$ and $B' = [K \mid N]$, yielding matrices $N$0, $N$1, and $N$2, with the contraction becoming the matrix product $N$3 (Pan et al., 1 Jun 2026). GETT expresses the same principle by selecting three index groups, $N$4, $N$5, and $N$6, so that $N$7, and then defining permutations $N$8, $N$9, and $K$0 that reorder the modes of $K$1, $K$2, and $K$3 as $K$4, $K$5, and $K$6 respectively (Springer et al., 2016).

The same GEMM-centric abstraction appears outside tensor contraction. In transformer layers, the fixed GEMM mainloop stages $K$7 and $K$8 tiles into shared memory and registers, accumulates into FP32 accumulators using Tensor Cores, and produces register-resident output fragments $K$9 per thread or warp; surrounding operators are then expressed as constrained epilogue primitives attached to that mainloop (Guo et al., 19 May 2026). In LLM pruning, the GEMM-centric taxonomy assigns token dimension to $M/N/K$0, input feature or reduction dimension to $M/N/K$1, and output feature dimension to $M/N/K$2, with a general linear projection written as $M/N/K$3 with $M/N/K$4, $M/N/K$5, and $M/N/K$6, mapping $M/N/K$7, $M/N/K$8, and $M/N/K$9 (Hu et al., 8 Jun 2026).

A common analytical rationale is roofline behavior. Arithmetic intensity for one GEMM is $\{d_x\}$0, and throughput follows $\{d_x\}$1 (Pan et al., 1 Jun 2026). In the pruning taxonomy, the same style of model is written as $\{d_x\}$2, with $\{d_x\}$3, emphasizing that realized speedup depends on whether reduced shapes remain well-tiled and whether workloads remain compute-bound or become memory-bound (Hu et al., 8 Jun 2026). This suggests that GEMM-oriented mode reordering is not only a syntactic conversion to matrix multiplication, but also a deliberate reshaping of dataflow so that the GEMM abstraction remains the performance-dominant one.

2. Tensor-network contraction and path-wide lifetime ordering

In large-scale tensor-network contraction, GEMM-oriented mode reordering is introduced as a deterministic backward pass over a fixed binary contraction path. The final output’s mode order is fixed first. Then, for each contraction $\{d_x\}$4 with partition $\{d_x\}$5, the output $\{d_x\}$6’s mode order is already fixed by its consumer; the retained modes of $\{d_x\}$7 are ordered exactly as they appear in $\{d_x\}$8 and reduced modes $\{d_x\}$9 are appended, giving $m = \prod_{x\in M} d_x$0 layout $m = \prod_{x\in M} d_x$1; the retained modes of $m = \prod_{x\in M} d_x$2 are ordered as in $m = \prod_{x\in M} d_x$3 and reduced modes $m = \prod_{x\in M} d_x$4 are appended, giving $m = \prod_{x\in M} d_x$5 layout $m = \prod_{x\in M} d_x$6; and the same permutation is propagated back to the producers of $m = \prod_{x\in M} d_x$7 and $m = \prod_{x\in M} d_x$8 if they are intermediates (Pan et al., 1 Jun 2026). After the pass, every tensor has modes sorted by remaining lifetime, with longest-lived modes leftmost and shortest-lived modes rightmost (Pan et al., 1 Jun 2026).

This path-wide ordering has several stated objectives. It produces $m = \prod_{x\in M} d_x$9 and $n = \prod_{x\in N} d_x$0 layouts so every contraction maps to a GEMM or strided batched GEMM; it ensures $n = \prod_{x\in N} d_x$1 is sufficiently large to keep GEMM kernels in their performance sweet spot; it balances $n = \prod_{x\in N} d_x$2 and $n = \prod_{x\in N} d_x$3 where possible; it minimizes non-coalesced strides and avoids runtime transposes; and it respects cuBLAS and cuTENSOR preferences for contiguous leading dimensions and regular strides while limiting temporary workspace (Pan et al., 1 Jun 2026). The scheme is deterministic and linear-time in the total number of modes across the path: each intermediate is visited at most once, and each permutation is applied once (Pan et al., 1 Jun 2026).

The implementation consequence is that TTGT is avoided at runtime because the offline backward pass produces globally consistent mode orders. The executor calls cuTENSORMp’s contraction kernels, which map to GEMM-like microkernels, without emitting auxiliary transpose kernels. All tensors are in row-major order, and GEMM-oriented reordering ensures operands are laid out as $n = \prod_{x\in N} d_x$4, yielding contiguous leading dimensions for GEMM, with row-major leading dimensions given as $n = \prod_{x\in N} d_x$5, $n = \prod_{x\in N} d_x$6, and $n = \prod_{x\in N} d_x$7 (Pan et al., 1 Jun 2026).

A worked example contracts two rank-4 tensors $n = \prod_{x\in N} d_x$8 and $n = \prod_{x\in N} d_x$9 over $k = \prod_{x\in K} d_x$0, with $k = \prod_{x\in K} d_x$1, $k = \prod_{x\in K} d_x$2, $k = \prod_{x\in K} d_x$3, $k = \prod_{x\in K} d_x$4, $k = \prod_{x\in K} d_x$5, and $k = \prod_{x\in K} d_x$6. The mode partition is $k = \prod_{x\in K} d_x$7, $k = \prod_{x\in K} d_x$8, and $k = \prod_{x\in K} d_x$9, and the reordered layouts are $A' = [M \mid K]$0, $A' = [M \mid K]$1, and $A' = [M \mid K]$2. This yields $A' = [M \mid K]$3, $A' = [M \mid K]$4, $A' = [M \mid K]$5, and the matrix product $A' = [M \mid K]$6. The reported approximate counts are $A' = [M \mid K]$7, $A' = [M \mid K]$8, $A' = [M \mid K]$9, $B' = [K \mid N]$0 flop/byte, and $B' = [K \mid N]$1 flop/byte; on H100 with FP32 peak $B' = [K \mid N]$2 TFLOP/s and HBM bandwidth $B' = [K \mid N]$3 TB/s, the GEMM is FLOP-limited (Pan et al., 1 Jun 2026).

The reported impact is tied directly to this reordering and its co-design with distribution. On a single DGX H100 node with $B' = [K \mid N]$4 GPUs and NVLink, distribution delivers $B' = [K \mid N]$5–$B' = [K \mid N]$6 extra speedup beyond embarrassingly parallel slicing and captures $B' = [K \mid N]$7–$B' = [K \mid N]$8 of the available compute reduction, sustaining $B' = [K \mid N]$9–$N$00 TFLOP/s per GPU (Pan et al., 1 Jun 2026). On up to $N$01 H100 GPUs over InfiniBand, extra speedups beyond the $N$02 slicing baseline range from $N$03 to $N$04 (Pan et al., 1 Jun 2026). The paper states that GEMM-oriented mode reordering contributes by eliminating runtime transposes, ensuring GEMM-friendly shapes, and stabilizing leading distributed modes, thereby converting compute reduction into realized throughput (Pan et al., 1 Jun 2026).

3. GEMM-plus-epilogue reordering in transformer systems

In transformer training systems, GEMM-oriented mode reordering is formulated not as a permutation of contraction indices but as an algebraic reparameterization of operator sequences and tensor modes. CODA keeps the tiled GEMM mainloop unchanged and exposes a constrained, composable epilogue interface with elementwise or pairwise maps, vector loads and stores, tile loads and stores, tile reductions, and stateful transforms (Guo et al., 19 May 2026). The practical objective is to execute normalization, activations, residual and bias updates, dropout, and reductions while the GEMM output tile is still resident in registers or shared memory, so that large intermediates are never materialized as standalone tensors (Guo et al., 19 May 2026).

The reordering has several specific forms. It fuses residual and bias add, GELU or SwiGLU, RoPE, dropout mask application, type casting, and row- or column-wise reductions into GEMM epilogues; it reorganizes tensor modes, dimensions, and strides to align with GEMM tile residency; it collapses multi-dimensional tensors into GEMM-friendly $N$05D shapes by flattening batch and sequence into $N$06 and keeping hidden $N$07 as $N$08; and it delays or commutes row-wise normalization scales so that LayerNorm or RMSNorm scale can be applied in the following GEMM’s epilogue (Guo et al., 19 May 2026). The stated effect is to eliminate explicit transpose and permutation kernels and to place memory-bound work in the shadow of matmul compute (Guo et al., 19 May 2026).

The mapped operators are described explicitly. In forward propagation, residual and bias add are expressed as $N$09; GELU is applied per element in registers; SwiGLU interleaves gate and value along $N$10, splits $N$11 into $N$12, and computes $N$13; RoPE operates on adjacent feature pairs; dropout applies $N$14 in registers; and LayerNorm computes tile-local partial reductions for $N$15 and $N$16, with a lightweight auxiliary kernel completing the row-wise reduction across tiles (Guo et al., 19 May 2026). In backward propagation, the paper states that elementwise epilogues preserve GEMM–epilogue structure, and local multiplication by $N$17 is fused into the GEMM that produces $N$18 (Guo et al., 19 May 2026).

The memory-traffic argument is concrete. In a naïve sequence such as GEMM $N$19 write $N$20 LN $N$21 write $N$22 activation $N$23 write $N$24 dropout $N$25 write $N$26 next GEMM, each memory-bound operator reads and writes an activation-sized tensor. Fusing bias, residual, activation, and dropout into the GEMM epilogue eliminates at least one full read and one full write per fused operator. For LN and RMSNorm, CODA replaces an activation-sized kernel with a few scalars per tile plus a tiny reduction (Guo et al., 19 May 2026). The example given is $N$27 in BF16, where a single extra read or write per op adds approximately $N$28 MB of traffic, and eliminating two such passes can save more than $N$29 MB per block (Guo et al., 19 May 2026).

The paper reports kernel-level speedups relative to cuBLAS with torch.compile when fusing RoPE, SwiGLU, and cross-entropy epilogues, and block-level speedups on hidden sizes $N$30, covering LLaMA-style scales, with both human-written and LLM-authored kernels achieving high utilization (Guo et al., 19 May 2026). It also states that CODA’s numerics study indicates that deferring row-wise normalization scale to the next GEMM epilogue can reduce error compared to the standard path, contingent on a high-accuracy mainloop (Guo et al., 19 May 2026). A plausible implication is that, in this setting, “mode reordering” denotes a tile-residency-preserving reorganization of computation rather than only a permutation of tensor indices.

4. GEMM-centric reordering for pruning and reduced dense inference

In LLM pruning, GEMM-oriented mode reordering is defined as reorganizing the computational modes of an LLM—tokens, channels or neurons, heads, and layers—so that the remaining work maps cleanly to the logical $N$31, $N$32, and $N$33 dimensions of GEMM (Hu et al., 8 Jun 2026). The key claim is that by compacting pruned modes and aligning the shapes that kernels actually see, reordering converts sparsity into reduced dense GEMMs that are well-tiled and fast on the target hardware (Hu et al., 8 Jun 2026).

The taxonomy classifies pruning families by the GEMM axis they reduce. Token or sequence pruning is an $N$34-axis reduction and removes rows of $N$35 across operators. Width or hidden-dimension pruning reduces $N$36 in the current GEMM when output channels are pruned, and this propagates as $N$37 reduction in the next GEMM; pruning input channels reduces $N$38 in the current GEMM and shrinks the reduction loop (Hu et al., 8 Jun 2026). Head pruning is an $N$39-axis reduction for $N$40 and $N$41, and by $N$42 propagation it becomes $N$43 pruning for the next operator (Hu et al., 8 Jun 2026). Depth pruning instead reduces GEMM invocation count while leaving the dimensions of each GEMM unchanged (Hu et al., 8 Jun 2026).

The reordering mechanisms are correspondingly axis-specific. For $N$44 compaction, active tokens are reordered so that kept rows occupy contiguous blocks in $N$45, only tiles containing active tokens are executed, and tiles with all pruned rows are skipped; the paper states that this avoids per-row gather and scatter inside the GEMM and preserves coalesced memory access and kernel tiling (Hu et al., 8 Jun 2026). For $N$46 packing, retained channels are packed contiguously along the pruned axis in both weights and activations, and packed sizes are aligned to hardware tiling constraints such as multiples of $N$47 for tensor cores (Hu et al., 8 Jun 2026). For head packing, retained heads are made contiguous, with consistent packing across $N$48 and $N$49, and $N$50 blocks are packed contiguously for each head (Hu et al., 8 Jun 2026). Depth scheduling replaces pruned layers with identity mappings for static depth, or gates execution per token or per layer for dynamic depth (Hu et al., 8 Jun 2026).

The theoretical bound is written with $N$51 and pruned sizes $N$52, $N$53, $N$54, giving a theoretical compute speedup bound $N$55 (Hu et al., 8 Jun 2026). The same paper emphasizes that realized speedup depends on arithmetic intensity, alignment, and non-GEMM overheads. It reports that static depth pruning at $N$56 sparsity reaches $N$57 prefill and $N$58 decode speedup; static NK at $N$59 sparsity reaches $N$60 prefill and $N$61 decode; static K low-rank reaches approximately $N$62 prefill and approximately $N$63 decode; static K semi-structured achieves approximately $N$64–$N$65; and dynamic M at $N$66 sparsity reaches approximately $N$67 prefill but only approximately $N$68 decode (Hu et al., 8 Jun 2026). It also states that misalignment penalties are severe, with up to $N$69 speed loss from misalignment and recovery when dimensions are aligned to $N$70 (Hu et al., 8 Jun 2026).

The broader empirical conclusion is that static depth pruning is the strongest Pareto-optimal baseline and stays closest to its theoretical acceleration upper bound in memory-bounded scenarios, while the prefill frontier transitions from static depth at low quality loss, to dynamic depth at moderate loss, and finally to static width pruning at higher loss levels (Hu et al., 8 Jun 2026). The paper further reports that dynamic M overheads rise by $N$71 in prefill and $N$72 in decode, that static NK cross-layer decode non-GEMM overhead is $N$73, and that low-rank K decode non-GEMM overhead is $N$74 (Hu et al., 8 Jun 2026). This suggests that GEMM-oriented mode reordering is most effective when it preserves large, contiguous dense kernels and minimizes gather, scatter, routing, and shape fragmentation.

5. Reordering for locality, communication avoidance, and accelerator mapping

A distinct line of work applies GEMM-oriented reordering directly to the traversal and partitioning of GEMM itself. Generalized space filling curves reorder the $N$75D $N$76 output modes of GEMM so that tiles of $N$77 are visited in generalized Hilbert order rather than row-major or column-major order (Georganas et al., 22 Jan 2026). For blocked dimensions $N$78 with $N$79 and $N$80, a generalized Hilbert map is built over the $N$81 grid of $N$82 tiles, and the outer loops traverse the $N$83 tiles in this SFC order while performing the $N$84 accumulation via BRGEMM over $N$85 panels (Georganas et al., 22 Jan 2026). The stated effect is that adjacent SFC indices correspond to neighboring $N$86 tiles, enhancing spatial and temporal reuse of $N$87 and $N$88 panels across successive tiles and reducing misses and traffic (Georganas et al., 22 Jan 2026).

This traversal is then combined with $N$89D communication-avoiding replication by splitting the $N$90 dimension into $N$91 layers, so that each layer owns a copy of $N$92 and processes a $N$93 fraction of the outer products, followed by a final reduction over the $N$94 copies (Georganas et al., 22 Jan 2026). In the square case with $N$95 cores, the paper states that the SFC-CA algorithm moves per core

$N$96

words on the critical path before the final $N$97-reduction, matching the $N$98D lower bound up to constants (Georganas et al., 22 Jan 2026). Empirically, on Intel Emerald Rapids, Intel Granite Rapids, AMD Zen5, and AWS Graviton4, the reported geometric-mean speedup over vendor libraries is $N$99, $K$00, $K$01, and $K$02 respectively, and for selected shapes SFC-CA reduces total L2 misses by $K$03 or $K$04 while increasing TFLOPs correspondingly (Georganas et al., 22 Jan 2026).

A second hardware-oriented formulation appears in GOMA, which maps GEMM axes as $K$05, $K$06, and $K$07, and treats mode reordering as the choice of walking axis at multiple hierarchy levels together with per-axis bypass and spatial unrolling (Yang et al., 9 Mar 2026). The walking axes $K$08 and $K$09 determine which dimension advances temporally at the DRAM↔SRAM and SRAM↔array scales, while bypass variables decide whether data with normal $K$10 resides at SRAM or regfile (Yang et al., 9 Mar 2026). Because $K$11 corresponds to $K$12, $K$13 to $K$14, and $K$15 to $K$16, the choice of walking axis determines which matrix projection remains unchanged and can be reused at a given level (Yang et al., 9 Mar 2026).

GOMA provides closed-form traffic counts and an analytical $K$17 energy objective, then formulates mapping selection as an integer optimization problem under capacity, divisibility, and PE-utilization constraints (Yang et al., 9 Mar 2026). The paper states that the closed-form energy matches timeloop-model with approximately $K$18 consistency across $K$19 mappings, with mean relative error $K$20 and energy-weighted overall relative error $K$21, and that Gurobi solves the mixed integer problem to zero optimality gap (Yang et al., 9 Mar 2026). Across four accelerator templates and $K$22 prefill workloads, the reported improvement in energy–delay product is $K$23–$K$24 over state-of-the-art mappers, with $K$25–$K$26 faster time-to-solution (Yang et al., 9 Mar 2026).

These communication-oriented and mapping-oriented formulations connect back to tensor-network contraction. In the multi-GPU tensor-network framework, GEMM-oriented reordering makes the longest-lived modes the leading dimensions, and the distribution planner chooses a minimum leading prefix whose product of extents covers $K$27 devices; because leading dimensions are contiguous in row-major layout, each device receives a single contiguous shard, and longest-lived modes being distributed reduces forced redistributions later (Pan et al., 1 Jun 2026). The planner therefore prioritizes distributing leading free modes over $K$28, since splitting $K$29 requires summing partial results across devices, which is generally more communication-heavy than splitting free modes (Pan et al., 1 Jun 2026). In all three cases—SFC traversal, GOMA mapping, and tensor-network distribution—the reordering target is communication as much as compute.

6. Related methods, limitations, and recurring trade-offs

The literature situates GEMM-oriented mode reordering against several established alternatives. In tensor networks, slicing-based parallelism exposes embarrassingly parallel tasks but repeats computation across slices and scales exponentially with sliced bonds; TTGT-style methods and libraries such as cuTENSOR and TBLIS insert transposes locally per contraction to reach GEMM shapes; Cyclops Tensor Framework and exaTN support distributed tensor operations but typically rely on per-operation planning or slicing-like strategies for tensor networks (Pan et al., 1 Jun 2026). In tensor contraction more broadly, GETT contrasts with TTGT and Loops-over-GEMM by embedding permutations inside packing, thereby avoiding global read and write traffic for transposes and avoiding many small GEMMs or strided submatrices (Springer et al., 2016). In transformer systems, framework graphs typically materialize boundaries between operators, and compiler fusion helps but is limited by global schedules and evolving hardware-specific details; CODA instead keeps the expert GEMM mainloop intact and expresses surrounding operators as tile-local epilogues (Guo et al., 19 May 2026).

The principal limitations are also recurrent across domains. In tensor-network contraction, very small $K$30 produces low arithmetic intensity and underutilized GEMM kernels; if path structure forces frequent changes in retained versus reduced modes, reordering cannot eliminate all redistributions; pathological stride patterns can still produce skinny or tall matrices less favorable to tiling; and very small batches remain latency-sensitive even with strided batched GEMM (Pan et al., 1 Jun 2026). In CODA, attention softmax and global reductions over entire sequences are not covered by the tile-local epilogue interface, very irregular sparsity or exotic shapes may not align with tile residency, and distributed multi-GPU coordination is out of scope (Guo et al., 19 May 2026). In pruning, dynamic pruning overheads, misalignment, metadata, extra launches, hardware variability, and small-$K$31 decode regimes can prevent approaching $K$32 (Hu et al., 8 Jun 2026). In SFC-CA GEMM, very small matrices can make mapping overhead and kernel launch costs dominant, and extreme aspect ratios can limit the benefit of replication (Georganas et al., 22 Jan 2026).

Several mitigation strategies are stated explicitly. The tensor-network framework uses lifetime ordering to make leading modes longest-lived and contiguous, reducing redistribution frequency and keeping transfers bandwidth-bound, and the distribution planner concentrates redistributions at tensor “valleys” and avoids latency-bound micro-transfers (Pan et al., 1 Jun 2026). CODA uses deterministic auxiliary reductions rather than atomics, FP32 accumulation for numerical stability, Welford updates for variance estimation, and careful management of masks, barriers, and async pathways (Guo et al., 19 May 2026). The pruning study recommends packing kept channels or heads contiguously, aligning pruned dimensions to multiples of $K$33, sharing masks across coupled GEMMs, fusing elementary ops, caching descriptors, and using CUDA graphs (Hu et al., 8 Jun 2026). SFC-CA recommends choosing $K$34 and a small set of $K$35 values, and GOMA recommends aligning walking axes with the matrix one wishes to keep stationary at each stage, often favoring output-stationary accumulation along $K$36 when partial sums should be completed close to compute (Georganas et al., 22 Jan 2026, Yang et al., 9 Mar 2026).

Taken together, these results define GEMM-oriented mode reordering as a general systems principle rather than a single algorithm. It includes path-wide lifetime ordering for tensor contractions, GEMM-plus-epilogue restructuring for transformer blocks, compaction and packing under a GEMM-centric pruning taxonomy, locality-preserving traversal for communication-avoiding GEMM, and analytical selection of $K$37 walking axes and residency on spatial accelerators. The shared aim is consistent: transform computation so that the hardware sees contiguous, regular, GEMM-friendly work, while transpose traffic, fragmented kernels, and unnecessary communication are either eliminated or pushed out of the critical path (Pan et al., 1 Jun 2026, Guo et al., 19 May 2026, Hu et al., 8 Jun 2026, Georganas et al., 22 Jan 2026, Springer et al., 2016, Yang et al., 9 Mar 2026).