diff --git a/docs/envvars.rst b/docs/envvars.rst index 3ced6a8e87..0fee105fd0 100644 --- a/docs/envvars.rst +++ b/docs/envvars.rst @@ -190,18 +190,18 @@ backend-selection overview. :Default: ``1`` :Description: Enable or disable UnfusedDotProductAttention backend (native PyTorch). When set to ``0``, UnfusedDotProductAttention will not be used. -.. envvar:: NVTE_FUSED_ATTN_BACKEND - - :Type: ``int`` (1 or 2) - :Default: Auto-selected - :Description: Request a cuDNN FusedAttention backend when that request is supported by the active fused-attention path. ``1`` = F16_arbitrary_seqlen (cuDNN, any seq len), ``2`` = FP8 backend. If not set, the backend is automatically selected based on the input configuration. BF16/FP16 attention uses sub-backend ``1`` when eligible. FP8 attention uses sub-backend ``2`` when FP8 DPA is enabled and supported by the architecture, cuDNN version, and input configuration. - .. envvar:: NVTE_FUSED_ATTN_USE_FAv2_BWD :Type: ``int`` (0 or 1) :Default: ``0`` :Description: When using FusedAttention, use FlashAttention-2 implementation for the backward pass instead of the cuDNN implementation. This can be useful due to performance differences between various versions of flash-attn and FusedAttention. +.. envvar:: NVTE_FUSED_ATTN_CACHE_DEBUG + + :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. + .. envvar:: NVTE_ALLOW_NONDETERMINISTIC_ALGO :Type: ``int`` (0 or 1) diff --git a/docs/examples/attention/attention.ipynb b/docs/examples/attention/attention.ipynb index 989661b543..c10ce17391 100644 --- a/docs/examples/attention/attention.ipynb +++ b/docs/examples/attention/attention.ipynb @@ -249,11 +249,7 @@ "NVTE_DEBUG = 0/1 # disables/enables debugging\n", "NVTE_DEBUG_LEVEL = 0/1/2 # enables logging.WARNING/INFO/DEBUG-level messages\n", "```\n", - "
\n", - "Note:\n", - " \n", - "These flags are supported in PyTorch only as of Transformer Engine 2.0. JAX support is expected to be added in the future.\n", - "
" + "These flags are available in both PyTorch and Jax in TE v2.20 onwards." ] }, { @@ -330,6 +326,51 @@ "!NVTE_DEBUG=1 NVTE_DEBUG_LEVEL=2 python example_attention.py" ] }, + { + "cell_type": "markdown", + "id": "fc0c9e3d", + "metadata": {}, + "source": [ + "For the FusedAttention backend, cuDNN graphs are cached and reused, and the following flag helps understand the activities taking place in the cache. `NVTE_FUSED_ATTN_CACHE_DEBUG` is supported in both PyTorch and JAX.\n", + "```\n", + "NVTE_FUSED_ATTN_CACHE_DEBUG = 0 # disables cache diagnostics\n", + "NVTE_FUSED_ATTN_CACHE_DEBUG = 1/2:ranks # enables cache diagnostics (events)/(tracing):select ranks for information collection\n", + "```\n", + "At level 1, these cache events are collected and reported in the end-of-run summary: `create_graph` (number of graphs created), `cache_graph` (number of graphs cached), `hit` (number of cache hits), `miss` (number of cache misses), `build_plans` (number of graphs whose plans are built), and `execute` (number of plan executions). These event counters are tallied up per backend (f16 and fp8), per pass (forward and backward), and per thread. They help paint a picture of the graph reuse rate and caching effectiveness. Another part of the end-of-run summary are the build times of various stages of the cuDNN graph. Out of the 5 cuDNN calls, `validate`, `build_operation_graph`, `create_execution_plans`, `check_support`, and `build_plans`, `build_plans` is the most expensive one, because it is where the kernels get compiled. It happens once per eligible graph and is expected to execute many times in a real-life run.\n", + "\n", + "For more verbose diagnostics, level 2 enables tracing for all the events taking place in the cache. For example, for every `miss` event, a \"MISS\" line is printed in the logs, as well as the cache key that triggered it. This helps identify if the cache is processing the right config as we intended.\n", + "\n", + "By default, only rank 0 is enabled for both level 1 and level 2. Users can use the \"level:ranks\" format to specify a select set of ranks to collect diagnostics on. For example, \"2:0,3\" means to enable level 2 diagnostics on ranks 0 and 3. An example of the level 2 diagnostics, which also includes the level 1 summary, is as follows, for a `tests/pytorch/attention/test_attention.py::test_dot_product_attention` test.\n", + "```\n", + "[FUSED-ATTN-CACHE] rank=0 | tid=0 dev=0 | f16 fwd MISS | train=1 det=0 cg=0 maxlogit=0 mask=0 bias=0 wl=-1 wr=-1 brd=0 softmax=0 scale_mode=0 dropout=0 attn_scale=1 qkv_dt=5 o_dt=5 do_dt=6 dqkv_dt=6 qkv_lay=0 o_fmt=0 do_fmt=8 dqkv_lay=26 qkv_sif=8 do_sif=8 b=8 h=16 hg=16 dqk=64 dv=64 sq=128 skv=128 tq=0 tkv=0 bb=0 btq=0 btkv=0 npk=0 npv=0 psk=0 psv=0 mppk=0 mppv=0 bias_b=0 bias_h=0 bias_sq=0 bias_skv=0\n", + "[FUSED-ATTN-CACHE] rank=0 | tid=0 dev=0 | f16 fwd CREATE_GRAPH | hit= 0, miss= 1, create_graph= 1, cache_graph= 0, build_plans= 0, execute= 0\n", + "[FUSED-ATTN-CACHE] rank=0 | tid=0 dev=0 | f16 fwd CACHE_GRAPH | hit= 0, miss= 1, create_graph= 1, cache_graph= 1, build_plans= 0, execute= 0\n", + "...\n", + "[FUSED-ATTN-CACHE] rank=0 | tid=0 dev=0 | f16 fwd HIT | train=1 det=0 cg=0 maxlogit=0 mask=0 bias=0 wl=-1 wr=-1 brd=0 softmax=0 scale_mode=0 dropout=0 attn_scale=1 qkv_dt=5 o_dt=5 do_dt=6 dqkv_dt=6 qkv_lay=0 o_fmt=0 do_fmt=8 dqkv_lay=26 qkv_sif=8 do_sif=8 b=8 h=16 hg=16 dqk=64 dv=64 sq=128 skv=128 tq=0 tkv=0 bb=0 btq=0 btkv=0 npk=0 npv=0 psk=0 psv=0 mppk=0 mppv=0 bias_b=0 bias_h=0 bias_sq=0 bias_skv=0\n", + "[FUSED-ATTN-CACHE] rank=0 | tid=0 dev=0 | f16 fwd BUILD_PLANS | hit= 3, miss= 1, create_graph= 1, cache_graph= 1, build_plans= 1, execute= 0\n", + "...\n", + "[FUSED-ATTN-CACHE] rank=0 | tid=0 dev=0 | f16 fwd EXECUTE | hit= 5, miss= 1, create_graph= 1, cache_graph= 1, build_plans= 1, execute= 1\n", + "...\n", + "[FUSED-ATTN-CACHE] rank=0 | ===== summary begin =====\n", + "[FUSED-ATTN-CACHE] rank=0 | tid=0 dev=0 | f16 fwd | hit= 5, miss= 1, create_graph= 1, cache_graph= 1, build_plans= 1, execute= 1\n", + "[FUSED-ATTN-CACHE] rank=0 | tid=0 dev=0 | f16 bwd | hit= 1, miss= 1, create_graph= 1, cache_graph= 1, build_plans= 0, execute= 0\n", + "[FUSED-ATTN-CACHE] rank=0 | tid=1 dev=0 | f16 bwd | hit= 4, miss= 0, create_graph= 0, cache_graph= 0, build_plans= 1, execute= 1\n", + "[FUSED-ATTN-CACHE] rank=0 | tid=all dev=all | f16 fwd | hit= 5, miss= 1, create_graph= 1, cache_graph= 1, build_plans= 1, execute= 1\n", + "[FUSED-ATTN-CACHE] rank=0 | tid=all dev=all | f16 bwd | hit= 5, miss= 1, create_graph= 1, cache_graph= 1, build_plans= 1, execute= 1\n", + "[FUSED-ATTN-CACHE] rank=0 | f16 fwd validate | calls=1 | time= 0.168 ms/call\n", + "[FUSED-ATTN-CACHE] rank=0 | f16 fwd build_operation_graph | calls=1 | time= 11.405 ms/call\n", + "[FUSED-ATTN-CACHE] rank=0 | f16 fwd create_execution_plans | calls=1 | time= 2.054 ms/call\n", + "[FUSED-ATTN-CACHE] rank=0 | f16 fwd check_support | calls=1 | time= 0.044 ms/call\n", + "[FUSED-ATTN-CACHE] rank=0 | f16 fwd build_plans | calls=1 | time= 327.083 ms/call\n", + "[FUSED-ATTN-CACHE] rank=0 | f16 bwd validate | calls=1 | time= 0.009 ms/call\n", + "[FUSED-ATTN-CACHE] rank=0 | f16 bwd build_operation_graph | calls=1 | time= 1.026 ms/call\n", + "[FUSED-ATTN-CACHE] rank=0 | f16 bwd create_execution_plans | calls=1 | time= 2.874 ms/call\n", + "[FUSED-ATTN-CACHE] rank=0 | f16 bwd check_support | calls=1 | time= 0.016 ms/call\n", + "[FUSED-ATTN-CACHE] rank=0 | f16 bwd build_plans | calls=1 | time= 514.599 ms/call\n", + "[FUSED-ATTN-CACHE] rank=0 | ===== summary end =====\n", + "```" + ] + }, { "cell_type": "markdown", "id": "611d8fdb", @@ -346,18 +387,7 @@ "NVTE_FUSED_ATTN = 0 # disables cuDNN attention; default = 1\n", "```\n", "\n", - "**cuDNN attention sub-backends:**\n", - "This environment variable allows users to express their preference of cuDNN attention sub-backends. However, the elected sub-backend will only be used *if* it is eligible, i.e. if it has support for the provided inputs and runtime environment.\n", - "```\n", - "NVTE_FUSED_ATTN_BACKEND = 1/2 # user preference of cuDNN sub-backend\n", - "```\n", - "\n", - "```\n", - "
\n", - "Note\n", - " \n", - "Environment variables NVTE_FLASH_ATTN, NVTE_UNFUSED_ATTN, NVTE_FUSED_ATTN_BACKEND, and NVTE_FUSED_ATTN_USE_FAv2_BWD are supported in PyTorch. NVTE_FUSED_ATTN and NVTE_ALLOW_NONDETERMINISTIC_ALGO are supported in both PyTorch and JAX.\n", - "
\n", + "Environment variables NVTE_FLASH_ATTN, NVTE_UNFUSED_ATTN, and NVTE_FUSED_ATTN_USE_FAv2_BWD are supported in PyTorch. NVTE_FUSED_ATTN and NVTE_ALLOW_NONDETERMINISTIC_ALGO are supported in both PyTorch and JAX.\n", "\n", "### 2.3 Example Tests\n", "\n", diff --git a/docs/examples/jax/attention_context_parallel.py b/docs/examples/jax/attention_context_parallel.py index 1557a30b7c..982ecc11b2 100644 --- a/docs/examples/jax/attention_context_parallel.py +++ b/docs/examples/jax/attention_context_parallel.py @@ -245,21 +245,22 @@ def context_parallel_supported() -> Tuple[bool, str]: return False, f"needs {cp_size} GPUs" has_kernel = is_fused_attn_kernel_available( - True, - dtype, - dtype, - QKVLayout.THD_THD_THD, - AttnBiasType.NO_BIAS, - AttnMaskType.PADDING_CAUSAL_MASK, - AttnSoftmaxType.VANILLA_SOFTMAX, - 0.0, - num_query_heads, - num_kv_heads, - seq, - seq, - head_dim, - head_dim, - window_size, + is_training=True, + batch_size=batch, + q_dtype=dtype, + kv_dtype=dtype, + qkv_layout=QKVLayout.THD_THD_THD, + attn_bias_type=AttnBiasType.NO_BIAS, + attn_mask_type=AttnMaskType.PADDING_CAUSAL_MASK, + softmax_type=AttnSoftmaxType.VANILLA_SOFTMAX, + dropout_probability=0.0, + q_num_heads=num_query_heads, + kv_num_heads=num_kv_heads, + q_max_seqlen=seq, + kv_max_seqlen=seq, + head_dim_qk=head_dim, + head_dim_v=head_dim, + window_size=window_size, ) if not has_kernel: return False, "no fused attention kernel for the THD SWA shape" diff --git a/docs/examples/jax/test_attention.py b/docs/examples/jax/test_attention.py index 3cc08271dc..b5db8c3535 100644 --- a/docs/examples/jax/test_attention.py +++ b/docs/examples/jax/test_attention.py @@ -87,21 +87,22 @@ def _context_parallel_supported(): return False, f"needs {cp_size} GPUs" has_kernel = is_fused_attn_kernel_available( - True, - jnp.bfloat16, - jnp.bfloat16, - QKVLayout.THD_THD_THD, - AttnBiasType.NO_BIAS, - AttnMaskType.PADDING_CAUSAL_MASK, - AttnSoftmaxType.VANILLA_SOFTMAX, - 0.0, - 128, - 8, - 65536, - 65536, - 128, - 128, - (8192, 0), + is_training=True, + batch_size=2, + q_dtype=jnp.bfloat16, + kv_dtype=jnp.bfloat16, + qkv_layout=QKVLayout.THD_THD_THD, + attn_bias_type=AttnBiasType.NO_BIAS, + attn_mask_type=AttnMaskType.PADDING_CAUSAL_MASK, + softmax_type=AttnSoftmaxType.VANILLA_SOFTMAX, + dropout_probability=0.0, + q_num_heads=128, + kv_num_heads=8, + q_max_seqlen=65536, + kv_max_seqlen=65536, + head_dim_qk=128, + head_dim_v=128, + window_size=(8192, 0), ) if not has_kernel: return False, "no fused attention kernel for the THD SWA shape" diff --git a/tests/jax/test_distributed_fused_attn.py b/tests/jax/test_distributed_fused_attn.py index b079d48aff..4813df730d 100644 --- a/tests/jax/test_distributed_fused_attn.py +++ b/tests/jax/test_distributed_fused_attn.py @@ -82,25 +82,6 @@ def impl_test_self_attn( is_training = True batch, seqlen, num_head, hidden = data_shape - if not is_fused_attn_kernel_available( - is_training, - dtype, - dtype, - QKVLayout.BS3HD, - attn_bias_type, - attn_mask_type, - softmax_type, - dropout_prob, - num_head, - num_head, - seqlen, - seqlen, - hidden, - hidden, - None, # no window - ): - pytest.skip("No FusedAttn backend found") - col_ref = self.generate_collectives_count_ref( mesh_shape, mesh_axes, @@ -234,25 +215,6 @@ def test_cross_attn( batch, seqlen, num_head, hidden = data_shape - if not is_fused_attn_kernel_available( - is_training, - dtype, - dtype, - QKVLayout.BSHD_BS2HD, - attn_bias_type, - attn_mask_type, - softmax_type, - dropout_prob, - num_head, - num_head, - seqlen, - seqlen, - hidden, - hidden, - None, # no window - ): - pytest.skip("No FusedAttn backend found") - col_ref = self.generate_collectives_count_ref() runner = FusedAttnRunner( batch, @@ -491,6 +453,7 @@ def impl_test_context_parallel_attn( def check_has_backend_for_mask(mask_type): return is_fused_attn_kernel_available( is_training, + batch, dtype, dtype, qkv_layout, diff --git a/tests/jax/test_fused_attn.py b/tests/jax/test_fused_attn.py index d68e409331..a6ba34be6a 100644 --- a/tests/jax/test_fused_attn.py +++ b/tests/jax/test_fused_attn.py @@ -4,7 +4,7 @@ """Tests for fused attention""" import os from enum import Enum, auto -from dataclasses import dataclass, field +from dataclasses import dataclass, field, replace from functools import partial from math import sqrt from typing import Any, Callable, Mapping, Tuple, Optional, Dict @@ -48,11 +48,14 @@ ) from distributed_test_base import assert_equal_collectives -from utils import assert_allclose, print_debug_tensor_stats +from utils import assert_allclose, get_test_level, print_debug_tensor_stats # Get determinism _deterministic = not bool(int(os.getenv("NVTE_ALLOW_NONDETERMINISTIC_ALGO", "1"))) +# CI test level +_TEST_LEVEL = get_test_level() + @pytest.fixture(autouse=True, scope="module") def init(): @@ -383,6 +386,48 @@ def score_mod(_graph, score, _tensors): ) +def test_fused_attn_backend_message(): + """Test the error messaging of the fused attention backend query.""" + baseline = 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=8, + kv_num_heads=8, + q_max_seqlen=128, + kv_max_seqlen=128, + head_dim_qk=64, + head_dim_v=64, + window_size=(-1, -1), + attn_scale=0.125, + ) + + # One of TE's rules is violated and the error message is surfaced + backend, message = replace( + baseline, attn_bias_type=AttnBiasType.PRE_SCALE_BIAS + ).get_fused_attn_backend() + assert backend == NVTE_Fused_Attn_Backend.NVTE_No_Backend + assert message == "Fused attention does not support pre-scale bias." + + # No error message if supported; otherwise skip the test + backend, message = baseline.get_fused_attn_backend() + if backend == NVTE_Fused_Attn_Backend.NVTE_No_Backend: + pytest.skip(f"FusedAttention does not support the baseline config: {message}") + assert message == "" + + # All TE rules have cleared; now gets rejected by cuDNN's support check + # cuDNN's error message might change across cuDNN versions, so only verify the presence of the string + backend, message = replace(baseline, head_dim_qk=1024, head_dim_v=1024).get_fused_attn_backend() + assert backend == NVTE_Fused_Attn_Backend.NVTE_No_Backend + assert message != "" + + class BiasShape(Enum): """ Enum class to represent the different bias shapes used in the fused attention. @@ -482,6 +527,27 @@ def _get_max_segments_per_sequence(self): return 1 def _check_configs(self): + # Trim SWA configs for L0 and L1 to reduce test time; need to trim more in future test refactoring. + if self.window_size is not None and ( + self.dropout_prob != 0.0 or self.attn_bias_type is not AttnBiasType.NO_BIAS + ): + if _TEST_LEVEL == "L0" and ( + self.softmax_type != AttnSoftmaxType.VANILLA_SOFTMAX + or self.dtype != jnp.bfloat16 + or self.attn_bias_type is not AttnBiasType.POST_SCALE_BIAS + or self.attn_mask_type is not AttnMaskType.NO_MASK + ): + pytest.skip( + "Trimmed SWA+bias/dropout config: only vanilla-softmax + bf16 + post_scale_bias" + " + no-mask runs at L0" + ) + if _TEST_LEVEL == "L1" and ( + self.dtype != jnp.float16 or self.softmax_type != AttnSoftmaxType.LEARNABLE_SOFTMAX + ): + pytest.skip( + "Trimmed SWA+bias/dropout config: only float16 + learnable-softmax runs at L1" + ) + # TODO(KshitijLakhani): probably add/move this to is_fused_attn_available if self.qkv_layout.is_thd() and not self.attn_mask_type.is_padding(): pytest.skip("THD format requires padding masks.") @@ -584,8 +650,21 @@ def _check_configs(self): "is either BSHD_BSHD_BSHD or THD_THD_THD" ) - self.backend = FusedAttnHelper( + bias_batch = bias_heads = bias_seqlen_q = bias_seqlen_kv = None + if self.attn_bias_type == AttnBiasType.POST_SCALE_BIAS: + if self.bias_shape == BiasShape._1HSS: + bias_batch, bias_heads = 1, self.num_heads_q + elif self.bias_shape == BiasShape._B1SS: + bias_batch, bias_heads = self.batch_size, 1 + elif self.bias_shape == BiasShape._BHSS: + bias_batch, bias_heads = self.batch_size, self.num_heads_q + elif self.bias_shape == BiasShape._11SS: + bias_batch, bias_heads = 1, 1 + bias_seqlen_q, bias_seqlen_kv = self.max_seqlen_q, self.max_seqlen_kv + + self.backend, message = FusedAttnHelper( self.is_training, + self.batch_size, self.dtype, self.dtype, self.qkv_layout, @@ -600,9 +679,14 @@ def _check_configs(self): self.head_dim_qk, self.head_dim_v, (-1, -1) if self.window_size is None else self.window_size, + bottom_right_diagonal=self.attn_mask_type.is_bottom_right(), + bias_batch=bias_batch, + bias_heads=bias_heads, + bias_seqlen_q=bias_seqlen_q, + bias_seqlen_kv=bias_seqlen_kv, ).get_fused_attn_backend() if self.backend != NVTE_Fused_Attn_Backend.NVTE_F16_arbitrary_seqlen: - pytest.skip("Unsupported inputs combination or device compute capability.") + pytest.skip(message) if ( self.attn_bias_type == AttnBiasType.POST_SCALE_BIAS diff --git a/tests/jax/test_fused_attn_score_mod.py b/tests/jax/test_fused_attn_score_mod.py index 7d1137b21f..4a9e730f55 100644 --- a/tests/jax/test_fused_attn_score_mod.py +++ b/tests/jax/test_fused_attn_score_mod.py @@ -18,7 +18,7 @@ ) from transformer_engine.jax.cpp_extensions import make_fused_attn_score_mod_config from transformer_engine.jax.flax import transformer as flax_transformer -from transformer_engine_jax import get_device_compute_capability +from transformer_engine_jax import get_device_compute_capability, NVTE_Fused_Attn_Backend from test_fused_attn import FusedAttnRunner, SeqDescFormat @@ -397,9 +397,17 @@ def _identity_score_mod(_graph, score, _tensors): def _install_fake_flax_fused_attn(monkeypatch, *, kernel_available=True): captured = {} - def fake_fused_attn_kernel_check(*args, **kwargs): - captured.setdefault("kernel_checks", []).append((args, kwargs)) - return kernel_available + class FakeFusedAttnHelper: + def __init__(self, *args, **kwargs): + captured.setdefault("kernel_checks", []).append((args, kwargs)) + + def get_fused_attn_backend(self): + if kernel_available: + return NVTE_Fused_Attn_Backend.NVTE_F16_arbitrary_seqlen, "" + return ( + NVTE_Fused_Attn_Backend.NVTE_No_Backend, + "fake FusedAttnHelper: no fused attention backend available for this configuration", + ) def fake_fused_attn( qkv, @@ -455,11 +463,7 @@ def fake_fused_attn( ) return qkv[0] - monkeypatch.setattr( - flax_transformer, - "is_fused_attn_kernel_available", - fake_fused_attn_kernel_check, - ) + monkeypatch.setattr(flax_transformer, "FusedAttnHelper", FakeFusedAttnHelper) monkeypatch.setattr(flax_transformer, "fused_attn", fake_fused_attn) return captured @@ -534,7 +538,7 @@ def test_dot_product_attention_plumbs_score_mod_to_fused_attn(monkeypatch): assert captured["attn_bias_type"] is AttnBiasType.NO_BIAS assert captured["qkv_layout"] is QKVLayout.BSHD_BSHD_BSHD assert captured["softmax_type"] is AttnSoftmaxType.VANILLA_SOFTMAX - assert captured["kernel_checks"][0][0][3] is QKVLayout.BSHD_BSHD_BSHD + assert captured["kernel_checks"][0][0][4] is QKVLayout.BSHD_BSHD_BSHD def test_dot_product_attention_unpacks_packed_score_mod_to_separate_layout(monkeypatch): @@ -558,7 +562,7 @@ def test_dot_product_attention_unpacks_packed_score_mod_to_separate_layout(monke assert captured["qkv"][0].shape == (1, 8, 1, 16) assert captured["qkv_layout"] is QKVLayout.BSHD_BSHD_BSHD assert captured["score_mod"] is _identity_score_mod - assert captured["kernel_checks"][0][0][3] is QKVLayout.BSHD_BSHD_BSHD + assert captured["kernel_checks"][0][0][4] is QKVLayout.BSHD_BSHD_BSHD def test_multi_head_attention_plumbs_score_mod_to_dot_product_attention(monkeypatch): diff --git a/tests/jax/utils.py b/tests/jax/utils.py index c5e564dbc7..f0a6928f1b 100644 --- a/tests/jax/utils.py +++ b/tests/jax/utils.py @@ -119,13 +119,27 @@ def combine_biases(*masks: Optional[Array]): return mask +TEST_LEVELS = ("L0", "L1", "L2") + + +def get_test_level(): + """ + Returns the test level specified in the environment variable, NVTE_JAX_UNITTEST_LEVEL. + """ + test_level = os.environ.get("NVTE_JAX_UNITTEST_LEVEL", TEST_LEVELS[0]) + if test_level not in TEST_LEVELS: + raise ValueError( + f"Unsupported test level {test_level!r}, expected one of {', '.join(TEST_LEVELS)}" + ) + return test_level + + def get_parameters_for_test_level(param_dict: dict): """ Takes an input dictionary of parameters keyed by test type "L0", etc. Returns the parameters for the test level specified in the environment variable """ - DEFAULT_TEST_LEVEL = "L0" - test_level = os.environ.get("NVTE_JAX_UNITTEST_LEVEL", DEFAULT_TEST_LEVEL) + test_level = get_test_level() if test_level not in param_dict: raise ValueError("Unsupported test level") return param_dict[test_level] diff --git a/tests/pytorch/attention/run_graph_cache.py b/tests/pytorch/attention/run_graph_cache.py new file mode 100644 index 0000000000..534711515c --- /dev/null +++ b/tests/pytorch/attention/run_graph_cache.py @@ -0,0 +1,107 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""Worker for test_attention.py::test_fused_attn_graph_cache. + +Runs a fixed sequence of support queries and executions against the FusedAttention graph cache and +collects level-2 diagnostics with NVTE_FUSED_ATTN_CACHE_DEBUG=2 for different phases: + query the first support query for a config -- the miss that creates its fwd/bwd graphs + requery the identical query again -- hits the cache + exec forward and backward of that config -- hits the cache and builds plans before execution + rescale the same execution with only softmax_scale changed -- should hit the cache and reuse the plans, + since attn_scale is normalized out of the cache key + reshape a query differing in max_seqlen -- should miss and recreate the graphs +""" + +import os +import pathlib +import sys + +import torch + +_current_file = pathlib.Path(__file__).resolve() +sys.path = [str(_current_file.parent.parent)] + sys.path + +from transformer_engine.pytorch import DotProductAttention +from transformer_engine.pytorch.attention.dot_product_attention import _attention_backends +from utils import ModelConfig, get_available_attention_backends + +DTYPE = torch.bfloat16 +QKV_FORMAT = "bshd" +QKV_LAYOUT = "bshd_bshd_bshd" +DETERMINISTIC = ( + not bool(int(os.getenv("NVTE_ALLOW_NONDETERMINISTIC_ALGO", "1"))) + or torch.are_deterministic_algorithms_enabled() +) + + +def mark_phase(name: str) -> None: + """Delimit the cache events of one phase from the next one's.""" + sys.stderr.write(f"[CACHE-TEST] phase={name}\n") + sys.stderr.flush() + + +def query(config: ModelConfig) -> bool: + """Run one backend support query and report whether FusedAttention supports + the configuration. If so, it populates the cache without executing anything.""" + available_backends, _, fused_attn_backends = get_available_attention_backends( + config, qkv_dtype=DTYPE, qkv_layout=QKV_LAYOUT, deterministic=DETERMINISTIC + ) + _, fused_attn_supported, _ = available_backends + return fused_attn_supported and len(fused_attn_backends) > 0 + + +def execute(config: ModelConfig, softmax_scale: float) -> None: + """Run a forward and backward pass of `config` on the FusedAttention backend.""" + block = DotProductAttention( + config.num_heads, + config.head_dim_qk, + attention_dropout=config.dropout_p, + qkv_format=QKV_FORMAT, + attn_mask_type=config.attn_mask_type, + softmax_scale=softmax_scale, + layer_number=1, + attention_type=config.attn_type, + ).to(dtype=DTYPE, device="cuda") + shape = (config.batch_size, config.max_seqlen_q, config.num_heads, config.head_dim_qk) + q, k, v = [torch.randn(shape, dtype=DTYPE, device="cuda", requires_grad=True) for _ in range(3)] + out = block(q, k, v, core_attention_bias_type=config.attn_bias_type) + out.backward(torch.randn_like(out)) + torch.cuda.synchronize() + + +def main() -> int: + torch.manual_seed(1234) + config = ModelConfig(2, 512, 8, 64) + reshaped = ModelConfig(2, 256, 8, 64) + + mark_phase("query") + fused_available = query(config) + print(f"[CACHE-TEST] fused={int(fused_available)}", flush=True) + if not fused_available: + return 0 + + mark_phase("requery") + query(config) + + os.environ["NVTE_FLASH_ATTN"] = "0" + os.environ["NVTE_FUSED_ATTN"] = "1" + os.environ["NVTE_UNFUSED_ATTN"] = "0" + _attention_backends["backend_selection_requires_update"] = True + + mark_phase("exec") + execute(config, softmax_scale=0.125) + + mark_phase("rescale") + execute(config, softmax_scale=0.25) + + mark_phase("reshape") + query(reshaped) + + mark_phase("done") + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/tests/pytorch/attention/test_attention.py b/tests/pytorch/attention/test_attention.py index 0f5cdfa647..2770520a4b 100644 --- a/tests/pytorch/attention/test_attention.py +++ b/tests/pytorch/attention/test_attention.py @@ -1,12 +1,17 @@ # Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. +import collections +import copy import logging import os +import re +import subprocess import sys import pathlib import copy -from typing import Any, Dict, Tuple, Union +from dataclasses import replace +from typing import Any, Dict, List, NamedTuple, Tuple, Union from packaging.version import Version as PkgVersion import pytest @@ -30,13 +35,17 @@ ) from transformer_engine.pytorch.attention.dot_product_attention.utils import ( FlashAttentionUtils, + FusedAttentionParams, _get_supported_versions, check_set_window_size, + get_fused_attn_spec, ) from transformer_engine.pytorch.attention import RotaryPositionEmbedding import transformer_engine.pytorch.cpp_extensions as ext from transformer_engine.pytorch.cpp_extensions.fused_attn import ( FusedAttnBackend, + QKVFormat, + QKVLayout, fused_attn_bwd, fused_attn_fwd, ) @@ -129,6 +138,81 @@ def test_flash_attention_supported_version_message(): ) +def test_fused_attn_backend_message(): + """Test the error messaging of the fused attention backend query.""" + baseline = FusedAttentionParams( + qkv_layout=tex.NVTE_QKV_Layout.NVTE_BSHD_BSHD_BSHD, + dqkv_layout=tex.NVTE_QKV_Layout.NVTE_BSHD_BSHD_BSHD, + o_format=tex.NVTE_QKV_Format.NVTE_BSHD, + do_format=tex.NVTE_QKV_Format.NVTE_BSHD, + batch_size=2, + num_attn_heads=8, + num_gqa_groups=8, + head_dim_qk=64, + head_dim_v=64, + max_seqlen_q=128, + max_seqlen_kv=128, + attn_scale=0.125, + ) + + # One of TE's rules is violated and the error message is surfaced + backend, message = tex.get_fused_attn_backend( + replace(baseline, bias_type=tex.NVTE_Bias_Type.NVTE_PRE_SCALE_BIAS) + ) + assert backend == tex.NVTE_Fused_Attn_Backend.NVTE_No_Backend + assert message == "Fused attention does not support pre-scale bias." + + # No error message if supported; otherwise skip the test + backend, message = tex.get_fused_attn_backend(baseline) + if backend == tex.NVTE_Fused_Attn_Backend.NVTE_No_Backend: + pytest.skip(f"FusedAttention does not support the baseline config: {message}") + assert message == "" + + # All TE rules have cleared; now gets rejected by cuDNN's support check + # cuDNN's error message might change across cuDNN versions, so only verify the presence of the string + backend, message = tex.get_fused_attn_backend( + replace(baseline, head_dim_qk=1024, head_dim_v=1024) + ) + assert backend == tex.NVTE_Fused_Attn_Backend.NVTE_No_Backend + assert message != "" + + if not fp8_attn_available: + return + + # The same checks in FP8; the dtypes, formats and scaling mode come from get_fused_attn_spec so + # that the config matches what FusedAttnFunc feeds the kernels for this recipe + fp8_recipe = recipe.DelayedScaling(fp8_format=recipe.Format.HYBRID, fp8_dpa=True) + spec = get_fused_attn_spec(fp8_recipe, torch.float8_e4m3fn, "bshd_bshd_bshd", cs_o_in_f16=True) + fp8_baseline = replace( + baseline, + scaling_mode=spec.scaling_mode, + qkv_dtype=spec.qkv, + o_dtype=spec.o, + do_dtype=spec.do, + dqkv_dtype=spec.dqkv, + qkv_layout=QKVLayout[spec.qkv_layout], + o_format=QKVFormat[spec.o_format], + do_format=QKVFormat[spec.do_format], + dqkv_layout=QKVLayout[spec.dqkv_layout], + qkv_scale_inv_format=QKVFormat[spec.scale_inv_format], + do_scale_inv_format=QKVFormat[spec.scale_inv_format], + ) + + # No error message if supported; otherwise skip the test + backend, message = tex.get_fused_attn_backend(fp8_baseline) + if backend == tex.NVTE_Fused_Attn_Backend.NVTE_No_Backend: + pytest.skip(f"FusedAttention does not support the baseline FP8 config: {message}") + assert backend == tex.NVTE_Fused_Attn_Backend.NVTE_FP8 + assert message == "" + + # All TE FP8 rules have cleared; now gets rejected by cuDNN's FP8 support check + backend, message = tex.get_fused_attn_backend( + replace(fp8_baseline, head_dim_qk=1024, head_dim_v=1024) + ) + assert backend == tex.NVTE_Fused_Attn_Backend.NVTE_No_Backend + assert message != "" + + # Define F16 data types to test param_types = [torch.float16] if is_bf16_available(): @@ -212,6 +296,13 @@ def test_dot_product_attention( "Setting is_training to False as cuDNN does not support dbias for" f" {config.bias_shape=} " ) + # Generate token counts for THD so the support query and execution will use the same values + num_tokens_q, num_tokens_kv = None, None + if qkv_format == "thd": + reset_rng_states() + seqlens = _generate_seqlens(config, qkv_format, pad_between_seqs) + num_tokens_q = int(seqlens.cu_q_after_pad[-1]) + num_tokens_kv = int(seqlens.cu_kv_after_pad[-1]) available_backends, _, fused_attn_backends = get_available_attention_backends( config, qkv_dtype=dtype, @@ -219,6 +310,8 @@ def test_dot_product_attention( pad_between_seqs=pad_between_seqs, is_training=is_training, deterministic=_deterministic, + num_tokens_q=num_tokens_q, + num_tokens_kv=num_tokens_kv, ) flash_attn_supported, fused_attn_supported, unfused_attn_supported = available_backends @@ -231,6 +324,8 @@ def test_dot_product_attention( pad_between_seqs=pad_between_seqs, is_training=is_training, deterministic=_deterministic, + num_tokens_q=num_tokens_q, + num_tokens_kv=num_tokens_kv, ) flash_attn_supported, fused_attn_supported, unfused_attn_supported = available_backends @@ -297,6 +392,97 @@ def test_dot_product_attention( torch.testing.assert_close(fused_attn_bwd[i], flash_attn_bwd[i], **tols) +_CACHE_EVENT = re.compile( + r"\[FUSED-ATTN-CACHE\]\s+(?:rank=\d+\s+\|\s+)?tid=\d+\s+dev=-?\d+\s+\|\s+" + r"(?Pf16|fp8)\s+(?Pfwd|bwd)\s+" + r"(?PCREATE_GRAPH|CACHE_GRAPH|BUILD_PLANS|EXECUTE|MISS|HIT)\b(?P.*)" +) +_CACHE_PHASE = re.compile(r"\[CACHE-TEST\] phase=(?P\w+)") + + +@pytest.mark.skipif(get_cudnn_version() < (8, 9, 1), reason="cuDNN 8.9.1+ is required.") +def test_fused_attn_graph_cache(): + """Test FusedAttention graph cache with level 2 diagnostics. It runs a subprocess + to avoid contamination from other tests as the counters are process-wide. + """ + worker = _current_file.parent / "run_graph_cache.py" + result = subprocess.run( + [sys.executable, str(worker)], + env={ + **os.environ, + "NVTE_FUSED_ATTN_CACHE_DEBUG": "2", + "PYTHONUNBUFFERED": "1", + }, + capture_output=True, + text=True, + timeout=900, + check=False, + ) + assert result.returncode == 0, f"Graph cache worker exited with {result.returncode}." + if "[CACHE-TEST] fused=1" not in result.stdout: + pytest.skip("No FusedAttention backend for the graph cache test config.") + + # Group the cache diagnostics: + # events[phase][(pass, event)] = count of events + # miss_keys[phase][pass] = set of distinct cache keys that missed + events = collections.defaultdict(collections.Counter) + miss_keys = collections.defaultdict(lambda: collections.defaultdict(set)) + phase = None + for line in result.stderr.splitlines(): + phase_match = _CACHE_PHASE.search(line) + if phase_match is not None: + phase = phase_match.group("name") + continue + event_match = _CACHE_EVENT.search(line) + if event_match is None or phase is None: + continue + event_pass, event = event_match.group("pass"), event_match.group("event") + events[phase][(event_pass, event)] += 1 + if event == "MISS": + miss_keys[phase][event_pass].add(event_match.group("rest").split("|")[-1].strip()) + + context = f"\n--- stderr ---\n{result.stderr[-8000:]}" + + for pass_name in ("fwd", "bwd"): + + def expect(phase, event, expected, reason, actual=None, pass_name=pass_name): + if actual is None: + actual = events[phase][(pass_name, event)] + ok = actual >= 1 if expected == "1+" else actual == expected + assert ( + ok + ), f"{pass_name} {phase}: expected {event}={expected}, got {actual} ({reason}){context}" + + expect("query", "MISS", 1, "expected one cold miss") + expect("query", "CREATE_GRAPH", 1, "expected one build") + expect("query", "CACHE_GRAPH", 1, "build was not cached") + expect("query", "BUILD_PLANS", 0, "query compiled kernels") + + expect("requery", "MISS", 0, "repeated query missed") + expect("requery", "CREATE_GRAPH", 0, "repeated query rebuilt") + expect("requery", "HIT", "1+", "repeated query never looked") + + expect("exec", "MISS", 0, "execution missed the query's graph") + expect("exec", "CREATE_GRAPH", 0, "execution rebuilt the graph") + expect("exec", "EXECUTE", "1+", "fused attention never ran") + expect("exec", "BUILD_PLANS", 1, "expected one plan build") + + expect("rescale", "MISS", 0, "attn_scale changed the key") + expect("rescale", "CREATE_GRAPH", 0, "attn_scale forced a build") + expect("rescale", "BUILD_PLANS", 0, "attn_scale recompiled") + expect("rescale", "EXECUTE", "1+", "rescaled run did not execute") + + expect("reshape", "MISS", 1, "expected one miss") + expect("reshape", "CREATE_GRAPH", 1, "expected one build") + expect( + "reshape", + "MISS keys", + 1, + "more than one new cache key", + actual=len(miss_keys["reshape"][pass_name]), + ) + + @pytest.mark.skipif(get_cudnn_version() < (8, 9, 1), reason="cuDNN 8.9.1+ is required.") @pytest.mark.parametrize("dtype", param_types) @pytest.mark.parametrize("model_configs", [model_configs_base]) @@ -642,6 +828,9 @@ def test_dpa_softmax(dtype, model_configs, model): @pytest.mark.parametrize("model", model_configs_softmax.keys()) def test_dpa_softmax_thd(dtype, model_configs, model): """Test DotProductAttention module with different softmax types""" + config = model_configs[model] + if "padding" not in config.attn_mask_type: + pytest.skip(f"Duplicate test to others with THD and padding mask.") test_dot_product_attention(dtype, model_configs, model, True, "thd_thd_thd", False, False) @@ -912,6 +1101,9 @@ def test_dpa_bias_shapes(dtype, model_configs, model): @pytest.mark.parametrize("qkv_layout", ["thd_thd_thd", "sbhd_sbhd_sbhd"]) def test_dpa_sliding_window(dtype, model_configs, model, qkv_layout): """Test DotProductAttention module with sliding window attention""" + config = model_configs[model] + if qkv_layout == "thd_thd_thd" and "padding" not in config.attn_mask_type: + pytest.skip(f"Duplicate test to others with THD and padding mask.") test_dot_product_attention(dtype, model_configs, model, False, qkv_layout, True, False) @@ -1189,38 +1381,26 @@ def make_dot_product_attention( return block -def run_dot_product_attention( - dtype: torch.dtype, - config: ModelConfig, - backend: str, - ckpt_attn: bool, - qkv_layout: str, - pad_between_seqs: bool, - is_training: bool, - declarative_packed: bool = False, - forward_kwargs: Dict[str, Any] = None, -) -> Tuple[torch.Tensor, Tuple[torch.Tensor, torch.Tensor, torch.Tensor]]: - """Run DotProductAttention module with one forward pass and one backward pass. +class _Seqlens(NamedTuple): + """Sequence lengths for one test case, before and after inter-sequence padding.""" - With declarative_packed=True (packed qkv_layout only), the packed buffer is - passed to DotProductAttention directly via qkv_layer/kv_layer instead of - slicing it into q/k/v views, and input gradients are read off the packed - buffer itself.""" - # Set RNG and environment varables - reset_rng_states() - os.environ["NVTE_FLASH_ATTN"] = "0" - os.environ["NVTE_FUSED_ATTN"] = "0" - os.environ["NVTE_UNFUSED_ATTN"] = "0" - if backend == "FlashAttention": - os.environ["NVTE_FLASH_ATTN"] = "1" - if backend == "FusedAttention": - os.environ["NVTE_FUSED_ATTN"] = "1" - if backend == "UnfusedDotProductAttention": - os.environ["NVTE_UNFUSED_ATTN"] = "1" - _attention_backends["backend_selection_requires_update"] = True + q: torch.Tensor + kv: torch.Tensor + cu_q: torch.Tensor + cu_kv: torch.Tensor + q_after_pad: torch.Tensor + kv_after_pad: torch.Tensor + cu_q_after_pad: torch.Tensor + cu_kv_after_pad: torch.Tensor + pad_len: Union[List[int], torch.Tensor] - # Create seqlens - qkv_format = "".join([i for i in qkv_layout.split("_")[0] if i.isalpha()]) + +def _generate_seqlens( + config: ModelConfig, + qkv_format: str, + pad_between_seqs: bool, +) -> _Seqlens: + """Draw the sequence lengths for one test case.""" if "padding" in config.attn_mask_type or qkv_format == "thd": if config.attn_type == "self": seqlens_q = torch.randint( @@ -1260,6 +1440,63 @@ def run_dot_product_attention( cu_seqlens_q_after_pad[1:] = torch.cumsum(seqlens_q_after_pad, dim=0) cu_seqlens_kv_after_pad[1:] = torch.cumsum(seqlens_kv_after_pad, dim=0) + return _Seqlens( + seqlens_q, + seqlens_kv, + cu_seqlens_q, + cu_seqlens_kv, + seqlens_q_after_pad, + seqlens_kv_after_pad, + cu_seqlens_q_after_pad, + cu_seqlens_kv_after_pad, + pad_len, + ) + + +def run_dot_product_attention( + dtype: torch.dtype, + config: ModelConfig, + backend: str, + ckpt_attn: bool, + qkv_layout: str, + pad_between_seqs: bool, + is_training: bool, + declarative_packed: bool = False, + forward_kwargs: Dict[str, Any] = None, +) -> Tuple[torch.Tensor, Tuple[torch.Tensor, torch.Tensor, torch.Tensor]]: + """Run DotProductAttention module with one forward pass and one backward pass. + + With declarative_packed=True (packed qkv_layout only), the packed buffer is + passed to DotProductAttention directly via qkv_layer/kv_layer instead of + slicing it into q/k/v views, and input gradients are read off the packed + buffer itself.""" + # Set RNG and environment varables + reset_rng_states() + os.environ["NVTE_FLASH_ATTN"] = "0" + os.environ["NVTE_FUSED_ATTN"] = "0" + os.environ["NVTE_UNFUSED_ATTN"] = "0" + if backend == "FlashAttention": + os.environ["NVTE_FLASH_ATTN"] = "1" + if backend == "FusedAttention": + os.environ["NVTE_FUSED_ATTN"] = "1" + if backend == "UnfusedDotProductAttention": + os.environ["NVTE_UNFUSED_ATTN"] = "1" + _attention_backends["backend_selection_requires_update"] = True + + # Create seqlens + qkv_format = "".join([i for i in qkv_layout.split("_")[0] if i.isalpha()]) + ( + seqlens_q, + seqlens_kv, + cu_seqlens_q, + cu_seqlens_kv, + seqlens_q_after_pad, + seqlens_kv_after_pad, + cu_seqlens_q_after_pad, + cu_seqlens_kv_after_pad, + pad_len, + ) = _generate_seqlens(config, qkv_format, pad_between_seqs) + # Create attention mask if padding attention_mask = None if "padding" in config.attn_mask_type: @@ -1639,6 +1876,13 @@ def test_transformer_layer( # Test backend availability is_training = True + # Get the token counts for THD and use them for both support query and actual execution + num_tokens_q, num_tokens_kv = None, None + if qkv_format == "thd": + reset_rng_states() + seqlens = _generate_seqlens(config, qkv_format, pad_between_seqs=False) + num_tokens_q = int(seqlens.cu_q[-1]) + num_tokens_kv = int(seqlens.cu_kv[-1]) available_backends, _, fused_attn_backends = get_available_attention_backends( config, qkv_dtype=dtype, @@ -1647,6 +1891,8 @@ def test_transformer_layer( ), is_training=is_training, deterministic=_deterministic, + num_tokens_q=num_tokens_q, + num_tokens_kv=num_tokens_kv, ) flash_attn_supported, fused_attn_supported, unfused_attn_supported = available_backends if not fused_attn_supported: @@ -1661,6 +1907,8 @@ def test_transformer_layer( ), is_training=is_training, deterministic=_deterministic, + num_tokens_q=num_tokens_q, + num_tokens_kv=num_tokens_kv, ) flash_attn_supported, fused_attn_supported, unfused_attn_supported = available_backends @@ -1979,12 +2227,24 @@ def test_dpa_fp8_extra_state(model, dtype): config = model_configs_fp8_extra_state[model] # Test backend availability is_training = True + fp8_recipe = recipe.DelayedScaling( + margin=0, + fp8_format=recipe.Format.HYBRID, + amax_history_len=1, + amax_compute_algo="most_recent", + fp8_dpa=True, + ) + fp8_meta = {} + fp8_meta["recipe"] = fp8_recipe available_backends, _, fused_attn_backends = get_available_attention_backends( config, qkv_dtype=torch.float8_e4m3fn, + nominal_dtype=dtype, qkv_layout="sb3hd", is_training=is_training, deterministic=_deterministic, + fp8=True, + fp8_meta=fp8_meta, ) flash_attn_supported, fused_attn_supported, unfused_attn_supported = available_backends if not fused_attn_supported and not flash_attn_supported: @@ -2171,6 +2431,70 @@ def _get_fp8_vs_f16_config(model, qkv_layout): return config +def _dpa_fp8_vs_f16_seqlens(config, qkv_format): + """Draw the sequence lengths for one test_dpa_fp8_vs_f16 case.""" + if "padding" in config.attn_mask_type or qkv_format == "thd": + if config.attn_type == "self": + seqlens_q = torch.randint( + 1, config.max_seqlen_q, [config.batch_size], dtype=torch.int32, device="cuda" + ) + seqlens_kv = seqlens_q + if config.attn_type == "cross": + seqlens_q = torch.randint( + 1, config.max_seqlen_q, [config.batch_size], dtype=torch.int32, device="cuda" + ) + seqlens_kv = torch.randint( + 1, config.max_seqlen_kv, [config.batch_size], dtype=torch.int32, device="cuda" + ) + else: + seqlens_q = torch.full( + [config.batch_size], config.max_seqlen_q, dtype=torch.int32, device="cuda" + ) + seqlens_kv = torch.full( + [config.batch_size], config.max_seqlen_kv, dtype=torch.int32, device="cuda" + ) + return seqlens_q, seqlens_kv + + +def _mha_fp8_vs_f16_seqlens(config, qkv_format): + """Draw the sequence lengths for one test_mha_fp8_vs_f16 case. + + These come from a dedicated generator rather than the default one, because the module the run + builds beforehand consumes randomness the backend probe cannot replay.""" + gen = torch.Generator(device="cuda") + gen.manual_seed(seed) + + def draw(high, count): + return torch.randint(1, high, [count], dtype=torch.int32, device="cuda", generator=gen) + + if "padding" in config.attn_mask_type or qkv_format == "thd": + + def random_seqlens(max_seqlen): + if qkv_format != "thd": + return draw(max_seqlen, config.batch_size) + # Reserve seven positions so total-token alignment only increases the final length. + return torch.cat((draw(max_seqlen, config.batch_size - 1), draw(max_seqlen - 6, 1))) + + if config.attn_type == "self": + seqlens_q = random_seqlens(config.max_seqlen_q) + seqlens_kv = seqlens_q + if config.attn_type == "cross": + seqlens_q = random_seqlens(config.max_seqlen_q) + seqlens_kv = random_seqlens(config.max_seqlen_kv) + else: + seqlens_q = torch.full( + [config.batch_size], config.max_seqlen_q, dtype=torch.int32, device="cuda" + ) + seqlens_kv = torch.full( + [config.batch_size], config.max_seqlen_kv, dtype=torch.int32, device="cuda" + ) + if qkv_format == "thd": + # FP8 Linear flattens THD input to [t, h*d], so align total tokens for cuBLAS. + seqlens_q[-1] += -seqlens_q.sum() % 8 + seqlens_kv[-1] += -seqlens_kv.sum() % 8 + return seqlens_q, seqlens_kv + + @pytest.mark.skipif(get_cudnn_version() < (9, 2, 1), reason="cuDNN 9.2.1+ is required.") @pytest.mark.skipif(not fp8_attn_available, reason=reason_for_no_fp8_attn) @pytest.mark.parametrize("dtype", param_types_fp8_vs_f16) @@ -2192,6 +2516,9 @@ def test_mha_fp8_vs_f16( scaling_mode, ): """Test MultiHeadAttention module in FP8""" + if not is_training and fp8_dpa_bwd: + pytest.skip("fp8_dpa_bwd=True not applicable for inference") + os.environ["NVTE_FP8_DPA_BWD"] = "1" if fp8_dpa_bwd else "0" config = _get_fp8_vs_f16_config(model, qkv_format) @@ -2219,14 +2546,23 @@ def test_mha_fp8_vs_f16( ) fp8_meta = {} fp8_meta["recipe"] = fp8_recipe + # Get the token counts for THD and use them for both support query and actual execution + num_tokens_q, num_tokens_kv = None, None + if qkv_format == "thd": + seqlens_q, seqlens_kv = _mha_fp8_vs_f16_seqlens(config, qkv_format) + num_tokens_q = int(seqlens_q.sum()) + num_tokens_kv = int(seqlens_kv.sum()) available_backends, _, _ = get_available_attention_backends( config, qkv_dtype=torch.float8_e4m3fn, + nominal_dtype=dtype, qkv_layout=qkv_format.replace("hd", "h3d"), fp8=True, fp8_meta=fp8_meta, is_training=is_training, deterministic=_deterministic, + num_tokens_q=num_tokens_q, + num_tokens_kv=num_tokens_kv, ) flash_attn_supported, fused_attn_supported_fp8, unfused_attn_supported = available_backends available_backends, _, fused_attn_backends = get_available_attention_backends( @@ -2235,6 +2571,8 @@ def test_mha_fp8_vs_f16( qkv_layout=qkv_format.replace("hd", "h3d"), is_training=is_training, deterministic=_deterministic, + num_tokens_q=num_tokens_q, + num_tokens_kv=num_tokens_kv, ) _, fused_attn_supported_f16, _ = available_backends if flash_attn_supported + fused_attn_supported_fp8 < 1: @@ -2349,47 +2687,13 @@ def get_dummy_cuda_rng_tracker() -> CudaRNGStatesTracker: attention_type="self", qkv_weight_interleaved=True, qkv_format=qkv_format, + window_size=config.window_size, + softmax_type=config.softmax_type, ).to(dtype=dtype, device="cuda") if not is_training: mha = mha.eval() - def random_seqlens(max_seqlen): - if qkv_format != "thd": - return torch.randint( - 1, max_seqlen, [config.batch_size], dtype=torch.int32, device="cuda" - ) - # Reserve seven positions so total-token alignment only increases the final length. - return torch.cat( - ( - torch.randint( - 1, - max_seqlen, - [config.batch_size - 1], - dtype=torch.int32, - device="cuda", - ), - torch.randint(1, max_seqlen - 6, [1], dtype=torch.int32, device="cuda"), - ) - ) - - if "padding" in config.attn_mask_type or qkv_format == "thd": - if config.attn_type == "self": - seqlens_q = random_seqlens(config.max_seqlen_q) - seqlens_kv = seqlens_q - if config.attn_type == "cross": - seqlens_q = random_seqlens(config.max_seqlen_q) - seqlens_kv = random_seqlens(config.max_seqlen_kv) - else: - seqlens_q = torch.full( - [config.batch_size], config.max_seqlen_q, dtype=torch.int32, device="cuda" - ) - seqlens_kv = torch.full( - [config.batch_size], config.max_seqlen_kv, dtype=torch.int32, device="cuda" - ) - if qkv_format == "thd": - # FP8 Linear flattens THD input to [t, h*d], so align total tokens for cuBLAS. - seqlens_q[-1] += -seqlens_q.sum() % 8 - seqlens_kv[-1] += -seqlens_kv.sum() % 8 + seqlens_q, seqlens_kv = _mha_fp8_vs_f16_seqlens(config, qkv_format) cu_seqlens_q = torch.zeros(config.batch_size + 1, dtype=torch.int32, device="cuda") cu_seqlens_kv = torch.zeros(config.batch_size + 1, dtype=torch.int32, device="cuda") cu_seqlens_q[1:] = torch.cumsum(seqlens_q, dim=0) @@ -2457,6 +2761,10 @@ def random_seqlens(max_seqlen): def test_dpa_fp8_vs_f16(dtype, model, qkv_layout, fp8_dpa_bwd, is_training, scaling_mode): """Test DotProductAttention module in FP8""" config = _get_fp8_vs_f16_config(model, qkv_layout) + if config.num_heads != config.num_gqa_groups and "3" in qkv_layout: + pytest.skip("qkv_layout not applicable for MQA/GQA") + if not is_training and fp8_dpa_bwd: + pytest.skip("fp8_dpa_bwd=True not applicable for inference") # TODO(cyang): think of another way to verify dropout results # test cuDNN FP8 dropout @@ -2493,14 +2801,25 @@ def test_dpa_fp8_vs_f16(dtype, model, qkv_layout, fp8_dpa_bwd, is_training, scal ) fp8_meta = {} fp8_meta["recipe"] = fp8_recipe + # Get the token counts for THD and use them for both support query and actual execution + qkv_format = "".join([i for i in qkv_layout.split("_")[0] if i.isalpha()]) + num_tokens_q, num_tokens_kv = None, None + if qkv_format == "thd": + reset_rng_states() + seqlens_q, seqlens_kv = _dpa_fp8_vs_f16_seqlens(config, qkv_format) + num_tokens_q = int(seqlens_q.sum()) + num_tokens_kv = int(seqlens_kv.sum()) available_backends, _, _ = get_available_attention_backends( config, qkv_dtype=torch.float8_e4m3fn, + nominal_dtype=dtype, qkv_layout=qkv_layout, fp8=True, fp8_meta=fp8_meta, is_training=is_training, deterministic=_deterministic, + num_tokens_q=num_tokens_q, + num_tokens_kv=num_tokens_kv, ) flash_attn_supported, fused_attn_supported_fp8, unfused_attn_supported = available_backends available_backends, _, _ = get_available_attention_backends( @@ -2509,14 +2828,14 @@ def test_dpa_fp8_vs_f16(dtype, model, qkv_layout, fp8_dpa_bwd, is_training, scal qkv_layout=qkv_layout, is_training=is_training, deterministic=_deterministic, + num_tokens_q=num_tokens_q, + num_tokens_kv=num_tokens_kv, ) _, fused_attn_supported_f16, _ = available_backends if flash_attn_supported + fused_attn_supported_fp8 < 1: pytest.skip("No FP8 attention backend available.") if not fused_attn_supported_f16: pytest.skip("No reference backend available.") - if config.num_heads != config.num_gqa_groups and "3" in qkv_layout: - pytest.skip("qkv_layout not applicable for MQA/GQA") if flash_attn_supported: os.environ["NVTE_FLASH_ATTN"] = "1" @@ -2666,26 +2985,7 @@ def get_dummy_cuda_rng_tracker() -> CudaRNGStatesTracker: if not is_training: dpa = dpa.eval() - if "padding" in config.attn_mask_type or qkv_format == "thd": - if config.attn_type == "self": - seqlens_q = torch.randint( - 1, config.max_seqlen_q, [config.batch_size], dtype=torch.int32, device="cuda" - ) - seqlens_kv = seqlens_q - if config.attn_type == "cross": - seqlens_q = torch.randint( - 1, config.max_seqlen_q, [config.batch_size], dtype=torch.int32, device="cuda" - ) - seqlens_kv = torch.randint( - 1, config.max_seqlen_kv, [config.batch_size], dtype=torch.int32, device="cuda" - ) - else: - seqlens_q = torch.full( - [config.batch_size], config.max_seqlen_q, dtype=torch.int32, device="cuda" - ) - seqlens_kv = torch.full( - [config.batch_size], config.max_seqlen_kv, dtype=torch.int32, device="cuda" - ) + seqlens_q, seqlens_kv = _dpa_fp8_vs_f16_seqlens(config, qkv_format) cu_seqlens_q = torch.zeros(config.batch_size + 1, dtype=torch.int32, device="cuda") cu_seqlens_kv = torch.zeros(config.batch_size + 1, dtype=torch.int32, device="cuda") cu_seqlens_q[1:] = torch.cumsum(seqlens_q, dim=0) @@ -2806,10 +3106,22 @@ def test_custom_mha_fp8_vs_f16(dtype, model): # Test backend availability is_training = True + fp8_meta = {} + fp8_recipe = recipe.DelayedScaling( + margin=0, + fp8_format=recipe.Format.HYBRID, + amax_history_len=1, + amax_compute_algo="most_recent", + fp8_dpa=True, + ) + fp8_meta["recipe"] = fp8_recipe available_backends, _, fused_attn_backends = get_available_attention_backends( config, qkv_dtype=torch.float8_e4m3fn, + nominal_dtype=dtype, qkv_layout="bs3hd", + fp8=True, + fp8_meta=fp8_meta, is_training=is_training, deterministic=_deterministic, ) @@ -2887,6 +3199,7 @@ def _run_custom_mha_fp8(dtype, config, backend): fp8_format=recipe.Format.HYBRID, amax_history_len=1, amax_compute_algo="most_recent", + fp8_dpa=True, ) mha = Custom_MHA_FP8(config).to(dtype=dtype, device="cuda") diff --git a/tests/pytorch/attention/test_attention_with_cp.py b/tests/pytorch/attention/test_attention_with_cp.py index af22c42016..367fe7271e 100644 --- a/tests/pytorch/attention/test_attention_with_cp.py +++ b/tests/pytorch/attention/test_attention_with_cp.py @@ -401,6 +401,8 @@ def test_cp_with_flash_attention(cp_pool, dtype, model, qkv_format, cp_comm_type config, qkv_dtype=dtypes[dtype], qkv_layout="_".join([qkv_format] * 3), + cp_size=num_gpus, + cp_size_a2a=2 if cp_comm_type == "a2a+p2p" else 1, ) flash_attn_supported, *_ = available_backends if not flash_attn_supported: @@ -668,7 +670,7 @@ def test_cp_with_fused_attention( # For 111s, dbias calculation is not supported as of cuDNN 9.18, hence, test fwd only for 111s. is_training = False if config.bias_shape == "111s" else True - available_backends, _, fused_attn_backends = get_available_attention_backends( + available_backends, *_ = get_available_attention_backends( config, qkv_dtype=dtypes[dtype] if dtype != "fp8" else torch.float8_e4m3fn, qkv_layout="_".join([qkv_format] * 3), @@ -676,23 +678,11 @@ def test_cp_with_fused_attention( fp8_meta=fp8_meta, is_training=is_training, deterministic=_deterministic, + cp_size=num_gpus, + cp_size_a2a=2 if cp_comm_type == "a2a+p2p" else 1, ) _, fused_attn_supported, _ = available_backends - if fused_attn_supported and config.attn_mask_type in ["causal", "padding_causal"]: - config_copy = copy.deepcopy(config) - config_copy.context_parallel = False - config_copy.attn_mask_type = config.attn_mask_type + "_bottom_right" - available_backends, _, fused_attn_backends = get_available_attention_backends( - config_copy, - qkv_dtype=dtypes[dtype] if dtype != "fp8" else torch.float8_e4m3fn, - qkv_layout="_".join([qkv_format] * 3), - fp8=fp8, - fp8_meta=fp8_meta, - is_training=is_training, - deterministic=_deterministic, - ) - _, fused_attn_supported, _ = available_backends if not fused_attn_supported: pytest.skip("No attention backend available.") diff --git a/tests/pytorch/attention/test_kv_cache.py b/tests/pytorch/attention/test_kv_cache.py index 2a857a10dc..46cd2eb58c 100644 --- a/tests/pytorch/attention/test_kv_cache.py +++ b/tests/pytorch/attention/test_kv_cache.py @@ -4,6 +4,7 @@ from collections import OrderedDict from typing import List +import copy import os import sys import pathlib @@ -468,13 +469,15 @@ def test_kv_cache(dtype, model, qkv_format, is_paged, backend, module, is_cuda_g for layer_number in range(1, num_layers + 1): inference_params.allocate_memory(layer_number) - # figure out supported backends + # figure out supported backends for both reference and inference models inference_params_qkv_format = "bshd" qkv_layout = qkv_format + "_" + "_".join([inference_params_qkv_format] * 2) if is_paged: qkv_layout = "paged_kv_" + qkv_layout - available_backends, _, fused_attn_backends = get_available_attention_backends( - config, + inference_config = copy.deepcopy(config) + inference_config.attn_mask_type = "padding_causal" + inference_backends, _, _ = get_available_attention_backends( + inference_config, qkv_dtype=dtype, qkv_layout=qkv_layout, pad_between_seqs=False, @@ -483,13 +486,24 @@ def test_kv_cache(dtype, model, qkv_format, is_paged, backend, module, is_cuda_g fp8_meta=fp8_meta, inference_params=inference_params, ) - flash_attn_supported, fused_attn_supported, unfused_attn_supported = available_backends - if backend == "FlashAttention" and not flash_attn_supported: - pytest.skip("FlashAttention backend is not supported") - if backend == "FusedAttention" and not fused_attn_supported: - pytest.skip("FusedAttention backend is not supported") - if backend == "UnfusedAttention" and not unfused_attn_supported: - pytest.skip("UnfusedAttention backend is not supported") + reference_config = copy.deepcopy(config) + reference_config.attn_mask_type = "causal" + reference_config.batch_size = config.total_requests + reference_config.max_seqlen_q = config.max_seqlen_kv + reference_backends, _, _ = get_available_attention_backends( + reference_config, + qkv_dtype=dtype, + qkv_layout="bshd_bshd_bshd", + pad_between_seqs=False, + is_training=False, + ) + backend_index = ("FlashAttention", "FusedAttention", "UnfusedAttention").index(backend) + for probe, supported in ( + ("inference", inference_backends), + ("reference", reference_backends), + ): + if not supported[backend_index]: + pytest.skip(f"{backend} backend is not supported for the {probe} config") os.environ["NVTE_FLASH_ATTN"] = str(int(backend == "FlashAttention")) os.environ["NVTE_FUSED_ATTN"] = str(int(backend == "FusedAttention")) os.environ["NVTE_UNFUSED_ATTN"] = str(int(backend == "UnfusedAttention")) diff --git a/tests/pytorch/test_torch_compile.py b/tests/pytorch/test_torch_compile.py index f8e09d5ce1..38722bb51b 100644 --- a/tests/pytorch/test_torch_compile.py +++ b/tests/pytorch/test_torch_compile.py @@ -1356,7 +1356,7 @@ def fn(x, params): monkeypatch.setattr( dpa_utils.tex, "get_fused_attn_backend", - lambda *args: dpa_utils.FusedAttnBackend["No_Backend"], + lambda *args: (tex.NVTE_Fused_Attn_Backend.NVTE_No_Backend, "disabled by test"), ) def fn_no_backend(x, params): diff --git a/tests/pytorch/utils.py b/tests/pytorch/utils.py index 21601d8cdd..1d50318a36 100644 --- a/tests/pytorch/utils.py +++ b/tests/pytorch/utils.py @@ -23,6 +23,7 @@ from transformer_engine.pytorch.attention.dot_product_attention import _attention_backends from transformer_engine.pytorch.attention.dot_product_attention.utils import ( get_attention_backend, + get_qkv_format, AttentionParams, AttentionLogging, check_set_window_size, @@ -312,6 +313,10 @@ def __init__( self.attn_type = "self" if (self.max_seqlen_q == self.max_seqlen_kv) else "cross" self.bias_shape = bias_shape self.window_size = check_set_window_size(self.attn_mask_type, window_size) + self.bottom_right_diagonal = self.attn_mask_type not in { + "causal", + "padding_causal", + } self.context_parallel = context_parallel self.cp_comm_type = cp_comm_type self.return_max_logit = return_max_logit @@ -336,6 +341,7 @@ def get_available_attention_backends( config: ModelConfig, qkv_dtype: torch.dtype, qkv_layout: str, + nominal_dtype: Optional[torch.dtype] = None, pad_between_seqs: bool = False, deterministic: bool = False, fp8: bool = False, @@ -344,9 +350,23 @@ def get_available_attention_backends( inference_params: Optional[InferenceParams] = None, score_mod: bool = False, score_mod_bprop: bool = False, + cp_size: int = 1, + cp_size_a2a: int = 1, + num_tokens_q: Optional[int] = None, + num_tokens_kv: Optional[int] = None, ) -> Tuple[List, List]: """Check for all available attention backends that support a model configuration""" + _, q_format, kv_format = get_qkv_format(qkv_layout, inference_params) + if num_tokens_q is None: + num_tokens_q = ( + max(config.batch_size * config.max_seqlen_q // cp_size, 1) if q_format == "thd" else 0 + ) + if num_tokens_kv is None: + num_tokens_kv = ( + max(config.batch_size * config.max_seqlen_kv // cp_size, 1) if kv_format == "thd" else 0 + ) + os.environ["NVTE_FLASH_ATTN"] = "1" os.environ["NVTE_FUSED_ATTN"] = "1" os.environ["NVTE_UNFUSED_ATTN"] = "1" @@ -358,9 +378,15 @@ def get_available_attention_backends( if config.bias_shape == "bhss": alibi_slopes_shape = [config.batch_size, config.num_heads] - core_attention_bias_shape = ( - config.bias_shape if config.attn_bias_type == "post_scale_bias" else None - ) + core_attention_bias_shape = None + if config.attn_bias_type == "post_scale_bias": + b_dim, h_dim, sq_dim, skv_dim = config.bias_shape + core_attention_bias_shape = ( + config.batch_size if b_dim == "b" else 1, + config.num_heads if h_dim == "h" else 1, + config.max_seqlen_q if sq_dim == "s" else 1, + config.max_seqlen_kv if skv_dim == "s" else 1, + ) core_attention_bias_requires_grad = False # d=256 is supported by cuDNN 9.0+ for inference but not training if ( @@ -369,7 +395,7 @@ def get_available_attention_backends( and config.head_dim_v <= 128 ): # TODO(KshitijLakhani): Remove this guard when cuDNN starts support dbias calculation for bias shape 111s - if core_attention_bias_shape != "111s": + if config.bias_shape != "111s": core_attention_bias_requires_grad = True fused_attn_backends = [] @@ -380,6 +406,7 @@ def get_available_attention_backends( def test(): attention_params = AttentionParams( qkv_dtype=qkv_dtype, + nominal_dtype=nominal_dtype, qkv_layout=qkv_layout, batch_size=config.batch_size, num_heads=config.num_heads, @@ -388,8 +415,11 @@ def test(): max_seqlen_kv=config.max_seqlen_kv, head_dim_qk=config.head_dim_qk, head_dim_v=config.head_dim_v, + num_tokens_q=num_tokens_q, + num_tokens_kv=num_tokens_kv, attn_mask_type=config.attn_mask_type, window_size=config.window_size, + bottom_right_diagonal=config.bottom_right_diagonal, alibi_slopes_shape=alibi_slopes_shape, core_attention_bias_type=config.attn_bias_type, core_attention_bias_shape=core_attention_bias_shape, @@ -398,6 +428,8 @@ def test(): attention_dropout=config.dropout_p, context_parallel=config.context_parallel, cp_comm_type=config.cp_comm_type, + cp_size=cp_size, + cp_size_a2a=cp_size_a2a, deterministic=deterministic, fp8=fp8, fp8_meta=fp8_meta, @@ -437,12 +469,10 @@ def test(): if AttentionLogging._is_logging_setup is False: AttentionLogging.setup_logging() - for i in backends: - os.environ["NVTE_FUSED_ATTN_BACKEND"] = str(i) - _attention_backends["backend_selection_requires_update"] = True - available_backends, flash_attention_backend, fused_attention_backend = test() - if fused_attention_backend == FusedAttnBackend[backends[i]]: - fused_attn_backends.append(fused_attention_backend) + _attention_backends["backend_selection_requires_update"] = True + available_backends, flash_attention_backend, fused_attention_backend = test() + if fused_attention_backend in (FusedAttnBackend[name] for name in backends.values()): + fused_attn_backends.append(fused_attention_backend) return available_backends, flash_attention_backend, fused_attn_backends diff --git a/transformer_engine/common/CMakeLists.txt b/transformer_engine/common/CMakeLists.txt index c9d4558f9c..9d4c8a9dc8 100644 --- a/transformer_engine/common/CMakeLists.txt +++ b/transformer_engine/common/CMakeLists.txt @@ -184,6 +184,7 @@ list(APPEND transformer_engine_cpp_sources cudnn_utils.cpp transformer_engine.cpp fused_attn/fused_attn.cpp + fused_attn/config_and_params.cpp gemm/config.cpp normalization/common.cpp normalization/rtc_dispatch.cpp diff --git a/transformer_engine/common/fused_attn/config_and_params.cpp b/transformer_engine/common/fused_attn/config_and_params.cpp new file mode 100644 index 0000000000..8c0029542e --- /dev/null +++ b/transformer_engine/common/fused_attn/config_and_params.cpp @@ -0,0 +1,1237 @@ +/************************************************************************* + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * + * See LICENSE for license information. + ************************************************************************/ + +#include "config_and_params.h" + +#include +#include + +#include +#include +#include +#include + +#include "../common.h" +#include "../util/cuda_runtime.h" + +namespace { + +void bool_to_uint8(bool in, void *out) { + *reinterpret_cast(out) = static_cast(in); +} + +void uint8_to_bool(const void *in, bool &out) { + out = static_cast(*reinterpret_cast(in)); +} + +} // namespace + +namespace transformer_engine { + +namespace fused_attn { + +// Forward declarations +size_t get_max_batch_size(size_t batch_size); +size_t get_max_tokens(size_t num_tokens); +DType get_ragged_offset_dtype(NVTE_QKV_Layout_Group layout_group, int64_t num_attn_heads, + int64_t num_gqa_groups, int64_t max_seqlen_q, int64_t max_seqlen_kv, + int64_t head_dim_qk, int64_t head_dim_v); + +void FusedAttnConfig::derive() { + if (is_derived) return; + + // Common attributes + qkv_format = nvte_get_qkv_format(qkv_layout); + q_format = nvte_get_q_format(qkv_layout); + kv_format = nvte_get_kv_format(qkv_layout); + const NVTE_QKV_Layout_Group layout_group = nvte_get_qkv_layout_group(qkv_layout); + is_paged_kv = (layout_group == NVTE_QKV_Layout_Group::NVTE_Paged_KV_HD_HD_HD); + is_ragged_q = (q_format == NVTE_QKV_Format::NVTE_THD); + is_ragged_kv = (kv_format == NVTE_QKV_Format::NVTE_THD); + is_padding = (attn_mask_type == NVTE_Mask_Type::NVTE_PADDING_MASK) || + (attn_mask_type == NVTE_Mask_Type::NVTE_PADDING_CAUSAL_MASK) || + (attn_mask_type == NVTE_Mask_Type::NVTE_PADDING_CAUSAL_BOTTOM_RIGHT_MASK); + is_causal = (attn_mask_type == NVTE_Mask_Type::NVTE_CAUSAL_MASK) || + (attn_mask_type == NVTE_Mask_Type::NVTE_PADDING_CAUSAL_MASK); + is_causal_bottom_right = + (attn_mask_type == NVTE_Mask_Type::NVTE_CAUSAL_BOTTOM_RIGHT_MASK) || + (attn_mask_type == NVTE_Mask_Type::NVTE_PADDING_CAUSAL_BOTTOM_RIGHT_MASK); + if (is_causal_bottom_right && !bottom_right_diagonal) { + bottom_right_diagonal = true; + } + if (is_causal && bottom_right_diagonal) { + bottom_right_diagonal = false; + } + is_bias = (bias_type == NVTE_Bias_Type::NVTE_POST_SCALE_BIAS); + is_alibi = (bias_type == NVTE_Bias_Type::NVTE_ALIBI); + is_softmax_offset = (softmax_type != NVTE_Softmax_Type::NVTE_VANILLA_SOFTMAX); + is_dropout = is_training && dropout != 0.0f; + + // Determine the FP8 recipe + is_o_in_fp8 = (o_dtype == kNVTEFloat8E4M3 || o_dtype == kNVTEFloat8E5M2); + is_dqkv_in_fp8 = (dqkv_dtype == kNVTEFloat8E4M3 || dqkv_dtype == kNVTEFloat8E5M2); + is_o_in_f16 = (o_dtype == kNVTEFloat16 || o_dtype == kNVTEBFloat16); + is_dqkv_in_f16 = (dqkv_dtype == kNVTEFloat16 || dqkv_dtype == kNVTEBFloat16); + is_tensor_scaling = (scaling_mode == NVTE_DELAYED_TENSOR_SCALING); + is_mxfp8 = (scaling_mode == NVTE_MXFP8_1D_SCALING); + is_delayed_scaling_fwd = is_tensor_scaling && is_o_in_fp8; + is_delayed_scaling_bwd = is_tensor_scaling && is_dqkv_in_fp8; + is_current_scaling_fwd = is_tensor_scaling && is_o_in_f16; + is_current_scaling_bwd = is_tensor_scaling && is_dqkv_in_f16; + is_mxfp8_fwd = is_mxfp8 && is_o_in_f16; + is_mxfp8_bwd = is_mxfp8 && is_dqkv_in_f16; + + // Use cu_seqlens vs actual_seqlens for THD or padding masks + const size_t cudnn_runtime_version = cudnnGetVersion(); + const bool is_fp8_dtype = (qkv_dtype == kNVTEFloat8E4M3 || qkv_dtype == kNVTEFloat8E5M2); + const size_t min_frontend_version = is_fp8_dtype ? 12600 : 12500; + const size_t min_cudnn_version = is_fp8_dtype ? 92500 : 92400; + uses_cu_seqlens_directly = + CUDNN_FRONTEND_VERSION >= min_frontend_version && + (CUDNN_VERSION >= min_cudnn_version && cudnn_runtime_version >= min_cudnn_version) && + !is_dropout; + + // Bucket the batch size and token counts for THD + bucketed_batch_size = + (is_ragged_q || is_ragged_kv) ? fused_attn::get_max_batch_size(batch_size) : 0; + bucketed_num_tokens_q = is_ragged_q ? fused_attn::get_max_tokens(num_tokens_q) : 0; + bucketed_num_tokens_kv = is_ragged_kv ? fused_attn::get_max_tokens(num_tokens_kv) : 0; + + // Use ragged (TH1) or dense (BHS1) graphs and stats + const int sm_arch = cuda::sm_arch(cuda::current_device()); + uses_ragged_graph = cudnn_runtime_version >= 90600 && sm_arch >= 90 && sm_arch != 120; + uses_ragged_stats = is_ragged_q && uses_ragged_graph; + const bool buckets_the_batch = (is_ragged_q || is_ragged_kv) && uses_ragged_graph; + graph_batch_size_fwd = + (buckets_the_batch && !uses_cu_seqlens_directly) ? bucketed_batch_size : batch_size; + graph_batch_size_bwd = buckets_the_batch ? bucketed_batch_size : batch_size; + graph_max_seqlen_q = (is_ragged_q && uses_ragged_graph) ? bucketed_num_tokens_q : max_seqlen_q; + graph_max_seqlen_kv = + (is_ragged_kv && uses_ragged_graph) ? bucketed_num_tokens_kv : max_seqlen_kv; + + // Set up ragged offset widths and multipliers + needs_64bit_ragged_offset = + (is_ragged_q || is_ragged_kv) && + fused_attn::get_ragged_offset_dtype( + layout_group, static_cast(num_attn_heads), static_cast(num_gqa_groups), + static_cast(max_seqlen_q), static_cast(max_seqlen_kv), + static_cast(head_dim_qk), static_cast(head_dim_v)) == DType::kInt64; + const DType wide_ragged_offsets = cudnn_runtime_version >= 90500 ? DType::kInt64 : DType::kInt32; + ragged_offset_type_fwd = uses_cu_seqlens_directly ? DType::kInt32 : wide_ragged_offsets; + ragged_offset_type_bwd = wide_ragged_offsets; + ragged_offset_mults = RaggedOffsetMultipliers( + layout_group, static_cast(num_attn_heads), static_cast(num_gqa_groups), + static_cast(head_dim_qk), static_cast(head_dim_v)); + + // Mark as derived + is_derived = true; +} + +FusedAttnConfig FusedAttnConfig::make_cache_key(Pass pass) const { + check_derived(); + FusedAttnConfig cache_cfg = *this; + cache_cfg.device_id = cuda::current_device(); + + // Normalize bottom_right_diagonal + const bool has_window = cache_cfg.window_size_left != -1 || cache_cfg.window_size_right != -1; + if (!cache_cfg.is_causal && !cache_cfg.is_causal_bottom_right && !has_window) { + cache_cfg.bottom_right_diagonal = false; + } else if (cache_cfg.is_causal_bottom_right && + cache_cfg.max_seqlen_q == cache_cfg.max_seqlen_kv && !cache_cfg.is_padding) { + cache_cfg.bottom_right_diagonal = false; + } + + // Normalize sequence lengths the graph is built at + cache_cfg.max_seqlen_q = cache_cfg.graph_max_seqlen_q; + cache_cfg.max_seqlen_kv = cache_cfg.graph_max_seqlen_kv; + + // Normalize the batch size the graph is built at, and drop the token counts. + cache_cfg.num_tokens_q = 0; + cache_cfg.num_tokens_kv = 0; + if ((cache_cfg.is_ragged_q || cache_cfg.is_ragged_kv) && cache_cfg.uses_ragged_graph) { + cache_cfg.batch_size = + pass == Pass::Fwd ? cache_cfg.graph_batch_size_fwd : cache_cfg.graph_batch_size_bwd; + } + + // attn_scale is a pass-by-value graph input and different scales can share the same cached graph + cache_cfg.attn_scale = 1.0f; + + // cuda_graph never reaches a graph builder. Its one use is the cuDNN <= 9.15 rejection in + // nvte_get_fused_attn_backend_v2(). + cache_cfg.cuda_graph = false; + + // Normalize the fields its graph actually consumes + if (pass == Pass::Fwd) { + cache_cfg.do_dtype = kNVTEBFloat16; + cache_cfg.dqkv_dtype = kNVTEBFloat16; + cache_cfg.do_format = NVTE_QKV_Format_NOT_SET; + cache_cfg.dqkv_layout = NVTE_QKV_Layout_NOT_SET; + cache_cfg.do_scale_inv_format = NVTE_QKV_Format_NOT_SET; + cache_cfg.deterministic = false; + } else { + cache_cfg.return_max_logit = false; + } + + return cache_cfg; +} + +std::string FusedAttnConfig::to_string() const { + char buf[1024]; + std::snprintf( + buf, sizeof(buf), + "train=%d det=%d cg=%d maxlogit=%d mask=%" PRId64 " bias=%" PRId64 " wl=%" PRId64 + " wr=%" PRId64 " brd=%d softmax=%" PRId64 " scale_mode=%" PRId64 + " dropout=%g attn_scale=%g qkv_dt=%" PRId64 " o_dt=%" PRId64 " do_dt=%" PRId64 + " dqkv_dt=%" PRId64 " qkv_lay=%" PRId64 " o_fmt=%" PRId64 " do_fmt=%" PRId64 + " dqkv_lay=%" PRId64 " qkv_sif=%" PRId64 " do_sif=%" PRId64 " b=%" PRId64 " h=%" PRId64 + " hg=%" PRId64 " dqk=%" PRId64 " dv=%" PRId64 " sq=%" PRId64 " skv=%" PRId64 " tq=%" PRId64 + " tkv=%" PRId64 " bb=%" PRId64 " btq=%" PRId64 " btkv=%" PRId64 " npk=%" PRId64 + " npv=%" PRId64 " psk=%" PRId64 " psv=%" PRId64 " mppk=%" PRId64 " mppv=%" PRId64 + " bias_b=%" PRId64 " bias_h=%" PRId64 " bias_sq=%" PRId64 " bias_skv=%" PRId64, + static_cast(is_training), static_cast(deterministic), static_cast(cuda_graph), + static_cast(return_max_logit), static_cast(attn_mask_type), + static_cast(bias_type), static_cast(window_size_left), + static_cast(window_size_right), static_cast(bottom_right_diagonal), + static_cast(softmax_type), static_cast(scaling_mode), + static_cast(dropout), static_cast(attn_scale), + static_cast(qkv_dtype), static_cast(o_dtype), + static_cast(do_dtype), static_cast(dqkv_dtype), + static_cast(qkv_layout), static_cast(o_format), + static_cast(do_format), static_cast(dqkv_layout), + static_cast(qkv_scale_inv_format), static_cast(do_scale_inv_format), + static_cast(batch_size), static_cast(num_attn_heads), + static_cast(num_gqa_groups), static_cast(head_dim_qk), + static_cast(head_dim_v), static_cast(max_seqlen_q), + static_cast(max_seqlen_kv), static_cast(num_tokens_q), + static_cast(num_tokens_kv), static_cast(bucketed_batch_size), + static_cast(bucketed_num_tokens_q), static_cast(bucketed_num_tokens_kv), + static_cast(num_pages_k), static_cast(num_pages_v), + static_cast(page_size_k), static_cast(page_size_v), + static_cast(max_pages_per_seq_k), static_cast(max_pages_per_seq_v), + static_cast(bias_batch_size), static_cast(bias_num_heads), + static_cast(bias_seqlen_q), static_cast(bias_seqlen_kv)); + return std::string(buf); +} + +FusedAttnConfig FusedAttnFwdParams::make_config() const { + const FusedAttnFwdParams ¶ms = *this; + FusedAttnConfig cfg{}; + // Forward execution: only the forward graph is run, so do not pay for a backward support + // check whose graph this call will never execute. + cfg.check_for_forward_support = true; + cfg.check_for_backward_support = false; + cfg.is_training = params.is_training; + cfg.deterministic = false; + cfg.cuda_graph = params.cuda_graph; + cfg.return_max_logit = params.return_max_logit; + cfg.attn_mask_type = params.attn_mask_type; + cfg.bias_type = params.bias_type; + cfg.window_size_left = params.window_size_left; + cfg.window_size_right = params.window_size_right; + cfg.bottom_right_diagonal = params.bottom_right_diagonal; + cfg.softmax_type = params.softmax_type; + cfg.dropout = params.dropout; + cfg.attn_scale = params.attn_scale; + cfg.qkv_layout = params.qkv_layout; + cfg.o_format = params.o_format; + cfg.qkv_scale_inv_format = params.qkv_scale_inv_format; + cfg.max_seqlen_q = params.max_seqlen_q; + cfg.max_seqlen_kv = params.max_seqlen_kv; + + const Tensor *input_cu_seqlens_q = convertNVTETensorCheck(params.cu_seqlens_q); + const Tensor *input_cu_seqlens_kv = convertNVTETensorCheck(params.cu_seqlens_kv); + const Tensor *input_page_table_k = convertNVTETensorCheck(params.page_table_k); + const Tensor *input_page_table_v = convertNVTETensorCheck(params.page_table_v); + const Tensor *input_Q = convertNVTETensorCheck(params.Q); + const Tensor *input_K = convertNVTETensorCheck(params.K); + const Tensor *input_V = convertNVTETensorCheck(params.V); + const Tensor *input_Bias = convertNVTETensorCheck(params.Bias); + const Tensor *output_O = convertNVTETensorCheck(params.O); + + const NVTE_QKV_Format q_format = nvte_get_q_format(params.qkv_layout); + const NVTE_QKV_Format kv_format = nvte_get_kv_format(params.qkv_layout); + auto *q_dims = input_Q->data.shape.data(); + auto *k_dims = input_K->data.shape.data(); + auto *v_dims = input_V->scaling_mode != NVTE_MXFP8_1D_SCALING + ? input_V->data.shape.data() + : input_V->columnwise_data.shape.data(); + AttentionShape q_shape(q_format, q_dims); + AttentionShape k_shape(kv_format, k_dims); + AttentionShape v_shape(kv_format, v_dims); + size_t b = q_shape.b(), h_q = q_shape.h(), d_qk = q_shape.d(), t_q = q_shape.t(); + size_t h_kv = k_shape.h(), t_kv = k_shape.t(), d_v = v_shape.d(); + if (q_format == NVTE_QKV_Format::NVTE_THD) { + b = input_cu_seqlens_q->data.shape[0] - 1; + } else if (kv_format == NVTE_QKV_Format::NVTE_THD) { + b = input_cu_seqlens_kv->data.shape[0] - 1; + } + + int64_t num_pages_k = 0, num_pages_v = 0, page_size_k = 0, page_size_v = 0; + int64_t max_pages_per_seq_k = 0, max_pages_per_seq_v = 0; + if (input_page_table_k->data.dptr != nullptr) { + max_pages_per_seq_k = input_page_table_k->data.shape[1]; + } + if (input_page_table_v->data.dptr != nullptr) { + max_pages_per_seq_v = input_page_table_v->data.shape[1]; + } + const NVTE_QKV_Layout_Group layout_group = nvte_get_qkv_layout_group(params.qkv_layout); + if (layout_group == NVTE_QKV_Layout_Group::NVTE_Paged_KV_HD_HD_HD) { + const NVTE_QKV_Format paged_kv_format = nvte_get_kv_format(params.qkv_layout); + if (paged_kv_format == NVTE_QKV_Format::NVTE_BSHD) { + num_pages_k = input_K->data.shape[0]; + page_size_k = input_K->data.shape[1]; + num_pages_v = input_V->data.shape[0]; + page_size_v = input_V->data.shape[1]; + } else if (paged_kv_format == NVTE_QKV_Format::NVTE_SBHD) { + num_pages_k = input_K->data.shape[1]; + page_size_k = input_K->data.shape[0]; + num_pages_v = input_V->data.shape[1]; + page_size_v = input_V->data.shape[0]; + } + } + + const NVTEDType Q_type = static_cast(input_Q->data.dtype); + const NVTEDType KV_type = static_cast(input_K->data.dtype); + NVTE_CHECK(Q_type == KV_type, "Q and KV must have the same data type."); + + cfg.scaling_mode = input_Q->scaling_mode; + cfg.qkv_dtype = Q_type; + cfg.o_dtype = static_cast(output_O->data.dtype); + cfg.batch_size = b; + cfg.num_attn_heads = h_q; + cfg.num_gqa_groups = h_kv; + cfg.head_dim_qk = d_qk; + cfg.head_dim_v = d_v; + cfg.num_pages_k = static_cast(num_pages_k); + cfg.num_pages_v = static_cast(num_pages_v); + cfg.page_size_k = static_cast(page_size_k); + cfg.page_size_v = static_cast(page_size_v); + cfg.max_pages_per_seq_k = static_cast(max_pages_per_seq_k); + cfg.max_pages_per_seq_v = static_cast(max_pages_per_seq_v); + cfg.num_tokens_q = t_q; + cfg.num_tokens_kv = t_kv; + + if ((params.bias_type != NVTE_NO_BIAS) && (params.bias_type != NVTE_ALIBI) && + input_Bias->data.shape.size() >= 4) { + cfg.bias_batch_size = input_Bias->data.shape[0]; + cfg.bias_num_heads = input_Bias->data.shape[1]; + cfg.bias_seqlen_q = input_Bias->data.shape[2]; + cfg.bias_seqlen_kv = input_Bias->data.shape[3]; + } + return cfg; +} + +FusedAttnConfig FusedAttnBwdParams::make_config() const { + const FusedAttnBwdParams ¶ms = *this; + FusedAttnConfig cfg{}; + // Backward execution: only the backward graph is run, so do not pay for a forward support + // check whose graph this call will never execute. + cfg.check_for_forward_support = false; + cfg.check_for_backward_support = true; + cfg.is_training = true; + cfg.deterministic = params.deterministic; + cfg.cuda_graph = params.cuda_graph; + cfg.return_max_logit = false; + cfg.attn_mask_type = params.attn_mask_type; + cfg.bias_type = params.bias_type; + cfg.window_size_left = params.window_size_left; + cfg.window_size_right = params.window_size_right; + cfg.bottom_right_diagonal = params.bottom_right_diagonal; + cfg.softmax_type = params.softmax_type; + cfg.dropout = params.dropout; + cfg.attn_scale = params.attn_scale; + cfg.qkv_layout = params.qkv_layout; + cfg.o_format = params.o_format; + cfg.do_format = params.do_format; + cfg.dqkv_layout = params.dqkv_layout; + cfg.qkv_scale_inv_format = params.qkv_scale_inv_format; + cfg.do_scale_inv_format = params.do_scale_inv_format; + cfg.max_seqlen_q = params.max_seqlen_q; + cfg.max_seqlen_kv = params.max_seqlen_kv; + + const Tensor *input_cu_seqlens_q = convertNVTETensorCheck(params.cu_seqlens_q); + const Tensor *input_cu_seqlens_kv = convertNVTETensorCheck(params.cu_seqlens_kv); + const Tensor *input_Q = convertNVTETensorCheck(params.Q); + const Tensor *input_K = convertNVTETensorCheck(params.K); + const Tensor *input_V = convertNVTETensorCheck(params.V); + const Tensor *input_O = convertNVTETensorCheck(params.O); + const Tensor *input_dO = convertNVTETensorCheck(params.dO); + const Tensor *output_dQ = convertNVTETensorCheck(params.dQ); + const Tensor *output_dBias = convertNVTETensorCheck(params.dBias); + + const NVTE_QKV_Format q_format = nvte_get_q_format(params.qkv_layout); + const NVTE_QKV_Format kv_format = nvte_get_kv_format(params.qkv_layout); + auto *q_dims = input_Q->data.shape.data(); + auto *k_dims = input_K->data.shape.data(); + auto *v_dims = input_V->data.shape.data(); + AttentionShape q_shape(q_format, q_dims); + AttentionShape k_shape(kv_format, k_dims); + AttentionShape v_shape(kv_format, v_dims); + size_t b = q_shape.b(), h_q = q_shape.h(), d_qk = q_shape.d(), t_q = q_shape.t(); + size_t h_kv = k_shape.h(), t_kv = k_shape.t(), d_v = v_shape.d(); + if (q_format == NVTE_QKV_Format::NVTE_THD) { + b = input_cu_seqlens_q->data.shape[0] - 1; + } else if (kv_format == NVTE_QKV_Format::NVTE_THD) { + b = input_cu_seqlens_kv->data.shape[0] - 1; + } + + const NVTEDType Q_type = static_cast(input_Q->data.dtype); + const NVTEDType KV_type = static_cast(input_K->data.dtype); + NVTE_CHECK(Q_type == KV_type, "Q and KV must have the same data type."); + + cfg.scaling_mode = input_Q->scaling_mode; + cfg.qkv_dtype = Q_type; + cfg.o_dtype = static_cast(input_O->data.dtype); + cfg.do_dtype = static_cast(input_dO->data.dtype); + cfg.dqkv_dtype = static_cast(output_dQ->data.dtype); + cfg.batch_size = b; + cfg.num_attn_heads = h_q; + cfg.num_gqa_groups = h_kv; + cfg.head_dim_qk = d_qk; + cfg.head_dim_v = d_v; + cfg.num_tokens_q = t_q; + cfg.num_tokens_kv = t_kv; + + if ((params.bias_type != NVTE_NO_BIAS) && (params.bias_type != NVTE_ALIBI) && + output_dBias->data.shape.size() >= 4) { + cfg.bias_batch_size = output_dBias->data.shape[0]; + cfg.bias_num_heads = output_dBias->data.shape[1]; + cfg.bias_seqlen_q = output_dBias->data.shape[2]; + cfg.bias_seqlen_kv = output_dBias->data.shape[3]; + } + return cfg; +} + +} // namespace fused_attn +} // namespace transformer_engine + +NVTEFusedAttnConfig nvte_create_fused_attn_config() { + return new transformer_engine::fused_attn::FusedAttnConfig{}; +} + +void nvte_destroy_fused_attn_config(NVTEFusedAttnConfig config) { + delete transformer_engine::fused_attn::get_fused_attn_config_mutable(config); +} + +void nvte_get_fused_attn_config_attribute(NVTEFusedAttnConfig config, + NVTEFusedAttnConfigAttribute attr, void *buf, + size_t size_in_bytes, size_t *size_written) { + using namespace transformer_engine; + using namespace transformer_engine::fused_attn; + + NVTE_CHECK(attr < kNVTEFusedAttnConfigNumAttributes, "Invalid NVTEFusedAttnConfigAttribute (got ", + static_cast(attr), ")"); + const auto &attr_size = FusedAttnConfig::attr_sizes[attr]; + if (size_written != nullptr) { + *size_written = attr_size; + } + if (buf == nullptr) { + return; + } + NVTE_CHECK(size_in_bytes >= attr_size, + "Buffer is too small for fused attention config attribute (attribute ", + static_cast(attr), " needs ", attr_size, " bytes, but buffer has ", size_in_bytes, + " bytes)"); + + const auto &cfg = *get_fused_attn_config(config); + switch (attr) { + case kNVTEFusedAttnConfigIsTraining: + bool_to_uint8(cfg.is_training, buf); + break; + case kNVTEFusedAttnConfigDeterministic: + bool_to_uint8(cfg.deterministic, buf); + break; + case kNVTEFusedAttnConfigCudaGraph: + bool_to_uint8(cfg.cuda_graph, buf); + break; + case kNVTEFusedAttnConfigReturnMaxLogit: + bool_to_uint8(cfg.return_max_logit, buf); + break; + case kNVTEFusedAttnConfigAttnMaskType: + std::memcpy(buf, &cfg.attn_mask_type, attr_size); + break; + case kNVTEFusedAttnConfigBiasType: + std::memcpy(buf, &cfg.bias_type, attr_size); + break; + case kNVTEFusedAttnConfigWindowSizeLeft: + std::memcpy(buf, &cfg.window_size_left, attr_size); + break; + case kNVTEFusedAttnConfigWindowSizeRight: + std::memcpy(buf, &cfg.window_size_right, attr_size); + break; + case kNVTEFusedAttnConfigBottomRightDiagonal: + bool_to_uint8(cfg.bottom_right_diagonal, buf); + break; + case kNVTEFusedAttnConfigSoftmaxType: + std::memcpy(buf, &cfg.softmax_type, attr_size); + break; + case kNVTEFusedAttnConfigScalingMode: + std::memcpy(buf, &cfg.scaling_mode, attr_size); + break; + case kNVTEFusedAttnConfigDropout: + std::memcpy(buf, &cfg.dropout, attr_size); + break; + case kNVTEFusedAttnConfigAttnScale: + std::memcpy(buf, &cfg.attn_scale, attr_size); + break; + case kNVTEFusedAttnConfigQKVDtype: + std::memcpy(buf, &cfg.qkv_dtype, attr_size); + break; + case kNVTEFusedAttnConfigODtype: + std::memcpy(buf, &cfg.o_dtype, attr_size); + break; + case kNVTEFusedAttnConfigDODtype: + std::memcpy(buf, &cfg.do_dtype, attr_size); + break; + case kNVTEFusedAttnConfigDQKVDtype: + std::memcpy(buf, &cfg.dqkv_dtype, attr_size); + break; + case kNVTEFusedAttnConfigQKVLayout: + std::memcpy(buf, &cfg.qkv_layout, attr_size); + break; + case kNVTEFusedAttnConfigOFormat: + std::memcpy(buf, &cfg.o_format, attr_size); + break; + case kNVTEFusedAttnConfigDOFormat: + std::memcpy(buf, &cfg.do_format, attr_size); + break; + case kNVTEFusedAttnConfigDQKVLayout: + std::memcpy(buf, &cfg.dqkv_layout, attr_size); + break; + case kNVTEFusedAttnConfigQKVScaleInvFormat: + std::memcpy(buf, &cfg.qkv_scale_inv_format, attr_size); + break; + case kNVTEFusedAttnConfigDOScaleInvFormat: + std::memcpy(buf, &cfg.do_scale_inv_format, attr_size); + break; + case kNVTEFusedAttnConfigBatchSize: + std::memcpy(buf, &cfg.batch_size, attr_size); + break; + case kNVTEFusedAttnConfigNumAttnHeads: + std::memcpy(buf, &cfg.num_attn_heads, attr_size); + break; + case kNVTEFusedAttnConfigNumGQAGroups: + std::memcpy(buf, &cfg.num_gqa_groups, attr_size); + break; + case kNVTEFusedAttnConfigHeadDimQK: + std::memcpy(buf, &cfg.head_dim_qk, attr_size); + break; + case kNVTEFusedAttnConfigHeadDimV: + std::memcpy(buf, &cfg.head_dim_v, attr_size); + break; + case kNVTEFusedAttnConfigMaxSeqlenQ: + std::memcpy(buf, &cfg.max_seqlen_q, attr_size); + break; + case kNVTEFusedAttnConfigMaxSeqlenKV: + std::memcpy(buf, &cfg.max_seqlen_kv, attr_size); + break; + case kNVTEFusedAttnConfigNumTokensQ: + std::memcpy(buf, &cfg.num_tokens_q, attr_size); + break; + case kNVTEFusedAttnConfigNumTokensKV: + std::memcpy(buf, &cfg.num_tokens_kv, attr_size); + break; + case kNVTEFusedAttnConfigNumPagesK: + std::memcpy(buf, &cfg.num_pages_k, attr_size); + break; + case kNVTEFusedAttnConfigNumPagesV: + std::memcpy(buf, &cfg.num_pages_v, attr_size); + break; + case kNVTEFusedAttnConfigPageSizeK: + std::memcpy(buf, &cfg.page_size_k, attr_size); + break; + case kNVTEFusedAttnConfigPageSizeV: + std::memcpy(buf, &cfg.page_size_v, attr_size); + break; + case kNVTEFusedAttnConfigMaxPagesPerSeqK: + std::memcpy(buf, &cfg.max_pages_per_seq_k, attr_size); + break; + case kNVTEFusedAttnConfigMaxPagesPerSeqV: + std::memcpy(buf, &cfg.max_pages_per_seq_v, attr_size); + break; + case kNVTEFusedAttnConfigBiasBatchSize: + std::memcpy(buf, &cfg.bias_batch_size, attr_size); + break; + case kNVTEFusedAttnConfigBiasNumHeads: + std::memcpy(buf, &cfg.bias_num_heads, attr_size); + break; + case kNVTEFusedAttnConfigBiasSeqlenQ: + std::memcpy(buf, &cfg.bias_seqlen_q, attr_size); + break; + case kNVTEFusedAttnConfigBiasSeqlenKV: + std::memcpy(buf, &cfg.bias_seqlen_kv, attr_size); + break; + default: + NVTE_ERROR("Unsupported NVTEFusedAttnConfigAttribute (got ", static_cast(attr), ")"); + } +} + +void nvte_set_fused_attn_config_attribute(NVTEFusedAttnConfig config, + NVTEFusedAttnConfigAttribute attr, const void *buf, + size_t size_in_bytes) { + using namespace transformer_engine; + using namespace transformer_engine::fused_attn; + + NVTE_CHECK(attr < kNVTEFusedAttnConfigNumAttributes, "Invalid NVTEFusedAttnConfigAttribute (got ", + static_cast(attr), ")"); + const auto &attr_size = FusedAttnConfig::attr_sizes[attr]; + NVTE_CHECK(size_in_bytes >= attr_size, + "Buffer is too small for fused attention config attribute (attribute ", + static_cast(attr), " needs ", attr_size, " bytes, but buffer has ", size_in_bytes, + " bytes)"); + NVTE_CHECK(buf != nullptr, "Invalid buffer (got NULL)"); + + auto &cfg = *get_fused_attn_config_mutable(config); + cfg.is_derived = false; + switch (attr) { + case kNVTEFusedAttnConfigIsTraining: + uint8_to_bool(buf, cfg.is_training); + break; + case kNVTEFusedAttnConfigDeterministic: + uint8_to_bool(buf, cfg.deterministic); + break; + case kNVTEFusedAttnConfigCudaGraph: + uint8_to_bool(buf, cfg.cuda_graph); + break; + case kNVTEFusedAttnConfigReturnMaxLogit: + uint8_to_bool(buf, cfg.return_max_logit); + break; + case kNVTEFusedAttnConfigAttnMaskType: + std::memcpy(&cfg.attn_mask_type, buf, attr_size); + break; + case kNVTEFusedAttnConfigBiasType: + std::memcpy(&cfg.bias_type, buf, attr_size); + break; + case kNVTEFusedAttnConfigWindowSizeLeft: + std::memcpy(&cfg.window_size_left, buf, attr_size); + break; + case kNVTEFusedAttnConfigWindowSizeRight: + std::memcpy(&cfg.window_size_right, buf, attr_size); + break; + case kNVTEFusedAttnConfigBottomRightDiagonal: + uint8_to_bool(buf, cfg.bottom_right_diagonal); + break; + case kNVTEFusedAttnConfigSoftmaxType: + std::memcpy(&cfg.softmax_type, buf, attr_size); + break; + case kNVTEFusedAttnConfigScalingMode: + std::memcpy(&cfg.scaling_mode, buf, attr_size); + break; + case kNVTEFusedAttnConfigDropout: + std::memcpy(&cfg.dropout, buf, attr_size); + break; + case kNVTEFusedAttnConfigAttnScale: + std::memcpy(&cfg.attn_scale, buf, attr_size); + break; + case kNVTEFusedAttnConfigQKVDtype: + std::memcpy(&cfg.qkv_dtype, buf, attr_size); + break; + case kNVTEFusedAttnConfigODtype: + std::memcpy(&cfg.o_dtype, buf, attr_size); + break; + case kNVTEFusedAttnConfigDODtype: + std::memcpy(&cfg.do_dtype, buf, attr_size); + break; + case kNVTEFusedAttnConfigDQKVDtype: + std::memcpy(&cfg.dqkv_dtype, buf, attr_size); + break; + case kNVTEFusedAttnConfigQKVLayout: + std::memcpy(&cfg.qkv_layout, buf, attr_size); + break; + case kNVTEFusedAttnConfigOFormat: + std::memcpy(&cfg.o_format, buf, attr_size); + break; + case kNVTEFusedAttnConfigDOFormat: + std::memcpy(&cfg.do_format, buf, attr_size); + break; + case kNVTEFusedAttnConfigDQKVLayout: + std::memcpy(&cfg.dqkv_layout, buf, attr_size); + break; + case kNVTEFusedAttnConfigQKVScaleInvFormat: + std::memcpy(&cfg.qkv_scale_inv_format, buf, attr_size); + break; + case kNVTEFusedAttnConfigDOScaleInvFormat: + std::memcpy(&cfg.do_scale_inv_format, buf, attr_size); + break; + case kNVTEFusedAttnConfigBatchSize: + std::memcpy(&cfg.batch_size, buf, attr_size); + break; + case kNVTEFusedAttnConfigNumAttnHeads: + std::memcpy(&cfg.num_attn_heads, buf, attr_size); + break; + case kNVTEFusedAttnConfigNumGQAGroups: + std::memcpy(&cfg.num_gqa_groups, buf, attr_size); + break; + case kNVTEFusedAttnConfigHeadDimQK: + std::memcpy(&cfg.head_dim_qk, buf, attr_size); + break; + case kNVTEFusedAttnConfigHeadDimV: + std::memcpy(&cfg.head_dim_v, buf, attr_size); + break; + case kNVTEFusedAttnConfigMaxSeqlenQ: + std::memcpy(&cfg.max_seqlen_q, buf, attr_size); + break; + case kNVTEFusedAttnConfigMaxSeqlenKV: + std::memcpy(&cfg.max_seqlen_kv, buf, attr_size); + break; + case kNVTEFusedAttnConfigNumTokensQ: + std::memcpy(&cfg.num_tokens_q, buf, attr_size); + break; + case kNVTEFusedAttnConfigNumTokensKV: + std::memcpy(&cfg.num_tokens_kv, buf, attr_size); + break; + case kNVTEFusedAttnConfigNumPagesK: + std::memcpy(&cfg.num_pages_k, buf, attr_size); + break; + case kNVTEFusedAttnConfigNumPagesV: + std::memcpy(&cfg.num_pages_v, buf, attr_size); + break; + case kNVTEFusedAttnConfigPageSizeK: + std::memcpy(&cfg.page_size_k, buf, attr_size); + break; + case kNVTEFusedAttnConfigPageSizeV: + std::memcpy(&cfg.page_size_v, buf, attr_size); + break; + case kNVTEFusedAttnConfigMaxPagesPerSeqK: + std::memcpy(&cfg.max_pages_per_seq_k, buf, attr_size); + break; + case kNVTEFusedAttnConfigMaxPagesPerSeqV: + std::memcpy(&cfg.max_pages_per_seq_v, buf, attr_size); + break; + case kNVTEFusedAttnConfigBiasBatchSize: + std::memcpy(&cfg.bias_batch_size, buf, attr_size); + break; + case kNVTEFusedAttnConfigBiasNumHeads: + std::memcpy(&cfg.bias_num_heads, buf, attr_size); + break; + case kNVTEFusedAttnConfigBiasSeqlenQ: + std::memcpy(&cfg.bias_seqlen_q, buf, attr_size); + break; + case kNVTEFusedAttnConfigBiasSeqlenKV: + std::memcpy(&cfg.bias_seqlen_kv, buf, attr_size); + break; + default: + NVTE_ERROR("Unsupported NVTEFusedAttnConfigAttribute (got ", static_cast(attr), ")"); + } +} + +NVTEFusedAttnFwdParams nvte_create_fused_attn_fwd_params() { + return new transformer_engine::fused_attn::FusedAttnFwdParams{}; +} + +void nvte_destroy_fused_attn_fwd_params(NVTEFusedAttnFwdParams params) { + delete transformer_engine::fused_attn::get_fused_attn_fwd_params_mutable(params); +} + +void nvte_get_fused_attn_fwd_params_attribute(NVTEFusedAttnFwdParams params, + NVTEFusedAttnFwdParamsAttribute attr, void *buf, + size_t size_in_bytes, size_t *size_written) { + using namespace transformer_engine; + using namespace transformer_engine::fused_attn; + NVTE_CHECK(attr < kNVTEFusedAttnFwdParamsNumAttributes, + "Invalid NVTEFusedAttnFwdParamsAttribute (got ", static_cast(attr), ")"); + const auto &attr_size = FusedAttnFwdParams::attr_sizes[attr]; + if (size_written != nullptr) { + *size_written = attr_size; + } + if (buf == nullptr) { + return; + } + NVTE_CHECK(size_in_bytes >= attr_size, "Buffer is too small for attribute (need ", attr_size, + ", got ", size_in_bytes, ")"); + const auto &p = *get_fused_attn_fwd_params(params); + switch (attr) { + case kNVTEFusedAttnFwdParamsQ: + std::memcpy(buf, &p.Q, attr_size); + break; + case kNVTEFusedAttnFwdParamsK: + std::memcpy(buf, &p.K, attr_size); + break; + case kNVTEFusedAttnFwdParamsV: + std::memcpy(buf, &p.V, attr_size); + break; + case kNVTEFusedAttnFwdParamsBias: + std::memcpy(buf, &p.Bias, attr_size); + break; + case kNVTEFusedAttnFwdParamsSoftmaxOffset: + std::memcpy(buf, &p.SoftmaxOffset, attr_size); + break; + case kNVTEFusedAttnFwdParamsS: + std::memcpy(buf, &p.S, attr_size); + break; + case kNVTEFusedAttnFwdParamsO: + std::memcpy(buf, &p.O, attr_size); + break; + case kNVTEFusedAttnFwdParamsAuxCtxTensors: + std::memcpy(buf, &p.Aux_CTX_Tensors, attr_size); + break; + case kNVTEFusedAttnFwdParamsCuSeqlensQ: + std::memcpy(buf, &p.cu_seqlens_q, attr_size); + break; + case kNVTEFusedAttnFwdParamsCuSeqlensKV: + std::memcpy(buf, &p.cu_seqlens_kv, attr_size); + break; + case kNVTEFusedAttnFwdParamsCuSeqlensQPadded: + std::memcpy(buf, &p.cu_seqlens_q_padded, attr_size); + break; + case kNVTEFusedAttnFwdParamsCuSeqlensKVPadded: + std::memcpy(buf, &p.cu_seqlens_kv_padded, attr_size); + break; + case kNVTEFusedAttnFwdParamsPageTableK: + std::memcpy(buf, &p.page_table_k, attr_size); + break; + case kNVTEFusedAttnFwdParamsPageTableV: + std::memcpy(buf, &p.page_table_v, attr_size); + break; + case kNVTEFusedAttnFwdParamsRngState: + std::memcpy(buf, &p.rng_state, attr_size); + break; + case kNVTEFusedAttnFwdParamsIsTraining: + bool_to_uint8(p.is_training, buf); + break; + case kNVTEFusedAttnFwdParamsCudaGraph: + bool_to_uint8(p.cuda_graph, buf); + break; + case kNVTEFusedAttnFwdParamsReturnMaxLogit: + bool_to_uint8(p.return_max_logit, buf); + break; + case kNVTEFusedAttnFwdParamsAttnMaskType: + std::memcpy(buf, &p.attn_mask_type, attr_size); + break; + case kNVTEFusedAttnFwdParamsBiasType: + std::memcpy(buf, &p.bias_type, attr_size); + break; + case kNVTEFusedAttnFwdParamsWindowSizeLeft: + std::memcpy(buf, &p.window_size_left, attr_size); + break; + case kNVTEFusedAttnFwdParamsWindowSizeRight: + std::memcpy(buf, &p.window_size_right, attr_size); + break; + case kNVTEFusedAttnFwdParamsBottomRightDiagonal: + bool_to_uint8(p.bottom_right_diagonal, buf); + break; + case kNVTEFusedAttnFwdParamsSoftmaxType: + std::memcpy(buf, &p.softmax_type, attr_size); + break; + case kNVTEFusedAttnFwdParamsDropout: + std::memcpy(buf, &p.dropout, attr_size); + break; + case kNVTEFusedAttnFwdParamsAttnScale: + std::memcpy(buf, &p.attn_scale, attr_size); + break; + case kNVTEFusedAttnFwdParamsQKVLayout: + std::memcpy(buf, &p.qkv_layout, attr_size); + break; + case kNVTEFusedAttnFwdParamsOFormat: + std::memcpy(buf, &p.o_format, attr_size); + break; + case kNVTEFusedAttnFwdParamsQKVScaleInvFormat: + std::memcpy(buf, &p.qkv_scale_inv_format, attr_size); + break; + case kNVTEFusedAttnFwdParamsMaxSeqlenQ: + std::memcpy(buf, &p.max_seqlen_q, attr_size); + break; + case kNVTEFusedAttnFwdParamsMaxSeqlenKV: + std::memcpy(buf, &p.max_seqlen_kv, attr_size); + break; + case kNVTEFusedAttnFwdParamsWorkspace: + std::memcpy(buf, &p.workspace, attr_size); + break; + case kNVTEFusedAttnFwdParamsStream: + std::memcpy(buf, &p.stream, attr_size); + break; + default: + NVTE_ERROR("Unsupported NVTEFusedAttnFwdParamsAttribute (got ", static_cast(attr), ")"); + } +} + +void nvte_set_fused_attn_fwd_params_attribute(NVTEFusedAttnFwdParams params, + NVTEFusedAttnFwdParamsAttribute attr, const void *buf, + size_t size_in_bytes) { + using namespace transformer_engine; + using namespace transformer_engine::fused_attn; + NVTE_CHECK(attr < kNVTEFusedAttnFwdParamsNumAttributes, + "Invalid NVTEFusedAttnFwdParamsAttribute (got ", static_cast(attr), ")"); + const auto &attr_size = FusedAttnFwdParams::attr_sizes[attr]; + NVTE_CHECK(buf != nullptr, "Input buffer must not be NULL."); + NVTE_CHECK(size_in_bytes >= attr_size, "Buffer is too small for attribute (need ", attr_size, + ", got ", size_in_bytes, ")"); + auto &p = *get_fused_attn_fwd_params_mutable(params); + switch (attr) { + case kNVTEFusedAttnFwdParamsQ: + std::memcpy(&p.Q, buf, attr_size); + break; + case kNVTEFusedAttnFwdParamsK: + std::memcpy(&p.K, buf, attr_size); + break; + case kNVTEFusedAttnFwdParamsV: + std::memcpy(&p.V, buf, attr_size); + break; + case kNVTEFusedAttnFwdParamsBias: + std::memcpy(&p.Bias, buf, attr_size); + break; + case kNVTEFusedAttnFwdParamsSoftmaxOffset: + std::memcpy(&p.SoftmaxOffset, buf, attr_size); + break; + case kNVTEFusedAttnFwdParamsS: + std::memcpy(&p.S, buf, attr_size); + break; + case kNVTEFusedAttnFwdParamsO: + std::memcpy(&p.O, buf, attr_size); + break; + case kNVTEFusedAttnFwdParamsAuxCtxTensors: + std::memcpy(&p.Aux_CTX_Tensors, buf, attr_size); + break; + case kNVTEFusedAttnFwdParamsCuSeqlensQ: + std::memcpy(&p.cu_seqlens_q, buf, attr_size); + break; + case kNVTEFusedAttnFwdParamsCuSeqlensKV: + std::memcpy(&p.cu_seqlens_kv, buf, attr_size); + break; + case kNVTEFusedAttnFwdParamsCuSeqlensQPadded: + std::memcpy(&p.cu_seqlens_q_padded, buf, attr_size); + break; + case kNVTEFusedAttnFwdParamsCuSeqlensKVPadded: + std::memcpy(&p.cu_seqlens_kv_padded, buf, attr_size); + break; + case kNVTEFusedAttnFwdParamsPageTableK: + std::memcpy(&p.page_table_k, buf, attr_size); + break; + case kNVTEFusedAttnFwdParamsPageTableV: + std::memcpy(&p.page_table_v, buf, attr_size); + break; + case kNVTEFusedAttnFwdParamsRngState: + std::memcpy(&p.rng_state, buf, attr_size); + break; + case kNVTEFusedAttnFwdParamsIsTraining: + uint8_to_bool(buf, p.is_training); + break; + case kNVTEFusedAttnFwdParamsCudaGraph: + uint8_to_bool(buf, p.cuda_graph); + break; + case kNVTEFusedAttnFwdParamsReturnMaxLogit: + uint8_to_bool(buf, p.return_max_logit); + break; + case kNVTEFusedAttnFwdParamsAttnMaskType: + std::memcpy(&p.attn_mask_type, buf, attr_size); + break; + case kNVTEFusedAttnFwdParamsBiasType: + std::memcpy(&p.bias_type, buf, attr_size); + break; + case kNVTEFusedAttnFwdParamsWindowSizeLeft: + std::memcpy(&p.window_size_left, buf, attr_size); + break; + case kNVTEFusedAttnFwdParamsWindowSizeRight: + std::memcpy(&p.window_size_right, buf, attr_size); + break; + case kNVTEFusedAttnFwdParamsBottomRightDiagonal: + uint8_to_bool(buf, p.bottom_right_diagonal); + break; + case kNVTEFusedAttnFwdParamsSoftmaxType: + std::memcpy(&p.softmax_type, buf, attr_size); + break; + case kNVTEFusedAttnFwdParamsDropout: + std::memcpy(&p.dropout, buf, attr_size); + break; + case kNVTEFusedAttnFwdParamsAttnScale: + std::memcpy(&p.attn_scale, buf, attr_size); + break; + case kNVTEFusedAttnFwdParamsQKVLayout: + std::memcpy(&p.qkv_layout, buf, attr_size); + break; + case kNVTEFusedAttnFwdParamsOFormat: + std::memcpy(&p.o_format, buf, attr_size); + break; + case kNVTEFusedAttnFwdParamsQKVScaleInvFormat: + std::memcpy(&p.qkv_scale_inv_format, buf, attr_size); + break; + case kNVTEFusedAttnFwdParamsMaxSeqlenQ: + std::memcpy(&p.max_seqlen_q, buf, attr_size); + break; + case kNVTEFusedAttnFwdParamsMaxSeqlenKV: + std::memcpy(&p.max_seqlen_kv, buf, attr_size); + break; + case kNVTEFusedAttnFwdParamsWorkspace: + std::memcpy(&p.workspace, buf, attr_size); + break; + case kNVTEFusedAttnFwdParamsStream: + std::memcpy(&p.stream, buf, attr_size); + break; + default: + NVTE_ERROR("Unsupported NVTEFusedAttnFwdParamsAttribute (got ", static_cast(attr), ")"); + } +} + +NVTEFusedAttnBwdParams nvte_create_fused_attn_bwd_params() { + return new transformer_engine::fused_attn::FusedAttnBwdParams{}; +} + +void nvte_destroy_fused_attn_bwd_params(NVTEFusedAttnBwdParams params) { + delete transformer_engine::fused_attn::get_fused_attn_bwd_params_mutable(params); +} + +void nvte_get_fused_attn_bwd_params_attribute(NVTEFusedAttnBwdParams params, + NVTEFusedAttnBwdParamsAttribute attr, void *buf, + size_t size_in_bytes, size_t *size_written) { + using namespace transformer_engine; + using namespace transformer_engine::fused_attn; + NVTE_CHECK(attr < kNVTEFusedAttnBwdParamsNumAttributes, + "Invalid NVTEFusedAttnBwdParamsAttribute (got ", static_cast(attr), ")"); + const auto &attr_size = FusedAttnBwdParams::attr_sizes[attr]; + if (size_written != nullptr) { + *size_written = attr_size; + } + if (buf == nullptr) { + return; + } + NVTE_CHECK(size_in_bytes >= attr_size, "Buffer is too small for attribute (need ", attr_size, + ", got ", size_in_bytes, ")"); + const auto &p = *get_fused_attn_bwd_params(params); + switch (attr) { + case kNVTEFusedAttnBwdParamsQ: + std::memcpy(buf, &p.Q, attr_size); + break; + case kNVTEFusedAttnBwdParamsK: + std::memcpy(buf, &p.K, attr_size); + break; + case kNVTEFusedAttnBwdParamsV: + std::memcpy(buf, &p.V, attr_size); + break; + case kNVTEFusedAttnBwdParamsO: + std::memcpy(buf, &p.O, attr_size); + break; + case kNVTEFusedAttnBwdParamsDO: + std::memcpy(buf, &p.dO, attr_size); + break; + case kNVTEFusedAttnBwdParamsS: + std::memcpy(buf, &p.S, attr_size); + break; + case kNVTEFusedAttnBwdParamsDP: + std::memcpy(buf, &p.dP, attr_size); + break; + case kNVTEFusedAttnBwdParamsAuxCtxTensors: + std::memcpy(buf, &p.Aux_CTX_Tensors, attr_size); + break; + case kNVTEFusedAttnBwdParamsDQ: + std::memcpy(buf, &p.dQ, attr_size); + break; + case kNVTEFusedAttnBwdParamsDK: + std::memcpy(buf, &p.dK, attr_size); + break; + case kNVTEFusedAttnBwdParamsDV: + std::memcpy(buf, &p.dV, attr_size); + break; + case kNVTEFusedAttnBwdParamsDBias: + std::memcpy(buf, &p.dBias, attr_size); + break; + case kNVTEFusedAttnBwdParamsDSoftmaxOffset: + std::memcpy(buf, &p.dSoftmaxOffset, attr_size); + break; + case kNVTEFusedAttnBwdParamsCuSeqlensQ: + std::memcpy(buf, &p.cu_seqlens_q, attr_size); + break; + case kNVTEFusedAttnBwdParamsCuSeqlensKV: + std::memcpy(buf, &p.cu_seqlens_kv, attr_size); + break; + case kNVTEFusedAttnBwdParamsCuSeqlensQPadded: + std::memcpy(buf, &p.cu_seqlens_q_padded, attr_size); + break; + case kNVTEFusedAttnBwdParamsCuSeqlensKVPadded: + std::memcpy(buf, &p.cu_seqlens_kv_padded, attr_size); + break; + case kNVTEFusedAttnBwdParamsDeterministic: + bool_to_uint8(p.deterministic, buf); + break; + case kNVTEFusedAttnBwdParamsCudaGraph: + bool_to_uint8(p.cuda_graph, buf); + break; + case kNVTEFusedAttnBwdParamsAttnMaskType: + std::memcpy(buf, &p.attn_mask_type, attr_size); + break; + case kNVTEFusedAttnBwdParamsBiasType: + std::memcpy(buf, &p.bias_type, attr_size); + break; + case kNVTEFusedAttnBwdParamsWindowSizeLeft: + std::memcpy(buf, &p.window_size_left, attr_size); + break; + case kNVTEFusedAttnBwdParamsWindowSizeRight: + std::memcpy(buf, &p.window_size_right, attr_size); + break; + case kNVTEFusedAttnBwdParamsBottomRightDiagonal: + bool_to_uint8(p.bottom_right_diagonal, buf); + break; + case kNVTEFusedAttnBwdParamsSoftmaxType: + std::memcpy(buf, &p.softmax_type, attr_size); + break; + case kNVTEFusedAttnBwdParamsDropout: + std::memcpy(buf, &p.dropout, attr_size); + break; + case kNVTEFusedAttnBwdParamsAttnScale: + std::memcpy(buf, &p.attn_scale, attr_size); + break; + case kNVTEFusedAttnBwdParamsQKVLayout: + std::memcpy(buf, &p.qkv_layout, attr_size); + break; + case kNVTEFusedAttnBwdParamsOFormat: + std::memcpy(buf, &p.o_format, attr_size); + break; + case kNVTEFusedAttnBwdParamsDOFormat: + std::memcpy(buf, &p.do_format, attr_size); + break; + case kNVTEFusedAttnBwdParamsDQKVLayout: + std::memcpy(buf, &p.dqkv_layout, attr_size); + break; + case kNVTEFusedAttnBwdParamsQKVScaleInvFormat: + std::memcpy(buf, &p.qkv_scale_inv_format, attr_size); + break; + case kNVTEFusedAttnBwdParamsDOScaleInvFormat: + std::memcpy(buf, &p.do_scale_inv_format, attr_size); + break; + case kNVTEFusedAttnBwdParamsMaxSeqlenQ: + std::memcpy(buf, &p.max_seqlen_q, attr_size); + break; + case kNVTEFusedAttnBwdParamsMaxSeqlenKV: + std::memcpy(buf, &p.max_seqlen_kv, attr_size); + break; + case kNVTEFusedAttnBwdParamsWorkspace: + std::memcpy(buf, &p.workspace, attr_size); + break; + case kNVTEFusedAttnBwdParamsStream: + std::memcpy(buf, &p.stream, attr_size); + break; + default: + NVTE_ERROR("Unsupported NVTEFusedAttnBwdParamsAttribute (got ", static_cast(attr), ")"); + } +} + +void nvte_set_fused_attn_bwd_params_attribute(NVTEFusedAttnBwdParams params, + NVTEFusedAttnBwdParamsAttribute attr, const void *buf, + size_t size_in_bytes) { + using namespace transformer_engine; + using namespace transformer_engine::fused_attn; + NVTE_CHECK(attr < kNVTEFusedAttnBwdParamsNumAttributes, + "Invalid NVTEFusedAttnBwdParamsAttribute (got ", static_cast(attr), ")"); + const auto &attr_size = FusedAttnBwdParams::attr_sizes[attr]; + NVTE_CHECK(buf != nullptr, "Input buffer must not be NULL."); + NVTE_CHECK(size_in_bytes >= attr_size, "Buffer is too small for attribute (need ", attr_size, + ", got ", size_in_bytes, ")"); + auto &p = *get_fused_attn_bwd_params_mutable(params); + switch (attr) { + case kNVTEFusedAttnBwdParamsQ: + std::memcpy(&p.Q, buf, attr_size); + break; + case kNVTEFusedAttnBwdParamsK: + std::memcpy(&p.K, buf, attr_size); + break; + case kNVTEFusedAttnBwdParamsV: + std::memcpy(&p.V, buf, attr_size); + break; + case kNVTEFusedAttnBwdParamsO: + std::memcpy(&p.O, buf, attr_size); + break; + case kNVTEFusedAttnBwdParamsDO: + std::memcpy(&p.dO, buf, attr_size); + break; + case kNVTEFusedAttnBwdParamsS: + std::memcpy(&p.S, buf, attr_size); + break; + case kNVTEFusedAttnBwdParamsDP: + std::memcpy(&p.dP, buf, attr_size); + break; + case kNVTEFusedAttnBwdParamsAuxCtxTensors: + std::memcpy(&p.Aux_CTX_Tensors, buf, attr_size); + break; + case kNVTEFusedAttnBwdParamsDQ: + std::memcpy(&p.dQ, buf, attr_size); + break; + case kNVTEFusedAttnBwdParamsDK: + std::memcpy(&p.dK, buf, attr_size); + break; + case kNVTEFusedAttnBwdParamsDV: + std::memcpy(&p.dV, buf, attr_size); + break; + case kNVTEFusedAttnBwdParamsDBias: + std::memcpy(&p.dBias, buf, attr_size); + break; + case kNVTEFusedAttnBwdParamsDSoftmaxOffset: + std::memcpy(&p.dSoftmaxOffset, buf, attr_size); + break; + case kNVTEFusedAttnBwdParamsCuSeqlensQ: + std::memcpy(&p.cu_seqlens_q, buf, attr_size); + break; + case kNVTEFusedAttnBwdParamsCuSeqlensKV: + std::memcpy(&p.cu_seqlens_kv, buf, attr_size); + break; + case kNVTEFusedAttnBwdParamsCuSeqlensQPadded: + std::memcpy(&p.cu_seqlens_q_padded, buf, attr_size); + break; + case kNVTEFusedAttnBwdParamsCuSeqlensKVPadded: + std::memcpy(&p.cu_seqlens_kv_padded, buf, attr_size); + break; + case kNVTEFusedAttnBwdParamsDeterministic: + uint8_to_bool(buf, p.deterministic); + break; + case kNVTEFusedAttnBwdParamsCudaGraph: + uint8_to_bool(buf, p.cuda_graph); + break; + case kNVTEFusedAttnBwdParamsAttnMaskType: + std::memcpy(&p.attn_mask_type, buf, attr_size); + break; + case kNVTEFusedAttnBwdParamsBiasType: + std::memcpy(&p.bias_type, buf, attr_size); + break; + case kNVTEFusedAttnBwdParamsWindowSizeLeft: + std::memcpy(&p.window_size_left, buf, attr_size); + break; + case kNVTEFusedAttnBwdParamsWindowSizeRight: + std::memcpy(&p.window_size_right, buf, attr_size); + break; + case kNVTEFusedAttnBwdParamsBottomRightDiagonal: + uint8_to_bool(buf, p.bottom_right_diagonal); + break; + case kNVTEFusedAttnBwdParamsSoftmaxType: + std::memcpy(&p.softmax_type, buf, attr_size); + break; + case kNVTEFusedAttnBwdParamsDropout: + std::memcpy(&p.dropout, buf, attr_size); + break; + case kNVTEFusedAttnBwdParamsAttnScale: + std::memcpy(&p.attn_scale, buf, attr_size); + break; + case kNVTEFusedAttnBwdParamsQKVLayout: + std::memcpy(&p.qkv_layout, buf, attr_size); + break; + case kNVTEFusedAttnBwdParamsOFormat: + std::memcpy(&p.o_format, buf, attr_size); + break; + case kNVTEFusedAttnBwdParamsDOFormat: + std::memcpy(&p.do_format, buf, attr_size); + break; + case kNVTEFusedAttnBwdParamsDQKVLayout: + std::memcpy(&p.dqkv_layout, buf, attr_size); + break; + case kNVTEFusedAttnBwdParamsQKVScaleInvFormat: + std::memcpy(&p.qkv_scale_inv_format, buf, attr_size); + break; + case kNVTEFusedAttnBwdParamsDOScaleInvFormat: + std::memcpy(&p.do_scale_inv_format, buf, attr_size); + break; + case kNVTEFusedAttnBwdParamsMaxSeqlenQ: + std::memcpy(&p.max_seqlen_q, buf, attr_size); + break; + case kNVTEFusedAttnBwdParamsMaxSeqlenKV: + std::memcpy(&p.max_seqlen_kv, buf, attr_size); + break; + case kNVTEFusedAttnBwdParamsWorkspace: + std::memcpy(&p.workspace, buf, attr_size); + break; + case kNVTEFusedAttnBwdParamsStream: + std::memcpy(&p.stream, buf, attr_size); + break; + default: + NVTE_ERROR("Unsupported NVTEFusedAttnBwdParamsAttribute (got ", static_cast(attr), ")"); + } +} diff --git a/transformer_engine/common/fused_attn/config_and_params.h b/transformer_engine/common/fused_attn/config_and_params.h new file mode 100644 index 0000000000..bf7e904af9 --- /dev/null +++ b/transformer_engine/common/fused_attn/config_and_params.h @@ -0,0 +1,472 @@ +/************************************************************************* + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * + * See LICENSE for license information. + ************************************************************************/ + +/*! \file config_and_params.h + * \brief Internal objects for fused-attention config and parameter handles. + */ + +#ifndef TRANSFORMER_ENGINE_COMMON_FUSED_ATTN_CONFIG_AND_PARAMS_H_ +#define TRANSFORMER_ENGINE_COMMON_FUSED_ATTN_CONFIG_AND_PARAMS_H_ + +#include +#include + +#include "common/common.h" +#include "transformer_engine/fused_attn.h" +#include "utils.h" + +namespace transformer_engine { +namespace fused_attn { + +enum class Backend { F16, FP8 }; +enum class Pass { Fwd, Bwd }; + +inline constexpr const char *backend_name(Backend b) { return b == Backend::F16 ? "f16" : "fp8"; } +inline constexpr const char *pass_name(Pass p) { return p == Pass::Fwd ? "fwd" : "bwd"; } + +struct FusedAttnConfig { + // Basic attention settings + bool is_training = true; + bool deterministic = false; + bool cuda_graph = false; + bool return_max_logit = false; + NVTE_Mask_Type attn_mask_type = NVTE_NO_MASK; + NVTE_Bias_Type bias_type = NVTE_NO_BIAS; + int64_t window_size_left = -1; + int64_t window_size_right = -1; + bool bottom_right_diagonal = true; + NVTE_Softmax_Type softmax_type = NVTE_VANILLA_SOFTMAX; + NVTEScalingMode scaling_mode = NVTE_DELAYED_TENSOR_SCALING; + float dropout = 0.0f; + float attn_scale = 1.0f; + + // Tensor types + NVTEDType qkv_dtype = kNVTEBFloat16; + NVTEDType o_dtype = kNVTEBFloat16; + NVTEDType do_dtype = kNVTEBFloat16; + NVTEDType dqkv_dtype = kNVTEBFloat16; + + // Tensor layouts + NVTE_QKV_Layout qkv_layout = NVTE_QKV_Layout_NOT_SET; + NVTE_QKV_Format o_format = NVTE_QKV_Format_NOT_SET; + NVTE_QKV_Format do_format = NVTE_QKV_Format_NOT_SET; + NVTE_QKV_Layout dqkv_layout = NVTE_QKV_Layout_NOT_SET; + NVTE_QKV_Format qkv_scale_inv_format = NVTE_QKV_Format_NOT_SET; + NVTE_QKV_Format do_scale_inv_format = NVTE_QKV_Format_NOT_SET; + + // Tensor dimensions + size_t batch_size = 0; + size_t num_attn_heads = 0; + size_t num_gqa_groups = 0; + size_t head_dim_qk = 0; + size_t head_dim_v = 0; + size_t max_seqlen_q = 0; + size_t max_seqlen_kv = 0; + size_t num_tokens_q = 0; + size_t num_tokens_kv = 0; + + // Paged KV dimensions + size_t num_pages_k = 0; + size_t num_pages_v = 0; + size_t page_size_k = 0; + size_t page_size_v = 0; + size_t max_pages_per_seq_k = 0; + size_t max_pages_per_seq_v = 0; + + // Bias dimensions + size_t bias_batch_size = 0; + size_t bias_num_heads = 0; + size_t bias_seqlen_q = 0; + size_t bias_seqlen_kv = 0; + + // ============================================================================ + // Internal fields: not part of attr_sizes[] or NVTEFusedAttnConfigAttribute, and + // unreachable from nvte_set_fused_attn_config_attribute, hence "internal". + // + // - Keyed: do participate in operator<, and can distinguish graphs in the cache, + // e.g. device_id; + // - Unkeyed: do not participate in operator<, and can not distinguish graphs + // in the cache; used for convenience purposes only; run derive() to fill them + // ============================================================================ + + // Keyed: + // - distinguish graphs on different GPUs in a single-process run + int device_id = -1; + + // Unkeyed: + // - support query directions + bool check_for_forward_support = true; + bool check_for_backward_support = true; + // - whether derived fields have been filled + bool is_derived = false; + // - common attributes + NVTE_QKV_Format qkv_format = NVTE_QKV_Format_NOT_SET; + NVTE_QKV_Format q_format = NVTE_QKV_Format_NOT_SET; + NVTE_QKV_Format kv_format = NVTE_QKV_Format_NOT_SET; + bool is_ragged_q = false; + bool is_ragged_kv = false; + bool is_paged_kv = false; + bool is_padding = false; + bool is_causal = false; + bool is_causal_bottom_right = false; + bool is_bias = false; + bool is_alibi = false; + bool is_softmax_offset = false; + bool is_dropout = false; + // - FP8 recipes + bool is_o_in_fp8 = false; + bool is_dqkv_in_fp8 = false; + bool is_o_in_f16 = false; + bool is_dqkv_in_f16 = false; + bool is_tensor_scaling = false; + bool is_mxfp8 = false; + bool is_delayed_scaling_fwd = false; + bool is_delayed_scaling_bwd = false; + bool is_current_scaling_fwd = false; + bool is_current_scaling_bwd = false; + bool is_mxfp8_fwd = false; + bool is_mxfp8_bwd = false; + // - cu_seqlens vs actual_seqlens for THD or padding masks + bool uses_cu_seqlens_directly = false; + // - bucket the batch size and token counts for THD + size_t bucketed_batch_size = 0; + size_t bucketed_num_tokens_q = 0; + size_t bucketed_num_tokens_kv = 0; + // - packed (TH1) vs dense (BHS1) graphs and stats + bool uses_ragged_graph = false; + bool uses_ragged_stats = false; + size_t graph_batch_size_fwd = 0; + size_t graph_batch_size_bwd = 0; + size_t graph_max_seqlen_q = 0; + size_t graph_max_seqlen_kv = 0; + // - ragged offset widths and multipliers + bool needs_64bit_ragged_offset = false; + DType ragged_offset_type_fwd = DType::kInt32; + DType ragged_offset_type_bwd = DType::kInt32; + RaggedOffsetMultipliers ragged_offset_mults; + + static constexpr size_t attr_sizes[] = { + // Basic attention settings + sizeof(uint8_t), // is_training + sizeof(uint8_t), // deterministic + sizeof(uint8_t), // cuda_graph + sizeof(uint8_t), // return_max_logit + sizeof(NVTE_Mask_Type), // attn_mask_type + sizeof(NVTE_Bias_Type), // bias_type + sizeof(int64_t), // window_size_left + sizeof(int64_t), // window_size_right + sizeof(uint8_t), // bottom_right_diagonal + sizeof(NVTE_Softmax_Type), // softmax_type + sizeof(NVTEScalingMode), // scaling_mode + sizeof(float), // dropout + sizeof(float), // attn_scale + // Tensor types + sizeof(NVTEDType), // qkv_dtype + sizeof(NVTEDType), // o_dtype + sizeof(NVTEDType), // do_dtype + sizeof(NVTEDType), // dqkv_dtype + // Tensor layouts + sizeof(NVTE_QKV_Layout), // qkv_layout + sizeof(NVTE_QKV_Format), // o_format + sizeof(NVTE_QKV_Format), // do_format + sizeof(NVTE_QKV_Layout), // dqkv_layout + sizeof(NVTE_QKV_Format), // qkv_scale_inv_format + sizeof(NVTE_QKV_Format), // do_scale_inv_format + // Tensor dimensions + sizeof(size_t), // batch_size + sizeof(size_t), // num_attn_heads + sizeof(size_t), // num_gqa_groups + sizeof(size_t), // head_dim_qk + sizeof(size_t), // head_dim_v + sizeof(size_t), // max_seqlen_q + sizeof(size_t), // max_seqlen_kv + sizeof(size_t), // num_tokens_q + sizeof(size_t), // num_tokens_kv + // Paged KV dimensions + sizeof(size_t), // num_pages_k + sizeof(size_t), // num_pages_v + sizeof(size_t), // page_size_k + sizeof(size_t), // page_size_v + sizeof(size_t), // max_pages_per_seq_k + sizeof(size_t), // max_pages_per_seq_v + // Bias dimensions + sizeof(size_t), // bias_batch_size + sizeof(size_t), // bias_num_heads + sizeof(size_t), // bias_seqlen_q + sizeof(size_t), // bias_seqlen_kv + }; + + static_assert(sizeof(attr_sizes) / sizeof(attr_sizes[0]) == kNVTEFusedAttnConfigNumAttributes, + "attr_sizes must have one entry per NVTEFusedAttnConfigAttribute; add the size of " + "the new attribute alongside its enumerator."); + + bool operator<(const FusedAttnConfig &rhs) const { + return std::tie(is_training, deterministic, cuda_graph, return_max_logit, attn_mask_type, + bias_type, window_size_left, window_size_right, bottom_right_diagonal, + softmax_type, scaling_mode, dropout, attn_scale, qkv_dtype, o_dtype, do_dtype, + dqkv_dtype, qkv_layout, o_format, do_format, dqkv_layout, qkv_scale_inv_format, + do_scale_inv_format, batch_size, num_attn_heads, num_gqa_groups, head_dim_qk, + head_dim_v, max_seqlen_q, max_seqlen_kv, num_tokens_q, num_tokens_kv, + num_pages_k, num_pages_v, page_size_k, page_size_v, max_pages_per_seq_k, + max_pages_per_seq_v, bias_batch_size, bias_num_heads, bias_seqlen_q, + bias_seqlen_kv, device_id) < + std::tie(rhs.is_training, rhs.deterministic, rhs.cuda_graph, rhs.return_max_logit, + rhs.attn_mask_type, rhs.bias_type, rhs.window_size_left, rhs.window_size_right, + rhs.bottom_right_diagonal, rhs.softmax_type, rhs.scaling_mode, rhs.dropout, + rhs.attn_scale, rhs.qkv_dtype, rhs.o_dtype, rhs.do_dtype, rhs.dqkv_dtype, + rhs.qkv_layout, rhs.o_format, rhs.do_format, rhs.dqkv_layout, + rhs.qkv_scale_inv_format, rhs.do_scale_inv_format, rhs.batch_size, + rhs.num_attn_heads, rhs.num_gqa_groups, rhs.head_dim_qk, rhs.head_dim_v, + rhs.max_seqlen_q, rhs.max_seqlen_kv, rhs.num_tokens_q, rhs.num_tokens_kv, + rhs.num_pages_k, rhs.num_pages_v, rhs.page_size_k, rhs.page_size_v, + rhs.max_pages_per_seq_k, rhs.max_pages_per_seq_v, rhs.bias_batch_size, + rhs.bias_num_heads, rhs.bias_seqlen_q, rhs.bias_seqlen_kv, rhs.device_id); + } + + // Derive relevant fields based on input fields that have been set by the caller. They are + // read by the graph build, cache lookup, and support query. + void derive(); + + // Assert that derive() has run, for code about to read a derived field. + void check_derived() const { + NVTE_CHECK(is_derived, + "FusedAttnConfig's derived fields are not set. Please run " + "FusedAttnConfig::derive() first."); + } + + // Return a normalized copy of this config to be used as a key for the cuDNN graph cache. + // It drops fields that are either invariant (e.g. attn_scale) or irrelevant (e.g. dO/dQKV dtypes + // and `deterministic` for forward, and `return_max_logit` for backward). + FusedAttnConfig make_cache_key(Pass pass) const; + + // Return a string representation of this config for level-2 cache diagnostics. + std::string to_string() const; +}; + +inline const FusedAttnConfig *get_fused_attn_config(NVTEFusedAttnConfig config) { + NVTE_CHECK(config != nullptr, "NVTEFusedAttnConfig must not be NULL."); + return reinterpret_cast(config); +} + +inline FusedAttnConfig *get_fused_attn_config_mutable(NVTEFusedAttnConfig config) { + NVTE_CHECK(config != nullptr, "NVTEFusedAttnConfig must not be NULL."); + return reinterpret_cast(config); +} + +struct FusedAttnFwdParams { + // Input tensors + NVTETensor Q = nullptr; + NVTETensor K = nullptr; + NVTETensor V = nullptr; + NVTETensor Bias = nullptr; + NVTETensor SoftmaxOffset = nullptr; + // Intermediate tensors + NVTETensor S = nullptr; + // Output tensor + NVTETensor O = nullptr; + // Auxiliary context tensor pack + NVTETensorPack *Aux_CTX_Tensors = nullptr; + // Miscellaneous tensors + NVTETensor cu_seqlens_q = nullptr; + NVTETensor cu_seqlens_kv = nullptr; + NVTETensor cu_seqlens_q_padded = nullptr; + NVTETensor cu_seqlens_kv_padded = nullptr; + NVTETensor page_table_k = nullptr; + NVTETensor page_table_v = nullptr; + NVTETensor rng_state = nullptr; + // Scalars + bool is_training = true; + bool cuda_graph = false; + bool return_max_logit = false; + NVTE_Mask_Type attn_mask_type = NVTE_NO_MASK; + NVTE_Bias_Type bias_type = NVTE_NO_BIAS; + int64_t window_size_left = -1; + int64_t window_size_right = -1; + bool bottom_right_diagonal = true; + NVTE_Softmax_Type softmax_type = NVTE_VANILLA_SOFTMAX; + float dropout = 0.0f; + float attn_scale = 1.0f; + NVTE_QKV_Layout qkv_layout = NVTE_QKV_Layout_NOT_SET; + NVTE_QKV_Format o_format = NVTE_QKV_Format_NOT_SET; + NVTE_QKV_Format qkv_scale_inv_format = NVTE_QKV_Format_NOT_SET; + size_t max_seqlen_q = 0; + size_t max_seqlen_kv = 0; + // Workspace and stream + NVTETensor workspace = nullptr; + cudaStream_t stream = nullptr; + + static constexpr size_t attr_sizes[] = { + // Tensor handles + sizeof(NVTETensor), // Q + sizeof(NVTETensor), // K + sizeof(NVTETensor), // V + sizeof(NVTETensor), // Bias + sizeof(NVTETensor), // SoftmaxOffset + sizeof(NVTETensor), // S + sizeof(NVTETensor), // O + sizeof(NVTETensorPack *), // Aux_CTX_Tensors + sizeof(NVTETensor), // cu_seqlens_q + sizeof(NVTETensor), // cu_seqlens_kv + sizeof(NVTETensor), // cu_seqlens_q_padded + sizeof(NVTETensor), // cu_seqlens_kv_padded + sizeof(NVTETensor), // page_table_k + sizeof(NVTETensor), // page_table_v + sizeof(NVTETensor), // rng_state + // Configuration knobs + sizeof(uint8_t), // is_training + sizeof(uint8_t), // cuda_graph + sizeof(uint8_t), // return_max_logit + sizeof(NVTE_Mask_Type), // attn_mask_type + sizeof(NVTE_Bias_Type), // bias_type + sizeof(int64_t), // window_size_left + sizeof(int64_t), // window_size_right + sizeof(uint8_t), // bottom_right_diagonal + sizeof(NVTE_Softmax_Type), // softmax_type + sizeof(float), // dropout + sizeof(float), // attn_scale + sizeof(NVTE_QKV_Layout), // qkv_layout + sizeof(NVTE_QKV_Format), // o_format + sizeof(NVTE_QKV_Format), // qkv_scale_inv_format + sizeof(size_t), // max_seqlen_q + sizeof(size_t), // max_seqlen_kv + // Workspace and stream + sizeof(NVTETensor), // workspace + sizeof(cudaStream_t), // stream + }; + + static_assert(sizeof(attr_sizes) / sizeof(attr_sizes[0]) == kNVTEFusedAttnFwdParamsNumAttributes, + "attr_sizes must have one entry per NVTEFusedAttnFwdParamsAttribute; add the size " + "of the new attribute alongside its enumerator."); + + // Build a FusedAttnConfig from the scalar "knobs" carried here (e.g. attn_mask_type, bias_type) + // and the fields derived from the tensor handles (dtypes, dims, scaling mode, paged-KV and bias + // broadcast shapes). Returns the real execution config; call FusedAttnConfig::make_cache_key on + // it to obtain the normalized cuDNN graph-cache key. + FusedAttnConfig make_config() const; +}; + +inline const FusedAttnFwdParams *get_fused_attn_fwd_params(NVTEFusedAttnFwdParams params) { + NVTE_CHECK(params != nullptr, "NVTEFusedAttnFwdParams must not be NULL."); + return reinterpret_cast(params); +} + +inline FusedAttnFwdParams *get_fused_attn_fwd_params_mutable(NVTEFusedAttnFwdParams params) { + NVTE_CHECK(params != nullptr, "NVTEFusedAttnFwdParams must not be NULL."); + return reinterpret_cast(params); +} + +struct FusedAttnBwdParams { + // Input tensors + NVTETensor Q = nullptr; + NVTETensor K = nullptr; + NVTETensor V = nullptr; + NVTETensor O = nullptr; + NVTETensor dO = nullptr; + NVTETensor S = nullptr; + NVTETensor dP = nullptr; + const NVTETensorPack *Aux_CTX_Tensors = nullptr; + // Output tensors + NVTETensor dQ = nullptr; + NVTETensor dK = nullptr; + NVTETensor dV = nullptr; + NVTETensor dBias = nullptr; + NVTETensor dSoftmaxOffset = nullptr; + // Miscellaneous tensors + NVTETensor cu_seqlens_q = nullptr; + NVTETensor cu_seqlens_kv = nullptr; + NVTETensor cu_seqlens_q_padded = nullptr; + NVTETensor cu_seqlens_kv_padded = nullptr; + // Scalars + bool deterministic = false; + bool cuda_graph = false; + NVTE_Mask_Type attn_mask_type = NVTE_NO_MASK; + NVTE_Bias_Type bias_type = NVTE_NO_BIAS; + int64_t window_size_left = -1; + int64_t window_size_right = -1; + bool bottom_right_diagonal = true; + NVTE_Softmax_Type softmax_type = NVTE_VANILLA_SOFTMAX; + float dropout = 0.0f; + float attn_scale = 1.0f; + NVTE_QKV_Layout qkv_layout = NVTE_QKV_Layout_NOT_SET; + NVTE_QKV_Format o_format = NVTE_QKV_Format_NOT_SET; + NVTE_QKV_Format do_format = NVTE_QKV_Format_NOT_SET; + NVTE_QKV_Layout dqkv_layout = NVTE_QKV_Layout_NOT_SET; + NVTE_QKV_Format qkv_scale_inv_format = NVTE_QKV_Format_NOT_SET; + NVTE_QKV_Format do_scale_inv_format = NVTE_QKV_Format_NOT_SET; + size_t max_seqlen_q = 0; + size_t max_seqlen_kv = 0; + // Workspace and stream + NVTETensor workspace = nullptr; + cudaStream_t stream = nullptr; + + static constexpr size_t attr_sizes[] = { + // Tensor handles + sizeof(NVTETensor), // Q + sizeof(NVTETensor), // K + sizeof(NVTETensor), // V + sizeof(NVTETensor), // O + sizeof(NVTETensor), // dO + sizeof(NVTETensor), // S + sizeof(NVTETensor), // dP + sizeof(const NVTETensorPack *), // Aux_CTX_Tensors + sizeof(NVTETensor), // dQ + sizeof(NVTETensor), // dK + sizeof(NVTETensor), // dV + sizeof(NVTETensor), // dBias + sizeof(NVTETensor), // dSoftmaxOffset + sizeof(NVTETensor), // cu_seqlens_q + sizeof(NVTETensor), // cu_seqlens_kv + sizeof(NVTETensor), // cu_seqlens_q_padded + sizeof(NVTETensor), // cu_seqlens_kv_padded + // Configuration knobs + sizeof(uint8_t), // deterministic + sizeof(uint8_t), // cuda_graph + sizeof(NVTE_Mask_Type), // attn_mask_type + sizeof(NVTE_Bias_Type), // bias_type + sizeof(int64_t), // window_size_left + sizeof(int64_t), // window_size_right + sizeof(uint8_t), // bottom_right_diagonal + sizeof(NVTE_Softmax_Type), // softmax_type + sizeof(float), // dropout + sizeof(float), // attn_scale + sizeof(NVTE_QKV_Layout), // qkv_layout + sizeof(NVTE_QKV_Format), // o_format + sizeof(NVTE_QKV_Format), // do_format + sizeof(NVTE_QKV_Layout), // dqkv_layout + sizeof(NVTE_QKV_Format), // qkv_scale_inv_format + sizeof(NVTE_QKV_Format), // do_scale_inv_format + sizeof(size_t), // max_seqlen_q + sizeof(size_t), // max_seqlen_kv + // Workspace and stream + sizeof(NVTETensor), // workspace + sizeof(cudaStream_t), // stream + }; + + static_assert(sizeof(attr_sizes) / sizeof(attr_sizes[0]) == kNVTEFusedAttnBwdParamsNumAttributes, + "attr_sizes must have one entry per NVTEFusedAttnBwdParamsAttribute; add the size " + "of the new attribute alongside its enumerator."); + + // Build a FusedAttnConfig from the scalar "knobs" carried here (e.g. attn_mask_type, bias_type) + // and the fields derived from the tensor handles (e.g. dtypes, dims, scaling mode and bias broadcast + // shape). Returns the real execution config; call FusedAttnConfig::make_cache_key on it to + // obtain the normalized cuDNN graph-cache key. + FusedAttnConfig make_config() const; +}; + +inline const FusedAttnBwdParams *get_fused_attn_bwd_params(NVTEFusedAttnBwdParams params) { + NVTE_CHECK(params != nullptr, "NVTEFusedAttnBwdParams must not be NULL."); + return reinterpret_cast(params); +} + +inline FusedAttnBwdParams *get_fused_attn_bwd_params_mutable(NVTEFusedAttnBwdParams params) { + NVTE_CHECK(params != nullptr, "NVTEFusedAttnBwdParams must not be NULL."); + return reinterpret_cast(params); +} + +} // namespace fused_attn +} // namespace transformer_engine + +#endif // TRANSFORMER_ENGINE_COMMON_FUSED_ATTN_CONFIG_AND_PARAMS_H_ diff --git a/transformer_engine/common/fused_attn/fused_attn.cpp b/transformer_engine/common/fused_attn/fused_attn.cpp index 0e773fdc84..6dee715dae 100644 --- a/transformer_engine/common/fused_attn/fused_attn.cpp +++ b/transformer_engine/common/fused_attn/fused_attn.cpp @@ -10,6 +10,7 @@ #include "../cudnn_utils.h" #include "../util/cuda_runtime.h" #include "../util/system.h" +#include "config_and_params.h" #include "fused_attn_f16_arbitrary_seqlen.h" #include "fused_attn_fp8.h" #include "utils.h" @@ -183,7 +184,7 @@ NVTE_QKV_Format nvte_get_qkv_format(NVTE_QKV_Layout qkv_layout) { // map NVTE_QKV_Layout to NVTE_QKV_Format for Q NVTE_QKV_Format nvte_get_q_format(NVTE_QKV_Layout qkv_layout) { - NVTE_QKV_Format qkv_format = nvte_get_qkv_format(qkv_layout); + const NVTE_QKV_Format qkv_format = nvte_get_qkv_format(qkv_layout); switch (qkv_format) { case NVTE_QKV_Format::NVTE_SBHD: case NVTE_QKV_Format::NVTE_SBHD_2BSHD: @@ -205,7 +206,7 @@ NVTE_QKV_Format nvte_get_q_format(NVTE_QKV_Layout qkv_layout) { // map NVTE_QKV_Layout to NVTE_QKV_Format for KV NVTE_QKV_Format nvte_get_kv_format(NVTE_QKV_Layout qkv_layout) { - NVTE_QKV_Format qkv_format = nvte_get_qkv_format(qkv_layout); + const NVTE_QKV_Format qkv_format = nvte_get_qkv_format(qkv_layout); switch (qkv_format) { case NVTE_QKV_Format::NVTE_SBHD: case NVTE_QKV_Format::NVTE_BSHD_2SBHD: @@ -225,335 +226,296 @@ NVTE_QKV_Format nvte_get_kv_format(NVTE_QKV_Layout qkv_layout) { } } -// select a backend for fused attention -NVTE_Fused_Attn_Backend nvte_get_fused_attn_backend( - bool is_training, NVTEDType q_dtype, NVTEDType kv_dtype, NVTE_QKV_Layout qkv_layout, - NVTE_Bias_Type bias_type, NVTE_Mask_Type attn_mask_type, NVTE_Softmax_Type softmax_type, - float dropout, size_t num_attn_heads, size_t num_gqa_groups, size_t max_seqlen_q, - size_t max_seqlen_kv, size_t head_dim_qk, size_t head_dim_v, int64_t window_size_left, - int64_t window_size_right, bool return_max_logit, bool cuda_graph, bool deterministic) { +namespace { + +// The per-thread storage for the diagnostic string +thread_local std::string fused_attn_backend_message_buffer; + +// Records `reason` in the per-thread buffer and `message`; returns with "no backend" +[[nodiscard]] NVTE_Fused_Attn_Backend reject(const char **message, std::string reason) { + if (message != nullptr) { + fused_attn_backend_message_buffer = std::move(reason); + *message = fused_attn_backend_message_buffer.c_str(); + } + return NVTE_Fused_Attn_Backend::NVTE_No_Backend; +} + +} // namespace + +// Fused attention backend query: runs TE's specific rules first, then cuDNN's. First rejection wins. +NVTE_Fused_Attn_Backend nvte_get_fused_attn_backend_v2(NVTEFusedAttnConfig config, + const char **message) { + NVTE_API_CALL(nvte_get_fused_attn_backend_v2); using namespace transformer_engine; - NVTE_Fused_Attn_Backend backend = NVTE_Fused_Attn_Backend::NVTE_No_Backend; - const int device_id = cuda::current_device(); - const int sm_arch_ = cuda::sm_arch(device_id); - NVTE_CHECK(q_dtype == kv_dtype, "Q and KV must have the same data type."); - NVTE_QKV_Format qkv_format = nvte_get_qkv_format(qkv_layout); - NVTE_QKV_Format q_format = nvte_get_q_format(qkv_layout); - NVTE_QKV_Format kv_format = nvte_get_kv_format(qkv_layout); - const bool is_thd_layout = - q_format == NVTE_QKV_Format::NVTE_THD || kv_format == NVTE_QKV_Format::NVTE_THD; - NVTE_QKV_Layout_Group layout_group = nvte_get_qkv_layout_group(qkv_layout); - auto cudnn_runtime_version = cudnnGetVersion(); - - // For ragged offsets we only support 32-bit prior to cuDNN 9.5 - // Only used when THD format is requested. - const bool requires_64bit_ragged_offset = - (qkv_format == NVTE_THD && fused_attn::get_ragged_offset_dtype( - layout_group, num_attn_heads, num_gqa_groups, max_seqlen_q, - max_seqlen_kv, head_dim_qk, head_dim_v) == DType::kInt64); - const bool supported_ragged_offset_size = - (!requires_64bit_ragged_offset || cudnn_runtime_version >= 90500); - - if ((q_dtype == NVTEDType::kNVTEFloat8E4M3 || q_dtype == NVTEDType::kNVTEFloat8E5M2) && - sm_arch_ >= 90 && bias_type == NVTE_Bias_Type::NVTE_NO_BIAS && - ( - // 9.2.1: {bshd, sbhd}, any seqlen, d=128, {no_mask, causal} - (cudnn_runtime_version >= 90201 && sm_arch_ < 100 && max_seqlen_q % 128 == 0 && - max_seqlen_kv % 128 == 0 && head_dim_qk == 128 && head_dim_v == 128 && - (attn_mask_type == NVTE_Mask_Type::NVTE_CAUSAL_MASK || - attn_mask_type == NVTE_Mask_Type::NVTE_NO_MASK)) || - // 9.7: {bshd, sbhd}, any seqlen, d<=256 for sm90 and d<=128 for sm100, {padding, padding_causal} - (cudnn_runtime_version >= 90700 && - // TODO (cyang): add is_training to nvte_get_fused_attn_backend - // sm90: fwd d<=256, bwd d=128 only - // sm100: fwd d<=128, bwd d<=128 - ((sm_arch_ < 100 && (!is_training) && head_dim_qk <= 256 && head_dim_v <= 256) || - (sm_arch_ < 100 && is_training && head_dim_qk == 128 && head_dim_v == 128) || - (sm_arch_ >= 100 && head_dim_qk <= 128 && head_dim_v <= 128)) && - head_dim_qk % 16 == 0 && head_dim_v % 16 == 0 && - (attn_mask_type == NVTE_Mask_Type::NVTE_NO_MASK || - attn_mask_type == NVTE_Mask_Type::NVTE_CAUSAL_MASK || - attn_mask_type == NVTE_Mask_Type::NVTE_PADDING_MASK || - attn_mask_type == NVTE_Mask_Type::NVTE_PADDING_CAUSAL_MASK || - (sm_arch_ >= 100 && - attn_mask_type == NVTE_Mask_Type::NVTE_PADDING_CAUSAL_BOTTOM_RIGHT_MASK))) || - // 9.21: d_qk=192, d_v=128 - (cudnn_runtime_version >= 92100 && sm_arch_ >= 100 && head_dim_qk <= 192 && - head_dim_v <= 128 && head_dim_qk % 16 == 0 && head_dim_v % 16 == 0 && - (attn_mask_type == NVTE_Mask_Type::NVTE_NO_MASK || - attn_mask_type == NVTE_Mask_Type::NVTE_CAUSAL_MASK || - attn_mask_type == NVTE_Mask_Type::NVTE_CAUSAL_BOTTOM_RIGHT_MASK))) && - // pre-9.21: {bshd, sbhd}, {vanilla} - // 9.21+: {bshd, sbhd, bhsd}, {vanilla, off-by-one, learnable} - // 9.23+: {thd}; sm90 fwd only, sm100+ fwd/bwd - ((cudnn_runtime_version < 92100 && - (qkv_format == NVTE_QKV_Format::NVTE_BSHD || qkv_format == NVTE_QKV_Format::NVTE_SBHD) && - softmax_type == NVTE_Softmax_Type::NVTE_VANILLA_SOFTMAX) || - (cudnn_runtime_version >= 92100 && - (qkv_format == NVTE_QKV_Format::NVTE_BSHD || qkv_format == NVTE_QKV_Format::NVTE_SBHD || - qkv_format == NVTE_QKV_Format::NVTE_BHSD)) || - ((cudnn_runtime_version >= 92300 && (sm_arch_ >= 100 || (sm_arch_ >= 90 && !is_training))) && - qkv_format == NVTE_QKV_Format::NVTE_THD && supported_ragged_offset_size && - // cuDNN before 9.26 misindexes ragged FP8 Stats during sink-token backward. - (!is_training || softmax_type == NVTE_Softmax_Type::NVTE_VANILLA_SOFTMAX || - cudnn_runtime_version >= 92600) && - (attn_mask_type == NVTE_Mask_Type::NVTE_PADDING_MASK || - attn_mask_type == NVTE_Mask_Type::NVTE_PADDING_CAUSAL_MASK || - attn_mask_type == NVTE_Mask_Type::NVTE_PADDING_CAUSAL_BOTTOM_RIGHT_MASK))) && - // 9.10.0: known bugs with SDPA FP8 - (cudnn_runtime_version != 91000) && !return_max_logit) { - backend = NVTE_Fused_Attn_Backend::NVTE_FP8; - } else if ((q_dtype == NVTEDType::kNVTEFloat16) || (q_dtype == NVTEDType::kNVTEBFloat16)) { - bool flag_arb = false; - if ( - // TODO(cyang): replace with cudnn-frontend check_support for cleaner logic and better error messaging - // architecture - ((cudnn_runtime_version < 8903 && (sm_arch_ == 80 || sm_arch_ == 90)) || - (cudnn_runtime_version >= 8903 && sm_arch_ >= 80 && sm_arch_ < 100) || - (cudnn_runtime_version >= 90700 && sm_arch_ >= 100)) && - // sequence length - ((cudnn_runtime_version < 90000 && max_seqlen_q % 64 == 0 && max_seqlen_kv % 64 == 0) || - (cudnn_runtime_version >= 90000)) && - // number of heads - ((cudnn_runtime_version < 8907 && num_attn_heads == num_gqa_groups) || - (cudnn_runtime_version >= 8907)) && - // head dimension - // multiples of 8 - (head_dim_qk % 8 == 0 && head_dim_v % 8 == 0 && - // <= 128 - ((head_dim_qk <= 128 && head_dim_v <= 128) || - // 9.1: <= 256 + Hopper + fprop - // 9.5: <= 256 + Hopper + bprop - (head_dim_qk <= 256 && head_dim_v <= 256 && - ((!is_training && sm_arch_ == 90 && cudnn_runtime_version >= 90100) || - (is_training && sm_arch_ == 90 && cudnn_runtime_version >= 90500))) || - // 9.9: any head_dim + Blackwell + fprop + non_paged + sq > 1 - (!is_training && sm_arch_ >= 100 && cudnn_runtime_version >= 90900 && max_seqlen_q > 1 && - layout_group != NVTE_QKV_Layout_Group::NVTE_Paged_KV_HD_HD_HD) || - // 9.10.2: any head_dim + any arch + fprop + paged - // 9.10.2: any head_dim + any arch + fprop + non_paged + sq > 1 - // 9.10.2: any head_dim + any arch + fprop + non_paged + sq = 1 + {no_mask, padding, BRCM, padding_BRCM} - (!is_training && cudnn_runtime_version >= 91002 && - (layout_group == NVTE_QKV_Layout_Group::NVTE_Paged_KV_HD_HD_HD || max_seqlen_q > 1 || - (max_seqlen_q == 1 && attn_mask_type != NVTE_Mask_Type::NVTE_CAUSAL_MASK && - attn_mask_type != NVTE_Mask_Type::NVTE_PADDING_CAUSAL_MASK))) || - // 9.11: d_qk = 192, d_v = 128 + Blackwell + bprop + non-paged - (head_dim_qk == 192 && head_dim_v == 128 && is_training && sm_arch_ >= 100 && - cudnn_runtime_version >= 91100) || - // 9.23: d_qk = d_v = 256 + SM10x (cuDNN FE 1.24 / BE 9.23+) + bprop + non-paged. - // THD layouts require cuDNN FE 1.26 / BE 9.25+ for execution-plan support. - (head_dim_qk == 256 && head_dim_v == 256 && is_training && sm_arch_ >= 100 && - sm_arch_ < 110 && cudnn_runtime_version >= (is_thd_layout ? 92500 : 92300) && - layout_group != NVTE_QKV_Layout_Group::NVTE_Paged_KV_HD_HD_HD && - // The FE forces this path onto the deterministic bprop algorithm, which on - // Blackwell rejects dBias, dropout, and ALiBi (and supports vanilla softmax only). - bias_type == NVTE_Bias_Type::NVTE_NO_BIAS && dropout == 0.0 && - softmax_type == NVTE_Softmax_Type::NVTE_VANILLA_SOFTMAX && - // Non-causal D=256 supports only full-window attention; SWA is allowed only for causal masks. - ((window_size_left == -1 && window_size_right == -1) || - ((attn_mask_type == NVTE_Mask_Type::NVTE_CAUSAL_MASK || - attn_mask_type == NVTE_Mask_Type::NVTE_PADDING_CAUSAL_MASK || - attn_mask_type == NVTE_Mask_Type::NVTE_CAUSAL_BOTTOM_RIGHT_MASK || - attn_mask_type == NVTE_Mask_Type::NVTE_PADDING_CAUSAL_BOTTOM_RIGHT_MASK) && - (window_size_right == -1 || window_size_right == 0))))) && - // 9.11+ bug: 128 < d_qk <= 256, 128 < d_v <= 256 + Hopper + bprop + MLA - // Conditional to temporarily use blanket cudnn_runtime_version >= 9.11 until fixed - (!((cudnn_runtime_version >= 91100) && is_training && sm_arch_ == 90 && - head_dim_qk >= 128 && head_dim_v >= 128 && !(head_dim_qk == 192 && head_dim_v == 128) && - head_dim_qk != head_dim_v))) && - // bias type - ((cudnn_runtime_version < 8906 && bias_type == NVTE_Bias_Type::NVTE_NO_BIAS) || - (cudnn_runtime_version >= 8906 && - (bias_type == NVTE_Bias_Type::NVTE_NO_BIAS || - (bias_type == NVTE_Bias_Type::NVTE_ALIBI && - attn_mask_type != NVTE_Mask_Type::NVTE_NO_MASK && - attn_mask_type != NVTE_Mask_Type::NVTE_PADDING_MASK && - attn_mask_type != NVTE_Mask_Type::NVTE_PADDING_CAUSAL_MASK && - attn_mask_type != NVTE_Mask_Type::NVTE_PADDING_CAUSAL_BOTTOM_RIGHT_MASK && - sm_arch_ >= 90) || - (bias_type == NVTE_Bias_Type::NVTE_POST_SCALE_BIAS && sm_arch_ >= 90))) || - (cudnn_runtime_version >= 90000 && - (bias_type == NVTE_Bias_Type::NVTE_POST_SCALE_BIAS && sm_arch_ >= 80))) && - // mask type - // pre-8.9.6: causal - ((cudnn_runtime_version < 8906 && attn_mask_type == NVTE_Mask_Type::NVTE_CAUSAL_MASK) || - // 8.9.6: {bshd, sbhd} + {no_mask, causal, padding, padding_causal} - (cudnn_runtime_version >= 8906 && - (qkv_format == NVTE_QKV_Format::NVTE_SBHD || qkv_format == NVTE_QKV_Format::NVTE_BSHD) && - (attn_mask_type == NVTE_Mask_Type::NVTE_CAUSAL_MASK || - attn_mask_type == NVTE_Mask_Type::NVTE_PADDING_MASK || - attn_mask_type == NVTE_Mask_Type::NVTE_PADDING_CAUSAL_MASK || - attn_mask_type == NVTE_Mask_Type::NVTE_NO_MASK)) || - // 9.1: adds thd + {padding, padding_causal} - (cudnn_runtime_version >= 90100 && qkv_format == NVTE_QKV_Format::NVTE_THD && - (attn_mask_type == NVTE_Mask_Type::NVTE_PADDING_MASK || - attn_mask_type == NVTE_Mask_Type::NVTE_PADDING_CAUSAL_MASK)) || - // 9.3: adds {bshd, sbhd} + causal_bottom_right + self/cross-attn (sq <= skv) - (cudnn_runtime_version >= 90300 && - (qkv_format == NVTE_QKV_Format::NVTE_SBHD || qkv_format == NVTE_QKV_Format::NVTE_BSHD) && - attn_mask_type == NVTE_Mask_Type::NVTE_CAUSAL_BOTTOM_RIGHT_MASK && - max_seqlen_q % 64 == 0 && max_seqlen_kv % 64 == 0 && max_seqlen_q <= max_seqlen_kv && - bias_type == NVTE_Bias_Type::NVTE_NO_BIAS && dropout == 0.0) || - // 9.5: adds {paged_kv_bshd, paged_kv_sbhd} + {padding, padding_causal, padding_causal_bottom_right} - (cudnn_runtime_version >= 90500 && - layout_group == NVTE_QKV_Layout_Group::NVTE_Paged_KV_HD_HD_HD && - (attn_mask_type == NVTE_Mask_Type::NVTE_PADDING_MASK || - attn_mask_type == NVTE_Mask_Type::NVTE_PADDING_CAUSAL_MASK || - (attn_mask_type == NVTE_Mask_Type::NVTE_PADDING_CAUSAL_BOTTOM_RIGHT_MASK && - max_seqlen_q % 64 == 0 && max_seqlen_kv % 64 == 0 && max_seqlen_q <= max_seqlen_kv)) && - bias_type == NVTE_Bias_Type::NVTE_NO_BIAS && dropout == 0.0) || - // 9.6: adds {bshd, sbhd, thd} + padding_causal_bottom_right + self/cross-attn (sq <= skv) - (cudnn_runtime_version >= 90600 && - attn_mask_type == NVTE_Mask_Type::NVTE_PADDING_CAUSAL_BOTTOM_RIGHT_MASK && - max_seqlen_q % 64 == 0 && max_seqlen_kv % 64 == 0 && max_seqlen_q <= max_seqlen_kv && - bias_type == NVTE_Bias_Type::NVTE_NO_BIAS && dropout == 0.0) || - // 9.7: removes s_q/s_kv % 64 = 0 for {causal_bottom_right, padding_causal_bottom_right} - // for any q_format/kv_format, and paged/non-paged - (cudnn_runtime_version >= 90700 && - (attn_mask_type == NVTE_Mask_Type::NVTE_NO_MASK || - attn_mask_type == NVTE_Mask_Type::NVTE_CAUSAL_MASK || - ((attn_mask_type == NVTE_Mask_Type::NVTE_PADDING_MASK || - attn_mask_type == NVTE_Mask_Type::NVTE_PADDING_CAUSAL_MASK || - attn_mask_type == NVTE_Mask_Type::NVTE_PADDING_CAUSAL_BOTTOM_RIGHT_MASK) && - bias_type == NVTE_Bias_Type::NVTE_NO_BIAS && dropout == 0.0) || - ((attn_mask_type == NVTE_Mask_Type::NVTE_CAUSAL_BOTTOM_RIGHT_MASK || - attn_mask_type == NVTE_Mask_Type::NVTE_PADDING_CAUSAL_BOTTOM_RIGHT_MASK) && - max_seqlen_q <= max_seqlen_kv)))) && - // bias + mask combination - (!(cudnn_runtime_version >= 8906 && - (attn_mask_type == NVTE_Mask_Type::NVTE_PADDING_MASK || - attn_mask_type == NVTE_Mask_Type::NVTE_PADDING_CAUSAL_MASK) && - bias_type == NVTE_Bias_Type::NVTE_POST_SCALE_BIAS)) && - // qkv format - (qkv_format == NVTE_QKV_Format::NVTE_SBHD || qkv_format == NVTE_QKV_Format::NVTE_BSHD || - qkv_format == NVTE_QKV_Format::NVTE_BHSD || - (qkv_format == NVTE_QKV_Format::NVTE_THD && sm_arch_ >= 90 && - ((cudnn_runtime_version >= 90100 && num_attn_heads == num_gqa_groups) || - cudnn_runtime_version >= 90600)) || - ((q_format == NVTE_QKV_Format::NVTE_SBHD || q_format == NVTE_QKV_Format::NVTE_BSHD || - q_format == NVTE_QKV_Format::NVTE_BHSD || - (q_format == NVTE_QKV_Format::NVTE_THD && sm_arch_ >= 90) || - kv_format == NVTE_QKV_Format::NVTE_SBHD || kv_format == NVTE_QKV_Format::NVTE_BSHD || - kv_format == NVTE_QKV_Format::NVTE_BHSD || - (kv_format == NVTE_QKV_Format::NVTE_THD && sm_arch_ >= 90)) && - cudnn_runtime_version >= 90700)) && - // sliding window - // pre-9.2: full attn, causal - ((cudnn_runtime_version < 90200 && window_size_left == -1 && - (window_size_right == -1 || window_size_right == 0)) || - // 9.2: SWA (left, 0) + top-left diagonal + {bshd, sbhd} - (cudnn_runtime_version >= 90200 && - ((window_size_left == -1 && window_size_right == -1 && - attn_mask_type == NVTE_Mask_Type::NVTE_NO_MASK) || - ((window_size_left == -1 || window_size_left >= 0) && window_size_right == 0 && - (attn_mask_type == NVTE_Mask_Type::NVTE_NO_MASK || - attn_mask_type == NVTE_Mask_Type::NVTE_CAUSAL_MASK || - (attn_mask_type == NVTE_Mask_Type::NVTE_CAUSAL_BOTTOM_RIGHT_MASK && - max_seqlen_q == max_seqlen_kv)) && - max_seqlen_q <= max_seqlen_kv && dropout == 0.0 && - bias_type == NVTE_Bias_Type::NVTE_NO_BIAS && - (qkv_format == NVTE_QKV_Format::NVTE_BSHD || - qkv_format == NVTE_QKV_Format::NVTE_SBHD)))) || - // 9.6: SWA (left, 0) + top-left/bottom-right diagonal + {bshd, sbhd, thd} - (cudnn_runtime_version >= 90600 && - ((window_size_left == -1 && (window_size_right == -1 || window_size_right == 0)) || - ((window_size_left >= 0 || window_size_left == -1) && - (window_size_right >= 0 || window_size_right == -1) && - ((attn_mask_type == NVTE_Mask_Type::NVTE_CAUSAL_BOTTOM_RIGHT_MASK && - // TODO(cyang): fix bug for BRCM + cross-attention on sm100 - (sm_arch_ < 100 || (sm_arch_ >= 100 && ((max_seqlen_q == max_seqlen_kv && - cudnn_runtime_version <= 90700) || - cudnn_runtime_version > 90700)))) || - attn_mask_type == NVTE_Mask_Type::NVTE_NO_MASK || - attn_mask_type == NVTE_Mask_Type::NVTE_PADDING_MASK || - attn_mask_type == NVTE_Mask_Type::NVTE_PADDING_CAUSAL_MASK || - (attn_mask_type == NVTE_Mask_Type::NVTE_PADDING_CAUSAL_BOTTOM_RIGHT_MASK && - (sm_arch_ < 100 || (sm_arch_ >= 100 && ((max_seqlen_q == max_seqlen_kv && - cudnn_runtime_version <= 90700) || - cudnn_runtime_version > 90700))))) && - max_seqlen_q <= max_seqlen_kv && bias_type == NVTE_Bias_Type::NVTE_NO_BIAS && - dropout == 0.0)))) && - // check 64-bit ragged offset support - (supported_ragged_offset_size) && - // 9.10.0/9.10.1: known bugs with SDPA F16 - (cudnn_runtime_version != 91000) && (cudnn_runtime_version != 91001) && - // softmax type - // pre-9.13.1: vanilla - // 9.13.1+: vanilla, off-by-one, learnable - (cudnn_runtime_version >= 91301 || - softmax_type == NVTE_Softmax_Type::NVTE_VANILLA_SOFTMAX) && - // max_logit - // pre-9.21: no (the composite softmax node rejects the Stats + Max output combination) - // 9.21+: yes (Stats + Max via the unified softmax node) - (!return_max_logit || cudnn_runtime_version >= 92100) && - // determinism on Blackwell - // pre-9.18.1: fwd: deterministic; bwd: non-deterministic - // 9.18.1+: fwd: deterministic; bwd: non-deterministic/deterministic - (sm_arch_ < 100 || - (sm_arch_ >= 100 && (!is_training || - (is_training && !deterministic && - (dropout == 0.0 || bias_type == NVTE_Bias_Type::NVTE_NO_BIAS)) || - (is_training && deterministic && cudnn_runtime_version >= 91801 && - dropout == 0.0 && bias_type == NVTE_Bias_Type::NVTE_NO_BIAS))))) { - flag_arb = true; + using namespace transformer_engine::fused_attn; + FusedAttnConfig cfg = *get_fused_attn_config(config); + cfg.derive(); + if (message != nullptr) *message = ""; + + cudnnHandle_t handle = cudnnExecutionPlanManager::Instance().GetHandle(); + const auto cudnn_runtime_version = cudnnGetVersion(); + const int sm_arch = cuda::sm_arch(cuda::current_device()); + + // THD + 64-bit ragged offsets require cuDNN >= 9.5 + if (cfg.needs_64bit_ragged_offset && cudnn_runtime_version < 90500) { + return reject( + message, + "This config requires 64-bit ragged offsets, which is only supported by cuDNN >= 9.5."); + } + + // THD input requires a padding mask + if ((cfg.is_ragged_q || cfg.is_ragged_kv) && !cfg.is_padding) { + return reject( + message, + "THD format requires PADDING / PADDING_CAUSAL / PADDING_CAUSAL_BOTTOM_RIGHT mask."); + } + + if ((cfg.is_ragged_q && cfg.num_tokens_q == 0) || (cfg.is_ragged_kv && cfg.num_tokens_kv == 0)) { + return reject(message, + "THD format requires num_tokens_q / num_tokens_kv to be set for the ragged " + "inputs."); + } + + // Paged KV requires a padding mask + if (cfg.is_paged_kv && !cfg.is_padding) { + return reject(message, + "Paged KV requires PADDING / PADDING_CAUSAL / PADDING_CAUSAL_BOTTOM_RIGHT mask."); + } + + // Paged KV requires cache dimensions to be set + if (cfg.is_paged_kv && + (cfg.num_pages_k == 0 || cfg.num_pages_v == 0 || cfg.page_size_k == 0 || + cfg.page_size_v == 0 || cfg.max_pages_per_seq_k == 0 || cfg.max_pages_per_seq_v == 0)) { + return reject(message, + "Paged KV requires num_pages, page_size and max_pages_per_seq to be set for both " + "K and V."); + } + + // Fused-attention does not support pre-scale bias + if (cfg.bias_type == NVTE_Bias_Type::NVTE_PRE_SCALE_BIAS) { + return reject(message, "Fused attention does not support pre-scale bias."); + } + + const bool is_fp8 = + (cfg.qkv_dtype == NVTEDType::kNVTEFloat8E4M3 || cfg.qkv_dtype == NVTEDType::kNVTEFloat8E5M2); + const bool is_f16_or_bf16 = + (cfg.qkv_dtype == NVTEDType::kNVTEFloat16 || cfg.qkv_dtype == NVTEDType::kNVTEBFloat16); + + auto each_pass = [&](auto &&verdict) -> std::string { + if (cfg.check_for_forward_support) { + std::string reason = verdict(Pass::Fwd); + if (!reason.empty()) return reason; } - if (flag_arb) { - backend = NVTE_Fused_Attn_Backend::NVTE_F16_arbitrary_seqlen; + if (cfg.is_training && cfg.check_for_backward_support) { + std::string reason = verdict(Pass::Bwd); + if (!reason.empty()) return reason; } - if (cudnn_runtime_version < 8900 && - backend == NVTE_Fused_Attn_Backend::NVTE_F16_arbitrary_seqlen) { - backend = NVTE_Fused_Attn_Backend::NVTE_No_Backend; - std::cout << "Warning: FP16/BF16 fused attention is supported by cuDNN 8.9.0+." - " Please upgrade your cuDNN version if possible." - << std::endl; + return ""; + }; + + // F16/BF16 support checks + if (is_f16_or_bf16) { + if (cfg.is_ragged_q && cfg.is_ragged_kv && sm_arch < 90) { + return reject(message, "F16/BF16 fused attention with THD format requires sm90 or later."); + } + if ((cfg.is_ragged_q || cfg.is_ragged_kv) && sm_arch < 90 && cudnn_runtime_version < 90700) { + return reject(message, + "F16/BF16 fused attention with a ragged Q and non-ragged KV requires cuDNN " + ">= 9.7 before sm90."); + } + const bool has_sliding_window = !(cfg.window_size_left == -1 && + (cfg.window_size_right == -1 || cfg.window_size_right == 0)); + if (cfg.is_causal_bottom_right && has_sliding_window && cfg.max_seqlen_q != cfg.max_seqlen_kv && + cudnn_runtime_version <= 90700 && sm_arch >= 100) { + return reject(message, + "Known cuDNN <= 9.7.0 issue with bottom-right causal masking and a sliding " + "window for cross-attention on sm100. Please upgrade cuDNN."); + } + if (cudnn_runtime_version <= 91500 && cfg.is_training && + (cfg.qkv_format == NVTE_QKV_Format::NVTE_BSHD || + cfg.qkv_format == NVTE_QKV_Format::NVTE_SBHD) && + (cfg.max_seqlen_kv % 128 != 0) && cfg.cuda_graph && !cfg.is_padding) { + return reject(message, "Known cuDNN <= 9.15 issue with CUDA graph. Please upgrade cuDNN."); + } + + // Run cuDNN support checks + std::string cudnn_reason = + each_pass([&](Pass pass) { return support_verdict_f16(cfg, pass, handle); }); + if (!cudnn_reason.empty()) return reject(message, std::move(cudnn_reason)); + return NVTE_Fused_Attn_Backend::NVTE_F16_arbitrary_seqlen; + } + + // FP8 support checks + if (is_fp8) { + if (cfg.return_max_logit) { + return reject(message, "FP8 fused attention does not support return_max_logit=True."); } - if ((cudnn_runtime_version == 91400) && (max_seqlen_kv > 1024) && (window_size_left != -1) && - (attn_mask_type != NVTE_Mask_Type::NVTE_CAUSAL_MASK) && - (attn_mask_type != NVTE_Mask_Type::NVTE_CAUSAL_BOTTOM_RIGHT_MASK)) { - backend = NVTE_Fused_Attn_Backend::NVTE_No_Backend; - std::cout << "Warning: Given combination of attention mask (non-causal) and " - "max_seqlen_kv (> 1024) does not support fused attention for cuDNN 9.14.0. " - " Please upgrade your cuDNN version if possible." - << std::endl; + if (cfg.qkv_format != NVTE_QKV_Format::NVTE_BSHD && + cfg.qkv_format != NVTE_QKV_Format::NVTE_SBHD && + cfg.qkv_format != NVTE_QKV_Format::NVTE_BHSD && + cfg.qkv_format != NVTE_QKV_Format::NVTE_THD) { + return reject(message, "FP8 fused attention supports BSHD/SBHD/BHSD/THD formats, found " + + std::to_string(static_cast(cfg.qkv_format)) + "."); } - if ((cudnn_runtime_version <= 91500) && is_training && - (qkv_format == NVTE_QKV_Format::NVTE_BSHD || qkv_format == NVTE_QKV_Format::NVTE_SBHD) && - (max_seqlen_kv % 128 != 0) && cuda_graph && - (attn_mask_type != NVTE_Mask_Type::NVTE_PADDING_MASK) && - (attn_mask_type != NVTE_Mask_Type::NVTE_PADDING_CAUSAL_MASK) && - (attn_mask_type != NVTE_Mask_Type::NVTE_PADDING_CAUSAL_BOTTOM_RIGHT_MASK)) { - backend = NVTE_Fused_Attn_Backend::NVTE_No_Backend; - std::cout << "Warning: Given combination of attention mask (non-padding)," - " max_seqlen_kv (not divisible by 128), and qkv_format (BSHD/SBHD) for" - " backward fused attention with graph capture requires cuDNN 9.15.1+. " - "Please upgrade your cuDNN version if possible." - << std::endl; + if (cfg.is_training && cfg.check_for_backward_support && + cfg.dqkv_layout != NVTE_QKV_Layout_NOT_SET) { + const NVTE_QKV_Format dqkv_format = nvte_get_qkv_format(cfg.dqkv_layout); + if (dqkv_format != NVTE_QKV_Format::NVTE_BSHD && dqkv_format != NVTE_QKV_Format::NVTE_SBHD && + dqkv_format != NVTE_QKV_Format::NVTE_BHSD && dqkv_format != NVTE_QKV_Format::NVTE_THD) { + return reject(message, + "FP8 fused attention supports BSHD/SBHD/BHSD/THD gradient formats, found " + + std::to_string(static_cast(dqkv_format)) + "."); + } } - if (backend == NVTE_Fused_Attn_Backend::NVTE_F16_arbitrary_seqlen && sm_arch_ == 120) { - if (cudnn_runtime_version < 91801) { - backend = NVTE_Fused_Attn_Backend::NVTE_No_Backend; - std::cout << "Warning: Given combination of sm_arch_ == 120 and cudnn_runtime_version < " - "91801 is not supported. " - << " Please upgrade your cuDNN version if possible." << std::endl; - } else if (deterministic && is_training) { - backend = NVTE_Fused_Attn_Backend::NVTE_No_Backend; - std::cout << "Warning: Deterministic fused attention on SM120 is not supported." - << std::endl; - } else { - // Known missing support for T3HD/TH3D layouts on SM120 - const bool is_t3hd_or_th3d = - (qkv_layout == NVTE_QKV_Layout::NVTE_T3HD || qkv_layout == NVTE_QKV_Layout::NVTE_TH3D); - if (is_t3hd_or_th3d) { - backend = NVTE_Fused_Attn_Backend::NVTE_No_Backend; - std::cout << "Warning: Given combination of T3HD/TH3D layouts on SM120 is not supported. " - << " Please consider using other THD layouts if possible." << std::endl; - } + if (cfg.qkv_format == NVTE_QKV_Format::NVTE_THD) { + if (cudnn_runtime_version < 92300) { + return reject(message, + "FP8 fused attention with THD format requires cuDNN 9.23.0 or later!"); + } + if (cfg.is_training && cfg.check_for_backward_support && sm_arch < 100) { + return reject(message, + "FP8 fused attention with THD format supports backward on sm100+ only!"); + } + if (cfg.is_training && cfg.check_for_backward_support && + cfg.softmax_type != NVTE_Softmax_Type::NVTE_VANILLA_SOFTMAX && + cudnn_runtime_version < 92600) { + return reject(message, + "FP8 fused attention with THD format and a sink token requires cuDNN 9.26.0 " + "or later for backward!"); } + if (sm_arch >= 100 && (cfg.head_dim_qk > 128 || cfg.head_dim_v > 128)) { + return reject(message, + "FP8 fused attention with THD format supports head dimensions up to 128 on " + "sm100+ only!"); + } + } + if (cfg.is_bias) { + return reject(message, "FP8 fused attention does not support pre/post_scale_bias yet!"); + } + if (cfg.is_alibi) { + return reject(message, "FP8 fused attention does not support ALiBi yet!"); } + const char *const recipe_reason = + "FP8 fused attention only supports FP8DelayedScaling or FP8CurrentScaling or MXFP8 " + "recipes!"; + if (cfg.check_for_forward_support && + !(cfg.is_delayed_scaling_fwd || cfg.is_current_scaling_fwd || cfg.is_mxfp8_fwd)) { + return reject(message, recipe_reason); + } + if (cfg.is_training && cfg.check_for_backward_support && + !(cfg.is_delayed_scaling_bwd || cfg.is_current_scaling_bwd || cfg.is_mxfp8_bwd)) { + return reject(message, recipe_reason); + } + if (cfg.is_mxfp8 && cudnn_runtime_version < 92100) { + return reject(message, "MXFP8 fused attention requires cuDNN 9.21.0 or later!"); + } + + // Run cuDNN support checks + std::string cudnn_reason = + each_pass([&](Pass pass) { return support_verdict_fp8(cfg, pass, handle); }); + if (!cudnn_reason.empty()) return reject(message, std::move(cudnn_reason)); + return NVTE_Fused_Attn_Backend::NVTE_FP8; + } + + // Unsupported dtype + return reject(message, "Unsupported QKV dtype qkv_dtype=" + std::to_string(cfg.qkv_dtype) + " ."); +} + +// Select a backend for fused attention +NVTE_Fused_Attn_Backend nvte_get_fused_attn_backend( + bool is_training, NVTEDType q_dtype, NVTEDType kv_dtype, NVTE_QKV_Layout qkv_layout, + NVTE_Bias_Type bias_type, NVTE_Mask_Type attn_mask_type, NVTE_Softmax_Type softmax_type, + float dropout, size_t num_attn_heads, size_t num_gqa_groups, size_t max_seqlen_q, + size_t max_seqlen_kv, size_t head_dim_qk, size_t head_dim_v, int64_t window_size_left, + int64_t window_size_right, bool return_max_logit, bool cuda_graph, bool deterministic) { + NVTE_API_CALL(nvte_get_fused_attn_backend); + transformer_engine::fused_attn::FusedAttnConfig cfg{}; + cfg.qkv_layout = qkv_layout; + cfg.bias_type = bias_type; + cfg.attn_mask_type = attn_mask_type; + cfg.softmax_type = softmax_type; + cfg.dropout = dropout; + cfg.max_seqlen_q = max_seqlen_q; + cfg.max_seqlen_kv = max_seqlen_kv; + cfg.window_size_left = window_size_left; + cfg.window_size_right = window_size_right; + cfg.cuda_graph = cuda_graph; + NVTE_CHECK(q_dtype == kv_dtype, "Q and KV must have the same data type."); + cfg.qkv_dtype = q_dtype; + cfg.o_dtype = q_dtype; + cfg.do_dtype = q_dtype; + cfg.dqkv_dtype = q_dtype; + cfg.num_attn_heads = num_attn_heads; + cfg.num_gqa_groups = num_gqa_groups; + cfg.head_dim_qk = head_dim_qk; + cfg.head_dim_v = head_dim_v; + cfg.is_training = is_training; + cfg.return_max_logit = return_max_logit; + cfg.deterministic = deterministic; + // fill in the missing fields with the most common use case; + // otherwise it would return NVTE_No_Backend always + cfg.batch_size = 1; + cfg.num_tokens_q = cfg.batch_size * max_seqlen_q; + cfg.num_tokens_kv = cfg.batch_size * max_seqlen_kv; + cfg.o_format = nvte_get_q_format(qkv_layout); + cfg.do_format = cfg.o_format; + cfg.dqkv_layout = qkv_layout; + if (bias_type == NVTE_Bias_Type::NVTE_POST_SCALE_BIAS) { + cfg.bias_batch_size = cfg.batch_size; + cfg.bias_num_heads = num_attn_heads; + cfg.bias_seqlen_q = max_seqlen_q; + cfg.bias_seqlen_kv = max_seqlen_kv; + } + + return nvte_get_fused_attn_backend_v2(reinterpret_cast(&cfg), + /*message=*/nullptr); +} + +// Fused attention forward: create a config based on the params, check which backend supports it, +// and run that backend's implementation. +void nvte_fused_attn_fwd_v2(NVTEFusedAttnFwdParams params) { + NVTE_API_CALL(nvte_fused_attn_fwd_v2); + using namespace transformer_engine; + using namespace transformer_engine::fused_attn; + const FusedAttnFwdParams &p = *get_fused_attn_fwd_params(params); + const Tensor *input_cu_seqlens_q = convertNVTETensorCheck(p.cu_seqlens_q); + const Tensor *input_cu_seqlens_kv = convertNVTETensorCheck(p.cu_seqlens_kv); + const Tensor *input_cu_seqlens_q_padded = convertNVTETensorCheck(p.cu_seqlens_q_padded); + const Tensor *input_cu_seqlens_kv_padded = convertNVTETensorCheck(p.cu_seqlens_kv_padded); + const Tensor *input_page_table_k = convertNVTETensorCheck(p.page_table_k); + const Tensor *input_page_table_v = convertNVTETensorCheck(p.page_table_v); + const Tensor *input_rng_state = convertNVTETensorCheck(p.rng_state); + const Tensor *input_Q = convertNVTETensorCheck(p.Q); + const Tensor *input_K = convertNVTETensorCheck(p.K); + const Tensor *input_V = convertNVTETensorCheck(p.V); + const Tensor *input_Bias = convertNVTETensorCheck(p.Bias); + const Tensor *input_SoftmaxOffset = convertNVTETensorCheck(p.SoftmaxOffset); + Tensor *input_output_S = convertNVTETensorCheck(p.S); + Tensor *output_O = convertNVTETensorCheck(p.O); + Tensor *wkspace = convertNVTETensor(p.workspace); + + auto handle = cudnnExecutionPlanManager::Instance().GetHandle(); + FusedAttnConfig cfg = p.make_config(); + cfg.derive(); + const char *fused_attn_reject_reason = nullptr; + NVTE_Fused_Attn_Backend fused_attention_backend = nvte_get_fused_attn_backend_v2( + reinterpret_cast(&cfg), &fused_attn_reject_reason); + + if (fused_attention_backend == NVTE_Fused_Attn_Backend::NVTE_F16_arbitrary_seqlen) { + fused_attn_arbitrary_seqlen_fwd(cfg, input_Q, input_K, input_V, input_Bias, input_SoftmaxOffset, + output_O, p.Aux_CTX_Tensors, input_cu_seqlens_q, + input_cu_seqlens_kv, input_cu_seqlens_q_padded, + input_cu_seqlens_kv_padded, input_page_table_k, + input_page_table_v, input_rng_state, wkspace, p.stream, handle); + } else if (fused_attention_backend == NVTE_Fused_Attn_Backend::NVTE_FP8) { + fused_attn_fp8_fwd(cfg, input_Q, input_K, input_V, input_SoftmaxOffset, input_output_S, + output_O, p.Aux_CTX_Tensors, input_cu_seqlens_q, input_cu_seqlens_kv, + input_cu_seqlens_q_padded, input_cu_seqlens_kv_padded, input_rng_state, + wkspace, p.stream, handle); } else { - backend = NVTE_Fused_Attn_Backend::NVTE_No_Backend; + NVTE_ERROR("Fused attention is not supported for the user configuration: ", + fused_attn_reject_reason); } - return backend; } // NVTE fused attention FWD with separate Q, K and V @@ -571,100 +533,117 @@ void nvte_fused_attn_fwd(const NVTETensor Q, const NVTETensor K, const NVTETenso NVTE_Mask_Type attn_mask_type, NVTE_Softmax_Type softmax_type, int64_t window_size_left, int64_t window_size_right, bool bottom_right_diagonal, NVTETensor workspace, cudaStream_t stream) { - NVTE_API_CALL(nvte_flash_attn_fwd); - using namespace transformer_engine; - const Tensor *input_cu_seqlens_q = convertNVTETensorCheck(cu_seqlens_q); - const Tensor *input_cu_seqlens_kv = convertNVTETensorCheck(cu_seqlens_kv); - const Tensor *input_cu_seqlens_q_padded = convertNVTETensorCheck(cu_seqlens_q_padded); - const Tensor *input_cu_seqlens_kv_padded = convertNVTETensorCheck(cu_seqlens_kv_padded); - const Tensor *input_page_table_k = convertNVTETensorCheck(page_table_k); - const Tensor *input_page_table_v = convertNVTETensorCheck(page_table_v); - const Tensor *input_rng_state = convertNVTETensorCheck(rng_state); - const Tensor *input_Q = convertNVTETensorCheck(Q); - const Tensor *input_K = convertNVTETensorCheck(K); - const Tensor *input_V = convertNVTETensorCheck(V); - const Tensor *input_Bias = convertNVTETensorCheck(Bias); - const Tensor *input_SoftmaxOffset = convertNVTETensorCheck(SoftmaxOffset); - Tensor *input_output_S = convertNVTETensorCheck(S); - Tensor *output_O = convertNVTETensorCheck(O); - Tensor *wkspace = convertNVTETensor(workspace); - - NVTE_QKV_Format q_format = nvte_get_q_format(qkv_layout); - NVTE_QKV_Format kv_format = nvte_get_kv_format(qkv_layout); - auto *q_dims = input_Q->data.shape.data(); - auto *k_dims = input_K->data.shape.data(); - auto *v_dims = input_V->scaling_mode != NVTE_MXFP8_1D_SCALING - ? input_V->data.shape.data() - : input_V->columnwise_data.shape.data(); - AttentionShape q_shape(q_format, q_dims); - AttentionShape k_shape(kv_format, k_dims); - AttentionShape v_shape(kv_format, v_dims); - size_t b = q_shape.b(), h_q = q_shape.h(), d_qk = q_shape.d(), t_q = q_shape.t(); - size_t h_kv = k_shape.h(), t_kv = k_shape.t(), d_v = v_shape.d(); - if (q_format == NVTE_QKV_Format::NVTE_THD) { - b = input_cu_seqlens_q->data.shape[0] - 1; - } else if (kv_format == NVTE_QKV_Format::NVTE_THD) { - b = input_cu_seqlens_kv->data.shape[0] - 1; - } + NVTE_API_CALL(nvte_fused_attn_fwd); + transformer_engine::fused_attn::FusedAttnFwdParams p{}; + p.Q = Q; + p.K = K; + p.V = V; + p.Bias = Bias; + p.SoftmaxOffset = SoftmaxOffset; + p.S = S; + p.O = O; + p.Aux_CTX_Tensors = Aux_CTX_Tensors; + p.cu_seqlens_q = cu_seqlens_q; + p.cu_seqlens_kv = cu_seqlens_kv; + p.cu_seqlens_q_padded = cu_seqlens_q_padded; + p.cu_seqlens_kv_padded = cu_seqlens_kv_padded; + p.page_table_k = page_table_k; + p.page_table_v = page_table_v; + p.rng_state = rng_state; + p.is_training = is_training; + p.cuda_graph = cuda_graph; + p.return_max_logit = return_max_logit; + p.attn_mask_type = attn_mask_type; + p.bias_type = bias_type; + p.window_size_left = window_size_left; + p.window_size_right = window_size_right; + p.bottom_right_diagonal = bottom_right_diagonal; + p.softmax_type = softmax_type; + p.dropout = dropout; + p.attn_scale = attn_scale; + p.qkv_layout = qkv_layout; + p.o_format = o_format; + p.qkv_scale_inv_format = qkv_scale_inv_format; + p.max_seqlen_q = max_seqlen_q; + p.max_seqlen_kv = max_seqlen_kv; + p.workspace = workspace; + p.stream = stream; + nvte_fused_attn_fwd_v2(reinterpret_cast(&p)); +} - int64_t num_pages_k = 0; - int64_t num_pages_v = 0; - int64_t page_size_k = 0; - int64_t page_size_v = 0; - int64_t max_pages_per_seq_k = 0; - int64_t max_pages_per_seq_v = 0; - if (input_page_table_k->data.dptr != nullptr) { - max_pages_per_seq_k = input_page_table_k->data.shape[1]; - } - if (input_page_table_v->data.dptr != nullptr) { - max_pages_per_seq_v = input_page_table_v->data.shape[1]; - } - NVTE_QKV_Layout_Group layout_group = nvte_get_qkv_layout_group(qkv_layout); - if (layout_group == NVTE_QKV_Layout_Group::NVTE_Paged_KV_HD_HD_HD) { - NVTE_QKV_Format kv_format = nvte_get_kv_format(qkv_layout); - if (kv_format == NVTE_QKV_Format::NVTE_BSHD) { - num_pages_k = input_K->data.shape[0]; - page_size_k = input_K->data.shape[1]; - num_pages_v = input_V->data.shape[0]; - page_size_v = input_V->data.shape[1]; - } else if (kv_format == NVTE_QKV_Format::NVTE_SBHD) { - num_pages_k = input_K->data.shape[1]; - page_size_k = input_K->data.shape[0]; - num_pages_v = input_V->data.shape[1]; - page_size_v = input_V->data.shape[0]; - } - } +// Fused attention backward. Same shape as nvte_fused_attn_fwd_v2, whose comment sketches the path; +// this one asks the selector for backward support and probes the backward builders. +void nvte_fused_attn_bwd_v2(NVTEFusedAttnBwdParams params) { + NVTE_API_CALL(nvte_fused_attn_bwd_v2); + using namespace transformer_engine; + using namespace transformer_engine::fused_attn; + const FusedAttnBwdParams &p = *get_fused_attn_bwd_params(params); + const Tensor *input_cu_seqlens_q = convertNVTETensorCheck(p.cu_seqlens_q); + const Tensor *input_cu_seqlens_kv = convertNVTETensorCheck(p.cu_seqlens_kv); + const Tensor *input_cu_seqlens_q_padded = convertNVTETensorCheck(p.cu_seqlens_q_padded); + const Tensor *input_cu_seqlens_kv_padded = convertNVTETensorCheck(p.cu_seqlens_kv_padded); + const Tensor *input_Q = convertNVTETensorCheck(p.Q); + const Tensor *input_K = convertNVTETensorCheck(p.K); + const Tensor *input_V = convertNVTETensorCheck(p.V); + const Tensor *input_O = convertNVTETensorCheck(p.O); + const Tensor *input_dO = convertNVTETensorCheck(p.dO); + const Tensor *input_S = convertNVTETensorCheck(p.S); + Tensor *input_output_dP = convertNVTETensorCheck(p.dP); + Tensor *output_dQ = convertNVTETensorCheck(p.dQ); + Tensor *output_dK = convertNVTETensorCheck(p.dK); + Tensor *output_dV = convertNVTETensorCheck(p.dV); + Tensor *output_dBias = convertNVTETensorCheck(p.dBias); + Tensor *output_dSoftmaxOffset = convertNVTETensorCheck(p.dSoftmaxOffset); + Tensor *wkspace = convertNVTETensor(p.workspace); auto handle = cudnnExecutionPlanManager::Instance().GetHandle(); - const NVTEDType Q_type = static_cast(input_Q->data.dtype); - const NVTEDType KV_type = static_cast(input_K->data.dtype); - - NVTE_Fused_Attn_Backend fused_attention_backend = nvte_get_fused_attn_backend( - is_training, Q_type, KV_type, qkv_layout, bias_type, attn_mask_type, softmax_type, dropout, - h_q, h_kv, max_seqlen_q, max_seqlen_kv, d_qk, d_v, window_size_left, window_size_right, - return_max_logit, cuda_graph, false); + FusedAttnConfig cfg = p.make_config(); + // Derived here, not by the query below: the query works on its own copy, and it is this config + // that goes on to the backend and must arrive with its derived fields filled in. + cfg.derive(); + const char *fused_attn_reject_reason = nullptr; + NVTE_Fused_Attn_Backend fused_attention_backend = nvte_get_fused_attn_backend_v2( + reinterpret_cast(&cfg), &fused_attn_reject_reason); if (fused_attention_backend == NVTE_Fused_Attn_Backend::NVTE_F16_arbitrary_seqlen) { - fused_attn_arbitrary_seqlen_fwd( - b, h_q, h_kv, max_seqlen_q, max_seqlen_kv, d_qk, d_v, t_q, t_kv, num_pages_k, num_pages_v, - page_size_k, page_size_v, max_pages_per_seq_k, max_pages_per_seq_v, is_training, - return_max_logit, attn_scale, dropout, qkv_layout, o_format, bias_type, attn_mask_type, - softmax_type, window_size_left, window_size_right, bottom_right_diagonal, input_Q, input_K, - input_V, input_Bias, input_SoftmaxOffset, output_O, Aux_CTX_Tensors, input_cu_seqlens_q, - input_cu_seqlens_kv, input_cu_seqlens_q_padded, input_cu_seqlens_kv_padded, - input_page_table_k, input_page_table_v, input_rng_state, wkspace, stream, handle); + size_t i = 0; + Tensor *output_S = convertNVTETensorCheck(p.Aux_CTX_Tensors->tensors[i++]); + Tensor *input_rng_state = convertNVTETensorCheck(p.Aux_CTX_Tensors->tensors[i++]); + Tensor *input_Bias = nullptr, *input_SoftmaxOffset = nullptr; + if ((p.bias_type != NVTE_NO_BIAS) && (p.bias_type != NVTE_ALIBI)) { + input_Bias = convertNVTETensorCheck(p.Aux_CTX_Tensors->tensors[i++]); + } + if (p.softmax_type != NVTE_VANILLA_SOFTMAX) { + input_SoftmaxOffset = convertNVTETensorCheck(p.Aux_CTX_Tensors->tensors[i++]); + } + fused_attn_arbitrary_seqlen_bwd( + cfg, input_Q, input_K, input_V, input_O, input_dO, input_Bias, input_SoftmaxOffset, + output_S, output_dQ, output_dK, output_dV, output_dBias, output_dSoftmaxOffset, + input_cu_seqlens_q, input_cu_seqlens_kv, input_cu_seqlens_q_padded, + input_cu_seqlens_kv_padded, input_rng_state, wkspace, p.stream, handle); } else if (fused_attention_backend == NVTE_Fused_Attn_Backend::NVTE_FP8) { - fused_attn_fp8_fwd(b, h_q, h_kv, max_seqlen_q, max_seqlen_kv, d_qk, d_v, t_q, t_kv, is_training, - attn_scale, dropout, qkv_layout, o_format, qkv_scale_inv_format, bias_type, - attn_mask_type, softmax_type, window_size_left, window_size_right, - bottom_right_diagonal, input_Q, input_K, input_V, input_SoftmaxOffset, - input_output_S, output_O, Aux_CTX_Tensors, input_cu_seqlens_q, - input_cu_seqlens_kv, input_cu_seqlens_q_padded, input_cu_seqlens_kv_padded, - input_rng_state, wkspace, stream, handle); + size_t i = 0; + const Tensor *input_M = convertNVTETensorCheck(p.Aux_CTX_Tensors->tensors[i++]); + const Tensor *input_rng_state = convertNVTETensorCheck(p.Aux_CTX_Tensors->tensors[i++]); + const Tensor *input_SoftmaxOffset = nullptr; + if (p.softmax_type != NVTE_VANILLA_SOFTMAX) { + input_SoftmaxOffset = convertNVTETensorCheck(p.Aux_CTX_Tensors->tensors[i++]); + } + const Tensor *input_dO_f16 = nullptr; + if (input_dO->scaling_mode == NVTE_MXFP8_1D_SCALING) { + input_dO_f16 = convertNVTETensorCheck(p.Aux_CTX_Tensors->tensors[i++]); + } + fused_attn_fp8_bwd(cfg, input_Q, input_K, input_V, input_O, input_dO, input_dO_f16, input_M, + input_S, input_SoftmaxOffset, input_output_dP, output_dQ, output_dK, + output_dV, output_dSoftmaxOffset, input_cu_seqlens_q, input_cu_seqlens_kv, + input_cu_seqlens_q_padded, input_cu_seqlens_kv_padded, input_rng_state, + wkspace, p.stream, handle); } else { - NVTE_ERROR("Invalid combination of data type and sequence length for fused attention. \n"); + NVTE_ERROR("Fused attention is not supported for this configuration: ", + fused_attn_reject_reason); } } + // NVTE fused attention BWD with separate Q, K and V void nvte_fused_attn_bwd(const NVTETensor Q, const NVTETensor K, const NVTETensor V, const NVTETensor O, const NVTETensor dO, const NVTETensor S, NVTETensor dP, @@ -681,94 +660,46 @@ void nvte_fused_attn_bwd(const NVTETensor Q, const NVTETensor K, const NVTETenso NVTE_Softmax_Type softmax_type, int64_t window_size_left, int64_t window_size_right, bool bottom_right_diagonal, bool deterministic, bool cuda_graph, NVTETensor workspace, cudaStream_t stream) { - NVTE_API_CALL(nvte_flash_attn_bwd); - using namespace transformer_engine; - const Tensor *input_cu_seqlens_q = convertNVTETensorCheck(cu_seqlens_q); - const Tensor *input_cu_seqlens_kv = convertNVTETensorCheck(cu_seqlens_kv); - const Tensor *input_cu_seqlens_q_padded = convertNVTETensorCheck(cu_seqlens_q_padded); - const Tensor *input_cu_seqlens_kv_padded = convertNVTETensorCheck(cu_seqlens_kv_padded); - const Tensor *input_Q = convertNVTETensorCheck(Q); - const Tensor *input_K = convertNVTETensorCheck(K); - const Tensor *input_V = convertNVTETensorCheck(V); - const Tensor *input_O = convertNVTETensorCheck(O); - const Tensor *input_dO = convertNVTETensorCheck(dO); - const Tensor *input_S = convertNVTETensorCheck(S); - Tensor *input_output_dP = convertNVTETensorCheck(dP); - Tensor *output_dQ = convertNVTETensorCheck(dQ); - Tensor *output_dK = convertNVTETensorCheck(dK); - Tensor *output_dV = convertNVTETensorCheck(dV); - Tensor *output_dBias = convertNVTETensorCheck(dBias); - Tensor *output_dSoftmaxOffset = convertNVTETensorCheck(dSoftmaxOffset); - Tensor *wkspace = convertNVTETensor(workspace); - - NVTE_QKV_Format q_format = nvte_get_q_format(qkv_layout); - NVTE_QKV_Format kv_format = nvte_get_kv_format(qkv_layout); - auto *q_dims = input_Q->data.shape.data(); - auto *k_dims = input_K->data.shape.data(); - auto *v_dims = input_V->data.shape.data(); - AttentionShape q_shape(q_format, q_dims); - AttentionShape k_shape(kv_format, k_dims); - AttentionShape v_shape(kv_format, v_dims); - size_t b = q_shape.b(), h_q = q_shape.h(), d_qk = q_shape.d(), t_q = q_shape.t(); - size_t h_kv = k_shape.h(), t_kv = k_shape.t(), d_v = v_shape.d(); - if (q_format == NVTE_QKV_Format::NVTE_THD) { - b = input_cu_seqlens_q->data.shape[0] - 1; - } else if (kv_format == NVTE_QKV_Format::NVTE_THD) { - b = input_cu_seqlens_kv->data.shape[0] - 1; - } - - auto handle = cudnnExecutionPlanManager::Instance().GetHandle(); - const NVTEDType Q_type = static_cast(input_Q->data.dtype); - const NVTEDType KV_type = static_cast(input_K->data.dtype); - - NVTE_Fused_Attn_Backend fused_attention_backend = nvte_get_fused_attn_backend( - true, Q_type, KV_type, qkv_layout, bias_type, attn_mask_type, softmax_type, dropout, h_q, - h_kv, max_seqlen_q, max_seqlen_kv, d_qk, d_v, window_size_left, window_size_right, false, - cuda_graph, deterministic); - - if (fused_attention_backend == NVTE_Fused_Attn_Backend::NVTE_F16_arbitrary_seqlen) { - size_t i = 0; - Tensor *output_S = convertNVTETensorCheck(Aux_CTX_Tensors->tensors[i++]); - Tensor *input_rng_state = convertNVTETensorCheck(Aux_CTX_Tensors->tensors[i++]); - Tensor *input_Bias, *input_SoftmaxOffset; - if ((bias_type != NVTE_NO_BIAS) && (bias_type != NVTE_ALIBI)) { - input_Bias = convertNVTETensorCheck(Aux_CTX_Tensors->tensors[i++]); - } - if (softmax_type != NVTE_VANILLA_SOFTMAX) { - input_SoftmaxOffset = convertNVTETensorCheck(Aux_CTX_Tensors->tensors[i++]); - } - fused_attn_arbitrary_seqlen_bwd( - b, h_q, h_kv, max_seqlen_q, max_seqlen_kv, d_qk, d_v, t_q, t_kv, attn_scale, dropout, - qkv_layout, o_format, do_format, dqkv_layout, bias_type, attn_mask_type, softmax_type, - window_size_left, window_size_right, bottom_right_diagonal, deterministic, input_Q, input_K, - input_V, input_O, input_dO, input_Bias, input_SoftmaxOffset, output_S, output_dQ, output_dK, - output_dV, output_dBias, output_dSoftmaxOffset, input_cu_seqlens_q, input_cu_seqlens_kv, - input_cu_seqlens_q_padded, input_cu_seqlens_kv_padded, input_rng_state, wkspace, stream, - handle); - } else if (fused_attention_backend == NVTE_Fused_Attn_Backend::NVTE_FP8) { - size_t i = 0; - const Tensor *input_M = convertNVTETensorCheck(Aux_CTX_Tensors->tensors[i++]); - const Tensor *input_rng_state = convertNVTETensorCheck(Aux_CTX_Tensors->tensors[i++]); - const Tensor *input_SoftmaxOffset = nullptr; - if (softmax_type != NVTE_VANILLA_SOFTMAX) { - input_SoftmaxOffset = convertNVTETensorCheck(Aux_CTX_Tensors->tensors[i++]); - } - const Tensor *input_dO_f16 = nullptr; - if (input_dO->scaling_mode == NVTE_MXFP8_1D_SCALING) { - input_dO_f16 = convertNVTETensorCheck(Aux_CTX_Tensors->tensors[i++]); - } - fused_attn_fp8_bwd(b, h_q, h_kv, max_seqlen_q, max_seqlen_kv, d_qk, d_v, t_q, t_kv, attn_scale, - dropout, qkv_layout, o_format, do_format, dqkv_layout, qkv_scale_inv_format, - do_scale_inv_format, bias_type, attn_mask_type, softmax_type, - window_size_left, window_size_right, bottom_right_diagonal, deterministic, - input_Q, input_K, input_V, input_O, input_dO, input_dO_f16, input_M, input_S, - input_SoftmaxOffset, input_output_dP, output_dQ, output_dK, output_dV, - output_dSoftmaxOffset, input_cu_seqlens_q, input_cu_seqlens_kv, - input_cu_seqlens_q_padded, input_cu_seqlens_kv_padded, input_rng_state, - wkspace, stream, handle); - } else { - NVTE_ERROR("Invalid combination of data type and sequence length for fused attention. \n"); - } + NVTE_API_CALL(nvte_fused_attn_bwd); + transformer_engine::fused_attn::FusedAttnBwdParams p{}; + p.Q = Q; + p.K = K; + p.V = V; + p.O = O; + p.dO = dO; + p.S = S; + p.dP = dP; + p.Aux_CTX_Tensors = Aux_CTX_Tensors; + p.dQ = dQ; + p.dK = dK; + p.dV = dV; + p.dBias = dBias; + p.dSoftmaxOffset = dSoftmaxOffset; + p.cu_seqlens_q = cu_seqlens_q; + p.cu_seqlens_kv = cu_seqlens_kv; + p.cu_seqlens_q_padded = cu_seqlens_q_padded; + p.cu_seqlens_kv_padded = cu_seqlens_kv_padded; + p.deterministic = deterministic; + p.cuda_graph = cuda_graph; + p.attn_mask_type = attn_mask_type; + p.bias_type = bias_type; + p.window_size_left = window_size_left; + p.window_size_right = window_size_right; + p.bottom_right_diagonal = bottom_right_diagonal; + p.softmax_type = softmax_type; + p.dropout = dropout; + p.attn_scale = attn_scale; + p.qkv_layout = qkv_layout; + p.o_format = o_format; + p.do_format = do_format; + p.dqkv_layout = dqkv_layout; + p.qkv_scale_inv_format = qkv_scale_inv_format; + p.do_scale_inv_format = do_scale_inv_format; + p.max_seqlen_q = max_seqlen_q; + p.max_seqlen_kv = max_seqlen_kv; + p.workspace = workspace; + p.stream = stream; + nvte_fused_attn_bwd_v2(reinterpret_cast(&p)); } uint32_t nvte_get_runtime_num_segments(NVTETensor cu_seqlen, NVTETensor workspace, size_t len, diff --git a/transformer_engine/common/fused_attn/fused_attn_f16_arbitrary_seqlen.cu b/transformer_engine/common/fused_attn/fused_attn_f16_arbitrary_seqlen.cu index bf34758a35..fcbde95549 100644 --- a/transformer_engine/common/fused_attn/fused_attn_f16_arbitrary_seqlen.cu +++ b/transformer_engine/common/fused_attn/fused_attn_f16_arbitrary_seqlen.cu @@ -9,7 +9,6 @@ #include #include -#include #include #include "../common.h" @@ -17,486 +16,380 @@ #include "../util/cuda_runtime.h" #include "../util/system.h" #include "fused_attn_f16_arbitrary_seqlen.h" +#include "graph_cache.h" +#include "graph_cache_debug.h" #include "utils.h" -#define Q_ID 1 -#define K_ID 2 -#define V_ID 3 -#define O_ID 4 -#define S_ID 5 -#define B_ID 6 -#define D_CONST_ID 7 -#define S_CONST_ID 8 -#define Q_SEQLEN_ID 9 -#define K_SEQLEN_ID 10 -#define dQ_ID 11 -#define dK_ID 12 -#define dV_ID 13 -#define dO_ID 14 -#define MASK_VAL_ID 15 -#define dS_ID 16 -#define D_SEED_ID 17 -#define D_OFFSET_ID 18 -#define S_STATS_ID 19 -#define S_SUM_ID 20 -#define SCALE_PROB 21 -#define K_TRANSPOSE_ID 22 -#define dQ_ACCUM_ID 23 - -#define VIRTUAL_ID 30 - namespace transformer_engine { namespace fused_attn { -void fused_attn_arbitrary_seqlen_fwd_impl( - int64_t b, int64_t h, int64_t hg, int64_t s_q, int64_t s_kv, int64_t d_qk, int64_t d_v, - int64_t max_b, int64_t max_t_q, int64_t max_t_kv, int64_t num_pages_k, int64_t num_pages_v, - int64_t page_size_k, int64_t page_size_v, int64_t max_pages_per_seq_k, - int64_t max_pages_per_seq_v, int64_t bias_b, int64_t bias_h, int64_t bias_sq, int64_t bias_skv, - bool is_training, bool return_max_logit, float scaling_factor, float dropout_probability, - NVTE_QKV_Layout qkv_layout, NVTE_QKV_Format o_format, NVTE_Bias_Type bias_type, - NVTE_Mask_Type mask_type, NVTE_Softmax_Type softmax_type, int64_t window_size_left, - int64_t window_size_right, bool bottom_right_diagonal, void *devPtrQ, void *devPtrK, - void *devPtrV, void *devPtrBias, void *devPtrSoftmaxOffset, void *devPtrS1, void *devPtrS2, - void *devPtrO, void *devPtrDropoutSeed, void *devPtrDropoutOffset, void *devPtrCuSeqlensQ, - void *devPtrCuSeqlensKV, void *devPtrPageTableK, void *devPtrPageTableV, - void *devPtrSeqOffsetsQ, void *devPtrSeqOffsetsKV, cudnn_frontend::DataType_t tensorType, - void *workspace, size_t *workspace_size, cudaStream_t stream, cudnnHandle_t handle) { - using namespace transformer_engine; - bool is_bias = (bias_type == NVTE_Bias_Type::NVTE_POST_SCALE_BIAS); - bool is_alibi = (bias_type == NVTE_Bias_Type::NVTE_ALIBI); - bool is_causal = ((mask_type == NVTE_Mask_Type::NVTE_CAUSAL_MASK) || - (mask_type == NVTE_Mask_Type::NVTE_PADDING_CAUSAL_MASK)); - bool is_bottom_right = ((mask_type == NVTE_Mask_Type::NVTE_CAUSAL_BOTTOM_RIGHT_MASK) || - (mask_type == NVTE_Mask_Type::NVTE_PADDING_CAUSAL_BOTTOM_RIGHT_MASK)); - bool is_padding = ((mask_type == NVTE_Mask_Type::NVTE_PADDING_MASK) || - (mask_type == NVTE_Mask_Type::NVTE_PADDING_CAUSAL_MASK) || - (mask_type == NVTE_Mask_Type::NVTE_PADDING_CAUSAL_BOTTOM_RIGHT_MASK)); - if (is_bottom_right && s_q == s_kv && !is_padding) { - is_causal = true; - is_bottom_right = false; - bottom_right_diagonal = false; - } - bool is_softmax_offset = (softmax_type != NVTE_Softmax_Type::NVTE_VANILLA_SOFTMAX); - bool is_dropout = (is_training && dropout_probability != 0.0f); - NVTE_QKV_Format q_format = nvte_get_q_format(qkv_layout); - NVTE_QKV_Format kv_format = nvte_get_kv_format(qkv_layout); - bool is_ragged_q = (q_format == NVTE_QKV_Format::NVTE_THD); - bool is_ragged_kv = (kv_format == NVTE_QKV_Format::NVTE_THD); +namespace fe = cudnn_frontend; + +using F16FwdGraphAndTensors = + std::tuple, + std::shared_ptr, // Q + std::shared_ptr, // K + std::shared_ptr, // V + std::shared_ptr, // attn_scale + std::shared_ptr, // O + std::shared_ptr, // S1 + std::shared_ptr, // S2 + std::shared_ptr, // bias + std::shared_ptr, // softmax_offset + std::shared_ptr, // seq_q / cu_seq_len_q + std::shared_ptr, // seq_kv / cu_seq_len_kv + std::shared_ptr, // page_table_k + std::shared_ptr, // page_table_v + std::shared_ptr, // offset_q + std::shared_ptr, // offset_k + std::shared_ptr, // offset_v + std::shared_ptr, // offset_o + std::shared_ptr, // offset_stats + std::shared_ptr, // dropout_seed + std::shared_ptr>; // dropout_offset + +static F16FwdGraphAndTensors create_graph_f16_fwd(const FusedAttnConfig &cfg) { + const int64_t b = static_cast(cfg.graph_batch_size_fwd); + const int64_t s_q = static_cast(cfg.graph_max_seqlen_q); + const int64_t s_kv = static_cast(cfg.graph_max_seqlen_kv); + const cudnn_frontend::DataType_t tensorType = + get_cudnn_fe_dtype(static_cast(cfg.qkv_dtype)); + const int64_t h = static_cast(cfg.num_attn_heads); + const int64_t hg = static_cast(cfg.num_gqa_groups); + const int64_t d_qk = static_cast(cfg.head_dim_qk); + const int64_t d_v = static_cast(cfg.head_dim_v); + const int64_t num_pages_k = static_cast(cfg.num_pages_k); + const int64_t num_pages_v = static_cast(cfg.num_pages_v); + const int64_t page_size_k = static_cast(cfg.page_size_k); + const int64_t page_size_v = static_cast(cfg.page_size_v); + const int64_t max_pages_per_seq_k = static_cast(cfg.max_pages_per_seq_k); + const int64_t max_pages_per_seq_v = static_cast(cfg.max_pages_per_seq_v); + const int64_t bias_b = static_cast(cfg.bias_batch_size); + const int64_t bias_h = static_cast(cfg.bias_num_heads); + const int64_t bias_sq = static_cast(cfg.bias_seqlen_q); + const int64_t bias_skv = static_cast(cfg.bias_seqlen_kv); + const int64_t window_size_left = cfg.window_size_left; + const int64_t window_size_right = cfg.window_size_right; + const bool return_max_logit = cfg.return_max_logit; + const float dropout_probability = cfg.dropout; + const NVTE_QKV_Layout qkv_layout = cfg.qkv_layout; + const bool bottom_right_diagonal = cfg.bottom_right_diagonal; + const bool is_bias = cfg.is_bias; + const bool is_alibi = cfg.is_alibi; + const bool is_causal = cfg.is_causal; + const bool is_causal_bottom_right = cfg.is_causal_bottom_right; + const bool is_padding = cfg.is_padding; + const bool is_paged_kv = cfg.is_paged_kv; + const bool is_softmax_offset = cfg.is_softmax_offset; + const bool is_dropout = cfg.is_dropout; + const bool is_ragged_q = cfg.is_ragged_q; + const bool is_ragged_kv = cfg.is_ragged_kv; + const bool use_cu_seqlens_directly = cfg.uses_cu_seqlens_directly; const auto cudnn_runtime_version = cudnnGetVersion(); - const int device_id = cuda::current_device(); - const int sm_arch_ = cuda::sm_arch(device_id); - bool use_ragged_stats = is_ragged_q && cudnn_runtime_version >= 90600 && sm_arch_ != 120; - - NVTE_QKV_Layout_Group layout_group = nvte_get_qkv_layout_group(qkv_layout); - bool is_paged_kv = (layout_group == NVTE_QKV_Layout_Group::NVTE_Paged_KV_HD_HD_HD); + const bool use_ragged_stats = cfg.uses_ragged_stats; + const DType ragged_offset_type = cfg.ragged_offset_type_fwd; + const RaggedOffsetMultipliers offset_mults = cfg.ragged_offset_mults; + const bool generate_stats = true; // Always return stats + + auto mha_graph = std::make_shared(); + mha_graph->set_io_data_type(tensorType) + .set_intermediate_data_type(fe::DataType_t::FLOAT) + .set_compute_data_type(fe::DataType_t::FLOAT); + + std::shared_ptr Q, K, V, attn_scale, softmax_offset; + std::shared_ptr bias, seq_q, seq_kv; + std::shared_ptr page_table_k, page_table_v; + std::shared_ptr offset_q, offset_k, offset_v, offset_o, + offset_stats; + std::shared_ptr dropout_seed, dropout_offset; + + std::vector q_stride(4); + std::vector k_stride(4); + std::vector v_stride(4); + generateMatrixStrides(b, h, s_q, s_kv, d_qk, q_stride.data(), qkv_layout, + NVTE_QKV_Matrix::NVTE_Q_Matrix); if (is_paged_kv) { - NVTE_CHECK(is_padding, "Paged attention requires padding mask!"); + generateMatrixStrides(num_pages_k, hg, page_size_k, page_size_v, d_qk, k_stride.data(), + qkv_layout, NVTE_QKV_Matrix::NVTE_K_Matrix); + generateMatrixStrides(num_pages_v, hg, page_size_k, page_size_v, d_v, v_stride.data(), + qkv_layout, NVTE_QKV_Matrix::NVTE_V_Matrix); + } else { + generateMatrixStrides(b, hg, s_q, s_kv, d_qk, k_stride.data(), qkv_layout, + NVTE_QKV_Matrix::NVTE_K_Matrix); + generateMatrixStrides(b, hg, s_q, s_kv, d_v, v_stride.data(), qkv_layout, + NVTE_QKV_Matrix::NVTE_V_Matrix); } - // Newer versions of cuDNN SDPA can accept sequence lengths directly as a cumulative - // tensor, and can accept ragged offsets in arbitrary units (such as tokens) instead - // of elements. Take advantage of this if possible to avoid 2 extra kernel calls. - const bool use_cu_seqlens_directly = - CUDNN_FRONTEND_VERSION >= 12500 && - // The frontend gates cu_seq_len support on min(compile-time, runtime) cuDNN - // version, so we'll do the same. - (CUDNN_VERSION >= 92400 && cudnn_runtime_version >= 92400) && - // This extra restriction is needed because cuDNN frontend doesn't yet allow - // the combination of dropout and stats generation for the fprop unified engine, - // so any such request would always get routed to the old composite SDPA engine - // (which doesn't support cu_seqlens). Remove this restriction when possible. - !is_dropout; - - // keep original batch size because cu_seqlens are created with [b+1] shape - int64_t actual_b = b; - if ((is_ragged_q || is_ragged_kv) && cudnn_runtime_version >= 90600) { - NVTE_CHECK(is_padding, "Ragged QKV input requires padding or padding_causal mask!"); - // On SM 120, cuDNN support check treats layouts with stride[0] > dim[1]*dim[2]*dim[3] - // as interleaved and rejects them. Use BHSD-like dimensions/strides with max_seqlen at plan build - // so the check passes; ragged offset still provides variable-length boundaries. - if (sm_arch_ != 120) { - // replace batch size and maximum sequence lengths with maximum token counts - // for query and key/value so the graph is static within each quantization bucket. - // When passing cu_seqlens* directly to cuDNN SDPA, keep the true batch size: - // cuDNN reads the user's [actual_b+1] cu_seqlens buffers, so a quantized batch - // would read out of bounds. - if (!use_cu_seqlens_directly) { - b = max_b; - } - s_q = is_ragged_q ? max_t_q : s_q; - s_kv = is_ragged_kv ? max_t_kv : s_kv; + Q = mha_graph->tensor( + fe::graph::Tensor_attributes().set_name("Q").set_dim({b, h, s_q, d_qk}).set_stride(q_stride)); + if (is_ragged_q) { + offset_q = mha_graph->tensor(fe::graph::Tensor_attributes() + .set_name("offset_q") + .set_dim({b + 1, 1, 1, 1}) + .set_stride({1, 1, 1, 1}) + .set_data_type(get_cudnn_fe_dtype(ragged_offset_type))); + Q->set_ragged_offset(offset_q); + if (use_cu_seqlens_directly) { + Q->set_ragged_offset_multiplier(offset_mults.q); } } + K = mha_graph->tensor(fe::graph::Tensor_attributes().set_name("K").set_stride(k_stride)); + V = mha_graph->tensor(fe::graph::Tensor_attributes().set_name("V").set_stride(v_stride)); + if (is_paged_kv) { + K->set_dim({num_pages_k, hg, page_size_k, d_qk}); + V->set_dim({num_pages_v, hg, page_size_v, d_v}); + } else if (is_ragged_kv) { + offset_k = mha_graph->tensor(fe::graph::Tensor_attributes() + .set_name("offset_k") + .set_dim({b + 1, 1, 1, 1}) + .set_stride({1, 1, 1, 1}) + .set_data_type(get_cudnn_fe_dtype(ragged_offset_type))); + offset_v = mha_graph->tensor(fe::graph::Tensor_attributes() + .set_name("offset_v") + .set_dim({b + 1, 1, 1, 1}) + .set_stride({1, 1, 1, 1}) + .set_data_type(get_cudnn_fe_dtype(ragged_offset_type))); + K->set_dim({b, hg, s_kv, d_qk}).set_ragged_offset(offset_k); + V->set_dim({b, hg, s_kv, d_v}).set_ragged_offset(offset_v); + if (use_cu_seqlens_directly) { + K->set_ragged_offset_multiplier(offset_mults.k); + V->set_ragged_offset_multiplier(offset_mults.v); + } + } else { + K->set_dim({b, hg, s_kv, d_qk}); + V->set_dim({b, hg, s_kv, d_v}); + } - const DType ragged_offset_type = - use_cu_seqlens_directly - ? DType::kInt32 // cu_seqlens* are given to us as int32; keep it that way. - : (cudnn_runtime_version >= 90500 ? DType::kInt64 : DType::kInt32); - - // Ragged offset multipliers (elements per token); shared with the legacy conversion - // kernel (cu_seqlens_padded_to_offsets) so the two paths cannot drift apart. - const RaggedOffsetMultipliers offset_mults(layout_group, h, hg, d_qk, d_v); + attn_scale = mha_graph->tensor(fe::graph::Tensor_attributes() + .set_name("attn_scale") + .set_dim({1, 1, 1, 1}) + .set_stride({1, 1, 1, 1}) + .set_is_pass_by_value(true) + .set_data_type(fe::DataType_t::FLOAT)); + + fe::graph::SDPA_attributes sdpa_options; + sdpa_options = fe::graph::SDPA_attributes() + .set_name("flash_attention") + .set_generate_stats(generate_stats) + .set_attn_scale(attn_scale); + + fe::DiagonalAlignment_t const &diagonal_alignment = bottom_right_diagonal + ? fe::DiagonalAlignment_t::BOTTOM_RIGHT + : fe::DiagonalAlignment_t::TOP_LEFT; + sdpa_options.set_diagonal_alignment(diagonal_alignment); + if (cudnn_runtime_version >= 90200 && window_size_left != -1) { + sdpa_options.set_diagonal_band_left_bound(window_size_left + 1); + } + if (cudnn_runtime_version >= 90600 && window_size_right != -1) { + sdpa_options.set_diagonal_band_right_bound(window_size_right); + } + if (is_causal || is_causal_bottom_right) { + sdpa_options.set_diagonal_band_right_bound(0); + } - bool generate_stats = true; // Always return stats - try { - FADescriptor_v1 descriptor{ - b, - h, - hg, - s_q, - s_kv, - d_qk, - d_v, - num_pages_k, - num_pages_v, - page_size_k, - page_size_v, - max_pages_per_seq_k, - max_pages_per_seq_v, - bias_b, - bias_h, - bias_sq, - bias_skv, - scaling_factor, - is_training, - dropout_probability, - qkv_layout, - o_format, - NVTE_QKV_Format_NOT_SET, - NVTE_QKV_Layout_NOT_SET, - NVTE_QKV_Format_NOT_SET, - NVTE_QKV_Format_NOT_SET, - bias_type, - mask_type, - softmax_type, - window_size_left, - window_size_right, - bottom_right_diagonal, - true, - tensorType, - cudnn_frontend::DataType_t::NOT_SET, - cudnn_frontend::DataType_t::NOT_SET, - cudnn_frontend::DataType_t::NOT_SET, - return_max_logit, - }; + sdpa_options.set_alibi_mask(is_alibi); - namespace fe = cudnn_frontend; - using graph_and_tensors = - std::tuple, - std::shared_ptr, // Q - std::shared_ptr, // K - std::shared_ptr, // V - std::shared_ptr, // attn_scale - std::shared_ptr, // O - std::shared_ptr, // S1 - std::shared_ptr, // S2 - std::shared_ptr, // bias - std::shared_ptr, // softmax_offset - std::shared_ptr, // seq_q / cu_seq_len_q - std::shared_ptr, // seq_kv / cu_seq_len_kv - std::shared_ptr, // page_table_k - std::shared_ptr, // page_table_v - std::shared_ptr, // offset_q - std::shared_ptr, // offset_k - std::shared_ptr, // offset_v - std::shared_ptr, // offset_o - std::shared_ptr, // offset_stats - std::shared_ptr, // dropout_seed - std::shared_ptr>; // dropout_offset - - using CacheType = std::map; - static thread_local CacheType sdpa_f16_fprop_cache; - - // Get plan from cache if cache is available, otherwise create one - auto get_graph = [&](CacheType &cache, const FADescriptor_v1 &descriptor) -> graph_and_tensors { - // if hit, return - auto it = cache.find(descriptor); - if (it != cache.end()) { - auto graph = it->second; - return graph; - } + if (is_bias) { + bias = mha_graph->tensor( + fe::graph::Tensor_attributes() + .set_name("bias") + .set_dim({bias_b, bias_h, bias_sq, bias_skv}) + .set_stride({bias_h * bias_sq * bias_skv, bias_sq * bias_skv, bias_skv, 1})); + sdpa_options.set_bias(bias); + } - // otherwise, build the op_graph and the plan. Then update cache - auto mha_graph = std::make_shared(); - mha_graph->set_io_data_type(tensorType) - .set_intermediate_data_type(fe::DataType_t::FLOAT) - .set_compute_data_type(fe::DataType_t::FLOAT); - - std::shared_ptr Q, K, V, attn_scale, softmax_offset; - std::shared_ptr bias, seq_q, seq_kv; - std::shared_ptr page_table_k, page_table_v; - std::shared_ptr offset_q, offset_k, offset_v, offset_o, - offset_stats; - std::shared_ptr dropout_seed, dropout_offset; - - std::vector q_stride(4); - std::vector k_stride(4); - std::vector v_stride(4); - generateMatrixStrides(b, h, s_q, s_kv, d_qk, q_stride.data(), qkv_layout, - NVTE_QKV_Matrix::NVTE_Q_Matrix); - if (is_paged_kv) { - generateMatrixStrides(num_pages_k, hg, page_size_k, page_size_v, d_qk, k_stride.data(), - qkv_layout, NVTE_QKV_Matrix::NVTE_K_Matrix); - generateMatrixStrides(num_pages_v, hg, page_size_k, page_size_v, d_v, v_stride.data(), - qkv_layout, NVTE_QKV_Matrix::NVTE_V_Matrix); - } else { - generateMatrixStrides(b, hg, s_q, s_kv, d_qk, k_stride.data(), qkv_layout, - NVTE_QKV_Matrix::NVTE_K_Matrix); - generateMatrixStrides(b, hg, s_q, s_kv, d_v, v_stride.data(), qkv_layout, - NVTE_QKV_Matrix::NVTE_V_Matrix); - } + if (is_padding) { + if (use_cu_seqlens_directly) { + // seq_q/seq_kv keep their tuple slots but hold (b+1)-shaped cu_seqlen tensors. + seq_q = mha_graph->tensor(fe::graph::Tensor_attributes() + .set_name("cu_seq_len_q") + .set_dim({b + 1, 1, 1, 1}) + .set_stride({1, 1, 1, 1}) + .set_data_type(fe::DataType_t::INT32)); + seq_kv = mha_graph->tensor(fe::graph::Tensor_attributes() + .set_name("cu_seq_len_kv") + .set_dim({b + 1, 1, 1, 1}) + .set_stride({1, 1, 1, 1}) + .set_data_type(fe::DataType_t::INT32)); + sdpa_options.set_padding_mask(is_padding).set_cu_seq_len_q(seq_q).set_cu_seq_len_kv(seq_kv); + // cu_seq_len (and the ragged offset multiplier) are unified-engine-only. + // Pin the implementation so an unsupported config fails with the unified + // engine's specific error instead of auto-selection's generic failure. + sdpa_options.set_implementation(fe::AttentionImplementation_t::UNIFIED); + } else { + seq_q = mha_graph->tensor(fe::graph::Tensor_attributes() + .set_name("seq_q") + .set_dim({b, 1, 1, 1}) + .set_stride({1, 1, 1, 1}) + .set_data_type(fe::DataType_t::INT32)); + seq_kv = mha_graph->tensor(fe::graph::Tensor_attributes() + .set_name("seq_kv") + .set_dim({b, 1, 1, 1}) + .set_stride({1, 1, 1, 1}) + .set_data_type(fe::DataType_t::INT32)); + sdpa_options.set_padding_mask(is_padding).set_seq_len_q(seq_q).set_seq_len_kv(seq_kv); + } + } - Q = mha_graph->tensor(fe::graph::Tensor_attributes() - .set_name("Q") - .set_dim({b, h, s_q, d_qk}) - .set_stride(q_stride)); - if (is_ragged_q) { - offset_q = mha_graph->tensor(fe::graph::Tensor_attributes() - .set_name("offset_q") - .set_dim({b + 1, 1, 1, 1}) - .set_stride({1, 1, 1, 1}) - .set_data_type(get_cudnn_fe_dtype(ragged_offset_type))); - Q->set_ragged_offset(offset_q); - if (use_cu_seqlens_directly) { - Q->set_ragged_offset_multiplier(offset_mults.q); - } - } - K = mha_graph->tensor(fe::graph::Tensor_attributes().set_name("K").set_stride(k_stride)); - V = mha_graph->tensor(fe::graph::Tensor_attributes().set_name("V").set_stride(v_stride)); - if (is_paged_kv) { - K->set_dim({num_pages_k, hg, page_size_k, d_qk}); - V->set_dim({num_pages_v, hg, page_size_v, d_v}); - } else if (is_ragged_kv) { - offset_k = mha_graph->tensor(fe::graph::Tensor_attributes() - .set_name("offset_k") - .set_dim({b + 1, 1, 1, 1}) - .set_stride({1, 1, 1, 1}) - .set_data_type(get_cudnn_fe_dtype(ragged_offset_type))); - offset_v = mha_graph->tensor(fe::graph::Tensor_attributes() - .set_name("offset_v") - .set_dim({b + 1, 1, 1, 1}) - .set_stride({1, 1, 1, 1}) - .set_data_type(get_cudnn_fe_dtype(ragged_offset_type))); - K->set_dim({b, hg, s_kv, d_qk}).set_ragged_offset(offset_k); - V->set_dim({b, hg, s_kv, d_v}).set_ragged_offset(offset_v); - if (use_cu_seqlens_directly) { - K->set_ragged_offset_multiplier(offset_mults.k); - V->set_ragged_offset_multiplier(offset_mults.v); - } - } else { - K->set_dim({b, hg, s_kv, d_qk}); - V->set_dim({b, hg, s_kv, d_v}); - } + if (is_paged_kv) { + page_table_k = + mha_graph->tensor(fe::graph::Tensor_attributes() + .set_name("page_table_k") + .set_dim({b, 1, max_pages_per_seq_k, 1}) + .set_stride({{max_pages_per_seq_k, max_pages_per_seq_v, 1, 1}}) + .set_data_type(fe::DataType_t::INT32)); + page_table_v = + mha_graph->tensor(fe::graph::Tensor_attributes() + .set_name("page_table_v") + .set_dim({b, 1, max_pages_per_seq_v, 1}) + .set_stride({{max_pages_per_seq_v, max_pages_per_seq_v, 1, 1}}) + .set_data_type(fe::DataType_t::INT32)); + sdpa_options.set_paged_attention_k_table(page_table_k); + sdpa_options.set_paged_attention_v_table(page_table_v); + sdpa_options.set_paged_attention_max_seq_len_kv(static_cast(s_kv)); + } - attn_scale = mha_graph->tensor(fe::graph::Tensor_attributes() - .set_name("attn_scale") + if (is_dropout) { + dropout_seed = mha_graph->tensor(fe::graph::Tensor_attributes() + .set_name("Seed") .set_dim({1, 1, 1, 1}) .set_stride({1, 1, 1, 1}) - .set_is_pass_by_value(true) - .set_data_type(fe::DataType_t::FLOAT)); - - fe::graph::SDPA_attributes sdpa_options; - sdpa_options = fe::graph::SDPA_attributes() - .set_name("flash_attention") - .set_generate_stats(generate_stats) - .set_causal_mask(is_causal) - .set_causal_mask_bottom_right(is_bottom_right) - .set_attn_scale(attn_scale); - - fe::DiagonalAlignment_t const &diagonal_alignment = - bottom_right_diagonal ? fe::DiagonalAlignment_t::BOTTOM_RIGHT - : fe::DiagonalAlignment_t::TOP_LEFT; - sdpa_options.set_diagonal_alignment(diagonal_alignment); - if (cudnn_runtime_version >= 90200 && window_size_left != -1) { - sdpa_options.set_diagonal_band_left_bound(window_size_left + 1); - } - if (cudnn_runtime_version >= 90600 && window_size_right != -1) { - sdpa_options.set_diagonal_band_right_bound(window_size_right); - } - - sdpa_options.set_alibi_mask(is_alibi); + .set_data_type(fe::DataType_t::INT64)); + dropout_offset = mha_graph->tensor(fe::graph::Tensor_attributes() + .set_name("Offset") + .set_dim({1, 1, 1, 1}) + .set_stride({1, 1, 1, 1}) + .set_data_type(fe::DataType_t::INT64)); + sdpa_options.set_dropout(dropout_probability, dropout_seed, dropout_offset); + } - if (is_bias) { - bias = mha_graph->tensor( - fe::graph::Tensor_attributes() - .set_name("bias") - .set_dim({bias_b, bias_h, bias_sq, bias_skv}) - .set_stride({bias_h * bias_sq * bias_skv, bias_sq * bias_skv, bias_skv, 1})); - sdpa_options.set_bias(bias); - } + if (is_softmax_offset) { + softmax_offset = mha_graph->tensor(fe::graph::Tensor_attributes() + .set_name("softmax_offset") + .set_dim({1, h, 1, 1}) + .set_stride({h, 1, 1, 1}) + .set_data_type(fe::DataType_t::FLOAT)); + sdpa_options.set_sink_token(softmax_offset); + } - if (is_padding) { - if (use_cu_seqlens_directly) { - // seq_q/seq_kv keep their tuple slots but hold (b+1)-shaped cu_seqlen tensors. - seq_q = mha_graph->tensor(fe::graph::Tensor_attributes() - .set_name("cu_seq_len_q") - .set_dim({b + 1, 1, 1, 1}) - .set_stride({1, 1, 1, 1}) - .set_data_type(fe::DataType_t::INT32)); - seq_kv = mha_graph->tensor(fe::graph::Tensor_attributes() - .set_name("cu_seq_len_kv") + std::shared_ptr Max; + if (use_ragged_stats) { + offset_stats = mha_graph->tensor(fe::graph::Tensor_attributes() + .set_name("offset_stats") .set_dim({b + 1, 1, 1, 1}) .set_stride({1, 1, 1, 1}) - .set_data_type(fe::DataType_t::INT32)); - sdpa_options.set_padding_mask(is_padding) - .set_cu_seq_len_q(seq_q) - .set_cu_seq_len_kv(seq_kv); - // cu_seq_len (and the ragged offset multiplier) are unified-engine-only. - // Pin the implementation so an unsupported config fails with the unified - // engine's specific error instead of auto-selection's generic failure. - sdpa_options.set_implementation(fe::AttentionImplementation_t::UNIFIED); - } else { - seq_q = mha_graph->tensor(fe::graph::Tensor_attributes() - .set_name("seq_q") - .set_dim({b, 1, 1, 1}) - .set_stride({1, 1, 1, 1}) - .set_data_type(fe::DataType_t::INT32)); - seq_kv = mha_graph->tensor(fe::graph::Tensor_attributes() - .set_name("seq_kv") - .set_dim({b, 1, 1, 1}) - .set_stride({1, 1, 1, 1}) - .set_data_type(fe::DataType_t::INT32)); - sdpa_options.set_padding_mask(is_padding).set_seq_len_q(seq_q).set_seq_len_kv(seq_kv); - } - } - - if (is_paged_kv) { - page_table_k = - mha_graph->tensor(fe::graph::Tensor_attributes() - .set_name("page_table_k") - .set_dim({b, 1, max_pages_per_seq_k, 1}) - .set_stride({{max_pages_per_seq_k, max_pages_per_seq_v, 1, 1}}) - .set_data_type(fe::DataType_t::INT32)); - page_table_v = - mha_graph->tensor(fe::graph::Tensor_attributes() - .set_name("page_table_v") - .set_dim({b, 1, max_pages_per_seq_v, 1}) - .set_stride({{max_pages_per_seq_v, max_pages_per_seq_v, 1, 1}}) - .set_data_type(fe::DataType_t::INT32)); - sdpa_options.set_paged_attention_k_table(page_table_k); - sdpa_options.set_paged_attention_v_table(page_table_v); - sdpa_options.set_paged_attention_max_seq_len_kv(static_cast(s_kv)); - } - - if (is_dropout) { - dropout_seed = mha_graph->tensor(fe::graph::Tensor_attributes() - .set_name("Seed") - .set_dim({1, 1, 1, 1}) - .set_stride({1, 1, 1, 1}) - .set_data_type(fe::DataType_t::INT64)); - dropout_offset = mha_graph->tensor(fe::graph::Tensor_attributes() - .set_name("Offset") - .set_dim({1, 1, 1, 1}) - .set_stride({1, 1, 1, 1}) - .set_data_type(fe::DataType_t::INT64)); - sdpa_options.set_dropout(dropout_probability, dropout_seed, dropout_offset); - } - - if (is_softmax_offset) { - softmax_offset = mha_graph->tensor(fe::graph::Tensor_attributes() - .set_name("softmax_offset") - .set_dim({1, h, 1, 1}) - .set_stride({h, 1, 1, 1}) - .set_data_type(fe::DataType_t::FLOAT)); - sdpa_options.set_sink_token(softmax_offset); + .set_data_type(get_cudnn_fe_dtype(ragged_offset_type))); + } + if (return_max_logit) { + Max = mha_graph->tensor(fe::graph::Tensor_attributes() + .set_name("Max") + .set_dim({b, h, s_q, 1}) + .set_data_type(fe::DataType_t::FLOAT)); + if (use_ragged_stats) { + Max->set_stride({h * s_q, 1, h, 1}).set_ragged_offset(offset_stats); + if (use_cu_seqlens_directly) { + Max->set_ragged_offset_multiplier(offset_mults.stats); } + } else { + Max->set_stride({h * s_q, s_q, 1, 1}); + } + sdpa_options.set_logit_max(Max); + } - std::shared_ptr Max; - if (use_ragged_stats) { - offset_stats = - mha_graph->tensor(fe::graph::Tensor_attributes() - .set_name("offset_stats") - .set_dim({b + 1, 1, 1, 1}) - .set_stride({1, 1, 1, 1}) - .set_data_type(get_cudnn_fe_dtype(ragged_offset_type))); - } - if (return_max_logit) { - Max = mha_graph->tensor(fe::graph::Tensor_attributes() - .set_name("Max") - .set_dim({b, h, s_q, 1}) - .set_data_type(fe::DataType_t::FLOAT)); - if (use_ragged_stats) { - Max->set_stride({h * s_q, 1, h, 1}).set_ragged_offset(offset_stats); - if (use_cu_seqlens_directly) { - Max->set_ragged_offset_multiplier(offset_mults.stats); - } - } else { - Max->set_stride({h * s_q, s_q, 1, 1}); - } - sdpa_options.set_logit_max(Max); - } + auto [O, Stats] = mha_graph->sdpa(Q, K, V, std::move(sdpa_options)); + + std::vector o_stride(4); + generateMatrixStrides(b, h, s_q, s_kv, d_v, o_stride.data(), qkv_layout, + NVTE_QKV_Matrix::NVTE_O_Matrix); + O->set_output(true).set_dim({b, h, s_q, d_v}).set_stride(o_stride); + if (is_ragged_q) { + offset_o = mha_graph->tensor(fe::graph::Tensor_attributes() + .set_name("offset_o") + .set_dim({b + 1, 1, 1, 1}) + .set_stride({1, 1, 1, 1}) + .set_data_type(get_cudnn_fe_dtype(ragged_offset_type))); + O->set_ragged_offset(offset_o); + if (use_cu_seqlens_directly) { + O->set_ragged_offset_multiplier(offset_mults.o); + } + } - auto [O, Stats] = mha_graph->sdpa(Q, K, V, std::move(sdpa_options)); + Stats->set_output(true).set_data_type(fe::DataType_t::FLOAT).set_dim({b, h, s_q, 1}); + if (use_ragged_stats) { + Stats->set_stride({h * s_q, 1, h, 1}).set_ragged_offset(offset_stats); + if (use_cu_seqlens_directly) { + Stats->set_ragged_offset_multiplier(offset_mults.stats); + } + } else { + Stats->set_stride({h * s_q, s_q, 1, 1}); + } - std::vector o_stride(4); - generateMatrixStrides(b, h, s_q, s_kv, d_v, o_stride.data(), qkv_layout, - NVTE_QKV_Matrix::NVTE_O_Matrix); - O->set_output(true).set_dim({b, h, s_q, d_v}).set_stride(o_stride); - if (is_ragged_q) { - offset_o = mha_graph->tensor(fe::graph::Tensor_attributes() - .set_name("offset_o") - .set_dim({b + 1, 1, 1, 1}) - .set_stride({1, 1, 1, 1}) - .set_data_type(get_cudnn_fe_dtype(ragged_offset_type))); - O->set_ragged_offset(offset_o); - if (use_cu_seqlens_directly) { - O->set_ragged_offset_multiplier(offset_mults.o); - } - } + std::tuple, // Q + std::shared_ptr, // K + std::shared_ptr, // V + std::shared_ptr, // attn_scale + std::shared_ptr> // O + key_tensors_tuple = std::make_tuple(Q, K, V, attn_scale, O); + auto Stats_tuple = + return_max_logit ? std::make_tuple(Stats, Max) : std::make_tuple(Stats, nullptr); + auto bias_tuple = is_bias ? std::make_tuple(bias) : std::make_tuple(nullptr); + auto softmax_offset_tuple = + is_softmax_offset ? std::make_tuple(softmax_offset) : std::make_tuple(nullptr); + auto padding_tuple = + is_padding ? std::make_tuple(seq_q, seq_kv) : std::make_tuple(nullptr, nullptr); + auto page_table_tuple = + is_paged_kv ? std::make_tuple(page_table_k, page_table_v) : std::make_tuple(nullptr, nullptr); + auto offset_qo_tuple = + is_ragged_q ? std::make_tuple(offset_q, offset_o) : std::make_tuple(nullptr, nullptr); + auto offset_kv_tuple = + is_ragged_kv ? std::make_tuple(offset_k, offset_v) : std::make_tuple(nullptr, nullptr); + auto offset_s_tuple = use_ragged_stats ? std::make_tuple(offset_stats) : std::make_tuple(nullptr); + auto dropout_tuple = is_dropout ? std::make_tuple(dropout_seed, dropout_offset) + : std::make_tuple(nullptr, nullptr); + + return std::tuple_cat(std::make_tuple(mha_graph), key_tensors_tuple, Stats_tuple, bias_tuple, + softmax_offset_tuple, padding_tuple, page_table_tuple, offset_qo_tuple, + offset_kv_tuple, offset_s_tuple, dropout_tuple); +} - Stats->set_output(true).set_data_type(fe::DataType_t::FLOAT).set_dim({b, h, s_q, 1}); - if (use_ragged_stats) { - Stats->set_stride({h * s_q, 1, h, 1}).set_ragged_offset(offset_stats); - if (use_cu_seqlens_directly) { - Stats->set_ragged_offset_multiplier(offset_mults.stats); - } - } else { - Stats->set_stride({h * s_q, s_q, 1, 1}); - } +void fused_attn_arbitrary_seqlen_fwd_impl( + const FusedAttnConfig &cfg, void *devPtrQ, void *devPtrK, void *devPtrV, void *devPtrBias, + void *devPtrSoftmaxOffset, void *devPtrS1, void *devPtrS2, void *devPtrO, + void *devPtrDropoutSeed, void *devPtrDropoutOffset, void *devPtrCuSeqlensQ, + void *devPtrCuSeqlensKV, void *devPtrPageTableK, void *devPtrPageTableV, + void *devPtrSeqOffsetsQ, void *devPtrSeqOffsetsKV, void *workspace, size_t *workspace_size, + cudaStream_t stream, cudnnHandle_t handle) { + using namespace transformer_engine; - std::tuple, // Q - std::shared_ptr, // K - std::shared_ptr, // V - std::shared_ptr, // attn_scale - std::shared_ptr> // O - key_tensors_tuple = std::make_tuple(Q, K, V, attn_scale, O); - auto Stats_tuple = - return_max_logit ? std::make_tuple(Stats, Max) : std::make_tuple(Stats, nullptr); - auto bias_tuple = is_bias ? std::make_tuple(bias) : std::make_tuple(nullptr); - auto softmax_offset_tuple = - is_softmax_offset ? std::make_tuple(softmax_offset) : std::make_tuple(nullptr); - auto padding_tuple = - is_padding ? std::make_tuple(seq_q, seq_kv) : std::make_tuple(nullptr, nullptr); - auto page_table_tuple = is_paged_kv ? std::make_tuple(page_table_k, page_table_v) - : std::make_tuple(nullptr, nullptr); - auto offset_qo_tuple = - is_ragged_q ? std::make_tuple(offset_q, offset_o) : std::make_tuple(nullptr, nullptr); - auto offset_kv_tuple = - is_ragged_kv ? std::make_tuple(offset_k, offset_v) : std::make_tuple(nullptr, nullptr); - auto offset_s_tuple = - use_ragged_stats ? std::make_tuple(offset_stats) : std::make_tuple(nullptr); - auto dropout_tuple = is_dropout ? std::make_tuple(dropout_seed, dropout_offset) - : std::make_tuple(nullptr, nullptr); - - NVTE_CHECK_CUDNN_FE(mha_graph->validate()); - NVTE_CHECK_CUDNN_FE(mha_graph->build_operation_graph(handle)); - NVTE_CHECK_CUDNN_FE(mha_graph->create_execution_plans({fe::HeurMode_t::A})); - NVTE_CHECK_CUDNN_FE(mha_graph->check_support(handle)); - NVTE_CHECK_CUDNN_FE(mha_graph->build_plans(handle)); - - auto return_tuple = - std::tuple_cat(std::make_tuple(mha_graph), key_tensors_tuple, Stats_tuple, bias_tuple, - softmax_offset_tuple, padding_tuple, page_table_tuple, offset_qo_tuple, - offset_kv_tuple, offset_s_tuple, dropout_tuple); - cache.insert({descriptor, return_tuple}); - - return return_tuple; - }; + cfg.check_derived(); + const int64_t b = static_cast(cfg.graph_batch_size_fwd); + const DType ragged_offset_type = cfg.ragged_offset_type_fwd; + const int64_t actual_b = static_cast(cfg.batch_size); + const bool use_ragged_stats = cfg.uses_ragged_stats; + const RaggedOffsetMultipliers offset_mults = cfg.ragged_offset_mults; + + const bool return_max_logit = cfg.return_max_logit; + float scaling_factor = cfg.attn_scale; + const bool is_bias = cfg.is_bias; + const bool is_padding = cfg.is_padding; + const bool is_softmax_offset = cfg.is_softmax_offset; + const bool is_dropout = cfg.is_dropout; + const bool is_ragged_q = cfg.is_ragged_q; + const bool is_ragged_kv = cfg.is_ragged_kv; + const bool is_paged_kv = cfg.is_paged_kv; + // Newer versions of cuDNN SDPA can accept sequence lengths directly as a cumulative + // tensor, and can accept ragged offsets in arbitrary units (such as tokens) instead + // of elements. Take advantage of this if possible to avoid 2 extra kernel calls. + const bool use_cu_seqlens_directly = cfg.uses_cu_seqlens_directly; + try { + auto cache_entry = get_graph(cfg, handle); auto [mha_graph, Q, K, V, attn_scale, O, S1, S2, bias, softmax_offset, seq_q, seq_kv, page_table_k, page_table_v, offset_q, offset_o, offset_k, offset_v, offset_stats, - dropout_seed, dropout_offset] = get_graph(sdpa_f16_fprop_cache, descriptor); + dropout_seed, dropout_offset] = cache_entry->graph_and_tensors; + + // This graph is going to be used, so finish the build the cache deferred. + build_plans(Backend::F16, Pass::Fwd, *cache_entry); // Exit to request upper level API to allocate memory if needed // n.b. Care should be taken to align each of the added worksapce tensors to their type. @@ -525,7 +418,6 @@ void fused_attn_arbitrary_seqlen_fwd_impl( plan_workspace_size + actual_seqlen_workspace_size + seqlen_offsets_workspace_size; return; } - // cuDNN stream check needs to be moved here to support dummy kernel calls with // null streams for sizing the cuDNN workspace. NVTE_CHECK_CUDNN(cudnnSetStream(handle, stream)); @@ -634,407 +526,332 @@ void fused_attn_arbitrary_seqlen_fwd_impl( } NVTE_CHECK_CUDNN_FE(mha_graph->execute(handle, variant_pack, workspace)); + graph_cache_debug::record_execute(Backend::F16, Pass::Fwd); } catch (cudnn_frontend::cudnnException &e) { NVTE_ERROR(e.what()); } -} // NOLINT(readability/fn_size) - -void fused_attn_arbitrary_seqlen_bwd_impl( - int64_t b, int64_t h, int64_t hg, int64_t s_q, int64_t s_kv, int64_t d_qk, int64_t d_v, - int64_t max_b, int64_t max_t_q, int64_t max_t_kv, int64_t bias_b, int64_t bias_h, - int64_t bias_sq, int64_t bias_skv, float scaling_factor, float dropout_probability, - NVTE_QKV_Layout qkv_layout, NVTE_QKV_Format o_format, NVTE_QKV_Format do_format, - NVTE_QKV_Layout dqkv_layout, NVTE_Bias_Type bias_type, NVTE_Mask_Type mask_type, - NVTE_Softmax_Type softmax_type, int64_t window_size_left, int64_t window_size_right, - bool bottom_right_diagonal, bool deterministic, void *devPtrQ, void *devPtrKTranspose, - void *devPtrVTranspose, void *devPtrO, void *devPtrSoftmaxStats, void *devPtrBias, - void *devPtrSoftmaxOffset, void *devPtrdQ, void *devPtrdK, void *devPtrdV, void *devPtrdO, - void *devPtrdBias, void *devPtrdSoftmaxOffset, void *devPtrDropoutSeed, - void *devPtrDropoutOffset, void *devPtrCuSeqlensQ, void *devPtrCuSeqlensKV, - void *devPtrSeqOffsetsQ, void *devPtrSeqOffsetsKV, cudnn_frontend::DataType_t tensorType, - void *workspace, size_t *workspace_size, cudaStream_t stream, cudnnHandle_t handle) { - using namespace transformer_engine; +} - bool is_bias = (bias_type == NVTE_Bias_Type::NVTE_POST_SCALE_BIAS); - bool is_alibi = (bias_type == NVTE_Bias_Type::NVTE_ALIBI); - bool is_causal = ((mask_type == NVTE_Mask_Type::NVTE_CAUSAL_MASK) || - (mask_type == NVTE_Mask_Type::NVTE_PADDING_CAUSAL_MASK)); - bool is_bottom_right = ((mask_type == NVTE_Mask_Type::NVTE_CAUSAL_BOTTOM_RIGHT_MASK) || - (mask_type == NVTE_Mask_Type::NVTE_PADDING_CAUSAL_BOTTOM_RIGHT_MASK)); - bool is_padding = ((mask_type == NVTE_Mask_Type::NVTE_PADDING_MASK) || - (mask_type == NVTE_Mask_Type::NVTE_PADDING_CAUSAL_MASK) || - (mask_type == NVTE_Mask_Type::NVTE_PADDING_CAUSAL_BOTTOM_RIGHT_MASK)); - if (is_bottom_right && s_q == s_kv && !is_padding) { - is_causal = true; - is_bottom_right = false; - bottom_right_diagonal = false; - } - bool is_softmax_offset = (softmax_type != NVTE_Softmax_Type::NVTE_VANILLA_SOFTMAX); - bool is_dropout = (dropout_probability != 0.0f); - NVTE_QKV_Format q_format = nvte_get_q_format(qkv_layout); - NVTE_QKV_Format kv_format = nvte_get_kv_format(qkv_layout); - bool is_ragged_q = (q_format == NVTE_QKV_Format::NVTE_THD); - bool is_ragged_kv = (kv_format == NVTE_QKV_Format::NVTE_THD); +using F16BwdGraphAndTensors = + std::tuple, + std::shared_ptr, // q + std::shared_ptr, // k + std::shared_ptr, // v + std::shared_ptr, // o + std::shared_ptr, // dO + std::shared_ptr, // stats + std::shared_ptr, // attn_scale + std::shared_ptr, // dQ + std::shared_ptr, // dK + std::shared_ptr, // dV + std::shared_ptr, // bias + std::shared_ptr, // dBias + std::shared_ptr, // softmax_offset + std::shared_ptr, // d_softmax_offset + std::shared_ptr, // seq_q + std::shared_ptr, // seq_kv + std::shared_ptr, // offset_q + std::shared_ptr, // offset_k + std::shared_ptr, // offset_v + std::shared_ptr, // offset_o + std::shared_ptr, // offset_stats + std::shared_ptr, // dropout_seed + std::shared_ptr>; // dropout_offset + +static F16BwdGraphAndTensors create_graph_f16_bwd(const FusedAttnConfig &cfg) { + const int64_t b = static_cast(cfg.graph_batch_size_bwd); + const int64_t s_q = static_cast(cfg.graph_max_seqlen_q); + const int64_t s_kv = static_cast(cfg.graph_max_seqlen_kv); + const cudnn_frontend::DataType_t tensorType = + get_cudnn_fe_dtype(static_cast(cfg.qkv_dtype)); + const int64_t h = static_cast(cfg.num_attn_heads); + const int64_t hg = static_cast(cfg.num_gqa_groups); + const int64_t d_qk = static_cast(cfg.head_dim_qk); + const int64_t d_v = static_cast(cfg.head_dim_v); + const int64_t bias_b = static_cast(cfg.bias_batch_size); + const int64_t bias_h = static_cast(cfg.bias_num_heads); + const int64_t bias_sq = static_cast(cfg.bias_seqlen_q); + const int64_t bias_skv = static_cast(cfg.bias_seqlen_kv); + const int64_t window_size_left = cfg.window_size_left; + const int64_t window_size_right = cfg.window_size_right; + const float dropout_probability = cfg.dropout; + const NVTE_QKV_Layout qkv_layout = cfg.qkv_layout; + const bool bottom_right_diagonal = cfg.bottom_right_diagonal; + const bool deterministic = cfg.deterministic; + const bool is_bias = cfg.is_bias; + const bool is_alibi = cfg.is_alibi; + const bool is_causal = cfg.is_causal; + const bool is_causal_bottom_right = cfg.is_causal_bottom_right; + const bool is_padding = cfg.is_padding; + const bool is_softmax_offset = cfg.is_softmax_offset; + const bool is_dropout = cfg.is_dropout; + const bool is_ragged_q = cfg.is_ragged_q; + const bool is_ragged_kv = cfg.is_ragged_kv; const auto cudnn_runtime_version = cudnnGetVersion(); - const int device_id = cuda::current_device(); - const int sm_arch_ = cuda::sm_arch(device_id); - bool use_ragged_stats = is_ragged_q && cudnn_runtime_version >= 90600 && sm_arch_ != 120; - - NVTE_QKV_Layout_Group layout_group = nvte_get_qkv_layout_group(qkv_layout); - bool is_paged_kv = (layout_group == NVTE_QKV_Layout_Group::NVTE_Paged_KV_HD_HD_HD); - if (is_paged_kv) { - NVTE_CHECK(is_padding, "Paged attention requires padding mask!"); + const bool use_ragged_graph = cfg.uses_ragged_graph; + const bool use_ragged_stats = cfg.uses_ragged_stats; + const DType ragged_offset_type = cfg.ragged_offset_type_bwd; + + auto mha_graph = std::make_shared(); + mha_graph->set_io_data_type(tensorType) + .set_intermediate_data_type(fe::DataType_t::FLOAT) + .set_compute_data_type(fe::DataType_t::FLOAT); + + std::shared_ptr q, k, v, o, dO, stats, attn_scale; + std::shared_ptr bias, dBias, softmax_offset, d_softmax_offset, + seq_q, seq_kv; + std::shared_ptr offset_q, offset_k, offset_v, offset_o, + offset_stats; + std::shared_ptr dropout_seed, dropout_offset; + + std::vector q_stride(4); + std::vector k_stride(4); + std::vector v_stride(4); + std::vector o_stride(4); + generateMatrixStrides(b, h, s_q, s_kv, d_qk, q_stride.data(), qkv_layout, + NVTE_QKV_Matrix::NVTE_Q_Matrix); + generateMatrixStrides(b, hg, s_q, s_kv, d_qk, k_stride.data(), qkv_layout, + NVTE_QKV_Matrix::NVTE_K_Matrix); + generateMatrixStrides(b, hg, s_q, s_kv, d_v, v_stride.data(), qkv_layout, + NVTE_QKV_Matrix::NVTE_V_Matrix); + generateMatrixStrides(b, h, s_q, s_kv, d_v, o_stride.data(), qkv_layout, + NVTE_QKV_Matrix::NVTE_O_Matrix); + + q = mha_graph->tensor( + fe::graph::Tensor_attributes().set_name("Q").set_dim({b, h, s_q, d_qk}).set_stride(q_stride)); + k = mha_graph->tensor(fe::graph::Tensor_attributes() + .set_name("K") + .set_dim({b, hg, s_kv, d_qk}) + .set_stride(k_stride)); + v = mha_graph->tensor(fe::graph::Tensor_attributes() + .set_name("V") + .set_dim({b, hg, s_kv, d_v}) + .set_stride(v_stride)); + o = mha_graph->tensor( + fe::graph::Tensor_attributes().set_name("O").set_dim({b, h, s_q, d_v}).set_stride(o_stride)); + dO = mha_graph->tensor( + fe::graph::Tensor_attributes().set_name("dO").set_dim({b, h, s_q, d_v}).set_stride(o_stride)); + if (is_ragged_q) { + offset_q = mha_graph->tensor(fe::graph::Tensor_attributes() + .set_name("offset_q") + .set_dim({b + 1, 1, 1, 1}) + .set_stride({1, 1, 1, 1}) + .set_data_type(get_cudnn_fe_dtype(ragged_offset_type))); + offset_o = mha_graph->tensor(fe::graph::Tensor_attributes() + .set_name("offset_o") + .set_dim({b + 1, 1, 1, 1}) + .set_stride({1, 1, 1, 1}) + .set_data_type(get_cudnn_fe_dtype(ragged_offset_type))); + q->set_ragged_offset(offset_q); + o->set_ragged_offset(offset_o); + dO->set_ragged_offset(offset_o); } - - // keep original batch size because cu_seqlens are created with [b+1] shape - int64_t actual_b = b; - if ((is_ragged_q || is_ragged_kv) && cudnn_runtime_version >= 90600) { - NVTE_CHECK(is_padding, "Ragged QKV input requires padding or padding_causal mask!"); - // On SM 120, cuDNN support check requires BHSD-like strides with max_seqlen (see fwd). - if (sm_arch_ != 120) { - // replace batch size and maximum sequence lengths with maximum token counts - // for query and key/value so the graph is static within each quantization bucket - b = max_b; - s_q = is_ragged_q ? max_t_q : s_q; - s_kv = is_ragged_kv ? max_t_kv : s_kv; - } + if (is_ragged_kv) { + offset_k = mha_graph->tensor(fe::graph::Tensor_attributes() + .set_name("offset_k") + .set_dim({b + 1, 1, 1, 1}) + .set_stride({1, 1, 1, 1}) + .set_data_type(get_cudnn_fe_dtype(ragged_offset_type))); + offset_v = mha_graph->tensor(fe::graph::Tensor_attributes() + .set_name("offset_v") + .set_dim({b + 1, 1, 1, 1}) + .set_stride({1, 1, 1, 1}) + .set_data_type(get_cudnn_fe_dtype(ragged_offset_type))); + k->set_ragged_offset(offset_k); + v->set_ragged_offset(offset_v); } - // We choose between 32-bit and 64-bit offsets depending on need. - // This allows us to support older cuDNN runtimes gracefully. - const DType ragged_offset_type = cudnn_runtime_version >= 90500 ? DType::kInt64 : DType::kInt32; - - try { - FADescriptor_v1 descriptor{ - b, - h, - hg, - s_q, - s_kv, - d_qk, - d_v, - 0, - 0, - 0, - 0, - 0, - 0, - bias_b, - bias_h, - bias_sq, - bias_skv, - scaling_factor, - true, - dropout_probability, - qkv_layout, - o_format, - do_format, - dqkv_layout, - NVTE_QKV_Format_NOT_SET, - NVTE_QKV_Format_NOT_SET, - bias_type, - mask_type, - softmax_type, - window_size_left, - window_size_right, - bottom_right_diagonal, - deterministic, - tensorType, - cudnn_frontend::DataType_t::NOT_SET, - cudnn_frontend::DataType_t::NOT_SET, - cudnn_frontend::DataType_t::NOT_SET, - false, - }; - namespace fe = cudnn_frontend; - using graph_and_tensors = - std::tuple, - std::shared_ptr, // q - std::shared_ptr, // k - std::shared_ptr, // v - std::shared_ptr, // o - std::shared_ptr, // dO - std::shared_ptr, // stats - std::shared_ptr, // attn_scale - std::shared_ptr, // dQ - std::shared_ptr, // dK - std::shared_ptr, // dV - std::shared_ptr, // bias - std::shared_ptr, // dBias - std::shared_ptr, // softmax_offset - std::shared_ptr, // d_softmax_offset - std::shared_ptr, // seq_q - std::shared_ptr, // seq_kv - std::shared_ptr, // offset_q - std::shared_ptr, // offset_k - std::shared_ptr, // offset_v - std::shared_ptr, // offset_o - std::shared_ptr, // offset_stats - std::shared_ptr, // dropout_seed - std::shared_ptr>; // dropout_offset - - using CacheType = std::map; - static thread_local CacheType sdpa_f16_bprop_cache; - - // Get plan from cache if cache is available, otherwise create one - auto get_graph = [&](CacheType &cache, const FADescriptor_v1 &descriptor) -> graph_and_tensors { - // if hit, return - auto it = cache.find(descriptor); - if (it != cache.end()) { - auto graph = it->second; - return graph; - } - - // otherwise, build the op_graph and the plan. Then update cache - auto mha_graph = std::make_shared(); - mha_graph->set_io_data_type(tensorType) - .set_intermediate_data_type(fe::DataType_t::FLOAT) - .set_compute_data_type(fe::DataType_t::FLOAT); - - std::shared_ptr q, k, v, o, dO, stats, attn_scale; - std::shared_ptr bias, dBias, softmax_offset, d_softmax_offset, - seq_q, seq_kv; - std::shared_ptr offset_q, offset_k, offset_v, offset_o, - offset_stats; - std::shared_ptr dropout_seed, dropout_offset; - - std::vector q_stride(4); - std::vector k_stride(4); - std::vector v_stride(4); - std::vector o_stride(4); - generateMatrixStrides(b, h, s_q, s_kv, d_qk, q_stride.data(), qkv_layout, - NVTE_QKV_Matrix::NVTE_Q_Matrix); - generateMatrixStrides(b, hg, s_q, s_kv, d_qk, k_stride.data(), qkv_layout, - NVTE_QKV_Matrix::NVTE_K_Matrix); - generateMatrixStrides(b, hg, s_q, s_kv, d_v, v_stride.data(), qkv_layout, - NVTE_QKV_Matrix::NVTE_V_Matrix); - generateMatrixStrides(b, h, s_q, s_kv, d_v, o_stride.data(), qkv_layout, - NVTE_QKV_Matrix::NVTE_O_Matrix); - - q = mha_graph->tensor(fe::graph::Tensor_attributes() - .set_name("Q") - .set_dim({b, h, s_q, d_qk}) - .set_stride(q_stride)); - k = mha_graph->tensor(fe::graph::Tensor_attributes() - .set_name("K") - .set_dim({b, hg, s_kv, d_qk}) - .set_stride(k_stride)); - v = mha_graph->tensor(fe::graph::Tensor_attributes() - .set_name("V") - .set_dim({b, hg, s_kv, d_v}) - .set_stride(v_stride)); - o = mha_graph->tensor(fe::graph::Tensor_attributes() - .set_name("O") - .set_dim({b, h, s_q, d_v}) - .set_stride(o_stride)); - dO = mha_graph->tensor(fe::graph::Tensor_attributes() - .set_name("dO") - .set_dim({b, h, s_q, d_v}) - .set_stride(o_stride)); - if (is_ragged_q) { - offset_q = mha_graph->tensor(fe::graph::Tensor_attributes() - .set_name("offset_q") - .set_dim({b + 1, 1, 1, 1}) - .set_stride({1, 1, 1, 1}) - .set_data_type(get_cudnn_fe_dtype(ragged_offset_type))); - offset_o = mha_graph->tensor(fe::graph::Tensor_attributes() - .set_name("offset_o") + stats = mha_graph->tensor(fe::graph::Tensor_attributes() + .set_name("stats") + .set_dim({b, h, s_q, 1}) + .set_data_type(fe::DataType_t::FLOAT)); + if (use_ragged_stats) { + offset_stats = mha_graph->tensor(fe::graph::Tensor_attributes() + .set_name("offset_stats") .set_dim({b + 1, 1, 1, 1}) .set_stride({1, 1, 1, 1}) .set_data_type(get_cudnn_fe_dtype(ragged_offset_type))); - q->set_ragged_offset(offset_q); - o->set_ragged_offset(offset_o); - dO->set_ragged_offset(offset_o); - } - if (is_ragged_kv) { - offset_k = mha_graph->tensor(fe::graph::Tensor_attributes() - .set_name("offset_k") - .set_dim({b + 1, 1, 1, 1}) - .set_stride({1, 1, 1, 1}) - .set_data_type(get_cudnn_fe_dtype(ragged_offset_type))); - offset_v = mha_graph->tensor(fe::graph::Tensor_attributes() - .set_name("offset_v") - .set_dim({b + 1, 1, 1, 1}) - .set_stride({1, 1, 1, 1}) - .set_data_type(get_cudnn_fe_dtype(ragged_offset_type))); - k->set_ragged_offset(offset_k); - v->set_ragged_offset(offset_v); - } + stats->set_stride({h * s_q, 1, h, 1}).set_ragged_offset(offset_stats); + } else { + stats->set_stride({h * s_q, s_q, 1, 1}); + } - stats = mha_graph->tensor(fe::graph::Tensor_attributes() - .set_name("stats") - .set_dim({b, h, s_q, 1}) - .set_data_type(fe::DataType_t::FLOAT)); - if (use_ragged_stats) { - offset_stats = - mha_graph->tensor(fe::graph::Tensor_attributes() - .set_name("offset_stats") - .set_dim({b + 1, 1, 1, 1}) - .set_stride({1, 1, 1, 1}) - .set_data_type(get_cudnn_fe_dtype(ragged_offset_type))); - stats->set_stride({h * s_q, 1, h, 1}).set_ragged_offset(offset_stats); - } else { - stats->set_stride({h * s_q, s_q, 1, 1}); - } + attn_scale = mha_graph->tensor(fe::graph::Tensor_attributes() + .set_name("attn_scale") + .set_dim({1, 1, 1, 1}) + .set_stride({1, 1, 1, 1}) + .set_is_pass_by_value(true) + .set_data_type(fe::DataType_t::FLOAT)); - attn_scale = mha_graph->tensor(fe::graph::Tensor_attributes() - .set_name("attn_scale") - .set_dim({1, 1, 1, 1}) - .set_stride({1, 1, 1, 1}) - .set_is_pass_by_value(true) - .set_data_type(fe::DataType_t::FLOAT)); + fe::graph::SDPA_backward_attributes sdpa_backward_options; + sdpa_backward_options = fe::graph::SDPA_backward_attributes() + .set_name("flash_attention_backward") + .set_attn_scale(attn_scale); - fe::graph::SDPA_backward_attributes sdpa_backward_options; - sdpa_backward_options = fe::graph::SDPA_backward_attributes() - .set_name("flash_attention_backward") - .set_causal_mask(is_causal) - .set_causal_mask_bottom_right(is_bottom_right) - .set_attn_scale(attn_scale); + if (use_ragged_stats) { + sdpa_backward_options.set_max_total_seq_len_q(s_q); + } + if (is_ragged_kv && use_ragged_graph) { + sdpa_backward_options.set_max_total_seq_len_kv(s_kv); + } - if (use_ragged_stats) { - sdpa_backward_options.set_max_total_seq_len_q(s_q); - } - if (is_ragged_kv && cudnn_runtime_version >= 90600 && sm_arch_ != 120) { - sdpa_backward_options.set_max_total_seq_len_kv(s_kv); - } + fe::DiagonalAlignment_t const &diagonal_alignment = bottom_right_diagonal + ? fe::DiagonalAlignment_t::BOTTOM_RIGHT + : fe::DiagonalAlignment_t::TOP_LEFT; + sdpa_backward_options.set_diagonal_alignment(diagonal_alignment); - fe::DiagonalAlignment_t const &diagonal_alignment = - bottom_right_diagonal ? fe::DiagonalAlignment_t::BOTTOM_RIGHT - : fe::DiagonalAlignment_t::TOP_LEFT; - sdpa_backward_options.set_diagonal_alignment(diagonal_alignment); + if (cudnn_runtime_version >= 90200 && window_size_left != -1) { + sdpa_backward_options.set_diagonal_band_left_bound(window_size_left + 1); + } + if (cudnn_runtime_version >= 90600 && window_size_right != -1) { + sdpa_backward_options.set_diagonal_band_right_bound(window_size_right); + } + if (is_causal || is_causal_bottom_right) { + sdpa_backward_options.set_diagonal_band_right_bound(0); + } - if (cudnn_runtime_version >= 90200 && window_size_left != -1) { - sdpa_backward_options.set_diagonal_band_left_bound(window_size_left + 1); - } - if (cudnn_runtime_version >= 90600 && window_size_right != -1) { - sdpa_backward_options.set_diagonal_band_right_bound(window_size_right); - } + if (cudnn_runtime_version >= 90000) { + sdpa_backward_options.set_deterministic_algorithm(deterministic); + } - if (cudnn_runtime_version >= 90000) { - sdpa_backward_options.set_deterministic_algorithm(deterministic); - } + sdpa_backward_options.set_alibi_mask(is_alibi); + + if (is_bias) { + bias = mha_graph->tensor( + fe::graph::Tensor_attributes() + .set_name("bias") + .set_dim({bias_b, bias_h, bias_sq, bias_skv}) + .set_stride({bias_h * bias_sq * bias_skv, bias_sq * bias_skv, bias_skv, 1})); + sdpa_backward_options.set_bias(bias); + // bias shapes [1, 1, s, s], [b, 1, s, s], [b, h, s, s], [1, h, s, s] are supported for dbias calculation + // bias shape [1, 1, 1, s] is not supported for dbias calculation as of cuDNN 9.18 + if (!((bias_b == 1) && (bias_h == 1) && (bias_sq == 1))) { + dBias = mha_graph->tensor( + fe::graph::Tensor_attributes() + .set_name("dBias") + .set_dim({bias_b, bias_h, bias_sq, bias_skv}) + .set_stride({bias_h * bias_sq * bias_skv, bias_sq * bias_skv, bias_skv, 1})); + sdpa_backward_options.set_dbias(dBias); + } + } - sdpa_backward_options.set_alibi_mask(is_alibi); - - if (is_bias) { - bias = mha_graph->tensor( - fe::graph::Tensor_attributes() - .set_name("bias") - .set_dim({bias_b, bias_h, bias_sq, bias_skv}) - .set_stride({bias_h * bias_sq * bias_skv, bias_sq * bias_skv, bias_skv, 1})); - sdpa_backward_options.set_bias(bias); - // bias shapes [1, 1, s, s], [b, 1, s, s], [b, h, s, s], [1, h, s, s] are supported for dbias calculation - // bias shape [1, 1, 1, s] is not supported for dbias calculation as of cuDNN 9.18 - if (!((bias_b == 1) && (bias_h == 1) && (bias_sq == 1))) { - dBias = mha_graph->tensor( - fe::graph::Tensor_attributes() - .set_name("dBias") - .set_dim({bias_b, bias_h, bias_sq, bias_skv}) - .set_stride({bias_h * bias_sq * bias_skv, bias_sq * bias_skv, bias_skv, 1})); - sdpa_backward_options.set_dbias(dBias); - } - } + if (is_padding) { + seq_q = mha_graph->tensor(fe::graph::Tensor_attributes() + .set_name("seq_q") + .set_dim({b, 1, 1, 1}) + .set_stride({1, 1, 1, 1}) + .set_data_type(fe::DataType_t::INT32)); + seq_kv = mha_graph->tensor(fe::graph::Tensor_attributes() + .set_name("seq_kv") + .set_dim({b, 1, 1, 1}) + .set_stride({1, 1, 1, 1}) + .set_data_type(fe::DataType_t::INT32)); + sdpa_backward_options.set_padding_mask(is_padding).set_seq_len_q(seq_q).set_seq_len_kv(seq_kv); + } - if (is_padding) { - seq_q = mha_graph->tensor(fe::graph::Tensor_attributes() - .set_name("seq_q") - .set_dim({b, 1, 1, 1}) - .set_stride({1, 1, 1, 1}) - .set_data_type(fe::DataType_t::INT32)); - seq_kv = mha_graph->tensor(fe::graph::Tensor_attributes() - .set_name("seq_kv") - .set_dim({b, 1, 1, 1}) - .set_stride({1, 1, 1, 1}) - .set_data_type(fe::DataType_t::INT32)); - sdpa_backward_options.set_padding_mask(is_padding) - .set_seq_len_q(seq_q) - .set_seq_len_kv(seq_kv); - } + if (is_dropout) { + dropout_seed = mha_graph->tensor(fe::graph::Tensor_attributes() + .set_name("Seed") + .set_dim({1, 1, 1, 1}) + .set_stride({1, 1, 1, 1}) + .set_data_type(fe::DataType_t::INT64)); + dropout_offset = mha_graph->tensor(fe::graph::Tensor_attributes() + .set_name("Offset") + .set_dim({1, 1, 1, 1}) + .set_stride({1, 1, 1, 1}) + .set_data_type(fe::DataType_t::INT64)); + sdpa_backward_options.set_dropout(dropout_probability, dropout_seed, dropout_offset); + } - if (is_dropout) { - dropout_seed = mha_graph->tensor(fe::graph::Tensor_attributes() - .set_name("Seed") - .set_dim({1, 1, 1, 1}) - .set_stride({1, 1, 1, 1}) - .set_data_type(fe::DataType_t::INT64)); - dropout_offset = mha_graph->tensor(fe::graph::Tensor_attributes() - .set_name("Offset") - .set_dim({1, 1, 1, 1}) - .set_stride({1, 1, 1, 1}) - .set_data_type(fe::DataType_t::INT64)); - sdpa_backward_options.set_dropout(dropout_probability, dropout_seed, dropout_offset); - } + if (is_softmax_offset) { + softmax_offset = mha_graph->tensor(fe::graph::Tensor_attributes() + .set_name("softmax_offset") + .set_dim({1, h, 1, 1}) + .set_stride({h, 1, 1, 1}) + .set_data_type(fe::DataType_t::FLOAT)); + sdpa_backward_options.set_sink_token(softmax_offset); + d_softmax_offset = mha_graph->tensor(fe::graph::Tensor_attributes() + .set_name("d_softmax_offset") + .set_dim({1, h, 1, 1}) + .set_stride({h, 1, 1, 1}) + .set_data_type(fe::DataType_t::FLOAT)); + sdpa_backward_options.set_dsink_token(d_softmax_offset); + } - if (is_softmax_offset) { - softmax_offset = mha_graph->tensor(fe::graph::Tensor_attributes() - .set_name("softmax_offset") - .set_dim({1, h, 1, 1}) - .set_stride({h, 1, 1, 1}) - .set_data_type(fe::DataType_t::FLOAT)); - sdpa_backward_options.set_sink_token(softmax_offset); - d_softmax_offset = mha_graph->tensor(fe::graph::Tensor_attributes() - .set_name("d_softmax_offset") - .set_dim({1, h, 1, 1}) - .set_stride({h, 1, 1, 1}) - .set_data_type(fe::DataType_t::FLOAT)); - sdpa_backward_options.set_dsink_token(d_softmax_offset); - } + auto [dQ, dK, dV] = mha_graph->sdpa_backward(q, k, v, o, dO, stats, sdpa_backward_options); - auto [dQ, dK, dV] = mha_graph->sdpa_backward(q, k, v, o, dO, stats, sdpa_backward_options); + dQ->set_output(true).set_dim({b, h, s_q, d_qk}).set_stride(q_stride); + dK->set_output(true).set_dim({b, hg, s_kv, d_qk}).set_stride(k_stride); + dV->set_output(true).set_dim({b, hg, s_kv, d_v}).set_stride(v_stride); + if (is_ragged_q) { + dQ->set_ragged_offset(offset_q); + } + if (is_ragged_kv) { + dK->set_ragged_offset(offset_k); + dV->set_ragged_offset(offset_v); + } - dQ->set_output(true).set_dim({b, h, s_q, d_qk}).set_stride(q_stride); - dK->set_output(true).set_dim({b, hg, s_kv, d_qk}).set_stride(k_stride); - dV->set_output(true).set_dim({b, hg, s_kv, d_v}).set_stride(v_stride); - if (is_ragged_q) { - dQ->set_ragged_offset(offset_q); - } - if (is_ragged_kv) { - dK->set_ragged_offset(offset_k); - dV->set_ragged_offset(offset_v); - } + std::tuple, // q + std::shared_ptr, // k + std::shared_ptr, // v + std::shared_ptr, // o + std::shared_ptr, // dO + std::shared_ptr, // stats + std::shared_ptr, // attn_scale + std::shared_ptr, // dQ + std::shared_ptr, // dK + std::shared_ptr> // dV + key_tensors_tuple = std::make_tuple(q, k, v, o, dO, stats, attn_scale, dQ, dK, dV); + auto bias_tuple = is_bias ? std::make_tuple(bias, dBias) : std::make_tuple(nullptr, nullptr); + auto softmax_offset_tuple = is_softmax_offset ? std::make_tuple(softmax_offset, d_softmax_offset) + : std::make_tuple(nullptr, nullptr); + auto padding_tuple = + is_padding ? std::make_tuple(seq_q, seq_kv) : std::make_tuple(nullptr, nullptr); + auto offset_qo_tuple = + is_ragged_q ? std::make_tuple(offset_q, offset_o) : std::make_tuple(nullptr, nullptr); + auto offset_kv_tuple = + is_ragged_kv ? std::make_tuple(offset_k, offset_v) : std::make_tuple(nullptr, nullptr); + auto offset_s_tuple = use_ragged_stats ? std::make_tuple(offset_stats) : std::make_tuple(nullptr); + auto dropout_tuple = is_dropout ? std::make_tuple(dropout_seed, dropout_offset) + : std::make_tuple(nullptr, nullptr); + + return std::tuple_cat(std::make_tuple(mha_graph), key_tensors_tuple, bias_tuple, + softmax_offset_tuple, padding_tuple, offset_qo_tuple, offset_kv_tuple, + offset_s_tuple, dropout_tuple); +} - std::tuple, // q - std::shared_ptr, // k - std::shared_ptr, // v - std::shared_ptr, // o - std::shared_ptr, // dO - std::shared_ptr, // stats - std::shared_ptr, // attn_scale - std::shared_ptr, // dQ - std::shared_ptr, // dK - std::shared_ptr> // dV - key_tensors_tuple = std::make_tuple(q, k, v, o, dO, stats, attn_scale, dQ, dK, dV); - auto bias_tuple = is_bias ? std::make_tuple(bias, dBias) : std::make_tuple(nullptr, nullptr); - auto softmax_offset_tuple = is_softmax_offset - ? std::make_tuple(softmax_offset, d_softmax_offset) - : std::make_tuple(nullptr, nullptr); - auto padding_tuple = - is_padding ? std::make_tuple(seq_q, seq_kv) : std::make_tuple(nullptr, nullptr); - auto offset_qo_tuple = - is_ragged_q ? std::make_tuple(offset_q, offset_o) : std::make_tuple(nullptr, nullptr); - auto offset_kv_tuple = - is_ragged_kv ? std::make_tuple(offset_k, offset_v) : std::make_tuple(nullptr, nullptr); - auto offset_s_tuple = - use_ragged_stats ? std::make_tuple(offset_stats) : std::make_tuple(nullptr); - auto dropout_tuple = is_dropout ? std::make_tuple(dropout_seed, dropout_offset) - : std::make_tuple(nullptr, nullptr); - - NVTE_CHECK_CUDNN_FE(mha_graph->validate()); - NVTE_CHECK_CUDNN_FE(mha_graph->build_operation_graph(handle)); - NVTE_CHECK_CUDNN_FE(mha_graph->create_execution_plans({fe::HeurMode_t::A})); - NVTE_CHECK_CUDNN_FE(mha_graph->check_support(handle)); - NVTE_CHECK_CUDNN_FE(mha_graph->build_plans(handle)); - - auto return_tuple = std::tuple_cat(std::make_tuple(mha_graph), key_tensors_tuple, bias_tuple, - softmax_offset_tuple, padding_tuple, offset_qo_tuple, - offset_kv_tuple, offset_s_tuple, dropout_tuple); - cache.insert({descriptor, return_tuple}); - - return return_tuple; - }; +void fused_attn_arbitrary_seqlen_bwd_impl( + const FusedAttnConfig &cfg, void *devPtrQ, void *devPtrKTranspose, void *devPtrVTranspose, + void *devPtrO, void *devPtrSoftmaxStats, void *devPtrBias, void *devPtrSoftmaxOffset, + void *devPtrdQ, void *devPtrdK, void *devPtrdV, void *devPtrdO, void *devPtrdBias, + void *devPtrdSoftmaxOffset, void *devPtrDropoutSeed, void *devPtrDropoutOffset, + void *devPtrCuSeqlensQ, void *devPtrCuSeqlensKV, void *devPtrSeqOffsetsQ, + void *devPtrSeqOffsetsKV, void *workspace, size_t *workspace_size, cudaStream_t stream, + cudnnHandle_t handle) { + using namespace transformer_engine; + + cfg.check_derived(); + const int64_t b = static_cast(cfg.graph_batch_size_bwd); + const DType ragged_offset_type = cfg.ragged_offset_type_bwd; + const int64_t actual_b = static_cast(cfg.batch_size); + const bool use_ragged_stats = cfg.uses_ragged_stats; + float scaling_factor = cfg.attn_scale; + const bool is_bias = cfg.is_bias; + const bool is_padding = cfg.is_padding; + const bool is_softmax_offset = cfg.is_softmax_offset; + const bool is_dropout = cfg.is_dropout; + const bool is_ragged_q = cfg.is_ragged_q; + const bool is_ragged_kv = cfg.is_ragged_kv; + + try { + auto cache_entry = get_graph(cfg, handle); auto [mha_graph, q, k, v, o, dO, stats, attn_scale, dQ, dK, dV, bias, dBias, softmax_offset, d_softmax_offset, seq_q, seq_kv, offset_q, offset_o, offset_k, offset_v, offset_stats, - dropout_seed, dropout_offset] = get_graph(sdpa_f16_bprop_cache, descriptor); + dropout_seed, dropout_offset] = cache_entry->graph_and_tensors; + + // This graph is going to be used, so finish the build the cache deferred. + build_plans(Backend::F16, Pass::Bwd, *cache_entry); // Exit to request upper level API to allocate memory if needed // n.b. Care should be taken to align each of the added worksapce tensors to their type. @@ -1058,7 +875,6 @@ void fused_attn_arbitrary_seqlen_bwd_impl( plan_workspace_size + actual_seqlen_workspace_size + seqlen_offsets_workspace_size; return; } - // cuDNN stream check needs to be moved here to support dummy kernel calls with // null streams for sizing the cuDNN workspace. NVTE_CHECK_CUDNN(cudnnSetStream(handle, stream)); @@ -1122,10 +938,8 @@ void fused_attn_arbitrary_seqlen_bwd_impl( (static_cast(is_ragged_q) + static_cast(is_ragged_kv)) * 2 * num_bytes_per_ragged_offset; } - const RaggedOffsetMultipliers offset_mults(nvte_get_qkv_layout_group(qkv_layout), h, hg, d_qk, - d_v); cu_seqlens_padded_to_offsets<<>>( - offset_mults, actual_b, b, static_cast(devPtrSeqOffsetsQ), + cfg.ragged_offset_mults, actual_b, b, static_cast(devPtrSeqOffsetsQ), static_cast(devPtrSeqOffsetsKV), ragged_offset_type, devOffsetsQ, devOffsetsK, devOffsetsV, devOffsetsO, devOffsetsS); NVTE_CHECK_CUDA(cudaGetLastError()); @@ -1153,6 +967,7 @@ void fused_attn_arbitrary_seqlen_bwd_impl( } NVTE_CHECK_CUDNN_FE(mha_graph->execute(handle, variant_pack, workspace)); + graph_cache_debug::record_execute(Backend::F16, Pass::Bwd); } catch (cudnn_frontend::cudnnException &e) { NVTE_ERROR(e.what()); } @@ -1160,25 +975,26 @@ void fused_attn_arbitrary_seqlen_bwd_impl( } // namespace fused_attn using namespace transformer_engine::fused_attn; -void fused_attn_arbitrary_seqlen_fwd( - size_t batch, size_t num_attn_heads, size_t num_gqa_groups, size_t max_seqlen_q, - size_t max_seqlen_kv, size_t head_dim_qk, size_t head_dim_v, size_t num_tokens_q, - size_t num_tokens_kv, size_t num_pages_k, size_t num_pages_v, size_t page_size_k, - size_t page_size_v, size_t max_pages_per_seq_k, size_t max_pages_per_seq_v, bool is_training, - bool return_max_logit, float attn_scale, float p_dropout, NVTE_QKV_Layout qkv_layout, - NVTE_QKV_Format o_format, NVTE_Bias_Type bias_type, NVTE_Mask_Type mask_type, - NVTE_Softmax_Type softmax_type, int64_t window_size_left, int64_t window_size_right, - bool bottom_right_diagonal, const Tensor *input_Q, const Tensor *input_K, const Tensor *input_V, - const Tensor *input_Bias, const Tensor *input_SoftmaxOffset, Tensor *output_O, - NVTETensorPack *Aux_CTX_Tensors, const Tensor *cu_seqlens_q, const Tensor *cu_seqlens_kv, - const Tensor *cu_seqlens_q_padded, const Tensor *cu_seqlens_kv_padded, - const Tensor *page_table_k, const Tensor *page_table_v, const Tensor *rng_state, - Tensor *workspace, cudaStream_t stream, cudnnHandle_t handle) { +void fused_attn_arbitrary_seqlen_fwd(const FusedAttnConfig &cfg, const Tensor *input_Q, + const Tensor *input_K, const Tensor *input_V, + const Tensor *input_Bias, const Tensor *input_SoftmaxOffset, + Tensor *output_O, NVTETensorPack *Aux_CTX_Tensors, + const Tensor *cu_seqlens_q, const Tensor *cu_seqlens_kv, + const Tensor *cu_seqlens_q_padded, + const Tensor *cu_seqlens_kv_padded, const Tensor *page_table_k, + const Tensor *page_table_v, const Tensor *rng_state, + Tensor *workspace, cudaStream_t stream, cudnnHandle_t handle) { using namespace transformer_engine; + const size_t batch = cfg.batch_size; + const size_t num_attn_heads = cfg.num_attn_heads; + const size_t max_seqlen_q = cfg.max_seqlen_q; + const size_t num_tokens_q = cfg.num_tokens_q; + const bool return_max_logit = cfg.return_max_logit; + const NVTE_Bias_Type bias_type = cfg.bias_type; + const NVTE_Softmax_Type softmax_type = cfg.softmax_type; + const auto QKV_type = input_Q->data.dtype; - NVTE_QKV_Format q_format = nvte_get_q_format(qkv_layout); - NVTE_QKV_Format kv_format = nvte_get_kv_format(qkv_layout); void *devPtrQ = input_Q->data.dptr; void *devPtrK = input_K->data.dptr; void *devPtrV = input_V->data.dptr; @@ -1186,25 +1002,14 @@ void fused_attn_arbitrary_seqlen_fwd( void *devPtrS1 = nullptr; void *devPtrS2 = nullptr; void *devPtrBias = nullptr; - size_t bias_b = 0; - size_t bias_h = 0; - size_t bias_sq = 0; - size_t bias_skv = 0; if ((bias_type != NVTE_Bias_Type::NVTE_NO_BIAS) && (bias_type != NVTE_Bias_Type::NVTE_ALIBI)) { devPtrBias = input_Bias->data.dptr; - bias_b = input_Bias->data.shape[0]; - bias_h = input_Bias->data.shape[1]; - bias_sq = input_Bias->data.shape[2]; - bias_skv = input_Bias->data.shape[3]; } void *devPtrSoftmaxOffset = nullptr; if (softmax_type != NVTE_VANILLA_SOFTMAX) { devPtrSoftmaxOffset = input_SoftmaxOffset->data.dptr; } - const int device_id = cuda::current_device(); - const int sm_arch_ = cuda::sm_arch(device_id); - void *devPtrCuSeqlensQ = cu_seqlens_q->data.dptr; void *devPtrCuSeqlensKV = cu_seqlens_kv->data.dptr; void *devPtrSeqOffsetsQ = cu_seqlens_q_padded->data.dptr; @@ -1212,29 +1017,13 @@ void fused_attn_arbitrary_seqlen_fwd( void *devPtrPageTableK = page_table_k ? page_table_k->data.dptr : nullptr; void *devPtrPageTableV = page_table_v ? page_table_v->data.dptr : nullptr; - size_t max_batch_size = 0; - size_t max_tokens_q = 0; - size_t max_tokens_kv = 0; - if (q_format == NVTE_QKV_Format::NVTE_THD || kv_format == NVTE_QKV_Format::NVTE_THD) { - max_batch_size = get_max_batch_size(batch); - } - if (q_format == NVTE_QKV_Format::NVTE_THD) { - max_tokens_q = get_max_tokens(num_tokens_q); - } - if (kv_format == NVTE_QKV_Format::NVTE_THD) { - max_tokens_kv = get_max_tokens(num_tokens_kv); - } - size_t i = 0; if (Aux_CTX_Tensors->size == 0) { - const auto cudnn_runtime_version = cudnnGetVersion(); + const bool use_ragged_stats = cfg.uses_ragged_stats; Tensor *output_S = convertNVTETensorCheck(Aux_CTX_Tensors->tensors[i++]); output_S->data.dptr = nullptr; - // sm120 does not use ragged stats: the graph declares a dense - // [b, h, s_q, 1] stats tensor, so allocate to match (same as Max below). - if ((q_format == NVTE_QKV_Format::NVTE_THD && cudnn_runtime_version >= 90600) && - (sm_arch_ != 120)) { + if (use_ragged_stats) { output_S->data.shape = {num_tokens_q, num_attn_heads, 1}; } else { output_S->data.shape = {batch, num_attn_heads, max_seqlen_q, 1}; @@ -1244,8 +1033,7 @@ void fused_attn_arbitrary_seqlen_fwd( if (return_max_logit) { Tensor *output_Max = convertNVTETensorCheck(Aux_CTX_Tensors->tensors[i++]); output_Max->data.dptr = nullptr; - if ((q_format == NVTE_QKV_Format::NVTE_THD && cudnn_runtime_version >= 90600) && - (sm_arch_ != 120)) { + if (use_ragged_stats) { output_Max->data.shape = {num_tokens_q, num_attn_heads, 1}; } else { output_Max->data.shape = {batch, num_attn_heads, max_seqlen_q, 1}; @@ -1261,7 +1049,8 @@ void fused_attn_arbitrary_seqlen_fwd( if ((bias_type != NVTE_NO_BIAS) && (bias_type != NVTE_ALIBI)) { Tensor *output_bias = convertNVTETensorCheck(Aux_CTX_Tensors->tensors[i++]); output_bias->data.dptr = nullptr; - output_bias->data.shape = {bias_b, bias_h, bias_sq, bias_skv}; + output_bias->data.shape = {cfg.bias_batch_size, cfg.bias_num_heads, cfg.bias_seqlen_q, + cfg.bias_seqlen_kv}; output_bias->data.dtype = QKV_type; } @@ -1302,15 +1091,10 @@ void fused_attn_arbitrary_seqlen_fwd( size_t workspace_size = 0; fused_attn_arbitrary_seqlen_fwd_impl( - batch, num_attn_heads, num_gqa_groups, max_seqlen_q, max_seqlen_kv, head_dim_qk, head_dim_v, - max_batch_size, max_tokens_q, max_tokens_kv, num_pages_k, num_pages_v, page_size_k, - page_size_v, max_pages_per_seq_k, max_pages_per_seq_v, bias_b, bias_h, bias_sq, bias_skv, - is_training, return_max_logit, attn_scale, p_dropout, qkv_layout, o_format, bias_type, - mask_type, softmax_type, window_size_left, window_size_right, bottom_right_diagonal, devPtrQ, - devPtrK, devPtrV, devPtrBias, devPtrSoftmaxOffset, devPtrS1, devPtrS2, devPtrO, + cfg, devPtrQ, devPtrK, devPtrV, devPtrBias, devPtrSoftmaxOffset, devPtrS1, devPtrS2, devPtrO, devPtrDropoutSeed, devPtrDropoutOffset, devPtrCuSeqlensQ, devPtrCuSeqlensKV, devPtrPageTableK, - devPtrPageTableV, devPtrSeqOffsetsQ, devPtrSeqOffsetsKV, get_cudnn_fe_dtype(QKV_type), - workspace->data.dptr, &workspace_size, stream, handle); + devPtrPageTableV, devPtrSeqOffsetsQ, devPtrSeqOffsetsKV, workspace->data.dptr, + &workspace_size, stream, handle); if (workspace_size > 0) { if (workspace->data.dptr == nullptr) { @@ -1327,22 +1111,21 @@ void fused_attn_arbitrary_seqlen_fwd( } } -void fused_attn_arbitrary_seqlen_bwd( - size_t batch, size_t num_attn_heads, size_t num_gqa_groups, size_t max_seqlen_q, - size_t max_seqlen_kv, size_t head_dim_qk, size_t head_dim_v, size_t num_tokens_q, - size_t num_tokens_kv, float attn_scale, float p_dropout, NVTE_QKV_Layout qkv_layout, - NVTE_QKV_Format o_format, NVTE_QKV_Format do_format, NVTE_QKV_Layout dqkv_layout, - NVTE_Bias_Type bias_type, NVTE_Mask_Type mask_type, NVTE_Softmax_Type softmax_type, - int64_t window_size_left, int64_t window_size_right, bool bottom_right_diagonal, - bool deterministic, const Tensor *input_Q, const Tensor *input_K, const Tensor *input_V, - const Tensor *input_O, const Tensor *input_dO, const Tensor *input_Bias, - const Tensor *input_SoftmaxOffset, Tensor *output_S, Tensor *output_dQ, Tensor *output_dK, - Tensor *output_dV, Tensor *output_dBias, Tensor *output_dSoftmaxOffset, - const Tensor *cu_seqlens_q, const Tensor *cu_seqlens_kv, const Tensor *cu_seqlens_q_padded, - const Tensor *cu_seqlens_kv_padded, const Tensor *rng_state, Tensor *workspace, - cudaStream_t stream, cudnnHandle_t handle) { +void fused_attn_arbitrary_seqlen_bwd(const FusedAttnConfig &cfg, const Tensor *input_Q, + const Tensor *input_K, const Tensor *input_V, + const Tensor *input_O, const Tensor *input_dO, + const Tensor *input_Bias, const Tensor *input_SoftmaxOffset, + Tensor *output_S, Tensor *output_dQ, Tensor *output_dK, + Tensor *output_dV, Tensor *output_dBias, + Tensor *output_dSoftmaxOffset, const Tensor *cu_seqlens_q, + const Tensor *cu_seqlens_kv, const Tensor *cu_seqlens_q_padded, + const Tensor *cu_seqlens_kv_padded, const Tensor *rng_state, + Tensor *workspace, cudaStream_t stream, cudnnHandle_t handle) { using namespace transformer_engine; - const auto QKV_type = input_Q->data.dtype; + + const NVTE_Bias_Type bias_type = cfg.bias_type; + const NVTE_Softmax_Type softmax_type = cfg.softmax_type; + void *devPtrQ = input_Q->data.dptr; void *devPtrK = input_K->data.dptr; void *devPtrV = input_V->data.dptr; @@ -1350,32 +1133,9 @@ void fused_attn_arbitrary_seqlen_bwd( void *devPtrdO = input_dO->data.dptr; void *devPtrBias = nullptr; void *devPtrdBias = nullptr; - size_t bias_b = 0; - size_t bias_h = 0; - size_t bias_sq = 0; - size_t bias_skv = 0; if ((bias_type != NVTE_Bias_Type::NVTE_NO_BIAS) && (bias_type != NVTE_Bias_Type::NVTE_ALIBI)) { devPtrBias = input_Bias->data.dptr; devPtrdBias = output_dBias->data.dptr; - bias_b = output_dBias->data.shape[0]; - bias_h = output_dBias->data.shape[1]; - bias_sq = output_dBias->data.shape[2]; - bias_skv = output_dBias->data.shape[3]; - } - - size_t max_batch_size = 0; - size_t max_tokens_q = 0; - size_t max_tokens_kv = 0; - NVTE_QKV_Format q_format = nvte_get_q_format(qkv_layout); - NVTE_QKV_Format kv_format = nvte_get_kv_format(qkv_layout); - if (q_format == NVTE_QKV_Format::NVTE_THD || kv_format == NVTE_QKV_Format::NVTE_THD) { - max_batch_size = get_max_batch_size(batch); - } - if (q_format == NVTE_QKV_Format::NVTE_THD) { - max_tokens_q = get_max_tokens(num_tokens_q); - } - if (kv_format == NVTE_QKV_Format::NVTE_THD) { - max_tokens_kv = get_max_tokens(num_tokens_kv); } void *devPtrdQ = output_dQ->data.dptr; @@ -1402,14 +1162,10 @@ void fused_attn_arbitrary_seqlen_bwd( size_t workspace_size = 0; fused_attn_arbitrary_seqlen_bwd_impl( - batch, num_attn_heads, num_gqa_groups, max_seqlen_q, max_seqlen_kv, head_dim_qk, head_dim_v, - max_batch_size, max_tokens_q, max_tokens_kv, bias_b, bias_h, bias_sq, bias_skv, attn_scale, - p_dropout, qkv_layout, o_format, do_format, dqkv_layout, bias_type, mask_type, softmax_type, - window_size_left, window_size_right, bottom_right_diagonal, deterministic, devPtrQ, devPtrK, - devPtrV, devPtrO, devPtrSoftmaxStats, devPtrBias, devPtrSoftmaxOffset, devPtrdQ, devPtrdK, - devPtrdV, devPtrdO, devPtrdBias, devPtrdSoftmaxOffset, devPtrDropoutSeed, devPtrDropoutOffset, - devPtrCuSeqlensQ, devPtrCuSeqlensKV, devPtrSeqOffsetsQ, devPtrSeqOffsetsKV, - get_cudnn_fe_dtype(QKV_type), workspace->data.dptr, &workspace_size, stream, handle); + cfg, devPtrQ, devPtrK, devPtrV, devPtrO, devPtrSoftmaxStats, devPtrBias, devPtrSoftmaxOffset, + devPtrdQ, devPtrdK, devPtrdV, devPtrdO, devPtrdBias, devPtrdSoftmaxOffset, devPtrDropoutSeed, + devPtrDropoutOffset, devPtrCuSeqlensQ, devPtrCuSeqlensKV, devPtrSeqOffsetsQ, + devPtrSeqOffsetsKV, workspace->data.dptr, &workspace_size, stream, handle); if (workspace_size > 0) { if (workspace->data.dptr == nullptr) { @@ -1425,4 +1181,13 @@ void fused_attn_arbitrary_seqlen_bwd( NVTE_ERROR("Unexpected workspace_size."); } } + +// Check whether cuDNN can support a given config, per forward/backward pass. +std::string support_verdict_f16(const FusedAttnConfig &cfg, Pass pass, cudnnHandle_t handle) { + if (pass == Pass::Fwd) { + return fused_attn::support_verdict(cfg, handle); + } + return fused_attn::support_verdict(cfg, handle); +} + } // namespace transformer_engine diff --git a/transformer_engine/common/fused_attn/fused_attn_f16_arbitrary_seqlen.h b/transformer_engine/common/fused_attn/fused_attn_f16_arbitrary_seqlen.h index 8f79b5bb4a..ac26f4dd12 100644 --- a/transformer_engine/common/fused_attn/fused_attn_f16_arbitrary_seqlen.h +++ b/transformer_engine/common/fused_attn/fused_attn_f16_arbitrary_seqlen.h @@ -8,45 +8,42 @@ * \brief Functions for fused attention with seqlen > 512 */ -#ifndef TRANSFORMER_ENGINE_COMMON_FUSED_ATTN_FUSED_ATTN_ARBITRARY_SEQLEN_H_ -#define TRANSFORMER_ENGINE_COMMON_FUSED_ATTN_FUSED_ATTN_ARBITRARY_SEQLEN_H_ +#ifndef TRANSFORMER_ENGINE_COMMON_FUSED_ATTN_FUSED_ATTN_F16_ARBITRARY_SEQLEN_H_ +#define TRANSFORMER_ENGINE_COMMON_FUSED_ATTN_FUSED_ATTN_F16_ARBITRARY_SEQLEN_H_ #include +#include + #include "common/common.h" +#include "config_and_params.h" #include "transformer_engine/fused_attn.h" namespace transformer_engine { -void fused_attn_arbitrary_seqlen_fwd( - size_t batch, size_t num_attn_heads, size_t num_gqa_groups, size_t max_seqlen_q, - size_t max_seqlen_kv, size_t head_dim_qk, size_t head_dim_v, size_t num_tokens_q, - size_t num_tokens_kv, size_t num_pages_k, size_t num_pages_v, size_t page_size_k, - size_t page_size_v, size_t max_pages_per_seq_k, size_t max_pages_per_seq_v, bool is_training, - bool return_max_logit, float attn_scale, float p_dropout, NVTE_QKV_Layout qkv_layout, - NVTE_QKV_Format o_format, NVTE_Bias_Type bias_type, NVTE_Mask_Type mask_type, - NVTE_Softmax_Type softmax_type, int64_t window_size_left, int64_t window_size_right, - bool bottom_right_diagonal, const Tensor *input_Q, const Tensor *input_K, const Tensor *input_V, - const Tensor *input_Bias, const Tensor *input_SoftmaxOffset, Tensor *output_O, - NVTETensorPack *Aux_CTX_Tensors, const Tensor *cu_seqlens_q, const Tensor *cu_seqlens_kv, - const Tensor *cu_seqlens_q_padded, const Tensor *cu_seqlens_kv_padded, - const Tensor *page_table_k, const Tensor *page_table_v, const Tensor *rng_state, - Tensor *workspace, cudaStream_t stream, cudnnHandle_t handle); - -void fused_attn_arbitrary_seqlen_bwd( - size_t batch, size_t num_attn_heads, size_t num_gqa_groups, size_t max_seqlen_q, - size_t max_seqlen_kv, size_t head_dim_qk, size_t head_dim_v, size_t num_tokens_q, - size_t num_tokens_kv, float attn_scale, float p_dropout, NVTE_QKV_Layout qkv_layout, - NVTE_QKV_Format o_format, NVTE_QKV_Format do_format, NVTE_QKV_Layout dqkv_layout, - NVTE_Bias_Type bias_type, NVTE_Mask_Type mask_type, NVTE_Softmax_Type softmax_type, - int64_t window_size_left, int64_t window_size_right, bool bottom_right_diagonal, - bool deterministic, const Tensor *input_Q, const Tensor *input_K, const Tensor *input_V, - const Tensor *input_O, const Tensor *input_dO, const Tensor *input_Bias, - const Tensor *input_SoftmaxOffset, Tensor *output_S, Tensor *output_dQ, Tensor *output_dK, - Tensor *output_dV, Tensor *output_dBias, Tensor *output_dSoftmaxOffset, - const Tensor *cu_seqlens_q, const Tensor *cu_seqlens_kv, const Tensor *cu_seqlens_q_padded, - const Tensor *cu_seqlens_kv_padded, const Tensor *rng_state, Tensor *workspace, - cudaStream_t stream, cudnnHandle_t handle); +void fused_attn_arbitrary_seqlen_fwd(const fused_attn::FusedAttnConfig &cfg, const Tensor *input_Q, + const Tensor *input_K, const Tensor *input_V, + const Tensor *input_Bias, const Tensor *input_SoftmaxOffset, + Tensor *output_O, NVTETensorPack *Aux_CTX_Tensors, + const Tensor *cu_seqlens_q, const Tensor *cu_seqlens_kv, + const Tensor *cu_seqlens_q_padded, + const Tensor *cu_seqlens_kv_padded, const Tensor *page_table_k, + const Tensor *page_table_v, const Tensor *rng_state, + Tensor *workspace, cudaStream_t stream, cudnnHandle_t handle); + +void fused_attn_arbitrary_seqlen_bwd(const fused_attn::FusedAttnConfig &cfg, const Tensor *input_Q, + const Tensor *input_K, const Tensor *input_V, + const Tensor *input_O, const Tensor *input_dO, + const Tensor *input_Bias, const Tensor *input_SoftmaxOffset, + Tensor *output_S, Tensor *output_dQ, Tensor *output_dK, + Tensor *output_dV, Tensor *output_dBias, + Tensor *output_dSoftmaxOffset, const Tensor *cu_seqlens_q, + const Tensor *cu_seqlens_kv, const Tensor *cu_seqlens_q_padded, + const Tensor *cu_seqlens_kv_padded, const Tensor *rng_state, + Tensor *workspace, cudaStream_t stream, cudnnHandle_t handle); + +std::string support_verdict_f16(const fused_attn::FusedAttnConfig &cfg, fused_attn::Pass pass, + cudnnHandle_t handle); } // namespace transformer_engine -#endif // TRANSFORMER_ENGINE_COMMON_FUSED_ATTN_FUSED_ATTN_ARBITRARY_SEQLEN_H_ +#endif // TRANSFORMER_ENGINE_COMMON_FUSED_ATTN_FUSED_ATTN_F16_ARBITRARY_SEQLEN_H_ diff --git a/transformer_engine/common/fused_attn/fused_attn_fp8.cu b/transformer_engine/common/fused_attn/fused_attn_fp8.cu index 2ef0ac3393..c1e938a724 100644 --- a/transformer_engine/common/fused_attn/fused_attn_fp8.cu +++ b/transformer_engine/common/fused_attn/fused_attn_fp8.cu @@ -4,525 +4,453 @@ * See LICENSE for license information. ************************************************************************/ +#include + #include "../common.h" #include "../cudnn_utils.h" -#include "../util/cuda_runtime.h" #include "../util/system.h" #include "fused_attn_fp8.h" +#include "graph_cache.h" +#include "graph_cache_debug.h" #include "utils.h" namespace transformer_engine { namespace fused_attn { using namespace transformer_engine; - -constexpr size_t kFP8THDRaggedCudnnVersion = 92300; - -// fused attention FWD FP8 with FE 1.0+ -void fused_attn_fp8_fwd_impl( - int64_t b, int64_t h, int64_t hg, int64_t s_q, int64_t s_kv, int64_t d_qk, int64_t d_v, - int64_t max_b, int64_t max_t_q, int64_t max_t_kv, bool is_training, float scaling_factor, - float dropout_probability, NVTE_QKV_Layout qkv_layout, NVTE_QKV_Format o_format, - NVTE_Bias_Type bias_type, NVTE_Mask_Type mask_type, NVTE_Softmax_Type softmax_type, - int64_t window_size_left, int64_t window_size_right, bool bottom_right_diagonal, void* devPtrQ, - void* devPtrK, void* devPtrV, void* devPtrSoftmaxOffset, void* devPtrM, void* devPtrO, - void* devPtrDescaleQ, void* devPtrDescaleK, void* devPtrDescaleV, void* devPtrDescaleS, - void* devPtrScaleS, void* devPtrScaleO, void* devPtrAmaxO, void* devPtrAmaxS, - void* devPtrcuSeqlensQ, void* devPtrcuSeqlensKV, void* devPtrSeqOffsetsQ, - void* devPtrSeqOffsetsKV, void* devPtrDropoutSeed, void* devPtrDropoutOffset, - cudnn_frontend::DataType_t qkv_tensor_type, cudnn_frontend::DataType_t o_tensor_type, - NVTEScalingMode scaling_mode, NVTE_QKV_Format qkv_scale_inv_format, void* workspace, - size_t* workspace_size, cudaStream_t stream, cudnnHandle_t handle) { - using namespace transformer_engine; +namespace fe = cudnn_frontend; + +using Fp8FwdGraphAndTensors = + std::tuple, + std::shared_ptr, // Q + std::shared_ptr, // K + std::shared_ptr, // V + std::shared_ptr, // descale_q + std::shared_ptr, // descale_k + std::shared_ptr, // descale_v + std::shared_ptr, // descale_s + std::shared_ptr, // scale_s + std::shared_ptr, // scale_o + std::shared_ptr, // attn_scale + std::shared_ptr, // O + std::shared_ptr, // amax_s + std::shared_ptr, // amax_o + std::shared_ptr, // Stats + std::shared_ptr, // bias + std::shared_ptr, // softmax_offset + std::shared_ptr, // seq_q + std::shared_ptr, // seq_kv + std::shared_ptr, // offset_q + std::shared_ptr, // offset_o + std::shared_ptr, // offset_k + std::shared_ptr, // offset_v + std::shared_ptr, // offset_stats + std::shared_ptr, // dropout_seed + std::shared_ptr>; // dropout_offset + +static Fp8FwdGraphAndTensors create_graph_fp8_fwd(const FusedAttnConfig& cfg) { const auto cudnn_runtime_version = cudnnGetVersion(); - bool is_bias = (bias_type == NVTE_Bias_Type::NVTE_POST_SCALE_BIAS); - bool is_alibi = (bias_type == NVTE_Bias_Type::NVTE_ALIBI); - bool is_causal = ((mask_type == NVTE_Mask_Type::NVTE_CAUSAL_MASK) || - (mask_type == NVTE_Mask_Type::NVTE_PADDING_CAUSAL_MASK) || - (mask_type == NVTE_Mask_Type::NVTE_PADDING_CAUSAL_BOTTOM_RIGHT_MASK)); - bool is_padding = ((mask_type == NVTE_Mask_Type::NVTE_PADDING_MASK) || - (mask_type == NVTE_Mask_Type::NVTE_PADDING_CAUSAL_MASK) || - (mask_type == NVTE_Mask_Type::NVTE_PADDING_CAUSAL_BOTTOM_RIGHT_MASK)); - bool is_dropout = (is_training && dropout_probability != 0.0f); - bool is_softmax_offset = (softmax_type != NVTE_Softmax_Type::NVTE_VANILLA_SOFTMAX); - auto bias_b = b; - auto bias_h = h; - auto bias_sq = s_q; - auto bias_skv = s_kv; - NVTE_CHECK(~is_bias, "FP8 fused attention does not support pre/post_scale_bias yet!"); - NVTE_CHECK(~is_alibi, "FP8 fused attention does not support ALiBi yet!"); - bool is_delayed_scaling = (scaling_mode == NVTE_DELAYED_TENSOR_SCALING) && - (o_tensor_type == cudnn_frontend::DataType_t::FP8_E4M3 || - o_tensor_type == cudnn_frontend::DataType_t::FP8_E5M2); - bool is_current_scaling = (scaling_mode == NVTE_DELAYED_TENSOR_SCALING) && - (o_tensor_type == cudnn_frontend::DataType_t::HALF || - o_tensor_type == cudnn_frontend::DataType_t::BFLOAT16); - bool is_mxfp8 = (scaling_mode == NVTE_MXFP8_1D_SCALING) && - (o_tensor_type == cudnn_frontend::DataType_t::HALF || - o_tensor_type == cudnn_frontend::DataType_t::BFLOAT16); - NVTE_CHECK( - is_delayed_scaling || is_current_scaling || is_mxfp8, - "FP8 fused attention only supports FP8DelayedScaling or FP8CurrentScaling or MXFP8 recipes!"); - NVTE_CHECK(!is_mxfp8 || cudnn_runtime_version >= 92100, - "MXFP8 fused attention requires cuDNN 9.21.0 or later!"); - - NVTE_QKV_Format q_format = nvte_get_q_format(qkv_layout); - NVTE_QKV_Format kv_format = nvte_get_kv_format(qkv_layout); - bool is_ragged_q = (q_format == NVTE_QKV_Format::NVTE_THD); - bool is_ragged_kv = (kv_format == NVTE_QKV_Format::NVTE_THD); - const int device_id = cuda::current_device(); - const int sm_arch_ = cuda::sm_arch(device_id); - bool use_ragged_stats = - is_ragged_q && cudnn_runtime_version >= kFP8THDRaggedCudnnVersion && sm_arch_ != 120; - - NVTE_QKV_Layout_Group layout_group = nvte_get_qkv_layout_group(qkv_layout); - const DType ragged_offset_type = DType::kInt64; - - // Newer versions of cuDNN SDPA can accept sequence lengths directly as a cumulative - // tensor. Take advantage of this if possible to avoid the actual-seqlen conversion; - // THD inputs still use their separate ragged-offset tensors. - const bool use_cu_seqlens_directly = - // Frontend 1.26 supports fp8+cu_seqlens (for the C++ API). - // Note: For the Python API, 1.27 is required. - CUDNN_FRONTEND_VERSION >= 12600 && - // The frontend gates cu_seq_len support on min(compile-time, runtime) cuDNN - // version, so we'll do the same. - (CUDNN_VERSION >= 92500 && cudnn_runtime_version >= 92500) && - // This extra restriction is needed because cuDNN frontend doesn't yet allow - // the combination of dropout and stats generation for the fprop unified engine, - // so any such request would always get routed to the old composite SDPA engine - // (which doesn't support cu_seqlens). Remove this restriction when possible. - !is_dropout; - - int64_t actual_b = b; - if ((is_ragged_q || is_ragged_kv) && cudnn_runtime_version >= kFP8THDRaggedCudnnVersion) { - NVTE_CHECK(is_padding, "Ragged QKV input requires padding or padding_causal mask!"); - if (sm_arch_ != 120) { - // cuDNN reads the user's [actual_b+1] cu_seqlens buffers directly, so a quantized - // batch dimension would read out of bounds on the direct path. - if (!use_cu_seqlens_directly) { - b = max_b; - } - s_q = is_ragged_q ? max_t_q : s_q; - s_kv = is_ragged_kv ? max_t_kv : s_kv; + const cudnn_frontend::DataType_t qkv_tensor_type = + get_cudnn_fe_dtype(static_cast(cfg.qkv_dtype)); + const cudnn_frontend::DataType_t o_tensor_type = + get_cudnn_fe_dtype(static_cast(cfg.o_dtype)); + const int64_t b = static_cast(cfg.graph_batch_size_fwd); + const int64_t h = static_cast(cfg.num_attn_heads); + const int64_t hg = static_cast(cfg.num_gqa_groups); + const int64_t s_q = static_cast(cfg.graph_max_seqlen_q); + const int64_t s_kv = static_cast(cfg.graph_max_seqlen_kv); + const int64_t d_qk = static_cast(cfg.head_dim_qk); + const int64_t d_v = static_cast(cfg.head_dim_v); + const int64_t window_size_left = cfg.window_size_left; + const int64_t window_size_right = cfg.window_size_right; + const float dropout_probability = cfg.dropout; + const NVTE_QKV_Layout qkv_layout = cfg.qkv_layout; + const NVTE_QKV_Format o_format = cfg.o_format; + const NVTE_QKV_Format qkv_scale_inv_format = cfg.qkv_scale_inv_format; + const bool bottom_right_diagonal = cfg.bottom_right_diagonal; + const bool is_bias = cfg.is_bias; + const bool is_causal = cfg.is_causal; + const bool is_causal_bottom_right = cfg.is_causal_bottom_right; + const bool is_padding = cfg.is_padding; + const bool is_dropout = cfg.is_dropout; + const bool is_softmax_offset = cfg.is_softmax_offset; + const bool is_mxfp8 = cfg.is_mxfp8; + const bool is_tensor_scaling = cfg.is_tensor_scaling; + const bool is_delayed_scaling = cfg.is_delayed_scaling_fwd; + const bool is_current_scaling = cfg.is_current_scaling_fwd; + const bool use_cu_seqlens_directly = cfg.uses_cu_seqlens_directly; + const bool is_ragged_q = cfg.is_ragged_q; + const bool is_ragged_kv = cfg.is_ragged_kv; + const bool use_ragged_stats = cfg.uses_ragged_stats; + const DType ragged_offset_type = cfg.ragged_offset_type_fwd; + const RaggedOffsetMultipliers offset_mults = cfg.ragged_offset_mults; + + auto mha_graph = std::make_shared(); + mha_graph->set_io_data_type(qkv_tensor_type) + .set_intermediate_data_type(fe::DataType_t::FLOAT) + .set_compute_data_type(fe::DataType_t::FLOAT); + + std::shared_ptr Q, K, V, attn_scale; + std::shared_ptr descale_q, descale_k, descale_v; + std::shared_ptr descale_s, scale_s, scale_o; + std::shared_ptr bias, softmax_offset, seq_q, seq_kv; + std::shared_ptr offset_q, offset_k, offset_v, offset_o, + offset_stats; + std::shared_ptr dropout_seed, dropout_offset; + + // Q, K, V, attn_scale + std::vector q_strides(4), k_strides(4), v_strides(4); + generateMatrixStridesWithLayout(b, h, hg, s_q, s_kv, d_qk, d_v, q_strides.data(), + k_strides.data(), v_strides.data(), qkv_layout); + Q = mha_graph->tensor(fe::graph::Tensor_attributes() + .set_name("Q") + .set_dim({b, h, s_q, d_qk}) + .set_stride(q_strides) + .set_data_type(qkv_tensor_type)); + if (is_ragged_q) { + offset_q = mha_graph->tensor(fe::graph::Tensor_attributes() + .set_name("offset_q") + .set_dim({b + 1, 1, 1, 1}) + .set_stride({1, 1, 1, 1}) + .set_data_type(get_cudnn_fe_dtype(ragged_offset_type))); + Q->set_ragged_offset(offset_q); + if (use_cu_seqlens_directly) { + Q->set_ragged_offset_multiplier(offset_mults.q); + } + } + K = mha_graph->tensor(fe::graph::Tensor_attributes() + .set_name("K") + .set_dim({b, hg, s_kv, d_qk}) + .set_stride(k_strides) + .set_data_type(qkv_tensor_type)); + V = mha_graph->tensor(fe::graph::Tensor_attributes() + .set_name("V") + .set_dim({b, hg, s_kv, d_v}) + .set_stride(v_strides) + .set_data_type(qkv_tensor_type)); + if (is_ragged_kv) { + offset_k = mha_graph->tensor(fe::graph::Tensor_attributes() + .set_name("offset_k") + .set_dim({b + 1, 1, 1, 1}) + .set_stride({1, 1, 1, 1}) + .set_data_type(get_cudnn_fe_dtype(ragged_offset_type))); + offset_v = mha_graph->tensor(fe::graph::Tensor_attributes() + .set_name("offset_v") + .set_dim({b + 1, 1, 1, 1}) + .set_stride({1, 1, 1, 1}) + .set_data_type(get_cudnn_fe_dtype(ragged_offset_type))); + K->set_ragged_offset(offset_k); + V->set_ragged_offset(offset_v); + if (use_cu_seqlens_directly) { + K->set_ragged_offset_multiplier(offset_mults.k); + V->set_ragged_offset_multiplier(offset_mults.v); } } + attn_scale = mha_graph->tensor(fe::graph::Tensor_attributes() + .set_name("attn_scale") + .set_dim({1, 1, 1, 1}) + .set_stride({1, 1, 1, 1}) + .set_is_pass_by_value(true) + .set_data_type(fe::DataType_t::FLOAT)); + + // Descale_q, Descale_k, Descale_v, Descale_s, Scale_s, Scale_o + if (is_tensor_scaling) { + descale_q = mha_graph->tensor(fe::graph::Tensor_attributes() + .set_name("Descale_q") + .set_dim({1, 1, 1, 1}) + .set_stride({1, 1, 1, 1}) + .set_data_type(fe::DataType_t::FLOAT)); + descale_k = mha_graph->tensor_like(descale_q, "Descale_q"); + descale_v = mha_graph->tensor_like(descale_q, "Descale_v"); + descale_s = mha_graph->tensor_like(descale_q, "Descale_s"); + scale_s = mha_graph->tensor_like(descale_q, "Scale_s"); + if (is_delayed_scaling) { + scale_o = mha_graph->tensor_like(descale_q, "Scale_o"); + } + if (is_current_scaling) { + scale_o = mha_graph->tensor(1.0f); + } + } else if (is_mxfp8) { + const NVTE_QKV_Format q_scale_inv_format = + (qkv_scale_inv_format != NVTE_QKV_Format_NOT_SET) ? qkv_scale_inv_format : cfg.q_format; + const NVTE_QKV_Format kv_scale_inv_format = + (qkv_scale_inv_format != NVTE_QKV_Format_NOT_SET) ? qkv_scale_inv_format : cfg.kv_format; + std::vector q_scale_strides(4); + std::vector k_scale_strides(4); + std::vector v_scale_strides(4); + auto padded = pad_s_d_for_mxfp8(s_q, s_kv, d_qk, d_v); + generateMatrixStridesWithFormat(b, h, padded.s_q_padded, padded.d_qk_scale_padded, + q_scale_strides.data(), q_scale_inv_format); + generateMatrixStridesWithFormat(b, hg, padded.s_kv_padded, padded.d_qk_scale_padded, + k_scale_strides.data(), kv_scale_inv_format); + generateMatrixStridesWithFormat(b, hg, padded.s_kv_scale_padded, padded.d_v_padded, + v_scale_strides.data(), kv_scale_inv_format); + descale_q = mha_graph->tensor(fe::graph::Tensor_attributes() + .set_name("Descale_q") + .set_dim({b, h, padded.s_q_padded, padded.d_qk_scale_padded}) + .set_stride(q_scale_strides) + .set_data_type(fe::DataType_t::FP8_E8M0) + .set_reordering_type(fe::TensorReordering_t::F8_128x4)); + descale_k = + mha_graph->tensor(fe::graph::Tensor_attributes() + .set_name("Descale_k") + .set_dim({b, hg, padded.s_kv_padded, padded.d_qk_scale_padded}) + .set_stride(k_scale_strides) + .set_data_type(fe::DataType_t::FP8_E8M0) + .set_reordering_type(fe::TensorReordering_t::F8_128x4)); + descale_v = mha_graph->tensor(fe::graph::Tensor_attributes() + .set_name("Descale_v") + .set_dim({b, hg, padded.s_kv_scale_padded, padded.d_v_padded}) + .set_stride(v_scale_strides) + .set_data_type(fe::DataType_t::FP8_E8M0) + .set_reordering_type(fe::TensorReordering_t::F8_128x4)); + } - try { - FADescriptor_v1 descriptor{b, - h, - hg, - s_q, - s_kv, - d_qk, - d_v, - 0, - 0, - 0, - 0, - 0, - 0, - bias_b, - bias_h, - bias_sq, - bias_skv, - scaling_factor, - is_training, - dropout_probability, - qkv_layout, - o_format, - NVTE_QKV_Format_NOT_SET, - NVTE_QKV_Layout_NOT_SET, - qkv_scale_inv_format, - NVTE_QKV_Format_NOT_SET, - bias_type, - mask_type, - softmax_type, - window_size_left, - window_size_right, - bottom_right_diagonal, - true, - qkv_tensor_type, - o_tensor_type, - cudnn_frontend::DataType_t::NOT_SET, - cudnn_frontend::DataType_t::NOT_SET, - false}; - - namespace fe = cudnn_frontend; - using graph_and_tensors = - std::tuple, - std::shared_ptr, // Q - std::shared_ptr, // K - std::shared_ptr, // V - std::shared_ptr, // descale_q - std::shared_ptr, // descale_k - std::shared_ptr, // descale_v - std::shared_ptr, // descale_s - std::shared_ptr, // scale_s - std::shared_ptr, // scale_o - std::shared_ptr, // attn_scale - std::shared_ptr, // O - std::shared_ptr, // amax_s - std::shared_ptr, // amax_o - std::shared_ptr, // Stats - std::shared_ptr, // bias - std::shared_ptr, // softmax_offset - std::shared_ptr, // seq_q - std::shared_ptr, // seq_kv - std::shared_ptr, // offset_q - std::shared_ptr, // offset_k - std::shared_ptr, // offset_v - std::shared_ptr, // offset_o - std::shared_ptr, // offset_stats - std::shared_ptr, // dropout_seed - std::shared_ptr>; // dropout_offset - - using CacheType = std::map; - static thread_local CacheType sdpa_fp8_fprop_cache; - - // Get plan from cache if cache is available, otherwise create one - auto get_graph = [&](CacheType& cache, const FADescriptor_v1& descriptor) -> graph_and_tensors { - // if hit, return - auto it = cache.find(descriptor); - if (it != cache.end()) { - auto graph = it->second; - return graph; - } + fe::graph::SDPA_fp8_attributes sdpa_options; + sdpa_options = fe::graph::SDPA_fp8_attributes() + .set_name("sdpa_fp8") + .set_generate_stats(true) + .set_causal_mask(is_causal) + .set_attn_scale(attn_scale); + + fe::DiagonalAlignment_t const& diagonal_alignment = bottom_right_diagonal + ? fe::DiagonalAlignment_t::BOTTOM_RIGHT + : fe::DiagonalAlignment_t::TOP_LEFT; + sdpa_options.set_diagonal_alignment(diagonal_alignment); + + if (cudnn_runtime_version >= 92100) { + if (window_size_left != -1) { + sdpa_options.set_diagonal_band_left_bound(window_size_left + 1); + } + if (window_size_right != -1) { + sdpa_options.set_diagonal_band_right_bound(window_size_right); + } + } + if (is_causal_bottom_right) { + sdpa_options.set_diagonal_band_right_bound(0); + } - // otherwise, build the op_graph and the plan. Then update cache - auto mha_graph = std::make_shared(); - mha_graph->set_io_data_type(qkv_tensor_type) - .set_intermediate_data_type(fe::DataType_t::FLOAT) - .set_compute_data_type(fe::DataType_t::FLOAT); - - std::shared_ptr Q, K, V, attn_scale; - std::shared_ptr descale_q, descale_k, descale_v; - std::shared_ptr descale_s, scale_s, scale_o; - std::shared_ptr bias, softmax_offset, seq_q, seq_kv; - std::shared_ptr offset_q, offset_k, offset_v, offset_o, - offset_stats; - std::shared_ptr dropout_seed, dropout_offset; - - // Q, K, V, attn_scale - std::vector q_strides(4), k_strides(4), v_strides(4); - generateMatrixStridesWithLayout(b, h, hg, s_q, s_kv, d_qk, d_v, q_strides.data(), - k_strides.data(), v_strides.data(), qkv_layout); - Q = mha_graph->tensor(fe::graph::Tensor_attributes() - .set_name("Q") - .set_dim({b, h, s_q, d_qk}) - .set_stride(q_strides) - .set_data_type(qkv_tensor_type)); - if (is_ragged_q) { - offset_q = mha_graph->tensor(fe::graph::Tensor_attributes() - .set_name("offset_q") - .set_dim({b + 1, 1, 1, 1}) - .set_stride({1, 1, 1, 1}) - .set_data_type(get_cudnn_fe_dtype(ragged_offset_type))); - Q->set_ragged_offset(offset_q); - } - K = mha_graph->tensor(fe::graph::Tensor_attributes() - .set_name("K") - .set_dim({b, hg, s_kv, d_qk}) - .set_stride(k_strides) - .set_data_type(qkv_tensor_type)); - V = mha_graph->tensor(fe::graph::Tensor_attributes() - .set_name("V") - .set_dim({b, hg, s_kv, d_v}) - .set_stride(v_strides) - .set_data_type(qkv_tensor_type)); - if (is_ragged_kv) { - offset_k = mha_graph->tensor(fe::graph::Tensor_attributes() - .set_name("offset_k") - .set_dim({b + 1, 1, 1, 1}) - .set_stride({1, 1, 1, 1}) - .set_data_type(get_cudnn_fe_dtype(ragged_offset_type))); - offset_v = mha_graph->tensor(fe::graph::Tensor_attributes() - .set_name("offset_v") - .set_dim({b + 1, 1, 1, 1}) - .set_stride({1, 1, 1, 1}) - .set_data_type(get_cudnn_fe_dtype(ragged_offset_type))); - K->set_ragged_offset(offset_k); - V->set_ragged_offset(offset_v); - } - attn_scale = mha_graph->tensor(fe::graph::Tensor_attributes() - .set_name("attn_scale") + // sdpa_options.set_alibi_mask(is_alibi); + // if (is_bias) { + // bias = mha_graph->tensor(fe::graph::Tensor_attributes() + // .set_name("bias") + // .set_dim({bias_b, bias_h, bias_sq, bias_skv}) + // .set_stride({bias_h * bias_sq * bias_skv, bias_sq * bias_skv, bias_skv, 1})); + // sdpa_options.set_bias(bias); + // } + + if (is_padding) { + if (use_cu_seqlens_directly) { + // seq_q/seq_kv keep their tuple slots but hold (b+1)-shaped cu_seqlen tensors. + seq_q = mha_graph->tensor(fe::graph::Tensor_attributes() + .set_name("cu_seq_len_q") + .set_dim({b + 1, 1, 1, 1}) + .set_stride({1, 1, 1, 1}) + .set_data_type(fe::DataType_t::INT32)); + seq_kv = mha_graph->tensor(fe::graph::Tensor_attributes() + .set_name("cu_seq_len_kv") + .set_dim({b + 1, 1, 1, 1}) + .set_stride({1, 1, 1, 1}) + .set_data_type(fe::DataType_t::INT32)); + sdpa_options.set_padding_mask(is_padding).set_cu_seq_len_q(seq_q).set_cu_seq_len_kv(seq_kv); + // cu_seq_len (and the ragged offset multiplier) are unified-engine-only. + // Pin the implementation so an unsupported config fails with the unified + // engine's specific error instead of auto-selection's generic failure. + sdpa_options.set_implementation(fe::AttentionImplementation_t::UNIFIED); + } else { + seq_q = mha_graph->tensor(fe::graph::Tensor_attributes() + .set_name("seq_q") + .set_dim({b, 1, 1, 1}) + .set_stride({1, 1, 1, 1}) + .set_data_type(fe::DataType_t::INT32)); + seq_kv = mha_graph->tensor(fe::graph::Tensor_attributes() + .set_name("seq_kv") + .set_dim({b, 1, 1, 1}) + .set_stride({1, 1, 1, 1}) + .set_data_type(fe::DataType_t::INT32)); + sdpa_options.set_padding_mask(is_padding).set_seq_len_q(seq_q).set_seq_len_kv(seq_kv); + } + } + + if (is_dropout) { + dropout_seed = mha_graph->tensor(fe::graph::Tensor_attributes() + .set_name("Seed") .set_dim({1, 1, 1, 1}) .set_stride({1, 1, 1, 1}) - .set_is_pass_by_value(true) - .set_data_type(fe::DataType_t::FLOAT)); - - // Descale_q, Descale_k, Descale_v, Descale_s, Scale_s, Scale_o - if (is_delayed_scaling || is_current_scaling) { - descale_q = mha_graph->tensor(fe::graph::Tensor_attributes() - .set_name("Descale_q") - .set_dim({1, 1, 1, 1}) - .set_stride({1, 1, 1, 1}) - .set_data_type(fe::DataType_t::FLOAT)); - descale_k = mha_graph->tensor_like(descale_q, "Descale_q"); - descale_v = mha_graph->tensor_like(descale_q, "Descale_v"); - descale_s = mha_graph->tensor_like(descale_q, "Descale_s"); - scale_s = mha_graph->tensor_like(descale_q, "Scale_s"); - if (is_delayed_scaling) { - scale_o = mha_graph->tensor_like(descale_q, "Scale_o"); - } - if (is_current_scaling) { - scale_o = mha_graph->tensor(1.0f); - } - } else if (is_mxfp8) { - NVTE_QKV_Format q_scale_inv_format = (qkv_scale_inv_format != NVTE_QKV_Format_NOT_SET) - ? qkv_scale_inv_format - : nvte_get_q_format(qkv_layout); - NVTE_QKV_Format kv_scale_inv_format = (qkv_scale_inv_format != NVTE_QKV_Format_NOT_SET) - ? qkv_scale_inv_format - : nvte_get_kv_format(qkv_layout); - std::vector q_scale_strides(4); - std::vector k_scale_strides(4); - std::vector v_scale_strides(4); - auto padded = pad_s_d_for_mxfp8(s_q, s_kv, d_qk, d_v); - generateMatrixStridesWithFormat(b, h, padded.s_q_padded, padded.d_qk_scale_padded, - q_scale_strides.data(), q_scale_inv_format); - generateMatrixStridesWithFormat(b, hg, padded.s_kv_padded, padded.d_qk_scale_padded, - k_scale_strides.data(), kv_scale_inv_format); - generateMatrixStridesWithFormat(b, hg, padded.s_kv_scale_padded, padded.d_v_padded, - v_scale_strides.data(), kv_scale_inv_format); - descale_q = - mha_graph->tensor(fe::graph::Tensor_attributes() - .set_name("Descale_q") - .set_dim({b, h, padded.s_q_padded, padded.d_qk_scale_padded}) - .set_stride(q_scale_strides) - .set_data_type(fe::DataType_t::FP8_E8M0) - .set_reordering_type(fe::TensorReordering_t::F8_128x4)); - descale_k = - mha_graph->tensor(fe::graph::Tensor_attributes() - .set_name("Descale_k") - .set_dim({b, hg, padded.s_kv_padded, padded.d_qk_scale_padded}) - .set_stride(k_scale_strides) - .set_data_type(fe::DataType_t::FP8_E8M0) - .set_reordering_type(fe::TensorReordering_t::F8_128x4)); - descale_v = - mha_graph->tensor(fe::graph::Tensor_attributes() - .set_name("Descale_v") - .set_dim({b, hg, padded.s_kv_scale_padded, padded.d_v_padded}) - .set_stride(v_scale_strides) - .set_data_type(fe::DataType_t::FP8_E8M0) - .set_reordering_type(fe::TensorReordering_t::F8_128x4)); - } + .set_data_type(fe::DataType_t::INT64)); + dropout_offset = mha_graph->tensor(fe::graph::Tensor_attributes() + .set_name("Offset") + .set_dim({1, 1, 1, 1}) + .set_stride({1, 1, 1, 1}) + .set_data_type(fe::DataType_t::INT64)); + sdpa_options.set_dropout(dropout_probability, dropout_seed, dropout_offset); + } - fe::graph::SDPA_fp8_attributes sdpa_options; - sdpa_options = fe::graph::SDPA_fp8_attributes() - .set_name("sdpa_fp8") - .set_generate_stats(true) - .set_causal_mask(is_causal) - .set_attn_scale(attn_scale); - - fe::DiagonalAlignment_t const& diagonal_alignment = - bottom_right_diagonal ? fe::DiagonalAlignment_t::BOTTOM_RIGHT - : fe::DiagonalAlignment_t::TOP_LEFT; - sdpa_options.set_diagonal_alignment(diagonal_alignment); - - if (cudnn_runtime_version >= 92100) { - if (window_size_left != -1) { - sdpa_options.set_diagonal_band_left_bound(window_size_left + 1); - } - if (window_size_right != -1) { - sdpa_options.set_diagonal_band_right_bound(window_size_right); - } - } + if (is_softmax_offset) { + softmax_offset = mha_graph->tensor(fe::graph::Tensor_attributes() + .set_name("softmax_offset") + .set_dim({1, h, 1, 1}) + .set_stride({h, 1, 1, 1}) + .set_data_type(fe::DataType_t::FLOAT)); + sdpa_options.set_sink_token(softmax_offset); + } - // sdpa_options.set_alibi_mask(is_alibi); - // if (is_bias) { - // bias = mha_graph->tensor(fe::graph::Tensor_attributes() - // .set_name("bias") - // .set_dim({bias_b, bias_h, bias_sq, bias_skv}) - // .set_stride({bias_h * bias_sq * bias_skv, bias_sq * bias_skv, bias_skv, 1})); - // sdpa_options.set_bias(bias); - // } + std::shared_ptr O, Stats, amax_s, amax_o; + if (is_tensor_scaling) { + auto outputs = mha_graph->sdpa_fp8(Q, K, V, descale_q, descale_k, descale_v, descale_s, scale_s, + scale_o, sdpa_options); + O = outputs[0]; + Stats = outputs[1]; + amax_s = outputs[2]; + amax_o = outputs[3]; + amax_s->set_output(true) + .set_dim({1, 1, 1, 1}) + .set_stride({1, 1, 1, 1}) + .set_data_type(fe::DataType_t::FLOAT); + } else if (is_mxfp8) { + auto outputs = mha_graph->sdpa_fp8(Q, K, V, descale_q, descale_k, descale_v, sdpa_options); + O = outputs[0]; + Stats = outputs[1]; + amax_o = outputs[2]; + } - if (is_padding) { - if (use_cu_seqlens_directly) { - // seq_q/seq_kv keep their tuple slots but hold (b+1)-shaped cu_seqlen tensors. - seq_q = mha_graph->tensor(fe::graph::Tensor_attributes() - .set_name("cu_seq_len_q") - .set_dim({b + 1, 1, 1, 1}) - .set_stride({1, 1, 1, 1}) - .set_data_type(fe::DataType_t::INT32)); - seq_kv = mha_graph->tensor(fe::graph::Tensor_attributes() - .set_name("cu_seq_len_kv") + std::vector o_strides(4); + generateMatrixStridesWithFormat(b, h, s_q, d_v, o_strides.data(), o_format); + O->set_output(true).set_dim({b, h, s_q, d_v}).set_stride(o_strides).set_data_type(o_tensor_type); + if (is_ragged_q) { + offset_o = mha_graph->tensor(fe::graph::Tensor_attributes() + .set_name("offset_o") + .set_dim({b + 1, 1, 1, 1}) + .set_stride({1, 1, 1, 1}) + .set_data_type(get_cudnn_fe_dtype(ragged_offset_type))); + O->set_ragged_offset(offset_o); + if (use_cu_seqlens_directly) { + O->set_ragged_offset_multiplier(offset_mults.o); + } + } + amax_o->set_output(!is_mxfp8) + .set_dim({1, 1, 1, 1}) + .set_stride({1, 1, 1, 1}) + .set_data_type(fe::DataType_t::FLOAT); + + if (use_ragged_stats) { + offset_stats = mha_graph->tensor(fe::graph::Tensor_attributes() + .set_name("offset_stats") .set_dim({b + 1, 1, 1, 1}) .set_stride({1, 1, 1, 1}) - .set_data_type(fe::DataType_t::INT32)); - sdpa_options.set_padding_mask(is_padding) - .set_cu_seq_len_q(seq_q) - .set_cu_seq_len_kv(seq_kv); - // cu_seq_len (and the ragged offset multiplier) are unified-engine-only. - // Pin the implementation so an unsupported config fails with the unified - // engine's specific error instead of auto-selection's generic failure. - sdpa_options.set_implementation(fe::AttentionImplementation_t::UNIFIED); - } else { - seq_q = mha_graph->tensor(fe::graph::Tensor_attributes() - .set_name("seq_q") - .set_dim({b, 1, 1, 1}) - .set_stride({1, 1, 1, 1}) - .set_data_type(fe::DataType_t::INT32)); - seq_kv = mha_graph->tensor(fe::graph::Tensor_attributes() - .set_name("seq_kv") - .set_dim({b, 1, 1, 1}) - .set_stride({1, 1, 1, 1}) - .set_data_type(fe::DataType_t::INT32)); - sdpa_options.set_padding_mask(is_padding).set_seq_len_q(seq_q).set_seq_len_kv(seq_kv); - } - } - - if (is_dropout) { - dropout_seed = mha_graph->tensor(fe::graph::Tensor_attributes() - .set_name("Seed") - .set_dim({1, 1, 1, 1}) - .set_stride({1, 1, 1, 1}) - .set_data_type(fe::DataType_t::INT64)); - dropout_offset = mha_graph->tensor(fe::graph::Tensor_attributes() - .set_name("Offset") - .set_dim({1, 1, 1, 1}) - .set_stride({1, 1, 1, 1}) - .set_data_type(fe::DataType_t::INT64)); - sdpa_options.set_dropout(dropout_probability, dropout_seed, dropout_offset); - } + .set_data_type(get_cudnn_fe_dtype(ragged_offset_type))); + } + Stats->set_output(true).set_data_type(fe::DataType_t::FLOAT).set_dim({b, h, s_q, 1}); + if (use_ragged_stats) { + Stats->set_stride({h * s_q, 1, h, 1}).set_ragged_offset(offset_stats); + if (use_cu_seqlens_directly) { + Stats->set_ragged_offset_multiplier(offset_mults.stats); + } + } else { + Stats->set_stride({h * s_q, s_q, 1, 1}); + } - if (is_softmax_offset) { - softmax_offset = mha_graph->tensor(fe::graph::Tensor_attributes() - .set_name("softmax_offset") - .set_dim({1, h, 1, 1}) - .set_stride({h, 1, 1, 1}) - .set_data_type(fe::DataType_t::FLOAT)); - sdpa_options.set_sink_token(softmax_offset); - } + std::tuple, // Q + std::shared_ptr, // K + std::shared_ptr, // V + std::shared_ptr, // descale_q + std::shared_ptr, // descale_k + std::shared_ptr, // descale_v + std::shared_ptr, // descale_s + std::shared_ptr, // scale_s + std::shared_ptr, // scale_o + std::shared_ptr, // attn_scale + std::shared_ptr, // O + std::shared_ptr, // amax_s + std::shared_ptr> // amax_o + key_tensors_tuple = + is_mxfp8 ? std::make_tuple(Q, K, V, descale_q, descale_k, descale_v, nullptr, nullptr, + nullptr, attn_scale, O, nullptr, amax_o) + : std::make_tuple(Q, K, V, descale_q, descale_k, descale_v, descale_s, scale_s, + scale_o, attn_scale, O, amax_s, amax_o); + auto Stats_tuple = std::make_tuple(Stats); + auto bias_tuple = is_bias ? std::make_tuple(bias) : std::make_tuple(nullptr); + auto softmax_offset_tuple = + is_softmax_offset ? std::make_tuple(softmax_offset) : std::make_tuple(nullptr); + auto padding_tuple = + is_padding ? std::make_tuple(seq_q, seq_kv) : std::make_tuple(nullptr, nullptr); + auto offset_qo_tuple = + is_ragged_q ? std::make_tuple(offset_q, offset_o) : std::make_tuple(nullptr, nullptr); + auto offset_kv_tuple = + is_ragged_kv ? std::make_tuple(offset_k, offset_v) : std::make_tuple(nullptr, nullptr); + auto offset_s_tuple = use_ragged_stats ? std::make_tuple(offset_stats) : std::make_tuple(nullptr); + auto dropout_tuple = is_dropout ? std::make_tuple(dropout_seed, dropout_offset) + : std::make_tuple(nullptr, nullptr); + + return std::tuple_cat(std::make_tuple(mha_graph), key_tensors_tuple, Stats_tuple, bias_tuple, + softmax_offset_tuple, padding_tuple, offset_qo_tuple, offset_kv_tuple, + offset_s_tuple, dropout_tuple); +} - std::shared_ptr O, Stats, amax_s, amax_o; - if (is_delayed_scaling || is_current_scaling) { - auto outputs = mha_graph->sdpa_fp8(Q, K, V, descale_q, descale_k, descale_v, descale_s, - scale_s, scale_o, sdpa_options); - O = outputs[0]; - Stats = outputs[1]; - amax_s = outputs[2]; - amax_o = outputs[3]; - amax_s->set_output(true) - .set_dim({1, 1, 1, 1}) - .set_stride({1, 1, 1, 1}) - .set_data_type(fe::DataType_t::FLOAT); - } else if (is_mxfp8) { - auto outputs = mha_graph->sdpa_fp8(Q, K, V, descale_q, descale_k, descale_v, sdpa_options); - O = outputs[0]; - Stats = outputs[1]; - amax_o = outputs[2]; - } +void fused_attn_fp8_fwd_impl(const FusedAttnConfig& cfg, void* devPtrQ, void* devPtrK, + void* devPtrV, void* devPtrSoftmaxOffset, void* devPtrM, void* devPtrO, + void* devPtrDescaleQ, void* devPtrDescaleK, void* devPtrDescaleV, + void* devPtrDescaleS, void* devPtrScaleS, void* devPtrScaleO, + void* devPtrAmaxO, void* devPtrAmaxS, void* devPtrcuSeqlensQ, + void* devPtrcuSeqlensKV, void* devPtrSeqOffsetsQ, + void* devPtrSeqOffsetsKV, void* devPtrDropoutSeed, + void* devPtrDropoutOffset, void* workspace, size_t* workspace_size, + cudaStream_t stream, cudnnHandle_t handle) { + using namespace transformer_engine; - std::vector o_strides(4); - generateMatrixStridesWithFormat(b, h, s_q, d_v, o_strides.data(), o_format); - O->set_output(true) - .set_dim({b, h, s_q, d_v}) - .set_stride(o_strides) - .set_data_type(o_tensor_type); - if (is_ragged_q) { - offset_o = mha_graph->tensor(fe::graph::Tensor_attributes() - .set_name("offset_o") - .set_dim({b + 1, 1, 1, 1}) - .set_stride({1, 1, 1, 1}) - .set_data_type(get_cudnn_fe_dtype(ragged_offset_type))); - O->set_ragged_offset(offset_o); - } - amax_o->set_output(!is_mxfp8) - .set_dim({1, 1, 1, 1}) - .set_stride({1, 1, 1, 1}) - .set_data_type(fe::DataType_t::FLOAT); + cfg.check_derived(); - if (use_ragged_stats) { - offset_stats = - mha_graph->tensor(fe::graph::Tensor_attributes() - .set_name("offset_stats") - .set_dim({b + 1, 1, 1, 1}) - .set_stride({1, 1, 1, 1}) - .set_data_type(get_cudnn_fe_dtype(ragged_offset_type))); - } - Stats->set_output(true).set_data_type(fe::DataType_t::FLOAT).set_dim({b, h, s_q, 1}); - if (use_ragged_stats) { - Stats->set_stride({h * s_q, 1, h, 1}).set_ragged_offset(offset_stats); - } else { - Stats->set_stride({h * s_q, s_q, 1, 1}); - } + const bool is_tensor_scaling = cfg.is_tensor_scaling; + const bool is_delayed_scaling = cfg.is_delayed_scaling_fwd; + const bool use_cu_seqlens_directly = cfg.uses_cu_seqlens_directly; - std::tuple, // Q - std::shared_ptr, // K - std::shared_ptr, // V - std::shared_ptr, // descale_q - std::shared_ptr, // descale_k - std::shared_ptr, // descale_v - std::shared_ptr, // descale_s - std::shared_ptr, // scale_s - std::shared_ptr, // scale_o - std::shared_ptr, // attn_scale - std::shared_ptr, // O - std::shared_ptr, // amax_s - std::shared_ptr> // amax_o - key_tensors_tuple = - is_mxfp8 ? std::make_tuple(Q, K, V, descale_q, descale_k, descale_v, nullptr, nullptr, - nullptr, attn_scale, O, nullptr, amax_o) - : std::make_tuple(Q, K, V, descale_q, descale_k, descale_v, descale_s, - scale_s, scale_o, attn_scale, O, amax_s, amax_o); - auto Stats_tuple = std::make_tuple(Stats); - auto bias_tuple = is_bias ? std::make_tuple(bias) : std::make_tuple(nullptr); - auto softmax_offset_tuple = - is_softmax_offset ? std::make_tuple(softmax_offset) : std::make_tuple(nullptr); - auto padding_tuple = - is_padding ? std::make_tuple(seq_q, seq_kv) : std::make_tuple(nullptr, nullptr); - auto offset_q_tuple = is_ragged_q ? std::make_tuple(offset_q) : std::make_tuple(nullptr); - auto offset_kv_tuple = - is_ragged_kv ? std::make_tuple(offset_k, offset_v) : std::make_tuple(nullptr, nullptr); - auto offset_o_tuple = is_ragged_q ? std::make_tuple(offset_o) : std::make_tuple(nullptr); - auto offset_s_tuple = - use_ragged_stats ? std::make_tuple(offset_stats) : std::make_tuple(nullptr); - auto dropout_tuple = is_dropout ? std::make_tuple(dropout_seed, dropout_offset) - : std::make_tuple(nullptr, nullptr); - - NVTE_CHECK_CUDNN_FE(mha_graph->validate()); - NVTE_CHECK_CUDNN_FE(mha_graph->build_operation_graph(handle)); - NVTE_CHECK_CUDNN_FE(mha_graph->create_execution_plans({fe::HeurMode_t::A})); - NVTE_CHECK_CUDNN_FE(mha_graph->check_support(handle)); - NVTE_CHECK_CUDNN_FE(mha_graph->build_plans(handle)); - auto return_tuple = - std::tuple_cat(std::make_tuple(mha_graph), key_tensors_tuple, Stats_tuple, bias_tuple, - softmax_offset_tuple, padding_tuple, offset_q_tuple, offset_kv_tuple, - offset_o_tuple, offset_s_tuple, dropout_tuple); - cache.insert({descriptor, return_tuple}); - - return return_tuple; - }; + const int64_t b = static_cast(cfg.graph_batch_size_fwd); + const int64_t actual_b = static_cast(cfg.batch_size); + float scaling_factor = cfg.attn_scale; + const bool is_padding = cfg.is_padding; + const bool is_dropout = cfg.is_dropout; + const bool is_softmax_offset = cfg.is_softmax_offset; + const bool is_ragged_q = cfg.is_ragged_q; + const bool is_ragged_kv = cfg.is_ragged_kv; + const bool use_ragged_stats = cfg.uses_ragged_stats; + const DType ragged_offset_type = cfg.ragged_offset_type_fwd; + const RaggedOffsetMultipliers offset_mults = cfg.ragged_offset_mults; + try { + auto cache_entry = get_graph(cfg, handle); auto [mha_graph, Q, K, V, descale_q, descale_k, descale_v, descale_s, scale_s, scale_o, attn_scale, O, amax_s, amax_o, Stats, bias, softmax_offset, seq_q, seq_kv, offset_q, - offset_k, offset_v, offset_o, offset_stats, dropout_seed, dropout_offset] = - get_graph(sdpa_fp8_fprop_cache, descriptor); + offset_o, offset_k, offset_v, offset_stats, dropout_seed, dropout_offset] = + cache_entry->graph_and_tensors; + + // This graph is going to be used, so finish the build the cache deferred. + build_plans(Backend::FP8, Pass::Fwd, *cache_entry); + // Exit to request upper level API to allocate memory if needed. + // When passing cu_seqlens* directly to cuDNN SDPA, no conversion workspace is + // needed: cuDNN consumes the user's cu_seqlens buffers as-is. auto plan_workspace_size = alignTo<16>(mha_graph->get_workspace_size()); const size_t num_bytes_per_seqlen = alignTo<16>(b * sizeof(int32_t)); - const size_t actual_seqlen_workspace_size = - (is_padding && !use_cu_seqlens_directly) ? 2 * num_bytes_per_seqlen : 0; const size_t num_bytes_per_ragged_offset = alignTo<16>(((b + 1) * typeToNumBits(ragged_offset_type)) / 8); + size_t actual_seqlen_workspace_size = 0; size_t seqlen_offsets_workspace_size = 0; - if (is_ragged_q || is_ragged_kv) { - size_t count = 2 * (static_cast(is_ragged_q) + static_cast(is_ragged_kv)); - if (use_ragged_stats) { - seqlen_offsets_workspace_size = (count + 1) * num_bytes_per_ragged_offset; - } else { - seqlen_offsets_workspace_size = count * num_bytes_per_ragged_offset; + if (!use_cu_seqlens_directly) { + if (is_padding) { + actual_seqlen_workspace_size = 2 * num_bytes_per_seqlen; + } + if (is_ragged_q || is_ragged_kv) { + const size_t count = + 2 * (static_cast(is_ragged_q) + static_cast(is_ragged_kv)); + seqlen_offsets_workspace_size = + (use_ragged_stats ? count + 1 : count) * num_bytes_per_ragged_offset; } } - if (workspace == nullptr) { *workspace_size = plan_workspace_size + actual_seqlen_workspace_size + seqlen_offsets_workspace_size; return; } - // cuDNN stream check needs to be moved here to support dummy kernel calls with // null streams for sizing the cuDNN workspace. NVTE_CHECK_CUDNN(cudnnSetStream(handle, stream)); @@ -542,14 +470,14 @@ void fused_attn_fp8_fwd_impl( if (is_delayed_scaling) { variant_pack[scale_o] = devPtrScaleO; } - if (is_delayed_scaling || is_current_scaling) { + if (is_tensor_scaling) { variant_pack[descale_s] = devPtrDescaleS; variant_pack[scale_s] = devPtrScaleS; variant_pack[amax_s] = devPtrAmaxS; variant_pack[amax_o] = devPtrAmaxO; } - /* if (is_bias) { + /* if (cfg.is_bias) { variant_pack[bias] = devPtrBias; } */ @@ -572,7 +500,22 @@ void fused_attn_fp8_fwd_impl( } } - if (is_ragged_q || is_ragged_kv) { + if (use_cu_seqlens_directly) { + // The token-unit cu_seqlens_padded buffers serve as the ragged offsets; the engine + // applies the per-tensor multipliers set at graph build time. + if (is_ragged_q) { + variant_pack[offset_q] = devPtrSeqOffsetsQ; + variant_pack[offset_o] = devPtrSeqOffsetsQ; + } + if (is_ragged_kv) { + void* devOffsetsKV = offset_mults.kv_from_q ? devPtrSeqOffsetsQ : devPtrSeqOffsetsKV; + variant_pack[offset_k] = devOffsetsKV; + variant_pack[offset_v] = devOffsetsKV; + } + if (use_ragged_stats) { + variant_pack[offset_stats] = devPtrSeqOffsetsQ; + } + } else if (is_ragged_q || is_ragged_kv) { constexpr size_t nthreads_per_block = 128; const size_t grid = (b + nthreads_per_block) / nthreads_per_block; void* devOffsets = @@ -596,7 +539,6 @@ void fused_attn_fp8_fwd_impl( (static_cast(is_ragged_q) + static_cast(is_ragged_kv)) * 2 * num_bytes_per_ragged_offset; } - const RaggedOffsetMultipliers offset_mults(layout_group, h, hg, d_qk, d_v); cu_seqlens_padded_to_offsets<<>>( offset_mults, actual_b, b, static_cast(devPtrSeqOffsetsQ), static_cast(devPtrSeqOffsetsKV), ragged_offset_type, devOffsetsQ, devOffsetsK, @@ -625,668 +567,596 @@ void fused_attn_fp8_fwd_impl( } NVTE_CHECK_CUDNN_FE(mha_graph->execute(handle, variant_pack, workspace)); + graph_cache_debug::record_execute(Backend::FP8, Pass::Fwd); } catch (cudnn_frontend::cudnnException& e) { NVTE_ERROR(e.what()); } -} // NOLINT(readability/fn_size) +} -// fused attention BWD FP8 with FE 1.0+ -void fused_attn_fp8_bwd_impl( - int64_t b, int64_t h, int64_t hg, int64_t s_q, int64_t s_kv, int64_t d_qk, int64_t d_v, - int64_t max_b, int64_t max_t_q, int64_t max_t_kv, float scaling_factor, - float dropout_probability, NVTE_QKV_Layout qkv_layout, NVTE_QKV_Format o_format, - NVTE_QKV_Format do_format, NVTE_QKV_Layout dqkv_layout, NVTE_Bias_Type bias_type, - NVTE_Mask_Type mask_type, NVTE_Softmax_Type softmax_type, int64_t window_size_left, - int64_t window_size_right, bool bottom_right_diagonal, bool deterministic, void* devPtrQ, - void* devPtrK, void* devPtrV, void* devPtrM, void* devPtrO, void* devPtrdO, - void* devPtrSoftmaxOffset, void* devPtrdQ, void* devPtrdK, void* devPtrdV, - void* devPtrdSoftmaxOffset, void* devPtrDescaleQ, void* devPtrDescaleK, void* devPtrDescaleV, - void* devPtrDescaleO, void* devPtrDescaledO, void* devPtrDescaleS, void* devPtrDescaledP, - void* devPtrScaleS, void* devPtrScaledP, void* devPtrScaledQ, void* devPtrScaledK, - void* devPtrScaledV, void* devPtrAmaxdP, void* devPtrAmaxdQ, void* devPtrAmaxdK, - void* devPtrAmaxdV, void* devPtrQ_t, void* devPtrK_t, void* devPtrdO_f16, void* devPtrdO_t, - void* devPtrDescaleQ_t, void* devPtrDescaleK_t, void* devPtrDescaledO_t, void* devPtrcuSeqlensQ, - void* devPtrcuSeqlensKV, void* devPtrSeqOffsetsQ, void* devPtrSeqOffsetsKV, - void* devPtrDropoutSeed, void* devPtrDropoutOffset, cudnn_frontend::DataType_t qkv_tensor_type, - cudnn_frontend::DataType_t o_tensor_type, cudnn_frontend::DataType_t do_tensor_type, - cudnn_frontend::DataType_t dqkv_tensor_type, NVTEScalingMode scaling_mode, - NVTE_QKV_Format qkv_scale_inv_format, NVTE_QKV_Format do_scale_inv_format, void* workspace, - size_t* workspace_size, cudaStream_t stream, cudnnHandle_t handle) { - using namespace transformer_engine; +using Fp8BwdGraphAndTensors = + std::tuple, + std::shared_ptr, // Q + std::shared_ptr, // Q_t + std::shared_ptr, // K + std::shared_ptr, // K_t + std::shared_ptr, // V + std::shared_ptr, // O + std::shared_ptr, // Stats + std::shared_ptr, // dO + std::shared_ptr, // dO_t + std::shared_ptr, // dO_f16 + std::shared_ptr, // attn_scale + std::shared_ptr, // descale_q + std::shared_ptr, // descale_q_t + std::shared_ptr, // descale_k + std::shared_ptr, // descale_k_t + std::shared_ptr, // descale_v + std::shared_ptr, // descale_o + std::shared_ptr, // descale_dO + std::shared_ptr, // descale_dO_t + std::shared_ptr, // descale_s + std::shared_ptr, // descale_dP + std::shared_ptr, // scale_dQ + std::shared_ptr, // scale_dK + std::shared_ptr, // scale_dV + std::shared_ptr, // scale_s + std::shared_ptr, // scale_dP + std::shared_ptr, // dQ + std::shared_ptr, // dK + std::shared_ptr, // dV + std::shared_ptr, // amax_dQ + std::shared_ptr, // amax_dK + std::shared_ptr, // amax_dV + std::shared_ptr, // amax_dP + std::shared_ptr, // bias + std::shared_ptr, // dBias + std::shared_ptr, // softmax_offset + std::shared_ptr, // d_softmax_offset + std::shared_ptr, // seq_q + std::shared_ptr, // seq_kv + std::shared_ptr, // offset_q + std::shared_ptr, // offset_o + std::shared_ptr, // offset_k + std::shared_ptr, // offset_v + std::shared_ptr, // offset_stats + std::shared_ptr, // dropout_seed + std::shared_ptr>; // dropout_offset + +static Fp8BwdGraphAndTensors create_graph_fp8_bwd(const FusedAttnConfig& cfg) { const auto cudnn_runtime_version = cudnnGetVersion(); - bool is_bias = (bias_type == NVTE_Bias_Type::NVTE_POST_SCALE_BIAS); - bool is_alibi = (bias_type == NVTE_Bias_Type::NVTE_ALIBI); - bool is_causal = ((mask_type == NVTE_Mask_Type::NVTE_CAUSAL_MASK) || - (mask_type == NVTE_Mask_Type::NVTE_PADDING_CAUSAL_MASK) || - (mask_type == NVTE_Mask_Type::NVTE_PADDING_CAUSAL_BOTTOM_RIGHT_MASK)); - bool is_padding = ((mask_type == NVTE_Mask_Type::NVTE_PADDING_MASK) || - (mask_type == NVTE_Mask_Type::NVTE_PADDING_CAUSAL_MASK) || - (mask_type == NVTE_Mask_Type::NVTE_PADDING_CAUSAL_BOTTOM_RIGHT_MASK)); - bool is_dropout = (dropout_probability != 0.0f); - bool is_softmax_offset = (softmax_type != NVTE_Softmax_Type::NVTE_VANILLA_SOFTMAX); - auto bias_b = b; - auto bias_h = h; - auto bias_sq = s_q; - auto bias_skv = s_kv; - NVTE_CHECK(~is_bias, "FP8 fused attention does not support pre/post_scale_bias yet!"); - NVTE_CHECK(~is_alibi, "FP8 fused attention does not support ALiBi yet!"); - bool is_delayed_scaling = (scaling_mode == NVTE_DELAYED_TENSOR_SCALING) && - (dqkv_tensor_type == cudnn_frontend::DataType_t::FP8_E4M3 || - dqkv_tensor_type == cudnn_frontend::DataType_t::FP8_E5M2); - bool is_current_scaling = (scaling_mode == NVTE_DELAYED_TENSOR_SCALING) && - (dqkv_tensor_type == cudnn_frontend::DataType_t::HALF || - dqkv_tensor_type == cudnn_frontend::DataType_t::BFLOAT16); - bool is_mxfp8 = (scaling_mode == NVTE_MXFP8_1D_SCALING) && - (dqkv_tensor_type == cudnn_frontend::DataType_t::HALF || - dqkv_tensor_type == cudnn_frontend::DataType_t::BFLOAT16); - NVTE_CHECK( - is_delayed_scaling || is_current_scaling || is_mxfp8, - "FP8 fused attention only supports FP8DelayedScaling or FP8CurrentScaling or MXFP8 recipes!"); - NVTE_CHECK(!is_mxfp8 || cudnn_runtime_version >= 92100, - "MXFP8 fused attention requires cuDNN 9.21.0 or later!"); - - NVTE_QKV_Format q_format = nvte_get_q_format(qkv_layout); - NVTE_QKV_Format kv_format = nvte_get_kv_format(qkv_layout); - bool is_ragged_q = (q_format == NVTE_QKV_Format::NVTE_THD); - bool is_ragged_kv = (kv_format == NVTE_QKV_Format::NVTE_THD); - const int device_id = cuda::current_device(); - const int sm_arch_ = cuda::sm_arch(device_id); - bool use_ragged_stats = - is_ragged_q && cudnn_runtime_version >= kFP8THDRaggedCudnnVersion && sm_arch_ != 120; - - NVTE_QKV_Layout_Group layout_group = nvte_get_qkv_layout_group(qkv_layout); - const DType ragged_offset_type = DType::kInt64; - - int64_t actual_b = b; - if ((is_ragged_q || is_ragged_kv) && cudnn_runtime_version >= kFP8THDRaggedCudnnVersion) { - NVTE_CHECK(is_padding, "Ragged QKV input requires padding or padding_causal mask!"); - if (sm_arch_ != 120) { - b = max_b; - s_q = is_ragged_q ? max_t_q : s_q; - s_kv = is_ragged_kv ? max_t_kv : s_kv; - } + const cudnn_frontend::DataType_t qkv_tensor_type = + get_cudnn_fe_dtype(static_cast(cfg.qkv_dtype)); + const cudnn_frontend::DataType_t o_tensor_type = + get_cudnn_fe_dtype(static_cast(cfg.o_dtype)); + const cudnn_frontend::DataType_t do_tensor_type = + get_cudnn_fe_dtype(static_cast(cfg.do_dtype)); + const cudnn_frontend::DataType_t dqkv_tensor_type = + get_cudnn_fe_dtype(static_cast(cfg.dqkv_dtype)); + const int64_t b = static_cast(cfg.graph_batch_size_bwd); + const int64_t h = static_cast(cfg.num_attn_heads); + const int64_t hg = static_cast(cfg.num_gqa_groups); + const int64_t s_q = static_cast(cfg.graph_max_seqlen_q); + const int64_t s_kv = static_cast(cfg.graph_max_seqlen_kv); + const int64_t d_qk = static_cast(cfg.head_dim_qk); + const int64_t d_v = static_cast(cfg.head_dim_v); + const int64_t window_size_left = cfg.window_size_left; + const int64_t window_size_right = cfg.window_size_right; + const float dropout_probability = cfg.dropout; + const NVTE_QKV_Layout qkv_layout = cfg.qkv_layout; + const NVTE_QKV_Layout dqkv_layout = cfg.dqkv_layout; + const NVTE_QKV_Format o_format = cfg.o_format; + const NVTE_QKV_Format do_format = cfg.do_format; + const NVTE_QKV_Format qkv_scale_inv_format = cfg.qkv_scale_inv_format; + const NVTE_QKV_Format do_scale_inv_format = cfg.do_scale_inv_format; + const bool bottom_right_diagonal = cfg.bottom_right_diagonal; + const bool deterministic = cfg.deterministic; + const bool is_bias = cfg.is_bias; + const bool is_causal = cfg.is_causal; + const bool is_causal_bottom_right = cfg.is_causal_bottom_right; + const bool is_padding = cfg.is_padding; + const bool is_dropout = cfg.is_dropout; + const bool is_softmax_offset = cfg.is_softmax_offset; + const bool is_mxfp8 = cfg.is_mxfp8; + const bool is_tensor_scaling = cfg.is_tensor_scaling; + const bool is_delayed_scaling = cfg.is_delayed_scaling_bwd; + const bool is_current_scaling = cfg.is_current_scaling_bwd; + const bool is_O_in_F16 = !cfg.is_o_in_fp8; + const bool is_ragged_q = cfg.is_ragged_q; + const bool is_ragged_kv = cfg.is_ragged_kv; + const bool use_ragged_stats = cfg.uses_ragged_stats; + const DType ragged_offset_type = cfg.ragged_offset_type_bwd; + + auto mha_graph = std::make_shared(); + + mha_graph->set_io_data_type(qkv_tensor_type) + .set_intermediate_data_type(fe::DataType_t::FLOAT) + .set_compute_data_type(fe::DataType_t::FLOAT); + + std::shared_ptr offset_q, offset_k, offset_v, offset_o, + offset_stats; + std::shared_ptr Q, Q_t, K, K_t, V, O, dO, dO_t, dO_f16, Stats, + attn_scale; + std::shared_ptr descale_q, descale_q_t, descale_k, descale_k_t, + descale_v; + std::shared_ptr descale_s, descale_o; + std::shared_ptr descale_dP, descale_dO, descale_dO_t; + std::shared_ptr scale_s, scale_dP; + std::shared_ptr scale_dQ, scale_dK, scale_dV; + std::shared_ptr bias, dBias, softmax_offset, d_softmax_offset; + std::shared_ptr seq_q, seq_kv; + std::shared_ptr dropout_seed, dropout_offset; + + // Q, K, V, O, dO, stats, attn_scale + std::vector q_strides(4), k_strides(4), v_strides(4), o_strides(4), dO_strides(4); + generateMatrixStridesWithLayout(b, h, hg, s_q, s_kv, d_qk, d_v, q_strides.data(), + k_strides.data(), v_strides.data(), qkv_layout); + generateMatrixStridesWithFormat(b, h, s_q, d_v, o_strides.data(), o_format); + generateMatrixStridesWithFormat(b, h, s_q, d_v, dO_strides.data(), do_format); + Q = mha_graph->tensor(fe::graph::Tensor_attributes() + .set_name("Q") + .set_dim({b, h, s_q, d_qk}) + .set_stride(q_strides) + .set_data_type(qkv_tensor_type)); + K = mha_graph->tensor(fe::graph::Tensor_attributes() + .set_name("K") + .set_dim({b, hg, s_kv, d_qk}) + .set_stride(k_strides) + .set_data_type(qkv_tensor_type)); + V = mha_graph->tensor(fe::graph::Tensor_attributes() + .set_name("V") + .set_dim({b, hg, s_kv, d_v}) + .set_stride(v_strides) + .set_data_type(qkv_tensor_type)); + O = mha_graph->tensor(fe::graph::Tensor_attributes() + .set_name("O") + .set_dim({b, h, s_q, d_v}) + .set_stride(o_strides) + .set_data_type(o_tensor_type)); + dO = mha_graph->tensor(fe::graph::Tensor_attributes() + .set_name("dO") + .set_dim({b, h, s_q, d_v}) + .set_stride(dO_strides) + .set_data_type(do_tensor_type)); + if (is_ragged_q) { + offset_q = mha_graph->tensor(fe::graph::Tensor_attributes() + .set_name("offset_q") + .set_dim({b + 1, 1, 1, 1}) + .set_stride({1, 1, 1, 1}) + .set_data_type(get_cudnn_fe_dtype(ragged_offset_type))); + offset_o = mha_graph->tensor(fe::graph::Tensor_attributes() + .set_name("offset_o") + .set_dim({b + 1, 1, 1, 1}) + .set_stride({1, 1, 1, 1}) + .set_data_type(get_cudnn_fe_dtype(ragged_offset_type))); + Q->set_ragged_offset(offset_q); + O->set_ragged_offset(offset_o); + dO->set_ragged_offset(offset_o); } - - bool is_O_in_F16 = (o_tensor_type == cudnn_frontend::DataType_t::HALF || - o_tensor_type == cudnn_frontend::DataType_t::BFLOAT16); - - try { - FADescriptor_v1 descriptor{b, - h, - hg, - s_q, - s_kv, - d_qk, - d_v, - 0, - 0, - 0, - 0, - 0, - 0, - bias_b, - bias_h, - bias_sq, - bias_skv, - scaling_factor, - true, - dropout_probability, - qkv_layout, - o_format, - do_format, - dqkv_layout, - qkv_scale_inv_format, - do_scale_inv_format, - bias_type, - mask_type, - softmax_type, - window_size_left, - window_size_right, - bottom_right_diagonal, - deterministic, - qkv_tensor_type, - o_tensor_type, - do_tensor_type, - dqkv_tensor_type, - false}; - - namespace fe = cudnn_frontend; - using graph_and_tensors = - std::tuple, - std::shared_ptr, // Q - std::shared_ptr, // Q_t - std::shared_ptr, // K - std::shared_ptr, // K_t - std::shared_ptr, // V - std::shared_ptr, // O - std::shared_ptr, // Stats - std::shared_ptr, // dO - std::shared_ptr, // dO_t - std::shared_ptr, // dO_f16 - std::shared_ptr, // attn_scale - std::shared_ptr, // descale_q - std::shared_ptr, // descale_q_t - std::shared_ptr, // descale_k - std::shared_ptr, // descale_k_t - std::shared_ptr, // descale_v - std::shared_ptr, // descale_o - std::shared_ptr, // descale_dO - std::shared_ptr, // descale_dO_t - std::shared_ptr, // descale_s - std::shared_ptr, // descale_dP - std::shared_ptr, // scale_dQ - std::shared_ptr, // scale_dK - std::shared_ptr, // scale_dV - std::shared_ptr, // scale_s - std::shared_ptr, // scale_dP - std::shared_ptr, // dQ - std::shared_ptr, // dK - std::shared_ptr, // dV - std::shared_ptr, // amax_dQ - std::shared_ptr, // amax_dK - std::shared_ptr, // amax_dV - std::shared_ptr, // amax_dP - std::shared_ptr, // bias - std::shared_ptr, // dBias - std::shared_ptr, // softmax_offset - std::shared_ptr, // d_softmax_offset - std::shared_ptr, // seq_q - std::shared_ptr, // seq_kv - std::shared_ptr, // offset_q - std::shared_ptr, // offset_k - std::shared_ptr, // offset_v - std::shared_ptr, // offset_o - std::shared_ptr, // offset_stats - std::shared_ptr, // dropout_seed - std::shared_ptr>; // dropout_offset - - using CacheType = std::map; - static thread_local CacheType sdpa_fp8_bprop_cache; - - // Get plan from cache if cache is available, otherwise create one - auto get_graph = [&](CacheType& cache, const FADescriptor_v1& descriptor) -> graph_and_tensors { - // if hit, return - auto it = cache.find(descriptor); - if (it != cache.end()) { - auto graph = it->second; - return graph; - } - - // otherwise, build the op_graph and the plan. Then update cache - auto mha_graph = std::make_shared(); - - mha_graph->set_io_data_type(qkv_tensor_type) - .set_intermediate_data_type(fe::DataType_t::FLOAT) - .set_compute_data_type(fe::DataType_t::FLOAT); - - std::shared_ptr Q, Q_t, K, K_t, V, O, dO, dO_t, dO_f16, Stats, - attn_scale; - std::shared_ptr descale_q, descale_q_t, descale_k, descale_k_t, - descale_v; - std::shared_ptr descale_s, descale_o; - std::shared_ptr descale_dP, descale_dO, descale_dO_t; - std::shared_ptr scale_s, scale_dP; - std::shared_ptr scale_dQ, scale_dK, scale_dV; - std::shared_ptr bias, dBias, softmax_offset, d_softmax_offset; - std::shared_ptr seq_q, seq_kv; - std::shared_ptr offset_q, offset_k, offset_v, offset_o, - offset_stats; - std::shared_ptr dropout_seed, dropout_offset; - - // Q, K, V, O, dO, stats, attn_scale - std::vector q_strides(4), k_strides(4), v_strides(4), o_strides(4), dO_strides(4); - generateMatrixStridesWithLayout(b, h, hg, s_q, s_kv, d_qk, d_v, q_strides.data(), - k_strides.data(), v_strides.data(), qkv_layout); - generateMatrixStridesWithFormat(b, h, s_q, d_v, o_strides.data(), o_format); - generateMatrixStridesWithFormat(b, h, s_q, d_v, dO_strides.data(), do_format); - Q = mha_graph->tensor(fe::graph::Tensor_attributes() - .set_name("Q") - .set_dim({b, h, s_q, d_qk}) - .set_stride(q_strides) - .set_data_type(qkv_tensor_type)); - if (is_ragged_q) { - offset_q = mha_graph->tensor(fe::graph::Tensor_attributes() - .set_name("offset_q") - .set_dim({b + 1, 1, 1, 1}) - .set_stride({1, 1, 1, 1}) - .set_data_type(get_cudnn_fe_dtype(ragged_offset_type))); - offset_o = mha_graph->tensor(fe::graph::Tensor_attributes() - .set_name("offset_o") + if (is_ragged_kv) { + offset_k = mha_graph->tensor(fe::graph::Tensor_attributes() + .set_name("offset_k") + .set_dim({b + 1, 1, 1, 1}) + .set_stride({1, 1, 1, 1}) + .set_data_type(get_cudnn_fe_dtype(ragged_offset_type))); + offset_v = mha_graph->tensor(fe::graph::Tensor_attributes() + .set_name("offset_v") + .set_dim({b + 1, 1, 1, 1}) + .set_stride({1, 1, 1, 1}) + .set_data_type(get_cudnn_fe_dtype(ragged_offset_type))); + K->set_ragged_offset(offset_k); + V->set_ragged_offset(offset_v); + } + if (use_ragged_stats) { + offset_stats = mha_graph->tensor(fe::graph::Tensor_attributes() + .set_name("offset_stats") .set_dim({b + 1, 1, 1, 1}) .set_stride({1, 1, 1, 1}) .set_data_type(get_cudnn_fe_dtype(ragged_offset_type))); - Q->set_ragged_offset(offset_q); - } - K = mha_graph->tensor(fe::graph::Tensor_attributes() - .set_name("K") - .set_dim({b, hg, s_kv, d_qk}) - .set_stride(k_strides) + } + Stats = mha_graph->tensor(fe::graph::Tensor_attributes() + .set_name("Stats") + .set_dim({b, h, s_q, 1}) + .set_data_type(fe::DataType_t::FLOAT)); + if (use_ragged_stats) { + Stats->set_stride({h * s_q, 1, h, 1}).set_ragged_offset(offset_stats); + } else { + Stats->set_stride({h * s_q, s_q, 1, 1}); + } + attn_scale = mha_graph->tensor(fe::graph::Tensor_attributes() + .set_name("attn_scale") + .set_dim({1, 1, 1, 1}) + .set_stride({1, 1, 1, 1}) + .set_is_pass_by_value(true) + .set_data_type(fe::DataType_t::FLOAT)); + + // Descale_q, Descale_k, Descale_v, Descale_s, Scale_s, Descale_dP, Scale_dP, Descale_o, Descale_dO, Scale_dQ, Scale_dK, Scale_dV + if (is_tensor_scaling) { + descale_q = mha_graph->tensor(fe::graph::Tensor_attributes() + .set_name("Descale_q") + .set_dim({1, 1, 1, 1}) + .set_stride({1, 1, 1, 1}) + .set_data_type(fe::DataType_t::FLOAT)); + descale_k = mha_graph->tensor_like(descale_q, "Descale_q"); + descale_v = mha_graph->tensor_like(descale_q, "Descale_v"); + descale_s = mha_graph->tensor_like(descale_q, "Descale_s"); + scale_s = mha_graph->tensor_like(descale_q, "Scale_s"); + descale_dP = mha_graph->tensor_like(descale_q, "Descale_dP"); + scale_dP = mha_graph->tensor_like(descale_q, "Scale_dP"); + if (is_current_scaling && is_O_in_F16) { + descale_o = mha_graph->tensor(1.0f); + } else { + descale_o = mha_graph->tensor_like(descale_q, "Descale_O"); + } + descale_dO = mha_graph->tensor_like(descale_q, "Descale_dO"); + if (is_delayed_scaling) { + scale_dQ = mha_graph->tensor_like(descale_q, "Scale_dQ"); + scale_dK = mha_graph->tensor_like(descale_q, "Scale_dK"); + scale_dV = mha_graph->tensor_like(descale_q, "Scale_dV"); + } + if (is_current_scaling) { + scale_dQ = mha_graph->tensor(1.0f); + scale_dK = mha_graph->tensor(1.0f); + scale_dV = mha_graph->tensor(1.0f); + } + } else if (is_mxfp8) { + const NVTE_QKV_Format q_format = cfg.q_format; + const NVTE_QKV_Format kv_format = cfg.kv_format; + const NVTE_QKV_Format q_scale_inv_format = + (qkv_scale_inv_format != NVTE_QKV_Format_NOT_SET) ? qkv_scale_inv_format : q_format; + const NVTE_QKV_Format kv_scale_inv_format = + (qkv_scale_inv_format != NVTE_QKV_Format_NOT_SET) ? qkv_scale_inv_format : kv_format; + const NVTE_QKV_Format do_scale_format_ = + (do_scale_inv_format != NVTE_QKV_Format_NOT_SET) ? do_scale_inv_format : do_format; + // Q_t, K_t, dO_t, dO_f16 + std::vector q_t_strides(4), k_t_strides(4), dO_t_strides(4); + generateMatrixStridesWithFormat(b, h, s_q, d_qk, q_t_strides.data(), q_format); + generateMatrixStridesWithFormat(b, hg, s_kv, d_qk, k_t_strides.data(), kv_format); + generateMatrixStridesWithFormat(b, h, s_q, d_v, dO_t_strides.data(), do_format); + Q_t = mha_graph->tensor(fe::graph::Tensor_attributes() + .set_name("Q_t") + .set_dim({b, h, s_q, d_qk}) + .set_stride(q_t_strides) .set_data_type(qkv_tensor_type)); - V = mha_graph->tensor(fe::graph::Tensor_attributes() - .set_name("V") - .set_dim({b, hg, s_kv, d_v}) - .set_stride(v_strides) + K_t = mha_graph->tensor(fe::graph::Tensor_attributes() + .set_name("K_t") + .set_dim({b, hg, s_kv, d_qk}) + .set_stride(k_t_strides) .set_data_type(qkv_tensor_type)); - if (is_ragged_kv) { - offset_k = mha_graph->tensor(fe::graph::Tensor_attributes() - .set_name("offset_k") - .set_dim({b + 1, 1, 1, 1}) - .set_stride({1, 1, 1, 1}) - .set_data_type(get_cudnn_fe_dtype(ragged_offset_type))); - offset_v = mha_graph->tensor(fe::graph::Tensor_attributes() - .set_name("offset_v") - .set_dim({b + 1, 1, 1, 1}) - .set_stride({1, 1, 1, 1}) - .set_data_type(get_cudnn_fe_dtype(ragged_offset_type))); - K->set_ragged_offset(offset_k); - V->set_ragged_offset(offset_v); - } - O = mha_graph->tensor(fe::graph::Tensor_attributes() - .set_name("O") - .set_dim({b, h, s_q, d_v}) - .set_stride(o_strides) - .set_data_type(o_tensor_type)); - if (is_ragged_q) { - O->set_ragged_offset(offset_o); - } - dO = mha_graph->tensor(fe::graph::Tensor_attributes() - .set_name("dO") + dO_t = mha_graph->tensor(fe::graph::Tensor_attributes() + .set_name("dO_t") .set_dim({b, h, s_q, d_v}) - .set_stride(dO_strides) + .set_stride(dO_t_strides) .set_data_type(do_tensor_type)); - if (is_ragged_q) { - dO->set_ragged_offset(offset_o); - } - if (use_ragged_stats) { - offset_stats = - mha_graph->tensor(fe::graph::Tensor_attributes() - .set_name("offset_stats") - .set_dim({b + 1, 1, 1, 1}) + dO_f16 = mha_graph->tensor(fe::graph::Tensor_attributes() + .set_name("dO_f16") + .set_dim({b, h, s_q, d_v}) + .set_stride(dO_strides) + .set_data_type(o_tensor_type)); + // Descale_q, Descale_q_t, Descale_k, Descale_k_t, Descale_v, Descale_dO, Descale_dO_t + auto padded = pad_s_d_for_mxfp8(s_q, s_kv, d_qk, d_v); + std::vector q_scale_strides(4), q_t_scale_strides(4), k_scale_strides(4), + k_t_scale_strides(4), v_scale_strides(4), dO_scale_strides(4), dO_t_scale_strides(4); + generateMatrixStridesWithFormat(b, h, padded.s_q_padded, padded.d_qk_scale_padded, + q_scale_strides.data(), q_scale_inv_format); + generateMatrixStridesWithFormat(b, h, padded.s_q_scale_padded, padded.d_qk_padded, + q_t_scale_strides.data(), q_scale_inv_format); + generateMatrixStridesWithFormat(b, hg, padded.s_kv_padded, padded.d_qk_scale_padded, + k_scale_strides.data(), kv_scale_inv_format); + generateMatrixStridesWithFormat(b, hg, padded.s_kv_scale_padded, padded.d_qk_padded, + k_t_scale_strides.data(), kv_scale_inv_format); + generateMatrixStridesWithFormat(b, hg, padded.s_kv_padded, padded.d_v_scale_padded, + v_scale_strides.data(), kv_scale_inv_format); + generateMatrixStridesWithFormat(b, h, padded.s_q_padded, padded.d_v_scale_padded, + dO_scale_strides.data(), do_scale_format_); + generateMatrixStridesWithFormat(b, h, padded.s_q_scale_padded, padded.d_v_padded, + dO_t_scale_strides.data(), do_scale_format_); + descale_q = mha_graph->tensor(fe::graph::Tensor_attributes() + .set_name("Descale_q") + .set_dim({b, h, padded.s_q_padded, padded.d_qk_scale_padded}) + .set_stride(q_scale_strides) + .set_data_type(fe::DataType_t::FP8_E8M0) + .set_reordering_type(fe::TensorReordering_t::F8_128x4)); + descale_q_t = + mha_graph->tensor(fe::graph::Tensor_attributes() + .set_name("Descale_q_t") + .set_dim({b, h, padded.s_q_scale_padded, padded.d_qk_padded}) + .set_stride(q_t_scale_strides) + .set_data_type(fe::DataType_t::FP8_E8M0) + .set_reordering_type(fe::TensorReordering_t::F8_128x4)); + descale_k = + mha_graph->tensor(fe::graph::Tensor_attributes() + .set_name("Descale_k") + .set_dim({b, hg, padded.s_kv_padded, padded.d_qk_scale_padded}) + .set_stride(k_scale_strides) + .set_data_type(fe::DataType_t::FP8_E8M0) + .set_reordering_type(fe::TensorReordering_t::F8_128x4)); + descale_k_t = + mha_graph->tensor(fe::graph::Tensor_attributes() + .set_name("Descale_k_t") + .set_dim({b, hg, padded.s_kv_scale_padded, padded.d_qk_padded}) + .set_stride(k_t_scale_strides) + .set_data_type(fe::DataType_t::FP8_E8M0) + .set_reordering_type(fe::TensorReordering_t::F8_128x4)); + descale_v = mha_graph->tensor(fe::graph::Tensor_attributes() + .set_name("Descale_v") + .set_dim({b, hg, padded.s_kv_padded, padded.d_v_scale_padded}) + .set_stride(v_scale_strides) + .set_data_type(fe::DataType_t::FP8_E8M0) + .set_reordering_type(fe::TensorReordering_t::F8_128x4)); + descale_dO = mha_graph->tensor(fe::graph::Tensor_attributes() + .set_name("Descale_dO") + .set_dim({b, h, padded.s_q_padded, padded.d_v_scale_padded}) + .set_stride(dO_scale_strides) + .set_data_type(fe::DataType_t::FP8_E8M0) + .set_reordering_type(fe::TensorReordering_t::F8_128x4)); + descale_dO_t = + mha_graph->tensor(fe::graph::Tensor_attributes() + .set_name("Descale_dO_t") + .set_dim({b, h, padded.s_q_scale_padded, padded.d_v_padded}) + .set_stride(dO_t_scale_strides) + .set_data_type(fe::DataType_t::FP8_E8M0) + .set_reordering_type(fe::TensorReordering_t::F8_128x4)); + } + + fe::graph::SDPA_fp8_backward_attributes sdpa_backward_options; + sdpa_backward_options = fe::graph::SDPA_fp8_backward_attributes() + .set_name("sdpa_fp8_backward") + .set_causal_mask(is_causal) + .set_attn_scale(attn_scale); + + fe::DiagonalAlignment_t const& diagonal_alignment = bottom_right_diagonal + ? fe::DiagonalAlignment_t::BOTTOM_RIGHT + : fe::DiagonalAlignment_t::TOP_LEFT; + sdpa_backward_options.set_diagonal_alignment(diagonal_alignment); + + if (cudnn_runtime_version >= 92100) { + if (window_size_left != -1) { + sdpa_backward_options.set_diagonal_band_left_bound(window_size_left + 1); + } + if (window_size_right != -1) { + sdpa_backward_options.set_diagonal_band_right_bound(window_size_right); + } + } + if (is_causal_bottom_right) { + sdpa_backward_options.set_diagonal_band_right_bound(0); + } + + // sdpa_backward_options.set_alibi_mask(is_alibi); + + // if (is_bias) { + // bias = mha_graph->tensor(fe::graph::Tensor_attributes() + // .set_name("bias") + // .set_dim({bias_b, bias_h, bias_sq, bias_skv}) + // .set_stride({bias_h * bias_sq * bias_skv, bias_sq * bias_skv, bias_skv, 1})); + // dBias = mha_graph->tensor(fe::graph::Tensor_attributes() + // .set_name("dBias") + // .set_dim({bias_b, bias_h, bias_sq, bias_skv}) + // .set_stride({bias_h * bias_sq * bias_skv, bias_sq * bias_skv, bias_skv, 1})); + // sdpa_backward_options.set_bias(bias); + // bias shapes [1, 1, s, s], [b, 1, s, s], [b, h, s, s], [1, h, s, s] are supported for dbias calculation + // bias shape [1, 1, 1, s] is not supported for dbias calculation as of cuDNN 9.18 + // if (!((bias_b == 1) && (bias_h == 1) && (bias_sq == 1))) { + // sdpa_backward_options.set_dbias(dBias); + // } + // } + + if (cudnn_runtime_version >= 91900) { + sdpa_backward_options.set_deterministic_algorithm(deterministic); + } + + if (is_padding) { + seq_q = mha_graph->tensor(fe::graph::Tensor_attributes() + .set_name("seq_q") + .set_dim({b, 1, 1, 1}) .set_stride({1, 1, 1, 1}) - .set_data_type(get_cudnn_fe_dtype(ragged_offset_type))); - } - Stats = mha_graph->tensor(fe::graph::Tensor_attributes() - .set_name("Stats") - .set_dim({b, h, s_q, 1}) - .set_data_type(fe::DataType_t::FLOAT)); - if (use_ragged_stats) { - Stats->set_stride({h * s_q, 1, h, 1}).set_ragged_offset(offset_stats); - } else { - Stats->set_stride({h * s_q, s_q, 1, 1}); - } - attn_scale = mha_graph->tensor(fe::graph::Tensor_attributes() - .set_name("attn_scale") + .set_data_type(fe::DataType_t::INT32)); + seq_kv = mha_graph->tensor(fe::graph::Tensor_attributes() + .set_name("seq_kv") + .set_dim({b, 1, 1, 1}) + .set_stride({1, 1, 1, 1}) + .set_data_type(fe::DataType_t::INT32)); + sdpa_backward_options.set_padding_mask(is_padding).set_seq_len_q(seq_q).set_seq_len_kv(seq_kv); + } + + if (is_dropout) { + dropout_seed = mha_graph->tensor(fe::graph::Tensor_attributes() + .set_name("Seed") .set_dim({1, 1, 1, 1}) .set_stride({1, 1, 1, 1}) - .set_is_pass_by_value(true) - .set_data_type(fe::DataType_t::FLOAT)); - - // Descale_q, Descale_k, Descale_v, Descale_s, Scale_s, Descale_dP, Scale_dP, Descale_o, Descale_dO, Scale_dQ, Scale_dK, Scale_dV - if (is_delayed_scaling || is_current_scaling) { - descale_q = mha_graph->tensor(fe::graph::Tensor_attributes() - .set_name("Descale_q") - .set_dim({1, 1, 1, 1}) - .set_stride({1, 1, 1, 1}) - .set_data_type(fe::DataType_t::FLOAT)); - descale_k = mha_graph->tensor_like(descale_q, "Descale_q"); - descale_v = mha_graph->tensor_like(descale_q, "Descale_v"); - descale_s = mha_graph->tensor_like(descale_q, "Descale_s"); - scale_s = mha_graph->tensor_like(descale_q, "Scale_s"); - descale_dP = mha_graph->tensor_like(descale_q, "Descale_dP"); - scale_dP = mha_graph->tensor_like(descale_q, "Scale_dP"); - if (is_current_scaling && is_O_in_F16) { - descale_o = mha_graph->tensor(1.0f); - } else { - descale_o = mha_graph->tensor_like(descale_q, "Descale_O"); - } - descale_dO = mha_graph->tensor_like(descale_q, "Descale_dO"); - if (is_delayed_scaling) { - scale_dQ = mha_graph->tensor_like(descale_q, "Scale_dQ"); - scale_dK = mha_graph->tensor_like(descale_q, "Scale_dK"); - scale_dV = mha_graph->tensor_like(descale_q, "Scale_dV"); - } - if (is_current_scaling) { - scale_dQ = mha_graph->tensor(1.0f); - scale_dK = mha_graph->tensor(1.0f); - scale_dV = mha_graph->tensor(1.0f); - } - } else if (is_mxfp8) { - 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 q_scale_inv_format = - (qkv_scale_inv_format != NVTE_QKV_Format_NOT_SET) ? qkv_scale_inv_format : q_format; - NVTE_QKV_Format kv_scale_inv_format = - (qkv_scale_inv_format != NVTE_QKV_Format_NOT_SET) ? qkv_scale_inv_format : kv_format; - NVTE_QKV_Format do_scale_format_ = - (do_scale_inv_format != NVTE_QKV_Format_NOT_SET) ? do_scale_inv_format : do_format; - // Q_t, K_t, dO_t, dO_f16 - std::vector q_t_strides(4), k_t_strides(4), dO_t_strides(4); - generateMatrixStridesWithFormat(b, h, s_q, d_qk, q_t_strides.data(), q_format); - generateMatrixStridesWithFormat(b, hg, s_kv, d_qk, k_t_strides.data(), kv_format); - generateMatrixStridesWithFormat(b, h, s_q, d_v, dO_t_strides.data(), do_format); - Q_t = mha_graph->tensor(fe::graph::Tensor_attributes() - .set_name("Q_t") - .set_dim({b, h, s_q, d_qk}) - .set_stride(q_t_strides) - .set_data_type(qkv_tensor_type)); - K_t = mha_graph->tensor(fe::graph::Tensor_attributes() - .set_name("K_t") - .set_dim({b, hg, s_kv, d_qk}) - .set_stride(k_t_strides) - .set_data_type(qkv_tensor_type)); - dO_t = mha_graph->tensor(fe::graph::Tensor_attributes() - .set_name("dO_t") - .set_dim({b, h, s_q, d_v}) - .set_stride(dO_t_strides) - .set_data_type(do_tensor_type)); - dO_f16 = mha_graph->tensor(fe::graph::Tensor_attributes() - .set_name("dO_f16") - .set_dim({b, h, s_q, d_v}) - .set_stride(dO_strides) - .set_data_type(o_tensor_type)); - // Descale_q, Descale_q_t, Descale_k, Descale_k_t, Descale_v, Descale_dO, Descale_dO_t - auto padded = pad_s_d_for_mxfp8(s_q, s_kv, d_qk, d_v); - std::vector q_scale_strides(4), q_t_scale_strides(4), k_scale_strides(4), - k_t_scale_strides(4), v_scale_strides(4), dO_scale_strides(4), dO_t_scale_strides(4); - generateMatrixStridesWithFormat(b, h, padded.s_q_padded, padded.d_qk_scale_padded, - q_scale_strides.data(), q_scale_inv_format); - generateMatrixStridesWithFormat(b, h, padded.s_q_scale_padded, padded.d_qk_padded, - q_t_scale_strides.data(), q_scale_inv_format); - generateMatrixStridesWithFormat(b, hg, padded.s_kv_padded, padded.d_qk_scale_padded, - k_scale_strides.data(), kv_scale_inv_format); - generateMatrixStridesWithFormat(b, hg, padded.s_kv_scale_padded, padded.d_qk_padded, - k_t_scale_strides.data(), kv_scale_inv_format); - generateMatrixStridesWithFormat(b, hg, padded.s_kv_padded, padded.d_v_scale_padded, - v_scale_strides.data(), kv_scale_inv_format); - generateMatrixStridesWithFormat(b, h, padded.s_q_padded, padded.d_v_scale_padded, - dO_scale_strides.data(), do_scale_format_); - generateMatrixStridesWithFormat(b, h, padded.s_q_scale_padded, padded.d_v_padded, - dO_t_scale_strides.data(), do_scale_format_); - descale_q = - mha_graph->tensor(fe::graph::Tensor_attributes() - .set_name("Descale_q") - .set_dim({b, h, padded.s_q_padded, padded.d_qk_scale_padded}) - .set_stride(q_scale_strides) - .set_data_type(fe::DataType_t::FP8_E8M0) - .set_reordering_type(fe::TensorReordering_t::F8_128x4)); - descale_q_t = - mha_graph->tensor(fe::graph::Tensor_attributes() - .set_name("Descale_q_t") - .set_dim({b, h, padded.s_q_scale_padded, padded.d_qk_padded}) - .set_stride(q_t_scale_strides) - .set_data_type(fe::DataType_t::FP8_E8M0) - .set_reordering_type(fe::TensorReordering_t::F8_128x4)); - descale_k = - mha_graph->tensor(fe::graph::Tensor_attributes() - .set_name("Descale_k") - .set_dim({b, hg, padded.s_kv_padded, padded.d_qk_scale_padded}) - .set_stride(k_scale_strides) - .set_data_type(fe::DataType_t::FP8_E8M0) - .set_reordering_type(fe::TensorReordering_t::F8_128x4)); - descale_k_t = - mha_graph->tensor(fe::graph::Tensor_attributes() - .set_name("Descale_k_t") - .set_dim({b, hg, padded.s_kv_scale_padded, padded.d_qk_padded}) - .set_stride(k_t_scale_strides) - .set_data_type(fe::DataType_t::FP8_E8M0) - .set_reordering_type(fe::TensorReordering_t::F8_128x4)); - descale_v = - mha_graph->tensor(fe::graph::Tensor_attributes() - .set_name("Descale_v") - .set_dim({b, hg, padded.s_kv_padded, padded.d_v_scale_padded}) - .set_stride(v_scale_strides) - .set_data_type(fe::DataType_t::FP8_E8M0) - .set_reordering_type(fe::TensorReordering_t::F8_128x4)); - descale_dO = - mha_graph->tensor(fe::graph::Tensor_attributes() - .set_name("Descale_dO") - .set_dim({b, h, padded.s_q_padded, padded.d_v_scale_padded}) - .set_stride(dO_scale_strides) - .set_data_type(fe::DataType_t::FP8_E8M0) - .set_reordering_type(fe::TensorReordering_t::F8_128x4)); - descale_dO_t = - mha_graph->tensor(fe::graph::Tensor_attributes() - .set_name("Descale_dO_t") - .set_dim({b, h, padded.s_q_scale_padded, padded.d_v_padded}) - .set_stride(dO_t_scale_strides) - .set_data_type(fe::DataType_t::FP8_E8M0) - .set_reordering_type(fe::TensorReordering_t::F8_128x4)); - } - - fe::graph::SDPA_fp8_backward_attributes sdpa_backward_options; - sdpa_backward_options = fe::graph::SDPA_fp8_backward_attributes() - .set_name("sdpa_fp8_backward") - .set_causal_mask(is_causal) - .set_attn_scale(attn_scale); - - fe::DiagonalAlignment_t const& diagonal_alignment = - bottom_right_diagonal ? fe::DiagonalAlignment_t::BOTTOM_RIGHT - : fe::DiagonalAlignment_t::TOP_LEFT; - sdpa_backward_options.set_diagonal_alignment(diagonal_alignment); - - if (cudnn_runtime_version >= 92100) { - if (window_size_left != -1) { - sdpa_backward_options.set_diagonal_band_left_bound(window_size_left + 1); - } - if (window_size_right != -1) { - sdpa_backward_options.set_diagonal_band_right_bound(window_size_right); - } - } + .set_data_type(fe::DataType_t::INT64)); + dropout_offset = mha_graph->tensor(fe::graph::Tensor_attributes() + .set_name("Offset") + .set_dim({1, 1, 1, 1}) + .set_stride({1, 1, 1, 1}) + .set_data_type(fe::DataType_t::INT64)); + sdpa_backward_options.set_dropout(dropout_probability, dropout_seed, dropout_offset); + } - // sdpa_backward_options.set_alibi_mask(is_alibi); - - // if (is_bias) { - // bias = mha_graph->tensor(fe::graph::Tensor_attributes() - // .set_name("bias") - // .set_dim({bias_b, bias_h, bias_sq, bias_skv}) - // .set_stride({bias_h * bias_sq * bias_skv, bias_sq * bias_skv, bias_skv, 1})); - // dBias = mha_graph->tensor(fe::graph::Tensor_attributes() - // .set_name("dBias") - // .set_dim({bias_b, bias_h, bias_sq, bias_skv}) - // .set_stride({bias_h * bias_sq * bias_skv, bias_sq * bias_skv, bias_skv, 1})); - // sdpa_backward_options.set_bias(bias); - // bias shapes [1, 1, s, s], [b, 1, s, s], [b, h, s, s], [1, h, s, s] are supported for dbias calculation - // bias shape [1, 1, 1, s] is not supported for dbias calculation as of cuDNN 9.18 - // if (!((bias_b == 1) && (bias_h == 1) && (bias_sq == 1))) { - // sdpa_backward_options.set_dbias(dBias); - // } - // } - - if (cudnn_runtime_version >= 91900) { - sdpa_backward_options.set_deterministic_algorithm(deterministic); - } + if (is_softmax_offset) { + softmax_offset = mha_graph->tensor(fe::graph::Tensor_attributes() + .set_name("softmax_offset") + .set_dim({1, h, 1, 1}) + .set_stride({h, 1, 1, 1}) + .set_data_type(fe::DataType_t::FLOAT)); + sdpa_backward_options.set_sink_token(softmax_offset); + d_softmax_offset = mha_graph->tensor(fe::graph::Tensor_attributes() + .set_name("d_softmax_offset") + .set_dim({1, h, 1, 1}) + .set_stride({h, 1, 1, 1}) + .set_data_type(fe::DataType_t::FLOAT)); + sdpa_backward_options.set_dsink_token(d_softmax_offset); + } - if (is_padding) { - seq_q = mha_graph->tensor(fe::graph::Tensor_attributes() - .set_name("seq_q") - .set_dim({b, 1, 1, 1}) - .set_stride({1, 1, 1, 1}) - .set_data_type(fe::DataType_t::INT32)); - seq_kv = mha_graph->tensor(fe::graph::Tensor_attributes() - .set_name("seq_kv") - .set_dim({b, 1, 1, 1}) - .set_stride({1, 1, 1, 1}) - .set_data_type(fe::DataType_t::INT32)); - sdpa_backward_options.set_padding_mask(is_padding) - .set_seq_len_q(seq_q) - .set_seq_len_kv(seq_kv); - } + std::shared_ptr dQ, dK, dV, amax_dQ, amax_dK, amax_dV, amax_dP; + if (is_tensor_scaling) { + std::tie(dQ, dK, dV, amax_dQ, amax_dK, amax_dV, amax_dP) = + std::apply([](const auto&... elems) { return std::make_tuple(elems...); }, + mha_graph->sdpa_fp8_backward(Q, K, V, O, dO, Stats, descale_q, descale_k, + descale_v, descale_o, descale_dO, descale_s, + descale_dP, scale_s, scale_dQ, scale_dK, scale_dV, + scale_dP, sdpa_backward_options)); + } else if (is_mxfp8) { + std::tie(dQ, dK, dV, amax_dQ, amax_dK, amax_dV) = std::apply( + [](const auto&... elems) { return std::make_tuple(elems...); }, + mha_graph->sdpa_fp8_backward(Q, Q_t, K, K_t, V, O, dO_f16, dO, dO_t, Stats, descale_q, + descale_q_t, descale_k, descale_k_t, descale_v, descale_dO, + descale_dO_t, sdpa_backward_options)); + } + std::vector dq_strides(4), dk_strides(4), dv_strides(4); + generateMatrixStridesWithLayout(b, h, hg, s_q, s_kv, d_qk, d_v, dq_strides.data(), + dk_strides.data(), dv_strides.data(), dqkv_layout); + dQ->set_output(true) + .set_dim({b, h, s_q, d_qk}) + .set_stride(dq_strides) + .set_data_type(dqkv_tensor_type); + dK->set_output(true) + .set_dim({b, hg, s_kv, d_qk}) + .set_stride(dk_strides) + .set_data_type(dqkv_tensor_type); + dV->set_output(true) + .set_dim({b, hg, s_kv, d_v}) + .set_stride(dv_strides) + .set_data_type(dqkv_tensor_type); + if (is_ragged_q) { + dQ->set_ragged_offset(offset_q); + } + if (is_ragged_kv) { + dK->set_ragged_offset(offset_k); + dV->set_ragged_offset(offset_v); + } + amax_dQ->set_output(!is_mxfp8) + .set_dim({1, 1, 1, 1}) + .set_stride({1, 1, 1, 1}) + .set_data_type(fe::DataType_t::FLOAT); + amax_dK->set_output(!is_mxfp8) + .set_dim({1, 1, 1, 1}) + .set_stride({1, 1, 1, 1}) + .set_data_type(fe::DataType_t::FLOAT); + amax_dV->set_output(!is_mxfp8) + .set_dim({1, 1, 1, 1}) + .set_stride({1, 1, 1, 1}) + .set_data_type(fe::DataType_t::FLOAT); + if (is_tensor_scaling) { + amax_dP->set_output(true) + .set_dim({1, 1, 1, 1}) + .set_stride({1, 1, 1, 1}) + .set_data_type(fe::DataType_t::FLOAT); + } - if (is_dropout) { - dropout_seed = mha_graph->tensor(fe::graph::Tensor_attributes() - .set_name("Seed") - .set_dim({1, 1, 1, 1}) - .set_stride({1, 1, 1, 1}) - .set_data_type(fe::DataType_t::INT64)); - dropout_offset = mha_graph->tensor(fe::graph::Tensor_attributes() - .set_name("Offset") - .set_dim({1, 1, 1, 1}) - .set_stride({1, 1, 1, 1}) - .set_data_type(fe::DataType_t::INT64)); - sdpa_backward_options.set_dropout(dropout_probability, dropout_seed, dropout_offset); - } + std::tuple, // Q + std::shared_ptr, // K + std::shared_ptr, // V + std::shared_ptr, // O + std::shared_ptr, // Stats + std::shared_ptr, // dO + std::shared_ptr, // attn_scale + std::shared_ptr, // descale_q + std::shared_ptr, // descale_k + std::shared_ptr, // descale_v + std::shared_ptr, // descale_o + std::shared_ptr, // descale_dO + std::shared_ptr, // descale_s + std::shared_ptr, // descale_dP + std::shared_ptr, // scale_dQ + std::shared_ptr, // scale_dK + std::shared_ptr, // scale_dV + std::shared_ptr, // scale_s + std::shared_ptr, // scale_dP + std::shared_ptr, // dQ + std::shared_ptr, // dK + std::shared_ptr, // dV + std::shared_ptr, // amax_dQ + std::shared_ptr, // amax_dK + std::shared_ptr, // amax_dV + std::shared_ptr> // amax_dP + key_tensors_tuple = + std::make_tuple(Q, K, V, O, Stats, dO, attn_scale, descale_q, descale_k, descale_v, + descale_o, descale_dO, descale_s, descale_dP, scale_s, scale_dQ, scale_dK, + scale_dV, scale_dP, dQ, dK, dV, amax_dQ, amax_dK, amax_dV, amax_dP); + auto mxfp8_tensors_tuple = + is_mxfp8 ? std::make_tuple(Q_t, K_t, dO_f16, dO_t, descale_q_t, descale_k_t, descale_dO_t) + : std::make_tuple(nullptr, nullptr, nullptr, nullptr, nullptr, nullptr, nullptr); + auto bias_tuple = is_bias ? std::make_tuple(bias, dBias) : std::make_tuple(nullptr, nullptr); + auto softmax_offset_tuple = is_softmax_offset ? std::make_tuple(softmax_offset, d_softmax_offset) + : std::make_tuple(nullptr, nullptr); + auto padding_tuple = + is_padding ? std::make_tuple(seq_q, seq_kv) : std::make_tuple(nullptr, nullptr); + auto offset_qo_tuple = + is_ragged_q ? std::make_tuple(offset_q, offset_o) : std::make_tuple(nullptr, nullptr); + auto offset_kv_tuple = + is_ragged_kv ? std::make_tuple(offset_k, offset_v) : std::make_tuple(nullptr, nullptr); + auto offset_s_tuple = use_ragged_stats ? std::make_tuple(offset_stats) : std::make_tuple(nullptr); + auto dropout_tuple = is_dropout ? std::make_tuple(dropout_seed, dropout_offset) + : std::make_tuple(nullptr, nullptr); + + return std::tuple_cat(std::make_tuple(mha_graph), key_tensors_tuple, mxfp8_tensors_tuple, + bias_tuple, softmax_offset_tuple, padding_tuple, offset_qo_tuple, + offset_kv_tuple, offset_s_tuple, dropout_tuple); +} - if (is_softmax_offset) { - softmax_offset = mha_graph->tensor(fe::graph::Tensor_attributes() - .set_name("softmax_offset") - .set_dim({1, h, 1, 1}) - .set_stride({h, 1, 1, 1}) - .set_data_type(fe::DataType_t::FLOAT)); - sdpa_backward_options.set_sink_token(softmax_offset); - d_softmax_offset = mha_graph->tensor(fe::graph::Tensor_attributes() - .set_name("d_softmax_offset") - .set_dim({1, h, 1, 1}) - .set_stride({h, 1, 1, 1}) - .set_data_type(fe::DataType_t::FLOAT)); - sdpa_backward_options.set_dsink_token(d_softmax_offset); - } +void fused_attn_fp8_bwd_impl( + const FusedAttnConfig& cfg, void* devPtrQ, void* devPtrK, void* devPtrV, void* devPtrM, + void* devPtrO, void* devPtrdO, void* devPtrSoftmaxOffset, void* devPtrdQ, void* devPtrdK, + void* devPtrdV, void* devPtrdSoftmaxOffset, void* devPtrDescaleQ, void* devPtrDescaleK, + void* devPtrDescaleV, void* devPtrDescaleO, void* devPtrDescaledO, void* devPtrDescaleS, + void* devPtrDescaledP, void* devPtrScaleS, void* devPtrScaledP, void* devPtrScaledQ, + void* devPtrScaledK, void* devPtrScaledV, void* devPtrAmaxdP, void* devPtrAmaxdQ, + void* devPtrAmaxdK, void* devPtrAmaxdV, void* devPtrQ_t, void* devPtrK_t, void* devPtrdO_f16, + void* devPtrdO_t, void* devPtrDescaleQ_t, void* devPtrDescaleK_t, void* devPtrDescaledO_t, + void* devPtrcuSeqlensQ, void* devPtrcuSeqlensKV, void* devPtrSeqOffsetsQ, + void* devPtrSeqOffsetsKV, void* devPtrDropoutSeed, void* devPtrDropoutOffset, void* workspace, + size_t* workspace_size, cudaStream_t stream, cudnnHandle_t handle) { + using namespace transformer_engine; - std::shared_ptr dQ, dK, dV, amax_dQ, amax_dK, amax_dV, amax_dP; - if (is_delayed_scaling || is_current_scaling) { - std::tie(dQ, dK, dV, amax_dQ, amax_dK, amax_dV, amax_dP) = - std::apply([](const auto&... elems) { return std::make_tuple(elems...); }, - mha_graph->sdpa_fp8_backward(Q, K, V, O, dO, Stats, descale_q, descale_k, - descale_v, descale_o, descale_dO, descale_s, - descale_dP, scale_s, scale_dQ, scale_dK, - scale_dV, scale_dP, sdpa_backward_options)); - } else if (is_mxfp8) { - std::tie(dQ, dK, dV, amax_dQ, amax_dK, amax_dV) = std::apply( - [](const auto&... elems) { return std::make_tuple(elems...); }, - mha_graph->sdpa_fp8_backward(Q, Q_t, K, K_t, V, O, dO_f16, dO, dO_t, Stats, descale_q, - descale_q_t, descale_k, descale_k_t, descale_v, descale_dO, - descale_dO_t, sdpa_backward_options)); - } - std::vector dq_strides(4), dk_strides(4), dv_strides(4); - generateMatrixStridesWithLayout(b, h, hg, s_q, s_kv, d_qk, d_v, dq_strides.data(), - dk_strides.data(), dv_strides.data(), dqkv_layout); - dQ->set_output(true) - .set_dim({b, h, s_q, d_qk}) - .set_stride(dq_strides) - .set_data_type(dqkv_tensor_type); - if (is_ragged_q) { - dQ->set_ragged_offset(offset_q); - } - dK->set_output(true) - .set_dim({b, hg, s_kv, d_qk}) - .set_stride(dk_strides) - .set_data_type(dqkv_tensor_type); - dV->set_output(true) - .set_dim({b, hg, s_kv, d_v}) - .set_stride(dv_strides) - .set_data_type(dqkv_tensor_type); - if (is_ragged_kv) { - dK->set_ragged_offset(offset_k); - dV->set_ragged_offset(offset_v); - } - amax_dQ->set_output(!is_mxfp8) - .set_dim({1, 1, 1, 1}) - .set_stride({1, 1, 1, 1}) - .set_data_type(fe::DataType_t::FLOAT); - amax_dK->set_output(!is_mxfp8) - .set_dim({1, 1, 1, 1}) - .set_stride({1, 1, 1, 1}) - .set_data_type(fe::DataType_t::FLOAT); - amax_dV->set_output(!is_mxfp8) - .set_dim({1, 1, 1, 1}) - .set_stride({1, 1, 1, 1}) - .set_data_type(fe::DataType_t::FLOAT); - if (is_delayed_scaling || is_current_scaling) { - amax_dP->set_output(true) - .set_dim({1, 1, 1, 1}) - .set_stride({1, 1, 1, 1}) - .set_data_type(fe::DataType_t::FLOAT); - } + cfg.check_derived(); + + const bool is_mxfp8 = cfg.is_mxfp8; + const bool is_tensor_scaling = cfg.is_tensor_scaling; + const bool is_delayed_scaling = cfg.is_delayed_scaling_bwd; + const bool is_current_scaling = cfg.is_current_scaling_bwd; + const bool is_O_in_F16 = !cfg.is_o_in_fp8; + + const int64_t b = static_cast(cfg.graph_batch_size_bwd); + const int64_t actual_b = static_cast(cfg.batch_size); + float scaling_factor = cfg.attn_scale; + const bool is_padding = cfg.is_padding; + const bool is_dropout = cfg.is_dropout; + const bool is_softmax_offset = cfg.is_softmax_offset; + const bool is_ragged_q = cfg.is_ragged_q; + const bool is_ragged_kv = cfg.is_ragged_kv; + const bool use_ragged_stats = cfg.uses_ragged_stats; + const DType ragged_offset_type = cfg.ragged_offset_type_bwd; - std::tuple, // Q - std::shared_ptr, // K - std::shared_ptr, // V - std::shared_ptr, // O - std::shared_ptr, // Stats - std::shared_ptr, // dO - std::shared_ptr, // attn_scale - std::shared_ptr, // descale_q - std::shared_ptr, // descale_k - std::shared_ptr, // descale_v - std::shared_ptr, // descale_o - std::shared_ptr, // descale_dO - std::shared_ptr, // descale_s - std::shared_ptr, // descale_dP - std::shared_ptr, // scale_dQ - std::shared_ptr, // scale_dK - std::shared_ptr, // scale_dV - std::shared_ptr, // scale_s - std::shared_ptr, // scale_dP - std::shared_ptr, // dQ - std::shared_ptr, // dK - std::shared_ptr, // dV - std::shared_ptr, // amax_dQ - std::shared_ptr, // amax_dK - std::shared_ptr, // amax_dV - std::shared_ptr> // amax_dP - key_tensors_tuple = std::make_tuple( - Q, K, V, O, Stats, dO, attn_scale, descale_q, descale_k, descale_v, descale_o, - descale_dO, descale_s, descale_dP, scale_s, scale_dQ, scale_dK, scale_dV, scale_dP, - dQ, dK, dV, amax_dQ, amax_dK, amax_dV, amax_dP); - auto mxfp8_tensors_tuple = - is_mxfp8 ? std::make_tuple(Q_t, K_t, dO_f16, dO_t, descale_q_t, descale_k_t, descale_dO_t) - : std::make_tuple(nullptr, nullptr, nullptr, nullptr, nullptr, nullptr, nullptr); - auto bias_tuple = is_bias ? std::make_tuple(bias, dBias) : std::make_tuple(nullptr, nullptr); - auto softmax_offset_tuple = is_softmax_offset - ? std::make_tuple(softmax_offset, d_softmax_offset) - : std::make_tuple(nullptr, nullptr); - auto padding_tuple = - is_padding ? std::make_tuple(seq_q, seq_kv) : std::make_tuple(nullptr, nullptr); - auto offset_q_tuple = is_ragged_q ? std::make_tuple(offset_q) : std::make_tuple(nullptr); - auto offset_kv_tuple = - is_ragged_kv ? std::make_tuple(offset_k, offset_v) : std::make_tuple(nullptr, nullptr); - auto offset_o_tuple = is_ragged_q ? std::make_tuple(offset_o) : std::make_tuple(nullptr); - auto offset_s_tuple = - use_ragged_stats ? std::make_tuple(offset_stats) : std::make_tuple(nullptr); - auto dropout_tuple = is_dropout ? std::make_tuple(dropout_seed, dropout_offset) - : std::make_tuple(nullptr, nullptr); - - NVTE_CHECK_CUDNN_FE(mha_graph->validate()); - NVTE_CHECK_CUDNN_FE(mha_graph->build_operation_graph(handle)); - NVTE_CHECK_CUDNN_FE(mha_graph->create_execution_plans({fe::HeurMode_t::A})); - NVTE_CHECK_CUDNN_FE(mha_graph->check_support(handle)); - NVTE_CHECK_CUDNN_FE(mha_graph->build_plans(handle)); - - auto return_tuple = - std::tuple_cat(std::make_tuple(mha_graph), key_tensors_tuple, mxfp8_tensors_tuple, - bias_tuple, softmax_offset_tuple, padding_tuple, offset_q_tuple, - offset_kv_tuple, offset_o_tuple, offset_s_tuple, dropout_tuple); - cache.insert({descriptor, return_tuple}); - - return return_tuple; - }; + try { + auto cache_entry = get_graph(cfg, handle); auto [mha_graph, Q, K, V, O, Stats, dO, attn_scale, descale_q, descale_k, descale_v, descale_o, descale_dO, descale_s, descale_dP, scale_s, scale_dQ, scale_dK, scale_dV, scale_dP, dQ, dK, dV, amax_dQ, amax_dK, amax_dV, amax_dP, Q_t, K_t, dO_f16, dO_t, descale_q_t, descale_k_t, descale_dO_t, bias, dBias, softmax_offset, d_softmax_offset, seq_q, seq_kv, - offset_q, offset_k, offset_v, offset_o, offset_stats, dropout_seed, dropout_offset] = - get_graph(sdpa_fp8_bprop_cache, descriptor); + offset_q, offset_o, offset_k, offset_v, offset_stats, dropout_seed, dropout_offset] = + cache_entry->graph_and_tensors; + // This graph is going to be used, so finish the build the cache deferred. + build_plans(Backend::FP8, Pass::Bwd, *cache_entry); + + // Exit to request upper level API to allocate memory if needed auto plan_workspace_size = alignTo<16>(mha_graph->get_workspace_size()); const size_t num_bytes_per_seqlen = alignTo<16>(b * sizeof(int32_t)); - const size_t actual_seqlen_workspace_size = is_padding ? 2 * num_bytes_per_seqlen : 0; const size_t num_bytes_per_ragged_offset = alignTo<16>(((b + 1) * typeToNumBits(ragged_offset_type)) / 8); + const size_t actual_seqlen_workspace_size = 2 * num_bytes_per_seqlen; size_t seqlen_offsets_workspace_size = 0; if (is_ragged_q || is_ragged_kv) { - size_t count = 2 * (static_cast(is_ragged_q) + static_cast(is_ragged_kv)); - if (use_ragged_stats) { - seqlen_offsets_workspace_size = (count + 1) * num_bytes_per_ragged_offset; - } else { - seqlen_offsets_workspace_size = count * num_bytes_per_ragged_offset; - } + const size_t count = + 2 * (static_cast(is_ragged_q) + static_cast(is_ragged_kv)); + seqlen_offsets_workspace_size = + (use_ragged_stats ? count + 1 : count) * num_bytes_per_ragged_offset; } - if (workspace == nullptr) { *workspace_size = plan_workspace_size + actual_seqlen_workspace_size + seqlen_offsets_workspace_size; return; } - // cuDNN stream check needs to be moved here to support dummy kernel calls with // null streams for sizing the cuDNN workspace. NVTE_CHECK_CUDNN(cudnnSetStream(handle, stream)); @@ -1308,7 +1178,7 @@ void fused_attn_fp8_bwd_impl( {dK, devPtrdK}, {dV, devPtrdV}, }; - if (is_delayed_scaling || is_current_scaling) { + if (is_tensor_scaling) { variant_pack[descale_s] = devPtrDescaleS; variant_pack[descale_dP] = devPtrDescaledP; variant_pack[scale_s] = devPtrScaleS; @@ -1336,9 +1206,9 @@ void fused_attn_fp8_bwd_impl( variant_pack[descale_dO_t] = devPtrDescaledO_t; } - /* if (is_bias) { + /* if (cfg.is_bias) { variant_pack[bias] = devPtrBias; - if ((bias_b == 1) && (bias_h == h)) { + if ((bias_b == 1) && (bias_h == cfg.num_attn_heads)) { variant_pack[dBias] = devPtrdBias; } else { variant_pack[dBias] = nullptr; @@ -1383,9 +1253,8 @@ void fused_attn_fp8_bwd_impl( (static_cast(is_ragged_q) + static_cast(is_ragged_kv)) * 2 * num_bytes_per_ragged_offset; } - const RaggedOffsetMultipliers offset_mults(layout_group, h, hg, d_qk, d_v); cu_seqlens_padded_to_offsets<<>>( - offset_mults, actual_b, b, static_cast(devPtrSeqOffsetsQ), + cfg.ragged_offset_mults, actual_b, b, static_cast(devPtrSeqOffsetsQ), static_cast(devPtrSeqOffsetsKV), ragged_offset_type, devOffsetsQ, devOffsetsK, devOffsetsV, devOffsetsO, devOffsetsS); NVTE_CHECK_CUDA(cudaGetLastError()); @@ -1413,27 +1282,32 @@ void fused_attn_fp8_bwd_impl( } NVTE_CHECK_CUDNN_FE(mha_graph->execute(handle, variant_pack, workspace)); + graph_cache_debug::record_execute(Backend::FP8, Pass::Bwd); } catch (cudnn_frontend::cudnnException& e) { NVTE_ERROR(e.what()); } -} // NOLINT(readability/fn_size) +} } // namespace fused_attn +using namespace transformer_engine::fused_attn; + // fused attention FWD FP8 with separate Q, K, V -void fused_attn_fp8_fwd( - size_t batch, size_t num_attn_heads, size_t num_gqa_groups, size_t max_seqlen_q, - size_t max_seqlen_kv, size_t head_dim_qk, size_t head_dim_v, size_t num_tokens_q, - size_t num_tokens_kv, bool is_training, float attn_scale, float p_dropout, - NVTE_QKV_Layout qkv_layout, NVTE_QKV_Format o_format, NVTE_QKV_Format qkv_scale_inv_format, - NVTE_Bias_Type bias_type, NVTE_Mask_Type mask_type, NVTE_Softmax_Type softmax_type, - size_t window_size_left, size_t window_size_right, bool bottom_right_diagonal, - const Tensor* input_Q, const Tensor* input_K, const Tensor* input_V, - const Tensor* input_SoftmaxOffset, Tensor* input_output_S, Tensor* output_O, - NVTETensorPack* Aux_CTX_Tensors, const Tensor* cu_seqlens_q, const Tensor* cu_seqlens_kv, - const Tensor* cu_seqlens_q_padded, const Tensor* cu_seqlens_kv_padded, const Tensor* rng_state, - Tensor* workspace, cudaStream_t stream, cudnnHandle_t handle) { +void fused_attn_fp8_fwd(const FusedAttnConfig& cfg, const Tensor* input_Q, const Tensor* input_K, + const Tensor* input_V, const Tensor* input_SoftmaxOffset, + Tensor* input_output_S, Tensor* output_O, NVTETensorPack* Aux_CTX_Tensors, + const Tensor* cu_seqlens_q, const Tensor* cu_seqlens_kv, + const Tensor* cu_seqlens_q_padded, const Tensor* cu_seqlens_kv_padded, + const Tensor* rng_state, Tensor* workspace, cudaStream_t stream, + cudnnHandle_t handle) { using namespace transformer_engine; + + const size_t batch = cfg.batch_size; + const size_t num_attn_heads = cfg.num_attn_heads; + const size_t max_seqlen_q = cfg.max_seqlen_q; + const size_t num_tokens_q = cfg.num_tokens_q; + const NVTE_Softmax_Type softmax_type = cfg.softmax_type; + void *devPtrQ = nullptr, *devPtrK = nullptr, *devPtrV = nullptr; void *devPtrDescaleQ = nullptr, *devPtrDescaleK = nullptr, *devPtrDescaleV = nullptr; void *devPtrO = nullptr, *devPtrAmaxO = nullptr, *devPtrScaleO = nullptr; @@ -1459,35 +1333,12 @@ void fused_attn_fp8_fwd( if (softmax_type != NVTE_VANILLA_SOFTMAX) { devPtrSoftmaxOffset = input_SoftmaxOffset->data.dptr; } - NVTE_QKV_Format q_format = nvte_get_q_format(qkv_layout); - NVTE_QKV_Format kv_format = nvte_get_kv_format(qkv_layout); - const auto cudnn_runtime_version = cudnnGetVersion(); - const int sm_arch_ = cuda::sm_arch(cuda::current_device()); - - void* devPtrSeqOffsetsQ = cu_seqlens_q_padded->data.dptr; - void* devPtrSeqOffsetsKV = cu_seqlens_kv_padded->data.dptr; - - size_t max_batch_size = 0; - size_t max_tokens_q = 0; - size_t max_tokens_kv = 0; - if (q_format == NVTE_QKV_Format::NVTE_THD || kv_format == NVTE_QKV_Format::NVTE_THD) { - max_batch_size = fused_attn::get_max_batch_size(batch); - } - if (q_format == NVTE_QKV_Format::NVTE_THD) { - max_tokens_q = fused_attn::get_max_tokens(num_tokens_q); - } - if (kv_format == NVTE_QKV_Format::NVTE_THD) { - max_tokens_kv = fused_attn::get_max_tokens(num_tokens_kv); - } - void* devPtrM = nullptr; if (Aux_CTX_Tensors->size == 0) { int i = 0; Tensor* output_M = convertNVTETensorCheck(Aux_CTX_Tensors->tensors[i++]); output_M->data.dptr = nullptr; - // SM120 uses dense stats in the graph, so its allocation must remain [b, h, s_q, 1]. - if (q_format == NVTE_QKV_Format::NVTE_THD && - cudnn_runtime_version >= fused_attn::kFP8THDRaggedCudnnVersion && sm_arch_ != 120) { + if (cfg.uses_ragged_stats) { output_M->data.shape = {num_tokens_q, num_attn_heads, 1}; } else { output_M->data.shape = {batch, num_attn_heads, max_seqlen_q, 1}; @@ -1522,31 +1373,21 @@ void fused_attn_fp8_fwd( reinterpret_cast(reinterpret_cast(cu_seqlens_q->data.dptr)); void* devPtrcuSeqlensKV = reinterpret_cast(reinterpret_cast(cu_seqlens_kv->data.dptr)); + void* devPtrSeqOffsetsQ = cu_seqlens_q_padded->data.dptr; + void* devPtrSeqOffsetsKV = cu_seqlens_kv_padded->data.dptr; void* devPtrDropoutSeed = reinterpret_cast(reinterpret_cast(rng_state->data.dptr)); void* devPtrDropoutOffset = reinterpret_cast(reinterpret_cast(rng_state->data.dptr) + 1); - const DType QKV_type = input_Q->data.dtype; - const DType O_type = output_O->data.dtype; size_t workspace_size = 0; - NVTE_QKV_Format qkv_format = nvte_get_qkv_format(qkv_layout); - if ((qkv_format == NVTE_QKV_Format::NVTE_BSHD) || (qkv_format == NVTE_QKV_Format::NVTE_SBHD) || - (qkv_format == NVTE_QKV_Format::NVTE_BHSD) || (qkv_format == NVTE_QKV_Format::NVTE_THD)) { - fused_attn::fused_attn_fp8_fwd_impl( - batch, num_attn_heads, num_gqa_groups, max_seqlen_q, max_seqlen_kv, head_dim_qk, head_dim_v, - max_batch_size, max_tokens_q, max_tokens_kv, is_training, attn_scale, p_dropout, qkv_layout, - o_format, bias_type, mask_type, softmax_type, window_size_left, window_size_right, - bottom_right_diagonal, devPtrQ, devPtrK, devPtrV, devPtrSoftmaxOffset, devPtrM, devPtrO, - devPtrDescaleQ, devPtrDescaleK, devPtrDescaleV, devPtrDescaleS, devPtrScaleS, devPtrScaleO, - devPtrAmaxO, devPtrAmaxS, devPtrcuSeqlensQ, devPtrcuSeqlensKV, devPtrSeqOffsetsQ, - devPtrSeqOffsetsKV, devPtrDropoutSeed, devPtrDropoutOffset, get_cudnn_fe_dtype(QKV_type), - get_cudnn_fe_dtype(O_type), input_Q->scaling_mode, qkv_scale_inv_format, - workspace->data.dptr, &workspace_size, stream, handle); - } else { - NVTE_ERROR("FP8 fused attention only supports qkv_format=BSHD, SBHD, BHSD, or THD.\n"); - } + fused_attn::fused_attn_fp8_fwd_impl( + cfg, devPtrQ, devPtrK, devPtrV, devPtrSoftmaxOffset, devPtrM, devPtrO, devPtrDescaleQ, + devPtrDescaleK, devPtrDescaleV, devPtrDescaleS, devPtrScaleS, devPtrScaleO, devPtrAmaxO, + devPtrAmaxS, devPtrcuSeqlensQ, devPtrcuSeqlensKV, devPtrSeqOffsetsQ, devPtrSeqOffsetsKV, + devPtrDropoutSeed, devPtrDropoutOffset, workspace->data.dptr, &workspace_size, stream, + handle); if (workspace_size > 0) { if (workspace->data.dptr == nullptr) { @@ -1561,23 +1402,19 @@ void fused_attn_fp8_fwd( } } // fused attention BWD FP8 with separate Q, K, V -void fused_attn_fp8_bwd( - size_t batch, size_t num_attn_heads, size_t num_gqa_groups, size_t max_seqlen_q, - size_t max_seqlen_kv, size_t head_dim_qk, size_t head_dim_v, size_t num_tokens_q, - size_t num_tokens_kv, float attn_scale, float p_dropout, NVTE_QKV_Layout qkv_layout, - NVTE_QKV_Format o_format, NVTE_QKV_Format do_format, NVTE_QKV_Layout dqkv_layout, - NVTE_QKV_Format qkv_scale_inv_format, NVTE_QKV_Format do_scale_inv_format, - NVTE_Bias_Type bias_type, NVTE_Mask_Type mask_type, NVTE_Softmax_Type softmax_type, - size_t window_size_left, size_t window_size_right, bool bottom_right_diagonal, - bool deterministic, const Tensor* input_Q, const Tensor* input_K, const Tensor* input_V, - const Tensor* input_O, const Tensor* input_dO, const Tensor* input_dO_f16, - const Tensor* input_M, const Tensor* input_S, const Tensor* input_SoftmaxOffset, - Tensor* input_output_dP, const Tensor* output_dQ, const Tensor* output_dK, - const Tensor* output_dV, Tensor* output_dSoftmaxOffset, const Tensor* cu_seqlens_q, - const Tensor* cu_seqlens_kv, const Tensor* cu_seqlens_q_padded, - const Tensor* cu_seqlens_kv_padded, const Tensor* rng_state, Tensor* workspace, - cudaStream_t stream, cudnnHandle_t handle) { +void fused_attn_fp8_bwd(const FusedAttnConfig& cfg, const Tensor* input_Q, const Tensor* input_K, + const Tensor* input_V, const Tensor* input_O, const Tensor* input_dO, + const Tensor* input_dO_f16, const Tensor* input_M, const Tensor* input_S, + const Tensor* input_SoftmaxOffset, Tensor* input_output_dP, + const Tensor* output_dQ, const Tensor* output_dK, const Tensor* output_dV, + Tensor* output_dSoftmaxOffset, const Tensor* cu_seqlens_q, + const Tensor* cu_seqlens_kv, const Tensor* cu_seqlens_q_padded, + const Tensor* cu_seqlens_kv_padded, const Tensor* rng_state, + Tensor* workspace, cudaStream_t stream, cudnnHandle_t handle) { using namespace transformer_engine; + + const NVTE_Softmax_Type softmax_type = cfg.softmax_type; + void* devPtrQ = input_Q->data.dptr; void* devPtrK = input_K->data.dptr; void* devPtrV = input_V->data.dptr; @@ -1593,8 +1430,8 @@ void fused_attn_fp8_bwd( devPtrDescaleK_t = input_K->columnwise_scale_inv.dptr; } - void* devPtrO = input_O->data.dptr; const DType O_type = input_O->data.dtype; + void* devPtrO = input_O->data.dptr; void* devPtrDescaleO = nullptr; if (O_type == DType::kFloat8E4M3 || O_type == DType::kFloat8E5M2) { devPtrDescaleO = input_O->scale_inv.dptr; @@ -1641,60 +1478,28 @@ void fused_attn_fp8_bwd( devPtrScaledV = output_dV->scale.dptr; } - NVTE_QKV_Format q_format = nvte_get_q_format(qkv_layout); - NVTE_QKV_Format kv_format = nvte_get_kv_format(qkv_layout); - - void* devPtrSeqOffsetsQ = cu_seqlens_q_padded->data.dptr; - void* devPtrSeqOffsetsKV = cu_seqlens_kv_padded->data.dptr; - - size_t max_batch_size = 0; - size_t max_tokens_q = 0; - size_t max_tokens_kv = 0; - if (q_format == NVTE_QKV_Format::NVTE_THD || kv_format == NVTE_QKV_Format::NVTE_THD) { - max_batch_size = fused_attn::get_max_batch_size(batch); - } - if (q_format == NVTE_QKV_Format::NVTE_THD) { - max_tokens_q = fused_attn::get_max_tokens(num_tokens_q); - } - if (kv_format == NVTE_QKV_Format::NVTE_THD) { - max_tokens_kv = fused_attn::get_max_tokens(num_tokens_kv); - } - void* devPtrcuSeqlensQ = reinterpret_cast(reinterpret_cast(cu_seqlens_q->data.dptr)); void* devPtrcuSeqlensKV = reinterpret_cast(reinterpret_cast(cu_seqlens_kv->data.dptr)); + void* devPtrSeqOffsetsQ = cu_seqlens_q_padded->data.dptr; + void* devPtrSeqOffsetsKV = cu_seqlens_kv_padded->data.dptr; void* devPtrDropoutSeed = reinterpret_cast(reinterpret_cast(rng_state->data.dptr)); void* devPtrDropoutOffset = reinterpret_cast(reinterpret_cast(rng_state->data.dptr) + 1); - const DType QKV_type = input_Q->data.dtype; - const DType dO_type = input_dO->data.dtype; - const DType dQKV_type = output_dQ->data.dtype; size_t workspace_size = 0; - NVTE_QKV_Format dqkv_format = nvte_get_qkv_format(dqkv_layout); - if ((dqkv_format == NVTE_QKV_Format::NVTE_BSHD) || (dqkv_format == NVTE_QKV_Format::NVTE_SBHD) || - (dqkv_format == NVTE_QKV_Format::NVTE_BHSD) || (dqkv_format == NVTE_QKV_Format::NVTE_THD)) { - fused_attn::fused_attn_fp8_bwd_impl( - batch, num_attn_heads, num_gqa_groups, max_seqlen_q, max_seqlen_kv, head_dim_qk, head_dim_v, - max_batch_size, max_tokens_q, max_tokens_kv, attn_scale, p_dropout, qkv_layout, o_format, - do_format, dqkv_layout, bias_type, mask_type, softmax_type, window_size_left, - window_size_right, bottom_right_diagonal, deterministic, devPtrQ, devPtrK, devPtrV, devPtrM, - devPtrO, devPtrdO, devPtrSoftmaxOffset, devPtrdQ, devPtrdK, devPtrdV, devPtrdSoftmaxOffset, - devPtrDescaleQ, devPtrDescaleK, devPtrDescaleV, devPtrDescaleO, devPtrDescaledO, - devPtrDescaleS, devPtrDescaledP, devPtrScaleS, devPtrScaledP, devPtrScaledQ, devPtrScaledK, - devPtrScaledV, devPtrAmaxdP, devPtrAmaxdQ, devPtrAmaxdK, devPtrAmaxdV, devPtrQ_t, devPtrK_t, - devPtrdO_f16, devPtrdO_t, devPtrDescaleQ_t, devPtrDescaleK_t, devPtrDescaledO_t, - devPtrcuSeqlensQ, devPtrcuSeqlensKV, devPtrSeqOffsetsQ, devPtrSeqOffsetsKV, - devPtrDropoutSeed, devPtrDropoutOffset, get_cudnn_fe_dtype(QKV_type), - get_cudnn_fe_dtype(O_type), get_cudnn_fe_dtype(dO_type), get_cudnn_fe_dtype(dQKV_type), - input_dO->scaling_mode, qkv_scale_inv_format, do_scale_inv_format, workspace->data.dptr, - &workspace_size, stream, handle); - } else { - NVTE_ERROR("FP8 fused attention only supports dqkv_format=BSHD, SBHD, BHSD, or THD.\n"); - } + fused_attn::fused_attn_fp8_bwd_impl( + cfg, devPtrQ, devPtrK, devPtrV, devPtrM, devPtrO, devPtrdO, devPtrSoftmaxOffset, devPtrdQ, + devPtrdK, devPtrdV, devPtrdSoftmaxOffset, devPtrDescaleQ, devPtrDescaleK, devPtrDescaleV, + devPtrDescaleO, devPtrDescaledO, devPtrDescaleS, devPtrDescaledP, devPtrScaleS, devPtrScaledP, + devPtrScaledQ, devPtrScaledK, devPtrScaledV, devPtrAmaxdP, devPtrAmaxdQ, devPtrAmaxdK, + devPtrAmaxdV, devPtrQ_t, devPtrK_t, devPtrdO_f16, devPtrdO_t, devPtrDescaleQ_t, + devPtrDescaleK_t, devPtrDescaledO_t, devPtrcuSeqlensQ, devPtrcuSeqlensKV, devPtrSeqOffsetsQ, + devPtrSeqOffsetsKV, devPtrDropoutSeed, devPtrDropoutOffset, workspace->data.dptr, + &workspace_size, stream, handle); if (workspace_size > 0) { if (workspace->data.dptr == nullptr) { @@ -1708,4 +1513,13 @@ void fused_attn_fp8_bwd( return; } } + +// Check whether cuDNN can support a given config, per forward/backward pass. +std::string support_verdict_fp8(const FusedAttnConfig& cfg, Pass pass, cudnnHandle_t handle) { + if (pass == Pass::Fwd) { + return fused_attn::support_verdict(cfg, handle); + } + return fused_attn::support_verdict(cfg, handle); +} + } // namespace transformer_engine diff --git a/transformer_engine/common/fused_attn/fused_attn_fp8.h b/transformer_engine/common/fused_attn/fused_attn_fp8.h index 5b3b2fff14..9696beb18f 100644 --- a/transformer_engine/common/fused_attn/fused_attn_fp8.h +++ b/transformer_engine/common/fused_attn/fused_attn_fp8.h @@ -8,39 +8,42 @@ * \brief Functions for fused attention for FP8 */ +#ifndef TRANSFORMER_ENGINE_COMMON_FUSED_ATTN_FUSED_ATTN_FP8_H_ +#define TRANSFORMER_ENGINE_COMMON_FUSED_ATTN_FUSED_ATTN_FP8_H_ + +#include + +#include + +#include "config_and_params.h" #include "transformer_engine/fused_attn.h" #include "transformer_engine/transformer_engine.h" namespace transformer_engine { // fused attention FWD FP8 with separate Q, K, V -void fused_attn_fp8_fwd( - size_t batch, size_t num_attn_heads, size_t num_gqa_groups, size_t max_seqlen_q, - size_t max_seqlen_kv, size_t head_dim_qk, size_t head_dim_v, size_t num_tokens_q, - size_t num_tokens_kv, bool is_training, float attn_scale, float p_dropout, - NVTE_QKV_Layout qkv_layout, NVTE_QKV_Format o_format, NVTE_QKV_Format qkv_scale_inv_format, - NVTE_Bias_Type bias_type, NVTE_Mask_Type mask_type, NVTE_Softmax_Type softmax_type, - size_t window_size_left, size_t window_size_right, bool bottom_right_diagonal, - const Tensor *input_Q, const Tensor *input_K, const Tensor *input_V, - const Tensor *input_SoftmaxOffset, Tensor *input_output_S, Tensor *output_O, - NVTETensorPack *Aux_CTX_Tensors, const Tensor *cu_seqlens_q, const Tensor *cu_seqlens_kv, - const Tensor *cu_seqlens_q_padded, const Tensor *cu_seqlens_kv_padded, const Tensor *rng_state, - Tensor *workspace, cudaStream_t stream, cudnnHandle_t handle); +void fused_attn_fp8_fwd(const fused_attn::FusedAttnConfig &cfg, const Tensor *input_Q, + const Tensor *input_K, const Tensor *input_V, + const Tensor *input_SoftmaxOffset, Tensor *input_output_S, Tensor *output_O, + NVTETensorPack *Aux_CTX_Tensors, const Tensor *cu_seqlens_q, + const Tensor *cu_seqlens_kv, const Tensor *cu_seqlens_q_padded, + const Tensor *cu_seqlens_kv_padded, const Tensor *rng_state, + Tensor *workspace, cudaStream_t stream, cudnnHandle_t handle); // fused attention BWD FP8 with separate Q, K, V -void fused_attn_fp8_bwd( - size_t batch, size_t num_attn_heads, size_t num_gqa_groups, size_t max_seqlen_q, - size_t max_seqlen_kv, size_t head_dim_qk, size_t head_dim_v, size_t num_tokens_q, - size_t num_tokens_kv, float attn_scale, float p_dropout, NVTE_QKV_Layout qkv_layout, - NVTE_QKV_Format o_format, NVTE_QKV_Format do_format, NVTE_QKV_Layout dqkv_layout, - NVTE_QKV_Format qkv_scale_inv_format, NVTE_QKV_Format do_scale_inv_format, - NVTE_Bias_Type bias_type, NVTE_Mask_Type mask_type, NVTE_Softmax_Type softmax_type, - size_t window_size_left, size_t window_size_right, bool bottom_right_diagonal, - bool deterministic, const Tensor *input_Q, const Tensor *input_K, const Tensor *input_V, - const Tensor *input_O, const Tensor *input_dO, const Tensor *input_dO_f16, - const Tensor *input_M, const Tensor *input_S, const Tensor *input_SoftmaxOffset, - Tensor *input_output_dP, const Tensor *output_dQ, const Tensor *output_dK, - const Tensor *output_dV, Tensor *output_dSoftmaxOffset, const Tensor *cu_seqlens_q, - const Tensor *cu_seqlens_kv, const Tensor *cu_seqlens_q_padded, - const Tensor *cu_seqlens_kv_padded, const Tensor *rng_state, Tensor *workspace, - cudaStream_t stream, cudnnHandle_t handle); +void fused_attn_fp8_bwd(const fused_attn::FusedAttnConfig &cfg, const Tensor *input_Q, + const Tensor *input_K, const Tensor *input_V, const Tensor *input_O, + const Tensor *input_dO, const Tensor *input_dO_f16, const Tensor *input_M, + const Tensor *input_S, const Tensor *input_SoftmaxOffset, + Tensor *input_output_dP, const Tensor *output_dQ, const Tensor *output_dK, + const Tensor *output_dV, Tensor *output_dSoftmaxOffset, + const Tensor *cu_seqlens_q, const Tensor *cu_seqlens_kv, + const Tensor *cu_seqlens_q_padded, const Tensor *cu_seqlens_kv_padded, + const Tensor *rng_state, Tensor *workspace, cudaStream_t stream, + cudnnHandle_t handle); + +std::string support_verdict_fp8(const fused_attn::FusedAttnConfig &cfg, fused_attn::Pass pass, + cudnnHandle_t handle); + } // namespace transformer_engine + +#endif // TRANSFORMER_ENGINE_COMMON_FUSED_ATTN_FUSED_ATTN_FP8_H_ diff --git a/transformer_engine/common/fused_attn/graph_cache.h b/transformer_engine/common/fused_attn/graph_cache.h new file mode 100644 index 0000000000..b3020f25ac --- /dev/null +++ b/transformer_engine/common/fused_attn/graph_cache.h @@ -0,0 +1,170 @@ +/************************************************************************* + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * + * See LICENSE for license information. + ************************************************************************/ + +// Fused-attention graph cache. The four fused-attention implementation sites, (f16/fp8 + fwd/bwd), +// each create a different graph. They differ in the operations in the graph and the input/output tensors +// that bind to the graph, but the mechanism used for their graph caching, support queries, error messaging, +// and plan building is the same. They all call these functions in this file: get_graph(), support_verdict(), +// and build_plans(). + +#ifndef TRANSFORMER_ENGINE_COMMON_FUSED_ATTN_GRAPH_CACHE_H_ +#define TRANSFORMER_ENGINE_COMMON_FUSED_ATTN_GRAPH_CACHE_H_ + +#include +#include +#include +#include +#include +#include +#include + +#include "../common.h" +#include "../cudnn_utils.h" +#include "config_and_params.h" +#include "graph_cache_debug.h" + +namespace transformer_engine { +namespace fused_attn { + +// An entry in graph cache; contains a cuDNN graph, its input/output tensors, and +// a once_flag that guards its plan build +template +struct CacheEntry { + explicit CacheEntry(GraphAndTensors graph_and_tensors) + : graph_and_tensors(std::move(graph_and_tensors)) {} + + GraphAndTensors graph_and_tensors; + std::once_flag build_plans_once; +}; + +// The graph cache; a process-wide map that maps a normalized FusedAttnConfig to a CacheEntry +template +struct GraphCache { + std::mutex mutex; + std::map>> entries; +}; + +namespace detail { + +// Thrown only when cuDNN declines a graph, so that support_verdict() can tell a config cuDNN does +// not support from a genuine failure -- a violated TE invariant, a CUDA error, a failed allocation +// -- which must surface as an error rather than be reported as an unsupported config. +struct UnsupportedByCudnn : std::runtime_error { + explicit UnsupportedByCudnn(const std::string &reason) : std::runtime_error(reason) {} +}; + +// Query if a cuDNN graph can be supported or not; if so, safely return; if not, throw with +// cuDNN frontend's original error message; times for the four stages are also recorded. +inline void query_support(Backend backend, Pass pass, cudnn_frontend::graph::Graph &graph, + cudnnHandle_t handle) { + using graph_cache_debug::BuildStage; + + auto run = [&](BuildStage stage, const char *call_name, auto &&call) { + const cudnn_frontend::error_t error = + graph_cache_debug::record_time(backend, pass, stage, [&] { return call(); }); + if (error.is_good()) return; + throw UnsupportedByCudnn(error.err_msg.empty() ? std::string(call_name) + " failed." + : error.err_msg); + }; + + run(BuildStage::Validate, "validate", [&] { return graph.validate(); }); + run(BuildStage::BuildOpGraph, "build_operation_graph", + [&] { return graph.build_operation_graph(handle); }); + run(BuildStage::CreatePlans, "create_execution_plans", + [&] { return graph.create_execution_plans({cudnn_frontend::HeurMode_t::A}); }); + run(BuildStage::CheckSupport, "check_support", [&] { return graph.check_support(); }); +} + +// Look up the key in the cache and if +// hit -> record HIT, return the cached entry +// miss -> record MISS, run build() to get a new graph, run query_support() on the new graph, +// if supported, insert it to the cache; if not, throw cuDNN frontend's original +// error message +// +// The cache lookup and insert are guarded by mutex, not the graph builds. Multiple threads +// may build for the same key concurrently, but only the first successful build will be inserted. +template +std::shared_ptr> cache_graph(GraphCache &cache, + const FusedAttnConfig &key, + Backend backend, Pass pass, + cudnnHandle_t handle, BuildFn &&build) { + using graph_cache_debug::LookupResult; + + auto find = [&]() -> std::shared_ptr> { + std::lock_guard lock(cache.mutex); + auto it = cache.entries.find(key); + return it != cache.entries.end() ? it->second : nullptr; + }; + + if (std::shared_ptr> cached = find()) { + graph_cache_debug::record_hit_miss(backend, pass, LookupResult::Hit, key); + return cached; + } + + graph_cache_debug::record_hit_miss(backend, pass, LookupResult::Miss, key); + + auto entry = std::make_shared>(build()); + graph_cache_debug::record_create_graph(backend, pass); + + query_support(backend, pass, *std::get<0>(entry->graph_and_tensors), handle); + + // Concurrent builders of the same key all arrive here, but only the one whose insert wins ends + // up cached, so only that one counts as CACHE_GRAPH. The others discard their graph. + std::shared_ptr> cached; + bool inserted = false; + { + std::lock_guard lock(cache.mutex); + const auto result = cache.entries.insert({key, std::move(entry)}); + cached = result.first->second; + inserted = result.second; + } + if (inserted) graph_cache_debug::record_cache_graph(backend, pass); + return cached; +} + +} // namespace detail + +// Create a cache for each (backend, pass) pair, and either get cached entry or build anew. +template +auto get_graph(const FusedAttnConfig &cfg, cudnnHandle_t handle) { + static GraphCache cache; + cfg.check_derived(); + return detail::cache_graph(cache, cfg.make_cache_key(kPass), kBackend, kPass, handle, + [&] { return kCreateGraphFn(cfg); }); +} + +// Check if cuDNN supports a given config; if yes, return an empty string; if not, return a diagnostic +// string with the reason that get_graph() throws. Anything other than a cuDNN rejection propagates, +// so that a real failure is not reported as an unsupported config. +template +std::string support_verdict(const FusedAttnConfig &cfg, cudnnHandle_t handle) { + try { + get_graph(cfg, handle); + return ""; + } catch (const detail::UnsupportedByCudnn &e) { + // An empty reason would read as supported, so name the site instead. + if (e.what()[0] != '\0') return e.what(); + return std::string("support_verdict<") + backend_name(kBackend) + ", " + pass_name(kPass) + + ">: rejected without a reason."; + } +} + +// Compile the kernels for the graph before execution; once per cache entry; most expensive cuDNN +// frontend call in the pre-execution, preparation process. +template +void build_plans(Backend backend, Pass pass, CacheEntry &entry) { + std::call_once(entry.build_plans_once, [&] { + cudnn_frontend::graph::Graph &graph = *std::get<0>(entry.graph_and_tensors); + graph_cache_debug::record_time(backend, pass, graph_cache_debug::BuildStage::BuildPlans, + [&] { NVTE_CHECK_CUDNN_FE(graph.build_plans()); }); + graph_cache_debug::record_build_plans(backend, pass); + }); +} + +} // namespace fused_attn +} // namespace transformer_engine + +#endif // TRANSFORMER_ENGINE_COMMON_FUSED_ATTN_GRAPH_CACHE_H_ diff --git a/transformer_engine/common/fused_attn/graph_cache_debug.h b/transformer_engine/common/fused_attn/graph_cache_debug.h new file mode 100644 index 0000000000..3ee774bacf --- /dev/null +++ b/transformer_engine/common/fused_attn/graph_cache_debug.h @@ -0,0 +1,531 @@ +/************************************************************************* + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * + * See LICENSE for license information. + ************************************************************************/ + +// Fused-attention graph cache diagnostics. Written to stderr and prefixed with [FUSED-ATTN-CACHE]. +// Enable with NVTE_FUSED_ATTN_CACHE_DEBUG=[:] at runtime, e.g., "1:all" for level 1 +// on all ranks, and "2:0,3" for level 2 on ranks 0 and 3 only. By default, only rank 0 is enabled +// when diagnostics are on. +// +// Level 0: off (default). +// Level 1: prints a summary at process exit, of cache counters and cuDNN build times. Cache +// events include: HIT/MISS, CREATE_GRAPH, CACHE_GRAPH, BUILD_PLANS, and EXECUTE, and +// counters are accumulated per backend (f16 vs fp8), per pass (forward vs backward), per +// thread. Average, CPU walltimes are also recorded for these cuDNN frontend calls: +// validate(), build_operation_graph(), create_execution_plans(), check_support(), and +// build_plans(). +// Level 2: all Level 1 diagnostics, plus one log message per event as it happens. HIT/MISS +// messages also include the full normalized cache key to identify the config that +// triggered it. +// +// For an example of Level 2 diagnostics, which also includes the summary provided by Level 1, +// please refer to docs/examples/attention/attention.ipynb. + +#ifndef TRANSFORMER_ENGINE_COMMON_FUSED_ATTN_GRAPH_CACHE_DEBUG_H_ +#define TRANSFORMER_ENGINE_COMMON_FUSED_ATTN_GRAPH_CACHE_DEBUG_H_ + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "../util/cuda_runtime.h" +#include "config_and_params.h" + +namespace transformer_engine { +namespace fused_attn { +namespace graph_cache_debug { + +enum class BuildStage { Validate, BuildOpGraph, CreatePlans, CheckSupport, BuildPlans, kCount }; +enum class LookupResult { Miss, Hit }; + +namespace detail { + +// ============================================================================ +// Debug level and rank selection +// ============================================================================ + +// NVTE_FUSED_ATTN_CACHE_DEBUG: 0=off, 1=summary, 2=trace. +inline int debug_level() { + static const int lvl = [] { + const char *e = std::getenv("NVTE_FUSED_ATTN_CACHE_DEBUG"); + if (e == nullptr || e[0] == '\0' || e[0] == '0') return 0; + const int v = std::atoi(e); + return v > 0 ? v : 1; + }(); + return lvl; +} + +// Identify which rank this process is. +inline int launcher_rank() { + static const int rank = []() -> int { + for (const char *var : {"RANK", "LOCAL_RANK", "OMPI_COMM_WORLD_RANK", "SLURM_PROCID"}) { + const char *v = std::getenv(var); + if (v != nullptr && v[0] != '\0') return std::atoi(v); + } + return -1; + }(); + return rank; +} + +// Enable diagnostics for level >= 1; on ranks specified by ":". +inline bool enabled() { + static const bool on = [] { + if (debug_level() < 1) return false; + const int rank = launcher_rank(); + if (rank < 0) return true; + const char *e = std::getenv("NVTE_FUSED_ATTN_CACHE_DEBUG"); + const char *sep = (e != nullptr) ? std::strchr(e, ':') : nullptr; + if (sep == nullptr) return rank == 0; + const std::string list(sep + 1); + if (list == "all") return true; + for (size_t pos = 0; pos <= list.size();) { + const size_t comma = list.find(',', pos); + const std::string tok = + list.substr(pos, comma == std::string::npos ? std::string::npos : comma - pos); + if (!tok.empty() && std::atoi(tok.c_str()) == rank) return true; + if (comma == std::string::npos) break; + pos = comma + 1; + } + return false; + }(); + return on; +} + +// Enable tracing for level >= 2. +inline bool enabled_with_trace() { return enabled() && debug_level() >= 2; } + +// ============================================================================ +// Build-site indexing and labels +// ============================================================================ + +// (backend, pass) pair for each build site: (f16/fp8, fwd/bwd). +constexpr size_t kSiteCount = 4; +inline constexpr size_t site_index(Backend b, Pass p) { + return (b == Backend::F16 ? 0u : 2u) + (p == Pass::Fwd ? 0u : 1u); +} + +// Name the emitting rank. +inline const std::string &rank_tag() { + static const std::string *tag = [] { + const int rank = launcher_rank(); + if (rank < 0) return new std::string(); + return new std::string("rank=" + std::to_string(rank) + " | "); + }(); + return *tag; +} + +// Short thread IDs (0, 1, 2, ...) for convenience, not pthread IDs; +// numbered by assignment order, and incremented sequentially by each thread. +inline unsigned thread_seq_id() { + static std::atomic next{0}; + static thread_local unsigned id = next.fetch_add(1, std::memory_order_relaxed); + return id; +} + +// ============================================================================ +// Event counters +// ============================================================================ + +// Cache event counters, one block per build site (f16/fp8, fwd/bwd). +struct EventCounters { + std::atomic create_graph{0}; + std::atomic cache_graph{0}; + std::atomic build_plans{0}; + std::atomic execute{0}; + std::atomic hit{0}; + std::atomic miss{0}; +}; + +inline EventCounters &counters(Backend b, Pass p) { + static std::array table{}; + return table[site_index(b, p)]; +} + +// One counter block read out into plain values, so the summary can sum blocks for +// its per-backend and all-backends rows. +struct CounterSnapshot { + uint64_t create_graph = 0; + uint64_t cache_graph = 0; + uint64_t build_plans = 0; + uint64_t execute = 0; + uint64_t hit = 0; + uint64_t miss = 0; + + CounterSnapshot &operator+=(const CounterSnapshot &other) { + create_graph += other.create_graph; + cache_graph += other.cache_graph; + build_plans += other.build_plans; + execute += other.execute; + hit += other.hit; + miss += other.miss; + return *this; + } + + bool empty() const { + return (create_graph | cache_graph | build_plans | execute | hit | miss) == 0; + } +}; + +inline CounterSnapshot snapshot(const EventCounters &c) { + CounterSnapshot s; + s.create_graph = c.create_graph.load(std::memory_order_relaxed); + s.cache_graph = c.cache_graph.load(std::memory_order_relaxed); + s.build_plans = c.build_plans.load(std::memory_order_relaxed); + s.execute = c.execute.load(std::memory_order_relaxed); + s.hit = c.hit.load(std::memory_order_relaxed); + s.miss = c.miss.load(std::memory_order_relaxed); + return s; +} + +// Per-thread counters, so the summary can break every column down by thread and backend. +// A thread that drives more than one device accumulates all of its events in one block, so record +// that rather than name a single device; the per-event lines at level 2 carry the live device. +struct ThreadCounters { + unsigned tid = 0; + std::atomic device{-1}; + std::atomic multi_device{false}; + std::array sites; + + void note_device(int current) { + const int seen = device.load(std::memory_order_relaxed); + if (seen < 0) { + device.store(current, std::memory_order_relaxed); + } else if (seen != current) { + multi_device.store(true, std::memory_order_relaxed); + } + } +}; + +// The registry and its mutex are heap-allocated and deliberately never freed. Static destructors +// and atexit handlers run as one sequence in reverse order of construction, and this registry is +// built lazily, so it can be constructed after the summary handler is registered -- and would then +// be destroyed before it runs, leaving the handler to lock a destroyed mutex and walk a destroyed +// vector. Leaking removes the ordering question, at a cost of one mutex and one vector. +inline std::mutex &thread_registry_mutex() { + static std::mutex *m = new std::mutex(); + return *m; +} +inline std::vector &thread_registry() { + static std::vector *v = new std::vector(); + return *v; +} + +// This thread's block, leaked for a related but distinct reason: a worker thread can exit long +// before the process does, while the registry holds a pointer to its block for the exit summary. +inline ThreadCounters &thread_counters() { + static thread_local ThreadCounters *tc = [] { + auto *p = new ThreadCounters(); + p->tid = thread_seq_id(); + { + std::lock_guard lock(thread_registry_mutex()); + thread_registry().push_back(p); + } + return p; + }(); + return *tc; +} + +inline EventCounters &thread_counters(Backend b, Pass p) { + return thread_counters().sites[site_index(b, p)]; +} + +// ============================================================================ +// Build-stage timing buckets +// ============================================================================ + +inline constexpr const char *kStageNames[] = { + "validate", "build_operation_graph", "create_execution_plans", "check_support", "build_plans"}; + +// Totals for one (pass, stage) pair. +struct StageTiming { + std::atomic calls{0}; + std::atomic time_ns{0}; +}; + +// Bucketed by build site, fp8 vs f16. +constexpr size_t kStageBuckets = kSiteCount * static_cast(BuildStage::kCount); +inline StageTiming &stage_timing(Backend b, Pass p, BuildStage s) { + static std::array table{}; + const size_t idx = + site_index(b, p) * static_cast(BuildStage::kCount) + static_cast(s); + return table[idx]; +} + +// ============================================================================ +// Writing to stderr +// ============================================================================ + +// The one place diagnostics reach stderr. +inline void write_stderr(const std::string &text) { + static std::atomic first_line{true}; + if (first_line.exchange(false, std::memory_order_relaxed)) { + const std::string first = "\n" + text; + std::fwrite(first.data(), 1, first.size(), stderr); + } else { + std::fwrite(text.data(), 1, text.size(), stderr); + } + std::fflush(stderr); +} + +// Format one counter block -- one pass of one backend -- as one line. +inline std::string format_counter_line(const char *tid_field, const char *dev_field, + const char *label, const CounterSnapshot &c) { + char buf[512]; + std::snprintf(buf, sizeof(buf), + "[FUSED-ATTN-CACHE] %s%-7s %-7s | %s | hit=%4" PRIu64 ", miss=%4" PRIu64 + ", create_graph=%4" PRIu64 ", cache_graph=%4" PRIu64 ", build_plans=%4" PRIu64 + ", execute=%4" PRIu64 "\n", + rank_tag().c_str(), tid_field, dev_field, label, c.hit, c.miss, c.create_graph, + c.cache_graph, c.build_plans, c.execute); + return std::string(buf); +} + +// ============================================================================ +// Exit summary +// ============================================================================ + +inline constexpr Backend kSummaryBackends[] = {Backend::F16, Backend::FP8}; + +// Names one build site for a summary row. +inline std::string site_label(Backend b, Pass p) { + return std::string(backend_name(b)) + " " + pass_name(p); +} + +// How many backends the run used. +inline size_t active_backend_count() { + size_t active = 0; + for (const Backend b : kSummaryBackends) { + if (!snapshot(counters(b, Pass::Fwd)).empty() || !snapshot(counters(b, Pass::Bwd)).empty()) { + ++active; + } + } + return active; +} + +// Per-thread breakdown, sorted by tid, one row per build site that thread used. +inline void append_thread_rows(std::string &block) { + std::lock_guard lock(thread_registry_mutex()); + std::vector blocks = thread_registry(); + std::sort(blocks.begin(), blocks.end(), + [](const ThreadCounters *a, const ThreadCounters *b) { return a->tid < b->tid; }); + for (const ThreadCounters *tc : blocks) { + char tid_field[16]; + char dev_field[16]; + std::snprintf(tid_field, sizeof(tid_field), "tid=%u", tc->tid); + if (tc->multi_device.load(std::memory_order_relaxed)) { + std::snprintf(dev_field, sizeof(dev_field), "dev=mixed"); + } else { + std::snprintf(dev_field, sizeof(dev_field), "dev=%d", + tc->device.load(std::memory_order_relaxed)); + } + for (const Backend b : kSummaryBackends) { + for (const Pass p : {Pass::Fwd, Pass::Bwd}) { + const CounterSnapshot c = snapshot(tc->sites[site_index(b, p)]); + if (c.empty()) continue; + block += format_counter_line(tid_field, dev_field, site_label(b, p).c_str(), c); + } + } + } +} + +// Totals, printed after the per-thread rows so they read as their sum: one row per build site, +// then one per pass across the backends when the run used more than one. +inline void append_total_rows(std::string &block) { + CounterSnapshot all_fwd; + CounterSnapshot all_bwd; + for (const Backend b : kSummaryBackends) { + for (const Pass p : {Pass::Fwd, Pass::Bwd}) { + const CounterSnapshot c = snapshot(counters(b, p)); + (p == Pass::Fwd ? all_fwd : all_bwd) += c; + if (c.empty()) continue; + block += format_counter_line("tid=all", "dev=all", site_label(b, p).c_str(), c); + } + } + if (active_backend_count() <= 1) return; + for (const Pass p : {Pass::Fwd, Pass::Bwd}) { + const CounterSnapshot &c = (p == Pass::Fwd ? all_fwd : all_bwd); + if (c.empty()) continue; + block += + format_counter_line("tid=all", "dev=all", (std::string("all ") + pass_name(p)).c_str(), c); + } +} + +// Mean time per call for each stage of each build site, skipping stages that no build reached. +inline void append_stage_rows(std::string &block) { + for (const Backend b : kSummaryBackends) { + for (const Pass p : {Pass::Fwd, Pass::Bwd}) { + for (int i = 0; i < static_cast(BuildStage::kCount); ++i) { + const StageTiming &t = stage_timing(b, p, static_cast(i)); + const uint64_t n = t.calls.load(std::memory_order_relaxed); + if (n == 0) continue; + const double total_ms = + static_cast(t.time_ns.load(std::memory_order_relaxed)) / 1e6; + char line[288]; + std::snprintf( + line, sizeof(line), + "[FUSED-ATTN-CACHE] %s%-3s %-3s %-22s | calls=%" PRIu64 " | time=%9.3f ms/call\n", + rank_tag().c_str(), backend_name(b), pass_name(p), kStageNames[i], n, total_ms / n); + block += line; + } + } + } +} + +// Registered at first use, by whichever recorder or stage timer fires first. +inline void register_summary_once() { + static const bool registered = [] { + std::atexit([] { + if (!enabled()) return; + // Built in memory and emitted with one write, so that concurrently-exiting processes stay + // grouped. + const std::string marker = "[FUSED-ATTN-CACHE] " + rank_tag() + "===== summary "; + std::string block = marker + "begin =====\n"; + append_thread_rows(block); + append_total_rows(block); + append_stage_rows(block); + block += marker + "end =====\n"; + write_stderr(block); + }); + return true; + }(); + (void)registered; +} + +// ============================================================================ +// Recorder internals +// ============================================================================ + +// The body every recorder shares: gate, register the exit summary, and add one to `column` in +// both the process-wide block and this thread's. +inline bool record_counter(Backend b, Pass p, std::atomic EventCounters::*column) { + if (!enabled()) return false; + register_summary_once(); + (counters(b, p).*column).fetch_add(1, std::memory_order_relaxed); + (thread_counters(b, p).*column).fetch_add(1, std::memory_order_relaxed); + thread_counters().note_device(cuda::current_device()); + return enabled_with_trace(); +} + +// One event line, from the thread the event happened on, carrying the running totals of the build +// site that raised it. +inline void print_counters(Backend b, Pass p, const char *event) { + const int device = cuda::current_device(); + char label[32]; + char tid_field[16]; + char dev_field[16]; + std::snprintf(label, sizeof(label), "%s %s %-12s", backend_name(b), pass_name(p), event); + std::snprintf(tid_field, sizeof(tid_field), "tid=%u", thread_seq_id()); + std::snprintf(dev_field, sizeof(dev_field), "dev=%d", device); + write_stderr(format_counter_line(tid_field, dev_field, label, snapshot(counters(b, p)))); +} + +// The column a lookup lands in. +inline std::atomic EventCounters::*lookup_column(LookupResult result) { + switch (result) { + case LookupResult::Hit: + return &EventCounters::hit; + case LookupResult::Miss: + break; + } + return &EventCounters::miss; +} + +inline const char *lookup_name(LookupResult result) { + switch (result) { + case LookupResult::Hit: + return "HIT"; + case LookupResult::Miss: + break; + } + return "MISS"; +} + +// Times one stage: clock read in the constructor, accumulated in the destructor. +struct ScopedBuildTimer { + BuildStage stage; + bool on; + Backend backend; + Pass pass; + std::chrono::steady_clock::time_point start; + ScopedBuildTimer(Backend b, Pass p, BuildStage s) : stage(s), on(enabled()), backend(b), pass(p) { + if (!on) return; + register_summary_once(); + start = std::chrono::steady_clock::now(); + } + ~ScopedBuildTimer() { + if (!on) return; + const uint64_t elapsed_ns = + static_cast(std::chrono::duration_cast( + std::chrono::steady_clock::now() - start) + .count()); + StageTiming &t = stage_timing(backend, pass, stage); + t.time_ns.fetch_add(elapsed_ns, std::memory_order_relaxed); + t.calls.fetch_add(1, std::memory_order_relaxed); + } +}; + +} // namespace detail + +// ============================================================================ +// Recorders: everything a call site (f16/fp8 + fwd/bwd) uses. Each recorder +// takes the calling (backend, pass) pair, adds one to that call site's counter, +// and prints a line when the level asks for it. +// ============================================================================ + +inline void record_create_graph(Backend b, Pass p) { + if (detail::record_counter(b, p, &detail::EventCounters::create_graph)) { + detail::print_counters(b, p, "CREATE_GRAPH"); + } +} + +inline void record_cache_graph(Backend b, Pass p) { + if (detail::record_counter(b, p, &detail::EventCounters::cache_graph)) { + detail::print_counters(b, p, "CACHE_GRAPH"); + } +} + +inline void record_build_plans(Backend b, Pass p) { + if (detail::record_counter(b, p, &detail::EventCounters::build_plans)) { + detail::print_counters(b, p, "BUILD_PLANS"); + } +} + +inline void record_execute(Backend b, Pass p) { + if (detail::record_counter(b, p, &detail::EventCounters::execute)) { + detail::print_counters(b, p, "EXECUTE"); + } +} + +inline void record_hit_miss(Backend b, Pass p, LookupResult result, const FusedAttnConfig &key) { + if (!detail::record_counter(b, p, detail::lookup_column(result))) return; + char prefix[128]; + std::snprintf(prefix, sizeof(prefix), + "[FUSED-ATTN-CACHE] %stid=%-3u dev=%-3d | %-3s %-3s %-12s | ", + detail::rank_tag().c_str(), detail::thread_seq_id(), key.device_id, backend_name(b), + pass_name(p), detail::lookup_name(result)); + detail::write_stderr(prefix + key.to_string() + "\n"); +} + +template +inline decltype(auto) record_time(Backend b, Pass p, BuildStage stage, Fn &&fn) { + detail::ScopedBuildTimer scoped(b, p, stage); + return fn(); +} + +} // namespace graph_cache_debug +} // namespace fused_attn +} // namespace transformer_engine + +#endif // TRANSFORMER_ENGINE_COMMON_FUSED_ATTN_GRAPH_CACHE_DEBUG_H_ diff --git a/transformer_engine/common/fused_attn/utils.cu b/transformer_engine/common/fused_attn/utils.cu index 9b54a64cbe..c6c8957b1b 100644 --- a/transformer_engine/common/fused_attn/utils.cu +++ b/transformer_engine/common/fused_attn/utils.cu @@ -8,7 +8,7 @@ #include #include "../common.h" -#include "../cudnn_utils.h" +#include "../util/cuda_runtime.h" #include "transformer_engine/fused_attn.h" #include "utils.h" @@ -324,93 +324,6 @@ void generateMatrixStrides(int64_t b, int64_t h, int64_t s_q, int64_t s_kv, int6 } } -bool allowAllConfig(cudnnBackendDescriptor_t engine_config) { - (void)engine_config; - return false; -} - -cudnn_frontend::Tensor tensor_create(cudnnDataType_t type, int64_t id, int64_t const *dim, - int64_t const *stride, bool is_virtual, bool is_value) { - int nbDims = 4; - auto tensor_created = - cudnn_frontend::TensorBuilder() - .setDim(nbDims, dim) - .setStride(nbDims, stride) - .setId(id) - .setAlignment(16) // 16B alignment is needed to run a tensor core engine - .setDataType(type) - .setVirtual(is_virtual) - .setByValue(is_value) - .build(); - return tensor_created; -} - -cudnn_frontend::Tensor tensor_create_with_offset( - cudnnDataType_t type, int64_t id, int64_t const *dim, int64_t const *stride, bool is_virtual, - bool is_value, std::shared_ptr raggedOffset) { - int nbDims = 4; - auto tensor_created = - cudnn_frontend::TensorBuilder() - .setDim(nbDims, dim) - .setStride(nbDims, stride) - .setId(id) - .setAlignment(16) // 16B alignment is needed to run a tensor core engine - .setDataType(type) - .setVirtual(is_virtual) - .setByValue(is_value) - .setRaggedOffset(raggedOffset) - .build(); - return tensor_created; -} - -cudnn_frontend::PointWiseDesc pw_desc_create(cudnnDataType_t type, cudnnPointwiseMode_t mode) { - auto pw_desc_created = - cudnn_frontend::PointWiseDescBuilder().setMode(mode).setComputeType(type).build(); - return pw_desc_created; -} - -cudnn_frontend::Operation unary_pw_op_create(cudnn_frontend::Tensor const &xDesc, - cudnn_frontend::Tensor const &yDesc, - cudnn_frontend::PointWiseDesc const &pwDesc) { - auto pw_op_created = - cudnn_frontend::OperationBuilder(CUDNN_BACKEND_OPERATION_POINTWISE_DESCRIPTOR) - .setxDesc(xDesc) - .setyDesc(yDesc) - .setpwDesc(pwDesc) - .build(); - return pw_op_created; -} - -cudnn_frontend::Operation binary_pw_op_create(cudnn_frontend::Tensor const &xDesc, - cudnn_frontend::Tensor const &bDesc, - cudnn_frontend::Tensor const &yDesc, - cudnn_frontend::PointWiseDesc const &pwDesc) { - auto pw_op_created = - cudnn_frontend::OperationBuilder(CUDNN_BACKEND_OPERATION_POINTWISE_DESCRIPTOR) - .setxDesc(xDesc) - .setbDesc(bDesc) - .setyDesc(yDesc) - .setpwDesc(pwDesc) - .build(); - return pw_op_created; -} - -cudnn_frontend::Operation ternary_pw_op_create(cudnn_frontend::Tensor const &xDesc, - cudnn_frontend::Tensor const &bDesc, - cudnn_frontend::Tensor const &tDesc, - cudnn_frontend::Tensor const &yDesc, - cudnn_frontend::PointWiseDesc const &pwDesc) { - auto pw_op_created = - cudnn_frontend::OperationBuilder(CUDNN_BACKEND_OPERATION_POINTWISE_DESCRIPTOR) - .setxDesc(xDesc) - .setbDesc(bDesc) - .settDesc(tDesc) - .setyDesc(yDesc) - .setpwDesc(pwDesc) - .build(); - return pw_op_created; -} - // convert cu_seqlens to actual_seqlens __global__ void cu_seqlens_to_actual_seqlens(int64_t actual_b, int64_t max_b, int32_t const *const q_cu_seqlens, @@ -509,6 +422,7 @@ DType get_ragged_offset_dtype(NVTE_QKV_Layout_Group layout_group, int64_t num_at // quantize batch size size_t get_max_batch_size(size_t batch_size) { + if (batch_size == 0) return 0; // guard: log2(0) = -inf, casting to size_t is UB size_t max_b = batch_size; size_t log2_b = ceil(log2(batch_size)); // batch size is expected to be 10s-100s @@ -527,6 +441,7 @@ size_t get_max_batch_size(size_t batch_size) { // quantize token count size_t get_max_tokens(size_t num_tokens) { + if (num_tokens == 0) return 0; // guard: log2(0) = -inf, casting to size_t is UB // token count is expected to be 1k's-100k's // t = 0, ..., 1024 -> max_t = 1024 // t = 1025, ..., 32k -> max_t = next power of 2 diff --git a/transformer_engine/common/fused_attn/utils.h b/transformer_engine/common/fused_attn/utils.h index 1864f9417d..237507f30b 100644 --- a/transformer_engine/common/fused_attn/utils.h +++ b/transformer_engine/common/fused_attn/utils.h @@ -4,15 +4,10 @@ * See LICENSE for license information. ************************************************************************/ -#ifndef TRANSFORMER_ENGINE_FUSED_ATTN_UTILS_H_ -#define TRANSFORMER_ENGINE_FUSED_ATTN_UTILS_H_ - -#include -#include -#include +#ifndef TRANSFORMER_ENGINE_COMMON_FUSED_ATTN_UTILS_H_ +#define TRANSFORMER_ENGINE_COMMON_FUSED_ATTN_UTILS_H_ #include -#include #include "../common.h" #include "transformer_engine/fused_attn.h" @@ -223,97 +218,13 @@ inline void generateMatrixStridesWithLayout(int64_t b, int64_t h, int64_t hg, in void generateMatrixStrides(int64_t b, int64_t h, int64_t s_q, int64_t s_kv, int64_t d, int64_t *strideA, NVTE_QKV_Layout layout, NVTE_QKV_Matrix matrix); -bool allowAllConfig(cudnnBackendDescriptor_t engine_config); - -cudnn_frontend::Tensor tensor_create(cudnnDataType_t type, int64_t id, int64_t const *dim, - int64_t const *stride, bool is_virtual, bool is_value); - -cudnn_frontend::Tensor tensor_create_with_offset( - cudnnDataType_t type, int64_t id, int64_t const *dim, int64_t const *stride, bool is_virtual, - bool is_value, std::shared_ptr raggedOffset); - -cudnn_frontend::PointWiseDesc pw_desc_create(cudnnDataType_t type, cudnnPointwiseMode_t mode); - -cudnn_frontend::Operation unary_pw_op_create(cudnn_frontend::Tensor const &xDesc, - cudnn_frontend::Tensor const &yDesc, - cudnn_frontend::PointWiseDesc const &pwDesc); - -cudnn_frontend::Operation binary_pw_op_create(cudnn_frontend::Tensor const &xDesc, - cudnn_frontend::Tensor const &bDesc, - cudnn_frontend::Tensor const &yDesc, - cudnn_frontend::PointWiseDesc const &pwDesc); - -cudnn_frontend::Operation ternary_pw_op_create(cudnn_frontend::Tensor const &xDesc, - cudnn_frontend::Tensor const &bDesc, - cudnn_frontend::Tensor const &tDesc, - cudnn_frontend::Tensor const &yDesc, - cudnn_frontend::PointWiseDesc const &pwDesc); - -struct FADescriptor_v1 { - std::int64_t b; - std::int64_t h; - std::int64_t hg; - std::int64_t s_q; - std::int64_t s_kv; - std::int64_t d_qk; - std::int64_t d_v; - std::int64_t num_pages_k; - std::int64_t num_pages_v; - std::int64_t page_size_k; - std::int64_t page_size_v; - std::int64_t max_pages_per_seq_k; - std::int64_t max_pages_per_seq_v; - std::int64_t bias_b; - std::int64_t bias_h; - std::int64_t bias_sq; - std::int64_t bias_skv; - float attnScale; - bool isTraining; - float dropoutProbability; - 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 mask_type; - NVTE_Softmax_Type softmax_type; - std::int64_t window_size_left; - std::int64_t window_size_right; - bool bottom_right_diagonal; - bool deterministic; - cudnn_frontend::DataType_t qkv_tensor_type; - cudnn_frontend::DataType_t o_tensor_type; - cudnn_frontend::DataType_t do_tensor_type; - cudnn_frontend::DataType_t dqkv_tensor_type; - bool return_max_logit; - - bool operator<(const FADescriptor_v1 &rhs) const { - return std::tie(b, h, hg, s_q, s_kv, d_qk, d_v, num_pages_k, num_pages_v, page_size_k, - page_size_v, max_pages_per_seq_k, max_pages_per_seq_v, bias_b, bias_h, bias_sq, - bias_skv, attnScale, isTraining, dropoutProbability, qkv_layout, o_format, - do_format, dqkv_layout, qkv_scale_inv_format, do_scale_inv_format, mask_type, - softmax_type, window_size_left, window_size_right, bottom_right_diagonal, - deterministic, bias_type, qkv_tensor_type, o_tensor_type, do_tensor_type, - dqkv_tensor_type, return_max_logit) < - std::tie(rhs.b, rhs.h, rhs.hg, rhs.s_q, rhs.s_kv, rhs.d_qk, rhs.d_v, rhs.num_pages_k, - rhs.num_pages_v, rhs.page_size_k, rhs.page_size_v, rhs.max_pages_per_seq_k, - rhs.max_pages_per_seq_v, rhs.bias_b, rhs.bias_h, rhs.bias_sq, rhs.bias_skv, - rhs.attnScale, rhs.isTraining, rhs.dropoutProbability, rhs.qkv_layout, - rhs.o_format, rhs.do_format, rhs.dqkv_layout, rhs.qkv_scale_inv_format, - rhs.do_scale_inv_format, rhs.mask_type, rhs.softmax_type, rhs.window_size_left, - rhs.window_size_right, rhs.bottom_right_diagonal, rhs.deterministic, - rhs.bias_type, rhs.qkv_tensor_type, rhs.o_tensor_type, rhs.do_tensor_type, - rhs.dqkv_tensor_type, rhs.return_max_logit); - } -}; - // Per-tensor scale factors relating cu_seqlens_padded (token units) to tensor-element // ragged offsets, as a function of the QKV layout group. Single source of truth shared // by the cu_seqlens_padded_to_offsets conversion kernel and the direct-seqlens path // (which passes them to cuDNN as ragged offset multipliers). struct RaggedOffsetMultipliers { + RaggedOffsetMultipliers() = default; + RaggedOffsetMultipliers(NVTE_QKV_Layout_Group layout_group, int64_t h, int64_t hg, int64_t d_qk, int64_t d_v) : q(h * d_qk), k(hg * d_qk), v(hg * d_v), o(h * d_v), stats(h), kv_from_q(false) { @@ -332,13 +243,13 @@ struct RaggedOffsetMultipliers { } } - int64_t q; - int64_t k; - int64_t v; - int64_t o; - int64_t stats; + int64_t q = 0; + int64_t k = 0; + int64_t v = 0; + int64_t o = 0; + int64_t stats = 0; // K/V offsets scale the Q-side cu_seqlens_padded (interleaved QKV layouts) - bool kv_from_q; + bool kv_from_q = false; }; __global__ void cu_seqlens_to_actual_seqlens(int64_t actual_b, int64_t max_b, @@ -394,4 +305,4 @@ uint32_t GetRuntimeNumSegments(void *cu_seqlen, void *workspace, size_t len, cud } // namespace fused_attn } // namespace transformer_engine -#endif +#endif // TRANSFORMER_ENGINE_COMMON_FUSED_ATTN_UTILS_H_ diff --git a/transformer_engine/common/include/transformer_engine/fused_attn.h b/transformer_engine/common/include/transformer_engine/fused_attn.h index 4a7fd55f9d..e889c28219 100644 --- a/transformer_engine/common/include/transformer_engine/fused_attn.h +++ b/transformer_engine/common/include/transformer_engine/fused_attn.h @@ -8,8 +8,8 @@ * \brief Enums and functions for fused attention. */ -#ifndef TRANSFORMER_ENGINE_FUSED_ATTN_FP8_H_ -#define TRANSFORMER_ENGINE_FUSED_ATTN_FP8_H_ +#ifndef TRANSFORMER_ENGINE_FUSED_ATTN_H_ +#define TRANSFORMER_ENGINE_FUSED_ATTN_H_ #include "stdint.h" #include "transformer_engine.h" @@ -194,6 +194,239 @@ NVTE_QKV_Format nvte_get_q_format(NVTE_QKV_Layout qkv_layout); */ NVTE_QKV_Format nvte_get_kv_format(NVTE_QKV_Layout qkv_layout); +/*! \brief Opaque fused-attention configuration handle. */ +typedef void *NVTEFusedAttnConfig; + +/*! \enum NVTEFusedAttnConfigAttribute + * \brief Attributes for `NVTEFusedAttnConfig`. + * + * This enum is used to index the `FusedAttnConfig` struct. The order of its fields matches + * that of the declaration fields and the `attr_sizes` array of `FusedAttnConfig`. New fields + * may be appended at the end, and existing fields are never reordered, removed, or resized. + */ +enum NVTEFusedAttnConfigAttribute { + // Basic attention settings + kNVTEFusedAttnConfigIsTraining = 0, + kNVTEFusedAttnConfigDeterministic, + kNVTEFusedAttnConfigCudaGraph, + kNVTEFusedAttnConfigReturnMaxLogit, + kNVTEFusedAttnConfigAttnMaskType, + kNVTEFusedAttnConfigBiasType, + kNVTEFusedAttnConfigWindowSizeLeft, + kNVTEFusedAttnConfigWindowSizeRight, + kNVTEFusedAttnConfigBottomRightDiagonal, + kNVTEFusedAttnConfigSoftmaxType, + kNVTEFusedAttnConfigScalingMode, + kNVTEFusedAttnConfigDropout, + kNVTEFusedAttnConfigAttnScale, + // Tensor types + kNVTEFusedAttnConfigQKVDtype, + kNVTEFusedAttnConfigODtype, + kNVTEFusedAttnConfigDODtype, + kNVTEFusedAttnConfigDQKVDtype, + // Tensor layouts + kNVTEFusedAttnConfigQKVLayout, + kNVTEFusedAttnConfigOFormat, + kNVTEFusedAttnConfigDOFormat, + kNVTEFusedAttnConfigDQKVLayout, + kNVTEFusedAttnConfigQKVScaleInvFormat, + kNVTEFusedAttnConfigDOScaleInvFormat, + // Tensor dimensions + kNVTEFusedAttnConfigBatchSize, + kNVTEFusedAttnConfigNumAttnHeads, + kNVTEFusedAttnConfigNumGQAGroups, + kNVTEFusedAttnConfigHeadDimQK, + kNVTEFusedAttnConfigHeadDimV, + kNVTEFusedAttnConfigMaxSeqlenQ, + kNVTEFusedAttnConfigMaxSeqlenKV, + kNVTEFusedAttnConfigNumTokensQ, + kNVTEFusedAttnConfigNumTokensKV, + // Paged KV dimensions + kNVTEFusedAttnConfigNumPagesK, + kNVTEFusedAttnConfigNumPagesV, + kNVTEFusedAttnConfigPageSizeK, + kNVTEFusedAttnConfigPageSizeV, + kNVTEFusedAttnConfigMaxPagesPerSeqK, + kNVTEFusedAttnConfigMaxPagesPerSeqV, + // Bias dimensions + kNVTEFusedAttnConfigBiasBatchSize, + kNVTEFusedAttnConfigBiasNumHeads, + kNVTEFusedAttnConfigBiasSeqlenQ, + kNVTEFusedAttnConfigBiasSeqlenKV, + // Number of attributes + kNVTEFusedAttnConfigNumAttributes +}; + +/*! \brief Create a fused-attention configuration. */ +NVTEFusedAttnConfig nvte_create_fused_attn_config(void); + +/*! \brief Destroy a fused-attention configuration. */ +void nvte_destroy_fused_attn_config(NVTEFusedAttnConfig config); + +/*! \brief Query an attribute in a fused-attention configuration. */ +void nvte_get_fused_attn_config_attribute(NVTEFusedAttnConfig config, + NVTEFusedAttnConfigAttribute attr, void *buf, + size_t size_in_bytes, size_t *size_written); + +/*! \brief Set an attribute in a fused-attention configuration. */ +void nvte_set_fused_attn_config_attribute(NVTEFusedAttnConfig config, + NVTEFusedAttnConfigAttribute attr, const void *buf, + size_t size_in_bytes); + +/*! \brief Opaque fused-attention forward-parameter handle. */ +typedef void *NVTEFusedAttnFwdParams; + +/*! \enum NVTEFusedAttnFwdParamsAttribute + * \brief Attributes for `NVTEFusedAttnFwdParams`. + * + * This enum is used to index the `FusedAttnFwdParams` struct. The order of its fields matches + * that of the declaration fields and the `attr_sizes` array of `FusedAttnFwdParams`. New fields + * may be appended at the end, and existing fields are never reordered, removed, or resized. + */ +enum NVTEFusedAttnFwdParamsAttribute { + // Tensor handles + kNVTEFusedAttnFwdParamsQ = 0, + kNVTEFusedAttnFwdParamsK, + kNVTEFusedAttnFwdParamsV, + kNVTEFusedAttnFwdParamsBias, + kNVTEFusedAttnFwdParamsSoftmaxOffset, + kNVTEFusedAttnFwdParamsS, + kNVTEFusedAttnFwdParamsO, + kNVTEFusedAttnFwdParamsAuxCtxTensors, + kNVTEFusedAttnFwdParamsCuSeqlensQ, + kNVTEFusedAttnFwdParamsCuSeqlensKV, + kNVTEFusedAttnFwdParamsCuSeqlensQPadded, + kNVTEFusedAttnFwdParamsCuSeqlensKVPadded, + kNVTEFusedAttnFwdParamsPageTableK, + kNVTEFusedAttnFwdParamsPageTableV, + kNVTEFusedAttnFwdParamsRngState, + // Configuration knobs + kNVTEFusedAttnFwdParamsIsTraining, + kNVTEFusedAttnFwdParamsCudaGraph, + kNVTEFusedAttnFwdParamsReturnMaxLogit, + kNVTEFusedAttnFwdParamsAttnMaskType, + kNVTEFusedAttnFwdParamsBiasType, + kNVTEFusedAttnFwdParamsWindowSizeLeft, + kNVTEFusedAttnFwdParamsWindowSizeRight, + kNVTEFusedAttnFwdParamsBottomRightDiagonal, + kNVTEFusedAttnFwdParamsSoftmaxType, + kNVTEFusedAttnFwdParamsDropout, + kNVTEFusedAttnFwdParamsAttnScale, + kNVTEFusedAttnFwdParamsQKVLayout, + kNVTEFusedAttnFwdParamsOFormat, + kNVTEFusedAttnFwdParamsQKVScaleInvFormat, + kNVTEFusedAttnFwdParamsMaxSeqlenQ, + kNVTEFusedAttnFwdParamsMaxSeqlenKV, + // Workspace and stream + kNVTEFusedAttnFwdParamsWorkspace, + kNVTEFusedAttnFwdParamsStream, + // Number of attributes + kNVTEFusedAttnFwdParamsNumAttributes +}; + +/*! \brief Create a fused-attention forward-parameter object. */ +NVTEFusedAttnFwdParams nvte_create_fused_attn_fwd_params(void); + +/*! \brief Destroy a fused-attention forward-parameter object. */ +void nvte_destroy_fused_attn_fwd_params(NVTEFusedAttnFwdParams params); + +/*! \brief Query an attribute in a fused-attention forward-parameter object. */ +void nvte_get_fused_attn_fwd_params_attribute(NVTEFusedAttnFwdParams params, + NVTEFusedAttnFwdParamsAttribute attr, void *buf, + size_t size_in_bytes, size_t *size_written); + +/*! \brief Set an attribute in a fused-attention forward-parameter object. */ +void nvte_set_fused_attn_fwd_params_attribute(NVTEFusedAttnFwdParams params, + NVTEFusedAttnFwdParamsAttribute attr, const void *buf, + size_t size_in_bytes); + +/*! \brief Opaque fused-attention backward-parameter handle. */ +typedef void *NVTEFusedAttnBwdParams; + +/*! \enum NVTEFusedAttnBwdParamsAttribute + * \brief Attributes for `NVTEFusedAttnBwdParams`. + * + * This enum is used to index the `FusedAttnBwdParams` struct. The order of its fields matches + * that of the declaration fields and the `attr_sizes` array of `FusedAttnBwdParams`. New fields + * may be appended at the end, and existing fields are never reordered, removed, or resized. + */ +enum NVTEFusedAttnBwdParamsAttribute { + // Tensor handles + kNVTEFusedAttnBwdParamsQ = 0, + kNVTEFusedAttnBwdParamsK, + kNVTEFusedAttnBwdParamsV, + kNVTEFusedAttnBwdParamsO, + kNVTEFusedAttnBwdParamsDO, + kNVTEFusedAttnBwdParamsS, + kNVTEFusedAttnBwdParamsDP, + kNVTEFusedAttnBwdParamsAuxCtxTensors, + kNVTEFusedAttnBwdParamsDQ, + kNVTEFusedAttnBwdParamsDK, + kNVTEFusedAttnBwdParamsDV, + kNVTEFusedAttnBwdParamsDBias, + kNVTEFusedAttnBwdParamsDSoftmaxOffset, + kNVTEFusedAttnBwdParamsCuSeqlensQ, + kNVTEFusedAttnBwdParamsCuSeqlensKV, + kNVTEFusedAttnBwdParamsCuSeqlensQPadded, + kNVTEFusedAttnBwdParamsCuSeqlensKVPadded, + // Configuration knobs + kNVTEFusedAttnBwdParamsDeterministic, + kNVTEFusedAttnBwdParamsCudaGraph, + kNVTEFusedAttnBwdParamsAttnMaskType, + kNVTEFusedAttnBwdParamsBiasType, + kNVTEFusedAttnBwdParamsWindowSizeLeft, + kNVTEFusedAttnBwdParamsWindowSizeRight, + kNVTEFusedAttnBwdParamsBottomRightDiagonal, + kNVTEFusedAttnBwdParamsSoftmaxType, + kNVTEFusedAttnBwdParamsDropout, + kNVTEFusedAttnBwdParamsAttnScale, + kNVTEFusedAttnBwdParamsQKVLayout, + kNVTEFusedAttnBwdParamsOFormat, + kNVTEFusedAttnBwdParamsDOFormat, + kNVTEFusedAttnBwdParamsDQKVLayout, + kNVTEFusedAttnBwdParamsQKVScaleInvFormat, + kNVTEFusedAttnBwdParamsDOScaleInvFormat, + kNVTEFusedAttnBwdParamsMaxSeqlenQ, + kNVTEFusedAttnBwdParamsMaxSeqlenKV, + // Workspace and stream + kNVTEFusedAttnBwdParamsWorkspace, + kNVTEFusedAttnBwdParamsStream, + // Number of attributes + kNVTEFusedAttnBwdParamsNumAttributes +}; + +/*! \brief Create a fused-attention backward-parameter object. */ +NVTEFusedAttnBwdParams nvte_create_fused_attn_bwd_params(void); + +/*! \brief Destroy a fused-attention backward-parameter object. */ +void nvte_destroy_fused_attn_bwd_params(NVTEFusedAttnBwdParams params); + +/*! \brief Query an attribute in a fused-attention backward-parameter object. */ +void nvte_get_fused_attn_bwd_params_attribute(NVTEFusedAttnBwdParams params, + NVTEFusedAttnBwdParamsAttribute attr, void *buf, + size_t size_in_bytes, size_t *size_written); + +/*! \brief Set an attribute in a fused-attention backward-parameter object. */ +void nvte_set_fused_attn_bwd_params_attribute(NVTEFusedAttnBwdParams params, + NVTEFusedAttnBwdParamsAttribute attr, const void *buf, + size_t size_in_bytes); + +/*! \brief Get fused-attention backend based on the user configuration. + * + * \param[in] cfg Fused-attention configuration created by + * `nvte_create_fused_attn_config()`. + * \param[out] message If the configuration is supported, an empty string. If not supported, + * a diagnostic message explaining why the configuration is rejected. Pass + * `NULL` to skip the diagnostics. The library maintains a per-thread buffer, + * and callers do not need to allocate memory for the message. The buffer + * content stays valid until the next call of `nvte_get_fused_attn_backend_v2`. + * + * \return Fused-attention backend, `NVTE_F16_arbitrary_seqlen` or `NVTE_FP8`, + * if the configuration is supported; otherwise, `NVTE_No_Backend`. + */ +NVTE_Fused_Attn_Backend nvte_get_fused_attn_backend_v2(NVTEFusedAttnConfig cfg, + const char **message); + /*! \brief Get fused attention backend based on input parameters. * * \param[in] is_training Whether the model is in training mode. @@ -215,6 +448,16 @@ NVTE_QKV_Format nvte_get_kv_format(NVTE_QKV_Layout qkv_layout); * \param[in] return_max_logit Whether to produce Max along with Stats. * \param[in] cuda_graph Whether cuda graph capture is enabled or not. * \param[in] deterministic Whether determinism is required or not. + * + * \deprecated This function has been deprecated in favor of `nvte_get_fused_attn_backend_v2`. + * + * \note `nvte_get_fused_attn_backend` has a narrower input signature compared to + * `nvte_get_fused_attn_backend_v2`. For the fields it cannot express, it fills them with + * the default values of `nvte_get_fused_attn_backend_v2`. This includes setting + * `batch_size` = 1, deriving output/gradient formats from `qkv_layout`, assuming a standard + * bias shape [b, h, sq, skv] for `NVTE_POST_SCALE_BIAS`, using delayed scaling for all FP8, + * and not supporting paged-KV attention checks. Users who need more precise control should + * use `nvte_get_fused_attn_backend_v2` instead. */ NVTE_Fused_Attn_Backend nvte_get_fused_attn_backend( bool is_training, NVTEDType q_dtype, NVTEDType kv_dtype, NVTE_QKV_Layout qkv_layout, @@ -223,6 +466,17 @@ NVTE_Fused_Attn_Backend nvte_get_fused_attn_backend( size_t max_seqlen_kv, size_t head_dim_qk, size_t head_dim_v, int64_t window_size_left, int64_t window_size_right, bool return_max_logit, bool cuda_graph, bool deterministic); +/*! \brief Compute dot product attention with Q, K, and V. + * + * All inputs and outputs are carried by the opaque \p params handle. Create it with + * `nvte_create_fused_attn_fwd_params()`, populate it with + * `nvte_set_fused_attn_fwd_params_attribute()` or `FusedAttnFwdParamsWrapper` setters, and + * destroy it with `nvte_destroy_fused_attn_fwd_params()`. + * + * \param[in,out] params Fused-attention forward-parameter handle. + */ +void nvte_fused_attn_fwd_v2(NVTEFusedAttnFwdParams params); + /*! \brief Compute dot product attention with separate Q, K and V. * * Computes: @@ -272,7 +526,8 @@ NVTE_Fused_Attn_Backend nvte_get_fused_attn_backend( * \param[in] qkv_layout QKV tensors' layout. * \param[in] o_format Output format. * \param[in] qkv_scale_inv_format Format of scale-inverse tensors for QKV; - * if NVTE_QKV_Format_NOT_SET, inferred from qkv_layout. + * if `NVTE_QKV_Format_NOT_SET`, inferred from + * `qkv_layout`. * \param[in] bias_type Bias type. * \param[in] attn_mask_type Attention mask type. * \param[in] softmax_type Attention softmax type. @@ -281,6 +536,8 @@ NVTE_Fused_Attn_Backend nvte_get_fused_attn_backend( * \param[in] bottom_right_diagonal Whether to align sliding window and ALiBi diagonal to the bottom right corner of the softmax matrix. * \param[in] workspace Workspace tensor. * \param[in] stream CUDA stream used for this operation. + * + * \deprecated This function has been deprecated in favor of `nvte_fused_attn_fwd_v2`. */ void nvte_fused_attn_fwd(const NVTETensor Q, const NVTETensor K, const NVTETensor V, const NVTETensor Bias, const NVTETensor SoftmaxOffset, NVTETensor S, @@ -297,6 +554,17 @@ void nvte_fused_attn_fwd(const NVTETensor Q, const NVTETensor K, const NVTETenso int64_t window_size_left, int64_t window_size_right, bool bottom_right_diagonal, NVTETensor workspace, cudaStream_t stream); +/*! \brief Compute the backward of the dot product attention with Q, K and V. + * + * All inputs and outputs are carried by the opaque \p params handle. Create it with + * `nvte_create_fused_attn_bwd_params()`, populate it with + * `nvte_set_fused_attn_bwd_params_attribute()` or `FusedAttnBwdParamsWrapper` setters, and + * destroy it with `nvte_destroy_fused_attn_bwd_params()`. + * + * \param[in,out] params Fused-attention backward-parameter handle. + */ +void nvte_fused_attn_bwd_v2(NVTEFusedAttnBwdParams params); + /*! \brief Compute the backward of the dot product attention with separate Q, K and V. * * Notes: @@ -341,9 +609,11 @@ void nvte_fused_attn_fwd(const NVTETensor Q, const NVTETensor K, const NVTETenso * \param[in] do_format Output gradient's format. * \param[in] dqkv_layout QKV gradient tensors' layout. * \param[in] qkv_scale_inv_format Format of scale-inverse tensors for QKV; - * if NVTE_QKV_Format_NOT_SET, inferred from qkv_layout. + * if `NVTE_QKV_Format_NOT_SET`, inferred from + * `qkv_layout`. * \param[in] do_scale_inv_format Format of scale-inverse tensors for dO; - * if NVTE_QKV_Format_NOT_SET, inferred from the output layout. + * if `NVTE_QKV_Format_NOT_SET`, inferred from the + * output layout. * \param[in] bias_type Bias type. * \param[in] attn_mask_type Attention mask type. * \param[in] softmax_type Attention softmax type. @@ -354,6 +624,8 @@ void nvte_fused_attn_fwd(const NVTETensor Q, const NVTETensor K, const NVTETenso * \param[in] cuda_graph Whether cuda graph capture is enabled or not. * \param[in] workspace Workspace tensor. * \param[in] stream CUDA stream used for this operation. + * + * \deprecated This function has been deprecated in favor of `nvte_fused_attn_bwd_v2`. */ void nvte_fused_attn_bwd(const NVTETensor Q, const NVTETensor K, const NVTETensor V, const NVTETensor O, const NVTETensor dO, const NVTETensor S, NVTETensor dP, @@ -714,6 +986,488 @@ class AttentionShape { size_t canonical_[5] = {}; }; +/*! \class FusedAttnConfigWrapper + * \brief C++ helper for constructing an `NVTEFusedAttnConfig`. + * + * It owns an opaque `NVTEFusedAttnConfig` handle created by + * `nvte_create_fused_attn_config()`, and provides a convenient, + * chainable interface for setting every field in `FusedAttnConfig`. + */ +class FusedAttnConfigWrapper { + public: + FusedAttnConfigWrapper() : cfg_{nvte_create_fused_attn_config()} {} + + FusedAttnConfigWrapper(const FusedAttnConfigWrapper &) = delete; + FusedAttnConfigWrapper &operator=(const FusedAttnConfigWrapper &) = delete; + + FusedAttnConfigWrapper(FusedAttnConfigWrapper &&other) noexcept : cfg_{other.cfg_} { + other.cfg_ = nullptr; + } + + FusedAttnConfigWrapper &operator=(FusedAttnConfigWrapper &&other) noexcept { + if (this != &other) { + if (cfg_ != nullptr) { + nvte_destroy_fused_attn_config(cfg_); + } + cfg_ = other.cfg_; + other.cfg_ = nullptr; + } + return *this; + } + + ~FusedAttnConfigWrapper() { + if (cfg_ != nullptr) { + nvte_destroy_fused_attn_config(cfg_); + } + } + + operator NVTEFusedAttnConfig() const noexcept { return cfg_; } + NVTEFusedAttnConfig get() const noexcept { return cfg_; } + + FusedAttnConfigWrapper &set_is_training(bool val) noexcept { + return set_attr(kNVTEFusedAttnConfigIsTraining, static_cast(val)); + } + FusedAttnConfigWrapper &set_deterministic(bool val) noexcept { + return set_attr(kNVTEFusedAttnConfigDeterministic, static_cast(val)); + } + FusedAttnConfigWrapper &set_cuda_graph(bool val) noexcept { + return set_attr(kNVTEFusedAttnConfigCudaGraph, static_cast(val)); + } + FusedAttnConfigWrapper &set_return_max_logit(bool val) noexcept { + return set_attr(kNVTEFusedAttnConfigReturnMaxLogit, static_cast(val)); + } + FusedAttnConfigWrapper &set_attn_mask_type(NVTE_Mask_Type val) noexcept { + return set_attr(kNVTEFusedAttnConfigAttnMaskType, val); + } + FusedAttnConfigWrapper &set_bias_type(NVTE_Bias_Type val) noexcept { + return set_attr(kNVTEFusedAttnConfigBiasType, val); + } + FusedAttnConfigWrapper &set_window_size_left(int64_t val) noexcept { + return set_attr(kNVTEFusedAttnConfigWindowSizeLeft, val); + } + FusedAttnConfigWrapper &set_window_size_right(int64_t val) noexcept { + return set_attr(kNVTEFusedAttnConfigWindowSizeRight, val); + } + FusedAttnConfigWrapper &set_bottom_right_diagonal(bool val) noexcept { + return set_attr(kNVTEFusedAttnConfigBottomRightDiagonal, static_cast(val)); + } + FusedAttnConfigWrapper &set_softmax_type(NVTE_Softmax_Type val) noexcept { + return set_attr(kNVTEFusedAttnConfigSoftmaxType, val); + } + FusedAttnConfigWrapper &set_scaling_mode(NVTEScalingMode val) noexcept { + return set_attr(kNVTEFusedAttnConfigScalingMode, val); + } + FusedAttnConfigWrapper &set_dropout(float val) noexcept { + return set_attr(kNVTEFusedAttnConfigDropout, val); + } + FusedAttnConfigWrapper &set_attn_scale(float val) noexcept { + return set_attr(kNVTEFusedAttnConfigAttnScale, val); + } + FusedAttnConfigWrapper &set_qkv_dtype(NVTEDType val) noexcept { + return set_attr(kNVTEFusedAttnConfigQKVDtype, val); + } + FusedAttnConfigWrapper &set_o_dtype(NVTEDType val) noexcept { + return set_attr(kNVTEFusedAttnConfigODtype, val); + } + FusedAttnConfigWrapper &set_do_dtype(NVTEDType val) noexcept { + return set_attr(kNVTEFusedAttnConfigDODtype, val); + } + FusedAttnConfigWrapper &set_dqkv_dtype(NVTEDType val) noexcept { + return set_attr(kNVTEFusedAttnConfigDQKVDtype, val); + } + FusedAttnConfigWrapper &set_qkv_layout(NVTE_QKV_Layout val) noexcept { + return set_attr(kNVTEFusedAttnConfigQKVLayout, val); + } + FusedAttnConfigWrapper &set_o_format(NVTE_QKV_Format val) noexcept { + return set_attr(kNVTEFusedAttnConfigOFormat, val); + } + FusedAttnConfigWrapper &set_do_format(NVTE_QKV_Format val) noexcept { + return set_attr(kNVTEFusedAttnConfigDOFormat, val); + } + FusedAttnConfigWrapper &set_dqkv_layout(NVTE_QKV_Layout val) noexcept { + return set_attr(kNVTEFusedAttnConfigDQKVLayout, val); + } + FusedAttnConfigWrapper &set_qkv_scale_inv_format(NVTE_QKV_Format val) noexcept { + return set_attr(kNVTEFusedAttnConfigQKVScaleInvFormat, val); + } + FusedAttnConfigWrapper &set_do_scale_inv_format(NVTE_QKV_Format val) noexcept { + return set_attr(kNVTEFusedAttnConfigDOScaleInvFormat, val); + } + FusedAttnConfigWrapper &set_batch_size(size_t val) noexcept { + return set_attr(kNVTEFusedAttnConfigBatchSize, val); + } + FusedAttnConfigWrapper &set_num_attn_heads(size_t val) noexcept { + return set_attr(kNVTEFusedAttnConfigNumAttnHeads, val); + } + FusedAttnConfigWrapper &set_num_gqa_groups(size_t val) noexcept { + return set_attr(kNVTEFusedAttnConfigNumGQAGroups, val); + } + FusedAttnConfigWrapper &set_head_dim_qk(size_t val) noexcept { + return set_attr(kNVTEFusedAttnConfigHeadDimQK, val); + } + FusedAttnConfigWrapper &set_head_dim_v(size_t val) noexcept { + return set_attr(kNVTEFusedAttnConfigHeadDimV, val); + } + FusedAttnConfigWrapper &set_max_seqlen_q(size_t val) noexcept { + return set_attr(kNVTEFusedAttnConfigMaxSeqlenQ, val); + } + FusedAttnConfigWrapper &set_max_seqlen_kv(size_t val) noexcept { + return set_attr(kNVTEFusedAttnConfigMaxSeqlenKV, val); + } + FusedAttnConfigWrapper &set_num_tokens_q(size_t val) noexcept { + return set_attr(kNVTEFusedAttnConfigNumTokensQ, val); + } + FusedAttnConfigWrapper &set_num_tokens_kv(size_t val) noexcept { + return set_attr(kNVTEFusedAttnConfigNumTokensKV, val); + } + FusedAttnConfigWrapper &set_num_pages_k(size_t val) noexcept { + return set_attr(kNVTEFusedAttnConfigNumPagesK, val); + } + FusedAttnConfigWrapper &set_num_pages_v(size_t val) noexcept { + return set_attr(kNVTEFusedAttnConfigNumPagesV, val); + } + FusedAttnConfigWrapper &set_page_size_k(size_t val) noexcept { + return set_attr(kNVTEFusedAttnConfigPageSizeK, val); + } + FusedAttnConfigWrapper &set_page_size_v(size_t val) noexcept { + return set_attr(kNVTEFusedAttnConfigPageSizeV, val); + } + FusedAttnConfigWrapper &set_max_pages_per_seq_k(size_t val) noexcept { + return set_attr(kNVTEFusedAttnConfigMaxPagesPerSeqK, val); + } + FusedAttnConfigWrapper &set_max_pages_per_seq_v(size_t val) noexcept { + return set_attr(kNVTEFusedAttnConfigMaxPagesPerSeqV, val); + } + FusedAttnConfigWrapper &set_bias_batch_size(size_t val) noexcept { + return set_attr(kNVTEFusedAttnConfigBiasBatchSize, val); + } + FusedAttnConfigWrapper &set_bias_num_heads(size_t val) noexcept { + return set_attr(kNVTEFusedAttnConfigBiasNumHeads, val); + } + FusedAttnConfigWrapper &set_bias_seqlen_q(size_t val) noexcept { + return set_attr(kNVTEFusedAttnConfigBiasSeqlenQ, val); + } + FusedAttnConfigWrapper &set_bias_seqlen_kv(size_t val) noexcept { + return set_attr(kNVTEFusedAttnConfigBiasSeqlenKV, val); + } + + private: + template + FusedAttnConfigWrapper &set_attr(NVTEFusedAttnConfigAttribute attr, T val) noexcept { + nvte_set_fused_attn_config_attribute(cfg_, attr, &val, sizeof(val)); + return *this; + } + + NVTEFusedAttnConfig cfg_ = nullptr; +}; + +/*! \class FusedAttnFwdParamsWrapper + * \brief C++ helper for constructing an `NVTEFusedAttnFwdParams`. + * + * It owns an opaque `NVTEFusedAttnFwdParams` handle created by + * `nvte_create_fused_attn_fwd_params()`, and provides a convenient, + * chainable interface for setting every field in `FusedAttnFwdParams`. + */ +class FusedAttnFwdParamsWrapper { + public: + FusedAttnFwdParamsWrapper() : params_{nvte_create_fused_attn_fwd_params()} {} + + FusedAttnFwdParamsWrapper(const FusedAttnFwdParamsWrapper &) = delete; + FusedAttnFwdParamsWrapper &operator=(const FusedAttnFwdParamsWrapper &) = delete; + + FusedAttnFwdParamsWrapper(FusedAttnFwdParamsWrapper &&other) noexcept : params_{other.params_} { + other.params_ = nullptr; + } + + FusedAttnFwdParamsWrapper &operator=(FusedAttnFwdParamsWrapper &&other) noexcept { + if (this != &other) { + if (params_ != nullptr) { + nvte_destroy_fused_attn_fwd_params(params_); + } + params_ = other.params_; + other.params_ = nullptr; + } + return *this; + } + + ~FusedAttnFwdParamsWrapper() { + if (params_ != nullptr) { + nvte_destroy_fused_attn_fwd_params(params_); + } + } + + operator NVTEFusedAttnFwdParams() const noexcept { return params_; } + NVTEFusedAttnFwdParams get() const noexcept { return params_; } + + FusedAttnFwdParamsWrapper &set_Q(NVTETensor val) noexcept { + return set_attr(kNVTEFusedAttnFwdParamsQ, val); + } + FusedAttnFwdParamsWrapper &set_K(NVTETensor val) noexcept { + return set_attr(kNVTEFusedAttnFwdParamsK, val); + } + FusedAttnFwdParamsWrapper &set_V(NVTETensor val) noexcept { + return set_attr(kNVTEFusedAttnFwdParamsV, val); + } + FusedAttnFwdParamsWrapper &set_Bias(NVTETensor val) noexcept { + return set_attr(kNVTEFusedAttnFwdParamsBias, val); + } + FusedAttnFwdParamsWrapper &set_SoftmaxOffset(NVTETensor val) noexcept { + return set_attr(kNVTEFusedAttnFwdParamsSoftmaxOffset, val); + } + FusedAttnFwdParamsWrapper &set_S(NVTETensor val) noexcept { + return set_attr(kNVTEFusedAttnFwdParamsS, val); + } + FusedAttnFwdParamsWrapper &set_O(NVTETensor val) noexcept { + return set_attr(kNVTEFusedAttnFwdParamsO, val); + } + FusedAttnFwdParamsWrapper &set_Aux_CTX_Tensors(NVTETensorPack *val) noexcept { + return set_attr(kNVTEFusedAttnFwdParamsAuxCtxTensors, val); + } + FusedAttnFwdParamsWrapper &set_cu_seqlens_q(NVTETensor val) noexcept { + return set_attr(kNVTEFusedAttnFwdParamsCuSeqlensQ, val); + } + FusedAttnFwdParamsWrapper &set_cu_seqlens_kv(NVTETensor val) noexcept { + return set_attr(kNVTEFusedAttnFwdParamsCuSeqlensKV, val); + } + FusedAttnFwdParamsWrapper &set_cu_seqlens_q_padded(NVTETensor val) noexcept { + return set_attr(kNVTEFusedAttnFwdParamsCuSeqlensQPadded, val); + } + FusedAttnFwdParamsWrapper &set_cu_seqlens_kv_padded(NVTETensor val) noexcept { + return set_attr(kNVTEFusedAttnFwdParamsCuSeqlensKVPadded, val); + } + FusedAttnFwdParamsWrapper &set_page_table_k(NVTETensor val) noexcept { + return set_attr(kNVTEFusedAttnFwdParamsPageTableK, val); + } + FusedAttnFwdParamsWrapper &set_page_table_v(NVTETensor val) noexcept { + return set_attr(kNVTEFusedAttnFwdParamsPageTableV, val); + } + FusedAttnFwdParamsWrapper &set_rng_state(NVTETensor val) noexcept { + return set_attr(kNVTEFusedAttnFwdParamsRngState, val); + } + FusedAttnFwdParamsWrapper &set_is_training(bool val) noexcept { + return set_attr(kNVTEFusedAttnFwdParamsIsTraining, static_cast(val)); + } + FusedAttnFwdParamsWrapper &set_cuda_graph(bool val) noexcept { + return set_attr(kNVTEFusedAttnFwdParamsCudaGraph, static_cast(val)); + } + FusedAttnFwdParamsWrapper &set_return_max_logit(bool val) noexcept { + return set_attr(kNVTEFusedAttnFwdParamsReturnMaxLogit, static_cast(val)); + } + FusedAttnFwdParamsWrapper &set_attn_mask_type(NVTE_Mask_Type val) noexcept { + return set_attr(kNVTEFusedAttnFwdParamsAttnMaskType, val); + } + FusedAttnFwdParamsWrapper &set_bias_type(NVTE_Bias_Type val) noexcept { + return set_attr(kNVTEFusedAttnFwdParamsBiasType, val); + } + FusedAttnFwdParamsWrapper &set_window_size_left(int64_t val) noexcept { + return set_attr(kNVTEFusedAttnFwdParamsWindowSizeLeft, val); + } + FusedAttnFwdParamsWrapper &set_window_size_right(int64_t val) noexcept { + return set_attr(kNVTEFusedAttnFwdParamsWindowSizeRight, val); + } + FusedAttnFwdParamsWrapper &set_bottom_right_diagonal(bool val) noexcept { + return set_attr(kNVTEFusedAttnFwdParamsBottomRightDiagonal, static_cast(val)); + } + FusedAttnFwdParamsWrapper &set_softmax_type(NVTE_Softmax_Type val) noexcept { + return set_attr(kNVTEFusedAttnFwdParamsSoftmaxType, val); + } + FusedAttnFwdParamsWrapper &set_dropout(float val) noexcept { + return set_attr(kNVTEFusedAttnFwdParamsDropout, val); + } + FusedAttnFwdParamsWrapper &set_attn_scale(float val) noexcept { + return set_attr(kNVTEFusedAttnFwdParamsAttnScale, val); + } + FusedAttnFwdParamsWrapper &set_qkv_layout(NVTE_QKV_Layout val) noexcept { + return set_attr(kNVTEFusedAttnFwdParamsQKVLayout, val); + } + FusedAttnFwdParamsWrapper &set_o_format(NVTE_QKV_Format val) noexcept { + return set_attr(kNVTEFusedAttnFwdParamsOFormat, val); + } + FusedAttnFwdParamsWrapper &set_qkv_scale_inv_format(NVTE_QKV_Format val) noexcept { + return set_attr(kNVTEFusedAttnFwdParamsQKVScaleInvFormat, val); + } + FusedAttnFwdParamsWrapper &set_max_seqlen_q(size_t val) noexcept { + return set_attr(kNVTEFusedAttnFwdParamsMaxSeqlenQ, val); + } + FusedAttnFwdParamsWrapper &set_max_seqlen_kv(size_t val) noexcept { + return set_attr(kNVTEFusedAttnFwdParamsMaxSeqlenKV, val); + } + FusedAttnFwdParamsWrapper &set_workspace(NVTETensor val) noexcept { + return set_attr(kNVTEFusedAttnFwdParamsWorkspace, val); + } + FusedAttnFwdParamsWrapper &set_stream(cudaStream_t val) noexcept { + return set_attr(kNVTEFusedAttnFwdParamsStream, val); + } + + private: + template + FusedAttnFwdParamsWrapper &set_attr(NVTEFusedAttnFwdParamsAttribute attr, T val) noexcept { + nvte_set_fused_attn_fwd_params_attribute(params_, attr, &val, sizeof(val)); + return *this; + } + + NVTEFusedAttnFwdParams params_ = nullptr; +}; + +/*! \class FusedAttnBwdParamsWrapper + * \brief C++ helper for constructing an `NVTEFusedAttnBwdParams`. + * + * It owns an opaque `NVTEFusedAttnBwdParams` handle created by + * `nvte_create_fused_attn_bwd_params()`, and provides a convenient, + * chainable interface for setting every field in `FusedAttnBwdParams`. + */ +class FusedAttnBwdParamsWrapper { + public: + FusedAttnBwdParamsWrapper() : params_{nvte_create_fused_attn_bwd_params()} {} + + FusedAttnBwdParamsWrapper(const FusedAttnBwdParamsWrapper &) = delete; + FusedAttnBwdParamsWrapper &operator=(const FusedAttnBwdParamsWrapper &) = delete; + + FusedAttnBwdParamsWrapper(FusedAttnBwdParamsWrapper &&other) noexcept : params_{other.params_} { + other.params_ = nullptr; + } + + FusedAttnBwdParamsWrapper &operator=(FusedAttnBwdParamsWrapper &&other) noexcept { + if (this != &other) { + if (params_ != nullptr) { + nvte_destroy_fused_attn_bwd_params(params_); + } + params_ = other.params_; + other.params_ = nullptr; + } + return *this; + } + + ~FusedAttnBwdParamsWrapper() { + if (params_ != nullptr) { + nvte_destroy_fused_attn_bwd_params(params_); + } + } + + operator NVTEFusedAttnBwdParams() const noexcept { return params_; } + NVTEFusedAttnBwdParams get() const noexcept { return params_; } + + FusedAttnBwdParamsWrapper &set_Q(NVTETensor val) noexcept { + return set_attr(kNVTEFusedAttnBwdParamsQ, val); + } + FusedAttnBwdParamsWrapper &set_K(NVTETensor val) noexcept { + return set_attr(kNVTEFusedAttnBwdParamsK, val); + } + FusedAttnBwdParamsWrapper &set_V(NVTETensor val) noexcept { + return set_attr(kNVTEFusedAttnBwdParamsV, val); + } + FusedAttnBwdParamsWrapper &set_O(NVTETensor val) noexcept { + return set_attr(kNVTEFusedAttnBwdParamsO, val); + } + FusedAttnBwdParamsWrapper &set_dO(NVTETensor val) noexcept { + return set_attr(kNVTEFusedAttnBwdParamsDO, val); + } + FusedAttnBwdParamsWrapper &set_S(NVTETensor val) noexcept { + return set_attr(kNVTEFusedAttnBwdParamsS, val); + } + FusedAttnBwdParamsWrapper &set_dP(NVTETensor val) noexcept { + return set_attr(kNVTEFusedAttnBwdParamsDP, val); + } + FusedAttnBwdParamsWrapper &set_Aux_CTX_Tensors(const NVTETensorPack *val) noexcept { + return set_attr(kNVTEFusedAttnBwdParamsAuxCtxTensors, val); + } + FusedAttnBwdParamsWrapper &set_dQ(NVTETensor val) noexcept { + return set_attr(kNVTEFusedAttnBwdParamsDQ, val); + } + FusedAttnBwdParamsWrapper &set_dK(NVTETensor val) noexcept { + return set_attr(kNVTEFusedAttnBwdParamsDK, val); + } + FusedAttnBwdParamsWrapper &set_dV(NVTETensor val) noexcept { + return set_attr(kNVTEFusedAttnBwdParamsDV, val); + } + FusedAttnBwdParamsWrapper &set_dBias(NVTETensor val) noexcept { + return set_attr(kNVTEFusedAttnBwdParamsDBias, val); + } + FusedAttnBwdParamsWrapper &set_dSoftmaxOffset(NVTETensor val) noexcept { + return set_attr(kNVTEFusedAttnBwdParamsDSoftmaxOffset, val); + } + FusedAttnBwdParamsWrapper &set_cu_seqlens_q(NVTETensor val) noexcept { + return set_attr(kNVTEFusedAttnBwdParamsCuSeqlensQ, val); + } + FusedAttnBwdParamsWrapper &set_cu_seqlens_kv(NVTETensor val) noexcept { + return set_attr(kNVTEFusedAttnBwdParamsCuSeqlensKV, val); + } + FusedAttnBwdParamsWrapper &set_cu_seqlens_q_padded(NVTETensor val) noexcept { + return set_attr(kNVTEFusedAttnBwdParamsCuSeqlensQPadded, val); + } + FusedAttnBwdParamsWrapper &set_cu_seqlens_kv_padded(NVTETensor val) noexcept { + return set_attr(kNVTEFusedAttnBwdParamsCuSeqlensKVPadded, val); + } + FusedAttnBwdParamsWrapper &set_deterministic(bool val) noexcept { + return set_attr(kNVTEFusedAttnBwdParamsDeterministic, static_cast(val)); + } + FusedAttnBwdParamsWrapper &set_cuda_graph(bool val) noexcept { + return set_attr(kNVTEFusedAttnBwdParamsCudaGraph, static_cast(val)); + } + FusedAttnBwdParamsWrapper &set_attn_mask_type(NVTE_Mask_Type val) noexcept { + return set_attr(kNVTEFusedAttnBwdParamsAttnMaskType, val); + } + FusedAttnBwdParamsWrapper &set_bias_type(NVTE_Bias_Type val) noexcept { + return set_attr(kNVTEFusedAttnBwdParamsBiasType, val); + } + FusedAttnBwdParamsWrapper &set_window_size_left(int64_t val) noexcept { + return set_attr(kNVTEFusedAttnBwdParamsWindowSizeLeft, val); + } + FusedAttnBwdParamsWrapper &set_window_size_right(int64_t val) noexcept { + return set_attr(kNVTEFusedAttnBwdParamsWindowSizeRight, val); + } + FusedAttnBwdParamsWrapper &set_bottom_right_diagonal(bool val) noexcept { + return set_attr(kNVTEFusedAttnBwdParamsBottomRightDiagonal, static_cast(val)); + } + FusedAttnBwdParamsWrapper &set_softmax_type(NVTE_Softmax_Type val) noexcept { + return set_attr(kNVTEFusedAttnBwdParamsSoftmaxType, val); + } + FusedAttnBwdParamsWrapper &set_dropout(float val) noexcept { + return set_attr(kNVTEFusedAttnBwdParamsDropout, val); + } + FusedAttnBwdParamsWrapper &set_attn_scale(float val) noexcept { + return set_attr(kNVTEFusedAttnBwdParamsAttnScale, val); + } + FusedAttnBwdParamsWrapper &set_qkv_layout(NVTE_QKV_Layout val) noexcept { + return set_attr(kNVTEFusedAttnBwdParamsQKVLayout, val); + } + FusedAttnBwdParamsWrapper &set_o_format(NVTE_QKV_Format val) noexcept { + return set_attr(kNVTEFusedAttnBwdParamsOFormat, val); + } + FusedAttnBwdParamsWrapper &set_do_format(NVTE_QKV_Format val) noexcept { + return set_attr(kNVTEFusedAttnBwdParamsDOFormat, val); + } + FusedAttnBwdParamsWrapper &set_dqkv_layout(NVTE_QKV_Layout val) noexcept { + return set_attr(kNVTEFusedAttnBwdParamsDQKVLayout, val); + } + FusedAttnBwdParamsWrapper &set_qkv_scale_inv_format(NVTE_QKV_Format val) noexcept { + return set_attr(kNVTEFusedAttnBwdParamsQKVScaleInvFormat, val); + } + FusedAttnBwdParamsWrapper &set_do_scale_inv_format(NVTE_QKV_Format val) noexcept { + return set_attr(kNVTEFusedAttnBwdParamsDOScaleInvFormat, val); + } + FusedAttnBwdParamsWrapper &set_max_seqlen_q(size_t val) noexcept { + return set_attr(kNVTEFusedAttnBwdParamsMaxSeqlenQ, val); + } + FusedAttnBwdParamsWrapper &set_max_seqlen_kv(size_t val) noexcept { + return set_attr(kNVTEFusedAttnBwdParamsMaxSeqlenKV, val); + } + FusedAttnBwdParamsWrapper &set_workspace(NVTETensor val) noexcept { + return set_attr(kNVTEFusedAttnBwdParamsWorkspace, val); + } + FusedAttnBwdParamsWrapper &set_stream(cudaStream_t val) noexcept { + return set_attr(kNVTEFusedAttnBwdParamsStream, val); + } + + private: + template + FusedAttnBwdParamsWrapper &set_attr(NVTEFusedAttnBwdParamsAttribute attr, T val) noexcept { + nvte_set_fused_attn_bwd_params_attribute(params_, attr, &val, sizeof(val)); + return *this; + } + + NVTEFusedAttnBwdParams params_ = nullptr; +}; #endif // __cplusplus #endif diff --git a/transformer_engine/common/util/pybind_helper.h b/transformer_engine/common/util/pybind_helper.h index f7ffb5ad8d..d739965163 100644 --- a/transformer_engine/common/util/pybind_helper.h +++ b/transformer_engine/common/util/pybind_helper.h @@ -86,11 +86,19 @@ .value("NVTE_Paged_KV_SBHD_SBHD_SBHD", NVTE_QKV_Layout::NVTE_Paged_KV_SBHD_SBHD_SBHD) \ .value("NVTE_Paged_KV_THD_BSHD_BSHD", NVTE_QKV_Layout::NVTE_Paged_KV_THD_BSHD_BSHD) \ .value("NVTE_Paged_KV_THD_SBHD_SBHD", NVTE_QKV_Layout::NVTE_Paged_KV_THD_SBHD_SBHD) \ - .value("NVTE_BHSD_BHSD_BHSD", NVTE_QKV_Layout::NVTE_BHSD_BHSD_BHSD); \ + .value("NVTE_BHSD_BHSD_BHSD", NVTE_QKV_Layout::NVTE_BHSD_BHSD_BHSD) \ + .value("NVTE_QKV_Layout_NOT_SET", NVTE_QKV_Layout::NVTE_QKV_Layout_NOT_SET); \ pybind11::enum_(m, "NVTE_Fused_Attn_Backend", pybind11::module_local()) \ .value("NVTE_F16_arbitrary_seqlen", NVTE_Fused_Attn_Backend::NVTE_F16_arbitrary_seqlen) \ .value("NVTE_FP8", NVTE_Fused_Attn_Backend::NVTE_FP8) \ .value("NVTE_No_Backend", NVTE_Fused_Attn_Backend::NVTE_No_Backend); \ + pybind11::enum_(m, "NVTEScalingMode", pybind11::module_local()) \ + .value("NVTE_DELAYED_TENSOR_SCALING", NVTEScalingMode::NVTE_DELAYED_TENSOR_SCALING) \ + .value("NVTE_MXFP8_1D_SCALING", NVTEScalingMode::NVTE_MXFP8_1D_SCALING) \ + .value("NVTE_BLOCK_SCALING_1D", NVTEScalingMode::NVTE_BLOCK_SCALING_1D) \ + .value("NVTE_BLOCK_SCALING_2D", NVTEScalingMode::NVTE_BLOCK_SCALING_2D) \ + .value("NVTE_NVFP4_1D_SCALING", NVTEScalingMode::NVTE_NVFP4_1D_SCALING) \ + .value("NVTE_INVALID_SCALING", NVTEScalingMode::NVTE_INVALID_SCALING); \ pybind11::enum_( \ m, "Float8BlockScaleTensorFormat", pybind11::module_local()) \ .value("GEMM_READY", transformer_engine::Float8BlockScaleTensorFormat::GEMM_READY) \ diff --git a/transformer_engine/jax/attention.py b/transformer_engine/jax/attention.py index 6d7c823e12..4e7f097d7a 100644 --- a/transformer_engine/jax/attention.py +++ b/transformer_engine/jax/attention.py @@ -325,6 +325,7 @@ def canonicalize_attn_mask_type(attn_mask_type: str): def is_fused_attn_kernel_available( is_training, + batch_size, q_dtype, kv_dtype, qkv_layout, @@ -340,15 +341,26 @@ def is_fused_attn_kernel_available( head_dim_v, window_size: Optional[Tuple[int, int]] = None, return_max_logit: bool = False, + bottom_right_diagonal: Optional[bool] = None, + bias_batch: Optional[int] = None, + bias_heads: Optional[int] = None, + bias_seqlen_q: Optional[int] = None, + bias_seqlen_kv: Optional[int] = None, ): """ - To check whether the fused attention kernel is supported + To check whether the fused attention kernel is supported. """ window_size_tuple = (-1, -1) if window_size is None else window_size def make_helper(attn_mask_type): + bottom_right = ( + attn_mask_type.is_bottom_right() + if bottom_right_diagonal is None + else bottom_right_diagonal + ) return tex.FusedAttnHelper( is_training, + batch_size, q_dtype, kv_dtype, qkv_layout, @@ -363,10 +375,16 @@ def make_helper(attn_mask_type): head_dim_qk, head_dim_v, window_size_tuple, - return_max_logit, + return_max_logit=return_max_logit, + bottom_right_diagonal=bottom_right, + bias_batch=bias_batch, + bias_heads=bias_heads, + bias_seqlen_q=bias_seqlen_q, + bias_seqlen_kv=bias_seqlen_kv, ) - return make_helper(attn_mask_type).is_fused_attn_kernel_available() + helper = make_helper(attn_mask_type) + return helper.is_fused_attn_kernel_available() def _obtain_batch_and_max_seqlen(qkv, qkv_layout): diff --git a/transformer_engine/jax/cpp_extensions/attention.py b/transformer_engine/jax/cpp_extensions/attention.py index 6ea54195ab..b8e9475de8 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 import warnings @@ -16,7 +17,16 @@ from jax.experimental.custom_partitioning import SdyShardingRule import transformer_engine_jax -from transformer_engine_jax import NVTE_Fused_Attn_Backend +from transformer_engine_jax import ( + DType, + JAXX_Scaling_Mode, + NVTE_Bias_Type, + NVTE_Fused_Attn_Backend, + NVTE_Mask_Type, + NVTE_QKV_Format, + NVTE_QKV_Layout, + NVTE_Softmax_Type, +) from transformer_engine.jax.attention import ( AttnBiasType, AttnMaskType, @@ -57,6 +67,37 @@ ] +# NVTE_DEBUG = 0/1 # disables/enables debug mode, default = 0 +_NVTE_DEBUG = int(os.getenv("NVTE_DEBUG", "0")) +# NVTE_DEBUG_LEVEL = 0/1/2 # enables increasingly verbose debug messages, default = 0 +_NVTE_DEBUG_LEVEL = int(os.getenv("NVTE_DEBUG_LEVEL", "0")) + + +class AttentionLogging: + """Logging for the JAX attention module""" + + _log_level = _NVTE_DEBUG * _NVTE_DEBUG_LEVEL + _formatter = logging.Formatter("[%(levelname)-8s | %(name)-19s]: %(message)s") + _stream_handler = logging.StreamHandler() + logger = logging.getLogger(__name__) + _is_logging_setup = False + + @staticmethod + def setup_logging(): + """Set up log levels, logger and handlers (idempotent).""" + if AttentionLogging._is_logging_setup: + return + _log_levels = {0: logging.WARNING, 1: logging.INFO, 2: logging.DEBUG} + AttentionLogging._log_level = _log_levels[ + AttentionLogging._log_level if AttentionLogging._log_level in [0, 1, 2] else 2 + ] + AttentionLogging._stream_handler.setFormatter(AttentionLogging._formatter) + AttentionLogging.logger.setLevel(AttentionLogging._log_level) + if not AttentionLogging.logger.hasHandlers(): + AttentionLogging.logger.addHandler(AttentionLogging._stream_handler) + AttentionLogging._is_logging_setup = True + + @partial( jax.tree_util.register_dataclass, data_fields=[], @@ -103,6 +144,64 @@ class _FusedAttnConfig: return_max_logit: bool = False +@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) class FusedAttnHelper: """ @@ -110,6 +209,7 @@ class FusedAttnHelper: """ is_training: bool + batch_size: int q_dtype: jnp.dtype kv_dtype: jnp.dtype qkv_layout: QKVLayout @@ -125,34 +225,91 @@ class FusedAttnHelper: head_dim_v: int window_size: Tuple[int, int] return_max_logit: bool = False + bottom_right_diagonal: bool = True + attn_scale: float = 1.0 + bias_batch: Optional[int] = None + bias_heads: Optional[int] = None + bias_seqlen_q: Optional[int] = None + bias_seqlen_kv: Optional[int] = None def is_fused_attn_kernel_available(self): """Check if there is available fused attention kernel""" - return self.get_fused_attn_backend() != NVTE_Fused_Attn_Backend.NVTE_No_Backend + backend, _ = self.get_fused_attn_backend() + return backend != NVTE_Fused_Attn_Backend.NVTE_No_Backend def get_fused_attn_backend(self): - """Get the fused attention kernel backend""" - return transformer_engine_jax.get_fused_attn_backend( - self.is_training, - jax_dtype_to_te_dtype(self.q_dtype), - jax_dtype_to_te_dtype(self.kv_dtype), - self.qkv_layout.value, - self.attn_bias_type.value, - self.attn_mask_type.value, - self.softmax_type.value, - self.dropout_probability, - self.q_num_heads, - self.kv_num_heads, - self.q_max_seqlen, - self.kv_max_seqlen, - self.head_dim_qk, - self.head_dim_v, - self.window_size[0], - self.window_size[1], - self.return_max_logit, - not self.is_non_deterministic_allowed(), + """Get the fused attention kernel backend. + + Returns a ``(backend, message)`` tuple. ``message`` is empty on success, otherwise a + diagnostic string explaining why the configuration was rejected. + + When ``NVTE_DEBUG=1``, ``NVTE_DEBUG_LEVEL=1`` logs the outcome (the selected backend, or + that no fused backend is available), and ``NVTE_DEBUG_LEVEL=2`` additionally logs the + resolved config and the reason fused attention was rejected. + """ + q_type = jax_dtype_to_te_dtype(self.q_dtype) + kv_type = jax_dtype_to_te_dtype(self.kv_dtype) + if q_type != kv_type: + raise ValueError("Q and KV must have the same data type.") + bias_batch = bias_heads = bias_seqlen_q = bias_seqlen_kv = 0 + if self.attn_bias_type == AttnBiasType.POST_SCALE_BIAS: + bias_batch = self.bias_batch or 0 + bias_heads = self.bias_heads or 0 + bias_seqlen_q = self.bias_seqlen_q or 0 + bias_seqlen_kv = self.bias_seqlen_kv or 0 + num_tokens_q = num_tokens_kv = 0 + if self.qkv_layout.is_thd(): + num_tokens_q = self.batch_size * self.q_max_seqlen + num_tokens_kv = self.batch_size * self.kv_max_seqlen + backend, message = transformer_engine_jax.get_fused_attn_backend( + FusedAttnParams( + is_training=self.is_training, + deterministic=not self.is_non_deterministic_allowed(), + return_max_logit=self.return_max_logit, + attn_mask_type=self.attn_mask_type.value, + bias_type=self.attn_bias_type.value, + window_size_left=self.window_size[0], + window_size_right=self.window_size[1], + bottom_right_diagonal=self.bottom_right_diagonal, + softmax_type=self.softmax_type.value, + dropout=self.dropout_probability, + attn_scale=self.attn_scale, + qkv_dtype=q_type, + o_dtype=q_type, + do_dtype=q_type, + dqkv_dtype=q_type, + qkv_layout=self.qkv_layout.value, + batch_size=self.batch_size, + 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, + ) ) + AttentionLogging.setup_logging() + logger = AttentionLogging.logger + logger.debug("Running fused attention backend selection with config=%s", self) + if backend == NVTE_Fused_Attn_Backend.NVTE_No_Backend: + logger.info("No fused attention backend available; falling back to unfused attention.") + logger.debug( + "Reason fused attention was rejected: %s", + message or "(no diagnostic message available)", + ) + else: + logger.info("Selected fused attention backend: %s", backend) + if message: + logger.debug("Fused attention backend diagnostic message: %s", message) + return backend, message + @staticmethod def is_non_deterministic_allowed(): """Check if non-deterministic kernels are allowed""" @@ -339,8 +496,17 @@ def abstract( out_aval = q_aval.update(shape=output_shape, dtype=q_dtype) # backend determines the softmax buffer shape/dtype - backend = FusedAttnHelper( + input_batch = reduce(operator.mul, batch_shape) + bias_batch = bias_heads = bias_seqlen_q = bias_seqlen_kv = None + if config.attn_bias_type == AttnBiasType.POST_SCALE_BIAS: + *bias_batch_shape, bias_heads, bias_seqlen_q, bias_seqlen_kv = bias_aval.shape + bias_batch = reduce(operator.mul, bias_batch_shape) + num_segments = input_batch + if config.qkv_layout.is_thd(): + num_segments = input_batch * config.max_segments_per_seq + backend, message = FusedAttnHelper( config.is_training, + num_segments, q_dtype, k_dtype, config.qkv_layout, @@ -355,7 +521,13 @@ def abstract( q_head_dim, v_head_dim, config.window_size, - config.return_max_logit, + return_max_logit=config.return_max_logit, + bottom_right_diagonal=config.bottom_right_diagonal, + attn_scale=float(config.scaling_factor), + bias_batch=bias_batch, + bias_heads=bias_heads, + bias_seqlen_q=bias_seqlen_q, + bias_seqlen_kv=bias_seqlen_kv, ).get_fused_attn_backend() if backend == NVTE_Fused_Attn_Backend.NVTE_F16_arbitrary_seqlen: @@ -374,7 +546,7 @@ def abstract( ) softmax_dtype = dtypes.canonicalize_dtype(jnp.float32) else: - raise ValueError(f"Unsupported {backend=}") + 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 @@ -404,11 +576,6 @@ def abstract( *bias_batch_shape, bias_heads, _, _ = bias_aval.shape bias_batch = reduce(operator.mul, bias_batch_shape) - bottom_right_diagonal = config.attn_mask_type in [ - AttnMaskType.CAUSAL_BOTTOM_RIGHT_MASK, - AttnMaskType.PADDING_CAUSAL_BOTTOM_RIGHT_MASK, - ] - # 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) @@ -434,7 +601,7 @@ def abstract( config.window_size[0], config.window_size[1], config.return_max_logit, - bottom_right_diagonal, + config.bottom_right_diagonal, ) wkspace_aval = q_aval.update( shape=wkspace_info[0], dtype=te_dtype_to_jax_dtype(wkspace_info[1]) diff --git a/transformer_engine/jax/csrc/extensions.h b/transformer_engine/jax/csrc/extensions.h index 580219baf2..a7bf8d21d5 100644 --- a/transformer_engine/jax/csrc/extensions.h +++ b/transformer_engine/jax/csrc/extensions.h @@ -23,6 +23,7 @@ #include #include #include +#include #include #include "common/common.h" @@ -151,12 +152,8 @@ XLA_FFI_DECLARE_HANDLER_SYMBOL(FusedAttnScoreModForwardHandler); XLA_FFI_DECLARE_HANDLER_SYMBOL(FusedAttnScoreModBackwardHandler); -NVTE_Fused_Attn_Backend GetFusedAttnBackend( - bool is_training, DType q_dtype, DType kv_dtype, NVTE_QKV_Layout qkv_layout, - NVTE_Bias_Type bias_type, NVTE_Mask_Type mask_type, NVTE_Softmax_Type softmax_type, - float dropout_probability, size_t q_attn_heads, size_t kv_attn_heads, size_t q_max_seqlen, - size_t kv_max_seqlen, size_t qk_head_dim, size_t v_head_dim, int64_t window_size_left, - int64_t window_size_right, bool return_max_logit, bool deterministic); +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, diff --git a/transformer_engine/jax/csrc/extensions/attention.cpp b/transformer_engine/jax/csrc/extensions/attention.cpp index 4af463f236..3bed7d994b 100644 --- a/transformer_engine/jax/csrc/extensions/attention.cpp +++ b/transformer_engine/jax/csrc/extensions/attention.cpp @@ -24,18 +24,98 @@ namespace transformer_engine { namespace jax { -NVTE_Fused_Attn_Backend GetFusedAttnBackend( - bool is_training, DType q_dtype, DType kv_dtype, NVTE_QKV_Layout qkv_layout, - NVTE_Bias_Type bias_type, NVTE_Mask_Type mask_type, NVTE_Softmax_Type softmax_type, - float dropout_probability, size_t q_attn_heads, size_t kv_attn_heads, size_t q_max_seqlen, - size_t kv_max_seqlen, size_t qk_head_dim, size_t v_head_dim, int64_t window_size_left, - int64_t window_size_right, bool return_max_logit, bool deterministic) { - auto backend = nvte_get_fused_attn_backend( - is_training, static_cast(q_dtype), static_cast(kv_dtype), qkv_layout, - bias_type, mask_type, softmax_type, dropout_probability, q_attn_heads, kv_attn_heads, - q_max_seqlen, kv_max_seqlen, qk_head_dim, v_head_dim, window_size_left, window_size_right, - return_max_logit, false, deterministic); - return backend; +static std::tuple GetFusedAttnBackendImpl( + bool is_training, bool deterministic, bool cuda_graph, bool return_max_logit, + NVTE_Mask_Type mask_type, NVTE_Bias_Type bias_type, int64_t window_size_left, + int64_t window_size_right, bool bottom_right_diagonal, NVTE_Softmax_Type softmax_type, + JAXX_Scaling_Mode scaling_mode, float dropout_probability, float attn_scale, DType q_dtype, + DType o_dtype, DType do_dtype, DType dqkv_dtype, 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, size_t batch_size, + size_t q_attn_heads, size_t kv_attn_heads, size_t qk_head_dim, size_t v_head_dim, + size_t q_max_seqlen, size_t kv_max_seqlen, size_t num_tokens_q, size_t num_tokens_kv, + size_t bias_batch, size_t bias_heads, size_t bias_seqlen_q, size_t bias_seqlen_kv) { + FusedAttnConfigWrapper cfg; + cfg.set_is_training(is_training) + .set_deterministic(deterministic) + .set_cuda_graph(cuda_graph) + .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_scaling_mode(get_nvte_scaling_mode(scaling_mode)) + .set_dropout(dropout_probability) + .set_attn_scale(attn_scale) + .set_qkv_dtype(static_cast(q_dtype)) + .set_o_dtype(static_cast(o_dtype)) + .set_do_dtype(static_cast(do_dtype)) + .set_dqkv_dtype(static_cast(dqkv_dtype)) + .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_batch_size(batch_size) + .set_num_attn_heads(q_attn_heads) + .set_num_gqa_groups(kv_attn_heads) + .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(num_tokens_q) + .set_num_tokens_kv(num_tokens_kv) + .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); + + 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; + } + + return GetFusedAttnBackendImpl( + params.attr("is_training").cast(), params.attr("deterministic").cast(), + params.attr("cuda_graph").cast(), params.attr("return_max_logit").cast(), + params.attr("attn_mask_type").cast(), + params.attr("bias_type").cast(), + params.attr("window_size_left").cast(), + params.attr("window_size_right").cast(), + params.attr("bottom_right_diagonal").cast(), + params.attr("softmax_type").cast(), + params.attr("scaling_mode").cast(), params.attr("dropout").cast(), + params.attr("attn_scale").cast(), params.attr("qkv_dtype").cast(), + params.attr("o_dtype").cast(), params.attr("do_dtype").cast(), + params.attr("dqkv_dtype").cast(), qkv_layout, o_format, do_format, dqkv_layout, + params.attr("qkv_scale_inv_format").cast(), + params.attr("do_scale_inv_format").cast(), + params.attr("batch_size").cast(), params.attr("num_attn_heads").cast(), + params.attr("num_gqa_groups").cast(), params.attr("head_dim_qk").cast(), + params.attr("head_dim_v").cast(), params.attr("max_seqlen_q").cast(), + params.attr("max_seqlen_kv").cast(), params.attr("num_tokens_q").cast(), + params.attr("num_tokens_kv").cast(), params.attr("bias_batch_size").cast(), + params.attr("bias_num_heads").cast(), params.attr("bias_seqlen_q").cast(), + params.attr("bias_seqlen_kv").cast()); } /* @@ -211,15 +291,41 @@ pybind11::tuple GetFusedAttnForwardWorkspaceSizes( TensorWrapper(nullptr, std::vector{num_segments + 1}, DType::kInt32); auto ragged_offset_tensor = TensorWrapper(nullptr, std::vector{num_segments + 1}, DType::kInt32); - nvte_fused_attn_fwd( - q_tensor.data(), k_tensor.data(), v_tensor.data(), bias_tensor.data(), - dummy_softmax_offset_tensor.data(), s_tensor.data(), o_tensor.data(), &aux_output_tensors, - q_cu_seqlens_tensor.data(), kv_cu_seqlens_tensor.data(), ragged_offset_tensor.data(), - ragged_offset_tensor.data(), dummy_page_table_tensor.data(), dummy_page_table_tensor.data(), - dummy_rng_state_tensor.data(), q_max_seqlen, kv_max_seqlen, is_training, return_max_logit, - false, scaling_factor, dropout_probability, qkv_layout, nvte_get_q_format(qkv_layout), - NVTE_QKV_Format_NOT_SET, bias_type, mask_type, softmax_type, window_size_left, - window_size_right, bottom_right_diagonal, query_workspace_tensor.data(), nullptr); + 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); @@ -231,6 +337,10 @@ pybind11::tuple GetFusedAttnForwardWorkspaceSizes( #define FUSED_ATTN_IMPL_COMMON_BLOCK \ auto is_ragged = nvte_get_qkv_format(qkv_layout) == NVTE_QKV_Format::NVTE_THD; \ auto bias_shape = std::vector{bias_batch, bias_heads, q_max_seqlen, kv_max_seqlen}; \ + const bool has_bias_tensor = \ + bias_type != NVTE_Bias_Type::NVTE_NO_BIAS && bias_type != NVTE_Bias_Type::NVTE_ALIBI; \ + const size_t bias_seqlen_q = has_bias_tensor ? q_max_seqlen : 0; \ + const size_t bias_seqlen_kv = has_bias_tensor ? kv_max_seqlen : 0; \ size_t num_segments = input_batch; \ if (is_ragged) { \ auto cudnn_runtime_version = cudnnGetVersion(); \ @@ -294,11 +404,16 @@ static void FusedAttnForwardImpl( /* Prepare RNG state */ auto rng_state_tensor = TensorWrapper(rng_state, std::vector{2}, DType::kInt64); - auto backend = nvte_get_fused_attn_backend( - is_training, static_cast(dtype), static_cast(dtype), qkv_layout, - bias_type, mask_type, softmax_type, dropout_probability, attn_heads, num_gqa_groups, - q_max_seqlen, kv_max_seqlen, qk_head_dim, v_head_dim, window_size_left, window_size_right, - return_max_logit, false, deterministic); + auto [backend, _fwd_msg] = GetFusedAttnBackendImpl( + is_training, deterministic, false, return_max_logit, mask_type, bias_type, window_size_left, + window_size_right, bottom_right_diagonal, softmax_type, JAXX_Scaling_Mode::NO_SCALING, + dropout_probability, scaling_factor, dtype, dtype, dtype, dtype, qkv_layout, + nvte_get_q_format(qkv_layout), nvte_get_q_format(qkv_layout), qkv_layout, + NVTE_QKV_Format::NVTE_QKV_Format_NOT_SET, NVTE_QKV_Format::NVTE_QKV_Format_NOT_SET, + num_segments, attn_heads, num_gqa_groups, qk_head_dim, v_head_dim, q_max_seqlen, + kv_max_seqlen, is_ragged ? input_batch * q_max_seqlen : 0, + is_ragged ? input_batch * kv_max_seqlen : 0, bias_batch, bias_heads, bias_seqlen_q, + bias_seqlen_kv); 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) */ @@ -360,15 +475,41 @@ static void FusedAttnForwardImpl( auto k_tensor = TensorWrapper(k_ptr, k_shape, dtype); auto v_tensor = TensorWrapper(v_ptr, v_shape, dtype); - nvte_fused_attn_fwd( - q_tensor.data(), k_tensor.data(), v_tensor.data(), bias_tensor.data(), - softmax_offset_tensor.data(), s_tensor.data(), o_tensor.data(), &aux_output_tensors, - q_cu_seqlens_tensor.data(), kv_cu_seqlens_tensor.data(), q_seq_offsets_tensor.data(), - k_seq_offsets_tensor.data(), dummy_page_table_tensor.data(), dummy_page_table_tensor.data(), - rng_state_tensor.data(), q_max_seqlen, kv_max_seqlen, is_training, return_max_logit, false, - scaling_factor, dropout_probability, qkv_layout, nvte_get_q_format(qkv_layout), - NVTE_QKV_Format_NOT_SET, bias_type, mask_type, softmax_type, window_size_left, - window_size_right, bottom_right_diagonal, workspace_tensor.data(), stream); + 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); } @@ -520,19 +661,45 @@ pybind11::tuple GetFusedAttnBackwardWorkspaceSizes( auto dummy_ragged_offset_tensor = TensorWrapper(nullptr, std::vector{num_segments + 1}, DType::kInt32); - nvte_fused_attn_bwd( - q_tensor.data(), k_tensor.data(), v_tensor.data(), output_tensor.data(), - doutput_tensor.data(), - s_tensor.data(), // not used for F16 - s_tensor.data(), // not used for F16 - &aux_input_tensors, dq_tensor.data(), dk_tensor.data(), dv_tensor.data(), - dbias_tensor.data(), dummy_d_softmax_offset_tensor.data(), q_cu_seqlens_tensor.data(), - kv_cu_seqlens_tensor.data(), dummy_ragged_offset_tensor.data(), - dummy_ragged_offset_tensor.data(), q_max_seqlen, kv_max_seqlen, scaling_factor, - dropout_probability, qkv_layout, nvte_get_q_format(qkv_layout), - nvte_get_q_format(qkv_layout), qkv_layout, NVTE_QKV_Format_NOT_SET, NVTE_QKV_Format_NOT_SET, - bias_type, mask_type, softmax_type, window_size_left, window_size_right, - bottom_right_diagonal, deterministic, false, query_workspace_tensor.data(), nullptr); + 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); @@ -575,11 +742,16 @@ static void FusedAttnBackwardImpl( /* Auxiliary tensors (propagated from the forward pass) */ NVTETensorPack aux_input_tensors; nvte_tensor_pack_create(&aux_input_tensors); - auto backend = nvte_get_fused_attn_backend( - is_training, static_cast(dtype), static_cast(dtype), qkv_layout, - bias_type, mask_type, softmax_type, dropout_probability, attn_heads, num_gqa_groups, - q_max_seqlen, kv_max_seqlen, qk_head_dim, v_head_dim, window_size_left, window_size_right, - false, false, deterministic); + auto [backend, _bwd_msg] = GetFusedAttnBackendImpl( + is_training, deterministic, false, false, mask_type, bias_type, window_size_left, + window_size_right, bottom_right_diagonal, softmax_type, JAXX_Scaling_Mode::NO_SCALING, + dropout_probability, scaling_factor, dtype, dtype, dtype, dtype, qkv_layout, + nvte_get_q_format(qkv_layout), nvte_get_q_format(qkv_layout), qkv_layout, + NVTE_QKV_Format::NVTE_QKV_Format_NOT_SET, NVTE_QKV_Format::NVTE_QKV_Format_NOT_SET, + num_segments, attn_heads, num_gqa_groups, qk_head_dim, v_head_dim, q_max_seqlen, + kv_max_seqlen, is_ragged ? input_batch * q_max_seqlen : 0, + is_ragged ? input_batch * kv_max_seqlen : 0, bias_batch, bias_heads, bias_seqlen_q, + bias_seqlen_kv); 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); @@ -656,18 +828,45 @@ static void FusedAttnBackwardImpl( } } - nvte_fused_attn_bwd( - q_tensor.data(), k_tensor.data(), v_tensor.data(), output_tensor.data(), - doutput_tensor.data(), - s_tensor.data(), // not used for F16 - s_tensor.data(), // not used for F16 - &aux_input_tensors, dq_tensor.data(), dk_tensor.data(), dv_tensor.data(), dbias_tensor.data(), - dsoftmax_offset_tensor.data(), q_cu_seqlens_tensor.data(), kv_cu_seqlens_tensor.data(), - q_seq_offsets_tensor.data(), k_seq_offsets_tensor.data(), q_max_seqlen, kv_max_seqlen, - scaling_factor, dropout_probability, qkv_layout, nvte_get_q_format(qkv_layout), - nvte_get_q_format(qkv_layout), qkv_layout, NVTE_QKV_Format_NOT_SET, NVTE_QKV_Format_NOT_SET, - bias_type, mask_type, softmax_type, window_size_left, window_size_right, - bottom_right_diagonal, deterministic, false, workspace_tensor.data(), 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); } diff --git a/transformer_engine/jax/csrc/extensions/pybind.cpp b/transformer_engine/jax/csrc/extensions/pybind.cpp index 3927e2686e..bc31a74d4f 100644 --- a/transformer_engine/jax/csrc/extensions/pybind.cpp +++ b/transformer_engine/jax/csrc/extensions/pybind.cpp @@ -134,7 +134,8 @@ pybind11::dict Registrations() { PYBIND11_MODULE(transformer_engine_jax, m) { m.def("registrations", &Registrations); - m.def("get_fused_attn_backend", &GetFusedAttnBackend); + 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); @@ -200,12 +201,14 @@ PYBIND11_MODULE(transformer_engine_jax, m) { .value("NVTE_BSHD_BSHD_BSHD", NVTE_QKV_Layout::NVTE_BSHD_BSHD_BSHD) .value("NVTE_T3HD", NVTE_QKV_Layout::NVTE_T3HD) .value("NVTE_THD_T2HD", NVTE_QKV_Layout::NVTE_THD_T2HD) - .value("NVTE_THD_THD_THD", NVTE_QKV_Layout::NVTE_THD_THD_THD); + .value("NVTE_THD_THD_THD", NVTE_QKV_Layout::NVTE_THD_THD_THD) + .value("NVTE_QKV_Layout_NOT_SET", NVTE_QKV_Layout::NVTE_QKV_Layout_NOT_SET); pybind11::enum_(m, "NVTE_QKV_Format", pybind11::module_local()) .value("NVTE_SBHD", NVTE_QKV_Format::NVTE_SBHD) .value("NVTE_BSHD", NVTE_QKV_Format::NVTE_BSHD) - .value("NVTE_THD", NVTE_QKV_Format::NVTE_THD); + .value("NVTE_THD", NVTE_QKV_Format::NVTE_THD) + .value("NVTE_QKV_Format_NOT_SET", NVTE_QKV_Format::NVTE_QKV_Format_NOT_SET); pybind11::enum_(m, "NVTE_Softmax_Type", pybind11::module_local()) .value("NVTE_VANILLA_SOFTMAX", NVTE_Softmax_Type::NVTE_VANILLA_SOFTMAX) diff --git a/transformer_engine/jax/flax/transformer.py b/transformer_engine/jax/flax/transformer.py index 4b497826cc..a048a9c443 100644 --- a/transformer_engine/jax/flax/transformer.py +++ b/transformer_engine/jax/flax/transformer.py @@ -5,6 +5,7 @@ Wrapper module for Transformer related layers with FP8 support. """ import functools +import operator from enum import Enum from math import sqrt import os @@ -20,6 +21,7 @@ from jax import random as jax_random from jax import lax, vmap from jax.ad_checkpoint import checkpoint_name +from transformer_engine_jax import NVTE_Fused_Attn_Backend from .module import DenseGeneral, LayerNormDenseGeneral, LayerNormMLP from .module import LayerNorm, Softmax @@ -30,9 +32,10 @@ QKVLayout, SequenceDescriptor, ) -from ..attention import is_fused_attn_kernel_available, make_swa_mask, canonicalize_attn_mask_type +from ..attention import make_swa_mask, canonicalize_attn_mask_type from ..attention import fused_attn from ..attention import CPStrategy +from ..cpp_extensions import FusedAttnHelper from ..softmax import SoftmaxFusionType from ..sharding import num_of_devices from ..sharding import get_sharding_map_logic_axis_to_mesh_axis @@ -792,6 +795,8 @@ def __call__( raise ValueError("return_max_logit requires fused attention, but NVTE_FUSED_ATTN=0.") sequence_dim = 0 if self.transpose_batch_sequence else 1 + batch_dim = 1 - sequence_dim + batch_size = query.shape[batch_dim] seqlen_q = query.shape[sequence_dim] if qkv_layout == QKVLayout.BS3HD: seqlen_kv = seqlen_q @@ -808,10 +813,18 @@ def __call__( if not enable_fused_attn: raise ValueError("score_mod requires fused attention, but NVTE_FUSED_ATTN=0.") kernel_qkv_layout = qkv_layout.to_separate() if score_mod_requested else qkv_layout - has_fused_attn_kernel = is_fused_attn_kernel_available( + bias_batch = bias_heads = bias_seqlen_q = bias_seqlen_kv = None + if attn_bias_type == AttnBiasType.POST_SCALE_BIAS: + *bias_batch_shape, bias_heads, bias_seqlen_q, bias_seqlen_kv = bias.shape + bias_batch = functools.reduce(operator.mul, bias_batch_shape) + num_segments = batch_size + if kernel_qkv_layout.is_thd(): + num_segments = batch_size * self.max_segments_per_seq + fused_attn_helper = FusedAttnHelper( # This needs to be fixed: TE-Jax has historically correlated training mode # with deterministic mode. not deterministic, + num_segments, input_dtype, # self._assert_dtypes enforces Q, K, V, bias to have the same dtype, so # using input_dtype as kv dtype is sufficient. @@ -827,9 +840,16 @@ def __call__( seqlen_kv, head_dim_qk, head_dim_v, - self.window_size, + (-1, -1) if self.window_size is None else self.window_size, return_max_logit=self.return_max_logit, + bottom_right_diagonal=attn_mask_type.is_bottom_right(), + bias_batch=bias_batch, + bias_heads=bias_heads, + bias_seqlen_q=bias_seqlen_q, + bias_seqlen_kv=bias_seqlen_kv, ) + fused_attn_backend, _ = fused_attn_helper.get_fused_attn_backend() + has_fused_attn_kernel = fused_attn_backend != NVTE_Fused_Attn_Backend.NVTE_No_Backend if score_mod_requested and not has_fused_attn_kernel: raise ValueError( "score_mod requires fused attention, but no fused attention kernel is available." @@ -844,12 +864,9 @@ def __call__( if enable_fused_attn and not has_fused_attn_kernel: warnings.warn( - "Fused attention is not enabled because there is no available kernel.\n" - "Fall back to the unfused attention.\n" - "Please try to update the cuDNN and TE to the latest version.\n" - f"{qkv_layout=}\n{attn_bias_type=}\n{attn_mask_type=}\n" - f"{self.attention_dropout=}\n{self.num_attention_heads=}\n{self.window_size=}\n" - f"{self.num_gqa_groups=}\n{seqlen_q=}\n{seqlen_kv=}\n{head_dim_qk=}\n{head_dim_v=}\n" + "Falling back to the unfused attention backend as fused attention does not support" + " this config. Set NVTE_DEBUG=1 and NVTE_DEBUG_LEVEL=2 to see the detailed" + " rejection reason.\n" ) dropout_rng = None diff --git a/transformer_engine/pytorch/attention/dot_product_attention/backends.py b/transformer_engine/pytorch/attention/dot_product_attention/backends.py index bca1f3200d..29181a882d 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/backends.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/backends.py @@ -2045,12 +2045,29 @@ def backward(ctx, d_out, *_args): class FusedAttention(torch.nn.Module): - """Dot product attention using cuDNN attention: + """Dot product attention using `cuDNN attention `_: FusedAttnBackend["F16_arbitrary_seqlen"] cuDNN attention for FP16/BF16 with any sequence length. FusedAttnBackend["FP8"] - cuDNN attention for FP8 with any sequence length. + cuDNN attention for FP8 with any sequence length. It supports the following recipes, where + "Inputs", "Intermediates" and "Outputs" are in the format of "tensor: quantizer". The recipes + are implemented in transformer_engine.pytorch.cpp_extension.fused_attn.fused_attn_fwd and + transformer_engine.pytorch.cpp_extension.fused_attn.fused_attn_bwd. + + Direction Inputs Intermediates Outputs + DelayedScaling (DS) forward Q/K/V: DS S: DS O: DS + backward Q/K/V/O (from forward), dO: DS dP: DS dQ/dK/dV: DS + Float8CurrentScaling (CS) forward Q/K/V: CS S: DS O: F16 + backward Q/K/V (from forward), dO: CS, + O: F16 (or CS if NVTE_DPA_FP8CS_O_in_F16=0) dP: DS dQ/dK/dV: F16 + MXFP8BlockScaling (MXFP8) forward Q/K row, V col: MXFP8 S: None O: F16 + backward Q/K row+col, V row: MXFP8, + O/dO: F16, dO row+col: MXFP8 dP: None dQ/dK/dV: F16 + + For MXFP8, "row" and "col" are the quantization directions, which align with the contraction axes + of the matmuls that consume the tensor. For more details, please refer to + `How Scales Are Applied in MXFP8 Attention `_. """ def __init__( diff --git a/transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py b/transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py index e9948cb9d8..1bc8478538 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py @@ -5411,6 +5411,103 @@ def backward(ctx, dout, *_args): ) +def cp_per_step_configs( + cp_comm_type, + cp_size, + cp_size_a2a, + *, + max_seqlen_q, + max_seqlen_kv, + num_tokens_q, + num_tokens_kv, + num_heads, + num_gqa_groups, + attn_mask_type, + window_size, + bottom_right_diagonal, +): + """Per-step attention configs a context-parallel run dispatches to its attention backend. + + CP runs attention in multiple steps, each with a distinct config (e.g. mask, and seqlens) + that differs from the single global config. This function returns the list of those distinct + per-step configs so `get_attention_backend` can check if the backend supports all of them. + """ + is_causal = "causal" in attn_mask_type + padding_or_no_mask = "padding" if "padding" in attn_mask_type else "no_mask" + window_left, window_right = window_size + + def config(mask, s_q, s_kv, heads, gqa, bottom_right, t_q, t_kv, window=None): + w_left, w_right = window if window is not None else (window_left, window_right) + return { + "attn_mask_type": mask, + "max_seqlen_q": s_q, + "max_seqlen_kv": s_kv, + "num_tokens_q": t_q, + "num_tokens_kv": t_kv, + "num_attn_heads": heads, + "num_gqa_groups": gqa, + "window_size_left": w_left, + "window_size_right": w_right, + "bottom_right_diagonal": bottom_right, + } + + if cp_comm_type == "a2a": + # split heads across the cp ranks + return [ + config( + attn_mask_type, + max_seqlen_q, + max_seqlen_kv, + num_heads // cp_size, + num_gqa_groups // cp_size, + bottom_right_diagonal, + num_tokens_q * cp_size, + num_tokens_kv * cp_size, + ) + ] + + if cp_comm_type == "all_gather": + # one short Q chunk vs a growing KV chunk; causal -> causal_bottom_right + s_q = max_seqlen_q // (2 * cp_size) + s_kv_chunk = max_seqlen_kv // (2 * cp_size) + mask, br = attn_mask_type, bottom_right_diagonal + if is_causal and "bottom_right" not in attn_mask_type: + mask, br = attn_mask_type + "_bottom_right", True + # Each step narrows max_seqlen_*, but the token counts it dispatches with are the + # rank's full Q tokens and the all-gathered KV tokens, unchanged across steps. + # Scaling them per step would key the probe's graph differently from the one the + # step looks up, and rebuild every graph this probes at execution time. + t_q = num_tokens_q + t_kv = num_tokens_kv * cp_size + # s_kv ranges from s_kv_chunk, i*s_kv_chunk, ..., max_seqlen_kv + # check a single chunk and the full KV + return [ + config(mask, s_q, s_kv, num_heads, num_gqa_groups, br, t_q, t_kv) + for s_kv in dict.fromkeys([s_kv_chunk, max_seqlen_kv]) + ] + + # p2p and a2a+p2p: split heads across the a2a subgroup, and ring over the p2p subgroup + p2p_size = cp_size // cp_size_a2a + heads = num_heads // cp_size_a2a + gqa = num_gqa_groups // cp_size_a2a + r_q = max_seqlen_q // p2p_size + r_kv = max_seqlen_kv // p2p_size + # The tensors handed to this rank already correspond to (r_q, r_kv), so the token counts + # need no rescaling here; they only follow the halving below. + t_q, t_kv = num_tokens_q, num_tokens_kv + if not is_causal: + return [config(attn_mask_type, r_q, r_kv, heads, gqa, bottom_right_diagonal, t_q, t_kv)] + return [ + config(attn_mask_type, r_q, r_kv, heads, gqa, bottom_right_diagonal, t_q, t_kv), # diagonal + config( + padding_or_no_mask, r_q, r_kv // 2, heads, gqa, bottom_right_diagonal, t_q, t_kv // 2 + ), # lower-triangle + config( + padding_or_no_mask, r_q // 2, r_kv, heads, gqa, bottom_right_diagonal, t_q // 2, t_kv + ), # upper-triangle + ] + + def attn_forward_func_with_cp( is_training, q, diff --git a/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py b/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py index dabfd10de4..82767beb49 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py @@ -923,6 +923,7 @@ def __init__( softmax_scale = 1.0 / math.sqrt( kv_channels if isinstance(kv_channels, int) else kv_channels[0] ) + self.softmax_scale = softmax_scale self.deterministic = ( not bool(int(os.getenv("NVTE_ALLOW_NONDETERMINISTIC_ALGO", "1"))) @@ -1964,14 +1965,13 @@ def forward( .. note:: Users can use environment variables :attr:`NVTE_FLASH_ATTN`, :attr:`NVTE_FUSED_ATTN`, - and :attr:`NVTE_FUSED_ATTN_BACKEND` to control which DotProductAttention backend, - and FusedAttention backend if applicable, to use. Transformer Engine first filters - backends by support for the runtime environment and input configuration, then applies - a performance-based preference order. On supported pre-Hopper GPUs, FlashAttention is - preferred over FusedAttention and UnfusedDotProductAttention when both optimized - backends are eligible. On Hopper and newer GPUs, including Blackwell, FusedAttention is - preferred over FlashAttention and UnfusedDotProductAttention when both optimized - backends are eligible. + and :attr:`NVTE_UNFUSED_ATTN` to control which DotProductAttention backend to use. + Transformer Engine first filters backends by support for the runtime environment + and input configuration, then applies a performance-based preference order. + On supported pre-Hopper GPUs, FlashAttention is preferred over FusedAttention and + UnfusedDotProductAttention when both optimized backends are eligible. On Hopper and + newer GPUs, including Blackwell, FusedAttention is preferred over FlashAttention and + UnfusedDotProductAttention when both optimized backends are eligible. If FusedAttention is being used, users can also choose to switch to flash-attn's implementation for backward by setting :attr:`NVTE_FUSED_ATTN_USE_FAv2_BWD=1` (default: 0), because of the performance differences between various versions of @@ -2589,11 +2589,14 @@ def forward( # adjust max_seqlen and cu_seqlens for CP cp_size = 1 + cp_size_a2a = 1 if isinstance(self.cp_group, dist_group_type): cp_size = get_distributed_world_size(self.cp_group) elif isinstance(self.cp_group, list): for group in self.cp_group: cp_size *= get_distributed_world_size(group) + if self.cp_comm_type == "a2a+p2p" and len(self.cp_group) > 0: + cp_size_a2a = get_distributed_world_size(self.cp_group[0]) context_parallel = cp_size > 1 if thd_mask_policies is not None and context_parallel: raise ValueError("Mixed THD policies do not support context parallelism.") @@ -2655,32 +2658,11 @@ def forward( _alibi_cache["_alibi_slopes_require_update"] = True _alibi_cache["_alibi_bias_require_update"] = True - # detect bias shape - core_attention_bias_shape = None - if core_attention_bias is not None: - if ( - core_attention_bias.shape[0] == batch_size - and core_attention_bias.shape[1] == query_layer.shape[-2] - ): - core_attention_bias_shape = "bhss" - elif ( - core_attention_bias.shape[0] == 1 - and core_attention_bias.shape[1] == query_layer.shape[-2] - ): - core_attention_bias_shape = "1hss" - elif ( - core_attention_bias.shape[0] == batch_size and core_attention_bias.shape[1] == 1 - ): - core_attention_bias_shape = "b1ss" - elif core_attention_bias.shape[0] == 1 and core_attention_bias.shape[1] == 1: - if core_attention_bias.shape[2] == 1: - core_attention_bias_shape = "111s" - else: - core_attention_bias_shape = "11ss" - else: - assert ( - False - ), "core_attention_bias must be in one of {bhss, 1hss, b1ss, 11ss, 111s} shapes" + core_attention_bias_shape = ( + tuple(core_attention_bias.shape) + if core_attention_bias_type != "no_bias" and core_attention_bias is not None + else None + ) # Default pad_between_seqs auto-detect. For THD, infer presence of # inter-sequence padding from whether padded cu_seqlens were supplied -- @@ -2746,18 +2728,27 @@ def forward( num_gqa_groups=num_gqa_groups, max_seqlen_q=max_seqlen_q, max_seqlen_kv=max_seqlen_kv, + num_tokens_q=(query_layer.shape[0] if q_format == "thd" else 0), + num_tokens_kv=(key_layer.shape[0] if kv_format == "thd" else 0), head_dim_qk=head_dim_qk, head_dim_v=head_dim_v, - alibi_slopes_shape=alibi_slopes.shape if alibi_slopes is not None else None, + alibi_slopes_shape=( + alibi_slopes.shape + if core_attention_bias_type == "alibi" and alibi_slopes is not None + else None + ), core_attention_bias_type=core_attention_bias_type, core_attention_bias_shape=core_attention_bias_shape, core_attention_bias_requires_grad=( - core_attention_bias.requires_grad if core_attention_bias is not None else False + core_attention_bias.requires_grad + if core_attention_bias_type != "no_bias" and core_attention_bias is not None + else False ), attention_dropout=self.attention_dropout, context_parallel=context_parallel, cp_comm_type=self.cp_comm_type, cp_size=cp_size, + cp_size_a2a=cp_size_a2a, deterministic=self.deterministic, is_training=self.training, fp8=self.fp8, @@ -2767,6 +2758,7 @@ def forward( return_max_logit=self.return_max_logit, cuda_graph=is_graph_capturing(), num_splits=num_splits, + softmax_scale=self.softmax_scale, fp8_output=fp8_output, checkpoint_core_attention=checkpoint_core_attention, has_score_mod=score_mod is not None, diff --git a/transformer_engine/pytorch/attention/dot_product_attention/utils.py b/transformer_engine/pytorch/attention/dot_product_attention/utils.py index 33d612b4f4..ca0675ab42 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/utils.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/utils.py @@ -24,6 +24,7 @@ import transformer_engine as te from transformer_engine.pytorch.cpp_extensions.fused_attn import ( QKVLayout, + QKVFormat, AttnBiasType, AttnMaskType, SoftmaxType, @@ -48,7 +49,7 @@ from transformer_engine.pytorch.tensor.storage.mxfp8_tensor_storage import MXFP8TensorStorage from transformer_engine.pytorch.quantization import get_fp8_te_dtype -from transformer_engine.pytorch.constants import TE_DType, MXFP8_BLOCK_SCALING_SIZE +from transformer_engine.pytorch.constants import TE_DType, DType, MXFP8_BLOCK_SCALING_SIZE from transformer_engine.pytorch.utils import ( @@ -205,6 +206,9 @@ class AttentionParams: Type of query/key/value tensors, {`torch.Tensor`, `Float8Tensor`}. qkv_dtype : torch.dtype, default = torch.bfloat16 Data type of query/key/value tensors. + nominal_dtype : Optional[torch.dtype], default = None + Model precision (F16/BF16) of the unquantized tensors (O, and dQ/dK/dV under + current/mxfp8) when `qkv_dtype` itself is FP8. qkv_layout : str, default = "sbh3d" Query/key/value tensor memory layout. batch_size : int, default = 1 @@ -217,6 +221,10 @@ class AttentionParams: Maximum sequence length of the query tensor. max_seqlen_kv : int, default = 128 Maximum sequence length of the key and value tensors. + num_tokens_q : int, default = 0 + Total number of query tokens in a batch, when `qkv_format=thd`. + num_tokens_kv : int, default = 0 + Total number of key/value tokens in a batch, when `qkv_format=thd`. head_dim_qk : int, default = 64 The size of each attention head in query and key tensors. head_dim_v : int, default = 64 @@ -226,15 +234,15 @@ class AttentionParams: `causal_bottom_right`, `padding_causal_bottom_right`, `arbitrary`} window_size : Tuple[int, int], default = None Sliding window attention size. - bottom_right_diagonal: bool, default = `None` + bottom_right_diagonal: bool, default = `True` Whether to align sliding window and ALiBi diagonal to the bottom right corner of the softmax matrix. alibi_slopes_shape : Optional[Union[torch.Size, List]], default = None Tensor shape of :attr:`alibi_slopes` in `DotProductAttention`. core_attention_bias_type : str, default = no_bias Attention bias type, {`no_bias`, `pre_scale_bias`, `post_scale_bias`, `alibi`}. - core_attention_bias_shape : str, default = 1hss - Attention bias shape, {`1hss`, `b1ss`, `bhss`}. + core_attention_bias_shape : Optional[Tuple[int, int, int, int]], default = None + Attention bias shape, (b, h, sq, skv). core_attention_bias_requires_grad : bool, default = True Whether attention bias requires gradient. pad_between_seqs : bool, default = False @@ -247,7 +255,9 @@ class AttentionParams: cp_comm_type : str, default = "p2p" The communication type of context parallelism. cp_size : int, default = 1 - The group size of context parallelism. + The (total) group size of context parallelism. + cp_size_a2a : int, default = 1 + The all-to-all subgroup size when `cp_comm_type == "a2a+p2p"`. deterministic : bool, default = False Whether to run `DotProductAttention` with determinism or not. is_training : bool, default = True @@ -266,6 +276,8 @@ class AttentionParams: Whether support for cuda graph capture is needed or not. num_splits : int, default = 1 The number of kernels to split attention to. + softmax_scale : float, default = 1.0 + Pre-softmax attention scale. fp8_output : bool, default = False Whether output is requested in FP8. checkpoint_core_attention : bool, default = False @@ -278,12 +290,15 @@ class AttentionParams: qkv_type: Union[torch.Tensor, Float8Tensor] = torch.Tensor qkv_dtype: torch.dtype = torch.bfloat16 + nominal_dtype: Optional[torch.dtype] = None qkv_layout: str = "sbh3d" batch_size: int = 1 num_heads: int = 16 num_gqa_groups: int = 16 max_seqlen_q: int = 128 max_seqlen_kv: int = 128 + num_tokens_q: int = 0 + num_tokens_kv: int = 0 head_dim_qk: int = 64 head_dim_v: int = 64 attn_mask_type: str = "no_mask" @@ -291,13 +306,14 @@ class AttentionParams: bottom_right_diagonal: bool = True alibi_slopes_shape: Union[torch.Size, List, None] = None core_attention_bias_type: str = "no_bias" - core_attention_bias_shape: str = "1hss" + core_attention_bias_shape: Union[Tuple[int, int, int, int], None] = None core_attention_bias_requires_grad: bool = True pad_between_seqs: bool = False attention_dropout: float = 0.0 context_parallel: bool = False cp_comm_type: str = "p2p" cp_size: int = 1 + cp_size_a2a: int = 1 deterministic: bool = False is_training: bool = True fp8: bool = False @@ -307,6 +323,7 @@ class AttentionParams: return_max_logit: bool = False cuda_graph: bool = False num_splits: int = 1 + softmax_scale: float = 1.0 fp8_output: bool = False checkpoint_core_attention: bool = False has_score_mod: bool = False @@ -331,6 +348,67 @@ def __eq__(self, other): return True +@dataclass(eq=True) +class FusedAttentionParams: + """ + Attention parameters used by the `FusedAttention` backend. + """ + + # basic attention settings + is_training: bool = True + deterministic: bool = False + cuda_graph: bool = False + return_max_logit: bool = False + attn_mask_type: tex.NVTE_Mask_Type = tex.NVTE_Mask_Type.NVTE_NO_MASK + bias_type: tex.NVTE_Bias_Type = tex.NVTE_Bias_Type.NVTE_NO_BIAS + window_size_left: int = -1 + window_size_right: int = -1 + bottom_right_diagonal: bool = True + softmax_type: tex.NVTE_Softmax_Type = tex.NVTE_Softmax_Type.NVTE_VANILLA_SOFTMAX + scaling_mode: tex.NVTEScalingMode = tex.NVTEScalingMode.NVTE_DELAYED_TENSOR_SCALING + dropout: float = 0.0 + attn_scale: float = 1.0 + + # tensor types + qkv_dtype: DType = DType.kBFloat16 + o_dtype: DType = DType.kBFloat16 + do_dtype: DType = DType.kBFloat16 + dqkv_dtype: DType = DType.kBFloat16 + + # tensor layouts + qkv_layout: tex.NVTE_QKV_Layout = tex.NVTE_QKV_Layout.NVTE_QKV_Layout_NOT_SET + o_format: tex.NVTE_QKV_Format = tex.NVTE_QKV_Format.NVTE_QKV_Format_NOT_SET + do_format: tex.NVTE_QKV_Format = tex.NVTE_QKV_Format.NVTE_QKV_Format_NOT_SET + dqkv_layout: tex.NVTE_QKV_Layout = tex.NVTE_QKV_Layout.NVTE_QKV_Layout_NOT_SET + qkv_scale_inv_format: tex.NVTE_QKV_Format = tex.NVTE_QKV_Format.NVTE_QKV_Format_NOT_SET + do_scale_inv_format: tex.NVTE_QKV_Format = tex.NVTE_QKV_Format.NVTE_QKV_Format_NOT_SET + + # tensor dimensions + batch_size: int = 0 + num_attn_heads: int = 0 + num_gqa_groups: int = 0 + head_dim_qk: int = 0 + head_dim_v: int = 0 + max_seqlen_q: int = 0 + max_seqlen_kv: int = 0 + num_tokens_q: int = 0 + num_tokens_kv: int = 0 + + # paged KV dimensions + num_pages_k: int = 0 + num_pages_v: int = 0 + page_size_k: int = 0 + page_size_v: int = 0 + max_pages_per_seq_k: int = 0 + max_pages_per_seq_v: int = 0 + + # bias dimensions + bias_batch_size: int = 0 + bias_num_heads: int = 0 + bias_seqlen_q: int = 0 + bias_seqlen_kv: int = 0 + + class _NoOpLogger: """ Stand-in for the "DotProductAttention" logger used when get_attention_backend @@ -355,20 +433,9 @@ def error(self, *args, **kwargs): @torch.compiler.assume_constant_result -def _get_fused_attn_backend( - is_training, - q_type, - kv_type, - qkv_layout, - bias_type, - attn_mask_type, - softmax_type, - *args, -): +def _get_fused_attn_backend(**fused_attn_kwargs): """Constant-foldable tex.get_fused_attn_backend: the result depends only on - the attention config. Layout/bias/mask/softmax are taken as their string - keys and resolved to the pybind enums here, so that every argument is a - python literal or a python enum. + the attention config. Returns a plain int rather than a FusedAttnBackend member: dynamo reconstructs the result of an assume_constant_result call by re-emitting the @@ -376,18 +443,10 @@ def _get_fused_attn_backend( 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).""" - return int( - tex.get_fused_attn_backend( - is_training, - q_type, - kv_type, - QKVLayout[qkv_layout], - AttnBiasType[bias_type], - AttnMaskType[attn_mask_type], - SoftmaxType[softmax_type], - *args, - ) + fused_attention_backend, reject_message = tex.get_fused_attn_backend( + FusedAttentionParams(**fused_attn_kwargs) ) + return int(fused_attention_backend), reject_message def get_attention_backend( @@ -422,12 +481,15 @@ def get_attention_backend( # is shifted over to the caller of this function qkv_type = attention_params.qkv_type qkv_dtype = attention_params.qkv_dtype + nominal_dtype = attention_params.nominal_dtype qkv_layout = attention_params.qkv_layout batch_size = attention_params.batch_size num_heads = attention_params.num_heads num_gqa_groups = attention_params.num_gqa_groups max_seqlen_q = attention_params.max_seqlen_q max_seqlen_kv = attention_params.max_seqlen_kv + num_tokens_q = attention_params.num_tokens_q + num_tokens_kv = attention_params.num_tokens_kv head_dim_qk = attention_params.head_dim_qk head_dim_v = attention_params.head_dim_v attn_mask_type = attention_params.attn_mask_type @@ -441,7 +503,8 @@ def get_attention_backend( attention_dropout = attention_params.attention_dropout context_parallel = attention_params.context_parallel cp_comm_type = attention_params.cp_comm_type - cp_size = attention_params.cp_size # pylint: disable=unused-variable + cp_size = attention_params.cp_size + cp_size_a2a = attention_params.cp_size_a2a deterministic = attention_params.deterministic is_training = attention_params.is_training fp8 = attention_params.fp8 @@ -451,6 +514,7 @@ def get_attention_backend( return_max_logit = attention_params.return_max_logit cuda_graph = attention_params.cuda_graph num_splits = attention_params.num_splits + softmax_scale = attention_params.softmax_scale fp8_output = attention_params.fp8_output checkpoint_core_attention = attention_params.checkpoint_core_attention has_score_mod = attention_params.has_score_mod @@ -1069,7 +1133,7 @@ def _is_fa3_supported(num_heads, num_gqa_groups, head_dim_qk, head_dim_v, qkv_dt use_flash_attention_4 = False # Filter: QKV layout - if qkv_format == "thd": + if "thd" in (q_format, kv_format): if pad_between_seqs: if ( # pylint: disable=too-many-boolean-expressions use_flash_attention_2 and FlashAttentionUtils.is_installed @@ -1087,7 +1151,7 @@ def _is_fa3_supported(num_heads, num_gqa_groups, head_dim_qk, head_dim_v, qkv_dt if cudnn_version < (9, 18, 1): if use_fused_attention: logger.debug( - "Disabling FusedAttention as qkv_format = thd is" + "Disabling FusedAttention when Q or KV uses THD because it is" " not supported for compute capability = sm120 and cuDNN version < 9.18.1" ) use_fused_attention = False @@ -1437,65 +1501,186 @@ def _is_fa3_supported(num_heads, num_gqa_groups, head_dim_qk, head_dim_v, qkv_dt fu_core_attention_bias_requires_grad = False if len(alibi_slopes_shape) == 1 and alibi_slopes_shape[0] == num_heads: - fu_core_attention_bias_shape = "1hss" + fu_core_attention_bias_shape = (1, num_heads, max_seqlen_q, max_seqlen_kv) elif ( len(alibi_slopes_shape) == 2 and alibi_slopes_shape[0] == batch_size and alibi_slopes_shape[1] == num_heads ): - fu_core_attention_bias_shape = "bhss" + fu_core_attention_bias_shape = (batch_size, num_heads, max_seqlen_q, max_seqlen_kv) + fu_core_attention_bias_shape_type = None + if ( + fu_core_attention_bias_type == "post_scale_bias" + and fu_core_attention_bias_shape is not None + ): + b, h, sq, _skv = fu_core_attention_bias_shape + if b == batch_size and h == num_heads: + fu_core_attention_bias_shape_type = "bhss" + elif b == 1 and h == num_heads: + fu_core_attention_bias_shape_type = "1hss" + elif b == batch_size and h == 1: + fu_core_attention_bias_shape_type = "b1ss" + elif b == 1 and h == 1: + fu_core_attention_bias_shape_type = "111s" if sq == 1 and max_seqlen_q != 1 else "11ss" + else: + raise ValueError( + "core_attention_bias tensor must be in one of " + '{"bhss", "1hss", "b1ss", "11ss", "111s"} shapes. ' + f"Found (b,h,sq,skv) = ({b},{h},{sq},{_skv})" + ) if ( use_fused_attention and fu_core_attention_bias_type == "post_scale_bias" - and fu_core_attention_bias_shape != "1hss" + and fu_core_attention_bias_shape_type != "1hss" ): # dbias calculation is not supported for 111s as of cuDNN 9.18. So, use fused attention backend only if bias does not require grad. - if fu_core_attention_bias_requires_grad and fu_core_attention_bias_shape == "111s": + if fu_core_attention_bias_requires_grad and fu_core_attention_bias_shape_type == "111s": logger.warning( "Disabling FusedAttention as dbias calculation is not supported for 111s" ) use_fused_attention = False + # Filter: cuDNN support fused_attention_backend = None if use_fused_attention: - # ``DType`` is implicitly convertible to ``transformer_engine::DType`` - # on the C++ side, so pass it straight to the pybind function. - q_type = TE_DType[qkv_dtype] - kv_type = q_type - if fp8 and fp8_meta["recipe"].fp8_dpa: - q_type = get_fp8_te_dtype(fp8_meta["recipe"], fprop_tensor=True) - kv_type = q_type - # NOTE: under torch.compile the numeric args below must not be symbolic - # (assume_constant_result requires concrete values); ints/floats made - # dynamic by automatic dynamic currently graph break here. - fused_attention_backend = _get_fused_attn_backend( - is_training, - q_type, - kv_type, - qkv_layout, - fu_core_attention_bias_type, - attn_mask_type, - softmax_type, - attention_dropout, - num_heads, - num_gqa_groups, - max_seqlen_q, - max_seqlen_kv, - head_dim_qk, - head_dim_v, - window_size[0], - window_size[1], - return_max_logit, - cuda_graph, - deterministic, + recipe = fp8_meta["recipe"] if (fp8 and fp8_meta["recipe"].fp8_dpa) else None + cs_o_in_f16 = os.getenv("NVTE_DPA_FP8CS_O_in_F16", "1") == "1" + spec = get_fused_attn_spec( + recipe, qkv_dtype, qkv_layout, cs_o_in_f16=cs_o_in_f16, nominal_dtype=nominal_dtype ) - if fused_attention_backend == FusedAttnBackend.No_Backend.value: - logger.debug("Disabling FusedAttention as no backend supports the provided input") - use_fused_attention = False - fused_attention_backend = None - elif ( - has_score_mod and fused_attention_backend != FusedAttnBackend.F16_arbitrary_seqlen.value + qkv_type, o_type, do_type, dqkv_type = spec.qkv, spec.o, spec.do, spec.dqkv + scaling_mode = spec.scaling_mode + qkv_scale_inv_format = spec.scale_inv_format + do_scale_inv_format = spec.scale_inv_format + o_format = spec.o_format + do_format = spec.do_format + dqkv_layout = spec.dqkv_layout + num_pages_k = num_pages_v = 0 + page_size_k = page_size_v = 0 + max_pages_per_seq_k = max_pages_per_seq_v = 0 + if inference_params is not None and getattr(inference_params, "is_paged", False): + num_pages_k = num_pages_v = inference_params.total_num_pages + page_size_k = page_size_v = inference_params.page_size + max_pages_per_seq_k = max_pages_per_seq_v = ( + inference_params.cache_manager.max_pages_per_seq + ) + bias_batch_size = bias_num_heads = bias_seqlen_q = bias_seqlen_kv = 0 + if fu_core_attention_bias_shape is not None: + bias_batch_size, bias_num_heads, bias_seqlen_q, bias_seqlen_kv = ( + fu_core_attention_bias_shape + ) + base_fused_attn_kwargs = { + "is_training": is_training, + "deterministic": deterministic, + "cuda_graph": cuda_graph, + "return_max_logit": return_max_logit, + "attn_mask_type": AttnMaskType[attn_mask_type], + "bias_type": AttnBiasType[fu_core_attention_bias_type], + "window_size_left": window_size[0], + "window_size_right": window_size[1], + "bottom_right_diagonal": bottom_right_diagonal, + "softmax_type": SoftmaxType[softmax_type], + "scaling_mode": scaling_mode, + "dropout": attention_dropout, + "attn_scale": softmax_scale, + "qkv_dtype": qkv_type, + "o_dtype": o_type, + "do_dtype": do_type, + "dqkv_dtype": dqkv_type, + "qkv_layout": QKVLayout[spec.qkv_layout], + "o_format": QKVFormat[o_format], + "do_format": QKVFormat[do_format], + "dqkv_layout": QKVLayout[dqkv_layout], + "qkv_scale_inv_format": QKVFormat[qkv_scale_inv_format], + "do_scale_inv_format": QKVFormat[do_scale_inv_format], + "batch_size": batch_size, + "num_attn_heads": num_heads, + "num_gqa_groups": num_gqa_groups, + "head_dim_qk": head_dim_qk, + "head_dim_v": head_dim_v, + "max_seqlen_q": max_seqlen_q, + "max_seqlen_kv": max_seqlen_kv, + "num_tokens_q": num_tokens_q, + "num_tokens_kv": num_tokens_kv, + "num_pages_k": num_pages_k, + "num_pages_v": num_pages_v, + "page_size_k": page_size_k, + "page_size_v": page_size_v, + "max_pages_per_seq_k": max_pages_per_seq_k, + "max_pages_per_seq_v": max_pages_per_seq_v, + "bias_batch_size": bias_batch_size, + "bias_num_heads": bias_num_heads, + "bias_seqlen_q": bias_seqlen_q, + "bias_seqlen_kv": bias_seqlen_kv, + } + # Context-parallel per-step configs + if context_parallel: + from transformer_engine.pytorch.attention.dot_product_attention.context_parallel import ( + cp_per_step_configs, + ) + + per_step_configs = cp_per_step_configs( + cp_comm_type, + cp_size, + cp_size_a2a, + max_seqlen_q=max_seqlen_q, + max_seqlen_kv=max_seqlen_kv, + num_tokens_q=num_tokens_q, + num_tokens_kv=num_tokens_kv, + num_heads=num_heads, + num_gqa_groups=num_gqa_groups, + attn_mask_type=attn_mask_type, + window_size=window_size, + bottom_right_diagonal=bottom_right_diagonal, + ) + else: + per_step_configs = [None] + + for step_config in per_step_configs: + fused_attn_kwargs = dict(base_fused_attn_kwargs) + if step_config is not None: + step_seqlen_q = step_config["max_seqlen_q"] + step_seqlen_kv = step_config["max_seqlen_kv"] + fused_attn_kwargs.update( + attn_mask_type=AttnMaskType[step_config["attn_mask_type"]], + max_seqlen_q=step_seqlen_q, + max_seqlen_kv=step_seqlen_kv, + num_tokens_q=step_config["num_tokens_q"], + num_tokens_kv=step_config["num_tokens_kv"], + num_attn_heads=step_config["num_attn_heads"], + num_gqa_groups=step_config["num_gqa_groups"], + window_size_left=step_config["window_size_left"], + window_size_right=step_config["window_size_right"], + bottom_right_diagonal=step_config["bottom_right_diagonal"], + ) + if fu_core_attention_bias_shape is not None: + if bias_seqlen_q != 1: + fused_attn_kwargs["bias_seqlen_q"] = step_seqlen_q + if bias_seqlen_kv != 1: + fused_attn_kwargs["bias_seqlen_kv"] = step_seqlen_kv + # NOTE: under torch.compile the numeric entries of fused_attn_kwargs must not be + # symbolic (assume_constant_result requires concrete values); ints/floats made + # dynamic by automatic dynamic currently graph break here. + fused_attention_backend, reject_message = _get_fused_attn_backend(**fused_attn_kwargs) + if fused_attention_backend == FusedAttnBackend.No_Backend.value: + logger.debug( + "Disabling FusedAttention: %s%s", + reject_message, + ( + f" (context-parallel per-step config {step_config})" + if step_config is not None + else "" + ), + ) + use_fused_attention = False + fused_attention_backend = None + break + + if ( + use_fused_attention + and has_score_mod + and fused_attention_backend != FusedAttnBackend.F16_arbitrary_seqlen.value ): logger.debug( "Disabling FusedAttention for score_mod because sub-backend %s is not " @@ -2440,6 +2625,102 @@ def get_qkv_format( return qkv_format, q_format, kv_format +@dataclass(frozen=True) +class FusedAttnSpec: + """Fused-attention spec for a given config. + + Mirrors what `FusedAttnFunc` feeds `fused_attn_fwd`/`fused_attn_bwd` (backends.py), + so the availability probe (`get_attention_backend`) cannot drift from runtime. + """ + + scaling_mode: Any + qkv: Any + o: Any + do: Any + dqkv: Any + scale_inv_format: Optional[str] + qkv_layout: str + o_format: str + do_format: str + dqkv_layout: str + + +def get_fused_attn_spec(recipe, qkv_dtype, qkv_layout, *, cs_o_in_f16, nominal_dtype=None): + """Resolve fused-attention specs, e.g. tensor dtypes, formats, for a given config""" + q_format = get_qkv_format(qkv_layout)[1] + eff_qkv_layout = qkv_layout # FP16/BF16 + if recipe is not None: + if not recipe.mxfp8(): + # Delayed/current scaling + eff_qkv_layout = qkv_layout.replace("paged_kv_", "") + elif qkv_layout in ("bshd_bshd_bshd", "sbhd_sbhd_sbhd"): + eff_qkv_layout = qkv_layout # MXFP8 fast path + else: + eff_qkv_layout = "bhsd_bhsd_bhsd" # MXFP8 slow path + layout_kwargs = { + "qkv_layout": eff_qkv_layout, + "o_format": q_format, + "do_format": q_format, + "dqkv_layout": qkv_layout, + } + + if qkv_dtype in (torch.float8_e4m3fn, torch.float8_e5m2): + ref = TE_DType[nominal_dtype if nominal_dtype is not None else torch.bfloat16] + else: + ref = TE_DType[qkv_dtype] + + # FP16/BF16: every tensor is in model precision; scaling_mode is a placeholder + if recipe is None: + return FusedAttnSpec( + tex.NVTEScalingMode.NVTE_DELAYED_TENSOR_SCALING, + ref, + ref, + ref, + ref, + None, + **layout_kwargs, + ) + + fprop_fp8 = get_fp8_te_dtype(recipe, fprop_tensor=True) + grad_fp8 = get_fp8_te_dtype(recipe, fprop_tensor=False) + + # MXFP8 block scaling: Q/K/V/dO are in MXFP8; O/dQ/dK/dV stay in model precision + if recipe.mxfp8(): + return FusedAttnSpec( + tex.NVTEScalingMode.NVTE_MXFP8_1D_SCALING, + fprop_fp8, + ref, + grad_fp8, + ref, + "bhsd", + **layout_kwargs, + ) + + # FP8 current scaling: Q/K/V/dO are in FP8; O in model precision if `cs_o_in_f16` (default), otherwise FP8; + # dQ/dK/dV in model precision + if recipe.float8_current_scaling(): + return FusedAttnSpec( + tex.NVTEScalingMode.NVTE_DELAYED_TENSOR_SCALING, + fprop_fp8, + ref if cs_o_in_f16 else fprop_fp8, + grad_fp8, + ref, + None, + **layout_kwargs, + ) + + # FP8 delayed scaling: Q/K/V/O are in FP8 (e.g. E4M3); dO/dQ/dK/dV in FP8 (e.g. E5M2) + return FusedAttnSpec( + tex.NVTEScalingMode.NVTE_DELAYED_TENSOR_SCALING, + fprop_fp8, + fprop_fp8, + grad_fp8, + grad_fp8, + None, + **layout_kwargs, + ) + + def qkv_layout_needs_detection(*qkv: Optional[torch.Tensor]) -> bool: """Whether the layout of these q/k/v can only be told by inspecting memory. diff --git a/transformer_engine/pytorch/csrc/extensions.h b/transformer_engine/pytorch/csrc/extensions.h index 51f8b636b2..f1ef80ca2e 100644 --- a/transformer_engine/pytorch/csrc/extensions.h +++ b/transformer_engine/pytorch/csrc/extensions.h @@ -94,12 +94,8 @@ std::tuple moe_unpermute_bwd(at::Tensor input_bwd, at::T * Attention **************************************************************************************************/ -NVTE_Fused_Attn_Backend get_fused_attn_backend( - bool is_training, const DType q_dtype, const DType kv_dtype, NVTE_QKV_Layout qkv_layout, - NVTE_Bias_Type bias_type, NVTE_Mask_Type attn_mask_type, NVTE_Softmax_Type softmax_type, - float p_dropout, size_t num_attn_heads, size_t num_gqa_groups, size_t max_seqlen_q, - size_t max_seqlen_kv, size_t head_dim_qk, size_t head_dim_v, int64_t window_size_left, - int64_t window_size_right, bool return_max_logit, bool cuda_graph, bool deterministic); +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, diff --git a/transformer_engine/pytorch/csrc/extensions/attention.cpp b/transformer_engine/pytorch/csrc/extensions/attention.cpp index 82d18eb1d9..9cda33474f 100644 --- a/transformer_engine/pytorch/csrc/extensions/attention.cpp +++ b/transformer_engine/pytorch/csrc/extensions/attention.cpp @@ -11,18 +11,55 @@ namespace transformer_engine::pytorch { // get the fused attention backend -NVTE_Fused_Attn_Backend get_fused_attn_backend( - bool is_training, const DType q_dtype, const DType kv_dtype, NVTE_QKV_Layout qkv_layout, - NVTE_Bias_Type bias_type, NVTE_Mask_Type attn_mask_type, NVTE_Softmax_Type softmax_type, - float p_dropout, size_t num_attn_heads, size_t num_gqa_groups, size_t max_seqlen_q, - size_t max_seqlen_kv, size_t head_dim_qk, size_t head_dim_v, int64_t window_size_left, - int64_t window_size_right, bool return_max_logit, bool cuda_graph, bool deterministic) { - NVTE_Fused_Attn_Backend fused_attention_backend = nvte_get_fused_attn_backend( - is_training, static_cast(q_dtype), static_cast(kv_dtype), qkv_layout, - bias_type, attn_mask_type, softmax_type, p_dropout, num_attn_heads, num_gqa_groups, - max_seqlen_q, max_seqlen_kv, head_dim_qk, head_dim_v, window_size_left, window_size_right, - return_max_logit, cuda_graph, deterministic); - return 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 @@ -210,22 +247,50 @@ std::vector fused_attn_fwd( // 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( - te_Q.data(), te_K.data(), te_V.data(), te_Bias.data(), te_SoftmaxOffset.data(), te_S.data(), - te_O.data(), &nvte_aux_tensor_pack, te_cu_seqlens_q.data(), te_cu_seqlens_kv.data(), - te_cu_seqlens_q_padded.data(), te_cu_seqlens_kv_padded.data(), te_page_table_k.data(), - te_page_table_v.data(), te_rng_state.data(), max_seqlen_q, max_seqlen_kv, is_training, - return_max_logit, cuda_graph, attn_scale, p_dropout, qkv_layout, o_format, - qkv_scale_inv_format, bias_type, attn_mask_type, softmax_type, window_size[0], - window_size[1], bottom_right_diagonal, workspace.data(), at::cuda::getCurrentCUDAStream()); - }); + 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; @@ -268,16 +333,7 @@ std::vector fused_attn_fwd( } // execute the kernel - NVTE_SCOPED_GIL_RELEASE({ - nvte_fused_attn_fwd( - te_Q.data(), te_K.data(), te_V.data(), te_Bias.data(), te_SoftmaxOffset.data(), te_S.data(), - te_O.data(), &nvte_aux_tensor_pack, te_cu_seqlens_q.data(), te_cu_seqlens_kv.data(), - te_cu_seqlens_q_padded.data(), te_cu_seqlens_kv_padded.data(), te_page_table_k.data(), - te_page_table_v.data(), te_rng_state.data(), max_seqlen_q, max_seqlen_kv, is_training, - return_max_logit, cuda_graph, attn_scale, p_dropout, qkv_layout, o_format, - qkv_scale_inv_format, bias_type, attn_mask_type, softmax_type, window_size[0], - window_size[1], bottom_right_diagonal, workspace.data(), at::cuda::getCurrentCUDAStream()); - }); + 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); @@ -532,36 +588,57 @@ std::vector fused_attn_bwd( // 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( - te_Q.data(), te_K.data(), te_V.data(), te_O.data(), te_dO.data(), te_S.data(), te_dP.data(), - &nvte_aux_tensor_pack, te_dQ.data(), te_dK.data(), te_dV.data(), te_dBias.data(), - te_dSoftmaxOffset.data(), te_cu_seqlens_q.data(), te_cu_seqlens_kv.data(), - te_cu_seqlens_q_padded.data(), te_cu_seqlens_kv_padded.data(), max_seqlen_q, max_seqlen_kv, - attn_scale, p_dropout, qkv_layout, o_format, do_format, dqkv_layout, qkv_scale_inv_format, - do_scale_inv_format, bias_type, attn_mask_type, softmax_type, window_size[0], - window_size[1], bottom_right_diagonal, deterministic, cuda_graph, workspace.data(), - at::cuda::getCurrentCUDAStream()); - }); + 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( - te_Q.data(), te_K.data(), te_V.data(), te_O.data(), te_dO.data(), te_S.data(), te_dP.data(), - &nvte_aux_tensor_pack, te_dQ.data(), te_dK.data(), te_dV.data(), te_dBias.data(), - te_dSoftmaxOffset.data(), te_cu_seqlens_q.data(), te_cu_seqlens_kv.data(), - te_cu_seqlens_q_padded.data(), te_cu_seqlens_kv_padded.data(), max_seqlen_q, max_seqlen_kv, - attn_scale, p_dropout, qkv_layout, o_format, do_format, dqkv_layout, qkv_scale_inv_format, - do_scale_inv_format, bias_type, attn_mask_type, softmax_type, window_size[0], - window_size[1], bottom_right_diagonal, deterministic, cuda_graph, workspace.data(), - at::cuda::getCurrentCUDAStream()); - }); + NVTE_SCOPED_GIL_RELEASE({ nvte_fused_attn_bwd_v2(params); }); // destroy tensor wrappers nvte_tensor_pack_destroy(&nvte_aux_tensor_pack); diff --git a/transformer_engine/pytorch/csrc/extensions/pybind.cpp b/transformer_engine/pytorch/csrc/extensions/pybind.cpp index c24e95e951..eff5b9aafe 100644 --- a/transformer_engine/pytorch/csrc/extensions/pybind.cpp +++ b/transformer_engine/pytorch/csrc/extensions/pybind.cpp @@ -470,7 +470,7 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { "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::call_guard()); + "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());