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