RankMixer: Hardware Efficiency Optimization

📝 Before You Continue: You have read 7.3 MTGR (feature compatibility) and 3.2 Feature Crossing (cross-feature expressiveness). This chapter switches perspective — looking at scaling not from "can it be modeled" but from "is it worth what the GPU pays", with hardware-aware architecture design at its core.

A traditional DLRM's MFU (Model FLOPs Utilization) on GPU is typically only 4–5%, while large language models reach 40–60%. This ten-fold efficiency gap directly prevents recommendation models from enjoying the Scaling Law's dividends — even as parameter counts grow, most of the added compute is wasted on inefficient memory access.

The root cause is that the traditional DLRM inherited its architecture from the CPU era, exposing three fundamental problems on GPU: (1) the core operations are predominantly memory-bound — embedding lookups, feature crossing, and sequence modeling move far more memory than they compute; (2) the computation graph is highly fragmented — many independent hand-crafted modules chained together, with kernel launch overhead and global memory transfers accumulating into a bottleneck; (3) Tensor Cores cannot be fully utilized — most operations are small vector ops or irregular memory accesses that cannot leverage the matrix-multiply acceleration units.

RankMixer solves this at the root through hardware-aware architecture design. The core principle: derive the architecture from hardware characteristics, restructuring the recommendation model as a unified, GPU-friendly computation graph — Token Mixing replaces Self-Attention to cut complexity, Per-Token FFN captures feature heterogeneity, and Sparse MoE enables parameter-efficient scaling.

7.4.0 The RankMixer Architecture

The model's core is $L$ stacked RankMixer Blocks, each containing Multi-head Token Mixing (replacing Self-Attention) and a Per-Token FFN (capturing feature heterogeneity). Input features are tokenized into $T$ tokens of uniform dimension, pass through $L$ blocks, and produce an output via mean pooling.

After input features are tokenized, they pass through multiple RankMixer Blocks (each = Token Mixing + Per-Token FFN), and mean pooling finally produces the logit. All core operations are matrix multiplications.

Each RankMixer Block's forward pass:

$$\begin{array}{rl}{\mathbf{S}}_{n-1}& =\text{LN}(\text{TokenMixing}({\mathbf{X}}_{n-1})+{\mathbf{X}}_{n-1})\\ {\mathbf{X}}_n& =\text{LN}(\text{PFFN}({\mathbf{S}}_{n-1})+{\mathbf{S}}_{n-1})\end{array}$$

Overall complexity is $O(LT{D}^2)$. In the Sparse MoE version, the Per-Token FFN can be replaced with expert networks, expanding parameters while keeping inference cost. The design follows three principles: (1) all core operations are matrix multiplications, fully exploiting Tensor Cores; (2) the computation graph stays as simple as possible to reduce kernel launch overhead; (3) the expressiveness required by recommendation tasks is preserved.

7.4.1 The Token Mixing Mechanism

Self-Attention complexity is $O(T^{2}D)$ (from computing the full token-pair similarity matrix $Q{K}^T$). In recommendation, feature counts can reach hundreds or thousands, making the $T^2$ term a significant bottleneck. RankMixer's core insight: what recommendation tasks need is information mixing between tokens, not similarity-based dynamic weighting (attention). For example, learning high-order interactions like "young users in tier-1 cities prefer tech items" is essentially fusing information from multiple tokens into new representations — it does not necessarily require explicitly computing token-pair similarities.

Token Mixing's core idea: mix along the feature dimension rather than the token dimension. Given input $\mathbf{X}\in {\mathbb{R}}^{T\times D}$, there are two steps:

Step 1, Multi-head decomposition — each token is decomposed into $H$ heads: $[{\mathbf{x}}_t^{(1)}\Vert \cdots \Vert {\mathbf{x}}_t^{(H)}]=\text{SplitHead}({\mathbf{x}}_t)$, where ${\mathbf{x}}_t^{(h)}\in {\mathbb{R}}^{D/H}$ is head $h$ of token $t$.

Step 2, Token-wise mixing — within each head, concatenate that head's portion across all tokens: ${\mathbf{s}}^{(h)}=\text{Concat}({\mathbf{x}}_1^{(h)},\dots ,{\mathbf{x}}_T^{(h)})\in {\mathbb{R}}^{TD/H}$.

The key: change the data layout from "per token" to "per head". The original $\mathbf{X}$ is $T$ vectors of length $D$; after SplitHead+Concat it becomes $H$ vectors of length $TD/H$, with different tokens' features densely packed within each head, creating the conditions for mixing. In practice $H=T$, so each "head" holds a slice of every token's features. After mixing the token count is unchanged, enabling residual connections.

