@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 8b5014d3dd

Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git

Synced Nov 7, 11:30 UTC

[Attention] FA4 integration (#32974)

Signed-off-by: Lucas Wilkinson <lwilkins@redhat.com> Signed-off-by: Matthew Bonanni <mbonanni@redhat.com> Signed-off-by: Lucas Wilkinson <LucasWilkinson@users.noreply.github.com> Co-authored-by: Matthew Bonanni <mbonanni@redhat.com> Co-authored-by: Tyler Michael Smith <tyler@neuralmagic.com>

Lucas Wilkinson committed Mar 1, 2026 at 18:44 UTC 8b5014d3dd343736ccf3e26cd44a0bb7700d205c

15 files changed +817 -54

.buildkite/test_areas/misc.yaml

@@ -9,6 +9,7 @@ steps:

9 - tests/v1

10 commands:

11 - uv pip install --system -r /vllm-workspace/requirements/kv_connectors.txt

12 + - export VLLM_WORKER_MULTIPROC_METHOD=spawn

13 # split the test to avoid interference

14 - pytest -v -s -m 'not cpu_test' v1/core

15 - pytest -v -s v1/executor

.gitignore

@@ -3,6 +3,8 @@

4 # vllm-flash-attn built from source

5 vllm/vllm_flash_attn/*

6 +!vllm/vllm_flash_attn/__init__.py

7 +!vllm/vllm_flash_attn/flash_attn_interface.py

9 # OpenAI triton kernels copied from source

10 vllm/third_party/triton_kernels/*

cmake/external_projects/vllm_flash_attn.cmake

+55 -30

@@ -17,7 +17,8 @@ endif()

17 # They should be identical but if they aren't, this is a massive footgun.

18 #

19 # The vllm-flash-attn install rules are nested under vllm to make sure the library gets installed in the correct place.

20 -# To only install vllm-flash-attn, use --component _vllm_fa2_C (for FA2) or --component _vllm_fa3_C (for FA3).

20 +# To only install vllm-flash-attn, use --component _vllm_fa2_C (for FA2), --component _vllm_fa3_C (for FA3),

21 +# or --component _vllm_fa4_cutedsl_C (for FA4 CuteDSL Python files).

22 # If no component is specified, vllm-flash-attn is still installed.

24 # If VLLM_FLASH_ATTN_SRC_DIR is set, vllm-flash-attn is installed from that directory instead of downloading.

@@ -38,7 +39,7 @@ else()

39 FetchContent_Declare(

40 vllm-flash-attn

41 GIT_REPOSITORY https://github.com/vllm-project/flash-attention.git

41 - GIT_TAG 5824e6e2008271063c3229ab3e7032bd74abbbc6

42 + GIT_TAG 140c00c0241bb60cc6e44e7c1be9998d4b20d8d2

43 GIT_PROGRESS TRUE

44 # Don't share the vllm-flash-attn build between build types

45 BINARY_DIR ${CMAKE_BINARY_DIR}/vllm-flash-attn

@@ -46,38 +47,62 @@ else()

47 endif()

49 -# Ensure the vllm/vllm_flash_attn directory exists before installation

50 -install(CODE "file(MAKE_DIRECTORY \"\${CMAKE_INSTALL_PREFIX}/vllm/vllm_flash_attn\")" ALL_COMPONENTS)

51 -

52 -# Make sure vllm-flash-attn install rules are nested under vllm/

53 -# This is here to support installing all components under the same prefix with cmake --install.

54 -# setup.py installs every component separately but uses the same prefix for all.

55 -# ALL_COMPONENTS is used to avoid duplication for FA2 and FA3,

56 -# and these statements don't hurt when installing neither component.

57 -install(CODE "set(CMAKE_INSTALL_LOCAL_ONLY FALSE)" ALL_COMPONENTS)

58 -install(CODE "set(OLD_CMAKE_INSTALL_PREFIX \"\${CMAKE_INSTALL_PREFIX}\")" ALL_COMPONENTS)

59 -install(CODE "set(CMAKE_INSTALL_PREFIX \"\${CMAKE_INSTALL_PREFIX}/vllm/\")" ALL_COMPONENTS)

50 +# Install rules for FA components need the install prefix nested under vllm/

51 +# These run at install time, before the FA library's own install rules

52 +foreach(_FA_COMPONENT _vllm_fa2_C _vllm_fa3_C)

53 + install(CODE "set(CMAKE_INSTALL_LOCAL_ONLY FALSE)" COMPONENT ${_FA_COMPONENT})

54 + install(CODE "set(OLD_CMAKE_INSTALL_PREFIX \"\${CMAKE_INSTALL_PREFIX}\")" COMPONENT ${_FA_COMPONENT})

55 + install(CODE "set(CMAKE_INSTALL_PREFIX \"\${CMAKE_INSTALL_PREFIX}/vllm/\")" COMPONENT ${_FA_COMPONENT})

56 +endforeach()

58 # Fetch the vllm-flash-attn library

59 FetchContent_MakeAvailable(vllm-flash-attn)

60 message(STATUS "vllm-flash-attn is available at ${vllm-flash-attn_SOURCE_DIR}")

65 -# Restore the install prefix

66 -install(CODE "set(CMAKE_INSTALL_PREFIX \"\${OLD_CMAKE_INSTALL_PREFIX}\")" ALL_COMPONENTS)

67 -install(CODE "set(CMAKE_INSTALL_LOCAL_ONLY TRUE)" ALL_COMPONENTS)

62 +# Restore the install prefix after FA's install rules

63 +foreach(_FA_COMPONENT _vllm_fa2_C _vllm_fa3_C)

64 + install(CODE "set(CMAKE_INSTALL_PREFIX \"\${OLD_CMAKE_INSTALL_PREFIX}\")" COMPONENT ${_FA_COMPONENT})

65 + install(CODE "set(CMAKE_INSTALL_LOCAL_ONLY TRUE)" COMPONENT ${_FA_COMPONENT})

66 +endforeach()

67 +

68 +# Install shared Python files for both FA2 and FA3 components

69 +foreach(_FA_COMPONENT _vllm_fa2_C _vllm_fa3_C)

70 + # Ensure the vllm/vllm_flash_attn directory exists before installation

71 + install(CODE "file(MAKE_DIRECTORY \"\${CMAKE_INSTALL_PREFIX}/vllm/vllm_flash_attn\")"

72 + COMPONENT ${_FA_COMPONENT})

73 +

74 + # Copy vllm_flash_attn python files (except __init__.py and flash_attn_interface.py

75 + # which are source-controlled in vllm)

76 + install(

77 + DIRECTORY ${vllm-flash-attn_SOURCE_DIR}/vllm_flash_attn/

78 + DESTINATION vllm/vllm_flash_attn

79 + COMPONENT ${_FA_COMPONENT}

80 + FILES_MATCHING PATTERN "*.py"

81 + PATTERN "__init__.py" EXCLUDE

82 + PATTERN "flash_attn_interface.py" EXCLUDE

83 + )

84 +

85 +endforeach()

69 -# Copy over the vllm-flash-attn python files (duplicated for fa2 and fa3, in

70 -# case only one is built, in the case both are built redundant work is done)

71 -install(

72 - DIRECTORY ${vllm-flash-attn_SOURCE_DIR}/vllm_flash_attn/

73 - DESTINATION vllm/vllm_flash_attn

74 - COMPONENT _vllm_fa2_C

75 - FILES_MATCHING PATTERN "*.py"

76 -)

87 +#

88 +# FA4 CuteDSL component

89 +# This is a Python-only component that copies the flash_attn/cute directory

90 +# and transforms imports to match our package structure.

91 +#

92 +add_custom_target(_vllm_fa4_cutedsl_C)

78 -install(

79 - DIRECTORY ${vllm-flash-attn_SOURCE_DIR}/vllm_flash_attn/

80 - DESTINATION vllm/vllm_flash_attn

81 - COMPONENT _vllm_fa3_C

82 - FILES_MATCHING PATTERN "*.py"

83 -)

94 +# Copy flash_attn/cute directory (needed for FA4) and transform imports

95 +# The cute directory uses flash_attn.cute imports internally, which we replace

96 +# with vllm.vllm_flash_attn.cute to match our package structure.

97 +install(CODE "

98 + file(GLOB_RECURSE CUTE_PY_FILES \"${vllm-flash-attn_SOURCE_DIR}/flash_attn/cute/*.py\")

99 + foreach(SRC_FILE \${CUTE_PY_FILES})

100 + file(RELATIVE_PATH REL_PATH \"${vllm-flash-attn_SOURCE_DIR}/flash_attn/cute\" \${SRC_FILE})

101 + set(DST_FILE \"\${CMAKE_INSTALL_PREFIX}/vllm/vllm_flash_attn/cute/\${REL_PATH}\")

102 + get_filename_component(DST_DIR \${DST_FILE} DIRECTORY)

103 + file(MAKE_DIRECTORY \${DST_DIR})

104 + file(READ \${SRC_FILE} FILE_CONTENTS)

105 + string(REPLACE \"flash_attn.cute\" \"vllm.vllm_flash_attn.cute\" FILE_CONTENTS \"\${FILE_CONTENTS}\")

106 + file(WRITE \${DST_FILE} \"\${FILE_CONTENTS}\")

107 + endforeach()

108 +" COMPONENT _vllm_fa4_cutedsl_C)

docs/design/attention_backends.md

+2 -1

@@ -168,6 +168,7 @@ Priority is **1 = highest** (tried first).

168 | `FLASHINFER` | TRTLLM† | fp16, bf16 | `auto`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | 16, 32, 64 | 64, 128, 256 | ✅ | ❌ | ✅ | Decoder | 10.x |

169 | `FLASH_ATTN` | FA2* | fp16, bf16 | `auto`, `bfloat16` | %16 | Any | ❌ | ❌ | ✅ | All | ≥8.0 |

170 | `FLASH_ATTN` | FA3* | fp16, bf16 | `auto`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | %16 | Any | ✅ | ❌ | ✅ | All | 9.x |

171 +| `FLASH_ATTN` | FA4* | fp16, bf16 | `auto`, `bfloat16` | %16 | Any | ❌ | ❌ | ✅ | All | ≥10.0 |

172 | `FLASH_ATTN_DIFFKV` | | fp16, bf16 | `auto` | Any | Any | ❌ | ❌ | ✅ | Decoder | Any |

173 | `FLEX_ATTENTION` | | fp16, bf16, fp32 | `auto`, `bfloat16` | Any | Any | ❌ | ✅ | ❌ | Decoder, Encoder Only | Any |

174 | `ROCM_AITER_FA` | | fp16, bf16 | `auto` | 16, 32 | 64, 128, 256 | ❌ | ❌ | ❌ | Decoder | N/A |

@@ -178,7 +179,7 @@ Priority is **1 = highest** (tried first).

179

180 > **†** FlashInfer uses TRTLLM attention on Blackwell (SM100), which supports sinks. Disable via `--attention-config.use_trtllm_attention=0`.

181 >

181 -> **\*** Specify the FlashAttention version via `--attention-config.flash_attn_version=2` or `3`. Default is FA3 on SM90, FA2 otherwise.

182 +> **\*** Specify the FlashAttention version via `--attention-config.flash_attn_version=2`, `3`, or `4`. Default is FA4 on SM100+ (Blackwell), FA3 on SM90 (Hopper), FA2 otherwise.

183

184 ## MLA (Multi-head Latent Attention) Backends

185

requirements/cuda.txt

@@ -11,3 +11,7 @@ torchaudio==2.10.0

11 torchvision==0.25.0 # Required for phi3v processor. See https://github.com/pytorch/vision?tab=readme-ov-file#installation for corresponding version

12 # FlashInfer should be updated together with the Dockerfile

13 flashinfer-python==0.6.4

14 +

15 +# QuACK and Cutlass DSL for FA4 (cute-DSL implementation)

16 +nvidia-cutlass-dsl>=4.4.0.dev1

17 +quack-kernels>=0.2.7

setup.py

@@ -976,6 +976,11 @@ if _is_cuda():

976 ):

977 # FA3 requires CUDA 12.3 or later

978 ext_modules.append(CMakeExtension(name="vllm.vllm_flash_attn._vllm_fa3_C"))

979 + # FA4 CuteDSL - Python-only component for FA4's cute DSL support

980 + # Optional since this doesn't produce a .so file, just copies Python files

981 + ext_modules.append(

982 + CMakeExtension(name="vllm.vllm_flash_attn._vllm_fa4_cutedsl_C", optional=True)

983 + )

984 if envs.VLLM_USE_PRECOMPILED or (

985 CUDA_HOME and get_nvcc_cuda_version() >= Version("12.9")

986 ):

tools/pre_commit/generate_attention_backend_docs.py

+98 -10

@@ -563,14 +563,53 @@ def analyze_backend(backend_name: str, class_path: str) -> dict[str, Any] | None

563

564

565 # ---------------------------------------------------------------------------

566 -# Special backend variant parsers (FA2/FA3, FlashInfer TRTLLM, MLA prefill)

566 +# Special backend variant parsers (FA2/FA3/FA4, FlashInfer TRTLLM, MLA prefill)

567 # ---------------------------------------------------------------------------

568

569

570 +def _parse_fa4_supported_caps() -> str | None:

571 + """Parse flash_attn_interface.py for FA4 supported compute capabilities.

572 +

573 + Looks for `cc not in [9, 10, 11]` pattern in _is_fa4_supported().

574 + """

575 + fa_interface_file = (

576 + REPO_ROOT / "vllm" / "vllm_flash_attn" / "flash_attn_interface.py"

577 + )

578 + if not fa_interface_file.exists():

579 + return None

580 +

581 + try:

582 + tree = ast.parse(fa_interface_file.read_text())

583 + except Exception:

584 + return None

585 +

586 + for node in ast.walk(tree):

587 + if not isinstance(node, ast.FunctionDef) or node.name != "_is_fa4_supported":

588 + continue

589 + for n in ast.walk(node):

590 + if not (

591 + isinstance(n, ast.Compare)

592 + and len(n.ops) == 1

593 + and isinstance(n.ops[0], ast.NotIn)

594 + and isinstance(n.comparators[0], ast.List)

595 + ):

596 + continue

597 + caps: list[int] = [

598 + e.value

599 + for e in n.comparators[0].elts

600 + if isinstance(e, ast.Constant) and isinstance(e.value, int)

601 + ]

602 + if caps:

603 + caps.sort()

604 + return f"{caps[0]}.x-{caps[-1]}.x"

605 +

606 + return None

607 +

608 +

609 def parse_flash_attn_features() -> dict[str, dict[str, Any]]:

571 - """Parse fa_utils.py to detect FA2 vs FA3 feature differences.

610 + """Parse fa_utils.py to detect FA2 vs FA3 vs FA4 feature differences.

611

573 - Returns a dict with 'fa2' and 'fa3' keys containing their respective

612 + Returns a dict with 'fa2', 'fa3', and 'fa4' keys containing their respective

613 feature overrides for compute capability, KV cache dtypes, and sink support.

614 """

615 if not FA_UTILS_FILE.exists():

@@ -585,6 +624,7 @@ def parse_flash_attn_features() -> dict[str, dict[str, Any]]:

624 fa3_supports_fp8 = False

625 fa3_supports_sinks = False

626 fa3_compute_cap: str | None = None

627 + fa4_compute_cap: str | None = None

628

629 for node in ast.walk(tree):

630 if not isinstance(node, ast.FunctionDef):

@@ -614,14 +654,12 @@ def parse_flash_attn_features() -> dict[str, dict[str, Any]]:

654 fa3_supports_sinks = True

655 break

656

617 - # Check get_flash_attn_version for FA3 compute capability

618 - # Look for the ternary: 3 if (device_capability.major == 9 ...) else 2

657 + # Check get_flash_attn_version for FA3/FA4 compute capability

658 if node.name == "get_flash_attn_version":

659 for n in ast.walk(node):

621 - # Look for IfExp (ternary) with `device_capability.major == 9`

660 + # Handle IfExp (ternary) with `device_capability.major == 9`

661 if isinstance(n, ast.IfExp):

662 test = n.test

624 - # Check if test is a BoolOp (and) containing the major check

663 if isinstance(test, ast.BoolOp):

664 for val in test.values:

665 if (

@@ -634,6 +672,38 @@ def parse_flash_attn_features() -> dict[str, dict[str, Any]]:

672 fa3_compute_cap = f"{val.comparators[0].value}.x"

673 break

674

675 + # Handle If statements for FA3/FA4 detection

676 + # e.g. `if device_capability.major == 9` -> FA3

677 + # `elif device_capability.major >= 10` -> FA4

678 + if isinstance(n, ast.If):

679 + test = n.test

680 + comparisons = (

681 + [v for v in test.values if isinstance(v, ast.Compare)]

682 + if isinstance(test, ast.BoolOp)

683 + else [test]

684 + if isinstance(test, ast.Compare)

685 + else []

686 + )

687 + for comp in comparisons:

688 + if not (

689 + isinstance(comp.left, ast.Attribute)

690 + and comp.left.attr == "major"

691 + and comp.comparators

692 + and isinstance(comp.comparators[0], ast.Constant)

693 + and isinstance(comp.comparators[0].value, int)

694 + ):

695 + continue

696 + op = comp.ops[0]

697 + val = comp.comparators[0].value

698 + if isinstance(op, ast.Eq) and fa3_compute_cap is None:

699 + fa3_compute_cap = f"{val}.x"

700 + elif isinstance(op, ast.GtE) and fa4_compute_cap is None:

701 + fa4_compute_cap = f"≥{val}.0"

702 +

703 + # Fallback: try to parse FA4 compute caps from flash_attn_interface.py

704 + if fa4_compute_cap is None:

705 + fa4_compute_cap = _parse_fa4_supported_caps()

706 +

707 return {

708 "fa2": {

709 "supports_fp8": False,

@@ -644,6 +714,11 @@ def parse_flash_attn_features() -> dict[str, dict[str, Any]]:

714 "supports_fp8": fa3_supports_fp8,

715 "supports_sink": fa3_supports_sinks,

716 },

717 + "fa4": {

718 + "compute_capability": fa4_compute_cap,

719 + "supports_fp8": False,

720 + "supports_sink": False,

721 + },

722 }

723

724

@@ -760,7 +835,7 @@ def parse_mla_prefill_backends() -> list[dict[str, Any]]:

835

836

837 # ---------------------------------------------------------------------------

763 -# Backend variant expansion (FA2/FA3, FlashInfer native/TRTLLM)

838 +# Backend variant expansion (FA2/FA3/FA4, FlashInfer native/TRTLLM)

839 # ---------------------------------------------------------------------------

840

841

@@ -768,7 +843,7 @@ def _expand_flash_attn_variants(

843 all_backends: list[dict[str, Any]],

844 fa_features: dict[str, dict[str, Any]],

845 ) -> list[dict[str, Any]]:

771 - """Expand FLASH_ATTN into FA2 and FA3 variants with different capabilities."""

846 + """Expand FLASH_ATTN into FA2, FA3, and FA4 variants."""

847 expanded = []

848 for backend in all_backends:

849 if backend["name"] != "FLASH_ATTN":

@@ -801,6 +876,18 @@ def _expand_flash_attn_variants(

876

877 expanded.append(fa2)

878 expanded.append(fa3)

879 +

880 + # Create FA4 entry if FA4 features are available

881 + if "fa4" in fa_features:

882 + fa4 = backend.copy()

883 + fa4["version"] = "FA4*"

884 + fa4["_sort_key"] = "FLASH_ATTN"

885 + fa4["_sort_order"] = 2

886 + if fa_features["fa4"].get("compute_capability"):

887 + fa4["compute_capability"] = fa_features["fa4"]["compute_capability"]

888 + fa4["supports_sink"] = fa_features["fa4"]["supports_sink"]

889 + expanded.append(fa4)

890 +

891 return expanded

892

893

@@ -1360,7 +1447,8 @@ def generate_docs() -> str:

1447 if fa_features:

1448 footnotes.append(

1449 "> **\\*** Specify the FlashAttention version via "

1363 - "`--attention-config.flash_attn_version=2` or `3`. Default is FA3 on SM90, "

1450 + "`--attention-config.flash_attn_version=2`, `3`, or `4`. "

1451 + "Default is FA4 on SM100+ (Blackwell), FA3 on SM90 (Hopper), "

1452 "FA2 otherwise."

1453 )

1454 if footnotes:

vllm/config/attention.py

+2 -2

@@ -16,8 +16,8 @@ class AttentionConfig:

16 backend: AttentionBackendEnum | None = None

17 """Attention backend to use. If None, will be selected automatically."""

19 - flash_attn_version: Literal[2, 3] | None = None

20 - """Force vllm to use a specific flash-attention version (2 or 3).

19 + flash_attn_version: Literal[2, 3, 4] | None = None

20 + """Force vllm to use a specific flash-attention version (2, 3, or 4).

21 Only valid when using the flash-attention backend."""

23 use_prefill_decode_attention: bool = False

vllm/model_executor/layers/attention/mla_attention.py

+3 -1

@@ -2014,7 +2014,9 @@ class MLACommonImpl(MLAAttentionImpl[M], Generic[M]):

2014 # RoCM and the latter has an additional parameter to control

2015 # FA2 vs FA3

2016 self.flash_attn_varlen_func = flash_attn_varlen_func

2017 - self.vllm_flash_attn_version = get_flash_attn_version()

2017 + self.vllm_flash_attn_version = get_flash_attn_version(

2018 + head_size=self.qk_head_dim

2019 + )

2020 if self.vllm_flash_attn_version is not None:

2021 self.flash_attn_varlen_func = functools.partial(

2022 flash_attn_varlen_func, fa_version=self.vllm_flash_attn_version

vllm/model_executor/layers/attention/mm_encoder_attention.py

+3 -1

@@ -204,7 +204,9 @@ class MMEncoderAttention(CustomOp):

204 }

205

206 self._fa_version = (

207 - get_flash_attn_version() if self.is_flash_attn_backend else None

207 + get_flash_attn_version(head_size=head_size)

208 + if self.is_flash_attn_backend

209 + else None

210 )

211

212 if self.attn_backend == AttentionBackendEnum.FLASHINFER:

vllm/v1/attention/backends/fa_utils.py

+41 -7

@@ -52,7 +52,9 @@ elif current_platform.is_rocm():

52 reshape_and_cache_flash = ops.reshape_and_cache_flash

55 -def get_flash_attn_version(requires_alibi: bool = False) -> int | None:

55 +def get_flash_attn_version(

56 + requires_alibi: bool = False, head_size: int | None = None

57 +) -> int | None:

58 # import here to avoid circular dependencies

59 from vllm.platforms import current_platform

@@ -72,9 +74,15 @@ def get_flash_attn_version(requires_alibi: bool = False) -> int | None:

74 assert device_capability is not None

76 # 1. default version depending on platform

75 - fa_version = (

76 - 3 if (device_capability.major == 9 and is_fa_version_supported(3)) else 2

77 - )

77 + if device_capability.major == 9 and is_fa_version_supported(3):

78 + # Hopper (SM90): prefer FA3

79 + fa_version = 3

80 + elif device_capability.major == 10 and is_fa_version_supported(4):

81 + # Blackwell (SM100+, restrict to SM100 for now): prefer FA4

82 + fa_version = 4

83 + else:

84 + # Fallback to FA2

85 + fa_version = 2

87 # 2. override if passed by environment or config

88 from vllm.config import get_current_vllm_config_or_none

@@ -87,12 +95,12 @@ def get_flash_attn_version(requires_alibi: bool = False) -> int | None:

95 fa_version = vllm_config.attention_config.flash_attn_version

97 # 3. fallback for unsupported combinations

90 - if device_capability.major == 10 and fa_version == 3:

98 + if device_capability.major >= 10 and fa_version == 3:

99 logger.warning_once(

100 "Cannot use FA version 3 on Blackwell platform, "

93 - "defaulting to FA version 2."

101 + "defaulting to FA version 4 if supported, otherwise FA2."

102 )

95 - fa_version = 2

103 + fa_version = 4 if is_fa_version_supported(4) else 2

104

105 if requires_alibi and fa_version == 3:

106 logger.warning_once(

@@ -100,6 +108,28 @@ def get_flash_attn_version(requires_alibi: bool = False) -> int | None:

108 )

109 fa_version = 2

110

111 + if requires_alibi and fa_version == 4:

112 + logger.warning_once(

113 + "Cannot use FA version 4 with ALiBi, defaulting to FA version 2."

114 + )

115 + fa_version = 2

116 +

117 + # FA4 on SM100 (Blackwell) has TMEM capacity limits that restrict

118 + # supported head dimensions.

119 + # See: https://github.com/Dao-AILab/flash-attention/issues/1959

120 + if (

121 + fa_version == 4

122 + and device_capability.major >= 10

123 + and head_size is not None

124 + and head_size > 128

125 + ):

126 + logger.warning_once(

127 + "FA4 on Blackwell does not support head_size=%d due to TMEM "

128 + "capacity limits, defaulting to FA version 2.",

129 + head_size,

130 + )

131 + fa_version = 2

132 +

133 if not is_fa_version_supported(fa_version):

134 logger.error(

135 "Cannot use FA version %d is not supported due to %s",

@@ -139,6 +169,10 @@ def flash_attn_supports_mla():

169 return is_fa_version_supported(

170 3

171 ) and current_platform.is_device_capability_family(90)

172 +

173 + # NOTE(Lucas): FA4 CuteDSL does NOT currently support MLA's non-standard

174 + # head dimensions (576 for qk, 512 for v) due to TMEM capacity limits.

175 +

176 except (ImportError, AssertionError):

177 pass

178 return False

vllm/v1/attention/backends/flash_attn.py

+9 -1

@@ -580,7 +580,15 @@ class FlashAttentionImpl(AttentionImpl):

580 self.num_queries_per_kv = self.num_heads // self.num_kv_heads

581

582 self.attn_type = attn_type

583 - self.vllm_flash_attn_version = get_flash_attn_version()

583 + self.vllm_flash_attn_version = get_flash_attn_version(

584 + requires_alibi=alibi_slopes is not None,

585 + head_size=head_size,

586 + )

587 + logger.info_once(

588 + "Using FlashAttention version %s",

589 + self.vllm_flash_attn_version,

590 + scope="local",

591 + )

592 # Cache the batch invariant result for use in forward passes

593 self.batch_invariant_enabled = vllm_is_batch_invariant()

594

vllm/v1/cudagraph_dispatcher.py

+1 -1

@@ -137,7 +137,7 @@ class CudagraphDispatcher:

137 num_tokens_padded = self._bs_to_padded_graph_size[num_tokens]

138

139 if uniform_decode and self.cudagraph_mode.has_mode(CUDAGraphMode.FULL):

140 - num_reqs = num_tokens_padded // uniform_decode_query_len

140 + num_reqs = min(num_tokens_padded // uniform_decode_query_len, max_num_seqs)

141 assert num_tokens_padded % uniform_decode_query_len == 0

142 else:

143 uniform_decode = False

vllm/vllm_flash_attn/__init__.py new

+24

@@ -0,0 +1,24 @@

1 +# SPDX-License-Identifier: Apache-2.0

2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 +

4 +from vllm.vllm_flash_attn.flash_attn_interface import (

5 + FA2_AVAILABLE,

6 + FA3_AVAILABLE,

7 + fa_version_unsupported_reason,

8 + flash_attn_varlen_func,

9 + get_scheduler_metadata,

10 + is_fa_version_supported,

11 +)

12 +

13 +if not (FA2_AVAILABLE or FA3_AVAILABLE):

14 + raise ImportError(

15 + "vllm.vllm_flash_attn requires the CUDA flash attention extensions "

16 + "(_vllm_fa2_C or _vllm_fa3_C). On ROCm, use upstream flash_attn."

17 + )

18 +

19 +__all__ = [

20 + "fa_version_unsupported_reason",

21 + "flash_attn_varlen_func",

22 + "get_scheduler_metadata",

23 + "is_fa_version_supported",

24 +]

vllm/vllm_flash_attn/flash_attn_interface.py new

+567

@@ -0,0 +1,567 @@

1 +# SPDX-License-Identifier: Apache-2.0

2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

4 +# ruff: noqa: E501

5 +

6 +

7 +import torch

8 +

9 +# isort: off

10 +# We need to import the CUDA kernels after importing torch

11 +# Use relative import to support build-from-source installation in vLLM

12 +

13 +try:

14 + from . import _vllm_fa2_C # type: ignore[attr-defined] # noqa: F401

15 +

16 + FA2_UNAVAILABLE_REASON = None

17 + FA2_AVAILABLE = True

18 +except ImportError as e:

19 + FA2_UNAVAILABLE_REASON = str(e)

20 + FA2_AVAILABLE = False

21 +

22 +try:

23 + from . import _vllm_fa3_C # type: ignore[attr-defined] # noqa: F401

24 +

25 + FA3_UNAVAILABLE_REASON = None

26 + FA3_AVAILABLE = True

27 +except ImportError as e:

28 + FA3_UNAVAILABLE_REASON = str(e)

29 + FA3_AVAILABLE = False

30 +

31 +

32 +try:

33 + import os

34 +

35 + _cute_interface_path = os.path.join(

36 + os.path.dirname(__file__), "cute", "interface.py"

37 + )

38 + if not os.path.exists(_cute_interface_path):

39 + raise ImportError("vllm.vllm_flash_attn.cute.interface not found")

40 +

41 + FA4_UNAVAILABLE_REASON = None

42 + FA4_AVAILABLE = True

43 +except (ImportError, ModuleNotFoundError) as e:

44 + FA4_UNAVAILABLE_REASON = str(e)

45 + FA4_AVAILABLE = False

46 +

47 +# isort: on

48 +

49 +DEFAULT_FA_VERSION = 2

50 +

51 +

52 +def _is_fa2_supported() -> tuple[bool, str | None]:

53 + if not FA2_AVAILABLE:

54 + return False, f"FA2 is unavailable due to: {FA2_UNAVAILABLE_REASON}"

55 + from vllm.platforms import current_platform

56 +

57 + if not current_platform.has_device_capability(80):

58 + return False, "FA2 is only supported on devices with compute capability >= 8"

59 + return True, None

60 +

61 +

62 +def _is_fa3_supported() -> tuple[bool, str | None]:

63 + if not FA3_AVAILABLE:

64 + return False, f"FA3 is unavailable due to: {FA3_UNAVAILABLE_REASON}"

65 + from vllm.platforms import current_platform

66 +

67 + if not current_platform.is_device_capability_family(90):

68 + return False, "FA3 is only supported on devices with compute capability 9.x"

69 + return True, None

70 +

71 +

72 +def _is_fa4_supported() -> tuple[bool, str | None]:

73 + if not FA4_AVAILABLE:

74 + return False, f"FA4 is unavailable due to: {FA4_UNAVAILABLE_REASON}"

75 + from vllm.platforms import current_platform

76 +

77 + if not (

78 + current_platform.is_device_capability_family(90)

79 + or current_platform.is_device_capability_family(100)

80 + or current_platform.is_device_capability_family(110)

81 + ):

82 + return (

83 + False,

84 + "FA4 is only supported on devices with compute capability 9.x, 10.x, or 11.x",

85 + )

86 + return True, None

87 +

88 +

89 +def is_fa_version_supported(fa_version: int) -> bool:

90 + if fa_version == 2:

91 + return _is_fa2_supported()[0]

92 + elif fa_version == 3:

93 + return _is_fa3_supported()[0]

94 + elif fa_version == 4:

95 + return _is_fa4_supported()[0]

96 + else:

97 + raise ValueError(f"Unsupported FA version: {fa_version}")

98 +

99 +

100 +def fa_version_unsupported_reason(fa_version: int) -> str | None:

101 + if fa_version == 2:

102 + return _is_fa2_supported()[1]

103 + elif fa_version == 3:

104 + return _is_fa3_supported()[1]

105 + elif fa_version == 4:

106 + return _is_fa4_supported()[1]

107 + else:

108 + raise ValueError(f"Unsupported FA version: {fa_version}")

109 +

110 +

111 +#

112 +# For vLLM we only care about `flash_attn_varlen_func` and

113 +# `flash_attn_with_kvcache` so we only maintain wrappers for these two.

114 +#

115 +

116 +

117 +def maybe_contiguous(x):

118 + return x.contiguous() if x is not None and x.stride(-1) != 1 else x

119 +

120 +

121 +# NOTE only used in FA3

122 +def get_scheduler_metadata(

123 + batch_size,

124 + max_seqlen_q,

125 + max_seqlen_k,

126 + num_heads_q,

127 + num_heads_kv,

128 + headdim,

129 + cache_seqlens: torch.Tensor,

130 + qkv_dtype=torch.bfloat16,

131 + headdim_v=None,

132 + cu_seqlens_q: torch.Tensor | None = None,

133 + cu_seqlens_k_new: torch.Tensor | None = None,

134 + cache_leftpad: torch.Tensor | None = None,

135 + page_size: int | None = None,

136 + max_seqlen_k_new=0,

137 + causal=False,

138 + window_size=(-1, -1), # -1 means infinite context window

139 + has_softcap=False,

140 + num_splits=0, # Can be tuned for speed

141 + pack_gqa=None, # Can be tuned for speed

142 + sm_margin=0, # Can be tuned if some SMs are used for communication

143 +):

144 + cache_seqlens = maybe_contiguous(cache_seqlens)

145 + if headdim_v is None:

146 + headdim_v = headdim

147 + scheduler_metadata = torch.ops._vllm_fa3_C.get_scheduler_metadata(

148 + batch_size,

149 + max_seqlen_q,

150 + max_seqlen_k,

151 + num_heads_q,

152 + num_heads_kv,

153 + headdim,

154 + headdim_v,

155 + qkv_dtype,

156 + cache_seqlens,

157 + cu_seqlens_q,

158 + None, # cu_seqlens_k

159 + cu_seqlens_k_new,

160 + None, # seqused_q

161 + cache_leftpad,

162 + page_size,

163 + max_seqlen_k_new,

164 + causal,

165 + window_size[0],

166 + window_size[1],

167 + has_softcap,

168 + num_splits,

169 + pack_gqa,

170 + sm_margin,

171 + )

172 +

173 + return scheduler_metadata

174 +

175 +

176 +def flash_attn_varlen_func(

177 + q,

178 + k,

179 + v,

180 + max_seqlen_q,

181 + cu_seqlens_q,

182 + max_seqlen_k,

183 + cu_seqlens_k=None, # only used for non-paged prefill

184 + seqused_k=None,

185 + q_v=None,

186 + dropout_p=0.0,

187 + softmax_scale=None,

188 + causal=False,

189 + window_size: list[int] | None = None,

190 + softcap=0.0, # 0.0 means deactivated

191 + alibi_slopes=None,

192 + deterministic=False,

193 + return_attn_probs=False,

194 + block_table=None,

195 + return_softmax_lse=False,

196 + out=None,

197 + # FA3 Only

198 + scheduler_metadata=None,

199 + q_descale=None,

200 + k_descale=None,

201 + v_descale=None,

202 + num_splits: int = 0,

203 + # Version selector

204 + fa_version: int = DEFAULT_FA_VERSION,

205 + s_aux=None,

206 + cp_world_size=1,

207 + cp_rank=0,

208 + cp_tot_seqused_k=None,

209 +):

210 + """dropout_p should be set to 0.0 during evaluation

211 + Supports multi-query and grouped-query attention (MQA/GQA) by passing in K, V with fewer heads

212 + than Q. Note that the number of heads in Q must be divisible by the number of heads in KV.

213 + For example, if Q has 6 heads and K, V have 2 heads, head 0, 1, 2 of Q will attention to head

214 + 0 of K, V, and head 3, 4, 5 of Q will attention to head 1 of K, V.

215 +

216 + If causal=True, the causal mask is aligned to the bottom right corner of the attention matrix.

217 + For example, if seqlen_q = 2 and seqlen_k = 5, the causal mask (1 = keep, 0 = masked out) is:

218 + 1 1 1 1 0

219 + 1 1 1 1 1

220 + If seqlen_q = 5 and seqlen_k = 2, the causal mask is:

221 + 0 0

222 + 0 0

223 + 0 0

224 + 1 0

225 + 1 1

226 + If the row of the mask is all zero, the output will be zero.

227 +

228 + If window_size != (-1, -1), implements sliding window local attention. Query at position i

229 + will only attend to keys between

230 + [i + seqlen_k - seqlen_q - window_size[0], i + seqlen_k - seqlen_q + window_size[1]] inclusive.

231 +

232 + Arguments:

233 + q: (total_q, nheads, headdim), where total_q = total number of query tokens in the batch.

234 + k: (total_k, nheads_k, headdim), where total_k = total number of key tokens in the batch.

235 + v: (total_k, nheads_k, headdim), where total_k = total number of key tokens in the batch.

236 + cu_seqlens_q: (batch_size + 1,), dtype torch.int32. The cumulative sequence lengths

237 + of the sequences in the batch, used to index into q.

238 + cu_seqlens_k: (batch_size + 1,), dtype torch.int32. The cumulative sequence lengths

239 + of the sequences in the batch, used to index into kv.

240 + max_seqlen_q: int. Maximum query sequence length in the batch.

241 + max_seqlen_k: int. Maximum key sequence length in the batch.

242 + dropout_p: float. Dropout probability.

243 + softmax_scale: float. The scaling of QK^T before applying softmax.

244 + Default to 1 / sqrt(headdim).

245 + causal: bool. Whether to apply causal attention mask (e.g., for auto-regressive modeling).

246 + window_size: (left, right). If not (-1, -1), implements sliding window local attention.

247 + softcap: float. Anything > 0 activates softcapping attention.

248 + alibi_slopes: (nheads,) or (batch_size, nheads), fp32. A bias of

249 + (-alibi_slope * |i + seqlen_k - seqlen_q - j|)

250 + is added to the attention score of query i and key j.

251 + deterministic: bool. Whether to use the deterministic implementation of the backward pass,

252 + which is slightly slower and uses more memory. The forward pass is always deterministic.

253 + return_attn_probs: bool. Whether to return the attention probabilities. This option is for

254 + testing only. The returned probabilities are not guaranteed to be correct

255 + (they might not have the right scaling).

256 + Return:

257 + out: (total, nheads, headdim).

258 + softmax_lse [optional, if return_softmax_lse=True]: (nheads, total_q_seqlen). The

259 + logsumexp of each row of the matrix QK^T * scaling (e.g., log of the softmax

260 + normalization factor).

261 + """

262 + assert cu_seqlens_k is not None or seqused_k is not None, (

263 + "cu_seqlens_k or seqused_k must be provided"

264 + )

265 + assert cu_seqlens_k is None or seqused_k is None, (

266 + "cu_seqlens_k and seqused_k cannot be provided at the same time"

267 + )

268 + assert block_table is None or seqused_k is not None, (

269 + "seqused_k must be provided if block_table is provided"

270 + )

271 +

272 + if softmax_scale is None:

273 + softmax_scale = q.shape[-1] ** (-0.5)

274 + # custom op does not support non-tuple input

275 + real_window_size: tuple[int, int]

276 + if window_size is None:

277 + real_window_size = (-1, -1)

278 + else:

279 + assert len(window_size) == 2

280 + real_window_size = (window_size[0], window_size[1])

281 + q, k, v = [maybe_contiguous(x) for x in (q, k, v)]

282 +

283 + dummy_cu_seqlens_k = torch.empty_like(cu_seqlens_q)

284 +

285 + if fa_version == 2:

286 + if (

287 + scheduler_metadata is not None

288 + and q_descale is not None

289 + and k_descale is not None

290 + and v_descale is not None

291 + ):

292 + raise NotImplementedError(

293 + "FA2 does not support scheduler_metadata, q_descale, "

294 + "k_descale, v_descale"

295 + )

296 + if s_aux is not None:

297 + raise NotImplementedError("FA2 does not support s_aux")

298 + if num_splits > 1:

299 + raise NotImplementedError("FA2 does not support num_splits > 1")

300 + out, softmax_lse = torch.ops._vllm_fa2_C.varlen_fwd(

301 + q,

302 + k,

303 + v,

304 + out,

305 + cu_seqlens_q,

306 + # cu_seqlens_k not used since we use seqused_k, but flash_api.cpp

307 + # still wants it so we pass all zeros

308 + dummy_cu_seqlens_k if cu_seqlens_k is None else cu_seqlens_k,

309 + seqused_k,

310 + None,

311 + block_table,

312 + alibi_slopes,

313 + max_seqlen_q,

314 + max_seqlen_k,

315 + dropout_p,

316 + softmax_scale,

317 + False,

318 + causal,

319 + real_window_size[0],

320 + real_window_size[1],

321 + softcap,

322 + return_softmax_lse and dropout_p > 0,

323 + num_splits,

324 + None,

325 + )

326 + elif fa_version == 3:

327 + assert alibi_slopes is None, "Alibi is not supported in FA3"

328 + out, softmax_lse, _, _ = torch.ops._vllm_fa3_C.fwd(

329 + q,

330 + k,

331 + v,

332 + None,

333 + None, # k_new, v_new

334 + q_v,

335 + out,

336 + cu_seqlens_q,

337 + cu_seqlens_k, # cu_seqlens_k

338 + None, # cu_seqlens_k_new

339 + None,

340 + seqused_k, # seqused_q, seqused_k

341 + max_seqlen_q,

342 + max_seqlen_k,

343 + block_table,

344 + None, # kv_batch_idx

345 + None, # leftpad_k

346 + None,

347 + None,

348 + None, # rotary_cos, rotary_sin, seqlens_rotary

349 + q_descale,

350 + k_descale,

351 + v_descale,

352 + softmax_scale,

353 + causal,

354 + real_window_size[0],

355 + real_window_size[1],

356 + softcap,

357 + True, # rotary_interleaved

358 + scheduler_metadata,

359 + num_splits,

360 + None, # pack_gqa

361 + 0, # sm_margin

362 + s_aux, # s_aux

363 + cp_world_size,

364 + cp_rank,

365 + cp_tot_seqused_k,

366 + )

367 + elif fa_version == 4:

368 + assert alibi_slopes is None, "Alibi is not supported in FA4"

369 + # FA4 on SM90 doesn't support paged KV; SM100+ does

370 + from vllm.platforms import current_platform

371 +

372 + if block_table is not None and current_platform.is_device_capability_family(90):

373 + raise NotImplementedError(

374 + "FA4 with paged KV is not supported on SM90 (Hopper). "

375 + "Use FA3 or upgrade to Blackwell (SM100+)."

376 + )

377 + from vllm.vllm_flash_attn.cute.interface import _flash_attn_fwd

378 +

379 + out, softmax_lse = _flash_attn_fwd(

380 + q,

381 + k,

382 + v,

383 + cu_seqlens_q=cu_seqlens_q,

384 + cu_seqlens_k=cu_seqlens_k,

385 + seqused_k=seqused_k,

386 + max_seqlen_q=max_seqlen_q,

387 + max_seqlen_k=max_seqlen_k,

388 + page_table=block_table,

389 + softmax_scale=softmax_scale,

390 + causal=causal,

391 + softcap=softcap,

392 + window_size_left=real_window_size[0] if real_window_size[0] >= 0 else None,

393 + window_size_right=real_window_size[1] if real_window_size[1] >= 0 else None,

394 + num_splits=num_splits,

395 + return_lse=return_softmax_lse,

396 + out=out,

397 + )

398 + else:

399 + raise ValueError(f"Unsupported FA version: {fa_version}")

400 + return (out, softmax_lse) if return_softmax_lse else out

401 +

402 +

403 +def sparse_attn_func(

404 + q,

405 + k,

406 + v,

407 + block_count,

408 + block_offset,

409 + column_count,

410 + column_index,

411 + dropout_p=0.0,

412 + softmax_scale=None,

413 + causal=False,

414 + softcap=0.0, # 0.0 means deactivated

415 + alibi_slopes=None,

416 + deterministic=False,

417 + return_attn_probs=False,

418 + *,

419 + return_softmax_lse=False,

420 + out=None,

421 +):

422 + """Compute attention with vertical and slash sparsity patterns.

423 + Most Arguments are the same with the flash_attn_func interface, except for 4 extra args:

424 + block_count and block_offset for slash sparsity patterns, and

425 + column_count and column_index for vertical sparsity patterns.

426 + For more details please refer to Appendix C.4.2 of paper https://arxiv.org/abs/2407.02490.

427 +

428 + Arguments:

429 + q: (batch_size, seqlen, nheads, headdim)

430 + k: (batch_size, seqlen, nheads_k, headdim)

431 + v: (batch_size, seqlen, nheads_k, headdim)

432 + block_count: (batch_size, nheads, cdiv(seqlen, BLOCK_M))

433 + block_offset: (batch_size, nheads, cdiv(seqlen, BLOCK_M), NNZ_S)

434 + column_count: (batch_size, nheads, cdiv(seqlen, BLOCK_M))

435 + column_index: (batch_size, nheads, cdiv(seqlen, BLOCK_M), NNZ_V)

436 + dropout_p: float. Dropout probability.

437 + softmax_scale: float. The scaling of QK^T before applying softmax.

438 + Default to 1 / sqrt(headdim).

439 + causal: bool. Whether to apply causal attention mask (e.g., for auto-regressive modeling).

440 + alibi_slopes: (nheads,) or (batch_size, nheads), fp32. A bias of

441 + (-alibi_slope * |i + seqlen_k - seqlen_q - j|)

442 + is added to the attention score of query i and key j.

443 + deterministic: bool. Whether to use the deterministic implementation of the backward pass,

444 + which is slightly slower and uses more memory. The forward pass is always deterministic.

445 + return_attn_probs: bool. Whether to return the attention probabilities. This option is for

446 + testing only. The returned probabilities are not guaranteed to be correct

447 + (they might not have the right scaling).

448 + Return:

449 + out: (batch_size, seqlen, nheads, headdim).

450 + softmax_lse [optional, if return_softmax_lse=True]: (batch_size, nheads, seqlen). The

451 + logsumexp of each row of the matrix QK^T * scaling (e.g., log of the softmax

452 + normalization factor).

453 + """

454 + if softmax_scale is None:

455 + softmax_scale = q.shape[-1] ** (-0.5)

456 +

457 + q, k, v = [maybe_contiguous(x) for x in (q, k, v)]

458 + out, softmax_lse = torch.ops._vllm_fa2_C.fwd_sparse(

459 + q,

460 + k,

461 + v,

462 + block_count,

463 + block_offset,

464 + column_count,

465 + column_index,

466 + out,

467 + alibi_slopes,

468 + dropout_p,

469 + softmax_scale,

470 + causal,

471 + softcap,

472 + return_attn_probs and dropout_p > 0,

473 + None,

474 + )

475 + return (out, softmax_lse) if return_softmax_lse else out

476 +

477 +

478 +def sparse_attn_varlen_func(

479 + q,

480 + k,

481 + v,

482 + block_count,

483 + block_offset,

484 + column_count,

485 + column_index,

486 + cu_seqlens_q,

487 + cu_seqlens_k,

488 + max_seqlen_q,

489 + max_seqlen_k,

490 + dropout_p=0.0,

491 + softmax_scale=None,

492 + causal=False,

493 + softcap=0.0, # 0.0 means deactivated

494 + alibi_slopes=None,

495 + deterministic=False,

496 + return_attn_probs=False,

497 + *,

498 + return_softmax_lse=False,

499 + out=None,

500 +):

501 + """Compute attention with vertical and slash sparsity patterns.

502 + Most Arguments are the same with the flash_attn_varlen_func interface, except for 4 extra args:

503 + block_count and block_offset for slash sparsity patterns, and

504 + column_count and column_index for vertical sparsity patterns.

505 + For more details please refer to Appendix C.4.2 of paper https://arxiv.org/abs/2407.02490.

506 +

507 + Arguments:

508 + q: (total_q, nheads, headdim), where total_q = total number of query tokens in the batch.

509 + k: (total_k, nheads_k, headdim), where total_k = total number of key tokens in the batch.

510 + v: (total_k, nheads_k, headdim), where total_k = total number of key tokens in the batch.

511 + block_count: (batch_size, nheads, cdiv(seqlen, BLOCK_M))

512 + block_offset: (batch_size, nheads, cdiv(seqlen, BLOCK_M), NNZ_S)

513 + column_count: (batch_size, nheads, cdiv(seqlen, BLOCK_M))

514 + column_index: (batch_size, nheads, cdiv(seqlen, BLOCK_M), NNZ_V)

515 + cu_seqlens_q: (batch_size + 1,), dtype torch.int32. The cumulative sequence lengths

516 + of the sequences in the batch, used to index into q.

517 + cu_seqlens_k: (batch_size + 1,), dtype torch.int32. The cumulative sequence lengths

518 + of the sequences in the batch, used to index into kv.

519 + max_seqlen_q: int. Maximum query sequence length in the batch.

520 + max_seqlen_k: int. Maximum key sequence length in the batch.

521 + dropout_p: float. Dropout probability.

522 + softmax_scale: float. The scaling of QK^T before applying softmax.

523 + Default to 1 / sqrt(headdim).

524 + causal: bool. Whether to apply causal attention mask (e.g., for auto-regressive modeling).

525 + softcap: float. Anything > 0 activates softcapping attention.

526 + alibi_slopes: (nheads,) or (batch_size, nheads), fp32. A bias of

527 + (-alibi_slope * |i + seqlen_k - seqlen_q - j|)

528 + is added to the attention score of query i and key j.

529 + deterministic: bool. Whether to use the deterministic implementation of the backward pass,

530 + which is slightly slower and uses more memory. The forward pass is always deterministic.

531 + return_attn_probs: bool. Whether to return the attention probabilities. This option is for

532 + testing only. The returned probabilities are not guaranteed to be correct

533 + (they might not have the right scaling).

534 + Return:

535 + out: (total, nheads, headdim).

536 + softmax_lse [optional, if return_softmax_lse=True]: (nheads, total_q_seqlen). The

537 + logsumexp of each row of the matrix QK^T * scaling (e.g., log of the softmax

538 + normalization factor).

539 + """

540 + if softmax_scale is None:

541 + softmax_scale = q.shape[-1] ** (-0.5)

542 +

543 + q, k, v = [maybe_contiguous(x) for x in (q, k, v)]

544 + out, softmax_lse = torch.ops._vllm_fa2_C.varlen_fwd_sparse(

545 + q,

546 + k,

547 + v,

548 + block_count,

549 + block_offset,

550 + column_count,

551 + column_index,

552 + out,

553 + cu_seqlens_q,

554 + cu_seqlens_k,

555 + None,

556 + alibi_slopes,

557 + max_seqlen_q,

558 + max_seqlen_k,

559 + dropout_p,

560 + softmax_scale,

561 + False,

562 + causal,

563 + softcap,

564 + return_attn_probs and dropout_p > 0,

565 + None,

566 + )

567 + return (out, softmax_lse) if return_softmax_lse else out