> This page is for version Latest.
> For other versions, use one of these documentation indexes:
> - Nightly (default): https://docs.nvidia.com/nemo/automodel/nightly/llms.txt
> - Latest: https://docs.nvidia.com/nemo/automodel/latest/llms.txt
> - 0.7.0 · 26.10: https://docs.nvidia.com/nemo/automodel/v0.7/llms.txt
> - 0.5.0 · 26.06: https://docs.nvidia.com/nemo/automodel/v0.5/llms.txt
> - 0.4.0 · 26.04: https://docs.nvidia.com/nemo/automodel/v0.4/llms.txt

> For clean Markdown of any page, append .md to the page URL.
> For a complete documentation index, see https://docs.nvidia.com/nemo/automodel/llms.txt.

# nemo_automodel.components.models.inkling2.flex_attention

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

## Module Contents

### Classes

| Name                                                                                                       | Description                                                                    |
| ---------------------------------------------------------------------------------------------------------- | ------------------------------------------------------------------------------ |
| [`Inkling2FlexAttention`](#nemo_automodel-components-models-inkling2-flex_attention-Inkling2FlexAttention) | Own compiled callables, with masks rebuilt from each forward's token validity. |

### API

```python
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`

---

```python
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:**

**`q`** `torch.Tensor`

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

---

**`k`** `torch.Tensor`

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

---

**`v`** `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 | None` — default: None

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

---

**`positions`** `torch.Tensor | None` — default: None

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

---

**`key_positions`** `torch.Tensor | None` — default: None

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

---

**`query_offset`** `int` — default: 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