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

Co-authored-by: Claude Fable 5 <noreply@anthropic.com>

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 <MODEL_PATH_OR_ID> \
+  --prompt "..." \
+  --attention-backend-override sage_attn
+```
+
+```python
+client.images.generate(
+    model="<MODEL_PATH_OR_ID>",
+    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

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

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 <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

@@ -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

@@ -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

@@ -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

@@ -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

@@ -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

@@ -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

@@ -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

@@ -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

@@ -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

@@ -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

@@ -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

@@ -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

@@ -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

@@ -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,
         )
     )