LLaDA2: EOS early stopping runs every remaining block to the step limit in batched generation #15242

Describe the bug

In batched LLaDA2Pipeline calls, EOS early stopping can make generation slower instead of faster. Two things cause it:

1. A padded row is never marked finished. BlockRefinementScheduler.check_eos_finished looks for the first EOS in the whole row, prompt included, and skips the row if that EOS is before prompt_length. The LLaDA2 tokenizer uses <|endoftext|> as both EOS and pad, so any shorter prompt in a batch has an EOS in its padding, and that row never finishes.
2. A finished row keeps later blocks running. The pipeline stops updating finished rows (final_transfer & ~finished[:, None]), so their masks are never filled. masks_remaining is computed over the whole batch, so it stays True, and every later block runs all num_inference_steps.

The decoded text is still correct. Only compute is wasted. In the repro below, one row finishing makes the batch take 130 model calls instead of 10.

Proposed fix: search for EOS only after prompt_length, and compute masks_remaining over unfinished rows only. I'd like to open a PR for this, with regression tests. #14741 moves check_eos_finished but keeps the same logic.

Reproduction

I couldn't run inclusionAI/LLaDA2.1-mini itself (30 GiB of weights, 16 GiB of RAM). This uses the real tokenizer and the real pipeline, with a stub model that confidently predicts EOS for some rows and an ordinary token for the others.

import torch
from transformers import AutoTokenizer

from diffusers import BlockRefinementScheduler, LLaDA2Pipeline


tok = AutoTokenizer.from_pretrained("inclusionAI/LLaDA2.1-mini")
EOS, OTHER = tok.eos_token_id, tok.convert_tokens_to_ids("Hello")


class StubModel(torch.nn.Module):
    """Stands in for LLaDA2.1-mini: confidently predicts EOS for rows in `eos_rows`, "Hello" for the others."""

    dtype = torch.float32

    def __init__(self, eos_rows):
        super().__init__()
        self.eos_rows, self.calls = eos_rows, 0
        self.register_buffer("_anchor", torch.empty(0))

    @property
    def device(self):
        return self._anchor.device

    def forward(self, input_ids, **kwargs):
        self.calls += 1
        logits = torch.zeros(*input_ids.shape, len(tok))
        logits[..., OTHER] = 20.0
        logits[self.eos_rows, :, OTHER] = 0.0
        logits[self.eos_rows, :, EOS] = 20.0
        return type("Out", (), {"logits": logits})()


def count_calls(prompts, eos_rows):
    model = StubModel(eos_rows)
    pipe = LLaDA2Pipeline(model=model, scheduler=BlockRefinementScheduler(), tokenizer=tok)
    pipe.set_progress_bar_config(disable=True)
    pipe(prompt=prompts, use_chat_template=False, gen_length=128, block_length=32, num_inference_steps=32)
    return model.calls


print(f"pad_token == eos_token: {tok.pad_token_id == EOS}")
print("1) both rows predict EOS, equal-length prompts:  ", count_calls(["Hello there", "Hi you"], [0, 1]))
print("2) both rows predict EOS, one prompt padded:     ", count_calls(["Hello there", "Hi"], [0, 1]))
print("3) no row predicts EOS:                          ", count_calls(["Hello there", "Hi you"], []))
print("4) only row 0 predicts EOS:                      ", count_calls(["Hello there", "Hi you"], [0]))

Context

I was reading the LLaDA2 pipeline and BlockRefinementScheduler to learn how block-wise text diffusion works, and noticed this while following how eos_early_stop decides when a row is finished. I haven't hit it in a real workload, but it affects any batched LLaDA2Pipeline call with the default eos_early_stop=True. With the LLaDA2 tokenizer, a batch of prompts with different lengths is enough to trigger it.

Logs

pad_token == eos_token: True
1) both rows predict EOS, equal-length prompts:   1
2) both rows predict EOS, one prompt padded:      130
3) no row predicts EOS:                           10
4) only row 0 predicts EOS:                       130

Expected: 1, 1, 10, and at most 10 for case 4.

System Info

• 🤗 Diffusers version: 0.42.0.dev0 (main @ f598e6e)
• Platform: macOS-26.6.2-arm64-arm-64bit
• Python version: 3.11.14
• PyTorch version (GPU?): 2.15.0 (False)
• Transformers version: 5.22.0
• Accelerator: Apple M1 Pro
• Using GPU in script?: No (CPU)
• Using distributed or parallel set-up in script?: No

Who can help?

@yiyixuxu @kashif