diff --git a/docs/envvars.rst b/docs/envvars.rst index b3765a06bd..d6ffa241ef 100644 --- a/docs/envvars.rst +++ b/docs/envvars.rst @@ -171,18 +171,18 @@ backend-selection overview. :Default: ``1`` :Description: Enable or disable UnfusedDotProductAttention backend (native PyTorch). When set to ``0``, UnfusedDotProductAttention will not be used. -.. envvar:: NVTE_FUSED_ATTN_BACKEND - - :Type: ``int`` (1 or 2) - :Default: Auto-selected - :Description: Request a cuDNN FusedAttention backend when that request is supported by the active fused-attention path. ``1`` = F16_arbitrary_seqlen (cuDNN, any seq len), ``2`` = FP8 backend. If not set, the backend is automatically selected based on the input configuration. BF16/FP16 attention uses sub-backend ``1`` when eligible. FP8 attention uses sub-backend ``2`` when FP8 DPA is enabled and supported by the architecture, cuDNN version, and input configuration. - .. envvar:: NVTE_FUSED_ATTN_USE_FAv2_BWD :Type: ``int`` (0 or 1) :Default: ``0`` :Description: When using FusedAttention, use FlashAttention-2 implementation for the backward pass instead of the cuDNN implementation. This can be useful due to performance differences between various versions of flash-attn and FusedAttention. +.. envvar:: NVTE_FUSED_ATTN_CACHE_DEBUG + + :Type: ``int`` (0 or 1) + :Default: ``0`` + :Description: Enable diagnostic logging for the FusedAttention graph cache (covers both the F16 and FP8 kernels, forward and backward). When set to ``1``, prints to stderr (prefixed ``[FUSED-ATTN-CACHE]``) a per-lookup ``HIT``/``MISS`` line with the full graph-cache key, a ``BUILD`` line whenever a new graph is constructed, an ``EXEC`` line whenever a graph is executed, a ``SUMMARY`` of graph builds vs. executions at process exit, and a breakdown of cuDNN graph-build timings. Useful for diagnosing redundant graph rebuilds or stale-cache reuse, and for profiling graph-build cost. 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..4ffa804401 100644 --- a/docs/examples/attention/attention.ipynb +++ b/docs/examples/attention/attention.ipynb @@ -346,17 +346,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/tests/jax/test_distributed_fused_attn.py b/tests/jax/test_distributed_fused_attn.py index 2abd9824b6..6657962e93 100644 --- a/tests/jax/test_distributed_fused_attn.py +++ b/tests/jax/test_distributed_fused_attn.py @@ -81,25 +81,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, @@ -233,25 +214,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, @@ -425,6 +387,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 0d1db0b9e1..30f29c12d4 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(rewang): probably adds this in 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.") @@ -520,8 +544,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, @@ -536,9 +573,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/test_attention.py b/tests/pytorch/attention/test_attention.py index 7bcdd10daf..6c876ad7cc 100644 --- a/tests/pytorch/attention/test_attention.py +++ b/tests/pytorch/attention/test_attention.py @@ -1,6 +1,7 @@ # Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. +import copy import logging import os import sys @@ -140,7 +141,7 @@ def test_dot_product_attention( tols = dict(atol=1e-3, rtol=1e-3) if dtype == torch.bfloat16: tols = dict(atol=1.5e-2, rtol=1.5e-2) - config = model_configs[model] + config = copy.deepcopy(model_configs[model]) is_mla = config.head_dim_qk != config.head_dim_v is_mqa_gqa = config.num_heads != config.num_gqa_groups if qkv_layout is None: @@ -555,6 +556,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) @@ -825,6 +829,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) @@ -1867,12 +1874,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: @@ -2068,6 +2087,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] @@ -2098,6 +2120,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, @@ -2316,6 +2339,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 @@ -2355,6 +2382,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, @@ -2374,8 +2402,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" @@ -2664,10 +2690,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, ) @@ -2745,6 +2783,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 681ee5e6e0..9ffdb865fa 100644 --- a/tests/pytorch/attention/test_attention_with_cp.py +++ b/tests/pytorch/attention/test_attention_with_cp.py @@ -378,6 +378,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: @@ -638,6 +640,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 cdd98d2445..ad57d584f3 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 @@ -472,8 +473,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/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 13880e0d61..742d37a2f8 100644 --- a/transformer_engine/common/CMakeLists.txt +++ b/transformer_engine/common/CMakeLists.txt @@ -181,6 +181,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..ca4214dac3 --- /dev/null +++ b/transformer_engine/common/fused_attn/config_and_params.cpp @@ -0,0 +1,1160 @@ +/************************************************************************* + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * + * See LICENSE for license information. + ************************************************************************/ + +#include "config_and_params.h" + +#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); + +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 + 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); + + // 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 of cu_seqlens vs actual_seqlens + const size_t cudnn_runtime_version = cudnnGetVersion(); + const bool is_dropout = is_training && dropout != 0.0f; + uses_cu_seqlens_directly = CUDNN_FRONTEND_VERSION >= 12500 && + (CUDNN_VERSION >= 92400 && cudnn_runtime_version >= 92400) && + !is_dropout; + + // 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; + } + } +} + +FusedAttnConfig FusedAttnConfig::make_cache_key() const { + 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; + } + + // Bucket THD (ragged) batch and token counts + if (cache_cfg.is_ragged_q || cache_cfg.is_ragged_kv) { + const auto cudnn_runtime_version = cudnnGetVersion(); + const int sm_arch_ = cuda::sm_arch(cuda::current_device()); + if (cudnn_runtime_version >= 90600 && sm_arch_ != 120) { + if (cache_cfg.is_ragged_q) { + cache_cfg.max_seqlen_q = cache_cfg.bucketed_num_tokens_q; + } + if (cache_cfg.is_ragged_kv) { + cache_cfg.max_seqlen_kv = cache_cfg.bucketed_num_tokens_kv; + } + cache_cfg.num_tokens_q = 0; + cache_cfg.num_tokens_kv = 0; + const bool bucket_batch = !is_forward || !cache_cfg.uses_cu_seqlens_directly; + if (bucket_batch) { + cache_cfg.batch_size = cache_cfg.bucketed_batch_size; + } + } + } + + // attn_scale is a pass-by-value graph input and different scales can share the same cached graph + cache_cfg.attn_scale = 1.0f; + + // Restrict each direction's key to the fields its graph actually consumes, so + // no redundant graphs are built and no cache misses either + if (is_forward) { + 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; +} + +FusedAttnConfig FusedAttnFwdParams::make_config() const { + const FusedAttnFwdParams ¶ms = *this; + FusedAttnConfig cfg{}; + cfg.is_forward = true; + 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{}; + 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 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 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 kNVTEFusedAttnFwdParamsIsTraining: + bool_to_uint8(p.is_training, buf); + break; + case kNVTEFusedAttnFwdParamsCudaGraph: + bool_to_uint8(p.cuda_graph, buf); + break; + case kNVTEFusedAttnFwdParamsReturnMaxLogit: + bool_to_uint8(p.return_max_logit, buf); + break; + case kNVTEFusedAttnFwdParamsAttnMaskType: + std::memcpy(buf, &p.attn_mask_type, attr_size); + break; + case kNVTEFusedAttnFwdParamsBiasType: + std::memcpy(buf, &p.bias_type, attr_size); + break; + case 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 kNVTEFusedAttnFwdParamsDropout: + std::memcpy(buf, &p.dropout, attr_size); + break; + case kNVTEFusedAttnFwdParamsAttnScale: + std::memcpy(buf, &p.attn_scale, attr_size); + break; + case kNVTEFusedAttnFwdParamsQKVLayout: + std::memcpy(buf, &p.qkv_layout, attr_size); + break; + case kNVTEFusedAttnFwdParamsOFormat: + std::memcpy(buf, &p.o_format, attr_size); + break; + case kNVTEFusedAttnFwdParamsQKVScaleInvFormat: + std::memcpy(buf, &p.qkv_scale_inv_format, attr_size); + break; + case kNVTEFusedAttnFwdParamsMaxSeqlenQ: + std::memcpy(buf, &p.max_seqlen_q, attr_size); + break; + case kNVTEFusedAttnFwdParamsMaxSeqlenKV: + std::memcpy(buf, &p.max_seqlen_kv, attr_size); + break; + case kNVTEFusedAttnFwdParamsWorkspace: + std::memcpy(buf, &p.workspace, attr_size); + break; + case kNVTEFusedAttnFwdParamsStream: + std::memcpy(buf, &p.stream, attr_size); + break; + default: + NVTE_ERROR("Unsupported NVTEFusedAttnFwdParamsAttribute (got ", static_cast(attr), ")"); + } +} + +void nvte_set_fused_attn_fwd_params_attribute(NVTEFusedAttnFwdParams params, + NVTEFusedAttnFwdParamsAttribute attr, const void *buf, + size_t size_in_bytes) { + using namespace transformer_engine; + using namespace transformer_engine::fused_attn; + NVTE_CHECK(attr < kNVTEFusedAttnFwdParamsNumAttributes, + "Invalid NVTEFusedAttnFwdParamsAttribute (got ", static_cast(attr), ")"); + const auto &attr_size = FusedAttnFwdParams::attr_sizes[attr]; + NVTE_CHECK(buf != nullptr, "Input buffer must not be NULL."); + NVTE_CHECK(size_in_bytes >= attr_size, "Buffer is too small for attribute (need ", attr_size, + ", got ", size_in_bytes, ")"); + auto &p = *get_fused_attn_fwd_params_mutable(params); + switch (attr) { + case kNVTEFusedAttnFwdParamsQ: + std::memcpy(&p.Q, buf, attr_size); + break; + case kNVTEFusedAttnFwdParamsK: + std::memcpy(&p.K, buf, attr_size); + break; + case kNVTEFusedAttnFwdParamsV: + std::memcpy(&p.V, buf, attr_size); + break; + case kNVTEFusedAttnFwdParamsBias: + std::memcpy(&p.Bias, buf, attr_size); + break; + case kNVTEFusedAttnFwdParamsSoftmaxOffset: + std::memcpy(&p.SoftmaxOffset, buf, attr_size); + break; + case 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 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 kNVTEFusedAttnFwdParamsIsTraining: + uint8_to_bool(buf, p.is_training); + break; + case kNVTEFusedAttnFwdParamsCudaGraph: + uint8_to_bool(buf, p.cuda_graph); + break; + case kNVTEFusedAttnFwdParamsReturnMaxLogit: + uint8_to_bool(buf, p.return_max_logit); + break; + case kNVTEFusedAttnFwdParamsAttnMaskType: + std::memcpy(&p.attn_mask_type, buf, attr_size); + break; + case kNVTEFusedAttnFwdParamsBiasType: + std::memcpy(&p.bias_type, buf, attr_size); + break; + case 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 kNVTEFusedAttnFwdParamsDropout: + std::memcpy(&p.dropout, buf, attr_size); + break; + case kNVTEFusedAttnFwdParamsAttnScale: + std::memcpy(&p.attn_scale, buf, attr_size); + break; + case kNVTEFusedAttnFwdParamsQKVLayout: + std::memcpy(&p.qkv_layout, buf, attr_size); + break; + case kNVTEFusedAttnFwdParamsOFormat: + std::memcpy(&p.o_format, buf, attr_size); + break; + case kNVTEFusedAttnFwdParamsQKVScaleInvFormat: + std::memcpy(&p.qkv_scale_inv_format, buf, attr_size); + break; + case kNVTEFusedAttnFwdParamsMaxSeqlenQ: + std::memcpy(&p.max_seqlen_q, buf, attr_size); + break; + case kNVTEFusedAttnFwdParamsMaxSeqlenKV: + std::memcpy(&p.max_seqlen_kv, buf, attr_size); + break; + case kNVTEFusedAttnFwdParamsWorkspace: + std::memcpy(&p.workspace, buf, attr_size); + break; + case kNVTEFusedAttnFwdParamsStream: + std::memcpy(&p.stream, buf, attr_size); + break; + default: + NVTE_ERROR("Unsupported NVTEFusedAttnFwdParamsAttribute (got ", static_cast(attr), ")"); + } +} + +NVTEFusedAttnBwdParams nvte_create_fused_attn_bwd_params() { + return new transformer_engine::fused_attn::FusedAttnBwdParams{}; +} + +void nvte_destroy_fused_attn_bwd_params(NVTEFusedAttnBwdParams params) { + delete transformer_engine::fused_attn::get_fused_attn_bwd_params_mutable(params); +} + +void nvte_get_fused_attn_bwd_params_attribute(NVTEFusedAttnBwdParams params, + NVTEFusedAttnBwdParamsAttribute attr, void *buf, + size_t size_in_bytes, size_t *size_written) { + using namespace transformer_engine; + using namespace transformer_engine::fused_attn; + NVTE_CHECK(attr < kNVTEFusedAttnBwdParamsNumAttributes, + "Invalid NVTEFusedAttnBwdParamsAttribute (got ", static_cast(attr), ")"); + const auto &attr_size = FusedAttnBwdParams::attr_sizes[attr]; + if (size_written != nullptr) { + *size_written = attr_size; + } + if (buf == nullptr) { + return; + } + NVTE_CHECK(size_in_bytes >= attr_size, "Buffer is too small for attribute (need ", attr_size, + ", got ", size_in_bytes, ")"); + const auto &p = *get_fused_attn_bwd_params(params); + switch (attr) { + case kNVTEFusedAttnBwdParamsQ: + std::memcpy(buf, &p.Q, attr_size); + break; + case kNVTEFusedAttnBwdParamsK: + std::memcpy(buf, &p.K, attr_size); + break; + case kNVTEFusedAttnBwdParamsV: + std::memcpy(buf, &p.V, attr_size); + break; + case kNVTEFusedAttnBwdParamsO: + std::memcpy(buf, &p.O, attr_size); + break; + case kNVTEFusedAttnBwdParamsDO: + std::memcpy(buf, &p.dO, attr_size); + break; + case kNVTEFusedAttnBwdParamsS: + std::memcpy(buf, &p.S, attr_size); + break; + case kNVTEFusedAttnBwdParamsDP: + std::memcpy(buf, &p.dP, attr_size); + break; + case kNVTEFusedAttnBwdParamsAuxCtxTensors: + std::memcpy(buf, &p.Aux_CTX_Tensors, attr_size); + break; + case kNVTEFusedAttnBwdParamsDQ: + std::memcpy(buf, &p.dQ, attr_size); + break; + case kNVTEFusedAttnBwdParamsDK: + std::memcpy(buf, &p.dK, attr_size); + break; + case kNVTEFusedAttnBwdParamsDV: + std::memcpy(buf, &p.dV, attr_size); + break; + case kNVTEFusedAttnBwdParamsDBias: + std::memcpy(buf, &p.dBias, attr_size); + break; + case kNVTEFusedAttnBwdParamsDSoftmaxOffset: + std::memcpy(buf, &p.dSoftmaxOffset, attr_size); + break; + case kNVTEFusedAttnBwdParamsCuSeqlensQ: + std::memcpy(buf, &p.cu_seqlens_q, attr_size); + break; + case kNVTEFusedAttnBwdParamsCuSeqlensKV: + std::memcpy(buf, &p.cu_seqlens_kv, attr_size); + break; + case kNVTEFusedAttnBwdParamsCuSeqlensQPadded: + std::memcpy(buf, &p.cu_seqlens_q_padded, attr_size); + break; + case kNVTEFusedAttnBwdParamsCuSeqlensKVPadded: + std::memcpy(buf, &p.cu_seqlens_kv_padded, attr_size); + break; + case kNVTEFusedAttnBwdParamsCudaGraph: + bool_to_uint8(p.cuda_graph, buf); + break; + case kNVTEFusedAttnBwdParamsDeterministic: + bool_to_uint8(p.deterministic, buf); + break; + case kNVTEFusedAttnBwdParamsAttnMaskType: + std::memcpy(buf, &p.attn_mask_type, attr_size); + break; + case kNVTEFusedAttnBwdParamsBiasType: + std::memcpy(buf, &p.bias_type, attr_size); + break; + case 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 kNVTEFusedAttnBwdParamsDropout: + std::memcpy(buf, &p.dropout, attr_size); + break; + case kNVTEFusedAttnBwdParamsAttnScale: + std::memcpy(buf, &p.attn_scale, attr_size); + break; + case kNVTEFusedAttnBwdParamsQKVLayout: + std::memcpy(buf, &p.qkv_layout, attr_size); + break; + case kNVTEFusedAttnBwdParamsOFormat: + std::memcpy(buf, &p.o_format, attr_size); + break; + case kNVTEFusedAttnBwdParamsDOFormat: + std::memcpy(buf, &p.do_format, attr_size); + break; + case kNVTEFusedAttnBwdParamsDQKVLayout: + std::memcpy(buf, &p.dqkv_layout, attr_size); + break; + case kNVTEFusedAttnBwdParamsQKVScaleInvFormat: + std::memcpy(buf, &p.qkv_scale_inv_format, attr_size); + break; + case kNVTEFusedAttnBwdParamsDOScaleInvFormat: + std::memcpy(buf, &p.do_scale_inv_format, attr_size); + break; + case kNVTEFusedAttnBwdParamsMaxSeqlenQ: + std::memcpy(buf, &p.max_seqlen_q, attr_size); + break; + case kNVTEFusedAttnBwdParamsMaxSeqlenKV: + std::memcpy(buf, &p.max_seqlen_kv, attr_size); + break; + case kNVTEFusedAttnBwdParamsWorkspace: + std::memcpy(buf, &p.workspace, attr_size); + break; + case kNVTEFusedAttnBwdParamsStream: + std::memcpy(buf, &p.stream, attr_size); + break; + default: + NVTE_ERROR("Unsupported NVTEFusedAttnBwdParamsAttribute (got ", static_cast(attr), ")"); + } +} + +void nvte_set_fused_attn_bwd_params_attribute(NVTEFusedAttnBwdParams params, + NVTEFusedAttnBwdParamsAttribute attr, const void *buf, + size_t size_in_bytes) { + using namespace transformer_engine; + using namespace transformer_engine::fused_attn; + NVTE_CHECK(attr < kNVTEFusedAttnBwdParamsNumAttributes, + "Invalid NVTEFusedAttnBwdParamsAttribute (got ", static_cast(attr), ")"); + const auto &attr_size = FusedAttnBwdParams::attr_sizes[attr]; + NVTE_CHECK(buf != nullptr, "Input buffer must not be NULL."); + NVTE_CHECK(size_in_bytes >= attr_size, "Buffer is too small for attribute (need ", attr_size, + ", got ", size_in_bytes, ")"); + auto &p = *get_fused_attn_bwd_params_mutable(params); + switch (attr) { + case kNVTEFusedAttnBwdParamsQ: + std::memcpy(&p.Q, buf, attr_size); + break; + case kNVTEFusedAttnBwdParamsK: + std::memcpy(&p.K, buf, attr_size); + break; + case kNVTEFusedAttnBwdParamsV: + std::memcpy(&p.V, buf, attr_size); + break; + case kNVTEFusedAttnBwdParamsO: + std::memcpy(&p.O, buf, attr_size); + break; + case kNVTEFusedAttnBwdParamsDO: + std::memcpy(&p.dO, buf, attr_size); + break; + case kNVTEFusedAttnBwdParamsS: + std::memcpy(&p.S, buf, attr_size); + break; + case kNVTEFusedAttnBwdParamsDP: + std::memcpy(&p.dP, buf, attr_size); + break; + case kNVTEFusedAttnBwdParamsAuxCtxTensors: + std::memcpy(&p.Aux_CTX_Tensors, buf, attr_size); + break; + case kNVTEFusedAttnBwdParamsDQ: + std::memcpy(&p.dQ, buf, attr_size); + break; + case kNVTEFusedAttnBwdParamsDK: + std::memcpy(&p.dK, buf, attr_size); + break; + case kNVTEFusedAttnBwdParamsDV: + std::memcpy(&p.dV, buf, attr_size); + break; + case kNVTEFusedAttnBwdParamsDBias: + std::memcpy(&p.dBias, buf, attr_size); + break; + case kNVTEFusedAttnBwdParamsDSoftmaxOffset: + std::memcpy(&p.dSoftmaxOffset, buf, attr_size); + break; + case kNVTEFusedAttnBwdParamsCuSeqlensQ: + std::memcpy(&p.cu_seqlens_q, buf, attr_size); + break; + case kNVTEFusedAttnBwdParamsCuSeqlensKV: + std::memcpy(&p.cu_seqlens_kv, buf, attr_size); + break; + case kNVTEFusedAttnBwdParamsCuSeqlensQPadded: + std::memcpy(&p.cu_seqlens_q_padded, buf, attr_size); + break; + case kNVTEFusedAttnBwdParamsCuSeqlensKVPadded: + std::memcpy(&p.cu_seqlens_kv_padded, buf, attr_size); + break; + case kNVTEFusedAttnBwdParamsCudaGraph: + uint8_to_bool(buf, p.cuda_graph); + break; + case kNVTEFusedAttnBwdParamsDeterministic: + uint8_to_bool(buf, p.deterministic); + break; + case kNVTEFusedAttnBwdParamsAttnMaskType: + std::memcpy(&p.attn_mask_type, buf, attr_size); + break; + case kNVTEFusedAttnBwdParamsBiasType: + std::memcpy(&p.bias_type, buf, attr_size); + break; + case 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 kNVTEFusedAttnBwdParamsDropout: + std::memcpy(&p.dropout, buf, attr_size); + break; + case kNVTEFusedAttnBwdParamsAttnScale: + std::memcpy(&p.attn_scale, buf, attr_size); + break; + case kNVTEFusedAttnBwdParamsQKVLayout: + std::memcpy(&p.qkv_layout, buf, attr_size); + break; + case kNVTEFusedAttnBwdParamsOFormat: + std::memcpy(&p.o_format, buf, attr_size); + break; + case kNVTEFusedAttnBwdParamsDOFormat: + std::memcpy(&p.do_format, buf, attr_size); + break; + case kNVTEFusedAttnBwdParamsDQKVLayout: + std::memcpy(&p.dqkv_layout, buf, attr_size); + break; + case kNVTEFusedAttnBwdParamsQKVScaleInvFormat: + std::memcpy(&p.qkv_scale_inv_format, buf, attr_size); + break; + case kNVTEFusedAttnBwdParamsDOScaleInvFormat: + std::memcpy(&p.do_scale_inv_format, buf, attr_size); + break; + case kNVTEFusedAttnBwdParamsMaxSeqlenQ: + std::memcpy(&p.max_seqlen_q, buf, attr_size); + break; + case kNVTEFusedAttnBwdParamsMaxSeqlenKV: + std::memcpy(&p.max_seqlen_kv, buf, attr_size); + break; + case kNVTEFusedAttnBwdParamsWorkspace: + std::memcpy(&p.workspace, buf, attr_size); + break; + case kNVTEFusedAttnBwdParamsStream: + std::memcpy(&p.stream, buf, attr_size); + break; + default: + NVTE_ERROR("Unsupported NVTEFusedAttnBwdParamsAttribute (got ", static_cast(attr), ")"); + } +} diff --git a/transformer_engine/common/fused_attn/config_and_params.h b/transformer_engine/common/fused_attn/config_and_params.h new file mode 100644 index 0000000000..ebc5b3eb07 --- /dev/null +++ b/transformer_engine/common/fused_attn/config_and_params.h @@ -0,0 +1,385 @@ +/************************************************************************* + * 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 "common/common.h" +#include "transformer_engine/fused_attn.h" + +namespace transformer_engine { +namespace fused_attn { + +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; + + // device ID: not part of attribute serialization, but part of operator< and used to + // differentiate graphs built for different devices in multi-GPU single-process runs + int device_id = -1; + + // Internal-only fields: never part of attribute serialization, operator<, or the graph cache key. + // Filled by derive() or set by caller (i.e. is_forward). Added for convinence purposes and do not + // represent any graph properties. + + // Direction to build the cuDNN graph for; steers make_cache_key() normalization. + bool is_forward = false; + // THD batch/token counts; make_cache_key() folds these into batch_size/max_seqlen_*. + size_t bucketed_batch_size = 0; + size_t bucketed_num_tokens_q = 0; + size_t bucketed_num_tokens_kv = 0; + // Uses cu_seqlens or actual_seqlens. + bool uses_cu_seqlens_directly = false; + // Convinence fields to avoid recompute. + 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; + + 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 + }; + + 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 fields such as bucketed batch_size or num_tokens for THD, based on input fields + // that have been set by the caller. + void derive(); + + // Return a normalized copy of this config to be used as a key for the cuDNN graph cache. + // It drops fields that are invariant (e.g. attn_scale) or irrelevant (e.g. dO/dQKV dtypes + // and `deterministic` for forward, and `return_max_logit` for backward) to the corresponding graph. + // This helps avoid redundant graph builds and cache misses. + FusedAttnConfig make_cache_key() 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 { + NVTETensor Q = nullptr; + NVTETensor K = nullptr; + NVTETensor V = nullptr; + NVTETensor Bias = nullptr; + NVTETensor SoftmaxOffset = nullptr; + 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; + NVTETensor S = nullptr; + NVTETensor O = nullptr; + NVTETensorPack *Aux_CTX_Tensors = nullptr; + bool is_training = true; + bool cuda_graph = false; + bool return_max_logit = false; + NVTE_Mask_Type attn_mask_type = NVTE_NO_MASK; + NVTE_Bias_Type bias_type = NVTE_NO_BIAS; + 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; + float dropout = 0.0f; + float attn_scale = 1.0f; + NVTE_QKV_Layout qkv_layout = NVTE_QKV_Layout_NOT_SET; + NVTE_QKV_Format o_format = NVTE_QKV_Format_NOT_SET; + NVTE_QKV_Format qkv_scale_inv_format = NVTE_QKV_Format_NOT_SET; + size_t max_seqlen_q = 0; + size_t max_seqlen_kv = 0; + 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), // 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(NVTETensor), // S + sizeof(NVTETensor), // O + sizeof(NVTETensorPack *), // Aux_CTX_Tensors + sizeof(uint8_t), // is_training + sizeof(uint8_t), // cuda_graph + sizeof(uint8_t), // return_max_logit + sizeof(NVTE_Mask_Type), // attn_mask_type + sizeof(NVTE_Bias_Type), // bias_type + sizeof(NVTE_Softmax_Type), // softmax_type + sizeof(int64_t), // window_size_left + sizeof(int64_t), // window_size_right + sizeof(uint8_t), // bottom_right_diagonal + sizeof(float), // dropout + sizeof(float), // attn_scale + sizeof(NVTE_QKV_Layout), // qkv_layout + sizeof(NVTE_QKV_Format), // o_format + sizeof(NVTE_QKV_Format), // qkv_scale_inv_format + sizeof(size_t), // max_seqlen_q + sizeof(size_t), // max_seqlen_kv + sizeof(NVTETensor), // workspace + sizeof(cudaStream_t), // stream + }; + + // 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 { + 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; + NVTETensor dQ = nullptr; + NVTETensor dK = nullptr; + NVTETensor dV = nullptr; + NVTETensor dBias = nullptr; + NVTETensor dSoftmaxOffset = nullptr; + NVTETensor cu_seqlens_q = nullptr; + NVTETensor cu_seqlens_kv = nullptr; + NVTETensor cu_seqlens_q_padded = nullptr; + NVTETensor cu_seqlens_kv_padded = nullptr; + bool cuda_graph = false; + bool deterministic = false; + NVTE_Mask_Type attn_mask_type = NVTE_NO_MASK; + NVTE_Bias_Type bias_type = NVTE_NO_BIAS; + 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; + float dropout = 0.0f; + float attn_scale = 1.0f; + NVTE_QKV_Layout qkv_layout = NVTE_QKV_Layout_NOT_SET; + NVTE_QKV_Format o_format = NVTE_QKV_Format_NOT_SET; + NVTE_QKV_Format do_format = NVTE_QKV_Format_NOT_SET; + NVTE_QKV_Layout dqkv_layout = NVTE_QKV_Layout_NOT_SET; + NVTE_QKV_Format qkv_scale_inv_format = NVTE_QKV_Format_NOT_SET; + NVTE_QKV_Format do_scale_inv_format = NVTE_QKV_Format_NOT_SET; + size_t max_seqlen_q = 0; + size_t max_seqlen_kv = 0; + 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(uint8_t), // cuda_graph + sizeof(uint8_t), // deterministic + sizeof(NVTE_Mask_Type), // attn_mask_type + sizeof(NVTE_Bias_Type), // bias_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(float), // dropout + sizeof(float), // attn_scale + sizeof(NVTE_QKV_Layout), // qkv_layout + sizeof(NVTE_QKV_Format), // o_format + sizeof(NVTE_QKV_Format), // do_format + sizeof(NVTE_QKV_Layout), // dqkv_layout + sizeof(NVTE_QKV_Format), // qkv_scale_inv_format + sizeof(NVTE_QKV_Format), // do_scale_inv_format + sizeof(size_t), // max_seqlen_q + sizeof(size_t), // max_seqlen_kv + sizeof(NVTETensor), // workspace + sizeof(cudaStream_t), // stream + }; + + // 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 fc21771297..b3b9922abf 100644 --- a/transformer_engine/common/fused_attn/fused_attn.cpp +++ b/transformer_engine/common/fused_attn/fused_attn.cpp @@ -10,6 +10,7 @@ #include "../cudnn_utils.h" #include "../util/cuda_runtime.h" #include "../util/system.h" +#include "config_and_params.h" #include "fused_attn_f16_arbitrary_seqlen.h" #include "fused_attn_fp8.h" #include "utils.h" @@ -225,308 +226,215 @@ 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; it's re-used (cleared + re-populated) +// on every call to nvte_get_fused_attn_backend_v2 on the same thread. +thread_local std::string fused_attn_backend_message_buffer; + +// Stash `reason` in the thread-local buffer and, if the caller asked for a diagnostic, +// publish a NUL-terminated pointer to it via `*message`. Safe to call with `message == nullptr`. +void set_message(const char **message, std::string reason) { + if (message == nullptr) return; + fused_attn_backend_message_buffer = std::move(reason); + *message = fused_attn_backend_message_buffer.c_str(); +} + +} // namespace + +// select a backend for fused attention; the diagnostic message is based on the first failure, not cumulative. +NVTE_Fused_Attn_Backend nvte_get_fused_attn_backend_v2(NVTEFusedAttnConfig config, + const char **message) { 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); - NVTE_QKV_Layout_Group layout_group = nvte_get_qkv_layout_group(qkv_layout); - auto cudnn_runtime_version = cudnnGetVersion(); + using namespace transformer_engine::fused_attn; + const FusedAttnConfig &cfg = *get_fused_attn_config(config); + set_message(message, ""); - // For ragged offsets we only support 32-bit prior to cuDNN 9.5 - // Only used when THD format is requested. + cudnnHandle_t handle = cudnnExecutionPlanManager::Instance().GetHandle(); + const NVTE_QKV_Format qkv_format = nvte_get_qkv_format(cfg.qkv_layout); + const NVTE_QKV_Layout_Group layout_group = nvte_get_qkv_layout_group(cfg.qkv_layout); + const auto cudnn_runtime_version = cudnnGetVersion(); + + // THD + 64-bit ragged offsets require cuDNN >= 9.5 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.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; + (qkv_format == NVTE_THD && + fused_attn::get_ragged_offset_dtype(layout_group, cfg.num_attn_heads, cfg.num_gqa_groups, + cfg.max_seqlen_q, cfg.max_seqlen_kv, cfg.head_dim_qk, + cfg.head_dim_v) == DType::kInt64); + if (requires_64bit_ragged_offset && cudnn_runtime_version < 90500) { + set_message(message, + "Configuration requires 64-bit ragged offsets, which require " + "cuDNN >= 9.5."); + return NVTE_Fused_Attn_Backend::NVTE_No_Backend; + } + + // THD requires padding-style mask + if (qkv_format == NVTE_QKV_Format::NVTE_THD && + cfg.attn_mask_type != NVTE_Mask_Type::NVTE_PADDING_MASK && + cfg.attn_mask_type != NVTE_Mask_Type::NVTE_PADDING_CAUSAL_MASK && + cfg.attn_mask_type != NVTE_Mask_Type::NVTE_PADDING_CAUSAL_BOTTOM_RIGHT_MASK) { + set_message(message, + "THD format requires PADDING / PADDING_CAUSAL / PADDING_CAUSAL_BOTTOM_RIGHT mask."); + return NVTE_Fused_Attn_Backend::NVTE_No_Backend; + } + + // cuDNN does not support pre-scale bias + if (cfg.bias_type == NVTE_Bias_Type::NVTE_PRE_SCALE_BIAS) { + set_message(message, "Fused attention does not support pre-scale bias."); + return NVTE_Fused_Attn_Backend::NVTE_No_Backend; + } + + 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); + + if (is_fp8) { + if (cfg.return_max_logit) { + set_message(message, "FP8 fused attention does not support return_max_logit=True."); + return NVTE_Fused_Attn_Backend::NVTE_No_Backend; } - if (flag_arb) { - backend = NVTE_Fused_Attn_Backend::NVTE_F16_arbitrary_seqlen; + if (qkv_format != NVTE_QKV_Format::NVTE_BSHD && qkv_format != NVTE_QKV_Format::NVTE_SBHD && + qkv_format != NVTE_QKV_Format::NVTE_BHSD) { + set_message(message, "FP8 fused attention supports BSHD/SBHD/BHSD formats, found " + + std::to_string(static_cast(qkv_format)) + "."); + return NVTE_Fused_Attn_Backend::NVTE_No_Backend; } - 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; + std::string fwd_reason = is_supported_fp8_fwd(cfg, handle); + if (!fwd_reason.empty()) { + set_message(message, std::move(fwd_reason)); + return NVTE_Fused_Attn_Backend::NVTE_No_Backend; } - 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_training && !cfg.is_forward) { + std::string bwd_reason = is_supported_fp8_bwd(cfg, handle); + if (!bwd_reason.empty()) { + set_message(message, std::move(bwd_reason)); + return NVTE_Fused_Attn_Backend::NVTE_No_Backend; + } } - if ((cudnn_runtime_version <= 91500) && is_training && + return NVTE_Fused_Attn_Backend::NVTE_FP8; + } + + if (is_f16_or_bf16) { + if (cudnn_runtime_version <= 91500 && cfg.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; + (cfg.max_seqlen_kv % 128 != 0) && cfg.cuda_graph && + cfg.attn_mask_type != NVTE_Mask_Type::NVTE_PADDING_MASK && + cfg.attn_mask_type != NVTE_Mask_Type::NVTE_PADDING_CAUSAL_MASK && + cfg.attn_mask_type != NVTE_Mask_Type::NVTE_PADDING_CAUSAL_BOTTOM_RIGHT_MASK) { + set_message(message, "Known cuDNN <= 9.15 issue with CUDA graph. Please upgrade cuDNN."); + return NVTE_Fused_Attn_Backend::NVTE_No_Backend; + } + std::string fwd_reason = is_supported_f16_fwd(cfg, handle); + if (!fwd_reason.empty()) { + set_message(message, std::move(fwd_reason)); + return NVTE_Fused_Attn_Backend::NVTE_No_Backend; } - if (backend == NVTE_Fused_Attn_Backend::NVTE_F16_arbitrary_seqlen && sm_arch_ == 120) { - if (cudnn_runtime_version < 91801) { - backend = NVTE_Fused_Attn_Backend::NVTE_No_Backend; - std::cout << "Warning: Given combination of sm_arch_ == 120 and cudnn_runtime_version < " - "91801 is not supported. " - << " Please upgrade your cuDNN version if possible." << std::endl; - } else if (deterministic && is_training) { - backend = NVTE_Fused_Attn_Backend::NVTE_No_Backend; - std::cout << "Warning: Deterministic fused attention on SM120 is not supported." - << std::endl; - } else { - // Known missing support for T3HD/TH3D layouts on SM120 - const bool is_t3hd_or_th3d = - (qkv_layout == NVTE_QKV_Layout::NVTE_T3HD || qkv_layout == NVTE_QKV_Layout::NVTE_TH3D); - if (is_t3hd_or_th3d) { - backend = NVTE_Fused_Attn_Backend::NVTE_No_Backend; - std::cout << "Warning: Given combination of T3HD/TH3D layouts on SM120 is not supported. " - << " Please consider using other THD layouts if possible." << std::endl; - } + if (cfg.is_training && !cfg.is_forward) { + std::string bwd_reason = is_supported_f16_bwd(cfg, handle); + if (!bwd_reason.empty()) { + set_message(message, std::move(bwd_reason)); + return NVTE_Fused_Attn_Backend::NVTE_No_Backend; } } + return NVTE_Fused_Attn_Backend::NVTE_F16_arbitrary_seqlen; + } + + set_message(message, "Unsupported QKV dtype qkv_dtype=" + std::to_string(cfg.qkv_dtype) + " ."); + return NVTE_Fused_Attn_Backend::NVTE_No_Backend; +} + +// 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) { + 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 missing fields so it doesn't always return NVTE_No_Backend + 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 +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(); + 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 *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 @@ -545,98 +453,114 @@ 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) { 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; - } + 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 +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(); + 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 *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, @@ -654,92 +578,45 @@ void nvte_fused_attn_bwd(const NVTETensor Q, const NVTETensor K, const NVTETenso 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"); - } + 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..b7c7a349af 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 @@ -10,6 +10,7 @@ #include #include +#include #include #include "../common.h" @@ -17,79 +18,73 @@ #include "../util/cuda_runtime.h" #include "../util/system.h" #include "fused_attn_f16_arbitrary_seqlen.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, + 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, cudnn_frontend::DataType_t tensorType, - void *workspace, size_t *workspace_size, cudaStream_t stream, cudnnHandle_t handle) { + void *devPtrSeqOffsetsQ, void *devPtrSeqOffsetsKV, void *workspace, size_t *workspace_size, + cudaStream_t stream, cudnnHandle_t handle) { using namespace transformer_engine; + const cudnn_frontend::DataType_t tensorType = + get_cudnn_fe_dtype(static_cast(cfg.qkv_dtype)); + + 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); + int64_t s_q = static_cast(cfg.max_seqlen_q); + int64_t s_kv = static_cast(cfg.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); + int64_t bucketed_batch_size = static_cast(cfg.bucketed_batch_size); + int64_t bucketed_num_tokens_q = static_cast(cfg.bucketed_num_tokens_q); + int64_t bucketed_num_tokens_kv = static_cast(cfg.bucketed_num_tokens_kv); + int64_t num_pages_k = static_cast(cfg.num_pages_k); + int64_t num_pages_v = static_cast(cfg.num_pages_v); + int64_t page_size_k = static_cast(cfg.page_size_k); + int64_t page_size_v = static_cast(cfg.page_size_v); + int64_t max_pages_per_seq_k = static_cast(cfg.max_pages_per_seq_k); + int64_t max_pages_per_seq_v = static_cast(cfg.max_pages_per_seq_v); + int64_t bias_b = static_cast(cfg.bias_batch_size); + int64_t bias_h = static_cast(cfg.bias_num_heads); + int64_t bias_sq = static_cast(cfg.bias_seqlen_q); + int64_t bias_skv = static_cast(cfg.bias_seqlen_kv); + const bool is_training = cfg.is_training; + const bool return_max_logit = cfg.return_max_logit; + float scaling_factor = cfg.attn_scale; + const float dropout_probability = cfg.dropout; + const NVTE_QKV_Layout qkv_layout = cfg.qkv_layout; + const NVTE_Bias_Type bias_type = cfg.bias_type; + const NVTE_Mask_Type mask_type = cfg.attn_mask_type; + const NVTE_Softmax_Type softmax_type = cfg.softmax_type; + const int64_t window_size_left = cfg.window_size_left; + const int64_t window_size_right = cfg.window_size_right; + bool bottom_right_diagonal = cfg.bottom_right_diagonal; + 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_causal_bottom_right = cfg.is_causal_bottom_right; + bool is_padding = cfg.is_padding; 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); + bool is_ragged_q = cfg.is_ragged_q; + 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); + bool is_paged_kv = cfg.is_paged_kv; if (is_paged_kv) { NVTE_CHECK(is_padding, "Paged attention requires padding mask!"); } @@ -97,16 +92,7 @@ void fused_attn_arbitrary_seqlen_fwd_impl( // 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; + const bool use_cu_seqlens_directly = cfg.uses_cu_seqlens_directly; // keep original batch size because cu_seqlens are created with [b+1] shape int64_t actual_b = b; @@ -122,10 +108,10 @@ void fused_attn_arbitrary_seqlen_fwd_impl( // 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; + b = bucketed_batch_size; } - s_q = is_ragged_q ? max_t_q : s_q; - s_kv = is_ragged_kv ? max_t_kv : s_kv; + s_q = is_ragged_q ? bucketed_num_tokens_q : s_q; + s_kv = is_ragged_kv ? bucketed_num_tokens_kv : s_kv; } } @@ -139,48 +125,8 @@ void fused_attn_arbitrary_seqlen_fwd_impl( const RaggedOffsetMultipliers offset_mults(layout_group, h, hg, d_qk, d_v); bool generate_stats = true; // Always return stats + const FusedAttnConfig cache_cfg = cfg.make_cache_key(); 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, - }; - namespace fe = cudnn_frontend; using graph_and_tensors = std::tuple, @@ -205,16 +151,26 @@ void fused_attn_arbitrary_seqlen_fwd_impl( std::shared_ptr, // dropout_seed std::shared_ptr>; // dropout_offset - using CacheType = std::map; - static thread_local CacheType sdpa_f16_fprop_cache; + using CacheType = std::map; + // Process-wide graph cache so a compiled graph is reused across threads instead of rebuilt per thread. + // Safe because cuDNN >= 9.0 allows concurrent execution of a shared plan and cudnn-frontend >= 1.25.0 has a thread-safe execute(). + static CacheType sdpa_f16_fprop_cache; + static std::mutex sdpa_f16_fprop_cache_mutex; // 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; + auto get_graph = [&](CacheType &cache, const FusedAttnConfig &descriptor) -> graph_and_tensors { + // Lock the map lookup, not the build, so different graphs can build in parallel + graph_and_tensors cached_graph{}; + bool cache_hit = false; + { + std::lock_guard shared_cache_lock(sdpa_f16_fprop_cache_mutex); + auto it = cache.find(descriptor); + cache_hit = (it != cache.end()); + if (cache_hit) cached_graph = it->second; + } + graph_cache_debug::record_cache_lookup("fwd", cache_hit, cfg); + if (cache_hit) { + return cached_graph; } // otherwise, build the op_graph and the plan. Then update cache @@ -300,8 +256,6 @@ void fused_attn_arbitrary_seqlen_fwd_impl( 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 = @@ -314,6 +268,9 @@ void fused_attn_arbitrary_seqlen_fwd_impl( if (cudnn_runtime_version >= 90600 && window_size_right != -1) { sdpa_options.set_diagonal_band_right_bound(window_size_right); } + if (is_causal || is_causal_bottom_right) { + sdpa_options.set_diagonal_band_right_bound(0); + } sdpa_options.set_alibi_mask(is_alibi); @@ -479,24 +436,36 @@ void fused_attn_arbitrary_seqlen_fwd_impl( 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)); + graph_cache_debug::timer("fwd", graph_cache_debug::BuildStage::Validate, + [&] { NVTE_CHECK_CUDNN_FE(mha_graph->validate()); }); + graph_cache_debug::timer("fwd", graph_cache_debug::BuildStage::BuildOpGraph, [&] { + NVTE_CHECK_CUDNN_FE(mha_graph->build_operation_graph(handle)); + }); + graph_cache_debug::timer("fwd", graph_cache_debug::BuildStage::CreatePlans, [&] { + NVTE_CHECK_CUDNN_FE(mha_graph->create_execution_plans({fe::HeurMode_t::A})); + }); + graph_cache_debug::timer("fwd", graph_cache_debug::BuildStage::CheckSupport, + [&] { NVTE_CHECK_CUDNN_FE(mha_graph->check_support()); }); + graph_cache_debug::timer("fwd", graph_cache_debug::BuildStage::BuildPlans, + [&] { NVTE_CHECK_CUDNN_FE(mha_graph->build_plans()); }); 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; + graph_cache_debug::record_build("fwd"); + // Lock the insert. If another thread inserted a graph for the same key while we were building, + // use their graph (it's the same as ours) and discard our graph. + { + std::lock_guard shared_cache_lock(sdpa_f16_fprop_cache_mutex); + auto inserted = cache.insert({descriptor, return_tuple}); + return inserted.first->second; + } }; 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] = get_graph(sdpa_f16_fprop_cache, cache_cfg); // 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,6 +494,7 @@ void fused_attn_arbitrary_seqlen_fwd_impl( plan_workspace_size + actual_seqlen_workspace_size + seqlen_offsets_workspace_size; return; } + graph_cache_debug::record_exec("fwd"); // cuDNN stream check needs to be moved here to support dummy kernel calls with // null streams for sizing the cuDNN workspace. @@ -640,52 +610,58 @@ void fused_attn_arbitrary_seqlen_fwd_impl( } // 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) { + 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; + const cudnn_frontend::DataType_t tensorType = + get_cudnn_fe_dtype(static_cast(cfg.qkv_dtype)); + + 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); + int64_t s_q = static_cast(cfg.max_seqlen_q); + int64_t s_kv = static_cast(cfg.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); + int64_t bucketed_batch_size = static_cast(cfg.bucketed_batch_size); + int64_t bucketed_num_tokens_q = static_cast(cfg.bucketed_num_tokens_q); + int64_t bucketed_num_tokens_kv = static_cast(cfg.bucketed_num_tokens_kv); + int64_t bias_b = static_cast(cfg.bias_batch_size); + int64_t bias_h = static_cast(cfg.bias_num_heads); + int64_t bias_sq = static_cast(cfg.bias_seqlen_q); + int64_t bias_skv = static_cast(cfg.bias_seqlen_kv); + float scaling_factor = cfg.attn_scale; + const float dropout_probability = cfg.dropout; + const NVTE_QKV_Layout qkv_layout = cfg.qkv_layout; + const NVTE_Bias_Type bias_type = cfg.bias_type; + const NVTE_Mask_Type mask_type = cfg.attn_mask_type; + const NVTE_Softmax_Type softmax_type = cfg.softmax_type; + const int64_t window_size_left = cfg.window_size_left; + const int64_t window_size_right = cfg.window_size_right; + bool bottom_right_diagonal = cfg.bottom_right_diagonal; + const bool deterministic = cfg.deterministic; + 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_causal_bottom_right = cfg.is_causal_bottom_right; + bool is_padding = cfg.is_padding; 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); + bool is_ragged_q = cfg.is_ragged_q; + 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!"); - } - // 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) { @@ -694,57 +670,17 @@ void fused_attn_arbitrary_seqlen_bwd_impl( 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; + b = bucketed_batch_size; + s_q = is_ragged_q ? bucketed_num_tokens_q : s_q; + s_kv = is_ragged_kv ? bucketed_num_tokens_kv : s_kv; } } // 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; + const FusedAttnConfig cache_cfg = cfg.make_cache_key(); 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, @@ -772,16 +708,24 @@ void fused_attn_arbitrary_seqlen_bwd_impl( std::shared_ptr, // dropout_seed std::shared_ptr>; // dropout_offset - using CacheType = std::map; - static thread_local CacheType sdpa_f16_bprop_cache; + using CacheType = std::map; + static CacheType sdpa_f16_bprop_cache; + static std::mutex sdpa_f16_bprop_cache_mutex; // 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; + auto get_graph = [&](CacheType &cache, const FusedAttnConfig &descriptor) -> graph_and_tensors { + // Lock the map lookup, not the build, so different graphs can build in parallel + graph_and_tensors cached_graph{}; + bool cache_hit = false; + { + std::lock_guard shared_cache_lock(sdpa_f16_bprop_cache_mutex); + auto it = cache.find(descriptor); + cache_hit = (it != cache.end()); + if (cache_hit) cached_graph = it->second; + } + graph_cache_debug::record_cache_lookup("bwd", cache_hit, cfg); + if (cache_hit) { + return cached_graph; } // otherwise, build the op_graph and the plan. Then update cache @@ -886,8 +830,6 @@ void fused_attn_arbitrary_seqlen_bwd_impl( 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) { @@ -908,6 +850,9 @@ void fused_attn_arbitrary_seqlen_bwd_impl( if (cudnn_runtime_version >= 90600 && window_size_right != -1) { sdpa_backward_options.set_diagonal_band_right_bound(window_size_right); } + if (is_causal || is_causal_bottom_right) { + sdpa_backward_options.set_diagonal_band_right_bound(0); + } if (cudnn_runtime_version >= 90000) { sdpa_backward_options.set_deterministic_algorithm(deterministic); @@ -1018,23 +963,35 @@ void fused_attn_arbitrary_seqlen_bwd_impl( 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)); + graph_cache_debug::timer("bwd", graph_cache_debug::BuildStage::Validate, + [&] { NVTE_CHECK_CUDNN_FE(mha_graph->validate()); }); + graph_cache_debug::timer("bwd", graph_cache_debug::BuildStage::BuildOpGraph, [&] { + NVTE_CHECK_CUDNN_FE(mha_graph->build_operation_graph(handle)); + }); + graph_cache_debug::timer("bwd", graph_cache_debug::BuildStage::CreatePlans, [&] { + NVTE_CHECK_CUDNN_FE(mha_graph->create_execution_plans({fe::HeurMode_t::A})); + }); + graph_cache_debug::timer("bwd", graph_cache_debug::BuildStage::CheckSupport, + [&] { NVTE_CHECK_CUDNN_FE(mha_graph->check_support()); }); + graph_cache_debug::timer("bwd", graph_cache_debug::BuildStage::BuildPlans, + [&] { NVTE_CHECK_CUDNN_FE(mha_graph->build_plans()); }); 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; + graph_cache_debug::record_build("bwd"); + // Lock the insert. If another thread inserted a graph for the same key while we were building, + // use their graph (it's the same as ours) and discard our graph. + { + std::lock_guard shared_cache_lock(sdpa_f16_bprop_cache_mutex); + auto inserted = cache.insert({descriptor, return_tuple}); + return inserted.first->second; + } }; 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] = get_graph(sdpa_f16_bprop_cache, cache_cfg); // 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,6 +1015,7 @@ void fused_attn_arbitrary_seqlen_bwd_impl( plan_workspace_size + actual_seqlen_workspace_size + seqlen_offsets_workspace_size; return; } + graph_cache_debug::record_exec("bwd"); // cuDNN stream check needs to be moved here to support dummy kernel calls with // null streams for sizing the cuDNN workspace. @@ -1160,25 +1118,28 @@ 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_QKV_Layout qkv_layout = cfg.qkv_layout; + 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,16 +1147,8 @@ 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) { @@ -1212,18 +1165,8 @@ 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); - } + FusedAttnConfig graph_cfg = cfg; + graph_cfg.derive(); size_t i = 0; if (Aux_CTX_Tensors->size == 0) { @@ -1261,7 +1204,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 = {graph_cfg.bias_batch_size, graph_cfg.bias_num_heads, + graph_cfg.bias_seqlen_q, graph_cfg.bias_seqlen_kv}; output_bias->data.dtype = QKV_type; } @@ -1302,14 +1246,9 @@ 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, - devPtrDropoutSeed, devPtrDropoutOffset, devPtrCuSeqlensQ, devPtrCuSeqlensKV, devPtrPageTableK, - devPtrPageTableV, devPtrSeqOffsetsQ, devPtrSeqOffsetsKV, get_cudnn_fe_dtype(QKV_type), + graph_cfg, devPtrQ, devPtrK, devPtrV, devPtrBias, devPtrSoftmaxOffset, devPtrS1, devPtrS2, + devPtrO, devPtrDropoutSeed, devPtrDropoutOffset, devPtrCuSeqlensQ, devPtrCuSeqlensKV, + devPtrPageTableK, devPtrPageTableV, devPtrSeqOffsetsQ, devPtrSeqOffsetsKV, workspace->data.dptr, &workspace_size, stream, handle); if (workspace_size > 0) { @@ -1327,22 +1266,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,33 +1288,13 @@ 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); - } + FusedAttnConfig graph_cfg = cfg; + graph_cfg.derive(); void *devPtrdQ = output_dQ->data.dptr; void *devPtrdK = output_dK->data.dptr; @@ -1402,14 +1320,11 @@ 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); + graph_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 +1340,57 @@ void fused_attn_arbitrary_seqlen_bwd( NVTE_ERROR("Unexpected workspace_size."); } } + +std::string is_supported_f16_fwd(const FusedAttnConfig &cfg, cudnnHandle_t handle) { + FusedAttnConfig graph_cfg = cfg; + graph_cfg.is_forward = true; + graph_cfg.derive(); + + size_t workspace_size = 0; + try { + fused_attn::fused_attn_arbitrary_seqlen_fwd_impl( + graph_cfg, + /*devPtrQ=*/nullptr, /*devPtrK=*/nullptr, /*devPtrV=*/nullptr, /*devPtrBias=*/nullptr, + /*devPtrSoftmaxOffset=*/nullptr, /*devPtrS1=*/nullptr, /*devPtrS2=*/nullptr, + /*devPtrO=*/nullptr, /*devPtrDropoutSeed=*/nullptr, /*devPtrDropoutOffset=*/nullptr, + /*devPtrCuSeqlensQ=*/nullptr, /*devPtrCuSeqlensKV=*/nullptr, + /*devPtrPageTableK=*/nullptr, /*devPtrPageTableV=*/nullptr, + /*devPtrSeqOffsetsQ=*/nullptr, /*devPtrSeqOffsetsKV=*/nullptr, + /*workspace=*/nullptr, &workspace_size, + /*stream=*/static_cast(0), handle); + return ""; + } catch (const std::exception &e) { + return e.what(); + } catch (...) { + return "is_supported_f16_fwd: unknown failure."; + } +} + +std::string is_supported_f16_bwd(const FusedAttnConfig &cfg, cudnnHandle_t handle) { + FusedAttnConfig graph_cfg = cfg; + graph_cfg.is_forward = false; + graph_cfg.derive(); + + size_t workspace_size = 0; + try { + fused_attn::fused_attn_arbitrary_seqlen_bwd_impl( + graph_cfg, + /*devPtrQ=*/nullptr, /*devPtrKTranspose=*/nullptr, + /*devPtrVTranspose=*/nullptr, /*devPtrO=*/nullptr, /*devPtrSoftmaxStats=*/nullptr, + /*devPtrBias=*/nullptr, /*devPtrSoftmaxOffset=*/nullptr, /*devPtrdQ=*/nullptr, + /*devPtrdK=*/nullptr, /*devPtrdV=*/nullptr, /*devPtrdO=*/nullptr, + /*devPtrdBias=*/nullptr, /*devPtrdSoftmaxOffset=*/nullptr, + /*devPtrDropoutSeed=*/nullptr, /*devPtrDropoutOffset=*/nullptr, + /*devPtrCuSeqlensQ=*/nullptr, /*devPtrCuSeqlensKV=*/nullptr, + /*devPtrSeqOffsetsQ=*/nullptr, /*devPtrSeqOffsetsKV=*/nullptr, + /*workspace=*/nullptr, &workspace_size, + /*stream=*/static_cast(0), handle); + return ""; + } catch (const std::exception &e) { + return e.what(); + } catch (...) { + return "is_supported_f16_bwd: unknown failure."; + } +} + } // 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..c493b5ee1b 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,49 @@ * \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); + +// check if a given configuration is supported for F16/BF16 forward; +// if it is, cache the graph built for this config, and return an empty string; +// if not, return a diagnostic message explaining why it is not supported. +std::string is_supported_f16_fwd(const fused_attn::FusedAttnConfig &cfg, cudnnHandle_t handle); + +// check if a given configuration is supported for F16/BF16 backward; +// if it is, cache the graph built for this config, and return an empty string; +// if not, return a diagnostic message explaining why it is not supported. +std::string is_supported_f16_bwd(const fused_attn::FusedAttnConfig &cfg, 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..9449f74206 100644 --- a/transformer_engine/common/fused_attn/fused_attn_fp8.cu +++ b/transformer_engine/common/fused_attn/fused_attn_fp8.cu @@ -4,10 +4,14 @@ * See LICENSE for license information. ************************************************************************/ +#include +#include + #include "../common.h" #include "../cudnn_utils.h" #include "../util/system.h" #include "fused_attn_fp8.h" +#include "graph_cache_debug.h" #include "utils.h" namespace transformer_engine { @@ -16,33 +20,51 @@ 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) { +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; const auto cudnn_runtime_version = cudnnGetVersion(); + + const cudnn_frontend::DataType_t qkv_tensor_type = + get_cudnn_fe_dtype(static_cast(cfg.qkv_dtype)); + const cudnn_frontend::DataType_t o_tensor_type = + get_cudnn_fe_dtype(static_cast(cfg.o_dtype)); + + const int64_t b = static_cast(cfg.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.max_seqlen_q); + const int64_t s_kv = static_cast(cfg.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 bool is_training = cfg.is_training; + float scaling_factor = cfg.attn_scale; + 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_Bias_Type bias_type = cfg.bias_type; + const NVTE_Mask_Type mask_type = cfg.attn_mask_type; + const NVTE_Softmax_Type softmax_type = cfg.softmax_type; + const int64_t window_size_left = cfg.window_size_left; + const int64_t window_size_right = cfg.window_size_right; + const bool bottom_right_diagonal = cfg.bottom_right_diagonal; + const NVTEScalingMode scaling_mode = cfg.scaling_mode; + const NVTE_QKV_Format qkv_scale_inv_format = cfg.qkv_scale_inv_format; + 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_causal_bottom_right = cfg.is_causal_bottom_right; + bool is_padding = cfg.is_padding; 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) && @@ -78,46 +100,8 @@ void fused_attn_fp8_fwd_impl( // (which doesn't support cu_seqlens). Remove this restriction when possible. !is_dropout; + const FusedAttnConfig cache_cfg = cfg.make_cache_key(); 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, @@ -142,16 +126,26 @@ void fused_attn_fp8_fwd_impl( std::shared_ptr, // dropout_seed std::shared_ptr>; // dropout_offset - using CacheType = std::map; - static thread_local CacheType sdpa_fp8_fprop_cache; + using CacheType = std::map; + // Process-wide graph cache so a compiled graph is reused across threads instead of rebuilt per thread. + // Safe because cuDNN >= 9.0 allows concurrent execution of a shared plan and cudnn-frontend >= 1.25.0 has a thread-safe execute(). + static CacheType sdpa_fp8_fprop_cache; + static std::mutex sdpa_fp8_fprop_cache_mutex; // 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; + auto get_graph = [&](CacheType& cache, const FusedAttnConfig& descriptor) -> graph_and_tensors { + // Lock the map lookup, not the build, so different graphs can build in parallel + graph_and_tensors cached_graph{}; + bool cache_hit = false; + { + std::lock_guard shared_cache_lock(sdpa_fp8_fprop_cache_mutex); + auto it = cache.find(descriptor); + cache_hit = (it != cache.end()); + if (cache_hit) cached_graph = it->second; + } + graph_cache_debug::record_cache_lookup("fwd", cache_hit, cfg); + if (cache_hit) { + return cached_graph; } // otherwise, build the op_graph and the plan. Then update cache @@ -210,12 +204,11 @@ void fused_attn_fp8_fwd_impl( 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 q_scale_inv_format = + (qkv_scale_inv_format != NVTE_QKV_Format_NOT_SET) ? qkv_scale_inv_format : cfg.q_format; 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); + : cfg.kv_format; std::vector q_scale_strides(4); std::vector k_scale_strides(4); std::vector v_scale_strides(4); @@ -269,6 +262,9 @@ void fused_attn_fp8_fwd_impl( sdpa_options.set_diagonal_band_right_bound(window_size_right); } } + if (is_causal_bottom_right) { + sdpa_options.set_diagonal_band_right_bound(0); + } // sdpa_options.set_alibi_mask(is_alibi); // if (is_bias) { @@ -399,22 +395,34 @@ void fused_attn_fp8_fwd_impl( 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)); + graph_cache_debug::timer("fwd", graph_cache_debug::BuildStage::Validate, + [&] { NVTE_CHECK_CUDNN_FE(mha_graph->validate()); }); + graph_cache_debug::timer("fwd", graph_cache_debug::BuildStage::BuildOpGraph, [&] { + NVTE_CHECK_CUDNN_FE(mha_graph->build_operation_graph(handle)); + }); + graph_cache_debug::timer("fwd", graph_cache_debug::BuildStage::CreatePlans, [&] { + NVTE_CHECK_CUDNN_FE(mha_graph->create_execution_plans({fe::HeurMode_t::A})); + }); + graph_cache_debug::timer("fwd", graph_cache_debug::BuildStage::CheckSupport, + [&] { NVTE_CHECK_CUDNN_FE(mha_graph->check_support()); }); + graph_cache_debug::timer("fwd", graph_cache_debug::BuildStage::BuildPlans, + [&] { NVTE_CHECK_CUDNN_FE(mha_graph->build_plans()); }); 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; + graph_cache_debug::record_build("fwd"); + // Lock the insert. If another thread inserted a graph for the same key while we were building, + // use their graph (it's the same as ours) and discard our graph. + { + std::lock_guard shared_cache_lock(sdpa_fp8_fprop_cache_mutex); + auto inserted = cache.insert({descriptor, return_tuple}); + return inserted.first->second; + } }; 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] = get_graph(sdpa_fp8_fprop_cache, cache_cfg); auto plan_workspace_size = mha_graph->get_workspace_size(); @@ -426,6 +434,7 @@ void fused_attn_fp8_fwd_impl( *workspace_size = plan_workspace_size + actual_seqlen_workspace_size; return; } + graph_cache_debug::record_exec("fwd"); // cuDNN stream check needs to be moved here to support dummy kernel calls with // null streams for sizing the cuDNN workspace. @@ -467,7 +476,7 @@ void fused_attn_fp8_fwd_impl( void* devActualSeqlenQ = static_cast(workspace) + plan_workspace_size; void* devActualSeqlenKV = static_cast(devActualSeqlenQ) + b * sizeof(int32_t); cu_seqlens_to_actual_seqlens<<>>( - b, b, static_cast(devPtrcuSeqlensQ), // TODO(pass max_b) + b, b, static_cast(devPtrcuSeqlensQ), // TODO(pass bucketed_batch_size) static_cast(devPtrcuSeqlensKV), static_cast(devActualSeqlenQ), static_cast(devActualSeqlenKV)); NVTE_CHECK_CUDA(cudaGetLastError()); @@ -493,39 +502,61 @@ void fused_attn_fp8_fwd_impl( // 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) { + 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; const auto cudnn_runtime_version = cudnnGetVersion(); + + const cudnn_frontend::DataType_t qkv_tensor_type = + get_cudnn_fe_dtype(static_cast(cfg.qkv_dtype)); + const cudnn_frontend::DataType_t o_tensor_type = + get_cudnn_fe_dtype(static_cast(cfg.o_dtype)); + const cudnn_frontend::DataType_t do_tensor_type = + get_cudnn_fe_dtype(static_cast(cfg.do_dtype)); + const cudnn_frontend::DataType_t dqkv_tensor_type = + get_cudnn_fe_dtype(static_cast(cfg.dqkv_dtype)); + + const int64_t b = static_cast(cfg.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.max_seqlen_q); + const int64_t s_kv = static_cast(cfg.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); + float scaling_factor = cfg.attn_scale; + 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 do_format = cfg.do_format; + const NVTE_QKV_Layout dqkv_layout = cfg.dqkv_layout; + const NVTE_Bias_Type bias_type = cfg.bias_type; + const NVTE_Mask_Type mask_type = cfg.attn_mask_type; + const NVTE_Softmax_Type softmax_type = cfg.softmax_type; + const int64_t window_size_left = cfg.window_size_left; + const int64_t window_size_right = cfg.window_size_right; + const bool bottom_right_diagonal = cfg.bottom_right_diagonal; + const bool deterministic = cfg.deterministic; + const NVTEScalingMode scaling_mode = cfg.scaling_mode; + 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; + 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_causal_bottom_right = cfg.is_causal_bottom_right; + bool is_padding = cfg.is_padding; 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) && @@ -546,46 +577,8 @@ void fused_attn_fp8_bwd_impl( bool is_O_in_F16 = (o_tensor_type == cudnn_frontend::DataType_t::HALF || o_tensor_type == cudnn_frontend::DataType_t::BFLOAT16); + const FusedAttnConfig cache_cfg = cfg.make_cache_key(); 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, @@ -631,16 +624,24 @@ void fused_attn_fp8_bwd_impl( std::shared_ptr, // dropout_seed std::shared_ptr>; // dropout_offset - using CacheType = std::map; - static thread_local CacheType sdpa_fp8_bprop_cache; + using CacheType = std::map; + static CacheType sdpa_fp8_bprop_cache; + static std::mutex sdpa_fp8_bprop_cache_mutex; // 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; + auto get_graph = [&](CacheType& cache, const FusedAttnConfig& descriptor) -> graph_and_tensors { + // Lock the map lookup, not the build, so different graphs can build in parallel + graph_and_tensors cached_graph{}; + bool cache_hit = false; + { + std::lock_guard shared_cache_lock(sdpa_fp8_bprop_cache_mutex); + auto it = cache.find(descriptor); + cache_hit = (it != cache.end()); + if (cache_hit) cached_graph = it->second; + } + graph_cache_debug::record_cache_lookup("bwd", cache_hit, cfg); + if (cache_hit) { + return cached_graph; } // otherwise, build the op_graph and the plan. Then update cache @@ -735,8 +736,8 @@ void fused_attn_fp8_bwd_impl( 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_format = cfg.q_format; + NVTE_QKV_Format kv_format = cfg.kv_format; 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 = @@ -856,6 +857,9 @@ void fused_attn_fp8_bwd_impl( 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); @@ -1016,24 +1020,36 @@ void fused_attn_fp8_bwd_impl( 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)); + graph_cache_debug::timer("bwd", graph_cache_debug::BuildStage::Validate, + [&] { NVTE_CHECK_CUDNN_FE(mha_graph->validate()); }); + graph_cache_debug::timer("bwd", graph_cache_debug::BuildStage::BuildOpGraph, [&] { + NVTE_CHECK_CUDNN_FE(mha_graph->build_operation_graph(handle)); + }); + graph_cache_debug::timer("bwd", graph_cache_debug::BuildStage::CreatePlans, [&] { + NVTE_CHECK_CUDNN_FE(mha_graph->create_execution_plans({fe::HeurMode_t::A})); + }); + graph_cache_debug::timer("bwd", graph_cache_debug::BuildStage::CheckSupport, + [&] { NVTE_CHECK_CUDNN_FE(mha_graph->check_support()); }); + graph_cache_debug::timer("bwd", graph_cache_debug::BuildStage::BuildPlans, + [&] { NVTE_CHECK_CUDNN_FE(mha_graph->build_plans()); }); 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; + graph_cache_debug::record_build("bwd"); + // Lock the insert. If another thread inserted a graph for the same key while we were building, + // use their graph (it's the same as ours) and discard our graph. + { + std::lock_guard shared_cache_lock(sdpa_fp8_bprop_cache_mutex); + auto inserted = cache.insert({descriptor, return_tuple}); + return inserted.first->second; + } }; 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] = get_graph(sdpa_fp8_bprop_cache, cache_cfg); auto plan_workspace_size = mha_graph->get_workspace_size(); @@ -1043,6 +1059,7 @@ void fused_attn_fp8_bwd_impl( *workspace_size = plan_workspace_size + actual_seqlen_workspace_size; return; } + graph_cache_debug::record_exec("bwd"); // cuDNN stream check needs to be moved here to support dummy kernel calls with // null streams for sizing the cuDNN workspace. @@ -1108,7 +1125,7 @@ void fused_attn_fp8_bwd_impl( void* devActualSeqlenQ = static_cast(workspace) + plan_workspace_size; void* devActualSeqlenKV = static_cast(devActualSeqlenQ) + b * sizeof(int32_t); cu_seqlens_to_actual_seqlens<<>>( - b, b, static_cast(devPtrcuSeqlensQ), // TODO(pass max_b) + b, b, static_cast(devPtrcuSeqlensQ), // TODO(pass bucketed_batch_size) static_cast(devPtrcuSeqlensKV), static_cast(devActualSeqlenQ), static_cast(devActualSeqlenKV)); NVTE_CHECK_CUDA(cudaGetLastError()); @@ -1134,18 +1151,23 @@ void fused_attn_fp8_bwd_impl( } // 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 +1234,19 @@ 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; + FusedAttnConfig graph_cfg = cfg; + graph_cfg.derive(); + 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); + graph_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 +1264,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 +1292,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 +1349,23 @@ 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; + FusedAttnConfig graph_cfg = cfg; + graph_cfg.derive(); + 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, - 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, + graph_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 +1382,64 @@ void fused_attn_fp8_bwd( return; } } + +std::string is_supported_fp8_fwd(const FusedAttnConfig& cfg, cudnnHandle_t handle) { + FusedAttnConfig graph_cfg = cfg; + graph_cfg.is_forward = true; + graph_cfg.derive(); + + size_t workspace_size = 0; + try { + fused_attn::fused_attn_fp8_fwd_impl( + graph_cfg, + /*devPtrQ=*/nullptr, /*devPtrK=*/nullptr, /*devPtrV=*/nullptr, + /*devPtrSoftmaxOffset=*/nullptr, /*devPtrM=*/nullptr, /*devPtrO=*/nullptr, + /*devPtrDescaleQ=*/nullptr, /*devPtrDescaleK=*/nullptr, /*devPtrDescaleV=*/nullptr, + /*devPtrDescaleS=*/nullptr, /*devPtrScaleS=*/nullptr, /*devPtrScaleO=*/nullptr, + /*devPtrAmaxO=*/nullptr, /*devPtrAmaxS=*/nullptr, /*devPtrcuSeqlensQ=*/nullptr, + /*devPtrcuSeqlensKV=*/nullptr, /*devPtrDropoutSeed=*/nullptr, + /*devPtrDropoutOffset=*/nullptr, + /*workspace=*/nullptr, &workspace_size, + /*stream=*/static_cast(0), handle); + return ""; + } catch (const std::exception& e) { + return e.what(); + } catch (...) { + return "is_supported_fp8_fwd: unknown failure."; + } +} + +std::string is_supported_fp8_bwd(const FusedAttnConfig& cfg, cudnnHandle_t handle) { + FusedAttnConfig graph_cfg = cfg; + graph_cfg.is_forward = false; + graph_cfg.derive(); + + size_t workspace_size = 0; + try { + fused_attn::fused_attn_fp8_bwd_impl( + graph_cfg, + /*devPtrQ=*/nullptr, /*devPtrK=*/nullptr, /*devPtrV=*/nullptr, /*devPtrM=*/nullptr, + /*devPtrO=*/nullptr, /*devPtrdO=*/nullptr, /*devPtrSoftmaxOffset=*/nullptr, + /*devPtrdQ=*/nullptr, /*devPtrdK=*/nullptr, /*devPtrdV=*/nullptr, + /*devPtrdSoftmaxOffset=*/nullptr, /*devPtrDescaleQ=*/nullptr, + /*devPtrDescaleK=*/nullptr, /*devPtrDescaleV=*/nullptr, /*devPtrDescaleO=*/nullptr, + /*devPtrDescaledO=*/nullptr, /*devPtrDescaleS=*/nullptr, /*devPtrDescaledP=*/nullptr, + /*devPtrScaleS=*/nullptr, /*devPtrScaledP=*/nullptr, /*devPtrScaledQ=*/nullptr, + /*devPtrScaledK=*/nullptr, /*devPtrScaledV=*/nullptr, /*devPtrAmaxdP=*/nullptr, + /*devPtrAmaxdQ=*/nullptr, /*devPtrAmaxdK=*/nullptr, /*devPtrAmaxdV=*/nullptr, + /*devPtrQ_t=*/nullptr, /*devPtrK_t=*/nullptr, /*devPtrdO_f16=*/nullptr, + /*devPtrdO_t=*/nullptr, /*devPtrDescaleQ_t=*/nullptr, /*devPtrDescaleK_t=*/nullptr, + /*devPtrDescaledO_t=*/nullptr, /*devPtrcuSeqlensQ=*/nullptr, + /*devPtrcuSeqlensKV=*/nullptr, /*devPtrDropoutSeed=*/nullptr, + /*devPtrDropoutOffset=*/nullptr, + /*workspace=*/nullptr, &workspace_size, + /*stream=*/static_cast(0), handle); + return ""; + } catch (const std::exception& e) { + return e.what(); + } catch (...) { + return "is_supported_fp8_bwd: unknown failure."; + } +} + } // 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..79b279a833 100644 --- a/transformer_engine/common/fused_attn/fused_attn_fp8.h +++ b/transformer_engine/common/fused_attn/fused_attn_fp8.h @@ -8,35 +8,44 @@ * \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 "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); + +// check if a given configuration is supported for FP8 forward; +// if it is, cache the graph built for this config, and return an empty string; +// if not, return a diagnostic message explaining why it is not supported. +std::string is_supported_fp8_fwd(const fused_attn::FusedAttnConfig &cfg, cudnnHandle_t handle); + +// check if a given configuration is supported for FP8 backward; +// if it is, cache the graph built for this config, and return an empty string; +// if not, return a diagnostic message explaining why it is not supported. +std::string is_supported_fp8_bwd(const fused_attn::FusedAttnConfig &cfg, 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_debug.h b/transformer_engine/common/fused_attn/graph_cache_debug.h new file mode 100644 index 0000000000..033edf2ead --- /dev/null +++ b/transformer_engine/common/fused_attn/graph_cache_debug.h @@ -0,0 +1,235 @@ +/************************************************************************* + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * + * See LICENSE for license information. + ************************************************************************/ + +// ============================================================================ +// Fused-attention graph cache diagnostics. +// +// Enable at runtime with NVTE_FUSED_ATTN_CACHE_DEBUG=1 to get the cache event +// counters and graph build timings, to help diagnose redundant graph rebuilds +// or stale-cache reuse, and to profile graph-build cost. +// ============================================================================ + +#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 "config_and_params.h" + +namespace transformer_engine { +namespace fused_attn { +namespace graph_cache_debug { + +// Enable diagnostics with NVTE_FUSED_ATTN_CACHE_DEBUG=1. Single read at startup, cached. +// Negligible overhead when unset. +inline bool enabled() { + static const bool on = [] { + const char *e = std::getenv("NVTE_FUSED_ATTN_CACHE_DEBUG"); + return e != nullptr && e[0] != '\0' && e[0] != '0'; + }(); + return on; +} + +// More readable, shorter thread IDs (0, 1, 2, ...). +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 overall event counters and +// graph build timings. +inline void register_summary_once(); + +// ============================================================================ +// Cache event counters (forward/backward): +// - BUILD: a successful graph build; triggered by a cache miss +// - EXEC: a graph execution call with valid runtime tensors +// - HIT: a cache lookup that hit; may not trigger an EXEC, and may only be +// a backend availability check or from the first shape-probing call of +// nvte_fused_attn_fwd/bwd which has no runtime tensors +// - MISS: a cache lookup that missed; triggers a graph build +// ============================================================================ + +struct EventCounters { + std::atomic built{0}; + std::atomic exec{0}; + std::atomic hit{0}; + std::atomic miss{0}; +}; + +inline EventCounters &counters(bool is_fwd) { + static EventCounters fwd; + static EventCounters bwd; + return is_fwd ? fwd : bwd; +} + +inline void print_counters(const char *event) { + const EventCounters &f = counters(/*is_fwd=*/true); + const EventCounters &b = counters(/*is_fwd=*/false); + std::fprintf(stderr, + "[FUSED-ATTN-CACHE] %-10s | tid=%u | fwd built=%" PRIu64 " exec=%" PRIu64 + " hit=%" PRIu64 " miss=%" PRIu64 " | bwd built=%" PRIu64 " exec=%" PRIu64 + " hit=%" PRIu64 " miss=%" PRIu64 "\n", + event, thread_seq_id(), f.built.load(std::memory_order_relaxed), + f.exec.load(std::memory_order_relaxed), f.hit.load(std::memory_order_relaxed), + f.miss.load(std::memory_order_relaxed), b.built.load(std::memory_order_relaxed), + b.exec.load(std::memory_order_relaxed), b.hit.load(std::memory_order_relaxed), + b.miss.load(std::memory_order_relaxed)); + std::fflush(stderr); +} + +inline void record_build(const char *pass) { + if (!enabled()) return; + register_summary_once(); + const bool is_fwd = std::strcmp(pass, "fwd") == 0; + counters(is_fwd).built.fetch_add(1, std::memory_order_relaxed); + print_counters(is_fwd ? "fwd BUILD" : "bwd BUILD"); +} + +inline void record_exec(const char *pass) { + if (!enabled()) return; + register_summary_once(); + const bool is_fwd = std::strcmp(pass, "fwd") == 0; + counters(is_fwd).exec.fetch_add(1, std::memory_order_relaxed); + print_counters(is_fwd ? "fwd EXEC" : "bwd EXEC"); +} + +inline void record_cache_lookup(const char *pass, bool hit, const FusedAttnConfig &c) { + if (!enabled()) return; + register_summary_once(); + EventCounters &pc = counters(std::strcmp(pass, "fwd") == 0); + (hit ? pc.hit : pc.miss).fetch_add(1, std::memory_order_relaxed); + std::fprintf( + stderr, + "[FUSED-ATTN-CACHE] %-3s %-4s | tid=%u | train=%d det=%d cg=%d maxlogit=%d fwd=%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 "\n", + pass, hit ? "HIT" : "MISS", thread_seq_id(), static_cast(c.is_training), + static_cast(c.deterministic), static_cast(c.cuda_graph), + static_cast(c.return_max_logit), static_cast(c.is_forward), + static_cast(c.attn_mask_type), static_cast(c.bias_type), + static_cast(c.window_size_left), static_cast(c.window_size_right), + static_cast(c.bottom_right_diagonal), static_cast(c.softmax_type), + static_cast(c.scaling_mode), static_cast(c.dropout), + static_cast(c.attn_scale), static_cast(c.qkv_dtype), + static_cast(c.o_dtype), static_cast(c.do_dtype), + static_cast(c.dqkv_dtype), static_cast(c.qkv_layout), + static_cast(c.o_format), static_cast(c.do_format), + static_cast(c.dqkv_layout), static_cast(c.qkv_scale_inv_format), + static_cast(c.do_scale_inv_format), static_cast(c.batch_size), + static_cast(c.num_attn_heads), static_cast(c.num_gqa_groups), + static_cast(c.head_dim_qk), static_cast(c.head_dim_v), + static_cast(c.max_seqlen_q), static_cast(c.max_seqlen_kv), + static_cast(c.num_tokens_q), static_cast(c.num_tokens_kv), + static_cast(c.bucketed_batch_size), static_cast(c.bucketed_num_tokens_q), + static_cast(c.bucketed_num_tokens_kv), static_cast(c.num_pages_k), + static_cast(c.num_pages_v), static_cast(c.page_size_k), + static_cast(c.page_size_v), static_cast(c.max_pages_per_seq_k), + static_cast(c.max_pages_per_seq_v), static_cast(c.bias_batch_size), + static_cast(c.bias_num_heads), static_cast(c.bias_seqlen_q), + static_cast(c.bias_seqlen_kv)); + std::fflush(stderr); +} + +// ============================================================================ +// Graph build timings for individual cuDNN-frontend calls in forward/backward: +// e.g. `validate`, `build_operation_graph`, `create_execution_plans`, +// `check_support`, `build_plans` +// ============================================================================ + +enum class BuildStage { Validate, BuildOpGraph, CreatePlans, CheckSupport, BuildPlans, kCount }; +inline constexpr const char *kStageNames[] = { + "validate", "build_operation_graph", "create_execution_plans", "check_support", "build_plans"}; +struct StageTiming { + std::atomic calls{0}; + std::atomic time_ns{0}; +}; +constexpr size_t kStageBuckets = 2 * static_cast(BuildStage::kCount); +inline StageTiming &stage_timing(bool is_fwd, BuildStage s) { + static std::array table{}; + const size_t idx = + (is_fwd ? 0u : 1u) * static_cast(BuildStage::kCount) + static_cast(s); + return table[idx]; +} + +struct ScopedBuildTimer { + BuildStage stage; + bool on; + bool is_fwd; + std::chrono::steady_clock::time_point start; + ScopedBuildTimer(bool is_fwd_, BuildStage s) : stage(s), on(enabled()), is_fwd(is_fwd_) { + 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(is_fwd, stage); + t.time_ns.fetch_add(elapsed_ns, std::memory_order_relaxed); + t.calls.fetch_add(1, std::memory_order_relaxed); + } +}; + +template +inline void timer(const char *pass, BuildStage stage, Fn &&fn) { + ScopedBuildTimer scoped(std::strcmp(pass, "fwd") == 0, stage); + fn(); +} + +// ============================================================================ +// Summary: on process exit, print cache event counters and graph build timings. +// ============================================================================ +inline void register_summary_once() { + static const bool registered = [] { + std::atexit([] { + if (!enabled()) return; + print_counters("SUMMARY"); + for (int p = 0; p < 2; ++p) { + const bool is_fwd = (p == 0); + const char *pass = is_fwd ? "fwd" : "bwd"; + for (int i = 0; i < static_cast(BuildStage::kCount); ++i) { + const BuildStage s = static_cast(i); + const StageTiming &t = stage_timing(is_fwd, s); + 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; + std::fprintf(stderr, + "[FUSED-ATTN-CACHE] %-3s %-22s | calls=%" PRIu64 + " | time=%9.1f ms | avg=%9.3f ms/call\n", + pass, kStageNames[i], n, total_ms, total_ms / n); + } + } + std::fflush(stderr); + }); + return true; + }(); + (void)registered; +} + +} // 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..e240e2a421 100644 --- a/transformer_engine/common/fused_attn/utils.h +++ b/transformer_engine/common/fused_attn/utils.h @@ -4,12 +4,8 @@ * 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 @@ -223,92 +219,6 @@ 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 @@ -394,4 +304,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..938fa1747e 100644 --- a/transformer_engine/common/include/transformer_engine/fused_attn.h +++ b/transformer_engine/common/include/transformer_engine/fused_attn.h @@ -8,8 +8,10 @@ * \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 #include "stdint.h" #include "transformer_engine.h" @@ -194,6 +196,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, + kNVTEFusedAttnFwdParamsCuSeqlensQ, + kNVTEFusedAttnFwdParamsCuSeqlensKV, + kNVTEFusedAttnFwdParamsCuSeqlensQPadded, + kNVTEFusedAttnFwdParamsCuSeqlensKVPadded, + kNVTEFusedAttnFwdParamsPageTableK, + kNVTEFusedAttnFwdParamsPageTableV, + kNVTEFusedAttnFwdParamsRngState, + kNVTEFusedAttnFwdParamsS, + kNVTEFusedAttnFwdParamsO, + kNVTEFusedAttnFwdParamsAuxCtxTensors, + // configuration knobs + kNVTEFusedAttnFwdParamsIsTraining, + kNVTEFusedAttnFwdParamsCudaGraph, + kNVTEFusedAttnFwdParamsReturnMaxLogit, + kNVTEFusedAttnFwdParamsAttnMaskType, + kNVTEFusedAttnFwdParamsBiasType, + kNVTEFusedAttnFwdParamsSoftmaxType, + kNVTEFusedAttnFwdParamsWindowSizeLeft, + kNVTEFusedAttnFwdParamsWindowSizeRight, + kNVTEFusedAttnFwdParamsBottomRightDiagonal, + kNVTEFusedAttnFwdParamsDropout, + kNVTEFusedAttnFwdParamsAttnScale, + kNVTEFusedAttnFwdParamsQKVLayout, + kNVTEFusedAttnFwdParamsOFormat, + kNVTEFusedAttnFwdParamsQKVScaleInvFormat, + kNVTEFusedAttnFwdParamsMaxSeqlenQ, + kNVTEFusedAttnFwdParamsMaxSeqlenKV, + // workspace and stream + kNVTEFusedAttnFwdParamsWorkspace, + kNVTEFusedAttnFwdParamsStream, + // number of attributes + kNVTEFusedAttnFwdParamsNumAttributes +}; + +/*! \brief Create a fused-attention forward-parameter object. */ +NVTEFusedAttnFwdParams nvte_create_fused_attn_fwd_params(void); + +/*! \brief Destroy a fused-attention forward-parameter object. */ +void nvte_destroy_fused_attn_fwd_params(NVTEFusedAttnFwdParams params); + +/*! \brief Query an attribute in a fused-attention forward-parameter object. */ +void nvte_get_fused_attn_fwd_params_attribute(NVTEFusedAttnFwdParams params, + NVTEFusedAttnFwdParamsAttribute attr, void *buf, + size_t size_in_bytes, size_t *size_written); + +/*! \brief Set an attribute in a fused-attention forward-parameter object. */ +void nvte_set_fused_attn_fwd_params_attribute(NVTEFusedAttnFwdParams params, + NVTEFusedAttnFwdParamsAttribute attr, const void *buf, + size_t size_in_bytes); + +/*! \brief Opaque fused-attention backward-parameter handle. */ +typedef void *NVTEFusedAttnBwdParams; + +/*! \enum NVTEFusedAttnBwdParamsAttribute + * \brief Attributes for ``NVTEFusedAttnBwdParams``. + * + * This enum is used to index the ``FusedAttnBwdParams`` struct. The order of its fields 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 + kNVTEFusedAttnBwdParamsCudaGraph, + kNVTEFusedAttnBwdParamsDeterministic, + kNVTEFusedAttnBwdParamsAttnMaskType, + kNVTEFusedAttnBwdParamsBiasType, + kNVTEFusedAttnBwdParamsSoftmaxType, + kNVTEFusedAttnBwdParamsWindowSizeLeft, + kNVTEFusedAttnBwdParamsWindowSizeRight, + kNVTEFusedAttnBwdParamsBottomRightDiagonal, + kNVTEFusedAttnBwdParamsDropout, + kNVTEFusedAttnBwdParamsAttnScale, + kNVTEFusedAttnBwdParamsQKVLayout, + kNVTEFusedAttnBwdParamsOFormat, + kNVTEFusedAttnBwdParamsDOFormat, + kNVTEFusedAttnBwdParamsDQKVLayout, + kNVTEFusedAttnBwdParamsQKVScaleInvFormat, + kNVTEFusedAttnBwdParamsDOScaleInvFormat, + kNVTEFusedAttnBwdParamsMaxSeqlenQ, + kNVTEFusedAttnBwdParamsMaxSeqlenKV, + // workspace and stream + kNVTEFusedAttnBwdParamsWorkspace, + kNVTEFusedAttnBwdParamsStream, + // number of attributes + kNVTEFusedAttnBwdParamsNumAttributes +}; + +/*! \brief Create a fused-attention backward-parameter object. */ +NVTEFusedAttnBwdParams nvte_create_fused_attn_bwd_params(void); + +/*! \brief Destroy a fused-attention backward-parameter object. */ +void nvte_destroy_fused_attn_bwd_params(NVTEFusedAttnBwdParams params); + +/*! \brief Query an attribute in a fused-attention backward-parameter object. */ +void nvte_get_fused_attn_bwd_params_attribute(NVTEFusedAttnBwdParams params, + NVTEFusedAttnBwdParamsAttribute attr, void *buf, + size_t size_in_bytes, size_t *size_written); + +/*! \brief Set an attribute in a fused-attention backward-parameter object. */ +void nvte_set_fused_attn_bwd_params_attribute(NVTEFusedAttnBwdParams params, + NVTEFusedAttnBwdParamsAttribute attr, const void *buf, + size_t size_in_bytes); + +/*! \brief Get fused-attention backend based on user configuration. + * + * This function passes the user configuration to cuDNN frontend, runs its support checks, + * attempts to build the necessary graphs, and if successful, caches the graphs (if not, returns + * ``NVTE_No_Backend``). + * + * \param[in] cfg Fused-attention configuration created by + * ``nvte_create_fused_attn_config()``. + * \param[out] message If cuDNN graphs are built successfully, an empty string; + * if not, a diagnostic message explaining why there is no support. + * Pass NULL to skip the diagnostics. 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 +456,15 @@ 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 signature than nvte_get_fused_attn_backend_v2, + * and it fills the fields that it cannot express with default values. For example, it sets + * batch_size = 1, derives output/gradient formats from qkv_layout, assumes a standard + * bias shape [b, h, sq, skv] for NVTE_POST_SCALE_BIAS, uses delayed scaling for all FP8, + * and does not support paged-KV attention. Users who need more precise control should + * use nvte_get_fused_attn_backend_v2 directly. */ 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: @@ -281,6 +542,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 +560,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: @@ -354,6 +628,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, @@ -708,6 +984,488 @@ class AttentionShape { size_t canonical_[5] = {}; }; +/*! \class FusedAttnConfigWrapper + * \brief C++ helper for constructing an ``NVTEFusedAttnConfig``. + * + * It owns an opaque ``NVTEFusedAttnConfig`` handle created by + * ``nvte_create_fused_attn_config()``, and provides a convenient, + * chainable interface for setting every field in ``FusedAttnConfig``. + */ +class FusedAttnConfigWrapper { + public: + FusedAttnConfigWrapper() : cfg_{nvte_create_fused_attn_config()} {} + + FusedAttnConfigWrapper(const FusedAttnConfigWrapper &) = delete; + FusedAttnConfigWrapper &operator=(const FusedAttnConfigWrapper &) = delete; + + FusedAttnConfigWrapper(FusedAttnConfigWrapper &&other) noexcept : cfg_{other.cfg_} { + other.cfg_ = nullptr; + } + + FusedAttnConfigWrapper &operator=(FusedAttnConfigWrapper &&other) noexcept { + if (this != &other) { + 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) { + 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_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_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_is_training(bool val) noexcept { + return set_attr(kNVTEFusedAttnFwdParamsIsTraining, static_cast(val)); + } + FusedAttnFwdParamsWrapper &set_cuda_graph(bool val) noexcept { + return set_attr(kNVTEFusedAttnFwdParamsCudaGraph, static_cast(val)); + } + FusedAttnFwdParamsWrapper &set_return_max_logit(bool val) noexcept { + return set_attr(kNVTEFusedAttnFwdParamsReturnMaxLogit, static_cast(val)); + } + FusedAttnFwdParamsWrapper &set_attn_mask_type(NVTE_Mask_Type val) noexcept { + return set_attr(kNVTEFusedAttnFwdParamsAttnMaskType, val); + } + FusedAttnFwdParamsWrapper &set_bias_type(NVTE_Bias_Type val) noexcept { + return set_attr(kNVTEFusedAttnFwdParamsBiasType, val); + } + FusedAttnFwdParamsWrapper &set_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_dropout(float val) noexcept { + return set_attr(kNVTEFusedAttnFwdParamsDropout, val); + } + FusedAttnFwdParamsWrapper &set_attn_scale(float val) noexcept { + return set_attr(kNVTEFusedAttnFwdParamsAttnScale, val); + } + FusedAttnFwdParamsWrapper &set_qkv_layout(NVTE_QKV_Layout val) noexcept { + return set_attr(kNVTEFusedAttnFwdParamsQKVLayout, val); + } + FusedAttnFwdParamsWrapper &set_o_format(NVTE_QKV_Format val) noexcept { + return set_attr(kNVTEFusedAttnFwdParamsOFormat, val); + } + FusedAttnFwdParamsWrapper &set_qkv_scale_inv_format(NVTE_QKV_Format val) noexcept { + return set_attr(kNVTEFusedAttnFwdParamsQKVScaleInvFormat, val); + } + FusedAttnFwdParamsWrapper &set_max_seqlen_q(size_t val) noexcept { + return set_attr(kNVTEFusedAttnFwdParamsMaxSeqlenQ, val); + } + FusedAttnFwdParamsWrapper &set_max_seqlen_kv(size_t val) noexcept { + return set_attr(kNVTEFusedAttnFwdParamsMaxSeqlenKV, val); + } + FusedAttnFwdParamsWrapper &set_workspace(NVTETensor val) noexcept { + return set_attr(kNVTEFusedAttnFwdParamsWorkspace, val); + } + FusedAttnFwdParamsWrapper &set_stream(cudaStream_t val) noexcept { + return set_attr(kNVTEFusedAttnFwdParamsStream, val); + } + + private: + // 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) { + 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_cuda_graph(bool val) noexcept { + return set_attr(kNVTEFusedAttnBwdParamsCudaGraph, static_cast(val)); + } + FusedAttnBwdParamsWrapper &set_deterministic(bool val) noexcept { + return set_attr(kNVTEFusedAttnBwdParamsDeterministic, static_cast(val)); + } + FusedAttnBwdParamsWrapper &set_attn_mask_type(NVTE_Mask_Type val) noexcept { + return set_attr(kNVTEFusedAttnBwdParamsAttnMaskType, val); + } + FusedAttnBwdParamsWrapper &set_bias_type(NVTE_Bias_Type val) noexcept { + return set_attr(kNVTEFusedAttnBwdParamsBiasType, val); + } + FusedAttnBwdParamsWrapper &set_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_dropout(float val) noexcept { + return set_attr(kNVTEFusedAttnBwdParamsDropout, val); + } + FusedAttnBwdParamsWrapper &set_attn_scale(float val) noexcept { + return set_attr(kNVTEFusedAttnBwdParamsAttnScale, val); + } + FusedAttnBwdParamsWrapper &set_qkv_layout(NVTE_QKV_Layout val) noexcept { + return set_attr(kNVTEFusedAttnBwdParamsQKVLayout, val); + } + FusedAttnBwdParamsWrapper &set_o_format(NVTE_QKV_Format val) noexcept { + return set_attr(kNVTEFusedAttnBwdParamsOFormat, val); + } + FusedAttnBwdParamsWrapper &set_do_format(NVTE_QKV_Format val) noexcept { + return set_attr(kNVTEFusedAttnBwdParamsDOFormat, val); + } + FusedAttnBwdParamsWrapper &set_dqkv_layout(NVTE_QKV_Layout val) noexcept { + return set_attr(kNVTEFusedAttnBwdParamsDQKVLayout, val); + } + FusedAttnBwdParamsWrapper &set_qkv_scale_inv_format(NVTE_QKV_Format val) noexcept { + return set_attr(kNVTEFusedAttnBwdParamsQKVScaleInvFormat, val); + } + FusedAttnBwdParamsWrapper &set_do_scale_inv_format(NVTE_QKV_Format val) noexcept { + return set_attr(kNVTEFusedAttnBwdParamsDOScaleInvFormat, val); + } + FusedAttnBwdParamsWrapper &set_max_seqlen_q(size_t val) noexcept { + return set_attr(kNVTEFusedAttnBwdParamsMaxSeqlenQ, val); + } + FusedAttnBwdParamsWrapper &set_max_seqlen_kv(size_t val) noexcept { + return set_attr(kNVTEFusedAttnBwdParamsMaxSeqlenKV, val); + } + FusedAttnBwdParamsWrapper &set_workspace(NVTETensor val) noexcept { + return set_attr(kNVTEFusedAttnBwdParamsWorkspace, val); + } + FusedAttnBwdParamsWrapper &set_stream(cudaStream_t val) noexcept { + return set_attr(kNVTEFusedAttnBwdParamsStream, val); + } + + private: + // 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..f8cd1308cb 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,12 @@ 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 ( + JAXX_Scaling_Mode, + NVTE_Fused_Attn_Backend, + NVTE_QKV_Format, + NVTE_QKV_Layout, +) from transformer_engine.jax.attention import ( AttnBiasType, AttnMaskType, @@ -57,6 +63,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=[], @@ -108,6 +145,7 @@ class FusedAttnHelper: """ is_training: bool + batch_size: int q_dtype: jnp.dtype kv_dtype: jnp.dtype qkv_layout: QKVLayout @@ -122,21 +160,54 @@ 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( + """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) + 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( self.is_training, - jax_dtype_to_te_dtype(self.q_dtype), + self.batch_size, + q_type, jax_dtype_to_te_dtype(self.kv_dtype), + q_type, + q_type, + q_type, + JAXX_Scaling_Mode.NO_SCALING, self.qkv_layout.value, + NVTE_QKV_Format.NVTE_QKV_Format_NOT_SET, + NVTE_QKV_Format.NVTE_QKV_Format_NOT_SET, + NVTE_QKV_Layout.NVTE_QKV_Layout_NOT_SET, + NVTE_QKV_Format.NVTE_QKV_Format_NOT_SET, + NVTE_QKV_Format.NVTE_QKV_Format_NOT_SET, self.attn_bias_type.value, self.attn_mask_type.value, self.softmax_type.value, + self.attn_scale, self.dropout_probability, self.q_num_heads, self.kv_num_heads, @@ -146,9 +217,29 @@ def get_fused_attn_backend(self): self.head_dim_v, self.window_size[0], self.window_size[1], + self.bottom_right_diagonal, not self.is_non_deterministic_allowed(), + bias_batch, + bias_heads, + bias_seqlen_q, + 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 +426,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 +448,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 +472,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 b9c7c849f2..7e66b35f5b 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,17 @@ 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, +std::tuple GetFusedAttnBackend( + bool is_training, size_t batch_size, DType q_dtype, DType kv_dtype, DType o_dtype, + DType do_dtype, DType dqkv_dtype, JAXX_Scaling_Mode scaling_mode, 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, - 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); + float attn_scale, 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 bottom_right_diagonal, + bool deterministic, size_t bias_batch, size_t bias_heads, size_t bias_seqlen_q, + size_t bias_seqlen_kv); 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..30ff61b013 100644 --- a/transformer_engine/jax/csrc/extensions/attention.cpp +++ b/transformer_engine/jax/csrc/extensions/attention.cpp @@ -24,18 +24,67 @@ namespace transformer_engine { namespace jax { -NVTE_Fused_Attn_Backend GetFusedAttnBackend( - bool is_training, DType q_dtype, DType kv_dtype, NVTE_QKV_Layout qkv_layout, +std::tuple GetFusedAttnBackend( + bool is_training, size_t batch_size, DType q_dtype, DType kv_dtype, DType o_dtype, + DType do_dtype, DType dqkv_dtype, JAXX_Scaling_Mode scaling_mode, 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, - 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; + float attn_scale, 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 bottom_right_diagonal, + bool deterministic, size_t bias_batch, size_t bias_heads, size_t bias_seqlen_q, + size_t bias_seqlen_kv) { + 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; + } + NVTE_CHECK(q_dtype == kv_dtype, "Q and KV must have the same data type."); + + FusedAttnConfigWrapper cfg; + cfg.set_is_training(is_training) + .set_deterministic(deterministic) + .set_cuda_graph(false) + .set_return_max_logit(false) + .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(mask_type) + .set_softmax_type(softmax_type) + .set_scaling_mode(get_nvte_scaling_mode(scaling_mode)) + .set_attn_scale(attn_scale) + .set_dropout(dropout_probability) + .set_max_seqlen_q(q_max_seqlen) + .set_max_seqlen_kv(kv_max_seqlen) + .set_window_size_left(window_size_left) + .set_window_size_right(window_size_right) + .set_bottom_right_diagonal(bottom_right_diagonal) + .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_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_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()}; } /* @@ -195,15 +244,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 +349,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, + auto [backend, _fwd_msg] = GetFusedAttnBackend( + is_training, input_batch, dtype, dtype, dtype, dtype, dtype, JAXX_Scaling_Mode::NO_SCALING, + qkv_layout, NVTE_QKV_Format::NVTE_QKV_Format_NOT_SET, + NVTE_QKV_Format::NVTE_QKV_Format_NOT_SET, NVTE_QKV_Layout::NVTE_QKV_Layout_NOT_SET, + NVTE_QKV_Format::NVTE_QKV_Format_NOT_SET, NVTE_QKV_Format::NVTE_QKV_Format_NOT_SET, bias_type, + mask_type, softmax_type, scaling_factor, 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); + bottom_right_diagonal, deterministic, 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 +417,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 +599,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 +680,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, + auto [backend, _bwd_msg] = GetFusedAttnBackend( + is_training, input_batch, dtype, dtype, dtype, dtype, dtype, JAXX_Scaling_Mode::NO_SCALING, + qkv_layout, NVTE_QKV_Format::NVTE_QKV_Format_NOT_SET, + NVTE_QKV_Format::NVTE_QKV_Format_NOT_SET, NVTE_QKV_Layout::NVTE_QKV_Layout_NOT_SET, + NVTE_QKV_Format::NVTE_QKV_Format_NOT_SET, NVTE_QKV_Format::NVTE_QKV_Format_NOT_SET, bias_type, + mask_type, softmax_type, scaling_factor, 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); + bottom_right_diagonal, deterministic, 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 +764,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 db5468afe6..bfb7b2a826 100644 --- a/transformer_engine/jax/csrc/extensions/pybind.cpp +++ b/transformer_engine/jax/csrc/extensions/pybind.cpp @@ -200,12 +200,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 6a60160f77..9f81c78f9c 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/backends.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/backends.py @@ -1965,7 +1965,19 @@ class FusedAttention(torch.nn.Module): 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. The supported recipes are as follows. Inputs, + Intermediates, and Outputs are in the format of "tensor: quantizer", and are used by function calls, + tex.fused_attn_fwd and tex.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 """ 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 995dd2e90a..0c611051f7 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py @@ -4918,6 +4918,87 @@ 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_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): + return { + "attn_mask_type": mask, + "max_seqlen_q": s_q, + "max_seqlen_kv": s_kv, + "num_attn_heads": heads, + "num_gqa_groups": gqa, + "window_size_left": window_left, + "window_size_right": window_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, + ) + ] + + 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 + # 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) + 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 + if not is_causal: + return [config(attn_mask_type, r_q, r_kv, heads, gqa, bottom_right_diagonal)] + return [ + config(attn_mask_type, r_q, r_kv, heads, gqa, bottom_right_diagonal), # diagonal + config( + padding_or_no_mask, r_q, r_kv // 2, heads, gqa, bottom_right_diagonal + ), # lower-triangle + config( + padding_or_no_mask, r_q // 2, r_kv, heads, gqa, bottom_right_diagonal + ), # 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 4cc4cab1b8..6c347d4186 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 @@ -551,6 +551,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"))) @@ -1146,14 +1147,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 @@ -1603,11 +1603,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 @@ -1667,32 +1670,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 -- @@ -1759,22 +1741,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, @@ -1784,6 +1775,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 6c47d0f1cb..79377ea0c0 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 ( @@ -200,6 +201,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 @@ -212,6 +216,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 @@ -228,8 +236,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 @@ -242,7 +250,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 @@ -261,6 +271,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 @@ -273,12 +285,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" @@ -286,13 +301,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 @@ -302,6 +318,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 @@ -326,6 +343,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 @@ -350,33 +428,12 @@ 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_params): """Constant-foldable tex.get_fused_attn_backend: the result depends only on the attention config, and the python-side enum keeps it traceable by - torch.compile (see the FusedAttnBackend docstring). 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.""" - return FusedAttnBackend.cast( - 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, - ) - ) + torch.compile (see the FusedAttnBackend docstring).""" + fused_attention_backend, reject_message = tex.get_fused_attn_backend(fused_attn_params) + return FusedAttnBackend.cast(fused_attention_backend), reject_message def get_attention_backend( @@ -408,12 +465,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 @@ -427,7 +487,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 @@ -437,6 +498,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 @@ -1402,64 +1464,184 @@ 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"]: - 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"]: + 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_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_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 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 + fused_attn_params = FusedAttentionParams(**fused_attn_kwargs) + + # NOTE: under torch.compile the numeric fields of fused_attn_params 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_params) + 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 " "F16/BF16 arbitrary-seqlen", @@ -2379,6 +2561,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 get_qkv_layout( q: torch.Tensor, k: torch.Tensor, diff --git a/transformer_engine/pytorch/csrc/extensions.h b/transformer_engine/pytorch/csrc/extensions.h index 6edfbdc00e..83707a2e16 100644 --- a/transformer_engine/pytorch/csrc/extensions.h +++ b/transformer_engine/pytorch/csrc/extensions.h @@ -86,12 +86,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..5bfbc79d53 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_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_attn_scale(p.attr("attn_scale").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 7e9d114be8..d0e46e565f 100644 --- a/transformer_engine/pytorch/csrc/extensions/pybind.cpp +++ b/transformer_engine/pytorch/csrc/extensions/pybind.cpp @@ -411,7 +411,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());