Struct MultiHeadAttention 

pub struct MultiHeadAttention {
    pub w_q: Tensor,
    pub w_k: Tensor,
    pub w_v: Tensor,
    pub w_o: Tensor,
    pub b_q: Option<Tensor>,
    pub b_k: Option<Tensor>,
    pub b_v: Option<Tensor>,
    pub q_norm: Option<Tensor>,
    pub k_norm: Option<Tensor>,
    /* private fields */
}

Multi-head self-attention layer

Fields

w_q: Tensor

Query projection weight (hidden_size x hidden_size)

w_k: Tensor

Key projection weight (hidden_size x kv_hidden_size)

w_v: Tensor

Value projection weight (hidden_size x kv_hidden_size)

w_o: Tensor

Output projection weight (hidden_size x hidden_size)

b_q: Option<Tensor>

Optional query bias (Qwen2 uses attention biases)

b_k: Option<Tensor>

Optional key bias

b_v: Option<Tensor>

Optional value bias

q_norm: Option<Tensor>

Optional Q RMSNorm weight (Qwen3 uses QK-norm, shape=[head_dim])

k_norm: Option<Tensor>

Optional K RMSNorm weight (Qwen3 uses QK-norm, shape=[head_dim])

Implementations

impl MultiHeadAttention

pub fn new(config: &TransformerConfig) -> Self

Create new attention layer with initialized weights.

When config.use_bias == true (Qwen2 family), allocates Q/K/V projection biases as zero tensors. The forward pass already honors Option<Tensor> biases (lines 388-395 add via add_bias when Some); without allocating them here, biases stay None and populate_trainer_from_init_tensors silently drops the corresponding init tensors during fine-tune from a Qwen APR checkpoint — see FALSIFY-APR-PRETRAIN-INIT-POPULATE- COVERAGE-001/002 in transformer::config::tests for the 290-vs-218 named-parameters gap that surfaced this bug.

Zero-init for biases matches HuggingFace LLaMA / Qwen convention (PyTorch nn.Linear(bias=True) initializes the weight with kaiming_uniform_ but the bias as the all-zeros tensor — see torch.nn.modules.linear.Linear.reset_parameters).

pub fn from_params( config: &TransformerConfig, params: &HashMap<String, Tensor>, prefix: &str, ) -> Option<Self>

Create attention layer from parameter map

Expected parameter names (following HuggingFace convention):

• {prefix}.q_proj.weight
• {prefix}.k_proj.weight
• {prefix}.v_proj.weight
• {prefix}.o_proj.weight

Contract (PMAT-331)

Validates Q/K/V/O projection shapes against config dimensions. Returns None if any key is missing or shape is wrong.

pub fn forward(&self, x: &Tensor, seq_len: usize) -> Tensor

Forward pass

Arguments

• x - Input tensor (seq_len * hidden_size, flattened)
• seq_len - Sequence length

Returns

Output tensor (seq_len * hidden_size, flattened)

pub fn forward_with_lora( &self, x: &Tensor, seq_len: usize, lora_a_q: &Tensor, lora_b_q: &Tensor, lora_a_v: &Tensor, lora_b_v: &Tensor, lora_rank: usize, lora_scale: f32, ) -> Tensor

Forward pass with LoRA adjusts on Q and V projections (KAIZEN-010).

Applies LoRA adapters to Q and V during the forward pass so that gradients flow through LoRA A/B matrices on non-CUDA paths.

Arguments

• x - Input tensor (seq_len * hidden_size)
• seq_len - Sequence length
• lora_a_q, lora_b_q - Q projection LoRA matrices (rank×d_in, d_out×rank)
• lora_a_v, lora_b_v - V projection LoRA matrices (rank×d_in, d_out×rank)
• lora_rank - LoRA rank
• lora_scale - LoRA scaling factor (alpha/rank)

pub fn parameters(&self) -> Vec<&Tensor>

Get all parameters as a vector

pub fn parameters_mut(&mut self) -> Vec<&mut Tensor>

Get all parameters as mutable references for optimizer

pub fn has_biases(&self) -> bool

Whether this attention layer has QKV biases

pub fn named_parameters(&self, prefix: &str) -> Vec<(String, &Tensor)>

Get named parameters for checkpoint serialization

pub fn set_named_parameter(&mut self, suffix: &str, value: Tensor) -> bool

ENT-282: Set a named parameter by suffix (after “self_attn.”).

Bias suffixes route to b_q / b_k / b_v only when those fields are already Some (i.e., MultiHeadAttention::new allocated them because config.use_bias == true). If the caller asks to set a bias on an attention that doesn’t have one, return false — same semantic as setting an unrecognized suffix. This keeps populate_trainer_from_init_tensors honest: a Qwen-init APR’s biases populate iff the target Transformer was built from a use_bias=true config.