Is this page helpful?

core.transformer.mla_qk_norm_config#

Resolve MLA and DSA Q/KV norm configuration from a layer specification.

Module Contents#

Classes#

QKNormConfigResolver

Validate and resolve Q/KV norm placement for MLA and DSA.

Data#

__all__

_QKNormResolvedConfig

API#

core.transformer.mla_qk_norm_config.__all__#

[]

core.transformer.mla_qk_norm_config._QKNormResolvedConfig#

None

class core.transformer.mla_qk_norm_config.QKNormConfigResolver(
config: megatron.core.transformer.transformer_config.MLATransformerConfig,
submodules,
)#

Validate and resolve Q/KV norm placement for MLA and DSA.

Q/KV norm can be represented either by a standalone norm module or by a fused norm+linear projection. MLA can use the fused form; DSA cannot because it needs the normalized Q/KV values outside the projection.

Constraints:

  • qk_l2_norm is unsupported for MLA/DSA.

  • A standalone Q norm is only usable when q_lora_rank is set.

  • Explicit norm modules cannot be paired with fused norm+linear projections.

  • Disabled QK norm rejects both explicit norms and fused norm+linear projections.

  • DSA with QK norm requires non-fused projections and standalone Q/KV norms.

Initialization

Capture the configuration, requested modules, and backend implementations.

resolve() → core.transformer.mla_qk_norm_config._QKNormResolvedConfig#

Validate the specification and return the modules to instantiate.

Returns:

The Q/KV norms and projections after applying the MLA or DSA constraints.

Raises:

ValueError – If the requested norm placement is unsupported or conflicting.

_resolve_disabled_qk_layernorm() → core.transformer.mla_qk_norm_config._QKNormResolvedConfig#

Resolve projections when Q/KV normalization is disabled.

Explicit norm modules and fused norm-linear projections are rejected because they would still introduce Q/KV normalization.

_resolve_dsa_qk_layernorm() → core.transformer.mla_qk_norm_config._QKNormResolvedConfig#

Resolve DSA’s standalone Q/KV norms and non-fused projections.

DSA consumes the normalized Q/KV values outside the projection, so it cannot use fused norm-linear projections.

_resolve_mla_qk_layernorm() → core.transformer.mla_qk_norm_config._QKNormResolvedConfig#

Resolve MLA norms, fusing them into projections when no norm is explicit.

_reject_common_spec_conflicts() → None#

Reject conflicts that apply regardless of the selected attention variant.

_reject_disabled_norm(
module_spec,
norm_spec,
module_name,
norm_name,
) → None#

Reject a norm module or fused projection when Q/KV norm is disabled.

_reject_explicit_norm_with_fused_linear(
module_spec,
norm_spec,
module_name,
norm_name,
) → None#

Reject specifying the same norm both explicitly and inside a projection.

_non_fused_or_default(module_spec, module_name)#

Return a linear implementation, requiring it not to fuse normalization.

_dsa_linear_or_default(module_spec, module_name)#

Return DSA’s non-fused projection implementation.

This uses a DSA-specific diagnostic so the rejected constraint is clear.

_mla_fused_linear_or_default(module_spec, module_name)#

Return a fused MLA projection, using the backend default when available.

_require_linear(module_spec, module_name)#

Return a configured projection or report that no viable implementation exists.

_raise_unused_q_norm() → NoReturn#

Report an explicit Q norm that has no Q-LoRA projection to consume it.

_is_fused_norm_linear(module_spec) → bool#

Return whether a module specification selects the backend fused projection.

static _is_trivial(module_spec) → bool#

Return whether a norm slot is unset or explicitly an identity operation.

classmethod _default_if_trivial(module_spec, default)#

Replace an unset or identity specification with the supplied default.

static _result(
*,
linear_q_proj,
linear_q_up_proj,
linear_kv_up_proj,
q_layernorm,
kv_layernorm,
) → core.transformer.mla_qk_norm_config._QKNormResolvedConfig#

Package the resolved Q/KV norms and projections in the caller’s schema.