[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
co-authored by Tiwei Bie
parent 5795da5e83
commit 21b0582d4b
13 changed files with 1286 additions and 6 deletions
@@ -281,6 +281,36 @@ class SchedulerOutputProcessorMixin:
return predict_tokens
def process_batch_result_dllm(
self: Scheduler,
batch: ScheduleBatch,
result: GenerationBatchResult,
):
if result.copy_done is not None:
result.copy_done.synchronize()
next_token_ids = result.next_token_ids.tolist()
self.num_generated_tokens += len(next_token_ids)
self.token_to_kv_pool_allocator.free_group_begin()
assert len(batch.reqs) == 1, "batch size is currently expected to be 1"
req = batch.reqs[0]
for next_token_id in next_token_ids:
req.output_ids.append(next_token_id)
req.check_finished()
if req.finished():
release_kv_cache(req, self.tree_cache)
req.time_stats.completion_time = time.perf_counter()
break
self.tree_cache.cache_unfinished_req(req)
self.stream_output(batch.reqs, batch.return_logprob)
self.token_to_kv_pool_allocator.free_group_end()
def process_batch_result_decode(
self: Scheduler,
batch: ScheduleBatch,