diff --git a/docs/envvars.rst b/docs/envvars.rst index 97eaed5ddc..4490320030 100644 --- a/docs/envvars.rst +++ b/docs/envvars.rst @@ -190,18 +190,32 @@ 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: Enable diagnostic logging for the FusedAttention graph cache (covers both the F16 and FP8 kernels, forward and backward). Output goes to stderr, prefixed ``[FUSED-ATTN-CACHE]``. + + ``1`` emits one line per event that happens once per distinct cache key -- ``CREATE_GRAPH`` when a graph is constructed, ``CACHE_GRAPH`` when cuDNN has agreed to run it, ``BUILD_PLANS`` when its kernels are compiled on first execution -- plus an end-of-run summary block (one row per build site, per thread and in total, plus a row per pass across the backends if a run used both) and a breakdown of cuDNN graph-build timings. This is enough to diagnose redundant graph rebuilds and to profile build cost. Every event name is also the counter column it increments, so each line can be read against the running totals it carries. + + Every line names the thread and device it came from, then the build site behind it -- ``f16`` or ``fp8``, then the pass -- and carries the counters of that site alone, so a process that uses both backends can still tell which of them built what. One line is one pass; the forward and the backward read as adjacent rows. A summary row is an event line without the event name, the block is delimited by ``===== summary begin =====`` and ``===== summary end =====``, and ``tid=all dev=all`` marks the totals. A build site the run never reached is left out rather than shown as a row of zeros. + + ``hit`` and ``miss`` account for every lookup. A miss builds a graph, counted in ``create_graph``, and keeps it only if cuDNN agrees to run it, counted in ``cache_graph``; so the columns fall ``miss`` >= ``create_graph`` >= ``cache_graph``. ``create_graph`` minus ``cache_graph`` is graphs cuDNN refused to run. Nothing is cached for them, so this counts rejected builds rather than rejected configurations: a configuration that is queried again is built and rejected again. A site whose ``cache_graph`` stays put while ``miss`` climbs never runs fused and keeps paying to find that out, which makes these the columns to read when attention is slower than expected and nothing raised an error -- and at level 1, a ``CREATE_GRAPH`` line with no ``CACHE_GRAPH`` after it is one such rejection as it happens. The reason cuDNN gave is not logged here; it reaches the framework as the message explaining why the fused backend was not selected. ``miss`` and ``create_graph`` should agree: configurations FusedAttention itself does not serve are refused before any graph is built, and reported as that same message, so a gap between those two columns means a graph build failed where none was expected to. + + ``2`` additionally emits a per-lookup ``HIT``/``MISS`` line carrying the full cache key, and a per-execution ``EXECUTE`` line. Diffing two ``MISS`` lines names the fields that cost the extra build. These fire on every lookup and execution, so at test-suite scale they add I/O and serialize threads on the stderr lock; prefer ``1`` unless you need to see which shapes are missing. + + Each line is written after the cache lock is released rather than under it, so that no thread waits on stderr while holding the cache. With several threads active this means the lines can appear in a different order than the lookups they report; the counters each line carries still increase in event order, and lines from a single thread are still in that thread's order. + + By default only rank 0 emits, so that output does not scale with the world size. Append ``:`` to override -- ``1:all`` for every rank, ``2:0,3`` for a specific set. Worth overriding under context parallelism, where the ranks genuinely run different configurations. + + Has negligible overhead when unset. + .. 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..7f5deee722 100644 --- a/docs/examples/attention/attention.ipynb +++ b/docs/examples/attention/attention.ipynb @@ -253,7 +253,21 @@ "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", - "" + "\n", + "\n", + "Once cuDNN attention has been selected, a separate flag reports on the cuDNN graph cache underneath it. Transformer Engine builds a cuDNN graph per distinct attention configuration and reuses it, so a workload that unexpectedly rebuilds graphs pays for it on every step.\n", + "```\n", + "NVTE_FUSED_ATTN_CACHE_DEBUG = 0/1/2 # disables/enables graph cache diagnostics\n", + "```\n", + "At `1`, every line is prefixed `[FUSED-ATTN-CACHE]`, names the thread and device it came from and then the build site behind it -- `f16` or `fp8`, then the pass -- and there is one per event that happens once per configuration: `CREATE_GRAPH` when a graph is constructed, and `BUILD_PLANS` when its kernels are compiled on first execution. Each event name is also the counter column it increments. A configuration cuDNN declines raises a miss and no build, so `miss` minus `create_graph` counts refused builds; nothing is cached for a refusal, so a configuration that is queried again is built and refused again. An end-of-run summary block gives one row per build site, per thread and in total (`tid=all dev=all`), followed by where the build time went:\n", + "```\n", + "[FUSED-ATTN-CACHE] tid=0 dev=0 | f16 fwd CREATE_GRAPH | hit=0, miss=1, create_graph=1, ...\n", + "[FUSED-ATTN-CACHE] tid=all dev=all | f16 fwd | hit=5, miss=1, create_graph=1, ...\n", + "[FUSED-ATTN-CACHE] f16 fwd check_support | calls=1 | time= 42.135 ms/call\n", + "```\n", + "The number to read first is `create_graph`. It should settle at the number of distinct configurations the model uses and then stop growing; if it keeps climbing step after step, something in the configuration is varying that need not be. At `2`, each cache lookup adds a `HIT`/`MISS` line carrying the full key, and diffing two `MISS` lines names the fields that cost the extra build. Level `2` fires on every lookup, so use it to answer a specific question rather than leaving it on.\n", + "\n", + "This flag is supported in both PyTorch and JAX, since the cache it reports on lives in the common C++ layer. By default only rank 0 emits; see [NVTE_FUSED_ATTN_CACHE_DEBUG](../../envvars.rst) for selecting other ranks." ] }, { @@ -346,17 +360,11 @@ "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", + "Environment variables NVTE_FLASH_ATTN, NVTE_UNFUSED_ATTN, and NVTE_FUSED_ATTN_USE_FAv2_BWD are supported in PyTorch. NVTE_FUSED_ATTN and NVTE_ALLOW_NONDETERMINISTIC_ALGO are supported in both PyTorch and JAX.\n", "
\n", "\n", "### 2.3 Example Tests\n", 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 b6e11b8bea..205f9671ab 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, @@ -479,6 +441,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 352ab64a0d..ddfb2c5fd0 100644 --- a/tests/jax/test_fused_attn.py +++ b/tests/jax/test_fused_attn.py @@ -53,6 +53,9 @@ # Get determinism _deterministic = not bool(int(os.getenv("NVTE_ALLOW_NONDETERMINISTIC_ALGO", "1"))) +# CI test level +_TEST_LEVEL = os.getenv("NVTE_JAX_UNITTEST_LEVEL", "L0") + @pytest.fixture(autouse=True, scope="module") def init(): @@ -469,6 +472,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.") @@ -571,8 +595,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, @@ -587,9 +624,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, + 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 b1f165f491..f854f4b16a 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, @@ -454,11 +462,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 @@ -533,7 +537,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): @@ -557,7 +561,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/pytorch/attention/run_graph_cache.py b/tests/pytorch/attention/run_graph_cache.py new file mode 100644 index 0000000000..dcf1f1a0a5 --- /dev/null +++ b/tests/pytorch/attention/run_graph_cache.py @@ -0,0 +1,135 @@ +# 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 cuDNN graph cache and +marks each phase boundary on stderr, so that the parent can attribute the +[FUSED-ATTN-CACHE] lines NVTE_FUSED_ATTN_CACHE_DEBUG=2 emits to the phase that produced +them. + +This runs as its own process because the cache is process-wide and its counters only +accumulate: inside the pytest process, the graphs every earlier test built would be mixed +into the counts, and the cache would already be warm for whatever this test asked about. + +The phases, in order, and what each one is for: + query the first support query for a config -- the miss that builds its graphs + requery the identical query again -- must be answered from the cache + exec forward and backward of that config -- must reuse the graphs the query built, + and is where the plan build the query deferred happens + rescale the same execution with only softmax_scale changed -- must still reuse them, + since attn_scale is normalized out of the cache key + reshape a query differing in max_seqlen -- must build again, once per pass + +Prints ``[CACHE-TEST] fused=1`` (or 0) on stdout so the parent can skip rather than fail on +a GPU or cuDNN version with no fused-attention backend for the config. +""" + +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" + +# Derived exactly as DotProductAttention derives it, because the query below has to ask about +# the configuration the execution phases will run: deterministic is part of the backward cache +# key, so a query that assumed the default would build a graph the execution then misses. +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. + + Written to stderr, which is where the diagnostics go, so that the marker keeps its + place in the stream instead of racing them on a second file descriptor. + """ + sys.stderr.write(f"[CACHE-TEST] phase={name}\n") + sys.stderr.flush() + + +def query(config: ModelConfig) -> bool: + """Run one backend support query, as the test suite does, and report whether cuDNN + took the configuration. This is the call that 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 fused 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)) + # The counters are incremented from the launching thread, but the graphs are not + # necessarily done with; synchronize so that nothing lands in the next phase. + torch.cuda.synchronize() + + +def main() -> int: + torch.manual_seed(1234) + # No mask, no bias, no dropout: the simplest configuration cuDNN supports, so that the + # counts this produces are about the cache rather than about which graph got built. + 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) + + # get_available_attention_backends() enables every backend so it can report on all of + # them; the execution phases have to land on the fused one for their cache events to + # exist at all, so leave it the only one available. + 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 bfd2cdf9fd..0202c1899d 100644 --- a/tests/pytorch/attention/test_attention.py +++ b/tests/pytorch/attention/test_attention.py @@ -1,8 +1,12 @@ # 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 @@ -306,6 +310,151 @@ def test_dpa_checkpoint(dtype, model_configs, model): test_dot_product_attention(dtype, model_configs, model, True, None, False, False) +# One [FUSED-ATTN-CACHE] event, as either a counter line ("f16 fwd CREATE_GRAPH") or a level-2 +# trace line ("f16 fwd MISS"). Both name the thread and device first and then the build site, of +# which the backend half is matched but not counted on: the worker below drives whichever one its +# dtype selects, and every assertion here holds of either. The pass and the event name are what +# this test reads, plus the trace line's cache key, kept so that distinct keys can be counted. +# Every event name is also the counter column it increments. Requiring an event name is also what +# excludes the end-of-run summary rows, which are otherwise the same shape. The rank prefix is +# optional because it is emitted only when the launcher exports a rank, which a plain subprocess +# like the worker does not. +_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+)") + + +def _parse_cache_events(stderr: str): + """Group the worker's cache diagnostics by the phase that produced them. + + Returns (events, miss_keys): events[phase][(pass, event)] is a count, and + miss_keys[phase][pass] is the set of distinct cache keys that missed, so that "one extra + graph" can be told apart from "the same graph rebuilt". + """ + events = collections.defaultdict(collections.Counter) + miss_keys = collections.defaultdict(lambda: collections.defaultdict(set)) + phase = None + for line in 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 + pass_name, event = event_match.group("pass"), event_match.group("event") + events[phase][(pass_name, event)] += 1 + if event == "MISS": + miss_keys[phase][pass_name].add(event_match.group("rest").split("|")[-1].strip()) + return events, miss_keys + + +@pytest.mark.skipif(get_cudnn_version() < (8, 9, 1), reason="cuDNN 8.9.1+ is required.") +def test_fused_attn_graph_cache(): + """Test that the cuDNN graph cache is hit when it should be, and missed when it must be. + + A cuDNN graph build is the most expensive thing in a fused-attention call, so what this + checks is that each distinct configuration pays for one and no more: that a support query + builds the graph an execution then reuses, that a field the graph does not read (here + softmax_scale, which the key normalizes away) does not multiply the cache, and that a + field it does read still gets its own graph. The counters come from + NVTE_FUSED_ATTN_CACHE_DEBUG, which is also what a user would reach for to answer the + same question about their own model. + + The work runs in a subprocess (run_graph_cache.py): the cache is process-wide with + accumulating counters, so within pytest the graphs built by other tests would be + indistinguishable from this test's own. + """ + if torch.cuda.device_count() == 0: + pytest.skip("No CUDA device available.") + + worker = _current_file.parent / "run_graph_cache.py" + result = subprocess.run( + [sys.executable, str(worker)], + env={ + **os.environ, + # Level 2: the per-lookup HIT/MISS lines are what make the hits visible, and + # the volume is trivial for the handful of configurations below. + "NVTE_FUSED_ATTN_CACHE_DEBUG": "2", + "PYTHONUNBUFFERED": "1", + }, + capture_output=True, + text=True, + timeout=900, + check=False, + ) + assert result.returncode == 0, ( + f"{worker.name} failed with exit code {result.returncode}\n" + f"--- stdout ---\n{result.stdout}\n--- stderr (tail) ---\n{result.stderr[-4000:]}" + ) + if "[CACHE-TEST] fused=1" not in result.stdout: + pytest.skip("No cuDNN fused attention backend for the graph cache test config.") + + events, miss_keys = _parse_cache_events(result.stderr) + context = f"\n--- stderr ---\n{result.stderr[-8000:]}" + for phase in ("query", "requery", "exec", "rescale", "reshape"): + assert phase in events, f"worker emitted no cache events for phase {phase}{context}" + + # Both passes are queried by one call and executed by one forward/backward pair, so each + # of them sees the same sequence of events. + for pass_name in ("fwd", "bwd"): + + def count(phase, event, pass_name=pass_name): + return events[phase][(pass_name, event)] + + # The first query builds each pass's graph, and no more than its graph: a support + # query stops at check_support(), leaving the kernel compilation (BUILD_PLANS) to + # whoever executes it. A build cuDNN refused would show up here as the miss without + # the build, since nothing is recorded for a refusal. + assert count("query", "MISS") == 1, f"{pass_name}: expected one cold miss{context}" + assert count("query", "CREATE_GRAPH") == 1, f"{pass_name}: expected one build{context}" + assert count("query", "BUILD_PLANS") == 0, f"{pass_name}: query compiled kernels{context}" + # The graph cuDNN took, which is the one the execution phases below go on to find. A + # build refused by check_support() would show up here as CREATE_GRAPH without this. + assert count("query", "CACHE_GRAPH") == 1, f"{pass_name}: build was not cached{context}" + + # Asking the identical question again must cost nothing. + assert count("requery", "MISS") == 0, f"{pass_name}: repeated query missed{context}" + assert ( + count("requery", "CREATE_GRAPH") == 0 + ), f"{pass_name}: repeated query rebuilt{context}" + assert count("requery", "HIT") >= 1, f"{pass_name}: repeated query never looked{context}" + + # The execution must find the graph the query left behind -- a miss here is the + # probe/execute key drift this cache is most likely to develop -- and it is what + # finishes the build, exactly once. + assert ( + count("exec", "MISS") == 0 + ), f"{pass_name}: execution missed the query's graph{context}" + assert ( + count("exec", "CREATE_GRAPH") == 0 + ), f"{pass_name}: execution rebuilt the graph{context}" + assert count("exec", "EXECUTE") >= 1, f"{pass_name}: fused attention never ran{context}" + assert count("exec", "BUILD_PLANS") == 1, f"{pass_name}: expected one plan build{context}" + + # softmax_scale reaches the graph as a pointer, not as a shape, so the key drops it: + # a different scale has to reuse everything, down to the compiled kernels. + assert count("rescale", "MISS") == 0, f"{pass_name}: attn_scale changed the key{context}" + assert ( + count("rescale", "CREATE_GRAPH") == 0 + ), f"{pass_name}: attn_scale forced a build{context}" + assert count("rescale", "BUILD_PLANS") == 0, f"{pass_name}: attn_scale recompiled{context}" + assert ( + count("rescale", "EXECUTE") >= 1 + ), f"{pass_name}: rescaled run did not execute{context}" + + # max_seqlen is a dimension the graph is built at, so it must miss -- once, for one + # new graph, rather than invalidating what is already cached. + assert count("reshape", "MISS") == 1, f"{pass_name}: expected one miss{context}" + assert count("reshape", "CREATE_GRAPH") == 1, f"{pass_name}: expected one build{context}" + assert ( + len(miss_keys["reshape"][pass_name]) == 1 + ), f"{pass_name}: more than one new cache key{context}" + + model_configs_max_logit = { # test: ModelConfig(b, sq, hq, dqk) "max_logit_1": ModelConfig(1, 2048, 24, 128, max_seqlen_kv=4096), @@ -642,6 +791,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 +1064,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) @@ -1954,12 +2109,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: @@ -2155,6 +2322,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 = model_configs_fp8_vs_f16[model] @@ -2185,6 +2355,7 @@ def test_mha_fp8_vs_f16( 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, @@ -2312,6 +2483,8 @@ 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() @@ -2403,6 +2576,10 @@ def get_dummy_cuda_rng_tracker() -> CudaRNGStatesTracker: def test_dpa_fp8_vs_f16(dtype, model, qkv_layout, fp8_dpa_bwd, is_training, scaling_mode): """Test DotProductAttention module in FP8""" config = model_configs_fp8_vs_f16[model] + 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 @@ -2442,6 +2619,7 @@ def test_dpa_fp8_vs_f16(dtype, model, qkv_layout, fp8_dpa_bwd, is_training, scal 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, @@ -2461,8 +2639,6 @@ def test_dpa_fp8_vs_f16(dtype, model, qkv_layout, fp8_dpa_bwd, is_training, scal 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" @@ -2751,10 +2927,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, ) @@ -2832,6 +3020,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 d7eb16b862..6f2efaf83a 100644 --- a/tests/pytorch/attention/test_attention_with_cp.py +++ b/tests/pytorch/attention/test_attention_with_cp.py @@ -375,6 +375,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: @@ -650,6 +652,8 @@ 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 diff --git a/tests/pytorch/attention/test_kv_cache.py b/tests/pytorch/attention/test_kv_cache.py index 2a857a10dc..7bc20310b4 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 @@ -473,8 +474,11 @@ def test_kv_cache(dtype, model, qkv_format, is_paged, backend, module, is_cuda_g 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, + # probe inference configs only; reference configs are widely supported + probe_config = copy.deepcopy(config) + probe_config.attn_mask_type = "padding_causal" + available_backends, _, _ = get_available_attention_backends( + probe_config, qkv_dtype=dtype, qkv_layout=qkv_layout, pad_between_seqs=False, diff --git a/tests/pytorch/test_torch_compile.py b/tests/pytorch/test_torch_compile.py index 6ebc408dcb..ca8f2b63fd 100644 --- a/tests/pytorch/test_torch_compile.py +++ b/tests/pytorch/test_torch_compile.py @@ -1269,7 +1269,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..cdab36b2c8 100644 --- a/tests/pytorch/utils.py +++ b/tests/pytorch/utils.py @@ -312,6 +312,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 +340,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,6 +349,8 @@ 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, ) -> Tuple[List, List]: """Check for all available attention backends that support a model configuration""" @@ -358,9 +365,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 +382,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 +393,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, @@ -390,6 +404,7 @@ def test(): head_dim_v=config.head_dim_v, 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 +413,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 +454,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 6ceadc7405..4fbd33480a 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..751f075a47 --- /dev/null +++ b/transformer_engine/common/fused_attn/config_and_params.cpp @@ -0,0 +1,1264 @@ +/************************************************************************* + * 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() { + const int64_t b = static_cast(batch_size); + const int64_t sq = static_cast(max_seqlen_q); + const int64_t skv = static_cast(max_seqlen_kv); + + // Convenience fields + 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); + 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_mxfp8 = (scaling_mode == NVTE_MXFP8_1D_SCALING); + is_dropout = is_training && dropout != 0.0f; + + // Bucket the THD (ragged) batch and token counts + const size_t tokens_q = num_tokens_q != 0 ? num_tokens_q : static_cast(b * sq); + const size_t tokens_kv = num_tokens_kv != 0 ? num_tokens_kv : static_cast(b * skv); + 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(tokens_q) : 0; + bucketed_num_tokens_kv = is_ragged_kv ? fused_attn::get_max_tokens(tokens_kv) : 0; + + // Use cu_seqlens vs actual_seqlens + const size_t cudnn_runtime_version = cudnnGetVersion(); + uses_cu_seqlens_directly = CUDNN_FRONTEND_VERSION >= 12500 && + (CUDNN_VERSION >= 92400 && cudnn_runtime_version >= 92400) && + !is_dropout; + fp8_uses_cu_seqlens_directly = CUDNN_FRONTEND_VERSION >= 12600 && + (CUDNN_VERSION >= 92500 && cudnn_runtime_version >= 92500) && + !is_dropout; + + is_o_in_fp8 = (o_dtype == kNVTEFloat8E4M3 || o_dtype == kNVTEFloat8E5M2); + is_dqkv_in_fp8 = (dqkv_dtype == kNVTEFloat8E4M3 || dqkv_dtype == kNVTEFloat8E5M2); + const bool is_o_in_f16 = (o_dtype == kNVTEFloat16 || o_dtype == kNVTEBFloat16); + const bool is_dqkv_in_f16 = (dqkv_dtype == kNVTEFloat16 || dqkv_dtype == kNVTEBFloat16); + + // Determine the FP8 recipe + is_tensor_scaling = (scaling_mode == NVTE_DELAYED_TENSOR_SCALING); + is_delayed_scaling_fwd = is_tensor_scaling && is_o_in_fp8; + is_current_scaling_fwd = is_tensor_scaling && is_o_in_f16; + is_delayed_scaling_bwd = is_tensor_scaling && is_dqkv_in_fp8; + 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; + + // Whether packed graphs exist for THD + const int sm_arch = cuda::sm_arch(cuda::current_device()); + uses_packed_ragged_graph = cudnn_runtime_version >= 90600 && sm_arch >= 90 && sm_arch != 120; + uses_ragged_stats = is_ragged_q && uses_packed_ragged_graph; + + // Sequence lengths the graph is built at + graph_max_seqlen_q = + (is_ragged_q && uses_packed_ragged_graph) ? bucketed_num_tokens_q : max_seqlen_q; + graph_max_seqlen_kv = + (is_ragged_kv && uses_packed_ragged_graph) ? bucketed_num_tokens_kv : max_seqlen_kv; + + // Ragged-offset width that this config needs + 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; + + // Batch size the graph is built at + const bool buckets_the_batch = (is_ragged_q || is_ragged_kv) && uses_packed_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; + + // Elements per token for each ragged tensor + 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)); + + // Paged KV dimensions + if (is_paged_kv) { + if (num_pages_k == 0) { + num_pages_k = static_cast(b); + } + if (num_pages_v == 0) { + num_pages_v = static_cast(b); + } + if (page_size_k == 0) { + page_size_k = static_cast(skv); + } + if (page_size_v == 0) { + page_size_v = static_cast(skv); + } + if (max_pages_per_seq_k == 0) { + max_pages_per_seq_k = 1; + } + if (max_pages_per_seq_v == 0) { + max_pages_per_seq_v = 1; + } + } + + is_derived = true; +} + +FusedAttnConfig FusedAttnConfig::make_cache_key(Pass pass) const { + check_derived(); + FusedAttnConfig cache_cfg = *this; + + // Key the device ID for multi-GPU single-process runs + 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 batch size the graph is built at, and drop the token counts the bucketing replaced. + if ((cache_cfg.is_ragged_q || cache_cfg.is_ragged_kv) && cache_cfg.uses_packed_ragged_graph) { + cache_cfg.num_tokens_q = 0; + cache_cfg.num_tokens_kv = 0; + 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.qkv_layout = params.qkv_layout; + cfg.o_format = params.o_format; + cfg.qkv_scale_inv_format = params.qkv_scale_inv_format; + cfg.bias_type = params.bias_type; + cfg.attn_mask_type = params.attn_mask_type; + cfg.softmax_type = params.softmax_type; + cfg.attn_scale = params.attn_scale; + cfg.dropout = params.dropout; + cfg.max_seqlen_q = params.max_seqlen_q; + cfg.max_seqlen_kv = params.max_seqlen_kv; + cfg.window_size_left = params.window_size_left; + cfg.window_size_right = params.window_size_right; + cfg.bottom_right_diagonal = params.bottom_right_diagonal; + + 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.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.bias_type = params.bias_type; + cfg.attn_mask_type = params.attn_mask_type; + cfg.softmax_type = params.softmax_type; + cfg.attn_scale = params.attn_scale; + cfg.dropout = params.dropout; + cfg.max_seqlen_q = params.max_seqlen_q; + cfg.max_seqlen_kv = params.max_seqlen_kv; + cfg.window_size_left = params.window_size_left; + cfg.window_size_right = params.window_size_right; + cfg.bottom_right_diagonal = params.bottom_right_diagonal; + + 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); + 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 kNVTEFusedAttnFwdParamsMaxSeqlenQ: + std::memcpy(buf, &p.max_seqlen_q, attr_size); + break; + case kNVTEFusedAttnFwdParamsMaxSeqlenKV: + std::memcpy(buf, &p.max_seqlen_kv, attr_size); + break; + case kNVTEFusedAttnFwdParamsIsTraining: + bool_to_uint8(p.is_training, buf); + break; + case kNVTEFusedAttnFwdParamsReturnMaxLogit: + bool_to_uint8(p.return_max_logit, buf); + break; + case kNVTEFusedAttnFwdParamsCudaGraph: + bool_to_uint8(p.cuda_graph, buf); + break; + case kNVTEFusedAttnFwdParamsAttnScale: + std::memcpy(buf, &p.attn_scale, attr_size); + break; + case kNVTEFusedAttnFwdParamsDropout: + std::memcpy(buf, &p.dropout, 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 kNVTEFusedAttnFwdParamsBiasType: + std::memcpy(buf, &p.bias_type, attr_size); + break; + case kNVTEFusedAttnFwdParamsAttnMaskType: + std::memcpy(buf, &p.attn_mask_type, attr_size); + break; + case kNVTEFusedAttnFwdParamsSoftmaxType: + std::memcpy(buf, &p.softmax_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 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 kNVTEFusedAttnFwdParamsMaxSeqlenQ: + std::memcpy(&p.max_seqlen_q, buf, attr_size); + break; + case kNVTEFusedAttnFwdParamsMaxSeqlenKV: + std::memcpy(&p.max_seqlen_kv, buf, attr_size); + break; + case kNVTEFusedAttnFwdParamsIsTraining: + uint8_to_bool(buf, p.is_training); + break; + case kNVTEFusedAttnFwdParamsReturnMaxLogit: + uint8_to_bool(buf, p.return_max_logit); + break; + case kNVTEFusedAttnFwdParamsCudaGraph: + uint8_to_bool(buf, p.cuda_graph); + break; + case kNVTEFusedAttnFwdParamsAttnScale: + std::memcpy(&p.attn_scale, buf, attr_size); + break; + case kNVTEFusedAttnFwdParamsDropout: + std::memcpy(&p.dropout, 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 kNVTEFusedAttnFwdParamsBiasType: + std::memcpy(&p.bias_type, buf, attr_size); + break; + case kNVTEFusedAttnFwdParamsAttnMaskType: + std::memcpy(&p.attn_mask_type, buf, attr_size); + break; + case kNVTEFusedAttnFwdParamsSoftmaxType: + std::memcpy(&p.softmax_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 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 kNVTEFusedAttnBwdParamsMaxSeqlenQ: + std::memcpy(buf, &p.max_seqlen_q, attr_size); + break; + case kNVTEFusedAttnBwdParamsMaxSeqlenKV: + std::memcpy(buf, &p.max_seqlen_kv, attr_size); + break; + case kNVTEFusedAttnBwdParamsAttnScale: + std::memcpy(buf, &p.attn_scale, attr_size); + break; + case kNVTEFusedAttnBwdParamsDropout: + std::memcpy(buf, &p.dropout, 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 kNVTEFusedAttnBwdParamsBiasType: + std::memcpy(buf, &p.bias_type, attr_size); + break; + case kNVTEFusedAttnBwdParamsAttnMaskType: + std::memcpy(buf, &p.attn_mask_type, attr_size); + break; + case kNVTEFusedAttnBwdParamsSoftmaxType: + std::memcpy(buf, &p.softmax_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 kNVTEFusedAttnBwdParamsDeterministic: + bool_to_uint8(p.deterministic, buf); + break; + case kNVTEFusedAttnBwdParamsCudaGraph: + bool_to_uint8(p.cuda_graph, buf); + 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 kNVTEFusedAttnBwdParamsMaxSeqlenQ: + std::memcpy(&p.max_seqlen_q, buf, attr_size); + break; + case kNVTEFusedAttnBwdParamsMaxSeqlenKV: + std::memcpy(&p.max_seqlen_kv, buf, attr_size); + break; + case kNVTEFusedAttnBwdParamsAttnScale: + std::memcpy(&p.attn_scale, buf, attr_size); + break; + case kNVTEFusedAttnBwdParamsDropout: + std::memcpy(&p.dropout, 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 kNVTEFusedAttnBwdParamsBiasType: + std::memcpy(&p.bias_type, buf, attr_size); + break; + case kNVTEFusedAttnBwdParamsAttnMaskType: + std::memcpy(&p.attn_mask_type, buf, attr_size); + break; + case kNVTEFusedAttnBwdParamsSoftmaxType: + std::memcpy(&p.softmax_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 kNVTEFusedAttnBwdParamsDeterministic: + uint8_to_bool(buf, p.deterministic); + break; + case kNVTEFusedAttnBwdParamsCudaGraph: + uint8_to_bool(buf, p.cuda_graph); + 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..7bfe3dc1ad --- /dev/null +++ b/transformer_engine/common/fused_attn/config_and_params.h @@ -0,0 +1,494 @@ +/************************************************************************* + * 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 }; + +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: keyed + // + // device ID is not part of attribute serialization, i.e. internal, but it participates in + // operator< and is used to differentiate graphs built for different devices in multi-GPU + // single-process runs + int device_id = -1; + + // Internal fields: not keyed. The following fields are not part of attribute serialization, + // operator<, or the cache key; they are filled by derive() or set by caller such as with + // check_for_forward_support, and are used for convinence purposes + // + // run query_support() for forward or backward + bool check_for_forward_support = true; + bool check_for_backward_support = true; + // whether derive() has been run + bool is_derived = false; + // bucketed batch size/token counts for THD + size_t bucketed_batch_size = 0; + size_t bucketed_num_tokens_q = 0; + size_t bucketed_num_tokens_kv = 0; + // whether to use cu_seqlens or actual_seqlens for THD or padding masks + bool uses_cu_seqlens_directly = false; + bool fp8_uses_cu_seqlens_directly = false; + // Whether packed graphs exist for THD. A memory optimization, not a correctness gate: where this + // is false, ragged input still builds a correct graph, just the dense one, whose dimensions are + // max_seqlen rather than the token total and whose Stats is BHS1 rather than TH1. Nothing may + // gate support on it -- which architectures run ragged attention at all is cuDNN's answer. + bool uses_packed_ragged_graph = false; + bool uses_ragged_stats = false; + // sequence lengths the graph is built at + size_t graph_max_seqlen_q = 0; + size_t graph_max_seqlen_kv = 0; + // batch size the graph is built at + size_t graph_batch_size_fwd = 0; + size_t graph_batch_size_bwd = 0; + // ragged offset type for THD + DType ragged_offset_type_fwd = DType::kInt32; + DType ragged_offset_type_bwd = DType::kInt32; + // Whether this config's ragged offsets overflow 32 bits. The counterpart to the two fields above + // rather than a third of them: those are the width TE will use, already capped by the running + // cuDNN, while this is the width the config needs. nvte_get_fused_attn_backend_v2 compares the + // two and refuses the config whose need outruns the cuDNN it is running on. + bool needs_64bit_ragged_offset = false; + // elements per token for each ragged tensor + RaggedOffsetMultipliers ragged_offset_mults; + // convenience fields + // + // qkv_format is the combined format: it says what Q and KV each are and whether they agree, with + // the mixed layouts keeping their own enumerators (NVTE_THD_2BSHD and friends) rather than + // collapsing onto either side. Ask it only where a rule means "Q and KV are the same dense + // layout". Anything about raggedness belongs to is_ragged_q/is_ragged_kv instead, because + // NVTE_THD names only the fully ragged layouts, so a test against it passes THD_BSHD_BSHD + // straight through -- a rule here did exactly that until it was found. + 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_mxfp8 = false; + bool is_dropout = false; + bool is_o_in_fp8 = false; + bool is_dqkv_in_fp8 = false; + // Whether the FP8 recipe is tensor scaling, i.e. delayed or current rather than MXFP8. The + // graphs need this on its own, wherever a tensor is per-tensor scaled and it does not matter + // which of the two put the scale there. + bool is_tensor_scaling = false; + // Which recipe serves each pass, one flag per recipe per pass. Each means "this recipe is in + // effect and can write this pass's output dtype", so at most one of a pass's three holds, and all + // three false is a configuration no FP8 graph is written for -- which is what lets + // nvte_get_fused_attn_backend_v2 refuse it by asking three booleans and nothing else. + // + // Delayed against current is told apart by that output dtype rather than by scaling_mode, because + // NVTEScalingMode has no current-scaling enumerator: both arrive as NVTE_DELAYED_TENSOR_SCALING, + // and what separates them is that delayed knows the output scale before the graph is built while + // current has the graph compute it. So delayed writes FP8 and current writes F16/BF16, and an + // output dtype neither can write (FP32, say) leaves both false rather than defaulting to one. + bool is_delayed_scaling_fwd = false; + bool is_current_scaling_fwd = false; + bool is_mxfp8_fwd = false; + bool is_delayed_scaling_bwd = false; + bool is_current_scaling_bwd = false; + bool is_mxfp8_bwd = false; + + 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. + // + // Computes only; it validates nothing. Rules about which configurations are legal belong in + // nvte_get_fused_attn_backend_v2, which derives first and then states them once, so that a + // violation comes back as an unsupported configuration instead of being thrown from a query. + void derive(); + + // Assert that derive() has run, for code about to read a derived field. Worth asserting rather + // than assuming because the failure is silent: an unset derived field reads as zero, which is a + // legal value that yields a graph of the wrong shape and a key that collides with unrelated + // configs. + 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 + size_t max_seqlen_q = 0; + size_t max_seqlen_kv = 0; + bool is_training = true; + bool return_max_logit = false; + bool cuda_graph = false; + float attn_scale = 1.0f; + float dropout = 0.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; + NVTE_Bias_Type bias_type = NVTE_NO_BIAS; + NVTE_Mask_Type attn_mask_type = NVTE_NO_MASK; + NVTE_Softmax_Type softmax_type = NVTE_VANILLA_SOFTMAX; + int64_t window_size_left = -1; + int64_t window_size_right = -1; + bool bottom_right_diagonal = true; + // Workspace and stream + NVTETensor workspace = nullptr; + cudaStream_t stream = nullptr; + + static constexpr size_t attr_sizes[] = { + 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 + sizeof(size_t), // max_seqlen_q + sizeof(size_t), // max_seqlen_kv + sizeof(uint8_t), // is_training + sizeof(uint8_t), // return_max_logit + sizeof(uint8_t), // cuda_graph + sizeof(float), // attn_scale + sizeof(float), // dropout + sizeof(NVTE_QKV_Layout), // qkv_layout + sizeof(NVTE_QKV_Format), // o_format + sizeof(NVTE_QKV_Format), // qkv_scale_inv_format + sizeof(NVTE_Bias_Type), // bias_type + sizeof(NVTE_Mask_Type), // attn_mask_type + sizeof(NVTE_Softmax_Type), // softmax_type + sizeof(int64_t), // window_size_left + sizeof(int64_t), // window_size_right + sizeof(uint8_t), // bottom_right_diagonal + 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 + size_t max_seqlen_q = 0; + size_t max_seqlen_kv = 0; + float attn_scale = 1.0f; + float dropout = 0.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; + NVTE_Bias_Type bias_type = NVTE_NO_BIAS; + NVTE_Mask_Type attn_mask_type = NVTE_NO_MASK; + NVTE_Softmax_Type softmax_type = NVTE_VANILLA_SOFTMAX; + int64_t window_size_left = -1; + int64_t window_size_right = -1; + bool bottom_right_diagonal = true; + bool deterministic = false; + bool cuda_graph = false; + // Workspace and stream + NVTETensor workspace = nullptr; + cudaStream_t stream = nullptr; + + static constexpr size_t attr_sizes[] = { + 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 + sizeof(size_t), // max_seqlen_q + sizeof(size_t), // max_seqlen_kv + sizeof(float), // attn_scale + sizeof(float), // dropout + 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(NVTE_Bias_Type), // bias_type + sizeof(NVTE_Mask_Type), // attn_mask_type + sizeof(NVTE_Softmax_Type), // softmax_type + sizeof(int64_t), // window_size_left + sizeof(int64_t), // window_size_right + sizeof(uint8_t), // bottom_right_diagonal + sizeof(uint8_t), // deterministic + sizeof(uint8_t), // cuda_graph + 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 1ac7a36383..dc8a5be41d 100644 --- a/transformer_engine/common/fused_attn/fused_attn.cpp +++ b/transformer_engine/common/fused_attn/fused_attn.cpp @@ -6,10 +6,13 @@ #include "transformer_engine/fused_attn.h" +#include + #include "../common.h" #include "../cudnn_utils.h" #include "../util/cuda_runtime.h" #include "../util/system.h" +#include "config_and_params.h" #include "fused_attn_f16_arbitrary_seqlen.h" #include "fused_attn_fp8.h" #include "utils.h" @@ -183,7 +186,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 +208,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,326 +228,266 @@ 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, re-populated by every rejection on this +// thread. `*message` is handed a pointer into it, which is what limits how long that pointer +// stays good: only until the next rejection on the same thread. +thread_local std::string fused_attn_backend_message_buffer; + +// Records `reason` and answers with the backend that means "none", so that a rejection reads as +// the one statement it is: `if (cond) return reject(message, "why");`. +[[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: returns the backend that supports the given configuration; +// otherwise, returns NVTE_No_Backend and a diagnostic message. It performs several TE-specific +// checks before running the cuDNN support query. +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)) || - // 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} - ((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))) && - !requires_64bit_ragged_offset && - // 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 || - (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; + const FusedAttnConfig &caller_cfg = *get_fused_attn_config(config); + if (message != nullptr) *message = ""; + std::optional derived_cfg; + if (!caller_cfg.is_derived) { + derived_cfg = caller_cfg; + derived_cfg->derive(); + } + const FusedAttnConfig &cfg = caller_cfg.is_derived ? caller_cfg : *derived_cfg; + + 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."); + } + + // Ragged (THD) input requires a padding-style 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."); + } + + // Paged KV requires a padding-style mask, for the same reason ragged input does. + if (cfg.is_paged_kv && !cfg.is_padding) { + return reject(message, + "Paged KV requires PADDING / PADDING_CAUSAL / PADDING_CAUSAL_BOTTOM_RIGHT mask."); + } + + // cuDNN attention graphs do 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); + + // Ask `verdict` about each direction the caller wants, and report the first refusal + 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; + } + return ""; + }; + + 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 < 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; + 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) { + return reject(message, "FP8 fused attention supports BSHD/SBHD/BHSD formats, found " + + std::to_string(static_cast(cfg.qkv_format)) + "."); } - 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.is_bias) { + return reject(message, "FP8 fused attention does not support pre/post_scale_bias yet!"); } - 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_alibi) { + return reject(message, "FP8 fused attention does not support ALiBi yet!"); } - 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; - } + + std::string recipe_reason = each_pass([&](Pass pass) -> std::string { + const bool serves_this_output = + (pass == Pass::Fwd) + ? (cfg.is_delayed_scaling_fwd || cfg.is_current_scaling_fwd || cfg.is_mxfp8_fwd) + : (cfg.is_delayed_scaling_bwd || cfg.is_current_scaling_bwd || cfg.is_mxfp8_bwd); + if (!serves_this_output) { + return "FP8 fused attention only supports FP8DelayedScaling or FP8CurrentScaling or MXFP8 " + "recipes!"; } + return ""; + }); + if (!recipe_reason.empty()) return reject(message, std::move(recipe_reason)); + + if (cfg.is_mxfp8 && cudnn_runtime_version < 92100) { + return reject(message, "MXFP8 fused attention requires cuDNN 9.21.0 or later!"); } + + 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; + } + + if (is_f16_or_bf16) { + 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; + } + + 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.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: derive the config, ask the selector which backend can run it, and run +// that backend's implementation. +// +// Both reach the same cache through the same accessor, which is what the HIT below means -- the +// entry execution needs was built by the probe that selected the backend, so what was checked is +// what runs. See graph_cache.h for the cache, graph_cache_debug.h for the events it emits. +// +// nvte_fused_attn_fwd_v2 +// | +// +-- cfg = p.make_config(); cfg.derive() +// | +// +-- nvte_get_fused_attn_backend_v2 TE's own rules, then a probe per backend +// | `-- support_verdict_f16 / support_verdict_fp8, with Pass::Fwd +// | `-- get_graph(): builds and inserts the entry, or refuses with cuDNN's reason +// | +// `-- fused_attn_arbitrary_seqlen_fwd -> ..._fwd_impl the selected backend +// `-- get_graph() HIT -> build_plans() -> bind device pointers, graph.execute() +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_rng_state, wkspace, p.stream, handle); } else { - backend = NVTE_Fused_Attn_Backend::NVTE_No_Backend; + const char *const reject_reason = + (fused_attn_reject_reason != nullptr && fused_attn_reject_reason[0] != '\0') + ? fused_attn_reject_reason + : "no cuDNN fused-attention backend supports the requested parameters"; + NVTE_ERROR("Fused attention is not supported for this configuration: ", reject_reason); } - return backend; } // NVTE fused attention FWD with separate Q, K and V @@ -562,99 +505,119 @@ 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.max_seqlen_q = max_seqlen_q; + p.max_seqlen_kv = max_seqlen_kv; + p.is_training = is_training; + p.return_max_logit = return_max_logit; + p.cuda_graph = cuda_graph; + p.attn_scale = attn_scale; + p.dropout = dropout; + p.qkv_layout = qkv_layout; + p.o_format = o_format; + p.qkv_scale_inv_format = qkv_scale_inv_format; + p.bias_type = bias_type; + p.attn_mask_type = attn_mask_type; + p.softmax_type = softmax_type; + p.window_size_left = window_size_left; + p.window_size_right = window_size_right; + p.bottom_right_diagonal = bottom_right_diagonal; + 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, 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_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_rng_state, wkspace, p.stream, handle); } else { - NVTE_ERROR("Invalid combination of data type and sequence length for fused attention. \n"); + const char *const reject_reason = + (fused_attn_reject_reason != nullptr && fused_attn_reject_reason[0] != '\0') + ? fused_attn_reject_reason + : "no cuDNN fused-attention backend supports the requested parameters"; + NVTE_ERROR("Fused attention is not supported for this configuration: ", 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, @@ -671,93 +634,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, 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_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.max_seqlen_q = max_seqlen_q; + p.max_seqlen_kv = max_seqlen_kv; + p.attn_scale = attn_scale; + p.dropout = dropout; + 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.bias_type = bias_type; + p.attn_mask_type = attn_mask_type; + p.softmax_type = softmax_type; + p.window_size_left = window_size_left; + p.window_size_right = window_size_right; + p.bottom_right_diagonal = bottom_right_diagonal; + p.deterministic = deterministic; + p.cuda_graph = cuda_graph; + 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..175c1c327c 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_packed_ragged_graph = cfg.uses_packed_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_packed_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..a373bb5e85 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,48 @@ * \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); + +// cuDNN's verdict on this config's F16/BF16 graph for `pass`: an empty string if it can run, +// otherwise a diagnostic message explaining why not. A verdict of "supported" leaves the graph in +// the cache, where the execution path finds it. +// +// The direction is a runtime argument, not two functions, because the graph builder it selects is +// local to this translation unit -- so this is the only place that can map one to the other. +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 000af41aee..8d05a53ecf 100644 --- a/transformer_engine/common/fused_attn/fused_attn_fp8.cu +++ b/transformer_engine/common/fused_attn/fused_attn_fp8.cu @@ -4,417 +4,349 @@ * See LICENSE for license information. ************************************************************************/ +#include + #include "../common.h" #include "../cudnn_utils.h" #include "../util/system.h" #include "fused_attn_fp8.h" +#include "graph_cache.h" +#include "graph_cache_debug.h" #include "utils.h" namespace transformer_engine { namespace fused_attn { using namespace transformer_engine; - -// 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, - 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* 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, // 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)); - bool is_padding = ((mask_type == NVTE_Mask_Type::NVTE_PADDING_MASK) || - (mask_type == NVTE_Mask_Type::NVTE_PADDING_CAUSAL_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!"); - - // Newer versions of cuDNN SDPA can accept sequence lengths directly as a cumulative - // tensor. Take advantage of this if possible to avoid 1 extra kernel call. (Unlike - // the F16 path, the FP8 path has no THD/ragged-offset support, so only the - // cu_seqlens_to_actual_seqlens conversion applies here. Also note that the - // needed versions of cuDNN backend and frontend are higher than for F16.) - 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; + 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.batch_size); + 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.fp8_uses_cu_seqlens_directly; + + 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 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)); + 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)); + 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, // 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 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)); - 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)); - 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); - // } - - 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); - } - } + 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_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); - } + 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); + 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); + + Stats->set_output(true) + .set_data_type(fe::DataType_t::FLOAT) + .set_dim({b, h, s_q, 1}) + .set_stride({h * s_q, s_q, 1, 1}); + + std::tuple, // Q + std::shared_ptr, // K + std::shared_ptr, // V + std::shared_ptr, // descale_q + std::shared_ptr, // descale_k + std::shared_ptr, // descale_v + std::shared_ptr, // descale_s + std::shared_ptr, // scale_s + std::shared_ptr, // scale_o + std::shared_ptr, // attn_scale + std::shared_ptr, // O + std::shared_ptr, // amax_s + std::shared_ptr> // amax_o + key_tensors_tuple = + is_mxfp8 ? std::make_tuple(Q, K, V, descale_q, descale_k, descale_v, nullptr, nullptr, + nullptr, attn_scale, O, nullptr, amax_o) + : std::make_tuple(Q, K, V, descale_q, descale_k, descale_v, descale_s, scale_s, + scale_o, attn_scale, O, amax_s, amax_o); + auto Stats_tuple = std::make_tuple(Stats); + auto bias_tuple = is_bias ? std::make_tuple(bias) : std::make_tuple(nullptr); + auto softmax_offset_tuple = + is_softmax_offset ? std::make_tuple(softmax_offset) : std::make_tuple(nullptr); + auto padding_tuple = + is_padding ? std::make_tuple(seq_q, seq_kv) : std::make_tuple(nullptr, nullptr); + auto 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, 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_options.set_sink_token(softmax_offset); - } +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* devPtrDropoutSeed, + void* devPtrDropoutOffset, void* workspace, size_t* workspace_size, + cudaStream_t stream, cudnnHandle_t handle) { + using namespace transformer_engine; - 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]; - } + cfg.check_derived(); - 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); - 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); - - Stats->set_output(true) - .set_data_type(fe::DataType_t::FLOAT) - .set_dim({b, h, s_q, 1}) - .set_stride({h * s_q, s_q, 1, 1}); - - std::tuple, // Q - std::shared_ptr, // K - std::shared_ptr, // V - std::shared_ptr, // descale_q - std::shared_ptr, // descale_k - std::shared_ptr, // descale_v - std::shared_ptr, // descale_s - std::shared_ptr, // scale_s - std::shared_ptr, // scale_o - std::shared_ptr, // attn_scale - std::shared_ptr, // O - std::shared_ptr, // amax_s - std::shared_ptr> // amax_o - key_tensors_tuple = - is_mxfp8 ? std::make_tuple(Q, K, V, descale_q, descale_k, descale_v, nullptr, nullptr, - nullptr, attn_scale, O, nullptr, amax_o) - : std::make_tuple(Q, K, V, descale_q, descale_k, descale_v, descale_s, - scale_s, scale_o, attn_scale, O, amax_s, amax_o); - auto Stats_tuple = std::make_tuple(Stats); - auto bias_tuple = is_bias ? std::make_tuple(bias) : std::make_tuple(nullptr); - auto softmax_offset_tuple = - is_softmax_offset ? std::make_tuple(softmax_offset) : std::make_tuple(nullptr); - auto padding_tuple = - is_padding ? std::make_tuple(seq_q, seq_kv) : std::make_tuple(nullptr, nullptr); - auto 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, dropout_tuple); - cache.insert({descriptor, return_tuple}); - - return return_tuple; - }; + 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.fp8_uses_cu_seqlens_directly; + + const int64_t 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; + 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, dropout_seed, - dropout_offset] = get_graph(sdpa_fp8_fprop_cache, descriptor); + 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); auto plan_workspace_size = mha_graph->get_workspace_size(); @@ -426,7 +358,6 @@ void fused_attn_fp8_fwd_impl( *workspace_size = plan_workspace_size + actual_seqlen_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)); @@ -446,14 +377,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; } */ @@ -466,8 +397,9 @@ void fused_attn_fp8_fwd_impl( const size_t grid = (b + nthreads_per_block - 1) / nthreads_per_block; void* devActualSeqlenQ = static_cast(workspace) + plan_workspace_size; void* devActualSeqlenKV = static_cast(devActualSeqlenQ) + b * sizeof(int32_t); + // TODO(cyanguwa): pass bucketed_batch_size cu_seqlens_to_actual_seqlens<<>>( - b, b, static_cast(devPtrcuSeqlensQ), // TODO(pass max_b) + b, b, static_cast(devPtrcuSeqlensQ), static_cast(devPtrcuSeqlensKV), static_cast(devActualSeqlenQ), static_cast(devActualSeqlenKV)); NVTE_CHECK_CUDA(cudaGetLastError()); @@ -486,554 +418,507 @@ 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()); } } -// 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, - 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* 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, // 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)); - bool is_padding = ((mask_type == NVTE_Mask_Type::NVTE_PADDING_MASK) || - (mask_type == NVTE_Mask_Type::NVTE_PADDING_CAUSAL_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!"); - - 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, // 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 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") + 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.batch_size); + 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; + + 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 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)); + Stats = mha_graph->tensor(fe::graph::Tensor_attributes() + .set_name("Stats") + .set_dim({b, h, s_q, 1}) + .set_stride({h * s_q, s_q, 1, 1}) + .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)); + + // 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_strides) + .set_stride(q_t_strides) .set_data_type(qkv_tensor_type)); - K = mha_graph->tensor(fe::graph::Tensor_attributes() - .set_name("K") + K_t = mha_graph->tensor(fe::graph::Tensor_attributes() + .set_name("K_t") .set_dim({b, hg, s_kv, d_qk}) - .set_stride(k_strides) + .set_stride(k_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) - .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") + 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)); - Stats = mha_graph->tensor(fe::graph::Tensor_attributes() - .set_name("Stats") - .set_dim({b, h, s_q, 1}) - .set_stride({h * s_q, s_q, 1, 1}) - .set_data_type(fe::DataType_t::FLOAT)); - attn_scale = mha_graph->tensor(fe::graph::Tensor_attributes() - .set_name("attn_scale") + dO_f16 = mha_graph->tensor(fe::graph::Tensor_attributes() + .set_name("dO_f16") + .set_dim({b, h, s_q, d_v}) + .set_stride(dO_strides) + .set_data_type(o_tensor_type)); + // Descale_q, Descale_q_t, Descale_k, Descale_k_t, Descale_v, Descale_dO, Descale_dO_t + auto padded = pad_s_d_for_mxfp8(s_q, s_kv, d_qk, d_v); + std::vector q_scale_strides(4), q_t_scale_strides(4), k_scale_strides(4), + k_t_scale_strides(4), v_scale_strides(4), dO_scale_strides(4), dO_t_scale_strides(4); + generateMatrixStridesWithFormat(b, h, padded.s_q_padded, padded.d_qk_scale_padded, + q_scale_strides.data(), q_scale_inv_format); + generateMatrixStridesWithFormat(b, h, padded.s_q_scale_padded, padded.d_qk_padded, + q_t_scale_strides.data(), q_scale_inv_format); + generateMatrixStridesWithFormat(b, hg, padded.s_kv_padded, padded.d_qk_scale_padded, + k_scale_strides.data(), kv_scale_inv_format); + generateMatrixStridesWithFormat(b, hg, padded.s_kv_scale_padded, padded.d_qk_padded, + k_t_scale_strides.data(), kv_scale_inv_format); + generateMatrixStridesWithFormat(b, hg, padded.s_kv_padded, padded.d_v_scale_padded, + v_scale_strides.data(), kv_scale_inv_format); + generateMatrixStridesWithFormat(b, h, padded.s_q_padded, padded.d_v_scale_padded, + dO_scale_strides.data(), do_scale_format_); + generateMatrixStridesWithFormat(b, h, padded.s_q_scale_padded, padded.d_v_padded, + dO_t_scale_strides.data(), do_scale_format_); + descale_q = mha_graph->tensor(fe::graph::Tensor_attributes() + .set_name("Descale_q") + .set_dim({b, h, padded.s_q_padded, padded.d_qk_scale_padded}) + .set_stride(q_scale_strides) + .set_data_type(fe::DataType_t::FP8_E8M0) + .set_reordering_type(fe::TensorReordering_t::F8_128x4)); + descale_q_t = + mha_graph->tensor(fe::graph::Tensor_attributes() + .set_name("Descale_q_t") + .set_dim({b, h, padded.s_q_scale_padded, padded.d_qk_padded}) + .set_stride(q_t_scale_strides) + .set_data_type(fe::DataType_t::FP8_E8M0) + .set_reordering_type(fe::TensorReordering_t::F8_128x4)); + descale_k = + mha_graph->tensor(fe::graph::Tensor_attributes() + .set_name("Descale_k") + .set_dim({b, hg, padded.s_kv_padded, padded.d_qk_scale_padded}) + .set_stride(k_scale_strides) + .set_data_type(fe::DataType_t::FP8_E8M0) + .set_reordering_type(fe::TensorReordering_t::F8_128x4)); + descale_k_t = + mha_graph->tensor(fe::graph::Tensor_attributes() + .set_name("Descale_k_t") + .set_dim({b, hg, padded.s_kv_scale_padded, padded.d_qk_padded}) + .set_stride(k_t_scale_strides) + .set_data_type(fe::DataType_t::FP8_E8M0) + .set_reordering_type(fe::TensorReordering_t::F8_128x4)); + descale_v = mha_graph->tensor(fe::graph::Tensor_attributes() + .set_name("Descale_v") + .set_dim({b, hg, padded.s_kv_padded, padded.d_v_scale_padded}) + .set_stride(v_scale_strides) + .set_data_type(fe::DataType_t::FP8_E8M0) + .set_reordering_type(fe::TensorReordering_t::F8_128x4)); + descale_dO = mha_graph->tensor(fe::graph::Tensor_attributes() + .set_name("Descale_dO") + .set_dim({b, h, padded.s_q_padded, padded.d_v_scale_padded}) + .set_stride(dO_scale_strides) + .set_data_type(fe::DataType_t::FP8_E8M0) + .set_reordering_type(fe::TensorReordering_t::F8_128x4)); + descale_dO_t = + mha_graph->tensor(fe::graph::Tensor_attributes() + .set_name("Descale_dO_t") + .set_dim({b, h, padded.s_q_scale_padded, padded.d_v_padded}) + .set_stride(dO_t_scale_strides) + .set_data_type(fe::DataType_t::FP8_E8M0) + .set_reordering_type(fe::TensorReordering_t::F8_128x4)); + } + + fe::graph::SDPA_fp8_backward_attributes sdpa_backward_options; + sdpa_backward_options = fe::graph::SDPA_fp8_backward_attributes() + .set_name("sdpa_fp8_backward") + .set_causal_mask(is_causal) + .set_attn_scale(attn_scale); + + fe::DiagonalAlignment_t const& diagonal_alignment = bottom_right_diagonal + ? fe::DiagonalAlignment_t::BOTTOM_RIGHT + : fe::DiagonalAlignment_t::TOP_LEFT; + sdpa_backward_options.set_diagonal_alignment(diagonal_alignment); + + if (cudnn_runtime_version >= 92100) { + if (window_size_left != -1) { + sdpa_backward_options.set_diagonal_band_left_bound(window_size_left + 1); + } + if (window_size_right != -1) { + sdpa_backward_options.set_diagonal_band_right_bound(window_size_right); + } + } + if (is_causal_bottom_right) { + sdpa_backward_options.set_diagonal_band_right_bound(0); + } + + // sdpa_backward_options.set_alibi_mask(is_alibi); + + // if (is_bias) { + // bias = mha_graph->tensor(fe::graph::Tensor_attributes() + // .set_name("bias") + // .set_dim({bias_b, bias_h, bias_sq, bias_skv}) + // .set_stride({bias_h * bias_sq * bias_skv, bias_sq * bias_skv, bias_skv, 1})); + // dBias = mha_graph->tensor(fe::graph::Tensor_attributes() + // .set_name("dBias") + // .set_dim({bias_b, bias_h, bias_sq, bias_skv}) + // .set_stride({bias_h * bias_sq * bias_skv, bias_sq * bias_skv, bias_skv, 1})); + // sdpa_backward_options.set_bias(bias); + // bias shapes [1, 1, s, s], [b, 1, s, s], [b, h, s, s], [1, h, s, s] are supported for dbias calculation + // bias shape [1, 1, 1, s] is not supported for dbias calculation as of cuDNN 9.18 + // if (!((bias_b == 1) && (bias_h == 1) && (bias_sq == 1))) { + // sdpa_backward_options.set_dbias(dBias); + // } + // } + + if (cudnn_runtime_version >= 91900) { + sdpa_backward_options.set_deterministic_algorithm(deterministic); + } + + if (is_padding) { + seq_q = mha_graph->tensor(fe::graph::Tensor_attributes() + .set_name("seq_q") + .set_dim({b, 1, 1, 1}) + .set_stride({1, 1, 1, 1}) + .set_data_type(fe::DataType_t::INT32)); + seq_kv = mha_graph->tensor(fe::graph::Tensor_attributes() + .set_name("seq_kv") + .set_dim({b, 1, 1, 1}) + .set_stride({1, 1, 1, 1}) + .set_data_type(fe::DataType_t::INT32)); + sdpa_backward_options.set_padding_mask(is_padding).set_seq_len_q(seq_q).set_seq_len_kv(seq_kv); + } + + if (is_dropout) { + dropout_seed = mha_graph->tensor(fe::graph::Tensor_attributes() + .set_name("Seed") .set_dim({1, 1, 1, 1}) .set_stride({1, 1, 1, 1}) - .set_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)); - } + .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); + } - 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_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); + } - // 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); - } + 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); + 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_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::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 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, dropout_tuple); +} - 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); - } +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* devPtrDropoutSeed, + void* devPtrDropoutOffset, void* workspace, size_t* workspace_size, cudaStream_t stream, + cudnnHandle_t handle) { + using namespace transformer_engine; - 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); - } + cfg.check_derived(); - 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); - 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); - 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); - } + 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; - 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 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, dropout_tuple); - cache.insert({descriptor, return_tuple}); - - return return_tuple; - }; + const int64_t 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; + + 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, - dropout_seed, dropout_offset] = get_graph(sdpa_fp8_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::FP8, Pass::Bwd, *cache_entry); auto plan_workspace_size = mha_graph->get_workspace_size(); @@ -1043,7 +928,6 @@ void fused_attn_fp8_bwd_impl( *workspace_size = plan_workspace_size + actual_seqlen_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)); @@ -1065,7 +949,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; @@ -1093,9 +977,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; @@ -1107,8 +991,9 @@ void fused_attn_fp8_bwd_impl( const size_t grid = (b + nthreads_per_block - 1) / nthreads_per_block; void* devActualSeqlenQ = static_cast(workspace) + plan_workspace_size; void* devActualSeqlenKV = static_cast(devActualSeqlenQ) + b * sizeof(int32_t); + // TODO(cyanguwa): pass bucketed_batch_size cu_seqlens_to_actual_seqlens<<>>( - b, b, static_cast(devPtrcuSeqlensQ), // TODO(pass max_b) + b, b, static_cast(devPtrcuSeqlensQ), static_cast(devPtrcuSeqlensKV), static_cast(devActualSeqlenQ), static_cast(devActualSeqlenKV)); NVTE_CHECK_CUDA(cudaGetLastError()); @@ -1127,25 +1012,31 @@ 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, 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* 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* 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 NVTE_QKV_Layout qkv_layout = cfg.qkv_layout; + 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; @@ -1212,22 +1103,16 @@ void fused_attn_fp8_fwd( 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); + const 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)) { 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, - 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, 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); + cfg, devPtrQ, devPtrK, devPtrV, devPtrSoftmaxOffset, devPtrM, devPtrO, devPtrDescaleQ, + devPtrDescaleK, devPtrDescaleV, devPtrDescaleS, devPtrScaleS, devPtrScaleO, devPtrAmaxO, + devPtrAmaxS, devPtrcuSeqlensQ, devPtrcuSeqlensKV, devPtrDropoutSeed, devPtrDropoutOffset, + workspace->data.dptr, &workspace_size, stream, handle); } else { NVTE_ERROR("FP8 fused attention only supports qkv_format=BSHD, SBHD, or BHSD.\n"); } @@ -1245,21 +1130,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, 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* 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* rng_state, Tensor* workspace, + cudaStream_t stream, cudnnHandle_t handle) { using namespace transformer_engine; + + const NVTE_QKV_Layout dqkv_layout = cfg.dqkv_layout; + 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; @@ -1275,8 +1158,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; @@ -1332,28 +1215,20 @@ void fused_attn_fp8_bwd( 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); + const 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)) { 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, - 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, + 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, - 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); + devPtrDropoutSeed, devPtrDropoutOffset, workspace->data.dptr, &workspace_size, stream, + handle); } else { NVTE_ERROR("FP8 fused attention only supports dqkv_format=BSHD, SBHD, or BHSD.\n"); } @@ -1370,4 +1245,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 b9660128ca..2749aa1fc1 100644 --- a/transformer_engine/common/fused_attn/fused_attn_fp8.h +++ b/transformer_engine/common/fused_attn/fused_attn_fp8.h @@ -8,35 +8,40 @@ * \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, 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 *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 *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, 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 *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 *rng_state, Tensor *workspace, cudaStream_t stream, + cudnnHandle_t handle); + +// The FP8 counterpart of support_verdict_f16; see there. +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..a694e7590b --- /dev/null +++ b/transformer_engine/common/fused_attn/graph_cache.h @@ -0,0 +1,181 @@ +/************************************************************************* + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * + * See LICENSE for license information. + ************************************************************************/ + +// The fused-attention graph cache: what a cache entry is, and how one is built, cached and found +// again. The four build sites for cuDNN graphs -- f16/fp8 crossed with fwd/bwd -- differ only in +// what the graph computes and which tensors it binds; caching, lookup, locking, the support query +// and the plan build are the same for all four and live here. +// +// The pieces below elide the `backend, pass` pair most of them also take: it never steers the +// logic, only attributing debug counters and stage timings to a build site. +// +// - CacheEntry: a graph, the tensors it binds as inputs and outputs, and a once_flag guarding its +// plan build. +// - GraphCache: process-wide map from a normalized FusedAttnConfig to a CacheEntry. +// - get_graph(cfg, handle): the execution path's way in. Keys `cfg` +// and owns the cache for its one triple; kCreateGraphFn is a create_graph_f16/fp8_fwd/bwd from a +// .cu file, the only piece a build site supplies. +// - support_verdict<...>(cfg, handle): the backend selector's way in. get_graph() in a try, +// returning the empty string when cuDNN accepts the graph and its complaint when it does not. +// - cache_graph(cache, key, handle, build): a hit, or a build under frontend_build_mutex() and an +// insert. The work behind both of the above. +// - query_support(graph, handle): takes a constructed graph through validate, +// build_operation_graph, create_execution_plans and check_support; throws cuDNN's message on +// refusal. +// - build_plans(entry): the kernel compilation cache_graph() deferred, once per entry, no handle. + +#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 + +#include "../common.h" +#include "../cudnn_utils.h" +#include "config_and_params.h" +#include "graph_cache_debug.h" + +namespace transformer_engine { +namespace fused_attn { + +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; +}; + +template +struct GraphCache { + std::mutex mutex; + std::map>> entries; +}; + +inline std::mutex &frontend_build_mutex() { + static std::mutex mutex; + return mutex; +} + +// Takes a constructed graph through the frontend calls that decide whether cuDNN can run it. +// `backend` and `pass` only name the build site the stage timers attribute the calls to. +// +// Reports by throwing, carrying cuDNN's message alone: that message is what support_verdict() +// returns as the reason a backend was refused. +inline void query_support(Backend backend, Pass pass, cudnn_frontend::graph::Graph &graph, + cudnnHandle_t handle) { + auto run = [&](graph_cache_debug::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 std::runtime_error(error.err_msg.empty() ? std::string(call_name) + " failed." + : error.err_msg); + }; + + run(graph_cache_debug::BuildStage::Validate, "validate", [&] { return graph.validate(); }); + run(graph_cache_debug::BuildStage::BuildOpGraph, "build_operation_graph", + [&] { return graph.build_operation_graph(handle); }); + run(graph_cache_debug::BuildStage::CreatePlans, "create_execution_plans", + [&] { return graph.create_execution_plans({cudnn_frontend::HeurMode_t::A}); }); + run(graph_cache_debug::BuildStage::CheckSupport, "check_support", + [&] { return graph.check_support(); }); +} + +// Cache for the entry `key`; build it first if absent. Record the lookup result and return the entry. +// hit -> record HIT, return the entry +// miss -> take frontend_build_mutex(), look again (a thread that raced us has finished by now), +// record MISS, build(), query_support(), insert +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; + } + + std::lock_guard build_lock(frontend_build_mutex()); + 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); + graph_cache_debug::record_cache_graph(backend, pass); + + std::lock_guard lock(cache.mutex); + return cache.entries.insert({key, std::move(entry)}).first->second; +} + +// Each backend's graph cache per forward/backward pass, called by support_verdict() and the +// execution path. The cache is this instantiation's static local, so the callers naming one +// triple share one cache and each triple gets its own. +template +auto get_graph(const FusedAttnConfig &cfg, cudnnHandle_t handle) { + static GraphCache cache; + cfg.check_derived(); + return cache_graph(cache, cfg.make_cache_key(kPass), kBackend, kPass, handle, + [&] { return kCreateGraphFn(cfg); }); +} + +// Check whether cuDNN can support a given config, per forward/backward pass. +// Returns an empty string if can; otherwise, a diagnostic string for the reason. +template +std::string support_verdict(const FusedAttnConfig &cfg, cudnnHandle_t handle) { + auto label = [] { + return std::string("support_verdict<") + graph_cache_debug::backend_name(kBackend) + ", " + + graph_cache_debug::pass_name(kPass) + ">"; + }; + try { + get_graph(cfg, handle); + return ""; + } catch (const std::exception &e) { + const char *reason = e.what(); + if (reason != nullptr && reason[0] != '\0') return reason; + return label() + ": rejected without a reason."; + } catch (...) { + return label() + ": unknown failure."; + } +} + +// Previous calls only create the graph, caches it if verified to be supported. This function +// compiles the kernels via graph.build_plans(). It is the most expensive frontend call, and +// done only once per cache entry. +template +void build_plans(Backend backend, Pass pass, CacheEntry &entry) { + std::call_once(entry.build_plans_once, [&] { + std::lock_guard build_lock(frontend_build_mutex()); + 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..7e5de19d58 --- /dev/null +++ b/transformer_engine/common/fused_attn/graph_cache_debug.h @@ -0,0 +1,528 @@ +/************************************************************************* + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * + * See LICENSE for license information. + ************************************************************************/ + +// Fused-attention graph cache diagnostics. +// +// Enable with NVTE_FUSED_ATTN_CACHE_DEBUG=[:]. The output format, how to read it and +// the rank suffix are documented for users in docs/envvars.rst; what follows is what maintaining +// this file needs. +// +// level 1 (events) : one line per event that happens once per distinct cache key (CREATE_GRAPH, +// CACHE_GRAPH, BUILD_PLANS), plus the exit summary and its stage timings. +// level 2 (trace) : adds a line per lookup (HIT/MISS, with the normalized key) and per execution +// (EXECUTE). High volume, and it serializes threads on the stderr lock, which +// the stage timings are then measured under -- no timed region writes to +// stderr, so they stay sound, but they read a little high. +// +// Counters are kept per build site -- f16/fp8 crossed with fwd/bwd -- since one process can drive +// both backends, and every event name is also the counter column it increments. Where the events +// sit on the path nvte_fused_attn_fwd_v2 sketches: HIT/MISS on every get_graph() lookup, +// CREATE_GRAPH and CACHE_GRAPH inside it on a miss, BUILD_PLANS on an entry's first execution, and +// EXECUTE on every call. +// +// One level-1 training step, line prefixes and trailing columns elided: +// +// tid=0 dev=0 | f16 fwd CREATE_GRAPH | hit=0, miss=1, create_graph=1, cache_graph=0, ... +// tid=0 dev=0 | f16 fwd CACHE_GRAPH | hit=0, miss=1, create_graph=1, cache_graph=1, ... +// ===== summary begin ===== +// tid=0 dev=0 | f16 fwd | hit=5, miss=1, create_graph=1, cache_graph=1, ... +// tid=1 dev=0 | f16 bwd | hit=4, build_plans=1, execute=1, ... +// tid=all dev=all | f16 fwd | hit=5, miss=1, create_graph=1, cache_graph=1, ... +// f16 fwd build_plans | calls=1 | time= 262.104 ms/call +// ===== summary end ===== +// +// Those first two lines are one graph before and after cuDNN was asked to support it, which is why +// a CREATE_GRAPH with no CACHE_GRAPH after it is a refusal. Rows for a site a thread never reached +// are left out rather than zeroed, hence tid=1 having a backward row and no forward one: a PyTorch +// step runs the forward and the backward's probe on the main thread and the backward itself on the +// autograd thread, which finds the graph that probe left behind. That split is why the build +// identities hold on the totals rows and not on any single thread's. + +#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 { + +// ============================================================================ +// Vocabulary: which build site an event came from, and which build stage or lookup outcome it +// reports. These four names and the recorders at the bottom are the whole interface. Backend and +// Pass are fused_attn's own, so a recorder and the key it prints share one notion of a site, and a +// mistake at a call site is a compile error rather than a mistyped "fwd". +// ============================================================================ + +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"; } + +enum class BuildStage { Validate, BuildOpGraph, CreatePlans, CheckSupport, BuildPlans, kCount }; +enum class LookupResult { Miss, Hit }; + +namespace detail { + +// Verbosity parsed once from NVTE_FUSED_ATTN_CACHE_DEBUG: 0=off, 1=events, 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); // stops at the optional ":" suffix + return v > 0 ? v : 1; // any non-empty, non-"0" value enables at least level 1 + }(); + return lvl; +} + +// Rank of this process as reported by the launcher, or -1 when there is none. First variable that +// is set wins. +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; +} + +// On at level >= 1, and only for the ranks the ":" suffix selects. Rank 0 only by default. +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; +} + +// The gate on the trace printouts. Tests enabled() rather than just the level. +inline bool enabled_with_trace() { return enabled() && debug_level() >= 2; } + +// Names 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, ...) in assignment order. tid=0 is whichever thread touched this cache first. +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; +} + +// Registered at first use. On process exit, prints event counters and build timings. Only used by the summary handler. +inline void register_summary_once(); + +// Backend major, pass minor, so that the two passes of one backend are adjacent. +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); +} + +// Cache event counters, one block per build site. Each name is both the event tag on the line that +// records it and the column carrying its running total: +// - create_graph: a graph constructed for a miss, counted before cuDNN is asked about it. +// - cache_graph: one of those graphs cleared check_support(), so this is the graphs cuDNN agreed +// to run. Counted on that verdict rather than on the insert that follows, so it says what cuDNN +// accepted and not how many entries the map holds. +// - build_plans: a cached graph finished with graph.build_plans(), the kernel compilation +// cache_graph deferred, paid by that graph's first execution rather than by the probe. +// - execute: an execution cuDNN accepted, counted once the enqueue returns. Not a completed +// execution: the work is asynchronous, so a device-side fault is not reflected here. +// - hit: a lookup answered from the cache. Need not lead to an execution -- it can be a backend +// availability check, or a workspace-sizing call, which has no tensors to run with. +// - miss: a lookup the cache did not answer; triggers a build. + +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; + } + + // Whether this block saw nothing at all. + 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. +struct ThreadCounters { + unsigned tid = 0; + std::atomic device{-1}; + std::array sites; +}; + +// 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(); + p->device.store(cuda::current_device(), std::memory_order_relaxed); + { + 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)]; +} + +// 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); +} + +// 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(); + thread_counters().device.store(device, std::memory_order_relaxed); + 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 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); + return true; +} + +// 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"; +} + +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]; +} + +// 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); + } +}; + +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); + 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; + } + } + } +} + +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; +} + +} // namespace detail + +// The recorders: everything a call site calls. Each takes the build site it is reporting for, adds +// one to that site's column, 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::enabled_with_trace()) { + 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)) || + !detail::enabled_with_trace()) { + 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..d1bbeae4ad 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,16 @@ 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 { + // Zeroed, for a FusedAttnConfig that has not been through derive() yet. Every multiplier is + // a per-token element count, so zero is not a usable value; it is only ever read after + // derive() has replaced it, which check_derived() is what enforces. + 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 +246,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 +308,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 41e4b136bd..691881dcb2 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" @@ -136,9 +136,10 @@ enum NVTE_Mask_Type { * \brief Attention softmax types as described in * Efficient Streaming Language Models with Attention Sinks (https://arxiv.org/pdf/2309.17453v3). * For a given attention score S = Q*K^T, different softmax types perform different operations on S, - * NVTE_VANILLA_SOFTMAX: S[:,:,:,i] = exp(S[:,:,:,i])/sum(exp(S[:,:,:,:]), dim=-1), - * NVTE_OFF_BY_ONE_SOFTMAX: S[:,:,:,i] = exp(S[:,:,:,i])/(1 + sum(exp(S[:,:,:,:]), dim=-1)), and - * NVTE_LEARNABLE_SOFTMAX: S[:,j,:,i] = exp(S[:,j,:,i])/(exp(alpha[j]) + sum(exp(S[:,j,:,:]), dim=-1)), + * `NVTE_VANILLA_SOFTMAX`: S[:,:,:,i] = exp(S[:,:,:,i])/sum(exp(S[:,:,:,:]), dim=-1), + * `NVTE_OFF_BY_ONE_SOFTMAX`: S[:,:,:,i] = exp(S[:,:,:,i])/(1 + sum(exp(S[:,:,:,:]), dim=-1)), and + * `NVTE_LEARNABLE_SOFTMAX`: S[:,j,:,i] = exp(S[:,j,:,i])/(exp(alpha[j]) + + * sum(exp(S[:,j,:,:]), dim=-1)), * where alpha is a learnable parameter of shape [H]. */ enum NVTE_Softmax_Type { @@ -194,6 +195,245 @@ 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 must match + * that of the declaration fields and `attr_sizes` array of `FusedAttnConfig`. New fields may + * only 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 must match + * that of the declaration fields and `attr_sizes` array of `FusedAttnFwdParams`. New fields may + * only 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 + kNVTEFusedAttnFwdParamsMaxSeqlenQ, + kNVTEFusedAttnFwdParamsMaxSeqlenKV, + kNVTEFusedAttnFwdParamsIsTraining, + kNVTEFusedAttnFwdParamsReturnMaxLogit, + kNVTEFusedAttnFwdParamsCudaGraph, + kNVTEFusedAttnFwdParamsAttnScale, + kNVTEFusedAttnFwdParamsDropout, + kNVTEFusedAttnFwdParamsQKVLayout, + kNVTEFusedAttnFwdParamsOFormat, + kNVTEFusedAttnFwdParamsQKVScaleInvFormat, + kNVTEFusedAttnFwdParamsBiasType, + kNVTEFusedAttnFwdParamsAttnMaskType, + kNVTEFusedAttnFwdParamsSoftmaxType, + kNVTEFusedAttnFwdParamsWindowSizeLeft, + kNVTEFusedAttnFwdParamsWindowSizeRight, + kNVTEFusedAttnFwdParamsBottomRightDiagonal, + // 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 must match + * that of the declaration fields and `attr_sizes` array of `FusedAttnBwdParams`. New fields may + * only 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 + kNVTEFusedAttnBwdParamsMaxSeqlenQ, + kNVTEFusedAttnBwdParamsMaxSeqlenKV, + kNVTEFusedAttnBwdParamsAttnScale, + kNVTEFusedAttnBwdParamsDropout, + kNVTEFusedAttnBwdParamsQKVLayout, + kNVTEFusedAttnBwdParamsOFormat, + kNVTEFusedAttnBwdParamsDOFormat, + kNVTEFusedAttnBwdParamsDQKVLayout, + kNVTEFusedAttnBwdParamsQKVScaleInvFormat, + kNVTEFusedAttnBwdParamsDOScaleInvFormat, + kNVTEFusedAttnBwdParamsBiasType, + kNVTEFusedAttnBwdParamsAttnMaskType, + kNVTEFusedAttnBwdParamsSoftmaxType, + kNVTEFusedAttnBwdParamsWindowSizeLeft, + kNVTEFusedAttnBwdParamsWindowSizeRight, + kNVTEFusedAttnBwdParamsBottomRightDiagonal, + kNVTEFusedAttnBwdParamsDeterministic, + kNVTEFusedAttnBwdParamsCudaGraph, + // 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 user configuration. + * + * This function passes the user configuration to cuDNN frontend, runs its support checks, + * and returns a backend if supported, or `NVTE_No_Backend` and a message explaining why. + * If the configuration is supported, the backend is cached and reused for future calls. + * + * \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 there is no support. Pass `NULL` to + * skip the diagnostics. Also, note that the string pointer refers to a + * per-thread buffer owned by the library and remains valid only until the + * next call to `nvte_get_fused_attn_backend_v2` on the same thread. + * Callers that need to retain the message across further calls must + * copy it. + * + * \return Fused-attention backend, `NVTE_F16_arbitrary_seqlen` or `NVTE_FP8`, + * if the given 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 +455,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 than + * `nvte_get_fused_attn_backend_v2`. It fills the fields that it cannot express with + * `nvte_get_fused_attn_backend_v2`'s default values. 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. Users who need more precise control should + * switch to `nvte_get_fused_attn_backend_v2`. */ 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 +473,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 Opaque 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 +533,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 +543,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 +561,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 Opaque 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 +616,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 +631,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, @@ -462,7 +741,7 @@ void nvte_cp_thd_read_half_tensor(const NVTETensor &tensor, const NVTETensor &cu * \param[out] lse Output tensor. * \param[in] lse_per_step Input tensor. * \param[in] cu_seqlens Cumulative sequence lengths, [batch_size + 1]. - * \param[in] lse_packed Whether or not lse_per_step is packed. + * \param[in] lse_packed Whether or not `lse_per_step` is packed. * \param[in] stream CUDA stream used for this operation. */ void nvte_cp_thd_second_half_lse_correction(NVTETensor lse, const NVTETensor &lse_per_step, @@ -511,7 +790,7 @@ void nvte_cp_thd_out_correction(NVTETensor out, const NVTETensor &out_per_step, * \param[in] cu_seqlens Cumulative sequence lengths, [batch_size + 1]. * \param[in] first_half One of ("add", "copy", "none") correction op for first half. * \param[in] second_half One of ("add", "copy", "none") correction op for second half. - Must be different from first_half. + * Must be different from `first_half`. * \param[in] stream CUDA stream used for this operation. */ void nvte_cp_thd_grad_correction(NVTETensor grad, const NVTETensor &grad_per_step, @@ -540,9 +819,9 @@ void nvte_cp_thd_get_partitioned_indices(const NVTETensor &cu_seqlens, NVTETenso * * \param[in] inp Input THD tensor [total_tokens, ...]. * \param[in] cu_seqlens Padded cumulative sequence lengths, [batch_size + 1], int32. - * \param[out] out Output tensor, same shape/dtype as inp. + * \param[out] out Output tensor, same shape/dtype as `inp`. * \param[in] world_size Context-parallel size. - * \param[in] total_tokens Total padded tokens (= inp.shape[0]). + * \param[in] total_tokens Total padded tokens (= `inp.shape[0]`). * \param[in] stream CUDA stream used for this operation. */ void nvte_thd_sequence_order_to_cp_rank_order(const NVTETensor &inp, const NVTETensor &cu_seqlens, @@ -556,9 +835,9 @@ void nvte_thd_sequence_order_to_cp_rank_order(const NVTETensor &inp, const NVTET * * \param[in] inp Input THD tensor [total_tokens, ...]. * \param[in] cu_seqlens Padded cumulative sequence lengths, [batch_size + 1], int32. - * \param[out] out Output tensor, same shape/dtype as inp. + * \param[out] out Output tensor, same shape/dtype as `inp`. * \param[in] world_size Context-parallel size. - * \param[in] total_tokens Total padded tokens (= inp.shape[0]). + * \param[in] total_tokens Total padded tokens (= `inp.shape[0]`). * \param[in] stream CUDA stream used for this operation. */ void nvte_thd_cp_rank_order_to_sequence_order(const NVTETensor &inp, const NVTETensor &cu_seqlens, @@ -573,8 +852,8 @@ void nvte_thd_cp_rank_order_to_sequence_order(const NVTETensor &inp, const NVTET * \param[in] inp Per-split THD source tensor [total_tokens, ...]. * \param[in] cu_seqlens_padded Padded cumulative sequence lengths, [batch_size + 1], int32. * \param[in] cu_seqlens Valid cumulative sequence lengths, [batch_size + 1], int32. - * \param[in,out] out Rank-local accumulator, same shape/dtype as inp. - * \param[in] total_tokens Total padded tokens (= inp.shape[0]). + * \param[in,out] out Rank-local accumulator, same shape/dtype as `inp`. + * \param[in] total_tokens Total padded tokens (= `inp.shape[0]`). * \param[in] stream CUDA stream used for this operation. */ void nvte_thd_copy_valid_tokens_from_per_split_to_rank_local(const NVTETensor &inp, @@ -641,7 +920,7 @@ void nvte_prepare_flash_attn_bwd(NVTETensor q, NVTETensor k, NVTETensor v, NVTET * \param[in] inputs List of input tensors. * \param[in,out] outputs List of output tensors. * \param[in] num_tensors Number of tensors in the list. - * \param[in] original_format Original QKV format (NVTE_BSHD or NVTE_SBHD). + * \param[in] original_format Original QKV format (`NVTE_BSHD` or `NVTE_SBHD`). * \param[in] stream CUDA stream. */ void nvte_multi_tensor_transpose_to_bhsd(NVTETensor *inputs, NVTETensor *outputs, @@ -708,6 +987,494 @@ 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: + // Common implementation for every setter: copy the value to a local variable, + // forward its address and size to the C API, and return *this for chaining. + 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_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_is_training(bool val) noexcept { + return set_attr(kNVTEFusedAttnFwdParamsIsTraining, static_cast(val)); + } + FusedAttnFwdParamsWrapper &set_return_max_logit(bool val) noexcept { + return set_attr(kNVTEFusedAttnFwdParamsReturnMaxLogit, static_cast(val)); + } + FusedAttnFwdParamsWrapper &set_cuda_graph(bool val) noexcept { + return set_attr(kNVTEFusedAttnFwdParamsCudaGraph, static_cast(val)); + } + FusedAttnFwdParamsWrapper &set_attn_scale(float val) noexcept { + return set_attr(kNVTEFusedAttnFwdParamsAttnScale, val); + } + FusedAttnFwdParamsWrapper &set_dropout(float val) noexcept { + return set_attr(kNVTEFusedAttnFwdParamsDropout, 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_bias_type(NVTE_Bias_Type val) noexcept { + return set_attr(kNVTEFusedAttnFwdParamsBiasType, val); + } + FusedAttnFwdParamsWrapper &set_attn_mask_type(NVTE_Mask_Type val) noexcept { + return set_attr(kNVTEFusedAttnFwdParamsAttnMaskType, val); + } + FusedAttnFwdParamsWrapper &set_softmax_type(NVTE_Softmax_Type val) noexcept { + return set_attr(kNVTEFusedAttnFwdParamsSoftmaxType, 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_workspace(NVTETensor val) noexcept { + return set_attr(kNVTEFusedAttnFwdParamsWorkspace, val); + } + FusedAttnFwdParamsWrapper &set_stream(cudaStream_t val) noexcept { + return set_attr(kNVTEFusedAttnFwdParamsStream, val); + } + + private: + // Common implementation for every setter: copy the value to a local variable, + // forward its address and size to the C API, and return *this for chaining. + 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_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_attn_scale(float val) noexcept { + return set_attr(kNVTEFusedAttnBwdParamsAttnScale, val); + } + FusedAttnBwdParamsWrapper &set_dropout(float val) noexcept { + return set_attr(kNVTEFusedAttnBwdParamsDropout, 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_bias_type(NVTE_Bias_Type val) noexcept { + return set_attr(kNVTEFusedAttnBwdParamsBiasType, val); + } + FusedAttnBwdParamsWrapper &set_attn_mask_type(NVTE_Mask_Type val) noexcept { + return set_attr(kNVTEFusedAttnBwdParamsAttnMaskType, val); + } + FusedAttnBwdParamsWrapper &set_softmax_type(NVTE_Softmax_Type val) noexcept { + return set_attr(kNVTEFusedAttnBwdParamsSoftmaxType, 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_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_workspace(NVTETensor val) noexcept { + return set_attr(kNVTEFusedAttnBwdParamsWorkspace, val); + } + FusedAttnBwdParamsWrapper &set_stream(cudaStream_t val) noexcept { + return set_attr(kNVTEFusedAttnBwdParamsStream, val); + } + + private: + // Common implementation for every setter: copy the value to a local variable, + // forward its address and size to the C API, and return *this for chaining. + 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 ecca4a3871..af1eda478d 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, @@ -339,15 +340,26 @@ def is_fused_attn_kernel_available( head_dim_qk, head_dim_v, window_size: Optional[Tuple[int, int]] = None, + 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, @@ -362,9 +374,15 @@ def make_helper(attn_mask_type): head_dim_qk, head_dim_v, window_size_tuple, + 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 489bfde997..a312d39b42 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=[], @@ -101,6 +142,63 @@ class _FusedAttnConfig: ) # Only for CP + Striped. For Ring P2P, stripe_size=1 only.For AG, stripe_size>=1. +@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, namely the paged-KV + dimensions and the ragged token counts, are omitted and keep 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 + + # 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: """ @@ -108,6 +206,7 @@ class FusedAttnHelper: """ is_training: bool + batch_size: int q_dtype: jnp.dtype kv_dtype: jnp.dtype qkv_layout: QKVLayout @@ -122,33 +221,84 @@ class FusedAttnHelper: head_dim_qk: int head_dim_v: int window_size: Tuple[int, int] + bottom_right_diagonal: bool + 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], - 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 + backend, message = transformer_engine_jax.get_fused_attn_backend( + FusedAttnParams( + is_training=self.is_training, + deterministic=not self.is_non_deterministic_allowed(), + 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, + 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""" @@ -335,8 +485,14 @@ 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) + backend, message = FusedAttnHelper( config.is_training, + input_batch, q_dtype, k_dtype, config.qkv_layout, @@ -351,6 +507,12 @@ def abstract( q_head_dim, v_head_dim, config.window_size, + 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: @@ -369,7 +531,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) # JAX does not enable 64-bit int by default so we get XLA to allocate x8 memory with diff --git a/transformer_engine/jax/csrc/extensions.h b/transformer_engine/jax/csrc/extensions.h index 9bd0940c4b..b7ba1c6af5 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,11 @@ 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 deterministic); +// Select the fused attention backend for the configuration carried by a FusedAttnParams object +// (see jax/cpp_extensions/attention.py). Returns the backend and, when no backend supports the +// configuration, a diagnostic message explaining why it was rejected. +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 3fd6780d6d..f24900410e 100644 --- a/transformer_engine/jax/csrc/extensions/attention.cpp +++ b/transformer_engine/jax/csrc/extensions/attention.cpp @@ -24,18 +24,95 @@ 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 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, - false, 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 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_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, message != nullptr ? std::string(message) : std::string()}; +} + +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("bias_batch_size").cast(), + params.attr("bias_num_heads").cast(), params.attr("bias_seqlen_q").cast(), + params.attr("bias_seqlen_kv").cast()); } /* @@ -195,15 +272,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, false, 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_max_seqlen_q(q_max_seqlen) + .set_max_seqlen_kv(kv_max_seqlen) + .set_is_training(is_training) + .set_return_max_logit(false) + .set_cuda_graph(false) + .set_attn_scale(scaling_factor) + .set_dropout(dropout_probability) + .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_bias_type(bias_type) + .set_attn_mask_type(mask_type) + .set_softmax_type(softmax_type) + .set_window_size_left(window_size_left) + .set_window_size_right(window_size_right) + .set_bottom_right_diagonal(bottom_right_diagonal) + .set_workspace(query_workspace_tensor.data()) + .set_stream(nullptr); + nvte_fused_attn_fwd_v2(params); } nvte_tensor_pack_destroy(&aux_output_tensors); @@ -274,11 +377,14 @@ 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, - false, false, deterministic); + auto [backend, _fwd_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, + input_batch, attn_heads, num_gqa_groups, qk_head_dim, v_head_dim, q_max_seqlen, kv_max_seqlen, + bias_batch, bias_heads, q_max_seqlen, kv_max_seqlen); 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) */ @@ -339,15 +445,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, false, 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_max_seqlen_q(q_max_seqlen) + .set_max_seqlen_kv(kv_max_seqlen) + .set_is_training(is_training) + .set_return_max_logit(false) + .set_cuda_graph(false) + .set_attn_scale(scaling_factor) + .set_dropout(dropout_probability) + .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_bias_type(bias_type) + .set_attn_mask_type(mask_type) + .set_softmax_type(softmax_type) + .set_window_size_left(window_size_left) + .set_window_size_right(window_size_right) + .set_bottom_right_diagonal(bottom_right_diagonal) + .set_workspace(workspace_tensor.data()) + .set_stream(stream); + nvte_fused_attn_fwd_v2(params); nvte_tensor_pack_destroy(&aux_output_tensors); } @@ -495,19 +627,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_max_seqlen_q(q_max_seqlen) + .set_max_seqlen_kv(kv_max_seqlen) + .set_attn_scale(scaling_factor) + .set_dropout(dropout_probability) + .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_bias_type(bias_type) + .set_attn_mask_type(mask_type) + .set_softmax_type(softmax_type) + .set_window_size_left(window_size_left) + .set_window_size_right(window_size_right) + .set_bottom_right_diagonal(bottom_right_diagonal) + .set_deterministic(deterministic) + .set_cuda_graph(false) + .set_workspace(query_workspace_tensor.data()) + .set_stream(nullptr); + nvte_fused_attn_bwd_v2(params); } nvte_tensor_pack_destroy(&aux_input_tensors); @@ -550,11 +708,14 @@ 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, + input_batch, attn_heads, num_gqa_groups, qk_head_dim, v_head_dim, q_max_seqlen, kv_max_seqlen, + bias_batch, bias_heads, q_max_seqlen, kv_max_seqlen); 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); @@ -631,18 +792,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_max_seqlen_q(q_max_seqlen) + .set_max_seqlen_kv(kv_max_seqlen) + .set_attn_scale(scaling_factor) + .set_dropout(dropout_probability) + .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_bias_type(bias_type) + .set_attn_mask_type(mask_type) + .set_softmax_type(softmax_type) + .set_window_size_left(window_size_left) + .set_window_size_right(window_size_right) + .set_bottom_right_diagonal(bottom_right_diagonal) + .set_deterministic(deterministic) + .set_cuda_graph(false) + .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 76922d2b55..9578219230 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 @@ -779,6 +782,8 @@ def __call__( enable_fused_attn = int(os.getenv("NVTE_FUSED_ATTN", "1")) 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 @@ -795,10 +800,15 @@ 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) + fused_attn_helper = FusedAttnHelper( # This needs to be fixed: TE-Jax has historically correlated training mode # with deterministic mode. not deterministic, + batch_size, input_dtype, # self._assert_dtypes enforces Q, K, V, bias to have the same dtype, so # using input_dtype as kv dtype is sufficient. @@ -814,8 +824,15 @@ 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, + 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." @@ -825,12 +842,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 8a219a6a4d..b0d24ac1bc 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/backends.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/backends.py @@ -2023,12 +2023,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 ea89ca97eb..cd9555da52 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py @@ -4921,6 +4921,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 d5adbbcadf..65f268815b 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 @@ -757,6 +757,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"))) @@ -1433,14 +1434,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 @@ -1913,11 +1913,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 q_format in ["sbhd", "bshd"]: max_seqlen_q *= cp_size @@ -1977,32 +1980,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 -- @@ -2069,22 +2051,31 @@ 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, attn_mask_type=attn_mask_type, window_size=window_size, bottom_right_diagonal=bottom_right_diagonal, - 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 ), pad_between_seqs=pad_between_seqs, 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, @@ -2094,6 +2085,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 ba049c9aef..a1774c9df2 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 @@ -233,8 +241,8 @@ class AttentionParams: Tensor shape of :attr:`alibi_slopes` in `DotProductAttention`. core_attention_bias_type : str, default = no_bias Attention bias type, {`no_bias`, `pre_scale_bias`, `post_scale_bias`, `alibi`}. - core_attention_bias_shape : 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,39 +433,25 @@ 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. - - Returns a plain int rather than a FusedAttnBackend member: dynamo - reconstructs the result of an assume_constant_result call by re-emitting the - call, which is only valid inside the frame that made it. An int survives a - graph break because it is baked into the graph as a literal, while an enum - member comes out of the reconstruction corrupted (see the cast at the call - 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, - ) + the attention config. + + The config comes in as keyword arguments rather than as a FusedAttentionParams: + torch.compile materializes the arguments of a constant-folded call, and a + dataclass built by traced code arrives here with its fields reset to defaults. + + The backend comes back as a plain int rather than a FusedAttnBackend member: + dynamo reconstructs the result of an assume_constant_result call by + re-emitting the call, which is only valid inside the frame that made it. An + int survives a graph break because it is baked into the graph as a literal, + while an enum member comes out of the reconstruction corrupted (see the cast + at the call site, which restores the enum). + """ + 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 +486,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 +508,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 +519,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,13 +1138,13 @@ 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 ) or (use_flash_attention_4 and FlashAttentionUtils.v4_is_installed): logger.debug( - "Disabling FlashAttention 2 and 4 for qkv_format = thd when there is " + "Disabling FlashAttention 2 and 4 when Q or KV uses THD and there is " "padding between sequences, i.e. [a, a, PAD, b, b, b, PAD, c, PAD]" ) use_flash_attention_2 = False @@ -1088,7 +1157,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 @@ -1434,65 +1503,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: + 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 ): logger.debug( "Disabling FusedAttention for score_mod because sub-backend %s is not " @@ -1543,7 +1733,7 @@ def _is_fa3_supported(num_heads, num_gqa_groups, head_dim_qk, head_dim_v, qkv_dt use_fused_attention = False fused_attention_backend = None if ( - fused_attention_backend == FusedAttnBackend.FP8.value + fused_attention_backend == FusedAttnBackend.FP8 and is_training and (device_compute_capability < (9, 0) or cudnn_version < (9, 19, 0)) ): @@ -1554,7 +1744,7 @@ def _is_fa3_supported(num_heads, num_gqa_groups, head_dim_qk, head_dim_v, qkv_dt use_fused_attention = False fused_attention_backend = None if ( - fused_attention_backend == FusedAttnBackend.F16_arbitrary_seqlen.value + fused_attention_backend == FusedAttnBackend.F16_arbitrary_seqlen and is_training and ( device_compute_capability < (9, 0) @@ -2404,6 +2594,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/cpp_extensions/fused_attn.py b/transformer_engine/pytorch/cpp_extensions/fused_attn.py index 9a33df7634..04f005522a 100644 --- a/transformer_engine/pytorch/cpp_extensions/fused_attn.py +++ b/transformer_engine/pytorch/cpp_extensions/fused_attn.py @@ -109,7 +109,11 @@ class FusedAttnBackend(IntEnum): share the same integer value. Unlike the pybind enum, a plain-python ``IntEnum`` is traceable by ``torch.compile``: comparisons against a member constant-fold cleanly. Lookup by name (``FusedAttnBackend["FP8"]``) works - the same way as with the dict this used to be. + the same way as with the dict this used to be. Adding ``__eq__``/``__ne__`` + overrides is unnecessary (the inherited ``int`` comparisons already match + the pybind enum in both operand orders) and harmful: a python ``__eq__`` + would push dynamo from constant-folding the comparison to + inline-with-guard, breaking tracing. Members do not survive a graph break, though, so ``get_attention_backend`` returns the sub-backend as a plain int and ``cast`` turns it back into a @@ -134,32 +138,15 @@ def cast( return backend return cls(int(backend)) - def __eq__(self, other: object) -> bool: - # ``FusedAttnBackend`` is an ``IntEnum`` while ``NVTE_Fused_Attn_Backend`` - # is a pybind11 enum. Compare by integer value so the two enums stay - # equivalent regardless of the pybind11 version (the pybind ``__eq__`` - # handles the reverse order). - if isinstance(other, NVTE_Fused_Attn_Backend): - return int(self) == int(other) - return int.__eq__(self, other) - - def __ne__(self, other: object) -> bool: - result = self.__eq__(other) - if result is NotImplemented: - return result - return not result - - def __hash__(self) -> int: - return int.__hash__(self) - # Fail fast at import time if a new enumerator is added on the C++ side -# without being mirrored above. -assert {f"NVTE_{m.name}" for m in FusedAttnBackend} == set(NVTE_Fused_Attn_Backend.__members__), ( - "FusedAttnBackend in python is out of sync with" - " transformer_engine_torch.NVTE_Fused_Attn_Backend defined on the C++ side." - " Please make sure TE C++ and python are in sync." -) +# without being mirrored above. Not an assert, so that the check survives -O. +if {f"NVTE_{m.name}" for m in FusedAttnBackend} != set(NVTE_Fused_Attn_Backend.__members__): + raise RuntimeError( + "FusedAttnBackend in python is out of sync with" + " transformer_engine_torch.NVTE_Fused_Attn_Backend defined on the C++ side." + " Please make sure TE C++ and python are in sync." + ) BACKEND_FP8_THREADS_PER_CTA = 128 BACKEND_F16arb_ELTS_PER_THREADS = 16 diff --git a/transformer_engine/pytorch/csrc/extensions.h b/transformer_engine/pytorch/csrc/extensions.h index 4a7971c7dc..c8a3c7642f 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 eb8813d4a0..767c93140f 100644 --- a/transformer_engine/pytorch/csrc/extensions/attention.cpp +++ b/transformer_engine/pytorch/csrc/extensions/attention.cpp @@ -40,18 +40,55 @@ void mha_fill(const transformer_engine::TensorWrapper &self, const at::Tensor &s 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, message != nullptr ? std::string(message) : std::string()}; } // helper function for S and dP quantizers @@ -243,22 +280,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_max_seqlen_q(max_seqlen_q) + .set_max_seqlen_kv(max_seqlen_kv) + .set_is_training(is_training) + .set_return_max_logit(return_max_logit) + .set_cuda_graph(cuda_graph) + .set_attn_scale(attn_scale) + .set_dropout(p_dropout) + .set_qkv_layout(qkv_layout) + .set_o_format(o_format) + .set_qkv_scale_inv_format(qkv_scale_inv_format) + .set_bias_type(bias_type) + .set_attn_mask_type(attn_mask_type) + .set_softmax_type(softmax_type) + .set_window_size_left(window_size[0]) + .set_window_size_right(window_size[1]) + .set_bottom_right_diagonal(bottom_right_diagonal) + .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; @@ -301,16 +366,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); @@ -571,36 +627,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_max_seqlen_q(max_seqlen_q) + .set_max_seqlen_kv(max_seqlen_kv) + .set_attn_scale(attn_scale) + .set_dropout(p_dropout) + .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_bias_type(bias_type) + .set_attn_mask_type(attn_mask_type) + .set_softmax_type(softmax_type) + .set_window_size_left(window_size[0]) + .set_window_size_right(window_size[1]) + .set_bottom_right_diagonal(bottom_right_diagonal) + .set_deterministic(deterministic) + .set_cuda_graph(cuda_graph) + .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 7259cfc6d9..f6cdaf5ffc 100644 --- a/transformer_engine/pytorch/csrc/extensions/pybind.cpp +++ b/transformer_engine/pytorch/csrc/extensions/pybind.cpp @@ -447,7 +447,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());