M6.8 - Conditional Compute Beyond FFN Experts: Attention, Depth, and Modalities

0x00 Overview: A Router Does Not Define What Becomes Sparse

M6.1 and M6.2 route token rows among FFN experts. M6.6 turns those logical routes into distributed execution. The source Transformer Series also points to a broader design space: attention heads can be experts, tokens can be routed around whole blocks, and image/text tokens can share a sparsely activated multimodal expert pool.

These systems all use words such as router, expert, top-k, capacity, and load balance. Their tensor contracts are not the same.

What changes when routing selects attention components, token participation across depth, or conditional paths over multimodal token populations rather than FFN experts?

0x00A Historical Lineage: The Routed Object Keeps Changing

Conditional computation did not arrive as one mechanism. Each lineage changed a different object in the Transformer graph.

Early adaptive-depth work changed how long an example or token remains in the stack. The Universal Transformer paired a recurrent self-attention transition with dynamic per-position halting; Depth-Adaptive Transformer learned to predict at multiple depths and choose sequence-level decoder effort; and DeeBERT let an encoder example leave at an intermediate classifier. LayerDrop changed the training distribution over whole-layer subsets, but selecting a shallower subnetwork after training is not automatically input-dependent routing.

The next distinction was which token rows remain in the graph. PoWER-BERT progressively eliminated encoder word vectors, so later attention layers saw a smaller source set. That is not the same operation as preserving every row while choosing different FFN parameters for it.

The modern methods owned here change still different coordinates:

Lineage	First changed contract	What stays structurally fixed
Mixture of Attention Heads, arXiv 2210.05144v1	token-to-Q/O attention-expert identity	shared K/V projections and full block participation
SwitchHead, arXiv 2312.07987v3	source V mixture and destination O mixture	a smaller fixed set of Q/K attention maps
Mixture-of-Depths, arXiv 2404.02258v1	token membership in a whole block	an exact top-$C$ training shape per routed layer
MoE-LLaVA, arXiv 2401.15947v5	staged transition from dense multimodal alignment to sparse FFN experts	multimodal token sequence and ordinary attention path

This last row needs particular care. MoE-LLaVA does not simply read an image or text label and choose a modality-specific path. Its router dispatches multimodal token states among FFN experts. Modality is a token population over which load, specialization, quality, and stage-transition behavior can be measured.

LayerSkip later combined an early-exit training recipe with self-speculative verification. It is useful lineage, but M7 owns the draft/verify exactness contract; M6.8 uses it only to keep early exit distinct from Mixture-of-Depths membership.

0x00B From A Fixed Network To An Input-Dependent Path

A static Transformer applies the same scheduled block graph to every input. A conditional model starts from a supernetwork and lets a decision function choose a path through it:

$$y=\operatorname{Model}(x;g(x)).$$

The notation is intentionally broad. g(x) may return expert ids, a stopping depth, surviving token rows, active attention projections, or a layer-by-token execution mask. These choices cannot be reduced to one generic “router.”

For depth routing, define:

$$R_{l,t}\in\{0,1\}, \qquad R_{l,t}=1 \iff \text{token position }t\text{ executes block }l.$$

A monotone per-token stopping depth $D_t$ is the special case:

$$R_{l,t}=\mathbf{1}[l\le D_t].$$

Once token t exits, it never re-enters a deeper block. Arbitrary layer skipping and Mixture-of-Depths are more general: their columns of R need not be monotone. A token may bypass one block and execute a later one.

Expert routing changes a different object. Ordinary FFN MoE usually keeps R[l, t] = 1 while choosing an expert set $\mathcal{E}_{l,t}$. Routed attention can keep the token in the block while changing which projection or readout parameters are active. Progressive token pruning is different again: it removes token rows from the later sequence, changing the attention set as well as the amount of block work.

This yields a broader map than the modern mechanisms developed later in the article:

Family	Decision unit	What becomes conditional	Path property
sample/request early exit	one example or request	suffix of model depth	one monotone stopping depth
progressive token pruning	token row at a block	membership in later encoder layers	usually monotone removal
FFN MoE	token-expert pair	FFN parameter subset	full depth, different parameters
sequence-level adaptive depth	input sequence	common executed depth	one depth shared by its tokens
token-level halting	token-layer pair	remaining recurrent/depth steps	monotone depth per token
Mixture-of-Depths	token-layer pair	whole block versus residual path	non-monotone layer-by-token matrix
layer skipping	model, request, or token, depending on method	selected block subset	may be non-monotone

The comparison layout is reconstructed from the method taxonomy in the book’s Figure 6.16. The depth-by-token matrices and path labels are a corrected Transformer System interpretation grounded in the primary sources linked below.

DeeBERT and the Depth-Adaptive Transformer are examples of depth exit; PoWER-BERT progressively eliminates encoder token vectors; the Universal Transformer includes dynamic per-position halting; and LayerSkip combines an early-exit training recipe with self-speculative verification. LayerDrop instead trains a model from which shallower block subsets can be selected on demand; that does not by itself make the subset input-adaptive. Those papers anchor different rows, not one interchangeable algorithm.

Stopping output generation through EOS or a stopping criterion is also not “early exit” in this table. It shortens the generated sequence; it does not choose a shallower Transformer path for an otherwise continuing token.

The routed object determines:

which tensor dimension is gathered or masked
what the fixed compute budget counts
whether future tokens influence a route
which weights are active
how many attention matrices are formed
which K/V rows must persist
whether the runtime shape is static
what collapse or imbalance means

Conditional computation is therefore better classified by routing axis than by the presence of a top-k operator.

0x00C Four Routing Axes

Family	Routed choice	Fixed budget	Dynamic identity
FFN MoE	expert FFN per token	top-k routes or expert capacity	expert ids
routed attention	query/output or value/output expert components	active experts and/or attention maps	attention expert ids
Mixture-of-Depths	block computation versus residual path	tokens entering a block	token positions per layer
multimodal MoE	expert FFN for image/text token rows	top-k routes	expert and modality mix

0x00D Ownership And Boundaries

Relation	Owner	Exact contract
consumes	M3.1/M3.2/M3.4	dense score path, Q/K/V/O head ownership, residual branch semantics, and position-preserving block assembly
consumes	M5.1/M5.3	per-layer K/V append identity, original positions, prefix compatibility, exact reuse, and rollback state
consumes	M6.1	router scores, selected slots, gates, capacity vocabulary, and logical route metadata
consumes	M6.2	load denominators, balance controls, and the distinction between route affinity and physical execution
consumes	M6.6	logical-to-physical runtime boundary; M6.8 does not choose collectives or expert kernels
owns	M6.8	routing-axis taxonomy; MoA coupled Q/O experts; SwitchHead source/destination routing; MoD membership and gather/block/scatter; train/decode route transition; sparse per-layer K/V consequences; MoD+MoE composition; multimodal sparse-training stages
defers	M7	LayerSkip and other self-speculative draft/verify exactness
defers	X series	general image encoding, projectors, multimodal sequence construction, and modality foundations
defers	M6.11	static layer removal and exported structural pruning

The article may derive the new state created by a changed route, but it does not reteach ordinary MHA, MQA/GQA, generic KV lifecycle, FFN MoE dispatch, or multimodal architecture. Those are upstream contracts, not background to copy again.

0x00E Evidence Ladder And Primary Pins

No one source proves the complete route from an abstract conditional graph to causal K/V reuse and a deployable runtime. Evidence is therefore kept in five separate classes:

A29 and the author's book
    source-local lineage, original explanation, and figure navigation

version-pinned primary papers
    each paper's stated operator, training protocol, and source-reported result

M3/M5/M6 upstream articles
    Transformer System tensor, residual, route, balance, and cache contracts

later M6.8 CPU oracles
    only the frozen toy tensors and failure mutations they execute

browser and runtime evidence
    only the exact rendered page or measured environment named by that loop

The primary pins used by this article are 1807.03819v3, 1909.11556v1, 1910.10073v4, 2001.08950v5, 2004.12993v1, 2210.05144v1, 2312.07987v3, 2404.02258v1, 2401.15947v5, and 2404.16710v4. Paper benchmark numbers remain source-reported. The sparse per-layer K/V contract developed later is this program’s own systems inference unless a primary source states that exact cache implementation.

0x00F Immutable State Manifest: Axes Before Algorithms

The word token is not a sufficient tensor coordinate. A token row belongs to a batch item, sequence, original position, and layer. Depending on the method, it may also carry an attention-map head, an attention-expert identity, a source- or destination-router slot, a block-membership bit, or a modality population label. Those coordinates must be named before any gather, top-k, cache append, or FLOP count.

This article uses the following symbols consistently:

Symbol	Coordinate	Contract
$B$	batch items	independent sequences or requests
$S$	token rows in the current dense sequence	includes visual and textual rows after projection when the input is multimodal
$S_q$	destination/query rows	equals $S$ for square training self-attention and $1$ for one-token decode
$S_{kv}$	source K/V rows visible to those queries	may be a dense sequence, a prefix cache, or a layer-specific sparse subset
$d$	model width	the residual-stream width consumed from M3
$L$	transformer layers	layer_id is part of every depth and cache identity
$H_A$	computed attention maps	a map axis, not automatically an expert axis
$H_{kv}$	stored K/V heads	consumed from the ordinary M3/M5 attention contract
$d_h$	one projected head width	projection-local channel axis
$E$	candidates in one router pool	attention experts or FFN experts, depending on the declared method
$k_{qo},k_S,k_D,k_{ffn}$	selected candidate slots	MoA Q/O, SwitchHead source, SwitchHead destination, and FFN budgets are separate contracts
$C$	MoD training capacity per sequence and routed block	exactly $C$ original token identities enter the block
$T_{r,g,l}$	committed K/V rows for request $r$, branch generation $g$, and layer $l$	may differ across routed layers and from total generated length

The four records are deliberately not shape-isomorphic.

MoA Record: One Selection Couples Two Projection Banks

The source and destination axes are written separately even when a training fixture happens to use square self-attention:

destination/source states      X_q,X_kv   [B,S_q,d] / [B,S_kv,d]
router logits / probabilities z,p        [B,S_q,E]
selected expert ids           expert_id  [B,S_q,k_qo]  integer
selected route weights        w          [B,S_q,k_qo]  router numeric type
selected queries              Q_sel      [B,S_q,k_qo,d_h]
shared keys / values          K,V        [B,S_kv,d_h]
attention maps                A_sel      [B,S_q,k_qo,S_kv]
output                         Y          [B,S_q,d]

The same expert_id[b,t,r] indexes $W^Q$ and $W^O$. Its candidate pool is the MoA Q/O pool, not a global expert namespace. Neither $K$ nor $V$ has an $E$ axis. The selected-slot axis $k_{qo}$ creates up to $k_{qo}$ query-specific attention maps per destination token; it does not create $E$ historical K/V representations. The selected weights use the paper’s selected-probability normalization with a stopped-gradient denominator; that gradient policy is part of saved forward state, not an incidental softmax implementation detail.

SwitchHead Record: Map Identity Survives Expert Selection

SwitchHead’s reported V/O design carries two router coordinates around a fixed set of computed attention maps:

source/destination states      X_src,X_dst        [B,S_kv,d] / [B,S_q,d]
source scores / ids / weights sS,idS,wS          [B,S_kv,H_A,E] / [B,S_kv,H_A,k_S]
destination scores/ids/weights sD,idD,wD          [B,S_q,H_A,E] / [B,S_q,H_A,k_D]
shared queries                 Q                 [B,S_q,H_A,d_h]
shared keys / mixed values     K,V_mix           [B,S_kv,H_A,d_h]
attention maps                 A                 [B,H_A,S_q,S_kv]
output                         Y                 [B,S_q,d]

source and destination name the two sides of attention, not two physical devices. Their selected identities are separate unless a declared variant shares them. Expertizing and mixing $V$ before the attention core removes the $E$ axis before $A$ is formed; therefore map count is governed by $H_A$, not by $E$, $k_S$, or $k_D$. Each selected component ID is scoped by (attention_map_id, route_side, projection_family, candidate_pool_id). SwitchHead uses noncompetitive sigmoid selection, so selected route weights need not sum to one.

MoD Record: Membership Changes The Sequence Seen By One Layer

At routed layer $l$ during training:

hidden, positions, valid rows  X,pos,valid       [B,S,d] / [B,S] / [B,S]
router scalar                  r_l               [B,S]
top-C dense row indices        selected_row_l    [B,C] integer
gathered semantic positions    selected_pos_l    [B,C] integer
membership target              z_l               [B,S] Boolean, sum=C per sequence
gathered hidden                X_sel_l           [B,C,d]
causal relation                M_sel_l           [B,C,C], derived from original positions
block delta                    Delta_l           [B,C,d]
scattered output               X_next            [B,S,d]

selected_row_l drives gather and inverse scatter. selected_pos_l drives RoPE and causal visibility. They can differ under left padding, packing, or non-zero position offsets and must never be substituted for one another. Top-C is computed per sequence over valid rows; a deterministic tie policy is an implementation contract, not something fixed by the paper. The block helper returns a delta. Selected rows receive $x_i+r_i^l\Delta_i^l$; unselected rows remain $x_i$. A residualized block output cannot be scattered as though it were a delta.

At autoregressive decode, the training target $z_t^l$ is unavailable. The manifest therefore stores two distinct scalar roles:

membership predictor output   mhat_t^l   decides whether the block executes
original router scalar        r_t^l      weights the executed block delta

Conflating them changes the trained model path. If membership is false, that layer appends no K/V row. If it is true, a production representation must bind the appended row to request_id, branch_generation, layer_id, original_position, K/V projection identity, router/predictor/policy revision, and rollback generation. The paper does not freeze whether those ragged rows live in packed, paged, or padded storage.

Multimodal Record: Modality Is An Observation Coordinate

After the visual projector, both token populations inhabit the same residual width:

visual token states            X_visual           [B,N_v,d]
text token states              X_text             [B,N_t,d]
concatenated model sequence    X                  [B,S,d], S=N_v+N_t
population labels              modality_label     [B,S] metadata
FFN router scores              p                  [B,S,E]
selected FFN ids / gates       expert_id,w        [B,S,k_ffn]
training stage                 stage_id            {align,dense_multimodal,sparse_moe}

modality_label supports stratified load and specialization measurements. It is not an input to the published MoE-LLaVA router contract and is not an expert-selection result. The router consumes token states after multimodal projection and selects ordinary FFN experts.

The stage record also carries checkpoint lineage and parameter ownership: Stage I introduces the aligned projector state, Stage II produces the dense multimodal parent checkpoint, and Stage III replicates each selected dense FFN into its initial expert bank before introducing sparse routing. Random expert initialization is not equivalent to that transition. The paper reports BF16 for its experiments and notes instability encountered with FP16; this is an experiment configuration, not a universal architectural dtype requirement.

Dtype And Lifecycle Contract

The papers define mathematical operators more consistently than kernel-level precision. Therefore this article freezes roles, not invented FP16/BF16/ FP32 choices:

State	Required class	Lifecycle
hidden, projected Q/K/V, block delta	model numerical type	layer-local except committed K/V
router logits, probabilities, route weights	floating numeric state; exact accumulation dtype is implementation-specific here	logits/probabilities are ephemeral; route identity may be saved or replayed for backward
expert ids, dense row indices, semantic positions, layer ids	integer identity state	retained as long as inverse scatter, masking, cache lookup, or rollback needs it
membership	Boolean decision state	training target or decode decision, never a residual weight
modality label	categorical observation metadata	evaluation/telemetry coordinate, not model routing state
committed K/V	model-compatible cache state plus representation identity	request/prefix/branch lifetime owned by M5

The immutable identity for one conditional event is consequently more than (token_id, expert_id). It is at least:

(batch_id, sequence_id, layer_id, token_occurrence_id, original_position,
 routing_axis, candidate_pool_id, attention_map_id, route_side,
 projection_family, selected_component_ids, membership,
 decision_policy_revision, model_revision)

Floating route weights, logits, normalization denominators, selection ranks, and inverse permutations are saved forward payload, not stable identity. Backward or deterministic replay may require them, but cache compatibility should bind their producing revisions and policies rather than treating a floating value as a namespace key.

Physical rank, collective, kernel, packed-page address, and measured latency are intentionally absent. M6.8 emits logical changed-axis state; M6.6 and M5 own its physical execution and persistent representation.

0x01 Dense Attention As The Typed Reference Contract

M3 derives dense attention. This section does not repeat that derivation; it freezes the exact tensor and accounting coordinates that conditional attention will change. Source and destination lengths are deliberately asymmetric:

$$X_q\in\mathbb{R}^{B\times S_q\times d}, \qquad X_{kv}\in\mathbb{R}^{B\times S_{kv}\times d}.$$

For $H$ dense heads of width $d_h$ and $Hd_h=d$:

$$Q_h=X_qW_h^Q, \quad K_h=X_{kv}W_h^K, \quad V_h=X_{kv}W_h^V,$$

with $Q_h\in\mathbb{R}^{B\times S_q\times d_h}$ and $K_h,V_h\in\mathbb{R}^{B\times S_{kv}\times d_h}$. The score, row-wise probability, value reduction, and readout are:

$$S_h=\frac{Q_hK_h^{\top}}{\sqrt{d_h}}+M_h, \qquad A_h=\operatorname{softmax}_{S_{kv}}(S_h),$$ $$U_h=A_hV_h, \qquad Y=\sum_{h=1}^{H}U_hW_h^O.$$

The softmax axis is the source axis, not the head or destination axis. The mask $M_h\in\mathbb{R}^{B\times S_q\times S_{kv}}$ must already encode causal, padding, and any structural visibility restrictions before normalization.

0x01A Logical Graph Versus Materialized State

The equations name a logical graph. They do not require every intermediate to be written to HBM:

State	Logical shape	Needed by the mathematics	May be fused or recomputed
projected query	[B,H,S_q,d_h]	yes	sometimes
projected key/value	[B,H,S_kv,d_h]	yes	decode K/V persists
scores	[B,H,S_q,S_kv]	yes	commonly tiled
probabilities	[B,H,S_q,S_kv]	yes	commonly tiled/recomputed
value reduction	[B,H,S_q,d_h]	yes	may feed output projection directly

M4 owns the fused-kernel schedule. Here, a state is listed because routing can add an expert axis, remove a head axis, or change its identity. “Not materialized” does not mean “not computed.”

0x01B Dense Parameter And Named-Operation MAC Ledger

Ignoring biases, dense Q/K/V/O parameters are:

$$P_{\mathrm{dense}} = 4Hd d_h = 4d^2.$$

Counting one multiply-accumulate as one MAC and naming every matrix product:

$$\begin{aligned} C_Q &= BS_qd^2, & C_O &= BS_qd^2,\\ C_K+C_V &= 2BS_{kv}d^2,\\ C_{QK^\top} &= BHS_qS_{kv}d_h, & C_{AV} &= BHS_qS_{kv}d_h. \end{aligned}$$

Therefore:

$$C_{\mathrm{dense}} = 2BS_qd^2+2BS_{kv}d^2+2BS_qS_{kv}d.$$

This ledger counts both quadratic matrix products. A paper or profiler may use a different convention; comparisons are valid only after the included operations, batch, lengths, and MAC-versus-FLOP convention are reconciled.

During one-token autoregressive decode, $S_q=1$ and the historical source length is $S_{kv}=T$. K/V persistence is determined by the projection path used when each source token entered the layer. Routed attention must therefore state which of $Q,K,V,O$ carries expert identity. Without that ownership statement, neither active compute nor cache compatibility is defined.

0x02 Mixture Of Attention Heads: Couple Query And Output Experts

The EMNLP 2022 Mixture of Attention Heads paper replaces a fixed bank of heads with $E$ candidate attention experts and selects $k$ for each destination token. Its specific construction matters: expert $i$ owns $W_i^Q$ and $W_i^O$, while all experts share $W^K$ and $W^V$.

Source: Mixture of Attention Heads, Figure 2. One token-level top-k decision indexes matching query and output projections. Shared K/V keeps the historical source representation independent of the current destination route.

The two expert banks in the figure are not independent MoEs. A selected slot uses the same expert id for its query map and output readout. Uncoupling those ids defines a different operator.

0x02A Complete Forward Tensor Trace

Let:

$$W_g\in\mathbb{R}^{d\times E}, \quad W^Q\in\mathbb{R}^{E\times d\times d_h}, \quad W^O\in\mathbb{R}^{E\times d_h\times d},$$ $$W^K,W^V\in\mathbb{R}^{d\times d_h}.$$

The destination router is:

$$L=X_qW_g \in\mathbb{R}^{B\times S_q\times E}, \qquad P=\operatorname{softmax}_{E}(L).$$

For destination row $(b,t)$, stable top-k selection returns an ordered expert set $G_{b,t}$ and selected probabilities $P_G[b,t,s]$ for slot $s\in\{0,\ldots,k-1\}$. Define:

$$G_{b,t} = \operatorname{StableTopK}_{i\in[0,E)}(P[b,t,i],k).$$ $$Z_{b,t}=\sum_{i\in G_{b,t}}P[b,t,i], \qquad w_{b,t,i} = \frac{P[b,t,i]}{\operatorname{StopGrad}(Z_{b,t})}.$$

The forward values sum to one, but the backward is not the derivative of an ordinary selected softmax. That distinction is derived in 0x02C.

Shared source projections are computed once:

$$K=X_{kv}W^K, \qquad V=X_{kv}W^V, \qquad K,V\in\mathbb{R}^{B\times S_{kv}\times d_h}.$$

Gathering the coupled Q/O expert ids produces:

$$Q_G[b,t,s,:] = X_q[b,t,:]W^Q_{G_{b,t}[s]} \in\mathbb{R}^{B\times S_q\times k\times d_h}.$$

Each selected query expert forms its own source distribution:

$$S_G[b,t,s,j] = \frac{\langle Q_G[b,t,s,:],K[b,j,:]\rangle}{\sqrt{d_h}} +M[b,t,j],$$ $$A_G=\operatorname{softmax}_{S_{kv}}(S_G) \in\mathbb{R}^{B\times S_q\times k\times S_{kv}}.$$

The selected value reductions and paired output experts are:

$$U_G[b,t,s,:] = \sum_{j=1}^{S_{kv}}A_G[b,t,s,j]V[b,j,:],$$ $$E_G[b,t,s,:] = U_G[b,t,s,:]W^O_{G_{b,t}[s]},$$ $$Y[b,t,:] = \sum_{s=1}^{k}w_G[b,t,s]E_G[b,t,s,:].$$

The map count is $BS_qk$, not $BS_qE$, but all $E$ Q/O parameter pairs remain resident unless sharded or offloaded by a separate runtime contract.

0x02B What The Forward Must Save

Top-k is piecewise constant, not globally differentiable. A backward or exact replay contract needs more than the final output:

stable identity:
    batch, destination row, semantic position, layer
    ordered selected expert ids
    Q/O candidate-pool and policy revisions

floating forward payload:
    full router probabilities P or sufficient softmax state
    selected probabilities P_G and stopped denominator Z
    selected attention row state or a recomputation contract
    dropout mask when attention dropout is active

permutation payload for grouped projections:
    origin row, selected slot, expert id, expert-local row
    inverse mapping back to [B, S_q, k]

Selected ids are integer identity. Probabilities and denominators are saved floating state. Treating either one as the other breaks deterministic replay.

0x02C Why Stop-Gradient Changes Router Learning

Fix one destination row and assume its top-k membership $G$ does not change in the local derivative. Write $p=\operatorname{softmax}(\ell)$, $Z=\sum_{j\in G}p_j$, and $w_i=p_i/\operatorname{StopGrad}(Z)$ for $i\in G$. Then:

$$\frac{\partial w_i}{\partial p_j} = \frac{\mathbf{1}[i=j]}{Z}.$$

Let upstream output gradient be $g_y$ and define the scalar contribution of expert $i$ as $c_i=\langle g_y,E_i\rangle$. The gradient entering the full softmax probability vector is:

$$u_i = \begin{cases} c_i/Z,&i\in G,\\ 0,&i\notin G. \end{cases}$$

Using the softmax Jacobian $\partial p_j/\partial\ell_m=p_j(\mathbf{1}[j=m]-p_m)$ gives:

$$\frac{\partial\mathcal{L}}{\partial\ell_m} = p_m\left(u_m-\sum_jp_ju_j\right).$$

The subtle consequence is that a non-selected logit can receive a nonzero gradient through the full softmax normalization even though top-k membership is held fixed.

Ordinary differentiable renormalization would instead use $w_i=p_i/Z$ with gradient through $Z$:

$$\frac{\partial w_i}{\partial p_j} = \frac{\mathbf{1}[i=j]Z-p_i\mathbf{1}[j\in G]}{Z^2}.$$

That forward is equivalent to a softmax restricted to selected logits. Its non-selected logit gradients are zero, and its selected logit gradients sum to zero. The two backward rules are therefore not interchangeable even though their numerical forward weights are identical.

