Struct std::Mla

#include <std/ml/attention_advanced.h> source

struct Mla {
    &FloatTensor Wq;
    &FloatTensor Wdkv;
    &FloatTensor Wuk;
    &FloatTensor Wuv;
    unsigned long numHeads;
}

Compresses K/V into a shared low-rank latent c_kv = X @ Wdkv before up-projecting back to full K/V width — the KV-cache-memory-saving idea MLA is built around. (Decoupled rotary keys, the other half of the real DeepSeek-V2 design, are out of scope — no positional encoding is baked into this library's attention at all, see attention.h.) Q = X @ Wq [seqLen, dModel] c_kv = X @ Wdkv [seqLen, dLatent] K = c_kv @ Wuk [seqLen, dModel] V = c_kv @ Wuv [seqLen, dModel] out = mha_forward(Q, K, V, numHeads) (Mla::forward).

Fields

§&FloatTensor Wqsource

[dModel, dModel]

§&FloatTensor Wdkvsource

[dModel, dLatent]

§&FloatTensor Wuksource

[dLatent, dModel]

§&FloatTensor Wuvsource

[dLatent, dModel]

§unsigned long numHeadssource

Methods

§&FloatTensor forward(const &FloatTensor X) constsource · definition

[seqLen, dModel] -> [seqLen, dModel]

§void collect_params(&Vec<&FloatTensor> out) constsource · definition

Parameterized (module.h): Wq, Wdkv, Wuk, Wuv.

§void collect_state(&Vec<&FloatTensor> out) constsource · definition

§void drop()source · definition

§void free() consumesource · definition

= drop()

§

static struct Mla new(
    unsigned long dModel,
    unsigned long dLatent,
    unsigned long numHeads,
    &Rng rng,
)

source · definition

Weights U(-1/sqrt(in), 1/sqrt(in)) from rng; numHeads divides dModel.