Function qk_norm_expand_bias 

pub fn qk_norm_expand_bias<const D: usize, const DP1: usize>(
    raw_mgr: Tensor<D>,
    norm: &RmsNorm,
    bias_hmr: Tensor<3>,
    group_dim: usize,
    nheads: usize,
) -> Tensor<D>

QK-Norm → GQA-expand groups→heads → add per-(head, mimo-rank) bias.

The input is the raw B/C projection already reshaped to expose the group dim, with last dim = state_rank. The output replaces the group dim with the head dim, leaving the last dim untouched.

DP1 = D + 1 (required by gqa_expand_to_heads’s intermediate rank).