nemo_automodel.components.models.inkling.text

Native Inkling text backbone, attention, masks, and decoding cache.

Module Contents

Classes

Name	Description
`InklingAttention`	Inkling grouped-query attention with learned relative logits.
`InklingDecoderLayer`	One native Inkling decoder layer.
`InklingDynamicCache`	Model-owned dynamic cache for Inkling attention and short convolutions.
`InklingRMSNorm`	RMS normalization with fp32 variance accumulation.
`InklingRelativeLogits`	Project token-conditioned relative-position profiles into attention bias.
`InklingTextModel`	Native Inkling text backbone with pipeline-stage support.
`_InklingCacheLayer`	Dynamic key/value and short-convolution state for one decoder layer.

Functions

Name	Description
`_build_attention_mask`	Build an additive causal mask for Inkling eager attention.
`_build_conv_mask`	Extract a local two-dimensional padding mask for short convolutions.
`_repeat_key_value`	Repeat grouped key/value heads to match query heads.

Data

__all__

API

class nemo_automodel.components.models.inkling.text.InklingAttention(
    config: nemo_automodel.components.models.inkling.configuration.InklingTextConfig,
    layer_idx: int,
    backend: nemo_automodel.components.models.common.BackendConfig
)

Bases: Module

Inkling grouped-query attention with learned relative logits.

attention_backend

attention_dropout

= config.attention_dropout

head_dim

is_sliding

= config.layer_types[layer_idx] == 'hybrid_sliding'

k_norm

k_proj

k_sconv

num_heads

num_key_value_groups

= self.num_heads // self.num_key_value_heads

num_key_value_heads

o_proj

q_norm

q_proj

r_proj

rel_extent

rel_logits_proj

scaling

= 1.0 / self.head_dim

sliding_window

v_proj

v_sconv

nemo_automodel.components.models.inkling.text.InklingAttention.forward(
    hidden_states: torch.Tensor,
    attention_mask: torch.Tensor,
    conv_mask: torch.Tensor | None = None,
    past_key_values: nemo_automodel.components.models.inkling.text.InklingDynamicCache | None = None,
    kwargs: typing.Any = {}
) -> tuple[torch.Tensor, torch.Tensor]

Apply self-attention.

Parameters:

hidden_states

torch.Tensor

Tensor of shape [batch, sequence, hidden].

attention_mask

torch.Tensor

Additive tensor of shape [batch, 1, query, key].

conv_mask

torch.Tensor | NoneDefaults to None

Optional boolean tensor of shape [batch, sequence] where True marks a valid token.

past_key_values

InklingDynamicCache | NoneDefaults to None

Optional model-owned decoding cache.

**kwargs

AnyDefaults to {}

Backend-compatible attention arguments; currently ignored by eager attention.

Returns: torch.Tensor

A pair containing the output tensor of shape [batch, sequence, hidden]

nemo_automodel.components.models.inkling.text.InklingAttention.init_weights(
    init_std: float
) -> None

Initialize projections, norms, relative logits, and convolution weights.

class nemo_automodel.components.models.inkling.text.InklingDecoderLayer(
    config: nemo_automodel.components.models.inkling.configuration.InklingTextConfig,
    layer_idx: int,
    backend: nemo_automodel.components.models.common.BackendConfig,
    moe_config: nemo_automodel.components.moe.config.MoEConfig
)

Bases: Module

One native Inkling decoder layer.

attention_type

attn_sconv

input_layernorm

layer_type

= config.layer_types[layer_idx]

mlp

mlp_sconv

post_attention_layernorm

self_attn

= InklingAttention(config, layer_idx, backend)

nemo_automodel.components.models.inkling.text.InklingDecoderLayer.forward(
    hidden_states: torch.Tensor,
    attention_mask: torch.Tensor,
    conv_mask: torch.Tensor | None = None,
    past_key_values: nemo_automodel.components.models.inkling.text.InklingDynamicCache | None = None,
    kwargs: typing.Any = {}
) -> torch.Tensor

Apply attention and feed-forward residual blocks.

