Bump xDiT stack

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

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

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,21 +11,20 @@ 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"
 +    AITER_SPARGE_FP8 = "AITER Sparge FP8"
      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,18 +106,18 @@ 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,
 +                                 AttentionBackendType.AITER_SPARGE_FP8,
                                   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

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"