RoPE into the KDA q/k path — findings (2026-09-30)
Lane: wt/rope-kda2, second attempt (the first died silently: no commits, no
findings file, empty worktree). Worktree /home/sehaxe/dormouse-wt/rope-kda2,
branch wt/rope-kda2, off 9d571c2.
Provenance of every source below, one pin:
github.com/fla-org/flash-linear-attention @ 9f38d24980c46d46bd38614e743cdacd21906578
(2026-09-29, HEAD of main at fetch; /tmp/opencode/fla, verified
git rev-parse HEAD = the full sha). arXiv:2510.26692 (Kimi Linear) names
fla/ops/kda as the official KDA implementation in its own footnote 1, and
fla/ops/kda/gate.py:8 carries “This file is modified and supported by the
Moonshot AI Team” — so the ops are Moonshot’s own file, and the layer is
FLA’s official KDA layer. This file is a transcription of that pin: the
sha is named, the line numbers are named, no external reference implementation
is compared numerically here (that is the second half of the lane).
The answer, in one line
Section titled “The answer, in one line”FLA’s official KDA has no RoPE at all — the layer is NoPE, positionally. RoPE in a KDA model comes from the other layers of the hybrid, and our arm is a cross-family transplant, exactly as the Qwen3.8 playbook predicted.
1. Does the official KDA use RoPE? No. Zero occurrences.
Section titled “1. Does the official KDA use RoPE? No. Zero occurrences.”fla/layers/kda.py (316 lines, the whole file) contains no occurrence of
rotary or rope. The class is KimiDeltaAttention (fla/layers/kda.py:28),
and its forward is, in order:
| line | what |
|---|---|
fla/layers/kda.py:248 |
q = F.silu(self.q_proj(hidden_states)) |
fla/layers/kda.py:249 |
k = F.silu(self.k_proj(hidden_states)) |
fla/layers/kda.py:255 |
q, k = (rearrange(x, "... (h d) -> ... h d", d=self.head_k_dim) ...) |
fla/layers/kda.py:262 |
o, recurrent_state = chunk_kda(...) |
fla/layers/kda.py:272, :292 |
use_qk_l2norm_in_kernel=True (chunk and fused_recurrent arms) |
No self.rotary is ever constructed — contrast the layers in the same repo that
do have one: fla/layers/abc.py:141-142 (if self.use_rope: self.rotary = RotaryEmbedding(self.head_k_dim)), fla/layers/attn.py:83,
fla/layers/bitattn.py:74, fla/layers/deltaformer.py:98. GatedDeltaNet also
has no use_rope (grep -n "use_rope" fla/layers/gated_deltanet.py → no
hits), so this is a property of the KDA/GDN family in FLA, not an oversight in
one file.
2. Where the L2 norm lives — and this is the part that matters for placement
Section titled “2. Where the L2 norm lives — and this is the part that matters for placement”The q/k L2 normalization is not in the layer and not in the naive reference; it is a kernel-side step:
fla/ops/kda/chunk.py:56-60—# Apply l2norm/if use_qk_l2norm_in_kernel: q, q_rstd = l2norm_fwd(q); k, k_rstd = l2norm_fwd(k). The backward ischunk.py:184-186(l2norm_bwd), so the rotation we add lands before this point in the dataflow — inside the kernel we are reimplementing.fla/ops/kda/naive.pyhas no l2norm at all. The pure-PyTorch references our existing fixture compares against takeq/kexactly as the layer hands them over. Consequence for the oracle: our forward normalizes outside, so a rope variant must apply the rotation in the same relative position — rope before our L2 norm — or the fixture is comparing two different functions and the red would mean nothing.- FLA’s
use_qk_l2norm_in_kerneldefaults toFalse(fla/ops/kda/chunk.py:40,fla/ops/kda/fused_recurrent.py:253) and the layer turns it on. So the norm is KDA-specific and optional, not part of the recurrence’s algebra.
3. The placement question — rope before or after the L2 norm
Section titled “3. The placement question — rope before or after the L2 norm”Our constraint, stated first: q/k are L2-normalized (K3 stability choice,
AGENTS.md §3.5 item 5: “L2-normalized q/k”). RoPE is a block-diagonal rotation
whose 2×2 blocks are [[cos, -sin], [sin, cos]], so it preserves the L2 norm
exactly.
Therefore, at full-head rotation, the two placements are the same function:
l2norm(rope(x)) = rope(x)/‖rope(x)‖ = rope(x)/‖x‖ = rope(l2norm(x))This is arithmetic, not a citation, and it is pinned by a test in this lane
(rope_commutes_with_l2norm_at_full_fraction). It is the answer to the
question as asked: for a full-head rope the placement is free, and no source
can be said to “prefer” one.
Where the placement stops being free: a partial rotation. FLA’s own
reference supports one — rotary_embedding_ref (fla/modules/rotary.py:30-36):
ro_dim = cos.shape[-1] * 2assert ro_dim <= x.shape[-1]...return torch.cat([x[..., :ro_dim] * cos + rotate_half(x[..., :ro_dim], interleaved) * sin, x[..., ro_dim:]], -1)The un-rotated tail x[..., ro_dim:] is carried through untouched, and once the
L2 norm is over all head dims it mixes rotated and unrotated coordinates — so
l2norm(rope_partial(x)) ≠ rope_partial(l2norm(x)). FLA treats that fraction as
an explicit knob and keeps the rotation before the scan:
fla/layers/mamba3.py:60 (“blockwise rotary on Q/K before the SSM scan”),
fla/layers/mamba3.py:71 (rope_fraction: float = 0.5),
fla/layers/mamba3.py:129-130 (refuses anything but 0.5 or 1.0).
The one upstream site that has both a rope and a per-head norm on the same
q/k is the ABC layer family, and it applies the rope to the raw,
un-normalized q/k: fla/layers/abc.py:202 (head rearrange) →
fla/layers/abc.py:204 (if self.use_rope:) → fla/layers/abc.py:208
(q, k = self.rotary(q, k, seqlen_offset=seqlen_offset)). use_rope defaults
to True there (fla/layers/abc.py:43). ABC attention has no
use_qk_l2norm_in_kernel flag, so this is not literally a rope-vs-norm
ordering — but it is upstream’s convention: rotate first, normalize never
(because the norm is the kernel’s), and let the kernel normalize after.
Our choice: rope before the L2 norm, at fraction 1.0. It matches the
“before the scan” convention (mamba3.py:60) and the ABC placement
(abc.py:202-208), and at fraction 1.0 it is provably the same function as the
other order — so the choice costs nothing and the equivalence is a test rather
than a belief. Fraction is not wired; a future partial rope must re-derive the
order rather than inherit it.
4. What Kimi Linear’s hybrid does instead — the transplant, cited
Section titled “4. What Kimi Linear’s hybrid does instead — the transplant, cited”FLA ships the hybrid machinery, and it is the answer to “where does position come from in a KDA model”:
fla/models/hybrid.py:17-23—HybridAttentionSpec: “JSON-serializable settings for standard attention at selected model layers”, withrope_theta: floatas one of the fields (:23), defaulted and validated at:103-115(10000.).fla/models/hybrid.py:190-200—get_hybrid_attention_spec’s docstring: “Unassigned layers returnNoneand retain their model’s native mixer.” That is the NoPE statement for the KDA half of a hybrid, in upstream’s own words.- The layers that are assigned get their rope from the standard-attention
layer:
fla/layers/attn.py:83—self.rotary = RotaryEmbedding(dim=self.head_dim, base=self.rope_theta), applied atfla/layers/attn.py:125(q, k = self.rotary(q, k, seqlen_offset=seqlen_offset, max_seqlen=..., cu_seqlens=...)). Note it is applied to the raw q/k, before the optionalqk_norm(fla/layers/attn.py:51-53, an RMSNorm, off by default) — upstream’s habit again: rotate, then normalize.
So: KDA layers NoPE, interleaved full-attention layers RoPE, rope_theta
10000. The transplant into a pure-KDA loop is the interesting move, and it is
required regardless of pretrain parity: AGENTS.md §3.5 item 5 records the
playbook finding as “NoPE → endless generation after post-training”, and the
owner’s goal is the full SFT→RLVR→distill loop (docs/architecture/post-training.md). A model
that cannot be steered after pretraining is not a candidate for that pipeline
whatever it scores at step 2000. This is the justification, and it is a
cross-family transplant, not a Kimi Linear reproduction.
5. RoPE is parameter-free — from the source, not from our arithmetic
Section titled “5. RoPE is parameter-free — from the source, not from our arithmetic”RotaryEmbedding holds inv_freq and scale as non-persistent buffers:
fla/modules/rotary.py:379 (self.register_buffer("inv_freq", ..., persistent=False)) and :384 (same for scale). The frequency formula is
fla/modules/rotary.py:410-414 — 1.0 / base ** (arange(0, dim, 2) / dim),
base = rope_theta. Zero nn.Parameter, zero persistent state, so:
- the arm adds no parameters and changes no param count;
- no checkpoint format change (nothing to save, nothing to version);
- nothing to load in an old checkpoint, nothing that breaks one.
That is a property of the mechanism, stated so a reader does not have to take our word for it: rotary carries no weights because the rotation is a function of position, not a learned parameter.
6. The oracle, the wiring, the gate — what was built
Section titled “6. The oracle, the wiring, the gate — what was built”-
Wiring:
KdaConfig::use_rope+#[module(skip)] pub use_ropeonKdaModule, off by default, applied to q/k before the L2 norm inproject, full-head only,ROPE_THETA = 10000.0as a named const (upstream’s default, two citations in §4). Zero parameters — rotary is parameter-free per §5, the flag is#[module(skip)], so the derivedRecordis byte-identical with the flag either way and no checkpoint moves. -
Tier-(a) oracle extension:
fla/modules/rotary.pypinned byte-identical (6e9d1051…),rotate_half(:21) androtary_embedding_ref(:30) extracted from the bytes andexec’d exactly as the four KDA references are.gen_kda_oracle.pyasserts the new sha256 before importing anything and the three existing pins re-verify against the clone (--verify-clone). The fixture gained threeropeblocks and zero numeric edits to any pre-existing block — every existing red/green verdict intests/kda_oracle.rsis unmoved. -
The three shapes, and why each exists:
rope_short(T=16, one tile, no boundary),rope_t38(T=38 at our BT=16, so the last chunk is zero-padded —burn-gdn2/src/forward.rs:177-178— against a BT=2 reference that is nearly exact, which makes the measured difference our f32 noise at a non-multiple T; and the only two-head case, which is what separates “rotate within each head” from “rotate across the head boundary”), andrope_t64(T=64, three exact boundaries — the multiple-of-BT case the padding path does not cover). -
Vacuity guards 4a/4b, both in the generator and both firing. 4a: the rotation must move q/k by ≥ 1e-2 (measured 4.7-5.5). 4b is the one that matters: the OUTPUT must move, or a no-op rope would pass — measured 0.24 / 0.69 / 0.43 relative across the three cases, against a chunked tolerance of 1e-3.
tests/kda_rope.rs::rope_fixture_separates_from_the_unrotated_answerre-asserts 4b on the Rust side, because a guard that only exists in the generator cannot see a fixture edited by hand. -
Red→green, three mutants, and it was RUN (2026-10-01,
tests/oracle/falsify.sh, whole script; the tree is restored byte-identical at the end and the script’s owncmpsays so):mutant what it perturbs what went red A6 the angle’s frequency loses its factor of 2 rope_rotation_matches_fla_reference,chunked_wy_with_rope_matches_fla_chunk— andrope_commutes_with_l2normstayed green, exactly as the comment predicted: a wrong angle is still orthogonal, which is why the commutation test cannot be the gateA7 the rotation is applied to q and not k only module_projects_both_q_and_k_rotated; both FLA-comparison tests stayed green, which is the whole reason that test existsA8 ROPE_THETA10000 → 1000rope_rotation_matches_fla_reference,chunked_wy_with_rope_matches_fla_chunkAll three behaved as their in-place reasoning predicted, which is the form of the demonstration worth having: the comment and the observation agree.
-
Baselines unchanged.
tests/kda_oracle.rsis 7 green / 3 red, the three being the pre-existing on-purpose ones, andtests/kda_rope.rsis 6 / 0 — so the generator change moved no existing verdict. -
Tolerances, not fitted numbers.
TOL_ROPE = 1e-5/ATOL_ROPE = 1e-6on the rotation (both sides evaluate onecos, onesin, two products and a difference on the same f32 inputs, so a few ulp is the whole expectation and 1e-5 is ~85 ulp);TOL_CHUNK = 1e-3/ATOL_CHUNK = 1e-6inherited unchanged fromkda_oracle.rsand for its already-derived reason (the chunked form formsk / exp(cumsum(g)), and over a 16-token tile withg ∈ (-5, 0)the cumsum reaches -80). The margins are printed by the tests and re-printed byfalsify.sh’srun_rope, not asserted as literals: the margin is an observation that moves with the backend, and the assertion is the normalised ratio against a derived tolerance.Measured, 2026-10-01, ndarray (worst normalised ratio of the six comparisons in each arm; the assertion is
<= 1.0):arm worst across rotation vs FLA rotary_embedding_ref1.10e-1 rope_short2.3e-2 / 2.7e-2,rope_t384.9e-2 / 5.0e-2,rope_t645.4e-2 / 1.10e-1 (q/k)chunked end-to-end vs FLA naive_chunk_kda1.38e-1 Batched o3.8e-2 … 1.38e-1,S1.3e-2 … 6.4e-2; Loop the sameo,S1.1e-2 … 1.28e-1So both arms sit roughly 7-9x inside their tolerance, and the rotation’s 1.10e-1 of a 1e-5 relative tolerance is ~1.1e-6 absolute on an O(1) quantity — about 9 ulp, which is the “a few ulp” the class predicts and not the “far above” the decay gate reports. The two arms are not the same comparison and are not given one number: the rotation is formula-vs-formula and the chunked path is two evaluation orders of one algebra in f32.
-
Zero-default identity:
rope_off_is_bitwise_identicalcompares a default-built module against an explicit-falseone on the raw f32 bytes, after handing one the other’sRecord. -
Placement test:
rope_commutes_with_l2normasserts the commutation at our norm (flatH*HD,burn-gdn2/src/l2norm.rs:4-7) and at the reference’s per-head norm, plus the mechanism on its own (the rotation preserves each head’s norm).
Two bugs this lane’s own tests had, and what they are worth
Section titled “Two bugs this lane’s own tests had, and what they are worth”Both were found by the gate, not by reading, and both are the same defect class
as kda_oracle.rs’s lesson that a self-comparison cannot see a choice:
- Two builds of one config are not the same model. Both module tests were
red at first because
KdaModule::newdraws every projection from the device RNG andDevice::seed()does not rewind a consumed stream (AGENTS.md §3.7, thetwo_models_one_seed_are_bit_identicalwithdrawal). The fix istry_load_record, not a seed. - A permuted tensor’s
into_data()is in the permuted order.projectreturns[B, H, T, D]andapply_ropetakes[B, T, H*HD], so comparing a projection’s flat data againstapply_rope’s input silently reads the head axis as the time axis. This bit the rotation comparison once and the module identity once.
7. What this lane does NOT establish
Section titled “7. What this lane does NOT establish”- No A/B. Pretrain parity is unmeasured and unmeasurable from here; the
justification in §4 is a post-training argument, not a loss curve. The
docs/protocols/AB-PROTOCOL.mdrow 8 is written and not run. - No gradient-flow claim for the rope arm. The fused chunked backward is a
separate question in this project. A rope that is not differentiated through
changes nothing; a rope that is supposed to be differentiated and is not is
the same class of defect as
8fa5d4c. The oracle here is a forward oracle and says so inPROVENANCE.md. - The partial-rope equivalence is not implemented, so the commutation test
only certifies fraction 1.0, and
apply_ropetakes no fraction. - No GVA case in the fixture (all three have
HV == H). The flag sits above the GVA repeat, which is correct — q/k are on the key-head axis there — but nothing in the fixture checks that placement. - Not wired to the trainer’s config or CLI. The flag is on
KdaConfig, anddocs/protocols/AB-PROTOCOL.mdrow 8 says so where it is quoted. Plumbing it throughdormouse-core’sDormouseConfigand a CLI flag is the owner’s call about whether the arm is a candidate at all.