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