From 038c9b948fbf10e99d3167c18a6a22c20049f7f5 Mon Sep 17 00:00:00 2001 From: NumberWan Date: Wed, 9 Sep 2026 18:20:34 +0800 Subject: [PATCH] [Bugfix] Resolve #6931 hub FA3 on torch 2.13 via kernels 0.16.1 (#7185) Signed-off-by: NumberWan --- .../attention_backends/huggingface_hub.md | 2 +- recipes/Boogu/Boogu-Image.md | 2 +- requirements/common.txt | 10 ++- tests/diffusion/attention/test_kernels_hub.py | 74 +++++++++++++++++ .../models/ltx2/test_ltx25_pipeline.py | 4 +- tests/e2e/accuracy/test_qwen_image.py | 83 ++++++++----------- .../attention/backends/flash_attn_hub.py | 6 +- .../diffusion/models/ltx2/ltx2_components.py | 4 +- vllm_omni/platforms/cuda/platform.py | 20 +++++ 9 files changed, 142 insertions(+), 63 deletions(-) diff --git a/docs/user_guide/diffusion/attention_backends/huggingface_hub.md b/docs/user_guide/diffusion/attention_backends/huggingface_hub.md index 574ff625a..cfd220e82 100644 --- a/docs/user_guide/diffusion/attention_backends/huggingface_hub.md +++ b/docs/user_guide/diffusion/attention_backends/huggingface_hub.md @@ -12,7 +12,7 @@ implementations to reduce numerical drift and sampling divergence. ## Installation ```bash -pip install kernels==0.15.2 +pip install kernels==0.16.1 ``` If `kernels` is unavailable, an explicit `FLASH_ATTN_HUB` or diff --git a/recipes/Boogu/Boogu-Image.md b/recipes/Boogu/Boogu-Image.md index 0d832b0ae..0353b9778 100644 --- a/recipes/Boogu/Boogu-Image.md +++ b/recipes/Boogu/Boogu-Image.md @@ -152,7 +152,7 @@ The following configuration was validated: - vLLM-Omni version or commit: Use a commit that contains TorchAO FP8 checkpoint loading for diffusion models - TorchAO version: 0.17.0 -- `kernels` / `kernels-data`: 0.15.2 / 0.16.1 +- `kernels` / `kernels-data`: 0.16.1 / matching kernels-data #### Command diff --git a/requirements/common.txt b/requirements/common.txt index bb79e0a6a..855fb6b0a 100644 --- a/requirements/common.txt +++ b/requirements/common.txt @@ -9,13 +9,15 @@ # That breakage is now handled in code instead of by pinning: # MossTTSNanoForGeneration wraps weight loading in # transformers_keys_to_ignore_compat (vllm_omni/model_executor/models/utils.py). -transformers >= 5.10.1, < 5.15 +# Floor is 5.13.0 so transformers.integrations.hub_kernels LayerRepository +# entries always pass version=/revision= required by kernels>=0.16 (#6971). +transformers >= 5.13.0, < 5.15 av>=14.0.0 omegaconf>=2.3.0 diffusers==0.40.0 -# Hub kernels, including LTX-2.5 NATTEN. Keep this pin in sync with the -# supported transformers range above. -kernels==0.15.2 +# Hub kernels, including LTX-2.5 NATTEN. 0.16+ resolves torch-stable-abi* +# FA3 builds on torch 2.13 / cu130; keep in sync with transformers above. +kernels==0.16.1 safetensors>=0.8.0 accelerate==1.12.0 soundfile>=0.13.1 diff --git a/tests/diffusion/attention/test_kernels_hub.py b/tests/diffusion/attention/test_kernels_hub.py index 3f7451ca3..c85fb58f3 100644 --- a/tests/diffusion/attention/test_kernels_hub.py +++ b/tests/diffusion/attention/test_kernels_hub.py @@ -13,6 +13,12 @@ from vllm_omni.diffusion.attention.backends.sdpa import SDPAImpl from vllm_omni.platforms import current_omni_platform +def _clear_hub_module_cache(monkeypatch: pytest.MonkeyPatch) -> None: + from vllm_omni.diffusion.attention.backends import flash_attn_hub + + monkeypatch.setattr(flash_attn_hub, "_hub_modules", {}) + + @pytest.mark.core_model @pytest.mark.cpu def test_explicit_kernels_hub_selection_does_not_fallback(monkeypatch: pytest.MonkeyPatch): @@ -54,6 +60,74 @@ def test_explicit_kernels_hub_selection_does_not_fallback(monkeypatch: pytest.Mo CudaOmniPlatform.get_diffusion_attn_backend_cls("FLASH_ATTN_3_HUB", head_size=64) +@pytest.mark.core_model +@pytest.mark.cpu +def test_explicit_kernels_hub_variant_miss_does_not_fallback(monkeypatch: pytest.MonkeyPatch): + """Explicit Hub backends must fail loudly when no build variant resolves (#6971).""" + from vllm.platforms.interface import DeviceCapability + + from vllm_omni.diffusion.envs import PACKAGES_CHECKER + from vllm_omni.platforms.cuda.platform import CudaOmniPlatform + + monkeypatch.setattr( + CudaOmniPlatform, + "get_device_capability", + classmethod(lambda cls, device_id=0: DeviceCapability(9, 0)), + ) + monkeypatch.setattr(PACKAGES_CHECKER, "get_packages_info", lambda: {"has_flash_attn": True}) + _clear_hub_module_cache(monkeypatch) + + kernels_module = types.ModuleType("kernels") + + def _missing_variant(repo_id, version=None, **kwargs): + raise FileNotFoundError(f"Cannot find a build variant for this system in {repo_id}") + + kernels_module.get_kernel = _missing_variant # type: ignore[attr-defined] + monkeypatch.setitem(sys.modules, "kernels", kernels_module) + + with pytest.raises(RuntimeError, match="no compatible build"): + CudaOmniPlatform.get_diffusion_attn_backend_cls("FLASH_ATTN_3_HUB", head_size=64) + with pytest.raises(RuntimeError, match="no compatible build"): + CudaOmniPlatform.get_diffusion_attn_backend_cls("FLASH_ATTN_HUB", head_size=64) + + +@pytest.mark.core_model +@pytest.mark.cpu +def test_explicit_kernels_hub_preflight_uses_v2_when_v1_misses(monkeypatch: pytest.MonkeyPatch): + """Preflight must share _load_hub_module's (1, 2) policy, not pin version=1.""" + from vllm.platforms.interface import DeviceCapability + + from vllm_omni.diffusion.attention.backends.registry import DiffusionAttentionBackendEnum + from vllm_omni.diffusion.envs import PACKAGES_CHECKER + from vllm_omni.platforms.cuda.platform import CudaOmniPlatform + + monkeypatch.setattr( + CudaOmniPlatform, + "get_device_capability", + classmethod(lambda cls, device_id=0: DeviceCapability(9, 0)), + ) + monkeypatch.setattr(PACKAGES_CHECKER, "get_packages_info", lambda: {"has_flash_attn": True}) + _clear_hub_module_cache(monkeypatch) + + kernels_module = types.ModuleType("kernels") + versions_tried: list[int | None] = [] + + def _v1_miss_v2_hit(repo_id, version=None, **kwargs): + versions_tried.append(version) + if version == 1: + raise FileNotFoundError(f"Cannot find a build variant for this system in {repo_id}") + return types.SimpleNamespace(name=f"{repo_id}@v{version}") + + kernels_module.get_kernel = _v1_miss_v2_hit # type: ignore[attr-defined] + monkeypatch.setitem(sys.modules, "kernels", kernels_module) + + fa3_path = CudaOmniPlatform.get_diffusion_attn_backend_cls("FLASH_ATTN_3_HUB", head_size=64) + assert fa3_path == DiffusionAttentionBackendEnum.FLASH_ATTN_3_HUB.get_path() + fa2_path = CudaOmniPlatform.get_diffusion_attn_backend_cls("FLASH_ATTN_HUB", head_size=64) + assert fa2_path == DiffusionAttentionBackendEnum.FLASH_ATTN_HUB.get_path() + assert versions_tried == [1, 2, 1, 2] + + @pytest.mark.core_model @pytest.mark.cpu def test_explicit_flash_attention_unavailable_does_not_fallback(monkeypatch: pytest.MonkeyPatch): diff --git a/tests/diffusion/models/ltx2/test_ltx25_pipeline.py b/tests/diffusion/models/ltx2/test_ltx25_pipeline.py index ae90ea1a3..a96fa8a4e 100644 --- a/tests/diffusion/models/ltx2/test_ltx25_pipeline.py +++ b/tests/diffusion/models/ltx2/test_ltx25_pipeline.py @@ -186,7 +186,7 @@ def test_ltx25_missing_gemma4_recommends_supported_transformers_range(monkeypatc ltx2_components.initialize_pipeline_components(pipe, od_config) assert str(exc_info.value) == ( - "LTX-2.5 requires Gemma4UnifiedForConditionalGeneration; install transformers>=5.10.1,<5.15." + "LTX-2.5 requires Gemma4UnifiedForConditionalGeneration; install transformers>=5.13.0,<5.15." ) @@ -387,7 +387,7 @@ def test_ltx25_natten_failure_has_actionable_omni_remedy(monkeypatch): MissingNattenProcessor, ) - with pytest.raises(RuntimeError, match=r"kernels==0\.15\.2.*supported GPU.*allow Hub access"): + with pytest.raises(RuntimeError, match=r"kernels==0\.16\.1.*supported GPU.*allow Hub access"): ltx2_components._create_ltx25_natten_processor() diff --git a/tests/e2e/accuracy/test_qwen_image.py b/tests/e2e/accuracy/test_qwen_image.py index a09014c8d..45fd20934 100644 --- a/tests/e2e/accuracy/test_qwen_image.py +++ b/tests/e2e/accuracy/test_qwen_image.py @@ -34,11 +34,7 @@ TRUE_CFG_SCALE = 4.0 SEED = 42 SSIM_THRESHOLD = 0.94 # With --fa-deterministic, H100 nightlies stably land near SSIM 0.958 / PSNR 27.8 -# vs Diffusers (see #5734 / #5963), so the primary gates are 0.94 / 27.0. -# The SDPA-matched fallback (hub kernel unavailable for the image's torch) -# landed at SSIM 0.9687 in build 2954 — keep a slightly tighter SSIM gate for -# that path only. -SSIM_THRESHOLD_SDPA_FALLBACK = 0.96 +# vs Diffusers on the FA3 hub path (see #5734 / #5963 / #6971). PSNR_THRESHOLD = 27.0 MODEL_2512_ID = "Qwen/Qwen-Image-2512" @@ -65,44 +61,6 @@ MEAN_ABS_DIFF_THRESHOLD_2512 = 3e-2 P99_ABS_DIFF_THRESHOLD_2512 = 4e-1 -_FLASH_ATTN3_HUB_AVAILABLE: bool | None = None - - -def _flash_attn3_hub_available() -> bool: - """Whether kernels-community/flash-attn3 has a build for this torch/CUDA. - - The hub repo only ships prebuilt variants for specific torch versions; - a base-image torch bump can outrun it (build 2952 failed with "Cannot - find a build variant for this system"). When it is unavailable, both - sides of the comparison fall back to torch SDPA together — the test - compares implementations, so the attention math must stay matched. - """ - global _FLASH_ATTN3_HUB_AVAILABLE - if _FLASH_ATTN3_HUB_AVAILABLE is None: - try: - from kernels import get_kernel - - get_kernel("kernels-community/flash-attn3") - _FLASH_ATTN3_HUB_AVAILABLE = True - except Exception as exc: - print(f"kernels-community/flash-attn3 unavailable ({exc}); using matched torch SDPA on both sides.") - _FLASH_ATTN3_HUB_AVAILABLE = False - return _FLASH_ATTN3_HUB_AVAILABLE - - -def _omni_server_env() -> dict[str, str] | None: - if _flash_attn3_hub_available(): - return None - return {"DIFFUSION_ATTENTION_BACKEND": "TORCH_SDPA"} - - -def _set_reference_attention_backend(pipe: DiffusionPipeline) -> None: - if _flash_attn3_hub_available(): - pipe.transformer.set_attention_backend("_flash_3_hub") - else: - pipe.transformer.set_attention_backend("native") - - def _model_name() -> str: return os.environ.get(MODEL_ENV_VAR, MODEL_ID) @@ -115,7 +73,22 @@ def _local_files_only(model: str) -> bool: return Path(model).exists() +_OMNI_FA3_HUB_BACKEND = "FLASH_ATTN_3_HUB" +_OMNI_FA3_HUB_ENV = {"DIFFUSION_ATTENTION_BACKEND": _OMNI_FA3_HUB_BACKEND} + + +def _assert_omni_fa3_hub_resolves() -> None: + from vllm_omni.diffusion.attention.backends.registry import DiffusionAttentionBackendEnum + from vllm_omni.platforms.cuda.platform import CudaOmniPlatform + + path = CudaOmniPlatform.get_diffusion_attn_backend_cls(_OMNI_FA3_HUB_BACKEND, head_size=128) + expected = DiffusionAttentionBackendEnum.FLASH_ATTN_3_HUB.get_path() + assert path == expected, path + print(f"Resolved Omni attention backend: {_OMNI_FA3_HUB_BACKEND} -> {path}") + + def _run_vllm_omni_qwen_image(*, model: str, output_path: Path) -> Image.Image: + _assert_omni_fa3_hub_resolves() server_args = [ "--num-gpus", "1", @@ -124,8 +97,10 @@ def _run_vllm_omni_qwen_image(*, model: str, output_path: Path) -> Image.Image: "--init-timeout", "900", "--fa-deterministic", + "--diffusion-attention-backend", + _OMNI_FA3_HUB_BACKEND, ] - with OmniServer(model, server_args, use_omni=True, env_dict=_omni_server_env()) as omni_server: + with OmniServer(model, server_args, use_omni=True, env_dict=_OMNI_FA3_HUB_ENV) as omni_server: response = requests.post( f"http://{omni_server.host}:{omni_server.port}/v1/images/generations", json={ @@ -161,7 +136,7 @@ def _run_diffusers_qwen_image(*, model: str, output_path: Path) -> Image.Image: trust_remote_code=True, local_files_only=_local_files_only(model), ).to("cuda") - _set_reference_attention_backend(pipe) + pipe.transformer.set_attention_backend("_flash_3_hub") generator = torch.Generator(device="cuda").manual_seed(SEED) result = pipe( # pyright: ignore[reportCallIssue] prompt=PROMPT, @@ -186,8 +161,18 @@ def _run_diffusers_qwen_image(*, model: str, output_path: Path) -> Image.Image: def _run_vllm_omni_qwen_image_2512(*, model: str, output_path: Path) -> Image.Image: - server_args = ["--num-gpus", "1", "--stage-init-timeout", "300", "--init-timeout", "900"] - with OmniServer(model, server_args, use_omni=True, env_dict=_omni_server_env()) as omni_server: + _assert_omni_fa3_hub_resolves() + server_args = [ + "--num-gpus", + "1", + "--stage-init-timeout", + "300", + "--init-timeout", + "900", + "--diffusion-attention-backend", + _OMNI_FA3_HUB_BACKEND, + ] + with OmniServer(model, server_args, use_omni=True, env_dict=_OMNI_FA3_HUB_ENV) as omni_server: response = requests.post( f"http://{omni_server.host}:{omni_server.port}/v1/images/generations", json={ @@ -223,7 +208,7 @@ def _run_diffusers_qwen_image_2512(*, model: str, output_path: Path) -> Image.Im trust_remote_code=True, local_files_only=_local_files_only(model), ).to("cuda") - _set_reference_attention_backend(pipe) + pipe.transformer.set_attention_backend("_flash_3_hub") generator = torch.Generator(device="cuda").manual_seed(SEED_2512) result = pipe( # pyright: ignore[reportCallIssue] prompt=PROMPT_2512, @@ -262,7 +247,7 @@ def test_qwen_image_matches_diffusers(accuracy_artifact_root: Path) -> None: diffusers_image=diffusers_output, width=WIDTH, height=HEIGHT, - ssim_threshold=SSIM_THRESHOLD if _flash_attn3_hub_available() else SSIM_THRESHOLD_SDPA_FALLBACK, + ssim_threshold=SSIM_THRESHOLD, psnr_threshold=PSNR_THRESHOLD, ) diff --git a/vllm_omni/diffusion/attention/backends/flash_attn_hub.py b/vllm_omni/diffusion/attention/backends/flash_attn_hub.py index 7f9c78bf7..b24c99694 100644 --- a/vllm_omni/diffusion/attention/backends/flash_attn_hub.py +++ b/vllm_omni/diffusion/attention/backends/flash_attn_hub.py @@ -28,11 +28,9 @@ def _load_hub_module(repo_id: str): logger.info("Loading %s kernel from HuggingFace Hub...", repo_id) last_error = None - for version in (1, 2, None): + for version in (1, 2): try: - if version is not None: - return get_kernel(repo_id, version=version) - return get_kernel(repo_id) + return get_kernel(repo_id, version=version) except Exception as exc: logger.info("Failed to load %s version %s: %s", repo_id, version, exc) last_error = exc diff --git a/vllm_omni/diffusion/models/ltx2/ltx2_components.py b/vllm_omni/diffusion/models/ltx2/ltx2_components.py index f1dcc38bb..5826f273e 100644 --- a/vllm_omni/diffusion/models/ltx2/ltx2_components.py +++ b/vllm_omni/diffusion/models/ltx2/ltx2_components.py @@ -298,7 +298,7 @@ def _create_ltx25_natten_processor() -> LTX2VideoVaeNeighborhoodNattenProcessor: except (ImportError, OSError, RuntimeError, ValueError) as exc: raise RuntimeError( "LTX-2.5 DiffVAE requires the shi-labs/natten Hub kernel. " - "Install kernels==0.15.2, use a supported GPU, leave " + "Install kernels==0.16.1, use a supported GPU, leave " "DIFFUSERS_DISABLE_REMOTE_CODE unset, and allow Hub access during kernel initialization." ) from exc @@ -662,7 +662,7 @@ def initialize_pipeline_components(pipeline: Any, od_config: Any) -> None: revision=revision, ) if profile.text_encoder_cls is None: - raise ImportError("LTX-2.5 requires Gemma4UnifiedForConditionalGeneration; install transformers>=5.10.1,<5.15.") + raise ImportError("LTX-2.5 requires Gemma4UnifiedForConditionalGeneration; install transformers>=5.13.0,<5.15.") with torch.device("cpu"): pipeline.text_encoder = _load_component( profile.text_encoder_cls, diff --git a/vllm_omni/platforms/cuda/platform.py b/vllm_omni/platforms/cuda/platform.py index c635877a9..25124b4e0 100644 --- a/vllm_omni/platforms/cuda/platform.py +++ b/vllm_omni/platforms/cuda/platform.py @@ -231,6 +231,26 @@ class CudaOmniPlatform(OmniPlatform, CudaPlatformBase): "with compute capability 9.x. Select a compatible backend." ) + if backend_upper in ("FLASH_ATTN_HUB", "FLASH_ATTN_3_HUB"): + # Same loader/cache as execution (Hub versions 1 then 2). A missing + # torch/CUDA variant used to degrade silently to native/SDPA (#6971). + hub_repo = ( + "kernels-community/flash-attn3" + if backend_upper == "FLASH_ATTN_3_HUB" + else "kernels-community/flash-attn2" + ) + try: + from vllm_omni.diffusion.attention.backends.flash_attn_hub import _get_hub_module + + _get_hub_module(hub_repo) + except Exception as e: + raise RuntimeError( + f"{backend_upper} was explicitly selected, but {hub_repo} has no " + "compatible build for this torch/CUDA (or failed to load). " + "Install kernels>=0.16.1 so stable-ABI variants resolve, or select " + "a different attention backend." + ) from e + if backend_upper == "FLASH_ATTN" and not flash_attn_supported: if is_blackwell: reason = (