anodyine/vllm-omni
1
0 Fork 0
Code Issues 4 Pull Requests 2 Actions Packages Projects Releases Wiki Activity

[Kernel] Enable LTX DiffVAE fusions on SM100 and SM103 (#7350)

Browse Source 
Signed-off-by: mglyn <1203789601@qq.com>
This commit is contained in:
Mu GuanLin authored and GitHub committed 2026-09-10 03:07:35 +00:00
1 parent 35a83a0885
commit 4b8dcac681
2 files changed
+12 -8
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

tests/diffusion/models/ltx2/test_ltx2_ops_platform.py
+4 -2
View File
@@ -37,8 +37,10 @@ def clear_platform_cache():
(90, True, True, True, True),
(90, True, False, True, False),
(90, False, True, False, False),
(100, True, True, False, False),
(103, True, True, False, False),
(100, True, True, True, False),
(100, True, False, True, False),
(103, True, True, True, False),
(103, True, False, True, False),
(110, True, True, False, False),
(120, True, True, False, False),
(121, True, True, False, False),
vllm_omni/diffusion/models/ltx2/ops/__init__.py
+8 -6
View File
@@ -17,9 +17,10 @@ from vllm_omni.platforms import current_omni_platform
# Pointwise fusions are launch/memory-traffic reductions whose first execution
# for every runtime key is checked bit-for-bit; failures use the ordinary model
# math. Both paths remain SM90-only until other architectures are validated.
# The FNA schedule is performance-tuned for SM90 and has no silent fallback.
_CUDA_COMPUTE_CAPABILITIES = frozenset({90})
# math. The FNA schedule is performance-tuned for SM90 and has no silent
# fallback, so qualify the two paths independently.
_FUSION_CUDA_COMPUTE_CAPABILITIES = frozenset({90, 100, 103})
_FNA_CUDA_COMPUTE_CAPABILITIES = frozenset({90})
@dataclass(frozen=True)
@@ -66,9 +67,10 @@ def resolve_ltx2_vae_operators(device: torch.device) -> LTX2VAEOperatorSet | Non
if device.type != "cuda" or not current_omni_platform.is_cuda() or not current_omni_platform.is_available():
return None
device_index = device.index if device.index is not None else torch.accelerator.current_device_index()
if _cuda_compute_capability(int(device_index)) not in _CUDA_COMPUTE_CAPABILITIES:
capability = _cuda_compute_capability(int(device_index))
if capability not in _FUSION_CUDA_COMPUTE_CAPABILITIES:
return None
return _operator_set(has_tilelang())
return _operator_set(capability in _FNA_CUDA_COMPUTE_CAPABILITIES and has_tilelang())
def _eligible_operators(tensor: torch.Tensor) -> LTX2VAEOperatorSet | None:
@@ -78,7 +80,7 @@ def _eligible_operators(tensor: torch.Tensor) -> LTX2VAEOperatorSet | None:
def is_ltx2_fusion_eligible(tensor: torch.Tensor) -> bool:
"""Allow self-verifying pointwise fusions on validated SM90 devices."""
"""Allow self-verifying pointwise fusions on qualified CUDA devices."""
return _eligible_operators(tensor) is not None