bridge.diffusion.models.llada15.inference_llada15#

Block-diffusion generation for LLaDA1.5 loaded into a Megatron GPTModel.

Mirrors the official sampling loop in the ML-GSAI/LLaDA repo (generate.py): the prompt is concatenated with a sequence of <MASK> tokens, and the model is repeatedly invoked on the full sequence (with fully bidirectional attention — see :class:LLaDA15TEDotProductAttention) to predict the masked positions. Each iteration unmasks the most confident predictions inside the current block; once a block is fully unmasked the loop advances to the next block.

Note: unlike LLaDA2, no block-diagonal attention mask is constructed. The “block” structure is purely a sampling-time choice (which positions to unmask per step). The model itself sees the full sequence with zero attention bias.

Module Contents#

Functions#

_unwrap

Unwrap Float16Module / DDP / VLM wrappers to reach the raw GPTModel.

_iter_llada15_attentions

Yield each layer’s LLaDA15TEDotProductAttention core-attention module.

_set_padding_mask

Broadcast a boolean key-padding mask [B, S] to every attention layer.

_clear_attention_state

Drop any stored mask state so it does not leak into the next batch.

generate_block_diffusion

Sample tokens from a LLaDA1.5 Megatron GPTModel via block diffusion.

API#

bridge.diffusion.models.llada15.inference_llada15._unwrap(model)#

Unwrap Float16Module / DDP / VLM wrappers to reach the raw GPTModel.

bridge.diffusion.models.llada15.inference_llada15._iter_llada15_attentions(model)#

Yield each layer’s LLaDA15TEDotProductAttention core-attention module.

bridge.diffusion.models.llada15.inference_llada15._set_padding_mask(model, mask: Optional[torch.Tensor]) → None#

Broadcast a boolean key-padding mask [B, S] to every attention layer.

bridge.diffusion.models.llada15.inference_llada15._clear_attention_state(model) → None#

Drop any stored mask state so it does not leak into the next batch.

bridge.diffusion.models.llada15.inference_llada15.generate_block_diffusion(
model,
input_ids: torch.Tensor,
*,
gen_length: int = 256,
block_length: int = 32,
steps: int = 32,
temperature: float = 0.0,
remasking: str = 'low_confidence',
neg_entropy: bool = False,
threshold: Optional[float] = None,
mask_token_id: int = 126336,
eos_token_id: Optional[int] = 126081,
eos_early_stop: bool = False,
pad_token_id: Optional[int] = None,
) → torch.Tensor#

Sample tokens from a LLaDA1.5 Megatron GPTModel via block diffusion.

The model attends fully bidirectionally over the entire sequence at every step (LLaDA1.5’s reference uses a zero attention bias), so each iteration re-forwards the whole x. The block structure governs only which positions are eligible to be unmasked per step, not the attention pattern. Per-step token selection is delegated to the shared diffusion sampler in

Mod:

megatron.bridge.diffusion.common.dllm, the same primitives used by NemotronLabsDiffusion and verified by the generation-parity test.

The defaults (remasking="low_confidence", neg_entropy=False, threshold=None, temperature=0) reproduce greedy ML-GSAI sampling exactly. Set temperature > 0 for Gumbel sampling, neg_entropy=True to rank by distribution entropy, or threshold for confidence-gated transfer.

Parameters:

  • model – Megatron GPTModel built with :class:LLaDA15ModelProvider.

  • input_ids – Prompt tokens [B, prompt_len].

  • gen_length – Number of new tokens to generate.

  • block_length – Tokens unmasked per outer block iteration.

  • steps – Total denoising steps (split evenly across blocks).

  • temperature – Gumbel sampling temperature (0 = greedy).

  • remasking – Confidence source, "low_confidence" or "random".

  • neg_entropy – Rank by negative entropy instead of chosen-token probability.

  • threshold – Optional confidence threshold for gated transfer.

  • mask_token_id – LLaDA1.5 mask token id (default 126336).

  • eos_token_id – EOS id (default 126081) for early stopping.

  • eos_early_stop – Stop generation once every sample in the batch has emitted at least one EOS in its generated region. Evaluated at block boundaries (not per step), so the current block is always fully unmasked before stopping and no mask id survives before a sample’s EOS. Disable for fixed-length outputs.

  • pad_token_id – If given and the batched prompt contains padding, a boolean key-padding mask is installed so padded positions are never attended to. Required for correct mixed-length batched generation (LLaDA1.5 attends fully bidirectionally and has no implicit padding mask).

Returns:

Token ids [B, prompt_len + gen_length] — always full width. When eos_early_stop triggers, generation halts but the tensor is not truncated; positions past each sample’s first EOS keep their generated ids, so callers should trim at the first eos_token_id per row.