[Bug] Fix incorrect assertion in FA4 and add UT. (#11182)

2025-10-06 14:58:39 -07:00
parent 73ea484af1
commit 748f86f3de
2 changed files with 576 additions and 5 deletions
--- a/sgl-kernel/python/sgl_kernel/flash_attn.py
+++ b/sgl-kernel/python/sgl_kernel/flash_attn.py
@@ -161,10 +161,7 @@ def flash_attn_with_kvcache(
            k is None and v is None
        ), "FA4 does not support updating KV cache in-place."
        assert (
-            rotary_cos is None
-            and rotary_sin is None
-            and rotary_interleaved is None
-            and rotary_seqlens is None
+            rotary_cos is None and rotary_sin is None and rotary_seqlens is None
        ), "FA4 does not support rotary embedding."
        assert (
            cache_batch_idx is None and cache_leftpad is None