`core.transformer.multi_latent_attention`#

Module Contents#

Classes#

`_QuantizeKVForFusedAttn`	Bridge the autograd graph across mxfp8_quantize_only’s tensor reconstruction.
`_QuantizeQueryForFusedAttn`	Quantize the RoPE’d query to MXFP8 as attention would; its gradient passes through.
`_FusedKVRoPEQuantForFusedAttn`	Split kv, apply RoPE to the key and quantize key and value to MXFP8 in one TE kernel.
`MLASelfAttentionSubmodules`	Submodules for the MLA self-attention layer.
`MultiLatentAttention`	Multi-Latent Attention layer abstract class.
`MLASelfAttention`	MLA Self-attention layer class
`FusedMLASelfAttention`	MLA self-attention with fused q/kv down projection.

Functions#

`_is_mxfp8_attention_active`	Whether attention runs in MXFP8 here: in an FP8 autocast with an MXFP8 recipe and FP8 DPA (layers kept in BF16 run outside it), with an FP8 backward. With NVTE_FP8_DPA_BWD=0, attention keeps the BF16 q/k/v for its backward, which pre-quantized inputs would replace by their dequantized values.
`_prepare_mla_core_attention_value`	Prepare value tensor for MLA core attention THD execution.
`_trim_mla_core_attention_output`	Trim THD MLA core attention output back to the original V dimension.
`_up_proj_input`	Return the up-projection input, copied to contiguous memory if `projection` needs it.

API#

class core.transformer.multi_latent_attention._QuantizeKVForFusedAttn#

Bases: torch.autograd.Function

Bridge the autograd graph across mxfp8_quantize_only’s tensor reconstruction.

static forward(ctx, key, value)#: Quantize the k anv v tensors

static backward(ctx, dk, dv)#: Attention returns bf16 tensors, so no need to quantize/unquantize in bwd

class core.transformer.multi_latent_attention._QuantizeQueryForFusedAttn#

Bases: torch.autograd.Function

Quantize the RoPE’d query to MXFP8 as attention would; its gradient passes through.

static forward(ctx, query)#: Quantize the query

static backward(ctx, dq)#: Attention returns a bf16 query gradient

class core.transformer.multi_latent_attention._FusedKVRoPEQuantForFusedAttn#

Bases: torch.autograd.Function

Split kv, apply RoPE to the key and quantize key and value to MXFP8 in one TE kernel.

The outputs are those of fused_apply_mla_rope_for_kv quantized by mxfp8_quantize_only. The backward is that of the key/value split RoPE, given the same saved state.

static forward(ctx, kv, k_pos_emb, cos, sin, emb_dim, k_dim, v_dim)#: Split kv, apply RoPE to k_pos_emb and quantize key and value

static backward(ctx, dk, dv)#: Attention returns bf16 key and value gradients

core.transformer.multi_latent_attention._is_mxfp8_attention_active() → bool#: Whether attention runs in MXFP8 here: in an FP8 autocast with an MXFP8 recipe and FP8 DPA (layers kept in BF16 run outside it), with an FP8 backward. With NVTE_FP8_DPA_BWD=0, attention keeps the BF16 q/k/v for its backward, which pre-quantized inputs would replace by their dequantized values.

core.transformer.multi_latent_attention._prepare_mla_core_attention_value( parallel_attention, query, value, packed_seq_params, )#: Prepare value tensor for MLA core attention THD execution.

core.transformer.multi_latent_attention._trim_mla_core_attention_output( core_attn_out, need_v_pad, orig_v_dim, padded_v_dim, )#: Trim THD MLA core attention output back to the original V dimension.

core.transformer.multi_latent_attention._up_proj_input( projection: torch.nn.Module, x: torch.Tensor, ) → torch.Tensor#

Return the up-projection input, copied to contiguous memory if projection needs it.

