"""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()