[feature] Initial block diffusion language model support (#12588)

Co-authored-by: Tiwei Bie <tiwei.btw@antgroup.com>
This commit is contained in:
Zehuan Li
2025-11-26 17:57:54 +08:00
committed by GitHub
parent 5795da5e83
commit 21b0582d4b
13 changed files with 1286 additions and 6 deletions

View File

@@ -441,8 +441,17 @@ class ForwardBatch:
)
return ret
# Override the positions with spec_info
if (
# Override the positions with diffusion LLM or spec_info
if batch.dllm_config is not None:
block_size = batch.dllm_config.block_size
ret.positions = torch.tensor(
[
[i for i in range(block_offset, block_offset + block_size)]
for block_offset in batch.dllm_block_offsets
],
dtype=torch.int32,
).to(device, non_blocking=True)
elif (
ret.spec_info is not None
and getattr(ret.spec_info, "positions", None) is not None
):