Module mla

Expand description

2026-09-25: Re-exports the Kimi K3 gated NoPE MLA CPU reference, metrale_model_weights::kimi_k3_host::mla.

Decode runs the CUDA launch in super::mla_cuda instead, unless K3_CUDA_MLA=0.

Owner: model-arch, Kimi K3. Invariants: none beyond the types.

Structs§

apply_output_gate: Apply sigmoid(g) ⊙ attn when the output gate is on; otherwise identity.
cuda_mla_enabled: FullAttention BoundLayer uses CUDA mla_decode unless K3_CUDA_MLA=0. Same polarity as K3_CUDA_KDA. Projections stay on the host either way.
gated_mla_attend: Gated NoPE attend: optionally skip RoPE, SDPA, optional output gate.
maybe_rope: Optional RoPE on the rope slice of a packed [nope | rope] head. NoPE leaves q/k unchanged.
mla_decode_token: One decode token: optional RoPE, append K/V, SDPA, optional output gate. Projections stay in mla_mixer. CUDA k3_mla_* matches this order.
mla_from
sdpa: Causal scaled-dot-product attention. q/k: [T, H, dq], v: [T, H, dv].
sdpa_one: Decode SDPA: one query [H, dq] against cached K/V of length t.