For a concrete router row:

p = [0.50, 0.30, 0.20]       G = [0, 1]       Z = 0.80
w = [0.625, 0.375]
upstream expert contributions c = [2.0, -1.0]

stopped-denominator dL/dlogit:
    [ 0.8125, -0.6375, -0.1750 ]

ordinary selected-renormalization dL/dlogit:
    [ 0.703125, -0.703125, 0.0 ]

This is a custom-gradient contract. A finite-difference test of the numerical forward $p_i/Z$ recovers the ordinary quotient derivative, not the stop-gradient VJP. The executable oracle consequently tests the stopped rule against its declared VJP and finite-differences only the smooth attention and projection paths.

0x02D Full Selected-Expert Backward Trace

For one selected expert and one destination row, abbreviate:

$$q_i=xW_i^Q, \quad s_i=\frac{q_iK^\top}{\sqrt{d_h}}+m, \quad a_i=\operatorname{softmax}(s_i),$$ $$u_i=a_iV, \quad e_i=u_iW_i^O, \quad y=\sum_{i\in G}w_ie_i.$$

With upstream $g=\partial\mathcal{L}/\partial y$, the expert path begins:

$$g_{e_i}=w_i g, \qquad \nabla W_i^O=u_i^\top g_{e_i}, \qquad g_{u_i}=g_{e_i}(W_i^O)^\top.$$

The shared value and attention-probability gradients are:

$$\nabla V\mathrel{+}=a_i^\top g_{u_i}, \qquad g_{a_i}=g_{u_i}V^\top.$$

The row-wise softmax VJP is:

$$g_{s_i} = a_i\odot \left(g_{a_i}-\langle a_i,g_{a_i}\rangle\mathbf{1}\right).$$

Then:

$$g_{q_i}=\frac{g_{s_i}K}{\sqrt{d_h}}, \qquad \nabla K\mathrel{+}=\frac{g_{s_i}^\top q_i}{\sqrt{d_h}},$$ $$\nabla W_i^Q=x^\top g_{q_i}, \qquad g_x^{\mathrm{expert}}\mathrel{+}=g_{q_i}(W_i^Q)^\top.$$

Finally, shared projection gradients are:

$$\nabla W^K=X_{kv}^\top\nabla K, \qquad \nabla W^V=X_{kv}^\top\nabla V.$$

Both $\nabla K$ and $\nabla V$ accumulate contributions from every selected expert and destination row that read the same source state. The router supplies a second input-gradient path through $L=X_qW_g$. A correct implementation sums the expert-query path and router path before returning $\nabla X_q$.

0x02E Parameters: Resident Is Not Active

Ignoring biases, the paper’s Q/O expert bank plus shared K/V projections has:

$$P_{\mathrm{MoA,paper}} = (2E+2)dd_h.$$

This is Equation 20 of the paper. It excludes the router matrix. An explicit checkpoint ledger adds:

$$P_{\mathrm{router}}=dE, \qquad P_{\mathrm{MoA,checkpoint}}=(2E+2)dd_h+dE,$$

plus any projection, router, or noise biases actually enabled by an implementation. For one destination token, only $k$ Q/O pairs are active:

$$P_{\mathrm{active/token}} = 2kdd_h+2dd_h+dE,$$

where the router is evaluated densely over $E$. This active-parameter count is not resident memory and is not a byte-traffic measurement.

0x02F Compute: Reconcile Two Conventions Instead Of Hiding One

A named-operation logical MAC ledger for asymmetric lengths is:

$$\begin{aligned} C_{\mathrm{router}} &= BS_qdE,\\ C_{KV} &= 2BS_{kv}dd_h,\\ C_{QO} &= 2BS_qkdd_h,\\ C_{QK^\top} &= BkS_qS_{kv}d_h,\\ C_{AV} &= BkS_qS_{kv}d_h. \end{aligned}$$

Therefore:

$$C_{\mathrm{MoA,named}} = BS_qdE +2BS_{kv}dd_h +2BS_qkdd_h +2BkS_qS_{kv}d_h.$$

The paper prints, for $B=1$ and $S_q=S_{kv}=T$:

$$C_{\mathrm{MoA,paper}} = kT^2d_h+2(k+1)Tdd_h,$$

and compares it with $C_{\mathrm{MHA,paper}}=T^2d+4Td^2$. Those expressions exclude router cost and place one $T^2$ attention term in the ledger, whereas the named-operation derivation above counts both $QK^\top$ and $AV$. This article preserves the source formula verbatim and does not silently call it the same accounting convention. The paper’s reported MAC tables were measured with PTFLOPS on fixed input shapes; they are source-reported profiler outputs, not a substitute for the operator ledger.

0x02G Why K And V Are Shared

Suppose every expert instead owned $W_i^K,W_i^V$. A future destination could select expert $i$ even if an earlier source token did not. Serving would then need either:

1. every expert-specific K/V row precomputed and cached, multiplying persistent state by $E$;
2. historical hidden states retained so missing expert K/V can be projected on demand; or
3. a source-side routing policy that freezes which expert K/V identities exist, producing a different conditional operator.

MoA chooses none of these. Shared $W^K,W^V$ yields one reusable source representation:

append for a new source token:
    K_new = x_new @ W_k       [B, 1, d_h]
    V_new = x_new @ W_v       [B, 1, d_h]

read for the current destination:
    choose k coupled Q/O expert ids
    form k query-specific score rows over the same K cache
    reduce the same V cache and apply the paired O experts

The route changes current map and readout computation, not historical K/V projection identity. During cross-attention with static_kv, the same shared source K/V may be reused across destination steps.

0x02H Paper Contract Versus The Pinned Public Code

The official public repository is useful implementation evidence, but it is not byte-for-byte the paper equations. At commit 96c33ca:

Paper contract	Pinned Fairseq implementation
Equation 9 divides selected probabilities by a stopped denominator	the corresponding normalization lines in moe.py are commented out; batch_gates are raw selected full-softmax probabilities
one top-k decision couples Q and O expert ids	q_proj.map(...) records sorted expert membership; q_proj.reduce(...) reuses that saved membership and gate vector
K/V are shared across attention experts	k_proj and v_proj are ordinary Linear(..., head_dim) modules with no expert axis
conceptual selected tensor [B,S_q,k,d_h]	map flattens rows, groups them by expert, then reconstructs [B,S_q,k,d_h] through an index-add permutation
incremental shared K/V	prev_key and prev_value are appended or reused for static_kv; the saved shape has no expert axis

The code also supports noisy gating and sampled top-k options not required by the core paper equations. Its map/attention/reduce schedule is a historical Fairseq implementation, not a claim about a fused modern serving kernel. The paper itself reports extra memory-copy overhead and states that its implementation did not fully exploit GPU parallelism.

The paper-faithful forward can be written as:

def moa_forward(x_q, x_kv, mask, params, k):
    # x_q: [B, S_q, d]; x_kv: [B, S_kv, d]
    logits = x_q @ params.W_gate                    # [B, S_q, E]
    probs = softmax(logits, dim=-1)
    selected_p, selected_id = stable_topk(probs, k) # [B, S_q, k]

    denominator = selected_p.sum(-1, keepdim=True)
    route_weight = selected_p / stop_gradient(denominator)

    K = x_kv @ params.W_k                           # [B, S_kv, d_h]
    V = x_kv @ params.W_v                           # [B, S_kv, d_h]

    Q = gather_expert_linear(x_q, params.W_q, selected_id)
    # Q: [B, S_q, k, d_h]
    score = einsum("btkh,bjh->btkj", Q, K) / sqrt(d_h)
    attention = softmax(score + mask[:, :, None, :], dim=-1)
    U = einsum("btkj,bjh->btkh", attention, V)
    expert_y = gather_expert_linear(U, params.W_o, selected_id)

    return (route_weight[..., None] * expert_y).sum(dim=2)

Grouped or segmented expert projection must preserve the route tuple (origin_row, selected_slot, expert_id, expert_local_row, route_weight) through both map and reduce. Merely retaining an expert id is insufficient to restore destination order.

0x02I Router Objectives And Source-Reported Training Contract

For $T$ valid destination queries, the paper defines expert assignment count $f_i$ and full-softmax mass $P_i$:

$$f_i=\sum_{t=1}^{T}\mathbf{1}[i\in G_t], \qquad P_i=\sum_{t=1}^{T}p_{i,t}.$$

After the paper’s stated column-wise norm-1 normalization, its balancing loss is:

$$\mathcal{L}_a=E\sum_{i=1}^{E}f_iP_i.$$

The hard count $f_i$ supplies no ordinary gradient; $P_i$ supplies the smooth router path. Its router z-loss is:

$$\mathcal{L}_z = \frac{1}{T}\sum_{t=1}^{T} \left(\log\sum_{i=1}^{E}e^{\ell_{t,i}}\right)^2.$$

The source-reported experiment objective is:

$$\mathcal{L} = \mathcal{L}_{\mathrm{model}} +\sum_{\text{MoA modules}} \left(0.01\mathcal{L}_a+0.001\mathcal{L}_z\right).$$

Those coefficients are experimental choices, not universal defaults. The paper trained its reported models on 32 V100 GPUs with Adam $(\beta_1,\beta_2,\epsilon)=(0.9,0.98,10^{-9})$, label smoothing $0.1$, and task-specific schedules and batch sizes. Appendix D reports an auxiliary-loss ablation rather than proving that one coefficient pair transfers to another model or workload.

Attention-expert collapse has two coupled effects: few Q/O pairs receive useful gradient, and the actual query-map workload becomes skewed. Generic load denominators and distributed balancing remain M6.2’s ownership; M6.8 records only the method-local objective and the changed attention axis.

0x02J Design Alternatives And Their State Consequences

Design	Active maps per destination	Historical K/V identity	Main tradeoff
fixed dense MHA	$H$	one K/V pair per fixed head	no conditional parameter choice
MoA Q/O experts, shared K/V	$k$	one shared K/V pair	conditional maps/readouts without expert K/V cache
full Q/K/V/O expert heads	$k$ if source history is available	potentially $E$ expert-specific K/V families	larger or lazily reconstructed persistent state
$k=E$ MoA	$E$	one shared K/V pair	dense mixture; no projection sparsity
SwitchHead	fixed $H_A$	per-map shared K plus source-routed mixed V	separates map count from V/O parameter capacity

MoA’s shared K/V resembles the state consequence of aggressive K/V sharing, but it is not ordinary MQA/GQA: its destination token still chooses among Q/O experts dynamically. Any combination with grouped K/V projections must declare an additional fixed or routed group axis rather than inheriting compatibility by analogy.

0x02K Loop 3 Executable Closure

The checked-in Loop 3 oracle uses an asymmetric fixture B=1, S_q=2, S_kv=3, d=2, d_h=2, E=3, k=2. It verifies:

dense and MoA tensor shapes
stable coupled Q/O selection
shared K/V without an expert axis
paper stopped-denominator router VJP
ordinary-renormalization counterfactual
selected Q/O and shared K/V backward equations
finite differences for smooth projection parameters
resident/active parameter ledgers
named-operation and paper MAC conventions
decode cache identity and official-code divergence

Negative fixtures reject uncoupled Q/O ids, expertized K/V, wrong softmax axes, collapsed source/destination lengths, an unstopped paper denominator, omitted inverse route metadata, router-cost convention collapse, unstable top-k ties, and claims that the historical public code implements Equation 9 exactly. The oracle executes no GPU kernel and measures no latency.

0x03 SwitchHead: Move Conditional Projection Outside The Attention Core

The NeurIPS 2024 SwitchHead paper starts from a narrower systems question than “can attention be an MoE?” At long sequence length, the expensive state is not only projection capacity. It is the number of full source-by-destination attention maps that must be computed and, under a materialized training schedule, stored. SwitchHead therefore asks:

Can a model retain a large bank of attention projection choices while computing only a small fixed number of attention matrices?

The answer depends on where expert mixtures occur. Source-side value experts are mixed before attention; destination-side output experts are mixed after attention. The expert axes therefore surround the quadratic core rather than entering its map identity.

Source: SwitchHead, Figure 1, revision v3. The figure uses “head” for one computed attention matrix. Each such head owns multiple V and O projection experts.

0x03A Why The Dense Output Can Be A Sum Of Head-Local Readouts

Let $H_A$ denote the number of computed attention maps. The subscript is deliberate: later, $E$ will count projection experts and must not be mistaken for another map axis. Split the ordinary output projection into head-specific row blocks $W_h^O\in\mathbb{R}^{d_h\times d}$:

$$W^O = \begin{bmatrix} W_1^O\\ \vdots\\ W_{H_A}^O \end{bmatrix}.$$

If $U_h=A_hV_h\in\mathbb{R}^{B\times S_q\times d_h}$, concatenation followed by one output matrix is algebraically the same as a sum of head-local projections:

$$Y = \operatorname{Concat}(U_1,\ldots,U_{H_A})W^O = \sum_{h=1}^{H_A}U_hW_h^O.$$

This decomposition exposes the design coordinate. Each $h$ can retain one attention matrix $A_h$ while choosing among several local projection matrices. For dense attention, the two quadratic products are:

$$Q_hK_h^{\top} \quad\text{and}\quad A_hV_h,$$

so their MAC count is $2BH_AS_qS_{kv}d_h$. A materialized training schedule also has an $O(BH_AS_qS_{kv})$ score/probability state. M4 may tile or recompute that state; changing materialization does not change the map identity counted here.

0x03B Why Naive Destination-Only Head Routing Does Not Remove Source Work

A tempting first design is to score a pool of complete heads from each destination token:

$$g_q=\sigma(X_qW_R), \qquad \mathcal{I}_q=\operatorname{TopK}(g_q,k),$$ $$Y_q = \sum_{e\in\mathcal{I}_q} g_{q,e}\,(A_eV_eW_e^O)_q.$$

This makes destination work sparse, but an attention entry couples a destination $q$ to a source $s$. The source K/V projections required by a batch are the union of every destination’s selected heads:

$$\mathcal{U}_{KV} = \bigcup_{q=1}^{S_q}\mathcal{I}_q.$$

For example:

q0 selects heads {0, 1}
q1 selects heads {1, 2}
q2 selects heads {2, 3}

union of required source K/V projections = {0, 1, 2, 3}

Only two heads were selected per destination, yet every source token may need all four K/V projections. Irregularly gathering a different head set for each destination also complicates a dense attention kernel. This is the paper’s motivation for giving source and destination projections their own routing decisions outside the attention core.

0x03C General Source/Destination Projection Family

For each computed map head $h\in\{1,\ldots,H_A\}$, let there be $E$ projection experts. Source tokens produce one noncompetitive sigmoid route:

$$z^S_{b,s,h,e} = \sum_{c=1}^{d}X_{kv}[b,s,c]R^S[h,c,e], \qquad g^S_{b,s,h,e}=\sigma(z^S_{b,s,h,e}),$$ $$\mathcal{I}^S_{b,s,h} = \operatorname{TopK}_{e} \left(g^S_{b,s,h,:},k_S\right).$$

Destination tokens analogously produce $g^D$ and $\mathcal{I}^D_{b,q,h}$ from $X_q$ and $R^D$. The full candidate family can expertize K/V on the source side and Q/O on the destination side:

$$\begin{aligned} K_{b,s,h} &=\sum_{e\in\mathcal{I}^S_{b,s,h}}g^S_{b,s,h,e} X_{kv}[b,s]W^K_{h,e},\\ V_{b,s,h} &=\sum_{e\in\mathcal{I}^S_{b,s,h}}g^S_{b,s,h,e} X_{kv}[b,s]W^V_{h,e},\\ Q_{b,q,h} &=\sum_{e\in\mathcal{I}^D_{b,q,h}}g^D_{b,q,h,e} X_q[b,q]W^Q_{h,e}. \end{aligned}$$

The selected sigmoid weights are gathered but not renormalized to sum to one. For the Loop 4 fixture, selected source sums range from about 1.10 to 1.21; destination sums range from about 1.10 to 1.15. Their amplitude is part of the learned operator, not a probability-simplex invariant.

0x03D Why The Reported Final Design Expertizes V And O

The source paper exhaustively evaluated all 16 choices of expert versus fixed V/K/Q/O projections in one parameter-matched setting: 47M-parameter models, WikiText-103, $E=5$, and $H_A=2$. The complete source-reported sweep is:

V	K	Q	O	perplexity $\downarrow$
expert	fixed	fixed	expert	12.27
fixed	fixed	fixed	expert	12.30
fixed	expert	fixed	expert	12.36
expert	expert	fixed	expert	12.37
expert	fixed	expert	expert	12.42
expert	fixed	fixed	fixed	12.45
fixed	fixed	expert	expert	12.45
expert	fixed	expert	fixed	12.51
expert	expert	expert	expert	12.57
fixed	expert	expert	expert	12.59
expert	expert	expert	fixed	12.61
expert	expert	fixed	fixed	12.69
fixed	fixed	expert	fixed	12.75
fixed	expert	fixed	fixed	12.79
fixed	expert	expert	fixed	12.90

The same table reports 12.31 for a parameter-matched dense 10-head baseline and 12.74 for a dense 2-head baseline. V/O therefore gave the best result in that experiment, while O-only was close. But this is not a theorem about Q/K:

1. The comparison holds total parameters fixed, so assigning experts to Q/K removes parameter budget from other model components.
2. The paper reports that preliminary experiments without the fixed parameter budget found that adding more experts helped.
3. Different scale, data, kernels, and parameter allocation may change the ordering.

The final SwitchHead operator used below has head-specific shared Q/K and expert V/O.

0x03E Final V/O Tensor Trace

With asymmetric destination and source lengths:

$$X_q\in\mathbb{R}^{B\times S_q\times d}, \qquad X_{kv}\in\mathbb{R}^{B\times S_{kv}\times d}.$$

Shared head-specific projections are:

$$Q_{b,q,h,:}=X_q[b,q,:]W_h^Q, \qquad K_{b,s,h,:}=X_{kv}[b,s,:]W_h^K,$$

with $Q\in\mathbb{R}^{B\times S_q\times H_A\times d_h}$ and $K\in\mathbb{R}^{B\times S_{kv}\times H_A\times d_h}$. Selected source experts first produce:

$$P^V_{b,s,h,j,:} = X_{kv}[b,s,:]W^V_{h,e_j}, \qquad e_j=\mathcal{I}^S_{b,s,h}[j],$$ $$V^{\mathrm{mix}}_{b,s,h,:} = \sum_{j=1}^{k_S} g^S_{b,s,h,e_j}P^V_{b,s,h,j,:}.$$

The attention core has no E or selected-slot axis:

$$S_{b,h,q,s} = \frac{Q_{b,q,h,:}K_{b,s,h,:}^{\top}}{\sqrt{d_h}} +M_{b,q,s},$$ $$A_{b,h,q,:} = \operatorname{softmax}_{S_{kv}}(S_{b,h,q,:}), \qquad U_{b,q,h,:} = \sum_{s=1}^{S_{kv}}A_{b,h,q,s}V^{\mathrm{mix}}_{b,s,h,:}.$$

Finally, destination experts project the common readout:

$$P^O_{b,q,h,j,:} = U_{b,q,h,:}W^O_{h,e_j}, \qquad e_j=\mathcal{I}^D_{b,q,h}[j],$$ $$Y_{b,q,:} = \sum_{h=1}^{H_A}\sum_{j=1}^{k_D} g^D_{b,q,h,e_j}P^O_{b,q,h,j,:}.$$

An implementation trace must preserve that ordering:

def switchhead_vo(X_q, X_kv, W_q, W_k, W_v, W_o, R_s, R_d, k_s, k_d, mask):
    # X_q:  [B, S_q,  d]
    # X_kv: [B, S_kv, d]
    Q = project_shared_heads(X_q, W_q)       # [B, H_A, S_q, d_h]
    K = project_shared_heads(X_kv, W_k)      # [B, H_A, S_kv, d_h]

    src_logits = route(X_kv, R_s)            # [B, S_kv, H_A, E]
    src_ids = stable_topk(src_logits, k_s)
    src_gates = gather(sigmoid(src_logits), src_ids)
    V_selected = selected_project(X_kv, W_v, src_ids)
    V_mixed = weighted_sum(V_selected, src_gates, axis="selected_slot")

    scores = Q @ transpose_source(K) / sqrt(d_h)
    A = softmax(scores + mask, axis="S_kv") # [B, H_A, S_q, S_kv]
    U = A @ V_mixed                           # [B, H_A, S_q, d_h]

    dst_logits = route(X_q, R_d)             # [B, S_q, H_A, E]
    dst_ids = stable_topk(dst_logits, k_d)
    dst_gates = gather(sigmoid(dst_logits), dst_ids)
    O_selected = selected_project(U, W_o, dst_ids)
    return weighted_sum(O_selected, dst_gates, axes=("head", "selected_slot"))

If V is projected after attention, E becomes entangled with map/readout work. If O is mixed before the readout, the destination expert no longer acts on the head result described by the paper. These are different operators even when all final shapes still match.

0x03F Backward: Selected Sigmoid Gates Are Noncompetitive

Treat top-k membership as fixed inside one smooth region. For a selected gate $g_e=\sigma(z_e)$:

$$\frac{\partial\mathcal{L}}{\partial z_e} = \mathbf{1}[e\in\mathcal{I}] \frac{\partial\mathcal{L}}{\partial g_e} g_e(1-g_e).$$

There is no selected-softmax cross term $-g_eg_j\,\partial\mathcal{L}/\partial g_j$. Changing one selected sigmoid logit changes its gate directly; it does not force another selected gate down. Unselected logits have zero local derivative through the hard gather. At a top-k membership boundary the ordinary derivative is undefined, so the piecewise-smooth formula must not be advertised as differentiating through the discrete ordering.

For source expert $e$ at $(b,s,h)$:

$$\frac{\partial\mathcal{L}}{\partial W^V_{h,e}} += g^S_{b,s,h,e} X_{kv}[b,s,:]^{\top} \frac{\partial\mathcal{L}}{\partial V^{\mathrm{mix}}_{b,s,h,:}},$$ $$\frac{\partial\mathcal{L}}{\partial g^S_{b,s,h,e}} = \left\langle \frac{\partial\mathcal{L}}{\partial V^{\mathrm{mix}}_{b,s,h,:}}, P^V_{b,s,h,e,:} \right\rangle.$$

Destination O gradients have the same structure with $U$, $W^O$, and $g^D$. Router matrices then receive:

$$\frac{\partial\mathcal{L}}{\partial R^S_{h,:,e}} += X_{kv}[b,s,:]^{\top} \frac{\partial\mathcal{L}}{\partial z^S_{b,s,h,e}},$$

and analogously for $R^D$ from $X_q$. The remaining $U\rightarrow A,V^{\mathrm{mix}}$ and $A\rightarrow Q,K$ paths are the dense attention backward from 0x01. The checked Loop 4 oracle closes all six smooth parameter banks W_q, W_k, W_v, W_o, R_s, and R_d by central finite difference; the largest fixture error is below $5\times10^{-11}$.

0x03G Map Count: H_A Is Independent Of E And k

SwitchHead’s logical attention tensor is:

$$A\in \mathbb{R}^{B\times H_A\times S_q\times S_{kv}}.$$

Therefore:

$$N_{\mathrm{map\ elements}}^{\mathrm{SH}} = BH_AS_qS_{kv}.$$

A hypothetical dense model with one map for every $(h,e)$ pair has $BH_AES_qS_{kv}$ elements. The expert pool increases resident projection capacity, and $k_S,k_D$ increase selected projection work, but neither changes the SwitchHead map tensor. In the asymmetric oracle B=1,H_A=2,E=3,S_q=2,S_kv=3, SwitchHead computes 12 map elements; the virtual $H_AE$ dense bank computes 36.

MoA requires a different ledger. Its selected expert rows contain $BS_qkS_{kv}$ elements, but each row has an expert-specific Q projection and map identity. A concrete kernel may materialize only gathered rows, the union of selected experts, or a denser expert tensor. “MoA has k maps” is therefore insufficient without saying whether the count is per destination, per unique expert, or per materialized tensor.

0x03H Parameters And Compute: Reconcile Local And Paper Conventions

Ignoring bias and using equal input/output width $d$, final V/O SwitchHead has:

$$\begin{aligned} P_{QK} &= 2H_Add_h,\\ P_{VO} &= 2H_AEdd_h,\\ P_{R} &= 2H_AdE,\\ P_{\mathrm{checkpoint}} &= 2H_Add_h+2H_AEdd_h+2H_AdE. \end{aligned}$$

For one source token, projection-plus-router active parameters are $H_Add_h+H_Ak_Sdd_h+H_AdE$. The analogous destination ledger replaces K/V with Q/O and $k_S$ with $k_D$. These are conceptual active-work ledgers; resident memory still contains the whole V/O expert bank unless another system shards or offloads it.

For asymmetric lengths, a named-operation MAC ledger is:

