Qwen/Qwen3.8-27B @1.0.0

qwen3_5 · 64 layers · hidden 5120 · 26.90 B parameters · text-only language-model path; the vision encoder and the MTP head are out of scope

Trust

Compute graph

Domains

Theorems (Lean kernel-checked)

Versions

Advisories

No advisories.

Lean 4 definitions

TheFormalEra/Ops/Attention.lean

import TheFormalEra.Ops.Basic

/-!
# Op Specs: softmax attention
-/

namespace TheFormalEra.Ops

open Finset

/-- Softmax over the reals, normalized only over visible positions -/
noncomputable def softmax {L : ℕ} (s : Fin L → ℝ) (visible : Fin L → Prop)
    [DecidablePred visible] : Fin L → ℝ :=
  fun l => if visible l then
    Real.exp (s l) / ∑ m ∈ Finset.univ.filter visible, Real.exp (s m) else 0

/-- Causal GQA: each KV head serves G query heads; pos s is the position of query s -/
noncomputable def causalGQA {S L Hkv G D : ℕ} (scale : ℝ) (pos : Fin S → ℕ)
    (q : Fin S → Fin Hkv → Fin G → Fin D → ℝ)
    (k v : Fin L → Fin Hkv → Fin D → ℝ) : Fin S → Fin Hkv → Fin G → Fin D → ℝ :=
  fun s c g d =>
    let logit : Fin L → ℝ := fun l => scale * ∑ e, q s c g e * k l c e
    ∑ l, softmax logit (fun l => l.val ≤ pos s) l * v l c d

/-- Sliding-window causal GQA: query at position `p` sees keys `l` with `p - w < l ≤ p`. -/
noncomputable def slidingGQA {S L Hkv G D : ℕ} (scale : ℝ) (window : ℕ) (pos : Fin S → ℕ)
    (q : Fin S → Fin Hkv → Fin G → Fin D → ℝ)
    (k v : Fin L → Fin Hkv → Fin D → ℝ) : Fin S → Fin Hkv → Fin G → Fin D → ℝ :=
  fun s c g d =>
    let logit : Fin L → ℝ := fun l => scale * ∑ e, q s c g e * k l c e
    ∑ l, softmax logit (fun l => l.val ≤ pos s ∧ pos s < l.val + window) l * v l c d

/-- Decode attention for one sequence of `L` cached tokens: the output of the query at the last
position `L - 1` (every cached key is visible).  Query head `c·G + g` reads KV head `c`. -/
noncomputable def decodeGQA {L Hkv G D : ℕ} (scale : ℝ) (q : Fin Hkv → Fin G → Fin D → ℝ)
    (k v : Fin L → Fin Hkv → Fin D → ℝ) : Fin Hkv → Fin G → Fin D → ℝ :=
  fun c g d => causalGQA (S := 1) scale (fun _ => L - 1) (fun _ => q) k v 0 c g d

/-- Sigmoid output gate of Gated Attention: `o ⊙ σ(gate)`. -/
noncomputable def sigmoidGate {ι : Type*} {N : ℕ} (o gate : ι → Fin N → ℝ) : ι → Fin N → ℝ :=
  fun t i => o t i * sigmoid (gate t i)

end TheFormalEra.Ops

TheFormalEra/Ops/DeltaNet.lean

import Mathlib.Data.Matrix.Mul
import TheFormalEra.Ops.Basic

/-!
# Op Specs: the gated delta rule
-/

namespace TheFormalEra.Ops

open Matrix

/-- State of the gated delta rule (single head):
    S_t = α_t · S_{t-1} · (I − β_t k_t k_tᵀ) + β_t v_t k_tᵀ -/
noncomputable def gdnState {T Dk Dv : ℕ} (α β : Fin T → ℝ)
    (k : Fin T → Fin Dk → ℝ) (v : Fin T → Fin Dv → ℝ) :
    (t : ℕ) → t ≤ T → Matrix (Fin Dv) (Fin Dk) ℝ
  | 0, _ => 0
  | t + 1, h =>
    let i : Fin T := ⟨t, h⟩
    α i • (gdnState α β k v t (Nat.le_of_succ_le h) * (1 - β i • Matrix.vecMulVec (k i) (k i)))
      + β i • Matrix.vecMulVec (v i) (k i)

/-- Output at step t: o_t = S_t q_t -/
noncomputable def gatedDeltaRule {T Dk Dv : ℕ} (α β : Fin T → ℝ)
    (q k : Fin T → Fin Dk → ℝ) (v : Fin T → Fin Dv → ℝ) : Fin T → Fin Dv → ℝ :=
  fun i => (gdnState α β k v (i.val + 1) i.isLt).mulVec (q i)

