bridge.diffusion.models.llada15.llada15_attention

LLaDA15TEDotProductAttention: TE-backed core attention for LLaDA1.5.

LLaDA1.5 is a block-diffusion masked diffusion language model. Unlike LLaDA2, the reference implementation in modeling_llada.py uses fully bidirectional attention at both training and inference time — the block structure exists only in the sampling schedule (which positions get unmasked per iteration), not in the attention pattern.

This shim therefore needs only one job: override Megatron’s default AttnMaskType.causal to AttnMaskType.no_mask so attention is bidirectional. RoPE is full and handled upstream by Megatron’s SelfAttention — this module does not touch RoPE.

The set_block_mask / reset_inference_state hooks are kept as no-op extension points for users who want to experiment with block-diagonal attention (the LLaDA2-style design) — they are not wired into the default inference loop.

Module Contents

Classes

LLaDA15TEDotProductAttention

TE-backed core attention for LLaDA1.5 masked-diffusion training/inference.

API

class bridge.diffusion.models.llada15.llada15_attention.LLaDA15TEDotProductAttention(

config: megatron.core.transformer.transformer_config.TransformerConfig,

layer_number: int,

attn_mask_type: megatron.core.transformer.enums.AttnMaskType,

attention_type: str,

attention_dropout: Optional[float] = None,

**kwargs,

)

Bases: megatron.core.extensions.transformer_engine.TEDotProductAttention

TE-backed core attention for LLaDA1.5 masked-diffusion training/inference.

Overrides attn_mask_type to AttnMaskType.no_mask so the model sees fully bidirectional attention, matching modeling_llada.py’s get_bidirectional_attention_bias (a zero tensor) at line 1273.

Initialization

set_block_mask(mask: Optional[torch.Tensor]) → None

Install a boolean block-diagonal mask (True = blocked) for experimental use.

Not used by the default LLaDA1.5 inference loop; reserved for users who want to layer LLaDA2-style block-diagonal attention on top.

set_padding_mask(mask: Optional[torch.Tensor]) → None

Install a boolean key-padding mask [B, S] (True = padding token).

Required for correct batched generation: LLaDA1.5 attends fully bidirectionally, so without this every query would attend to left-pad tokens and corrupt short prompts in a mixed-length batch.

reset_inference_state() → None

Clear any stored inference state. Safe to call between generations.

forward(

query: torch.Tensor,

key: torch.Tensor,

value: torch.Tensor,

attention_mask: Optional[torch.Tensor] = None,

attn_mask_type: Optional[megatron.core.transformer.enums.AttnMaskType] = None,

attention_bias: Optional[torch.Tensor] = None,

packed_seq_params: Optional[megatron.core.packed_seq_params.PackedSeqParams] = None,

) → torch.Tensor