bump fa4 to official released fa4 pkg (#20303)

This commit is contained in:
Rain Jiang
2026-03-17 17:22:56 -07:00
committed by GitHub
parent c77d7c629e
commit cb1e63aba4
4 changed files with 11 additions and 25 deletions

View File

@@ -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)]