diff --git a/docker/Dockerfile.ci b/docker/Dockerfile.ci index 927d544..09050bd 100644 --- a/docker/Dockerfile.ci +++ b/docker/Dockerfile.ci @@ -39,9 +39,9 @@ ARG PYTORCH_CODEC_COMMIT="6df7fc8e81c9509e86833ca48695609c583e2953" # deps image arguments ARG AITER_COMMIT="0c2b0f77b2ff6d13c677d12466abf87299f8b260" ARG LONGCONTEXTATTENTION_COMMIT="631bdfd" -ARG DIFFUSERS_COMMIT="f6e1c4d1e7a6aebaeeaaa0880a82ad036517909e" -ARG DISTVAE_COMMIT="6d8025c96b9975d45badae12d8ed9b6422e34e2c" -ARG XDIT_COMMIT="2f74acb12830163d1fe0dd1abb5422f0bff7acf0" +ARG DIFFUSERS_COMMIT="9f7aee48215c5d01b1d84b52058669710596c7ee" +ARG DISTVAE_COMMIT="dce484185dc11d7406076117771e36acd796aa6a" +ARG XDIT_COMMIT="3aa7b34ec3f1d4f86f0befc5f09384333067846c" ARG AO_COMMIT="92dcc96162f21928092f8e8e35bac29d3b180e34" ARG ARBITER_COMMIT="252b0301d072f9d6ea895e32f6898fed951cf295" @@ -526,7 +526,7 @@ RUN --mount=type=bind,source=patches,target=/patches \ git apply /patches/yunchang_aiter_round_mode.patch && \ pip install -e . # install transformers with [sentencepiece] extras that brings in protobuf, required for T5 tokenizers -RUN pip install --no-cache-dir "transformers[sentencepiece]==5.5.4" +RUN pip install --no-cache-dir "transformers[sentencepiece]==5.16.1" # install diffusers # ftfy is an optional dependency, we install it as needed for StableDiffusion 3.5 (CLIP-based encoder pipelines) # final rm step is needed presently to pass a release process image scan diff --git a/docker/Dockerfile.cuda b/docker/Dockerfile.cuda index e8f4256..7070f4d 100644 --- a/docker/Dockerfile.cuda +++ b/docker/Dockerfile.cuda @@ -12,9 +12,9 @@ # deps image arguments ARG FLASHATTENTION_COMMIT="fffabc3" ARG LONGCONTEXTATTENTION_COMMIT="631bdfd" -ARG DIFFUSERS_COMMIT="f6e1c4d1e7a6aebaeeaaa0880a82ad036517909e" -ARG DISTVAE_COMMIT="6d8025c96b9975d45badae12d8ed9b6422e34e2c" -ARG XDIT_COMMIT="2f74acb12830163d1fe0dd1abb5422f0bff7acf0" +ARG DIFFUSERS_COMMIT="9f7aee48215c5d01b1d84b52058669710596c7ee" +ARG DISTVAE_COMMIT="dce484185dc11d7406076117771e36acd796aa6a" +ARG XDIT_COMMIT="3aa7b34ec3f1d4f86f0befc5f09384333067846c" ARG SAGEATTENTION_COMMIT="d1a57a5" ARG MSLK_COMMIT="e14925cede32a31ffdd5abf7a9702579d268238c" ARG FLEX_BLOCK_ATTENTION_COMMIT="bc2ec1e" @@ -100,7 +100,7 @@ RUN git clone https://github.com/feifeibear/long-context-attention.git && \ git reset --hard ${LONGCONTEXTATTENTION_COMMIT} && \ pip install -e . # install transformers with [sentencepiece] extras that brings in protobuf, required for T5 tokenizers -RUN pip install --no-cache-dir "transformers[sentencepiece]==5.5.4" +RUN pip install --no-cache-dir "transformers[sentencepiece]==5.16.1" # install diffusers # ftfy is an optional dependency, we install it as needed for StableDiffusion 3.5 (CLIP-based encoder pipelines) RUN git clone https://github.com/huggingface/diffusers.git && \ diff --git a/patches/xfuser_aiter_sparge_fp8_asm_backend.patch b/patches/xfuser_aiter_sparge_fp8_asm_backend.patch index 7ec2e67..f968026 100644 --- a/patches/xfuser_aiter_sparge_fp8_asm_backend.patch +++ b/patches/xfuser_aiter_sparge_fp8_asm_backend.patch @@ -1,8 +1,8 @@ diff --git a/xfuser/core/distributed/attention_backend.py b/xfuser/core/distributed/attention_backend.py -index 9a08501..982b9a6 100644 +index 3c978b1..0834211 100644 --- a/xfuser/core/distributed/attention_backend.py +++ b/xfuser/core/distributed/attention_backend.py -@@ -372,6 +372,10 @@ if env_info["has_aiter"]: +@@ -408,6 +408,10 @@ if env_info["has_aiter"]: from aiter.ops.triton.attention.utils import block_attn_mask_to_ragged_lut except ImportError: pass # Error is rasied in runtime_state.py if AITER_SPARSE_SAGE is not available. @@ -11,9 +11,9 @@ index 9a08501..982b9a6 100644 + except ImportError: + pass # Error is raised in runtime_state.py if AITER_SPARGE_FP8 is not available. - AITER_FP8_STATIC_SCALE_WITH_DESCALE, AITER_FP8_STATIC_SCALE_NO_DESCALE, AITER_SAGE_V2_BLOCK_R = _setup_aiter_environment_variables() - AITER_HAS_ROUND_MODE, HOW_V3_BF16_CVT = _check_aiter_round_mode() -@@ -548,6 +552,7 @@ class AttentionBackendType(Enum): + try: + from aiter.ops.mha_v4 import ( +@@ -643,6 +647,7 @@ class AttentionBackendType(Enum): AITER_SPARSE_SAGE_V2 = "AITER Sparse Sage V2" AITER_SPARGE = "AITER Sparge" AITER_SPARGE_V2 = "AITER Sparge V2" @@ -21,11 +21,10 @@ index 9a08501..982b9a6 100644 AITER_VSA = "AITER VSA CK" FLEX_BLOCK_SPARGE = "Flex Block Sparge" AITER_FLYDSL = "AITER FlyDSL" -@@ -1339,6 +1344,84 @@ def _aiter_sparge_v2_attn_call(query, key, value, dropout_p, is_causal, attentio +@@ -2000,6 +2005,83 @@ def _aiter_sparge_v2_attn_call(query, key, value, dropout_p, is_causal, attentio ) return restore_sparge_output(output, state), None -+ + +# Hand-written gfx950 sparse ASM kernels are hard-wired to (kTileQ, kTileKV) = +# (256, 128); force the LUT block sizes to match. @@ -107,10 +106,10 @@ index 9a08501..982b9a6 100644 def _flex_block_sparge_attn_call(query, key, value, dropout_p, is_causal, attention_kwargs=None): config = {"BLOCK_M": 256, "BLOCK_N": 256} diff --git a/xfuser/core/distributed/runtime_state.py b/xfuser/core/distributed/runtime_state.py -index e36df3c..598c4b8 100644 +index c7d574a..631dd50 100644 --- a/xfuser/core/distributed/runtime_state.py +++ b/xfuser/core/distributed/runtime_state.py -@@ -225,6 +225,7 @@ class RuntimeState(metaclass=ABCMeta): +@@ -241,6 +241,7 @@ class RuntimeState(metaclass=ABCMeta): AttentionBackendType.AITER_SAGE_V2, AttentionBackendType.AITER_SPARSE_SAGE_V2, AttentionBackendType.AITER_SPARGE_V2, @@ -118,7 +117,7 @@ index e36df3c..598c4b8 100644 AttentionBackendType.AITER_VSA, AttentionBackendType.AITER_FLYDSL, AttentionBackendType.AITER_FLYDSL_FP8, -@@ -321,6 +322,20 @@ class RuntimeState(metaclass=ABCMeta): +@@ -383,6 +384,20 @@ class RuntimeState(metaclass=ABCMeta): raise RuntimeError(msg) from None except ImportError: raise RuntimeError(msg) from None @@ -140,26 +139,26 @@ index e36df3c..598c4b8 100644 try: from aiter.ops.jenga_sparse_attention import vsa_sparse_attention diff --git a/xfuser/model_executor/layers/usp.py b/xfuser/model_executor/layers/usp.py -index 341a3c4..fc37433 100644 +index e700f26..962159d 100644 --- a/xfuser/model_executor/layers/usp.py +++ b/xfuser/model_executor/layers/usp.py -@@ -39,6 +39,7 @@ from xfuser.core.sparge_attention.head_balance import ( +@@ -40,6 +40,7 @@ from xfuser.core.sparge_attention.head_balance import ( _HEAD_BALANCE_BACKENDS = frozenset({ AttentionBackendType.AITER_SPARGE, AttentionBackendType.AITER_SPARGE_V2, + AttentionBackendType.AITER_SPARGE_FP8, AttentionBackendType.FLEX_BLOCK_SPARGE, - }) + }) | AITER_MHA_V4_SPARGE_BACKEND_SET diff --git a/xfuser/model_executor/models/runner_models/base_model.py b/xfuser/model_executor/models/runner_models/base_model.py -index 5ed1d07..ebb3e22 100644 +index 7356a2d..457628c 100644 --- a/xfuser/model_executor/models/runner_models/base_model.py +++ b/xfuser/model_executor/models/runner_models/base_model.py -@@ -77,6 +77,7 @@ _SPARSE_ATTENTION_BACKENDS = frozenset({ +@@ -79,6 +79,7 @@ _SPARSE_ATTENTION_BACKENDS = frozenset({ _SPARGE_ATTENTION_BACKENDS = frozenset({ AttentionBackendType.AITER_SPARGE, AttentionBackendType.AITER_SPARGE_V2, + AttentionBackendType.AITER_SPARGE_FP8, AttentionBackendType.AITER_VSA, AttentionBackendType.FLEX_BLOCK_SPARGE, - }) + }) | AITER_MHA_V4_SPARGE_BACKEND_SET diff --git a/patches/xfuser_flux.patch b/patches/xfuser_flux.patch index f0c235f..ce60912 100644 --- a/patches/xfuser_flux.patch +++ b/patches/xfuser_flux.patch @@ -1,10 +1,10 @@ diff --git a/xfuser/model_executor/models/runner_models/flux.py b/xfuser/model_executor/models/runner_models/flux.py -index 2e8d32d..a9f58f8 100644 +index 144e8b9..5bf8e0d 100644 --- a/xfuser/model_executor/models/runner_models/flux.py +++ b/xfuser/model_executor/models/runner_models/flux.py -@@ -333,7 +333,7 @@ class xFuserFlux2Model(xFuserModel): - # and cause pathological re-captures on RDNA4. - if self.config.use_fbcache or PACKAGES_CHECKER._on_rdna4(): +@@ -341,7 +341,7 @@ class xFuserFlux2Model(xFuserModel): + # cause pathological re-captures on RDNA4. + if (self.config.cache_method or PACKAGES_CHECKER._on_rdna4()): return "default" - return "reduce-overhead" + return "max-autotune-no-cudagraphs"