Upgrade fa3 hash (#28394)

Co-authored-by: Shijin <dovis.zhang02@gmail.com>

python/sglang/jit_kernel/flash_attention.py

@@ -35,6 +35,7 @@ def flash_attn_with_kvcache(
     scheduler_metadata=None,
     num_splits=0,  # Can be tuned for speed
     pack_gqa=None,  # Can be tuned for speed
+    only_qv=False,  # ver=3 only: skip K matmul when qk rope dim is 0
     sm_margin=0,  # Can be tuned if some SMs are used for communication
     return_softmax_lse=False,
     sinks=None,
@@ -162,6 +163,7 @@ def flash_attn_with_kvcache(
             scheduler_metadata=scheduler_metadata,
             num_splits=num_splits,
             pack_gqa=pack_gqa,
+            only_qv=only_qv,
             sm_margin=sm_margin,
             return_softmax_lse=return_softmax_lse,
             sinks=sinks,

python/sglang/jit_kernel/flash_attention_v3.py

@@ -128,6 +128,7 @@ def flash_attn_with_kvcache(
     scheduler_metadata=None,
     num_splits=0,  # Can be tuned for speed
     pack_gqa=None,  # Can be tuned for speed
+    only_qv=False,  # Skip K matmul when qk rope dim is 0 (requires qv)
     sm_margin=0,  # Can be tuned if some SMs are used for communication
     return_softmax_lse=False,
     sinks=None,
@@ -138,7 +139,11 @@ def flash_attn_with_kvcache(
             "flash_attn at sgl-kernel is only supported on sm90 and above"
         )
 
-    assert k_cache.stride(-1) == 1, "k_cache must have contiguous last dimension"
+    # When only_qv=True the caller may pass k_cache=None (synthetic K is
+    # allocated inside the sgl-kernel wrapper). Skip the stride check in that
+    # case so the rope=0 path doesn't trip the assertion.
+    if k_cache is not None:
+        assert k_cache.stride(-1) == 1, "k_cache must have contiguous last dimension"
     assert v_cache.stride(-1) == 1, "v_cache must have contiguous last dimension"
 
     return _call_fa3_kernel(
@@ -171,9 +176,10 @@ def flash_attn_with_kvcache(
         scheduler_metadata,
         num_splits,
         pack_gqa,
-        sm_margin,
-        return_softmax_lse,
-        sinks,
+        sm_margin=sm_margin,
+        only_qv=only_qv,
+        return_softmax_lse=return_softmax_lse,
+        sinks=sinks,
         out=out,
     )
 
@@ -201,6 +207,7 @@ def flash_attn_varlen_func(
     softcap=0.0,
     num_splits=1,
     pack_gqa=None,
+    only_qv=False,
     sm_margin=0,
     return_softmax_lse=False,
     sinks=None,
@@ -265,6 +272,7 @@ def flash_attn_varlen_func(
         softcap=softcap,
         num_splits=num_splits,
         pack_gqa=pack_gqa,
+        only_qv=only_qv,
         sm_margin=sm_margin,
         return_softmax_lse=return_softmax_lse,
         sinks=sinks,