minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

[diffusion] feat: support cache-dit, cfg gating, attention backend override as per-request param (#35339)

Browse Source 
Co-authored-by: Claude Fable 5 <noreply@anthropic.com>
This commit is contained in:
Mick
2026-08-19 08:24:03 +08:00
committed by GitHub
co-authored by Claude Fable 5
parent 77fc5c128e
commit e73201e462
17 changed files with 1040 additions and 82 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/multimodal_gen/.claude/skills/sglang-diffusion-performance/SKILL.md
+3 -2
View File
@@ -57,8 +57,9 @@ These options **trade output quality** for speed or VRAM savings. Results will d
| Option | CLI Flag / Env Var | What It Does | Speedup | Quality Impact / Limitations |
|---|---|---|---|---|
| **Request Quality Fast Paths** | `--quality high` (`lossless` is default) | Mounts model-owned accelerated DiT/VAE paths that are validated for high quality but are not bit-exact to the reference path. | Model- and shape-specific | Support is per model and may be a no-op. Keep `--quality lossless` as the A/B ground truth. Do not confuse this with `--output-quality`, which controls file compression. |
| **Approximate Attention** | `--attention-backend sage_attn` / `sage_attn_3` / `sliding_tile_attn` / `video_sparse_attn` / `sparse_video_gen_2_attn` / `vmoba_attn` / `sla_attn` / `sage_sla_attn` | Replaces exact attention with approximate or sparse variants. `sage_attn`: INT8/FP8 quantized Q·K; `sliding_tile_attn`: spatial-temporal tile skipping; others: model-specific sparse patterns. | ~1.5–2x on attention (varies by backend) | Quality degradation varies by backend and model. `sage_attn` is the most general; sparse backends (`sliding_tile_attn`, `video_sparse_attn`, etc.) are video-model-specific and may require config files (e.g. `--mask-strategy-file-path` for STA). Requires corresponding packages installed. |
| **Cache-DiT** | Native: `SGLANG_CACHE_DIT_ENABLED=true` plus `SGLANG_CACHE_DIT_*` env vars. Diffusers backend: `--backend diffusers --cache-dit-config <yaml-or-json>` | Caches intermediate residuals across denoising steps and skips redundant computations via DBCache, TaylorSeer, and optional SCM. | ~1.5-2x on supported models | Quality depends on cache policy. Incompatible with `--dit-layerwise-offload`. Do not pass `--cache-dit-config` for native SGLang tuning unless you are intentionally using the diffusers backend flow. |
| **Approximate Attention** | Server-wide: `--attention-backend sage_attn` / `sage_attn_3` / `sliding_tile_attn` / `video_sparse_attn` / `sparse_video_gen_2_attn` / `vmoba_attn` / `sla_attn` / `sage_sla_attn`. Per-request (dense drop-ins only): `--attention-backend-override sage_attn` sampling param / API `extra_body` — valid values `fa`, `torch_sdpa`, `sage_attn`, `sage_attn_3`; rejected (with a log) under BCG, torch.compile, sparse server backends, or a non-ring-capable target with ring parallelism. | Replaces exact attention with approximate or sparse variants. `sage_attn`: INT8/FP8 quantized Q·K; `sliding_tile_attn`: spatial-temporal tile skipping; others: model-specific sparse patterns. | ~1.5–2x on attention (varies by backend) | Quality degradation varies by backend and model. `sage_attn` is the most general; sparse backends (`sliding_tile_attn`, `video_sparse_attn`, etc.) are video-model-specific, may require config files (e.g. `--mask-strategy-file-path` for STA), and are server-level only. Requires corresponding packages installed. |
| **Cache-DiT** | Native: per-request `--enable-cache-dit true\|false` + `--cache-dit-params <json>` (sampling params; also via API `extra_body`). `SGLANG_CACHE_DIT_ENABLED` / `SGLANG_CACHE_DIT_*` env vars are the server-wide defaults for requests that leave them unset. Diffusers backend: `--backend diffusers --cache-dit-config <yaml-or-json>` | Caches intermediate residuals across denoising steps and skips redundant computations via DBCache, TaylorSeer, and optional SCM. | ~1.5-2x on supported models | Quality depends on cache policy. Incompatible with `--dit-layerwise-offload`. Do not pass `--cache-dit-config` for native SGLang tuning unless you are intentionally using the diffusers backend flow. |
| **CFG Gating** | Per-request `--cfg-gate-step 0.5` (sampling param; also via API `extra_body`). `SGLANG_DIFFUSION_CFG_GATE_STEP` is the server-wide default (1.0 = off). | After the given fraction of denoising steps, reuses the cached cond-uncond residual instead of running the unconditional branch each step. | Up to ~2x on the gated tail of CFG models (skips one of two branches) | Lossy; no-op without classifier-free guidance or with `--enable-cfg-parallel`. Lower fractions gate earlier and drift more. |
| **TeaCache** | `--enable-teacache` (uses model sampling presets) | Reuses residuals when adjacent denoising steps are sufficiently similar. | Model- and threshold-dependent | Approximate and model-specific. Mutually exclusive with Spectrum. Fix prompt/seed/shape/steps and validate temporal consistency, not only single frames. |
| **Spectrum** | `--enable-spectrum` plus optional `--spectrum-*` controls | Forecasts DiT features and skips selected denoising steps. | Defaults target an accuracy/speed tradeoff; aggressive windows can be much faster | Native `sglang generate` only for FLUX.1, Wan, HunyuanVideo, and SD3; not FLUX.2 or server requests. Mutually exclusive with TeaCache. `--debug` adds shadow validation and is not representative latency. |
| **Progressive Resolution** | `--progressive-mode dct_rewind --progressive-levels N --progressive-delta D` | Runs early denoising at lower latent resolution, then spectrally upsamples and switches to the target resolution. | Model- and schedule-dependent | Approximate and pipeline-specific. Keep the switch schedule fixed and compare detail, composition, and temporal stability. |
python/sglang/multimodal_gen/configs/sample/sampling_params.py
+32
View File
@@ -206,6 +206,10 @@ class SamplingParams:
guidance_rescale: float = 0.0
cfg_normalization: float | bool = 0.0
boundary_ratio: float | None = None
# CFG gating (lossy): reuse the cached cond-uncond residual after this
# fraction of the steps. None = follow the SGLANG_DIFFUSION_CFG_GATE_STEP
# server default; 1.0 = off for this request.
cfg_gate_step: float | None = None
progressive_mode: str = "fullres"
progressive_levels: int = 1
@@ -222,6 +226,18 @@ class SamplingParams:
None # TeaCacheParams or WanTeaCacheParams, set by model-specific subclass
)
# Cache-DiT (lossy). None = follow the SGLANG_CACHE_DIT_ENABLED server
# default; True/False = explicit per-request opt-in/out.
enable_cache_dit: bool | None = None
# Per-request knob overrides on top of the SGLANG_CACHE_DIT_* defaults.
# Valid keys: CACHE_DIT_REQUEST_PARAM_KEYS in cache_dit_integration.py.
cache_dit_params: dict[str, Any] | None = None
# Per-request DiT attention backend ("fa", "torch_sdpa", "sage_attn",
# "sage_attn_3"; sage is lossy). Incompatible server settings reject the
# request; see DenoisingStage._maybe_override_attention_backend.
attention_backend_override: str | None = None
# Spectrum parameters
enable_spectrum: bool = False
spectrum_params: Any = None # SpectrumParams
@@ -908,6 +924,22 @@ class SamplingParams:
"--enable-teacache",
action="store_true",
)
add_argument(
"--enable-cache-dit",
action=StoreBoolean,
)
add_argument(
"--cache-dit-params",
type=json.loads,
)
add_argument(
"--cfg-gate-step",
type=float,
)
add_argument(
"--attention-backend-override",
type=str,
)
add_argument(
"--enable-spectrum",
action="store_true",
python/sglang/multimodal_gen/runtime/cache/cache_dit_integration.py
Vendored
+69
View File
@@ -192,6 +192,75 @@ def get_scm_mask(
return mask
# Keys accepted in SamplingParams.cache_dit_params; "secondary" nests the
# DBCache knobs for the second transformer of dual-DiT models.
CACHE_DIT_REQUEST_KNOB_KEYS = frozenset(
{
"Fn_compute_blocks",
"Bn_compute_blocks",
"max_warmup_steps",
"residual_diff_threshold",
"max_continuous_cached_steps",
"enable_taylorseer",
"taylorseer_order",
}
)
CACHE_DIT_REQUEST_SCM_KEYS = frozenset(
{
"scm_preset",
"scm_compute_bins",
"scm_cache_bins",
"scm_policy",
}
)
CACHE_DIT_REQUEST_PARAM_KEYS = (
CACHE_DIT_REQUEST_KNOB_KEYS | CACHE_DIT_REQUEST_SCM_KEYS | {"secondary"}
)
def resolve_cache_dit_request_overrides(raw: dict | None) -> dict:
"""Validate cache_dit_params and return a copy; unknown keys fail the request."""
if raw is None:
return {}
if not isinstance(raw, dict):
raise ValueError(f"cache_dit_params must be a dict, got {type(raw).__name__}.")
unknown = set(raw) - CACHE_DIT_REQUEST_PARAM_KEYS
if unknown:
raise ValueError(
f"Unknown cache_dit_params keys: {sorted(unknown)}. "
f"Valid keys: {sorted(CACHE_DIT_REQUEST_PARAM_KEYS)}."
)
overrides = dict(raw)
secondary = overrides.get("secondary")
if secondary is not None:
if not isinstance(secondary, dict):
raise ValueError(
"cache_dit_params['secondary'] must be a dict, got "
f"{type(secondary).__name__}."
)
unknown = set(secondary) - CACHE_DIT_REQUEST_KNOB_KEYS
if unknown:
raise ValueError(
f"Unknown cache_dit_params['secondary'] keys: {sorted(unknown)}. "
f"Valid keys: {sorted(CACHE_DIT_REQUEST_KNOB_KEYS)}."
)
overrides["secondary"] = dict(secondary)
return overrides
def cache_dit_overrides_key(overrides: dict) -> tuple:
"""Hashable snapshot of request overrides, for mount-change detection."""
def _freeze(value):
if isinstance(value, dict):
return tuple(sorted((k, _freeze(v)) for k, v in value.items()))
if isinstance(value, (list, tuple)):
return tuple(_freeze(v) for v in value)
return value
return _freeze(overrides)
@dataclass
class CacheDitConfig:
"""Configuration for cache-dit integration.
python/sglang/multimodal_gen/runtime/entrypoints/openai/image_api.py
+6
View File
@@ -314,6 +314,12 @@ async def generations(
use_system_prompt=_get_extra_field(request, "use_system_prompt"),
use_guardrails=_get_extra_field(request, "use_guardrails"),
enable_teacache=request.enable_teacache,
enable_cache_dit=_get_extra_field(request, "enable_cache_dit"),
cache_dit_params=_get_extra_field(request, "cache_dit_params"),
cfg_gate_step=_get_extra_field(request, "cfg_gate_step"),
attention_backend_override=_get_extra_field(
request, "attention_backend_override"
),
quality=_runtime_sampling_quality(request.quality),
output_compression=request.output_compression,
output_quality=request.output_quality,
python/sglang/multimodal_gen/runtime/entrypoints/openai/video_api.py
+6
View File
@@ -390,6 +390,12 @@ def _build_video_sampling_params(request_id: str, request: VideoGenerationsReque
"use_system_prompt": _extra_value(request, "use_system_prompt"),
"use_guardrails": _extra_value(request, "use_guardrails"),
"enable_teacache": request.enable_teacache,
"enable_cache_dit": _extra_value(request, "enable_cache_dit"),
"cache_dit_params": _extra_value(request, "cache_dit_params"),
"cfg_gate_step": _extra_value(request, "cfg_gate_step"),
"attention_backend_override": _extra_value(
request, "attention_backend_override"
),
"enable_frame_interpolation": request.enable_frame_interpolation,
"frame_interpolation_exp": request.frame_interpolation_exp,
"frame_interpolation_scale": request.frame_interpolation_scale,
python/sglang/multimodal_gen/runtime/layers/attention/layer.py
+50 -3
View File
@@ -284,6 +284,41 @@ class DynamicVarlenMaskMeta:
return self._meta
def prepare_attention_backend_override(
layer: nn.Module, target: AttentionBackendEnum
) -> None:
"""Build and cache the impl for ``target``; may raise, mutates nothing."""
if target in layer._attn_impl_by_backend:
return
backend_cls = get_attn_backend(
layer.head_size,
layer.dtype,
supported_attention_backends=layer._supported_attention_backends,
selected_attention_backend=target,
)
resolved = backend_cls.get_enum()
if resolved is not target:
raise ValueError(
f"Attention backend override '{target}' resolved to '{resolved}' on "
f"{type(layer).__name__}; refusing the request instead of silently "
"falling back."
)
impl = backend_cls.get_impl_cls()(**layer._attn_impl_ctor_kwargs)
wrap_attention_impl_forward(impl)
layer._attn_impl_by_backend[target] = impl
def apply_attention_backend_override(
layer: nn.Module, target: AttentionBackendEnum | None
) -> None:
"""Flip to a prepared impl (None = construction default); cannot fail."""
target = target or layer._default_attn_backend
if target is layer.backend:
return
layer.attn_impl = layer._attn_impl_by_backend[target]
layer.backend = target
class UlyssesAttention(nn.Module):
"""Ulysses-style SequenceParallelism attention layer."""
@@ -321,7 +356,7 @@ class UlyssesAttention(nn.Module):
)
impl_cls = attn_backend.get_impl_cls()
self.attn_impl = impl_cls(
self._attn_impl_ctor_kwargs = dict(
num_heads=num_heads,
head_size=head_size,
causal=causal,
@@ -330,11 +365,15 @@ class UlyssesAttention(nn.Module):
prefix=f"{prefix}.impl",
**extra_impl_args,
)
self.attn_impl = impl_cls(**self._attn_impl_ctor_kwargs)
wrap_attention_impl_forward(self.attn_impl)
self.num_heads = num_heads
self.head_size = head_size
self.num_kv_heads = num_kv_heads
self.backend = attn_backend.get_enum()
self._default_attn_backend = self.backend
self._attn_impl_by_backend = {self.backend: self.attn_impl}
self._supported_attention_backends = supported_attention_backends
self.dtype = dtype
self.causal = causal
self.sp_attention_mode, self.sp_attention_mode_is_auto = (
@@ -575,7 +614,7 @@ class LocalAttention(nn.Module):
)
impl_cls = attn_backend.get_impl_cls()
self.allow_cudnn_sdp = bool(extra_impl_args.get("allow_cudnn_sdp", False))
self.attn_impl = impl_cls(
self._attn_impl_ctor_kwargs = dict(
num_heads=num_heads,
head_size=head_size,
softmax_scale=self.softmax_scale,
@@ -583,11 +622,15 @@ class LocalAttention(nn.Module):
causal=causal,
**extra_impl_args,
)
self.attn_impl = impl_cls(**self._attn_impl_ctor_kwargs)
wrap_attention_impl_forward(self.attn_impl)
self.num_heads = num_heads
self.head_size = head_size
self.num_kv_heads = num_kv_heads
self.backend = attn_backend.get_enum()
self._default_attn_backend = self.backend
self._attn_impl_by_backend = {self.backend: self.attn_impl}
self._supported_attention_backends = supported_attention_backends
self.dtype = dtype
def forward(
@@ -729,7 +772,7 @@ class USPAttention(nn.Module):
)
impl_cls: Type[AttentionImpl] = attn_backend.get_impl_cls()
self.allow_cudnn_sdp = bool(extra_impl_args.get("allow_cudnn_sdp", False))
self.attn_impl = impl_cls(
self._attn_impl_ctor_kwargs = dict(
num_heads=num_heads,
head_size=head_size,
causal=causal,
@@ -738,11 +781,15 @@ class USPAttention(nn.Module):
prefix=f"{prefix}.impl",
**extra_impl_args,
)
self.attn_impl = impl_cls(**self._attn_impl_ctor_kwargs)
wrap_attention_impl_forward(self.attn_impl)
self.num_heads = num_heads
self.head_size = head_size
self.num_kv_heads = num_kv_heads
self.backend = attn_backend.get_enum()
self._default_attn_backend = self.backend
self._attn_impl_by_backend = {self.backend: self.attn_impl}
self._supported_attention_backends = supported_attention_backends
self.dtype = dtype
self.causal = causal
self.dropout_p = dropout_rate
python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising.py
+249 -46
View File
@@ -44,11 +44,14 @@ from sglang.multimodal_gen.runtime.breakable_cuda_graph import (
)
from sglang.multimodal_gen.runtime.cache.cache_dit_integration import (
CacheDitConfig,
cache_dit_overrides_key,
disable_cache_on_transformer,
enable_cache_on_dual_transformer,
enable_cache_on_transformer,
get_scm_mask,
refresh_context_on_dual_transformer,
refresh_context_on_transformer,
resolve_cache_dit_request_overrides,
)
from sglang.multimodal_gen.runtime.disaggregation.roles import RoleType
from sglang.multimodal_gen.runtime.distributed import (
@@ -76,6 +79,13 @@ from sglang.multimodal_gen.runtime.distributed.parallel_state import (
get_classifier_free_guidance_world_size,
world_group_is_initialized,
)
from sglang.multimodal_gen.runtime.layers.attention.layer import (
LocalAttention,
UlyssesAttention,
USPAttention,
apply_attention_backend_override,
prepare_attention_backend_override,
)
from sglang.multimodal_gen.runtime.layers.attention.selector import get_attn_backend
from sglang.multimodal_gen.runtime.layers.attention.STA_configuration import (
configure_sta,
@@ -232,6 +242,18 @@ class DenoisingStepState:
attn_metadata: Any | None
# Only exact/drop-in dense kernels: the sparse family needs per-model mask
# configs and per-step metadata, and stays a server-level choice.
REQUEST_SWITCHABLE_ATTENTION_BACKENDS = frozenset(
{
AttentionBackendEnum.FA,
AttentionBackendEnum.TORCH_SDPA,
AttentionBackendEnum.SAGE_ATTN,
AttentionBackendEnum.SAGE_ATTN_3,
}
)
class DualTransformerExecutionMode(str, Enum):
"""How a denoising stage uses a second DiT.
@@ -264,6 +286,11 @@ class DenoisingStage(PipelineStage, RolloutDenoisingMixin):
# cache-dit state (for delayed mounting and idempotent control)
self._cache_dit_enabled = False
self._cached_num_steps = None
# Per-request Cache-DiT overrides for the batch being executed
# (stashed by _maybe_enable_cache_dit; read by the config builders).
self._cache_dit_request_overrides: dict[str, Any] = {}
# Overrides key the mounted hooks were built from; None when unmounted.
self._cache_dit_active_key: tuple | None = None
# Whether request-scoped quality="high" fusions are currently mounted.
self._quality_fusions_mounted = False
self._torch_compile_registry = CompiledModuleRegistry()
@@ -300,6 +327,10 @@ class DenoisingStage(PipelineStage, RolloutDenoisingMixin):
dtype=torch.float16,
selected_attention_backend=selected_attention_backend,
)
# head size kept so the metadata backend can be re-resolved per batch
self._attn_backend_default = self.attn_backend
self._attn_metadata_head_size = attn_head_size
self._attention_backend_active_override: AttentionBackendEnum | None = None
# cfg
self.guidance = None
@@ -487,11 +518,118 @@ class DenoisingStage(PipelineStage, RolloutDenoisingMixin):
self, num_inference_steps: int | tuple[int, int], batch: Req
) -> None:
"""Apply request-dependent transformer acceleration in trace-safe order."""
self._maybe_override_attention_backend(batch)
self._maybe_toggle_quality_fusions(batch)
self._maybe_enable_cache_dit(num_inference_steps, batch)
for transformer in filter(None, [self.transformer, self.transformer_2]):
self._maybe_torch_compile(transformer)
def _maybe_override_attention_backend(self, batch: Req) -> None:
"""Two-phase per-request backend switch: prepare all layers (may
raise, mutates nothing), then flip all — a rejected request leaves the
transformers untouched. Safe at this batch boundary because the field
is in the dynamic-batch signature."""
target = self._parse_attention_backend_override(
batch.sampling_params.attention_backend_override
)
if target == self._attention_backend_active_override:
return
layers = self._request_switchable_attention_layers()
stage_backend = self._attn_backend_default
if target is not None:
stage_backend = self._validate_attention_backend_override(target, layers)
for layer in layers:
prepare_attention_backend_override(layer, target)
for layer in layers:
apply_attention_backend_override(layer, target)
self.attn_backend = stage_backend
self._attention_backend_active_override = target
logger.info(
"Attention backend for this batch: %s (%d layers switched)",
target.name.lower() if target else "server default",
len(layers),
)
def _parse_attention_backend_override(
self, name: str | None
) -> AttentionBackendEnum | None:
if name is None:
return None
try:
target = AttentionBackendEnum[name.upper()]
except KeyError:
raise ValueError(
f"Unknown attention_backend_override {name!r}. Valid values: "
f"{sorted(b.name.lower() for b in REQUEST_SWITCHABLE_ATTENTION_BACKENDS)}."
) from None
if target not in REQUEST_SWITCHABLE_ATTENTION_BACKENDS:
raise ValueError(
f"attention_backend_override {name!r} is not switchable per "
f"request. Valid values: "
f"{sorted(b.name.lower() for b in REQUEST_SWITCHABLE_ATTENTION_BACKENDS)}."
)
return target
def _request_switchable_attention_layers(self) -> list[nn.Module]:
return [
module
for transformer in filter(None, [self.transformer, self.transformer_2])
for module in transformer.modules()
if isinstance(module, (LocalAttention, UlyssesAttention, USPAttention))
]
def _validate_attention_backend_override(
self, target: AttentionBackendEnum, layers: list[nn.Module]
) -> type:
"""Reject incompatible server settings; returns the resolved backend cls."""
args = self.server_args
reasons: list[str] = []
if args.enable_breakable_cuda_graph:
reasons.append("breakable CUDA graphs bake the attention kernel in")
if args.enable_torch_compile:
reasons.append("torch.compile traces the attention kernel in")
if not layers:
reasons.append("this model exposes no switchable attention layers")
sparse_defaults = sorted(
{
layer._default_attn_backend.name.lower()
for layer in layers
if layer._default_attn_backend.is_sparse
}
)
if sparse_defaults:
reasons.append(
f"the server-selected sparse backend(s) {sparse_defaults} cannot "
"be mixed with per-request dense switching"
)
stage_backend = None
if not reasons:
try:
stage_backend = get_attn_backend(
head_size=self._attn_metadata_head_size,
dtype=torch.float16,
selected_attention_backend=target,
)
except ValueError as exc:
reasons.append(str(exc))
if stage_backend is not None:
if (
args.ring_degree or 1
) > 1 and not stage_backend.supports_ring_rotation():
reasons.append(
f"ring parallelism requires a ring-capable backend; "
f"{target.name.lower()} is not"
)
if reasons:
message = (
f"Rejecting attention_backend_override={target.name.lower()!r}: "
+ "; ".join(reasons)
+ "."
)
logger.warning(message)
raise ValueError(message)
return stage_backend
def _maybe_toggle_quality_fusions(self, batch: Req) -> None:
"""Mount/unmount the ``quality="high"`` fusions for this batch.
@@ -522,8 +660,24 @@ class DenoisingStage(PipelineStage, RolloutDenoisingMixin):
return "wan2.2"
def _cache_dit_requested(self) -> bool:
"""Request-independent server default (init-time decisions only)."""
return envs.SGLANG_CACHE_DIT_ENABLED
def _cache_dit_requested_for_batch(self, batch: Req) -> bool:
"""Per-request Cache-DiT switch; the server default applies when unset."""
enable = batch.sampling_params.enable_cache_dit
if enable is None:
return self._cache_dit_requested()
return enable
def _unmount_cache_dit(self) -> None:
"""Remove Cache-DiT hooks so subsequent batches run the native forward."""
for transformer in filter(None, [self.transformer, self.transformer_2]):
disable_cache_on_transformer(transformer)
self._cache_dit_enabled = False
self._cached_num_steps = None
self._cache_dit_active_key = None
def _cache_dit_secondary_uses_primary_config(self) -> bool:
return False
@@ -546,9 +700,22 @@ class DenoisingStage(PipelineStage, RolloutDenoisingMixin):
return steps, steps
raise ValueError("Boundary-expert dual transformers require split step counts.")
@staticmethod
def _parse_cache_dit_scm_bins() -> tuple[list[int] | None, list[int] | None, str]:
scm_preset = envs.SGLANG_CACHE_DIT_SCM_PRESET
def _parse_cache_dit_scm_bins(
self,
) -> tuple[list[int] | None, list[int] | None, str]:
overrides = self._cache_dit_request_overrides
scm_preset = overrides.get("scm_preset", envs.SGLANG_CACHE_DIT_SCM_PRESET)
request_compute_bins = overrides.get("scm_compute_bins")
request_cache_bins = overrides.get("scm_cache_bins")
if request_compute_bins is not None or request_cache_bins is not None:
if request_compute_bins is None or request_cache_bins is None:
raise ValueError(
"cache_dit_params SCM custom bins require both "
"scm_compute_bins and scm_cache_bins."
)
compute_bins = [int(x) for x in request_compute_bins]
cache_bins = [int(x) for x in request_cache_bins]
return compute_bins, cache_bins, scm_preset
compute_bins_str = envs.SGLANG_CACHE_DIT_SCM_COMPUTE_BINS
cache_bins_str = envs.SGLANG_CACHE_DIT_SCM_CACHE_BINS
compute_bins = None
@@ -574,7 +741,9 @@ class DenoisingStage(PipelineStage, RolloutDenoisingMixin):
self, primary_num_steps: int, secondary_num_steps: int | None = None
) -> tuple[str, str, list[int] | None, list[int] | None]:
scm_compute_bins, scm_cache_bins, scm_preset = self._parse_cache_dit_scm_bins()
scm_policy = envs.SGLANG_CACHE_DIT_SCM_POLICY
scm_policy = self._cache_dit_request_overrides.get(
"scm_policy", envs.SGLANG_CACHE_DIT_SCM_POLICY
)
steps_computation_mask = get_scm_mask(
preset=scm_preset,
num_inference_steps=primary_num_steps,
@@ -598,50 +767,70 @@ class DenoisingStage(PipelineStage, RolloutDenoisingMixin):
)
return scm_preset, scm_policy, steps_computation_mask, steps_computation_mask_2
@staticmethod
def _cache_dit_knob(
self, key: str, env_value: Any, env_secondary_value: Any, *, secondary: bool
) -> Any:
"""One DBCache knob: request > env; secondary inherits request primary."""
overrides = self._cache_dit_request_overrides
if not secondary:
return overrides.get(key, env_value)
secondary_overrides = overrides.get("secondary") or {}
if key in secondary_overrides:
return secondary_overrides[key]
return overrides.get(key, env_secondary_value)
def _build_cache_dit_config(
self,
num_inference_steps: int,
steps_computation_mask: list[int] | None,
scm_policy: str,
*,
secondary: bool = False,
) -> CacheDitConfig:
knob = self._cache_dit_knob
return CacheDitConfig(
enabled=True,
Fn_compute_blocks=(
envs.SGLANG_CACHE_DIT_SECONDARY_FN
if secondary
else envs.SGLANG_CACHE_DIT_FN
Fn_compute_blocks=knob(
"Fn_compute_blocks",
envs.SGLANG_CACHE_DIT_FN,
envs.SGLANG_CACHE_DIT_SECONDARY_FN,
secondary=secondary,
),
Bn_compute_blocks=(
envs.SGLANG_CACHE_DIT_SECONDARY_BN
if secondary
else envs.SGLANG_CACHE_DIT_BN
Bn_compute_blocks=knob(
"Bn_compute_blocks",
envs.SGLANG_CACHE_DIT_BN,
envs.SGLANG_CACHE_DIT_SECONDARY_BN,
secondary=secondary,
),
max_warmup_steps=(
envs.SGLANG_CACHE_DIT_SECONDARY_WARMUP
if secondary
else envs.SGLANG_CACHE_DIT_WARMUP
max_warmup_steps=knob(
"max_warmup_steps",
envs.SGLANG_CACHE_DIT_WARMUP,
envs.SGLANG_CACHE_DIT_SECONDARY_WARMUP,
secondary=secondary,
),
residual_diff_threshold=(
envs.SGLANG_CACHE_DIT_SECONDARY_RDT
if secondary
else envs.SGLANG_CACHE_DIT_RDT
residual_diff_threshold=knob(
"residual_diff_threshold",
envs.SGLANG_CACHE_DIT_RDT,
envs.SGLANG_CACHE_DIT_SECONDARY_RDT,
secondary=secondary,
),
max_continuous_cached_steps=(
envs.SGLANG_CACHE_DIT_SECONDARY_MC
if secondary
else envs.SGLANG_CACHE_DIT_MC
max_continuous_cached_steps=knob(
"max_continuous_cached_steps",
envs.SGLANG_CACHE_DIT_MC,
envs.SGLANG_CACHE_DIT_SECONDARY_MC,
secondary=secondary,
),
enable_taylorseer=(
envs.SGLANG_CACHE_DIT_SECONDARY_TAYLORSEER
if secondary
else envs.SGLANG_CACHE_DIT_TAYLORSEER
enable_taylorseer=knob(
"enable_taylorseer",
envs.SGLANG_CACHE_DIT_TAYLORSEER,
envs.SGLANG_CACHE_DIT_SECONDARY_TAYLORSEER,
secondary=secondary,
),
taylorseer_order=(
envs.SGLANG_CACHE_DIT_SECONDARY_TS_ORDER
if secondary
else envs.SGLANG_CACHE_DIT_TS_ORDER
taylorseer_order=knob(
"taylorseer_order",
envs.SGLANG_CACHE_DIT_TS_ORDER,
envs.SGLANG_CACHE_DIT_SECONDARY_TS_ORDER,
secondary=secondary,
),
num_inference_steps=num_inference_steps,
steps_computation_mask=steps_computation_mask,
@@ -651,24 +840,35 @@ class DenoisingStage(PipelineStage, RolloutDenoisingMixin):
def _maybe_enable_cache_dit(
self, num_inference_steps: int | tuple[int, int], batch: Req
) -> None:
"""Enable cache-dit on the transformers if configured (idempotent).
This method should be called after the transformer is fully loaded
and before torch.compile is applied.
For dual-transformer models (e.g., Wan2.2), this enables cache-dit on both
transformers with (potentially) different configurations.
"""Enable cache-dit on the transformers for this batch (idempotent).
Must run after the transformer is fully loaded and before
torch.compile. Per-request switch/knobs (env values are the defaults);
both fields are in the dynamic-batch signature, so mount/unmount
transitions are safe at this batch boundary. Dual-transformer models
(e.g. Wan2.2) get per-transformer configs.
"""
requested = self._cache_dit_requested_for_batch(batch)
if self.server_args.enable_breakable_cuda_graph:
# Cache-DiT wraps transformer.forward with step-skipping control
# flow that must not be baked into a captured CUDA graph.
if self._cache_dit_requested():
if requested:
logger.warning_once(
"Cache-DiT was requested but is disabled because breakable "
"CUDA graphs are enabled."
)
return
self._cache_dit_request_overrides = resolve_cache_dit_request_overrides(
batch.sampling_params.cache_dit_params
)
desired_key = (
cache_dit_overrides_key(self._cache_dit_request_overrides)
if requested
else None
)
# opt-out or knob change: unmount, then remount below (or return)
if self._cache_dit_enabled and desired_key != self._cache_dit_active_key:
self._unmount_cache_dit()
# NOTE: When a new request arrives, we need to refresh the cache-dit context.
if self._cache_dit_enabled:
primary_num_steps, secondary_num_steps = self._cache_dit_step_counts(
@@ -697,10 +897,10 @@ class DenoisingStage(PipelineStage, RolloutDenoisingMixin):
)
return
if not requested:
return
# Keep cache-dit disabled for ordinary warmup, but allow torch.compile
# warmup to mount cache-dit before Dynamo traces the transformer.
if not self._cache_dit_requested():
return
if batch.is_warmup and not getattr(
self.server_args, "enable_torch_compile", False
):
@@ -794,6 +994,7 @@ class DenoisingStage(PipelineStage, RolloutDenoisingMixin):
self._cache_dit_enabled = True
self._cached_num_steps = num_inference_steps
self._cache_dit_active_key = desired_key
@lru_cache(maxsize=8)
def _build_guidance(self, batch_size, target_dtype, device, guidance_val):
@@ -1107,11 +1308,13 @@ class DenoisingStage(PipelineStage, RolloutDenoisingMixin):
self, ctx: DenoisingContext, batch: Req, server_args: ServerArgs
) -> None:
"""Initialize optional CFG residual reuse for the current denoising loop."""
fraction = envs.SGLANG_DIFFUSION_CFG_GATE_STEP
fraction = batch.sampling_params.cfg_gate_step
if fraction is None:
fraction = envs.SGLANG_DIFFUSION_CFG_GATE_STEP
if not 0.0 <= fraction <= 1.0:
raise ValueError(
"SGLANG_DIFFUSION_CFG_GATE_STEP must be between 0.0 and 1.0, "
f"got {fraction}."
"cfg_gate_step (SGLANG_DIFFUSION_CFG_GATE_STEP) must be between "
f"0.0 and 1.0, got {fraction}."
)
num_steps = len(ctx.timesteps)
python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/ltx_2/denoising.py
+8 -8
View File
@@ -178,6 +178,8 @@ class LTX2DenoisingStage(DenoisingStage):
transformer=transformer, scheduler=scheduler, vae=vae, **kwargs
)
self.sampler_name = sampler_name
# set per request by _prepare_denoising_loop before the cache-dit hook
self._disable_cache_dit_for_request = False
def _scheduler_step_kwargs(self, batch: Req, scheduler) -> dict:
return self.prepare_extra_func_kwargs(
@@ -444,14 +446,12 @@ class LTX2DenoisingStage(DenoisingStage):
)
return latents[:, :orig_s, :].contiguous()
def _maybe_enable_cache_dit(self, num_inference_steps: int, batch: Req) -> None:
"""Disable cache-dit for TI2V-style requests to avoid stale activations.
NOTE: base denoising stage calls this hook with (num_inference_steps, batch).
"""
if getattr(self, "_disable_cache_dit_for_request", False):
return
return super()._maybe_enable_cache_dit(num_inference_steps, batch)
def _cache_dit_requested_for_batch(self, batch: Req) -> bool:
"""TI2V requests must not cache stale activations; reporting "not
requested" lets the base stage unmount hooks a prior request left."""
if self._disable_cache_dit_for_request:
return False
return super()._cache_dit_requested_for_batch(batch)
def _get_ltx2_stage1_guider_params(
self, batch: Req, server_args: ServerArgs, stage: str
python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/minimax_h3/stages/denoising.py
+15 -8
View File
@@ -404,12 +404,20 @@ class MiniMaxH3DenoisingStage(DenoisingStage):
) -> None:
quality = getattr(batch.sampling_params, "quality", "lossless")
explicit_fields = getattr(batch.sampling_params, "_explicit_fields", ())
generic_requested = (
super()._cache_dit_requested() and "quality" not in explicit_fields
)
desired_mode = (
"high" if quality == "high" else ("generic" if generic_requested else None)
enable_override = batch.sampling_params.enable_cache_dit
generic_enabled = (
super()._cache_dit_requested()
if enable_override is None
else enable_override
)
generic_requested = generic_enabled and "quality" not in explicit_fields
if enable_override is False:
# The per-request kill switch wins over quality="high".
desired_mode = None
elif quality == "high":
desired_mode = "high"
else:
desired_mode = "generic" if generic_requested else None
current_mode = getattr(self, "_minimax_h3_cache_mode", None)
self._minimax_h3_quality = quality
@@ -426,9 +434,7 @@ class MiniMaxH3DenoisingStage(DenoisingStage):
# Cache-DiT still holds references to their inputs. Settle the state
# fields before restoring the in-place path, so a failure there
# costs throughput rather than leaving the stage inconsistent.
self.transformer = disable_cache_on_transformer(self.transformer)
self._cache_dit_enabled = False
self._cached_num_steps = None
self._unmount_cache_dit()
self._minimax_h3_cache_mode = None
self._set_cache_dit_input_preservation(False)
@@ -484,6 +490,7 @@ class MiniMaxH3DenoisingStage(DenoisingStage):
# mounted.
self._cache_dit_enabled = False
self._cached_num_steps = None
self._cache_dit_active_key = None
self._minimax_h3_cache_mode = None
self._set_cache_dit_input_preservation(False)
python/sglang/multimodal_gen/runtime/pipelines_core/stages/progressive_resolution/denoising.py
+10 -6
View File
@@ -219,11 +219,6 @@ class ProgressiveDenoisingStageRouter(PipelineStage):
raise ValueError(f"Unsupported progressive_mode: {mode!r}")
def _get_scm_preset() -> str | None:
preset = envs.SGLANG_CACHE_DIT_SCM_PRESET
return None if (preset is None or preset == "none") else preset
class ProgressiveDenoisingStage(DenoisingStage):
"""DenoisingStage extended with progressive resolution growing.
@@ -300,6 +295,13 @@ class ProgressiveDenoisingStage(DenoisingStage):
"""Called after each stage transition. Update resolution-dependent state."""
pass
def _effective_scm_preset(self) -> str | None:
"""SCM preset for this request: per-request override, then env default."""
preset = self._cache_dit_request_overrides.get(
"scm_preset", envs.SGLANG_CACHE_DIT_SCM_PRESET
)
return None if (preset is None or preset == "none") else preset
def _refresh_cache_dit_context(
self, n_remaining: int, scm_preset: str | None
) -> None:
@@ -612,7 +614,9 @@ class ProgressiveDenoisingStage(DenoisingStage):
# residual-diff decision for the first full-res steps.
if self._cache_dit_enabled:
n_remaining = n_steps - stage_end
self._refresh_cache_dit_context(n_remaining, _get_scm_preset())
self._refresh_cache_dit_context(
n_remaining, self._effective_scm_preset()
)
logger.info(
"cache-dit context refreshed at stage transition "
"(step %d, %d steps remaining)",
python/sglang/multimodal_gen/test/unit/test_attention_backend_override.py
+228
View File
@@ -0,0 +1,228 @@
# SPDX-License-Identifier: Apache-2.0
"""Per-request attention backend override: fail-fast validation and the
two-phase layer switching (CPU-only, layer/selector boundaries patched)."""
import unittest
from types import SimpleNamespace
from unittest.mock import patch
from sglang.multimodal_gen.runtime.layers.attention import layer as layer_module
from sglang.multimodal_gen.runtime.pipelines_core.stages import (
denoising as denoising_module,
)
from sglang.multimodal_gen.runtime.pipelines_core.stages.denoising import (
DenoisingStage,
)
from sglang.multimodal_gen.runtime.platforms import AttentionBackendEnum
def _batch(override=None) -> SimpleNamespace:
return SimpleNamespace(
sampling_params=SimpleNamespace(attention_backend_override=override)
)
def _fake_backend_cls(enum, *, ring_capable=True):
return SimpleNamespace(
get_enum=lambda: enum,
supports_ring_rotation=lambda: ring_capable,
get_impl_cls=lambda: (lambda **kwargs: f"{enum.name.lower()}_impl"),
)
def _fake_layer(default=AttentionBackendEnum.FA) -> SimpleNamespace:
return SimpleNamespace(
backend=default,
_default_attn_backend=default,
_attn_impl_by_backend={default: f"{default.name.lower()}_impl"},
_supported_attention_backends=None,
_attn_impl_ctor_kwargs={"num_heads": 2},
attn_impl=f"{default.name.lower()}_impl",
head_size=64,
dtype="bf16",
)
class TestMaybeOverrideAttentionBackend(unittest.TestCase):
def setUp(self):
self.default_backend_cls = _fake_backend_cls(AttentionBackendEnum.FA)
self.stage = DenoisingStage.__new__(DenoisingStage)
self.stage.server_args = SimpleNamespace(
enable_breakable_cuda_graph=False,
enable_torch_compile=False,
ring_degree=None,
)
self.stage.transformer = object()
self.stage.transformer_2 = None
self.stage.attn_backend = self.default_backend_cls
self.stage._attn_backend_default = self.default_backend_cls
self.stage._attn_metadata_head_size = 64
self.stage._attention_backend_active_override = None
self.layers = [_fake_layer(), _fake_layer()]
self.prepare_calls = []
self.apply_calls = []
self.resolved_backend_cls = _fake_backend_cls(AttentionBackendEnum.SAGE_ATTN)
patchers = [
patch.object(
DenoisingStage,
"_request_switchable_attention_layers",
lambda stage: self.layers,
),
patch.object(
denoising_module,
"prepare_attention_backend_override",
lambda layer, target: self.prepare_calls.append((layer, target)),
),
patch.object(
denoising_module,
"apply_attention_backend_override",
lambda layer, target: self.apply_calls.append((layer, target)),
),
patch.object(
denoising_module,
"get_attn_backend",
lambda **kwargs: self.resolved_backend_cls,
),
]
for patcher in patchers:
patcher.start()
self.addCleanup(patcher.stop)
def test_default_is_a_no_op(self):
self.stage._maybe_override_attention_backend(_batch(None))
self.assertEqual(self.prepare_calls, [])
self.assertEqual(self.apply_calls, [])
self.assertIs(self.stage.attn_backend, self.default_backend_cls)
def test_override_switches_all_layers_then_fast_paths(self):
self.stage._maybe_override_attention_backend(_batch("sage_attn"))
self.assertEqual(len(self.prepare_calls), 2)
self.assertEqual(len(self.apply_calls), 2)
self.assertTrue(
all(t is AttentionBackendEnum.SAGE_ATTN for _, t in self.apply_calls)
)
self.assertIs(self.stage.attn_backend, self.resolved_backend_cls)
self.assertIs(
self.stage._attention_backend_active_override,
AttentionBackendEnum.SAGE_ATTN,
)
self.stage._maybe_override_attention_backend(_batch("sage_attn"))
self.assertEqual(len(self.apply_calls), 2) # unchanged: fast path
def test_default_batch_restores_server_backend(self):
self.stage._maybe_override_attention_backend(_batch("sage_attn"))
self.stage._maybe_override_attention_backend(_batch(None))
self.assertEqual(
self.apply_calls[-2:], [(self.layers[0], None), (self.layers[1], None)]
)
self.assertIs(self.stage.attn_backend, self.default_backend_cls)
self.assertIsNone(self.stage._attention_backend_active_override)
def test_unknown_backend_name_rejected(self):
with self.assertRaisesRegex(ValueError, "Unknown attention_backend_override"):
self.stage._maybe_override_attention_backend(_batch("bogus_attn"))
def test_non_switchable_backend_rejected(self):
with self.assertRaisesRegex(ValueError, "not switchable per"):
self.stage._maybe_override_attention_backend(_batch("video_sparse_attn"))
def test_rejected_under_breakable_cuda_graph(self):
self.stage.server_args.enable_breakable_cuda_graph = True
with self.assertRaisesRegex(ValueError, "breakable CUDA graphs"):
self.stage._maybe_override_attention_backend(_batch("sage_attn"))
self.assertEqual(self.apply_calls, [])
def test_rejected_under_torch_compile(self):
self.stage.server_args.enable_torch_compile = True
with self.assertRaisesRegex(ValueError, "torch.compile"):
self.stage._maybe_override_attention_backend(_batch("sage_attn"))
def test_rejected_when_server_backend_is_sparse(self):
self.layers[1] = _fake_layer(default=AttentionBackendEnum.VIDEO_SPARSE_ATTN)
with self.assertRaisesRegex(ValueError, "sparse backend"):
self.stage._maybe_override_attention_backend(_batch("sage_attn"))
def test_rejected_when_no_switchable_layers(self):
self.layers.clear()
with self.assertRaisesRegex(ValueError, "no switchable attention layers"):
self.stage._maybe_override_attention_backend(_batch("sage_attn"))
def test_rejected_when_target_not_ring_capable_under_ring(self):
self.stage.server_args.ring_degree = 2
self.resolved_backend_cls = _fake_backend_cls(
AttentionBackendEnum.SAGE_ATTN_3, ring_capable=False
)
with self.assertRaisesRegex(ValueError, "ring-capable"):
self.stage._maybe_override_attention_backend(_batch("sage_attn_3"))
def test_prepare_failure_leaves_layers_unswitched(self):
def failing_prepare(layer, target):
if layer is self.layers[1]:
raise ValueError("unsupported on this layer")
self.prepare_calls.append((layer, target))
with patch.object(
denoising_module, "prepare_attention_backend_override", failing_prepare
):
with self.assertRaisesRegex(ValueError, "unsupported on this layer"):
self.stage._maybe_override_attention_backend(_batch("sage_attn"))
self.assertEqual(self.apply_calls, [])
self.assertIsNone(self.stage._attention_backend_active_override)
class TestLayerPrepareApply(unittest.TestCase):
def _patch_selector(self, backend_cls):
return patch.object(
layer_module, "get_attn_backend", lambda *args, **kwargs: backend_cls
)
def test_prepare_builds_and_caches_impl(self):
layer = _fake_layer()
backend_cls = _fake_backend_cls(AttentionBackendEnum.SAGE_ATTN)
with (
self._patch_selector(backend_cls),
patch.object(
layer_module, "wrap_attention_impl_forward", lambda impl: impl
),
):
layer_module.prepare_attention_backend_override(
layer, AttentionBackendEnum.SAGE_ATTN
)
self.assertEqual(
layer._attn_impl_by_backend[AttentionBackendEnum.SAGE_ATTN],
"sage_attn_impl",
)
# prepare must not mutate the active impl
self.assertEqual(layer.attn_impl, "fa_impl")
self.assertIs(layer.backend, AttentionBackendEnum.FA)
def test_prepare_rejects_silent_fallback(self):
layer = _fake_layer()
fallback_cls = _fake_backend_cls(AttentionBackendEnum.TORCH_SDPA)
with self._patch_selector(fallback_cls):
with self.assertRaisesRegex(ValueError, "refusing the request"):
layer_module.prepare_attention_backend_override(
layer, AttentionBackendEnum.SAGE_ATTN
)
self.assertNotIn(AttentionBackendEnum.SAGE_ATTN, layer._attn_impl_by_backend)
def test_apply_flips_and_restores(self):
layer = _fake_layer()
layer._attn_impl_by_backend[AttentionBackendEnum.SAGE_ATTN] = "sage_attn_impl"
layer_module.apply_attention_backend_override(
layer, AttentionBackendEnum.SAGE_ATTN
)
self.assertEqual(layer.attn_impl, "sage_attn_impl")
self.assertIs(layer.backend, AttentionBackendEnum.SAGE_ATTN)
layer_module.apply_attention_backend_override(layer, None)
self.assertEqual(layer.attn_impl, "fa_impl")
self.assertIs(layer.backend, AttentionBackendEnum.FA)
if __name__ == "__main__":
unittest.main()
python/sglang/multimodal_gen/test/unit/test_cache_dit_per_request.py
+236
View File
@@ -0,0 +1,236 @@
# SPDX-License-Identifier: Apache-2.0
"""Per-request Cache-DiT: request-param validation and the mount/refresh/
unmount transitions in DenoisingStage (CPU-only, mount boundary patched)."""
import unittest
from types import SimpleNamespace
from unittest.mock import patch
from sglang.multimodal_gen.runtime.cache.cache_dit_integration import (
cache_dit_overrides_key,
resolve_cache_dit_request_overrides,
)
from sglang.multimodal_gen.runtime.pipelines_core.stages import (
denoising as denoising_module,
)
from sglang.multimodal_gen.runtime.pipelines_core.stages.denoising import (
DenoisingStage,
)
class TestResolveCacheDitRequestOverrides(unittest.TestCase):
def test_none_returns_empty_dict(self):
self.assertEqual(resolve_cache_dit_request_overrides(None), {})
def test_valid_overrides_are_copied(self):
raw = {
"residual_diff_threshold": 0.12,
"scm_preset": "fast",
"secondary": {"max_warmup_steps": 2},
}
resolved = resolve_cache_dit_request_overrides(raw)
self.assertEqual(resolved, raw)
self.assertIsNot(resolved, raw)
self.assertIsNot(resolved["secondary"], raw["secondary"])
def test_non_dict_raises(self):
with self.assertRaisesRegex(ValueError, "must be a dict"):
resolve_cache_dit_request_overrides("fast")
def test_unknown_key_raises(self):
with self.assertRaisesRegex(ValueError, "Unknown cache_dit_params keys"):
resolve_cache_dit_request_overrides({"residual_diff_thresh": 0.1})
def test_secondary_unknown_key_raises(self):
with self.assertRaisesRegex(ValueError, "secondary"):
resolve_cache_dit_request_overrides({"secondary": {"scm_preset": "fast"}})
def test_secondary_non_dict_raises(self):
with self.assertRaisesRegex(ValueError, "secondary"):
resolve_cache_dit_request_overrides({"secondary": 3})
def test_overrides_key_detects_changes(self):
key_a = cache_dit_overrides_key({"Fn_compute_blocks": 1, "scm_preset": "fast"})
key_b = cache_dit_overrides_key({"scm_preset": "fast", "Fn_compute_blocks": 1})
key_c = cache_dit_overrides_key({"Fn_compute_blocks": 2, "scm_preset": "fast"})
self.assertEqual(key_a, key_b)
self.assertNotEqual(key_a, key_c)
def test_overrides_key_freezes_nested_values(self):
key = cache_dit_overrides_key(
{"scm_compute_bins": [4, 2], "secondary": {"Bn_compute_blocks": 1}}
)
hash(key) # must be hashable / comparable
def _batch(
*,
enable_cache_dit=None,
cache_dit_params=None,
is_warmup=False,
) -> SimpleNamespace:
return SimpleNamespace(
is_warmup=is_warmup,
do_classifier_free_guidance=False,
sampling_params=SimpleNamespace(
enable_cache_dit=enable_cache_dit,
cache_dit_params=cache_dit_params,
),
)
class TestPerRequestCacheDitTransitions(unittest.TestCase):
def setUp(self):
self.stage = DenoisingStage.__new__(DenoisingStage)
self.stage.server_args = SimpleNamespace(
enable_breakable_cuda_graph=False,
enable_torch_compile=False,
)
self.stage.transformer = object()
self.stage.transformer_2 = None
self.stage._cache_dit_enabled = False
self.stage._cached_num_steps = None
self.stage._cache_dit_request_overrides = {}
self.stage._cache_dit_active_key = None
self.enable_calls = []
self.disable_calls = []
self.refresh_calls = []
def fake_enable(transformer, config, **kwargs):
self.enable_calls.append(config)
return transformer
def fake_disable(transformer):
self.disable_calls.append(transformer)
return transformer
def fake_refresh(transformer, num_inference_steps, scm_preset=None, **kwargs):
self.refresh_calls.append(num_inference_steps)
patchers = [
patch.object(denoising_module, "enable_cache_on_transformer", fake_enable),
patch.object(
denoising_module, "disable_cache_on_transformer", fake_disable
),
patch.object(
denoising_module, "refresh_context_on_transformer", fake_refresh
),
patch.object(denoising_module, "get_world_size", return_value=1),
]
for patcher in patchers:
patcher.start()
self.addCleanup(patcher.stop)
def test_request_enable_mounts_without_env_default(self):
self.stage._maybe_enable_cache_dit(8, _batch(enable_cache_dit=True))
self.assertEqual(len(self.enable_calls), 1)
self.assertTrue(self.stage._cache_dit_enabled)
self.assertEqual(self.stage._cache_dit_active_key, cache_dit_overrides_key({}))
def test_server_default_off_unmounts_after_request_enable(self):
self.stage._maybe_enable_cache_dit(8, _batch(enable_cache_dit=True))
self.stage._maybe_enable_cache_dit(8, _batch(enable_cache_dit=None))
self.assertEqual(len(self.disable_calls), 1)
self.assertFalse(self.stage._cache_dit_enabled)
self.assertIsNone(self.stage._cache_dit_active_key)
def test_explicit_disable_wins_over_env_default(self):
with patch.object(self.stage, "_cache_dit_requested", return_value=True):
self.stage._maybe_enable_cache_dit(8, _batch(enable_cache_dit=False))
self.assertEqual(self.enable_calls, [])
self.stage._maybe_enable_cache_dit(8, _batch(enable_cache_dit=None))
self.assertEqual(len(self.enable_calls), 1)
def test_same_overrides_refresh_without_remount(self):
params = {"residual_diff_threshold": 0.12}
self.stage._maybe_enable_cache_dit(
8, _batch(enable_cache_dit=True, cache_dit_params=dict(params))
)
self.stage._maybe_enable_cache_dit(
12, _batch(enable_cache_dit=True, cache_dit_params=dict(params))
)
self.assertEqual(len(self.enable_calls), 1)
self.assertEqual(self.disable_calls, [])
self.assertEqual(self.refresh_calls, [12])
def test_changed_overrides_unmount_and_remount(self):
self.stage._maybe_enable_cache_dit(
8,
_batch(
enable_cache_dit=True,
cache_dit_params={"residual_diff_threshold": 0.3},
),
)
self.stage._maybe_enable_cache_dit(
8,
_batch(
enable_cache_dit=True,
cache_dit_params={"residual_diff_threshold": 0.1},
),
)
self.assertEqual(len(self.disable_calls), 1)
self.assertEqual(len(self.enable_calls), 2)
self.assertEqual(self.refresh_calls, [])
self.assertEqual(self.enable_calls[0].residual_diff_threshold, 0.3)
self.assertEqual(self.enable_calls[1].residual_diff_threshold, 0.1)
def test_request_knobs_reach_cache_dit_config(self):
self.stage._maybe_enable_cache_dit(
8,
_batch(
enable_cache_dit=True,
cache_dit_params={
"Fn_compute_blocks": 4,
"max_warmup_steps": 2,
"scm_policy": "static",
},
),
)
(config,) = self.enable_calls
self.assertEqual(config.Fn_compute_blocks, 4)
self.assertEqual(config.max_warmup_steps, 2)
self.assertEqual(config.steps_computation_policy, "static")
self.assertEqual(config.num_inference_steps, 8)
def test_invalid_request_params_raise(self):
with self.assertRaisesRegex(ValueError, "Unknown cache_dit_params keys"):
self.stage._maybe_enable_cache_dit(
8, _batch(enable_cache_dit=True, cache_dit_params={"bogus": 1})
)
self.assertEqual(self.enable_calls, [])
def test_ordinary_warmup_does_not_mount(self):
self.stage._maybe_enable_cache_dit(
8, _batch(enable_cache_dit=True, is_warmup=True)
)
self.assertEqual(self.enable_calls, [])
self.assertFalse(self.stage._cache_dit_enabled)
def test_secondary_inherits_request_primary_overrides(self):
self.stage._cache_dit_request_overrides = resolve_cache_dit_request_overrides(
{
"Fn_compute_blocks": 5,
"secondary": {"Bn_compute_blocks": 7},
}
)
primary = self.stage._build_cache_dit_config(
10, steps_computation_mask=None, scm_policy="dynamic"
)
secondary = self.stage._build_cache_dit_config(
10, steps_computation_mask=None, scm_policy="dynamic", secondary=True
)
self.assertEqual(primary.Fn_compute_blocks, 5)
self.assertEqual(secondary.Fn_compute_blocks, 5) # inherited from primary
self.assertEqual(secondary.Bn_compute_blocks, 7)
def test_request_scm_bins_require_both(self):
self.stage._cache_dit_request_overrides = resolve_cache_dit_request_overrides(
{"scm_compute_bins": [4, 2]}
)
with self.assertRaisesRegex(ValueError, "scm_compute_bins and scm_cache_bins"):
self.stage._parse_cache_dit_scm_bins()
if __name__ == "__main__":
unittest.main()
python/sglang/multimodal_gen/test/unit/test_cfg_gating.py
+38 -1
View File
@@ -29,12 +29,13 @@ class TestCFGGating(unittest.TestCase):
pipeline_config=_PipelineConfig(),
)
def _make_batch(self):
def _make_batch(self, cfg_gate_step=None):
return SimpleNamespace(
cfg_normalization=0,
guidance_rescale=0,
do_classifier_free_guidance=True,
is_cfg_negative=False,
sampling_params=SimpleNamespace(cfg_gate_step=cfg_gate_step),
)
def _make_gate_state(self, gate_step=1, model_id=None, delta=None):
@@ -175,6 +176,42 @@ class TestCFGGating(unittest.TestCase):
with self.assertRaises(ValueError):
stage._init_cfg_gate_state(ctx, batch, server_args)
def test_request_fraction_overrides_env_default(self):
stage = DenoisingStage.__new__(DenoisingStage)
ctx = SimpleNamespace(timesteps=torch.arange(10), extra={}, is_warmup=True)
batch = self._make_batch(cfg_gate_step=0.5)
server_args = self._make_server_args()
# env default is 1.0 (off); the request opts in.
stage._init_cfg_gate_state(ctx, batch, server_args)
state = ctx.extra["cfg_gate_state"]
self.assertTrue(state["requested"])
self.assertTrue(state["active"])
self.assertEqual(state["gate_step"], 5)
def test_request_fraction_disables_env_default(self):
stage = DenoisingStage.__new__(DenoisingStage)
ctx = SimpleNamespace(timesteps=torch.arange(10), extra={}, is_warmup=True)
batch = self._make_batch(cfg_gate_step=1.0)
server_args = self._make_server_args()
with patch.dict(os.environ, {"SGLANG_DIFFUSION_CFG_GATE_STEP": "0.5"}):
stage._init_cfg_gate_state(ctx, batch, server_args)
state = ctx.extra["cfg_gate_state"]
self.assertFalse(state["requested"])
self.assertFalse(state["active"])
def test_rejects_invalid_request_fraction(self):
stage = DenoisingStage.__new__(DenoisingStage)
ctx = SimpleNamespace(timesteps=torch.arange(10), extra={}, is_warmup=True)
batch = self._make_batch(cfg_gate_step=1.5)
server_args = self._make_server_args()
with self.assertRaises(ValueError):
stage._init_cfg_gate_state(ctx, batch, server_args)
if __name__ == "__main__":
unittest.main()
python/sglang/multimodal_gen/test/unit/test_diffusion_bcg_padding.py
+13 -4
View File
@@ -50,6 +50,13 @@ class ZImageTransformer2DModel(torch.nn.Module):
return torch.zeros(pos_ids.shape[0], 8, device=pos_ids.device)
def _fake_cache_dit_batch(*, is_warmup: bool) -> SimpleNamespace:
return SimpleNamespace(
is_warmup=is_warmup,
sampling_params=SimpleNamespace(enable_cache_dit=None, cache_dit_params=None),
)
class TestDiffusionBCGPadding(unittest.TestCase):
def setUp(self):
self.stage = DenoisingStage.__new__(DenoisingStage)
@@ -443,7 +450,7 @@ class TestDiffusionBCGPadding(unittest.TestCase):
self.assertIsNone(self.stage._maybe_get_bcg_runner(self.qwen_model))
self.stage._maybe_torch_compile(self.qwen_model)
self.stage._maybe_enable_cache_dit(1, SimpleNamespace(is_warmup=True))
self.stage._maybe_enable_cache_dit(1, _fake_cache_dit_batch(is_warmup=True))
self.assertEqual(self.stage._bcg_runners, {})
def test_bcg_warns_when_cache_dit_is_requested(self):
@@ -455,8 +462,10 @@ class TestDiffusionBCGPadding(unittest.TestCase):
patch.object(self.stage, "_cache_dit_requested", return_value=True),
patch.object(denoising_module.logger, "warning") as warning,
):
self.stage._maybe_enable_cache_dit(1, SimpleNamespace(is_warmup=True))
self.stage._maybe_enable_cache_dit(1, SimpleNamespace(is_warmup=False))
self.stage._maybe_enable_cache_dit(1, _fake_cache_dit_batch(is_warmup=True))
self.stage._maybe_enable_cache_dit(
1, _fake_cache_dit_batch(is_warmup=False)
)
warning.assert_called_once_with(
"Cache-DiT was requested but is disabled because breakable CUDA "
@@ -474,7 +483,7 @@ class TestDiffusionBCGPadding(unittest.TestCase):
"logger.warning_once"
) as warning_once,
):
self.stage._maybe_enable_cache_dit(1, SimpleNamespace(is_warmup=True))
self.stage._maybe_enable_cache_dit(1, _fake_cache_dit_batch(is_warmup=True))
warning_once.assert_not_called()
python/sglang/multimodal_gen/test/unit/test_minimax_h3_admission.py
+2
View File
@@ -267,6 +267,8 @@ def test_high_quality_request_warns_when_bcg_suppresses_cache_dit():
sampling_params=SimpleNamespace(
quality="high",
_explicit_fields={"quality"},
enable_cache_dit=None,
cache_dit_params=None,
)
)