Left: Self-Attention computes the full $T\times T$ similarity matrix ($O(T^{2}D)$); right: Token Mixing rearranges by head and applies a linear transformation along the feature dimension ($O(TD)$), with no softmax.

Complexity-wise, Token Mixing costs $O(TD)$ (mainly memory rearrangement). Compared with self-attention's $O(T^{2}D+T{D}^2)$, this avoids the $T^2$ term when $T$ is large (hundreds to thousands in recommendation) — a significant reduction. And Token Mixing has no softmax normalization (which requires an extra reduction kernel), further cutting kernel overhead.

The key question: without explicitly computing token-pair similarities, how are interactions still captured? The answer is stacking multiple layers. A single Token Mixing layer is "feature-level mixing" — the same feature dimension across tokens influences each other (because they are concatenated into one vector). With multiple stacked layers, layer 1's per-token output fuses first-order information from all tokens; layer 2's input is already a mixed result where each token contains other tokens' information, so layer 2's mixing achieves second-order interactions:

$${\mathbf{X}}_1=f({\mathbf{X}}_0),{\mathbf{X}}_2=f(f({\mathbf{X}}_0)),\dots ,{\mathbf{X}}_L=\underset{L\text{ times}}{\underbrace{f\circ \cdots \circ f}}({\mathbf{X}}_0)$$

Each Token Mixing layer applies a linear transformation along the feature dimension (followed by FFNs), so stacking $L$ layers can model $L$-th order polynomial interactions between tokens. In practice $L$ is usually 6–12 layers, enough to capture the needed high-order crossings.

From the hardware's perspective, Token Mixing's core is data rearrangement (SplitHead, Concat) that can be implemented with efficient kernels, designed as contiguous memory reads/writes (coalesced memory access) to fully use bandwidth. Compared with attention's softmax (a global normalization), Token Mixing is local and parallel. More importantly, SplitHead, Concat, and the FFN can be fused into a single kernel, reducing kernel launch overhead — the key to the MFU improvement.

Analysis: Token Mixing replaces "token-pair similarity + softmax" with "feature-dimension mixing + multi-layer stacking", dropping $O(T^{2}D)$ to $O(T{D}^2)$ overall, with all operations fusable into matrix-multiply kernels. What is sacrificed is attention's "dynamic similarity routing"; what is gained is GPU utilization — for recommendation, where features are numerous and interaction patterns fairly fixed, this trade is well worth it.

7.4.2 Per-Token FFN

The standard Transformer's FFN uses the same weights for all tokens: $\text{FFN}(\mathbf{x})={\mathbf{W}}_2\cdot \text{GELU}({\mathbf{W}}_1\mathbf{x}+{\mathbf{b}}_1)+{\mathbf{b}}_2$. This is reasonable for LLMs (all tokens share one semantic space). But recommendation features live in entirely different semantic spaces: user IDs imply implicit preferences, item categories are coarse-grained classes, click rates follow long-tail distributions, timestamps have periodicity. Forcing the same FFN on them loses parameter efficiency.

RankMixer's core design: each token gets its own FFN parameters. For the $t$-th token:

$${\mathbf{v}}_t={\mathbf{W}}_{\text{pffn}}^{t,2}\cdot \text{GELU}({\mathbf{W}}_{\text{pffn}}^{t,1}{\mathbf{s}}_t+{\mathbf{b}}_{\text{pffn}}^{t,1})+{\mathbf{b}}_{\text{pffn}}^{t,2}$$

Each token's ${\mathbf{W}}_{\text{pffn}}^{t,i}$ is independent, so: (1) each token learns a transformation specific to its semantic space; (2) high-information tokens are automatically allocated more parameter capacity; (3) different semantic spaces no longer interfere.

Per-Token FFN is fundamentally different from MMoE. In MMoE, multiple experts share the same input, and gating dynamically weights a combination of expert outputs: $\mathbf{y}={\sum }_i{g}_i(\mathbf{x})\cdot {\text{Expert}}_i(\mathbf{x})$ (all experts see the same $\mathbf{x}$). Per-Token FFN gives each token its own input and its own FFN: ${\mathbf{v}}_t={\text{FFN}}_t({\mathbf{s}}_t)$ (each FFN sees a different ${\mathbf{s}}_t$). Parameter isolation ensures learning in different feature spaces stays independent, preventing high-frequency features from dominating low-frequency ones.

