nemo_automodel.components.models.inkling2.attention

Learned-relative GQA with eager reference and compiled training attention.

Module Contents

Classes

Name	Description
`AttentionSources`	Q/K/V inputs, each [batch, sequence, hidden], after its respective norm.
`Inkling2Attention`	Local or full GQA with relative features, convolutions, gate and sink.

API

class nemo_automodel.components.models.inkling2.attention.AttentionSources()

Bases: NamedTuple

Q/K/V inputs, each [batch, sequence, hidden], after its respective norm.

Tensor

class nemo_automodel.components.models.inkling2.attention.Inkling2Attention(
    config: nemo_automodel.components.models.inkling2.config.Inkling2TextConfig,
    layer_idx: int,
    dtype: torch.dtype,
    backend: typing.Literal['eager', 'flex'] = 'eager'
)

Bases: Module

Local or full GQA with relative features, convolutions, gate and sink.

cp_mesh

DeviceMesh | None = None

d_rel

= config.d_rel

eps

= config.rms_norm_eps

flex

head_dim

heads

is_local

= layer_idx in config.local_layer_ids

k_l2_clip

k_sconv

kv_heads

log_scaling_alpha

= config.log_scaling_alpha

log_scaling_n_floor

= config.log_scaling_n_floor

q_l2_clip

qkvr

rel_extent

rel_l2_clip

relative_table

split_sizes

v_sconv

value_norm

= config.use_value_norm_no_gamma

window

= config.sliding_window_size

wo_ud

nemo_automodel.components.models.inkling2.attention.Inkling2Attention._eager(
    q: torch.Tensor,
    k: torch.Tensor,
    v: torch.Tensor,
    profiles: torch.Tensor,
    sink: torch.Tensor,
    valid: torch.Tensor,
    positions: torch.Tensor,
    key_valid: torch.Tensor,
    key_positions: torch.Tensor,
    query_offset: int = 0
) -> torch.Tensor

Evaluate the quadratic reference without mutating inputs.

Parameters:

torch.Tensor

Tensor of shape [batch, sequence, heads, head_dim].

torch.Tensor

Tensor of shape [batch, key_sequence, kv_heads, head_dim].

torch.Tensor

Tensor of shape [batch, key_sequence, kv_heads, head_dim].

profiles

torch.Tensor

Tensor of shape [batch, heads, sequence, relative_extent].

sink

torch.Tensor

Tensor of shape [batch, sequence, heads].

valid

torch.Tensor

Boolean tensor of shape [batch, sequence].

positions

torch.Tensor

Integer logical positions [batch, local_sequence].

key_valid

torch.Tensor

Boolean validity [batch, key_sequence], global under CP.

key_positions

torch.Tensor

Integer logical positions [batch, key_sequence].

query_offset

intDefaults to 0

First local query’s column in the global token stream.

Returns: torch.Tensor

FP32 context of shape [batch, sequence, heads, head_dim].

nemo_automodel.components.models.inkling2.attention.Inkling2Attention.forward(
    sources: nemo_automodel.components.models.inkling2.attention.AttentionSources,
    valid: torch.Tensor,
    positions: torch.Tensor
) -> torch.Tensor

Evaluate the configured attention backend with a zero-valued virtual sink.

Parameters:

sources

AttentionSources

Q/K/V tensors of shape [batch, sequence, hidden].

valid

torch.Tensor

Boolean tensor of shape [batch, sequence], true for real tokens.

positions

torch.Tensor

Integer tensor of shape [batch, sequence], logical token positions.

Returns: torch.Tensor

Tensor of shape [batch, sequence, hidden] in sources.q’s dtype.

nemo_automodel.components.models.inkling2.attention.Inkling2Attention.reset_parameters(
    init_std: float
) -> None

Initialize fused projections, relative table, scales and convolution weights.

nemo_automodel.components.models.inkling2.attention.Inkling2Attention.setup_cp_attention(
    cp_mesh: torch.distributed.device_mesh.DeviceMesh
) -> None

Configure model-owned KV transport through the shared MoE CP hook.