diff --git a/docs/docs/sglang-diffusion/attention_backends.mdx b/docs/docs/sglang-diffusion/attention_backends.mdx index e0be8b6f5..1433feb30 100644 --- a/docs/docs/sglang-diffusion/attention_backends.mdx +++ b/docs/docs/sglang-diffusion/attention_backends.mdx @@ -632,6 +632,37 @@ sglang generate \ Component keys match pipeline module names from `model_index.json`, such as `text_encoder`, `text_encoder_2`, `transformer`, `transformer_2`, or `connectors`. +### Per-request override (denoise loop) + +A single server can serve exact and approximate attention side by side: requests +may switch the DiT denoise attention backend via the `attention_backend_override` +sampling param. Valid values are the exact/drop-in dense kernels — `fa`, +`torch_sdpa`, `sage_attn`, `sage_attn_3`. The field participates in the +dynamic-batch signature, so requests with different backends never share a batch. + +```bash +sglang generate \ + --model-path \ + --prompt "..." \ + --attention-backend-override sage_attn +``` + +```python +client.images.generate( + model="", + prompt="...", + extra_body={"attention_backend_override": "sage_attn"}, +) +``` + +Incompatible server settings **reject the request** (with a server log) instead +of silently falling back: breakable CUDA graphs and `torch.compile` bake the +attention kernel into a captured/traced graph; sparse server backends +(`sliding_tile_attn`, `video_sparse_attn`, ...) cannot be mixed with per-request +dense switching; under ring parallelism the target must be ring-capable. Note +`sage_attn` / `sage_attn_3` are lossy (quantized attention) — validate quality +on your workload. + ### Sage then Sol hybrid `sol_attn` keeps the first `dense_steps` steps dense. Set diff --git a/docs/docs/sglang-diffusion/cache_dit.mdx b/docs/docs/sglang-diffusion/cache_dit.mdx index c8880ae55..14d62e0ca 100644 --- a/docs/docs/sglang-diffusion/cache_dit.mdx +++ b/docs/docs/sglang-diffusion/cache_dit.mdx @@ -15,12 +15,50 @@ SGLang integrates [Cache-DiT](https://github.com/vipshop/cache-dit), a caching a ## Basic Usage -Enable Cache-DiT by exporting the environment variable and using `sglang generate` or `sglang serve` : +Cache-DiT is a **per-request** switch: each request decides whether to run +cached or lossless, and requests with different Cache-DiT settings never share +a batch. The `SGLANG_CACHE_DIT_*` environment variables remain available as +server-wide defaults for requests that leave the switch unset. + +Enable it for a single generation: + +```bash +sglang generate --model-path Qwen/Qwen-Image \ + --prompt "A beautiful sunset over the mountains" \ + --enable-cache-dit true +``` + +Or per request against a running server, via the OpenAI-compatible API: + +```python +client.images.generate( + model="Qwen/Qwen-Image", + prompt="A beautiful sunset over the mountains", + extra_body={ + "enable_cache_dit": True, + # optional knob overrides for this request only + "cache_dit_params": {"residual_diff_threshold": 0.12, "scm_preset": "fast"}, + }, +) +``` + +`enable_cache_dit` accepts three states: `true` (on for this request), `false` +(off for this request, overriding the server default), and unset (follow the +`SGLANG_CACHE_DIT_ENABLED` server default). `cache_dit_params` accepts the +DBCache knobs (`Fn_compute_blocks`, `Bn_compute_blocks`, `max_warmup_steps`, +`residual_diff_threshold`, `max_continuous_cached_steps`, `enable_taylorseer`, +`taylorseer_order`), the SCM knobs (`scm_preset`, `scm_compute_bins`, +`scm_cache_bins`, `scm_policy`), and a nested `secondary` dict with the DBCache +knobs for the second transformer of dual-DiT models (unset secondary keys +inherit the request's primary values, then the +`SGLANG_CACHE_DIT_SECONDARY_*` defaults). + +To make Cache-DiT the default for every request instead, export the +environment variable when launching: ```bash SGLANG_CACHE_DIT_ENABLED=true \ -sglang generate --model-path Qwen/Qwen-Image \ - --prompt "A beautiful sunset over the mountains" +sglang serve --model-path Qwen/Qwen-Image ``` ## Diffusers Backend @@ -508,7 +546,9 @@ sglang generate --model-path Qwen/Qwen-Image \ ## Environment Variables -All Cache-DiT parameters can be configured via environment variables. +All Cache-DiT parameters can also be configured via environment variables, +which act as the server-wide defaults for requests that don't set +`enable_cache_dit` / `cache_dit_params`. See [Environment Variables](./environment_variables) for the complete list. ## Supported Models diff --git a/python/sglang/multimodal_gen/.claude/skills/sglang-diffusion-performance/SKILL.md b/python/sglang/multimodal_gen/.claude/skills/sglang-diffusion-performance/SKILL.md index 4e4862bec..db4703345 100644 --- a/python/sglang/multimodal_gen/.claude/skills/sglang-diffusion-performance/SKILL.md +++ b/python/sglang/multimodal_gen/.claude/skills/sglang-diffusion-performance/SKILL.md @@ -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 ` | 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 ` (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 ` | 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. | diff --git a/python/sglang/multimodal_gen/configs/sample/sampling_params.py b/python/sglang/multimodal_gen/configs/sample/sampling_params.py index 32f713450..f28227fec 100644 --- a/python/sglang/multimodal_gen/configs/sample/sampling_params.py +++ b/python/sglang/multimodal_gen/configs/sample/sampling_params.py @@ -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", diff --git a/python/sglang/multimodal_gen/runtime/cache/cache_dit_integration.py b/python/sglang/multimodal_gen/runtime/cache/cache_dit_integration.py index 2616c9a37..6b6bf2db4 100644 --- a/python/sglang/multimodal_gen/runtime/cache/cache_dit_integration.py +++ b/python/sglang/multimodal_gen/runtime/cache/cache_dit_integration.py @@ -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. diff --git a/python/sglang/multimodal_gen/runtime/entrypoints/openai/image_api.py b/python/sglang/multimodal_gen/runtime/entrypoints/openai/image_api.py index 73bb468f3..ee0881813 100644 --- a/python/sglang/multimodal_gen/runtime/entrypoints/openai/image_api.py +++ b/python/sglang/multimodal_gen/runtime/entrypoints/openai/image_api.py @@ -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, diff --git a/python/sglang/multimodal_gen/runtime/entrypoints/openai/video_api.py b/python/sglang/multimodal_gen/runtime/entrypoints/openai/video_api.py index ffe213bb5..319632a5a 100644 --- a/python/sglang/multimodal_gen/runtime/entrypoints/openai/video_api.py +++ b/python/sglang/multimodal_gen/runtime/entrypoints/openai/video_api.py @@ -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, diff --git a/python/sglang/multimodal_gen/runtime/layers/attention/layer.py b/python/sglang/multimodal_gen/runtime/layers/attention/layer.py index 0aaadfc4c..74f03edf0 100644 --- a/python/sglang/multimodal_gen/runtime/layers/attention/layer.py +++ b/python/sglang/multimodal_gen/runtime/layers/attention/layer.py @@ -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 diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising.py index 5bd497015..abbd6f6a4 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising.py @@ -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) diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/ltx_2/denoising.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/ltx_2/denoising.py index 8b3ebf601..e6c60bd96 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/ltx_2/denoising.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/ltx_2/denoising.py @@ -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 diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/minimax_h3/stages/denoising.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/minimax_h3/stages/denoising.py index 099901840..69cd37352 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/minimax_h3/stages/denoising.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/minimax_h3/stages/denoising.py @@ -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) diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/progressive_resolution/denoising.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/progressive_resolution/denoising.py index 4db6b3a33..32c3285e0 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/progressive_resolution/denoising.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/progressive_resolution/denoising.py @@ -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)", diff --git a/python/sglang/multimodal_gen/test/unit/test_attention_backend_override.py b/python/sglang/multimodal_gen/test/unit/test_attention_backend_override.py new file mode 100644 index 000000000..468a81d7b --- /dev/null +++ b/python/sglang/multimodal_gen/test/unit/test_attention_backend_override.py @@ -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() diff --git a/python/sglang/multimodal_gen/test/unit/test_cache_dit_per_request.py b/python/sglang/multimodal_gen/test/unit/test_cache_dit_per_request.py new file mode 100644 index 000000000..86e8c03b8 --- /dev/null +++ b/python/sglang/multimodal_gen/test/unit/test_cache_dit_per_request.py @@ -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() diff --git a/python/sglang/multimodal_gen/test/unit/test_cfg_gating.py b/python/sglang/multimodal_gen/test/unit/test_cfg_gating.py index 37f66c9c7..9627e17b3 100644 --- a/python/sglang/multimodal_gen/test/unit/test_cfg_gating.py +++ b/python/sglang/multimodal_gen/test/unit/test_cfg_gating.py @@ -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() diff --git a/python/sglang/multimodal_gen/test/unit/test_diffusion_bcg_padding.py b/python/sglang/multimodal_gen/test/unit/test_diffusion_bcg_padding.py index 8ce7abb43..689fbecf9 100644 --- a/python/sglang/multimodal_gen/test/unit/test_diffusion_bcg_padding.py +++ b/python/sglang/multimodal_gen/test/unit/test_diffusion_bcg_padding.py @@ -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() diff --git a/python/sglang/multimodal_gen/test/unit/test_minimax_h3_admission.py b/python/sglang/multimodal_gen/test/unit/test_minimax_h3_admission.py index 170f12bbd..978df8ba6 100644 --- a/python/sglang/multimodal_gen/test/unit/test_minimax_h3_admission.py +++ b/python/sglang/multimodal_gen/test/unit/test_minimax_h3_admission.py @@ -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, ) )