Operation	MACs
shared Q projection	$BS_qH_Add_h$
shared K projection	$BS_{kv}H_Add_h$
source router	$BS_{kv}H_AdE$
selected V projection	$BS_{kv}H_Ak_Sdd_h$
source weighted reduction	$BS_{kv}H_Ak_Sd_h$
$QK^{\top}$	$BH_AS_qS_{kv}d_h$
$AV$	$BH_AS_qS_{kv}d_h$
destination router	$BS_qH_AdE$
selected O projection	$BS_qH_Ak_Dd_hd$
destination weighted reduction	$BS_qH_Ak_Dd$

The Loop 4 teaching fixture totals 268 MACs under that complete convention. Its 88 checkpoint parameters split into 16 shared Q/K, 48 resident V/O expert, and 24 router parameters.

The paper uses a different, symmetric Transformer-XL convention. Let the current chunk length be $T$, the source context length be $CT$, and one common active count be $k$. It reports:

$$N_{\mathrm{MAC}}^{\mathrm{SH}} = H_A\left[ 2T d_hd +2Tk d_h(d+1) +2CT^2d_h +2CTd_hd \right].$$

The terms are shared Q/K, selected V/O plus their weighted reductions, two attention products, and the paper’s Transformer-XL positional-projection work. The paper explicitly ignores the selector cost as small. Its materialized training-state formula is:

$$N_{\mathrm{mem}}^{\mathrm{SH}} = H_A\left(4Td_h+2CT^2+2CTd_h\right) \quad\text{floats}.$$

That memory expression assumes a “smart kernel” avoids a $k$-scaled stored projection state. It is not a universal peak-memory result and not an HBM traffic model. FlashAttention can tile the quadratic state and is described by the paper as orthogonal to SwitchHead; M4 owns that fused schedule.

0x03I Decode Cache ABI: Route V At Insertion, Route O At Query Time

At one decode insertion for token $t$, each computed head appends:

$$K_h(t)=x_tW_h^K \in\mathbb{R}^{d_h},$$ $$V_h^{\mathrm{mix}}(t) = \sum_{e\in\mathcal{I}_{t,h}^S} g^S_{t,h,e}\,x_tW^V_{h,e} \in\mathbb{R}^{d_h}.$$

The append shapes are [B,H_A,1,d_h] for both K and mixed V. Inference does not need to cache E raw V rows or source route ids once the mixture has been materialized. Training may still save or deterministically recompute route ids and weights for backward.

The historical V row depends on:

W_K and the K position transform
W_V expert bank and R_s
k_S, sigmoid, top-k, and tie semantics
projection and accumulation dtype
checkpoint and cache generation

Changing any of those invalidates exact reuse. By contrast, $W_Q$, $R^D$, and $W^O$ change how the current destination reads history but do not define the already stored K/V rows. A deployment will normally key the cache by the whole model generation anyway; the finer dependency list explains why.

The pinned user-facing implementation computes selected V with CVMM, then concatenates only kv_cache["v"] and kv_cache["k"]. Its output selection is computed from q_src and applied after attention. That static code trace agrees with the operator above; it does not prove cache correctness for every runtime integration.

0x03J Pinned Training And User-Facing Code Contracts

Two official repositories serve different purposes:

Artifact	Pinned revision	What it establishes
moe_attention	7169ad370e68185ecddb592c623b8d550de725df	training implementation, K/Q/V/O ablation flags, separate/shared selection options, CVMM path
switchhead	0bb2f61d55bed44ce2e99b5164c9d84554a3b020	compact V/O implementation, RoPE/XL variants, SDPA path, and explicit K/V cache dictionary

In the compact implementation, get_sel performs this exact sequence:

float router logits
reshape to [token, H_A, E]
sigmoid
top-k under no_grad
gather unnormalized sigmoid weights
shift expert ids by head
prepare CVMM selection metadata

The training implementation exposes q_expert, k_expert, v_expert, o_expert, same_sel, and separate_kq_sel. Its create_selection_logic can remap Q to O selection and K to V selection, while same_sel can reuse a selection between source and destination roles. Those flags are evidence that the paper’s projection sweep is a real implementation family, not four names added after training.

0x03K Design Alternatives Are Different Contracts

Design	Projection choice	Map identity	Main tradeoff
naive destination head top-k	whole Q/K/V/O heads	selected expert head	source union can approach full pool
full projection MoE	source K/V, destination Q/O	$H_A$ mixed-projection maps	maximum projection capacity and routing work
final SwitchHead	source V, destination O; shared Q/K	$H_A$ maps	best source-reported parameter-matched ablation
O-only	destination O; shared Q/K/V	$H_A$ maps	near-best ablation with less source routing
shared selection	reuse source/destination choice	$H_A$ maps	fewer sort/top-k operations, source-reported minor quality loss
MoA	expert Q/O; shared K/V	selected expert maps	more map identities for comparable quality in source experiments

Shared selection is valid for autoregressive self-attention when a token’s selection is computed from that token state and used as destination now and as source after insertion. It saves selector and top-k work, but does not make source and destination tensor roles identical. Cross-attention must still state which sequence produced the shared choice.

0x03L Training And Source-Reported Evaluation Contract

The paper’s main experiments are parameter matched, not simply FLOP matched. Its search procedure first sets $H_AE$ equal to the dense baseline head count, starts with $H_A=2,k=2$, increases $k$ if quality is insufficient, then tries $H_A=4,k=2$, while changing $d_h$ to match total parameters. This procedure means small perplexity differences in the main tables can be imperfect matching artifacts; the stated target is baseline quality with maximum resource reduction.

Models were trained for 100k batches with batch size 64 and Adam at learning rate 0.00025; the larger configurations use 4k warmup steps. The main study uses Transformer-XL with one previous chunk of context, while an appendix tests RoPE without the XL cache. The training code adapts a Triton CVMM kernel from sigma-MoE.

The paper reports these end-to-end measurements, including MLPs, optimizer, and, for multi-GPU runs, gradient synchronization:

scale and hardware	dense	SwitchHead	source-reported relative result
47M, one RTX 3090	473 ms/iter, 20.5 GB	342 ms/iter, 13.5 GB	0.72 time, 0.65 memory
262M, eight V100	670 ms/iter, 20.5 GB/GPU	442 ms/iter, 12.5 GB/GPU	0.65 time, 0.61 memory

These are evidence for that training pipeline and hardware, not a serving latency prediction. The paper also reports its selected-expert Triton kernel at about 60% of an expert-sized cuBLAS GEMM and labels 80-90% as an estimate for a better kernel. It identifies model-parallel load balancing and dynamic expert movement as unfinished work. Those limitations are part of the result.

0x03M Failure Lab

Failure	Same-shape symptom	Detection	Repair
softmax selected gates substituted for sigmoid	output shape unchanged, gradients become competitive	selected gate sums forced to one; cross-gate VJP appears	preserve gathered sigmoid gates
V experts applied after attention	final [B,S_q,d] still matches	E or k enters the readout/map path	mix selected V before $A_hV_h$
O experts mixed before attention	final shape still matches	destination route changes source value state	apply O experts to $U_h$
E treated as map axis	model runs with inflated attention tensor	map tensor is [B,H_A,E,S_q,S_kv]	keep map identity on $H_A$ only
raw expert V cached	cache grows by E or k	cache V shape has an expert/slot axis	cache the mixed V row
source route recomputed after parameter update	old and new tokens use different V semantics	cache generation spans route/weight versions	invalidate or rebuild cache
paper MAC compared to fused profiler bytes	incompatible quantities look contradictory	formula includes no router or traffic terms	reconcile operation and materialization conventions
V/O ablation generalized as theorem	design choice appears universal	no parameter/data/kernel scope beside claim	retain the source-reported parameter-matched scope

0x03N Loop 4 Executable Closure

The checked-in pure-Python fixture uses B=1,S_q=2,S_kv=3,d=2,H_A=2,E=3,k_S=k_D=2,d_h=2. It executes:

independent source and destination sigmoid top-k routes
selected V projection and pre-attention mixture
one map per H_A with source-axis softmax
selected O projection and post-attention mixture
full shared-Q/K and selected-V/O/router backward
central finite differences for all six parameter banks
resident, active, named-MAC, and paper-MAC ledgers
map-count and decode-cache identities

Twenty-two one-field negative fixtures reject collapsed source/destination lengths, invalid active counts, softmax routing, normalized selected gates, differentiable top-k claims, shared-selection substitution, wrong map or softmax axes, misplaced V/O mixtures, expert-axis cache payloads, and malformed tensor banks. The oracle executes no GPU kernel, framework autograd, training run, latency benchmark, or paper evaluation.

0x03O MoA Versus SwitchHead

Property	MoA	SwitchHead
shared projections	K and V across attention experts	Q and K within each computed map head
routed projections	coupled Q and O	V and O in reported final design
route coordinate	destination token to attention expert	source token to V expert; destination token to O expert
attention-map identity	selected Q expert	fixed computed head $h\in H_A$
cache	shared K/V without attention-expert axis	head K plus source-routed mixed V
selected gate	paper softmax with stopped selected denominator	gathered, unnormalized sigmoid
source training control	balance, z-loss, dropout in MoA paper	source reports no extra balance regularizer for final SwitchHead
scale idea	enlarge selectable attention-head pool	reduce map count while retaining projection capacity

The distinction is not “both are attention MoEs.” It is whether expert identity survives into the attention-map coordinate. MoA keeps it there; SwitchHead removes it before the quadratic core and reintroduces output choice after the readout.

0x04 Mixture-Of-Depths: Route Token Identities Through A Block

Mixture-of-Depths changes the conditional axis again. MoA and SwitchHead ask which projection bank should process a token. Mixture-of-Depths asks whether a token-layer pair should execute an entire Transformer block at all.

That distinction changes the accounting contract. The selectable capacity is not a bank of resident parameters. It is a fixed number of token identities that may enter self-attention and the subsequent MLP at one routed layer. The paper fixes that number before training, so tensor extents and aggregate block work are predictable even though the selected positions vary with the input.

Loop 5 owner question. How does a routed layer score every valid token, select exactly $C$ identities per sequence, execute a causal block on a gathered $[B,C,d]$ tensor, weight selected deltas by the original router scalar, and scatter those deltas into an otherwise unchanged residual stream?

This loop owns the fixed-capacity training operator, its local gradient, and its MAC/movement derivation. Loop 6 owns the noncausal-training versus causal-decode decision. Loop 7 owns sparse per-layer K/V and replay state.

Source: Mixture-of-Depths, Figure 1. Purple cells execute a block; orange cells route around it. Unlike early exit, a token can skip an earlier block and execute a later one, so the route is a layer-by-token matrix rather than one stopping depth per token.

0x04A Capacity Is A Contract On Rows, Not An Observed Router Outcome

Let a routed layer receive:

$$X^l\in\mathbb{R}^{B\times S\times d}, \qquad V\in\{0,1\}^{B\times S},$$

where $V_{b,i}=1$ means row $i$ is a valid member of sequence $b$. The paper’s routing equation defines an integer capacity $C<S$ per batch element. For a fixed-length, unpadded training batch, every sequence therefore contributes exactly $C$ rows and the block input has shape:

$$\widetilde X^l\in\mathbb{R}^{B\times C\times d}.$$

The paper’s earlier compute-budget discussion also uses $T$ for tokens across sequence and batch. These statements are compatible only after the unit is named: total selected rows are $BC$, while the routing decision and the Section 3.4 capacity $C$ are per sequence. A global top-$BC$ across the whole minibatch is a different algorithm; one sequence could consume another sequence’s budget.

For padded or ragged data, an executable implementation needs an additional policy. It may:

1. bucket examples so each sequence exposes at least $C$ valid rows;
2. derive $C_b$ from each valid length and accept shape variation; or
3. pad selected rows while carrying a selected-row validity mask.

The paper’s static-graph claim is cleanest for fixed $S$ and fixed $C$. The Loop 5 oracle keeps the shape $[B,C,d]$ fixed while proving that a padded row with a very large numerical value is still excluded from selection.

At the limiting settings:

$$C=S \Rightarrow \text{dense block membership},$$ $$C=0 \Rightarrow \text{identity residual path at that routed layer}.$$

The local executable fixture requires $0<C<S$ because it tests the genuinely conditional case. Those endpoint identities remain useful conceptual checks.

0x04B Why This Is Expert-Choice With One Expensive Path

The paper reaches its fixed budget by comparing token-choice and expert-choice routing. The source figure makes the difference visible.

Source: Mixture-of-Depths, Figure 2, arXiv v1 source artifact. Token-choice first lets tokens request paths and then resolves capacity overflow. Expert-choice lets each path select a fixed top-$k$ set. MoD has one expensive self-attention-plus-MLP path; every non-selected identity takes its complementary residual path.

The three contracts are not interchangeable:

policy	decision	selected count	failure or complement
token-choice	each token chooses a path	input-dependent	overloaded paths need dropping, padding, or rerouting
multi-expert expert-choice	each expert chooses its top-$k$ tokens	fixed per expert	a token can enter several experts or none
single-path MoD expert-choice	one expensive block chooses top-$C$ tokens	exactly $C$ per sequence	every other valid token takes identity residual

With only one expensive path, selection and complement are mutually exclusive. That removes the ordinary multi-expert load-balancing problem: the block always receives $C$ rows. It does not make routing free, smooth, or causal.

The paper also trains a stochastic top-$C$ control by sampling router scores from a Gaussian. That control has the same selected shape and a different information policy. Fixed work therefore does not prove that learned routing is useful; it merely makes an equal-compute comparison possible.

0x04C Router Scalar, Executable Top-$C$, And The Percentile Limit

For sequence $b$, layer $l$, and row $i$, the paper uses one linear router:

$$r^l_{b,i} = \left(w^l_\theta\right)^\top x^l_{b,i}.$$

An implementation first excludes invalid rows:

$$\widehat r^l_{b,i} = \begin{cases} r^l_{b,i}, & V_{b,i}=1,\\ -\infty, & V_{b,i}=0. \end{cases}$$

Then it selects exactly $C$ row identities per sequence:

$$J^{\text{score}}_{b,:} = \operatorname{StableTopC} \left( \widehat r^l_{b,:}; \text{score descending, row id ascending} \right).$$

The paper writes the same idea using percentile $P_\beta(R^l)$ with $\beta=1-C/S$ and strict inequalities around the threshold. That notation explains the capacity fraction, but it does not by itself define exact membership when scores tie the percentile. An executable contract needs top-$C$ plus a tie policy. Backend-dependent ties would change token identities, gradients, and later per-layer state while preserving all tensor shapes.

Top-$C$ returns score order, not necessarily sequence order. Let $p_{b,i}$ be the original position id. Define a second permutation:

$$J^{\text{exec}}_{b,:} = \operatorname{SortByPosition} \left(J^{\text{score}}_{b,:},p_{b,:}\right).$$

The selected router scalars must follow the same permutation:

$$g^l_{b,u} = r^l_{b,J^{\text{exec}}_{b,u}}.$$

These are the raw linear scores from the paper’s main block path. They are not a softmax over selected rows, and Loop 5 does not replace them by a sigmoid. The sigmoid appears later when the paper constructs a causal membership classifier; that is Loop 6.

0x04D A Two-Sequence Identity Trace

The executable fixture uses $B=2$, $S=6$, $C=3$, and $d=4$. Its scores are:

sequence	valid score vector	top-$C$ score order	chronological execution order	selected positions
0	[.50,-.30,.20,.10,.30,.70]	[5,0,4]	[0,4,5]	[10,14,15]
1	[-.10,.60,.40,.20,1.00,PAD]	[4,1,2]	[1,2,4]	[21,22,24]

The padded row in sequence 1 contains [100,100,100,100]. It would dominate the router if validity were ignored. The oracle proves it is never selected.

Gather the chronological rows:

$$G^l_{b,u,:} = X^l_{b,J^{\text{exec}}_{b,u},:}, \qquad G^l\in\mathbb{R}^{B\times C\times d}.$$

For decoder self-attention, causal visibility is a relation over original positions, not an accident of gather order:

$$M^l_{b,u,v} = \begin{cases} 0, & p_{b,J^{\text{exec}}_{b,v}} \le p_{b,J^{\text{exec}}_{b,u}},\\ -\infty, & \text{otherwise}. \end{cases}$$

Because the fixture sorts selected rows chronologically, this relation becomes a $C\times C$ lower triangle. A kernel is also allowed to retain score order and construct the mask from position comparisons. Applying an ordinary lower triangle directly to score-ordered rows is wrong: it lets router rank redefine causal time.

0x04E Gather, Delta-Only Block, Gate, And Scatter

The paper’s selected update is:

$$x_{b,i}^{l+1} = \begin{cases} x_{b,i}^{l} + r^l_{b,i} f^l_i\!\left(\widetilde X^l\right), & i\in J_b,\\ x_{b,i}^{l}, & i\notin J_b. \end{cases}$$

Here $f^l$ comprises self-attention and the subsequent MLP on the selected rows. To make residual ownership executable, define a helper that returns only the block delta:

$$F^l = \operatorname{SelectedBlockDelta} \left(G^l,M^l\right) \in \mathbb{R}^{B\times C\times d}.$$

If a framework block returns a complete residual output $\operatorname{Block}(G)$, the adapter must construct $F=\operatorname{Block}(G)-G$. Passing a full residual output to the paper’s outer $X+rF$ equation would add the gathered identity twice.

Weight only selected deltas:

$$D^l_{b,u,:} = g^l_{b,u}F^l_{b,u,:}.$$

Scatter into a zero delta tensor using original row identities:

$$\Delta X^l = \operatorname{ScatterAdd} \left( J^{\text{exec}},D^l; \text{shape}=[B,S,d] \right),$$ $$X^{l+1}=X^l+\Delta X^l.$$

Each row appears at most once in this single-path top-$C$ contract, so indexed assignment and scatter-add have the same forward value. Scatter-add states the transpose relation needed by backward and remains correct if a later design composes several routed deltas.

An implementation trace is:

def mod_training_block(x, valid, position_ids, router, block_delta, C):
    # x: [B,S,d]; capacity is per sequence, not global over B*S.
    scores = einsum("bsd,d->bs", x, router)
    scores = where(valid, scores, -inf)

    score_order = stable_topk(
        scores,
        k=C,
        per_sequence=True,
        tie_break="row_index_ascending",
    ).indices                                      # [B,C]

    ids = sort_ids_by_position(score_order, position_ids)  # [B,C]
    selected_x = gather_rows(x, ids)                       # [B,C,d]
    selected_pos = gather_rows(position_ids, ids)          # [B,C]
    selected_scores = gather_rows(scores, ids)             # [B,C]

    causal = selected_pos[:, None, :] <= selected_pos[:, :, None]
    delta = block_delta(selected_x, causal)                 # [B,C,d]
    weighted_delta = selected_scores[..., None] * delta

    return x + scatter_add_rows(
        zeros_like(x),
        ids,
        weighted_delta,
    )

The code intentionally keeps four objects separate:

membership ids       [B,C]      discontinuous identity
execution positions  [B,C]      causal and positional identity
router scalars        [B,C]      differentiable selected gate
block deltas          [B,C,d]    selected computation

0x04F Backward: Membership Is Hard, The Selected Gate Still Learns

Let $\overline Y_{b,i,:}=\partial\mathcal L/\partial Y_{b,i,:}$. Inside a region where top-$C$ membership is unchanged, selected router scalar $g_{b,u}$ has derivative:

$$\overline g_{b,u} = \sum_{j=1}^{d} \overline Y_{b,J_{b,u},j} F_{b,u,j}.$$

Because $g_{b,u}=w_\theta^\top x_{b,J_{b,u}}$, the main language-model path gives the router vector:

$$\overline w_\theta = \sum_{b=1}^{B} \sum_{u=1}^{C} \overline g_{b,u} x_{b,J_{b,u}}.$$

For a selected input row, the router branch also contributes:

$$\left.\overline x_{b,J_{b,u}}\right|_{\text{router}} = \overline g_{b,u}w_\theta.$$

The block-transpose and outer identity paths contribute additional input gradient. An unselected row still receives its identity residual gradient, but it contributes zero router gradient from this main block path because its score neither gates a block delta nor enters a differentiable selected normalization.

Top-$C$ membership itself is piecewise constant:

$$\frac{\partial J}{\partial r}=0 \quad\text{almost everywhere, and undefined at route boundaries.}$$

Therefore, “the router is on the gradient path” does not mean gradient descent differentiates through token rank. It means the selected raw scalar multiplies the executed delta. In the fixture, the smallest selected-versus-unselected score margin is $0.1$, so central perturbations keep membership fixed. The analytic router-vector gradient agrees with finite differences to $3.2\times10^{-9}$ maximum absolute error.

0x04G Derive The Block MACs Instead Of Quoting One Ratio

Use one named-MAC convention first: one multiply-accumulate counts as one MAC. A FLOP ledger that counts multiply and add separately multiplies the matmul terms by approximately two. Let the decoder block use model width $d$, $H$ heads with $Hd_h=d$, and MLP width $d_{ff}$.

For a dense block at sequence length $S$:

$$M_{\text{QKVO,dense}} = 4BSd^2.$$

The score product and value reduction each cost $BS^2d$ MACs after summing over heads:

$$M_{\text{attn,dense}} = \underbrace{BS^2d}_{QK^\top} + \underbrace{BS^2d}_{AV} = 2BS^2d.$$

Let $p_{mlp}$ count model-to-hidden or hidden-to-model projections:

$$M_{\text{MLP,dense}} = p_{mlp}BSdd_{ff}, \qquad p_{mlp} = \begin{cases} 2,&\text{two-matrix GELU-style MLP},\\ 3,&\text{SwiGLU-style MLP}. \end{cases}$$

At a routed MoD block, the router still scores all $S$ rows:

$$M_{\text{router}} = BSd.$$

But the selected block’s linear projections see only $C$ rows:

$$M_{\text{QKVO,MoD}} = 4BCd^2,$$ $$M_{\text{MLP,MoD}} = p_{mlp}BCdd_{ff}.$$

Both attention axes are selected, so the quadratic products are $C\times C$:

$$M_{\text{attn,MoD}} = 2BC^2d.$$

Finally, weighting the selected delta performs $BCd$ scalar multiplies. Top-$C$ selection, gather/scatter, masks, softmax scalar work, normalization, residual adds, and launch overhead are not matmul MACs and must not be silently hidden inside these equations.

The named dense-block and routed-block matmul sums are therefore:

$$M_{\text{dense block}} = 4BSd^2 +2BS^2d +p_{mlp}BSdd_{ff},$$ $$M_{\text{MoD routed block}} = BSd +4BCd^2 +2BC^2d +p_{mlp}BCdd_{ff},$$

before the selected $BCd$ gate multiplies and named non-matmul work.

Set the capacity ratio:

$$\rho=\frac{C}{S}.$$

Then:

$$\frac{M_{\text{QKVO,MoD}}}{M_{\text{QKVO,dense}}} = \frac{M_{\text{MLP,MoD}}}{M_{\text{MLP,dense}}} = \rho,$$

while:

$$\frac{M_{\text{attn,MoD}}}{M_{\text{attn,dense}}} = \rho^2.$$

This is the precise origin of the familiar “linear MLP savings, quadratic attention savings” statement. It holds because both queries and keys/values come from the selected set. A variant with selected queries but dense historical K/V has $C\times S$ maps and ratio $\rho$, not $\rho^2$. That variant may be useful, but it is not the paper’s routed-block operator.

For $L_f$ full blocks and $L_r$ routed blocks, the network-level named MAC ledger is:

$$M_{\text{network}} = L_fM_{\text{dense block}} + L_r \left( M_{\text{MoD selected block}} +M_{\text{router}} +M_{\text{named non-matmul}} \right).$$

Consequently, a capacity ratio does not determine a whole-model speedup unless the routed-layer frequency is also specified. The paper’s source-reported best small-budget configuration routes every other block at $C/S=12.5\%$; that is an empirical configuration, not an architecture law.

0x04H Numerical MAC And Movement Reconciliation

The fixture uses $B=2$, $S=6$, $C=3$, $d=4$, $d_{ff}=6$, two attention heads, and a two-matrix MLP. Its exact named ledger is:

term	dense block	MoD routed block
Q/K/V/O projections	768 MAC	384 MAC
$QK^\top$ and $AV$	576 MAC	144 MAC
two-matrix MLP	576 MAC	288 MAC
router over all rows	0 in baseline	48 MAC
selected raw-gate multiply	0	24 multiplies
block matmul total	1,920 MAC	816 MAC

Here $\rho=0.5$. Linear terms halve, map products fall to one quarter, and the mixture of those terms gives a selected-block ratio of $816/1920=0.425$. Adding the router gives $(816+48)/1920=0.45$ before route movement and scalar overhead.

For two-byte activations and four-byte indices, a minimal logical route payload is:

state transition	logical payload
materialized router scores $[B,S]$	24 B
selected row ids $[B,C]$	24 B
gather read plus gathered write	96 B
scatter delta read plus selected output write	96 B

