Op Specs and Model Specs

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

/-- 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

def qwen3_8_27b : Qwen35Hyper where
  vocab_size := 248320
  hidden_size := 5120
  intermediate_size := 17408
  layer_types := cycle 16 [.linear, .linear, .linear, .full]
  num_attention_heads := 24
  num_key_value_heads := 4
  head_dim := 256
  partial_rotary_factor := 1 / 4
  rope_theta := 10000000
  rms_norm_eps := 1e-6
  linear_num_key_heads := 16
  linear_num_value_heads := 48
  linear_key_head_dim := 128
  linear_value_head_dim := 128
  linear_conv_kernel_dim := 4
  tie_word_embeddings := false

$ theformalera spec show 'Qwen/Qwen3.8-27B@1.0:layers[*].mlp.act'
mlp.act   TheFormalEra.Ops.siluMul  ·  64 layer(s)
├── y[t,i] = silu(gate_up[t,i]) · gate_up[t,17408+i]
├── constants: I=17408, act=silu
├── tensors
│   ├── gate_up  in      [T, 34816]  bf16
│   └── y        out     [T, 17408]  bf16
└── shape domain: T ∈ [1, 262144]

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⟩