[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

@@ -240,6 +240,10 @@ class ServerArgs:
revision: Optional[str] = None
model_impl: str = "auto"
# Diffusion LLM
dllm_algorithm: Optional[str] = None
dllm_block_size: Optional[int] = None
# HTTP server
host: str = "127.0.0.1"
port: int = 30000
@@ -663,6 +667,9 @@ class ServerArgs:
# Handle exporting request-level metrics.
self._handle_request_metrics_exporters()
# Handle diffusion LLM inference.
self._handle_dllm_inference()
# Handle any other necessary validations.
self._handle_other_validations()
@@ -1974,6 +1981,30 @@ class ServerArgs:
"--export-metrics-to-file-dir is required when --export-metrics-to-file is enabled"
)
def _handle_dllm_inference(self):
if self.dllm_algorithm is None:
return
if not self.disable_cuda_graph:
logger.warning(
"Cuda graph is disabled because of using diffusion LLM inference"
)
self.disable_cuda_graph = True
if not self.disable_overlap_schedule:
logger.warning(
"Overlap schedule is disabled because of using diffusion LLM inference"
)
self.disable_overlap_schedule = True
if not self.disable_radix_cache:
logger.warning(
"Radix cache is disabled because of using diffusion LLM inference"
)
self.disable_radix_cache = True
if not self.pp_size > 1:
logger.warning(
"Pipeline parallelism is disabled because of using diffusion LLM inference"
)
self.pp_size = 1
def _handle_other_validations(self):
# Handle model inference tensor dump.
if self.debug_tensor_dump_output_folder is not None:
@@ -2093,6 +2124,20 @@ class ServerArgs:
"implementation.\n",
)
# Diffusion LLM
parser.add_argument(
"--dllm-algorithm",
type=str,
default=ServerArgs.dllm_algorithm,
help="The diffusion LLM algorithm.",
)
parser.add_argument(
"--dllm-block-size",
type=int,
default=ServerArgs.dllm_block_size,
help="The number of tokens processed in each iteration of the block diffusion LLM.",
)
# HTTP server
parser.add_argument(
"--host",