These are not measured HBM bytes. They omit cache-line amplification, top-$C$ workspace, allocator traffic, hidden copies needed by a concrete schedule, block-internal tensors, backward state, launches, and synchronization. The accounting firewall from M4 still applies.

0x04I Static Shapes Do Not Select A Kernel Or Guarantee Latency

Fixing $C$ gives useful compile-time extents:

scores                 [B,S]
selected row ids       [B,C]
selected positions     [B,C]
selected hidden        [B,C,d]
causal relation        [B,C,C]
attention maps         [B,H,C,C]
block delta            [B,C,d]
scatter destinations   [B,C]
output                 [B,S,d]

But the runtime still has to perform or fuse:

score -> mask padding -> stable top-C -> optional chronological sort
      -> gather -> selected block -> raw gate -> scatter-add

Several schedules satisfy the same logical operator:

schedule	advantage	cost or risk
explicit gather and dense $C\times C$ block	simple fixed shapes; reuses dense kernels	route movement and extra launches
score-order gather plus position-aware mask	avoids an explicit chronological reorder	arbitrary mask may be less kernel-friendly
chronological gather plus ordinary causal kernel	simple lower-triangular selected attention	sort and permutation metadata are required
fused select/gather or block/scatter	less materialized route traffic	custom backward, tails, and compiler constraints
dense block followed by output masking	trivial integration	does not save block compute and is not MoD execution

The paper reports that normalized FLOPs and step time are tightly correlated in its evaluated training setup. That source-reported result does not prove that every implementation realizes the same ratio. For small $C$, top-$C$, movement, launch latency, and underfilled kernels can dominate.

0x04J Alternative Contracts That Must Be Named

alternative	shape contract	semantic difference from paper MoD
global top-$BC$ over a minibatch	fixed total rows, variable rows per sequence	budgets compete across unrelated sequences
independent score threshold	variable selected count	causal online decision is possible, but exact training capacity is lost
random top-$C$	same fixed shape	equal-compute control without learned allocation
soft dense residual mixture	dense $[B,S,d]$ block	differentiable gates but no conditional block-compute saving
selected queries with dense K/V	$[B,H,C,S]$ maps	updates fewer queries but keeps all source rows visible
fixed position pattern	fixed identities and shape	deterministic sparse depth, not context-sensitive routing

The first four can look nearly identical in pseudocode if membership, materialized block domain, and residual weighting are omitted. Transformer System treats those as first-class contract fields.

0x04K Failure Lab And Executable Closure

failure	violated invariant	consequence
top-$C$ across flattened $B\times S$	capacity is per sequence	one example consumes another’s compute budget
padding score left finite	selection domain is valid rows	padding enters the block and a valid token is displaced
percentile with unspecified equality	exact $C$ membership needs a tie policy	backend-dependent row identity
score order passed to triangular causal mask	causal time is original position	future selected token becomes visible
selected weights softmax-normalized	paper uses raw router scalar	different forward and cross-token router gradients
sigmoid replaces raw block gate	decode classifier confused with model path	training operator changes
full residual output used as $f$	helper must return a delta	selected identity is added twice
skipped rows multiplied by router	complement is exact residual	skip path changes and extra router gradient appears
attention uses selected Q but all K/V	paper block domain is selected rows	attention ratio is $C/S$, not $(C/S)^2$
gather/scatter omitted from cost model	MAC and movement ledgers differ	theoretical saving is overclaimed as latency

Loop 5 adds a pure-Python oracle with an actual two-head causal attention operator, a subsequent MLP, raw selected routing, and exact scatter. It checks:

B=2, S=6, C=3, d=4, H=2, d_h=2, d_ff=6
13 named tensor/state shapes
per-sequence top-C and padding exclusion
score order versus chronological execution order
original-position causal relation
selected raw-gate update and exact identity skip
router-vector central finite differences
dense versus MoD MACs and logical route payloads
Twenty-two one-field invalid fixtures

The checked report is exact; the maximum router-gradient finite-difference error is $3.2\times10^{-9}$. It executes no GPU kernel, loads no Transformer checkpoint, measures no latency, and does not reproduce the paper’s benchmark. No author-maintained code artifact is assumed: the local oracle tests the stated paper operator and Transformer System implementation invariants.

0x04L Loop 6 Handoff: Fixed Top-$C$ Is A Training Teacher

Loop 5 freezes the training membership operator. It deliberately does not claim that sequence-global top-$C$ can be used during token-by-token generation. That decision depends on future router scores. Loop 6 will separate:

training membership target       top-C over the complete sequence
selected model-path gate         original raw router scalar
causal decode membership         router classifier or auxiliary predictor
capacity at training             exactly C rows per sequence
selected count during decoding   threshold-dependent and generally variable

Without that separation, an implementation can preserve every Loop 5 tensor shape in training and still leak future information or execute the wrong model at inference.

0x05 Training Membership And Causal Decode Are Different Operators

Loop 5 fixed the training operator: score every valid row in a complete sequence and let one expensive path choose exactly $C$ rows. Autoregressive decode cannot replay that choice online. The current token’s final rank depends on router scores that do not exist yet.

This section owns that mismatch. It does not yet own the sparse per-layer K/V ABI; Loop 7 starts only after the online membership decision is frozen.

The predictor answers a narrower question than the training router: should this current causal row execute the expensive path? It does not reconstruct the future-dependent rank that produced the teacher label.

0x05A A Concrete Noncausal Witness

Let $R_{b,l}=(r_{b,l,0},\ldots,r_{b,l,S-1})$ be the valid router logits for sequence $b$ and routed layer $l$. Training membership is:

$$z_{b,l,i} = \mathbb{1} \left[ i\in\operatorname{StableTopC}(R_{b,l},C) \right].$$

The local two-sequence fixture contains this witness:

batch 0, token 1 router logit = -0.14

prefix through token 1:
    token 1 is in prefix top-C because only two valid rows exist

complete sequence:
    later scores 0.81 and 0.79 occupy both selected slots
    token 1 is not in full-sequence top-C

No tie or implementation accident is involved. Future competitors change the rank. A prefix top-$C$ policy is therefore causal but is not the training teacher.

0x05B Freeze Three Different Variables

The paper’s sampling workaround is easy to misread because one scalar can play several roles. Keep these variables separate:

Variable	Meaning	Gradient and lifetime
$r_{b,l,i}$	original router logit	ranks teacher rows and gates the selected block delta
$z_{b,l,i}$	complete-sequence top-$C$ teacher label	hard stopped target; available during training only
$a_{b,l,i}$	optional auxiliary predictor logit	trained to imitate $z$ from a causal current-row input
$\widehat z_{b,l,t}$	online binary decision	threshold result used to execute or skip at decode

The hard target is derived from $r$ but is not differentiated through top-C. Away from a route boundary its VJP is zero; at a boundary the discrete membership derivative is undefined. The original selected-path gate remains $r$. A predictor probability is not silently substituted into the trained block equation.

0x05C Method 1: Put BCE Directly On The Router Logit

The first paper method uses router logits as binary-classification logits and the stopped top-$C$ labels as targets. For $N_{valid}$ valid rows:

$$\mathcal{L}_{direct} = \mathcal{L}_{LM} + \frac{\lambda_r}{N_{valid}} \sum_{b,i\in\mathcal V_b} \operatorname{BCEWithLogits}(r_{b,l,i},z_{b,l,i}).$$

The paper does not fully specify every experiment’s loss coefficient in the main text, so $\lambda_r$ is an implementation-state field, not a universal constant. The useful derivative is:

$$\frac{\partial\mathcal{L}_{direct}} {\partial r_{b,l,i}} = \frac{\lambda_r}{N_{valid}} \left(\sigma(r_{b,l,i})-z_{b,l,i}\right).$$

For a linear router $r=x^\top w_r+b_r$:

$$\frac{\partial\mathcal{L}_{direct}}{\partial w_r} = \sum_{b,i} \frac{\partial\mathcal{L}_{direct}}{\partial r_{b,l,i}} x_{b,l,i}, \qquad \left.\frac{\partial\mathcal{L}_{direct}}{\partial x_{b,l,i}} \right|_{aux} = \frac{\partial\mathcal{L}_{direct}}{\partial r_{b,l,i}}w_r.$$

So this method does more than calibrate a post-hoc threshold. Its auxiliary gradient changes the router and enters the main hidden representation through the router input path. The paper reports an approximately $0.2$-$0.3\%$ effect on the primary language-modeling objective for this method. That is a source-reported result, not a bound.

At decode, probability threshold $0.5$ is logit threshold zero:

$$\widehat z^{direct}_{b,l,t} = \mathbb{1}[\sigma(r_{b,l,t})>0.5] = \mathbb{1}[r_{b,l,t}>0].$$

0x05D Method 2: A Stopped-Input Auxiliary Predictor

The second method introduces a small auxiliary MLP receiving the router’s input through stop-gradient:

$$a_{b,l,i} = g_{\phi,l} \left(\operatorname{StopGrad}(x_{b,l,i})\right), \qquad \widehat p_{b,l,i}=\sigma(a_{b,l,i}).$$

Its objective is:

$$\mathcal{L}_{pred} = \frac{\lambda_p}{N_{valid}} \sum_{b,i\in\mathcal V_b} \operatorname{BCEWithLogits}(a_{b,l,i},z_{b,l,i}).$$

Now the backward ownership differs:

$$\frac{\partial\mathcal{L}_{pred}}{\partial a_{b,l,i}} = \frac{\lambda_p}{N_{valid}} \left(\sigma(a_{b,l,i})-z_{b,l,i}\right), \qquad \left.\frac{\partial\mathcal{L}_{pred}}{\partial x_{b,l,i}} \right|_{aux} =0.$$

The forward predictor is mathematically sensitive to its input payload, but autodiff does not return this auxiliary VJP into $x$ because the payload is detached. Predictor parameters $\phi$ still learn normally. The local oracle uses one linear logit only as an executable gradient surrogate; the paper says “small auxiliary MLP” and does not make that local teaching architecture a paper claim.

0x05E Thresholding Removes The Fixed Count

For either causal logit $q_{b,l,t}\in\{r_{b,l,t},a_{b,l,t}\}$ and probability threshold $\tau_l$:

$$\widehat z_{b,l,t}(\tau_l) = \mathbb{1} \left[ q_{b,l,t}>\operatorname{logit}(\tau_l) \right].$$

The number of executed rows becomes:

$$\widehat N_{b,l}(\tau_l) = \sum_t\widehat z_{b,l,t}(\tau_l),$$

which is generally not $C$. Training capacity is exact by construction; decode work is a threshold-dependent observation. A validation-calibrated threshold can target an expected rate on one distribution, but cannot guarantee exactly $C$ rows per generated sequence under distribution shift.

When a row executes, the trained selected path is still:

$$x_{b,l,t}^{l+1} = x_{b,l,t}^{l} + r_{b,l,t}\,f_l(\widetilde X_{b,l}), \qquad \widehat z_{b,l,t}=1.$$

The threshold logit answers whether to execute. The original router scalar answers how the selected block delta is weighted. Replacing $r$ by $\widehat p$ defines a different model.

0x05F Confusion Counts Are Also A Compute Ledger

Against retrospective teacher labels, define TP, FP, TN, and FN. Then:

$$C_{total}=TP+FN, \qquad \widehat N=TP+FP, \qquad \widehat N-C_{total}=FP-FN.$$

A false positive executes block work and applies a delta where the teacher skipped. A false negative skips a teacher-selected delta. Their compute costs and model effects are asymmetric; one accuracy number hides both.

The local fixture has $4$ positives among $15$ valid rows. An always-skip classifier already reaches:

$$\operatorname{accuracy}_{skip} = \frac{11}{15} \approx 73.3\%.$$

The auxiliary-predictor threshold sweep is:

$\tau$	TP	FP	TN	FN	precision	recall	accuracy	selected	drift from teacher
$0.50$	4	4	7	0	$0.50$	$1.00$	$0.733$	8	$+4$
$0.55$	4	1	10	0	$0.80$	$1.00$	$0.933$	5	$+1$
$0.60$	1	0	11	3	$1.00$	$0.25$	$0.800$	1	$-3$

One mismatch also changes later hidden states, so later predictor inputs are no longer the same as in a teacher-forced offline table. Route accuracy on static training rows is useful, but rollout quality and realized compute must be measured end to end.

0x05G What The Paper’s Evaluation Does And Does Not Establish

Source: Mixture-of-Depths, autoregressive evaluation figure from arXiv v1. The left plot compares source-reported held-out log-likelihood under training top-k and causal predictor routing; the right plot reports top-k prediction accuracy over training. The paper text says the auxiliary task quickly reaches $99\%$ accuracy, while the figure caption says upwards of $97\%$ soon into training. These are separate source surfaces, not a Transformer System reproduction.

The paper evaluates the same held-out data for each model, reporting $256{,}000$ sequences and $500$M tokens, little degradation when switching to predictor routing, and no significant step-speed impact for the auxiliary predictor in its setup. Those observations support feasibility for the tested models. They do not prove:

exact teacher-route parity;
exact C rows at decode;
preservation of an arbitrary output distribution;
the same confusion matrix across layers, prompts, or thresholds;
proportional latency from selected-count reduction;
a production sparse-KV cache contract.

This is approximate causalization of a noncausal training policy, not exact speculative decoding. M7 owns target-distribution-preserving draft/verify algorithms.

0x05H Alternative Causalization Contracts

Policy	What it changes	What it does not guarantee
direct router BCE	modifies router and hidden-state auxiliary gradients	exact $C$ or teacher parity at decode
stopped auxiliary predictor	adds predictor state without auxiliary VJP into the backbone	exact $C$ or zero rollout drift
per-layer calibrated threshold	matches expected rate on calibration data	per-sequence budget under shift
online hard cap	refuses later admissions after a causal budget is exhausted	teacher route; also needs a horizon/reservation policy
rolling-window top-$C$	compares within a bounded past window	complete-sequence teacher route without delay
execute both and mix softly	avoids hard membership	skipped-block compute savings
fixed schedule	gives exact predictable work	input-adaptive allocation

These alternatives should not share one benchmark row unless their teacher, causal inputs, threshold/cap policy, selected-count contract, and executed operator are all named.

0x05I Training And Decode Pseudocode

# Training: complete sequence is available.
router_logits = router(hidden)                       # [B,S]
teacher_ids = stable_top_c(router_logits, valid, C) # noncausal labels
teacher = membership_mask(teacher_ids).detach()

lm_loss = routed_model_loss(
    hidden,
    teacher_ids,
    selected_delta_gate=router_logits,
)

if mode == "direct_router_bce":
    predictor_loss = bce_with_logits(router_logits[valid], teacher[valid])
elif mode == "stopped_aux_predictor":
    predictor_logits = predictor(stop_gradient(hidden))
    predictor_loss = bce_with_logits(predictor_logits[valid], teacher[valid])

loss = lm_loss + lambda_predictor * predictor_loss

# Decode: only the current row and causal state are available.
router_logit = router(current_hidden)
decision_logit = (
    router_logit
    if mode == "direct_router_bce"
    else predictor(stop_gradient(current_hidden))
)

execute = decision_logit > logit(layer_threshold)
if execute:
    delta = run_selected_block(current_hidden, causal_layer_state)
    next_hidden = current_hidden + router_logit * delta
else:
    next_hidden = current_hidden

The threshold, predictor revision, valid-row policy, and routed-layer set are runtime state. Their K/V consequences are intentionally deferred.

0x05J Failure Lab And Executable Closure

Failure	Same-shape symptom	Consequence
prefix top-$C$ called the teacher	binary decision has expected shape	future competitors would change training membership
top-$C$ target differentiated as smooth	optimizer accepts a surrogate gradient	no longer the hard paper teacher
auxiliary predictor input not detached	predictor loss decreases normally	backbone receives an unintended auxiliary gradient
predictor probability used as block gate	output shape unchanged	executed model differs from the trained paper equation
$0.5$ threshold assumed to preserve $C$	decisions remain binary	selected count drifts by $FP-FN$
accuracy reported without prevalence	metric can look high	always-skip baseline may already dominate
offline route accuracy called rollout parity	static labels match often	early mismatches alter later states and decisions
source-reported $97$-$99\%$ called exact	rounded headline looks strong	no route, quality, or distribution identity follows
selected count called latency	fewer block rows are observed	launches, movement, K/V, and utilization remain unmeasured

mod_causal_predictor_oracle.py executes the complete-sequence teacher, a direct-router BCE path, a stopped-input predictor surrogate, smooth parameter VJPs, central finite differences, three thresholds, confusion/count ledgers, and twenty-seven one-field invalid fixtures. The checked report is exact; the maximum smooth-gradient error is below $3\times10^{-11}$. It does not load a checkpoint, run an autoregressive Transformer rollout, execute a GPU kernel, reproduce paper accuracy, or define sparse K/V storage.

0x05K Loop 7 Handoff: Membership Is Frozen, Cache Semantics Are Not

Loop 6 freezes these inputs to the next loop:

training membership: per-sequence complete-sequence stable top-C
decode membership: strict threshold on a causal router or predictor logit
selected count: variable at decode
selected delta gate: original raw router scalar
route provenance: predictor/policy revision plus threshold and routed layer

Loop 7 owns what happens after execute is known: per-layer K/V append, original positions, empty and skipped layers, prefix reuse, deterministic replay, branch state, rollback, and cache-key compatibility.

0x05L The Paper Consequence Is Not Yet A Runtime ABI

The Mixture-of-Depths paper makes one semantic consequence explicit: when a token skips self-attention at layer $l$, that token does not produce a key or value at that layer, so later selected queries at layer $l$ cannot attend to it. The discussion notes the possible inference-time K/V-cache benefit and also names decoupled query, key, and value routing as a possible extension. Those are model statements. The paper does not specify a production sparse cache layout, allocator, prefix index, transaction protocol, or replay ABI.

This section therefore has two evidence levels:

Contract	Status	What it establishes
skipped rows are absent from later same-layer attention	externally verified from Mixture-of-Depths, arXiv 2404.02258v1	the model’s layer-local visibility rule
one global token prefix with a different compact K/V row stream at every routed layer	rewrite derivation consuming M5.1 and M5.3	a sufficient logical representation of that rule
atomic present-or-absent receipts, cache compatibility fields, fork generations, and rollback	local executable contract	one auditable way to integrate sparse rows into serving
page tables, kernels, allocator fragmentation, prefix-index throughput, and latency	not measured here	implementation work, not a paper result

The distinction matters. “MoD may reduce K/V state” does not imply “an ordinary dense K/V cache can omit rows without changing any other runtime contract.” Sparse rows alter position lookup, cache length, prefix identity, fork state, rollback, and batching metadata at the same time.

0x05M One Global Prefix, Layer-Local Row Sets

Let $m_{l,t}\in\{0,1\}$ be the frozen online membership decision for global token index $t$ at routed layer $l$, and let $p_t$ be the token’s semantic position. The cache does not own a second model position. It stores the subsequence of original positions selected by that layer:

$$\mathcal P_l(T) = \left[p_t\;\middle|\;0\le t<T,\;m_{l,t}=1\right].$$

Its row count is therefore

$$n_l(T)=\lvert\mathcal P_l(T)\rvert = \sum_{t=0}^{T-1}m_{l,t},$$

not the global prefix length $T$. For one request, a compact logical layout is

$$K_l,V_l\in \mathbb R^{H_{kv}\times n_l(T)\times d_h}, \qquad P_l\in\mathbb N^{n_l(T)}.$$

The storage slot for a selected token is its exclusive rank among earlier selected rows:

$$j_l(t)=\sum_{u=0}^{t-1}m_{l,u} \quad\text{only when }m_{l,t}=1.$$

There is deliberately no $j_l(t)$ for a skipped row. Inventing a zero-filled slot would conflate model absence with a stored vector whose value happens to be zero.

The Loop 7 fixture uses four routed layers and six positions $[10,11,12,13,14,15]$:

Layer	Membership over six tokens	Stored semantic positions	Compact slots
$0$	[1,0,1,1,0,1]	[10,12,13,15]	[0,1,2,3]
$1$	[0,1,1,0,1,1]	[11,12,14,15]	[0,1,2,3]
$2$	[1,1,0,0,1,1]	[10,11,14,15]	[0,1,2,3]
$3$	[0,0,0,0,0,0]	[]	[]

Layer $3$ is not uninitialized and it does not lag behind the request. It has six closed route outcomes and a valid empty K/V stream. The request’s global committed frontier is six tokens even though the layer row counts are $[4,4,4,0]$.

Native Transformer System reconstruction. The row number in a compact layer buffer is an address coordinate; the recorded semantic position remains the model coordinate.

Production batches cannot generally materialize this ragged state as one rectangular $[B,H_{kv},T,d_h]$ tensor without padding. They need segmented offsets, pages, indirection, or another ragged representation. That physical choice belongs to M5.3. M6.8 owns only the logical requirement that every K/V row remains paired with its layer, request generation, and original position.

0x05N Attention Uses Semantic Positions, Not Compact Slots

Compaction changes addresses; it must not silently renumber the model. For a selected token $t$, projection and positional encoding use $p_t$:

$$q_{l,t}=\operatorname{RoPE}_{p_t}(W_l^Qx_{l,t}), \qquad k_{l,t}=\operatorname{RoPE}_{p_t}(W_l^Kx_{l,t}), \qquad v_{l,t}=W_l^Vx_{l,t}.$$

Define the visible selected-token indices for that query as

$$\mathcal A_l(t)=\{u\mid 0\le u\le t,\;m_{l,u}=1\}.$$

Then the attention row is

$$a_{l,t} = \operatorname{softmax}_{u\in\mathcal A_l(t)} \left( \frac{q_{l,t}k_{l,u}^{\mathsf T}}{\sqrt{d_h}} \right), \qquad o_{l,t}=\sum_{u\in\mathcal A_l(t)}a_{l,t,u}v_{l,u}.$$

The use of $u\le t$ includes the current selected row. A decode implementation may stage the new K/V, read the prior cache plus the staged row, then publish the row only after the complete token transaction succeeds. “Not yet published” and “not visible to its own query” are different statements.

For the final fixture token at semantic position $15$:

layer 0 executes -> visible positions [10,12,13,15]
layer 1 executes -> visible positions [11,12,14,15]
layer 2 executes -> visible positions [10,11,14,15]
layer 3 skips    -> no query, no K/V row, residual identity only

Using compact slots [0,1,2,3] as RoPE positions would make all three executed layers represent the last token as position $3$ instead of $15$ and would also compress every historical distance differently. The shapes remain valid while the model changes.

This is also why the route mask cannot be discarded after packing. A segmented offset says where rows live; semantic positions say what those rows mean; the route trace says why a row exists or is absent.

0x05O A Token Closes As Present Or Provably Absent At Every Routed Layer

M5.1 advances a request only after its serving state is internally coherent. With layer-local sparse streams, “coherent” cannot mean that every layer appended one row. Instead, each routed layer must publish exactly one outcome for global token $t$:

$$o_{l,t}\in \{\texttt{PRESENT\_KV}(p_t,j_l(t)),\; \texttt{ABSENT\_BY\_ROUTE}(p_t)\}.$$

Let $c$ be the global committed-token frontier. Token $t=c$ can advance it to $c+1$ only when all $L_r$ routed layers have staged one outcome:

$$\operatorname{publish}(t) \iff \sum_{l=0}^{L_r-1} \mathbf 1[o_{l,t}\text{ is complete}] =L_r.$$

This is not a distributed-consensus claim. It is a request-local atomicity contract: a failure after layer $1$ must not leave layers $0$ and $1$ at token $t$ while later layers and the global prefix remain at $t-1$.

def stage_mod_token(snapshot, token_id, semantic_position, hidden):
    assert snapshot.token_count == token_id.index
    txn = snapshot.begin_token(token_id, semantic_position)

    for layer in routed_layers:
        decision = route_decision(layer, hidden[layer])

        if not decision.execute:
            txn.stage_absence(
                layer=layer,
                position=semantic_position,
                reason="ABSENT_BY_ROUTE",
            )
            hidden[layer + 1] = hidden[layer]  # residual identity only
            continue

        update, new_k, new_v = layer.block_delta.decode(
            hidden[layer],
            position_id=semantic_position,
            sparse_kv=snapshot.layer_cache(layer),
        )
        txn.stage_kv(
            layer=layer,
            position=semantic_position,
            key=new_k,
            value=new_v,
        )
        hidden[layer + 1] = hidden[layer] + decision.raw_router * update

    # Atomic publication checks one outcome per routed layer.
    return txn.commit_all_layers()

An explicit absence receipt serves three purposes. It proves that the layer was evaluated under the pinned route policy, distinguishes a skipped row from missing work after a fault, and lets replay validate a complete token frontier without fabricating K/V payload.

Native Transformer System state-machine trace. Payload rows and absence receipts are different objects, but both are required to close one global token position.

