SiTU-GLU as the FFN activation arm — formula fidelity, the wiring, and what the gates do not cover

Lane: situ-ffn, 2026-09-30. Worktree wt/situ-ffn off 9d571c2. Owner picked SiTU-GLU over SwiGLU; Kimi K3 (arXiv:2607.24653) reports better stability than SwiGLU for deep models and friendliness to low-bit quantization (MXFP4), which maps onto our two documented pains (the NaN-under-overfit class, AGENTS.md §2.3; Fp8 TSCT factor quantization, §2.3 max_ortho).

1. Verdict and tier

The form is tier (a): a named external reference exists and it agrees with the crate. The arithmetic is not ours — Moonshot ships it:

what	where	what it settles
SituAndMul, lines 64–82	huggingface.co/moonshotai/Kimi-K3@main modeling_kimi_linear.py	the exact form, and that the sigmoid is on the RAW gate
text_config.hidden_act = "situ", activation_situ_beta = 4.0, activation_situ_linear_beta = 25.0	same repo, config.json	the β the released model actually uses
Eq (12), §2.3.2; Eq (18)/(19), App. B	arXiv:2607.24653v2	the equation, the near-origin match to SwiGLU, the bound

The pinned numbers are a transcription. SituAndMul.forward needs torch and torch is not installed on this box, and nothing gets installed without the owner saying so. tools/ generator: vendor/dormouse-fused/crates/burn-situ/tools/gen_ref.py (numpy only). It keeps the reference’s own precision (.to(torch.float32) on both halves → f32, not f64) and its op order, so the goldens are f32-faithful up to a last-ulp reassociation. No external golden file exists and no byte-exact comparison is possible: the reference computes beta * tanh(gate/beta) * sigmoid(gate), the crate computes (gate/beta).tanh() * beta, and f32 multiplication is not associative. The gates assert 1e-6 relative, which is the f32 ulp band, not a loose bound.

The reference’s backward is not in the file; the crate’s fused backward is (fused_situ.rs:129-146) and it is the exact analytic derivative of Eq (12) — see §5, where it is a third independent confirmation of the derivative.

2. Equation ↔ file:line ↔ citation

SiTU-GLU(x) = β₁·tanh(W_g x / β₁) ⊙ Sigmoid(W_g x) ⊙ β₂·tanh(W_u x / β₂)     Eq (12)

term	crate (post-change lines)	our wiring	note
softcap(x, β) = β·tanh(x/β)	burn-situ/src/lib.rs:53-55	—	one line, the whole cap
split gate = first half, up = second half	lib.rs:107-108	loop_block.rs:559-562 (gate_up.forward)	reference does the same: x[..., :d], x[..., d:] (modeling_kimi_linear.py:76-78)
cap the gate	lib.rs:110	via situ_glu	Eq (12) caps the Swish gate’s LINEAR factor
Sigmoid(W_g x) on the raw gate	lib.rs:112	via situ_glu	the one place a transcription goes wrong; the crate has a test, see §6
cap the up branch	lib.rs:111	via situ_glu	independent β
gate_cap ⊙ σ(gate) ⊙ up_cap	lib.rs:112	loop_block.rs:564-570 (the use_situ branch, probe::note(SITU) at 564)
fused kernel, same form	fused_situ.rs:109	not enabled, see §7	bg·tanh(g/bg)·σ(g)·bu·tanh(u/bu) — the same expression
fused backward, exact derivative	fused_situ.rs:139-143	not enabled	(1-tanh²)·σ + β₁·tanh·σ·(1-σ), 1-tanh² — see §5

β is fixed, not learned: β₁ = 4 (gate), β₂ = 25 (up) for K3 (§2.3.2 “we set the soft-cap hyperparameters to β1 = 4 … and β2 = 25”, and the released config.json). Pinned as burn_situ::K3_GATE_BETA (lib.rs:29) / K3_UP_BETA (lib.rs:40). No learnable β and no config field — one mechanism per arm; a β sweep is a follow-up, and the class signature defaults to beta=1.0 (modeling_kimi_linear.py:70) which is NOT what K3 runs.

Correction to the crate’s own docs. burn-situ/src/lib.rs’s situ_glu doc and the README.md quick start both said the default is “1.0”. K3 runs 4 and 25. A reader who took the crate’s default would wire a cap 25× too tight on the up branch and would never find out from a loss curve — the cap is silent by construction. Both documents are fixed by this lane, the constants are named rather than defaulted, and the reason is on the constants.

3. What the arm actually is in OUR model (not a pure activation swap)

