bump fa4 to official released fa4 pkg (#20303)

python/pyproject.toml

@@ -38,7 +38,7 @@ dependencies = [
   "msgspec",
   "ninja",
   "numpy",
-  "nvidia-cutlass-dsl==4.3.5",
+  "nvidia-cutlass-dsl>=4.4.1",
   "nvidia-ml-py",
   "openai-harmony==0.0.4",
   "openai==2.6.1",
@@ -54,12 +54,12 @@ dependencies = [
   "pydantic",
   "python-multipart",
   "pyzmq>=25.1.2",
-  "quack-kernels==0.2.4",
+  "quack-kernels>=0.3.0",
   "requests",
   "scipy",
   "sentencepiece",
   "setproctitle",
-  "sgl-fa4==4.0.5",
+  "flash-attn-4>=4.0.0b4",
   "sglang-kernel==0.4.0",
   "soundfile==0.13.1",
   "tiktoken",
@@ -77,7 +77,6 @@ dependencies = [
   "uvloop",
   "watchfiles",
   "xgrammar==0.1.27",
-
   "smg-grpc-servicer>=0.5.0",
 ]
 

python/sglang/jit_kernel/flash_attention_v4.py

@@ -5,7 +5,7 @@ from typing import Callable, Optional, Tuple, Union
 import torch
 
 try:
-    from sgl_fa4.cute import flash_attn_varlen_func as _flash_attn_varlen_func
+    from flash_attn.cute import flash_attn_varlen_func as _flash_attn_varlen_func
 except Exception as _e:  # pragma: no cover
     _flash_attn_varlen_func = None
     _flash_attn_import_error = _e
@@ -44,7 +44,7 @@ def flash_attn_varlen_func(
     if _flash_attn_varlen_func is None:  # pragma: no cover
         raise ImportError(
             "Vendored FlashAttention CUTE is not available (cannot import "
-            "sgl_fa4.cute). Please check your source tree."
+            "flash_attn.cute). Please check your source tree."
         ) from _flash_attn_import_error
 
     q, k, v = [_maybe_contiguous(t) for t in (q, k, v)]