0x05P Prefill And Decode Must Declare One Route Policy

There are now three superficially similar prefill computations:

Mode	Membership rule	Fixed count?	Valid cache consumer
paper training/evaluation	complete-sequence stable top-$C$	yes, per sequence and routed layer	the same training-model contract
serving prefill	pinned causal router or predictor at each token/layer state	no	later decode under the identical serving route policy
one-token decode	the same pinned causal decision at the frontier	no	continuation of that exact serving prefix

The first row cannot seed the second merely because both return binary masks. Its decision for an early token depends on later competitors. The cache would then encode a route history that online decode could not have produced.

A causal predictor may still be evaluated in a batched prefill kernel. “Causal” does not mean “must launch once per token.” It means the computed decision for $(l,t)$ consumes only state available at that causal point. Batched and stepwise evaluation are route-equivalent only under a pinned numerical and tie/comparator contract; changing precision, threshold equality, or fused operator order can flip a boundary decision and therefore every downstream state.

The safe prefill handoff is:

same model revision
same predictor and route-policy revisions
same strict comparator and threshold coordinates
same original-position rule
same selected-delta gate
same cache ABI
same deterministic numerical contract, or a frozen accepted route trace

0x05Q Prefix Identity Includes The Routing Program

Ordinary token equality is necessary but no longer sufficient for exact prefix reuse. Two requests with the same token ids can produce different layer-local K/V streams if their predictor, threshold, comparator, side input, transformed weights, or route policy differs.

One sufficient compatibility manifest is

$$I_{\text{prefix}}= \bigl( I_{\text{model}},I_{\text{weights}},I_{\text{predictor}},I_{\text{policy}}, \tau,\bowtie,I_{\text{position}},I_{\text{cache ABI}},I_{\text{tokens}}, I_{\text{side input}},I_{\text{isolation}},I_{\text{route trace}} \bigr).$$

Here $\bowtie$ names the comparator, for example strict >. The route-trace digest binds all token-major membership decisions already committed. It does not replace model, predictor, or policy revision fields: the same bit pattern produced by an incompatible program is not automatically a reusable execution artifact.

def sparse_prefix_compatible(producer, consumer):
    required = (
        "model_revision",
        "weight_transform_revision",
        "predictor_revision",
        "route_policy_revision",
        "logit_threshold",
        "comparator",
        "position_contract",
        "cache_abi_sha256",
        "token_prefix_sha256",
        "side_input_sha256",
        "isolation_namespace",
        "route_trace_sha256",
    )
    mismatches = [name for name in required
                  if producer[name] != consumer[name]]
    return len(mismatches) == 0, mismatches

The executable fixture accepts an exact consumer and rejects three controlled counterexamples with mismatch fields logit_threshold, predictor_revision, and route_trace_sha256. This is stricter than checking that both caches have the same number of rows.

0x05R Fork, Replay, Rollback, And Generation Fences

A prefix fork shares an immutable predecessor. It does not let parent and child append into the same mutable tail. At fork token count $T_f$, both branches must agree on

$$\left( x_{<T_f},p_{<T_f}, \{\mathcal P_l(T_f),K_l(T_f),V_l(T_f)\}_{l=0}^{L_r-1}, m_{:,<T_f} \right).$$

After the fork, route decisions may diverge. The fixture shares five tokens, then gives the parent final membership [1,1,1,0] and the child [0,1,1,0]. Their shared-prefix semantic digest is equal before the append; their successor snapshot digests differ afterward.

Three replay modes must remain distinct:

Replay mode	Recomputed state	Required evidence	Exactness limit
deterministic model replay	hidden states, route logits, decisions, and K/V	pinned model, policy, numerical contract, tokens, positions, and side inputs	exact only if those operations reproduce the same boundary decisions
frozen-decision replay	model state and K/V, while committed membership bits are supplied	immutable route trace plus compatible model/cache ABI	preserves the recorded path; it does not prove the current predictor would choose it
cache-only reuse	no earlier model work	complete compatibility manifest and immutable cache generation	valid only for the exact producer identity

On fault, discard every staged K/V row and absence receipt for the token and retain the predecessor generation digest. On rollback to an earlier token, truncate each layer to its own selected-row count at that frontier, not to the global token count. On cancellation, never publish a partially closed token. On retry, use a generation fence so a late completion from the abandoned transaction cannot attach to the new branch.

The fixture injects a failure at global token index $4$ after layer $1$. Its staged outcomes are observable for audit, publication is false, and the predecessor snapshot digest remains byte-for-byte unchanged.

0x05S Payload, Metadata, And Read Ledgers

For $H_{kv}$ K/V heads, head dimension $d_h$, and $b$ bytes per element, one selected layer-token stores

$$B_{\text{row}}=2H_{kv}d_hb$$

payload bytes. Across $L_r$ routed layers and $T$ tokens, dense and sparse logical payloads are

$$B_{\text{dense}}=L_rT\,B_{\text{row}}, \qquad B_{\text{sparse}} = B_{\text{row}} \sum_{l=0}^{L_r-1}\sum_{t=0}^{T-1}m_{l,t}.$$

The payload ratio is therefore the mean layer-token membership rate:

$$\frac{B_{\text{sparse}}}{B_{\text{dense}}} = \frac{1}{L_rT} \sum_{l,t}m_{l,t}.$$

That equality covers payload only. If every selected row stores a $b_p$-byte semantic position, a packed route bitmap uses one bit per layer-token, and an offset table uses $(L_r+1)b_o$ bytes, one illustrative named ledger is

$$B_{\text{named}} = B_{\text{sparse}} +b_p\sum_{l,t}m_{l,t} +\left\lceil\frac{L_rT}{8}\right\rceil +(L_r+1)b_o.$$

For the fixture, $H_{kv}=2$, $d_h=2$, $b=2$, $L_r=4$, $T=6$, and twelve of twenty-four layer-token pairs execute:

Ledger	Bytes
one K/V pair row	$16$
dense payload, $24$ rows	$384$
sparse payload, $12$ rows	$192$
semantic-position metadata, $12\times2$	$24$
route bitmap, $\lceil24/8\rceil$	$3$
five 4-byte layer offsets	$20$
sparse payload plus these named metadata	$239$

The $0.5$ payload ratio is exact for this fixture. It is not a measured HBM ratio or latency speedup. Alignment, page tables, block descriptors, reference counts, padding, staging buffers, replication, and transport are excluded.

Reads are layer- and query-dependent too. A selected query at $(l,t)$ reads $|\mathcal A_l(t)|$ historical/current rows; a skipped query performs no attention read at that layer. Fewer logical rows do not guarantee proportionate latency because indirection, short segments, page fragmentation, batch heterogeneity, and reduced kernel occupancy can dominate.

0x05T Alternative Cache Contracts Change Either Storage Or The Model

Design	Later selected query can see a skipped token?	Storage consequence	Model relation
sparse layer-local K/V, owned here	no	only selected rows plus metadata	matches the paper’s stated visibility consequence
zero placeholder plus exact mask	no	dense row count remains; compressibility depends on runtime	semantically possible if the mask is never omitted
dense shadow K/V for all tokens	yes	dense payload and ordinary positions	changes the MoD attention graph
selected queries but independently routed/dense K/V memory	policy-dependent	may trade query work against memory	a different architecture; related to the paper’s proposed decoupling direction
K/V shared across layers	policy-dependent	can reduce resident cache by architectural reuse	changes projection and layer-ownership contracts
static layer-token schedule	fixed by deployment artifact	can pre-plan allocation and kernels	removes input-dependent online membership

The alternatives are valuable precisely because they expose separate design axes. Query execution, K/V production, K/V retention, and K/V visibility need not share one gate. But once they are decoupled, each gate needs its own training objective, position semantics, exactness claim, and cache manifest. Calling all of them “MoD cache optimization” hides those changes.

0x05U Sparse-State Failure Lab And Executable Closure

Failure	Same-shape symptom	Consequence
compact slot used as RoPE position	every packed tensor is rectangular	positional geometry changes after the first skip
skipped row stored as unmasked zeros	cache length matches $T$	zero key/value can still receive probability mass or alter normalization
no absence receipt	layer row count simply does not increase	skip is indistinguishable from incomplete work after a fault
per-layer immediate publication	early layer append succeeds	request exposes a torn global token frontier
training top-$C$ cache reused by predictor decode	membership masks have equal shape	online continuation inherits a noncausal route history
token hash checked without route identity	lexical prefix matches	layer-local K/V streams may encode another predictor or threshold
route trace replaces raw model-path gate	replay follows the same binary path	selected block delta can still use the wrong scalar
rollback truncates every layer to token count	API call succeeds for dense cache	sparse layers overrun, fabricate, or delete the wrong rows
empty layer treated as uninitialized	no K/V payload exists	valid all-skip history is rejected or recomputed
logical payload reduction called measured speedup	byte equation is correct	allocator, kernels, batching, and reads remain unmeasured

mod_sparse_kv_oracle.py executes the layer-local position streams, compact slot maps, last-token visibility, all-layer transaction closure, injected rollback, immutable-prefix fork, compatibility rejection, and payload/metadata ledger. It rejects thirty-two one-field mutations. It does not load a model checkpoint, execute a paper or official sparse-cache implementation, run a GPU kernel, instantiate a production allocator or prefix index, or measure memory or latency.

0x05V Loop 8 Handoff: The Depth Path Now Has Stateful Semantics

Loop 7 freezes the state consumed by composition:

membership decision      pinned causal policy from Loop 6
selected model-path gate original raw router scalar
per-layer K/V             selected rows only, with original positions
skipped layer outcome     explicit ABSENT_BY_ROUTE receipt
token publication         atomic across all routed-layer outcomes
prefix identity           tokens + model + predictor + policy + route trace + cache ABI
fork/replay/rollback      immutable predecessor plus generation fence

Loop 8 may now place FFN expert routing inside or beside this depth decision. It must account separately for no-attention paths, attention-plus-MoE paths, integrated no-op experts, expert communication, and which route state enters the cache identity. It may not reopen the sparse-position or transaction contracts established here.

0x06 MoD And MoE Compose In Two Different Ways

The Mixture-of-Depths paper asks whether depth sparsity and expert sparsity can coexist in one block. This is not merely a matter of placing two routers in the same diagram. The order of the decisions determines which tokens execute attention, which tokens create K/V state, which population reaches the expert router, which route can communicate across devices, and what a residual path actually means.

Loop 8 owner question. Given the fixed-capacity depth route from Loops 5—7 and the expert route owned by M6.1, how do staged and integrated MoDE define different tensors, gradients, compute ledgers, cache state, and failure semantics even when both expose experts and a no-op path?

This section consumes the logical expert contract from M6.1, the balance contract from M6.2, and the depth-state transaction from 0x05L—0x05U. It does not rederive generic expert capacity, all-to-all kernels, or sparse-cache position identity.

Source: Mixture-of-Depths, Figure 7. The staged variant first routes around or through the full block, then selects an FFN expert. The integrated variant places a no-op option beside FFN experts after attention. These variants therefore do not save the same work even when their normalized training FLOPs are similar.

0x06A What Figure 7 Establishes, And What It Leaves Open

The source establishes three architecture-level claims:

1. staged MoDE applies MoD machinery before self-attention and MoE machinery afterward, so a depth-skipped token can avoid both attention and expert FFN work;
2. integrated MoDE uses one routing operation over conventional FFN experts and one or more no-op operations after attention, simplifying the routing structure but retaining dense attention;
3. a learned no-op route performed better in the source experiment than reducing conventional expert capacity and treating capacity-dropped tokens as residual routes.

The paper does not publish a complete executable contract for Figure 7. In particular, Section 4.3 does not pin the expert top-k, number of no-op choices, normalization axis, expert-capacity factor, overflow ordering, balance loss, collective implementation, or autoregressive cache ABI. The equations below therefore have two levels:

level	status	use
staged before attention; integrated no-op after attention	externally verified	architecture boundary from Figure 7 and Section 4.3
tensor shapes, named residual ownership, MAC/byte ledgers	mechanism-derived	consequences of that boundary plus prior M6 contracts
top-1 teaching fixture, selected-gate normalization, tie rules	Transformer System executable instantiation	makes one auditable operator; not claimed as unreleased paper code
latency, quality, and production collective performance	unmeasured here	requires checkpoint- and runtime-specific evaluation

This firewall matters because several plausible implementations share the same picture and produce different gradients.

0x06B Freeze The Two Routing Relations Before Composing Them

Let the layer input be:

$$X^l\in\mathbb{R}^{B\times S\times d}.$$

The depth route from Loop 5 scores all valid rows and selects exactly $C_D$ rows per sequence during fixed-capacity training:

$$r^D_{b,i} = \left(w_D\right)^\top x^l_{b,i}, \qquad J^D_{b,:} = \operatorname{StableTopC_D}(r^D_{b,:}),$$ $$G_{b,u,:} = X^l_{b,J^D_{b,u},:} \in \mathbb{R}^{B\times C_D\times d}.$$

The integer relation $J^D$ answers which token-layer pairs execute the block. The selected raw scalar:

$$g^D_{b,u}=r^D_{b,J^D_{b,u}}$$

answers how the selected block delta is weighted. Membership and gate are not interchangeable.

The expert relation has a different coordinate. For $N$ candidate token rows, an $E$-expert, top-$k_E$ token-choice router requests:

$$Z^E\in\mathbb{R}^{N\times E}, \qquad I^E\in\mathbb{N}^{N\times k_E}, \qquad A^E\in\mathbb{R}^{N\times k_E}.$$

$I^E$ stores expert identities and $A^E$ stores selected combine weights. Expert capacity can later accept, reject, reroute, or execute those requested edges. Loop 8 keeps the requested relation separate from the executed relation exactly as M6.1 requires.

The two axes are therefore:

depth relation   token position x layer -> execute block or preserve identity
expert relation  admitted token row x expert -> execute selected FFN delta

Staged MoDE nests the expert relation inside the selected depth rows. Integrated MoDE removes the depth relation and extends the expert path set with a no-op.

0x06C Staged MoDE Is A Nested Conditional Operator

The staged path is:

token state
    -> depth router over all valid rows
        -> unselected: preserve block input exactly
        -> selected: gather C_D rows in semantic position order
            -> selected-row self-attention
            -> expert router over those C_D post-attention rows
            -> dispatch selected expert edges
            -> combine expert FFN deltas
            -> form one selected block delta
            -> multiply by the original depth-router scalar
            -> scatter to the original token identities

0x06C.1 Attention State Precedes Expert State

Let selected-row attention return a delta:

$$A^D = \operatorname{SelectedAttentionDelta} \left(G,P^D,M^D\right) \in \mathbb{R}^{B\times C_D\times d},$$

where $P^D$ contains original semantic positions and $M^D$ is the causal relation derived in Loop 5. The residual state seen by the expert router is:

$$U=G+A^D.$$

Flatten only the batch and selected-row axes for expert routing:

$$U^{flat} \in \mathbb{R}^{(BC_D)\times d}.$$

For requested expert edges $I^E_{n,s}$ and selected gates $A^E_{n,s}$, define the MoE FFN delta:

$$M^E_n = \sum_{s=1}^{k_E} A^E_{n,s} \operatorname{Expert}_{I^E_{n,s}} \left(\operatorname{Norm}(U_n)\right).$$

This sum is over executed edges after the named capacity policy. If a route is dropped, rerouted, or sent to a fallback, that policy changes the sum and must not be hidden inside $I^E$.

The complete selected block delta is:

$$F^D=A^D+M^E.$$

Finally, staged MoDE preserves the paper’s outer MoD gate:

$$Y = X^l + \operatorname{ScatterAdd} \left( J^D, g^D\odot F^D; [B,S,d] \right).$$

This equation owns each residual exactly once. A^D is an attention delta, M^E is an expert-FFN delta, and unselected rows receive zero scattered delta. If a library helper returns U + M^E rather than M^E, using it in $A^D+M^E$ would add $G$ a second time.

0x06C.2 The Expert Router Sees A Selected Population

The staged expert router does not see $BS$ rows. It sees $BC_D$ rows whose identity was already chosen by the depth router:

$$N_{stage}=BC_D=BS\rho_D, \qquad \rho_D=\frac{C_D}{S}.$$

Before overflow, the number of requested expert edges is:

$$R_{stage}^{req}=BC_Dk_E.$$

Fixed depth capacity therefore bounds expert-router input shape, but it does not balance experts. If all $BC_D$ rows request one expert, that expert still becomes a hotspot. Any expert-load fraction or auxiliary objective in staged MoDE must use the selected population as its denominator:

$$f^{stage}_e = \frac{1}{BC_Dk_E} \sum_{n=1}^{BC_D} \sum_{s=1}^{k_E} \mathbf{1}\!\left[I^E_{n,s}=e\right].$$

Dividing by $BSk_E$ would make all expert loads appear smaller by $\rho_D$ and would mix depth sparsity with expert imbalance.

0x06C.3 Tensor And Code Trace

def staged_mode_block(
    x, valid, position_ids,
    depth_router, depth_capacity, selected_attention_delta,
    expert_router, experts, expert_policy,
):
    # x: [B,S,d]
    depth_score = where(valid, depth_router(x), -inf)       # [B,S]
    depth_ids = stable_topc_per_sequence(
        depth_score,
        depth_capacity,
        tie_break="row_index_ascending",
    )                                                       # [B,C_D]
    depth_ids = sort_ids_by_position(depth_ids, position_ids)

    selected_x = gather_rows(x, depth_ids)                  # [B,C_D,d]
    selected_pos = gather_rows(position_ids, depth_ids)     # [B,C_D]
    depth_gate = gather_rows(depth_score, depth_ids)         # [B,C_D]

    attention_delta = selected_attention_delta(
        selected_x,
        causal_from_semantic_positions(selected_pos),
    )                                                       # [B,C_D,d]
    post_attention = selected_x + attention_delta

    flat = post_attention.reshape(B * C_D, d)
    requested = expert_router.request(flat)                 # ids/gates [B*C_D,k_E]
    executed = expert_policy.admit(requested)                # capacity is explicit
    expert_delta = dispatch_run_combine(flat, executed, experts)
    expert_delta = expert_delta.reshape(B, C_D, d)

    selected_block_delta = attention_delta + expert_delta
    weighted = depth_gate[..., None] * selected_block_delta
    return x + scatter_add_rows(zeros_like(x), depth_ids, weighted)

The trace contains two gathers with different purposes. The depth gather changes the self-attention row set. Expert dispatch then permutes already-selected rows by expert owner. Collapsing them into one undocumented permutation loses both causal position identity and inverse-combine metadata.

0x06C.4 Backward Has Two Piecewise-Smooth Boundaries

Hold both integer route relations fixed. Let $\overline Y^D_{b,u,:}$ be the upstream gradient gathered at selected depth rows. The original depth gate receives the composed block delta:

$$\overline g^D_{b,u} = \left\langle \overline Y^D_{b,u,:}, A^D_{b,u,:}+M^E_{b,u,:} \right\rangle.$$

The gradient entering attention and the expert sublayer is scaled by the outer depth gate:

$$\overline F^D_{b,u,:} = g^D_{b,u}\overline Y^D_{b,u,:}.$$

For one executed expert gate $a_{n,s}$:

$$\overline a_{n,s} = \left\langle \overline F^D_n, \operatorname{Expert}_{I^E_{n,s}} \left(\operatorname{Norm}(U_n)\right) \right\rangle.$$

If the selected expert gates are a softmax over $k_E$ selected logits, their local Jacobian is:

$$\overline z^E_{n,s} = a_{n,s} \left( \overline a_{n,s} - \sum_{q=1}^{k_E}a_{n,q}\overline a_{n,q} \right).$$

Rows skipped by the depth route receive the outer identity gradient but no task-path expert-router or expert-weight gradient at this layer. The depth top $C_D$ relation and the expert top $k_E$ relation are each piecewise constant; gradient descent differentiates the active graph, not the integer identities. There are consequently two route margins to monitor, not one.

0x06D Integrated MoDE Keeps Attention Dense And Routes The FFN Delta

Integrated MoDE first runs ordinary attention for every valid token:

$$A^{dense} = \operatorname{DenseAttentionDelta}(X^l,P,M), \qquad U=X^l+A^{dense}.$$

It then extends the expert path set with $N_\varnothing$ no-op choices:

$$\mathcal P = \{1,\ldots,E\} \cup \{\varnothing_1,\ldots,\varnothing_{N_\varnothing}\}.$$

A no-op FFN is a zero delta:

$$\operatorname{Expert}_{\varnothing_j}(u)=0.$$

Therefore, for one selected path $p_{b,i}\in\mathcal P$ and selected gate $a_{b,i}$:

$$Y_{b,i} = U_{b,i} + a_{b,i} \operatorname{Expert}_{p_{b,i}} \left(\operatorname{Norm}(U_{b,i})\right).$$

If $p_{b,i}=\varnothing_j$, the result is $Y_{b,i}=U_{b,i}$: attention remains, the FFN delta is absent, and the outer FFN residual is preserved once.

0x06D.1 Zero Delta Is Not Identity Output

Two library interfaces can describe a residual block:

delta interface:       y = u + f(u)       no-op returns 0
full-output interface: y = F(u)           no-op returns u

Mixing them is a correctness bug. Returning u from a no-op expert and then executing y = u + expert(u) produces 2u. Returning zero from a full-output mixture erases the residual. The executable contract in this article uses the delta interface.

A learned adapter is also not a no-op. If $E_\varnothing(u)=B A u$, the path owns parameters, MACs, optimizer state, and a nonzero function. It may be a useful light expert, but it changes Figure 7’s residual-path semantics.

0x06D.2 One Router Does Not Mean One Unambiguous Gradient

The paper states that integrated MoDE uses one routing operation, but does not publish its complete gate normalization. To expose the design choice, consider the Loop 8 teaching instantiation:

1. compute logits $z\in\mathbb{R}^{E+1}$ over $E$ real experts and one no-op;
2. compute $\pi=\operatorname{softmax}(z)$ over all paths;
3. select deterministic top-1 path $p=\arg\max_j z_j$;
4. use $a=\pi_p$ as the selected real-expert gate; a selected no-op contributes zero delta.

For a selected real expert $p$, let:

$$q_p = \left\langle \overline Y, \operatorname{Expert}_p(\operatorname{Norm}(U)) \right\rangle.$$

Within a fixed top-1 region:

$$\overline z_j = q_p\pi_p \left( \mathbf{1}[j=p]-\pi_j \right).$$

The no-op logit can therefore receive task gradient on a row that executes a real expert because it participates in the softmax denominator. On a row whose selected path is the zero-delta no-op, $q_p=0$ and this main-path gradient is zero. Such rows can still be affected indirectly through updates learned on other rows, auxiliary objectives, or route-boundary changes.

If instead top-1 gates are renormalized only over the selected set, then $a=1$. The main task path has no smooth derivative with respect to any routing logit inside that fixed membership region. A straight-through estimator, auxiliary objective, noisy gate, multi-path mixture, or another training rule would be required. These are different algorithms, not notation changes.

0x06D.3 Integrated Code Trace

def integrated_mode_block(
    x, position_ids, dense_attention_delta,
    combined_router, real_experts,
):
    attention_delta = dense_attention_delta(
        x,
        dense_causal_mask(position_ids),
    )                                                        # [B,S,d]
    post_attention = x + attention_delta

    logits = combined_router(norm(post_attention))           # [B,S,E+1]
    probabilities = softmax(logits, dim=-1)                  # teaching contract
    path_id = stable_argmax(logits, tie_break="path_id_ascending")
    selected_gate = gather(probabilities, path_id)

    expert_delta = zeros_like(post_attention)
    real = path_id < E
    expert_delta[real] = dispatch_run_combine_top1(
        post_attention[real],
        path_id[real],
        selected_gate[real],
        real_experts,
    )
    # path_id == E is the local zero-delta no-op and is never dispatched.
    return post_attention + expert_delta

The no-op identity belongs in route metadata but not in an expert send buffer. Encoding it as a fake remote expert would preserve output shape while wasting packing, communication, and launch work.

0x06E Three Residual-Looking Outcomes Have Different Causes

Dropping a route because an expert is over capacity means the token requested real computation and did not receive it. Selecting a no-op means the model’s router explicitly chose the zero-delta FFN path. Skipping a staged depth block means attention and FFN were both bypassed.

outcome	model request	attention executed	K/V row at this routed layer	FFN delta	training meaning
staged depth skip	skip whole block	no	absent by route	none	depth decision
integrated no-op	choose no-op after attention	yes	present	zero by definition	learned path choice
real expert overflow	choose real expert, then reject edge	yes unless outer depth also skipped	follows outer design	missing, fallback, or rerouted	execution-capacity policy
ordinary real expert	choose and execute expert	follows outer design	follows outer design	nonzero expert delta	selected expert path

