---
title: GEMM-Oriented Mode Reordering Techniques
url: https://www.emergentmind.com/topics/gemm-oriented-mode-reordering
type: topic
---

# GEMM-Oriented Mode Reordering Techniques

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 [2606.01852]. 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 [2605.19269]. 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 [2606.09080]. Related formulations also reorder GEMM traversal and mapping itself, including generalized space filling curves for communication-avoiding matrix multiplication [2601.16294], GEMM-like tensor contractions without global transposes [1607.00145], and analytical optimization of $M/N/K$ walking axes and residency on spatial accelerators [2603.07962].

## 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 $A_{M\times K}$, $B_{K\times N}$, and $C_{M\times N}$, with the contraction becoming the matrix product $C_{M\times N} = A_{M\times K} \times B_{K\times N}$ [2606.01852]. GETT expresses the same principle by selecting three index groups, $I_m$, $I_n$, and $I_k$, so that $C_{I_m,I_n} += \sum_{I_k} A_{I_m,I_k} B_{I_k,I_n}$, and then defining permutations $\Pi_A$, $\Pi_B$, and $\Pi_C$ that reorder the modes of $A$, $B$, and $C$ as $[I_m \mid I_k]$, $[I_k \mid I_n]$, and $[I_m \mid I_n]$ respectively [1607.00145].

The same GEMM-centric abstraction appears outside tensor contraction. In transformer layers, the fixed GEMM mainloop stages $A$ and $B$ tiles into shared memory and registers, accumulates into FP32 accumulators using Tensor Cores, and produces register-resident output fragments $rD$ per thread or warp; surrounding operators are then expressed as constrained epilogue primitives attached to that mainloop [2605.19269]. In LLM pruning, the GEMM-centric taxonomy assigns token dimension to $M$, input feature or reduction dimension to $K$, and output feature dimension to $N$, with a general linear projection written as $O = A W^\top$ with $A \in \mathbb{R}^{T\times D}$, $W \in \mathbb{R}^{I\times D}$, and $O \in \mathbb{R}^{T\times I}$, mapping $T \to M$, $D \to K$, and $I \to N$ [2606.09080].

A common analytical rationale is roofline behavior. Arithmetic intensity for one GEMM is $AI = (2MNK)/\text{bytes\_moved}$, and throughput follows $T(M,N,K) = \min(\text{PeakFLOPs}, BW \times AI)$ [2606.01852]. In the pruning taxonomy, the same style of model is written as $T \approx \max(\text{FLOPs}/P_{\text{FLOPs}}, \text{Bytes}/P_{BW})$, with $AI = \text{FLOPs}/\text{Bytes}$, emphasizing that realized speedup depends on whether reduced shapes remain well-tiled and whether workloads remain compute-bound or become memory-bound [2606.09080]. 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 $A \times B \to C$ with partition $(M,N,K)$, the output $C$’s mode order is already fixed by its consumer; the retained modes of $A$ are ordered exactly as they appear in $C$ and reduced modes $K$ are appended, giving $A$ layout $[M \mid K]$; the retained modes of $B$ are ordered as in $C$ and reduced modes $K$ are appended, giving $B$ layout $[K \mid N]$; and the same permutation is propagated back to the producers of $A$ and $B$ if they are intermediates [2606.01852]. After the pass, every tensor has modes sorted by remaining lifetime, with longest-lived modes leftmost and shortest-lived modes rightmost [2606.01852].

This path-wide ordering has several stated objectives. It produces $[M \mid K]$ and $[K \mid N]$ layouts so every contraction maps to a GEMM or strided batched GEMM; it ensures $K$ is sufficiently large to keep GEMM kernels in their performance sweet spot; it balances $M$ and $N$ 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 [2606.01852]. 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 [2606.01852].

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 $[\text{retained} \mid \text{reduced}]$, yielding contiguous leading dimensions for GEMM, with row-major leading dimensions given as $\text{lda}=K$, $\text{ldb}=N$, and $\text{ldc}=N$ [2606.01852].

A worked example contracts two rank-4 tensors $A_{abkl}$ and $B_{klcd}$ over $k,l$, with $d_a=8$, $d_b=16$, $d_c=16$, $d_d=8$, $d_k=64$, and $d_l=32$. The mode partition is $M=\{a,b\}$, $N=\{c,d\}$, and $K=\{k,l\}$, and the reordered layouts are $A'=[ab \mid kl]$, $B'=[kl \mid cd]$, and $C'=[ab \mid cd]$. This yields $m=128$, $n=128$, $k=2048$, and the matrix product $C_{128\times128} = A_{128\times2048} \times B_{2048\times128}$. The reported approximate counts are $\text{flops}_{\text{real}} \approx 67.1\text{M}$, $\text{flops}_{\text{complex}} \approx 268.4\text{M}$, $\text{bytes\_moved} \approx 4.28\text{ MB}$, $AI_{\text{real}} \approx 15.7$ flop/byte, and $AI_{\text{complex}} \approx 62.7$ flop/byte; on H100 with FP32 peak $67$ TFLOP/s and HBM bandwidth $\approx 3.35$ TB/s, the GEMM is FLOP-limited [2606.01852].

