Module mla

Module mla 

Source
Expand description

Gated NoPE MLA CPU reference.

Do not reuse qwen3_attention blindly: K3 sets mla_use_nope=true while still allocating qk_rope_head_dim slots, and mla_use_output_gate=true applies a full-rank g_proj sigmoid gate on the attention output.

RoPE dims stay in the head (prod 128 nope + 64 rope = 192). NoPE means those slots are not rotated, not that they are dropped.

CUDA decode is kernels/gb10/kimi-k3/bf16/mla_decode.cu (k3_mla_*). BoundLayer FullAttention uses CUDA mla_decode unless K3_CUDA_MLA=0.

Structs§

MlaConfig

Functions§

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.