x is a column slice of the down-projection output. Transformer Engine projections make such a row-strided input contiguous with Tensor.contiguous, so for them the slice is copied first with the vectorized contiguous_rows; other projections read it as it is.

class core.transformer.multi_latent_attention.MLASelfAttentionSubmodules#

Submodules for the MLA self-attention layer.

linear_proj: megatron.core.transformer.attention.LinearProjBuilder#: None

q_layernorm: megatron.core.transformer.torch_norm.LayerNormBuilder#: None

kv_layernorm: megatron.core.transformer.torch_norm.LayerNormBuilder#: None

linear_q_proj: Union[megatron.core.transformer.spec_utils.ModuleSpec, type]#: None

linear_q_down_proj: Union[megatron.core.transformer.spec_utils.ModuleSpec, type]#: None

linear_q_up_proj: Union[megatron.core.transformer.spec_utils.ModuleSpec, type]#: None

linear_kv_down_proj: Union[megatron.core.transformer.spec_utils.ModuleSpec, type]#: None

linear_kv_up_proj: Union[megatron.core.transformer.spec_utils.ModuleSpec, type]#: None

linear_qkv_down_proj: Union[megatron.core.transformer.spec_utils.ModuleSpec, type]#: None

core_attention: Union[megatron.core.transformer.spec_utils.ModuleSpec, type]#: None

class core.transformer.multi_latent_attention.MultiLatentAttention( config: megatron.core.transformer.transformer_config.MLATransformerConfig, submodules: core.transformer.multi_latent_attention.MLASelfAttentionSubmodules, layer_number: int, attn_mask_type: megatron.core.transformer.enums.AttnMaskType, attention_type: str, cp_comm_type: Optional[str] = None, pg_collection: Optional[megatron.core.process_groups_config.ProcessGroupCollection] = None, pp_layer_offset: Optional[int] = None, is_mtp_layer: bool = False, name: str | None = None, )#

Bases: megatron.core.transformer.attention.Attention

Multi-Latent Attention layer abstract class.

This layer only contains common modules required for the “self attn” and “cross attn” specializations.

Initialization

Parameters:: name (str | None) – module instance name passed top-down from its paranet module

_run_core_attention(

query,

key,

value,

attention_mask,

packed_seq_params=None,

attn_mask_type=None,

**extra_kwargs,

)#: Run MLA core attention with the THD value pad/trim workaround.

forward( hidden_states: torch.Tensor, attention_mask: torch.Tensor, key_value_states: torch.Tensor | None = None, inference_context: megatron.core.inference.contexts.BaseInferenceContext | None = None, rotary_pos_emb: torch.Tensor | tuple[torch.Tensor, torch.Tensor] | None = None, rotary_pos_cos: torch.Tensor | None = None, rotary_pos_sin: torch.Tensor | None = None, rotary_pos_cos_sin: torch.Tensor | None = None, attention_bias: torch.Tensor | None = None, packed_seq_params: megatron.core.packed_seq_params.PackedSeqParams | None = None, position_ids: torch.Tensor | None = None, sequence_len_offset: int | None = None, *, inference_params: megatron.core.inference.contexts.BaseInferenceContext | None = None, packed_sequence_cp_metadata=None, )#: Forward pass for multi-latent attention

class core.transformer.multi_latent_attention.MLASelfAttention( config: megatron.core.transformer.transformer_config.MLATransformerConfig, submodules: core.transformer.multi_latent_attention.MLASelfAttentionSubmodules, layer_number: int, attn_mask_type=AttnMaskType.padding, cp_comm_type: Optional[str] = None, pg_collection: Optional[megatron.core.process_groups_config.ProcessGroupCollection] = None, pp_layer_offset: Optional[int] = None, is_mtp_layer: bool = False, name: str | None = None, )#

Bases: core.transformer.multi_latent_attention.MultiLatentAttention

MLA Self-attention layer class

Self-attention layer takes input with size [s, b, h] and returns output of the same size.

Initialization