Parameters:

hidden_states

torch.Tensor

Tensor of shape [batch, sequence, hidden].

attention_mask

torch.Tensor

Additive tensor of shape [batch, 1, query, key].

conv_mask

torch.Tensor | NoneDefaults to None

Optional boolean tensor of shape [batch, sequence] where True marks a valid token.

past_key_values

InklingDynamicCache | NoneDefaults to None

Optional model-owned decoding cache.

**kwargs

AnyDefaults to {}

Additional eager-attention arguments.

Returns: torch.Tensor

Tensor of shape [batch, sequence, hidden].

nemo_automodel.components.models.inkling.text.InklingDecoderLayer.init_weights(
    init_std: float,
    buffer_device: torch.device
) -> None

Initialize all parameters owned by the decoder layer.

class nemo_automodel.components.models.inkling.text.InklingDynamicCache(
    config: nemo_automodel.components.models.inkling.configuration.InklingTextConfig
)

Model-owned dynamic cache for Inkling attention and short convolutions.

layers

nemo_automodel.components.models.inkling.text.InklingDynamicCache.get_mask_sizes(
    query_length: int,
    layer_idx: int
) -> tuple[int, int]

Return the key/value length and absolute offset for a decoder layer.

nemo_automodel.components.models.inkling.text.InklingDynamicCache.get_query_offset(
    layer_idx: int = 0
) -> int

Return the absolute starting position of the next query.

nemo_automodel.components.models.inkling.text.InklingDynamicCache.get_seq_length(
    layer_idx: int = 0
) -> int

Return the number of tokens observed at a decoder layer.

nemo_automodel.components.models.inkling.text.InklingDynamicCache.has_previous_state(
    layer_idx: int,
    state_idx: int
) -> bool

Return whether one short-convolution cache has been populated.

nemo_automodel.components.models.inkling.text.InklingDynamicCache.reorder_cache(
    beam_indices: torch.Tensor
) -> None

Reorder every layer cache during beam search.

Parameters:

beam_indices

torch.Tensor

Long tensor of shape [new_batch] indexing the old batch axis.

nemo_automodel.components.models.inkling.text.InklingDynamicCache.update(
    key_states: torch.Tensor,
    value_states: torch.Tensor,
    layer_idx: int
) -> tuple[torch.Tensor, torch.Tensor | None]

Append key/value states at one decoder layer.

Parameters:

key_states

torch.Tensor

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

value_states

torch.Tensor

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

layer_idx

int

Decoder-layer index.

Returns: tuple[torch.Tensor, torch.Tensor | None]

A pair of tensors with shape [batch, key_value_heads, visible_sequence, head_dim].

nemo_automodel.components.models.inkling.text.InklingDynamicCache.update_conv_state(
    conv_states: torch.Tensor,
    layer_idx: int,
    state_idx: int,
    conv_kernel_size: int
) -> torch.Tensor

Append and cache inputs for one short convolution.

Parameters:

conv_states

torch.Tensor

Tensor of shape [batch, hidden, sequence].

layer_idx

int

Decoder-layer index.

state_idx

int

Convolution index inside that layer.

conv_kernel_size

int

Number of trailing tokens retained.

Returns: torch.Tensor

Tensor of shape [batch, hidden, cached_sequence + sequence].

class nemo_automodel.components.models.inkling.text.InklingRMSNorm(
    hidden_size: int,
    eps: float,
    dtype: torch.dtype
)

Bases: Module

RMS normalization with fp32 variance accumulation.

weight

= nn.Parameter(torch.ones(hidden_size, dtype=dtype))

nemo_automodel.components.models.inkling.text.InklingRMSNorm.forward(
    hidden_states: torch.Tensor
) -> torch.Tensor

Normalize the final hidden axis.

Parameters:

hidden_states

torch.Tensor

Tensor of shape [..., hidden] with arbitrary leading dimensions.

Returns: torch.Tensor

Tensor with the same shape and dtype as hidden_states.

nemo_automodel.components.models.inkling.text.InklingRMSNorm.init_weights() -> None

