From 95a1f02bbc04179272bdcc02949bff96b21f0693 Mon Sep 17 00:00:00 2001 From: Mikko Lauri Date: Tue, 25 Aug 2026 12:54:44 +0000 Subject: [PATCH 01/13] Bump AITER_COMMIT to 4fa508ef for CI builds. Track upstream AITER at the commit validated in the aiter-bump test branch. Co-authored-by: Cursor --- docker/Dockerfile.ci | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docker/Dockerfile.ci b/docker/Dockerfile.ci index 8c04c35..aac7316 100644 --- a/docker/Dockerfile.ci +++ b/docker/Dockerfile.ci @@ -37,7 +37,7 @@ ARG PYTORCH_VISION_COMMIT="b919bd0c56abbb3c5ca056a3a458af9fd1cabf52" ARG PYTORCH_AUDIO_COMMIT="e3c6ee2b6588b7cd27a84182de74bf12fe043831" ARG PYTORCH_CODEC_COMMIT="6df7fc8e81c9509e86833ca48695609c583e2953" # deps image arguments -ARG AITER_COMMIT="3a5261496c591b4fd7cf280d35d20dee713098fd" +ARG AITER_COMMIT="4fa508ef2935110ff99adf2743ea93807dbd9c67" ARG LONGCONTEXTATTENTION_COMMIT="631bdfd" ARG DIFFUSERS_COMMIT="f6e1c4d1e7a6aebaeeaaa0880a82ad036517909e" ARG DISTVAE_COMMIT="6d8025c96b9975d45badae12d8ed9b6422e34e2c" From 2208972bbc56783b80f0899835749cd6bc438b56 Mon Sep 17 00:00:00 2001 From: Mikko Lauri Date: Tue, 25 Aug 2026 12:54:44 +0000 Subject: [PATCH 02/13] Refresh sparse FP8 ASM patch for bumped AITER. Regenerate against 4fa508ef: drop stale Union import hunk, use Optional[...] annotations for infer_schema, and retarget hunks to the new file layout. Co-authored-by: Cursor --- patches/aiter_fp8_sparse_asm_kernel.patch | 357 +++++++++++----------- 1 file changed, 174 insertions(+), 183 deletions(-) diff --git a/patches/aiter_fp8_sparse_asm_kernel.patch b/patches/aiter_fp8_sparse_asm_kernel.patch index ac32d3d..df4e11f 100644 --- a/patches/aiter_fp8_sparse_asm_kernel.patch +++ b/patches/aiter_fp8_sparse_asm_kernel.patch @@ -1,26 +1,17 @@ diff --git a/aiter/jit/core.py b/aiter/jit/core.py -index 0f2f20d1f..2c7c3abc0 100644 +index b2c62ad01..8b05ab33d 100644 --- a/aiter/jit/core.py +++ b/aiter/jit/core.py -@@ -14,7 +14,7 @@ import time - import traceback - import types - import typing --from typing import Any, Callable, List, Optional -+from typing import Any, Callable, List, Optional, Union - - from packaging.version import Version, parse - -@@ -1378,6 +1378,7 @@ def compile_ops( - gen_fake: Optional[Callable[..., Any]] = None, +@@ -1632,6 +1632,7 @@ def compile_ops( + gen_fake: Callable[..., Any] | None = None, ffi_type: str = "pybind", develop: bool = False, -+ mutates_args: Union[List[str], str] = "unknown", ++ mutates_args: list[str] | str = "unknown", ): def decorator(func): loadName = fc_name if fc_name is not None else func.__name__ -@@ -1639,7 +1640,12 @@ def compile_ops( - ) +@@ -1925,7 +1926,12 @@ def compile_ops( + module._set_current_hip_stream(raw_stream(current_device())) return op(*args, **kwargs) - @torch_compile_guard(device="cuda", gen_fake=gen_fake, calling_func_=func) @@ -34,10 +25,10 @@ index 0f2f20d1f..2c7c3abc0 100644 return wrapper(*args, **kwargs) diff --git a/aiter/jit/optCompilerConfig.json b/aiter/jit/optCompilerConfig.json -index 8cfa3b6ef..c0a1ff28c 100644 +index 171d82568..857a28c26 100644 --- a/aiter/jit/optCompilerConfig.json +++ b/aiter/jit/optCompilerConfig.json -@@ -1091,8 +1091,10 @@ +@@ -1298,8 +1298,10 @@ "srcs": [ "f'{AITER_CSRC_DIR}/kernels/mha_common.cu'", "f'{AITER_CSRC_DIR}/py_itfs_cu/asm_mha_fwd.cu'", @@ -49,170 +40,6 @@ index 8cfa3b6ef..c0a1ff28c 100644 ], "flags_extra_cc": [ "'-DFAV3_ON=1'", -diff --git a/aiter/ops/mha.py b/aiter/ops/mha.py -index 7e0379e28..ce10a3ee9 100644 ---- a/aiter/ops/mha.py -+++ b/aiter/ops/mha.py -@@ -270,6 +270,83 @@ def fmha_v3_fwd( - ) -> Tuple[Tensor, Tensor, Tensor, Tensor]: ... - - -+def _gen_fmha_v3_fwd_fp8_sparse_fake_tensors( -+ q: Tensor, -+ k: Tensor, -+ v: Tensor, -+ q_descale: Tensor, -+ k_descale: Tensor, -+ v_descale: Tensor, -+ kv_block_indices: Tensor, -+ lut_start: Tensor, -+ lut_count: Tensor, -+ work_table: Tensor, -+ softmax_scale: float, -+ out: Optional[Tensor] = None, -+) -> Tuple[Tensor]: -+ if out is not None: -+ return (out,) -+ b, sq, hq, _ = q.shape -+ head_dim_v = v.shape[-1] -+ return (q.new_empty((b, sq, hq, head_dim_v), dtype=dtypes.bf16),) -+ -+ -+@compile_ops( -+ "module_fmha_v3_fwd", -+ fc_name="fmha_v3_fwd_fp8_sparse", -+ gen_fake=_gen_fmha_v3_fwd_fp8_sparse_fake_tensors, -+ mutates_args=[], -+) -+def fmha_v3_fwd_fp8_sparse( -+ q: Tensor, # [b, sq, hq, 128], fp8 -+ k: Tensor, # [b, sk, hk, 128], fp8 -+ v: Tensor, # [b, sk, hk, 128], fp8 -+ q_descale: Tensor, # [1] or [b, hk], fp32 -+ k_descale: Tensor, # [1] or [b, hk], fp32 -+ v_descale: Tensor, # [1] or [b, hk], fp32 -+ kv_block_indices: Tensor, # int32 -+ lut_start: Tensor, # int32 [b*hq*num_q_blocks] -+ lut_count: Tensor, # int32 [b*hq*num_q_blocks] -+ work_table: Tensor, # int32 [b*hq*num_q_blocks] (packed q|h<<16|b<<24) -+ softmax_scale: float, -+ out: Optional[Tensor] = None, -+) -> Tuple[Tensor]: ... -+ -+ -+def build_sparse_work_table( -+ lut_count: Tensor, -+ batch: int, -+ num_heads: int, -+ num_q_blocks: int, -+) -> Tensor: -+ """Build the LPT-sorted work table for the fp8 sparse kernel. -+ -+ The flat tile index matches the kernel's lut_idx: -+ i = b * (HQ * num_q_blocks) + h * num_q_blocks + q -+ We sort those indices by lut_count (KV blocks per tile) descending so the -+ heaviest tiles are handed out first and spread across workgroups. Each entry -+ packs the decoded tile as -+ q | (h << 16) | (b << 24) -+ which the kernel unpacks into tgid_x/y/z (q < 65536, h < 256, b < 256). -+ -+ Returns an int32 [batch*HQ*num_q_blocks] tensor on lut_count's device. -+ """ -+ flat = lut_count.reshape(-1) -+ total = batch * num_heads * num_q_blocks -+ assert flat.numel() == total, ( -+ f"lut_count.numel()={flat.numel()} != batch*HQ*num_q_blocks={total}" -+ ) -+ # Descending LPT order over flat tile indices (tie order is irrelevant). -+ order = torch.argsort(flat, descending=True) # int64 tile indices -+ q = order % num_q_blocks -+ h = (order // num_q_blocks) % num_heads -+ b = order // (num_q_blocks * num_heads) -+ packed = ( -+ (q & 0xFFFF) | ((h & 0xFF) << 16) | ((b & 0xFF) << 24) -+ ).to(torch.int32) -+ return packed.contiguous() -+ -+ - def cmdGenFunc_mha_varlen_fwd( - q: torch.Tensor, - k: torch.Tensor, -@@ -3101,6 +3178,75 @@ def flash_attn_fp8_pertensor_func( - return out - - -+def flash_attn_fp8_sparse_pertensor_func( -+ q: torch.Tensor, -+ k: torch.Tensor, -+ v: torch.Tensor, -+ q_descale: torch.Tensor, -+ k_descale: torch.Tensor, -+ v_descale: torch.Tensor, -+ kv_block_indices: torch.Tensor, -+ lut_start: torch.Tensor, -+ lut_count: torch.Tensor, -+ softmax_scale: Optional[float] = None, -+): -+ """Block-sparse Sage fp8 FMHA forward (hd=128, gfx950). -+ -+ All-fp8 (E4M3) block-sparse attention routed to the LPT-sorted, intra-GPU -+ load-balanced kernel fwd_hd128_fp8_sparse.co. The host builds the LPT work -+ table from lut_count and launches on a flat grid (total_tiles, 1, 1). -+ -+ Args: -+ q: fp8 tensor [b, sq, hq, 128], bshd -+ k: fp8 tensor [b, sk, hk, 128], bshd -+ v: fp8 tensor [b, sk, hk, 128], bshd -+ q_descale, k_descale, v_descale: fp32 scalars or [b, hk] -+ kv_block_indices, lut_start, lut_count: int32 LUT produced by -+ aiter.ops.triton.attention.utils.block_attn_mask_to_ragged_lut -+ (return_none_if_dense=False; this kernel has no dense path). -+ softmax_scale: if None, defaults to head_dim**-0.5 -+ -+ Constraints (assert-checked C++-side, see asm_mha_fwd_sparse.cu:: -+ fmha_v3_fwd_fp8_sparse): -+ sq % 256 == 0, sk % 128 == 0, hq % hk == 0, (hq/hk) is a power of 2, -+ head_dim == 128, batch mode, non-causal, gfx950. -+ -+ Returns: -+ out: bf16 tensor [b, sq, hq, 128], bshd -+ """ -+ if softmax_scale is None: -+ softmax_scale = q.shape[-1] ** (-0.5) -+ head_size_q_og = q.size(3) -+ head_size_v_og = v.size(3) -+ if head_size_q_og % 8 != 0: -+ q = torch.nn.functional.pad(q, [0, 8 - head_size_q_og % 8]) -+ k = torch.nn.functional.pad(k, [0, 8 - head_size_q_og % 8]) -+ if head_size_v_og % 8 != 0: -+ v = torch.nn.functional.pad(v, [0, 8 - head_size_v_og % 8]) -+ batch, seqlen_q, num_heads = q.shape[0], q.shape[1], q.shape[2] -+ num_q_blocks = (seqlen_q + 255) // 256 # kTileQ = 256 -+ work_table = build_sparse_work_table( -+ lut_count, batch, num_heads, num_q_blocks -+ ) -+ outs = fmha_v3_fwd_fp8_sparse( -+ q, -+ k, -+ v, -+ q_descale, -+ k_descale, -+ v_descale, -+ kv_block_indices, -+ lut_start, -+ lut_count, -+ work_table, -+ float(softmax_scale), -+ None, -+ ) -+ out_padded = outs[0] -+ out = out_padded[..., :head_size_v_og] -+ return out -+ -+ - def flash_attn_varlen_fp8_pertensor_func( - q, - k, diff --git a/csrc/cpp_itfs/mha_fwd_sparse.cu b/csrc/cpp_itfs/mha_fwd_sparse.cu new file mode 100644 index 000000000..3491e675d @@ -476,10 +303,10 @@ index 000000000..7f20e2386 + +} // namespace aiter diff --git a/csrc/include/rocm_ops.hpp b/csrc/include/rocm_ops.hpp -index d42014c5d..0ab504e0f 100644 +index b090e610e..705e928da 100644 --- a/csrc/include/rocm_ops.hpp +++ b/csrc/include/rocm_ops.hpp -@@ -823,7 +823,21 @@ namespace py = pybind11; +@@ -1088,7 +1088,21 @@ namespace py = pybind11; py::arg("q_descale") = std::nullopt, \ py::arg("k_descale") = std::nullopt, \ py::arg("v_descale") = std::nullopt, \ @@ -785,3 +612,167 @@ index 298e61dbc..f2e0a1478 100644 PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { +diff --git aa/aiter/ops/mha.py bb/aiter/ops/mha.py +index 6f0ecbff5..13bf187d0 100644 +--- aa/aiter/ops/mha.py ++++ bb/aiter/ops/mha.py +@@ -726,6 +726,83 @@ def fmha_fwd_mxfp8_asm( + return out, lse + + ++def _gen_fmha_v3_fwd_fp8_sparse_fake_tensors( ++ q: Tensor, ++ k: Tensor, ++ v: Tensor, ++ q_descale: Tensor, ++ k_descale: Tensor, ++ v_descale: Tensor, ++ kv_block_indices: Tensor, ++ lut_start: Tensor, ++ lut_count: Tensor, ++ work_table: Tensor, ++ softmax_scale: float, ++ out: Optional[Tensor] = None, ++) -> tuple[Tensor]: ++ if out is not None: ++ return (out,) ++ b, sq, hq, _ = q.shape ++ head_dim_v = v.shape[-1] ++ return (q.new_empty((b, sq, hq, head_dim_v), dtype=dtypes.bf16),) ++ ++ ++@compile_ops( ++ "module_fmha_v3_fwd", ++ fc_name="fmha_v3_fwd_fp8_sparse", ++ gen_fake=_gen_fmha_v3_fwd_fp8_sparse_fake_tensors, ++ mutates_args=[], ++) ++def fmha_v3_fwd_fp8_sparse( ++ q: Tensor, # [b, sq, hq, 128], fp8 ++ k: Tensor, # [b, sk, hk, 128], fp8 ++ v: Tensor, # [b, sk, hk, 128], fp8 ++ q_descale: Tensor, # [1] or [b, hk], fp32 ++ k_descale: Tensor, # [1] or [b, hk], fp32 ++ v_descale: Tensor, # [1] or [b, hk], fp32 ++ kv_block_indices: Tensor, # int32 ++ lut_start: Tensor, # int32 [b*hq*num_q_blocks] ++ lut_count: Tensor, # int32 [b*hq*num_q_blocks] ++ work_table: Tensor, # int32 [b*hq*num_q_blocks] (packed q|h<<16|b<<24) ++ softmax_scale: float, ++ out: Optional[Tensor] = None, ++) -> tuple[Tensor]: ... ++ ++ ++def build_sparse_work_table( ++ lut_count: Tensor, ++ batch: int, ++ num_heads: int, ++ num_q_blocks: int, ++) -> Tensor: ++ """Build the LPT-sorted work table for the fp8 sparse kernel. ++ ++ The flat tile index matches the kernel's lut_idx: ++ i = b * (HQ * num_q_blocks) + h * num_q_blocks + q ++ We sort those indices by lut_count (KV blocks per tile) descending so the ++ heaviest tiles are handed out first and spread across workgroups. Each entry ++ packs the decoded tile as ++ q | (h << 16) | (b << 24) ++ which the kernel unpacks into tgid_x/y/z (q < 65536, h < 256, b < 256). ++ ++ Returns an int32 [batch*HQ*num_q_blocks] tensor on lut_count's device. ++ """ ++ flat = lut_count.reshape(-1) ++ total = batch * num_heads * num_q_blocks ++ assert flat.numel() == total, ( ++ f"lut_count.numel()={flat.numel()} != batch*HQ*num_q_blocks={total}" ++ ) ++ # Descending LPT order over flat tile indices (tie order is irrelevant). ++ order = torch.argsort(flat, descending=True) # int64 tile indices ++ q = order % num_q_blocks ++ h = (order // num_q_blocks) % num_heads ++ b = order // (num_q_blocks * num_heads) ++ packed = ( ++ (q & 0xFFFF) | ((h & 0xFF) << 16) | ((b & 0xFF) << 24) ++ ).to(torch.int32) ++ return packed.contiguous() ++ ++ + def cmdGenFunc_mha_varlen_fwd( + q: torch.Tensor, + k: torch.Tensor, +@@ -4032,6 +4109,75 @@ def flash_attn_fp8_pertensor_func( + return out + + ++def flash_attn_fp8_sparse_pertensor_func( ++ q: torch.Tensor, ++ k: torch.Tensor, ++ v: torch.Tensor, ++ q_descale: torch.Tensor, ++ k_descale: torch.Tensor, ++ v_descale: torch.Tensor, ++ kv_block_indices: torch.Tensor, ++ lut_start: torch.Tensor, ++ lut_count: torch.Tensor, ++ softmax_scale: Optional[float] = None, ++): ++ """Block-sparse Sage fp8 FMHA forward (hd=128, gfx950). ++ ++ All-fp8 (E4M3) block-sparse attention routed to the LPT-sorted, intra-GPU ++ load-balanced kernel fwd_hd128_fp8_sparse.co. The host builds the LPT work ++ table from lut_count and launches on a flat grid (total_tiles, 1, 1). ++ ++ Args: ++ q: fp8 tensor [b, sq, hq, 128], bshd ++ k: fp8 tensor [b, sk, hk, 128], bshd ++ v: fp8 tensor [b, sk, hk, 128], bshd ++ q_descale, k_descale, v_descale: fp32 scalars or [b, hk] ++ kv_block_indices, lut_start, lut_count: int32 LUT produced by ++ aiter.ops.triton.attention.utils.block_attn_mask_to_ragged_lut ++ (return_none_if_dense=False; this kernel has no dense path). ++ softmax_scale: if None, defaults to head_dim**-0.5 ++ ++ Constraints (assert-checked C++-side, see asm_mha_fwd_sparse.cu:: ++ fmha_v3_fwd_fp8_sparse): ++ sq % 256 == 0, sk % 128 == 0, hq % hk == 0, (hq/hk) is a power of 2, ++ head_dim == 128, batch mode, non-causal, gfx950. ++ ++ Returns: ++ out: bf16 tensor [b, sq, hq, 128], bshd ++ """ ++ if softmax_scale is None: ++ softmax_scale = q.shape[-1] ** (-0.5) ++ head_size_q_og = q.size(3) ++ head_size_v_og = v.size(3) ++ if head_size_q_og % 8 != 0: ++ q = torch.nn.functional.pad(q, [0, 8 - head_size_q_og % 8]) ++ k = torch.nn.functional.pad(k, [0, 8 - head_size_q_og % 8]) ++ if head_size_v_og % 8 != 0: ++ v = torch.nn.functional.pad(v, [0, 8 - head_size_v_og % 8]) ++ batch, seqlen_q, num_heads = q.shape[0], q.shape[1], q.shape[2] ++ num_q_blocks = (seqlen_q + 255) // 256 # kTileQ = 256 ++ work_table = build_sparse_work_table( ++ lut_count, batch, num_heads, num_q_blocks ++ ) ++ outs = fmha_v3_fwd_fp8_sparse( ++ q, ++ k, ++ v, ++ q_descale, ++ k_descale, ++ v_descale, ++ kv_block_indices, ++ lut_start, ++ lut_count, ++ work_table, ++ float(softmax_scale), ++ None, ++ ) ++ out_padded = outs[0] ++ out = out_padded[..., :head_size_v_og] ++ return out ++ ++ + def flash_attn_varlen_fp8_pertensor_func( + q, + k, From 27ada849a889c7a5fa78912cb005daa105086465 Mon Sep 17 00:00:00 2001 From: Mikko Lauri Date: Tue, 25 Aug 2026 12:54:44 +0000 Subject: [PATCH 03/13] Work around PyTorch 2.9 Dynamo PEP 604 graph break in AITER MHA. Nested _validate_cu in aiter/ops/mha.py breaks torch.compile when annotated with PEP 604 syntax. Apply aiter_mha_pep604_dynamo.patch in Dockerfile.ci; remove it after PyTorch is bumped above 2.9. Co-authored-by: Cursor --- docker/Dockerfile.ci | 3 +++ patches/aiter_mha_pep604_dynamo.patch | 22 ++++++++++++++++++++++ 2 files changed, 25 insertions(+) create mode 100644 patches/aiter_mha_pep604_dynamo.patch diff --git a/docker/Dockerfile.ci b/docker/Dockerfile.ci index aac7316..9c7b190 100644 --- a/docker/Dockerfile.ci +++ b/docker/Dockerfile.ci @@ -486,6 +486,9 @@ RUN --mount=type=bind,source=patches,target=/patches \ cd aiter && \ git reset --hard ${AITER_COMMIT} && \ git apply /patches/aiter_fp8_sparse_asm_kernel.patch && \ + # PyTorch 2.9 Dynamo graph-breaks on PEP 604 in nested _validate_cu (aiter/ops/mha.py). + # Remove this patch after PyTorch is bumped above 2.9. + git apply /patches/aiter_mha_pep604_dynamo.patch && \ cp /code_objects/fwd_hd128_fp8.co hsa/gfx950/fmha_v3_fwd/fwd_hd128_fp8.co && \ cp /code_objects/fwd_hd128_fp8_sparse.co hsa/gfx950/fmha_v3_fwd/fwd_hd128_fp8_sparse.co && \ git submodule update --init --recursive && \ diff --git a/patches/aiter_mha_pep604_dynamo.patch b/patches/aiter_mha_pep604_dynamo.patch new file mode 100644 index 0000000..9f8bab7 --- /dev/null +++ b/patches/aiter_mha_pep604_dynamo.patch @@ -0,0 +1,22 @@ +diff --git a/aiter/ops/mha.py b/aiter/ops/mha.py +index deef9d5..a1b2c3d 100644 +--- a/aiter/ops/mha.py ++++ b/aiter/ops/mha.py +@@ -2,7 +2,7 @@ + # Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved. + + import os +-from typing import Any ++from typing import Any, Optional + + import torch + from torch import Generator, Tensor +@@ -2060,7 +2060,7 @@ def _flash_attn_forward( + + # Validate newly added optional cumulative length / padded arrays if provided. + # They are currently only plumbed through for future CK support enabling per-batch padding. +- def _validate_cu(name: str, x: torch.Tensor | None): ++ def _validate_cu(name: str, x: Optional[torch.Tensor]): + if x is None: + return + assert x.dim() == 1, f"{name} must be 1D" From dae232e9e61a0ee029516c461c7043763748d5ed Mon Sep 17 00:00:00 2001 From: Mikko Lauri Date: Thu, 27 Aug 2026 09:10:35 +0000 Subject: [PATCH 04/13] Squashed commit of the following: commit 27ada849a889c7a5fa78912cb005daa105086465 Author: Mikko Lauri Date: Tue Aug 25 12:54:44 2026 +0000 Work around PyTorch 2.9 Dynamo PEP 604 graph break in AITER MHA. Nested _validate_cu in aiter/ops/mha.py breaks torch.compile when annotated with PEP 604 syntax. Apply aiter_mha_pep604_dynamo.patch in Dockerfile.ci; remove it after PyTorch is bumped above 2.9. Co-authored-by: Cursor commit 2208972bbc56783b80f0899835749cd6bc438b56 Author: Mikko Lauri Date: Tue Aug 25 12:54:44 2026 +0000 Refresh sparse FP8 ASM patch for bumped AITER. Regenerate against 4fa508ef: drop stale Union import hunk, use Optional[...] annotations for infer_schema, and retarget hunks to the new file layout. Co-authored-by: Cursor commit 95a1f02bbc04179272bdcc02949bff96b21f0693 Author: Mikko Lauri Date: Tue Aug 25 12:54:44 2026 +0000 Bump AITER_COMMIT to 4fa508ef for CI builds. Track upstream AITER at the commit validated in the aiter-bump test branch. Co-authored-by: Cursor --- docker/Dockerfile.ci | 5 +- patches/aiter_fp8_sparse_asm_kernel.patch | 357 +++++++++++----------- patches/aiter_mha_pep604_dynamo.patch | 22 ++ 3 files changed, 200 insertions(+), 184 deletions(-) create mode 100644 patches/aiter_mha_pep604_dynamo.patch diff --git a/docker/Dockerfile.ci b/docker/Dockerfile.ci index 8c04c35..9c7b190 100644 --- a/docker/Dockerfile.ci +++ b/docker/Dockerfile.ci @@ -37,7 +37,7 @@ ARG PYTORCH_VISION_COMMIT="b919bd0c56abbb3c5ca056a3a458af9fd1cabf52" ARG PYTORCH_AUDIO_COMMIT="e3c6ee2b6588b7cd27a84182de74bf12fe043831" ARG PYTORCH_CODEC_COMMIT="6df7fc8e81c9509e86833ca48695609c583e2953" # deps image arguments -ARG AITER_COMMIT="3a5261496c591b4fd7cf280d35d20dee713098fd" +ARG AITER_COMMIT="4fa508ef2935110ff99adf2743ea93807dbd9c67" ARG LONGCONTEXTATTENTION_COMMIT="631bdfd" ARG DIFFUSERS_COMMIT="f6e1c4d1e7a6aebaeeaaa0880a82ad036517909e" ARG DISTVAE_COMMIT="6d8025c96b9975d45badae12d8ed9b6422e34e2c" @@ -486,6 +486,9 @@ RUN --mount=type=bind,source=patches,target=/patches \ cd aiter && \ git reset --hard ${AITER_COMMIT} && \ git apply /patches/aiter_fp8_sparse_asm_kernel.patch && \ + # PyTorch 2.9 Dynamo graph-breaks on PEP 604 in nested _validate_cu (aiter/ops/mha.py). + # Remove this patch after PyTorch is bumped above 2.9. + git apply /patches/aiter_mha_pep604_dynamo.patch && \ cp /code_objects/fwd_hd128_fp8.co hsa/gfx950/fmha_v3_fwd/fwd_hd128_fp8.co && \ cp /code_objects/fwd_hd128_fp8_sparse.co hsa/gfx950/fmha_v3_fwd/fwd_hd128_fp8_sparse.co && \ git submodule update --init --recursive && \ diff --git a/patches/aiter_fp8_sparse_asm_kernel.patch b/patches/aiter_fp8_sparse_asm_kernel.patch index ac32d3d..df4e11f 100644 --- a/patches/aiter_fp8_sparse_asm_kernel.patch +++ b/patches/aiter_fp8_sparse_asm_kernel.patch @@ -1,26 +1,17 @@ diff --git a/aiter/jit/core.py b/aiter/jit/core.py -index 0f2f20d1f..2c7c3abc0 100644 +index b2c62ad01..8b05ab33d 100644 --- a/aiter/jit/core.py +++ b/aiter/jit/core.py -@@ -14,7 +14,7 @@ import time - import traceback - import types - import typing --from typing import Any, Callable, List, Optional -+from typing import Any, Callable, List, Optional, Union - - from packaging.version import Version, parse - -@@ -1378,6 +1378,7 @@ def compile_ops( - gen_fake: Optional[Callable[..., Any]] = None, +@@ -1632,6 +1632,7 @@ def compile_ops( + gen_fake: Callable[..., Any] | None = None, ffi_type: str = "pybind", develop: bool = False, -+ mutates_args: Union[List[str], str] = "unknown", ++ mutates_args: list[str] | str = "unknown", ): def decorator(func): loadName = fc_name if fc_name is not None else func.__name__ -@@ -1639,7 +1640,12 @@ def compile_ops( - ) +@@ -1925,7 +1926,12 @@ def compile_ops( + module._set_current_hip_stream(raw_stream(current_device())) return op(*args, **kwargs) - @torch_compile_guard(device="cuda", gen_fake=gen_fake, calling_func_=func) @@ -34,10 +25,10 @@ index 0f2f20d1f..2c7c3abc0 100644 return wrapper(*args, **kwargs) diff --git a/aiter/jit/optCompilerConfig.json b/aiter/jit/optCompilerConfig.json -index 8cfa3b6ef..c0a1ff28c 100644 +index 171d82568..857a28c26 100644 --- a/aiter/jit/optCompilerConfig.json +++ b/aiter/jit/optCompilerConfig.json -@@ -1091,8 +1091,10 @@ +@@ -1298,8 +1298,10 @@ "srcs": [ "f'{AITER_CSRC_DIR}/kernels/mha_common.cu'", "f'{AITER_CSRC_DIR}/py_itfs_cu/asm_mha_fwd.cu'", @@ -49,170 +40,6 @@ index 8cfa3b6ef..c0a1ff28c 100644 ], "flags_extra_cc": [ "'-DFAV3_ON=1'", -diff --git a/aiter/ops/mha.py b/aiter/ops/mha.py -index 7e0379e28..ce10a3ee9 100644 ---- a/aiter/ops/mha.py -+++ b/aiter/ops/mha.py -@@ -270,6 +270,83 @@ def fmha_v3_fwd( - ) -> Tuple[Tensor, Tensor, Tensor, Tensor]: ... - - -+def _gen_fmha_v3_fwd_fp8_sparse_fake_tensors( -+ q: Tensor, -+ k: Tensor, -+ v: Tensor, -+ q_descale: Tensor, -+ k_descale: Tensor, -+ v_descale: Tensor, -+ kv_block_indices: Tensor, -+ lut_start: Tensor, -+ lut_count: Tensor, -+ work_table: Tensor, -+ softmax_scale: float, -+ out: Optional[Tensor] = None, -+) -> Tuple[Tensor]: -+ if out is not None: -+ return (out,) -+ b, sq, hq, _ = q.shape -+ head_dim_v = v.shape[-1] -+ return (q.new_empty((b, sq, hq, head_dim_v), dtype=dtypes.bf16),) -+ -+ -+@compile_ops( -+ "module_fmha_v3_fwd", -+ fc_name="fmha_v3_fwd_fp8_sparse", -+ gen_fake=_gen_fmha_v3_fwd_fp8_sparse_fake_tensors, -+ mutates_args=[], -+) -+def fmha_v3_fwd_fp8_sparse( -+ q: Tensor, # [b, sq, hq, 128], fp8 -+ k: Tensor, # [b, sk, hk, 128], fp8 -+ v: Tensor, # [b, sk, hk, 128], fp8 -+ q_descale: Tensor, # [1] or [b, hk], fp32 -+ k_descale: Tensor, # [1] or [b, hk], fp32 -+ v_descale: Tensor, # [1] or [b, hk], fp32 -+ kv_block_indices: Tensor, # int32 -+ lut_start: Tensor, # int32 [b*hq*num_q_blocks] -+ lut_count: Tensor, # int32 [b*hq*num_q_blocks] -+ work_table: Tensor, # int32 [b*hq*num_q_blocks] (packed q|h<<16|b<<24) -+ softmax_scale: float, -+ out: Optional[Tensor] = None, -+) -> Tuple[Tensor]: ... -+ -+ -+def build_sparse_work_table( -+ lut_count: Tensor, -+ batch: int, -+ num_heads: int, -+ num_q_blocks: int, -+) -> Tensor: -+ """Build the LPT-sorted work table for the fp8 sparse kernel. -+ -+ The flat tile index matches the kernel's lut_idx: -+ i = b * (HQ * num_q_blocks) + h * num_q_blocks + q -+ We sort those indices by lut_count (KV blocks per tile) descending so the -+ heaviest tiles are handed out first and spread across workgroups. Each entry -+ packs the decoded tile as -+ q | (h << 16) | (b << 24) -+ which the kernel unpacks into tgid_x/y/z (q < 65536, h < 256, b < 256). -+ -+ Returns an int32 [batch*HQ*num_q_blocks] tensor on lut_count's device. -+ """ -+ flat = lut_count.reshape(-1) -+ total = batch * num_heads * num_q_blocks -+ assert flat.numel() == total, ( -+ f"lut_count.numel()={flat.numel()} != batch*HQ*num_q_blocks={total}" -+ ) -+ # Descending LPT order over flat tile indices (tie order is irrelevant). -+ order = torch.argsort(flat, descending=True) # int64 tile indices -+ q = order % num_q_blocks -+ h = (order // num_q_blocks) % num_heads -+ b = order // (num_q_blocks * num_heads) -+ packed = ( -+ (q & 0xFFFF) | ((h & 0xFF) << 16) | ((b & 0xFF) << 24) -+ ).to(torch.int32) -+ return packed.contiguous() -+ -+ - def cmdGenFunc_mha_varlen_fwd( - q: torch.Tensor, - k: torch.Tensor, -@@ -3101,6 +3178,75 @@ def flash_attn_fp8_pertensor_func( - return out - - -+def flash_attn_fp8_sparse_pertensor_func( -+ q: torch.Tensor, -+ k: torch.Tensor, -+ v: torch.Tensor, -+ q_descale: torch.Tensor, -+ k_descale: torch.Tensor, -+ v_descale: torch.Tensor, -+ kv_block_indices: torch.Tensor, -+ lut_start: torch.Tensor, -+ lut_count: torch.Tensor, -+ softmax_scale: Optional[float] = None, -+): -+ """Block-sparse Sage fp8 FMHA forward (hd=128, gfx950). -+ -+ All-fp8 (E4M3) block-sparse attention routed to the LPT-sorted, intra-GPU -+ load-balanced kernel fwd_hd128_fp8_sparse.co. The host builds the LPT work -+ table from lut_count and launches on a flat grid (total_tiles, 1, 1). -+ -+ Args: -+ q: fp8 tensor [b, sq, hq, 128], bshd -+ k: fp8 tensor [b, sk, hk, 128], bshd -+ v: fp8 tensor [b, sk, hk, 128], bshd -+ q_descale, k_descale, v_descale: fp32 scalars or [b, hk] -+ kv_block_indices, lut_start, lut_count: int32 LUT produced by -+ aiter.ops.triton.attention.utils.block_attn_mask_to_ragged_lut -+ (return_none_if_dense=False; this kernel has no dense path). -+ softmax_scale: if None, defaults to head_dim**-0.5 -+ -+ Constraints (assert-checked C++-side, see asm_mha_fwd_sparse.cu:: -+ fmha_v3_fwd_fp8_sparse): -+ sq % 256 == 0, sk % 128 == 0, hq % hk == 0, (hq/hk) is a power of 2, -+ head_dim == 128, batch mode, non-causal, gfx950. -+ -+ Returns: -+ out: bf16 tensor [b, sq, hq, 128], bshd -+ """ -+ if softmax_scale is None: -+ softmax_scale = q.shape[-1] ** (-0.5) -+ head_size_q_og = q.size(3) -+ head_size_v_og = v.size(3) -+ if head_size_q_og % 8 != 0: -+ q = torch.nn.functional.pad(q, [0, 8 - head_size_q_og % 8]) -+ k = torch.nn.functional.pad(k, [0, 8 - head_size_q_og % 8]) -+ if head_size_v_og % 8 != 0: -+ v = torch.nn.functional.pad(v, [0, 8 - head_size_v_og % 8]) -+ batch, seqlen_q, num_heads = q.shape[0], q.shape[1], q.shape[2] -+ num_q_blocks = (seqlen_q + 255) // 256 # kTileQ = 256 -+ work_table = build_sparse_work_table( -+ lut_count, batch, num_heads, num_q_blocks -+ ) -+ outs = fmha_v3_fwd_fp8_sparse( -+ q, -+ k, -+ v, -+ q_descale, -+ k_descale, -+ v_descale, -+ kv_block_indices, -+ lut_start, -+ lut_count, -+ work_table, -+ float(softmax_scale), -+ None, -+ ) -+ out_padded = outs[0] -+ out = out_padded[..., :head_size_v_og] -+ return out -+ -+ - def flash_attn_varlen_fp8_pertensor_func( - q, - k, diff --git a/csrc/cpp_itfs/mha_fwd_sparse.cu b/csrc/cpp_itfs/mha_fwd_sparse.cu new file mode 100644 index 000000000..3491e675d @@ -476,10 +303,10 @@ index 000000000..7f20e2386 + +} // namespace aiter diff --git a/csrc/include/rocm_ops.hpp b/csrc/include/rocm_ops.hpp -index d42014c5d..0ab504e0f 100644 +index b090e610e..705e928da 100644 --- a/csrc/include/rocm_ops.hpp +++ b/csrc/include/rocm_ops.hpp -@@ -823,7 +823,21 @@ namespace py = pybind11; +@@ -1088,7 +1088,21 @@ namespace py = pybind11; py::arg("q_descale") = std::nullopt, \ py::arg("k_descale") = std::nullopt, \ py::arg("v_descale") = std::nullopt, \ @@ -785,3 +612,167 @@ index 298e61dbc..f2e0a1478 100644 PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { +diff --git aa/aiter/ops/mha.py bb/aiter/ops/mha.py +index 6f0ecbff5..13bf187d0 100644 +--- aa/aiter/ops/mha.py ++++ bb/aiter/ops/mha.py +@@ -726,6 +726,83 @@ def fmha_fwd_mxfp8_asm( + return out, lse + + ++def _gen_fmha_v3_fwd_fp8_sparse_fake_tensors( ++ q: Tensor, ++ k: Tensor, ++ v: Tensor, ++ q_descale: Tensor, ++ k_descale: Tensor, ++ v_descale: Tensor, ++ kv_block_indices: Tensor, ++ lut_start: Tensor, ++ lut_count: Tensor, ++ work_table: Tensor, ++ softmax_scale: float, ++ out: Optional[Tensor] = None, ++) -> tuple[Tensor]: ++ if out is not None: ++ return (out,) ++ b, sq, hq, _ = q.shape ++ head_dim_v = v.shape[-1] ++ return (q.new_empty((b, sq, hq, head_dim_v), dtype=dtypes.bf16),) ++ ++ ++@compile_ops( ++ "module_fmha_v3_fwd", ++ fc_name="fmha_v3_fwd_fp8_sparse", ++ gen_fake=_gen_fmha_v3_fwd_fp8_sparse_fake_tensors, ++ mutates_args=[], ++) ++def fmha_v3_fwd_fp8_sparse( ++ q: Tensor, # [b, sq, hq, 128], fp8 ++ k: Tensor, # [b, sk, hk, 128], fp8 ++ v: Tensor, # [b, sk, hk, 128], fp8 ++ q_descale: Tensor, # [1] or [b, hk], fp32 ++ k_descale: Tensor, # [1] or [b, hk], fp32 ++ v_descale: Tensor, # [1] or [b, hk], fp32 ++ kv_block_indices: Tensor, # int32 ++ lut_start: Tensor, # int32 [b*hq*num_q_blocks] ++ lut_count: Tensor, # int32 [b*hq*num_q_blocks] ++ work_table: Tensor, # int32 [b*hq*num_q_blocks] (packed q|h<<16|b<<24) ++ softmax_scale: float, ++ out: Optional[Tensor] = None, ++) -> tuple[Tensor]: ... ++ ++ ++def build_sparse_work_table( ++ lut_count: Tensor, ++ batch: int, ++ num_heads: int, ++ num_q_blocks: int, ++) -> Tensor: ++ """Build the LPT-sorted work table for the fp8 sparse kernel. ++ ++ The flat tile index matches the kernel's lut_idx: ++ i = b * (HQ * num_q_blocks) + h * num_q_blocks + q ++ We sort those indices by lut_count (KV blocks per tile) descending so the ++ heaviest tiles are handed out first and spread across workgroups. Each entry ++ packs the decoded tile as ++ q | (h << 16) | (b << 24) ++ which the kernel unpacks into tgid_x/y/z (q < 65536, h < 256, b < 256). ++ ++ Returns an int32 [batch*HQ*num_q_blocks] tensor on lut_count's device. ++ """ ++ flat = lut_count.reshape(-1) ++ total = batch * num_heads * num_q_blocks ++ assert flat.numel() == total, ( ++ f"lut_count.numel()={flat.numel()} != batch*HQ*num_q_blocks={total}" ++ ) ++ # Descending LPT order over flat tile indices (tie order is irrelevant). ++ order = torch.argsort(flat, descending=True) # int64 tile indices ++ q = order % num_q_blocks ++ h = (order // num_q_blocks) % num_heads ++ b = order // (num_q_blocks * num_heads) ++ packed = ( ++ (q & 0xFFFF) | ((h & 0xFF) << 16) | ((b & 0xFF) << 24) ++ ).to(torch.int32) ++ return packed.contiguous() ++ ++ + def cmdGenFunc_mha_varlen_fwd( + q: torch.Tensor, + k: torch.Tensor, +@@ -4032,6 +4109,75 @@ def flash_attn_fp8_pertensor_func( + return out + + ++def flash_attn_fp8_sparse_pertensor_func( ++ q: torch.Tensor, ++ k: torch.Tensor, ++ v: torch.Tensor, ++ q_descale: torch.Tensor, ++ k_descale: torch.Tensor, ++ v_descale: torch.Tensor, ++ kv_block_indices: torch.Tensor, ++ lut_start: torch.Tensor, ++ lut_count: torch.Tensor, ++ softmax_scale: Optional[float] = None, ++): ++ """Block-sparse Sage fp8 FMHA forward (hd=128, gfx950). ++ ++ All-fp8 (E4M3) block-sparse attention routed to the LPT-sorted, intra-GPU ++ load-balanced kernel fwd_hd128_fp8_sparse.co. The host builds the LPT work ++ table from lut_count and launches on a flat grid (total_tiles, 1, 1). ++ ++ Args: ++ q: fp8 tensor [b, sq, hq, 128], bshd ++ k: fp8 tensor [b, sk, hk, 128], bshd ++ v: fp8 tensor [b, sk, hk, 128], bshd ++ q_descale, k_descale, v_descale: fp32 scalars or [b, hk] ++ kv_block_indices, lut_start, lut_count: int32 LUT produced by ++ aiter.ops.triton.attention.utils.block_attn_mask_to_ragged_lut ++ (return_none_if_dense=False; this kernel has no dense path). ++ softmax_scale: if None, defaults to head_dim**-0.5 ++ ++ Constraints (assert-checked C++-side, see asm_mha_fwd_sparse.cu:: ++ fmha_v3_fwd_fp8_sparse): ++ sq % 256 == 0, sk % 128 == 0, hq % hk == 0, (hq/hk) is a power of 2, ++ head_dim == 128, batch mode, non-causal, gfx950. ++ ++ Returns: ++ out: bf16 tensor [b, sq, hq, 128], bshd ++ """ ++ if softmax_scale is None: ++ softmax_scale = q.shape[-1] ** (-0.5) ++ head_size_q_og = q.size(3) ++ head_size_v_og = v.size(3) ++ if head_size_q_og % 8 != 0: ++ q = torch.nn.functional.pad(q, [0, 8 - head_size_q_og % 8]) ++ k = torch.nn.functional.pad(k, [0, 8 - head_size_q_og % 8]) ++ if head_size_v_og % 8 != 0: ++ v = torch.nn.functional.pad(v, [0, 8 - head_size_v_og % 8]) ++ batch, seqlen_q, num_heads = q.shape[0], q.shape[1], q.shape[2] ++ num_q_blocks = (seqlen_q + 255) // 256 # kTileQ = 256 ++ work_table = build_sparse_work_table( ++ lut_count, batch, num_heads, num_q_blocks ++ ) ++ outs = fmha_v3_fwd_fp8_sparse( ++ q, ++ k, ++ v, ++ q_descale, ++ k_descale, ++ v_descale, ++ kv_block_indices, ++ lut_start, ++ lut_count, ++ work_table, ++ float(softmax_scale), ++ None, ++ ) ++ out_padded = outs[0] ++ out = out_padded[..., :head_size_v_og] ++ return out ++ ++ + def flash_attn_varlen_fp8_pertensor_func( + q, + k, diff --git a/patches/aiter_mha_pep604_dynamo.patch b/patches/aiter_mha_pep604_dynamo.patch new file mode 100644 index 0000000..9f8bab7 --- /dev/null +++ b/patches/aiter_mha_pep604_dynamo.patch @@ -0,0 +1,22 @@ +diff --git a/aiter/ops/mha.py b/aiter/ops/mha.py +index deef9d5..a1b2c3d 100644 +--- a/aiter/ops/mha.py ++++ b/aiter/ops/mha.py +@@ -2,7 +2,7 @@ + # Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved. + + import os +-from typing import Any ++from typing import Any, Optional + + import torch + from torch import Generator, Tensor +@@ -2060,7 +2060,7 @@ def _flash_attn_forward( + + # Validate newly added optional cumulative length / padded arrays if provided. + # They are currently only plumbed through for future CK support enabling per-batch padding. +- def _validate_cu(name: str, x: torch.Tensor | None): ++ def _validate_cu(name: str, x: Optional[torch.Tensor]): + if x is None: + return + assert x.dim() == 1, f"{name} must be 1D" From 61acca8dd86dce8b1f62ab7b2ec8eabc60ae16fb Mon Sep 17 00:00:00 2001 From: Mikko Lauri Date: Thu, 27 Aug 2026 09:11:04 +0000 Subject: [PATCH 05/13] bump up xdit and patches --- docker/Dockerfile.ci | 2 +- .../xfuser_aiter_sparge_fp8_asm_backend.patch | 21 +++++++++---------- patches/xfuser_flux.patch | 8 +++---- 3 files changed, 15 insertions(+), 16 deletions(-) diff --git a/docker/Dockerfile.ci b/docker/Dockerfile.ci index 9c7b190..a070061 100644 --- a/docker/Dockerfile.ci +++ b/docker/Dockerfile.ci @@ -41,7 +41,7 @@ ARG AITER_COMMIT="4fa508ef2935110ff99adf2743ea93807dbd9c67" ARG LONGCONTEXTATTENTION_COMMIT="631bdfd" ARG DIFFUSERS_COMMIT="f6e1c4d1e7a6aebaeeaaa0880a82ad036517909e" ARG DISTVAE_COMMIT="6d8025c96b9975d45badae12d8ed9b6422e34e2c" -ARG XDIT_COMMIT="2f74acb12830163d1fe0dd1abb5422f0bff7acf0" +ARG XDIT_COMMIT="8b546dc3af70f1cf0432aa5059c3020c78428aaf" ARG AO_COMMIT="92dcc96162f21928092f8e8e35bac29d3b180e34" ARG ARBITER_COMMIT="252b0301d072f9d6ea895e32f6898fed951cf295" diff --git a/patches/xfuser_aiter_sparge_fp8_asm_backend.patch b/patches/xfuser_aiter_sparge_fp8_asm_backend.patch index 7ec2e67..552b5ce 100644 --- a/patches/xfuser_aiter_sparge_fp8_asm_backend.patch +++ b/patches/xfuser_aiter_sparge_fp8_asm_backend.patch @@ -1,5 +1,5 @@ diff --git a/xfuser/core/distributed/attention_backend.py b/xfuser/core/distributed/attention_backend.py -index 9a08501..982b9a6 100644 +index a7cf679..6c62236 100644 --- a/xfuser/core/distributed/attention_backend.py +++ b/xfuser/core/distributed/attention_backend.py @@ -372,6 +372,10 @@ if env_info["has_aiter"]: @@ -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 ( +@@ -594,6 +598,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 +@@ -1676,6 +1681,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 a66da14..519f7b2 100644 --- a/xfuser/core/distributed/runtime_state.py +++ b/xfuser/core/distributed/runtime_state.py -@@ -225,6 +225,7 @@ class RuntimeState(metaclass=ABCMeta): +@@ -236,6 +236,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): +@@ -348,6 +349,20 @@ class RuntimeState(metaclass=ABCMeta): raise RuntimeError(msg) from None except ImportError: raise RuntimeError(msg) from None @@ -152,10 +151,10 @@ index 341a3c4..fc37433 100644 }) 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 7948d99..febe00e 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({ +@@ -75,6 +75,7 @@ _SPARSE_ATTENTION_BACKENDS = frozenset({ _SPARGE_ATTENTION_BACKENDS = frozenset({ AttentionBackendType.AITER_SPARGE, AttentionBackendType.AITER_SPARGE_V2, 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" From f329623db69115d78fafd1bf1afde0be2f154320 Mon Sep 17 00:00:00 2001 From: Mikko Lauri Date: Thu, 27 Aug 2026 09:11:32 +0000 Subject: [PATCH 06/13] bump up xdit in cuda dockerfile --- docker/Dockerfile.cuda | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docker/Dockerfile.cuda b/docker/Dockerfile.cuda index e8f4256..44e0f98 100644 --- a/docker/Dockerfile.cuda +++ b/docker/Dockerfile.cuda @@ -14,7 +14,7 @@ ARG FLASHATTENTION_COMMIT="fffabc3" ARG LONGCONTEXTATTENTION_COMMIT="631bdfd" ARG DIFFUSERS_COMMIT="f6e1c4d1e7a6aebaeeaaa0880a82ad036517909e" ARG DISTVAE_COMMIT="6d8025c96b9975d45badae12d8ed9b6422e34e2c" -ARG XDIT_COMMIT="2f74acb12830163d1fe0dd1abb5422f0bff7acf0" +ARG XDIT_COMMIT="8b546dc3af70f1cf0432aa5059c3020c78428aaf" ARG SAGEATTENTION_COMMIT="d1a57a5" ARG MSLK_COMMIT="e14925cede32a31ffdd5abf7a9702579d268238c" ARG FLEX_BLOCK_ATTENTION_COMMIT="bc2ec1e" From 7aabb2f754d39ff10ed04eadf0dd3284c3f07f3a Mon Sep 17 00:00:00 2001 From: Mikko Lauri Date: Fri, 28 Aug 2026 10:12:58 +0000 Subject: [PATCH 07/13] bump aiter --- docker/Dockerfile.ci | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docker/Dockerfile.ci b/docker/Dockerfile.ci index a070061..d1f2550 100644 --- a/docker/Dockerfile.ci +++ b/docker/Dockerfile.ci @@ -37,7 +37,7 @@ ARG PYTORCH_VISION_COMMIT="b919bd0c56abbb3c5ca056a3a458af9fd1cabf52" ARG PYTORCH_AUDIO_COMMIT="e3c6ee2b6588b7cd27a84182de74bf12fe043831" ARG PYTORCH_CODEC_COMMIT="6df7fc8e81c9509e86833ca48695609c583e2953" # deps image arguments -ARG AITER_COMMIT="4fa508ef2935110ff99adf2743ea93807dbd9c67" +ARG AITER_COMMIT="4ce24d1e9a9c35af586279bcb525dd936ad49cad" ARG LONGCONTEXTATTENTION_COMMIT="631bdfd" ARG DIFFUSERS_COMMIT="f6e1c4d1e7a6aebaeeaaa0880a82ad036517909e" ARG DISTVAE_COMMIT="6d8025c96b9975d45badae12d8ed9b6422e34e2c" From 34eb2c5a64357c4968282afb133d056aff5c260d Mon Sep 17 00:00:00 2001 From: Mikko Lauri Date: Fri, 28 Aug 2026 10:14:04 +0000 Subject: [PATCH 08/13] bump diffusers --- docker/Dockerfile.ci | 2 +- docker/Dockerfile.cuda | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/docker/Dockerfile.ci b/docker/Dockerfile.ci index d1f2550..7aaf443 100644 --- a/docker/Dockerfile.ci +++ b/docker/Dockerfile.ci @@ -39,7 +39,7 @@ ARG PYTORCH_CODEC_COMMIT="6df7fc8e81c9509e86833ca48695609c583e2953" # deps image arguments ARG AITER_COMMIT="4ce24d1e9a9c35af586279bcb525dd936ad49cad" ARG LONGCONTEXTATTENTION_COMMIT="631bdfd" -ARG DIFFUSERS_COMMIT="f6e1c4d1e7a6aebaeeaaa0880a82ad036517909e" +ARG DIFFUSERS_COMMIT="9f7aee48215c5d01b1d84b52058669710596c7ee" ARG DISTVAE_COMMIT="6d8025c96b9975d45badae12d8ed9b6422e34e2c" ARG XDIT_COMMIT="8b546dc3af70f1cf0432aa5059c3020c78428aaf" ARG AO_COMMIT="92dcc96162f21928092f8e8e35bac29d3b180e34" diff --git a/docker/Dockerfile.cuda b/docker/Dockerfile.cuda index 44e0f98..d00851d 100644 --- a/docker/Dockerfile.cuda +++ b/docker/Dockerfile.cuda @@ -12,7 +12,7 @@ # deps image arguments ARG FLASHATTENTION_COMMIT="fffabc3" ARG LONGCONTEXTATTENTION_COMMIT="631bdfd" -ARG DIFFUSERS_COMMIT="f6e1c4d1e7a6aebaeeaaa0880a82ad036517909e" +ARG DIFFUSERS_COMMIT="9f7aee48215c5d01b1d84b52058669710596c7ee" ARG DISTVAE_COMMIT="6d8025c96b9975d45badae12d8ed9b6422e34e2c" ARG XDIT_COMMIT="8b546dc3af70f1cf0432aa5059c3020c78428aaf" ARG SAGEATTENTION_COMMIT="d1a57a5" From 0b235ec34f77bcea73ce88b8bf6232f26e8bc4bd Mon Sep 17 00:00:00 2001 From: Mikko Lauri Date: Fri, 28 Aug 2026 10:19:14 +0000 Subject: [PATCH 09/13] bump aiter --- docker/Dockerfile.ci | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docker/Dockerfile.ci b/docker/Dockerfile.ci index 9c7b190..b24dcc6 100644 --- a/docker/Dockerfile.ci +++ b/docker/Dockerfile.ci @@ -37,7 +37,7 @@ ARG PYTORCH_VISION_COMMIT="b919bd0c56abbb3c5ca056a3a458af9fd1cabf52" ARG PYTORCH_AUDIO_COMMIT="e3c6ee2b6588b7cd27a84182de74bf12fe043831" ARG PYTORCH_CODEC_COMMIT="6df7fc8e81c9509e86833ca48695609c583e2953" # deps image arguments -ARG AITER_COMMIT="4fa508ef2935110ff99adf2743ea93807dbd9c67" +ARG AITER_COMMIT="4ce24d1e9a9c35af586279bcb525dd936ad49cad" ARG LONGCONTEXTATTENTION_COMMIT="631bdfd" ARG DIFFUSERS_COMMIT="f6e1c4d1e7a6aebaeeaaa0880a82ad036517909e" ARG DISTVAE_COMMIT="6d8025c96b9975d45badae12d8ed9b6422e34e2c" From 13de42ac492151668bea460f67f8a202242839a8 Mon Sep 17 00:00:00 2001 From: Mikko Lauri Date: Mon, 31 Aug 2026 05:52:27 +0000 Subject: [PATCH 10/13] bump transformers --- docker/Dockerfile.ci | 2 +- docker/Dockerfile.cuda | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/docker/Dockerfile.ci b/docker/Dockerfile.ci index 7aaf443..47036f4 100644 --- a/docker/Dockerfile.ci +++ b/docker/Dockerfile.ci @@ -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 d00851d..c00fc23 100644 --- a/docker/Dockerfile.cuda +++ b/docker/Dockerfile.cuda @@ -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 && \ From f9099e7b95ca2e44890d0887c8ad0ad12adc55d9 Mon Sep 17 00:00:00 2001 From: Mikko Lauri Date: Mon, 31 Aug 2026 06:02:07 +0000 Subject: [PATCH 11/13] bump aiter to incl block sparse support --- docker/Dockerfile.ci | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docker/Dockerfile.ci b/docker/Dockerfile.ci index b24dcc6..f4cf265 100644 --- a/docker/Dockerfile.ci +++ b/docker/Dockerfile.ci @@ -37,7 +37,7 @@ ARG PYTORCH_VISION_COMMIT="b919bd0c56abbb3c5ca056a3a458af9fd1cabf52" ARG PYTORCH_AUDIO_COMMIT="e3c6ee2b6588b7cd27a84182de74bf12fe043831" ARG PYTORCH_CODEC_COMMIT="6df7fc8e81c9509e86833ca48695609c583e2953" # deps image arguments -ARG AITER_COMMIT="4ce24d1e9a9c35af586279bcb525dd936ad49cad" +ARG AITER_COMMIT="b8f52325255aa92156c43eb4e580dc07750850c9" ARG LONGCONTEXTATTENTION_COMMIT="631bdfd" ARG DIFFUSERS_COMMIT="f6e1c4d1e7a6aebaeeaaa0880a82ad036517909e" ARG DISTVAE_COMMIT="6d8025c96b9975d45badae12d8ed9b6422e34e2c" From d492022047acb96bce558ceb20f5f34a71e9150b Mon Sep 17 00:00:00 2001 From: Mikko Lauri Date: Mon, 31 Aug 2026 12:44:44 +0000 Subject: [PATCH 12/13] bump xdit to include mhav4 sparse support --- docker/Dockerfile.ci | 2 +- docker/Dockerfile.cuda | 2 +- .../xfuser_aiter_sparge_fp8_asm_backend.patch | 26 +++++++++---------- 3 files changed, 15 insertions(+), 15 deletions(-) diff --git a/docker/Dockerfile.ci b/docker/Dockerfile.ci index 279e538..3a2868a 100644 --- a/docker/Dockerfile.ci +++ b/docker/Dockerfile.ci @@ -41,7 +41,7 @@ ARG AITER_COMMIT="0c2b0f77b2ff6d13c677d12466abf87299f8b260" ARG LONGCONTEXTATTENTION_COMMIT="631bdfd" ARG DIFFUSERS_COMMIT="9f7aee48215c5d01b1d84b52058669710596c7ee" ARG DISTVAE_COMMIT="6d8025c96b9975d45badae12d8ed9b6422e34e2c" -ARG XDIT_COMMIT="8b546dc3af70f1cf0432aa5059c3020c78428aaf" +ARG XDIT_COMMIT="3aa7b34ec3f1d4f86f0befc5f09384333067846c" ARG AO_COMMIT="92dcc96162f21928092f8e8e35bac29d3b180e34" ARG ARBITER_COMMIT="252b0301d072f9d6ea895e32f6898fed951cf295" diff --git a/docker/Dockerfile.cuda b/docker/Dockerfile.cuda index c00fc23..14a910c 100644 --- a/docker/Dockerfile.cuda +++ b/docker/Dockerfile.cuda @@ -14,7 +14,7 @@ ARG FLASHATTENTION_COMMIT="fffabc3" ARG LONGCONTEXTATTENTION_COMMIT="631bdfd" ARG DIFFUSERS_COMMIT="9f7aee48215c5d01b1d84b52058669710596c7ee" ARG DISTVAE_COMMIT="6d8025c96b9975d45badae12d8ed9b6422e34e2c" -ARG XDIT_COMMIT="8b546dc3af70f1cf0432aa5059c3020c78428aaf" +ARG XDIT_COMMIT="3aa7b34ec3f1d4f86f0befc5f09384333067846c" ARG SAGEATTENTION_COMMIT="d1a57a5" ARG MSLK_COMMIT="e14925cede32a31ffdd5abf7a9702579d268238c" ARG FLEX_BLOCK_ATTENTION_COMMIT="bc2ec1e" diff --git a/patches/xfuser_aiter_sparge_fp8_asm_backend.patch b/patches/xfuser_aiter_sparge_fp8_asm_backend.patch index 552b5ce..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 a7cf679..6c62236 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. @@ -13,7 +13,7 @@ index a7cf679..6c62236 100644 try: from aiter.ops.mha_v4 import ( -@@ -594,6 +598,7 @@ class AttentionBackendType(Enum): +@@ -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,7 +21,7 @@ index a7cf679..6c62236 100644 AITER_VSA = "AITER VSA CK" FLEX_BLOCK_SPARGE = "Flex Block Sparge" AITER_FLYDSL = "AITER FlyDSL" -@@ -1676,6 +1681,83 @@ 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 @@ -106,10 +106,10 @@ index a7cf679..6c62236 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 a66da14..519f7b2 100644 +index c7d574a..631dd50 100644 --- a/xfuser/core/distributed/runtime_state.py +++ b/xfuser/core/distributed/runtime_state.py -@@ -236,6 +236,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, @@ -117,7 +117,7 @@ index a66da14..519f7b2 100644 AttentionBackendType.AITER_VSA, AttentionBackendType.AITER_FLYDSL, AttentionBackendType.AITER_FLYDSL_FP8, -@@ -348,6 +349,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 @@ -139,26 +139,26 @@ index a66da14..519f7b2 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 7948d99..febe00e 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 -@@ -75,6 +75,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 From 72d02191a93874c6afd8f5568d4ec6d7f17fa037 Mon Sep 17 00:00:00 2001 From: Mikko Lauri Date: Tue, 1 Sep 2026 12:58:43 +0000 Subject: [PATCH 13/13] bump distvae --- docker/Dockerfile.ci | 2 +- docker/Dockerfile.cuda | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/docker/Dockerfile.ci b/docker/Dockerfile.ci index 3a2868a..09050bd 100644 --- a/docker/Dockerfile.ci +++ b/docker/Dockerfile.ci @@ -40,7 +40,7 @@ ARG PYTORCH_CODEC_COMMIT="6df7fc8e81c9509e86833ca48695609c583e2953" ARG AITER_COMMIT="0c2b0f77b2ff6d13c677d12466abf87299f8b260" ARG LONGCONTEXTATTENTION_COMMIT="631bdfd" ARG DIFFUSERS_COMMIT="9f7aee48215c5d01b1d84b52058669710596c7ee" -ARG DISTVAE_COMMIT="6d8025c96b9975d45badae12d8ed9b6422e34e2c" +ARG DISTVAE_COMMIT="dce484185dc11d7406076117771e36acd796aa6a" ARG XDIT_COMMIT="3aa7b34ec3f1d4f86f0befc5f09384333067846c" ARG AO_COMMIT="92dcc96162f21928092f8e8e35bac29d3b180e34" ARG ARBITER_COMMIT="252b0301d072f9d6ea895e32f6898fed951cf295" diff --git a/docker/Dockerfile.cuda b/docker/Dockerfile.cuda index 14a910c..7070f4d 100644 --- a/docker/Dockerfile.cuda +++ b/docker/Dockerfile.cuda @@ -13,7 +13,7 @@ ARG FLASHATTENTION_COMMIT="fffabc3" ARG LONGCONTEXTATTENTION_COMMIT="631bdfd" ARG DIFFUSERS_COMMIT="9f7aee48215c5d01b1d84b52058669710596c7ee" -ARG DISTVAE_COMMIT="6d8025c96b9975d45badae12d8ed9b6422e34e2c" +ARG DISTVAE_COMMIT="dce484185dc11d7406076117771e36acd796aa6a" ARG XDIT_COMMIT="3aa7b34ec3f1d4f86f0befc5f09384333067846c" ARG SAGEATTENTION_COMMIT="d1a57a5" ARG MSLK_COMMIT="e14925cede32a31ffdd5abf7a9702579d268238c"