Skip to content

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).

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 is chunk.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.py has no l2norm at all. The pure-PyTorch references our existing fixture compares against take q/k exactly 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_kernel defaults to False (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] * 2
assert 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”, with rope_theta: float as 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 return None and 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 at fla/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 optional qk_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_rope on KdaModule, off by default, applied to q/k before the L2 norm in project, full-head only, ROPE_THETA = 10000.0 as 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 derived Record is byte-identical with the flag either way and no checkpoint moves.

  • Tier-(a) oracle extension: fla/modules/rotary.py pinned byte-identical (6e9d1051…), rotate_half (:21) and rotary_embedding_ref (:30) extracted from the bytes and exec’d exactly as the four KDA references are. gen_kda_oracle.py asserts the new sha256 before importing anything and the three existing pins re-verify against the clone (--verify-clone). The fixture gained three rope blocks and zero numeric edits to any pre-existing block — every existing red/green verdict in tests/kda_oracle.rs is 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”), and rope_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_answer re-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 own cmp says 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 — and rope_commutes_with_l2norm stayed green, exactly as the comment predicted: a wrong angle is still orthogonal, which is why the commutation test cannot be the gate
    A7 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 exists
    A8 ROPE_THETA 10000 → 1000 rope_rotation_matches_fla_reference, chunked_wy_with_rope_matches_fla_chunk

    All 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.rs is 7 green / 3 red, the three being the pre-existing on-purpose ones, and tests/kda_rope.rs is 6 / 0 — so the generator change moved no existing verdict.

  • Tolerances, not fitted numbers. TOL_ROPE = 1e-5 / ATOL_ROPE = 1e-6 on the rotation (both sides evaluate one cos, one sin, 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-6 inherited unchanged from kda_oracle.rs and for its already-derived reason (the chunked form forms k / exp(cumsum(g)), and over a 16-token tile with g ∈ (-5, 0) the cumsum reaches -80). The margins are printed by the tests and re-printed by falsify.sh’s run_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_ref 1.10e-1 rope_short 2.3e-2 / 2.7e-2, rope_t38 4.9e-2 / 5.0e-2, rope_t64 5.4e-2 / 1.10e-1 (q/k)
    chunked end-to-end vs FLA naive_chunk_kda 1.38e-1 Batched o 3.8e-2 … 1.38e-1, S 1.3e-2 … 6.4e-2; Loop the same o, S 1.1e-2 … 1.28e-1

    So 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_identical compares a default-built module against an explicit-false one on the raw f32 bytes, after handing one the other’s Record.

  • Placement test: rope_commutes_with_l2norm asserts the commutation at our norm (flat H*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:

  1. Two builds of one config are not the same model. Both module tests were red at first because KdaModule::new draws every projection from the device RNG and Device::seed() does not rewind a consumed stream (AGENTS.md §3.7, the two_models_one_seed_are_bit_identical withdrawal). The fix is try_load_record, not a seed.
  2. A permuted tensor’s into_data() is in the permuted order. project returns [B, H, T, D] and apply_rope takes [B, T, H*HD], so comparing a projection’s flat data against apply_rope’s input silently reads the head axis as the time axis. This bit the rotation comparison once and the module identity once.
  • 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.md row 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 in PROVENANCE.md.
  • The partial-rope equivalence is not implemented, so the commutation test only certifies fraction 1.0, and apply_rope takes 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, and docs/protocols/AB-PROTOCOL.md row 8 says so where it is quoted. Plumbing it through dormouse-core’s DormouseConfig and a CLI flag is the owner’s call about whether the arm is a candidate at all.