Module attention 

Rotary position embedding (RoPE, both the split-half apply_rope and interleaved apply_rope_interleaved conventions) and grouped-query causal attention (GQA). This is the “vanilla” attention path used as the correctness baseline. causal_mla_attention/causal_mla_attention_sparse add DeepSeek-style latent attention and its DSA sparse-selection variant (GLM-5.2, DeepSeek V3.2/V4); both mechanisms are now backed by real, public reference implementations (see docs/MODELS.md). ferrox_models::mla/ferrox_models::glm_dsa compose these primitives into full RoPE-carrying MLA forward passes.

Structs

YarnScaling

YaRN RoPE scaling exactly as a checkpoint declares it, in the shape the reference reads out of rope_scaling (FreeToken python/freetoken/layers/rotary.py:139, the "yarn" arm of _get_rope). beta_fast / beta_slow / truncate carry that arm’s own defaults, because a real YaRN checkpoint usually declares only factor and original_max_position_embeddings.

Functions

apply_rope

Applies rotary position embedding in place to a single head’s vector, split-half (GPT-NeoX / LLAMA_ROPE_TYPE_NEOX) style: each pair (i, i+half) is rotated together, for position pos with base theta. This is what llama.cpp calls NEOX-style RoPE (used by e.g. DeepSeek-V3.2’s lightning indexer); see apply_rope_interleaved for the other real convention.

apply_rope_back

Inverse of apply_rope (split-half / NeoX): rotates each pair by -angle. DeepSeek V4 applies this (“derope” / ggml_rope_ext_back) to the rope slice of attention output before the grouped wo_a projection — see .scratch/NOTES_DS4_INFERENCE.md.

apply_rope_interleaved

Applies rotary position embedding in place, interleaved (GPT-J / llama.cpp’s LLAMA_ROPE_TYPE_NORM) style: adjacent pairs (2*i, 2*i+1) are rotated together, rather than apply_rope’s split-half pairing. GLM-5.2 uses this convention for both its main attention (rope_interleave: true) and its lightning indexer (indexer_rope_interleave: true) per its real config.json (huggingface.co/zai-org/GLM-5.2) — confirmed against llama.cpp PR #25407, which rotates the indexer with LLAMA_ROPE_TYPE_NORM where DeepSeek-V3.2’s PR #23346 uses LLAMA_ROPE_TYPE_NEOX.

apply_rope_interleaved_back

Inverse of apply_rope_interleaved (adjacent-pair / Norm RoPE).

apply_rope_interleaved_with_freq_factors

Interleaved (GPT-J / LLAMA_ROPE_TYPE_NORM) RoPE with Llama 3/3.1/3.2’s per-band frequency correction – the combination real llama.cpp uses for general.architecture = "llama" checkpoints that carry rope_freqs.weight. Pairing is adjacent (2*i, 2*i+1) as in apply_rope_interleaved; each band’s angle is divided by freq_factors[i] as in apply_rope_with_freq_factors. freq_factors all-1.0 is mathematically identical to plain apply_rope_interleaved (pinned by rope_interleaved_with_all_ones_freq_factors_matches_plain_interleaved).

apply_rope_with_freq_factors

Same split-half rotation as apply_rope, but each frequency band i has its angle divided by freq_factors[i] before the rotation – Llama 3/3.1/3.2’s real per-band RoPE frequency correction (the rope_freqs.weight GGUF tensor, n_rot/2 elements, TENSOR_NOT_REQUIRED so most non-Llama-3 checkpoints don’t carry it). Confirmed against real llama.cpp source, not guessed: ggml_rope_cache_init (ggml/src/ggml-cpu/ops.cpp) computes theta/freq_factors[i0/2] per band before rope_yarn. freq_factors all-1.0 is mathematically identical to plain apply_rope (pinned by rope_with_all_ones_freq_factors_matches_plain_rope).

causal_gqa_attention

Single-token causal attention for one query against all cached key/value positions (0..=pos), grouped-query style: n_kv_heads may be fewer than n_heads, with each KV head shared by n_heads / n_kv_heads query heads.

causal_gqa_attention_paged

Same math as causal_gqa_attention, but K/V positions are read through a PagedKvStore block table instead of one contiguous slice: position t lives in block block_table[t / block_size] at offset t % block_size, so blocks need not be physically adjacent or in order. Must match causal_gqa_attention given the same logical K/V contents (float noise only) — the block table is a storage-layout detail, not a math change.

