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