For parameter efficiency, with $T$ tokens, Per-Token FFN has total parameters $\text{Param}=2Tk{D}^2$ — $T$ times the shared FFN's ($2k{D}^2$). But computational complexity is unchanged: $\text{FLOPs}=2Tk{D}^2$, identical to the shared FFN (which must also compute once per token for $T$ tokens). The added parameters are "specialized" — each block serves only one token, giving higher learning efficiency.

Cross-feature-space interaction happens through the Token Mixing layers. Per-Token FFN focuses on deep modeling within each space; Token Mixing ensures information flows between tokens. This "mixing + per-token processing" combination preserves parameter isolation while achieving thorough cross-space interaction through multi-layer stacking.

7.4.3 Sparse MoE Scaling

With Token Mixing and Per-Token FFN in place, how do we scale to billions or even tens of billions of parameters? Directly adding depth/width scales compute linearly, and inference latency grows proportionally — industrially unacceptable. Sparse MoE (Sparse Mixture of Experts) provides the solution: not all parameters participate in every sample's computation; a subset of experts is dynamically selected per sample. The model can have a huge parameter count while per-sample compute stays fixed (activating only a few experts).

The ideal MoE pattern has each expert specialize in some sample pattern. But achieving effective expert specialization in recommendation faces three challenges: (1) the input is high-dimensional sparse features whose combinatorial space is exponential; representations are scattered across the high-dimensional space without clear cluster structure, making stable routing hard for the gating to learn; (2) data is extremely imbalanced — head users account for 50% of samples; if gating routes many head samples to one expert early on, that expert receives more gradients, gating keeps sending it more, and a few experts end up processing most samples (expert overload) while the rest go nearly unused (expert underutilization); (3) even if training is load-balanced, request distributions at inference may differ, making some experts bottlenecks and increasing latency variance.

RankMixer counters with two complementary training strategies. First, ReLU Routing — standard MoE uses Top-$k$ + Softmax routing, activating a fixed $k$ experts per token. RankMixer's ReLU Routing lets each token activate a variable number of experts:

$$G_{i,j}=\text{ReLU}(h({\mathbf{s}}_i)),{\mathbf{v}}_i=\sum _{j=1}^{N_e}{G}_{i,j}\cdot e_{i,j}({\mathbf{s}}_i)$$

ReLU outputs can be 0 (not activated) or positive (activated), so high-information tokens may activate more experts. To control sparsity, a regularization is added: $\mathcal{L}={\mathcal{L}}_{\text{task}}+\lambda {\mathcal{L}}_{\text{reg}}$, where ${\mathcal{L}}_{\text{reg}}={\sum }_{i,j}{G}_{i,j}$ and $\lambda$ controls the average number of activated experts.

Second, Dense-Training / Sparse-Inference (DTSI-MoE) — Per-Token FFN already multiplies parameters by $T$; adding MoE on top expands expert count further, easily causing expert under-training. DTSI-MoE uses two routers: during training, a dense router $h_{\text{train}}$ activates all or most experts to ensure sufficient training; at inference, a sparse router $h_{\text{infer}}$ activates only a few experts to cut compute. Both routers train simultaneously; only $h_{\text{infer}}$ is constrained by ${\mathcal{L}}_{\text{reg}}$:

$$G_{i,j}^{\text{train}}=\text{ReLU}(h_{\text{train}}({\mathbf{s}}_i)),G_{i,j}^{\text{infer}}=\text{ReLU}(h_{\text{infer}}({\mathbf{s}}_i))$$

During training, the forward pass uses $G^{\text{train}}$ while $G^{\text{infer}}$ is computed alongside and given the sparsity regularization; at inference, only $h_{\text{infer}}$ is used. This achieves sufficient training, efficient inference, and consistent strategy.

Load balancing is achieved through the soft constraint of ${\mathcal{L}}_{\text{reg}}$. Expert $j$'s total activation in a batch is $A_j={\sum }_i{G}_{i,j}$, and the regularization can be rewritten as ${\mathcal{L}}_{\text{reg}}={\sum }_j{A}_j$. When some expert's $A_j$ grows too large, the gradient $\partial {\mathcal{L}}_{\text{reg}}/\partial h_{\text{infer}}$ suppresses its activation probability, achieving load balance. Compared with hard constraints (capacity limits), a soft constraint never force-assigns suboptimal experts when one is saturated, preserving routing flexibility.

Left: the traditional DLRM's fragmented computation graph (memory-bound embedding lookups, small kernels, launch overhead), with effective GEMM at only 5%; right: RankMixer's core operations are all GEMMs — Token Mixing + PFFN take ~85% of compute, and MFU reaches 45%.