Parameters:: name (str | None) – module instance name passed top-down from its paranet module

_resolve_qk_norm_config( submodules, ) → dict[str, megatron.core.transformer.spec_utils.ModuleSpec | type | megatron.core.transformer.torch_norm.LayerNormBuilder]#: Resolve which Q/KV norm and up-projection implementations to build.

_qkv_down_projection(hidden_states)#: Unfused q/kv down projection path.

get_query_key_value_tensors( hidden_states, key_value_states=None, position_ids=None, packed_seq_params=None, inference_context=None, *, inference_params=None, )#: Derives query, key and value tensors from hidden_states.

uncompress_kv_from_cache(kv_cached)#: Take a compressed kv and uncompress them

prepare_for_absorption()#

Prepare the model for absorption optimization in MLA (Multi-Latent Attention).

This method sets up the necessary components for the absorption technique, which optimizes memory during inference by caching compressed KV latents instead of full KV states. The absorption technique allows efficient decode-only operations by pre-computing certain matrix multiplications.

Note (Peter): Right now we are not doing true absorption. We will add this support at a later time.

The method performs the following operations:

Splits the fused layernorm + linear layer (linear_kv_up_proj) into separate components.
Extracts and stores the up-projection weights for K and V separately, which are used during the absorption process
Replaces the identity kv_layernorm with the actual layernorm from the split
Stores the linear component separately for uncompressing KV cache during prefill/mixed stages

This is a one-time setup that should only be called once at initialization when cache_mla_latents is enabled.

backward_dw() → NoReturn#: Execute weight gradient computation

_backward_kv_proj()#: Computes weight gradients of KV projection layers

_backward_q_proj()#: Computes weight gradients of Q projection layers

_backward_output_proj()#: Computes weight gradients of output projection layer

set_for_recompute_input_layernorm()#: Set the attention layer for recompute input_layernorm. Only needed for fp8/fp4.

clip_qk()#: QK Clipping is a technique to clip the query and key attention logits to prevent the attention logits from exploding. Per MuonClip usage, we update the weight by calling this function after Muon optimizer step.

_clip_q_proj_weight(weight)#: Clip q_proj_weight

_clip_kv_proj_weight(weight)#: Clip kv_proj_weight

class core.transformer.multi_latent_attention.FusedMLASelfAttention( config: megatron.core.transformer.transformer_config.MLATransformerConfig, submodules: core.transformer.multi_latent_attention.MLASelfAttentionSubmodules, layer_number: int, attn_mask_type=AttnMaskType.padding, cp_comm_type: Optional[str] = None, pg_collection: Optional[megatron.core.process_groups_config.ProcessGroupCollection] = None, is_mtp_layer: bool = False, pp_layer_offset: Optional[int] = None, name: str | None = None, )#

Bases: core.transformer.multi_latent_attention.MLASelfAttention

MLA self-attention with fused q/kv down projection.

Initialization

Parameters:: name (str | None) – module instance name passed top-down from its paranet module

_qkv_down_projection(hidden_states)#: Fused q/kv down projection path.

backward_dw() → NoReturn#: Execute weight gradient computation.

set_for_recompute_input_layernorm()#: Set the attention layer for recompute input_layernorm. Only needed for fp8/fp4.

sharded_state_dict( prefix: str = '', sharded_offsets: tuple = (), metadata=None, )#: Return a sharded state dict compatible with pre-fusion checkpoints.

_synthetic_state_dict_key_suffixes()#: Return source checkpoint keys used to locate this module in a state dict.

_synthesize_fused_qkv_down_weight(state_dict, prefix)#: Materialize fused qkv-down weight from old separate q/kv checkpoint keys.

_load_from_state_dict(state_dict, prefix, *args, **kwargs)#: Load state dict with automatic unfused->fused conversion.

core.transformer.multi_latent_attention#

Module Contents#

Classes#

Functions#

API#

`core.transformer.multi_latent_attention`#