The reported impact is tied directly to this reordering and its co-design with distribution. On a single DGX H100 node with $8$ GPUs and NVLink, distribution delivers $7$–$173\times$ extra speedup beyond embarrassingly parallel slicing and captures $87$–$101\%$ of the available compute reduction, sustaining $28$–$33$ TFLOP/s per GPU [2606.01852]. On up to $1024$ H100 GPUs over InfiniBand, extra speedups beyond the $1024\times$ slicing baseline range from $42\times$ to $67{,}869\times$ [2606.01852]. 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 [2606.01852].

## 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 [2605.19269]. 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 [2605.19269].

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 $2$D shapes by flattening batch and sequence into $M = N \times S$ and keeping hidden $H$ as $N$; and it delays or commutes row-wise normalization scales so that LayerNorm or RMSNorm scale can be applied in the following GEMM’s epilogue [2605.19269]. The stated effect is to eliminate explicit transpose and permutation kernels and to place memory-bound work in the shadow of matmul compute [2605.19269].

The mapped operators are described explicitly. In forward propagation, residual and bias add are expressed as $D = \alpha \cdot rD + C + b$; GELU is applied per element in registers; SwiGLU interleaves gate and value along $N$, splits $D$ into $(G,V)$, and computes $y = \text{silu}(G) \odot V$; RoPE operates on adjacent feature pairs; dropout applies $y = (\text{mask} \odot x)/(1-p)$ in registers; and LayerNorm computes tile-local partial reductions for $\mu$ and $\sigma^2$, with a lightweight auxiliary kernel completing the row-wise reduction across tiles [2605.19269]. In backward propagation, the paper states that elementwise epilogues preserve GEMM–epilogue structure, and local multiplication by $f'(\cdot)$ is fused into the GEMM that produces $\nabla D_{\text{out}}$ [2605.19269].

The memory-traffic argument is concrete. In a naïve sequence such as GEMM $\to$ write $\to$ LN $\to$ write $\to$ activation $\to$ write $\to$ dropout $\to$ write $\to$ 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 [2605.19269]. The example given is $M\times N = 4096\times4096$ in BF16, where a single extra read or write per op adds approximately $137$ MB of traffic, and eliminating two such passes can save more than $250$ MB per block [2605.19269].

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 $\{2048,4096,8192\}$, covering LLaMA-style scales, with both human-written and LLM-authored kernels achieving high utilization [2605.19269]. 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 [2605.19269]. 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 $M$, $N$, and $K$ dimensions of GEMM [2606.09080]. 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 [2606.09080].

The taxonomy classifies pruning families by the GEMM axis they reduce. Token or sequence pruning is an $M$-axis reduction and removes rows of $A$ across operators. Width or hidden-dimension pruning reduces $N$ in the current GEMM when output channels are pruned, and this propagates as $K$ reduction in the next GEMM; pruning input channels reduces $K$ in the current GEMM and shrinks the reduction loop [2606.09080]. Head pruning is an $N$-axis reduction for $W_q$ and $W_o$, and by $N \to K$ propagation it becomes $K$ pruning for the next operator [2606.09080]. Depth pruning instead reduces GEMM invocation count while leaving the dimensions of each GEMM unchanged [2606.09080].

The reordering mechanisms are correspondingly axis-specific. For $M$ compaction, active tokens are reordered so that kept rows occupy contiguous blocks in $A$, 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 [2606.09080]. For $N/K$ 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 $16$ for tensor cores [2606.09080]. For head packing, retained heads are made contiguous, with consistent packing across $Q/K/V$ and $O$, and $d_h$ blocks are packed contiguously for each head [2606.09080]. Depth scheduling replaces pruned layers with identity mappings for static depth, or gates execution per token or per layer for dynamic depth [2606.09080].

