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

API

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'

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