[Diffusion] Support select fa2 backend in hopper (#17514)
This commit is contained in:
@@ -246,6 +246,13 @@ class CudaPlatformBase(Platform):
|
|||||||
elif selected_backend == AttentionBackendEnum.SAGE_SLA_ATTN:
|
elif selected_backend == AttentionBackendEnum.SAGE_SLA_ATTN:
|
||||||
logger.info("Using Sage Sparse Linear Attention backend")
|
logger.info("Using Sage Sparse Linear Attention backend")
|
||||||
return "sglang.multimodal_gen.runtime.layers.attention.backends.sparse_linear_attn.SageSparseLinearAttentionBackend"
|
return "sglang.multimodal_gen.runtime.layers.attention.backends.sparse_linear_attn.SageSparseLinearAttentionBackend"
|
||||||
|
elif selected_backend == AttentionBackendEnum.FA2:
|
||||||
|
from sglang.multimodal_gen.runtime.layers.attention.backends.flash_attn_2 import ( # noqa: F401
|
||||||
|
FlashAttention2Backend,
|
||||||
|
)
|
||||||
|
|
||||||
|
logger.info("Using FlashAttention2 backend")
|
||||||
|
return "sglang.multimodal_gen.runtime.layers.attention.backends.flash_attn_2.FlashAttention2Backend"
|
||||||
elif selected_backend in [
|
elif selected_backend in [
|
||||||
AttentionBackendEnum.FA,
|
AttentionBackendEnum.FA,
|
||||||
]:
|
]:
|
||||||
|
|||||||
Reference in New Issue
Block a user