diff --git a/README.rst b/README.rst index 144bae37141..e5ed6ac971e 100644 --- a/README.rst +++ b/README.rst @@ -164,7 +164,7 @@ System Requirements * Compiler: GCC 9+ or Clang 10+ with C++17 support * Python: 3.12 recommended -* **Source Build Requirements:** CMake 3.18+, Ninja, Git 2.17+, pybind11 2.6.0+, nvidia-cudnn-frontend 1.25.0+ +* **Source Build Requirements:** CMake 3.18+, Ninja, Git 2.17+, pybind11 2.6.0+, nvidia-cudnn-frontend 1.27.0+ * **Notes:** FP8 features require Compute Capability 8.9+ (Ada/Hopper/Blackwell) diff --git a/build_tools/build_ext.py b/build_tools/build_ext.py index 4b7ec153be1..4542b0ce5d6 100644 --- a/build_tools/build_ext.py +++ b/build_tools/build_ext.py @@ -213,8 +213,9 @@ def run(self) -> None: shutil.rmtree(nccl_ep_dir) def build_extensions(self): - # For core lib + JAX install, fix build_ext from pybind11.setup_helpers - # to handle CUDA files correctly. + # For JAX builds, fix build_ext from pybind11.setup_helpers to handle + # CUDA files correctly. This is also required by the standalone JAX + # wheel, where framework_extension_only is true. if "pytorch" not in get_frameworks(): # Ensure at least an empty list of flags for 'cxx' and 'nvcc' when # extra_compile_args is a dict. @@ -226,8 +227,7 @@ def build_extensions(self): # Define new _compile method that redirects to NVCC for .cu and .cuh files. original_compile_fn = self.compiler._compile - if not framework_extension_only: - self.compiler.src_extensions += [".cu", ".cuh"] + self.compiler.src_extensions += [".cu", ".cuh"] def _compile_fn(obj, src, ext, cc_args, extra_postargs, pp_opts) -> None: # Copy before we make any modifications. @@ -236,10 +236,7 @@ def _compile_fn(obj, src, ext, cc_args, extra_postargs, pp_opts) -> None: try: original_compiler = self.compiler.compiler_so - if ( - os.path.splitext(src)[1] in [".cu", ".cuh"] - and not framework_extension_only - ): + if os.path.splitext(src)[1] in [".cu", ".cuh"]: nvcc_bin = nvcc_path() if nvcc_bin is None: raise RuntimeError( diff --git a/build_tools/jax.py b/build_tools/jax.py index 031432e6f90..bd89164aa24 100644 --- a/build_tools/jax.py +++ b/build_tools/jax.py @@ -5,27 +5,35 @@ """JAX related extensions.""" import os +from importlib.metadata import version as get_package_version from pathlib import Path -from packaging import version +from typing import List import setuptools +from packaging import version from .utils import ( - get_cuda_include_dirs, all_files_in_dir, cudnn_frontend_include_path, debug_build_enabled, - setup_mpi_flags, + get_cuda_include_dirs, + nccl_ep_enabled, nccl_include_path, nccl_lib_path, - nccl_ep_enabled, + setup_mpi_flags, ) -from typing import List def install_requirements() -> List[str]: """Install dependencies for TE/JAX extensions.""" - return ["jax", "flax>=0.7.1", "nvidia-cudnn-frontend>=1.25.0"] + # Serialized cuDNN graphs use a version-specific wire format, so the Python + # frontend used at runtime must match the headers used to build the extension. + frontend_version = get_package_version("nvidia-cudnn-frontend") + return [ + "jax", + "flax>=0.7.1", + f"nvidia-cudnn-frontend=={frontend_version}", + ] def test_requirements() -> List[str]: @@ -89,6 +97,7 @@ def setup_jax_extension( csrc_source_files = Path(csrc_source_files) extensions_dir = csrc_source_files / "extensions" sources = all_files_in_dir(extensions_dir, name_extension="cpp") + sources += all_files_in_dir(extensions_dir, name_extension="cu") # Header files include_dirs = get_cuda_include_dirs() diff --git a/build_tools/pytorch.py b/build_tools/pytorch.py index 918b5d941c4..da10c7e2713 100644 --- a/build_tools/pytorch.py +++ b/build_tools/pytorch.py @@ -31,6 +31,8 @@ def install_requirements() -> List[str]: "packaging", "pydantic", "nvdlfw-inspect", + # PyTorch cuDNN attention is built and executed with the Python graph + # API; FP8/MXFP8 graph capture requires Frontend 1.28 or newer. "nvidia-cudnn-frontend>=1.28.0", ] diff --git a/build_tools/wheel_utils/build_wheels.sh b/build_tools/wheel_utils/build_wheels.sh index 8fe2b629bed..42174ebe688 100644 --- a/build_tools/wheel_utils/build_wheels.sh +++ b/build_tools/wheel_utils/build_wheels.sh @@ -23,7 +23,7 @@ git checkout $TARGET_BRANCH git submodule update --init --recursive # Install deps -/opt/python/cp310-cp310/bin/pip install cmake pybind11[global] ninja setuptools wheel 'nvidia-cudnn-frontend>=1.25.0' +/opt/python/cp310-cp310/bin/pip install cmake pybind11[global] ninja setuptools wheel 'nvidia-cudnn-frontend>=1.27.0' if $BUILD_METAPACKAGE ; then cd /TransformerEngine diff --git a/docs/envvars.rst b/docs/envvars.rst index 0fee105fd05..e6433f5733a 100644 --- a/docs/envvars.rst +++ b/docs/envvars.rst @@ -200,7 +200,7 @@ backend-selection overview. :Type: ``int`` (0, 1 or 2), optionally followed by ``:`` :Default: ``0`` - :Description: Log FusedAttention graph cache activity to stderr, prefixed with ``[FUSED-ATTN-CACHE]``. ``1`` prints an end-of-run summary of the cache counters and the mean time of each cuDNN build stage. ``2`` additionally traces every event as it happens: each graph built, each graph cuDNN accepts and the cache keeps, each lookup and whether it hit or missed, each first execution that compiles kernels, and each execution. When the launcher exports a rank, only rank 0 logs; append ``:`` to override, as in ``1:all`` for level 1 on every rank or ``2:0,3`` for level 2 on ranks 0 and 3. + :Description: Log FusedAttention Python graph-cache activity to stderr, prefixed with ``[FUSED-ATTN-CACHE]``. ``1`` prints an end-of-run summary of cache counters and the mean CPU wall time of each cuDNN graph-build stage. ``2`` additionally traces every cache event, including the cache key on hits and misses. When the launcher exports a rank, only rank 0 logs by default; append ``:`` to override this, for example ``1:all`` or ``2:0,3``. Supported by PyTorch and JAX. .. envvar:: NVTE_ALLOW_NONDETERMINISTIC_ALGO diff --git a/docs/examples/attention/attention.ipynb b/docs/examples/attention/attention.ipynb index c10ce173917..c7d75daabd8 100644 --- a/docs/examples/attention/attention.ipynb +++ b/docs/examples/attention/attention.ipynb @@ -249,7 +249,7 @@ "NVTE_DEBUG = 0/1 # disables/enables debugging\n", "NVTE_DEBUG_LEVEL = 0/1/2 # enables logging.WARNING/INFO/DEBUG-level messages\n", "```\n", - "These flags are available in both PyTorch and Jax in TE v2.20 onwards." + "These flags are available in both PyTorch and JAX." ] }, { @@ -331,44 +331,12 @@ "id": "fc0c9e3d", "metadata": {}, "source": [ - "For the FusedAttention backend, cuDNN graphs are cached and reused, and the following flag helps understand the activities taking place in the cache. `NVTE_FUSED_ATTN_CACHE_DEBUG` is supported in both PyTorch and JAX.\n", + "For the FusedAttention backend, cuDNN Python graphs are cached and reused. `NVTE_FUSED_ATTN_CACHE_DEBUG` reports activity in these caches in both PyTorch and JAX.\n", "```\n", - "NVTE_FUSED_ATTN_CACHE_DEBUG = 0 # disables cache diagnostics\n", - "NVTE_FUSED_ATTN_CACHE_DEBUG = 1/2:ranks # enables cache diagnostics (events)/(tracing):select ranks for information collection\n", + "NVTE_FUSED_ATTN_CACHE_DEBUG = 0 # disable cache diagnostics\n", + "NVTE_FUSED_ATTN_CACHE_DEBUG = 1/2:ranks # summary/trace level and optional ranks\n", "```\n", - "At level 1, these cache events are collected and reported in the end-of-run summary: `create_graph` (number of graphs created), `cache_graph` (number of graphs cached), `hit` (number of cache hits), `miss` (number of cache misses), `build_plans` (number of graphs whose plans are built), and `execute` (number of plan executions). These event counters are tallied up per backend (f16 and fp8), per pass (forward and backward), and per thread. They help paint a picture of the graph reuse rate and caching effectiveness. Another part of the end-of-run summary are the build times of various stages of the cuDNN graph. Out of the 5 cuDNN calls, `validate`, `build_operation_graph`, `create_execution_plans`, `check_support`, and `build_plans`, `build_plans` is the most expensive one, because it is where the kernels get compiled. It happens once per eligible graph and is expected to execute many times in a real-life run.\n", - "\n", - "For more verbose diagnostics, level 2 enables tracing for all the events taking place in the cache. For example, for every `miss` event, a \"MISS\" line is printed in the logs, as well as the cache key that triggered it. This helps identify if the cache is processing the right config as we intended.\n", - "\n", - "By default, only rank 0 is enabled for both level 1 and level 2. Users can use the \"level:ranks\" format to specify a select set of ranks to collect diagnostics on. For example, \"2:0,3\" means to enable level 2 diagnostics on ranks 0 and 3. An example of the level 2 diagnostics, which also includes the level 1 summary, is as follows, for a `tests/pytorch/attention/test_attention.py::test_dot_product_attention` test.\n", - "```\n", - "[FUSED-ATTN-CACHE] rank=0 | tid=0 dev=0 | f16 fwd MISS | train=1 det=0 cg=0 maxlogit=0 mask=0 bias=0 wl=-1 wr=-1 brd=0 softmax=0 scale_mode=0 dropout=0 attn_scale=1 qkv_dt=5 o_dt=5 do_dt=6 dqkv_dt=6 qkv_lay=0 o_fmt=0 do_fmt=8 dqkv_lay=26 qkv_sif=8 do_sif=8 b=8 h=16 hg=16 dqk=64 dv=64 sq=128 skv=128 tq=0 tkv=0 bb=0 btq=0 btkv=0 npk=0 npv=0 psk=0 psv=0 mppk=0 mppv=0 bias_b=0 bias_h=0 bias_sq=0 bias_skv=0\n", - "[FUSED-ATTN-CACHE] rank=0 | tid=0 dev=0 | f16 fwd CREATE_GRAPH | hit= 0, miss= 1, create_graph= 1, cache_graph= 0, build_plans= 0, execute= 0\n", - "[FUSED-ATTN-CACHE] rank=0 | tid=0 dev=0 | f16 fwd CACHE_GRAPH | hit= 0, miss= 1, create_graph= 1, cache_graph= 1, build_plans= 0, execute= 0\n", - "...\n", - "[FUSED-ATTN-CACHE] rank=0 | tid=0 dev=0 | f16 fwd HIT | train=1 det=0 cg=0 maxlogit=0 mask=0 bias=0 wl=-1 wr=-1 brd=0 softmax=0 scale_mode=0 dropout=0 attn_scale=1 qkv_dt=5 o_dt=5 do_dt=6 dqkv_dt=6 qkv_lay=0 o_fmt=0 do_fmt=8 dqkv_lay=26 qkv_sif=8 do_sif=8 b=8 h=16 hg=16 dqk=64 dv=64 sq=128 skv=128 tq=0 tkv=0 bb=0 btq=0 btkv=0 npk=0 npv=0 psk=0 psv=0 mppk=0 mppv=0 bias_b=0 bias_h=0 bias_sq=0 bias_skv=0\n", - "[FUSED-ATTN-CACHE] rank=0 | tid=0 dev=0 | f16 fwd BUILD_PLANS | hit= 3, miss= 1, create_graph= 1, cache_graph= 1, build_plans= 1, execute= 0\n", - "...\n", - "[FUSED-ATTN-CACHE] rank=0 | tid=0 dev=0 | f16 fwd EXECUTE | hit= 5, miss= 1, create_graph= 1, cache_graph= 1, build_plans= 1, execute= 1\n", - "...\n", - "[FUSED-ATTN-CACHE] rank=0 | ===== summary begin =====\n", - "[FUSED-ATTN-CACHE] rank=0 | tid=0 dev=0 | f16 fwd | hit= 5, miss= 1, create_graph= 1, cache_graph= 1, build_plans= 1, execute= 1\n", - "[FUSED-ATTN-CACHE] rank=0 | tid=0 dev=0 | f16 bwd | hit= 1, miss= 1, create_graph= 1, cache_graph= 1, build_plans= 0, execute= 0\n", - "[FUSED-ATTN-CACHE] rank=0 | tid=1 dev=0 | f16 bwd | hit= 4, miss= 0, create_graph= 0, cache_graph= 0, build_plans= 1, execute= 1\n", - "[FUSED-ATTN-CACHE] rank=0 | tid=all dev=all | f16 fwd | hit= 5, miss= 1, create_graph= 1, cache_graph= 1, build_plans= 1, execute= 1\n", - "[FUSED-ATTN-CACHE] rank=0 | tid=all dev=all | f16 bwd | hit= 5, miss= 1, create_graph= 1, cache_graph= 1, build_plans= 1, execute= 1\n", - "[FUSED-ATTN-CACHE] rank=0 | f16 fwd validate | calls=1 | time= 0.168 ms/call\n", - "[FUSED-ATTN-CACHE] rank=0 | f16 fwd build_operation_graph | calls=1 | time= 11.405 ms/call\n", - "[FUSED-ATTN-CACHE] rank=0 | f16 fwd create_execution_plans | calls=1 | time= 2.054 ms/call\n", - "[FUSED-ATTN-CACHE] rank=0 | f16 fwd check_support | calls=1 | time= 0.044 ms/call\n", - "[FUSED-ATTN-CACHE] rank=0 | f16 fwd build_plans | calls=1 | time= 327.083 ms/call\n", - "[FUSED-ATTN-CACHE] rank=0 | f16 bwd validate | calls=1 | time= 0.009 ms/call\n", - "[FUSED-ATTN-CACHE] rank=0 | f16 bwd build_operation_graph | calls=1 | time= 1.026 ms/call\n", - "[FUSED-ATTN-CACHE] rank=0 | f16 bwd create_execution_plans | calls=1 | time= 2.874 ms/call\n", - "[FUSED-ATTN-CACHE] rank=0 | f16 bwd check_support | calls=1 | time= 0.016 ms/call\n", - "[FUSED-ATTN-CACHE] rank=0 | f16 bwd build_plans | calls=1 | time= 514.599 ms/call\n", - "[FUSED-ATTN-CACHE] rank=0 | ===== summary end =====\n", - "```" + "Level 1 prints an end-of-run summary with graph creation, insertion, hit, miss, plan-build, and execution counters for each backend and pass. It also reports the mean CPU wall time of `validate`, `build_operation_graph`, `create_execution_plans`, `check_support`, and `build_plans`. Level 2 additionally prints each event as it occurs and includes the normalized cache key on hits and misses.\n" ] }, { @@ -387,7 +355,12 @@ "NVTE_FUSED_ATTN = 0 # disables cuDNN attention; default = 1\n", "```\n", "\n", + "```\n", + "
\n", + "Note\n", + " \n", "Environment variables NVTE_FLASH_ATTN, NVTE_UNFUSED_ATTN, and NVTE_FUSED_ATTN_USE_FAv2_BWD are supported in PyTorch. NVTE_FUSED_ATTN and NVTE_ALLOW_NONDETERMINISTIC_ALGO are supported in both PyTorch and JAX.\n", + "
\n", "\n", "### 2.3 Example Tests\n", "\n", diff --git a/docs/examples/jax/attention_context_parallel.py b/docs/examples/jax/attention_context_parallel.py index c41aa6c425f..1557a30b7cc 100644 --- a/docs/examples/jax/attention_context_parallel.py +++ b/docs/examples/jax/attention_context_parallel.py @@ -245,23 +245,21 @@ def context_parallel_supported() -> Tuple[bool, str]: return False, f"needs {cp_size} GPUs" has_kernel = is_fused_attn_kernel_available( - is_training=True, - batch_size=batch, - q_dtype=dtype, - kv_dtype=dtype, - qkv_layout=QKVLayout.THD_THD_THD, - attn_bias_type=AttnBiasType.NO_BIAS, - attn_mask_type=AttnMaskType.PADDING_CAUSAL_MASK, - softmax_type=AttnSoftmaxType.VANILLA_SOFTMAX, - dropout_probability=0.0, - q_num_heads=num_query_heads, - kv_num_heads=num_kv_heads, - q_max_seqlen=seq, - kv_max_seqlen=seq, - head_dim_qk=head_dim, - head_dim_v=head_dim, - window_size=window_size, - max_segments_per_seq=max_segments_per_seq, + True, + dtype, + dtype, + QKVLayout.THD_THD_THD, + AttnBiasType.NO_BIAS, + AttnMaskType.PADDING_CAUSAL_MASK, + AttnSoftmaxType.VANILLA_SOFTMAX, + 0.0, + num_query_heads, + num_kv_heads, + seq, + seq, + head_dim, + head_dim, + window_size, ) if not has_kernel: return False, "no fused attention kernel for the THD SWA shape" diff --git a/docs/examples/jax/test_attention.py b/docs/examples/jax/test_attention.py index 15321e85da3..3cc08271dc7 100644 --- a/docs/examples/jax/test_attention.py +++ b/docs/examples/jax/test_attention.py @@ -87,23 +87,21 @@ def _context_parallel_supported(): return False, f"needs {cp_size} GPUs" has_kernel = is_fused_attn_kernel_available( - is_training=True, - batch_size=2, - q_dtype=jnp.bfloat16, - kv_dtype=jnp.bfloat16, - qkv_layout=QKVLayout.THD_THD_THD, - attn_bias_type=AttnBiasType.NO_BIAS, - attn_mask_type=AttnMaskType.PADDING_CAUSAL_MASK, - softmax_type=AttnSoftmaxType.VANILLA_SOFTMAX, - dropout_probability=0.0, - q_num_heads=128, - kv_num_heads=8, - q_max_seqlen=65536, - kv_max_seqlen=65536, - head_dim_qk=128, - head_dim_v=128, - window_size=(8192, 0), - max_segments_per_seq=4, + True, + jnp.bfloat16, + jnp.bfloat16, + QKVLayout.THD_THD_THD, + AttnBiasType.NO_BIAS, + AttnMaskType.PADDING_CAUSAL_MASK, + AttnSoftmaxType.VANILLA_SOFTMAX, + 0.0, + 128, + 8, + 65536, + 65536, + 128, + 128, + (8192, 0), ) if not has_kernel: return False, "no fused attention kernel for the THD SWA shape" diff --git a/qa/L0_jax_unittest/test.sh b/qa/L0_jax_unittest/test.sh index ad1157cbadc..897c2c7bd2c 100644 --- a/qa/L0_jax_unittest/test.sh +++ b/qa/L0_jax_unittest/test.sh @@ -28,7 +28,9 @@ pip3 install pytest==8.2.1 pytest-timeout==2.4.0 || error_exit "Failed to instal : ${XML_LOG_DIR:=/logs} mkdir -p "$XML_LOG_DIR" -python3 -m pytest -c $TE_PATH/tests/jax/pytest.ini -v --junitxml=$XML_LOG_DIR/pytest_jax_not_distributed.xml $TE_PATH/tests/jax --ignore=$TE_PATH/tests/jax/test_multi_process_ep.py -k 'not distributed' || test_fail "tests/jax/*not_distributed_*" +python3 -m pytest --import-mode=importlib -v --junitxml=$XML_LOG_DIR/pytest_test_common_attention_helpers.xml $TE_PATH/tests/test_common_attention_helpers.py || test_fail "test_common_attention_helpers.py" +python3 -m pytest -c $TE_PATH/tests/jax/pytest.ini -v --junitxml=$XML_LOG_DIR/pytest_jax_not_distributed.xml $TE_PATH/tests/jax --ignore=$TE_PATH/tests/jax/test_multi_process_ep.py --ignore=$TE_PATH/tests/jax/test_fp8_fused_attn.py -k 'not distributed' || test_fail "tests/jax/*not_distributed_*" +python3 -m pytest -c $TE_PATH/tests/jax/pytest.ini -v --junitxml=$XML_LOG_DIR/pytest_jax_fp8_fused_attn.xml $TE_PATH/tests/jax/test_fp8_fused_attn.py || test_fail "tests/jax/test_fp8_fused_attn.py" python3 -m pytest -c $TE_PATH/tests/jax/pytest.ini -v --junitxml=$XML_LOG_DIR/pytest_jax_fused_attn_score_mod.xml $TE_PATH/tests/jax/test_fused_attn_score_mod.py || test_fail "tests/jax/test_fused_attn_score_mod.py" NVTE_ALLOW_NONDETERMINISTIC_ALGO=0 python3 -m pytest -c $TE_PATH/tests/jax/pytest.ini -v --junitxml=$XML_LOG_DIR/pytest_jax_fused_attn_with_determinism.xml $TE_PATH/tests/jax/test_fused_attn.py -k "TestFusedAttnWithDeterminism" || test_fail "tests/jax/test_fused_attn.py" diff --git a/qa/L0_pytorch_unittest/test.sh b/qa/L0_pytorch_unittest/test.sh index a78a99d7f95..bb868197ea4 100644 --- a/qa/L0_pytorch_unittest/test.sh +++ b/qa/L0_pytorch_unittest/test.sh @@ -29,6 +29,7 @@ export NVTE_FLASH_ATTN_V4=0 pip3 install pytest==8.2.1 || error_exit "Failed to install pytest" +python3 -m pytest --import-mode=importlib --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_common_attention_helpers.xml $TE_PATH/tests/test_common_attention_helpers.py || test_fail "test_common_attention_helpers.py" NVTE_GROUPED_LINEAR_SINGLE_PARAM=1 python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_sanity.xml $TE_PATH/tests/pytorch/test_sanity.py || test_fail "test_sanity.py" python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_recipe.xml $TE_PATH/tests/pytorch/test_recipe.py || test_fail "test_recipe.py" python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_custom_recipe.xml $TE_PATH/tests/pytorch/test_custom_recipe.py || test_fail "test_custom_recipe.py" diff --git a/tests/jax/test_distributed_fused_attn.py b/tests/jax/test_distributed_fused_attn.py index 2648088b678..b079d48aff0 100644 --- a/tests/jax/test_distributed_fused_attn.py +++ b/tests/jax/test_distributed_fused_attn.py @@ -82,6 +82,25 @@ def impl_test_self_attn( is_training = True batch, seqlen, num_head, hidden = data_shape + if not is_fused_attn_kernel_available( + is_training, + dtype, + dtype, + QKVLayout.BS3HD, + attn_bias_type, + attn_mask_type, + softmax_type, + dropout_prob, + num_head, + num_head, + seqlen, + seqlen, + hidden, + hidden, + None, # no window + ): + pytest.skip("No FusedAttn backend found") + col_ref = self.generate_collectives_count_ref( mesh_shape, mesh_axes, @@ -215,6 +234,25 @@ def test_cross_attn( batch, seqlen, num_head, hidden = data_shape + if not is_fused_attn_kernel_available( + is_training, + dtype, + dtype, + QKVLayout.BSHD_BS2HD, + attn_bias_type, + attn_mask_type, + softmax_type, + dropout_prob, + num_head, + num_head, + seqlen, + seqlen, + hidden, + hidden, + None, # no window + ): + pytest.skip("No FusedAttn backend found") + col_ref = self.generate_collectives_count_ref() runner = FusedAttnRunner( batch, @@ -450,17 +488,9 @@ def impl_test_context_parallel_attn( cp_load_balanced=load_balanced, ) - # Mirror _FusedAttnCPWithAllGatherHelper.get_adjusted_max_segments_per_seq() - runner_segments = runner._get_max_segments_per_sequence() - if stripe_size and cp_strategy in (CPStrategy.DEFAULT, CPStrategy.ALL_GATHER): - max_segments_per_seq = runner_segments + seqlen // (stripe_size * cp_size) - else: - max_segments_per_seq = runner_segments - def check_has_backend_for_mask(mask_type): return is_fused_attn_kernel_available( is_training, - batch, dtype, dtype, qkv_layout, @@ -474,9 +504,8 @@ def check_has_backend_for_mask(mask_type): seqlen, hidden, hidden, - None, # no SWA for CP - max_segments_per_seq=max_segments_per_seq, - ) + None, + ) # no SWA for CP # For causal masking we depend on having bottom right support also. # The API does not check this and instead we rely on lower level checks to raise diff --git a/tests/jax/test_fp8_fused_attn.py b/tests/jax/test_fp8_fused_attn.py new file mode 100644 index 00000000000..6d689cbd36e --- /dev/null +++ b/tests/jax/test_fp8_fused_attn.py @@ -0,0 +1,314 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""Coverage for JAX FP8 DPA and attention features shared with PyTorch.""" + +from math import sqrt +from types import SimpleNamespace + +import jax +import jax.numpy as jnp +import numpy as np +import pytest +from transformer_engine_jax import get_cudnn_version, get_device_compute_capability + +from transformer_engine.common import recipe +from transformer_engine.jax import autocast +from transformer_engine.jax.attention import ( + AttnBiasType, + AttnMaskType, + AttnSoftmaxType, + QKVLayout, + SequenceDescriptor, + fused_attn, +) +from transformer_engine.jax.cpp_extensions import FusedAttnHelper +from transformer_engine.jax.cpp_extensions.fp8_attention import ( + _mx_scale, + _mxfp8_scale_inv, + _validate_quantizer_modes, +) +from transformer_engine.jax.flax import DotProductAttention +from transformer_engine.jax.quantize import ( + AttentionQuantizerSet, + BlockScaleQuantizer, + QuantizeLayout, + ScalingMode, +) +from transformer_engine.jax.sharding import MeshResource + + +def _require_gpu(min_arch=90, min_cudnn=90700, max_arch=None): + try: + if not any(device.platform == "gpu" for device in jax.devices()): + pytest.skip("A CUDA device is required.") + arch = get_device_compute_capability(0) + except RuntimeError as exc: + pytest.skip(f"A usable CUDA device is required: {exc}") + if arch < min_arch: + pytest.skip(f"This test requires SM{min_arch} or newer, found SM{arch}.") + if max_arch is not None and arch >= max_arch: + pytest.skip(f"This test requires an architecture older than SM{max_arch}, found SM{arch}.") + cudnn_version = get_cudnn_version() + if cudnn_version < min_cudnn: + pytest.skip(f"This test requires cuDNN {min_cudnn}, found {cudnn_version}.") + if cudnn_version == 91000: + pytest.skip("cuDNN 9.10.0 has known FP8 SDPA issues.") + return arch + + +def _reference_attention(q, k, v, *, bottom_right=False, alibi=False): + q_seqlen, kv_seqlen = q.shape[1], k.shape[1] + scores = jnp.einsum("bqhd,bkhd->bhqk", q.astype(jnp.float32), k.astype(jnp.float32)) / sqrt( + q.shape[-1] + ) + q_pos = jnp.arange(q_seqlen)[:, None] + kv_pos = jnp.arange(kv_seqlen)[None, :] + shift = kv_seqlen - q_seqlen if bottom_right else 0 + if alibi: + heads = q.shape[-2] + power_of_two_heads = 2 ** int(np.floor(np.log2(heads))) + base = 2.0 ** (-8.0 / power_of_two_heads) + slopes = base ** jnp.arange(1, power_of_two_heads + 1, dtype=jnp.float32) + if power_of_two_heads < heads: + extra_base = 2.0 ** (-4.0 / power_of_two_heads) + extra = extra_base ** jnp.arange( + 1, 1 + 2 * (heads - power_of_two_heads), 2, dtype=jnp.float32 + ) + slopes = jnp.concatenate((slopes, extra)) + distance = jnp.abs(q_pos + shift - kv_pos).astype(jnp.float32) + scores -= slopes[None, :, None, None] * distance[None, None, :, :] + allowed = kv_pos <= q_pos + shift + scores = jnp.where(allowed[None, None, :, :], scores, -jnp.inf) + probabilities = jax.nn.softmax(scores, axis=-1) + return jnp.einsum("bhqk,bkhd->bqhd", probabilities, v.astype(jnp.float32)).astype(q.dtype) + + +def _assert_fp8_close(actual, expected): + actual = np.asarray(actual, dtype=np.float32) + expected = np.asarray(expected, dtype=np.float32) + np.testing.assert_allclose(actual, expected, atol=0.5, rtol=0.05) + assert np.sqrt(np.mean(np.square(actual - expected))) < 0.11 + + +def _mode_quantizer(mode): + return SimpleNamespace(scaling_mode=mode) + + +def test_fp8_attention_quantizer_mode_assignments(): + """Current-scaling DPA uses delayed scaling for its internal S and dP tensors.""" + + current = _mode_quantizer(ScalingMode.CURRENT_TENSOR_SCALING) + delayed = _mode_quantizer(ScalingMode.DELAYED_TENSOR_SCALING) + quantizers = AttentionQuantizerSet( + qkv=current, + s=delayed, + o=current, + do=current, + dp=delayed, + dqkv=current, + ) + assert _validate_quantizer_modes(quantizers) == "current" + + quantizers.s = current + with pytest.raises(ValueError, match=r"s=current \(expected delayed\)"): + _validate_quantizer_modes(quantizers) + + +@pytest.mark.parametrize( + "fp8_recipe,min_arch,min_cudnn", + ( + pytest.param( + recipe.DelayedScaling(amax_history_len=1, fp8_dpa=True), + 90, + 90700, + id="delayed", + ), + pytest.param( + recipe.Float8CurrentScaling(fp8_dpa=True), + 100, + 91400, + id="current", + ), + pytest.param( + recipe.MXFP8BlockScaling(fp8_dpa=True), + 100, + 92100, + id="mxfp8", + ), + ), +) +@pytest.mark.parametrize("input_dtype", (jnp.float16, jnp.bfloat16), ids=("float16", "bfloat16")) +def test_fp8_dpa_forward_backward(fp8_recipe, min_arch, min_cudnn, input_dtype): + """Each supported recipe executes FP8 DPA behind FP16/BF16 module boundaries.""" + + _require_gpu(min_arch, min_cudnn, max_arch=120) + if fp8_recipe.mxfp8() and get_cudnn_version() in (92300, 92301): + pytest.skip("cuDNN 9.23.0 and 9.23.1 have known MXFP8 SDPA correctness issues.") + batch, seqlen, heads, dim = 2, 128, 8, 128 + q_key, k_key, v_key, do_key = jax.random.split(jax.random.PRNGKey(1234), 4) + shape = (batch, seqlen, heads, dim) + q = jax.random.uniform(q_key, shape, input_dtype, minval=-0.5, maxval=0.5) + k = jax.random.uniform(k_key, shape, input_dtype, minval=-0.5, maxval=0.5) + v = jax.random.uniform(v_key, shape, input_dtype, minval=-0.5, maxval=0.5) + doutput = jax.random.uniform(do_key, shape, input_dtype, minval=-0.5, maxval=0.5) + seqlens = jnp.full((batch,), seqlen, dtype=jnp.int32) + descriptor = SequenceDescriptor.from_seqlens((seqlens, seqlens)) + module = DotProductAttention( + head_dim=dim, + num_attention_heads=heads, + num_gqa_groups=heads, + attn_mask_type="causal", + qkv_layout="bshd_bshd_bshd", + transpose_batch_sequence=False, + ) + + def loss_fn(variables, query, key, value): + output = module.apply(variables, query, key, value, descriptor, deterministic=False) + loss = jnp.sum(output.astype(jnp.float32) * doutput.astype(jnp.float32)) + return loss, output + + def reference_loss(query, key, value): + output = _reference_attention(query, key, value) + loss = jnp.sum(output.astype(jnp.float32) * doutput.astype(jnp.float32)) + return loss, output + + with autocast(enabled=True, recipe=fp8_recipe, mesh_resource=MeshResource()): + variables = module.init(jax.random.PRNGKey(0), q, k, v, descriptor, deterministic=False) + (_, output), (_, dq, dk, dv) = jax.value_and_grad( + loss_fn, argnums=(0, 1, 2, 3), has_aux=True + )(variables, q, k, v) + (_, reference), (dq_ref, dk_ref, dv_ref) = jax.value_and_grad( + reference_loss, argnums=(0, 1, 2), has_aux=True + )(q, k, v) + + assert output.dtype == q.dtype + _assert_fp8_close(output, reference) + for actual, expected in zip((dq, dk, dv), (dq_ref, dk_ref, dv_ref)): + _assert_fp8_close(actual, expected) + + +def test_mxfp8_attention_scale_layout(): + """Attention pads, permutes, and swizzles compact JAX MXFP8 scales for cuDNN.""" + + quantizer = BlockScaleQuantizer( + q_dtype=jnp.float8_e4m3fn, + scaling_mode=ScalingMode.MXFP8_1D_SCALING, + q_layout=QuantizeLayout.ROWWISE_COLWISE, + data_layout="NN", + ) + tensor = quantizer.quantize(jnp.ones((2, 64, 8, 64), dtype=jnp.bfloat16), flatten_axis=-2) + assert tensor.rowwise_tensor.scale_inv.shape == (2, 64, 8, 2) + assert tensor.colwise_tensor.scale_inv.shape == (2, 2, 8, 64) + assert _mxfp8_scale_inv(tensor).shape == (2, 8, 128, 4) + assert _mxfp8_scale_inv(tensor, colwise=True).shape == (2, 8, 4, 128) + + +def test_mxfp8_attention_scale_graph_stride(): + """The cuDNN scale descriptor matches the contiguous BHSD swizzle buffer.""" + + class FakeTensor: + def set_reordering_type(self, reordering): + self.reordering = reordering + return self + + class FakeGraph: + def tensor(self, **kwargs): + self.kwargs = kwargs + return FakeTensor() + + class FakeCudnn: + class data_type: + FP8_E8M0 = "fp8_e8m0" + + class tensor_reordering: + F8_128x4 = "f8_128x4" + + graph = FakeGraph() + tensor = _mx_scale( + graph, + FakeCudnn, + name="descale_q", + uid=101, + batch=2, + heads=8, + seqlen=128, + dim=4, + ) + + assert graph.kwargs["dim"] == (2, 8, 128, 4) + assert graph.kwargs["stride"] == (4096, 512, 4, 1) + assert tensor.reordering == FakeCudnn.tensor_reordering.F8_128x4 + + +@pytest.mark.parametrize("feature", ("alibi", "bottom_right")) +def test_fused_attention_parity_features(feature): + """JAX executes ALiBi and explicit bottom-right diagonal attention.""" + + _require_gpu(90, 90700) + batch, heads, dim = 2, 8, 64 + q_seqlen, kv_seqlen = (128, 128) if feature == "alibi" else (64, 128) + q_key, k_key, v_key, do_key = jax.random.split(jax.random.PRNGKey(4321), 4) + q = jax.random.normal(q_key, (batch, q_seqlen, heads, dim), jnp.bfloat16) * 0.25 + k = jax.random.normal(k_key, (batch, kv_seqlen, heads, dim), jnp.bfloat16) * 0.25 + v = jax.random.normal(v_key, (batch, kv_seqlen, heads, dim), jnp.bfloat16) * 0.25 + doutput = jax.random.normal(do_key, (batch, q_seqlen, heads, dim), jnp.bfloat16) * 0.25 + q_lengths = jnp.full((batch,), q_seqlen, dtype=jnp.int32) + kv_lengths = jnp.full((batch,), kv_seqlen, dtype=jnp.int32) + descriptor = SequenceDescriptor.from_seqlens((q_lengths, kv_lengths)) + bias_type = AttnBiasType.ALIBI if feature == "alibi" else AttnBiasType.NO_BIAS + helper = FusedAttnHelper( + True, + q.dtype, + k.dtype, + QKVLayout.BSHD_BSHD_BSHD, + bias_type, + AttnMaskType.CAUSAL_MASK, + AttnSoftmaxType.VANILLA_SOFTMAX, + 0.0, + heads, + heads, + q_seqlen, + kv_seqlen, + dim, + dim, + (-1, -1), + ) + if not helper.is_fused_attn_kernel_available(): + pytest.skip("No fused-attention kernel supports this configuration.") + + def te_loss(query, key, value): + output = fused_attn( + (query, key, value), + None, + descriptor, + None, + bias_type, + AttnMaskType.CAUSAL_MASK, + QKVLayout.BSHD_BSHD_BSHD, + AttnSoftmaxType.VANILLA_SOFTMAX, + 1.0 / sqrt(dim), + 0.0, + True, + bottom_right_diagonal=feature == "bottom_right", + ) + return jnp.sum(output.astype(jnp.float32) * doutput.astype(jnp.float32)), output + + def reference_loss(query, key, value): + output = _reference_attention( + query, + key, + value, + bottom_right=feature == "bottom_right", + alibi=feature == "alibi", + ) + return jnp.sum(output.astype(jnp.float32) * doutput.astype(jnp.float32)), output + + (_, output), grads = jax.value_and_grad(te_loss, argnums=(0, 1, 2), has_aux=True)(q, k, v) + (_, reference), reference_grads = jax.value_and_grad( + reference_loss, argnums=(0, 1, 2), has_aux=True + )(q, k, v) + np.testing.assert_allclose(output, reference, atol=0.02, rtol=0.02) + for actual, expected in zip(grads, reference_grads): + np.testing.assert_allclose(actual, expected, atol=0.03, rtol=0.03) diff --git a/tests/jax/test_fused_attn.py b/tests/jax/test_fused_attn.py index 6544c63716d..bf74cbbfdc1 100644 --- a/tests/jax/test_fused_attn.py +++ b/tests/jax/test_fused_attn.py @@ -40,7 +40,7 @@ CPStrategy, ReorderStrategy, ) -from transformer_engine.jax.cpp_extensions import FusedAttnHelper +from transformer_engine.jax.cpp_extensions import FusedAttnHelper, cudnn_attention from transformer_engine_jax import ( NVTE_Fused_Attn_Backend, get_cudnn_version, @@ -48,14 +48,11 @@ ) from distributed_test_base import assert_equal_collectives -from utils import assert_allclose, get_test_level, print_debug_tensor_stats +from utils import assert_allclose, print_debug_tensor_stats # Get determinism _deterministic = not bool(int(os.getenv("NVTE_ALLOW_NONDETERMINISTIC_ALGO", "1"))) -# CI test level -_TEST_LEVEL = get_test_level() - @pytest.fixture(autouse=True, scope="module") def init(): @@ -386,11 +383,48 @@ def score_mod(_graph, score, _tensors): ) -def test_fused_attn_backend_message(): - """Test the error messaging of the fused attention backend query.""" +@pytest.mark.parametrize( + "cudnn_version, expected_dimensions", + [ + ((9, 5, 1), (8, 768, 640, False)), + ((9, 6, 0), (32, 2048, 2048, True)), + ], +) +def test_thd_graph_bucketing_requires_cudnn_9_6(monkeypatch, cudnn_version, expected_dimensions): + """Pre-9.6 THD graphs retain dense dimensions and metadata extents.""" + monkeypatch.setattr(cudnn_attention, "get_cudnn_version", lambda: cudnn_version) + monkeypatch.setattr(cudnn_attention, "_device_arch", lambda: 90) + info = cudnn_attention._LayoutInfo( + batch_shape=(2,), + input_batch=2, + q_max_seqlen=768, + kv_max_seqlen=640, + q_heads=8, + kv_heads=8, + qk_dim=128, + v_dim=128, + ) + + class Config: + qkv_layout = QKVLayout.THD_THD_THD + max_segments_per_seq = 4 + return_max_logit = False + + dimensions = cudnn_attention._graph_dimensions(info, Config()) + + assert dimensions[:4] == expected_dimensions + if cudnn_version < (9, 6, 0): + assert dimensions[4] == (2, 8, 768, 4) + else: + assert dimensions[4] == (2, 768, 8, 1) + + +def test_fused_attn_backend_message(monkeypatch): + """The JAX selector returns the shared policy's rejection reason.""" + monkeypatch.setattr(cudnn_attention, "get_cudnn_version", lambda: (9, 25, 0)) + monkeypatch.setattr(cudnn_attention, "_device_arch", lambda: 90) baseline = FusedAttnHelper( is_training=True, - batch_size=2, q_dtype=jnp.bfloat16, kv_dtype=jnp.bfloat16, qkv_layout=QKVLayout.BSHD_BSHD_BSHD, @@ -405,27 +439,21 @@ def test_fused_attn_backend_message(): head_dim_qk=64, head_dim_v=64, window_size=(-1, -1), - attn_scale=0.125, ) - # One of TE's rules is violated and the error message is surfaced + backend, message = baseline.get_fused_attn_backend() + assert backend == NVTE_Fused_Attn_Backend.NVTE_F16_arbitrary_seqlen + assert message == "" + backend, message = replace( baseline, attn_bias_type=AttnBiasType.PRE_SCALE_BIAS ).get_fused_attn_backend() assert backend == NVTE_Fused_Attn_Backend.NVTE_No_Backend - assert message == "Fused attention does not support pre-scale bias." - - # No error message if supported; otherwise skip the test - backend, message = baseline.get_fused_attn_backend() - if backend == NVTE_Fused_Attn_Backend.NVTE_No_Backend: - pytest.skip(f"FusedAttention does not support the baseline config: {message}") - assert message == "" + assert message == "attention bias is not supported" - # All TE rules have cleared; now gets rejected by cuDNN's support check - # cuDNN's error message might change across cuDNN versions, so only verify the presence of the string backend, message = replace(baseline, head_dim_qk=1024, head_dim_v=1024).get_fused_attn_backend() assert backend == NVTE_Fused_Attn_Backend.NVTE_No_Backend - assert message != "" + assert message == "head dimensions are not supported" class BiasShape(Enum): @@ -527,27 +555,6 @@ def _get_max_segments_per_sequence(self): return 1 def _check_configs(self): - # Trim SWA configs for L0 and L1 to reduce test time; need to trim more in future test refactoring. - if self.window_size is not None and ( - self.dropout_prob != 0.0 or self.attn_bias_type is not AttnBiasType.NO_BIAS - ): - if _TEST_LEVEL == "L0" and ( - self.softmax_type != AttnSoftmaxType.VANILLA_SOFTMAX - or self.dtype != jnp.bfloat16 - or self.attn_bias_type is not AttnBiasType.POST_SCALE_BIAS - or self.attn_mask_type is not AttnMaskType.NO_MASK - ): - pytest.skip( - "Trimmed SWA+bias/dropout config: only vanilla-softmax + bf16 + post_scale_bias" - " + no-mask runs at L0" - ) - if _TEST_LEVEL == "L1" and ( - self.dtype != jnp.float16 or self.softmax_type != AttnSoftmaxType.LEARNABLE_SOFTMAX - ): - pytest.skip( - "Trimmed SWA+bias/dropout config: only float16 + learnable-softmax runs at L1" - ) - # TODO(KshitijLakhani): probably add/move this to is_fused_attn_available if self.qkv_layout.is_thd() and not self.attn_mask_type.is_padding(): pytest.skip("THD format requires padding masks.") @@ -650,44 +657,25 @@ def _check_configs(self): "is either BSHD_BSHD_BSHD or THD_THD_THD" ) - bias_batch = bias_heads = bias_seqlen_q = bias_seqlen_kv = None - if self.attn_bias_type == AttnBiasType.POST_SCALE_BIAS: - if self.bias_shape == BiasShape._1HSS: - bias_batch, bias_heads = 1, self.num_heads_q - elif self.bias_shape == BiasShape._B1SS: - bias_batch, bias_heads = self.batch_size, 1 - elif self.bias_shape == BiasShape._BHSS: - bias_batch, bias_heads = self.batch_size, self.num_heads_q - elif self.bias_shape == BiasShape._11SS: - bias_batch, bias_heads = 1, 1 - bias_seqlen_q, bias_seqlen_kv = self.max_seqlen_q, self.max_seqlen_kv - - self.backend, message = FusedAttnHelper( - is_training=self.is_training, - batch_size=self.batch_size, - q_dtype=self.dtype, - kv_dtype=self.dtype, - qkv_layout=self.qkv_layout, - attn_bias_type=self.attn_bias_type, - attn_mask_type=self.attn_mask_type, - softmax_type=self.softmax_type, - dropout_probability=self.dropout_prob, - q_num_heads=self.num_heads_q, - kv_num_heads=self.num_heads_kv, - q_max_seqlen=self.max_seqlen_q, - kv_max_seqlen=self.max_seqlen_kv, - head_dim_qk=self.head_dim_qk, - head_dim_v=self.head_dim_v, - window_size=(-1, -1) if self.window_size is None else self.window_size, - bottom_right_diagonal=self.attn_mask_type.is_bottom_right(), - bias_batch=bias_batch, - bias_heads=bias_heads, - bias_seqlen_q=bias_seqlen_q, - bias_seqlen_kv=bias_seqlen_kv, - max_segments_per_seq=self._get_max_segments_per_sequence(), + self.backend, _ = FusedAttnHelper( + self.is_training, + self.dtype, + self.dtype, + self.qkv_layout, + self.attn_bias_type, + self.attn_mask_type, + self.softmax_type, + self.dropout_prob, + self.num_heads_q, + self.num_heads_kv, + self.max_seqlen_q, + self.max_seqlen_kv, + self.head_dim_qk, + self.head_dim_v, + (-1, -1) if self.window_size is None else self.window_size, ).get_fused_attn_backend() if self.backend != NVTE_Fused_Attn_Backend.NVTE_F16_arbitrary_seqlen: - pytest.skip(message) + pytest.skip("Unsupported inputs combination or device compute capability.") if ( self.attn_bias_type == AttnBiasType.POST_SCALE_BIAS diff --git a/tests/jax/test_fused_attn_score_mod.py b/tests/jax/test_fused_attn_score_mod.py index 561aa9b851a..181a6344499 100644 --- a/tests/jax/test_fused_attn_score_mod.py +++ b/tests/jax/test_fused_attn_score_mod.py @@ -8,7 +8,10 @@ import jax.numpy as jnp import numpy as np import pytest +from test_fused_attn import FusedAttnRunner, SeqDescFormat +from transformer_engine_jax import get_device_compute_capability +import transformer_engine.jax.cpp_extensions.cudnn_graph as cudnn_graph import transformer_engine.jax.cpp_extensions.flex_attention as tex_attention from transformer_engine.jax.attention import ( AttnBiasType, @@ -18,9 +21,6 @@ ) from transformer_engine.jax.cpp_extensions import make_fused_attn_score_mod_config from transformer_engine.jax.flax import transformer as flax_transformer -from transformer_engine_jax import get_device_compute_capability, NVTE_Fused_Attn_Backend -from test_fused_attn import FusedAttnRunner, SeqDescFormat - _CONFIG_TEST_HEAD_DIM = 128 _CONFIG_TEST_SCALING_FACTOR = 1.0 / sqrt(_CONFIG_TEST_HEAD_DIM) @@ -397,17 +397,9 @@ def _identity_score_mod(_graph, score, _tensors): def _install_fake_flax_fused_attn(monkeypatch, *, kernel_available=True): captured = {} - class FakeFusedAttnHelper: - def __init__(self, *args, **kwargs): - captured.setdefault("kernel_checks", []).append((args, kwargs)) - - def get_fused_attn_backend(self): - if kernel_available: - return NVTE_Fused_Attn_Backend.NVTE_F16_arbitrary_seqlen, "" - return ( - NVTE_Fused_Attn_Backend.NVTE_No_Backend, - "fake FusedAttnHelper: no fused attention backend available for this configuration", - ) + def fake_fused_attn_kernel_check(*args, **kwargs): + captured.setdefault("kernel_checks", []).append((args, kwargs)) + return kernel_available def fake_fused_attn( qkv, @@ -435,6 +427,7 @@ def fake_fused_attn( score_mod_tensors=None, score_mod_bprop_tensors=None, return_max_logit=False, + bottom_right_diagonal=None, ): captured.update( qkv=qkv, @@ -460,10 +453,15 @@ def fake_fused_attn( score_mod_bprop=score_mod_bprop, score_mod_tensors=score_mod_tensors, score_mod_bprop_tensors=score_mod_bprop_tensors, + bottom_right_diagonal=bottom_right_diagonal, ) return qkv[0] - monkeypatch.setattr(flax_transformer, "FusedAttnHelper", FakeFusedAttnHelper) + monkeypatch.setattr( + flax_transformer, + "is_fused_attn_kernel_available", + fake_fused_attn_kernel_check, + ) monkeypatch.setattr(flax_transformer, "fused_attn", fake_fused_attn) return captured @@ -538,7 +536,7 @@ def test_dot_product_attention_plumbs_score_mod_to_fused_attn(monkeypatch): assert captured["attn_bias_type"] is AttnBiasType.NO_BIAS assert captured["qkv_layout"] is QKVLayout.BSHD_BSHD_BSHD assert captured["softmax_type"] is AttnSoftmaxType.VANILLA_SOFTMAX - assert captured["kernel_checks"][0][1]["qkv_layout"] is QKVLayout.BSHD_BSHD_BSHD + assert captured["kernel_checks"][0][0][3] is QKVLayout.BSHD_BSHD_BSHD def test_dot_product_attention_unpacks_packed_score_mod_to_separate_layout(monkeypatch): @@ -562,7 +560,7 @@ def test_dot_product_attention_unpacks_packed_score_mod_to_separate_layout(monke assert captured["qkv"][0].shape == (1, 8, 1, 16) assert captured["qkv_layout"] is QKVLayout.BSHD_BSHD_BSHD assert captured["score_mod"] is _identity_score_mod - assert captured["kernel_checks"][0][1]["qkv_layout"] is QKVLayout.BSHD_BSHD_BSHD + assert captured["kernel_checks"][0][0][3] is QKVLayout.BSHD_BSHD_BSHD def test_multi_head_attention_plumbs_score_mod_to_dot_product_attention(monkeypatch): @@ -659,19 +657,19 @@ class FakeCudnn: __version__ = "1.22.0" monkeypatch.setattr( - tex_attention.transformer_engine_jax, + cudnn_graph.transformer_engine_jax, "get_cudnn_frontend_version", lambda: 12200, ) - assert tex_attention._check_cudnn_frontend_version_match(FakeCudnn) == 12200 + assert cudnn_graph.check_cudnn_frontend_version_match(FakeCudnn) == 12200 monkeypatch.setattr( - tex_attention.transformer_engine_jax, + cudnn_graph.transformer_engine_jax, "get_cudnn_frontend_version", lambda: 12100, ) with pytest.raises(RuntimeError, match="Python/C\\+\\+ version mismatch"): - tex_attention._check_cudnn_frontend_version_match(FakeCudnn) + cudnn_graph.check_cudnn_frontend_version_match(FakeCudnn) def test_fused_attn_score_mod_config_stabilizes_bound_method_cache_keys(): @@ -735,6 +733,27 @@ def forward(self, _graph, score, _tensors): assert tex_attention._graph_cache_key("fwd", config_1, ()) is None +def test_fused_attn_score_mod_module_lambda_cache_keys_do_not_collide(): + """Different module-level lambdas must not reuse the same cuDNN graph.""" + score_mod_1 = lambda _graph, score, _tensors: score + score_mod_2 = lambda _graph, score, _tensors: score + score_mod_1.__module__ = __name__ + score_mod_2.__module__ = __name__ + score_mod_1.__qualname__ = "" + score_mod_2.__qualname__ = "" + + config_1, _, _ = make_fused_attn_score_mod_config( + score_mod_1, None, None, None, _CONFIG_TEST_SCALING_FACTOR, True + ) + config_2, _, _ = make_fused_attn_score_mod_config( + score_mod_2, None, None, None, _CONFIG_TEST_SCALING_FACTOR, True + ) + + assert config_1 != config_2 + assert tex_attention._graph_cache_key("fwd", config_1, ()) is not None + assert tex_attention._graph_cache_key("fwd", config_2, ()) is not None + + @pytest.mark.skipif(not _has_cudnn_frontend_python(), reason="cuDNN Python frontend is required") def test_fused_attn_score_mod_post_scale_bias_optional_bprop(): """Post-scale-bias score_mod matches the JAX reference without explicit bprop.""" diff --git a/tests/jax/utils.py b/tests/jax/utils.py index f0a6928f1b2..c5e564dbc7b 100644 --- a/tests/jax/utils.py +++ b/tests/jax/utils.py @@ -119,27 +119,13 @@ def combine_biases(*masks: Optional[Array]): return mask -TEST_LEVELS = ("L0", "L1", "L2") - - -def get_test_level(): - """ - Returns the test level specified in the environment variable, NVTE_JAX_UNITTEST_LEVEL. - """ - test_level = os.environ.get("NVTE_JAX_UNITTEST_LEVEL", TEST_LEVELS[0]) - if test_level not in TEST_LEVELS: - raise ValueError( - f"Unsupported test level {test_level!r}, expected one of {', '.join(TEST_LEVELS)}" - ) - return test_level - - def get_parameters_for_test_level(param_dict: dict): """ Takes an input dictionary of parameters keyed by test type "L0", etc. Returns the parameters for the test level specified in the environment variable """ - test_level = get_test_level() + DEFAULT_TEST_LEVEL = "L0" + test_level = os.environ.get("NVTE_JAX_UNITTEST_LEVEL", DEFAULT_TEST_LEVEL) if test_level not in param_dict: raise ValueError("Unsupported test level") return param_dict[test_level] diff --git a/tests/pytorch/attention/run_graph_cache.py b/tests/pytorch/attention/run_graph_cache.py deleted file mode 100644 index 534711515c6..00000000000 --- a/tests/pytorch/attention/run_graph_cache.py +++ /dev/null @@ -1,107 +0,0 @@ -# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# -# See LICENSE for license information. - -"""Worker for test_attention.py::test_fused_attn_graph_cache. - -Runs a fixed sequence of support queries and executions against the FusedAttention graph cache and -collects level-2 diagnostics with NVTE_FUSED_ATTN_CACHE_DEBUG=2 for different phases: - query the first support query for a config -- the miss that creates its fwd/bwd graphs - requery the identical query again -- hits the cache - exec forward and backward of that config -- hits the cache and builds plans before execution - rescale the same execution with only softmax_scale changed -- should hit the cache and reuse the plans, - since attn_scale is normalized out of the cache key - reshape a query differing in max_seqlen -- should miss and recreate the graphs -""" - -import os -import pathlib -import sys - -import torch - -_current_file = pathlib.Path(__file__).resolve() -sys.path = [str(_current_file.parent.parent)] + sys.path - -from transformer_engine.pytorch import DotProductAttention -from transformer_engine.pytorch.attention.dot_product_attention import _attention_backends -from utils import ModelConfig, get_available_attention_backends - -DTYPE = torch.bfloat16 -QKV_FORMAT = "bshd" -QKV_LAYOUT = "bshd_bshd_bshd" -DETERMINISTIC = ( - not bool(int(os.getenv("NVTE_ALLOW_NONDETERMINISTIC_ALGO", "1"))) - or torch.are_deterministic_algorithms_enabled() -) - - -def mark_phase(name: str) -> None: - """Delimit the cache events of one phase from the next one's.""" - sys.stderr.write(f"[CACHE-TEST] phase={name}\n") - sys.stderr.flush() - - -def query(config: ModelConfig) -> bool: - """Run one backend support query and report whether FusedAttention supports - the configuration. If so, it populates the cache without executing anything.""" - available_backends, _, fused_attn_backends = get_available_attention_backends( - config, qkv_dtype=DTYPE, qkv_layout=QKV_LAYOUT, deterministic=DETERMINISTIC - ) - _, fused_attn_supported, _ = available_backends - return fused_attn_supported and len(fused_attn_backends) > 0 - - -def execute(config: ModelConfig, softmax_scale: float) -> None: - """Run a forward and backward pass of `config` on the FusedAttention backend.""" - block = DotProductAttention( - config.num_heads, - config.head_dim_qk, - attention_dropout=config.dropout_p, - qkv_format=QKV_FORMAT, - attn_mask_type=config.attn_mask_type, - softmax_scale=softmax_scale, - layer_number=1, - attention_type=config.attn_type, - ).to(dtype=DTYPE, device="cuda") - shape = (config.batch_size, config.max_seqlen_q, config.num_heads, config.head_dim_qk) - q, k, v = [torch.randn(shape, dtype=DTYPE, device="cuda", requires_grad=True) for _ in range(3)] - out = block(q, k, v, core_attention_bias_type=config.attn_bias_type) - out.backward(torch.randn_like(out)) - torch.cuda.synchronize() - - -def main() -> int: - torch.manual_seed(1234) - config = ModelConfig(2, 512, 8, 64) - reshaped = ModelConfig(2, 256, 8, 64) - - mark_phase("query") - fused_available = query(config) - print(f"[CACHE-TEST] fused={int(fused_available)}", flush=True) - if not fused_available: - return 0 - - mark_phase("requery") - query(config) - - os.environ["NVTE_FLASH_ATTN"] = "0" - os.environ["NVTE_FUSED_ATTN"] = "1" - os.environ["NVTE_UNFUSED_ATTN"] = "0" - _attention_backends["backend_selection_requires_update"] = True - - mark_phase("exec") - execute(config, softmax_scale=0.125) - - mark_phase("rescale") - execute(config, softmax_scale=0.25) - - mark_phase("reshape") - query(reshaped) - - mark_phase("done") - return 0 - - -if __name__ == "__main__": - sys.exit(main()) diff --git a/tests/pytorch/attention/test_attention.py b/tests/pytorch/attention/test_attention.py index 1b934af323d..8897f717568 100644 --- a/tests/pytorch/attention/test_attention.py +++ b/tests/pytorch/attention/test_attention.py @@ -1,16 +1,12 @@ # Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. -import collections import logging import os -import re -import subprocess import sys import pathlib import copy -from dataclasses import replace -from typing import Any, Dict, List, NamedTuple, Tuple, Union +from typing import Any, Dict, Tuple, Union from packaging.version import Version as PkgVersion import pytest @@ -34,28 +30,26 @@ ) from transformer_engine.pytorch.attention.dot_product_attention.utils import ( FlashAttentionUtils, - FusedAttentionParams, _get_supported_versions, check_set_window_size, - get_fused_attn_spec, ) from transformer_engine.pytorch.attention import RotaryPositionEmbedding import transformer_engine.pytorch.cpp_extensions as ext from transformer_engine.pytorch.cpp_extensions.fused_attn import ( FusedAttnBackend, - QKVFormat, - QKVLayout, fused_attn_bwd, fused_attn_fwd, ) from transformer_engine.pytorch.distributed import CudaRNGStatesTracker +from transformer_engine.pytorch._extra_state import UNSAFE_PICKLE_EXTRA_STATE_ENV from transformer_engine.pytorch.module.base import TransformerEngineBaseModule from transformer_engine.pytorch.utils import ( init_method_normal, scaled_init_method_normal, ) from transformer_engine.pytorch.utils import get_cudnn_version -from transformer_engine.pytorch.constants import FP8BwdTensorIdx, FP8FwdTensorIdx +from transformer_engine.pytorch.constants import DType, FP8BwdTensorIdx, FP8FwdTensorIdx +from transformer_engine.pytorch.attention.dot_product_attention import _cudnn_backend import transformer_engine_torch as tex from transformer_engine.pytorch.quantized_tensor import ( Quantizer, @@ -137,79 +131,53 @@ def test_flash_attention_supported_version_message(): ) -def test_fused_attn_backend_message(): - """Test the error messaging of the fused attention backend query.""" - baseline = FusedAttentionParams( - qkv_layout=tex.NVTE_QKV_Layout.NVTE_BSHD_BSHD_BSHD, - dqkv_layout=tex.NVTE_QKV_Layout.NVTE_BSHD_BSHD_BSHD, - o_format=tex.NVTE_QKV_Format.NVTE_BSHD, - do_format=tex.NVTE_QKV_Format.NVTE_BSHD, - batch_size=2, +def test_fused_attn_backend_message(monkeypatch): + """The PyTorch selector returns the shared policy's rejection reason.""" + monkeypatch.setattr(_cudnn_backend, "get_cudnn_version", lambda: (9, 25, 0)) + monkeypatch.setattr(_cudnn_backend, "get_device_compute_capability", lambda: (9, 0)) + baseline = dict( + is_training=True, + q_dtype=DType.kBFloat16, + kv_dtype=DType.kBFloat16, + qkv_layout="bshd_bshd_bshd", + bias_type="no_bias", + attn_mask_type="no_mask", + softmax_type="vanilla", + dropout=0.0, num_attn_heads=8, num_gqa_groups=8, - head_dim_qk=64, - head_dim_v=64, max_seqlen_q=128, max_seqlen_kv=128, - attn_scale=0.125, - ) - - # One of TE's rules is violated and the error message is surfaced - backend, message = tex.get_fused_attn_backend( - replace(baseline, bias_type=tex.NVTE_Bias_Type.NVTE_PRE_SCALE_BIAS) + head_dim_qk=64, + head_dim_v=64, + window_size_left=-1, + window_size_right=-1, + return_max_logit=False, + cuda_graph=False, + deterministic=False, ) - assert backend == tex.NVTE_Fused_Attn_Backend.NVTE_No_Backend - assert message == "Fused attention does not support pre-scale bias." - # No error message if supported; otherwise skip the test - backend, message = tex.get_fused_attn_backend(baseline) - if backend == tex.NVTE_Fused_Attn_Backend.NVTE_No_Backend: - pytest.skip(f"FusedAttention does not support the baseline config: {message}") + backend, message = _cudnn_backend.get_fused_attn_backend(**baseline) + assert backend == FusedAttnBackend.F16_arbitrary_seqlen assert message == "" - # All TE rules have cleared; now gets rejected by cuDNN's support check - # cuDNN's error message might change across cuDNN versions, so only verify the presence of the string - backend, message = tex.get_fused_attn_backend( - replace(baseline, head_dim_qk=1024, head_dim_v=1024) - ) - assert backend == tex.NVTE_Fused_Attn_Backend.NVTE_No_Backend - assert message != "" - - if not fp8_attn_available: - return - - # The same checks in FP8; the dtypes, formats and scaling mode come from get_fused_attn_spec so - # that the config matches what FusedAttnFunc feeds the kernels for this recipe - fp8_recipe = recipe.DelayedScaling(fp8_format=recipe.Format.HYBRID, fp8_dpa=True) - spec = get_fused_attn_spec(fp8_recipe, torch.float8_e4m3fn, "bshd_bshd_bshd", cs_o_in_f16=True) - fp8_baseline = replace( - baseline, - scaling_mode=spec.scaling_mode, - qkv_dtype=spec.qkv, - o_dtype=spec.o, - do_dtype=spec.do, - dqkv_dtype=spec.dqkv, - qkv_layout=QKVLayout[spec.qkv_layout], - o_format=QKVFormat[spec.o_format], - do_format=QKVFormat[spec.do_format], - dqkv_layout=QKVLayout[spec.dqkv_layout], - qkv_scale_inv_format=QKVFormat[spec.scale_inv_format], - do_scale_inv_format=QKVFormat[spec.scale_inv_format], + backend, message = _cudnn_backend.get_fused_attn_backend( + **{**baseline, "bias_type": "pre_scale_bias"} ) + assert backend == FusedAttnBackend.No_Backend + assert message == "attention bias is not supported" - # No error message if supported; otherwise skip the test - backend, message = tex.get_fused_attn_backend(fp8_baseline) - if backend == tex.NVTE_Fused_Attn_Backend.NVTE_No_Backend: - pytest.skip(f"FusedAttention does not support the baseline FP8 config: {message}") - assert backend == tex.NVTE_Fused_Attn_Backend.NVTE_FP8 - assert message == "" + backend, message = _cudnn_backend.get_fused_attn_backend( + **{**baseline, "head_dim_qk": 1024, "head_dim_v": 1024} + ) + assert backend == FusedAttnBackend.No_Backend + assert message == "head dimensions are not supported" - # All TE FP8 rules have cleared; now gets rejected by cuDNN's FP8 support check - backend, message = tex.get_fused_attn_backend( - replace(fp8_baseline, head_dim_qk=1024, head_dim_v=1024) + backend, message = _cudnn_backend.get_fused_attn_backend( + **{**baseline, "q_dtype": DType.kFloat16} ) - assert backend == tex.NVTE_Fused_Attn_Backend.NVTE_No_Backend - assert message != "" + assert backend == FusedAttnBackend.No_Backend + assert message == "Q and KV must have the same data type" # Define F16 data types to test @@ -296,13 +264,6 @@ def test_dot_product_attention( "Setting is_training to False as cuDNN does not support dbias for" f" {config.bias_shape=} " ) - # Generate token counts for THD so the support query and execution will use the same values - num_tokens_q, num_tokens_kv = None, None - if qkv_format == "thd": - reset_rng_states() - seqlens = _generate_seqlens(config, qkv_format, pad_between_seqs) - num_tokens_q = int(seqlens.cu_q_after_pad[-1]) - num_tokens_kv = int(seqlens.cu_kv_after_pad[-1]) available_backends, _, fused_attn_backends = get_available_attention_backends( config, qkv_dtype=dtype, @@ -310,8 +271,6 @@ def test_dot_product_attention( pad_between_seqs=pad_between_seqs, is_training=is_training, deterministic=_deterministic, - num_tokens_q=num_tokens_q, - num_tokens_kv=num_tokens_kv, ) flash_attn_supported, fused_attn_supported, unfused_attn_supported = available_backends @@ -328,8 +287,6 @@ def test_dot_product_attention( pad_between_seqs=pad_between_seqs, is_training=is_training, deterministic=_deterministic, - num_tokens_q=num_tokens_q, - num_tokens_kv=num_tokens_kv, ) flash_attn_supported, fused_attn_supported, unfused_attn_supported = available_backends @@ -396,97 +353,6 @@ def test_dot_product_attention( torch.testing.assert_close(fused_attn_bwd[i], flash_attn_bwd[i], **tols) -_CACHE_EVENT = re.compile( - r"\[FUSED-ATTN-CACHE\]\s+(?:rank=\d+\s+\|\s+)?tid=\d+\s+dev=-?\d+\s+\|\s+" - r"(?Pf16|fp8)\s+(?Pfwd|bwd)\s+" - r"(?PCREATE_GRAPH|CACHE_GRAPH|BUILD_PLANS|EXECUTE|MISS|HIT)\b(?P.*)" -) -_CACHE_PHASE = re.compile(r"\[CACHE-TEST\] phase=(?P\w+)") - - -@pytest.mark.skipif(get_cudnn_version() < (8, 9, 1), reason="cuDNN 8.9.1+ is required.") -def test_fused_attn_graph_cache(): - """Test FusedAttention graph cache with level 2 diagnostics. It runs a subprocess - to avoid contamination from other tests as the counters are process-wide. - """ - worker = _current_file.parent / "run_graph_cache.py" - result = subprocess.run( - [sys.executable, str(worker)], - env={ - **os.environ, - "NVTE_FUSED_ATTN_CACHE_DEBUG": "2", - "PYTHONUNBUFFERED": "1", - }, - capture_output=True, - text=True, - timeout=900, - check=False, - ) - assert result.returncode == 0, f"Graph cache worker exited with {result.returncode}." - if "[CACHE-TEST] fused=1" not in result.stdout: - pytest.skip("No FusedAttention backend for the graph cache test config.") - - # Group the cache diagnostics: - # events[phase][(pass, event)] = count of events - # miss_keys[phase][pass] = set of distinct cache keys that missed - events = collections.defaultdict(collections.Counter) - miss_keys = collections.defaultdict(lambda: collections.defaultdict(set)) - phase = None - for line in result.stderr.splitlines(): - phase_match = _CACHE_PHASE.search(line) - if phase_match is not None: - phase = phase_match.group("name") - continue - event_match = _CACHE_EVENT.search(line) - if event_match is None or phase is None: - continue - event_pass, event = event_match.group("pass"), event_match.group("event") - events[phase][(event_pass, event)] += 1 - if event == "MISS": - miss_keys[phase][event_pass].add(event_match.group("rest").split("|")[-1].strip()) - - context = f"\n--- stderr ---\n{result.stderr[-8000:]}" - - for pass_name in ("fwd", "bwd"): - - def expect(phase, event, expected, reason, actual=None, pass_name=pass_name): - if actual is None: - actual = events[phase][(pass_name, event)] - ok = actual >= 1 if expected == "1+" else actual == expected - assert ( - ok - ), f"{pass_name} {phase}: expected {event}={expected}, got {actual} ({reason}){context}" - - expect("query", "MISS", 1, "expected one cold miss") - expect("query", "CREATE_GRAPH", 1, "expected one build") - expect("query", "CACHE_GRAPH", 1, "build was not cached") - expect("query", "BUILD_PLANS", 0, "query compiled kernels") - - expect("requery", "MISS", 0, "repeated query missed") - expect("requery", "CREATE_GRAPH", 0, "repeated query rebuilt") - expect("requery", "HIT", "1+", "repeated query never looked") - - expect("exec", "MISS", 0, "execution missed the query's graph") - expect("exec", "CREATE_GRAPH", 0, "execution rebuilt the graph") - expect("exec", "EXECUTE", "1+", "fused attention never ran") - expect("exec", "BUILD_PLANS", 1, "expected one plan build") - - expect("rescale", "MISS", 0, "attn_scale changed the key") - expect("rescale", "CREATE_GRAPH", 0, "attn_scale forced a build") - expect("rescale", "BUILD_PLANS", 0, "attn_scale recompiled") - expect("rescale", "EXECUTE", "1+", "rescaled run did not execute") - - expect("reshape", "MISS", 1, "expected one miss") - expect("reshape", "CREATE_GRAPH", 1, "expected one build") - expect( - "reshape", - "MISS keys", - 1, - "more than one new cache key", - actual=len(miss_keys["reshape"][pass_name]), - ) - - @pytest.mark.skipif(get_cudnn_version() < (8, 9, 1), reason="cuDNN 8.9.1+ is required.") @pytest.mark.parametrize("dtype", param_types) @pytest.mark.parametrize("model_configs", [model_configs_base]) @@ -832,14 +698,6 @@ def test_dpa_softmax(dtype, model_configs, model): @pytest.mark.parametrize("model", model_configs_softmax.keys()) def test_dpa_softmax_thd(dtype, model_configs, model): """Test DotProductAttention module with different softmax types""" - config = model_configs[model] - if "padding" not in config.attn_mask_type: - promoted = dict(vars(config)) - promoted["attn_mask_type"] = ( - "padding" if config.attn_mask_type == "no_mask" else "padding_" + config.attn_mask_type - ) - if any(name != model and vars(other) == promoted for name, other in model_configs.items()): - pytest.skip("Duplicate test to others with THD and padding mask.") test_dot_product_attention(dtype, model_configs, model, True, "thd_thd_thd", False, False) @@ -1486,9 +1344,6 @@ def test_dpa_bias_shapes(dtype, model_configs, model): @pytest.mark.parametrize("qkv_layout", ["thd_thd_thd", "sbhd_sbhd_sbhd"]) def test_dpa_sliding_window(dtype, model_configs, model, qkv_layout): """Test DotProductAttention module with sliding window attention""" - config = model_configs[model] - if qkv_layout == "thd_thd_thd" and "padding" not in config.attn_mask_type: - pytest.skip("Duplicate test to others with THD and padding mask.") test_dot_product_attention(dtype, model_configs, model, False, qkv_layout, True, False) @@ -1767,26 +1622,38 @@ def make_dot_product_attention( return block -class _Seqlens(NamedTuple): - """Sequence lengths for one test case, before and after inter-sequence padding.""" - - q: torch.Tensor - kv: torch.Tensor - cu_q: torch.Tensor - cu_kv: torch.Tensor - q_after_pad: torch.Tensor - kv_after_pad: torch.Tensor - cu_q_after_pad: torch.Tensor - cu_kv_after_pad: torch.Tensor - pad_len: Union[List[int], torch.Tensor] - - -def _generate_seqlens( +def run_dot_product_attention( + dtype: torch.dtype, config: ModelConfig, - qkv_format: str, + backend: str, + ckpt_attn: bool, + qkv_layout: str, pad_between_seqs: bool, -) -> _Seqlens: - """Draw the sequence lengths for one test case.""" + is_training: bool, + declarative_packed: bool = False, + forward_kwargs: Dict[str, Any] = None, +) -> Tuple[torch.Tensor, Tuple[torch.Tensor, torch.Tensor, torch.Tensor]]: + """Run DotProductAttention module with one forward pass and one backward pass. + + With declarative_packed=True (packed qkv_layout only), the packed buffer is + passed to DotProductAttention directly via qkv_layer/kv_layer instead of + slicing it into q/k/v views, and input gradients are read off the packed + buffer itself.""" + # Set RNG and environment varables + reset_rng_states() + os.environ["NVTE_FLASH_ATTN"] = "0" + os.environ["NVTE_FUSED_ATTN"] = "0" + os.environ["NVTE_UNFUSED_ATTN"] = "0" + if backend == "FlashAttention": + os.environ["NVTE_FLASH_ATTN"] = "1" + if backend == "FusedAttention": + os.environ["NVTE_FUSED_ATTN"] = "1" + if backend == "UnfusedDotProductAttention": + os.environ["NVTE_UNFUSED_ATTN"] = "1" + _attention_backends["backend_selection_requires_update"] = True + + # Create seqlens + qkv_format = "".join([i for i in qkv_layout.split("_")[0] if i.isalpha()]) if "padding" in config.attn_mask_type or qkv_format == "thd": if config.attn_type == "self": seqlens_q = torch.randint( @@ -1826,63 +1693,6 @@ def _generate_seqlens( cu_seqlens_q_after_pad[1:] = torch.cumsum(seqlens_q_after_pad, dim=0) cu_seqlens_kv_after_pad[1:] = torch.cumsum(seqlens_kv_after_pad, dim=0) - return _Seqlens( - seqlens_q, - seqlens_kv, - cu_seqlens_q, - cu_seqlens_kv, - seqlens_q_after_pad, - seqlens_kv_after_pad, - cu_seqlens_q_after_pad, - cu_seqlens_kv_after_pad, - pad_len, - ) - - -def run_dot_product_attention( - dtype: torch.dtype, - config: ModelConfig, - backend: str, - ckpt_attn: bool, - qkv_layout: str, - pad_between_seqs: bool, - is_training: bool, - declarative_packed: bool = False, - forward_kwargs: Dict[str, Any] = None, -) -> Tuple[torch.Tensor, Tuple[torch.Tensor, torch.Tensor, torch.Tensor]]: - """Run DotProductAttention module with one forward pass and one backward pass. - - With declarative_packed=True (packed qkv_layout only), the packed buffer is - passed to DotProductAttention directly via qkv_layer/kv_layer instead of - slicing it into q/k/v views, and input gradients are read off the packed - buffer itself.""" - # Set RNG and environment varables - reset_rng_states() - os.environ["NVTE_FLASH_ATTN"] = "0" - os.environ["NVTE_FUSED_ATTN"] = "0" - os.environ["NVTE_UNFUSED_ATTN"] = "0" - if backend == "FlashAttention": - os.environ["NVTE_FLASH_ATTN"] = "1" - if backend == "FusedAttention": - os.environ["NVTE_FUSED_ATTN"] = "1" - if backend == "UnfusedDotProductAttention": - os.environ["NVTE_UNFUSED_ATTN"] = "1" - _attention_backends["backend_selection_requires_update"] = True - - # Create seqlens - qkv_format = "".join([i for i in qkv_layout.split("_")[0] if i.isalpha()]) - ( - seqlens_q, - seqlens_kv, - cu_seqlens_q, - cu_seqlens_kv, - seqlens_q_after_pad, - seqlens_kv_after_pad, - cu_seqlens_q_after_pad, - cu_seqlens_kv_after_pad, - pad_len, - ) = _generate_seqlens(config, qkv_format, pad_between_seqs) - # Create attention mask if padding attention_mask = None if "padding" in config.attn_mask_type: @@ -2262,13 +2072,6 @@ def test_transformer_layer( # Test backend availability is_training = True - # Get the token counts for THD and use them for both support query and actual execution - num_tokens_q, num_tokens_kv = None, None - if qkv_format == "thd": - reset_rng_states() - seqlens = _generate_seqlens(config, qkv_format, pad_between_seqs=False) - num_tokens_q = int(seqlens.cu_q[-1]) - num_tokens_kv = int(seqlens.cu_kv[-1]) available_backends, _, fused_attn_backends = get_available_attention_backends( config, qkv_dtype=dtype, @@ -2277,8 +2080,6 @@ def test_transformer_layer( ), is_training=is_training, deterministic=_deterministic, - num_tokens_q=num_tokens_q, - num_tokens_kv=num_tokens_kv, ) flash_attn_supported, fused_attn_supported, unfused_attn_supported = available_backends if not fused_attn_supported: @@ -2293,8 +2094,6 @@ def test_transformer_layer( ), is_training=is_training, deterministic=_deterministic, - num_tokens_q=num_tokens_q, - num_tokens_kv=num_tokens_kv, ) flash_attn_supported, fused_attn_supported, unfused_attn_supported = available_backends @@ -2608,35 +2407,25 @@ def _run_transformer_layer( @pytest.mark.skipif(get_cudnn_version() < (9, 3, 0), reason="cuDNN 9.3.0+ is required.") @pytest.mark.parametrize("model", ["large"]) @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) -def test_dpa_fp8_extra_state(model, dtype): +def test_dpa_fp8_extra_state(model, dtype, monkeypatch): """Test DotProductAttention module in FP8 with checkpointing""" config = model_configs_fp8_extra_state[model] # Test backend availability is_training = True - fp8_recipe = recipe.DelayedScaling( - margin=0, - fp8_format=recipe.Format.HYBRID, - amax_history_len=1, - amax_compute_algo="most_recent", - fp8_dpa=True, - ) - fp8_meta = {} - fp8_meta["recipe"] = fp8_recipe available_backends, _, fused_attn_backends = get_available_attention_backends( config, qkv_dtype=torch.float8_e4m3fn, - nominal_dtype=dtype, qkv_layout="sb3hd", is_training=is_training, deterministic=_deterministic, - fp8=True, - fp8_meta=fp8_meta, ) flash_attn_supported, fused_attn_supported, unfused_attn_supported = available_backends if not fused_attn_supported and not flash_attn_supported: pytest.skip("No attention backend available.") outputs = _run_dpa_fp8_extra_state(dtype, config, checkpoint=False) + # The checkpoints are generated locally by this test and are therefore trusted. + monkeypatch.setenv(UNSAFE_PICKLE_EXTRA_STATE_ENV, "1") outputs_checkpoint = _run_dpa_fp8_extra_state(dtype, config, checkpoint=True) outputs_checkpoint_v1_6 = _run_dpa_fp8_extra_state( dtype, config, mimic_v1_6=True, checkpoint=True @@ -2817,70 +2606,6 @@ def _get_fp8_vs_f16_config(model, qkv_layout): return config -def _dpa_fp8_vs_f16_seqlens(config, qkv_format): - """Draw the sequence lengths for one test_dpa_fp8_vs_f16 case.""" - if "padding" in config.attn_mask_type or qkv_format == "thd": - if config.attn_type == "self": - seqlens_q = torch.randint( - 1, config.max_seqlen_q, [config.batch_size], dtype=torch.int32, device="cuda" - ) - seqlens_kv = seqlens_q - if config.attn_type == "cross": - seqlens_q = torch.randint( - 1, config.max_seqlen_q, [config.batch_size], dtype=torch.int32, device="cuda" - ) - seqlens_kv = torch.randint( - 1, config.max_seqlen_kv, [config.batch_size], dtype=torch.int32, device="cuda" - ) - else: - seqlens_q = torch.full( - [config.batch_size], config.max_seqlen_q, dtype=torch.int32, device="cuda" - ) - seqlens_kv = torch.full( - [config.batch_size], config.max_seqlen_kv, dtype=torch.int32, device="cuda" - ) - return seqlens_q, seqlens_kv - - -def _mha_fp8_vs_f16_seqlens(config, qkv_format): - """Draw the sequence lengths for one test_mha_fp8_vs_f16 case. - - These come from a dedicated generator rather than the default one, because the module the run - builds beforehand consumes randomness the backend probe cannot replay.""" - gen = torch.Generator(device="cuda") - gen.manual_seed(seed) - - def draw(high, count): - return torch.randint(1, high, [count], dtype=torch.int32, device="cuda", generator=gen) - - if "padding" in config.attn_mask_type or qkv_format == "thd": - - def random_seqlens(max_seqlen): - if qkv_format != "thd": - return draw(max_seqlen, config.batch_size) - # Reserve seven positions so total-token alignment only increases the final length. - return torch.cat((draw(max_seqlen, config.batch_size - 1), draw(max_seqlen - 6, 1))) - - if config.attn_type == "self": - seqlens_q = random_seqlens(config.max_seqlen_q) - seqlens_kv = seqlens_q - if config.attn_type == "cross": - seqlens_q = random_seqlens(config.max_seqlen_q) - seqlens_kv = random_seqlens(config.max_seqlen_kv) - else: - seqlens_q = torch.full( - [config.batch_size], config.max_seqlen_q, dtype=torch.int32, device="cuda" - ) - seqlens_kv = torch.full( - [config.batch_size], config.max_seqlen_kv, dtype=torch.int32, device="cuda" - ) - if qkv_format == "thd": - # FP8 Linear flattens THD input to [t, h*d], so align total tokens for cuBLAS. - seqlens_q[-1] += -seqlens_q.sum() % 8 - seqlens_kv[-1] += -seqlens_kv.sum() % 8 - return seqlens_q, seqlens_kv - - @pytest.mark.skipif(get_cudnn_version() < (9, 2, 1), reason="cuDNN 9.2.1+ is required.") @pytest.mark.skipif(not fp8_attn_available, reason=reason_for_no_fp8_attn) @pytest.mark.parametrize("dtype", param_types_fp8_vs_f16) @@ -2902,9 +2627,6 @@ def test_mha_fp8_vs_f16( scaling_mode, ): """Test MultiHeadAttention module in FP8""" - if not is_training and fp8_dpa_bwd: - pytest.skip("fp8_dpa_bwd=True not applicable for inference") - os.environ["NVTE_FP8_DPA_BWD"] = "1" if fp8_dpa_bwd else "0" config = _get_fp8_vs_f16_config(model, qkv_format) @@ -2932,23 +2654,14 @@ def test_mha_fp8_vs_f16( ) fp8_meta = {} fp8_meta["recipe"] = fp8_recipe - # Get the token counts for THD and use them for both support query and actual execution - num_tokens_q, num_tokens_kv = None, None - if qkv_format == "thd": - seqlens_q, seqlens_kv = _mha_fp8_vs_f16_seqlens(config, qkv_format) - num_tokens_q = int(seqlens_q.sum()) - num_tokens_kv = int(seqlens_kv.sum()) available_backends, _, _ = get_available_attention_backends( config, qkv_dtype=torch.float8_e4m3fn, - nominal_dtype=dtype, qkv_layout=qkv_format.replace("hd", "h3d"), fp8=True, fp8_meta=fp8_meta, is_training=is_training, deterministic=_deterministic, - num_tokens_q=num_tokens_q, - num_tokens_kv=num_tokens_kv, ) flash_attn_supported, fused_attn_supported_fp8, unfused_attn_supported = available_backends available_backends, _, fused_attn_backends = get_available_attention_backends( @@ -2957,8 +2670,6 @@ def test_mha_fp8_vs_f16( qkv_layout=qkv_format.replace("hd", "h3d"), is_training=is_training, deterministic=_deterministic, - num_tokens_q=num_tokens_q, - num_tokens_kv=num_tokens_kv, ) _, fused_attn_supported_f16, _ = available_backends if flash_attn_supported + fused_attn_supported_fp8 < 1: @@ -3073,13 +2784,47 @@ def get_dummy_cuda_rng_tracker() -> CudaRNGStatesTracker: attention_type="self", qkv_weight_interleaved=True, qkv_format=qkv_format, - window_size=config.window_size, - softmax_type=config.softmax_type, ).to(dtype=dtype, device="cuda") if not is_training: mha = mha.eval() - seqlens_q, seqlens_kv = _mha_fp8_vs_f16_seqlens(config, qkv_format) + def random_seqlens(max_seqlen): + if qkv_format != "thd": + return torch.randint( + 1, max_seqlen, [config.batch_size], dtype=torch.int32, device="cuda" + ) + # Reserve seven positions so total-token alignment only increases the final length. + return torch.cat( + ( + torch.randint( + 1, + max_seqlen, + [config.batch_size - 1], + dtype=torch.int32, + device="cuda", + ), + torch.randint(1, max_seqlen - 6, [1], dtype=torch.int32, device="cuda"), + ) + ) + + if "padding" in config.attn_mask_type or qkv_format == "thd": + if config.attn_type == "self": + seqlens_q = random_seqlens(config.max_seqlen_q) + seqlens_kv = seqlens_q + if config.attn_type == "cross": + seqlens_q = random_seqlens(config.max_seqlen_q) + seqlens_kv = random_seqlens(config.max_seqlen_kv) + else: + seqlens_q = torch.full( + [config.batch_size], config.max_seqlen_q, dtype=torch.int32, device="cuda" + ) + seqlens_kv = torch.full( + [config.batch_size], config.max_seqlen_kv, dtype=torch.int32, device="cuda" + ) + if qkv_format == "thd": + # FP8 Linear flattens THD input to [t, h*d], so align total tokens for cuBLAS. + seqlens_q[-1] += -seqlens_q.sum() % 8 + seqlens_kv[-1] += -seqlens_kv.sum() % 8 cu_seqlens_q = torch.zeros(config.batch_size + 1, dtype=torch.int32, device="cuda") cu_seqlens_kv = torch.zeros(config.batch_size + 1, dtype=torch.int32, device="cuda") cu_seqlens_q[1:] = torch.cumsum(seqlens_q, dim=0) @@ -3147,10 +2892,6 @@ def get_dummy_cuda_rng_tracker() -> CudaRNGStatesTracker: def test_dpa_fp8_vs_f16(dtype, model, qkv_layout, fp8_dpa_bwd, is_training, scaling_mode): """Test DotProductAttention module in FP8""" config = _get_fp8_vs_f16_config(model, qkv_layout) - if config.num_heads != config.num_gqa_groups and "3" in qkv_layout: - pytest.skip("qkv_layout not applicable for MQA/GQA") - if not is_training and fp8_dpa_bwd: - pytest.skip("fp8_dpa_bwd=True not applicable for inference") # TODO(cyang): think of another way to verify dropout results # test cuDNN FP8 dropout @@ -3187,25 +2928,14 @@ def test_dpa_fp8_vs_f16(dtype, model, qkv_layout, fp8_dpa_bwd, is_training, scal ) fp8_meta = {} fp8_meta["recipe"] = fp8_recipe - # Get the token counts for THD and use them for both support query and actual execution - qkv_format = "".join([i for i in qkv_layout.split("_")[0] if i.isalpha()]) - num_tokens_q, num_tokens_kv = None, None - if qkv_format == "thd": - reset_rng_states() - seqlens_q, seqlens_kv = _dpa_fp8_vs_f16_seqlens(config, qkv_format) - num_tokens_q = int(seqlens_q.sum()) - num_tokens_kv = int(seqlens_kv.sum()) available_backends, _, _ = get_available_attention_backends( config, qkv_dtype=torch.float8_e4m3fn, - nominal_dtype=dtype, qkv_layout=qkv_layout, fp8=True, fp8_meta=fp8_meta, is_training=is_training, deterministic=_deterministic, - num_tokens_q=num_tokens_q, - num_tokens_kv=num_tokens_kv, ) flash_attn_supported, fused_attn_supported_fp8, unfused_attn_supported = available_backends available_backends, _, _ = get_available_attention_backends( @@ -3214,14 +2944,14 @@ def test_dpa_fp8_vs_f16(dtype, model, qkv_layout, fp8_dpa_bwd, is_training, scal qkv_layout=qkv_layout, is_training=is_training, deterministic=_deterministic, - num_tokens_q=num_tokens_q, - num_tokens_kv=num_tokens_kv, ) _, fused_attn_supported_f16, _ = available_backends if flash_attn_supported + fused_attn_supported_fp8 < 1: pytest.skip("No FP8 attention backend available.") if not fused_attn_supported_f16: pytest.skip("No reference backend available.") + if config.num_heads != config.num_gqa_groups and "3" in qkv_layout: + pytest.skip("qkv_layout not applicable for MQA/GQA") if flash_attn_supported: os.environ["NVTE_FLASH_ATTN"] = "1" @@ -3371,7 +3101,26 @@ def get_dummy_cuda_rng_tracker() -> CudaRNGStatesTracker: if not is_training: dpa = dpa.eval() - seqlens_q, seqlens_kv = _dpa_fp8_vs_f16_seqlens(config, qkv_format) + if "padding" in config.attn_mask_type or qkv_format == "thd": + if config.attn_type == "self": + seqlens_q = torch.randint( + 1, config.max_seqlen_q, [config.batch_size], dtype=torch.int32, device="cuda" + ) + seqlens_kv = seqlens_q + if config.attn_type == "cross": + seqlens_q = torch.randint( + 1, config.max_seqlen_q, [config.batch_size], dtype=torch.int32, device="cuda" + ) + seqlens_kv = torch.randint( + 1, config.max_seqlen_kv, [config.batch_size], dtype=torch.int32, device="cuda" + ) + else: + seqlens_q = torch.full( + [config.batch_size], config.max_seqlen_q, dtype=torch.int32, device="cuda" + ) + seqlens_kv = torch.full( + [config.batch_size], config.max_seqlen_kv, dtype=torch.int32, device="cuda" + ) cu_seqlens_q = torch.zeros(config.batch_size + 1, dtype=torch.int32, device="cuda") cu_seqlens_kv = torch.zeros(config.batch_size + 1, dtype=torch.int32, device="cuda") cu_seqlens_q[1:] = torch.cumsum(seqlens_q, dim=0) @@ -3492,22 +3241,10 @@ def test_custom_mha_fp8_vs_f16(dtype, model): # Test backend availability is_training = True - fp8_meta = {} - fp8_recipe = recipe.DelayedScaling( - margin=0, - fp8_format=recipe.Format.HYBRID, - amax_history_len=1, - amax_compute_algo="most_recent", - fp8_dpa=True, - ) - fp8_meta["recipe"] = fp8_recipe available_backends, _, fused_attn_backends = get_available_attention_backends( config, qkv_dtype=torch.float8_e4m3fn, - nominal_dtype=dtype, qkv_layout="bs3hd", - fp8=True, - fp8_meta=fp8_meta, is_training=is_training, deterministic=_deterministic, ) @@ -3585,7 +3322,6 @@ def _run_custom_mha_fp8(dtype, config, backend): fp8_format=recipe.Format.HYBRID, amax_history_len=1, amax_compute_algo="most_recent", - fp8_dpa=True, ) mha = Custom_MHA_FP8(config).to(dtype=dtype, device="cuda") diff --git a/tests/pytorch/attention/test_attention_with_cp.py b/tests/pytorch/attention/test_attention_with_cp.py index 8b85c300577..4d6369dbeee 100644 --- a/tests/pytorch/attention/test_attention_with_cp.py +++ b/tests/pytorch/attention/test_attention_with_cp.py @@ -401,8 +401,6 @@ def test_cp_with_flash_attention(cp_pool, dtype, model, qkv_format, cp_comm_type config, qkv_dtype=dtypes[dtype], qkv_layout="_".join([qkv_format] * 3), - cp_size=num_gpus, - cp_size_a2a=2 if cp_comm_type == "a2a+p2p" else 1, ) flash_attn_supported, *_ = available_backends if not flash_attn_supported: @@ -709,7 +707,7 @@ def test_cp_with_fused_attention( # For 111s, dbias calculation is not supported as of cuDNN 9.18, hence, test fwd only for 111s. is_training = False if config.bias_shape == "111s" else True - available_backends, *_ = get_available_attention_backends( + available_backends, _, fused_attn_backends = get_available_attention_backends( config, qkv_dtype=dtypes[dtype] if dtype != "fp8" else torch.float8_e4m3fn, qkv_layout="_".join([qkv_format] * 3), @@ -717,11 +715,23 @@ def test_cp_with_fused_attention( fp8_meta=fp8_meta, is_training=is_training, deterministic=_deterministic, - cp_size=num_gpus, - cp_size_a2a=2 if cp_comm_type == "a2a+p2p" else 1, ) _, fused_attn_supported, _ = available_backends + if fused_attn_supported and config.attn_mask_type in ["causal", "padding_causal"]: + config_copy = copy.deepcopy(config) + config_copy.context_parallel = False + config_copy.attn_mask_type = config.attn_mask_type + "_bottom_right" + available_backends, _, fused_attn_backends = get_available_attention_backends( + config_copy, + qkv_dtype=dtypes[dtype] if dtype != "fp8" else torch.float8_e4m3fn, + qkv_layout="_".join([qkv_format] * 3), + fp8=fp8, + fp8_meta=fp8_meta, + is_training=is_training, + deterministic=_deterministic, + ) + _, fused_attn_supported, _ = available_backends if not fused_attn_supported: pytest.skip("No attention backend available.") diff --git a/tests/pytorch/attention/test_cudnn_graph.py b/tests/pytorch/attention/test_cudnn_graph.py new file mode 100644 index 00000000000..599b6cc8ce7 --- /dev/null +++ b/tests/pytorch/attention/test_cudnn_graph.py @@ -0,0 +1,91 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""Tests for the PyTorch cuDNN graph runtime.""" + +import weakref + +import torch + +from transformer_engine.pytorch.attention.dot_product_attention import ( + _cudnn_graph, + cudnn_attention, +) + + +def test_page_table_uses_cudnn_logical_layout(): + """cuDNN paged attention requires a four-dimensional table descriptor.""" + + class Graph: + @staticmethod + def tensor(**kwargs): + return kwargs + + page_table = torch.empty_strided((3, 4), (6, 1), dtype=torch.int32) + graph_tensor = cudnn_attention._make_page_table_graph_tensor( + Graph(), page_table, batch=2, name="page_table_k" + ) + + assert graph_tensor == { + "name": "page_table_k", + "dim": (2, 1, 4, 1), + "stride": (6, 6, 1, 1), + "data_type": torch.int32, + } + + +def test_graph_entry_does_not_retain_execution_workspaces(monkeypatch): + """Cached graph entries should not own per-execution scratch allocations.""" + + active_devices = [] + entered_devices = [] + + class DeviceGuard: + def __init__(self, device): + self.device = device + + def __enter__(self): + active_devices.append(self.device) + entered_devices.append(self.device) + + def __exit__(self, *_args): + active_devices.pop() + + class Workspace: + pass + + class Graph: + def __init__(self): + self.workspace_refs = [] + + def execute(self, variant_pack, workspace, handle): + assert variant_pack == {"q": "tensor"} + assert handle == "handle" + assert active_devices == [torch.device("cpu")] + self.workspace_refs.append(weakref.ref(workspace)) + + allocated_workspace_refs = [] + + def allocate_workspace(size, *, dtype, device): + assert size == 123 + assert dtype == torch.uint8 + assert device == torch.device("cpu") + workspace = Workspace() + allocated_workspace_refs.append(weakref.ref(workspace)) + return workspace + + monkeypatch.setattr(_cudnn_graph.torch, "empty", allocate_workspace) + monkeypatch.setattr(_cudnn_graph.torch.cuda, "device", DeviceGuard) + monkeypatch.setattr(_cudnn_graph, "current_stream_handle", lambda _device: "handle") + + graph = Graph() + entry = _cudnn_graph.GraphEntry(graph=graph, tensors={}, workspace_size=123) + entry.execute({"q": "tensor"}, torch.device("cpu")) + entry.execute({"q": "tensor"}, torch.device("cpu")) + + assert entered_devices == [torch.device("cpu"), torch.device("cpu")] + assert not active_devices + assert len(allocated_workspace_refs) == 2 + assert all(ref() is None for ref in allocated_workspace_refs) + assert all(ref() is None for ref in graph.workspace_refs) diff --git a/tests/pytorch/attention/test_flex_attention.py b/tests/pytorch/attention/test_flex_attention.py index 55b8b3b9aa6..3d5c9731c4d 100644 --- a/tests/pytorch/attention/test_flex_attention.py +++ b/tests/pytorch/attention/test_flex_attention.py @@ -446,8 +446,8 @@ class FakeEntry: score_mod_graph_tensors = {"softcap": object()} workspace_size = 1 - def fake_execute(graph, variant_pack, workspace_size, device): - del graph, variant_pack, workspace_size, device + def fake_execute(graph, variant_pack, workspace_size, device, cache_site): + del graph, variant_pack, workspace_size, device, cache_site q, k, v, _, _ = _score_mod_cache_cpu_inputs() q = q.requires_grad_() diff --git a/tests/pytorch/attention/test_kv_cache.py b/tests/pytorch/attention/test_kv_cache.py index 46cd2eb58cf..2a857a10dcb 100644 --- a/tests/pytorch/attention/test_kv_cache.py +++ b/tests/pytorch/attention/test_kv_cache.py @@ -4,7 +4,6 @@ from collections import OrderedDict from typing import List -import copy import os import sys import pathlib @@ -469,15 +468,13 @@ def test_kv_cache(dtype, model, qkv_format, is_paged, backend, module, is_cuda_g for layer_number in range(1, num_layers + 1): inference_params.allocate_memory(layer_number) - # figure out supported backends for both reference and inference models + # figure out supported backends inference_params_qkv_format = "bshd" qkv_layout = qkv_format + "_" + "_".join([inference_params_qkv_format] * 2) if is_paged: qkv_layout = "paged_kv_" + qkv_layout - inference_config = copy.deepcopy(config) - inference_config.attn_mask_type = "padding_causal" - inference_backends, _, _ = get_available_attention_backends( - inference_config, + available_backends, _, fused_attn_backends = get_available_attention_backends( + config, qkv_dtype=dtype, qkv_layout=qkv_layout, pad_between_seqs=False, @@ -486,24 +483,13 @@ def test_kv_cache(dtype, model, qkv_format, is_paged, backend, module, is_cuda_g fp8_meta=fp8_meta, inference_params=inference_params, ) - reference_config = copy.deepcopy(config) - reference_config.attn_mask_type = "causal" - reference_config.batch_size = config.total_requests - reference_config.max_seqlen_q = config.max_seqlen_kv - reference_backends, _, _ = get_available_attention_backends( - reference_config, - qkv_dtype=dtype, - qkv_layout="bshd_bshd_bshd", - pad_between_seqs=False, - is_training=False, - ) - backend_index = ("FlashAttention", "FusedAttention", "UnfusedAttention").index(backend) - for probe, supported in ( - ("inference", inference_backends), - ("reference", reference_backends), - ): - if not supported[backend_index]: - pytest.skip(f"{backend} backend is not supported for the {probe} config") + flash_attn_supported, fused_attn_supported, unfused_attn_supported = available_backends + if backend == "FlashAttention" and not flash_attn_supported: + pytest.skip("FlashAttention backend is not supported") + if backend == "FusedAttention" and not fused_attn_supported: + pytest.skip("FusedAttention backend is not supported") + if backend == "UnfusedAttention" and not unfused_attn_supported: + pytest.skip("UnfusedAttention backend is not supported") os.environ["NVTE_FLASH_ATTN"] = str(int(backend == "FlashAttention")) os.environ["NVTE_FUSED_ATTN"] = str(int(backend == "FusedAttention")) os.environ["NVTE_UNFUSED_ATTN"] = str(int(backend == "UnfusedAttention")) diff --git a/tests/pytorch/test_torch_compile.py b/tests/pytorch/test_torch_compile.py index e5d7169da12..e790f51b701 100644 --- a/tests/pytorch/test_torch_compile.py +++ b/tests/pytorch/test_torch_compile.py @@ -1274,7 +1274,7 @@ def test_get_attention_backend_traceable(monkeypatch): """get_attention_backend must trace under torch.compile(fullgraph=True) without graph breaks. The compiled selection must stay consistent with eager when NVTE_* env vars flip (dynamo guards on os.environ) and when - attention params change, and the baked tex.get_fused_attn_backend result + attention params change, and the baked Python cuDNN backend result must drive the selection.""" from transformer_engine.pytorch.attention.dot_product_attention import utils as dpa_utils @@ -1354,9 +1354,9 @@ def fn(x, params): # guard on the wrapped function). monkeypatch.setenv("NVTE_FLASH_ATTN", "0") monkeypatch.setattr( - dpa_utils.tex, - "get_fused_attn_backend", - lambda *args: (tex.NVTE_Fused_Attn_Backend.NVTE_No_Backend, "disabled by test"), + dpa_utils, + "get_cudnn_fused_attn_backend", + lambda *args: (dpa_utils.FusedAttnBackend["No_Backend"], "disabled by test"), ) def fn_no_backend(x, params): diff --git a/tests/pytorch/utils.py b/tests/pytorch/utils.py index 6b66458985d..9ec93a3c603 100644 --- a/tests/pytorch/utils.py +++ b/tests/pytorch/utils.py @@ -23,12 +23,10 @@ from transformer_engine.pytorch.attention.dot_product_attention import _attention_backends from transformer_engine.pytorch.attention.dot_product_attention.utils import ( get_attention_backend, - get_qkv_format, AttentionParams, AttentionLogging, check_set_window_size, ) -from transformer_engine.pytorch.cpp_extensions.fused_attn import FusedAttnBackend from transformer_engine.pytorch.module.base import get_dummy_wgrad @@ -314,10 +312,6 @@ def __init__( self.attn_type = "self" if (self.max_seqlen_q == self.max_seqlen_kv) else "cross" self.bias_shape = bias_shape self.window_size = check_set_window_size(self.attn_mask_type, window_size) - self.bottom_right_diagonal = self.attn_mask_type not in { - "causal", - "padding_causal", - } self.softcap = softcap self.context_parallel = context_parallel self.cp_comm_type = cp_comm_type @@ -343,7 +337,6 @@ def get_available_attention_backends( config: ModelConfig, qkv_dtype: torch.dtype, qkv_layout: str, - nominal_dtype: Optional[torch.dtype] = None, pad_between_seqs: bool = False, deterministic: bool = False, fp8: bool = False, @@ -352,23 +345,9 @@ def get_available_attention_backends( inference_params: Optional[InferenceParams] = None, score_mod: bool = False, score_mod_bprop: bool = False, - cp_size: int = 1, - cp_size_a2a: int = 1, - num_tokens_q: Optional[int] = None, - num_tokens_kv: Optional[int] = None, ) -> Tuple[List, List]: """Check for all available attention backends that support a model configuration""" - _, q_format, kv_format = get_qkv_format(qkv_layout, inference_params) - if num_tokens_q is None: - num_tokens_q = ( - max(config.batch_size * config.max_seqlen_q // cp_size, 1) if q_format == "thd" else 0 - ) - if num_tokens_kv is None: - num_tokens_kv = ( - max(config.batch_size * config.max_seqlen_kv // cp_size, 1) if kv_format == "thd" else 0 - ) - os.environ["NVTE_FLASH_ATTN"] = "1" os.environ["NVTE_FUSED_ATTN"] = "1" os.environ["NVTE_UNFUSED_ATTN"] = "1" @@ -380,15 +359,9 @@ def get_available_attention_backends( if config.bias_shape == "bhss": alibi_slopes_shape = [config.batch_size, config.num_heads] - core_attention_bias_shape = None - if config.attn_bias_type == "post_scale_bias": - b_dim, h_dim, sq_dim, skv_dim = config.bias_shape - core_attention_bias_shape = ( - config.batch_size if b_dim == "b" else 1, - config.num_heads if h_dim == "h" else 1, - config.max_seqlen_q if sq_dim == "s" else 1, - config.max_seqlen_kv if skv_dim == "s" else 1, - ) + core_attention_bias_shape = ( + config.bias_shape if config.attn_bias_type == "post_scale_bias" else None + ) core_attention_bias_requires_grad = False # d=256 is supported by cuDNN 9.0+ for inference but not training if ( @@ -397,18 +370,12 @@ def get_available_attention_backends( and config.head_dim_v <= 128 ): # TODO(KshitijLakhani): Remove this guard when cuDNN starts support dbias calculation for bias shape 111s - if config.bias_shape != "111s": + if core_attention_bias_shape != "111s": core_attention_bias_requires_grad = True - fused_attn_backends = [] - available_backends = None - flash_attention_backend = None - fused_attention_backend = None - def test(): attention_params = AttentionParams( qkv_dtype=qkv_dtype, - nominal_dtype=nominal_dtype, qkv_layout=qkv_layout, batch_size=config.batch_size, num_heads=config.num_heads, @@ -417,11 +384,8 @@ def test(): max_seqlen_kv=config.max_seqlen_kv, head_dim_qk=config.head_dim_qk, head_dim_v=config.head_dim_v, - num_tokens_q=num_tokens_q, - num_tokens_kv=num_tokens_kv, attn_mask_type=config.attn_mask_type, window_size=config.window_size, - bottom_right_diagonal=config.bottom_right_diagonal, softcap=config.softcap, alibi_slopes_shape=alibi_slopes_shape, core_attention_bias_type=config.attn_bias_type, @@ -431,8 +395,6 @@ def test(): attention_dropout=config.dropout_p, context_parallel=config.context_parallel, cp_comm_type=config.cp_comm_type, - cp_size=cp_size, - cp_size_a2a=cp_size_a2a, deterministic=deterministic, fp8=fp8, fp8_meta=fp8_meta, @@ -468,13 +430,12 @@ def test(): _attention_backends["backend_selection_requires_update"] = False return available_backends, flash_attention_backend, fused_attention_backend - backends = {1: "F16_arbitrary_seqlen", 2: "FP8"} if AttentionLogging._is_logging_setup is False: AttentionLogging.setup_logging() - _attention_backends["backend_selection_requires_update"] = True available_backends, flash_attention_backend, fused_attention_backend = test() - if fused_attention_backend in (FusedAttnBackend[name] for name in backends.values()): + fused_attn_backends = [] + if fused_attention_backend is not None: fused_attn_backends.append(fused_attention_backend) return available_backends, flash_attention_backend, fused_attn_backends diff --git a/tests/test_common_attention_helpers.py b/tests/test_common_attention_helpers.py new file mode 100644 index 00000000000..8b19acbc6c4 --- /dev/null +++ b/tests/test_common_attention_helpers.py @@ -0,0 +1,550 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""GPU-independent tests for shared cuDNN attention helpers.""" + +from dataclasses import replace +from types import SimpleNamespace + +import pytest + +from transformer_engine.common.attention import cache_debug +from transformer_engine.common.attention.cudnn import ( + AttentionLayout, + FusedAttentionConfig, + check_f16_fused_attention_support, + check_fp8_fused_attention_support, + cudnn_mask_options, + encode_cudnn_version, + normalize_attention_mask, + parse_attention_layout, + ragged_batch_bucket, + ragged_token_bucket, +) +from transformer_engine.common.attention.fp8 import ( + FP8AttentionGraphConfig, + attention_format_stride, + build_fp8_backward_operation, + build_fp8_forward_operation, + mxfp8_padded_sizes, +) +from transformer_engine.common.attention.score_mod import ( + UNCACHEABLE_SCORE_MOD, + freeze_score_mod_cache_key, + score_mod_callback_cache_key, +) +from transformer_engine.common.cudnn_frontend import build_cudnn_graph, make_cudnn_graph + + +def _attention_config(**kwargs): + config = FusedAttentionConfig( + is_training=False, + q_dtype="float16", + kv_dtype="float16", + layout=AttentionLayout("bshd", "bshd", "bshd", "separate"), + bias_type="no_bias", + mask_type="no_mask", + softmax_type="vanilla", + dropout=0.0, + num_attn_heads=16, + num_gqa_groups=16, + max_seqlen_q=128, + max_seqlen_kv=256, + head_dim_qk=128, + head_dim_v=128, + window_size=(-1, -1), + return_max_logit=False, + cuda_graph=False, + deterministic=False, + cudnn_version=(9, 25, 0), + sm_arch=90, + ) + return replace(config, **kwargs) + + +def test_cudnn_version_encoding(): + assert encode_cudnn_version((8, 9, 7)) == 8907 + assert encode_cudnn_version((9, 25, 1)) == 92501 + + +@pytest.mark.parametrize( + "value, expected", + [(1, 32), (32, 32), (33, 64), (513, 1024), (1025, 1536)], +) +def test_ragged_batch_bucket(value, expected): + assert ragged_batch_bucket(value) == expected + + +@pytest.mark.parametrize( + "value, expected", + [(1, 1024), (1024, 1024), (1025, 2048), (32769, 65536), (65537, 98304)], +) +def test_ragged_token_bucket(value, expected): + assert ragged_token_bucket(value) == expected + + +def test_bottom_right_self_attention_normalizes_to_causal(): + mask = normalize_attention_mask( + causal=False, + bottom_right=True, + padding=False, + bottom_right_diagonal=True, + window_size=(-1, 0), + max_seqlen_q=128, + max_seqlen_kv=128, + ) + assert mask.causal + assert not mask.bottom_right + assert not mask.bottom_right_diagonal + + +def test_shared_f16_policy_basic_support_and_rejection(): + assert check_f16_fused_attention_support(_attention_config()).supported + unsupported = check_f16_fused_attention_support(_attention_config(sm_arch=75)) + assert not unsupported.supported + assert "architecture" in unsupported.reason + + +def test_shared_f16_policy_framework_parity_features(): + alibi = _attention_config(bias_type="alibi", mask_type="causal") + assert check_f16_fused_attention_support(alibi).supported + + extended_causal = _attention_config( + mask_type="causal", + window_size=(128, 64), + sm_arch=100, + ) + assert check_f16_fused_attention_support(extended_causal).supported + + modern_padding = _attention_config( + mask_type="padding", + dropout=0.1, + cudnn_version=(9, 7, 0), + ) + assert check_f16_fused_attention_support(modern_padding).supported + + +def test_shared_fp8_policy(): + fp8 = _attention_config(q_dtype="float8_e4m3", kv_dtype="float8_e4m3", sm_arch=100) + assert check_fp8_fused_attention_support(fp8).supported + assert check_fp8_fused_attention_support(replace(fp8, cudnn_version=(9, 10, 1))).supported + assert not check_fp8_fused_attention_support(replace(fp8, cudnn_version=(9, 10, 0))).supported + assert not check_fp8_fused_attention_support(replace(fp8, bias_type="alibi")).supported + assert not check_fp8_fused_attention_support(replace(fp8, return_max_logit=True)).supported + assert not check_fp8_fused_attention_support(replace(fp8, head_dim_qk=200)).supported + + +def test_shared_fp8_scaling_mode_policy(): + fp8 = _attention_config(q_dtype="float8_e4m3", kv_dtype="float8_e4m3", sm_arch=100) + + assert check_fp8_fused_attention_support(fp8, scaling_mode="delayed").supported + assert check_fp8_fused_attention_support(fp8, scaling_mode="current").supported + assert check_fp8_fused_attention_support(fp8, scaling_mode="mxfp8").supported + + hopper = replace(fp8, sm_arch=90) + assert check_fp8_fused_attention_support(hopper, scaling_mode="delayed").supported + assert not check_fp8_fused_attention_support(hopper, scaling_mode="current").supported + assert not check_fp8_fused_attention_support(hopper, scaling_mode="mxfp8").supported + + assert not check_fp8_fused_attention_support( + replace(fp8, cudnn_version=(9, 13, 9)), scaling_mode="current" + ).supported + assert not check_fp8_fused_attention_support( + replace(fp8, cudnn_version=(9, 20, 9)), scaling_mode="mxfp8" + ).supported + assert not check_fp8_fused_attention_support( + replace(fp8, cudnn_version=(9, 23, 0)), scaling_mode="mxfp8" + ).supported + assert check_fp8_fused_attention_support( + replace(fp8, cudnn_version=(9, 23, 2)), scaling_mode="mxfp8" + ).supported + + assert not check_fp8_fused_attention_support( + replace(fp8, sm_arch=120), scaling_mode="delayed" + ).supported + + +def test_shared_fp8_deterministic_backward_policy(): + fp8 = _attention_config( + is_training=True, + deterministic=True, + q_dtype="float8_e4m3", + kv_dtype="float8_e4m3", + sm_arch=100, + ) + + assert not check_fp8_fused_attention_support( + replace(fp8, cudnn_version=(9, 18, 9)), scaling_mode="delayed" + ).supported + assert check_fp8_fused_attention_support( + replace(fp8, cudnn_version=(9, 19, 0)), scaling_mode="delayed" + ).supported + + +def test_shared_fp8_thd_policy(): + thd = _attention_config( + q_dtype="float8_e4m3", + kv_dtype="float8_e4m3", + layout=AttentionLayout("thd", "thd", "thd", "separate"), + mask_type="padding", + cudnn_version=(9, 23, 0), + sm_arch=100, + ) + assert check_fp8_fused_attention_support(thd).supported + assert check_fp8_fused_attention_support(replace(thd, sm_arch=90)).supported + assert check_fp8_fused_attention_support( + replace(thd, mask_type="padding_causal_bottom_right") + ).supported + assert not check_fp8_fused_attention_support(replace(thd, cudnn_version=(9, 22, 9))).supported + assert not check_fp8_fused_attention_support(replace(thd, mask_type="no_mask")).supported + assert not check_fp8_fused_attention_support( + replace(thd, is_training=True, sm_arch=90) + ).supported + assert not check_fp8_fused_attention_support(replace(thd, head_dim_qk=144)).supported + + sink_backward = replace(thd, is_training=True, softmax_type="learnable") + assert not check_fp8_fused_attention_support( + replace(sink_backward, cudnn_version=(9, 25, 1)) + ).supported + assert check_fp8_fused_attention_support( + replace(sink_backward, cudnn_version=(9, 26, 0)) + ).supported + + +def test_shared_fp8_thd_policy_allows_64bit_offsets(): + thd = _attention_config( + q_dtype="float8_e4m3", + kv_dtype="float8_e4m3", + layout=AttentionLayout("thd", "thd", "thd", "separate"), + mask_type="padding", + max_seqlen_q=1_048_577, + max_seqlen_kv=1_048_577, + cudnn_version=(9, 23, 0), + sm_arch=100, + ) + assert check_fp8_fused_attention_support(thd).supported + + +@pytest.mark.parametrize( + "layout, expected", + [ + ("bs3hd", ("bshd", "bshd", "3hd")), + ("bshd_bs2hd", ("bshd", "bshd", "hd_2hd")), + ("bhsd_bhsd_bhsd", ("bhsd", "bhsd", "sd_sd_sd")), + ("paged_kv_bshd_bshd_bshd", ("bshd", "bshd", "paged_separate")), + ], +) +def test_parse_attention_layout(layout, expected): + parsed = parse_attention_layout(layout) + assert (parsed.q_format, parsed.kv_format, parsed.layout_group) == expected + + +def test_shared_mask_options_use_modern_band_api(): + options = cudnn_mask_options( + causal=True, + bottom_right=False, + padding=False, + bottom_right_diagonal=True, + window_size=(32, -1), + max_seqlen_q=64, + max_seqlen_kv=128, + cudnn_version=(9, 6, 0), + ) + assert options == { + "diagonal_alignment": "bottom_right", + "is_padding": False, + "diagonal_band_left_bound": 33, + "diagonal_band_right_bound": 0, + } + + +def test_shared_fp8_shape_helpers(): + assert attention_format_stride(2, 8, 128, 64, "bshd") == (65536, 64, 512, 1) + assert attention_format_stride(2, 8, 128, 64, "bhsd") == (65536, 8192, 64, 1) + assert mxfp8_padded_sizes(129, 33, 160, 96) == { + "s_q_padded": 256, + "s_kv_padded": 128, + "s_q_scale_padded": 8, + "s_kv_scale_padded": 4, + "d_qk_padded": 256, + "d_v_padded": 128, + "d_qk_scale_padded": 8, + "d_v_scale_padded": 4, + } + + +class _FakeFP8Graph: + def __init__(self): + self.call = None + + def sdpa_fp8(self, *args, **kwargs): + self.call = ("fp8_fwd", args, kwargs) + return "o", "stats", "amax_s", "amax_o" + + def sdpa_mxfp8(self, *args, **kwargs): + self.call = ("mx_fwd", args, kwargs) + return "o", "stats", "amax_o" + + def sdpa_fp8_backward(self, *args, **kwargs): + self.call = ("fp8_bwd", args, kwargs) + return "dq", "dk", "dv", "aq", "ak", "av", "ap" + + def sdpa_mxfp8_backward(self, *args, **kwargs): + self.call = ("mx_bwd", args, kwargs) + return "dq", "dk", "dv", "aq", "ak", "av" + + +def test_shared_fp8_graph_operation_dispatch(): + graph = _FakeFP8Graph() + forward = build_fp8_forward_operation( + graph, + { + "q": 1, + "k": 2, + "v": 3, + "descale_q": 4, + "descale_k": 5, + "descale_v": 6, + "descale_s": 7, + "scale_s": 8, + "scale_o": 9, + }, + {"attn_scale": 0.125}, + FP8AttentionGraphConfig("delayed", "forward"), + ) + assert forward == { + "output": "o", + "stats": "stats", + "amax_s": "amax_s", + "amax_o": "amax_o", + } + assert graph.call[0] == "fp8_fwd" + + backward = build_fp8_backward_operation( + graph, + { + **{name: name for name in ("q", "k", "v", "o", "do", "stats")}, + **{ + name: name + for name in ( + "descale_q", + "descale_k", + "descale_v", + "descale_o", + "descale_do", + "descale_s", + "descale_dp", + "scale_s", + "scale_dq", + "scale_dk", + "scale_dv", + "scale_dp", + ) + }, + }, + {"attn_scale": 0.125}, + FP8AttentionGraphConfig("current", "backward"), + ) + assert backward["amax_dp"] == "ap" + assert graph.call[0] == "fp8_bwd" + + +def _not_an_array(_value): + return False + + +def test_score_mod_module_lambda_keys_do_not_collide(): + score_mod_0 = lambda _graph, score, _tensors: score + score_mod_1 = lambda _graph, score, _tensors: score + score_mod_0.__module__ = __name__ + score_mod_1.__module__ = __name__ + score_mod_0.__qualname__ = "" + score_mod_1.__qualname__ = "" + + key_0 = score_mod_callback_cache_key(score_mod_0, is_array=_not_an_array) + key_1 = score_mod_callback_cache_key(score_mod_1, is_array=_not_an_array) + assert key_0 is not UNCACHEABLE_SCORE_MOD + assert key_1 is not UNCACHEABLE_SCORE_MOD + assert key_0 != key_1 + + +def test_score_mod_bound_method_cache_policy(): + class Unkeyed: + def forward(self, _graph, score, _tensors): + return score + + class Keyed: + def score_mod_graph_cache_key(self): + return {"layers": [1, 2]} + + def forward(self, _graph, score, _tensors): + return score + + assert ( + score_mod_callback_cache_key(Unkeyed().forward, is_array=_not_an_array) + is UNCACHEABLE_SCORE_MOD + ) + assert score_mod_callback_cache_key( + Keyed().forward, is_array=_not_an_array + ) == score_mod_callback_cache_key(Keyed().forward, is_array=_not_an_array) + + +def test_score_mod_key_rejects_runtime_arrays(): + marker = object() + with pytest.raises(TypeError, match="must not include tensors"): + freeze_score_mod_cache_key( + {"nested": [marker]}, + is_array=lambda value: value is marker, + ) + + +class _FakeGraph: + def __init__(self, workspace_size=0, unsupported=False): + self.workspace_size = workspace_size + self.unsupported = unsupported + self.calls = [] + + def validate(self): + self.calls.append("validate") + + def build_operation_graph(self): + self.calls.append("build_operation_graph") + + def create_execution_plans(self, modes): + self.calls.append(("create_execution_plans", modes)) + + def check_support(self): + self.calls.append("check_support") + if self.unsupported: + raise _FakeCudnn.cudnnGraphNotSupportedError("unsupported") + + def build_plans(self, policy): + self.calls.append(("build_plans", policy)) + + def get_workspace_size(self): + return self.workspace_size + + +class _FakeCudnn: + class cudnnGraphNotSupportedError(Exception): + pass + + data_type = SimpleNamespace(FLOAT="float") + heur_mode = SimpleNamespace(A="a", FALLBACK="fallback") + build_plan_policy = SimpleNamespace(HEURISTICS_CHOICE="heuristics") + + def __init__(self): + self.graph_kwargs = None + + def pygraph(self, **kwargs): + self.graph_kwargs = kwargs + return "graph" + + +def test_shared_cudnn_graph_creation_and_finalization(): + cudnn = _FakeCudnn() + assert make_cudnn_graph(cudnn, "half", name="attention", handle=7) == "graph" + assert cudnn.graph_kwargs == { + "io_data_type": "half", + "intermediate_data_type": "float", + "compute_data_type": "float", + "name": "attention", + "handle": 7, + } + + graph = _FakeGraph(workspace_size=0) + assert build_cudnn_graph(cudnn, graph, description="attention") == 1 + assert graph.calls == [ + "validate", + "build_operation_graph", + ("create_execution_plans", ["a", "fallback"]), + "check_support", + ("build_plans", "heuristics"), + ] + + +def test_shared_cudnn_graph_reports_build_diagnostics(): + events = [] + graph = _FakeGraph() + build_cudnn_graph( + _FakeCudnn(), + graph, + description="attention", + debug_callback=lambda event, elapsed_ns: events.append((event, elapsed_ns)), + ) + assert [event for event, _ in events] == [ + "CREATE_GRAPH", + "validate", + "build_operation_graph", + "create_execution_plans", + "check_support", + "build_plans", + "BUILD_PLANS", + ] + assert all(elapsed_ns >= 0 for _, elapsed_ns in events) + + +def test_shared_cudnn_graph_support_error_has_context(): + with pytest.raises(RuntimeError, match="cuDNN test graph is not supported"): + build_cudnn_graph(_FakeCudnn(), _FakeGraph(unsupported=True), description="test") + + +@pytest.fixture +def cache_debug_environment(monkeypatch): + for variable in ( + "NVTE_FUSED_ATTN_CACHE_DEBUG", + "RANK", + "LOCAL_RANK", + "OMPI_COMM_WORLD_RANK", + "SLURM_PROCID", + ): + monkeypatch.delenv(variable, raising=False) + cache_debug._reset_for_tests() + yield monkeypatch + cache_debug._reset_for_tests() + + +def test_cache_debug_rank_selection(cache_debug_environment): + monkeypatch = cache_debug_environment + monkeypatch.setenv("NVTE_FUSED_ATTN_CACHE_DEBUG", "2") + monkeypatch.setenv("RANK", "1") + cache_debug._reset_for_tests() + assert not cache_debug.enabled() + + monkeypatch.setenv("NVTE_FUSED_ATTN_CACHE_DEBUG", "2:1,3") + cache_debug._reset_for_tests() + assert cache_debug.enabled(trace=True) + + monkeypatch.setenv("NVTE_FUSED_ATTN_CACHE_DEBUG", "1:all") + cache_debug._reset_for_tests() + assert cache_debug.enabled() + assert not cache_debug.enabled(trace=True) + + +def test_cache_debug_summary(cache_debug_environment): + cache_debug_environment.setenv("NVTE_FUSED_ATTN_CACHE_DEBUG", "1") + cache_debug.record_lookup("f16", "fwd", hit=False, key=("graph", 1)) + cache_debug.record_event("f16", "fwd", "create_graph") + cache_debug.record_event("f16", "fwd", "cache_graph") + cache_debug.record_event("f16", "fwd", "execute", device=0) + cache_debug.record_lookup("f16", "fwd", hit=True, key=("graph", 1)) + cache_debug.record_build_time("f16", "fwd", "validate", 2_000_000) + + summary = cache_debug.render_summary() + assert "summary begin" in summary + assert "f16 fwd" in summary + assert "hit= 1" in summary + assert "miss= 1" in summary + assert "execute= 1" in summary + assert "validate" in summary + assert "2.000 ms/call" in summary + + +def test_cache_debug_level_two_traces_lookup_key(cache_debug_environment, capsys): + cache_debug_environment.setenv("NVTE_FUSED_ATTN_CACHE_DEBUG", "2") + cache_debug.record_lookup("fp8", "bwd", hit=False, key=("shape", 128)) + trace = capsys.readouterr().err + assert "[FUSED-ATTN-CACHE]" in trace + assert "fp8 bwd MISS" in trace + assert "('shape', 128)" in trace diff --git a/transformer_engine/common/CMakeLists.txt b/transformer_engine/common/CMakeLists.txt index 4436af1955d..1452d4b49f9 100644 --- a/transformer_engine/common/CMakeLists.txt +++ b/transformer_engine/common/CMakeLists.txt @@ -184,7 +184,6 @@ list(APPEND transformer_engine_cpp_sources cudnn_utils.cpp transformer_engine.cpp fused_attn/fused_attn.cpp - fused_attn/config_and_params.cpp gemm/config.cpp normalization/common.cpp normalization/rtc_dispatch.cpp @@ -218,9 +217,6 @@ list(APPEND transformer_engine_cuda_sources dropout/dropout.cu fused_attn/context_parallel.cu fused_attn/kv_cache.cu - fused_attn/fused_attn_f16_arbitrary_seqlen.cu - fused_attn/fused_attn_fp8.cu - fused_attn/utils.cu gemm/cublaslt_gemm.cu gemm/cublaslt_grouped_gemm.cu normalization/layernorm/ln_bwd_semi_cuda_kernel.cu diff --git a/transformer_engine/common/attention/__init__.py b/transformer_engine/common/attention/__init__.py new file mode 100644 index 00000000000..da628165044 --- /dev/null +++ b/transformer_engine/common/attention/__init__.py @@ -0,0 +1,5 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""Framework-independent attention helpers.""" diff --git a/transformer_engine/common/attention/cache_debug.py b/transformer_engine/common/attention/cache_debug.py new file mode 100644 index 00000000000..7a0eb11522c --- /dev/null +++ b/transformer_engine/common/attention/cache_debug.py @@ -0,0 +1,272 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""Fused-attention graph-cache diagnostics for Python graph runtimes.""" + +from __future__ import annotations + +import atexit +import os +import sys +import threading +from collections import defaultdict +from functools import lru_cache +from time import perf_counter_ns +from typing import Callable, Optional + +_PREFIX = "[FUSED-ATTN-CACHE]" +_RANK_ENV_VARS = ("RANK", "LOCAL_RANK", "OMPI_COMM_WORLD_RANK", "SLURM_PROCID") +_COUNTER_NAMES = ("hit", "miss", "create_graph", "cache_graph", "build_plans", "execute") +_BUILD_STAGES = ( + "validate", + "build_operation_graph", + "create_execution_plans", + "check_support", + "build_plans", +) + +_lock = threading.Lock() +_counters = defaultdict(lambda: defaultdict(int)) +_thread_counters = defaultdict(lambda: defaultdict(int)) +_thread_devices = defaultdict(set) +_stage_timings = defaultdict(lambda: [0, 0]) +_thread_ids: dict[int, int] = {} +_summary_registered = False + + +@lru_cache(maxsize=1) +def _configuration() -> tuple[int, bool, Optional[int]]: + """Return ``(level, selected, rank)`` for the current process.""" + + value = os.getenv("NVTE_FUSED_ATTN_CACHE_DEBUG", "0") + level_text, separator, rank_text = value.partition(":") + try: + level = int(level_text) + except ValueError: + level = 1 if level_text else 0 + if level <= 0: + return 0, False, None + + rank = None + for variable in _RANK_ENV_VARS: + rank_value = os.getenv(variable) + if rank_value: + try: + rank = int(rank_value) + except ValueError: + rank = 0 + break + if rank is None: + return level, True, None + if not separator: + return level, rank == 0, rank + if rank_text == "all": + return level, True, rank + selected_ranks = set() + for token in rank_text.split(","): + try: + selected_ranks.add(int(token)) + except ValueError: + continue + return level, rank in selected_ranks, rank + + +def enabled(*, trace: bool = False) -> bool: + """Whether cache diagnostics are enabled for this process and level.""" + + level, selected, _ = _configuration() + return selected and level >= (2 if trace else 1) + + +def _thread_id() -> int: + native_id = threading.get_ident() + with _lock: + thread_id = _thread_ids.get(native_id) + if thread_id is None: + thread_id = len(_thread_ids) + _thread_ids[native_id] = thread_id + return thread_id + + +def _rank_tag() -> str: + _, _, rank = _configuration() + return "" if rank is None else f"rank={rank} | " + + +def _write(text: str) -> None: + sys.stderr.write(text) + sys.stderr.flush() + + +def _counter_line( + thread_field: str, + device_field: str, + backend: str, + direction: str, + counters: dict[str, int], + event: str = "", +) -> str: + label = f"{backend} {direction}" + if event: + label += f" {event}" + values = ", ".join(f"{name}={counters.get(name, 0):4d}" for name in _COUNTER_NAMES) + return f"{_PREFIX} {_rank_tag()}{thread_field:<7} {device_field:<9} | {label:<24} | {values}\n" + + +def _register_summary() -> None: + global _summary_registered + with _lock: + if _summary_registered: + return + atexit.register(print_summary) + _summary_registered = True + + +def record_event( + backend: str, + direction: str, + event: str, + *, + device: Optional[int] = None, + key=None, +) -> None: + """Record one graph cache event and optionally emit its level-2 trace.""" + + if not enabled(): + return + event = event.lower() + if event not in _COUNTER_NAMES: + raise ValueError(f"Unknown fused-attention cache event {event!r}.") + _register_summary() + thread_id = _thread_id() + site = (backend, direction) + thread_site = (thread_id, backend, direction) + with _lock: + _counters[site][event] += 1 + _thread_counters[thread_site][event] += 1 + if device is not None: + _thread_devices[thread_id].add(device) + snapshot = dict(_counters[site]) + if not enabled(trace=True): + return + if event in ("hit", "miss") and key is not None: + _write( + f"{_PREFIX} {_rank_tag()}tid={thread_id:<3} dev={str(device):<3} | " + f"{backend} {direction} {event.upper():<12} | {key!r}\n" + ) + else: + _write( + _counter_line( + f"tid={thread_id}", + f"dev={device}", + backend, + direction, + snapshot, + event.upper(), + ) + ) + + +def record_lookup( + backend: str, + direction: str, + *, + hit: bool, + device: Optional[int] = None, + key=None, +) -> None: + """Record a cache hit or miss.""" + + record_event(backend, direction, "hit" if hit else "miss", device=device, key=key) + + +def record_build_time(backend: str, direction: str, stage: str, elapsed_ns: int) -> None: + """Accumulate CPU wall time for a cuDNN graph build stage.""" + + if not enabled(): + return + if stage not in _BUILD_STAGES: + raise ValueError(f"Unknown fused-attention graph build stage {stage!r}.") + _register_summary() + with _lock: + timing = _stage_timings[(backend, direction, stage)] + timing[0] += 1 + timing[1] += elapsed_ns + + +def build_recorder(backend: str, direction: str) -> Callable[[str, int], None]: + """Create the callback consumed by ``build_cudnn_graph``.""" + + def record(name: str, elapsed_ns: int) -> None: + if name in _BUILD_STAGES: + record_build_time(backend, direction, name, elapsed_ns) + else: + record_event(backend, direction, name.lower()) + + return record + + +def render_summary() -> str: + """Render the current diagnostic summary without writing it.""" + + if not enabled(): + return "" + with _lock: + counters = {site: dict(values) for site, values in _counters.items()} + thread_counters = {site: dict(values) for site, values in _thread_counters.items()} + thread_devices = {thread_id: set(values) for thread_id, values in _thread_devices.items()} + stage_timings = {site: tuple(values) for site, values in _stage_timings.items()} + + marker = f"{_PREFIX} {_rank_tag()}===== summary" + lines = [marker + " begin =====\n"] + for (thread_id, backend, direction), values in sorted(thread_counters.items()): + devices = thread_devices.get(thread_id, set()) + if not devices: + device_field = "dev=None" + elif len(devices) == 1: + device_field = f"dev={next(iter(devices))}" + else: + device_field = "dev=mixed" + lines.append(_counter_line(f"tid={thread_id}", device_field, backend, direction, values)) + for (backend, direction), values in sorted(counters.items()): + lines.append(_counter_line("tid=all", "dev=all", backend, direction, values)) + for (backend, direction, stage), (calls, elapsed_ns) in sorted(stage_timings.items()): + lines.append( + f"{_PREFIX} {_rank_tag()}{backend:<3} {direction:<3} {stage:<22} | " + f"calls={calls} | time={elapsed_ns / calls / 1e6:9.3f} ms/call\n" + ) + lines.append(marker + " end =====\n") + return "".join(lines) + + +def print_summary() -> None: + """Write the current summary to stderr if diagnostics are enabled.""" + + summary = render_summary() + if summary: + _write(summary) + + +def time_call(callback: Optional[Callable[[str, int], None]], stage: str, function): + """Call a graph-build stage and report its elapsed CPU wall time.""" + + if callback is None: + return function() + start = perf_counter_ns() + try: + return function() + finally: + callback(stage, perf_counter_ns() - start) + + +def _reset_for_tests() -> None: + """Clear diagnostic state. Intended only for GPU-independent tests.""" + + with _lock: + _counters.clear() + _thread_counters.clear() + _thread_devices.clear() + _stage_timings.clear() + _thread_ids.clear() + _configuration.cache_clear() diff --git a/transformer_engine/common/attention/cudnn.py b/transformer_engine/common/attention/cudnn.py new file mode 100644 index 00000000000..f78d01101b7 --- /dev/null +++ b/transformer_engine/common/attention/cudnn.py @@ -0,0 +1,688 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""Framework-independent cuDNN attention policy and graph-shape helpers.""" + +from __future__ import annotations + +from dataclasses import dataclass + + +@dataclass(frozen=True) +class AttentionLayout: + """Normalized layout properties used by the cuDNN support policy.""" + + qkv_format: str + q_format: str + kv_format: str + layout_group: str + is_qkvpacked: bool = False + + @property + def is_thd(self) -> bool: + """Return whether either Q or KV uses packed-token storage.""" + + return self.q_format == "thd" or self.kv_format == "thd" + + +@dataclass(frozen=True) +class FusedAttentionConfig: + """Normalized inputs to the shared FP16/BF16 cuDNN support policy.""" + + is_training: bool + q_dtype: str + kv_dtype: str + layout: AttentionLayout + bias_type: str + mask_type: str + softmax_type: str + dropout: float + num_attn_heads: int + num_gqa_groups: int + max_seqlen_q: int + max_seqlen_kv: int + head_dim_qk: int + head_dim_v: int + window_size: tuple[int, int] + return_max_logit: bool + cuda_graph: bool + deterministic: bool + cudnn_version: tuple[int, int, int] + sm_arch: int + + +@dataclass(frozen=True) +class FusedAttentionSupport: + """Result of checking a normalized attention configuration.""" + + supported: bool + reason: str = "" + warning: str | None = None + + +@dataclass(frozen=True) +class AttentionMask: + """Framework-neutral interpretation of mask and sliding-window options.""" + + causal: bool + bottom_right: bool + padding: bool + bottom_right_diagonal: bool + window_left: int + window_right: int + + +def parse_attention_layout(qkv_layout: str) -> AttentionLayout: + """Normalize a TE QKV layout string for framework-independent policy checks.""" + + paged = qkv_layout.startswith("paged_kv_") + layout = qkv_layout.removeprefix("paged_kv_") + components = layout.split("_") + + def tensor_format(component: str) -> str: + return "".join(char for char in component if char.isalpha()) + + q_format = tensor_format(components[0]) + kv_format = tensor_format(components[-1]) if len(components) > 1 else q_format + qkv_format = q_format if q_format == kv_format else f"{q_format}_2{kv_format}" + if paged: + layout_group = "paged_separate" + elif len(components) == 1 and "3" in components[0]: + layout_group = "h3d" if "h3d" in components[0] else "3hd" + elif len(components) == 2: + layout_group = "hd_h2d" if "h2d" in components[1] else "hd_2hd" + elif q_format == "bhsd": + layout_group = "sd_sd_sd" + else: + layout_group = "separate" + return AttentionLayout( + qkv_format=qkv_format, + q_format=q_format, + kv_format=kv_format, + layout_group=layout_group, + is_qkvpacked=layout_group in ("3hd", "h3d"), + ) + + +def encode_cudnn_version(version: tuple[int, int, int]) -> int: + """Encode a cuDNN backend version using its native integer convention.""" + + major, minor, patch = (int(part) for part in version) + magnitude = 1000 if major < 9 else 10000 + return major * magnitude + minor * 100 + patch + + +def round_up(value: int, multiple: int) -> int: + """Round ``value`` up to a positive multiple.""" + + return (int(value) + int(multiple) - 1) // int(multiple) * int(multiple) + + +def ragged_token_bucket(tokens: int) -> int: + """Return the cuDNN graph bucket for a packed-token extent.""" + + tokens = int(tokens) + if tokens <= 1024: + return 1024 + if tokens <= 32768: + return 1 << (tokens - 1).bit_length() + return round_up(tokens, 32768) + + +def ragged_batch_bucket(batch: int) -> int: + """Return the cuDNN graph bucket for a ragged batch extent.""" + + batch = int(batch) + if batch <= 32: + return 32 + if batch <= 512: + return 1 << (batch - 1).bit_length() + return round_up(batch, 512) + + +def normalize_attention_mask( + *, + causal: bool, + bottom_right: bool, + padding: bool, + bottom_right_diagonal: bool, + window_size: tuple[int, int], + max_seqlen_q: int, + max_seqlen_kv: int, +) -> AttentionMask: + """Normalize equivalent causal and bottom-right mask configurations.""" + + if bottom_right and max_seqlen_q == max_seqlen_kv and not padding: + causal = True + bottom_right = False + bottom_right_diagonal = False + return AttentionMask( + causal=bool(causal), + bottom_right=bool(bottom_right), + padding=bool(padding), + bottom_right_diagonal=bool(bottom_right_diagonal), + window_left=int(window_size[0]), + window_right=int(window_size[1]), + ) + + +def cudnn_mask_options( + *, + causal: bool, + bottom_right: bool, + padding: bool, + bottom_right_diagonal: bool, + window_size: tuple[int, int], + max_seqlen_q: int, + max_seqlen_kv: int, + cudnn_version: tuple[int, int, int], +) -> dict[str, bool | int | str]: + """Return canonical cuDNN SDPA mask options using framework-neutral values.""" + + mask = normalize_attention_mask( + causal=causal, + bottom_right=bottom_right, + padding=padding, + bottom_right_diagonal=bottom_right_diagonal, + window_size=window_size, + max_seqlen_q=max_seqlen_q, + max_seqlen_kv=max_seqlen_kv, + ) + version = encode_cudnn_version(cudnn_version) + options: dict[str, bool | int | str] = { + "diagonal_alignment": ("bottom_right" if mask.bottom_right_diagonal else "top_left"), + "is_padding": mask.padding, + } + if version < 90600: + options["use_causal_mask"] = mask.causal + options["use_causal_mask_bottom_right"] = mask.bottom_right + if version >= 90200 and mask.window_left != -1: + options["diagonal_band_left_bound"] = mask.window_left + 1 + if version >= 90600: + if mask.window_right != -1: + options["diagonal_band_right_bound"] = mask.window_right + elif mask.causal or mask.bottom_right: + options["diagonal_band_right_bound"] = 0 + return options + + +def requires_64bit_ragged_offset( + layout: AttentionLayout, + num_attn_heads: int, + num_gqa_groups: int, + max_seqlen_q: int, + max_seqlen_kv: int, + head_dim_qk: int, + head_dim_v: int, +) -> bool: + """Return whether legacy THD element offsets can overflow signed int32.""" + + if layout.qkv_format != "thd": + return False + if layout.layout_group in ("3hd", "h3d", "qkv_packed"): + q = k = v = 3 * num_attn_heads * head_dim_qk * max_seqlen_q + elif layout.layout_group in ("hd_2hd", "hd_h2d", "kv_packed"): + q = num_attn_heads * head_dim_qk * max_seqlen_q + k = v = 2 * num_gqa_groups * head_dim_qk * max_seqlen_kv + else: + q = num_attn_heads * head_dim_qk * max_seqlen_q + k = num_gqa_groups * head_dim_qk * max_seqlen_kv + v = num_gqa_groups * head_dim_v * max_seqlen_kv + output = num_attn_heads * head_dim_qk * max_seqlen_q + return max(q, k, v, output) > 2**31 - 1 + + +def _unsupported(reason: str, warning: str | None = None) -> FusedAttentionSupport: + return FusedAttentionSupport(False, reason, warning) + + +def check_f16_fused_attention_support( + config: FusedAttentionConfig, +) -> FusedAttentionSupport: + """Check the shared FP16/BF16 cuDNN fused-attention compatibility policy.""" + + if config.q_dtype != config.kv_dtype: + return _unsupported("Q and KV must have the same data type") + if config.q_dtype not in ("float16", "bfloat16"): + return _unsupported("only FP16 and BF16 are supported") + + version = encode_cudnn_version(config.cudnn_version) + arch = int(config.sm_arch) + layout = config.layout + is_thd = layout.is_thd + is_training = bool(config.is_training) + sq = int(config.max_seqlen_q) + skv = int(config.max_seqlen_kv) + h = int(config.num_attn_heads) + hg = int(config.num_gqa_groups) + dqk = int(config.head_dim_qk) + dv = int(config.head_dim_v) + dropout = float(config.dropout) + bias = config.bias_type + mask = config.mask_type + softmax = config.softmax_type + left, right = config.window_size + + if version < 8900: + return _unsupported( + "cuDNN is older than 8.9.0", + "FP16/BF16 fused attention requires cuDNN 8.9.0 or newer", + ) + + architecture_ok = ( + (version < 8903 and arch in (80, 90)) + or (version >= 8903 and 80 <= arch < 100) + or (version >= 90700 and arch >= 100) + ) + if not architecture_ok: + return _unsupported("device architecture is not supported") + if version < 90000 and (sq % 64 or skv % 64): + return _unsupported("sequence lengths must be multiples of 64") + if version < 8907 and h != hg: + return _unsupported("GQA requires cuDNN 8.9.7 or newer") + if dqk % 8 or dv % 8: + return _unsupported("head dimensions must be multiples of 8") + + standard_dim = dqk <= 128 and dv <= 128 + hopper_large_dim = ( + dqk <= 256 + and dv <= 256 + and ( + (not is_training and arch == 90 and version >= 90100) + or (is_training and arch == 90 and version >= 90500) + ) + ) + blackwell_fwd_any_dim = ( + not is_training + and arch >= 100 + and version >= 90900 + and sq > 1 + and layout.layout_group != "paged_separate" + ) + generic_fwd_any_dim = ( + not is_training + and version >= 91002 + and ( + layout.layout_group == "paged_separate" + or sq > 1 + or (sq == 1 and mask not in ("causal", "padding_causal")) + ) + ) + blackwell_mla_bwd = ( + dqk == 192 and dv == 128 and is_training and arch >= 100 and version >= 91100 + ) + blackwell_d256_bwd = ( + dqk == 256 + and dv == 256 + and is_training + and 100 <= arch < 110 + and version >= (92500 if is_thd else 92300) + and layout.layout_group != "paged_separate" + and bias == "no_bias" + and dropout == 0.0 + and softmax == "vanilla" + and ( + (left == -1 and right == -1) + or ( + mask + in ( + "causal", + "padding_causal", + "causal_bottom_right", + "padding_causal_bottom_right", + ) + and right in (-1, 0) + ) + ) + ) + if not ( + standard_dim + or hopper_large_dim + or blackwell_fwd_any_dim + or generic_fwd_any_dim + or blackwell_mla_bwd + or blackwell_d256_bwd + ): + return _unsupported("head dimensions are not supported") + unsupported_hopper_bwd_dims = ( + version >= 91100 + and is_training + and arch == 90 + and dqk >= 128 + and dv >= 128 + and (dqk, dv) != (192, 128) + and dqk != dv + ) + if unsupported_hopper_bwd_dims: + return _unsupported("this Hopper backward head-dimension combination is unsupported") + + alibi_supported = ( + bias == "alibi" + and version >= 8906 + and arch >= 90 + and mask + not in ( + "no_mask", + "padding", + "padding_causal", + "padding_causal_bottom_right", + ) + ) + post_scale_bias_supported = bias == "post_scale_bias" and ( + (version >= 8906 and arch >= 90) or (version >= 90000 and arch >= 80) + ) + if bias != "no_bias" and not alibi_supported and not post_scale_bias_supported: + return _unsupported("attention bias is not supported") + + standard_format = layout.qkv_format in ("sbhd", "bshd") + basic_masks = mask in ("no_mask", "causal", "padding", "padding_causal") + aligned_bottom_right = sq % 64 == 0 and skv % 64 == 0 and sq <= skv + no_bias_no_dropout = bias == "no_bias" and dropout == 0.0 + mask_ok = version < 8906 and mask == "causal" + if version >= 8906 and standard_format and basic_masks: + mask_ok = True + if ( + version >= 90100 + and layout.qkv_format == "thd" + and mask + in ( + "padding", + "padding_causal", + ) + ): + mask_ok = True + if ( + version >= 90300 + and standard_format + and mask == "causal_bottom_right" + and aligned_bottom_right + and no_bias_no_dropout + ): + mask_ok = True + paged_mask_supported = mask in ("padding", "padding_causal") or ( + mask == "padding_causal_bottom_right" and aligned_bottom_right + ) + if ( + version >= 90500 + and layout.layout_group == "paged_separate" + and paged_mask_supported + and no_bias_no_dropout + ): + mask_ok = True + if ( + version >= 90600 + and mask == "padding_causal_bottom_right" + and aligned_bottom_right + and no_bias_no_dropout + ): + mask_ok = True + if version >= 90700: + modern_mask_ok = ( + mask in ("no_mask", "causal") + or ( + mask in ("padding", "padding_causal", "padding_causal_bottom_right") + and bias == "no_bias" + and dropout == 0.0 + ) + or (mask in ("causal_bottom_right", "padding_causal_bottom_right") and sq <= skv) + ) + mask_ok = mask_ok or modern_mask_ok + if not mask_ok: + return _unsupported("attention mask is not supported") + if mask in ("padding", "padding_causal") and bias == "post_scale_bias": + return _unsupported("post-scale bias cannot be combined with this padding mask") + + format_ok = ( + layout.qkv_format in ("sbhd", "bshd", "bhsd") + or ( + layout.qkv_format == "thd" + and arch >= 90 + and ((version >= 90100 and h == hg) or version >= 90600) + ) + or ( + layout.q_format in ("sbhd", "bshd", "bhsd", "thd") + and layout.kv_format in ("sbhd", "bshd", "bhsd", "thd") + and (layout.q_format != "thd" or arch >= 90) + and (layout.kv_format != "thd" or arch >= 90) + and version >= 90700 + ) + ) + if not format_ok: + return _unsupported("QKV format is not supported") + + pre_902_window = version < 90200 and left == -1 and right in (-1, 0) + v902_window = version >= 90200 and ( + (left == -1 and right == -1 and mask == "no_mask") + or ( + left >= -1 + and right == 0 + and (mask in ("no_mask", "causal") or (mask == "causal_bottom_right" and sq == skv)) + and sq <= skv + and dropout == 0.0 + and bias == "no_bias" + and standard_format + ) + ) + bottom_right_swa_supported = ( + mask not in ("causal_bottom_right", "padding_causal_bottom_right") + or arch < 100 + or sq == skv + or version > 90700 + ) + v906_window = version >= 90600 and ( + (left == -1 and right in (-1, 0)) + or ( + left >= -1 + and right >= -1 + and ( + mask + in ( + "no_mask", + "padding", + "padding_causal", + "causal_bottom_right", + "padding_causal_bottom_right", + ) + or mask == "causal" + ) + ) + and sq <= skv + and bias == "no_bias" + and dropout == 0.0 + and bottom_right_swa_supported + ) + window_ok = pre_902_window or v902_window or v906_window + if not window_ok: + return _unsupported("sliding-window configuration is not supported") + + requires_i64 = requires_64bit_ragged_offset( + layout, + h, + hg, + sq, + skv, + dqk, + dv, + ) + if requires_i64 and version < 90500: + return _unsupported("ragged offsets require int64 support") + if version == 91000: + return _unsupported("cuDNN 9.10.0 has known SDPA issues") + if version < 91301 and softmax != "vanilla": + return _unsupported("this softmax type requires cuDNN 9.13.1 or newer") + if config.return_max_logit and version < 92100: + return _unsupported("returning max logits requires cuDNN 9.21 or newer") + if arch >= 100 and is_training: + if config.deterministic: + if version < 91801 or dropout != 0.0 or bias != "no_bias": + return _unsupported("deterministic Blackwell backward is not supported") + elif dropout != 0.0 and bias != "no_bias": + return _unsupported("Blackwell backward does not support dropout with bias") + + if ( + version == 91400 + and skv > 1024 + and left != -1 + and mask not in ("causal", "causal_bottom_right") + ): + return _unsupported( + "cuDNN 9.14.0 does not support this non-causal sliding window", + "This non-causal sliding-window configuration requires cuDNN > 9.14.0", + ) + unsupported_cuda_graph_bwd = ( + version <= 91500 + and is_training + and standard_format + and skv % 128 != 0 + and config.cuda_graph + and mask not in ("padding", "padding_causal", "padding_causal_bottom_right") + ) + if unsupported_cuda_graph_bwd: + return _unsupported( + "this backward CUDA-graph configuration requires cuDNN 9.15.1", + "This backward CUDA-graph configuration requires cuDNN 9.15.1 or newer", + ) + if arch == 120: + if version < 91801: + return _unsupported( + "SM120 requires cuDNN 9.18.1", + "SM120 fused attention requires cuDNN 9.18.1 or newer", + ) + if config.deterministic and is_training: + return _unsupported( + "deterministic backward is not supported on SM120", + "Deterministic fused-attention backward is not supported on SM120", + ) + if is_thd and layout.is_qkvpacked: + return _unsupported( + "QKV-packed THD attention is not supported on SM120", + "T3HD/TH3D fused attention is not supported on SM120", + ) + + return FusedAttentionSupport(True) + + +def check_fp8_fused_attention_support( + config: FusedAttentionConfig, + *, + scaling_mode: str | None = None, +) -> FusedAttentionSupport: + """Check the shared cuDNN FP8/MXFP8 fused-attention compatibility policy.""" + + if config.q_dtype != config.kv_dtype: + return _unsupported("Q and KV must have the same data type") + if config.q_dtype not in ("float8_e4m3", "float8_e5m2"): + return _unsupported("only FP8 E4M3 and E5M2 are supported") + + version = encode_cudnn_version(config.cudnn_version) + arch = int(config.sm_arch) + layout = config.layout + sq = int(config.max_seqlen_q) + skv = int(config.max_seqlen_kv) + dqk = int(config.head_dim_qk) + dv = int(config.head_dim_v) + mask = config.mask_type + + if scaling_mode not in (None, "delayed", "current", "mxfp8"): + return _unsupported(f"unknown FP8 attention scaling mode {scaling_mode!r}") + if arch < 90: + return _unsupported("FP8 attention requires SM90 or newer") + if arch >= 120: + return _unsupported("FP8 attention is not supported on SM120 or newer") + if config.is_training and config.deterministic and version < 91900: + return _unsupported("deterministic FP8 attention backward requires cuDNN 9.19 or newer") + if scaling_mode == "current": + if arch < 100: + return _unsupported("FP8 current-scaling attention requires SM100 or newer") + if version < 91400: + return _unsupported("FP8 current-scaling attention requires cuDNN 9.14 or newer") + if scaling_mode == "mxfp8": + if arch < 100: + return _unsupported("MXFP8 attention requires SM100 or newer") + if version < 92100: + return _unsupported("MXFP8 attention requires cuDNN 9.21 or newer") + if version in (92300, 92301): + return _unsupported("cuDNN 9.23.0 and 9.23.1 have known MXFP8 SDPA issues") + if config.bias_type != "no_bias": + return _unsupported("FP8 attention does not support attention bias") + if config.return_max_logit: + return _unsupported("FP8 attention does not support returning max logits") + if version == 91000: + return _unsupported("cuDNN 9.10.0 has known SDPA issues") + if ( + requires_64bit_ragged_offset( + layout, + config.num_attn_heads, + config.num_gqa_groups, + sq, + skv, + dqk, + dv, + ) + and version < 90500 + ): + return _unsupported("FP8 attention requires cuDNN 9.5 for 64-bit ragged offsets") + + is_thd = layout.qkv_format == "thd" + if is_thd: + if version < 92300: + return _unsupported("FP8 THD attention requires cuDNN 9.23 or newer") + if mask not in ("padding", "padding_causal", "padding_causal_bottom_right"): + return _unsupported("FP8 THD attention requires a padding mask") + if config.is_training and arch < 100: + return _unsupported("FP8 THD attention backward requires SM100 or newer") + if config.is_training and config.softmax_type != "vanilla" and version < 92600: + return _unsupported("FP8 THD sink-token backward requires cuDNN 9.26 or newer") + if arch >= 100 and (dqk > 128 or dv > 128): + return _unsupported("FP8 THD attention supports head dimensions up to 128 on SM100+") + + shape_mask_ok = ( + ( + version >= 90201 + and arch < 100 + and sq % 128 == 0 + and skv % 128 == 0 + and dqk == 128 + and dv == 128 + and mask in ("causal", "no_mask") + ) + or ( + version >= 90700 + and ( + (arch < 100 and not config.is_training and dqk <= 256 and dv <= 256) + or (arch < 100 and config.is_training and dqk == 128 and dv == 128) + or (arch >= 100 and dqk <= 128 and dv <= 128) + ) + and dqk % 16 == 0 + and dv % 16 == 0 + and ( + mask in ("no_mask", "causal", "padding", "padding_causal") + or (arch >= 100 and mask == "padding_causal_bottom_right") + ) + ) + or ( + version >= 92100 + and arch >= 100 + and dqk <= 192 + and dv <= 128 + and dqk % 16 == 0 + and dv % 16 == 0 + and mask in ("no_mask", "causal", "causal_bottom_right") + ) + ) + if not shape_mask_ok: + return _unsupported("FP8 attention shape or mask is not supported") + + format_softmax_ok = ( + ( + version < 92100 + and layout.qkv_format in ("bshd", "sbhd") + and config.softmax_type == "vanilla" + ) + or (version >= 92100 and layout.qkv_format in ("bshd", "sbhd", "bhsd")) + or is_thd + ) + if not format_softmax_ok: + return _unsupported("FP8 attention layout or softmax type is not supported") + return FusedAttentionSupport(True) diff --git a/transformer_engine/common/attention/fp8.py b/transformer_engine/common/attention/fp8.py new file mode 100644 index 00000000000..5675a68d82c --- /dev/null +++ b/transformer_engine/common/attention/fp8.py @@ -0,0 +1,171 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""Framework-neutral helpers for cuDNN FP8 attention graph construction.""" + +from __future__ import annotations + +from collections.abc import Mapping +from dataclasses import dataclass +from typing import Any + +from .cudnn import round_up + + +@dataclass(frozen=True) +class FP8AttentionGraphConfig: + """Static choices that select a cuDNN FP8 attention graph family.""" + + mode: str + name: str + + def __post_init__(self): + if self.mode not in ("delayed", "current", "mxfp8"): + raise ValueError(f"Unknown FP8 attention scaling mode {self.mode!r}.") + + @property + def is_mxfp8(self) -> bool: + """Return whether the graph uses microscaling FP8 nodes.""" + + return self.mode == "mxfp8" + + +def attention_format_stride( + batch: int, heads: int, seqlen: int, dim: int, tensor_format: str +) -> tuple[int, int, int, int]: + """Describe a contiguous TE attention tensor as logical BHSD.""" + + if tensor_format in ("bshd", "thd"): + return (seqlen * heads * dim, dim, heads * dim, 1) + if tensor_format == "sbhd": + return (heads * dim, dim, batch * heads * dim, 1) + if tensor_format == "bhsd": + return (heads * seqlen * dim, seqlen * dim, dim, 1) + raise ValueError(f"Unsupported FP8 tensor format {tensor_format!r}.") + + +def mxfp8_padded_sizes(s_q: int, s_kv: int, d_qk: int, d_v: int) -> dict[str, int]: + """Return the padded data and E8M0-scale dimensions required by cuDNN MXFP8.""" + + return { + "s_q_padded": round_up(s_q, 128), + "s_kv_padded": round_up(s_kv, 128), + "s_q_scale_padded": round_up((s_q + 31) // 32, 4), + "s_kv_scale_padded": round_up((s_kv + 31) // 32, 4), + "d_qk_padded": round_up(d_qk, 128), + "d_v_padded": round_up(d_v, 128), + "d_qk_scale_padded": round_up((d_qk + 31) // 32, 4), + "d_v_scale_padded": round_up((d_v + 31) // 32, 4), + } + + +def build_fp8_forward_operation( + graph: Any, + tensors: Mapping[str, Any], + options: Mapping[str, Any], + config: FP8AttentionGraphConfig, +) -> dict[str, Any]: + """Add the selected FP8 SDPA forward operation to a cuDNN frontend graph.""" + + kwargs = dict(options) + if config.is_mxfp8: + output, stats, amax_o = graph.sdpa_mxfp8( + tensors["q"], + tensors["k"], + tensors["v"], + tensors["descale_q"], + tensors["descale_k"], + tensors["descale_v"], + name=config.name, + **kwargs, + ) + return {"output": output, "stats": stats, "amax_o": amax_o} + + output, stats, amax_s, amax_o = graph.sdpa_fp8( + tensors["q"], + tensors["k"], + tensors["v"], + tensors["descale_q"], + tensors["descale_k"], + tensors["descale_v"], + tensors["descale_s"], + tensors["scale_s"], + tensors["scale_o"], + name=config.name, + **kwargs, + ) + return { + "output": output, + "stats": stats, + "amax_s": amax_s, + "amax_o": amax_o, + } + + +def build_fp8_backward_operation( + graph: Any, + tensors: Mapping[str, Any], + options: Mapping[str, Any], + config: FP8AttentionGraphConfig, +) -> dict[str, Any]: + """Add the selected FP8 SDPA backward operation to a cuDNN frontend graph.""" + + kwargs = dict(options) + if config.is_mxfp8: + outputs = graph.sdpa_mxfp8_backward( + tensors["q"], + tensors["q_t"], + tensors["k"], + tensors["k_t"], + tensors["v"], + tensors["o"], + tensors["do_f16"], + tensors["do"], + tensors["do_t"], + tensors["stats"], + tensors["descale_q"], + tensors["descale_q_t"], + tensors["descale_k"], + tensors["descale_k_t"], + tensors["descale_v"], + tensors["descale_do"], + tensors["descale_do_t"], + name=config.name, + **kwargs, + ) + dq, dk, dv, *amax = outputs + return {"dq": dq, "dk": dk, "dv": dv, "amax": tuple(amax)} + + outputs = graph.sdpa_fp8_backward( + tensors["q"], + tensors["k"], + tensors["v"], + tensors["o"], + tensors["do"], + tensors["stats"], + tensors["descale_q"], + tensors["descale_k"], + tensors["descale_v"], + tensors["descale_o"], + tensors["descale_do"], + tensors["descale_s"], + tensors["descale_dp"], + tensors["scale_s"], + tensors["scale_dq"], + tensors["scale_dk"], + tensors["scale_dv"], + tensors["scale_dp"], + name=config.name, + **kwargs, + ) + dq, dk, dv, amax_dq, amax_dk, amax_dv, amax_dp = outputs + return { + "dq": dq, + "dk": dk, + "dv": dv, + "amax_dq": amax_dq, + "amax_dk": amax_dk, + "amax_dv": amax_dv, + "amax_dp": amax_dp, + } diff --git a/transformer_engine/common/attention/score_mod.py b/transformer_engine/common/attention/score_mod.py new file mode 100644 index 00000000000..dc9c16f71bc --- /dev/null +++ b/transformer_engine/common/attention/score_mod.py @@ -0,0 +1,126 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""Framework-independent cache-key policy for cuDNN score-modification graphs.""" + +from __future__ import annotations + +import inspect +from collections.abc import Callable, Mapping +from typing import Any + + +class UncacheableScoreModKey: + """Identity key for callbacks whose graph topology cannot be cached safely.""" + + def __hash__(self): + return id(self) + + def __eq__(self, other): + return self is other + + +UNCACHEABLE_SCORE_MOD = UncacheableScoreModKey() + + +def is_uncacheable_score_mod_key(key: Any) -> bool: + """Return whether a score-modification graph key disables caching.""" + + return isinstance(key, UncacheableScoreModKey) + + +def freeze_score_mod_cache_key(value: Any, *, is_array: Callable[[Any], bool]) -> Any: + """Convert a user-provided score-modification key into a hashable structure.""" + + if is_array(value): + raise TypeError( + "score_mod_graph_cache_key() must not include tensors. Pass runtime tensors " + "through score_mod_tensors or score_mod_bprop_tensors instead." + ) + if isinstance(value, Mapping): + items = ( + ( + freeze_score_mod_cache_key(key, is_array=is_array), + freeze_score_mod_cache_key(item, is_array=is_array), + ) + for key, item in value.items() + ) + return tuple(sorted(items, key=repr)) + if isinstance(value, (list, tuple)): + return tuple(freeze_score_mod_cache_key(item, is_array=is_array) for item in value) + if isinstance(value, (set, frozenset)): + items = (freeze_score_mod_cache_key(item, is_array=is_array) for item in value) + return tuple(sorted(items, key=repr)) + try: + hash(value) + except TypeError as exc: + raise TypeError( + "score_mod_graph_cache_key() must return a hashable value or a nested " + "combination of mapping/list/tuple/set values." + ) from exc + return value + + +def _explicit_cache_key(callback_owner: Any, *, is_array: Callable[[Any], bool]) -> Any | None: + explicit_key = getattr(callback_owner, "score_mod_graph_cache_key", None) + if explicit_key is None: + return None + explicit_key = explicit_key() if callable(explicit_key) else explicit_key + return freeze_score_mod_cache_key(explicit_key, is_array=is_array) + + +def score_mod_callback_cache_key( + callback: Callable | None, + *, + is_array: Callable[[Any], bool], + uncacheable_key_factory: Callable[[], Any] | None = None, +) -> Any: + """Create a stable graph key for a score-modification callable. + + Stateful callables must provide ``score_mod_graph_cache_key``. Stateless named + functions use their qualified name, while lambdas use their code object so two + lambdas in the same module cannot collide. + """ + + def uncacheable_key(): + if uncacheable_key_factory is None: + return UNCACHEABLE_SCORE_MOD + return uncacheable_key_factory() + + if callback is None: + return None + self_obj = getattr(callback, "__self__", None) + func_obj = getattr(callback, "__func__", None) + if self_obj is not None and func_obj is not None: + explicit_key = _explicit_cache_key(self_obj, is_array=is_array) + if explicit_key is None: + return uncacheable_key() + return ( + "bound_method", + type(self_obj), + func_obj.__module__, + func_obj.__qualname__, + explicit_key, + ) + + explicit_key = _explicit_cache_key(callback, is_array=is_array) + if explicit_key is not None: + return ( + "callable", + type(callback), + getattr(callback, "__module__", None), + getattr(callback, "__qualname__", None), + explicit_key, + ) + + if ( + inspect.isfunction(callback) + and callback.__closure__ is None + and "" not in callback.__qualname__ + ): + if callback.__name__ == "" or not callback.__qualname__: + return ("function", callback.__module__, callback.__code__) + return ("function", callback.__module__, callback.__qualname__) + + return uncacheable_key() diff --git a/transformer_engine/common/cudnn_frontend.py b/transformer_engine/common/cudnn_frontend.py new file mode 100644 index 00000000000..c735025b5dc --- /dev/null +++ b/transformer_engine/common/cudnn_frontend.py @@ -0,0 +1,76 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""Framework-independent helpers for constructing cuDNN frontend graphs.""" + +from __future__ import annotations + +import importlib +from typing import Any, Callable, Optional + +from transformer_engine.common.attention.cache_debug import time_call + + +def import_cudnn_frontend(*, feature: str, requirement: str): + """Import the cuDNN frontend Python package lazily.""" + + try: + return importlib.import_module("cudnn") + except ImportError as exc: + raise ImportError( + f"{feature} requires the cuDNN frontend Python package. Install {requirement}." + ) from exc + + +def make_cudnn_graph( + cudnn, + io_dtype: Any, + *, + name: str | None = None, + handle: Any = None, +): + """Create a cuDNN graph with TE's standard compute and intermediate types.""" + + kwargs = { + "io_data_type": io_dtype, + "intermediate_data_type": cudnn.data_type.FLOAT, + "compute_data_type": cudnn.data_type.FLOAT, + } + if name is not None: + kwargs["name"] = name + if handle is not None: + kwargs["handle"] = handle + return cudnn.pygraph(**kwargs) + + +def build_cudnn_graph( + cudnn, + graph, + *, + description: str, + debug_callback: Optional[Callable[[str, int], None]] = None, +) -> int: + """Validate and plan a graph, returning a nonzero workspace size.""" + + if debug_callback is not None: + debug_callback("CREATE_GRAPH", 0) + time_call(debug_callback, "validate", graph.validate) + time_call(debug_callback, "build_operation_graph", graph.build_operation_graph) + try: + time_call( + debug_callback, + "create_execution_plans", + lambda: graph.create_execution_plans([cudnn.heur_mode.A, cudnn.heur_mode.FALLBACK]), + ) + time_call(debug_callback, "check_support", graph.check_support) + except cudnn.cudnnGraphNotSupportedError as exc: + raise RuntimeError(f"cuDNN {description} graph is not supported: {exc}") from exc + time_call( + debug_callback, + "build_plans", + lambda: graph.build_plans(cudnn.build_plan_policy.HEURISTICS_CHOICE), + ) + if debug_callback is not None: + debug_callback("BUILD_PLANS", 0) + return max(int(graph.get_workspace_size()), 1) diff --git a/transformer_engine/common/cudnn_utils.cpp b/transformer_engine/common/cudnn_utils.cpp index 05ee35ccc7f..cf864a4b075 100644 --- a/transformer_engine/common/cudnn_utils.cpp +++ b/transformer_engine/common/cudnn_utils.cpp @@ -11,29 +11,6 @@ namespace transformer_engine { -// get cuDNN data type -cudnnDataType_t get_cudnn_dtype(const transformer_engine::DType t) { - using namespace transformer_engine; - switch (t) { - case DType::kInt32: - return CUDNN_DATA_INT32; - case DType::kInt64: - return CUDNN_DATA_INT64; - case DType::kFloat16: - return CUDNN_DATA_HALF; - case DType::kFloat32: - return CUDNN_DATA_FLOAT; - case DType::kBFloat16: - return CUDNN_DATA_BFLOAT16; - case DType::kFloat8E4M3: - return CUDNN_DATA_FP8_E4M3; - case DType::kFloat8E5M2: - return CUDNN_DATA_FP8_E5M2; - default: - NVTE_ERROR("Invalid cuDNN data type. \n"); - } -} - // get cuDNN data type cudnn_frontend::DataType_t get_cudnn_fe_dtype(const transformer_engine::DType t) { using namespace transformer_engine; diff --git a/transformer_engine/common/cudnn_utils.h b/transformer_engine/common/cudnn_utils.h index 0777d1e03d0..beba7de0916 100644 --- a/transformer_engine/common/cudnn_utils.h +++ b/transformer_engine/common/cudnn_utils.h @@ -23,8 +23,6 @@ void CreateCuDNNHandle(cudnnHandle_t* handle); } // namespace detail -cudnnDataType_t get_cudnn_dtype(const transformer_engine::DType t); - cudnn_frontend::DataType_t get_cudnn_fe_dtype(const transformer_engine::DType t); using cudnnExecutionPlanManager = detail::HandleManager; diff --git a/transformer_engine/common/fused_attn/config_and_params.cpp b/transformer_engine/common/fused_attn/config_and_params.cpp deleted file mode 100644 index edea00b275f..00000000000 --- a/transformer_engine/common/fused_attn/config_and_params.cpp +++ /dev/null @@ -1,1232 +0,0 @@ -/************************************************************************* - * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. - * - * See LICENSE for license information. - ************************************************************************/ - -#include "config_and_params.h" - -#include -#include - -#include -#include -#include -#include - -#include "../common.h" -#include "../util/cuda_runtime.h" - -namespace { - -void bool_to_uint8(bool in, void *out) { - *reinterpret_cast(out) = static_cast(in); -} - -void uint8_to_bool(const void *in, bool &out) { - out = static_cast(*reinterpret_cast(in)); -} - -} // namespace - -namespace transformer_engine { - -namespace fused_attn { - -// Forward declarations -size_t get_max_batch_size(size_t batch_size); -size_t get_max_tokens(size_t num_tokens); -DType get_ragged_offset_dtype(NVTE_QKV_Layout_Group layout_group, int64_t num_attn_heads, - int64_t num_gqa_groups, int64_t max_seqlen_q, int64_t max_seqlen_kv, - int64_t head_dim_qk, int64_t head_dim_v); - -void FusedAttnConfig::derive() { - if (is_derived) return; - - // Common attributes - qkv_format = nvte_get_qkv_format(qkv_layout); - q_format = nvte_get_q_format(qkv_layout); - kv_format = nvte_get_kv_format(qkv_layout); - const NVTE_QKV_Layout_Group layout_group = nvte_get_qkv_layout_group(qkv_layout); - is_paged_kv = (layout_group == NVTE_QKV_Layout_Group::NVTE_Paged_KV_HD_HD_HD); - is_ragged_q = (q_format == NVTE_QKV_Format::NVTE_THD); - is_ragged_kv = (kv_format == NVTE_QKV_Format::NVTE_THD); - is_padding = (attn_mask_type == NVTE_Mask_Type::NVTE_PADDING_MASK) || - (attn_mask_type == NVTE_Mask_Type::NVTE_PADDING_CAUSAL_MASK) || - (attn_mask_type == NVTE_Mask_Type::NVTE_PADDING_CAUSAL_BOTTOM_RIGHT_MASK); - is_causal = (attn_mask_type == NVTE_Mask_Type::NVTE_CAUSAL_MASK) || - (attn_mask_type == NVTE_Mask_Type::NVTE_PADDING_CAUSAL_MASK); - is_causal_bottom_right = - (attn_mask_type == NVTE_Mask_Type::NVTE_CAUSAL_BOTTOM_RIGHT_MASK) || - (attn_mask_type == NVTE_Mask_Type::NVTE_PADDING_CAUSAL_BOTTOM_RIGHT_MASK); - if (is_causal_bottom_right && !bottom_right_diagonal) { - bottom_right_diagonal = true; - } - if (is_causal && bottom_right_diagonal) { - bottom_right_diagonal = false; - } - const bool has_window = window_size_left != -1 || window_size_right != -1; - if (!is_causal && !is_causal_bottom_right && !has_window) { - bottom_right_diagonal = false; - } - is_bias = (bias_type == NVTE_Bias_Type::NVTE_POST_SCALE_BIAS); - is_alibi = (bias_type == NVTE_Bias_Type::NVTE_ALIBI); - is_softmax_offset = (softmax_type != NVTE_Softmax_Type::NVTE_VANILLA_SOFTMAX); - is_dropout = is_training && dropout != 0.0f; - - // Determine the FP8 recipe - const bool is_o_in_fp8 = (o_dtype == kNVTEFloat8E4M3 || o_dtype == kNVTEFloat8E5M2); - const bool is_dqkv_in_fp8 = (dqkv_dtype == kNVTEFloat8E4M3 || dqkv_dtype == kNVTEFloat8E5M2); - is_o_in_f16 = (o_dtype == kNVTEFloat16 || o_dtype == kNVTEBFloat16); - const bool is_dqkv_in_f16 = (dqkv_dtype == kNVTEFloat16 || dqkv_dtype == kNVTEBFloat16); - is_tensor_scaling = (scaling_mode == NVTE_DELAYED_TENSOR_SCALING); - is_mxfp8 = (scaling_mode == NVTE_MXFP8_1D_SCALING); - is_delayed_scaling_fwd = is_tensor_scaling && is_o_in_fp8; - is_delayed_scaling_bwd = is_tensor_scaling && is_dqkv_in_fp8; - is_current_scaling_fwd = is_tensor_scaling && is_o_in_f16; - is_current_scaling_bwd = is_tensor_scaling && is_dqkv_in_f16; - is_mxfp8_fwd = is_mxfp8 && is_o_in_f16; - is_mxfp8_bwd = is_mxfp8 && is_dqkv_in_f16; - - // Use cu_seqlens vs actual_seqlens for THD or padding masks - const size_t cudnn_runtime_version = cudnnGetVersion(); - const bool is_fp8_dtype = (qkv_dtype == kNVTEFloat8E4M3 || qkv_dtype == kNVTEFloat8E5M2); - const size_t min_frontend_version = is_fp8_dtype ? 12600 : 12500; - const size_t min_cudnn_version = is_fp8_dtype ? 92500 : 92400; - uses_cu_seqlens_directly = - CUDNN_FRONTEND_VERSION >= min_frontend_version && - (CUDNN_VERSION >= min_cudnn_version && cudnn_runtime_version >= min_cudnn_version) && - !is_dropout; - - // Bucket the batch size and token counts for THD - bucketed_batch_size = - (is_ragged_q || is_ragged_kv) ? fused_attn::get_max_batch_size(batch_size) : 0; - bucketed_num_tokens_q = is_ragged_q ? fused_attn::get_max_tokens(num_tokens_q) : 0; - bucketed_num_tokens_kv = is_ragged_kv ? fused_attn::get_max_tokens(num_tokens_kv) : 0; - - // Use ragged (TH1) or dense (BHS1) graphs and stats - const int sm_arch = cuda::sm_arch(cuda::current_device()); - uses_ragged_graph = cudnn_runtime_version >= 90600 && sm_arch >= 90 && sm_arch != 120; - uses_ragged_stats = is_ragged_q && uses_ragged_graph; - const bool buckets_the_batch = (is_ragged_q || is_ragged_kv) && uses_ragged_graph; - graph_batch_size_fwd = - (buckets_the_batch && !uses_cu_seqlens_directly) ? bucketed_batch_size : batch_size; - graph_batch_size_bwd = buckets_the_batch ? bucketed_batch_size : batch_size; - graph_max_seqlen_q = (is_ragged_q && uses_ragged_graph) ? bucketed_num_tokens_q : max_seqlen_q; - graph_max_seqlen_kv = - (is_ragged_kv && uses_ragged_graph) ? bucketed_num_tokens_kv : max_seqlen_kv; - - // Set up ragged offset widths and multipliers - needs_64bit_ragged_offset = - (is_ragged_q || is_ragged_kv) && - fused_attn::get_ragged_offset_dtype( - layout_group, static_cast(num_attn_heads), static_cast(num_gqa_groups), - static_cast(max_seqlen_q), static_cast(max_seqlen_kv), - static_cast(head_dim_qk), static_cast(head_dim_v)) == DType::kInt64; - const DType wide_ragged_offsets = cudnn_runtime_version >= 90500 ? DType::kInt64 : DType::kInt32; - ragged_offset_type_fwd = uses_cu_seqlens_directly ? DType::kInt32 : wide_ragged_offsets; - ragged_offset_type_bwd = wide_ragged_offsets; - ragged_offset_mults = RaggedOffsetMultipliers( - layout_group, static_cast(num_attn_heads), static_cast(num_gqa_groups), - static_cast(head_dim_qk), static_cast(head_dim_v)); - - // Mark as derived - is_derived = true; -} - -FusedAttnConfig FusedAttnConfig::make_cache_key(Pass pass) const { - check_derived(); - FusedAttnConfig cache_cfg = *this; - cache_cfg.device_id = cuda::current_device(); - - // Normalize sequence lengths the graph is built at - cache_cfg.max_seqlen_q = cache_cfg.graph_max_seqlen_q; - cache_cfg.max_seqlen_kv = cache_cfg.graph_max_seqlen_kv; - - // Normalize the batch size the graph is built at, and drop the token counts. - cache_cfg.num_tokens_q = 0; - cache_cfg.num_tokens_kv = 0; - if ((cache_cfg.is_ragged_q || cache_cfg.is_ragged_kv) && cache_cfg.uses_ragged_graph) { - cache_cfg.batch_size = - pass == Pass::Fwd ? cache_cfg.graph_batch_size_fwd : cache_cfg.graph_batch_size_bwd; - } - - // attn_scale is a pass-by-value graph input and different scales can share the same cached graph - cache_cfg.attn_scale = 1.0f; - - // cuda_graph never reaches a graph builder. Its one use is the cuDNN <= 9.15 rejection in - // nvte_get_fused_attn_backend_v2(). - cache_cfg.cuda_graph = false; - - // Normalize the fields its graph actually consumes - if (pass == Pass::Fwd) { - cache_cfg.do_dtype = kNVTEBFloat16; - cache_cfg.dqkv_dtype = kNVTEBFloat16; - cache_cfg.do_format = NVTE_QKV_Format_NOT_SET; - cache_cfg.dqkv_layout = NVTE_QKV_Layout_NOT_SET; - cache_cfg.do_scale_inv_format = NVTE_QKV_Format_NOT_SET; - cache_cfg.deterministic = false; - } else { - cache_cfg.return_max_logit = false; - } - - return cache_cfg; -} - -std::string FusedAttnConfig::to_string() const { - char buf[1024]; - std::snprintf( - buf, sizeof(buf), - "train=%d det=%d cg=%d maxlogit=%d mask=%" PRId64 " bias=%" PRId64 " wl=%" PRId64 - " wr=%" PRId64 " brd=%d softmax=%" PRId64 " scale_mode=%" PRId64 - " dropout=%g attn_scale=%g qkv_dt=%" PRId64 " o_dt=%" PRId64 " do_dt=%" PRId64 - " dqkv_dt=%" PRId64 " qkv_lay=%" PRId64 " o_fmt=%" PRId64 " do_fmt=%" PRId64 - " dqkv_lay=%" PRId64 " qkv_sif=%" PRId64 " do_sif=%" PRId64 " b=%" PRId64 " h=%" PRId64 - " hg=%" PRId64 " dqk=%" PRId64 " dv=%" PRId64 " sq=%" PRId64 " skv=%" PRId64 " tq=%" PRId64 - " tkv=%" PRId64 " bb=%" PRId64 " btq=%" PRId64 " btkv=%" PRId64 " npk=%" PRId64 - " npv=%" PRId64 " psk=%" PRId64 " psv=%" PRId64 " mppk=%" PRId64 " mppv=%" PRId64 - " bias_b=%" PRId64 " bias_h=%" PRId64 " bias_sq=%" PRId64 " bias_skv=%" PRId64, - static_cast(is_training), static_cast(deterministic), static_cast(cuda_graph), - static_cast(return_max_logit), static_cast(attn_mask_type), - static_cast(bias_type), static_cast(window_size_left), - static_cast(window_size_right), static_cast(bottom_right_diagonal), - static_cast(softmax_type), static_cast(scaling_mode), - static_cast(dropout), static_cast(attn_scale), - static_cast(qkv_dtype), static_cast(o_dtype), - static_cast(do_dtype), static_cast(dqkv_dtype), - static_cast(qkv_layout), static_cast(o_format), - static_cast(do_format), static_cast(dqkv_layout), - static_cast(qkv_scale_inv_format), static_cast(do_scale_inv_format), - static_cast(batch_size), static_cast(num_attn_heads), - static_cast(num_gqa_groups), static_cast(head_dim_qk), - static_cast(head_dim_v), static_cast(max_seqlen_q), - static_cast(max_seqlen_kv), static_cast(num_tokens_q), - static_cast(num_tokens_kv), static_cast(bucketed_batch_size), - static_cast(bucketed_num_tokens_q), static_cast(bucketed_num_tokens_kv), - static_cast(num_pages_k), static_cast(num_pages_v), - static_cast(page_size_k), static_cast(page_size_v), - static_cast(max_pages_per_seq_k), static_cast(max_pages_per_seq_v), - static_cast(bias_batch_size), static_cast(bias_num_heads), - static_cast(bias_seqlen_q), static_cast(bias_seqlen_kv)); - return std::string(buf); -} - -FusedAttnConfig FusedAttnFwdParams::make_config() const { - const FusedAttnFwdParams ¶ms = *this; - FusedAttnConfig cfg{}; - // Forward execution: only the forward graph is run, so do not pay for a backward support - // check whose graph this call will never execute. - cfg.check_for_forward_support = true; - cfg.check_for_backward_support = false; - cfg.is_training = params.is_training; - cfg.deterministic = false; - cfg.cuda_graph = params.cuda_graph; - cfg.return_max_logit = params.return_max_logit; - cfg.attn_mask_type = params.attn_mask_type; - cfg.bias_type = params.bias_type; - cfg.window_size_left = params.window_size_left; - cfg.window_size_right = params.window_size_right; - cfg.bottom_right_diagonal = params.bottom_right_diagonal; - cfg.softmax_type = params.softmax_type; - cfg.dropout = params.dropout; - cfg.attn_scale = params.attn_scale; - cfg.qkv_layout = params.qkv_layout; - cfg.o_format = params.o_format; - cfg.qkv_scale_inv_format = params.qkv_scale_inv_format; - cfg.max_seqlen_q = params.max_seqlen_q; - cfg.max_seqlen_kv = params.max_seqlen_kv; - - const Tensor *input_cu_seqlens_q = convertNVTETensorCheck(params.cu_seqlens_q); - const Tensor *input_cu_seqlens_kv = convertNVTETensorCheck(params.cu_seqlens_kv); - const Tensor *input_page_table_k = convertNVTETensorCheck(params.page_table_k); - const Tensor *input_page_table_v = convertNVTETensorCheck(params.page_table_v); - const Tensor *input_Q = convertNVTETensorCheck(params.Q); - const Tensor *input_K = convertNVTETensorCheck(params.K); - const Tensor *input_V = convertNVTETensorCheck(params.V); - const Tensor *input_Bias = convertNVTETensorCheck(params.Bias); - const Tensor *output_O = convertNVTETensorCheck(params.O); - - const NVTE_QKV_Format q_format = nvte_get_q_format(params.qkv_layout); - const NVTE_QKV_Format kv_format = nvte_get_kv_format(params.qkv_layout); - auto *q_dims = input_Q->data.shape.data(); - auto *k_dims = input_K->data.shape.data(); - auto *v_dims = input_V->scaling_mode != NVTE_MXFP8_1D_SCALING - ? input_V->data.shape.data() - : input_V->columnwise_data.shape.data(); - AttentionShape q_shape(q_format, q_dims); - AttentionShape k_shape(kv_format, k_dims); - AttentionShape v_shape(kv_format, v_dims); - size_t b = q_shape.b(), h_q = q_shape.h(), d_qk = q_shape.d(), t_q = q_shape.t(); - size_t h_kv = k_shape.h(), t_kv = k_shape.t(), d_v = v_shape.d(); - if (q_format == NVTE_QKV_Format::NVTE_THD) { - b = input_cu_seqlens_q->data.shape[0] - 1; - } else if (kv_format == NVTE_QKV_Format::NVTE_THD) { - b = input_cu_seqlens_kv->data.shape[0] - 1; - } - - int64_t num_pages_k = 0, num_pages_v = 0, page_size_k = 0, page_size_v = 0; - int64_t max_pages_per_seq_k = 0, max_pages_per_seq_v = 0; - if (input_page_table_k->data.dptr != nullptr) { - max_pages_per_seq_k = input_page_table_k->data.shape[1]; - } - if (input_page_table_v->data.dptr != nullptr) { - max_pages_per_seq_v = input_page_table_v->data.shape[1]; - } - const NVTE_QKV_Layout_Group layout_group = nvte_get_qkv_layout_group(params.qkv_layout); - if (layout_group == NVTE_QKV_Layout_Group::NVTE_Paged_KV_HD_HD_HD) { - const NVTE_QKV_Format paged_kv_format = nvte_get_kv_format(params.qkv_layout); - if (paged_kv_format == NVTE_QKV_Format::NVTE_BSHD) { - num_pages_k = input_K->data.shape[0]; - page_size_k = input_K->data.shape[1]; - num_pages_v = input_V->data.shape[0]; - page_size_v = input_V->data.shape[1]; - } else if (paged_kv_format == NVTE_QKV_Format::NVTE_SBHD) { - num_pages_k = input_K->data.shape[1]; - page_size_k = input_K->data.shape[0]; - num_pages_v = input_V->data.shape[1]; - page_size_v = input_V->data.shape[0]; - } - } - - const NVTEDType Q_type = static_cast(input_Q->data.dtype); - const NVTEDType KV_type = static_cast(input_K->data.dtype); - NVTE_CHECK(Q_type == KV_type, "Q and KV must have the same data type."); - - cfg.scaling_mode = input_Q->scaling_mode; - cfg.qkv_dtype = Q_type; - cfg.o_dtype = static_cast(output_O->data.dtype); - cfg.batch_size = b; - cfg.num_attn_heads = h_q; - cfg.num_gqa_groups = h_kv; - cfg.head_dim_qk = d_qk; - cfg.head_dim_v = d_v; - cfg.num_pages_k = static_cast(num_pages_k); - cfg.num_pages_v = static_cast(num_pages_v); - cfg.page_size_k = static_cast(page_size_k); - cfg.page_size_v = static_cast(page_size_v); - cfg.max_pages_per_seq_k = static_cast(max_pages_per_seq_k); - cfg.max_pages_per_seq_v = static_cast(max_pages_per_seq_v); - cfg.num_tokens_q = t_q; - cfg.num_tokens_kv = t_kv; - - if ((params.bias_type != NVTE_NO_BIAS) && (params.bias_type != NVTE_ALIBI) && - input_Bias->data.shape.size() >= 4) { - cfg.bias_batch_size = input_Bias->data.shape[0]; - cfg.bias_num_heads = input_Bias->data.shape[1]; - cfg.bias_seqlen_q = input_Bias->data.shape[2]; - cfg.bias_seqlen_kv = input_Bias->data.shape[3]; - } - return cfg; -} - -FusedAttnConfig FusedAttnBwdParams::make_config() const { - const FusedAttnBwdParams ¶ms = *this; - FusedAttnConfig cfg{}; - // Backward execution: only the backward graph is run, so do not pay for a forward support - // check whose graph this call will never execute. - cfg.check_for_forward_support = false; - cfg.check_for_backward_support = true; - cfg.is_training = true; - cfg.deterministic = params.deterministic; - cfg.cuda_graph = params.cuda_graph; - cfg.return_max_logit = false; - cfg.attn_mask_type = params.attn_mask_type; - cfg.bias_type = params.bias_type; - cfg.window_size_left = params.window_size_left; - cfg.window_size_right = params.window_size_right; - cfg.bottom_right_diagonal = params.bottom_right_diagonal; - cfg.softmax_type = params.softmax_type; - cfg.dropout = params.dropout; - cfg.attn_scale = params.attn_scale; - cfg.qkv_layout = params.qkv_layout; - cfg.o_format = params.o_format; - cfg.do_format = params.do_format; - cfg.dqkv_layout = params.dqkv_layout; - cfg.qkv_scale_inv_format = params.qkv_scale_inv_format; - cfg.do_scale_inv_format = params.do_scale_inv_format; - cfg.max_seqlen_q = params.max_seqlen_q; - cfg.max_seqlen_kv = params.max_seqlen_kv; - - const Tensor *input_cu_seqlens_q = convertNVTETensorCheck(params.cu_seqlens_q); - const Tensor *input_cu_seqlens_kv = convertNVTETensorCheck(params.cu_seqlens_kv); - const Tensor *input_Q = convertNVTETensorCheck(params.Q); - const Tensor *input_K = convertNVTETensorCheck(params.K); - const Tensor *input_V = convertNVTETensorCheck(params.V); - const Tensor *input_O = convertNVTETensorCheck(params.O); - const Tensor *input_dO = convertNVTETensorCheck(params.dO); - const Tensor *output_dQ = convertNVTETensorCheck(params.dQ); - const Tensor *output_dBias = convertNVTETensorCheck(params.dBias); - - const NVTE_QKV_Format q_format = nvte_get_q_format(params.qkv_layout); - const NVTE_QKV_Format kv_format = nvte_get_kv_format(params.qkv_layout); - auto *q_dims = input_Q->data.shape.data(); - auto *k_dims = input_K->data.shape.data(); - auto *v_dims = input_V->data.shape.data(); - AttentionShape q_shape(q_format, q_dims); - AttentionShape k_shape(kv_format, k_dims); - AttentionShape v_shape(kv_format, v_dims); - size_t b = q_shape.b(), h_q = q_shape.h(), d_qk = q_shape.d(), t_q = q_shape.t(); - size_t h_kv = k_shape.h(), t_kv = k_shape.t(), d_v = v_shape.d(); - if (q_format == NVTE_QKV_Format::NVTE_THD) { - b = input_cu_seqlens_q->data.shape[0] - 1; - } else if (kv_format == NVTE_QKV_Format::NVTE_THD) { - b = input_cu_seqlens_kv->data.shape[0] - 1; - } - - const NVTEDType Q_type = static_cast(input_Q->data.dtype); - const NVTEDType KV_type = static_cast(input_K->data.dtype); - NVTE_CHECK(Q_type == KV_type, "Q and KV must have the same data type."); - - cfg.scaling_mode = input_Q->scaling_mode; - cfg.qkv_dtype = Q_type; - cfg.o_dtype = static_cast(input_O->data.dtype); - cfg.do_dtype = static_cast(input_dO->data.dtype); - cfg.dqkv_dtype = static_cast(output_dQ->data.dtype); - cfg.batch_size = b; - cfg.num_attn_heads = h_q; - cfg.num_gqa_groups = h_kv; - cfg.head_dim_qk = d_qk; - cfg.head_dim_v = d_v; - cfg.num_tokens_q = t_q; - cfg.num_tokens_kv = t_kv; - - if ((params.bias_type != NVTE_NO_BIAS) && (params.bias_type != NVTE_ALIBI) && - output_dBias->data.shape.size() >= 4) { - cfg.bias_batch_size = output_dBias->data.shape[0]; - cfg.bias_num_heads = output_dBias->data.shape[1]; - cfg.bias_seqlen_q = output_dBias->data.shape[2]; - cfg.bias_seqlen_kv = output_dBias->data.shape[3]; - } - return cfg; -} - -} // namespace fused_attn -} // namespace transformer_engine - -NVTEFusedAttnConfig nvte_create_fused_attn_config() { - return new transformer_engine::fused_attn::FusedAttnConfig{}; -} - -void nvte_destroy_fused_attn_config(NVTEFusedAttnConfig config) { - delete transformer_engine::fused_attn::get_fused_attn_config_mutable(config); -} - -void nvte_get_fused_attn_config_attribute(NVTEFusedAttnConfig config, - NVTEFusedAttnConfigAttribute attr, void *buf, - size_t size_in_bytes, size_t *size_written) { - using namespace transformer_engine; - using namespace transformer_engine::fused_attn; - - NVTE_CHECK(attr < kNVTEFusedAttnConfigNumAttributes, "Invalid NVTEFusedAttnConfigAttribute (got ", - static_cast(attr), ")"); - const auto &attr_size = FusedAttnConfig::attr_sizes[attr]; - if (size_written != nullptr) { - *size_written = attr_size; - } - if (buf == nullptr) { - return; - } - NVTE_CHECK(size_in_bytes >= attr_size, - "Buffer is too small for fused attention config attribute (attribute ", - static_cast(attr), " needs ", attr_size, " bytes, but buffer has ", size_in_bytes, - " bytes)"); - - const auto &cfg = *get_fused_attn_config(config); - switch (attr) { - case kNVTEFusedAttnConfigIsTraining: - bool_to_uint8(cfg.is_training, buf); - break; - case kNVTEFusedAttnConfigDeterministic: - bool_to_uint8(cfg.deterministic, buf); - break; - case kNVTEFusedAttnConfigCudaGraph: - bool_to_uint8(cfg.cuda_graph, buf); - break; - case kNVTEFusedAttnConfigReturnMaxLogit: - bool_to_uint8(cfg.return_max_logit, buf); - break; - case kNVTEFusedAttnConfigAttnMaskType: - std::memcpy(buf, &cfg.attn_mask_type, attr_size); - break; - case kNVTEFusedAttnConfigBiasType: - std::memcpy(buf, &cfg.bias_type, attr_size); - break; - case kNVTEFusedAttnConfigWindowSizeLeft: - std::memcpy(buf, &cfg.window_size_left, attr_size); - break; - case kNVTEFusedAttnConfigWindowSizeRight: - std::memcpy(buf, &cfg.window_size_right, attr_size); - break; - case kNVTEFusedAttnConfigBottomRightDiagonal: - bool_to_uint8(cfg.bottom_right_diagonal, buf); - break; - case kNVTEFusedAttnConfigSoftmaxType: - std::memcpy(buf, &cfg.softmax_type, attr_size); - break; - case kNVTEFusedAttnConfigScalingMode: - std::memcpy(buf, &cfg.scaling_mode, attr_size); - break; - case kNVTEFusedAttnConfigDropout: - std::memcpy(buf, &cfg.dropout, attr_size); - break; - case kNVTEFusedAttnConfigAttnScale: - std::memcpy(buf, &cfg.attn_scale, attr_size); - break; - case kNVTEFusedAttnConfigQKVDtype: - std::memcpy(buf, &cfg.qkv_dtype, attr_size); - break; - case kNVTEFusedAttnConfigODtype: - std::memcpy(buf, &cfg.o_dtype, attr_size); - break; - case kNVTEFusedAttnConfigDODtype: - std::memcpy(buf, &cfg.do_dtype, attr_size); - break; - case kNVTEFusedAttnConfigDQKVDtype: - std::memcpy(buf, &cfg.dqkv_dtype, attr_size); - break; - case kNVTEFusedAttnConfigQKVLayout: - std::memcpy(buf, &cfg.qkv_layout, attr_size); - break; - case kNVTEFusedAttnConfigOFormat: - std::memcpy(buf, &cfg.o_format, attr_size); - break; - case kNVTEFusedAttnConfigDOFormat: - std::memcpy(buf, &cfg.do_format, attr_size); - break; - case kNVTEFusedAttnConfigDQKVLayout: - std::memcpy(buf, &cfg.dqkv_layout, attr_size); - break; - case kNVTEFusedAttnConfigQKVScaleInvFormat: - std::memcpy(buf, &cfg.qkv_scale_inv_format, attr_size); - break; - case kNVTEFusedAttnConfigDOScaleInvFormat: - std::memcpy(buf, &cfg.do_scale_inv_format, attr_size); - break; - case kNVTEFusedAttnConfigBatchSize: - std::memcpy(buf, &cfg.batch_size, attr_size); - break; - case kNVTEFusedAttnConfigNumAttnHeads: - std::memcpy(buf, &cfg.num_attn_heads, attr_size); - break; - case kNVTEFusedAttnConfigNumGQAGroups: - std::memcpy(buf, &cfg.num_gqa_groups, attr_size); - break; - case kNVTEFusedAttnConfigHeadDimQK: - std::memcpy(buf, &cfg.head_dim_qk, attr_size); - break; - case kNVTEFusedAttnConfigHeadDimV: - std::memcpy(buf, &cfg.head_dim_v, attr_size); - break; - case kNVTEFusedAttnConfigMaxSeqlenQ: - std::memcpy(buf, &cfg.max_seqlen_q, attr_size); - break; - case kNVTEFusedAttnConfigMaxSeqlenKV: - std::memcpy(buf, &cfg.max_seqlen_kv, attr_size); - break; - case kNVTEFusedAttnConfigNumTokensQ: - std::memcpy(buf, &cfg.num_tokens_q, attr_size); - break; - case kNVTEFusedAttnConfigNumTokensKV: - std::memcpy(buf, &cfg.num_tokens_kv, attr_size); - break; - case kNVTEFusedAttnConfigNumPagesK: - std::memcpy(buf, &cfg.num_pages_k, attr_size); - break; - case kNVTEFusedAttnConfigNumPagesV: - std::memcpy(buf, &cfg.num_pages_v, attr_size); - break; - case kNVTEFusedAttnConfigPageSizeK: - std::memcpy(buf, &cfg.page_size_k, attr_size); - break; - case kNVTEFusedAttnConfigPageSizeV: - std::memcpy(buf, &cfg.page_size_v, attr_size); - break; - case kNVTEFusedAttnConfigMaxPagesPerSeqK: - std::memcpy(buf, &cfg.max_pages_per_seq_k, attr_size); - break; - case kNVTEFusedAttnConfigMaxPagesPerSeqV: - std::memcpy(buf, &cfg.max_pages_per_seq_v, attr_size); - break; - case kNVTEFusedAttnConfigBiasBatchSize: - std::memcpy(buf, &cfg.bias_batch_size, attr_size); - break; - case kNVTEFusedAttnConfigBiasNumHeads: - std::memcpy(buf, &cfg.bias_num_heads, attr_size); - break; - case kNVTEFusedAttnConfigBiasSeqlenQ: - std::memcpy(buf, &cfg.bias_seqlen_q, attr_size); - break; - case kNVTEFusedAttnConfigBiasSeqlenKV: - std::memcpy(buf, &cfg.bias_seqlen_kv, attr_size); - break; - default: - NVTE_ERROR("Unsupported NVTEFusedAttnConfigAttribute (got ", static_cast(attr), ")"); - } -} - -void nvte_set_fused_attn_config_attribute(NVTEFusedAttnConfig config, - NVTEFusedAttnConfigAttribute attr, const void *buf, - size_t size_in_bytes) { - using namespace transformer_engine; - using namespace transformer_engine::fused_attn; - - NVTE_CHECK(attr < kNVTEFusedAttnConfigNumAttributes, "Invalid NVTEFusedAttnConfigAttribute (got ", - static_cast(attr), ")"); - const auto &attr_size = FusedAttnConfig::attr_sizes[attr]; - NVTE_CHECK(size_in_bytes >= attr_size, - "Buffer is too small for fused attention config attribute (attribute ", - static_cast(attr), " needs ", attr_size, " bytes, but buffer has ", size_in_bytes, - " bytes)"); - NVTE_CHECK(buf != nullptr, "Invalid buffer (got NULL)"); - - auto &cfg = *get_fused_attn_config_mutable(config); - cfg.is_derived = false; - switch (attr) { - case kNVTEFusedAttnConfigIsTraining: - uint8_to_bool(buf, cfg.is_training); - break; - case kNVTEFusedAttnConfigDeterministic: - uint8_to_bool(buf, cfg.deterministic); - break; - case kNVTEFusedAttnConfigCudaGraph: - uint8_to_bool(buf, cfg.cuda_graph); - break; - case kNVTEFusedAttnConfigReturnMaxLogit: - uint8_to_bool(buf, cfg.return_max_logit); - break; - case kNVTEFusedAttnConfigAttnMaskType: - std::memcpy(&cfg.attn_mask_type, buf, attr_size); - break; - case kNVTEFusedAttnConfigBiasType: - std::memcpy(&cfg.bias_type, buf, attr_size); - break; - case kNVTEFusedAttnConfigWindowSizeLeft: - std::memcpy(&cfg.window_size_left, buf, attr_size); - break; - case kNVTEFusedAttnConfigWindowSizeRight: - std::memcpy(&cfg.window_size_right, buf, attr_size); - break; - case kNVTEFusedAttnConfigBottomRightDiagonal: - uint8_to_bool(buf, cfg.bottom_right_diagonal); - break; - case kNVTEFusedAttnConfigSoftmaxType: - std::memcpy(&cfg.softmax_type, buf, attr_size); - break; - case kNVTEFusedAttnConfigScalingMode: - std::memcpy(&cfg.scaling_mode, buf, attr_size); - break; - case kNVTEFusedAttnConfigDropout: - std::memcpy(&cfg.dropout, buf, attr_size); - break; - case kNVTEFusedAttnConfigAttnScale: - std::memcpy(&cfg.attn_scale, buf, attr_size); - break; - case kNVTEFusedAttnConfigQKVDtype: - std::memcpy(&cfg.qkv_dtype, buf, attr_size); - break; - case kNVTEFusedAttnConfigODtype: - std::memcpy(&cfg.o_dtype, buf, attr_size); - break; - case kNVTEFusedAttnConfigDODtype: - std::memcpy(&cfg.do_dtype, buf, attr_size); - break; - case kNVTEFusedAttnConfigDQKVDtype: - std::memcpy(&cfg.dqkv_dtype, buf, attr_size); - break; - case kNVTEFusedAttnConfigQKVLayout: - std::memcpy(&cfg.qkv_layout, buf, attr_size); - break; - case kNVTEFusedAttnConfigOFormat: - std::memcpy(&cfg.o_format, buf, attr_size); - break; - case kNVTEFusedAttnConfigDOFormat: - std::memcpy(&cfg.do_format, buf, attr_size); - break; - case kNVTEFusedAttnConfigDQKVLayout: - std::memcpy(&cfg.dqkv_layout, buf, attr_size); - break; - case kNVTEFusedAttnConfigQKVScaleInvFormat: - std::memcpy(&cfg.qkv_scale_inv_format, buf, attr_size); - break; - case kNVTEFusedAttnConfigDOScaleInvFormat: - std::memcpy(&cfg.do_scale_inv_format, buf, attr_size); - break; - case kNVTEFusedAttnConfigBatchSize: - std::memcpy(&cfg.batch_size, buf, attr_size); - break; - case kNVTEFusedAttnConfigNumAttnHeads: - std::memcpy(&cfg.num_attn_heads, buf, attr_size); - break; - case kNVTEFusedAttnConfigNumGQAGroups: - std::memcpy(&cfg.num_gqa_groups, buf, attr_size); - break; - case kNVTEFusedAttnConfigHeadDimQK: - std::memcpy(&cfg.head_dim_qk, buf, attr_size); - break; - case kNVTEFusedAttnConfigHeadDimV: - std::memcpy(&cfg.head_dim_v, buf, attr_size); - break; - case kNVTEFusedAttnConfigMaxSeqlenQ: - std::memcpy(&cfg.max_seqlen_q, buf, attr_size); - break; - case kNVTEFusedAttnConfigMaxSeqlenKV: - std::memcpy(&cfg.max_seqlen_kv, buf, attr_size); - break; - case kNVTEFusedAttnConfigNumTokensQ: - std::memcpy(&cfg.num_tokens_q, buf, attr_size); - break; - case kNVTEFusedAttnConfigNumTokensKV: - std::memcpy(&cfg.num_tokens_kv, buf, attr_size); - break; - case kNVTEFusedAttnConfigNumPagesK: - std::memcpy(&cfg.num_pages_k, buf, attr_size); - break; - case kNVTEFusedAttnConfigNumPagesV: - std::memcpy(&cfg.num_pages_v, buf, attr_size); - break; - case kNVTEFusedAttnConfigPageSizeK: - std::memcpy(&cfg.page_size_k, buf, attr_size); - break; - case kNVTEFusedAttnConfigPageSizeV: - std::memcpy(&cfg.page_size_v, buf, attr_size); - break; - case kNVTEFusedAttnConfigMaxPagesPerSeqK: - std::memcpy(&cfg.max_pages_per_seq_k, buf, attr_size); - break; - case kNVTEFusedAttnConfigMaxPagesPerSeqV: - std::memcpy(&cfg.max_pages_per_seq_v, buf, attr_size); - break; - case kNVTEFusedAttnConfigBiasBatchSize: - std::memcpy(&cfg.bias_batch_size, buf, attr_size); - break; - case kNVTEFusedAttnConfigBiasNumHeads: - std::memcpy(&cfg.bias_num_heads, buf, attr_size); - break; - case kNVTEFusedAttnConfigBiasSeqlenQ: - std::memcpy(&cfg.bias_seqlen_q, buf, attr_size); - break; - case kNVTEFusedAttnConfigBiasSeqlenKV: - std::memcpy(&cfg.bias_seqlen_kv, buf, attr_size); - break; - default: - NVTE_ERROR("Unsupported NVTEFusedAttnConfigAttribute (got ", static_cast(attr), ")"); - } -} - -NVTEFusedAttnFwdParams nvte_create_fused_attn_fwd_params() { - return new transformer_engine::fused_attn::FusedAttnFwdParams{}; -} - -void nvte_destroy_fused_attn_fwd_params(NVTEFusedAttnFwdParams params) { - delete transformer_engine::fused_attn::get_fused_attn_fwd_params_mutable(params); -} - -void nvte_get_fused_attn_fwd_params_attribute(NVTEFusedAttnFwdParams params, - NVTEFusedAttnFwdParamsAttribute attr, void *buf, - size_t size_in_bytes, size_t *size_written) { - using namespace transformer_engine; - using namespace transformer_engine::fused_attn; - NVTE_CHECK(attr < kNVTEFusedAttnFwdParamsNumAttributes, - "Invalid NVTEFusedAttnFwdParamsAttribute (got ", static_cast(attr), ")"); - const auto &attr_size = FusedAttnFwdParams::attr_sizes[attr]; - if (size_written != nullptr) { - *size_written = attr_size; - } - if (buf == nullptr) { - return; - } - NVTE_CHECK(size_in_bytes >= attr_size, "Buffer is too small for attribute (need ", attr_size, - ", got ", size_in_bytes, ")"); - const auto &p = *get_fused_attn_fwd_params(params); - switch (attr) { - case kNVTEFusedAttnFwdParamsQ: - std::memcpy(buf, &p.Q, attr_size); - break; - case kNVTEFusedAttnFwdParamsK: - std::memcpy(buf, &p.K, attr_size); - break; - case kNVTEFusedAttnFwdParamsV: - std::memcpy(buf, &p.V, attr_size); - break; - case kNVTEFusedAttnFwdParamsBias: - std::memcpy(buf, &p.Bias, attr_size); - break; - case kNVTEFusedAttnFwdParamsSoftmaxOffset: - std::memcpy(buf, &p.SoftmaxOffset, attr_size); - break; - case kNVTEFusedAttnFwdParamsS: - std::memcpy(buf, &p.S, attr_size); - break; - case kNVTEFusedAttnFwdParamsO: - std::memcpy(buf, &p.O, attr_size); - break; - case kNVTEFusedAttnFwdParamsAuxCtxTensors: - std::memcpy(buf, &p.Aux_CTX_Tensors, attr_size); - break; - case kNVTEFusedAttnFwdParamsCuSeqlensQ: - std::memcpy(buf, &p.cu_seqlens_q, attr_size); - break; - case kNVTEFusedAttnFwdParamsCuSeqlensKV: - std::memcpy(buf, &p.cu_seqlens_kv, attr_size); - break; - case kNVTEFusedAttnFwdParamsCuSeqlensQPadded: - std::memcpy(buf, &p.cu_seqlens_q_padded, attr_size); - break; - case kNVTEFusedAttnFwdParamsCuSeqlensKVPadded: - std::memcpy(buf, &p.cu_seqlens_kv_padded, attr_size); - break; - case kNVTEFusedAttnFwdParamsPageTableK: - std::memcpy(buf, &p.page_table_k, attr_size); - break; - case kNVTEFusedAttnFwdParamsPageTableV: - std::memcpy(buf, &p.page_table_v, attr_size); - break; - case kNVTEFusedAttnFwdParamsRngState: - std::memcpy(buf, &p.rng_state, attr_size); - break; - case kNVTEFusedAttnFwdParamsIsTraining: - bool_to_uint8(p.is_training, buf); - break; - case kNVTEFusedAttnFwdParamsCudaGraph: - bool_to_uint8(p.cuda_graph, buf); - break; - case kNVTEFusedAttnFwdParamsReturnMaxLogit: - bool_to_uint8(p.return_max_logit, buf); - break; - case kNVTEFusedAttnFwdParamsAttnMaskType: - std::memcpy(buf, &p.attn_mask_type, attr_size); - break; - case kNVTEFusedAttnFwdParamsBiasType: - std::memcpy(buf, &p.bias_type, attr_size); - break; - case kNVTEFusedAttnFwdParamsWindowSizeLeft: - std::memcpy(buf, &p.window_size_left, attr_size); - break; - case kNVTEFusedAttnFwdParamsWindowSizeRight: - std::memcpy(buf, &p.window_size_right, attr_size); - break; - case kNVTEFusedAttnFwdParamsBottomRightDiagonal: - bool_to_uint8(p.bottom_right_diagonal, buf); - break; - case kNVTEFusedAttnFwdParamsSoftmaxType: - std::memcpy(buf, &p.softmax_type, attr_size); - break; - case kNVTEFusedAttnFwdParamsDropout: - std::memcpy(buf, &p.dropout, attr_size); - break; - case kNVTEFusedAttnFwdParamsAttnScale: - std::memcpy(buf, &p.attn_scale, attr_size); - break; - case kNVTEFusedAttnFwdParamsQKVLayout: - std::memcpy(buf, &p.qkv_layout, attr_size); - break; - case kNVTEFusedAttnFwdParamsOFormat: - std::memcpy(buf, &p.o_format, attr_size); - break; - case kNVTEFusedAttnFwdParamsQKVScaleInvFormat: - std::memcpy(buf, &p.qkv_scale_inv_format, attr_size); - break; - case kNVTEFusedAttnFwdParamsMaxSeqlenQ: - std::memcpy(buf, &p.max_seqlen_q, attr_size); - break; - case kNVTEFusedAttnFwdParamsMaxSeqlenKV: - std::memcpy(buf, &p.max_seqlen_kv, attr_size); - break; - case kNVTEFusedAttnFwdParamsWorkspace: - std::memcpy(buf, &p.workspace, attr_size); - break; - case kNVTEFusedAttnFwdParamsStream: - std::memcpy(buf, &p.stream, attr_size); - break; - default: - NVTE_ERROR("Unsupported NVTEFusedAttnFwdParamsAttribute (got ", static_cast(attr), ")"); - } -} - -void nvte_set_fused_attn_fwd_params_attribute(NVTEFusedAttnFwdParams params, - NVTEFusedAttnFwdParamsAttribute attr, const void *buf, - size_t size_in_bytes) { - using namespace transformer_engine; - using namespace transformer_engine::fused_attn; - NVTE_CHECK(attr < kNVTEFusedAttnFwdParamsNumAttributes, - "Invalid NVTEFusedAttnFwdParamsAttribute (got ", static_cast(attr), ")"); - const auto &attr_size = FusedAttnFwdParams::attr_sizes[attr]; - NVTE_CHECK(buf != nullptr, "Input buffer must not be NULL."); - NVTE_CHECK(size_in_bytes >= attr_size, "Buffer is too small for attribute (need ", attr_size, - ", got ", size_in_bytes, ")"); - auto &p = *get_fused_attn_fwd_params_mutable(params); - switch (attr) { - case kNVTEFusedAttnFwdParamsQ: - std::memcpy(&p.Q, buf, attr_size); - break; - case kNVTEFusedAttnFwdParamsK: - std::memcpy(&p.K, buf, attr_size); - break; - case kNVTEFusedAttnFwdParamsV: - std::memcpy(&p.V, buf, attr_size); - break; - case kNVTEFusedAttnFwdParamsBias: - std::memcpy(&p.Bias, buf, attr_size); - break; - case kNVTEFusedAttnFwdParamsSoftmaxOffset: - std::memcpy(&p.SoftmaxOffset, buf, attr_size); - break; - case kNVTEFusedAttnFwdParamsS: - std::memcpy(&p.S, buf, attr_size); - break; - case kNVTEFusedAttnFwdParamsO: - std::memcpy(&p.O, buf, attr_size); - break; - case kNVTEFusedAttnFwdParamsAuxCtxTensors: - std::memcpy(&p.Aux_CTX_Tensors, buf, attr_size); - break; - case kNVTEFusedAttnFwdParamsCuSeqlensQ: - std::memcpy(&p.cu_seqlens_q, buf, attr_size); - break; - case kNVTEFusedAttnFwdParamsCuSeqlensKV: - std::memcpy(&p.cu_seqlens_kv, buf, attr_size); - break; - case kNVTEFusedAttnFwdParamsCuSeqlensQPadded: - std::memcpy(&p.cu_seqlens_q_padded, buf, attr_size); - break; - case kNVTEFusedAttnFwdParamsCuSeqlensKVPadded: - std::memcpy(&p.cu_seqlens_kv_padded, buf, attr_size); - break; - case kNVTEFusedAttnFwdParamsPageTableK: - std::memcpy(&p.page_table_k, buf, attr_size); - break; - case kNVTEFusedAttnFwdParamsPageTableV: - std::memcpy(&p.page_table_v, buf, attr_size); - break; - case kNVTEFusedAttnFwdParamsRngState: - std::memcpy(&p.rng_state, buf, attr_size); - break; - case kNVTEFusedAttnFwdParamsIsTraining: - uint8_to_bool(buf, p.is_training); - break; - case kNVTEFusedAttnFwdParamsCudaGraph: - uint8_to_bool(buf, p.cuda_graph); - break; - case kNVTEFusedAttnFwdParamsReturnMaxLogit: - uint8_to_bool(buf, p.return_max_logit); - break; - case kNVTEFusedAttnFwdParamsAttnMaskType: - std::memcpy(&p.attn_mask_type, buf, attr_size); - break; - case kNVTEFusedAttnFwdParamsBiasType: - std::memcpy(&p.bias_type, buf, attr_size); - break; - case kNVTEFusedAttnFwdParamsWindowSizeLeft: - std::memcpy(&p.window_size_left, buf, attr_size); - break; - case kNVTEFusedAttnFwdParamsWindowSizeRight: - std::memcpy(&p.window_size_right, buf, attr_size); - break; - case kNVTEFusedAttnFwdParamsBottomRightDiagonal: - uint8_to_bool(buf, p.bottom_right_diagonal); - break; - case kNVTEFusedAttnFwdParamsSoftmaxType: - std::memcpy(&p.softmax_type, buf, attr_size); - break; - case kNVTEFusedAttnFwdParamsDropout: - std::memcpy(&p.dropout, buf, attr_size); - break; - case kNVTEFusedAttnFwdParamsAttnScale: - std::memcpy(&p.attn_scale, buf, attr_size); - break; - case kNVTEFusedAttnFwdParamsQKVLayout: - std::memcpy(&p.qkv_layout, buf, attr_size); - break; - case kNVTEFusedAttnFwdParamsOFormat: - std::memcpy(&p.o_format, buf, attr_size); - break; - case kNVTEFusedAttnFwdParamsQKVScaleInvFormat: - std::memcpy(&p.qkv_scale_inv_format, buf, attr_size); - break; - case kNVTEFusedAttnFwdParamsMaxSeqlenQ: - std::memcpy(&p.max_seqlen_q, buf, attr_size); - break; - case kNVTEFusedAttnFwdParamsMaxSeqlenKV: - std::memcpy(&p.max_seqlen_kv, buf, attr_size); - break; - case kNVTEFusedAttnFwdParamsWorkspace: - std::memcpy(&p.workspace, buf, attr_size); - break; - case kNVTEFusedAttnFwdParamsStream: - std::memcpy(&p.stream, buf, attr_size); - break; - default: - NVTE_ERROR("Unsupported NVTEFusedAttnFwdParamsAttribute (got ", static_cast(attr), ")"); - } -} - -NVTEFusedAttnBwdParams nvte_create_fused_attn_bwd_params() { - return new transformer_engine::fused_attn::FusedAttnBwdParams{}; -} - -void nvte_destroy_fused_attn_bwd_params(NVTEFusedAttnBwdParams params) { - delete transformer_engine::fused_attn::get_fused_attn_bwd_params_mutable(params); -} - -void nvte_get_fused_attn_bwd_params_attribute(NVTEFusedAttnBwdParams params, - NVTEFusedAttnBwdParamsAttribute attr, void *buf, - size_t size_in_bytes, size_t *size_written) { - using namespace transformer_engine; - using namespace transformer_engine::fused_attn; - NVTE_CHECK(attr < kNVTEFusedAttnBwdParamsNumAttributes, - "Invalid NVTEFusedAttnBwdParamsAttribute (got ", static_cast(attr), ")"); - const auto &attr_size = FusedAttnBwdParams::attr_sizes[attr]; - if (size_written != nullptr) { - *size_written = attr_size; - } - if (buf == nullptr) { - return; - } - NVTE_CHECK(size_in_bytes >= attr_size, "Buffer is too small for attribute (need ", attr_size, - ", got ", size_in_bytes, ")"); - const auto &p = *get_fused_attn_bwd_params(params); - switch (attr) { - case kNVTEFusedAttnBwdParamsQ: - std::memcpy(buf, &p.Q, attr_size); - break; - case kNVTEFusedAttnBwdParamsK: - std::memcpy(buf, &p.K, attr_size); - break; - case kNVTEFusedAttnBwdParamsV: - std::memcpy(buf, &p.V, attr_size); - break; - case kNVTEFusedAttnBwdParamsO: - std::memcpy(buf, &p.O, attr_size); - break; - case kNVTEFusedAttnBwdParamsDO: - std::memcpy(buf, &p.dO, attr_size); - break; - case kNVTEFusedAttnBwdParamsS: - std::memcpy(buf, &p.S, attr_size); - break; - case kNVTEFusedAttnBwdParamsDP: - std::memcpy(buf, &p.dP, attr_size); - break; - case kNVTEFusedAttnBwdParamsAuxCtxTensors: - std::memcpy(buf, &p.Aux_CTX_Tensors, attr_size); - break; - case kNVTEFusedAttnBwdParamsDQ: - std::memcpy(buf, &p.dQ, attr_size); - break; - case kNVTEFusedAttnBwdParamsDK: - std::memcpy(buf, &p.dK, attr_size); - break; - case kNVTEFusedAttnBwdParamsDV: - std::memcpy(buf, &p.dV, attr_size); - break; - case kNVTEFusedAttnBwdParamsDBias: - std::memcpy(buf, &p.dBias, attr_size); - break; - case kNVTEFusedAttnBwdParamsDSoftmaxOffset: - std::memcpy(buf, &p.dSoftmaxOffset, attr_size); - break; - case kNVTEFusedAttnBwdParamsCuSeqlensQ: - std::memcpy(buf, &p.cu_seqlens_q, attr_size); - break; - case kNVTEFusedAttnBwdParamsCuSeqlensKV: - std::memcpy(buf, &p.cu_seqlens_kv, attr_size); - break; - case kNVTEFusedAttnBwdParamsCuSeqlensQPadded: - std::memcpy(buf, &p.cu_seqlens_q_padded, attr_size); - break; - case kNVTEFusedAttnBwdParamsCuSeqlensKVPadded: - std::memcpy(buf, &p.cu_seqlens_kv_padded, attr_size); - break; - case kNVTEFusedAttnBwdParamsDeterministic: - bool_to_uint8(p.deterministic, buf); - break; - case kNVTEFusedAttnBwdParamsCudaGraph: - bool_to_uint8(p.cuda_graph, buf); - break; - case kNVTEFusedAttnBwdParamsAttnMaskType: - std::memcpy(buf, &p.attn_mask_type, attr_size); - break; - case kNVTEFusedAttnBwdParamsBiasType: - std::memcpy(buf, &p.bias_type, attr_size); - break; - case kNVTEFusedAttnBwdParamsWindowSizeLeft: - std::memcpy(buf, &p.window_size_left, attr_size); - break; - case kNVTEFusedAttnBwdParamsWindowSizeRight: - std::memcpy(buf, &p.window_size_right, attr_size); - break; - case kNVTEFusedAttnBwdParamsBottomRightDiagonal: - bool_to_uint8(p.bottom_right_diagonal, buf); - break; - case kNVTEFusedAttnBwdParamsSoftmaxType: - std::memcpy(buf, &p.softmax_type, attr_size); - break; - case kNVTEFusedAttnBwdParamsDropout: - std::memcpy(buf, &p.dropout, attr_size); - break; - case kNVTEFusedAttnBwdParamsAttnScale: - std::memcpy(buf, &p.attn_scale, attr_size); - break; - case kNVTEFusedAttnBwdParamsQKVLayout: - std::memcpy(buf, &p.qkv_layout, attr_size); - break; - case kNVTEFusedAttnBwdParamsOFormat: - std::memcpy(buf, &p.o_format, attr_size); - break; - case kNVTEFusedAttnBwdParamsDOFormat: - std::memcpy(buf, &p.do_format, attr_size); - break; - case kNVTEFusedAttnBwdParamsDQKVLayout: - std::memcpy(buf, &p.dqkv_layout, attr_size); - break; - case kNVTEFusedAttnBwdParamsQKVScaleInvFormat: - std::memcpy(buf, &p.qkv_scale_inv_format, attr_size); - break; - case kNVTEFusedAttnBwdParamsDOScaleInvFormat: - std::memcpy(buf, &p.do_scale_inv_format, attr_size); - break; - case kNVTEFusedAttnBwdParamsMaxSeqlenQ: - std::memcpy(buf, &p.max_seqlen_q, attr_size); - break; - case kNVTEFusedAttnBwdParamsMaxSeqlenKV: - std::memcpy(buf, &p.max_seqlen_kv, attr_size); - break; - case kNVTEFusedAttnBwdParamsWorkspace: - std::memcpy(buf, &p.workspace, attr_size); - break; - case kNVTEFusedAttnBwdParamsStream: - std::memcpy(buf, &p.stream, attr_size); - break; - default: - NVTE_ERROR("Unsupported NVTEFusedAttnBwdParamsAttribute (got ", static_cast(attr), ")"); - } -} - -void nvte_set_fused_attn_bwd_params_attribute(NVTEFusedAttnBwdParams params, - NVTEFusedAttnBwdParamsAttribute attr, const void *buf, - size_t size_in_bytes) { - using namespace transformer_engine; - using namespace transformer_engine::fused_attn; - NVTE_CHECK(attr < kNVTEFusedAttnBwdParamsNumAttributes, - "Invalid NVTEFusedAttnBwdParamsAttribute (got ", static_cast(attr), ")"); - const auto &attr_size = FusedAttnBwdParams::attr_sizes[attr]; - NVTE_CHECK(buf != nullptr, "Input buffer must not be NULL."); - NVTE_CHECK(size_in_bytes >= attr_size, "Buffer is too small for attribute (need ", attr_size, - ", got ", size_in_bytes, ")"); - auto &p = *get_fused_attn_bwd_params_mutable(params); - switch (attr) { - case kNVTEFusedAttnBwdParamsQ: - std::memcpy(&p.Q, buf, attr_size); - break; - case kNVTEFusedAttnBwdParamsK: - std::memcpy(&p.K, buf, attr_size); - break; - case kNVTEFusedAttnBwdParamsV: - std::memcpy(&p.V, buf, attr_size); - break; - case kNVTEFusedAttnBwdParamsO: - std::memcpy(&p.O, buf, attr_size); - break; - case kNVTEFusedAttnBwdParamsDO: - std::memcpy(&p.dO, buf, attr_size); - break; - case kNVTEFusedAttnBwdParamsS: - std::memcpy(&p.S, buf, attr_size); - break; - case kNVTEFusedAttnBwdParamsDP: - std::memcpy(&p.dP, buf, attr_size); - break; - case kNVTEFusedAttnBwdParamsAuxCtxTensors: - std::memcpy(&p.Aux_CTX_Tensors, buf, attr_size); - break; - case kNVTEFusedAttnBwdParamsDQ: - std::memcpy(&p.dQ, buf, attr_size); - break; - case kNVTEFusedAttnBwdParamsDK: - std::memcpy(&p.dK, buf, attr_size); - break; - case kNVTEFusedAttnBwdParamsDV: - std::memcpy(&p.dV, buf, attr_size); - break; - case kNVTEFusedAttnBwdParamsDBias: - std::memcpy(&p.dBias, buf, attr_size); - break; - case kNVTEFusedAttnBwdParamsDSoftmaxOffset: - std::memcpy(&p.dSoftmaxOffset, buf, attr_size); - break; - case kNVTEFusedAttnBwdParamsCuSeqlensQ: - std::memcpy(&p.cu_seqlens_q, buf, attr_size); - break; - case kNVTEFusedAttnBwdParamsCuSeqlensKV: - std::memcpy(&p.cu_seqlens_kv, buf, attr_size); - break; - case kNVTEFusedAttnBwdParamsCuSeqlensQPadded: - std::memcpy(&p.cu_seqlens_q_padded, buf, attr_size); - break; - case kNVTEFusedAttnBwdParamsCuSeqlensKVPadded: - std::memcpy(&p.cu_seqlens_kv_padded, buf, attr_size); - break; - case kNVTEFusedAttnBwdParamsDeterministic: - uint8_to_bool(buf, p.deterministic); - break; - case kNVTEFusedAttnBwdParamsCudaGraph: - uint8_to_bool(buf, p.cuda_graph); - break; - case kNVTEFusedAttnBwdParamsAttnMaskType: - std::memcpy(&p.attn_mask_type, buf, attr_size); - break; - case kNVTEFusedAttnBwdParamsBiasType: - std::memcpy(&p.bias_type, buf, attr_size); - break; - case kNVTEFusedAttnBwdParamsWindowSizeLeft: - std::memcpy(&p.window_size_left, buf, attr_size); - break; - case kNVTEFusedAttnBwdParamsWindowSizeRight: - std::memcpy(&p.window_size_right, buf, attr_size); - break; - case kNVTEFusedAttnBwdParamsBottomRightDiagonal: - uint8_to_bool(buf, p.bottom_right_diagonal); - break; - case kNVTEFusedAttnBwdParamsSoftmaxType: - std::memcpy(&p.softmax_type, buf, attr_size); - break; - case kNVTEFusedAttnBwdParamsDropout: - std::memcpy(&p.dropout, buf, attr_size); - break; - case kNVTEFusedAttnBwdParamsAttnScale: - std::memcpy(&p.attn_scale, buf, attr_size); - break; - case kNVTEFusedAttnBwdParamsQKVLayout: - std::memcpy(&p.qkv_layout, buf, attr_size); - break; - case kNVTEFusedAttnBwdParamsOFormat: - std::memcpy(&p.o_format, buf, attr_size); - break; - case kNVTEFusedAttnBwdParamsDOFormat: - std::memcpy(&p.do_format, buf, attr_size); - break; - case kNVTEFusedAttnBwdParamsDQKVLayout: - std::memcpy(&p.dqkv_layout, buf, attr_size); - break; - case kNVTEFusedAttnBwdParamsQKVScaleInvFormat: - std::memcpy(&p.qkv_scale_inv_format, buf, attr_size); - break; - case kNVTEFusedAttnBwdParamsDOScaleInvFormat: - std::memcpy(&p.do_scale_inv_format, buf, attr_size); - break; - case kNVTEFusedAttnBwdParamsMaxSeqlenQ: - std::memcpy(&p.max_seqlen_q, buf, attr_size); - break; - case kNVTEFusedAttnBwdParamsMaxSeqlenKV: - std::memcpy(&p.max_seqlen_kv, buf, attr_size); - break; - case kNVTEFusedAttnBwdParamsWorkspace: - std::memcpy(&p.workspace, buf, attr_size); - break; - case kNVTEFusedAttnBwdParamsStream: - std::memcpy(&p.stream, buf, attr_size); - break; - default: - NVTE_ERROR("Unsupported NVTEFusedAttnBwdParamsAttribute (got ", static_cast(attr), ")"); - } -} diff --git a/transformer_engine/common/fused_attn/config_and_params.h b/transformer_engine/common/fused_attn/config_and_params.h deleted file mode 100644 index ec4e95d1552..00000000000 --- a/transformer_engine/common/fused_attn/config_and_params.h +++ /dev/null @@ -1,469 +0,0 @@ -/************************************************************************* - * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. - * - * See LICENSE for license information. - ************************************************************************/ - -/*! \file config_and_params.h - * \brief Internal objects for fused-attention config and parameter handles. - */ - -#ifndef TRANSFORMER_ENGINE_COMMON_FUSED_ATTN_CONFIG_AND_PARAMS_H_ -#define TRANSFORMER_ENGINE_COMMON_FUSED_ATTN_CONFIG_AND_PARAMS_H_ - -#include -#include - -#include "common/common.h" -#include "transformer_engine/fused_attn.h" -#include "utils.h" - -namespace transformer_engine { -namespace fused_attn { - -enum class Backend { F16, FP8 }; -enum class Pass { Fwd, Bwd }; - -inline constexpr const char *backend_name(Backend b) { return b == Backend::F16 ? "f16" : "fp8"; } -inline constexpr const char *pass_name(Pass p) { return p == Pass::Fwd ? "fwd" : "bwd"; } - -struct FusedAttnConfig { - // Basic attention settings - bool is_training = true; - bool deterministic = false; - bool cuda_graph = false; - bool return_max_logit = false; - NVTE_Mask_Type attn_mask_type = NVTE_NO_MASK; - NVTE_Bias_Type bias_type = NVTE_NO_BIAS; - int64_t window_size_left = -1; - int64_t window_size_right = -1; - bool bottom_right_diagonal = true; - NVTE_Softmax_Type softmax_type = NVTE_VANILLA_SOFTMAX; - NVTEScalingMode scaling_mode = NVTE_DELAYED_TENSOR_SCALING; - float dropout = 0.0f; - float attn_scale = 1.0f; - - // Tensor types - NVTEDType qkv_dtype = kNVTEBFloat16; - NVTEDType o_dtype = kNVTEBFloat16; - NVTEDType do_dtype = kNVTEBFloat16; - NVTEDType dqkv_dtype = kNVTEBFloat16; - - // Tensor layouts - NVTE_QKV_Layout qkv_layout = NVTE_QKV_Layout_NOT_SET; - NVTE_QKV_Format o_format = NVTE_QKV_Format_NOT_SET; - NVTE_QKV_Format do_format = NVTE_QKV_Format_NOT_SET; - NVTE_QKV_Layout dqkv_layout = NVTE_QKV_Layout_NOT_SET; - NVTE_QKV_Format qkv_scale_inv_format = NVTE_QKV_Format_NOT_SET; - NVTE_QKV_Format do_scale_inv_format = NVTE_QKV_Format_NOT_SET; - - // Tensor dimensions - size_t batch_size = 0; - size_t num_attn_heads = 0; - size_t num_gqa_groups = 0; - size_t head_dim_qk = 0; - size_t head_dim_v = 0; - size_t max_seqlen_q = 0; - size_t max_seqlen_kv = 0; - size_t num_tokens_q = 0; - size_t num_tokens_kv = 0; - - // Paged KV dimensions - size_t num_pages_k = 0; - size_t num_pages_v = 0; - size_t page_size_k = 0; - size_t page_size_v = 0; - size_t max_pages_per_seq_k = 0; - size_t max_pages_per_seq_v = 0; - - // Bias dimensions - size_t bias_batch_size = 0; - size_t bias_num_heads = 0; - size_t bias_seqlen_q = 0; - size_t bias_seqlen_kv = 0; - - // ============================================================================ - // Internal fields: not part of attr_sizes[] or NVTEFusedAttnConfigAttribute, and - // unreachable from nvte_set_fused_attn_config_attribute, hence "internal". - // - // - Keyed: do participate in operator<, and can distinguish graphs in the cache, - // e.g. device_id; - // - Unkeyed: do not participate in operator<, and can not distinguish graphs - // in the cache; used for convenience purposes only; run derive() to fill them - // ============================================================================ - - // Keyed: - // - distinguish graphs on different GPUs in a single-process run - int device_id = -1; - - // Unkeyed: - // - support query directions - bool check_for_forward_support = true; - bool check_for_backward_support = true; - // - whether derived fields have been filled - bool is_derived = false; - // - common attributes - NVTE_QKV_Format qkv_format = NVTE_QKV_Format_NOT_SET; - NVTE_QKV_Format q_format = NVTE_QKV_Format_NOT_SET; - NVTE_QKV_Format kv_format = NVTE_QKV_Format_NOT_SET; - bool is_ragged_q = false; - bool is_ragged_kv = false; - bool is_paged_kv = false; - bool is_padding = false; - bool is_causal = false; - bool is_causal_bottom_right = false; - bool is_bias = false; - bool is_alibi = false; - bool is_softmax_offset = false; - bool is_dropout = false; - // - FP8 recipes - bool is_o_in_f16 = false; - bool is_tensor_scaling = false; - bool is_mxfp8 = false; - bool is_delayed_scaling_fwd = false; - bool is_delayed_scaling_bwd = false; - bool is_current_scaling_fwd = false; - bool is_current_scaling_bwd = false; - bool is_mxfp8_fwd = false; - bool is_mxfp8_bwd = false; - // - cu_seqlens vs actual_seqlens for THD or padding masks - bool uses_cu_seqlens_directly = false; - // - bucket the batch size and token counts for THD - size_t bucketed_batch_size = 0; - size_t bucketed_num_tokens_q = 0; - size_t bucketed_num_tokens_kv = 0; - // - packed (TH1) vs dense (BHS1) graphs and stats - bool uses_ragged_graph = false; - bool uses_ragged_stats = false; - size_t graph_batch_size_fwd = 0; - size_t graph_batch_size_bwd = 0; - size_t graph_max_seqlen_q = 0; - size_t graph_max_seqlen_kv = 0; - // - ragged offset widths and multipliers - bool needs_64bit_ragged_offset = false; - DType ragged_offset_type_fwd = DType::kInt32; - DType ragged_offset_type_bwd = DType::kInt32; - RaggedOffsetMultipliers ragged_offset_mults; - - static constexpr size_t attr_sizes[] = { - // Basic attention settings - sizeof(uint8_t), // is_training - sizeof(uint8_t), // deterministic - sizeof(uint8_t), // cuda_graph - sizeof(uint8_t), // return_max_logit - sizeof(NVTE_Mask_Type), // attn_mask_type - sizeof(NVTE_Bias_Type), // bias_type - sizeof(int64_t), // window_size_left - sizeof(int64_t), // window_size_right - sizeof(uint8_t), // bottom_right_diagonal - sizeof(NVTE_Softmax_Type), // softmax_type - sizeof(NVTEScalingMode), // scaling_mode - sizeof(float), // dropout - sizeof(float), // attn_scale - // Tensor types - sizeof(NVTEDType), // qkv_dtype - sizeof(NVTEDType), // o_dtype - sizeof(NVTEDType), // do_dtype - sizeof(NVTEDType), // dqkv_dtype - // Tensor layouts - sizeof(NVTE_QKV_Layout), // qkv_layout - sizeof(NVTE_QKV_Format), // o_format - sizeof(NVTE_QKV_Format), // do_format - sizeof(NVTE_QKV_Layout), // dqkv_layout - sizeof(NVTE_QKV_Format), // qkv_scale_inv_format - sizeof(NVTE_QKV_Format), // do_scale_inv_format - // Tensor dimensions - sizeof(size_t), // batch_size - sizeof(size_t), // num_attn_heads - sizeof(size_t), // num_gqa_groups - sizeof(size_t), // head_dim_qk - sizeof(size_t), // head_dim_v - sizeof(size_t), // max_seqlen_q - sizeof(size_t), // max_seqlen_kv - sizeof(size_t), // num_tokens_q - sizeof(size_t), // num_tokens_kv - // Paged KV dimensions - sizeof(size_t), // num_pages_k - sizeof(size_t), // num_pages_v - sizeof(size_t), // page_size_k - sizeof(size_t), // page_size_v - sizeof(size_t), // max_pages_per_seq_k - sizeof(size_t), // max_pages_per_seq_v - // Bias dimensions - sizeof(size_t), // bias_batch_size - sizeof(size_t), // bias_num_heads - sizeof(size_t), // bias_seqlen_q - sizeof(size_t), // bias_seqlen_kv - }; - - static_assert(sizeof(attr_sizes) / sizeof(attr_sizes[0]) == kNVTEFusedAttnConfigNumAttributes, - "attr_sizes must have one entry per NVTEFusedAttnConfigAttribute; add the size of " - "the new attribute alongside its enumerator."); - - bool operator<(const FusedAttnConfig &rhs) const { - return std::tie(is_training, deterministic, cuda_graph, return_max_logit, attn_mask_type, - bias_type, window_size_left, window_size_right, bottom_right_diagonal, - softmax_type, scaling_mode, dropout, attn_scale, qkv_dtype, o_dtype, do_dtype, - dqkv_dtype, qkv_layout, o_format, do_format, dqkv_layout, qkv_scale_inv_format, - do_scale_inv_format, batch_size, num_attn_heads, num_gqa_groups, head_dim_qk, - head_dim_v, max_seqlen_q, max_seqlen_kv, num_tokens_q, num_tokens_kv, - num_pages_k, num_pages_v, page_size_k, page_size_v, max_pages_per_seq_k, - max_pages_per_seq_v, bias_batch_size, bias_num_heads, bias_seqlen_q, - bias_seqlen_kv, device_id) < - std::tie(rhs.is_training, rhs.deterministic, rhs.cuda_graph, rhs.return_max_logit, - rhs.attn_mask_type, rhs.bias_type, rhs.window_size_left, rhs.window_size_right, - rhs.bottom_right_diagonal, rhs.softmax_type, rhs.scaling_mode, rhs.dropout, - rhs.attn_scale, rhs.qkv_dtype, rhs.o_dtype, rhs.do_dtype, rhs.dqkv_dtype, - rhs.qkv_layout, rhs.o_format, rhs.do_format, rhs.dqkv_layout, - rhs.qkv_scale_inv_format, rhs.do_scale_inv_format, rhs.batch_size, - rhs.num_attn_heads, rhs.num_gqa_groups, rhs.head_dim_qk, rhs.head_dim_v, - rhs.max_seqlen_q, rhs.max_seqlen_kv, rhs.num_tokens_q, rhs.num_tokens_kv, - rhs.num_pages_k, rhs.num_pages_v, rhs.page_size_k, rhs.page_size_v, - rhs.max_pages_per_seq_k, rhs.max_pages_per_seq_v, rhs.bias_batch_size, - rhs.bias_num_heads, rhs.bias_seqlen_q, rhs.bias_seqlen_kv, rhs.device_id); - } - - // Derive relevant fields based on input fields that have been set by the caller. They are - // read by the graph build, cache lookup, and support query. - void derive(); - - // Assert that derive() has run, for code about to read a derived field. - void check_derived() const { - NVTE_CHECK(is_derived, - "FusedAttnConfig's derived fields are not set. Please run " - "FusedAttnConfig::derive() first."); - } - - // Return a normalized copy of this config to be used as a key for the cuDNN graph cache. - // It drops fields that are either invariant (e.g. attn_scale) or irrelevant (e.g. dO/dQKV dtypes - // and `deterministic` for forward, and `return_max_logit` for backward). - FusedAttnConfig make_cache_key(Pass pass) const; - - // Return a string representation of this config for level-2 cache diagnostics. - std::string to_string() const; -}; - -inline const FusedAttnConfig *get_fused_attn_config(NVTEFusedAttnConfig config) { - NVTE_CHECK(config != nullptr, "NVTEFusedAttnConfig must not be NULL."); - return reinterpret_cast(config); -} - -inline FusedAttnConfig *get_fused_attn_config_mutable(NVTEFusedAttnConfig config) { - NVTE_CHECK(config != nullptr, "NVTEFusedAttnConfig must not be NULL."); - return reinterpret_cast(config); -} - -struct FusedAttnFwdParams { - // Input tensors - NVTETensor Q = nullptr; - NVTETensor K = nullptr; - NVTETensor V = nullptr; - NVTETensor Bias = nullptr; - NVTETensor SoftmaxOffset = nullptr; - // Intermediate tensors - NVTETensor S = nullptr; - // Output tensor - NVTETensor O = nullptr; - // Auxiliary context tensor pack - NVTETensorPack *Aux_CTX_Tensors = nullptr; - // Miscellaneous tensors - NVTETensor cu_seqlens_q = nullptr; - NVTETensor cu_seqlens_kv = nullptr; - NVTETensor cu_seqlens_q_padded = nullptr; - NVTETensor cu_seqlens_kv_padded = nullptr; - NVTETensor page_table_k = nullptr; - NVTETensor page_table_v = nullptr; - NVTETensor rng_state = nullptr; - // Scalars - bool is_training = true; - bool cuda_graph = false; - bool return_max_logit = false; - NVTE_Mask_Type attn_mask_type = NVTE_NO_MASK; - NVTE_Bias_Type bias_type = NVTE_NO_BIAS; - int64_t window_size_left = -1; - int64_t window_size_right = -1; - bool bottom_right_diagonal = true; - NVTE_Softmax_Type softmax_type = NVTE_VANILLA_SOFTMAX; - float dropout = 0.0f; - float attn_scale = 1.0f; - NVTE_QKV_Layout qkv_layout = NVTE_QKV_Layout_NOT_SET; - NVTE_QKV_Format o_format = NVTE_QKV_Format_NOT_SET; - NVTE_QKV_Format qkv_scale_inv_format = NVTE_QKV_Format_NOT_SET; - size_t max_seqlen_q = 0; - size_t max_seqlen_kv = 0; - // Workspace and stream - NVTETensor workspace = nullptr; - cudaStream_t stream = nullptr; - - static constexpr size_t attr_sizes[] = { - // Tensor handles - sizeof(NVTETensor), // Q - sizeof(NVTETensor), // K - sizeof(NVTETensor), // V - sizeof(NVTETensor), // Bias - sizeof(NVTETensor), // SoftmaxOffset - sizeof(NVTETensor), // S - sizeof(NVTETensor), // O - sizeof(NVTETensorPack *), // Aux_CTX_Tensors - sizeof(NVTETensor), // cu_seqlens_q - sizeof(NVTETensor), // cu_seqlens_kv - sizeof(NVTETensor), // cu_seqlens_q_padded - sizeof(NVTETensor), // cu_seqlens_kv_padded - sizeof(NVTETensor), // page_table_k - sizeof(NVTETensor), // page_table_v - sizeof(NVTETensor), // rng_state - // Configuration knobs - sizeof(uint8_t), // is_training - sizeof(uint8_t), // cuda_graph - sizeof(uint8_t), // return_max_logit - sizeof(NVTE_Mask_Type), // attn_mask_type - sizeof(NVTE_Bias_Type), // bias_type - sizeof(int64_t), // window_size_left - sizeof(int64_t), // window_size_right - sizeof(uint8_t), // bottom_right_diagonal - sizeof(NVTE_Softmax_Type), // softmax_type - sizeof(float), // dropout - sizeof(float), // attn_scale - sizeof(NVTE_QKV_Layout), // qkv_layout - sizeof(NVTE_QKV_Format), // o_format - sizeof(NVTE_QKV_Format), // qkv_scale_inv_format - sizeof(size_t), // max_seqlen_q - sizeof(size_t), // max_seqlen_kv - // Workspace and stream - sizeof(NVTETensor), // workspace - sizeof(cudaStream_t), // stream - }; - - static_assert(sizeof(attr_sizes) / sizeof(attr_sizes[0]) == kNVTEFusedAttnFwdParamsNumAttributes, - "attr_sizes must have one entry per NVTEFusedAttnFwdParamsAttribute; add the size " - "of the new attribute alongside its enumerator."); - - // Build a FusedAttnConfig from the scalar "knobs" carried here (e.g. attn_mask_type, bias_type) - // and the fields derived from the tensor handles (dtypes, dims, scaling mode, paged-KV and bias - // broadcast shapes). Returns the real execution config; call FusedAttnConfig::make_cache_key on - // it to obtain the normalized cuDNN graph-cache key. - FusedAttnConfig make_config() const; -}; - -inline const FusedAttnFwdParams *get_fused_attn_fwd_params(NVTEFusedAttnFwdParams params) { - NVTE_CHECK(params != nullptr, "NVTEFusedAttnFwdParams must not be NULL."); - return reinterpret_cast(params); -} - -inline FusedAttnFwdParams *get_fused_attn_fwd_params_mutable(NVTEFusedAttnFwdParams params) { - NVTE_CHECK(params != nullptr, "NVTEFusedAttnFwdParams must not be NULL."); - return reinterpret_cast(params); -} - -struct FusedAttnBwdParams { - // Input tensors - NVTETensor Q = nullptr; - NVTETensor K = nullptr; - NVTETensor V = nullptr; - NVTETensor O = nullptr; - NVTETensor dO = nullptr; - NVTETensor S = nullptr; - NVTETensor dP = nullptr; - const NVTETensorPack *Aux_CTX_Tensors = nullptr; - // Output tensors - NVTETensor dQ = nullptr; - NVTETensor dK = nullptr; - NVTETensor dV = nullptr; - NVTETensor dBias = nullptr; - NVTETensor dSoftmaxOffset = nullptr; - // Miscellaneous tensors - NVTETensor cu_seqlens_q = nullptr; - NVTETensor cu_seqlens_kv = nullptr; - NVTETensor cu_seqlens_q_padded = nullptr; - NVTETensor cu_seqlens_kv_padded = nullptr; - // Scalars - bool deterministic = false; - bool cuda_graph = false; - NVTE_Mask_Type attn_mask_type = NVTE_NO_MASK; - NVTE_Bias_Type bias_type = NVTE_NO_BIAS; - int64_t window_size_left = -1; - int64_t window_size_right = -1; - bool bottom_right_diagonal = true; - NVTE_Softmax_Type softmax_type = NVTE_VANILLA_SOFTMAX; - float dropout = 0.0f; - float attn_scale = 1.0f; - NVTE_QKV_Layout qkv_layout = NVTE_QKV_Layout_NOT_SET; - NVTE_QKV_Format o_format = NVTE_QKV_Format_NOT_SET; - NVTE_QKV_Format do_format = NVTE_QKV_Format_NOT_SET; - NVTE_QKV_Layout dqkv_layout = NVTE_QKV_Layout_NOT_SET; - NVTE_QKV_Format qkv_scale_inv_format = NVTE_QKV_Format_NOT_SET; - NVTE_QKV_Format do_scale_inv_format = NVTE_QKV_Format_NOT_SET; - size_t max_seqlen_q = 0; - size_t max_seqlen_kv = 0; - // Workspace and stream - NVTETensor workspace = nullptr; - cudaStream_t stream = nullptr; - - static constexpr size_t attr_sizes[] = { - // Tensor handles - sizeof(NVTETensor), // Q - sizeof(NVTETensor), // K - sizeof(NVTETensor), // V - sizeof(NVTETensor), // O - sizeof(NVTETensor), // dO - sizeof(NVTETensor), // S - sizeof(NVTETensor), // dP - sizeof(const NVTETensorPack *), // Aux_CTX_Tensors - sizeof(NVTETensor), // dQ - sizeof(NVTETensor), // dK - sizeof(NVTETensor), // dV - sizeof(NVTETensor), // dBias - sizeof(NVTETensor), // dSoftmaxOffset - sizeof(NVTETensor), // cu_seqlens_q - sizeof(NVTETensor), // cu_seqlens_kv - sizeof(NVTETensor), // cu_seqlens_q_padded - sizeof(NVTETensor), // cu_seqlens_kv_padded - // Configuration knobs - sizeof(uint8_t), // deterministic - sizeof(uint8_t), // cuda_graph - sizeof(NVTE_Mask_Type), // attn_mask_type - sizeof(NVTE_Bias_Type), // bias_type - sizeof(int64_t), // window_size_left - sizeof(int64_t), // window_size_right - sizeof(uint8_t), // bottom_right_diagonal - sizeof(NVTE_Softmax_Type), // softmax_type - sizeof(float), // dropout - sizeof(float), // attn_scale - sizeof(NVTE_QKV_Layout), // qkv_layout - sizeof(NVTE_QKV_Format), // o_format - sizeof(NVTE_QKV_Format), // do_format - sizeof(NVTE_QKV_Layout), // dqkv_layout - sizeof(NVTE_QKV_Format), // qkv_scale_inv_format - sizeof(NVTE_QKV_Format), // do_scale_inv_format - sizeof(size_t), // max_seqlen_q - sizeof(size_t), // max_seqlen_kv - // Workspace and stream - sizeof(NVTETensor), // workspace - sizeof(cudaStream_t), // stream - }; - - static_assert(sizeof(attr_sizes) / sizeof(attr_sizes[0]) == kNVTEFusedAttnBwdParamsNumAttributes, - "attr_sizes must have one entry per NVTEFusedAttnBwdParamsAttribute; add the size " - "of the new attribute alongside its enumerator."); - - // Build a FusedAttnConfig from the scalar "knobs" carried here (e.g. attn_mask_type, bias_type) - // and the fields derived from the tensor handles (e.g. dtypes, dims, scaling mode and bias broadcast - // shape). Returns the real execution config; call FusedAttnConfig::make_cache_key on it to - // obtain the normalized cuDNN graph-cache key. - FusedAttnConfig make_config() const; -}; - -inline const FusedAttnBwdParams *get_fused_attn_bwd_params(NVTEFusedAttnBwdParams params) { - NVTE_CHECK(params != nullptr, "NVTEFusedAttnBwdParams must not be NULL."); - return reinterpret_cast(params); -} - -inline FusedAttnBwdParams *get_fused_attn_bwd_params_mutable(NVTEFusedAttnBwdParams params) { - NVTE_CHECK(params != nullptr, "NVTEFusedAttnBwdParams must not be NULL."); - return reinterpret_cast(params); -} - -} // namespace fused_attn -} // namespace transformer_engine - -#endif // TRANSFORMER_ENGINE_COMMON_FUSED_ATTN_CONFIG_AND_PARAMS_H_ diff --git a/transformer_engine/common/fused_attn/fused_attn.cpp b/transformer_engine/common/fused_attn/fused_attn.cpp index 0c31936e2f7..c6caf156486 100644 --- a/transformer_engine/common/fused_attn/fused_attn.cpp +++ b/transformer_engine/common/fused_attn/fused_attn.cpp @@ -6,14 +6,9 @@ #include "transformer_engine/fused_attn.h" +#include + #include "../common.h" -#include "../cudnn_utils.h" -#include "../util/cuda_runtime.h" -#include "../util/system.h" -#include "config_and_params.h" -#include "fused_attn_f16_arbitrary_seqlen.h" -#include "fused_attn_fp8.h" -#include "utils.h" namespace transformer_engine { @@ -184,7 +179,7 @@ NVTE_QKV_Format nvte_get_qkv_format(NVTE_QKV_Layout qkv_layout) { // map NVTE_QKV_Layout to NVTE_QKV_Format for Q NVTE_QKV_Format nvte_get_q_format(NVTE_QKV_Layout qkv_layout) { - const NVTE_QKV_Format qkv_format = nvte_get_qkv_format(qkv_layout); + NVTE_QKV_Format qkv_format = nvte_get_qkv_format(qkv_layout); switch (qkv_format) { case NVTE_QKV_Format::NVTE_SBHD: case NVTE_QKV_Format::NVTE_SBHD_2BSHD: @@ -206,7 +201,7 @@ NVTE_QKV_Format nvte_get_q_format(NVTE_QKV_Layout qkv_layout) { // map NVTE_QKV_Layout to NVTE_QKV_Format for KV NVTE_QKV_Format nvte_get_kv_format(NVTE_QKV_Layout qkv_layout) { - const NVTE_QKV_Format qkv_format = nvte_get_qkv_format(qkv_layout); + NVTE_QKV_Format qkv_format = nvte_get_qkv_format(qkv_layout); switch (qkv_format) { case NVTE_QKV_Format::NVTE_SBHD: case NVTE_QKV_Format::NVTE_BSHD_2SBHD: @@ -225,501 +220,3 @@ NVTE_QKV_Format nvte_get_kv_format(NVTE_QKV_Layout qkv_layout) { " in nvte_get_kv_format."); } } - -namespace { - -// The per-thread storage for the diagnostic string -thread_local std::string fused_attn_backend_message_buffer; - -// Records `reason` in the per-thread buffer and `message`; returns with "no backend" -[[nodiscard]] NVTE_Fused_Attn_Backend reject(const char **message, std::string reason) { - if (message != nullptr) { - fused_attn_backend_message_buffer = std::move(reason); - *message = fused_attn_backend_message_buffer.c_str(); - } - return NVTE_Fused_Attn_Backend::NVTE_No_Backend; -} - -} // namespace - -// Fused attention backend query: runs TE's specific rules first, then cuDNN's. First rejection wins. -NVTE_Fused_Attn_Backend nvte_get_fused_attn_backend_v2(NVTEFusedAttnConfig config, - const char **message) { - NVTE_API_CALL(nvte_get_fused_attn_backend_v2); - using namespace transformer_engine; - using namespace transformer_engine::fused_attn; - FusedAttnConfig cfg = *get_fused_attn_config(config); - cfg.derive(); - if (message != nullptr) *message = ""; - - cudnnHandle_t handle = cudnnExecutionPlanManager::Instance().GetHandle(); - const auto cudnn_runtime_version = cudnnGetVersion(); - const int sm_arch = cuda::sm_arch(cuda::current_device()); - - // THD + 64-bit ragged offsets require cuDNN >= 9.5 - if (cfg.needs_64bit_ragged_offset && cudnn_runtime_version < 90500) { - return reject( - message, - "This config requires 64-bit ragged offsets, which is only supported by cuDNN >= 9.5."); - } - - // THD input requires a padding mask - if ((cfg.is_ragged_q || cfg.is_ragged_kv) && !cfg.is_padding) { - return reject( - message, - "THD format requires PADDING / PADDING_CAUSAL / PADDING_CAUSAL_BOTTOM_RIGHT mask."); - } - - if ((cfg.is_ragged_q && cfg.num_tokens_q == 0) || (cfg.is_ragged_kv && cfg.num_tokens_kv == 0)) { - return reject(message, - "THD format requires num_tokens_q / num_tokens_kv to be set for the ragged " - "inputs."); - } - - // Paged KV requires a padding mask - if (cfg.is_paged_kv && !cfg.is_padding) { - return reject(message, - "Paged KV requires PADDING / PADDING_CAUSAL / PADDING_CAUSAL_BOTTOM_RIGHT mask."); - } - - // Paged KV requires cache dimensions to be set - if (cfg.is_paged_kv && - (cfg.num_pages_k == 0 || cfg.num_pages_v == 0 || cfg.page_size_k == 0 || - cfg.page_size_v == 0 || cfg.max_pages_per_seq_k == 0 || cfg.max_pages_per_seq_v == 0)) { - return reject(message, - "Paged KV requires num_pages, page_size and max_pages_per_seq to be set for both " - "K and V."); - } - - // Fused-attention does not support pre-scale bias - if (cfg.bias_type == NVTE_Bias_Type::NVTE_PRE_SCALE_BIAS) { - return reject(message, "Fused attention does not support pre-scale bias."); - } - - const bool is_fp8 = - (cfg.qkv_dtype == NVTEDType::kNVTEFloat8E4M3 || cfg.qkv_dtype == NVTEDType::kNVTEFloat8E5M2); - const bool is_f16_or_bf16 = - (cfg.qkv_dtype == NVTEDType::kNVTEFloat16 || cfg.qkv_dtype == NVTEDType::kNVTEBFloat16); - - auto each_pass = [&](auto &&verdict) -> std::string { - if (cfg.check_for_forward_support) { - std::string reason = verdict(Pass::Fwd); - if (!reason.empty()) return reason; - } - if (cfg.is_training && cfg.check_for_backward_support) { - std::string reason = verdict(Pass::Bwd); - if (!reason.empty()) return reason; - } - return ""; - }; - - // F16/BF16 support checks - if (is_f16_or_bf16) { - if (cfg.is_ragged_q && cfg.is_ragged_kv && sm_arch < 90) { - return reject(message, "F16/BF16 fused attention with THD format requires sm90 or later."); - } - if ((cfg.is_ragged_q || cfg.is_ragged_kv) && sm_arch < 90 && cudnn_runtime_version < 90700) { - return reject(message, - "F16/BF16 fused attention with a ragged Q and non-ragged KV requires cuDNN " - ">= 9.7 before sm90."); - } - const bool has_sliding_window = !(cfg.window_size_left == -1 && - (cfg.window_size_right == -1 || cfg.window_size_right == 0)); - if (cfg.is_causal_bottom_right && has_sliding_window && cfg.max_seqlen_q != cfg.max_seqlen_kv && - cudnn_runtime_version <= 90700 && sm_arch >= 100) { - return reject(message, - "Known cuDNN <= 9.7.0 issue with bottom-right causal masking and a sliding " - "window for cross-attention on sm100. Please upgrade cuDNN."); - } - if (cudnn_runtime_version <= 91500 && cfg.is_training && - (cfg.qkv_format == NVTE_QKV_Format::NVTE_BSHD || - cfg.qkv_format == NVTE_QKV_Format::NVTE_SBHD) && - (cfg.max_seqlen_kv % 128 != 0) && cfg.cuda_graph && !cfg.is_padding) { - return reject(message, "Known cuDNN <= 9.15 issue with CUDA graph. Please upgrade cuDNN."); - } - if (cfg.is_training && cfg.check_for_backward_support && cfg.uses_ragged_stats && - cfg.softmax_type == NVTE_Softmax_Type::NVTE_LEARNABLE_SOFTMAX && - cudnn_runtime_version < 92600) { - return reject( - message, - "Known cuDNN < 9.26.0 issue with THD learnable softmax backward. Please upgrade cuDNN."); - } - - // Run cuDNN support checks - std::string cudnn_reason = - each_pass([&](Pass pass) { return support_verdict_f16(cfg, pass, handle); }); - if (!cudnn_reason.empty()) return reject(message, std::move(cudnn_reason)); - return NVTE_Fused_Attn_Backend::NVTE_F16_arbitrary_seqlen; - } - - // FP8 support checks - if (is_fp8) { - if (cfg.return_max_logit) { - return reject(message, "FP8 fused attention does not support return_max_logit=True."); - } - if (cfg.qkv_format != NVTE_QKV_Format::NVTE_BSHD && - cfg.qkv_format != NVTE_QKV_Format::NVTE_SBHD && - cfg.qkv_format != NVTE_QKV_Format::NVTE_BHSD && - cfg.qkv_format != NVTE_QKV_Format::NVTE_THD) { - return reject(message, "FP8 fused attention supports BSHD/SBHD/BHSD/THD formats, found " + - std::to_string(static_cast(cfg.qkv_format)) + "."); - } - if (cfg.is_training && cfg.check_for_backward_support && - cfg.dqkv_layout != NVTE_QKV_Layout_NOT_SET) { - const NVTE_QKV_Format dqkv_format = nvte_get_qkv_format(cfg.dqkv_layout); - if (dqkv_format != NVTE_QKV_Format::NVTE_BSHD && dqkv_format != NVTE_QKV_Format::NVTE_SBHD && - dqkv_format != NVTE_QKV_Format::NVTE_BHSD && dqkv_format != NVTE_QKV_Format::NVTE_THD) { - return reject(message, - "FP8 fused attention supports BSHD/SBHD/BHSD/THD gradient formats, found " + - std::to_string(static_cast(dqkv_format)) + "."); - } - } - if (cfg.qkv_format == NVTE_QKV_Format::NVTE_THD) { - if (cudnn_runtime_version < 92300) { - return reject(message, - "FP8 fused attention with THD format requires cuDNN 9.23.0 or later!"); - } - if (cfg.is_training && cfg.check_for_backward_support && sm_arch < 100) { - return reject(message, - "FP8 fused attention with THD format supports backward on sm100+ only!"); - } - if (cfg.is_training && cfg.check_for_backward_support && - cfg.softmax_type != NVTE_Softmax_Type::NVTE_VANILLA_SOFTMAX && - cudnn_runtime_version < 92600) { - return reject(message, - "FP8 fused attention with THD format and a sink token requires cuDNN 9.26.0 " - "or later for backward!"); - } - if (sm_arch >= 100 && (cfg.head_dim_qk > 128 || cfg.head_dim_v > 128)) { - return reject(message, - "FP8 fused attention with THD format supports head dimensions up to 128 on " - "sm100+ only!"); - } - } - if (cfg.is_bias) { - return reject(message, "FP8 fused attention does not support pre/post_scale_bias yet!"); - } - if (cfg.is_alibi) { - return reject(message, "FP8 fused attention does not support ALiBi yet!"); - } - const char *const recipe_reason = - "FP8 fused attention only supports FP8DelayedScaling or FP8CurrentScaling or MXFP8 " - "recipes!"; - if (cfg.check_for_forward_support && - !(cfg.is_delayed_scaling_fwd || cfg.is_current_scaling_fwd || cfg.is_mxfp8_fwd)) { - return reject(message, recipe_reason); - } - if (cfg.is_training && cfg.check_for_backward_support && - !(cfg.is_delayed_scaling_bwd || cfg.is_current_scaling_bwd || cfg.is_mxfp8_bwd)) { - return reject(message, recipe_reason); - } - if (cfg.is_mxfp8 && cudnn_runtime_version < 92100) { - return reject(message, "MXFP8 fused attention requires cuDNN 9.21.0 or later!"); - } - - // Run cuDNN support checks - std::string cudnn_reason = - each_pass([&](Pass pass) { return support_verdict_fp8(cfg, pass, handle); }); - if (!cudnn_reason.empty()) return reject(message, std::move(cudnn_reason)); - return NVTE_Fused_Attn_Backend::NVTE_FP8; - } - - // Unsupported dtype - return reject(message, "Unsupported QKV dtype qkv_dtype=" + std::to_string(cfg.qkv_dtype) + " ."); -} - -// Select a backend for fused attention -NVTE_Fused_Attn_Backend nvte_get_fused_attn_backend( - bool is_training, NVTEDType q_dtype, NVTEDType kv_dtype, NVTE_QKV_Layout qkv_layout, - NVTE_Bias_Type bias_type, NVTE_Mask_Type attn_mask_type, NVTE_Softmax_Type softmax_type, - float dropout, size_t num_attn_heads, size_t num_gqa_groups, size_t max_seqlen_q, - size_t max_seqlen_kv, size_t head_dim_qk, size_t head_dim_v, int64_t window_size_left, - int64_t window_size_right, bool return_max_logit, bool cuda_graph, bool deterministic) { - NVTE_API_CALL(nvte_get_fused_attn_backend); - transformer_engine::fused_attn::FusedAttnConfig cfg{}; - cfg.qkv_layout = qkv_layout; - cfg.bias_type = bias_type; - cfg.attn_mask_type = attn_mask_type; - cfg.softmax_type = softmax_type; - cfg.dropout = dropout; - cfg.max_seqlen_q = max_seqlen_q; - cfg.max_seqlen_kv = max_seqlen_kv; - cfg.window_size_left = window_size_left; - cfg.window_size_right = window_size_right; - cfg.cuda_graph = cuda_graph; - NVTE_CHECK(q_dtype == kv_dtype, "Q and KV must have the same data type."); - cfg.qkv_dtype = q_dtype; - cfg.o_dtype = q_dtype; - cfg.do_dtype = q_dtype; - cfg.dqkv_dtype = q_dtype; - cfg.num_attn_heads = num_attn_heads; - cfg.num_gqa_groups = num_gqa_groups; - cfg.head_dim_qk = head_dim_qk; - cfg.head_dim_v = head_dim_v; - cfg.is_training = is_training; - cfg.return_max_logit = return_max_logit; - cfg.deterministic = deterministic; - // fill in the missing fields with the most common use case; - // otherwise it would return NVTE_No_Backend always - cfg.batch_size = 1; - cfg.num_tokens_q = cfg.batch_size * max_seqlen_q; - cfg.num_tokens_kv = cfg.batch_size * max_seqlen_kv; - cfg.o_format = nvte_get_q_format(qkv_layout); - cfg.do_format = cfg.o_format; - cfg.dqkv_layout = qkv_layout; - if (bias_type == NVTE_Bias_Type::NVTE_POST_SCALE_BIAS) { - cfg.bias_batch_size = cfg.batch_size; - cfg.bias_num_heads = num_attn_heads; - cfg.bias_seqlen_q = max_seqlen_q; - cfg.bias_seqlen_kv = max_seqlen_kv; - } - - return nvte_get_fused_attn_backend_v2(reinterpret_cast(&cfg), - /*message=*/nullptr); -} - -// Fused attention forward: create a config based on the params, check which backend supports it, -// and run that backend's implementation. -void nvte_fused_attn_fwd_v2(NVTEFusedAttnFwdParams params) { - NVTE_API_CALL(nvte_fused_attn_fwd_v2); - using namespace transformer_engine; - using namespace transformer_engine::fused_attn; - const FusedAttnFwdParams &p = *get_fused_attn_fwd_params(params); - const Tensor *input_cu_seqlens_q = convertNVTETensorCheck(p.cu_seqlens_q); - const Tensor *input_cu_seqlens_kv = convertNVTETensorCheck(p.cu_seqlens_kv); - const Tensor *input_cu_seqlens_q_padded = convertNVTETensorCheck(p.cu_seqlens_q_padded); - const Tensor *input_cu_seqlens_kv_padded = convertNVTETensorCheck(p.cu_seqlens_kv_padded); - const Tensor *input_page_table_k = convertNVTETensorCheck(p.page_table_k); - const Tensor *input_page_table_v = convertNVTETensorCheck(p.page_table_v); - const Tensor *input_rng_state = convertNVTETensorCheck(p.rng_state); - const Tensor *input_Q = convertNVTETensorCheck(p.Q); - const Tensor *input_K = convertNVTETensorCheck(p.K); - const Tensor *input_V = convertNVTETensorCheck(p.V); - const Tensor *input_Bias = convertNVTETensorCheck(p.Bias); - const Tensor *input_SoftmaxOffset = convertNVTETensorCheck(p.SoftmaxOffset); - Tensor *input_output_S = convertNVTETensorCheck(p.S); - Tensor *output_O = convertNVTETensorCheck(p.O); - Tensor *wkspace = convertNVTETensor(p.workspace); - - auto handle = cudnnExecutionPlanManager::Instance().GetHandle(); - FusedAttnConfig cfg = p.make_config(); - cfg.derive(); - const char *fused_attn_reject_reason = nullptr; - NVTE_Fused_Attn_Backend fused_attention_backend = nvte_get_fused_attn_backend_v2( - reinterpret_cast(&cfg), &fused_attn_reject_reason); - - if (fused_attention_backend == NVTE_Fused_Attn_Backend::NVTE_F16_arbitrary_seqlen) { - fused_attn_arbitrary_seqlen_fwd(cfg, input_Q, input_K, input_V, input_Bias, input_SoftmaxOffset, - output_O, p.Aux_CTX_Tensors, input_cu_seqlens_q, - input_cu_seqlens_kv, input_cu_seqlens_q_padded, - input_cu_seqlens_kv_padded, input_page_table_k, - input_page_table_v, input_rng_state, wkspace, p.stream, handle); - } else if (fused_attention_backend == NVTE_Fused_Attn_Backend::NVTE_FP8) { - fused_attn_fp8_fwd(cfg, input_Q, input_K, input_V, input_SoftmaxOffset, input_output_S, - output_O, p.Aux_CTX_Tensors, input_cu_seqlens_q, input_cu_seqlens_kv, - input_cu_seqlens_q_padded, input_cu_seqlens_kv_padded, input_rng_state, - wkspace, p.stream, handle); - } else { - NVTE_ERROR("Fused attention is not supported for the user configuration: ", - fused_attn_reject_reason); - } -} - -// NVTE fused attention FWD with separate Q, K and V -void nvte_fused_attn_fwd(const NVTETensor Q, const NVTETensor K, const NVTETensor V, - const NVTETensor Bias, const NVTETensor SoftmaxOffset, NVTETensor S, - NVTETensor O, NVTETensorPack *Aux_CTX_Tensors, - const NVTETensor cu_seqlens_q, const NVTETensor cu_seqlens_kv, - const NVTETensor cu_seqlens_q_padded, - const NVTETensor cu_seqlens_kv_padded, const NVTETensor page_table_k, - const NVTETensor page_table_v, const NVTETensor rng_state, - size_t max_seqlen_q, size_t max_seqlen_kv, bool is_training, - bool return_max_logit, bool cuda_graph, float attn_scale, float dropout, - NVTE_QKV_Layout qkv_layout, NVTE_QKV_Format o_format, - NVTE_QKV_Format qkv_scale_inv_format, NVTE_Bias_Type bias_type, - NVTE_Mask_Type attn_mask_type, NVTE_Softmax_Type softmax_type, - int64_t window_size_left, int64_t window_size_right, - bool bottom_right_diagonal, NVTETensor workspace, cudaStream_t stream) { - NVTE_API_CALL(nvte_fused_attn_fwd); - transformer_engine::fused_attn::FusedAttnFwdParams p{}; - p.Q = Q; - p.K = K; - p.V = V; - p.Bias = Bias; - p.SoftmaxOffset = SoftmaxOffset; - p.S = S; - p.O = O; - p.Aux_CTX_Tensors = Aux_CTX_Tensors; - p.cu_seqlens_q = cu_seqlens_q; - p.cu_seqlens_kv = cu_seqlens_kv; - p.cu_seqlens_q_padded = cu_seqlens_q_padded; - p.cu_seqlens_kv_padded = cu_seqlens_kv_padded; - p.page_table_k = page_table_k; - p.page_table_v = page_table_v; - p.rng_state = rng_state; - p.is_training = is_training; - p.cuda_graph = cuda_graph; - p.return_max_logit = return_max_logit; - p.attn_mask_type = attn_mask_type; - p.bias_type = bias_type; - p.window_size_left = window_size_left; - p.window_size_right = window_size_right; - p.bottom_right_diagonal = bottom_right_diagonal; - p.softmax_type = softmax_type; - p.dropout = dropout; - p.attn_scale = attn_scale; - p.qkv_layout = qkv_layout; - p.o_format = o_format; - p.qkv_scale_inv_format = qkv_scale_inv_format; - p.max_seqlen_q = max_seqlen_q; - p.max_seqlen_kv = max_seqlen_kv; - p.workspace = workspace; - p.stream = stream; - nvte_fused_attn_fwd_v2(reinterpret_cast(&p)); -} - -// Fused attention backward. Same shape as nvte_fused_attn_fwd_v2, whose comment sketches the path; -// this one asks the selector for backward support and probes the backward builders. -void nvte_fused_attn_bwd_v2(NVTEFusedAttnBwdParams params) { - NVTE_API_CALL(nvte_fused_attn_bwd_v2); - using namespace transformer_engine; - using namespace transformer_engine::fused_attn; - const FusedAttnBwdParams &p = *get_fused_attn_bwd_params(params); - const Tensor *input_cu_seqlens_q = convertNVTETensorCheck(p.cu_seqlens_q); - const Tensor *input_cu_seqlens_kv = convertNVTETensorCheck(p.cu_seqlens_kv); - const Tensor *input_cu_seqlens_q_padded = convertNVTETensorCheck(p.cu_seqlens_q_padded); - const Tensor *input_cu_seqlens_kv_padded = convertNVTETensorCheck(p.cu_seqlens_kv_padded); - const Tensor *input_Q = convertNVTETensorCheck(p.Q); - const Tensor *input_K = convertNVTETensorCheck(p.K); - const Tensor *input_V = convertNVTETensorCheck(p.V); - const Tensor *input_O = convertNVTETensorCheck(p.O); - const Tensor *input_dO = convertNVTETensorCheck(p.dO); - const Tensor *input_S = convertNVTETensorCheck(p.S); - Tensor *input_output_dP = convertNVTETensorCheck(p.dP); - Tensor *output_dQ = convertNVTETensorCheck(p.dQ); - Tensor *output_dK = convertNVTETensorCheck(p.dK); - Tensor *output_dV = convertNVTETensorCheck(p.dV); - Tensor *output_dBias = convertNVTETensorCheck(p.dBias); - Tensor *output_dSoftmaxOffset = convertNVTETensorCheck(p.dSoftmaxOffset); - Tensor *wkspace = convertNVTETensor(p.workspace); - - auto handle = cudnnExecutionPlanManager::Instance().GetHandle(); - FusedAttnConfig cfg = p.make_config(); - // Derived here, not by the query below: the query works on its own copy, and it is this config - // that goes on to the backend and must arrive with its derived fields filled in. - cfg.derive(); - const char *fused_attn_reject_reason = nullptr; - NVTE_Fused_Attn_Backend fused_attention_backend = nvte_get_fused_attn_backend_v2( - reinterpret_cast(&cfg), &fused_attn_reject_reason); - - if (fused_attention_backend == NVTE_Fused_Attn_Backend::NVTE_F16_arbitrary_seqlen) { - size_t i = 0; - Tensor *output_S = convertNVTETensorCheck(p.Aux_CTX_Tensors->tensors[i++]); - Tensor *input_rng_state = convertNVTETensorCheck(p.Aux_CTX_Tensors->tensors[i++]); - Tensor *input_Bias = nullptr, *input_SoftmaxOffset = nullptr; - if ((p.bias_type != NVTE_NO_BIAS) && (p.bias_type != NVTE_ALIBI)) { - input_Bias = convertNVTETensorCheck(p.Aux_CTX_Tensors->tensors[i++]); - } - if (p.softmax_type != NVTE_VANILLA_SOFTMAX) { - input_SoftmaxOffset = convertNVTETensorCheck(p.Aux_CTX_Tensors->tensors[i++]); - } - fused_attn_arbitrary_seqlen_bwd( - cfg, input_Q, input_K, input_V, input_O, input_dO, input_Bias, input_SoftmaxOffset, - output_S, output_dQ, output_dK, output_dV, output_dBias, output_dSoftmaxOffset, - input_cu_seqlens_q, input_cu_seqlens_kv, input_cu_seqlens_q_padded, - input_cu_seqlens_kv_padded, input_rng_state, wkspace, p.stream, handle); - } else if (fused_attention_backend == NVTE_Fused_Attn_Backend::NVTE_FP8) { - size_t i = 0; - const Tensor *input_M = convertNVTETensorCheck(p.Aux_CTX_Tensors->tensors[i++]); - const Tensor *input_rng_state = convertNVTETensorCheck(p.Aux_CTX_Tensors->tensors[i++]); - const Tensor *input_SoftmaxOffset = nullptr; - if (p.softmax_type != NVTE_VANILLA_SOFTMAX) { - input_SoftmaxOffset = convertNVTETensorCheck(p.Aux_CTX_Tensors->tensors[i++]); - } - const Tensor *input_dO_f16 = nullptr; - if (input_dO->scaling_mode == NVTE_MXFP8_1D_SCALING) { - input_dO_f16 = convertNVTETensorCheck(p.Aux_CTX_Tensors->tensors[i++]); - } - fused_attn_fp8_bwd(cfg, input_Q, input_K, input_V, input_O, input_dO, input_dO_f16, input_M, - input_S, input_SoftmaxOffset, input_output_dP, output_dQ, output_dK, - output_dV, output_dSoftmaxOffset, input_cu_seqlens_q, input_cu_seqlens_kv, - input_cu_seqlens_q_padded, input_cu_seqlens_kv_padded, input_rng_state, - wkspace, p.stream, handle); - } else { - NVTE_ERROR("Fused attention is not supported for this configuration: ", - fused_attn_reject_reason); - } -} - -// NVTE fused attention BWD with separate Q, K and V -void nvte_fused_attn_bwd(const NVTETensor Q, const NVTETensor K, const NVTETensor V, - const NVTETensor O, const NVTETensor dO, const NVTETensor S, NVTETensor dP, - const NVTETensorPack *Aux_CTX_Tensors, NVTETensor dQ, NVTETensor dK, - NVTETensor dV, NVTETensor dBias, NVTETensor dSoftmaxOffset, - const NVTETensor cu_seqlens_q, const NVTETensor cu_seqlens_kv, - const NVTETensor cu_seqlens_q_padded, - const NVTETensor cu_seqlens_kv_padded, size_t max_seqlen_q, - size_t max_seqlen_kv, float attn_scale, float dropout, - NVTE_QKV_Layout qkv_layout, NVTE_QKV_Format o_format, - NVTE_QKV_Format do_format, NVTE_QKV_Layout dqkv_layout, - NVTE_QKV_Format qkv_scale_inv_format, NVTE_QKV_Format do_scale_inv_format, - NVTE_Bias_Type bias_type, NVTE_Mask_Type attn_mask_type, - NVTE_Softmax_Type softmax_type, int64_t window_size_left, - int64_t window_size_right, bool bottom_right_diagonal, bool deterministic, - bool cuda_graph, NVTETensor workspace, cudaStream_t stream) { - NVTE_API_CALL(nvte_fused_attn_bwd); - transformer_engine::fused_attn::FusedAttnBwdParams p{}; - p.Q = Q; - p.K = K; - p.V = V; - p.O = O; - p.dO = dO; - p.S = S; - p.dP = dP; - p.Aux_CTX_Tensors = Aux_CTX_Tensors; - p.dQ = dQ; - p.dK = dK; - p.dV = dV; - p.dBias = dBias; - p.dSoftmaxOffset = dSoftmaxOffset; - p.cu_seqlens_q = cu_seqlens_q; - p.cu_seqlens_kv = cu_seqlens_kv; - p.cu_seqlens_q_padded = cu_seqlens_q_padded; - p.cu_seqlens_kv_padded = cu_seqlens_kv_padded; - p.deterministic = deterministic; - p.cuda_graph = cuda_graph; - p.attn_mask_type = attn_mask_type; - p.bias_type = bias_type; - p.window_size_left = window_size_left; - p.window_size_right = window_size_right; - p.bottom_right_diagonal = bottom_right_diagonal; - p.softmax_type = softmax_type; - p.dropout = dropout; - p.attn_scale = attn_scale; - p.qkv_layout = qkv_layout; - p.o_format = o_format; - p.do_format = do_format; - p.dqkv_layout = dqkv_layout; - p.qkv_scale_inv_format = qkv_scale_inv_format; - p.do_scale_inv_format = do_scale_inv_format; - p.max_seqlen_q = max_seqlen_q; - p.max_seqlen_kv = max_seqlen_kv; - p.workspace = workspace; - p.stream = stream; - nvte_fused_attn_bwd_v2(reinterpret_cast(&p)); -} - -uint32_t nvte_get_runtime_num_segments(NVTETensor cu_seqlen, NVTETensor workspace, size_t len, - cudaStream_t stream) { - NVTE_API_CALL(nvte_get_runtime_num_segments); - using namespace transformer_engine::fused_attn; - return GetRuntimeNumSegments(cu_seqlen, workspace, len, stream); -} - -void nvte_populate_rng_state_async(NVTETensor rng_state_dst, const NVTETensor seed, - size_t q_max_seqlen, size_t kv_max_seqlen, - NVTE_Fused_Attn_Backend backend, cudaStream_t stream) { - NVTE_API_CALL(nvte_populate_rng_state_async); - using namespace transformer_engine::fused_attn; - PopulateRngStateAsync(rng_state_dst, seed, q_max_seqlen, kv_max_seqlen, backend, stream); -} diff --git a/transformer_engine/common/fused_attn/fused_attn_f16_arbitrary_seqlen.cu b/transformer_engine/common/fused_attn/fused_attn_f16_arbitrary_seqlen.cu deleted file mode 100644 index fcbde955496..00000000000 --- a/transformer_engine/common/fused_attn/fused_attn_f16_arbitrary_seqlen.cu +++ /dev/null @@ -1,1193 +0,0 @@ -/************************************************************************* - * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. - * - * See LICENSE for license information. - ************************************************************************/ - -#include -#include -#include -#include - -#include - -#include "../common.h" -#include "../cudnn_utils.h" -#include "../util/cuda_runtime.h" -#include "../util/system.h" -#include "fused_attn_f16_arbitrary_seqlen.h" -#include "graph_cache.h" -#include "graph_cache_debug.h" -#include "utils.h" - -namespace transformer_engine { -namespace fused_attn { - -namespace fe = cudnn_frontend; - -using F16FwdGraphAndTensors = - std::tuple, - std::shared_ptr, // Q - std::shared_ptr, // K - std::shared_ptr, // V - std::shared_ptr, // attn_scale - std::shared_ptr, // O - std::shared_ptr, // S1 - std::shared_ptr, // S2 - std::shared_ptr, // bias - std::shared_ptr, // softmax_offset - std::shared_ptr, // seq_q / cu_seq_len_q - std::shared_ptr, // seq_kv / cu_seq_len_kv - std::shared_ptr, // page_table_k - std::shared_ptr, // page_table_v - std::shared_ptr, // offset_q - std::shared_ptr, // offset_k - std::shared_ptr, // offset_v - std::shared_ptr, // offset_o - std::shared_ptr, // offset_stats - std::shared_ptr, // dropout_seed - std::shared_ptr>; // dropout_offset - -static F16FwdGraphAndTensors create_graph_f16_fwd(const FusedAttnConfig &cfg) { - const int64_t b = static_cast(cfg.graph_batch_size_fwd); - const int64_t s_q = static_cast(cfg.graph_max_seqlen_q); - const int64_t s_kv = static_cast(cfg.graph_max_seqlen_kv); - const cudnn_frontend::DataType_t tensorType = - get_cudnn_fe_dtype(static_cast(cfg.qkv_dtype)); - const int64_t h = static_cast(cfg.num_attn_heads); - const int64_t hg = static_cast(cfg.num_gqa_groups); - const int64_t d_qk = static_cast(cfg.head_dim_qk); - const int64_t d_v = static_cast(cfg.head_dim_v); - const int64_t num_pages_k = static_cast(cfg.num_pages_k); - const int64_t num_pages_v = static_cast(cfg.num_pages_v); - const int64_t page_size_k = static_cast(cfg.page_size_k); - const int64_t page_size_v = static_cast(cfg.page_size_v); - const int64_t max_pages_per_seq_k = static_cast(cfg.max_pages_per_seq_k); - const int64_t max_pages_per_seq_v = static_cast(cfg.max_pages_per_seq_v); - const int64_t bias_b = static_cast(cfg.bias_batch_size); - const int64_t bias_h = static_cast(cfg.bias_num_heads); - const int64_t bias_sq = static_cast(cfg.bias_seqlen_q); - const int64_t bias_skv = static_cast(cfg.bias_seqlen_kv); - const int64_t window_size_left = cfg.window_size_left; - const int64_t window_size_right = cfg.window_size_right; - const bool return_max_logit = cfg.return_max_logit; - const float dropout_probability = cfg.dropout; - const NVTE_QKV_Layout qkv_layout = cfg.qkv_layout; - const bool bottom_right_diagonal = cfg.bottom_right_diagonal; - const bool is_bias = cfg.is_bias; - const bool is_alibi = cfg.is_alibi; - const bool is_causal = cfg.is_causal; - const bool is_causal_bottom_right = cfg.is_causal_bottom_right; - const bool is_padding = cfg.is_padding; - const bool is_paged_kv = cfg.is_paged_kv; - const bool is_softmax_offset = cfg.is_softmax_offset; - const bool is_dropout = cfg.is_dropout; - const bool is_ragged_q = cfg.is_ragged_q; - const bool is_ragged_kv = cfg.is_ragged_kv; - const bool use_cu_seqlens_directly = cfg.uses_cu_seqlens_directly; - const auto cudnn_runtime_version = cudnnGetVersion(); - const bool use_ragged_stats = cfg.uses_ragged_stats; - const DType ragged_offset_type = cfg.ragged_offset_type_fwd; - const RaggedOffsetMultipliers offset_mults = cfg.ragged_offset_mults; - const bool generate_stats = true; // Always return stats - - auto mha_graph = std::make_shared(); - mha_graph->set_io_data_type(tensorType) - .set_intermediate_data_type(fe::DataType_t::FLOAT) - .set_compute_data_type(fe::DataType_t::FLOAT); - - std::shared_ptr Q, K, V, attn_scale, softmax_offset; - std::shared_ptr bias, seq_q, seq_kv; - std::shared_ptr page_table_k, page_table_v; - std::shared_ptr offset_q, offset_k, offset_v, offset_o, - offset_stats; - std::shared_ptr dropout_seed, dropout_offset; - - std::vector q_stride(4); - std::vector k_stride(4); - std::vector v_stride(4); - generateMatrixStrides(b, h, s_q, s_kv, d_qk, q_stride.data(), qkv_layout, - NVTE_QKV_Matrix::NVTE_Q_Matrix); - if (is_paged_kv) { - generateMatrixStrides(num_pages_k, hg, page_size_k, page_size_v, d_qk, k_stride.data(), - qkv_layout, NVTE_QKV_Matrix::NVTE_K_Matrix); - generateMatrixStrides(num_pages_v, hg, page_size_k, page_size_v, d_v, v_stride.data(), - qkv_layout, NVTE_QKV_Matrix::NVTE_V_Matrix); - } else { - generateMatrixStrides(b, hg, s_q, s_kv, d_qk, k_stride.data(), qkv_layout, - NVTE_QKV_Matrix::NVTE_K_Matrix); - generateMatrixStrides(b, hg, s_q, s_kv, d_v, v_stride.data(), qkv_layout, - NVTE_QKV_Matrix::NVTE_V_Matrix); - } - - Q = mha_graph->tensor( - fe::graph::Tensor_attributes().set_name("Q").set_dim({b, h, s_q, d_qk}).set_stride(q_stride)); - if (is_ragged_q) { - offset_q = mha_graph->tensor(fe::graph::Tensor_attributes() - .set_name("offset_q") - .set_dim({b + 1, 1, 1, 1}) - .set_stride({1, 1, 1, 1}) - .set_data_type(get_cudnn_fe_dtype(ragged_offset_type))); - Q->set_ragged_offset(offset_q); - if (use_cu_seqlens_directly) { - Q->set_ragged_offset_multiplier(offset_mults.q); - } - } - K = mha_graph->tensor(fe::graph::Tensor_attributes().set_name("K").set_stride(k_stride)); - V = mha_graph->tensor(fe::graph::Tensor_attributes().set_name("V").set_stride(v_stride)); - if (is_paged_kv) { - K->set_dim({num_pages_k, hg, page_size_k, d_qk}); - V->set_dim({num_pages_v, hg, page_size_v, d_v}); - } else if (is_ragged_kv) { - offset_k = mha_graph->tensor(fe::graph::Tensor_attributes() - .set_name("offset_k") - .set_dim({b + 1, 1, 1, 1}) - .set_stride({1, 1, 1, 1}) - .set_data_type(get_cudnn_fe_dtype(ragged_offset_type))); - offset_v = mha_graph->tensor(fe::graph::Tensor_attributes() - .set_name("offset_v") - .set_dim({b + 1, 1, 1, 1}) - .set_stride({1, 1, 1, 1}) - .set_data_type(get_cudnn_fe_dtype(ragged_offset_type))); - K->set_dim({b, hg, s_kv, d_qk}).set_ragged_offset(offset_k); - V->set_dim({b, hg, s_kv, d_v}).set_ragged_offset(offset_v); - if (use_cu_seqlens_directly) { - K->set_ragged_offset_multiplier(offset_mults.k); - V->set_ragged_offset_multiplier(offset_mults.v); - } - } else { - K->set_dim({b, hg, s_kv, d_qk}); - V->set_dim({b, hg, s_kv, d_v}); - } - - attn_scale = mha_graph->tensor(fe::graph::Tensor_attributes() - .set_name("attn_scale") - .set_dim({1, 1, 1, 1}) - .set_stride({1, 1, 1, 1}) - .set_is_pass_by_value(true) - .set_data_type(fe::DataType_t::FLOAT)); - - fe::graph::SDPA_attributes sdpa_options; - sdpa_options = fe::graph::SDPA_attributes() - .set_name("flash_attention") - .set_generate_stats(generate_stats) - .set_attn_scale(attn_scale); - - fe::DiagonalAlignment_t const &diagonal_alignment = bottom_right_diagonal - ? fe::DiagonalAlignment_t::BOTTOM_RIGHT - : fe::DiagonalAlignment_t::TOP_LEFT; - sdpa_options.set_diagonal_alignment(diagonal_alignment); - if (cudnn_runtime_version >= 90200 && window_size_left != -1) { - sdpa_options.set_diagonal_band_left_bound(window_size_left + 1); - } - if (cudnn_runtime_version >= 90600 && window_size_right != -1) { - sdpa_options.set_diagonal_band_right_bound(window_size_right); - } - if (is_causal || is_causal_bottom_right) { - sdpa_options.set_diagonal_band_right_bound(0); - } - - sdpa_options.set_alibi_mask(is_alibi); - - if (is_bias) { - bias = mha_graph->tensor( - fe::graph::Tensor_attributes() - .set_name("bias") - .set_dim({bias_b, bias_h, bias_sq, bias_skv}) - .set_stride({bias_h * bias_sq * bias_skv, bias_sq * bias_skv, bias_skv, 1})); - sdpa_options.set_bias(bias); - } - - if (is_padding) { - if (use_cu_seqlens_directly) { - // seq_q/seq_kv keep their tuple slots but hold (b+1)-shaped cu_seqlen tensors. - seq_q = mha_graph->tensor(fe::graph::Tensor_attributes() - .set_name("cu_seq_len_q") - .set_dim({b + 1, 1, 1, 1}) - .set_stride({1, 1, 1, 1}) - .set_data_type(fe::DataType_t::INT32)); - seq_kv = mha_graph->tensor(fe::graph::Tensor_attributes() - .set_name("cu_seq_len_kv") - .set_dim({b + 1, 1, 1, 1}) - .set_stride({1, 1, 1, 1}) - .set_data_type(fe::DataType_t::INT32)); - sdpa_options.set_padding_mask(is_padding).set_cu_seq_len_q(seq_q).set_cu_seq_len_kv(seq_kv); - // cu_seq_len (and the ragged offset multiplier) are unified-engine-only. - // Pin the implementation so an unsupported config fails with the unified - // engine's specific error instead of auto-selection's generic failure. - sdpa_options.set_implementation(fe::AttentionImplementation_t::UNIFIED); - } else { - seq_q = mha_graph->tensor(fe::graph::Tensor_attributes() - .set_name("seq_q") - .set_dim({b, 1, 1, 1}) - .set_stride({1, 1, 1, 1}) - .set_data_type(fe::DataType_t::INT32)); - seq_kv = mha_graph->tensor(fe::graph::Tensor_attributes() - .set_name("seq_kv") - .set_dim({b, 1, 1, 1}) - .set_stride({1, 1, 1, 1}) - .set_data_type(fe::DataType_t::INT32)); - sdpa_options.set_padding_mask(is_padding).set_seq_len_q(seq_q).set_seq_len_kv(seq_kv); - } - } - - if (is_paged_kv) { - page_table_k = - mha_graph->tensor(fe::graph::Tensor_attributes() - .set_name("page_table_k") - .set_dim({b, 1, max_pages_per_seq_k, 1}) - .set_stride({{max_pages_per_seq_k, max_pages_per_seq_v, 1, 1}}) - .set_data_type(fe::DataType_t::INT32)); - page_table_v = - mha_graph->tensor(fe::graph::Tensor_attributes() - .set_name("page_table_v") - .set_dim({b, 1, max_pages_per_seq_v, 1}) - .set_stride({{max_pages_per_seq_v, max_pages_per_seq_v, 1, 1}}) - .set_data_type(fe::DataType_t::INT32)); - sdpa_options.set_paged_attention_k_table(page_table_k); - sdpa_options.set_paged_attention_v_table(page_table_v); - sdpa_options.set_paged_attention_max_seq_len_kv(static_cast(s_kv)); - } - - if (is_dropout) { - dropout_seed = mha_graph->tensor(fe::graph::Tensor_attributes() - .set_name("Seed") - .set_dim({1, 1, 1, 1}) - .set_stride({1, 1, 1, 1}) - .set_data_type(fe::DataType_t::INT64)); - dropout_offset = mha_graph->tensor(fe::graph::Tensor_attributes() - .set_name("Offset") - .set_dim({1, 1, 1, 1}) - .set_stride({1, 1, 1, 1}) - .set_data_type(fe::DataType_t::INT64)); - sdpa_options.set_dropout(dropout_probability, dropout_seed, dropout_offset); - } - - if (is_softmax_offset) { - softmax_offset = mha_graph->tensor(fe::graph::Tensor_attributes() - .set_name("softmax_offset") - .set_dim({1, h, 1, 1}) - .set_stride({h, 1, 1, 1}) - .set_data_type(fe::DataType_t::FLOAT)); - sdpa_options.set_sink_token(softmax_offset); - } - - std::shared_ptr Max; - if (use_ragged_stats) { - offset_stats = mha_graph->tensor(fe::graph::Tensor_attributes() - .set_name("offset_stats") - .set_dim({b + 1, 1, 1, 1}) - .set_stride({1, 1, 1, 1}) - .set_data_type(get_cudnn_fe_dtype(ragged_offset_type))); - } - if (return_max_logit) { - Max = mha_graph->tensor(fe::graph::Tensor_attributes() - .set_name("Max") - .set_dim({b, h, s_q, 1}) - .set_data_type(fe::DataType_t::FLOAT)); - if (use_ragged_stats) { - Max->set_stride({h * s_q, 1, h, 1}).set_ragged_offset(offset_stats); - if (use_cu_seqlens_directly) { - Max->set_ragged_offset_multiplier(offset_mults.stats); - } - } else { - Max->set_stride({h * s_q, s_q, 1, 1}); - } - sdpa_options.set_logit_max(Max); - } - - auto [O, Stats] = mha_graph->sdpa(Q, K, V, std::move(sdpa_options)); - - std::vector o_stride(4); - generateMatrixStrides(b, h, s_q, s_kv, d_v, o_stride.data(), qkv_layout, - NVTE_QKV_Matrix::NVTE_O_Matrix); - O->set_output(true).set_dim({b, h, s_q, d_v}).set_stride(o_stride); - if (is_ragged_q) { - offset_o = mha_graph->tensor(fe::graph::Tensor_attributes() - .set_name("offset_o") - .set_dim({b + 1, 1, 1, 1}) - .set_stride({1, 1, 1, 1}) - .set_data_type(get_cudnn_fe_dtype(ragged_offset_type))); - O->set_ragged_offset(offset_o); - if (use_cu_seqlens_directly) { - O->set_ragged_offset_multiplier(offset_mults.o); - } - } - - Stats->set_output(true).set_data_type(fe::DataType_t::FLOAT).set_dim({b, h, s_q, 1}); - if (use_ragged_stats) { - Stats->set_stride({h * s_q, 1, h, 1}).set_ragged_offset(offset_stats); - if (use_cu_seqlens_directly) { - Stats->set_ragged_offset_multiplier(offset_mults.stats); - } - } else { - Stats->set_stride({h * s_q, s_q, 1, 1}); - } - - std::tuple, // Q - std::shared_ptr, // K - std::shared_ptr, // V - std::shared_ptr, // attn_scale - std::shared_ptr> // O - key_tensors_tuple = std::make_tuple(Q, K, V, attn_scale, O); - auto Stats_tuple = - return_max_logit ? std::make_tuple(Stats, Max) : std::make_tuple(Stats, nullptr); - auto bias_tuple = is_bias ? std::make_tuple(bias) : std::make_tuple(nullptr); - auto softmax_offset_tuple = - is_softmax_offset ? std::make_tuple(softmax_offset) : std::make_tuple(nullptr); - auto padding_tuple = - is_padding ? std::make_tuple(seq_q, seq_kv) : std::make_tuple(nullptr, nullptr); - auto page_table_tuple = - is_paged_kv ? std::make_tuple(page_table_k, page_table_v) : std::make_tuple(nullptr, nullptr); - auto offset_qo_tuple = - is_ragged_q ? std::make_tuple(offset_q, offset_o) : std::make_tuple(nullptr, nullptr); - auto offset_kv_tuple = - is_ragged_kv ? std::make_tuple(offset_k, offset_v) : std::make_tuple(nullptr, nullptr); - auto offset_s_tuple = use_ragged_stats ? std::make_tuple(offset_stats) : std::make_tuple(nullptr); - auto dropout_tuple = is_dropout ? std::make_tuple(dropout_seed, dropout_offset) - : std::make_tuple(nullptr, nullptr); - - return std::tuple_cat(std::make_tuple(mha_graph), key_tensors_tuple, Stats_tuple, bias_tuple, - softmax_offset_tuple, padding_tuple, page_table_tuple, offset_qo_tuple, - offset_kv_tuple, offset_s_tuple, dropout_tuple); -} - -void fused_attn_arbitrary_seqlen_fwd_impl( - const FusedAttnConfig &cfg, void *devPtrQ, void *devPtrK, void *devPtrV, void *devPtrBias, - void *devPtrSoftmaxOffset, void *devPtrS1, void *devPtrS2, void *devPtrO, - void *devPtrDropoutSeed, void *devPtrDropoutOffset, void *devPtrCuSeqlensQ, - void *devPtrCuSeqlensKV, void *devPtrPageTableK, void *devPtrPageTableV, - void *devPtrSeqOffsetsQ, void *devPtrSeqOffsetsKV, void *workspace, size_t *workspace_size, - cudaStream_t stream, cudnnHandle_t handle) { - using namespace transformer_engine; - - cfg.check_derived(); - const int64_t b = static_cast(cfg.graph_batch_size_fwd); - const DType ragged_offset_type = cfg.ragged_offset_type_fwd; - const int64_t actual_b = static_cast(cfg.batch_size); - const bool use_ragged_stats = cfg.uses_ragged_stats; - const RaggedOffsetMultipliers offset_mults = cfg.ragged_offset_mults; - - const bool return_max_logit = cfg.return_max_logit; - float scaling_factor = cfg.attn_scale; - const bool is_bias = cfg.is_bias; - const bool is_padding = cfg.is_padding; - const bool is_softmax_offset = cfg.is_softmax_offset; - const bool is_dropout = cfg.is_dropout; - const bool is_ragged_q = cfg.is_ragged_q; - const bool is_ragged_kv = cfg.is_ragged_kv; - const bool is_paged_kv = cfg.is_paged_kv; - // Newer versions of cuDNN SDPA can accept sequence lengths directly as a cumulative - // tensor, and can accept ragged offsets in arbitrary units (such as tokens) instead - // of elements. Take advantage of this if possible to avoid 2 extra kernel calls. - const bool use_cu_seqlens_directly = cfg.uses_cu_seqlens_directly; - - try { - auto cache_entry = get_graph(cfg, handle); - auto [mha_graph, Q, K, V, attn_scale, O, S1, S2, bias, softmax_offset, seq_q, seq_kv, - page_table_k, page_table_v, offset_q, offset_o, offset_k, offset_v, offset_stats, - dropout_seed, dropout_offset] = cache_entry->graph_and_tensors; - - // This graph is going to be used, so finish the build the cache deferred. - build_plans(Backend::F16, Pass::Fwd, *cache_entry); - - // Exit to request upper level API to allocate memory if needed - // n.b. Care should be taken to align each of the added worksapce tensors to their type. - // We do this by adding padding at the end of each separate allocation. - // When passing cu_seqlens* directly to cuDNN SDPA, no conversion workspace is - // needed: cuDNN consumes the user's cu_seqlens buffers as-is. - auto plan_workspace_size = alignTo<16>(mha_graph->get_workspace_size()); - const size_t num_bytes_per_seqlen = alignTo<16>(b * sizeof(int32_t)); - const size_t num_bytes_per_ragged_offset = - alignTo<16>(((b + 1) * typeToNumBits(ragged_offset_type)) / 8); - size_t actual_seqlen_workspace_size = 0; - size_t seqlen_offsets_workspace_size = 0; - if (!use_cu_seqlens_directly) { - if (is_padding) { - actual_seqlen_workspace_size = 2 * num_bytes_per_seqlen; - } - if (is_ragged_q || is_ragged_kv) { - const size_t count = - 2 * (static_cast(is_ragged_q) + static_cast(is_ragged_kv)); - seqlen_offsets_workspace_size = - (use_ragged_stats ? count + 1 : count) * num_bytes_per_ragged_offset; - } - } - if (workspace == nullptr) { - *workspace_size = - plan_workspace_size + actual_seqlen_workspace_size + seqlen_offsets_workspace_size; - return; - } - // cuDNN stream check needs to be moved here to support dummy kernel calls with - // null streams for sizing the cuDNN workspace. - NVTE_CHECK_CUDNN(cudnnSetStream(handle, stream)); - - // Build variant pack - std::unordered_map, void *> variant_pack = { - {Q, devPtrQ}, {K, devPtrK}, {V, devPtrV}, {attn_scale, &scaling_factor}, - {O, devPtrO}, {S1, devPtrS1}}; - - if (return_max_logit) { - variant_pack[S2] = devPtrS2; - } - - if (is_bias) { - variant_pack[bias] = devPtrBias; - } - - if (is_padding) { - if (use_cu_seqlens_directly) { - variant_pack[seq_q] = devPtrCuSeqlensQ; - variant_pack[seq_kv] = devPtrCuSeqlensKV; - } else { - constexpr size_t nthreads_per_block = 128; - const size_t grid = (b + nthreads_per_block - 1) / nthreads_per_block; - void *devActualSeqlenQ = static_cast(workspace) + plan_workspace_size; - void *devActualSeqlenKV = static_cast(devActualSeqlenQ) + num_bytes_per_seqlen; - cu_seqlens_to_actual_seqlens<<>>( - actual_b, b, static_cast(devPtrCuSeqlensQ), - static_cast(devPtrCuSeqlensKV), - static_cast(devActualSeqlenQ), static_cast(devActualSeqlenKV)); - NVTE_CHECK_CUDA(cudaGetLastError()); - variant_pack[seq_q] = devActualSeqlenQ; - variant_pack[seq_kv] = devActualSeqlenKV; - } - } - - if (is_paged_kv) { - variant_pack[page_table_k] = devPtrPageTableK; - variant_pack[page_table_v] = devPtrPageTableV; - } - - if (use_cu_seqlens_directly) { - // The token-unit cu_seqlens_padded buffers serve as the ragged offsets; the engine - // applies the per-tensor multipliers set at graph build time. - if (is_ragged_q) { - variant_pack[offset_q] = devPtrSeqOffsetsQ; - variant_pack[offset_o] = devPtrSeqOffsetsQ; - } - if (is_ragged_kv) { - void *devOffsetsKV = offset_mults.kv_from_q ? devPtrSeqOffsetsQ : devPtrSeqOffsetsKV; - variant_pack[offset_k] = devOffsetsKV; - variant_pack[offset_v] = devOffsetsKV; - } - if (use_ragged_stats) { - variant_pack[offset_stats] = devPtrSeqOffsetsQ; - } - } else if (is_ragged_q || is_ragged_kv) { - constexpr size_t nthreads_per_block = 128; - const size_t grid = (b + nthreads_per_block) / nthreads_per_block; - void *devOffsets = - static_cast(workspace) + plan_workspace_size + actual_seqlen_workspace_size; - void *devOffsetsQ = nullptr; - void *devOffsetsO = nullptr; - if (is_ragged_q) { - devOffsetsQ = devOffsets; - devOffsetsO = static_cast(devOffsetsQ) + num_bytes_per_ragged_offset; - } - void *devOffsetsK = nullptr; - void *devOffsetsV = nullptr; - if (is_ragged_kv) { - devOffsetsK = static_cast(devOffsets) + - static_cast(is_ragged_q) * 2 * num_bytes_per_ragged_offset; - devOffsetsV = static_cast(devOffsetsK) + num_bytes_per_ragged_offset; - } - void *devOffsetsS = nullptr; - if (use_ragged_stats) { - devOffsetsS = static_cast(devOffsets) + - (static_cast(is_ragged_q) + static_cast(is_ragged_kv)) * 2 * - num_bytes_per_ragged_offset; - } - cu_seqlens_padded_to_offsets<<>>( - offset_mults, actual_b, b, static_cast(devPtrSeqOffsetsQ), - static_cast(devPtrSeqOffsetsKV), ragged_offset_type, devOffsetsQ, devOffsetsK, - devOffsetsV, devOffsetsO, devOffsetsS); - NVTE_CHECK_CUDA(cudaGetLastError()); - if (is_ragged_q) { - variant_pack[offset_q] = devOffsetsQ; - variant_pack[offset_o] = devOffsetsO; - } - if (is_ragged_kv) { - variant_pack[offset_k] = devOffsetsK; - variant_pack[offset_v] = devOffsetsV; - } - if (use_ragged_stats) { - variant_pack[offset_stats] = devOffsetsS; - } - } - - if (is_dropout) { - variant_pack[dropout_seed] = devPtrDropoutSeed; - variant_pack[dropout_offset] = devPtrDropoutOffset; - } - - if (is_softmax_offset) { - variant_pack[softmax_offset] = devPtrSoftmaxOffset; - } - - NVTE_CHECK_CUDNN_FE(mha_graph->execute(handle, variant_pack, workspace)); - graph_cache_debug::record_execute(Backend::F16, Pass::Fwd); - } catch (cudnn_frontend::cudnnException &e) { - NVTE_ERROR(e.what()); - } -} - -using F16BwdGraphAndTensors = - std::tuple, - std::shared_ptr, // q - std::shared_ptr, // k - std::shared_ptr, // v - std::shared_ptr, // o - std::shared_ptr, // dO - std::shared_ptr, // stats - std::shared_ptr, // attn_scale - std::shared_ptr, // dQ - std::shared_ptr, // dK - std::shared_ptr, // dV - std::shared_ptr, // bias - std::shared_ptr, // dBias - std::shared_ptr, // softmax_offset - std::shared_ptr, // d_softmax_offset - std::shared_ptr, // seq_q - std::shared_ptr, // seq_kv - std::shared_ptr, // offset_q - std::shared_ptr, // offset_k - std::shared_ptr, // offset_v - std::shared_ptr, // offset_o - std::shared_ptr, // offset_stats - std::shared_ptr, // dropout_seed - std::shared_ptr>; // dropout_offset - -static F16BwdGraphAndTensors create_graph_f16_bwd(const FusedAttnConfig &cfg) { - const int64_t b = static_cast(cfg.graph_batch_size_bwd); - const int64_t s_q = static_cast(cfg.graph_max_seqlen_q); - const int64_t s_kv = static_cast(cfg.graph_max_seqlen_kv); - const cudnn_frontend::DataType_t tensorType = - get_cudnn_fe_dtype(static_cast(cfg.qkv_dtype)); - const int64_t h = static_cast(cfg.num_attn_heads); - const int64_t hg = static_cast(cfg.num_gqa_groups); - const int64_t d_qk = static_cast(cfg.head_dim_qk); - const int64_t d_v = static_cast(cfg.head_dim_v); - const int64_t bias_b = static_cast(cfg.bias_batch_size); - const int64_t bias_h = static_cast(cfg.bias_num_heads); - const int64_t bias_sq = static_cast(cfg.bias_seqlen_q); - const int64_t bias_skv = static_cast(cfg.bias_seqlen_kv); - const int64_t window_size_left = cfg.window_size_left; - const int64_t window_size_right = cfg.window_size_right; - const float dropout_probability = cfg.dropout; - const NVTE_QKV_Layout qkv_layout = cfg.qkv_layout; - const bool bottom_right_diagonal = cfg.bottom_right_diagonal; - const bool deterministic = cfg.deterministic; - const bool is_bias = cfg.is_bias; - const bool is_alibi = cfg.is_alibi; - const bool is_causal = cfg.is_causal; - const bool is_causal_bottom_right = cfg.is_causal_bottom_right; - const bool is_padding = cfg.is_padding; - const bool is_softmax_offset = cfg.is_softmax_offset; - const bool is_dropout = cfg.is_dropout; - const bool is_ragged_q = cfg.is_ragged_q; - const bool is_ragged_kv = cfg.is_ragged_kv; - const auto cudnn_runtime_version = cudnnGetVersion(); - const bool use_ragged_graph = cfg.uses_ragged_graph; - const bool use_ragged_stats = cfg.uses_ragged_stats; - const DType ragged_offset_type = cfg.ragged_offset_type_bwd; - - auto mha_graph = std::make_shared(); - mha_graph->set_io_data_type(tensorType) - .set_intermediate_data_type(fe::DataType_t::FLOAT) - .set_compute_data_type(fe::DataType_t::FLOAT); - - std::shared_ptr q, k, v, o, dO, stats, attn_scale; - std::shared_ptr bias, dBias, softmax_offset, d_softmax_offset, - seq_q, seq_kv; - std::shared_ptr offset_q, offset_k, offset_v, offset_o, - offset_stats; - std::shared_ptr dropout_seed, dropout_offset; - - std::vector q_stride(4); - std::vector k_stride(4); - std::vector v_stride(4); - std::vector o_stride(4); - generateMatrixStrides(b, h, s_q, s_kv, d_qk, q_stride.data(), qkv_layout, - NVTE_QKV_Matrix::NVTE_Q_Matrix); - generateMatrixStrides(b, hg, s_q, s_kv, d_qk, k_stride.data(), qkv_layout, - NVTE_QKV_Matrix::NVTE_K_Matrix); - generateMatrixStrides(b, hg, s_q, s_kv, d_v, v_stride.data(), qkv_layout, - NVTE_QKV_Matrix::NVTE_V_Matrix); - generateMatrixStrides(b, h, s_q, s_kv, d_v, o_stride.data(), qkv_layout, - NVTE_QKV_Matrix::NVTE_O_Matrix); - - q = mha_graph->tensor( - fe::graph::Tensor_attributes().set_name("Q").set_dim({b, h, s_q, d_qk}).set_stride(q_stride)); - k = mha_graph->tensor(fe::graph::Tensor_attributes() - .set_name("K") - .set_dim({b, hg, s_kv, d_qk}) - .set_stride(k_stride)); - v = mha_graph->tensor(fe::graph::Tensor_attributes() - .set_name("V") - .set_dim({b, hg, s_kv, d_v}) - .set_stride(v_stride)); - o = mha_graph->tensor( - fe::graph::Tensor_attributes().set_name("O").set_dim({b, h, s_q, d_v}).set_stride(o_stride)); - dO = mha_graph->tensor( - fe::graph::Tensor_attributes().set_name("dO").set_dim({b, h, s_q, d_v}).set_stride(o_stride)); - if (is_ragged_q) { - offset_q = mha_graph->tensor(fe::graph::Tensor_attributes() - .set_name("offset_q") - .set_dim({b + 1, 1, 1, 1}) - .set_stride({1, 1, 1, 1}) - .set_data_type(get_cudnn_fe_dtype(ragged_offset_type))); - offset_o = mha_graph->tensor(fe::graph::Tensor_attributes() - .set_name("offset_o") - .set_dim({b + 1, 1, 1, 1}) - .set_stride({1, 1, 1, 1}) - .set_data_type(get_cudnn_fe_dtype(ragged_offset_type))); - q->set_ragged_offset(offset_q); - o->set_ragged_offset(offset_o); - dO->set_ragged_offset(offset_o); - } - if (is_ragged_kv) { - offset_k = mha_graph->tensor(fe::graph::Tensor_attributes() - .set_name("offset_k") - .set_dim({b + 1, 1, 1, 1}) - .set_stride({1, 1, 1, 1}) - .set_data_type(get_cudnn_fe_dtype(ragged_offset_type))); - offset_v = mha_graph->tensor(fe::graph::Tensor_attributes() - .set_name("offset_v") - .set_dim({b + 1, 1, 1, 1}) - .set_stride({1, 1, 1, 1}) - .set_data_type(get_cudnn_fe_dtype(ragged_offset_type))); - k->set_ragged_offset(offset_k); - v->set_ragged_offset(offset_v); - } - - stats = mha_graph->tensor(fe::graph::Tensor_attributes() - .set_name("stats") - .set_dim({b, h, s_q, 1}) - .set_data_type(fe::DataType_t::FLOAT)); - if (use_ragged_stats) { - offset_stats = mha_graph->tensor(fe::graph::Tensor_attributes() - .set_name("offset_stats") - .set_dim({b + 1, 1, 1, 1}) - .set_stride({1, 1, 1, 1}) - .set_data_type(get_cudnn_fe_dtype(ragged_offset_type))); - stats->set_stride({h * s_q, 1, h, 1}).set_ragged_offset(offset_stats); - } else { - stats->set_stride({h * s_q, s_q, 1, 1}); - } - - attn_scale = mha_graph->tensor(fe::graph::Tensor_attributes() - .set_name("attn_scale") - .set_dim({1, 1, 1, 1}) - .set_stride({1, 1, 1, 1}) - .set_is_pass_by_value(true) - .set_data_type(fe::DataType_t::FLOAT)); - - fe::graph::SDPA_backward_attributes sdpa_backward_options; - sdpa_backward_options = fe::graph::SDPA_backward_attributes() - .set_name("flash_attention_backward") - .set_attn_scale(attn_scale); - - if (use_ragged_stats) { - sdpa_backward_options.set_max_total_seq_len_q(s_q); - } - if (is_ragged_kv && use_ragged_graph) { - sdpa_backward_options.set_max_total_seq_len_kv(s_kv); - } - - fe::DiagonalAlignment_t const &diagonal_alignment = bottom_right_diagonal - ? fe::DiagonalAlignment_t::BOTTOM_RIGHT - : fe::DiagonalAlignment_t::TOP_LEFT; - sdpa_backward_options.set_diagonal_alignment(diagonal_alignment); - - if (cudnn_runtime_version >= 90200 && window_size_left != -1) { - sdpa_backward_options.set_diagonal_band_left_bound(window_size_left + 1); - } - if (cudnn_runtime_version >= 90600 && window_size_right != -1) { - sdpa_backward_options.set_diagonal_band_right_bound(window_size_right); - } - if (is_causal || is_causal_bottom_right) { - sdpa_backward_options.set_diagonal_band_right_bound(0); - } - - if (cudnn_runtime_version >= 90000) { - sdpa_backward_options.set_deterministic_algorithm(deterministic); - } - - sdpa_backward_options.set_alibi_mask(is_alibi); - - if (is_bias) { - bias = mha_graph->tensor( - fe::graph::Tensor_attributes() - .set_name("bias") - .set_dim({bias_b, bias_h, bias_sq, bias_skv}) - .set_stride({bias_h * bias_sq * bias_skv, bias_sq * bias_skv, bias_skv, 1})); - sdpa_backward_options.set_bias(bias); - // bias shapes [1, 1, s, s], [b, 1, s, s], [b, h, s, s], [1, h, s, s] are supported for dbias calculation - // bias shape [1, 1, 1, s] is not supported for dbias calculation as of cuDNN 9.18 - if (!((bias_b == 1) && (bias_h == 1) && (bias_sq == 1))) { - dBias = mha_graph->tensor( - fe::graph::Tensor_attributes() - .set_name("dBias") - .set_dim({bias_b, bias_h, bias_sq, bias_skv}) - .set_stride({bias_h * bias_sq * bias_skv, bias_sq * bias_skv, bias_skv, 1})); - sdpa_backward_options.set_dbias(dBias); - } - } - - if (is_padding) { - seq_q = mha_graph->tensor(fe::graph::Tensor_attributes() - .set_name("seq_q") - .set_dim({b, 1, 1, 1}) - .set_stride({1, 1, 1, 1}) - .set_data_type(fe::DataType_t::INT32)); - seq_kv = mha_graph->tensor(fe::graph::Tensor_attributes() - .set_name("seq_kv") - .set_dim({b, 1, 1, 1}) - .set_stride({1, 1, 1, 1}) - .set_data_type(fe::DataType_t::INT32)); - sdpa_backward_options.set_padding_mask(is_padding).set_seq_len_q(seq_q).set_seq_len_kv(seq_kv); - } - - if (is_dropout) { - dropout_seed = mha_graph->tensor(fe::graph::Tensor_attributes() - .set_name("Seed") - .set_dim({1, 1, 1, 1}) - .set_stride({1, 1, 1, 1}) - .set_data_type(fe::DataType_t::INT64)); - dropout_offset = mha_graph->tensor(fe::graph::Tensor_attributes() - .set_name("Offset") - .set_dim({1, 1, 1, 1}) - .set_stride({1, 1, 1, 1}) - .set_data_type(fe::DataType_t::INT64)); - sdpa_backward_options.set_dropout(dropout_probability, dropout_seed, dropout_offset); - } - - if (is_softmax_offset) { - softmax_offset = mha_graph->tensor(fe::graph::Tensor_attributes() - .set_name("softmax_offset") - .set_dim({1, h, 1, 1}) - .set_stride({h, 1, 1, 1}) - .set_data_type(fe::DataType_t::FLOAT)); - sdpa_backward_options.set_sink_token(softmax_offset); - d_softmax_offset = mha_graph->tensor(fe::graph::Tensor_attributes() - .set_name("d_softmax_offset") - .set_dim({1, h, 1, 1}) - .set_stride({h, 1, 1, 1}) - .set_data_type(fe::DataType_t::FLOAT)); - sdpa_backward_options.set_dsink_token(d_softmax_offset); - } - - auto [dQ, dK, dV] = mha_graph->sdpa_backward(q, k, v, o, dO, stats, sdpa_backward_options); - - dQ->set_output(true).set_dim({b, h, s_q, d_qk}).set_stride(q_stride); - dK->set_output(true).set_dim({b, hg, s_kv, d_qk}).set_stride(k_stride); - dV->set_output(true).set_dim({b, hg, s_kv, d_v}).set_stride(v_stride); - if (is_ragged_q) { - dQ->set_ragged_offset(offset_q); - } - if (is_ragged_kv) { - dK->set_ragged_offset(offset_k); - dV->set_ragged_offset(offset_v); - } - - std::tuple, // q - std::shared_ptr, // k - std::shared_ptr, // v - std::shared_ptr, // o - std::shared_ptr, // dO - std::shared_ptr, // stats - std::shared_ptr, // attn_scale - std::shared_ptr, // dQ - std::shared_ptr, // dK - std::shared_ptr> // dV - key_tensors_tuple = std::make_tuple(q, k, v, o, dO, stats, attn_scale, dQ, dK, dV); - auto bias_tuple = is_bias ? std::make_tuple(bias, dBias) : std::make_tuple(nullptr, nullptr); - auto softmax_offset_tuple = is_softmax_offset ? std::make_tuple(softmax_offset, d_softmax_offset) - : std::make_tuple(nullptr, nullptr); - auto padding_tuple = - is_padding ? std::make_tuple(seq_q, seq_kv) : std::make_tuple(nullptr, nullptr); - auto offset_qo_tuple = - is_ragged_q ? std::make_tuple(offset_q, offset_o) : std::make_tuple(nullptr, nullptr); - auto offset_kv_tuple = - is_ragged_kv ? std::make_tuple(offset_k, offset_v) : std::make_tuple(nullptr, nullptr); - auto offset_s_tuple = use_ragged_stats ? std::make_tuple(offset_stats) : std::make_tuple(nullptr); - auto dropout_tuple = is_dropout ? std::make_tuple(dropout_seed, dropout_offset) - : std::make_tuple(nullptr, nullptr); - - return std::tuple_cat(std::make_tuple(mha_graph), key_tensors_tuple, bias_tuple, - softmax_offset_tuple, padding_tuple, offset_qo_tuple, offset_kv_tuple, - offset_s_tuple, dropout_tuple); -} - -void fused_attn_arbitrary_seqlen_bwd_impl( - const FusedAttnConfig &cfg, void *devPtrQ, void *devPtrKTranspose, void *devPtrVTranspose, - void *devPtrO, void *devPtrSoftmaxStats, void *devPtrBias, void *devPtrSoftmaxOffset, - void *devPtrdQ, void *devPtrdK, void *devPtrdV, void *devPtrdO, void *devPtrdBias, - void *devPtrdSoftmaxOffset, void *devPtrDropoutSeed, void *devPtrDropoutOffset, - void *devPtrCuSeqlensQ, void *devPtrCuSeqlensKV, void *devPtrSeqOffsetsQ, - void *devPtrSeqOffsetsKV, void *workspace, size_t *workspace_size, cudaStream_t stream, - cudnnHandle_t handle) { - using namespace transformer_engine; - - cfg.check_derived(); - const int64_t b = static_cast(cfg.graph_batch_size_bwd); - const DType ragged_offset_type = cfg.ragged_offset_type_bwd; - const int64_t actual_b = static_cast(cfg.batch_size); - const bool use_ragged_stats = cfg.uses_ragged_stats; - - float scaling_factor = cfg.attn_scale; - const bool is_bias = cfg.is_bias; - const bool is_padding = cfg.is_padding; - const bool is_softmax_offset = cfg.is_softmax_offset; - const bool is_dropout = cfg.is_dropout; - const bool is_ragged_q = cfg.is_ragged_q; - const bool is_ragged_kv = cfg.is_ragged_kv; - - try { - auto cache_entry = get_graph(cfg, handle); - auto [mha_graph, q, k, v, o, dO, stats, attn_scale, dQ, dK, dV, bias, dBias, softmax_offset, - d_softmax_offset, seq_q, seq_kv, offset_q, offset_o, offset_k, offset_v, offset_stats, - dropout_seed, dropout_offset] = cache_entry->graph_and_tensors; - - // This graph is going to be used, so finish the build the cache deferred. - build_plans(Backend::F16, Pass::Bwd, *cache_entry); - - // Exit to request upper level API to allocate memory if needed - // n.b. Care should be taken to align each of the added worksapce tensors to their type. - // We do this by adding padding at the end of each separate allocation. - auto plan_workspace_size = alignTo<16>(mha_graph->get_workspace_size()); - const size_t num_bytes_per_seqlen = alignTo<16>(b * sizeof(int32_t)); - const size_t actual_seqlen_workspace_size = is_padding ? 2 * num_bytes_per_seqlen : 0; - const size_t num_bytes_per_ragged_offset = - alignTo<16>(((b + 1) * typeToNumBits(ragged_offset_type)) / 8); - size_t seqlen_offsets_workspace_size = 0; - if (is_ragged_q || is_ragged_kv) { - size_t count = 2 * (static_cast(is_ragged_q) + static_cast(is_ragged_kv)); - if (use_ragged_stats) { - seqlen_offsets_workspace_size = (count + 1) * num_bytes_per_ragged_offset; - } else { - seqlen_offsets_workspace_size = count * num_bytes_per_ragged_offset; - } - } - if (workspace == nullptr) { - *workspace_size = - plan_workspace_size + actual_seqlen_workspace_size + seqlen_offsets_workspace_size; - return; - } - // cuDNN stream check needs to be moved here to support dummy kernel calls with - // null streams for sizing the cuDNN workspace. - NVTE_CHECK_CUDNN(cudnnSetStream(handle, stream)); - - // build variant pack - std::unordered_map, void *> variant_pack = { - {q, devPtrQ}, - {k, devPtrKTranspose}, - {v, devPtrVTranspose}, - {o, devPtrO}, - {dO, devPtrdO}, - {stats, devPtrSoftmaxStats}, - {attn_scale, &scaling_factor}, - {dQ, devPtrdQ}, - {dK, devPtrdK}, - {dV, devPtrdV}, - }; - - if (is_bias) { - variant_pack[bias] = devPtrBias; - if (dBias != nullptr) { - variant_pack[dBias] = devPtrdBias; - } - } - - if (is_padding) { - constexpr size_t nthreads_per_block = 128; - const size_t grid = (b + nthreads_per_block - 1) / nthreads_per_block; - void *devActualSeqlenQ = static_cast(workspace) + plan_workspace_size; - void *devActualSeqlenKV = static_cast(devActualSeqlenQ) + num_bytes_per_seqlen; - cu_seqlens_to_actual_seqlens<<>>( - actual_b, b, static_cast(devPtrCuSeqlensQ), - static_cast(devPtrCuSeqlensKV), static_cast(devActualSeqlenQ), - static_cast(devActualSeqlenKV)); - NVTE_CHECK_CUDA(cudaGetLastError()); - variant_pack[seq_q] = devActualSeqlenQ; - variant_pack[seq_kv] = devActualSeqlenKV; - } - - if (is_ragged_q || is_ragged_kv) { - constexpr size_t nthreads_per_block = 128; - const size_t grid = (b + nthreads_per_block) / nthreads_per_block; - void *devOffsets = - static_cast(workspace) + plan_workspace_size + actual_seqlen_workspace_size; - void *devOffsetsQ = nullptr; - void *devOffsetsO = nullptr; - if (is_ragged_q) { - devOffsetsQ = devOffsets; - devOffsetsO = static_cast(devOffsetsQ) + num_bytes_per_ragged_offset; - } - void *devOffsetsK = nullptr; - void *devOffsetsV = nullptr; - if (is_ragged_kv) { - devOffsetsK = static_cast(devOffsets) + - static_cast(is_ragged_q) * 2 * num_bytes_per_ragged_offset; - devOffsetsV = static_cast(devOffsetsK) + num_bytes_per_ragged_offset; - } - void *devOffsetsS = nullptr; - if (use_ragged_stats) { - devOffsetsS = static_cast(devOffsets) + - (static_cast(is_ragged_q) + static_cast(is_ragged_kv)) * 2 * - num_bytes_per_ragged_offset; - } - cu_seqlens_padded_to_offsets<<>>( - cfg.ragged_offset_mults, actual_b, b, static_cast(devPtrSeqOffsetsQ), - static_cast(devPtrSeqOffsetsKV), ragged_offset_type, devOffsetsQ, devOffsetsK, - devOffsetsV, devOffsetsO, devOffsetsS); - NVTE_CHECK_CUDA(cudaGetLastError()); - if (is_ragged_q) { - variant_pack[offset_q] = devOffsetsQ; - variant_pack[offset_o] = devOffsetsO; - } - if (is_ragged_kv) { - variant_pack[offset_k] = devOffsetsK; - variant_pack[offset_v] = devOffsetsV; - } - if (use_ragged_stats) { - variant_pack[offset_stats] = devOffsetsS; - } - } - - if (is_dropout) { - variant_pack[dropout_seed] = devPtrDropoutSeed; - variant_pack[dropout_offset] = devPtrDropoutOffset; - } - - if (is_softmax_offset) { - variant_pack[softmax_offset] = devPtrSoftmaxOffset; - variant_pack[d_softmax_offset] = devPtrdSoftmaxOffset; - } - - NVTE_CHECK_CUDNN_FE(mha_graph->execute(handle, variant_pack, workspace)); - graph_cache_debug::record_execute(Backend::F16, Pass::Bwd); - } catch (cudnn_frontend::cudnnException &e) { - NVTE_ERROR(e.what()); - } -} -} // namespace fused_attn - -using namespace transformer_engine::fused_attn; -void fused_attn_arbitrary_seqlen_fwd(const FusedAttnConfig &cfg, const Tensor *input_Q, - const Tensor *input_K, const Tensor *input_V, - const Tensor *input_Bias, const Tensor *input_SoftmaxOffset, - Tensor *output_O, NVTETensorPack *Aux_CTX_Tensors, - const Tensor *cu_seqlens_q, const Tensor *cu_seqlens_kv, - const Tensor *cu_seqlens_q_padded, - const Tensor *cu_seqlens_kv_padded, const Tensor *page_table_k, - const Tensor *page_table_v, const Tensor *rng_state, - Tensor *workspace, cudaStream_t stream, cudnnHandle_t handle) { - using namespace transformer_engine; - - const size_t batch = cfg.batch_size; - const size_t num_attn_heads = cfg.num_attn_heads; - const size_t max_seqlen_q = cfg.max_seqlen_q; - const size_t num_tokens_q = cfg.num_tokens_q; - const bool return_max_logit = cfg.return_max_logit; - const NVTE_Bias_Type bias_type = cfg.bias_type; - const NVTE_Softmax_Type softmax_type = cfg.softmax_type; - - const auto QKV_type = input_Q->data.dtype; - void *devPtrQ = input_Q->data.dptr; - void *devPtrK = input_K->data.dptr; - void *devPtrV = input_V->data.dptr; - void *devPtrO = output_O->data.dptr; - void *devPtrS1 = nullptr; - void *devPtrS2 = nullptr; - void *devPtrBias = nullptr; - if ((bias_type != NVTE_Bias_Type::NVTE_NO_BIAS) && (bias_type != NVTE_Bias_Type::NVTE_ALIBI)) { - devPtrBias = input_Bias->data.dptr; - } - void *devPtrSoftmaxOffset = nullptr; - if (softmax_type != NVTE_VANILLA_SOFTMAX) { - devPtrSoftmaxOffset = input_SoftmaxOffset->data.dptr; - } - - void *devPtrCuSeqlensQ = cu_seqlens_q->data.dptr; - void *devPtrCuSeqlensKV = cu_seqlens_kv->data.dptr; - void *devPtrSeqOffsetsQ = cu_seqlens_q_padded->data.dptr; - void *devPtrSeqOffsetsKV = cu_seqlens_kv_padded->data.dptr; - void *devPtrPageTableK = page_table_k ? page_table_k->data.dptr : nullptr; - void *devPtrPageTableV = page_table_v ? page_table_v->data.dptr : nullptr; - - size_t i = 0; - if (Aux_CTX_Tensors->size == 0) { - const bool use_ragged_stats = cfg.uses_ragged_stats; - - Tensor *output_S = convertNVTETensorCheck(Aux_CTX_Tensors->tensors[i++]); - output_S->data.dptr = nullptr; - if (use_ragged_stats) { - output_S->data.shape = {num_tokens_q, num_attn_heads, 1}; - } else { - output_S->data.shape = {batch, num_attn_heads, max_seqlen_q, 1}; - } - output_S->data.dtype = DType::kFloat32; - - if (return_max_logit) { - Tensor *output_Max = convertNVTETensorCheck(Aux_CTX_Tensors->tensors[i++]); - output_Max->data.dptr = nullptr; - if (use_ragged_stats) { - output_Max->data.shape = {num_tokens_q, num_attn_heads, 1}; - } else { - output_Max->data.shape = {batch, num_attn_heads, max_seqlen_q, 1}; - } - output_Max->data.dtype = DType::kFloat32; - } - - Tensor *output_rng_state = convertNVTETensorCheck(Aux_CTX_Tensors->tensors[i++]); - output_rng_state->data.dptr = nullptr; - output_rng_state->data.shape = {2}; - output_rng_state->data.dtype = DType::kInt64; - - if ((bias_type != NVTE_NO_BIAS) && (bias_type != NVTE_ALIBI)) { - Tensor *output_bias = convertNVTETensorCheck(Aux_CTX_Tensors->tensors[i++]); - output_bias->data.dptr = nullptr; - output_bias->data.shape = {cfg.bias_batch_size, cfg.bias_num_heads, cfg.bias_seqlen_q, - cfg.bias_seqlen_kv}; - output_bias->data.dtype = QKV_type; - } - - if (softmax_type != NVTE_VANILLA_SOFTMAX) { - Tensor *output_softmax_offset = convertNVTETensorCheck(Aux_CTX_Tensors->tensors[i++]); - output_softmax_offset->data.dptr = nullptr; - output_softmax_offset->data.shape = {1, num_attn_heads, 1, 1}; - output_softmax_offset->data.dtype = DType::kFloat32; - } - - Aux_CTX_Tensors->size = i; - } else if (Aux_CTX_Tensors->size >= 2) { - Tensor *output_S = convertNVTETensorCheck(Aux_CTX_Tensors->tensors[i++]); - devPtrS1 = output_S->data.dptr; - - if (return_max_logit) { - Tensor *output_Max = convertNVTETensorCheck(Aux_CTX_Tensors->tensors[i++]); - devPtrS2 = output_Max->data.dptr; - } - Tensor *output_rng_state = convertNVTETensorCheck(Aux_CTX_Tensors->tensors[i++]); - output_rng_state->data.dptr = rng_state->data.dptr; - if ((bias_type != NVTE_NO_BIAS) && (bias_type != NVTE_ALIBI)) { - Tensor *output_bias = convertNVTETensorCheck(Aux_CTX_Tensors->tensors[i++]); - output_bias->data.dptr = devPtrBias; - } - if (softmax_type != NVTE_VANILLA_SOFTMAX) { - Tensor *output_softmax_offset = convertNVTETensorCheck(Aux_CTX_Tensors->tensors[i++]); - output_softmax_offset->data.dptr = devPtrSoftmaxOffset; - } - } else { - NVTE_ERROR("Unexpected Aux_CTX_Tensors->size."); - } - - void *devPtrDropoutSeed = rng_state->data.dptr; - void *devPtrDropoutOffset = - reinterpret_cast(reinterpret_cast(rng_state->data.dptr) + 1); - - size_t workspace_size = 0; - - fused_attn_arbitrary_seqlen_fwd_impl( - cfg, devPtrQ, devPtrK, devPtrV, devPtrBias, devPtrSoftmaxOffset, devPtrS1, devPtrS2, devPtrO, - devPtrDropoutSeed, devPtrDropoutOffset, devPtrCuSeqlensQ, devPtrCuSeqlensKV, devPtrPageTableK, - devPtrPageTableV, devPtrSeqOffsetsQ, devPtrSeqOffsetsKV, workspace->data.dptr, - &workspace_size, stream, handle); - - if (workspace_size > 0) { - if (workspace->data.dptr == nullptr) { - workspace->data.shape = {workspace_size}; - workspace->data.dtype = DType::kByte; - return; - } - } else if (workspace_size == 0) { - workspace->data.shape = {1}; - workspace->data.dtype = DType::kByte; - return; - } else { - NVTE_ERROR("Unexpected workspace_size."); - } -} - -void fused_attn_arbitrary_seqlen_bwd(const FusedAttnConfig &cfg, const Tensor *input_Q, - const Tensor *input_K, const Tensor *input_V, - const Tensor *input_O, const Tensor *input_dO, - const Tensor *input_Bias, const Tensor *input_SoftmaxOffset, - Tensor *output_S, Tensor *output_dQ, Tensor *output_dK, - Tensor *output_dV, Tensor *output_dBias, - Tensor *output_dSoftmaxOffset, const Tensor *cu_seqlens_q, - const Tensor *cu_seqlens_kv, const Tensor *cu_seqlens_q_padded, - const Tensor *cu_seqlens_kv_padded, const Tensor *rng_state, - Tensor *workspace, cudaStream_t stream, cudnnHandle_t handle) { - using namespace transformer_engine; - - const NVTE_Bias_Type bias_type = cfg.bias_type; - const NVTE_Softmax_Type softmax_type = cfg.softmax_type; - - void *devPtrQ = input_Q->data.dptr; - void *devPtrK = input_K->data.dptr; - void *devPtrV = input_V->data.dptr; - void *devPtrO = input_O->data.dptr; - void *devPtrdO = input_dO->data.dptr; - void *devPtrBias = nullptr; - void *devPtrdBias = nullptr; - if ((bias_type != NVTE_Bias_Type::NVTE_NO_BIAS) && (bias_type != NVTE_Bias_Type::NVTE_ALIBI)) { - devPtrBias = input_Bias->data.dptr; - devPtrdBias = output_dBias->data.dptr; - } - - void *devPtrdQ = output_dQ->data.dptr; - void *devPtrdK = output_dK->data.dptr; - void *devPtrdV = output_dV->data.dptr; - void *devPtrSoftmaxStats = nullptr; - devPtrSoftmaxStats = output_S->data.dptr; - void *devPtrSoftmaxOffset = nullptr; - void *devPtrdSoftmaxOffset = nullptr; - if (softmax_type != NVTE_VANILLA_SOFTMAX) { - devPtrSoftmaxOffset = input_SoftmaxOffset->data.dptr; - devPtrdSoftmaxOffset = output_dSoftmaxOffset->data.dptr; - } - - void *devPtrCuSeqlensQ = cu_seqlens_q->data.dptr; - void *devPtrCuSeqlensKV = cu_seqlens_kv->data.dptr; - void *devPtrSeqOffsetsQ = cu_seqlens_q_padded->data.dptr; - void *devPtrSeqOffsetsKV = cu_seqlens_kv_padded->data.dptr; - - void *devPtrDropoutSeed = rng_state->data.dptr; - void *devPtrDropoutOffset = - reinterpret_cast(reinterpret_cast(rng_state->data.dptr) + 1); - - size_t workspace_size = 0; - - fused_attn_arbitrary_seqlen_bwd_impl( - cfg, devPtrQ, devPtrK, devPtrV, devPtrO, devPtrSoftmaxStats, devPtrBias, devPtrSoftmaxOffset, - devPtrdQ, devPtrdK, devPtrdV, devPtrdO, devPtrdBias, devPtrdSoftmaxOffset, devPtrDropoutSeed, - devPtrDropoutOffset, devPtrCuSeqlensQ, devPtrCuSeqlensKV, devPtrSeqOffsetsQ, - devPtrSeqOffsetsKV, workspace->data.dptr, &workspace_size, stream, handle); - - if (workspace_size > 0) { - if (workspace->data.dptr == nullptr) { - workspace->data.shape = {workspace_size}; - workspace->data.dtype = DType::kByte; - return; - } - } else if (workspace_size == 0) { - workspace->data.shape = {1}; - workspace->data.dtype = DType::kByte; - return; - } else { - NVTE_ERROR("Unexpected workspace_size."); - } -} - -// Check whether cuDNN can support a given config, per forward/backward pass. -std::string support_verdict_f16(const FusedAttnConfig &cfg, Pass pass, cudnnHandle_t handle) { - if (pass == Pass::Fwd) { - return fused_attn::support_verdict(cfg, handle); - } - return fused_attn::support_verdict(cfg, handle); -} - -} // namespace transformer_engine diff --git a/transformer_engine/common/fused_attn/fused_attn_f16_arbitrary_seqlen.h b/transformer_engine/common/fused_attn/fused_attn_f16_arbitrary_seqlen.h deleted file mode 100644 index 3f196538c4d..00000000000 --- a/transformer_engine/common/fused_attn/fused_attn_f16_arbitrary_seqlen.h +++ /dev/null @@ -1,49 +0,0 @@ -/************************************************************************* - * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. - * - * See LICENSE for license information. - ************************************************************************/ - -/*! \file fused_attn_arbitrary_seqlen.h - * \brief Functions for fused attention with seqlen > 512 - */ - -#ifndef TRANSFORMER_ENGINE_COMMON_FUSED_ATTN_FUSED_ATTN_F16_ARBITRARY_SEQLEN_H_ -#define TRANSFORMER_ENGINE_COMMON_FUSED_ATTN_FUSED_ATTN_F16_ARBITRARY_SEQLEN_H_ - -#include - -#include - -#include "common/common.h" -#include "config_and_params.h" -#include "transformer_engine/transformer_engine.h" - -namespace transformer_engine { -void fused_attn_arbitrary_seqlen_fwd(const fused_attn::FusedAttnConfig &cfg, const Tensor *input_Q, - const Tensor *input_K, const Tensor *input_V, - const Tensor *input_Bias, const Tensor *input_SoftmaxOffset, - Tensor *output_O, NVTETensorPack *Aux_CTX_Tensors, - const Tensor *cu_seqlens_q, const Tensor *cu_seqlens_kv, - const Tensor *cu_seqlens_q_padded, - const Tensor *cu_seqlens_kv_padded, const Tensor *page_table_k, - const Tensor *page_table_v, const Tensor *rng_state, - Tensor *workspace, cudaStream_t stream, cudnnHandle_t handle); - -void fused_attn_arbitrary_seqlen_bwd(const fused_attn::FusedAttnConfig &cfg, const Tensor *input_Q, - const Tensor *input_K, const Tensor *input_V, - const Tensor *input_O, const Tensor *input_dO, - const Tensor *input_Bias, const Tensor *input_SoftmaxOffset, - Tensor *output_S, Tensor *output_dQ, Tensor *output_dK, - Tensor *output_dV, Tensor *output_dBias, - Tensor *output_dSoftmaxOffset, const Tensor *cu_seqlens_q, - const Tensor *cu_seqlens_kv, const Tensor *cu_seqlens_q_padded, - const Tensor *cu_seqlens_kv_padded, const Tensor *rng_state, - Tensor *workspace, cudaStream_t stream, cudnnHandle_t handle); - -std::string support_verdict_f16(const fused_attn::FusedAttnConfig &cfg, fused_attn::Pass pass, - cudnnHandle_t handle); - -} // namespace transformer_engine - -#endif // TRANSFORMER_ENGINE_COMMON_FUSED_ATTN_FUSED_ATTN_F16_ARBITRARY_SEQLEN_H_ diff --git a/transformer_engine/common/fused_attn/fused_attn_fp8.cu b/transformer_engine/common/fused_attn/fused_attn_fp8.cu deleted file mode 100644 index 8a346408e20..00000000000 --- a/transformer_engine/common/fused_attn/fused_attn_fp8.cu +++ /dev/null @@ -1,1525 +0,0 @@ -/************************************************************************* - * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. - * - * See LICENSE for license information. - ************************************************************************/ - -#include - -#include "../common.h" -#include "../cudnn_utils.h" -#include "../util/system.h" -#include "fused_attn_fp8.h" -#include "graph_cache.h" -#include "graph_cache_debug.h" -#include "utils.h" - -namespace transformer_engine { -namespace fused_attn { - -using namespace transformer_engine; -namespace fe = cudnn_frontend; - -using Fp8FwdGraphAndTensors = - std::tuple, - std::shared_ptr, // Q - std::shared_ptr, // K - std::shared_ptr, // V - std::shared_ptr, // descale_q - std::shared_ptr, // descale_k - std::shared_ptr, // descale_v - std::shared_ptr, // descale_s - std::shared_ptr, // scale_s - std::shared_ptr, // scale_o - std::shared_ptr, // attn_scale - std::shared_ptr, // O - std::shared_ptr, // amax_s - std::shared_ptr, // amax_o - std::shared_ptr, // Stats - std::shared_ptr, // bias - std::shared_ptr, // softmax_offset - std::shared_ptr, // seq_q - std::shared_ptr, // seq_kv - std::shared_ptr, // offset_q - std::shared_ptr, // offset_o - std::shared_ptr, // offset_k - std::shared_ptr, // offset_v - std::shared_ptr, // offset_stats - std::shared_ptr, // dropout_seed - std::shared_ptr>; // dropout_offset - -static Fp8FwdGraphAndTensors create_graph_fp8_fwd(const FusedAttnConfig& cfg) { - const auto cudnn_runtime_version = cudnnGetVersion(); - const cudnn_frontend::DataType_t qkv_tensor_type = - get_cudnn_fe_dtype(static_cast(cfg.qkv_dtype)); - const cudnn_frontend::DataType_t o_tensor_type = - get_cudnn_fe_dtype(static_cast(cfg.o_dtype)); - const int64_t b = static_cast(cfg.graph_batch_size_fwd); - const int64_t h = static_cast(cfg.num_attn_heads); - const int64_t hg = static_cast(cfg.num_gqa_groups); - const int64_t s_q = static_cast(cfg.graph_max_seqlen_q); - const int64_t s_kv = static_cast(cfg.graph_max_seqlen_kv); - const int64_t d_qk = static_cast(cfg.head_dim_qk); - const int64_t d_v = static_cast(cfg.head_dim_v); - const int64_t window_size_left = cfg.window_size_left; - const int64_t window_size_right = cfg.window_size_right; - const float dropout_probability = cfg.dropout; - const NVTE_QKV_Layout qkv_layout = cfg.qkv_layout; - const NVTE_QKV_Format o_format = cfg.o_format; - const NVTE_QKV_Format qkv_scale_inv_format = cfg.qkv_scale_inv_format; - const bool bottom_right_diagonal = cfg.bottom_right_diagonal; - const bool is_bias = cfg.is_bias; - const bool is_causal = cfg.is_causal; - const bool is_causal_bottom_right = cfg.is_causal_bottom_right; - const bool is_padding = cfg.is_padding; - const bool is_dropout = cfg.is_dropout; - const bool is_softmax_offset = cfg.is_softmax_offset; - const bool is_mxfp8 = cfg.is_mxfp8; - const bool is_tensor_scaling = cfg.is_tensor_scaling; - const bool is_delayed_scaling = cfg.is_delayed_scaling_fwd; - const bool is_current_scaling = cfg.is_current_scaling_fwd; - const bool use_cu_seqlens_directly = cfg.uses_cu_seqlens_directly; - const bool is_ragged_q = cfg.is_ragged_q; - const bool is_ragged_kv = cfg.is_ragged_kv; - const bool use_ragged_stats = cfg.uses_ragged_stats; - const DType ragged_offset_type = cfg.ragged_offset_type_fwd; - const RaggedOffsetMultipliers offset_mults = cfg.ragged_offset_mults; - - auto mha_graph = std::make_shared(); - mha_graph->set_io_data_type(qkv_tensor_type) - .set_intermediate_data_type(fe::DataType_t::FLOAT) - .set_compute_data_type(fe::DataType_t::FLOAT); - - std::shared_ptr Q, K, V, attn_scale; - std::shared_ptr descale_q, descale_k, descale_v; - std::shared_ptr descale_s, scale_s, scale_o; - std::shared_ptr bias, softmax_offset, seq_q, seq_kv; - std::shared_ptr offset_q, offset_k, offset_v, offset_o, - offset_stats; - std::shared_ptr dropout_seed, dropout_offset; - - // Q, K, V, attn_scale - std::vector q_strides(4), k_strides(4), v_strides(4); - generateMatrixStridesWithLayout(b, h, hg, s_q, s_kv, d_qk, d_v, q_strides.data(), - k_strides.data(), v_strides.data(), qkv_layout); - Q = mha_graph->tensor(fe::graph::Tensor_attributes() - .set_name("Q") - .set_dim({b, h, s_q, d_qk}) - .set_stride(q_strides) - .set_data_type(qkv_tensor_type)); - if (is_ragged_q) { - offset_q = mha_graph->tensor(fe::graph::Tensor_attributes() - .set_name("offset_q") - .set_dim({b + 1, 1, 1, 1}) - .set_stride({1, 1, 1, 1}) - .set_data_type(get_cudnn_fe_dtype(ragged_offset_type))); - Q->set_ragged_offset(offset_q); - if (use_cu_seqlens_directly) { - Q->set_ragged_offset_multiplier(offset_mults.q); - } - } - K = mha_graph->tensor(fe::graph::Tensor_attributes() - .set_name("K") - .set_dim({b, hg, s_kv, d_qk}) - .set_stride(k_strides) - .set_data_type(qkv_tensor_type)); - V = mha_graph->tensor(fe::graph::Tensor_attributes() - .set_name("V") - .set_dim({b, hg, s_kv, d_v}) - .set_stride(v_strides) - .set_data_type(qkv_tensor_type)); - if (is_ragged_kv) { - offset_k = mha_graph->tensor(fe::graph::Tensor_attributes() - .set_name("offset_k") - .set_dim({b + 1, 1, 1, 1}) - .set_stride({1, 1, 1, 1}) - .set_data_type(get_cudnn_fe_dtype(ragged_offset_type))); - offset_v = mha_graph->tensor(fe::graph::Tensor_attributes() - .set_name("offset_v") - .set_dim({b + 1, 1, 1, 1}) - .set_stride({1, 1, 1, 1}) - .set_data_type(get_cudnn_fe_dtype(ragged_offset_type))); - K->set_ragged_offset(offset_k); - V->set_ragged_offset(offset_v); - if (use_cu_seqlens_directly) { - K->set_ragged_offset_multiplier(offset_mults.k); - V->set_ragged_offset_multiplier(offset_mults.v); - } - } - attn_scale = mha_graph->tensor(fe::graph::Tensor_attributes() - .set_name("attn_scale") - .set_dim({1, 1, 1, 1}) - .set_stride({1, 1, 1, 1}) - .set_is_pass_by_value(true) - .set_data_type(fe::DataType_t::FLOAT)); - - // Descale_q, Descale_k, Descale_v, Descale_s, Scale_s, Scale_o - if (is_tensor_scaling) { - descale_q = mha_graph->tensor(fe::graph::Tensor_attributes() - .set_name("Descale_q") - .set_dim({1, 1, 1, 1}) - .set_stride({1, 1, 1, 1}) - .set_data_type(fe::DataType_t::FLOAT)); - descale_k = mha_graph->tensor_like(descale_q, "Descale_q"); - descale_v = mha_graph->tensor_like(descale_q, "Descale_v"); - descale_s = mha_graph->tensor_like(descale_q, "Descale_s"); - scale_s = mha_graph->tensor_like(descale_q, "Scale_s"); - if (is_delayed_scaling) { - scale_o = mha_graph->tensor_like(descale_q, "Scale_o"); - } - if (is_current_scaling) { - scale_o = mha_graph->tensor(1.0f); - } - } else if (is_mxfp8) { - const NVTE_QKV_Format q_scale_inv_format = - (qkv_scale_inv_format != NVTE_QKV_Format_NOT_SET) ? qkv_scale_inv_format : cfg.q_format; - const NVTE_QKV_Format kv_scale_inv_format = - (qkv_scale_inv_format != NVTE_QKV_Format_NOT_SET) ? qkv_scale_inv_format : cfg.kv_format; - std::vector q_scale_strides(4); - std::vector k_scale_strides(4); - std::vector v_scale_strides(4); - auto padded = pad_s_d_for_mxfp8(s_q, s_kv, d_qk, d_v); - generateMatrixStridesWithFormat(b, h, padded.s_q_padded, padded.d_qk_scale_padded, - q_scale_strides.data(), q_scale_inv_format); - generateMatrixStridesWithFormat(b, hg, padded.s_kv_padded, padded.d_qk_scale_padded, - k_scale_strides.data(), kv_scale_inv_format); - generateMatrixStridesWithFormat(b, hg, padded.s_kv_scale_padded, padded.d_v_padded, - v_scale_strides.data(), kv_scale_inv_format); - descale_q = mha_graph->tensor(fe::graph::Tensor_attributes() - .set_name("Descale_q") - .set_dim({b, h, padded.s_q_padded, padded.d_qk_scale_padded}) - .set_stride(q_scale_strides) - .set_data_type(fe::DataType_t::FP8_E8M0) - .set_reordering_type(fe::TensorReordering_t::F8_128x4)); - descale_k = - mha_graph->tensor(fe::graph::Tensor_attributes() - .set_name("Descale_k") - .set_dim({b, hg, padded.s_kv_padded, padded.d_qk_scale_padded}) - .set_stride(k_scale_strides) - .set_data_type(fe::DataType_t::FP8_E8M0) - .set_reordering_type(fe::TensorReordering_t::F8_128x4)); - descale_v = mha_graph->tensor(fe::graph::Tensor_attributes() - .set_name("Descale_v") - .set_dim({b, hg, padded.s_kv_scale_padded, padded.d_v_padded}) - .set_stride(v_scale_strides) - .set_data_type(fe::DataType_t::FP8_E8M0) - .set_reordering_type(fe::TensorReordering_t::F8_128x4)); - } - - fe::graph::SDPA_fp8_attributes sdpa_options; - sdpa_options = fe::graph::SDPA_fp8_attributes() - .set_name("sdpa_fp8") - .set_generate_stats(true) - .set_causal_mask(is_causal) - .set_attn_scale(attn_scale); - - fe::DiagonalAlignment_t const& diagonal_alignment = bottom_right_diagonal - ? fe::DiagonalAlignment_t::BOTTOM_RIGHT - : fe::DiagonalAlignment_t::TOP_LEFT; - sdpa_options.set_diagonal_alignment(diagonal_alignment); - - if (cudnn_runtime_version >= 92100) { - if (window_size_left != -1) { - sdpa_options.set_diagonal_band_left_bound(window_size_left + 1); - } - if (window_size_right != -1) { - sdpa_options.set_diagonal_band_right_bound(window_size_right); - } - } - if (is_causal_bottom_right) { - sdpa_options.set_diagonal_band_right_bound(0); - } - - // sdpa_options.set_alibi_mask(is_alibi); - // if (is_bias) { - // bias = mha_graph->tensor(fe::graph::Tensor_attributes() - // .set_name("bias") - // .set_dim({bias_b, bias_h, bias_sq, bias_skv}) - // .set_stride({bias_h * bias_sq * bias_skv, bias_sq * bias_skv, bias_skv, 1})); - // sdpa_options.set_bias(bias); - // } - - if (is_padding) { - if (use_cu_seqlens_directly) { - // seq_q/seq_kv keep their tuple slots but hold (b+1)-shaped cu_seqlen tensors. - seq_q = mha_graph->tensor(fe::graph::Tensor_attributes() - .set_name("cu_seq_len_q") - .set_dim({b + 1, 1, 1, 1}) - .set_stride({1, 1, 1, 1}) - .set_data_type(fe::DataType_t::INT32)); - seq_kv = mha_graph->tensor(fe::graph::Tensor_attributes() - .set_name("cu_seq_len_kv") - .set_dim({b + 1, 1, 1, 1}) - .set_stride({1, 1, 1, 1}) - .set_data_type(fe::DataType_t::INT32)); - sdpa_options.set_padding_mask(is_padding).set_cu_seq_len_q(seq_q).set_cu_seq_len_kv(seq_kv); - // cu_seq_len (and the ragged offset multiplier) are unified-engine-only. - // Pin the implementation so an unsupported config fails with the unified - // engine's specific error instead of auto-selection's generic failure. - sdpa_options.set_implementation(fe::AttentionImplementation_t::UNIFIED); - } else { - seq_q = mha_graph->tensor(fe::graph::Tensor_attributes() - .set_name("seq_q") - .set_dim({b, 1, 1, 1}) - .set_stride({1, 1, 1, 1}) - .set_data_type(fe::DataType_t::INT32)); - seq_kv = mha_graph->tensor(fe::graph::Tensor_attributes() - .set_name("seq_kv") - .set_dim({b, 1, 1, 1}) - .set_stride({1, 1, 1, 1}) - .set_data_type(fe::DataType_t::INT32)); - sdpa_options.set_padding_mask(is_padding).set_seq_len_q(seq_q).set_seq_len_kv(seq_kv); - } - } - - if (is_dropout) { - dropout_seed = mha_graph->tensor(fe::graph::Tensor_attributes() - .set_name("Seed") - .set_dim({1, 1, 1, 1}) - .set_stride({1, 1, 1, 1}) - .set_data_type(fe::DataType_t::INT64)); - dropout_offset = mha_graph->tensor(fe::graph::Tensor_attributes() - .set_name("Offset") - .set_dim({1, 1, 1, 1}) - .set_stride({1, 1, 1, 1}) - .set_data_type(fe::DataType_t::INT64)); - sdpa_options.set_dropout(dropout_probability, dropout_seed, dropout_offset); - } - - if (is_softmax_offset) { - softmax_offset = mha_graph->tensor(fe::graph::Tensor_attributes() - .set_name("softmax_offset") - .set_dim({1, h, 1, 1}) - .set_stride({h, 1, 1, 1}) - .set_data_type(fe::DataType_t::FLOAT)); - sdpa_options.set_sink_token(softmax_offset); - } - - std::shared_ptr O, Stats, amax_s, amax_o; - if (is_tensor_scaling) { - auto outputs = mha_graph->sdpa_fp8(Q, K, V, descale_q, descale_k, descale_v, descale_s, scale_s, - scale_o, sdpa_options); - O = outputs[0]; - Stats = outputs[1]; - amax_s = outputs[2]; - amax_o = outputs[3]; - amax_s->set_output(true) - .set_dim({1, 1, 1, 1}) - .set_stride({1, 1, 1, 1}) - .set_data_type(fe::DataType_t::FLOAT); - } else if (is_mxfp8) { - auto outputs = mha_graph->sdpa_fp8(Q, K, V, descale_q, descale_k, descale_v, sdpa_options); - O = outputs[0]; - Stats = outputs[1]; - amax_o = outputs[2]; - } - - std::vector o_strides(4); - generateMatrixStridesWithFormat(b, h, s_q, d_v, o_strides.data(), o_format); - O->set_output(true).set_dim({b, h, s_q, d_v}).set_stride(o_strides).set_data_type(o_tensor_type); - if (is_ragged_q) { - offset_o = mha_graph->tensor(fe::graph::Tensor_attributes() - .set_name("offset_o") - .set_dim({b + 1, 1, 1, 1}) - .set_stride({1, 1, 1, 1}) - .set_data_type(get_cudnn_fe_dtype(ragged_offset_type))); - O->set_ragged_offset(offset_o); - if (use_cu_seqlens_directly) { - O->set_ragged_offset_multiplier(offset_mults.o); - } - } - amax_o->set_output(!is_mxfp8) - .set_dim({1, 1, 1, 1}) - .set_stride({1, 1, 1, 1}) - .set_data_type(fe::DataType_t::FLOAT); - - if (use_ragged_stats) { - offset_stats = mha_graph->tensor(fe::graph::Tensor_attributes() - .set_name("offset_stats") - .set_dim({b + 1, 1, 1, 1}) - .set_stride({1, 1, 1, 1}) - .set_data_type(get_cudnn_fe_dtype(ragged_offset_type))); - } - Stats->set_output(true).set_data_type(fe::DataType_t::FLOAT).set_dim({b, h, s_q, 1}); - if (use_ragged_stats) { - Stats->set_stride({h * s_q, 1, h, 1}).set_ragged_offset(offset_stats); - if (use_cu_seqlens_directly) { - Stats->set_ragged_offset_multiplier(offset_mults.stats); - } - } else { - Stats->set_stride({h * s_q, s_q, 1, 1}); - } - - std::tuple, // Q - std::shared_ptr, // K - std::shared_ptr, // V - std::shared_ptr, // descale_q - std::shared_ptr, // descale_k - std::shared_ptr, // descale_v - std::shared_ptr, // descale_s - std::shared_ptr, // scale_s - std::shared_ptr, // scale_o - std::shared_ptr, // attn_scale - std::shared_ptr, // O - std::shared_ptr, // amax_s - std::shared_ptr> // amax_o - key_tensors_tuple = - is_mxfp8 ? std::make_tuple(Q, K, V, descale_q, descale_k, descale_v, nullptr, nullptr, - nullptr, attn_scale, O, nullptr, amax_o) - : std::make_tuple(Q, K, V, descale_q, descale_k, descale_v, descale_s, scale_s, - scale_o, attn_scale, O, amax_s, amax_o); - auto Stats_tuple = std::make_tuple(Stats); - auto bias_tuple = is_bias ? std::make_tuple(bias) : std::make_tuple(nullptr); - auto softmax_offset_tuple = - is_softmax_offset ? std::make_tuple(softmax_offset) : std::make_tuple(nullptr); - auto padding_tuple = - is_padding ? std::make_tuple(seq_q, seq_kv) : std::make_tuple(nullptr, nullptr); - auto offset_qo_tuple = - is_ragged_q ? std::make_tuple(offset_q, offset_o) : std::make_tuple(nullptr, nullptr); - auto offset_kv_tuple = - is_ragged_kv ? std::make_tuple(offset_k, offset_v) : std::make_tuple(nullptr, nullptr); - auto offset_s_tuple = use_ragged_stats ? std::make_tuple(offset_stats) : std::make_tuple(nullptr); - auto dropout_tuple = is_dropout ? std::make_tuple(dropout_seed, dropout_offset) - : std::make_tuple(nullptr, nullptr); - - return std::tuple_cat(std::make_tuple(mha_graph), key_tensors_tuple, Stats_tuple, bias_tuple, - softmax_offset_tuple, padding_tuple, offset_qo_tuple, offset_kv_tuple, - offset_s_tuple, dropout_tuple); -} - -void fused_attn_fp8_fwd_impl(const FusedAttnConfig& cfg, void* devPtrQ, void* devPtrK, - void* devPtrV, void* devPtrSoftmaxOffset, void* devPtrM, void* devPtrO, - void* devPtrDescaleQ, void* devPtrDescaleK, void* devPtrDescaleV, - void* devPtrDescaleS, void* devPtrScaleS, void* devPtrScaleO, - void* devPtrAmaxO, void* devPtrAmaxS, void* devPtrcuSeqlensQ, - void* devPtrcuSeqlensKV, void* devPtrSeqOffsetsQ, - void* devPtrSeqOffsetsKV, void* devPtrDropoutSeed, - void* devPtrDropoutOffset, void* workspace, size_t* workspace_size, - cudaStream_t stream, cudnnHandle_t handle) { - using namespace transformer_engine; - - cfg.check_derived(); - - const bool is_tensor_scaling = cfg.is_tensor_scaling; - const bool is_delayed_scaling = cfg.is_delayed_scaling_fwd; - const bool use_cu_seqlens_directly = cfg.uses_cu_seqlens_directly; - - const int64_t b = static_cast(cfg.graph_batch_size_fwd); - const int64_t actual_b = static_cast(cfg.batch_size); - float scaling_factor = cfg.attn_scale; - const bool is_padding = cfg.is_padding; - const bool is_dropout = cfg.is_dropout; - const bool is_softmax_offset = cfg.is_softmax_offset; - const bool is_ragged_q = cfg.is_ragged_q; - const bool is_ragged_kv = cfg.is_ragged_kv; - const bool use_ragged_stats = cfg.uses_ragged_stats; - const DType ragged_offset_type = cfg.ragged_offset_type_fwd; - const RaggedOffsetMultipliers offset_mults = cfg.ragged_offset_mults; - - try { - auto cache_entry = get_graph(cfg, handle); - auto [mha_graph, Q, K, V, descale_q, descale_k, descale_v, descale_s, scale_s, scale_o, - attn_scale, O, amax_s, amax_o, Stats, bias, softmax_offset, seq_q, seq_kv, offset_q, - offset_o, offset_k, offset_v, offset_stats, dropout_seed, dropout_offset] = - cache_entry->graph_and_tensors; - - // This graph is going to be used, so finish the build the cache deferred. - build_plans(Backend::FP8, Pass::Fwd, *cache_entry); - - // Exit to request upper level API to allocate memory if needed. - // When passing cu_seqlens* directly to cuDNN SDPA, no conversion workspace is - // needed: cuDNN consumes the user's cu_seqlens buffers as-is. - auto plan_workspace_size = alignTo<16>(mha_graph->get_workspace_size()); - const size_t num_bytes_per_seqlen = alignTo<16>(b * sizeof(int32_t)); - const size_t num_bytes_per_ragged_offset = - alignTo<16>(((b + 1) * typeToNumBits(ragged_offset_type)) / 8); - size_t actual_seqlen_workspace_size = 0; - size_t seqlen_offsets_workspace_size = 0; - if (!use_cu_seqlens_directly) { - if (is_padding) { - actual_seqlen_workspace_size = 2 * num_bytes_per_seqlen; - } - if (is_ragged_q || is_ragged_kv) { - const size_t count = - 2 * (static_cast(is_ragged_q) + static_cast(is_ragged_kv)); - seqlen_offsets_workspace_size = - (use_ragged_stats ? count + 1 : count) * num_bytes_per_ragged_offset; - } - } - if (workspace == nullptr) { - *workspace_size = - plan_workspace_size + actual_seqlen_workspace_size + seqlen_offsets_workspace_size; - return; - } - // cuDNN stream check needs to be moved here to support dummy kernel calls with - // null streams for sizing the cuDNN workspace. - NVTE_CHECK_CUDNN(cudnnSetStream(handle, stream)); - - // Build variant pack - std::unordered_map, void*> variant_pack = { - {Q, devPtrQ}, - {K, devPtrK}, - {V, devPtrV}, - {descale_q, devPtrDescaleQ}, - {descale_k, devPtrDescaleK}, - {descale_v, devPtrDescaleV}, - {attn_scale, &scaling_factor}, - {O, devPtrO}, - {Stats, devPtrM}}; - - if (is_delayed_scaling) { - variant_pack[scale_o] = devPtrScaleO; - } - if (is_tensor_scaling) { - variant_pack[descale_s] = devPtrDescaleS; - variant_pack[scale_s] = devPtrScaleS; - variant_pack[amax_s] = devPtrAmaxS; - variant_pack[amax_o] = devPtrAmaxO; - } - - /* if (cfg.is_bias) { - variant_pack[bias] = devPtrBias; - } */ - - if (is_padding) { - if (use_cu_seqlens_directly) { - variant_pack[seq_q] = devPtrcuSeqlensQ; - variant_pack[seq_kv] = devPtrcuSeqlensKV; - } else { - constexpr size_t nthreads_per_block = 128; - const size_t grid = (b + nthreads_per_block - 1) / nthreads_per_block; - void* devActualSeqlenQ = static_cast(workspace) + plan_workspace_size; - void* devActualSeqlenKV = static_cast(devActualSeqlenQ) + num_bytes_per_seqlen; - cu_seqlens_to_actual_seqlens<<>>( - actual_b, b, static_cast(devPtrcuSeqlensQ), - static_cast(devPtrcuSeqlensKV), static_cast(devActualSeqlenQ), - static_cast(devActualSeqlenKV)); - NVTE_CHECK_CUDA(cudaGetLastError()); - variant_pack[seq_q] = devActualSeqlenQ; - variant_pack[seq_kv] = devActualSeqlenKV; - } - } - - if (use_cu_seqlens_directly) { - // The token-unit cu_seqlens_padded buffers serve as the ragged offsets; the engine - // applies the per-tensor multipliers set at graph build time. - if (is_ragged_q) { - variant_pack[offset_q] = devPtrSeqOffsetsQ; - variant_pack[offset_o] = devPtrSeqOffsetsQ; - } - if (is_ragged_kv) { - void* devOffsetsKV = offset_mults.kv_from_q ? devPtrSeqOffsetsQ : devPtrSeqOffsetsKV; - variant_pack[offset_k] = devOffsetsKV; - variant_pack[offset_v] = devOffsetsKV; - } - if (use_ragged_stats) { - variant_pack[offset_stats] = devPtrSeqOffsetsQ; - } - } else if (is_ragged_q || is_ragged_kv) { - constexpr size_t nthreads_per_block = 128; - const size_t grid = (b + nthreads_per_block) / nthreads_per_block; - void* devOffsets = - static_cast(workspace) + plan_workspace_size + actual_seqlen_workspace_size; - void* devOffsetsQ = nullptr; - void* devOffsetsO = nullptr; - if (is_ragged_q) { - devOffsetsQ = devOffsets; - devOffsetsO = static_cast(devOffsetsQ) + num_bytes_per_ragged_offset; - } - void* devOffsetsK = nullptr; - void* devOffsetsV = nullptr; - if (is_ragged_kv) { - devOffsetsK = static_cast(devOffsets) + - static_cast(is_ragged_q) * 2 * num_bytes_per_ragged_offset; - devOffsetsV = static_cast(devOffsetsK) + num_bytes_per_ragged_offset; - } - void* devOffsetsS = nullptr; - if (use_ragged_stats) { - devOffsetsS = static_cast(devOffsets) + - (static_cast(is_ragged_q) + static_cast(is_ragged_kv)) * 2 * - num_bytes_per_ragged_offset; - } - cu_seqlens_padded_to_offsets<<>>( - offset_mults, actual_b, b, static_cast(devPtrSeqOffsetsQ), - static_cast(devPtrSeqOffsetsKV), ragged_offset_type, devOffsetsQ, devOffsetsK, - devOffsetsV, devOffsetsO, devOffsetsS); - NVTE_CHECK_CUDA(cudaGetLastError()); - if (is_ragged_q) { - variant_pack[offset_q] = devOffsetsQ; - variant_pack[offset_o] = devOffsetsO; - } - if (is_ragged_kv) { - variant_pack[offset_k] = devOffsetsK; - variant_pack[offset_v] = devOffsetsV; - } - if (use_ragged_stats) { - variant_pack[offset_stats] = devOffsetsS; - } - } - - if (is_dropout) { - variant_pack[dropout_seed] = devPtrDropoutSeed; - variant_pack[dropout_offset] = devPtrDropoutOffset; - } - - if (is_softmax_offset) { - variant_pack[softmax_offset] = devPtrSoftmaxOffset; - } - - NVTE_CHECK_CUDNN_FE(mha_graph->execute(handle, variant_pack, workspace)); - graph_cache_debug::record_execute(Backend::FP8, Pass::Fwd); - } catch (cudnn_frontend::cudnnException& e) { - NVTE_ERROR(e.what()); - } -} - -using Fp8BwdGraphAndTensors = - std::tuple, - std::shared_ptr, // Q - std::shared_ptr, // Q_t - std::shared_ptr, // K - std::shared_ptr, // K_t - std::shared_ptr, // V - std::shared_ptr, // O - std::shared_ptr, // Stats - std::shared_ptr, // dO - std::shared_ptr, // dO_t - std::shared_ptr, // dO_f16 - std::shared_ptr, // attn_scale - std::shared_ptr, // descale_q - std::shared_ptr, // descale_q_t - std::shared_ptr, // descale_k - std::shared_ptr, // descale_k_t - std::shared_ptr, // descale_v - std::shared_ptr, // descale_o - std::shared_ptr, // descale_dO - std::shared_ptr, // descale_dO_t - std::shared_ptr, // descale_s - std::shared_ptr, // descale_dP - std::shared_ptr, // scale_dQ - std::shared_ptr, // scale_dK - std::shared_ptr, // scale_dV - std::shared_ptr, // scale_s - std::shared_ptr, // scale_dP - std::shared_ptr, // dQ - std::shared_ptr, // dK - std::shared_ptr, // dV - std::shared_ptr, // amax_dQ - std::shared_ptr, // amax_dK - std::shared_ptr, // amax_dV - std::shared_ptr, // amax_dP - std::shared_ptr, // bias - std::shared_ptr, // dBias - std::shared_ptr, // softmax_offset - std::shared_ptr, // d_softmax_offset - std::shared_ptr, // seq_q - std::shared_ptr, // seq_kv - std::shared_ptr, // offset_q - std::shared_ptr, // offset_o - std::shared_ptr, // offset_k - std::shared_ptr, // offset_v - std::shared_ptr, // offset_stats - std::shared_ptr, // dropout_seed - std::shared_ptr>; // dropout_offset - -static Fp8BwdGraphAndTensors create_graph_fp8_bwd(const FusedAttnConfig& cfg) { - const auto cudnn_runtime_version = cudnnGetVersion(); - const cudnn_frontend::DataType_t qkv_tensor_type = - get_cudnn_fe_dtype(static_cast(cfg.qkv_dtype)); - const cudnn_frontend::DataType_t o_tensor_type = - get_cudnn_fe_dtype(static_cast(cfg.o_dtype)); - const cudnn_frontend::DataType_t do_tensor_type = - get_cudnn_fe_dtype(static_cast(cfg.do_dtype)); - const cudnn_frontend::DataType_t dqkv_tensor_type = - get_cudnn_fe_dtype(static_cast(cfg.dqkv_dtype)); - const int64_t b = static_cast(cfg.graph_batch_size_bwd); - const int64_t h = static_cast(cfg.num_attn_heads); - const int64_t hg = static_cast(cfg.num_gqa_groups); - const int64_t s_q = static_cast(cfg.graph_max_seqlen_q); - const int64_t s_kv = static_cast(cfg.graph_max_seqlen_kv); - const int64_t d_qk = static_cast(cfg.head_dim_qk); - const int64_t d_v = static_cast(cfg.head_dim_v); - const int64_t window_size_left = cfg.window_size_left; - const int64_t window_size_right = cfg.window_size_right; - const float dropout_probability = cfg.dropout; - const NVTE_QKV_Layout qkv_layout = cfg.qkv_layout; - const NVTE_QKV_Layout dqkv_layout = cfg.dqkv_layout; - const NVTE_QKV_Format o_format = cfg.o_format; - const NVTE_QKV_Format do_format = cfg.do_format; - const NVTE_QKV_Format qkv_scale_inv_format = cfg.qkv_scale_inv_format; - const NVTE_QKV_Format do_scale_inv_format = cfg.do_scale_inv_format; - const bool bottom_right_diagonal = cfg.bottom_right_diagonal; - const bool deterministic = cfg.deterministic; - const bool is_bias = cfg.is_bias; - const bool is_causal = cfg.is_causal; - const bool is_causal_bottom_right = cfg.is_causal_bottom_right; - const bool is_padding = cfg.is_padding; - const bool is_dropout = cfg.is_dropout; - const bool is_softmax_offset = cfg.is_softmax_offset; - const bool is_mxfp8 = cfg.is_mxfp8; - const bool is_tensor_scaling = cfg.is_tensor_scaling; - const bool is_delayed_scaling = cfg.is_delayed_scaling_bwd; - const bool is_current_scaling = cfg.is_current_scaling_bwd; - const bool is_O_in_F16 = cfg.is_o_in_f16; - const bool is_ragged_q = cfg.is_ragged_q; - const bool is_ragged_kv = cfg.is_ragged_kv; - const bool use_ragged_stats = cfg.uses_ragged_stats; - const DType ragged_offset_type = cfg.ragged_offset_type_bwd; - - auto mha_graph = std::make_shared(); - - mha_graph->set_io_data_type(qkv_tensor_type) - .set_intermediate_data_type(fe::DataType_t::FLOAT) - .set_compute_data_type(fe::DataType_t::FLOAT); - - std::shared_ptr offset_q, offset_k, offset_v, offset_o, - offset_stats; - std::shared_ptr Q, Q_t, K, K_t, V, O, dO, dO_t, dO_f16, Stats, - attn_scale; - std::shared_ptr descale_q, descale_q_t, descale_k, descale_k_t, - descale_v; - std::shared_ptr descale_s, descale_o; - std::shared_ptr descale_dP, descale_dO, descale_dO_t; - std::shared_ptr scale_s, scale_dP; - std::shared_ptr scale_dQ, scale_dK, scale_dV; - std::shared_ptr bias, dBias, softmax_offset, d_softmax_offset; - std::shared_ptr seq_q, seq_kv; - std::shared_ptr dropout_seed, dropout_offset; - - // Q, K, V, O, dO, stats, attn_scale - std::vector q_strides(4), k_strides(4), v_strides(4), o_strides(4), dO_strides(4); - generateMatrixStridesWithLayout(b, h, hg, s_q, s_kv, d_qk, d_v, q_strides.data(), - k_strides.data(), v_strides.data(), qkv_layout); - generateMatrixStridesWithFormat(b, h, s_q, d_v, o_strides.data(), o_format); - generateMatrixStridesWithFormat(b, h, s_q, d_v, dO_strides.data(), do_format); - Q = mha_graph->tensor(fe::graph::Tensor_attributes() - .set_name("Q") - .set_dim({b, h, s_q, d_qk}) - .set_stride(q_strides) - .set_data_type(qkv_tensor_type)); - K = mha_graph->tensor(fe::graph::Tensor_attributes() - .set_name("K") - .set_dim({b, hg, s_kv, d_qk}) - .set_stride(k_strides) - .set_data_type(qkv_tensor_type)); - V = mha_graph->tensor(fe::graph::Tensor_attributes() - .set_name("V") - .set_dim({b, hg, s_kv, d_v}) - .set_stride(v_strides) - .set_data_type(qkv_tensor_type)); - O = mha_graph->tensor(fe::graph::Tensor_attributes() - .set_name("O") - .set_dim({b, h, s_q, d_v}) - .set_stride(o_strides) - .set_data_type(o_tensor_type)); - dO = mha_graph->tensor(fe::graph::Tensor_attributes() - .set_name("dO") - .set_dim({b, h, s_q, d_v}) - .set_stride(dO_strides) - .set_data_type(do_tensor_type)); - if (is_ragged_q) { - offset_q = mha_graph->tensor(fe::graph::Tensor_attributes() - .set_name("offset_q") - .set_dim({b + 1, 1, 1, 1}) - .set_stride({1, 1, 1, 1}) - .set_data_type(get_cudnn_fe_dtype(ragged_offset_type))); - offset_o = mha_graph->tensor(fe::graph::Tensor_attributes() - .set_name("offset_o") - .set_dim({b + 1, 1, 1, 1}) - .set_stride({1, 1, 1, 1}) - .set_data_type(get_cudnn_fe_dtype(ragged_offset_type))); - Q->set_ragged_offset(offset_q); - O->set_ragged_offset(offset_o); - dO->set_ragged_offset(offset_o); - } - if (is_ragged_kv) { - offset_k = mha_graph->tensor(fe::graph::Tensor_attributes() - .set_name("offset_k") - .set_dim({b + 1, 1, 1, 1}) - .set_stride({1, 1, 1, 1}) - .set_data_type(get_cudnn_fe_dtype(ragged_offset_type))); - offset_v = mha_graph->tensor(fe::graph::Tensor_attributes() - .set_name("offset_v") - .set_dim({b + 1, 1, 1, 1}) - .set_stride({1, 1, 1, 1}) - .set_data_type(get_cudnn_fe_dtype(ragged_offset_type))); - K->set_ragged_offset(offset_k); - V->set_ragged_offset(offset_v); - } - if (use_ragged_stats) { - offset_stats = mha_graph->tensor(fe::graph::Tensor_attributes() - .set_name("offset_stats") - .set_dim({b + 1, 1, 1, 1}) - .set_stride({1, 1, 1, 1}) - .set_data_type(get_cudnn_fe_dtype(ragged_offset_type))); - } - Stats = mha_graph->tensor(fe::graph::Tensor_attributes() - .set_name("Stats") - .set_dim({b, h, s_q, 1}) - .set_data_type(fe::DataType_t::FLOAT)); - if (use_ragged_stats) { - Stats->set_stride({h * s_q, 1, h, 1}).set_ragged_offset(offset_stats); - } else { - Stats->set_stride({h * s_q, s_q, 1, 1}); - } - attn_scale = mha_graph->tensor(fe::graph::Tensor_attributes() - .set_name("attn_scale") - .set_dim({1, 1, 1, 1}) - .set_stride({1, 1, 1, 1}) - .set_is_pass_by_value(true) - .set_data_type(fe::DataType_t::FLOAT)); - - // Descale_q, Descale_k, Descale_v, Descale_s, Scale_s, Descale_dP, Scale_dP, Descale_o, Descale_dO, Scale_dQ, Scale_dK, Scale_dV - if (is_tensor_scaling) { - descale_q = mha_graph->tensor(fe::graph::Tensor_attributes() - .set_name("Descale_q") - .set_dim({1, 1, 1, 1}) - .set_stride({1, 1, 1, 1}) - .set_data_type(fe::DataType_t::FLOAT)); - descale_k = mha_graph->tensor_like(descale_q, "Descale_q"); - descale_v = mha_graph->tensor_like(descale_q, "Descale_v"); - descale_s = mha_graph->tensor_like(descale_q, "Descale_s"); - scale_s = mha_graph->tensor_like(descale_q, "Scale_s"); - descale_dP = mha_graph->tensor_like(descale_q, "Descale_dP"); - scale_dP = mha_graph->tensor_like(descale_q, "Scale_dP"); - if (is_current_scaling && is_O_in_F16) { - descale_o = mha_graph->tensor(1.0f); - } else { - descale_o = mha_graph->tensor_like(descale_q, "Descale_O"); - } - descale_dO = mha_graph->tensor_like(descale_q, "Descale_dO"); - if (is_delayed_scaling) { - scale_dQ = mha_graph->tensor_like(descale_q, "Scale_dQ"); - scale_dK = mha_graph->tensor_like(descale_q, "Scale_dK"); - scale_dV = mha_graph->tensor_like(descale_q, "Scale_dV"); - } - if (is_current_scaling) { - scale_dQ = mha_graph->tensor(1.0f); - scale_dK = mha_graph->tensor(1.0f); - scale_dV = mha_graph->tensor(1.0f); - } - } else if (is_mxfp8) { - const NVTE_QKV_Format q_format = cfg.q_format; - const NVTE_QKV_Format kv_format = cfg.kv_format; - const NVTE_QKV_Format q_scale_inv_format = - (qkv_scale_inv_format != NVTE_QKV_Format_NOT_SET) ? qkv_scale_inv_format : q_format; - const NVTE_QKV_Format kv_scale_inv_format = - (qkv_scale_inv_format != NVTE_QKV_Format_NOT_SET) ? qkv_scale_inv_format : kv_format; - const NVTE_QKV_Format do_scale_format_ = - (do_scale_inv_format != NVTE_QKV_Format_NOT_SET) ? do_scale_inv_format : do_format; - // Q_t, K_t, dO_t, dO_f16 - std::vector q_t_strides(4), k_t_strides(4), dO_t_strides(4); - generateMatrixStridesWithFormat(b, h, s_q, d_qk, q_t_strides.data(), q_format); - generateMatrixStridesWithFormat(b, hg, s_kv, d_qk, k_t_strides.data(), kv_format); - generateMatrixStridesWithFormat(b, h, s_q, d_v, dO_t_strides.data(), do_format); - Q_t = mha_graph->tensor(fe::graph::Tensor_attributes() - .set_name("Q_t") - .set_dim({b, h, s_q, d_qk}) - .set_stride(q_t_strides) - .set_data_type(qkv_tensor_type)); - K_t = mha_graph->tensor(fe::graph::Tensor_attributes() - .set_name("K_t") - .set_dim({b, hg, s_kv, d_qk}) - .set_stride(k_t_strides) - .set_data_type(qkv_tensor_type)); - dO_t = mha_graph->tensor(fe::graph::Tensor_attributes() - .set_name("dO_t") - .set_dim({b, h, s_q, d_v}) - .set_stride(dO_t_strides) - .set_data_type(do_tensor_type)); - dO_f16 = mha_graph->tensor(fe::graph::Tensor_attributes() - .set_name("dO_f16") - .set_dim({b, h, s_q, d_v}) - .set_stride(dO_strides) - .set_data_type(o_tensor_type)); - // Descale_q, Descale_q_t, Descale_k, Descale_k_t, Descale_v, Descale_dO, Descale_dO_t - auto padded = pad_s_d_for_mxfp8(s_q, s_kv, d_qk, d_v); - std::vector q_scale_strides(4), q_t_scale_strides(4), k_scale_strides(4), - k_t_scale_strides(4), v_scale_strides(4), dO_scale_strides(4), dO_t_scale_strides(4); - generateMatrixStridesWithFormat(b, h, padded.s_q_padded, padded.d_qk_scale_padded, - q_scale_strides.data(), q_scale_inv_format); - generateMatrixStridesWithFormat(b, h, padded.s_q_scale_padded, padded.d_qk_padded, - q_t_scale_strides.data(), q_scale_inv_format); - generateMatrixStridesWithFormat(b, hg, padded.s_kv_padded, padded.d_qk_scale_padded, - k_scale_strides.data(), kv_scale_inv_format); - generateMatrixStridesWithFormat(b, hg, padded.s_kv_scale_padded, padded.d_qk_padded, - k_t_scale_strides.data(), kv_scale_inv_format); - generateMatrixStridesWithFormat(b, hg, padded.s_kv_padded, padded.d_v_scale_padded, - v_scale_strides.data(), kv_scale_inv_format); - generateMatrixStridesWithFormat(b, h, padded.s_q_padded, padded.d_v_scale_padded, - dO_scale_strides.data(), do_scale_format_); - generateMatrixStridesWithFormat(b, h, padded.s_q_scale_padded, padded.d_v_padded, - dO_t_scale_strides.data(), do_scale_format_); - descale_q = mha_graph->tensor(fe::graph::Tensor_attributes() - .set_name("Descale_q") - .set_dim({b, h, padded.s_q_padded, padded.d_qk_scale_padded}) - .set_stride(q_scale_strides) - .set_data_type(fe::DataType_t::FP8_E8M0) - .set_reordering_type(fe::TensorReordering_t::F8_128x4)); - descale_q_t = - mha_graph->tensor(fe::graph::Tensor_attributes() - .set_name("Descale_q_t") - .set_dim({b, h, padded.s_q_scale_padded, padded.d_qk_padded}) - .set_stride(q_t_scale_strides) - .set_data_type(fe::DataType_t::FP8_E8M0) - .set_reordering_type(fe::TensorReordering_t::F8_128x4)); - descale_k = - mha_graph->tensor(fe::graph::Tensor_attributes() - .set_name("Descale_k") - .set_dim({b, hg, padded.s_kv_padded, padded.d_qk_scale_padded}) - .set_stride(k_scale_strides) - .set_data_type(fe::DataType_t::FP8_E8M0) - .set_reordering_type(fe::TensorReordering_t::F8_128x4)); - descale_k_t = - mha_graph->tensor(fe::graph::Tensor_attributes() - .set_name("Descale_k_t") - .set_dim({b, hg, padded.s_kv_scale_padded, padded.d_qk_padded}) - .set_stride(k_t_scale_strides) - .set_data_type(fe::DataType_t::FP8_E8M0) - .set_reordering_type(fe::TensorReordering_t::F8_128x4)); - descale_v = mha_graph->tensor(fe::graph::Tensor_attributes() - .set_name("Descale_v") - .set_dim({b, hg, padded.s_kv_padded, padded.d_v_scale_padded}) - .set_stride(v_scale_strides) - .set_data_type(fe::DataType_t::FP8_E8M0) - .set_reordering_type(fe::TensorReordering_t::F8_128x4)); - descale_dO = mha_graph->tensor(fe::graph::Tensor_attributes() - .set_name("Descale_dO") - .set_dim({b, h, padded.s_q_padded, padded.d_v_scale_padded}) - .set_stride(dO_scale_strides) - .set_data_type(fe::DataType_t::FP8_E8M0) - .set_reordering_type(fe::TensorReordering_t::F8_128x4)); - descale_dO_t = - mha_graph->tensor(fe::graph::Tensor_attributes() - .set_name("Descale_dO_t") - .set_dim({b, h, padded.s_q_scale_padded, padded.d_v_padded}) - .set_stride(dO_t_scale_strides) - .set_data_type(fe::DataType_t::FP8_E8M0) - .set_reordering_type(fe::TensorReordering_t::F8_128x4)); - } - - fe::graph::SDPA_fp8_backward_attributes sdpa_backward_options; - sdpa_backward_options = fe::graph::SDPA_fp8_backward_attributes() - .set_name("sdpa_fp8_backward") - .set_causal_mask(is_causal) - .set_attn_scale(attn_scale); - - fe::DiagonalAlignment_t const& diagonal_alignment = bottom_right_diagonal - ? fe::DiagonalAlignment_t::BOTTOM_RIGHT - : fe::DiagonalAlignment_t::TOP_LEFT; - sdpa_backward_options.set_diagonal_alignment(diagonal_alignment); - - if (cudnn_runtime_version >= 92100) { - if (window_size_left != -1) { - sdpa_backward_options.set_diagonal_band_left_bound(window_size_left + 1); - } - if (window_size_right != -1) { - sdpa_backward_options.set_diagonal_band_right_bound(window_size_right); - } - } - if (is_causal_bottom_right) { - sdpa_backward_options.set_diagonal_band_right_bound(0); - } - - // sdpa_backward_options.set_alibi_mask(is_alibi); - - // if (is_bias) { - // bias = mha_graph->tensor(fe::graph::Tensor_attributes() - // .set_name("bias") - // .set_dim({bias_b, bias_h, bias_sq, bias_skv}) - // .set_stride({bias_h * bias_sq * bias_skv, bias_sq * bias_skv, bias_skv, 1})); - // dBias = mha_graph->tensor(fe::graph::Tensor_attributes() - // .set_name("dBias") - // .set_dim({bias_b, bias_h, bias_sq, bias_skv}) - // .set_stride({bias_h * bias_sq * bias_skv, bias_sq * bias_skv, bias_skv, 1})); - // sdpa_backward_options.set_bias(bias); - // bias shapes [1, 1, s, s], [b, 1, s, s], [b, h, s, s], [1, h, s, s] are supported for dbias calculation - // bias shape [1, 1, 1, s] is not supported for dbias calculation as of cuDNN 9.18 - // if (!((bias_b == 1) && (bias_h == 1) && (bias_sq == 1))) { - // sdpa_backward_options.set_dbias(dBias); - // } - // } - - if (cudnn_runtime_version >= 91900) { - sdpa_backward_options.set_deterministic_algorithm(deterministic); - } - - if (is_padding) { - seq_q = mha_graph->tensor(fe::graph::Tensor_attributes() - .set_name("seq_q") - .set_dim({b, 1, 1, 1}) - .set_stride({1, 1, 1, 1}) - .set_data_type(fe::DataType_t::INT32)); - seq_kv = mha_graph->tensor(fe::graph::Tensor_attributes() - .set_name("seq_kv") - .set_dim({b, 1, 1, 1}) - .set_stride({1, 1, 1, 1}) - .set_data_type(fe::DataType_t::INT32)); - sdpa_backward_options.set_padding_mask(is_padding).set_seq_len_q(seq_q).set_seq_len_kv(seq_kv); - } - - if (is_dropout) { - dropout_seed = mha_graph->tensor(fe::graph::Tensor_attributes() - .set_name("Seed") - .set_dim({1, 1, 1, 1}) - .set_stride({1, 1, 1, 1}) - .set_data_type(fe::DataType_t::INT64)); - dropout_offset = mha_graph->tensor(fe::graph::Tensor_attributes() - .set_name("Offset") - .set_dim({1, 1, 1, 1}) - .set_stride({1, 1, 1, 1}) - .set_data_type(fe::DataType_t::INT64)); - sdpa_backward_options.set_dropout(dropout_probability, dropout_seed, dropout_offset); - } - - if (is_softmax_offset) { - softmax_offset = mha_graph->tensor(fe::graph::Tensor_attributes() - .set_name("softmax_offset") - .set_dim({1, h, 1, 1}) - .set_stride({h, 1, 1, 1}) - .set_data_type(fe::DataType_t::FLOAT)); - sdpa_backward_options.set_sink_token(softmax_offset); - d_softmax_offset = mha_graph->tensor(fe::graph::Tensor_attributes() - .set_name("d_softmax_offset") - .set_dim({1, h, 1, 1}) - .set_stride({h, 1, 1, 1}) - .set_data_type(fe::DataType_t::FLOAT)); - sdpa_backward_options.set_dsink_token(d_softmax_offset); - } - - std::shared_ptr dQ, dK, dV, amax_dQ, amax_dK, amax_dV, amax_dP; - if (is_tensor_scaling) { - std::tie(dQ, dK, dV, amax_dQ, amax_dK, amax_dV, amax_dP) = - std::apply([](const auto&... elems) { return std::make_tuple(elems...); }, - mha_graph->sdpa_fp8_backward(Q, K, V, O, dO, Stats, descale_q, descale_k, - descale_v, descale_o, descale_dO, descale_s, - descale_dP, scale_s, scale_dQ, scale_dK, scale_dV, - scale_dP, sdpa_backward_options)); - } else if (is_mxfp8) { - std::tie(dQ, dK, dV, amax_dQ, amax_dK, amax_dV) = std::apply( - [](const auto&... elems) { return std::make_tuple(elems...); }, - mha_graph->sdpa_fp8_backward(Q, Q_t, K, K_t, V, O, dO_f16, dO, dO_t, Stats, descale_q, - descale_q_t, descale_k, descale_k_t, descale_v, descale_dO, - descale_dO_t, sdpa_backward_options)); - } - std::vector dq_strides(4), dk_strides(4), dv_strides(4); - generateMatrixStridesWithLayout(b, h, hg, s_q, s_kv, d_qk, d_v, dq_strides.data(), - dk_strides.data(), dv_strides.data(), dqkv_layout); - dQ->set_output(true) - .set_dim({b, h, s_q, d_qk}) - .set_stride(dq_strides) - .set_data_type(dqkv_tensor_type); - dK->set_output(true) - .set_dim({b, hg, s_kv, d_qk}) - .set_stride(dk_strides) - .set_data_type(dqkv_tensor_type); - dV->set_output(true) - .set_dim({b, hg, s_kv, d_v}) - .set_stride(dv_strides) - .set_data_type(dqkv_tensor_type); - if (is_ragged_q) { - dQ->set_ragged_offset(offset_q); - } - if (is_ragged_kv) { - dK->set_ragged_offset(offset_k); - dV->set_ragged_offset(offset_v); - } - amax_dQ->set_output(!is_mxfp8) - .set_dim({1, 1, 1, 1}) - .set_stride({1, 1, 1, 1}) - .set_data_type(fe::DataType_t::FLOAT); - amax_dK->set_output(!is_mxfp8) - .set_dim({1, 1, 1, 1}) - .set_stride({1, 1, 1, 1}) - .set_data_type(fe::DataType_t::FLOAT); - amax_dV->set_output(!is_mxfp8) - .set_dim({1, 1, 1, 1}) - .set_stride({1, 1, 1, 1}) - .set_data_type(fe::DataType_t::FLOAT); - if (is_tensor_scaling) { - amax_dP->set_output(true) - .set_dim({1, 1, 1, 1}) - .set_stride({1, 1, 1, 1}) - .set_data_type(fe::DataType_t::FLOAT); - } - - std::tuple, // Q - std::shared_ptr, // K - std::shared_ptr, // V - std::shared_ptr, // O - std::shared_ptr, // Stats - std::shared_ptr, // dO - std::shared_ptr, // attn_scale - std::shared_ptr, // descale_q - std::shared_ptr, // descale_k - std::shared_ptr, // descale_v - std::shared_ptr, // descale_o - std::shared_ptr, // descale_dO - std::shared_ptr, // descale_s - std::shared_ptr, // descale_dP - std::shared_ptr, // scale_dQ - std::shared_ptr, // scale_dK - std::shared_ptr, // scale_dV - std::shared_ptr, // scale_s - std::shared_ptr, // scale_dP - std::shared_ptr, // dQ - std::shared_ptr, // dK - std::shared_ptr, // dV - std::shared_ptr, // amax_dQ - std::shared_ptr, // amax_dK - std::shared_ptr, // amax_dV - std::shared_ptr> // amax_dP - key_tensors_tuple = - std::make_tuple(Q, K, V, O, Stats, dO, attn_scale, descale_q, descale_k, descale_v, - descale_o, descale_dO, descale_s, descale_dP, scale_s, scale_dQ, scale_dK, - scale_dV, scale_dP, dQ, dK, dV, amax_dQ, amax_dK, amax_dV, amax_dP); - auto mxfp8_tensors_tuple = - is_mxfp8 ? std::make_tuple(Q_t, K_t, dO_f16, dO_t, descale_q_t, descale_k_t, descale_dO_t) - : std::make_tuple(nullptr, nullptr, nullptr, nullptr, nullptr, nullptr, nullptr); - auto bias_tuple = is_bias ? std::make_tuple(bias, dBias) : std::make_tuple(nullptr, nullptr); - auto softmax_offset_tuple = is_softmax_offset ? std::make_tuple(softmax_offset, d_softmax_offset) - : std::make_tuple(nullptr, nullptr); - auto padding_tuple = - is_padding ? std::make_tuple(seq_q, seq_kv) : std::make_tuple(nullptr, nullptr); - auto offset_qo_tuple = - is_ragged_q ? std::make_tuple(offset_q, offset_o) : std::make_tuple(nullptr, nullptr); - auto offset_kv_tuple = - is_ragged_kv ? std::make_tuple(offset_k, offset_v) : std::make_tuple(nullptr, nullptr); - auto offset_s_tuple = use_ragged_stats ? std::make_tuple(offset_stats) : std::make_tuple(nullptr); - auto dropout_tuple = is_dropout ? std::make_tuple(dropout_seed, dropout_offset) - : std::make_tuple(nullptr, nullptr); - - return std::tuple_cat(std::make_tuple(mha_graph), key_tensors_tuple, mxfp8_tensors_tuple, - bias_tuple, softmax_offset_tuple, padding_tuple, offset_qo_tuple, - offset_kv_tuple, offset_s_tuple, dropout_tuple); -} - -void fused_attn_fp8_bwd_impl( - const FusedAttnConfig& cfg, void* devPtrQ, void* devPtrK, void* devPtrV, void* devPtrM, - void* devPtrO, void* devPtrdO, void* devPtrSoftmaxOffset, void* devPtrdQ, void* devPtrdK, - void* devPtrdV, void* devPtrdSoftmaxOffset, void* devPtrDescaleQ, void* devPtrDescaleK, - void* devPtrDescaleV, void* devPtrDescaleO, void* devPtrDescaledO, void* devPtrDescaleS, - void* devPtrDescaledP, void* devPtrScaleS, void* devPtrScaledP, void* devPtrScaledQ, - void* devPtrScaledK, void* devPtrScaledV, void* devPtrAmaxdP, void* devPtrAmaxdQ, - void* devPtrAmaxdK, void* devPtrAmaxdV, void* devPtrQ_t, void* devPtrK_t, void* devPtrdO_f16, - void* devPtrdO_t, void* devPtrDescaleQ_t, void* devPtrDescaleK_t, void* devPtrDescaledO_t, - void* devPtrcuSeqlensQ, void* devPtrcuSeqlensKV, void* devPtrSeqOffsetsQ, - void* devPtrSeqOffsetsKV, void* devPtrDropoutSeed, void* devPtrDropoutOffset, void* workspace, - size_t* workspace_size, cudaStream_t stream, cudnnHandle_t handle) { - using namespace transformer_engine; - - cfg.check_derived(); - - const bool is_mxfp8 = cfg.is_mxfp8; - const bool is_tensor_scaling = cfg.is_tensor_scaling; - const bool is_delayed_scaling = cfg.is_delayed_scaling_bwd; - const bool is_current_scaling = cfg.is_current_scaling_bwd; - const bool is_O_in_F16 = cfg.is_o_in_f16; - - const int64_t b = static_cast(cfg.graph_batch_size_bwd); - const int64_t actual_b = static_cast(cfg.batch_size); - float scaling_factor = cfg.attn_scale; - const bool is_padding = cfg.is_padding; - const bool is_dropout = cfg.is_dropout; - const bool is_softmax_offset = cfg.is_softmax_offset; - const bool is_ragged_q = cfg.is_ragged_q; - const bool is_ragged_kv = cfg.is_ragged_kv; - const bool use_ragged_stats = cfg.uses_ragged_stats; - const DType ragged_offset_type = cfg.ragged_offset_type_bwd; - - try { - auto cache_entry = get_graph(cfg, handle); - auto [mha_graph, Q, K, V, O, Stats, dO, attn_scale, descale_q, descale_k, descale_v, descale_o, - descale_dO, descale_s, descale_dP, scale_s, scale_dQ, scale_dK, scale_dV, scale_dP, dQ, - dK, dV, amax_dQ, amax_dK, amax_dV, amax_dP, Q_t, K_t, dO_f16, dO_t, descale_q_t, - descale_k_t, descale_dO_t, bias, dBias, softmax_offset, d_softmax_offset, seq_q, seq_kv, - offset_q, offset_o, offset_k, offset_v, offset_stats, dropout_seed, dropout_offset] = - cache_entry->graph_and_tensors; - - // This graph is going to be used, so finish the build the cache deferred. - build_plans(Backend::FP8, Pass::Bwd, *cache_entry); - - // Exit to request upper level API to allocate memory if needed - auto plan_workspace_size = alignTo<16>(mha_graph->get_workspace_size()); - const size_t num_bytes_per_seqlen = alignTo<16>(b * sizeof(int32_t)); - const size_t num_bytes_per_ragged_offset = - alignTo<16>(((b + 1) * typeToNumBits(ragged_offset_type)) / 8); - const size_t actual_seqlen_workspace_size = 2 * num_bytes_per_seqlen; - size_t seqlen_offsets_workspace_size = 0; - if (is_ragged_q || is_ragged_kv) { - const size_t count = - 2 * (static_cast(is_ragged_q) + static_cast(is_ragged_kv)); - seqlen_offsets_workspace_size = - (use_ragged_stats ? count + 1 : count) * num_bytes_per_ragged_offset; - } - if (workspace == nullptr) { - *workspace_size = - plan_workspace_size + actual_seqlen_workspace_size + seqlen_offsets_workspace_size; - return; - } - // cuDNN stream check needs to be moved here to support dummy kernel calls with - // null streams for sizing the cuDNN workspace. - NVTE_CHECK_CUDNN(cudnnSetStream(handle, stream)); - - // build variant pack - std::unordered_map, void*> variant_pack = { - {Q, devPtrQ}, - {K, devPtrK}, - {V, devPtrV}, - {O, devPtrO}, - {Stats, devPtrM}, - {dO, devPtrdO}, - {attn_scale, &scaling_factor}, - {descale_q, devPtrDescaleQ}, - {descale_k, devPtrDescaleK}, - {descale_v, devPtrDescaleV}, - {descale_dO, devPtrDescaledO}, - {dQ, devPtrdQ}, - {dK, devPtrdK}, - {dV, devPtrdV}, - }; - if (is_tensor_scaling) { - variant_pack[descale_s] = devPtrDescaleS; - variant_pack[descale_dP] = devPtrDescaledP; - variant_pack[scale_s] = devPtrScaleS; - variant_pack[scale_dP] = devPtrScaledP; - variant_pack[amax_dP] = devPtrAmaxdP; - variant_pack[amax_dQ] = devPtrAmaxdQ; - variant_pack[amax_dK] = devPtrAmaxdK; - variant_pack[amax_dV] = devPtrAmaxdV; - } - if (is_delayed_scaling || (is_current_scaling && !is_O_in_F16)) { - variant_pack[descale_o] = devPtrDescaleO; - } - if (is_delayed_scaling) { - variant_pack[scale_dQ] = devPtrScaledQ; - variant_pack[scale_dK] = devPtrScaledK; - variant_pack[scale_dV] = devPtrScaledV; - } - if (is_mxfp8) { - variant_pack[Q_t] = devPtrQ_t; - variant_pack[K_t] = devPtrK_t; - variant_pack[dO_f16] = devPtrdO_f16; - variant_pack[dO_t] = devPtrdO_t; - variant_pack[descale_q_t] = devPtrDescaleQ_t; - variant_pack[descale_k_t] = devPtrDescaleK_t; - variant_pack[descale_dO_t] = devPtrDescaledO_t; - } - - /* if (cfg.is_bias) { - variant_pack[bias] = devPtrBias; - if ((bias_b == 1) && (bias_h == cfg.num_attn_heads)) { - variant_pack[dBias] = devPtrdBias; - } else { - variant_pack[dBias] = nullptr; - } - } */ - - if (is_padding) { - constexpr size_t nthreads_per_block = 128; - const size_t grid = (b + nthreads_per_block - 1) / nthreads_per_block; - void* devActualSeqlenQ = static_cast(workspace) + plan_workspace_size; - void* devActualSeqlenKV = static_cast(devActualSeqlenQ) + num_bytes_per_seqlen; - cu_seqlens_to_actual_seqlens<<>>( - actual_b, b, static_cast(devPtrcuSeqlensQ), - static_cast(devPtrcuSeqlensKV), static_cast(devActualSeqlenQ), - static_cast(devActualSeqlenKV)); - NVTE_CHECK_CUDA(cudaGetLastError()); - variant_pack[seq_q] = devActualSeqlenQ; - variant_pack[seq_kv] = devActualSeqlenKV; - } - - if (is_ragged_q || is_ragged_kv) { - constexpr size_t nthreads_per_block = 128; - const size_t grid = (b + nthreads_per_block) / nthreads_per_block; - void* devOffsets = - static_cast(workspace) + plan_workspace_size + actual_seqlen_workspace_size; - void* devOffsetsQ = nullptr; - void* devOffsetsO = nullptr; - if (is_ragged_q) { - devOffsetsQ = devOffsets; - devOffsetsO = static_cast(devOffsetsQ) + num_bytes_per_ragged_offset; - } - void* devOffsetsK = nullptr; - void* devOffsetsV = nullptr; - if (is_ragged_kv) { - devOffsetsK = static_cast(devOffsets) + - static_cast(is_ragged_q) * 2 * num_bytes_per_ragged_offset; - devOffsetsV = static_cast(devOffsetsK) + num_bytes_per_ragged_offset; - } - void* devOffsetsS = nullptr; - if (use_ragged_stats) { - devOffsetsS = static_cast(devOffsets) + - (static_cast(is_ragged_q) + static_cast(is_ragged_kv)) * 2 * - num_bytes_per_ragged_offset; - } - cu_seqlens_padded_to_offsets<<>>( - cfg.ragged_offset_mults, actual_b, b, static_cast(devPtrSeqOffsetsQ), - static_cast(devPtrSeqOffsetsKV), ragged_offset_type, devOffsetsQ, devOffsetsK, - devOffsetsV, devOffsetsO, devOffsetsS); - NVTE_CHECK_CUDA(cudaGetLastError()); - if (is_ragged_q) { - variant_pack[offset_q] = devOffsetsQ; - variant_pack[offset_o] = devOffsetsO; - } - if (is_ragged_kv) { - variant_pack[offset_k] = devOffsetsK; - variant_pack[offset_v] = devOffsetsV; - } - if (use_ragged_stats) { - variant_pack[offset_stats] = devOffsetsS; - } - } - - if (is_dropout) { - variant_pack[dropout_seed] = devPtrDropoutSeed; - variant_pack[dropout_offset] = devPtrDropoutOffset; - } - - if (is_softmax_offset) { - variant_pack[softmax_offset] = devPtrSoftmaxOffset; - variant_pack[d_softmax_offset] = devPtrdSoftmaxOffset; - } - - NVTE_CHECK_CUDNN_FE(mha_graph->execute(handle, variant_pack, workspace)); - graph_cache_debug::record_execute(Backend::FP8, Pass::Bwd); - } catch (cudnn_frontend::cudnnException& e) { - NVTE_ERROR(e.what()); - } -} - -} // namespace fused_attn - -using namespace transformer_engine::fused_attn; - -// fused attention FWD FP8 with separate Q, K, V -void fused_attn_fp8_fwd(const FusedAttnConfig& cfg, const Tensor* input_Q, const Tensor* input_K, - const Tensor* input_V, const Tensor* input_SoftmaxOffset, - Tensor* input_output_S, Tensor* output_O, NVTETensorPack* Aux_CTX_Tensors, - const Tensor* cu_seqlens_q, const Tensor* cu_seqlens_kv, - const Tensor* cu_seqlens_q_padded, const Tensor* cu_seqlens_kv_padded, - const Tensor* rng_state, Tensor* workspace, cudaStream_t stream, - cudnnHandle_t handle) { - using namespace transformer_engine; - - const size_t batch = cfg.batch_size; - const size_t num_attn_heads = cfg.num_attn_heads; - const size_t max_seqlen_q = cfg.max_seqlen_q; - const size_t num_tokens_q = cfg.num_tokens_q; - const NVTE_Softmax_Type softmax_type = cfg.softmax_type; - - void *devPtrQ = nullptr, *devPtrK = nullptr, *devPtrV = nullptr; - void *devPtrDescaleQ = nullptr, *devPtrDescaleK = nullptr, *devPtrDescaleV = nullptr; - void *devPtrO = nullptr, *devPtrAmaxO = nullptr, *devPtrScaleO = nullptr; - void *devPtrAmaxS = nullptr, *devPtrScaleS = nullptr, *devPtrDescaleS = nullptr; - devPtrQ = input_Q->data.dptr; - devPtrDescaleQ = input_Q->scale_inv.dptr; - devPtrK = input_K->data.dptr; - devPtrDescaleK = input_K->scale_inv.dptr; - devPtrO = output_O->data.dptr; - if (input_Q->scaling_mode == NVTE_DELAYED_TENSOR_SCALING) { - devPtrV = input_V->data.dptr; - devPtrDescaleV = input_V->scale_inv.dptr; - devPtrScaleO = output_O->scale.dptr; - devPtrAmaxS = input_output_S->amax.dptr; - devPtrScaleS = input_output_S->scale.dptr; - devPtrDescaleS = input_output_S->scale_inv.dptr; - devPtrAmaxO = output_O->amax.dptr; - } else if (input_Q->scaling_mode == NVTE_MXFP8_1D_SCALING) { - devPtrV = input_V->columnwise_data.dptr; - devPtrDescaleV = input_V->columnwise_scale_inv.dptr; - } - void* devPtrSoftmaxOffset = nullptr; - if (softmax_type != NVTE_VANILLA_SOFTMAX) { - devPtrSoftmaxOffset = input_SoftmaxOffset->data.dptr; - } - void* devPtrM = nullptr; - if (Aux_CTX_Tensors->size == 0) { - int i = 0; - Tensor* output_M = convertNVTETensorCheck(Aux_CTX_Tensors->tensors[i++]); - output_M->data.dptr = nullptr; - if (cfg.uses_ragged_stats) { - output_M->data.shape = {num_tokens_q, num_attn_heads, 1}; - } else { - output_M->data.shape = {batch, num_attn_heads, max_seqlen_q, 1}; - } - output_M->data.dtype = DType::kFloat32; - Tensor* output_rng_state = convertNVTETensorCheck(Aux_CTX_Tensors->tensors[i++]); - output_rng_state->data.dptr = nullptr; - output_rng_state->data.shape = {2}; - output_rng_state->data.dtype = DType::kInt64; - if (softmax_type != NVTE_VANILLA_SOFTMAX) { - Tensor* output_softmax_offset = convertNVTETensorCheck(Aux_CTX_Tensors->tensors[i++]); - output_softmax_offset->data.dptr = nullptr; - output_softmax_offset->data.shape = {1, num_attn_heads, 1, 1}; - output_softmax_offset->data.dtype = DType::kFloat32; - } - Aux_CTX_Tensors->size = i; - } else if (Aux_CTX_Tensors->size >= 2) { - int i = 0; - Tensor* output_M = convertNVTETensorCheck(Aux_CTX_Tensors->tensors[i++]); - devPtrM = output_M->data.dptr; - Tensor* output_rng_state = convertNVTETensorCheck(Aux_CTX_Tensors->tensors[i++]); - output_rng_state->data.dptr = rng_state->data.dptr; - if (softmax_type != NVTE_VANILLA_SOFTMAX) { - Tensor* output_softmax_offset = convertNVTETensorCheck(Aux_CTX_Tensors->tensors[i++]); - output_softmax_offset->data.dptr = devPtrSoftmaxOffset; - } - } else { - NVTE_ERROR("Unexpected Aux_CTX_Tensors->size."); - } - - void* devPtrcuSeqlensQ = - reinterpret_cast(reinterpret_cast(cu_seqlens_q->data.dptr)); - void* devPtrcuSeqlensKV = - reinterpret_cast(reinterpret_cast(cu_seqlens_kv->data.dptr)); - void* devPtrSeqOffsetsQ = cu_seqlens_q_padded->data.dptr; - void* devPtrSeqOffsetsKV = cu_seqlens_kv_padded->data.dptr; - void* devPtrDropoutSeed = - reinterpret_cast(reinterpret_cast(rng_state->data.dptr)); - void* devPtrDropoutOffset = - reinterpret_cast(reinterpret_cast(rng_state->data.dptr) + 1); - - size_t workspace_size = 0; - - fused_attn::fused_attn_fp8_fwd_impl( - cfg, devPtrQ, devPtrK, devPtrV, devPtrSoftmaxOffset, devPtrM, devPtrO, devPtrDescaleQ, - devPtrDescaleK, devPtrDescaleV, devPtrDescaleS, devPtrScaleS, devPtrScaleO, devPtrAmaxO, - devPtrAmaxS, devPtrcuSeqlensQ, devPtrcuSeqlensKV, devPtrSeqOffsetsQ, devPtrSeqOffsetsKV, - devPtrDropoutSeed, devPtrDropoutOffset, workspace->data.dptr, &workspace_size, stream, - handle); - - if (workspace_size > 0) { - if (workspace->data.dptr == nullptr) { - workspace->data.shape = {workspace_size}; - workspace->data.dtype = DType::kByte; - return; - } - } else if (workspace_size == 0) { - workspace->data.shape = {1}; - workspace->data.dtype = DType::kByte; - return; - } -} -// fused attention BWD FP8 with separate Q, K, V -void fused_attn_fp8_bwd(const FusedAttnConfig& cfg, const Tensor* input_Q, const Tensor* input_K, - const Tensor* input_V, const Tensor* input_O, const Tensor* input_dO, - const Tensor* input_dO_f16, const Tensor* input_M, const Tensor* input_S, - const Tensor* input_SoftmaxOffset, Tensor* input_output_dP, - const Tensor* output_dQ, const Tensor* output_dK, const Tensor* output_dV, - Tensor* output_dSoftmaxOffset, const Tensor* cu_seqlens_q, - const Tensor* cu_seqlens_kv, const Tensor* cu_seqlens_q_padded, - const Tensor* cu_seqlens_kv_padded, const Tensor* rng_state, - Tensor* workspace, cudaStream_t stream, cudnnHandle_t handle) { - using namespace transformer_engine; - - const NVTE_Softmax_Type softmax_type = cfg.softmax_type; - - void* devPtrQ = input_Q->data.dptr; - void* devPtrK = input_K->data.dptr; - void* devPtrV = input_V->data.dptr; - void* devPtrDescaleQ = input_Q->scale_inv.dptr; - void* devPtrDescaleK = input_K->scale_inv.dptr; - void* devPtrDescaleV = input_V->scale_inv.dptr; - void *devPtrQ_t = nullptr, *devPtrK_t = nullptr, *devPtrDescaleQ_t = nullptr, - *devPtrDescaleK_t = nullptr; - if (input_Q->scaling_mode == NVTE_MXFP8_1D_SCALING) { - devPtrQ_t = input_Q->columnwise_data.dptr; - devPtrDescaleQ_t = input_Q->columnwise_scale_inv.dptr; - devPtrK_t = input_K->columnwise_data.dptr; - devPtrDescaleK_t = input_K->columnwise_scale_inv.dptr; - } - - const DType O_type = input_O->data.dtype; - void* devPtrO = input_O->data.dptr; - void* devPtrDescaleO = nullptr; - if (O_type == DType::kFloat8E4M3 || O_type == DType::kFloat8E5M2) { - devPtrDescaleO = input_O->scale_inv.dptr; - } - void* devPtrdO = input_dO->data.dptr; - void* devPtrDescaledO = input_dO->scale_inv.dptr; - void *devPtrdO_t = nullptr, *devPtrdO_f16 = nullptr, *devPtrDescaledO_t = nullptr; - if (input_dO->scaling_mode == NVTE_MXFP8_1D_SCALING) { - devPtrdO_t = input_dO->columnwise_data.dptr; - devPtrdO_f16 = input_dO_f16->data.dptr; - devPtrDescaledO_t = input_dO->columnwise_scale_inv.dptr; - } - - void* devPtrM = input_M->data.dptr; - - void *devPtrScaleS = nullptr, *devPtrDescaleS = nullptr, *devPtrAmaxdP = nullptr, - *devPtrScaledP = nullptr, *devPtrDescaledP = nullptr; - if (input_Q->scaling_mode == NVTE_DELAYED_TENSOR_SCALING) { - devPtrScaleS = input_S->scale.dptr; - devPtrDescaleS = input_S->scale_inv.dptr; - devPtrAmaxdP = input_output_dP->amax.dptr; - devPtrScaledP = input_output_dP->scale.dptr; - devPtrDescaledP = input_output_dP->scale_inv.dptr; - } - - void* devPtrSoftmaxOffset = nullptr; - void* devPtrdSoftmaxOffset = nullptr; - if (softmax_type != NVTE_VANILLA_SOFTMAX) { - devPtrSoftmaxOffset = input_SoftmaxOffset->data.dptr; - devPtrdSoftmaxOffset = output_dSoftmaxOffset->data.dptr; - } - - void* devPtrdQ = output_dQ->data.dptr; - void* devPtrdK = output_dK->data.dptr; - void* devPtrdV = output_dV->data.dptr; - void *devPtrAmaxdQ = nullptr, *devPtrAmaxdK = nullptr, *devPtrAmaxdV = nullptr, - *devPtrScaledQ = nullptr, *devPtrScaledK = nullptr, *devPtrScaledV = nullptr; - if (input_Q->scaling_mode == NVTE_DELAYED_TENSOR_SCALING) { - devPtrAmaxdQ = output_dQ->amax.dptr; - devPtrAmaxdK = output_dK->amax.dptr; - devPtrAmaxdV = output_dV->amax.dptr; - devPtrScaledQ = output_dQ->scale.dptr; - devPtrScaledK = output_dK->scale.dptr; - devPtrScaledV = output_dV->scale.dptr; - } - - void* devPtrcuSeqlensQ = - reinterpret_cast(reinterpret_cast(cu_seqlens_q->data.dptr)); - void* devPtrcuSeqlensKV = - reinterpret_cast(reinterpret_cast(cu_seqlens_kv->data.dptr)); - void* devPtrSeqOffsetsQ = cu_seqlens_q_padded->data.dptr; - void* devPtrSeqOffsetsKV = cu_seqlens_kv_padded->data.dptr; - void* devPtrDropoutSeed = - reinterpret_cast(reinterpret_cast(rng_state->data.dptr)); - void* devPtrDropoutOffset = - reinterpret_cast(reinterpret_cast(rng_state->data.dptr) + 1); - - size_t workspace_size = 0; - - fused_attn::fused_attn_fp8_bwd_impl( - cfg, devPtrQ, devPtrK, devPtrV, devPtrM, devPtrO, devPtrdO, devPtrSoftmaxOffset, devPtrdQ, - devPtrdK, devPtrdV, devPtrdSoftmaxOffset, devPtrDescaleQ, devPtrDescaleK, devPtrDescaleV, - devPtrDescaleO, devPtrDescaledO, devPtrDescaleS, devPtrDescaledP, devPtrScaleS, devPtrScaledP, - devPtrScaledQ, devPtrScaledK, devPtrScaledV, devPtrAmaxdP, devPtrAmaxdQ, devPtrAmaxdK, - devPtrAmaxdV, devPtrQ_t, devPtrK_t, devPtrdO_f16, devPtrdO_t, devPtrDescaleQ_t, - devPtrDescaleK_t, devPtrDescaledO_t, devPtrcuSeqlensQ, devPtrcuSeqlensKV, devPtrSeqOffsetsQ, - devPtrSeqOffsetsKV, devPtrDropoutSeed, devPtrDropoutOffset, workspace->data.dptr, - &workspace_size, stream, handle); - - if (workspace_size > 0) { - if (workspace->data.dptr == nullptr) { - workspace->data.shape = {workspace_size}; - workspace->data.dtype = DType::kByte; - return; - } - } else if (workspace_size == 0) { - workspace->data.shape = {1}; - workspace->data.dtype = DType::kByte; - return; - } -} - -// Check whether cuDNN can support a given config, per forward/backward pass. -std::string support_verdict_fp8(const FusedAttnConfig& cfg, Pass pass, cudnnHandle_t handle) { - if (pass == Pass::Fwd) { - return fused_attn::support_verdict(cfg, handle); - } - return fused_attn::support_verdict(cfg, handle); -} - -} // namespace transformer_engine diff --git a/transformer_engine/common/fused_attn/fused_attn_fp8.h b/transformer_engine/common/fused_attn/fused_attn_fp8.h deleted file mode 100644 index 9211311df3c..00000000000 --- a/transformer_engine/common/fused_attn/fused_attn_fp8.h +++ /dev/null @@ -1,48 +0,0 @@ -/************************************************************************* - * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. - * - * See LICENSE for license information. - ************************************************************************/ - -/*! \file fused_attn_fp8.h - * \brief Functions for fused attention for FP8 - */ - -#ifndef TRANSFORMER_ENGINE_COMMON_FUSED_ATTN_FUSED_ATTN_FP8_H_ -#define TRANSFORMER_ENGINE_COMMON_FUSED_ATTN_FUSED_ATTN_FP8_H_ - -#include - -#include - -#include "config_and_params.h" -#include "transformer_engine/transformer_engine.h" - -namespace transformer_engine { -// fused attention FWD FP8 with separate Q, K, V -void fused_attn_fp8_fwd(const fused_attn::FusedAttnConfig &cfg, const Tensor *input_Q, - const Tensor *input_K, const Tensor *input_V, - const Tensor *input_SoftmaxOffset, Tensor *input_output_S, Tensor *output_O, - NVTETensorPack *Aux_CTX_Tensors, const Tensor *cu_seqlens_q, - const Tensor *cu_seqlens_kv, const Tensor *cu_seqlens_q_padded, - const Tensor *cu_seqlens_kv_padded, const Tensor *rng_state, - Tensor *workspace, cudaStream_t stream, cudnnHandle_t handle); - -// fused attention BWD FP8 with separate Q, K, V -void fused_attn_fp8_bwd(const fused_attn::FusedAttnConfig &cfg, const Tensor *input_Q, - const Tensor *input_K, const Tensor *input_V, const Tensor *input_O, - const Tensor *input_dO, const Tensor *input_dO_f16, const Tensor *input_M, - const Tensor *input_S, const Tensor *input_SoftmaxOffset, - Tensor *input_output_dP, const Tensor *output_dQ, const Tensor *output_dK, - const Tensor *output_dV, Tensor *output_dSoftmaxOffset, - const Tensor *cu_seqlens_q, const Tensor *cu_seqlens_kv, - const Tensor *cu_seqlens_q_padded, const Tensor *cu_seqlens_kv_padded, - const Tensor *rng_state, Tensor *workspace, cudaStream_t stream, - cudnnHandle_t handle); - -std::string support_verdict_fp8(const fused_attn::FusedAttnConfig &cfg, fused_attn::Pass pass, - cudnnHandle_t handle); - -} // namespace transformer_engine - -#endif // TRANSFORMER_ENGINE_COMMON_FUSED_ATTN_FUSED_ATTN_FP8_H_ diff --git a/transformer_engine/common/fused_attn/graph_cache.h b/transformer_engine/common/fused_attn/graph_cache.h deleted file mode 100644 index b3020f25acf..00000000000 --- a/transformer_engine/common/fused_attn/graph_cache.h +++ /dev/null @@ -1,170 +0,0 @@ -/************************************************************************* - * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. - * - * See LICENSE for license information. - ************************************************************************/ - -// Fused-attention graph cache. The four fused-attention implementation sites, (f16/fp8 + fwd/bwd), -// each create a different graph. They differ in the operations in the graph and the input/output tensors -// that bind to the graph, but the mechanism used for their graph caching, support queries, error messaging, -// and plan building is the same. They all call these functions in this file: get_graph(), support_verdict(), -// and build_plans(). - -#ifndef TRANSFORMER_ENGINE_COMMON_FUSED_ATTN_GRAPH_CACHE_H_ -#define TRANSFORMER_ENGINE_COMMON_FUSED_ATTN_GRAPH_CACHE_H_ - -#include -#include -#include -#include -#include -#include -#include - -#include "../common.h" -#include "../cudnn_utils.h" -#include "config_and_params.h" -#include "graph_cache_debug.h" - -namespace transformer_engine { -namespace fused_attn { - -// An entry in graph cache; contains a cuDNN graph, its input/output tensors, and -// a once_flag that guards its plan build -template -struct CacheEntry { - explicit CacheEntry(GraphAndTensors graph_and_tensors) - : graph_and_tensors(std::move(graph_and_tensors)) {} - - GraphAndTensors graph_and_tensors; - std::once_flag build_plans_once; -}; - -// The graph cache; a process-wide map that maps a normalized FusedAttnConfig to a CacheEntry -template -struct GraphCache { - std::mutex mutex; - std::map>> entries; -}; - -namespace detail { - -// Thrown only when cuDNN declines a graph, so that support_verdict() can tell a config cuDNN does -// not support from a genuine failure -- a violated TE invariant, a CUDA error, a failed allocation -// -- which must surface as an error rather than be reported as an unsupported config. -struct UnsupportedByCudnn : std::runtime_error { - explicit UnsupportedByCudnn(const std::string &reason) : std::runtime_error(reason) {} -}; - -// Query if a cuDNN graph can be supported or not; if so, safely return; if not, throw with -// cuDNN frontend's original error message; times for the four stages are also recorded. -inline void query_support(Backend backend, Pass pass, cudnn_frontend::graph::Graph &graph, - cudnnHandle_t handle) { - using graph_cache_debug::BuildStage; - - auto run = [&](BuildStage stage, const char *call_name, auto &&call) { - const cudnn_frontend::error_t error = - graph_cache_debug::record_time(backend, pass, stage, [&] { return call(); }); - if (error.is_good()) return; - throw UnsupportedByCudnn(error.err_msg.empty() ? std::string(call_name) + " failed." - : error.err_msg); - }; - - run(BuildStage::Validate, "validate", [&] { return graph.validate(); }); - run(BuildStage::BuildOpGraph, "build_operation_graph", - [&] { return graph.build_operation_graph(handle); }); - run(BuildStage::CreatePlans, "create_execution_plans", - [&] { return graph.create_execution_plans({cudnn_frontend::HeurMode_t::A}); }); - run(BuildStage::CheckSupport, "check_support", [&] { return graph.check_support(); }); -} - -// Look up the key in the cache and if -// hit -> record HIT, return the cached entry -// miss -> record MISS, run build() to get a new graph, run query_support() on the new graph, -// if supported, insert it to the cache; if not, throw cuDNN frontend's original -// error message -// -// The cache lookup and insert are guarded by mutex, not the graph builds. Multiple threads -// may build for the same key concurrently, but only the first successful build will be inserted. -template -std::shared_ptr> cache_graph(GraphCache &cache, - const FusedAttnConfig &key, - Backend backend, Pass pass, - cudnnHandle_t handle, BuildFn &&build) { - using graph_cache_debug::LookupResult; - - auto find = [&]() -> std::shared_ptr> { - std::lock_guard lock(cache.mutex); - auto it = cache.entries.find(key); - return it != cache.entries.end() ? it->second : nullptr; - }; - - if (std::shared_ptr> cached = find()) { - graph_cache_debug::record_hit_miss(backend, pass, LookupResult::Hit, key); - return cached; - } - - graph_cache_debug::record_hit_miss(backend, pass, LookupResult::Miss, key); - - auto entry = std::make_shared>(build()); - graph_cache_debug::record_create_graph(backend, pass); - - query_support(backend, pass, *std::get<0>(entry->graph_and_tensors), handle); - - // Concurrent builders of the same key all arrive here, but only the one whose insert wins ends - // up cached, so only that one counts as CACHE_GRAPH. The others discard their graph. - std::shared_ptr> cached; - bool inserted = false; - { - std::lock_guard lock(cache.mutex); - const auto result = cache.entries.insert({key, std::move(entry)}); - cached = result.first->second; - inserted = result.second; - } - if (inserted) graph_cache_debug::record_cache_graph(backend, pass); - return cached; -} - -} // namespace detail - -// Create a cache for each (backend, pass) pair, and either get cached entry or build anew. -template -auto get_graph(const FusedAttnConfig &cfg, cudnnHandle_t handle) { - static GraphCache cache; - cfg.check_derived(); - return detail::cache_graph(cache, cfg.make_cache_key(kPass), kBackend, kPass, handle, - [&] { return kCreateGraphFn(cfg); }); -} - -// Check if cuDNN supports a given config; if yes, return an empty string; if not, return a diagnostic -// string with the reason that get_graph() throws. Anything other than a cuDNN rejection propagates, -// so that a real failure is not reported as an unsupported config. -template -std::string support_verdict(const FusedAttnConfig &cfg, cudnnHandle_t handle) { - try { - get_graph(cfg, handle); - return ""; - } catch (const detail::UnsupportedByCudnn &e) { - // An empty reason would read as supported, so name the site instead. - if (e.what()[0] != '\0') return e.what(); - return std::string("support_verdict<") + backend_name(kBackend) + ", " + pass_name(kPass) + - ">: rejected without a reason."; - } -} - -// Compile the kernels for the graph before execution; once per cache entry; most expensive cuDNN -// frontend call in the pre-execution, preparation process. -template -void build_plans(Backend backend, Pass pass, CacheEntry &entry) { - std::call_once(entry.build_plans_once, [&] { - cudnn_frontend::graph::Graph &graph = *std::get<0>(entry.graph_and_tensors); - graph_cache_debug::record_time(backend, pass, graph_cache_debug::BuildStage::BuildPlans, - [&] { NVTE_CHECK_CUDNN_FE(graph.build_plans()); }); - graph_cache_debug::record_build_plans(backend, pass); - }); -} - -} // namespace fused_attn -} // namespace transformer_engine - -#endif // TRANSFORMER_ENGINE_COMMON_FUSED_ATTN_GRAPH_CACHE_H_ diff --git a/transformer_engine/common/fused_attn/graph_cache_debug.h b/transformer_engine/common/fused_attn/graph_cache_debug.h deleted file mode 100644 index 3ee774bacf2..00000000000 --- a/transformer_engine/common/fused_attn/graph_cache_debug.h +++ /dev/null @@ -1,531 +0,0 @@ -/************************************************************************* - * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. - * - * See LICENSE for license information. - ************************************************************************/ - -// Fused-attention graph cache diagnostics. Written to stderr and prefixed with [FUSED-ATTN-CACHE]. -// Enable with NVTE_FUSED_ATTN_CACHE_DEBUG=[:] at runtime, e.g., "1:all" for level 1 -// on all ranks, and "2:0,3" for level 2 on ranks 0 and 3 only. By default, only rank 0 is enabled -// when diagnostics are on. -// -// Level 0: off (default). -// Level 1: prints a summary at process exit, of cache counters and cuDNN build times. Cache -// events include: HIT/MISS, CREATE_GRAPH, CACHE_GRAPH, BUILD_PLANS, and EXECUTE, and -// counters are accumulated per backend (f16 vs fp8), per pass (forward vs backward), per -// thread. Average, CPU walltimes are also recorded for these cuDNN frontend calls: -// validate(), build_operation_graph(), create_execution_plans(), check_support(), and -// build_plans(). -// Level 2: all Level 1 diagnostics, plus one log message per event as it happens. HIT/MISS -// messages also include the full normalized cache key to identify the config that -// triggered it. -// -// For an example of Level 2 diagnostics, which also includes the summary provided by Level 1, -// please refer to docs/examples/attention/attention.ipynb. - -#ifndef TRANSFORMER_ENGINE_COMMON_FUSED_ATTN_GRAPH_CACHE_DEBUG_H_ -#define TRANSFORMER_ENGINE_COMMON_FUSED_ATTN_GRAPH_CACHE_DEBUG_H_ - -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include - -#include "../util/cuda_runtime.h" -#include "config_and_params.h" - -namespace transformer_engine { -namespace fused_attn { -namespace graph_cache_debug { - -enum class BuildStage { Validate, BuildOpGraph, CreatePlans, CheckSupport, BuildPlans, kCount }; -enum class LookupResult { Miss, Hit }; - -namespace detail { - -// ============================================================================ -// Debug level and rank selection -// ============================================================================ - -// NVTE_FUSED_ATTN_CACHE_DEBUG: 0=off, 1=summary, 2=trace. -inline int debug_level() { - static const int lvl = [] { - const char *e = std::getenv("NVTE_FUSED_ATTN_CACHE_DEBUG"); - if (e == nullptr || e[0] == '\0' || e[0] == '0') return 0; - const int v = std::atoi(e); - return v > 0 ? v : 1; - }(); - return lvl; -} - -// Identify which rank this process is. -inline int launcher_rank() { - static const int rank = []() -> int { - for (const char *var : {"RANK", "LOCAL_RANK", "OMPI_COMM_WORLD_RANK", "SLURM_PROCID"}) { - const char *v = std::getenv(var); - if (v != nullptr && v[0] != '\0') return std::atoi(v); - } - return -1; - }(); - return rank; -} - -// Enable diagnostics for level >= 1; on ranks specified by ":". -inline bool enabled() { - static const bool on = [] { - if (debug_level() < 1) return false; - const int rank = launcher_rank(); - if (rank < 0) return true; - const char *e = std::getenv("NVTE_FUSED_ATTN_CACHE_DEBUG"); - const char *sep = (e != nullptr) ? std::strchr(e, ':') : nullptr; - if (sep == nullptr) return rank == 0; - const std::string list(sep + 1); - if (list == "all") return true; - for (size_t pos = 0; pos <= list.size();) { - const size_t comma = list.find(',', pos); - const std::string tok = - list.substr(pos, comma == std::string::npos ? std::string::npos : comma - pos); - if (!tok.empty() && std::atoi(tok.c_str()) == rank) return true; - if (comma == std::string::npos) break; - pos = comma + 1; - } - return false; - }(); - return on; -} - -// Enable tracing for level >= 2. -inline bool enabled_with_trace() { return enabled() && debug_level() >= 2; } - -// ============================================================================ -// Build-site indexing and labels -// ============================================================================ - -// (backend, pass) pair for each build site: (f16/fp8, fwd/bwd). -constexpr size_t kSiteCount = 4; -inline constexpr size_t site_index(Backend b, Pass p) { - return (b == Backend::F16 ? 0u : 2u) + (p == Pass::Fwd ? 0u : 1u); -} - -// Name the emitting rank. -inline const std::string &rank_tag() { - static const std::string *tag = [] { - const int rank = launcher_rank(); - if (rank < 0) return new std::string(); - return new std::string("rank=" + std::to_string(rank) + " | "); - }(); - return *tag; -} - -// Short thread IDs (0, 1, 2, ...) for convenience, not pthread IDs; -// numbered by assignment order, and incremented sequentially by each thread. -inline unsigned thread_seq_id() { - static std::atomic next{0}; - static thread_local unsigned id = next.fetch_add(1, std::memory_order_relaxed); - return id; -} - -// ============================================================================ -// Event counters -// ============================================================================ - -// Cache event counters, one block per build site (f16/fp8, fwd/bwd). -struct EventCounters { - std::atomic create_graph{0}; - std::atomic cache_graph{0}; - std::atomic build_plans{0}; - std::atomic execute{0}; - std::atomic hit{0}; - std::atomic miss{0}; -}; - -inline EventCounters &counters(Backend b, Pass p) { - static std::array table{}; - return table[site_index(b, p)]; -} - -// One counter block read out into plain values, so the summary can sum blocks for -// its per-backend and all-backends rows. -struct CounterSnapshot { - uint64_t create_graph = 0; - uint64_t cache_graph = 0; - uint64_t build_plans = 0; - uint64_t execute = 0; - uint64_t hit = 0; - uint64_t miss = 0; - - CounterSnapshot &operator+=(const CounterSnapshot &other) { - create_graph += other.create_graph; - cache_graph += other.cache_graph; - build_plans += other.build_plans; - execute += other.execute; - hit += other.hit; - miss += other.miss; - return *this; - } - - bool empty() const { - return (create_graph | cache_graph | build_plans | execute | hit | miss) == 0; - } -}; - -inline CounterSnapshot snapshot(const EventCounters &c) { - CounterSnapshot s; - s.create_graph = c.create_graph.load(std::memory_order_relaxed); - s.cache_graph = c.cache_graph.load(std::memory_order_relaxed); - s.build_plans = c.build_plans.load(std::memory_order_relaxed); - s.execute = c.execute.load(std::memory_order_relaxed); - s.hit = c.hit.load(std::memory_order_relaxed); - s.miss = c.miss.load(std::memory_order_relaxed); - return s; -} - -// Per-thread counters, so the summary can break every column down by thread and backend. -// A thread that drives more than one device accumulates all of its events in one block, so record -// that rather than name a single device; the per-event lines at level 2 carry the live device. -struct ThreadCounters { - unsigned tid = 0; - std::atomic device{-1}; - std::atomic multi_device{false}; - std::array sites; - - void note_device(int current) { - const int seen = device.load(std::memory_order_relaxed); - if (seen < 0) { - device.store(current, std::memory_order_relaxed); - } else if (seen != current) { - multi_device.store(true, std::memory_order_relaxed); - } - } -}; - -// The registry and its mutex are heap-allocated and deliberately never freed. Static destructors -// and atexit handlers run as one sequence in reverse order of construction, and this registry is -// built lazily, so it can be constructed after the summary handler is registered -- and would then -// be destroyed before it runs, leaving the handler to lock a destroyed mutex and walk a destroyed -// vector. Leaking removes the ordering question, at a cost of one mutex and one vector. -inline std::mutex &thread_registry_mutex() { - static std::mutex *m = new std::mutex(); - return *m; -} -inline std::vector &thread_registry() { - static std::vector *v = new std::vector(); - return *v; -} - -// This thread's block, leaked for a related but distinct reason: a worker thread can exit long -// before the process does, while the registry holds a pointer to its block for the exit summary. -inline ThreadCounters &thread_counters() { - static thread_local ThreadCounters *tc = [] { - auto *p = new ThreadCounters(); - p->tid = thread_seq_id(); - { - std::lock_guard lock(thread_registry_mutex()); - thread_registry().push_back(p); - } - return p; - }(); - return *tc; -} - -inline EventCounters &thread_counters(Backend b, Pass p) { - return thread_counters().sites[site_index(b, p)]; -} - -// ============================================================================ -// Build-stage timing buckets -// ============================================================================ - -inline constexpr const char *kStageNames[] = { - "validate", "build_operation_graph", "create_execution_plans", "check_support", "build_plans"}; - -// Totals for one (pass, stage) pair. -struct StageTiming { - std::atomic calls{0}; - std::atomic time_ns{0}; -}; - -// Bucketed by build site, fp8 vs f16. -constexpr size_t kStageBuckets = kSiteCount * static_cast(BuildStage::kCount); -inline StageTiming &stage_timing(Backend b, Pass p, BuildStage s) { - static std::array table{}; - const size_t idx = - site_index(b, p) * static_cast(BuildStage::kCount) + static_cast(s); - return table[idx]; -} - -// ============================================================================ -// Writing to stderr -// ============================================================================ - -// The one place diagnostics reach stderr. -inline void write_stderr(const std::string &text) { - static std::atomic first_line{true}; - if (first_line.exchange(false, std::memory_order_relaxed)) { - const std::string first = "\n" + text; - std::fwrite(first.data(), 1, first.size(), stderr); - } else { - std::fwrite(text.data(), 1, text.size(), stderr); - } - std::fflush(stderr); -} - -// Format one counter block -- one pass of one backend -- as one line. -inline std::string format_counter_line(const char *tid_field, const char *dev_field, - const char *label, const CounterSnapshot &c) { - char buf[512]; - std::snprintf(buf, sizeof(buf), - "[FUSED-ATTN-CACHE] %s%-7s %-7s | %s | hit=%4" PRIu64 ", miss=%4" PRIu64 - ", create_graph=%4" PRIu64 ", cache_graph=%4" PRIu64 ", build_plans=%4" PRIu64 - ", execute=%4" PRIu64 "\n", - rank_tag().c_str(), tid_field, dev_field, label, c.hit, c.miss, c.create_graph, - c.cache_graph, c.build_plans, c.execute); - return std::string(buf); -} - -// ============================================================================ -// Exit summary -// ============================================================================ - -inline constexpr Backend kSummaryBackends[] = {Backend::F16, Backend::FP8}; - -// Names one build site for a summary row. -inline std::string site_label(Backend b, Pass p) { - return std::string(backend_name(b)) + " " + pass_name(p); -} - -// How many backends the run used. -inline size_t active_backend_count() { - size_t active = 0; - for (const Backend b : kSummaryBackends) { - if (!snapshot(counters(b, Pass::Fwd)).empty() || !snapshot(counters(b, Pass::Bwd)).empty()) { - ++active; - } - } - return active; -} - -// Per-thread breakdown, sorted by tid, one row per build site that thread used. -inline void append_thread_rows(std::string &block) { - std::lock_guard lock(thread_registry_mutex()); - std::vector blocks = thread_registry(); - std::sort(blocks.begin(), blocks.end(), - [](const ThreadCounters *a, const ThreadCounters *b) { return a->tid < b->tid; }); - for (const ThreadCounters *tc : blocks) { - char tid_field[16]; - char dev_field[16]; - std::snprintf(tid_field, sizeof(tid_field), "tid=%u", tc->tid); - if (tc->multi_device.load(std::memory_order_relaxed)) { - std::snprintf(dev_field, sizeof(dev_field), "dev=mixed"); - } else { - std::snprintf(dev_field, sizeof(dev_field), "dev=%d", - tc->device.load(std::memory_order_relaxed)); - } - for (const Backend b : kSummaryBackends) { - for (const Pass p : {Pass::Fwd, Pass::Bwd}) { - const CounterSnapshot c = snapshot(tc->sites[site_index(b, p)]); - if (c.empty()) continue; - block += format_counter_line(tid_field, dev_field, site_label(b, p).c_str(), c); - } - } - } -} - -// Totals, printed after the per-thread rows so they read as their sum: one row per build site, -// then one per pass across the backends when the run used more than one. -inline void append_total_rows(std::string &block) { - CounterSnapshot all_fwd; - CounterSnapshot all_bwd; - for (const Backend b : kSummaryBackends) { - for (const Pass p : {Pass::Fwd, Pass::Bwd}) { - const CounterSnapshot c = snapshot(counters(b, p)); - (p == Pass::Fwd ? all_fwd : all_bwd) += c; - if (c.empty()) continue; - block += format_counter_line("tid=all", "dev=all", site_label(b, p).c_str(), c); - } - } - if (active_backend_count() <= 1) return; - for (const Pass p : {Pass::Fwd, Pass::Bwd}) { - const CounterSnapshot &c = (p == Pass::Fwd ? all_fwd : all_bwd); - if (c.empty()) continue; - block += - format_counter_line("tid=all", "dev=all", (std::string("all ") + pass_name(p)).c_str(), c); - } -} - -// Mean time per call for each stage of each build site, skipping stages that no build reached. -inline void append_stage_rows(std::string &block) { - for (const Backend b : kSummaryBackends) { - for (const Pass p : {Pass::Fwd, Pass::Bwd}) { - for (int i = 0; i < static_cast(BuildStage::kCount); ++i) { - const StageTiming &t = stage_timing(b, p, static_cast(i)); - const uint64_t n = t.calls.load(std::memory_order_relaxed); - if (n == 0) continue; - const double total_ms = - static_cast(t.time_ns.load(std::memory_order_relaxed)) / 1e6; - char line[288]; - std::snprintf( - line, sizeof(line), - "[FUSED-ATTN-CACHE] %s%-3s %-3s %-22s | calls=%" PRIu64 " | time=%9.3f ms/call\n", - rank_tag().c_str(), backend_name(b), pass_name(p), kStageNames[i], n, total_ms / n); - block += line; - } - } - } -} - -// Registered at first use, by whichever recorder or stage timer fires first. -inline void register_summary_once() { - static const bool registered = [] { - std::atexit([] { - if (!enabled()) return; - // Built in memory and emitted with one write, so that concurrently-exiting processes stay - // grouped. - const std::string marker = "[FUSED-ATTN-CACHE] " + rank_tag() + "===== summary "; - std::string block = marker + "begin =====\n"; - append_thread_rows(block); - append_total_rows(block); - append_stage_rows(block); - block += marker + "end =====\n"; - write_stderr(block); - }); - return true; - }(); - (void)registered; -} - -// ============================================================================ -// Recorder internals -// ============================================================================ - -// The body every recorder shares: gate, register the exit summary, and add one to `column` in -// both the process-wide block and this thread's. -inline bool record_counter(Backend b, Pass p, std::atomic EventCounters::*column) { - if (!enabled()) return false; - register_summary_once(); - (counters(b, p).*column).fetch_add(1, std::memory_order_relaxed); - (thread_counters(b, p).*column).fetch_add(1, std::memory_order_relaxed); - thread_counters().note_device(cuda::current_device()); - return enabled_with_trace(); -} - -// One event line, from the thread the event happened on, carrying the running totals of the build -// site that raised it. -inline void print_counters(Backend b, Pass p, const char *event) { - const int device = cuda::current_device(); - char label[32]; - char tid_field[16]; - char dev_field[16]; - std::snprintf(label, sizeof(label), "%s %s %-12s", backend_name(b), pass_name(p), event); - std::snprintf(tid_field, sizeof(tid_field), "tid=%u", thread_seq_id()); - std::snprintf(dev_field, sizeof(dev_field), "dev=%d", device); - write_stderr(format_counter_line(tid_field, dev_field, label, snapshot(counters(b, p)))); -} - -// The column a lookup lands in. -inline std::atomic EventCounters::*lookup_column(LookupResult result) { - switch (result) { - case LookupResult::Hit: - return &EventCounters::hit; - case LookupResult::Miss: - break; - } - return &EventCounters::miss; -} - -inline const char *lookup_name(LookupResult result) { - switch (result) { - case LookupResult::Hit: - return "HIT"; - case LookupResult::Miss: - break; - } - return "MISS"; -} - -// Times one stage: clock read in the constructor, accumulated in the destructor. -struct ScopedBuildTimer { - BuildStage stage; - bool on; - Backend backend; - Pass pass; - std::chrono::steady_clock::time_point start; - ScopedBuildTimer(Backend b, Pass p, BuildStage s) : stage(s), on(enabled()), backend(b), pass(p) { - if (!on) return; - register_summary_once(); - start = std::chrono::steady_clock::now(); - } - ~ScopedBuildTimer() { - if (!on) return; - const uint64_t elapsed_ns = - static_cast(std::chrono::duration_cast( - std::chrono::steady_clock::now() - start) - .count()); - StageTiming &t = stage_timing(backend, pass, stage); - t.time_ns.fetch_add(elapsed_ns, std::memory_order_relaxed); - t.calls.fetch_add(1, std::memory_order_relaxed); - } -}; - -} // namespace detail - -// ============================================================================ -// Recorders: everything a call site (f16/fp8 + fwd/bwd) uses. Each recorder -// takes the calling (backend, pass) pair, adds one to that call site's counter, -// and prints a line when the level asks for it. -// ============================================================================ - -inline void record_create_graph(Backend b, Pass p) { - if (detail::record_counter(b, p, &detail::EventCounters::create_graph)) { - detail::print_counters(b, p, "CREATE_GRAPH"); - } -} - -inline void record_cache_graph(Backend b, Pass p) { - if (detail::record_counter(b, p, &detail::EventCounters::cache_graph)) { - detail::print_counters(b, p, "CACHE_GRAPH"); - } -} - -inline void record_build_plans(Backend b, Pass p) { - if (detail::record_counter(b, p, &detail::EventCounters::build_plans)) { - detail::print_counters(b, p, "BUILD_PLANS"); - } -} - -inline void record_execute(Backend b, Pass p) { - if (detail::record_counter(b, p, &detail::EventCounters::execute)) { - detail::print_counters(b, p, "EXECUTE"); - } -} - -inline void record_hit_miss(Backend b, Pass p, LookupResult result, const FusedAttnConfig &key) { - if (!detail::record_counter(b, p, detail::lookup_column(result))) return; - char prefix[128]; - std::snprintf(prefix, sizeof(prefix), - "[FUSED-ATTN-CACHE] %stid=%-3u dev=%-3d | %-3s %-3s %-12s | ", - detail::rank_tag().c_str(), detail::thread_seq_id(), key.device_id, backend_name(b), - pass_name(p), detail::lookup_name(result)); - detail::write_stderr(prefix + key.to_string() + "\n"); -} - -template -inline decltype(auto) record_time(Backend b, Pass p, BuildStage stage, Fn &&fn) { - detail::ScopedBuildTimer scoped(b, p, stage); - return fn(); -} - -} // namespace graph_cache_debug -} // namespace fused_attn -} // namespace transformer_engine - -#endif // TRANSFORMER_ENGINE_COMMON_FUSED_ATTN_GRAPH_CACHE_DEBUG_H_ diff --git a/transformer_engine/common/fused_attn/utils.cu b/transformer_engine/common/fused_attn/utils.cu deleted file mode 100644 index b74dcfe4344..00000000000 --- a/transformer_engine/common/fused_attn/utils.cu +++ /dev/null @@ -1,535 +0,0 @@ -/************************************************************************* - * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. - * - * See LICENSE for license information. - ************************************************************************/ - -#include -#include -#include - -#include "../common.h" -#include "../util/cuda_runtime.h" -#include "transformer_engine/fused_attn.h" -#include "utils.h" - -namespace transformer_engine { -namespace fused_attn { - -using namespace transformer_engine; - -// get matrix strides based on matrix type -void generateMatrixStrides(int64_t b, int64_t h, int64_t s_q, int64_t s_kv, int64_t d, - int64_t *strideA, NVTE_QKV_Layout layout, NVTE_QKV_Matrix matrix) { - constexpr int batch_dim_idx = 0; - constexpr int head_dim_idx = 1; - constexpr int seqlen_dim_idx = 2; - constexpr int hidden_dim_idx = 3; - - constexpr int seqlen_transpose_dim_idx = 3; - constexpr int hidden_transpose_dim_idx = 2; - - constexpr int seqlen_q_dim_idx = 2; - constexpr int seqlen_kv_dim_idx = 3; - - switch (layout) { - case NVTE_QKV_Layout::NVTE_SB3HD: - if ((matrix == NVTE_QKV_Matrix::NVTE_Q_Matrix) || - (matrix == NVTE_QKV_Matrix::NVTE_K_Matrix) || - (matrix == NVTE_QKV_Matrix::NVTE_V_Matrix)) { - strideA[batch_dim_idx] = 3 * h * d; - strideA[head_dim_idx] = d; - strideA[seqlen_dim_idx] = b * 3 * h * d; - strideA[hidden_dim_idx] = 1; - } else if ((matrix == NVTE_QKV_Matrix::NVTE_K_Matrix_Transpose) || - (matrix == NVTE_QKV_Matrix::NVTE_V_Matrix_Transpose)) { - strideA[batch_dim_idx] = 3 * h * d; - strideA[head_dim_idx] = d; - strideA[seqlen_transpose_dim_idx] = b * 3 * h * d; - strideA[hidden_transpose_dim_idx] = 1; - } else if (matrix == NVTE_QKV_Matrix::NVTE_O_Matrix) { - strideA[batch_dim_idx] = h * d; - strideA[head_dim_idx] = d; - strideA[seqlen_dim_idx] = b * h * d; - strideA[hidden_dim_idx] = 1; - } - break; - case NVTE_QKV_Layout::NVTE_SBH3D: - if ((matrix == NVTE_QKV_Matrix::NVTE_Q_Matrix) || - (matrix == NVTE_QKV_Matrix::NVTE_K_Matrix) || - (matrix == NVTE_QKV_Matrix::NVTE_V_Matrix)) { - strideA[batch_dim_idx] = 3 * h * d; - strideA[head_dim_idx] = 3 * d; - strideA[seqlen_dim_idx] = b * 3 * h * d; - strideA[hidden_dim_idx] = 1; - } else if ((matrix == NVTE_QKV_Matrix::NVTE_K_Matrix_Transpose) || - (matrix == NVTE_QKV_Matrix::NVTE_V_Matrix_Transpose)) { - strideA[batch_dim_idx] = 3 * h * d; - strideA[head_dim_idx] = 3 * d; - strideA[seqlen_transpose_dim_idx] = b * 3 * h * d; - strideA[hidden_transpose_dim_idx] = 1; - } else if (matrix == NVTE_QKV_Matrix::NVTE_O_Matrix) { - strideA[batch_dim_idx] = h * d; - strideA[head_dim_idx] = d; - strideA[seqlen_dim_idx] = b * h * d; - strideA[hidden_dim_idx] = 1; - } - break; - case NVTE_QKV_Layout::NVTE_SBHD_SB2HD: - if ((matrix == NVTE_QKV_Matrix::NVTE_K_Matrix) || - (matrix == NVTE_QKV_Matrix::NVTE_V_Matrix)) { - strideA[batch_dim_idx] = 2 * h * d; - strideA[head_dim_idx] = d; - strideA[seqlen_dim_idx] = b * 2 * h * d; - strideA[hidden_dim_idx] = 1; - } else if ((matrix == NVTE_QKV_Matrix::NVTE_K_Matrix_Transpose) || - (matrix == NVTE_QKV_Matrix::NVTE_V_Matrix_Transpose)) { - strideA[batch_dim_idx] = 2 * h * d; - strideA[head_dim_idx] = d; - strideA[seqlen_transpose_dim_idx] = b * 2 * h * d; - strideA[hidden_transpose_dim_idx] = 1; - } else if ((matrix == NVTE_QKV_Matrix::NVTE_Q_Matrix) || - (matrix == NVTE_QKV_Matrix::NVTE_O_Matrix)) { - strideA[batch_dim_idx] = h * d; - strideA[head_dim_idx] = d; - strideA[seqlen_dim_idx] = b * h * d; - strideA[hidden_dim_idx] = 1; - } - break; - case NVTE_QKV_Layout::NVTE_SBHD_SBH2D: - if ((matrix == NVTE_QKV_Matrix::NVTE_K_Matrix) || - (matrix == NVTE_QKV_Matrix::NVTE_V_Matrix)) { - strideA[batch_dim_idx] = 2 * h * d; - strideA[head_dim_idx] = 2 * d; - strideA[seqlen_dim_idx] = b * 2 * h * d; - strideA[hidden_dim_idx] = 1; - } else if ((matrix == NVTE_QKV_Matrix::NVTE_K_Matrix_Transpose) || - (matrix == NVTE_QKV_Matrix::NVTE_V_Matrix_Transpose)) { - strideA[batch_dim_idx] = 2 * h * d; - strideA[head_dim_idx] = 2 * d; - strideA[seqlen_transpose_dim_idx] = b * 2 * h * d; - strideA[hidden_transpose_dim_idx] = 1; - } else if ((matrix == NVTE_QKV_Matrix::NVTE_Q_Matrix) || - (matrix == NVTE_QKV_Matrix::NVTE_O_Matrix)) { - strideA[batch_dim_idx] = h * d; - strideA[head_dim_idx] = d; - strideA[seqlen_dim_idx] = b * h * d; - strideA[hidden_dim_idx] = 1; - } - break; - case NVTE_QKV_Layout::NVTE_SBHD_SBHD_SBHD: - case NVTE_QKV_Layout::NVTE_Paged_KV_SBHD_SBHD_SBHD: - if ((matrix == NVTE_QKV_Matrix::NVTE_Q_Matrix) || - (matrix == NVTE_QKV_Matrix::NVTE_K_Matrix) || - (matrix == NVTE_QKV_Matrix::NVTE_V_Matrix) || - (matrix == NVTE_QKV_Matrix::NVTE_O_Matrix)) { - strideA[batch_dim_idx] = h * d; - strideA[head_dim_idx] = d; - strideA[seqlen_dim_idx] = b * h * d; - strideA[hidden_dim_idx] = 1; - } else if ((matrix == NVTE_QKV_Matrix::NVTE_K_Matrix_Transpose) || - (matrix == NVTE_QKV_Matrix::NVTE_V_Matrix_Transpose)) { - strideA[batch_dim_idx] = h * d; - strideA[head_dim_idx] = d; - strideA[seqlen_transpose_dim_idx] = b * h * d; - strideA[hidden_transpose_dim_idx] = 1; - } - break; - case NVTE_QKV_Layout::NVTE_BS3HD: - case NVTE_QKV_Layout::NVTE_T3HD: - if ((matrix == NVTE_QKV_Matrix::NVTE_Q_Matrix) || - (matrix == NVTE_QKV_Matrix::NVTE_K_Matrix) || - (matrix == NVTE_QKV_Matrix::NVTE_V_Matrix)) { - strideA[batch_dim_idx] = s_q * 3 * h * d; - strideA[head_dim_idx] = d; - strideA[seqlen_dim_idx] = 3 * h * d; - strideA[hidden_dim_idx] = 1; - } else if ((matrix == NVTE_QKV_Matrix::NVTE_K_Matrix_Transpose) || - (matrix == NVTE_QKV_Matrix::NVTE_V_Matrix_Transpose)) { - strideA[batch_dim_idx] = s_q * 3 * h * d; - strideA[head_dim_idx] = d; - strideA[seqlen_transpose_dim_idx] = 3 * h * d; - strideA[hidden_transpose_dim_idx] = 1; - } else if (matrix == NVTE_QKV_Matrix::NVTE_O_Matrix) { - strideA[batch_dim_idx] = s_q * h * d; - strideA[head_dim_idx] = d; - strideA[seqlen_dim_idx] = h * d; - strideA[hidden_dim_idx] = 1; - } - break; - case NVTE_QKV_Layout::NVTE_BSH3D: - case NVTE_QKV_Layout::NVTE_TH3D: - if ((matrix == NVTE_QKV_Matrix::NVTE_Q_Matrix) || - (matrix == NVTE_QKV_Matrix::NVTE_K_Matrix) || - (matrix == NVTE_QKV_Matrix::NVTE_V_Matrix)) { - strideA[batch_dim_idx] = s_q * 3 * h * d; - strideA[head_dim_idx] = 3 * d; - strideA[seqlen_dim_idx] = 3 * h * d; - strideA[hidden_dim_idx] = 1; - } else if ((matrix == NVTE_QKV_Matrix::NVTE_K_Matrix_Transpose) || - (matrix == NVTE_QKV_Matrix::NVTE_V_Matrix_Transpose)) { - strideA[batch_dim_idx] = s_q * 3 * h * d; - strideA[head_dim_idx] = 3 * d; - strideA[seqlen_transpose_dim_idx] = 3 * h * d; - strideA[hidden_transpose_dim_idx] = 1; - } else if (matrix == NVTE_QKV_Matrix::NVTE_O_Matrix) { - strideA[batch_dim_idx] = s_q * h * d; - strideA[head_dim_idx] = d; - strideA[seqlen_dim_idx] = h * d; - strideA[hidden_dim_idx] = 1; - } - break; - case NVTE_QKV_Layout::NVTE_BSHD_BS2HD: - case NVTE_QKV_Layout::NVTE_THD_T2HD: - if ((matrix == NVTE_QKV_Matrix::NVTE_K_Matrix) || - (matrix == NVTE_QKV_Matrix::NVTE_V_Matrix)) { - strideA[batch_dim_idx] = s_kv * 2 * h * d; - strideA[head_dim_idx] = d; - strideA[seqlen_dim_idx] = 2 * h * d; - strideA[hidden_dim_idx] = 1; - } else if ((matrix == NVTE_QKV_Matrix::NVTE_K_Matrix_Transpose) || - (matrix == NVTE_QKV_Matrix::NVTE_V_Matrix_Transpose)) { - strideA[batch_dim_idx] = s_kv * 2 * h * d; - strideA[head_dim_idx] = d; - strideA[seqlen_transpose_dim_idx] = 2 * h * d; - strideA[hidden_transpose_dim_idx] = 1; - } else if ((matrix == NVTE_QKV_Matrix::NVTE_Q_Matrix) || - (matrix == NVTE_QKV_Matrix::NVTE_O_Matrix)) { - strideA[batch_dim_idx] = s_q * h * d; - strideA[head_dim_idx] = d; - strideA[seqlen_dim_idx] = h * d; - strideA[hidden_dim_idx] = 1; - } - break; - case NVTE_QKV_Layout::NVTE_BSHD_BSH2D: - case NVTE_QKV_Layout::NVTE_THD_TH2D: - if ((matrix == NVTE_QKV_Matrix::NVTE_K_Matrix) || - (matrix == NVTE_QKV_Matrix::NVTE_V_Matrix)) { - strideA[batch_dim_idx] = s_kv * 2 * h * d; - strideA[head_dim_idx] = 2 * d; - strideA[seqlen_dim_idx] = 2 * h * d; - strideA[hidden_dim_idx] = 1; - } else if ((matrix == NVTE_QKV_Matrix::NVTE_K_Matrix_Transpose) || - (matrix == NVTE_QKV_Matrix::NVTE_V_Matrix_Transpose)) { - strideA[batch_dim_idx] = s_kv * 2 * h * d; - strideA[head_dim_idx] = 2 * d; - strideA[seqlen_transpose_dim_idx] = 2 * h * d; - strideA[hidden_transpose_dim_idx] = 1; - } else if ((matrix == NVTE_QKV_Matrix::NVTE_Q_Matrix) || - (matrix == NVTE_QKV_Matrix::NVTE_O_Matrix)) { - strideA[batch_dim_idx] = s_q * h * d; - strideA[head_dim_idx] = d; - strideA[seqlen_dim_idx] = h * d; - strideA[hidden_dim_idx] = 1; - } - break; - case NVTE_QKV_Layout::NVTE_BSHD_BSHD_BSHD: - case NVTE_QKV_Layout::NVTE_THD_THD_THD: - case NVTE_QKV_Layout::NVTE_THD_BSHD_BSHD: - case NVTE_QKV_Layout::NVTE_Paged_KV_BSHD_BSHD_BSHD: - case NVTE_QKV_Layout::NVTE_Paged_KV_THD_BSHD_BSHD: - if ((matrix == NVTE_QKV_Matrix::NVTE_Q_Matrix) || - (matrix == NVTE_QKV_Matrix::NVTE_O_Matrix)) { - strideA[batch_dim_idx] = s_q * h * d; - strideA[head_dim_idx] = d; - strideA[seqlen_dim_idx] = h * d; - strideA[hidden_dim_idx] = 1; - } else if ((matrix == NVTE_QKV_Matrix::NVTE_K_Matrix) || - (matrix == NVTE_QKV_Matrix::NVTE_V_Matrix)) { - strideA[batch_dim_idx] = s_kv * h * d; - strideA[head_dim_idx] = d; - strideA[seqlen_dim_idx] = h * d; - strideA[hidden_dim_idx] = 1; - } else if ((matrix == NVTE_QKV_Matrix::NVTE_K_Matrix_Transpose) || - (matrix == NVTE_QKV_Matrix::NVTE_V_Matrix_Transpose)) { - strideA[batch_dim_idx] = s_kv * h * d; - strideA[head_dim_idx] = d; - strideA[seqlen_transpose_dim_idx] = h * d; - strideA[hidden_transpose_dim_idx] = 1; - } - break; - case NVTE_QKV_Layout::NVTE_SBHD_BSHD_BSHD: - case NVTE_QKV_Layout::NVTE_Paged_KV_SBHD_BSHD_BSHD: - if ((matrix == NVTE_QKV_Matrix::NVTE_K_Matrix) || - (matrix == NVTE_QKV_Matrix::NVTE_V_Matrix)) { - strideA[batch_dim_idx] = s_kv * h * d; - strideA[head_dim_idx] = d; - strideA[seqlen_dim_idx] = h * d; - strideA[hidden_dim_idx] = 1; - } else if ((matrix == NVTE_QKV_Matrix::NVTE_K_Matrix_Transpose) || - (matrix == NVTE_QKV_Matrix::NVTE_V_Matrix_Transpose)) { - strideA[batch_dim_idx] = s_kv * h * d; - strideA[head_dim_idx] = d; - strideA[seqlen_transpose_dim_idx] = h * d; - strideA[hidden_transpose_dim_idx] = 1; - } else if ((matrix == NVTE_QKV_Matrix::NVTE_Q_Matrix) || - (matrix == NVTE_QKV_Matrix::NVTE_O_Matrix)) { - strideA[batch_dim_idx] = h * d; - strideA[head_dim_idx] = d; - strideA[seqlen_dim_idx] = b * h * d; - strideA[hidden_dim_idx] = 1; - } - break; - case NVTE_QKV_Layout::NVTE_BSHD_SBHD_SBHD: - case NVTE_QKV_Layout::NVTE_THD_SBHD_SBHD: - case NVTE_QKV_Layout::NVTE_Paged_KV_BSHD_SBHD_SBHD: - case NVTE_QKV_Layout::NVTE_Paged_KV_THD_SBHD_SBHD: - if ((matrix == NVTE_QKV_Matrix::NVTE_K_Matrix) || - (matrix == NVTE_QKV_Matrix::NVTE_V_Matrix)) { - strideA[batch_dim_idx] = h * d; - strideA[head_dim_idx] = d; - strideA[seqlen_dim_idx] = b * h * d; - strideA[hidden_dim_idx] = 1; - } else if ((matrix == NVTE_QKV_Matrix::NVTE_K_Matrix_Transpose) || - (matrix == NVTE_QKV_Matrix::NVTE_V_Matrix_Transpose)) { - strideA[batch_dim_idx] = h * d; - strideA[head_dim_idx] = d; - strideA[seqlen_transpose_dim_idx] = b * h * d; - strideA[hidden_transpose_dim_idx] = 1; - } else if ((matrix == NVTE_QKV_Matrix::NVTE_Q_Matrix) || - (matrix == NVTE_QKV_Matrix::NVTE_O_Matrix)) { - strideA[batch_dim_idx] = s_q * h * d; - strideA[head_dim_idx] = d; - strideA[seqlen_dim_idx] = h * d; - strideA[hidden_dim_idx] = 1; - } - break; - case NVTE_QKV_Layout::NVTE_BHSD_BHSD_BHSD: - if ((matrix == NVTE_QKV_Matrix::NVTE_Q_Matrix) || - (matrix == NVTE_QKV_Matrix::NVTE_O_Matrix)) { - strideA[batch_dim_idx] = h * s_q * d; - strideA[head_dim_idx] = s_q * d; - strideA[seqlen_dim_idx] = d; - strideA[hidden_dim_idx] = 1; - } else if ((matrix == NVTE_QKV_Matrix::NVTE_K_Matrix) || - (matrix == NVTE_QKV_Matrix::NVTE_V_Matrix)) { - strideA[batch_dim_idx] = h * s_kv * d; - strideA[head_dim_idx] = s_kv * d; - strideA[seqlen_dim_idx] = d; - strideA[hidden_dim_idx] = 1; - } else if ((matrix == NVTE_QKV_Matrix::NVTE_K_Matrix_Transpose) || - (matrix == NVTE_QKV_Matrix::NVTE_V_Matrix_Transpose)) { - strideA[batch_dim_idx] = h * s_kv * d; - strideA[head_dim_idx] = s_kv * d; - strideA[seqlen_transpose_dim_idx] = d; - strideA[hidden_transpose_dim_idx] = 1; - } - break; - } - - if (matrix == NVTE_QKV_Matrix::NVTE_S_Matrix) { - strideA[seqlen_kv_dim_idx] = 1; - strideA[seqlen_q_dim_idx] = s_kv; - strideA[head_dim_idx] = s_q * s_kv; - strideA[batch_dim_idx] = h * s_q * s_kv; - } -} - -// convert cu_seqlens to actual_seqlens -__global__ void cu_seqlens_to_actual_seqlens(int64_t actual_b, int64_t max_b, - int32_t const *const q_cu_seqlens, - int32_t const *const kv_cu_seqlens, int32_t *q_seqlens, - int32_t *kv_seqlens) { - size_t tid = blockIdx.x * blockDim.x + threadIdx.x; - if (tid < actual_b) { - q_seqlens[tid] = q_cu_seqlens[tid + 1] - q_cu_seqlens[tid]; - kv_seqlens[tid] = kv_cu_seqlens[tid + 1] - kv_cu_seqlens[tid]; - } else if (tid < max_b) { - q_seqlens[tid] = 0; - kv_seqlens[tid] = 0; - } -} - -// convert cu_seqlens_padded to offsets -template -__device__ void cu_seqlens_padded_to_offsets_impl( - const RaggedOffsetMultipliers &mults, int64_t actual_b, int64_t max_b, - const int32_t *cu_seqlens_q_padded, const int32_t *cu_seqlens_kv_padded, OFFSETS_T *offsets_q, - OFFSETS_T *offsets_k, OFFSETS_T *offsets_v, OFFSETS_T *offsets_o, OFFSETS_T *offsets_s) { - size_t tid = blockIdx.x * blockDim.x + threadIdx.x; - auto cu_seqlens_id = min(tid, actual_b); - if (tid <= max_b) { - if (offsets_s != nullptr) { - offsets_s[tid] = mults.stats * cu_seqlens_q_padded[cu_seqlens_id]; - } - if (offsets_q != nullptr && offsets_o != nullptr) { - offsets_q[tid] = mults.q * cu_seqlens_q_padded[cu_seqlens_id]; - offsets_o[tid] = mults.o * cu_seqlens_q_padded[cu_seqlens_id]; - } - if (offsets_k != nullptr && offsets_v != nullptr) { - const int32_t *cu_seqlens_kv_src = - mults.kv_from_q ? cu_seqlens_q_padded : cu_seqlens_kv_padded; - offsets_k[tid] = mults.k * cu_seqlens_kv_src[cu_seqlens_id]; - offsets_v[tid] = mults.v * cu_seqlens_kv_src[cu_seqlens_id]; - } - } -} - -__global__ void cu_seqlens_padded_to_offsets(RaggedOffsetMultipliers mults, int64_t actual_b, - int64_t max_b, const int32_t *cu_seqlens_q_padded, - const int32_t *cu_seqlens_kv_padded, - DType offset_dtype, void *offsets_q, void *offsets_k, - void *offsets_v, void *offsets_o, void *offsets_s) { - if (offset_dtype == DType::kInt32) { - cu_seqlens_padded_to_offsets_impl( - mults, actual_b, max_b, cu_seqlens_q_padded, cu_seqlens_kv_padded, - reinterpret_cast(offsets_q), reinterpret_cast(offsets_k), - reinterpret_cast(offsets_v), reinterpret_cast(offsets_o), - reinterpret_cast(offsets_s)); - } else { - assert(offset_dtype == DType::kInt64 && "expect int64"); - cu_seqlens_padded_to_offsets_impl( - mults, actual_b, max_b, cu_seqlens_q_padded, cu_seqlens_kv_padded, - reinterpret_cast(offsets_q), reinterpret_cast(offsets_k), - reinterpret_cast(offsets_v), reinterpret_cast(offsets_o), - reinterpret_cast(offsets_s)); - } -} - -DType get_ragged_offset_dtype(NVTE_QKV_Layout_Group layout_group, int64_t num_attn_heads, - int64_t num_gqa_groups, int64_t max_seqlen_q, int64_t max_seqlen_kv, - int64_t head_dim_qk, int64_t head_dim_v) { - std::array offsets_qkvo{}; - switch (layout_group) { - case NVTE_QKV_Layout_Group::NVTE_HD_HD_HD: - case NVTE_QKV_Layout_Group::NVTE_Paged_KV_HD_HD_HD: - offsets_qkvo[0] = num_attn_heads * head_dim_qk * max_seqlen_q; - offsets_qkvo[1] = num_gqa_groups * head_dim_qk * max_seqlen_kv; - offsets_qkvo[2] = num_gqa_groups * head_dim_v * max_seqlen_kv; - break; - case NVTE_QKV_Layout_Group::NVTE_3HD: - case NVTE_QKV_Layout_Group::NVTE_H3D: - offsets_qkvo[0] = 3 * num_attn_heads * head_dim_qk * max_seqlen_q; - offsets_qkvo[1] = offsets_qkvo[0]; - offsets_qkvo[2] = offsets_qkvo[0]; - break; - case NVTE_QKV_Layout_Group::NVTE_HD_2HD: - case NVTE_QKV_Layout_Group::NVTE_HD_H2D: - offsets_qkvo[0] = num_attn_heads * head_dim_qk * max_seqlen_q; - offsets_qkvo[1] = 2 * num_gqa_groups * head_dim_qk * max_seqlen_kv; - offsets_qkvo[2] = offsets_qkvo[1]; - break; - } - - offsets_qkvo[3] = num_attn_heads * head_dim_qk * max_seqlen_q; - - size_t max_offset = *std::max_element(offsets_qkvo.begin(), offsets_qkvo.end()); - if (max_offset > std::numeric_limits::max()) { - return DType::kInt64; - } - - return DType::kInt32; -} - -// quantize batch size -size_t get_max_batch_size(size_t batch_size) { - if (batch_size == 0) return 0; // guard: log2(0) = -inf, casting to size_t is UB - size_t max_b = batch_size; - size_t log2_b = ceil(log2(batch_size)); - // batch size is expected to be 10s-100s - // b = 1, ..., 32 -> max_b = 32 - // b = 33, ..., 512 -> max_b = next power of 2 - // b = 513, ... -> max_b = increment by 512 - if (log2_b <= 5) { - max_b = 32; - } else if (log2_b <= 9) { - max_b = pow(2, log2_b); - } else { - max_b = (batch_size + 511) / 512 * 512; - } - return max_b; -} - -// quantize token count -size_t get_max_tokens(size_t num_tokens) { - if (num_tokens == 0) return 0; // guard: log2(0) = -inf, casting to size_t is UB - // token count is expected to be 1k's-100k's - // t = 0, ..., 1024 -> max_t = 1024 - // t = 1025, ..., 32k -> max_t = next power of 2 - // t = 32k+1, ... -> max_t = increment by 32k - size_t log2_t = ceil(log2(num_tokens)); - size_t max_t = 0; - if (log2_t <= 10) { - max_t = 1024; - } else if (log2_t <= 15) { - max_t = pow(2, log2_t); - } else { - max_t = (num_tokens + 32767) / 32768 * 32768; - } - return max_t; -} - -__global__ void populate_rng_state_kernel(int64_t *rng_state_dst, const int64_t *const seed, - int64_t offset) { - int tid = blockIdx.x * blockDim.x + threadIdx.x; - if (tid > 0) return; - rng_state_dst[0] = seed[0]; - rng_state_dst[1] = offset; -} - -__global__ void get_runtime_num_segments_kernel(int32_t *cu_seqlen, size_t len, uint32_t *out) { - int tid = blockDim.x * blockIdx.x + threadIdx.x; - if (tid >= len) return; - - if (cu_seqlen[tid] > 0) { - // atomicAdd only support 32 bits dtype - atomicAdd(out, 1); - } -} - -void PopulateRngStateAsync(void *rng_state_dst, const void *seed, size_t q_max_seqlen, - size_t kv_max_seqlen, NVTE_Fused_Attn_Backend backend, - cudaStream_t stream) { - size_t increment = 0; - if (backend == NVTE_Fused_Attn_Backend::NVTE_F16_arbitrary_seqlen) { - increment = 16; - } else { - constexpr int threads_per_cta = 128; - increment = (q_max_seqlen * kv_max_seqlen + threads_per_cta - 1) / threads_per_cta; - } - auto offset = FusedAttnOffsetManager::Instance().GetAndUpdateOffset(increment); - populate_rng_state_kernel<<<1, 1, 0, stream>>>(reinterpret_cast(rng_state_dst), - reinterpret_cast(seed), offset); - NVTE_CHECK_CUDA(cudaGetLastError()); -} - -uint32_t GetRuntimeNumSegments(void *cu_seqlen, void *workspace, size_t len, cudaStream_t stream) { - // workspace size requires 4 bytes - uint32_t *dout = static_cast(workspace); - uint32_t hout{}; - NVTE_CHECK_CUDA(cudaMemsetAsync(dout, 0, sizeof(uint32_t), stream)); - constexpr int threads = 128; - const int blocks = (len - 1) / threads + 1; - get_runtime_num_segments_kernel<<>>(static_cast(cu_seqlen), - len, dout); - NVTE_CHECK_CUDA(cudaGetLastError()); - NVTE_CHECK_CUDA(cudaMemcpyAsync(&hout, dout, sizeof(uint32_t), cudaMemcpyDeviceToHost, stream)); - NVTE_CHECK_CUDA(cudaStreamSynchronize(stream)); - return hout; -} - -__global__ void extract_seed_and_offset(int64_t *rng_state_ptr, bool captured, int64_t *seed_ptr, - uint64_t seed_val, int64_t *offset_ptr, uint64_t offset_val, - uint32_t offset_intragraph) { - if (captured) { - rng_state_ptr[0] = *seed_ptr; - rng_state_ptr[1] = static_cast(*offset_ptr + static_cast(offset_intragraph)); - } else { - rng_state_ptr[0] = static_cast(seed_val); - rng_state_ptr[1] = static_cast(offset_val); - } -} - -} // namespace fused_attn -} // namespace transformer_engine - -void nvte_extract_seed_and_offset(int64_t *rng_state_ptr, int captured, int64_t *seed_ptr, - uint64_t seed_val, int64_t *offset_ptr, uint64_t offset_val, - uint32_t offset_intragraph, cudaStream_t stream) { - NVTE_API_CALL(nvte_extract_seed_and_offset); - using namespace transformer_engine; - - fused_attn::extract_seed_and_offset<<<1, 1, 0, stream>>>( - rng_state_ptr, captured, seed_ptr, seed_val, offset_ptr, offset_val, offset_intragraph); - NVTE_CHECK_CUDA(cudaGetLastError()); -} diff --git a/transformer_engine/common/fused_attn/utils.h b/transformer_engine/common/fused_attn/utils.h deleted file mode 100644 index 237507f30b9..00000000000 --- a/transformer_engine/common/fused_attn/utils.h +++ /dev/null @@ -1,308 +0,0 @@ -/************************************************************************* - * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. - * - * See LICENSE for license information. - ************************************************************************/ - -#ifndef TRANSFORMER_ENGINE_COMMON_FUSED_ATTN_UTILS_H_ -#define TRANSFORMER_ENGINE_COMMON_FUSED_ATTN_UTILS_H_ - -#include - -#include "../common.h" -#include "transformer_engine/fused_attn.h" -#include "transformer_engine/transformer_engine.h" - -namespace transformer_engine { -namespace fused_attn { - -using namespace transformer_engine; - -enum NVTE_QKV_Matrix { - NVTE_Q_Matrix = 0, // queries - NVTE_K_Matrix = 1, // keys - NVTE_K_Matrix_Transpose = 2, // keys transposed - NVTE_V_Matrix = 3, // values - NVTE_V_Matrix_Transpose = 4, // values transposed - NVTE_S_Matrix = 5, // output of GEMM1 - NVTE_O_Matrix = 6, // final output -}; - -// Padded sizes for MXFP8 layout (s_q/s_kv/d_qk/d_v and their scaled dimensions) -struct MXFP8PaddedSizes { - int64_t s_q_padded; - int64_t s_kv_padded; - int64_t s_q_scale; - int64_t s_kv_scale; - int64_t s_q_scale_padded; - int64_t s_kv_scale_padded; - int64_t d_qk_padded; - int64_t d_v_padded; - int64_t d_qk_scale; - int64_t d_v_scale; - int64_t d_qk_scale_padded; - int64_t d_v_scale_padded; -}; - -// Pad s and d for MXFP8 quantization -inline MXFP8PaddedSizes pad_s_d_for_mxfp8(int64_t s_q, int64_t s_kv, int64_t d_qk, int64_t d_v) { - constexpr int64_t block_size = 32; - MXFP8PaddedSizes p; - p.s_q_padded = DIVUP_TO_MULTIPLE(s_q, 128); - p.s_kv_padded = DIVUP_TO_MULTIPLE(s_kv, 128); - p.s_q_scale = DIVUP(s_q, block_size); - p.s_kv_scale = DIVUP(s_kv, block_size); - p.s_q_scale_padded = DIVUP_TO_MULTIPLE(p.s_q_scale, 4); - p.s_kv_scale_padded = DIVUP_TO_MULTIPLE(p.s_kv_scale, 4); - p.d_qk_padded = DIVUP_TO_MULTIPLE(d_qk, 128); - p.d_v_padded = DIVUP_TO_MULTIPLE(d_v, 128); - p.d_qk_scale = DIVUP(d_qk, block_size); - p.d_v_scale = DIVUP(d_v, block_size); - p.d_qk_scale_padded = DIVUP_TO_MULTIPLE(p.d_qk_scale, 4); - p.d_v_scale_padded = DIVUP_TO_MULTIPLE(p.d_v_scale, 4); - return p; -} - -// Get matrix strides for a 4D tensor [batch_size, num_heads, sequence_len, head_dim] given a QKV format. -// strides must point to at least 4 int64_t elements. -inline void generateMatrixStridesWithFormat(int64_t b, int64_t h, int64_t s, int64_t d, - int64_t *strides, NVTE_QKV_Format format) { - constexpr int b_dim = 0; - constexpr int h_dim = 1; - constexpr int s_dim = 2; - constexpr int d_dim = 3; - - switch (format) { - case NVTE_QKV_Format::NVTE_BSHD: - case NVTE_QKV_Format::NVTE_THD: - strides[b_dim] = s * h * d; - strides[h_dim] = d; - strides[s_dim] = h * d; - strides[d_dim] = 1; - break; - case NVTE_QKV_Format::NVTE_SBHD: - strides[b_dim] = h * d; - strides[h_dim] = d; - strides[s_dim] = b * h * d; - strides[d_dim] = 1; - break; - case NVTE_QKV_Format::NVTE_BHSD: - strides[b_dim] = h * s * d; - strides[h_dim] = s * d; - strides[s_dim] = d; - strides[d_dim] = 1; - break; - default: - NVTE_CHECK(false, "Invalid format."); - break; - } -} - -// get matrix strides based on layout and matrix type -inline void generateMatrixStridesWithLayout(int64_t b, int64_t h, int64_t hg, int64_t s_q, - int64_t s_kv, int64_t d_qk, int64_t d_v, - int64_t *q_strides, int64_t *k_strides, - int64_t *v_strides, NVTE_QKV_Layout layout) { - constexpr int b_dim = 0; - constexpr int h_dim = 1; - constexpr int s_dim = 2; - constexpr int d_dim = 3; - const NVTE_QKV_Format q_format = nvte_get_q_format(layout); - const NVTE_QKV_Format kv_format = nvte_get_kv_format(layout); - - switch (layout) { - case NVTE_QKV_Layout::NVTE_SB3HD: - q_strides[b_dim] = 3 * h * d_qk; - q_strides[h_dim] = d_qk; - q_strides[s_dim] = b * 3 * h * d_qk; - q_strides[d_dim] = 1; - for (int i = 0; i < 4; i++) { - k_strides[i] = v_strides[i] = q_strides[i]; - } - break; - case NVTE_QKV_Layout::NVTE_SBH3D: - q_strides[b_dim] = 3 * h * d_qk; - q_strides[h_dim] = 3 * d_qk; - q_strides[s_dim] = b * 3 * h * d_qk; - q_strides[d_dim] = 1; - for (int i = 0; i < 4; i++) { - k_strides[i] = v_strides[i] = q_strides[i]; - } - break; - case NVTE_QKV_Layout::NVTE_SBHD_SB2HD: - generateMatrixStridesWithFormat(b, h, s_q, d_qk, q_strides, q_format); - k_strides[b_dim] = 2 * hg * d_qk; - k_strides[h_dim] = d_qk; - k_strides[s_dim] = b * 2 * hg * d_qk; - k_strides[d_dim] = 1; - for (int i = 0; i < 4; i++) { - v_strides[i] = k_strides[i]; - } - break; - case NVTE_QKV_Layout::NVTE_SBHD_SBH2D: - generateMatrixStridesWithFormat(b, h, s_q, d_qk, q_strides, q_format); - k_strides[b_dim] = 2 * hg * d_qk; - k_strides[h_dim] = 2 * d_qk; - k_strides[s_dim] = b * 2 * hg * d_qk; - k_strides[d_dim] = 1; - for (int i = 0; i < 4; i++) { - v_strides[i] = k_strides[i]; - } - break; - case NVTE_QKV_Layout::NVTE_BS3HD: - case NVTE_QKV_Layout::NVTE_T3HD: - q_strides[b_dim] = s_q * 3 * h * d_qk; - q_strides[h_dim] = d_qk; - q_strides[s_dim] = 3 * h * d_qk; - q_strides[d_dim] = 1; - for (int i = 0; i < 4; i++) { - k_strides[i] = v_strides[i] = q_strides[i]; - } - break; - case NVTE_QKV_Layout::NVTE_BSH3D: - case NVTE_QKV_Layout::NVTE_TH3D: - q_strides[b_dim] = s_q * 3 * h * d_qk; - q_strides[h_dim] = 3 * d_qk; - q_strides[s_dim] = 3 * h * d_qk; - q_strides[d_dim] = 1; - for (int i = 0; i < 4; i++) { - k_strides[i] = v_strides[i] = q_strides[i]; - } - break; - case NVTE_QKV_Layout::NVTE_BSHD_BS2HD: - case NVTE_QKV_Layout::NVTE_THD_T2HD: - generateMatrixStridesWithFormat(b, h, s_q, d_qk, q_strides, q_format); - k_strides[b_dim] = s_kv * 2 * hg * d_qk; - k_strides[h_dim] = d_qk; - k_strides[s_dim] = 2 * hg * d_qk; - k_strides[d_dim] = 1; - for (int i = 0; i < 4; i++) { - v_strides[i] = k_strides[i]; - } - break; - case NVTE_QKV_Layout::NVTE_BSHD_BSH2D: - case NVTE_QKV_Layout::NVTE_THD_TH2D: - generateMatrixStridesWithFormat(b, h, s_q, d_qk, q_strides, q_format); - k_strides[b_dim] = s_kv * 2 * hg * d_qk; - k_strides[h_dim] = 2 * d_qk; - k_strides[s_dim] = 2 * hg * d_qk; - k_strides[d_dim] = 1; - for (int i = 0; i < 4; i++) { - v_strides[i] = k_strides[i]; - } - break; - case NVTE_QKV_Layout::NVTE_SBHD_SBHD_SBHD: - case NVTE_QKV_Layout::NVTE_Paged_KV_SBHD_SBHD_SBHD: - case NVTE_QKV_Layout::NVTE_BSHD_BSHD_BSHD: - case NVTE_QKV_Layout::NVTE_THD_THD_THD: - case NVTE_QKV_Layout::NVTE_THD_BSHD_BSHD: - case NVTE_QKV_Layout::NVTE_Paged_KV_BSHD_BSHD_BSHD: - case NVTE_QKV_Layout::NVTE_Paged_KV_THD_BSHD_BSHD: - case NVTE_QKV_Layout::NVTE_SBHD_BSHD_BSHD: - case NVTE_QKV_Layout::NVTE_Paged_KV_SBHD_BSHD_BSHD: - case NVTE_QKV_Layout::NVTE_BSHD_SBHD_SBHD: - case NVTE_QKV_Layout::NVTE_THD_SBHD_SBHD: - case NVTE_QKV_Layout::NVTE_Paged_KV_BSHD_SBHD_SBHD: - case NVTE_QKV_Layout::NVTE_Paged_KV_THD_SBHD_SBHD: - case NVTE_QKV_Layout::NVTE_BHSD_BHSD_BHSD: - generateMatrixStridesWithFormat(b, h, s_q, d_qk, q_strides, q_format); - generateMatrixStridesWithFormat(b, hg, s_kv, d_qk, k_strides, kv_format); - generateMatrixStridesWithFormat(b, hg, s_kv, d_v, v_strides, kv_format); - break; - default: - NVTE_CHECK(false, "Invalid layout."); - break; - } -} - -void generateMatrixStrides(int64_t b, int64_t h, int64_t s_q, int64_t s_kv, int64_t d, - int64_t *strideA, NVTE_QKV_Layout layout, NVTE_QKV_Matrix matrix); - -// Per-tensor scale factors relating cu_seqlens_padded (token units) to tensor-element -// ragged offsets, as a function of the QKV layout group. Single source of truth shared -// by the cu_seqlens_padded_to_offsets conversion kernel and the direct-seqlens path -// (which passes them to cuDNN as ragged offset multipliers). -struct RaggedOffsetMultipliers { - RaggedOffsetMultipliers() = default; - - RaggedOffsetMultipliers(NVTE_QKV_Layout_Group layout_group, int64_t h, int64_t hg, int64_t d_qk, - int64_t d_v) - : q(h * d_qk), k(hg * d_qk), v(hg * d_v), o(h * d_v), stats(h), kv_from_q(false) { - switch (layout_group) { - case NVTE_QKV_Layout_Group::NVTE_3HD: - case NVTE_QKV_Layout_Group::NVTE_H3D: - q = k = v = 3 * h * d_qk; - kv_from_q = true; - break; - case NVTE_QKV_Layout_Group::NVTE_HD_2HD: - case NVTE_QKV_Layout_Group::NVTE_HD_H2D: - k = v = 2 * hg * d_qk; - break; - default: - break; - } - } - - int64_t q = 0; - int64_t k = 0; - int64_t v = 0; - int64_t o = 0; - int64_t stats = 0; - // K/V offsets scale the Q-side cu_seqlens_padded (interleaved QKV layouts) - bool kv_from_q = false; -}; - -__global__ void cu_seqlens_to_actual_seqlens(int64_t actual_b, int64_t max_b, - int32_t const *const q_cu_seqlens, - int32_t const *const kv_cu_seqlens, int32_t *q_seqlens, - int32_t *kv_seqlens); - -__global__ void cu_seqlens_padded_to_offsets(RaggedOffsetMultipliers mults, int64_t actual_b, - int64_t max_b, const int32_t *cu_seqlens_q_padded, - const int32_t *cu_seqlens_kv_padded, - DType offset_dtype, void *offsets_q, void *offsets_k, - void *offsets_v, void *offsets_o, void *offsets_s); - -DType get_ragged_offset_dtype(NVTE_QKV_Layout_Group layout_group, int64_t num_attn_heads, - int64_t num_gqa_groups, int64_t max_seqlen_q, int64_t max_seqlen_kv, - int64_t head_dim_qk, int64_t head_dim_v); - -size_t get_max_batch_size(size_t batch_size); -size_t get_max_tokens(size_t num_tokens); - -class FusedAttnOffsetManager { - public: - static FusedAttnOffsetManager &Instance() { - static thread_local FusedAttnOffsetManager instance; - return instance; - } - - size_t GetAndUpdateOffset(size_t increment) { - size_t ret = offset_; - offset_ += increment; - return ret; - } - - FusedAttnOffsetManager(FusedAttnOffsetManager const &) = delete; - void operator=(FusedAttnOffsetManager const &) = delete; - - private: - FusedAttnOffsetManager() {} - size_t offset_ = 0; -}; - -__global__ void populate_rng_state_kernel(int64_t *rng_state_dst, const int64_t *const seed, - int64_t offset); - -__global__ void get_runtime_num_segments_kernel(int32_t *cu_seqlen, size_t len, uint32_t *out); - -void PopulateRngStateAsync(void *rng_state_dst, const void *const seed, size_t q_max_seqlen, - size_t kv_max_seqlen, NVTE_Fused_Attn_Backend backend, - cudaStream_t stream); - -uint32_t GetRuntimeNumSegments(void *cu_seqlen, void *workspace, size_t len, cudaStream_t stream); - -} // namespace fused_attn -} // namespace transformer_engine - -#endif // TRANSFORMER_ENGINE_COMMON_FUSED_ATTN_UTILS_H_ diff --git a/transformer_engine/common/include/transformer_engine/fused_attn.h b/transformer_engine/common/include/transformer_engine/fused_attn.h index e889c28219c..5c2474700c9 100644 --- a/transformer_engine/common/include/transformer_engine/fused_attn.h +++ b/transformer_engine/common/include/transformer_engine/fused_attn.h @@ -5,7 +5,7 @@ ************************************************************************/ /*! \file fused_attn.h - * \brief Enums and functions for fused attention. + * \brief Attention enums and framework support functions. */ #ifndef TRANSFORMER_ENGINE_FUSED_ATTN_H_ @@ -194,501 +194,6 @@ NVTE_QKV_Format nvte_get_q_format(NVTE_QKV_Layout qkv_layout); */ NVTE_QKV_Format nvte_get_kv_format(NVTE_QKV_Layout qkv_layout); -/*! \brief Opaque fused-attention configuration handle. */ -typedef void *NVTEFusedAttnConfig; - -/*! \enum NVTEFusedAttnConfigAttribute - * \brief Attributes for `NVTEFusedAttnConfig`. - * - * This enum is used to index the `FusedAttnConfig` struct. The order of its fields matches - * that of the declaration fields and the `attr_sizes` array of `FusedAttnConfig`. New fields - * may be appended at the end, and existing fields are never reordered, removed, or resized. - */ -enum NVTEFusedAttnConfigAttribute { - // Basic attention settings - kNVTEFusedAttnConfigIsTraining = 0, - kNVTEFusedAttnConfigDeterministic, - kNVTEFusedAttnConfigCudaGraph, - kNVTEFusedAttnConfigReturnMaxLogit, - kNVTEFusedAttnConfigAttnMaskType, - kNVTEFusedAttnConfigBiasType, - kNVTEFusedAttnConfigWindowSizeLeft, - kNVTEFusedAttnConfigWindowSizeRight, - kNVTEFusedAttnConfigBottomRightDiagonal, - kNVTEFusedAttnConfigSoftmaxType, - kNVTEFusedAttnConfigScalingMode, - kNVTEFusedAttnConfigDropout, - kNVTEFusedAttnConfigAttnScale, - // Tensor types - kNVTEFusedAttnConfigQKVDtype, - kNVTEFusedAttnConfigODtype, - kNVTEFusedAttnConfigDODtype, - kNVTEFusedAttnConfigDQKVDtype, - // Tensor layouts - kNVTEFusedAttnConfigQKVLayout, - kNVTEFusedAttnConfigOFormat, - kNVTEFusedAttnConfigDOFormat, - kNVTEFusedAttnConfigDQKVLayout, - kNVTEFusedAttnConfigQKVScaleInvFormat, - kNVTEFusedAttnConfigDOScaleInvFormat, - // Tensor dimensions - kNVTEFusedAttnConfigBatchSize, - kNVTEFusedAttnConfigNumAttnHeads, - kNVTEFusedAttnConfigNumGQAGroups, - kNVTEFusedAttnConfigHeadDimQK, - kNVTEFusedAttnConfigHeadDimV, - kNVTEFusedAttnConfigMaxSeqlenQ, - kNVTEFusedAttnConfigMaxSeqlenKV, - kNVTEFusedAttnConfigNumTokensQ, - kNVTEFusedAttnConfigNumTokensKV, - // Paged KV dimensions - kNVTEFusedAttnConfigNumPagesK, - kNVTEFusedAttnConfigNumPagesV, - kNVTEFusedAttnConfigPageSizeK, - kNVTEFusedAttnConfigPageSizeV, - kNVTEFusedAttnConfigMaxPagesPerSeqK, - kNVTEFusedAttnConfigMaxPagesPerSeqV, - // Bias dimensions - kNVTEFusedAttnConfigBiasBatchSize, - kNVTEFusedAttnConfigBiasNumHeads, - kNVTEFusedAttnConfigBiasSeqlenQ, - kNVTEFusedAttnConfigBiasSeqlenKV, - // Number of attributes - kNVTEFusedAttnConfigNumAttributes -}; - -/*! \brief Create a fused-attention configuration. */ -NVTEFusedAttnConfig nvte_create_fused_attn_config(void); - -/*! \brief Destroy a fused-attention configuration. */ -void nvte_destroy_fused_attn_config(NVTEFusedAttnConfig config); - -/*! \brief Query an attribute in a fused-attention configuration. */ -void nvte_get_fused_attn_config_attribute(NVTEFusedAttnConfig config, - NVTEFusedAttnConfigAttribute attr, void *buf, - size_t size_in_bytes, size_t *size_written); - -/*! \brief Set an attribute in a fused-attention configuration. */ -void nvte_set_fused_attn_config_attribute(NVTEFusedAttnConfig config, - NVTEFusedAttnConfigAttribute attr, const void *buf, - size_t size_in_bytes); - -/*! \brief Opaque fused-attention forward-parameter handle. */ -typedef void *NVTEFusedAttnFwdParams; - -/*! \enum NVTEFusedAttnFwdParamsAttribute - * \brief Attributes for `NVTEFusedAttnFwdParams`. - * - * This enum is used to index the `FusedAttnFwdParams` struct. The order of its fields matches - * that of the declaration fields and the `attr_sizes` array of `FusedAttnFwdParams`. New fields - * may be appended at the end, and existing fields are never reordered, removed, or resized. - */ -enum NVTEFusedAttnFwdParamsAttribute { - // Tensor handles - kNVTEFusedAttnFwdParamsQ = 0, - kNVTEFusedAttnFwdParamsK, - kNVTEFusedAttnFwdParamsV, - kNVTEFusedAttnFwdParamsBias, - kNVTEFusedAttnFwdParamsSoftmaxOffset, - kNVTEFusedAttnFwdParamsS, - kNVTEFusedAttnFwdParamsO, - kNVTEFusedAttnFwdParamsAuxCtxTensors, - kNVTEFusedAttnFwdParamsCuSeqlensQ, - kNVTEFusedAttnFwdParamsCuSeqlensKV, - kNVTEFusedAttnFwdParamsCuSeqlensQPadded, - kNVTEFusedAttnFwdParamsCuSeqlensKVPadded, - kNVTEFusedAttnFwdParamsPageTableK, - kNVTEFusedAttnFwdParamsPageTableV, - kNVTEFusedAttnFwdParamsRngState, - // Configuration knobs - kNVTEFusedAttnFwdParamsIsTraining, - kNVTEFusedAttnFwdParamsCudaGraph, - kNVTEFusedAttnFwdParamsReturnMaxLogit, - kNVTEFusedAttnFwdParamsAttnMaskType, - kNVTEFusedAttnFwdParamsBiasType, - kNVTEFusedAttnFwdParamsWindowSizeLeft, - kNVTEFusedAttnFwdParamsWindowSizeRight, - kNVTEFusedAttnFwdParamsBottomRightDiagonal, - kNVTEFusedAttnFwdParamsSoftmaxType, - kNVTEFusedAttnFwdParamsDropout, - kNVTEFusedAttnFwdParamsAttnScale, - kNVTEFusedAttnFwdParamsQKVLayout, - kNVTEFusedAttnFwdParamsOFormat, - kNVTEFusedAttnFwdParamsQKVScaleInvFormat, - kNVTEFusedAttnFwdParamsMaxSeqlenQ, - kNVTEFusedAttnFwdParamsMaxSeqlenKV, - // Workspace and stream - kNVTEFusedAttnFwdParamsWorkspace, - kNVTEFusedAttnFwdParamsStream, - // Number of attributes - kNVTEFusedAttnFwdParamsNumAttributes -}; - -/*! \brief Create a fused-attention forward-parameter object. */ -NVTEFusedAttnFwdParams nvte_create_fused_attn_fwd_params(void); - -/*! \brief Destroy a fused-attention forward-parameter object. */ -void nvte_destroy_fused_attn_fwd_params(NVTEFusedAttnFwdParams params); - -/*! \brief Query an attribute in a fused-attention forward-parameter object. */ -void nvte_get_fused_attn_fwd_params_attribute(NVTEFusedAttnFwdParams params, - NVTEFusedAttnFwdParamsAttribute attr, void *buf, - size_t size_in_bytes, size_t *size_written); - -/*! \brief Set an attribute in a fused-attention forward-parameter object. */ -void nvte_set_fused_attn_fwd_params_attribute(NVTEFusedAttnFwdParams params, - NVTEFusedAttnFwdParamsAttribute attr, const void *buf, - size_t size_in_bytes); - -/*! \brief Opaque fused-attention backward-parameter handle. */ -typedef void *NVTEFusedAttnBwdParams; - -/*! \enum NVTEFusedAttnBwdParamsAttribute - * \brief Attributes for `NVTEFusedAttnBwdParams`. - * - * This enum is used to index the `FusedAttnBwdParams` struct. The order of its fields matches - * that of the declaration fields and the `attr_sizes` array of `FusedAttnBwdParams`. New fields - * may be appended at the end, and existing fields are never reordered, removed, or resized. - */ -enum NVTEFusedAttnBwdParamsAttribute { - // Tensor handles - kNVTEFusedAttnBwdParamsQ = 0, - kNVTEFusedAttnBwdParamsK, - kNVTEFusedAttnBwdParamsV, - kNVTEFusedAttnBwdParamsO, - kNVTEFusedAttnBwdParamsDO, - kNVTEFusedAttnBwdParamsS, - kNVTEFusedAttnBwdParamsDP, - kNVTEFusedAttnBwdParamsAuxCtxTensors, - kNVTEFusedAttnBwdParamsDQ, - kNVTEFusedAttnBwdParamsDK, - kNVTEFusedAttnBwdParamsDV, - kNVTEFusedAttnBwdParamsDBias, - kNVTEFusedAttnBwdParamsDSoftmaxOffset, - kNVTEFusedAttnBwdParamsCuSeqlensQ, - kNVTEFusedAttnBwdParamsCuSeqlensKV, - kNVTEFusedAttnBwdParamsCuSeqlensQPadded, - kNVTEFusedAttnBwdParamsCuSeqlensKVPadded, - // Configuration knobs - kNVTEFusedAttnBwdParamsDeterministic, - kNVTEFusedAttnBwdParamsCudaGraph, - kNVTEFusedAttnBwdParamsAttnMaskType, - kNVTEFusedAttnBwdParamsBiasType, - kNVTEFusedAttnBwdParamsWindowSizeLeft, - kNVTEFusedAttnBwdParamsWindowSizeRight, - kNVTEFusedAttnBwdParamsBottomRightDiagonal, - kNVTEFusedAttnBwdParamsSoftmaxType, - kNVTEFusedAttnBwdParamsDropout, - kNVTEFusedAttnBwdParamsAttnScale, - kNVTEFusedAttnBwdParamsQKVLayout, - kNVTEFusedAttnBwdParamsOFormat, - kNVTEFusedAttnBwdParamsDOFormat, - kNVTEFusedAttnBwdParamsDQKVLayout, - kNVTEFusedAttnBwdParamsQKVScaleInvFormat, - kNVTEFusedAttnBwdParamsDOScaleInvFormat, - kNVTEFusedAttnBwdParamsMaxSeqlenQ, - kNVTEFusedAttnBwdParamsMaxSeqlenKV, - // Workspace and stream - kNVTEFusedAttnBwdParamsWorkspace, - kNVTEFusedAttnBwdParamsStream, - // Number of attributes - kNVTEFusedAttnBwdParamsNumAttributes -}; - -/*! \brief Create a fused-attention backward-parameter object. */ -NVTEFusedAttnBwdParams nvte_create_fused_attn_bwd_params(void); - -/*! \brief Destroy a fused-attention backward-parameter object. */ -void nvte_destroy_fused_attn_bwd_params(NVTEFusedAttnBwdParams params); - -/*! \brief Query an attribute in a fused-attention backward-parameter object. */ -void nvte_get_fused_attn_bwd_params_attribute(NVTEFusedAttnBwdParams params, - NVTEFusedAttnBwdParamsAttribute attr, void *buf, - size_t size_in_bytes, size_t *size_written); - -/*! \brief Set an attribute in a fused-attention backward-parameter object. */ -void nvte_set_fused_attn_bwd_params_attribute(NVTEFusedAttnBwdParams params, - NVTEFusedAttnBwdParamsAttribute attr, const void *buf, - size_t size_in_bytes); - -/*! \brief Get fused-attention backend based on the user configuration. - * - * \param[in] cfg Fused-attention configuration created by - * `nvte_create_fused_attn_config()`. - * \param[out] message If the configuration is supported, an empty string. If not supported, - * a diagnostic message explaining why the configuration is rejected. Pass - * `NULL` to skip the diagnostics. The library maintains a per-thread buffer, - * and callers do not need to allocate memory for the message. The buffer - * content stays valid until the next call of `nvte_get_fused_attn_backend_v2`. - * - * \return Fused-attention backend, `NVTE_F16_arbitrary_seqlen` or `NVTE_FP8`, - * if the configuration is supported; otherwise, `NVTE_No_Backend`. - */ -NVTE_Fused_Attn_Backend nvte_get_fused_attn_backend_v2(NVTEFusedAttnConfig cfg, - const char **message); - -/*! \brief Get fused attention backend based on input parameters. - * - * \param[in] is_training Whether the model is in training mode. - * \param[in] q_dtype The data type of Tensor Q. - * \param[in] kv_dtype The data type of Tensors K, V. - * \param[in] qkv_layout The layout of Tensors Q, K, V. - * \param[in] bias_type The attention bias type. - * \param[in] attn_mask_type The attention mask type. - * \param[in] softmax_type The attention softmax type. - * \param[in] dropout The dropout probability. - * \param[in] num_attn_heads The number of heads in Q. - * \param[in] num_gqa_groups The number of heads in K, V. - * \param[in] max_seqlen_q The sequence length of Q. - * \param[in] max_seqlen_kv The sequence length of K, V. - * \param[in] head_dim_qk The head dimension of Q, K. - * \param[in] head_dim_v The head dimension of V. - * \param[in] window_size_left Sliding window size (the left half). - * \param[in] window_size_right Sliding window size (the right half). - * \param[in] return_max_logit Whether to produce Max along with Stats. - * \param[in] cuda_graph Whether cuda graph capture is enabled or not. - * \param[in] deterministic Whether determinism is required or not. - * - * \deprecated This function has been deprecated in favor of `nvte_get_fused_attn_backend_v2`. - * - * \note `nvte_get_fused_attn_backend` has a narrower input signature compared to - * `nvte_get_fused_attn_backend_v2`. For the fields it cannot express, it fills them with - * the default values of `nvte_get_fused_attn_backend_v2`. This includes setting - * `batch_size` = 1, deriving output/gradient formats from `qkv_layout`, assuming a standard - * bias shape [b, h, sq, skv] for `NVTE_POST_SCALE_BIAS`, using delayed scaling for all FP8, - * and not supporting paged-KV attention checks. Users who need more precise control should - * use `nvte_get_fused_attn_backend_v2` instead. - */ -NVTE_Fused_Attn_Backend nvte_get_fused_attn_backend( - bool is_training, NVTEDType q_dtype, NVTEDType kv_dtype, NVTE_QKV_Layout qkv_layout, - NVTE_Bias_Type bias_type, NVTE_Mask_Type attn_mask_type, NVTE_Softmax_Type softmax_type, - float dropout, size_t num_attn_heads, size_t num_gqa_groups, size_t max_seqlen_q, - size_t max_seqlen_kv, size_t head_dim_qk, size_t head_dim_v, int64_t window_size_left, - int64_t window_size_right, bool return_max_logit, bool cuda_graph, bool deterministic); - -/*! \brief Compute dot product attention with Q, K, and V. - * - * All inputs and outputs are carried by the opaque \p params handle. Create it with - * `nvte_create_fused_attn_fwd_params()`, populate it with - * `nvte_set_fused_attn_fwd_params_attribute()` or `FusedAttnFwdParamsWrapper` setters, and - * destroy it with `nvte_destroy_fused_attn_fwd_params()`. - * - * \param[in,out] params Fused-attention forward-parameter handle. - */ -void nvte_fused_attn_fwd_v2(NVTEFusedAttnFwdParams params); - -/*! \brief Compute dot product attention with separate Q, K and V. - * - * Computes: - * - P = Q * Transpose(K) + Bias - * - S = ScaleMaskSoftmax(P) - * - D = Dropout(S) - * - O = D * Transpose(V) - * - * Notes: - * - * Tensors `cu_seqlens_q_padded` and `cu_seqlens_kv_padded` - * help identify the correct offsets of different sequences in tensors Q, K, V and O. - * When the QKV format (`nvte_get_qkv_format(qkv_layout)`) is `bshd` or `sbhd`, - * offset tensors are not used in the attention calculation and can be set to empty `NVTETensor`s. - * When the QKV format is `thd`, these tensors should follow the following rules. - * When there is no padding between sequences, the offset tensors should be equal to - * `cu_seqlens_q` and `cu_seqlens_kv` respectively. - * When there is padding between sequences, users are responsible to adjust the offsets as needed. - * For example, a tensor of 4 sequences `[a, PAD, b, b, c, PAD, PAD, d, d]` should have - * `cu_seqlens = [0, 1, 3, 4, 6]` and `cu_seqlens_padded= [0, 2, 4, 7, 9]`. - * - * \param[in] Q The Q tensor. - * \param[in] K The K tensor. - * \param[in] V The V tensor. - * \param[in] Bias The Bias tensor. - * \param[in] SoftmaxOffset The SoftmaxOffset tensor. - * \param[in,out] S The S tensor. - * \param[out] O The output O tensor. - * \param[out] Aux_CTX_Tensors Auxiliary output tensors when training, - * e.g. softmax stats, optional Max, rng_state. - * \param[in] cu_seqlens_q Cumulative sequence lengths for Q, [batch_size + 1]. - * \param[in] cu_seqlens_kv Cumulative sequence lengths for K and V, [batch_size + 1]. - * \param[in] cu_seqlens_q_padded Cumulative sequence offsets for Q, [batch_size + 1]. - * \param[in] cu_seqlens_kv_padded Cumulative sequence offsets for KV, [batch_size + 1]. - * \param[in] page_table_k Page table for K cache, [batch_size, max_pages_per_seq_k]. - * \param[in] page_table_v Page table for V cache, [batch_size, max_pages_per_seq_v]. - * \param[in] rng_state Seed and offset of CUDA random number generator. - * \param[in] max_seqlen_q Max sequence length used for computing for Q. - * it may be >= max(seqlen_q_i) for i=0,...batch_size-1. - * \param[in] max_seqlen_kv Max sequence length used for computing for K and V. - * it may be >= max(seqlen_kv_i) for i=0,...batch_size-1. - * \param[in] is_training Whether this is in training mode or inference. - * \param[in] return_max_logit Whether to produce Max along with Stats. - * \param[in] cuda_graph Whether cuda graph capture is enabled or not. - * \param[in] attn_scale Scaling factor for Q * K.T. - * \param[in] dropout Dropout probability. - * \param[in] qkv_layout QKV tensors' layout. - * \param[in] o_format Output format. - * \param[in] qkv_scale_inv_format Format of scale-inverse tensors for QKV; - * if `NVTE_QKV_Format_NOT_SET`, inferred from - * `qkv_layout`. - * \param[in] bias_type Bias type. - * \param[in] attn_mask_type Attention mask type. - * \param[in] softmax_type Attention softmax type. - * \param[in] window_size_left Sliding window size (the left half). - * \param[in] window_size_right Sliding window size (the right half). - * \param[in] bottom_right_diagonal Whether to align sliding window and ALiBi diagonal to the bottom right corner of the softmax matrix. - * \param[in] workspace Workspace tensor. - * \param[in] stream CUDA stream used for this operation. - * - * \deprecated This function has been deprecated in favor of `nvte_fused_attn_fwd_v2`. - */ -void nvte_fused_attn_fwd(const NVTETensor Q, const NVTETensor K, const NVTETensor V, - const NVTETensor Bias, const NVTETensor SoftmaxOffset, NVTETensor S, - NVTETensor O, NVTETensorPack *Aux_CTX_Tensors, - const NVTETensor cu_seqlens_q, const NVTETensor cu_seqlens_kv, - const NVTETensor cu_seqlens_q_padded, - const NVTETensor cu_seqlens_kv_padded, const NVTETensor page_table_k, - const NVTETensor page_table_v, const NVTETensor rng_state, - size_t max_seqlen_q, size_t max_seqlen_kv, bool is_training, - bool return_max_logit, bool cuda_graph, float attn_scale, float dropout, - NVTE_QKV_Layout qkv_layout, NVTE_QKV_Format o_format, - NVTE_QKV_Format qkv_scale_inv_format, NVTE_Bias_Type bias_type, - NVTE_Mask_Type attn_mask_type, NVTE_Softmax_Type softmax_type, - int64_t window_size_left, int64_t window_size_right, - bool bottom_right_diagonal, NVTETensor workspace, cudaStream_t stream); - -/*! \brief Compute the backward of the dot product attention with Q, K and V. - * - * All inputs and outputs are carried by the opaque \p params handle. Create it with - * `nvte_create_fused_attn_bwd_params()`, populate it with - * `nvte_set_fused_attn_bwd_params_attribute()` or `FusedAttnBwdParamsWrapper` setters, and - * destroy it with `nvte_destroy_fused_attn_bwd_params()`. - * - * \param[in,out] params Fused-attention backward-parameter handle. - */ -void nvte_fused_attn_bwd_v2(NVTEFusedAttnBwdParams params); - -/*! \brief Compute the backward of the dot product attention with separate Q, K and V. - * - * Notes: - * - * Tensors `cu_seqlens_q_padded` and `cu_seqlens_kv_padded` - * help identify the correct offsets of different sequences in tensors Q, K, V and O. - * When the QKV format (`nvte_get_qkv_format(qkv_layout)`) is `bshd` or `sbhd`, - * offset tensors are not used in the attention calculation and can be set to empty `NVTETensor`s. - * When the QKV format is `thd`, these tensors should follow the following rules. - * When there is no padding between sequences, the offset tensors should be equal to - * `cu_seqlens_q` and `cu_seqlens_kv` respectively. - * When there is padding between sequences, users are responsible to adjust the offsets as needed. - * For example, a tensor of 4 sequences `[a, PAD, b, b, c, PAD, PAD, d, d]` should have - * `cu_seqlens = [0, 1, 3, 4, 6]` and `cu_seqlens_padded= [0, 2, 4, 7, 9]`. - * - * \param[in] Q The Q tensor. - * \param[in] K The K tensor. - * \param[in] V The V tensor. - * \param[in] O The O tensor from forward. - * \param[in] dO The gradient of the O tensor. - * \param[in] S The S tensor. - * \param[in,out] dP The gradient of the P tensor. - * \param[in] Aux_CTX_Tensors Auxiliary tensors from context when in training mode, - * e.g. softmax stats, optional Max, rng_state. - * \param[out] dQ The gradient of the Q tensor. - * \param[out] dK The gradient of the K tensor. - * \param[out] dV The gradient of the V tensor. - * \param[out] dBias The gradient of the Bias tensor. - * \param[out] dSoftmaxOffset The gradient of the SoftmaxOffset tensor. - * \param[in] cu_seqlens_q Cumulative sequence lengths for Q, [batch_size + 1]. - * \param[in] cu_seqlens_kv Cumulative sequence lengths for K and V, [batch_size + 1]. - * \param[in] cu_seqlens_q_padded Cumulative sequence offsets for Q, [batch_size + 1]. - * \param[in] cu_seqlens_kv_padded Cumulative sequence offsets for KV, [batch_size + 1]. - * \param[in] max_seqlen_q Max sequence length used for computing for Q. - * it may be >= max(seqlen_q_i) for i=0,...batch_size-1. - * \param[in] max_seqlen_kv Max sequence length used for computing for K and V. - * it may be >= max(seqlen_kv_i) for i=0,...batch_size-1. - * \param[in] attn_scale Scaling factor for Q * K.T. - * \param[in] dropout Dropout probability. - * \param[in] qkv_layout QKV tensors' layout. - * \param[in] o_format Output format. - * \param[in] do_format Output gradient's format. - * \param[in] dqkv_layout QKV gradient tensors' layout. - * \param[in] qkv_scale_inv_format Format of scale-inverse tensors for QKV; - * if `NVTE_QKV_Format_NOT_SET`, inferred from - * `qkv_layout`. - * \param[in] do_scale_inv_format Format of scale-inverse tensors for dO; - * if `NVTE_QKV_Format_NOT_SET`, inferred from the - * output layout. - * \param[in] bias_type Bias type. - * \param[in] attn_mask_type Attention mask type. - * \param[in] softmax_type Attention softmax type. - * \param[in] window_size_left Sliding window size (the left half). - * \param[in] window_size_right Sliding window size (the right half). - * \param[in] bottom_right_diagonal Whether to align sliding window and ALiBi diagonal to the bottom right corner of the softmax matrix. - * \param[in] deterministic Whether to execute with deterministic behaviours. - * \param[in] cuda_graph Whether cuda graph capture is enabled or not. - * \param[in] workspace Workspace tensor. - * \param[in] stream CUDA stream used for this operation. - * - * \deprecated This function has been deprecated in favor of `nvte_fused_attn_bwd_v2`. - */ -void nvte_fused_attn_bwd(const NVTETensor Q, const NVTETensor K, const NVTETensor V, - const NVTETensor O, const NVTETensor dO, const NVTETensor S, NVTETensor dP, - const NVTETensorPack *Aux_CTX_Tensors, NVTETensor dQ, NVTETensor dK, - NVTETensor dV, NVTETensor dBias, NVTETensor dSoftmaxOffset, - const NVTETensor cu_seqlens_q, const NVTETensor cu_seqlens_kv, - const NVTETensor cu_seqlens_q_padded, - const NVTETensor cu_seqlens_kv_padded, size_t max_seqlen_q, - size_t max_seqlen_kv, float attn_scale, float dropout, - NVTE_QKV_Layout qkv_layout, NVTE_QKV_Format o_format, - NVTE_QKV_Format do_format, NVTE_QKV_Layout dqkv_layout, - NVTE_QKV_Format qkv_scale_inv_format, NVTE_QKV_Format do_scale_inv_format, - NVTE_Bias_Type bias_type, NVTE_Mask_Type attn_mask_type, - NVTE_Softmax_Type softmax_type, int64_t window_size_left, - int64_t window_size_right, bool bottom_right_diagonal, bool deterministic, - bool cuda_graph, NVTETensor workspace, cudaStream_t stream); - -/*! \brief Update the RNG state with the seed and calculated offset. - * - * \warning This API is **experimental** and subject to change. - * - * \param[in] rng_state_dst RNG state to store seed and offset. - * \param[in] seed Seed for RNG state. - * \param[in] q_max_seqlen Max sequence length used for computing for Q. - * it may be >= max(seqlen_q_i) for i=0,...batch_size-1. - * \param[in] kv_max_seqlen Max sequence length used for computing for K and V. - * it may be >= max(seqlen_kv_i) for i=0,...batch_size-1. - * \param[in] backend Fused attention backend. - * \param[in] stream CUDA stream used for this operation. - */ -void nvte_populate_rng_state_async(NVTETensor rng_state_dst, const NVTETensor seed, - size_t q_max_seqlen, size_t kv_max_seqlen, - NVTE_Fused_Attn_Backend backend, cudaStream_t stream); - -/*! \brief Get KV format for a given QKV layout. - * - * \warning This API is **experimental** and subject to change. - * - * \param[in] cu_seqlens Cumulative sequence lengths, [batch_size + 1]. - * \param[in] workspace Workspace tensor. - * \param[in] len batch_size x sequence_length. - * \param[in] stream CUDA stream used for this operation. - */ -uint32_t nvte_get_runtime_num_segments(NVTETensor cu_seqlens, NVTETensor workspace, size_t len, - cudaStream_t stream); - -/*! \brief Set the seed and offset for RNG state. - * - * \warning This API is **experimental** and subject to change. - * - * \param[out] rng_state_ptr A size 2 array storing the RNG's seed and offset respectively. - * \param[in] captured Whether a CUDA graph is being captured. - * \param[in] seed_ptr Seed pointer. - * \param[in] seed_val Seed value. - * \param[in] offset_ptr Offset pointer. - * \param[in] offset_val Offset value. - * \param[in] offset_intragraph Intragraph offset in RNG states. For use with CUDA Graphs. - * \param[in] stream CUDA stream used for this operation. - */ -void nvte_extract_seed_and_offset(int64_t *rng_state_ptr, int captured, int64_t *seed_ptr, - uint64_t seed_val, int64_t *offset_ptr, uint64_t offset_val, - uint32_t offset_intragraph, cudaStream_t stream); - /*! \brief Copy keys and values into the KV cache. * * \warning This API is **experimental** and subject to change. @@ -942,532 +447,6 @@ void nvte_multi_tensor_pad_last_dim(NVTETensor *inputs, NVTETensor *outputs, siz #ifdef __cplusplus } // extern "C" - -#include -#include -#include - -/*! \brief Parses a QKV tensor shape into canonical (b, h, s, d, t) dimensions - * and converts between QKV formats. - */ -class AttentionShape { - public: - inline AttentionShape(NVTE_QKV_Format fmt, const size_t *shape) : canonical_{} { - auto [ndim, order] = dim_order(fmt); - for (size_t i = 0; i < ndim; ++i) canonical_[order[i]] = shape[i]; - } - - size_t b() const { return canonical_[0]; } - size_t h() const { return canonical_[1]; } - size_t s() const { return canonical_[2]; } - size_t d() const { return canonical_[3]; } - size_t t() const { return canonical_[4]; } - - inline void to_format(NVTE_QKV_Format dst_fmt, size_t *dst_shape) const { - auto [ndim, order] = dim_order(dst_fmt); - for (size_t i = 0; i < ndim; ++i) dst_shape[i] = canonical_[order[i]]; - } - - private: - static inline std::pair> dim_order(NVTE_QKV_Format fmt) { - switch (fmt) { - case NVTE_QKV_Format::NVTE_BSHD: - return {4, {0, 2, 1, 3}}; // b s h d - case NVTE_QKV_Format::NVTE_SBHD: - return {4, {2, 0, 1, 3}}; // s b h d - case NVTE_QKV_Format::NVTE_BHSD: - return {4, {0, 1, 2, 3}}; // b h s d - case NVTE_QKV_Format::NVTE_THD: - return {3, {4, 1, 3, -1}}; // t h d - default: - return {0, {}}; - } - } - size_t canonical_[5] = {}; -}; - -/*! \class FusedAttnConfigWrapper - * \brief C++ helper for constructing an `NVTEFusedAttnConfig`. - * - * It owns an opaque `NVTEFusedAttnConfig` handle created by - * `nvte_create_fused_attn_config()`, and provides a convenient, - * chainable interface for setting every field in `FusedAttnConfig`. - */ -class FusedAttnConfigWrapper { - public: - FusedAttnConfigWrapper() : cfg_{nvte_create_fused_attn_config()} {} - - FusedAttnConfigWrapper(const FusedAttnConfigWrapper &) = delete; - FusedAttnConfigWrapper &operator=(const FusedAttnConfigWrapper &) = delete; - - FusedAttnConfigWrapper(FusedAttnConfigWrapper &&other) noexcept : cfg_{other.cfg_} { - other.cfg_ = nullptr; - } - - FusedAttnConfigWrapper &operator=(FusedAttnConfigWrapper &&other) noexcept { - if (this != &other) { - if (cfg_ != nullptr) { - nvte_destroy_fused_attn_config(cfg_); - } - cfg_ = other.cfg_; - other.cfg_ = nullptr; - } - return *this; - } - - ~FusedAttnConfigWrapper() { - if (cfg_ != nullptr) { - nvte_destroy_fused_attn_config(cfg_); - } - } - - operator NVTEFusedAttnConfig() const noexcept { return cfg_; } - NVTEFusedAttnConfig get() const noexcept { return cfg_; } - - FusedAttnConfigWrapper &set_is_training(bool val) noexcept { - return set_attr(kNVTEFusedAttnConfigIsTraining, static_cast(val)); - } - FusedAttnConfigWrapper &set_deterministic(bool val) noexcept { - return set_attr(kNVTEFusedAttnConfigDeterministic, static_cast(val)); - } - FusedAttnConfigWrapper &set_cuda_graph(bool val) noexcept { - return set_attr(kNVTEFusedAttnConfigCudaGraph, static_cast(val)); - } - FusedAttnConfigWrapper &set_return_max_logit(bool val) noexcept { - return set_attr(kNVTEFusedAttnConfigReturnMaxLogit, static_cast(val)); - } - FusedAttnConfigWrapper &set_attn_mask_type(NVTE_Mask_Type val) noexcept { - return set_attr(kNVTEFusedAttnConfigAttnMaskType, val); - } - FusedAttnConfigWrapper &set_bias_type(NVTE_Bias_Type val) noexcept { - return set_attr(kNVTEFusedAttnConfigBiasType, val); - } - FusedAttnConfigWrapper &set_window_size_left(int64_t val) noexcept { - return set_attr(kNVTEFusedAttnConfigWindowSizeLeft, val); - } - FusedAttnConfigWrapper &set_window_size_right(int64_t val) noexcept { - return set_attr(kNVTEFusedAttnConfigWindowSizeRight, val); - } - FusedAttnConfigWrapper &set_bottom_right_diagonal(bool val) noexcept { - return set_attr(kNVTEFusedAttnConfigBottomRightDiagonal, static_cast(val)); - } - FusedAttnConfigWrapper &set_softmax_type(NVTE_Softmax_Type val) noexcept { - return set_attr(kNVTEFusedAttnConfigSoftmaxType, val); - } - FusedAttnConfigWrapper &set_scaling_mode(NVTEScalingMode val) noexcept { - return set_attr(kNVTEFusedAttnConfigScalingMode, val); - } - FusedAttnConfigWrapper &set_dropout(float val) noexcept { - return set_attr(kNVTEFusedAttnConfigDropout, val); - } - FusedAttnConfigWrapper &set_attn_scale(float val) noexcept { - return set_attr(kNVTEFusedAttnConfigAttnScale, val); - } - FusedAttnConfigWrapper &set_qkv_dtype(NVTEDType val) noexcept { - return set_attr(kNVTEFusedAttnConfigQKVDtype, val); - } - FusedAttnConfigWrapper &set_o_dtype(NVTEDType val) noexcept { - return set_attr(kNVTEFusedAttnConfigODtype, val); - } - FusedAttnConfigWrapper &set_do_dtype(NVTEDType val) noexcept { - return set_attr(kNVTEFusedAttnConfigDODtype, val); - } - FusedAttnConfigWrapper &set_dqkv_dtype(NVTEDType val) noexcept { - return set_attr(kNVTEFusedAttnConfigDQKVDtype, val); - } - FusedAttnConfigWrapper &set_qkv_layout(NVTE_QKV_Layout val) noexcept { - return set_attr(kNVTEFusedAttnConfigQKVLayout, val); - } - FusedAttnConfigWrapper &set_o_format(NVTE_QKV_Format val) noexcept { - return set_attr(kNVTEFusedAttnConfigOFormat, val); - } - FusedAttnConfigWrapper &set_do_format(NVTE_QKV_Format val) noexcept { - return set_attr(kNVTEFusedAttnConfigDOFormat, val); - } - FusedAttnConfigWrapper &set_dqkv_layout(NVTE_QKV_Layout val) noexcept { - return set_attr(kNVTEFusedAttnConfigDQKVLayout, val); - } - FusedAttnConfigWrapper &set_qkv_scale_inv_format(NVTE_QKV_Format val) noexcept { - return set_attr(kNVTEFusedAttnConfigQKVScaleInvFormat, val); - } - FusedAttnConfigWrapper &set_do_scale_inv_format(NVTE_QKV_Format val) noexcept { - return set_attr(kNVTEFusedAttnConfigDOScaleInvFormat, val); - } - FusedAttnConfigWrapper &set_batch_size(size_t val) noexcept { - return set_attr(kNVTEFusedAttnConfigBatchSize, val); - } - FusedAttnConfigWrapper &set_num_attn_heads(size_t val) noexcept { - return set_attr(kNVTEFusedAttnConfigNumAttnHeads, val); - } - FusedAttnConfigWrapper &set_num_gqa_groups(size_t val) noexcept { - return set_attr(kNVTEFusedAttnConfigNumGQAGroups, val); - } - FusedAttnConfigWrapper &set_head_dim_qk(size_t val) noexcept { - return set_attr(kNVTEFusedAttnConfigHeadDimQK, val); - } - FusedAttnConfigWrapper &set_head_dim_v(size_t val) noexcept { - return set_attr(kNVTEFusedAttnConfigHeadDimV, val); - } - FusedAttnConfigWrapper &set_max_seqlen_q(size_t val) noexcept { - return set_attr(kNVTEFusedAttnConfigMaxSeqlenQ, val); - } - FusedAttnConfigWrapper &set_max_seqlen_kv(size_t val) noexcept { - return set_attr(kNVTEFusedAttnConfigMaxSeqlenKV, val); - } - FusedAttnConfigWrapper &set_num_tokens_q(size_t val) noexcept { - return set_attr(kNVTEFusedAttnConfigNumTokensQ, val); - } - FusedAttnConfigWrapper &set_num_tokens_kv(size_t val) noexcept { - return set_attr(kNVTEFusedAttnConfigNumTokensKV, val); - } - FusedAttnConfigWrapper &set_num_pages_k(size_t val) noexcept { - return set_attr(kNVTEFusedAttnConfigNumPagesK, val); - } - FusedAttnConfigWrapper &set_num_pages_v(size_t val) noexcept { - return set_attr(kNVTEFusedAttnConfigNumPagesV, val); - } - FusedAttnConfigWrapper &set_page_size_k(size_t val) noexcept { - return set_attr(kNVTEFusedAttnConfigPageSizeK, val); - } - FusedAttnConfigWrapper &set_page_size_v(size_t val) noexcept { - return set_attr(kNVTEFusedAttnConfigPageSizeV, val); - } - FusedAttnConfigWrapper &set_max_pages_per_seq_k(size_t val) noexcept { - return set_attr(kNVTEFusedAttnConfigMaxPagesPerSeqK, val); - } - FusedAttnConfigWrapper &set_max_pages_per_seq_v(size_t val) noexcept { - return set_attr(kNVTEFusedAttnConfigMaxPagesPerSeqV, val); - } - FusedAttnConfigWrapper &set_bias_batch_size(size_t val) noexcept { - return set_attr(kNVTEFusedAttnConfigBiasBatchSize, val); - } - FusedAttnConfigWrapper &set_bias_num_heads(size_t val) noexcept { - return set_attr(kNVTEFusedAttnConfigBiasNumHeads, val); - } - FusedAttnConfigWrapper &set_bias_seqlen_q(size_t val) noexcept { - return set_attr(kNVTEFusedAttnConfigBiasSeqlenQ, val); - } - FusedAttnConfigWrapper &set_bias_seqlen_kv(size_t val) noexcept { - return set_attr(kNVTEFusedAttnConfigBiasSeqlenKV, val); - } - - private: - template - FusedAttnConfigWrapper &set_attr(NVTEFusedAttnConfigAttribute attr, T val) noexcept { - nvte_set_fused_attn_config_attribute(cfg_, attr, &val, sizeof(val)); - return *this; - } - - NVTEFusedAttnConfig cfg_ = nullptr; -}; - -/*! \class FusedAttnFwdParamsWrapper - * \brief C++ helper for constructing an `NVTEFusedAttnFwdParams`. - * - * It owns an opaque `NVTEFusedAttnFwdParams` handle created by - * `nvte_create_fused_attn_fwd_params()`, and provides a convenient, - * chainable interface for setting every field in `FusedAttnFwdParams`. - */ -class FusedAttnFwdParamsWrapper { - public: - FusedAttnFwdParamsWrapper() : params_{nvte_create_fused_attn_fwd_params()} {} - - FusedAttnFwdParamsWrapper(const FusedAttnFwdParamsWrapper &) = delete; - FusedAttnFwdParamsWrapper &operator=(const FusedAttnFwdParamsWrapper &) = delete; - - FusedAttnFwdParamsWrapper(FusedAttnFwdParamsWrapper &&other) noexcept : params_{other.params_} { - other.params_ = nullptr; - } - - FusedAttnFwdParamsWrapper &operator=(FusedAttnFwdParamsWrapper &&other) noexcept { - if (this != &other) { - if (params_ != nullptr) { - nvte_destroy_fused_attn_fwd_params(params_); - } - params_ = other.params_; - other.params_ = nullptr; - } - return *this; - } - - ~FusedAttnFwdParamsWrapper() { - if (params_ != nullptr) { - nvte_destroy_fused_attn_fwd_params(params_); - } - } - - operator NVTEFusedAttnFwdParams() const noexcept { return params_; } - NVTEFusedAttnFwdParams get() const noexcept { return params_; } - - FusedAttnFwdParamsWrapper &set_Q(NVTETensor val) noexcept { - return set_attr(kNVTEFusedAttnFwdParamsQ, val); - } - FusedAttnFwdParamsWrapper &set_K(NVTETensor val) noexcept { - return set_attr(kNVTEFusedAttnFwdParamsK, val); - } - FusedAttnFwdParamsWrapper &set_V(NVTETensor val) noexcept { - return set_attr(kNVTEFusedAttnFwdParamsV, val); - } - FusedAttnFwdParamsWrapper &set_Bias(NVTETensor val) noexcept { - return set_attr(kNVTEFusedAttnFwdParamsBias, val); - } - FusedAttnFwdParamsWrapper &set_SoftmaxOffset(NVTETensor val) noexcept { - return set_attr(kNVTEFusedAttnFwdParamsSoftmaxOffset, val); - } - FusedAttnFwdParamsWrapper &set_S(NVTETensor val) noexcept { - return set_attr(kNVTEFusedAttnFwdParamsS, val); - } - FusedAttnFwdParamsWrapper &set_O(NVTETensor val) noexcept { - return set_attr(kNVTEFusedAttnFwdParamsO, val); - } - FusedAttnFwdParamsWrapper &set_Aux_CTX_Tensors(NVTETensorPack *val) noexcept { - return set_attr(kNVTEFusedAttnFwdParamsAuxCtxTensors, val); - } - FusedAttnFwdParamsWrapper &set_cu_seqlens_q(NVTETensor val) noexcept { - return set_attr(kNVTEFusedAttnFwdParamsCuSeqlensQ, val); - } - FusedAttnFwdParamsWrapper &set_cu_seqlens_kv(NVTETensor val) noexcept { - return set_attr(kNVTEFusedAttnFwdParamsCuSeqlensKV, val); - } - FusedAttnFwdParamsWrapper &set_cu_seqlens_q_padded(NVTETensor val) noexcept { - return set_attr(kNVTEFusedAttnFwdParamsCuSeqlensQPadded, val); - } - FusedAttnFwdParamsWrapper &set_cu_seqlens_kv_padded(NVTETensor val) noexcept { - return set_attr(kNVTEFusedAttnFwdParamsCuSeqlensKVPadded, val); - } - FusedAttnFwdParamsWrapper &set_page_table_k(NVTETensor val) noexcept { - return set_attr(kNVTEFusedAttnFwdParamsPageTableK, val); - } - FusedAttnFwdParamsWrapper &set_page_table_v(NVTETensor val) noexcept { - return set_attr(kNVTEFusedAttnFwdParamsPageTableV, val); - } - FusedAttnFwdParamsWrapper &set_rng_state(NVTETensor val) noexcept { - return set_attr(kNVTEFusedAttnFwdParamsRngState, val); - } - FusedAttnFwdParamsWrapper &set_is_training(bool val) noexcept { - return set_attr(kNVTEFusedAttnFwdParamsIsTraining, static_cast(val)); - } - FusedAttnFwdParamsWrapper &set_cuda_graph(bool val) noexcept { - return set_attr(kNVTEFusedAttnFwdParamsCudaGraph, static_cast(val)); - } - FusedAttnFwdParamsWrapper &set_return_max_logit(bool val) noexcept { - return set_attr(kNVTEFusedAttnFwdParamsReturnMaxLogit, static_cast(val)); - } - FusedAttnFwdParamsWrapper &set_attn_mask_type(NVTE_Mask_Type val) noexcept { - return set_attr(kNVTEFusedAttnFwdParamsAttnMaskType, val); - } - FusedAttnFwdParamsWrapper &set_bias_type(NVTE_Bias_Type val) noexcept { - return set_attr(kNVTEFusedAttnFwdParamsBiasType, val); - } - FusedAttnFwdParamsWrapper &set_window_size_left(int64_t val) noexcept { - return set_attr(kNVTEFusedAttnFwdParamsWindowSizeLeft, val); - } - FusedAttnFwdParamsWrapper &set_window_size_right(int64_t val) noexcept { - return set_attr(kNVTEFusedAttnFwdParamsWindowSizeRight, val); - } - FusedAttnFwdParamsWrapper &set_bottom_right_diagonal(bool val) noexcept { - return set_attr(kNVTEFusedAttnFwdParamsBottomRightDiagonal, static_cast(val)); - } - FusedAttnFwdParamsWrapper &set_softmax_type(NVTE_Softmax_Type val) noexcept { - return set_attr(kNVTEFusedAttnFwdParamsSoftmaxType, val); - } - FusedAttnFwdParamsWrapper &set_dropout(float val) noexcept { - return set_attr(kNVTEFusedAttnFwdParamsDropout, val); - } - FusedAttnFwdParamsWrapper &set_attn_scale(float val) noexcept { - return set_attr(kNVTEFusedAttnFwdParamsAttnScale, val); - } - FusedAttnFwdParamsWrapper &set_qkv_layout(NVTE_QKV_Layout val) noexcept { - return set_attr(kNVTEFusedAttnFwdParamsQKVLayout, val); - } - FusedAttnFwdParamsWrapper &set_o_format(NVTE_QKV_Format val) noexcept { - return set_attr(kNVTEFusedAttnFwdParamsOFormat, val); - } - FusedAttnFwdParamsWrapper &set_qkv_scale_inv_format(NVTE_QKV_Format val) noexcept { - return set_attr(kNVTEFusedAttnFwdParamsQKVScaleInvFormat, val); - } - FusedAttnFwdParamsWrapper &set_max_seqlen_q(size_t val) noexcept { - return set_attr(kNVTEFusedAttnFwdParamsMaxSeqlenQ, val); - } - FusedAttnFwdParamsWrapper &set_max_seqlen_kv(size_t val) noexcept { - return set_attr(kNVTEFusedAttnFwdParamsMaxSeqlenKV, val); - } - FusedAttnFwdParamsWrapper &set_workspace(NVTETensor val) noexcept { - return set_attr(kNVTEFusedAttnFwdParamsWorkspace, val); - } - FusedAttnFwdParamsWrapper &set_stream(cudaStream_t val) noexcept { - return set_attr(kNVTEFusedAttnFwdParamsStream, val); - } - - private: - template - FusedAttnFwdParamsWrapper &set_attr(NVTEFusedAttnFwdParamsAttribute attr, T val) noexcept { - nvte_set_fused_attn_fwd_params_attribute(params_, attr, &val, sizeof(val)); - return *this; - } - - NVTEFusedAttnFwdParams params_ = nullptr; -}; - -/*! \class FusedAttnBwdParamsWrapper - * \brief C++ helper for constructing an `NVTEFusedAttnBwdParams`. - * - * It owns an opaque `NVTEFusedAttnBwdParams` handle created by - * `nvte_create_fused_attn_bwd_params()`, and provides a convenient, - * chainable interface for setting every field in `FusedAttnBwdParams`. - */ -class FusedAttnBwdParamsWrapper { - public: - FusedAttnBwdParamsWrapper() : params_{nvte_create_fused_attn_bwd_params()} {} - - FusedAttnBwdParamsWrapper(const FusedAttnBwdParamsWrapper &) = delete; - FusedAttnBwdParamsWrapper &operator=(const FusedAttnBwdParamsWrapper &) = delete; - - FusedAttnBwdParamsWrapper(FusedAttnBwdParamsWrapper &&other) noexcept : params_{other.params_} { - other.params_ = nullptr; - } - - FusedAttnBwdParamsWrapper &operator=(FusedAttnBwdParamsWrapper &&other) noexcept { - if (this != &other) { - if (params_ != nullptr) { - nvte_destroy_fused_attn_bwd_params(params_); - } - params_ = other.params_; - other.params_ = nullptr; - } - return *this; - } - - ~FusedAttnBwdParamsWrapper() { - if (params_ != nullptr) { - nvte_destroy_fused_attn_bwd_params(params_); - } - } - - operator NVTEFusedAttnBwdParams() const noexcept { return params_; } - NVTEFusedAttnBwdParams get() const noexcept { return params_; } - - FusedAttnBwdParamsWrapper &set_Q(NVTETensor val) noexcept { - return set_attr(kNVTEFusedAttnBwdParamsQ, val); - } - FusedAttnBwdParamsWrapper &set_K(NVTETensor val) noexcept { - return set_attr(kNVTEFusedAttnBwdParamsK, val); - } - FusedAttnBwdParamsWrapper &set_V(NVTETensor val) noexcept { - return set_attr(kNVTEFusedAttnBwdParamsV, val); - } - FusedAttnBwdParamsWrapper &set_O(NVTETensor val) noexcept { - return set_attr(kNVTEFusedAttnBwdParamsO, val); - } - FusedAttnBwdParamsWrapper &set_dO(NVTETensor val) noexcept { - return set_attr(kNVTEFusedAttnBwdParamsDO, val); - } - FusedAttnBwdParamsWrapper &set_S(NVTETensor val) noexcept { - return set_attr(kNVTEFusedAttnBwdParamsS, val); - } - FusedAttnBwdParamsWrapper &set_dP(NVTETensor val) noexcept { - return set_attr(kNVTEFusedAttnBwdParamsDP, val); - } - FusedAttnBwdParamsWrapper &set_Aux_CTX_Tensors(const NVTETensorPack *val) noexcept { - return set_attr(kNVTEFusedAttnBwdParamsAuxCtxTensors, val); - } - FusedAttnBwdParamsWrapper &set_dQ(NVTETensor val) noexcept { - return set_attr(kNVTEFusedAttnBwdParamsDQ, val); - } - FusedAttnBwdParamsWrapper &set_dK(NVTETensor val) noexcept { - return set_attr(kNVTEFusedAttnBwdParamsDK, val); - } - FusedAttnBwdParamsWrapper &set_dV(NVTETensor val) noexcept { - return set_attr(kNVTEFusedAttnBwdParamsDV, val); - } - FusedAttnBwdParamsWrapper &set_dBias(NVTETensor val) noexcept { - return set_attr(kNVTEFusedAttnBwdParamsDBias, val); - } - FusedAttnBwdParamsWrapper &set_dSoftmaxOffset(NVTETensor val) noexcept { - return set_attr(kNVTEFusedAttnBwdParamsDSoftmaxOffset, val); - } - FusedAttnBwdParamsWrapper &set_cu_seqlens_q(NVTETensor val) noexcept { - return set_attr(kNVTEFusedAttnBwdParamsCuSeqlensQ, val); - } - FusedAttnBwdParamsWrapper &set_cu_seqlens_kv(NVTETensor val) noexcept { - return set_attr(kNVTEFusedAttnBwdParamsCuSeqlensKV, val); - } - FusedAttnBwdParamsWrapper &set_cu_seqlens_q_padded(NVTETensor val) noexcept { - return set_attr(kNVTEFusedAttnBwdParamsCuSeqlensQPadded, val); - } - FusedAttnBwdParamsWrapper &set_cu_seqlens_kv_padded(NVTETensor val) noexcept { - return set_attr(kNVTEFusedAttnBwdParamsCuSeqlensKVPadded, val); - } - FusedAttnBwdParamsWrapper &set_deterministic(bool val) noexcept { - return set_attr(kNVTEFusedAttnBwdParamsDeterministic, static_cast(val)); - } - FusedAttnBwdParamsWrapper &set_cuda_graph(bool val) noexcept { - return set_attr(kNVTEFusedAttnBwdParamsCudaGraph, static_cast(val)); - } - FusedAttnBwdParamsWrapper &set_attn_mask_type(NVTE_Mask_Type val) noexcept { - return set_attr(kNVTEFusedAttnBwdParamsAttnMaskType, val); - } - FusedAttnBwdParamsWrapper &set_bias_type(NVTE_Bias_Type val) noexcept { - return set_attr(kNVTEFusedAttnBwdParamsBiasType, val); - } - FusedAttnBwdParamsWrapper &set_window_size_left(int64_t val) noexcept { - return set_attr(kNVTEFusedAttnBwdParamsWindowSizeLeft, val); - } - FusedAttnBwdParamsWrapper &set_window_size_right(int64_t val) noexcept { - return set_attr(kNVTEFusedAttnBwdParamsWindowSizeRight, val); - } - FusedAttnBwdParamsWrapper &set_bottom_right_diagonal(bool val) noexcept { - return set_attr(kNVTEFusedAttnBwdParamsBottomRightDiagonal, static_cast(val)); - } - FusedAttnBwdParamsWrapper &set_softmax_type(NVTE_Softmax_Type val) noexcept { - return set_attr(kNVTEFusedAttnBwdParamsSoftmaxType, val); - } - FusedAttnBwdParamsWrapper &set_dropout(float val) noexcept { - return set_attr(kNVTEFusedAttnBwdParamsDropout, val); - } - FusedAttnBwdParamsWrapper &set_attn_scale(float val) noexcept { - return set_attr(kNVTEFusedAttnBwdParamsAttnScale, val); - } - FusedAttnBwdParamsWrapper &set_qkv_layout(NVTE_QKV_Layout val) noexcept { - return set_attr(kNVTEFusedAttnBwdParamsQKVLayout, val); - } - FusedAttnBwdParamsWrapper &set_o_format(NVTE_QKV_Format val) noexcept { - return set_attr(kNVTEFusedAttnBwdParamsOFormat, val); - } - FusedAttnBwdParamsWrapper &set_do_format(NVTE_QKV_Format val) noexcept { - return set_attr(kNVTEFusedAttnBwdParamsDOFormat, val); - } - FusedAttnBwdParamsWrapper &set_dqkv_layout(NVTE_QKV_Layout val) noexcept { - return set_attr(kNVTEFusedAttnBwdParamsDQKVLayout, val); - } - FusedAttnBwdParamsWrapper &set_qkv_scale_inv_format(NVTE_QKV_Format val) noexcept { - return set_attr(kNVTEFusedAttnBwdParamsQKVScaleInvFormat, val); - } - FusedAttnBwdParamsWrapper &set_do_scale_inv_format(NVTE_QKV_Format val) noexcept { - return set_attr(kNVTEFusedAttnBwdParamsDOScaleInvFormat, val); - } - FusedAttnBwdParamsWrapper &set_max_seqlen_q(size_t val) noexcept { - return set_attr(kNVTEFusedAttnBwdParamsMaxSeqlenQ, val); - } - FusedAttnBwdParamsWrapper &set_max_seqlen_kv(size_t val) noexcept { - return set_attr(kNVTEFusedAttnBwdParamsMaxSeqlenKV, val); - } - FusedAttnBwdParamsWrapper &set_workspace(NVTETensor val) noexcept { - return set_attr(kNVTEFusedAttnBwdParamsWorkspace, val); - } - FusedAttnBwdParamsWrapper &set_stream(cudaStream_t val) noexcept { - return set_attr(kNVTEFusedAttnBwdParamsStream, val); - } - - private: - template - FusedAttnBwdParamsWrapper &set_attr(NVTEFusedAttnBwdParamsAttribute attr, T val) noexcept { - nvte_set_fused_attn_bwd_params_attribute(params_, attr, &val, sizeof(val)); - return *this; - } - - NVTEFusedAttnBwdParams params_ = nullptr; -}; -#endif // __cplusplus - #endif + +#endif // TRANSFORMER_ENGINE_FUSED_ATTN_H_ diff --git a/transformer_engine/common/include/transformer_engine/utils.h b/transformer_engine/common/include/transformer_engine/utils.h index fda49dd5499..69ac9676f34 100644 --- a/transformer_engine/common/include/transformer_engine/utils.h +++ b/transformer_engine/common/include/transformer_engine/utils.h @@ -40,6 +40,24 @@ void nvte_copy_host_to_device_via_kernel(const void *host_ptr, void *device_ptr, void nvte_convert_pointers_to_tensor(const uint64_t *host_ptrs, NVTETensor output, int64_t count, cudaStream_t stream); +/*! \brief Extract a CUDA generator's seed and offset into a device RNG-state buffer. + * + * When a CUDA graph is being captured, the seed and offset are read from device pointers and + * the graph-local offset is added. Otherwise the provided host values are stored directly. + * + * \param[out] rng_state_ptr A two-element device array containing seed and offset. + * \param[in] captured Whether CUDA graph capture is active. + * \param[in] seed_ptr Device pointer to the seed used during capture. + * \param[in] seed_val Seed value used outside capture. + * \param[in] offset_ptr Device pointer to the offset used during capture. + * \param[in] offset_val Offset value used outside capture. + * \param[in] offset_intragraph Offset to add within a captured graph. + * \param[in] stream CUDA stream for the operation. + */ +void nvte_extract_seed_and_offset(int64_t *rng_state_ptr, int captured, int64_t *seed_ptr, + uint64_t seed_val, int64_t *offset_ptr, uint64_t offset_val, + uint32_t offset_intragraph, cudaStream_t stream); + #ifdef __cplusplus } // extern "C" #endif diff --git a/transformer_engine/common/util/pybind_helper.h b/transformer_engine/common/util/pybind_helper.h index d7399651634..f7ffb5ad8db 100644 --- a/transformer_engine/common/util/pybind_helper.h +++ b/transformer_engine/common/util/pybind_helper.h @@ -86,19 +86,11 @@ .value("NVTE_Paged_KV_SBHD_SBHD_SBHD", NVTE_QKV_Layout::NVTE_Paged_KV_SBHD_SBHD_SBHD) \ .value("NVTE_Paged_KV_THD_BSHD_BSHD", NVTE_QKV_Layout::NVTE_Paged_KV_THD_BSHD_BSHD) \ .value("NVTE_Paged_KV_THD_SBHD_SBHD", NVTE_QKV_Layout::NVTE_Paged_KV_THD_SBHD_SBHD) \ - .value("NVTE_BHSD_BHSD_BHSD", NVTE_QKV_Layout::NVTE_BHSD_BHSD_BHSD) \ - .value("NVTE_QKV_Layout_NOT_SET", NVTE_QKV_Layout::NVTE_QKV_Layout_NOT_SET); \ + .value("NVTE_BHSD_BHSD_BHSD", NVTE_QKV_Layout::NVTE_BHSD_BHSD_BHSD); \ pybind11::enum_(m, "NVTE_Fused_Attn_Backend", pybind11::module_local()) \ .value("NVTE_F16_arbitrary_seqlen", NVTE_Fused_Attn_Backend::NVTE_F16_arbitrary_seqlen) \ .value("NVTE_FP8", NVTE_Fused_Attn_Backend::NVTE_FP8) \ .value("NVTE_No_Backend", NVTE_Fused_Attn_Backend::NVTE_No_Backend); \ - pybind11::enum_(m, "NVTEScalingMode", pybind11::module_local()) \ - .value("NVTE_DELAYED_TENSOR_SCALING", NVTEScalingMode::NVTE_DELAYED_TENSOR_SCALING) \ - .value("NVTE_MXFP8_1D_SCALING", NVTEScalingMode::NVTE_MXFP8_1D_SCALING) \ - .value("NVTE_BLOCK_SCALING_1D", NVTEScalingMode::NVTE_BLOCK_SCALING_1D) \ - .value("NVTE_BLOCK_SCALING_2D", NVTEScalingMode::NVTE_BLOCK_SCALING_2D) \ - .value("NVTE_NVFP4_1D_SCALING", NVTEScalingMode::NVTE_NVFP4_1D_SCALING) \ - .value("NVTE_INVALID_SCALING", NVTEScalingMode::NVTE_INVALID_SCALING); \ pybind11::enum_( \ m, "Float8BlockScaleTensorFormat", pybind11::module_local()) \ .value("GEMM_READY", transformer_engine::Float8BlockScaleTensorFormat::GEMM_READY) \ diff --git a/transformer_engine/common/util/utils.cu b/transformer_engine/common/util/utils.cu index 39d82624632..f659759163d 100644 --- a/transformer_engine/common/util/utils.cu +++ b/transformer_engine/common/util/utils.cu @@ -14,6 +14,22 @@ #include "../util/logging.h" namespace transformer_engine { +namespace extract_seed_and_offset { +namespace { + +__global__ void kernel(int64_t *rng_state_ptr, bool captured, int64_t *seed_ptr, uint64_t seed_val, + int64_t *offset_ptr, uint64_t offset_val, uint32_t offset_intragraph) { + if (captured) { + rng_state_ptr[0] = *seed_ptr; + rng_state_ptr[1] = static_cast(*offset_ptr + static_cast(offset_intragraph)); + } else { + rng_state_ptr[0] = static_cast(seed_val); + rng_state_ptr[1] = static_cast(offset_val); + } +} + +} // namespace +} // namespace extract_seed_and_offset namespace copy_host_to_device_via_kernel { namespace { @@ -80,3 +96,12 @@ void nvte_convert_pointers_to_tensor(const uint64_t *host_ptrs, NVTETensor outpu nvte_copy_host_to_device_via_kernel(host_ptrs, out_tensor->data.dptr, static_cast(count) * sizeof(uint64_t), stream); } + +void nvte_extract_seed_and_offset(int64_t *rng_state_ptr, int captured, int64_t *seed_ptr, + uint64_t seed_val, int64_t *offset_ptr, uint64_t offset_val, + uint32_t offset_intragraph, cudaStream_t stream) { + NVTE_API_CALL(nvte_extract_seed_and_offset); + transformer_engine::extract_seed_and_offset::kernel<<<1, 1, 0, stream>>>( + rng_state_ptr, captured, seed_ptr, seed_val, offset_ptr, offset_val, offset_intragraph); + NVTE_CHECK_CUDA(cudaGetLastError()); +} diff --git a/transformer_engine/jax/attention.py b/transformer_engine/jax/attention.py index dcf3a955054..2a0e8c2fd36 100644 --- a/transformer_engine/jax/attention.py +++ b/transformer_engine/jax/attention.py @@ -20,6 +20,7 @@ from transformer_engine_jax import NVTE_Softmax_Type from . import cpp_extensions as tex +from .quantize import AttentionQuantizerSet class AttnBiasType(Enum): @@ -32,6 +33,8 @@ class AttnBiasType(Enum): NO_BIAS = NVTE_Bias_Type.NVTE_NO_BIAS PRE_SCALE_BIAS = NVTE_Bias_Type.NVTE_PRE_SCALE_BIAS POST_SCALE_BIAS = NVTE_Bias_Type.NVTE_POST_SCALE_BIAS + # Keep source-tree imports working before the local extension is rebuilt. + ALIBI = getattr(NVTE_Bias_Type, "NVTE_ALIBI", 3) class AttnMaskType(Enum): @@ -325,7 +328,6 @@ def canonicalize_attn_mask_type(attn_mask_type: str): def is_fused_attn_kernel_available( is_training, - batch_size, q_dtype, kv_dtype, qkv_layout, @@ -341,52 +343,33 @@ def is_fused_attn_kernel_available( head_dim_v, window_size: Optional[Tuple[int, int]] = None, return_max_logit: bool = False, - bottom_right_diagonal: Optional[bool] = None, - bias_batch: Optional[int] = None, - bias_heads: Optional[int] = None, - bias_seqlen_q: Optional[int] = None, - bias_seqlen_kv: Optional[int] = None, - max_segments_per_seq: int = 1, ): """ - To check whether the fused attention kernel is supported. + To check whether the fused attention kernel is supported """ window_size_tuple = (-1, -1) if window_size is None else window_size def make_helper(attn_mask_type): - bottom_right = ( - attn_mask_type.is_bottom_right() - if bottom_right_diagonal is None - else bottom_right_diagonal - ) return tex.FusedAttnHelper( - is_training=is_training, - batch_size=batch_size, - q_dtype=q_dtype, - kv_dtype=kv_dtype, - qkv_layout=qkv_layout, - attn_bias_type=attn_bias_type, - attn_mask_type=attn_mask_type, - softmax_type=softmax_type, - dropout_probability=dropout_probability, - q_num_heads=q_num_heads, - kv_num_heads=kv_num_heads, - q_max_seqlen=q_max_seqlen, - kv_max_seqlen=kv_max_seqlen, - head_dim_qk=head_dim_qk, - head_dim_v=head_dim_v, - window_size=window_size_tuple, - return_max_logit=return_max_logit, - bottom_right_diagonal=bottom_right, - bias_batch=bias_batch, - bias_heads=bias_heads, - bias_seqlen_q=bias_seqlen_q, - bias_seqlen_kv=bias_seqlen_kv, - max_segments_per_seq=max_segments_per_seq, + is_training, + q_dtype, + kv_dtype, + qkv_layout, + attn_bias_type, + attn_mask_type, + softmax_type, + dropout_probability, + q_num_heads, + kv_num_heads, + q_max_seqlen, + kv_max_seqlen, + head_dim_qk, + head_dim_v, + window_size_tuple, + return_max_logit, ) - helper = make_helper(attn_mask_type) - return helper.is_fused_attn_kernel_available() + return make_helper(attn_mask_type).is_fused_attn_kernel_available() def _obtain_batch_and_max_seqlen(qkv, qkv_layout): @@ -1076,6 +1059,7 @@ def _legacy_fused_attn( context_parallel_axis: str = "", softmax_offset: Optional[jnp.ndarray] = None, return_max_logit: bool = False, + bottom_right_diagonal: Optional[bool] = None, ): """ Perform non-THD (non-packed) cuDNN fused attention. @@ -1170,6 +1154,11 @@ def _legacy_fused_attn( context_parallel_causal_load_balanced=context_parallel_causal_load_balanced, context_parallel_axis=context_parallel_axis, return_max_logit=return_max_logit, + bottom_right_diagonal=( + attn_mask_type.is_bottom_right() + if bottom_right_diagonal is None + else bottom_right_diagonal + ), ) return output @@ -1258,7 +1247,7 @@ def fused_attn_thd( @partial( jax.custom_vjp, - nondiff_argnums=(5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19), + nondiff_argnums=(5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20), ) def _fused_attn( qkv: Tuple[jnp.ndarray, ...], @@ -1281,6 +1270,7 @@ def _fused_attn( context_checkpoint_name: str = "context", stripe_size: int | None = None, return_max_logit: bool = False, + bottom_right_diagonal: bool = False, ): output, _ = _fused_attn_fwd_rule( qkv, @@ -1303,6 +1293,7 @@ def _fused_attn( context_checkpoint_name=context_checkpoint_name, stripe_size=stripe_size, return_max_logit=return_max_logit, + bottom_right_diagonal=bottom_right_diagonal, ) return output @@ -1328,6 +1319,7 @@ def _fused_attn_fwd_rule( context_checkpoint_name, stripe_size, return_max_logit, + bottom_right_diagonal, ): output, softmax_aux, rng_state, max_logit = tex.fused_attn_fwd( qkv, @@ -1349,6 +1341,7 @@ def _fused_attn_fwd_rule( context_parallel_axis=context_parallel_axis, stripe_size=stripe_size, return_max_logit=return_max_logit, + bottom_right_diagonal=bottom_right_diagonal, ) output = checkpoint_name(output, context_checkpoint_name) softmax_aux = checkpoint_name(softmax_aux, context_checkpoint_name) @@ -1382,6 +1375,7 @@ def _fused_attn_bwd_rule( context_checkpoint_name, stripe_size, return_max_logit, + bottom_right_diagonal, ctx, dz, ): @@ -1419,8 +1413,9 @@ def _fused_attn_bwd_rule( context_parallel_causal_load_balanced=context_parallel_causal_load_balanced, context_parallel_axis=context_parallel_axis, stripe_size=stripe_size, + bottom_right_diagonal=bottom_right_diagonal, ) - if attn_bias_type == AttnBiasType.NO_BIAS: + if attn_bias_type in (AttnBiasType.NO_BIAS, AttnBiasType.ALIBI): grad_bias = None if softmax_type != AttnSoftmaxType.LEARNABLE_SOFTMAX: grad_softmax_offset = None @@ -1436,6 +1431,24 @@ def _fused_attn_bwd_rule( _fused_attn.defvjp(_fused_attn_fwd_rule, _fused_attn_bwd_rule) +@partial(jax.custom_vjp, nondiff_argnums=(4,)) +def _fused_attn_fp8(qkv, sequence_descriptor, seed, quantizer_set, config): + output, _ = tex.fused_attn_fp8_fwd(qkv, sequence_descriptor, seed, quantizer_set, config) + return output + + +def _fused_attn_fp8_fwd_rule(qkv, sequence_descriptor, seed, quantizer_set, config): + return tex.fused_attn_fp8_fwd(qkv, sequence_descriptor, seed, quantizer_set, config) + + +def _fused_attn_fp8_bwd_rule(config, ctx, doutput): + grad_qkv, quantizer_set = tex.fused_attn_fp8_bwd(ctx, doutput, config) + return grad_qkv, None, None, quantizer_set + + +_fused_attn_fp8.defvjp(_fused_attn_fp8_fwd_rule, _fused_attn_fp8_bwd_rule) + + @partial(jax.custom_vjp, nondiff_argnums=(3, 4)) def _fused_attn_score_mod( qkv: Tuple[jnp.ndarray, jnp.ndarray, jnp.ndarray], @@ -1514,6 +1527,8 @@ def fused_attn( score_mod_tensors: Optional[Mapping[str, Any]] = None, score_mod_bprop_tensors: Optional[Mapping[str, Any]] = None, return_max_logit: bool = False, + bottom_right_diagonal: Optional[bool] = None, + quantizer_set: Optional[AttentionQuantizerSet] = None, ): """ Perform cuDNN fused attention. @@ -1572,6 +1587,11 @@ def fused_attn( Python/NumPy scalars made available to `score_mod_bprop`. return_max_logit (bool): If True, also return per-head maximum attention logits with shape ``[h]``. + bottom_right_diagonal (Optional[bool]): Explicitly select bottom-right diagonal + alignment independently of the mask type. By default it follows bottom-right masks. + quantizer_set (Optional[AttentionQuantizerSet]): Quantizers for FP8 DPA. When + provided, Q/K/V and dO are quantized internally while public inputs and outputs + remain FP16/BF16. Returns: jnp.ndarray: Attention output when ``return_max_logit`` is False. @@ -1620,6 +1640,8 @@ def fused_attn( if score_mod_only_args: raise ValueError(f"{', '.join(score_mod_only_args)} require score_mod to be provided.") else: + if quantizer_set is not None: + raise NotImplementedError("JAX FP8 attention does not support score_mod.") if return_max_logit: raise ValueError("return_max_logit is not supported with score_mod fused_attn.") tex.validate_fused_attn_score_mod( @@ -1657,6 +1679,10 @@ def fused_attn( ) if sequence_descriptor is None or isinstance(sequence_descriptor, jnp.ndarray): + if quantizer_set is not None: + raise NotImplementedError( + "JAX FP8 attention requires a SequenceDescriptor instead of a legacy mask." + ) warnings.warn( "Pass mask to fused_attn is deprecated, please use SequenceDescriptor instead. " + "See help(transformer_engine.jax.attention.SequenceDescriptor) for details.", @@ -1682,6 +1708,7 @@ def fused_attn( context_parallel_axis=context_parallel_axis, softmax_offset=softmax_offset, return_max_logit=return_max_logit, + bottom_right_diagonal=bottom_right_diagonal, ) if max_segments_per_seq > 1 and not qkv_layout.is_thd(): warnings.warn( @@ -1692,6 +1719,36 @@ def fused_attn( UserWarning, stacklevel=2, ) + if quantizer_set is not None: + if not isinstance(quantizer_set, AttentionQuantizerSet): + raise TypeError("quantizer_set must be an AttentionQuantizerSet.") + if qkv_layout.is_thd() or max_segments_per_seq != 1: + raise NotImplementedError("JAX FP8 attention does not support THD sequence packing.") + if context_parallel_axis or context_parallel_strategy != CPStrategy.DEFAULT: + raise NotImplementedError("JAX FP8 attention does not support context parallelism.") + if bias is not None or attn_bias_type != AttnBiasType.NO_BIAS: + raise NotImplementedError("JAX FP8 attention does not support attention bias.") + if return_max_logit: + raise NotImplementedError("JAX FP8 attention does not support return_max_logit.") + if softmax_offset is not None: + raise NotImplementedError("JAX FP8 attention does not support a softmax offset.") + diagonal = ( + attn_mask_type.is_bottom_right() + if bottom_right_diagonal is None + else bottom_right_diagonal + ) + config = tex.FP8AttentionConfig( + attn_bias_type=attn_bias_type, + attn_mask_type=attn_mask_type, + softmax_type=softmax_type, + qkv_layout=qkv_layout, + scaling_factor=scaling_factor, + dropout_probability=dropout_probability, + is_training=is_training, + window_size=(-1, -1) if window_size is None else window_size, + bottom_right_diagonal=diagonal, + ) + return _fused_attn_fp8(qkv, sequence_descriptor, seed, quantizer_set, config) output = _fused_attn( qkv, bias, @@ -1713,5 +1770,10 @@ def fused_attn( context_checkpoint_name=context_checkpoint_name, stripe_size=stripe_size, return_max_logit=return_max_logit, + bottom_right_diagonal=( + attn_mask_type.is_bottom_right() + if bottom_right_diagonal is None + else bottom_right_diagonal + ), ) return output diff --git a/transformer_engine/jax/cpp_extensions/__init__.py b/transformer_engine/jax/cpp_extensions/__init__.py index c9647afb826..68eaa7d6776 100644 --- a/transformer_engine/jax/cpp_extensions/__init__.py +++ b/transformer_engine/jax/cpp_extensions/__init__.py @@ -5,6 +5,7 @@ from .activation import * from .amax import * from .attention import * +from .fp8_attention import * from .flex_attention import * from .normalization import * from .quantization import * diff --git a/transformer_engine/jax/cpp_extensions/attention.py b/transformer_engine/jax/cpp_extensions/attention.py index c380159a46f..c9745675aaf 100644 --- a/transformer_engine/jax/cpp_extensions/attention.py +++ b/transformer_engine/jax/cpp_extensions/attention.py @@ -3,62 +3,56 @@ # See LICENSE for license information. """JAX/TE custom ops for attention""" import logging -import operator import os import warnings from dataclasses import dataclass, replace -from functools import partial, reduce +from functools import partial from typing import Optional, Tuple import jax import jax.numpy as jnp -from jax import dtypes, lax, ffi -from jax.sharding import PartitionSpec, NamedSharding +from jax import dtypes, ffi, lax from jax.experimental.custom_partitioning import SdyShardingRule +from jax.sharding import NamedSharding, PartitionSpec +from transformer_engine_jax import NVTE_Fused_Attn_Backend -import transformer_engine_jax -from transformer_engine_jax import ( - DType, - JAXX_Scaling_Mode, - NVTE_Bias_Type, - NVTE_Fused_Attn_Backend, - NVTE_Mask_Type, - NVTE_QKV_Format, - NVTE_QKV_Layout, - NVTE_Softmax_Type, -) from transformer_engine.jax.attention import ( AttnBiasType, AttnMaskType, AttnSoftmaxType, - QKVLayout, - QKVFormat, CPStrategy, + QKVFormat, + QKVLayout, SequenceDescriptor, ) -from ..sharding import with_sharding_constraint_by_logical_axes, HEAD_AXES, is_mesh_available -from .base import BasePrimitive, register_primitive -from .misc import ( - check_valid_batch_dims, - jax_dtype_to_te_dtype, - te_dtype_to_jax_dtype, - get_padded_spec, - get_cudnn_version, - get_all_device_compute_capability, -) from ..sharding import ( - global_mesh_resource, - lax_paral_op, + HEAD_AXES, all_reduce_sum_along_dp_fsdp, - get_mesh_axis_size, + get_all_mesh_axes, get_mesh_axis_rank, get_mesh_axis_rank_host, - get_all_mesh_axes, + get_mesh_axis_size, + global_mesh_resource, + is_mesh_available, + lax_paral_op, num_of_devices, with_sharding_constraint, + with_sharding_constraint_by_logical_axes, +) +from .base import BasePrimitive, register_primitive +from .cudnn_attention import ( + build_bwd_graph, + build_fwd_graph, + get_fused_attn_support, + ragged_graph_batch_size, +) +from .misc import ( + check_valid_batch_dims, + get_all_device_compute_capability, + get_cudnn_version, + get_padded_spec, ) - __all__ = [ "FusedAttnHelper", @@ -74,7 +68,7 @@ class AttentionLogging: - """Logging for the JAX attention module""" + """Logging for the JAX attention module.""" _log_level = _NVTE_DEBUG * _NVTE_DEBUG_LEVEL _formatter = logging.Formatter("[%(levelname)-8s | %(name)-19s]: %(message)s") @@ -84,15 +78,13 @@ class AttentionLogging: @staticmethod def setup_logging(): - """Set up log levels, logger and handlers (idempotent).""" + """Set up log levels, logger, and handlers.""" if AttentionLogging._is_logging_setup: return - _log_levels = {0: logging.WARNING, 1: logging.INFO, 2: logging.DEBUG} - AttentionLogging._log_level = _log_levels[ - AttentionLogging._log_level if AttentionLogging._log_level in [0, 1, 2] else 2 - ] + log_levels = {0: logging.WARNING, 1: logging.INFO, 2: logging.DEBUG} + level = AttentionLogging._log_level if AttentionLogging._log_level in log_levels else 2 AttentionLogging._stream_handler.setFormatter(AttentionLogging._formatter) - AttentionLogging.logger.setLevel(AttentionLogging._log_level) + AttentionLogging.logger.setLevel(log_levels[level]) if not AttentionLogging.logger.hasHandlers(): AttentionLogging.logger.addHandler(AttentionLogging._stream_handler) AttentionLogging._is_logging_setup = True @@ -143,80 +135,14 @@ class _FusedAttnConfig: ) # Only for CP + Striped. For Ring P2P, stripe_size=1 only.For AG, stripe_size>=1. return_max_logit: bool = False - @property - def effective_window_size(self) -> Tuple[int, int]: - """Derive the effective window size that the kernel runs in CP + Ring + THD + SWA case.""" - if self.cp_striped_window_size is not None: - return self.cp_striped_window_size - return self.window_size - -@dataclass -class FusedAttnParams: - """ - Attention parameters used to select the fused attention backend. - - Fields are declared in the order of the ``FusedAttnConfig`` struct in - ``common/fused_attn/config_and_params.h``, which is the order the C++ binding reads them in - and the order it fills the config with. Fields JAX does not use, are left at their ``FusedAttnConfig`` - defaults. - """ - - # basic attention settings - is_training: bool = True - deterministic: bool = False - cuda_graph: bool = False - return_max_logit: bool = False - attn_mask_type: NVTE_Mask_Type = NVTE_Mask_Type.NVTE_NO_MASK - bias_type: NVTE_Bias_Type = NVTE_Bias_Type.NVTE_NO_BIAS - window_size_left: int = -1 - window_size_right: int = -1 - bottom_right_diagonal: bool = True - softmax_type: NVTE_Softmax_Type = NVTE_Softmax_Type.NVTE_VANILLA_SOFTMAX - scaling_mode: JAXX_Scaling_Mode = JAXX_Scaling_Mode.NO_SCALING - dropout: float = 0.0 - attn_scale: float = 1.0 - - # tensor types - qkv_dtype: DType = DType.kBFloat16 - o_dtype: DType = DType.kBFloat16 - do_dtype: DType = DType.kBFloat16 - dqkv_dtype: DType = DType.kBFloat16 - - # tensor layouts - qkv_layout: NVTE_QKV_Layout = NVTE_QKV_Layout.NVTE_QKV_Layout_NOT_SET - o_format: NVTE_QKV_Format = NVTE_QKV_Format.NVTE_QKV_Format_NOT_SET - do_format: NVTE_QKV_Format = NVTE_QKV_Format.NVTE_QKV_Format_NOT_SET - dqkv_layout: NVTE_QKV_Layout = NVTE_QKV_Layout.NVTE_QKV_Layout_NOT_SET - qkv_scale_inv_format: NVTE_QKV_Format = NVTE_QKV_Format.NVTE_QKV_Format_NOT_SET - do_scale_inv_format: NVTE_QKV_Format = NVTE_QKV_Format.NVTE_QKV_Format_NOT_SET - - # tensor dimensions - batch_size: int = 0 - num_attn_heads: int = 0 - num_gqa_groups: int = 0 - head_dim_qk: int = 0 - head_dim_v: int = 0 - max_seqlen_q: int = 0 - max_seqlen_kv: int = 0 - num_tokens_q: int = 0 - num_tokens_kv: int = 0 - - # bias dimensions - bias_batch_size: int = 0 - bias_num_heads: int = 0 - bias_seqlen_q: int = 0 - bias_seqlen_kv: int = 0 - - -@dataclass(frozen=True, kw_only=True) +@dataclass(frozen=True) class FusedAttnHelper: """ Helper for the fused attention backend """ is_training: bool - batch_size: int q_dtype: jnp.dtype kv_dtype: jnp.dtype qkv_layout: QKVLayout @@ -232,13 +158,6 @@ class FusedAttnHelper: head_dim_v: int window_size: Tuple[int, int] return_max_logit: bool = False - bottom_right_diagonal: bool = True - attn_scale: float = 1.0 - bias_batch: Optional[int] = None - bias_heads: Optional[int] = None - bias_seqlen_q: Optional[int] = None - bias_seqlen_kv: Optional[int] = None - max_segments_per_seq: int = 1 def is_fused_attn_kernel_available(self): """Check if there is available fused attention kernel""" @@ -246,78 +165,23 @@ def is_fused_attn_kernel_available(self): return backend != NVTE_Fused_Attn_Backend.NVTE_No_Backend def get_fused_attn_backend(self): - """Get the fused attention kernel backend. - - Returns a ``(backend, message)`` tuple. ``message`` is empty on success, otherwise a - diagnostic string explaining why the configuration was rejected. - - When ``NVTE_DEBUG=1``, ``NVTE_DEBUG_LEVEL=1`` logs the outcome (the selected backend, or - that no fused backend is available), and ``NVTE_DEBUG_LEVEL=2`` additionally logs the - resolved config and the reason fused attention was rejected. - """ - q_type = jax_dtype_to_te_dtype(self.q_dtype) - kv_type = jax_dtype_to_te_dtype(self.kv_dtype) - if q_type != kv_type: - raise ValueError("Q and KV must have the same data type.") - bias_batch = bias_heads = bias_seqlen_q = bias_seqlen_kv = 0 - if self.attn_bias_type == AttnBiasType.POST_SCALE_BIAS: - bias_batch = self.bias_batch or 0 - bias_heads = self.bias_heads or 0 - bias_seqlen_q = self.bias_seqlen_q or 0 - bias_seqlen_kv = self.bias_seqlen_kv or 0 - num_segments = self.batch_size - num_tokens_q = num_tokens_kv = 0 - if self.qkv_layout.is_thd(): - num_segments = self.batch_size * self.max_segments_per_seq - num_tokens_q = self.batch_size * self.q_max_seqlen - num_tokens_kv = self.batch_size * self.kv_max_seqlen - backend, message = transformer_engine_jax.get_fused_attn_backend( - FusedAttnParams( - is_training=self.is_training, - deterministic=not self.is_non_deterministic_allowed(), - return_max_logit=self.return_max_logit, - attn_mask_type=self.attn_mask_type.value, - bias_type=self.attn_bias_type.value, - window_size_left=self.window_size[0], - window_size_right=self.window_size[1], - bottom_right_diagonal=self.bottom_right_diagonal, - softmax_type=self.softmax_type.value, - dropout=self.dropout_probability, - attn_scale=self.attn_scale, - qkv_dtype=q_type, - o_dtype=q_type, - do_dtype=q_type, - dqkv_dtype=q_type, - qkv_layout=self.qkv_layout.value, - batch_size=num_segments, - num_attn_heads=self.q_num_heads, - num_gqa_groups=self.kv_num_heads, - head_dim_qk=self.head_dim_qk, - head_dim_v=self.head_dim_v, - max_seqlen_q=self.q_max_seqlen, - max_seqlen_kv=self.kv_max_seqlen, - num_tokens_q=num_tokens_q, - num_tokens_kv=num_tokens_kv, - bias_batch_size=bias_batch, - bias_num_heads=bias_heads, - bias_seqlen_q=bias_seqlen_q, - bias_seqlen_kv=bias_seqlen_kv, - ) + """Get the fused attention backend and a rejection reason when unavailable.""" + support = get_fused_attn_support(self) + backend = ( + NVTE_Fused_Attn_Backend.NVTE_F16_arbitrary_seqlen + if support.supported + else NVTE_Fused_Attn_Backend.NVTE_No_Backend ) + message = "" if support.supported else support.reason AttentionLogging.setup_logging() logger = AttentionLogging.logger logger.debug("Running fused attention backend selection with config=%s", self) if backend == NVTE_Fused_Attn_Backend.NVTE_No_Backend: logger.info("No fused attention backend available; falling back to unfused attention.") - logger.debug( - "Reason fused attention was rejected: %s", - message or "(no diagnostic message available)", - ) + logger.debug("Reason fused attention was rejected: %s", message) else: logger.info("Selected fused attention backend: %s", backend) - if message: - logger.debug("Fused attention backend diagnostic message: %s", message) return backend, message @staticmethod @@ -438,13 +302,73 @@ def check_seed(self, seed, dropout_probability, is_training): return seed -def generate_cu_seqlen(actual_seqlen): - """ - Generating cumsum seqlen for a batch - """ - actual_seqlen = jnp.where(actual_seqlen < 0, 0, actual_seqlen) - cu_seqlen = jnp.cumulative_sum(actual_seqlen, include_initial=True) - return cu_seqlen +def _multiply_offsets_as_uint64_words(offsets, multiplier): + """Return unsigned 64-bit products as interleaved low/high uint32 words.""" + values = offsets.astype(jnp.uint32) + scale = jnp.asarray(multiplier, dtype=jnp.uint32) + mask = jnp.asarray(0xFFFF, dtype=jnp.uint32) + + value_lo = values & mask + value_hi = values >> 16 + scale_lo = scale & mask + scale_hi = scale >> 16 + + product_lo = value_lo * scale_lo + product_mid_lo = value_lo * scale_hi + product_mid_hi = value_hi * scale_lo + product_hi = value_hi * scale_hi + carry = (product_lo >> 16) + (product_mid_lo & mask) + (product_mid_hi & mask) + + low_word = (product_lo & mask) | ((carry & mask) << 16) + high_word = product_hi + (product_mid_lo >> 16) + (product_mid_hi >> 16) + (carry >> 16) + return jnp.stack((low_word, high_word), axis=-1) + + +def _pack_ragged_offsets( + q_seq_offsets, + k_seq_offsets, + qkv_layout, + attn_heads, + num_gqa_groups, + q_head_dim, + v_head_dim, +): + """Pack Q/K/V/O/Stats element offsets into one JAX buffer.""" + q_multiplier = attn_heads * q_head_dim + if qkv_layout.is_qkvpacked(): + q_multiplier *= 3 + k_multiplier = v_multiplier = q_multiplier + elif qkv_layout.is_kvpacked(): + k_multiplier = v_multiplier = 2 * num_gqa_groups * q_head_dim + else: + k_multiplier = num_gqa_groups * q_head_dim + v_multiplier = num_gqa_groups * v_head_dim + offsets_and_multipliers = ( + (q_seq_offsets, q_multiplier), + (k_seq_offsets, k_multiplier), + (k_seq_offsets, v_multiplier), + (q_seq_offsets, attn_heads * v_head_dim), + (q_seq_offsets, attn_heads), + ) + if get_cudnn_version() < (9, 5, 0): + return jnp.stack( + tuple( + offsets.astype(jnp.int32) * jnp.asarray(multiplier, dtype=jnp.int32) + for offsets, multiplier in offsets_and_multipliers + ) + ) + return jnp.stack( + tuple( + _multiply_offsets_as_uint64_words(offsets, multiplier) + for offsets, multiplier in offsets_and_multipliers + ) + ) + + +def _pad_ragged_metadata(values, size, fill_value): + """Pad compact sequence metadata to the bucketed cuDNN graph extent.""" + values = values.flatten()[:size] + return jnp.pad(values, (0, size - values.size), constant_values=fill_value) class FusedAttnFwdPrimitive(BasePrimitive): @@ -505,67 +429,31 @@ def abstract( output_shape = (*batch_shape, q_max_seqlen, attn_heads, v_head_dim) out_aval = q_aval.update(shape=output_shape, dtype=q_dtype) - # backend determines the softmax buffer shape/dtype - input_batch = reduce(operator.mul, batch_shape) - bias_batch = bias_heads = bias_seqlen_q = bias_seqlen_kv = None - if config.attn_bias_type == AttnBiasType.POST_SCALE_BIAS: - *bias_batch_shape, bias_heads, bias_seqlen_q, bias_seqlen_kv = bias_aval.shape - bias_batch = reduce(operator.mul, bias_batch_shape) backend, message = FusedAttnHelper( - is_training=config.is_training, - batch_size=input_batch, - q_dtype=q_dtype, - kv_dtype=k_dtype, - qkv_layout=config.qkv_layout, - attn_bias_type=config.attn_bias_type, - attn_mask_type=config.attn_mask_type, - softmax_type=config.softmax_type, - dropout_probability=config.dropout_probability, - q_num_heads=attn_heads, - kv_num_heads=num_gqa_groups, - q_max_seqlen=q_max_seqlen, - kv_max_seqlen=kv_max_seqlen, - head_dim_qk=q_head_dim, - head_dim_v=v_head_dim, - window_size=config.effective_window_size, - return_max_logit=config.return_max_logit, - bottom_right_diagonal=config.bottom_right_diagonal, - attn_scale=float(config.scaling_factor), - bias_batch=bias_batch, - bias_heads=bias_heads, - bias_seqlen_q=bias_seqlen_q, - bias_seqlen_kv=bias_seqlen_kv, - max_segments_per_seq=config.max_segments_per_seq, + config.is_training, + q_dtype, + k_dtype, + config.qkv_layout, + config.attn_bias_type, + config.attn_mask_type, + config.softmax_type, + config.dropout_probability, + attn_heads, + num_gqa_groups, + q_max_seqlen, + kv_max_seqlen, + q_head_dim, + v_head_dim, + config.window_size, + config.return_max_logit, ).get_fused_attn_backend() + if backend != NVTE_Fused_Attn_Backend.NVTE_F16_arbitrary_seqlen: + raise ValueError(f"Unsupported {backend=}: {message}") - if backend == NVTE_Fused_Attn_Backend.NVTE_F16_arbitrary_seqlen: - # cuDNN 9.6 reduces the required softmax shape - if get_cudnn_version() >= (9, 6, 0): - if config.qkv_layout.is_thd(): - softmax_shape = (*batch_shape, q_max_seqlen, attn_heads, 1) - else: - softmax_shape = (*batch_shape, attn_heads, q_max_seqlen, 1) - else: - softmax_shape = ( - *batch_shape, - attn_heads, - q_max_seqlen, - config.max_segments_per_seq, - ) - softmax_dtype = dtypes.canonicalize_dtype(jnp.float32) - else: - raise ValueError(f"Unsupported backend: {message}") - softmax_aux_aval = q_aval.update(shape=softmax_shape, dtype=softmax_dtype) - if config.return_max_logit: - # cuDNN Max is row-wise over S_kv. Dense and SM120 THD use - # [..., H, S_q, 1]; cuDNN >= 9.6 non-SM120 THD uses [..., S_q, H, 1]. - # Both raw layouts are reduced to the public per-head [H] result below. - if FusedAttnFwdPrimitive._uses_thd_ragged_max_tensor(config): - max_tensor_shape = (*batch_shape, q_max_seqlen, attn_heads, 1) - else: - max_tensor_shape = (*batch_shape, attn_heads, q_max_seqlen, 1) - else: - max_tensor_shape = (0,) + graph_info = build_fwd_graph(q_aval, k_aval, v_aval, bias_aval, config) + softmax_dtype = dtypes.canonicalize_dtype(jnp.float32) + softmax_aux_aval = q_aval.update(shape=graph_info.stats_shape, dtype=softmax_dtype) + max_tensor_shape = graph_info.max_shape max_tensor_aval = q_aval.update(shape=max_tensor_shape, dtype=softmax_dtype) # JAX does not enable 64-bit int by default so we get XLA to allocate x8 memory with @@ -578,42 +466,7 @@ def abstract( rng_state_shape = (seed_aval.shape[0], checker.rng_state_size) rng_state_aval = seed_aval.update(shape=rng_state_shape, dtype=checker.rng_state_dtype) - if config.attn_bias_type == AttnBiasType.NO_BIAS: - bias_batch = bias_heads = 0 - else: - *bias_batch_shape, bias_heads, _, _ = bias_aval.shape - bias_batch = reduce(operator.mul, bias_batch_shape) - - # do a dummy kernel call here to get workspace buffer shapes/dtypes that XLA needs to - # prepare for the active fused-attn backend - input_batch = reduce(operator.mul, batch_shape) - wkspace_info = transformer_engine_jax.get_fused_attn_fwd_workspace_sizes( - input_batch, - bias_batch, - q_max_seqlen, - kv_max_seqlen, - attn_heads, - num_gqa_groups, - bias_heads, - q_head_dim, - v_head_dim, - config.scaling_factor, - config.dropout_probability, - config.attn_bias_type.value, - config.attn_mask_type.value, - config.softmax_type.value, - config.qkv_layout.value, - jax_dtype_to_te_dtype(q_aval.dtype), - config.is_training, - config.max_segments_per_seq, - config.effective_window_size[0], - config.effective_window_size[1], - config.return_max_logit, - config.bottom_right_diagonal, - ) - wkspace_aval = q_aval.update( - shape=wkspace_info[0], dtype=te_dtype_to_jax_dtype(wkspace_info[1]) - ) + wkspace_aval = q_aval.update(shape=(graph_info.graph.workspace_size,), dtype=jnp.uint8) assert ( softmax_offset_aval.dtype == jnp.float32 @@ -668,25 +521,7 @@ def lowering( """ q_aval, k_aval, v_aval, bias_aval, *_ = ctx.avals_in - ( - batch_shape, - q_max_seqlen, - kv_max_seqlen, - attn_heads, - num_gqa_groups, - q_head_dim, - v_head_dim, - ) = FusedAttnHelper.parse_qkv_aval(q_aval, k_aval, v_aval, config.qkv_layout) - - input_batch = reduce(operator.mul, batch_shape) - - if config.attn_bias_type == AttnBiasType.NO_BIAS: - bias_batch = bias_heads = 0 - else: - *bias_batch_shape, bias_heads, _, _ = bias_aval.shape - bias_batch = reduce(operator.mul, bias_batch_shape) - - window_size_left, window_size_right = config.effective_window_size + graph = build_fwd_graph(q_aval, k_aval, v_aval, bias_aval, config).graph return ffi.ffi_lowering(FusedAttnFwdPrimitive.name)( ctx, @@ -704,28 +539,9 @@ def lowering( _kv_segment_ids, _q_segment_pos, _kv_segment_pos, # ffi_lowering needs number of parameters meets primitive.lowering - input_batch=input_batch, - bias_batch=bias_batch, - q_max_seqlen=q_max_seqlen, - kv_max_seqlen=kv_max_seqlen, - attn_heads=attn_heads, - num_gqa_groups=num_gqa_groups, - bias_heads=bias_heads, - qk_head_dim=q_head_dim, - v_head_dim=v_head_dim, - max_segments_per_seq=config.max_segments_per_seq, - scaling_factor=float(config.scaling_factor), - dropout_probability=float(config.dropout_probability), - bias_type=int(config.attn_bias_type.value), - mask_type=int(config.attn_mask_type.value), - qkv_layout=int(config.qkv_layout.value), - is_training=config.is_training, - return_max_logit=config.return_max_logit, - deterministic=not FusedAttnHelper.is_non_deterministic_allowed(), - window_size_left=window_size_left, - window_size_right=window_size_right, - bottom_right_diagonal=config.bottom_right_diagonal, - softmax_type=int(config.softmax_type.value), + is_ragged=config.qkv_layout.is_thd(), + rng_offset_increment=16, + **graph.ffi_attrs(), ) @staticmethod @@ -785,9 +601,15 @@ def convert_to_2d(offsets, batch, max_seqlen): ) return offsets_2d - batch, q_max_seqlen, kv_max_seqlen, *_ = FusedAttnHelper.parse_qkv_aval( - q, k, v, config.qkv_layout - ) + ( + batch, + q_max_seqlen, + kv_max_seqlen, + attn_heads, + num_gqa_groups, + q_head_dim, + v_head_dim, + ) = FusedAttnHelper.parse_qkv_aval(q, k, v, config.qkv_layout) assert len(batch) == 1, f"Expected len(batch) == 1, but got {len(batch)=}" kv_batch = q_batch = batch[0] @@ -820,8 +642,28 @@ def convert_to_2d(offsets, batch, max_seqlen): k_seq_offsets, k_seq_offsets >= 0, fill_value=kv_batch * kv_max_seqlen ) - q_cu_seqlen = generate_cu_seqlen(q_seqlen.flatten()) - kv_cu_seqlen = generate_cu_seqlen(kv_seqlen.flatten()) + graph_batch = ragged_graph_batch_size(q_batch, config.max_segments_per_seq) + q_seqlen = _pad_ragged_metadata(q_seqlen, graph_batch, fill_value) + kv_seqlen = _pad_ragged_metadata(kv_seqlen, graph_batch, fill_value) + q_seq_offsets = _pad_ragged_metadata( + q_seq_offsets, graph_batch + 1, q_batch * q_max_seqlen + ) + k_seq_offsets = _pad_ragged_metadata( + k_seq_offsets, graph_batch + 1, kv_batch * kv_max_seqlen + ) + + # Supported cuDNN ragged graphs require external element offsets. JAX disables + # x64 by default, so represent INT64 offsets as pairs of uint32 words and pack + # all five graph offsets into one otherwise-unused inner operand. + _q_segment_ids = _pack_ragged_offsets( + q_seq_offsets, + k_seq_offsets, + config.qkv_layout, + attn_heads, + num_gqa_groups, + q_head_dim, + v_head_dim, + ) output, softmax_aux, max_tensor, rng_state, _ = FusedAttnFwdPrimitive.inner_primitive.bind( q, @@ -830,8 +672,8 @@ def convert_to_2d(offsets, batch, max_seqlen): bias, softmax_offset, seed, - q_cu_seqlen, - kv_cu_seqlen, + q_seqlen.flatten(), + kv_seqlen.flatten(), q_seq_offsets, k_seq_offsets, _q_segment_ids, @@ -1099,7 +941,7 @@ def abstract( """ Fused attention bwd abstract """ - del softmax_aux_aval, rng_state_aval, output_aval + del rng_state_aval q_dtype = dtypes.canonicalize_dtype(q_aval.dtype) k_dtype = dtypes.canonicalize_dtype(k_aval.dtype) @@ -1117,7 +959,7 @@ def abstract( ) ( - batch_shape, + _batch_shape, q_max_seqlen, kv_max_seqlen, attn_heads, @@ -1126,47 +968,43 @@ def abstract( v_head_dim, ) = FusedAttnHelper.parse_qkv_aval(q_aval, k_aval, v_aval, config.qkv_layout) - if config.attn_bias_type == AttnBiasType.NO_BIAS: - bias_batch = bias_heads = 0 - else: - *bias_batch_shape, bias_heads, _, _ = bias_aval.shape - bias_batch = reduce(operator.mul, bias_batch_shape) - - deterministic = not FusedAttnHelper.is_non_deterministic_allowed() - - input_batch = reduce(operator.mul, batch_shape) - wkspace_shape, wkspace_dtype = transformer_engine_jax.get_fused_attn_bwd_workspace_sizes( - input_batch, - bias_batch, - q_max_seqlen, - kv_max_seqlen, + backend, message = FusedAttnHelper( + config.is_training, + q_dtype, + k_dtype, + config.qkv_layout, + config.attn_bias_type, + config.attn_mask_type, + config.softmax_type, + config.dropout_probability, attn_heads, num_gqa_groups, - bias_heads, + q_max_seqlen, + kv_max_seqlen, qk_head_dim, v_head_dim, - config.scaling_factor, - config.dropout_probability, - config.attn_bias_type.value, - config.attn_mask_type.value, - config.softmax_type.value, - config.qkv_layout.value, - jax_dtype_to_te_dtype(q_aval.dtype), - config.is_training, - deterministic, - config.max_segments_per_seq, - config.effective_window_size[0], - config.effective_window_size[1], - config.bottom_right_diagonal, + config.window_size, + False, + ).get_fused_attn_backend() + if backend != NVTE_Fused_Attn_Backend.NVTE_F16_arbitrary_seqlen: + raise ValueError(f"Unsupported {backend=}: {message}") + + graph_info = build_bwd_graph( + q_aval, + k_aval, + v_aval, + bias_aval, + softmax_aux_aval, + output_aval, + doutput_aval, + config, ) dq_aval = q_aval.update(shape=q_aval.shape, dtype=q_dtype) dk_aval = k_aval.update(shape=k_aval.shape, dtype=k_dtype) dv_aval = v_aval.update(shape=v_aval.shape, dtype=v_dtype) dbias_aval = bias_aval.update(shape=bias_aval.shape, dtype=bias_dtype) - wkspace_aval = q_aval.update( - shape=wkspace_shape, dtype=te_dtype_to_jax_dtype(wkspace_dtype) - ) + wkspace_aval = q_aval.update(shape=(graph_info.graph.workspace_size,), dtype=jnp.uint8) # Validate incoming softmax_offset shape and dtype assert ( @@ -1230,25 +1068,16 @@ def lowering( """ q_aval, k_aval, v_aval, bias_aval, *_ = ctx.avals_in - ( - batch_shape, - q_max_seqlen, - kv_max_seqlen, - attn_heads, - num_gqa_groups, - qk_head_dim, - v_head_dim, - ) = FusedAttnHelper.parse_qkv_aval(q_aval, k_aval, v_aval, config.qkv_layout) - - input_batch = reduce(operator.mul, batch_shape) - - if config.attn_bias_type == AttnBiasType.NO_BIAS: - bias_batch = bias_heads = 0 - else: - *bias_batch_shape, bias_heads, _, _ = bias_aval.shape - bias_batch = reduce(operator.mul, bias_batch_shape) - - window_size_left, window_size_right = config.effective_window_size + graph = build_bwd_graph( + q_aval, + k_aval, + v_aval, + bias_aval, + ctx.avals_in[5], + ctx.avals_in[7], + ctx.avals_in[8], + config, + ).graph return ffi.ffi_lowering(FusedAttnBwdPrimitive.name)( ctx, @@ -1269,27 +1098,8 @@ def lowering( kv_segment_ids, q_segment_pos, kv_segment_pos, # ffi_lowering needs number of parameters meets primitive.lowering - input_batch=input_batch, - bias_batch=bias_batch, - q_max_seqlen=q_max_seqlen, - kv_max_seqlen=kv_max_seqlen, - attn_heads=attn_heads, - num_gqa_groups=num_gqa_groups, - bias_heads=bias_heads, - qk_head_dim=qk_head_dim, - v_head_dim=v_head_dim, - max_segments_per_seq=config.max_segments_per_seq, - scaling_factor=float(config.scaling_factor), - dropout_probability=float(config.dropout_probability), - bias_type=int(config.attn_bias_type.value), - mask_type=int(config.attn_mask_type.value), - qkv_layout=int(config.qkv_layout.value), - is_training=config.is_training, - deterministic=not FusedAttnHelper.is_non_deterministic_allowed(), - window_size_left=window_size_left, - window_size_right=window_size_right, - bottom_right_diagonal=config.bottom_right_diagonal, - softmax_type=int(config.softmax_type.value), + is_ragged=config.qkv_layout.is_thd(), + **graph.ffi_attrs(), ) @staticmethod @@ -1352,9 +1162,15 @@ def convert_to_2d(offsets, batch, max_seqlen): ) return offsets_2d - batch, q_max_seqlen, kv_max_seqlen, *_ = FusedAttnHelper.parse_qkv_aval( - q, k, v, config.qkv_layout - ) + ( + batch, + q_max_seqlen, + kv_max_seqlen, + attn_heads, + num_gqa_groups, + q_head_dim, + v_head_dim, + ) = FusedAttnHelper.parse_qkv_aval(q, k, v, config.qkv_layout) assert ( len(batch) == 1 ), f"Expected len(batch) == 1, but got len(batch)={len(batch)}, batch={batch}" @@ -1388,8 +1204,24 @@ def convert_to_2d(offsets, batch, max_seqlen): k_seq_offsets, k_seq_offsets >= 0, fill_value=kv_batch * kv_max_seqlen ) - q_cu_seqlen = generate_cu_seqlen(q_seqlen.flatten()) - kv_cu_seqlen = generate_cu_seqlen(kv_seqlen.flatten()) + graph_batch = ragged_graph_batch_size(q_batch, config.max_segments_per_seq) + q_seqlen = _pad_ragged_metadata(q_seqlen, graph_batch, fill_value) + kv_seqlen = _pad_ragged_metadata(kv_seqlen, graph_batch, fill_value) + q_seq_offsets = _pad_ragged_metadata( + q_seq_offsets, graph_batch + 1, q_batch * q_max_seqlen + ) + k_seq_offsets = _pad_ragged_metadata( + k_seq_offsets, graph_batch + 1, kv_batch * kv_max_seqlen + ) + _q_segment_ids = _pack_ragged_offsets( + q_seq_offsets, + k_seq_offsets, + config.qkv_layout, + attn_heads, + num_gqa_groups, + q_head_dim, + v_head_dim, + ) dq, dk, dv, dbias, dsoftmax_offset, _ = FusedAttnBwdPrimitive.inner_primitive.bind( q, @@ -1401,8 +1233,8 @@ def convert_to_2d(offsets, batch, max_seqlen): rng_state, output, doutput, - q_cu_seqlen, - kv_cu_seqlen, + q_seqlen.flatten(), + kv_seqlen.flatten(), q_seq_offsets, k_seq_offsets, _q_segment_ids, @@ -3743,6 +3575,7 @@ def fused_attn_fwd( context_parallel_axis: str = "", stripe_size: int | None = None, return_max_logit: bool = False, + bottom_right_diagonal: bool | None = None, ) -> jnp.ndarray: """ Perform the forward pass of with cuDNN fused attention implementations. @@ -3783,6 +3616,8 @@ def fused_attn_fwd( context_parallel_axis (str): The name of the context parallel axis. stripe_size (int | None): Indicates the striping height to be used for ReorderStrategy.Striped Load Balancing return_max_logit (bool): Whether to return the per-head maximum attention logit. + bottom_right_diagonal (bool | None): Explicit diagonal alignment. When unset, it + follows whether ``attn_mask_type`` is a bottom-right mask. Returns: (jnp.ndarray): The output tensor from the fused attention. """ @@ -3806,10 +3641,10 @@ def fused_attn_fwd( else: raise ValueError(f"Unknown {qkv_layout=}") - if attn_bias_type == AttnBiasType.NO_BIAS: + if attn_bias_type in (AttnBiasType.NO_BIAS, AttnBiasType.ALIBI): assert ( bias is None - ), f"bias must be None when attn_bias_type is NO_BIAS, but got bias={bias}" + ), f"bias must be None when attn_bias_type is {attn_bias_type}, but got bias={bias}" bias = jnp.zeros(0, dtype=qkv[0].dtype) if softmax_offset is None: @@ -3853,7 +3688,11 @@ def fused_attn_fwd( is_training=is_training, max_segments_per_seq=max_segments_per_seq, window_size=(-1, -1) if window_size is None else window_size, - bottom_right_diagonal=attn_mask_type.is_bottom_right(), + bottom_right_diagonal=( + attn_mask_type.is_bottom_right() + if bottom_right_diagonal is None + else bottom_right_diagonal + ), context_parallel_load_balanced=context_parallel_causal_load_balanced, cp_axis=_maybe_context_parallel_axis(context_parallel_axis), cp_striped_window_size=None, @@ -3910,6 +3749,7 @@ def fused_attn_bwd( context_parallel_causal_load_balanced: bool = False, context_parallel_axis: str = "", stripe_size: int | None = None, + bottom_right_diagonal: bool | None = None, ): """ Perform the backward pass of the cuDNN fused attention implementations. @@ -3950,6 +3790,8 @@ def fused_attn_bwd( Indicates the sequences are ordered for causal mask load balancing when running context parallelism. context_parallel_axis (str): The name of the context parallel axis. stripe_size (int | None): Indicates the striping height to be used for ReorderStrategy.Striped Load Balancing + bottom_right_diagonal (bool | None): Explicit diagonal alignment. When unset, it + follows whether ``attn_mask_type`` is a bottom-right mask. Returns: Tuple[jnp.ndarray, ...], jnp.ndarray: - The first tuple contains the gradients with respect to the input `qkv` tensors in the @@ -3975,10 +3817,11 @@ def fused_attn_bwd( else: raise ValueError(f"Unknown {qkv_layout=}") - if attn_bias_type == AttnBiasType.NO_BIAS: - assert ( - bias is None - ), f"bias must be None when attn_bias_type is NO_BIAS, but got bias with type={type(bias)}" + if attn_bias_type in (AttnBiasType.NO_BIAS, AttnBiasType.ALIBI): + assert bias is None, ( + f"bias must be None when attn_bias_type is {attn_bias_type}, but got bias with" + f" type={type(bias)}" + ) bias = jnp.zeros(0, dtype=qkv[0].dtype) if softmax_offset is None: @@ -4029,7 +3872,11 @@ def fused_attn_bwd( is_training=is_training, max_segments_per_seq=max_segments_per_seq, window_size=(-1, -1) if window_size is None else window_size, - bottom_right_diagonal=attn_mask_type.is_bottom_right(), + bottom_right_diagonal=( + attn_mask_type.is_bottom_right() + if bottom_right_diagonal is None + else bottom_right_diagonal + ), context_parallel_load_balanced=context_parallel_causal_load_balanced, cp_axis=_maybe_context_parallel_axis(context_parallel_axis), cp_striped_window_size=None, diff --git a/transformer_engine/jax/cpp_extensions/cudnn_attention.py b/transformer_engine/jax/cpp_extensions/cudnn_attention.py new file mode 100644 index 00000000000..47be240dd8d --- /dev/null +++ b/transformer_engine/jax/cpp_extensions/cudnn_attention.py @@ -0,0 +1,959 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. +"""Python cuDNN frontend graphs for the standard JAX fused-attention path.""" + +from __future__ import annotations + +import operator +import os +from dataclasses import dataclass +from functools import reduce +from typing import Any + +import jax.numpy as jnp +import numpy as np + +from transformer_engine.common.attention.cudnn import ( + AttentionLayout, + FusedAttentionConfig, + FusedAttentionSupport, + check_f16_fused_attention_support, + cudnn_mask_options, + ragged_batch_bucket, + ragged_token_bucket, +) + +from .cudnn_graph import ( + GraphBinding, + SerializedGraph, + cudnn_data_type, + dtype_name, + finalize_graph, + import_cudnn, + make_graph, + record_cache_event, + record_cache_lookup, + serialized_graph, +) +from .misc import get_all_device_compute_capability, get_cudnn_version + +# Stable graph tensor UIDs. They are deliberately shared by forward and backward +# so serialized graphs and variant packs remain easy to inspect. +_UID_Q = 1 +_UID_K = 2 +_UID_V = 3 +_UID_O = 4 +_UID_STATS = 5 +_UID_MAX = 6 +_UID_BIAS = 7 +_UID_SINK = 8 +_UID_SEQ_Q = 9 +_UID_SEQ_KV = 10 +_UID_OFFSET_Q = 11 +_UID_OFFSET_K = 12 +_UID_OFFSET_V = 13 +_UID_OFFSET_O = 14 +_UID_OFFSET_STATS = 15 +_UID_DROPOUT_SEED = 16 +_UID_DROPOUT_OFFSET = 17 +_UID_DO = 18 +_UID_DQ = 19 +_UID_DK = 20 +_UID_DV = 21 +_UID_DBIAS = 22 +_UID_DSINK = 23 +_UID_ATTN_SCALE = 24 + + +@dataclass(frozen=True) +class AttentionGraphInfo: + """Serialized graph plus abstract-output information used by JAX lowering.""" + + graph: SerializedGraph + stats_shape: tuple[int, ...] + max_shape: tuple[int, ...] + + +@dataclass(frozen=True) +class _LayoutInfo: + batch_shape: tuple[int, ...] + input_batch: int + q_max_seqlen: int + kv_max_seqlen: int + q_heads: int + kv_heads: int + qk_dim: int + v_dim: int + + +_graph_cache: dict[tuple[Any, ...], AttentionGraphInfo] = {} + + +def _layout_info(q_aval, k_aval, v_aval, layout) -> _LayoutInfo: + """Parse TE's three supported JAX QKV layout groups.""" + if layout.is_qkvpacked(): + *batch_shape, q_seqlen, packed, q_heads, qk_dim = q_aval.shape + if packed != 3: + raise ValueError(f"QKV-packed fused attention expects dimension 3, got {q_aval.shape}.") + kv_seqlen = q_seqlen + kv_heads = q_heads + v_dim = qk_dim + elif layout.is_kvpacked(): + *batch_shape, q_seqlen, q_heads, qk_dim = q_aval.shape + *kv_batch_shape, kv_seqlen, packed, kv_heads, v_dim = k_aval.shape + if tuple(batch_shape) != tuple(kv_batch_shape) or packed != 2: + raise ValueError(f"Invalid KV-packed fused-attention shapes: {q_aval}, {k_aval}.") + if qk_dim != v_dim: + raise ValueError("KV-packed fused attention requires equal QK and V head dimensions.") + elif layout.is_separate(): + *batch_shape, q_seqlen, q_heads, qk_dim = q_aval.shape + *k_batch_shape, kv_seqlen, kv_heads, k_dim = k_aval.shape + *v_batch_shape, v_seqlen, v_heads, v_dim = v_aval.shape + if tuple(batch_shape) != tuple(k_batch_shape) or tuple(batch_shape) != tuple(v_batch_shape): + raise ValueError("Separate Q, K and V tensors must have matching batch shapes.") + if qk_dim != k_dim or kv_seqlen != v_seqlen or kv_heads != v_heads: + raise ValueError("Separate fused-attention K and V shapes are inconsistent.") + else: + raise ValueError(f"Unsupported JAX fused-attention layout: {layout}.") + return _LayoutInfo( + batch_shape=tuple(int(dim) for dim in batch_shape), + input_batch=reduce(operator.mul, batch_shape, 1), + q_max_seqlen=int(q_seqlen), + kv_max_seqlen=int(kv_seqlen), + q_heads=int(q_heads), + kv_heads=int(kv_heads), + qk_dim=int(qk_dim), + v_dim=int(v_dim), + ) + + +def _matrix_stride(info: _LayoutInfo, layout, matrix: str, graph_sq: int, graph_skv: int): + """Port generateMatrixStrides for JAX's BSHD/THD layout subset.""" + if matrix in ("q", "o"): + heads = info.q_heads + dim = info.qk_dim if matrix == "q" else info.v_dim + seqlen = graph_sq + else: + heads = info.kv_heads + dim = info.qk_dim if matrix == "k" else info.v_dim + seqlen = graph_skv + + if matrix in ("q", "k", "v") and layout.is_qkvpacked(): + return ( + graph_sq * 3 * info.q_heads * info.qk_dim, + dim, + 3 * info.q_heads * info.qk_dim, + 1, + ) + if matrix in ("k", "v") and layout.is_kvpacked(): + return ( + graph_skv * 2 * info.kv_heads * info.qk_dim, + dim, + 2 * info.kv_heads * info.qk_dim, + 1, + ) + return (seqlen * heads * dim, dim, heads * dim, 1) + + +def _qkv_bindings(info: _LayoutInfo, layout, itemsize: int, *, outputs: bool = False): + """Return UID bindings for separate or physically packed QKV buffers.""" + if outputs: + uids = (_UID_DQ, _UID_DK, _UID_DV) + else: + uids = (_UID_Q, _UID_K, _UID_V) + if layout.is_qkvpacked(): + stride = info.q_heads * info.qk_dim * itemsize + return ( + GraphBinding(uids[0], 0, 0), + GraphBinding(uids[1], 0, stride), + GraphBinding(uids[2], 0, 2 * stride), + ) + if layout.is_kvpacked(): + stride = info.kv_heads * info.qk_dim * itemsize + return ( + GraphBinding(uids[0], 0, 0), + GraphBinding(uids[1], 1, 0), + GraphBinding(uids[2], 1, stride), + ) + return tuple(GraphBinding(uid, index, 0) for index, uid in enumerate(uids)) + + +def _is_bias(config) -> bool: + return getattr(config.attn_bias_type, "name", "") == "POST_SCALE_BIAS" + + +def _is_padding(config) -> bool: + return bool(config.attn_mask_type.is_padding()) + + +def _is_causal(config) -> bool: + name = getattr(config.attn_mask_type, "name", "") + return name in ("CAUSAL_MASK", "PADDING_CAUSAL_MASK") + + +def _is_bottom_right(config) -> bool: + return bool(config.attn_mask_type.is_bottom_right()) + + +def _has_sink(config) -> bool: + return getattr(config.softmax_type, "name", "") != "VANILLA_SOFTMAX" + + +def _is_dropout(config) -> bool: + return bool(config.is_training and config.dropout_probability != 0.0) + + +def _device_arch() -> int: + capabilities = get_all_device_compute_capability() + return int(capabilities[0]) if capabilities else 0 + + +def ragged_graph_batch_size(input_batch: int, max_segments_per_seq: int) -> int: + """Preserve the legacy cuDNN graph batch-size bucket for ragged attention.""" + batch = int(input_batch) * int(max_segments_per_seq) + # Bucketing is part of cuDNN's ragged-stats layout, introduced in 9.6. + # Older versions use dense stats and require the physical metadata extent. + if get_cudnn_version() < (9, 6, 0) or _device_arch() == 120: + return batch + return ragged_batch_bucket(batch) + + +def _ragged_graph_token_count(tokens: int) -> int: + """Preserve the legacy cuDNN graph token-count bucket for ragged attention.""" + return ragged_token_bucket(tokens) + + +def _graph_dimensions(info: _LayoutInfo, config): + """Return logical cuDNN B/H/S dimensions and physical auxiliary shapes.""" + is_ragged = config.qkv_layout.is_thd() + cudnn_version = get_cudnn_version() + arch = _device_arch() + use_ragged_stats = is_ragged and cudnn_version >= (9, 6, 0) and arch != 120 + if is_ragged: + graph_batch = ragged_graph_batch_size(info.input_batch, config.max_segments_per_seq) + if cudnn_version < (9, 6, 0) or arch == 120: + graph_sq = info.q_max_seqlen + graph_skv = info.kv_max_seqlen + else: + graph_sq = _ragged_graph_token_count(info.input_batch * info.q_max_seqlen) + graph_skv = _ragged_graph_token_count(info.input_batch * info.kv_max_seqlen) + else: + graph_batch = info.input_batch + graph_sq = info.q_max_seqlen + graph_skv = info.kv_max_seqlen + + if is_ragged and cudnn_version >= (9, 6, 0): + stats_shape = (*info.batch_shape, info.q_max_seqlen, info.q_heads, 1) + elif cudnn_version >= (9, 6, 0): + stats_shape = (*info.batch_shape, info.q_heads, info.q_max_seqlen, 1) + else: + stats_shape = ( + *info.batch_shape, + info.q_heads, + info.q_max_seqlen, + int(config.max_segments_per_seq), + ) + if config.return_max_logit: + max_shape = ( + (*info.batch_shape, info.q_max_seqlen, info.q_heads, 1) + if use_ragged_stats + else (*info.batch_shape, info.q_heads, info.q_max_seqlen, 1) + ) + else: + max_shape = (0,) + return graph_batch, graph_sq, graph_skv, use_ragged_stats, stats_shape, max_shape + + +def _tensor(graph, _cudnn, *, name, dim, stride, dtype, uid): + return graph.tensor( + name=name, + dim=tuple(int(x) for x in dim), + stride=tuple(int(x) for x in stride), + data_type=dtype, + uid=uid, + ) + + +def _ragged_offset(graph, cudnn, name: str, uid: int, graph_batch: int, dtype): + return _tensor( + graph, + cudnn, + name=name, + dim=(graph_batch + 1, 1, 1, 1), + stride=(1, 1, 1, 1), + dtype=dtype, + uid=uid, + ) + + +def _ragged_offset_spec(cudnn): + """Return the cuDNN datatype and byte size for external element offsets.""" + use_int64 = get_cudnn_version() >= (9, 5, 0) + dtype = cudnn.data_type.INT64 if use_int64 else cudnn.data_type.INT32 + itemsize = np.dtype(np.int64 if use_int64 else np.int32).itemsize + return dtype, itemsize + + +def _mask_options(cudnn, info: _LayoutInfo, config): + cp_striped_window_size = getattr(config, "cp_striped_window_size", None) + window_left, window_right = ( + cp_striped_window_size if cp_striped_window_size is not None else config.window_size + ) + options = cudnn_mask_options( + causal=_is_causal(config), + bottom_right=_is_bottom_right(config), + padding=_is_padding(config), + bottom_right_diagonal=bool(config.bottom_right_diagonal), + window_size=(window_left, window_right), + max_seqlen_q=info.q_max_seqlen, + max_seqlen_kv=info.kv_max_seqlen, + cudnn_version=get_cudnn_version(), + ) + options.pop("is_padding") + options["diagonal_alignment"] = ( + cudnn.diagonal_alignment.BOTTOM_RIGHT + if options["diagonal_alignment"] == "bottom_right" + else cudnn.diagonal_alignment.TOP_LEFT + ) + return options + + +def _scalar_tensor(graph, _cudnn, name: str, uid: int, dtype): + return graph.tensor( + name=name, + dim=(1, 1, 1, 1), + stride=(1, 1, 1, 1), + data_type=dtype, + is_pass_by_value=True, + uid=uid, + ) + + +def _cache_key(direction: str, q_aval, k_aval, v_aval, bias_aval, config, *extra_avals): + avals = (q_aval, k_aval, v_aval, bias_aval, *extra_avals) + return ( + direction, + config, + tuple((tuple(aval.shape), dtype_name(aval.dtype)) for aval in avals), + get_cudnn_version(), + _device_arch(), + ) + + +def build_fwd_graph(q_aval, k_aval, v_aval, bias_aval, config) -> AttentionGraphInfo: + """Build or retrieve the standard fused-attention forward graph.""" + key = _cache_key("fwd", q_aval, k_aval, v_aval, bias_aval, config) + graph_info = _graph_cache.get(key) + record_cache_lookup(("f16", "fwd"), hit=graph_info is not None, key=key) + if graph_info is None: + _graph_cache[key] = _build_fwd_graph(q_aval, k_aval, v_aval, bias_aval, config) + record_cache_event(("f16", "fwd"), "cache_graph") + graph_info = _graph_cache[key] + return graph_info + + +def _build_fwd_graph(q_aval, k_aval, v_aval, bias_aval, config) -> AttentionGraphInfo: + cudnn = import_cudnn() + info = _layout_info(q_aval, k_aval, v_aval, config.qkv_layout) + graph_batch, graph_sq, graph_skv, ragged_stats, stats_shape, max_shape = _graph_dimensions( + info, config + ) + io_dtype = cudnn_data_type(cudnn, q_aval.dtype) + graph = make_graph(cudnn, io_dtype) + + q = _tensor( + graph, + cudnn, + name="q", + dim=(graph_batch, info.q_heads, graph_sq, info.qk_dim), + stride=_matrix_stride(info, config.qkv_layout, "q", graph_sq, graph_skv), + dtype=io_dtype, + uid=_UID_Q, + ) + k = _tensor( + graph, + cudnn, + name="k", + dim=(graph_batch, info.kv_heads, graph_skv, info.qk_dim), + stride=_matrix_stride(info, config.qkv_layout, "k", graph_sq, graph_skv), + dtype=io_dtype, + uid=_UID_K, + ) + v = _tensor( + graph, + cudnn, + name="v", + dim=(graph_batch, info.kv_heads, graph_skv, info.v_dim), + stride=_matrix_stride(info, config.qkv_layout, "v", graph_sq, graph_skv), + dtype=io_dtype, + uid=_UID_V, + ) + + input_bindings = list(_qkv_bindings(info, config.qkv_layout, jnp.dtype(q_aval.dtype).itemsize)) + output_bindings = [GraphBinding(_UID_O, 0), GraphBinding(_UID_STATS, 1)] + scale = _scalar_tensor(graph, cudnn, "attn_scale", _UID_ATTN_SCALE, cudnn.data_type.FLOAT) + scalar_uids = [_UID_ATTN_SCALE] + scalar_values = [np.asarray(config.scaling_factor, dtype=np.float32).tobytes()] + + kwargs = { + "name": "te_fused_attention", + "q": q, + "k": k, + "v": v, + "generate_stats": True, + "attn_scale": scale, + **_mask_options(cudnn, info, config), + } + if getattr(config.attn_bias_type, "name", "") == "ALIBI": + kwargs["use_alibi_mask"] = True + + if _is_bias(config): + *bias_batch_shape, bias_heads, bias_sq, bias_skv = bias_aval.shape + bias_batch = reduce(operator.mul, bias_batch_shape, 1) + bias = _tensor( + graph, + cudnn, + name="bias", + dim=(bias_batch, bias_heads, bias_sq, bias_skv), + stride=(bias_heads * bias_sq * bias_skv, bias_sq * bias_skv, bias_skv, 1), + dtype=io_dtype, + uid=_UID_BIAS, + ) + kwargs["bias"] = bias + input_bindings.append(GraphBinding(_UID_BIAS, 3)) + + if _has_sink(config): + sink = _tensor( + graph, + cudnn, + name="softmax_offset", + dim=(1, info.q_heads, 1, 1), + stride=(info.q_heads, 1, 1, 1), + dtype=cudnn.data_type.FLOAT, + uid=_UID_SINK, + ) + kwargs["sink_token"] = sink + input_bindings.append(GraphBinding(_UID_SINK, 4)) + + if _is_padding(config): + seq_q = _tensor( + graph, + cudnn, + name="seq_len_q", + dim=(graph_batch, 1, 1, 1), + stride=(1, 1, 1, 1), + dtype=cudnn.data_type.INT32, + uid=_UID_SEQ_Q, + ) + seq_kv = _tensor( + graph, + cudnn, + name="seq_len_kv", + dim=(graph_batch, 1, 1, 1), + stride=(1, 1, 1, 1), + dtype=cudnn.data_type.INT32, + uid=_UID_SEQ_KV, + ) + kwargs.update(use_padding_mask=True, seq_len_q=seq_q, seq_len_kv=seq_kv) + input_bindings.extend((GraphBinding(_UID_SEQ_Q, 6), GraphBinding(_UID_SEQ_KV, 7))) + + offset_q = offset_k = offset_v = offset_o = offset_stats = None + if config.qkv_layout.is_thd(): + offset_dtype, offset_itemsize = _ragged_offset_spec(cudnn) + offset_bytes = (graph_batch + 1) * offset_itemsize + offset_q = _ragged_offset( + graph, cudnn, "offset_q", _UID_OFFSET_Q, graph_batch, offset_dtype + ) + offset_k = _ragged_offset( + graph, cudnn, "offset_k", _UID_OFFSET_K, graph_batch, offset_dtype + ) + offset_v = _ragged_offset( + graph, cudnn, "offset_v", _UID_OFFSET_V, graph_batch, offset_dtype + ) + offset_o = _ragged_offset( + graph, cudnn, "offset_o", _UID_OFFSET_O, graph_batch, offset_dtype + ) + q.set_ragged_offset(offset_q) + k.set_ragged_offset(offset_k) + v.set_ragged_offset(offset_v) + input_bindings.extend( + ( + GraphBinding(_UID_OFFSET_Q, 10, 0), + GraphBinding(_UID_OFFSET_K, 10, offset_bytes), + GraphBinding(_UID_OFFSET_V, 10, 2 * offset_bytes), + GraphBinding(_UID_OFFSET_O, 10, 3 * offset_bytes), + ) + ) + + if _is_dropout(config): + seed = _tensor( + graph, + cudnn, + name="dropout_seed", + dim=(1, 1, 1, 1), + stride=(1, 1, 1, 1), + dtype=cudnn.data_type.INT64, + uid=_UID_DROPOUT_SEED, + ) + offset = _tensor( + graph, + cudnn, + name="dropout_offset", + dim=(1, 1, 1, 1), + stride=(1, 1, 1, 1), + dtype=cudnn.data_type.INT64, + uid=_UID_DROPOUT_OFFSET, + ) + kwargs["dropout"] = (float(config.dropout_probability), seed, offset) + # rng_state is forward result #3: two int64 values stored as four uint32 values. + output_bindings.extend( + ( + GraphBinding(_UID_DROPOUT_SEED, 3, 0), + GraphBinding(_UID_DROPOUT_OFFSET, 3, 8), + ) + ) + + if config.return_max_logit: + max_tensor = _tensor( + graph, + cudnn, + name="max_logit", + dim=(graph_batch, info.q_heads, graph_sq, 1), + stride=( + (info.q_heads * graph_sq, 1, info.q_heads, 1) + if ragged_stats + else (info.q_heads * graph_sq, graph_sq, 1, 1) + ), + dtype=cudnn.data_type.FLOAT, + uid=_UID_MAX, + ) + if ragged_stats: + offset_stats = _ragged_offset( + graph, + cudnn, + "offset_stats", + _UID_OFFSET_STATS, + graph_batch, + dtype=offset_dtype, + ) + max_tensor.set_ragged_offset(offset_stats) + input_bindings.append(GraphBinding(_UID_OFFSET_STATS, 10, 4 * offset_bytes)) + max_tensor.set_output(True) + kwargs["score_max"] = max_tensor + output_bindings.append(GraphBinding(_UID_MAX, 2)) + + output, stats = graph.sdpa(**kwargs) + output.set_output(True).set_uid(_UID_O).set_dim( + (graph_batch, info.q_heads, graph_sq, info.v_dim) + ).set_stride(_matrix_stride(info, config.qkv_layout, "o", graph_sq, graph_skv)) + output.set_data_type(io_dtype) + if config.qkv_layout.is_thd(): + output.set_ragged_offset(offset_o) + + stats.set_output(True).set_uid(_UID_STATS).set_data_type(cudnn.data_type.FLOAT) + stats.set_dim((graph_batch, info.q_heads, graph_sq, 1)) + if ragged_stats: + if offset_stats is None: + offset_stats = _ragged_offset( + graph, + cudnn, + "offset_stats", + _UID_OFFSET_STATS, + graph_batch, + dtype=offset_dtype, + ) + input_bindings.append(GraphBinding(_UID_OFFSET_STATS, 10, 4 * offset_bytes)) + stats.set_stride((info.q_heads * graph_sq, 1, info.q_heads, 1)) + stats.set_ragged_offset(offset_stats) + else: + stats.set_stride((info.q_heads * graph_sq, graph_sq, 1, 1)) + + cache_site = ("f16", "fwd") + workspace, data, version = finalize_graph( + cudnn, + graph, + description="fused-attention forward", + cache_site=cache_site, + ) + result = serialized_graph( + serialized_graph_data=data, + cudnn_frontend_version=version, + workspace_size=workspace, + cache_site=cache_site, + input_bindings=input_bindings, + output_bindings=output_bindings, + scalar_uids=scalar_uids, + scalar_values=scalar_values, + ) + return AttentionGraphInfo(result, stats_shape, max_shape) + + +def build_bwd_graph( + q_aval, + k_aval, + v_aval, + bias_aval, + stats_aval, + output_aval, + doutput_aval, + config, +) -> AttentionGraphInfo: + """Build or retrieve the standard fused-attention backward graph.""" + key = _cache_key( + "bwd", + q_aval, + k_aval, + v_aval, + bias_aval, + config, + stats_aval, + output_aval, + doutput_aval, + ) + graph_info = _graph_cache.get(key) + record_cache_lookup(("f16", "bwd"), hit=graph_info is not None, key=key) + if graph_info is None: + _graph_cache[key] = _build_bwd_graph( + q_aval, + k_aval, + v_aval, + bias_aval, + stats_aval, + output_aval, + doutput_aval, + config, + ) + record_cache_event(("f16", "bwd"), "cache_graph") + graph_info = _graph_cache[key] + return graph_info + + +def _build_bwd_graph( + q_aval, + k_aval, + v_aval, + bias_aval, + stats_aval, + output_aval, + doutput_aval, + config, +) -> AttentionGraphInfo: + del stats_aval, output_aval, doutput_aval + + cudnn = import_cudnn() + info = _layout_info(q_aval, k_aval, v_aval, config.qkv_layout) + graph_batch, graph_sq, graph_skv, ragged_stats, stats_shape, max_shape = _graph_dimensions( + info, config + ) + io_dtype = cudnn_data_type(cudnn, q_aval.dtype) + graph = make_graph(cudnn, io_dtype) + + def io_tensor(name, dim, stride, uid, dtype=io_dtype): + return _tensor(graph, cudnn, name=name, dim=dim, stride=stride, dtype=dtype, uid=uid) + + q = io_tensor( + "q", + (graph_batch, info.q_heads, graph_sq, info.qk_dim), + _matrix_stride(info, config.qkv_layout, "q", graph_sq, graph_skv), + _UID_Q, + ) + k = io_tensor( + "k", + (graph_batch, info.kv_heads, graph_skv, info.qk_dim), + _matrix_stride(info, config.qkv_layout, "k", graph_sq, graph_skv), + _UID_K, + ) + v = io_tensor( + "v", + (graph_batch, info.kv_heads, graph_skv, info.v_dim), + _matrix_stride(info, config.qkv_layout, "v", graph_sq, graph_skv), + _UID_V, + ) + output = io_tensor( + "o", + (graph_batch, info.q_heads, graph_sq, info.v_dim), + _matrix_stride(info, config.qkv_layout, "o", graph_sq, graph_skv), + _UID_O, + ) + doutput = io_tensor( + "dO", + (graph_batch, info.q_heads, graph_sq, info.v_dim), + _matrix_stride(info, config.qkv_layout, "o", graph_sq, graph_skv), + _UID_DO, + ) + stats = io_tensor( + "stats", + (graph_batch, info.q_heads, graph_sq, 1), + ( + (info.q_heads * graph_sq, 1, info.q_heads, 1) + if ragged_stats + else (info.q_heads * graph_sq, graph_sq, 1, 1) + ), + _UID_STATS, + cudnn.data_type.FLOAT, + ) + + itemsize = jnp.dtype(q_aval.dtype).itemsize + input_bindings = list(_qkv_bindings(info, config.qkv_layout, itemsize)) + input_bindings.extend( + ( + GraphBinding(_UID_STATS, 5), + GraphBinding(_UID_O, 7), + GraphBinding(_UID_DO, 8), + ) + ) + output_bindings = list(_qkv_bindings(info, config.qkv_layout, itemsize, outputs=True)) + scale = _scalar_tensor(graph, cudnn, "attn_scale", _UID_ATTN_SCALE, cudnn.data_type.FLOAT) + scalar_uids = [_UID_ATTN_SCALE] + scalar_values = [np.asarray(config.scaling_factor, dtype=np.float32).tobytes()] + + kwargs = { + "name": "te_fused_attention_backward", + "q": q, + "k": k, + "v": v, + "o": output, + "dO": doutput, + "stats": stats, + "attn_scale": scale, + **_mask_options(cudnn, info, config), + } + if getattr(config.attn_bias_type, "name", "") == "ALIBI": + kwargs["use_alibi_mask"] = True + if get_cudnn_version() >= (9, 0, 0): + kwargs["use_deterministic_algorithm"] = not bool( + int(os.getenv("NVTE_ALLOW_NONDETERMINISTIC_ALGO", "1")) + ) + if ragged_stats: + kwargs["max_total_seq_len_q"] = graph_sq + if config.qkv_layout.is_thd() and get_cudnn_version() >= (9, 6, 0) and _device_arch() != 120: + kwargs["max_total_seq_len_kv"] = graph_skv + + if _is_bias(config): + *bias_batch_shape, bias_heads, bias_sq, bias_skv = bias_aval.shape + bias_batch = reduce(operator.mul, bias_batch_shape, 1) + bias = io_tensor( + "bias", + (bias_batch, bias_heads, bias_sq, bias_skv), + (bias_heads * bias_sq * bias_skv, bias_sq * bias_skv, bias_skv, 1), + _UID_BIAS, + ) + kwargs["bias"] = bias + input_bindings.append(GraphBinding(_UID_BIAS, 3)) + if not (bias_batch == 1 and bias_heads == 1 and bias_sq == 1): + dbias = io_tensor( + "dBias", + (bias_batch, bias_heads, bias_sq, bias_skv), + (bias_heads * bias_sq * bias_skv, bias_sq * bias_skv, bias_skv, 1), + _UID_DBIAS, + ) + dbias.set_output(True) + kwargs["dBias"] = dbias + output_bindings.append(GraphBinding(_UID_DBIAS, 3)) + + if _has_sink(config): + sink = io_tensor( + "softmax_offset", + (1, info.q_heads, 1, 1), + (info.q_heads, 1, 1, 1), + _UID_SINK, + cudnn.data_type.FLOAT, + ) + dsink = io_tensor( + "dsoftmax_offset", + (1, info.q_heads, 1, 1), + (info.q_heads, 1, 1, 1), + _UID_DSINK, + cudnn.data_type.FLOAT, + ) + dsink.set_output(True) + kwargs.update(sink_token=sink, dSink_token=dsink) + input_bindings.append(GraphBinding(_UID_SINK, 4)) + output_bindings.append(GraphBinding(_UID_DSINK, 4)) + + if _is_padding(config): + seq_q = io_tensor( + "seq_len_q", + (graph_batch, 1, 1, 1), + (1, 1, 1, 1), + _UID_SEQ_Q, + cudnn.data_type.INT32, + ) + seq_kv = io_tensor( + "seq_len_kv", + (graph_batch, 1, 1, 1), + (1, 1, 1, 1), + _UID_SEQ_KV, + cudnn.data_type.INT32, + ) + kwargs.update(use_padding_mask=True, seq_len_q=seq_q, seq_len_kv=seq_kv) + input_bindings.extend((GraphBinding(_UID_SEQ_Q, 9), GraphBinding(_UID_SEQ_KV, 10))) + + if config.qkv_layout.is_thd(): + offset_dtype, offset_itemsize = _ragged_offset_spec(cudnn) + offset_bytes = (graph_batch + 1) * offset_itemsize + offset_q = _ragged_offset( + graph, cudnn, "offset_q", _UID_OFFSET_Q, graph_batch, offset_dtype + ) + offset_k = _ragged_offset( + graph, cudnn, "offset_k", _UID_OFFSET_K, graph_batch, offset_dtype + ) + offset_v = _ragged_offset( + graph, cudnn, "offset_v", _UID_OFFSET_V, graph_batch, offset_dtype + ) + offset_o = _ragged_offset( + graph, cudnn, "offset_o", _UID_OFFSET_O, graph_batch, offset_dtype + ) + q.set_ragged_offset(offset_q) + k.set_ragged_offset(offset_k) + v.set_ragged_offset(offset_v) + output.set_ragged_offset(offset_o) + doutput.set_ragged_offset(offset_o) + input_bindings.extend( + ( + GraphBinding(_UID_OFFSET_Q, 13, 0), + GraphBinding(_UID_OFFSET_K, 13, offset_bytes), + GraphBinding(_UID_OFFSET_V, 13, 2 * offset_bytes), + GraphBinding(_UID_OFFSET_O, 13, 3 * offset_bytes), + ) + ) + if ragged_stats: + offset_stats = _ragged_offset( + graph=graph, + cudnn=cudnn, + name="offset_stats", + uid=_UID_OFFSET_STATS, + graph_batch=graph_batch, + dtype=offset_dtype, + ) + stats.set_ragged_offset(offset_stats) + input_bindings.append(GraphBinding(_UID_OFFSET_STATS, 13, 4 * offset_bytes)) + + if _is_dropout(config): + seed = io_tensor( + "dropout_seed", + (1, 1, 1, 1), + (1, 1, 1, 1), + _UID_DROPOUT_SEED, + cudnn.data_type.INT64, + ) + offset = io_tensor( + "dropout_offset", + (1, 1, 1, 1), + (1, 1, 1, 1), + _UID_DROPOUT_OFFSET, + cudnn.data_type.INT64, + ) + kwargs["dropout"] = (float(config.dropout_probability), seed, offset) + input_bindings.extend( + ( + GraphBinding(_UID_DROPOUT_SEED, 6, 0), + GraphBinding(_UID_DROPOUT_OFFSET, 6, 8), + ) + ) + + dq, dk, dv = graph.sdpa_backward(**kwargs) + q_stride = _matrix_stride(info, config.qkv_layout, "q", graph_sq, graph_skv) + k_stride = _matrix_stride(info, config.qkv_layout, "k", graph_sq, graph_skv) + v_stride = _matrix_stride(info, config.qkv_layout, "v", graph_sq, graph_skv) + dq.set_output(True).set_uid(_UID_DQ).set_dim( + (graph_batch, info.q_heads, graph_sq, info.qk_dim) + ).set_stride(q_stride) + dk.set_output(True).set_uid(_UID_DK).set_dim( + (graph_batch, info.kv_heads, graph_skv, info.qk_dim) + ).set_stride(k_stride) + dv.set_output(True).set_uid(_UID_DV).set_dim( + (graph_batch, info.kv_heads, graph_skv, info.v_dim) + ).set_stride(v_stride) + if config.qkv_layout.is_thd(): + dq.set_ragged_offset(offset_q) + dk.set_ragged_offset(offset_k) + dv.set_ragged_offset(offset_v) + + cache_site = ("f16", "bwd") + workspace, data, version = finalize_graph( + cudnn, + graph, + description="fused-attention backward", + cache_site=cache_site, + ) + result = serialized_graph( + serialized_graph_data=data, + cudnn_frontend_version=version, + workspace_size=workspace, + cache_site=cache_site, + input_bindings=input_bindings, + output_bindings=output_bindings, + scalar_uids=scalar_uids, + scalar_values=scalar_values, + ) + return AttentionGraphInfo(result, stats_shape, max_shape) + + +def clear_graph_cache(): + """Clear the process-local serialized graph cache (primarily for tests).""" + _graph_cache.clear() + + +def _policy_layout(layout) -> AttentionLayout: + qkv_format = layout.get_qkv_format().name.lower() + if layout.is_qkvpacked(): + layout_group = "qkv_packed" + elif layout.is_kvpacked(): + layout_group = "kv_packed" + else: + layout_group = "separate" + return AttentionLayout( + qkv_format=qkv_format, + q_format=qkv_format, + kv_format=qkv_format, + layout_group=layout_group, + is_qkvpacked=layout.is_qkvpacked(), + ) + + +def _policy_mask_name(mask) -> str: + return { + "NO_MASK": "no_mask", + "CAUSAL_MASK": "causal", + "PADDING_MASK": "padding", + "PADDING_CAUSAL_MASK": "padding_causal", + "CAUSAL_BOTTOM_RIGHT_MASK": "causal_bottom_right", + "PADDING_CAUSAL_BOTTOM_RIGHT_MASK": "padding_causal_bottom_right", + }[mask.name] + + +def get_fused_attn_support(helper) -> FusedAttentionSupport: + """Return the shared F16/BF16 cuDNN attention compatibility result.""" + + return check_f16_fused_attention_support( + FusedAttentionConfig( + is_training=bool(helper.is_training), + q_dtype=str(jnp.dtype(helper.q_dtype)), + kv_dtype=str(jnp.dtype(helper.kv_dtype)), + layout=_policy_layout(helper.qkv_layout), + bias_type=helper.attn_bias_type.name.lower(), + mask_type=_policy_mask_name(helper.attn_mask_type), + softmax_type=helper.softmax_type.name.lower().removesuffix("_softmax"), + dropout=float(helper.dropout_probability), + num_attn_heads=int(helper.q_num_heads), + num_gqa_groups=int(helper.kv_num_heads), + max_seqlen_q=int(helper.q_max_seqlen), + max_seqlen_kv=int(helper.kv_max_seqlen), + head_dim_qk=int(helper.head_dim_qk), + head_dim_v=int(helper.head_dim_v), + window_size=tuple(int(value) for value in helper.window_size), + return_max_logit=bool(helper.return_max_logit), + cuda_graph=False, + deterministic=not bool(int(os.getenv("NVTE_ALLOW_NONDETERMINISTIC_ALGO", "1"))), + cudnn_version=get_cudnn_version(), + sm_arch=_device_arch(), + ) + ) + + +def is_fused_attn_supported(helper) -> bool: + """Apply the shared F16/BF16 cuDNN attention compatibility policy.""" + + return get_fused_attn_support(helper).supported diff --git a/transformer_engine/jax/cpp_extensions/cudnn_graph.py b/transformer_engine/jax/cpp_extensions/cudnn_graph.py new file mode 100644 index 00000000000..c13d9ed7296 --- /dev/null +++ b/transformer_engine/jax/cpp_extensions/cudnn_graph.py @@ -0,0 +1,314 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. +"""Shared cuDNN frontend graph serialization support for JAX custom calls. + +cuDNN frontend graphs are constructed and planned in Python while XLA executes a +serialized graph through a small, graph-agnostic FFI runtime. A binding maps a +cuDNN tensor UID to an XLA operand/result and an optional byte offset. Offsets +are required for TE's packed QKV layouts, where several logical cuDNN tensors +share one JAX buffer. +""" + +from __future__ import annotations + +import hashlib +from collections.abc import Sequence +from dataclasses import dataclass +from typing import Any + +import jax +import jax.numpy as jnp +import numpy as np +import transformer_engine_jax + +from transformer_engine.common.attention.cache_debug import enabled as cache_debug_enabled +from transformer_engine.common.cudnn_frontend import ( + build_cudnn_graph, + import_cudnn_frontend, + make_cudnn_graph, +) + + +@dataclass(frozen=True) +class GraphBinding: + """Bind a cuDNN tensor UID to a JAX buffer and byte offset.""" + + uid: int + buffer_index: int + byte_offset: int = 0 + + +@dataclass(frozen=True) +class SerializedGraph: + """Serialized cuDNN graph and static metadata for the generic FFI executor.""" + + serialized_graph: bytes + graph_hash: tuple[int, int] + cudnn_frontend_version: int + workspace_size: int + attention_backend: str + attention_direction: str + input_uids: np.ndarray + input_buffer_indices: np.ndarray + input_byte_offsets: np.ndarray + output_uids: np.ndarray + output_buffer_indices: np.ndarray + output_byte_offsets: np.ndarray + scalar_uids: np.ndarray + scalar_sizes: np.ndarray + scalar_values: np.ndarray + + def ffi_attrs(self) -> dict[str, Any]: + """Return the static attributes consumed by the generic C++ executor.""" + return { + "serialized_graph": self.serialized_graph, + "graph_hash0": self.graph_hash[0], + "graph_hash1": self.graph_hash[1], + "cudnn_frontend_version": self.cudnn_frontend_version, + "attention_backend": self.attention_backend, + "attention_direction": self.attention_direction, + "input_uids": self.input_uids, + "input_buffer_indices": self.input_buffer_indices, + "input_byte_offsets": self.input_byte_offsets, + "output_uids": self.output_uids, + "output_buffer_indices": self.output_buffer_indices, + "output_byte_offsets": self.output_byte_offsets, + "scalar_uids": self.scalar_uids, + "scalar_sizes": self.scalar_sizes, + "scalar_values": self.scalar_values, + } + + +def row_major_stride(shape: Sequence[int]) -> tuple[int, ...]: + """Return element strides for a contiguous row-major tensor.""" + stride = [] + running = 1 + for dim in reversed(tuple(shape)): + stride.append(running) + running *= int(dim) + return tuple(reversed(stride)) + + +def bshd_as_bhsd_dim_stride( + shape: Sequence[int], +) -> tuple[tuple[int, ...], tuple[int, ...]]: + """Describe a contiguous BSHD buffer as cuDNN's logical BHSD tensor.""" + if len(shape) != 4: + raise ValueError(f"Expected a rank-4 BSHD tensor, got shape={shape}.") + batch, seqlen, heads, head_dim = (int(dim) for dim in shape) + return ( + (batch, heads, seqlen, head_dim), + (seqlen * heads * head_dim, head_dim, heads * head_dim, 1), + ) + + +def dtype_name(dtype) -> str: + """Stable dtype name for graph cache keys.""" + return str(jnp.dtype(dtype)) + + +def cudnn_data_type(cudnn, dtype): + """Convert a NumPy/JAX dtype to a cuDNN frontend data type.""" + dtype = jnp.dtype(dtype) + if dtype == jnp.float16: + return cudnn.data_type.HALF + if dtype == jnp.bfloat16: + return cudnn.data_type.BFLOAT16 + if dtype == jnp.float8_e4m3fn: + return cudnn.data_type.FP8_E4M3 + if dtype == jnp.float8_e5m2: + return cudnn.data_type.FP8_E5M2 + if hasattr(jnp, "float8_e8m0fnu") and dtype == jnp.float8_e8m0fnu: + return cudnn.data_type.FP8_E8M0 + if dtype == jnp.float32: + return cudnn.data_type.FLOAT + if dtype == jnp.float64: + return cudnn.data_type.DOUBLE + if dtype == jnp.int32: + return cudnn.data_type.INT32 + if dtype == jnp.int64: + return cudnn.data_type.INT64 + if dtype == jnp.uint8: + return cudnn.data_type.UINT8 + if dtype == jnp.bool_: + return cudnn.data_type.BOOLEAN + raise ValueError(f"Unsupported cuDNN graph tensor dtype: {dtype}.") + + +def cudnn_data_type_from_name(cudnn, dtype_name_: str): + """Convert a serialized NumPy dtype name to a cuDNN frontend dtype.""" + if dtype_name_ == "bfloat16": + return cudnn.data_type.BFLOAT16 + return cudnn_data_type(cudnn, np.dtype(dtype_name_)) + + +def graph_tensor_from_aval(cudnn, graph, name: str, aval, uid: int): + """Create a contiguous graph tensor from a JAX abstract value.""" + shape = tuple(int(dim) for dim in aval.shape) + return graph.tensor( + name=name, + dim=shape, + stride=row_major_stride(shape), + data_type=cudnn_data_type(cudnn, aval.dtype), + uid=uid, + ) + + +def encode_cudnn_frontend_version(version: str) -> int: + """Encode a PEP-440 cuDNN frontend version as MMmmpp.""" + public_version = version.split("+", 1)[0].split("-", 1)[0] + parts = public_version.split(".") + if len(parts) < 3: + raise RuntimeError(f"Could not parse cuDNN frontend Python version: {version!r}.") + major, minor, patch = (int(part) for part in parts[:3]) + return major * 10000 + minor * 100 + patch + + +def check_cudnn_frontend_version_match(cudnn) -> int: + """Ensure Python and C++ frontend versions use a compatible wire format.""" + python_version_string = getattr(cudnn, "__version__", None) + if python_version_string is None: + raise RuntimeError("cuDNN frontend Python package does not expose __version__.") + python_version = encode_cudnn_frontend_version(python_version_string) + cpp_version = int(transformer_engine_jax.get_cudnn_frontend_version()) + if python_version != cpp_version: + raise RuntimeError( + "cuDNN frontend Python/C++ version mismatch for graph serialization: " + f"Python cudnn.__version__={python_version_string!r} encodes to {python_version}, " + f"but Transformer Engine C++ was built with CUDNN_FRONTEND_VERSION={cpp_version}. " + "Use matching cuDNN frontend Python package and C++ headers." + ) + return python_version + + +def import_cudnn(): + """Import and validate the cuDNN frontend Python binding.""" + cudnn = import_cudnn_frontend( + feature="JAX fused attention", + requirement="nvidia-cudnn-frontend", + ) + check_cudnn_frontend_version_match(cudnn) + return cudnn + + +def make_graph(cudnn, io_dtype): + """Create a JAX cuDNN graph with TE's standard compute types.""" + + return make_cudnn_graph(cudnn, io_dtype) + + +def graph_hash(graph_data: bytes) -> tuple[int, int]: + """Return two signed int64 values used as the C++ graph-cache key.""" + digest = hashlib.sha256(graph_data).digest() + return ( + int.from_bytes(digest[0:8], byteorder="little", signed=True), + int.from_bytes(digest[8:16], byteorder="little", signed=True), + ) + + +def pack_scalar_values(scalar_values: Sequence[bytes]) -> tuple[np.ndarray, np.ndarray]: + """Pack pass-by-value scalars into fixed, aligned 16-byte records.""" + scalar_sizes = np.asarray([len(value) for value in scalar_values], dtype=np.int64) + packed_values = np.zeros((len(scalar_values), 16), dtype=np.uint8) + for index, value in enumerate(scalar_values): + if len(value) > 16: + raise ValueError("cuDNN pass-by-value scalars must be at most 16 bytes.") + packed_values[index, : len(value)] = np.frombuffer(value, dtype=np.uint8) + return scalar_sizes, packed_values.reshape(-1) + + +def serialized_graph( + *, + serialized_graph_data: bytes, + cudnn_frontend_version: int, + workspace_size: int, + cache_site: tuple[str, str], + input_bindings: Sequence[GraphBinding], + output_bindings: Sequence[GraphBinding], + scalar_uids: Sequence[int] = (), + scalar_values: Sequence[bytes] = (), +) -> SerializedGraph: + """Construct normalized, NumPy-backed metadata for an FFI graph call.""" + scalar_sizes, packed_scalar_values = pack_scalar_values(scalar_values) + + def binding_array(bindings, field): + return np.asarray([getattr(binding, field) for binding in bindings], dtype=np.int64) + + return SerializedGraph( + serialized_graph=serialized_graph_data, + graph_hash=graph_hash(serialized_graph_data), + cudnn_frontend_version=int(cudnn_frontend_version), + workspace_size=max(int(workspace_size), 1), + attention_backend=cache_site[0], + attention_direction=cache_site[1], + input_uids=binding_array(input_bindings, "uid"), + input_buffer_indices=binding_array(input_bindings, "buffer_index"), + input_byte_offsets=binding_array(input_bindings, "byte_offset"), + output_uids=binding_array(output_bindings, "uid"), + output_buffer_indices=binding_array(output_bindings, "buffer_index"), + output_byte_offsets=binding_array(output_bindings, "byte_offset"), + scalar_uids=np.asarray(scalar_uids, dtype=np.int64), + scalar_sizes=scalar_sizes, + scalar_values=packed_scalar_values, + ) + + +def record_cache_event( + cache_site: tuple[str, str], + event: str, + *, + key=None, + elapsed_ns: int = 0, +) -> None: + """Record an event in the JAX native cache-diagnostic state.""" + + if not cache_debug_enabled(): + return + native_event = "plans_built" if event == "BUILD_PLANS" else event.lower() + transformer_engine_jax.record_fused_attn_cache_event( + *cache_site, + native_event, + -1, + "" if key is None else repr(key), + elapsed_ns, + ) + + +def record_cache_lookup(cache_site: tuple[str, str], *, hit: bool, key=None) -> None: + """Record a JAX Python graph-cache lookup.""" + + record_cache_event(cache_site, "hit" if hit else "miss", key=key) + + +def _build_recorder(cache_site: tuple[str, str]): + def record(event: str, elapsed_ns: int) -> None: + record_cache_event(cache_site, event, elapsed_ns=elapsed_ns) + + return record + + +def finalize_graph( + cudnn, + graph, + *, + description: str, + cache_site: tuple[str, str], +) -> tuple[int, bytes, int]: + """Validate, plan and serialize a cuDNN frontend graph.""" + workspace_size = build_cudnn_graph( + cudnn, + graph, + description=description, + debug_callback=_build_recorder(cache_site), + ) + return ( + workspace_size, + bytes(graph.serialize()), + check_cudnn_frontend_version_match(cudnn), + ) + + +def shape_dtype(value) -> jax.ShapeDtypeStruct: + """Return a hashable-enough static shape/dtype descriptor for graph construction.""" + return jax.ShapeDtypeStruct(tuple(value.shape), value.dtype) diff --git a/transformer_engine/jax/cpp_extensions/flex_attention.py b/transformer_engine/jax/cpp_extensions/flex_attention.py index ff9f5fb1460..b1700c27514 100644 --- a/transformer_engine/jax/cpp_extensions/flex_attention.py +++ b/transformer_engine/jax/cpp_extensions/flex_attention.py @@ -3,9 +3,6 @@ # See LICENSE for license information. """cuDNN frontend score_mod fused attention helpers.""" -import hashlib -import importlib -import inspect import os from dataclasses import dataclass from typing import Any, Callable, Dict, Mapping, Optional, Sequence, Tuple @@ -15,7 +12,45 @@ import numpy as np from jax import ffi -import transformer_engine_jax +from transformer_engine.common.attention.score_mod import ( + UncacheableScoreModKey, + is_uncacheable_score_mod_key, + score_mod_callback_cache_key, +) + +from .cudnn_graph import ( + GraphBinding, + SerializedGraph, + finalize_graph, + import_cudnn, + make_graph, + record_cache_event, + record_cache_lookup, +) +from .cudnn_graph import ( + bshd_as_bhsd_dim_stride as _bshd_as_bhsd_dim_stride, +) +from .cudnn_graph import ( + cudnn_data_type as _cudnn_data_type, +) +from .cudnn_graph import ( + cudnn_data_type_from_name as _cudnn_data_type_from_name, +) +from .cudnn_graph import ( + dtype_name as _dtype_name, +) +from .cudnn_graph import ( + graph_tensor_from_aval as _graph_tensor_from_aval, +) +from .cudnn_graph import ( + row_major_stride as _row_major_stride, +) +from .cudnn_graph import ( + serialized_graph as make_serialized_graph, +) +from .cudnn_graph import ( + shape_dtype as _shape_dtype, +) __all__ = [ "FusedAttnScoreModHelper", @@ -110,103 +145,6 @@ class _ScoreModScalarSpec: stride: Tuple[int, ...] = (1, 1, 1, 1) -class _UncacheableScoreModKey: - """Unique static key for callbacks that must not share compiled score_mod graphs.""" - - def __hash__(self): - return id(self) - - def __eq__(self, other): - return self is other - - -def _score_mod_key_is_uncacheable(key: Any) -> bool: - return isinstance(key, _UncacheableScoreModKey) - - -def _freeze_score_mod_cache_key(value: Any) -> Any: - """Convert a user-provided score_mod graph key into a hashable structure.""" - if _is_array_operand(value): - raise TypeError( - "score_mod_graph_cache_key() must not include tensors. Pass runtime tensors " - "through score_mod_tensors or score_mod_bprop_tensors instead." - ) - if isinstance(value, Mapping): - items = ( - ( - _freeze_score_mod_cache_key(key), - _freeze_score_mod_cache_key(val), - ) - for key, val in value.items() - ) - return tuple(sorted(items, key=repr)) - if isinstance(value, (list, tuple)): - return tuple(_freeze_score_mod_cache_key(item) for item in value) - if isinstance(value, (set, frozenset)): - items = (_freeze_score_mod_cache_key(item) for item in value) - return tuple(sorted(items, key=repr)) - try: - hash(value) - except TypeError as exc: - raise TypeError( - "score_mod_graph_cache_key() must return a hashable value or a nested " - "combination of mapping/list/tuple/set values." - ) from exc - return value - - -def _score_mod_explicit_cache_key(callback_owner: Any) -> Optional[Any]: - """Return a user-provided structural graph key for a score_mod callback.""" - explicit_key = getattr(callback_owner, "score_mod_graph_cache_key", None) - if explicit_key is None: - return None - explicit_key = explicit_key() if callable(explicit_key) else explicit_key - return _freeze_score_mod_cache_key(explicit_key) - - -def _score_mod_callback_cache_key(callback: Optional[Callable]) -> Any: - """Create a stable graph cache key for a score_mod callable. - - Module-level functions are assumed to have stable topology. Stateful bound methods and - callable instances need an explicit score_mod_graph_cache_key(); otherwise their graphs - are left uncached to avoid reusing stale graphs after Python object address reuse. - """ - if callback is None: - return None - self_obj = getattr(callback, "__self__", None) - func_obj = getattr(callback, "__func__", None) - if self_obj is not None and func_obj is not None: - explicit_key = _score_mod_explicit_cache_key(self_obj) - if explicit_key is None: - return _UncacheableScoreModKey() - return ( - "bound_method", - type(self_obj), - func_obj.__module__, - func_obj.__qualname__, - explicit_key, - ) - - explicit_key = _score_mod_explicit_cache_key(callback) - if explicit_key is not None: - return ( - "callable", - type(callback), - getattr(callback, "__module__", None), - getattr(callback, "__qualname__", None), - explicit_key, - ) - - if ( - inspect.isfunction(callback) - and callback.__closure__ is None - and "" not in callback.__qualname__ - ): - return ("function", callback.__module__, callback.__qualname__) - - return _UncacheableScoreModKey() - - @dataclass(frozen=True) class _FusedAttnScoreModConfig: """Static configuration for cuDNN frontend score_mod SDPA graphs.""" @@ -254,19 +192,7 @@ def __eq__(self, other): ) -@dataclass(frozen=True) -class _SerializedScoreModGraph: - """Serialized cuDNN frontend graph and static metadata for C++ execution.""" - - serialized_graph: bytes - graph_hash: Tuple[int, int] - cudnn_frontend_version: int - workspace_size: int - input_uids: np.ndarray - output_uids: np.ndarray - scalar_uids: np.ndarray - scalar_sizes: np.ndarray - scalar_values: np.ndarray +_SerializedScoreModGraph = SerializedGraph # cuDNN frontend tensor UIDs are arbitrary, but assigning stable values makes serialized @@ -288,29 +214,6 @@ class _SerializedScoreModGraph: _score_mod_graph_cache: Dict[Tuple[Any, ...], _SerializedScoreModGraph] = {} -def _row_major_stride(shape: Sequence[int]) -> Tuple[int, ...]: - stride = [] - running = 1 - for dim in reversed(tuple(shape)): - stride.append(running) - running *= dim - return tuple(reversed(stride)) - - -def _bshd_as_bhsd_dim_stride(shape: Sequence[int]) -> Tuple[Tuple[int, ...], Tuple[int, ...]]: - if len(shape) != 4: - raise ValueError(f"score_mod requires rank-4 BSHD tensors, got shape={shape}.") - batch, seqlen, heads, head_dim = tuple(shape) - return ( - (batch, heads, seqlen, head_dim), - (seqlen * heads * head_dim, head_dim, heads * head_dim, 1), - ) - - -def _dtype_name(dtype) -> str: - return str(jnp.dtype(dtype)) - - def _is_array_operand(value: Any) -> bool: return ( hasattr(value, "shape") @@ -319,6 +222,20 @@ def _is_array_operand(value: Any) -> bool: ) +def _score_mod_callback_cache_key(callback: Optional[Callable]) -> Any: + """Compatibility wrapper around the shared score-modification key policy.""" + + return score_mod_callback_cache_key( + callback, + is_array=_is_array_operand, + uncacheable_key_factory=UncacheableScoreModKey, + ) + + +def _score_mod_key_is_uncacheable(key: Any) -> bool: + return is_uncacheable_score_mod_key(key) + + def _scalar_to_spec(name: str, value: Any) -> _ScoreModScalarSpec: if isinstance(value, bool): dtype = np.bool_ @@ -405,44 +322,6 @@ def _make_fused_attn_score_mod_config( return config, tensor_operands, bprop_tensor_operands -def _cudnn_data_type(cudnn, dtype): - dtype = jnp.dtype(dtype) - if dtype == jnp.float16: - return cudnn.data_type.HALF - if dtype == jnp.bfloat16: - return cudnn.data_type.BFLOAT16 - if dtype == jnp.float32: - return cudnn.data_type.FLOAT - if dtype == jnp.float64: - return cudnn.data_type.DOUBLE - if dtype == jnp.int32: - return cudnn.data_type.INT32 - if dtype == jnp.int64: - return cudnn.data_type.INT64 - if dtype == jnp.uint8: - return cudnn.data_type.UINT8 - if dtype == jnp.bool_: - return cudnn.data_type.BOOLEAN - raise ValueError(f"Unsupported score_mod tensor dtype: {dtype}.") - - -def _cudnn_data_type_from_name(cudnn, dtype_name: str): - if dtype_name == "bfloat16": - return cudnn.data_type.BFLOAT16 - return _cudnn_data_type(cudnn, np.dtype(dtype_name)) - - -def _graph_tensor_from_aval(cudnn, graph, name: str, aval, uid: int): - shape = tuple(int(dim) for dim in aval.shape) - return graph.tensor( - name=name, - dim=shape, - stride=_row_major_stride(shape), - data_type=_cudnn_data_type(cudnn, aval.dtype), - uid=uid, - ) - - def _score_mod_graph_tensors( cudnn, graph, @@ -477,51 +356,6 @@ def _score_mod_graph_tensors( return graph_tensors, tuple(tensor_uids), tuple(scalar_uids), tuple(scalar_values) -def _encode_cudnn_frontend_version(version: str) -> int: - public_version = version.split("+", 1)[0].split("-", 1)[0] - parts = public_version.split(".") - if len(parts) < 3: - raise RuntimeError(f"Could not parse cuDNN frontend Python version: {version!r}.") - major, minor, patch = (int(part) for part in parts[:3]) - return major * 10000 + minor * 100 + patch - - -def _check_cudnn_frontend_version_match(cudnn) -> int: - python_version_string = getattr(cudnn, "__version__", None) - if python_version_string is None: - raise RuntimeError("cuDNN frontend Python package does not expose __version__.") - python_version = _encode_cudnn_frontend_version(python_version_string) - cpp_version = int(transformer_engine_jax.get_cudnn_frontend_version()) - if python_version != cpp_version: - raise RuntimeError( - "cuDNN frontend Python/C++ version mismatch for score_mod graph serialization: " - f"Python cudnn.__version__={python_version_string!r} encodes to {python_version}, " - f"but Transformer Engine C++ was built with CUDNN_FRONTEND_VERSION={cpp_version}. " - "Use matching cuDNN frontend Python package and C++ headers." - ) - return python_version - - -def _score_mod_graph_hash(serialized_graph: bytes) -> Tuple[int, int]: - digest = hashlib.sha256(serialized_graph).digest() - return ( - int.from_bytes(digest[0:8], byteorder="little", signed=True), - int.from_bytes(digest[8:16], byteorder="little", signed=True), - ) - - -def _pack_score_mod_scalar_values( - scalar_values: Sequence[bytes], -) -> Tuple[np.ndarray, np.ndarray]: - scalar_sizes = np.asarray([len(value) for value in scalar_values], dtype=np.int64) - packed_values = np.zeros((len(scalar_values), 16), dtype=np.uint8) - for index, value in enumerate(scalar_values): - if len(value) > 16: - raise ValueError("score_mod pass-by-value scalars must be at most 16 bytes.") - packed_values[index, : len(value)] = np.frombuffer(value, dtype=np.uint8) - return scalar_sizes, packed_values.reshape(-1) - - def _serialized_score_mod_graph( *, serialized_graph: bytes, @@ -531,18 +365,21 @@ def _serialized_score_mod_graph( output_uids: Sequence[int], scalar_uids: Sequence[int], scalar_values: Sequence[bytes], + cache_site: Tuple[str, str], ) -> _SerializedScoreModGraph: - scalar_sizes, packed_scalar_values = _pack_score_mod_scalar_values(scalar_values) - return _SerializedScoreModGraph( - serialized_graph=serialized_graph, - graph_hash=_score_mod_graph_hash(serialized_graph), + return make_serialized_graph( + serialized_graph_data=serialized_graph, cudnn_frontend_version=int(cudnn_frontend_version), workspace_size=int(workspace_size), - input_uids=np.asarray(input_uids, dtype=np.int64), - output_uids=np.asarray(output_uids, dtype=np.int64), + cache_site=cache_site, + input_bindings=[ + GraphBinding(uid=int(uid), buffer_index=index) for index, uid in enumerate(input_uids) + ], + output_bindings=[ + GraphBinding(uid=int(uid), buffer_index=index) for index, uid in enumerate(output_uids) + ], scalar_uids=np.asarray(scalar_uids, dtype=np.int64), - scalar_sizes=scalar_sizes, - scalar_values=packed_scalar_values, + scalar_values=scalar_values, ) @@ -556,20 +393,12 @@ def wrapped_score_mod(sdpa_graph, score_tensor): return wrapped_score_mod -def _finalize_score_mod_graph(cudnn, graph) -> Tuple[int, bytes, int]: - graph.validate() - graph.build_operation_graph() - try: - graph.create_execution_plans([cudnn.heur_mode.A, cudnn.heur_mode.FALLBACK]) - graph.check_support() - except cudnn.cudnnGraphNotSupportedError as exc: - raise RuntimeError(f"cuDNN score_mod SDPA graph is not supported: {exc}") from exc - graph.build_plans(cudnn.build_plan_policy.HEURISTICS_CHOICE) - serialized_graph = bytes(graph.serialize()) - return ( - max(int(graph.get_workspace_size()), 1), - serialized_graph, - _check_cudnn_frontend_version_match(cudnn), +def _finalize_score_mod_graph(cudnn, graph, cache_site: Tuple[str, str]) -> Tuple[int, bytes, int]: + return finalize_graph( + cudnn, + graph, + description="score_mod SDPA", + cache_site=cache_site, ) @@ -589,30 +418,15 @@ def _graph_cache_key( ) -def _shape_dtype(value) -> jax.ShapeDtypeStruct: - return jax.ShapeDtypeStruct(tuple(value.shape), value.dtype) - - def _import_cudnn_for_score_mod(): - try: - cudnn = importlib.import_module("cudnn") - except ImportError as exc: - raise ImportError( - "score_mod fused_attn requires the cuDNN frontend Python package (`cudnn`)." - ) from exc - _check_cudnn_frontend_version_match(cudnn) - return cudnn + return import_cudnn() def _build_score_mod_fwd_graph(q_aval, k_aval, v_aval, score_mod_avals, config): cudnn = _import_cudnn_for_score_mod() io_data_type = _cudnn_data_type(cudnn, q_aval.dtype) - graph = cudnn.pygraph( - io_data_type=io_data_type, - intermediate_data_type=cudnn.data_type.FLOAT, - compute_data_type=cudnn.data_type.FLOAT, - ) + graph = make_graph(cudnn, io_data_type) q_dim, q_stride = _bshd_as_bhsd_dim_stride(q_aval.shape) k_dim, k_stride = _bshd_as_bhsd_dim_stride(k_aval.shape) @@ -663,11 +477,15 @@ def _build_score_mod_fwd_graph(q_aval, k_aval, v_aval, score_mod_avals, config): stats.set_data_type(cudnn.data_type.FLOAT) output_uids.append(_SCORE_MOD_UID_STATS) - workspace_size, serialized_graph, frontend_version = _finalize_score_mod_graph(cudnn, graph) + cache_site = ("f16", "fwd") + workspace_size, serialized_graph, frontend_version = _finalize_score_mod_graph( + cudnn, graph, cache_site + ) return _serialized_score_mod_graph( serialized_graph=serialized_graph, cudnn_frontend_version=frontend_version, workspace_size=workspace_size, + cache_site=cache_site, input_uids=[_SCORE_MOD_UID_Q, _SCORE_MOD_UID_K, _SCORE_MOD_UID_V, *tensor_uids], output_uids=output_uids, scalar_uids=scalar_uids, @@ -689,11 +507,7 @@ def _build_score_mod_bwd_graph( cudnn = _import_cudnn_for_score_mod() io_data_type = _cudnn_data_type(cudnn, q_aval.dtype) - graph = cudnn.pygraph( - io_data_type=io_data_type, - intermediate_data_type=cudnn.data_type.FLOAT, - compute_data_type=cudnn.data_type.FLOAT, - ) + graph = make_graph(cudnn, io_data_type) q_dim, q_stride = _bshd_as_bhsd_dim_stride(q_aval.shape) k_dim, k_stride = _bshd_as_bhsd_dim_stride(k_aval.shape) @@ -766,11 +580,15 @@ def _build_score_mod_bwd_graph( dk.set_output(True).set_uid(_SCORE_MOD_UID_DK).set_dim(k_dim).set_stride(k_stride) dv.set_output(True).set_uid(_SCORE_MOD_UID_DV).set_dim(v_dim).set_stride(v_stride) - workspace_size, serialized_graph, frontend_version = _finalize_score_mod_graph(cudnn, graph) + cache_site = ("f16", "bwd") + workspace_size, serialized_graph, frontend_version = _finalize_score_mod_graph( + cudnn, graph, cache_site + ) return _serialized_score_mod_graph( serialized_graph=serialized_graph, cudnn_frontend_version=frontend_version, workspace_size=workspace_size, + cache_site=cache_site, input_uids=[ _SCORE_MOD_UID_Q, _SCORE_MOD_UID_K, @@ -798,13 +616,17 @@ def _fused_attn_score_mod_fwd( score_mod_avals = tuple(_shape_dtype(arg) for arg in score_mod_tensors) key = _graph_cache_key("fwd", config, (q_aval, k_aval, v_aval, *score_mod_avals)) if key is None: + record_cache_lookup(("f16", "fwd"), hit=False, key="uncacheable score_mod") graph = _build_score_mod_fwd_graph(q_aval, k_aval, v_aval, score_mod_avals, config) else: - if key not in _score_mod_graph_cache: + graph = _score_mod_graph_cache.get(key) + record_cache_lookup(("f16", "fwd"), hit=graph is not None, key=key) + if graph is None: _score_mod_graph_cache[key] = _build_score_mod_fwd_graph( q_aval, k_aval, v_aval, score_mod_avals, config ) - graph = _score_mod_graph_cache[key] + record_cache_event(("f16", "fwd"), "cache_graph") + graph = _score_mod_graph_cache[key] batch, q_seqlen, q_heads, _ = q.shape _, _, _, v_head_dim = v.shape @@ -820,15 +642,7 @@ def _fused_attn_score_mod_fwd( k, v, *score_mod_tensors, - serialized_graph=graph.serialized_graph, - graph_hash0=graph.graph_hash[0], - graph_hash1=graph.graph_hash[1], - cudnn_frontend_version=graph.cudnn_frontend_version, - input_uids=graph.input_uids, - output_uids=graph.output_uids, - scalar_uids=graph.scalar_uids, - scalar_sizes=graph.scalar_sizes, - scalar_values=graph.scalar_values, + **graph.ffi_attrs(), ) return output, softmax_stats @@ -852,6 +666,7 @@ def _fused_attn_score_mod_bwd( avals = tuple(_shape_dtype(arg) for arg in all_inputs) key = _graph_cache_key("bwd", config, avals) if key is None: + record_cache_lookup(("f16", "bwd"), hit=False, key="uncacheable score_mod") graph = _build_score_mod_bwd_graph( *avals[:6], avals[6 : 6 + len(score_mod_tensors)], @@ -859,14 +674,17 @@ def _fused_attn_score_mod_bwd( config, ) else: - if key not in _score_mod_graph_cache: + graph = _score_mod_graph_cache.get(key) + record_cache_lookup(("f16", "bwd"), hit=graph is not None, key=key) + if graph is None: _score_mod_graph_cache[key] = _build_score_mod_bwd_graph( *avals[:6], avals[6 : 6 + len(score_mod_tensors)], avals[6 + len(score_mod_tensors) :], config, ) - graph = _score_mod_graph_cache[key] + record_cache_event(("f16", "bwd"), "cache_graph") + graph = _score_mod_graph_cache[key] dq = jax.ShapeDtypeStruct(q.shape, q.dtype) dk = jax.ShapeDtypeStruct(k.shape, k.dtype) @@ -884,15 +702,7 @@ def _fused_attn_score_mod_bwd( softmax_stats, *score_mod_tensors, *score_mod_bprop_tensors, - serialized_graph=graph.serialized_graph, - graph_hash0=graph.graph_hash[0], - graph_hash1=graph.graph_hash[1], - cudnn_frontend_version=graph.cudnn_frontend_version, - input_uids=graph.input_uids, - output_uids=graph.output_uids, - scalar_uids=graph.scalar_uids, - scalar_sizes=graph.scalar_sizes, - scalar_values=graph.scalar_values, + **graph.ffi_attrs(), ) return dq, dk, dv diff --git a/transformer_engine/jax/cpp_extensions/fp8_attention.py b/transformer_engine/jax/cpp_extensions/fp8_attention.py new file mode 100644 index 00000000000..8fd5e79eb6d --- /dev/null +++ b/transformer_engine/jax/cpp_extensions/fp8_attention.py @@ -0,0 +1,1249 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""JAX execution adapter for common cuDNN FP8 attention graph construction.""" + +from __future__ import annotations + +import copy +import operator +import os +from dataclasses import dataclass +from functools import reduce +from typing import Any + +import jax +import jax.numpy as jnp +from jax import ffi + +from transformer_engine.common.attention.cudnn import ( + FusedAttentionConfig, + check_fp8_fused_attention_support, +) +from transformer_engine.common.attention.fp8 import ( + FP8AttentionGraphConfig, + attention_format_stride, + build_fp8_backward_operation, + build_fp8_forward_operation, + mxfp8_padded_sizes, +) + +from ..quantize import ScalingMode, TensorUsage, swizzle_mxfp8_scale +from .attention import _FusedAttnRNGStateChecker +from .cudnn_attention import ( + _UID_DK, + _UID_DO, + _UID_DQ, + _UID_DROPOUT_OFFSET, + _UID_DROPOUT_SEED, + _UID_DV, + _UID_K, + _UID_O, + _UID_Q, + _UID_SEQ_KV, + _UID_SEQ_Q, + _UID_STATS, + _UID_V, + _device_arch, + _is_dropout, + _is_padding, + _layout_info, + _mask_options, + _matrix_stride, + _policy_layout, + _policy_mask_name, + _qkv_bindings, +) +from .cudnn_graph import ( + GraphBinding, + SerializedGraph, + cudnn_data_type, + dtype_name, + finalize_graph, + import_cudnn, + make_graph, + record_cache_event, + record_cache_lookup, + serialized_graph, +) +from .misc import get_cudnn_version +from .quantization import quantize + +__all__ = ["FP8AttentionConfig", "fused_attn_fp8_bwd", "fused_attn_fp8_fwd"] + + +@dataclass(frozen=True) +class FP8AttentionConfig: + """Static configuration for JAX FP8 dot-product attention.""" + + attn_bias_type: Any + attn_mask_type: Any + softmax_type: Any + qkv_layout: Any + scaling_factor: float + dropout_probability: float + is_training: bool + window_size: tuple[int, int] + bottom_right_diagonal: bool = False + + +_UID_DESCALE_Q = 101 +_UID_DESCALE_K = 102 +_UID_DESCALE_V = 103 +_UID_DESCALE_S = 104 +_UID_SCALE_S = 105 +_UID_SCALE_O = 106 +_UID_AMAX_S = 107 +_UID_AMAX_O = 108 +_UID_DESCALE_O = 109 +_UID_DESCALE_DO = 110 +_UID_DESCALE_DP = 111 +_UID_SCALE_DQ = 112 +_UID_SCALE_DK = 113 +_UID_SCALE_DV = 114 +_UID_SCALE_DP = 115 +_UID_AMAX_DQ = 116 +_UID_AMAX_DK = 117 +_UID_AMAX_DV = 118 +_UID_AMAX_DP = 119 +_UID_Q_T = 120 +_UID_K_T = 121 +_UID_DO_T = 122 +_UID_DESCALE_Q_T = 123 +_UID_DESCALE_K_T = 124 +_UID_DESCALE_DO_T = 125 +_UID_DO_F16 = 126 + + +@dataclass(frozen=True) +class FP8AttentionGraphInfo: + """Serialized graph and result metadata for an FP8 attention call.""" + + graph: SerializedGraph + output_shape: tuple[int, ...] + q_shape: tuple[int, ...] + k_shape: tuple[int, ...] + v_shape: tuple[int, ...] + + +_graph_cache: dict[tuple[Any, ...], FP8AttentionGraphInfo] = {} + + +def _cache_key(direction, mode, avals, config, output_dtype): + return ( + direction, + mode, + config, + dtype_name(output_dtype), + tuple((tuple(aval.shape), dtype_name(aval.dtype)) for aval in avals), + get_cudnn_version(), + _device_arch(), + ) + + +def _tensor(graph, *, name, dim, stride, dtype, uid): + return graph.tensor( + name=name, + dim=tuple(int(value) for value in dim), + stride=tuple(int(value) for value in stride), + data_type=dtype, + uid=uid, + ) + + +def _scalar(graph, cudnn, name, uid): + return _tensor( + graph, + name=name, + dim=(1, 1, 1, 1), + stride=(1, 1, 1, 1), + dtype=cudnn.data_type.FLOAT, + uid=uid, + ) + + +def _mx_scale( + graph, + cudnn, + *, + name, + uid, + batch, + heads, + seqlen, + dim, +): + # _mxfp8_scale_inv transposes every compact scale buffer to contiguous BHSD + # before applying cuDNN's F8_128x4 physical reordering. + return _tensor( + graph, + name=name, + dim=(batch, heads, seqlen, dim), + stride=attention_format_stride(batch, heads, seqlen, dim, "bhsd"), + dtype=cudnn.data_type.FP8_E8M0, + uid=uid, + ).set_reordering_type(cudnn.tensor_reordering.F8_128x4) + + +def _logical_shapes(info): + q = (*info.batch_shape, info.q_max_seqlen, info.q_heads, info.qk_dim) + k = (*info.batch_shape, info.kv_max_seqlen, info.kv_heads, info.qk_dim) + v = (*info.batch_shape, info.kv_max_seqlen, info.kv_heads, info.v_dim) + o = (*info.batch_shape, info.q_max_seqlen, info.q_heads, info.v_dim) + return q, k, v, o + + +def _graph_io_tensors(graph, cudnn, q_aval, k_aval, v_aval, config): + info = _layout_info(q_aval, k_aval, v_aval, config.qkv_layout) + io_dtype = cudnn_data_type(cudnn, q_aval.dtype) + q = _tensor( + graph, + name="Q", + dim=(info.input_batch, info.q_heads, info.q_max_seqlen, info.qk_dim), + stride=_matrix_stride(info, config.qkv_layout, "q", info.q_max_seqlen, info.kv_max_seqlen), + dtype=io_dtype, + uid=_UID_Q, + ) + k = _tensor( + graph, + name="K", + dim=(info.input_batch, info.kv_heads, info.kv_max_seqlen, info.qk_dim), + stride=_matrix_stride(info, config.qkv_layout, "k", info.q_max_seqlen, info.kv_max_seqlen), + dtype=io_dtype, + uid=_UID_K, + ) + v = _tensor( + graph, + name="V", + dim=(info.input_batch, info.kv_heads, info.kv_max_seqlen, info.v_dim), + stride=_matrix_stride(info, config.qkv_layout, "v", info.q_max_seqlen, info.kv_max_seqlen), + dtype=io_dtype, + uid=_UID_V, + ) + return info, io_dtype, q, k, v + + +def _fp8_options(cudnn, info, config): + options = _mask_options(cudnn, info, config) + is_padding = options.pop("is_padding", _is_padding(config)) + if "diagonal_band_left_bound" in options: + options["left_bound"] = options.pop("diagonal_band_left_bound") + if "diagonal_band_right_bound" in options: + options["right_bound"] = options.pop("diagonal_band_right_bound") + options.update( + attn_scale=float(config.scaling_factor), + use_padding_mask=is_padding, + ) + return options, is_padding + + +def build_fp8_fwd_graph( + q_aval, + k_aval, + v_aval, + q_scale_aval, + k_scale_aval, + v_scale_aval, + config, + mode: str, + output_dtype, +) -> FP8AttentionGraphInfo: + """Build or retrieve a dense JAX FP8 attention forward graph.""" + + avals = (q_aval, k_aval, v_aval, q_scale_aval, k_scale_aval, v_scale_aval) + key = _cache_key("fwd", mode, avals, config, output_dtype) + graph_info = _graph_cache.get(key) + record_cache_lookup(("fp8", "fwd"), hit=graph_info is not None, key=key) + if graph_info is not None: + return graph_info + + cudnn = import_cudnn() + graph = make_graph(cudnn, cudnn_data_type(cudnn, q_aval.dtype)) + info, _, q, k, v = _graph_io_tensors(graph, cudnn, q_aval, k_aval, v_aval, config) + q_shape, k_shape, v_shape, o_shape = _logical_shapes(info) + input_bindings = list(_qkv_bindings(info, config.qkv_layout, jnp.dtype(q_aval.dtype).itemsize)) + tensors = {"q": q, "k": k, "v": v} + options, is_padding = _fp8_options(cudnn, info, config) + options["generate_stats"] = True + + if mode == "mxfp8": + if is_padding: + raise ValueError("JAX MXFP8 attention does not support padding masks.") + options.pop("use_padding_mask", None) + # The current cuDNN MXFP8 forward binding uses diagonal_band_* names. + if "left_bound" in options: + options["diagonal_band_left_bound"] = options.pop("left_bound") + if "right_bound" in options: + options["diagonal_band_right_bound"] = options.pop("right_bound") + padded = mxfp8_padded_sizes(info.q_max_seqlen, info.kv_max_seqlen, info.qk_dim, info.v_dim) + scale_specs = ( + ( + "descale_q", + _UID_DESCALE_Q, + info.q_heads, + "s_q_padded", + "d_qk_scale_padded", + 3, + ), + ( + "descale_k", + _UID_DESCALE_K, + info.kv_heads, + "s_kv_padded", + "d_qk_scale_padded", + 4, + ), + ( + "descale_v", + _UID_DESCALE_V, + info.kv_heads, + "s_kv_scale_padded", + "d_v_padded", + 6, + ), + ) + for name, uid, heads, s_key, d_key, buffer_index in scale_specs: + tensors[name] = _mx_scale( + graph, + cudnn, + name=name, + uid=uid, + batch=info.input_batch, + heads=heads, + seqlen=padded[s_key], + dim=padded[d_key], + ) + input_bindings.append(GraphBinding(uid, buffer_index)) + else: + for name, uid, index in ( + ("descale_q", _UID_DESCALE_Q, 3), + ("descale_k", _UID_DESCALE_K, 4), + ("descale_v", _UID_DESCALE_V, 6), + ("descale_s", _UID_DESCALE_S, 7), + ("scale_s", _UID_SCALE_S, 8), + ("scale_o", _UID_SCALE_O, 9), + ): + tensors[name] = _scalar(graph, cudnn, name, uid) + input_bindings.append(GraphBinding(uid, index)) + + if is_padding: + seq_q = _tensor( + graph, + name="seq_len_q", + dim=(info.input_batch, 1, 1, 1), + stride=(1, 1, 1, 1), + dtype=cudnn.data_type.INT32, + uid=_UID_SEQ_Q, + ) + seq_kv = _tensor( + graph, + name="seq_len_kv", + dim=(info.input_batch, 1, 1, 1), + stride=(1, 1, 1, 1), + dtype=cudnn.data_type.INT32, + uid=_UID_SEQ_KV, + ) + options.update(seq_len_q=seq_q, seq_len_kv=seq_kv) + input_bindings.extend((GraphBinding(_UID_SEQ_Q, 10), GraphBinding(_UID_SEQ_KV, 11))) + + output_bindings = [GraphBinding(_UID_O, 0), GraphBinding(_UID_STATS, 1)] + if _is_dropout(config): + seed = _tensor( + graph, + name="dropout_seed", + dim=(1, 1, 1, 1), + stride=(1, 1, 1, 1), + dtype=cudnn.data_type.INT64, + uid=_UID_DROPOUT_SEED, + ) + offset = _tensor( + graph, + name="dropout_offset", + dim=(1, 1, 1, 1), + stride=(1, 1, 1, 1), + dtype=cudnn.data_type.INT64, + uid=_UID_DROPOUT_OFFSET, + ) + options["dropout"] = (float(config.dropout_probability), seed, offset) + output_bindings.extend( + ( + GraphBinding(_UID_DROPOUT_SEED, 3), + GraphBinding(_UID_DROPOUT_OFFSET, 3, 8), + ) + ) + + op = build_fp8_forward_operation( + graph, + tensors, + options, + FP8AttentionGraphConfig(mode, "te_jax_fp8_sdpa_forward"), + ) + output = op["output"] + output.set_output(True).set_uid(_UID_O).set_data_type( + cudnn_data_type(cudnn, output_dtype) + ).set_dim((info.input_batch, info.q_heads, info.q_max_seqlen, info.v_dim)).set_stride( + attention_format_stride( + info.input_batch, info.q_heads, info.q_max_seqlen, info.v_dim, "bshd" + ) + ) + stats = op["stats"] + stats.set_output(True).set_uid(_UID_STATS).set_data_type(cudnn.data_type.FLOAT).set_dim( + (info.input_batch, info.q_heads, info.q_max_seqlen, 1) + ).set_stride((info.q_heads * info.q_max_seqlen, info.q_max_seqlen, 1, 1)) + if mode != "mxfp8": + for name, uid, offset in ( + ("amax_s", _UID_AMAX_S, 0), + ("amax_o", _UID_AMAX_O, 4), + ): + op[name].set_output(True).set_uid(uid).set_data_type(cudnn.data_type.FLOAT).set_dim( + (1, 1, 1, 1) + ).set_stride((1, 1, 1, 1)) + output_bindings.append(GraphBinding(uid, 2, offset)) + else: + op["amax_o"].set_output(False).set_data_type(cudnn.data_type.FLOAT).set_dim( + (1, 1, 1, 1) + ).set_stride((1, 1, 1, 1)) + + cache_site = ("fp8", "fwd") + workspace, data, version = finalize_graph( + cudnn, + graph, + description=f"JAX {mode} FP8 attention forward", + cache_site=cache_site, + ) + result = serialized_graph( + serialized_graph_data=data, + cudnn_frontend_version=version, + workspace_size=workspace, + cache_site=cache_site, + input_bindings=input_bindings, + output_bindings=output_bindings, + ) + info_result = FP8AttentionGraphInfo(result, o_shape, q_shape, k_shape, v_shape) + _graph_cache[key] = info_result + record_cache_event(cache_site, "cache_graph") + return info_result + + +def _mx_bwd_scales(graph, cudnn, info, padded, input_bindings): + specs = ( + ( + "descale_q", + _UID_DESCALE_Q, + info.q_heads, + "s_q_padded", + "d_qk_scale_padded", + 11, + ), + ( + "descale_q_t", + _UID_DESCALE_Q_T, + info.q_heads, + "s_q_scale_padded", + "d_qk_padded", + 25, + ), + ( + "descale_k", + _UID_DESCALE_K, + info.kv_heads, + "s_kv_padded", + "d_qk_scale_padded", + 12, + ), + ( + "descale_k_t", + _UID_DESCALE_K_T, + info.kv_heads, + "s_kv_scale_padded", + "d_qk_padded", + 26, + ), + ( + "descale_v", + _UID_DESCALE_V, + info.kv_heads, + "s_kv_padded", + "d_v_scale_padded", + 13, + ), + ( + "descale_do", + _UID_DESCALE_DO, + info.q_heads, + "s_q_padded", + "d_v_scale_padded", + 15, + ), + ( + "descale_do_t", + _UID_DESCALE_DO_T, + info.q_heads, + "s_q_scale_padded", + "d_v_padded", + 24, + ), + ) + tensors = {} + for name, uid, heads, s_key, d_key, index in specs: + tensors[name] = _mx_scale( + graph, + cudnn, + name=name, + uid=uid, + batch=info.input_batch, + heads=heads, + seqlen=padded[s_key], + dim=padded[d_key], + ) + input_bindings.append(GraphBinding(uid, index)) + return tensors + + +def build_fp8_bwd_graph( + q_aval, + k_aval, + v_aval, + stats_aval, + output_aval, + doutput_aval, + config, + mode: str, + grad_dtype, +) -> FP8AttentionGraphInfo: + """Build or retrieve a dense JAX FP8 attention backward graph.""" + + avals = (q_aval, k_aval, v_aval, stats_aval, output_aval, doutput_aval) + key = _cache_key("bwd", mode, avals, config, grad_dtype) + graph_info = _graph_cache.get(key) + record_cache_lookup(("fp8", "bwd"), hit=graph_info is not None, key=key) + if graph_info is not None: + return graph_info + + cudnn = import_cudnn() + graph = make_graph(cudnn, cudnn_data_type(cudnn, q_aval.dtype)) + info, io_dtype, q, k, v = _graph_io_tensors(graph, cudnn, q_aval, k_aval, v_aval, config) + q_shape, k_shape, v_shape, o_shape = _logical_shapes(info) + input_bindings = list(_qkv_bindings(info, config.qkv_layout, jnp.dtype(q_aval.dtype).itemsize)) + o = _tensor( + graph, + name="O", + dim=(info.input_batch, info.q_heads, info.q_max_seqlen, info.v_dim), + stride=attention_format_stride( + info.input_batch, info.q_heads, info.q_max_seqlen, info.v_dim, "bshd" + ), + dtype=cudnn_data_type(cudnn, output_aval.dtype), + uid=_UID_O, + ) + do = _tensor( + graph, + name="dO", + dim=(info.input_batch, info.q_heads, info.q_max_seqlen, info.v_dim), + stride=attention_format_stride( + info.input_batch, info.q_heads, info.q_max_seqlen, info.v_dim, "bshd" + ), + dtype=cudnn_data_type(cudnn, doutput_aval.dtype), + uid=_UID_DO, + ) + stats = _tensor( + graph, + name="Stats", + dim=(info.input_batch, info.q_heads, info.q_max_seqlen, 1), + stride=(info.q_heads * info.q_max_seqlen, info.q_max_seqlen, 1, 1), + dtype=cudnn.data_type.FLOAT, + uid=_UID_STATS, + ) + input_bindings.extend( + (GraphBinding(_UID_STATS, 5), GraphBinding(_UID_O, 7), GraphBinding(_UID_DO, 8)) + ) + tensors = {"q": q, "k": k, "v": v, "o": o, "do": do, "stats": stats} + options, is_padding = _fp8_options(cudnn, info, config) + options["use_deterministic_algorithm"] = not bool( + int(os.getenv("NVTE_ALLOW_NONDETERMINISTIC_ALGO", "1")) + ) + if is_padding: + seq_q = _tensor( + graph, + name="seq_len_q", + dim=(info.input_batch, 1, 1, 1), + stride=(1, 1, 1, 1), + dtype=cudnn.data_type.INT32, + uid=_UID_SEQ_Q, + ) + seq_kv = _tensor( + graph, + name="seq_len_kv", + dim=(info.input_batch, 1, 1, 1), + stride=(1, 1, 1, 1), + dtype=cudnn.data_type.INT32, + uid=_UID_SEQ_KV, + ) + options.update(seq_len_q=seq_q, seq_len_kv=seq_kv) + input_bindings.extend((GraphBinding(_UID_SEQ_Q, 9), GraphBinding(_UID_SEQ_KV, 10))) + + if _is_dropout(config): + seed = _tensor( + graph, + name="dropout_seed", + dim=(1, 1, 1, 1), + stride=(1, 1, 1, 1), + dtype=cudnn.data_type.INT64, + uid=_UID_DROPOUT_SEED, + ) + offset = _tensor( + graph, + name="dropout_offset", + dim=(1, 1, 1, 1), + stride=(1, 1, 1, 1), + dtype=cudnn.data_type.INT64, + uid=_UID_DROPOUT_OFFSET, + ) + options["dropout"] = (float(config.dropout_probability), seed, offset) + input_bindings.extend( + ( + GraphBinding(_UID_DROPOUT_SEED, 6), + GraphBinding(_UID_DROPOUT_OFFSET, 6, 8), + ) + ) + + if mode == "mxfp8": + if is_padding: + raise ValueError("JAX MXFP8 attention does not support padding masks.") + q_t = _tensor( + graph, + name="Q_T", + dim=(info.input_batch, info.q_heads, info.q_max_seqlen, info.qk_dim), + stride=attention_format_stride( + info.input_batch, info.q_heads, info.q_max_seqlen, info.qk_dim, "bshd" + ), + dtype=io_dtype, + uid=_UID_Q_T, + ) + k_t = _tensor( + graph, + name="K_T", + dim=(info.input_batch, info.kv_heads, info.kv_max_seqlen, info.qk_dim), + stride=attention_format_stride( + info.input_batch, info.kv_heads, info.kv_max_seqlen, info.qk_dim, "bshd" + ), + dtype=io_dtype, + uid=_UID_K_T, + ) + do_t = _tensor( + graph, + name="dO_T", + dim=(info.input_batch, info.q_heads, info.q_max_seqlen, info.v_dim), + stride=attention_format_stride( + info.input_batch, info.q_heads, info.q_max_seqlen, info.v_dim, "bshd" + ), + dtype=cudnn_data_type(cudnn, doutput_aval.dtype), + uid=_UID_DO_T, + ) + do_f16 = _tensor( + graph, + name="dO_f16", + dim=(info.input_batch, info.q_heads, info.q_max_seqlen, info.v_dim), + stride=attention_format_stride( + info.input_batch, info.q_heads, info.q_max_seqlen, info.v_dim, "bshd" + ), + dtype=cudnn_data_type(cudnn, output_aval.dtype), + uid=_UID_DO_F16, + ) + tensors.update(q_t=q_t, k_t=k_t, do_t=do_t, do_f16=do_f16) + input_bindings.extend( + ( + GraphBinding(_UID_Q_T, 3), + GraphBinding(_UID_K_T, 4), + GraphBinding(_UID_DO_T, 23), + GraphBinding(_UID_DO_F16, 27), + ) + ) + padded = mxfp8_padded_sizes(info.q_max_seqlen, info.kv_max_seqlen, info.qk_dim, info.v_dim) + tensors.update(_mx_bwd_scales(graph, cudnn, info, padded, input_bindings)) + else: + for name, uid, index in ( + ("descale_q", _UID_DESCALE_Q, 11), + ("descale_k", _UID_DESCALE_K, 12), + ("descale_v", _UID_DESCALE_V, 13), + ("descale_o", _UID_DESCALE_O, 14), + ("descale_do", _UID_DESCALE_DO, 15), + ("descale_s", _UID_DESCALE_S, 16), + ("descale_dp", _UID_DESCALE_DP, 17), + ("scale_s", _UID_SCALE_S, 18), + ("scale_dq", _UID_SCALE_DQ, 19), + ("scale_dk", _UID_SCALE_DK, 20), + ("scale_dv", _UID_SCALE_DV, 21), + ("scale_dp", _UID_SCALE_DP, 22), + ): + tensors[name] = _scalar(graph, cudnn, name, uid) + input_bindings.append(GraphBinding(uid, index)) + + op = build_fp8_backward_operation( + graph, + tensors, + options, + FP8AttentionGraphConfig(mode, "te_jax_fp8_sdpa_backward"), + ) + output_bindings = [] + for name, uid, index, shape in ( + ("dq", _UID_DQ, 0, q_shape), + ("dk", _UID_DK, 1, k_shape), + ("dv", _UID_DV, 2, v_shape), + ): + batch = reduce(operator.mul, shape[:-3], 1) + seqlen, heads, dim = shape[-3:] + op[name].set_output(True).set_uid(uid).set_data_type( + cudnn_data_type(cudnn, grad_dtype) + ).set_dim((batch, heads, seqlen, dim)).set_stride( + attention_format_stride(batch, heads, seqlen, dim, "bshd") + ) + output_bindings.append(GraphBinding(uid, index)) + + if mode != "mxfp8": + for name, uid, offset in ( + ("amax_dq", _UID_AMAX_DQ, 0), + ("amax_dk", _UID_AMAX_DK, 4), + ("amax_dv", _UID_AMAX_DV, 8), + ("amax_dp", _UID_AMAX_DP, 12), + ): + op[name].set_output(True).set_uid(uid).set_data_type(cudnn.data_type.FLOAT).set_dim( + (1, 1, 1, 1) + ).set_stride((1, 1, 1, 1)) + output_bindings.append(GraphBinding(uid, 3, offset)) + else: + for amax in op["amax"]: + amax.set_output(False).set_data_type(cudnn.data_type.FLOAT).set_dim( + (1, 1, 1, 1) + ).set_stride((1, 1, 1, 1)) + + cache_site = ("fp8", "bwd") + workspace, data, version = finalize_graph( + cudnn, + graph, + description=f"JAX {mode} FP8 attention backward", + cache_site=cache_site, + ) + result = serialized_graph( + serialized_graph_data=data, + cudnn_frontend_version=version, + workspace_size=workspace, + cache_site=cache_site, + input_bindings=input_bindings, + output_bindings=output_bindings, + ) + graph_info = FP8AttentionGraphInfo(result, o_shape, q_shape, k_shape, v_shape) + _graph_cache[key] = graph_info + record_cache_event(cache_site, "cache_graph") + return graph_info + + +def execute_fp8_fwd( + q, + k, + v, + q_scale_inv, + k_scale_inv, + v_scale_inv, + s_scale_inv, + s_scale, + o_scale, + q_seqlen, + kv_seqlen, + seed, + *, + config, + mode, + output_dtype, +): + """Execute a serialized dense FP8 attention forward graph.""" + + graph_info = build_fp8_fwd_graph( + jax.ShapeDtypeStruct(q.shape, q.dtype), + jax.ShapeDtypeStruct(k.shape, k.dtype), + jax.ShapeDtypeStruct(v.shape, v.dtype), + jax.ShapeDtypeStruct(q_scale_inv.shape, q_scale_inv.dtype), + jax.ShapeDtypeStruct(k_scale_inv.shape, k_scale_inv.dtype), + jax.ShapeDtypeStruct(v_scale_inv.shape, v_scale_inv.dtype), + config, + mode, + output_dtype, + ) + result_specs = ( + jax.ShapeDtypeStruct(graph_info.output_shape, output_dtype), + jax.ShapeDtypeStruct( + ( + *graph_info.output_shape[:-3], + graph_info.output_shape[-2], + graph_info.output_shape[-3], + 1, + ), + jnp.float32, + ), + jax.ShapeDtypeStruct((2,), jnp.float32), + jax.ShapeDtypeStruct((seed.shape[0], 4), jnp.uint32), + jax.ShapeDtypeStruct((graph_info.graph.workspace_size,), jnp.uint8), + ) + return ffi.ffi_call("te_fused_attn_forward_ffi", result_specs)( + q, + k, + v, + q_scale_inv, + k_scale_inv, + seed, + v_scale_inv, + s_scale_inv, + s_scale, + o_scale, + q_seqlen, + kv_seqlen, + is_ragged=False, + rng_offset_increment=16, + **graph_info.graph.ffi_attrs(), + ) + + +def execute_fp8_bwd( + q, + k, + v, + q_t, + k_t, + stats, + rng_state, + output, + doutput, + q_seqlen, + kv_seqlen, + q_scale_inv, + k_scale_inv, + v_scale_inv, + o_scale_inv, + do_scale_inv, + s_scale_inv, + dp_scale_inv, + s_scale, + dq_scale, + dk_scale, + dv_scale, + dp_scale, + do_t, + do_scale_inv_t, + q_scale_inv_t, + k_scale_inv_t, + doutput_f16, + *, + config, + mode, + grad_dtype, +): + """Execute a serialized dense FP8 attention backward graph.""" + + graph_info = build_fp8_bwd_graph( + jax.ShapeDtypeStruct(q.shape, q.dtype), + jax.ShapeDtypeStruct(k.shape, k.dtype), + jax.ShapeDtypeStruct(v.shape, v.dtype), + jax.ShapeDtypeStruct(stats.shape, stats.dtype), + jax.ShapeDtypeStruct(output.shape, output.dtype), + jax.ShapeDtypeStruct(doutput.shape, doutput.dtype), + config, + mode, + grad_dtype, + ) + result_specs = ( + jax.ShapeDtypeStruct(graph_info.q_shape, grad_dtype), + jax.ShapeDtypeStruct(graph_info.k_shape, grad_dtype), + jax.ShapeDtypeStruct(graph_info.v_shape, grad_dtype), + jax.ShapeDtypeStruct((4,), jnp.float32), + jax.ShapeDtypeStruct((0,), output.dtype), + jax.ShapeDtypeStruct((graph_info.graph.workspace_size,), jnp.uint8), + ) + return ffi.ffi_call("te_fused_attn_backward_ffi", result_specs)( + q, + k, + v, + q_t, + k_t, + stats, + rng_state, + output, + doutput, + q_seqlen, + kv_seqlen, + q_scale_inv, + k_scale_inv, + v_scale_inv, + o_scale_inv, + do_scale_inv, + s_scale_inv, + dp_scale_inv, + s_scale, + dq_scale, + dk_scale, + dv_scale, + dp_scale, + do_t, + do_scale_inv_t, + q_scale_inv_t, + k_scale_inv_t, + doutput_f16, + is_ragged=False, + **graph_info.graph.ffi_attrs(), + ) + + +def _scaling_mode(quantizer) -> str: + mode = quantizer.scaling_mode + if mode == ScalingMode.DELAYED_TENSOR_SCALING: + return "delayed" + if mode == ScalingMode.CURRENT_TENSOR_SCALING: + return "current" + if mode == ScalingMode.MXFP8_1D_SCALING: + return "mxfp8" + raise ValueError(f"FP8 attention does not support scaling mode {mode}.") + + +def _validate_quantizer_modes(quantizers) -> str: + """Validate the supported scaling-mode assignment for each DPA tensor role.""" + + roles = ("qkv", "s", "o", "do", "dp", "dqkv") + mode = _scaling_mode(quantizers.qkv) + expected = { + "delayed": {role: "delayed" for role in roles}, + "current": { + "qkv": "current", + "s": "delayed", + "o": "current", + "do": "current", + "dp": "delayed", + "dqkv": "current", + }, + "mxfp8": {role: "mxfp8" for role in roles}, + }[mode] + actual = {role: _scaling_mode(getattr(quantizers, role)) for role in roles} + mismatches = [ + f"{role}={actual[role]} (expected {expected[role]})" + for role in roles + if actual[role] != expected[role] + ] + if mismatches: + raise ValueError("Unsupported FP8 attention quantizer modes: " + ", ".join(mismatches)) + return mode + + +def _rowwise(tensor): + return tensor.get_tensor(TensorUsage.LHS) + + +def _colwise(tensor): + return tensor.get_tensor(TensorUsage.RHS) + + +def _quantize_many(values, quantizer, *, both=False): + tensors = [] + amaxes = [] + for value in values: + local_quantizer = copy.copy(quantizer) + tensor = quantize(value, quantizer=local_quantizer, flatten_axis=-2) + tensors.append(tensor) + rowwise_amax = _rowwise(tensor).amax + if rowwise_amax is not None: + amaxes.append(rowwise_amax) + if quantizer.scaling_mode == ScalingMode.DELAYED_TENSOR_SCALING and amaxes: + quantizer.update(jnp.max(jnp.stack(amaxes))) + if both: + # Access both layouts here so an invalid recipe/layout fails before graph construction. + for tensor in tensors: + _rowwise(tensor) + _colwise(tensor) + return tuple(tensors) + + +def _quantized_operands(qkv, layout, quantizer, *, both=False): + quantized = _quantize_many(qkv, quantizer, both=both) + empty = jnp.zeros((0,), dtype=quantizer.q_dtype) + if layout.is_qkvpacked(): + tensor = quantized[0] + return (tensor, tensor, tensor), (_rowwise(tensor).data, empty, empty) + if layout.is_kvpacked(): + q_tensor, kv_tensor = quantized[0], quantized[1] + return ( + q_tensor, + kv_tensor, + kv_tensor, + ), (_rowwise(q_tensor).data, _rowwise(kv_tensor).data, empty) + if layout.is_separate(): + return quantized, tuple(_rowwise(tensor).data for tensor in quantized) + raise ValueError(f"FP8 attention does not support layout {layout}.") + + +def _scale_inv(tensor, *, colwise=False): + return (_colwise(tensor) if colwise else _rowwise(tensor)).scale_inv + + +def _mxfp8_scale_inv(tensor, *, colwise=False): + """Prepare a compact BSHD MXFP8 scale tensor for cuDNN's F8_128x4 layout.""" + + component = _colwise(tensor) if colwise else _rowwise(tensor) + *batch_shape, seqlen, heads, dim = component.data.shape + batch = reduce(operator.mul, batch_shape, 1) + if colwise: + scale = component.scale_inv.reshape(batch, seqlen // 32, heads, dim) + target_seqlen = ((seqlen + 127) // 128) * 4 + target_dim = ((dim + 127) // 128) * 128 + else: + scale = component.scale_inv.reshape(batch, seqlen, heads, dim // 32) + target_seqlen = ((seqlen + 127) // 128) * 128 + target_dim = ((dim + 127) // 128) * 4 + scale = jnp.pad( + scale, + ( + (0, 0), + (0, target_seqlen - scale.shape[1]), + (0, 0), + (0, target_dim - scale.shape[3]), + ), + mode="constant", + constant_values=2**-127, + ) + scale = jnp.transpose(scale, (0, 2, 1, 3)) + return swizzle_mxfp8_scale(scale, -1, colwise) + + +def _tensor_scale(quantizer): + if quantizer.scaling_mode == ScalingMode.DELAYED_TENSOR_SCALING: + return quantizer.scale + return jnp.ones((1,), dtype=jnp.float32) + + +def _tensor_scale_inv(quantizer): + return jnp.reciprocal(_tensor_scale(quantizer)) + + +def _graph_scale_inv(tensor, mode, *, colwise=False): + if mode == "mxfp8": + return _mxfp8_scale_inv(tensor, colwise=colwise) + return _scale_inv(tensor, colwise=colwise) + + +def _sequence_lengths(sequence_descriptor, config): + (q_seqlen, kv_seqlen), _ = sequence_descriptor.get_seqlens_and_offsets( + config.attn_mask_type, + config.qkv_layout, + config.window_size, + 1, + ) + return q_seqlen.flatten(), kv_seqlen.flatten() + + +def _validate_fp8_support(qkv, quantizers, config, mode): + if config.qkv_layout.is_qkvpacked(): + q = k = v = qkv[0] + elif config.qkv_layout.is_kvpacked(): + q, k = qkv + v = k + else: + q, k, v = qkv + info = _layout_info(q, k, v, config.qkv_layout) + q_dtype = jnp.dtype(quantizers.qkv.q_dtype) + dtype_name_ = { + jnp.dtype(jnp.float8_e4m3fn): "float8_e4m3", + jnp.dtype(jnp.float8_e5m2): "float8_e5m2", + }.get(q_dtype, str(q_dtype)) + cudnn_version = get_cudnn_version() + device_arch = _device_arch() + support = check_fp8_fused_attention_support( + FusedAttentionConfig( + is_training=bool(config.is_training), + q_dtype=dtype_name_, + kv_dtype=dtype_name_, + layout=_policy_layout(config.qkv_layout), + bias_type=config.attn_bias_type.name.lower(), + mask_type=_policy_mask_name(config.attn_mask_type), + softmax_type=config.softmax_type.name.lower().removesuffix("_softmax"), + dropout=float(config.dropout_probability), + num_attn_heads=info.q_heads, + num_gqa_groups=info.kv_heads, + max_seqlen_q=info.q_max_seqlen, + max_seqlen_kv=info.kv_max_seqlen, + head_dim_qk=info.qk_dim, + head_dim_v=info.v_dim, + window_size=tuple(int(value) for value in config.window_size), + return_max_logit=False, + cuda_graph=False, + deterministic=not bool(int(os.getenv("NVTE_ALLOW_NONDETERMINISTIC_ALGO", "1"))), + cudnn_version=cudnn_version, + sm_arch=device_arch, + ), + scaling_mode=mode, + ) + if not support.supported: + raise ValueError(f"Unsupported JAX FP8 attention configuration: {support.reason}.") + + +def fused_attn_fp8_fwd(qkv, sequence_descriptor, seed, quantizers, config): + """Quantize high-precision inputs and execute dense FP8 attention forward.""" + + mode = _validate_quantizer_modes(quantizers) + if config.qkv_layout.is_thd(): + raise NotImplementedError("FP8 attention does not support THD layouts in JAX.") + if mode == "mxfp8" and not config.qkv_layout.is_separate(): + raise NotImplementedError("JAX MXFP8 attention currently requires separate BSHD Q/K/V.") + if getattr(config.attn_bias_type, "name", "") != "NO_BIAS": + raise NotImplementedError("FP8 attention does not support attention bias.") + if getattr(config.softmax_type, "name", "") != "VANILLA_SOFTMAX": + raise NotImplementedError("JAX FP8 attention currently supports vanilla softmax only.") + _validate_fp8_support(qkv, quantizers, config, mode) + if mode == "mxfp8" and _is_padding(config): + raise NotImplementedError("JAX MXFP8 attention does not support padding masks.") + + quantized, data = _quantized_operands( + qkv, config.qkv_layout, quantizers.qkv, both=mode == "mxfp8" + ) + q_tensor, k_tensor, v_tensor = quantized + q_data, k_data, v_data = data + # MXFP8 forward consumes V in the columnwise orientation. + if mode == "mxfp8": + v_data = _colwise(v_tensor).data + q_seqlen, kv_seqlen = _sequence_lengths(sequence_descriptor, config) + seed = _FusedAttnRNGStateChecker().check_seed( + seed, config.dropout_probability, config.is_training + ) + output_is_delayed = _scaling_mode(quantizers.o) == "delayed" + output_dtype = quantizers.o.q_dtype if output_is_delayed else qkv[0].dtype + s_scale = _tensor_scale(quantizers.s) + s_scale_inv = jnp.reciprocal(s_scale) + output_scale_inv = _tensor_scale_inv(quantizers.o) + raw_output, stats, amax, rng_state, _ = execute_fp8_fwd( + q_data, + k_data, + v_data, + _graph_scale_inv(q_tensor, mode), + _graph_scale_inv(k_tensor, mode), + _graph_scale_inv(v_tensor, mode, colwise=mode == "mxfp8"), + s_scale_inv, + s_scale, + _tensor_scale(quantizers.o), + q_seqlen, + kv_seqlen, + seed, + config=config, + mode=mode, + output_dtype=output_dtype, + ) + if _scaling_mode(quantizers.s) == "delayed": + quantizers.s.update(amax[0:1]) + if output_is_delayed: + quantizers.o.update(amax[1:2]) + output = (raw_output.astype(qkv[0].dtype) * output_scale_inv).astype(qkv[0].dtype) + else: + output = raw_output + return output, ( + quantized, + raw_output, + stats, + rng_state, + q_seqlen, + kv_seqlen, + quantizers, + s_scale, + s_scale_inv, + output_scale_inv, + ) + + +def _split_gradient_outputs(dq, dk, dv, layout): + if layout.is_qkvpacked(): + return (jnp.stack((dq, dk, dv), axis=-3),) + if layout.is_kvpacked(): + return dq, jnp.stack((dk, dv), axis=-3) + return dq, dk, dv + + +def fused_attn_fp8_bwd(ctx, doutput, config): + """Execute FP8 attention backward and update delayed-scaling quantizers.""" + + ( + quantized, + raw_output, + stats, + rng_state, + q_seqlen, + kv_seqlen, + quantizers, + s_scale, + s_scale_inv, + output_scale_inv, + ) = ctx + q_tensor, k_tensor, v_tensor = quantized + mode = _scaling_mode(quantizers.qkv) + input_dtype = _rowwise(q_tensor).dq_dtype + do_tensor = _quantize_many((doutput,), quantizers.do, both=mode == "mxfp8")[0] + q_data, k_data, v_data = ( + _rowwise(q_tensor).data, + _rowwise(k_tensor).data, + _rowwise(v_tensor).data, + ) + empty_data = jnp.zeros((0,), dtype=q_data.dtype) + empty_scale = jnp.ones((1,), dtype=jnp.float32) + if config.qkv_layout.is_qkvpacked(): + k_data = v_data = empty_data + elif config.qkv_layout.is_kvpacked(): + v_data = empty_data + if mode == "mxfp8": + q_t, k_t = _colwise(q_tensor).data, _colwise(k_tensor).data + do_t = _colwise(do_tensor).data + q_scale_t, k_scale_t = ( + _graph_scale_inv(q_tensor, mode, colwise=True), + _graph_scale_inv(k_tensor, mode, colwise=True), + ) + do_scale_t = _graph_scale_inv(do_tensor, mode, colwise=True) + else: + q_t = k_t = do_t = empty_data + q_scale_t = k_scale_t = do_scale_t = empty_scale + + grad_is_delayed = _scaling_mode(quantizers.dqkv) == "delayed" + grad_dtype = quantizers.dqkv.q_dtype if grad_is_delayed else input_dtype + grad_scale_inv = _tensor_scale_inv(quantizers.dqkv) + dq, dk, dv, amax, _, _ = execute_fp8_bwd( + q_data, + k_data, + v_data, + q_t, + k_t, + stats, + rng_state, + raw_output, + _rowwise(do_tensor).data, + q_seqlen, + kv_seqlen, + _graph_scale_inv(q_tensor, mode), + _graph_scale_inv(k_tensor, mode), + _graph_scale_inv(v_tensor, mode), + output_scale_inv, + _graph_scale_inv(do_tensor, mode), + s_scale_inv, + _tensor_scale_inv(quantizers.dp), + s_scale, + _tensor_scale(quantizers.dqkv), + _tensor_scale(quantizers.dqkv), + _tensor_scale(quantizers.dqkv), + _tensor_scale(quantizers.dp), + do_t, + do_scale_t, + q_scale_t, + k_scale_t, + doutput, + config=config, + mode=mode, + grad_dtype=grad_dtype, + ) + if grad_is_delayed: + quantizers.dqkv.update(jnp.max(amax[:3])) + dq, dk, dv = ( + (tensor.astype(input_dtype) * grad_scale_inv).astype(input_dtype) + for tensor in (dq, dk, dv) + ) + if _scaling_mode(quantizers.dp) == "delayed": + quantizers.dp.update(amax[3:4]) + return _split_gradient_outputs(dq, dk, dv, config.qkv_layout), quantizers diff --git a/transformer_engine/jax/cpp_extensions/gemm.py b/transformer_engine/jax/cpp_extensions/gemm.py index 7cb48f5c218..f4a8d2c5eea 100644 --- a/transformer_engine/jax/cpp_extensions/gemm.py +++ b/transformer_engine/jax/cpp_extensions/gemm.py @@ -47,6 +47,7 @@ noop_quantizer_set, is_fp8_gemm_with_all_layouts_supported, apply_padding_to_scale_inv, + swizzle_mxfp8_scale, QuantizeLayout, ) from .misc import get_padded_spec, is_all_reduce_in_float32, get_min_device_compute_capability @@ -398,21 +399,6 @@ def create(forward_collective_op: CollectiveOp): noop_collective_op_set = CollectiveOpSet.create(forward_collective_op=CollectiveOp.NONE) -@partial(jax.jit, static_argnums=(1, 2)) -def swizzled_scale(scale_inv, flatten_axis, is_colwise): - "Swizzle scale_inv via JAX transpose ops" - original_shape = scale_inv.shape - shape_2d = (math.prod(original_shape[:flatten_axis]), math.prod(original_shape[flatten_axis:])) - if is_colwise: - scale_inv = jnp.transpose(scale_inv.reshape(shape_2d)) - cols, rows = shape_2d - else: - rows, cols = shape_2d - reshape = scale_inv.reshape(rows // 128, 4, 32, cols // 4, 4) - swizzled = jnp.transpose(reshape, (0, 3, 2, 1, 4)) - return swizzled.reshape(original_shape) - - def get_lhs_axis_boundary(lhs_cdims, is_transposed): """Get the axis boundary for the LHS operand.""" return max(lhs_cdims) + 1 if is_transposed else min(lhs_cdims) @@ -759,8 +745,8 @@ def impl( # Only perform JAX-based swizzle for MXFP8, NVFP4 swizzle will go though nvte kernel if scaling_mode.is_mxfp8_scaling: - lhs_scale_inv = swizzled_scale(lhs_scale_inv, lhs_flatten_axis, lhs_transposed) - rhs_scale_inv = swizzled_scale(rhs_scale_inv, rhs_flatten_axis, not rhs_transposed) + lhs_scale_inv = swizzle_mxfp8_scale(lhs_scale_inv, lhs_flatten_axis, lhs_transposed) + rhs_scale_inv = swizzle_mxfp8_scale(rhs_scale_inv, rhs_flatten_axis, not rhs_transposed) # Determine if we need to reorder the tensor so that the input/output are in the correct layout for the collective operation need_reorder = not transpose_batch_sequence and not is_outer and not collective_op.is_none diff --git a/transformer_engine/jax/csrc/extensions.h b/transformer_engine/jax/csrc/extensions.h index a7bf8d21d55..c3e80b4202b 100644 --- a/transformer_engine/jax/csrc/extensions.h +++ b/transformer_engine/jax/csrc/extensions.h @@ -23,7 +23,6 @@ #include #include #include -#include #include #include "common/common.h" @@ -152,24 +151,8 @@ XLA_FFI_DECLARE_HANDLER_SYMBOL(FusedAttnScoreModForwardHandler); XLA_FFI_DECLARE_HANDLER_SYMBOL(FusedAttnScoreModBackwardHandler); -std::tuple GetFusedAttnBackend( - const pybind11::object ¶ms); - -pybind11::tuple GetFusedAttnForwardWorkspaceSizes( - size_t input_batch, size_t bias_batch, size_t q_max_seqlen, size_t kv_max_seqlen, - size_t attn_heads, size_t num_gqa_groups, size_t bias_heads, size_t qk_head_dim, - size_t v_head_dim, float scaling_factor, float dropout_probability, NVTE_Bias_Type bias_type, - NVTE_Mask_Type mask_type, NVTE_Softmax_Type softmax_type, NVTE_QKV_Layout qkv_layout, - DType dtype, bool is_training, size_t max_segments_per_seq, int64_t window_size_left, - int64_t window_size_right, bool return_max_logit, bool bottom_right_diagonal); - -pybind11::tuple GetFusedAttnBackwardWorkspaceSizes( - size_t input_batch, size_t bias_batch, size_t q_max_seqlen, size_t kv_max_seqlen, - size_t attn_heads, size_t num_gqa_groups, size_t bias_heads, size_t qk_head_dim, - size_t v_head_dim, float scaling_factor, float dropout_probability, NVTE_Bias_Type bias_type, - NVTE_Mask_Type mask_type, NVTE_Softmax_Type softmax_type, NVTE_QKV_Layout qkv_layout, - DType dtype, bool is_training, bool deterministic, size_t max_segments_per_seq, - int64_t window_size_left, int64_t window_size_right, bool bottom_right_diagonal); +void PopulateFusedAttnRngState(void *rng_state, const void *seed, uint64_t offset, + cudaStream_t stream); // GEMM XLA_FFI_DECLARE_HANDLER_SYMBOL(GemmHandler); diff --git a/transformer_engine/jax/csrc/extensions/attention.cpp b/transformer_engine/jax/csrc/extensions/attention.cpp index 81712c8a71f..da6989bcabb 100644 --- a/transformer_engine/jax/csrc/extensions/attention.cpp +++ b/transformer_engine/jax/csrc/extensions/attention.cpp @@ -18,932 +18,32 @@ #include #include "../extensions.h" -#include "transformer_engine/fused_attn.h" -#include "transformer_engine/transformer_engine.h" +#include "attention_cache_debug.h" namespace transformer_engine { namespace jax { - -static std::tuple GetFusedAttnBackendImpl( - const FusedAttnConfigWrapper &cfg) { - const char *message = nullptr; - auto backend = nvte_get_fused_attn_backend_v2(cfg, &message); - return {backend, std::string(message)}; -} - -std::tuple GetFusedAttnBackend( - const pybind11::object ¶ms) { - const auto qkv_layout = params.attr("qkv_layout").cast(); - auto o_format = params.attr("o_format").cast(); - auto do_format = params.attr("do_format").cast(); - auto dqkv_layout = params.attr("dqkv_layout").cast(); - if (o_format == NVTE_QKV_Format::NVTE_QKV_Format_NOT_SET) { - o_format = nvte_get_q_format(qkv_layout); - } - if (do_format == NVTE_QKV_Format::NVTE_QKV_Format_NOT_SET) { - do_format = o_format; - } - if (dqkv_layout == NVTE_QKV_Layout::NVTE_QKV_Layout_NOT_SET) { - dqkv_layout = qkv_layout; - } - - FusedAttnConfigWrapper cfg; - cfg.set_is_training(params.attr("is_training").cast()) - .set_deterministic(params.attr("deterministic").cast()) - .set_cuda_graph(params.attr("cuda_graph").cast()) - .set_return_max_logit(params.attr("return_max_logit").cast()) - .set_attn_mask_type(params.attr("attn_mask_type").cast()) - .set_bias_type(params.attr("bias_type").cast()) - .set_window_size_left(params.attr("window_size_left").cast()) - .set_window_size_right(params.attr("window_size_right").cast()) - .set_bottom_right_diagonal(params.attr("bottom_right_diagonal").cast()) - .set_softmax_type(params.attr("softmax_type").cast()) - .set_scaling_mode( - get_nvte_scaling_mode(params.attr("scaling_mode").cast())) - .set_dropout(params.attr("dropout").cast()) - .set_attn_scale(params.attr("attn_scale").cast()) - .set_qkv_dtype(static_cast(params.attr("qkv_dtype").cast())) - .set_o_dtype(static_cast(params.attr("o_dtype").cast())) - .set_do_dtype(static_cast(params.attr("do_dtype").cast())) - .set_dqkv_dtype(static_cast(params.attr("dqkv_dtype").cast())) - .set_qkv_layout(qkv_layout) - .set_o_format(o_format) - .set_do_format(do_format) - .set_dqkv_layout(dqkv_layout) - .set_qkv_scale_inv_format(params.attr("qkv_scale_inv_format").cast()) - .set_do_scale_inv_format(params.attr("do_scale_inv_format").cast()) - .set_batch_size(params.attr("batch_size").cast()) - .set_num_attn_heads(params.attr("num_attn_heads").cast()) - .set_num_gqa_groups(params.attr("num_gqa_groups").cast()) - .set_head_dim_qk(params.attr("head_dim_qk").cast()) - .set_head_dim_v(params.attr("head_dim_v").cast()) - .set_max_seqlen_q(params.attr("max_seqlen_q").cast()) - .set_max_seqlen_kv(params.attr("max_seqlen_kv").cast()) - .set_num_tokens_q(params.attr("num_tokens_q").cast()) - .set_num_tokens_kv(params.attr("num_tokens_kv").cast()) - .set_bias_batch_size(params.attr("bias_batch_size").cast()) - .set_bias_num_heads(params.attr("bias_num_heads").cast()) - .set_bias_seqlen_q(params.attr("bias_seqlen_q").cast()) - .set_bias_seqlen_kv(params.attr("bias_seqlen_kv").cast()); - return GetFusedAttnBackendImpl(cfg); -} - -/* - NOTE: PrepareFusedAttnForwardAuxTensors unifies the auxiliary tensor pack logic from the fused - attention forward kernels in: - - common/fused_attn/fused_attn_f16_arbitrary_seqlen.cu lines 1270-1281 and 1348-1359 -*/ -void PrepareFusedAttnForwardAuxTensors(NVTETensorPack *tensor_pack, const size_t input_batch, - const size_t bias_batch, const size_t attn_heads, - const size_t bias_heads, const size_t q_max_seqlen, - const size_t kv_max_seqlen, DType dtype, - NVTE_Bias_Type bias_type, NVTE_Fused_Attn_Backend backend, - void *softmax_buf, void *max_logits_buf = nullptr, - void *rng_state_buf = nullptr, void *bias_buf = nullptr, - void *softmax_offset_buf = nullptr) { - // all backends need softmax but expect different shapes/dtypes - tensor_pack->size = 1; - NVTETensor &softmax_aux = tensor_pack->tensors[0]; - NVTEBasicTensor softmax_aux_data; - softmax_aux_data.data_ptr = softmax_buf; - softmax_aux_data.shape.ndim = 4; - softmax_aux_data.shape.data[0] = input_batch; - softmax_aux_data.shape.data[1] = attn_heads; - softmax_aux_data.shape.data[2] = q_max_seqlen; - softmax_aux_data.shape.data[3] = kv_max_seqlen; - softmax_aux_data.dtype = static_cast(dtype); - - // arbitrary sequence length backend needs the RNG state and a different shape/dtype softmax - if (backend == NVTE_Fused_Attn_Backend::NVTE_F16_arbitrary_seqlen) { - int size = 1; // Start after softmax. - auto next_aux_tensor = [&]() -> NVTETensor & { - NVTE_CHECK(size < NVTETensorPack::MAX_SIZE, - "Fused attention auxiliary tensor pack capacity exceeded."); - return tensor_pack->tensors[size++]; - }; - - if (max_logits_buf != nullptr) { - NVTETensor &max_aux = next_aux_tensor(); - NVTEBasicTensor max_aux_data; - max_aux_data.data_ptr = max_logits_buf; - max_aux_data.shape = {}; - max_aux_data.shape.ndim = 4; - max_aux_data.shape.data[0] = input_batch; - max_aux_data.shape.data[1] = attn_heads; - max_aux_data.shape.data[2] = q_max_seqlen; - max_aux_data.shape.data[3] = 1; - max_aux_data.dtype = static_cast(DType::kFloat32); - nvte_set_tensor_param(&max_aux, kNVTERowwiseData, &max_aux_data); - } - - NVTETensor &rng_state_aux = next_aux_tensor(); - NVTEBasicTensor rng_state_aux_data; - rng_state_aux_data.data_ptr = rng_state_buf; - rng_state_aux_data.shape = {}; - rng_state_aux_data.shape.ndim = 2; - rng_state_aux_data.dtype = static_cast(DType::kInt64); - nvte_set_tensor_param(&rng_state_aux, kNVTERowwiseData, &rng_state_aux_data); - // correct softmax shape/dtype - softmax_aux_data.shape.data[3] = 1; // {B,H,Qs,Ks} -> {B,H,Qs,1} - softmax_aux_data.dtype = static_cast(DType::kFloat32); - - // include bias if enabled - if (bias_type != NVTE_Bias_Type::NVTE_NO_BIAS && bias_type != NVTE_Bias_Type::NVTE_ALIBI) { - NVTETensor &bias_aux = next_aux_tensor(); - NVTEBasicTensor bias_aux_data; - bias_aux_data.data_ptr = bias_buf; - bias_aux_data.shape.ndim = 4; - bias_aux_data.shape.data[0] = bias_batch; - bias_aux_data.shape.data[1] = bias_heads; - bias_aux_data.shape.data[2] = q_max_seqlen; - bias_aux_data.shape.data[3] = kv_max_seqlen; - bias_aux_data.dtype = static_cast(dtype); - nvte_set_tensor_param(&bias_aux, kNVTERowwiseData, &bias_aux_data); - } - - // include softmax_offset if provided - if (softmax_offset_buf != nullptr) { - NVTETensor &softmax_offset_aux = next_aux_tensor(); - NVTEBasicTensor softmax_offset_aux_data; - softmax_offset_aux_data.data_ptr = softmax_offset_buf; - softmax_offset_aux_data.shape.ndim = 4; - softmax_offset_aux_data.shape.data[0] = 1; - softmax_offset_aux_data.shape.data[1] = attn_heads; - softmax_offset_aux_data.shape.data[2] = 1; - softmax_offset_aux_data.shape.data[3] = 1; - softmax_offset_aux_data.dtype = static_cast(DType::kFloat32); - nvte_set_tensor_param(&softmax_offset_aux, kNVTERowwiseData, &softmax_offset_aux_data); - } - - // Set final size - tensor_pack->size = size; - } - nvte_set_tensor_param(&softmax_aux, kNVTERowwiseData, &softmax_aux_data); -} - -/* - NOTE: Backward fused attention kernels accept auxiliary tensors as explicit function arguments - instead of an NVTETensorPack and nvte_fused_attn_bwd() API does all the logic for pulling the - necessary tensors out of the tensor pack for the active kernel. That means we can just dump - everything we got into the tensor pack and not worry about its sizing for the backward pass. - - TODO(Alp): Refactor the nvte_fused_attn_fwd() to work like nvte_fused_attn_bwd()? -*/ -void PrepareFusedAttnBackwardAuxTensors(NVTETensorPack *tensor_pack, const size_t input_batch, - const size_t bias_batch, const size_t attn_heads, - const size_t bias_heads, const size_t q_max_seqlen, - const size_t kv_max_seqlen, DType dtype, - NVTE_Fused_Attn_Backend backend, void *softmax_buf, - void *rng_state_buf, void *bias_buf, - void *softmax_offset_buf = nullptr) { - // Backward calls put everything into the tensor pack for every backend - // so we set dummy bias_type and backend choices here to follow the correct code path - auto dummy_bias_type = NVTE_Bias_Type::NVTE_POST_SCALE_BIAS; - auto dummy_backend = NVTE_Fused_Attn_Backend::NVTE_F16_arbitrary_seqlen; - PrepareFusedAttnForwardAuxTensors(tensor_pack, input_batch, bias_batch, attn_heads, bias_heads, - q_max_seqlen, kv_max_seqlen, dtype, dummy_bias_type, - dummy_backend, softmax_buf, nullptr, rng_state_buf, bias_buf, - softmax_offset_buf); -} - -pybind11::tuple GetFusedAttnForwardWorkspaceSizes( - size_t input_batch, size_t bias_batch, size_t q_max_seqlen, size_t kv_max_seqlen, - size_t attn_heads, size_t num_gqa_groups, size_t bias_heads, size_t qk_head_dim, - size_t v_head_dim, float scaling_factor, float dropout_probability, NVTE_Bias_Type bias_type, - NVTE_Mask_Type mask_type, NVTE_Softmax_Type softmax_type, NVTE_QKV_Layout qkv_layout, - DType dtype, bool is_training, size_t max_segments_per_seq, int64_t window_size_left, - int64_t window_size_right, bool return_max_logit, bool bottom_right_diagonal) { - auto is_ragged = nvte_get_qkv_format(qkv_layout) == NVTE_QKV_Format::NVTE_THD; - auto q_shape = is_ragged - ? std::vector{input_batch * q_max_seqlen, attn_heads, qk_head_dim} - : std::vector{input_batch, q_max_seqlen, attn_heads, qk_head_dim}; - auto q_tensor = TensorWrapper(nullptr, q_shape, dtype); - auto k_shape = is_ragged - ? std::vector{input_batch * kv_max_seqlen, num_gqa_groups, qk_head_dim} - : std::vector{input_batch, kv_max_seqlen, num_gqa_groups, qk_head_dim}; - auto k_tensor = TensorWrapper(nullptr, k_shape, dtype); - auto v_shape = is_ragged - ? std::vector{input_batch * kv_max_seqlen, num_gqa_groups, v_head_dim} - : std::vector{input_batch, kv_max_seqlen, num_gqa_groups, v_head_dim}; - auto v_tensor = TensorWrapper(nullptr, v_shape, dtype); - auto o_shape = is_ragged ? std::vector{input_batch * q_max_seqlen, attn_heads, v_head_dim} - : std::vector{input_batch, q_max_seqlen, attn_heads, v_head_dim}; - - auto bias_shape = std::vector{bias_batch, bias_heads, q_max_seqlen, kv_max_seqlen}; - auto bias_tensor = TensorWrapper(nullptr, bias_shape, dtype); - - // F16 doesn't use this tensor - auto s_tensor = TensorWrapper(nullptr, std::vector{1}, dtype); - auto o_tensor = TensorWrapper(nullptr, o_shape, dtype); - - auto dummy_rng_state_tensor = TensorWrapper(nullptr, std::vector{2}, DType::kInt64); - auto dummy_page_table_tensor = TensorWrapper(nullptr, std::vector{1}, DType::kInt32); - auto dummy_softmax_offset_tensor = - TensorWrapper(nullptr, std::vector{1}, DType::kFloat32); - - NVTETensorPack aux_output_tensors; - nvte_tensor_pack_create(&aux_output_tensors); - - TensorWrapper query_workspace_tensor; - // It is a WAR to pre-create all possible cuDNN graph at the JIT compile time - size_t max_num_segments = is_ragged ? input_batch * max_segments_per_seq : input_batch; - size_t min_num_segments = input_batch; - auto cudnn_runtime_version = cudnnGetVersion(); - if (is_ragged && cudnn_runtime_version >= 90300) { - // For cuDNN < 9.3.0, it requires to run all possible seqlens to address act_seqlen = 0 - min_num_segments = input_batch * max_segments_per_seq; - } - for (auto num_segments = min_num_segments; num_segments <= max_num_segments; ++num_segments) { - // the last one is the largest which will be the returned workspace size - auto q_cu_seqlens_tensor = - TensorWrapper(nullptr, std::vector{num_segments + 1}, DType::kInt32); - auto kv_cu_seqlens_tensor = - TensorWrapper(nullptr, std::vector{num_segments + 1}, DType::kInt32); - auto ragged_offset_tensor = - TensorWrapper(nullptr, std::vector{num_segments + 1}, DType::kInt32); - FusedAttnFwdParamsWrapper params; - params.set_Q(q_tensor.data()) - .set_K(k_tensor.data()) - .set_V(v_tensor.data()) - .set_Bias(bias_tensor.data()) - .set_SoftmaxOffset(dummy_softmax_offset_tensor.data()) - .set_S(s_tensor.data()) - .set_O(o_tensor.data()) - .set_Aux_CTX_Tensors(&aux_output_tensors) - .set_cu_seqlens_q(q_cu_seqlens_tensor.data()) - .set_cu_seqlens_kv(kv_cu_seqlens_tensor.data()) - .set_cu_seqlens_q_padded(ragged_offset_tensor.data()) - .set_cu_seqlens_kv_padded(ragged_offset_tensor.data()) - .set_page_table_k(dummy_page_table_tensor.data()) - .set_page_table_v(dummy_page_table_tensor.data()) - .set_rng_state(dummy_rng_state_tensor.data()) - .set_is_training(is_training) - .set_cuda_graph(false) - .set_return_max_logit(return_max_logit) - .set_attn_mask_type(mask_type) - .set_bias_type(bias_type) - .set_window_size_left(window_size_left) - .set_window_size_right(window_size_right) - .set_bottom_right_diagonal(bottom_right_diagonal) - .set_softmax_type(softmax_type) - .set_dropout(dropout_probability) - .set_attn_scale(scaling_factor) - .set_qkv_layout(qkv_layout) - .set_o_format(nvte_get_q_format(qkv_layout)) - .set_qkv_scale_inv_format(NVTE_QKV_Format_NOT_SET) - .set_max_seqlen_q(q_max_seqlen) - .set_max_seqlen_kv(kv_max_seqlen) - .set_workspace(query_workspace_tensor.data()) - .set_stream(nullptr); - nvte_fused_attn_fwd_v2(params); - } - - nvte_tensor_pack_destroy(&aux_output_tensors); - - auto workspace_shape = MakeShapeVector(query_workspace_tensor.shape()); - return pybind11::make_tuple(workspace_shape, query_workspace_tensor.dtype()); -} - -#define FUSED_ATTN_IMPL_COMMON_BLOCK \ - auto is_ragged = nvte_get_qkv_format(qkv_layout) == NVTE_QKV_Format::NVTE_THD; \ - auto bias_shape = std::vector{bias_batch, bias_heads, q_max_seqlen, kv_max_seqlen}; \ - const bool has_bias_tensor = \ - bias_type != NVTE_Bias_Type::NVTE_NO_BIAS && bias_type != NVTE_Bias_Type::NVTE_ALIBI; \ - const size_t bias_seqlen_q = has_bias_tensor ? q_max_seqlen : 0; \ - const size_t bias_seqlen_kv = has_bias_tensor ? kv_max_seqlen : 0; \ - size_t num_segments = input_batch; \ - if (is_ragged) { \ - auto cudnn_runtime_version = cudnnGetVersion(); \ - if (cudnn_runtime_version >= 90300) { \ - num_segments = input_batch * max_segments_per_seq; \ - } else { \ - size_t runtime_num_segments_q = nvte_get_runtime_num_segments( \ - q_cu_seqlens, workspace, input_batch * q_max_seqlen, stream); \ - size_t runtime_num_segments_kv = nvte_get_runtime_num_segments( \ - kv_cu_seqlens, workspace, input_batch * kv_max_seqlen, stream); \ - NVTE_CHECK(runtime_num_segments_q == runtime_num_segments_kv); \ - NVTE_CHECK(runtime_num_segments_q <= input_batch * max_segments_per_seq); \ - num_segments = runtime_num_segments_q; \ - } \ - } \ - std::vector seq_shape{num_segments + 1}; \ - auto q_cu_seqlens_tensor = TensorWrapper(q_cu_seqlens, seq_shape, DType::kInt32); \ - auto kv_cu_seqlens_tensor = TensorWrapper(kv_cu_seqlens, seq_shape, DType::kInt32); \ - auto q_seq_offsets_tensor = TensorWrapper(q_seq_offsets, seq_shape, DType::kInt32); \ - auto k_seq_offsets_tensor = TensorWrapper(k_seq_offsets, seq_shape, DType::kInt32); \ - auto workspace_tensor = \ - TensorWrapper(workspace, std::vector{wkspace_size}, wkspace_dtype); \ - auto layout_group = nvte_get_qkv_layout_group(qkv_layout); \ - FusedAttnConfigWrapper cfg; \ - cfg.set_is_training(is_training) \ - .set_deterministic(deterministic) \ - .set_cuda_graph(false) \ - .set_return_max_logit(false) \ - .set_attn_mask_type(mask_type) \ - .set_bias_type(bias_type) \ - .set_window_size_left(window_size_left) \ - .set_window_size_right(window_size_right) \ - .set_bottom_right_diagonal(bottom_right_diagonal) \ - .set_softmax_type(softmax_type) \ - .set_scaling_mode(get_nvte_scaling_mode(JAXX_Scaling_Mode::NO_SCALING)) \ - .set_dropout(dropout_probability) \ - .set_attn_scale(scaling_factor) \ - .set_qkv_dtype(static_cast(dtype)) \ - .set_o_dtype(static_cast(dtype)) \ - .set_do_dtype(static_cast(dtype)) \ - .set_dqkv_dtype(static_cast(dtype)) \ - .set_qkv_layout(qkv_layout) \ - .set_o_format(nvte_get_q_format(qkv_layout)) \ - .set_do_format(nvte_get_q_format(qkv_layout)) \ - .set_dqkv_layout(qkv_layout) \ - .set_qkv_scale_inv_format(NVTE_QKV_Format::NVTE_QKV_Format_NOT_SET) \ - .set_do_scale_inv_format(NVTE_QKV_Format::NVTE_QKV_Format_NOT_SET) \ - .set_batch_size(num_segments) \ - .set_num_attn_heads(attn_heads) \ - .set_num_gqa_groups(num_gqa_groups) \ - .set_head_dim_qk(qk_head_dim) \ - .set_head_dim_v(v_head_dim) \ - .set_max_seqlen_q(q_max_seqlen) \ - .set_max_seqlen_kv(kv_max_seqlen) \ - .set_num_tokens_q(is_ragged ? input_batch *q_max_seqlen : 0) \ - .set_num_tokens_kv(is_ragged ? input_batch *kv_max_seqlen : 0) \ - .set_bias_batch_size(bias_batch) \ - .set_bias_num_heads(bias_heads) \ - .set_bias_seqlen_q(bias_seqlen_q) \ - .set_bias_seqlen_kv(bias_seqlen_kv); - -static void FusedAttnForwardImpl( - cudaStream_t stream, void *q, void *k, void *v, void *bias, void *softmax_offset, void *seed, - void *q_cu_seqlens, void *kv_cu_seqlens, void *q_seq_offsets, void *k_seq_offsets, void *output, - void *softmax_aux, void *max_tensor, void *rng_state, void *workspace, size_t input_batch, - size_t bias_batch, size_t q_max_seqlen, size_t kv_max_seqlen, size_t attn_heads, - size_t num_gqa_groups, size_t bias_heads, size_t qk_head_dim, size_t v_head_dim, - size_t max_segments_per_seq, size_t wkspace_size, float scaling_factor, - float dropout_probability, NVTE_Bias_Type bias_type, NVTE_Mask_Type mask_type, - NVTE_Softmax_Type softmax_type, NVTE_QKV_Layout qkv_layout, DType dtype, DType wkspace_dtype, - bool is_training, bool return_max_logit, bool deterministic, int64_t window_size_left, - int64_t window_size_right, bool bottom_right_diagonal) { - FUSED_ATTN_IMPL_COMMON_BLOCK; - - /* Input tensors */ - auto bias_tensor = TensorWrapper(bias, bias_shape, dtype); - auto softmax_offset_tensor = - TensorWrapper(softmax_offset, std::vector{1, attn_heads, 1, 1}, DType::kFloat32); - - if (is_ragged) { - auto output_size = input_batch * q_max_seqlen * attn_heads * v_head_dim; - cudaMemsetAsync(output, 0, output_size * typeToSize(dtype), stream); - - // Memset to 0xF0 for filling large negative numbers - auto softmax_aux_size = input_batch * q_max_seqlen * attn_heads; - cudaMemsetAsync(softmax_aux, 0xF0, softmax_aux_size * sizeof(float), stream); - if (return_max_logit) { - cudaMemsetAsync(max_tensor, 0xF0, softmax_aux_size * sizeof(float), stream); - } - } - - /* Output tensors */ - auto s_tensor = TensorWrapper(nullptr, std::vector{1}, dtype); // not used in F16 - auto o_shape = is_ragged ? std::vector{input_batch * q_max_seqlen, attn_heads, v_head_dim} - : std::vector{input_batch, q_max_seqlen, attn_heads, v_head_dim}; - auto o_tensor = TensorWrapper(output, o_shape, dtype); - - /* Prepare RNG state */ - auto rng_state_tensor = TensorWrapper(rng_state, std::vector{2}, DType::kInt64); - - cfg.set_return_max_logit(return_max_logit); - auto [backend, fwd_msg] = GetFusedAttnBackendImpl(cfg); - NVTE_CHECK(backend != NVTE_Fused_Attn_Backend::NVTE_No_Backend, - "Fused attention is not supported for this configuration: ", fwd_msg); - nvte_populate_rng_state_async(rng_state, seed, q_max_seqlen, kv_max_seqlen, backend, stream); - - /* Auxiliary tensors (to be propagated to the backward pass later) */ - NVTETensorPack aux_output_tensors; - nvte_tensor_pack_create(&aux_output_tensors); - PrepareFusedAttnForwardAuxTensors(&aux_output_tensors, input_batch, bias_batch, attn_heads, - bias_heads, q_max_seqlen, kv_max_seqlen, dtype, bias_type, - backend, softmax_aux, return_max_logit ? max_tensor : nullptr, - rng_state, bias, softmax_offset); - - /* Call the underlying NVTE API */ - auto dummy_page_table_tensor = TensorWrapper(nullptr, std::vector{1}, DType::kInt32); - - // Prepare Q, K, V pointers and shapes based on layout - // Python passes dummy tensors for unused slots, so we extract from the actual packed data - void *q_ptr = q; - void *k_ptr = k; - void *v_ptr = v; - auto q_shape = is_ragged - ? std::vector{input_batch * q_max_seqlen, attn_heads, qk_head_dim} - : std::vector{input_batch, q_max_seqlen, attn_heads, qk_head_dim}; - auto k_shape = is_ragged - ? std::vector{input_batch * kv_max_seqlen, num_gqa_groups, qk_head_dim} - : std::vector{input_batch, kv_max_seqlen, num_gqa_groups, qk_head_dim}; - auto v_shape = is_ragged - ? std::vector{input_batch * kv_max_seqlen, num_gqa_groups, v_head_dim} - : std::vector{input_batch, kv_max_seqlen, num_gqa_groups, v_head_dim}; - - if (layout_group == NVTE_QKV_Layout_Group::NVTE_3HD) { - // QKV packed in q: [batch*seqlen, 3, heads, dim] - // Python passes: q=packed_qkv, k=dummy, v=dummy - // Extract K and V pointers from the packed q data - NVTE_CHECK(q_max_seqlen == kv_max_seqlen, "q_max_seqlen must equal kv_max_seqlen"); - NVTE_CHECK(qk_head_dim == v_head_dim, - "For QKV packed layout, qk_head_dim must equal v_head_dim"); - size_t stride = (typeToSize(dtype) * attn_heads * qk_head_dim); - q_ptr = q; - k_ptr = static_cast(static_cast(q) + stride); - v_ptr = static_cast(static_cast(q) + 2 * stride); - // For packed QKV, all have same shape since they're views into the same packed tensor - k_shape = q_shape; - v_shape = q_shape; - } else if (layout_group == NVTE_QKV_Layout_Group::NVTE_HD_2HD) { - // Q separate, KV packed in k: [batch*seqlen, 2, num_gqa_groups, dim] - // Python passes: q=query, k=packed_kv, v=dummy - // Extract V pointer from the packed k data - NVTE_CHECK(qk_head_dim == v_head_dim, - "For KV packed layout, qk_head_dim must equal v_head_dim"); - size_t stride = (typeToSize(dtype) * num_gqa_groups * qk_head_dim); - q_ptr = q; - k_ptr = k; - v_ptr = static_cast(static_cast(k) + stride); - // V has same shape as K since they're packed together - v_shape = k_shape; - } - // else NVTE_HD_HD_HD: pointers and shapes already correct - - auto q_tensor = TensorWrapper(q_ptr, q_shape, dtype); - auto k_tensor = TensorWrapper(k_ptr, k_shape, dtype); - auto v_tensor = TensorWrapper(v_ptr, v_shape, dtype); - - FusedAttnFwdParamsWrapper params; - params.set_Q(q_tensor.data()) - .set_K(k_tensor.data()) - .set_V(v_tensor.data()) - .set_Bias(bias_tensor.data()) - .set_SoftmaxOffset(softmax_offset_tensor.data()) - .set_S(s_tensor.data()) - .set_O(o_tensor.data()) - .set_Aux_CTX_Tensors(&aux_output_tensors) - .set_cu_seqlens_q(q_cu_seqlens_tensor.data()) - .set_cu_seqlens_kv(kv_cu_seqlens_tensor.data()) - .set_cu_seqlens_q_padded(q_seq_offsets_tensor.data()) - .set_cu_seqlens_kv_padded(k_seq_offsets_tensor.data()) - .set_page_table_k(dummy_page_table_tensor.data()) - .set_page_table_v(dummy_page_table_tensor.data()) - .set_rng_state(rng_state_tensor.data()) - .set_is_training(is_training) - .set_cuda_graph(false) - .set_return_max_logit(return_max_logit) - .set_attn_mask_type(mask_type) - .set_bias_type(bias_type) - .set_window_size_left(window_size_left) - .set_window_size_right(window_size_right) - .set_bottom_right_diagonal(bottom_right_diagonal) - .set_softmax_type(softmax_type) - .set_dropout(dropout_probability) - .set_attn_scale(scaling_factor) - .set_qkv_layout(qkv_layout) - .set_o_format(nvte_get_q_format(qkv_layout)) - .set_qkv_scale_inv_format(NVTE_QKV_Format_NOT_SET) - .set_max_seqlen_q(q_max_seqlen) - .set_max_seqlen_kv(kv_max_seqlen) - .set_workspace(workspace_tensor.data()) - .set_stream(stream); - nvte_fused_attn_fwd_v2(params); - - nvte_tensor_pack_destroy(&aux_output_tensors); -} - -#define FUSED_ATTN_FFI_GET_ATTRS \ - size_t input_batch = get_attr_value(attrs, "input_batch"); \ - size_t bias_batch = get_attr_value(attrs, "bias_batch"); \ - size_t q_max_seqlen = get_attr_value(attrs, "q_max_seqlen"); \ - size_t kv_max_seqlen = get_attr_value(attrs, "kv_max_seqlen"); \ - size_t attn_heads = get_attr_value(attrs, "attn_heads"); \ - size_t num_gqa_groups = get_attr_value(attrs, "num_gqa_groups"); \ - size_t bias_heads = get_attr_value(attrs, "bias_heads"); \ - size_t qk_head_dim = get_attr_value(attrs, "qk_head_dim"); \ - size_t v_head_dim = get_attr_value(attrs, "v_head_dim"); \ - size_t max_segments_per_seq = get_attr_value(attrs, "max_segments_per_seq"); \ - auto window_size_left = get_attr_value(attrs, "window_size_left"); \ - auto window_size_right = get_attr_value(attrs, "window_size_right"); \ - bool bottom_right_diagonal = get_attr_value(attrs, "bottom_right_diagonal"); \ - float scaling_factor = get_attr_value(attrs, "scaling_factor"); \ - float dropout_probability = get_attr_value(attrs, "dropout_probability"); \ - NVTE_Bias_Type bias_type = \ - static_cast(get_attr_value(attrs, "bias_type")); \ - NVTE_Mask_Type mask_type = \ - static_cast(get_attr_value(attrs, "mask_type")); \ - NVTE_Softmax_Type softmax_type = \ - static_cast(get_attr_value_or_default( \ - attrs, "softmax_type", static_cast(NVTE_Softmax_Type::NVTE_VANILLA_SOFTMAX))); \ - NVTE_QKV_Layout qkv_layout = \ - static_cast(get_attr_value(attrs, "qkv_layout")); \ - bool is_training = get_attr_value(attrs, "is_training"); \ - bool deterministic = get_attr_value(attrs, "deterministic"); \ - auto is_ragged = nvte_get_qkv_format(qkv_layout) == NVTE_QKV_Format::NVTE_THD; \ - size_t wkspace_size = product(workspace_buf->dimensions()); \ - DType dtype = convert_ffi_datatype_to_te_dtype(q_buf.element_type()); \ - DType wkspace_dtype = convert_ffi_datatype_to_te_dtype(workspace_buf->element_type()); - -Error_Type FusedAttnForwardFFI(cudaStream_t stream, Buffer_Type q_buf, Buffer_Type k_buf, - Buffer_Type v_buf, Buffer_Type bias_buf, - Buffer_Type softmax_offset_buf, Buffer_Type seed_buf, - Buffer_Type q_cu_seqlens_buf, Buffer_Type kv_cu_seqlens_buf, - Buffer_Type q_seq_offsets_buf, Buffer_Type k_seq_offsets_buf, - Variadic_Buffer_Type _unused_args, Result_Type output_buf, - Result_Type softmax_aux_buf, Result_Type max_tensor_buf, - Result_Type rng_state_buf, Result_Type workspace_buf, - Dictionary attrs) { - FUSED_ATTN_FFI_GET_ATTRS; - bool return_max_logit = get_attr_value_or_default(attrs, "return_max_logit", false); - - FusedAttnForwardImpl( - stream, q_buf.untyped_data(), k_buf.untyped_data(), v_buf.untyped_data(), - bias_buf.untyped_data(), softmax_offset_buf.untyped_data(), seed_buf.untyped_data(), - q_cu_seqlens_buf.untyped_data(), kv_cu_seqlens_buf.untyped_data(), - is_ragged ? q_seq_offsets_buf.untyped_data() : nullptr, - is_ragged ? k_seq_offsets_buf.untyped_data() : nullptr, output_buf->untyped_data(), - softmax_aux_buf->untyped_data(), max_tensor_buf->untyped_data(), - rng_state_buf->untyped_data(), workspace_buf->untyped_data(), input_batch, bias_batch, - q_max_seqlen, kv_max_seqlen, attn_heads, num_gqa_groups, bias_heads, qk_head_dim, v_head_dim, - max_segments_per_seq, wkspace_size, scaling_factor, dropout_probability, bias_type, mask_type, - softmax_type, qkv_layout, dtype, wkspace_dtype, is_training, return_max_logit, deterministic, - window_size_left, window_size_right, bottom_right_diagonal); - return ffi_with_cuda_error_check(); -} - -XLA_FFI_DEFINE_HANDLER_SYMBOL(FusedAttnForwardHandler, FusedAttnForwardFFI, - FFI::Bind() - .Ctx() // stream - .Arg() // q - .Arg() // k - .Arg() // v - .Arg() // bias - .Arg() // softmax_offset - .Arg() // seed_buf - .Arg() // q_cu_seqlens - .Arg() // kv_cu_seqlens - .Arg() // q_seq_offsets - .Arg() // k_seq_offsets - .RemainingArgs() // _cp_aux_args unused - .Ret() // output - .Ret() // softmax_aux - .Ret() // max_tensor - .Ret() // rng_state - .Ret() // workspace - .Attrs(), - FFI_CudaGraph_Traits); - -pybind11::tuple GetFusedAttnBackwardWorkspaceSizes( - size_t input_batch, size_t bias_batch, size_t q_max_seqlen, size_t kv_max_seqlen, - size_t attn_heads, size_t num_gqa_groups, size_t bias_heads, size_t qk_head_dim, - size_t v_head_dim, float scaling_factor, float dropout_probability, NVTE_Bias_Type bias_type, - NVTE_Mask_Type mask_type, NVTE_Softmax_Type softmax_type, NVTE_QKV_Layout qkv_layout, - DType dtype, bool is_training, bool deterministic, size_t max_segments_per_seq, - int64_t window_size_left, int64_t window_size_right, bool bottom_right_diagonal) { - auto is_ragged = nvte_get_qkv_format(qkv_layout) == NVTE_QKV_Format::NVTE_THD; - auto q_shape = is_ragged - ? std::vector{input_batch * q_max_seqlen, attn_heads, qk_head_dim} - : std::vector{input_batch, q_max_seqlen, attn_heads, qk_head_dim}; - auto q_tensor = TensorWrapper(nullptr, q_shape, dtype); - auto dq_tensor = TensorWrapper(nullptr, q_shape, dtype); - auto k_shape = is_ragged - ? std::vector{input_batch * kv_max_seqlen, num_gqa_groups, qk_head_dim} - : std::vector{input_batch, kv_max_seqlen, num_gqa_groups, qk_head_dim}; - auto k_tensor = TensorWrapper(nullptr, k_shape, dtype); - auto dk_tensor = TensorWrapper(nullptr, k_shape, dtype); - auto v_shape = is_ragged - ? std::vector{input_batch * kv_max_seqlen, num_gqa_groups, v_head_dim} - : std::vector{input_batch, kv_max_seqlen, num_gqa_groups, v_head_dim}; - auto v_tensor = TensorWrapper(nullptr, v_shape, dtype); - auto dv_tensor = TensorWrapper(nullptr, v_shape, dtype); - - auto output_shape = is_ragged - ? std::vector{input_batch * q_max_seqlen, attn_heads, v_head_dim} - : std::vector{input_batch, q_max_seqlen, attn_heads, v_head_dim}; - auto doutput_tensor = TensorWrapper(nullptr, output_shape, dtype); - auto output_tensor = TensorWrapper(nullptr, output_shape, dtype); - - // F16 doesn't use this tensor - auto s_tensor = TensorWrapper(nullptr, std::vector{1}, dtype); - - auto bias_shape = std::vector{bias_batch, bias_heads, q_max_seqlen, kv_max_seqlen}; - auto dbias_tensor = TensorWrapper(nullptr, bias_shape, dtype); - - NVTETensorPack aux_input_tensors; - nvte_tensor_pack_create(&aux_input_tensors); - - TensorWrapper query_workspace_tensor; - - // It is a WAR to pre-create all possible cuDNN graph at the JIT compile time - size_t max_num_segments = is_ragged ? input_batch * max_segments_per_seq : input_batch; - size_t min_num_segments = input_batch; - auto cudnn_runtime_version = cudnnGetVersion(); - if (is_ragged && cudnn_runtime_version >= 90300) { - // For cuDNN < 9.3.0, it requires to run all possible seqlens to address act_seqlen = 0 - min_num_segments = input_batch * max_segments_per_seq; - } - - TensorWrapper dummy_d_softmax_offset_tensor; - if (softmax_type == NVTE_Softmax_Type::NVTE_OFF_BY_ONE_SOFTMAX || - softmax_type == NVTE_Softmax_Type::NVTE_LEARNABLE_SOFTMAX) { - dummy_d_softmax_offset_tensor = - TensorWrapper(nullptr, std::vector{1, attn_heads, 1, 1}, DType::kFloat32); - } - - for (auto num_segments = min_num_segments; num_segments <= max_num_segments; ++num_segments) { - // the last one is the largest which will be the returned workspace size - auto q_cu_seqlens_tensor = - TensorWrapper(nullptr, std::vector{num_segments + 1}, DType::kInt32); - auto kv_cu_seqlens_tensor = - TensorWrapper(nullptr, std::vector{num_segments + 1}, DType::kInt32); - auto dummy_ragged_offset_tensor = - TensorWrapper(nullptr, std::vector{num_segments + 1}, DType::kInt32); - - FusedAttnBwdParamsWrapper params; - params.set_Q(q_tensor.data()) - .set_K(k_tensor.data()) - .set_V(v_tensor.data()) - .set_O(output_tensor.data()) - .set_dO(doutput_tensor.data()) - .set_S(s_tensor.data()) // not used for F16 - .set_dP(s_tensor.data()) // not used for F16 - .set_Aux_CTX_Tensors(&aux_input_tensors) - .set_dQ(dq_tensor.data()) - .set_dK(dk_tensor.data()) - .set_dV(dv_tensor.data()) - .set_dBias(dbias_tensor.data()) - .set_dSoftmaxOffset(dummy_d_softmax_offset_tensor.data()) - .set_cu_seqlens_q(q_cu_seqlens_tensor.data()) - .set_cu_seqlens_kv(kv_cu_seqlens_tensor.data()) - .set_cu_seqlens_q_padded(dummy_ragged_offset_tensor.data()) - .set_cu_seqlens_kv_padded(dummy_ragged_offset_tensor.data()) - .set_deterministic(deterministic) - .set_cuda_graph(false) - .set_attn_mask_type(mask_type) - .set_bias_type(bias_type) - .set_window_size_left(window_size_left) - .set_window_size_right(window_size_right) - .set_bottom_right_diagonal(bottom_right_diagonal) - .set_softmax_type(softmax_type) - .set_dropout(dropout_probability) - .set_attn_scale(scaling_factor) - .set_qkv_layout(qkv_layout) - .set_o_format(nvte_get_q_format(qkv_layout)) - .set_do_format(nvte_get_q_format(qkv_layout)) - .set_dqkv_layout(qkv_layout) - .set_qkv_scale_inv_format(NVTE_QKV_Format_NOT_SET) - .set_do_scale_inv_format(NVTE_QKV_Format_NOT_SET) - .set_max_seqlen_q(q_max_seqlen) - .set_max_seqlen_kv(kv_max_seqlen) - .set_workspace(query_workspace_tensor.data()) - .set_stream(nullptr); - nvte_fused_attn_bwd_v2(params); - } - - nvte_tensor_pack_destroy(&aux_input_tensors); - - auto work_shape = MakeShapeVector(query_workspace_tensor.shape()); - return pybind11::make_tuple(work_shape, query_workspace_tensor.dtype()); -} - -static void FusedAttnBackwardImpl( - cudaStream_t stream, void *q, void *k, void *v, void *bias, void *softmax_offset, - void *softmax_aux, void *rng_state, void *output, void *doutput, void *q_cu_seqlens, - void *kv_cu_seqlens, void *q_seq_offsets, void *k_seq_offsets, void *dq, void *dk, void *dv, - void *dbias, void *dsoftmax_offset, void *workspace, size_t input_batch, size_t bias_batch, - size_t q_max_seqlen, size_t kv_max_seqlen, size_t attn_heads, size_t num_gqa_groups, - size_t bias_heads, size_t qk_head_dim, size_t v_head_dim, size_t max_segments_per_seq, - size_t wkspace_size, float scaling_factor, float dropout_probability, NVTE_Bias_Type bias_type, - NVTE_Mask_Type mask_type, NVTE_Softmax_Type softmax_type, NVTE_QKV_Layout qkv_layout, - DType dtype, DType wkspace_dtype, bool is_training, bool deterministic, - int64_t window_size_left, int64_t window_size_right, bool bottom_right_diagonal) { - FUSED_ATTN_IMPL_COMMON_BLOCK; - - /* Input tensors */ - auto output_shape = is_ragged - ? std::vector{input_batch * q_max_seqlen, attn_heads, v_head_dim} - : std::vector{input_batch, q_max_seqlen, attn_heads, v_head_dim}; - auto output_tensor = TensorWrapper(output, output_shape, dtype); - auto doutput_tensor = TensorWrapper(doutput, output_shape, dtype); - - /* Output tensors */ - auto s_tensor = TensorWrapper(nullptr, std::vector{1}, dtype); // not used in F16 - auto dbias_tensor = TensorWrapper(dbias, bias_shape, dtype); - - TensorWrapper dsoftmax_offset_tensor; - if (softmax_type == NVTE_Softmax_Type::NVTE_OFF_BY_ONE_SOFTMAX || - softmax_type == NVTE_Softmax_Type::NVTE_LEARNABLE_SOFTMAX) { - dsoftmax_offset_tensor = - TensorWrapper(dsoftmax_offset, std::vector{1, attn_heads, 1, 1}, DType::kFloat32); - } - - /* Auxiliary tensors (propagated from the forward pass) */ - NVTETensorPack aux_input_tensors; - nvte_tensor_pack_create(&aux_input_tensors); - auto [backend, bwd_msg] = GetFusedAttnBackendImpl(cfg); - NVTE_CHECK(backend != NVTE_Fused_Attn_Backend::NVTE_No_Backend, - "Fused attention is not supported for this configuration: ", bwd_msg); - PrepareFusedAttnBackwardAuxTensors(&aux_input_tensors, input_batch, bias_batch, attn_heads, - bias_heads, q_max_seqlen, kv_max_seqlen, dtype, backend, - softmax_aux, rng_state, bias, softmax_offset); - - /* Call the underlying NVTE API */ - // Prepare Q, K, V pointers and shapes based on layout - void *q_ptr = q; - void *k_ptr = k; - void *v_ptr = v; - void *dq_ptr = dq; - void *dk_ptr = dk; - void *dv_ptr = dv; - auto q_shape = is_ragged - ? std::vector{input_batch * q_max_seqlen, attn_heads, qk_head_dim} - : std::vector{input_batch, q_max_seqlen, attn_heads, qk_head_dim}; - auto k_shape = is_ragged - ? std::vector{input_batch * kv_max_seqlen, num_gqa_groups, qk_head_dim} - : std::vector{input_batch, kv_max_seqlen, num_gqa_groups, qk_head_dim}; - auto v_shape = is_ragged - ? std::vector{input_batch * kv_max_seqlen, num_gqa_groups, v_head_dim} - : std::vector{input_batch, kv_max_seqlen, num_gqa_groups, v_head_dim}; - - if (layout_group == NVTE_QKV_Layout_Group::NVTE_3HD) { - // QKV packed in q: [batch*seqlen, 3, heads, dim] - NVTE_CHECK(q_max_seqlen == kv_max_seqlen, "q_max_seqlen must equal kv_max_seqlen"); - NVTE_CHECK(qk_head_dim == v_head_dim, - "For QKV packed layout, qk_head_dim must equal v_head_dim"); - size_t stride = (typeToSize(dtype) * attn_heads * qk_head_dim); - q_ptr = q; - k_ptr = static_cast(static_cast(q) + stride); - v_ptr = static_cast(static_cast(q) + 2 * stride); - dq_ptr = dq; - dk_ptr = static_cast(static_cast(dq) + stride); - dv_ptr = static_cast(static_cast(dq) + 2 * stride); - k_shape = q_shape; - v_shape = q_shape; - } else if (layout_group == NVTE_QKV_Layout_Group::NVTE_HD_2HD) { - // Q separate, KV packed in k: [batch*seqlen, 2, num_gqa_groups, dim] - NVTE_CHECK(qk_head_dim == v_head_dim, - "For KV packed layout, qk_head_dim must equal v_head_dim"); - size_t stride = (typeToSize(dtype) * num_gqa_groups * qk_head_dim); - q_ptr = q; - k_ptr = k; - v_ptr = static_cast(static_cast(k) + stride); - dq_ptr = dq; - dk_ptr = dk; - dv_ptr = static_cast(static_cast(dk) + stride); - // V has same shape as K since they're packed together - v_shape = k_shape; - } - - auto q_tensor = TensorWrapper(q_ptr, q_shape, dtype); - auto k_tensor = TensorWrapper(k_ptr, k_shape, dtype); - auto v_tensor = TensorWrapper(v_ptr, v_shape, dtype); - auto dq_tensor = TensorWrapper(dq_ptr, q_shape, dtype); - auto dk_tensor = TensorWrapper(dk_ptr, k_shape, dtype); - auto dv_tensor = TensorWrapper(dv_ptr, v_shape, dtype); - - if (is_ragged) { - size_t dtype_size = typeToSize(dtype); - if (layout_group == NVTE_QKV_Layout_Group::NVTE_3HD) { - // For packed QKV, dq contains all gradients (dq, dk, dv) - clear all at once - cudaMemsetAsync(dq, 0, 3 * transformer_engine::jax::product(q_shape) * dtype_size, stream); - } else if (layout_group == NVTE_QKV_Layout_Group::NVTE_HD_2HD) { - // Clear dq - cudaMemsetAsync(dq, 0, transformer_engine::jax::product(q_shape) * dtype_size, stream); - // For packed KV, dk contains both dk and dv - clear all at once - cudaMemsetAsync(dk, 0, 2 * transformer_engine::jax::product(k_shape) * dtype_size, stream); - } else { - // All separate - clear each individually - cudaMemsetAsync(dq, 0, transformer_engine::jax::product(q_shape) * dtype_size, stream); - cudaMemsetAsync(dk, 0, transformer_engine::jax::product(k_shape) * dtype_size, stream); - cudaMemsetAsync(dv, 0, transformer_engine::jax::product(v_shape) * dtype_size, stream); - } - } - - FusedAttnBwdParamsWrapper params; - params.set_Q(q_tensor.data()) - .set_K(k_tensor.data()) - .set_V(v_tensor.data()) - .set_O(output_tensor.data()) - .set_dO(doutput_tensor.data()) - .set_S(s_tensor.data()) // not used for F16 - .set_dP(s_tensor.data()) // not used for F16 - .set_Aux_CTX_Tensors(&aux_input_tensors) - .set_dQ(dq_tensor.data()) - .set_dK(dk_tensor.data()) - .set_dV(dv_tensor.data()) - .set_dBias(dbias_tensor.data()) - .set_dSoftmaxOffset(dsoftmax_offset_tensor.data()) - .set_cu_seqlens_q(q_cu_seqlens_tensor.data()) - .set_cu_seqlens_kv(kv_cu_seqlens_tensor.data()) - .set_cu_seqlens_q_padded(q_seq_offsets_tensor.data()) - .set_cu_seqlens_kv_padded(k_seq_offsets_tensor.data()) - .set_deterministic(deterministic) - .set_cuda_graph(false) - .set_attn_mask_type(mask_type) - .set_bias_type(bias_type) - .set_window_size_left(window_size_left) - .set_window_size_right(window_size_right) - .set_bottom_right_diagonal(bottom_right_diagonal) - .set_softmax_type(softmax_type) - .set_dropout(dropout_probability) - .set_attn_scale(scaling_factor) - .set_qkv_layout(qkv_layout) - .set_o_format(nvte_get_q_format(qkv_layout)) - .set_do_format(nvte_get_q_format(qkv_layout)) - .set_dqkv_layout(qkv_layout) - .set_qkv_scale_inv_format(NVTE_QKV_Format_NOT_SET) - .set_do_scale_inv_format(NVTE_QKV_Format_NOT_SET) - .set_max_seqlen_q(q_max_seqlen) - .set_max_seqlen_kv(kv_max_seqlen) - .set_workspace(workspace_tensor.data()) - .set_stream(stream); - nvte_fused_attn_bwd_v2(params); - - nvte_tensor_pack_destroy(&aux_input_tensors); -} - -Error_Type FusedAttnBackwardFFI(cudaStream_t stream, Buffer_Type q_buf, Buffer_Type k_buf, - Buffer_Type v_buf, Buffer_Type bias_buf, - Buffer_Type softmax_offset_buf, Buffer_Type softmax_aux_buf, - Buffer_Type rng_state_buf, Buffer_Type output_buf, - Buffer_Type doutput_buf, Buffer_Type q_cu_seqlens_buf, - Buffer_Type kv_cu_seqlens_buf, Buffer_Type q_seq_offsets_buf, - Buffer_Type k_seq_offsets_buf, Variadic_Buffer_Type _unused_args, - Result_Type dq_buf, Result_Type dk_buf, Result_Type dv_buf, - Result_Type dbias_buf, Result_Type dsoftmax_offset_buf, - Result_Type workspace_buf, Dictionary attrs) { - FUSED_ATTN_FFI_GET_ATTRS; - - FusedAttnBackwardImpl( - stream, q_buf.untyped_data(), k_buf.untyped_data(), v_buf.untyped_data(), - bias_buf.untyped_data(), softmax_offset_buf.untyped_data(), softmax_aux_buf.untyped_data(), - rng_state_buf.untyped_data(), output_buf.untyped_data(), doutput_buf.untyped_data(), - q_cu_seqlens_buf.untyped_data(), kv_cu_seqlens_buf.untyped_data(), - is_ragged ? q_seq_offsets_buf.untyped_data() : nullptr, - is_ragged ? k_seq_offsets_buf.untyped_data() : nullptr, dq_buf->untyped_data(), - dk_buf->untyped_data(), dv_buf->untyped_data(), dbias_buf->untyped_data(), - dsoftmax_offset_buf->untyped_data(), workspace_buf->untyped_data(), input_batch, bias_batch, - q_max_seqlen, kv_max_seqlen, attn_heads, num_gqa_groups, bias_heads, qk_head_dim, v_head_dim, - max_segments_per_seq, wkspace_size, scaling_factor, dropout_probability, bias_type, mask_type, - softmax_type, qkv_layout, dtype, wkspace_dtype, is_training, deterministic, window_size_left, - window_size_right, bottom_right_diagonal); - - return ffi_with_cuda_error_check(); -} - -XLA_FFI_DEFINE_HANDLER_SYMBOL(FusedAttnBackwardHandler, FusedAttnBackwardFFI, - FFI::Bind() - .Ctx() // stream - .Arg() // q - .Arg() // k - .Arg() // v - .Arg() // bias - .Arg() // softmax_offset - .Arg() // softmax_aux - .Arg() // rng_state - .Arg() // output - .Arg() // doutput - .Arg() // q_cu_seqlens - .Arg() // kv_cu_seqlens - .Arg() // q_seq_offsets - .Arg() // k_seq_offsets - .RemainingArgs() // _cp_aux_args unused - .Ret() // dq - .Ret() // dk - .Ret() // dv - .Ret() // dbias - .Ret() // dsoftmax_offset - .Ret() // workspace - .Attrs(), - FFI_CudaGraph_Traits); - namespace { -struct ScoreModScalarStorage { +struct CudnnGraphScalarStorage { alignas(16) std::array data{}; - size_t size = 0; }; -struct ScoreModGraphCacheKey { +struct CudnnGraphCacheKey { int device_id = 0; int64_t hash0 = 0; int64_t hash1 = 0; int64_t frontend_version = 0; - bool operator==(const ScoreModGraphCacheKey &other) const { + bool operator==(const CudnnGraphCacheKey &other) const { return device_id == other.device_id && hash0 == other.hash0 && hash1 == other.hash1 && frontend_version == other.frontend_version; } }; -struct ScoreModGraphCacheKeyHash { - size_t operator()(const ScoreModGraphCacheKey &key) const { +struct CudnnGraphCacheKeyHash { + size_t operator()(const CudnnGraphCacheKey &key) const { size_t seed = std::hash{}(key.device_id); auto combine = [&seed](int64_t value) { - // 64-bit golden ratio constant from boost::hash_combine to spread mixed keys. seed ^= std::hash{}(value) + 0x9e3779b97f4a7c15ULL + (seed << 6) + (seed >> 2); }; combine(key.hash0); @@ -953,21 +53,20 @@ struct ScoreModGraphCacheKeyHash { } }; -using ScoreModGraphPtr = std::shared_ptr; +using CudnnGraphPtr = std::shared_ptr; -std::unordered_map & -getScoreModeGraphCache() { - static std::unordered_map - cache; +std::unordered_map & +GetCudnnGraphCache() { + static std::unordered_map cache; return cache; } -std::mutex &getScoreModGraphCacheMutex() { +std::mutex &GetCudnnGraphCacheMutex() { static std::mutex mutex; return mutex; } -struct ScoreModCudnnHandleCache { +struct CudnnHandleCache { std::unordered_map handles; cudnnHandle_t GetHandle() { @@ -982,30 +81,29 @@ struct ScoreModCudnnHandleCache { return it->second; } - ~ScoreModCudnnHandleCache() { + ~CudnnHandleCache() { for (auto &[_, handle] : handles) { cudnnDestroy(handle); } } }; -cudnnHandle_t GetScoreModCudnnHandle() { - static thread_local ScoreModCudnnHandleCache cache; +cudnnHandle_t GetCudnnHandle() { + static thread_local CudnnHandleCache cache; return cache.GetHandle(); } -ScoreModGraphCacheKey GetScoreModGraphCacheKey(Dictionary &attrs) { +CudnnGraphCacheKey GetCudnnGraphCacheKey(Dictionary &attrs) { const int64_t frontend_version = get_attr_value(attrs, "cudnn_frontend_version"); NVTE_CHECK(frontend_version == CUDNN_FRONTEND_VERSION, - "cuDNN frontend version mismatch for score_mod graph deserialization: graph was " - "serialized with Python cuDNN frontend version ", - frontend_version, - ", but Transformer Engine C++ was built with CUDNN_FRONTEND_VERSION ", + "cuDNN frontend version mismatch for graph deserialization: graph was serialized " + "with Python frontend version ", + frontend_version, ", but Transformer Engine C++ was built with version ", CUDNN_FRONTEND_VERSION, "."); int device_id = 0; NVTE_CHECK_CUDA(cudaGetDevice(&device_id)); - return ScoreModGraphCacheKey{ + return CudnnGraphCacheKey{ device_id, get_attr_value(attrs, "graph_hash0"), get_attr_value(attrs, "graph_hash1"), @@ -1013,11 +111,11 @@ ScoreModGraphCacheKey GetScoreModGraphCacheKey(Dictionary &attrs) { }; } -ScoreModGraphPtr GetScoreModGraph(cudaStream_t stream, Dictionary &attrs) { - const auto key = GetScoreModGraphCacheKey(attrs); +CudnnGraphPtr GetCudnnGraph(cudaStream_t stream, Dictionary &attrs) { + const auto key = GetCudnnGraphCacheKey(attrs); { - std::lock_guard lock(getScoreModGraphCacheMutex()); - auto &cache = getScoreModeGraphCache(); + std::lock_guard lock(GetCudnnGraphCacheMutex()); + auto &cache = GetCudnnGraphCache(); auto it = cache.find(key); if (it != cache.end()) { return it->second; @@ -1026,17 +124,16 @@ ScoreModGraphPtr GetScoreModGraph(cudaStream_t stream, Dictionary &attrs) { const auto serialized_graph = get_attr_value(attrs, "serialized_graph"); std::vector serialized_data(serialized_graph.begin(), serialized_graph.end()); - - auto handle = GetScoreModCudnnHandle(); + auto handle = GetCudnnHandle(); NVTE_CHECK_CUDNN(cudnnSetStream(handle, stream)); auto graph = std::make_shared(); auto status = graph->deserialize(handle, serialized_data); NVTE_CHECK(status.is_good(), - "Failed to deserialize cuDNN score_mod SDPA graph: ", status.get_message()); + "Failed to deserialize cuDNN frontend graph: ", status.get_message()); - std::lock_guard lock(getScoreModGraphCacheMutex()); - auto &cache = getScoreModeGraphCache(); + std::lock_guard lock(GetCudnnGraphCacheMutex()); + auto &cache = GetCudnnGraphCache(); auto it = cache.find(key); if (it != cache.end()) { return it->second; @@ -1045,47 +142,71 @@ ScoreModGraphPtr GetScoreModGraph(cudaStream_t stream, Dictionary &attrs) { return graph; } -Error_Type ExecuteScoreModGraph(cudaStream_t stream, Dictionary &attrs, - const std::vector &input_ptrs, - const std::vector &output_ptrs, void *workspace) { - auto graph = GetScoreModGraph(stream, attrs); +Error_Type ExecuteCudnnGraph(cudaStream_t stream, Dictionary &attrs, + const std::vector &input_ptrs, + const std::vector &output_ptrs, void *workspace) { + auto graph = GetCudnnGraph(stream, attrs); auto input_uids = get_attr_value>(attrs, "input_uids"); + auto input_buffer_indices = + get_attr_value>(attrs, "input_buffer_indices"); + auto input_byte_offsets = + get_attr_value>(attrs, "input_byte_offsets"); auto output_uids = get_attr_value>(attrs, "output_uids"); + auto output_buffer_indices = + get_attr_value>(attrs, "output_buffer_indices"); + auto output_byte_offsets = + get_attr_value>(attrs, "output_byte_offsets"); auto scalar_uids = get_attr_value>(attrs, "scalar_uids"); auto scalar_sizes = get_attr_value>(attrs, "scalar_sizes"); auto scalar_values = get_attr_value>(attrs, "scalar_values"); - NVTE_CHECK(input_ptrs.size() == input_uids.size(), "cuDNN score_mod graph expected ", - input_uids.size(), " inputs but got ", input_ptrs.size()); - NVTE_CHECK(output_ptrs.size() >= output_uids.size(), "cuDNN score_mod graph expected at least ", - output_uids.size(), " outputs but got ", output_ptrs.size()); + NVTE_CHECK(input_uids.size() == input_buffer_indices.size() && + input_uids.size() == input_byte_offsets.size(), + "Mismatched cuDNN graph input binding metadata."); + NVTE_CHECK(output_uids.size() == output_buffer_indices.size() && + output_uids.size() == output_byte_offsets.size(), + "Mismatched cuDNN graph output binding metadata."); NVTE_CHECK(scalar_uids.size() == scalar_sizes.size(), - "Mismatched score_mod scalar uid/value-size counts."); + "Mismatched cuDNN graph scalar uid/value-size counts."); NVTE_CHECK(scalar_values.size() == scalar_uids.size() * 16, - "Mismatched score_mod packed scalar value size."); + "Mismatched cuDNN graph packed scalar value size."); std::unordered_map variant_pack; for (size_t i = 0; i < input_uids.size(); ++i) { - variant_pack.emplace(input_uids[i], input_ptrs[i]); + NVTE_CHECK(input_buffer_indices[i] >= 0 && + static_cast(input_buffer_indices[i]) < input_ptrs.size(), + "cuDNN graph input binding index is out of range."); + NVTE_CHECK(input_byte_offsets[i] >= 0, "cuDNN graph input byte offset must be non-negative."); + auto *ptr = static_cast(input_ptrs[input_buffer_indices[i]]) + input_byte_offsets[i]; + variant_pack.emplace(input_uids[i], ptr); } for (size_t i = 0; i < output_uids.size(); ++i) { - variant_pack.emplace(output_uids[i], output_ptrs[i]); + NVTE_CHECK(output_buffer_indices[i] >= 0 && + static_cast(output_buffer_indices[i]) < output_ptrs.size(), + "cuDNN graph output binding index is out of range."); + NVTE_CHECK(output_byte_offsets[i] >= 0, "cuDNN graph output byte offset must be non-negative."); + auto *ptr = + static_cast(output_ptrs[output_buffer_indices[i]]) + output_byte_offsets[i]; + variant_pack.emplace(output_uids[i], ptr); } - std::vector scalar_storage(scalar_uids.size()); + std::vector scalar_storage(scalar_uids.size()); for (size_t i = 0; i < scalar_uids.size(); ++i) { NVTE_CHECK(scalar_sizes[i] >= 0 && scalar_sizes[i] <= 16, - "score_mod pass-by-value scalars must be at most 16 bytes."); - scalar_storage[i].size = static_cast(scalar_sizes[i]); + "cuDNN graph pass-by-value scalars must be at most 16 bytes."); std::copy_n(scalar_values.begin() + i * 16, 16, scalar_storage[i].data.begin()); variant_pack.emplace(scalar_uids[i], scalar_storage[i].data.data()); } - auto handle = GetScoreModCudnnHandle(); + auto handle = GetCudnnHandle(); NVTE_CHECK_CUDNN(cudnnSetStream(handle, stream)); + int device_id = 0; + NVTE_CHECK_CUDA(cudaGetDevice(&device_id)); + attention_cache_debug::Record(get_attr_value(attrs, "attention_backend"), + get_attr_value(attrs, "attention_direction"), + "execute", device_id); auto status = graph->execute(handle, variant_pack, workspace); - NVTE_CHECK(status.is_good(), - "cuDNN score_mod SDPA graph execution failed: ", status.get_message()); + NVTE_CHECK(status.is_good(), "cuDNN frontend graph execution failed: ", status.get_message()); return ffi_with_cuda_error_check(); } @@ -1093,13 +214,173 @@ void AppendRemainingBuffers(Variadic_Buffer_Type args, std::vector *ptrs ptrs->reserve(ptrs->size() + args.size()); for (size_t i = 0; i < args.size(); ++i) { auto maybe_buf = args.get(i); - NVTE_CHECK(!maybe_buf.has_error(), "Failed to decode variadic score_mod input buffer."); + NVTE_CHECK(!maybe_buf.has_error(), "Failed to decode variadic cuDNN graph input buffer."); ptrs->push_back(maybe_buf.value().untyped_data()); } } +size_t BufferBytes(const Buffer_Type &buffer) { return buffer.size_bytes(); } + +void MemsetResultAsync(cudaStream_t stream, Result_Type result, int value) { + NVTE_CHECK_CUDA(cudaMemsetAsync(result->untyped_data(), value, BufferBytes(*result), stream)); +} + +class FusedAttnOffsetManager { + public: + static FusedAttnOffsetManager &Instance() { + static thread_local FusedAttnOffsetManager manager; + return manager; + } + + uint64_t GetAndUpdate(uint64_t increment) { + uint64_t current = offset_; + offset_ += increment; + return current; + } + + private: + uint64_t offset_ = 0; +}; + +void PopulateRngStateAsync(cudaStream_t stream, const Buffer_Type &seed, Result_Type rng_state, + uint64_t increment) { + NVTE_CHECK(BufferBytes(seed) >= sizeof(uint64_t), "Fused-attention seed buffer is too small."); + NVTE_CHECK(BufferBytes(*rng_state) >= 2 * sizeof(uint64_t), + "Fused-attention RNG-state buffer is too small."); + const uint64_t offset = FusedAttnOffsetManager::Instance().GetAndUpdate(increment); + PopulateFusedAttnRngState(rng_state->untyped_data(), seed.untyped_data(), offset, stream); +} + } // namespace +Error_Type FusedAttnForwardFFI(cudaStream_t stream, Buffer_Type q_buf, Buffer_Type k_buf, + Buffer_Type v_buf, Buffer_Type bias_buf, + Buffer_Type softmax_offset_buf, Buffer_Type seed_buf, + Buffer_Type q_seqlens_buf, Buffer_Type kv_seqlens_buf, + Buffer_Type q_seq_offsets_buf, Buffer_Type k_seq_offsets_buf, + Variadic_Buffer_Type remaining_args, Result_Type output_buf, + Result_Type stats_buf, Result_Type max_buf, + Result_Type rng_state_buf, Result_Type workspace_buf, + Dictionary attrs) { + const bool is_ragged = get_attr_value(attrs, "is_ragged"); + const uint64_t rng_increment = + static_cast(get_attr_value(attrs, "rng_offset_increment")); + PopulateRngStateAsync(stream, seed_buf, rng_state_buf, rng_increment); + if (is_ragged) { + MemsetResultAsync(stream, output_buf, 0); + MemsetResultAsync(stream, stats_buf, 0xF0); + if (BufferBytes(*max_buf) != 0) { + MemsetResultAsync(stream, max_buf, 0xF0); + } + } + + std::vector input_ptrs = { + q_buf.untyped_data(), + k_buf.untyped_data(), + v_buf.untyped_data(), + bias_buf.untyped_data(), + softmax_offset_buf.untyped_data(), + seed_buf.untyped_data(), + q_seqlens_buf.untyped_data(), + kv_seqlens_buf.untyped_data(), + q_seq_offsets_buf.untyped_data(), + k_seq_offsets_buf.untyped_data(), + }; + AppendRemainingBuffers(remaining_args, &input_ptrs); + std::vector output_ptrs = {output_buf->untyped_data(), stats_buf->untyped_data(), + max_buf->untyped_data(), rng_state_buf->untyped_data()}; + return ExecuteCudnnGraph(stream, attrs, input_ptrs, output_ptrs, workspace_buf->untyped_data()); +} + +XLA_FFI_DEFINE_HANDLER_SYMBOL(FusedAttnForwardHandler, FusedAttnForwardFFI, + FFI::Bind() + .Ctx() + .Arg() + .Arg() + .Arg() + .Arg() + .Arg() + .Arg() + .Arg() + .Arg() + .Arg() + .Arg() + .RemainingArgs() + .Ret() + .Ret() + .Ret() + .Ret() + .Ret() + .Attrs(), + FFI_CudaGraph_Traits); + +Error_Type FusedAttnBackwardFFI(cudaStream_t stream, Buffer_Type q_buf, Buffer_Type k_buf, + Buffer_Type v_buf, Buffer_Type bias_buf, + Buffer_Type softmax_offset_buf, Buffer_Type stats_buf, + Buffer_Type rng_state_buf, Buffer_Type output_buf, + Buffer_Type doutput_buf, Buffer_Type q_seqlens_buf, + Buffer_Type kv_seqlens_buf, Buffer_Type q_seq_offsets_buf, + Buffer_Type k_seq_offsets_buf, Variadic_Buffer_Type remaining_args, + Result_Type dq_buf, Result_Type dk_buf, Result_Type dv_buf, + Result_Type dbias_buf, Result_Type dsoftmax_offset_buf, + Result_Type workspace_buf, Dictionary attrs) { + if (get_attr_value(attrs, "is_ragged")) { + MemsetResultAsync(stream, dq_buf, 0); + MemsetResultAsync(stream, dk_buf, 0); + MemsetResultAsync(stream, dv_buf, 0); + } + std::vector input_ptrs = { + q_buf.untyped_data(), + k_buf.untyped_data(), + v_buf.untyped_data(), + bias_buf.untyped_data(), + softmax_offset_buf.untyped_data(), + stats_buf.untyped_data(), + rng_state_buf.untyped_data(), + output_buf.untyped_data(), + doutput_buf.untyped_data(), + q_seqlens_buf.untyped_data(), + kv_seqlens_buf.untyped_data(), + q_seq_offsets_buf.untyped_data(), + k_seq_offsets_buf.untyped_data(), + }; + AppendRemainingBuffers(remaining_args, &input_ptrs); + std::vector output_ptrs = { + dq_buf->untyped_data(), + dk_buf->untyped_data(), + dv_buf->untyped_data(), + dbias_buf->untyped_data(), + dsoftmax_offset_buf->untyped_data(), + }; + return ExecuteCudnnGraph(stream, attrs, input_ptrs, output_ptrs, workspace_buf->untyped_data()); +} + +XLA_FFI_DEFINE_HANDLER_SYMBOL(FusedAttnBackwardHandler, FusedAttnBackwardFFI, + FFI::Bind() + .Ctx() + .Arg() + .Arg() + .Arg() + .Arg() + .Arg() + .Arg() + .Arg() + .Arg() + .Arg() + .Arg() + .Arg() + .Arg() + .Arg() + .RemainingArgs() + .Ret() + .Ret() + .Ret() + .Ret() + .Ret() + .Ret() + .Attrs(), + FFI_CudaGraph_Traits); + Error_Type FusedAttnScoreModForwardFFI(cudaStream_t stream, Buffer_Type q_buf, Buffer_Type k_buf, Buffer_Type v_buf, Variadic_Buffer_Type score_mod_args, Result_Type output_buf, Result_Type stats_buf, @@ -1107,23 +388,22 @@ Error_Type FusedAttnScoreModForwardFFI(cudaStream_t stream, Buffer_Type q_buf, B std::vector input_ptrs = {q_buf.untyped_data(), k_buf.untyped_data(), v_buf.untyped_data()}; AppendRemainingBuffers(score_mod_args, &input_ptrs); - std::vector output_ptrs = {output_buf->untyped_data(), stats_buf->untyped_data()}; - return ExecuteScoreModGraph(stream, attrs, input_ptrs, output_ptrs, - workspace_buf->untyped_data()); + return ExecuteCudnnGraph(stream, attrs, input_ptrs, output_ptrs, workspace_buf->untyped_data()); } XLA_FFI_DEFINE_HANDLER_SYMBOL(FusedAttnScoreModForwardHandler, FusedAttnScoreModForwardFFI, FFI::Bind() - .Ctx() // stream - .Arg() // q - .Arg() // k - .Arg() // v - .RemainingArgs() // score_mod tensor operands - .Ret() // output - .Ret() // stats - .Ret() // workspace - .Attrs()); + .Ctx() + .Arg() + .Arg() + .Arg() + .RemainingArgs() + .Ret() + .Ret() + .Ret() + .Attrs(), + FFI_CudaGraph_Traits); Error_Type FusedAttnScoreModBackwardFFI(cudaStream_t stream, Buffer_Type q_buf, Buffer_Type k_buf, Buffer_Type v_buf, Buffer_Type output_buf, @@ -1135,28 +415,27 @@ Error_Type FusedAttnScoreModBackwardFFI(cudaStream_t stream, Buffer_Type q_buf, v_buf.untyped_data(), output_buf.untyped_data(), doutput_buf.untyped_data(), stats_buf.untyped_data()}; AppendRemainingBuffers(score_mod_args, &input_ptrs); - std::vector output_ptrs = {dq_buf->untyped_data(), dk_buf->untyped_data(), dv_buf->untyped_data()}; - return ExecuteScoreModGraph(stream, attrs, input_ptrs, output_ptrs, - workspace_buf->untyped_data()); + return ExecuteCudnnGraph(stream, attrs, input_ptrs, output_ptrs, workspace_buf->untyped_data()); } XLA_FFI_DEFINE_HANDLER_SYMBOL(FusedAttnScoreModBackwardHandler, FusedAttnScoreModBackwardFFI, FFI::Bind() - .Ctx() // stream - .Arg() // q - .Arg() // k - .Arg() // v - .Arg() // output - .Arg() // doutput - .Arg() // stats - .RemainingArgs() // score_mod tensor operands - .Ret() // dq - .Ret() // dk - .Ret() // dv - .Ret() // workspace - .Attrs()); + .Ctx() + .Arg() + .Arg() + .Arg() + .Arg() + .Arg() + .Arg() + .RemainingArgs() + .Ret() + .Ret() + .Ret() + .Ret() + .Attrs(), + FFI_CudaGraph_Traits); } // namespace jax } // namespace transformer_engine diff --git a/transformer_engine/jax/csrc/extensions/attention_cache_debug.h b/transformer_engine/jax/csrc/extensions/attention_cache_debug.h new file mode 100644 index 00000000000..a0b71e34c85 --- /dev/null +++ b/transformer_engine/jax/csrc/extensions/attention_cache_debug.h @@ -0,0 +1,244 @@ +/************************************************************************* + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * + * See LICENSE for license information. + ************************************************************************/ + +#ifndef TRANSFORMER_ENGINE_JAX_CSRC_EXTENSIONS_ATTENTION_CACHE_DEBUG_H_ +#define TRANSFORMER_ENGINE_JAX_CSRC_EXTENSIONS_ATTENTION_CACHE_DEBUG_H_ + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace transformer_engine { +namespace jax { +namespace attention_cache_debug { +namespace detail { + +constexpr size_t kSiteCount = 4; +constexpr size_t kStageCount = 5; +inline constexpr std::array kStageNames = { + "validate", "build_operation_graph", "create_execution_plans", "check_support", "build_plans"}; + +inline int DebugLevel() { + static const int level = [] { + const char *value = std::getenv("NVTE_FUSED_ATTN_CACHE_DEBUG"); + if (value == nullptr || value[0] == '\0' || value[0] == '0') return 0; + const int parsed = std::atoi(value); + return parsed > 0 ? parsed : 1; + }(); + return level; +} + +inline int LauncherRank() { + static const int rank = [] { + for (const char *name : {"RANK", "LOCAL_RANK", "OMPI_COMM_WORLD_RANK", "SLURM_PROCID"}) { + const char *value = std::getenv(name); + if (value != nullptr && value[0] != '\0') return std::atoi(value); + } + return -1; + }(); + return rank; +} + +inline bool Enabled() { + static const bool enabled = [] { + if (DebugLevel() < 1) return false; + const int rank = LauncherRank(); + if (rank < 0) return true; + const char *value = std::getenv("NVTE_FUSED_ATTN_CACHE_DEBUG"); + const char *separator = value == nullptr ? nullptr : std::strchr(value, ':'); + if (separator == nullptr) return rank == 0; + const std::string ranks(separator + 1); + if (ranks == "all") return true; + for (size_t start = 0; start <= ranks.size();) { + const size_t end = ranks.find(',', start); + const std::string token = ranks.substr(start, end - start); + if (!token.empty() && std::atoi(token.c_str()) == rank) return true; + if (end == std::string::npos) break; + start = end + 1; + } + return false; + }(); + return enabled; +} + +inline bool TraceEnabled() { return Enabled() && DebugLevel() >= 2; } + +inline size_t SiteIndex(std::string_view backend, std::string_view direction) { + const size_t backend_index = backend == "f16" ? 0 : backend == "fp8" ? 1 : kSiteCount; + const size_t direction_index = direction == "fwd" ? 0 : direction == "bwd" ? 1 : 2; + if (backend_index > 1 || direction_index > 1) { + throw std::invalid_argument("Invalid fused-attention cache diagnostic site: " + + std::string(backend) + " " + std::string(direction)); + } + return backend_index * 2 + direction_index; +} + +struct Counters { + std::atomic hit{0}; + std::atomic miss{0}; + std::atomic create_graph{0}; + std::atomic cache_graph{0}; + std::atomic build_plans{0}; + std::atomic execute{0}; +}; + +struct Timing { + std::atomic calls{0}; + std::atomic elapsed_ns{0}; +}; + +inline std::array &AllCounters() { + static auto *counters = new std::array(); + return *counters; +} + +inline std::array &AllTimings() { + static auto *timings = new std::array(); + return *timings; +} + +inline std::mutex &OutputMutex() { + static auto *mutex = new std::mutex(); + return *mutex; +} + +inline const char *Backend(size_t site) { return site < 2 ? "f16" : "fp8"; } +inline const char *Direction(size_t site) { return site % 2 == 0 ? "fwd" : "bwd"; } + +inline std::string RankTag() { + const int rank = LauncherRank(); + return rank < 0 ? "" : "rank=" + std::to_string(rank) + " | "; +} + +inline void Write(const std::string &message) { + std::lock_guard lock(OutputMutex()); + std::fwrite(message.data(), 1, message.size(), stderr); + std::fflush(stderr); +} + +inline uint64_t Load(const std::atomic &value) { + return value.load(std::memory_order_relaxed); +} + +inline std::string CounterLine(size_t site, const char *event = nullptr, int device = -1, + bool all_devices = false) { + const Counters &counters = AllCounters()[site]; + const std::string device_name = all_devices ? "all" : std::to_string(device); + char line[640]; + std::snprintf(line, sizeof(line), + "[FUSED-ATTN-CACHE] %sdev=%-3s | %s %s %-12s | hit=%4" PRIu64 ", miss=%4" PRIu64 + ", create_graph=%4" PRIu64 ", cache_graph=%4" PRIu64 ", build_plans=%4" PRIu64 + ", execute=%4" PRIu64 "\n", + RankTag().c_str(), device_name.c_str(), Backend(site), Direction(site), + event == nullptr ? "" : event, Load(counters.hit), Load(counters.miss), + Load(counters.create_graph), Load(counters.cache_graph), Load(counters.build_plans), + Load(counters.execute)); + return line; +} + +inline void PrintSummary() { + if (!Enabled()) return; + const std::string marker = "[FUSED-ATTN-CACHE] " + RankTag() + "===== summary "; + std::string output = marker + "begin =====\n"; + for (size_t site = 0; site < kSiteCount; ++site) { + const Counters &counters = AllCounters()[site]; + if ((Load(counters.hit) | Load(counters.miss) | Load(counters.create_graph) | + Load(counters.cache_graph) | Load(counters.build_plans) | Load(counters.execute)) != 0) { + output += CounterLine(site, nullptr, -1, true); + } + } + for (size_t site = 0; site < kSiteCount; ++site) { + for (size_t stage = 0; stage < kStageCount; ++stage) { + const Timing &timing = AllTimings()[site * kStageCount + stage]; + const uint64_t calls = Load(timing.calls); + if (calls == 0) continue; + const double milliseconds = static_cast(Load(timing.elapsed_ns)) / calls / 1e6; + char line[320]; + std::snprintf(line, sizeof(line), + "[FUSED-ATTN-CACHE] %s%s %-3s %-22s | calls=%" PRIu64 " | time=%9.3f ms/call\n", + RankTag().c_str(), Backend(site), Direction(site), kStageNames[stage], calls, + milliseconds); + output += line; + } + } + output += marker + "end =====\n"; + Write(output); +} + +inline void RegisterSummary() { + static const bool registered = [] { + std::atexit(PrintSummary); + return true; + }(); + (void)registered; +} + +inline std::atomic *EventCounter(Counters &counters, std::string_view event) { + if (event == "hit") return &counters.hit; + if (event == "miss") return &counters.miss; + if (event == "create_graph") return &counters.create_graph; + if (event == "cache_graph") return &counters.cache_graph; + if (event == "plans_built") return &counters.build_plans; + if (event == "execute") return &counters.execute; + return nullptr; +} + +inline size_t StageIndex(std::string_view event) { + for (size_t index = 0; index < kStageCount; ++index) { + if (event == kStageNames[index]) return index; + } + return kStageCount; +} + +} // namespace detail + +inline void Record(std::string_view backend, std::string_view direction, std::string_view event, + int device = -1, std::string_view key = {}, uint64_t elapsed_ns = 0) { + if (!detail::Enabled()) return; + detail::RegisterSummary(); + const size_t site = detail::SiteIndex(backend, direction); + const size_t stage = detail::StageIndex(event); + if (stage < detail::kStageCount) { + detail::Timing &timing = detail::AllTimings()[site * detail::kStageCount + stage]; + timing.calls.fetch_add(1, std::memory_order_relaxed); + timing.elapsed_ns.fetch_add(elapsed_ns, std::memory_order_relaxed); + return; + } + + detail::Counters &counters = detail::AllCounters()[site]; + std::atomic *counter = detail::EventCounter(counters, event); + if (counter == nullptr) { + throw std::invalid_argument("Invalid fused-attention cache diagnostic event: " + + std::string(event)); + } + counter->fetch_add(1, std::memory_order_relaxed); + if (!detail::TraceEnabled()) return; + if ((event == "hit" || event == "miss") && !key.empty()) { + detail::Write("[FUSED-ATTN-CACHE] " + detail::RankTag() + "dev=" + std::to_string(device) + + " | " + std::string(backend) + " " + std::string(direction) + " " + + std::string(event) + " | " + std::string(key) + "\n"); + } else { + std::string uppercase(event == "plans_built" ? "build_plans" : event); + for (char &character : uppercase) { + if (character >= 'a' && character <= 'z') character -= 'a' - 'A'; + } + detail::Write(detail::CounterLine(site, uppercase.c_str(), device)); + } +} + +} // namespace attention_cache_debug +} // namespace jax +} // namespace transformer_engine + +#endif // TRANSFORMER_ENGINE_JAX_CSRC_EXTENSIONS_ATTENTION_CACHE_DEBUG_H_ diff --git a/transformer_engine/jax/csrc/extensions/attention_kernels.cu b/transformer_engine/jax/csrc/extensions/attention_kernels.cu new file mode 100644 index 00000000000..c2797292f30 --- /dev/null +++ b/transformer_engine/jax/csrc/extensions/attention_kernels.cu @@ -0,0 +1,29 @@ +/************************************************************************* + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * + * See LICENSE for license information. + ************************************************************************/ + +#include "../extensions.h" + +namespace transformer_engine { +namespace jax { +namespace { + +__global__ void PopulateFusedAttnRngStateKernel(int64_t *rng_state, const int64_t *seed, + uint64_t offset) { + rng_state[0] = seed[0]; + rng_state[1] = static_cast(offset); +} + +} // namespace + +void PopulateFusedAttnRngState(void *rng_state, const void *seed, uint64_t offset, + cudaStream_t stream) { + PopulateFusedAttnRngStateKernel<<<1, 1, 0, stream>>>(static_cast(rng_state), + static_cast(seed), offset); + NVTE_CHECK_CUDA(cudaGetLastError()); +} + +} // namespace jax +} // namespace transformer_engine diff --git a/transformer_engine/jax/csrc/extensions/pybind.cpp b/transformer_engine/jax/csrc/extensions/pybind.cpp index e65cea9fa0e..a5820f4b3b9 100644 --- a/transformer_engine/jax/csrc/extensions/pybind.cpp +++ b/transformer_engine/jax/csrc/extensions/pybind.cpp @@ -5,6 +5,7 @@ ************************************************************************/ #include "../extensions.h" +#include "attention_cache_debug.h" #include "cgemm_helper.h" #include "common/util/cuda_runtime.h" #include "transformer_engine/gemm.h" @@ -134,11 +135,17 @@ pybind11::dict Registrations() { PYBIND11_MODULE(transformer_engine_jax, m) { m.def("registrations", &Registrations); - m.def("get_fused_attn_backend", &GetFusedAttnBackend, "Get Fused Attention backend", - pybind11::arg("fused_attn_params")); m.def("get_cuda_version", &GetCudaRuntimeVersion); m.def("get_cudnn_version", &GetCudnnRuntimeVersion); m.def("get_cudnn_frontend_version", &GetCudnnFrontendVersion); + m.def( + "record_fused_attn_cache_event", + [](const std::string &backend, const std::string &direction, const std::string &event, + int device, const std::string &key, uint64_t elapsed_ns) { + attention_cache_debug::Record(backend, direction, event, device, key, elapsed_ns); + }, + pybind11::arg("backend"), pybind11::arg("direction"), pybind11::arg("event"), + pybind11::arg("device") = -1, pybind11::arg("key") = "", pybind11::arg("elapsed_ns") = 0); m.def("get_device_compute_capability", &GetDeviceComputeCapability); m.def("get_num_compute_streams", &nvte_get_num_compute_streams); m.def("get_cublasLt_version", &cublasLtGetVersion); @@ -146,8 +153,6 @@ PYBIND11_MODULE(transformer_engine_jax, m) { m.def("get_dbias_quantize_workspace_sizes", &GetDBiasQuantizeWorkspaceSizes); m.def("get_norm_fwd_workspace_sizes", &GetNormForwardWorkspaceSizes); m.def("get_norm_bwd_workspace_sizes", &GetNormBackwardWorkspaceSizes); - m.def("get_fused_attn_fwd_workspace_sizes", &GetFusedAttnForwardWorkspaceSizes); - m.def("get_fused_attn_bwd_workspace_sizes", &GetFusedAttnBackwardWorkspaceSizes); m.def("get_topk_workspace_sizes", &GetTopkWorkspaceSizes); m.def("nvte_get_qkv_format", &nvte_get_qkv_format); m.def("is_non_nt_fp8_gemm_supported", &nvte_is_non_tn_fp8_gemm_supported); @@ -184,7 +189,8 @@ PYBIND11_MODULE(transformer_engine_jax, m) { pybind11::enum_(m, "NVTE_Bias_Type", pybind11::module_local()) .value("NVTE_NO_BIAS", NVTE_Bias_Type::NVTE_NO_BIAS) .value("NVTE_PRE_SCALE_BIAS", NVTE_Bias_Type::NVTE_PRE_SCALE_BIAS) - .value("NVTE_POST_SCALE_BIAS", NVTE_Bias_Type::NVTE_POST_SCALE_BIAS); + .value("NVTE_POST_SCALE_BIAS", NVTE_Bias_Type::NVTE_POST_SCALE_BIAS) + .value("NVTE_ALIBI", NVTE_Bias_Type::NVTE_ALIBI); pybind11::enum_(m, "NVTE_Mask_Type", pybind11::module_local()) .value("NVTE_NO_MASK", NVTE_Mask_Type::NVTE_NO_MASK) @@ -201,14 +207,12 @@ PYBIND11_MODULE(transformer_engine_jax, m) { .value("NVTE_BSHD_BSHD_BSHD", NVTE_QKV_Layout::NVTE_BSHD_BSHD_BSHD) .value("NVTE_T3HD", NVTE_QKV_Layout::NVTE_T3HD) .value("NVTE_THD_T2HD", NVTE_QKV_Layout::NVTE_THD_T2HD) - .value("NVTE_THD_THD_THD", NVTE_QKV_Layout::NVTE_THD_THD_THD) - .value("NVTE_QKV_Layout_NOT_SET", NVTE_QKV_Layout::NVTE_QKV_Layout_NOT_SET); + .value("NVTE_THD_THD_THD", NVTE_QKV_Layout::NVTE_THD_THD_THD); pybind11::enum_(m, "NVTE_QKV_Format", pybind11::module_local()) .value("NVTE_SBHD", NVTE_QKV_Format::NVTE_SBHD) .value("NVTE_BSHD", NVTE_QKV_Format::NVTE_BSHD) - .value("NVTE_THD", NVTE_QKV_Format::NVTE_THD) - .value("NVTE_QKV_Format_NOT_SET", NVTE_QKV_Format::NVTE_QKV_Format_NOT_SET); + .value("NVTE_THD", NVTE_QKV_Format::NVTE_THD); pybind11::enum_(m, "NVTE_Softmax_Type", pybind11::module_local()) .value("NVTE_VANILLA_SOFTMAX", NVTE_Softmax_Type::NVTE_VANILLA_SOFTMAX) diff --git a/transformer_engine/jax/flax/module.py b/transformer_engine/jax/flax/module.py index 17c9a242f04..e0012bf5926 100644 --- a/transformer_engine/jax/flax/module.py +++ b/transformer_engine/jax/flax/module.py @@ -17,6 +17,7 @@ from jax.ad_checkpoint import checkpoint_name from transformer_engine.common.recipe import ( + DelayedScaling, MXFP8BlockScaling, ) @@ -37,6 +38,7 @@ jax_scaled_upper_triang_masked_softmax, ) from ..quantize import ( + AttentionQuantizerSet, QuantizerFactory, get_global_quantize_recipe, QuantizeMetaSet, @@ -417,6 +419,42 @@ def generate_quantizer_set( ) return quantizer_set + def generate_attention_quantizer_set(self, fp8_recipe=None): + """Generate independent quantizers for all FP8 DPA tensor roles.""" + + if fp8_recipe is None: + fp8_recipe = get_global_quantize_recipe() + first = self.generate_quantizer_set(postfix="_attention_qkv_s_do", fp8_recipe=fp8_recipe) + second = self.generate_quantizer_set(postfix="_attention_o_dp", fp8_recipe=fp8_recipe) + third = self.generate_quantizer_set(postfix="_attention_dqkv", fp8_recipe=fp8_recipe) + s_quantizer = first.kernel + dp_quantizer = second.dgrad + if fp8_recipe.float8_current_scaling(): + # S and dP are internal to the fused attention graph, so their current + # amax is not available before they are quantized. Match the PyTorch DPA + # recipe by maintaining one-step delayed scales for these two roles. + delayed_recipe = DelayedScaling( + fp8_format=fp8_recipe.fp8_format, + amax_history_len=1, + amax_compute_algo="most_recent", + fp8_dpa=fp8_recipe.fp8_dpa, + fp8_mha=fp8_recipe.fp8_mha, + ) + internal = self.generate_quantizer_set( + postfix="_attention_s_dp", + fp8_recipe=delayed_recipe, + ) + s_quantizer = internal.kernel + dp_quantizer = internal.dgrad + return AttentionQuantizerSet( + qkv=first.x, + s=s_quantizer, + o=second.x, + do=first.dgrad, + dp=dp_quantizer, + dqkv=third.dgrad, + ) + class DenseGeneral(TransformerEngineBase): r""" diff --git a/transformer_engine/jax/flax/transformer.py b/transformer_engine/jax/flax/transformer.py index 5a01572f1fa..d8ceab1c84d 100644 --- a/transformer_engine/jax/flax/transformer.py +++ b/transformer_engine/jax/flax/transformer.py @@ -5,7 +5,6 @@ Wrapper module for Transformer related layers with FP8 support. """ import functools -import operator from enum import Enum from math import sqrt import os @@ -21,9 +20,8 @@ from jax import random as jax_random from jax import lax, vmap from jax.ad_checkpoint import checkpoint_name -from transformer_engine_jax import NVTE_Fused_Attn_Backend -from .module import DenseGeneral, LayerNormDenseGeneral, LayerNormMLP +from .module import DenseGeneral, LayerNormDenseGeneral, LayerNormMLP, TransformerEngineBase from .module import LayerNorm, Softmax from ..attention import ( AttnBiasType, @@ -32,10 +30,10 @@ QKVLayout, SequenceDescriptor, ) -from ..attention import make_swa_mask, canonicalize_attn_mask_type +from ..attention import is_fused_attn_kernel_available, make_swa_mask, canonicalize_attn_mask_type from ..attention import fused_attn from ..attention import CPStrategy -from ..cpp_extensions import FusedAttnHelper +from ..quantize import get_global_quantize_recipe from ..softmax import SoftmaxFusionType from ..sharding import num_of_devices from ..sharding import get_sharding_map_logic_axis_to_mesh_axis @@ -294,7 +292,7 @@ def convert_to_softmax_fusion_type(attn_mask_type, mask): return jnp.einsum("bhqk,bkhd->bqhd", attn_weights, value) -class _FusedDotProductAttention(nn.Module): # pylint: disable=too-few-public-methods +class _FusedDotProductAttention(TransformerEngineBase): # pylint: disable=too-few-public-methods attention_dropout: float = 0.0 attn_mask_type: AttnMaskType = AttnMaskType.CAUSAL_MASK attn_bias_type: Optional[AttnBiasType] = None @@ -312,6 +310,7 @@ class _FusedDotProductAttention(nn.Module): # pylint: disable=too-few-public-me score_mod_bprop: Optional[Callable] = None score_mod_requested: bool = False return_max_logit: bool = False + bottom_right_diagonal: Optional[bool] = None @nn.compact def __call__( @@ -368,7 +367,16 @@ def __call__( "score_mod_tensors": score_mod_tensors, "score_mod_bprop_tensors": score_mod_bprop_tensors, "return_max_logit": self.return_max_logit, + "bottom_right_diagonal": self.bottom_right_diagonal, } + fp8_recipe = get_global_quantize_recipe() + if fp8_recipe is not None and getattr(fp8_recipe, "fp8_dpa", False): + if getattr(fp8_recipe, "fp8_mha", False): + raise NotImplementedError( + "JAX FP8 attention currently supports FP16/BF16 DPA boundaries only; " + "fp8_mha is not implemented." + ) + fused_attn_kwargs["quantizer_set"] = self.generate_attention_quantizer_set(fp8_recipe) if self.qkv_layout.is_qkvpacked(): """qkvpacked format, treat @@ -632,6 +640,9 @@ class DotProductAttention(nn.Module): # pylint: disable=too-few-public-methods return_max_logit: bool, default = False If True, return ``(output, max_logit)`` where ``max_logit`` contains the per-head maximum attention logits with shape ``[h]``. This path requires fused attention. + bottom_right_diagonal: Optional[bool], default = None + Explicit diagonal alignment for fused attention. When unset, bottom-right mask types + use bottom-right alignment and other masks use top-left alignment. Optimization parameters ----------------------- @@ -661,6 +672,7 @@ class DotProductAttention(nn.Module): # pylint: disable=too-few-public-methods score_mod: Optional[Callable] = None score_mod_bprop: Optional[Callable] = None return_max_logit: bool = False + bottom_right_diagonal: Optional[bool] = None def __post_init__(self): # TODO(KshitijLakhani): Remove warning in TransformerEngine v2.12 @@ -761,10 +773,10 @@ def __call__( or score_mod_bprop_tensors is not None ) - if attn_bias_type == AttnBiasType.NO_BIAS: + if attn_bias_type in (AttnBiasType.NO_BIAS, AttnBiasType.ALIBI): assert ( bias is None - ), f"bias must be None when attn_bias_type is NO_BIAS, but got bias={bias}" + ), f"bias must be None when attn_bias_type is {attn_bias_type}, but got bias={bias}" else: assert ( bias is not None @@ -795,8 +807,6 @@ def __call__( raise ValueError("return_max_logit requires fused attention, but NVTE_FUSED_ATTN=0.") sequence_dim = 0 if self.transpose_batch_sequence else 1 - batch_dim = 1 - sequence_dim - batch_size = query.shape[batch_dim] seqlen_q = query.shape[sequence_dim] if qkv_layout == QKVLayout.BS3HD: seqlen_kv = seqlen_q @@ -813,41 +823,28 @@ def __call__( if not enable_fused_attn: raise ValueError("score_mod requires fused attention, but NVTE_FUSED_ATTN=0.") kernel_qkv_layout = qkv_layout.to_separate() if score_mod_requested else qkv_layout - bias_batch = bias_heads = bias_seqlen_q = bias_seqlen_kv = None - if attn_bias_type == AttnBiasType.POST_SCALE_BIAS: - *bias_batch_shape, bias_heads, bias_seqlen_q, bias_seqlen_kv = bias.shape - bias_batch = functools.reduce(operator.mul, bias_batch_shape) - fused_attn_helper = FusedAttnHelper( + has_fused_attn_kernel = is_fused_attn_kernel_available( # This needs to be fixed: TE-Jax has historically correlated training mode # with deterministic mode. - is_training=not deterministic, - batch_size=batch_size, - q_dtype=input_dtype, + not deterministic, + input_dtype, # self._assert_dtypes enforces Q, K, V, bias to have the same dtype, so # using input_dtype as kv dtype is sufficient. - kv_dtype=input_dtype, - qkv_layout=kernel_qkv_layout, - attn_bias_type=attn_bias_type, - attn_mask_type=attn_mask_type, - softmax_type=softmax_type, - dropout_probability=self.attention_dropout, - q_num_heads=self.num_attention_heads, - kv_num_heads=self.num_gqa_groups, - q_max_seqlen=seqlen_q, - kv_max_seqlen=seqlen_kv, - head_dim_qk=head_dim_qk, - head_dim_v=head_dim_v, - window_size=(-1, -1) if self.window_size is None else self.window_size, + input_dtype, + kernel_qkv_layout, + attn_bias_type, + attn_mask_type, + softmax_type, + self.attention_dropout, + self.num_attention_heads, + self.num_gqa_groups, + seqlen_q, + seqlen_kv, + head_dim_qk, + head_dim_v, + self.window_size, return_max_logit=self.return_max_logit, - bottom_right_diagonal=attn_mask_type.is_bottom_right(), - bias_batch=bias_batch, - bias_heads=bias_heads, - bias_seqlen_q=bias_seqlen_q, - bias_seqlen_kv=bias_seqlen_kv, - max_segments_per_seq=self.max_segments_per_seq, ) - fused_attn_backend, _ = fused_attn_helper.get_fused_attn_backend() - has_fused_attn_kernel = fused_attn_backend != NVTE_Fused_Attn_Backend.NVTE_No_Backend if score_mod_requested and not has_fused_attn_kernel: raise ValueError( "score_mod requires fused attention, but no fused attention kernel is available." @@ -862,9 +859,12 @@ def __call__( if enable_fused_attn and not has_fused_attn_kernel: warnings.warn( - "Falling back to the unfused attention backend as fused attention does not support" - " this config. Set NVTE_DEBUG=1 and NVTE_DEBUG_LEVEL=2 to see the detailed" - " rejection reason.\n" + "Fused attention is not enabled because there is no available kernel.\n" + "Fall back to the unfused attention.\n" + "Please try to update the cuDNN and TE to the latest version.\n" + f"{qkv_layout=}\n{attn_bias_type=}\n{attn_mask_type=}\n" + f"{self.attention_dropout=}\n{self.num_attention_heads=}\n{self.window_size=}\n" + f"{self.num_gqa_groups=}\n{seqlen_q=}\n{seqlen_kv=}\n{head_dim_qk=}\n{head_dim_v=}\n" ) dropout_rng = None @@ -951,6 +951,7 @@ def __call__( score_mod_bprop=self.score_mod_bprop, score_mod_requested=score_mod_requested, return_max_logit=self.return_max_logit, + bottom_right_diagonal=self.bottom_right_diagonal, )( query, key, diff --git a/transformer_engine/jax/quantize/helper.py b/transformer_engine/jax/quantize/helper.py index 3a93af4a68d..0540e1edc2f 100644 --- a/transformer_engine/jax/quantize/helper.py +++ b/transformer_engine/jax/quantize/helper.py @@ -14,7 +14,7 @@ from enum import Enum import hashlib from typing import Optional, Tuple, Dict, Union, Sequence, Type, List -from functools import reduce +from functools import partial, reduce import operator import warnings @@ -58,6 +58,7 @@ "update_collections", "apply_padding_to_scale_inv", "remove_padding_from_scale_inv", + "swizzle_mxfp8_scale", "NVTE_FP8_COLLECTION_NAME", "TensorSource", ] @@ -69,6 +70,24 @@ NVTE_FP8_COLLECTION_NAME = "fp8_metas" +@partial(jax.jit, static_argnums=(1, 2)) +def swizzle_mxfp8_scale(scale_inv, flatten_axis, is_colwise): + """Convert a padded MXFP8 scale tensor to the F8_128x4 physical layout.""" + + original_shape = scale_inv.shape + shape_2d = ( + reduce(operator.mul, original_shape[:flatten_axis], 1), + reduce(operator.mul, original_shape[flatten_axis:], 1), + ) + if is_colwise: + scale_inv = jnp.transpose(scale_inv.reshape(shape_2d)) + cols, rows = shape_2d + else: + rows, cols = shape_2d + reshaped = scale_inv.reshape(rows // 128, 4, 32, cols // 4, 4) + return jnp.transpose(reshaped, (0, 3, 2, 1, 4)).reshape(original_shape) + + def _check_delayed_scaling_fp8_support(gpu_arch) -> Tuple[bool, str]: """Check if delayed scaling FP8 is supported on the given GPU architecture. diff --git a/transformer_engine/jax/quantize/quantizer.py b/transformer_engine/jax/quantize/quantizer.py index db56db935dc..85506475539 100644 --- a/transformer_engine/jax/quantize/quantizer.py +++ b/transformer_engine/jax/quantize/quantizer.py @@ -39,6 +39,7 @@ __all__ = [ "Quantizer", "QuantizerSet", + "AttentionQuantizerSet", "CurrentScaleQuantizer", "DelayedScaleQuantizer", "BlockScaleQuantizer", @@ -876,6 +877,29 @@ def tree_unflatten(cls, aux_data, children): return cls(*aux_data, *children) +@register_pytree_node_class +@dataclass +class AttentionQuantizerSet: + """Quantizers for the six independent FP8 dot-product-attention roles.""" + + qkv: Quantizer + s: Quantizer + o: Quantizer + do: Quantizer + dp: Quantizer + dqkv: Quantizer + + def tree_flatten(self): + """Flatten all quantizers so delayed-scaling state participates in autodiff.""" + + return (self.qkv, self.s, self.o, self.do, self.dp, self.dqkv), () + + @classmethod + def tree_unflatten(cls, aux_data, children): + del aux_data + return cls(*children) + + @register_pytree_node_class @dataclass class GroupedQuantizer(Quantizer): diff --git a/transformer_engine/pytorch/attention/dot_product_attention/_cudnn_backend.py b/transformer_engine/pytorch/attention/dot_product_attention/_cudnn_backend.py new file mode 100644 index 00000000000..d146752c125 --- /dev/null +++ b/transformer_engine/pytorch/attention/dot_product_attention/_cudnn_backend.py @@ -0,0 +1,127 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""cuDNN SDPA capability selection for the PyTorch frontend.""" + +from __future__ import annotations + +import warnings + +from transformer_engine.common.attention.cudnn import ( + FusedAttentionConfig, + check_f16_fused_attention_support, + check_fp8_fused_attention_support, + parse_attention_layout, +) +from transformer_engine.pytorch.constants import DType +from transformer_engine.pytorch.utils import ( + get_cudnn_version, + get_device_compute_capability, +) + + +def _dtype_name(dtype: DType) -> str: + if dtype == DType.kFloat16: + return "float16" + if dtype == DType.kBFloat16: + return "bfloat16" + if dtype == DType.kFloat8E4M3: + return "float8_e4m3" + if dtype == DType.kFloat8E5M2: + return "float8_e5m2" + return dtype.name + + +def get_fused_attn_backend( + is_training, + q_dtype, + kv_dtype, + qkv_layout, + bias_type, + attn_mask_type, + softmax_type, + dropout, + num_attn_heads, + num_gqa_groups, + max_seqlen_q, + max_seqlen_kv, + head_dim_qk, + head_dim_v, + window_size_left, + window_size_right, + return_max_logit, + cuda_graph, + deterministic, +): + """Return the Python cuDNN SDPA backend and rejection reason for a configuration.""" + + # Import lazily to avoid a circular import through dot_product_attention.utils. + from .cudnn_attention import FusedAttnBackend + + q_dtype = DType.cast(q_dtype) + kv_dtype = DType.cast(kv_dtype) + major, minor = get_device_compute_capability() + sm_arch = major * 10 + minor + cudnn_version_tuple = get_cudnn_version() + layout = parse_attention_layout(qkv_layout) + + fp8_dtype = q_dtype in (DType.kFloat8E4M3, DType.kFloat8E5M2) + if fp8_dtype: + support = check_fp8_fused_attention_support( + FusedAttentionConfig( + is_training=bool(is_training), + q_dtype=_dtype_name(q_dtype), + kv_dtype=_dtype_name(kv_dtype), + layout=layout, + bias_type=bias_type, + mask_type=attn_mask_type, + softmax_type=softmax_type, + dropout=float(dropout), + num_attn_heads=int(num_attn_heads), + num_gqa_groups=int(num_gqa_groups), + max_seqlen_q=int(max_seqlen_q), + max_seqlen_kv=int(max_seqlen_kv), + head_dim_qk=int(head_dim_qk), + head_dim_v=int(head_dim_v), + window_size=(int(window_size_left), int(window_size_right)), + return_max_logit=bool(return_max_logit), + cuda_graph=bool(cuda_graph), + deterministic=bool(deterministic), + cudnn_version=cudnn_version_tuple, + sm_arch=sm_arch, + ) + ) + if support.supported: + return FusedAttnBackend.FP8, "" + return FusedAttnBackend.No_Backend, support.reason + + support = check_f16_fused_attention_support( + FusedAttentionConfig( + is_training=bool(is_training), + q_dtype=_dtype_name(q_dtype), + kv_dtype=_dtype_name(kv_dtype), + layout=layout, + bias_type=bias_type, + mask_type=attn_mask_type, + softmax_type=softmax_type, + dropout=float(dropout), + num_attn_heads=int(num_attn_heads), + num_gqa_groups=int(num_gqa_groups), + max_seqlen_q=int(max_seqlen_q), + max_seqlen_kv=int(max_seqlen_kv), + head_dim_qk=int(head_dim_qk), + head_dim_v=int(head_dim_v), + window_size=(int(window_size_left), int(window_size_right)), + return_max_logit=bool(return_max_logit), + cuda_graph=bool(cuda_graph), + deterministic=bool(deterministic), + cudnn_version=cudnn_version_tuple, + sm_arch=sm_arch, + ) + ) + if support.warning is not None: + warnings.warn(support.warning) + if support.supported: + return FusedAttnBackend.F16_arbitrary_seqlen, "" + return FusedAttnBackend.No_Backend, support.reason diff --git a/transformer_engine/pytorch/attention/dot_product_attention/_cudnn_graph.py b/transformer_engine/pytorch/attention/dot_product_attention/_cudnn_graph.py new file mode 100644 index 00000000000..6d34b70042f --- /dev/null +++ b/transformer_engine/pytorch/attention/dot_product_attention/_cudnn_graph.py @@ -0,0 +1,199 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""Shared cuDNN Frontend Python graph runtime for PyTorch attention.""" + +from __future__ import annotations + +import threading +from collections.abc import Hashable +from dataclasses import dataclass +from typing import Any, Dict, Optional, Tuple + +import torch + +from transformer_engine.common.attention.cache_debug import ( + build_recorder, + record_event, + record_lookup, +) +from transformer_engine.common.cudnn_frontend import ( + build_cudnn_graph, + make_cudnn_graph, +) +from transformer_engine.common.cudnn_frontend import ( + import_cudnn_frontend as _import_cudnn_frontend, +) + +_thread_state = threading.local() + + +def import_cudnn_frontend(): + """Import the cuDNN Frontend Python package lazily. + + PyTorch FusedAttention is optional at import time, so importing Transformer + Engine must not eagerly initialize cuDNN or fail on CPU-only processes. + """ + + return _import_cudnn_frontend( + feature="PyTorch fused attention", + requirement="nvidia-cudnn-frontend>=1.28.0", + ) + + +def _device_key(device: torch.device) -> Tuple[str, Optional[int]]: + device = torch.device(device) + if device.type == "cuda" and device.index is None: + return (device.type, torch.cuda.current_device()) + return (device.type, device.index) + + +def current_stream_handle(device: torch.device): + """Return this thread's cuDNN handle, bound to PyTorch's current stream.""" + + device = torch.device(device) + if device.type != "cuda": + raise ValueError(f"cuDNN attention only supports CUDA tensors, got {device}.") + device = torch.device("cuda", _device_key(device)[1]) + cudnn = import_cudnn_frontend() + + handles = getattr(_thread_state, "handles", None) + if handles is None: + handles = {} + _thread_state.handles = handles + handle = handles.get(device) + with torch.cuda.device(device): + if handle is None: + handle = cudnn.create_handle() + handles[device] = handle + cudnn.set_stream( + handle=handle, + stream=torch.cuda.current_stream(device).cuda_stream, + ) + return handle + + +def torch_to_cudnn_dtype(dtype: torch.dtype): + """Map a PyTorch scalar dtype to a cuDNN Frontend dtype.""" + + cudnn = import_cudnn_frontend() + mapping = { + torch.float16: cudnn.data_type.HALF, + torch.bfloat16: cudnn.data_type.BFLOAT16, + torch.float32: cudnn.data_type.FLOAT, + torch.int32: cudnn.data_type.INT32, + torch.int64: cudnn.data_type.INT64, + torch.uint8: cudnn.data_type.UINT8, + torch.float8_e4m3fn: cudnn.data_type.FP8_E4M3, + torch.float8_e5m2: cudnn.data_type.FP8_E5M2, + } + try: + return mapping[dtype] + except KeyError as exc: + raise ValueError(f"Unsupported cuDNN graph tensor dtype {dtype}.") from exc + + +def make_graph(io_dtype: Any, device: torch.device, *, name: str): + """Create an SDPA graph using FP32 intermediate and compute types.""" + + cudnn = import_cudnn_frontend() + return make_cudnn_graph( + cudnn, + io_dtype, + name=name, + handle=current_stream_handle(device), + ) + + +def finalize_graph(graph, *, cache_site: Tuple[str, str]) -> int: + """Build a cuDNN graph and return its required workspace size.""" + + cudnn = import_cudnn_frontend() + return build_cudnn_graph( + cudnn, + graph, + description="attention", + debug_callback=build_recorder(*cache_site), + ) + + +@dataclass +class GraphEntry: + """Built graph plus named graph tensors.""" + + graph: Any + tensors: Dict[str, Any] + workspace_size: int + cache_site: Optional[Tuple[str, str]] = None + + def execute(self, variant_pack: Dict[Any, Any], device: torch.device) -> None: + """Execute the graph on PyTorch's current stream.""" + + if self.cache_site is not None: + record_event(*self.cache_site, "execute", device=_device_key(device)[1]) + # Workspaces are execution scratch. Keeping them in graph cache entries + # retains one potentially large allocation for every cached configuration. + workspace = torch.empty( + self.workspace_size, + dtype=torch.uint8, + device=device, + ) + with torch.cuda.device(device): + self.graph.execute( + variant_pack, + workspace, + handle=current_stream_handle(device), + ) + + +def graph_cache() -> Dict[Hashable, GraphEntry]: + """Return a thread-local graph cache.""" + + cache = getattr(_thread_state, "graph_cache", None) + if cache is None: + cache = {} + _thread_state.graph_cache = cache + return cache + + +def get_graph_entry(key: Hashable) -> Optional[GraphEntry]: + """Look up a graph in this thread's cache.""" + + entry = graph_cache().get(key) + cache_site = _cache_site(key) + if cache_site is not None: + record_lookup(*cache_site, hit=entry is not None, key=key) + return entry + + +def put_graph_entry(key: Hashable, entry: GraphEntry) -> GraphEntry: + """Insert and return a graph cache entry.""" + + graph_cache()[key] = entry + cache_site = _cache_site(key) + if cache_site is not None: + entry.cache_site = cache_site + record_event(*cache_site, "cache_graph") + return entry + + +def _cache_site(key: Hashable) -> Optional[Tuple[str, str]]: + """Extract a diagnostic build site from an attention graph cache key.""" + + if not isinstance(key, tuple) or not key or not isinstance(key[0], str): + return None + try: + backend, direction = key[0].split("_", 1) + except ValueError: + return None + if backend not in ("f16", "fp8") or direction not in ("fwd", "bwd"): + return None + return backend, direction + + +def clear_graph_cache() -> None: + """Clear thread-local handles and graphs. Intended for tests.""" + + _thread_state.handles = {} + _thread_state.graph_cache = {} diff --git a/transformer_engine/pytorch/attention/dot_product_attention/backends.py b/transformer_engine/pytorch/attention/dot_product_attention/backends.py index 895881fced1..f3525db5420 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/backends.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/backends.py @@ -42,7 +42,7 @@ QKVLayouts, dist_group_type, ) -from transformer_engine.pytorch.cpp_extensions.fused_attn import ( +from transformer_engine.pytorch.attention.dot_product_attention.cudnn_attention import ( fused_attn_fwd, fused_attn_bwd, FusedAttnBackend, @@ -2088,29 +2088,12 @@ def backward(ctx, d_out, *_args): class FusedAttention(torch.nn.Module): - """Dot product attention using `cuDNN attention `_: + """Dot product attention using cuDNN attention: FusedAttnBackend["F16_arbitrary_seqlen"] cuDNN attention for FP16/BF16 with any sequence length. FusedAttnBackend["FP8"] - cuDNN attention for FP8 with any sequence length. It supports the following recipes, where - "Inputs", "Intermediates" and "Outputs" are in the format of "tensor: quantizer". The recipes - are implemented in transformer_engine.pytorch.cpp_extension.fused_attn.fused_attn_fwd and - transformer_engine.pytorch.cpp_extension.fused_attn.fused_attn_bwd. - - Direction Inputs Intermediates Outputs - DelayedScaling (DS) forward Q/K/V: DS S: DS O: DS - backward Q/K/V/O (from forward), dO: DS dP: DS dQ/dK/dV: DS - Float8CurrentScaling (CS) forward Q/K/V: CS S: DS O: F16 - backward Q/K/V (from forward), dO: CS, - O: F16 (or CS if NVTE_DPA_FP8CS_O_in_F16=0) dP: DS dQ/dK/dV: F16 - MXFP8BlockScaling (MXFP8) forward Q/K row, V col: MXFP8 S: None O: F16 - backward Q/K row+col, V row: MXFP8, - O/dO: F16, dO row+col: MXFP8 dP: None dQ/dK/dV: F16 - - For MXFP8, "row" and "col" are the quantization directions, which align with the contraction axes - of the matmuls that consume the tensor. For more details, please refer to - `How Scales Are Applied in MXFP8 Attention `_. + cuDNN attention for FP8 with any sequence length. """ def __init__( diff --git a/transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py b/transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py index 510d14ac638..e5ddebdac66 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py @@ -14,7 +14,7 @@ nvtx_range_push, get_device_compute_capability, ) -from transformer_engine.pytorch.cpp_extensions.fused_attn import ( +from transformer_engine.pytorch.attention.dot_product_attention.cudnn_attention import ( fused_attn_fwd, fused_attn_bwd, FusedAttnBackend, @@ -1121,7 +1121,7 @@ def cp_p2p_fwd_fused_attn( attn_bias = rest[0] if len(rest) > 0 else None if return_max_logit: - return out_per_step, softmax_lse_per_step, rng_states, attn_bias, *max_logit + return out_per_step, softmax_lse_per_step, rng_states, attn_bias, max_logit[0] return out_per_step, softmax_lse_per_step, rng_states, attn_bias, None @@ -5420,119 +5420,6 @@ def backward(ctx, dout, *_args): ) -def cp_per_step_configs( - cp_comm_type, - cp_size, - cp_size_a2a, - *, - max_seqlen_q, - max_seqlen_kv, - num_tokens_q, - num_tokens_kv, - num_heads, - num_gqa_groups, - attn_mask_type, - window_size, - bottom_right_diagonal, -): - """Per-step attention configs a context-parallel run dispatches to its attention backend. - - CP runs attention in multiple steps, each with a distinct config (e.g. mask, and seqlens) - that differs from the single global config. This function returns the list of those distinct - per-step configs so `get_attention_backend` can check if the backend supports all of them. - """ - is_causal = "causal" in attn_mask_type - padding_or_no_mask = "padding" if "padding" in attn_mask_type else "no_mask" - window_left, window_right = window_size - - def config(mask, s_q, s_kv, heads, gqa, bottom_right, t_q, t_kv, window=None): - w_left, w_right = window if window is not None else (window_left, window_right) - return { - "attn_mask_type": mask, - "max_seqlen_q": s_q, - "max_seqlen_kv": s_kv, - "num_tokens_q": t_q, - "num_tokens_kv": t_kv, - "num_attn_heads": heads, - "num_gqa_groups": gqa, - "window_size_left": w_left, - "window_size_right": w_right, - "bottom_right_diagonal": bottom_right, - } - - if cp_comm_type == "a2a": - # split heads across the cp ranks - return [ - config( - attn_mask_type, - max_seqlen_q, - max_seqlen_kv, - num_heads // cp_size, - num_gqa_groups // cp_size, - bottom_right_diagonal, - num_tokens_q * cp_size, - num_tokens_kv * cp_size, - ) - ] - - if cp_comm_type == "all_gather": - # one short Q chunk vs a growing KV chunk; causal -> causal_bottom_right - s_q = max_seqlen_q // (2 * cp_size) - s_kv_chunk = max_seqlen_kv // (2 * cp_size) - mask, br = attn_mask_type, bottom_right_diagonal - if is_causal and "bottom_right" not in attn_mask_type: - mask, br = attn_mask_type + "_bottom_right", True - # Each step narrows max_seqlen_*, but the token counts it dispatches with are the - # rank's full Q tokens and the all-gathered KV tokens, unchanged across steps. - # Scaling them per step would key the probe's graph differently from the one the - # step looks up, and rebuild every graph this probes at execution time. - t_q = num_tokens_q - t_kv = num_tokens_kv * cp_size - # s_kv ranges from s_kv_chunk, i*s_kv_chunk, ..., max_seqlen_kv - # check a single chunk and the full KV - return [ - config(mask, s_q, s_kv, num_heads, num_gqa_groups, br, t_q, t_kv) - for s_kv in dict.fromkeys([s_kv_chunk, max_seqlen_kv]) - ] - - # p2p and a2a+p2p: split heads across the a2a subgroup, and ring over the p2p subgroup - p2p_size = cp_size // cp_size_a2a - heads = num_heads // cp_size_a2a - gqa = num_gqa_groups // cp_size_a2a - r_q = max_seqlen_q // p2p_size - r_kv = max_seqlen_kv // p2p_size - # The tensors handed to this rank already correspond to (r_q, r_kv), so the token counts - # need no rescaling here; they only follow the halving below. - t_q, t_kv = num_tokens_q, num_tokens_kv - if not is_causal: - return [config(attn_mask_type, r_q, r_kv, heads, gqa, bottom_right_diagonal, t_q, t_kv)] - return [ - config(attn_mask_type, r_q, r_kv, heads, gqa, bottom_right_diagonal, t_q, t_kv), # diagonal - config( - padding_or_no_mask, - r_q, - r_kv // 2, - heads, - gqa, - bottom_right_diagonal, - t_q, - t_kv // 2, - window=(-1, -1), - ), # lower-triangle - config( - padding_or_no_mask, - r_q // 2, - r_kv, - heads, - gqa, - bottom_right_diagonal, - t_q // 2, - t_kv, - window=(-1, -1), - ), # upper-triangle - ] - - def attn_forward_func_with_cp( is_training, q, diff --git a/transformer_engine/pytorch/attention/dot_product_attention/cudnn_attention.py b/transformer_engine/pytorch/attention/dot_product_attention/cudnn_attention.py new file mode 100644 index 00000000000..492cf4c1f91 --- /dev/null +++ b/transformer_engine/pytorch/attention/dot_product_attention/cudnn_attention.py @@ -0,0 +1,3072 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""PyTorch implementation of cuDNN-backed scaled dot-product attention. + +All cuDNN graph construction and execution in this module goes through the +``nvidia-cudnn-frontend`` Python API rather than a TE-common attention +implementation. +""" + +from __future__ import annotations + +import math +from enum import IntEnum +from typing import Any, Dict, List, Optional, Sequence, Tuple, Union + +import torch + +# The extension is used only to reserve PyTorch's graph-safe Philox state; +# graph construction, backend selection, and execution do not call TE common. +import transformer_engine_torch as tex + +from transformer_engine.common.attention.cudnn import cudnn_mask_options +from transformer_engine.common.attention.cudnn import ( + ragged_batch_bucket as _max_ragged_batch, +) +from transformer_engine.common.attention.cudnn import ( + ragged_token_bucket as _max_ragged_tokens, +) +from transformer_engine.common.attention.fp8 import ( + FP8AttentionGraphConfig, + attention_format_stride as _format_stride, + build_fp8_backward_operation, + build_fp8_forward_operation, + mxfp8_padded_sizes as _mxfp8_padded_sizes, +) +from transformer_engine.pytorch.constants import ( + DType, + FP8BwdTensorIdx, + FP8FwdTensorIdx, + TE_DType_To_Torch, +) +from transformer_engine.pytorch.quantized_tensor import QuantizedTensorStorage +from transformer_engine.pytorch.utils import get_cudnn_version +from transformer_engine.pytorch.tensor.float8_tensor import ( + Float8CurrentScalingQuantizer, + Float8Quantizer, +) +from transformer_engine.pytorch.tensor.mxfp8_tensor import MXFP8Quantizer +from transformer_engine.pytorch.tensor.storage.float8_tensor_storage import ( + Float8TensorStorage, +) +from transformer_engine.pytorch.tensor.storage.mxfp8_tensor_storage import ( + MXFP8TensorStorage, +) + +from ._cudnn_graph import ( + GraphEntry, + finalize_graph, + get_graph_entry, + import_cudnn_frontend, + make_graph, + put_graph_entry, + torch_to_cudnn_dtype, +) + +__all__ = [ + "FusedAttnBackend", + "fused_attn_fwd", + "fused_attn_bwd", + "META_QKV", + "META_DQKV", + "META_O", + "META_DO", + "META_S", + "META_DP", +] + + +class FusedAttnBackend(IntEnum): + """PyTorch cuDNN attention implementation families. + + The numeric values intentionally preserve the historical TE-common ABI so + cached/backend-selection state and external Python callers remain source + compatible after removal of the pybind enum. + """ + + No_Backend = -1 + F16_arbitrary_seqlen = 1 + FP8 = 2 + + @classmethod + def cast(cls, backend: Union["FusedAttnBackend", int, Any]) -> "FusedAttnBackend": + """Convert a legacy backend integer to a fused-attention backend.""" + + if isinstance(backend, cls): + return backend + return cls(int(backend)) + + +META_QKV = FP8FwdTensorIdx.GEMM1_OUTPUT +META_DQKV = FP8BwdTensorIdx.GRAD_OUTPUT1 +META_O = FP8FwdTensorIdx.GEMM2_INPUT +META_DO = FP8BwdTensorIdx.GRAD_INPUT2 +META_S = FP8FwdTensorIdx.GEMM3_OUTPUT +META_DP = FP8BwdTensorIdx.GRAD_INPUT3 + +_F16_RNG_ELTS_PER_THREAD = 16 +_FP8_THREADS_PER_CTA = 128 + + +def _is_float8_tensor(tensor: Any) -> bool: + return isinstance(tensor, Float8TensorStorage) + + +def _is_mxfp8_tensor(tensor: Any) -> bool: + return isinstance(tensor, MXFP8TensorStorage) + + +def _quantized_data(tensor: Any, *, columnwise: bool = False) -> torch.Tensor: + if _is_float8_tensor(tensor): + data = tensor._transpose if columnwise else tensor._data + elif _is_mxfp8_tensor(tensor): + data = tensor._columnwise_data if columnwise else tensor._rowwise_data + else: + data = tensor + if data is None: + orientation = "columnwise" if columnwise else "rowwise" + raise ValueError(f"Attention input has no {orientation} data buffer.") + return data + + +def _quantized_scale_inv(tensor: Any, *, columnwise: bool = False) -> torch.Tensor: + if _is_float8_tensor(tensor): + return tensor._scale_inv + if _is_mxfp8_tensor(tensor): + scale = tensor._columnwise_scale_inv if columnwise else tensor._rowwise_scale_inv + if scale is None: + orientation = "columnwise" if columnwise else "rowwise" + raise ValueError(f"MXFP8 attention input has no {orientation} scale-inverse buffer.") + return scale + raise TypeError(f"Expected an FP8 attention tensor, got {type(tensor).__name__}.") + + +def _fp8_cudnn_dtype(tensor: Any): + return torch_to_cudnn_dtype(TE_DType_To_Torch[DType.cast(tensor._fp8_dtype)]) + + +def _scalar_graph_tensor(graph, cudnn, name: str): + return graph.tensor( + name=name, + dim=(1, 1, 1, 1), + stride=(1, 1, 1, 1), + data_type=cudnn.data_type.FLOAT, + ) + + +def _constant_graph_tensor(graph, cudnn, name: str): + """Create a scalar graph input that callers bind to a constant value.""" + return _scalar_graph_tensor(graph, cudnn, name) + + +def _padded_sequence_lengths(cu_seqlens: torch.Tensor, batch: int) -> torch.Tensor: + lengths = _sequence_lengths(cu_seqlens) + if lengths.numel() == batch: + return lengths + padding = torch.zeros(batch - lengths.numel(), dtype=lengths.dtype, device=lengths.device) + return torch.cat((lengths, padding)) + + +def _element_ragged_offsets( + cu_seqlens_padded: torch.Tensor, + batch: int, + multiplier: int, +) -> torch.Tensor: + """Convert token offsets to padded int64 element offsets for legacy SDPA graphs.""" + + offsets = cu_seqlens_padded.to(dtype=torch.int64) + if offsets.numel() < batch + 1: + tail = offsets[-1:].expand(batch + 1 - offsets.numel()) + offsets = torch.cat((offsets, tail)) + return offsets * multiplier + + +def _make_mxfp8_scale_tensor( + graph, + cudnn, + *, + name: str, + batch: int, + heads: int, + seqlen: int, + dim: int, + tensor_format: str, +): + return graph.tensor( + name=name, + dim=(batch, heads, seqlen, dim), + stride=_format_stride(batch, heads, seqlen, dim, tensor_format), + data_type=cudnn.data_type.FP8_E8M0, + ).set_reordering_type(cudnn.tensor_reordering.F8_128x4) + + +def _make_float8_output(quantizer, shape, fake_dtype, device): + data = torch.empty(shape, dtype=torch.uint8, device=device) + return quantizer.create_tensor_from_data( + data, + fake_dtype=fake_dtype, + internal=bool(getattr(quantizer, "internal", False)), + ) + + +def _allocate_fp8_kernel_output(quantizer, shape, fake_dtype, device): + """Allocate the FP8 SDPA output and any hidden amax buffer. + + Delayed scaling asks cuDNN to produce FP8 directly. Current scaling and + MXFP8 preserve the historical attention contract and ask cuDNN for a + high-precision output, which the caller may quantize afterwards. + """ + if isinstance(quantizer, Float8Quantizer): + return _make_float8_output(quantizer, shape, fake_dtype, device), quantizer.amax + if isinstance(quantizer, Float8CurrentScalingQuantizer): + return torch.empty(shape, dtype=fake_dtype, device=device), torch.zeros( + 1, dtype=torch.float32, device=device + ) + if isinstance(quantizer, MXFP8Quantizer): + return torch.empty(shape, dtype=fake_dtype, device=device), None + raise TypeError(f"Unsupported FP8 attention output quantizer {type(quantizer).__name__}.") + + +def _format_from_layout_component(component: str) -> str: + return "".join(char for char in component if char.isalpha()) + + +def _q_kv_formats(qkv_layout: str) -> Tuple[str, str]: + layout = qkv_layout.removeprefix("paged_kv_") + components = layout.split("_") + q_format = _format_from_layout_component(components[0]) + kv_format = _format_from_layout_component(components[-1]) if len(components) > 1 else q_format + return q_format, kv_format + + +def _kv_ragged_offsets_source( + qkv_layout: str, + cu_seqlens_q_padded: torch.Tensor, + cu_seqlens_kv_padded: torch.Tensor, +) -> torch.Tensor: + """Choose the token offsets backing K/V for a ragged storage layout.""" + + layout = qkv_layout.removeprefix("paged_kv_") + if len(layout.split("_")) == 1 and "3" in layout: + return cu_seqlens_q_padded + return cu_seqlens_kv_padded + + +def _is_paged_layout(qkv_layout: str) -> bool: + return qkv_layout.startswith("paged_kv_") + + +def _make_page_table_graph_tensor(graph, page_table: torch.Tensor, *, batch: int, name: str): + """Describe a physical ``[batch, pages]`` table in cuDNN's logical layout.""" + + if page_table.ndim != 2: + raise ValueError( + f"Paged attention expects a 2D page table, got shape {tuple(page_table.shape)}." + ) + batch_stride, page_stride = page_table.stride() + return graph.tensor( + name=name, + dim=(batch, 1, page_table.shape[1], 1), + stride=(batch_stride, batch_stride, page_stride, page_stride), + data_type=page_table.dtype, + ) + + +def _tensor_metadata(tensor: Optional[torch.Tensor]) -> Optional[Tuple[Any, ...]]: + if tensor is None: + return None + return ( + tuple(tensor.shape), + tuple(tensor.stride()), + tensor.dtype, + tensor.device.type, + tensor.device.index, + ) + + +def _logical_bhsd_desc( + tensor: torch.Tensor, + tensor_format: str, + *, + batch: int, + max_seqlen: int, +) -> Tuple[Tuple[int, ...], Tuple[int, ...]]: + """Describe a physical TE attention tensor as logical cuDNN BHSD.""" + + shape = tuple(tensor.shape) + stride = tuple(tensor.stride()) + if tensor_format == "sbhd": + return ( + (batch, shape[2], shape[0], shape[3]), + (stride[1], stride[2], stride[0], stride[3]), + ) + if tensor_format == "bshd": + return ( + (batch, shape[2], shape[1], shape[3]), + (stride[0], stride[2], stride[1], stride[3]), + ) + if tensor_format == "bhsd": + return ((batch, shape[1], shape[2], shape[3]), stride) + if tensor_format == "thd": + # The ragged offset selects each batch's first token. The synthetic + # batch stride is only graph metadata; S/H/D use the real packed view. + return ( + (batch, shape[1], max_seqlen, shape[2]), + (max_seqlen * stride[0], stride[1], stride[0], stride[2]), + ) + raise ValueError(f"Unsupported attention tensor format {tensor_format!r}.") + + +def _make_bhsd_graph_tensor( + graph, + tensor: torch.Tensor, + tensor_format: str, + *, + batch: int, + max_seqlen: int, + data_type=None, + ragged_offset=None, + ragged_offset_multiplier: int = 1, + name: str, +): + dim, stride = _logical_bhsd_desc( + tensor, + tensor_format, + batch=batch, + max_seqlen=max_seqlen, + ) + graph_tensor = graph.tensor( + name=name, + dim=dim, + stride=stride, + data_type=data_type if data_type is not None else tensor.dtype, + ragged_offset=ragged_offset, + ragged_offset_multiplier=ragged_offset_multiplier, + ) + return graph_tensor + + +def _allocate_output( + q: torch.Tensor, + value_head_dim: int, + fake_dtype: torch.dtype, + fast_zero_fill: bool, +) -> torch.Tensor: + shape = (*q.shape[:-1], value_head_dim) + factory = torch.zeros if fast_zero_fill else torch.empty + return factory(shape, dtype=fake_dtype, device=q.device) + + +def _storage_span(tensor: torch.Tensor) -> int: + if tensor.numel() == 0: + return 0 + return 1 + sum((size - 1) * stride for size, stride in zip(tensor.shape, tensor.stride())) + + +def _allocate_grad_views( + inputs: Sequence[torch.Tensor], +) -> Tuple[torch.Tensor, ...]: + """Allocate zeroed gradients while preserving packed-QKV storage relationships.""" + + groups: Dict[Tuple[str, int, int], List[int]] = {} + for index, tensor in enumerate(inputs): + storage = tensor.untyped_storage() + key = (tensor.device.type, tensor.device.index or 0, storage.data_ptr()) + groups.setdefault(key, []).append(index) + + outputs: List[Optional[torch.Tensor]] = [None] * len(inputs) + for indices in groups.values(): + if len(indices) == 1: + inp = inputs[indices[0]] + out = torch.empty_strided(inp.shape, inp.stride(), dtype=inp.dtype, device=inp.device) + out.zero_() + outputs[indices[0]] = out + continue + + min_offset = min(inputs[index].storage_offset() for index in indices) + max_end = max( + inputs[index].storage_offset() + _storage_span(inputs[index]) for index in indices + ) + exemplar = inputs[indices[0]] + base = torch.zeros(max_end - min_offset, dtype=exemplar.dtype, device=exemplar.device) + for index in indices: + inp = inputs[index] + outputs[index] = torch.as_strided( + base, + size=inp.shape, + stride=inp.stride(), + storage_offset=inp.storage_offset() - min_offset, + ) + + return tuple(output for output in outputs if output is not None) + + +def _reserve_philox_state( + device: torch.device, + rng_gen: Optional[torch.Generator], + increment: int, +) -> torch.Tensor: + """Reserve a Philox counter range and return CUDA ``[seed, offset]``.""" + + device = torch.device(device) + helper = getattr(tex, "get_cudnn_attention_rng_state", None) + if helper is not None: + with torch.cuda.device(device): + return helper(rng_gen, increment) + # Development-tree compatibility before the local extension has been + # rebuilt. Installed packages always expose the graph-safe helper above. + if rng_gen is None: + index = device.index if device.index is not None else torch.cuda.current_device() + rng_gen = torch.cuda.default_generators[index] + seed = rng_gen.initial_seed() + offset = rng_gen.get_offset() + rng_gen.set_offset(offset + increment) + return torch.tensor((seed, offset), dtype=torch.int64, device=device) + + +def _mask_options( + cudnn, + attn_mask_type: str, + window_size: Tuple[int, int], + bottom_right_diagonal: bool, + max_seqlen_q: int, + max_seqlen_kv: int, +) -> Dict[str, Any]: + options = cudnn_mask_options( + causal=attn_mask_type in ("causal", "padding_causal"), + bottom_right=attn_mask_type in ("causal_bottom_right", "padding_causal_bottom_right"), + padding=attn_mask_type in ("padding", "padding_causal", "padding_causal_bottom_right"), + bottom_right_diagonal=bottom_right_diagonal, + window_size=window_size, + max_seqlen_q=max_seqlen_q, + max_seqlen_kv=max_seqlen_kv, + cudnn_version=get_cudnn_version(), + ) + options["diagonal_alignment"] = ( + cudnn.diagonal_alignment.BOTTOM_RIGHT + if options["diagonal_alignment"] == "bottom_right" + else cudnn.diagonal_alignment.TOP_LEFT + ) + return options + + +def _sequence_lengths(cu_seqlens: torch.Tensor) -> torch.Tensor: + return cu_seqlens[1:] - cu_seqlens[:-1] + + +def _ragged_offset_tensor( + graph, + cu_seqlens_padded: torch.Tensor, + *, + multiplier: int, + name: str, + length: Optional[int] = None, + data_type=None, +): + return ( + graph.tensor( + name=name, + dim=(cu_seqlens_padded.numel() if length is None else length, 1, 1, 1), + stride=(1, 1, 1, 1), + data_type=cu_seqlens_padded.dtype if data_type is None else data_type, + is_pass_by_value=False, + ), + multiplier, + ) + + +def _stats_layout( + *, + batch: int, + heads: int, + max_seqlen_q: int, + total_tokens_q: int, + ragged: bool, +) -> Tuple[Tuple[int, ...], Tuple[int, ...], Tuple[int, ...]]: + if ragged: + physical_shape = (total_tokens_q, heads, 1) + logical_dim = (batch, heads, max_seqlen_q, 1) + logical_stride = (heads * max_seqlen_q, 1, heads, 1) + else: + physical_shape = logical_dim = (batch, heads, max_seqlen_q, 1) + logical_stride = (heads * max_seqlen_q, max_seqlen_q, 1, 1) + return physical_shape, logical_dim, logical_stride + + +def _f16_fwd_key(**kwargs) -> Tuple[Any, ...]: + return ("f16_fwd",) + tuple(kwargs.items()) + + +def _build_f16_fwd_graph( + *, + is_training: bool, + max_seqlen_q: int, + max_seqlen_kv: int, + cu_seqlens_q: torch.Tensor, + cu_seqlens_kv: torch.Tensor, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + output: torch.Tensor, + stats: Optional[torch.Tensor], + max_scores: Optional[torch.Tensor], + attn_bias: Optional[torch.Tensor], + cu_seqlens_q_padded: torch.Tensor, + cu_seqlens_kv_padded: torch.Tensor, + page_table_k: Optional[torch.Tensor], + page_table_v: Optional[torch.Tensor], + softmax_offset: Optional[torch.Tensor], + attn_scale: float, + dropout: float, + qkv_layout: str, + o_format: str, + attn_bias_type: str, + attn_mask_type: str, + softmax_type: str, + window_size: Tuple[int, int], + bottom_right_diagonal: bool, +) -> GraphEntry: + cudnn = import_cudnn_frontend() + graph = make_graph(torch_to_cudnn_dtype(q.dtype), q.device, name="te_fused_attention_fwd") + q_format, kv_format = _q_kv_formats(qkv_layout) + batch = cu_seqlens_q.numel() - 1 + is_ragged_q = q_format == "thd" + is_ragged_kv = kv_format == "thd" + use_ragged_stats = is_ragged_q and cudnn.backend_version() >= 90600 + use_token_buckets = cudnn.backend_version() >= 90600 and torch.cuda.get_device_capability( + q.device + ) != (12, 0) + use_direct_offsets = cudnn.backend_version() >= 92400 and dropout == 0.0 + use_legacy_offsets = (is_ragged_q or is_ragged_kv) and not use_direct_offsets + graph_batch = _max_ragged_batch(batch) if use_legacy_offsets and use_token_buckets else batch + if not use_token_buckets: + use_ragged_stats = False + graph_seqlen_q = ( + _max_ragged_tokens(q.shape[0]) if is_ragged_q and use_token_buckets else max_seqlen_q + ) + graph_seqlen_kv = ( + _max_ragged_tokens(k.shape[0]) if is_ragged_kv and use_token_buckets else max_seqlen_kv + ) + + tensors: Dict[str, Any] = {} + tensors["_legacy_offsets"] = use_legacy_offsets + tensors["_graph_batch"] = graph_batch + offset_q = offset_o = offset_k = offset_v = offset_stats = None + if is_ragged_q: + offset_q, q_mult = _ragged_offset_tensor( + graph, + cu_seqlens_q_padded, + multiplier=1 if use_legacy_offsets else q.stride(0), + name="offset_q", + length=graph_batch + 1, + data_type=torch.int64 if use_legacy_offsets else None, + ) + offset_o, o_mult = _ragged_offset_tensor( + graph, + cu_seqlens_q_padded, + multiplier=1 if use_legacy_offsets else output.stride(0), + name="offset_o", + length=graph_batch + 1, + data_type=torch.int64 if use_legacy_offsets else None, + ) + tensors["offset_q"] = offset_q + tensors["offset_o"] = offset_o + else: + q_mult = o_mult = 1 + if is_ragged_kv: + offset_k, k_mult = _ragged_offset_tensor( + graph, + cu_seqlens_kv_padded, + multiplier=1 if use_legacy_offsets else k.stride(0), + name="offset_k", + length=graph_batch + 1, + data_type=torch.int64 if use_legacy_offsets else None, + ) + offset_v, v_mult = _ragged_offset_tensor( + graph, + cu_seqlens_kv_padded, + multiplier=1 if use_legacy_offsets else v.stride(0), + name="offset_v", + length=graph_batch + 1, + data_type=torch.int64 if use_legacy_offsets else None, + ) + tensors["offset_k"] = offset_k + tensors["offset_v"] = offset_v + else: + k_mult = v_mult = 1 + + q_t = _make_bhsd_graph_tensor( + graph, + q, + q_format, + batch=graph_batch, + max_seqlen=graph_seqlen_q, + ragged_offset=offset_q, + ragged_offset_multiplier=q_mult, + name="Q", + ) + + if _is_paged_layout(qkv_layout): + if kv_format == "bshd": + num_pages_k, page_size_k = k.shape[0], k.shape[1] + num_pages_v, page_size_v = v.shape[0], v.shape[1] + elif kv_format == "sbhd": + page_size_k, num_pages_k = k.shape[0], k.shape[1] + page_size_v, num_pages_v = v.shape[0], v.shape[1] + else: + raise ValueError(f"Paged attention does not support KV format {kv_format}.") + k_batch, k_seqlen = num_pages_k, page_size_k + v_batch, v_seqlen = num_pages_v, page_size_v + else: + k_batch = v_batch = graph_batch + k_seqlen = v_seqlen = max_seqlen_kv + + k_t = _make_bhsd_graph_tensor( + graph, + k, + kv_format, + batch=k_batch, + max_seqlen=graph_seqlen_kv if is_ragged_kv else k_seqlen, + ragged_offset=offset_k, + ragged_offset_multiplier=k_mult, + name="K", + ) + v_t = _make_bhsd_graph_tensor( + graph, + v, + kv_format, + batch=v_batch, + max_seqlen=graph_seqlen_kv if is_ragged_kv else v_seqlen, + ragged_offset=offset_v, + ragged_offset_multiplier=v_mult, + name="V", + ) + tensors.update(Q=q_t, K=k_t, V=v_t) + + options = _mask_options( + cudnn, + attn_mask_type, + window_size, + bottom_right_diagonal, + max_seqlen_q, + max_seqlen_kv, + ) + is_padding = options.pop("is_padding") + options.update( + generate_stats=True, + attn_scale=float(attn_scale), + use_padding_mask=is_padding, + use_alibi_mask=attn_bias_type == "alibi", + ) + + if attn_bias_type == "post_scale_bias": + bias_t = graph.tensor_like(attn_bias, name="Bias") + tensors["Bias"] = bias_t + options["bias"] = bias_t + + if is_padding: + seq_q = _padded_sequence_lengths(cu_seqlens_q, graph_batch) + seq_kv = _padded_sequence_lengths(cu_seqlens_kv, graph_batch) + seq_q_t = graph.tensor_like(seq_q, name="seq_len_q") + seq_kv_t = graph.tensor_like(seq_kv, name="seq_len_kv") + tensors["seq_len_q"] = seq_q_t + tensors["seq_len_kv"] = seq_kv_t + options["seq_len_q"] = seq_q_t + options["seq_len_kv"] = seq_kv_t + + if page_table_k is not None: + page_k_t = _make_page_table_graph_tensor( + graph, page_table_k, batch=graph_batch, name="page_table_k" + ) + page_v_t = _make_page_table_graph_tensor( + graph, page_table_v, batch=graph_batch, name="page_table_v" + ) + tensors["page_table_k"] = page_k_t + tensors["page_table_v"] = page_v_t + options["paged_attention_k_table"] = page_k_t + options["paged_attention_v_table"] = page_v_t + options["paged_attention_max_seq_len_kv"] = max_seqlen_kv + + if is_training and dropout != 0.0: + seed_t = graph.tensor( + name="dropout_seed", + dim=(1, 1, 1, 1), + stride=(1, 1, 1, 1), + data_type=cudnn.data_type.INT64, + ) + offset_t = graph.tensor( + name="dropout_offset", + dim=(1, 1, 1, 1), + stride=(1, 1, 1, 1), + data_type=cudnn.data_type.INT64, + ) + tensors["dropout_seed"] = seed_t + tensors["dropout_offset"] = offset_t + options["dropout"] = (float(dropout), seed_t, offset_t) + + if softmax_type != "vanilla": + if softmax_offset is None: + raise ValueError(f"softmax_type={softmax_type!r} requires softmax_offset.") + softmax_offset_t = graph.tensor_like(softmax_offset, name="softmax_offset") + tensors["softmax_offset"] = softmax_offset_t + options["sink_token"] = softmax_offset_t + + if max_scores is not None: + _, max_dim, max_stride = _stats_layout( + batch=graph_batch, + heads=q_t.get_dim()[1], + max_seqlen_q=graph_seqlen_q, + total_tokens_q=q.shape[0] if is_ragged_q else batch * max_seqlen_q, + ragged=use_ragged_stats, + ) + if use_ragged_stats: + offset_stats, stats_mult = _ragged_offset_tensor( + graph, + cu_seqlens_q_padded, + multiplier=1 if use_legacy_offsets else q_t.get_dim()[1], + name="offset_stats", + length=graph_batch + 1, + data_type=torch.int64 if use_legacy_offsets else None, + ) + tensors["offset_stats"] = offset_stats + else: + stats_mult = 1 + max_t = graph.tensor( + name="Max", + dim=max_dim, + stride=max_stride, + data_type=cudnn.data_type.FLOAT, + ragged_offset=offset_stats, + ragged_offset_multiplier=stats_mult, + ).set_output(True) + tensors["Max"] = max_t + options["score_max"] = max_t + + output_t, stats_t = graph.sdpa(name="te_sdpa", q=q_t, k=k_t, v=v_t, **options) + output_dim, output_stride = _logical_bhsd_desc( + output, + o_format, + batch=graph_batch, + max_seqlen=graph_seqlen_q, + ) + output_t.set_output(True).set_dim(output_dim).set_stride(output_stride) + if is_ragged_q: + output_t.set_ragged_offset(offset_o).set_ragged_offset_multiplier(o_mult) + tensors["O"] = output_t + + assert stats is not None + _, stats_dim, stats_stride = _stats_layout( + batch=graph_batch, + heads=output_dim[1], + max_seqlen_q=graph_seqlen_q, + total_tokens_q=q.shape[0] if is_ragged_q else batch * max_seqlen_q, + ragged=use_ragged_stats, + ) + if use_ragged_stats and offset_stats is None: + offset_stats, stats_mult = _ragged_offset_tensor( + graph, + cu_seqlens_q_padded, + multiplier=1 if use_legacy_offsets else output_dim[1], + name="offset_stats", + length=graph_batch + 1, + data_type=torch.int64 if use_legacy_offsets else None, + ) + tensors["offset_stats"] = offset_stats + stats_t.set_output(True).set_data_type(cudnn.data_type.FLOAT).set_dim(stats_dim).set_stride( + stats_stride + ) + if use_ragged_stats: + stats_t.set_ragged_offset(offset_stats).set_ragged_offset_multiplier(stats_mult) + tensors["Stats"] = stats_t + + return GraphEntry( + graph=graph, + tensors=tensors, + workspace_size=finalize_graph(graph, cache_site=("f16", "fwd")), + ) + + +def _f16_forward( + is_training: bool, + max_seqlen_q: int, + max_seqlen_kv: int, + cu_seqlens_q: torch.Tensor, + cu_seqlens_kv: torch.Tensor, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + fake_dtype: torch.dtype, + attn_bias: Optional[torch.Tensor], + cu_seqlens_q_padded: Optional[torch.Tensor], + cu_seqlens_kv_padded: Optional[torch.Tensor], + page_table_k: Optional[torch.Tensor], + page_table_v: Optional[torch.Tensor], + attn_scale: float, + dropout: float, + fast_zero_fill: bool, + qkv_layout: str, + o_format: str, + attn_bias_type: str, + attn_mask_type: str, + softmax_type: str, + window_size: Tuple[int, int], + bottom_right_diagonal: bool, + rng_gen: Optional[torch.Generator], + softmax_offset: Optional[torch.Tensor], + return_max_logit: bool, +) -> Tuple[torch.Tensor, List[torch.Tensor], Optional[torch.Tensor]]: + q_format, kv_format = _q_kv_formats(qkv_layout) + batch = cu_seqlens_q.numel() - 1 + heads = q.shape[1] if q_format == "bhsd" else q.shape[-2] + total_tokens_q = q.shape[0] if q_format == "thd" else batch * max_seqlen_q + cudnn = import_cudnn_frontend() + use_token_buckets = cudnn.backend_version() >= 90600 and torch.cuda.get_device_capability( + q.device + ) != (12, 0) + use_direct_offsets = cudnn.backend_version() >= 92400 and dropout == 0.0 + use_legacy_offsets = (q_format == "thd" or kv_format == "thd") and not use_direct_offsets + graph_batch = _max_ragged_batch(batch) if use_legacy_offsets and use_token_buckets else batch + ragged_stats = ( + q_format == "thd" + and cudnn.backend_version() >= 90600 + and torch.cuda.get_device_capability(q.device) != (12, 0) + ) + output_shape = ( + (q.shape[0], heads, v.shape[-1]) + if o_format == "thd" + else _fp8_output_shape(batch, max_seqlen_q, heads, v.shape[-1], o_format) + ) + output_factory = torch.zeros if fast_zero_fill or o_format == "thd" else torch.empty + output = output_factory(output_shape, dtype=fake_dtype, device=q.device) + stats_shape, _, _ = _stats_layout( + batch=batch, + heads=heads, + max_seqlen_q=max_seqlen_q, + total_tokens_q=total_tokens_q, + ragged=ragged_stats, + ) + stats = torch.empty(stats_shape, dtype=torch.float32, device=q.device) + max_scores = ( + torch.empty(stats_shape, dtype=torch.float32, device=q.device) if return_max_logit else None + ) + rng_state = _reserve_philox_state(q.device, rng_gen, _F16_RNG_ELTS_PER_THREAD) + cu_seqlens_q_padded = cu_seqlens_q if cu_seqlens_q_padded is None else cu_seqlens_q_padded + cu_seqlens_kv_padded = cu_seqlens_kv if cu_seqlens_kv_padded is None else cu_seqlens_kv_padded + + key = _f16_fwd_key( + is_training=is_training, + max_seqlen_q=max_seqlen_q, + max_seqlen_kv=max_seqlen_kv, + graph_batch=graph_batch, + q=_tensor_metadata(q), + k=_tensor_metadata(k), + v=_tensor_metadata(v), + output=_tensor_metadata(output), + stats=_tensor_metadata(stats), + bias=_tensor_metadata(attn_bias), + qkv_layout=qkv_layout, + o_format=o_format, + attn_scale=float(attn_scale), + dropout=float(dropout), + attn_bias_type=attn_bias_type, + attn_mask_type=attn_mask_type, + softmax_type=softmax_type, + window_size=tuple(window_size), + bottom_right_diagonal=bottom_right_diagonal, + return_max_logit=return_max_logit, + page_table_k=_tensor_metadata(page_table_k), + page_table_v=_tensor_metadata(page_table_v), + ) + entry = get_graph_entry(key) + if entry is None: + if torch.cuda.is_current_stream_capturing(): + raise RuntimeError("cuDNN attention graph must be built before CUDA graph capture.") + entry = _build_f16_fwd_graph( + is_training=is_training, + max_seqlen_q=max_seqlen_q, + max_seqlen_kv=max_seqlen_kv, + cu_seqlens_q=cu_seqlens_q, + cu_seqlens_kv=cu_seqlens_kv, + q=q, + k=k, + v=v, + output=output, + stats=stats, + max_scores=max_scores, + attn_bias=attn_bias, + cu_seqlens_q_padded=cu_seqlens_q_padded, + cu_seqlens_kv_padded=cu_seqlens_kv_padded, + page_table_k=page_table_k, + page_table_v=page_table_v, + softmax_offset=softmax_offset, + attn_scale=attn_scale, + dropout=dropout, + qkv_layout=qkv_layout, + o_format=o_format, + attn_bias_type=attn_bias_type, + attn_mask_type=attn_mask_type, + softmax_type=softmax_type, + window_size=window_size, + bottom_right_diagonal=bottom_right_diagonal, + ) + put_graph_entry(key, entry) + + tensors = entry.tensors + variant_pack: Dict[Any, Any] = { + tensors["Q"]: q, + tensors["K"]: k, + tensors["V"]: v, + tensors["O"]: output, + } + variant_pack[tensors["Stats"]] = stats + if attn_bias_type == "post_scale_bias": + variant_pack[tensors["Bias"]] = attn_bias + legacy_offsets = tensors["_legacy_offsets"] + graph_batch = tensors["_graph_batch"] + if "seq_len_q" in tensors: + variant_pack[tensors["seq_len_q"]] = _padded_sequence_lengths(cu_seqlens_q, graph_batch) + variant_pack[tensors["seq_len_kv"]] = _padded_sequence_lengths(cu_seqlens_kv, graph_batch) + if "offset_q" in tensors: + if legacy_offsets: + variant_pack[tensors["offset_q"]] = _element_ragged_offsets( + cu_seqlens_q_padded, graph_batch, q.stride(0) + ) + variant_pack[tensors["offset_o"]] = _element_ragged_offsets( + cu_seqlens_q_padded, graph_batch, output.stride(0) + ) + else: + variant_pack[tensors["offset_q"]] = cu_seqlens_q_padded + variant_pack[tensors["offset_o"]] = cu_seqlens_q_padded + if "offset_k" in tensors: + if legacy_offsets: + variant_pack[tensors["offset_k"]] = _element_ragged_offsets( + cu_seqlens_kv_padded, graph_batch, k.stride(0) + ) + variant_pack[tensors["offset_v"]] = _element_ragged_offsets( + cu_seqlens_kv_padded, graph_batch, v.stride(0) + ) + else: + variant_pack[tensors["offset_k"]] = cu_seqlens_kv_padded + variant_pack[tensors["offset_v"]] = cu_seqlens_kv_padded + if "offset_stats" in tensors: + variant_pack[tensors["offset_stats"]] = ( + _element_ragged_offsets(cu_seqlens_q_padded, graph_batch, heads) + if legacy_offsets + else cu_seqlens_q_padded + ) + if "page_table_k" in tensors: + variant_pack[tensors["page_table_k"]] = page_table_k + variant_pack[tensors["page_table_v"]] = page_table_v + if "dropout_seed" in tensors: + variant_pack[tensors["dropout_seed"]] = rng_state[:1] + variant_pack[tensors["dropout_offset"]] = rng_state[1:] + if "softmax_offset" in tensors: + variant_pack[tensors["softmax_offset"]] = softmax_offset + if "Max" in tensors: + variant_pack[tensors["Max"]] = max_scores + entry.execute(variant_pack, q.device) + + aux: List[torch.Tensor] = [stats, rng_state] + if is_training: + if attn_bias_type not in ("no_bias", "alibi"): + aux.append(attn_bias) + if softmax_type != "vanilla": + aux.append(softmax_offset) + + max_logit = None + if return_max_logit: + max_scores_for_reduce = max_scores + if q_format == "thd": + if max_scores.ndim == 4: + seqlens_q = _sequence_lengths(cu_seqlens_q).to(device=max_scores.device) + sq_idx = torch.arange(max_scores.shape[2], device=max_scores.device).view( + 1, 1, -1, 1 + ) + valid = sq_idx < seqlens_q.view(-1, 1, 1, 1) + max_scores_for_reduce = max_scores.masked_fill(~valid, float("-inf")) + elif max_scores.ndim == 3: + seqlens_q = _sequence_lengths(cu_seqlens_q).to(device=max_scores.device) + total_tokens = max_scores.shape[0] + starts = cu_seqlens_q_padded[:-1].to(device=max_scores.device) + ends = (starts + seqlens_q).clamp(max=total_tokens) + delta = torch.zeros(total_tokens + 1, dtype=torch.int32, device=max_scores.device) + updates = torch.ones_like(starts, dtype=torch.int32) + delta.scatter_add_(0, starts.clamp(max=total_tokens), updates) + delta.scatter_add_(0, ends, -updates) + valid = delta[:-1].cumsum(0) > 0 + max_scores_for_reduce = max_scores.masked_fill(~valid.view(-1, 1, 1), float("-inf")) + reduce_dims = (0, 2) if max_scores_for_reduce.ndim == 3 else (0, 2, 3) + max_logit = torch.amax(max_scores_for_reduce, dim=reduce_dims).to(output.dtype) + return output, aux, max_logit + + +def _fp8_output_shape( + batch: int, + seqlen: int, + heads: int, + dim: int, + tensor_format: str, + *, + total_tokens: Optional[int] = None, +): + if tensor_format == "bshd": + return (batch, seqlen, heads, dim) + if tensor_format == "sbhd": + return (seqlen, batch, heads, dim) + if tensor_format == "bhsd": + return (batch, heads, seqlen, dim) + if tensor_format == "thd": + if total_tokens is None: + raise ValueError("FP8 THD output allocation requires the total token count.") + return (total_tokens, heads, dim) + raise ValueError(f"FP8 attention does not support output format {tensor_format!r}.") + + +def _allocate_attention_grad_data( + *, + batch: int, + heads: int, + kv_heads: int, + max_seqlen_q: int, + max_seqlen_kv: int, + head_dim_qk: int, + head_dim_v: int, + dqkv_layout: str, + dtype: torch.dtype, + device: torch.device, + zero: bool, + total_tokens_q: Optional[int] = None, + total_tokens_kv: Optional[int] = None, +): + """Allocate dQ/dK/dV buffers with the requested packed storage layout.""" + + layout = dqkv_layout.removeprefix("paged_kv_") + components = layout.split("_") + q_format, kv_format = _q_kv_formats(layout) + q_shape = _fp8_output_shape( + batch, + max_seqlen_q, + heads, + head_dim_qk, + q_format, + total_tokens=total_tokens_q, + ) + k_shape = _fp8_output_shape( + batch, + max_seqlen_kv, + kv_heads, + head_dim_qk, + kv_format, + total_tokens=total_tokens_kv, + ) + v_shape = _fp8_output_shape( + batch, + max_seqlen_kv, + kv_heads, + head_dim_v, + kv_format, + total_tokens=total_tokens_kv, + ) + factory = torch.zeros if zero else torch.empty + + if len(components) == 1: + if head_dim_qk != head_dim_v or heads != kv_heads or max_seqlen_q != max_seqlen_kv: + raise ValueError(f"Packed QKV gradient layout {layout!r} requires matching Q/K/V.") + packed_dim = components[0].index("3") + packed_shape = list(q_shape) + packed_shape.insert(packed_dim, 3) + packed = factory(packed_shape, dtype=dtype, device=device) + return tuple(packed.select(packed_dim, index) for index in range(3)) + if len(components) == 2: + if head_dim_qk != head_dim_v: + raise ValueError(f"Packed KV gradient layout {layout!r} requires dQK == dV.") + q_out = factory(q_shape, dtype=dtype, device=device) + packed_dim = components[1].index("2") + packed_shape = list(k_shape) + packed_shape.insert(packed_dim, 2) + packed = factory(packed_shape, dtype=dtype, device=device) + return q_out, packed.select(packed_dim, 0), packed.select(packed_dim, 1) + return tuple( + factory(shape, dtype=dtype, device=device) for shape in (q_shape, k_shape, v_shape) + ) + + +def _wrap_float8_grad_outputs(quantizer, tensors, fake_dtype): + return tuple( + quantizer.create_tensor_from_data( + tensor, + fake_dtype=fake_dtype, + internal=bool(getattr(quantizer, "internal", False)), + ) + for tensor in tensors + ) + + +def _build_fp8_fwd_graph( + *, + max_seqlen_q, + max_seqlen_kv, + q, + k, + v, + output, + s_quantizer, + o_quantizer, + qkv_layout, + o_format, + qkv_scale_inv_format, + attn_scale, + dropout, + attn_mask_type, + softmax_type, + window_size, + bottom_right_diagonal, + softmax_offset, + cu_seqlens_q, + cu_seqlens_kv, + cu_seqlens_q_padded, + cu_seqlens_kv_padded, +): + cudnn = import_cudnn_frontend() + q_data = _quantized_data(q) + k_data = _quantized_data(k) + v_data = _quantized_data(v, columnwise=_is_mxfp8_tensor(v)) + output_data = _quantized_data(output) if _is_float8_tensor(output) else output + q_format, kv_format = _q_kv_formats(qkv_layout) + batch = cu_seqlens_q.numel() - 1 + heads = q_data.shape[-2] if q_format != "bhsd" else q_data.shape[1] + kv_heads = k_data.shape[-2] if kv_format != "bhsd" else k_data.shape[1] + d_qk = q_data.shape[-1] + d_v = v_data.shape[-1] + is_ragged_q = q_format == "thd" + is_ragged_kv = kv_format == "thd" + use_token_buckets = cudnn.backend_version() >= 90600 and torch.cuda.get_device_capability( + q.device + ) != (12, 0) + use_ragged_stats = is_ragged_q and use_token_buckets + use_legacy_offsets = is_ragged_q or is_ragged_kv + graph_batch = _max_ragged_batch(batch) if use_legacy_offsets and use_token_buckets else batch + graph_seqlen_q = ( + _max_ragged_tokens(q_data.shape[0]) if is_ragged_q and use_token_buckets else max_seqlen_q + ) + graph_seqlen_kv = ( + _max_ragged_tokens(k_data.shape[0]) if is_ragged_kv and use_token_buckets else max_seqlen_kv + ) + graph = make_graph(_fp8_cudnn_dtype(q), q.device, name="te_fp8_sdpa_fwd") + tensors: Dict[str, Any] = { + "_legacy_offsets": use_legacy_offsets, + "_graph_batch": graph_batch, + } + + offset_q = offset_o = offset_k = offset_v = offset_stats = None + if is_ragged_q: + offset_q, q_mult = _ragged_offset_tensor( + graph, + cu_seqlens_q_padded, + multiplier=1, + name="offset_q", + length=graph_batch + 1, + data_type=torch.int64, + ) + offset_o, o_mult = _ragged_offset_tensor( + graph, + cu_seqlens_q_padded, + multiplier=1, + name="offset_o", + length=graph_batch + 1, + data_type=torch.int64, + ) + tensors.update(offset_q=offset_q, offset_o=offset_o) + else: + q_mult = o_mult = 1 + if is_ragged_kv: + offset_k, k_mult = _ragged_offset_tensor( + graph, + cu_seqlens_kv_padded, + multiplier=1, + name="offset_k", + length=graph_batch + 1, + data_type=torch.int64, + ) + offset_v, v_mult = _ragged_offset_tensor( + graph, + cu_seqlens_kv_padded, + multiplier=1, + name="offset_v", + length=graph_batch + 1, + data_type=torch.int64, + ) + tensors.update(offset_k=offset_k, offset_v=offset_v) + else: + k_mult = v_mult = 1 + + q_t = _make_bhsd_graph_tensor( + graph, + q_data, + q_format, + batch=graph_batch, + max_seqlen=graph_seqlen_q, + data_type=_fp8_cudnn_dtype(q), + ragged_offset=offset_q, + ragged_offset_multiplier=q_mult, + name="Q", + ) + k_t = _make_bhsd_graph_tensor( + graph, + k_data, + kv_format, + batch=graph_batch, + max_seqlen=graph_seqlen_kv, + data_type=_fp8_cudnn_dtype(k), + ragged_offset=offset_k, + ragged_offset_multiplier=k_mult, + name="K", + ) + v_t = _make_bhsd_graph_tensor( + graph, + v_data, + kv_format, + batch=graph_batch, + max_seqlen=graph_seqlen_kv, + data_type=_fp8_cudnn_dtype(v), + ragged_offset=offset_v, + ragged_offset_multiplier=v_mult, + name="V", + ) + tensors.update(Q=q_t, K=k_t, V=v_t) + + options = _mask_options( + cudnn, + attn_mask_type, + window_size, + bottom_right_diagonal, + max_seqlen_q, + max_seqlen_kv, + ) + is_padding = options.pop("is_padding") + options.update( + generate_stats=True, + attn_scale=float(attn_scale), + use_padding_mask=is_padding, + ) + if is_padding: + seq_q = _padded_sequence_lengths(cu_seqlens_q, graph_batch) + seq_kv = _padded_sequence_lengths(cu_seqlens_kv, graph_batch) + seq_q_t = graph.tensor_like(seq_q, name="seq_len_q") + seq_kv_t = graph.tensor_like(seq_kv, name="seq_len_kv") + tensors.update(seq_len_q=seq_q_t, seq_len_kv=seq_kv_t) + options.update(seq_len_q=seq_q_t, seq_len_kv=seq_kv_t) + if dropout != 0.0: + seed_t = graph.tensor( + name="dropout_seed", + dim=(1, 1, 1, 1), + stride=(1, 1, 1, 1), + data_type=cudnn.data_type.INT64, + ) + offset_t = graph.tensor( + name="dropout_offset", + dim=(1, 1, 1, 1), + stride=(1, 1, 1, 1), + data_type=cudnn.data_type.INT64, + ) + tensors.update(dropout_seed=seed_t, dropout_offset=offset_t) + options["dropout"] = (float(dropout), seed_t, offset_t) + if softmax_type != "vanilla": + sink_t = graph.tensor_like(softmax_offset, name="softmax_offset") + tensors["softmax_offset"] = sink_t + options["sink_token"] = sink_t + + if _is_mxfp8_tensor(q): + # The MXFP8 binding has no padding-mask keyword. Avoid forwarding the + # false default through the generic graph capture layer. + options.pop("use_padding_mask", None) + if is_padding: + raise RuntimeError( + "The installed cuDNN Frontend Python MXFP8 graph API does not expose " + "padding sequence lengths." + ) + scale_format_q = qkv_scale_inv_format or q_format + scale_format_kv = qkv_scale_inv_format or kv_format + padded = _mxfp8_padded_sizes(max_seqlen_q, max_seqlen_kv, d_qk, d_v) + descale_q = _make_mxfp8_scale_tensor( + graph, + cudnn, + name="Descale_Q", + batch=batch, + heads=heads, + seqlen=padded["s_q_padded"], + dim=padded["d_qk_scale_padded"], + tensor_format=scale_format_q, + ) + descale_k = _make_mxfp8_scale_tensor( + graph, + cudnn, + name="Descale_K", + batch=batch, + heads=kv_heads, + seqlen=padded["s_kv_padded"], + dim=padded["d_qk_scale_padded"], + tensor_format=scale_format_kv, + ) + descale_v = _make_mxfp8_scale_tensor( + graph, + cudnn, + name="Descale_V", + batch=batch, + heads=kv_heads, + seqlen=padded["s_kv_scale_padded"], + dim=padded["d_v_padded"], + tensor_format=scale_format_kv, + ) + tensors.update(descale_q=descale_q, descale_k=descale_k, descale_v=descale_v) + op = build_fp8_forward_operation( + graph, + { + "q": q_t, + "k": k_t, + "v": v_t, + "descale_q": descale_q, + "descale_k": descale_k, + "descale_v": descale_v, + }, + options, + FP8AttentionGraphConfig("mxfp8", "te_sdpa_mxfp8"), + ) + output_t, stats_t, amax_o_t = op["output"], op["stats"], op["amax_o"] + amax_o_t.set_output(False).set_data_type(cudnn.data_type.FLOAT).set_dim( + (1, 1, 1, 1) + ).set_stride((1, 1, 1, 1)) + else: + if "diagonal_band_left_bound" in options: + options["left_bound"] = options.pop("diagonal_band_left_bound") + if "diagonal_band_right_bound" in options: + options["right_bound"] = options.pop("diagonal_band_right_bound") + descale_q = _scalar_graph_tensor(graph, cudnn, "Descale_Q") + descale_k = _scalar_graph_tensor(graph, cudnn, "Descale_K") + descale_v = _scalar_graph_tensor(graph, cudnn, "Descale_V") + tensors.update(descale_q=descale_q, descale_k=descale_k, descale_v=descale_v) + if isinstance(s_quantizer, Float8Quantizer): + descale_s = _scalar_graph_tensor(graph, cudnn, "Descale_S") + scale_s = _scalar_graph_tensor(graph, cudnn, "Scale_S") + tensors.update(descale_s=descale_s, scale_s=scale_s) + else: + descale_s = _constant_graph_tensor(graph, cudnn, "Current_Descale_S") + scale_s = _constant_graph_tensor(graph, cudnn, "Current_Scale_S") + tensors.update(constant_descale_s=descale_s, constant_scale_s=scale_s) + if isinstance(o_quantizer, Float8Quantizer): + scale_o = _scalar_graph_tensor(graph, cudnn, "Scale_O") + tensors["scale_o"] = scale_o + else: + scale_o = _constant_graph_tensor(graph, cudnn, "Current_Scale_O") + tensors["constant_scale_o"] = scale_o + op = build_fp8_forward_operation( + graph, + { + "q": q_t, + "k": k_t, + "v": v_t, + "descale_q": descale_q, + "descale_k": descale_k, + "descale_v": descale_v, + "descale_s": descale_s, + "scale_s": scale_s, + "scale_o": scale_o, + }, + options, + FP8AttentionGraphConfig( + "delayed" if isinstance(o_quantizer, Float8Quantizer) else "current", + "te_sdpa_fp8", + ), + ) + output_t, stats_t = op["output"], op["stats"] + amax_s_t, amax_o_t = op["amax_s"], op["amax_o"] + amax_s_t.set_output(True).set_data_type(cudnn.data_type.FLOAT).set_dim( + (1, 1, 1, 1) + ).set_stride((1, 1, 1, 1)) + amax_o_t.set_output(True).set_data_type(cudnn.data_type.FLOAT).set_dim( + (1, 1, 1, 1) + ).set_stride((1, 1, 1, 1)) + tensors.update(amax_s=amax_s_t, amax_o=amax_o_t) + + output_dim, output_stride = _logical_bhsd_desc( + output_data, + o_format, + batch=graph_batch, + max_seqlen=graph_seqlen_q, + ) + output_t.set_output(True).set_dim(output_dim).set_stride(output_stride) + if is_ragged_q: + output_t.set_ragged_offset(offset_o).set_ragged_offset_multiplier(o_mult) + output_dtype = _fp8_cudnn_dtype(output) if _is_float8_tensor(output) else output.dtype + output_t.set_data_type(output_dtype) + _, stats_dim, stats_stride = _stats_layout( + batch=graph_batch, + heads=heads, + max_seqlen_q=graph_seqlen_q, + total_tokens_q=q_data.shape[0] if is_ragged_q else batch * max_seqlen_q, + ragged=use_ragged_stats, + ) + if use_ragged_stats: + offset_stats, stats_mult = _ragged_offset_tensor( + graph, + cu_seqlens_q_padded, + multiplier=1, + name="offset_stats", + length=graph_batch + 1, + data_type=torch.int64, + ) + tensors["offset_stats"] = offset_stats + else: + stats_mult = 1 + stats_t.set_output(True).set_data_type(cudnn.data_type.FLOAT).set_dim(stats_dim).set_stride( + stats_stride + ) + if use_ragged_stats: + stats_t.set_ragged_offset(offset_stats).set_ragged_offset_multiplier(stats_mult) + tensors.update(O=output_t, Stats=stats_t) + return GraphEntry( + graph=graph, + tensors=tensors, + workspace_size=finalize_graph(graph, cache_site=("fp8", "fwd")), + ) + + +def _fp8_forward( + is_training, + max_seqlen_q, + max_seqlen_kv, + cu_seqlens_q, + cu_seqlens_kv, + cu_seqlens_q_padded, + cu_seqlens_kv_padded, + q, + k, + v, + fake_dtype, + s_quantizer, + o_quantizer, + attn_scale, + dropout, + fast_zero_fill, + qkv_layout, + o_format, + qkv_scale_inv_format, + attn_mask_type, + softmax_type, + window_size, + bottom_right_diagonal, + rng_gen, + softmax_offset, +): + if not isinstance(q, QuantizedTensorStorage): + raise TypeError("The FP8 cuDNN attention backend requires quantized Q/K/V tensors.") + if _is_mxfp8_tensor(q) and "padding" in attn_mask_type: + # cuDNN Frontend 1.27's Python sdpa_mxfp8 forward binding omits the + # seq_len inputs that its C++ graph API exposes. Preserve functional + # padding semantics through the same Python graph API by dequantizing + # MXFP8 inputs and using the BF16/FP16 SDPA node for this configuration. + # Remove this fallback once the Python MXFP8 forward signature exposes + # padding sequence lengths. + q_hp, k_hp, v_hp = (tensor.dequantize(dtype=fake_dtype) for tensor in (q, k, v)) + output, aux, _ = _f16_forward( + is_training, + max_seqlen_q, + max_seqlen_kv, + cu_seqlens_q, + cu_seqlens_kv, + q_hp, + k_hp, + v_hp, + fake_dtype, + None, + None, + None, + None, + None, + attn_scale, + dropout, + fast_zero_fill, + qkv_layout, + o_format, + "no_bias", + attn_mask_type, + softmax_type, + window_size, + bottom_right_diagonal, + rng_gen, + softmax_offset, + False, + ) + return output, aux + del is_training + q_format, _ = _q_kv_formats(qkv_layout) + batch = cu_seqlens_q.numel() - 1 + heads = q.shape[-2] if q_format != "bhsd" else q.shape[1] + d_v = v.shape[-1] + output_shape = _fp8_output_shape( + batch, + max_seqlen_q, + heads, + d_v, + o_format, + total_tokens=q.shape[0] if o_format == "thd" else None, + ) + output, amax_o = _allocate_fp8_kernel_output(o_quantizer, output_shape, fake_dtype, q.device) + if fast_zero_fill and o_format == "thd": + output_data = _quantized_data(output) if _is_float8_tensor(output) else output + output_data.zero_() + if amax_o is not None: + amax_o.zero_() + cudnn = import_cudnn_frontend() + use_ragged_stats = ( + q_format == "thd" + and cudnn.backend_version() >= 90600 + and torch.cuda.get_device_capability(q.device) != (12, 0) + ) + stats_shape, _, _ = _stats_layout( + batch=batch, + heads=heads, + max_seqlen_q=max_seqlen_q, + total_tokens_q=q.shape[0] if q_format == "thd" else batch * max_seqlen_q, + ragged=use_ragged_stats, + ) + stats = torch.empty(stats_shape, dtype=torch.float32, device=q.device) + amax_s = ( + s_quantizer.amax + if isinstance(s_quantizer, Float8Quantizer) + else ( + torch.zeros(1, dtype=torch.float32, device=q.device) + if not _is_mxfp8_tensor(q) + else None + ) + ) + rng_elts = (max_seqlen_q * max_seqlen_q + _FP8_THREADS_PER_CTA - 1) // _FP8_THREADS_PER_CTA + rng_state = _reserve_philox_state(q.device, rng_gen, rng_elts) + cu_seqlens_q_padded = cu_seqlens_q if cu_seqlens_q_padded is None else cu_seqlens_q_padded + cu_seqlens_kv_padded = cu_seqlens_kv if cu_seqlens_kv_padded is None else cu_seqlens_kv_padded + kv_offsets_source = _kv_ragged_offsets_source( + qkv_layout, + cu_seqlens_q_padded, + cu_seqlens_kv_padded, + ) + + q_data = _quantized_data(q) + k_data = _quantized_data(k) + v_data = _quantized_data(v, columnwise=_is_mxfp8_tensor(v)) + output_data = _quantized_data(output) if _is_float8_tensor(output) else output + key = ( + "fp8_fwd", + max_seqlen_q, + max_seqlen_kv, + _tensor_metadata(q_data), + _tensor_metadata(k_data), + _tensor_metadata(v_data), + _tensor_metadata(output_data), + type(q).__name__, + type(o_quantizer).__name__, + qkv_layout, + o_format, + qkv_scale_inv_format, + float(attn_scale), + float(dropout), + attn_mask_type, + softmax_type, + tuple(window_size), + bottom_right_diagonal, + ) + entry = get_graph_entry(key) + if entry is None: + if torch.cuda.is_current_stream_capturing(): + raise RuntimeError("cuDNN FP8 attention graph must be built before CUDA graph capture.") + entry = _build_fp8_fwd_graph( + max_seqlen_q=max_seqlen_q, + max_seqlen_kv=max_seqlen_kv, + q=q, + k=k, + v=v, + output=output, + s_quantizer=s_quantizer, + o_quantizer=o_quantizer, + qkv_layout=qkv_layout, + o_format=o_format, + qkv_scale_inv_format=qkv_scale_inv_format, + attn_scale=attn_scale, + dropout=dropout, + attn_mask_type=attn_mask_type, + softmax_type=softmax_type, + window_size=window_size, + bottom_right_diagonal=bottom_right_diagonal, + softmax_offset=softmax_offset, + cu_seqlens_q=cu_seqlens_q, + cu_seqlens_kv=cu_seqlens_kv, + cu_seqlens_q_padded=cu_seqlens_q_padded, + cu_seqlens_kv_padded=cu_seqlens_kv_padded, + ) + put_graph_entry(key, entry) + + t = entry.tensors + variant_pack = { + t["Q"]: q_data, + t["K"]: k_data, + t["V"]: v_data, + t["O"]: output_data, + t["Stats"]: stats, + t["descale_q"]: _quantized_scale_inv(q), + t["descale_k"]: _quantized_scale_inv(k), + t["descale_v"]: _quantized_scale_inv(v, columnwise=_is_mxfp8_tensor(v)), + } + if "descale_s" in t: + variant_pack[t["descale_s"]] = torch.reciprocal(s_quantizer.scale) + variant_pack[t["scale_s"]] = s_quantizer.scale + if "scale_o" in t: + variant_pack[t["scale_o"]] = o_quantizer.scale + one = None + for name in ("constant_descale_s", "constant_scale_s", "constant_scale_o"): + if name in t: + if one is None: + one = torch.ones(1, dtype=torch.float32, device=q.device) + variant_pack[t[name]] = one + if "amax_s" in t: + variant_pack[t["amax_s"]] = amax_s + variant_pack[t["amax_o"]] = amax_o + if "seq_len_q" in t: + graph_batch = t["_graph_batch"] + variant_pack[t["seq_len_q"]] = _padded_sequence_lengths(cu_seqlens_q, graph_batch) + variant_pack[t["seq_len_kv"]] = _padded_sequence_lengths(cu_seqlens_kv, graph_batch) + graph_batch = t["_graph_batch"] + if "offset_q" in t: + variant_pack[t["offset_q"]] = _element_ragged_offsets( + cu_seqlens_q_padded, graph_batch, q_data.stride(0) + ) + variant_pack[t["offset_o"]] = _element_ragged_offsets( + cu_seqlens_q_padded, graph_batch, output_data.stride(0) + ) + if "offset_k" in t: + variant_pack[t["offset_k"]] = _element_ragged_offsets( + kv_offsets_source, graph_batch, k_data.stride(0) + ) + variant_pack[t["offset_v"]] = _element_ragged_offsets( + kv_offsets_source, graph_batch, v_data.stride(0) + ) + if "offset_stats" in t: + variant_pack[t["offset_stats"]] = _element_ragged_offsets( + cu_seqlens_q_padded, graph_batch, stats.stride(0) + ) + if "dropout_seed" in t: + variant_pack[t["dropout_seed"]] = rng_state[:1] + variant_pack[t["dropout_offset"]] = rng_state[1:] + if "softmax_offset" in t: + variant_pack[t["softmax_offset"]] = softmax_offset + entry.execute(variant_pack, q.device) + aux = [stats, rng_state] + if softmax_type != "vanilla": + aux.append(softmax_offset) + return output, aux + + +def fused_attn_fwd( + is_training: bool, + max_seqlen_q: int, + max_seqlen_kv: int, + cu_seqlens_q: torch.Tensor, + cu_seqlens_kv: torch.Tensor, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + fake_dtype: torch.dtype, + fused_attention_backend: FusedAttnBackend, + attn_bias: torch.Tensor = None, + cu_seqlens_q_padded: torch.Tensor = None, + cu_seqlens_kv_padded: torch.Tensor = None, + page_table_k: torch.Tensor = None, + page_table_v: torch.Tensor = None, + s_quantizer=None, + o_quantizer=None, + attn_scale: float = None, + dropout: float = 0.0, + fast_zero_fill: bool = True, + qkv_layout: str = "sbh3d", + o_format: str = "sbhd", + qkv_scale_inv_format: str = None, + attn_bias_type: str = "no_bias", + attn_mask_type: str = "padding", + softmax_type: str = "vanilla", + window_size: Tuple[int, int] = (-1, -1), + bottom_right_diagonal: bool = None, + rng_gen: torch.Generator = None, + softmax_offset: torch.Tensor = None, + return_max_logit: bool = False, + cuda_graph: bool = False, +) -> Tuple[Union[torch.Tensor, None], ...]: + """Execute fused attention through cuDNN Frontend's Python graph API.""" + + del cuda_graph + backend = FusedAttnBackend.cast(fused_attention_backend) + if backend == FusedAttnBackend.No_Backend: + raise ValueError("No cuDNN fused-attention backend supports this configuration.") + if attn_scale is None: + attn_scale = 1.0 / math.sqrt(q.size(-1)) + if bottom_right_diagonal is None: + bottom_right_diagonal = attn_mask_type in ( + "causal_bottom_right", + "padding_causal_bottom_right", + ) + if backend == FusedAttnBackend.FP8: + if page_table_k is not None or page_table_v is not None: + raise ValueError("FP8 fused attention does not support paged KV cache.") + if attn_bias_type != "no_bias" or attn_bias is not None: + raise ValueError("FP8 fused attention does not support attention bias.") + if return_max_logit: + raise ValueError("FP8 fused attention does not support returning maximum logits.") + return _fp8_forward( + is_training, + max_seqlen_q, + max_seqlen_kv, + cu_seqlens_q, + cu_seqlens_kv, + cu_seqlens_q_padded, + cu_seqlens_kv_padded, + q, + k, + v, + fake_dtype, + s_quantizer, + o_quantizer, + attn_scale, + dropout, + fast_zero_fill, + qkv_layout, + o_format, + qkv_scale_inv_format, + attn_mask_type, + softmax_type, + window_size, + bottom_right_diagonal, + rng_gen, + softmax_offset, + ) + + output, aux, max_logit = _f16_forward( + is_training, + max_seqlen_q, + max_seqlen_kv, + cu_seqlens_q, + cu_seqlens_kv, + q, + k, + v, + fake_dtype, + attn_bias, + cu_seqlens_q_padded, + cu_seqlens_kv_padded, + page_table_k, + page_table_v, + attn_scale, + dropout, + fast_zero_fill, + qkv_layout, + o_format, + attn_bias_type, + attn_mask_type, + softmax_type, + window_size, + bottom_right_diagonal, + rng_gen, + softmax_offset, + return_max_logit, + ) + if return_max_logit: + return output, aux, max_logit + return output, aux + + +def _build_f16_bwd_graph( + *, + max_seqlen_q: int, + max_seqlen_kv: int, + cu_seqlens_q: torch.Tensor, + cu_seqlens_kv: torch.Tensor, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + o: torch.Tensor, + d_o: torch.Tensor, + d_q: torch.Tensor, + d_k: torch.Tensor, + d_v: torch.Tensor, + attn_bias: Optional[torch.Tensor], + d_bias: Optional[torch.Tensor], + softmax_offset: Optional[torch.Tensor], + d_softmax_offset: Optional[torch.Tensor], + cu_seqlens_q_padded: torch.Tensor, + cu_seqlens_kv_padded: torch.Tensor, + attn_scale: float, + dropout: float, + qkv_layout: str, + o_format: str, + do_format: str, + dqkv_layout: str, + attn_bias_type: str, + attn_mask_type: str, + softmax_type: str, + window_size: Tuple[int, int], + bottom_right_diagonal: bool, + deterministic: bool, +) -> GraphEntry: + cudnn = import_cudnn_frontend() + graph = make_graph(torch_to_cudnn_dtype(q.dtype), q.device, name="te_fused_attention_bwd") + q_format, kv_format = _q_kv_formats(qkv_layout) + dq_format, dkv_format = _q_kv_formats(dqkv_layout) + batch = cu_seqlens_q.numel() - 1 + is_ragged_q = q_format == "thd" + is_ragged_kv = kv_format == "thd" + use_ragged_stats = is_ragged_q and cudnn.backend_version() >= 90600 + use_token_buckets = cudnn.backend_version() >= 90600 and torch.cuda.get_device_capability( + q.device + ) != (12, 0) + use_legacy_offsets = is_ragged_q or is_ragged_kv + graph_batch = _max_ragged_batch(batch) if use_legacy_offsets and use_token_buckets else batch + if not use_token_buckets: + use_ragged_stats = False + graph_seqlen_q = ( + _max_ragged_tokens(q.shape[0]) if is_ragged_q and use_token_buckets else max_seqlen_q + ) + graph_seqlen_kv = ( + _max_ragged_tokens(k.shape[0]) if is_ragged_kv and use_token_buckets else max_seqlen_kv + ) + + tensors: Dict[str, Any] = { + "_legacy_offsets": use_legacy_offsets, + "_graph_batch": graph_batch, + } + offset_q = offset_o = offset_k = offset_v = offset_stats = None + if is_ragged_q: + offset_q, q_mult = _ragged_offset_tensor( + graph, + cu_seqlens_q_padded, + multiplier=1, + name="offset_q", + length=graph_batch + 1, + data_type=torch.int64, + ) + offset_o, o_mult = _ragged_offset_tensor( + graph, + cu_seqlens_q_padded, + multiplier=1, + name="offset_o", + length=graph_batch + 1, + data_type=torch.int64, + ) + tensors.update(offset_q=offset_q, offset_o=offset_o) + else: + q_mult = o_mult = 1 + if is_ragged_kv: + offset_k, k_mult = _ragged_offset_tensor( + graph, + cu_seqlens_kv_padded, + multiplier=1, + name="offset_k", + length=graph_batch + 1, + data_type=torch.int64, + ) + offset_v, v_mult = _ragged_offset_tensor( + graph, + cu_seqlens_kv_padded, + multiplier=1, + name="offset_v", + length=graph_batch + 1, + data_type=torch.int64, + ) + tensors.update(offset_k=offset_k, offset_v=offset_v) + else: + k_mult = v_mult = 1 + + q_t = _make_bhsd_graph_tensor( + graph, + q, + q_format, + batch=graph_batch, + max_seqlen=graph_seqlen_q, + ragged_offset=offset_q, + ragged_offset_multiplier=q_mult, + name="Q", + ) + k_t = _make_bhsd_graph_tensor( + graph, + k, + kv_format, + batch=graph_batch, + max_seqlen=graph_seqlen_kv, + ragged_offset=offset_k, + ragged_offset_multiplier=k_mult, + name="K", + ) + v_t = _make_bhsd_graph_tensor( + graph, + v, + kv_format, + batch=graph_batch, + max_seqlen=graph_seqlen_kv, + ragged_offset=offset_v, + ragged_offset_multiplier=v_mult, + name="V", + ) + o_t = _make_bhsd_graph_tensor( + graph, + o, + o_format, + batch=graph_batch, + max_seqlen=graph_seqlen_q, + ragged_offset=offset_o, + ragged_offset_multiplier=o_mult, + name="O", + ) + do_t = _make_bhsd_graph_tensor( + graph, + d_o, + do_format, + batch=graph_batch, + max_seqlen=graph_seqlen_q, + ragged_offset=offset_o, + ragged_offset_multiplier=1, + name="dO", + ) + tensors.update(Q=q_t, K=k_t, V=v_t, O=o_t, dO=do_t) + + stats_physical, stats_dim, stats_stride = _stats_layout( + batch=graph_batch, + heads=q_t.get_dim()[1], + max_seqlen_q=graph_seqlen_q, + total_tokens_q=q.shape[0] if is_ragged_q else batch * max_seqlen_q, + ragged=use_ragged_stats, + ) + del stats_physical + if use_ragged_stats: + offset_stats, stats_mult = _ragged_offset_tensor( + graph, + cu_seqlens_q_padded, + multiplier=1, + name="offset_stats", + length=graph_batch + 1, + data_type=torch.int64, + ) + tensors["offset_stats"] = offset_stats + else: + stats_mult = 1 + stats_t = graph.tensor( + name="Stats", + dim=stats_dim, + stride=stats_stride, + data_type=cudnn.data_type.FLOAT, + ragged_offset=offset_stats, + ragged_offset_multiplier=stats_mult, + ) + tensors["Stats"] = stats_t + + options = _mask_options( + cudnn, + attn_mask_type, + window_size, + bottom_right_diagonal, + max_seqlen_q, + max_seqlen_kv, + ) + is_padding = options.pop("is_padding") + options.update( + attn_scale=float(attn_scale), + use_padding_mask=is_padding, + use_alibi_mask=attn_bias_type == "alibi", + use_deterministic_algorithm=deterministic, + ) + if use_ragged_stats: + options["max_total_seq_len_q"] = graph_seqlen_q + if ( + is_ragged_kv + and cudnn.backend_version() >= 90600 + and torch.cuda.get_device_capability(q.device) != (12, 0) + ): + options["max_total_seq_len_kv"] = graph_seqlen_kv + + if attn_bias_type == "post_scale_bias": + bias_t = graph.tensor_like(attn_bias, name="Bias") + tensors["Bias"] = bias_t + options["bias"] = bias_t + if d_bias is not None: + d_bias_t = graph.tensor_like(d_bias, name="dBias").set_output(True) + tensors["dBias"] = d_bias_t + options["dBias"] = d_bias_t + + if is_padding: + seq_q = _padded_sequence_lengths(cu_seqlens_q, graph_batch) + seq_kv = _padded_sequence_lengths(cu_seqlens_kv, graph_batch) + seq_q_t = graph.tensor_like(seq_q, name="seq_len_q") + seq_kv_t = graph.tensor_like(seq_kv, name="seq_len_kv") + tensors.update(seq_len_q=seq_q_t, seq_len_kv=seq_kv_t) + options.update(seq_len_q=seq_q_t, seq_len_kv=seq_kv_t) + + if dropout != 0.0: + seed_t = graph.tensor( + name="dropout_seed", + dim=(1, 1, 1, 1), + stride=(1, 1, 1, 1), + data_type=cudnn.data_type.INT64, + ) + offset_t = graph.tensor( + name="dropout_offset", + dim=(1, 1, 1, 1), + stride=(1, 1, 1, 1), + data_type=cudnn.data_type.INT64, + ) + tensors.update(dropout_seed=seed_t, dropout_offset=offset_t) + options["dropout"] = (float(dropout), seed_t, offset_t) + + if softmax_type != "vanilla": + if softmax_offset is None or d_softmax_offset is None: + raise ValueError(f"softmax_type={softmax_type!r} requires sink tensors.") + sink_t = graph.tensor_like(softmax_offset, name="softmax_offset") + dsink_t = graph.tensor_like(d_softmax_offset, name="d_softmax_offset").set_output(True) + tensors.update(softmax_offset=sink_t, d_softmax_offset=dsink_t) + options.update(sink_token=sink_t, dSink_token=dsink_t) + + dq_t, dk_t, dv_t = graph.sdpa_backward( + name="te_sdpa_backward", + q=q_t, + k=k_t, + v=v_t, + o=o_t, + dO=do_t, + stats=stats_t, + **options, + ) + dq_dim, dq_stride = _logical_bhsd_desc( + d_q, dq_format, batch=graph_batch, max_seqlen=graph_seqlen_q + ) + dk_dim, dk_stride = _logical_bhsd_desc( + d_k, dkv_format, batch=graph_batch, max_seqlen=graph_seqlen_kv + ) + dv_dim, dv_stride = _logical_bhsd_desc( + d_v, dkv_format, batch=graph_batch, max_seqlen=graph_seqlen_kv + ) + dq_t.set_output(True).set_dim(dq_dim).set_stride(dq_stride) + dk_t.set_output(True).set_dim(dk_dim).set_stride(dk_stride) + dv_t.set_output(True).set_dim(dv_dim).set_stride(dv_stride) + if is_ragged_q: + dq_t.set_ragged_offset(offset_q).set_ragged_offset_multiplier(1) + if is_ragged_kv: + dk_t.set_ragged_offset(offset_k).set_ragged_offset_multiplier(1) + dv_t.set_ragged_offset(offset_v).set_ragged_offset_multiplier(1) + tensors.update(dQ=dq_t, dK=dk_t, dV=dv_t) + + return GraphEntry( + graph=graph, + tensors=tensors, + workspace_size=finalize_graph(graph, cache_site=("f16", "bwd")), + ) + + +def _build_fp8_bwd_graph( + *, + max_seqlen_q, + max_seqlen_kv, + q, + k, + v, + o, + d_o, + d_o_f16, + d_q, + d_k, + d_v, + s_quantizer, + dp_quantizer, + dqkv_quantizer, + qkv_layout, + o_format, + do_format, + dqkv_layout, + qkv_scale_inv_format, + do_scale_inv_format, + attn_scale, + dropout, + attn_mask_type, + softmax_type, + window_size, + bottom_right_diagonal, + deterministic, + softmax_offset, + d_softmax_offset, + cu_seqlens_q, + cu_seqlens_kv, + cu_seqlens_q_padded, + cu_seqlens_kv_padded, +): + cudnn = import_cudnn_frontend() + q_format, kv_format = _q_kv_formats(qkv_layout) + dq_format, dkv_format = _q_kv_formats(dqkv_layout) + batch = cu_seqlens_q.numel() - 1 + q_data = _quantized_data(q) + k_data = _quantized_data(k) + v_data = _quantized_data(v) + o_data = _quantized_data(o) if _is_float8_tensor(o) else o + do_data = _quantized_data(d_o) + dq_data = _quantized_data(d_q) if _is_float8_tensor(d_q) else d_q + dk_data = _quantized_data(d_k) if _is_float8_tensor(d_k) else d_k + dv_data = _quantized_data(d_v) if _is_float8_tensor(d_v) else d_v + heads = q_data.shape[-2] if q_format != "bhsd" else q_data.shape[1] + kv_heads = k_data.shape[-2] if kv_format != "bhsd" else k_data.shape[1] + d_qk, d_value = q_data.shape[-1], v_data.shape[-1] + is_ragged_q = q_format == "thd" + is_ragged_kv = kv_format == "thd" + use_token_buckets = cudnn.backend_version() >= 90600 and torch.cuda.get_device_capability( + q.device + ) != (12, 0) + use_ragged_stats = is_ragged_q and use_token_buckets + use_legacy_offsets = is_ragged_q or is_ragged_kv + graph_batch = _max_ragged_batch(batch) if use_legacy_offsets and use_token_buckets else batch + graph_seqlen_q = ( + _max_ragged_tokens(q_data.shape[0]) if is_ragged_q and use_token_buckets else max_seqlen_q + ) + graph_seqlen_kv = ( + _max_ragged_tokens(k_data.shape[0]) if is_ragged_kv and use_token_buckets else max_seqlen_kv + ) + graph = make_graph(_fp8_cudnn_dtype(q), q.device, name="te_fp8_sdpa_bwd") + tensors: Dict[str, Any] = { + "_legacy_offsets": use_legacy_offsets, + "_graph_batch": graph_batch, + } + + offset_q = offset_o = offset_k = offset_v = offset_stats = None + if is_ragged_q: + offset_q, q_mult = _ragged_offset_tensor( + graph, + cu_seqlens_q_padded, + multiplier=1, + name="offset_q", + length=graph_batch + 1, + data_type=torch.int64, + ) + offset_o, o_mult = _ragged_offset_tensor( + graph, + cu_seqlens_q_padded, + multiplier=1, + name="offset_o", + length=graph_batch + 1, + data_type=torch.int64, + ) + tensors.update(offset_q=offset_q, offset_o=offset_o) + else: + q_mult = o_mult = 1 + if is_ragged_kv: + offset_k, k_mult = _ragged_offset_tensor( + graph, + cu_seqlens_kv_padded, + multiplier=1, + name="offset_k", + length=graph_batch + 1, + data_type=torch.int64, + ) + offset_v, v_mult = _ragged_offset_tensor( + graph, + cu_seqlens_kv_padded, + multiplier=1, + name="offset_v", + length=graph_batch + 1, + data_type=torch.int64, + ) + tensors.update(offset_k=offset_k, offset_v=offset_v) + else: + k_mult = v_mult = 1 + + q_t = _make_bhsd_graph_tensor( + graph, + q_data, + q_format, + batch=graph_batch, + max_seqlen=graph_seqlen_q, + data_type=_fp8_cudnn_dtype(q), + ragged_offset=offset_q, + ragged_offset_multiplier=q_mult, + name="Q", + ) + k_t = _make_bhsd_graph_tensor( + graph, + k_data, + kv_format, + batch=graph_batch, + max_seqlen=graph_seqlen_kv, + data_type=_fp8_cudnn_dtype(k), + ragged_offset=offset_k, + ragged_offset_multiplier=k_mult, + name="K", + ) + v_t = _make_bhsd_graph_tensor( + graph, + v_data, + kv_format, + batch=graph_batch, + max_seqlen=graph_seqlen_kv, + data_type=_fp8_cudnn_dtype(v), + ragged_offset=offset_v, + ragged_offset_multiplier=v_mult, + name="V", + ) + o_dtype = _fp8_cudnn_dtype(o) if _is_float8_tensor(o) else o.dtype + o_t = _make_bhsd_graph_tensor( + graph, + o_data, + o_format, + batch=graph_batch, + max_seqlen=graph_seqlen_q, + data_type=o_dtype, + ragged_offset=offset_o, + ragged_offset_multiplier=o_mult, + name="O", + ) + do_t = _make_bhsd_graph_tensor( + graph, + do_data, + do_format, + batch=graph_batch, + max_seqlen=graph_seqlen_q, + data_type=_fp8_cudnn_dtype(d_o), + ragged_offset=offset_o, + ragged_offset_multiplier=o_mult, + name="dO", + ) + _, stats_dim, stats_stride = _stats_layout( + batch=graph_batch, + heads=heads, + max_seqlen_q=graph_seqlen_q, + total_tokens_q=q_data.shape[0] if is_ragged_q else batch * max_seqlen_q, + ragged=use_ragged_stats, + ) + if use_ragged_stats: + offset_stats, stats_mult = _ragged_offset_tensor( + graph, + cu_seqlens_q_padded, + multiplier=1, + name="offset_stats", + length=graph_batch + 1, + data_type=torch.int64, + ) + tensors["offset_stats"] = offset_stats + else: + stats_mult = 1 + stats_t = graph.tensor( + name="Stats", + dim=stats_dim, + stride=stats_stride, + data_type=cudnn.data_type.FLOAT, + ragged_offset=offset_stats, + ragged_offset_multiplier=stats_mult, + ) + tensors.update(Q=q_t, K=k_t, V=v_t, O=o_t, dO=do_t, Stats=stats_t) + + options = _mask_options( + cudnn, + attn_mask_type, + window_size, + bottom_right_diagonal, + max_seqlen_q, + max_seqlen_kv, + ) + is_padding = options.pop("is_padding") + if "diagonal_band_left_bound" in options: + options["left_bound"] = options.pop("diagonal_band_left_bound") + if "diagonal_band_right_bound" in options: + options["right_bound"] = options.pop("diagonal_band_right_bound") + options.update( + attn_scale=float(attn_scale), + use_padding_mask=is_padding, + use_deterministic_algorithm=deterministic, + ) + if is_padding: + seq_q = _padded_sequence_lengths(cu_seqlens_q, graph_batch) + seq_kv = _padded_sequence_lengths(cu_seqlens_kv, graph_batch) + seq_q_t = graph.tensor_like(seq_q, name="seq_len_q") + seq_kv_t = graph.tensor_like(seq_kv, name="seq_len_kv") + tensors.update(seq_len_q=seq_q_t, seq_len_kv=seq_kv_t) + options.update(seq_len_q=seq_q_t, seq_len_kv=seq_kv_t) + if dropout != 0.0: + seed_t = graph.tensor( + name="dropout_seed", + dim=(1, 1, 1, 1), + stride=(1, 1, 1, 1), + data_type=cudnn.data_type.INT64, + ) + offset_t = graph.tensor( + name="dropout_offset", + dim=(1, 1, 1, 1), + stride=(1, 1, 1, 1), + data_type=cudnn.data_type.INT64, + ) + tensors.update(dropout_seed=seed_t, dropout_offset=offset_t) + options["dropout"] = (float(dropout), seed_t, offset_t) + if softmax_type != "vanilla": + sink_t = graph.tensor_like(softmax_offset, name="softmax_offset") + dsink_t = graph.tensor_like(d_softmax_offset, name="d_softmax_offset").set_output(True) + tensors.update(softmax_offset=sink_t, d_softmax_offset=dsink_t) + options.update(sink_token=sink_t, dSink_token=dsink_t) + + if _is_mxfp8_tensor(q): + if d_o_f16 is None: + raise ValueError("MXFP8 attention backward requires the high-precision dO tensor.") + q_col = _quantized_data(q, columnwise=True) + k_col = _quantized_data(k, columnwise=True) + do_col = _quantized_data(d_o, columnwise=True) + q_col_t = _make_bhsd_graph_tensor( + graph, + q_col, + q_format, + batch=batch, + max_seqlen=max_seqlen_q, + data_type=_fp8_cudnn_dtype(q), + name="Q_T", + ) + k_col_t = _make_bhsd_graph_tensor( + graph, + k_col, + kv_format, + batch=batch, + max_seqlen=max_seqlen_kv, + data_type=_fp8_cudnn_dtype(k), + name="K_T", + ) + do_col_t = _make_bhsd_graph_tensor( + graph, + do_col, + do_format, + batch=batch, + max_seqlen=max_seqlen_q, + data_type=_fp8_cudnn_dtype(d_o), + name="dO_T", + ) + do_f16_t = _make_bhsd_graph_tensor( + graph, + d_o_f16, + do_format, + batch=batch, + max_seqlen=max_seqlen_q, + data_type=d_o_f16.dtype, + name="dO_f16", + ) + scale_format_q = qkv_scale_inv_format or q_format + scale_format_kv = qkv_scale_inv_format or kv_format + scale_format_do = do_scale_inv_format or do_format + padded = _mxfp8_padded_sizes(max_seqlen_q, max_seqlen_kv, d_qk, d_value) + + def mx_scale(name, h, s, d, fmt): + tensor = _make_mxfp8_scale_tensor( + graph, + cudnn, + name=name, + batch=batch, + heads=h, + seqlen=padded[s], + dim=padded[d], + tensor_format=fmt, + ) + tensors[name] = tensor + return tensor + + descale_q = mx_scale("descale_q", heads, "s_q_padded", "d_qk_scale_padded", scale_format_q) + descale_q_t = mx_scale( + "descale_q_t", heads, "s_q_scale_padded", "d_qk_padded", scale_format_q + ) + descale_k = mx_scale( + "descale_k", kv_heads, "s_kv_padded", "d_qk_scale_padded", scale_format_kv + ) + descale_k_t = mx_scale( + "descale_k_t", kv_heads, "s_kv_scale_padded", "d_qk_padded", scale_format_kv + ) + descale_v = mx_scale( + "descale_v", kv_heads, "s_kv_padded", "d_v_scale_padded", scale_format_kv + ) + descale_do = mx_scale( + "descale_do", heads, "s_q_padded", "d_v_scale_padded", scale_format_do + ) + descale_do_t = mx_scale( + "descale_do_t", heads, "s_q_scale_padded", "d_v_padded", scale_format_do + ) + op = build_fp8_backward_operation( + graph, + { + "q": q_t, + "q_t": q_col_t, + "k": k_t, + "k_t": k_col_t, + "v": v_t, + "o": o_t, + "do_f16": do_f16_t, + "do": do_t, + "do_t": do_col_t, + "stats": stats_t, + "descale_q": descale_q, + "descale_q_t": descale_q_t, + "descale_k": descale_k, + "descale_k_t": descale_k_t, + "descale_v": descale_v, + "descale_do": descale_do, + "descale_do_t": descale_do_t, + }, + options, + FP8AttentionGraphConfig("mxfp8", "te_sdpa_mxfp8_backward"), + ) + dq_t, dk_t, dv_t, amax_outputs = op["dq"], op["dk"], op["dv"], op["amax"] + for amax_t in amax_outputs: + amax_t.set_output(False).set_data_type(cudnn.data_type.FLOAT).set_dim( + (1, 1, 1, 1) + ).set_stride((1, 1, 1, 1)) + tensors.update(Q_T=q_col_t, K_T=k_col_t, dO_T=do_col_t, dO_f16=do_f16_t) + else: + scalar_names = ( + "descale_q", + "descale_k", + "descale_v", + "descale_o", + "descale_do", + ) + scalars = {name: _scalar_graph_tensor(graph, cudnn, name) for name in scalar_names} + tensors.update(scalars) + delayed = isinstance(dqkv_quantizer, Float8Quantizer) + if isinstance(s_quantizer, Float8Quantizer): + for name in ("descale_s", "scale_s"): + tensors[name] = _scalar_graph_tensor(graph, cudnn, name) + else: + for name in ("descale_s", "scale_s"): + tensors[f"constant_{name}"] = _constant_graph_tensor(graph, cudnn, name) + if isinstance(dp_quantizer, Float8Quantizer): + for name in ("descale_dp", "scale_dp"): + tensors[name] = _scalar_graph_tensor(graph, cudnn, name) + else: + for name in ("descale_dp", "scale_dp"): + tensors[f"constant_{name}"] = _constant_graph_tensor(graph, cudnn, name) + for name in ("scale_dq", "scale_dk", "scale_dv"): + key = name if delayed else f"constant_{name}" + tensors[key] = _scalar_graph_tensor(graph, cudnn, name) + descale_o_arg = tensors["descale_o"] + descale_s_arg = ( + tensors["descale_s"] if "descale_s" in tensors else tensors["constant_descale_s"] + ) + descale_dp_arg = ( + tensors["descale_dp"] if "descale_dp" in tensors else tensors["constant_descale_dp"] + ) + scale_s_arg = tensors["scale_s"] if "scale_s" in tensors else tensors["constant_scale_s"] + scale_dp_arg = ( + tensors["scale_dp"] if "scale_dp" in tensors else tensors["constant_scale_dp"] + ) + scale_dq_arg = ( + tensors["scale_dq"] if "scale_dq" in tensors else tensors["constant_scale_dq"] + ) + scale_dk_arg = ( + tensors["scale_dk"] if "scale_dk" in tensors else tensors["constant_scale_dk"] + ) + scale_dv_arg = ( + tensors["scale_dv"] if "scale_dv" in tensors else tensors["constant_scale_dv"] + ) + op = build_fp8_backward_operation( + graph, + { + "q": q_t, + "k": k_t, + "v": v_t, + "o": o_t, + "do": do_t, + "stats": stats_t, + "descale_q": tensors["descale_q"], + "descale_k": tensors["descale_k"], + "descale_v": tensors["descale_v"], + "descale_o": descale_o_arg, + "descale_do": tensors["descale_do"], + "descale_s": descale_s_arg, + "descale_dp": descale_dp_arg, + "scale_s": scale_s_arg, + "scale_dq": scale_dq_arg, + "scale_dk": scale_dk_arg, + "scale_dv": scale_dv_arg, + "scale_dp": scale_dp_arg, + }, + options, + FP8AttentionGraphConfig("delayed" if delayed else "current", "te_sdpa_fp8_backward"), + ) + dq_t, dk_t, dv_t = op["dq"], op["dk"], op["dv"] + amax_dq_t, amax_dk_t = op["amax_dq"], op["amax_dk"] + amax_dv_t, amax_dp_t = op["amax_dv"], op["amax_dp"] + for name, amax_t in ( + ("amax_dq", amax_dq_t), + ("amax_dk", amax_dk_t), + ("amax_dv", amax_dv_t), + ("amax_dp", amax_dp_t), + ): + amax_t.set_output(True).set_data_type(cudnn.data_type.FLOAT).set_dim( + (1, 1, 1, 1) + ).set_stride((1, 1, 1, 1)) + tensors[name] = amax_t + + dq_dim, dq_stride = _logical_bhsd_desc( + dq_data, + dq_format, + batch=graph_batch, + max_seqlen=graph_seqlen_q, + ) + dk_dim, dk_stride = _logical_bhsd_desc( + dk_data, + dkv_format, + batch=graph_batch, + max_seqlen=graph_seqlen_kv, + ) + dv_dim, dv_stride = _logical_bhsd_desc( + dv_data, + dkv_format, + batch=graph_batch, + max_seqlen=graph_seqlen_kv, + ) + dq_t.set_output(True).set_data_type( + _fp8_cudnn_dtype(d_q) if _is_float8_tensor(d_q) else d_q.dtype + ).set_dim(dq_dim).set_stride(dq_stride) + dk_t.set_output(True).set_data_type( + _fp8_cudnn_dtype(d_k) if _is_float8_tensor(d_k) else d_k.dtype + ).set_dim(dk_dim).set_stride(dk_stride) + dv_t.set_output(True).set_data_type( + _fp8_cudnn_dtype(d_v) if _is_float8_tensor(d_v) else d_v.dtype + ).set_dim(dv_dim).set_stride(dv_stride) + if is_ragged_q: + dq_t.set_ragged_offset(offset_q).set_ragged_offset_multiplier(q_mult) + if is_ragged_kv: + dk_t.set_ragged_offset(offset_k).set_ragged_offset_multiplier(k_mult) + dv_t.set_ragged_offset(offset_v).set_ragged_offset_multiplier(v_mult) + tensors.update(dQ=dq_t, dK=dk_t, dV=dv_t) + return GraphEntry( + graph=graph, + tensors=tensors, + workspace_size=finalize_graph(graph, cache_site=("fp8", "bwd")), + ) + + +def _fp8_backward( + max_seqlen_q, + max_seqlen_kv, + cu_seqlens_q, + cu_seqlens_kv, + cu_seqlens_q_padded, + cu_seqlens_kv_padded, + q, + k, + v, + o, + d_o, + fake_dtype, + aux_ctx_tensors, + s_quantizer, + dp_quantizer, + dqkv_quantizer, + attn_scale, + dropout, + fast_zero_fill, + qkv_layout, + o_format, + do_format, + dqkv_layout, + qkv_scale_inv_format, + do_scale_inv_format, + attn_mask_type, + softmax_type, + window_size, + bottom_right_diagonal, + deterministic, +): + if _is_mxfp8_tensor(q) and "padding" in attn_mask_type: + q_hp, k_hp, v_hp = (tensor.dequantize(dtype=fake_dtype) for tensor in (q, k, v)) + d_o_f16 = aux_ctx_tensors[-1] + aux_count = 3 if softmax_type != "vanilla" else 2 + return fused_attn_bwd( + max_seqlen_q, + max_seqlen_kv, + cu_seqlens_q, + cu_seqlens_kv, + q_hp, + k_hp, + v_hp, + o, + d_o_f16, + fake_dtype, + aux_ctx_tensors[:aux_count], + FusedAttnBackend.F16_arbitrary_seqlen, + s_quantizer=None, + dp_quantizer=None, + dqkv_quantizer=None, + attn_scale=attn_scale, + dropout=dropout, + fast_zero_fill=fast_zero_fill, + qkv_layout=qkv_layout, + o_format=o_format, + do_format=do_format, + dqkv_layout=dqkv_layout, + attn_bias_type="no_bias", + attn_mask_type=attn_mask_type, + softmax_type=softmax_type, + window_size=window_size, + bottom_right_diagonal=bottom_right_diagonal, + deterministic=deterministic, + ) + q_format, kv_format = _q_kv_formats(qkv_layout) + dq_format, dkv_format = _q_kv_formats(dqkv_layout) + batch = cu_seqlens_q.numel() - 1 + heads = q.shape[-2] if q_format != "bhsd" else q.shape[1] + kv_heads = k.shape[-2] if kv_format != "bhsd" else k.shape[1] + output_dtype = torch.uint8 if isinstance(dqkv_quantizer, Float8Quantizer) else fake_dtype + grad_data = _allocate_attention_grad_data( + batch=batch, + heads=heads, + kv_heads=kv_heads, + max_seqlen_q=max_seqlen_q, + max_seqlen_kv=max_seqlen_kv, + head_dim_qk=q.shape[-1], + head_dim_v=v.shape[-1], + dqkv_layout=dqkv_layout, + dtype=output_dtype, + device=q.device, + zero=fast_zero_fill + or ( + not isinstance(dqkv_quantizer, Float8Quantizer) + and (dq_format == "thd" or dkv_format == "thd") + ), + total_tokens_q=q.shape[0] if dq_format == "thd" else None, + total_tokens_kv=k.shape[0] if dkv_format == "thd" else None, + ) + if isinstance(dqkv_quantizer, Float8Quantizer): + d_q, d_k, d_v = _wrap_float8_grad_outputs(dqkv_quantizer, grad_data, fake_dtype) + if fast_zero_fill and (dq_format == "thd" or dkv_format == "thd"): + dqkv_quantizer.amax.zero_() + else: + d_q, d_k, d_v = grad_data + stats, rng_state = aux_ctx_tensors[:2] + softmax_offset = aux_ctx_tensors[2] if softmax_type != "vanilla" else None + d_o_f16 = aux_ctx_tensors[-1] if _is_mxfp8_tensor(q) else None + d_softmax_offset = torch.empty_like(softmax_offset) if softmax_offset is not None else None + hidden_amax = [torch.zeros(1, dtype=torch.float32, device=q.device) for _ in range(4)] + cu_seqlens_q_padded = cu_seqlens_q if cu_seqlens_q_padded is None else cu_seqlens_q_padded + cu_seqlens_kv_padded = cu_seqlens_kv if cu_seqlens_kv_padded is None else cu_seqlens_kv_padded + kv_offsets_source = _kv_ragged_offsets_source( + qkv_layout, + cu_seqlens_q_padded, + cu_seqlens_kv_padded, + ) + + key = ( + "fp8_bwd", + max_seqlen_q, + max_seqlen_kv, + *( + _tensor_metadata(_quantized_data(x) if _is_float8_tensor(x) else x) + for x in (q, k, v, o, d_o) + ), + *(getattr(x, "_fp8_dtype", None) for x in (q, k, v, o, d_o, d_q, d_k, d_v)), + type(q).__name__, + type(dqkv_quantizer).__name__, + qkv_layout, + o_format, + do_format, + dqkv_layout, + qkv_scale_inv_format, + do_scale_inv_format, + float(attn_scale), + float(dropout), + attn_mask_type, + softmax_type, + tuple(window_size), + bottom_right_diagonal, + deterministic, + ) + entry = get_graph_entry(key) + if entry is None: + if torch.cuda.is_current_stream_capturing(): + raise RuntimeError("cuDNN FP8 attention graph must be built before CUDA graph capture.") + entry = _build_fp8_bwd_graph( + max_seqlen_q=max_seqlen_q, + max_seqlen_kv=max_seqlen_kv, + q=q, + k=k, + v=v, + o=o, + d_o=d_o, + d_o_f16=d_o_f16, + d_q=d_q, + d_k=d_k, + d_v=d_v, + s_quantizer=s_quantizer, + dp_quantizer=dp_quantizer, + dqkv_quantizer=dqkv_quantizer, + qkv_layout=qkv_layout, + o_format=o_format, + do_format=do_format, + dqkv_layout=dqkv_layout, + qkv_scale_inv_format=qkv_scale_inv_format, + do_scale_inv_format=do_scale_inv_format, + attn_scale=attn_scale, + dropout=dropout, + attn_mask_type=attn_mask_type, + softmax_type=softmax_type, + window_size=window_size, + bottom_right_diagonal=bottom_right_diagonal, + deterministic=deterministic, + softmax_offset=softmax_offset, + d_softmax_offset=d_softmax_offset, + cu_seqlens_q=cu_seqlens_q, + cu_seqlens_kv=cu_seqlens_kv, + cu_seqlens_q_padded=cu_seqlens_q_padded, + cu_seqlens_kv_padded=cu_seqlens_kv_padded, + ) + put_graph_entry(key, entry) + + t = entry.tensors + variant_pack = { + t["Q"]: _quantized_data(q), + t["K"]: _quantized_data(k), + t["V"]: _quantized_data(v), + t["O"]: _quantized_data(o) if _is_float8_tensor(o) else o, + t["dO"]: _quantized_data(d_o), + t["Stats"]: stats, + t["dQ"]: _quantized_data(d_q) if _is_float8_tensor(d_q) else d_q, + t["dK"]: _quantized_data(d_k) if _is_float8_tensor(d_k) else d_k, + t["dV"]: _quantized_data(d_v) if _is_float8_tensor(d_v) else d_v, + } + if _is_mxfp8_tensor(q): + variant_pack.update( + { + t["Q_T"]: _quantized_data(q, columnwise=True), + t["K_T"]: _quantized_data(k, columnwise=True), + t["dO_T"]: _quantized_data(d_o, columnwise=True), + t["dO_f16"]: d_o_f16, + t["descale_q"]: _quantized_scale_inv(q), + t["descale_q_t"]: _quantized_scale_inv(q, columnwise=True), + t["descale_k"]: _quantized_scale_inv(k), + t["descale_k_t"]: _quantized_scale_inv(k, columnwise=True), + t["descale_v"]: _quantized_scale_inv(v), + t["descale_do"]: _quantized_scale_inv(d_o), + t["descale_do_t"]: _quantized_scale_inv(d_o, columnwise=True), + } + ) + else: + variant_pack.update( + { + t["descale_q"]: _quantized_scale_inv(q), + t["descale_k"]: _quantized_scale_inv(k), + t["descale_v"]: _quantized_scale_inv(v), + t["descale_o"]: ( + _quantized_scale_inv(o) + if _is_float8_tensor(o) + else torch.ones(1, dtype=torch.float32, device=q.device) + ), + t["descale_do"]: _quantized_scale_inv(d_o), + } + ) + if "descale_s" in t: + variant_pack[t["descale_s"]] = torch.reciprocal(s_quantizer.scale) + variant_pack[t["scale_s"]] = s_quantizer.scale + if "descale_dp" in t: + variant_pack[t["descale_dp"]] = torch.reciprocal(dp_quantizer.scale) + variant_pack[t["scale_dp"]] = dp_quantizer.scale + for name in ("scale_dq", "scale_dk", "scale_dv"): + if name in t: + variant_pack[t[name]] = dqkv_quantizer.scale + one = torch.ones(1, dtype=torch.float32, device=q.device) + for name, graph_tensor in t.items(): + if name.startswith("constant_"): + variant_pack[graph_tensor] = one + amax_values = ( + ( + dqkv_quantizer.amax, + dqkv_quantizer.amax, + dqkv_quantizer.amax, + dp_quantizer.amax, + ) + if isinstance(dqkv_quantizer, Float8Quantizer) + else hidden_amax + ) + for name, value in zip(("amax_dq", "amax_dk", "amax_dv", "amax_dp"), amax_values): + variant_pack[t[name]] = value + if "seq_len_q" in t: + graph_batch = t["_graph_batch"] + variant_pack[t["seq_len_q"]] = _padded_sequence_lengths(cu_seqlens_q, graph_batch) + variant_pack[t["seq_len_kv"]] = _padded_sequence_lengths(cu_seqlens_kv, graph_batch) + graph_batch = t["_graph_batch"] + q_data = _quantized_data(q) + k_data = _quantized_data(k) + v_data = _quantized_data(v) + o_data = _quantized_data(o) if _is_float8_tensor(o) else o + if "offset_q" in t: + variant_pack[t["offset_q"]] = _element_ragged_offsets( + cu_seqlens_q_padded, graph_batch, q_data.stride(0) + ) + variant_pack[t["offset_o"]] = _element_ragged_offsets( + cu_seqlens_q_padded, graph_batch, o_data.stride(0) + ) + if "offset_k" in t: + variant_pack[t["offset_k"]] = _element_ragged_offsets( + kv_offsets_source, graph_batch, k_data.stride(0) + ) + variant_pack[t["offset_v"]] = _element_ragged_offsets( + kv_offsets_source, graph_batch, v_data.stride(0) + ) + if "offset_stats" in t: + variant_pack[t["offset_stats"]] = _element_ragged_offsets( + cu_seqlens_q_padded, graph_batch, stats.stride(0) + ) + if "dropout_seed" in t: + variant_pack[t["dropout_seed"]] = rng_state[:1] + variant_pack[t["dropout_offset"]] = rng_state[1:] + if "softmax_offset" in t: + variant_pack[t["softmax_offset"]] = softmax_offset + variant_pack[t["d_softmax_offset"]] = d_softmax_offset + entry.execute(variant_pack, q.device) + return d_q, d_k, d_v, None, d_softmax_offset + + +def fused_attn_bwd( + max_seqlen_q: int, + max_seqlen_kv: int, + cu_seqlens_q: torch.Tensor, + cu_seqlens_kv: torch.Tensor, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + o: torch.Tensor, + d_o: torch.Tensor, + fake_dtype: torch.dtype, + aux_ctx_tensors: List[torch.Tensor], + fused_attention_backend: FusedAttnBackend, + cu_seqlens_q_padded: torch.Tensor = None, + cu_seqlens_kv_padded: torch.Tensor = None, + s_quantizer=None, + dp_quantizer=None, + dqkv_quantizer=None, + attn_scale: Optional[float] = None, + dropout: float = 0.0, + fast_zero_fill: bool = True, + qkv_layout: str = "sbh3d", + o_format: str = "sbhd", + do_format: str = "sbhd", + dqkv_layout: str = "sbh3d", + qkv_scale_inv_format: str = None, + do_scale_inv_format: str = None, + attn_bias_type: str = "no_bias", + attn_mask_type: str = "padding", + softmax_type: str = "vanilla", + window_size: Tuple[int, int] = (-1, -1), + bottom_right_diagonal: bool = None, + deterministic: bool = False, + cuda_graph: bool = False, +) -> Tuple[Union[torch.Tensor, None], ...]: + """Execute fused-attention backward through the Python graph API.""" + + del cuda_graph + backend = FusedAttnBackend.cast(fused_attention_backend) + if not aux_ctx_tensors: + raise ValueError("Fused-attention backward requires forward auxiliary tensors.") + if attn_scale is None: + attn_scale = 1.0 / math.sqrt(q.size(-1)) + if bottom_right_diagonal is None: + bottom_right_diagonal = attn_mask_type in ( + "causal_bottom_right", + "padding_causal_bottom_right", + ) + if backend == FusedAttnBackend.FP8: + if attn_bias_type != "no_bias": + raise ValueError("FP8 fused attention backward does not support attention bias.") + return _fp8_backward( + max_seqlen_q, + max_seqlen_kv, + cu_seqlens_q, + cu_seqlens_kv, + cu_seqlens_q_padded, + cu_seqlens_kv_padded, + q, + k, + v, + o, + d_o, + fake_dtype, + aux_ctx_tensors, + s_quantizer, + dp_quantizer, + dqkv_quantizer, + attn_scale, + dropout, + fast_zero_fill, + qkv_layout, + o_format, + do_format, + dqkv_layout, + qkv_scale_inv_format, + do_scale_inv_format, + attn_mask_type, + softmax_type, + window_size, + bottom_right_diagonal, + deterministic, + ) + if backend != FusedAttnBackend.F16_arbitrary_seqlen: + raise ValueError("No cuDNN fused-attention backend supports this backward configuration.") + + stats = aux_ctx_tensors[0] + rng_state = aux_ctx_tensors[1] + aux_index = 2 + attn_bias = None + if attn_bias_type not in ("no_bias", "alibi"): + attn_bias = aux_ctx_tensors[aux_index] + aux_index += 1 + softmax_offset = None + if softmax_type != "vanilla": + softmax_offset = aux_ctx_tensors[aux_index] + + q_format, kv_format = _q_kv_formats(qkv_layout) + batch = cu_seqlens_q.numel() - 1 + if q_format == "thd" or kv_format == "thd": + d_q, d_k, d_v = _allocate_grad_views((q, k, v)) + else: + heads = q.shape[1] if q_format == "bhsd" else q.shape[-2] + kv_heads = k.shape[1] if kv_format == "bhsd" else k.shape[-2] + d_q, d_k, d_v = _allocate_attention_grad_data( + batch=batch, + heads=heads, + kv_heads=kv_heads, + max_seqlen_q=max_seqlen_q, + max_seqlen_kv=max_seqlen_kv, + head_dim_qk=q.shape[-1], + head_dim_v=v.shape[-1], + dqkv_layout=dqkv_layout, + dtype=q.dtype, + device=q.device, + zero=fast_zero_fill, + ) + d_bias = None + if attn_bias_type == "post_scale_bias": + # cuDNN does not support the [1,1,1,S] reduction form. + if tuple(attn_bias.shape[:3]) != (1, 1, 1): + d_bias = torch.empty_like(attn_bias) + d_softmax_offset = torch.empty_like(softmax_offset) if softmax_type != "vanilla" else None + cu_seqlens_q_padded = cu_seqlens_q if cu_seqlens_q_padded is None else cu_seqlens_q_padded + cu_seqlens_kv_padded = cu_seqlens_kv if cu_seqlens_kv_padded is None else cu_seqlens_kv_padded + cudnn = import_cudnn_frontend() + use_token_buckets = cudnn.backend_version() >= 90600 and torch.cuda.get_device_capability( + q.device + ) != (12, 0) + use_legacy_offsets = q_format == "thd" or kv_format == "thd" + graph_batch = _max_ragged_batch(batch) if use_legacy_offsets and use_token_buckets else batch + + key = ( + "f16_bwd", + max_seqlen_q, + max_seqlen_kv, + graph_batch, + _tensor_metadata(q), + _tensor_metadata(k), + _tensor_metadata(v), + _tensor_metadata(o), + _tensor_metadata(d_o), + _tensor_metadata(stats), + _tensor_metadata(d_q), + _tensor_metadata(d_k), + _tensor_metadata(d_v), + _tensor_metadata(attn_bias), + qkv_layout, + o_format, + do_format, + dqkv_layout, + float(attn_scale), + float(dropout), + attn_bias_type, + attn_mask_type, + softmax_type, + tuple(window_size), + bottom_right_diagonal, + deterministic, + ) + entry = get_graph_entry(key) + if entry is None: + if torch.cuda.is_current_stream_capturing(): + raise RuntimeError("cuDNN attention graph must be built before CUDA graph capture.") + entry = _build_f16_bwd_graph( + max_seqlen_q=max_seqlen_q, + max_seqlen_kv=max_seqlen_kv, + cu_seqlens_q=cu_seqlens_q, + cu_seqlens_kv=cu_seqlens_kv, + q=q, + k=k, + v=v, + o=o, + d_o=d_o, + d_q=d_q, + d_k=d_k, + d_v=d_v, + attn_bias=attn_bias, + d_bias=d_bias, + softmax_offset=softmax_offset, + d_softmax_offset=d_softmax_offset, + cu_seqlens_q_padded=cu_seqlens_q_padded, + cu_seqlens_kv_padded=cu_seqlens_kv_padded, + attn_scale=attn_scale, + dropout=dropout, + qkv_layout=qkv_layout, + o_format=o_format, + do_format=do_format, + dqkv_layout=dqkv_layout, + attn_bias_type=attn_bias_type, + attn_mask_type=attn_mask_type, + softmax_type=softmax_type, + window_size=window_size, + bottom_right_diagonal=bottom_right_diagonal, + deterministic=deterministic, + ) + put_graph_entry(key, entry) + + tensors = entry.tensors + variant_pack: Dict[Any, Any] = { + tensors["Q"]: q, + tensors["K"]: k, + tensors["V"]: v, + tensors["O"]: o, + tensors["dO"]: d_o, + tensors["Stats"]: stats, + tensors["dQ"]: d_q, + tensors["dK"]: d_k, + tensors["dV"]: d_v, + } + if "Bias" in tensors: + variant_pack[tensors["Bias"]] = attn_bias + if "dBias" in tensors: + variant_pack[tensors["dBias"]] = d_bias + graph_batch = tensors["_graph_batch"] + if "seq_len_q" in tensors: + variant_pack[tensors["seq_len_q"]] = _padded_sequence_lengths(cu_seqlens_q, graph_batch) + variant_pack[tensors["seq_len_kv"]] = _padded_sequence_lengths(cu_seqlens_kv, graph_batch) + if "offset_q" in tensors: + variant_pack[tensors["offset_q"]] = _element_ragged_offsets( + cu_seqlens_q_padded, graph_batch, q.stride(0) + ) + variant_pack[tensors["offset_o"]] = _element_ragged_offsets( + cu_seqlens_q_padded, graph_batch, o.stride(0) + ) + if "offset_k" in tensors: + variant_pack[tensors["offset_k"]] = _element_ragged_offsets( + cu_seqlens_kv_padded, graph_batch, k.stride(0) + ) + variant_pack[tensors["offset_v"]] = _element_ragged_offsets( + cu_seqlens_kv_padded, graph_batch, v.stride(0) + ) + if "offset_stats" in tensors: + stats_multiplier = stats.stride(0) + variant_pack[tensors["offset_stats"]] = _element_ragged_offsets( + cu_seqlens_q_padded, graph_batch, stats_multiplier + ) + if "dropout_seed" in tensors: + variant_pack[tensors["dropout_seed"]] = rng_state[:1] + variant_pack[tensors["dropout_offset"]] = rng_state[1:] + if "softmax_offset" in tensors: + variant_pack[tensors["softmax_offset"]] = softmax_offset + variant_pack[tensors["d_softmax_offset"]] = d_softmax_offset + entry.execute(variant_pack, q.device) + return d_q, d_k, d_v, d_bias, d_softmax_offset diff --git a/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py b/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py index 658dab5d88d..e59cb8fb619 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py @@ -944,7 +944,6 @@ def __init__( softmax_scale = 1.0 / math.sqrt( kv_channels if isinstance(kv_channels, int) else kv_channels[0] ) - self.softmax_scale = softmax_scale self.deterministic = ( not bool(int(os.getenv("NVTE_ALLOW_NONDETERMINISTIC_ALGO", "1"))) @@ -1789,11 +1788,7 @@ def _forward_thd_mask_types_grouped( policy_value = value_layer.index_select(0, kv_token_indices) _get_thd_policy_attention_backend( policy, - { - **grouped_attention_params_kwargs, - "num_tokens_q": policy_query.shape[0], - "num_tokens_kv": policy_key.shape[0], - }, + grouped_attention_params_kwargs, False, ) policy_output = self.forward( @@ -1990,11 +1985,11 @@ def forward( .. note:: - Users can use environment variables :attr:`NVTE_FLASH_ATTN`, :attr:`NVTE_FUSED_ATTN`, - and :attr:`NVTE_UNFUSED_ATTN` to control which DotProductAttention backend to use. - Transformer Engine first filters backends by support for the runtime environment - and input configuration, then applies a performance-based preference order. - On supported pre-Hopper GPUs, FlashAttention is preferred over FusedAttention and + Users can use environment variables :attr:`NVTE_FLASH_ATTN` and + :attr:`NVTE_FUSED_ATTN` to control which DotProductAttention backend to use. + Transformer Engine first filters backends by support for the runtime environment and + input configuration, then applies a performance-based preference order. On supported + pre-Hopper GPUs, FlashAttention is preferred over FusedAttention and UnfusedDotProductAttention when both optimized backends are eligible. On Hopper and newer GPUs, including Blackwell, FusedAttention is preferred over FlashAttention and UnfusedDotProductAttention when both optimized backends are eligible. @@ -2631,14 +2626,11 @@ def forward( # adjust max_seqlen and cu_seqlens for CP cp_size = 1 - cp_size_a2a = 1 if isinstance(self.cp_group, dist_group_type): cp_size = get_distributed_world_size(self.cp_group) elif isinstance(self.cp_group, list): for group in self.cp_group: cp_size *= get_distributed_world_size(group) - if self.cp_comm_type == "a2a+p2p" and len(self.cp_group) > 0: - cp_size_a2a = get_distributed_world_size(self.cp_group[0]) context_parallel = cp_size > 1 if thd_mask_policies is not None and context_parallel: raise ValueError("Mixed THD policies do not support context parallelism.") @@ -2700,11 +2692,32 @@ def forward( _alibi_cache["_alibi_slopes_require_update"] = True _alibi_cache["_alibi_bias_require_update"] = True - core_attention_bias_shape = ( - tuple(core_attention_bias.shape) - if core_attention_bias_type != "no_bias" and core_attention_bias is not None - else None - ) + # detect bias shape + core_attention_bias_shape = None + if core_attention_bias is not None: + if ( + core_attention_bias.shape[0] == batch_size + and core_attention_bias.shape[1] == query_layer.shape[-2] + ): + core_attention_bias_shape = "bhss" + elif ( + core_attention_bias.shape[0] == 1 + and core_attention_bias.shape[1] == query_layer.shape[-2] + ): + core_attention_bias_shape = "1hss" + elif ( + core_attention_bias.shape[0] == batch_size and core_attention_bias.shape[1] == 1 + ): + core_attention_bias_shape = "b1ss" + elif core_attention_bias.shape[0] == 1 and core_attention_bias.shape[1] == 1: + if core_attention_bias.shape[2] == 1: + core_attention_bias_shape = "111s" + else: + core_attention_bias_shape = "11ss" + else: + assert ( + False + ), "core_attention_bias must be in one of {bhss, 1hss, b1ss, 11ss, 111s} shapes" # Default pad_between_seqs auto-detect. For THD, infer presence of # inter-sequence padding from whether padded cu_seqlens were supplied -- @@ -2770,28 +2783,19 @@ def forward( "num_gqa_groups": num_gqa_groups, "max_seqlen_q": max_seqlen_q, "max_seqlen_kv": max_seqlen_kv, - "num_tokens_q": (query_layer.shape[0] if q_format == "thd" else 0), - "num_tokens_kv": (key_layer.shape[0] if kv_format == "thd" else 0), "head_dim_qk": head_dim_qk, "head_dim_v": head_dim_v, "softcap": softcap, - "alibi_slopes_shape": ( - alibi_slopes.shape - if core_attention_bias_type == "alibi" and alibi_slopes is not None - else None - ), + "alibi_slopes_shape": alibi_slopes.shape if alibi_slopes is not None else None, "core_attention_bias_type": core_attention_bias_type, "core_attention_bias_shape": core_attention_bias_shape, "core_attention_bias_requires_grad": ( - core_attention_bias.requires_grad - if core_attention_bias_type != "no_bias" and core_attention_bias is not None - else False + core_attention_bias.requires_grad if core_attention_bias is not None else False ), "attention_dropout": self.attention_dropout, "context_parallel": context_parallel, "cp_comm_type": self.cp_comm_type, "cp_size": cp_size, - "cp_size_a2a": cp_size_a2a, "deterministic": self.deterministic, "is_training": self.training, "fp8": self.fp8, @@ -2801,7 +2805,6 @@ def forward( "return_max_logit": self.return_max_logit, "cuda_graph": is_graph_capturing(), "num_splits": num_splits, - "softmax_scale": self.softmax_scale, "fp8_output": fp8_output, "checkpoint_core_attention": checkpoint_core_attention, "has_score_mod": score_mod is not None, diff --git a/transformer_engine/pytorch/attention/dot_product_attention/flex_attention.py b/transformer_engine/pytorch/attention/dot_product_attention/flex_attention.py index b9593b42d9b..4cd4adf225a 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/flex_attention.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/flex_attention.py @@ -5,26 +5,31 @@ """cuDNN-backed Flex Attention helpers.""" from dataclasses import dataclass -import importlib -import inspect from typing import Any, Callable, Dict, Optional, Tuple import torch -_cudnn_score_mod_handles: Dict[torch.device, Any] = {} +from transformer_engine.common.attention.cache_debug import record_event, record_lookup +from transformer_engine.common.attention.score_mod import ( + UNCACHEABLE_SCORE_MOD, + score_mod_callback_cache_key, +) + +from ._cudnn_graph import ( + current_stream_handle, + finalize_graph, + import_cudnn_frontend, + make_graph, + torch_to_cudnn_dtype, +) + _cudnn_score_mod_graph_cache: Dict[Tuple[Any, ...], Any] = {} -_SCORE_MOD_UNCACHEABLE = object() +_SCORE_MOD_UNCACHEABLE = UNCACHEABLE_SCORE_MOD def _import_cudnn_frontend(): - """Import the cuDNN frontend Python package.""" - try: - return importlib.import_module("cudnn") - except ImportError as exc: - raise ImportError( - "cuDNN frontend Python package not found. " - "Install it with: pip install nvidia-cudnn-frontend" - ) from exc + """Compatibility wrapper around the shared attention graph runtime.""" + return import_cudnn_frontend() def _bhsd_dim_stride( @@ -51,91 +56,13 @@ def _bhsd_graph_tensor(graph, tensor: torch.Tensor, tensor_format: str): # score_mod graph cache helpers. -def _freeze_score_mod_cache_key(value: Any) -> Any: - """Convert a user-provided score_mod graph key into a hashable structure.""" - if isinstance(value, torch.Tensor): - raise TypeError( - "score_mod_graph_cache_key() must not include tensors. Pass runtime tensors " - "through score_mod_tensors or score_mod_bprop_tensors instead." - ) - if isinstance(value, dict): - items = ( - ( - _freeze_score_mod_cache_key(key), - _freeze_score_mod_cache_key(val), - ) - for key, val in value.items() - ) - return tuple(sorted(items, key=repr)) - if isinstance(value, (list, tuple)): - return tuple(_freeze_score_mod_cache_key(item) for item in value) - if isinstance(value, (set, frozenset)): - items = (_freeze_score_mod_cache_key(item) for item in value) - return tuple(sorted(items, key=repr)) - try: - hash(value) - except TypeError as exc: - raise TypeError( - "score_mod_graph_cache_key() must return a hashable value or a nested " - "combination of dict/list/tuple/set values." - ) from exc - return value - - -def _score_mod_explicit_cache_key(callback_owner: Any) -> Optional[Any]: - """Return a user-provided structural graph key for a score_mod callback.""" - explicit_key = getattr(callback_owner, "score_mod_graph_cache_key", None) - if explicit_key is None: - return None - explicit_key = explicit_key() if callable(explicit_key) else explicit_key - return _freeze_score_mod_cache_key(explicit_key) - - def _score_mod_callback_cache_key(callback: Optional[Callable]) -> Any: - """Create a stable graph cache key for a score_mod callable. + """Compatibility wrapper around the shared score-modification key policy.""" - Module-level named functions are assumed to have stable topology. Anonymous functions - are keyed by code object because lambdas in the same module can share the same - qualname. Stateful bound methods and callable instances need an explicit - score_mod_graph_cache_key(); otherwise their graphs are left uncached to avoid reusing - stale graphs after Python object address reuse. - """ - if callback is None: - return None - self_obj = getattr(callback, "__self__", None) - func_obj = getattr(callback, "__func__", None) - if self_obj is not None and func_obj is not None: - explicit_key = _score_mod_explicit_cache_key(self_obj) - if explicit_key is None: - return _SCORE_MOD_UNCACHEABLE - return ( - "bound_method", - type(self_obj), - func_obj.__module__, - func_obj.__qualname__, - explicit_key, - ) - - explicit_key = _score_mod_explicit_cache_key(callback) - if explicit_key is not None: - return ( - "callable", - type(callback), - getattr(callback, "__module__", None), - getattr(callback, "__qualname__", None), - explicit_key, - ) - - if ( - inspect.isfunction(callback) - and callback.__closure__ is None - and "" not in callback.__qualname__ - ): - if callback.__name__ == "" or not callback.__qualname__: - return ("function", callback.__module__, callback.__code__) - return ("function", callback.__module__, callback.__qualname__) - - return _SCORE_MOD_UNCACHEABLE + return score_mod_callback_cache_key( + callback, + is_array=lambda item: isinstance(item, torch.Tensor), + ) def _score_mod_device_key(device: torch.device) -> Tuple[Any, ...]: @@ -193,41 +120,16 @@ def _wrapped_score_mod(sdpa_graph, score_tensor): def _get_cudnn_current_stream_handle(cudnn, device: torch.device): - """Return a cuDNN handle for device, bound to PyTorch's current stream.""" - if device.type != "cuda": - raise ValueError(f"Flex Attention only supports CUDA tensors, got device {device}.") - if device.index is None: - device = torch.device("cuda", torch.cuda.current_device()) - - handle = _cudnn_score_mod_handles.get(device) - with torch.cuda.device(device): - if handle is None: - handle = cudnn.create_handle() - _cudnn_score_mod_handles[device] = handle - - stream = torch.cuda.current_stream(device).cuda_stream - cudnn.set_stream(handle=handle, stream=stream) - return handle + """Compatibility wrapper around the shared current-stream handle.""" + del cudnn + return current_stream_handle(device) def _build_cudnn_pygraph(dtype: torch.dtype, device: torch.device): """Create a cuDNN frontend Python graph for F16/BF16 SDPA.""" - cudnn = _import_cudnn_frontend() - - if dtype == torch.float16: - io_data_type = cudnn.data_type.HALF - elif dtype == torch.bfloat16: - io_data_type = cudnn.data_type.BFLOAT16 - else: + if dtype not in (torch.float16, torch.bfloat16): raise ValueError(f"Flex Attention only supports FP16/BF16 tensors, got {dtype}.") - - graph = cudnn.pygraph( - io_data_type=io_data_type, - intermediate_data_type=cudnn.data_type.FLOAT, - compute_data_type=cudnn.data_type.FLOAT, - handle=_get_cudnn_current_stream_handle(cudnn, device), - ) - return graph + return make_graph(torch_to_cudnn_dtype(dtype), device, name="te_flex_attention") @dataclass @@ -263,19 +165,9 @@ class _CudnnScoreModBwdGraphEntry: workspace_size: int -def _finalize_cudnn_graph(graph) -> int: - """Build a cuDNN frontend Python graph and return its workspace size.""" - cudnn = _import_cudnn_frontend() - - graph.validate() - graph.build_operation_graph() - try: - graph.create_execution_plans([cudnn.heur_mode.A, cudnn.heur_mode.FALLBACK]) - graph.check_support() - except cudnn.cudnnGraphNotSupportedError as exc: - raise RuntimeError(f"cuDNN Flex Attention SDPA graph is not supported: {exc}") from exc - graph.build_plans(cudnn.build_plan_policy.HEURISTICS_CHOICE) - return max(graph.get_workspace_size(), 1) +def _finalize_cudnn_graph(graph, cache_site: Tuple[str, str]) -> int: + """Compatibility wrapper around shared graph finalization.""" + return finalize_graph(graph, cache_site=cache_site) def _execute_cudnn_graph( @@ -283,6 +175,7 @@ def _execute_cudnn_graph( variant_pack: Dict[Any, torch.Tensor], workspace_size: int, device: torch.device, + cache_site: Tuple[str, str], ): """Execute a built cuDNN frontend Python graph.""" cudnn = _import_cudnn_frontend() @@ -294,6 +187,7 @@ def _execute_cudnn_graph( device=device, dtype=torch.uint8, ) + record_event(*cache_site, "execute", device=_score_mod_device_key(device)[1]) graph.execute( variant_pack, workspace, @@ -316,8 +210,8 @@ def _cudnn_score_mod_fwd_cache_key( ) -> Optional[Tuple[Any, ...]]: """Pre-build cache key for score_mod fprop execution plans. - cuDNN exposes graph.key(), but only after graph construction has run the user callback. - This key avoids rebuilding the Python graph on cache hits. + cuDNN exposes graph.key(), but only after graph construction has run the user + callback. This key avoids rebuilding the Python graph on cache hits. """ score_mod_key = _score_mod_callback_cache_key(score_mod) if score_mod_key is _SCORE_MOD_UNCACHEABLE: @@ -423,7 +317,7 @@ def _build_cudnn_score_mod_fwd_graph( else: stats_tensor = None - workspace_size = _finalize_cudnn_graph(graph) + workspace_size = _finalize_cudnn_graph(graph, ("f16", "fwd")) return _CudnnScoreModFwdGraphEntry( graph=graph, q=q, @@ -465,11 +359,14 @@ def _get_cudnn_score_mod_fwd_graph( ) key = _cudnn_score_mod_fwd_cache_key(*build_args) if key is None: + record_lookup("f16", "fwd", hit=False, key="uncacheable score_mod") return _build_cudnn_score_mod_fwd_graph(*build_args) entry = _cudnn_score_mod_graph_cache.get(key) + record_lookup("f16", "fwd", hit=entry is not None, key=key) if entry is None: entry = _build_cudnn_score_mod_fwd_graph(*build_args) _cudnn_score_mod_graph_cache[key] = entry + record_event("f16", "fwd", "cache_graph") return entry @@ -531,7 +428,7 @@ def _build_cudnn_score_mod_bwd_graph( dk.set_output(True).set_dim(dk_dim).set_stride(dk_stride) dv.set_output(True).set_dim(dv_dim).set_stride(dv_stride) - workspace_size = _finalize_cudnn_graph(graph) + workspace_size = _finalize_cudnn_graph(graph, ("f16", "bwd")) return _CudnnScoreModBwdGraphEntry( graph=graph, q=q, @@ -584,11 +481,14 @@ def _get_cudnn_score_mod_bwd_graph( ) key = _cudnn_score_mod_bwd_cache_key(*build_args) if key is None: + record_lookup("f16", "bwd", hit=False, key="uncacheable score_mod") return _build_cudnn_score_mod_bwd_graph(*build_args) entry = _cudnn_score_mod_graph_cache.get(key) + record_lookup("f16", "bwd", hit=entry is not None, key=key) if entry is None: entry = _build_cudnn_score_mod_bwd_graph(*build_args) _cudnn_score_mod_graph_cache[key] = entry + record_event("f16", "bwd", "cache_graph") return entry @@ -655,6 +555,7 @@ def forward( variant_pack, entry.workspace_size, query_layer.device, + ("f16", "fwd"), ) ctx.is_training = is_training @@ -742,6 +643,7 @@ def backward(ctx, d_out: torch.Tensor): variant_pack, entry.workspace_size, query_layer.device, + ("f16", "bwd"), ) return ( diff --git a/transformer_engine/pytorch/attention/dot_product_attention/utils.py b/transformer_engine/pytorch/attention/dot_product_attention/utils.py index 23b89287dfd..01a8d2a237e 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/utils.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/utils.py @@ -22,12 +22,7 @@ import torch.nn.functional as F import transformer_engine_torch as tex import transformer_engine as te -from transformer_engine.pytorch.cpp_extensions.fused_attn import ( - QKVLayout, - QKVFormat, - AttnBiasType, - AttnMaskType, - SoftmaxType, +from transformer_engine.pytorch.attention.dot_product_attention.cudnn_attention import ( FusedAttnBackend, META_QKV, META_DQKV, @@ -36,6 +31,9 @@ META_S, META_DP, ) +from transformer_engine.pytorch.attention.dot_product_attention._cudnn_backend import ( + get_fused_attn_backend as get_cudnn_fused_attn_backend, +) from transformer_engine.pytorch.attention.inference import InferenceParams from transformer_engine.pytorch.cpu_offload import is_cpu_offload_enabled from transformer_engine.pytorch.quantized_tensor import QuantizedTensorStorage @@ -49,7 +47,7 @@ from transformer_engine.pytorch.tensor.storage.mxfp8_tensor_storage import MXFP8TensorStorage from transformer_engine.pytorch.quantization import get_fp8_te_dtype -from transformer_engine.pytorch.constants import TE_DType, DType, MXFP8_BLOCK_SCALING_SIZE +from transformer_engine.pytorch.constants import TE_DType, MXFP8_BLOCK_SCALING_SIZE from transformer_engine.pytorch.utils import ( @@ -208,9 +206,6 @@ class AttentionParams: Type of query/key/value tensors, {`torch.Tensor`, `Float8Tensor`}. qkv_dtype : torch.dtype, default = torch.bfloat16 Data type of query/key/value tensors. - nominal_dtype : Optional[torch.dtype], default = None - Model precision (F16/BF16) of the unquantized tensors (O, and dQ/dK/dV under - current/mxfp8) when `qkv_dtype` itself is FP8. qkv_layout : str, default = "sbh3d" Query/key/value tensor memory layout. batch_size : int, default = 1 @@ -223,10 +218,6 @@ class AttentionParams: Maximum sequence length of the query tensor. max_seqlen_kv : int, default = 128 Maximum sequence length of the key and value tensors. - num_tokens_q : int, default = 0 - Total number of query tokens in a batch, when `qkv_format=thd`. - num_tokens_kv : int, default = 0 - Total number of key/value tokens in a batch, when `qkv_format=thd`. head_dim_qk : int, default = 64 The size of each attention head in query and key tensors. head_dim_v : int, default = 64 @@ -236,7 +227,7 @@ class AttentionParams: `causal_bottom_right`, `padding_causal_bottom_right`, `arbitrary`} window_size : Tuple[int, int], default = None Sliding window attention size. - bottom_right_diagonal: bool, default = `True` + bottom_right_diagonal: bool, default = `None` Whether to align sliding window and ALiBi diagonal to the bottom right corner of the softmax matrix. softcap : float, default = 0.0 @@ -246,8 +237,8 @@ class AttentionParams: Tensor shape of :attr:`alibi_slopes` in `DotProductAttention`. core_attention_bias_type : str, default = no_bias Attention bias type, {`no_bias`, `pre_scale_bias`, `post_scale_bias`, `alibi`}. - core_attention_bias_shape : Optional[Tuple[int, int, int, int]], default = None - Attention bias shape, (b, h, sq, skv). + core_attention_bias_shape : str, default = 1hss + Attention bias shape, {`1hss`, `b1ss`, `bhss`}. core_attention_bias_requires_grad : bool, default = True Whether attention bias requires gradient. pad_between_seqs : bool, default = False @@ -260,9 +251,7 @@ class AttentionParams: cp_comm_type : str, default = "p2p" The communication type of context parallelism. cp_size : int, default = 1 - The (total) group size of context parallelism. - cp_size_a2a : int, default = 1 - The all-to-all subgroup size when `cp_comm_type == "a2a+p2p"`. + The group size of context parallelism. deterministic : bool, default = False Whether to run `DotProductAttention` with determinism or not. is_training : bool, default = True @@ -281,8 +270,6 @@ class AttentionParams: Whether support for cuda graph capture is needed or not. num_splits : int, default = 1 The number of kernels to split attention to. - softmax_scale : float, default = 1.0 - Pre-softmax attention scale. fp8_output : bool, default = False Whether output is requested in FP8. checkpoint_core_attention : bool, default = False @@ -295,15 +282,12 @@ class AttentionParams: qkv_type: Union[torch.Tensor, Float8Tensor] = torch.Tensor qkv_dtype: torch.dtype = torch.bfloat16 - nominal_dtype: Optional[torch.dtype] = None qkv_layout: str = "sbh3d" batch_size: int = 1 num_heads: int = 16 num_gqa_groups: int = 16 max_seqlen_q: int = 128 max_seqlen_kv: int = 128 - num_tokens_q: int = 0 - num_tokens_kv: int = 0 head_dim_qk: int = 64 head_dim_v: int = 64 attn_mask_type: str = "no_mask" @@ -312,14 +296,13 @@ class AttentionParams: softcap: float = 0.0 alibi_slopes_shape: Union[torch.Size, List, None] = None core_attention_bias_type: str = "no_bias" - core_attention_bias_shape: Union[Tuple[int, int, int, int], None] = None + core_attention_bias_shape: str = "1hss" core_attention_bias_requires_grad: bool = True pad_between_seqs: bool = False attention_dropout: float = 0.0 context_parallel: bool = False cp_comm_type: str = "p2p" cp_size: int = 1 - cp_size_a2a: int = 1 deterministic: bool = False is_training: bool = True fp8: bool = False @@ -329,7 +312,6 @@ class AttentionParams: return_max_logit: bool = False cuda_graph: bool = False num_splits: int = 1 - softmax_scale: float = 1.0 fp8_output: bool = False checkpoint_core_attention: bool = False has_score_mod: bool = False @@ -354,67 +336,6 @@ def __eq__(self, other): return True -@dataclass(eq=True) -class FusedAttentionParams: - """ - Attention parameters used by the `FusedAttention` backend. - """ - - # basic attention settings - is_training: bool = True - deterministic: bool = False - cuda_graph: bool = False - return_max_logit: bool = False - attn_mask_type: tex.NVTE_Mask_Type = tex.NVTE_Mask_Type.NVTE_NO_MASK - bias_type: tex.NVTE_Bias_Type = tex.NVTE_Bias_Type.NVTE_NO_BIAS - window_size_left: int = -1 - window_size_right: int = -1 - bottom_right_diagonal: bool = True - softmax_type: tex.NVTE_Softmax_Type = tex.NVTE_Softmax_Type.NVTE_VANILLA_SOFTMAX - scaling_mode: tex.NVTEScalingMode = tex.NVTEScalingMode.NVTE_DELAYED_TENSOR_SCALING - dropout: float = 0.0 - attn_scale: float = 1.0 - - # tensor types - qkv_dtype: DType = DType.kBFloat16 - o_dtype: DType = DType.kBFloat16 - do_dtype: DType = DType.kBFloat16 - dqkv_dtype: DType = DType.kBFloat16 - - # tensor layouts - qkv_layout: tex.NVTE_QKV_Layout = tex.NVTE_QKV_Layout.NVTE_QKV_Layout_NOT_SET - o_format: tex.NVTE_QKV_Format = tex.NVTE_QKV_Format.NVTE_QKV_Format_NOT_SET - do_format: tex.NVTE_QKV_Format = tex.NVTE_QKV_Format.NVTE_QKV_Format_NOT_SET - dqkv_layout: tex.NVTE_QKV_Layout = tex.NVTE_QKV_Layout.NVTE_QKV_Layout_NOT_SET - qkv_scale_inv_format: tex.NVTE_QKV_Format = tex.NVTE_QKV_Format.NVTE_QKV_Format_NOT_SET - do_scale_inv_format: tex.NVTE_QKV_Format = tex.NVTE_QKV_Format.NVTE_QKV_Format_NOT_SET - - # tensor dimensions - batch_size: int = 0 - num_attn_heads: int = 0 - num_gqa_groups: int = 0 - head_dim_qk: int = 0 - head_dim_v: int = 0 - max_seqlen_q: int = 0 - max_seqlen_kv: int = 0 - num_tokens_q: int = 0 - num_tokens_kv: int = 0 - - # paged KV dimensions - num_pages_k: int = 0 - num_pages_v: int = 0 - page_size_k: int = 0 - page_size_v: int = 0 - max_pages_per_seq_k: int = 0 - max_pages_per_seq_v: int = 0 - - # bias dimensions - bias_batch_size: int = 0 - bias_num_heads: int = 0 - bias_seqlen_q: int = 0 - bias_seqlen_kv: int = 0 - - class _NoOpLogger: """ Stand-in for the "DotProductAttention" logger used when get_attention_backend @@ -439,20 +360,38 @@ def error(self, *args, **kwargs): @torch.compiler.assume_constant_result -def _get_fused_attn_backend(**fused_attn_kwargs): - """Constant-foldable tex.get_fused_attn_backend: the result depends only on - the attention config. - - Returns a plain int rather than a FusedAttnBackend member: dynamo - reconstructs the result of an assume_constant_result call by re-emitting the - call, which is only valid inside the frame that made it. An int survives a - graph break because it is baked into the graph as a literal, while an enum - member comes out of the reconstruction corrupted (see the cast at the call +def _get_fused_attn_backend( + is_training, + q_type, + kv_type, + qkv_layout, + bias_type, + attn_mask_type, + softmax_type, + *args, +): + """Constant-foldable Python backend selector: the result depends only on + the attention config. Layout/bias/mask/softmax are taken as their string + keys and resolved to the pybind enums here, so that every argument is a + python literal or a python enum. + + Returns a plain int rather than a FusedAttnBackend member alongside the rejection + reason. Dynamo reconstructs the result of an assume_constant_result call by + re-emitting the call, which is only valid inside the frame that made it. An int + survives a graph break because it is baked into the graph as a literal, while an + enum member comes out of the reconstruction corrupted (see the cast at the call site, which restores the enum).""" - fused_attention_backend, reject_message = tex.get_fused_attn_backend( - FusedAttentionParams(**fused_attn_kwargs) + backend, reason = get_cudnn_fused_attn_backend( + is_training, + q_type, + kv_type, + qkv_layout, + bias_type, + attn_mask_type, + softmax_type, + *args, ) - return int(fused_attention_backend), reject_message + return int(backend), reason def get_attention_backend( @@ -487,15 +426,12 @@ def get_attention_backend( # is shifted over to the caller of this function qkv_type = attention_params.qkv_type qkv_dtype = attention_params.qkv_dtype - nominal_dtype = attention_params.nominal_dtype qkv_layout = attention_params.qkv_layout batch_size = attention_params.batch_size num_heads = attention_params.num_heads num_gqa_groups = attention_params.num_gqa_groups max_seqlen_q = attention_params.max_seqlen_q max_seqlen_kv = attention_params.max_seqlen_kv - num_tokens_q = attention_params.num_tokens_q - num_tokens_kv = attention_params.num_tokens_kv head_dim_qk = attention_params.head_dim_qk head_dim_v = attention_params.head_dim_v attn_mask_type = attention_params.attn_mask_type @@ -510,8 +446,7 @@ def get_attention_backend( attention_dropout = attention_params.attention_dropout context_parallel = attention_params.context_parallel cp_comm_type = attention_params.cp_comm_type - cp_size = attention_params.cp_size - cp_size_a2a = attention_params.cp_size_a2a + cp_size = attention_params.cp_size # pylint: disable=unused-variable deterministic = attention_params.deterministic is_training = attention_params.is_training fp8 = attention_params.fp8 @@ -521,7 +456,6 @@ def get_attention_backend( return_max_logit = attention_params.return_max_logit cuda_graph = attention_params.cuda_graph num_splits = attention_params.num_splits - softmax_scale = attention_params.softmax_scale fp8_output = attention_params.fp8_output checkpoint_core_attention = attention_params.checkpoint_core_attention has_score_mod = attention_params.has_score_mod @@ -1170,7 +1104,7 @@ def _is_fa3_supported(num_heads, num_gqa_groups, head_dim_qk, head_dim_v, qkv_dt use_flash_attention_4 = False # Filter: QKV layout - if "thd" in (q_format, kv_format): + if qkv_format == "thd": if pad_between_seqs: if ( # pylint: disable=too-many-boolean-expressions use_flash_attention_2 and FlashAttentionUtils.is_installed @@ -1188,7 +1122,7 @@ def _is_fa3_supported(num_heads, num_gqa_groups, head_dim_qk, head_dim_v, qkv_dt if cudnn_version < (9, 18, 1): if use_fused_attention: logger.debug( - "Disabling FusedAttention when Q or KV uses THD because it is" + "Disabling FusedAttention as qkv_format = thd is" " not supported for compute capability = sm120 and cuDNN version < 9.18.1" ) use_fused_attention = False @@ -1538,186 +1472,68 @@ def _is_fa3_supported(num_heads, num_gqa_groups, head_dim_qk, head_dim_v, qkv_dt fu_core_attention_bias_requires_grad = False if len(alibi_slopes_shape) == 1 and alibi_slopes_shape[0] == num_heads: - fu_core_attention_bias_shape = (1, num_heads, max_seqlen_q, max_seqlen_kv) + fu_core_attention_bias_shape = "1hss" elif ( len(alibi_slopes_shape) == 2 and alibi_slopes_shape[0] == batch_size and alibi_slopes_shape[1] == num_heads ): - fu_core_attention_bias_shape = (batch_size, num_heads, max_seqlen_q, max_seqlen_kv) + fu_core_attention_bias_shape = "bhss" - fu_core_attention_bias_shape_type = None - if ( - fu_core_attention_bias_type == "post_scale_bias" - and fu_core_attention_bias_shape is not None - ): - b, h, sq, _skv = fu_core_attention_bias_shape - if b == batch_size and h == num_heads: - fu_core_attention_bias_shape_type = "bhss" - elif b == 1 and h == num_heads: - fu_core_attention_bias_shape_type = "1hss" - elif b == batch_size and h == 1: - fu_core_attention_bias_shape_type = "b1ss" - elif b == 1 and h == 1: - fu_core_attention_bias_shape_type = "111s" if sq == 1 and max_seqlen_q != 1 else "11ss" - else: - raise ValueError( - "core_attention_bias tensor must be in one of " - '{"bhss", "1hss", "b1ss", "11ss", "111s"} shapes. ' - f"Found (b,h,sq,skv) = ({b},{h},{sq},{_skv})" - ) if ( use_fused_attention and fu_core_attention_bias_type == "post_scale_bias" - and fu_core_attention_bias_shape_type != "1hss" + and fu_core_attention_bias_shape != "1hss" ): # dbias calculation is not supported for 111s as of cuDNN 9.18. So, use fused attention backend only if bias does not require grad. - if fu_core_attention_bias_requires_grad and fu_core_attention_bias_shape_type == "111s": + if fu_core_attention_bias_requires_grad and fu_core_attention_bias_shape == "111s": logger.warning( "Disabling FusedAttention as dbias calculation is not supported for 111s" ) use_fused_attention = False - # Filter: cuDNN support fused_attention_backend = None if use_fused_attention: - recipe = fp8_meta["recipe"] if (fp8 and fp8_meta["recipe"].fp8_dpa) else None - cs_o_in_f16 = os.getenv("NVTE_DPA_FP8CS_O_in_F16", "1") == "1" - spec = get_fused_attn_spec( - recipe, qkv_dtype, qkv_layout, cs_o_in_f16=cs_o_in_f16, nominal_dtype=nominal_dtype + # ``DType`` is implicitly convertible to ``transformer_engine::DType`` + # on the C++ side, so pass it straight to the pybind function. + q_type = TE_DType[qkv_dtype] + kv_type = q_type + if fp8 and fp8_meta["recipe"].fp8_dpa: + q_type = get_fp8_te_dtype(fp8_meta["recipe"], fprop_tensor=True) + kv_type = q_type + # NOTE: under torch.compile the numeric args below must not be symbolic + # (assume_constant_result requires concrete values); ints/floats made + # dynamic by automatic dynamic currently graph break here. + fused_attention_backend, fused_attention_rejection_reason = _get_fused_attn_backend( + is_training, + q_type, + kv_type, + qkv_layout, + fu_core_attention_bias_type, + attn_mask_type, + softmax_type, + attention_dropout, + num_heads, + num_gqa_groups, + max_seqlen_q, + max_seqlen_kv, + head_dim_qk, + head_dim_v, + window_size[0], + window_size[1], + return_max_logit, + cuda_graph, + deterministic, ) - qkv_type, o_type, do_type, dqkv_type = spec.qkv, spec.o, spec.do, spec.dqkv - scaling_mode = spec.scaling_mode - qkv_scale_inv_format = spec.scale_inv_format - do_scale_inv_format = spec.scale_inv_format - o_format = spec.o_format - do_format = spec.do_format - dqkv_layout = spec.dqkv_layout - num_pages_k = num_pages_v = 0 - page_size_k = page_size_v = 0 - max_pages_per_seq_k = max_pages_per_seq_v = 0 - if inference_params is not None and getattr(inference_params, "is_paged", False): - num_pages_k = num_pages_v = inference_params.total_num_pages - page_size_k = page_size_v = inference_params.page_size - max_pages_per_seq_k = max_pages_per_seq_v = ( - inference_params.cache_manager.max_pages_per_seq - ) - bias_batch_size = bias_num_heads = bias_seqlen_q = bias_seqlen_kv = 0 - if fu_core_attention_bias_shape is not None: - bias_batch_size, bias_num_heads, bias_seqlen_q, bias_seqlen_kv = ( - fu_core_attention_bias_shape - ) - base_fused_attn_kwargs = { - "is_training": is_training, - "deterministic": deterministic, - "cuda_graph": cuda_graph, - "return_max_logit": return_max_logit, - "attn_mask_type": AttnMaskType[attn_mask_type], - "bias_type": AttnBiasType[fu_core_attention_bias_type], - "window_size_left": window_size[0], - "window_size_right": window_size[1], - "bottom_right_diagonal": bottom_right_diagonal, - "softmax_type": SoftmaxType[softmax_type], - "scaling_mode": scaling_mode, - "dropout": attention_dropout, - "attn_scale": softmax_scale, - "qkv_dtype": qkv_type, - "o_dtype": o_type, - "do_dtype": do_type, - "dqkv_dtype": dqkv_type, - "qkv_layout": QKVLayout[spec.qkv_layout], - "o_format": QKVFormat[o_format], - "do_format": QKVFormat[do_format], - "dqkv_layout": QKVLayout[dqkv_layout], - "qkv_scale_inv_format": QKVFormat[qkv_scale_inv_format], - "do_scale_inv_format": QKVFormat[do_scale_inv_format], - "batch_size": batch_size, - "num_attn_heads": num_heads, - "num_gqa_groups": num_gqa_groups, - "head_dim_qk": head_dim_qk, - "head_dim_v": head_dim_v, - "max_seqlen_q": max_seqlen_q, - "max_seqlen_kv": max_seqlen_kv, - "num_tokens_q": num_tokens_q, - "num_tokens_kv": num_tokens_kv, - "num_pages_k": num_pages_k, - "num_pages_v": num_pages_v, - "page_size_k": page_size_k, - "page_size_v": page_size_v, - "max_pages_per_seq_k": max_pages_per_seq_k, - "max_pages_per_seq_v": max_pages_per_seq_v, - "bias_batch_size": bias_batch_size, - "bias_num_heads": bias_num_heads, - "bias_seqlen_q": bias_seqlen_q, - "bias_seqlen_kv": bias_seqlen_kv, - } - # Context-parallel per-step configs - if context_parallel: - from transformer_engine.pytorch.attention.dot_product_attention.context_parallel import ( - cp_per_step_configs, - ) - - per_step_configs = cp_per_step_configs( - cp_comm_type, - cp_size, - cp_size_a2a, - max_seqlen_q=max_seqlen_q, - max_seqlen_kv=max_seqlen_kv, - num_tokens_q=num_tokens_q, - num_tokens_kv=num_tokens_kv, - num_heads=num_heads, - num_gqa_groups=num_gqa_groups, - attn_mask_type=attn_mask_type, - window_size=window_size, - bottom_right_diagonal=bottom_right_diagonal, + if fused_attention_backend == FusedAttnBackend.No_Backend.value: + logger.debug( + "Disabling FusedAttention: %s", + fused_attention_rejection_reason or "no backend supports the provided input", ) - else: - per_step_configs = [None] - - for step_config in per_step_configs: - fused_attn_kwargs = dict(base_fused_attn_kwargs) - if step_config is not None: - step_seqlen_q = step_config["max_seqlen_q"] - step_seqlen_kv = step_config["max_seqlen_kv"] - fused_attn_kwargs.update( - attn_mask_type=AttnMaskType[step_config["attn_mask_type"]], - max_seqlen_q=step_seqlen_q, - max_seqlen_kv=step_seqlen_kv, - num_tokens_q=step_config["num_tokens_q"], - num_tokens_kv=step_config["num_tokens_kv"], - num_attn_heads=step_config["num_attn_heads"], - num_gqa_groups=step_config["num_gqa_groups"], - window_size_left=step_config["window_size_left"], - window_size_right=step_config["window_size_right"], - bottom_right_diagonal=step_config["bottom_right_diagonal"], - ) - if fu_core_attention_bias_shape is not None: - if bias_seqlen_q != 1: - fused_attn_kwargs["bias_seqlen_q"] = step_seqlen_q - if bias_seqlen_kv != 1: - fused_attn_kwargs["bias_seqlen_kv"] = step_seqlen_kv - # NOTE: under torch.compile the numeric entries of fused_attn_kwargs must not be - # symbolic (assume_constant_result requires concrete values); ints/floats made - # dynamic by automatic dynamic currently graph break here. - fused_attention_backend, reject_message = _get_fused_attn_backend(**fused_attn_kwargs) - if fused_attention_backend == FusedAttnBackend.No_Backend.value: - logger.debug( - "Disabling FusedAttention: %s%s", - reject_message, - ( - f" (context-parallel per-step config {step_config})" - if step_config is not None - else "" - ), - ) - use_fused_attention = False - fused_attention_backend = None - break - - if ( - use_fused_attention - and has_score_mod - and fused_attention_backend != FusedAttnBackend.F16_arbitrary_seqlen.value + use_fused_attention = False + fused_attention_backend = None + elif ( + has_score_mod and fused_attention_backend != FusedAttnBackend.F16_arbitrary_seqlen.value ): logger.debug( "Disabling FusedAttention for score_mod because sub-backend %s is not " @@ -2662,102 +2478,6 @@ def get_qkv_format( return qkv_format, q_format, kv_format -@dataclass(frozen=True) -class FusedAttnSpec: - """Fused-attention spec for a given config. - - Mirrors what `FusedAttnFunc` feeds `fused_attn_fwd`/`fused_attn_bwd` (backends.py), - so the availability probe (`get_attention_backend`) cannot drift from runtime. - """ - - scaling_mode: Any - qkv: Any - o: Any - do: Any - dqkv: Any - scale_inv_format: Optional[str] - qkv_layout: str - o_format: str - do_format: str - dqkv_layout: str - - -def get_fused_attn_spec(recipe, qkv_dtype, qkv_layout, *, cs_o_in_f16, nominal_dtype=None): - """Resolve fused-attention specs, e.g. tensor dtypes, formats, for a given config""" - q_format = get_qkv_format(qkv_layout)[1] - eff_qkv_layout = qkv_layout # FP16/BF16 - if recipe is not None: - if not recipe.mxfp8(): - # Delayed/current scaling - eff_qkv_layout = qkv_layout.replace("paged_kv_", "") - elif qkv_layout in ("bshd_bshd_bshd", "sbhd_sbhd_sbhd"): - eff_qkv_layout = qkv_layout # MXFP8 fast path - else: - eff_qkv_layout = "bhsd_bhsd_bhsd" # MXFP8 slow path - layout_kwargs = { - "qkv_layout": eff_qkv_layout, - "o_format": q_format, - "do_format": q_format, - "dqkv_layout": qkv_layout, - } - - if qkv_dtype in (torch.float8_e4m3fn, torch.float8_e5m2): - ref = TE_DType[nominal_dtype if nominal_dtype is not None else torch.bfloat16] - else: - ref = TE_DType[qkv_dtype] - - # FP16/BF16: every tensor is in model precision; scaling_mode is a placeholder - if recipe is None: - return FusedAttnSpec( - tex.NVTEScalingMode.NVTE_DELAYED_TENSOR_SCALING, - ref, - ref, - ref, - ref, - None, - **layout_kwargs, - ) - - fprop_fp8 = get_fp8_te_dtype(recipe, fprop_tensor=True) - grad_fp8 = get_fp8_te_dtype(recipe, fprop_tensor=False) - - # MXFP8 block scaling: Q/K/V/dO are in MXFP8; O/dQ/dK/dV stay in model precision - if recipe.mxfp8(): - return FusedAttnSpec( - tex.NVTEScalingMode.NVTE_MXFP8_1D_SCALING, - fprop_fp8, - ref, - grad_fp8, - ref, - "bhsd", - **layout_kwargs, - ) - - # FP8 current scaling: Q/K/V/dO are in FP8; O in model precision if `cs_o_in_f16` (default), otherwise FP8; - # dQ/dK/dV in model precision - if recipe.float8_current_scaling(): - return FusedAttnSpec( - tex.NVTEScalingMode.NVTE_DELAYED_TENSOR_SCALING, - fprop_fp8, - ref if cs_o_in_f16 else fprop_fp8, - grad_fp8, - ref, - None, - **layout_kwargs, - ) - - # FP8 delayed scaling: Q/K/V/O are in FP8 (e.g. E4M3); dO/dQ/dK/dV in FP8 (e.g. E5M2) - return FusedAttnSpec( - tex.NVTEScalingMode.NVTE_DELAYED_TENSOR_SCALING, - fprop_fp8, - fprop_fp8, - grad_fp8, - grad_fp8, - None, - **layout_kwargs, - ) - - def qkv_layout_needs_detection(*qkv: Optional[torch.Tensor]) -> bool: """Whether the layout of these q/k/v can only be told by inspecting memory. diff --git a/transformer_engine/pytorch/attention/fused_mla_q_uproj.py b/transformer_engine/pytorch/attention/fused_mla_q_uproj.py index b15e53e6ee3..2e55ddedb35 100644 --- a/transformer_engine/pytorch/attention/fused_mla_q_uproj.py +++ b/transformer_engine/pytorch/attention/fused_mla_q_uproj.py @@ -203,7 +203,9 @@ def run( from cuda.bindings import driver as cuda - stream = cuda.CUstream(torch.cuda.current_stream(x.device).cuda_stream) + stream = cuda.CUstream( # pylint: disable=c-extension-no-member + torch.cuda.current_stream(x.device).cuda_stream + ) wrapper = cls._kernel() if isinstance(w, QuantizedTensor): @@ -343,7 +345,7 @@ def wrap_mxfp8( blk = MXFP8_BLOCK_SCALING_SIZE # Both rowwise and columnwise Q are required: # - Forward QK^T uses rowwise - # - cuDNN backward (fused_attn_fp8_bwd_impl) requires columnwise for dK gradient + # - The cuDNN backward graph requires columnwise for the dK gradient quantizer = MXFP8Quantizer(fp8_dtype=tex.DType.kFloat8E4M3, rowwise=True, columnwise=True) return MXFP8Tensor( shape=(s, b, nh, d), diff --git a/transformer_engine/pytorch/cpp_extensions/fused_attn.py b/transformer_engine/pytorch/cpp_extensions/fused_attn.py index 9a33df7634b..c509dadfd2e 100644 --- a/transformer_engine/pytorch/cpp_extensions/fused_attn.py +++ b/transformer_engine/pytorch/cpp_extensions/fused_attn.py @@ -2,30 +2,20 @@ # # See LICENSE for license information. -"""Python interface for fused attention extensions""" +"""Compatibility imports for the Python cuDNN attention implementation.""" -import math from enum import IntEnum -from typing import Tuple, List, Union, Optional + import torch -import transformer_engine_torch as tex -from transformer_engine_torch import ( - NVTE_QKV_Layout, - NVTE_QKV_Format, - NVTE_Bias_Type, - NVTE_Mask_Type, - NVTE_Softmax_Type, - NVTE_Fused_Attn_Backend, -) -from ..quantized_tensor import Quantizer -from ..constants import FP8BwdTensorIdx, FP8FwdTensorIdx, DType +from transformer_engine_torch import NVTE_QKV_Format + +from ..constants import DType, FP8BwdTensorIdx, FP8FwdTensorIdx -__all__ = [ - "fused_attn_fwd", - "fused_attn_bwd", -] +__all__ = ["fused_attn_fwd", "fused_attn_bwd"] +# Retained for rotary-position and custom-op compatibility. These operations +# still use TE's format enum even though attention execution itself does not. TORCH_DType = { DType.kFloat8E4M3: torch.uint8, DType.kFloat8E5M2: torch.uint8, @@ -47,122 +37,21 @@ "bhsd": NVTE_QKV_Format.NVTE_BHSD, } -QKVLayout = { - "sb3hd": NVTE_QKV_Layout.NVTE_SB3HD, - "sbh3d": NVTE_QKV_Layout.NVTE_SBH3D, - "sbhd_sb2hd": NVTE_QKV_Layout.NVTE_SBHD_SB2HD, - "sbhd_sbh2d": NVTE_QKV_Layout.NVTE_SBHD_SBH2D, - "sbhd_sbhd_sbhd": NVTE_QKV_Layout.NVTE_SBHD_SBHD_SBHD, - "bs3hd": NVTE_QKV_Layout.NVTE_BS3HD, - "bsh3d": NVTE_QKV_Layout.NVTE_BSH3D, - "bshd_bs2hd": NVTE_QKV_Layout.NVTE_BSHD_BS2HD, - "bshd_bsh2d": NVTE_QKV_Layout.NVTE_BSHD_BSH2D, - "bshd_bshd_bshd": NVTE_QKV_Layout.NVTE_BSHD_BSHD_BSHD, - "t3hd": NVTE_QKV_Layout.NVTE_T3HD, - "th3d": NVTE_QKV_Layout.NVTE_TH3D, - "thd_t2hd": NVTE_QKV_Layout.NVTE_THD_T2HD, - "thd_th2d": NVTE_QKV_Layout.NVTE_THD_TH2D, - "thd_thd_thd": NVTE_QKV_Layout.NVTE_THD_THD_THD, - "sbhd_bshd_bshd": NVTE_QKV_Layout.NVTE_SBHD_BSHD_BSHD, - "bshd_sbhd_sbhd": NVTE_QKV_Layout.NVTE_BSHD_SBHD_SBHD, - "thd_bshd_bshd": NVTE_QKV_Layout.NVTE_THD_BSHD_BSHD, - "thd_sbhd_sbhd": NVTE_QKV_Layout.NVTE_THD_SBHD_SBHD, - "paged_kv_bshd_bshd_bshd": NVTE_QKV_Layout.NVTE_Paged_KV_BSHD_BSHD_BSHD, - "paged_kv_bshd_sbhd_sbhd": NVTE_QKV_Layout.NVTE_Paged_KV_BSHD_SBHD_SBHD, - "paged_kv_sbhd_bshd_bshd": NVTE_QKV_Layout.NVTE_Paged_KV_SBHD_BSHD_BSHD, - "paged_kv_sbhd_sbhd_sbhd": NVTE_QKV_Layout.NVTE_Paged_KV_SBHD_SBHD_SBHD, - "paged_kv_thd_bshd_bshd": NVTE_QKV_Layout.NVTE_Paged_KV_THD_BSHD_BSHD, - "paged_kv_thd_sbhd_sbhd": NVTE_QKV_Layout.NVTE_Paged_KV_THD_SBHD_SBHD, - "bhsd_bhsd_bhsd": NVTE_QKV_Layout.NVTE_BHSD_BHSD_BHSD, -} - -AttnBiasType = { - "no_bias": NVTE_Bias_Type.NVTE_NO_BIAS, - "pre_scale_bias": NVTE_Bias_Type.NVTE_PRE_SCALE_BIAS, - "post_scale_bias": NVTE_Bias_Type.NVTE_POST_SCALE_BIAS, - "alibi": NVTE_Bias_Type.NVTE_ALIBI, -} - -AttnMaskType = { - "no_mask": NVTE_Mask_Type.NVTE_NO_MASK, - "padding": NVTE_Mask_Type.NVTE_PADDING_MASK, - "causal": NVTE_Mask_Type.NVTE_CAUSAL_MASK, - "padding_causal": NVTE_Mask_Type.NVTE_PADDING_CAUSAL_MASK, - "causal_bottom_right": NVTE_Mask_Type.NVTE_CAUSAL_BOTTOM_RIGHT_MASK, - "padding_causal_bottom_right": NVTE_Mask_Type.NVTE_PADDING_CAUSAL_BOTTOM_RIGHT_MASK, -} - -SoftmaxType = { - "vanilla": NVTE_Softmax_Type.NVTE_VANILLA_SOFTMAX, - "off-by-one": NVTE_Softmax_Type.NVTE_OFF_BY_ONE_SOFTMAX, - "learnable": NVTE_Softmax_Type.NVTE_LEARNABLE_SOFTMAX, -} - class FusedAttnBackend(IntEnum): - """Fused attention sub-backends. - - This is the canonical fused-attention backend enum for - ``transformer_engine.pytorch``. It mirrors the backend - ``transformer_engine_torch.NVTE_Fused_Attn_Backend`` (pybind11) enum - value-for-value, and instances of the two enums compare equal when they - share the same integer value. Unlike the pybind enum, a plain-python - ``IntEnum`` is traceable by ``torch.compile``: comparisons against a member - constant-fold cleanly. Lookup by name (``FusedAttnBackend["FP8"]``) works - the same way as with the dict this used to be. - - Members do not survive a graph break, though, so ``get_attention_backend`` - returns the sub-backend as a plain int and ``cast`` turns it back into a - member. - """ + """Legacy import-path mirror of the Python cuDNN attention backend enum.""" - No_Backend = int(NVTE_Fused_Attn_Backend.NVTE_No_Backend) - F16_arbitrary_seqlen = int(NVTE_Fused_Attn_Backend.NVTE_F16_arbitrary_seqlen) - FP8 = int(NVTE_Fused_Attn_Backend.NVTE_FP8) + No_Backend = -1 + F16_arbitrary_seqlen = 1 + FP8 = 2 @classmethod - def cast( - cls, backend: "Union[FusedAttnBackend, NVTE_Fused_Attn_Backend]" - ) -> "FusedAttnBackend": - """Normalize a backend value to the canonical ``FusedAttnBackend`` member. - - The pybind ``transformer_engine_torch.NVTE_Fused_Attn_Backend`` enum is - accepted as input for backward compatibility and mapped to the matching - ``FusedAttnBackend`` member. - """ + def cast(cls, backend): + """Convert an integer or another compatible enum to this enum.""" if isinstance(backend, cls): return backend return cls(int(backend)) - def __eq__(self, other: object) -> bool: - # ``FusedAttnBackend`` is an ``IntEnum`` while ``NVTE_Fused_Attn_Backend`` - # is a pybind11 enum. Compare by integer value so the two enums stay - # equivalent regardless of the pybind11 version (the pybind ``__eq__`` - # handles the reverse order). - if isinstance(other, NVTE_Fused_Attn_Backend): - return int(self) == int(other) - return int.__eq__(self, other) - - def __ne__(self, other: object) -> bool: - result = self.__eq__(other) - if result is NotImplemented: - return result - return not result - - def __hash__(self) -> int: - return int.__hash__(self) - - -# Fail fast at import time if a new enumerator is added on the C++ side -# without being mirrored above. -assert {f"NVTE_{m.name}" for m in FusedAttnBackend} == set(NVTE_Fused_Attn_Backend.__members__), ( - "FusedAttnBackend in python is out of sync with" - " transformer_engine_torch.NVTE_Fused_Attn_Backend defined on the C++ side." - " Please make sure TE C++ and python are in sync." -) - -BACKEND_FP8_THREADS_PER_CTA = 128 -BACKEND_F16arb_ELTS_PER_THREADS = 16 META_QKV = FP8FwdTensorIdx.GEMM1_OUTPUT META_DQKV = FP8BwdTensorIdx.GRAD_OUTPUT1 @@ -172,482 +61,21 @@ def __hash__(self) -> int: META_DP = FP8BwdTensorIdx.GRAD_INPUT3 -def fused_attn_fwd( - is_training: bool, - max_seqlen_q: int, - max_seqlen_kv: int, - cu_seqlens_q: torch.Tensor, - cu_seqlens_kv: torch.Tensor, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - fake_dtype: torch.dtype, - fused_attention_backend: FusedAttnBackend, - attn_bias: torch.Tensor = None, - cu_seqlens_q_padded: torch.Tensor = None, - cu_seqlens_kv_padded: torch.Tensor = None, - page_table_k: torch.Tensor = None, - page_table_v: torch.Tensor = None, - s_quantizer: Quantizer = None, - o_quantizer: Quantizer = None, - attn_scale: float = None, - dropout: float = 0.0, - fast_zero_fill: bool = True, - qkv_layout: str = "sbh3d", - o_format: str = "sbhd", - qkv_scale_inv_format: str = None, - attn_bias_type: str = "no_bias", - attn_mask_type: str = "padding", - softmax_type: str = "vanilla", - window_size: Tuple[int, int] = (-1, -1), - bottom_right_diagonal: bool = None, - rng_gen: torch.Generator = None, - softmax_offset: torch.Tensor = None, - return_max_logit: bool = False, - cuda_graph: bool = False, -) -> Tuple[Union[torch.Tensor, None], ...]: - """Fused Attention FWD for separate QKV input. - - Parameters - ---------- - is_training : bool - if True, runs training and produces auxiliary tensors aux_ctx_tensors - for the backward; if False, runs inference and doesn't produce aux_ctx_tensors - max_seqlen_q : int - max sequence length for Q, used for padding; - may be larger than max(seqlens_q), - seqlens_q = cu_seqlens_q[1:] - cu_seqlens_q[:-1] - max_seqlen_kv : int - max sequence length for K and V, used for padding; - may be larger than max(seqlens_kv), - seqlens_kv = cu_seqlens_kv[1:] - cu_seqlens_kv[:-1] - cu_seqlens_q : torch.Tensor - cumulative sequence lengths for Q; shape [batch_size + 1] - cu_seqlens_kv : torch.Tensor - cumulative sequence lengths for K and V; shape [batch_size + 1] - q : torch.Tensor - input tensor Q; shape sbhd, bshd or thd (see `qkv_layout` for details) - k : torch.Tensor - input tensor K; shape sbhd, bshd or thd (see `qkv_layout` for details) - v : torch.Tensor - input tensor V; shape sbhd, bshd or thd (see `qkv_layout` for details) - fake_dtype : DType - data type of Q, K and V - in case of high precision, fake dtype in case of FP8; - in torch.dtype - fused_attention_backend : FusedAttnBackend - please see FusedAttention module for details on supported backends. - attn_bias : torch.Tensor, default = None - input tensor Bias when attn_bias_type is "pre_scale_bias" or "post_scale_bias"; - shape [1, num_heads, max_seqlen_q, max_seqlen_kv], same data type as q, k and v - cu_seqlens_q_padded : torch.Tensor, default = None - cumulative sequence offsets for Q; shape [batch_size + 1] - cu_seqlens_kv_padded : torch.Tensor, default = None - cumulative sequence offsets for KV; shape [batch_size + 1] - page_table_k : torch.Tensor, default = None - page table for K cache; shape [batch_size, max_pages_per_seq_k] - page_table_v : torch.Tensor, default = None - page table for V cache; shape [batch_size, max_pages_per_seq_v] - s_quantizer : Quantizer, default = None - Quantizer object for the intermediate value S. - o_quantizer : Quantizer, default = None - Quantizer object for the output of the attention. - attn_scale : float, default = None - if not None, use attn_scale as the attention scale for Q*K.T BMM; - if None, use 1.0/sqrt(head_dim_qk) as the default - dropout : float, default = 0.0 - dropout probability, 0.0 means no dropout, 1.0 means no output; - dropout must be 0.0 if is_training is False - fast_zero_fill : bool, default = True - if True, initializes the output tensor O to zero using the fast filling method; - if False, uses PyTorch's .fill_() method - qkv_layout : str, default = "sbh3d" - layout of Q, K and V; - {"sb3hd", "sbh3d", "sbhd_sb2hd", "sbhd_sbh2d", "sbhd_sbhd_sbhd", - "bs3hd", "bsh3d", "bshd_bs2hd", "bshd_bsh2d", "bshd_bshd_bshd", - "t3hd", "th3d", "thd_t2hd", "thd_th2d", "thd_thd_thd"} - o_format : str, default = "sbhd" - format of O; {"sbhd", "bshd", "thd"} - qkv_scale_inv_format : str, default = None - format of the scale-inverse tensors for QKV; {"sbhd", "bshd", "thd", "bhsd"}; - if None, defaults to the format inferred from qkv_layout. - attn_bias_type : str, default = "no_bias" - type of the bias; {"no_bias", "pre_scale_bias", "post_scale_bias", "alibi"} - attn_mask_type : str, default = "padding" - type of the attention mask; {"padding", "causal", "padding_causal", "no_mask"} - softmax_type : str, default = "vanilla" - type of the attention softmax; {"vanilla", "off-by-one", "learnable"} - window_size : Tuple[int, int], default = (-1, -1) - sliding window size for local attention, where query at position i attends to keys - in [i + seqlen_k - seqlen_q - window_size[0], i + seqlen_k - seqlen_q - + window_size[1]] inclusive. Special cases (-1, -1) and (-1, 0) mean no sliding - window and causal mask specifically. - bottom_right_diagonal: bool, default = None - whether to align sliding window and ALiBi diagonal to the top left (False) or - bottom right (True) corner of the softmax matrix. - rng_gen : torch.Generator, default = None - random number generator; - if None, uses the default CUDA generator from PyTorch; otherwise, uses rng_gen - softmax_offset : torch.Tensor, default = None - softmax offset tensor of shape [1, h_q, 1, 1]. - See softmax_type in DotProductAttention for details. - return_max_logit : bool, default = False - whether to return the maximum attention score - cuda_graph : bool, default = False - whether or not cuda graph capture is enabled. - - Returns - ---------- - o : torch.Tensor - output tensor O, of the attention calculation; same data type as Q, K and V; - same shape as Q - aux_ctx_tensors : List[torch.Tensor] - auxiliary output tensors used for the backward; - if is_training is True, aux_ctx_tensors = [softmax-related tensors, rng_state] - if is_training is False, aux_ctx_tensors = None - - softmax-related tensors: - 1. if fused_attention_backend == FusedAttnBackend["F16_arbitrary_seqlen"] - softmaxStats: torch.Tensor - log(sum(e^(x - max(x)))), where x=Q*K.T - shape [batch_size, num_heads, max_seqlen_q, 1], dtype float32 - Max: torch.Tensor, only when return_max_logit is True - shape [batch_size, num_heads, max_seqlen_q, 1], dtype float32 - 2. if fused_attention_backend == FusedAttnBackend["FP8"] - softmaxStats: torch.Tensor - log(sum(e^(x - max(x)))), where x=Q*K.T - shape [batch_size, num_heads, max_seqlen_q, 1], dtype float32 - rng_state: torch.Tensor - state of the random number generator; - [seed, offset], dtype uint64 - max_logit : if return_max_logit = True, shape [h] and same data type as O; otherwise None - """ - - if bottom_right_diagonal is None: - bottom_right_diagonal = attn_mask_type in { - "causal_bottom_right", - "padding_causal_bottom_right", - } - - if attn_scale is None: - d = q.size(-1) - attn_scale = 1.0 / math.sqrt(d) - - if attn_bias_type not in ["no_bias", "alibi"]: - if attn_bias is None: - raise ValueError( - f"attn_bias tensor cannot be None when attn_bias_type={attn_bias_type!r}." - ) - if attn_bias.dtype != q.dtype: - raise ValueError( - "attn_bias tensor must have the same dtype as q and kv: " - f"attn_bias.dtype={attn_bias.dtype} but q.dtype={q.dtype}." - ) - - # Accept the pybind enum for backward compatibility. - fused_attention_backend = FusedAttnBackend.cast(fused_attention_backend) - if fused_attention_backend == FusedAttnBackend["No_Backend"]: - raise ValueError( - "Fused attention does not support this input combination:" - f" qkv_layout={qkv_layout!r}, attn_bias_type={attn_bias_type!r}," - f" attn_mask_type={attn_mask_type!r}, q.shape={list(q.shape)}," - f" q.dtype={q.dtype}, backend={fused_attention_backend}." - ) - - if fused_attention_backend == FusedAttnBackend["F16_arbitrary_seqlen"]: - rng_elts_per_thread = BACKEND_F16arb_ELTS_PER_THREADS - # FP8 fused attention API from fmha_v2 - elif fused_attention_backend == FusedAttnBackend["FP8"]: - rng_elts_per_thread = ( - max_seqlen_q * max_seqlen_q + BACKEND_FP8_THREADS_PER_CTA - 1 - ) // BACKEND_FP8_THREADS_PER_CTA - else: - raise ValueError(f"Unsupported backend {fused_attention_backend}") - - # execute kernel - output_tensors = tex.fused_attn_fwd( - max_seqlen_q, - max_seqlen_kv, - is_training, - attn_scale, - dropout, - fast_zero_fill, - QKVLayout[qkv_layout], - QKVFormat[o_format], - QKVFormat[qkv_scale_inv_format], - AttnBiasType[attn_bias_type], - AttnMaskType[attn_mask_type], - SoftmaxType[softmax_type], - window_size, - bottom_right_diagonal, - cu_seqlens_q, - cu_seqlens_kv, - q, - k, - v, - fake_dtype, - cu_seqlens_q_padded, - cu_seqlens_kv_padded, - page_table_k, - page_table_v, - s_quantizer, - o_quantizer, - attn_bias, - softmax_offset, - rng_gen, - rng_elts_per_thread, - return_max_logit, - cuda_graph, +# Resolve lazily because cpp_extensions is imported while transformer_engine.pytorch +# itself is still initializing. +def fused_attn_fwd(*args, **kwargs): + """Run fused attention forward through the cuDNN frontend Python API.""" + from transformer_engine.pytorch.attention.dot_product_attention.cudnn_attention import ( + fused_attn_fwd as implementation, ) - if return_max_logit: - qkv_format = qkv_layout.replace("3", "").replace("2", "").split("_")[0] - # thd (newer cuDNN runtimes, non-sm120): output_tensors: out [tq, h, d], Stats [tq, h, 1], Max [tq, h, 1] - # thd (older cuDNN runtimes or sm120): output_tensors: out [tq, h, d], Stats [b, h, sq, 1], Max [b, h, sq, 1] - # bshd: output_tensors: out [b, sq, h, d], Stats [b, h, sq, 1], Max [b, h, sq, 1] - # sbhd: output_tensors: out [sq, b, h, d], Stats [b, h, sq, 1], Max [b, h, sq, 1] - aux_ctx_tensors = [output_tensors[1]] + list( - output_tensors[3:] - ) # Stats + rng_state + optional tensors - max_tensor = output_tensors[2] - amax_dims = (0, 2) if max_tensor.ndim == 3 else (0, 2, 3) - - if qkv_format == "thd": - if max_tensor.ndim == 4: - # For THD on cuDNN <= 9.6 or THD on sm120, Max tensor can be [b, h, sq, 1] - # with padded sequence positions. Exclude those padded positions when computing max_logit. - seqlens_q = (cu_seqlens_q[1:] - cu_seqlens_q[:-1]).to(device=max_tensor.device) - sq_idx = torch.arange(max_tensor.shape[2], device=max_tensor.device).view( - 1, 1, -1, 1 - ) - valid = sq_idx < seqlens_q.view(-1, 1, 1, 1) - max_tensor = max_tensor.masked_fill(~valid, float("-inf")) - elif max_tensor.ndim == 3: - if cu_seqlens_q_padded is not None: - # Exclude padding; CP may pass nonzero padded offsets. - actual_seqlens = (cu_seqlens_q[1:] - cu_seqlens_q[:-1]).to( - device=max_tensor.device - ) - tq = max_tensor.shape[0] - starts = cu_seqlens_q_padded[:-1].to(device=max_tensor.device) - ends = (starts + actual_seqlens).clamp(max=tq) - delta = torch.zeros(tq + 1, dtype=torch.int32, device=max_tensor.device) - updates = torch.ones_like(starts, dtype=torch.int32) - delta.scatter_add_(0, starts.clamp(max=tq), updates) - delta.scatter_add_(0, ends, -updates) - valid = delta[:-1].cumsum(0) > 0 - max_tensor = max_tensor.masked_fill(~valid.view(-1, 1, 1), float("-inf")) - - # Max -> max_logit [h] - max_logit = torch.amax(max_tensor, dim=amax_dims).to(dtype=output_tensors[0].dtype) - return output_tensors[0], aux_ctx_tensors, max_logit - - # out, aux_ctx_tensors - return output_tensors[0], output_tensors[1:] - - -def fused_attn_bwd( - max_seqlen_q: int, - max_seqlen_kv: int, - cu_seqlens_q: torch.Tensor, - cu_seqlens_kv: torch.Tensor, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - o: torch.Tensor, - d_o: torch.Tensor, - fake_dtype: torch.dtype, - aux_ctx_tensors: List[torch.Tensor], - fused_attention_backend: FusedAttnBackend, - cu_seqlens_q_padded: torch.Tensor = None, - cu_seqlens_kv_padded: torch.Tensor = None, - s_quantizer: Quantizer = None, - dp_quantizer: Quantizer = None, - dqkv_quantizer: Quantizer = None, - attn_scale: Optional[float] = None, - dropout: float = 0.0, - fast_zero_fill: bool = True, - qkv_layout: str = "sbh3d", - o_format: str = "sbhd", - do_format: str = "sbhd", - dqkv_layout: str = "sbh3d", - qkv_scale_inv_format: str = None, - do_scale_inv_format: str = None, - attn_bias_type: str = "no_bias", - attn_mask_type: str = "padding", - softmax_type: str = "vanilla", - window_size: Tuple[int, int] = (-1, -1), - bottom_right_diagonal: bool = None, - deterministic: bool = False, - cuda_graph: bool = False, -) -> Tuple[Union[torch.Tensor, None], ...]: - """Fused Attention BWD for packed KV input. - - Parameters - ---------- - max_seqlen_q : int - max sequence length for Q, used for padding; may be larger than max(seqlens_q), - seqlens_q = cu_seqlens_q[1:] - cu_seqlens_q[:-1] - max_seqlen_kv : int - max sequence length for K and V, used for padding; - may be larger than max(seqlens_kv), - seqlens_kv = cu_seqlens_kv[1:] - cu_seqlens_kv[:-1] - cu_seqlens_q : torch.Tensor - cumulative sequence lengths for Q; shape [batch_size + 1] - cu_seqlens_kv : torch.Tensor - cumulative sequence lengths for K and V; shape [batch_size + 1] - q : torch.Tensor - input tensor Q; shape sbhd, bshd or thd (see `qkv_layout` for details) - k : torch.Tensor - input tensor K; shape sbhd, bshd or thd (see `qkv_layout` for details) - v : torch.Tensor - input tensor V; shape sbhd, bshd or thd (see `qkv_layout` for details) - o : torch.Tensor - input tensor O (output of forward); same data type as Q, K and V; - same shape as Q - d_o : torch.Tensor - input tensor dO (gradient of O); same data type as Q, K and V; - same shape as Q - fake_dtype : DType - data type of Q, K and V - in case of high precision, fake dtype in case of FP8; - in torch.dtype - aux_ctx_tensors : List[torch.Tensor] - auxiliary output tensors of the forward pass when its is_training is True, - e.g. aux_ctx_tensors = [S, Max, rng_state] - fused_attention_backend : FusedAttnBackend - please see FusedAttention module for details on supported backends. - cu_seqlens_q_padded : torch.Tensor, default = None - cumulative sequence offsets for Q; shape [batch_size + 1] - cu_seqlens_kv_padded : torch.Tensor, default = None - cumulative sequence offsets for KV; shape [batch_size + 1] - s_quantizer : Quantizer, default = None - Quantizer object for the intermediate value S. - dp_quantizer : Quantizer, default = None - Quantizer object for the intermediate value dP. - dqkv_quantizer : Quantizer, default = None - Quantizer object for the output values of the fused_attn_bwd. - attn_scale : float, default = None - if not None, use attn_scale as the attention scale for Q*K.T BMM; - if None, use 1.0/sqrt(head_dim_qk) as the default - dropout : float, default = 0.0 - dropout probability, 0.0 means no dropout, 1.0 means no output; - dropout must be 0.0 if is_training is False - fast_zero_fill : bool, default = True - if True, initializes the output tensor O to zero using the fast filling method; - if False, uses PyTorch's .fill_() method - qkv_layout : str, default = "sbh3d" - layout of Q, K and V; - {"sb3hd", "sbh3d", "sbhd_sb2hd", "sbhd_sbh2d", "sbhd_sbhd_sbhd", - "bs3hd", "bsh3d", "bshd_bs2hd", "bshd_bsh2d", "bshd_bshd_bshd", - "t3hd", "th3d", "thd_t2hd", "thd_th2d", "thd_thd_thd"} - o_format : str, default = "sbhd" - format of O; {"sbhd", "bshd", "thd"} - do_format : str, default = "sbhd" - format of dO; {"sbhd", "bshd", "thd"} - dqkv_layout : str, default = "sbh3d" - layout of dQ, dK and dV; - {"sb3hd", "sbh3d", "sbhd_sb2hd", "sbhd_sbh2d", "sbhd_sbhd_sbhd", - "bs3hd", "bsh3d", "bshd_bs2hd", "bshd_bsh2d", "bshd_bshd_bshd", - "t3hd", "th3d", "thd_t2hd", "thd_th2d", "thd_thd_thd"} - qkv_scale_inv_format : str, default = None - format of the scale-inverse tensors for QKV; {"sbhd", "bshd", "thd", "bhsd"}; - if None, defaults to the format inferred from qkv_layout. - do_scale_inv_format : str, default = None - format of the scale-inverse tensors for dO; {"sbhd", "bshd", "thd", "bhsd"}; - if None, defaults to the format inferred from the output layout. - attn_bias_type : str, default = "no_bias" - type of the bias; {"no_bias", "pre_scale_bias", "post_scale_bias", "alibi"} - attn_mask_type : str, default = "padding" - type of the attention mask; {"padding", "causal", "padding_causal", "no_mask"} - softmax_type : str, default = "vanilla" - type of the attention softmax; {"vanilla", "off-by-one", "learnable"} - window_size : Tuple[int, int], default = (-1, -1) - sliding window size for local attention, where query at position i attends to keys - in [i + seqlen_k - seqlen_q - window_size[0], i + seqlen_k - seqlen_q - + window_size[1]] inclusive. Special cases (-1, -1) and (-1, 0) mean no sliding - window and causal mask specifically. - bottom_right_diagonal: bool, default = None - whether to align sliding window and ALiBi diagonal to the top left (False) or - bottom right (True) corner of the softmax matrix. - deterministic : bool, default = False - whether to execute the backward pass with deterministic behaviours. - cuda_graph : bool, default = False - whether or not cuda graph capture is enabled. - - Returns - ---------- - d_q : torch.Tensor - gradient tensor of Q; same data type and shape as Q - d_k : torch.Tensor - gradient tensor of K; same data type and shape as K - d_v : torch.Tensor - gradient tensor of V; same data type and shape as V - d_bias : torch.Tensor, optional - gradient tensor of Bias when attn_bias_type is "pre_scale_bias" - or "post_scale_bias"; same data type and shape as Bias - d_softmax_offset : torch.Tensor, optional - gradient tensor of softmax offset of shape [1, h_q, 1, 1]. - See softmax_type in DotProductAttention for details. - """ - if bottom_right_diagonal is None: - bottom_right_diagonal = attn_mask_type in { - "causal_bottom_right", - "padding_causal_bottom_right", - } - - if attn_scale is None: - d = q.size(-1) - attn_scale = 1.0 / math.sqrt(d) - - # Accept the pybind enum for backward compatibility. - fused_attention_backend = FusedAttnBackend.cast(fused_attention_backend) - if fused_attention_backend == FusedAttnBackend["No_Backend"]: - raise ValueError( - "Fused attention backward does not support this input combination:" - f" qkv_layout={qkv_layout!r}, attn_bias_type={attn_bias_type!r}," - f" attn_mask_type={attn_mask_type!r}, q.shape={list(q.shape)}," - f" q.dtype={q.dtype}, backend={fused_attention_backend}." - ) + return implementation(*args, **kwargs) - if len(aux_ctx_tensors) < 1: - raise ValueError( - "aux_ctx_tensors must contain rng_state as its last element," - f" but got len(aux_ctx_tensors)={len(aux_ctx_tensors)}" - f" for backend={fused_attention_backend}." - ) - output_tensors = tex.fused_attn_bwd( - max_seqlen_q, - max_seqlen_kv, - attn_scale, - dropout, - fast_zero_fill, - QKVLayout[qkv_layout], - QKVFormat[o_format], - QKVFormat[do_format], - QKVLayout[dqkv_layout], - QKVFormat[qkv_scale_inv_format], - QKVFormat[do_scale_inv_format], - AttnBiasType[attn_bias_type], - AttnMaskType[attn_mask_type], - SoftmaxType[softmax_type], - window_size, - bottom_right_diagonal, - deterministic, - cu_seqlens_q, - cu_seqlens_kv, - q, - k, - v, - o, - d_o, - fake_dtype, - aux_ctx_tensors, - cu_seqlens_q_padded, - cu_seqlens_kv_padded, - s_quantizer, - dp_quantizer, - dqkv_quantizer, - cuda_graph, +def fused_attn_bwd(*args, **kwargs): + """Run fused attention backward through the cuDNN frontend Python API.""" + from transformer_engine.pytorch.attention.dot_product_attention.cudnn_attention import ( + fused_attn_bwd as implementation, ) - return output_tensors + return implementation(*args, **kwargs) diff --git a/transformer_engine/pytorch/csrc/extensions.h b/transformer_engine/pytorch/csrc/extensions.h index f1ef80ca2ee..6e368b6b692 100644 --- a/transformer_engine/pytorch/csrc/extensions.h +++ b/transformer_engine/pytorch/csrc/extensions.h @@ -94,36 +94,8 @@ std::tuple moe_unpermute_bwd(at::Tensor input_bwd, at::T * Attention **************************************************************************************************/ -std::tuple get_fused_attn_backend( - const py::object &fused_attn_params); - -std::vector fused_attn_fwd( - size_t max_seqlen_q, size_t max_seqlen_kv, bool is_training, float attn_scale, float p_dropout, - bool set_zero, NVTE_QKV_Layout qkv_layout, NVTE_QKV_Format o_format, - NVTE_QKV_Format qkv_scale_inv_format, NVTE_Bias_Type bias_type, NVTE_Mask_Type attn_mask_type, - NVTE_Softmax_Type softmax_type, const std::vector window_size, - bool bottom_right_diagonal, const at::Tensor cu_seqlens_q, const at::Tensor cu_seqlens_kv, - const py::handle Q, const py::handle K, const py::handle V, const at::ScalarType fake_dtype, - const std::optional cu_seqlens_q_padded, - const std::optional cu_seqlens_kv_padded, - const std::optional page_table_k, const std::optional page_table_v, - py::handle s_quantizer, py::handle o_quantizer, const std::optional Bias, - const std::optional SoftmaxOffset, const std::optional rng_gen, - size_t rng_elts_per_thread, bool return_max_logit, bool cuda_graph); - -std::vector fused_attn_bwd( - size_t max_seqlen_q, size_t max_seqlen_kv, float attn_scale, float p_dropout, bool set_zero, - NVTE_QKV_Layout qkv_layout, NVTE_QKV_Format o_format, NVTE_QKV_Format do_format, - NVTE_QKV_Layout dqkv_layout, NVTE_QKV_Format qkv_scale_inv_format, - NVTE_QKV_Format do_scale_inv_format, NVTE_Bias_Type bias_type, NVTE_Mask_Type attn_mask_type, - NVTE_Softmax_Type softmax_type, const std::vector window_size, - bool bottom_right_diagonal, bool deterministic, const at::Tensor cu_seqlens_q, - const at::Tensor cu_seqlens_kv, const py::handle Q, const py::handle K, const py::handle V, - const py::handle O, const py::handle dO, const at::ScalarType fake_dtype, - const std::vector Aux_CTX_Tensors, - const std::optional cu_seqlens_q_padded, - const std::optional cu_seqlens_kv_padded, py::handle s_quantizer, - py::handle dp_quantizer, py::handle dqkv_quantizer, bool cuda_graph); +at::Tensor get_cudnn_attention_rng_state(const std::optional rng_gen, + size_t increment); at::Tensor fa_prepare_fwd(at::Tensor qkvi); at::Tensor fa_prepare_bwd(at::Tensor q, at::Tensor k, at::Tensor v); diff --git a/transformer_engine/pytorch/csrc/extensions/attention.cpp b/transformer_engine/pytorch/csrc/extensions/attention.cpp index 9cda33474fa..0bd33a4e5a0 100644 --- a/transformer_engine/pytorch/csrc/extensions/attention.cpp +++ b/transformer_engine/pytorch/csrc/extensions/attention.cpp @@ -10,642 +10,6 @@ namespace transformer_engine::pytorch { -// get the fused attention backend -std::tuple get_fused_attn_backend(const py::object &p) { - FusedAttnConfigWrapper cfg; - cfg.set_is_training(p.attr("is_training").cast()) - .set_deterministic(p.attr("deterministic").cast()) - .set_cuda_graph(p.attr("cuda_graph").cast()) - .set_return_max_logit(p.attr("return_max_logit").cast()) - .set_attn_mask_type(p.attr("attn_mask_type").cast()) - .set_bias_type(p.attr("bias_type").cast()) - .set_window_size_left(p.attr("window_size_left").cast()) - .set_window_size_right(p.attr("window_size_right").cast()) - .set_bottom_right_diagonal(p.attr("bottom_right_diagonal").cast()) - .set_softmax_type(p.attr("softmax_type").cast()) - .set_scaling_mode(p.attr("scaling_mode").cast()) - .set_dropout(p.attr("dropout").cast()) - .set_attn_scale(p.attr("attn_scale").cast()) - .set_qkv_dtype(static_cast(p.attr("qkv_dtype").cast())) - .set_o_dtype(static_cast(p.attr("o_dtype").cast())) - .set_do_dtype(static_cast(p.attr("do_dtype").cast())) - .set_dqkv_dtype(static_cast(p.attr("dqkv_dtype").cast())) - .set_qkv_layout(p.attr("qkv_layout").cast()) - .set_o_format(p.attr("o_format").cast()) - .set_do_format(p.attr("do_format").cast()) - .set_dqkv_layout(p.attr("dqkv_layout").cast()) - .set_qkv_scale_inv_format(p.attr("qkv_scale_inv_format").cast()) - .set_do_scale_inv_format(p.attr("do_scale_inv_format").cast()) - .set_batch_size(p.attr("batch_size").cast()) - .set_num_attn_heads(p.attr("num_attn_heads").cast()) - .set_num_gqa_groups(p.attr("num_gqa_groups").cast()) - .set_head_dim_qk(p.attr("head_dim_qk").cast()) - .set_head_dim_v(p.attr("head_dim_v").cast()) - .set_max_seqlen_q(p.attr("max_seqlen_q").cast()) - .set_max_seqlen_kv(p.attr("max_seqlen_kv").cast()) - .set_num_tokens_q(p.attr("num_tokens_q").cast()) - .set_num_tokens_kv(p.attr("num_tokens_kv").cast()) - .set_num_pages_k(p.attr("num_pages_k").cast()) - .set_num_pages_v(p.attr("num_pages_v").cast()) - .set_page_size_k(p.attr("page_size_k").cast()) - .set_page_size_v(p.attr("page_size_v").cast()) - .set_max_pages_per_seq_k(p.attr("max_pages_per_seq_k").cast()) - .set_max_pages_per_seq_v(p.attr("max_pages_per_seq_v").cast()) - .set_bias_batch_size(p.attr("bias_batch_size").cast()) - .set_bias_num_heads(p.attr("bias_num_heads").cast()) - .set_bias_seqlen_q(p.attr("bias_seqlen_q").cast()) - .set_bias_seqlen_kv(p.attr("bias_seqlen_kv").cast()); - - py::gil_scoped_release nogil; - const char *message = nullptr; - NVTE_Fused_Attn_Backend fused_attention_backend = nvte_get_fused_attn_backend_v2(cfg, &message); - return {fused_attention_backend, std::string(message)}; -} - -// helper function for S and dP quantizers -std::tuple> quantizer_helper( - py::handle quantizer, const std::vector &shape, DType dtype, bool create_hp_tensor, - std::optional data) { - std::unique_ptr T_quantizer = convert_quantizer(quantizer); - TensorWrapper te_T; - py::object py_T; - std::optional amax_buf; - if (quantizer.is_none()) { - // high precision - auto *none_quantizer = dynamic_cast(T_quantizer.get()); - if (data.has_value()) { - std::tie(te_T, py_T) = none_quantizer->create_tensor(shape, dtype, data.value()); - } else { - std::tie(te_T, py_T) = none_quantizer->create_tensor(shape, dtype); - } - } else if (detail::IsFloat8Quantizers(quantizer.ptr())) { - // delayed scaling; this helps initialize scale_inv - auto *T_quantizer_fp8 = dynamic_cast(T_quantizer.get()); - std::tie(te_T, py_T) = - T_quantizer_fp8->create_tensor(shape, dtype, data, std::nullopt, std::nullopt); - } else if (detail::IsFloat8CurrentScalingQuantizers(quantizer.ptr())) { - // current scaling - auto *T_quantizer_fp8 = dynamic_cast(T_quantizer.get()); - if (create_hp_tensor) { - if (data.has_value()) { - std::tie(te_T, py_T, amax_buf) = - T_quantizer_fp8->create_unquantized_tensor_with_amax(shape, dtype, data.value()); - } else { - std::tie(te_T, py_T, amax_buf) = - T_quantizer_fp8->create_unquantized_tensor_with_amax(shape, dtype); - } - } else { - std::tie(te_T, py_T) = T_quantizer_fp8->create_tensor(shape, dtype); - NVTE_CHECK( - !data.has_value(), - "Float8CurrentScalingQuantizer::create_tensor() does not take data tensor as input!"); - } - } else if (detail::IsMXFP8Quantizers(quantizer.ptr())) { - // MXFP8 - if (create_hp_tensor) { - if (data.has_value()) { - std::tie(te_T, py_T) = NoneQuantizer(py::none()).create_tensor(shape, dtype, data.value()); - } else { - std::tie(te_T, py_T) = NoneQuantizer(py::none()).create_tensor(shape, dtype); - } - } else { - auto *T_quantizer_fp8 = dynamic_cast(T_quantizer.get()); - std::tie(te_T, py_T) = T_quantizer_fp8->create_tensor(shape, dtype); - NVTE_CHECK(!data.has_value(), - "MXFP8Quantizer::create_tensor() does not take data tensor as input!"); - } - } - return {std::move(te_T), std::move(py_T), std::move(amax_buf)}; -} - -// fused attention FWD with separate Q, K and V tensors -std::vector fused_attn_fwd( - size_t max_seqlen_q, size_t max_seqlen_kv, bool is_training, float attn_scale, float p_dropout, - bool set_zero, NVTE_QKV_Layout qkv_layout, NVTE_QKV_Format o_format, - NVTE_QKV_Format qkv_scale_inv_format, NVTE_Bias_Type bias_type, NVTE_Mask_Type attn_mask_type, - NVTE_Softmax_Type softmax_type, const std::vector window_size, - bool bottom_right_diagonal, const at::Tensor cu_seqlens_q, const at::Tensor cu_seqlens_kv, - const py::handle Q, const py::handle K, const py::handle V, const at::ScalarType fake_dtype, - const std::optional cu_seqlens_q_padded, - const std::optional cu_seqlens_kv_padded, - const std::optional page_table_k, const std::optional page_table_v, - py::handle s_quantizer, py::handle o_quantizer, const std::optional Bias, - const std::optional SoftmaxOffset, const std::optional rng_gen, - size_t rng_elts_per_thread, bool return_max_logit, bool cuda_graph) { - // Ensure that cuDNN handle is created on the correct device, - // overriding torch.cuda.set_device calls from user side. - // Assumes all tensors passed are on the same device. - at::cuda::CUDAGuard device_guard(cu_seqlens_q.device()); - - auto none = py::none(); - - // create QKV tensor wrappers - TensorWrapper te_Q, te_K, te_V; - te_Q = makeTransformerEngineTensor(Q, none); - te_K = makeTransformerEngineTensor(K, none); - te_V = makeTransformerEngineTensor(V, none); - const DType qkv_type = te_Q.dtype(); - - // create S tensor - auto [te_S, py_S, _] = quantizer_helper(s_quantizer, {0}, DType::kFloat32, false, std::nullopt); - - // create O tensor - std::unique_ptr O_quantizer = convert_quantizer(o_quantizer); - std::vector q_shape = convertShape(te_Q.shape()); - std::vector v_shape = convertShape(te_V.shape()); - auto o_shape_tmp = std::vector{q_shape.begin(), q_shape.end()}; - o_shape_tmp[o_shape_tmp.size() - 1] = v_shape[v_shape.size() - 1]; - auto o_shape = std::vector{o_shape_tmp.begin(), o_shape_tmp.end()}; - NVTE_QKV_Format q_format = nvte_get_q_format(qkv_layout); - AttentionShape o_parsed(q_format, o_shape_tmp.data()); - o_parsed.to_format(o_format, o_shape.data()); - const DType fake_dtype_te = GetTransformerEngineDType(fake_dtype); - auto [te_O, py_O, o_amax_buf] = - quantizer_helper(o_quantizer, o_shape, fake_dtype_te, true, std::nullopt); - - // construct NVTE tensors - TensorWrapper te_Bias; - TensorWrapper te_cu_seqlens_q, te_cu_seqlens_kv; - TensorWrapper te_cu_seqlens_q_padded, te_cu_seqlens_kv_padded; - TensorWrapper te_page_table_k, te_page_table_v; - if (qkv_type == DType::kFloat8E4M3 || qkv_type == DType::kFloat8E5M2) { - // FP8 - if (set_zero && (o_format == NVTE_QKV_Format::NVTE_THD)) { - // Initialize both the output data and its amax metadata. - te_O.zero_(at::cuda::getCurrentCUDAStream()); - } - } else if (qkv_type == DType::kBFloat16 || qkv_type == DType::kFloat16) { - if (o_format == NVTE_QKV_Format::NVTE_THD) { - te_O.zero_(at::cuda::getCurrentCUDAStream()); - } - } else { - NVTE_ERROR("Fused attention only supports FP8 and BF16/FP16 data types. \n"); - } - if ((bias_type != NVTE_NO_BIAS) && (bias_type != NVTE_ALIBI) && (Bias.has_value())) { - auto bias_sizes = Bias.value().sizes().vec(); - std::vector bias_shape{bias_sizes.begin(), bias_sizes.end()}; - te_Bias = makeTransformerEngineTensor(Bias.value().data_ptr(), bias_shape, DType::kFloat32); - } - auto cu_seqlens_q_sizes = cu_seqlens_q.sizes().vec(); - std::vector cu_seqlens_q_shape{cu_seqlens_q_sizes.begin(), cu_seqlens_q_sizes.end()}; - auto cu_seqlens_kv_sizes = cu_seqlens_kv.sizes().vec(); - std::vector cu_seqlens_kv_shape{cu_seqlens_kv_sizes.begin(), cu_seqlens_kv_sizes.end()}; - te_cu_seqlens_q = - makeTransformerEngineTensor(cu_seqlens_q.data_ptr(), cu_seqlens_q_shape, DType::kInt32); - te_cu_seqlens_kv = - makeTransformerEngineTensor(cu_seqlens_kv.data_ptr(), cu_seqlens_kv_shape, DType::kInt32); - - if ((cu_seqlens_q_padded.has_value()) && (cu_seqlens_kv_padded.has_value())) { - auto cu_seqlens_q_padded_sizes = cu_seqlens_q_padded.value().sizes().vec(); - std::vector cu_seqlens_q_padded_shape{cu_seqlens_q_padded_sizes.begin(), - cu_seqlens_q_padded_sizes.end()}; - auto cu_seqlens_kv_padded_sizes = cu_seqlens_kv_padded.value().sizes().vec(); - std::vector cu_seqlens_kv_padded_shape{cu_seqlens_kv_padded_sizes.begin(), - cu_seqlens_kv_padded_sizes.end()}; - te_cu_seqlens_q_padded = makeTransformerEngineTensor(cu_seqlens_q_padded.value().data_ptr(), - cu_seqlens_q_padded_shape, DType::kInt32); - te_cu_seqlens_kv_padded = makeTransformerEngineTensor( - cu_seqlens_kv_padded.value().data_ptr(), cu_seqlens_kv_padded_shape, DType::kInt32); - } - - if ((page_table_k.has_value()) && (page_table_v.has_value())) { - auto page_table_k_sizes = page_table_k.value().sizes().vec(); - std::vector page_table_k_shape{page_table_k_sizes.begin(), page_table_k_sizes.end()}; - auto page_table_v_sizes = page_table_v.value().sizes().vec(); - std::vector page_table_v_shape{page_table_v_sizes.begin(), page_table_v_sizes.end()}; - te_page_table_k = - makeTransformerEngineTensor(page_table_k.value().data_ptr(), page_table_k_shape, - DType::kInt32, nullptr, nullptr, nullptr); - te_page_table_v = - makeTransformerEngineTensor(page_table_v.value().data_ptr(), page_table_v_shape, - DType::kInt32, nullptr, nullptr, nullptr); - } - - // softmax offset - TensorWrapper te_SoftmaxOffset; - if ((softmax_type != NVTE_VANILLA_SOFTMAX) && (SoftmaxOffset.has_value())) { - auto SoftmaxOffset_sizes = SoftmaxOffset.value().sizes().vec(); - std::vector SoftmaxOffset_shape{SoftmaxOffset_sizes.begin(), SoftmaxOffset_sizes.end()}; - te_SoftmaxOffset = - makeTransformerEngineTensor(SoftmaxOffset.value().data_ptr(), SoftmaxOffset_shape, - DType::kFloat32, nullptr, nullptr, nullptr); - } - - // extract rng seed and offset - auto gen = at::get_generator_or_default( - rng_gen, at::cuda::detail::getDefaultCUDAGenerator()); - at::PhiloxCudaState philox_args = init_philox_state(gen, rng_elts_per_thread); - auto options = torch::TensorOptions().dtype(torch::kInt64).device(torch::kCUDA); - auto rng_state = torch::empty({2}, options); - philox_unpack(philox_args, static_cast(rng_state.data_ptr())); - auto te_rng_state = makeTransformerEngineTensor(rng_state); - - // create auxiliary output tensors - NVTETensorPack nvte_aux_tensor_pack; - nvte_tensor_pack_create(&nvte_aux_tensor_pack); - - // create workspace - TensorWrapper workspace; - - // build the parameter object - FusedAttnFwdParamsWrapper params; - params.set_Q(te_Q.data()) - .set_K(te_K.data()) - .set_V(te_V.data()) - .set_Bias(te_Bias.data()) - .set_SoftmaxOffset(te_SoftmaxOffset.data()) - .set_S(te_S.data()) - .set_O(te_O.data()) - .set_Aux_CTX_Tensors(&nvte_aux_tensor_pack) - .set_cu_seqlens_q(te_cu_seqlens_q.data()) - .set_cu_seqlens_kv(te_cu_seqlens_kv.data()) - .set_cu_seqlens_q_padded(te_cu_seqlens_q_padded.data()) - .set_cu_seqlens_kv_padded(te_cu_seqlens_kv_padded.data()) - .set_page_table_k(te_page_table_k.data()) - .set_page_table_v(te_page_table_v.data()) - .set_rng_state(te_rng_state.data()) - .set_is_training(is_training) - .set_cuda_graph(cuda_graph) - .set_return_max_logit(return_max_logit) - .set_attn_mask_type(attn_mask_type) - .set_bias_type(bias_type) - .set_window_size_left(window_size[0]) - .set_window_size_right(window_size[1]) - .set_bottom_right_diagonal(bottom_right_diagonal) - .set_softmax_type(softmax_type) - .set_dropout(p_dropout) - .set_attn_scale(attn_scale) - .set_qkv_layout(qkv_layout) - .set_o_format(o_format) - .set_qkv_scale_inv_format(qkv_scale_inv_format) - .set_max_seqlen_q(max_seqlen_q) - .set_max_seqlen_kv(max_seqlen_kv) - .set_workspace(workspace.data()) - .set_stream(at::cuda::getCurrentCUDAStream()); - - // populate tensors with appropriate shapes and dtypes - NVTE_SCOPED_GIL_RELEASE({ nvte_fused_attn_fwd_v2(params); }); - - // allocate memory for workspace and auxiliary output tensors - auto workspace_data = allocateSpace(workspace.shape(), workspace.dtype()); - workspace = - makeTransformerEngineTensor(workspace_data.data_ptr(), workspace.shape(), workspace.dtype()); - params.set_workspace(workspace.data()); - - // output_tensors = [O, nvte_aux_tensor_pack.tensors] - std::vector output_tensors; - output_tensors.push_back(py_O); - auto set_tensor_param = [&](size_t i, const at::Tensor &output_tensor) { - output_tensors.push_back(py::cast(output_tensor)); - NVTEBasicTensor temp_data = {output_tensor.data_ptr(), - nvte_tensor_type(nvte_aux_tensor_pack.tensors[i]), - nvte_tensor_shape(nvte_aux_tensor_pack.tensors[i])}; - nvte_set_tensor_param(&nvte_aux_tensor_pack.tensors[i], kNVTERowwiseData, &temp_data); - }; - // allocate memory for nvte_aux_tensor_pack.tensors - // f16_arbitrary: S [b, h, sq, 1]/[tq, h, 1], (optional) Max [b, h, sq, 1]/[tq, h, 1], rng_state [2], (optional) Bias [1, h, sq, skv], (optional) SoftmaxOffset [1, h, 1, 1] - // fp8 : S [b, h, sq, 1], rng_state [2] - size_t i = 0; - at::Tensor output_tensor; - // intermediate softmax stats tensor S - output_tensor = - allocateSpace(nvte_shape_to_vector(nvte_tensor_shape(nvte_aux_tensor_pack.tensors[i])), - static_cast(nvte_tensor_type(nvte_aux_tensor_pack.tensors[i])), false); - set_tensor_param(i++, output_tensor); - // return_max_logit=true allocates Max after S - if (return_max_logit) { - output_tensor = - allocateSpace(nvte_shape_to_vector(nvte_tensor_shape(nvte_aux_tensor_pack.tensors[i])), - static_cast(nvte_tensor_type(nvte_aux_tensor_pack.tensors[i])), false); - set_tensor_param(i++, output_tensor); - } - // rng_state - if (i < nvte_aux_tensor_pack.size) { - set_tensor_param(i++, rng_state); - } - // bias (optional) - if ((bias_type != NVTE_NO_BIAS) && (bias_type != NVTE_ALIBI) && (Bias.has_value())) { - set_tensor_param(i++, Bias.value()); - } - // softmax_offset (optional) - if ((softmax_type != NVTE_VANILLA_SOFTMAX) && (SoftmaxOffset.has_value())) { - set_tensor_param(i++, SoftmaxOffset.value()); - } - - // execute the kernel - NVTE_SCOPED_GIL_RELEASE({ nvte_fused_attn_fwd_v2(params); }); - - // destroy tensor wrappers, but not allocated memory - nvte_tensor_pack_destroy(&nvte_aux_tensor_pack); - - // if training, [O, softmax-related tensors, rng_state]; if inference, [O] - return output_tensors; -} - -// fused attention BWD with separate Q, K and V -std::vector fused_attn_bwd( - size_t max_seqlen_q, size_t max_seqlen_kv, float attn_scale, float p_dropout, bool set_zero, - NVTE_QKV_Layout qkv_layout, NVTE_QKV_Format o_format, NVTE_QKV_Format do_format, - NVTE_QKV_Layout dqkv_layout, NVTE_QKV_Format qkv_scale_inv_format, - NVTE_QKV_Format do_scale_inv_format, NVTE_Bias_Type bias_type, NVTE_Mask_Type attn_mask_type, - NVTE_Softmax_Type softmax_type, const std::vector window_size, - bool bottom_right_diagonal, bool deterministic, const at::Tensor cu_seqlens_q, - const at::Tensor cu_seqlens_kv, const py::handle Q, const py::handle K, const py::handle V, - const py::handle O, const py::handle dO, const at::ScalarType fake_dtype, - const std::vector Aux_CTX_Tensors, - const std::optional cu_seqlens_q_padded, - const std::optional cu_seqlens_kv_padded, py::handle s_quantizer, - py::handle dp_quantizer, py::handle dqkv_quantizer, bool cuda_graph) { - auto none = py::none(); - - // create QKV, O, dO tensor wrappers - TensorWrapper te_Q, te_K, te_V, te_O, te_dO; - te_Q = makeTransformerEngineTensor(Q, none); - te_K = makeTransformerEngineTensor(K, none); - te_V = makeTransformerEngineTensor(V, none); - te_O = makeTransformerEngineTensor(O, none); - te_dO = makeTransformerEngineTensor(dO, none); - - // create S and dP tensors - auto [te_S, py_S, _s] = quantizer_helper(s_quantizer, {0}, DType::kFloat32, false, std::nullopt); - auto [te_dP, py_dP, _dp] = - quantizer_helper(dp_quantizer, {0}, DType::kFloat32, false, std::nullopt); - - // create dQ, dK, dV tensors - TensorWrapper te_dQ, te_dK, te_dV; - py::object py_dQ, py_dK, py_dV; - std::optional dq_amax_buf, dk_amax_buf, dv_amax_buf; - std::unique_ptr dQKV_quantizer = convert_quantizer(dqkv_quantizer); - std::vector q_shape = convertShape(te_Q.shape()); - std::vector k_shape = convertShape(te_K.shape()); - std::vector v_shape = convertShape(te_V.shape()); - const DType dqkv_fake_dtype = GetTransformerEngineDType(fake_dtype); - size_t ndim_q = q_shape.size(); - size_t ndim_kv = k_shape.size(); - std::vector dQ_shape(ndim_q), dK_shape(ndim_kv), dV_shape(ndim_kv); - NVTE_QKV_Format q_format = nvte_get_q_format(qkv_layout); - NVTE_QKV_Format kv_format = nvte_get_kv_format(qkv_layout); - NVTE_QKV_Format dq_format = nvte_get_q_format(dqkv_layout); - NVTE_QKV_Format dkv_format = nvte_get_kv_format(dqkv_layout); - AttentionShape q_parsed(q_format, q_shape.data()); - size_t h_q = q_parsed.h(); - q_parsed.to_format(dq_format, dQ_shape.data()); - AttentionShape k_parsed(kv_format, k_shape.data()); - k_parsed.to_format(dkv_format, dK_shape.data()); - AttentionShape v_parsed(kv_format, v_shape.data()); - v_parsed.to_format(dkv_format, dV_shape.data()); - at::Tensor dQ, dK, dV, dQKV, dKV; - // FP16/BF16: dqkv_fake_dtype = kFloat16/kBFloat16, dQ/dK/dV.dtype = torch.float16/torch.bfloat16 - // FP8DS: dqkv_fake_dtype = kFloat16/kBFloat16, dQ/dK/dV.dtype = torch.uint8 - // FP8CS/MXFP8: dqkv_fake_dtype = kFloat16/kBFloat16, dQ/dK/dV.dtype = torch.float16/torch.bfloat16 - auto options = torch::TensorOptions().dtype(fake_dtype).device(torch::kCUDA); - if (detail::IsFloat8Quantizers(dqkv_quantizer.ptr())) { - options = options.dtype(torch::kUInt8); - } - - NVTE_QKV_Layout_Group layout_group = nvte_get_qkv_layout_group(dqkv_layout); - std::vector tmp_shape; - switch (layout_group) { - case NVTE_QKV_Layout_Group::NVTE_3HD: - tmp_shape = std::vector{dQ_shape.begin(), dQ_shape.end()}; - tmp_shape.insert(tmp_shape.begin() + tmp_shape.size() - 2, int64_t(3)); - dQKV = torch::empty(c10::IntArrayRef(tmp_shape), options); - dQ = dQKV.index({"...", torch::indexing::Slice(0, 1, 1), - torch::indexing::Slice(0, torch::indexing::None, 1), - torch::indexing::Slice(0, torch::indexing::None, 1)}) - .squeeze(tmp_shape.size() - 3); - dK = dQKV.index({"...", torch::indexing::Slice(1, 2, 1), - torch::indexing::Slice(0, torch::indexing::None, 1), - torch::indexing::Slice(0, torch::indexing::None, 1)}) - .squeeze(tmp_shape.size() - 3); - dV = dQKV.index({"...", torch::indexing::Slice(2, torch::indexing::None, 1), - torch::indexing::Slice(0, torch::indexing::None, 1), - torch::indexing::Slice(0, torch::indexing::None, 1)}) - .squeeze(tmp_shape.size() - 3); - break; - case NVTE_QKV_Layout_Group::NVTE_H3D: - tmp_shape = std::vector{dQ_shape.begin(), dQ_shape.end()}; - tmp_shape.insert(tmp_shape.begin() + tmp_shape.size() - 1, int64_t(3)); - dQKV = torch::empty(c10::IntArrayRef(tmp_shape), options); - dQ = dQKV.index({"...", torch::indexing::Slice(0, 1, 1), - torch::indexing::Slice(0, torch::indexing::None, 1)}) - .squeeze(tmp_shape.size() - 2); - dK = dQKV.index({"...", torch::indexing::Slice(1, 2, 1), - torch::indexing::Slice(0, torch::indexing::None, 1)}) - .squeeze(tmp_shape.size() - 2); - dV = dQKV.index({"...", torch::indexing::Slice(2, torch::indexing::None, 1), - torch::indexing::Slice(0, torch::indexing::None, 1)}) - .squeeze(tmp_shape.size() - 2); - break; - case NVTE_QKV_Layout_Group::NVTE_HD_2HD: - tmp_shape = std::vector(dQ_shape.begin(), dQ_shape.end()); - dQ = torch::empty(tmp_shape, options); - tmp_shape = std::vector{dK_shape.begin(), dK_shape.end()}; - tmp_shape.insert(tmp_shape.begin() + tmp_shape.size() - 2, int64_t(2)); - dKV = torch::empty(c10::IntArrayRef(tmp_shape), options); - dK = dKV.index({"...", torch::indexing::Slice(0, 1, 1), - torch::indexing::Slice(0, torch::indexing::None, 1), - torch::indexing::Slice(0, torch::indexing::None, 1)}) - .squeeze(tmp_shape.size() - 3); - dV = dKV.index({"...", torch::indexing::Slice(1, torch::indexing::None, 1), - torch::indexing::Slice(0, torch::indexing::None, 1), - torch::indexing::Slice(0, torch::indexing::None, 1)}) - .squeeze(tmp_shape.size() - 3); - break; - case NVTE_QKV_Layout_Group::NVTE_HD_H2D: - tmp_shape = std::vector(dQ_shape.begin(), dQ_shape.end()); - dQ = torch::empty(tmp_shape, options); - tmp_shape = std::vector{dK_shape.begin(), dK_shape.end()}; - tmp_shape.insert(tmp_shape.begin() + tmp_shape.size() - 1, int64_t(2)); - dKV = torch::empty(c10::IntArrayRef(tmp_shape), options); - dK = dKV.index({"...", torch::indexing::Slice(0, 1, 1), - torch::indexing::Slice(0, torch::indexing::None, 1)}) - .squeeze(tmp_shape.size() - 2); - dV = dKV.index({"...", torch::indexing::Slice(1, torch::indexing::None, 1), - torch::indexing::Slice(0, torch::indexing::None, 1)}) - .squeeze(tmp_shape.size() - 2); - break; - case NVTE_QKV_Layout_Group::NVTE_HD_HD_HD: - case NVTE_QKV_Layout_Group::NVTE_SD_SD_SD: - tmp_shape = std::vector(dQ_shape.begin(), dQ_shape.end()); - dQ = torch::empty(tmp_shape, options); - tmp_shape = std::vector(dK_shape.begin(), dK_shape.end()); - dK = torch::empty(tmp_shape, options); - tmp_shape = std::vector(dV_shape.begin(), dV_shape.end()); - dV = torch::empty(tmp_shape, options); - break; - default: - NVTE_ERROR("QKV layout not supported!"); - } - - std::tie(te_dQ, py_dQ, dq_amax_buf) = - quantizer_helper(dqkv_quantizer, dQ_shape, dqkv_fake_dtype, true, dQ); - std::tie(te_dK, py_dK, dk_amax_buf) = - quantizer_helper(dqkv_quantizer, dK_shape, dqkv_fake_dtype, true, dK); - std::tie(te_dV, py_dV, dv_amax_buf) = - quantizer_helper(dqkv_quantizer, dV_shape, dqkv_fake_dtype, true, dV); - - // construct NVTE tensors - if (detail::IsFloat8Quantizers(dqkv_quantizer.ptr())) { - // FP8 - if (set_zero) { - // dQ/dK/dV may be strided views, so zero data through ATen and reset the shared amax - // separately. - if (dq_format == NVTE_QKV_Format::NVTE_THD) { - dQ.fill_(0); - } - if (dkv_format == NVTE_QKV_Format::NVTE_THD) { - dK.fill_(0); - dV.fill_(0); - } - if (dq_format == NVTE_QKV_Format::NVTE_THD || dkv_format == NVTE_QKV_Format::NVTE_THD) { - auto *fp8_quantizer = dynamic_cast(dQKV_quantizer.get()); - fp8_quantizer->amax.zero_(); - } - } - } else if (dqkv_quantizer.is_none() || - detail::IsFloat8CurrentScalingQuantizers(dqkv_quantizer.ptr()) || - detail::IsMXFP8Quantizers(dqkv_quantizer.ptr())) { - if (dq_format == NVTE_QKV_Format::NVTE_THD) { - dQ.fill_(0); - } - if (dkv_format == NVTE_QKV_Format::NVTE_THD) { - dK.fill_(0); - dV.fill_(0); - } - } else { - NVTE_ERROR("Fused attention only supports FP8 and BF16/FP16 data types. \n"); - } - - // create cu_seqlens tensorwrappers - auto cu_seqlens_q_sizes = cu_seqlens_q.sizes().vec(); - std::vector cu_seqlens_q_shape{cu_seqlens_q_sizes.begin(), cu_seqlens_q_sizes.end()}; - auto cu_seqlens_kv_sizes = cu_seqlens_kv.sizes().vec(); - std::vector cu_seqlens_kv_shape{cu_seqlens_kv_sizes.begin(), cu_seqlens_kv_sizes.end()}; - TensorWrapper te_cu_seqlens_q, te_cu_seqlens_kv; - te_cu_seqlens_q = makeTransformerEngineTensor(cu_seqlens_q.data_ptr(), cu_seqlens_q_shape, - DType::kInt32, nullptr, nullptr, nullptr); - te_cu_seqlens_kv = makeTransformerEngineTensor(cu_seqlens_kv.data_ptr(), cu_seqlens_kv_shape, - DType::kInt32, nullptr, nullptr, nullptr); - - TensorWrapper te_cu_seqlens_q_padded, te_cu_seqlens_kv_padded; - if ((cu_seqlens_q_padded.has_value()) && (cu_seqlens_kv_padded.has_value())) { - auto cu_seqlens_q_padded_sizes = cu_seqlens_q_padded.value().sizes().vec(); - std::vector cu_seqlens_q_padded_shape{cu_seqlens_q_padded_sizes.begin(), - cu_seqlens_q_padded_sizes.end()}; - auto cu_seqlens_kv_padded_sizes = cu_seqlens_kv_padded.value().sizes().vec(); - std::vector cu_seqlens_kv_padded_shape{cu_seqlens_kv_padded_sizes.begin(), - cu_seqlens_kv_padded_sizes.end()}; - te_cu_seqlens_q_padded = makeTransformerEngineTensor(cu_seqlens_q_padded.value().data_ptr(), - cu_seqlens_q_padded_shape, DType::kInt32); - te_cu_seqlens_kv_padded = makeTransformerEngineTensor( - cu_seqlens_kv_padded.value().data_ptr(), cu_seqlens_kv_padded_shape, DType::kInt32); - } - - // convert auxiliary tensors from forward to NVTETensors - NVTETensorPack nvte_aux_tensor_pack; - nvte_tensor_pack_create(&nvte_aux_tensor_pack); - nvte_aux_tensor_pack.size = Aux_CTX_Tensors.size(); - for (size_t i = 0; i < nvte_aux_tensor_pack.size; ++i) { - const std::vector &signed_shape = Aux_CTX_Tensors[i].sizes().vec(); - const std::vector tmp(signed_shape.begin(), signed_shape.end()); - - NVTEBasicTensor temp_data = { - Aux_CTX_Tensors[i].data_ptr(), - static_cast(GetTransformerEngineDType(Aux_CTX_Tensors[i].scalar_type())), - nvte_make_shape(tmp.data(), tmp.size())}; - nvte_set_tensor_param(&nvte_aux_tensor_pack.tensors[i], kNVTERowwiseData, &temp_data); - } - - // create dBias the same shape as Bias - at::Tensor dBias; - TensorWrapper te_dBias; - if ((bias_type != NVTE_NO_BIAS) && (bias_type != NVTE_ALIBI)) { - if (nvte_aux_tensor_pack.size >= 2) { - std::vector bias_shape(Aux_CTX_Tensors[nvte_aux_tensor_pack.size - 1].sizes().vec()); - dBias = torch::empty(bias_shape, options); - te_dBias = makeTransformerEngineTensor(dBias); - } else { - dBias = torch::empty({1, static_cast(h_q), static_cast(max_seqlen_q), - static_cast(max_seqlen_kv)}, - options); - te_dBias = makeTransformerEngineTensor(dBias); - } - if (nvte_get_qkv_format(qkv_layout) == NVTE_QKV_Format::NVTE_THD) { - dBias.fill_(0); - } - } - - // create dSoftmaxOffset in the same shape as SoftmaxOffset - at::Tensor dSoftmaxOffset; - TensorWrapper te_dSoftmaxOffset; - if (softmax_type != NVTE_VANILLA_SOFTMAX) { - options = torch::TensorOptions().dtype(at::kFloat).device(torch::kCUDA); - dSoftmaxOffset = torch::empty({1, static_cast(h_q), 1, 1}, options); - te_dSoftmaxOffset = makeTransformerEngineTensor(dSoftmaxOffset); - } - - // create workspace - TensorWrapper workspace; - - // build the parameter object - FusedAttnBwdParamsWrapper params; - params.set_Q(te_Q.data()) - .set_K(te_K.data()) - .set_V(te_V.data()) - .set_O(te_O.data()) - .set_dO(te_dO.data()) - .set_S(te_S.data()) - .set_dP(te_dP.data()) - .set_Aux_CTX_Tensors(&nvte_aux_tensor_pack) - .set_dQ(te_dQ.data()) - .set_dK(te_dK.data()) - .set_dV(te_dV.data()) - .set_dBias(te_dBias.data()) - .set_dSoftmaxOffset(te_dSoftmaxOffset.data()) - .set_cu_seqlens_q(te_cu_seqlens_q.data()) - .set_cu_seqlens_kv(te_cu_seqlens_kv.data()) - .set_cu_seqlens_q_padded(te_cu_seqlens_q_padded.data()) - .set_cu_seqlens_kv_padded(te_cu_seqlens_kv_padded.data()) - .set_deterministic(deterministic) - .set_cuda_graph(cuda_graph) - .set_attn_mask_type(attn_mask_type) - .set_bias_type(bias_type) - .set_window_size_left(window_size[0]) - .set_window_size_right(window_size[1]) - .set_bottom_right_diagonal(bottom_right_diagonal) - .set_softmax_type(softmax_type) - .set_dropout(p_dropout) - .set_attn_scale(attn_scale) - .set_qkv_layout(qkv_layout) - .set_o_format(o_format) - .set_do_format(do_format) - .set_dqkv_layout(dqkv_layout) - .set_qkv_scale_inv_format(qkv_scale_inv_format) - .set_do_scale_inv_format(do_scale_inv_format) - .set_max_seqlen_q(max_seqlen_q) - .set_max_seqlen_kv(max_seqlen_kv) - .set_workspace(workspace.data()) - .set_stream(at::cuda::getCurrentCUDAStream()); - - // populate tensors with appropriate shapes and dtypes - NVTE_SCOPED_GIL_RELEASE({ nvte_fused_attn_bwd_v2(params); }); - - // allocate memory for workspace - auto workspace_data = allocateSpace(workspace.shape(), workspace.dtype()); - workspace = - makeTransformerEngineTensor(workspace_data.data_ptr(), workspace.shape(), workspace.dtype()); - params.set_workspace(workspace.data()); - - // execute kernel - NVTE_SCOPED_GIL_RELEASE({ nvte_fused_attn_bwd_v2(params); }); - - // destroy tensor wrappers - nvte_tensor_pack_destroy(&nvte_aux_tensor_pack); - - return {py_dQ, py_dK, py_dV, py::cast(dBias), py::cast(dSoftmaxOffset)}; -} - at::Tensor fa_prepare_fwd(at::Tensor qkvi) { NVTE_CHECK(qkvi.dim() == 4, "Expected 4-dim tensor."); NVTE_CHECK(qkvi.scalar_type() == at::ScalarType::Half || diff --git a/transformer_engine/pytorch/csrc/extensions/attention_rng.cpp b/transformer_engine/pytorch/csrc/extensions/attention_rng.cpp new file mode 100644 index 00000000000..ad3410a63a7 --- /dev/null +++ b/transformer_engine/pytorch/csrc/extensions/attention_rng.cpp @@ -0,0 +1,26 @@ +/************************************************************************* + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * + * See LICENSE for license information. + ************************************************************************/ + +#include "../extensions.h" +#include "common.h" + +namespace transformer_engine::pytorch { + +at::Tensor get_cudnn_attention_rng_state(const std::optional rng_gen, + size_t increment) { + auto gen = at::get_generator_or_default( + rng_gen, at::cuda::detail::getDefaultCUDAGenerator()); + at::PhiloxCudaState philox_args = init_philox_state(gen, increment); + auto options = torch::TensorOptions().dtype(torch::kInt64).device(torch::kCUDA); + auto rng_state = torch::empty({2}, options); + // PhiloxCudaState contains device pointers while CUDA graph capture is + // active. Unpacking on the current stream therefore preserves PyTorch's + // graph-safe intragraph offset semantics. + philox_unpack(philox_args, static_cast(rng_state.data_ptr())); + return rng_state; +} + +} // namespace transformer_engine::pytorch diff --git a/transformer_engine/pytorch/csrc/extensions/pybind.cpp b/transformer_engine/pytorch/csrc/extensions/pybind.cpp index eff5b9aafec..e4fe8b32b50 100644 --- a/transformer_engine/pytorch/csrc/extensions/pybind.cpp +++ b/transformer_engine/pytorch/csrc/extensions/pybind.cpp @@ -469,8 +469,6 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { m.def("swap_first_dims", &transformer_engine::pytorch::swap_first_dims, "Swap first two tensor dimensions", py::arg("tensor"), py::kw_only(), py::arg("out"), py::call_guard()); - m.def("get_fused_attn_backend", &transformer_engine::pytorch::get_fused_attn_backend, - "Get Fused Attention backend", py::arg("fused_attn_params")); m.def("compute_amax", &transformer_engine::pytorch::compute_amax, "Compute absolute max value in tensor", py::arg("input"), py::arg("amax"), py::call_guard()); @@ -548,6 +546,11 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { py::call_guard()); // attention kernels + m.def("get_cudnn_attention_rng_state", + &transformer_engine::pytorch::get_cudnn_attention_rng_state, + "Reserve and unpack a graph-safe Philox state for Python cuDNN attention", + py::arg("rng_gen") = py::none(), py::arg("increment"), + py::call_guard()); m.def("fa_prepare_fwd", &transformer_engine::pytorch::fa_prepare_fwd, "Prepare QKV for Flash Attention", py::call_guard()); m.def("fa_prepare_bwd", &transformer_engine::pytorch::fa_prepare_bwd, @@ -561,10 +564,6 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { m.def("multi_tensor_pad_last_dim", &transformer_engine::pytorch::multi_tensor_pad_last_dim, "Pad multiple tensors' last dimension to a common alignment.", py::arg("inputs"), py::arg("alignment"), py::call_guard()); - m.def("fused_attn_fwd", &transformer_engine::pytorch::fused_attn_fwd, - "Fused Attention FP8/BF16/FP16 FWD with separate Q, K and V"); - m.def("fused_attn_bwd", &transformer_engine::pytorch::fused_attn_bwd, - "Fused Attention FP8/BF16/FP16 BWD with separate Q, K and V"); m.def("copy_to_kv_cache", &transformer_engine::pytorch::copy_to_kv_cache, "Copy new KV tokens to KV cache", py::call_guard()); m.def("convert_thd_to_bshd", &transformer_engine::pytorch::convert_thd_to_bshd,