/-- One recurrent (decode) step from a state `S`: the new state. -/
noncomputable def gdnStep {Dk Dv : ℕ} (α β : ℝ) (S : Matrix (Fin Dv) (Fin Dk) ℝ)
    (k : Fin Dk → ℝ) (v : Fin Dv → ℝ) : Matrix (Fin Dv) (Fin Dk) ℝ :=
  α • (S * (1 - β • Matrix.vecMulVec k k)) + β • Matrix.vecMulVec v k

/-- The state recurrence is the iterated decode step. -/
theorem gdnState_succ {T Dk Dv : ℕ} (α β : Fin T → ℝ) (k : Fin T → Fin Dk → ℝ)
    (v : Fin T → Fin Dv → ℝ) (t : ℕ) (h : t + 1 ≤ T) :
    gdnState α β k v (t + 1) h =
      gdnStep (α ⟨t, h⟩) (β ⟨t, h⟩) (gdnState α β k v t (Nat.le_of_succ_le h)) (k ⟨t, h⟩) (v ⟨t, h⟩) := by
  rfl

/-- Depthwise causal convolution of width `K` followed by SiLU (zero left padding). -/
noncomputable def causalConvSilu {T C K : ℕ} (w : Fin C → Fin K → ℝ) (x : Fin T → Fin C → ℝ) :
    Fin T → Fin C → ℝ :=
  fun t c => silu (∑ j : Fin K, if h : K - 1 ≤ t.val + j.val ∧ t.val + j.val - (K - 1) < T then
      w c j * x ⟨t.val + j.val - (K - 1), h.2⟩ c else 0)

end TheFormalEra.Ops

TheFormalEra/Ops/Linear.lean

import TheFormalEra.Ops.Basic

/-!
# Op Specs: linear maps, feed-forward activations, residuals, embeddings
-/

namespace TheFormalEra.Ops

open Finset

variable {ι : Type*}

/-- `y = x Wᵀ` (bias-free `nn.Linear(K, N)`; `W` is stored `[N, K]`). -/
noncomputable def linear {K N : ℕ} (W : Fin N → Fin K → ℝ) (x : ι → Fin K → ℝ) : ι → Fin N → ℝ :=
  fun t n => ∑ k, x t k * W n k

/-- The fused gate/up projection: `[x W_gateᵀ ‖ x W_upᵀ]`, gate first. -/
noncomputable def gateUp {K I : ℕ} (Wg Wu : Fin I → Fin K → ℝ) (x : ι → Fin K → ℝ) :
    ι → Fin (2 * I) → ℝ :=
  fun t n => if h : n.val < I then linear Wg x t ⟨n.val, h⟩
    else linear Wu x t ⟨n.val - I, by omega⟩

/-- SwiGLU activation of a fused `[gate ‖ up]` row: `silu(gate) ⊙ up`. -/
noncomputable def siluMul {I : ℕ} (gu : ι → Fin (2 * I) → ℝ) : ι → Fin I → ℝ :=
  fun t i => silu (gu t ⟨i.val, by omega⟩) * gu t ⟨I + i.val, by omega⟩

/-- GeGLU (tanh) activation of a fused `[gate ‖ up]` row. -/
noncomputable def geluMul {I : ℕ} (gu : ι → Fin (2 * I) → ℝ) : ι → Fin I → ℝ :=
  fun t i => geluTanh (gu t ⟨i.val, by omega⟩) * gu t ⟨I + i.val, by omega⟩

/-- Residual connection with multiplier `r` (`r = 1` except for Granite). -/
noncomputable def residual {H : ℕ} (r : ℝ) (x h : ι → Fin H → ℝ) : ι → Fin H → ℝ :=
  fun t i => x t i + r * h t i

/-- Token embedding with multiplier `c`. -/
noncomputable def embed {V H T : ℕ} (c : ℝ) (E : Fin V → Fin H → ℝ) (tokens : Fin T → Fin V) :
    Fin T → Fin H → ℝ :=
  fun t i => c * E (tokens t) i

/-- Language-model head with logit multiplier and optional soft-capping. -/
noncomputable def lmHead {V H : ℕ} (mul : ℝ) (cap : Option ℝ) (W : Fin V → Fin H → ℝ)
    (x : ι → Fin H → ℝ) : ι → Fin V → ℝ :=
  fun t v => match cap with
    | none => mul * linear W x t v
    | some c => softcap c (mul * linear W x t v)

end TheFormalEra.Ops

TheFormalEra/Ops/Norm.lean

