Struct MlaWeights 

pub struct MlaWeights {

    pub wq_a: DenseWeight,
    pub wq_a_nvfp4: Option<QuantizedWeight>,
    pub wq_a_fp8: Option<Fp8Weight>,
    pub wq_b: DenseWeight,
    pub wq_b_nvfp4: Option<QuantizedWeight>,
    pub wq_b_fp8: Option<Fp8Weight>,
    pub q_a_norm: DenseWeight,
    pub wkv_a: DenseWeight,
    pub wkv_a_nvfp4: Option<QuantizedWeight>,
    pub wkv_a_fp8: Option<Fp8Weight>,
    pub wkv_b: DenseWeight,
    pub kv_a_norm: DenseWeight,
    pub wkv_a_rope: DenseWeight,
    pub wkv_a_merged: DenseWeight,
    pub wo: DenseWeight,
    pub wo_nvfp4: Option<QuantizedWeight>,
    pub wo_a: DenseWeight,
    pub wo_a_nvfp4: Option<QuantizedWeight>,
    pub wo_a_fp8: Option<Fp8Weight>,
    pub wo_b: DenseWeight,
    pub wo_b_nvfp4: Option<QuantizedWeight>,
    pub wo_b_fp8: Option<Fp8Weight>,
    pub w_uk_t: DenseWeight,
    pub w_uv: DenseWeight,
    pub wq_b_rope: DenseWeight,
    pub w_qk_absorbed: DenseWeight,
    pub w_uk_block_diag: DenseWeight,
    pub w_uv_block_diag: DenseWeight,
    pub yarn_inv_freq: DevicePtr,
    pub main_inv_freq: DevicePtr,
    pub q_lora_rank: usize,
    pub kv_lora_rank: usize,
    pub o_lora_rank: usize,
    pub nope: usize,
    pub rope: usize,
    pub v_dim: usize,
    pub compressor: Option<CompressorWeights>,
    pub attn_sink: DevicePtr,
}

MLA (Multi-head Latent Attention) weight components for 2-step decode.

Instead of a single Q GEMV: input × Q_expanded → Q[n_heads*hd], MLA does: input × wq_a → latent[q_lora] → norm → latent × wq_b → Q. This preserves the latent normalization that’s critical for output quality.

Fields

wq_a: DenseWeight

wq_a_nvfp4: Option<QuantizedWeight>

wq_a_fp8: Option<Fp8Weight>

Native block-scaled FP8 weight (the checkpoint ships these projections as FP8-E4M3 + 128×128 block scales). Used by the decode GEMV (w8a16_gemv) so the hot path reads 1 byte/elem instead of the BF16-dequant’s 2 — lossless (the in-kernel dequant keeps F32 precision before the BF16 activation MAC).

wq_b: DenseWeight

wq_b_nvfp4: Option<QuantizedWeight>

wq_b_fp8: Option<Fp8Weight>

q_a_norm: DenseWeight

wkv_a: DenseWeight

wkv_a_nvfp4: Option<QuantizedWeight>

wkv_a_fp8: Option<Fp8Weight>

wkv_b: DenseWeight

kv_a_norm: DenseWeight

wkv_a_rope: DenseWeight

wkv_a_merged: DenseWeight

Merged wkv_a + wkv_a_rope for prefill: [kv_lora+rope, h] — single GEMM replaces 2

wo: DenseWeight

wo_nvfp4: Option<QuantizedWeight>

wo_a: DenseWeight

Grouped low-rank O down-projection (wo_a → wo_b) for DeepSeek-V4-Flash. When o_lora_rank > 0, the decode/prefill paths use wo_a→wo_b instead of wo.

wo_a_nvfp4: Option<QuantizedWeight>

wo_a_fp8: Option<Fp8Weight>

Native block-scaled FP8 wo_a for the grouped decode O-projection. Sliced per o_group (block-diagonal) into w8a16_gemv calls.

wo_b: DenseWeight

wo_b_nvfp4: Option<QuantizedWeight>

wo_b_fp8: Option<Fp8Weight>

w_uk_t: DenseWeight

Absorbed MLA weights for decode (avoid full K/V expansion, preserve precision). W_UK_T: [n_heads, nope, kv_lora] — Q_nope absorption: Q_absorbed = Q_nope @ W_UK_T

w_uv: DenseWeight

W_UV: [n_heads, kv_lora, v_dim] — V extraction: v_out = attn_latent @ W_UV

wq_b_rope: DenseWeight

Q rope projection: wq_b_rope[nqrope, q_lora] — Q_rope = wq_b_rope @ Q_latent Extracted from wq_b rows [nhd+nope .. n*hd+nope+rope] for each head.

w_qk_absorbed: DenseWeight

Fused Q absorption: W_QK_absorbed[nq*kv_lora, q_lora] — Q_absorbed = W_QK @ Q_latent Precomputed as: W_QK[n, lkv, l] = sum_p wq_b_nope[n, p, l] * W_UK[n, p, lkv] Enables single GEMV: Q_absorbed[nq*kv_lora] = W_QK[nq*kv_lora, q_lora] @ Q_latent[q_lora]

w_uk_block_diag: DenseWeight

Block-diagonal W_UK for prefill batched GEMM: [nqkv_lora, nqnope] Single GEMM replaces 32*N per-head GEMV calls for Q absorption in prefill.

w_uv_block_diag: DenseWeight

Block-diagonal W_UV for prefill batched GEMM: [nqv_dim, nqkv_lora] Single GEMM replaces 32*N per-head GEMV calls for V extraction in prefill.

yarn_inv_freq: DevicePtr

Precomputed YaRN inv_freq table [rotary_dim/2] FP32 on GPU. NULL = use standard theta computation in the RoPE kernel.

main_inv_freq: DevicePtr

Plain θ=10000 inv_freq [rotary_dim/2] FP32 on GPU, NO YaRN. Used for the raw-arm Q/K rope on sliding_attention layers (compressor==None): the reference gives sliding layers the “main” rope (θ=rope_theta=10000, no yarn) while CSA/HCA layers use “compress” (θ=compress_rope_theta=160000

• yarn). Atlas previously applied the single yarn table to every layer.

q_lora_rank: usize

kv_lora_rank: usize

o_lora_rank: usize

nope: usize

rope: usize

v_dim: usize

compressor: Option<CompressorWeights>

DeepSeek Sparse Attention compressor (CSA ratio-4 / HCA ratio-128). None for full-attention layers (compress_ratios[L] == 0).

attn_sink: DevicePtr

Per-head attention sink logit [num_q_heads] BF16 (DeepSeek-V4 s_aux). NULL if the checkpoint has no attn_sink for this layer.