nemo_automodel.components.models.inkling2.flex_attention

Compiled learned-relative GQA, with differentiable per-query attention sinks.

Module Contents

Classes

Name	Description
`Inkling2FlexAttention`	Own compiled callables, with masks rebuilt from each forward’s token validity.

API

class nemo_automodel.components.models.inkling2.flex_attention.Inkling2FlexAttention(
    head_dim: int,
    window: int | None
)

Bases: Module

Own compiled callables, with masks rebuilt from each forward’s token validity.

Relative profiles have shape [batch, heads, sequence, relative_extent]; no sequence-by-sequence score/bias tensor is constructed. The Triton backend supports gradients to captured profiles and returned log-sum-exp values.

_attention

_make_mask

scale

= head_dim ** -0.5

nemo_automodel.components.models.inkling2.flex_attention.Inkling2FlexAttention.forward(
    q: torch.Tensor,
    k: torch.Tensor,
    v: torch.Tensor,
    profiles: torch.Tensor,
    sink: torch.Tensor,
    valid: torch.Tensor,
    key_valid: torch.Tensor | None = None,
    positions: torch.Tensor | None = None,
    key_positions: torch.Tensor | None = None,
    query_offset: int = 0
) -> torch.Tensor

Compute content plus learned relative attention, including its sink.

Parameters:

torch.Tensor

Query tensor of shape [batch, sequence, heads, head_dim].

torch.Tensor

Key tensor of shape [batch, key_sequence, kv_heads, head_dim].

torch.Tensor

Value tensor of shape [batch, key_sequence, kv_heads, head_dim].

profiles

torch.Tensor

Relative logits of shape [batch, heads, sequence, relative_extent]. Retains the projection’s FP32/BF16 rounding boundary.

sink

torch.Tensor

Sink logits of shape [batch, sequence, heads].

valid

torch.Tensor

Boolean tensor of shape [batch, sequence]. Each batch row has a local slice of one global contiguous real-token span.

key_valid

torch.Tensor | NoneDefaults to None

Boolean [batch, key_sequence] validity, global under CP; defaults to valid for CP1. Global real-token positions are consecutive.

positions

torch.Tensor | NoneDefaults to None

Logical query positions [batch, sequence], reset per document.

key_positions

torch.Tensor | NoneDefaults to None

Logical key positions [batch, key_sequence], global under CP.

query_offset

intDefaults to 0

First local query’s global token column, zero without CP.

Returns: torch.Tensor

FP32 tensor of shape [batch, sequence, heads, head_dim]. Inputs remain