[diffusion] feat: enable passing Cache‑DiT config for diffusers backend (#16662)

Signed-off-by: Chi <chixie.mcisaac@gmail.com>
Signed-off-by: qimcis <chixie.mcisaac@gmail.com>
Co-authored-by: Mick <mickjagger19@icloud.com>
This commit is contained in:
Chi McIsaac
2026-01-22 00:13:34 -05:00
committed by GitHub
parent 17807caf82
commit 71482dd171
7 changed files with 157 additions and 16 deletions

View File

@@ -235,7 +235,9 @@ class ServerArgs:
# Attention
attention_backend: str = None
diffusers_attention_backend: str = None # for diffusers backend only
cache_dit_config: str | dict[str, Any] | None = (
None # cache-dit config for diffusers
)
# Distributed executor backend
nccl_port: Optional[int] = None
@@ -452,15 +454,25 @@ class ServerArgs:
"--attention-backend",
type=str,
default=None,
choices=[e.name.lower() for e in AttentionBackendEnum] + ["fa3", "fa4"],
help="The attention backend to use. If not specified, the backend is automatically selected based on hardware and installed packages.",
help=(
"The attention backend to use. For SGLang-native pipelines, use "
"values like fa, torch_sdpa, sage_attn, etc. For diffusers pipelines, "
"use diffusers attention backend names such as flash, _flash_3_hub, "
"sage, or xformers."
),
)
parser.add_argument(
"--diffusers-attention-backend",
type=str,
dest="attention_backend",
default=None,
help="Attention backend for diffusers pipelines (e.g., flash, _flash_3_hub, sage, xformers). "
"See: https://huggingface.co/docs/diffusers/main/en/optimization/attention_backends",
help=argparse.SUPPRESS,
)
parser.add_argument(
"--cache-dit-config",
type=str,
default=ServerArgs.cache_dit_config,
help="Path to a Cache-DiT YAML/JSON config. Enables cache-dit for diffusers backend.",
)
# HuggingFace specific parameters
@@ -999,7 +1011,7 @@ class ServerArgs:
raise ValueError("pipeline_config is not set in ServerArgs")
self.pipeline_config.check_pipeline_config()
if self.attention_backend is None:
if self.attention_backend is None and self.backend != Backend.DIFFUSERS:
self._set_default_attention_backend()
# parallelism