The first two can both expose an identity-like residual in part of the block, but they are not interchangeable. Capacity dropping cannot be relabeled as a learned no-op after the fact: the requested route, gate gradient, balance statistics, and failure diagnosis are different.

0x06F Training Objectives Must Name Which Population They Control

Staged composition can carry at least three training signals:

$$\mathcal L_{stage} = \mathcal L_{LM} + \lambda_D\mathcal L_{depth\_predictor} + \lambda_E\mathcal L_{expert\_control}.$$

The depth predictor is owned by Loop 6. Expert balance is owned by M6.1/M6.2, but staged composition changes its candidate population from $BS$ rows to $BC_D$ selected rows. A route histogram should report both:

depth admission rate          C_D / S during fixed-capacity training
expert requested load         edges per expert / (B*C_D*k_E)
expert executed load          admitted edges per expert / executed edge count
compound path frequency       token -> depth selected -> expert e

Integrated composition has a more subtle control problem. If a balance loss treats the no-op as one ordinary expert among $E+1$ paths, equal balance targets a no-op fraction near $1/(E+1)$, not an independently chosen compute budget. If the no-op is excluded from balance, the real experts may balance while the total real-expert rate remains uncontrolled. Three explicit alternatives are:

integrated control	what it fixes	tradeoff
balance all $E+1$ paths	path frequencies	couples no-op rate to expert count
balance real experts only + compute-rate penalty	real-expert fairness and mean no-op rate separately	adds a second controller and coefficient
fixed path capacities / expert-choice	bounded rows per path	selection population may become noncausal or group-dependent
token-choice + runtime capacity	per-token preference first	overflow remains distinct from no-op choice

The Mixture-of-Depths paper does not specify which of these was used for every Figure 7 point. The table is a design analysis, not an attribution.

0x06G Parameter, MAC, And Communication Ledgers

Use one MAC convention: a multiply-accumulate counts as one MAC. Let every real expert have width $d_{ff}$ and $p_{mlp}$ projections, with $p_{mlp}=2$ for a two-matrix MLP and $3$ for SwiGLU. Ignore biases first.

0x06G.1 Stored Parameters

Let one expert own $P_E=p_{mlp}dd_{ff}$ parameters. Staged MoDE stores:

$$P_{stage} = EP_E + \underbrace{d}_{\text{depth router}} + \underbrace{Ed}_{\text{expert router}}.$$

Integrated MoDE with $N_\varnothing$ parameter-free no-op paths stores:

$$P_{integrated} = EP_E + (E+N_\varnothing)d.$$

With one no-op and the same expert bank, these named parameter totals are equal:

$$P_{stage}=P_{integrated} \quad\text{when}\quad N_\varnothing=1.$$

Equal parameter count does not imply equal active work. Staged composition has two decisions at different points; integrated composition has one larger decision after dense attention.

0x06G.2 Staged MACs

Before expert overflow, staged MoDE has:

$$M_{stage} = \underbrace{BSd}_{\text{depth router}} + \underbrace{4BC_Dd^2}_{\text{selected QKVO}} + \underbrace{2BC_D^2d}_{\text{selected }QK^\top\text{ and }AV} + \underbrace{BC_DEd}_{\text{expert router}} + \underbrace{p_{mlp}BC_Dk_Edd_{ff}}_{\text{requested real expert FFNs}}.$$

This excludes top-k comparisons, normalization scalar work, gate multiplies, gather/scatter, packing, communication, and overflow. If each token executes one full-width expert, $k_E=1$ recovers the selected dense-MLP term from Loop 5 while increasing stored capacity by $E$.

Relative to an otherwise identical dense-attention MoE, the selected QKVO, expert-router, and expert-FFN terms scale with $\rho_D=C_D/S$, while the attention products scale with $\rho_D^2$.

0x06G.3 Integrated MACs

Let $R_{real}$ be the number of executed real-expert edges after no-op choices and capacity policy. Integrated MoDE has:

$$M_{integrated} = \underbrace{4BSd^2}_{\text{dense QKVO}} + \underbrace{2BS^2d}_{\text{dense attention products}} + \underbrace{BS(E+N_\varnothing)d}_{\text{combined router}} + \underbrace{p_{mlp}R_{real}dd_{ff}}_{\text{real expert FFNs}}.$$

$R_{real}$ is semantic work, not automatically a static tensor extent. A fixed-capacity implementation may reserve or pad more expert rows than it executes. Conversely, a dropless implementation can execute all requested real edges with variable expert loads. The architecture diagram alone does not determine represented work or latency.

0x06G.4 Logical Movement And Expert-Parallel Payload

With activation element width $b_a$ bytes, staged expert routing requests:

$$B_{dispatch,stage}^{req} = BC_Dk_Edb_a,$$

and a simple remote dispatch-plus-return ledger is:

$$B_{EP,stage}^{roundtrip} \approx 2BC_Dk_Edb_a.$$

An integrated top-1 design dispatches only rows selecting real experts:

$$B_{EP,integrated}^{roundtrip} \approx 2R_{real}db_a.$$

These are logical activation payloads. They exclude route indices, gates, padding, headers, topology amplification, collective protocol, allocator traffic, and expert weights. M6.6 owns the physical mapping. Staged composition can reduce bytes by routing fewer rows, but it can also create smaller expert batches and poorer GEMM occupancy. Fewer bytes is not a latency proof.

0x06H Prefill, Decode, K/V, Prefix, And Replay Contracts

The two designs remain different after training:

state question	staged MoDE	integrated MoDE
training depth membership	fixed top-$C_D$ teacher	none
autoregressive depth decision	Loop 6 causal predictor or declared threshold policy	none
attention row set	depth-selected rows	all valid rows
per-layer K/V	sparse, with present/absent receipt	dense
expert-router population	only depth-selected post-attention rows	every post-attention row
no-op meaning	whole block skipped by depth relation	FFN zero delta after attention
expert communication	only after depth admission	only real-expert paths; no-op stays local

For staged decode, a token transaction follows this order:

predict depth membership
    skipped  -> stage ABSENT_BY_ROUTE for this layer; do not run expert router
    selected -> compute/stage K/V with semantic position
                run attention
                request/admit/execute expert routes
                stage hidden-state result
commit token only after every routed layer closes

Integrated decode stages a normal K/V row before its combined FFN route. A no-op selection cannot remove the attention K/V retrospectively.

Loop 7’s prefix identity must now bind the semantic composition policy:

staged:
    depth predictor/version + threshold or capacity policy
    expert router/version + top-k + normalization + capacity/overflow policy
    route trace needed by the chosen replay mode

integrated:
    combined router/version + no-op path count/ids
    gate normalization + top-k + capacity/overflow policy

Physical rank placement is not automatically model semantics, but changing expert parameters, route arbitration, or numerical execution can change hidden states and all later K/V rows. A prefix cache must reject incompatible route programs even if token ids and nominal architecture names match.

0x06I Alternatives Change The Operator, Not Just The Schedule

alternative	attention	FFN choice	important consequence
staged MoDE	selected depth rows	experts on selected rows	maximum logical skip scope; two route boundaries
integrated MoDE	dense	real expert or no-op	simpler route graph; dense K/V
dense attention + separate FFN skip router	dense	binary skip, then expert route	two post-attention routers; not Figure 7 integrated form
reduced expert capacity + drop	dense	requested expert may be rejected	overflow masquerades as residual if provenance is lost
light learned expert	design-dependent	cheap nonzero function	parameterized low-compute path, not a no-op
top-$k_E>1$ including no-op	dense	no-op can share a selected set with real experts	no-op gate mass can attenuate real expert contributions
selected queries + dense historical K/V	sparse query updates	design-dependent	attention maps are $C_D\times S$, not paper MoD’s $C_D\times C_D$

The last two deserve particular care. If a no-op and a real expert are both in a selected top-k set, the zero output may still consume normalized gate mass. Removing it and renormalizing real experts changes the output. Likewise, retaining dense historical K/V changes both the attention operator and Loop 7’s sparse-state contract.

0x06J A Four-Token Teaching Trace

The Loop 8 fixture uses one sequence, four rows, width two, depth capacity two, two real experts, one no-op, and top-1 expert paths.

stage	staged MoDE	integrated MoDE
candidate rows	positions [10,11,12,13]	positions [10,11,12,13]
depth result	rows [1,3] execute; [0,2] preserve input	no depth decision
attention domain	selected positions [11,13]	all four positions
expert candidates	two post-attention rows	four post-attention rows
no-op availability	none inside the selected expert bank	combined path id 2
K/V presence	rows [1,3] present; [0,2] absent	all rows present

The point is not the small values. It is that the same phrase, “this token took a residual path,” can mean no block state, attention-only state, or a rejected expert request. The fixture records these as different enums and rejects any mutation that collapses them.

0x06K Failure Lab And Executable Closure

substitution	shape-correct?	semantic failure
run expert router on all $BS$ rows in staged MoDE	yes	skipped rows receive FFN work and expert credit
feed pre-attention $G$ to the staged expert router	yes	changes Figure 7’s attention-then-MoE block order
multiply only $A^D$ by the depth gate	yes	expert delta escapes the outer depth decision
return full residual output from an expert-delta helper	yes	selected identity is added twice
return u from no-op under an outer residual	yes	integrated no-op doubles the post-attention state
remove attention K/V after integrated no-op	yes	no-op is incorrectly promoted to whole-block skip
dispatch no-op rows to a remote fake expert	yes	output may match while runtime accounting is false
count staged expert load over $BS$	yes	depth sparsity is mistaken for expert balance
label capacity overflow as no-op selection	yes	requested and executed routes become indistinguishable
claim integrated work is static from $E$ and $N_\varnothing$ alone	yes	real route count, capacity representation, and padding are omitted
reuse a prefix across different no-op ids or overflow policy	yes	later hidden states and K/V can disagree

The executable evidence is intentionally a teaching operator, not an unreleased Figure 7 reproduction:

python3 docs/numinens-evidence/m6-8-conditional-compute/mode_composition_oracle.py --write-report --fail-on-error
python3 -m unittest docs/numinens-evidence/m6-8-conditional-compute/test_loop8_mode_composition.py

It executes staged and integrated forward paths, route identities, zero-delta no-op semantics, nested-route gradients inside fixed membership regions, parameter/MAC/byte ledgers, K/V presence, and prefix-policy rejection. It also rejects one-field substitutions for residual ownership, route population, normalization, overflow provenance, cache state, and source scope. It does not load the paper checkpoint, reproduce Figure 7 loss points, run an expert collective, or measure latency.

0x06L Loop 9 Handoff

Loop 8 freezes a composition vocabulary that later modality routing must reuse:

depth skip            no attention, no K/V row, no FFN route
integrated no-op       dense attention and K/V, zero FFN delta
expert overflow        requested real route not executed as requested
staged route relation  depth membership followed by expert membership
integrated relation    one post-attention path over experts plus no-op
route state            selection, gate, admission, provenance, cache compatibility

Loop 9 may now introduce modality-conditioned expert populations and staged training transitions. It may not call every unexecuted FFN a no-op, reuse depth absence for an attention-present token, or hide modality-specific capacity inside generic expert load.

0x07 Multimodal Sparse Experts Need A Training Transition

MoE-LLaVA does not begin with a multimodal sparse model. It first aligns visual features to a language model, then trains a dense multimodal model, and only then converts selected dense FFNs into routed experts. That order is the mechanism. If the sparse router, expert copies, multimodal projector, and instruction objective are all introduced at once, a failed run cannot say whether it failed to align modalities, learn the task, preserve the dense checkpoint, admit routes, or balance experts.

Loop 9 owner question. How does one mixed visual/text sequence move through the three MoE-LLaVA stages, which parameters and loss populations exist at each stage, and under what exact admission and normalization conditions does cloning the Stage-II FFN preserve its function?

This section consumes the image-to-token and multimodal-sequence foundations from the X series. It consumes generic top-$k$, capacity, dispatch, and combine semantics from M6.1 and physical expert placement from M6.6. Its local owner is the training-state transition and the fact that modality is a population coordinate, not an input label consumed by the router.

Source: MoE-LLaVA v5, framework figure. Stage I aligns the visual projector while the language model remains frozen. Stage II trains a dense multimodal path. Stage III copies dense FFN weights into experts and introduces the router. The paper figure and the pinned public training code do not expose exactly the same Stage-III trainable set; 0x07C keeps them separate.

0x07A Three Evidence Surfaces, Not One Composite Story

The reconstruction pins three independent sources:

evidence surface	what it can establish	what it cannot silently establish
MoE-LLaVA arXiv v5	architecture equations, three-stage intent, paper objective, source-reported training recipe, capacity ablation, routing figures	exact behavior of a dependency version not printed in the equations
official repository commit 6cb5f66	input construction, parameter flags, layer replacement, loss accumulation, generation path, stage scripts	that every prose claim in the paper is implemented by this commit
DeepSpeed 0.9.5	actual pinned top-2 gate, capacity locations, selected-gate renormalization, auxiliary-loss code, dispatch/combine	that its cumulative row-order admission is the paper’s claimed BPR policy

The public project pins DeepSpeed 0.9.5 in pyproject.toml. That matters: modern MoE descriptions cannot be substituted for this route. In particular, the pinned top-2 implementation samples its second expert with a Gumbel-max choice, computes capacity locations in flattened row order, and renormalizes the surviving selected gates. The paper appendix, by contrast, says it uses Batch Priority Routing. The article will state both facts rather than invent a single implementation that satisfies both.

0x07B Build One Packed Sequence, Then Freeze Three Masks

For an image $I\in\mathbb{R}^{H\times W\times3}$, the vision encoder emits $P$ patch rows of width $C_v$:

$$Z=E_{vision}(I)\in\mathbb{R}^{P\times C_v}, \qquad P=\frac{HW}{14^2}$$

for the paper’s patch-14 encoder convention. A learned projector maps those rows into the language-model width $d$:

$$V=f_{proj}(Z)\in\mathbb{R}^{P\times d}.$$

Text token IDs become rows in the same width:

$$T=E_{text}(x_{1:N})\in\mathbb{R}^{N\times d}.$$

The paper writes a visual-then-text concatenation. The pinned code implements the more general operation required by prompts with image sentinels: split the text around each sentinel, replace the sentinel by projected image rows, then pad the resulting per-example sequences.

$$X_0^{(b)} = \operatorname{Interleave} \left(T^{(b)},V_1^{(b)},\ldots,V_J^{(b)}\right) \in\mathbb{R}^{S_b\times d}.$$

After padding to $S_{max}$, freeze three masks:

$$M_{valid}\in\{0,1\}^{B\times S_{max}}, \qquad M_{AR}\in\{0,1\}^{B\times S_{max}}, \qquad M_{mod}\in\{image,text,pad\}^{B\times S_{max}}.$$

• $M_{valid}$ determines which rows are semantic sequence state rather than rectangular padding;
• $M_{AR}$ marks shifted labels that are not IGNORE_INDEX and therefore enter the direct cross-entropy reduction;
• $M_{mod}$ is analysis metadata reconstructed from sequence assembly. It is not concatenated to the hidden state or passed to the gate.

The pinned input builder always writes IGNORE_INDEX for inserted image rows. Prompt text may also be ignored by the conversation-labeling procedure. It is therefore too broad to say “the text tokens receive language loss.” The exact direct support is

$$\mathcal{R}_{AR} = \left\{(b,s):M_{AR}[b,s]=1\right\},$$

At the paper/operator level, the natural content population is

$$\mathcal{R}_{content} = \left\{(b,s):M_{valid}[b,s]=1\right\}.$$

The pinned Phi path has a more surprising implementation contract. The decoder calls self.mlp(hidden_states) without passing an attention/used-token mask; DeepSpeed 0.9.5 then reshapes the complete rectangular tensor to [-1,d]. Its actual router population is therefore

$$\mathcal{R}_{route}^{code} = \left\{(b,s):0\le b<B,\ 0\le s<S_{max}\right\}, \qquad |\mathcal{R}_{route}^{code}|=BS_{max},$$

including materialized padding rows. M_valid still controls attention and semantic sequence validity, but it is not supplied to this pinned top-2 gate. An implementation that compacts valid rows or passes used_token would define a different route/capacity contract. Padding must consequently be a third analysis category rather than being silently counted as text.

These sets are not gradient-isolation boundaries. A later target-token loss can backpropagate through attention into earlier visual and prompt states. The claim is only that direct CE labels and direct router/auxiliary populations are different.

def interleave_multimodal_rows(text_ids, text_labels, image_features):
    pieces_x, pieces_y = [], []
    for text_span, image_rows in split_at_image_sentinels(
        text_ids, text_labels, image_features
    ):
        pieces_x.append(embed(text_span.ids))
        pieces_y.append(text_span.labels)
        if image_rows is not None:
            pieces_x.append(project(image_rows))
            pieces_y.append(full(image_rows.shape[0], IGNORE_INDEX))

    x = concat(pieces_x, axis=0)      # [S_b, d]
    labels = concat(pieces_y, axis=0) # [S_b]
    return x, labels

0x07C The Three Stages Are Checkpoint And Optimizer Boundaries

The word stage must name more than a dataset. It fixes the graph, checkpoint input, trainable set, optimizer state, and checkpoint output.

stage	checkpoint input	graph	paper-level trainable state	checkpoint output
I: alignment	base LLM + vision encoder	dense LLM, projected image rows	visual projector only	aligned projector / pretrain checkpoint
II: dense instruction tuning	base model + Stage-I projector	dense multimodal LLM	dense model except vision encoder	dense LVLM checkpoint
III: MoE tuning	Stage-II dense LVLM	selected FFNs replaced by routed expert copies	paper figure: experts and router	sparse multimodal checkpoint

0x07C.1 Stage I: Learn A Coordinate Bridge

The vision encoder and language model remain frozen. Gradients pass through the frozen language model to update the projector, but no optimizer state is allocated for the frozen weights. For projector parameters $\theta_P$:

$$\theta_P^{I+1} = \theta_P^I-\eta_I\nabla_{\theta_P}\mathcal{L}_{AR}, \qquad \Delta\theta_{vision}=\Delta\theta_{LLM}=0.$$

The official Phi script sets --tune_mm_mlp_adapter True, uses one epoch, BF16, gradient checkpointing, and a $10^{-3}$ learning rate. Those are source-reported settings, not a universal alignment recipe.

0x07C.2 Stage II: Solve Multimodal Instruction Tuning While Dense

The dense FFNs remain ordinary FFNs. The official stage-II script starts from the Stage-I projector, leaves the backbone unfrozen, freezes the vision tower, and tunes the multimodal model for one epoch with a $2\times10^{-5}$ learning rate. The output checkpoint is valuable for two independent reasons:

1. it already implements the multimodal task before routing is introduced;
2. every Stage-III expert can be initialized from a known dense FFN rather than from an unrelated random function.

0x07C.3 Stage III: Paper Ownership And Pinned-Code Ownership Diverge

The paper framework marks the copied experts and router trainable while the vision encoder, projector, embeddings, attention, and surrounding block state are locked. The pinned Phi script passes --train_modules fc1 fc2 wg, and initialize_moe_modules first freezes names outside that set. However, initialize_vision_modules runs later and explicitly sets every existing projector parameter back to requires_grad=True.

Therefore the effective pinned-code trainable set is:

copied expert fc1/fc2 parameters
DeepSpeed gate parameter wg
visual projector parameters re-enabled later

not merely “the MoE layers.” This is a code-order consequence, not a reinterpretation of the paper figure. The paper appendix also says Stage III uses ZeRO-2 offload, while the pinned Phi script references zero2.json rather than zero2_offload.json. Both artifacts are retained as a revision-level discrepancy.

0x07D Dense-To-Sparse Conversion: Derive The Continuity Condition

At a converted layer $\ell$, let the Stage-II dense FFN be $F_\ell$ with parameter state $\theta_\ell$. The pinned code stores that state dict, replaces the module by a DeepSpeed MoE, and asserts that every copied expert has exactly the same tensors:

$$\theta_{\ell,e}^{III,0} = \theta_\ell^{II}, \qquad e=1,\ldots,E.$$

The paper writes a gate matrix in $d\times E$ mathematical orientation. The pinned torch.nn.Linear(d,E,bias=False) stores it as $W_g\in\mathbb{R}^{E\times d}$. For each materialized hidden row $h_t$ in the pinned rectangular batch:

$$z_t=W_gh_t\in\mathbb{R}^{E}, \qquad p_{t,e}=\frac{\exp z_{t,e}}{\sum_{j=1}^{E}\exp z_{t,j}}.$$

DeepSpeed 0.9.5 chooses

$$e_t^{(1)}=\operatorname*{argmax}_{e}p_{t,e},$$

then samples a fixed replay realization of Gumbel noise $\gamma_{t,e}$ for the second choice and excludes the first expert:

$$e_t^{(2)} = \operatorname*{argmax}_{e\ne e_t^{(1)}} \left(z_{t,e}+\gamma_{t,e}\right).$$

That is not deterministic “take the two largest softmax values.” Tests and route replays must pin the random state or store the selected edges.

For $T=|\mathcal{R}_{route}^{code}|=BS_{max}$ flattened materialized rows, top-2 capacity in the pinned implementation is

$$C = \max\left( C_{min}, \left\lceil \frac{T}{E}\cdot 2c_f \right\rceil \right).$$

Let $a_t^{(j)}\in\{0,1\}$ indicate that selected route $j$ fits the expert’s capacity buffer. The combine weights are the surviving probabilities renormalized over the admitted selected set:

$$\widehat g_t^{(j)} = \frac{a_t^{(j)}p_{t,e_t^{(j)}}} {\sum_{q=1}^{2}a_t^{(q)}p_{t,e_t^{(q)}}+\varepsilon_t},$$

where $\varepsilon_t$ is only the implementation’s divide-by-zero guard. The MoE delta is

$$F_{MoE}(h_t) = \sum_{j=1}^{2} \widehat g_t^{(j)}F_{\ell,e_t^{(j)}}(h_t).$$

Now the warm-start claim can be proved rather than asserted. If all experts are exact clones and at least one selected route survives, then

$$F_{MoE}(h_t) = \left(\sum_j\widehat g_t^{(j)}\right)F_\ell(h_t) = F_\ell(h_t).$$

If both routes are capacity-dropped, the selected sum is zero and the MoE delta is zero. The surrounding residual may keep the block numerically finite, but the Stage-II FFN function is not preserved for that row. Clone equality alone is therefore insufficient; admission and combine normalization are part of the continuity contract.

def convert_dense_ffn_to_pinned_moe(layer, E, top_k, capacity_factor):
    dense_state = copy_state_dict(layer.mlp)
    layer.mlp = DeepSpeedMoE_0_9_5(
        expert=layer.mlp,
        num_experts=E,
        k=top_k,
        capacity_factor=capacity_factor,
    )
    for expert in layer.mlp.experts:
        assert state_dict_equal(expert, dense_state)
    return layer

0x07E The Language Loss And Router Loss Observe Different Populations

With logits $o_{b,s}$ and labels $y_{b,s}$, the direct autoregressive loss is

$$\mathcal{L}_{AR} = -\sum_{(b,s)\in\mathcal{R}_{AR}} \log p_\theta\left(y_{b,s}\mid X_{b,<s}\right).$$

The pinned model shifts logits and labels, applies cross entropy, collects one DeepSpeed auxiliary loss from each routed layer, scales their sum, and adds it to the language loss:

$$\mathcal{L}_{total} = \mathcal{L}_{AR} + \alpha\sum_{\ell\in\mathcal{L}_{MoE}} \mathcal{L}_{aux}^{(\ell)}.$$

The paper’s printed $\mathcal{F}$ and $\mathcal{G}$ equations have inconsistent summation indices. The prose and pinned DeepSpeed code support the following Switch-style reconstruction for the pinned route population $T=BS_{max}$:

$$f_e = \frac{1}{T}\sum_{t=1}^{T} \mathbf{1}\!\left[e_t^{(1)}=e\right], \qquad \bar p_e = \frac{1}{T}\sum_{t=1}^{T}p_{t,e},$$ $$\mathcal{L}_{aux} = E\sum_{e=1}^{E}f_e\bar p_e.$$

Three qualifications follow from the pinned implementation:

1. $f_e$ counts pre-capacity top-1 assignments, not executed top-2 route edges, and in this pinned Phi path includes padding rows;
2. $\bar p_e$ averages the full softmax over all experts;
3. the loss does not itself measure modality composition, bytes per rank, or post-capacity expert work.

In the official Phi stage-III script, $\alpha=0.01$. Visual rows have no direct CE label, but they enter the routed hidden population and hence the auxiliary loss. They also affect later text predictions through attention, so they can receive indirect task gradients. Padding rows have neither direct CE support nor semantic modality, yet the pinned rectangular MoE call still includes them in gate probabilities, top-1 counts, capacity, and auxiliary loss. That is an implementation fact, not a recommended masking policy.

0x07F The Paper Says BPR; The Pinned Top-2 Code Uses Row Order