Initialize the learned scale to one.

class nemo_automodel.components.models.inkling.text.InklingRelativeLogits(
    d_rel: int,
    rel_extent: int,
    dtype: torch.dtype
)

Bases: Module

Project token-conditioned relative-position profiles into attention bias.

proj

nemo_automodel.components.models.inkling.text.InklingRelativeLogits.forward(
    relative_states: torch.Tensor,
    query_positions: torch.Tensor,
    key_positions: torch.Tensor
) -> torch.Tensor

Materialize the relative-position bias.

Parameters:

relative_states

torch.Tensor

Tensor of shape [batch, query, heads, relative_dim].

query_positions

torch.Tensor

Long tensor of shape [query] containing absolute positions.

key_positions

torch.Tensor

Long tensor of shape [key] containing absolute positions.

Returns: torch.Tensor

Tensor of shape [batch, heads, query, key].

nemo_automodel.components.models.inkling.text.InklingRelativeLogits.init_weights(
    init_std: float
) -> None

Initialize the relative-position profile bank.

class nemo_automodel.components.models.inkling.text.InklingTextModel(
    config: nemo_automodel.components.models.inkling.configuration.InklingTextConfig,
    backend: nemo_automodel.components.models.common.BackendConfig,
    moe_config: nemo_automodel.components.moe.config.MoEConfig
)

Bases: Module

Native Inkling text backbone with pipeline-stage support.

embed_norm

embed_tokens

layers

norm

nemo_automodel.components.models.inkling.text.InklingTextModel.forward(
    input_ids: torch.LongTensor | None = None,
    attention_mask: torch.Tensor | dict[str, torch.Tensor] | None = None,
    position_ids: torch.LongTensor | None = None,
    past_key_values: nemo_automodel.components.models.inkling.text.InklingDynamicCache | None = None,
    inputs_embeds: torch.FloatTensor | None = None,
    use_cache: bool | None = None,
    kwargs: typing.Any = {}
) -> transformers.modeling_outputs.BaseModelOutputWithPast

Run the text backbone.

Parameters:

input_ids

torch.LongTensor | NoneDefaults to None

Optional tensor of shape [batch, sequence]. Pipeline stages without embeddings may instead receive floating hidden states of shape [batch, sequence, hidden] through this argument.

attention_mask

torch.Tensor | dict[str, torch.Tensor] | NoneDefaults to None

Optional padding tensor of shape [batch, total_sequence] or a mapping from attention type to additive masks of shape [batch, 1, query, key].

position_ids

torch.LongTensor | NoneDefaults to None

Optional long tensor of shape [batch, sequence].

past_key_values

InklingDynamicCache | NoneDefaults to None

Optional model-owned decoding cache.

inputs_embeds

torch.FloatTensor | NoneDefaults to None

Optional tensor of shape [batch, sequence, hidden].

use_cache

bool | NoneDefaults to None

Whether to allocate and return a decoding cache.

**kwargs

AnyDefaults to {}

Additional eager-attention arguments.

Returns: BaseModelOutputWithPast

A model output whose last_hidden_state has shape [batch, sequence, hidden].

nemo_automodel.components.models.inkling.text.InklingTextModel.get_input_embeddings() -> torch.nn.Module

Return the token-embedding module.

nemo_automodel.components.models.inkling.text.InklingTextModel.init_weights(
    buffer_device: torch.device
) -> None

Initialize the text backbone for checkpoint-free construction.

nemo_automodel.components.models.inkling.text.InklingTextModel.set_input_embeddings(
    embeddings: torch.nn.Module
) -> None

Replace the token-embedding module.

class nemo_automodel.components.models.inkling.text._InklingCacheLayer(
    sliding_window: int | None,
    number_of_conv_states: int
)

Dynamic key/value and short-convolution state for one decoder layer.

_has_previous_conv_state

dict[int, bool] = dict.fromkeys(range(number_of_conv_states), False)

conv_states

dict[int, Tensor | None] = dict.fromkeys(range(number_of_conv_states))

