> This page is for version 0.4.0 · 26.04.
> 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.attention

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

## Module Contents

### Classes

| Name                                                                                          | Description                                                               |
| --------------------------------------------------------------------------------------------- | ------------------------------------------------------------------------- |
| [`AttentionSources`](#nemo_automodel-components-models-inkling2-attention-AttentionSources)   | Q/K/V inputs, each \[batch, sequence, hidden], after its respective norm. |
| [`Inkling2Attention`](#nemo_automodel-components-models-inkling2-attention-Inkling2Attention) | Local or full GQA with relative features, convolutions, gate and sink.    |

### API

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

**Bases:** `NamedTuple`

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

**`k`** `Tensor`

---

**`q`** `Tensor`

---

**`v`** `Tensor`

---

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

---

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

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

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

---

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

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

---

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

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

---

**Returns:** `torch.Tensor`

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

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

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

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

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