Generic sol_h3_vae package and sol_h3_vae.safetensors; GPU-specific FA4 guard moves to Sol-H3 in ComfyUI

README.md

@@ -29,10 +29,12 @@ as Sol-H3 runs it.
 
 ## Usage
 
-Requirements: PyTorch with CUDA, Triton, [Diffusers](https://github.com/huggingface/diffusers)
-with MiniMax-H3 support, [FlashAttention-4](https://github.com/Dao-AILab/flash-attention)
-(`flash_attn.cute`) and [QuACK](https://github.com/Dao-AILab/quack). Tested on RTX 5090 with
-PyTorch 2.12 (CUDA 13), Triton 3.7, Diffusers 0.40, FlashAttention-4 4.0.0b19 and QuACK 0.6.3.
+Requirements: PyTorch with CUDA and [Diffusers](https://github.com/huggingface/diffusers) with
+MiniMax-H3 support. The fast engine (default) also needs Triton,
+[FlashAttention-4](https://github.com/Dao-AILab/flash-attention) (`flash_attn.cute`) and
+[QuACK](https://github.com/Dao-AILab/quack); `engine="ref"` runs in plain PyTorch without them.
+Tested on RTX 5090 with PyTorch 2.12 (CUDA 13), Triton 3.7, Diffusers 0.40, FlashAttention-4
+4.0.0b19 and QuACK 0.6.3.
 
 ```python
 import sys
@@ -40,10 +42,10 @@ from huggingface_hub import hf_hub_download, snapshot_download
 
 repo = snapshot_download("Efficient-Large-Model/Sol-H3-VAE")
 sys.path.insert(0, f"{repo}/code")
-from spark_decoder import SparkFastH3VAEDecoder
+from sol_h3_vae import SolH3VAEDecoder
 
 vae_config = hf_hub_download("MiniMaxAI/MiniMax-H3", "vae/config.json")
-decoder = SparkFastH3VAEDecoder(vae_config, f"{repo}/fast_h3_vae_noncausal_24l.safetensors")
+decoder = SolH3VAEDecoder(vae_config, f"{repo}/sol_h3_vae.safetensors")
 
 # z: normalized H3 video latents, (1, T, 24, H/16, W/16); T = 0 starts a native 5-latent block
 pixels = decoder.decode(z)          # (1, 3, F, H, W), the H3 VAE decoder's normalized pixel space
@@ -52,7 +54,12 @@ rgb = decoder.decode_rgb_uint8(z)   # (F, H, W, 3) uint8 on the GPU
 
 The decoder needs only the H3 VAE `config.json`, not the H3 VAE weights. The first call for each
 latent shape compiles kernels and captures a CUDA graph (about 25 s on RTX 5090); later calls
-reuse it. Pass `engine="ref"` for the plain PyTorch decoder (slower).
+reuse it.
+
+On SM120 GPUs (RTX 5090, DGX Spark), FlashAttention-4 4.0.0b19 to 4.0.0b22 has a varlen bug that
+makes the fast engine hit an illegal memory access. Apply the guard from
+[Sol-H3 in ComfyUI](https://github.com/NVlabs/Sana/blob/sol-engine/models/minimax_h3/Sol-H3-ComfyUI/sol_h3_engine/fa4_sm120.py)
+before the first decode, or use `engine="ref"`.
 
 ## Limitations
 
@@ -64,7 +71,6 @@ reuse it. Pass `engine="ref"` for the plain PyTorch decoder (slower).
 
 ## License
 
-The code is released under Apache-2.0; `code/_spark_flash_fwd_guarded.py` is adapted from
-FlashAttention (BSD-3-Clause). The weights are distilled from the MiniMax-H3 VAE and are
+The code is released under Apache-2.0. The weights are distilled from the MiniMax-H3 VAE and are
 subject to the
 [MiniMax-H3 license](https://huggingface.co/MiniMaxAI/MiniMax-H3/blob/main/LICENSE).

code/LICENSE.flash-attention

@@ -1,29 +0,0 @@
-BSD 3-Clause License
-
-Copyright (c) 2022, the respective contributors, as shown by the AUTHORS file.
-All rights reserved.
-
-Redistribution and use in source and binary forms, with or without
-modification, are permitted provided that the following conditions are met:
-
-* Redistributions of source code must retain the above copyright notice, this
-  list of conditions and the following disclaimer.
-
-* Redistributions in binary form must reproduce the above copyright notice,
-  this list of conditions and the following disclaimer in the documentation
-  and/or other materials provided with the distribution.
-
-* Neither the name of the copyright holder nor the names of its
-  contributors may be used to endorse or promote products derived from
-  this software without specific prior written permission.
-
-THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
-AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
-IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
-DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
-FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
-DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
-SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
-CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
-OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
-OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

code/_spark_flash_fwd_guarded.py

@@ -1,1246 +0,0 @@
-# Copyright (c) 2025, Jay Shah, Ganesh Bikshandi, Ying Zhang, Vijay Thakkar, Pradeep Ramani, Tri Dao.
-# A reimplementation of
-# https://github.com/Dao-AILab/flash-attention/blob/main/hopper/flash_fwd_kernel_sm80.h
-# and https://github.com/Dao-AILab/flash-attention/blob/main/hopper/flash_fwd_kernel_sm90.h
-# from Cutlass C++ to Cute-DSL.
-# Built on Cute-DSL example: https://github.com/NVIDIA/cutlass/blob/main/examples/python/CuTeDSL/ampere/flash_attention_v2.py
-
-import math
-from types import SimpleNamespace
-from typing import Type, Callable, Optional
-from functools import partial
-
-import cuda.bindings.driver as cuda
-
-import cutlass
-import cutlass.cute as cute
-from cutlass import Float32, Int32, const_expr
-from cutlass.cute.nvgpu import cpasync, warp
-import cutlass.utils as utils_basic
-from cutlass.base_dsl.arch import Arch
-from cutlass.cutlass_dsl import BaseDSL
-
-from quack import copy_utils
-from quack import layout_utils
-
-from flash_attn.cute import ampere_helpers as sm80_utils
-from flash_attn.cute.cute_dsl_utils import assume_tensor_aligned
-from flash_attn.cute import utils
-from flash_attn.cute.mask import AttentionMask
-from flash_attn.cute.softmax import Softmax, apply_score_mod_inner
-from flash_attn.cute.seqlen_info import SeqlenInfoQK
-from flash_attn.cute.block_info import BlockInfo
-from flash_attn.cute.pack_gqa import PackGQA, pack_gqa_layout
-from flash_attn.cute.named_barrier import NamedBarrierFwd
-from flash_attn.cute.block_sparsity import BlockSparseTensors
-from flash_attn.cute.tile_scheduler import SingleTileScheduler, SingleTileVarlenScheduler, TileSchedulerArguments
-from flash_attn.cute.utils import AuxData
-
-
-class FlashAttentionForwardBase:
-
-    def __init__(
-        self,
-        dtype: Type[cutlass.Numeric],
-        head_dim: int,
-        head_dim_v: Optional[int] = None,
-        qhead_per_kvhead: int = 1,
-        is_causal: bool = False,
-        is_local: bool = False,
-        pack_gqa: bool = True,
-        tile_m: int = 128,
-        tile_n: int = 128,
-        num_stages: int = 1,
-        num_threads: int = 128,
-        Q_in_regs: bool = False,
-        score_mod: Optional[cutlass.Constexpr] = None,
-        mask_mod: Optional[cutlass.Constexpr] = None,
-        has_aux_tensors: bool = False,
-        q_subtile_factor: int = 1,
-    ):
-        """Initializes the configuration for a flash attention kernel.
-
-        All contiguous dimensions must be at least 16 bytes aligned, which means that the head dimension
-        should be a multiple of 8.
-
-        :param head_dim: head dimension
-        :type head_dim: int
-        :param tile_m: m block size
-        :type tile_m: int
-        :param tile_n: n block size
-        :type tile_n: int
-        :param num_threads: number of threads
-        :type num_threads: int
-        :param is_causal: is causal
-        :param score_mod: A callable that takes the attention scores and applies a modification.
-            Callable signature: ``score_mod(scores, batch_idx, head_idx, q_idx, kv_idx, aux_tensors) -> Any``
-        :param mask_mod: A callable that takes the attention scores and returns a boolean representing whether that score should be masked.
-            Callable signature: ``mask_mod(batch_idx, head_idx, q_idx, kv_idx, aux_tensors) -> Boolean``
-        """
-        self.dtype = dtype
-        # padding head_dim to a multiple of 16 as k_block_size
-        hdim_multiple_of = 16
-        self.tile_hdim = int(math.ceil(head_dim / hdim_multiple_of) * hdim_multiple_of)
-        head_dim_v = head_dim_v if head_dim_v is not None else head_dim
-        self.same_hdim_kv = head_dim == head_dim_v
-        self.tile_hdimv = int(math.ceil(head_dim_v / hdim_multiple_of) * hdim_multiple_of)
-        # Can save registers (and hence be faster) if we don't have to check hdim predication
-        self.check_hdim_oob = head_dim != self.tile_hdim
-        self.check_hdim_v_oob = head_dim_v != self.tile_hdimv
-        self.qhead_per_kvhead = qhead_per_kvhead
-        self.is_causal = is_causal
-        self.is_local = is_local
-        self.pack_gqa = pack_gqa
-        self.tile_m = tile_m
-        self.tile_n = tile_n
-        self.num_threads = num_threads
-        self.num_stages = num_stages
-        self.q_subtile_factor = q_subtile_factor
-        self.Q_in_regs = Q_in_regs
-        self.score_mod = score_mod
-        self.mask_mod = mask_mod
-        self.qk_acc_dtype = Float32
-        self.score_vec_size: cutlass.Constexpr = getattr(
-            score_mod, "__vec_size__", 1 if cutlass.const_expr(has_aux_tensors) else 2
-        )
-        if self.score_vec_size > 2:
-            raise ValueError(
-                f"score_mod vec_size {self.score_vec_size} not supported on Sm80/90/120 "
-                "due to accumulator thread ownership pattern."
-            )
-        self.mask_vec_size: cutlass.Constexpr = getattr(mask_mod, "__vec_size__", 1)
-        if self.mask_vec_size > 1:
-            raise ValueError(
-                f"mask_mod vec_size {self.mask_vec_size} not supported on Sm80/90/120 "
-                "due to accumulator thread ownership pattern."
-            )
-        self.arch = BaseDSL._get_dsl().get_arch_enum()
-
-    @staticmethod
-    def can_implement(
-        dtype,
-        head_dim,
-        head_dim_v,
-        tile_m,
-        tile_n,
-        num_stages,
-        num_threads,
-        is_causal,
-        Q_in_regs=False,
-    ) -> bool:
-        """Check if the kernel can be implemented with the given parameters.
-
-        :param dtype: data type
-        :type dtype: cutlass.Numeric
-        :param head_dim: head dimension
-        :type head_dim: int
-        :param tile_m: m block size
-        :type tile_m: int
-        :param tile_n: n block size
-        :type tile_n: int
-        :param num_threads: number of threads
-        :type num_threads: int
-        :param is_causal: is causal
-        :type is_causal: bool
-
-        :return: True if the kernel can be implemented, False otherwise
-        :rtype: bool
-        """
-        if dtype not in [cutlass.Float16, cutlass.BFloat16]:
-            return False
-        if head_dim % 8 != 0:
-            return False
-        if head_dim_v % 8 != 0:
-            return False
-        if tile_n % 16 != 0:
-            return False
-        if num_threads % 32 != 0:
-            return False
-        # Check if block size setting is out of shared memory capacity
-        # Shared memory usage: Q tile + (K tile + V tile) where K and V use the same tile size
-        smem_usage_Q = tile_m * head_dim * 2
-        smem_usage_K = tile_n * head_dim * num_stages * 2
-        smem_usage_V = tile_n * head_dim_v * num_stages * 2
-        smem_usage_QV = (
-            (smem_usage_Q + smem_usage_V) if not Q_in_regs else max(smem_usage_Q, smem_usage_V)
-        )
-        smem_usage = smem_usage_QV + smem_usage_K
-        # TODO: sm86 and sm89
-        smem_capacity = utils_basic.get_smem_capacity_in_bytes("sm_80")
-        if smem_usage > smem_capacity:
-            return False
-        # Check if twice the block size is divisible by the number of threads
-        if (tile_m * 2) % num_threads != 0:
-            return False
-        return True
-
-    def _check_type(
-        self,
-        mQ_type: Type[cutlass.Numeric],
-        mK_type: Type[cutlass.Numeric],
-        mV_type: Type[cutlass.Numeric],
-        mO_type: Type[cutlass.Numeric],
-        mLSE_type: Type[cutlass.Numeric] | None,
-        mCuSeqlensQ_type: Type[cutlass.Numeric] | None,
-        mCuSeqlensK_type: Type[cutlass.Numeric] | None,
-        mSeqUsedQ_type: Type[cutlass.Numeric] | None,
-        mSeqUsedK_type: Type[cutlass.Numeric] | None,
-    ):
-        # Get the data type and check if it is fp16 or bf16
-        if const_expr(not (mQ_type == mK_type == mV_type == mO_type)):
-            raise TypeError("All tensors must have the same data type")
-        if const_expr(mQ_type not in [cutlass.Float16, cutlass.BFloat16]):
-            raise TypeError("Only Float16 or BFloat16 is supported")
-        if const_expr(mLSE_type not in [None, Float32]):
-            raise TypeError("LSE tensor must be Float32")
-        if const_expr(mCuSeqlensQ_type not in [None, Int32]):
-            raise TypeError("cu_seqlens_q tensor must be Int32")
-        if const_expr(mCuSeqlensK_type not in [None, Int32]):
-            raise TypeError("cu_seqlens_k tensor must be Int32")
-        if const_expr(mSeqUsedQ_type not in [None, Int32]):
-            raise TypeError("seqused_q tensor must be Int32")
-        if const_expr(mSeqUsedK_type not in [None, Int32]):
-            raise TypeError("seqused_k tensor must be Int32")
-        assert mQ_type == self.dtype
-
-    def _setup_attributes(self):
-        # ///////////////////////////////////////////////////////////////////////////////
-        # Shared memory layout: Q/K/V
-        # ///////////////////////////////////////////////////////////////////////////////
-        sQ_layout_atom, sK_layout_atom, sV_layout_atom, sO_layout_atom, sP_layout_atom = (
-            self._get_smem_layout_atom()
-        )
-        self.sQ_layout = cute.tile_to_shape(
-            sQ_layout_atom,
-            (self.tile_m, self.tile_hdim),
-            (0, 1),
-        )
-        self.sK_layout = cute.tile_to_shape(
-            sK_layout_atom,
-            (self.tile_n, self.tile_hdim, self.num_stages),
-            (0, 1, 2),
-        )
-        self.sV_layout = cute.tile_to_shape(
-            sV_layout_atom,
-            (self.tile_n, self.tile_hdimv, self.num_stages),
-            (0, 1, 2),
-        )
-        self.sO_layout = cute.tile_to_shape(
-            sO_layout_atom,
-            (self.tile_m, self.tile_hdimv),
-            (0, 1),
-        )
-        if const_expr(sP_layout_atom is not None):
-            self.sP_layout = cute.tile_to_shape(
-                sP_layout_atom,
-                (self.tile_m, self.tile_n),
-                (0, 1),
-            )
-        else:
-            self.sP_layout = None
-
-        # ///////////////////////////////////////////////////////////////////////////////
-        # GMEM Tiled copy:
-        # ///////////////////////////////////////////////////////////////////////////////
-        # Thread layouts for copies
-        universal_copy_bits = 128
-        async_copy_elems = universal_copy_bits // self.dtype.width
-        # atom_async_copy: async copy atom for QKV load
-        atom_async_copy = cute.make_copy_atom(
-            cpasync.CopyG2SOp(cache_mode=cpasync.LoadCacheMode.GLOBAL),
-            self.dtype,
-            num_bits_per_copy=universal_copy_bits,
-        )
-        # atom_universal_copy: universal copy atom for O store
-        atom_universal_copy = cute.make_copy_atom(
-            cute.nvgpu.CopyUniversalOp(),
-            self.dtype,
-            num_bits_per_copy=universal_copy_bits,
-        )
-        # tQ_layout and tK_layout: thread layout for QK load
-        tQK_shape_dim_1 = sQ_layout_atom.outer.shape[1] // async_copy_elems
-        assert self.num_Q_load_threads % tQK_shape_dim_1 == 0, (
-            "num_threads must be divisible by tQK_shape_dim_1"
-        )
-        assert self.num_producer_threads % tQK_shape_dim_1 == 0, (
-            "num_threads must be divisible by tQK_shape_dim_1"
-        )
-        tQ_layout = cute.make_ordered_layout(
-            (self.num_Q_load_threads // tQK_shape_dim_1, tQK_shape_dim_1),
-            order=(1, 0),
-        )
-        tK_layout = cute.make_ordered_layout(
-            (self.num_producer_threads // tQK_shape_dim_1, tQK_shape_dim_1),
-            order=(1, 0),
-        )
-        # So that we don't have to check if we overshoot kBlockM when we load Q
-        assert self.tile_m % tQ_layout.shape[0] == 0
-        tV_shape_dim_1 = sV_layout_atom.outer.shape[1] // async_copy_elems
-        tV_layout = cute.make_ordered_layout(
-            (self.num_producer_threads // tV_shape_dim_1, tV_shape_dim_1),
-            order=(1, 0),
-        )
-        # TODO: need a different layout for O if O dtype is not the same as V dtype
-        # tO_layout: thread layout for O store
-        tO_layout = cute.make_ordered_layout(
-            (self.num_epilogue_threads // tV_shape_dim_1, tV_shape_dim_1),
-            order=(1, 0),
-        )
-        # So that we don't have to check if we overshoot kBlockM when we store O
-        assert self.tile_m % tO_layout.shape[0] == 0
-
-        # Value layouts for copies
-        vQKV_layout = cute.make_layout((1, async_copy_elems))
-        vO_layout = vQKV_layout
-
-        self.gmem_tiled_copy_Q = cute.make_tiled_copy_tv(atom_async_copy, tQ_layout, vQKV_layout)
-        self.gmem_tiled_copy_K = cute.make_tiled_copy_tv(atom_async_copy, tK_layout, vQKV_layout)
-        self.gmem_tiled_copy_V = cute.make_tiled_copy_tv(atom_async_copy, tV_layout, vQKV_layout)
-        # gmem_tiled_copy_O: tiled copy for O store
-        self.gmem_tiled_copy_O = cute.make_tiled_copy_tv(atom_universal_copy, tO_layout, vO_layout)
-
-    def _get_smem_layout_atom(self):
-        raise NotImplementedError()
-
-    def _get_tiled_mma(self):
-        raise NotImplementedError()
-
-    def _get_shared_storage_cls(self):
-        raise NotImplementedError()
-
-    @cute.jit
-    def __call__(
-        self,
-        mQ: cute.Tensor,
-        mK: cute.Tensor,
-        mV: cute.Tensor,
-        mO: cute.Tensor,
-        mLSE: Optional[cute.Tensor],
-        softmax_scale: Float32,
-        # Always keep stream as the last parameter (EnvStream: obtained implicitly via TVM FFI).
-        stream: cuda.CUstream = None,
-    ):
-        """Configures and launches the flash attention kernel.
-
-        mQ/mK/mV/mO has same data types(supports fp16 and bf16) and same layout:
-        (batch_size, seqlen_q, num_head, head_dim):(_, _, _, 1)
-        """
-        raise NotImplementedError()
-
-    @cute.jit
-    def epilogue(
-        self,
-        acc_O: cute.Tensor,
-        lse: cute.Tensor,
-        mO: cute.Tensor,
-        mLSE: Optional[cute.Tensor],
-        sO: cute.Tensor,
-        seqlen: SeqlenInfoQK,
-        gmem_tiled_copy_O: cute.TiledCopy,
-        tma_atom_O: Optional[cute.CopyAtom],
-        tiled_mma: cute.TiledMma,
-        tidx: Int32,
-        m_block: Int32,
-        head_idx: Int32,
-        batch_idx: Int32,
-    ):
-        # store acc_O
-        rO = cute.make_fragment_like(acc_O, self.dtype)
-        rO.store(acc_O.load().to(self.dtype))
-        # Make sure all threads have finished reading V
-        cute.arch.barrier(
-            barrier_id=int(NamedBarrierFwd.Epilogue), number_of_threads=self.num_epilogue_threads
-        )
-        smem_copy_atom_O = utils.get_smem_store_atom(self.arch.major * 10 + self.arch.minor, self.dtype)
-        smem_thr_copy_O = cute.make_tiled_copy_C(smem_copy_atom_O, tiled_mma).get_slice(tidx)
-        taccOrO = smem_thr_copy_O.retile(rO)
-        taccOsO = smem_thr_copy_O.partition_D(sO)
-        # taccOsO = copy_utils.partition_D_position_independent(smem_thr_copy_O, sO)
-        # copy acc O from rmem to smem with the smem copy atom
-        cute.copy(smem_copy_atom_O, taccOrO, taccOsO)
-
-        cO = cute.make_identity_tensor((self.tile_m, self.tile_hdimv))
-        pack_gqa = PackGQA(
-            self.tile_m, self.tile_hdimv, self.check_hdim_v_oob, self.qhead_per_kvhead
-        )
-
-        # Write LSE from rmem -> gmem
-        if const_expr(mLSE is not None):
-            mLSE_cur = seqlen.offset_batch_Q(mLSE, batch_idx, dim=2)[None, head_idx]
-            if const_expr(not self.pack_gqa):
-                gLSE = cute.local_tile(mLSE_cur, (self.tile_m,), (m_block,))
-                gLSE_expanded_layout = cute.append(
-                    gLSE.layout, cute.make_layout((self.tile_hdimv,), stride=(0,))
-                )
-                gLSE_expanded = cute.make_tensor(gLSE.iterator, gLSE_expanded_layout)
-                thr_mma = tiled_mma.get_slice(tidx)
-                taccOgLSE = layout_utils.reshape_acc_to_mn(thr_mma.partition_C(gLSE_expanded))
-                assert cute.size(taccOgLSE, mode=[0]) == cute.size(lse)
-                taccOcO = layout_utils.reshape_acc_to_mn(thr_mma.partition_C(cO))
-                t0accOcO = layout_utils.reshape_acc_to_mn(thr_mma.get_slice(0).partition_C(cO))
-                # Only the thread corresponding to column 0 writes out the lse to gmem
-                if taccOcO[0][1] == 0:
-                    for m in cutlass.range(cute.size(taccOgLSE.shape[1]), unroll_full=True):
-                        if (
-                            t0accOcO[m, 0][0]
-                            < seqlen.seqlen_q - m_block * self.tile_m - taccOcO[0][0]
-                        ):
-                            taccOgLSE[m, 0] = lse[m]
-            else:
-                pack_gqa.store_LSE(mLSE_cur, lse, tiled_mma, tidx, m_block, seqlen.seqlen_q)
-
-        ragged = self.use_tma_O and (seqlen.has_cu_seqlens_q or seqlen.has_seqused_q)
-        mO_cur = seqlen.offset_batch_Q(mO, batch_idx, dim=3, ragged=ragged)[None, None, head_idx]
-        # thr_mma = tiled_mma.get_slice(tidx)
-        # taccOgO = thr_mma.partition_C(gO)
-        # cute.autovec_copy(rO, taccOgO)
-        # sync to make sure all smem stores are done
-        if const_expr(self.use_tma_O):
-            # ensure smem writes are visible to TMA
-            cute.arch.fence_view_async_shared()
-            cute.arch.barrier_arrive(
-                barrier_id=int(NamedBarrierFwd.Epilogue),
-                number_of_threads=self.num_epilogue_threads + cute.arch.WARP_SIZE,
-            )
-            gO = cute.local_tile(mO_cur, (self.tile_m, self.tile_hdimv), (m_block, 0))
-            store_O, _, _ = copy_utils.tma_get_copy_fn(
-                tma_atom_O, 0, cute.make_layout(1), sO, gO, single_stage=True
-            )
-            warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx())
-            if warp_idx == 4:
-                cute.arch.barrier(
-                    barrier_id=int(NamedBarrierFwd.Epilogue),
-                    number_of_threads=self.num_epilogue_threads + cute.arch.WARP_SIZE,
-                )
-                store_O()
-                cute.arch.cp_async_bulk_commit_group()
-                cute.arch.cp_async_bulk_wait_group(0, read=True)
-        else:
-            cute.arch.barrier(
-                barrier_id=int(NamedBarrierFwd.Epilogue),
-                number_of_threads=self.num_epilogue_threads,
-            )
-            gmem_thr_copy_O = gmem_tiled_copy_O.get_slice(tidx)
-            tOsO = gmem_thr_copy_O.partition_S(sO)
-            tOrO = cute.make_fragment_like(tOsO, self.dtype)
-            # load acc O from smem to rmem for wider vectorization
-            cute.autovec_copy(tOsO, tOrO)
-            if const_expr(not self.pack_gqa):
-                gO = cute.local_tile(mO_cur, (self.tile_m, self.tile_hdimv), (m_block, 0))
-                tOgO = gmem_thr_copy_O.partition_D(gO)
-                tOcO = gmem_thr_copy_O.partition_S(cO)
-                t0OcO = gmem_tiled_copy_O.get_slice(0).partition_S(cO)
-                tOpO = utils.predicate_k(tOcO, limit=mO.shape[1])
-                # copy acc O from rmem to gmem
-                for rest_m in cutlass.range_constexpr(cute.size(tOrO.shape[1])):
-                    if (
-                        t0OcO[0, rest_m, 0][0]
-                        < seqlen.seqlen_q - m_block * self.tile_m - tOcO[0][0]
-                    ):
-                        cute.copy(
-                            gmem_tiled_copy_O,
-                            tOrO[None, rest_m, None],
-                            tOgO[None, rest_m, None],
-                            pred=tOpO[None, rest_m, None]
-                            if const_expr(self.check_hdim_v_oob)
-                            else None,
-                        )
-            else:
-                pack_gqa.store_O(mO_cur, tOrO, gmem_tiled_copy_O, tidx, m_block, seqlen.seqlen_q)
-
-    @cute.jit
-    def advance_pipeline(self, pipeline_index):
-        return pipeline_index + 1 if pipeline_index < self.num_stages - 1 else 0
-
-    @cute.jit
-    def load_Q(
-        self,
-        gmem_thr_copy: cute.TiledCopy,
-        gQ: cute.Tensor,
-        sQ: cute.Tensor,
-        block: Int32,
-        seqlen: Int32,
-        headdim: Int32,
-    ):
-        tQsQ, tQgQ = gmem_thr_copy.partition_D(sQ), gmem_thr_copy.partition_S(gQ)
-        cQ = cute.make_identity_tensor((self.tile_m, self.tile_hdim))
-        tQcQ = gmem_thr_copy.partition_S(cQ)
-        t0QcQ = gmem_thr_copy.get_slice(0).partition_S(cQ)
-        tQpQ = utils.predicate_k(tQcQ, limit=headdim)
-        for m in cutlass.range_constexpr(cute.size(tQsQ.shape[1])):
-            # Instead of using tQcQ, we using t0QcQ and subtract the offset from the limit
-            # (seqlen - block * kBlockM). This is because the entries of t0QcQ are known at compile time.
-            if t0QcQ[0, m, 0][0] < seqlen - block * self.tile_m - tQcQ[0][0]:
-                cute.copy(
-                    gmem_thr_copy,
-                    tQgQ[None, m, None],
-                    tQsQ[None, m, None],
-                    pred=tQpQ[None, m, None] if const_expr(self.check_hdim_oob) else None,
-                )
-            # We don't need to clear the sQ smem tiles since we'll only write out the valid outputs
-
-    @cute.jit
-    def load_K(
-        self,
-        gmem_tiled_copy: cute.TiledCopy,
-        tKgK: cute.Tensor,
-        tKsK: cute.Tensor,
-        tKcK: cute.Tensor,
-        t0KcK: cute.Tensor,
-        tKpK: cute.Tensor,
-        block: Int32,
-        smem_pipe_write: Int32,
-        seqlen: Int32,
-        need_predicates: cutlass.Constexpr,
-    ):
-        # Do we need to check if we overshoot kBlockN when we load K?
-        is_even_n_smem_k = self.tile_n % gmem_tiled_copy.tiler_mn[0].shape == 0
-        if const_expr(need_predicates or not is_even_n_smem_k):
-            # Instead of using tKcK, we using t0KcK and subtract the offset from the limit
-            # (seqlen - block * kBlockN). This is because the entries of t0KcK are known at compile time.
-            if const_expr(is_even_n_smem_k):
-                seqlen_limit = seqlen - block * self.tile_n
-            else:
-                if const_expr(not need_predicates):
-                    seqlen_limit = self.tile_n
-                else:
-                    seqlen_limit = cutlass.min(seqlen - block * self.tile_n, self.tile_n)
-            seqlen_limit -= tKcK[0][0]
-            for n in cutlass.range_constexpr(cute.size(tKsK.shape[1])):
-                if t0KcK[0, n, 0][0] < seqlen_limit:
-                    cute.copy(
-                        gmem_tiled_copy,
-                        tKgK[None, n, None, block],
-                        tKsK[
-                            None, n, None, smem_pipe_write if const_expr(self.num_stages > 1) else 0
-                        ],
-                        pred=tKpK[None, n, None] if const_expr(self.check_hdim_oob) else None,
-                    )
-                # We don't need to clear the sK smem tiles since we'll mask out the scores anyway.
-        else:
-            cute.copy(
-                gmem_tiled_copy,
-                tKgK[None, None, None, block],
-                tKsK[None, None, None, smem_pipe_write if const_expr(self.num_stages > 1) else 0],
-                pred=tKpK if const_expr(self.check_hdim_oob) else None,
-            )
-
-    @cute.jit
-    def load_V(
-        self,
-        gmem_tiled_copy: cute.TiledCopy,
-        tVgV: cute.Tensor,
-        tVsV: cute.Tensor,
-        tVcV: cute.Tensor,
-        t0VcV: cute.Tensor,
-        tVpV: cute.Tensor,
-        block: Int32,
-        smem_pipe_write: Int32,
-        seqlen: Int32,
-        need_predicates: cutlass.Constexpr,
-    ):
-        # Do we need to check if we overshoot kBlockN when we load V?
-        is_even_n_smem_v = self.tile_n % gmem_tiled_copy.tiler_mn[0].shape == 0
-        if const_expr(need_predicates or not is_even_n_smem_v):
-            for n in cutlass.range_constexpr(cute.size(tVsV.shape[1])):
-                # If kBlockN doesn't evenly divide the tiled copy, only the last `n` needs to be checked
-                if (
-                    is_even_n_smem_v
-                    or n < cute.size(tVsV.shape[1]) - 1
-                    or tVcV[0, n, 0][0] < self.tile_n
-                ):
-                    predicate = tVpV[None, n, None] if const_expr(self.check_hdim_v_oob) else None
-                    if const_expr(need_predicates):
-                        seqlen_limit = seqlen - block * self.tile_n - tVcV[0][0]
-                        predicate_n = t0VcV[0, n, 0][0] < seqlen_limit
-                        predicate = cute.make_fragment_like(tVpV[None, 0, None])
-                        for k in cutlass.range_constexpr(cute.size(predicate.shape[1])):
-                            for i in cutlass.range_constexpr(cute.size(predicate.shape[0])):
-                                predicate[i, k] = (
-                                    tVpV[i, n, k] if const_expr(self.check_hdim_v_oob) else True
-                                ) and predicate_n
-                    cute.copy(
-                        gmem_tiled_copy,
-                        tVgV[None, n, None, block],
-                        tVsV[
-                            None, n, None, smem_pipe_write if const_expr(self.num_stages > 1) else 0
-                        ],
-                        pred=predicate,
-                    )
-        else:
-            cute.copy(
-                gmem_tiled_copy,
-                tVgV[None, None, None, block],
-                tVsV[None, None, None, smem_pipe_write if const_expr(self.num_stages > 1) else 0],
-                pred=tVpV if const_expr(self.check_hdim_v_oob) else None,
-            )
-
-
-class FlashAttentionForwardSm80(FlashAttentionForwardBase):
-    def _get_smem_layout_atom(self):
-        sQ_layout_atom = sm80_utils.get_smem_layout_atom(self.dtype, self.tile_hdim)
-        sK_layout_atom = sQ_layout_atom
-        sV_layout_atom = sm80_utils.get_smem_layout_atom(self.dtype, self.tile_hdimv)
-        sO_layout_atom = sV_layout_atom
-        sP_layout_atom = None
-        return sQ_layout_atom, sK_layout_atom, sV_layout_atom, sO_layout_atom, sP_layout_atom
-
-    def _get_tiled_mma(self):
-        tiled_mma_qk = cute.make_tiled_mma(
-            warp.MmaF16BF16Op(self.dtype, Float32, (16, 8, 16)),
-            (self.num_threads // 32, 1, 1),
-            permutation_mnk=(self.num_threads // 32 * 16, 16, 16),
-        )
-        tiled_mma_pv = cute.make_tiled_mma(
-            warp.MmaF16BF16Op(self.dtype, Float32, (16, 8, 16)),
-            (self.num_threads // 32, 1, 1),
-            permutation_mnk=(self.num_threads // 32 * 16, 16, 16),
-        )
-        return tiled_mma_qk, tiled_mma_pv
-
-    def _get_shared_storage_cls(self):
-        sQ_struct, sK_struct, sV_struct = [
-            cute.struct.Align[cute.struct.MemRange[self.dtype, cute.cosize(layout)], 1024]
-            for layout in (self.sQ_layout, self.sK_layout, self.sV_layout)
-        ]
-        cosize_sQV = max(cute.cosize(self.sQ_layout), cute.cosize(self.sV_layout))
-        sQV_struct = cute.struct.Align[cute.struct.MemRange[self.dtype, cosize_sQV], 1024]
-
-        @cute.struct
-        class SharedStorageQKV:
-            sV: sV_struct
-            sQ: sQ_struct
-            sK: sK_struct
-
-        @cute.struct
-        class SharedStorageSharedQV:
-            sQ: sQV_struct
-            sK: sK_struct
-
-        return SharedStorageQKV if const_expr(not self.Q_in_regs) else SharedStorageSharedQV
-
-    @cute.jit
-    def __call__(
-        self,
-        mQ: cute.Tensor,
-        mK: cute.Tensor,
-        mV: cute.Tensor,
-        mO: cute.Tensor,
-        mLSE: Optional[cute.Tensor],
-        softmax_scale: Float32,
-        mCuSeqlensQ: Optional[cute.Tensor] = None,
-        mCuSeqlensK: Optional[cute.Tensor] = None,
-        mSeqUsedQ: Optional[cute.Tensor] = None,
-        mSeqUsedK: Optional[cute.Tensor] = None,
-        mPageTable: Optional[cute.Tensor] = None,
-        window_size_left: Int32 | int | None = None,
-        window_size_right: Int32 | int | None = None,
-        learnable_sink: Optional[cute.Tensor] = None,
-        blocksparse_tensors: Optional[BlockSparseTensors] = None,
-        aux_data: AuxData = AuxData(),
-        # Always keep stream as the last parameter (EnvStream: obtained implicitly via TVM FFI).
-        stream: cuda.CUstream = None,
-    ):
-        """Configures and launches the flash attention kernel.
-
-        mQ/mK/mV/mO has same data types(supports fp16 and bf16) and same layout:
-        (batch_size, seqlen_q, num_head, head_dim):(_, _, _, 1)
-        """
-        assert learnable_sink is None, "Learnable sink is not supported in this kernel"
-        self._check_type(
-            *(t.element_type if t is not None else None for t in (mQ, mK, mV, mO, mLSE, mCuSeqlensQ, mCuSeqlensK, mSeqUsedQ, mSeqUsedK))
-        )
-        tiled_mma_qk, tiled_mma_pv = self._get_tiled_mma()
-        self.num_mma_threads = tiled_mma_pv.size
-        self.num_producer_threads = self.num_threads
-        self.num_Q_load_threads = self.num_threads
-        self.num_epilogue_threads = self.num_threads
-        self.use_tma_O = Arch.sm_90 <= self.arch < Arch.sm_120
-        self._setup_attributes()
-        SharedStorage = self._get_shared_storage_cls()
-        mQ, mK, mV, mO = [assume_tensor_aligned(t) for t in (mQ, mK, mV, mO)]
-        # Layout permutation: 4D non-varlen vs 3D varlen
-        QO_layout_transpose = [1, 3, 2, 0] if const_expr(mCuSeqlensQ is None) else [0, 2, 1]
-        KV_layout_transpose = [1, 3, 2, 0] if const_expr(mCuSeqlensK is None) else [0, 2, 1]
-        mQ, mO = [
-            cute.make_tensor(t.iterator, cute.select(t.layout, mode=QO_layout_transpose))
-            for t in (mQ, mO)
-        ]
-        mK, mV = [
-            cute.make_tensor(t.iterator, cute.select(t.layout, mode=KV_layout_transpose))
-            for t in (mK, mV)
-        ]
-        if const_expr(mLSE is not None):
-            LSE_layout_transpose = [2, 1, 0] if const_expr(mCuSeqlensQ is None) else [1, 0]
-            mLSE = cute.make_tensor(mLSE.iterator, cute.select(mLSE.layout, mode=LSE_layout_transpose))
-        if const_expr(self.pack_gqa):
-            nheads_kv = mK.shape[2]
-            mQ = pack_gqa_layout(mQ, self.qhead_per_kvhead, nheads_kv, head_idx=2)
-            mO = pack_gqa_layout(mO, self.qhead_per_kvhead, nheads_kv, head_idx=2)
-            if const_expr(mLSE is not None):
-                mLSE = pack_gqa_layout(mLSE, self.qhead_per_kvhead, nheads_kv, head_idx=1)
-        # TileScheduler for varlen, simple grid for non-varlen
-        if const_expr(mCuSeqlensQ is not None or mSeqUsedQ is not None):
-            TileScheduler = SingleTileVarlenScheduler
-        else:
-            TileScheduler = SingleTileScheduler
-        num_batch = (
-            mCuSeqlensQ.shape[0] - 1
-            if const_expr(mCuSeqlensQ is not None)
-            else cute.size(mQ.shape[3])
-        )
-        tile_sched_args = TileSchedulerArguments(
-            num_block=cute.ceil_div(cute.size(mQ.shape[0]), self.tile_m),
-            num_head=cute.size(mQ.shape[2]),
-            num_batch=num_batch,
-            num_splits=1,
-            seqlen_k=0,
-            headdim=mQ.shape[1],
-            headdim_v=mV.shape[1],
-            total_q=cute.size(mQ.shape[0])
-            if const_expr(mCuSeqlensQ is not None)
-            else cute.size(mQ.shape[0]) * cute.size(mQ.shape[3]),
-            tile_shape_mn=(self.tile_m, self.tile_n),
-            qhead_per_kvhead_packgqa=self.qhead_per_kvhead if const_expr(self.pack_gqa) else 1,
-            mCuSeqlensQ=mCuSeqlensQ,
-            mSeqUsedQ=mSeqUsedQ,
-        )
-        tile_sched_params = TileScheduler.to_underlying_arguments(tile_sched_args)
-        grid_dim = TileScheduler.get_grid_shape(tile_sched_params)
-        softmax_scale_log2, softmax_scale = utils.compute_softmax_scale_log2(softmax_scale, self.score_mod)
-        fastdiv_mods = utils.compute_fastdiv_mods(mQ, mK, self.qhead_per_kvhead, self.pack_gqa, aux_data.tensors)
-
-        self.kernel(
-            mQ,
-            mK,
-            mV,
-            mO,
-            mLSE,
-            mCuSeqlensQ,
-            mCuSeqlensK,
-            mSeqUsedQ,
-            mSeqUsedK,
-            softmax_scale_log2,
-            softmax_scale,
-            window_size_left,
-            window_size_right,
-            self.sQ_layout,
-            self.sK_layout,
-            self.sV_layout,
-            self.sO_layout,
-            self.sP_layout,
-            self.gmem_tiled_copy_Q,
-            self.gmem_tiled_copy_K,
-            self.gmem_tiled_copy_V,
-            self.gmem_tiled_copy_O,
-            tiled_mma_qk,
-            tiled_mma_pv,
-            SharedStorage,
-            tile_sched_params,
-            TileScheduler,
-            aux_data,
-            fastdiv_mods,
-        ).launch(
-            grid=grid_dim,
-            block=[self.num_threads, 1, 1],
-            smem=SharedStorage.size_in_bytes(),
-            stream=stream,
-        )
-
-    @cute.kernel
-    def kernel(
-        self,
-        mQ: cute.Tensor,
-        mK: cute.Tensor,
-        mV: cute.Tensor,
-        mO: cute.Tensor,
-        mLSE: Optional[cute.Tensor],
-        mCuSeqlensQ: Optional[cute.Tensor],
-        mCuSeqlensK: Optional[cute.Tensor],
-        mSeqUsedQ: Optional[cute.Tensor],
-        mSeqUsedK: Optional[cute.Tensor],
-        softmax_scale_log2: Float32,
-        softmax_scale: Optional[Float32],
-        window_size_left: Optional[Int32],
-        window_size_right: Optional[Int32],
-        sQ_layout: cute.ComposedLayout,
-        sK_layout: cute.ComposedLayout,
-        sV_layout: cute.ComposedLayout,
-        sO_layout: cute.ComposedLayout,
-        sP_layout: cute.ComposedLayout | None,
-        gmem_tiled_copy_Q: cute.TiledCopy,
-        gmem_tiled_copy_K: cute.TiledCopy,
-        gmem_tiled_copy_V: cute.TiledCopy,
-        gmem_tiled_copy_O: cute.TiledCopy,
-        tiled_mma_qk: cute.TiledMma,
-        tiled_mma_pv: cute.TiledMma,
-        SharedStorage: cutlass.Constexpr,
-        tile_sched_params,
-        TileScheduler: cutlass.Constexpr[Callable],
-        aux_data: AuxData = AuxData(),
-        fastdiv_mods=None,
-    ):
-        # Thread index, block index
-        tidx, _, _ = cute.arch.thread_idx()
-
-        tile_scheduler = TileScheduler.create(tile_sched_params)
-        work_tile = tile_scheduler.initial_work_tile_info()
-        m_block, num_head, batch_size, _ = work_tile.tile_idx
-
-        # Spark fix: surplus scheduler tiles must not read sequence metadata.
-        if work_tile.is_valid_tile:
-
-            block_info = BlockInfo(
-                self.tile_m,
-                self.tile_n,
-                self.is_causal,
-                self.is_local,
-                False,  # is_split_kv
-                window_size_left,
-                window_size_right,
-                qhead_per_kvhead_packgqa=self.qhead_per_kvhead if const_expr(self.pack_gqa) else 1,
-            )
-            seqlen = SeqlenInfoQK.create(
-                batch_idx=batch_size,
-                seqlen_q_static=mQ.shape[0] if const_expr(not self.pack_gqa) else mQ.shape[0][1],
-                seqlen_k_static=mK.shape[0],
-                mCuSeqlensQ=mCuSeqlensQ,
-                mCuSeqlensK=mCuSeqlensK,
-                mSeqUsedQ=mSeqUsedQ,
-                mSeqUsedK=mSeqUsedK,
-            )
-            n_block_min, n_block_max = block_info.get_n_block_min_max(seqlen, m_block)
-            # For varlen, wasted grid tiles (where batch_idx >= num_batch) will have
-            # seqlen_q=seqlen_k=0 and n_block_max=0.  Clamp to 0 so we don't use a
-            # negative block index for K/V loads; the load/store predicates already
-            # guard all memory accesses when seqlen is 0.
-            n_block = cutlass.max(n_block_max - 1, 0)
-
-            # ///////////////////////////////////////////////////////////////////////////////
-            # Get the appropriate tiles for this thread block.
-            # ///////////////////////////////////////////////////////////////////////////////
-            blkQ_shape = (self.tile_m, self.tile_hdim)
-            blkK_shape = (self.tile_n, self.tile_hdim)
-            blkV_shape = (self.tile_n, self.tile_hdimv)
-            num_head_kv = num_head if const_expr(self.pack_gqa) else num_head // self.qhead_per_kvhead
-            mQ_cur = seqlen.offset_batch_Q(mQ, batch_size, dim=3)[None, None, num_head]
-            if const_expr(not seqlen.has_cu_seqlens_k):
-                mK_cur = mK[None, None, num_head_kv, batch_size]
-                mV_cur = mV[None, None, num_head_kv, batch_size]
-            else:
-                mK_cur = cute.domain_offset((seqlen.offset_k, 0), mK[None, None, num_head_kv])
-                mV_cur = cute.domain_offset((seqlen.offset_k, 0), mV[None, None, num_head_kv])
-            if const_expr(not self.pack_gqa):
-                gQ = cute.local_tile(mQ_cur, blkQ_shape, (m_block, 0))
-            gK = cute.local_tile(mK_cur, blkK_shape, (None, 0))
-            gV = cute.local_tile(mV_cur, blkV_shape, (None, 0))
-
-            # ///////////////////////////////////////////////////////////////////////////////
-            # Get shared memory buffer
-            # ///////////////////////////////////////////////////////////////////////////////
-            smem = cutlass.utils.SmemAllocator()
-            storage = smem.allocate(SharedStorage)
-            sQ = storage.sQ.get_tensor(sQ_layout)
-            sK = storage.sK.get_tensor(sK_layout)
-            if const_expr(not self.Q_in_regs):
-                sV = storage.sV.get_tensor(sV_layout)
-            else:
-                sV = cute.make_tensor(cute.recast_ptr(sQ.iterator, dtype=self.dtype), sV_layout)
-            # Transpose view of V to tensor with layout (head_dim_v, tile_n) for tiled mma
-            sVt = layout_utils.transpose_view(sV)
-
-            gmem_thr_copy_K = gmem_tiled_copy_K.get_slice(tidx)
-            gmem_thr_copy_V = gmem_tiled_copy_V.get_slice(tidx)
-            # (CPY_Atom, CPY_N, CPY_K, n_block)
-            tKsK, tKgK = gmem_thr_copy_K.partition_D(sK), gmem_thr_copy_K.partition_S(gK)
-            # (CPY_Atom, CPY_N, CPY_K, n_block)
-            tVsV, tVgV = gmem_thr_copy_V.partition_D(sV), gmem_thr_copy_V.partition_S(gV)
-
-            # ///////////////////////////////////////////////////////////////////////////////
-            # Tile MMA compute thread partitions and allocate accumulators
-            # ///////////////////////////////////////////////////////////////////////////////
-            thr_mma_qk = tiled_mma_qk.get_slice(tidx)
-            thr_mma_pv = tiled_mma_pv.get_slice(tidx)
-            tSrQ = thr_mma_qk.make_fragment_A(thr_mma_qk.partition_A(sQ))
-            tSrK = thr_mma_qk.make_fragment_B(thr_mma_qk.partition_B(sK[None, None, 0]))
-            tOrVt = thr_mma_pv.make_fragment_B(thr_mma_pv.partition_B(sVt[None, None, 0]))
-            acc_shape_O = thr_mma_pv.partition_shape_C((self.tile_m, self.tile_hdimv))
-            acc_O = cute.make_rmem_tensor(acc_shape_O, Float32)
-            acc_O.fill(0.0)
-
-            # ///////////////////////////////////////////////////////////////////////////////
-            # Smem copy atom tiling
-            # ///////////////////////////////////////////////////////////////////////////////
-            smem_copy_atom_QK = cute.make_copy_atom(
-                warp.LdMatrix8x8x16bOp(transpose=False, num_matrices=4),
-                self.dtype,
-            )
-            smem_copy_atom_V = cute.make_copy_atom(
-                warp.LdMatrix8x8x16bOp(transpose=True, num_matrices=4),
-                self.dtype,
-            )
-            smem_thr_copy_Q = utils.make_tiled_copy_A(smem_copy_atom_QK, tiled_mma_qk).get_slice(tidx)
-            smem_thr_copy_K = utils.make_tiled_copy_B(smem_copy_atom_QK, tiled_mma_qk).get_slice(tidx)
-            smem_thr_copy_V = utils.make_tiled_copy_B(smem_copy_atom_V, tiled_mma_pv).get_slice(tidx)
-
-            tSsQ = smem_thr_copy_Q.partition_S(sQ)
-            tSsK = smem_thr_copy_K.partition_S(sK)
-            tOsVt = smem_thr_copy_V.partition_S(sVt)
-
-            # ///////////////////////////////////////////////////////////////////////////////
-            # Predicate: Mark indices that need to copy when problem_shape isn't a multiple
-            # of tile_shape
-            # ///////////////////////////////////////////////////////////////////////////////
-            # Construct identity layout for KV
-            cK = cute.make_identity_tensor((self.tile_n, self.tile_hdim))
-            tKcK = gmem_thr_copy_K.partition_S(cK)
-            t0KcK = gmem_thr_copy_K.get_slice(0).partition_S(cK)
-            if const_expr(self.tile_hdim == self.tile_hdimv):
-                tVcV = tKcK
-                t0VcV = t0KcK
-            else:
-                cV = cute.make_identity_tensor((self.tile_n, self.tile_hdimv))
-                tVcV = gmem_thr_copy_V.partition_S(cV)
-                t0VcV = gmem_thr_copy_V.get_slice(0).partition_S(cV)
-            # Allocate predicate tensors for m and n, here we only allocate the tile of k, and
-            # use "if" on the mn dimension.
-            # This is to reduce register pressure and gets 2-3% performance gain.
-            tKpK = utils.predicate_k(tKcK, limit=mK.shape[1])
-            if const_expr(self.same_hdim_kv):
-                tVpV = tKpK
-            else:
-                tVpV = utils.predicate_k(tVcV, limit=mV.shape[1])
-
-            # shape: (atom_v_m * rest_m)
-            softmax = Softmax.create(
-                softmax_scale_log2,
-                num_rows=acc_O.shape[0][0] * acc_O.shape[1],
-                softmax_scale=softmax_scale,
-            )
-            softmax.reset()
-
-            # group parameters for compute_one_n_block
-            mma_params = SimpleNamespace(
-                thr_mma_qk=thr_mma_qk,
-                thr_mma_pv=thr_mma_pv,
-                tSrQ=tSrQ,
-                tSrK=tSrK,
-                tOrVt=tOrVt,
-                acc_O=acc_O,
-            )
-            smem_copy_params = SimpleNamespace(
-                smem_thr_copy_Q=smem_thr_copy_Q,
-                smem_thr_copy_K=smem_thr_copy_K,
-                smem_thr_copy_V=smem_thr_copy_V,
-                tSsQ=tSsQ,
-                tSsK=tSsK,
-                tOsVt=tOsVt,
-            )
-            load_K = partial(
-                self.load_K, gmem_tiled_copy_K, tKgK, tKsK, tKcK, t0KcK, tKpK, seqlen=seqlen.seqlen_k
-            )
-            load_V = partial(
-                self.load_V, gmem_tiled_copy_V, tVgV, tVsV, tVcV, t0VcV, tVpV, seqlen=seqlen.seqlen_k
-            )
-
-            compute_one_n_block = partial(
-                self.compute_one_n_block,
-                mma_params=mma_params,
-                smem_copy_params=smem_copy_params,
-                softmax=softmax,
-                load_K=load_K,
-                load_V=load_V,
-                score_mod=self.score_mod,
-                batch_idx=batch_size,
-                head_idx=num_head,
-                m_block=m_block,
-                aux_data=aux_data,
-                fastdiv_mods=fastdiv_mods,
-            )
-
-            # ///////////////////////////////////////////////////////////////////////////////
-            # Prologue
-            # ///////////////////////////////////////////////////////////////////////////////
-            # Start async loads of the last mn-tile, where we take care of the mn residue
-            gmem_thr_copy_Q = gmem_tiled_copy_Q.get_slice(tidx)
-            if const_expr(not self.pack_gqa):
-                self.load_Q(gmem_thr_copy_Q, gQ, sQ, m_block, seqlen=seqlen.seqlen_q, headdim=mQ.shape[1])
-            else:
-                pack_gqa = PackGQA(self.tile_m, self.tile_hdim, self.check_hdim_oob, self.qhead_per_kvhead)
-                pack_gqa.load_Q(mQ_cur, sQ, gmem_tiled_copy_Q, tidx, m_block, seqlen.seqlen_q)
-            cute.arch.cp_async_commit_group()
-
-            def preprocess_Q():
-                cute.arch.cp_async_wait_group(self.num_stages * 2 - 1)
-                if const_expr(self.Q_in_regs):
-                    cute.arch.barrier()
-                    tSrQ_copy_view = smem_thr_copy_Q.retile(tSrQ)
-                    cute.copy(smem_thr_copy_Q, tSsQ, tSrQ_copy_view)
-
-            # If Q_in_regs, we load Q, then load 1 stage of K, then (optionally) rotate Q and
-            # read from smem_q to registers, then load V.
-            # If !Q_in_regs, we load Q, load all stages of K & V, then (optionally) rotate Q.
-            if const_expr(self.Q_in_regs):
-                load_K(n_block, smem_pipe_write=0, need_predicates=True)
-                cute.arch.cp_async_commit_group()
-                preprocess_Q()
-                cute.arch.barrier()  # Make sure all threads have read smem_q before loading V
-
-            for stage in cutlass.range_constexpr(self.num_stages):
-                if const_expr(not self.Q_in_regs or stage > 0):
-                    if stage == 0 or n_block - stage >= 0:
-                        load_K(n_block - stage, smem_pipe_write=stage, need_predicates=stage == 0)
-                    cute.arch.cp_async_commit_group()
-                if const_expr(stage < self.num_stages - 1):
-                    if stage == 0 or n_block - stage >= 0:
-                        load_V(n_block - stage, smem_pipe_write=stage, need_predicates=stage == 0)
-                    cute.arch.cp_async_commit_group()
-            if const_expr(not self.Q_in_regs):
-                preprocess_Q()
-
-            # ///////////////////////////////////////////////////////////////////////////////
-            # Mainloop
-            # ///////////////////////////////////////////////////////////////////////////////
-            # Start processing of the first n-block.
-            # For performance reason, we separate out two kinds of iterations:
-            # those that need masking on S, and those that don't.
-            # We need masking on S for the very last block when K and V has length not multiple of tile_n.
-            # We also need masking on S if it's causal, for the last several blocks.
-            mask = AttentionMask(
-                self.tile_m,
-                self.tile_n,
-                seqlen,
-                window_size_left,
-                window_size_right,
-                self.qhead_per_kvhead if const_expr(self.pack_gqa) else 1,
-            )
-            mask_fn = partial(
-                mask.apply_mask,
-                batch_idx=batch_size,
-                head_idx=num_head,
-                m_block=m_block,
-                thr_mma=thr_mma_qk,
-                mask_causal=self.is_causal,
-                mask_local=self.is_local,
-                aux_data=aux_data,
-                fastdiv_mods=fastdiv_mods if const_expr(self.mask_mod is not None) else None,
-            )
-
-            # First iteration with seqlen masking
-            smem_pipe_read = Int32(0)
-            smem_pipe_write = Int32(self.num_stages - 1)
-            compute_one_n_block(
-                n_block,
-                smem_pipe_read,
-                smem_pipe_write,
-                is_first_n_block=True,
-                seqlen=seqlen,
-                mask_fn=partial(mask_fn, mask_mod=self.mask_mod, mask_seqlen=True),
-            )
-            smem_pipe_read = self.advance_pipeline(smem_pipe_read)
-            smem_pipe_write = self.advance_pipeline(smem_pipe_write)
-            # Next couple of iterations with causal masking
-            if const_expr(self.is_causal or self.is_local):
-                n_block_min_causal_local_mask = block_info.get_n_block_min_causal_local_mask(
-                    seqlen, m_block, n_block_min
-                )
-                for n_tile in cutlass.range(n_block_max - 1 - n_block_min_causal_local_mask, unroll=1):
-                    n_block = n_block_max - 2 - n_tile
-                    compute_one_n_block(
-                        n_block,
-                        smem_pipe_read,
-                        smem_pipe_write,
-                        seqlen=seqlen,
-                        mask_fn=partial(mask_fn, mask_mod=self.mask_mod, mask_seqlen=True),
-                    )
-                    smem_pipe_read = self.advance_pipeline(smem_pipe_read)
-                    smem_pipe_write = self.advance_pipeline(smem_pipe_write)
-            # The remaining iterations have no masking
-            for n_tile in cutlass.range(n_block, unroll=1):
-                compute_one_n_block(
-                    n_block - n_tile - 1, smem_pipe_read, smem_pipe_write,
-                    seqlen=seqlen, is_first_n_block=False,
-                    mask_fn=partial(mask_fn, mask_mod=self.mask_mod, mask_seqlen=False)
-                )
-                smem_pipe_read = self.advance_pipeline(smem_pipe_read)
-                smem_pipe_write = self.advance_pipeline(smem_pipe_write)
-            # TODO: local
-
-            # normalize acc_O by row_sum and calculate the lse
-            row_scale = softmax.finalize()
-            softmax.rescale_O(acc_O, row_scale)
-
-            # ///////////////////////////////////////////////////////////////////////////////
-            # Epilogue
-            # ///////////////////////////////////////////////////////////////////////////////
-            # reuse sQ's data iterator
-            sO = cute.make_tensor(sQ.iterator, sO_layout)
-            self.epilogue(
-                acc_O,
-                softmax.row_sum,
-                mO,
-                mLSE,
-                sO,
-                seqlen,
-                gmem_tiled_copy_O,
-                None,
-                tiled_mma_pv,
-                tidx,
-                m_block,
-                num_head,
-                batch_size,
-            )
-
-    @cute.jit
-    def compute_one_n_block(
-        self,
-        n_block: Int32,
-        smem_pipe_read: Int32,
-        smem_pipe_write: Int32,
-        mma_params: SimpleNamespace,
-        smem_copy_params: SimpleNamespace,
-        softmax: Softmax,
-        load_K: Callable,
-        load_V: Callable,
-        score_mod: Callable | None,
-        batch_idx: cutlass.Int32,
-        head_idx: cutlass.Int32,
-        m_block: cutlass.Int32,
-        seqlen: SeqlenInfoQK,
-        aux_data: AuxData = AuxData(),
-        fastdiv_mods=None,
-        mask_fn: Optional[Callable] = None,
-        is_first_n_block: cutlass.Constexpr = False,
-        check_inf: cutlass.Constexpr = True,
-    ):
-        """Compute one n_block of S/O.
-
-        This function provides different variants for processing the first n block versus
-        subsequent blocks.
-        """
-
-        def sync():
-            cute.arch.cp_async_wait_group(self.num_stages * 2 - 2)
-            cute.arch.barrier()
-
-        acc_shape_S = mma_params.thr_mma_qk.partition_shape_C((self.tile_m, self.tile_n))
-        acc_S = cute.make_rmem_tensor(acc_shape_S, Float32)
-        acc_S.fill(0.0)
-        # wait for smem tile QK before mma calculation for S
-        sync()
-
-        # need predicates for the first tile
-        def load_V_next():
-            if self.num_stages == 1 or n_block - self.num_stages + 1 >= 0:
-                load_V(
-                    n_block - self.num_stages + 1,
-                    smem_pipe_write,
-                    need_predicates=is_first_n_block and self.num_stages == 1,
-                )
-            cute.arch.cp_async_commit_group()
-
-        load_V_next()
-        sm80_utils.gemm(
-            mma_params.thr_mma_qk,
-            acc_S,
-            mma_params.tSrQ,
-            mma_params.tSrK,
-            smem_copy_params.tSsQ,
-            smem_copy_params.tSsK[
-                None, None, None, smem_pipe_read if const_expr(self.num_stages > 1) else 0
-            ],
-            smem_copy_params.smem_thr_copy_Q,
-            smem_copy_params.smem_thr_copy_K,
-            # hook_fn=load_V_next,
-            A_in_regs=self.Q_in_regs,
-        )
-        if const_expr(score_mod is not None):
-            self.apply_score_mod(
-                mma_params.thr_mma_qk,
-                batch_idx,
-                head_idx,
-                m_block,
-                acc_S,
-                n_block,
-                softmax_scale=softmax.softmax_scale,
-                seqlen=seqlen,
-                aux_data=aux_data,
-                fastdiv_mods=fastdiv_mods,
-            )
-
-        smem_pipe_write = self.advance_pipeline(smem_pipe_write)
-
-        def load_K_next():
-            if n_block - self.num_stages >= 0:
-                load_K(n_block - self.num_stages, smem_pipe_write, need_predicates=False)
-            cute.arch.cp_async_commit_group()
-
-        # wait for smem tile V for O
-        if const_expr(self.num_stages == 1):
-            sync()
-            load_K_next()
-        if const_expr(mask_fn is not None):
-            mask_fn(acc_S, n_block=n_block)
-        row_scale = softmax.online_softmax(acc_S, is_first=is_first_n_block, check_inf=check_inf)
-        softmax.rescale_O(mma_params.acc_O, row_scale)
-        rP = cute.make_fragment_like(acc_S, self.dtype)
-        rP.store(acc_S.load().to(self.dtype))
-        tOrP = layout_utils.reshape_acc_to_frgA(rP)
-        if const_expr(self.num_stages > 1):
-            sync()
-            load_K_next()
-        sm80_utils.gemm_rs(
-            mma_params.thr_mma_pv,
-            mma_params.acc_O,
-            tOrP,
-            mma_params.tOrVt,
-            smem_copy_params.tOsVt[
-                None, None, None, smem_pipe_read if const_expr(self.num_stages > 1) else 0
-            ],
-            smem_copy_params.smem_thr_copy_V,
-            # hook_fn=load_K_next,
-        )
-        # if const_expr(self.num_stages > 1):
-        #     load_K_next()
-    @cute.jit
-    def apply_score_mod(
-        self,
-        thr_mma_qk,
-        batch_idx,
-        head_idx,
-        m_block,
-        acc_S,
-        n_block,
-        softmax_scale,
-        seqlen,
-        aux_data: AuxData = AuxData(),
-        fastdiv_mods=None,
-    ):
-        # Prepare index tensor
-        cS = cute.make_identity_tensor((self.tile_m, self.tile_n))
-        cS = cute.domain_offset((m_block * self.tile_m, n_block * self.tile_n), cS)
-        tScS = thr_mma_qk.partition_C(cS)
-
-        apply_score_mod_inner(
-            acc_S,
-            tScS,
-            self.score_mod,
-            batch_idx,
-            head_idx,
-            softmax_scale,
-            self.score_vec_size,
-            self.qk_acc_dtype,
-            aux_data,
-            fastdiv_mods,
-            seqlen_info=seqlen,
-            constant_q_idx=None,
-            qhead_per_kvhead=self.qhead_per_kvhead if const_expr(self.pack_gqa) else 1,
-        )
-
-
-# SM90 forward pass moved to flash_fwd_sm90.py; re-export for backward compatibility
-def __getattr__(name):
-    if name == "FlashAttentionForwardSm90":
-        from flash_attn.cute.flash_fwd_sm90 import FlashAttentionForwardSm90
-        return FlashAttentionForwardSm90
-    raise AttributeError(f"module {__name__!r} has no attribute {name!r}")

code/sol_h3_vae/__init__.py

@@ -0,0 +1,4 @@
+"""Sol-H3-VAE: fast noncausal decoder for MiniMax-H3 video latents."""
+from .decoder import SolH3VAEDecoder
+
+__all__ = ["SolH3VAEDecoder"]

code/{fast_h3_vae.py → sol_h3_vae/decoder.py}

@@ -1,41 +1,39 @@
-"""Fast H3 VAE (non-causal): drop-in replacement for MiniMaxH3VideoVAE.decode_base on whole clips.
+"""Sol-H3-VAE: a 24-layer noncausal decoder for MiniMax-H3 video latents.
 
-    from fast_h3_vae import FastH3VAEDecoder
-    dec = FastH3VAEDecoder(vae, "fast_h3_vae_noncausal_24l.safetensors")      # vae: the released H3 VAE (for its
-    pixels = dec.decode(z)            # z: normalized latents (B=1, T, 24, H, W)   post_quant_conv / embedder / stats)
+    from sol_h3_vae import SolH3VAEDecoder
+    dec = SolH3VAEDecoder("MiniMax-H3/vae/config.json", "sol_h3_vae.safetensors")
+    pixels = dec.decode(z)            # z: normalized latents (1, T, 24, H, W) -> (1, 3, F, 16H, 16W)
     rgb_u8 = dec.decode_rgb_uint8(z)  # (F, 16H, 16W, 3) uint8 on the GPU
 
-* any clip length T (t = 0 must be the start of a native 5-latent block, as for the released decoder). The output has
-  the same frames as decode_base: 17 frames per full block, the first latent of every block keeps 1 of its 4 frames
-  (tested T = 3..15: frame count equal to decode_base, fast vs reference 60.7-61.4 dB). For T = 1 and 2 the released
-  decode_base raises an error; this decoder returns 1 and 5 frames by the same rule;
+* any clip length T (t = 0 must be the start of a native 5-latent block, as for the H3 VAE). The output has the same
+  frames as the H3 VAE decoder: 17 frames per full block, the first latent of every block keeps 1 of its 4 frames;
 * engine="fast" (default): NonCausalFast3 (FlashAttention-4 varlen windows, fused Triton kernels, one CUDA graph per
-  (T, H, W)); engine="ref": the plain PyTorch NonCausalViTDecoder (any GPU with SDPA);
+  (T, H, W)); engine="ref": the plain PyTorch NonCausalViTDecoder (any GPU with SDPA, no FlashAttention-4 needed);
 * the fast engine is built on first use of a shape and cached (`max_engines` most recent shapes).
 """
 from __future__ import annotations
 
+import json
 from collections import OrderedDict
 
 import torch
 from safetensors.torch import load_file
 
-from noncausal_vit_decoder import KEEP, NonCausalViTDecoder
+from .model import NonCausalViTDecoder
+from .template import RGBProcessor, make_template
 
 
-class FastH3VAEDecoder:
-    def __init__(self, vae, weights: str, *, layers: int = 24, ffn_keep: int | None = None, time_norm: float = 5,
-                 engine: str = "fast", max_engines: int = 4, device="cuda", offload_vae: bool = True):
-        """offload_vae: after building the decoder, move the released VAE (encoder + 36-layer decoder, ~10 GB) to the CPU;
-        decoding only needs its processor (pixel de-normalisation), which works on GPU tensors."""
-        self.vae = vae
-        m = NonCausalViTDecoder(vae, KEEP[layers], time_norm=time_norm, ffn_keep=ffn_keep)
-        r = m.load_state_dict(load_file(weights, device="cpu"), strict=False)
-        assert not r.missing_keys and not r.unexpected_keys, r
-        self.model = m.to(device).eval().requires_grad_(False)
-        if offload_vae:
-            vae.to("cpu")
-            torch.cuda.empty_cache()
+class SolH3VAEDecoder:
+    def __init__(self, vae_config, weights, *, engine: str = "fast", device="cuda", max_engines: int = 1):
+        """vae_config: the H3 VAE config.json (Diffusers layout); weights: sol_h3_vae.safetensors. The checkpoint holds
+        every decoder parameter, so the H3 VAE weights are not needed."""
+        with open(vae_config) as f:
+            config = json.load(f)
+        model = NonCausalViTDecoder(make_template(config), time_norm=5)
+        model.load_state_dict(load_file(str(weights), device="cpu"), strict=True, assign=True)
+        assert not any(t.is_meta for t in list(model.parameters()) + list(model.buffers()))
+        self.model = model.to(device).eval().requires_grad_(False)
+        self.processor = RGBProcessor()
         self.engine, self.max_engines = engine, max_engines
         self._engines: OrderedDict = OrderedDict()
 
@@ -44,7 +42,7 @@ class FastH3VAEDecoder:
         if key in self._engines:
             self._engines.move_to_end(key)
             return self._engines[key]
-        from noncausal_fast3 import NonCausalFast3
+        from .fast import NonCausalFast3
         eng = NonCausalFast3(self.model, T, H, W, cuda_graph=True)
         self._engines[key] = eng
         while len(self._engines) > self.max_engines:
@@ -54,7 +52,8 @@ class FastH3VAEDecoder:
 
     @torch.no_grad()
     def decode(self, z):
-        """z: normalized latents (1, T, 24, H, W). Returns decode_base-space pixels (1, 3, F, 16H, 16W), fp32."""
+        """z: normalized latents (1, T, 24, H, W). Returns the H3 VAE decoder's pixels (1, 3, F, 16H, 16W), fp32,
+        ImageNet-normalized like the H3 VAE output; `processor.revert_tensor` maps them to [0, 1]."""
         assert z.dim() == 5 and z.shape[0] == 1 and z.shape[2] == 24, z.shape
         if self.engine == "ref":
             with torch.autocast("cuda", dtype=torch.bfloat16):
@@ -64,6 +63,6 @@ class FastH3VAEDecoder:
 
     @torch.no_grad()
     def decode_rgb_uint8(self, z):
-        """Same as decode, then the VAE's revert_tensor, clamp and uint8: (F, 16H, 16W, 3) on the GPU."""
+        """Same as decode, then revert_tensor, clamp and uint8: (F, 16H, 16W, 3) on the GPU."""
         px = self.decode(z)
-        return self.vae.processor.revert_tensor(px).clamp_(0, 1).mul_(255).round_().to(torch.uint8)[0].permute(1, 2, 3, 0)
+        return self.processor.revert_tensor(px).clamp_(0, 1).mul_(255).round_().to(torch.uint8)[0].permute(1, 2, 3, 0)

code/{noncausal_fast3.py → sol_h3_vae/fast.py}

@@ -10,7 +10,7 @@ import torch
 import triton
 import triton.language as tl
 
-from noncausal_fast import BF16, NonCausalFast, _resid_norm_kernel, qgemm, qgemm_act, _flash_attn_fwd, _nc_qkv_post_kernel
+from .fast_base import BF16, NonCausalFast, _resid_norm_kernel, qgemm, qgemm_act, _flash_attn_fwd, _nc_qkv_post_kernel
 
 
 @triton.jit

code/{noncausal_fast.py → sol_h3_vae/fast_base.py}

@@ -6,7 +6,7 @@
   tokens of the block + its global tokens;
 * the 5 global rows of every block are appended to the token rows, so each linear layer is a single GEMM;
 * residual add + RMSNorm, SwiGLU in the up-projection GEMM (quack), fp32 head as one bf16 GEMM on a hi/lo split
-  (fused_kernels.py);
+  (kernels.py);
 * the whole decode is one CUDA graph.
 """
 from __future__ import annotations
@@ -24,8 +24,8 @@ if "cutlass" not in sys.modules:
 from flash_attn.cute.interface import _flash_attn_fwd
 from quack.gemm_interface import gemm as qgemm, gemm_act as qgemm_act
 
-from fused_kernels import _interleave, _resid_ln_split_kernel, _resid_norm_kernel
-from noncausal_vit_decoder import OUT_CH, PATCH, PATCH_T, NonCausalViTDecoder
+from .kernels import _interleave, _resid_ln_split_kernel, _resid_norm_kernel
+from .model import OUT_CH, PATCH, PATCH_T, NonCausalViTDecoder
 
 BF16 = torch.bfloat16
 import triton

code/{noncausal_vit_decoder.py → sol_h3_vae/model.py}

@@ -23,7 +23,7 @@ import torch
 from torch import nn
 from torch.nn import functional as F
 
-from spark_compat import _apply_qk_norm, _apply_rotary_pos_emb_impl
+from .template import _apply_qk_norm, _apply_rotary_pos_emb_impl
 
 PATCH, PATCH_T, OUT_CH = 16, 4, 3
 KEEP_24 = [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 19, 20, 21, 31, 32, 35]

code/spark_decoder.py

@@ -1,25 +0,0 @@
-"""Build the decoder from the Diffusers H3 VAE config.json alone; decoding is FastH3VAEDecoder's."""
-import json
-from collections import OrderedDict
-from types import SimpleNamespace
-
-from safetensors.torch import load_file
-
-from fast_h3_vae import FastH3VAEDecoder
-from noncausal_vit_decoder import NonCausalViTDecoder
-from spark_compat import RGBProcessor, make_template
-from spark_fa4 import enable_sm120_guard
-
-
-class SparkFastH3VAEDecoder(FastH3VAEDecoder):
-    def __init__(self, vae_config, weights, *, engine='fast', device='cuda', max_engines=1):
-        enable_sm120_guard()
-        with open(vae_config) as f:
-            config = json.load(f)
-        model = NonCausalViTDecoder(make_template(config), time_norm=5)
-        model.load_state_dict(load_file(str(weights), device='cpu'), strict=True, assign=True)
-        assert not any(t.is_meta for t in list(model.parameters()) + list(model.buffers()))
-        self.model = model.to(device).eval().requires_grad_(False)
-        self.vae = SimpleNamespace(processor=RGBProcessor())
-        self.engine, self.max_engines = engine, max_engines
-        self._engines = OrderedDict()

code/spark_fa4.py

@@ -1,22 +0,0 @@
-"""Process-local SM120 fix for FlashAttention-4 (4.0.0b19-b22) varlen forward.
-
-On SM120 (RTX 5090, DGX Spark) surplus scheduler work tiles could read sequence
-metadata past the batch. _spark_flash_fwd_guarded.py is FA4's flash_fwd.py with that
-guard added; the installed FA4 package and valid-tile computation stay unchanged.
-"""
-from flash_attn.cute import interface
-from flash_attn.cute.cache_utils import JITCache
-from flash_attn.cute.flash_fwd_sm120 import FlashAttentionForwardSm120
-from _spark_flash_fwd_guarded import FlashAttentionForwardSm80
-
-
-class SparkFlashAttentionForwardSm120(FlashAttentionForwardSm120):
-    kernel = FlashAttentionForwardSm80.kernel
-
-
-def enable_sm120_guard():
-    if interface.FlashAttentionForwardSm120 is SparkFlashAttentionForwardSm120:
-        return
-    interface.FlashAttentionForwardSm120 = SparkFlashAttentionForwardSm120
-    # Never reuse an unguarded binary from an existing in-memory/disk cache.
-    interface._flash_attn_fwd.compile_cache = JITCache()