crates/dormouse-core/src/loop_block.rs:69-80 (struct + ExpertFFN::new) and the expert loop at :558-577. Before this change:

pub struct ExpertFFN { pub gate_up: LinearLike, pub down: LinearLike }
// gate_up: LinearLike::with_tsct(d, f, ...)   // d -> f.  ONE projection.
let mid = self.expert_ffns[e].gate_up.forward::<B>(normed_ffn.clone());
let mid = activation::silu(...);               // elementwise on [b*t, f]
let out = self.expert_ffns[e].down.forward::<B>(mid);   // f -> d

There is no gate today. The field is named gate_up but it is a single d -> f up-projection and the “gating” is elementwise SiLU — the MoE expert blend (blend, loop_block.rs:575-576) is a different, per-expert scalar softmax, not a GLU gate. So SiTU-GLU is not a drop-in replacement for activation::silu of the same tensor: Eq (12) needs a [b*t, 2f] input, because it reads W_g x and W_u x separately.

Wiring choice: gate_up becomes d -> 2f when use_situ, and the name finally becomes true. The alternative — splitting the existing f-wide projection into two halves — keeps the parameter count and halves the FFN width, which is a different mechanism at a different budget and confounds the A/B twice over. The paper’s shape is [d, 2·hidden] → act → [hidden, d], which is also what K3’s own SituAndMul asserts on entry (d = x.shape[-1] // 2).

The N % 4 == 0 pad is harmless here, and the reason is worth one line. LinearLike::new pads out_features up to a multiple of 4 for cubek vectorization (param.rs:47-53) but stores the requested width (param.rs:56), so (a) the branch assert compares requested against requested and cannot fire on a padded model, and (b) situ_glu gets a tensor of width ≥ 2f and slices [0..n, 0..f] and [0..n, f..2f], which are the gate and up columns at their true positions — the pad columns past 2f are simply not read. (The OFF path is the one that could not tolerate a d_ffn that is not a multiple of 4: silu would hand down the padded width. Pre-existing, out of scope here.)

Param cost on small (d_model 768, d_ffn 2048, rank 64, n_experts 3): the TSCT factors of gate_up go from d·r + r·f to d·r + r·2f, i.e. +r·f per expert = +131 072, ×3 experts = +393 216 on 9 197 390 = +4.28%. (Under a dense gate_up the same change is +51% — the number is only small because TSCT binds, and a reader should not transfer it to a dense model.)

This is a confound the A/B row states, not hides: SiTU vs silu is also +4.28% parameters. The width-matched control is SwiGLU at the same 2f, and burn-swiglu is already vendored — the paper’s own §2.3.2 comparison is exactly SwiGLU vs SiTU. One row is registered (§8); the confound is named in it.

4. The quantization interaction point (stated, not measured)

The claimed mechanism is that the output bound ‖SiTU-GLU(x)‖∞ ≤ β₁β₂ = 100 (Eq 19) removes activation outliers, so the low-bit formats stop overflowing. Where that lands in our FFN, in order — four steps, no diagram:

1. normed_ffn → gate_up (TSCT factors, Fp8) → [b*t, 2f], held in f32 because that is the tensor the cap’s input is
2. softcap(·) ⊙ σ(·) ⊙ softcap(·), bounded to ±100 by Eq 19
3. [b*t, f]
4. down (TSCT) → [b*t, d]

Three things this arm does and does not do for our Fp8 story, to be read as hypotheses for the long run, NOT as results:

1. It bounds the tensor that enters down. The outlier suppression the paper reports is on the activation, and down is the consumer of exactly the activation this arm bounds. That is the interaction point.
2. It does not touch the factors. --quant selects the TSCT factor format (AGENTS.md §3.7); the tanh cap is applied to activations and cannot change max_ortho, the fp32 fallback latch, or anything LinearLike does. A quant claim about this arm is a claim about down’s input range, not about the factors.
3. Our scale conventions are not the paper’s, and β is absolute. normed_ffn is an RMSNorm output, so it is O(1) by construction — the same convention K3 uses before its up-projection (App. B’s local expansion assumes z + O(z³) near the origin, which is a statement about a normalized input). With an O(1) input, Wg x is O(1) only if Wg is O(1)-scaled; at |Wg x| ~ 4 = β₁ the cap is already at tanh(1) = 0.76 of its range, and at |Wg x| ~ 20 the gate gradient is 1.8e-4 of SwiGLU’s. β is therefore not scale-free, and a fixed β against an un-normalized Wg would cap the FFN to a constant. Ours is RMSNormed, so the risk is bounded — but “does our learned Wg stay in the O(1) band” is a measurement nobody has, and it is the thing that decides whether β=4/25 is a good setting for us or a copy of a number that only means something at K3’s scale. Follow-up, not tonight.

5. Gradient health through the cap (gate (b)’s numbers)

tanh saturates, so the honest question is where the cap’s gradient dies. Two measurements, because they answer different questions: the mechanism’s own derivative (f64, from the generator) and the end-to-end gradient on OUR path (situ_glu → the expert’s real down projection → backward, in f32 on the CPU backend, printed by the gate itself with --nocapture).

On OUR path (|d/dinput| / |out|, the numbers loop_block.rs’s situ_gradient_survives_the_cap_where_the_cap_operates prints):

| gate multiple of β₁ (up matched to β₂) | |d/dinput|/|out| | |out| | |—|—|—| | 0.25× | 3.427e+00 | 1.1023 | | 0.5× | 6.518e+00 | 4.7281 | | 1× | 5.647e+00 | 1.4318e+01 | | 2× | 1.103e+00 | 2.3353e+01 | | 5× | 2.887e-03 | 2.5132e+01 |

The cap costs 380× of gradient between 2×β₁ and 5×β₁ (1.103 → 2.887e-3), and at 5× it is still 4 orders of magnitude above f32 eps (1.19e-7), so it is alive there. The mechanism’s own derivative says where it stops being alive.

Gate factor β₁tanh(g/β₁)·σ(g) vs SwiGLU’s g·σ(g) (f64 central differences cross-checked against the closed form, asserted to agree to 1e-5 relative — the f64 CD floor at these magnitudes):

gate g	d/ dg SiTU gate factor	d/ dg Swish	ratio
0.5	7.297e-01	7.400e-01	9.9e-01
1	8.798e-01	9.277e-01	9.5e-01
4 (=β₁)	4.662e-01	1.053e+00	4.4e-01
8 (2β₁)	7.192e-02	1.002e+00	7.2e-02
12 (3β₁)	9.890e-03	1.000e+00	9.9e-03
20 (5β₁)	1.816e-04	1.000e+00	1.8e-04
40 (10β₁)	8.245e-09	1.000e+00	8.2e-09

Up factor β₂tanh(u/β₂) vs SwiGLU’s identity:

up u	d/ du	ratio to 1.0
7.5	9.151e-01	9.2e-01
25 (=β₂)	4.200e-01	4.2e-01
50	7.065e-02	7.1e-02
100 (4β₂)	1.341e-03	1.3e-03
200 (8β₂)	4.501e-07	4.5e-07

Verdict: the cap does not kill the gradient over the range where it is doing its job; it kills it past ~7β. Relative to the output, at the points that matter:

| gate, up | |out| | rel ∂/∂gate | rel ∂/∂up | vs f32 eps 1.19e-07 | |—|—|—|—|—| | 4, 25 | 56.96 | 1.56e-01 | 2.21e-02 | 6–7 orders alive | | 8, 50 | 92.90 | 1.87e-02 | 2.93e-03 | 5–6 orders | | 12, 75 | 99.01 | 2.49e-03 | 3.97e-04 | 4–5 orders | | 20, 100 | 99.92 | 4.54e-05 | 5.37e-05 | ~2.7 orders | | 40, 200 | 100.00 | 2.06e-09 | 1.80e-08 | below f32 eps — dead |

This is the mechanism’s price stated honestly: bounded forward and vanishing gradient are the same fact. The A/B therefore has to answer a question the forward cannot: does Wg x actually stay inside ~5β₁ once trained, or does the arm spend its budget saturating a dead branch? That is the stability/quality trade the --stress protocol (AGENTS.md §3.7) is built to observe.

The analytic derivative used here is independently confirmed by the crate’s own CUDA backward (fused_situ.rs:129-146), which computes exactly (1-tanh²)·σ + β₁·tanh·σ·(1-σ) for the gate and 1-tanh² for the up branch.

The first version of the generator hand-wrote those closed forms and got two of them wrong — a spurious β factor on the tanh term, and σ² where the derivative of a sigmoid is σ(1-σ) — which put the gate gradient at 4.0 where it is 1.8e-4, i.e. exactly the “the cap is harmless” conclusion the bug wanted. The cross-check against central differences is what caught it and it stays in the generator as an assertion. A second version of the same cross-check then fired on the closed form instead: it was the derivative of the gate FACTOR while the numeric one was the derivative of the whole OUTPUT, which differs by the up factor. Comparing factors to factors is what makes the check a check.

6. Does burn-situ’s test suite execute anything? (the burn-rmsnorm question)

Partly — and the part that runs is blind where this arm’s claim lives.

• cargo test --workspace in vendor/dormouse-fused runs default features, and burn-situ’s default is ["std"]. So fused_situ (#[cfg(any(cuda, autodiff))], lib.rs:22-23) is not compiled, and its seam test (fused_situ.rs:405-440) never runs.
• tools/lib_gate.sh has no --features autodiff cell (deliberately: it re-ran everything a second time for 781 s and added zero tests), so that is not an accident of this lane — it is the gate’s shape.
• vendor/dormouse-fused/tools/gpu-gate.sh:73-77 lists burn-situ among the CUDA crates whose GPU tests “have never run in any executed job”.
• What does run: the 4 tests in src/lib.rs:90-154, on ndarray.

Those 4, read as gates:

test	what it actually pins	blind to
softcap_bounded	`	β tanh(x/β)
situ_shape	output dims	the entire form
situ_finite	no NaN/Inf	the entire form
situ_negative_gate_tail_vanishes	sigmoid on the RAW gate — gate=-50 gives 0.0, and sigmoid(capped) would leave −0.269	the up cap, the β scale, the bound

So the crate’s suite separates SiTU from sigmoid(capped) and from nothing else. SwiGLU passes all four (situ_negative_gate_tail_vanishes too: SwiGLU’s gate is also raw). A perturbation of this crate to plain SwiGLU would leave the library green — the AttnRes blind-gate lesson, live in this crate. The new gate in §9(a) is built to be that perturbation’s detector, and it asserts its own separating power rather than only its own output.

7. Why the fused kernel is NOT enabled

fused_situ.rs:105-106 and :129-146 implement Eq (12) and its exact backward, and the seam probes both checkpointing strategies (lib.rs:52-76), so the trainer’s BalancedCheckpointing backend would reach it. It is still off, for two reasons in order:

1. It has never executed anywhere in this repo (§6), and the burn-rmsnorm precedent is a kernel that compiled for its whole life and never ran (AGENTS.md §3.3). Shipping an unexecuted kernel into the training path is the ADR-0019 SILENT class: it would compute something, plausibly, with a counter only in a crate that does not read it.
2. The cost of not having it is ~1% of a step. The tensor path is ~11 elementwise kernels per expert per iteration; at depth 2 × 3 experts that is ~66 extra launches, order 1–2 ms against a 245 ms warm step (§3.1). The arm’s claim is quality and stability, not speed, and §3.1 says the workload is launch-bound — but 66 launches is not where that fact bites.

Follow-up (one task, needs a CUDA gap): enable burn-situ/cuda from dormouse-core, add a CUDA gate that the fused output equals the tensor path on the form fixture, and only then let the trainer reach it. Reading the kernel is not a substitute; that is the whole rmsnorm lesson.

8. A/B row

docs/protocols/AB-PROTOCOL.md row 8, NOT run. 3 seeds, 2k steps, pure CE, one batch size, and the eval window printed on every line (§2.6). The two claims the LONG run has to check, and neither is checkable at 2k steps:

• stability: the NaN-under-overfit class (AGENTS.md §2.3) — needs a run that overfits hard, which 2k steps will not do;
• quantization: whether bounding the down input actually moves Fp8 range overflow / max_ortho drift (§2.3), which is a max_ortho trace, not a BPB.

The row names the width confound (+4.28% params on small) and says the width-matched control is SwiGLU at the same 2f — burn-swiglu is vendored, and it is the paper’s own §2.3.2 comparison, so without it row 8 cannot separate “bounded” from “gated”. The row also carries the measured gradient decay, because that is the question a 2k-step BPB cannot answer.

9. The gates, and what each one is blind to

gate	file	what it fails on
(a) form + self-declared separating power	burn-situ/src/lib.rs	the form drifting from Eq (12); and it asserts the fixture is 0.42–0.99 (symmetric relative) away from SwiGLU and 0.86–1.25 from silu, so a fixture that stopped separating would fail rather than quietly pass
(a2) Eq (18) near-origin	burn-situ/src/lib.rs	SiTU ceasing to be a small perturbation of SwiGLU near zero — green under a SwiGLU perturbation by design, and it is not a cap detector
(a3) the bound is β₁·β₂	burn-situ/src/lib.rs	the output bound not being β’s product, and the bound not being reachable
(b) gradient health	dormouse-core loop_block.rs	a cap whose gradient dies inside the band the arm operates in (floor 1e-6 relative; measured 2.9e-3 at the worst rung)
(c) zero-default identity	dormouse-core	use_situ = false widening the up-projection, or a config lacking the field failing to parse
(d) the arm is observable	dormouse-core	the arm running and being discarded — the GR 9b343d3 defect; the counter is exact at n_experts × iters and 0 when off
(d2) the flag is load-bearing	dormouse-core	situ_glu being computed on a discarded value: on the block’s own weights, SiTU is bounded and the unbounded comparator is not (measured 9.2× the bound)
(e) preset_exec 11/0	existing	the parameter budget moving with the flag off

Red → green, measured

The perturbation was applied to situ_glu’s tensor path and the suite re-run (both patches reverted; the diff is the 4 gates only).

Perturbation 1 — plain SwiGLU, both factors uncapped:

test	result
situ_shape (pre-existing)	ok
situ_finite (pre-existing)	ok
softcap_bounded (pre-existing)	ok
situ_negative_gate_tail_vanishes (pre-existing)	FAILED — but on its saturation row, not its tail row: at β=1 SiTU’s saturated value is 1 and SwiGLU’s is 2500. A scale check wearing a tail check’s name
situ_matches_the_released_reference_and_separates… (new)	FAILED — row 1: 0.7310586 vs 0.7158178, rel 2.085e-02
the_output_bound_is_beta1_times_beta2 (new)	FAILED — row 0 gave 100 000 000 against the bound 100
situ_is_a_small_perturbation_of_swiglu_near_the_origin (new)	ok, and it must be: 2.08e-2 < 5e-2, because Eq (18) says SiTU ≈ SwiGLU near zero. It is a paper claim, not a cap detector

Perturbation 2 — capped gate, UN-capped up (“un-capped GLU”):

test	result
the three structural pre-existing tests	ok
situ_negative_gate_tail_vanishes	FAILED, same reason (2500 vs 1)
form gate (new)	FAILED — row 1: 0.7161996 vs 0.7158178, rel 5.330e-04 — 40× tighter than perturbation 1, which is the point: the up cap is a small effect near the origin and a large one away from it, and a near-origin-only fixture would have missed it
bound gate (new)	FAILED — row 0 gave 40 000 against the bound 100

So: 3 of the 4 pre-existing gates are blind to both perturbations, the 4th catches them for a reason it does not name, and the new gates catch them at the exact number the claim is about. The form gate’s own threshold had to be corrected once by this exercise: the first version claimed ≥0.50 separation at row 4 and the symmetric measure is 0.420, so the gate went red on my own over-claim before it ever saw a bug.

10. Open, not done tonight

• The gate_up field name is a document-vs-code disagreement (a d -> f projection named after two projections). With use_situ it becomes true; with the flag off it stays a lie. docs/glossary.md’s disagreement list is the place to record it; NOT touched by this lane — the glossary is a shared file and two other lanes are adding arms to it tonight.
• A learnable or swept β, and a width-matched SwiGLU control row.
• The fused kernel (§7).
• Whether trained Wg x stays inside ~5β₁ (§4, §5) — needs a long run.
• The situ counter is not on the eval line. §1.1 says an arm must be able to show it ran, and this is the same omission the AttnRes counter already has (§3.3, “owed”). This lane added the counter and the gate on it; printing it is the same one-line change AttnRes is waiting for, and neither lane should do it separately.
• use_situ composes with use_gr and use_attnres and needs no mutual exclusion — it replaces the FFN activation, not the residual accumulation, so config::validate is untouched. Deliberate: a refusal here would be refusing a configuration that has an interpretation.
• use_situ = true with use_tsct = false is legal and costs +51% parameters (a dense gate_up at 2f), against +4.28% under TSCT. No refusal; the A/B row quotes both numbers so a reader does not transfer the cheap one.

11. Verification actually run

what	result
cargo test -p burn-situ (default features, ndarray)	7 passed, 0 failed — 4 pre-existing + 3 new
cargo test -p dormouse-core --lib	70 passed, 0 failed
cargo test -p dormouse-core --test preset_exec	11 passed, 0 failed, 2 ignored, 1976 s — the counts are unchanged, which is the off arm’s parameter-set claim measured on the shipped presets
red→green, perturbation 1 (plain SwiGLU)	3 of 4 pre-existing gates green; both new gates red (§9)
red→green, perturbation 2 (uncapped up)	the same, and the form gate 40× tighter
--features cuda	NOT RUN — a train process held the GPU for the whole window. Nothing in this change touches a CUDA path (the fused kernel is deliberately not enabled and burn-situ/cuda is not in the feature list), so the residual risk is a compile break in the feature graph. Reported as unverified, not as passing.