causal_gqa_attention_paged_sinks

causal_gqa_attention_paged with per-head attention sinks and an optional sliding window: the paged twin of causal_gqa_attention_sinks.

causal_gqa_attention_prefill

Prefill (multi-query) causal GQA: q/k_cache/v_cache are all length seq_len in the time dimension. Query at position t attends only to keys/values 0..=t (same math as looping causal_gqa_attention per token). Layout: q/out [seq_len, n_heads, head_dim]; k/v [seq_len, n_kv_heads, head_dim]. Metal prefill kernels must match.

causal_gqa_attention_prefill_shared_kv

Prefill attention parallelized over (query, head) slots. Same math as calling causal_gqa_attention_softcap per query; used by the decoder CPU pp path so Rayon owns the full [n_q × n_heads] grid instead of only the query axis (better for large-head models like Phi-4).

causal_gqa_attention_prefill_shared_kv_windowed

causal_gqa_attention_prefill_shared_kv with an optional sliding window, so SWA models (Gemma-2/3, Mistral, Qwen2-MoE) get the same blocked kernel instead of the per-query causal_gqa_attention_windowed_softcap fallback.

causal_gqa_attention_sinks

Single-query causal GQA with per-head attention sinks, optionally windowed.

causal_gqa_attention_softcap

causal_gqa_attention with optional Gemma-2 attention logit softcap.

causal_gqa_attention_windowed

Same computation as causal_gqa_attention, but each query only attends to the last window cached positions (inclusive of itself) instead of the full causal history – Mistral/Mixtral/Qwen2-family sliding-window attention. Confirmed against the real sliding_window config field used by those models (real transformers source for Qwen2MoeAttention/Mixtral’s equivalent) and against candle-transformers’ mixtral.rs/qwen2_moe.rs, which both mask scores where key_pos + sliding_window < query_pos – i.e. only the most recent window positions (including the query’s own) stay unmasked. window >= seq_len degenerates to exactly causal_gqa_attention’s full-causal behavior (pinned by windowed_attention_with_window_covering_full_history_matches_full_causal).

causal_gqa_attention_windowed_softcap

causal_gqa_attention_windowed with optional attention logit softcap.

causal_mla_attention

Single-token causal attention for DeepSeek/Kimi-style Multi-head Latent Attention (MLA): every query head has its own key/value (no GQA-style grouping – verified directly against Kimi K3’s real KimiMLAAttention.forward, where kv_b_proj expands to the full num_heads count and the num_key_value_heads/num_key_value_groups fields computed in __init__ go unused), but the key/query head dimension (qk_head_dim = qk_nope_head_dim + qk_rope_head_dim) can differ from the value head dimension (v_head_dim) – unlike causal_gqa_attention, which assumes one shared head_dim for both.

causal_mla_attention_sinks

causal_mla_attention with DeepSeek V4’s per-head attention sinks.

causal_mla_attention_sparse

Same as causal_mla_attention for a single query position, but attention is restricted to the explicit visible key positions (ascending, a subset of 0..seq_len) rather than the full causal history — the sparse-attention half of GLM-5.2/DeepSeek-V3.2’s DSA, applied after lightning_indexer_topk selects visible.

causal_mla_attention_sparse_sinks

causal_mla_attention_sparse with per-head attention sinks.

lightning_indexer_topk

The DeepSeek-V3.2 / GLM-5.2 “lightning indexer” (arXiv 2512.02556; real, merged, tested reference implementations in llama.cpp PR #23346 and PR #25407): scores every causally-visible key position against the query using a cheap multi-head dot-product indexer, then keeps only the top_k highest-scoring positions.

proportional_freq_factors

The reference’s "proportional" arm (rotary.py:103), in the same per-band divisor form as yarn_freq_factors.

yarn_correction_range

The [low, high] band range the YaRN ramp interpolates across, as the reference computes it (rotary.py:167-179): both ends from [yarn_correction_dim], floored / ceiled when truncate, low clamped up to 0, and – the load-bearing detail, called out in the reference’s own comment at rotary.py:176 – high clamped to rotary_dim - 1, not rotary_dim / 2 - 1.

yarn_freq_factors

YaRN’s frequency rewrite, expressed as the per-band divisors apply_rope_with_freq_factors already consumes: one entry per rotation band (rotary_dim / 2), each the number the band’s RoPE angle is divided by.