cumulative_length

= 0

keys

Tensor | None = None

values

Tensor | None = None

nemo_automodel.components.models.inkling.text._InklingCacheLayer.get_mask_sizes(
    query_length: int
) -> tuple[int, int]

Return key/value length and absolute offset before appending a query.

nemo_automodel.components.models.inkling.text._InklingCacheLayer.get_seq_length() -> int

Return the number of tokens observed by this layer.

nemo_automodel.components.models.inkling.text._InklingCacheLayer.has_previous_state(
    state_idx: int
) -> bool

Return whether a convolution state has seen at least one token.

nemo_automodel.components.models.inkling.text._InklingCacheLayer.reorder(
    beam_indices: torch.Tensor
) -> None

Reorder cached batch entries during beam search.

Parameters:

beam_indices

torch.Tensor

Long tensor of shape [new_batch] indexing the old batch axis.

nemo_automodel.components.models.inkling.text._InklingCacheLayer.update(
    key_states: torch.Tensor,
    value_states: torch.Tensor
) -> tuple[torch.Tensor, torch.Tensor]

Append key/value states and return the states visible to this query.

Parameters:

key_states

torch.Tensor

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

value_states

torch.Tensor

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

Returns: tuple[torch.Tensor, torch.Tensor]

A pair of tensors with shape [batch, key_value_heads, visible_sequence, head_dim].

nemo_automodel.components.models.inkling.text._InklingCacheLayer.update_conv_state(
    conv_states: torch.Tensor,
    state_idx: int,
    conv_kernel_size: int
) -> torch.Tensor

Append short-convolution inputs and update the fixed-size cache.

Parameters:

conv_states

torch.Tensor

Tensor of shape [batch, hidden, sequence].

state_idx

int

Index of one of the four convolution states owned by the layer.

conv_kernel_size

int

Number of trailing tokens retained in the cache.

Returns: torch.Tensor

Tensor of shape [batch, hidden, cached_sequence + sequence] used by

nemo_automodel.components.models.inkling.text._build_attention_mask(
    inputs_embeds: torch.Tensor,
    attention_mask: torch.Tensor | None,
    past_key_values: nemo_automodel.components.models.inkling.text.InklingDynamicCache | None,
    layer_idx: int,
    sliding_window: int | None
) -> torch.Tensor

Build an additive causal mask for Inkling eager attention.

Parameters:

inputs_embeds

torch.Tensor

Tensor of shape [batch, query, hidden].

attention_mask

torch.Tensor | None

Optional tensor of shape [batch, total_sequence] where nonzero entries are valid keys, or a prepared tensor of shape [batch, 1, query, key].

past_key_values

InklingDynamicCache | None

Optional model-owned decoding cache.

layer_idx

int

Decoder-layer index used to size a hybrid cache.

sliding_window

int | None

Optional number of visible tokens for local attention.

Returns: torch.Tensor

Additive tensor of shape [batch, 1, query, key] in inputs_embeds.dtype.

nemo_automodel.components.models.inkling.text._build_conv_mask(
    attention_mask: torch.Tensor | None,
    query_length: int
) -> torch.Tensor | None

Extract a local two-dimensional padding mask for short convolutions.

Parameters:

attention_mask

torch.Tensor | None

Optional tensor of shape [batch, total_sequence].

query_length

int

Length of the current local sequence.

Returns: torch.Tensor | None

Optional boolean tensor of shape [batch, query_length] where True

nemo_automodel.components.models.inkling.text._repeat_key_value(
    hidden_states: torch.Tensor,
    repeats: int
) -> torch.Tensor

Repeat grouped key/value heads to match query heads.

Parameters:

hidden_states

torch.Tensor

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

repeats

int

Number of query heads sharing each key/value head.

Returns: torch.Tensor

Tensor of shape [batch, key_value_heads * repeats, sequence, head_dim].

nemo_automodel.components.models.inkling.text.__all__ = ['InklingDynamicCache', 'InklingRMSNorm', 'InklingTextModel']