Beyond Block Boundaries: Multi-Block Editing for Diffusion Large Language Models

2026-06-29

AI Summary

This work addresses the degradation in generation quality of chunked diffusion language models at chunk boundaries due to missing cross-chunk context and their inability to correct errors in already generated chunks. The authors propose Multi-Chunk Editing (MBE), a training-free inference-time method that leverages cross-chunk context to revise previously decoded text. Key innovations include a supervised fine-tuning–based bidirectional attention mask, full-attention window reopening, progressive edit span expansion, multi-shape CUDA graph pooling, and fine-grained KV cache control. The approach extends SGLang to enable efficient variable-length editing. Evaluated across 12 benchmarks, MBE improves average performance from 61.45 to 64.24—yielding a notable gain of up to 13.33 points on AIME 2025—while retaining 87.3%–96.7% of standard decoding throughput.

Abstract

Block diffusion has emerged as the dominant paradigm for scaling discrete diffusion language models (dLLMs), because decoding text in fixed-size blocks preserves parallel generation within each block while keeping the quadratic attention cost tractable. However, this efficiency comes with a structural limitation: tokens near the end of a block are generated without access to future cross-block context, and once a block is finalized, its uncertain predictions become irreversible context for all subsequent blocks. This creates a block boundary problem, in which uncertainty accumulates toward block boundaries and early mistakes propagate throughout later generation. To address this issue, we propose Multi-Block Editing (MBE), to mitigate this problem by editing decoded tokens based on cross-block context. Following this principle, MBE first proposes a training-free decoding algorithm to edit the decoded tokens in previous blocks, which is achieved by re-opening a full-attention window over selected blocks. Given the mismatched attention mechanism between block diffusion training and MBE inference, MBE further introduces a supervised Fine-tuning strategy, which equips the model with bidirectional attention masks that progressively expands the editing span. Furthermore, it also extends SGLang with a multi-shape CUDA Graph pool and fine-grained KV cache control to make these variable-length editing passes efficient in practice. Experiments on LLaDA2.1-Mini across 13 benchmarks show that training-free MBE outperforms all existing decoding baselines while maintaining comparable throughput, and MBE SFT further brings a performance gain of 2.7. The largest improvements appear on tasks requiring strong long-range consistency, including +13.3 on AIME 2025 and +5.9 on ZebraLogic, demonstrating the effectiveness of MBE.

Problem

• block boundary problem
• discrete diffusion language models
• cross-block context
• parallel decoding
• quadratic attention

Innovation

• Multi-Block Editing
• block boundary problem
• diffusion language models
• training-free inference
• bidirectional attention