mstar.model.glm52.components.indexer

mstar.model.glm52.components.indexer#

GLM-5.2’s DSA indexer: the index keys and queries its FULL layers hand the DSA resource (mstar/engine/resources/dsa), which scores them and picks each query’s top-k.

Functions

is_full_indexer_layer(config, layer_idx)

FULL layers run the indexer; SHARED layers reuse the last FULL selection.

Classes

Glm52Indexer(config)

DSA indexer projections for one FULL layer: the key each token stores, and each query row's heads and head weights.
class mstar.model.glm52.components.indexer.Glm52Indexer(config)[source]#

Bases: Module

DSA indexer projections for one FULL layer: the key each token stores, and each query row’s heads and head weights.

Parameters:

config (Glm52ModelConfig)

bind_resources(resources)[source]#
Parameters:

resources (dict)

Return type:

None

compute_k(hidden_states, positions)[source]#

Per-token index key (T, head_dim) — what the indexer cache stores.

Parameters:
Return type:

Tensor

compute_query(q_c, hidden_states, positions)[source]#

Each row’s index query (T, H, D) and head weights (T, H), fp32: what the DSA resource scores the stored keys with.

Parameters:
Return type:

tuple[Tensor, Tensor]

mstar.model.glm52.components.indexer.is_full_indexer_layer(config, layer_idx)[source]#

FULL layers run the indexer; SHARED layers reuse the last FULL selection.

Parameters:
Return type:

bool