[feature] Initial block diffusion language model support (#12588)
Co-authored-by: Tiwei Bie <tiwei.btw@antgroup.com>
This commit is contained in:
@@ -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
|
||||
):
|
||||
|
||||
Reference in New Issue
Block a user