> 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.
> For AI client integration (Claude Code, Cursor, etc.), connect to the MCP server at https://docs.nvidia.com/nemo/automodel/_mcp/server.

# nemo_automodel.components.attention.idlm_mask

I-DLM block-diffusion attention mask (Yu et al., 2026; arXiv:2604.11035).

Built over the concatenated `[x_t (L) | x_0 (L)]` input, where `x_t` is the
noisy (masked) copy and `x_0` the clean copy. It combines three components
(paper Appendix E):

* `M_BD`  — causal self-attention within each noisy `x_t` block.
* `M_OBC` — each `x_t` token cross-attends clean `x_0` tokens in strictly
  earlier blocks (the clean ground-truth prefix the decode is conditioned on).
* `M_BC`  — strict token-causal attention within the clean `x_0` copy.

With `block_size == 1` each `x_t` token attends only to itself plus the clean
tokens strictly before it, and `x_0` is plain autoregressive.

## Module Contents

### Functions

| Name                                                                                              | Description                                                            |
| ------------------------------------------------------------------------------------------------- | ---------------------------------------------------------------------- |
| [`create_idlm_block_mask`](#nemo_automodel-components-attention-idlm_mask-create_idlm_block_mask) | Build a sparse FlexAttention :class:`BlockMask` for I-DLM training.    |
| [`create_idlm_sdpa_mask`](#nemo_automodel-components-attention-idlm_mask-create_idlm_sdpa_mask)   | Build the dense additive `[x_t \| x_0]` block-diffusion mask for SDPA. |

### API

```python
nemo_automodel.components.attention.idlm_mask.create_idlm_block_mask(
    seq_len: int,
    block_size: int,
    valid_mask: torch.Tensor,
    device: torch.device,
    use_compile: bool = True
) -> 'BlockMask'
```

Build a sparse FlexAttention :class:`BlockMask` for I-DLM training.

Same semantics as :func:`create_idlm_sdpa_mask` but avoids materialising the
dense `2L x 2L` mask — preferred at scale. Consumed by transformers'
`flex_attention` backend when `_attn_implementation="flex_attention"`;
pass it via the `attention_mask` kwarg.

**Parameters:**

**`seq_len`** `int`

Length `L` of one copy (concatenated length is `2L`).

---

**`block_size`** `int`

Diffusion block size.

---

**`valid_mask`** `torch.Tensor`

Padding-validity mask over one copy, shape `[B, L]`.

---

**`device`** `torch.device`

torch device.

---

**`use_compile`** `bool` — default: True

Reuse a cached `torch.compile`'d `create_block_mask`.

---

**Returns:** `'BlockMask'`

class:`torch.nn.attention.flex_attention.BlockMask`.

```python
nemo_automodel.components.attention.idlm_mask.create_idlm_sdpa_mask(
    seq_len: int,
    block_size: int,
    valid_mask: torch.Tensor,
    device: torch.device,
    dtype: torch.dtype
) -> torch.Tensor
```

Build the dense additive `[x_t | x_0]` block-diffusion mask for SDPA.

The `x_0` region is always strict token-causal, matching the paper (§3.1)
and every released I-DLM checkpoint.

**Parameters:**

**`seq_len`** `int`

Length `L` of one copy (the concatenated length is `2L`).

---

**`block_size`** `int`

Diffusion block size (`block_length` in the paper).

---

**`valid_mask`** `torch.Tensor`

Padding-validity mask over one copy, shape `[B, L]`
(1 = real token); padded keys are blocked in both copies.

---

**`device`** `torch.device`

torch device.

---

**`dtype`** `torch.dtype`

dtype for the additive mask (typically the model dtype).

---

**Returns:** `torch.Tensor`

`[B, 1, 2L, 2L]` float tensor: `0` at attended positions, `-inf`