The paper appendix says capacity overflow is resolved with Batch Priority Routing. BPR, introduced for V-MoE, ranks candidate tokens for an expert by router priority before retaining its capacity. The pinned MoE-LLaVA code does not define such a sort. It delegates to DeepSpeed 0.9.5, whose top-2 path uses cumulative locations over the flattened row order.

contract	first expert	second expert	admission priority	combine after drop
paper prose	top-$k$ route	top-$k$ route	BPR claimed	not pinned by prose
DeepSpeed 0.9.5 top-2	softmax argmax	Gumbel-max excluding first	cumulative flattened-row location	renormalize surviving selected gates
Loop 9 oracle	exact pinned equations with stored noise	exact pinned equations with stored noise	batch-then-row replay	exact pinned renormalization

The source-reported capacity ablation compares $c_f=1.0$ and $1.5$ and reports higher average scores for $1.5$ in its tested model/configuration rows. It does not prove that larger capacity is monotonically better, that no token is dropped, or that the public commit implements the named BPR policy.

Alternatives must be named because they change which modality loses work:

flattened row-order admission
    sensitive to packing and image/text order

BPR by route probability
    favors high-confidence rows; requires a priority sort

random token selection
    changes variance and replay state

no dropping / dynamic capacity
    preserves routes but changes memory and communication bounds

0x07G Parameters, Training State, And Logical Traffic Need Separate Ledgers

For one converted dense FFN with $P_F$ parameters, $E$ experts, $k$ active experts per token, and a bias-free $d\to E$ gate:

$$P_{Stage\ II}^{FFN}=P_F,$$ $$P_{Stage\ III}^{resident}=EP_F+dE,$$ $$P_{Stage\ III}^{active/token}=kP_F+dE.$$

The gate is evaluated for every routed row even though only $k$ expert FFNs run. For $L_M$ converted layers, the paper-figure Stage-III trainable state is

$$P_{train,paper}^{III} = L_M(EP_F+dE),$$

while the pinned Phi code also has trainable projector state $P_{proj}$:

$$P_{train,code}^{III} = L_M(EP_F+dE)+P_{proj}.$$

Optimizer bytes depend on dtype, master-weight policy, ZeRO partitioning, and offload. They cannot be inferred from the trainable count alone. The paper’s offload statement and script’s zero2.json reference must therefore remain separate release manifests.

With $R_{exec}$ admitted expert edges, hidden width $d$, and activation width $b_a$ bytes, a local logical payload ledger is

$$B_{dispatch}=R_{exec}db_a, \qquad B_{roundtrip}\approx2R_{exec}db_a.$$

This excludes route metadata, padding to capacity, topology amplification, collective protocol, expert weights, and allocator traffic. M6.6 owns those physical costs.

The source-reported common training settings are one epoch per stage, $336\times336$ inputs, CLIP-Large/336, a two-linear GeLU projector, BF16, gradient checkpointing, and eight A800-80GB GPUs. Stage I reports learning rate $10^{-3}$ and batch 256; Stages II and III report $2\times10^{-5}$ and batch 128. The pinned shell scripts are the executable configuration surface and must win over a paraphrased table when reproducing a run.

0x07H Expert Balance And Modality Composition Are Orthogonal Projections

Let $m(t)\in\{image,text,pad\}$ be metadata for routed row $t$. Define requested and executed modality-by-expert counts. Analyses of semantic modality should report the padding row separately or exclude it explicitly:

$$C_{m,e}^{req} = \sum_t\sum_{j=1}^{k} \mathbf{1}[m(t)=m]\mathbf{1}[e_t^{(j)}=e],$$ $$C_{m,e}^{exec} = \sum_t\sum_{j=1}^{k} \mathbf{1}[m(t)=m]\mathbf{1}[e_t^{(j)}=e]a_t^{(j)}.$$

Counts alone lose gate magnitude, so also record executed gate mass:

$$G_{m,e}^{exec} = \sum_t\sum_{j=1}^{k} \mathbf{1}[m(t)=m]\mathbf{1}[e_t^{(j)}=e] a_t^{(j)}\widehat g_t^{(j)}.$$

Now choose the denominator that answers the question:

$$P(e\mid m) = \frac{C_{m,e}}{\sum_j C_{m,j}} \quad\text{asks where a modality routes,}$$ $$P(m\mid e) = \frac{C_{m,e}}{\sum_n C_{n,e}} \quad\text{asks what an expert receives.}$$

Two contingency matrices can both have expert totals $[2,2]$:

$$C^{specialized} = \begin{bmatrix}2&0\\0&2\end{bmatrix}, \qquad C^{mixed} = \begin{bmatrix}1&1\\1&1\end{bmatrix}.$$

An expert-load metric sees them as identical. Their modality conditionals are opposites. Conversely, strong modality skew does not prove semantic expert specialization; it can arise from token-count imbalance, prompt packing, capacity order, dataset mix, or one checkpoint’s router state.

Source: MoE-LLaVA v5 appendix routing visualization for one model/evaluation slice. The left panel projects total expert loading; the expert panels project text-versus-image composition. The dashed percentages are descriptive, not an optimal-balance target.

Source: MoE-LLaVA v5 appendix. Text-conditioned and image-conditioned expert proportions vary by layer. The figure records a final checkpoint and selected evaluation population; it does not show a modality label entering the gate.

0x07I Token Pathways Are Observations, Not Stable Expert Names

The appendix also traces the top-1 and top-2 expert index through successive MoE layers for text and image populations:

Source: MoE-LLaVA v5 appendix token-pathway visualization. A line joins aggregate expert indices across layers; it is not a single token’s causal trajectory and does not prove that Expert 1 means the same function at every layer.

The paper reports different patterns across OpenChat, Phi, Qwen, and StableLM variants. It also observes that the larger OpenChat sparse model performed worse than expected despite a comparatively balanced-looking distribution and speculates that the multimodal tuning data may be insufficient. That causal story is explicitly speculative. The defensible conclusion is narrower: balance, specialization, model scale, data scale, and quality must be measured together; no one routing picture establishes the cause of a benchmark result.

0x07J Prefill, Decode, K/V, And Route State

MoE-LLaVA sparsifies FFNs, not attention depth. During multimodal prefill, projected image rows and text rows occupy one causal sequence and produce ordinary dense per-layer K/V. The attention mask excludes rectangular padding from semantic visibility, but the pinned MoE wrapper routes every materialized [B,S_max,d] row because it receives no used_token mask. Capacity can change a content or padding row’s FFN delta, but an FFN route does not erase the ordinary attention K/V representation. Rectangular cache slots and semantic content positions must remain separate M5 ledgers.

The pinned generation path handles a one-token decode step specially: prepare_inputs_for_generation keeps only the latest token ID, while the multimodal input builder detects sequence length one and reuses the existing past K/V instead of projecting the image again. It updates the attention mask and position ID to the next dense semantic position. The current hidden row still passes through each MoE FFN router.

multimodal prefill
    image projector runs
    image and text rows enter dense attention
    ordinary K/V rows are cached
    MoE FFNs route every materialized post-attention row, padding included

one-token decode
    image projector is not rerun
    past multimodal K/V is reused
    one new dense K/V row is appended
    one current hidden row is routed per MoE layer

This differs from Mixture-of-Depths Loop 7: no per-layer ABSENT_BY_ROUTE K/V state is created by an FFN expert decision. Prefix compatibility still needs the checkpoint, tokenizer, image preprocessing, projector, and model program, but this method does not add a modality-route label to the cache ABI.

0x07K Alternatives Change Which Problem Is Solved First

design	advantage	changed contract
sparse from Stage I	fewer named stages	alignment, task learning, and routing compete from initialization
dense Stage-II warm start	isolates sparsification after task acquisition	requires conversion and optimizer restart/transfer policy
freeze projector in Stage III	matches paper figure more closely	image coordinate bridge cannot adapt to expert divergence
re-enable projector	matches pinned Phi code order	Stage-III trainable state is larger than “MoE only”
deterministic top-2	replay is simpler	differs from pinned Gumbel second choice
BPR admission	preserves high-priority expert edges	needs per-expert priority ordering and changes dropped modalities
no-drop dynamic capacity	preserves all selected edges	loses the fixed capacity bound
modality-specific gates or experts	exposes modality as model input	no longer the MoE-LLaVA hidden-state-only router studied here
shared always-on expert	protects common function	changes resident/active parameters and is owned by M6.2 variants

0x07L A Twelve-Row Executable Trace

The Loop 9 fixture packs two sequences to $S_{max}=6$:

batch 0: IMAGE IMAGE TEXT_PROMPT TEXT_TARGET TEXT_TARGET PAD
batch 1: IMAGE TEXT_PROMPT TEXT_TARGET TEXT_TARGET PAD PAD

It has nine valid content rows, three image rows, six text rows, three padding rows, and four direct target-label rows. Because the pinned Phi/DeepSpeed path flattens the rectangular tensor without used_token, all twelve materialized rows enter the router. With $E=4$, top-2, and $c_f=1.5$:

materialized router rows       12
valid content rows              9
capacity per expert             9
requested routes               24
executed routes                24
complete route losses           0
max |cloned MoE - dense FFN|  2.22e-16
logical dispatch payload       96 bytes
logical round trip            192 bytes

With the same logits, experts, Gumbel realization, and row order but $c_f=0.4$:

capacity per expert             3
requested routes               24
executed routes                12
complete route losses           3
max |cloned MoE - dense FFN|    1.75
logical dispatch payload       48 bytes
logical round trip             96 bytes

The lower payload is not a speed or quality win; it is partly missing work. Relabeling every valid content row’s modality leaves logits, selected routes, and outputs unchanged, while it changes the modality contingency table. Padding remains a separate category. That executable invariant proves the router consumes hidden state, not the analysis label.

0x07M Failure Lab And Executable Closure

failure	collapsed distinction	observable consequence
Stage III called “MoE only”	paper figure and code order merged	projector optimizer state omitted
all text rows called AR targets	text modality and non-ignore labels merged	wrong loss denominator
visual rows called gradient-free	direct labels and indirect influence merged	projector/router training misunderstood
padding excluded from the pinned router	attention validity mask assumed to reach self.mlp	wrong capacity, auxiliary loss, and payload denominators
top-2 described as deterministic largest two	pinned Gumbel second choice omitted	route replay diverges
cloned experts assumed exactly continuous	route admission omitted	complete-drop rows lose dense FFN delta
BPR claimed for public code	paper prose and dependency path merged	overflow ordering misdocumented
top-1 counts called executed load	pre-capacity statistic and executed edges merged	device work misestimated
equal expert totals called modality balance	marginal and conditional distributions merged	false specialization conclusion
FFN route said to remove K/V	attention state and FFN delta merged	invalid cache ABI

Replay the complete contract with:

python3 docs/numinens-evidence/m6-8-conditional-compute/multimodal_stage_oracle.py --write-report --fail-on-error
python3 -m unittest docs/numinens-evidence/m6-8-conditional-compute/test_loop9_multimodal_stage.py

The oracle executes packed-row construction, the pinned top-2 equations with a stored Gumbel realization, materialized padding routes, capacity locations, surviving-gate normalization, cloned-expert continuity, stress overflow, auxiliary-loss populations, modality invariance, parameter/byte ledgers, and 36 one-field contract mutations. It does not train MoE-LLaVA, reproduce a benchmark, execute a GPU collective, or claim that the pinned public code implements BPR.

0x07N Loop 10 Handoff

Loop 9 freezes the last operator-specific route in this article:

input state        interleaved visual/text hidden rows plus valid/label metadata
router input       hidden state only; modality label excluded
stage transition  projector alignment -> dense LVLM -> copied experts + router
warm start         exact only for admitted cloned routes with normalized weights
loss populations  direct non-ignore labels vs all materialized routed rows
capacity evidence paper BPR claim != pinned DeepSpeed row-order implementation
cache state        ordinary dense multimodal K/V; FFN routing changes only delta

Loop 10 may compare MoA, SwitchHead, MoD, MoDE, and multimodal sparse experts paper by paper. It may not erase their different routed objects, training transitions, random state, capacity semantics, or K/V consequences.

0x08 Paper-By-Contract Comparison: Compatible Shapes Are Not Equivalent Operators

The operator-specific sections are now frozen. A useful comparison must do more than place the words top-k, router, and expert in one table.

When may one conditional-compute method replace another without changing the model, training procedure, cache state, or evidence claim?

The answer is deliberately strict: only when the complete method contract agrees. Similar tensor ranks, equal parameter counts, or the same output shape are necessary for some substitutions, but they are not sufficient.

0x08A One Method Is A Twenty-One-Field Contract

For method $m$, define a typed record:

$$\mathfrak{C}_m = \bigl( O_m,U_m,S_m,A_m,W_m,R_m, Q_m,F_m,B_m,T_m,K_m,Z_m,E_m \bigr).$$

Those symbols group the 21 stored fields rather than hiding them:

group	required fields
changed object	routed object; decision unit
candidates	selection population; selection owner
route rule	score transform; selection policy; admission policy; randomness state; capacity unit
weighted path	route weight; residual location; attention execution; FFN execution
differentiation	backward boundary
lifecycle	training-to-decode transition; persistent route metadata
serving state	K/V row policy; K/V payload identity
side populations	padding-router policy; modality-label role
authority	paper, official code, dependency, or explicitly local systems contract

Let $\mathcal{F}$ be this field set. Contract equivalence means:

$$m\equiv n \quad\Longleftrightarrow\quad \forall f\in\mathcal{F},\; \mathfrak{C}_m[f]=\mathfrak{C}_n[f].$$

For audit purposes, define a Hamming-style contract distance:

$$D(m,n) = \sum_{f\in\mathcal{F}} \mathbf{1} \left[ \mathfrak{C}_m[f]\ne\mathfrak{C}_n[f] \right].$$

This is not a scientific similarity metric. It is a substitution detector. In the seven checked records below, all $\binom{7}{2}=21$ pairs differ, with $14\le D(m,n)\le21$. A distance of zero would permit comparison at this field resolution; a nonzero distance names exactly what must be adapted.

0x08B Evidence Is A Field, Not A Footnote

The same operator description can have different authority at different layers of detail:

primary paper
    establishes the printed method, objective, experiment, and stated limits

pinned official implementation
    establishes one concrete tensor schedule and saved state

pinned dependency
    may decide routing, admission, or numeric behavior hidden by the caller

Transformer System executable fixture
    proves only the declared toy operator and its negative substitutions

M5-consuming systems contract
    derives cache identity, publication, rollback, and reuse requirements

MoA and SwitchHead have pinned public implementations. The Mixture-of-Depths paper does not publish the complete local sparse-cache ABI or every Figure 7 router choice, so those cells remain paper + local contract. MoE-LLaVA needs three separate surfaces because the paper, repository, and DeepSpeed 0.9.5 dependency disagree on Stage-III ownership, BPR, and effective routed rows.

0x08C The Seven Frozen Contracts

method	routed object and candidates	selection and admission	residual and persistent state
MoA Q/O	destination row selects coupled Q/O experts	full-$E$ softmax, top-$k$, stopped selected-mass denominator; no capacity drop in pinned operator	selected attention maps; shared historical K/V without expert axis
SwitchHead V/O	source rows select V components; destination rows select O components per fixed map head	independent sigmoid top-$k$ on each side; no selected-mass normalization	fixed $H_A$ maps; cache shared K plus source-routed mixed V
MoD training	valid rows compete for membership in one whole block	stable exact top-$C$ per sequence; padding excluded; chronological execution	selected block delta scattered into residual; no serving-cache claim in teacher fixture
MoD decode	current token-layer pair decides block membership	strict causal threshold; variable count; original router scalar still gates the delta	selected block or whole-block identity; per-layer PRESENT_KV or ABSENT_BY_ROUTE
staged MoDE	depth membership, then FFN expert over selected rows	top-$C$ depth route followed by a real-expert route	attention and K/V only for depth members; two hard route boundaries
integrated MoDE	post-attention row selects a real FFN expert or no-op	top-1 over $E+1$ paths in the local fixture	attention and K/V for every valid row; no-op removes only FFN delta
MoE-LLaVA top-2	materialized multimodal hidden rows select FFN experts	softmax top-1 plus Gumbel second; row-order capacity in pinned dependency	dense valid-row attention/K/V; pinned router also processes padding rows

This table is compact, but each row in routing_family_contract_fixture.json stores all 21 fields and a complete semantic signature.

0x08D Same Array Shape, Different Attention Semantics

The Loop 3 and Loop 4 teaching fixtures both produce:

$$Y\in\mathbb{R}^{1\times2\times2}, \qquad A\in\mathbb{R}^{1\times2\times2\times3}.$$

That shape equality hides a different second axis. For MoA:

$$A^{\text{MoA}} \in \mathbb{R}^{B\times S_q\times k_{qo}\times S_{kv}},$$

where each selected slot names a coupled Q/O expert for one destination. For SwitchHead:

$$A^{\text{SH}} \in \mathbb{R}^{B\times H_A\times S_q\times S_{kv}},$$

where $H_A$ is a fixed attention-map identity after the source V mixture and before the destination O mixture. The numeric fixture happens to have $S_q=k_{qo}=H_A=2$, so a positional array check cannot distinguish them.

The caches expose the substitution:

$$K^{\text{MoA}},V^{\text{MoA}} \in\mathbb{R}^{B\times S_{kv}\times d_h},$$ $$K^{\text{SH}},V_{\text{mix}}^{\text{SH}} \in\mathbb{R}^{B\times S_{kv}\times H_A\times d_h}.$$

MoA has one destination router and a stopped selected-mass normalization. SwitchHead has source and destination routers and independent sigmoid weights. Replacing one implementation with the other because the attention arrays are both [1,2,2,3] would preserve memory shape while changing the function, gradient, and historical state.

0x08E Same Output Shape, Different Residual And K/V State

Staged and integrated MoDE both return $Y\in\mathbb{R}^{1\times4\times2}$ in the Loop 8 fixture and both store 30 named expert-plus-router parameters. Their path equations differ.

For staged composition, with depth membership $z_t\in\{0,1\}$:

$$Y_t^{\text{staged}} = X_t + z_t\left( r_t\,\Delta A_t + r_t\,g_{t,e_t}\,\Delta F_{t,e_t} \right).$$

For integrated composition:

$$H_t=X_t+\Delta A_t,$$ $$Y_t^{\text{integrated}} = H_t + g_{t,p_t}\,\Delta F_{t,p_t}(H_t),$$

with $\Delta F_{t,\text{noop}}=0$. Therefore the K/V-presence indicator is:

$$\kappa_t^{\text{staged}}=z_t, \qquad \kappa_t^{\text{integrated}}=1.$$

The fixture consequently records:

staged KV     ABSENT, PRESENT, ABSENT, PRESENT
integrated KV PRESENT, PRESENT, PRESENT, PRESENT

Equal output shape and equal named parameter count do not imply equal active attention work, prefix identity, or cache ABI.

0x08F Candidate, Requested, And Admitted Populations

For each method, let $\mathcal{U}_m$ be the rows eligible for scoring, $\mathcal{R}_m$ the requested route edges, and $\mathcal{A}_m$ the admitted edges:

$$\mathcal{A}_m \subseteq \mathcal{R}_m \subseteq \mathcal{U}_m\times\mathcal{P}_m,$$

where $\mathcal{P}_m$ is the declared path pool. The sets mean different things:

method	$\mathcal{U}_m$	why $\mathcal{A}_m$ may differ from $\mathcal{R}_m$
MoA	destination rows	no capacity admission in pinned local operator
SwitchHead	source and destination rows separately	no capacity admission in pinned local operator
MoD training	valid rows inside each complete sequence	exact top-$C$ membership defines request and admission together
MoD decode	current row only	threshold decides a variable binary path
staged MoDE	valid rows, then depth-selected rows	depth capacity removes rows before expert routing
integrated MoDE	every post-attention valid row	learned no-op is a selected preference, not rejected admission
pinned MoE-LLaVA	every materialized rectangular row, including padding	expert capacity can reject one or both selected edges

This is why an expert-load denominator cannot be copied across rows. Staged MoDE expert balance sees the depth-selected population. MoE-LLaVA’s pinned auxiliary path sees all materialized router rows. Neither denominator is the direct autoregressive-label count.

0x08G Hard Decisions And Backward Contracts

Every hard method is piecewise smooth only after the discrete state is frozen. For a saved decision state $d_m$:

$$\left. \nabla_{\theta}L_m \right|_{d_m} = \frac{\partial L_m(X;\theta,d_m)}{\partial\theta}.$$

That notation does not differentiate through top-k, top-$C$, threshold comparison, Gumbel argmax, or capacity admission. The saved state is itself method-specific:

MoA          selected Q/O ids + full probabilities + stopped denominator state
SwitchHead   source ids/weights + destination ids/weights
MoD train    score-ranked ids + chronological ids + original positions
MoD decode   predictor policy + threshold + comparator + membership outcome
staged MoDE  depth state + nested expert state
integrated   real-expert/no-op state
MoE-LLaVA    top1 + Gumbel second + row order + capacity locations + survivors

Even the smooth route weights differ: stopped selected normalization, selected sigmoid, raw scalar gating, softmax path probability, or survivor renormalization are not interchangeable VJPs.

0x08H Training-To-Decode Is Part Of The Operator

A route rule used during teacher-forced training is not automatically a causal serving rule:

$$g_m^{\text{train}} \not\equiv g_m^{\text{decode}} \quad\text{unless the transition contract proves equivalence.}$$

method	lifecycle contract
MoA	destination route remains query-local; shared historical K/V is reusable
SwitchHead	source route executes when a token enters K/V; destination route executes at query time
MoD	complete-sequence top-$C$ teacher becomes a causal threshold predictor with no exact-$C$ guarantee
staged MoDE	inherits the MoD causalization problem before nested expert routing
integrated MoDE	dense attention decode remains ordinary; route acts on the current post-attention row
MoE-LLaVA	Stage I alignment, Stage II dense task learning, and Stage III copied experts/router are checkpoint and optimizer transitions

The transition row also names persistent compatibility state. A changed MoD predictor invalidates sparse K/V reuse. A changed SwitchHead source router changes mixed historical V. A changed FFN route need not add an expert axis to ordinary dense K/V, but a changed multimodal projector still changes the prefix representation.

0x08I K/V Consequences Form Three Classes, Not One Ranking

For one request and layer, define:

$$N_{KV}^{(m,l)} = \sum_t \kappa_{m,l,t},$$

where $\kappa_{m,l,t}=1$ only when method $m$ materializes a valid historical row at layer $l$. The contracts partition into:

1. Dense-attention K/V: MoA, SwitchHead, integrated MoDE, and multimodal FFN MoE retain all source or valid content rows, although their K/V payloads differ.
2. Sparse depth K/V: causal MoD and staged MoDE record selected rows and explicit absence for skips.
3. Training-only state: the fixed-capacity MoD teacher fixture stores activations and selected identities but makes no serving-cache claim.

Calling the first class “the same cache” would still be wrong. MoA stores shared K/V without a map-head axis; SwitchHead stores head-specific K and mixed V; integrated MoDE and MoE-LLaVA consume ordinary dense attention K/V under different model and side-input identities.

0x08J A Work Vector Cannot Be Collapsed Into One Number

The executable comparison reports a named vector:

$$\mathbf{w}_m = \left( N_{score}, N_{proj}, N_{map}, N_{block}, N_{ffn}, N_{KV} \right)_m.$$

Its coordinates count different units from different teaching fixtures. A null coordinate means the fixture did not materialize that kind of work; it does not mean zero runtime cost.

method	score scalars	selected projection edges	map elements	block rows	FFN edges	committed K/V rows
MoA Q/O	6	4	12	-	-	3
SwitchHead V/O	30	20	12	-	-	3
MoD training	12	-	36	6	-	-
MoD decode sparse-state fixture	24	-	-	12	-	12
staged MoDE	6	-	-	2	2	2
integrated MoDE	12	-	-	4	3	4
MoE-LLaVA main-capacity fixture	48	-	-	-	24	9

Do not sum these columns, divide one method by another, or call the result a speedup. The fixture dimensions differ, and projections, attention-map elements, FFN edges, logical payload, kernel work, and measured latency have different units. Equal-budget empirical comparison requires a new manifest that pins model dimensions, sequence distribution, checkpoint, phase, hardware, kernels, batching, and measurement protocol.

0x08K Executable Comparison And Failure Closure

The Loop 10 oracle does not rerun every earlier mathematical derivation. It consumes the exact semantic digests of all seven Loop 3-9 reports, checks one mechanism-defining witness from each, normalizes the 21 fields, and rejects same-shape substitutions:

for method in methods:
    require_all_21_fields(method)
    verify_semantic_signature(method)
    verify_upstream_report_witness(method)

for left, right in all_pairs(methods):
    differing = [f for f in fields if left[f] != right[f]]
    assert differing  # no silent contract equivalence

for case in same_shape_counterexamples:
    assert case.left.shape == case.right.shape
    assert all(case.left[f] != case.right[f] for f in case.must_differ)

