File size: 999 Bytes
775b8fb | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 | """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()
|