import TheFormalEra.Ops.Basic

/-!
# Op Specs: normalisation

Every norm normalises the last axis of a tensor whose rows are indexed by an arbitrary type
`ι` (tokens `Fin T`, or `(token, head)` pairs for per-head norms).
-/

namespace TheFormalEra.Ops

open Finset

variable {ι : Type*} {H : ℕ}

/-- `√(mean(x²) + ε)` of one row. -/
noncomputable def rms (eps : ℝ) (x : Fin H → ℝ) : ℝ := Real.sqrt ((∑ j, x j ^ 2) / (H : ℝ) + eps)

/-- RMSNorm with the weight applied as `w` (Qwen3, Granite, Gemma 4). -/
noncomputable def rmsnorm (eps : ℝ) (x : ι → Fin H → ℝ) (w : Fin H → ℝ) : ι → Fin H → ℝ :=
  fun t h => x t h / Real.sqrt ((∑ j, x t j ^ 2) / (H : ℝ) + eps) * w h

/-- RMSNorm in the Qwen3.5 family: the weight is applied as (1 + w) -/
noncomputable def rmsnormOffset (eps : ℝ) (x : ι → Fin H → ℝ) (w : Fin H → ℝ) : ι → Fin H → ℝ :=
  fun t h => x t h / Real.sqrt ((∑ j, x t j ^ 2) / (H : ℝ) + eps) * (1 + w h)

/-- RMSNorm without a weight (Gemma 4 value norm, Muse Glimmer q/k and embedding norms). -/
noncomputable def rmsnormUnit (eps : ℝ) (x : ι → Fin H → ℝ) : ι → Fin H → ℝ :=
  fun t h => x t h / Real.sqrt ((∑ j, x t j ^ 2) / (H : ℝ) + eps)

/-- Gated RMSNorm of the Gated DeltaNet output: `w ⊙ x̂ ⊙ silu(z)`. -/
noncomputable def rmsnormGated (eps : ℝ) (x z : ι → Fin H → ℝ) (w : Fin H → ℝ) : ι → Fin H → ℝ :=
  fun t h => w h * (x t h / Real.sqrt ((∑ j, x t j ^ 2) / (H : ℝ) + eps)) * silu (z t h)

/-- L2 normalisation `x / √(Σ x² + ε)` (queries and keys of the delta rule). -/
noncomputable def l2norm (eps : ℝ) (x : ι → Fin H → ℝ) : ι → Fin H → ℝ :=
  fun t h => x t h / Real.sqrt (∑ j, x t j ^ 2 + eps)

end TheFormalEra.Ops

TheFormalEra/Ops/Rope.lean

import TheFormalEra.Ops.Basic

/-!
# Op Specs: rotary position embeddings

Rotate-half RoPE on the first `rot` coordinates of every head (`rot = D` for full RoPE,
`rot = 64` of `256` for Qwen3.8's partial RoPE), with inverse frequencies
`ω_i = θ^(-2i / denom)` for `i < active` and `ω_i = 0` beyond (Gemma 4's proportional RoPE
rotates the whole head but keeps only the first `active` frequencies).
-/

namespace TheFormalEra.Ops

variable {ι : Type*} {D : ℕ}

/-- Inverse frequency of pair `i`. -/
noncomputable def invFreq (theta : ℝ) (denom active i : ℕ) : ℝ :=
  if i < active then theta ^ (-(2 * (i : ℝ) / denom)) else 0

/-- RoPE at position `pos t` of row `t`; pairs `(i, i + rot/2)` for `i < rot/2`. -/
noncomputable def rope (theta : ℝ) (rot denom active : ℕ) (pos : ι → ℕ) (x : ι → Fin D → ℝ) :
    ι → Fin D → ℝ :=
  fun t d =>
    let half := rot / 2
    if d.val < half then
      let a := (pos t : ℝ) * invFreq theta denom active d.val
      x t d * Real.cos a - (if h : d.val + half < D then x t ⟨d.val + half, h⟩ else 0) * Real.sin a
    else if d.val < 2 * half then
      let a := (pos t : ℝ) * invFreq theta denom active (d.val - half)
      x t d * Real.cos a + x t ⟨d.val - half, by omega⟩ * Real.sin a
    else x t d

/-- Standard (partial) RoPE over the first `rot` coordinates. -/
noncomputable def ropeStd (theta : ℝ) (rot : ℕ) (pos : ι → ℕ) (x : ι → Fin D → ℝ) : ι → Fin D → ℝ :=
  rope theta rot rot (rot / 2) pos x

end TheFormalEra.Ops