invalid substitution	shape check	contract failure
MoA map axis replaced by SwitchHead map axis	passes in the toy shape	routed object, route sides, weights, and K/V payload change
staged output replaced by integrated output	passes	attention execution, residual location, capacity, and K/V state change
MoD predictor probability used as delta gate	scalar-compatible	training objective replaces the model-path gate
paper BPR copied into pinned MoE-LLaVA runtime	route tensor-compatible	admission order and dropped edges change
padding removed from pinned router population	output can remain masked	capacity and auxiliary-loss denominators change
no-op encoded as ABSENT_BY_ROUTE	output delta can remain zero	historical attention state is deleted incorrectly

Replay the comparison with:

python3 docs/numinens-evidence/m6-8-conditional-compute/routing_family_contract_oracle.py --write-report --fail-on-error
python3 -m unittest docs/numinens-evidence/m6-8-conditional-compute/test_loop10_routing_family_contracts.py

The comparison proves contract consistency for the declared fixtures. It does not rank quality, reproduce a paper benchmark, execute a production cache or GPU kernel, or measure latency.

0x08L Loop 11 Handoff

Loop 10 freezes what a method is. Loop 11 may now vary batch composition, sequence length, selected fraction, route entropy, expert skew, threshold, capacity factor, padding ratio, prefill/decode phase, and cache history while holding the 21-field contract fixed. A sensitivity point that changes one of those fields is a different method configuration and must receive a new identity rather than being mixed into the same curve.

0x09 Implementation Blueprints

0x09A MoD Training Block

def mod_block_train(x, valid, position_ids, router, block_delta, capacity):
    # x: [B, S, d], valid: [B, S]; select exactly capacity valid rows per
    # sequence (each sequence must expose at least capacity valid rows).
    scores = router(x).squeeze(-1)                 # [B, S]
    scores = where(valid, scores, -inf)            # padding never competes
    top_scores, top_idx = stable_topk(             # score desc, row id asc
        scores, capacity, dim=1, tie_break="row_index_ascending",
    )

    x_sel = gather_rows(x, top_idx)                # [B, C, d]
    pos_sel = gather_rows(position_ids, top_idx)   # [B, C]

    # Sort by original position if the attention kernel requires it.
    order = argsort(pos_sel, dim=1)
    x_sel = gather_rows(x_sel, order)
    pos_sel = gather_rows(pos_sel, order)
    top_scores = gather_rows(top_scores, order)
    top_idx = gather_rows(top_idx, order)

    # Contract: block_delta returns F(x_sel), not x_sel + F(x_sel).
    update = block_delta(x_sel, position_ids=pos_sel)
    update = update * top_scores[..., None]

    y = x.clone()
    scatter_add_rows(y, top_idx, update)
    return y, stop_gradient(top_idx)

0x09B Predictor Training

teacher = zeros([B, S])
teacher.scatter_(dim=1, index=top_idx, value=1.0)

predictor_logits = predictor(x.detach()).squeeze(-1)
predictor_loss = binary_cross_entropy_with_logits(
    predictor_logits[valid],
    teacher[valid],
)

0x09C MoD Decode

def mod_block_decode(
    x_t,
    position_t,
    layer_id,
    router,
    membership_predictor,
    block_delta,
    cache_snapshot,
    token_txn,
):
    route_weight = router(x_t).squeeze(-1)  # Original r_t^l on model path.

    if membership_predictor is None:
        # Direct-BCE variant: the router output also predicts membership.
        membership_logit = route_weight
        threshold = router.membership_threshold
    else:
        # Auxiliary-MLP variant: predictor decides membership only.
        membership_logit = membership_predictor(
            stop_gradient(x_t)
        ).squeeze(-1)
        threshold = membership_predictor.threshold

    membership_prob = sigmoid(membership_logit)

    if membership_prob <= threshold:
        token_txn.stage_absence(
            layer=layer_id,
            position=position_t,
            reason="ABSENT_BY_ROUTE",
        )
        return x_t  # exact residual identity; no layer K/V row

    update, new_k, new_v = block_delta.decode(
        x_t,
        position_id=position_t,
        sparse_kv=cache_snapshot.layer_cache(layer_id),
    )
    token_txn.stage_kv(
        layer=layer_id,
        position=position_t,
        key=new_k,
        value=new_v,
    )
    return x_t + route_weight[..., None] * update

The caller commits token_txn only after every routed layer has staged either K/V or ABSENT_BY_ROUTE. This is a required logical state contract, not a claim that the paper ships this exact cache implementation.

0x0A Sensitivity And Evaluation: Vary Conditions, Not Method Identity

The operator derivations answer what executes. Evaluation must answer a different question: which conclusion changes when one operating condition changes while the operator remains the same? A table that changes route policy, padding population, cache state, sequence geometry, and kernel at once cannot attribute its result to any one of them.

Loop 10 therefore supplies the immutable method identity. For method $m$, let the twenty-one semantic fields be serialized in source order and hashed as:

$$\sigma_m = \operatorname{SHA256} \left( \operatorname{canonical} \left\{ f_j(m) \right\}_{j=1}^{21} \right)$$

A sensitivity point is admissible only when it preserves $\sigma_m$. The operating record is:

$$\mathcal{S}_{m,i} = \left( \sigma_m, x_i, \phi_i, \mathcal{L}_{m,i}, \mathcal{E}_{i} \right)$$

where $x_i$ is exactly one varied coordinate, $\phi_i$ is the complete frozen context, $\mathcal{L}_{m,i}$ is a vector of method-specific logical ledgers, and $\mathcal{E}_i$ is the evidence class. If $\sigma_m$ changes, the point is not another sample on the same curve. It is a different method configuration.

The local executable evidence class is deliberately narrow:

evidence class:       synthetic logical counterfactual
quality measured:     no
latency measured:     no
throughput measured:  no
training reproduced:  no
GPU runtime executed: no
cross-method ranking: no

It can prove arithmetic and state consequences of a frozen teaching fixture. It cannot establish model quality, wall time, kernel efficiency, or production stability.

0x0AA Four Evaluation Layers Must Not Be Collapsed

Layer	Experimental object	Valid evidence	Invalid promotion
logical operator	route relation, admission, work and state vectors	exact oracle and negative substitutions	latency or quality
trained model	checkpoint, dataset, objective, random seeds	held-out loss, task metrics, route statistics	universal method ranking
physical runtime	hardware, software, batch, sequence, placement, kernel	measured latency, throughput, bytes, utilization	model quality
serving lifecycle	scheduler, prefix/cache identity, failures, rollback	canary traces and fault injection	operator equivalence

The layers consume each other but do not replace one another. A logical oracle can specify which counters a runtime must expose. It cannot predict their timing. A benchmark can report an end-to-end speedup. It cannot reveal whether the route relation, cache identity, or training objective drifted unless those states were also recorded.

0x0AB Routed-Attention Sensitivity Has Different Source And Destination Axes

For MoA, define the normalized entropy of the selected route weights over $T_q$ destination rows and $k_A$ selected experts:

$$H_{sel} = -\frac{1}{T_q\log k_A} \sum_{t=1}^{T_q} \sum_{j=1}^{k_A} \bar g_{t,j}\log \bar g_{t,j}$$

where $\bar g_{t,j}$ is the selected weight after the paper operator’s normalization. In the fixture, the selected expert IDs remain [[0,1],[0,2]], while the two selected weights move from concentrated to uniform. The average normalized entropies are:

concentrated:      0.377696
mixed:             0.926121
uniform selected:  1.000000

All three points still create twelve attention-map elements and cache three shared source K/V rows. Entropy measures weight concentration. It does not reconstruct route identity, map count, or cache payload.

SwitchHead exposes a different sensitivity. With $T_q$ destination rows, $T_k$ source rows, $H_A$ computed attention maps, $k_V$ selected source-V components, and $k_O$ selected destination-O components:

$$N_{map}=T_qT_kH_A$$ $$N_{Vroute}=T_kH_Ak_V$$ $$N_{Oroute}=T_qH_Ak_O$$

Holding $T_q=2$, $H_A=2$, and $k_V=k_O=2$ while changing $T_k$ from 1 -> 3 -> 8 changes map elements 4 -> 12 -> 32, source route edges 4 -> 12 -> 32, and cached rows 1 -> 3 -> 8. Destination route edges stay at eight. One generic “tokens times top-k” counter cannot represent these source and destination ownership relations.

0x0AC MoD Training Fraction Has Linear And Quadratic Coordinates

For one valid sequence of length $S$ and exact per-layer capacity $C$, define:

$$f=\frac{C}{S}$$

The selected FFN row count is linear in $f$. The selected self-attention map, however, is quadratic in the selected sequence length:

$$\frac{N_{FFN,sel}}{N_{FFN,dense}}=f$$ $$\frac{N_{map,sel}}{N_{map,dense}}=f^2$$

For $S=8$, two attention heads, and capacities 2,4,6:

$C$	$f$	selected block rows	attention-map elements $H C^2$	FFN coordinate	map coordinate
2	0.25	2	8	0.25	0.0625
4	0.50	4	32	0.50	0.2500
6	0.75	6	72	0.75	0.5625

These are logical dimensions, not a latency prediction. Gather/scatter, projection shapes, kernel tails, launch count, memory traffic, and realized hardware occupancy remain unmeasured. Even within one operator, “selected fraction” is not one universal savings scalar.

0x0AD Decode Threshold Changes Membership And Sparse State Together

The causal decode policy selects the current token-layer pair when:

$$m_t(\tau) = \mathbf{1}\left[\sigma(a_t)>\tau\right]$$

For a fixed probability vector $p_t=\sigma(a_t)$ over a workload cell:

$$N_{sel}(\tau) = \sum_t\mathbf{1}[p_t>\tau]$$

With probabilities [0.12,0.31,0.49,0.51,0.73,0.94] and a strict > comparator:

threshold $\tau$	selected	skipped	new sparse K/V rows
0.25	5	1	5
0.50	3	3	3
0.75	1	5	1

This fixture holds the probability vector, comparator, layer, route-weight rule, and cache ABI fixed. It proves only that thresholding changes the count and the sparse state transition. It does not claim that one threshold preserves the training teacher’s exact $C$, quality, or calibration on another rollout.

A threshold study on a trained model must therefore report at least:

checkpoint and predictor revision
layer and workload cell
complete probability distribution or calibrated bins
threshold and comparator
selected count and fraction
teacher TP, FP, FN, TN where teacher labels exist
raw route-gate distribution on selected rows
new PRESENT_KV and ABSENT_BY_ROUTE counts
quality metric with exact denominator

0x0AE Staged And Integrated MoDE Need Different Sensitivity Ledgers

For staged MoDE with depth capacity $C_D$ and one real expert per selected row:

$$N_{attn}=C_D, \qquad N_{expert}=C_D, \qquad N_{KV,new}=C_D$$

Changing $C_D$ from 2 -> 4 -> 8 moves all three coordinates together because depth admission occurs before attention and FFN expert routing.

For integrated MoDE with $S$ valid rows, $N_0$ learned no-op selections, and one selected path per row:

$$N_{attn}=S$$ $$N_{expert,real}=S-N_0$$ $$N_{KV,new}=S$$

At $S=8$, changing no-op rows from 0 -> 2 -> 6 changes real-expert edges 8 -> 6 -> 2, while dense attention rows and K/V rows remain eight. This is the operational reason that the same output shape and a similar “skip fraction” do not make staged and integrated composition interchangeable.

0x0AF Capacity And Padding Alter Different Multimodal Populations

For $T_{mat}$ materialized rows, $k$ requested experts per row, $E$ experts, and expert request counts $n_e$:

$$R=T_{mat}k$$ $$C(\mathrm{CF}) = \left\lceil \mathrm{CF}\frac{R}{E} \right\rceil$$ $$A=\sum_{e=1}^{E}\min(n_e,C)$$ $$D=R-A, \qquad U=EC-A$$

For the frozen request histogram [9,7,5,3]:

capacity factor	$C$	requested	executed	dropped	unused fixed slots
0.5	3	24	12	12	0
1.0	6	24	20	4	4
1.5	9	24	24	0	12

This is an admission sensitivity: requested route identity remains frozen, while the executed relation changes. The model function can therefore change. It is not a pure memory or runtime sweep.

Padding sensitivity changes an earlier population. In the pinned Phi path, padding rows are materialized router rows. With nine valid rows and top-2:

padding rows	materialized rows	padding ratio	requested routes	balanced-mean capacity at CF=1
0	9	0.00	18	5
3	12	0.25	24	6
9	18	0.50	36	9

The valid content remains fixed. Router work, requested routes, capacity, and auxiliary-loss population do not. An evaluation that reports only valid token count cannot explain this change.

0x0AG Phase And Cache History Are Part Of The Evaluation Cell

The same new output row does not imply the same cache transition:

method / phase	history rows	new rows	cache class	rows after
MoA decode	3	1	dense shared K/V	4
SwitchHead decode	3	1	shared K plus mixed V	4
MoD training	0	0	training-only activation state	0
MoD selected decode	3	1	sparse depth K/V	4
MoD skipped decode	3	0	sparse depth K/V	3
staged MoDE selected decode	2	1	sparse depth K/V	3
integrated MoDE no-op decode	3	1	dense attention K/V	4
MoE-LLaVA decode	3	1	dense multimodal K/V	4

Integrated no-op routing still appends attention K/V because attention already executed. A skipped MoD depth path does not. Batch size, output shape, and FFN route count cannot reconstruct that difference.

0x0AH A Reproducible Model And Runtime Study Needs Cells

The local oracle supplies a schema, not a benchmark result. A real study should cross workload, model, route, state, and runtime coordinates explicitly:

model:
    checkpoint hash, tokenizer, precision, architecture, routing revision

workload:
    dataset and split, batch composition, source/destination lengths,
    modality and padding composition, prompt/prefill/decode phase

route:
    score transform, top-k or threshold, tie/noise state,
    requested/admitted IDs, entropy/skew, capacity and overflow

state:
    cache class, history length, semantic positions, prefix policy,
    route-policy generation, rollback/fork state

runtime:
    hardware, framework and commit, kernel, compilation mode,
    placement, warmup, repetitions, synchronization, timing boundary

outcomes kept separate:
    quality | logical work | persistent state | temporary state |
    measured traffic | latency | throughput | failure/recovery

Comparisons must also name the matching rule. Equal total parameters, equal active parameters, equal logical MACs, equal measured latency, and equal quality are different controls. Most paper tables match only a subset. Source results remain source-reported unless the complete cell is reproduced.

0x0AI Exact Replay And Negative Sensitivity Fixtures

Replay the 24 operating points and eight phase/cache rows with:

python3 docs/numinens-evidence/m6-8-conditional-compute/conditional_sensitivity_oracle.py --write-report --fail-on-error
python3 -m unittest docs/numinens-evidence/m6-8-conditional-compute/test_loop11_sensitivity_failure_closure.py

The twenty one-field invalid fixtures reject hidden method-signature changes, inapplicable coordinates, unlike-unit scalar ranking, synthetic latency or quality claims, wrong entropy and work ledgers, false fixed-capacity decode, staged/integrated K/V substitutions, capacity arithmetic drift, padding population drift, and cache-history corruption.

0x0B Typed Failure Closure, Sources, And Publication Surface

0x0BA A Failure Record Needs Four Coordinates

The old three-column failure list named a symptom but did not always say how to detect it or who owns the repair. Loop 11 closes each record as:

$$\mathcal{F}_i = \left( I_i, D_i, C_i, O_i \right)$$

where $I_i$ is the invariant, $D_i$ is an executable detector, $C_i$ is the downstream consequence, and $O_i$ is the recovery owner. “Tensor shape matches” is not a sufficient detector. “Fix the router” is not a sufficient owner when the violated state belongs to M5 cache publication or to a source evidence gate.

One historical example remains deliberately explicit: membership student replaces original router state is a forbidden collapse unless a named method contract explicitly says otherwise. Treating an auxiliary causal predictor as the original router would change the decision state, gradient owner, and training-to-decode transition; it is not a harmless implementation substitution. This preserves the original closure marker exactly: Loop 1: membership student replaces original router state.

0x0BB Same Shape Does Not Mean Same Failure Or Recovery

Two important substitutions can pass shape checks:

MoA attention        [1,2,2,3]
SwitchHead attention [1,2,2,3]

staged MoDE output   [1,4,2]
integrated output    [1,4,2]

The first pair differs in routed projection, source/destination population, score transform, selected weight, map identity, and cached V payload. The second differs in attention execution, residual location, capacity semantics, FFN population, and K/V presence. Recovery must return to the first violated semantic field, not reshape one output into another.

0x0BC Cross-Loop Failure Registry

Invariant	Detector	Consequence	Recovery owner
axis and stable identity do not collapse into payload shape	Loop 2 manifest counterexamples	route or cache provenance disappears	M6.8
MoA Q/O expert and stopped-denominator gradient are preserved	Loop 3 forward/backward oracle	wrong projection or router training	M6.8
SwitchHead source-V and destination-O populations stay separate	Loop 4 asymmetric oracle	wrong map, route work, or cache payload	M6.8 + M5
MoD top-C is per valid sequence and gathered rows retain positions	Loop 5 selection/position oracle	wrong membership, RoPE, or causal mask	M6.8 + M3
selected helper returns a delta, not a residualized block output	Loop 5 direct equivalence	residual is added twice	M6.8
noncausal teacher is not reused at decode	Loop 6 prefix witness	future leak or impossible route	M6.8
predictor membership does not replace the raw model-path gate	Loop 6 gradient oracle	causalization changes the model	M6.8
skipped depth row stages ABSENT_BY_ROUTE	Loop 7 state machine	fabricated K/V becomes visible	M5 runtime
sparse slot is not semantic position	Loop 7 slot/position oracle	position-dependent attention drifts	M6.8 + M5
route-aware token transaction closes atomically	Loop 7 fault injection	partial cross-layer prefix is published	M5 runtime
staged skip, integrated no-op, and overflow drop remain distinct	Loop 8 composition oracle	wrong attention, FFN, or K/V ledger	M6.8
paper, repository, and dependency are separate evidence surfaces	Loop 9 source trace	unexecuted policy is attributed to code	source gate
padding rows remain in the pinned Phi router population	Loop 9 packed-row oracle	route and objective denominator is understated	M6.8
cloned experts require admitted normalized routes for continuity	Loop 9 stress capacity	complete route loss is hidden	M6.8
all twenty-one fields define method identity	Loop 10 typed comparison	same-shape methods are substituted	M6.8
unlike work-vector units stay separate	Loop 10 work-vector assertion	fabricated speed ranking	M4/runtime measurement owner
a sensitivity point preserves the method signature	Loop 11 signature check	different methods enter one curve	M6.8
synthetic counts remain synthetic	Loop 11 evidence firewall	unsupported quality or latency claim	release gate

The machine-readable registry contains thirty records spanning Loops 2-11. Each record points to the exact test surface that detects the violation. The owner can be M6.8 mechanism prose, M3 positional semantics, M5 cache/runtime, M4 or a physical measurement owner, or the source/release gate. That is not administrative metadata; it prevents a logical failure from being “fixed” by a physical optimization that changes the operator.

0x0BD Complete Static Source And Visual Surface

The Loop 11 static closure checks:

35 unique source records
    17 local and executable records
    11 primary papers
     7 official implementation/project records

53 unique article image routes after this loop
11 cumulative provenance manifests
30 typed failure records
9 prior invalid-fixture families with 276 mutations
20 new one-field sensitivity mutations

Every local source is content-hash bound. Every external article URL must be covered by the source or artifact manifests. Every local post route must map to an article. Every image route must resolve, carry nonempty alternative text, and have provenance. Every SVG must expose role="img", a viewBox, a title, and a description.

Static closure does not prove browser legibility, MathJax layout, responsive source selection, live external availability, or publication readiness. Loop 12 owns normalized desktop/mobile, light/dark browser evidence and the final cross-article boundary audit.

0x0BE Loop 12 Handoff

Loop 11 freezes the evaluation schema, 24 sensitivity points, eight phase/cache rows, thirty failure records, complete static source coverage, and two new responsive visual pairs. Loop 12 may repair formula, source, link, rendering, accessibility, or boundary defects. It may not change the operator, evidence class, route semantics, or sensitivity conclusions without reopening the owning loop.

0x0C Claim Status And Evidence

Claim	Status	Qualification
MoA selects query/output attention experts per token while sharing K/V	externally verified	EMNLP 2022 MoA paper
SwitchHead’s reported best design routes V/O experts around shared Q/K per computed head	externally verified	SwitchHead paper; source design, not universal optimum
MoA and SwitchHead have the same KV contract	false	their routed projections differ
MoD enforces fixed training compute by selecting exactly $C$ token identities per routed block	externally verified	Mixture-of-Depths paper
MoD training top-C is noncausal for autoregressive sampling	externally verified	Mixture-of-Depths paper
auxiliary predictor exactly reproduces top-C and compute	false if generalized	source reports high accuracy, not identity or fixed count
MoD production KV ABI is fully specified by the paper	false	cache benefit is anticipated; runtime contract requires further engineering
staged MoDE can skip attention while integrated no-op experts cannot	externally verified and mechanism-derived	Mixture-of-Depths paper
MoE-LLaVA uses three stages before and during sparsification	externally verified	MoE-LLaVA v5 paper and official scripts
the Stage-III paper figure and pinned Phi code expose the same trainable set	false	pinned code re-enables the visual projector after MoE initialization freezes unmatched names
copied experts preserve the dense FFN for every token unconditionally	false	exact continuity also requires at least one admitted route and surviving selected weights summing to one
the pinned DeepSpeed 0.9.5 top-2 path implements the paper’s claimed BPR policy	false for the audited path	source uses cumulative flattened-row capacity locations and a Gumbel second choice
the pinned Phi top-2 router receives only valid non-padding rows	false	the decoder passes the full rectangular hidden tensor to the MoE and DeepSpeed flattens it without used_token
modality labels are inputs to the MoE-LLaVA router	false	modality is reconstructed analysis metadata; the gate consumes hidden state
sparse multimodal experts generally reduce hallucination	not established	paper result is source- and benchmark-specific

0x0D Handoffs

FFN router, capacity, and combine semantics
    -> M6.1

DeepSeek balance and group-limited routing
    -> M6.2

physical expert placement and all-to-all runtime
    -> M6.6

KV lifecycle and exact prefix reuse
    -> M5

routed LoRA and conditional adapter composition
    -> M6.9

0x0E Publish Gate

Check	Status
Routing axes are classified by expert, attention component, depth, and modality	pass
Dense attention establishes the Q/K/V/O and KV reference contract	pass
MoA router, stopped-gradient normalization, expert function, and decode path are derived	pass
SwitchHead Q/K and V/O ownership is separated from MoA	pass
Attention-map count and cache consequences are explicit	pass
MoD capacity, gather, attention/MLP FLOPs, scatter, and residual path are derived	pass
Expert-choice top-C training is distinguished from token-choice routing	pass
Noncausal training membership and causal predictor inference are separate	pass
Inference selected count is not claimed to remain exactly $C$	pass
Sparse per-layer KV membership and position state are explicit and source-scoped	pass
Staged MoDE, integrated no-op experts, and overflow dropping are distinct	pass
Materialized router rows, valid content rows, padding rows, direct-label rows, and modality metadata are separate	pass
All three checkpoint and optimizer boundaries are explicit	pass
Paper-figure and pinned-code Stage-III trainable state are separate	pass
Dense-to-MoE continuity is derived through admission and selected-gate normalization	pass
Paper BPR prose and pinned DeepSpeed row-order capacity are not equated	pass
Requested routes, executed routes, top-1 load, and modality distribution are separate	pass
Multimodal FFN routing retains ordinary dense attention K/V	pass
Runtime, evaluation, and failure contracts are included	pass
All benchmark and speed claims remain source-reported	pass

0xFF References

• Transformer Series A29 - DeepSeek MoE
• M6.1 - Routed Capacity
• M6.2 - Expert Parallelism and DeepSeek MoE
• M6.6 - MoE Runtime Systems
• Mixture of Attention Heads: Selecting Attention Heads Per Token
• SwitchHead: Accelerating Transformers with Mixture-of-Experts Attention
• Mixture-of-Depths: Dynamically Allocating Compute in Transformer-Based Language Models
• MoE-LLaVA: Mixture of Experts for Large Vision-Language Models
• MoE-LLaVA official repository at audited commit
• DeepSpeed 0.9.5 top-2 MoE gate
• Scaling Vision with Sparse Mixture of Experts
• DeeBERT: Dynamic Early Exiting for Accelerating BERT Inference
• Depth-Adaptive Transformer
• PoWER-BERT: Accelerating BERT Inference via Progressive Word-Vector Elimination
• Universal Transformers
• LayerSkip: Enabling Early Exit Inference and Self-Speculative Decoding
• Reducing Transformer Depth on Demand with Structured Dropout