Function metal_graph_encode_decode_layer

ds4.c:14842–16068

}

/* Final language-model head: HC collapse, RMSNorm, and Q8_0 vocab projection. */
static void output_logits_one(
        float             * logits,
        const ds4_model   * model,
        const ds4_weights * weights,
        const float       * inp_hc) {
    float *embd = xmalloc((size_t)DS4_N_EMBD * sizeof(embd[0]));
    float *norm = xmalloc((size_t)DS4_N_EMBD * sizeof(norm[0]));

    output_hc_head_one(embd, model, weights, inp_hc);
    rms_norm_weight(norm, embd, tensor_data(model, weights->output_norm), DS4_N_EMBD, DS4_RMS_EPS);

    matvec_q8_0(logits, model, weights->output, norm);

    free(norm);
    free(embd);
}

static void layer_glm_first_token_attention_one(
        float                   * out,
        const ds4_model         * model,
        const ds4_layer_weights * layer,
        const float             * x) {
    float *norm = xmalloc((size_t)DS4_N_EMBD * sizeof(norm[0]));
    float *kv_raw = xmalloc((size_t)layer->attn_kv_a_mqa->dim[1] * sizeof(kv_raw[0]));
    float *kv_norm = xmalloc((size_t)DS4_N_KV_LORA * sizeof(kv_norm[0]));
    float *heads = xmalloc((size_t)DS4_N_HEAD * DS4_N_VALUE_MLA * sizeof(heads[0]));
    const uint64_t kv_blocks = (DS4_N_KV_LORA + 31) / 32;
    int8_t *kvq = xmalloc((size_t)kv_blocks * 32);
    float *kvscale = xmalloc((size_t)kv_blocks * sizeof(kvscale[0]));

    if (layer->attn_kv_a_mqa->dim[1] < DS4_N_KV_LORA ||
        layer->attn_v_b->dim[0] != DS4_N_KV_LORA ||
        layer->attn_v_b->dim[1] != DS4_N_VALUE_MLA ||
        layer->attn_v_b->dim[2] != DS4_N_HEAD ||
        layer->attn_output->dim[0] != (uint64_t)DS4_N_HEAD * DS4_N_VALUE_MLA ||
        layer->attn_output->dim[1] != DS4_N_EMBD) {
        ds4_die("GLM attention tensors have an unexpected layout");
    }

    rms_norm_weight(norm, x, tensor_data(model, layer->attn_norm), DS4_N_EMBD, DS4_RMS_EPS);
    matvec_q8_0(kv_raw, model, layer->attn_kv_a_mqa, norm);
    rms_norm_weight(kv_norm, kv_raw, tensor_data(model, layer->attn_kv_a_norm),
                    DS4_N_KV_LORA, DS4_RMS_EPS);
    quantize_q8_0_activation(kv_norm, kvq, kvscale, DS4_N_KV_LORA);

    for (uint32_t h = 0; h < DS4_N_HEAD; h++) {
        matvec_q8_0_3d_slice_prequant(heads + (uint64_t)h * DS4_N_VALUE_MLA,
                                      model,
                                      layer->attn_v_b,
                                      kvq,
                                      kvscale,
                                      h);
    }
    matvec_q8_0(out, model, layer->attn_output, heads);

    free(kvscale);
    free(kvq);