The key to RankMixer's high MFU: all core operations are compute-bound large matrix multiplications. Token Mixing and PFFN take about 85% of compute time, all GEMMs that efficiently use Tensor Cores (a single GEMM kernel reaches 60–80% MFU). By contrast, in the traditional DLRM, embedding lookups (40% of time, memory-bound), small kernels (35% of time, MFU<10%), and launch overhead (20% of time) dominate, with effective GEMM at only 5% — which is exactly why a DLRM's MFU is 4–5% while RankMixer reaches 45%.

💡 Key Insight: RankMixer moves the recommendation model from fragmented design to a unified architectural paradigm. Algorithmically, Token Mixing drops complexity from $O(T^{2}D)$ to $O(T{D}^2)$, Per-Token FFN captures feature heterogeneity, and Sparse MoE achieves parameter-efficient scaling through ReLU Routing and DTSI-MoE. Systematically, unifying all core operations as matrix multiplications raises MFU from 4–5% to 45%, making the recommendation model a "first-class citizen" on GPU that can directly use Tensor Cores and the mature LLM toolchain, opening a path to sustained scaling. But RankMixer focuses on compute efficiency inside the model; the pipeline still has other fragmentation — sequence modeling separate from feature interaction, retrieval separate from ranking, multi-task fragmentation. The next section, OneTrans, breaks through these remaining walls.

⚠️ Common Mistakes in 7.4

Chapter Summary

📌 Key Takeaways

❓ FAQ

Q1: Without computing similarities, can Token Mixing really replace Self-Attention?

A: Yes. High-order crossings in recommendation are essentially "fusing multiple tokens' information into new representations"; a single layer mixes along the feature dimension, and stacking layers achieves $L$-th order polynomial interactions. The cost is losing dynamic similarity routing, but with many features and relatively fixed patterns in recommendation, the gained GPU utilization is the better deal.

Q2: Per-Token FFN multiplies parameters by T — why don't FLOPs change?

A: A shared FFN also computes once per token for T tokens (one FFN per token), so FLOPs were already $2Tk{D}^2$; Per-Token merely swaps the shared weights for T independent sets, still computing once per token — same FLOPs. What grows is "specialized" parameters with higher learning efficiency.

Q3: Why does recommendation MoE use ReLU Routing instead of Top-k?

A: Recommendation inputs are high-dimensional, sparse, and unevenly distributed; Top-k's fixed activation count easily overloads a few experts while the rest go underutilized. ReLU lets each token activate a variable number of experts based on information content, and with the regularization it achieves soft load balancing — better suited to recommendation data.

🔗 Connections to Later Chapters

• 7.3 (MTGR) — also handles heterogeneous features, but with GLN + Dynamic Masking; compare with RankMixer's Per-Token FFN (separate parameters) approach.
• 7.5 (OneTrans) — further breaks the fragmentation of "sequence modeling separate from feature interaction", extending RankMixer's hardware thinking to a unified architecture.
• 3.2 (Feature Crossing) — the high-order crossings of DCN/xDeepFM are re-implemented in RankMixer via multi-layer Token Mixing, and in a more GPU-friendly way.

Practice Problems

Work through all problems in order — they get progressively harder. Each has a complete solution you can reveal after trying it yourself.

Problem 7.4.1 — The Root Cause of Low MFU 🟢 Easy

A traditional DLRM's MFU is about 4–5%, an LLM's about 40–60%. Name the three architecture-level causes of the DLRM's low MFU.

Problem 7.4.2 — Token Mixing Complexity 🟢 Easy

Self-Attention complexity is $O(T^{2}D)$ and Token Mixing about $O(TD)$ (rearrangement). With $T=1000,D=256$, by about how many times do the orders of magnitude differ?

Problem 7.4.3 — Per-Token FFN vs MMoE 🟡 Medium

Why are Per-Token FFN and MMoE "fundamentally different"? Summarize each one's parameter organization in a sentence.

Problem 7.4.4 — DTSI-MoE Design 🔴 Hard

Explain why DTSI-MoE needs two routers ($h_{\text{train}}$ and $h_{\text{infer}}$), and what would happen if you trained with only $h_{\text{infer}}$.

🏆 Challenge: Hardware-Aware Restructuring

You must restructure a fragmented DLRM (embedding lookup + hand-crafted crossing + DIN + MLP) into a GPU-friendly architecture. Within 150 words, state which three RankMixer components you would substitute for the existing modules, and the expected MFU change plus its precondition.