The theoretical bound is written with $\text{FLOPs} = 2MKN$ and pruned sizes $M'=(1-p_M)M$, $K'=(1-p_K)K$, $N'=(1-p_N)N$, giving a theoretical compute speedup bound
$S_{\text{theory}} = 1/[(1-p_M)(1-p_K)(1-p_N)]$ [2606.09080]. The same paper emphasizes that realized speedup depends on arithmetic intensity, alignment, and non-GEMM overheads. It reports that static depth pruning at $50\%$ sparsity reaches $1.88\times$ prefill and $1.91\times$ decode speedup; static NK at $50\%$ sparsity reaches $1.77\times$ prefill and $1.70\times$ decode; static K low-rank reaches approximately $1.43\times$ prefill and approximately $1.46\times$ decode; static K semi-structured achieves approximately $1.08\times$–$1.14\times$; and dynamic M at $50\%$ sparsity reaches approximately $1.44\times$ prefill but only approximately $1.10\times$ decode [2606.09080]. It also states that misalignment penalties are severe, with up to $35\%$ speed loss from misalignment and recovery when dimensions are aligned to $16$ [2606.09080].

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 [2606.09080]. The paper further reports that dynamic M overheads rise by $42.4\%$ in prefill and $61.5\%$ in decode, that static NK cross-layer decode non-GEMM overhead is $+48.5\%$, and that low-rank K decode non-GEMM overhead is $+30.9\%$ [2606.09080]. 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 $2$D $(i,j)$ output modes of GEMM so that tiles of $C$ are visited in generalized Hilbert order rather than row-major or column-major order [2601.16294]. For blocked dimensions $M,N,K$ with $M_b=M/b_m$ and $N_b=N/b_n$, a generalized Hilbert map is built over the $M_b \times N_b$ grid of $C$ tiles, and the outer loops traverse the $C$ tiles in this SFC order while performing the $K$ accumulation via BRGEMM over $K$ panels [2601.16294]. The stated effect is that adjacent SFC indices correspond to neighboring $(i_m,i_n)$ tiles, enhancing spatial and temporal reuse of $A$ and $B$ panels across successive tiles and reducing misses and traffic [2601.16294].

This traversal is then combined with $2.5$D communication-avoiding replication by splitting the $K$ dimension into $c = K_{\text{layers}}$ layers, so that each layer owns a copy of $C$ and processes a $1/c$ fraction of the outer products, followed by a final reduction over the $c$ copies [2601.16294]. In the square case with $T$ cores, the paper states that the SFC-CA algorithm moves per core
$W_{\text{core}} = \Theta(n^2/\sqrt{Tc})$
words on the critical path before the final $C$-reduction, matching the $2.5$D lower bound up to constants [2601.16294]. Empirically, on Intel Emerald Rapids, Intel Granite Rapids, AMD Zen5, and AWS Graviton4, the reported geometric-mean speedup over vendor libraries is $1.4\times$, $2\times$, $1.4\times$, and $1.4\times$ respectively, and for selected shapes SFC-CA reduces total L2 misses by $1.61\times$ or $2.9\times$ while increasing TFLOPs correspondingly [2601.16294].

A second hardware-oriented formulation appears in GOMA, which maps GEMM axes as $M \to x$, $N \to y$, and $K \to z$, and treats mode reordering as the choice of walking axis at multiple hierarchy levels together with per-axis bypass and spatial unrolling [2603.07962]. The walking axes $\alpha_{0-1} \in \{x,y,z\}$ and $\alpha_{1-2} \in \{x,y,z\}$ determine which dimension advances temporally at the DRAM↔SRAM and SRAM↔array scales, while bypass variables decide whether data with normal $d$ resides at SRAM or regfile [2603.07962]. Because $d=y$ corresponds to $A(x,z)$, $d=x$ to $B(y,z)$, and $d=z$ to $P(x,y)$, the choice of walking axis determines which matrix projection remains unchanged and can be reused at a given level [2603.07962].

GOMA provides closed-form traffic counts and an analytical $O(1)$ energy objective, then formulates mapping selection as an integer optimization problem under capacity, divisibility, and PE-utilization constraints [2603.07962]. The paper states that the closed-form energy matches timeloop-model with approximately $99.9\%$ consistency across $8064$ mappings, with mean relative error $0.099\%$ and energy-weighted overall relative error $0.066\%$, and that Gurobi solves the mixed integer problem to zero optimality gap [2603.07962]. Across four accelerator templates and $12$ prefill workloads, the reported improvement in energy–delay product is $2.24$–$4.24\times$ over state-of-the-art mappers, with $3.83$–$73.6\times$ faster time-to-solution [2603.07962].

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 $P$ 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 [2606.01852]. The planner therefore prioritizes distributing leading free modes over $K$, since splitting $K$ requires summing partial results across devices, which is generally more communication-heavy than splitting free modes [2606.01852]. 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 [2606.01852]. 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 [1607.00145]. 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 [2605.19269].

The principal limitations are also recurrent across domains. In tensor-network contraction, very small $K$ 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 [2606.01852]. 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 [2605.19269]. In pruning, dynamic pruning overheads, misalignment, metadata, extra launches, hardware variability, and small-$M$ decode regimes can prevent approaching $S_{\text{theory}}$ [2606.09080]. 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 [2601.16294].

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 [2606.01852]. 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 [2605.19269]. The pruning study recommends packing kept channels or heads contiguously, aligning pruned dimensions to multiples of $16$, sharing masks across coupled GEMMs, fusing elementary ops, caching descriptors, and using CUDA graphs [2606.09080]. SFC-CA recommends choosing $c \in \{1,2,4\}$ and a small set of $k\_\text{block\_factor}$ 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$ when partial sums should be completed close to compute [2601.16294; 2603.07962].

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 $M/N/K$ 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 [2606.01852; 2605.19269; 2606.09080; 2601.16294; 1607.00145; 2603.07962].

Source: https://www.emergentmind.com/topics/gemm-oriented-mode-reordering