[diffusion] Reject unsafe quality=high BCG replay (#36008)
This commit is contained in:
@@ -97,7 +97,7 @@ Use `sglang generate --help` and `sglang serve --help` for the full argument lis
|
|||||||
- `--warmup-mode {off|request|server}`: control startup warmup for `sglang serve`; `off` skips warmup, `request` primes the request path, and `server` runs a full synthetic server warmup before serving traffic
|
- `--warmup-mode {off|request|server}`: control startup warmup for `sglang serve`; `off` skips warmup, `request` primes the request path, and `server` runs a full synthetic server warmup before serving traffic
|
||||||
- `--enable-torch-compile {true|false}`: compile native diffusion hot paths. When no warmup mode is configured, this also enables server warmup so first real requests do not pay compile latency.
|
- `--enable-torch-compile {true|false}`: compile native diffusion hot paths. When no warmup mode is configured, this also enables server warmup so first real requests do not pay compile latency.
|
||||||
- `--offload-during-compile {true|false}`: when compile warmup is active, temporarily layerwise-offload DiT weights and move resident non-DiT components off-device so `max-autotune` fits on tighter-memory GPUs; the configured serving residency is restored before real traffic. Skipped under existing layerwise offload, Cache-DiT, or FSDP.
|
- `--offload-during-compile {true|false}`: when compile warmup is active, temporarily layerwise-offload DiT weights and move resident non-DiT components off-device so `max-autotune` fits on tighter-memory GPUs; the configured serving residency is restored before real traffic. Skipped under existing layerwise offload, Cache-DiT, or FSDP.
|
||||||
- `--enable-breakable-cuda-graph {true|false}`: capture supported DiT forwards as breakable CUDA graph segments to reduce launch overhead. Requires `--warmup-resolutions` for every served resolution because each resolution is captured separately.
|
- `--enable-breakable-cuda-graph {true|false}`: capture supported DiT forwards as breakable CUDA graph segments to reduce launch overhead. Requires `--warmup-resolutions` for every served resolution because each resolution is captured separately. A `quality=high` request is rejected when it would mount request-scoped DiT fusions that were not present during lossless graph capture; VAE-only high-quality paths remain compatible.
|
||||||
- `--bcg-text-buckets {N...}`: prompt-length padding buckets for breakable CUDA graph capture/replay reuse.
|
- `--bcg-text-buckets {N...}`: prompt-length padding buckets for breakable CUDA graph capture/replay reuse.
|
||||||
- `--attention-backend {BACKEND}`: attention backend for native SGLang and diffusers pipelines
|
- `--attention-backend {BACKEND}`: attention backend for native SGLang and diffusers pipelines
|
||||||
- `--component-attention-backends {MAP}`: per-component attention backend overrides, for example `text_encoder=torch_sdpa,transformer=fa`
|
- `--component-attention-backends {MAP}`: per-component attention backend overrides, for example `text_encoder=torch_sdpa,transformer=fa`
|
||||||
|
|||||||
@@ -48,6 +48,15 @@ The `quality` field in a **video response** body is unrelated. It is Sora-compat
|
|||||||
|
|
||||||
`quality` participates in the dynamic-batch signature, so mixed-quality traffic is batched separately and the transition happens safely at a batch boundary. Mounting is all-or-nothing: if any marked site on a transformer fails its static guards, no site on that transformer is fused.
|
`quality` participates in the dynamic-batch signature, so mixed-quality traffic is batched separately and the transition happens safely at a batch boundary. Mounting is all-or-nothing: if any marked site on a transformer fails its static guards, no site on that transformer is fused.
|
||||||
|
|
||||||
|
<Warning>
|
||||||
|
Do not combine request-gated DiT fusions with `--enable-breakable-cuda-graph`.
|
||||||
|
BCG warmup captures the lossless module branches before a high-quality request
|
||||||
|
mounts its DiT fusions, so replay would bypass the requested kernels. SGLang
|
||||||
|
rejects this combination for models with eligible DiT quality sites. Models
|
||||||
|
whose high-quality path changes only VAE decode remain allowed because BCG
|
||||||
|
captures the DiT only.
|
||||||
|
</Warning>
|
||||||
|
|
||||||
These fusion families mount under `quality="high"`:
|
These fusion families mount under `quality="high"`:
|
||||||
|
|
||||||
| Fusion | What it folds |
|
| Fusion | What it folds |
|
||||||
|
|||||||
@@ -659,6 +659,19 @@ class DenoisingStage(PipelineStage, RolloutDenoisingMixin):
|
|||||||
mounted_fusions.add(description)
|
mounted_fusions.add(description)
|
||||||
else:
|
else:
|
||||||
unmount(transformer)
|
unmount(transformer)
|
||||||
|
|
||||||
|
if want and mounted_fusions and self.server_args.enable_breakable_cuda_graph:
|
||||||
|
for transformer in filter(None, [self.transformer, self.transformer_2]):
|
||||||
|
for _, _, unmount in _QUALITY_FUSION_HANDLERS:
|
||||||
|
unmount(transformer)
|
||||||
|
descriptions = ", ".join(sorted(mounted_fusions))
|
||||||
|
raise ValueError(
|
||||||
|
"quality='high' cannot be used with breakable CUDA graphs for "
|
||||||
|
f"this model because its request-scoped DiT fusions "
|
||||||
|
f"({descriptions}) do not match the lossless warmup graphs. "
|
||||||
|
"Disable breakable CUDA graphs or use quality='lossless'."
|
||||||
|
)
|
||||||
|
|
||||||
self._quality_fusions_mounted = want
|
self._quality_fusions_mounted = want
|
||||||
for description in sorted(mounted_fusions):
|
for description in sorted(mounted_fusions):
|
||||||
logger.info("Mounted %s for quality=high", description)
|
logger.info("Mounted %s for quality=high", description)
|
||||||
|
|||||||
@@ -58,6 +58,44 @@ class SanaVideoTransformer3DModel(torch.nn.Module):
|
|||||||
pass
|
pass
|
||||||
|
|
||||||
|
|
||||||
|
class TestQualityFusionBCGCompatibility(unittest.TestCase):
|
||||||
|
def setUp(self):
|
||||||
|
self.stage = DenoisingStage.__new__(DenoisingStage)
|
||||||
|
self.stage.server_args = SimpleNamespace(enable_breakable_cuda_graph=True)
|
||||||
|
self.stage.transformer = OtherTransformer2DModel()
|
||||||
|
self.stage.transformer_2 = None
|
||||||
|
self.stage._quality_fusions_mounted = False
|
||||||
|
|
||||||
|
@staticmethod
|
||||||
|
def _batch(quality: str):
|
||||||
|
return SimpleNamespace(sampling_params=SimpleNamespace(quality=quality))
|
||||||
|
|
||||||
|
def test_rejects_high_when_dit_fusion_would_replace_captured_graph(self):
|
||||||
|
unmounted = []
|
||||||
|
handlers = (
|
||||||
|
(
|
||||||
|
"test fusion",
|
||||||
|
lambda _: True,
|
||||||
|
lambda transformer: unmounted.append(transformer),
|
||||||
|
),
|
||||||
|
)
|
||||||
|
|
||||||
|
with patch.object(denoising_module, "_QUALITY_FUSION_HANDLERS", handlers):
|
||||||
|
with self.assertRaisesRegex(ValueError, "lossless warmup graphs"):
|
||||||
|
self.stage._maybe_toggle_quality_fusions(self._batch("high"))
|
||||||
|
|
||||||
|
self.assertEqual(unmounted, [self.stage.transformer])
|
||||||
|
self.assertFalse(self.stage._quality_fusions_mounted)
|
||||||
|
|
||||||
|
def test_allows_high_when_model_has_no_dit_quality_fusions(self):
|
||||||
|
handlers = (("test fusion", lambda _: False, lambda _: None),)
|
||||||
|
|
||||||
|
with patch.object(denoising_module, "_QUALITY_FUSION_HANDLERS", handlers):
|
||||||
|
self.stage._maybe_toggle_quality_fusions(self._batch("high"))
|
||||||
|
|
||||||
|
self.assertTrue(self.stage._quality_fusions_mounted)
|
||||||
|
|
||||||
|
|
||||||
def _fake_cache_dit_batch(*, is_warmup: bool) -> SimpleNamespace:
|
def _fake_cache_dit_batch(*, is_warmup: bool) -> SimpleNamespace:
|
||||||
return SimpleNamespace(
|
return SimpleNamespace(
|
||||||
is_warmup=is_warmup,
|
is_warmup=is_warmup,
|
||||||
|
|||||||
Reference in New Issue
Block a user