diff --git a/build_tools/jax.py b/build_tools/jax.py index 891fb63eec8..56843e63c12 100644 --- a/build_tools/jax.py +++ b/build_tools/jax.py @@ -6,6 +6,7 @@ import os import warnings +from importlib.metadata import version as distribution_version from pathlib import Path from packaging import version @@ -27,7 +28,12 @@ def install_requirements() -> List[str]: """Install dependencies for TE/JAX extensions.""" - return ["jax", "flax>=0.7.1", "nvidia-cudnn-frontend>=1.25.0", "triton"] + # Python serializes graphs for C++ built with headers from this distribution. + # Keep the runtime serializer at the exact version used by this build. + frontend_version = distribution_version("nvidia-cudnn-frontend") + if version.parse(frontend_version) < version.parse("1.29.0"): + raise RuntimeError("Building TE/JAX requires nvidia-cudnn-frontend>=1.29.0.") + return ["jax", "flax>=0.7.1", f"nvidia-cudnn-frontend=={frontend_version}", "triton"] def test_requirements() -> List[str]: diff --git a/build_tools/wheel_utils/build_wheels.sh b/build_tools/wheel_utils/build_wheels.sh index 761f12390e1..0e49c2c2e7a 100644 --- a/build_tools/wheel_utils/build_wheels.sh +++ b/build_tools/wheel_utils/build_wheels.sh @@ -24,7 +24,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' 'apache-tvm-ffi>=0.1.12' +/opt/python/cp310-cp310/bin/pip install cmake pybind11[global] ninja setuptools wheel 'nvidia-cudnn-frontend>=1.29.0' 'apache-tvm-ffi>=0.1.12' if $BUILD_METAPACKAGE ; then cd /TransformerEngine diff --git a/docs/api/jax.rst b/docs/api/jax.rst index 7a31c9d3797..130e66bd5c1 100644 --- a/docs/api/jax.rst +++ b/docs/api/jax.rst @@ -59,3 +59,9 @@ Modules :members: __call__ .. autoapifunction:: transformer_engine.jax.flax.extend_logical_axis_rules + +Attention functions +------------------- + +.. autoapifunction:: transformer_engine.jax.attention.fused_attn +.. autoapifunction:: transformer_engine.jax.attention.is_fused_attn_kernel_available diff --git a/docs/envvars.rst b/docs/envvars.rst index 0d559030216..96c3335efa7 100644 --- a/docs/envvars.rst +++ b/docs/envvars.rst @@ -238,7 +238,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 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 plan compilation, and each execution. JAX reports Python graph-cache and build events during tracing; execution events apply to the PyTorch and common C++ paths. 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. .. envvar:: NVTE_ALLOW_NONDETERMINISTIC_ALGO @@ -280,7 +280,7 @@ FP8 Configuration :Type: ``int`` (0 or 1) :Default: ``1`` - :Description: Enable FP8 in the backward pass of DotProductAttention. ``1`` = FP8 forward and backward, ``0`` = FP8 forward and FP16/BF16 backward. + :Description: Enable FP8 in the backward pass of DotProductAttention in PyTorch and JAX. ``1`` = FP8 forward and backward, ``0`` = FP8 forward and FP16/BF16 backward. In JAX this applies when attention is configured with ``fp8=True``. .. envvar:: NVTE_DPA_FP8CS_O_in_F16 diff --git a/docs/examples/jax/attention.rst b/docs/examples/jax/attention.rst index f82a6ffe621..23f52b105d2 100644 --- a/docs/examples/jax/attention.rst +++ b/docs/examples/jax/attention.rst @@ -16,6 +16,42 @@ FlexAttention-style customization. The API also supports different Q/K and V head dimensions, as used by the attention operation after the projection stages in DeepSeek-style MLA. +cuDNN execution and FP8 +---------------------- + +JAX fused attention constructs and selects cuDNN execution plans with the +``nvidia-cudnn-frontend`` Python package. The compiled JAX FFI executor caches +those plans and its own cuDNN handles per device. Ordinary attention and +``score_mod`` share this executor; neither calls TE common's attention APIs. +The Python frontend and the headers used to build the JAX extension must have +the same version, since execution plans are serialized between the two. + +Pass ``fp8=True`` to ``transformer_engine.jax.attention.fused_attn`` or to the +Flax ``DotProductAttention``, ``MultiHeadAttention`` or ``TransformerLayer`` +module to enable cuDNN FP8 attention. Inputs remain FP16 or BF16. Query, key +and value are quantized to E4M3 using their current tensor maximum; backward +quantizes the output gradient to E5M2. Outputs and input gradients retain the +input dtype, and softmax statistics use FP32. Packed layouts use one scale per +packed tensor. This option is independent of the GEMM quantization recipe and +does not maintain delayed-scaling history. + +``NVTE_FP8_DPA_BWD=0`` selects FP16/BF16 backward while keeping FP8 forward. +The backward pass uses dequantized forward operands, so both passes operate on +the same quantized query, key and value. The default, ``1``, uses FP8 for both. + +FP8 availability depends on cuDNN, GPU architecture, layout and the selected +features. Unsupported combinations produce a diagnostic rather than silently +using a different precision. For example, cuDNN 9.26 rejects FP8 backward with +dropout on Blackwell; use ``NVTE_FP8_DPA_BWD=0`` for that combination. +Tensor-scaled FP8 is supported; MXFP8 attention, +paged KV-cache APIs and quantized module inputs/outputs require further JAX API +and frontend-binding work. ``score_mod`` currently supports FP16/BF16 only. + +``AttnBiasType.ALIBI`` (``attn_bias_type="alibi"`` in Flax) enables cuDNN's +ALiBi mask without an explicit bias tensor. Explicit dropout keys have JAX's +usual replay semantics: using the same key repeats the mask, including when +recomputing a forward pass for its backward pass. + For long contexts, selected BSHD and THD configurations support context parallelism with Ring or AllGather collectives. Exact fused-kernel availability depends on the input shape, dtype, GPU architecture, and feature combination; diff --git a/docs/installation.rst b/docs/installation.rst index 9d3eada827c..9fac9ee56a6 100644 --- a/docs/installation.rst +++ b/docs/installation.rst @@ -16,6 +16,30 @@ Prerequisites 3. |driver link|_ supporting CUDA 12.1 or later. 4. `cuDNN 9.12 `__ or later. +The PyTorch and JAX cuDNN attention backends use shared graph builders through +the cuDNN frontend Python API. Both installations require +``nvidia-cudnn-frontend`` version 1.29.0 or later. JAX builds pin their runtime +dependency to the exact frontend version that supplies the build headers. +Graph descriptions are cached by tensor metadata +and target device or architecture. Support queries select supported plans without +compiling their kernels. PyTorch compiles on first execution and uses its CUDA +generator for dropout, including CUDA graph capture and replay. JAX compiles +during tracing when workspace sizes are needed, then serializes plans for its +native FFI executor. The JAX extension and the frontend Python package must use +matching cuDNN frontend versions; rebuild the extension after upgrading that +package. + +Automatic ``score_mod`` graph caching includes the callback's code and snapshots +of builtin scalar and container defaults, including keyword-only defaults. +Callbacks with tensor or opaque stateful defaults bypass automatic caching. +Stateful callbacks can provide ``score_mod_graph_cache_key`` to explicitly +describe their graph structure; pass runtime data through ``score_mod_tensors`` +and ``score_mod_bprop_tensors``. + +On SM12x GPUs, including SM120 and SM121, packed attention uses dense +per-sequence softmax statistics internally. Graph construction and context +parallelism use the same statistics layout. + If the CUDA Toolkit headers are not available at runtime in a standard installation path, e.g. within `CUDA_HOME`, set `NVTE_CUDA_INCLUDE_PATH` in the environment. diff --git a/tests/jax/test_build_tools.py b/tests/jax/test_build_tools.py new file mode 100644 index 00000000000..714aafec106 --- /dev/null +++ b/tests/jax/test_build_tools.py @@ -0,0 +1,49 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""Dependency metadata for builds that deserialize cuDNN frontend graphs.""" + +from email.parser import Parser +from io import StringIO + +import pytest +from packaging.requirements import Requirement +from setuptools import Distribution + +from build_tools import jax as jax_build + + +@pytest.mark.parametrize("frontend_version", ("1.29.0", "1.30.0")) +def test_jax_metadata_pins_build_frontend(monkeypatch, frontend_version): + """Built metadata must reject runtime serializers other than the build version.""" + + def distribution_version(package): + assert package == "nvidia-cudnn-frontend" + return frontend_version + + monkeypatch.setattr(jax_build, "distribution_version", distribution_version) + dist = Distribution( + { + "name": "transformer_engine_jax", + "version": "0.0.0", + "install_requires": jax_build.install_requirements(), + } + ) + metadata = StringIO() + dist.metadata.write_pkg_file(metadata) + requirements = [ + Requirement(value) + for value in Parser().parsestr(metadata.getvalue()).get_all("Requires-Dist") + ] + frontend = next(req for req in requirements if req.name == "nvidia-cudnn-frontend") + assert frontend.specifier.contains(frontend_version) + for other_version in ("1.28.0", "1.29.0", "1.29.1", "1.30.0", "1.30.1", "1.31.0"): + assert frontend.specifier.contains(other_version) == (other_version == frontend_version) + + +def test_jax_build_requires_minimum_frontend(monkeypatch): + """Pinning the installed version must preserve the minimum build requirement.""" + monkeypatch.setattr(jax_build, "distribution_version", lambda _package: "1.28.0") + with pytest.raises(RuntimeError, match="nvidia-cudnn-frontend>=1.29.0"): + jax_build.install_requirements() diff --git a/tests/jax/test_fused_attn.py b/tests/jax/test_fused_attn.py index 4e5c2948a21..33157f2b183 100644 --- a/tests/jax/test_fused_attn.py +++ b/tests/jax/test_fused_attn.py @@ -428,6 +428,123 @@ def test_fused_attn_backend_message(): assert message != "" +@pytest.mark.parametrize("setup_error", (None, RuntimeError("frontend version mismatch"))) +def test_fused_attn_support_query_defers_compilation(setup_error, monkeypatch): + """Support queries must avoid compilation and propagate unexpected setup failures.""" + from transformer_engine.jax.cpp_extensions import cudnn_attention + + def forbidden(*_args, **_kwargs): + raise AssertionError("Support query compiled an execution plan") + + monkeypatch.setattr(cudnn_attention, "compile_graph", forbidden) + cudnn_attention._supported_graph.cache_clear() + if setup_error is not None: + + def fail(*_args, **_kwargs): + raise setup_error + + monkeypatch.setattr(cudnn_attention, "check_attention_support", fail) + helper = FusedAttnHelper( + is_training=True, + batch_size=2, + q_dtype=jnp.bfloat16, + kv_dtype=jnp.bfloat16, + qkv_layout=QKVLayout.BSHD_BSHD_BSHD, + attn_bias_type=AttnBiasType.NO_BIAS, + attn_mask_type=AttnMaskType.NO_MASK, + softmax_type=AttnSoftmaxType.VANILLA_SOFTMAX, + dropout_probability=0.0, + q_num_heads=4, + kv_num_heads=4, + q_max_seqlen=128, + kv_max_seqlen=128, + head_dim_qk=64, + head_dim_v=64, + window_size=(-1, -1), + attn_scale=0.125, + ) + if setup_error is not None: + with pytest.raises(RuntimeError, match=str(setup_error)): + helper.get_fused_attn_backend() + else: + backend, message = helper.get_fused_attn_backend() + if backend == NVTE_Fused_Attn_Backend.NVTE_No_Backend: + pytest.skip(f"Baseline attention is unsupported: {message}") + + +@pytest.mark.parametrize("batch,seqlen", ((0, 128), (2, 0))) +def test_cudnn_attention_empty_ragged_capacities(batch, seqlen, monkeypatch): + """Empty storage must retain positive graph capacities for ragged metadata.""" + from types import SimpleNamespace + from transformer_engine.jax.cpp_extensions import cudnn_attention + + monkeypatch.setattr(cudnn_attention, "cudnn_version", lambda: 92700) + config = SimpleNamespace( + qkv_layout=QKVLayout.THD_THD_THD, + max_segments_per_seq=1, + fp8=False, + attn_mask_type=AttnMaskType.PADDING_MASK, + attn_bias_type=AttnBiasType.NO_BIAS, + softmax_type=AttnSoftmaxType.VANILLA_SOFTMAX, + effective_window_size=(-1, -1), + scaling_factor=0.125, + bottom_right_diagonal=False, + dropout_probability=0.0, + is_training=False, + return_max_logit=False, + ) + shapes = tuple( + (shape, jnp.dtype(jnp.bfloat16)) + for shape in ((batch, seqlen, 4, 64),) * 3 + ((0,), (1, 4, 1, 1)) + ) + description = cudnn_attention._describe_graph("fwd", shapes, config, (100,), True, True) + assert description.metadata_batch == 32 + for name in ("q", "k", "v"): + assert dict(description.spec.tensors)[name].dim == (32, 4, 1024, 64) + + +@pytest.mark.parametrize("sm_arch", (100, 120, 121, 129, 130)) +@pytest.mark.parametrize("direction", ("fwd", "bwd")) +def test_cudnn_attention_sm12x_dense_statistics(sm_arch, direction, monkeypatch): + """All SM12x variants use dense per-segment statistics and per-sequence dimensions.""" + from types import SimpleNamespace + from transformer_engine.jax.cpp_extensions import cudnn_attention + + monkeypatch.setattr(cudnn_attention, "cudnn_version", lambda: 92700) + config = SimpleNamespace( + qkv_layout=QKVLayout.THD_THD_THD, + max_segments_per_seq=2, + fp8=False, + attn_mask_type=AttnMaskType.PADDING_MASK, + attn_bias_type=AttnBiasType.NO_BIAS, + softmax_type=AttnSoftmaxType.VANILLA_SOFTMAX, + effective_window_size=(-1, -1), + scaling_factor=0.125, + bottom_right_diagonal=False, + dropout_probability=0.0, + is_training=True, + return_max_logit=False, + ) + shapes = tuple( + (shape, jnp.dtype(jnp.bfloat16)) for shape in ((2, 128, 4, 64),) * 3 + ((0,), (1, 4, 1, 1)) + ) + description = cudnn_attention._describe_graph(direction, shapes, config, (sm_arch,), True, True) + tensors = dict(description.spec.tensors) + if 120 <= sm_arch < 130: + assert not description.ragged_stats + assert description.metadata_batch == 4 + assert tensors["q"].dim == (4, 4, 128, 64) + assert tensors["stats"].dim == (4, 4, 128, 1) + assert tensors["stats"].ragged_offset == "" + assert description.spec.options.max_total_seqlen_q == 0 + assert description.spec.options.max_total_seqlen_kv == 0 + else: + assert description.ragged_stats + assert description.metadata_batch == 32 + assert tensors["q"].dim == (32, 4, 1024, 64) + assert tensors["stats"].ragged_offset == "offset_stats" + + class BiasShape(Enum): """ Enum class to represent the different bias shapes used in the fused attention. @@ -2106,3 +2223,451 @@ def test_backward( seq_desc_format, is_self_attention=None, ) + + +@pytest.mark.parametrize("dtype", [jnp.float16, jnp.bfloat16]) +@pytest.mark.parametrize("fp8_backward", [False, True]) +@pytest.mark.parametrize( + "layout", + [ + QKVLayout.BS3HD, + QKVLayout.BSHD_BS2HD, + QKVLayout.BSHD_BSHD_BSHD, + ], +) +def test_cudnn_python_fp8_forward_backward(dtype, layout, fp8_backward, monkeypatch): + """Exercise real FP8 SDPA with exactly representable operands and probabilities.""" + monkeypatch.setenv("NVTE_FP8_DPA_BWD", "1" if fp8_backward else "0") + q = jnp.full((2, 128, 4, 64), 0.5, dtype) + k = jnp.full_like(q, 0.5) + v = jnp.zeros_like(q).at[:, ::2, :, 0].set(1.0) + if layout.is_qkvpacked(): + operands = (jnp.stack((q, k, v), axis=-3),) + elif layout.is_kvpacked(): + operands = (q, jnp.stack((k, v), axis=-3)) + else: + operands = (q, k, v) + + def unpack(operands): + if layout.is_qkvpacked(): + return tuple(operands[0][:, :, index] for index in range(3)) + if layout.is_kvpacked(): + return operands[0], operands[1][:, :, 0], operands[1][:, :, 1] + return operands + + def reference(operands): + query, key, value = (tensor.astype(jnp.float32) for tensor in unpack(operands)) + scores = jnp.einsum("bqhd,bkhd->bhqk", query, key) * 0.125 + return jnp.einsum("bhqk,bkhd->bqhd", jax.nn.softmax(scores, axis=-1), value) + + def attention(operands): + return fused_attn( + operands, + None, + None, + None, + attn_bias_type=AttnBiasType.NO_BIAS, + attn_mask_type=AttnMaskType.NO_MASK, + qkv_layout=layout, + softmax_type=AttnSoftmaxType.VANILLA_SOFTMAX, + scaling_factor=0.125, + dropout_probability=0.0, + is_training=True, + fp8=True, + ) + + if get_device_compute_capability(0) < 90: + pytest.skip("cuDNN FP8 SDPA requires Hopper or newer.") + actual = jax.jit(attention)(operands) + expected = jax.jit(reference)(operands) + np.testing.assert_allclose(actual, expected, rtol=5e-3, atol=5e-3) + actual_grads = jax.jit(jax.grad(lambda arrays: jnp.sum(attention(arrays))))(operands) + expected_grads = jax.jit(jax.grad(lambda arrays: jnp.sum(reference(arrays))))(operands) + for actual_grad, expected_grad in zip(actual_grads, expected_grads): + assert actual_grad.dtype == dtype + np.testing.assert_allclose(actual_grad, expected_grad, rtol=1e-2, atol=1e-2) + + +def test_cudnn_python_attention_does_not_call_common(monkeypatch): + """Backend selection, workspace sizing and execution are independent of common.""" + import transformer_engine_jax + from transformer_engine.jax.cpp_extensions import cudnn_attention + + def forbidden(*_args, **_kwargs): + raise AssertionError("JAX attention called TE common") + + for name in ( + "get_fused_attn_backend", + "get_fused_attn_fwd_workspace_sizes", + "get_fused_attn_bwd_workspace_sizes", + "nvte_get_qkv_format", + ): + monkeypatch.setattr(transformer_engine_jax, name, forbidden, raising=False) + cudnn_attention._build_graph.cache_clear() + tensor = jnp.ones((1, 128, 2, 64), jnp.bfloat16) + output = jax.jit( + lambda q: fused_attn( + (q, q, q), + None, + None, + None, + attn_bias_type=AttnBiasType.NO_BIAS, + attn_mask_type=AttnMaskType.NO_MASK, + qkv_layout=QKVLayout.BSHD_BSHD_BSHD, + softmax_type=AttnSoftmaxType.VANILLA_SOFTMAX, + scaling_factor=0.125, + dropout_probability=0.0, + is_training=True, + ) + )(tensor) + np.testing.assert_array_equal(output, tensor) + gradients = jax.jit( + jax.grad( + lambda q: jnp.sum( + fused_attn( + (q, q, q), + None, + None, + None, + attn_bias_type=AttnBiasType.NO_BIAS, + attn_mask_type=AttnMaskType.NO_MASK, + qkv_layout=QKVLayout.BSHD_BSHD_BSHD, + softmax_type=AttnSoftmaxType.VANILLA_SOFTMAX, + scaling_factor=0.125, + dropout_probability=0.0, + is_training=True, + ) + ) + ) + )(tensor) + np.testing.assert_array_equal(gradients, tensor) + + +@pytest.mark.parametrize( + "kind,sm_arch", + [ + ("dense", None), + ("causal", None), + ("padding", None), + ("window", None), + ("ragged", None), + ("ragged_dense_stats", 120), + ("ragged_dense_stats_f16", 120), + ("ragged_dense_stats_f16", 121), + ("ragged_dense_stats_f16", 129), + ], +) +def test_cudnn_python_fp8_gqa_masks(kind, sm_arch, monkeypatch, request): + """Compare non-uniform FP8 attention and gradients against native FP32 GQA.""" + ragged = kind.startswith("ragged") + fp8 = not kind.endswith("_f16") + if get_device_compute_capability(0) < (100 if ragged else 90): + pytest.skip("FP8 ragged backward requires Blackwell; dense FP8 requires Hopper.") + if sm_arch is not None: + from transformer_engine.jax.cpp_extensions import cudnn_attention + + # Test the dense statistics conversion on the available GPU. + # Architecture is fixed in a real process, so discard JAX's abstract-value + # cache when mocking a different architecture's Max tensor layout. + jax.clear_caches() + request.addfinalizer(jax.clear_caches) + monkeypatch.setattr(cudnn_attention, "device_capabilities", lambda: (sm_arch,)) + keys = jax.random.split(jax.random.PRNGKey(123), 4) + q = jax.random.choice(keys[0], jnp.array([-0.5, 0.5], jnp.bfloat16), (2, 128, 4, 64)) + k = jax.random.choice(keys[1], jnp.array([-0.5, 0.5], jnp.bfloat16), (2, 128, 2, 64)) + v = jax.random.choice(keys[2], jnp.array([-1.0, 1.0], jnp.bfloat16), (2, 128, 2, 64)) + weight = jax.random.choice(keys[3], jnp.array([-1.0, 1.0], jnp.bfloat16), q.shape) + positions = jnp.arange(128) + mask = jnp.ones((2, 128, 128), jnp.bool_) + descriptor = None + attn_mask = AttnMaskType.NO_MASK + layout = QKVLayout.BSHD_BSHD_BSHD + window = None + if kind in ("causal", "window") or ragged: + mask &= positions[None, :, None] >= positions[None, None, :] + attn_mask = AttnMaskType.CAUSAL_MASK + if kind == "window": + window = (15, 0) + mask &= positions[None, :, None] - positions[None, None, :] <= 15 + if kind == "padding": + lengths = jnp.array([128, 91], jnp.int32) + descriptor = SequenceDescriptor.from_seqlens((lengths, lengths)) + mask &= positions[None, None, :] < lengths[:, None, None] + attn_mask = AttnMaskType.PADDING_MASK + if ragged: + ids = jnp.broadcast_to(positions // 64 + 1, (2, 128)) + descriptor = SequenceDescriptor.from_segment_ids_and_pos( + ids, jnp.broadcast_to(positions % 64, ids.shape) + ) + mask &= ids[:, :, None] == ids[:, None, :] + attn_mask = AttnMaskType.PADDING_CAUSAL_MASK + layout = QKVLayout.THD_THD_THD + + def reference(arrays): + query, key, value = (tensor.astype(jnp.float32) for tensor in arrays) + key, value = (jnp.repeat(tensor, 2, axis=-2) for tensor in (key, value)) + scores = jnp.einsum("bqhd,bkhd->bhqk", query, key) * 0.125 + scores = jnp.where(mask[:, None], scores, -jnp.inf) + result = jnp.einsum("bhqk,bkhd->bqhd", jax.nn.softmax(scores, axis=-1), value) + if kind == "padding": + result = jnp.where( + (positions[None, :] < lengths[:, None])[:, :, None, None], result, 0.0 + ) + return result + + def attention(arrays): + result = fused_attn( + arrays, + None, + descriptor, + None, + attn_bias_type=AttnBiasType.NO_BIAS, + attn_mask_type=attn_mask, + qkv_layout=layout, + softmax_type=AttnSoftmaxType.VANILLA_SOFTMAX, + scaling_factor=0.125, + dropout_probability=0.0, + is_training=True, + max_segments_per_seq=2 if ragged else 1, + window_size=window, + fp8=fp8, + return_max_logit=True, + ) + return result[0] + + arrays = (q, k, v) + actual, expected = jax.jit(attention)(arrays), jax.jit(reference)(arrays) + + # Operands are exactly representable. Error comes from FP8 softmax and dP, + # which use E4M3 and E5M2 internally, and the BF16 returned tensors. + def relative_error(actual, expected): + return float( + jnp.linalg.norm(actual.astype(jnp.float32) - expected.astype(jnp.float32)) + / jnp.linalg.norm(expected.astype(jnp.float32)) + ) + + assert relative_error(actual, expected) < (0.06 if fp8 else 0.01) + actual_grads = jax.jit(jax.grad(lambda xs: jnp.sum(attention(xs) * weight)))(arrays) + expected_grads = jax.jit(jax.grad(lambda xs: jnp.sum(reference(xs) * weight)))(arrays) + for actual_grad, expected_grad in zip(actual_grads, expected_grads): + assert relative_error(actual_grad, expected_grad) < (0.1 if fp8 else 0.01) + _, max_logit = jax.jit( + lambda xs: fused_attn( + xs, + None, + descriptor, + None, + attn_bias_type=AttnBiasType.NO_BIAS, + attn_mask_type=attn_mask, + qkv_layout=layout, + softmax_type=AttnSoftmaxType.VANILLA_SOFTMAX, + scaling_factor=0.125, + dropout_probability=0.0, + is_training=True, + max_segments_per_seq=2 if ragged else 1, + return_max_logit=True, + window_size=window, + fp8=fp8, + ) + )(arrays) + scores = ( + jnp.einsum( + "bqhd,bkhd->bhqk", + q.astype(jnp.float32), + jnp.repeat(k.astype(jnp.float32), 2, axis=-2), + ) + * 0.125 + ) + max_mask = mask + if kind == "padding": + max_mask &= positions[None, :, None] < lengths[:, None, None] + expected_max = jnp.max(jnp.where(max_mask[:, None], scores, -jnp.inf), axis=(0, 2, 3)) + np.testing.assert_allclose(max_logit, expected_max, rtol=1e-2, atol=1e-2) + + +def test_cudnn_python_alibi_forward_backward(): + """cuDNN ALiBi agrees with explicit slopes in native JAX attention.""" + keys = jax.random.split(jax.random.PRNGKey(42), 3) + arrays = tuple(jax.random.normal(key, (1, 128, 4, 64), jnp.float16) * 0.1 for key in keys) + positions = jnp.arange(128) + causal = positions[:, None] >= positions[None, :] + slopes = 2.0 ** (-2.0 * jnp.arange(1, 5, dtype=jnp.float32)) + + def reference(arrays): + q, k, v = (tensor.astype(jnp.float32) for tensor in arrays) + scores = jnp.einsum("bqhd,bkhd->bhqk", q, k) * 0.125 + scores += slopes[None, :, None, None] * positions[None, None, None, :] + scores = jnp.where(causal[None, None], scores, -jnp.inf) + return jnp.einsum("bhqk,bkhd->bqhd", jax.nn.softmax(scores, axis=-1), v) + + def attention(arrays): + return fused_attn( + arrays, + None, + None, + None, + attn_bias_type=AttnBiasType.ALIBI, + attn_mask_type=AttnMaskType.CAUSAL_MASK, + qkv_layout=QKVLayout.BSHD_BSHD_BSHD, + softmax_type=AttnSoftmaxType.VANILLA_SOFTMAX, + scaling_factor=0.125, + dropout_probability=0.0, + is_training=True, + ) + + np.testing.assert_allclose( + jax.jit(attention)(arrays), jax.jit(reference)(arrays), rtol=1e-2, atol=1e-3 + ) + actual_grads = jax.jit(jax.grad(lambda xs: jnp.sum(attention(xs))))(arrays) + expected_grads = jax.jit(jax.grad(lambda xs: jnp.sum(reference(xs))))(arrays) + for actual, expected in zip(actual_grads, expected_grads): + np.testing.assert_allclose(actual, expected, rtol=1e-2, atol=1e-3) + + +@pytest.mark.parametrize("softmax", ["vanilla", "learnable"]) +@pytest.mark.parametrize("deterministic", [False, True]) +@pytest.mark.parametrize("fp8", [False, True]) +def test_cudnn_python_flax_forward_backward(softmax, deterministic, fp8, monkeypatch): + """Flax attention preserves sink gradients even when dropout is disabled.""" + from transformer_engine.jax.flax import DotProductAttention + + if fp8 and get_device_compute_capability(0) < 90: + pytest.skip("cuDNN FP8 SDPA requires Hopper or newer.") + monkeypatch.setenv("NVTE_FUSED_ATTN", "1") + module = DotProductAttention( + head_dim=64, + num_attention_heads=2, + attn_mask_type="no_mask", + transpose_batch_sequence=False, + softmax_type=softmax, + fp8=fp8, + ) + q = jnp.zeros((1, 128, 2, 64), jnp.bfloat16) + v = jnp.zeros_like(q).at[:, ::2, :, 0].set(1.0) + variables = module.init(jax.random.PRNGKey(0), q, q, v, deterministic=deterministic) + + def forward(variables, value): + return module.apply(variables, q, q, value, deterministic=deterministic) + + actual = jax.jit(forward)(variables, v) + denominator = 129 if softmax == "learnable" else 128 + expected = jnp.zeros_like(q).at[:, :, :, 0].set(64 / denominator) + np.testing.assert_allclose(actual, expected, rtol=5e-3, atol=5e-3) + parameter_grad, value_grad = jax.jit( + jax.grad( + lambda params, value: jnp.sum(forward(params, value)), + argnums=(0, 1), + ) + )(variables, v) + assert value_grad.dtype == jnp.bfloat16 + np.testing.assert_allclose(value_grad, 128 / denominator, rtol=5e-2, atol=1e-2) + if softmax == "learnable": + expected_sink_grad = -128 * 64 / denominator**2 + leaves = jax.tree_util.tree_leaves(parameter_grad) + assert len(leaves) == 1 + np.testing.assert_allclose(leaves[0], expected_sink_grad, rtol=0.1, atol=1e-2) + + +@pytest.mark.parametrize("fp8,fp8_backward", [(False, False), (True, False), (True, True)]) +def test_cudnn_python_dropout_replay(fp8, fp8_backward, monkeypatch): + """Explicit keys replay dropout through JIT and forward recomputation.""" + if fp8 and get_device_compute_capability(0) < 90: + pytest.skip("cuDNN FP8 SDPA requires Hopper or newer.") + monkeypatch.setenv("NVTE_FP8_DPA_BWD", "1" if fp8_backward else "0") + helper = FusedAttnHelper( + is_training=True, + batch_size=1, + q_dtype=jnp.bfloat16, + kv_dtype=jnp.bfloat16, + qkv_layout=QKVLayout.BSHD_BSHD_BSHD, + attn_bias_type=AttnBiasType.NO_BIAS, + attn_mask_type=AttnMaskType.NO_MASK, + softmax_type=AttnSoftmaxType.VANILLA_SOFTMAX, + dropout_probability=0.2, + q_num_heads=2, + kv_num_heads=2, + q_max_seqlen=128, + kv_max_seqlen=128, + head_dim_qk=64, + head_dim_v=64, + window_size=(-1, -1), + attn_scale=0.125, + fp8=fp8, + ) + backend, reason = helper.get_fused_attn_backend() + if backend == NVTE_Fused_Attn_Backend.NVTE_No_Backend: + pytest.skip(reason) + arrays = tuple(jnp.ones((1, 128, 2, 64), jnp.bfloat16) for _ in range(3)) + + def attention(arrays, seed): + return fused_attn( + arrays, + None, + None, + seed, + attn_bias_type=AttnBiasType.NO_BIAS, + attn_mask_type=AttnMaskType.NO_MASK, + qkv_layout=QKVLayout.BSHD_BSHD_BSHD, + softmax_type=AttnSoftmaxType.VANILLA_SOFTMAX, + scaling_factor=0.125, + dropout_probability=0.2, + is_training=True, + fp8=fp8, + ) + + seed = jax.random.PRNGKey(123) + forward = jax.jit(attention) + actual = forward(arrays, seed) + np.testing.assert_array_equal(actual, forward(arrays, seed)) + assert not np.array_equal(actual, forward(arrays, jax.random.PRNGKey(456))) + ordinary = jax.jit(jax.grad(lambda xs: jnp.sum(attention(xs, seed))))(arrays) + recomputed = jax.jit(jax.grad(lambda xs: jnp.sum(jax.checkpoint(attention)(xs, seed))))(arrays) + for expected, actual in zip(ordinary, recomputed): + np.testing.assert_array_equal(actual, expected) + + +@pytest.mark.parametrize("strategy", [CPStrategy.ALL_GATHER, CPStrategy.RING]) +@pytest.mark.parametrize("fp8", [False, True]) +def test_cudnn_python_context_parallel(strategy, fp8): + """Two-device context parallel attention preserves FP8 forward and backward.""" + if jax.local_device_count() < 2: + pytest.skip("Context parallel attention requires two GPUs.") + if fp8 and get_device_compute_capability(0) < 90: + pytest.skip("cuDNN FP8 SDPA requires Hopper or newer.") + mesh = Mesh(np.asarray(jax.local_devices()[:2]), ("cp",)) + sharding = NamedSharding(mesh, PartitionSpec(None, "cp", None, None)) + q = jnp.full((2, 256, 2, 64), 0.5, jnp.bfloat16) + v = jnp.zeros_like(q).at[:, ::2, :, 0].set(1.0) + arrays = tuple(jax.device_put(value, sharding) for value in (q, q, v)) + + def attention(arrays): + return fused_attn( + arrays, + None, + None, + None, + attn_bias_type=AttnBiasType.NO_BIAS, + attn_mask_type=AttnMaskType.NO_MASK, + qkv_layout=QKVLayout.BSHD_BSHD_BSHD, + softmax_type=AttnSoftmaxType.VANILLA_SOFTMAX, + scaling_factor=0.125, + dropout_probability=0.0, + is_training=True, + context_parallel_axis="cp", + context_parallel_strategy=strategy, + fp8=fp8, + ) + + def reference(arrays): + query, key, value = (tensor.astype(jnp.float32) for tensor in arrays) + scores = jnp.einsum("bqhd,bkhd->bhqk", query, key) * 0.125 + return jnp.einsum("bhqk,bkhd->bqhd", jax.nn.softmax(scores, axis=-1), value) + + expected = jax.jit(reference)(arrays) + expected_grads = jax.jit(jax.grad(lambda xs: jnp.sum(reference(xs))))(arrays) + with mesh, autocast(mesh_resource=MeshResource(cp_resource="cp")): + actual = jax.jit(attention)(arrays) + actual_grads = jax.jit(jax.grad(lambda xs: jnp.sum(attention(xs))))(arrays) + np.testing.assert_allclose(actual, expected, rtol=5e-3, atol=5e-3) + for actual, expected in zip(actual_grads, expected_grads): + np.testing.assert_allclose(actual, expected, rtol=1e-2, atol=1e-2) diff --git a/tests/jax/test_fused_attn_score_mod.py b/tests/jax/test_fused_attn_score_mod.py index 5a613101585..c4050f3b0fd 100644 --- a/tests/jax/test_fused_attn_score_mod.py +++ b/tests/jax/test_fused_attn_score_mod.py @@ -2,7 +2,9 @@ # # See LICENSE for license information. """Tests for cuDNN frontend score_mod fused attention.""" + from math import sqrt +from types import FunctionType import jax import jax.numpy as jnp @@ -10,6 +12,7 @@ import pytest import transformer_engine.jax.cpp_extensions.flex_attention as tex_attention +import transformer_engine.jax.cpp_extensions.cudnn_graph as tex_cudnn_graph from transformer_engine.jax.attention import ( AttnBiasType, AttnMaskType, @@ -21,7 +24,6 @@ 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) _SCORE_MOD_MIN_CUDNN_VERSION = (9, 23) @@ -659,14 +661,14 @@ class FakeCudnn: __version__ = "1.22.0" monkeypatch.setattr( - tex_attention.transformer_engine_jax, + tex_cudnn_graph.transformer_engine_jax, "get_cudnn_frontend_version", lambda: 12200, ) assert tex_attention._check_cudnn_frontend_version_match(FakeCudnn) == 12200 monkeypatch.setattr( - tex_attention.transformer_engine_jax, + tex_cudnn_graph.transformer_engine_jax, "get_cudnn_frontend_version", lambda: 12100, ) @@ -674,6 +676,45 @@ class FakeCudnn: tex_attention._check_cudnn_frontend_version_match(FakeCudnn) +def test_fused_attn_score_mod_module_lambda_keys_do_not_collide(): + """Different anonymous functions cannot share a serialized attention graph.""" + first = lambda graph, score, tensors: score # noqa: E731 + second = lambda graph, score, tensors: graph.tanh(score) # noqa: E731 + first.__qualname__ = second.__qualname__ = "" + assert first.__qualname__ == second.__qualname__ == "" + assert tex_attention._score_mod_callback_cache_key(first) != ( + tex_attention._score_mod_callback_cache_key(second) + ) + + +@pytest.mark.parametrize("keyword", (False, True)) +def test_fused_attn_score_mod_config_separates_function_defaults(keyword): + """Defaults must participate in both the static config and serialized-graph cache key.""" + first = FunctionType(_default_mode_score_mod.__code__, globals(), argdefs=(False,)) + second = FunctionType(_default_mode_score_mod.__code__, globals(), argdefs=(False,)) + first.__kwdefaults__ = second.__kwdefaults__ = {"keyword_mode": False} + if keyword: + second.__kwdefaults__ = {"keyword_mode": True} + else: + first.__defaults__ = (False,) + second.__defaults__ = (True,) + configs = [ + make_fused_attn_score_mod_config( + callback, None, None, None, _CONFIG_TEST_SCALING_FACTOR, True + )[0] + for callback in (first, first, second) + ] + assert configs[0] == configs[1] + assert configs[0] != configs[2] + assert tex_attention._graph_cache_key("fwd", configs[0], ()) != ( + tex_attention._graph_cache_key("fwd", configs[2], ()) + ) + + +def _default_mode_score_mod(graph, score, _tensors, mode=False, *, keyword_mode=False): + return graph.tanh(input=score) if mode or keyword_mode else score + + def test_fused_attn_score_mod_config_stabilizes_bound_method_cache_keys(): """Bound score_mod methods with graph keys share stable cache keys.""" softcap_score_mod = _ScoreModSoftcap() diff --git a/tests/pytorch/attention/test_attention.py b/tests/pytorch/attention/test_attention.py index 86d46a56167..6f8eae3b5d6 100644 --- a/tests/pytorch/attention/test_attention.py +++ b/tests/pytorch/attention/test_attention.py @@ -10,7 +10,9 @@ import pathlib import copy from dataclasses import replace +from types import SimpleNamespace from typing import Any, Dict, List, NamedTuple, Tuple, Union +from unittest.mock import Mock from packaging.version import Version as PkgVersion import pytest @@ -46,11 +48,13 @@ import transformer_engine.pytorch.cpp_extensions as ext from transformer_engine.pytorch.cpp_extensions.fused_attn import ( FusedAttnBackend, + AttnMaskType, QKVFormat, QKVLayout, fused_attn_bwd, fused_attn_fwd, ) +from transformer_engine.pytorch.cpp_extensions.cudnn_fused_attn import get_fused_attn_backend from transformer_engine.pytorch.distributed import CudaRNGStatesTracker from transformer_engine.pytorch.module.base import TransformerEngineBaseModule from transformer_engine.pytorch.utils import ( @@ -58,7 +62,7 @@ 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 import transformer_engine_torch as tex from transformer_engine.pytorch.quantized_tensor import ( Quantizer, @@ -160,23 +164,21 @@ def test_fused_attn_backend_message(): ) # One of TE's rules is violated and the error message is surfaced - backend, message = tex.get_fused_attn_backend( + backend, message = get_fused_attn_backend( replace(baseline, bias_type=tex.NVTE_Bias_Type.NVTE_PRE_SCALE_BIAS) ) 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) + backend, message = 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}") 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) - ) + backend, message = 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 != "" @@ -203,20 +205,570 @@ def test_fused_attn_backend_message(): ) # No error message if supported; otherwise skip the test - backend, message = tex.get_fused_attn_backend(fp8_baseline) + backend, message = 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 == "" # All TE FP8 rules have cleared; now gets rejected by cuDNN's FP8 support check - backend, message = tex.get_fused_attn_backend( + backend, message = get_fused_attn_backend( replace(fp8_baseline, head_dim_qk=1024, head_dim_v=1024) ) assert backend == tex.NVTE_Fused_Attn_Backend.NVTE_No_Backend assert message != "" +@pytest.mark.parametrize("dtype", (torch.float16, torch.bfloat16)) +@pytest.mark.parametrize("qkv_format", ("bshd", "sbhd")) +def test_fused_attn_python_execution(monkeypatch, dtype, qkv_format): + """The public module must use Python cuDNN for selection, forward, and backward.""" + + def common_attention(*_args, **_kwargs): + raise AssertionError("PyTorch attention called the TE common attention API") + + for name in ("get_fused_attn_backend", "fused_attn_fwd", "fused_attn_bwd"): + monkeypatch.setattr(tex, name, common_attention, raising=False) + monkeypatch.setenv("NVTE_FLASH_ATTN", "0") + monkeypatch.setenv("NVTE_UNFUSED_ATTN", "0") + monkeypatch.setenv("NVTE_FUSED_ATTN", "1") + _attention_backends["backend_selection_requires_update"] = True + shape = (2, 128, 4, 64) if qkv_format == "bshd" else (128, 2, 4, 64) + q, k, v = [torch.randn(shape, device="cuda", dtype=dtype, requires_grad=True) for _ in range(3)] + block = DotProductAttention(4, 64, qkv_format=qkv_format, attn_mask_type="causal") + try: + actual = block(q, k, v).reshape(shape) + order = (0, 2, 1, 3) if qkv_format == "bshd" else (1, 2, 0, 3) + inverse = (0, 2, 1, 3) if qkv_format == "bshd" else (2, 0, 1, 3) + refs = [x.detach().float().permute(order).requires_grad_() for x in (q, k, v)] + reference = torch.nn.functional.scaled_dot_product_attention(*refs, is_causal=True) + grad = torch.randn_like(actual) + actual.backward(grad) + reference.backward(grad.float().permute(order)) + rtol, atol = (2e-2, 1e-2) if dtype == torch.bfloat16 else (3e-3, 2e-3) + torch.testing.assert_close(actual.float(), reference.permute(inverse), rtol=rtol, atol=atol) + for tensor, ref in zip((q, k, v), refs): + torch.testing.assert_close( + tensor.grad.float(), ref.grad.permute(inverse), rtol=rtol, atol=atol + ) + finally: + _attention_backends["backend_selection_requires_update"] = True + + +@pytest.mark.parametrize("layout", ("bs3hd", "bsh3d", "sb3hd", "sbh3d", "t3hd", "th3d")) +@pytest.mark.parametrize("backward", (False, True)) +def test_fused_attn_query_reuses_runtime_graph(layout, backward, monkeypatch): + """Packed storage and runtime scalar values must preserve support-query graph identity.""" + from transformer_engine.pytorch.cpp_extensions import cudnn_fused_attn as implementation + + batch, seqlen, heads, dim = 2, 128, 4, 64 + fmt = implementation._formats(layout)[0] + base_shape = implementation._shape((batch, heads, seqlen, dim), fmt, batch * seqlen) + axis = -2 if "h3d" in layout else -3 + packed_shape = list(base_shape) + packed_shape.insert(len(packed_shape) + axis + 1, 3) + q, k, v = torch.empty(packed_shape, device="cuda", dtype=torch.bfloat16).unbind(axis) + cu = torch.arange(batch + 1, device="cuda", dtype=torch.int32) * seqlen + params = dict( + q=q, + k=k, + v=v, + cu_seqlens_q=cu, + cu_seqlens_kv=cu, + cu_seqlens_q_padded=cu, + cu_seqlens_kv_padded=cu, + qkv_layout=layout, + max_seqlen_q=seqlen, + max_seqlen_kv=seqlen, + attn_mask_type="padding" if fmt == "thd" else "no_mask", + attn_bias_type="no_bias", + softmax_type="vanilla", + dropout=0.0, + attn_scale=0.125, + window_size=(-1, -1), + bottom_right_diagonal=True, + ) + config = FusedAttentionParams( + qkv_layout=QKVLayout[layout], + o_format=QKVFormat[fmt], + batch_size=batch, + num_attn_heads=heads, + num_gqa_groups=heads, + head_dim_qk=dim, + head_dim_v=dim, + max_seqlen_q=seqlen, + max_seqlen_kv=seqlen, + num_tokens_q=batch * seqlen if fmt == "thd" else 0, + num_tokens_kv=batch * seqlen if fmt == "thd" else 0, + attn_mask_type=AttnMaskType[params["attn_mask_type"]], + ) + query = implementation._query_graph(config, backward) + runtime = implementation._AttentionGraph(params, backward) + implementation._prepare_inputs(runtime) + out = torch.empty(base_shape, device="cuda", dtype=torch.bfloat16) + runtime.matrix("o", out, fmt, heads, runtime.sq, cu, output=not backward) + stats_shape = (batch * seqlen, heads, 1) if runtime.ragged_stats else (batch, heads, seqlen, 1) + runtime.stats("stats", torch.empty(stats_shape, device="cuda"), output=not backward) + if backward: + runtime.matrix("do", torch.empty_like(out), fmt, heads, runtime.sq, cu) + grads = torch.empty(packed_shape, device="cuda", dtype=torch.bfloat16).unbind(axis) + for name, grad in zip(("dq", "dk", "dv"), grads): + runtime.matrix(name, grad, fmt, heads, runtime.sq, cu, output=True) + assert query.description() == runtime.description() + + # Isolate the cache to demonstrate that a query never compiles, and that a + # later execution with another scalar value compiles exactly the cached plan. + monkeypatch.setattr(implementation, "_plans", {}) + builds, compiles = [], [] + + def build(_self): + entry = implementation._PlanEntry(object(), {name: name for name in runtime.specs}) + builds.append(entry) + return entry + + def compile_plan(_cudnn, graph, _run): + compiles.append(graph) + return 17 + + monkeypatch.setattr(implementation._AttentionGraph, "build", build) + monkeypatch.setattr(implementation, "compile_graph", compile_plan) + monkeypatch.setattr(implementation, "_execute_cudnn_graph", lambda *_args: None) + queried = query.entry() + assert queried.workspace_size is None and not compiles + runtime.scalar("attn_scale", torch.tensor([0.25]), by_value=True) + runtime.execute() + runtime.execute() + assert builds == [queried] + assert compiles == [queried.graph] + assert queried.workspace_size == 17 + + +@pytest.mark.parametrize("dtype", (DType.kFloat16, DType.kBFloat16)) +@pytest.mark.parametrize("capability", ((10, 0), (12, 0), (12, 1), (12, 9), (13, 0))) +@pytest.mark.parametrize("layout", ("bshd_bshd_bshd", "thd_bshd_bshd", "thd_thd_thd")) +def test_fused_attn_backward_token_bounds(monkeypatch, dtype, capability, layout): + """Bound backward workspace by each ragged input's bucketed token capacity.""" + import cudnn + from transformer_engine.pytorch.cpp_extensions import cudnn_fused_attn as implementation + + monkeypatch.setattr(torch.cuda, "get_device_capability", lambda *_args: capability) + config = FusedAttentionParams( + qkv_layout=QKVLayout[layout], + qkv_dtype=dtype, + o_dtype=dtype, + do_dtype=dtype, + dqkv_dtype=dtype, + batch_size=4, + num_attn_heads=4, + num_gqa_groups=4, + head_dim_qk=64, + head_dim_v=64, + max_seqlen_q=512, + max_seqlen_kv=1024, + num_tokens_q=1153, + num_tokens_kv=2305, + attn_mask_type=AttnMaskType["padding"], + ) + query = implementation._query_graph(config, True) + graph = Mock() + graph.sdpa_backward.return_value = (Mock(), Mock(), Mock()) + implementation.build_attention_graph(cudnn, graph, query.description()) + bounds = { + name: value + for name, value in graph.sdpa_backward.call_args.kwargs.items() + if name.startswith("max_total_seq_len_") + } + expected = {} + if capability[0] != 12 and layout.startswith("thd"): + expected["max_total_seq_len_q"] = 2048 + if layout == "thd_thd_thd": + expected["max_total_seq_len_kv"] = 4096 + assert bounds == expected + # cuDNN's graph layout and context-parallel LSE interpretation must agree. + from transformer_engine.pytorch.attention.dot_product_attention import context_parallel + + packed = context_parallel.uses_ragged_attention_graph(10 * capability[0] + capability[1]) + assert bool(query.specs["stats"].ragged_offset) == (packed and layout.startswith("thd")) + if capability[0] == 12: + assert query.sq == 512 + assert query.skv == 1024 + + +@pytest.mark.parametrize("capability", ((10, 0), (12, 0), (12, 1), (12, 9), (13, 0))) +@pytest.mark.parametrize("layout", ("bshd_bshd_bshd", "thd_bshd_bshd", "thd_thd_thd")) +def test_fused_attn_fake_stats_shape(capability, layout, monkeypatch): + """Fake saved statistics must match the runtime layout across the SM12x family.""" + from transformer_engine.pytorch.attention.dot_product_attention import backends + + monkeypatch.setattr(backends, "get_device_compute_capability", lambda: capability) + batch, seqlen, heads, dim, tokens = 3, 128, 4, 64, 224 + ragged = layout.startswith("thd") + shape = (tokens, heads, dim) if ragged else (batch, seqlen, heads, dim) + kv_shape = (tokens, heads, dim) if layout.endswith("thd") else (batch, seqlen, heads, dim) + q = torch.empty(shape, device="meta", dtype=torch.bfloat16) + kv = torch.empty(kv_shape, device="meta", dtype=torch.bfloat16) + args = SimpleNamespace( + q=q, + k=kv, + v=kv, + packed_qkv=None, + packed_kv=None, + qkv_layout=layout, + cu_seqlens_q=torch.empty(batch + 1, device="meta", dtype=torch.int32), + max_seqlen_q=seqlen, + return_max_logit=False, + attn_bias_type="no_bias", + softmax_type="vanilla", + fused_attention_backend=FusedAttnBackend.F16_arbitrary_seqlen, + ) + _, _, saved, _ = backends._fused_attn_forward_fake(args) + expected = (batch, heads, seqlen, 1) + if ragged and capability[0] != 12: + expected = (tokens, heads, 1) + assert saved[10].shape == expected + assert saved[10].dtype == torch.float32 + + +@pytest.mark.parametrize("dtype", (torch.float16, torch.bfloat16)) +@pytest.mark.parametrize("kv_format", ("thd", "bshd", "sbhd")) +def test_fused_attn_packed_backward(dtype, kv_format): + """Packed backward must match FP32 SDPA with unequal lengths and padded capacities.""" + if device_compute_capability < (9, 0): + pytest.skip("FP16/BF16 THD fused attention requires SM90 or later") + q_lengths, kv_lengths = (64, 128, 32), (96, 160, 64) + q_offsets = (0, 128, 384, 512) + kv_offsets = (0, 384, 768, 1152) if kv_format == "thd" else (0, 256, 512, 768) + cu_q = torch.tensor((0, 64, 192, 224), device="cuda", dtype=torch.int32) + cu_kv = torch.tensor((0, 96, 256, 320), device="cuda", dtype=torch.int32) + q = torch.randn(512, 4, 64, device="cuda", dtype=dtype) + kv_shape = {"thd": (1152, 4, 64), "bshd": (3, 256, 4, 64), "sbhd": (256, 3, 4, 64)} + k, v = [torch.randn(kv_shape[kv_format], device="cuda", dtype=dtype) for _ in range(2)] + layout = f"thd_{kv_format}_{kv_format}" + kwargs = dict( + max_seqlen_q=128, + max_seqlen_kv=256, + cu_seqlens_q=cu_q, + cu_seqlens_kv=cu_kv, + cu_seqlens_q_padded=torch.tensor(q_offsets, device="cuda", dtype=torch.int32), + cu_seqlens_kv_padded=torch.tensor(kv_offsets, device="cuda", dtype=torch.int32), + q=q, + k=k, + v=v, + fake_dtype=dtype, + fused_attention_backend=FusedAttnBackend.F16_arbitrary_seqlen, + qkv_layout=layout, + o_format="thd", + attn_mask_type="padding", + ) + out, aux = fused_attn_fwd(is_training=True, **kwargs) + grad = torch.randn_like(out) + dq, dk, dv, *_ = fused_attn_bwd( + o=out, d_o=grad, aux_ctx_tensors=aux, do_format="thd", dqkv_layout=layout, **kwargs + ) + refs = [tensor.float().requires_grad_() for tensor in (q, k, v)] + reference = torch.zeros_like(refs[0]) + for index, (q_len, kv_len) in enumerate(zip(q_lengths, kv_lengths)): + q_slice = slice(q_offsets[index], q_offsets[index] + q_len) + if kv_format == "thd": + kv_slice = slice(kv_offsets[index], kv_offsets[index] + kv_len) + kr, vr = (tensor[kv_slice] for tensor in refs[1:]) + elif kv_format == "bshd": + kr, vr = (tensor[index, :kv_len] for tensor in refs[1:]) + else: + kr, vr = (tensor[:kv_len, index] for tensor in refs[1:]) + result = torch.nn.functional.scaled_dot_product_attention( + refs[0][q_slice].transpose(0, 1), kr.transpose(0, 1), vr.transpose(0, 1) + ) + reference[q_slice] = result.transpose(0, 1) + reference.backward(grad.float()) + rtol, atol = (2e-2, 1e-2) if dtype == torch.bfloat16 else (3e-3, 2e-3) + torch.testing.assert_close(out.float(), reference, rtol=rtol, atol=atol) + for actual, ref in zip((dq, dk, dv), refs): + if actual.ndim == 3: + torch.testing.assert_close(actual.float(), ref.grad, rtol=rtol, atol=atol) + else: + # cuDNN may leave dense KV padding gradients unwritten. + for index, length in enumerate(kv_lengths): + active = ( + (index, slice(0, length)) if kv_format == "bshd" else (slice(0, length), index) + ) + torch.testing.assert_close( + actual[active].float(), ref.grad[active], rtol=rtol, atol=atol + ) + + +@pytest.mark.parametrize("error", (RuntimeError("device failure"), ImportError("frontend missing"))) +def test_fused_attn_query_propagates_setup_errors(error, monkeypatch): + """A failed setup must surface as an error instead of selecting a fallback backend.""" + from transformer_engine.pytorch.cpp_extensions import cudnn_fused_attn as implementation + + def fail(_self): + raise error + + monkeypatch.setattr(implementation._AttentionGraph, "build", fail) + monkeypatch.setattr(implementation, "_plans", {}) + config = FusedAttentionParams( + qkv_layout=QKVLayout["bshd_bshd_bshd"], + batch_size=2, + num_attn_heads=4, + num_gqa_groups=4, + head_dim_qk=64, + head_dim_v=64, + max_seqlen_q=128, + max_seqlen_kv=128, + ) + with pytest.raises(type(error), match=str(error)): + implementation.get_fused_attn_backend(config) + + +@pytest.mark.parametrize( + "changes,version,capability,reason", + ( + ({"qkv_dtype": DType.kFloat32}, 92700, (10, 0), "Unsupported QKV dtype"), + ( + {"cuda_graph": True, "max_seqlen_kv": 127}, + 91500, + (10, 0), + "Known cuDNN <= 9.15 issue with CUDA graph", + ), + ( + {"qkv_layout": QKVLayout["thd_thd_thd"], "num_tokens_q": 256, "num_tokens_kv": 256}, + 92700, + (10, 0), + "THD format requires PADDING", + ), + ( + { + "qkv_layout": QKVLayout["thd_thd_thd"], + "attn_mask_type": AttnMaskType["padding"], + "num_tokens_q": 256, + "num_tokens_kv": 256, + }, + 92700, + (8, 0), + "F16/BF16 fused attention with THD format requires sm90", + ), + ( + { + "qkv_layout": QKVLayout["thd_thd_thd"], + "attn_mask_type": AttnMaskType["padding"], + "softmax_type": tex.NVTE_Softmax_Type.NVTE_LEARNABLE_SOFTMAX, + "num_tokens_q": 256, + "num_tokens_kv": 256, + }, + 92500, + (10, 0), + "Known cuDNN < 9.26.0 issue with THD learnable softmax backward", + ), + ( + { + "qkv_dtype": DType.kFloat8E4M3, + "qkv_layout": QKVLayout["thd_thd_thd"], + "attn_mask_type": AttnMaskType["padding"], + "head_dim_qk": 192, + "num_tokens_q": 256, + "num_tokens_kv": 256, + }, + 92700, + (10, 0), + "FP8 THD attention supports head dimensions up to 128", + ), + ( + { + "qkv_dtype": DType.kFloat8E4M3, + "scaling_mode": tex.NVTEScalingMode.NVTE_MXFP8_1D_SCALING, + }, + 92000, + (10, 0), + "MXFP8 fused attention requires cuDNN 9.21.0", + ), + ( + {"qkv_layout": QKVLayout["paged_kv_bshd_bshd_bshd"]}, + 92700, + (10, 0), + "Paged KV requires PADDING", + ), + ), +) +def test_fused_attn_python_capability_exclusions(monkeypatch, changes, version, capability, reason): + """Preserve capability and tracked-defect exclusions before asking cuDNN to build.""" + import cudnn + from transformer_engine.pytorch.cpp_extensions.cudnn_fused_attn import _AttentionGraph + + monkeypatch.setattr(cudnn, "backend_version", lambda: version) + monkeypatch.setattr(torch.cuda, "get_device_capability", lambda *_args: capability) + + def build(_self): + raise AssertionError("Excluded configuration reached cuDNN graph construction") + + monkeypatch.setattr(_AttentionGraph, "build", build) + baseline = FusedAttentionParams( + qkv_layout=QKVLayout["bshd_bshd_bshd"], + batch_size=2, + num_attn_heads=4, + num_gqa_groups=4, + head_dim_qk=64, + head_dim_v=64, + max_seqlen_q=128, + max_seqlen_kv=128, + ) + backend, message = get_fused_attn_backend(replace(baseline, **changes)) + assert backend == FusedAttnBackend.No_Backend + assert message.startswith(reason) + + +@pytest.mark.skipif(not fp8_attn_available, reason=reason_for_no_fp8_attn) +def test_fused_attn_python_dropout_cuda_graph(): + """Dropout must reproduce seeded eager results and advance its mask on graph replay.""" + from transformer_engine.pytorch.tensor.float8_tensor import Float8Quantizer + + layout = "bshd_bshd_bshd" + config = FusedAttentionParams( + qkv_layout=QKVLayout[layout], + dqkv_layout=QKVLayout[layout], + o_format=QKVFormat["bshd"], + do_format=QKVFormat["bshd"], + qkv_dtype=DType.kFloat8E4M3, + o_dtype=DType.kFloat8E4M3, + do_dtype=DType.kFloat8E5M2, + dqkv_dtype=DType.kFloat8E5M2, + batch_size=2, + num_attn_heads=4, + num_gqa_groups=4, + head_dim_qk=64, + head_dim_v=64, + max_seqlen_q=128, + max_seqlen_kv=128, + dropout=0.125, + ) + backend, message = get_fused_attn_backend(config) + if backend == FusedAttnBackend.No_Backend: + pytest.skip(f"cuDNN FP8 dropout is unavailable: {message}") + + def quantizer(dtype, scale=1.0): + return Float8Quantizer( + torch.full((1,), scale, device="cuda"), + torch.zeros(1, device="cuda"), + dtype, + columnwise=False, + ) + + source = torch.ones(2, 128, 4, 64, dtype=torch.float16, device="cuda") + q, k, v = [quantizer(DType.kFloat8E4M3)(source) for _ in range(3)] + d_o = quantizer(DType.kFloat8E5M2)(source) + s_quantizer = quantizer(DType.kFloat8E4M3, 128.0) + dp_quantizer = quantizer(DType.kFloat8E5M2) + o_quantizer = quantizer(DType.kFloat8E4M3) + dqkv_quantizer = quantizer(DType.kFloat8E5M2) + cu = torch.tensor([0, 128, 256], dtype=torch.int32, device="cuda") + kwargs = { + "qkv_layout": layout, + "o_format": "bshd", + "attn_mask_type": "no_mask", + "dropout": config.dropout, + "s_quantizer": s_quantizer, + } + + def run(): + out, aux = fused_attn_fwd( + True, + 128, + 128, + cu, + cu, + q, + k, + v, + source.dtype, + backend, + o_quantizer=o_quantizer, + **kwargs, + ) + dq, dk, dv, *_ = fused_attn_bwd( + 128, + 128, + cu, + cu, + q, + k, + v, + out, + d_o, + source.dtype, + aux, + backend, + do_format="bshd", + dqkv_layout=layout, + dp_quantizer=dp_quantizer, + dqkv_quantizer=dqkv_quantizer, + **kwargs, + ) + return tuple(tensor.dequantize() for tensor in (out, dq, dk, dv)) + + torch.cuda.manual_seed(1234) + eager = run() + # With uniform scores and values, inverted dropout has expected output and dV equal to one. + torch.testing.assert_close( + eager[0].float().mean(), torch.ones((), device="cuda"), atol=0.05, rtol=0 + ) + torch.testing.assert_close( + eager[3].float().mean(), torch.ones((), device="cuda"), atol=0.05, rtol=0 + ) + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(stream): + for _ in range(3): + run() + torch.cuda.current_stream().wait_stream(stream) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph, stream=stream): + captured = run() + torch.cuda.manual_seed(1234) + graph.replay() + first = [tensor.clone() for tensor in captured] + for actual, expected in zip(first, eager): + torch.testing.assert_close(actual, expected, atol=1e-2, rtol=1e-2) + graph.replay() + assert not torch.equal(captured[0], first[0]) + torch.cuda.manual_seed(1234) + graph.replay() + for actual, expected in zip(captured, first): + torch.testing.assert_close(actual, expected, atol=1e-2, rtol=1e-2) + + +@pytest.mark.parametrize("custom_generator", (False, True)) +def test_fused_attn_rng_state_cuda_graph(custom_generator): + """Philox seeds and offsets must follow the selected generator through capture and replay.""" + if custom_generator and not hasattr(torch.cuda.CUDAGraph, "register_generator_state"): + pytest.skip( + "This PyTorch version cannot register custom CUDA generators for graph capture." + ) + generator = ( + torch.Generator(device="cuda") + if custom_generator + else torch.cuda.default_generators[torch.cuda.current_device()] + ) + rng_gen = generator if custom_generator else None + generator.manual_seed(1234) + first = tex.fused_attn_rng_state(rng_gen, 16) + second = tex.fused_attn_rng_state(rng_gen, 16) + assert first.tolist() == [1234, 0] + assert second.tolist() == [1234, 16] + graph = torch.cuda.CUDAGraph() + if custom_generator: + graph.register_generator_state(generator) + with torch.cuda.graph(graph): + captured = tex.fused_attn_rng_state(rng_gen, 16) + generator.manual_seed(1234) + graph.replay() + torch.testing.assert_close(captured, first) + graph.replay() + torch.testing.assert_close(captured, second) + generator.manual_seed(1234) + graph.replay() + torch.testing.assert_close(captured, first) + + @pytest.mark.parametrize( "cu_seqlens,cu_seqlens_padded,expected", [ @@ -3919,7 +4471,7 @@ def forward( @staticmethod def backward(ctx, grad_output: torch.Tensor) -> Tuple[Union[torch.Tensor, None], ...]: with torch.cuda.nvtx.range("_DPA"): - (q, k, v, inp_fp8, qkv_weight_fp8, out) = restore_from_func_ctx(ctx) + q, k, v, inp_fp8, qkv_weight_fp8, out = restore_from_func_ctx(ctx) proj_dgrad = ctx.dO_quantizer(grad_output) diff --git a/tests/pytorch/attention/test_attention_backend_selection.py b/tests/pytorch/attention/test_attention_backend_selection.py index bd587416a79..1721e41adca 100644 --- a/tests/pytorch/attention/test_attention_backend_selection.py +++ b/tests/pytorch/attention/test_attention_backend_selection.py @@ -5,10 +5,12 @@ """Focused FA4 eligibility and backend-cache regressions.""" from unittest.mock import Mock +from types import SimpleNamespace import pytest import torch from packaging.version import Version as PkgVersion +from transformer_engine.common import recipe from transformer_engine.pytorch import DotProductAttention from transformer_engine.pytorch.attention.dot_product_attention import ( @@ -84,6 +86,45 @@ def test_fa4_deterministic_backend_selection( assert available == [expect_flash, False, unfused] +@pytest.mark.parametrize("capability", ((12, 0), (12, 1), (12, 9))) +@pytest.mark.parametrize("case", ("fp8", "kv_cache", "mla", "thd_version", "thd_packed")) +def test_sm12x_fused_backend_restrictions(monkeypatch, capability, case): + """SM121 and other SM12x variants must retain the SM120 backend restrictions.""" + monkeypatch.setenv("NVTE_FUSED_ATTN", "1") + monkeypatch.setenv("NVTE_FLASH_ATTN", "0") + monkeypatch.setattr(dpa_utils, "get_device_compute_capability", lambda _device: capability) + version = (9, 18, 0) if case == "thd_version" else (9, 27, 0) + monkeypatch.setattr(dpa_utils, "get_cudnn_version", lambda: version) + probe = Mock(return_value=(dpa_utils.FusedAttnBackend.F16_arbitrary_seqlen.value, "")) + monkeypatch.setattr(dpa_utils, "_get_fused_attn_backend", probe) + options = dict( + device=torch.device("cuda:0"), + qkv_dtype=torch.bfloat16, + qkv_layout="bshd_bshd_bshd", + num_heads=4, + num_gqa_groups=4, + head_dim_qk=64, + head_dim_v=64, + deterministic=False, + ) + if case == "fp8": + options.update(fp8=True, fp8_meta={"recipe": recipe.DelayedScaling(fp8_dpa=True)}) + elif case == "kv_cache": + options.update(inference_params=SimpleNamespace(is_paged=False), is_training=False) + elif case == "mla": + options.update(head_dim_qk=192, head_dim_v=128) + else: + options.update( + qkv_layout="t3hd" if case == "thd_packed" else "thd_thd_thd", + attn_mask_type="padding", + num_tokens_q=128, + num_tokens_kv=128, + ) + result = dpa_utils.get_attention_backend(dpa_utils.AttentionParams(**options)) + assert not result[2] + probe.assert_not_called() + + @pytest.mark.parametrize( "version,backward,expect_flash", [ diff --git a/tests/pytorch/attention/test_flex_attention.py b/tests/pytorch/attention/test_flex_attention.py index beed4069917..c40aaf2d0af 100644 --- a/tests/pytorch/attention/test_flex_attention.py +++ b/tests/pytorch/attention/test_flex_attention.py @@ -5,6 +5,8 @@ import os import sys import pathlib +from types import FunctionType +from unittest.mock import Mock import pytest import torch @@ -156,6 +158,109 @@ def _score_mod_cache_cpu_inputs(): return q, k, v, o, stats +def _score_mod_with_defaults(graph, score, _tensors, mode=False, *, options=None): + if options is not None: + mode = options["modes"][0] + return graph.tanh(input=score) if mode else score + + +def _clone_score_mod(default=False, options=None): + callback = FunctionType(_score_mod_with_defaults.__code__, globals(), argdefs=(default,)) + callback.__kwdefaults__ = {"options": options} + return callback + + +@pytest.mark.parametrize("direction", ("fwd", "bwd")) +@pytest.mark.parametrize("difference", ("positional", "keyword", "code", "type")) +def test_score_mod_cache_separates_function_defaults(monkeypatch, direction, difference): + """Different code or defaults must build separate forward and backward graphs.""" + first = _clone_score_mod() + second = _clone_score_mod(True) + if difference == "keyword": + second = _clone_score_mod(options={"modes": [True]}) + elif difference == "code": + second = FunctionType(_score_mod_identity_bprop.__code__, globals()) + second.__qualname__ = first.__qualname__ + elif difference == "type": + second = _clone_score_mod(0) + q, k, v, out, stats = _score_mod_cache_cpu_inputs() + monkeypatch.setattr(flex_attention, "_cudnn_score_mod_graph_cache", {}) + build = Mock(side_effect=lambda *_args: object()) + monkeypatch.setattr(flex_attention, f"_build_cudnn_score_mod_{direction}_graph", build) + + def get(callback): + if direction == "fwd": + return flex_attention._get_cudnn_score_mod_fwd_graph( + True, q, k, v, "bshd", "bshd", 1.0, callback, {}, out, stats + ) + return flex_attention._get_cudnn_score_mod_bwd_graph( + q, + k, + v, + out, + out, + stats, + "bshd", + "bshd", + 1.0, + _score_mod_identity_bprop, + callback, + {}, + {}, + False, + ) + + entry = get(first) + assert get(first) is entry + other = get(second) + assert other is not entry + assert get(second) is other + assert build.call_count == 2 + + # Keys snapshot nested defaults rather than retaining mutable references. + first.__kwdefaults__["options"] = {"modes": [False]} + changed = get(first) + assert changed is not entry + first.__kwdefaults__["options"]["modes"].append(True) + assert get(first) is not changed + assert build.call_count == 4 + + +@pytest.mark.parametrize("kind", ("tensor", "opaque", "cycle")) +def test_score_mod_cache_unsafe_defaults_are_uncached(kind): + """Opaque state, runtime tensors and recursive defaults cannot imply a stable graph.""" + if kind == "tensor": + default = torch.tensor(1.0) + elif kind == "opaque": + default = object() + else: + default = [] + default.append(default) + callback = _clone_score_mod(default) + key = flex_attention._score_mod_callback_cache_key(callback) + assert key is flex_attention._SCORE_MOD_UNCACHEABLE + + # An explicit structural key remains the caller's opt-in caching contract. + callback.score_mod_graph_cache_key = lambda: ("explicit-mode", True) + assert flex_attention._score_mod_callback_cache_key(callback) is not key + + +def test_score_mod_cache_snapshots_positional_defaults(): + """Changing positional defaults invalidates the old key while unchanged NaNs stay stable.""" + callback = _clone_score_mod(False) + before = flex_attention._score_mod_callback_cache_key(callback) + callback.__defaults__ = (True,) + assert flex_attention._score_mod_callback_cache_key(callback) != before + callback.__defaults__ = (float("nan"),) + assert flex_attention._score_mod_callback_cache_key(callback) == ( + flex_attention._score_mod_callback_cache_key(callback) + ) + callback.__defaults__ = (0.0,) + positive_zero = flex_attention._score_mod_callback_cache_key(callback) + callback.__defaults__ = (-0.0,) + assert flex_attention._score_mod_callback_cache_key(callback) != positive_zero + + def test_score_mod_cache_bound_method_requires_explicit_key(): """Unkeyed bound methods should be uncached instead of keyed by object id.""" diff --git a/tests/pytorch/test_torch_compile.py b/tests/pytorch/test_torch_compile.py index 472917d56f0..cf0195e910d 100644 --- a/tests/pytorch/test_torch_compile.py +++ b/tests/pytorch/test_torch_compile.py @@ -1316,9 +1316,10 @@ 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 get_fused_attn_backend result must drive the selection.""" from transformer_engine.pytorch.attention.dot_product_attention import utils as dpa_utils + from transformer_engine.pytorch.cpp_extensions import cudnn_fused_attn def fn(x, params): ( @@ -1396,9 +1397,9 @@ def fn(x, params): # guard on the wrapped function). monkeypatch.setenv("NVTE_FLASH_ATTN", "0") monkeypatch.setattr( - dpa_utils.tex, + cudnn_fused_attn, "get_fused_attn_backend", - lambda *args: (tex.NVTE_Fused_Attn_Backend.NVTE_No_Backend, "disabled by test"), + lambda _config: (dpa_utils.FusedAttnBackend.No_Backend, "disabled by test"), ) def fn_no_backend(x, params): diff --git a/transformer_engine/common/cudnn/__init__.py b/transformer_engine/common/cudnn/__init__.py new file mode 100644 index 00000000000..7501e1ed649 --- /dev/null +++ b/transformer_engine/common/cudnn/__init__.py @@ -0,0 +1,5 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""Framework-independent helpers for the cuDNN Python graph API.""" diff --git a/transformer_engine/common/cudnn/attention.py b/transformer_engine/common/cudnn/attention.py new file mode 100644 index 00000000000..e14f8e377cf --- /dev/null +++ b/transformer_engine/common/cudnn/attention.py @@ -0,0 +1,220 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""Construct SDPA graphs from immutable descriptors, without framework tensors.""" + +from dataclasses import dataclass + +from .graph import TensorSpec, apply_tensor_spec, create_tensor + + +@dataclass(frozen=True) +class AttentionOptions: + """Normalized attributes that affect graph construction and support.""" + + backward: bool = False + fp8: bool = False + mxfp8: bool = False + attn_scale: str | float = "attn_scale" + bottom_right: bool = False + left_bound: int | None = None + right_bound: int | None = None + padding: bool = False + alibi: bool = False + dropout: float = 0.0 + deterministic: bool = False + generate_stats: bool = True + paged_max_seqlen_kv: int = 0 + max_total_seqlen_q: int = 0 + max_total_seqlen_kv: int = 0 + + +@dataclass(frozen=True) +class AttentionGraphSpec: + """A cacheable graph description; bindings and scalar values live in adapters.""" + + options: AttentionOptions + tensors: tuple + + @classmethod + def from_dict(cls, options, tensors): + """Canonicalize tensor order for both UIDs and graph cache keys.""" + return cls(options, tuple(sorted(tensors.items()))) + + +@dataclass +class BuiltAttentionGraph: + """A live graph and its tensor handles, without any runtime buffers.""" + + graph: object + tensors: dict + uids: dict + + +def build_attention_graph(cudnn, graph, spec, *, score_mod=None, score_mod_bprop=None): + """Build ordinary, FP8, MXFP8, or score-modified SDPA using logical bindings.""" + specs = dict(spec.tensors) + tensors = {} + uids = {name: index + 1 for index, (name, _) in enumerate(spec.tensors)} + + def tensor(name): + if name not in tensors: + desc = create_tensor(cudnn, graph, name, specs[name], uids[name]) + if specs[name].ragged_offset: + desc.set_ragged_offset(tensor(specs[name].ragged_offset)) + tensors[name] = desc + return tensors[name] + + for name, desc in spec.tensors: + if not desc.output: + tensor(name) + + options = spec.options + kwargs = { + "attn_scale": ( + tensor(options.attn_scale) + if isinstance(options.attn_scale, str) + else options.attn_scale + ), + "diagonal_alignment": ( + cudnn.diagonal_alignment.BOTTOM_RIGHT + if options.bottom_right + else cudnn.diagonal_alignment.TOP_LEFT + ), + } + bound_names = ( + ("left_bound", "right_bound") + if options.fp8 and (not options.mxfp8 or options.backward) + else ("diagonal_band_left_bound", "diagonal_band_right_bound") + ) + for name, bound in zip(bound_names, (options.left_bound, options.right_bound)): + if bound is not None: + kwargs[name] = bound + if options.padding: + kwargs.update(use_padding_mask=True, seq_len_q=tensor("seq_q"), seq_len_kv=tensor("seq_kv")) + if options.alibi: + kwargs["use_alibi_mask"] = True + if "bias" in specs: + kwargs["bias"] = tensor("bias") + if options.backward and "dbias" in specs: + kwargs["dBias"] = tensor("dbias") + if options.dropout: + kwargs["dropout"] = (options.dropout, tensor("seed"), tensor("offset")) + if "sink" in specs: + kwargs["sink_token"] = tensor("sink") + if options.backward: + kwargs["dSink_token"] = tensor("dsink") + if score_mod is not None: + kwargs["score_mod"] = score_mod + if score_mod_bprop is not None: + kwargs["score_mod_bprop"] = score_mod_bprop + if options.backward: + kwargs["use_deterministic_algorithm"] = options.deterministic + if options.max_total_seqlen_q: + kwargs["max_total_seq_len_q"] = options.max_total_seqlen_q + if options.max_total_seqlen_kv: + kwargs["max_total_seq_len_kv"] = options.max_total_seqlen_kv + else: + kwargs["generate_stats"] = options.generate_stats + if "max" in specs: + kwargs["score_max"] = tensor("max") + if options.paged_max_seqlen_kv: + kwargs.update( + paged_attention_k_table=tensor("page_table_k"), + paged_attention_v_table=tensor("page_table_v"), + paged_attention_max_seq_len_kv=options.paged_max_seqlen_kv, + ) + + q, k, v = (tensor(name) for name in ("q", "k", "v")) + if not options.fp8: + if options.backward: + outputs = graph.sdpa_backward( + q, k, v, tensor("o"), tensor("do"), tensor("stats"), **kwargs + ) + names = ("dq", "dk", "dv") + else: + outputs = graph.sdpa(q, k, v, **kwargs) + names = ("o", "stats") + elif not options.backward: + scales = [tensor(name) for name in ("descale_q", "descale_k", "descale_v")] + if options.mxfp8: + outputs = graph.sdpa_mxfp8(q, k, v, *scales, **kwargs) + names = ("o", "stats", "amax_o") + else: + outputs = graph.sdpa_fp8( + q, + k, + v, + *scales, + *(tensor(name) for name in ("descale_s", "scale_s", "scale_o")), + **kwargs + ) + names = ("o", "stats", "amax_s", "amax_o") + elif options.mxfp8: + outputs = graph.sdpa_mxfp8_backward( + q, + tensor("q_t"), + k, + tensor("k_t"), + v, + tensor("o"), + tensor("do_f16"), + tensor("do"), + tensor("do_t"), + tensor("stats"), + *( + tensor(name) + for name in ( + "descale_q", + "descale_q_t", + "descale_k", + "descale_k_t", + "descale_v", + "descale_do", + "descale_do_t", + ) + ), + **kwargs + ) + names = ("dq", "dk", "dv", "amax_dq", "amax_dk", "amax_dv") + else: + outputs = graph.sdpa_fp8_backward( + q, + k, + v, + tensor("o"), + tensor("do"), + tensor("stats"), + *( + tensor(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", + ) + ), + **kwargs + ) + names = ("dq", "dk", "dv", "amax_dq", "amax_dk", "amax_dv", "amax_dp") + + for name, desc in zip(names, outputs): + if name not in specs: + # Unbound MXFP8 amax outputs stay virtual. Their descriptors are still required. + if desc is not None: + apply_tensor_spec(cudnn, desc, TensorSpec((1, 1, 1, 1), (1, 1, 1, 1), "FLOAT")) + continue + apply_tensor_spec(cudnn, desc, specs[name], uids[name]) + if specs[name].ragged_offset: + desc.set_ragged_offset(tensor(specs[name].ragged_offset)) + tensors[name] = desc + return BuiltAttentionGraph(graph, tensors, uids) diff --git a/transformer_engine/common/cudnn/graph.py b/transformer_engine/common/cudnn/graph.py new file mode 100644 index 00000000000..6393a476924 --- /dev/null +++ b/transformer_engine/common/cudnn/graph.py @@ -0,0 +1,254 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""Data-free tensor descriptions and explicit cuDNN graph build stages.""" + +import atexit +from collections import Counter +from dataclasses import dataclass +import importlib +import math +import os +import sys +import threading +import time + + +class UnsupportedGraph(ValueError): + """A configuration rejected by cuDNN, rather than an execution/setup failure.""" + + +def uses_ragged_attention_graph(sm_arch): + """Use packed statistics on SM90+ except the SM12x family (encoded as 120-129).""" + return sm_arch >= 90 and not 120 <= sm_arch < 130 + + +def import_cudnn(): + """Import the optional graph frontend without importing either framework.""" + 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 + + +def create_graph(cudnn, dtype, handle=None): + """Create a graph with FP32 intermediate values and accumulation.""" + kwargs = {} if handle is None else {"handle": handle} + return cudnn.pygraph( + io_data_type=getattr(cudnn.data_type, dtype) if isinstance(dtype, str) else dtype, + intermediate_data_type=cudnn.data_type.FLOAT, + compute_data_type=cudnn.data_type.FLOAT, + **kwargs, + ) + + +def check_graph_support(cudnn, graph, run_stage=None): + """Validate and select supported plans, without compiling their kernels.""" + run = run_stage or (lambda _name, function: function()) + try: + run("validate", graph.validate) + run("build_operation_graph", graph.build_operation_graph) + run( + "create_execution_plans", + lambda: graph.create_execution_plans([cudnn.heur_mode.A, cudnn.heur_mode.FALLBACK]), + ) + run("check_support", graph.check_support) + except cudnn.cudnnGraphNotSupportedError as exc: + raise UnsupportedGraph(f"cuDNN SDPA graph is not supported: {exc}") from exc + + +def compile_graph(cudnn, graph, run_stage=None): + """Compile the selected plan and return its nonempty workspace size.""" + run = run_stage or (lambda _name, function: function()) + run("build_plans", lambda: graph.build_plans(cudnn.build_plan_policy.HEURISTICS_CHOICE)) + return max(int(graph.get_workspace_size()), 1) + + +@dataclass(frozen=True) +class TensorSpec: + """A logical graph tensor, independent of its runtime buffer and framework.""" + + dim: tuple + stride: tuple + dtype: str + output: bool = False + ragged_offset: str = "" + reordered: bool = False + by_value: bool = False + + +def apply_tensor_spec(cudnn, tensor, spec, uid=None): + """Apply the same descriptor to an input or an operation-produced output.""" + tensor.set_dim(spec.dim).set_stride(spec.stride).set_data_type( + getattr(cudnn.data_type, spec.dtype) + ) + if uid is not None: + tensor.set_uid(uid) + if spec.output: + tensor.set_output(True) + return tensor + + +def create_tensor(cudnn, graph, name, spec, uid=None): + """Create an input or an explicitly supplied output with a stable UID.""" + kwargs = {} if uid is None else {"uid": uid} + tensor = graph.tensor( + name=name, + dim=spec.dim, + stride=spec.stride, + data_type=getattr(cudnn.data_type, spec.dtype), + is_pass_by_value=spec.by_value, + **kwargs, + ) + if spec.reordered: + tensor.set_reordering_type(cudnn.tensor_reordering.F8_128x4) + return apply_tensor_spec(cudnn, tensor, spec, uid) + + +def row_major_stride(shape): + """Return contiguous storage strides without consulting framework objects.""" + return tuple(math.prod(shape[index + 1 :]) for index in range(len(shape))) + + +def round_up(value, multiple): + """Round an integer up to a multiple.""" + return (value + multiple - 1) // multiple * multiple + + +def bucket(value, minimum, maximum): + """Bucket positive capacities by powers of two, then multiples of maximum.""" + if value == 0: + return 0 + if value <= maximum: + return max(minimum, 1 << (value - 1).bit_length()) + return round_up(value, maximum) + + +def attention_stride(dim, fmt): + """Describe contiguous attention storage in logical BHSD order.""" + b, h, s, d = dim + if fmt in ("bshd", "thd"): + return (s * h * d, d, h * d, 1) + if fmt == "sbhd": + return (h * d, d, b * h * d, 1) + if fmt == "bhsd": + return (h * s * d, s * d, d, 1) + raise ValueError(f"Unsupported attention format: {fmt}") + + +def attention_matrix_spec( + dim, fmt, dtype, *, storage_stride=None, packed="", output=False, ragged_offset="" +): + """Describe either actual storage strides or a hypothetical support-query tensor.""" + if storage_stride is not None and fmt == "thd": + token, head, last = storage_stride + stride = (dim[2] * token, head, token, last) + elif storage_stride is not None and len(storage_stride) == 4: + order = {"bshd": (0, 2, 1, 3), "sbhd": (1, 2, 0, 3), "bhsd": (0, 1, 2, 3)}[fmt] + stride = tuple(storage_stride[index] for index in order) + else: + stride = list(attention_stride(dim, fmt)) + if "3" in packed or "2" in packed: + count = 3 if "3" in packed else 2 + stride[0] *= count + stride[2] *= count + if "h3d" in packed or "h2d" in packed: + stride[1] *= count + stride = tuple(stride) + return TensorSpec(tuple(dim), stride, dtype, output=output, ragged_offset=ragged_offset) + + +def attention_stats_spec(batch, heads, seqlen, *, ragged=False, output=False): + """Describe dense BHS1 or packed TH1 softmax statistics.""" + stride = (heads * seqlen, 1, heads, 1) if ragged else (heads * seqlen, seqlen, 1, 1) + return TensorSpec( + (batch, heads, seqlen, 1), + stride, + "FLOAT", + output=output, + ragged_offset="stats_offset" if ragged else "", + ) + + +def mxfp8_scale_spec(dim, fmt, columnwise=False): + """Describe the reordered block-scale storage used by MXFP8 SDPA.""" + b, h, s, d = dim + shape = ( + (b, h, round_up((s + 31) // 32, 4), round_up(d, 128)) + if columnwise + else (b, h, round_up(s, 128), round_up((d + 31) // 32, 4)) + ) + return TensorSpec(shape, attention_stride(shape, fmt), "FP8_E8M0", reordered=True) + + +class CacheDiagnostics: + """Optional cache counters and build timings shared by Python frontends.""" + + def __init__(self, framework): + self.framework = framework + self.events = Counter() + self.times = Counter() + atexit.register(self.summary) + + @staticmethod + def level(): + """Resolve the existing debug level and rank filter.""" + setting = os.getenv("NVTE_FUSED_ATTN_CACHE_DEBUG", "0") + level, _, ranks = setting.partition(":") + rank = next( + ( + os.environ[n] + for n in ("RANK", "LOCAL_RANK", "OMPI_COMM_WORLD_RANK", "SLURM_PROCID") + if n in os.environ + ), + None, + ) + enabled = rank is None or ranks == "all" or rank in (ranks or "0").split(",") + return int(level or 0) if enabled else 0 + + def event(self, backend, direction, device, event, key=None): + """Record an event without retaining runtime tensors.""" + level = self.level() + if not level: + return + self.events[(backend, direction, event)] += 1 + if level >= 2: + print( + f"[FUSED-ATTN-CACHE] tid={threading.get_ident()} dev={device} | " + f"{backend} {direction} {event}" + + (f" | {key!r}" if key is not None else ""), + file=sys.stderr, + flush=True, + ) + + def run(self, backend, direction, stage, function): + """Time a build stage only when diagnostics are enabled.""" + if not self.level(): + return function() + start = time.perf_counter() + try: + return function() + finally: + self.times[(backend, direction, stage, "seconds")] += time.perf_counter() - start + self.times[(backend, direction, stage, "calls")] += 1 + + def summary(self): + """Print this frontend's collected counters and mean build times.""" + if not self.events: + return + print(f"[FUSED-ATTN-CACHE] {self.framework} Python frontend summary", file=sys.stderr) + for (backend, direction, event), count in sorted(self.events.items()): + print(f" {backend} {direction} {event}={count}", file=sys.stderr) + for backend, direction, stage, kind in sorted(self.times): + if kind == "calls": + calls = self.times[(backend, direction, stage, kind)] + seconds = self.times[(backend, direction, stage, "seconds")] + print( + f" {backend} {direction} {stage}: {calls} calls, " + f"{1000 * seconds / calls:.3f} ms average", + file=sys.stderr, + ) diff --git a/transformer_engine/common/cudnn/score_mod.py b/transformer_engine/common/cudnn/score_mod.py new file mode 100644 index 00000000000..e48e35626cc --- /dev/null +++ b/transformer_engine/common/cudnn/score_mod.py @@ -0,0 +1,125 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""Structural callback keys shared by score-modified attention frontends.""" + +from collections.abc import Mapping +import inspect +import struct + + +class UncacheableKey: + """Unique static identity for callbacks whose topology may change.""" + + +def freeze_cache_key(value, is_tensor): + """Freeze nested structural keys while rejecting framework runtime tensors.""" + if is_tensor(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): + return tuple( + sorted( + ( + (freeze_cache_key(key, is_tensor), freeze_cache_key(val, is_tensor)) + for key, val in value.items() + ), + key=repr, + ) + ) + if isinstance(value, (list, tuple)): + return tuple(freeze_cache_key(item, is_tensor) for item in value) + if isinstance(value, (set, frozenset)): + return tuple(sorted((freeze_cache_key(item, is_tensor) for item in value), 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(owner, is_tensor): + """Read an optional structural key from a callback's owner.""" + key = getattr(owner, "score_mod_graph_cache_key", None) + if key is None: + return None + return freeze_cache_key(key() if callable(key) else key, is_tensor) + + +def _freeze_function_default(value, is_tensor): + """Snapshot builtin defaults without trusting opaque objects' identity hashes.""" + value_type = type(value) + if is_tensor(value): + raise TypeError("Runtime tensors cannot be implicit callback cache keys.") + if value_type in (type(None), bool, int, str, bytes): + return value_type, value + if value_type is float: + return value_type, struct.pack("!d", value) + if value_type is complex: + return value_type, struct.pack("!dd", value.real, value.imag) + if value_type is dict: + return value_type, tuple( + (_freeze_function_default(key, is_tensor), _freeze_function_default(item, is_tensor)) + for key, item in value.items() + ) + if value_type in (list, tuple, set, frozenset): + return value_type, tuple(_freeze_function_default(item, is_tensor) for item in value) + raise TypeError("Opaque callback defaults require an explicit structural cache key.") + + +def callback_cache_key(callback, is_tensor, uncacheable): + """Cache stateless functions by code/defaults and explicitly keyed callback owners.""" + if callback is None: + return None + owner = getattr(callback, "__self__", None) + function = getattr(callback, "__func__", None) + if owner is not None and function is not None: + key = explicit_cache_key(owner, is_tensor) + if key is None: + return uncacheable() + return ("bound_method", type(owner), function.__module__, function.__qualname__, key) + key = explicit_cache_key(callback, is_tensor) + if key is not None: + return ( + "callable", + type(callback), + getattr(callback, "__module__", None), + getattr(callback, "__qualname__", None), + key, + ) + if ( + inspect.isfunction(callback) + and callback.__closure__ is None + and "" not in callback.__qualname__ + ): + try: + defaults = _freeze_function_default(callback.__defaults__, is_tensor) + kwdefaults = _freeze_function_default(callback.__kwdefaults__, is_tensor) + except (TypeError, RecursionError): + return uncacheable() + return ( + "function", + callback.__module__, + callback.__qualname__, + callback.__code__, + defaults, + kwdefaults, + ) + return uncacheable() + + +def wrap_score_mod(callback, graph_tensors): + """Adapt TE's three-argument callback to cuDNN's two-argument callback.""" + if callback is None: + return None + + def wrapped(graph, scores): + return callback(graph, scores, graph_tensors) + + return wrapped diff --git a/transformer_engine/common/fused_attn/config_and_params.cpp b/transformer_engine/common/fused_attn/config_and_params.cpp index b12dbf30f61..f27859791e8 100644 --- a/transformer_engine/common/fused_attn/config_and_params.cpp +++ b/transformer_engine/common/fused_attn/config_and_params.cpp @@ -106,7 +106,8 @@ void FusedAttnConfig::derive() { // Use ragged (TH1) or dense (BHS1) graphs and stats const int sm_arch = cuda::sm_arch(cuda::current_device()); - uses_ragged_graph = sm_arch >= 90 && sm_arch != 120; + const bool is_sm12x = sm_arch >= 120 && sm_arch < 130; + uses_ragged_graph = sm_arch >= 90 && !is_sm12x; 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 = diff --git a/transformer_engine/jax/attention.py b/transformer_engine/jax/attention.py index edf418f15ea..2765e1fbb5b 100644 --- a/transformer_engine/jax/attention.py +++ b/transformer_engine/jax/attention.py @@ -17,7 +17,6 @@ from transformer_engine_jax import NVTE_Mask_Type from transformer_engine_jax import NVTE_QKV_Layout from transformer_engine_jax import NVTE_QKV_Format -from transformer_engine_jax import nvte_get_qkv_format from transformer_engine_jax import NVTE_Softmax_Type from . import cpp_extensions as tex @@ -33,6 +32,7 @@ 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 + ALIBI = NVTE_Bias_Type.NVTE_ALIBI class AttnMaskType(Enum): @@ -139,7 +139,7 @@ def get_qkv_format(self): """ Return the corresponding qkv_format (BSHD, SBHD, THD) """ - return QKVFormat(nvte_get_qkv_format(self.value)) + return QKVFormat.THD if self.is_thd() else QKVFormat.BSHD def is_qkvpacked(self): """ @@ -348,6 +348,7 @@ def is_fused_attn_kernel_available( bias_seqlen_q: Optional[int] = None, bias_seqlen_kv: Optional[int] = None, max_segments_per_seq: int = 1, + fp8: bool = False, ): """ To check whether the fused attention kernel is supported. @@ -384,6 +385,7 @@ def make_helper(attn_mask_type): bias_seqlen_q=bias_seqlen_q, bias_seqlen_kv=bias_seqlen_kv, max_segments_per_seq=max_segments_per_seq, + fp8=fp8, ) helper = make_helper(attn_mask_type) @@ -1142,6 +1144,7 @@ def _legacy_fused_attn( context_parallel_axis: str = "", softmax_offset: Optional[jnp.ndarray] = None, return_max_logit: bool = False, + fp8: bool = False, ): """ Perform non-THD (non-packed) cuDNN fused attention. @@ -1236,6 +1239,7 @@ 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, + fp8=fp8, ) return output @@ -1262,6 +1266,7 @@ def fused_attn_thd( context_parallel_axis: str = "", softmax_offset: Optional[jnp.ndarray] = None, return_max_logit: bool = False, + fp8: bool = False, ): """ Deprecated THD fused attn, please use fusd_attn with SequenceDescriptor @@ -1317,6 +1322,7 @@ def fused_attn_thd( context_parallel_causal_load_balanced=context_parallel_causal_load_balanced, context_parallel_axis=context_parallel_axis, return_max_logit=return_max_logit, + fp8=fp8, ) return output @@ -1324,7 +1330,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, ...], @@ -1347,6 +1353,7 @@ def _fused_attn( context_checkpoint_name: str = "context", stripe_size: int | None = None, return_max_logit: bool = False, + fp8: bool = False, ): output, _ = _fused_attn_fwd_rule( qkv, @@ -1369,6 +1376,7 @@ def _fused_attn( context_checkpoint_name=context_checkpoint_name, stripe_size=stripe_size, return_max_logit=return_max_logit, + fp8=fp8, ) return output @@ -1394,6 +1402,7 @@ def _fused_attn_fwd_rule( context_checkpoint_name, stripe_size, return_max_logit, + fp8, ): output, softmax_aux, rng_state, max_logit = tex.fused_attn_fwd( qkv, @@ -1415,6 +1424,7 @@ def _fused_attn_fwd_rule( context_parallel_axis=context_parallel_axis, stripe_size=stripe_size, return_max_logit=return_max_logit, + fp8=fp8, ) output = checkpoint_name(output, context_checkpoint_name) softmax_aux = checkpoint_name(softmax_aux, context_checkpoint_name) @@ -1448,6 +1458,7 @@ def _fused_attn_bwd_rule( context_checkpoint_name, stripe_size, return_max_logit, + fp8, ctx, dz, ): @@ -1485,8 +1496,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, + fp8=fp8, ) - 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 @@ -1580,6 +1592,7 @@ 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, + fp8: bool = False, ): """ Perform cuDNN fused attention. @@ -1636,6 +1649,9 @@ def fused_attn( non-differentiable auxiliary inputs. score_mod_bprop_tensors (Optional[Mapping[str, Any]]): Additional tensors or Python/NumPy scalars made available to `score_mod_bprop`. + fp8 (bool): Use tensor-scaled cuDNN FP8 SDPA with E4M3 query/key/value and + E5M2 output gradients. Scaling uses the current tensor maximum. Outputs + and gradients retain the input FP16/BF16 dtype. return_max_logit (bool): If True, also return per-head maximum attention logits with shape ``[h]``. Returns: @@ -1686,6 +1702,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 fp8: + raise ValueError("FP8 is not supported with score_mod fused_attn.") if return_max_logit: raise ValueError("return_max_logit is not supported with score_mod fused_attn.") tex.validate_fused_attn_score_mod( @@ -1748,6 +1766,7 @@ def fused_attn( context_parallel_axis=context_parallel_axis, softmax_offset=softmax_offset, return_max_logit=return_max_logit, + fp8=fp8, ) if max_segments_per_seq > 1 and not qkv_layout.is_thd(): warnings.warn( @@ -1779,5 +1798,6 @@ def fused_attn( context_checkpoint_name=context_checkpoint_name, stripe_size=stripe_size, return_max_logit=return_max_logit, + fp8=fp8, ) return output diff --git a/transformer_engine/jax/cpp_extensions/attention.py b/transformer_engine/jax/cpp_extensions/attention.py index 3a1de7bd677..a5b1a747725 100644 --- a/transformer_engine/jax/cpp_extensions/attention.py +++ b/transformer_engine/jax/cpp_extensions/attention.py @@ -2,6 +2,7 @@ # # See LICENSE for license information. """JAX/TE custom ops for attention""" + import logging import operator import os @@ -12,21 +13,15 @@ import jax import jax.numpy as jnp -from jax import dtypes, lax, ffi +from jax import dtypes, lax +from jax.interpreters import mlir from jax.sharding import PartitionSpec, NamedSharding from jax.experimental.custom_partitioning import SdyShardingRule -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.common.cudnn.graph import uses_ragged_attention_graph from transformer_engine.jax.attention import ( AttnBiasType, AttnMaskType, @@ -39,13 +34,10 @@ from ..sharding import with_sharding_constraint_by_logical_axes, HEAD_AXES, is_mesh_available from .base import BasePrimitive, register_primitive +from . import cudnn_attention 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, @@ -59,7 +51,6 @@ with_sharding_constraint, ) - __all__ = [ "FusedAttnHelper", "fused_attn_fwd", @@ -118,6 +109,7 @@ def setup_logging(): "stripe_size", "return_max_logit", "allow_fast_causal_path", + "fp8", ], ) @dataclass(frozen=True) @@ -144,6 +136,7 @@ class _FusedAttnConfig: ) # Only for CP + Striped. For Ring P2P, stripe_size=1 only.For AG, stripe_size>=1. return_max_logit: bool = False allow_fast_causal_path: bool = True + fp8: bool = False @property def effective_window_size(self) -> Tuple[int, int]: @@ -153,64 +146,6 @@ def effective_window_size(self) -> Tuple[int, int]: 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) class FusedAttnHelper: """ @@ -241,6 +176,7 @@ class FusedAttnHelper: bias_seqlen_q: Optional[int] = None bias_seqlen_kv: Optional[int] = None max_segments_per_seq: int = 1 + fp8: bool = False def is_fused_attn_kernel_available(self): """Check if there is available fused attention kernel""" @@ -257,55 +193,68 @@ def get_fused_attn_backend(self): 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: + if self.q_dtype != self.kv_dtype: 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( + backend = NVTE_Fused_Attn_Backend.NVTE_No_Backend + message = "" + if self.attn_bias_type == AttnBiasType.PRE_SCALE_BIAS: + message = "Fused attention does not support pre-scale bias." + elif self.q_dtype not in (jnp.float16, jnp.bfloat16): + message = "Fused attention requires FP16/BF16 inputs." + else: + config = _FusedAttnConfig( + attn_bias_type=self.attn_bias_type, + attn_mask_type=self.attn_mask_type, + softmax_type=self.softmax_type, + qkv_layout=self.qkv_layout, + scaling_factor=self.attn_scale, + dropout_probability=self.dropout_probability, 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], + max_segments_per_seq=self.max_segments_per_seq, + window_size=self.window_size, 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, + context_parallel_load_balanced=False, + cp_axis="", + cp_striped_window_size=None, + stripe_size=None, + return_max_logit=self.return_max_logit, + fp8=self.fp8, ) - ) + qshape = (self.batch_size, self.q_max_seqlen, self.q_num_heads, self.head_dim_qk) + kshape = (self.batch_size, self.kv_max_seqlen, self.kv_num_heads, self.head_dim_qk) + vshape = (self.batch_size, self.kv_max_seqlen, self.kv_num_heads, self.head_dim_v) + if self.qkv_layout.is_qkvpacked(): + qshape = (*qshape[:-2], 3, *qshape[-2:]) + kshape = vshape = (0,) + elif self.qkv_layout.is_kvpacked(): + kshape = (*kshape[:-2], 2, *kshape[-2:]) + vshape = (0,) + bias_shape = ( + ( + self.bias_batch or 1, + self.bias_heads or self.q_num_heads, + self.bias_seqlen_q or self.q_max_seqlen, + self.bias_seqlen_kv or self.kv_max_seqlen, + ) + if self.attn_bias_type == AttnBiasType.POST_SCALE_BIAS + else (0,) + ) + avals = tuple( + jax.ShapeDtypeStruct(shape, self.q_dtype) + for shape in (qshape, kshape, vshape, bias_shape) + ) + avals += (jax.ShapeDtypeStruct((1, self.q_num_heads, 1, 1), jnp.float32),) + try: + cudnn_attention.check_attention_support("fwd", avals, config) + if self.is_training: + cudnn_attention.check_attention_support("bwd", avals, config) + backend = ( + NVTE_Fused_Attn_Backend.NVTE_FP8 + if self.fp8 + else NVTE_Fused_Attn_Backend.NVTE_F16_arbitrary_seqlen + ) + except ValueError as exc: + message = str(exc) AttentionLogging.setup_logging() logger = AttentionLogging.logger @@ -401,7 +350,8 @@ class _FusedAttnRNGStateChecker: The fused attention backend requires a 64 bits seed and a 64 bits offset. However, JAX doesn't enable 64 bits by default, so we have to emulate seed as two 32 bits array. - The offset calculation is maintained in the backend. + The explicit JAX key supplies the seed; the saved zero offset makes + dropout replayable when a forward pass is recomputed. """ rng_state_dtype: jnp.dtype = jnp.uint32 @@ -538,9 +488,13 @@ def abstract( bias_seqlen_q=bias_seqlen_q, bias_seqlen_kv=bias_seqlen_kv, max_segments_per_seq=config.max_segments_per_seq, + fp8=config.fp8, ).get_fused_attn_backend() - if backend == NVTE_Fused_Attn_Backend.NVTE_F16_arbitrary_seqlen: + if backend in ( + NVTE_Fused_Attn_Backend.NVTE_F16_arbitrary_seqlen, + NVTE_Fused_Attn_Backend.NVTE_FP8, + ): if config.qkv_layout.is_thd(): softmax_shape = (*batch_shape, q_max_seqlen, attn_heads, 1) else: @@ -550,8 +504,8 @@ def abstract( 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]; non-SM120 THD uses [..., S_q, H, 1]. + # cuDNN Max is row-wise over S_kv. Dense and SM12x THD use + # [..., H, S_q, 1]; other THD layouts use [..., 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) @@ -571,42 +525,16 @@ 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: + if config.attn_bias_type in (AttnBiasType.NO_BIAS, AttnBiasType.ALIBI): 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]) + graph = cudnn_attention.get_graph( + "fwd", (q_aval, k_aval, v_aval, bias_aval, softmax_offset_aval), config ) + wkspace_aval = q_aval.update(shape=(graph.graph.workspace_size,), dtype=jnp.uint8) assert ( softmax_offset_aval.dtype == jnp.float32 @@ -656,32 +584,12 @@ def lowering( *, config: _FusedAttnConfig, ): - """ - Fused attention fwd lowering rules - """ - q_aval, k_aval, v_aval, bias_aval, *_ = ctx.avals_in + """Lower Python-built cuDNN attention through the JAX-owned graph executor.""" - ( - 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) + def execute(*args): + return cudnn_attention.execute("fwd", args, config, ctx.avals_out) - 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 - - return ffi.ffi_lowering(FusedAttnFwdPrimitive.name)( + return mlir.lower_fun(execute, multiple_results=True)( ctx, q, k, @@ -696,29 +604,7 @@ def lowering( _q_segment_ids, _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), + _kv_segment_pos, ) @staticmethod @@ -839,8 +725,8 @@ def convert_to_2d(offsets, batch, max_seqlen): def _reduce_max_logit(max_tensor, output, q_seqlen, q_seq_offsets, config): """Reduce cuDNN's row-wise Max tensor to the public per-head max_logit. - Dense and SM120 THD use ``[..., H, S_q, 1]``; non-SM120 - THD uses ``[..., S_q, H, 1]``. A rank-3 THD result is ``[T_q, H, 1]``. + Dense and SM12x THD use ``[..., H, S_q, 1]``; THD with packed statistics + uses ``[..., S_q, H, 1]``. A rank-3 THD result is ``[T_q, H, 1]``. All layouts reduce to ``[H]``. Static THD buffers can contain invalid query rows, so those rows are masked before reduction. """ @@ -879,7 +765,9 @@ def _reduce_max_logit(max_tensor, output, q_seqlen, q_seq_offsets, config): @staticmethod def _uses_thd_ragged_max_tensor(config): """Return whether cuDNN writes THD Max with BSH-like ragged-stats layout.""" - return config.qkv_layout.is_thd() and 120 not in get_all_device_compute_capability() + return config.qkv_layout.is_thd() and uses_ragged_attention_graph( + cudnn_attention.device_capabilities()[0] + ) @staticmethod def _empty_or_neg_inf_max_logit(head, dtype, config): @@ -1114,57 +1002,13 @@ def abstract( f" kv_seqlen_or_cu_seqlen_aval.dtype={kv_seqlen_or_cu_seqlen_aval.dtype}" ) - ( - 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) - - 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, - attn_heads, - num_gqa_groups, - bias_heads, - 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, - ) + attn_heads = FusedAttnHelper.parse_qkv_aval(q_aval, k_aval, v_aval, config.qkv_layout)[3] - 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) + graph = cudnn_attention.get_graph( + "bwd", (q_aval, k_aval, v_aval, bias_aval, softmax_offset_aval), config ) + dq_aval, dk_aval, dv_aval, dbias_aval = q_aval, k_aval, v_aval, bias_aval + wkspace_aval = q_aval.update(shape=(graph.graph.workspace_size,), dtype=jnp.uint8) # Validate incoming softmax_offset shape and dtype assert ( @@ -1223,32 +1067,12 @@ def lowering( *, config, ): - """ - Fused attention bwd lowering rules - """ - 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) + """Lower Python-built cuDNN attention through the JAX-owned graph executor.""" - input_batch = reduce(operator.mul, batch_shape) + def execute(*args): + return cudnn_attention.execute("bwd", args, config, ctx.avals_out) - 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 - - return ffi.ffi_lowering(FusedAttnBwdPrimitive.name)( + return mlir.lower_fun(execute, multiple_results=True)( ctx, q, k, @@ -1266,28 +1090,7 @@ def lowering( q_segment_ids, 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), + kv_segment_pos, ) @staticmethod @@ -1714,6 +1517,7 @@ def get_step_config(self) -> _FusedAttnConfig: cp_striped_window_size=None, stripe_size=self.config.stripe_size, return_max_logit=self.config.return_max_logit, + fp8=self.config.fp8, ) def get_step_config_for_striped(self, max_seqlen, cp_size) -> _FusedAttnConfig: @@ -1735,6 +1539,7 @@ def get_step_config_for_striped(self, max_seqlen, cp_size) -> _FusedAttnConfig: cp_striped_window_size=None, stripe_size=self.config.stripe_size, return_max_logit=self.config.return_max_logit, + fp8=self.config.fp8, ) def all_gather_kv(self, k, v): @@ -2807,6 +2612,7 @@ def get_step_config(self, attn_mask_type) -> _FusedAttnConfig: cp_striped_window_size=None, stripe_size=self.config.stripe_size, return_max_logit=self.config.return_max_logit, + fp8=self.config.fp8, ) def get_striped_thd_step_config(self) -> _FusedAttnConfig: @@ -3089,7 +2895,7 @@ def correction(output, softmax_aux, output_per_step, softmax_aux_per_step): else: for i in range(0, cp_size): carry = scan_kv_block(i, carry) - (kv, output, softmax_aux, max_logit) = carry + kv, output, softmax_aux, max_logit = carry output = output.astype(q.dtype) # Globalize the rank-local running [H] max across DP/CP. @@ -3314,7 +3120,7 @@ def jax_cond_wrap(): else: for i in range(0, cp_size): carry = scan_kv_block(i, carry) - (kv, dq, dk_dv, dbias) = carry + kv, dq, dk_dv, dbias = carry # Final permute to put gradients back to their final resting place. dk_dv = helper.permute_kv(dk_dv, cp_perm) @@ -3553,7 +3359,7 @@ def correction(output, softmax_aux, output_per_step, softmax_aux_per_step): else: for i in range(0, cp_size): carry = scan_kv_block(i, carry) - (_, _, _, output, softmax_aux, max_logit) = carry + _, _, _, output, softmax_aux, max_logit = carry output = output.astype(q.dtype) # Globalize the rank-local running [H] max across DP/CP. @@ -3685,7 +3491,7 @@ def compute(config): else: for idx in range(cp_size): carry = scan_kv_block(idx, carry) - (_, _, _, dq, dkv, dbias) = carry + _, _, _, dq, dkv, dbias = carry # Final permute to put gradients back to their final resting place. dkv = helper.permute_kv(dkv, cp_perm) @@ -3735,6 +3541,7 @@ def fused_attn_fwd( context_parallel_axis: str = "", stripe_size: int | None = None, return_max_logit: bool = False, + fp8: bool = False, ) -> jnp.ndarray: """ Perform the forward pass of with cuDNN fused attention implementations. @@ -3798,7 +3605,7 @@ 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}" @@ -3851,6 +3658,7 @@ def fused_attn_fwd( cp_striped_window_size=None, stripe_size=stripe_size, return_max_logit=return_max_logit, + fp8=fp8, ) primitive = None @@ -3902,6 +3710,7 @@ def fused_attn_bwd( context_parallel_causal_load_balanced: bool = False, context_parallel_axis: str = "", stripe_size: int | None = None, + fp8: bool = False, ): """ Perform the backward pass of the cuDNN fused attention implementations. @@ -3967,7 +3776,7 @@ def fused_attn_bwd( 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 with type={type(bias)}" @@ -3995,15 +3804,18 @@ def fused_attn_bwd( softmax_offset, (None, HEAD_AXES, None, None) ) - compute_capabilities = get_all_device_compute_capability() + compute_capabilities = cudnn_attention.device_capabilities() if any(x >= 100 for x in compute_capabilities) and is_training: assert ( FusedAttnHelper.is_non_deterministic_allowed() - and (attn_bias_type == AttnBiasType.NO_BIAS or dropout_probability == 0.0) + and ( + attn_bias_type in (AttnBiasType.NO_BIAS, AttnBiasType.ALIBI) + or dropout_probability == 0.0 + ) ) or ( not FusedAttnHelper.is_non_deterministic_allowed() - and get_cudnn_version() >= (9, 18, 1) - and attn_bias_type == AttnBiasType.NO_BIAS + and cudnn_attention.cudnn_version() >= 91801 + and attn_bias_type in (AttnBiasType.NO_BIAS, AttnBiasType.ALIBI) and dropout_probability == 0.0 ), ( "For sm100+, non-deterministic bprop does not support bias with dropout," @@ -4012,6 +3824,7 @@ def fused_attn_bwd( fused_config = _FusedAttnConfig( attn_bias_type=attn_bias_type, + fp8=fp8, attn_mask_type=attn_mask_type, qkv_layout=qkv_layout, softmax_type=softmax_type, 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..054649e8585 --- /dev/null +++ b/transformer_engine/jax/cpp_extensions/cudnn_attention.py @@ -0,0 +1,542 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. +"""JAX cuDNN attention graphs, metadata preparation and cached plan dispatch.""" + +import math +import os +import threading +from dataclasses import dataclass +from functools import lru_cache + +import jax +import jax.numpy as jnp +import numpy as np +from jax import ffi + +from transformer_engine.common.cudnn.attention import ( + AttentionOptions, + AttentionGraphSpec, + BuiltAttentionGraph, + build_attention_graph, +) +from transformer_engine.common.cudnn.graph import ( + TensorSpec, + CacheDiagnostics, + bucket as _bucket, + row_major_stride as _stride, + create_graph, + check_graph_support, + compile_graph, + uses_ragged_attention_graph, +) + +from .cudnn_graph import import_cudnn, check_frontend_version, serialize_graph, graph_attributes + + +def device_capabilities(): + """Query the JAX devices without entering TE common.""" + return tuple(int(float(device.compute_capability) * 10) for device in jax.local_devices()) + + +def cudnn_version(): + """Return the backend used by the Python frontend.""" + return int(import_cudnn().backend_version()) + + +def _dtype(dtype): + return { + jnp.dtype(jnp.float16): "HALF", + jnp.dtype(jnp.bfloat16): "BFLOAT16", + jnp.dtype(jnp.float32): "FLOAT", + jnp.dtype(jnp.float8_e4m3fn): "FP8_E4M3", + jnp.dtype(jnp.float8_e5m2): "FP8_E5M2", + }[jnp.dtype(dtype)] + + +def _dims(avals, config): + q, k, v = avals[:3] + layout = config.qkv_layout + if layout.is_qkvpacked(): + *batch, sq, _, h, dq = q.shape + sk, hg, dv = sq, h, dq + elif layout.is_kvpacked(): + *batch, sq, h, dq = q.shape + *_, sk, _, hg, dv = k.shape + else: + *batch, sq, h, dq = q.shape + sk, hg, dv = v.shape[-3:] + return math.prod(batch), sq, sk, h, hg, dq, dv + + +@dataclass(frozen=True) +class AttentionGraph: + """Plan and bindings, including views of packed QKV storage.""" + + graph: object + names: tuple + input_indices: tuple + input_offsets: tuple + output_indices: tuple + output_offsets: tuple + metadata_batch: int + ragged_stats: bool + multipliers: tuple + fp8: bool + + +@dataclass(frozen=True) +class _AttentionDescription: + """Static graph metadata and JAX's packed-storage buffer bindings.""" + + spec: AttentionGraphSpec + inputs: tuple + outputs: tuple + scalars: tuple + metadata_batch: int + ragged_stats: bool + multipliers: tuple + fp8: bool + + +@dataclass +class _SupportedGraph: + """A supported live graph and its lazily compiled serialization.""" + + built: BuiltAttentionGraph + workspace_size: int | None = None + serialized_graph: bytes | None = None + + +_diagnostics = CacheDiagnostics("JAX") +_plan_lock = threading.RLock() + + +def _site(spec): + return "fp8" if spec.options.fp8 else "f16", "bwd" if spec.options.backward else "fwd" + + +@lru_cache(maxsize=256) +def _supported_graph(spec, capabilities, frontend_version, backend_version): + # Versions and target architecture are part of the plan identity, even when + # two targets happen to accept identical logical tensor descriptors. + del frontend_version, backend_version + cudnn = import_cudnn() + backend, direction = _site(spec) + _diagnostics.event(backend, direction, capabilities, "MISS", spec) + graph = create_graph(cudnn, dict(spec.tensors)["o"].dtype) + built = build_attention_graph(cudnn, graph, spec) + _diagnostics.event(backend, direction, capabilities, "CREATE_GRAPH") + check_graph_support( + cudnn, graph, lambda stage, function: _diagnostics.run(backend, direction, stage, function) + ) + _diagnostics.event(backend, direction, capabilities, "CACHE_GRAPH") + return _SupportedGraph(built) + + +def _request(direction, avals, config): + shapes = tuple((tuple(aval.shape), jnp.dtype(aval.dtype)) for aval in avals[:5]) + return ( + direction, + shapes, + config, + device_capabilities(), + bool(int(os.getenv("NVTE_ALLOW_NONDETERMINISTIC_ALGO", "1"))), + bool(int(os.getenv("NVTE_FP8_DPA_BWD", "1"))), + ) + + +def _get_supported_graph(spec, capabilities, frontend_version, backend_version): + """Share support-query graphs with compilation, and report cache reuse.""" + # pylint infers the wrapped function's arguments for lru_cache.cache_info. + # pylint: disable=no-value-for-parameter + with _plan_lock: + hits = _supported_graph.cache_info().hits + entry = _supported_graph(spec, capabilities, frontend_version, backend_version) + if _supported_graph.cache_info().hits > hits: + _diagnostics.event(*_site(spec), capabilities, "HIT") + return entry + + +def check_attention_support(direction, avals, config): + """Check the common graph without compiling kernels or serializing a plan.""" + description = _describe_graph(*_request(direction, avals, config)) + cudnn = import_cudnn() + _get_supported_graph( + description.spec, + device_capabilities(), + check_frontend_version(cudnn), + int(cudnn.backend_version()), + ) + + +def get_graph(direction, avals, config): + """Compile when tracing requires workspace sizes and serialized FFI attributes.""" + return _build_graph(*_request(direction, avals, config)) + + +@lru_cache(maxsize=256) +def _build_graph(direction, shapes, config, capabilities, nondeterministic, fp8_backward): + description = _describe_graph( + direction, shapes, config, capabilities, nondeterministic, fp8_backward + ) + cudnn = import_cudnn() + version = check_frontend_version(cudnn) + entry = _get_supported_graph( + description.spec, capabilities, version, int(cudnn.backend_version()) + ) + backend, direction = _site(description.spec) + with _plan_lock: + if entry.workspace_size is None: + workspace_size = compile_graph( + cudnn, + entry.built.graph, + lambda stage, function: _diagnostics.run(backend, direction, stage, function), + ) + serialized_graph = bytes(entry.built.graph.serialize()) + entry.workspace_size = workspace_size + entry.serialized_graph = serialized_graph + _diagnostics.event(backend, direction, capabilities, "BUILD_PLANS") + uids = entry.built.uids + serialized = serialize_graph( + serialized_graph=entry.serialized_graph, + cudnn_frontend_version=version, + workspace_size=entry.workspace_size, + input_uids=[uids[name] for name, _, _ in description.inputs], + output_uids=[uids[name] for name, _, _ in description.outputs], + scalar_uids=[uids[name] for name, _ in description.scalars], + scalar_values=[value for _, value in description.scalars], + ) + return AttentionGraph( + serialized, + tuple(source for _, source, _ in description.inputs), + tuple(range(len(description.inputs))), + tuple(offset for _, _, offset in description.inputs), + tuple(slot for _, slot, _ in description.outputs), + tuple(offset for _, _, offset in description.outputs), + description.metadata_batch, + description.ragged_stats, + description.multipliers, + description.fp8, + ) + + +def _describe_graph(direction, shapes, config, capabilities, nondeterministic, fp8_backward): + """Translate JAX's API and storage conventions into a common graph description.""" + avals = tuple(jax.ShapeDtypeStruct(*shape) for shape in shapes) + batch, sq, sk, h, hg, dq, dv = _dims(avals, config) + backward = direction == "bwd" + ragged = config.qkv_layout.is_thd() + ragged_stats = ragged and uses_ragged_attention_graph(capabilities[0]) + segments = batch * config.max_segments_per_seq if ragged else batch + # JAX uses fixed metadata capacities, including for empty ragged storage. + # Keep positive graph capacities while runtime segment lengths may be zero. + graph_batch = _bucket(max(segments, 1), 32, 512) if ragged_stats else segments + graph_sq = _bucket(max(batch * sq, 1), 1024, 32768) if ragged_stats else sq + graph_sk = _bucket(max(batch * sk, 1), 1024, 32768) if ragged_stats else sk + fp8 = config.fp8 and (not backward or fp8_backward) + version = cudnn_version() + if ragged and capabilities[0] < 90: + raise ValueError("F16/BF16 fused attention with THD format requires sm90 or later.") + if ( + backward + and ragged_stats + and config.softmax_type.name == "LEARNABLE_SOFTMAX" + and version < 92600 + ): + raise ValueError("THD learnable softmax backward requires cuDNN 9.26.0 or later.") + if fp8 and ragged: + if version < 92300: + raise ValueError("FP8 THD attention requires cuDNN 9.23.0 or later.") + if backward and capabilities[0] < 100: + raise ValueError("FP8 THD backward requires sm100 or later.") + if fp8 and backward and config.attn_bias_type.name != "NO_BIAS": + raise ValueError("cuDNN FP8 backward does not support attention bias or ALiBi.") + io_dtype = _dtype(avals[0].dtype) + qkv_dtype = "FP8_E4M3" if fp8 else io_dtype + specs, inputs, outputs, scalars = {}, [], [], [] + + def tensor(name, dim, stride, dtype=io_dtype, source=None, offset=0, ragged_offset=""): + specs[name] = TensorSpec(tuple(dim), tuple(stride), dtype, ragged_offset=ragged_offset) + inputs.append((name, name if source is None else source, offset)) + + def constant(name, value): + specs[name] = TensorSpec((1, 1, 1, 1), (1, 1, 1, 1), "FLOAT", by_value=True) + scalars.append((name, np.asarray(value, np.float32).tobytes())) + + def output(name, dim, stride, slot, offset=0, dtype=io_dtype, ragged_offset=""): + specs[name] = TensorSpec( + tuple(dim), tuple(stride), dtype, output=True, ragged_offset=ragged_offset + ) + outputs.append((name, slot, offset)) + + qpack = 3 if config.qkv_layout.is_qkvpacked() else 1 + kvpack = 3 if qpack == 3 else (2 if config.qkv_layout.is_kvpacked() else 1) + itemsize = 1 if fp8 else avals[0].dtype.itemsize + qstride = (graph_sq * qpack * h * dq, dq, qpack * h * dq, 1) + kstride = (graph_sk * kvpack * hg * dq, dq, kvpack * hg * dq, 1) + vstride = (graph_sk * kvpack * hg * dv, dv, kvpack * hg * dv, 1) + ostride = (graph_sq * h * dv, dv, h * dv, 1) + qdim = (graph_batch, h, graph_sq, dq) + kdim = (graph_batch, hg, graph_sk, dq) + vdim = (graph_batch, hg, graph_sk, dv) + odim = (graph_batch, h, graph_sq, dv) + tensor("q", qdim, qstride, qkv_dtype, ragged_offset="offset_q" if ragged else "") + tensor( + "k", + kdim, + kstride, + qkv_dtype, + source="q" if qpack == 3 else "k", + offset=h * dq * itemsize if qpack == 3 else 0, + ragged_offset="offset_k" if ragged else "", + ) + tensor( + "v", + vdim, + vstride, + qkv_dtype, + source="q" if qpack == 3 else ("k" if kvpack == 2 else "v"), + offset=2 * h * dq * itemsize if qpack == 3 else (hg * dq * itemsize if kvpack == 2 else 0), + ragged_offset="offset_v" if ragged else "", + ) + multipliers = (qpack * h * dq, kvpack * hg * dq, kvpack * hg * dv, h * dv, h) + if ragged: + for name in ("offset_q", "offset_k", "offset_v", "offset_o", "offset_stats"): + tensor(name, (graph_batch + 1, 1, 1, 1), (1, 1, 1, 1), "INT64") + stats_dim = (graph_batch, h, graph_sq, 1) + stats_stride = (h * graph_sq, 1, h, 1) if ragged_stats else (h * graph_sq, graph_sq, 1, 1) + stats_offset = "offset_stats" if ragged_stats else "" + o_offset = "offset_o" if ragged else "" + left, right = config.effective_window_size + if config.attn_mask_type.is_causal(): + right = 0 + padding = config.attn_mask_type.is_padding() or ragged + if padding: + for name in ("seq_q", "seq_kv"): + tensor(name, (graph_batch, 1, 1, 1), (1, 1, 1, 1), "INT32") + dropout = float(config.dropout_probability) if config.is_training else 0.0 + if dropout: + tensor("seed", (1, 1, 1, 1), (1, 1, 1, 1), "INT64", source="rng") + tensor("offset", (1, 1, 1, 1), (1, 1, 1, 1), "INT64", source="rng", offset=8) + bias_type = config.attn_bias_type.name + if bias_type not in ("ALIBI", "NO_BIAS"): + bias_shape = (math.prod(avals[3].shape[:-3]), *avals[3].shape[-3:]) + tensor("bias", bias_shape, _stride(bias_shape)) + if backward and bias_shape[:3] != (1, 1, 1): + output("dbias", bias_shape, _stride(bias_shape), 3) + if config.softmax_type.name != "VANILLA_SOFTMAX": + tensor("sink", (1, h, 1, 1), (h, 1, 1, 1), "FLOAT") + if backward: + output("dsink", (1, h, 1, 1), (h, 1, 1, 1), 4, dtype="FLOAT") + constant("attn_scale", config.scaling_factor) + if fp8: + for name in ("descale_q", "descale_k", "descale_v"): + tensor(name, (1, 1, 1, 1), (1, 1, 1, 1), "FLOAT") + constant("scale_s", 448.0) + constant("descale_s", 1.0 / 448.0) + if backward: + tensor("o", odim, ostride, ragged_offset=o_offset) + tensor("do", odim, ostride, "FP8_E5M2" if fp8 else io_dtype, ragged_offset=o_offset) + tensor("stats", stats_dim, stats_stride, "FLOAT", ragged_offset=stats_offset) + grad_slots = (0, 0 if qpack == 3 else 1, 0 if qpack == 3 else (1 if kvpack == 2 else 2)) + grad_offsets = ( + 0, + h * dq * avals[0].dtype.itemsize if qpack == 3 else 0, + ( + 2 * h * dq * avals[0].dtype.itemsize + if qpack == 3 + else (hg * dq * avals[0].dtype.itemsize if kvpack == 2 else 0) + ), + ) + for name, dim, stride, slot, offset, ragged_offset in zip( + ("dq", "dk", "dv"), + (qdim, kdim, vdim), + (qstride, kstride, vstride), + grad_slots, + grad_offsets, + ("offset_q", "offset_k", "offset_v"), + ): + output(name, dim, stride, slot, offset, ragged_offset=ragged_offset if ragged else "") + if fp8: + tensor("descale_do", (1, 1, 1, 1), (1, 1, 1, 1), "FLOAT") + for name in ("descale_o", "scale_dp", "descale_dp", "scale_dq", "scale_dk", "scale_dv"): + constant(name, 1.0) + for slot, name in enumerate(("amax_dq", "amax_dk", "amax_dv", "amax_dp"), 5): + output(name, (1, 1, 1, 1), (1, 1, 1, 1), slot, dtype="FLOAT") + else: + output("o", odim, ostride, 0, ragged_offset=o_offset) + output("stats", stats_dim, stats_stride, 1, dtype="FLOAT", ragged_offset=stats_offset) + if config.return_max_logit: + output("max", stats_dim, stats_stride, 2, dtype="FLOAT", ragged_offset=stats_offset) + if fp8: + constant("scale_o", 1.0) + output("amax_s", (1, 1, 1, 1), (1, 1, 1, 1), 3, dtype="FLOAT") + output("amax_o", (1, 1, 1, 1), (1, 1, 1, 1), 4, dtype="FLOAT") + options = AttentionOptions( + backward=backward, + fp8=fp8, + bottom_right=config.bottom_right_diagonal, + left_bound=left + 1 if left >= 0 else None, + right_bound=right if right >= 0 else None, + padding=padding, + alibi=bias_type == "ALIBI", + dropout=dropout, + deterministic=backward and not nondeterministic, + max_total_seqlen_q=graph_sq if backward and ragged_stats and not fp8 else 0, + max_total_seqlen_kv=graph_sk if backward and ragged_stats and not fp8 else 0, + ) + return _AttentionDescription( + AttentionGraphSpec.from_dict(options, specs), + tuple(inputs), + tuple(outputs), + tuple(scalars), + graph_batch, + ragged_stats, + multipliers, + fp8, + ) + + +def _offset64(offsets, multiplier): + """Compute exact uint64 element offsets without enabling JAX x64 globally.""" + value = offsets.astype(jnp.uint32) + low, high = value & 65535, value >> 16 + ml, mh = multiplier & 65535, multiplier >> 16 + w0 = low * ml + carry = high * ml + (w0 >> 16) + w1 = (carry & 65535) + low * mh + upper = high * mh + (carry >> 16) + (w1 >> 16) + lower = (w1 << 16) | (w0 & 65535) + return jnp.stack((lower, upper), axis=-1) + + +def _quantize(value, dtype): + """Current tensor scaling implemented in JAX, independently of TE common.""" + maximum = jnp.max(jnp.abs(value.astype(jnp.float32))) + scale = jnp.where(maximum > 0, float(jnp.finfo(dtype).max) / maximum, 1.0) + scale = jnp.minimum(scale, jnp.finfo(jnp.float32).max) + return (value.astype(jnp.float32) * scale).astype(dtype), jnp.reciprocal(scale) + + +def _metadata(value, size, fill=0): + """Bind only the configured segment capacity from fixed-size JAX metadata.""" + value = value.reshape(-1)[:size] + return jnp.pad(value, (0, size - value.size), constant_values=fill) + + +def _dense_stats_to_packed(value, offsets, lengths, shape): + """Convert SM12x's per-segment BHS1 statistics to JAX's packed TH1 layout.""" + tokens, heads = math.prod(shape[:-2]), shape[-2] + rows = value.shape[-2] + indices = offsets[: lengths.size, None] + jnp.arange(rows)[None, :] + indices = jnp.where(jnp.arange(rows)[None, :] < lengths[:, None], indices, tokens) + values = value.transpose(0, 2, 1, 3).reshape(-1, heads, 1) + packed = jnp.full((tokens, heads, 1), -1e30, value.dtype) + return packed.at[indices.reshape(-1)].set(values, mode="drop").reshape(shape) + + +def _packed_stats_to_dense(value, offsets, lengths, rows): + """Reconstruct dense per-segment statistics for the cuDNN backward graph.""" + heads = value.shape[-2] + indices = offsets[: lengths.size, None] + jnp.arange(rows)[None, :] + valid = jnp.arange(rows)[None, :] < lengths[:, None] + dense = jnp.take(value.reshape(-1, heads, 1), indices, axis=0, mode="fill", fill_value=0) + dense = jnp.where(valid[:, :, None, None], dense, 0) + return dense.transpose(0, 2, 1, 3) + + +def execute(direction, args, config, result_avals): + """Prepare device metadata with JAX and dispatch a cached graph through FFI.""" + plan = get_graph(direction, args, config) + q, k, v, bias, sink = args[:5] + backward = direction == "bwd" + if backward: + stats, rng, o, do, q_cu, kv_cu, q_offsets, kv_offsets = args[5:13] + else: + seed, q_cu, kv_cu, q_offsets, kv_offsets = args[5:10] + # JAX PRNG keys are explicit: replaying a key replays the dropout mask. + rng = jnp.broadcast_to( + jnp.concatenate((seed.reshape(-1)[:2], jnp.zeros(2, jnp.uint32))), result_avals[3].shape + ) + values = {"q": q, "k": k, "v": v, "bias": bias, "sink": sink, "rng": rng} + values["seq_q"] = _metadata(jnp.diff(q_cu), plan.metadata_batch) + values["seq_kv"] = _metadata(jnp.diff(kv_cu), plan.metadata_batch) + dense_ragged_stats = config.qkv_layout.is_thd() and not plan.ragged_stats + if config.qkv_layout.is_thd(): + batch, sq, sk, *_ = _dims(args, config) + q_offsets = _metadata(q_offsets, plan.metadata_batch + 1, batch * sq) + kv_offsets = _metadata(kv_offsets, plan.metadata_batch + 1, batch * sk) + for name, offsets, multiplier in zip( + ("offset_q", "offset_k", "offset_v", "offset_o", "offset_stats"), + (q_offsets, kv_offsets, kv_offsets, q_offsets, q_offsets), + plan.multipliers, + ): + values[name] = _offset64(offsets, multiplier) + if config.fp8: + for name in ("q", "k", "v"): + # Empty dummy operands are not quantized. + if values[name].size: + values[name], values["descale_" + name] = _quantize(values[name], jnp.float8_e4m3fn) + if config.qkv_layout.is_qkvpacked(): + values["descale_k"] = values["descale_v"] = values["descale_q"] + elif config.qkv_layout.is_kvpacked(): + values["descale_v"] = values["descale_k"] + if not plan.fp8: + for name in ("q", "k", "v"): + if values[name].size: + values[name] = ( + values[name].astype(jnp.float32) * values["descale_" + name] + ).astype(q.dtype) + if backward: + values.update(stats=stats, o=o, do=do) + if dense_ragged_stats: + values["stats"] = _packed_stats_to_dense( + stats, + q_offsets, + values["seq_q"], + q.shape[-4] if config.qkv_layout.is_qkvpacked() else q.shape[-3], + ) + if plan.fp8: + values["do"], values["descale_do"] = _quantize(do, jnp.float8_e5m2) + specs = list(result_avals[:-1]) + if plan.fp8: + specs += [jax.ShapeDtypeStruct((1,), jnp.float32)] * 4 + else: + specs = list(result_avals[:3]) + if dense_ragged_stats: + _, sq, _, heads, *_ = _dims(args, config) + dense_shape = (plan.metadata_batch, heads, sq, 1) + specs[1] = jax.ShapeDtypeStruct(dense_shape, jnp.float32) + if config.return_max_logit: + specs[2] = jax.ShapeDtypeStruct(dense_shape, jnp.float32) + if plan.fp8: + specs += [jax.ShapeDtypeStruct((1,), jnp.float32)] * 2 + specs.append(jax.ShapeDtypeStruct((plan.graph.workspace_size,), jnp.uint8)) + attributes = graph_attributes(plan.graph) + attributes.update( + input_indices=np.asarray(plan.input_indices, np.int64), + input_offsets=np.asarray(plan.input_offsets, np.int64), + output_indices=np.asarray(plan.output_indices, np.int64), + output_offsets=np.asarray(plan.output_offsets, np.int64), + ) + if not backward and config.qkv_layout.is_thd(): + # Match the neutral finite statistics value for empty ragged ring blocks. + # Filling FP32 bytes with 0xF0 gives a large negative number without NaNs + # when the ring correction subtracts statistics for two empty blocks. + fills = np.zeros(len(specs) - 1, np.int64) + fills[1:3] = 0xF0 + attributes["output_fill_values"] = fills + results = ffi.ffi_call( + "te_cudnn_graph_ffi", tuple(jax.ShapeDtypeStruct(spec.shape, spec.dtype) for spec in specs) + )(*(values[name] for name in plan.names), **attributes) + if backward: + return (*results[:5], results[-1]) + output, stats, score_max = results[:3] + if dense_ragged_stats: + packed_shape = result_avals[1].shape + stats = _dense_stats_to_packed(stats, q_offsets, values["seq_q"], packed_shape) + if config.return_max_logit: + score_max = _dense_stats_to_packed(score_max, q_offsets, values["seq_q"], packed_shape) + score_max = jnp.swapaxes(score_max, -3, -2) + return output, stats, score_max, rng, results[-1] 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..b5cc35b7579 --- /dev/null +++ b/transformer_engine/jax/cpp_extensions/cudnn_graph.py @@ -0,0 +1,129 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. +"""Serialize Python-built cuDNN graphs for the JAX-owned FFI executor.""" + +import hashlib +from dataclasses import dataclass +from typing import Sequence, Tuple + +import numpy as np +import transformer_engine_jax + +from transformer_engine.common.cudnn.graph import ( + import_cudnn as _import_cudnn, + check_graph_support, + compile_graph, +) + + +@dataclass(frozen=True) +class SerializedGraph: + """A compiled graph and its buffer bindings.""" + + 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 + + +def encode_frontend_version(version: str) -> int: + """Encode the public version as CUDNN_FRONTEND_VERSION.""" + parts = version.split("+", 1)[0].split("-", 1)[0].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_frontend_version(cudnn) -> int: + """Require compatible Python and C++ serialization formats.""" + version = getattr(cudnn, "__version__", None) + if version is None: + raise RuntimeError("cuDNN frontend Python package does not expose __version__.") + python_version = encode_frontend_version(version) + 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__={version!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(): + """Load the required Python frontend lazily.""" + cudnn = _import_cudnn() + check_frontend_version(cudnn) + return cudnn + + +def graph_hash(serialized_graph: bytes) -> Tuple[int, int]: + """Identify a serialized plan independently of Python object lifetimes.""" + digest = hashlib.sha256(serialized_graph).digest() + return tuple(int.from_bytes(digest[i : i + 8], "little", signed=True) for i in (0, 8)) + + +def pack_scalar_values(values: Sequence[bytes]) -> Tuple[np.ndarray, np.ndarray]: + """Pack pass-by-value constants into aligned slots.""" + sizes = np.asarray([len(value) for value in values], dtype=np.int64) + packed = np.zeros((len(values), 16), dtype=np.uint8) + for index, value in enumerate(values): + if len(value) > 16: + raise ValueError("cuDNN pass-by-value scalars must be at most 16 bytes.") + packed[index, : len(value)] = np.frombuffer(value, dtype=np.uint8) + return sizes, packed.reshape(-1) + + +def serialize_graph( + *, + serialized_graph, + cudnn_frontend_version, + workspace_size, + input_uids, + output_uids, + scalar_uids=(), + scalar_values=(), +) -> SerializedGraph: + """Package a graph and its stable tensor UIDs.""" + sizes, values = pack_scalar_values(scalar_values) + return SerializedGraph( + serialized_graph, + graph_hash(serialized_graph), + int(cudnn_frontend_version), + int(workspace_size), + np.asarray(input_uids, dtype=np.int64), + np.asarray(output_uids, dtype=np.int64), + np.asarray(scalar_uids, dtype=np.int64), + sizes, + values, + ) + + +def finalize_graph(cudnn, graph): + """Validate, select and compile a supported execution plan.""" + check_graph_support(cudnn, graph) + workspace_size = compile_graph(cudnn, graph) + return workspace_size, bytes(graph.serialize()), check_frontend_version(cudnn) + + +def graph_attributes(graph): + """Attributes shared by ordinary and score-modified attention.""" + return { + "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, + } diff --git a/transformer_engine/jax/cpp_extensions/flex_attention.py b/transformer_engine/jax/cpp_extensions/flex_attention.py index ff9f5fb1460..86da72c80bf 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,36 @@ import numpy as np from jax import ffi -import transformer_engine_jax +from transformer_engine.common.cudnn.attention import ( + AttentionOptions, + AttentionGraphSpec, + build_attention_graph, +) +from transformer_engine.common.cudnn.graph import ( + TensorSpec, + create_graph, + attention_matrix_spec, + row_major_stride as _row_major_stride, +) + +from transformer_engine.common.cudnn.score_mod import ( + UncacheableKey, + callback_cache_key, + explicit_cache_key, + freeze_cache_key, + wrap_score_mod as _wrap_score_mod, +) + +from .cudnn_graph import ( + SerializedGraph as _SerializedScoreModGraph, + encode_frontend_version as _encode_cudnn_frontend_version, + check_frontend_version as _check_cudnn_frontend_version_match, + graph_hash as _score_mod_graph_hash, + pack_scalar_values as _pack_score_mod_scalar_values, + serialize_graph as _serialized_score_mod_graph, + finalize_graph as _finalize_score_mod_graph, + import_cudnn as _import_cudnn_for_score_mod, +) __all__ = [ "FusedAttnScoreModHelper", @@ -110,14 +136,7 @@ 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 +_UncacheableScoreModKey = UncacheableKey def _score_mod_key_is_uncacheable(key: Any) -> bool: @@ -125,86 +144,18 @@ def _score_mod_key_is_uncacheable(key: Any) -> bool: 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 + """Freeze a structural callback key using this framework's tensor predicate.""" + return freeze_cache_key(value, _is_array_operand) 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) + """Read a callback owner's optional structural key.""" + return explicit_cache_key(callback_owner, _is_array_operand) 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() + """Share stable callbacks and distinguish module-level anonymous functions.""" + return callback_cache_key(callback, _is_array_operand, _UncacheableScoreModKey) @dataclass(frozen=True) @@ -254,32 +205,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 - - -# cuDNN frontend tensor UIDs are arbitrary, but assigning stable values makes serialized -# graphs deterministic and simplifies debugging across the Python graph builder and C++ executor. -_SCORE_MOD_UID_Q = 1 -_SCORE_MOD_UID_K = 2 -_SCORE_MOD_UID_V = 3 -_SCORE_MOD_UID_O = 4 -_SCORE_MOD_UID_STATS = 5 -_SCORE_MOD_UID_DO = 6 -_SCORE_MOD_UID_DQ = 7 -_SCORE_MOD_UID_DK = 8 -_SCORE_MOD_UID_DV = 9 +# Callback operands use a separate UID range from the shared SDPA builder. _SCORE_MOD_FWD_TENSOR_UID_BASE = 1000 _SCORE_MOD_BPROP_TENSOR_UID_BASE = 2000 _SCORE_MOD_FWD_SCALAR_UID_BASE = 3000 @@ -288,22 +214,14 @@ 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 _matrix_spec(cudnn, aval, output=False): + """Translate JAX's contiguous BSHD storage into a shared logical descriptor.""" + batch, seqlen, heads, dim = aval.shape + return attention_matrix_spec( + (batch, heads, seqlen, dim), + "bshd", + _cudnn_data_type(cudnn, aval.dtype).name, + output=output, ) @@ -477,102 +395,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, - cudnn_frontend_version: int, - workspace_size: int, - input_uids: Sequence[int], - output_uids: Sequence[int], - scalar_uids: Sequence[int], - scalar_values: Sequence[bytes], -) -> _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), - 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), - scalar_uids=np.asarray(scalar_uids, dtype=np.int64), - scalar_sizes=scalar_sizes, - scalar_values=packed_scalar_values, - ) - - -def _wrap_score_mod(score_mod: Optional[Callable], graph_tensors: Dict[str, Any]): - if score_mod is None: - return None - - def wrapped_score_mod(sdpa_graph, score_tensor): - return score_mod(sdpa_graph, score_tensor, graph_tensors) - - 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 _graph_cache_key( direction: str, config: _FusedAttnScoreModConfig, @@ -593,41 +415,10 @@ 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 - - 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, - ) - - q_dim, q_stride = _bshd_as_bhsd_dim_stride(q_aval.shape) - k_dim, k_stride = _bshd_as_bhsd_dim_stride(k_aval.shape) - v_dim, v_stride = _bshd_as_bhsd_dim_stride(v_aval.shape) - q = graph.tensor( - name="q", dim=q_dim, stride=q_stride, data_type=io_data_type, uid=_SCORE_MOD_UID_Q - ) - k = graph.tensor( - name="k", dim=k_dim, stride=k_stride, data_type=io_data_type, uid=_SCORE_MOD_UID_K - ) - v = graph.tensor( - name="v", dim=v_dim, stride=v_stride, data_type=io_data_type, uid=_SCORE_MOD_UID_V - ) - - score_mod_graph_tensors, tensor_uids, scalar_uids, scalar_values = _score_mod_graph_tensors( + graph = create_graph(cudnn, _cudnn_data_type(cudnn, q_aval.dtype)) + extra, tensor_uids, scalar_uids, scalar_values = _score_mod_graph_tensors( cudnn, graph, config.score_mod_tensor_names, @@ -636,40 +427,35 @@ def _build_score_mod_fwd_graph(q_aval, k_aval, v_aval, score_mod_avals, config): _SCORE_MOD_FWD_TENSOR_UID_BASE, _SCORE_MOD_FWD_SCALAR_UID_BASE, ) - - output, stats = graph.sdpa( - name="te_score_mod_sdpa", - q=q, - k=k, - v=v, - generate_stats=config.is_training, - attn_scale=config.scaling_factor, - use_causal_mask=False, - score_mod=_wrap_score_mod(config.score_mod, score_mod_graph_tensors), - ) - - batch, q_seqlen, q_heads, _ = q_aval.shape - _, _, _, v_head_dim = v_aval.shape - output_dim, output_stride = _bshd_as_bhsd_dim_stride((batch, q_seqlen, q_heads, v_head_dim)) - output.set_output(True).set_uid(_SCORE_MOD_UID_O).set_dim(output_dim).set_stride(output_stride) - output.set_data_type(io_data_type) - - output_uids = [_SCORE_MOD_UID_O] + output_shape = (*q_aval.shape[:-1], v_aval.shape[-1]) + specs = { + "q": _matrix_spec(cudnn, q_aval), + "k": _matrix_spec(cudnn, k_aval), + "v": _matrix_spec(cudnn, v_aval), + "o": _matrix_spec(cudnn, jax.ShapeDtypeStruct(output_shape, q_aval.dtype), output=True), + } + output_names = ["o"] if config.is_training: - stats_shape = (batch, q_heads, q_seqlen, 1) - stats.set_output(True).set_uid(_SCORE_MOD_UID_STATS).set_dim(stats_shape).set_stride( - _row_major_stride(stats_shape) - ) - 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) + batch, seqlen, heads, _ = q_aval.shape + shape = (batch, heads, seqlen, 1) + specs["stats"] = TensorSpec(shape, _row_major_stride(shape), "FLOAT", output=True) + output_names.append("stats") + built = build_attention_graph( + cudnn, + graph, + AttentionGraphSpec.from_dict( + AttentionOptions(attn_scale=config.scaling_factor, generate_stats=config.is_training), + specs, + ), + score_mod=_wrap_score_mod(config.score_mod, extra), + ) + workspace_size, data, frontend_version = _finalize_score_mod_graph(cudnn, graph) return _serialized_score_mod_graph( - serialized_graph=serialized_graph, + serialized_graph=data, cudnn_frontend_version=frontend_version, workspace_size=workspace_size, - input_uids=[_SCORE_MOD_UID_Q, _SCORE_MOD_UID_K, _SCORE_MOD_UID_V, *tensor_uids], - output_uids=output_uids, + input_uids=[built.uids[name] for name in ("q", "k", "v")] + list(tensor_uids), + output_uids=[built.uids[name] for name in output_names], scalar_uids=scalar_uids, scalar_values=scalar_values, ) @@ -687,43 +473,8 @@ def _build_score_mod_bwd_graph( 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, - ) - - q_dim, q_stride = _bshd_as_bhsd_dim_stride(q_aval.shape) - k_dim, k_stride = _bshd_as_bhsd_dim_stride(k_aval.shape) - v_dim, v_stride = _bshd_as_bhsd_dim_stride(v_aval.shape) - o_dim, o_stride = _bshd_as_bhsd_dim_stride(output_aval.shape) - do_dim, do_stride = _bshd_as_bhsd_dim_stride(doutput_aval.shape) - q = graph.tensor( - name="q", dim=q_dim, stride=q_stride, data_type=io_data_type, uid=_SCORE_MOD_UID_Q - ) - k = graph.tensor( - name="k", dim=k_dim, stride=k_stride, data_type=io_data_type, uid=_SCORE_MOD_UID_K - ) - v = graph.tensor( - name="v", dim=v_dim, stride=v_stride, data_type=io_data_type, uid=_SCORE_MOD_UID_V - ) - output = graph.tensor( - name="o", dim=o_dim, stride=o_stride, data_type=io_data_type, uid=_SCORE_MOD_UID_O - ) - doutput = graph.tensor( - name="dO", dim=do_dim, stride=do_stride, data_type=io_data_type, uid=_SCORE_MOD_UID_DO - ) - stats = graph.tensor( - name="stats", - dim=tuple(int(dim) for dim in stats_aval.shape), - stride=_row_major_stride(stats_aval.shape), - data_type=cudnn.data_type.FLOAT, - uid=_SCORE_MOD_UID_STATS, - ) - - score_mod_graph_tensors, tensor_uids, scalar_uids, scalar_values = _score_mod_graph_tensors( + graph = create_graph(cudnn, _cudnn_data_type(cudnn, q_aval.dtype)) + extra, tensor_uids, scalar_uids, scalar_values = _score_mod_graph_tensors( cudnn, graph, config.score_mod_tensor_names, @@ -732,56 +483,49 @@ def _build_score_mod_bwd_graph( _SCORE_MOD_FWD_TENSOR_UID_BASE, _SCORE_MOD_FWD_SCALAR_UID_BASE, ) - ( - score_mod_bprop_graph_tensors, - bprop_tensor_uids, - bprop_scalar_uids, - bprop_scalar_values, - ) = _score_mod_graph_tensors( + bprop_extra, bprop_tensor_uids, bprop_scalar_uids, bprop_scalar_values = ( + _score_mod_graph_tensors( + cudnn, + graph, + config.score_mod_bprop_tensor_names, + score_mod_bprop_avals, + config.score_mod_bprop_scalars, + _SCORE_MOD_BPROP_TENSOR_UID_BASE, + _SCORE_MOD_BPROP_SCALAR_UID_BASE, + ) + ) + specs = { + "q": _matrix_spec(cudnn, q_aval), + "k": _matrix_spec(cudnn, k_aval), + "v": _matrix_spec(cudnn, v_aval), + "o": _matrix_spec(cudnn, output_aval), + "do": _matrix_spec(cudnn, doutput_aval), + "stats": TensorSpec(tuple(stats_aval.shape), _row_major_stride(stats_aval.shape), "FLOAT"), + "dq": _matrix_spec(cudnn, q_aval, output=True), + "dk": _matrix_spec(cudnn, k_aval, output=True), + "dv": _matrix_spec(cudnn, v_aval, output=True), + } + built = build_attention_graph( cudnn, graph, - config.score_mod_bprop_tensor_names, - score_mod_bprop_avals, - config.score_mod_bprop_scalars, - _SCORE_MOD_BPROP_TENSOR_UID_BASE, - _SCORE_MOD_BPROP_SCALAR_UID_BASE, - ) - - dq, dk, dv = graph.sdpa_backward( - name="te_score_mod_sdpa_backward", - q=q, - k=k, - v=v, - o=output, - dO=doutput, - stats=stats, - attn_scale=config.scaling_factor, - use_causal_mask=False, - score_mod=_wrap_score_mod(config.score_mod, score_mod_graph_tensors), - score_mod_bprop=_wrap_score_mod(config.score_mod_bprop, score_mod_bprop_graph_tensors), - use_deterministic_algorithm=config.deterministic, - ) - - dq.set_output(True).set_uid(_SCORE_MOD_UID_DQ).set_dim(q_dim).set_stride(q_stride) - 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) + AttentionGraphSpec.from_dict( + AttentionOptions( + backward=True, attn_scale=config.scaling_factor, deterministic=config.deterministic + ), + specs, + ), + score_mod=_wrap_score_mod(config.score_mod, extra), + score_mod_bprop=_wrap_score_mod(config.score_mod_bprop, bprop_extra), + ) + workspace_size, data, frontend_version = _finalize_score_mod_graph(cudnn, graph) return _serialized_score_mod_graph( - serialized_graph=serialized_graph, + serialized_graph=data, cudnn_frontend_version=frontend_version, workspace_size=workspace_size, - input_uids=[ - _SCORE_MOD_UID_Q, - _SCORE_MOD_UID_K, - _SCORE_MOD_UID_V, - _SCORE_MOD_UID_O, - _SCORE_MOD_UID_DO, - _SCORE_MOD_UID_STATS, - *tensor_uids, - *bprop_tensor_uids, - ], - output_uids=[_SCORE_MOD_UID_DQ, _SCORE_MOD_UID_DK, _SCORE_MOD_UID_DV], + input_uids=[built.uids[name] for name in ("q", "k", "v", "o", "do", "stats")] + + list(tensor_uids) + + list(bprop_tensor_uids), + output_uids=[built.uids[name] for name in ("dq", "dk", "dv")], scalar_uids=[*scalar_uids, *bprop_scalar_uids], scalar_values=[*scalar_values, *bprop_scalar_values], ) @@ -813,7 +557,7 @@ def _fused_attn_score_mod_fwd( stats = jax.ShapeDtypeStruct(stats_shape, jnp.float32) workspace = jax.ShapeDtypeStruct((graph.workspace_size,), jnp.uint8) output, softmax_stats, _ = ffi.ffi_call( - "te_fused_attn_score_mod_forward_ffi", + "te_cudnn_graph_ffi", (output_shape, stats, workspace), )( q, @@ -873,7 +617,7 @@ def _fused_attn_score_mod_bwd( dv = jax.ShapeDtypeStruct(v.shape, v.dtype) workspace = jax.ShapeDtypeStruct((graph.workspace_size,), jnp.uint8) dq, dk, dv, _ = ffi.ffi_call( - "te_fused_attn_score_mod_backward_ffi", + "te_cudnn_graph_ffi", (dq, dk, dv, workspace), )( q, diff --git a/transformer_engine/jax/csrc/extensions.h b/transformer_engine/jax/csrc/extensions.h index 2237b5c204e..2636a4b306d 100644 --- a/transformer_engine/jax/csrc/extensions.h +++ b/transformer_engine/jax/csrc/extensions.h @@ -150,32 +150,8 @@ XLA_FFI_DECLARE_HANDLER_SYMBOL(ScaledUpperTriangMaskedSoftmaxBackwardHandler); XLA_FFI_DECLARE_HANDLER_SYMBOL(ScaledUpperTriangMaskedSoftmaxBackwardInitializeHandler); // Attention -XLA_FFI_DECLARE_HANDLER_SYMBOL(FusedAttnForwardHandler); - -XLA_FFI_DECLARE_HANDLER_SYMBOL(FusedAttnBackwardHandler); - -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); +XLA_FFI_DECLARE_HANDLER_SYMBOL(CudnnGraphHandler); +XLA_FFI_DECLARE_HANDLER_SYMBOL(CudnnGraphInitializeHandler); // 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 deleted file mode 100644 index def5e0149b8..00000000000 --- a/transformer_engine/jax/csrc/extensions/attention.cpp +++ /dev/null @@ -1,1115 +0,0 @@ -/************************************************************************* - * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. - * - * See LICENSE for license information. - ************************************************************************/ - -#include - -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include - -#include "../extensions.h" -#include "common/cudnn_utils.h" -#include "transformer_engine/fused_attn.h" -#include "transformer_engine/transformer_engine.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 = is_ragged ? input_batch * max_segments_per_seq : input_batch; - 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) { \ - num_segments = input_batch * max_segments_per_seq; \ - } \ - 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 = is_ragged ? input_batch * max_segments_per_seq : input_batch; - - 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 { - alignas(16) std::array data{}; - size_t size = 0; -}; - -struct ScoreModGraphCacheKey { - int device_id = 0; - int64_t hash0 = 0; - int64_t hash1 = 0; - int64_t frontend_version = 0; - - bool operator==(const ScoreModGraphCacheKey &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 { - 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); - combine(key.hash1); - combine(key.frontend_version); - return seed; - } -}; - -using ScoreModGraphPtr = std::shared_ptr; - -std::unordered_map & -getScoreModeGraphCache() { - static std::unordered_map - cache; - return cache; -} - -std::mutex &getScoreModGraphCacheMutex() { - static std::mutex mutex; - return mutex; -} - -ScoreModGraphCacheKey GetScoreModGraphCacheKey(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, "."); - - int device_id = 0; - NVTE_CHECK_CUDA(cudaGetDevice(&device_id)); - return ScoreModGraphCacheKey{ - device_id, - get_attr_value(attrs, "graph_hash0"), - get_attr_value(attrs, "graph_hash1"), - frontend_version, - }; -} - -ScoreModGraphPtr GetScoreModGraph(cudaStream_t stream, Dictionary &attrs) { - const auto key = GetScoreModGraphCacheKey(attrs); - { - std::lock_guard lock(getScoreModGraphCacheMutex()); - auto &cache = getScoreModeGraphCache(); - auto it = cache.find(key); - if (it != cache.end()) { - return it->second; - } - } - - const auto serialized_graph = get_attr_value(attrs, "serialized_graph"); - std::vector serialized_data(serialized_graph.begin(), serialized_graph.end()); - - auto handle = nvte_get_cudnn_handle(); - 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()); - - std::lock_guard lock(getScoreModGraphCacheMutex()); - auto &cache = getScoreModeGraphCache(); - auto it = cache.find(key); - if (it != cache.end()) { - return it->second; - } - cache.emplace(key, graph); - 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); - auto input_uids = get_attr_value>(attrs, "input_uids"); - auto output_uids = get_attr_value>(attrs, "output_uids"); - 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(scalar_uids.size() == scalar_sizes.size(), - "Mismatched score_mod scalar uid/value-size counts."); - NVTE_CHECK(scalar_values.size() == scalar_uids.size() * 16, - "Mismatched score_mod 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]); - } - for (size_t i = 0; i < output_uids.size(); ++i) { - variant_pack.emplace(output_uids[i], output_ptrs[i]); - } - - 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]); - 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 = nvte_get_cudnn_handle(); - NVTE_CHECK_CUDNN(cudnnSetStream(handle, stream)); - auto status = graph->execute(handle, variant_pack, workspace); - NVTE_CHECK(status.is_good(), - "cuDNN score_mod SDPA graph execution failed: ", status.get_message()); - return ffi_with_cuda_error_check(); -} - -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."); - ptrs->push_back(maybe_buf.value().untyped_data()); - } -} - -} // namespace - -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, - Result_Type workspace_buf, Dictionary attrs) { - 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()); -} - -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()); - -Error_Type FusedAttnScoreModBackwardFFI(cudaStream_t stream, Buffer_Type q_buf, Buffer_Type k_buf, - Buffer_Type v_buf, Buffer_Type output_buf, - Buffer_Type doutput_buf, Buffer_Type stats_buf, - Variadic_Buffer_Type score_mod_args, Result_Type dq_buf, - Result_Type dk_buf, Result_Type dv_buf, - Result_Type workspace_buf, Dictionary attrs) { - std::vector input_ptrs = {q_buf.untyped_data(), k_buf.untyped_data(), - 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()); -} - -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()); - -} // namespace jax -} // namespace transformer_engine diff --git a/transformer_engine/jax/csrc/extensions/cudnn_graph.cpp b/transformer_engine/jax/csrc/extensions/cudnn_graph.cpp new file mode 100644 index 00000000000..40b9fa7cdee --- /dev/null +++ b/transformer_engine/jax/csrc/extensions/cudnn_graph.cpp @@ -0,0 +1,196 @@ +/************************************************************************* + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * + * See LICENSE for license information. + ************************************************************************/ + +// JAX-owned execution of generic graphs built with the cuDNN Python frontend. +// This translation unit deliberately has no TE common dependencies. +#include +#include +#include +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include + +namespace transformer_engine { +namespace jax { +namespace { + +namespace ffi = xla::ffi; + +void CheckCuda(cudaError_t status) { + if (status != cudaSuccess) throw std::runtime_error(cudaGetErrorString(status)); +} + +void CheckCudnn(cudnnStatus_t status) { + if (status != CUDNN_STATUS_SUCCESS) throw std::runtime_error(cudnnGetErrorString(status)); +} + +template +T Attribute(ffi::Dictionary &attrs, const char *name) { + auto value = attrs.get(name); + if (value.has_error()) { + throw std::runtime_error(std::string("Missing cuDNN graph attribute: ") + name); + } + return value.value(); +} + +struct Handle { + cudnnHandle_t value = nullptr; + Handle() { CheckCudnn(cudnnCreate(&value)); } + ~Handle() { cudnnDestroy(value); } +}; + +cudnnHandle_t GetHandle() { + // Handles are thread-local and device-local; assign the stream on each call. + thread_local std::unordered_map> handles; + int device; + CheckCuda(cudaGetDevice(&device)); + auto &handle = handles[device]; + if (!handle) handle = std::make_unique(); + return handle->value; +} + +using Graph = cudnn_frontend::graph::Graph; + +std::shared_ptr GetGraph(cudnnHandle_t handle, ffi::Dictionary &attrs) { + const auto version = Attribute(attrs, "cudnn_frontend_version"); + if (version != CUDNN_FRONTEND_VERSION) { + throw std::runtime_error( + "cuDNN frontend Python/C++ version mismatch for graph deserialization."); + } + int device; + CheckCuda(cudaGetDevice(&device)); + const auto key = std::to_string(device) + ":" + + std::to_string(Attribute(attrs, "graph_hash0")) + ":" + + std::to_string(Attribute(attrs, "graph_hash1")) + ":" + + std::to_string(version); + static std::mutex mutex; + static std::unordered_map> cache; + std::lock_guard lock(mutex); + auto &graph = cache[key]; + if (!graph) { + const auto data = Attribute(attrs, "serialized_graph"); + auto candidate = std::make_shared(); + auto status = candidate->deserialize(handle, std::vector(data.begin(), data.end())); + if (!status.is_good()) throw std::runtime_error(status.get_message()); + graph = std::move(candidate); + } + return graph; +} + +struct Scalar { + alignas(16) std::array value{}; +}; + +ffi::Error InitializeGraph(cudaStream_t stream, ffi::RemainingArgs, ffi::RemainingRets, + ffi::Dictionary attrs) { + try { + auto handle = GetHandle(); + CheckCudnn(cudnnSetStream(handle, stream)); + GetGraph(handle, attrs); + return ffi::Error::Success(); + } catch (const std::exception &error) { + return ffi::Error::Internal(std::string("JAX cuDNN attention initialization: ") + error.what()); + } +} + +ffi::Error ExecuteGraph(cudaStream_t stream, ffi::RemainingArgs args, ffi::RemainingRets rets, + ffi::Dictionary attrs) { + try { + auto handle = GetHandle(); + CheckCudnn(cudnnSetStream(handle, stream)); + auto graph = GetGraph(handle, attrs); + auto input_uids = Attribute>(attrs, "input_uids"); + auto output_uids = Attribute>(attrs, "output_uids"); + auto input_indices = attrs.get>("input_indices"); + auto input_offsets = attrs.get>("input_offsets"); + auto output_indices = attrs.get>("output_indices"); + auto output_offsets = attrs.get>("output_offsets"); + auto output_fill_values = attrs.get>("output_fill_values"); + if (rets.size() == 0 || + (output_fill_values.has_value() && output_fill_values.value().size() != rets.size() - 1)) { + throw std::runtime_error("Invalid cuDNN output initialization bindings."); + } + std::unordered_map pointers; + for (size_t i = 0; i < input_uids.size(); ++i) { + const auto index = input_indices.has_value() ? input_indices.value()[i] : i; + auto buf = args.get(index); + if (buf.has_error()) throw std::runtime_error("Invalid cuDNN input buffer binding."); + auto ptr = static_cast(buf.value().untyped_data()); + pointers.emplace(input_uids[i], + ptr + (input_offsets.has_value() ? input_offsets.value()[i] : 0)); + } + for (size_t i = 0; i < rets.size() - 1; ++i) { + auto buf = rets.get(i); + if (buf.has_error()) throw std::runtime_error("Invalid cuDNN output buffer binding."); + // Initialize ragged slots and optional gradient outputs. + if (buf.value()->size_bytes()) { + const int fill = output_fill_values.has_value() ? output_fill_values.value()[i] : 0; + CheckCuda( + cudaMemsetAsync(buf.value()->untyped_data(), fill, buf.value()->size_bytes(), stream)); + } + } + for (size_t i = 0; i < output_uids.size(); ++i) { + if (output_uids[i] < 0) continue; + const auto index = output_indices.has_value() ? output_indices.value()[i] : i; + auto buf = rets.get(index); + if (buf.has_error()) throw std::runtime_error("Invalid cuDNN output buffer binding."); + auto ptr = static_cast(buf.value()->untyped_data()); + pointers.emplace(output_uids[i], + ptr + (output_offsets.has_value() ? output_offsets.value()[i] : 0)); + } + auto scalar_uids = Attribute>(attrs, "scalar_uids"); + auto scalar_sizes = Attribute>(attrs, "scalar_sizes"); + auto scalar_values = Attribute>(attrs, "scalar_values"); + if (scalar_sizes.size() != scalar_uids.size() || + scalar_values.size() != scalar_uids.size() * 16) { + throw std::runtime_error("Invalid cuDNN pass-by-value scalar bindings."); + } + std::vector scalars(scalar_uids.size()); + for (size_t i = 0; i < scalars.size(); ++i) { + if (scalar_sizes[i] < 0 || scalar_sizes[i] > 16) { + throw std::runtime_error("Invalid cuDNN pass-by-value scalar size."); + } + std::copy_n(scalar_values.begin() + i * 16, 16, scalars[i].value.begin()); + pointers.emplace(scalar_uids[i], scalars[i].value.data()); + } + auto workspace = rets.get(rets.size() - 1); + if (workspace.has_error()) throw std::runtime_error("Missing cuDNN graph workspace."); + auto status = graph->execute(handle, pointers, workspace.value()->untyped_data()); + if (!status.is_good()) throw std::runtime_error(status.get_message()); + CheckCuda(cudaGetLastError()); + return ffi::Error::Success(); + } catch (const std::exception &error) { + return ffi::Error::Internal(std::string("JAX cuDNN attention: ") + error.what()); + } +} + +} // namespace + +XLA_FFI_DEFINE_HANDLER_SYMBOL(CudnnGraphHandler, ExecuteGraph, + xla::ffi::Ffi::Bind() + .Ctx>() + .RemainingArgs() + .RemainingRets() + .Attrs(), + {xla::ffi::Traits::kCmdBufferCompatible}); + +XLA_FFI_DEFINE_HANDLER_SYMBOL(CudnnGraphInitializeHandler, InitializeGraph, + xla::ffi::Ffi::Bind() + .Ctx>() + .RemainingArgs() + .RemainingRets() + .Attrs()); + +} // namespace jax +} // namespace transformer_engine diff --git a/transformer_engine/jax/csrc/extensions/pybind.cpp b/transformer_engine/jax/csrc/extensions/pybind.cpp index 69f304ec999..308b80e29de 100644 --- a/transformer_engine/jax/csrc/extensions/pybind.cpp +++ b/transformer_engine/jax/csrc/extensions/pybind.cpp @@ -69,19 +69,10 @@ pybind11::dict Registrations() { pybind11::arg("initialize") = EncapsulateFFI(NormBackwardInitializeHandler), pybind11::arg("execute") = EncapsulateFFI(NormBackwardHandler)); - // Attention - dict["te_fused_attn_forward_ffi"] = - pybind11::dict(pybind11::arg("prepare") = EncapsulateFFI(CudnnHandleInitHandler), - pybind11::arg("execute") = EncapsulateFFI(FusedAttnForwardHandler)); - dict["te_fused_attn_backward_ffi"] = - pybind11::dict(pybind11::arg("prepare") = EncapsulateFFI(CudnnHandleInitHandler), - pybind11::arg("execute") = EncapsulateFFI(FusedAttnBackwardHandler)); - dict["te_fused_attn_score_mod_forward_ffi"] = - pybind11::dict(pybind11::arg("prepare") = EncapsulateFFI(CudnnHandleInitHandler), - pybind11::arg("execute") = EncapsulateFFI(FusedAttnScoreModForwardHandler)); - dict["te_fused_attn_score_mod_backward_ffi"] = - pybind11::dict(pybind11::arg("prepare") = EncapsulateFFI(CudnnHandleInitHandler), - pybind11::arg("execute") = EncapsulateFFI(FusedAttnScoreModBackwardHandler)); + // Python-built cuDNN graphs use a JAX-owned executor and handle cache. + dict["te_cudnn_graph_ffi"] = + pybind11::dict(pybind11::arg("initialize") = EncapsulateFFI(CudnnGraphInitializeHandler), + pybind11::arg("execute") = EncapsulateFFI(CudnnGraphHandler)); // GEMM dict["te_gemm_ffi"] = @@ -157,8 +148,6 @@ 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); @@ -169,10 +158,7 @@ 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); m.def("nvte_built_with_cublasmp", &::nvte_built_with_cublasmp); m.def("initialize_cgemm_communicator", &InitializeCgemmCommunicator); @@ -208,7 +194,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) diff --git a/transformer_engine/jax/flax/transformer.py b/transformer_engine/jax/flax/transformer.py index 5a01572f1fa..453ceab49ed 100644 --- a/transformer_engine/jax/flax/transformer.py +++ b/transformer_engine/jax/flax/transformer.py @@ -310,6 +310,7 @@ class _FusedDotProductAttention(nn.Module): # pylint: disable=too-few-public-me softmax_type: AttnSoftmaxType = AttnSoftmaxType.VANILLA_SOFTMAX score_mod: Optional[Callable] = None score_mod_bprop: Optional[Callable] = None + fp8: bool = False score_mod_requested: bool = False return_max_logit: bool = False @@ -369,6 +370,8 @@ def __call__( "score_mod_bprop_tensors": score_mod_bprop_tensors, "return_max_logit": self.return_max_logit, } + if self.fp8: + fused_attn_kwargs["fp8"] = True if self.qkv_layout.is_qkvpacked(): """qkvpacked format, treat @@ -632,6 +635,10 @@ 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. + fp8: bool, default = False + Use cuDNN FP8 attention with current tensor scaling. Inputs, outputs and input + gradients remain FP16/BF16. Fused attention is required, and this option is + independent of the GEMM quantization recipe. Optimization parameters ----------------------- @@ -661,6 +668,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 + fp8: bool = False def __post_init__(self): # TODO(KshitijLakhani): Remove warning in TransformerEngine v2.12 @@ -761,7 +769,7 @@ 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}" @@ -791,6 +799,8 @@ def __call__( # Use fused attn (if kernel check below passes) by default enable_fused_attn = int(os.getenv("NVTE_FUSED_ATTN", "1")) + if self.fp8 and not enable_fused_attn: + raise ValueError("FP8 attention requires fused attention, but NVTE_FUSED_ATTN=0.") if self.return_max_logit and not enable_fused_attn: raise ValueError("return_max_logit requires fused attention, but NVTE_FUSED_ATTN=0.") @@ -832,7 +842,7 @@ def __call__( softmax_type=softmax_type, dropout_probability=self.attention_dropout, q_num_heads=self.num_attention_heads, - kv_num_heads=self.num_gqa_groups, + kv_num_heads=self.num_gqa_groups or self.num_attention_heads, q_max_seqlen=seqlen_q, kv_max_seqlen=seqlen_kv, head_dim_qk=head_dim_qk, @@ -845,9 +855,16 @@ def __call__( bias_seqlen_q=bias_seqlen_q, bias_seqlen_kv=bias_seqlen_kv, max_segments_per_seq=self.max_segments_per_seq, + fp8=self.fp8, ) 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 attn_bias_type == AttnBiasType.ALIBI and ( + not enable_fused_attn or not has_fused_attn_kernel + ): + raise ValueError("ALiBi requires an available cuDNN fused attention kernel.") + if self.fp8 and not has_fused_attn_kernel: + raise ValueError("No cuDNN FP8 attention kernel supports this configuration.") if score_mod_requested and not has_fused_attn_kernel: raise ValueError( "score_mod requires fused attention, but no fused attention kernel is available." @@ -949,6 +966,7 @@ def __call__( softmax_type=softmax_type, score_mod=self.score_mod, score_mod_bprop=self.score_mod_bprop, + fp8=self.fp8, score_mod_requested=score_mod_requested, return_max_logit=self.return_max_logit, )( @@ -1272,6 +1290,9 @@ class MultiHeadAttention(nn.Module): # pylint: disable=too-few-public-methods argument to keep tensor operands as normal JAX inputs. score_mod_bprop_tensors: Optional[Mapping[str, Any]], default = None Additional tensors or pass-by-value scalars for ``score_mod_bprop``. + fp8: bool, default = False + Use cuDNN FP8 dot product attention with current tensor scaling. Inputs and + outputs retain their FP16/BF16 dtype; the GEMM quantization recipe is independent. """ head_dim: int @@ -1314,6 +1335,8 @@ class MultiHeadAttention(nn.Module): # pylint: disable=too-few-public-methods apply_residual_connection_post_layernorm: Optional[bool] = None fuse_qkv: Optional[bool] = None + fp8: bool = False + def __post_init__(self): # Deal with changed defaults in API # TODO(KshitijLakhani): Remove warning in TransformerEngine v2.12 @@ -1735,6 +1758,7 @@ def generate_batch_seqlen_logical_axes(is_sharded_seq): softmax_type=self.softmax_type, score_mod=self.score_mod, score_mod_bprop=self.score_mod_bprop, + fp8=self.fp8, )( *dpa_args, mask, @@ -2061,6 +2085,10 @@ class TransformerLayer(nn.Module): # pylint: disable=too-few-public-methods argument to keep tensor operands as normal JAX inputs. score_mod_bprop_tensors: Optional[Mapping[str, Any]], default = None Additional tensors or pass-by-value scalars for ``score_mod_bprop``. + fp8: bool, default = False + Use cuDNN FP8 attention in the self-attention and encoder-decoder blocks. + Attention uses current tensor scaling and retains FP16/BF16 inputs and outputs, + independently of the GEMM quantization recipe. Optimization parameters ----------------------- @@ -2129,6 +2157,7 @@ class TransformerLayer(nn.Module): # pylint: disable=too-few-public-methods softmax_type: str = "vanilla" score_mod: Optional[Callable] = None score_mod_bprop: Optional[Callable] = None + fp8: bool = False def __post_init__(self): if self.mha_kernel_init is None: @@ -2300,6 +2329,7 @@ def generate_batch_seqlen_logical_axes(is_shared_seq=None): softmax_type=self.softmax_type, score_mod=self.score_mod, score_mod_bprop=self.score_mod_bprop, + fp8=self.fp8, )( inputs, inputs, @@ -2395,6 +2425,7 @@ def hidden_dropout(x, deterministic): name="encoder_decoder_attention", window_size=self.window_size, softmax_type=self.softmax_type, + fp8=self.fp8, )(x, encoded, encoder_decoder_mask, deterministic=deterministic) y = with_sharding_constraint_by_logical_axes( diff --git a/transformer_engine/pytorch/attention/dot_product_attention/backends.py b/transformer_engine/pytorch/attention/dot_product_attention/backends.py index bccb4f1e26f..c78e9794cb2 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/backends.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/backends.py @@ -18,6 +18,7 @@ import torch import torch.nn.functional as F +from transformer_engine.common.cudnn.graph import uses_ragged_attention_graph from transformer_engine.pytorch.utils import ( get_device_compute_capability, split_tensor_along_dim, @@ -2442,7 +2443,7 @@ def _fused_attn_stats_shape( if q_format == "thd": num_heads = q_shape[1] sm = get_device_compute_capability() - if sm >= (9, 0) and sm != (12, 0): + if uses_ragged_attention_graph(10 * sm[0] + sm[1]): return (q_shape[0], num_heads, 1) batch_size = args.cu_seqlens_q.shape[0] - 1 return (batch_size, num_heads, args.max_seqlen_q, 1) @@ -2600,13 +2601,16 @@ def _needs_eager_fused_attention(call: Dict[str, Any]) -> Optional[str]: class FusedAttention(torch.nn.Module): """Dot product attention using `cuDNN attention `_: + PyTorch constructs and executes attention graphs through the cuDNN frontend Python API. + Standard attention and experimental score modification share graph lifecycle helpers. + Plans are cached per device and tensor metadata; support queries do not compile plans. + 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. + are implemented in transformer_engine.pytorch.cpp_extensions.cudnn_fused_attn. Direction Inputs Intermediates Outputs DelayedScaling (DS) forward Q/K/V: DS S: DS O: DS 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 3a153a3f994..bf924de722f 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py @@ -8,6 +8,8 @@ import torch import transformer_engine_torch as tex +from transformer_engine.common.cudnn.graph import uses_ragged_attention_graph + from transformer_engine.pytorch.utils import ( get_cudnn_version, nvtx_range_pop, @@ -1837,7 +1839,10 @@ def forward( softmax_lse_in_packed_format = False if qkv_format == "thd": if use_fused_attention: - softmax_lse_in_packed_format = get_device_compute_capability() != (12, 0) + capability = get_device_compute_capability() + softmax_lse_in_packed_format = uses_ragged_attention_graph( + 10 * capability[0] + capability[1] + ) else: softmax_lse_in_packed_format = ( fa_utils.v2_6_0_plus or use_flash_attn_3 or use_flash_attn_4 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..54a72a31aa8 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/flex_attention.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/flex_attention.py @@ -5,28 +5,35 @@ """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.cudnn.attention import ( + AttentionOptions, + AttentionGraphSpec, + build_attention_graph, +) +from transformer_engine.common.cudnn.graph import TensorSpec + +from transformer_engine.common.cudnn.score_mod import ( + callback_cache_key, + explicit_cache_key, + freeze_cache_key, + wrap_score_mod as _wrap_score_mod, +) + +from transformer_engine.pytorch.cpp_extensions.cudnn import ( + _import_cudnn_frontend, + _build_cudnn_pygraph, + _finalize_cudnn_graph, + _execute_cudnn_graph, +) + _cudnn_score_mod_graph_cache: Dict[Tuple[Any, ...], Any] = {} _SCORE_MOD_UNCACHEABLE = object() -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 - - def _bhsd_dim_stride( tensor: torch.Tensor, tensor_format: str ) -> Tuple[Tuple[int, ...], Tuple[int, ...]]: @@ -44,98 +51,35 @@ def _bhsd_dim_stride( raise ValueError(f"Flex Attention only supports SBHD/BSHD tensor formats, got {tensor_format}.") -def _bhsd_graph_tensor(graph, tensor: torch.Tensor, tensor_format: str): - """Create a cuDNN graph tensor with BHSD dims and TE-layout strides.""" - dim, stride = _bhsd_dim_stride(tensor, tensor_format) - return graph.tensor(dim=dim, stride=stride, data_type=tensor.dtype) +def _graph_tensor_spec(tensor, tensor_format=None, output=False): + """Translate framework metadata into the shared graph descriptor.""" + dim, stride = ( + _bhsd_dim_stride(tensor, tensor_format) + if tensor_format is not None + else (tuple(tensor.shape), tuple(tensor.stride())) + ) + dtype = {torch.float16: "HALF", torch.bfloat16: "BFLOAT16", torch.float32: "FLOAT"}[ + tensor.dtype + ] + return TensorSpec(dim, stride, dtype, output=output) # 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 + """Freeze a structural callback key using this framework's tensor predicate.""" + return freeze_cache_key(value, lambda value: isinstance(value, torch.Tensor)) 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) + """Read a callback owner's optional structural key.""" + return explicit_cache_key(callback_owner, lambda value: isinstance(value, torch.Tensor)) def _score_mod_callback_cache_key(callback: Optional[Callable]) -> Any: - """Create a stable graph cache key for a score_mod callable. - - 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 + """Share stable callbacks and distinguish module-level anonymous functions.""" + return callback_cache_key( + callback, lambda value: isinstance(value, torch.Tensor), lambda: _SCORE_MOD_UNCACHEABLE + ) def _score_mod_device_key(device: torch.device) -> Tuple[Any, ...]: @@ -181,55 +125,6 @@ def _make_cudnn_graph_tensor_dict(graph, tensors: Optional[Dict[str, torch.Tenso # cuDNN frontend score_mod graph helpers. -def _wrap_score_mod(score_mod: Optional[Callable], graph_tensors: Dict[str, Any]): - """Adapt TE's score_mod signature to cuDNN frontend's two-argument callback.""" - if score_mod is None: - return None - - def _wrapped_score_mod(sdpa_graph, score_tensor): - return score_mod(sdpa_graph, score_tensor, graph_tensors) - - return _wrapped_score_mod - - -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 - - -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: - 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 - - @dataclass class _CudnnScoreModFwdGraphEntry: """Cached cuDNN frontend graph and graph tensor handles for score_mod fprop.""" @@ -263,44 +158,6 @@ 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 _execute_cudnn_graph( - graph, - variant_pack: Dict[Any, torch.Tensor], - workspace_size: int, - device: torch.device, -): - """Execute a built cuDNN frontend Python graph.""" - cudnn = _import_cudnn_frontend() - - if device.type == "cuda" and device.index is None: - device = torch.device("cuda", torch.cuda.current_device()) - workspace = torch.empty( - workspace_size, - device=device, - dtype=torch.uint8, - ) - graph.execute( - variant_pack, - workspace, - handle=_get_cudnn_current_stream_handle(cudnn, device), - ) - - def _cudnn_score_mod_fwd_cache_key( is_training: bool, query_layer: torch.Tensor, @@ -395,43 +252,33 @@ def _build_cudnn_score_mod_fwd_graph( cudnn = _import_cudnn_frontend() graph = _build_cudnn_pygraph(query_layer.dtype, query_layer.device) - q = _bhsd_graph_tensor(graph, query_layer, q_format) - k = _bhsd_graph_tensor(graph, key_layer, kv_format) - v = _bhsd_graph_tensor(graph, value_layer, kv_format) - - score_mod_graph_tensors = _make_cudnn_graph_tensor_dict(graph, score_mod_tensors) - wrapped_score_mod = _wrap_score_mod(score_mod, score_mod_graph_tensors) - - output_dim, output_stride = _bhsd_dim_stride(output_layer, q_format) - output, stats_tensor = graph.sdpa( - name="te_score_mod_sdpa", - q=q, - k=k, - v=v, - generate_stats=is_training, - attn_scale=attn_scale, - use_causal_mask=False, - score_mod=wrapped_score_mod, - ) - output.set_output(True).set_dim(output_dim).set_stride(output_stride) - + extra = _make_cudnn_graph_tensor_dict(graph, score_mod_tensors) + specs = { + "q": _graph_tensor_spec(query_layer, q_format), + "k": _graph_tensor_spec(key_layer, kv_format), + "v": _graph_tensor_spec(value_layer, kv_format), + "o": _graph_tensor_spec(output_layer, q_format, output=True), + } if is_training: assert stats is not None - stats_tensor.set_output(True).set_dim(stats.size()).set_stride( - stats.stride() - ).set_data_type(cudnn.data_type.FLOAT) - else: - stats_tensor = None - + specs["stats"] = _graph_tensor_spec(stats, output=True) + built = build_attention_graph( + cudnn, + graph, + AttentionGraphSpec.from_dict( + AttentionOptions(attn_scale=attn_scale, generate_stats=is_training), specs + ), + score_mod=_wrap_score_mod(score_mod, extra), + ) workspace_size = _finalize_cudnn_graph(graph) return _CudnnScoreModFwdGraphEntry( graph=graph, - q=q, - k=k, - v=v, - output=output, - stats=stats_tensor, - score_mod_graph_tensors=score_mod_graph_tensors, + q=built.tensors["q"], + k=built.tensors["k"], + v=built.tensors["v"], + output=built.tensors["o"], + stats=built.tensors.get("stats"), + score_mod_graph_tensors=extra, workspace_size=workspace_size, ) @@ -490,61 +337,49 @@ def _build_cudnn_score_mod_bwd_graph( deterministic: bool, ) -> _CudnnScoreModBwdGraphEntry: """Build a cached cuDNN frontend graph for score_mod bprop.""" + cudnn = _import_cudnn_frontend() graph = _build_cudnn_pygraph(query_layer.dtype, query_layer.device) - q = _bhsd_graph_tensor(graph, query_layer, q_format) - k = _bhsd_graph_tensor(graph, key_layer, kv_format) - v = _bhsd_graph_tensor(graph, value_layer, kv_format) - output = _bhsd_graph_tensor(graph, output_layer, q_format) - d_output = _bhsd_graph_tensor(graph, d_out, q_format) - stats_tensor = graph.tensor_like(stats) - - score_mod_graph_tensors = _make_cudnn_graph_tensor_dict(graph, score_mod_tensors) - score_mod_bprop_graph_tensors = ( + extra = _make_cudnn_graph_tensor_dict(graph, score_mod_tensors) + bprop_extra = ( _make_cudnn_graph_tensor_dict(graph, score_mod_bprop_tensors) if score_mod_bprop is not None else {} ) - wrapped_score_mod = _wrap_score_mod(score_mod, score_mod_graph_tensors) - wrapped_score_mod_bprop = _wrap_score_mod(score_mod_bprop, score_mod_bprop_graph_tensors) - - dq_layer = torch.empty_like(query_layer) - dk_layer = torch.empty_like(key_layer) - dv_layer = torch.empty_like(value_layer) - dq_dim, dq_stride = _bhsd_dim_stride(dq_layer, q_format) - dk_dim, dk_stride = _bhsd_dim_stride(dk_layer, kv_format) - dv_dim, dv_stride = _bhsd_dim_stride(dv_layer, kv_format) - dq, dk, dv = graph.sdpa_backward( - name="te_score_mod_sdpa_backward", - q=q, - k=k, - v=v, - o=output, - dO=d_output, - stats=stats_tensor, - attn_scale=attn_scale, - use_causal_mask=False, - score_mod=wrapped_score_mod, - score_mod_bprop=wrapped_score_mod_bprop, - use_deterministic_algorithm=deterministic, + specs = { + "q": _graph_tensor_spec(query_layer, q_format), + "k": _graph_tensor_spec(key_layer, kv_format), + "v": _graph_tensor_spec(value_layer, kv_format), + "o": _graph_tensor_spec(output_layer, q_format), + "do": _graph_tensor_spec(d_out, q_format), + "stats": _graph_tensor_spec(stats), + "dq": _graph_tensor_spec(torch.empty_like(query_layer), q_format, output=True), + "dk": _graph_tensor_spec(torch.empty_like(key_layer), kv_format, output=True), + "dv": _graph_tensor_spec(torch.empty_like(value_layer), kv_format, output=True), + } + built = build_attention_graph( + cudnn, + graph, + AttentionGraphSpec.from_dict( + AttentionOptions(backward=True, attn_scale=attn_scale, deterministic=deterministic), + specs, + ), + score_mod=_wrap_score_mod(score_mod, extra), + score_mod_bprop=_wrap_score_mod(score_mod_bprop, bprop_extra), ) - dq.set_output(True).set_dim(dq_dim).set_stride(dq_stride) - 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) return _CudnnScoreModBwdGraphEntry( graph=graph, - q=q, - k=k, - v=v, - output=output, - d_output=d_output, - stats=stats_tensor, - dq=dq, - dk=dk, - dv=dv, - score_mod_graph_tensors=score_mod_graph_tensors, - score_mod_bprop_graph_tensors=score_mod_bprop_graph_tensors, + q=built.tensors["q"], + k=built.tensors["k"], + v=built.tensors["v"], + output=built.tensors["o"], + d_output=built.tensors["do"], + stats=built.tensors["stats"], + dq=built.tensors["dq"], + dk=built.tensors["dk"], + dv=built.tensors["dv"], + score_mod_graph_tensors=extra, + score_mod_bprop_graph_tensors=bprop_extra, workspace_size=workspace_size, ) diff --git a/transformer_engine/pytorch/attention/dot_product_attention/utils.py b/transformer_engine/pytorch/attention/dot_product_attention/utils.py index 710ed8f7fbf..9ed286b44a6 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/utils.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/utils.py @@ -449,7 +449,7 @@ 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 + """Constant-foldable cuDNN Python support query: the result depends only on the attention config. Returns a plain int rather than a FusedAttnBackend member: dynamo @@ -458,7 +458,9 @@ def _get_fused_attn_backend(**fused_attn_kwargs): 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( + from transformer_engine.pytorch.cpp_extensions.cudnn_fused_attn import get_fused_attn_backend + + fused_attention_backend, reject_message = get_fused_attn_backend( FusedAttentionParams(**fused_attn_kwargs) ) return int(fused_attention_backend), reject_message @@ -553,6 +555,7 @@ def get_attention_backend( if not logger.hasHandlers(): logger.addHandler(AttentionLogging._stream_handler) device_compute_capability = get_device_compute_capability(attention_params.device) + is_sm12x = device_compute_capability[0] == 12 cudnn_version = get_cudnn_version() run_config = { "transformer_engine_version": te.__version__, @@ -802,16 +805,16 @@ def _disable_all_flash_attention() -> None: logger.debug("Disabling FusedAttention for %s", fp8_recipe.__class__.__name__) use_fused_attention = False - if device_compute_capability == (12, 0): + if is_sm12x: if use_flash_attention: logger.debug( "Disabling FlashAttention as FP8 is not supported" - " for compute capability = sm120" + " for compute capability = sm12x" ) if use_fused_attention: logger.debug( "Disabling FusedAttention as FP8 is not supported" - " for compute capability = sm120" + " for compute capability = sm12x" ) use_flash_attention = False use_fused_attention = False @@ -950,14 +953,14 @@ def _disable_all_flash_attention() -> None: # Flash v4 | FP16/BF16 | TODO | sm80+ | bshd,sbhd,thd | TODO # Unfused | FP32/FP16/BF16 | non-paged/paged | all | bshd,sbhd,thd | >= 1 if inference_params is not None: - # Temporarily disabling fused attention for kv caching for sm89/sm120 irrespective of + # Temporarily disabling fused attention for kv caching for sm89/sm12x irrespective of # cuDNN version until the cuDNN bug is resolved. - if device_compute_capability in ((8, 9), (12, 0)): - logger.debug("Disabling FusedAttention for KV caching for sm89/sm120") + if device_compute_capability == (8, 9) or is_sm12x: + logger.debug("Disabling FusedAttention for KV caching for sm89/sm12x") use_fused_attention = False - # Temporarily disable FlashAttention for KV caching on sm120 - if device_compute_capability == (12, 0): - logger.debug("Disabling FlashAttention for KV caching for sm120") + # Temporarily disable FlashAttention for KV caching on sm12x + if is_sm12x: + logger.debug("Disabling FlashAttention for KV caching for sm12x") use_flash_attention = False if context_parallel: logger.debug("Disabling all backends for KV caching with context parallelism") @@ -1019,15 +1022,11 @@ def _disable_all_flash_attention() -> None: qkv_layout, ) use_fused_attention = False - if ( - device_compute_capability == (12, 0) - and (head_dim_qk > 128 or head_dim_qk % 8 != 0) - and is_training - ): + if is_sm12x and (head_dim_qk > 128 or head_dim_qk % 8 != 0) and is_training: if use_fused_attention: logger.debug( "Disabling FusedAttention as MLA for backward pass is not supported for compute" - " capability = sm120 for a head_dim_qk > 128 or head_dim_qk %%8 != 0. Found:" + " capability = sm12x for a head_dim_qk > 128 or head_dim_qk %%8 != 0. Found:" " head_dim_qk = %s", head_dim_qk, ) @@ -1193,19 +1192,19 @@ def _is_fa3_supported(num_heads, num_gqa_groups, head_dim_qk, head_dim_v, qkv_dt if use_unfused_attention: logger.debug("Disabling UnfusedDotProductAttention for pad_between_seqs = True") use_unfused_attention = False - if device_compute_capability == (12, 0): + if is_sm12x: if cudnn_version < (9, 18, 1): if use_fused_attention: logger.debug( "Disabling FusedAttention when Q or KV uses THD because it is" - " not supported for compute capability = sm120 and cuDNN version < 9.18.1" + " not supported for compute capability = sm12x and cuDNN version < 9.18.1" ) use_fused_attention = False elif qkv_layout in {"t3hd", "th3d"}: if use_fused_attention: logger.debug( "Disabling FusedAttention as qkv_layout = %s is not supported for" - " compute capability = sm120", + " compute capability = sm12x", qkv_layout, ) use_fused_attention = False diff --git a/transformer_engine/pytorch/cpp_extensions/cudnn.py b/transformer_engine/pytorch/cpp_extensions/cudnn.py new file mode 100644 index 00000000000..10b36b2d53a --- /dev/null +++ b/transformer_engine/pytorch/cpp_extensions/cudnn.py @@ -0,0 +1,62 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""cuDNN Python frontend graph lifecycle shared by attention implementations.""" + +import threading + +import torch + +from transformer_engine.common.cudnn.graph import ( + import_cudnn as _import_cudnn_frontend, + create_graph, + check_graph_support, + compile_graph, +) + +_handles = threading.local() + + +def _get_cudnn_current_stream_handle(cudnn, device): + """Get a thread-local, device-specific handle on PyTorch's current stream.""" + if device.type != "cuda": + raise ValueError(f"cuDNN attention requires CUDA tensors, got {device}.") + if device.index is None: + device = torch.device("cuda", torch.cuda.current_device()) + if not hasattr(_handles, "by_device"): + _handles.by_device = {} + with torch.cuda.device(device): + handle = _handles.by_device.get(device) + if handle is None: + handle = cudnn.create_handle() + _handles.by_device[device] = handle + cudnn.set_stream(handle=handle, stream=torch.cuda.current_stream(device).cuda_stream) + return handle + + +def _build_cudnn_pygraph(dtype, device): + """Create an attention graph with FP32 accumulation.""" + cudnn = _import_cudnn_frontend() + dtype = {torch.float16: cudnn.data_type.HALF, torch.bfloat16: cudnn.data_type.BFLOAT16}.get( + dtype, dtype + ) + return create_graph(cudnn, dtype, _get_cudnn_current_stream_handle(cudnn, device)) + + +def _finalize_cudnn_graph(graph): + """Check support, build an execution plan, and return workspace bytes.""" + cudnn = _import_cudnn_frontend() + check_graph_support(cudnn, graph) + return compile_graph(cudnn, graph) + + +def _execute_cudnn_graph(graph, variant_pack, workspace_size, device): + """Execute on the current stream using framework-owned workspace.""" + cudnn = _import_cudnn_frontend() + workspace = torch.empty(workspace_size, device=device, dtype=torch.uint8) + graph.execute( + variant_pack, + workspace, + handle=_get_cudnn_current_stream_handle(cudnn, device), + ) diff --git a/transformer_engine/pytorch/cpp_extensions/cudnn_fused_attn.py b/transformer_engine/pytorch/cpp_extensions/cudnn_fused_attn.py new file mode 100644 index 00000000000..da5a0f75c18 --- /dev/null +++ b/transformer_engine/pytorch/cpp_extensions/cudnn_fused_attn.py @@ -0,0 +1,794 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""cuDNN attention graphs constructed and executed through the Python frontend. + +The extension is used only to obtain capture-safe Philox state. Tensor allocation, +graph construction, plan caching, and execution do not use the common attention API. +""" + +from dataclasses import dataclass +import threading + +import torch +import transformer_engine_torch as tex + +from transformer_engine.common.cudnn.attention import ( + AttentionOptions, + AttentionGraphSpec, + build_attention_graph, +) +from transformer_engine.common.cudnn.graph import ( + TensorSpec as _TensorSpec, + CacheDiagnostics, + attention_matrix_spec, + attention_stats_spec, + mxfp8_scale_spec, + check_graph_support, + compile_graph, + bucket, + uses_ragged_attention_graph, +) + +from .cudnn import ( + _import_cudnn_frontend, + _build_cudnn_pygraph, + _execute_cudnn_graph, +) +from ..tensor.float8_tensor import Float8Quantizer +from ..tensor.storage.mxfp8_tensor_storage import MXFP8TensorStorage +from ..constants import DType + +_plans = {} +_cache_lock = threading.RLock() +_diagnostics = CacheDiagnostics("PyTorch") + + +def _event(attn, event, key=None): + _diagnostics.event( + "fp8" if attn.fp8 else "f16", + "bwd" if attn.backward else "fwd", + attn.device.index, + event, + key, + ) + + +def _build_stage(attn, name, function): + return _diagnostics.run( + "fp8" if attn.fp8 else "f16", + "bwd" if attn.backward else "fwd", + name, + function, + ) + + +def _formats(layout): + parts = layout.removeprefix("paged_kv_").replace("2", "").replace("3", "").split("_") + return parts[0], parts[1] if len(parts) > 1 else parts[0] + + +def _shape(dim, fmt, tokens=0): + b, h, s, d = dim + if fmt == "bshd": + return (b, s, h, d) + if fmt == "sbhd": + return (s, b, h, d) + if fmt == "bhsd": + return dim + if fmt == "thd": + return (tokens, h, d) + raise ValueError(f"Unsupported attention format: {fmt}") + + +def _raw(tensor, columnwise=False): + if isinstance(tensor, MXFP8TensorStorage): + return tensor._columnwise_data if columnwise else tensor._rowwise_data + return getattr(tensor, "_data", tensor) + + +def _dtype(tensor): + fp8_dtype = getattr(tensor, "_fp8_dtype", None) + if fp8_dtype is not None: + return "FP8_E4M3" if fp8_dtype == DType.kFloat8E4M3 else "FP8_E5M2" + return { + torch.float16: "HALF", + torch.bfloat16: "BFLOAT16", + torch.float32: "FLOAT", + torch.int32: "INT32", + torch.int64: "INT64", + }[tensor.dtype] + + +def _descale(tensor, columnwise=False): + if isinstance(tensor, MXFP8TensorStorage): + return tensor._columnwise_scale_inv if columnwise else tensor._rowwise_scale_inv + return tensor._scale_inv + + +@dataclass +class _PlanEntry: + """A supported graph, compiled lazily under the cache lock.""" + + graph: object + tensors: dict + workspace_size: int | None = None + + +class _AttentionGraph: + """Runtime tensors and data-free descriptors for one attention invocation.""" + + def __init__( + self, params, backward=False, *, dimensions=None, device=None, fp8=None, mxfp8=None + ): + self.params = params + self.backward = backward + self.device = params["q"].device if device is None else device + self.q_format, self.kv_format = _formats(params["qkv_layout"]) + if dimensions is None: + dimensions = ( + params["cu_seqlens_q"].numel() - 1, + params["q"].shape[1] if self.q_format == "bhsd" else params["q"].shape[-2], + params["k"].shape[1] if self.kv_format == "bhsd" else params["k"].shape[-2], + params["q"].shape[-1], + params["v"].shape[-1], + params["q"].shape[0] if self.q_format == "thd" else 0, + params["k"].shape[0] if self.kv_format == "thd" else 0, + ) + self.b, self.h, self.hg, self.dq, self.dv, self.tq, self.tk = dimensions + capability = torch.cuda.get_device_capability(self.device) + self.ragged_graph = uses_ragged_attention_graph(10 * capability[0] + capability[1]) + self.ragged_stats = self.q_format == "thd" and self.ragged_graph + self.sq = bucket(self.tq, 1024, 32768) if self.ragged_stats else params["max_seqlen_q"] + self.skv = ( + bucket(self.tk, 1024, 32768) + if self.kv_format == "thd" and self.ragged_graph + else params["max_seqlen_kv"] + ) + self.fp8 = hasattr(params["q"], "_fp8_dtype") if fp8 is None else fp8 + self.mxfp8 = isinstance(params["q"], MXFP8TensorStorage) if mxfp8 is None else mxfp8 + self.specs = {} + self.tensors = {} + + def add(self, name, tensor, dim=None, stride=None, dtype=None, **kwargs): + """Bind a runtime tensor and record its address-independent descriptor.""" + self.tensors[name] = tensor + self.specs[name] = _TensorSpec( + tuple(dim if dim is not None else tensor.shape), + tuple(stride if stride is not None else tensor.stride()), + dtype or _dtype(tensor), + **kwargs, + ) + return name + + def scalar(self, name, tensor, output=False, by_value=False): + """Describe a scalar using cuDNN's four-dimensional convention.""" + return self.add(name, tensor, (1, 1, 1, 1), (1, 1, 1, 1), output=output, by_value=by_value) + + def matrix( + self, name, tensor, fmt, heads, seqlen, cu_padded=None, columnwise=False, output=False + ): + """Bind storage using the same matrix descriptors as support queries.""" + data = _raw(tensor, columnwise) + dim = (self.b, heads, seqlen, tensor.shape[-1]) + offset_name = name + "_offset" if fmt == "thd" else "" + if offset_name: + offsets = cu_padded.to(torch.int64) * data.stride(0) + self.add(offset_name, offsets, (self.b + 1, 1, 1, 1), (1, 1, 1, 1)) + elif name in ("k", "v") and self.params.get("page_table_k") is not None: + order = (0, 2, 1, 3) if fmt == "bshd" else (0, 1, 2, 3) + dim = tuple(data.shape[index] for index in order) + self.tensors[name] = data + self.specs[name] = attention_matrix_spec( + dim, + fmt, + _dtype(tensor), + storage_stride=data.stride(), + output=output, + ragged_offset=offset_name, + ) + return name + + def stats(self, name, tensor, output=False): + """Bind dense or packed softmax statistics.""" + if self.ragged_stats: + offsets = self.params["cu_seqlens_q_padded"].to(torch.int64) * self.h + self.add("stats_offset", offsets, (self.b + 1, 1, 1, 1), (1, 1, 1, 1)) + self.tensors[name] = tensor + self.specs[name] = attention_stats_spec( + self.b, self.h, self.sq, ragged=self.ragged_stats, output=output + ) + return name + + def scale(self, name, tensor, matrix_name, fmt, columnwise=False): + """Bind tensor scales or reordered MXFP8 block scales.""" + if not self.mxfp8: + return self.scalar(name, tensor) + self.tensors[name] = tensor + self.specs[name] = mxfp8_scale_spec(self.specs[matrix_name].dim, fmt, columnwise) + return name + + def options(self): + """Return graph attributes that affect support and plan caching.""" + p = self.params + mask = p["attn_mask_type"] + causal = "causal" in mask + bottom_right = p["bottom_right_diagonal"] + if causal: + bottom_right = "bottom_right" in mask + left, right = p["window_size"] or (-1, -1) + if not causal and (left, right) == (-1, -1): + bottom_right = False + bounded_backward = self.backward and not self.fp8 and self.ragged_graph + return AttentionOptions( + backward=self.backward, + fp8=self.fp8, + mxfp8=self.mxfp8, + bottom_right=bottom_right, + left_bound=left + 1 if left != -1 else None, + right_bound=0 if causal else (right if right != -1 else None), + padding="padding" in mask, + alibi=p["attn_bias_type"] == "alibi", + dropout=p["dropout"], + deterministic=p.get("deterministic", False), + paged_max_seqlen_kv=p["max_seqlen_kv"] if p.get("page_table_k") is not None else 0, + # Ragged sequence dimensions are bucketed total-token capacities. + # Bound cuDNN's backward scratch buffers by those same capacities. + max_total_seqlen_q=self.sq if bounded_backward and self.q_format == "thd" else 0, + max_total_seqlen_kv=self.skv if bounded_backward and self.kv_format == "thd" else 0, + ) + + def description(self): + """Describe the graph independently of buffer addresses and scalar values.""" + return AttentionGraphSpec.from_dict(self.options(), self.specs) + + def entry(self): + """Share support-query graphs with execution without compiling on queries.""" + key = (self.device, self.description()) + # Autograd may run backward on another thread. Share device-specific plans; + # each execution supplies its thread-local handle on the current stream. + with _cache_lock: + entry = _plans.get(key) + if entry is None: + _event(self, "MISS", key) + entry = self.build() + _plans[key] = entry + _event(self, "CACHE_GRAPH") + else: + _event(self, "HIT", key) + return entry + + def execute(self): + """Compile once, then bind this invocation's runtime tensor addresses.""" + entry = self.entry() + with _cache_lock: + if entry.workspace_size is None: + entry.workspace_size = compile_graph( + _import_cudnn_frontend(), + entry.graph, + lambda name, function: _build_stage(self, name, function), + ) + _event(self, "BUILD_PLANS") + _execute_cudnn_graph( + entry.graph, + {entry.tensors[n]: t for n, t in self.tensors.items()}, + entry.workspace_size, + self.device, + ) + _event(self, "EXECUTE") + + def build(self): + """Construct and check the common graph description without compiling it.""" + cudnn = _import_cudnn_frontend() + graph = _build_cudnn_pygraph(getattr(cudnn.data_type, self.specs["q"].dtype), self.device) + built = build_attention_graph(cudnn, graph, self.description()) + _event(self, "CREATE_GRAPH") + check_graph_support(cudnn, graph, lambda name, function: _build_stage(self, name, function)) + return _PlanEntry(graph, built.tensors) + + +def _prepare_inputs(attn): + p = attn.params + for name, fmt, h, s in ( + ("q", attn.q_format, attn.h, attn.sq), + ("k", attn.kv_format, attn.hg, attn.skv), + ("v", attn.kv_format, attn.hg, attn.skv), + ): + cu = p["cu_seqlens_q_padded"] if name == "q" else p["cu_seqlens_kv_padded"] + col = attn.mxfp8 and name == "v" and not attn.backward + attn.matrix(name, p[name], fmt, h, s, cu, col) + if attn.fp8: + scale_fmt = p.get("qkv_scale_inv_format") or fmt + attn.scale("descale_" + name, _descale(p[name], col), name, scale_fmt, col) + attn.scalar("attn_scale", torch.tensor([p["attn_scale"]], dtype=torch.float32), by_value=True) + if "padding" in p["attn_mask_type"]: + for name, cu in (("seq_q", p["cu_seqlens_q"]), ("seq_kv", p["cu_seqlens_kv"])): + attn.add(name, cu[1:] - cu[:-1], (attn.b, 1, 1, 1), (1, 1, 1, 1)) + if p["attn_bias_type"] not in ("no_bias", "alibi"): + attn.add("bias", p["attn_bias"]) + if p["softmax_type"] != "vanilla": + attn.add("sink", p["softmax_offset"]) + if p["dropout"]: + rng = p["rng_state"] + attn.scalar("seed", rng[:1]) + attn.scalar("offset", rng[1:]) + if p.get("page_table_k") is not None: + for name in ("page_table_k", "page_table_v"): + table = p[name] + attn.add( + name, + table, + (attn.b, 1, table.shape[-1], 1), + (table.shape[-1], table.shape[-1], 1, 1), + ) + + +def _allocate(shape, dtype, device, quantizer, data=None): + """Allocate an output, preserving delayed-scaling metadata and packed views.""" + delayed = isinstance(quantizer, Float8Quantizer) + if data is None: + data = torch.empty(shape, dtype=torch.uint8 if delayed else dtype, device=device) + if delayed: + return quantizer.create_tensor_from_data( + data, fake_dtype=dtype, internal=quantizer.internal + ) + return data + + +def _output_shape(attn, name, fmt): + is_q = name == "q" + return _shape( + ( + attn.b, + attn.h if is_q else attn.hg, + attn.params["max_seqlen_q"] if is_q else attn.params["max_seqlen_kv"], + attn.dq if name in ("q", "k") else attn.dv, + ), + fmt, + attn.tq if is_q else attn.tk, + ) + + +def _allocate_gradients(attn): + p = attn.params + qfmt, kfmt = _formats(p["dqkv_layout"]) + shapes = [_output_shape(attn, n, qfmt if n == "q" else kfmt) for n in ("q", "k", "v")] + delayed = isinstance(p["dqkv_quantizer"], Float8Quantizer) + dtype = torch.uint8 if delayed else p["fake_dtype"] + layout = p["dqkv_layout"] + if "3" in layout: + axis = -2 if "h3d" in layout else -3 + packed_shape = list(shapes[0]) + packed_shape.insert(len(packed_shape) + axis + 1, 3) + data = torch.empty(packed_shape, dtype=dtype, device=attn.device).unbind(axis) + elif "2" in layout: + axis = -2 if "h2d" in layout else -3 + packed_shape = list(shapes[1]) + packed_shape.insert(len(packed_shape) + axis + 1, 2) + kv = torch.empty(packed_shape, dtype=dtype, device=attn.device).unbind(axis) + data = (torch.empty(shapes[0], dtype=dtype, device=attn.device), *kv) + else: + data = [torch.empty(s, dtype=dtype, device=attn.device) for s in shapes] + return [ + _allocate(s, p["fake_dtype"], attn.device, p["dqkv_quantizer"], t) + for s, t in zip(shapes, data) + ] + + +def _normalize(params): + p = dict(params) + p["cu_seqlens_q_padded"] = p.get("cu_seqlens_q_padded") + if p["cu_seqlens_q_padded"] is None: + p["cu_seqlens_q_padded"] = p["cu_seqlens_q"] + p["cu_seqlens_kv_padded"] = p.get("cu_seqlens_kv_padded") + if p["cu_seqlens_kv_padded"] is None: + p["cu_seqlens_kv_padded"] = p["cu_seqlens_kv"] + return p + + +def _query_graph(config, backward): + """Describe a support query without allocating tensors or touching RNG state.""" + from .fused_attn import QKVLayout, QKVFormat, AttnMaskType, AttnBiasType, SoftmaxType + + def name(mapping, value): + return next(k for k, v in mapping.items() if v == value) + + layout = name(QKVLayout, config.qkv_layout) + grad_layout = ( + name(QKVLayout, config.dqkv_layout) + if (config.dqkv_layout != tex.NVTE_QKV_Layout.NVTE_QKV_Layout_NOT_SET) + else layout + ) + types = { + DType.kFloat16: "HALF", + DType.kBFloat16: "BFLOAT16", + DType.kFloat32: "FLOAT", + DType.kFloat8E4M3: "FP8_E4M3", + DType.kFloat8E5M2: "FP8_E5M2", + } + fp8 = types[DType.cast(config.qkv_dtype)].startswith("FP8") + mxfp8 = fp8 and config.scaling_mode == tex.NVTEScalingMode.NVTE_MXFP8_1D_SCALING + params = { + "attn_mask_type": name(AttnMaskType, config.attn_mask_type), + "attn_bias_type": name(AttnBiasType, config.bias_type), + "softmax_type": name(SoftmaxType, config.softmax_type), + "bottom_right_diagonal": config.bottom_right_diagonal, + "window_size": (config.window_size_left, config.window_size_right), + "dropout": config.dropout if config.is_training else 0.0, + "deterministic": config.deterministic if backward else False, + "return_max_logit": config.return_max_logit if not backward else False, + "max_seqlen_kv": config.max_seqlen_kv, + "page_table_k": True if layout.startswith("paged_kv_") else None, + } + params.update(qkv_layout=layout, max_seqlen_q=config.max_seqlen_q) + attn = _AttentionGraph( + params, + backward, + dimensions=( + config.batch_size, + config.num_attn_heads, + config.num_gqa_groups, + config.head_dim_qk, + config.head_dim_v, + config.num_tokens_q, + config.num_tokens_kv, + ), + device=torch.device("cuda", torch.cuda.current_device()), + fp8=fp8, + mxfp8=mxfp8, + ) + + def spec(n, dim=(1, 1, 1, 1), stride=(1, 1, 1, 1), dtype="FLOAT", **kwargs): + attn.specs[n] = _TensorSpec(tuple(dim), tuple(stride), dtype, **kwargs) + + def matrix(n, fmt, h, s, d, dtype, packed="", output=False): + dim = (attn.b, h, s, d) + if n in ("k", "v") and layout.startswith("paged_kv_"): + dim = (getattr(config, "num_pages_" + n), h, getattr(config, "page_size_" + n), d) + offset = n + "_offset" if fmt == "thd" else "" + if offset: + spec(offset, (attn.b + 1, 1, 1, 1), dtype="INT64") + attn.specs[n] = attention_matrix_spec( + dim, fmt, dtype, packed=packed, output=output, ragged_offset=offset + ) + + def stats(n, output=False): + if attn.ragged_stats: + spec("stats_offset", (attn.b + 1, 1, 1, 1), dtype="INT64") + attn.specs[n] = attention_stats_spec( + attn.b, attn.h, attn.sq, ragged=attn.ragged_stats, output=output + ) + + def scale(n, matrix_name, fmt, columnwise=False): + if not attn.mxfp8: + spec(n) + else: + attn.specs[n] = mxfp8_scale_spec(attn.specs[matrix_name].dim, fmt, columnwise) + + parts = layout.removeprefix("paged_kv_").split("_") + for n, fmt, h, s, d, part in ( + ("q", attn.q_format, attn.h, attn.sq, attn.dq, parts[0]), + ("k", attn.kv_format, attn.hg, attn.skv, attn.dq, parts[1] if len(parts) > 1 else parts[0]), + ("v", attn.kv_format, attn.hg, attn.skv, attn.dv, parts[-1]), + ): + matrix(n, fmt, h, s, d, types[DType.cast(config.qkv_dtype)], part) + if attn.fp8: + scale_fmt = name(QKVFormat, config.qkv_scale_inv_format) or fmt + scale("descale_" + n, n, scale_fmt, attn.mxfp8 and n == "v" and not backward) + spec("attn_scale", by_value=True) + if "padding" in attn.params["attn_mask_type"]: + for n in ("seq_q", "seq_kv"): + spec(n, (attn.b, 1, 1, 1), dtype="INT32") + if attn.params["dropout"]: + spec("seed", dtype="INT64") + spec("offset", dtype="INT64") + if attn.params["attn_bias_type"] not in ("no_bias", "alibi"): + dim = ( + config.bias_batch_size, + config.bias_num_heads, + config.bias_seqlen_q, + config.bias_seqlen_kv, + ) + b, h, sq, skv = dim + stride = (h * sq * skv, sq * skv, skv, 1) + spec("bias", dim, stride, types[DType.cast(config.qkv_dtype)]) + if backward and (b, h, sq) != (1, 1, 1): + spec("dbias", dim, stride, types[DType.cast(config.qkv_dtype)], output=True) + if attn.params["softmax_type"] != "vanilla": + spec("sink", (1, attn.h, 1, 1), (attn.h, 1, 1, 1)) + if backward: + spec("dsink", (1, attn.h, 1, 1), (attn.h, 1, 1, 1), output=True) + if layout.startswith("paged_kv_"): + for n in ("k", "v"): + pages = getattr(config, "max_pages_per_seq_" + n) + spec("page_table_" + n, (attn.b, 1, pages, 1), (pages, pages, 1, 1), "INT32") + ofmt = name(QKVFormat, config.o_format) or attn.q_format + matrix( + "o", ofmt, attn.h, attn.sq, attn.dv, types[DType.cast(config.o_dtype)], output=not backward + ) + stats("stats", output=not backward) + if not backward and config.return_max_logit: + stats("max", output=True) + if backward: + dofmt = name(QKVFormat, config.do_format) or ofmt + matrix("do", dofmt, attn.h, attn.sq, attn.dv, types[DType.cast(config.do_dtype)]) + qfmt, kvfmt = _formats(grad_layout) + parts = grad_layout.split("_") + for n, fmt, h, s, d, part in ( + ("dq", qfmt, attn.h, attn.sq, attn.dq, parts[0]), + ("dk", kvfmt, attn.hg, attn.skv, attn.dq, parts[1] if len(parts) > 1 else parts[0]), + ("dv", kvfmt, attn.hg, attn.skv, attn.dv, parts[-1]), + ): + matrix(n, fmt, h, s, d, types[DType.cast(config.dqkv_dtype)], part, output=True) + if attn.fp8: + do_scale_fmt = name(QKVFormat, config.do_scale_inv_format) or dofmt + scale("descale_do", "do", do_scale_fmt) + if attn.mxfp8: + matrix("do_f16", dofmt, attn.h, attn.sq, attn.dv, types[DType.cast(config.o_dtype)]) + for n, fmt, h, s, d, dtype in ( + ("q_t", attn.q_format, attn.h, attn.sq, attn.dq, config.qkv_dtype), + ("k_t", attn.kv_format, attn.hg, attn.skv, attn.dq, config.qkv_dtype), + ("do_t", dofmt, attn.h, attn.sq, attn.dv, config.do_dtype), + ): + matrix(n, fmt, h, s, d, types[DType.cast(dtype)]) + sfmt = ( + do_scale_fmt + if n == "do_t" + else name(QKVFormat, config.qkv_scale_inv_format) or fmt + ) + scale("descale_" + n, n, sfmt, True) + else: + for n in ( + "descale_o", + "descale_s", + "descale_dp", + "scale_s", + "scale_dp", + "scale_dq", + "scale_dk", + "scale_dv", + ): + spec(n) + for n in ("amax_dq", "amax_dk", "amax_dv", "amax_dp"): + spec(n, output=True) + elif attn.fp8 and not attn.mxfp8: + for n in ("descale_s", "scale_s", "scale_o"): + spec(n) + for n in ("amax_s", "amax_o"): + spec(n, output=True) + return attn + + +def _support_reason(config, attn, version, capability): + """Preserve TE's hardware, recipe, and known cuDNN defect exclusions.""" + p = attn.params + ragged = attn.q_format == "thd" or attn.kv_format == "thd" + padding = "padding" in p["attn_mask_type"] + if ragged and not padding: + return "THD format requires PADDING / PADDING_CAUSAL / PADDING_CAUSAL_BOTTOM_RIGHT mask." + if (attn.q_format == "thd" and not attn.tq) or (attn.kv_format == "thd" and not attn.tk): + return "THD format requires num_tokens_q / num_tokens_kv to be set for the ragged inputs." + if p["page_table_k"] is not None: + if not padding: + return "Paged KV requires PADDING / PADDING_CAUSAL / PADDING_CAUSAL_BOTTOM_RIGHT mask." + if not all( + getattr(config, field + "_" + name) + for field in ("num_pages", "page_size", "max_pages_per_seq") + for name in ("k", "v") + ): + return "Paged KV requires num_pages, page_size and max_pages_per_seq for both K and V." + if not attn.fp8: + if attn.q_format == attn.kv_format == "thd" and capability < (9, 0): + return "F16/BF16 fused attention with THD format requires sm90 or later." + affected_cuda_graph = version <= 91500 and config.cuda_graph and config.max_seqlen_kv % 128 + if ( + affected_cuda_graph + and config.is_training + and attn.q_format == attn.kv_format + and attn.q_format in ("bshd", "sbhd") + and not padding + ): + return "Known cuDNN <= 9.15 issue with CUDA graph. Please upgrade cuDNN." + if ( + config.is_training + and attn.ragged_stats + and p["softmax_type"] == "learnable" + and version < 92600 + ): + return ( + "Known cuDNN < 9.26.0 issue with THD learnable softmax backward. Please upgrade" + " cuDNN." + ) + return "" + if config.return_max_logit: + return "FP8 fused attention does not support return_max_logit=True." + if attn.q_format != attn.kv_format: + return "FP8 fused attention supports BSHD/SBHD/BHSD/THD formats." + if config.is_training and config.dqkv_layout != tex.NVTE_QKV_Layout.NVTE_QKV_Layout_NOT_SET: + from .fused_attn import QKVLayout + + grad_layout = next(k for k, v in QKVLayout.items() if v == config.dqkv_layout) + qfmt, kvfmt = _formats(grad_layout) + if qfmt != kvfmt: + return "FP8 fused attention supports BSHD/SBHD/BHSD/THD gradient formats." + if attn.q_format == "thd": + if version < 92300: + return "FP8 fused attention with THD format requires cuDNN 9.23.0 or later!" + if config.is_training and capability < (10, 0): + return "FP8 fused attention with THD format supports backward on sm100+ only!" + if config.is_training and p["softmax_type"] != "vanilla" and version < 92600: + return ( + "FP8 THD attention with a sink token requires cuDNN 9.26.0 or later for backward!" + ) + if capability >= (10, 0) and max(attn.dq, attn.dv) > 128: + return "FP8 THD attention supports head dimensions up to 128 on sm100+ only!" + if p["attn_bias_type"] != "no_bias": + return "FP8 fused attention does not support attention bias." + half_types = (DType.kFloat16, DType.kBFloat16) + fp8_types = (DType.kFloat8E4M3, DType.kFloat8E5M2) + allowed_types = half_types if attn.mxfp8 else half_types + fp8_types + tensor_scaling = config.scaling_mode == tex.NVTEScalingMode.NVTE_DELAYED_TENSOR_SCALING + if ( + not (tensor_scaling or attn.mxfp8) + or DType.cast(config.o_dtype) not in allowed_types + or (config.is_training and DType.cast(config.dqkv_dtype) not in allowed_types) + ): + return ( + "FP8 fused attention only supports FP8DelayedScaling, FP8CurrentScaling, or MXFP8" + " recipes!" + ) + if attn.mxfp8 and version < 92100: + return "MXFP8 fused attention requires cuDNN 9.21.0 or later!" + return "" + + +def get_fused_attn_backend(config): + """Select an attention backend by validating the Python frontend graphs.""" + from .fused_attn import FusedAttnBackend + + cudnn = _import_cudnn_frontend() + if config.bias_type == tex.NVTE_Bias_Type.NVTE_PRE_SCALE_BIAS: + return FusedAttnBackend.No_Backend, "Fused attention does not support pre-scale bias." + if DType.cast(config.qkv_dtype) not in ( + DType.kFloat16, + DType.kBFloat16, + DType.kFloat8E4M3, + DType.kFloat8E5M2, + ): + return FusedAttnBackend.No_Backend, f"Unsupported QKV dtype qkv_dtype={config.qkv_dtype}." + try: + fwd = _query_graph(config, False) + reason = _support_reason( + config, fwd, cudnn.backend_version(), torch.cuda.get_device_capability() + ) + if reason: + return FusedAttnBackend.No_Backend, reason + fwd.entry() + if config.is_training: + _query_graph(config, True).entry() + return FusedAttnBackend.FP8 if fwd.fp8 else FusedAttnBackend.F16_arbitrary_seqlen, "" + except (cudnn.cudnnGraphNotSupportedError, ValueError) as exc: + return FusedAttnBackend.No_Backend, str(exc) + + +def fused_attn_fwd(params): + """Execute forward; preserve the extension's output and auxiliary tensor order.""" + p = _normalize(params) + with torch.cuda.device(p["q"].device): + p["rng_state"] = tex.fused_attn_rng_state(p["rng_gen"], p["rng_elts_per_thread"]) + attn = _AttentionGraph(p) + _prepare_inputs(attn) + out_shape = list(_output_shape(attn, "q", p["o_format"])) + out_shape[-1] = attn.dv + out = _allocate(out_shape, p["fake_dtype"], attn.device, p["o_quantizer"]) + if p["o_format"] == "thd": + _raw(out).zero_() + attn.matrix("o", out, p["o_format"], attn.h, attn.sq, p["cu_seqlens_q_padded"], output=True) + stats_shape = (attn.tq, attn.h, 1) if attn.ragged_stats else (attn.b, attn.h, attn.sq, 1) + stats = torch.empty(stats_shape, dtype=torch.float32, device=attn.device) + attn.stats("stats", stats, output=True) + aux_ctx_tensors = [stats, p["rng_state"]] + maximum = None + if p["return_max_logit"]: + maximum = torch.empty_like(stats) + attn.stats("max", maximum, output=True) + if p["attn_bias_type"] not in ("no_bias", "alibi"): + aux_ctx_tensors.append(p["attn_bias"]) + if p["softmax_type"] != "vanilla": + aux_ctx_tensors.append(p["softmax_offset"]) + if attn.fp8 and not attn.mxfp8: + sq = p["s_quantizer"] + attn.scalar("scale_s", sq.scale) + attn.scalar("descale_s", sq.scale.reciprocal()) + attn.scalar("amax_s", sq.amax, output=True) + if isinstance(p["o_quantizer"], Float8Quantizer): + scale_o, amax_o = p["o_quantizer"].scale, p["o_quantizer"].amax + else: + scale_o = torch.ones(1, dtype=torch.float32, device=attn.device) + amax_o = torch.empty_like(scale_o) + attn.scalar("scale_o", scale_o) + attn.scalar("amax_o", amax_o, output=True) + attn.execute() + return out, aux_ctx_tensors, maximum + + +def fused_attn_bwd(params): + """Execute backward, including FP8/MXFP8 and shared packed gradient storage.""" + p = _normalize(params) + aux = iter(p["aux_ctx_tensors"]) + stats, p["rng_state"] = next(aux), next(aux) + p["attn_bias"] = next(aux) if p["attn_bias_type"] not in ("no_bias", "alibi") else None + p["softmax_offset"] = next(aux) if p["softmax_type"] != "vanilla" else None + with torch.cuda.device(p["q"].device): + attn = _AttentionGraph(p, backward=True) + _prepare_inputs(attn) + attn.matrix("o", p["o"], p["o_format"], attn.h, attn.sq, p["cu_seqlens_q_padded"]) + attn.matrix("do", p["d_o"], p["do_format"], attn.h, attn.sq, p["cu_seqlens_q_padded"]) + attn.stats("stats", stats) + dq, dk, dv = _allocate_gradients(attn) + qfmt, kfmt = _formats(p["dqkv_layout"]) + for name, t, fmt, h, s, cu in ( + ("dq", dq, qfmt, attn.h, attn.sq, p["cu_seqlens_q_padded"]), + ("dk", dk, kfmt, attn.hg, attn.skv, p["cu_seqlens_kv_padded"]), + ("dv", dv, kfmt, attn.hg, attn.skv, p["cu_seqlens_kv_padded"]), + ): + if fmt == "thd": + _raw(t).zero_() + attn.matrix(name, t, fmt, h, s, cu, output=True) + dbias = torch.empty(0, device=attn.device, dtype=p["fake_dtype"]) + if p["attn_bias"] is not None: + dbias = torch.empty_like(p["attn_bias"]) + if p["attn_bias"].shape[-2] != 1 or p["attn_bias"].shape[:2] != (1, 1): + attn.add("dbias", dbias, output=True) + dsoftmax = torch.empty(0, device=attn.device, dtype=torch.float32) + if p["softmax_type"] != "vanilla": + dsoftmax = torch.empty_like(p["softmax_offset"]) + attn.add("dsink", dsoftmax, output=True) + if attn.fp8: + do_scale_fmt = p.get("do_scale_inv_format") or p["do_format"] + attn.scale("descale_do", _descale(p["d_o"]), "do", do_scale_fmt) + if attn.mxfp8: + do_f16 = next(aux) + attn.matrix( + "do_f16", do_f16, p["do_format"], attn.h, attn.sq, p["cu_seqlens_q_padded"] + ) + for name, t, fmt, h, s, cu in ( + ("q_t", p["q"], attn.q_format, attn.h, attn.sq, p["cu_seqlens_q_padded"]), + ("k_t", p["k"], attn.kv_format, attn.hg, attn.skv, p["cu_seqlens_kv_padded"]), + ("do_t", p["d_o"], p["do_format"], attn.h, attn.sq, p["cu_seqlens_q_padded"]), + ): + attn.matrix(name, t, fmt, h, s, cu, columnwise=True) + scale_fmt = ( + do_scale_fmt if name == "do_t" else p.get("qkv_scale_inv_format") or fmt + ) + attn.scale("descale_" + name, _descale(t, True), name, scale_fmt, True) + else: + attn.scalar( + "descale_o", + ( + _descale(p["o"]) + if hasattr(p["o"], "_fp8_dtype") + else torch.ones(1, dtype=torch.float32, device=attn.device) + ), + ) + for short, quantizer in (("s", p["s_quantizer"]), ("dp", p["dp_quantizer"])): + attn.scalar("scale_" + short, quantizer.scale) + attn.scalar("descale_" + short, quantizer.scale.reciprocal()) + attn.scalar("amax_dp", p["dp_quantizer"].amax, output=True) + amax_grads = [] + for name in ("dq", "dk", "dv"): + quantizer = p["dqkv_quantizer"] + scale = ( + quantizer.scale + if isinstance(quantizer, Float8Quantizer) + else (torch.ones(1, dtype=torch.float32, device=attn.device)) + ) + attn.scalar("scale_" + name, scale) + amax = torch.empty(1, dtype=torch.float32, device=attn.device) + attn.scalar("amax_" + name, amax, output=True) + amax_grads.append(amax) + attn.execute() + if attn.fp8 and not attn.mxfp8 and isinstance(p["dqkv_quantizer"], Float8Quantizer): + p["dqkv_quantizer"].amax.copy_(torch.stack(amax_grads).amax().reshape(1)) + return dq, dk, dv, dbias, dsoftmax diff --git a/transformer_engine/pytorch/cpp_extensions/fused_attn.py b/transformer_engine/pytorch/cpp_extensions/fused_attn.py index 9a33df7634b..98628205782 100644 --- a/transformer_engine/pytorch/cpp_extensions/fused_attn.py +++ b/transformer_engine/pytorch/cpp_extensions/fused_attn.py @@ -8,7 +8,6 @@ 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, @@ -363,52 +362,53 @@ def fused_attn_fwd( 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, + # Import lazily: the frontend's quantized tensor classes also import this module. + from . import cudnn_fused_attn + + out, aux_ctx_tensors, max_tensor = cudnn_fused_attn.fused_attn_fwd( + { + "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, + "fake_dtype": fake_dtype, + "fused_attention_backend": fused_attention_backend, + "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, + "s_quantizer": s_quantizer, + "o_quantizer": o_quantizer, + "attn_scale": attn_scale, + "dropout": dropout, + "fast_zero_fill": fast_zero_fill, + "qkv_layout": qkv_layout, + "o_format": o_format, + "qkv_scale_inv_format": qkv_scale_inv_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, + "rng_gen": rng_gen, + "softmax_offset": softmax_offset, + "return_max_logit": return_max_logit, + "cuda_graph": cuda_graph, + "rng_elts_per_thread": rng_elts_per_thread, + } ) 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 (newer cuDNN runtimes, non-sm120): 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": @@ -438,11 +438,11 @@ def fused_attn_fwd( 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 + max_logit = torch.amax(max_tensor, dim=amax_dims).to(dtype=out.dtype) + return out, aux_ctx_tensors, max_logit # out, aux_ctx_tensors - return output_tensors[0], output_tensors[1:] + return out, aux_ctx_tensors def fused_attn_bwd( @@ -615,39 +615,44 @@ def fused_attn_bwd( 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, + from . import cudnn_fused_attn + + output_tensors = cudnn_fused_attn.fused_attn_bwd( + { + "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, + "fake_dtype": fake_dtype, + "aux_ctx_tensors": aux_ctx_tensors, + "fused_attention_backend": fused_attention_backend, + "cu_seqlens_q_padded": cu_seqlens_q_padded, + "cu_seqlens_kv_padded": cu_seqlens_kv_padded, + "s_quantizer": s_quantizer, + "dp_quantizer": dp_quantizer, + "dqkv_quantizer": dqkv_quantizer, + "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, + "qkv_scale_inv_format": qkv_scale_inv_format, + "do_scale_inv_format": do_scale_inv_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, + "deterministic": deterministic, + "cuda_graph": cuda_graph, + } ) return output_tensors diff --git a/transformer_engine/pytorch/csrc/extensions.h b/transformer_engine/pytorch/csrc/extensions.h index 5a4378ea328..f955e816fdb 100644 --- a/transformer_engine/pytorch/csrc/extensions.h +++ b/transformer_engine/pytorch/csrc/extensions.h @@ -94,36 +94,7 @@ 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 fused_attn_rng_state(const std::optional rng_gen, size_t elts_per_thread); 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 98d55d2b8b5..547cde9ee39 100644 --- a/transformer_engine/pytorch/csrc/extensions/attention.cpp +++ b/transformer_engine/pytorch/csrc/extensions/attention.cpp @@ -10,640 +10,16 @@ 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 +// Keep dropout on the framework's CUDA generator, including graph replay offsets. +// This helper does not construct or execute an attention graph. +at::Tensor fused_attn_rng_state(const std::optional rng_gen, + size_t elts_per_thread) { 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)}; + auto state = init_philox_state(gen, elts_per_thread); + auto result = at::empty({2}, at::TensorOptions().dtype(at::kLong).device(at::kCUDA)); + philox_unpack(state, result.data_ptr()); + return result; } at::Tensor fa_prepare_fwd(at::Tensor qkvi) { diff --git a/transformer_engine/pytorch/csrc/extensions/pybind.cpp b/transformer_engine/pytorch/csrc/extensions/pybind.cpp index c7f70182f36..0036d3f67aa 100644 --- a/transformer_engine/pytorch/csrc/extensions/pybind.cpp +++ b/transformer_engine/pytorch/csrc/extensions/pybind.cpp @@ -475,8 +475,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()); @@ -569,10 +567,9 @@ 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("fused_attn_rng_state", &transformer_engine::pytorch::fused_attn_rng_state, + "Reserve and unpack the CUDA generator's Philox state", py::arg("generator"), + py::arg("elts_per_thread")); 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,