From c105f9d44875d70d6a936fed46b8090261b72ee9 Mon Sep 17 00:00:00 2001 From: Vladimir Cherepanov Date: Sun, 11 Oct 2026 21:09:30 +0000 Subject: [PATCH 1/6] Move JAX cuDNN attention to Python graphs and add FP8 support Build and select cuDNN attention plans through the Python frontend for ordinary and score-modified attention. Share graph serialization and a JAX-owned C++ FFI executor with per-device plan and handle caches and CUDA command-buffer support. Remove JAX's calls to TE common attention APIs while leaving TE common and PyTorch unchanged. Expose current tensor-scaled FP8 forward and backward through fused_attn and Flax attention modules, retaining FP16/BF16 inputs, outputs, and input gradients. Add ALiBi and preserve packed and ragged layouts, GQA, masks, sliding windows, dropout replay, attention sinks, max logits, and Ring and AllGather context parallelism. Preserve neutral statistics for empty ragged ring blocks and differentiation with is_training=False. Document FP8 availability, scaling, remaining feature gaps, and the NVTE_FP8_DPA_BWD switch. Add independent numerical references and focused regression coverage for the new precision and execution paths. Validate with the editable JAX build, the requested attention suites (1,241 passed, 3,901 skipped), and additional FP8, Flax, deterministic, two-GPU context-parallel, and CUDA command-buffer checks. JAX L0 lint, pre-commit, and changed-file license checks pass. Signed-off-by: Vladimir Cherepanov --- docs/api/jax.rst | 6 + docs/envvars.rst | 2 +- docs/examples/jax/attention.rst | 36 + tests/jax/test_fused_attn.py | 446 ++++++ tests/jax/test_fused_attn_score_mod.py | 5 +- transformer_engine/jax/attention.py | 28 +- .../jax/cpp_extensions/attention.py | 404 ++---- .../jax/cpp_extensions/cudnn_attention.py | 511 +++++++ .../jax/cpp_extensions/cudnn_graph.py | 139 ++ .../jax/cpp_extensions/flex_attention.py | 129 +- transformer_engine/jax/csrc/extensions.h | 28 +- .../jax/csrc/extensions/attention.cpp | 1223 ++--------------- .../jax/csrc/extensions/pybind.cpp | 25 +- transformer_engine/jax/flax/transformer.py | 35 +- 14 files changed, 1478 insertions(+), 1539 deletions(-) create mode 100644 transformer_engine/jax/cpp_extensions/cudnn_attention.py create mode 100644 transformer_engine/jax/cpp_extensions/cudnn_graph.py diff --git a/docs/api/jax.rst b/docs/api/jax.rst index 7a31c9d3797..130e66bd5c1 100644 --- a/docs/api/jax.rst +++ b/docs/api/jax.rst @@ -59,3 +59,9 @@ Modules :members: __call__ .. autoapifunction:: transformer_engine.jax.flax.extend_logical_axis_rules + +Attention functions +------------------- + +.. autoapifunction:: transformer_engine.jax.attention.fused_attn +.. autoapifunction:: transformer_engine.jax.attention.is_fused_attn_kernel_available diff --git a/docs/envvars.rst b/docs/envvars.rst index 62b8ac5f4f0..511a59c7917 100644 --- a/docs/envvars.rst +++ b/docs/envvars.rst @@ -271,7 +271,7 @@ FP8 Configuration :Type: ``int`` (0 or 1) :Default: ``1`` - :Description: Enable FP8 in the backward pass of DotProductAttention. ``1`` = FP8 forward and backward, ``0`` = FP8 forward and FP16/BF16 backward. + :Description: Enable FP8 in the backward pass of DotProductAttention in PyTorch and JAX. ``1`` = FP8 forward and backward, ``0`` = FP8 forward and FP16/BF16 backward. In JAX this applies when attention is configured with ``fp8=True``. .. envvar:: NVTE_DPA_FP8CS_O_in_F16 diff --git a/docs/examples/jax/attention.rst b/docs/examples/jax/attention.rst index f82a6ffe621..23f52b105d2 100644 --- a/docs/examples/jax/attention.rst +++ b/docs/examples/jax/attention.rst @@ -16,6 +16,42 @@ FlexAttention-style customization. The API also supports different Q/K and V head dimensions, as used by the attention operation after the projection stages in DeepSeek-style MLA. +cuDNN execution and FP8 +---------------------- + +JAX fused attention constructs and selects cuDNN execution plans with the +``nvidia-cudnn-frontend`` Python package. The compiled JAX FFI executor caches +those plans and its own cuDNN handles per device. Ordinary attention and +``score_mod`` share this executor; neither calls TE common's attention APIs. +The Python frontend and the headers used to build the JAX extension must have +the same version, since execution plans are serialized between the two. + +Pass ``fp8=True`` to ``transformer_engine.jax.attention.fused_attn`` or to the +Flax ``DotProductAttention``, ``MultiHeadAttention`` or ``TransformerLayer`` +module to enable cuDNN FP8 attention. Inputs remain FP16 or BF16. Query, key +and value are quantized to E4M3 using their current tensor maximum; backward +quantizes the output gradient to E5M2. Outputs and input gradients retain the +input dtype, and softmax statistics use FP32. Packed layouts use one scale per +packed tensor. This option is independent of the GEMM quantization recipe and +does not maintain delayed-scaling history. + +``NVTE_FP8_DPA_BWD=0`` selects FP16/BF16 backward while keeping FP8 forward. +The backward pass uses dequantized forward operands, so both passes operate on +the same quantized query, key and value. The default, ``1``, uses FP8 for both. + +FP8 availability depends on cuDNN, GPU architecture, layout and the selected +features. Unsupported combinations produce a diagnostic rather than silently +using a different precision. For example, cuDNN 9.26 rejects FP8 backward with +dropout on Blackwell; use ``NVTE_FP8_DPA_BWD=0`` for that combination. +Tensor-scaled FP8 is supported; MXFP8 attention, +paged KV-cache APIs and quantized module inputs/outputs require further JAX API +and frontend-binding work. ``score_mod`` currently supports FP16/BF16 only. + +``AttnBiasType.ALIBI`` (``attn_bias_type="alibi"`` in Flax) enables cuDNN's +ALiBi mask without an explicit bias tensor. Explicit dropout keys have JAX's +usual replay semantics: using the same key repeats the mask, including when +recomputing a forward pass for its backward pass. + For long contexts, selected BSHD and THD configurations support context parallelism with Ring or AllGather collectives. Exact fused-kernel availability depends on the input shape, dtype, GPU architecture, and feature combination; diff --git a/tests/jax/test_fused_attn.py b/tests/jax/test_fused_attn.py index 4e5c2948a21..d8d6c088195 100644 --- a/tests/jax/test_fused_attn.py +++ b/tests/jax/test_fused_attn.py @@ -2106,3 +2106,449 @@ def test_backward( seq_desc_format, is_self_attention=None, ) + + +@pytest.mark.parametrize("dtype", [jnp.float16, jnp.bfloat16]) +@pytest.mark.parametrize("fp8_backward", [False, True]) +@pytest.mark.parametrize( + "layout", + [ + QKVLayout.BS3HD, + QKVLayout.BSHD_BS2HD, + QKVLayout.BSHD_BSHD_BSHD, + ], +) +def test_cudnn_python_fp8_forward_backward(dtype, layout, fp8_backward, monkeypatch): + """Exercise real FP8 SDPA with exactly representable operands and probabilities.""" + monkeypatch.setenv("NVTE_FP8_DPA_BWD", "1" if fp8_backward else "0") + q = jnp.full((2, 128, 4, 64), 0.5, dtype) + k = jnp.full_like(q, 0.5) + v = jnp.zeros_like(q).at[:, ::2, :, 0].set(1.0) + if layout.is_qkvpacked(): + operands = (jnp.stack((q, k, v), axis=-3),) + elif layout.is_kvpacked(): + operands = (q, jnp.stack((k, v), axis=-3)) + else: + operands = (q, k, v) + + def unpack(operands): + if layout.is_qkvpacked(): + return tuple(operands[0][:, :, index] for index in range(3)) + if layout.is_kvpacked(): + return operands[0], operands[1][:, :, 0], operands[1][:, :, 1] + return operands + + def reference(operands): + query, key, value = (tensor.astype(jnp.float32) for tensor in unpack(operands)) + scores = jnp.einsum("bqhd,bkhd->bhqk", query, key) * 0.125 + return jnp.einsum("bhqk,bkhd->bqhd", jax.nn.softmax(scores, axis=-1), value) + + def attention(operands): + return fused_attn( + operands, + None, + None, + None, + attn_bias_type=AttnBiasType.NO_BIAS, + attn_mask_type=AttnMaskType.NO_MASK, + qkv_layout=layout, + softmax_type=AttnSoftmaxType.VANILLA_SOFTMAX, + scaling_factor=0.125, + dropout_probability=0.0, + is_training=True, + fp8=True, + ) + + if get_device_compute_capability(0) < 90: + pytest.skip("cuDNN FP8 SDPA requires Hopper or newer.") + actual = jax.jit(attention)(operands) + expected = jax.jit(reference)(operands) + np.testing.assert_allclose(actual, expected, rtol=5e-3, atol=5e-3) + actual_grads = jax.jit(jax.grad(lambda arrays: jnp.sum(attention(arrays))))(operands) + expected_grads = jax.jit(jax.grad(lambda arrays: jnp.sum(reference(arrays))))(operands) + for actual_grad, expected_grad in zip(actual_grads, expected_grads): + assert actual_grad.dtype == dtype + np.testing.assert_allclose(actual_grad, expected_grad, rtol=1e-2, atol=1e-2) + + +def test_cudnn_python_attention_does_not_call_common(monkeypatch): + """Backend selection, workspace sizing and execution are independent of common.""" + import transformer_engine_jax + from transformer_engine.jax.cpp_extensions import cudnn_attention + + def forbidden(*_args, **_kwargs): + raise AssertionError("JAX attention called TE common") + + for name in ( + "get_fused_attn_backend", + "get_fused_attn_fwd_workspace_sizes", + "get_fused_attn_bwd_workspace_sizes", + "nvte_get_qkv_format", + ): + monkeypatch.setattr(transformer_engine_jax, name, forbidden, raising=False) + cudnn_attention._build_graph.cache_clear() + tensor = jnp.ones((1, 128, 2, 64), jnp.bfloat16) + output = jax.jit( + lambda q: fused_attn( + (q, q, q), + None, + None, + None, + attn_bias_type=AttnBiasType.NO_BIAS, + attn_mask_type=AttnMaskType.NO_MASK, + qkv_layout=QKVLayout.BSHD_BSHD_BSHD, + softmax_type=AttnSoftmaxType.VANILLA_SOFTMAX, + scaling_factor=0.125, + dropout_probability=0.0, + is_training=True, + ) + )(tensor) + np.testing.assert_array_equal(output, tensor) + gradients = jax.jit( + jax.grad( + lambda q: jnp.sum( + fused_attn( + (q, q, q), + None, + None, + None, + attn_bias_type=AttnBiasType.NO_BIAS, + attn_mask_type=AttnMaskType.NO_MASK, + qkv_layout=QKVLayout.BSHD_BSHD_BSHD, + softmax_type=AttnSoftmaxType.VANILLA_SOFTMAX, + scaling_factor=0.125, + dropout_probability=0.0, + is_training=True, + ) + ) + ) + )(tensor) + np.testing.assert_array_equal(gradients, tensor) + + +@pytest.mark.parametrize( + "kind", + [ + "dense", + "causal", + "padding", + "window", + "ragged", + "ragged_dense_stats", + "ragged_dense_stats_f16", + ], +) +def test_cudnn_python_fp8_gqa_masks(kind, monkeypatch, request): + """Compare non-uniform FP8 attention and gradients against native FP32 GQA.""" + ragged = kind.startswith("ragged") + fp8 = not kind.endswith("_f16") + if get_device_compute_capability(0) < (100 if ragged else 90): + pytest.skip("FP8 ragged backward requires Blackwell; dense FP8 requires Hopper.") + if kind.startswith("ragged_dense_stats"): + from transformer_engine.jax.cpp_extensions import cudnn_attention + + # Test the dense statistics conversion on the available GPU. + # Architecture is fixed in a real process, so discard JAX's abstract-value + # cache when mocking a different architecture's Max tensor layout. + jax.clear_caches() + request.addfinalizer(jax.clear_caches) + monkeypatch.setattr(cudnn_attention, "device_capabilities", lambda: (120,)) + keys = jax.random.split(jax.random.PRNGKey(123), 4) + q = jax.random.choice(keys[0], jnp.array([-0.5, 0.5], jnp.bfloat16), (2, 128, 4, 64)) + k = jax.random.choice(keys[1], jnp.array([-0.5, 0.5], jnp.bfloat16), (2, 128, 2, 64)) + v = jax.random.choice(keys[2], jnp.array([-1.0, 1.0], jnp.bfloat16), (2, 128, 2, 64)) + weight = jax.random.choice(keys[3], jnp.array([-1.0, 1.0], jnp.bfloat16), q.shape) + positions = jnp.arange(128) + mask = jnp.ones((2, 128, 128), jnp.bool_) + descriptor = None + attn_mask = AttnMaskType.NO_MASK + layout = QKVLayout.BSHD_BSHD_BSHD + window = None + if kind in ("causal", "window") or ragged: + mask &= positions[None, :, None] >= positions[None, None, :] + attn_mask = AttnMaskType.CAUSAL_MASK + if kind == "window": + window = (15, 0) + mask &= positions[None, :, None] - positions[None, None, :] <= 15 + if kind == "padding": + lengths = jnp.array([128, 91], jnp.int32) + descriptor = SequenceDescriptor.from_seqlens((lengths, lengths)) + mask &= positions[None, None, :] < lengths[:, None, None] + attn_mask = AttnMaskType.PADDING_MASK + if ragged: + ids = jnp.broadcast_to(positions // 64 + 1, (2, 128)) + descriptor = SequenceDescriptor.from_segment_ids_and_pos( + ids, jnp.broadcast_to(positions % 64, ids.shape) + ) + mask &= ids[:, :, None] == ids[:, None, :] + attn_mask = AttnMaskType.PADDING_CAUSAL_MASK + layout = QKVLayout.THD_THD_THD + + def reference(arrays): + query, key, value = (tensor.astype(jnp.float32) for tensor in arrays) + key, value = (jnp.repeat(tensor, 2, axis=-2) for tensor in (key, value)) + scores = jnp.einsum("bqhd,bkhd->bhqk", query, key) * 0.125 + scores = jnp.where(mask[:, None], scores, -jnp.inf) + result = jnp.einsum("bhqk,bkhd->bqhd", jax.nn.softmax(scores, axis=-1), value) + if kind == "padding": + result = jnp.where( + (positions[None, :] < lengths[:, None])[:, :, None, None], result, 0.0 + ) + return result + + def attention(arrays): + result = fused_attn( + arrays, + None, + descriptor, + None, + attn_bias_type=AttnBiasType.NO_BIAS, + attn_mask_type=attn_mask, + qkv_layout=layout, + softmax_type=AttnSoftmaxType.VANILLA_SOFTMAX, + scaling_factor=0.125, + dropout_probability=0.0, + is_training=True, + max_segments_per_seq=2 if ragged else 1, + window_size=window, + fp8=fp8, + return_max_logit=True, + ) + return result[0] + + arrays = (q, k, v) + actual, expected = jax.jit(attention)(arrays), jax.jit(reference)(arrays) + + # Operands are exactly representable. Error comes from FP8 softmax and dP, + # which use E4M3 and E5M2 internally, and the BF16 returned tensors. + def relative_error(actual, expected): + return float( + jnp.linalg.norm(actual.astype(jnp.float32) - expected.astype(jnp.float32)) + / jnp.linalg.norm(expected.astype(jnp.float32)) + ) + + assert relative_error(actual, expected) < (0.06 if fp8 else 0.01) + actual_grads = jax.jit(jax.grad(lambda xs: jnp.sum(attention(xs) * weight)))(arrays) + expected_grads = jax.jit(jax.grad(lambda xs: jnp.sum(reference(xs) * weight)))(arrays) + for actual_grad, expected_grad in zip(actual_grads, expected_grads): + assert relative_error(actual_grad, expected_grad) < (0.1 if fp8 else 0.01) + _, max_logit = jax.jit( + lambda xs: fused_attn( + xs, + None, + descriptor, + None, + attn_bias_type=AttnBiasType.NO_BIAS, + attn_mask_type=attn_mask, + qkv_layout=layout, + softmax_type=AttnSoftmaxType.VANILLA_SOFTMAX, + scaling_factor=0.125, + dropout_probability=0.0, + is_training=True, + max_segments_per_seq=2 if ragged else 1, + return_max_logit=True, + window_size=window, + fp8=fp8, + ) + )(arrays) + scores = ( + jnp.einsum( + "bqhd,bkhd->bhqk", + q.astype(jnp.float32), + jnp.repeat(k.astype(jnp.float32), 2, axis=-2), + ) + * 0.125 + ) + max_mask = mask + if kind == "padding": + max_mask &= positions[None, :, None] < lengths[:, None, None] + expected_max = jnp.max(jnp.where(max_mask[:, None], scores, -jnp.inf), axis=(0, 2, 3)) + np.testing.assert_allclose(max_logit, expected_max, rtol=1e-2, atol=1e-2) + + +def test_cudnn_python_alibi_forward_backward(): + """cuDNN ALiBi agrees with explicit slopes in native JAX attention.""" + keys = jax.random.split(jax.random.PRNGKey(42), 3) + arrays = tuple(jax.random.normal(key, (1, 128, 4, 64), jnp.float16) * 0.1 for key in keys) + positions = jnp.arange(128) + causal = positions[:, None] >= positions[None, :] + slopes = 2.0 ** (-2.0 * jnp.arange(1, 5, dtype=jnp.float32)) + + def reference(arrays): + q, k, v = (tensor.astype(jnp.float32) for tensor in arrays) + scores = jnp.einsum("bqhd,bkhd->bhqk", q, k) * 0.125 + scores += slopes[None, :, None, None] * positions[None, None, None, :] + scores = jnp.where(causal[None, None], scores, -jnp.inf) + return jnp.einsum("bhqk,bkhd->bqhd", jax.nn.softmax(scores, axis=-1), v) + + def attention(arrays): + return fused_attn( + arrays, + None, + None, + None, + attn_bias_type=AttnBiasType.ALIBI, + attn_mask_type=AttnMaskType.CAUSAL_MASK, + qkv_layout=QKVLayout.BSHD_BSHD_BSHD, + softmax_type=AttnSoftmaxType.VANILLA_SOFTMAX, + scaling_factor=0.125, + dropout_probability=0.0, + is_training=True, + ) + + np.testing.assert_allclose( + jax.jit(attention)(arrays), jax.jit(reference)(arrays), rtol=1e-2, atol=1e-3 + ) + actual_grads = jax.jit(jax.grad(lambda xs: jnp.sum(attention(xs))))(arrays) + expected_grads = jax.jit(jax.grad(lambda xs: jnp.sum(reference(xs))))(arrays) + for actual, expected in zip(actual_grads, expected_grads): + np.testing.assert_allclose(actual, expected, rtol=1e-2, atol=1e-3) + + +@pytest.mark.parametrize("softmax", ["vanilla", "learnable"]) +@pytest.mark.parametrize("deterministic", [False, True]) +@pytest.mark.parametrize("fp8", [False, True]) +def test_cudnn_python_flax_forward_backward(softmax, deterministic, fp8, monkeypatch): + """Flax attention preserves sink gradients even when dropout is disabled.""" + from transformer_engine.jax.flax import DotProductAttention + + if fp8 and get_device_compute_capability(0) < 90: + pytest.skip("cuDNN FP8 SDPA requires Hopper or newer.") + monkeypatch.setenv("NVTE_FUSED_ATTN", "1") + module = DotProductAttention( + head_dim=64, + num_attention_heads=2, + attn_mask_type="no_mask", + transpose_batch_sequence=False, + softmax_type=softmax, + fp8=fp8, + ) + q = jnp.zeros((1, 128, 2, 64), jnp.bfloat16) + v = jnp.zeros_like(q).at[:, ::2, :, 0].set(1.0) + variables = module.init(jax.random.PRNGKey(0), q, q, v, deterministic=deterministic) + + def forward(variables, value): + return module.apply(variables, q, q, value, deterministic=deterministic) + + actual = jax.jit(forward)(variables, v) + denominator = 129 if softmax == "learnable" else 128 + expected = jnp.zeros_like(q).at[:, :, :, 0].set(64 / denominator) + np.testing.assert_allclose(actual, expected, rtol=5e-3, atol=5e-3) + parameter_grad, value_grad = jax.jit( + jax.grad( + lambda params, value: jnp.sum(forward(params, value)), + argnums=(0, 1), + ) + )(variables, v) + assert value_grad.dtype == jnp.bfloat16 + np.testing.assert_allclose(value_grad, 128 / denominator, rtol=5e-2, atol=1e-2) + if softmax == "learnable": + expected_sink_grad = -128 * 64 / denominator**2 + leaves = jax.tree_util.tree_leaves(parameter_grad) + assert len(leaves) == 1 + np.testing.assert_allclose(leaves[0], expected_sink_grad, rtol=0.1, atol=1e-2) + + +@pytest.mark.parametrize("fp8,fp8_backward", [(False, False), (True, False), (True, True)]) +def test_cudnn_python_dropout_replay(fp8, fp8_backward, monkeypatch): + """Explicit keys replay dropout through JIT and forward recomputation.""" + if fp8 and get_device_compute_capability(0) < 90: + pytest.skip("cuDNN FP8 SDPA requires Hopper or newer.") + monkeypatch.setenv("NVTE_FP8_DPA_BWD", "1" if fp8_backward else "0") + helper = FusedAttnHelper( + is_training=True, + batch_size=1, + q_dtype=jnp.bfloat16, + kv_dtype=jnp.bfloat16, + qkv_layout=QKVLayout.BSHD_BSHD_BSHD, + attn_bias_type=AttnBiasType.NO_BIAS, + attn_mask_type=AttnMaskType.NO_MASK, + softmax_type=AttnSoftmaxType.VANILLA_SOFTMAX, + dropout_probability=0.2, + q_num_heads=2, + kv_num_heads=2, + q_max_seqlen=128, + kv_max_seqlen=128, + head_dim_qk=64, + head_dim_v=64, + window_size=(-1, -1), + attn_scale=0.125, + fp8=fp8, + ) + backend, reason = helper.get_fused_attn_backend() + if backend == NVTE_Fused_Attn_Backend.NVTE_No_Backend: + pytest.skip(reason) + arrays = tuple(jnp.ones((1, 128, 2, 64), jnp.bfloat16) for _ in range(3)) + + def attention(arrays, seed): + return fused_attn( + arrays, + None, + None, + seed, + attn_bias_type=AttnBiasType.NO_BIAS, + attn_mask_type=AttnMaskType.NO_MASK, + qkv_layout=QKVLayout.BSHD_BSHD_BSHD, + softmax_type=AttnSoftmaxType.VANILLA_SOFTMAX, + scaling_factor=0.125, + dropout_probability=0.2, + is_training=True, + fp8=fp8, + ) + + seed = jax.random.PRNGKey(123) + forward = jax.jit(attention) + actual = forward(arrays, seed) + np.testing.assert_array_equal(actual, forward(arrays, seed)) + assert not np.array_equal(actual, forward(arrays, jax.random.PRNGKey(456))) + ordinary = jax.jit(jax.grad(lambda xs: jnp.sum(attention(xs, seed))))(arrays) + recomputed = jax.jit(jax.grad(lambda xs: jnp.sum(jax.checkpoint(attention)(xs, seed))))(arrays) + for expected, actual in zip(ordinary, recomputed): + np.testing.assert_array_equal(actual, expected) + + +@pytest.mark.parametrize("strategy", [CPStrategy.ALL_GATHER, CPStrategy.RING]) +@pytest.mark.parametrize("fp8", [False, True]) +def test_cudnn_python_context_parallel(strategy, fp8): + """Two-device context parallel attention preserves FP8 forward and backward.""" + if jax.local_device_count() < 2: + pytest.skip("Context parallel attention requires two GPUs.") + if fp8 and get_device_compute_capability(0) < 90: + pytest.skip("cuDNN FP8 SDPA requires Hopper or newer.") + mesh = Mesh(np.asarray(jax.local_devices()[:2]), ("cp",)) + sharding = NamedSharding(mesh, PartitionSpec(None, "cp", None, None)) + q = jnp.full((2, 256, 2, 64), 0.5, jnp.bfloat16) + v = jnp.zeros_like(q).at[:, ::2, :, 0].set(1.0) + arrays = tuple(jax.device_put(value, sharding) for value in (q, q, v)) + + def attention(arrays): + return fused_attn( + arrays, + None, + None, + None, + attn_bias_type=AttnBiasType.NO_BIAS, + attn_mask_type=AttnMaskType.NO_MASK, + qkv_layout=QKVLayout.BSHD_BSHD_BSHD, + softmax_type=AttnSoftmaxType.VANILLA_SOFTMAX, + scaling_factor=0.125, + dropout_probability=0.0, + is_training=True, + context_parallel_axis="cp", + context_parallel_strategy=strategy, + fp8=fp8, + ) + + def reference(arrays): + query, key, value = (tensor.astype(jnp.float32) for tensor in arrays) + scores = jnp.einsum("bqhd,bkhd->bhqk", query, key) * 0.125 + return jnp.einsum("bhqk,bkhd->bqhd", jax.nn.softmax(scores, axis=-1), value) + + expected = jax.jit(reference)(arrays) + expected_grads = jax.jit(jax.grad(lambda xs: jnp.sum(reference(xs))))(arrays) + with mesh, autocast(mesh_resource=MeshResource(cp_resource="cp")): + actual = jax.jit(attention)(arrays) + actual_grads = jax.jit(jax.grad(lambda xs: jnp.sum(attention(xs))))(arrays) + np.testing.assert_allclose(actual, expected, rtol=5e-3, atol=5e-3) + for actual, expected in zip(actual_grads, expected_grads): + np.testing.assert_allclose(actual, expected, rtol=1e-2, atol=1e-2) diff --git a/tests/jax/test_fused_attn_score_mod.py b/tests/jax/test_fused_attn_score_mod.py index 5a613101585..5d9b1508589 100644 --- a/tests/jax/test_fused_attn_score_mod.py +++ b/tests/jax/test_fused_attn_score_mod.py @@ -10,6 +10,7 @@ import pytest import transformer_engine.jax.cpp_extensions.flex_attention as tex_attention +import transformer_engine.jax.cpp_extensions.cudnn_graph as tex_cudnn_graph from transformer_engine.jax.attention import ( AttnBiasType, AttnMaskType, @@ -659,14 +660,14 @@ class FakeCudnn: __version__ = "1.22.0" monkeypatch.setattr( - tex_attention.transformer_engine_jax, + tex_cudnn_graph.transformer_engine_jax, "get_cudnn_frontend_version", lambda: 12200, ) assert tex_attention._check_cudnn_frontend_version_match(FakeCudnn) == 12200 monkeypatch.setattr( - tex_attention.transformer_engine_jax, + tex_cudnn_graph.transformer_engine_jax, "get_cudnn_frontend_version", lambda: 12100, ) diff --git a/transformer_engine/jax/attention.py b/transformer_engine/jax/attention.py index edf418f15ea..2765e1fbb5b 100644 --- a/transformer_engine/jax/attention.py +++ b/transformer_engine/jax/attention.py @@ -17,7 +17,6 @@ from transformer_engine_jax import NVTE_Mask_Type from transformer_engine_jax import NVTE_QKV_Layout from transformer_engine_jax import NVTE_QKV_Format -from transformer_engine_jax import nvte_get_qkv_format from transformer_engine_jax import NVTE_Softmax_Type from . import cpp_extensions as tex @@ -33,6 +32,7 @@ class AttnBiasType(Enum): NO_BIAS = NVTE_Bias_Type.NVTE_NO_BIAS PRE_SCALE_BIAS = NVTE_Bias_Type.NVTE_PRE_SCALE_BIAS POST_SCALE_BIAS = NVTE_Bias_Type.NVTE_POST_SCALE_BIAS + ALIBI = NVTE_Bias_Type.NVTE_ALIBI class AttnMaskType(Enum): @@ -139,7 +139,7 @@ def get_qkv_format(self): """ Return the corresponding qkv_format (BSHD, SBHD, THD) """ - return QKVFormat(nvte_get_qkv_format(self.value)) + return QKVFormat.THD if self.is_thd() else QKVFormat.BSHD def is_qkvpacked(self): """ @@ -348,6 +348,7 @@ def is_fused_attn_kernel_available( bias_seqlen_q: Optional[int] = None, bias_seqlen_kv: Optional[int] = None, max_segments_per_seq: int = 1, + fp8: bool = False, ): """ To check whether the fused attention kernel is supported. @@ -384,6 +385,7 @@ def make_helper(attn_mask_type): bias_seqlen_q=bias_seqlen_q, bias_seqlen_kv=bias_seqlen_kv, max_segments_per_seq=max_segments_per_seq, + fp8=fp8, ) helper = make_helper(attn_mask_type) @@ -1142,6 +1144,7 @@ def _legacy_fused_attn( context_parallel_axis: str = "", softmax_offset: Optional[jnp.ndarray] = None, return_max_logit: bool = False, + fp8: bool = False, ): """ Perform non-THD (non-packed) cuDNN fused attention. @@ -1236,6 +1239,7 @@ def _legacy_fused_attn( context_parallel_causal_load_balanced=context_parallel_causal_load_balanced, context_parallel_axis=context_parallel_axis, return_max_logit=return_max_logit, + fp8=fp8, ) return output @@ -1262,6 +1266,7 @@ def fused_attn_thd( context_parallel_axis: str = "", softmax_offset: Optional[jnp.ndarray] = None, return_max_logit: bool = False, + fp8: bool = False, ): """ Deprecated THD fused attn, please use fusd_attn with SequenceDescriptor @@ -1317,6 +1322,7 @@ def fused_attn_thd( context_parallel_causal_load_balanced=context_parallel_causal_load_balanced, context_parallel_axis=context_parallel_axis, return_max_logit=return_max_logit, + fp8=fp8, ) return output @@ -1324,7 +1330,7 @@ def fused_attn_thd( @partial( jax.custom_vjp, - nondiff_argnums=(5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19), + nondiff_argnums=(5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20), ) def _fused_attn( qkv: Tuple[jnp.ndarray, ...], @@ -1347,6 +1353,7 @@ def _fused_attn( context_checkpoint_name: str = "context", stripe_size: int | None = None, return_max_logit: bool = False, + fp8: bool = False, ): output, _ = _fused_attn_fwd_rule( qkv, @@ -1369,6 +1376,7 @@ def _fused_attn( context_checkpoint_name=context_checkpoint_name, stripe_size=stripe_size, return_max_logit=return_max_logit, + fp8=fp8, ) return output @@ -1394,6 +1402,7 @@ def _fused_attn_fwd_rule( context_checkpoint_name, stripe_size, return_max_logit, + fp8, ): output, softmax_aux, rng_state, max_logit = tex.fused_attn_fwd( qkv, @@ -1415,6 +1424,7 @@ def _fused_attn_fwd_rule( context_parallel_axis=context_parallel_axis, stripe_size=stripe_size, return_max_logit=return_max_logit, + fp8=fp8, ) output = checkpoint_name(output, context_checkpoint_name) softmax_aux = checkpoint_name(softmax_aux, context_checkpoint_name) @@ -1448,6 +1458,7 @@ def _fused_attn_bwd_rule( context_checkpoint_name, stripe_size, return_max_logit, + fp8, ctx, dz, ): @@ -1485,8 +1496,9 @@ def _fused_attn_bwd_rule( context_parallel_causal_load_balanced=context_parallel_causal_load_balanced, context_parallel_axis=context_parallel_axis, stripe_size=stripe_size, + fp8=fp8, ) - if attn_bias_type == AttnBiasType.NO_BIAS: + if attn_bias_type in (AttnBiasType.NO_BIAS, AttnBiasType.ALIBI): grad_bias = None if softmax_type != AttnSoftmaxType.LEARNABLE_SOFTMAX: grad_softmax_offset = None @@ -1580,6 +1592,7 @@ def fused_attn( score_mod_tensors: Optional[Mapping[str, Any]] = None, score_mod_bprop_tensors: Optional[Mapping[str, Any]] = None, return_max_logit: bool = False, + fp8: bool = False, ): """ Perform cuDNN fused attention. @@ -1636,6 +1649,9 @@ def fused_attn( non-differentiable auxiliary inputs. score_mod_bprop_tensors (Optional[Mapping[str, Any]]): Additional tensors or Python/NumPy scalars made available to `score_mod_bprop`. + fp8 (bool): Use tensor-scaled cuDNN FP8 SDPA with E4M3 query/key/value and + E5M2 output gradients. Scaling uses the current tensor maximum. Outputs + and gradients retain the input FP16/BF16 dtype. return_max_logit (bool): If True, also return per-head maximum attention logits with shape ``[h]``. Returns: @@ -1686,6 +1702,8 @@ def fused_attn( if score_mod_only_args: raise ValueError(f"{', '.join(score_mod_only_args)} require score_mod to be provided.") else: + if fp8: + raise ValueError("FP8 is not supported with score_mod fused_attn.") if return_max_logit: raise ValueError("return_max_logit is not supported with score_mod fused_attn.") tex.validate_fused_attn_score_mod( @@ -1748,6 +1766,7 @@ def fused_attn( context_parallel_axis=context_parallel_axis, softmax_offset=softmax_offset, return_max_logit=return_max_logit, + fp8=fp8, ) if max_segments_per_seq > 1 and not qkv_layout.is_thd(): warnings.warn( @@ -1779,5 +1798,6 @@ def fused_attn( context_checkpoint_name=context_checkpoint_name, stripe_size=stripe_size, return_max_logit=return_max_logit, + fp8=fp8, ) return output diff --git a/transformer_engine/jax/cpp_extensions/attention.py b/transformer_engine/jax/cpp_extensions/attention.py index 3a1de7bd677..702b9457017 100644 --- a/transformer_engine/jax/cpp_extensions/attention.py +++ b/transformer_engine/jax/cpp_extensions/attention.py @@ -12,20 +12,13 @@ import jax import jax.numpy as jnp -from jax import dtypes, lax, ffi +from jax import dtypes, lax +from jax.interpreters import mlir from jax.sharding import PartitionSpec, NamedSharding from jax.experimental.custom_partitioning import SdyShardingRule -import transformer_engine_jax from transformer_engine_jax import ( - DType, - JAXX_Scaling_Mode, - NVTE_Bias_Type, NVTE_Fused_Attn_Backend, - NVTE_Mask_Type, - NVTE_QKV_Format, - NVTE_QKV_Layout, - NVTE_Softmax_Type, ) from transformer_engine.jax.attention import ( AttnBiasType, @@ -39,13 +32,10 @@ from ..sharding import with_sharding_constraint_by_logical_axes, HEAD_AXES, is_mesh_available from .base import BasePrimitive, register_primitive +from . import cudnn_attention from .misc import ( check_valid_batch_dims, - jax_dtype_to_te_dtype, - te_dtype_to_jax_dtype, get_padded_spec, - get_cudnn_version, - get_all_device_compute_capability, ) from ..sharding import ( global_mesh_resource, @@ -118,6 +108,7 @@ def setup_logging(): "stripe_size", "return_max_logit", "allow_fast_causal_path", + "fp8", ], ) @dataclass(frozen=True) @@ -144,6 +135,7 @@ class _FusedAttnConfig: ) # Only for CP + Striped. For Ring P2P, stripe_size=1 only.For AG, stripe_size>=1. return_max_logit: bool = False allow_fast_causal_path: bool = True + fp8: bool = False @property def effective_window_size(self) -> Tuple[int, int]: @@ -153,64 +145,6 @@ def effective_window_size(self) -> Tuple[int, int]: return self.window_size -@dataclass -class FusedAttnParams: - """ - Attention parameters used to select the fused attention backend. - - Fields are declared in the order of the ``FusedAttnConfig`` struct in - ``common/fused_attn/config_and_params.h``, which is the order the C++ binding reads them in - and the order it fills the config with. Fields JAX does not use, are left at their ``FusedAttnConfig`` - defaults. - """ - - # basic attention settings - is_training: bool = True - deterministic: bool = False - cuda_graph: bool = False - return_max_logit: bool = False - attn_mask_type: NVTE_Mask_Type = NVTE_Mask_Type.NVTE_NO_MASK - bias_type: NVTE_Bias_Type = NVTE_Bias_Type.NVTE_NO_BIAS - window_size_left: int = -1 - window_size_right: int = -1 - bottom_right_diagonal: bool = True - softmax_type: NVTE_Softmax_Type = NVTE_Softmax_Type.NVTE_VANILLA_SOFTMAX - scaling_mode: JAXX_Scaling_Mode = JAXX_Scaling_Mode.NO_SCALING - dropout: float = 0.0 - attn_scale: float = 1.0 - - # tensor types - qkv_dtype: DType = DType.kBFloat16 - o_dtype: DType = DType.kBFloat16 - do_dtype: DType = DType.kBFloat16 - dqkv_dtype: DType = DType.kBFloat16 - - # tensor layouts - qkv_layout: NVTE_QKV_Layout = NVTE_QKV_Layout.NVTE_QKV_Layout_NOT_SET - o_format: NVTE_QKV_Format = NVTE_QKV_Format.NVTE_QKV_Format_NOT_SET - do_format: NVTE_QKV_Format = NVTE_QKV_Format.NVTE_QKV_Format_NOT_SET - dqkv_layout: NVTE_QKV_Layout = NVTE_QKV_Layout.NVTE_QKV_Layout_NOT_SET - qkv_scale_inv_format: NVTE_QKV_Format = NVTE_QKV_Format.NVTE_QKV_Format_NOT_SET - do_scale_inv_format: NVTE_QKV_Format = NVTE_QKV_Format.NVTE_QKV_Format_NOT_SET - - # tensor dimensions - batch_size: int = 0 - num_attn_heads: int = 0 - num_gqa_groups: int = 0 - head_dim_qk: int = 0 - head_dim_v: int = 0 - max_seqlen_q: int = 0 - max_seqlen_kv: int = 0 - num_tokens_q: int = 0 - num_tokens_kv: int = 0 - - # bias dimensions - bias_batch_size: int = 0 - bias_num_heads: int = 0 - bias_seqlen_q: int = 0 - bias_seqlen_kv: int = 0 - - @dataclass(frozen=True, kw_only=True) class FusedAttnHelper: """ @@ -241,6 +175,7 @@ class FusedAttnHelper: bias_seqlen_q: Optional[int] = None bias_seqlen_kv: Optional[int] = None max_segments_per_seq: int = 1 + fp8: bool = False def is_fused_attn_kernel_available(self): """Check if there is available fused attention kernel""" @@ -257,55 +192,68 @@ def get_fused_attn_backend(self): that no fused backend is available), and ``NVTE_DEBUG_LEVEL=2`` additionally logs the resolved config and the reason fused attention was rejected. """ - q_type = jax_dtype_to_te_dtype(self.q_dtype) - kv_type = jax_dtype_to_te_dtype(self.kv_dtype) - if q_type != kv_type: + if self.q_dtype != self.kv_dtype: raise ValueError("Q and KV must have the same data type.") - bias_batch = bias_heads = bias_seqlen_q = bias_seqlen_kv = 0 - if self.attn_bias_type == AttnBiasType.POST_SCALE_BIAS: - bias_batch = self.bias_batch or 0 - bias_heads = self.bias_heads or 0 - bias_seqlen_q = self.bias_seqlen_q or 0 - bias_seqlen_kv = self.bias_seqlen_kv or 0 - num_segments = self.batch_size - num_tokens_q = num_tokens_kv = 0 - if self.qkv_layout.is_thd(): - num_segments = self.batch_size * self.max_segments_per_seq - num_tokens_q = self.batch_size * self.q_max_seqlen - num_tokens_kv = self.batch_size * self.kv_max_seqlen - backend, message = transformer_engine_jax.get_fused_attn_backend( - FusedAttnParams( + backend = NVTE_Fused_Attn_Backend.NVTE_No_Backend + message = "" + if self.attn_bias_type == AttnBiasType.PRE_SCALE_BIAS: + message = "Fused attention does not support pre-scale bias." + elif self.q_dtype not in (jnp.float16, jnp.bfloat16): + message = "Fused attention requires FP16/BF16 inputs." + else: + config = _FusedAttnConfig( + attn_bias_type=self.attn_bias_type, + attn_mask_type=self.attn_mask_type, + softmax_type=self.softmax_type, + qkv_layout=self.qkv_layout, + scaling_factor=self.attn_scale, + dropout_probability=self.dropout_probability, is_training=self.is_training, - deterministic=not self.is_non_deterministic_allowed(), - return_max_logit=self.return_max_logit, - attn_mask_type=self.attn_mask_type.value, - bias_type=self.attn_bias_type.value, - window_size_left=self.window_size[0], - window_size_right=self.window_size[1], + max_segments_per_seq=self.max_segments_per_seq, + window_size=self.window_size, bottom_right_diagonal=self.bottom_right_diagonal, - softmax_type=self.softmax_type.value, - dropout=self.dropout_probability, - attn_scale=self.attn_scale, - qkv_dtype=q_type, - o_dtype=q_type, - do_dtype=q_type, - dqkv_dtype=q_type, - qkv_layout=self.qkv_layout.value, - batch_size=num_segments, - num_attn_heads=self.q_num_heads, - num_gqa_groups=self.kv_num_heads, - head_dim_qk=self.head_dim_qk, - head_dim_v=self.head_dim_v, - max_seqlen_q=self.q_max_seqlen, - max_seqlen_kv=self.kv_max_seqlen, - num_tokens_q=num_tokens_q, - num_tokens_kv=num_tokens_kv, - bias_batch_size=bias_batch, - bias_num_heads=bias_heads, - bias_seqlen_q=bias_seqlen_q, - bias_seqlen_kv=bias_seqlen_kv, + context_parallel_load_balanced=False, + cp_axis="", + cp_striped_window_size=None, + stripe_size=None, + return_max_logit=self.return_max_logit, + fp8=self.fp8, ) - ) + qshape = (self.batch_size, self.q_max_seqlen, self.q_num_heads, self.head_dim_qk) + kshape = (self.batch_size, self.kv_max_seqlen, self.kv_num_heads, self.head_dim_qk) + vshape = (self.batch_size, self.kv_max_seqlen, self.kv_num_heads, self.head_dim_v) + if self.qkv_layout.is_qkvpacked(): + qshape = (*qshape[:-2], 3, *qshape[-2:]) + kshape = vshape = (0,) + elif self.qkv_layout.is_kvpacked(): + kshape = (*kshape[:-2], 2, *kshape[-2:]) + vshape = (0,) + bias_shape = ( + ( + self.bias_batch or 1, + self.bias_heads or self.q_num_heads, + self.bias_seqlen_q or self.q_max_seqlen, + self.bias_seqlen_kv or self.kv_max_seqlen, + ) + if self.attn_bias_type == AttnBiasType.POST_SCALE_BIAS + else (0,) + ) + avals = tuple( + jax.ShapeDtypeStruct(shape, self.q_dtype) + for shape in (qshape, kshape, vshape, bias_shape) + ) + avals += (jax.ShapeDtypeStruct((1, self.q_num_heads, 1, 1), jnp.float32),) + try: + cudnn_attention.get_graph("fwd", avals, config) + if self.is_training: + cudnn_attention.get_graph("bwd", avals, config) + backend = ( + NVTE_Fused_Attn_Backend.NVTE_FP8 + if self.fp8 + else NVTE_Fused_Attn_Backend.NVTE_F16_arbitrary_seqlen + ) + except (ValueError, RuntimeError, ImportError) as exc: + message = str(exc) AttentionLogging.setup_logging() logger = AttentionLogging.logger @@ -401,7 +349,8 @@ class _FusedAttnRNGStateChecker: The fused attention backend requires a 64 bits seed and a 64 bits offset. However, JAX doesn't enable 64 bits by default, so we have to emulate seed as two 32 bits array. - The offset calculation is maintained in the backend. + The explicit JAX key supplies the seed; the saved zero offset makes + dropout replayable when a forward pass is recomputed. """ rng_state_dtype: jnp.dtype = jnp.uint32 @@ -538,9 +487,13 @@ def abstract( bias_seqlen_q=bias_seqlen_q, bias_seqlen_kv=bias_seqlen_kv, max_segments_per_seq=config.max_segments_per_seq, + fp8=config.fp8, ).get_fused_attn_backend() - if backend == NVTE_Fused_Attn_Backend.NVTE_F16_arbitrary_seqlen: + if backend in ( + NVTE_Fused_Attn_Backend.NVTE_F16_arbitrary_seqlen, + NVTE_Fused_Attn_Backend.NVTE_FP8, + ): if config.qkv_layout.is_thd(): softmax_shape = (*batch_shape, q_max_seqlen, attn_heads, 1) else: @@ -571,42 +524,16 @@ def abstract( rng_state_shape = (seed_aval.shape[0], checker.rng_state_size) rng_state_aval = seed_aval.update(shape=rng_state_shape, dtype=checker.rng_state_dtype) - if config.attn_bias_type == AttnBiasType.NO_BIAS: + if config.attn_bias_type in (AttnBiasType.NO_BIAS, AttnBiasType.ALIBI): bias_batch = bias_heads = 0 else: *bias_batch_shape, bias_heads, _, _ = bias_aval.shape bias_batch = reduce(operator.mul, bias_batch_shape) - # do a dummy kernel call here to get workspace buffer shapes/dtypes that XLA needs to - # prepare for the active fused-attn backend - input_batch = reduce(operator.mul, batch_shape) - wkspace_info = transformer_engine_jax.get_fused_attn_fwd_workspace_sizes( - input_batch, - bias_batch, - q_max_seqlen, - kv_max_seqlen, - attn_heads, - num_gqa_groups, - bias_heads, - q_head_dim, - v_head_dim, - config.scaling_factor, - config.dropout_probability, - config.attn_bias_type.value, - config.attn_mask_type.value, - config.softmax_type.value, - config.qkv_layout.value, - jax_dtype_to_te_dtype(q_aval.dtype), - config.is_training, - config.max_segments_per_seq, - config.effective_window_size[0], - config.effective_window_size[1], - config.return_max_logit, - config.bottom_right_diagonal, - ) - wkspace_aval = q_aval.update( - shape=wkspace_info[0], dtype=te_dtype_to_jax_dtype(wkspace_info[1]) + graph = cudnn_attention.get_graph( + "fwd", (q_aval, k_aval, v_aval, bias_aval, softmax_offset_aval), config ) + wkspace_aval = q_aval.update(shape=(graph.graph.workspace_size,), dtype=jnp.uint8) assert ( softmax_offset_aval.dtype == jnp.float32 @@ -656,32 +583,12 @@ def lowering( *, config: _FusedAttnConfig, ): - """ - Fused attention fwd lowering rules - """ - q_aval, k_aval, v_aval, bias_aval, *_ = ctx.avals_in - - ( - batch_shape, - q_max_seqlen, - kv_max_seqlen, - attn_heads, - num_gqa_groups, - q_head_dim, - v_head_dim, - ) = FusedAttnHelper.parse_qkv_aval(q_aval, k_aval, v_aval, config.qkv_layout) + """Lower Python-built cuDNN attention through the JAX-owned graph executor.""" - input_batch = reduce(operator.mul, batch_shape) - - if config.attn_bias_type == AttnBiasType.NO_BIAS: - bias_batch = bias_heads = 0 - else: - *bias_batch_shape, bias_heads, _, _ = bias_aval.shape - bias_batch = reduce(operator.mul, bias_batch_shape) - - window_size_left, window_size_right = config.effective_window_size + def execute(*args): + return cudnn_attention.execute("fwd", args, config, ctx.avals_out) - return ffi.ffi_lowering(FusedAttnFwdPrimitive.name)( + return mlir.lower_fun(execute, multiple_results=True)( ctx, q, k, @@ -696,29 +603,7 @@ def lowering( _q_segment_ids, _kv_segment_ids, _q_segment_pos, - _kv_segment_pos, # ffi_lowering needs number of parameters meets primitive.lowering - input_batch=input_batch, - bias_batch=bias_batch, - q_max_seqlen=q_max_seqlen, - kv_max_seqlen=kv_max_seqlen, - attn_heads=attn_heads, - num_gqa_groups=num_gqa_groups, - bias_heads=bias_heads, - qk_head_dim=q_head_dim, - v_head_dim=v_head_dim, - max_segments_per_seq=config.max_segments_per_seq, - scaling_factor=float(config.scaling_factor), - dropout_probability=float(config.dropout_probability), - bias_type=int(config.attn_bias_type.value), - mask_type=int(config.attn_mask_type.value), - qkv_layout=int(config.qkv_layout.value), - is_training=config.is_training, - return_max_logit=config.return_max_logit, - deterministic=not FusedAttnHelper.is_non_deterministic_allowed(), - window_size_left=window_size_left, - window_size_right=window_size_right, - bottom_right_diagonal=config.bottom_right_diagonal, - softmax_type=int(config.softmax_type.value), + _kv_segment_pos, ) @staticmethod @@ -879,7 +764,7 @@ def _reduce_max_logit(max_tensor, output, q_seqlen, q_seq_offsets, config): @staticmethod def _uses_thd_ragged_max_tensor(config): """Return whether cuDNN writes THD Max with BSH-like ragged-stats layout.""" - return config.qkv_layout.is_thd() and 120 not in get_all_device_compute_capability() + return config.qkv_layout.is_thd() and 120 not in cudnn_attention.device_capabilities() @staticmethod def _empty_or_neg_inf_max_logit(head, dtype, config): @@ -1114,57 +999,13 @@ def abstract( f" kv_seqlen_or_cu_seqlen_aval.dtype={kv_seqlen_or_cu_seqlen_aval.dtype}" ) - ( - batch_shape, - q_max_seqlen, - kv_max_seqlen, - attn_heads, - num_gqa_groups, - qk_head_dim, - v_head_dim, - ) = FusedAttnHelper.parse_qkv_aval(q_aval, k_aval, v_aval, config.qkv_layout) - - if config.attn_bias_type == AttnBiasType.NO_BIAS: - bias_batch = bias_heads = 0 - else: - *bias_batch_shape, bias_heads, _, _ = bias_aval.shape - bias_batch = reduce(operator.mul, bias_batch_shape) - - deterministic = not FusedAttnHelper.is_non_deterministic_allowed() + attn_heads = FusedAttnHelper.parse_qkv_aval(q_aval, k_aval, v_aval, config.qkv_layout)[3] - input_batch = reduce(operator.mul, batch_shape) - wkspace_shape, wkspace_dtype = transformer_engine_jax.get_fused_attn_bwd_workspace_sizes( - input_batch, - bias_batch, - q_max_seqlen, - kv_max_seqlen, - attn_heads, - num_gqa_groups, - bias_heads, - qk_head_dim, - v_head_dim, - config.scaling_factor, - config.dropout_probability, - config.attn_bias_type.value, - config.attn_mask_type.value, - config.softmax_type.value, - config.qkv_layout.value, - jax_dtype_to_te_dtype(q_aval.dtype), - config.is_training, - deterministic, - config.max_segments_per_seq, - config.effective_window_size[0], - config.effective_window_size[1], - config.bottom_right_diagonal, - ) - - dq_aval = q_aval.update(shape=q_aval.shape, dtype=q_dtype) - dk_aval = k_aval.update(shape=k_aval.shape, dtype=k_dtype) - dv_aval = v_aval.update(shape=v_aval.shape, dtype=v_dtype) - dbias_aval = bias_aval.update(shape=bias_aval.shape, dtype=bias_dtype) - wkspace_aval = q_aval.update( - shape=wkspace_shape, dtype=te_dtype_to_jax_dtype(wkspace_dtype) + graph = cudnn_attention.get_graph( + "bwd", (q_aval, k_aval, v_aval, bias_aval, softmax_offset_aval), config ) + dq_aval, dk_aval, dv_aval, dbias_aval = q_aval, k_aval, v_aval, bias_aval + wkspace_aval = q_aval.update(shape=(graph.graph.workspace_size,), dtype=jnp.uint8) # Validate incoming softmax_offset shape and dtype assert ( @@ -1223,32 +1064,12 @@ def lowering( *, config, ): - """ - Fused attention bwd lowering rules - """ - q_aval, k_aval, v_aval, bias_aval, *_ = ctx.avals_in - - ( - batch_shape, - q_max_seqlen, - kv_max_seqlen, - attn_heads, - num_gqa_groups, - qk_head_dim, - v_head_dim, - ) = FusedAttnHelper.parse_qkv_aval(q_aval, k_aval, v_aval, config.qkv_layout) - - input_batch = reduce(operator.mul, batch_shape) - - if config.attn_bias_type == AttnBiasType.NO_BIAS: - bias_batch = bias_heads = 0 - else: - *bias_batch_shape, bias_heads, _, _ = bias_aval.shape - bias_batch = reduce(operator.mul, bias_batch_shape) + """Lower Python-built cuDNN attention through the JAX-owned graph executor.""" - window_size_left, window_size_right = config.effective_window_size + def execute(*args): + return cudnn_attention.execute("bwd", args, config, ctx.avals_out) - return ffi.ffi_lowering(FusedAttnBwdPrimitive.name)( + return mlir.lower_fun(execute, multiple_results=True)( ctx, q, k, @@ -1266,28 +1087,7 @@ def lowering( q_segment_ids, kv_segment_ids, q_segment_pos, - kv_segment_pos, # ffi_lowering needs number of parameters meets primitive.lowering - input_batch=input_batch, - bias_batch=bias_batch, - q_max_seqlen=q_max_seqlen, - kv_max_seqlen=kv_max_seqlen, - attn_heads=attn_heads, - num_gqa_groups=num_gqa_groups, - bias_heads=bias_heads, - qk_head_dim=qk_head_dim, - v_head_dim=v_head_dim, - max_segments_per_seq=config.max_segments_per_seq, - scaling_factor=float(config.scaling_factor), - dropout_probability=float(config.dropout_probability), - bias_type=int(config.attn_bias_type.value), - mask_type=int(config.attn_mask_type.value), - qkv_layout=int(config.qkv_layout.value), - is_training=config.is_training, - deterministic=not FusedAttnHelper.is_non_deterministic_allowed(), - window_size_left=window_size_left, - window_size_right=window_size_right, - bottom_right_diagonal=config.bottom_right_diagonal, - softmax_type=int(config.softmax_type.value), + kv_segment_pos, ) @staticmethod @@ -1714,6 +1514,7 @@ def get_step_config(self) -> _FusedAttnConfig: cp_striped_window_size=None, stripe_size=self.config.stripe_size, return_max_logit=self.config.return_max_logit, + fp8=self.config.fp8, ) def get_step_config_for_striped(self, max_seqlen, cp_size) -> _FusedAttnConfig: @@ -1735,6 +1536,7 @@ def get_step_config_for_striped(self, max_seqlen, cp_size) -> _FusedAttnConfig: cp_striped_window_size=None, stripe_size=self.config.stripe_size, return_max_logit=self.config.return_max_logit, + fp8=self.config.fp8, ) def all_gather_kv(self, k, v): @@ -2807,6 +2609,7 @@ def get_step_config(self, attn_mask_type) -> _FusedAttnConfig: cp_striped_window_size=None, stripe_size=self.config.stripe_size, return_max_logit=self.config.return_max_logit, + fp8=self.config.fp8, ) def get_striped_thd_step_config(self) -> _FusedAttnConfig: @@ -3735,6 +3538,7 @@ def fused_attn_fwd( context_parallel_axis: str = "", stripe_size: int | None = None, return_max_logit: bool = False, + fp8: bool = False, ) -> jnp.ndarray: """ Perform the forward pass of with cuDNN fused attention implementations. @@ -3798,7 +3602,7 @@ def fused_attn_fwd( else: raise ValueError(f"Unknown {qkv_layout=}") - if attn_bias_type == AttnBiasType.NO_BIAS: + if attn_bias_type in (AttnBiasType.NO_BIAS, AttnBiasType.ALIBI): assert ( bias is None ), f"bias must be None when attn_bias_type is NO_BIAS, but got bias={bias}" @@ -3851,6 +3655,7 @@ def fused_attn_fwd( cp_striped_window_size=None, stripe_size=stripe_size, return_max_logit=return_max_logit, + fp8=fp8, ) primitive = None @@ -3902,6 +3707,7 @@ def fused_attn_bwd( context_parallel_causal_load_balanced: bool = False, context_parallel_axis: str = "", stripe_size: int | None = None, + fp8: bool = False, ): """ Perform the backward pass of the cuDNN fused attention implementations. @@ -3967,7 +3773,7 @@ def fused_attn_bwd( else: raise ValueError(f"Unknown {qkv_layout=}") - if attn_bias_type == AttnBiasType.NO_BIAS: + if attn_bias_type in (AttnBiasType.NO_BIAS, AttnBiasType.ALIBI): assert ( bias is None ), f"bias must be None when attn_bias_type is NO_BIAS, but got bias with type={type(bias)}" @@ -3995,15 +3801,18 @@ def fused_attn_bwd( softmax_offset, (None, HEAD_AXES, None, None) ) - compute_capabilities = get_all_device_compute_capability() + compute_capabilities = cudnn_attention.device_capabilities() if any(x >= 100 for x in compute_capabilities) and is_training: assert ( FusedAttnHelper.is_non_deterministic_allowed() - and (attn_bias_type == AttnBiasType.NO_BIAS or dropout_probability == 0.0) + and ( + attn_bias_type in (AttnBiasType.NO_BIAS, AttnBiasType.ALIBI) + or dropout_probability == 0.0 + ) ) or ( not FusedAttnHelper.is_non_deterministic_allowed() - and get_cudnn_version() >= (9, 18, 1) - and attn_bias_type == AttnBiasType.NO_BIAS + and cudnn_attention.cudnn_version() >= 91801 + and attn_bias_type in (AttnBiasType.NO_BIAS, AttnBiasType.ALIBI) and dropout_probability == 0.0 ), ( "For sm100+, non-deterministic bprop does not support bias with dropout," @@ -4012,6 +3821,7 @@ def fused_attn_bwd( fused_config = _FusedAttnConfig( attn_bias_type=attn_bias_type, + fp8=fp8, attn_mask_type=attn_mask_type, qkv_layout=qkv_layout, softmax_type=softmax_type, diff --git a/transformer_engine/jax/cpp_extensions/cudnn_attention.py b/transformer_engine/jax/cpp_extensions/cudnn_attention.py new file mode 100644 index 00000000000..1701d7c0d68 --- /dev/null +++ b/transformer_engine/jax/cpp_extensions/cudnn_attention.py @@ -0,0 +1,511 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. +"""JAX cuDNN attention graphs, metadata preparation and cached plan dispatch.""" + +import math +import os +from dataclasses import dataclass +from functools import lru_cache + +import jax +import jax.numpy as jnp +import numpy as np +from jax import ffi + +from .cudnn_graph import import_cudnn, finalize_graph, serialize_graph, graph_attributes + + +def device_capabilities(): + """Query the JAX devices without entering TE common.""" + return tuple(int(float(device.compute_capability) * 10) for device in jax.local_devices()) + + +def cudnn_version(): + """Return the backend used by the Python frontend.""" + return int(import_cudnn().backend_version()) + + +def _dtype(cudnn, dtype): + return { + jnp.dtype(jnp.float16): cudnn.data_type.HALF, + jnp.dtype(jnp.bfloat16): cudnn.data_type.BFLOAT16, + jnp.dtype(jnp.float32): cudnn.data_type.FLOAT, + jnp.dtype(jnp.float8_e4m3fn): cudnn.data_type.FP8_E4M3, + jnp.dtype(jnp.float8_e5m2): cudnn.data_type.FP8_E5M2, + }[jnp.dtype(dtype)] + + +def _bucket(value, minimum, maximum): + if value <= maximum: + return max(minimum, 1 << (value - 1).bit_length()) + return (value + maximum - 1) // maximum * maximum + + +def _stride(shape): + return tuple(math.prod(shape[index + 1 :]) for index in range(len(shape))) + + +def _dims(avals, config): + q, k, v = avals[:3] + layout = config.qkv_layout + if layout.is_qkvpacked(): + *batch, sq, _, h, dq = q.shape + sk, hg, dv = sq, h, dq + elif layout.is_kvpacked(): + *batch, sq, h, dq = q.shape + *_, sk, _, hg, dv = k.shape + else: + *batch, sq, h, dq = q.shape + sk, hg, dv = v.shape[-3:] + return math.prod(batch), sq, sk, h, hg, dq, dv + + +@dataclass(frozen=True) +class AttentionGraph: + """Plan and bindings, including views of packed QKV storage.""" + + graph: object + names: tuple + input_indices: tuple + input_offsets: tuple + output_indices: tuple + output_offsets: tuple + metadata_batch: int + ragged_stats: bool + multipliers: tuple + fp8: bool + + +def get_graph(direction, avals, config): + """Cache by configuration, operand metadata and GPU architecture.""" + shapes = tuple((tuple(aval.shape), jnp.dtype(aval.dtype)) for aval in avals[:5]) + return _build_graph( + direction, + shapes, + config, + device_capabilities(), + bool(int(os.getenv("NVTE_ALLOW_NONDETERMINISTIC_ALGO", "1"))), + bool(int(os.getenv("NVTE_FP8_DPA_BWD", "1"))), + ) + + +@lru_cache(maxsize=256) +def _build_graph(direction, shapes, config, capabilities, nondeterministic, fp8_backward): + cudnn = import_cudnn() + avals = tuple(jax.ShapeDtypeStruct(*shape) for shape in shapes) + batch, sq, sk, h, hg, dq, dv = _dims(avals, config) + backward = direction == "bwd" + ragged = config.qkv_layout.is_thd() + ragged_stats = ragged and capabilities[0] >= 90 and capabilities[0] != 120 + segments = batch * config.max_segments_per_seq if ragged else batch + graph_batch = _bucket(segments, 32, 512) if ragged_stats else segments + graph_sq = _bucket(batch * sq, 1024, 32768) if ragged_stats else sq + graph_sk = _bucket(batch * sk, 1024, 32768) if ragged_stats else sk + fp8 = config.fp8 and (not backward or fp8_backward) + version = int(cudnn.backend_version()) + if ragged and capabilities[0] < 90: + raise ValueError("F16/BF16 fused attention with THD format requires sm90 or later.") + if ( + backward + and ragged_stats + and config.softmax_type.name == "LEARNABLE_SOFTMAX" + and version < 92600 + ): + raise ValueError("THD learnable softmax backward requires cuDNN 9.26.0 or later.") + if fp8 and ragged: + if version < 92300: + raise ValueError("FP8 THD attention requires cuDNN 9.23.0 or later.") + if backward and capabilities[0] < 100: + raise ValueError("FP8 THD backward requires sm100 or later.") + io_dtype = _dtype(cudnn, avals[0].dtype) + qkv_dtype = cudnn.data_type.FP8_E4M3 if fp8 else io_dtype + graph = cudnn.pygraph( + io_data_type=io_dtype, + intermediate_data_type=cudnn.data_type.FLOAT, + compute_data_type=cudnn.data_type.FLOAT, + ) + names, input_uids, input_indices, input_offsets = [], [], [], [] + scalar_uids, scalar_values = [], [] + uid = 0 + + def tensor(name, dim, stride, dtype=None, offset=0): + nonlocal uid + uid += 1 + names.append(name) + input_uids.append(uid) + input_indices.append(len(names) - 1) + input_offsets.append(offset) + return graph.tensor( + name=name, + dim=dim, + stride=stride, + data_type=io_dtype if dtype is None else dtype, + uid=uid, + ) + + def constant(name, value): + nonlocal uid + uid += 1 + scalar_uids.append(uid) + scalar_values.append(np.asarray(value, np.float32).tobytes()) + return graph.tensor( + name=name, + dim=(1, 1, 1, 1), + stride=(1, 1, 1, 1), + data_type=cudnn.data_type.FLOAT, + uid=uid, + is_pass_by_value=True, + ) + + def output(value, dim, stride, slot, offset=0, dtype=None): + nonlocal uid + uid += 1 + value.set_output(True).set_uid(uid).set_dim(dim).set_stride(stride) + value.set_data_type(io_dtype if dtype is None else dtype) + outputs.append((uid, slot, offset)) + return value + + outputs = [] + qpack = 3 if config.qkv_layout.is_qkvpacked() else 1 + kvpack = 3 if qpack == 3 else (2 if config.qkv_layout.is_kvpacked() else 1) + # FP8 conversion preserves the packed storage layout. + itemsize = 1 if fp8 else avals[0].dtype.itemsize + qstride = (graph_sq * qpack * h * dq, dq, qpack * h * dq, 1) + kstride = (graph_sk * kvpack * hg * dq, dq, kvpack * hg * dq, 1) + vstride = (graph_sk * kvpack * hg * dv, dv, kvpack * hg * dv, 1) + ostride = (graph_sq * h * dv, dv, h * dv, 1) + qdim, kdim, vdim = ( + (graph_batch, h, graph_sq, dq), + (graph_batch, hg, graph_sk, dq), + (graph_batch, hg, graph_sk, dv), + ) + odim = (graph_batch, h, graph_sq, dv) + q = tensor("q", qdim, qstride, qkv_dtype) + k = tensor( + "q" if qpack == 3 else "k", kdim, kstride, qkv_dtype, h * dq * itemsize if qpack == 3 else 0 + ) + v = tensor( + "q" if qpack == 3 else ("k" if kvpack == 2 else "v"), + vdim, + vstride, + qkv_dtype, + 2 * h * dq * itemsize if qpack == 3 else (hg * dq * itemsize if kvpack == 2 else 0), + ) + + multipliers = (qpack * h * dq, kvpack * hg * dq, kvpack * hg * dv, h * dv, h) + offsets = [] + if ragged: + for name in ("offset_q", "offset_k", "offset_v", "offset_o", "offset_stats"): + offsets.append( + tensor(name, (graph_batch + 1, 1, 1, 1), (1, 1, 1, 1), cudnn.data_type.INT64) + ) + q.set_ragged_offset(offsets[0]) + k.set_ragged_offset(offsets[1]) + v.set_ragged_offset(offsets[2]) + stats_dim = (graph_batch, h, graph_sq, 1) + stats_stride = (h * graph_sq, 1, h, 1) if ragged_stats else (h * graph_sq, graph_sq, 1, 1) + options = { + "attn_scale": float(config.scaling_factor), + "diagonal_alignment": ( + cudnn.diagonal_alignment.BOTTOM_RIGHT + if config.bottom_right_diagonal + else cudnn.diagonal_alignment.TOP_LEFT + ), + } + left, right = config.effective_window_size + if left >= 0: + options["left_bound" if fp8 else "diagonal_band_left_bound"] = left + 1 + if config.attn_mask_type.is_causal(): + right = 0 + if right >= 0: + options["right_bound" if fp8 else "diagonal_band_right_bound"] = right + if config.attn_mask_type.is_padding() or ragged: + options.update( + use_padding_mask=True, + seq_len_q=tensor("seq_q", (graph_batch, 1, 1, 1), (1, 1, 1, 1), cudnn.data_type.INT32), + seq_len_kv=tensor( + "seq_kv", (graph_batch, 1, 1, 1), (1, 1, 1, 1), cudnn.data_type.INT32 + ), + ) + if config.dropout_probability and config.is_training: + options["dropout"] = ( + float(config.dropout_probability), + tensor("rng", (1, 1, 1, 1), (1, 1, 1, 1), cudnn.data_type.INT64), + tensor("rng", (1, 1, 1, 1), (1, 1, 1, 1), cudnn.data_type.INT64, 8), + ) + bias_type = config.attn_bias_type.name + if bias_type == "ALIBI": + if fp8 and backward: + raise ValueError("cuDNN FP8 backward does not support ALiBi.") + options["use_alibi_mask"] = True + elif bias_type != "NO_BIAS": + if fp8 and backward: + raise ValueError("cuDNN FP8 backward does not support attention bias.") + bias_shape = (*avals[3].shape[-3:],) + bias_shape = (math.prod(avals[3].shape[:-3]), *bias_shape) + options["bias"] = tensor("bias", bias_shape, _stride(bias_shape)) + if backward and bias_shape[:3] != (1, 1, 1): + dbias = graph.tensor( + name="dbias", dim=bias_shape, stride=_stride(bias_shape), data_type=io_dtype + ) + options["dBias"] = output(dbias, bias_shape, _stride(bias_shape), 3) + if config.softmax_type.name != "VANILLA_SOFTMAX": + options["sink_token"] = tensor("sink", (1, h, 1, 1), (h, 1, 1, 1), cudnn.data_type.FLOAT) + if backward: + dsink = graph.tensor( + name="dsink", dim=(1, h, 1, 1), stride=(h, 1, 1, 1), data_type=cudnn.data_type.FLOAT + ) + options["dSink_token"] = output( + dsink, (1, h, 1, 1), (h, 1, 1, 1), 4, dtype=cudnn.data_type.FLOAT + ) + + if fp8: + for name in ("descale_q", "descale_k", "descale_v"): + options[name] = tensor(name, (1, 1, 1, 1), (1, 1, 1, 1), cudnn.data_type.FLOAT) + # Softmax probabilities are in [0, 1]; use the full E4M3 range. + options["scale_s"] = constant("scale_s", 448.0) + options["descale_s"] = constant("descale_s", 1.0 / 448.0) + if backward: + options["use_deterministic_algorithm"] = not nondeterministic + o = tensor("o", odim, ostride) + do = tensor("do", odim, ostride, cudnn.data_type.FP8_E5M2 if fp8 else io_dtype) + stats = tensor("stats", stats_dim, stats_stride, cudnn.data_type.FLOAT) + if ragged: + o.set_ragged_offset(offsets[3]) + do.set_ragged_offset(offsets[3]) + if ragged_stats: + stats.set_ragged_offset(offsets[4]) + if fp8: + options["descale_o"] = constant("descale_o", 1.0) + options["descale_dO"] = tensor( + "descale_do", (1, 1, 1, 1), (1, 1, 1, 1), cudnn.data_type.FLOAT + ) + options["scale_dP"] = constant("scale_dp", 1.0) + options["descale_dP"] = constant("descale_dp", 1.0) + for name in ("scale_dQ", "scale_dK", "scale_dV"): + options[name] = constant(name, 1.0) + gradients = graph.sdpa_fp8_backward(q=q, k=k, v=v, o=o, dO=do, stats=stats, **options) + else: + if ragged_stats: + options["max_total_seq_len_q"] = graph_sq + options["max_total_seq_len_kv"] = graph_sk + gradients = graph.sdpa_backward(q=q, k=k, v=v, o=o, dO=do, stats=stats, **options) + grad_slots = (0, 0 if qpack == 3 else 1, 0 if qpack == 3 else (1 if kvpack == 2 else 2)) + grad_offsets = ( + 0, + h * dq * avals[0].dtype.itemsize if qpack == 3 else 0, + ( + 2 * h * dq * avals[0].dtype.itemsize + if qpack == 3 + else (hg * dq * avals[0].dtype.itemsize if kvpack == 2 else 0) + ), + ) + for index, (grad, dim, stride, slot, offset) in enumerate( + zip( + gradients[:3], + (qdim, kdim, vdim), + (qstride, kstride, vstride), + grad_slots, + grad_offsets, + ) + ): + output(grad, dim, stride, slot, offset) + if ragged: + grad.set_ragged_offset(offsets[index]) + if fp8: + for slot, amax in enumerate(gradients[3:7], 5): + output(amax, (1, 1, 1, 1), (1, 1, 1, 1), slot, dtype=cudnn.data_type.FLOAT) + else: + if config.return_max_logit: + score_max = graph.tensor( + name="max", dim=stats_dim, stride=stats_stride, data_type=cudnn.data_type.FLOAT + ) + if ragged_stats: + score_max.set_ragged_offset(offsets[4]) + options["score_max"] = output( + score_max, stats_dim, stats_stride, 2, dtype=cudnn.data_type.FLOAT + ) + # The public JAX operation remains differentiable with dropout disabled + # by is_training=False, just as the previous common-backed path did. + options["generate_stats"] = True + if fp8: + options["scale_o"] = constant("scale_o", 1.0) + o, stats, amax_s, amax_o = graph.sdpa_fp8(q=q, k=k, v=v, **options) + output(amax_s, (1, 1, 1, 1), (1, 1, 1, 1), 3, dtype=cudnn.data_type.FLOAT) + output(amax_o, (1, 1, 1, 1), (1, 1, 1, 1), 4, dtype=cudnn.data_type.FLOAT) + else: + o, stats = graph.sdpa(q=q, k=k, v=v, **options) + output(o, odim, ostride, 0) + if ragged: + o.set_ragged_offset(offsets[3]) + output(stats, stats_dim, stats_stride, 1, dtype=cudnn.data_type.FLOAT) + if ragged_stats: + stats.set_ragged_offset(offsets[4]) + + workspace, data, version = finalize_graph(cudnn, graph) + serialized = serialize_graph( + serialized_graph=data, + cudnn_frontend_version=version, + workspace_size=workspace, + input_uids=input_uids, + output_uids=[value[0] for value in outputs], + scalar_uids=scalar_uids, + scalar_values=scalar_values, + ) + return AttentionGraph( + serialized, + tuple(names), + tuple(input_indices), + tuple(input_offsets), + tuple(value[1] for value in outputs), + tuple(value[2] for value in outputs), + graph_batch, + ragged_stats, + multipliers, + fp8, + ) + + +def _offset64(offsets, multiplier): + """Compute exact uint64 element offsets without enabling JAX x64 globally.""" + value = offsets.astype(jnp.uint32) + low, high = value & 65535, value >> 16 + ml, mh = multiplier & 65535, multiplier >> 16 + w0 = low * ml + carry = high * ml + (w0 >> 16) + w1 = (carry & 65535) + low * mh + upper = high * mh + (carry >> 16) + (w1 >> 16) + lower = (w1 << 16) | (w0 & 65535) + return jnp.stack((lower, upper), axis=-1) + + +def _quantize(value, dtype): + """Current tensor scaling implemented in JAX, independently of TE common.""" + maximum = jnp.max(jnp.abs(value.astype(jnp.float32))) + scale = jnp.where(maximum > 0, float(jnp.finfo(dtype).max) / maximum, 1.0) + scale = jnp.minimum(scale, jnp.finfo(jnp.float32).max) + return (value.astype(jnp.float32) * scale).astype(dtype), jnp.reciprocal(scale) + + +def _metadata(value, size, fill=0): + """Bind only the configured segment capacity from fixed-size JAX metadata.""" + value = value.reshape(-1)[:size] + return jnp.pad(value, (0, size - value.size), constant_values=fill) + + +def _dense_stats_to_packed(value, offsets, lengths, shape): + """Convert SM120's per-segment BHS1 statistics to JAX's packed TH1 layout.""" + tokens, heads = math.prod(shape[:-2]), shape[-2] + rows = value.shape[-2] + indices = offsets[: lengths.size, None] + jnp.arange(rows)[None, :] + indices = jnp.where(jnp.arange(rows)[None, :] < lengths[:, None], indices, tokens) + values = value.transpose(0, 2, 1, 3).reshape(-1, heads, 1) + packed = jnp.full((tokens, heads, 1), -1e30, value.dtype) + return packed.at[indices.reshape(-1)].set(values, mode="drop").reshape(shape) + + +def _packed_stats_to_dense(value, offsets, lengths, rows): + """Reconstruct dense per-segment statistics for the cuDNN backward graph.""" + heads = value.shape[-2] + indices = offsets[: lengths.size, None] + jnp.arange(rows)[None, :] + valid = jnp.arange(rows)[None, :] < lengths[:, None] + dense = jnp.take(value.reshape(-1, heads, 1), indices, axis=0, mode="fill", fill_value=0) + dense = jnp.where(valid[:, :, None, None], dense, 0) + return dense.transpose(0, 2, 1, 3) + + +def execute(direction, args, config, result_avals): + """Prepare device metadata with JAX and dispatch a cached graph through FFI.""" + plan = get_graph(direction, args, config) + q, k, v, bias, sink = args[:5] + backward = direction == "bwd" + if backward: + stats, rng, o, do, q_cu, kv_cu, q_offsets, kv_offsets = args[5:13] + else: + seed, q_cu, kv_cu, q_offsets, kv_offsets = args[5:10] + # JAX PRNG keys are explicit: replaying a key replays the dropout mask. + rng = jnp.broadcast_to( + jnp.concatenate((seed.reshape(-1)[:2], jnp.zeros(2, jnp.uint32))), result_avals[3].shape + ) + values = {"q": q, "k": k, "v": v, "bias": bias, "sink": sink, "rng": rng} + values["seq_q"] = _metadata(jnp.diff(q_cu), plan.metadata_batch) + values["seq_kv"] = _metadata(jnp.diff(kv_cu), plan.metadata_batch) + dense_ragged_stats = config.qkv_layout.is_thd() and not plan.ragged_stats + if config.qkv_layout.is_thd(): + batch, sq, sk, *_ = _dims(args, config) + q_offsets = _metadata(q_offsets, plan.metadata_batch + 1, batch * sq) + kv_offsets = _metadata(kv_offsets, plan.metadata_batch + 1, batch * sk) + for name, offsets, multiplier in zip( + ("offset_q", "offset_k", "offset_v", "offset_o", "offset_stats"), + (q_offsets, kv_offsets, kv_offsets, q_offsets, q_offsets), + plan.multipliers, + ): + values[name] = _offset64(offsets, multiplier) + if config.fp8: + for name in ("q", "k", "v"): + # Empty dummy operands are not quantized. + if values[name].size: + values[name], values["descale_" + name] = _quantize(values[name], jnp.float8_e4m3fn) + if config.qkv_layout.is_qkvpacked(): + values["descale_k"] = values["descale_v"] = values["descale_q"] + elif config.qkv_layout.is_kvpacked(): + values["descale_v"] = values["descale_k"] + if not plan.fp8: + for name in ("q", "k", "v"): + if values[name].size: + values[name] = ( + values[name].astype(jnp.float32) * values["descale_" + name] + ).astype(q.dtype) + if backward: + values.update(stats=stats, o=o, do=do) + if dense_ragged_stats: + values["stats"] = _packed_stats_to_dense( + stats, + q_offsets, + values["seq_q"], + q.shape[-4] if config.qkv_layout.is_qkvpacked() else q.shape[-3], + ) + if plan.fp8: + values["do"], values["descale_do"] = _quantize(do, jnp.float8_e5m2) + specs = list(result_avals[:-1]) + if plan.fp8: + specs += [jax.ShapeDtypeStruct((1,), jnp.float32)] * 4 + else: + specs = list(result_avals[:3]) + if dense_ragged_stats: + _, sq, _, heads, *_ = _dims(args, config) + dense_shape = (plan.metadata_batch, heads, sq, 1) + specs[1] = jax.ShapeDtypeStruct(dense_shape, jnp.float32) + if config.return_max_logit: + specs[2] = jax.ShapeDtypeStruct(dense_shape, jnp.float32) + if plan.fp8: + specs += [jax.ShapeDtypeStruct((1,), jnp.float32)] * 2 + specs.append(jax.ShapeDtypeStruct((plan.graph.workspace_size,), jnp.uint8)) + attributes = graph_attributes(plan.graph) + attributes.update( + input_indices=np.asarray(plan.input_indices, np.int64), + input_offsets=np.asarray(plan.input_offsets, np.int64), + output_indices=np.asarray(plan.output_indices, np.int64), + output_offsets=np.asarray(plan.output_offsets, np.int64), + ) + if not backward and config.qkv_layout.is_thd(): + # Match the neutral finite statistics value for empty ragged ring blocks. + # Filling FP32 bytes with 0xF0 gives a large negative number without NaNs + # when the ring correction subtracts statistics for two empty blocks. + fills = np.zeros(len(specs) - 1, np.int64) + fills[1:3] = 0xF0 + attributes["output_fill_values"] = fills + results = ffi.ffi_call( + "te_cudnn_graph_ffi", tuple(jax.ShapeDtypeStruct(spec.shape, spec.dtype) for spec in specs) + )(*(values[name] for name in plan.names), **attributes) + if backward: + return (*results[:5], results[-1]) + output, stats, score_max = results[:3] + if dense_ragged_stats: + packed_shape = result_avals[1].shape + stats = _dense_stats_to_packed(stats, q_offsets, values["seq_q"], packed_shape) + if config.return_max_logit: + score_max = _dense_stats_to_packed(score_max, q_offsets, values["seq_q"], packed_shape) + score_max = jnp.swapaxes(score_max, -3, -2) + return output, stats, score_max, rng, results[-1] diff --git a/transformer_engine/jax/cpp_extensions/cudnn_graph.py b/transformer_engine/jax/cpp_extensions/cudnn_graph.py new file mode 100644 index 00000000000..e5acc167c59 --- /dev/null +++ b/transformer_engine/jax/cpp_extensions/cudnn_graph.py @@ -0,0 +1,139 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. +"""Serialize Python-built cuDNN graphs for the JAX-owned FFI executor.""" + +import hashlib +import importlib +from dataclasses import dataclass +from typing import Sequence, Tuple + +import numpy as np +import transformer_engine_jax + + +@dataclass(frozen=True) +class SerializedGraph: + """A compiled graph and its buffer bindings.""" + + serialized_graph: bytes + graph_hash: Tuple[int, int] + cudnn_frontend_version: int + workspace_size: int + input_uids: np.ndarray + output_uids: np.ndarray + scalar_uids: np.ndarray + scalar_sizes: np.ndarray + scalar_values: np.ndarray + + +def encode_frontend_version(version: str) -> int: + """Encode the public version as CUDNN_FRONTEND_VERSION.""" + parts = version.split("+", 1)[0].split("-", 1)[0].split(".") + if len(parts) < 3: + raise RuntimeError(f"Could not parse cuDNN frontend Python version: {version!r}.") + major, minor, patch = (int(part) for part in parts[:3]) + return major * 10000 + minor * 100 + patch + + +def check_frontend_version(cudnn) -> int: + """Require compatible Python and C++ serialization formats.""" + version = getattr(cudnn, "__version__", None) + if version is None: + raise RuntimeError("cuDNN frontend Python package does not expose __version__.") + python_version = encode_frontend_version(version) + cpp_version = int(transformer_engine_jax.get_cudnn_frontend_version()) + if python_version != cpp_version: + raise RuntimeError( + "cuDNN frontend Python/C++ version mismatch for graph serialization: " + f"Python cudnn.__version__={version!r} encodes to {python_version}, " + f"but Transformer Engine C++ was built with CUDNN_FRONTEND_VERSION={cpp_version}. " + "Use matching cuDNN frontend Python package and C++ headers." + ) + return python_version + + +def import_cudnn(): + """Load the required Python frontend lazily.""" + try: + cudnn = importlib.import_module("cudnn") + except ImportError as exc: + raise ImportError( + "JAX fused_attn requires the cuDNN frontend Python package (`cudnn`)." + ) from exc + check_frontend_version(cudnn) + return cudnn + + +def graph_hash(serialized_graph: bytes) -> Tuple[int, int]: + """Identify a serialized plan independently of Python object lifetimes.""" + digest = hashlib.sha256(serialized_graph).digest() + return tuple(int.from_bytes(digest[i : i + 8], "little", signed=True) for i in (0, 8)) + + +def pack_scalar_values(values: Sequence[bytes]) -> Tuple[np.ndarray, np.ndarray]: + """Pack pass-by-value constants into aligned slots.""" + sizes = np.asarray([len(value) for value in values], dtype=np.int64) + packed = np.zeros((len(values), 16), dtype=np.uint8) + for index, value in enumerate(values): + if len(value) > 16: + raise ValueError("cuDNN pass-by-value scalars must be at most 16 bytes.") + packed[index, : len(value)] = np.frombuffer(value, dtype=np.uint8) + return sizes, packed.reshape(-1) + + +def serialize_graph( + *, + serialized_graph, + cudnn_frontend_version, + workspace_size, + input_uids, + output_uids, + scalar_uids=(), + scalar_values=(), +) -> SerializedGraph: + """Package a graph and its stable tensor UIDs.""" + sizes, values = pack_scalar_values(scalar_values) + return SerializedGraph( + serialized_graph, + graph_hash(serialized_graph), + int(cudnn_frontend_version), + int(workspace_size), + np.asarray(input_uids, dtype=np.int64), + np.asarray(output_uids, dtype=np.int64), + np.asarray(scalar_uids, dtype=np.int64), + sizes, + values, + ) + + +def finalize_graph(cudnn, graph): + """Validate, select and compile a supported execution plan.""" + try: + graph.validate() + graph.build_operation_graph() + graph.create_execution_plans([cudnn.heur_mode.A, cudnn.heur_mode.FALLBACK]) + graph.check_support() + graph.build_plans(cudnn.build_plan_policy.HEURISTICS_CHOICE) + except cudnn.cudnnGraphNotSupportedError as exc: + raise RuntimeError(f"cuDNN SDPA graph is not supported: {exc}") from exc + return ( + max(int(graph.get_workspace_size()), 1), + bytes(graph.serialize()), + check_frontend_version(cudnn), + ) + + +def graph_attributes(graph): + """Attributes shared by ordinary and score-modified attention.""" + return { + "serialized_graph": graph.serialized_graph, + "graph_hash0": graph.graph_hash[0], + "graph_hash1": graph.graph_hash[1], + "cudnn_frontend_version": graph.cudnn_frontend_version, + "input_uids": graph.input_uids, + "output_uids": graph.output_uids, + "scalar_uids": graph.scalar_uids, + "scalar_sizes": graph.scalar_sizes, + "scalar_values": graph.scalar_values, + } diff --git a/transformer_engine/jax/cpp_extensions/flex_attention.py b/transformer_engine/jax/cpp_extensions/flex_attention.py index ff9f5fb1460..d3789f39f20 100644 --- a/transformer_engine/jax/cpp_extensions/flex_attention.py +++ b/transformer_engine/jax/cpp_extensions/flex_attention.py @@ -3,8 +3,6 @@ # See LICENSE for license information. """cuDNN frontend score_mod fused attention helpers.""" -import hashlib -import importlib import inspect import os from dataclasses import dataclass @@ -15,7 +13,16 @@ import numpy as np from jax import ffi -import transformer_engine_jax +from .cudnn_graph import ( + SerializedGraph as _SerializedScoreModGraph, + encode_frontend_version as _encode_cudnn_frontend_version, + check_frontend_version as _check_cudnn_frontend_version_match, + graph_hash as _score_mod_graph_hash, + pack_scalar_values as _pack_score_mod_scalar_values, + serialize_graph as _serialized_score_mod_graph, + finalize_graph as _finalize_score_mod_graph, + import_cudnn as _import_cudnn_for_score_mod, +) __all__ = [ "FusedAttnScoreModHelper", @@ -254,21 +261,6 @@ def __eq__(self, other): ) -@dataclass(frozen=True) -class _SerializedScoreModGraph: - """Serialized cuDNN frontend graph and static metadata for C++ execution.""" - - serialized_graph: bytes - graph_hash: Tuple[int, int] - cudnn_frontend_version: int - workspace_size: int - input_uids: np.ndarray - output_uids: np.ndarray - scalar_uids: np.ndarray - scalar_sizes: np.ndarray - scalar_values: np.ndarray - - # cuDNN frontend tensor UIDs are arbitrary, but assigning stable values makes serialized # graphs deterministic and simplifies debugging across the Python graph builder and C++ executor. _SCORE_MOD_UID_Q = 1 @@ -477,75 +469,6 @@ def _score_mod_graph_tensors( return graph_tensors, tuple(tensor_uids), tuple(scalar_uids), tuple(scalar_values) -def _encode_cudnn_frontend_version(version: str) -> int: - public_version = version.split("+", 1)[0].split("-", 1)[0] - parts = public_version.split(".") - if len(parts) < 3: - raise RuntimeError(f"Could not parse cuDNN frontend Python version: {version!r}.") - major, minor, patch = (int(part) for part in parts[:3]) - return major * 10000 + minor * 100 + patch - - -def _check_cudnn_frontend_version_match(cudnn) -> int: - python_version_string = getattr(cudnn, "__version__", None) - if python_version_string is None: - raise RuntimeError("cuDNN frontend Python package does not expose __version__.") - python_version = _encode_cudnn_frontend_version(python_version_string) - cpp_version = int(transformer_engine_jax.get_cudnn_frontend_version()) - if python_version != cpp_version: - raise RuntimeError( - "cuDNN frontend Python/C++ version mismatch for score_mod graph serialization: " - f"Python cudnn.__version__={python_version_string!r} encodes to {python_version}, " - f"but Transformer Engine C++ was built with CUDNN_FRONTEND_VERSION={cpp_version}. " - "Use matching cuDNN frontend Python package and C++ headers." - ) - return python_version - - -def _score_mod_graph_hash(serialized_graph: bytes) -> Tuple[int, int]: - digest = hashlib.sha256(serialized_graph).digest() - return ( - int.from_bytes(digest[0:8], byteorder="little", signed=True), - int.from_bytes(digest[8:16], byteorder="little", signed=True), - ) - - -def _pack_score_mod_scalar_values( - scalar_values: Sequence[bytes], -) -> Tuple[np.ndarray, np.ndarray]: - scalar_sizes = np.asarray([len(value) for value in scalar_values], dtype=np.int64) - packed_values = np.zeros((len(scalar_values), 16), dtype=np.uint8) - for index, value in enumerate(scalar_values): - if len(value) > 16: - raise ValueError("score_mod pass-by-value scalars must be at most 16 bytes.") - packed_values[index, : len(value)] = np.frombuffer(value, dtype=np.uint8) - return scalar_sizes, packed_values.reshape(-1) - - -def _serialized_score_mod_graph( - *, - serialized_graph: bytes, - cudnn_frontend_version: int, - workspace_size: int, - input_uids: Sequence[int], - output_uids: Sequence[int], - scalar_uids: Sequence[int], - scalar_values: Sequence[bytes], -) -> _SerializedScoreModGraph: - scalar_sizes, packed_scalar_values = _pack_score_mod_scalar_values(scalar_values) - return _SerializedScoreModGraph( - serialized_graph=serialized_graph, - graph_hash=_score_mod_graph_hash(serialized_graph), - cudnn_frontend_version=int(cudnn_frontend_version), - workspace_size=int(workspace_size), - input_uids=np.asarray(input_uids, dtype=np.int64), - output_uids=np.asarray(output_uids, dtype=np.int64), - scalar_uids=np.asarray(scalar_uids, dtype=np.int64), - scalar_sizes=scalar_sizes, - scalar_values=packed_scalar_values, - ) - - def _wrap_score_mod(score_mod: Optional[Callable], graph_tensors: Dict[str, Any]): if score_mod is None: return None @@ -556,23 +479,6 @@ def wrapped_score_mod(sdpa_graph, score_tensor): return wrapped_score_mod -def _finalize_score_mod_graph(cudnn, graph) -> Tuple[int, bytes, int]: - graph.validate() - graph.build_operation_graph() - try: - graph.create_execution_plans([cudnn.heur_mode.A, cudnn.heur_mode.FALLBACK]) - graph.check_support() - except cudnn.cudnnGraphNotSupportedError as exc: - raise RuntimeError(f"cuDNN score_mod SDPA graph is not supported: {exc}") from exc - graph.build_plans(cudnn.build_plan_policy.HEURISTICS_CHOICE) - serialized_graph = bytes(graph.serialize()) - return ( - max(int(graph.get_workspace_size()), 1), - serialized_graph, - _check_cudnn_frontend_version_match(cudnn), - ) - - def _graph_cache_key( direction: str, config: _FusedAttnScoreModConfig, @@ -593,17 +499,6 @@ def _shape_dtype(value) -> jax.ShapeDtypeStruct: return jax.ShapeDtypeStruct(tuple(value.shape), value.dtype) -def _import_cudnn_for_score_mod(): - try: - cudnn = importlib.import_module("cudnn") - except ImportError as exc: - raise ImportError( - "score_mod fused_attn requires the cuDNN frontend Python package (`cudnn`)." - ) from exc - _check_cudnn_frontend_version_match(cudnn) - return cudnn - - def _build_score_mod_fwd_graph(q_aval, k_aval, v_aval, score_mod_avals, config): cudnn = _import_cudnn_for_score_mod() @@ -813,7 +708,7 @@ def _fused_attn_score_mod_fwd( stats = jax.ShapeDtypeStruct(stats_shape, jnp.float32) workspace = jax.ShapeDtypeStruct((graph.workspace_size,), jnp.uint8) output, softmax_stats, _ = ffi.ffi_call( - "te_fused_attn_score_mod_forward_ffi", + "te_cudnn_graph_ffi", (output_shape, stats, workspace), )( q, @@ -873,7 +768,7 @@ def _fused_attn_score_mod_bwd( dv = jax.ShapeDtypeStruct(v.shape, v.dtype) workspace = jax.ShapeDtypeStruct((graph.workspace_size,), jnp.uint8) dq, dk, dv, _ = ffi.ffi_call( - "te_fused_attn_score_mod_backward_ffi", + "te_cudnn_graph_ffi", (dq, dk, dv, workspace), )( q, diff --git a/transformer_engine/jax/csrc/extensions.h b/transformer_engine/jax/csrc/extensions.h index 2237b5c204e..2636a4b306d 100644 --- a/transformer_engine/jax/csrc/extensions.h +++ b/transformer_engine/jax/csrc/extensions.h @@ -150,32 +150,8 @@ XLA_FFI_DECLARE_HANDLER_SYMBOL(ScaledUpperTriangMaskedSoftmaxBackwardHandler); XLA_FFI_DECLARE_HANDLER_SYMBOL(ScaledUpperTriangMaskedSoftmaxBackwardInitializeHandler); // Attention -XLA_FFI_DECLARE_HANDLER_SYMBOL(FusedAttnForwardHandler); - -XLA_FFI_DECLARE_HANDLER_SYMBOL(FusedAttnBackwardHandler); - -XLA_FFI_DECLARE_HANDLER_SYMBOL(FusedAttnScoreModForwardHandler); - -XLA_FFI_DECLARE_HANDLER_SYMBOL(FusedAttnScoreModBackwardHandler); - -std::tuple GetFusedAttnBackend( - const pybind11::object ¶ms); - -pybind11::tuple GetFusedAttnForwardWorkspaceSizes( - size_t input_batch, size_t bias_batch, size_t q_max_seqlen, size_t kv_max_seqlen, - size_t attn_heads, size_t num_gqa_groups, size_t bias_heads, size_t qk_head_dim, - size_t v_head_dim, float scaling_factor, float dropout_probability, NVTE_Bias_Type bias_type, - NVTE_Mask_Type mask_type, NVTE_Softmax_Type softmax_type, NVTE_QKV_Layout qkv_layout, - DType dtype, bool is_training, size_t max_segments_per_seq, int64_t window_size_left, - int64_t window_size_right, bool return_max_logit, bool bottom_right_diagonal); - -pybind11::tuple GetFusedAttnBackwardWorkspaceSizes( - size_t input_batch, size_t bias_batch, size_t q_max_seqlen, size_t kv_max_seqlen, - size_t attn_heads, size_t num_gqa_groups, size_t bias_heads, size_t qk_head_dim, - size_t v_head_dim, float scaling_factor, float dropout_probability, NVTE_Bias_Type bias_type, - NVTE_Mask_Type mask_type, NVTE_Softmax_Type softmax_type, NVTE_QKV_Layout qkv_layout, - DType dtype, bool is_training, bool deterministic, size_t max_segments_per_seq, - int64_t window_size_left, int64_t window_size_right, bool bottom_right_diagonal); +XLA_FFI_DECLARE_HANDLER_SYMBOL(CudnnGraphHandler); +XLA_FFI_DECLARE_HANDLER_SYMBOL(CudnnGraphInitializeHandler); // GEMM XLA_FFI_DECLARE_HANDLER_SYMBOL(GemmHandler); diff --git a/transformer_engine/jax/csrc/extensions/attention.cpp b/transformer_engine/jax/csrc/extensions/attention.cpp index def5e0149b8..4d4c2c522b3 100644 --- a/transformer_engine/jax/csrc/extensions/attention.cpp +++ b/transformer_engine/jax/csrc/extensions/attention.cpp @@ -4,1111 +4,192 @@ * See LICENSE for license information. ************************************************************************/ +// JAX-owned execution of plans built with the cuDNN Python frontend. +// This translation unit deliberately has no TE common dependencies. +#include +#include #include +#include +#include #include #include -#include -#include #include #include -#include +#include +#include #include -#include #include -#include "../extensions.h" -#include "common/cudnn_utils.h" -#include "transformer_engine/fused_attn.h" -#include "transformer_engine/transformer_engine.h" - namespace transformer_engine { namespace jax { +namespace { -static std::tuple GetFusedAttnBackendImpl( - const FusedAttnConfigWrapper &cfg) { - const char *message = nullptr; - auto backend = nvte_get_fused_attn_backend_v2(cfg, &message); - return {backend, std::string(message)}; -} - -std::tuple GetFusedAttnBackend( - const pybind11::object ¶ms) { - const auto qkv_layout = params.attr("qkv_layout").cast(); - auto o_format = params.attr("o_format").cast(); - auto do_format = params.attr("do_format").cast(); - auto dqkv_layout = params.attr("dqkv_layout").cast(); - if (o_format == NVTE_QKV_Format::NVTE_QKV_Format_NOT_SET) { - o_format = nvte_get_q_format(qkv_layout); - } - if (do_format == NVTE_QKV_Format::NVTE_QKV_Format_NOT_SET) { - do_format = o_format; - } - if (dqkv_layout == NVTE_QKV_Layout::NVTE_QKV_Layout_NOT_SET) { - dqkv_layout = qkv_layout; - } - - FusedAttnConfigWrapper cfg; - cfg.set_is_training(params.attr("is_training").cast()) - .set_deterministic(params.attr("deterministic").cast()) - .set_cuda_graph(params.attr("cuda_graph").cast()) - .set_return_max_logit(params.attr("return_max_logit").cast()) - .set_attn_mask_type(params.attr("attn_mask_type").cast()) - .set_bias_type(params.attr("bias_type").cast()) - .set_window_size_left(params.attr("window_size_left").cast()) - .set_window_size_right(params.attr("window_size_right").cast()) - .set_bottom_right_diagonal(params.attr("bottom_right_diagonal").cast()) - .set_softmax_type(params.attr("softmax_type").cast()) - .set_scaling_mode( - get_nvte_scaling_mode(params.attr("scaling_mode").cast())) - .set_dropout(params.attr("dropout").cast()) - .set_attn_scale(params.attr("attn_scale").cast()) - .set_qkv_dtype(static_cast(params.attr("qkv_dtype").cast())) - .set_o_dtype(static_cast(params.attr("o_dtype").cast())) - .set_do_dtype(static_cast(params.attr("do_dtype").cast())) - .set_dqkv_dtype(static_cast(params.attr("dqkv_dtype").cast())) - .set_qkv_layout(qkv_layout) - .set_o_format(o_format) - .set_do_format(do_format) - .set_dqkv_layout(dqkv_layout) - .set_qkv_scale_inv_format(params.attr("qkv_scale_inv_format").cast()) - .set_do_scale_inv_format(params.attr("do_scale_inv_format").cast()) - .set_batch_size(params.attr("batch_size").cast()) - .set_num_attn_heads(params.attr("num_attn_heads").cast()) - .set_num_gqa_groups(params.attr("num_gqa_groups").cast()) - .set_head_dim_qk(params.attr("head_dim_qk").cast()) - .set_head_dim_v(params.attr("head_dim_v").cast()) - .set_max_seqlen_q(params.attr("max_seqlen_q").cast()) - .set_max_seqlen_kv(params.attr("max_seqlen_kv").cast()) - .set_num_tokens_q(params.attr("num_tokens_q").cast()) - .set_num_tokens_kv(params.attr("num_tokens_kv").cast()) - .set_bias_batch_size(params.attr("bias_batch_size").cast()) - .set_bias_num_heads(params.attr("bias_num_heads").cast()) - .set_bias_seqlen_q(params.attr("bias_seqlen_q").cast()) - .set_bias_seqlen_kv(params.attr("bias_seqlen_kv").cast()); - return GetFusedAttnBackendImpl(cfg); -} - -/* - NOTE: PrepareFusedAttnForwardAuxTensors unifies the auxiliary tensor pack logic from the fused - attention forward kernels in: - - common/fused_attn/fused_attn_f16_arbitrary_seqlen.cu lines 1270-1281 and 1348-1359 -*/ -void PrepareFusedAttnForwardAuxTensors(NVTETensorPack *tensor_pack, const size_t input_batch, - const size_t bias_batch, const size_t attn_heads, - const size_t bias_heads, const size_t q_max_seqlen, - const size_t kv_max_seqlen, DType dtype, - NVTE_Bias_Type bias_type, NVTE_Fused_Attn_Backend backend, - void *softmax_buf, void *max_logits_buf = nullptr, - void *rng_state_buf = nullptr, void *bias_buf = nullptr, - void *softmax_offset_buf = nullptr) { - // all backends need softmax but expect different shapes/dtypes - tensor_pack->size = 1; - NVTETensor &softmax_aux = tensor_pack->tensors[0]; - NVTEBasicTensor softmax_aux_data; - softmax_aux_data.data_ptr = softmax_buf; - softmax_aux_data.shape.ndim = 4; - softmax_aux_data.shape.data[0] = input_batch; - softmax_aux_data.shape.data[1] = attn_heads; - softmax_aux_data.shape.data[2] = q_max_seqlen; - softmax_aux_data.shape.data[3] = kv_max_seqlen; - softmax_aux_data.dtype = static_cast(dtype); - - // arbitrary sequence length backend needs the RNG state and a different shape/dtype softmax - if (backend == NVTE_Fused_Attn_Backend::NVTE_F16_arbitrary_seqlen) { - int size = 1; // Start after softmax. - auto next_aux_tensor = [&]() -> NVTETensor & { - NVTE_CHECK(size < NVTETensorPack::MAX_SIZE, - "Fused attention auxiliary tensor pack capacity exceeded."); - return tensor_pack->tensors[size++]; - }; - - if (max_logits_buf != nullptr) { - NVTETensor &max_aux = next_aux_tensor(); - NVTEBasicTensor max_aux_data; - max_aux_data.data_ptr = max_logits_buf; - max_aux_data.shape = {}; - max_aux_data.shape.ndim = 4; - max_aux_data.shape.data[0] = input_batch; - max_aux_data.shape.data[1] = attn_heads; - max_aux_data.shape.data[2] = q_max_seqlen; - max_aux_data.shape.data[3] = 1; - max_aux_data.dtype = static_cast(DType::kFloat32); - nvte_set_tensor_param(&max_aux, kNVTERowwiseData, &max_aux_data); - } - - NVTETensor &rng_state_aux = next_aux_tensor(); - NVTEBasicTensor rng_state_aux_data; - rng_state_aux_data.data_ptr = rng_state_buf; - rng_state_aux_data.shape = {}; - rng_state_aux_data.shape.ndim = 2; - rng_state_aux_data.dtype = static_cast(DType::kInt64); - nvte_set_tensor_param(&rng_state_aux, kNVTERowwiseData, &rng_state_aux_data); - // correct softmax shape/dtype - softmax_aux_data.shape.data[3] = 1; // {B,H,Qs,Ks} -> {B,H,Qs,1} - softmax_aux_data.dtype = static_cast(DType::kFloat32); - - // include bias if enabled - if (bias_type != NVTE_Bias_Type::NVTE_NO_BIAS && bias_type != NVTE_Bias_Type::NVTE_ALIBI) { - NVTETensor &bias_aux = next_aux_tensor(); - NVTEBasicTensor bias_aux_data; - bias_aux_data.data_ptr = bias_buf; - bias_aux_data.shape.ndim = 4; - bias_aux_data.shape.data[0] = bias_batch; - bias_aux_data.shape.data[1] = bias_heads; - bias_aux_data.shape.data[2] = q_max_seqlen; - bias_aux_data.shape.data[3] = kv_max_seqlen; - bias_aux_data.dtype = static_cast(dtype); - nvte_set_tensor_param(&bias_aux, kNVTERowwiseData, &bias_aux_data); - } - - // include softmax_offset if provided - if (softmax_offset_buf != nullptr) { - NVTETensor &softmax_offset_aux = next_aux_tensor(); - NVTEBasicTensor softmax_offset_aux_data; - softmax_offset_aux_data.data_ptr = softmax_offset_buf; - softmax_offset_aux_data.shape.ndim = 4; - softmax_offset_aux_data.shape.data[0] = 1; - softmax_offset_aux_data.shape.data[1] = attn_heads; - softmax_offset_aux_data.shape.data[2] = 1; - softmax_offset_aux_data.shape.data[3] = 1; - softmax_offset_aux_data.dtype = static_cast(DType::kFloat32); - nvte_set_tensor_param(&softmax_offset_aux, kNVTERowwiseData, &softmax_offset_aux_data); - } - - // Set final size - tensor_pack->size = size; - } - nvte_set_tensor_param(&softmax_aux, kNVTERowwiseData, &softmax_aux_data); -} - -/* - NOTE: Backward fused attention kernels accept auxiliary tensors as explicit function arguments - instead of an NVTETensorPack and nvte_fused_attn_bwd() API does all the logic for pulling the - necessary tensors out of the tensor pack for the active kernel. That means we can just dump - everything we got into the tensor pack and not worry about its sizing for the backward pass. - - TODO(Alp): Refactor the nvte_fused_attn_fwd() to work like nvte_fused_attn_bwd()? -*/ -void PrepareFusedAttnBackwardAuxTensors(NVTETensorPack *tensor_pack, const size_t input_batch, - const size_t bias_batch, const size_t attn_heads, - const size_t bias_heads, const size_t q_max_seqlen, - const size_t kv_max_seqlen, DType dtype, - NVTE_Fused_Attn_Backend backend, void *softmax_buf, - void *rng_state_buf, void *bias_buf, - void *softmax_offset_buf = nullptr) { - // Backward calls put everything into the tensor pack for every backend - // so we set dummy bias_type and backend choices here to follow the correct code path - auto dummy_bias_type = NVTE_Bias_Type::NVTE_POST_SCALE_BIAS; - auto dummy_backend = NVTE_Fused_Attn_Backend::NVTE_F16_arbitrary_seqlen; - PrepareFusedAttnForwardAuxTensors(tensor_pack, input_batch, bias_batch, attn_heads, bias_heads, - q_max_seqlen, kv_max_seqlen, dtype, dummy_bias_type, - dummy_backend, softmax_buf, nullptr, rng_state_buf, bias_buf, - softmax_offset_buf); -} - -pybind11::tuple GetFusedAttnForwardWorkspaceSizes( - size_t input_batch, size_t bias_batch, size_t q_max_seqlen, size_t kv_max_seqlen, - size_t attn_heads, size_t num_gqa_groups, size_t bias_heads, size_t qk_head_dim, - size_t v_head_dim, float scaling_factor, float dropout_probability, NVTE_Bias_Type bias_type, - NVTE_Mask_Type mask_type, NVTE_Softmax_Type softmax_type, NVTE_QKV_Layout qkv_layout, - DType dtype, bool is_training, size_t max_segments_per_seq, int64_t window_size_left, - int64_t window_size_right, bool return_max_logit, bool bottom_right_diagonal) { - auto is_ragged = nvte_get_qkv_format(qkv_layout) == NVTE_QKV_Format::NVTE_THD; - auto q_shape = is_ragged - ? std::vector{input_batch * q_max_seqlen, attn_heads, qk_head_dim} - : std::vector{input_batch, q_max_seqlen, attn_heads, qk_head_dim}; - auto q_tensor = TensorWrapper(nullptr, q_shape, dtype); - auto k_shape = is_ragged - ? std::vector{input_batch * kv_max_seqlen, num_gqa_groups, qk_head_dim} - : std::vector{input_batch, kv_max_seqlen, num_gqa_groups, qk_head_dim}; - auto k_tensor = TensorWrapper(nullptr, k_shape, dtype); - auto v_shape = is_ragged - ? std::vector{input_batch * kv_max_seqlen, num_gqa_groups, v_head_dim} - : std::vector{input_batch, kv_max_seqlen, num_gqa_groups, v_head_dim}; - auto v_tensor = TensorWrapper(nullptr, v_shape, dtype); - auto o_shape = is_ragged ? std::vector{input_batch * q_max_seqlen, attn_heads, v_head_dim} - : std::vector{input_batch, q_max_seqlen, attn_heads, v_head_dim}; - - auto bias_shape = std::vector{bias_batch, bias_heads, q_max_seqlen, kv_max_seqlen}; - auto bias_tensor = TensorWrapper(nullptr, bias_shape, dtype); - - // F16 doesn't use this tensor - auto s_tensor = TensorWrapper(nullptr, std::vector{1}, dtype); - auto o_tensor = TensorWrapper(nullptr, o_shape, dtype); - - auto dummy_rng_state_tensor = TensorWrapper(nullptr, std::vector{2}, DType::kInt64); - auto dummy_page_table_tensor = TensorWrapper(nullptr, std::vector{1}, DType::kInt32); - auto dummy_softmax_offset_tensor = - TensorWrapper(nullptr, std::vector{1}, DType::kFloat32); - - NVTETensorPack aux_output_tensors; - nvte_tensor_pack_create(&aux_output_tensors); - - TensorWrapper query_workspace_tensor; - // It is a WAR to pre-create all possible cuDNN graph at the JIT compile time - size_t max_num_segments = is_ragged ? input_batch * max_segments_per_seq : input_batch; - size_t min_num_segments = is_ragged ? input_batch * max_segments_per_seq : input_batch; - for (auto num_segments = min_num_segments; num_segments <= max_num_segments; ++num_segments) { - // the last one is the largest which will be the returned workspace size - auto q_cu_seqlens_tensor = - TensorWrapper(nullptr, std::vector{num_segments + 1}, DType::kInt32); - auto kv_cu_seqlens_tensor = - TensorWrapper(nullptr, std::vector{num_segments + 1}, DType::kInt32); - auto ragged_offset_tensor = - TensorWrapper(nullptr, std::vector{num_segments + 1}, DType::kInt32); - FusedAttnFwdParamsWrapper params; - params.set_Q(q_tensor.data()) - .set_K(k_tensor.data()) - .set_V(v_tensor.data()) - .set_Bias(bias_tensor.data()) - .set_SoftmaxOffset(dummy_softmax_offset_tensor.data()) - .set_S(s_tensor.data()) - .set_O(o_tensor.data()) - .set_Aux_CTX_Tensors(&aux_output_tensors) - .set_cu_seqlens_q(q_cu_seqlens_tensor.data()) - .set_cu_seqlens_kv(kv_cu_seqlens_tensor.data()) - .set_cu_seqlens_q_padded(ragged_offset_tensor.data()) - .set_cu_seqlens_kv_padded(ragged_offset_tensor.data()) - .set_page_table_k(dummy_page_table_tensor.data()) - .set_page_table_v(dummy_page_table_tensor.data()) - .set_rng_state(dummy_rng_state_tensor.data()) - .set_is_training(is_training) - .set_cuda_graph(false) - .set_return_max_logit(return_max_logit) - .set_attn_mask_type(mask_type) - .set_bias_type(bias_type) - .set_window_size_left(window_size_left) - .set_window_size_right(window_size_right) - .set_bottom_right_diagonal(bottom_right_diagonal) - .set_softmax_type(softmax_type) - .set_dropout(dropout_probability) - .set_attn_scale(scaling_factor) - .set_qkv_layout(qkv_layout) - .set_o_format(nvte_get_q_format(qkv_layout)) - .set_qkv_scale_inv_format(NVTE_QKV_Format_NOT_SET) - .set_max_seqlen_q(q_max_seqlen) - .set_max_seqlen_kv(kv_max_seqlen) - .set_workspace(query_workspace_tensor.data()) - .set_stream(nullptr); - nvte_fused_attn_fwd_v2(params); - } - - nvte_tensor_pack_destroy(&aux_output_tensors); - - auto workspace_shape = MakeShapeVector(query_workspace_tensor.shape()); - return pybind11::make_tuple(workspace_shape, query_workspace_tensor.dtype()); -} - -#define FUSED_ATTN_IMPL_COMMON_BLOCK \ - auto is_ragged = nvte_get_qkv_format(qkv_layout) == NVTE_QKV_Format::NVTE_THD; \ - auto bias_shape = std::vector{bias_batch, bias_heads, q_max_seqlen, kv_max_seqlen}; \ - const bool has_bias_tensor = \ - bias_type != NVTE_Bias_Type::NVTE_NO_BIAS && bias_type != NVTE_Bias_Type::NVTE_ALIBI; \ - const size_t bias_seqlen_q = has_bias_tensor ? q_max_seqlen : 0; \ - const size_t bias_seqlen_kv = has_bias_tensor ? kv_max_seqlen : 0; \ - size_t num_segments = input_batch; \ - if (is_ragged) { \ - num_segments = input_batch * max_segments_per_seq; \ - } \ - std::vector seq_shape{num_segments + 1}; \ - auto q_cu_seqlens_tensor = TensorWrapper(q_cu_seqlens, seq_shape, DType::kInt32); \ - auto kv_cu_seqlens_tensor = TensorWrapper(kv_cu_seqlens, seq_shape, DType::kInt32); \ - auto q_seq_offsets_tensor = TensorWrapper(q_seq_offsets, seq_shape, DType::kInt32); \ - auto k_seq_offsets_tensor = TensorWrapper(k_seq_offsets, seq_shape, DType::kInt32); \ - auto workspace_tensor = \ - TensorWrapper(workspace, std::vector{wkspace_size}, wkspace_dtype); \ - auto layout_group = nvte_get_qkv_layout_group(qkv_layout); \ - FusedAttnConfigWrapper cfg; \ - cfg.set_is_training(is_training) \ - .set_deterministic(deterministic) \ - .set_cuda_graph(false) \ - .set_return_max_logit(false) \ - .set_attn_mask_type(mask_type) \ - .set_bias_type(bias_type) \ - .set_window_size_left(window_size_left) \ - .set_window_size_right(window_size_right) \ - .set_bottom_right_diagonal(bottom_right_diagonal) \ - .set_softmax_type(softmax_type) \ - .set_scaling_mode(get_nvte_scaling_mode(JAXX_Scaling_Mode::NO_SCALING)) \ - .set_dropout(dropout_probability) \ - .set_attn_scale(scaling_factor) \ - .set_qkv_dtype(static_cast(dtype)) \ - .set_o_dtype(static_cast(dtype)) \ - .set_do_dtype(static_cast(dtype)) \ - .set_dqkv_dtype(static_cast(dtype)) \ - .set_qkv_layout(qkv_layout) \ - .set_o_format(nvte_get_q_format(qkv_layout)) \ - .set_do_format(nvte_get_q_format(qkv_layout)) \ - .set_dqkv_layout(qkv_layout) \ - .set_qkv_scale_inv_format(NVTE_QKV_Format::NVTE_QKV_Format_NOT_SET) \ - .set_do_scale_inv_format(NVTE_QKV_Format::NVTE_QKV_Format_NOT_SET) \ - .set_batch_size(num_segments) \ - .set_num_attn_heads(attn_heads) \ - .set_num_gqa_groups(num_gqa_groups) \ - .set_head_dim_qk(qk_head_dim) \ - .set_head_dim_v(v_head_dim) \ - .set_max_seqlen_q(q_max_seqlen) \ - .set_max_seqlen_kv(kv_max_seqlen) \ - .set_num_tokens_q(is_ragged ? input_batch *q_max_seqlen : 0) \ - .set_num_tokens_kv(is_ragged ? input_batch *kv_max_seqlen : 0) \ - .set_bias_batch_size(bias_batch) \ - .set_bias_num_heads(bias_heads) \ - .set_bias_seqlen_q(bias_seqlen_q) \ - .set_bias_seqlen_kv(bias_seqlen_kv); - -static void FusedAttnForwardImpl( - cudaStream_t stream, void *q, void *k, void *v, void *bias, void *softmax_offset, void *seed, - void *q_cu_seqlens, void *kv_cu_seqlens, void *q_seq_offsets, void *k_seq_offsets, void *output, - void *softmax_aux, void *max_tensor, void *rng_state, void *workspace, size_t input_batch, - size_t bias_batch, size_t q_max_seqlen, size_t kv_max_seqlen, size_t attn_heads, - size_t num_gqa_groups, size_t bias_heads, size_t qk_head_dim, size_t v_head_dim, - size_t max_segments_per_seq, size_t wkspace_size, float scaling_factor, - float dropout_probability, NVTE_Bias_Type bias_type, NVTE_Mask_Type mask_type, - NVTE_Softmax_Type softmax_type, NVTE_QKV_Layout qkv_layout, DType dtype, DType wkspace_dtype, - bool is_training, bool return_max_logit, bool deterministic, int64_t window_size_left, - int64_t window_size_right, bool bottom_right_diagonal) { - FUSED_ATTN_IMPL_COMMON_BLOCK; - - /* Input tensors */ - auto bias_tensor = TensorWrapper(bias, bias_shape, dtype); - auto softmax_offset_tensor = - TensorWrapper(softmax_offset, std::vector{1, attn_heads, 1, 1}, DType::kFloat32); - - if (is_ragged) { - auto output_size = input_batch * q_max_seqlen * attn_heads * v_head_dim; - cudaMemsetAsync(output, 0, output_size * typeToSize(dtype), stream); - - // Memset to 0xF0 for filling large negative numbers - auto softmax_aux_size = input_batch * q_max_seqlen * attn_heads; - cudaMemsetAsync(softmax_aux, 0xF0, softmax_aux_size * sizeof(float), stream); - if (return_max_logit) { - cudaMemsetAsync(max_tensor, 0xF0, softmax_aux_size * sizeof(float), stream); - } - } - - /* Output tensors */ - auto s_tensor = TensorWrapper(nullptr, std::vector{1}, dtype); // not used in F16 - auto o_shape = is_ragged ? std::vector{input_batch * q_max_seqlen, attn_heads, v_head_dim} - : std::vector{input_batch, q_max_seqlen, attn_heads, v_head_dim}; - auto o_tensor = TensorWrapper(output, o_shape, dtype); - - /* Prepare RNG state */ - auto rng_state_tensor = TensorWrapper(rng_state, std::vector{2}, DType::kInt64); - - cfg.set_return_max_logit(return_max_logit); - auto [backend, fwd_msg] = GetFusedAttnBackendImpl(cfg); - NVTE_CHECK(backend != NVTE_Fused_Attn_Backend::NVTE_No_Backend, - "Fused attention is not supported for this configuration: ", fwd_msg); - nvte_populate_rng_state_async(rng_state, seed, q_max_seqlen, kv_max_seqlen, backend, stream); - - /* Auxiliary tensors (to be propagated to the backward pass later) */ - NVTETensorPack aux_output_tensors; - nvte_tensor_pack_create(&aux_output_tensors); - PrepareFusedAttnForwardAuxTensors(&aux_output_tensors, input_batch, bias_batch, attn_heads, - bias_heads, q_max_seqlen, kv_max_seqlen, dtype, bias_type, - backend, softmax_aux, return_max_logit ? max_tensor : nullptr, - rng_state, bias, softmax_offset); - - /* Call the underlying NVTE API */ - auto dummy_page_table_tensor = TensorWrapper(nullptr, std::vector{1}, DType::kInt32); - - // Prepare Q, K, V pointers and shapes based on layout - // Python passes dummy tensors for unused slots, so we extract from the actual packed data - void *q_ptr = q; - void *k_ptr = k; - void *v_ptr = v; - auto q_shape = is_ragged - ? std::vector{input_batch * q_max_seqlen, attn_heads, qk_head_dim} - : std::vector{input_batch, q_max_seqlen, attn_heads, qk_head_dim}; - auto k_shape = is_ragged - ? std::vector{input_batch * kv_max_seqlen, num_gqa_groups, qk_head_dim} - : std::vector{input_batch, kv_max_seqlen, num_gqa_groups, qk_head_dim}; - auto v_shape = is_ragged - ? std::vector{input_batch * kv_max_seqlen, num_gqa_groups, v_head_dim} - : std::vector{input_batch, kv_max_seqlen, num_gqa_groups, v_head_dim}; - - if (layout_group == NVTE_QKV_Layout_Group::NVTE_3HD) { - // QKV packed in q: [batch*seqlen, 3, heads, dim] - // Python passes: q=packed_qkv, k=dummy, v=dummy - // Extract K and V pointers from the packed q data - NVTE_CHECK(q_max_seqlen == kv_max_seqlen, "q_max_seqlen must equal kv_max_seqlen"); - NVTE_CHECK(qk_head_dim == v_head_dim, - "For QKV packed layout, qk_head_dim must equal v_head_dim"); - size_t stride = (typeToSize(dtype) * attn_heads * qk_head_dim); - q_ptr = q; - k_ptr = static_cast(static_cast(q) + stride); - v_ptr = static_cast(static_cast(q) + 2 * stride); - // For packed QKV, all have same shape since they're views into the same packed tensor - k_shape = q_shape; - v_shape = q_shape; - } else if (layout_group == NVTE_QKV_Layout_Group::NVTE_HD_2HD) { - // Q separate, KV packed in k: [batch*seqlen, 2, num_gqa_groups, dim] - // Python passes: q=query, k=packed_kv, v=dummy - // Extract V pointer from the packed k data - NVTE_CHECK(qk_head_dim == v_head_dim, - "For KV packed layout, qk_head_dim must equal v_head_dim"); - size_t stride = (typeToSize(dtype) * num_gqa_groups * qk_head_dim); - q_ptr = q; - k_ptr = k; - v_ptr = static_cast(static_cast(k) + stride); - // V has same shape as K since they're packed together - v_shape = k_shape; - } - // else NVTE_HD_HD_HD: pointers and shapes already correct - - auto q_tensor = TensorWrapper(q_ptr, q_shape, dtype); - auto k_tensor = TensorWrapper(k_ptr, k_shape, dtype); - auto v_tensor = TensorWrapper(v_ptr, v_shape, dtype); - - FusedAttnFwdParamsWrapper params; - params.set_Q(q_tensor.data()) - .set_K(k_tensor.data()) - .set_V(v_tensor.data()) - .set_Bias(bias_tensor.data()) - .set_SoftmaxOffset(softmax_offset_tensor.data()) - .set_S(s_tensor.data()) - .set_O(o_tensor.data()) - .set_Aux_CTX_Tensors(&aux_output_tensors) - .set_cu_seqlens_q(q_cu_seqlens_tensor.data()) - .set_cu_seqlens_kv(kv_cu_seqlens_tensor.data()) - .set_cu_seqlens_q_padded(q_seq_offsets_tensor.data()) - .set_cu_seqlens_kv_padded(k_seq_offsets_tensor.data()) - .set_page_table_k(dummy_page_table_tensor.data()) - .set_page_table_v(dummy_page_table_tensor.data()) - .set_rng_state(rng_state_tensor.data()) - .set_is_training(is_training) - .set_cuda_graph(false) - .set_return_max_logit(return_max_logit) - .set_attn_mask_type(mask_type) - .set_bias_type(bias_type) - .set_window_size_left(window_size_left) - .set_window_size_right(window_size_right) - .set_bottom_right_diagonal(bottom_right_diagonal) - .set_softmax_type(softmax_type) - .set_dropout(dropout_probability) - .set_attn_scale(scaling_factor) - .set_qkv_layout(qkv_layout) - .set_o_format(nvte_get_q_format(qkv_layout)) - .set_qkv_scale_inv_format(NVTE_QKV_Format_NOT_SET) - .set_max_seqlen_q(q_max_seqlen) - .set_max_seqlen_kv(kv_max_seqlen) - .set_workspace(workspace_tensor.data()) - .set_stream(stream); - nvte_fused_attn_fwd_v2(params); +namespace ffi = xla::ffi; - nvte_tensor_pack_destroy(&aux_output_tensors); +void CheckCuda(cudaError_t status) { + if (status != cudaSuccess) throw std::runtime_error(cudaGetErrorString(status)); } -#define FUSED_ATTN_FFI_GET_ATTRS \ - size_t input_batch = get_attr_value(attrs, "input_batch"); \ - size_t bias_batch = get_attr_value(attrs, "bias_batch"); \ - size_t q_max_seqlen = get_attr_value(attrs, "q_max_seqlen"); \ - size_t kv_max_seqlen = get_attr_value(attrs, "kv_max_seqlen"); \ - size_t attn_heads = get_attr_value(attrs, "attn_heads"); \ - size_t num_gqa_groups = get_attr_value(attrs, "num_gqa_groups"); \ - size_t bias_heads = get_attr_value(attrs, "bias_heads"); \ - size_t qk_head_dim = get_attr_value(attrs, "qk_head_dim"); \ - size_t v_head_dim = get_attr_value(attrs, "v_head_dim"); \ - size_t max_segments_per_seq = get_attr_value(attrs, "max_segments_per_seq"); \ - auto window_size_left = get_attr_value(attrs, "window_size_left"); \ - auto window_size_right = get_attr_value(attrs, "window_size_right"); \ - bool bottom_right_diagonal = get_attr_value(attrs, "bottom_right_diagonal"); \ - float scaling_factor = get_attr_value(attrs, "scaling_factor"); \ - float dropout_probability = get_attr_value(attrs, "dropout_probability"); \ - NVTE_Bias_Type bias_type = \ - static_cast(get_attr_value(attrs, "bias_type")); \ - NVTE_Mask_Type mask_type = \ - static_cast(get_attr_value(attrs, "mask_type")); \ - NVTE_Softmax_Type softmax_type = \ - static_cast(get_attr_value_or_default( \ - attrs, "softmax_type", static_cast(NVTE_Softmax_Type::NVTE_VANILLA_SOFTMAX))); \ - NVTE_QKV_Layout qkv_layout = \ - static_cast(get_attr_value(attrs, "qkv_layout")); \ - bool is_training = get_attr_value(attrs, "is_training"); \ - bool deterministic = get_attr_value(attrs, "deterministic"); \ - auto is_ragged = nvte_get_qkv_format(qkv_layout) == NVTE_QKV_Format::NVTE_THD; \ - size_t wkspace_size = product(workspace_buf->dimensions()); \ - DType dtype = convert_ffi_datatype_to_te_dtype(q_buf.element_type()); \ - DType wkspace_dtype = convert_ffi_datatype_to_te_dtype(workspace_buf->element_type()); - -Error_Type FusedAttnForwardFFI(cudaStream_t stream, Buffer_Type q_buf, Buffer_Type k_buf, - Buffer_Type v_buf, Buffer_Type bias_buf, - Buffer_Type softmax_offset_buf, Buffer_Type seed_buf, - Buffer_Type q_cu_seqlens_buf, Buffer_Type kv_cu_seqlens_buf, - Buffer_Type q_seq_offsets_buf, Buffer_Type k_seq_offsets_buf, - Variadic_Buffer_Type _unused_args, Result_Type output_buf, - Result_Type softmax_aux_buf, Result_Type max_tensor_buf, - Result_Type rng_state_buf, Result_Type workspace_buf, - Dictionary attrs) { - FUSED_ATTN_FFI_GET_ATTRS; - bool return_max_logit = get_attr_value_or_default(attrs, "return_max_logit", false); - - FusedAttnForwardImpl( - stream, q_buf.untyped_data(), k_buf.untyped_data(), v_buf.untyped_data(), - bias_buf.untyped_data(), softmax_offset_buf.untyped_data(), seed_buf.untyped_data(), - q_cu_seqlens_buf.untyped_data(), kv_cu_seqlens_buf.untyped_data(), - is_ragged ? q_seq_offsets_buf.untyped_data() : nullptr, - is_ragged ? k_seq_offsets_buf.untyped_data() : nullptr, output_buf->untyped_data(), - softmax_aux_buf->untyped_data(), max_tensor_buf->untyped_data(), - rng_state_buf->untyped_data(), workspace_buf->untyped_data(), input_batch, bias_batch, - q_max_seqlen, kv_max_seqlen, attn_heads, num_gqa_groups, bias_heads, qk_head_dim, v_head_dim, - max_segments_per_seq, wkspace_size, scaling_factor, dropout_probability, bias_type, mask_type, - softmax_type, qkv_layout, dtype, wkspace_dtype, is_training, return_max_logit, deterministic, - window_size_left, window_size_right, bottom_right_diagonal); - return ffi_with_cuda_error_check(); +void CheckCudnn(cudnnStatus_t status) { + if (status != CUDNN_STATUS_SUCCESS) throw std::runtime_error(cudnnGetErrorString(status)); } -XLA_FFI_DEFINE_HANDLER_SYMBOL(FusedAttnForwardHandler, FusedAttnForwardFFI, - FFI::Bind() - .Ctx() // stream - .Arg() // q - .Arg() // k - .Arg() // v - .Arg() // bias - .Arg() // softmax_offset - .Arg() // seed_buf - .Arg() // q_cu_seqlens - .Arg() // kv_cu_seqlens - .Arg() // q_seq_offsets - .Arg() // k_seq_offsets - .RemainingArgs() // _cp_aux_args unused - .Ret() // output - .Ret() // softmax_aux - .Ret() // max_tensor - .Ret() // rng_state - .Ret() // workspace - .Attrs(), - FFI_CudaGraph_Traits); - -pybind11::tuple GetFusedAttnBackwardWorkspaceSizes( - size_t input_batch, size_t bias_batch, size_t q_max_seqlen, size_t kv_max_seqlen, - size_t attn_heads, size_t num_gqa_groups, size_t bias_heads, size_t qk_head_dim, - size_t v_head_dim, float scaling_factor, float dropout_probability, NVTE_Bias_Type bias_type, - NVTE_Mask_Type mask_type, NVTE_Softmax_Type softmax_type, NVTE_QKV_Layout qkv_layout, - DType dtype, bool is_training, bool deterministic, size_t max_segments_per_seq, - int64_t window_size_left, int64_t window_size_right, bool bottom_right_diagonal) { - auto is_ragged = nvte_get_qkv_format(qkv_layout) == NVTE_QKV_Format::NVTE_THD; - auto q_shape = is_ragged - ? std::vector{input_batch * q_max_seqlen, attn_heads, qk_head_dim} - : std::vector{input_batch, q_max_seqlen, attn_heads, qk_head_dim}; - auto q_tensor = TensorWrapper(nullptr, q_shape, dtype); - auto dq_tensor = TensorWrapper(nullptr, q_shape, dtype); - auto k_shape = is_ragged - ? std::vector{input_batch * kv_max_seqlen, num_gqa_groups, qk_head_dim} - : std::vector{input_batch, kv_max_seqlen, num_gqa_groups, qk_head_dim}; - auto k_tensor = TensorWrapper(nullptr, k_shape, dtype); - auto dk_tensor = TensorWrapper(nullptr, k_shape, dtype); - auto v_shape = is_ragged - ? std::vector{input_batch * kv_max_seqlen, num_gqa_groups, v_head_dim} - : std::vector{input_batch, kv_max_seqlen, num_gqa_groups, v_head_dim}; - auto v_tensor = TensorWrapper(nullptr, v_shape, dtype); - auto dv_tensor = TensorWrapper(nullptr, v_shape, dtype); - - auto output_shape = is_ragged - ? std::vector{input_batch * q_max_seqlen, attn_heads, v_head_dim} - : std::vector{input_batch, q_max_seqlen, attn_heads, v_head_dim}; - auto doutput_tensor = TensorWrapper(nullptr, output_shape, dtype); - auto output_tensor = TensorWrapper(nullptr, output_shape, dtype); - - // F16 doesn't use this tensor - auto s_tensor = TensorWrapper(nullptr, std::vector{1}, dtype); - - auto bias_shape = std::vector{bias_batch, bias_heads, q_max_seqlen, kv_max_seqlen}; - auto dbias_tensor = TensorWrapper(nullptr, bias_shape, dtype); - - NVTETensorPack aux_input_tensors; - nvte_tensor_pack_create(&aux_input_tensors); - - TensorWrapper query_workspace_tensor; - - // It is a WAR to pre-create all possible cuDNN graph at the JIT compile time - size_t max_num_segments = is_ragged ? input_batch * max_segments_per_seq : input_batch; - size_t min_num_segments = is_ragged ? input_batch * max_segments_per_seq : input_batch; - - TensorWrapper dummy_d_softmax_offset_tensor; - if (softmax_type == NVTE_Softmax_Type::NVTE_OFF_BY_ONE_SOFTMAX || - softmax_type == NVTE_Softmax_Type::NVTE_LEARNABLE_SOFTMAX) { - dummy_d_softmax_offset_tensor = - TensorWrapper(nullptr, std::vector{1, attn_heads, 1, 1}, DType::kFloat32); - } - - for (auto num_segments = min_num_segments; num_segments <= max_num_segments; ++num_segments) { - // the last one is the largest which will be the returned workspace size - auto q_cu_seqlens_tensor = - TensorWrapper(nullptr, std::vector{num_segments + 1}, DType::kInt32); - auto kv_cu_seqlens_tensor = - TensorWrapper(nullptr, std::vector{num_segments + 1}, DType::kInt32); - auto dummy_ragged_offset_tensor = - TensorWrapper(nullptr, std::vector{num_segments + 1}, DType::kInt32); - - FusedAttnBwdParamsWrapper params; - params.set_Q(q_tensor.data()) - .set_K(k_tensor.data()) - .set_V(v_tensor.data()) - .set_O(output_tensor.data()) - .set_dO(doutput_tensor.data()) - .set_S(s_tensor.data()) // not used for F16 - .set_dP(s_tensor.data()) // not used for F16 - .set_Aux_CTX_Tensors(&aux_input_tensors) - .set_dQ(dq_tensor.data()) - .set_dK(dk_tensor.data()) - .set_dV(dv_tensor.data()) - .set_dBias(dbias_tensor.data()) - .set_dSoftmaxOffset(dummy_d_softmax_offset_tensor.data()) - .set_cu_seqlens_q(q_cu_seqlens_tensor.data()) - .set_cu_seqlens_kv(kv_cu_seqlens_tensor.data()) - .set_cu_seqlens_q_padded(dummy_ragged_offset_tensor.data()) - .set_cu_seqlens_kv_padded(dummy_ragged_offset_tensor.data()) - .set_deterministic(deterministic) - .set_cuda_graph(false) - .set_attn_mask_type(mask_type) - .set_bias_type(bias_type) - .set_window_size_left(window_size_left) - .set_window_size_right(window_size_right) - .set_bottom_right_diagonal(bottom_right_diagonal) - .set_softmax_type(softmax_type) - .set_dropout(dropout_probability) - .set_attn_scale(scaling_factor) - .set_qkv_layout(qkv_layout) - .set_o_format(nvte_get_q_format(qkv_layout)) - .set_do_format(nvte_get_q_format(qkv_layout)) - .set_dqkv_layout(qkv_layout) - .set_qkv_scale_inv_format(NVTE_QKV_Format_NOT_SET) - .set_do_scale_inv_format(NVTE_QKV_Format_NOT_SET) - .set_max_seqlen_q(q_max_seqlen) - .set_max_seqlen_kv(kv_max_seqlen) - .set_workspace(query_workspace_tensor.data()) - .set_stream(nullptr); - nvte_fused_attn_bwd_v2(params); +template +T Attribute(ffi::Dictionary &attrs, const char *name) { + auto value = attrs.get(name); + if (value.has_error()) { + throw std::runtime_error(std::string("Missing cuDNN graph attribute: ") + name); } - - nvte_tensor_pack_destroy(&aux_input_tensors); - - auto work_shape = MakeShapeVector(query_workspace_tensor.shape()); - return pybind11::make_tuple(work_shape, query_workspace_tensor.dtype()); + return value.value(); } -static void FusedAttnBackwardImpl( - cudaStream_t stream, void *q, void *k, void *v, void *bias, void *softmax_offset, - void *softmax_aux, void *rng_state, void *output, void *doutput, void *q_cu_seqlens, - void *kv_cu_seqlens, void *q_seq_offsets, void *k_seq_offsets, void *dq, void *dk, void *dv, - void *dbias, void *dsoftmax_offset, void *workspace, size_t input_batch, size_t bias_batch, - size_t q_max_seqlen, size_t kv_max_seqlen, size_t attn_heads, size_t num_gqa_groups, - size_t bias_heads, size_t qk_head_dim, size_t v_head_dim, size_t max_segments_per_seq, - size_t wkspace_size, float scaling_factor, float dropout_probability, NVTE_Bias_Type bias_type, - NVTE_Mask_Type mask_type, NVTE_Softmax_Type softmax_type, NVTE_QKV_Layout qkv_layout, - DType dtype, DType wkspace_dtype, bool is_training, bool deterministic, - int64_t window_size_left, int64_t window_size_right, bool bottom_right_diagonal) { - FUSED_ATTN_IMPL_COMMON_BLOCK; - - /* Input tensors */ - auto output_shape = is_ragged - ? std::vector{input_batch * q_max_seqlen, attn_heads, v_head_dim} - : std::vector{input_batch, q_max_seqlen, attn_heads, v_head_dim}; - auto output_tensor = TensorWrapper(output, output_shape, dtype); - auto doutput_tensor = TensorWrapper(doutput, output_shape, dtype); - - /* Output tensors */ - auto s_tensor = TensorWrapper(nullptr, std::vector{1}, dtype); // not used in F16 - auto dbias_tensor = TensorWrapper(dbias, bias_shape, dtype); - - TensorWrapper dsoftmax_offset_tensor; - if (softmax_type == NVTE_Softmax_Type::NVTE_OFF_BY_ONE_SOFTMAX || - softmax_type == NVTE_Softmax_Type::NVTE_LEARNABLE_SOFTMAX) { - dsoftmax_offset_tensor = - TensorWrapper(dsoftmax_offset, std::vector{1, attn_heads, 1, 1}, DType::kFloat32); - } - - /* Auxiliary tensors (propagated from the forward pass) */ - NVTETensorPack aux_input_tensors; - nvte_tensor_pack_create(&aux_input_tensors); - auto [backend, bwd_msg] = GetFusedAttnBackendImpl(cfg); - NVTE_CHECK(backend != NVTE_Fused_Attn_Backend::NVTE_No_Backend, - "Fused attention is not supported for this configuration: ", bwd_msg); - PrepareFusedAttnBackwardAuxTensors(&aux_input_tensors, input_batch, bias_batch, attn_heads, - bias_heads, q_max_seqlen, kv_max_seqlen, dtype, backend, - softmax_aux, rng_state, bias, softmax_offset); - - /* Call the underlying NVTE API */ - // Prepare Q, K, V pointers and shapes based on layout - void *q_ptr = q; - void *k_ptr = k; - void *v_ptr = v; - void *dq_ptr = dq; - void *dk_ptr = dk; - void *dv_ptr = dv; - auto q_shape = is_ragged - ? std::vector{input_batch * q_max_seqlen, attn_heads, qk_head_dim} - : std::vector{input_batch, q_max_seqlen, attn_heads, qk_head_dim}; - auto k_shape = is_ragged - ? std::vector{input_batch * kv_max_seqlen, num_gqa_groups, qk_head_dim} - : std::vector{input_batch, kv_max_seqlen, num_gqa_groups, qk_head_dim}; - auto v_shape = is_ragged - ? std::vector{input_batch * kv_max_seqlen, num_gqa_groups, v_head_dim} - : std::vector{input_batch, kv_max_seqlen, num_gqa_groups, v_head_dim}; - - if (layout_group == NVTE_QKV_Layout_Group::NVTE_3HD) { - // QKV packed in q: [batch*seqlen, 3, heads, dim] - NVTE_CHECK(q_max_seqlen == kv_max_seqlen, "q_max_seqlen must equal kv_max_seqlen"); - NVTE_CHECK(qk_head_dim == v_head_dim, - "For QKV packed layout, qk_head_dim must equal v_head_dim"); - size_t stride = (typeToSize(dtype) * attn_heads * qk_head_dim); - q_ptr = q; - k_ptr = static_cast(static_cast(q) + stride); - v_ptr = static_cast(static_cast(q) + 2 * stride); - dq_ptr = dq; - dk_ptr = static_cast(static_cast(dq) + stride); - dv_ptr = static_cast(static_cast(dq) + 2 * stride); - k_shape = q_shape; - v_shape = q_shape; - } else if (layout_group == NVTE_QKV_Layout_Group::NVTE_HD_2HD) { - // Q separate, KV packed in k: [batch*seqlen, 2, num_gqa_groups, dim] - NVTE_CHECK(qk_head_dim == v_head_dim, - "For KV packed layout, qk_head_dim must equal v_head_dim"); - size_t stride = (typeToSize(dtype) * num_gqa_groups * qk_head_dim); - q_ptr = q; - k_ptr = k; - v_ptr = static_cast(static_cast(k) + stride); - dq_ptr = dq; - dk_ptr = dk; - dv_ptr = static_cast(static_cast(dk) + stride); - // V has same shape as K since they're packed together - v_shape = k_shape; - } - - auto q_tensor = TensorWrapper(q_ptr, q_shape, dtype); - auto k_tensor = TensorWrapper(k_ptr, k_shape, dtype); - auto v_tensor = TensorWrapper(v_ptr, v_shape, dtype); - auto dq_tensor = TensorWrapper(dq_ptr, q_shape, dtype); - auto dk_tensor = TensorWrapper(dk_ptr, k_shape, dtype); - auto dv_tensor = TensorWrapper(dv_ptr, v_shape, dtype); - - if (is_ragged) { - size_t dtype_size = typeToSize(dtype); - if (layout_group == NVTE_QKV_Layout_Group::NVTE_3HD) { - // For packed QKV, dq contains all gradients (dq, dk, dv) - clear all at once - cudaMemsetAsync(dq, 0, 3 * transformer_engine::jax::product(q_shape) * dtype_size, stream); - } else if (layout_group == NVTE_QKV_Layout_Group::NVTE_HD_2HD) { - // Clear dq - cudaMemsetAsync(dq, 0, transformer_engine::jax::product(q_shape) * dtype_size, stream); - // For packed KV, dk contains both dk and dv - clear all at once - cudaMemsetAsync(dk, 0, 2 * transformer_engine::jax::product(k_shape) * dtype_size, stream); - } else { - // All separate - clear each individually - cudaMemsetAsync(dq, 0, transformer_engine::jax::product(q_shape) * dtype_size, stream); - cudaMemsetAsync(dk, 0, transformer_engine::jax::product(k_shape) * dtype_size, stream); - cudaMemsetAsync(dv, 0, transformer_engine::jax::product(v_shape) * dtype_size, stream); - } - } - - FusedAttnBwdParamsWrapper params; - params.set_Q(q_tensor.data()) - .set_K(k_tensor.data()) - .set_V(v_tensor.data()) - .set_O(output_tensor.data()) - .set_dO(doutput_tensor.data()) - .set_S(s_tensor.data()) // not used for F16 - .set_dP(s_tensor.data()) // not used for F16 - .set_Aux_CTX_Tensors(&aux_input_tensors) - .set_dQ(dq_tensor.data()) - .set_dK(dk_tensor.data()) - .set_dV(dv_tensor.data()) - .set_dBias(dbias_tensor.data()) - .set_dSoftmaxOffset(dsoftmax_offset_tensor.data()) - .set_cu_seqlens_q(q_cu_seqlens_tensor.data()) - .set_cu_seqlens_kv(kv_cu_seqlens_tensor.data()) - .set_cu_seqlens_q_padded(q_seq_offsets_tensor.data()) - .set_cu_seqlens_kv_padded(k_seq_offsets_tensor.data()) - .set_deterministic(deterministic) - .set_cuda_graph(false) - .set_attn_mask_type(mask_type) - .set_bias_type(bias_type) - .set_window_size_left(window_size_left) - .set_window_size_right(window_size_right) - .set_bottom_right_diagonal(bottom_right_diagonal) - .set_softmax_type(softmax_type) - .set_dropout(dropout_probability) - .set_attn_scale(scaling_factor) - .set_qkv_layout(qkv_layout) - .set_o_format(nvte_get_q_format(qkv_layout)) - .set_do_format(nvte_get_q_format(qkv_layout)) - .set_dqkv_layout(qkv_layout) - .set_qkv_scale_inv_format(NVTE_QKV_Format_NOT_SET) - .set_do_scale_inv_format(NVTE_QKV_Format_NOT_SET) - .set_max_seqlen_q(q_max_seqlen) - .set_max_seqlen_kv(kv_max_seqlen) - .set_workspace(workspace_tensor.data()) - .set_stream(stream); - nvte_fused_attn_bwd_v2(params); - - nvte_tensor_pack_destroy(&aux_input_tensors); -} - -Error_Type FusedAttnBackwardFFI(cudaStream_t stream, Buffer_Type q_buf, Buffer_Type k_buf, - Buffer_Type v_buf, Buffer_Type bias_buf, - Buffer_Type softmax_offset_buf, Buffer_Type softmax_aux_buf, - Buffer_Type rng_state_buf, Buffer_Type output_buf, - Buffer_Type doutput_buf, Buffer_Type q_cu_seqlens_buf, - Buffer_Type kv_cu_seqlens_buf, Buffer_Type q_seq_offsets_buf, - Buffer_Type k_seq_offsets_buf, Variadic_Buffer_Type _unused_args, - Result_Type dq_buf, Result_Type dk_buf, Result_Type dv_buf, - Result_Type dbias_buf, Result_Type dsoftmax_offset_buf, - Result_Type workspace_buf, Dictionary attrs) { - FUSED_ATTN_FFI_GET_ATTRS; - - FusedAttnBackwardImpl( - stream, q_buf.untyped_data(), k_buf.untyped_data(), v_buf.untyped_data(), - bias_buf.untyped_data(), softmax_offset_buf.untyped_data(), softmax_aux_buf.untyped_data(), - rng_state_buf.untyped_data(), output_buf.untyped_data(), doutput_buf.untyped_data(), - q_cu_seqlens_buf.untyped_data(), kv_cu_seqlens_buf.untyped_data(), - is_ragged ? q_seq_offsets_buf.untyped_data() : nullptr, - is_ragged ? k_seq_offsets_buf.untyped_data() : nullptr, dq_buf->untyped_data(), - dk_buf->untyped_data(), dv_buf->untyped_data(), dbias_buf->untyped_data(), - dsoftmax_offset_buf->untyped_data(), workspace_buf->untyped_data(), input_batch, bias_batch, - q_max_seqlen, kv_max_seqlen, attn_heads, num_gqa_groups, bias_heads, qk_head_dim, v_head_dim, - max_segments_per_seq, wkspace_size, scaling_factor, dropout_probability, bias_type, mask_type, - softmax_type, qkv_layout, dtype, wkspace_dtype, is_training, deterministic, window_size_left, - window_size_right, bottom_right_diagonal); - - return ffi_with_cuda_error_check(); -} - -XLA_FFI_DEFINE_HANDLER_SYMBOL(FusedAttnBackwardHandler, FusedAttnBackwardFFI, - FFI::Bind() - .Ctx() // stream - .Arg() // q - .Arg() // k - .Arg() // v - .Arg() // bias - .Arg() // softmax_offset - .Arg() // softmax_aux - .Arg() // rng_state - .Arg() // output - .Arg() // doutput - .Arg() // q_cu_seqlens - .Arg() // kv_cu_seqlens - .Arg() // q_seq_offsets - .Arg() // k_seq_offsets - .RemainingArgs() // _cp_aux_args unused - .Ret() // dq - .Ret() // dk - .Ret() // dv - .Ret() // dbias - .Ret() // dsoftmax_offset - .Ret() // workspace - .Attrs(), - FFI_CudaGraph_Traits); - -namespace { - -struct ScoreModScalarStorage { - alignas(16) std::array data{}; - size_t size = 0; +struct Handle { + cudnnHandle_t value = nullptr; + Handle() { CheckCudnn(cudnnCreate(&value)); } + ~Handle() { cudnnDestroy(value); } }; -struct ScoreModGraphCacheKey { - int device_id = 0; - int64_t hash0 = 0; - int64_t hash1 = 0; - int64_t frontend_version = 0; - - bool operator==(const ScoreModGraphCacheKey &other) const { - return device_id == other.device_id && hash0 == other.hash0 && hash1 == other.hash1 && - frontend_version == other.frontend_version; - } -}; - -struct ScoreModGraphCacheKeyHash { - size_t operator()(const ScoreModGraphCacheKey &key) const { - size_t seed = std::hash{}(key.device_id); - auto combine = [&seed](int64_t value) { - // 64-bit golden ratio constant from boost::hash_combine to spread mixed keys. - seed ^= std::hash{}(value) + 0x9e3779b97f4a7c15ULL + (seed << 6) + (seed >> 2); - }; - combine(key.hash0); - combine(key.hash1); - combine(key.frontend_version); - return seed; - } -}; - -using ScoreModGraphPtr = std::shared_ptr; - -std::unordered_map & -getScoreModeGraphCache() { - static std::unordered_map - cache; - return cache; -} - -std::mutex &getScoreModGraphCacheMutex() { +cudnnHandle_t GetHandle() { + // Handles are thread-local and device-local; assign the stream on each call. + thread_local std::unordered_map> handles; + int device; + CheckCuda(cudaGetDevice(&device)); + auto &handle = handles[device]; + if (!handle) handle = std::make_unique(); + return handle->value; +} + +using Graph = cudnn_frontend::graph::Graph; + +std::shared_ptr GetGraph(cudnnHandle_t handle, ffi::Dictionary &attrs) { + const auto version = Attribute(attrs, "cudnn_frontend_version"); + if (version != CUDNN_FRONTEND_VERSION) { + throw std::runtime_error( + "cuDNN frontend Python/C++ version mismatch for graph deserialization."); + } + int device; + CheckCuda(cudaGetDevice(&device)); + const auto key = std::to_string(device) + ":" + + std::to_string(Attribute(attrs, "graph_hash0")) + ":" + + std::to_string(Attribute(attrs, "graph_hash1")) + ":" + + std::to_string(version); static std::mutex mutex; - return mutex; -} - -ScoreModGraphCacheKey GetScoreModGraphCacheKey(Dictionary &attrs) { - const int64_t frontend_version = get_attr_value(attrs, "cudnn_frontend_version"); - NVTE_CHECK(frontend_version == CUDNN_FRONTEND_VERSION, - "cuDNN frontend version mismatch for score_mod graph deserialization: graph was " - "serialized with Python cuDNN frontend version ", - frontend_version, - ", but Transformer Engine C++ was built with CUDNN_FRONTEND_VERSION ", - CUDNN_FRONTEND_VERSION, "."); - - int device_id = 0; - NVTE_CHECK_CUDA(cudaGetDevice(&device_id)); - return ScoreModGraphCacheKey{ - device_id, - get_attr_value(attrs, "graph_hash0"), - get_attr_value(attrs, "graph_hash1"), - frontend_version, - }; -} - -ScoreModGraphPtr GetScoreModGraph(cudaStream_t stream, Dictionary &attrs) { - const auto key = GetScoreModGraphCacheKey(attrs); - { - std::lock_guard lock(getScoreModGraphCacheMutex()); - auto &cache = getScoreModeGraphCache(); - auto it = cache.find(key); - if (it != cache.end()) { - return it->second; - } + static std::unordered_map> cache; + std::lock_guard lock(mutex); + auto &graph = cache[key]; + if (!graph) { + const auto data = Attribute(attrs, "serialized_graph"); + auto candidate = std::make_shared(); + auto status = candidate->deserialize(handle, std::vector(data.begin(), data.end())); + if (!status.is_good()) throw std::runtime_error(status.get_message()); + graph = std::move(candidate); } - - const auto serialized_graph = get_attr_value(attrs, "serialized_graph"); - std::vector serialized_data(serialized_graph.begin(), serialized_graph.end()); - - auto handle = nvte_get_cudnn_handle(); - NVTE_CHECK_CUDNN(cudnnSetStream(handle, stream)); - - auto graph = std::make_shared(); - auto status = graph->deserialize(handle, serialized_data); - NVTE_CHECK(status.is_good(), - "Failed to deserialize cuDNN score_mod SDPA graph: ", status.get_message()); - - std::lock_guard lock(getScoreModGraphCacheMutex()); - auto &cache = getScoreModeGraphCache(); - auto it = cache.find(key); - if (it != cache.end()) { - return it->second; - } - cache.emplace(key, graph); return graph; } -Error_Type ExecuteScoreModGraph(cudaStream_t stream, Dictionary &attrs, - const std::vector &input_ptrs, - const std::vector &output_ptrs, void *workspace) { - auto graph = GetScoreModGraph(stream, attrs); - auto input_uids = get_attr_value>(attrs, "input_uids"); - auto output_uids = get_attr_value>(attrs, "output_uids"); - auto scalar_uids = get_attr_value>(attrs, "scalar_uids"); - auto scalar_sizes = get_attr_value>(attrs, "scalar_sizes"); - auto scalar_values = get_attr_value>(attrs, "scalar_values"); - - NVTE_CHECK(input_ptrs.size() == input_uids.size(), "cuDNN score_mod graph expected ", - input_uids.size(), " inputs but got ", input_ptrs.size()); - NVTE_CHECK(output_ptrs.size() >= output_uids.size(), "cuDNN score_mod graph expected at least ", - output_uids.size(), " outputs but got ", output_ptrs.size()); - NVTE_CHECK(scalar_uids.size() == scalar_sizes.size(), - "Mismatched score_mod scalar uid/value-size counts."); - NVTE_CHECK(scalar_values.size() == scalar_uids.size() * 16, - "Mismatched score_mod packed scalar value size."); - - std::unordered_map variant_pack; - for (size_t i = 0; i < input_uids.size(); ++i) { - variant_pack.emplace(input_uids[i], input_ptrs[i]); - } - for (size_t i = 0; i < output_uids.size(); ++i) { - variant_pack.emplace(output_uids[i], output_ptrs[i]); - } - - std::vector scalar_storage(scalar_uids.size()); - for (size_t i = 0; i < scalar_uids.size(); ++i) { - NVTE_CHECK(scalar_sizes[i] >= 0 && scalar_sizes[i] <= 16, - "score_mod pass-by-value scalars must be at most 16 bytes."); - scalar_storage[i].size = static_cast(scalar_sizes[i]); - std::copy_n(scalar_values.begin() + i * 16, 16, scalar_storage[i].data.begin()); - variant_pack.emplace(scalar_uids[i], scalar_storage[i].data.data()); - } - - auto handle = nvte_get_cudnn_handle(); - NVTE_CHECK_CUDNN(cudnnSetStream(handle, stream)); - auto status = graph->execute(handle, variant_pack, workspace); - NVTE_CHECK(status.is_good(), - "cuDNN score_mod SDPA graph execution failed: ", status.get_message()); - return ffi_with_cuda_error_check(); -} +struct Scalar { + alignas(16) std::array value{}; +}; -void AppendRemainingBuffers(Variadic_Buffer_Type args, std::vector *ptrs) { - ptrs->reserve(ptrs->size() + args.size()); - for (size_t i = 0; i < args.size(); ++i) { - auto maybe_buf = args.get(i); - NVTE_CHECK(!maybe_buf.has_error(), "Failed to decode variadic score_mod input buffer."); - ptrs->push_back(maybe_buf.value().untyped_data()); +ffi::Error InitializeGraph(cudaStream_t stream, ffi::RemainingArgs, ffi::RemainingRets, + ffi::Dictionary attrs) { + try { + auto handle = GetHandle(); + CheckCudnn(cudnnSetStream(handle, stream)); + GetGraph(handle, attrs); + return ffi::Error::Success(); + } catch (const std::exception &error) { + return ffi::Error::Internal(std::string("JAX cuDNN attention initialization: ") + error.what()); + } +} + +ffi::Error ExecuteGraph(cudaStream_t stream, ffi::RemainingArgs args, ffi::RemainingRets rets, + ffi::Dictionary attrs) { + try { + auto handle = GetHandle(); + CheckCudnn(cudnnSetStream(handle, stream)); + auto graph = GetGraph(handle, attrs); + auto input_uids = Attribute>(attrs, "input_uids"); + auto output_uids = Attribute>(attrs, "output_uids"); + auto input_indices = attrs.get>("input_indices"); + auto input_offsets = attrs.get>("input_offsets"); + auto output_indices = attrs.get>("output_indices"); + auto output_offsets = attrs.get>("output_offsets"); + auto output_fill_values = attrs.get>("output_fill_values"); + if (rets.size() == 0 || + (output_fill_values.has_value() && output_fill_values.value().size() != rets.size() - 1)) { + throw std::runtime_error("Invalid cuDNN output initialization bindings."); + } + std::unordered_map pointers; + for (size_t i = 0; i < input_uids.size(); ++i) { + const auto index = input_indices.has_value() ? input_indices.value()[i] : i; + auto buf = args.get(index); + if (buf.has_error()) throw std::runtime_error("Invalid cuDNN input buffer binding."); + auto ptr = static_cast(buf.value().untyped_data()); + pointers.emplace(input_uids[i], + ptr + (input_offsets.has_value() ? input_offsets.value()[i] : 0)); + } + for (size_t i = 0; i < rets.size() - 1; ++i) { + auto buf = rets.get(i); + if (buf.has_error()) throw std::runtime_error("Invalid cuDNN output buffer binding."); + // Initialize ragged slots and optional gradient outputs. + if (buf.value()->size_bytes()) { + const int fill = output_fill_values.has_value() ? output_fill_values.value()[i] : 0; + CheckCuda( + cudaMemsetAsync(buf.value()->untyped_data(), fill, buf.value()->size_bytes(), stream)); + } + } + for (size_t i = 0; i < output_uids.size(); ++i) { + if (output_uids[i] < 0) continue; + const auto index = output_indices.has_value() ? output_indices.value()[i] : i; + auto buf = rets.get(index); + if (buf.has_error()) throw std::runtime_error("Invalid cuDNN output buffer binding."); + auto ptr = static_cast(buf.value()->untyped_data()); + pointers.emplace(output_uids[i], + ptr + (output_offsets.has_value() ? output_offsets.value()[i] : 0)); + } + auto scalar_uids = Attribute>(attrs, "scalar_uids"); + auto scalar_sizes = Attribute>(attrs, "scalar_sizes"); + auto scalar_values = Attribute>(attrs, "scalar_values"); + if (scalar_sizes.size() != scalar_uids.size() || + scalar_values.size() != scalar_uids.size() * 16) { + throw std::runtime_error("Invalid cuDNN pass-by-value scalar bindings."); + } + std::vector scalars(scalar_uids.size()); + for (size_t i = 0; i < scalars.size(); ++i) { + if (scalar_sizes[i] < 0 || scalar_sizes[i] > 16) { + throw std::runtime_error("Invalid cuDNN pass-by-value scalar size."); + } + std::copy_n(scalar_values.begin() + i * 16, 16, scalars[i].value.begin()); + pointers.emplace(scalar_uids[i], scalars[i].value.data()); + } + auto workspace = rets.get(rets.size() - 1); + if (workspace.has_error()) throw std::runtime_error("Missing cuDNN graph workspace."); + auto status = graph->execute(handle, pointers, workspace.value()->untyped_data()); + if (!status.is_good()) throw std::runtime_error(status.get_message()); + CheckCuda(cudaGetLastError()); + return ffi::Error::Success(); + } catch (const std::exception &error) { + return ffi::Error::Internal(std::string("JAX cuDNN attention: ") + error.what()); } } } // namespace -Error_Type FusedAttnScoreModForwardFFI(cudaStream_t stream, Buffer_Type q_buf, Buffer_Type k_buf, - Buffer_Type v_buf, Variadic_Buffer_Type score_mod_args, - Result_Type output_buf, Result_Type stats_buf, - Result_Type workspace_buf, Dictionary attrs) { - std::vector input_ptrs = {q_buf.untyped_data(), k_buf.untyped_data(), - v_buf.untyped_data()}; - AppendRemainingBuffers(score_mod_args, &input_ptrs); - - std::vector output_ptrs = {output_buf->untyped_data(), stats_buf->untyped_data()}; - return ExecuteScoreModGraph(stream, attrs, input_ptrs, output_ptrs, - workspace_buf->untyped_data()); -} - -XLA_FFI_DEFINE_HANDLER_SYMBOL(FusedAttnScoreModForwardHandler, FusedAttnScoreModForwardFFI, - FFI::Bind() - .Ctx() // stream - .Arg() // q - .Arg() // k - .Arg() // v - .RemainingArgs() // score_mod tensor operands - .Ret() // output - .Ret() // stats - .Ret() // workspace - .Attrs()); - -Error_Type FusedAttnScoreModBackwardFFI(cudaStream_t stream, Buffer_Type q_buf, Buffer_Type k_buf, - Buffer_Type v_buf, Buffer_Type output_buf, - Buffer_Type doutput_buf, Buffer_Type stats_buf, - Variadic_Buffer_Type score_mod_args, Result_Type dq_buf, - Result_Type dk_buf, Result_Type dv_buf, - Result_Type workspace_buf, Dictionary attrs) { - std::vector input_ptrs = {q_buf.untyped_data(), k_buf.untyped_data(), - v_buf.untyped_data(), output_buf.untyped_data(), - doutput_buf.untyped_data(), stats_buf.untyped_data()}; - AppendRemainingBuffers(score_mod_args, &input_ptrs); - - std::vector output_ptrs = {dq_buf->untyped_data(), dk_buf->untyped_data(), - dv_buf->untyped_data()}; - return ExecuteScoreModGraph(stream, attrs, input_ptrs, output_ptrs, - workspace_buf->untyped_data()); -} +XLA_FFI_DEFINE_HANDLER_SYMBOL(CudnnGraphHandler, ExecuteGraph, + xla::ffi::Ffi::Bind() + .Ctx>() + .RemainingArgs() + .RemainingRets() + .Attrs(), + {xla::ffi::Traits::kCmdBufferCompatible}); -XLA_FFI_DEFINE_HANDLER_SYMBOL(FusedAttnScoreModBackwardHandler, FusedAttnScoreModBackwardFFI, - FFI::Bind() - .Ctx() // stream - .Arg() // q - .Arg() // k - .Arg() // v - .Arg() // output - .Arg() // doutput - .Arg() // stats - .RemainingArgs() // score_mod tensor operands - .Ret() // dq - .Ret() // dk - .Ret() // dv - .Ret() // workspace +XLA_FFI_DEFINE_HANDLER_SYMBOL(CudnnGraphInitializeHandler, InitializeGraph, + xla::ffi::Ffi::Bind() + .Ctx>() + .RemainingArgs() + .RemainingRets() .Attrs()); } // namespace jax diff --git a/transformer_engine/jax/csrc/extensions/pybind.cpp b/transformer_engine/jax/csrc/extensions/pybind.cpp index 69f304ec999..308b80e29de 100644 --- a/transformer_engine/jax/csrc/extensions/pybind.cpp +++ b/transformer_engine/jax/csrc/extensions/pybind.cpp @@ -69,19 +69,10 @@ pybind11::dict Registrations() { pybind11::arg("initialize") = EncapsulateFFI(NormBackwardInitializeHandler), pybind11::arg("execute") = EncapsulateFFI(NormBackwardHandler)); - // Attention - dict["te_fused_attn_forward_ffi"] = - pybind11::dict(pybind11::arg("prepare") = EncapsulateFFI(CudnnHandleInitHandler), - pybind11::arg("execute") = EncapsulateFFI(FusedAttnForwardHandler)); - dict["te_fused_attn_backward_ffi"] = - pybind11::dict(pybind11::arg("prepare") = EncapsulateFFI(CudnnHandleInitHandler), - pybind11::arg("execute") = EncapsulateFFI(FusedAttnBackwardHandler)); - dict["te_fused_attn_score_mod_forward_ffi"] = - pybind11::dict(pybind11::arg("prepare") = EncapsulateFFI(CudnnHandleInitHandler), - pybind11::arg("execute") = EncapsulateFFI(FusedAttnScoreModForwardHandler)); - dict["te_fused_attn_score_mod_backward_ffi"] = - pybind11::dict(pybind11::arg("prepare") = EncapsulateFFI(CudnnHandleInitHandler), - pybind11::arg("execute") = EncapsulateFFI(FusedAttnScoreModBackwardHandler)); + // Python-built cuDNN graphs use a JAX-owned executor and handle cache. + dict["te_cudnn_graph_ffi"] = + pybind11::dict(pybind11::arg("initialize") = EncapsulateFFI(CudnnGraphInitializeHandler), + pybind11::arg("execute") = EncapsulateFFI(CudnnGraphHandler)); // GEMM dict["te_gemm_ffi"] = @@ -157,8 +148,6 @@ pybind11::dict Registrations() { PYBIND11_MODULE(transformer_engine_jax, m) { m.def("registrations", &Registrations); - m.def("get_fused_attn_backend", &GetFusedAttnBackend, "Get Fused Attention backend", - pybind11::arg("fused_attn_params")); m.def("get_cuda_version", &GetCudaRuntimeVersion); m.def("get_cudnn_version", &GetCudnnRuntimeVersion); m.def("get_cudnn_frontend_version", &GetCudnnFrontendVersion); @@ -169,10 +158,7 @@ PYBIND11_MODULE(transformer_engine_jax, m) { m.def("get_dbias_quantize_workspace_sizes", &GetDBiasQuantizeWorkspaceSizes); m.def("get_norm_fwd_workspace_sizes", &GetNormForwardWorkspaceSizes); m.def("get_norm_bwd_workspace_sizes", &GetNormBackwardWorkspaceSizes); - m.def("get_fused_attn_fwd_workspace_sizes", &GetFusedAttnForwardWorkspaceSizes); - m.def("get_fused_attn_bwd_workspace_sizes", &GetFusedAttnBackwardWorkspaceSizes); m.def("get_topk_workspace_sizes", &GetTopkWorkspaceSizes); - m.def("nvte_get_qkv_format", &nvte_get_qkv_format); m.def("is_non_nt_fp8_gemm_supported", &nvte_is_non_tn_fp8_gemm_supported); m.def("nvte_built_with_cublasmp", &::nvte_built_with_cublasmp); m.def("initialize_cgemm_communicator", &InitializeCgemmCommunicator); @@ -208,7 +194,8 @@ PYBIND11_MODULE(transformer_engine_jax, m) { pybind11::enum_(m, "NVTE_Bias_Type", pybind11::module_local()) .value("NVTE_NO_BIAS", NVTE_Bias_Type::NVTE_NO_BIAS) .value("NVTE_PRE_SCALE_BIAS", NVTE_Bias_Type::NVTE_PRE_SCALE_BIAS) - .value("NVTE_POST_SCALE_BIAS", NVTE_Bias_Type::NVTE_POST_SCALE_BIAS); + .value("NVTE_POST_SCALE_BIAS", NVTE_Bias_Type::NVTE_POST_SCALE_BIAS) + .value("NVTE_ALIBI", NVTE_Bias_Type::NVTE_ALIBI); pybind11::enum_(m, "NVTE_Mask_Type", pybind11::module_local()) .value("NVTE_NO_MASK", NVTE_Mask_Type::NVTE_NO_MASK) diff --git a/transformer_engine/jax/flax/transformer.py b/transformer_engine/jax/flax/transformer.py index 5a01572f1fa..453ceab49ed 100644 --- a/transformer_engine/jax/flax/transformer.py +++ b/transformer_engine/jax/flax/transformer.py @@ -310,6 +310,7 @@ class _FusedDotProductAttention(nn.Module): # pylint: disable=too-few-public-me softmax_type: AttnSoftmaxType = AttnSoftmaxType.VANILLA_SOFTMAX score_mod: Optional[Callable] = None score_mod_bprop: Optional[Callable] = None + fp8: bool = False score_mod_requested: bool = False return_max_logit: bool = False @@ -369,6 +370,8 @@ def __call__( "score_mod_bprop_tensors": score_mod_bprop_tensors, "return_max_logit": self.return_max_logit, } + if self.fp8: + fused_attn_kwargs["fp8"] = True if self.qkv_layout.is_qkvpacked(): """qkvpacked format, treat @@ -632,6 +635,10 @@ class DotProductAttention(nn.Module): # pylint: disable=too-few-public-methods return_max_logit: bool, default = False If True, return ``(output, max_logit)`` where ``max_logit`` contains the per-head maximum attention logits with shape ``[h]``. This path requires fused attention. + fp8: bool, default = False + Use cuDNN FP8 attention with current tensor scaling. Inputs, outputs and input + gradients remain FP16/BF16. Fused attention is required, and this option is + independent of the GEMM quantization recipe. Optimization parameters ----------------------- @@ -661,6 +668,7 @@ class DotProductAttention(nn.Module): # pylint: disable=too-few-public-methods score_mod: Optional[Callable] = None score_mod_bprop: Optional[Callable] = None return_max_logit: bool = False + fp8: bool = False def __post_init__(self): # TODO(KshitijLakhani): Remove warning in TransformerEngine v2.12 @@ -761,7 +769,7 @@ def __call__( or score_mod_bprop_tensors is not None ) - if attn_bias_type == AttnBiasType.NO_BIAS: + if attn_bias_type in (AttnBiasType.NO_BIAS, AttnBiasType.ALIBI): assert ( bias is None ), f"bias must be None when attn_bias_type is NO_BIAS, but got bias={bias}" @@ -791,6 +799,8 @@ def __call__( # Use fused attn (if kernel check below passes) by default enable_fused_attn = int(os.getenv("NVTE_FUSED_ATTN", "1")) + if self.fp8 and not enable_fused_attn: + raise ValueError("FP8 attention requires fused attention, but NVTE_FUSED_ATTN=0.") if self.return_max_logit and not enable_fused_attn: raise ValueError("return_max_logit requires fused attention, but NVTE_FUSED_ATTN=0.") @@ -832,7 +842,7 @@ def __call__( softmax_type=softmax_type, dropout_probability=self.attention_dropout, q_num_heads=self.num_attention_heads, - kv_num_heads=self.num_gqa_groups, + kv_num_heads=self.num_gqa_groups or self.num_attention_heads, q_max_seqlen=seqlen_q, kv_max_seqlen=seqlen_kv, head_dim_qk=head_dim_qk, @@ -845,9 +855,16 @@ def __call__( bias_seqlen_q=bias_seqlen_q, bias_seqlen_kv=bias_seqlen_kv, max_segments_per_seq=self.max_segments_per_seq, + fp8=self.fp8, ) fused_attn_backend, _ = fused_attn_helper.get_fused_attn_backend() has_fused_attn_kernel = fused_attn_backend != NVTE_Fused_Attn_Backend.NVTE_No_Backend + if attn_bias_type == AttnBiasType.ALIBI and ( + not enable_fused_attn or not has_fused_attn_kernel + ): + raise ValueError("ALiBi requires an available cuDNN fused attention kernel.") + if self.fp8 and not has_fused_attn_kernel: + raise ValueError("No cuDNN FP8 attention kernel supports this configuration.") if score_mod_requested and not has_fused_attn_kernel: raise ValueError( "score_mod requires fused attention, but no fused attention kernel is available." @@ -949,6 +966,7 @@ def __call__( softmax_type=softmax_type, score_mod=self.score_mod, score_mod_bprop=self.score_mod_bprop, + fp8=self.fp8, score_mod_requested=score_mod_requested, return_max_logit=self.return_max_logit, )( @@ -1272,6 +1290,9 @@ class MultiHeadAttention(nn.Module): # pylint: disable=too-few-public-methods argument to keep tensor operands as normal JAX inputs. score_mod_bprop_tensors: Optional[Mapping[str, Any]], default = None Additional tensors or pass-by-value scalars for ``score_mod_bprop``. + fp8: bool, default = False + Use cuDNN FP8 dot product attention with current tensor scaling. Inputs and + outputs retain their FP16/BF16 dtype; the GEMM quantization recipe is independent. """ head_dim: int @@ -1314,6 +1335,8 @@ class MultiHeadAttention(nn.Module): # pylint: disable=too-few-public-methods apply_residual_connection_post_layernorm: Optional[bool] = None fuse_qkv: Optional[bool] = None + fp8: bool = False + def __post_init__(self): # Deal with changed defaults in API # TODO(KshitijLakhani): Remove warning in TransformerEngine v2.12 @@ -1735,6 +1758,7 @@ def generate_batch_seqlen_logical_axes(is_sharded_seq): softmax_type=self.softmax_type, score_mod=self.score_mod, score_mod_bprop=self.score_mod_bprop, + fp8=self.fp8, )( *dpa_args, mask, @@ -2061,6 +2085,10 @@ class TransformerLayer(nn.Module): # pylint: disable=too-few-public-methods argument to keep tensor operands as normal JAX inputs. score_mod_bprop_tensors: Optional[Mapping[str, Any]], default = None Additional tensors or pass-by-value scalars for ``score_mod_bprop``. + fp8: bool, default = False + Use cuDNN FP8 attention in the self-attention and encoder-decoder blocks. + Attention uses current tensor scaling and retains FP16/BF16 inputs and outputs, + independently of the GEMM quantization recipe. Optimization parameters ----------------------- @@ -2129,6 +2157,7 @@ class TransformerLayer(nn.Module): # pylint: disable=too-few-public-methods softmax_type: str = "vanilla" score_mod: Optional[Callable] = None score_mod_bprop: Optional[Callable] = None + fp8: bool = False def __post_init__(self): if self.mha_kernel_init is None: @@ -2300,6 +2329,7 @@ def generate_batch_seqlen_logical_axes(is_shared_seq=None): softmax_type=self.softmax_type, score_mod=self.score_mod, score_mod_bprop=self.score_mod_bprop, + fp8=self.fp8, )( inputs, inputs, @@ -2395,6 +2425,7 @@ def hidden_dropout(x, deterministic): name="encoder_decoder_attention", window_size=self.window_size, softmax_type=self.softmax_type, + fp8=self.fp8, )(x, encoded, encoder_decoder_mask, deterministic=deterministic) y = with_sharding_constraint_by_logical_axes( From 4b40acd67c8a924d70715df444edb90df9b4ebf5 Mon Sep 17 00:00:00 2001 From: Vladimir Cherepanov Date: Sun, 11 Oct 2026 21:10:45 +0000 Subject: [PATCH 2/6] Port PyTorch cuDNN attention to the Python frontend Move cuDNN attention support queries, graph construction, plan caching, and forward/backward execution into the PyTorch Python frontend. Cover FP16/BF16, FP8 and MXFP8 while preserving layout, masking, softmax, ragged and paged KV, quantization, and context-parallel behavior. Share graph lifecycle helpers with experimental flex attention. Remove the obsolete PyTorch bindings to common attention execution and retain a capture-safe C++ Philox helper for PyTorch generator state. Keep the common and JAX attention implementations intact. Add numerical-reference, execution-path isolation, capability-exclusion, and CUDA generator replay regressions, and document the Python frontend dependency and graph caching behavior. Validation on GB200 with cuDNN 9.27 and frontend 1.30: - Editable PyTorch build succeeded. - Requested attention/flex suites: 2,869 passed, 2,425 skipped. - Final focused numerical coverage: 629 passed, 560 skipped. - Deterministic attention: 2,231 passed, 3,038 skipped with NVTE_ALLOW_NONDETERMINISTIC_ALGO=0 and NVTE_FLASH_ATTN_V4=0. - Compiled attention: 52 passed; CUDA graph and two-/four-GPU context-parallel coverage passed. - Default/custom generator replay passed; FP8 dropout replay was skipped because cuDNN does not support the tested Blackwell configuration. - Changed-code lint, formatting, and whitespace checks passed. Whole-repository checks still report an existing Python import-order violation in fused_mla_q_uproj.py and missing license headers in the untracked nccl_ep directory. Signed-off-by: Vladimir Cherepanov --- docs/installation.rst | 5 + tests/pytorch/attention/test_attention.py | 302 ++++- .../dot_product_attention/backends.py | 7 +- .../dot_product_attention/flex_attention.py | 96 +- .../attention/dot_product_attention/utils.py | 6 +- .../pytorch/cpp_extensions/cudnn.py | 76 ++ .../cpp_extensions/cudnn_fused_attn.py | 1041 +++++++++++++++++ .../pytorch/cpp_extensions/fused_attn.py | 157 +-- transformer_engine/pytorch/csrc/extensions.h | 31 +- .../pytorch/csrc/extensions/attention.cpp | 640 +--------- .../pytorch/csrc/extensions/pybind.cpp | 9 +- 11 files changed, 1524 insertions(+), 846 deletions(-) create mode 100644 transformer_engine/pytorch/cpp_extensions/cudnn.py create mode 100644 transformer_engine/pytorch/cpp_extensions/cudnn_fused_attn.py diff --git a/docs/installation.rst b/docs/installation.rst index 9d3eada827c..cb9ddccc16f 100644 --- a/docs/installation.rst +++ b/docs/installation.rst @@ -16,6 +16,11 @@ Prerequisites 3. |driver link|_ supporting CUDA 12.1 or later. 4. `cuDNN 9.12 `__ or later. +PyTorch's cuDNN attention backend uses the cuDNN frontend Python API. The PyTorch +installation includes ``nvidia-cudnn-frontend>=1.29.0``. Attention graphs and plans +are cached by device and tensor metadata; dropout uses PyTorch's CUDA generator, +including CUDA graph capture and replay. + If the CUDA Toolkit headers are not available at runtime in a standard installation path, e.g. within `CUDA_HOME`, set `NVTE_CUDA_INCLUDE_PATH` in the environment. diff --git a/tests/pytorch/attention/test_attention.py b/tests/pytorch/attention/test_attention.py index 6bcff5ce3f1..47350b7a6bb 100644 --- a/tests/pytorch/attention/test_attention.py +++ b/tests/pytorch/attention/test_attention.py @@ -46,11 +46,13 @@ import transformer_engine.pytorch.cpp_extensions as ext from transformer_engine.pytorch.cpp_extensions.fused_attn import ( FusedAttnBackend, + AttnMaskType, QKVFormat, QKVLayout, fused_attn_bwd, fused_attn_fwd, ) +from transformer_engine.pytorch.cpp_extensions.cudnn_fused_attn import get_fused_attn_backend from transformer_engine.pytorch.distributed import CudaRNGStatesTracker from transformer_engine.pytorch.module.base import TransformerEngineBaseModule from transformer_engine.pytorch.utils import ( @@ -58,7 +60,7 @@ scaled_init_method_normal, ) from transformer_engine.pytorch.utils import get_cudnn_version -from transformer_engine.pytorch.constants import FP8BwdTensorIdx, FP8FwdTensorIdx +from transformer_engine.pytorch.constants import DType, FP8BwdTensorIdx, FP8FwdTensorIdx import transformer_engine_torch as tex from transformer_engine.pytorch.quantized_tensor import ( Quantizer, @@ -158,23 +160,21 @@ def test_fused_attn_backend_message(): ) # One of TE's rules is violated and the error message is surfaced - backend, message = tex.get_fused_attn_backend( + backend, message = get_fused_attn_backend( replace(baseline, bias_type=tex.NVTE_Bias_Type.NVTE_PRE_SCALE_BIAS) ) assert backend == tex.NVTE_Fused_Attn_Backend.NVTE_No_Backend assert message == "Fused attention does not support pre-scale bias." # No error message if supported; otherwise skip the test - backend, message = tex.get_fused_attn_backend(baseline) + backend, message = get_fused_attn_backend(baseline) if backend == tex.NVTE_Fused_Attn_Backend.NVTE_No_Backend: pytest.skip(f"FusedAttention does not support the baseline config: {message}") assert message == "" # All TE rules have cleared; now gets rejected by cuDNN's support check # cuDNN's error message might change across cuDNN versions, so only verify the presence of the string - backend, message = tex.get_fused_attn_backend( - replace(baseline, head_dim_qk=1024, head_dim_v=1024) - ) + backend, message = get_fused_attn_backend(replace(baseline, head_dim_qk=1024, head_dim_v=1024)) assert backend == tex.NVTE_Fused_Attn_Backend.NVTE_No_Backend assert message != "" @@ -201,20 +201,304 @@ def test_fused_attn_backend_message(): ) # No error message if supported; otherwise skip the test - backend, message = tex.get_fused_attn_backend(fp8_baseline) + backend, message = get_fused_attn_backend(fp8_baseline) if backend == tex.NVTE_Fused_Attn_Backend.NVTE_No_Backend: pytest.skip(f"FusedAttention does not support the baseline FP8 config: {message}") assert backend == tex.NVTE_Fused_Attn_Backend.NVTE_FP8 assert message == "" # All TE FP8 rules have cleared; now gets rejected by cuDNN's FP8 support check - backend, message = tex.get_fused_attn_backend( + backend, message = get_fused_attn_backend( replace(fp8_baseline, head_dim_qk=1024, head_dim_v=1024) ) assert backend == tex.NVTE_Fused_Attn_Backend.NVTE_No_Backend assert message != "" +@pytest.mark.parametrize("dtype", (torch.float16, torch.bfloat16)) +@pytest.mark.parametrize("qkv_format", ("bshd", "sbhd")) +def test_fused_attn_python_execution(monkeypatch, dtype, qkv_format): + """The public module must use Python cuDNN for selection, forward, and backward.""" + + def common_attention(*_args, **_kwargs): + raise AssertionError("PyTorch attention called the TE common attention API") + + for name in ("get_fused_attn_backend", "fused_attn_fwd", "fused_attn_bwd"): + monkeypatch.setattr(tex, name, common_attention, raising=False) + monkeypatch.setenv("NVTE_FLASH_ATTN", "0") + monkeypatch.setenv("NVTE_UNFUSED_ATTN", "0") + monkeypatch.setenv("NVTE_FUSED_ATTN", "1") + _attention_backends["backend_selection_requires_update"] = True + shape = (2, 128, 4, 64) if qkv_format == "bshd" else (128, 2, 4, 64) + q, k, v = [torch.randn(shape, device="cuda", dtype=dtype, requires_grad=True) for _ in range(3)] + block = DotProductAttention(4, 64, qkv_format=qkv_format, attn_mask_type="causal") + try: + actual = block(q, k, v).reshape(shape) + order = (0, 2, 1, 3) if qkv_format == "bshd" else (1, 2, 0, 3) + inverse = (0, 2, 1, 3) if qkv_format == "bshd" else (2, 0, 1, 3) + refs = [x.detach().float().permute(order).requires_grad_() for x in (q, k, v)] + reference = torch.nn.functional.scaled_dot_product_attention(*refs, is_causal=True) + grad = torch.randn_like(actual) + actual.backward(grad) + reference.backward(grad.float().permute(order)) + rtol, atol = (2e-2, 1e-2) if dtype == torch.bfloat16 else (3e-3, 2e-3) + torch.testing.assert_close(actual.float(), reference.permute(inverse), rtol=rtol, atol=atol) + for tensor, ref in zip((q, k, v), refs): + torch.testing.assert_close( + tensor.grad.float(), ref.grad.permute(inverse), rtol=rtol, atol=atol + ) + finally: + _attention_backends["backend_selection_requires_update"] = True + + +@pytest.mark.parametrize( + "changes,version,capability,reason", + ( + ({"qkv_dtype": DType.kFloat32}, 92700, (10, 0), "Unsupported QKV dtype"), + ( + {"cuda_graph": True, "max_seqlen_kv": 127}, + 91500, + (10, 0), + "Known cuDNN <= 9.15 issue with CUDA graph", + ), + ( + {"qkv_layout": QKVLayout["thd_thd_thd"], "num_tokens_q": 256, "num_tokens_kv": 256}, + 92700, + (10, 0), + "THD format requires PADDING", + ), + ( + { + "qkv_layout": QKVLayout["thd_thd_thd"], + "attn_mask_type": AttnMaskType["padding"], + "num_tokens_q": 256, + "num_tokens_kv": 256, + }, + 92700, + (8, 0), + "F16/BF16 fused attention with THD format requires sm90", + ), + ( + { + "qkv_layout": QKVLayout["thd_thd_thd"], + "attn_mask_type": AttnMaskType["padding"], + "softmax_type": tex.NVTE_Softmax_Type.NVTE_LEARNABLE_SOFTMAX, + "num_tokens_q": 256, + "num_tokens_kv": 256, + }, + 92500, + (10, 0), + "Known cuDNN < 9.26.0 issue with THD learnable softmax backward", + ), + ( + { + "qkv_dtype": DType.kFloat8E4M3, + "qkv_layout": QKVLayout["thd_thd_thd"], + "attn_mask_type": AttnMaskType["padding"], + "head_dim_qk": 192, + "num_tokens_q": 256, + "num_tokens_kv": 256, + }, + 92700, + (10, 0), + "FP8 THD attention supports head dimensions up to 128", + ), + ( + { + "qkv_dtype": DType.kFloat8E4M3, + "scaling_mode": tex.NVTEScalingMode.NVTE_MXFP8_1D_SCALING, + }, + 92000, + (10, 0), + "MXFP8 fused attention requires cuDNN 9.21.0", + ), + ( + {"qkv_layout": QKVLayout["paged_kv_bshd_bshd_bshd"]}, + 92700, + (10, 0), + "Paged KV requires PADDING", + ), + ), +) +def test_fused_attn_python_capability_exclusions(monkeypatch, changes, version, capability, reason): + """Preserve capability and tracked-defect exclusions before asking cuDNN to build.""" + import cudnn + from transformer_engine.pytorch.cpp_extensions.cudnn_fused_attn import _AttentionGraph + + monkeypatch.setattr(cudnn, "backend_version", lambda: version) + monkeypatch.setattr(torch.cuda, "get_device_capability", lambda *_args: capability) + + def build(_self): + raise AssertionError("Excluded configuration reached cuDNN graph construction") + + monkeypatch.setattr(_AttentionGraph, "build", build) + baseline = FusedAttentionParams( + qkv_layout=QKVLayout["bshd_bshd_bshd"], + batch_size=2, + num_attn_heads=4, + num_gqa_groups=4, + head_dim_qk=64, + head_dim_v=64, + max_seqlen_q=128, + max_seqlen_kv=128, + ) + backend, message = get_fused_attn_backend(replace(baseline, **changes)) + assert backend == FusedAttnBackend.No_Backend + assert message.startswith(reason) + + +@pytest.mark.skipif(not fp8_attn_available, reason=reason_for_no_fp8_attn) +def test_fused_attn_python_dropout_cuda_graph(): + """Dropout must reproduce seeded eager results and advance its mask on graph replay.""" + from transformer_engine.pytorch.tensor.float8_tensor import Float8Quantizer + + layout = "bshd_bshd_bshd" + config = FusedAttentionParams( + qkv_layout=QKVLayout[layout], + dqkv_layout=QKVLayout[layout], + o_format=QKVFormat["bshd"], + do_format=QKVFormat["bshd"], + qkv_dtype=DType.kFloat8E4M3, + o_dtype=DType.kFloat8E4M3, + do_dtype=DType.kFloat8E5M2, + dqkv_dtype=DType.kFloat8E5M2, + batch_size=2, + num_attn_heads=4, + num_gqa_groups=4, + head_dim_qk=64, + head_dim_v=64, + max_seqlen_q=128, + max_seqlen_kv=128, + dropout=0.125, + ) + backend, message = get_fused_attn_backend(config) + if backend == FusedAttnBackend.No_Backend: + pytest.skip(f"cuDNN FP8 dropout is unavailable: {message}") + + def quantizer(dtype, scale=1.0): + return Float8Quantizer( + torch.full((1,), scale, device="cuda"), + torch.zeros(1, device="cuda"), + dtype, + columnwise=False, + ) + + source = torch.ones(2, 128, 4, 64, dtype=torch.float16, device="cuda") + q, k, v = [quantizer(DType.kFloat8E4M3)(source) for _ in range(3)] + d_o = quantizer(DType.kFloat8E5M2)(source) + s_quantizer = quantizer(DType.kFloat8E4M3, 128.0) + dp_quantizer = quantizer(DType.kFloat8E5M2) + o_quantizer = quantizer(DType.kFloat8E4M3) + dqkv_quantizer = quantizer(DType.kFloat8E5M2) + cu = torch.tensor([0, 128, 256], dtype=torch.int32, device="cuda") + kwargs = { + "qkv_layout": layout, + "o_format": "bshd", + "attn_mask_type": "no_mask", + "dropout": config.dropout, + "s_quantizer": s_quantizer, + } + + def run(): + out, aux = fused_attn_fwd( + True, + 128, + 128, + cu, + cu, + q, + k, + v, + source.dtype, + backend, + o_quantizer=o_quantizer, + **kwargs, + ) + dq, dk, dv, *_ = fused_attn_bwd( + 128, + 128, + cu, + cu, + q, + k, + v, + out, + d_o, + source.dtype, + aux, + backend, + do_format="bshd", + dqkv_layout=layout, + dp_quantizer=dp_quantizer, + dqkv_quantizer=dqkv_quantizer, + **kwargs, + ) + return tuple(tensor.dequantize() for tensor in (out, dq, dk, dv)) + + torch.cuda.manual_seed(1234) + eager = run() + # With uniform scores and values, inverted dropout has expected output and dV equal to one. + torch.testing.assert_close( + eager[0].float().mean(), torch.ones((), device="cuda"), atol=0.05, rtol=0 + ) + torch.testing.assert_close( + eager[3].float().mean(), torch.ones((), device="cuda"), atol=0.05, rtol=0 + ) + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(stream): + for _ in range(3): + run() + torch.cuda.current_stream().wait_stream(stream) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph, stream=stream): + captured = run() + torch.cuda.manual_seed(1234) + graph.replay() + first = [tensor.clone() for tensor in captured] + for actual, expected in zip(first, eager): + torch.testing.assert_close(actual, expected, atol=1e-2, rtol=1e-2) + graph.replay() + assert not torch.equal(captured[0], first[0]) + torch.cuda.manual_seed(1234) + graph.replay() + for actual, expected in zip(captured, first): + torch.testing.assert_close(actual, expected, atol=1e-2, rtol=1e-2) + + +@pytest.mark.parametrize("custom_generator", (False, True)) +def test_fused_attn_rng_state_cuda_graph(custom_generator): + """Philox seeds and offsets must follow the selected generator through capture and replay.""" + if custom_generator and not hasattr(torch.cuda.CUDAGraph, "register_generator_state"): + pytest.skip( + "This PyTorch version cannot register custom CUDA generators for graph capture." + ) + generator = ( + torch.Generator(device="cuda") + if custom_generator + else torch.cuda.default_generators[torch.cuda.current_device()] + ) + rng_gen = generator if custom_generator else None + generator.manual_seed(1234) + first = tex.fused_attn_rng_state(rng_gen, 16) + second = tex.fused_attn_rng_state(rng_gen, 16) + assert first.tolist() == [1234, 0] + assert second.tolist() == [1234, 16] + graph = torch.cuda.CUDAGraph() + if custom_generator: + graph.register_generator_state(generator) + with torch.cuda.graph(graph): + captured = tex.fused_attn_rng_state(rng_gen, 16) + generator.manual_seed(1234) + graph.replay() + torch.testing.assert_close(captured, first) + graph.replay() + torch.testing.assert_close(captured, second) + generator.manual_seed(1234) + graph.replay() + torch.testing.assert_close(captured, first) + + @pytest.mark.parametrize( "cu_seqlens,cu_seqlens_padded,expected", [ @@ -3917,7 +4201,7 @@ def forward( @staticmethod def backward(ctx, grad_output: torch.Tensor) -> Tuple[Union[torch.Tensor, None], ...]: with torch.cuda.nvtx.range("_DPA"): - (q, k, v, inp_fp8, qkv_weight_fp8, out) = restore_from_func_ctx(ctx) + q, k, v, inp_fp8, qkv_weight_fp8, out = restore_from_func_ctx(ctx) proj_dgrad = ctx.dO_quantizer(grad_output) diff --git a/transformer_engine/pytorch/attention/dot_product_attention/backends.py b/transformer_engine/pytorch/attention/dot_product_attention/backends.py index 6c4ac477781..697e6b7b501 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/backends.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/backends.py @@ -2576,13 +2576,16 @@ def _needs_eager_fused_attention(call: Dict[str, Any]) -> Optional[str]: class FusedAttention(torch.nn.Module): """Dot product attention using `cuDNN attention `_: + PyTorch constructs and executes attention graphs through the cuDNN frontend Python API. + Standard attention and experimental score modification share graph lifecycle helpers. + Plans are cached per device and tensor metadata; support queries do not compile plans. + FusedAttnBackend["F16_arbitrary_seqlen"] cuDNN attention for FP16/BF16 with any sequence length. FusedAttnBackend["FP8"] cuDNN attention for FP8 with any sequence length. It supports the following recipes, where "Inputs", "Intermediates" and "Outputs" are in the format of "tensor: quantizer". The recipes - are implemented in transformer_engine.pytorch.cpp_extension.fused_attn.fused_attn_fwd and - transformer_engine.pytorch.cpp_extension.fused_attn.fused_attn_bwd. + are implemented in transformer_engine.pytorch.cpp_extensions.cudnn_fused_attn. Direction Inputs Intermediates Outputs DelayedScaling (DS) forward Q/K/V: DS S: DS O: DS diff --git a/transformer_engine/pytorch/attention/dot_product_attention/flex_attention.py b/transformer_engine/pytorch/attention/dot_product_attention/flex_attention.py index b9593b42d9b..22f388dc4b4 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/flex_attention.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/flex_attention.py @@ -5,28 +5,22 @@ """cuDNN-backed Flex Attention helpers.""" from dataclasses import dataclass -import importlib import inspect from typing import Any, Callable, Dict, Optional, Tuple import torch -_cudnn_score_mod_handles: Dict[torch.device, Any] = {} +from transformer_engine.pytorch.cpp_extensions.cudnn import ( + _import_cudnn_frontend, + _build_cudnn_pygraph, + _finalize_cudnn_graph, + _execute_cudnn_graph, +) + _cudnn_score_mod_graph_cache: Dict[Tuple[Any, ...], Any] = {} _SCORE_MOD_UNCACHEABLE = object() -def _import_cudnn_frontend(): - """Import the cuDNN frontend Python package.""" - try: - return importlib.import_module("cudnn") - except ImportError as exc: - raise ImportError( - "cuDNN frontend Python package not found. " - "Install it with: pip install nvidia-cudnn-frontend" - ) from exc - - def _bhsd_dim_stride( tensor: torch.Tensor, tensor_format: str ) -> Tuple[Tuple[int, ...], Tuple[int, ...]]: @@ -192,44 +186,6 @@ def _wrapped_score_mod(sdpa_graph, score_tensor): return _wrapped_score_mod -def _get_cudnn_current_stream_handle(cudnn, device: torch.device): - """Return a cuDNN handle for device, bound to PyTorch's current stream.""" - if device.type != "cuda": - raise ValueError(f"Flex Attention only supports CUDA tensors, got device {device}.") - if device.index is None: - device = torch.device("cuda", torch.cuda.current_device()) - - handle = _cudnn_score_mod_handles.get(device) - with torch.cuda.device(device): - if handle is None: - handle = cudnn.create_handle() - _cudnn_score_mod_handles[device] = handle - - stream = torch.cuda.current_stream(device).cuda_stream - cudnn.set_stream(handle=handle, stream=stream) - return handle - - -def _build_cudnn_pygraph(dtype: torch.dtype, device: torch.device): - """Create a cuDNN frontend Python graph for F16/BF16 SDPA.""" - cudnn = _import_cudnn_frontend() - - if dtype == torch.float16: - io_data_type = cudnn.data_type.HALF - elif dtype == torch.bfloat16: - io_data_type = cudnn.data_type.BFLOAT16 - else: - raise ValueError(f"Flex Attention only supports FP16/BF16 tensors, got {dtype}.") - - graph = cudnn.pygraph( - io_data_type=io_data_type, - intermediate_data_type=cudnn.data_type.FLOAT, - compute_data_type=cudnn.data_type.FLOAT, - handle=_get_cudnn_current_stream_handle(cudnn, device), - ) - return graph - - @dataclass class _CudnnScoreModFwdGraphEntry: """Cached cuDNN frontend graph and graph tensor handles for score_mod fprop.""" @@ -263,44 +219,6 @@ class _CudnnScoreModBwdGraphEntry: workspace_size: int -def _finalize_cudnn_graph(graph) -> int: - """Build a cuDNN frontend Python graph and return its workspace size.""" - cudnn = _import_cudnn_frontend() - - graph.validate() - graph.build_operation_graph() - try: - graph.create_execution_plans([cudnn.heur_mode.A, cudnn.heur_mode.FALLBACK]) - graph.check_support() - except cudnn.cudnnGraphNotSupportedError as exc: - raise RuntimeError(f"cuDNN Flex Attention SDPA graph is not supported: {exc}") from exc - graph.build_plans(cudnn.build_plan_policy.HEURISTICS_CHOICE) - return max(graph.get_workspace_size(), 1) - - -def _execute_cudnn_graph( - graph, - variant_pack: Dict[Any, torch.Tensor], - workspace_size: int, - device: torch.device, -): - """Execute a built cuDNN frontend Python graph.""" - cudnn = _import_cudnn_frontend() - - if device.type == "cuda" and device.index is None: - device = torch.device("cuda", torch.cuda.current_device()) - workspace = torch.empty( - workspace_size, - device=device, - dtype=torch.uint8, - ) - graph.execute( - variant_pack, - workspace, - handle=_get_cudnn_current_stream_handle(cudnn, device), - ) - - def _cudnn_score_mod_fwd_cache_key( is_training: bool, query_layer: torch.Tensor, diff --git a/transformer_engine/pytorch/attention/dot_product_attention/utils.py b/transformer_engine/pytorch/attention/dot_product_attention/utils.py index 3e11039c1f4..06e016bfa3d 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/utils.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/utils.py @@ -447,7 +447,7 @@ def error(self, *args, **kwargs): @torch.compiler.assume_constant_result def _get_fused_attn_backend(**fused_attn_kwargs): - """Constant-foldable tex.get_fused_attn_backend: the result depends only on + """Constant-foldable cuDNN Python support query: the result depends only on the attention config. Returns a plain int rather than a FusedAttnBackend member: dynamo @@ -456,7 +456,9 @@ def _get_fused_attn_backend(**fused_attn_kwargs): graph break because it is baked into the graph as a literal, while an enum member comes out of the reconstruction corrupted (see the cast at the call site, which restores the enum).""" - fused_attention_backend, reject_message = tex.get_fused_attn_backend( + from transformer_engine.pytorch.cpp_extensions.cudnn_fused_attn import get_fused_attn_backend + + fused_attention_backend, reject_message = get_fused_attn_backend( FusedAttentionParams(**fused_attn_kwargs) ) return int(fused_attention_backend), reject_message diff --git a/transformer_engine/pytorch/cpp_extensions/cudnn.py b/transformer_engine/pytorch/cpp_extensions/cudnn.py new file mode 100644 index 00000000000..fee40d0e98d --- /dev/null +++ b/transformer_engine/pytorch/cpp_extensions/cudnn.py @@ -0,0 +1,76 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""cuDNN Python frontend graph lifecycle shared by attention implementations.""" + +import importlib +import threading + +import torch + +_handles = threading.local() + + +def _import_cudnn_frontend(): + """Import the cuDNN frontend Python package.""" + try: + return importlib.import_module("cudnn") + except ImportError as exc: + raise ImportError( + "cuDNN frontend Python package not found. " + "Install it with: pip install nvidia-cudnn-frontend" + ) from exc + + +def _get_cudnn_current_stream_handle(cudnn, device): + """Get a thread-local, device-specific handle on PyTorch's current stream.""" + if device.type != "cuda": + raise ValueError(f"cuDNN attention requires CUDA tensors, got {device}.") + if device.index is None: + device = torch.device("cuda", torch.cuda.current_device()) + if not hasattr(_handles, "by_device"): + _handles.by_device = {} + with torch.cuda.device(device): + handle = _handles.by_device.get(device) + if handle is None: + handle = cudnn.create_handle() + _handles.by_device[device] = handle + cudnn.set_stream(handle=handle, stream=torch.cuda.current_stream(device).cuda_stream) + return handle + + +def _build_cudnn_pygraph(dtype, device): + """Create an attention graph with FP32 accumulation.""" + cudnn = _import_cudnn_frontend() + dtype = {torch.float16: cudnn.data_type.HALF, torch.bfloat16: cudnn.data_type.BFLOAT16}.get( + dtype, dtype + ) + return cudnn.pygraph( + io_data_type=dtype, + intermediate_data_type=cudnn.data_type.FLOAT, + compute_data_type=cudnn.data_type.FLOAT, + handle=_get_cudnn_current_stream_handle(cudnn, device), + ) + + +def _finalize_cudnn_graph(graph): + """Check support, build an execution plan, and return workspace bytes.""" + cudnn = _import_cudnn_frontend() + graph.validate() + graph.build_operation_graph() + graph.create_execution_plans([cudnn.heur_mode.A, cudnn.heur_mode.FALLBACK]) + graph.check_support() + graph.build_plans(cudnn.build_plan_policy.HEURISTICS_CHOICE) + return max(graph.get_workspace_size(), 1) + + +def _execute_cudnn_graph(graph, variant_pack, workspace_size, device): + """Execute on the current stream using framework-owned workspace.""" + cudnn = _import_cudnn_frontend() + workspace = torch.empty(workspace_size, device=device, dtype=torch.uint8) + graph.execute( + variant_pack, + workspace, + handle=_get_cudnn_current_stream_handle(cudnn, device), + ) diff --git a/transformer_engine/pytorch/cpp_extensions/cudnn_fused_attn.py b/transformer_engine/pytorch/cpp_extensions/cudnn_fused_attn.py new file mode 100644 index 00000000000..647c2667392 --- /dev/null +++ b/transformer_engine/pytorch/cpp_extensions/cudnn_fused_attn.py @@ -0,0 +1,1041 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""cuDNN attention graphs constructed and executed through the Python frontend. + +The extension is used only to obtain capture-safe Philox state. Tensor allocation, +graph construction, plan caching, and execution do not use the common attention API. +""" + +from dataclasses import dataclass +import atexit +from collections import Counter +import os +import sys +import threading +import time + +import torch +import transformer_engine_torch as tex + +from .cudnn import ( + _import_cudnn_frontend, + _build_cudnn_pygraph, + _execute_cudnn_graph, +) +from ..tensor.float8_tensor import Float8Quantizer +from ..tensor.storage.mxfp8_tensor_storage import MXFP8TensorStorage +from ..constants import DType + +_plans = {} +_cache_lock = threading.RLock() +_diagnostics = Counter() +_build_times = Counter() + + +def _debug_enabled(): + setting = os.getenv("NVTE_FUSED_ATTN_CACHE_DEBUG", "0") + level, _, ranks = setting.partition(":") + rank = next( + ( + os.environ[n] + for n in ("RANK", "LOCAL_RANK", "OMPI_COMM_WORLD_RANK", "SLURM_PROCID") + if n in os.environ + ), + None, + ) + enabled = rank is None or ranks == "all" or rank in (ranks or "0").split(",") + return int(level or 0) if enabled else 0 + + +def _event(attn, event, key=None): + level = _debug_enabled() + if level == 0: + return + site = ("fp8" if attn.fp8 else "f16", "bwd" if attn.backward else "fwd") + _diagnostics[(*site, event)] += 1 + if level >= 2: + print( + f"[FUSED-ATTN-CACHE] tid={threading.get_ident()} dev={attn.device.index} | " + f"{site[0]} {site[1]} {event}" + + (f" | {key!r}" if key is not None else ""), + file=sys.stderr, + flush=True, + ) + + +def _build_stage(attn, name, function): + if not _debug_enabled(): + return function() + start = time.perf_counter() + try: + return function() + finally: + site = ("fp8" if attn.fp8 else "f16", "bwd" if attn.backward else "fwd", name) + _build_times[(*site, "seconds")] += time.perf_counter() - start + _build_times[(*site, "calls")] += 1 + + +def _print_diagnostics(): + if not _diagnostics: + return + print("[FUSED-ATTN-CACHE] Python frontend summary", file=sys.stderr) + for (backend, direction, event), count in sorted(_diagnostics.items()): + print(f" {backend} {direction} {event}={count}", file=sys.stderr) + for backend, direction, stage, kind in sorted(_build_times): + if kind == "calls": + calls = _build_times[(backend, direction, stage, kind)] + seconds = _build_times[(backend, direction, stage, "seconds")] + print( + f" {backend} {direction} {stage}: {calls} calls, " + f"{1000 * seconds / calls:.3f} ms average", + file=sys.stderr, + ) + + +atexit.register(_print_diagnostics) + + +def _formats(layout): + parts = layout.removeprefix("paged_kv_").replace("2", "").replace("3", "").split("_") + return parts[0], parts[1] if len(parts) > 1 else parts[0] + + +def _strides(dim, fmt): + b, h, s, d = dim + if fmt in ("bshd", "thd"): + return (s * h * d, d, h * d, 1) + if fmt == "sbhd": + return (h * d, d, b * h * d, 1) + if fmt == "bhsd": + return (h * s * d, s * d, d, 1) + raise ValueError(f"Unsupported attention format: {fmt}") + + +def _shape(dim, fmt, tokens=0): + b, h, s, d = dim + if fmt == "bshd": + return (b, s, h, d) + if fmt == "sbhd": + return (s, b, h, d) + if fmt == "bhsd": + return dim + if fmt == "thd": + return (tokens, h, d) + raise ValueError(f"Unsupported attention format: {fmt}") + + +def _raw(tensor, columnwise=False): + if isinstance(tensor, MXFP8TensorStorage): + return tensor._columnwise_data if columnwise else tensor._rowwise_data + return getattr(tensor, "_data", tensor) + + +def _dtype(tensor): + fp8_dtype = getattr(tensor, "_fp8_dtype", None) + if fp8_dtype is not None: + return "FP8_E4M3" if fp8_dtype == DType.kFloat8E4M3 else "FP8_E5M2" + return { + torch.float16: "HALF", + torch.bfloat16: "BFLOAT16", + torch.float32: "FLOAT", + torch.int32: "INT32", + torch.int64: "INT64", + }[tensor.dtype] + + +def _descale(tensor, columnwise=False): + if isinstance(tensor, MXFP8TensorStorage): + return tensor._columnwise_scale_inv if columnwise else tensor._rowwise_scale_inv + return tensor._scale_inv + + +def _round_up(value, multiple): + return (value + multiple - 1) // multiple * multiple + + +def _token_bucket(tokens): + if tokens == 0: + return 0 + if tokens <= 1024: + return 1024 + if tokens <= 32768: + return 1 << (tokens - 1).bit_length() + return _round_up(tokens, 32768) + + +@dataclass(frozen=True) +class _TensorSpec: + dim: tuple + stride: tuple + dtype: str + output: bool = False + ragged_offset: str = "" + reordered: bool = False + by_value: bool = False + + +class _AttentionGraph: + """Runtime tensors and data-free descriptors for one attention invocation.""" + + def __init__(self, params, backward=False): + self.params = params + self.backward = backward + self.device = params["q"].device + self.q_format, self.kv_format = _formats(params["qkv_layout"]) + self.b = params["cu_seqlens_q"].numel() - 1 + self.h = params["q"].shape[1] if self.q_format == "bhsd" else params["q"].shape[-2] + self.hg = params["k"].shape[1] if self.kv_format == "bhsd" else params["k"].shape[-2] + self.dq = params["q"].shape[-1] + self.dv = params["v"].shape[-1] + self.tq = params["q"].shape[0] if self.q_format == "thd" else 0 + self.tk = params["k"].shape[0] if self.kv_format == "thd" else 0 + capability = torch.cuda.get_device_capability(self.device) + ragged_graph = capability >= (9, 0) and capability != (12, 0) + self.ragged_stats = self.q_format == "thd" and ragged_graph + self.sq = _token_bucket(self.tq) if self.ragged_stats else params["max_seqlen_q"] + self.skv = ( + _token_bucket(self.tk) + if self.kv_format == "thd" and ragged_graph + else params["max_seqlen_kv"] + ) + self.fp8 = hasattr(params["q"], "_fp8_dtype") + self.mxfp8 = isinstance(params["q"], MXFP8TensorStorage) + self.specs = {} + self.tensors = {} + + def add(self, name, tensor, dim=None, stride=None, dtype=None, **kwargs): + """Bind a runtime tensor and record its address-independent descriptor.""" + self.tensors[name] = tensor + self.specs[name] = _TensorSpec( + tuple(dim if dim is not None else tensor.shape), + tuple(stride if stride is not None else tensor.stride()), + dtype or _dtype(tensor), + **kwargs, + ) + return name + + def scalar(self, name, tensor, output=False, by_value=False): + """Describe a scalar using cuDNN's four-dimensional convention.""" + return self.add(name, tensor, (1, 1, 1, 1), (1, 1, 1, 1), output=output, by_value=by_value) + + def matrix( + self, name, tensor, fmt, heads, seqlen, cu_padded=None, columnwise=False, output=False + ): + """Describe an attention matrix in logical BHSD order, retaining storage strides.""" + data = _raw(tensor, columnwise) + dim = (self.b, heads, seqlen, tensor.shape[-1]) + if fmt == "bshd": + order = (0, 2, 1, 3) + elif fmt == "sbhd": + order = (1, 2, 0, 3) + else: + order = (0, 1, 2, 3) + offset_name = "" + if fmt == "thd": + token_stride, head_stride, last_stride = data.stride() + stride = (seqlen * token_stride, head_stride, token_stride, last_stride) + offset_name = name + "_offset" + offsets = cu_padded.to(torch.int64) * token_stride + self.add(offset_name, offsets, (self.b + 1, 1, 1, 1), (1, 1, 1, 1)) + elif data.ndim == 4: + stride = tuple(data.stride(i) for i in order) + # Paged K/V have a page axis rather than a batch axis. + if name in ("k", "v") and self.params.get("page_table_k") is not None: + dim = tuple(data.shape[i] for i in order) + else: + stride = _strides(dim, fmt) + return self.add( + name, data, dim, stride, _dtype(tensor), output=output, ragged_offset=offset_name + ) + + def stats(self, name, tensor, output=False): + """Describe softmax statistics, including offsets for packed THD storage.""" + dim = (self.b, self.h, self.sq, 1) + if self.ragged_stats: + offsets = self.params["cu_seqlens_q_padded"].to(torch.int64) * self.h + self.add("stats_offset", offsets, (self.b + 1, 1, 1, 1), (1, 1, 1, 1)) + return self.add( + name, + tensor, + dim, + (self.h * self.sq, 1, self.h, 1), + output=output, + ragged_offset="stats_offset", + ) + return self.add(name, tensor, dim, (self.h * self.sq, self.sq, 1, 1), output=output) + + def scale(self, name, tensor, matrix_name, fmt, columnwise=False): + """Describe per-tensor FP8 scales or reordered MXFP8 block scales.""" + if not self.mxfp8: + return self.scalar(name, tensor) + b, h, s, d = self.specs[matrix_name].dim + if columnwise: + dim = (b, h, _round_up((s + 31) // 32, 4), _round_up(d, 128)) + else: + dim = (b, h, _round_up(s, 128), _round_up((d + 31) // 32, 4)) + return self.add(name, tensor, dim, _strides(dim, fmt), "FP8_E8M0", reordered=True) + + def options(self): + """Return graph attributes that affect support and plan caching.""" + p = self.params + mask = p["attn_mask_type"] + causal = "causal" in mask + bottom_right = p["bottom_right_diagonal"] + if causal: + bottom_right = "bottom_right" in mask + left, right = p["window_size"] or (-1, -1) + if not causal and (left, right) == (-1, -1): + bottom_right = False + return ( + mask, + p["attn_bias_type"], + p["softmax_type"], + left + 1 if left != -1 else None, + 0 if causal else (right if right != -1 else None), + bottom_right, + p["dropout"], + p.get("deterministic", False), + p.get("return_max_logit", False), + self.fp8, + self.mxfp8, + self.params["max_seqlen_kv"] if self.params.get("page_table_k") is not None else 0, + ) + + def entry(self): + """Share support-query graphs with execution without compiling on queries.""" + key = (self.device, self.backward, self.options(), tuple(sorted(self.specs.items()))) + # Autograd may run backward on another thread. Share device-specific plans; + # each execution supplies its thread-local handle on the current stream. + with _cache_lock: + entry = _plans.get(key) + if entry is None: + _event(self, "MISS", key) + entry = self.build() + _plans[key] = entry + _event(self, "CACHE_GRAPH") + else: + _event(self, "HIT", key) + return entry + + def execute(self): + """Look up the graph by metadata, then bind this call's tensor addresses.""" + entry = self.entry() + with _cache_lock: + if entry[2] is None: + cudnn = _import_cudnn_frontend() + _build_stage( + self, + "build_plans", + lambda: entry[0].build_plans(cudnn.build_plan_policy.HEURISTICS_CHOICE), + ) + entry[2] = max(entry[0].get_workspace_size(), 1) + _event(self, "BUILD_PLANS") + graph, tensors, workspace = entry + _execute_cudnn_graph( + graph, {tensors[n]: t for n, t in self.tensors.items()}, workspace, self.device + ) + _event(self, "EXECUTE") + + def build(self): + """Construct and check a frontend graph without compiling its execution plan.""" + cudnn = _import_cudnn_frontend() + graph = _build_cudnn_pygraph(getattr(cudnn.data_type, self.specs["q"].dtype), self.device) + tensors = {} + + def tensor(name): + if name not in tensors: + spec = self.specs[name] + desc = graph.tensor( + name=name, + dim=spec.dim, + stride=spec.stride, + data_type=getattr(cudnn.data_type, spec.dtype), + is_pass_by_value=spec.by_value, + ) + if spec.ragged_offset: + desc.set_ragged_offset(tensor(spec.ragged_offset)) + if spec.reordered: + desc.set_reordering_type(cudnn.tensor_reordering.F8_128x4) + if spec.output: + desc.set_output(True) + tensors[name] = desc + return tensors[name] + + for name, spec in self.specs.items(): + if not spec.output: + tensor(name) + + mask, bias, softmax, left, right, bottom_right, dropout, deterministic, _, _, _, pages = ( + self.options() + ) + kwargs = { + "attn_scale": tensor("attn_scale"), + "diagonal_alignment": ( + cudnn.diagonal_alignment.BOTTOM_RIGHT + if bottom_right + else cudnn.diagonal_alignment.TOP_LEFT + ), + } + bound_names = ( + ("left_bound", "right_bound") + if (self.fp8 and (not self.mxfp8 or self.backward)) + else ("diagonal_band_left_bound", "diagonal_band_right_bound") + ) + kwargs.update(zip(bound_names, (left, right))) + if "padding" in mask: + kwargs.update( + use_padding_mask=True, seq_len_q=tensor("seq_q"), seq_len_kv=tensor("seq_kv") + ) + if bias == "alibi": + kwargs["use_alibi_mask"] = True + elif bias != "no_bias": + kwargs["bias"] = tensor("bias") + if self.backward and "dbias" in self.specs: + kwargs["dBias"] = tensor("dbias") + if dropout: + kwargs["dropout"] = (dropout, tensor("seed"), tensor("offset")) + if softmax != "vanilla": + kwargs["sink_token"] = tensor("softmax_offset") + if self.backward: + kwargs["dSink_token"] = tensor("dsoftmax_offset") + if self.backward: + kwargs["use_deterministic_algorithm"] = deterministic + else: + kwargs["generate_stats"] = True + if "max" in self.specs: + kwargs["score_max"] = tensor("max") + if pages: + kwargs.update( + paged_attention_k_table=tensor("page_table_k"), + paged_attention_v_table=tensor("page_table_v"), + paged_attention_max_seq_len_kv=pages, + ) + + q, k, v = (tensor(n) for n in ("q", "k", "v")) + if not self.fp8: + if self.backward: + outputs = graph.sdpa_backward( + q, k, v, tensor("o"), tensor("do"), tensor("stats"), **kwargs + ) + names = ("dq", "dk", "dv") + else: + outputs = graph.sdpa(q, k, v, **kwargs) + names = ("o", "stats") + elif not self.backward: + scale_args = [tensor(n) for n in ("descale_q", "descale_k", "descale_v")] + if self.mxfp8: + outputs = graph.sdpa_mxfp8(q, k, v, *scale_args, **kwargs) + names = ("o", "stats", "amax_o") + else: + outputs = graph.sdpa_fp8( + q, + k, + v, + *scale_args, + *(tensor(n) for n in ("descale_s", "scale_s", "scale_o")), + **kwargs, + ) + names = ("o", "stats", "amax_s", "amax_o") + elif self.mxfp8: + outputs = graph.sdpa_mxfp8_backward( + q, + tensor("q_t"), + k, + tensor("k_t"), + v, + tensor("o"), + tensor("do_f16"), + tensor("do"), + tensor("do_t"), + tensor("stats"), + *( + tensor(n) + for n in ( + "descale_q", + "descale_q_t", + "descale_k", + "descale_k_t", + "descale_v", + "descale_do", + "descale_do_t", + ) + ), + **kwargs, + ) + names = ("dq", "dk", "dv", "amax_dq", "amax_dk", "amax_dv") + else: + outputs = graph.sdpa_fp8_backward( + q, + k, + v, + tensor("o"), + tensor("do"), + tensor("stats"), + *( + tensor(n) + for n in ( + "descale_q", + "descale_k", + "descale_v", + "descale_o", + "descale_do", + "descale_s", + "descale_dp", + "scale_s", + "scale_dq", + "scale_dk", + "scale_dv", + "scale_dp", + ) + ), + **kwargs, + ) + names = ("dq", "dk", "dv", "amax_dq", "amax_dk", "amax_dv", "amax_dp") + + for name, desc in zip(names, outputs): + if name not in self.specs: + # MXFP8 amax outputs stay virtual, matching the C++ graph. + desc.set_dim((1, 1, 1, 1)).set_stride((1, 1, 1, 1)).set_data_type( + cudnn.data_type.FLOAT + ) + continue + spec = self.specs[name] + desc.set_output(True).set_dim(spec.dim).set_stride(spec.stride).set_data_type( + getattr(cudnn.data_type, spec.dtype) + ) + if spec.ragged_offset: + desc.set_ragged_offset(tensor(spec.ragged_offset)) + tensors[name] = desc + _event(self, "CREATE_GRAPH") + _build_stage(self, "validate", graph.validate) + _build_stage(self, "build_operation_graph", graph.build_operation_graph) + _build_stage( + self, + "create_execution_plans", + lambda: graph.create_execution_plans([cudnn.heur_mode.A, cudnn.heur_mode.FALLBACK]), + ) + _build_stage(self, "check_support", graph.check_support) + return [graph, tensors, None] + + +def _prepare_inputs(attn): + p = attn.params + for name, fmt, h, s in ( + ("q", attn.q_format, attn.h, attn.sq), + ("k", attn.kv_format, attn.hg, attn.skv), + ("v", attn.kv_format, attn.hg, attn.skv), + ): + cu = p["cu_seqlens_q_padded"] if name == "q" else p["cu_seqlens_kv_padded"] + col = attn.mxfp8 and name == "v" and not attn.backward + attn.matrix(name, p[name], fmt, h, s, cu, col) + if attn.fp8: + scale_fmt = p.get("qkv_scale_inv_format") or fmt + attn.scale("descale_" + name, _descale(p[name], col), name, scale_fmt, col) + attn.scalar("attn_scale", torch.tensor([p["attn_scale"]], dtype=torch.float32), by_value=True) + if "padding" in p["attn_mask_type"]: + for name, cu in (("seq_q", p["cu_seqlens_q"]), ("seq_kv", p["cu_seqlens_kv"])): + attn.add(name, cu[1:] - cu[:-1], (attn.b, 1, 1, 1), (1, 1, 1, 1)) + if p["attn_bias_type"] not in ("no_bias", "alibi"): + attn.add("bias", p["attn_bias"]) + if p["softmax_type"] != "vanilla": + attn.add("softmax_offset", p["softmax_offset"]) + if p["dropout"]: + rng = p["rng_state"] + attn.scalar("seed", rng[:1]) + attn.scalar("offset", rng[1:]) + if p.get("page_table_k") is not None: + for name in ("page_table_k", "page_table_v"): + table = p[name] + attn.add( + name, + table, + (attn.b, 1, table.shape[-1], 1), + (table.shape[-1], table.shape[-1], 1, 1), + ) + + +def _allocate(shape, dtype, device, quantizer, data=None): + """Allocate an output, preserving delayed-scaling metadata and packed views.""" + delayed = isinstance(quantizer, Float8Quantizer) + if data is None: + data = torch.empty(shape, dtype=torch.uint8 if delayed else dtype, device=device) + if delayed: + return quantizer.create_tensor_from_data( + data, fake_dtype=dtype, internal=quantizer.internal + ) + return data + + +def _output_shape(attn, name, fmt): + is_q = name == "q" + return _shape( + ( + attn.b, + attn.h if is_q else attn.hg, + attn.params["max_seqlen_q"] if is_q else attn.params["max_seqlen_kv"], + attn.dq if name in ("q", "k") else attn.dv, + ), + fmt, + attn.tq if is_q else attn.tk, + ) + + +def _allocate_gradients(attn): + p = attn.params + qfmt, kfmt = _formats(p["dqkv_layout"]) + shapes = [_output_shape(attn, n, qfmt if n == "q" else kfmt) for n in ("q", "k", "v")] + delayed = isinstance(p["dqkv_quantizer"], Float8Quantizer) + dtype = torch.uint8 if delayed else p["fake_dtype"] + layout = p["dqkv_layout"] + if "3" in layout: + axis = -2 if "h3d" in layout else -3 + packed_shape = list(shapes[0]) + packed_shape.insert(len(packed_shape) + axis + 1, 3) + data = torch.empty(packed_shape, dtype=dtype, device=attn.device).unbind(axis) + elif "2" in layout: + axis = -2 if "h2d" in layout else -3 + packed_shape = list(shapes[1]) + packed_shape.insert(len(packed_shape) + axis + 1, 2) + kv = torch.empty(packed_shape, dtype=dtype, device=attn.device).unbind(axis) + data = (torch.empty(shapes[0], dtype=dtype, device=attn.device), *kv) + else: + data = [torch.empty(s, dtype=dtype, device=attn.device) for s in shapes] + return [ + _allocate(s, p["fake_dtype"], attn.device, p["dqkv_quantizer"], t) + for s, t in zip(shapes, data) + ] + + +def _normalize(params): + p = dict(params) + p["cu_seqlens_q_padded"] = p.get("cu_seqlens_q_padded") + if p["cu_seqlens_q_padded"] is None: + p["cu_seqlens_q_padded"] = p["cu_seqlens_q"] + p["cu_seqlens_kv_padded"] = p.get("cu_seqlens_kv_padded") + if p["cu_seqlens_kv_padded"] is None: + p["cu_seqlens_kv_padded"] = p["cu_seqlens_kv"] + return p + + +def _query_graph(config, backward): + """Describe a support query without allocating tensors or touching RNG state.""" + from .fused_attn import QKVLayout, QKVFormat, AttnMaskType, AttnBiasType, SoftmaxType + + def name(mapping, value): + return next(k for k, v in mapping.items() if v == value) + + layout = name(QKVLayout, config.qkv_layout) + grad_layout = ( + name(QKVLayout, config.dqkv_layout) + if (config.dqkv_layout != tex.NVTE_QKV_Layout.NVTE_QKV_Layout_NOT_SET) + else layout + ) + attn = _AttentionGraph.__new__(_AttentionGraph) + attn.backward = backward + attn.device = torch.device("cuda", torch.cuda.current_device()) + attn.q_format, attn.kv_format = _formats(layout) + attn.b, attn.h, attn.hg = config.batch_size, config.num_attn_heads, config.num_gqa_groups + attn.dq, attn.dv = config.head_dim_qk, config.head_dim_v + attn.tq, attn.tk = config.num_tokens_q, config.num_tokens_kv + ragged_graph = torch.cuda.get_device_capability() >= (9, 0) and ( + torch.cuda.get_device_capability() != (12, 0) + ) + attn.ragged_stats = attn.q_format == "thd" and ragged_graph + attn.sq = _token_bucket(attn.tq) if attn.ragged_stats else config.max_seqlen_q + attn.skv = ( + _token_bucket(attn.tk) if attn.kv_format == "thd" and ragged_graph else config.max_seqlen_kv + ) + types = { + DType.kFloat16: "HALF", + DType.kBFloat16: "BFLOAT16", + DType.kFloat32: "FLOAT", + DType.kFloat8E4M3: "FP8_E4M3", + DType.kFloat8E5M2: "FP8_E5M2", + } + attn.fp8 = types[DType.cast(config.qkv_dtype)].startswith("FP8") + attn.mxfp8 = attn.fp8 and config.scaling_mode == tex.NVTEScalingMode.NVTE_MXFP8_1D_SCALING + attn.params = { + "attn_mask_type": name(AttnMaskType, config.attn_mask_type), + "attn_bias_type": name(AttnBiasType, config.bias_type), + "softmax_type": name(SoftmaxType, config.softmax_type), + "bottom_right_diagonal": config.bottom_right_diagonal, + "window_size": (config.window_size_left, config.window_size_right), + "dropout": config.dropout if config.is_training else 0.0, + "deterministic": config.deterministic if backward else False, + "return_max_logit": config.return_max_logit if not backward else False, + "max_seqlen_kv": config.max_seqlen_kv, + "page_table_k": True if layout.startswith("paged_kv_") else None, + } + attn.specs, attn.tensors = {}, {} + + def spec(n, dim=(1, 1, 1, 1), stride=(1, 1, 1, 1), dtype="FLOAT", **kwargs): + attn.specs[n] = _TensorSpec(tuple(dim), tuple(stride), dtype, **kwargs) + + def matrix(n, fmt, h, s, d, dtype, packed="", output=False): + dim = (attn.b, h, s, d) + if n in ("k", "v") and layout.startswith("paged_kv_"): + dim = (getattr(config, "num_pages_" + n), h, getattr(config, "page_size_" + n), d) + stride = list(_strides(dim, fmt)) + if "3" in packed or "2" in packed: + count = 3 if "3" in packed else 2 + stride[0] *= count + stride[2] *= count + if "h3d" in packed or "h2d" in packed: + stride[1] *= count + offset = "" + if fmt == "thd": + offset = n + "_offset" + spec(offset, (attn.b + 1, 1, 1, 1), dtype="INT64") + spec(n, dim, stride, dtype, output=output, ragged_offset=offset) + + def stats(n, output=False): + stride = ( + (attn.h * attn.sq, 1, attn.h, 1) + if attn.ragged_stats + else (attn.h * attn.sq, attn.sq, 1, 1) + ) + offset = "stats_offset" if attn.ragged_stats else "" + if offset: + spec(offset, (attn.b + 1, 1, 1, 1), dtype="INT64") + spec(n, (attn.b, attn.h, attn.sq, 1), stride, output=output, ragged_offset=offset) + + def scale(n, matrix_name, fmt, columnwise=False): + if not attn.mxfp8: + spec(n) + return + b, h, s, d = attn.specs[matrix_name].dim + dim = ( + (b, h, _round_up((s + 31) // 32, 4), _round_up(d, 128)) + if columnwise + else (b, h, _round_up(s, 128), _round_up((d + 31) // 32, 4)) + ) + spec(n, dim, _strides(dim, fmt), "FP8_E8M0", reordered=True) + + parts = layout.removeprefix("paged_kv_").split("_") + for n, fmt, h, s, d, part in ( + ("q", attn.q_format, attn.h, attn.sq, attn.dq, parts[0]), + ("k", attn.kv_format, attn.hg, attn.skv, attn.dq, parts[1] if len(parts) > 1 else parts[0]), + ("v", attn.kv_format, attn.hg, attn.skv, attn.dv, parts[-1]), + ): + matrix(n, fmt, h, s, d, types[DType.cast(config.qkv_dtype)], part) + if attn.fp8: + scale_fmt = name(QKVFormat, config.qkv_scale_inv_format) or fmt + scale("descale_" + n, n, scale_fmt, attn.mxfp8 and n == "v" and not backward) + spec("attn_scale", by_value=True) + if "padding" in attn.params["attn_mask_type"]: + for n in ("seq_q", "seq_kv"): + spec(n, (attn.b, 1, 1, 1), dtype="INT32") + if attn.params["dropout"]: + spec("seed", dtype="INT64") + spec("offset", dtype="INT64") + if attn.params["attn_bias_type"] not in ("no_bias", "alibi"): + dim = ( + config.bias_batch_size, + config.bias_num_heads, + config.bias_seqlen_q, + config.bias_seqlen_kv, + ) + b, h, sq, skv = dim + stride = (h * sq * skv, sq * skv, skv, 1) + spec("bias", dim, stride, types[DType.cast(config.qkv_dtype)]) + if backward and (b, h, sq) != (1, 1, 1): + spec("dbias", dim, stride, types[DType.cast(config.qkv_dtype)], output=True) + if attn.params["softmax_type"] != "vanilla": + spec("softmax_offset", (1, attn.h, 1, 1), (attn.h, 1, 1, 1)) + if backward: + spec("dsoftmax_offset", (1, attn.h, 1, 1), (attn.h, 1, 1, 1), output=True) + if layout.startswith("paged_kv_"): + for n in ("k", "v"): + pages = getattr(config, "max_pages_per_seq_" + n) + spec("page_table_" + n, (attn.b, 1, pages, 1), (pages, pages, 1, 1), "INT32") + ofmt = name(QKVFormat, config.o_format) or attn.q_format + matrix( + "o", ofmt, attn.h, attn.sq, attn.dv, types[DType.cast(config.o_dtype)], output=not backward + ) + stats("stats", output=not backward) + if not backward and config.return_max_logit: + stats("max", output=True) + if backward: + dofmt = name(QKVFormat, config.do_format) or ofmt + matrix("do", dofmt, attn.h, attn.sq, attn.dv, types[DType.cast(config.do_dtype)]) + qfmt, kvfmt = _formats(grad_layout) + parts = grad_layout.split("_") + for n, fmt, h, s, d, part in ( + ("dq", qfmt, attn.h, attn.sq, attn.dq, parts[0]), + ("dk", kvfmt, attn.hg, attn.skv, attn.dq, parts[1] if len(parts) > 1 else parts[0]), + ("dv", kvfmt, attn.hg, attn.skv, attn.dv, parts[-1]), + ): + matrix(n, fmt, h, s, d, types[DType.cast(config.dqkv_dtype)], part, output=True) + if attn.fp8: + do_scale_fmt = name(QKVFormat, config.do_scale_inv_format) or dofmt + scale("descale_do", "do", do_scale_fmt) + if attn.mxfp8: + matrix("do_f16", dofmt, attn.h, attn.sq, attn.dv, types[DType.cast(config.o_dtype)]) + for n, fmt, h, s, d, dtype in ( + ("q_t", attn.q_format, attn.h, attn.sq, attn.dq, config.qkv_dtype), + ("k_t", attn.kv_format, attn.hg, attn.skv, attn.dq, config.qkv_dtype), + ("do_t", dofmt, attn.h, attn.sq, attn.dv, config.do_dtype), + ): + matrix(n, fmt, h, s, d, types[DType.cast(dtype)]) + sfmt = ( + do_scale_fmt + if n == "do_t" + else name(QKVFormat, config.qkv_scale_inv_format) or fmt + ) + scale("descale_" + n, n, sfmt, True) + else: + for n in ( + "descale_o", + "descale_s", + "descale_dp", + "scale_s", + "scale_dp", + "scale_dq", + "scale_dk", + "scale_dv", + ): + spec(n) + for n in ("amax_dq", "amax_dk", "amax_dv", "amax_dp"): + spec(n, output=True) + elif attn.fp8 and not attn.mxfp8: + for n in ("descale_s", "scale_s", "scale_o"): + spec(n) + for n in ("amax_s", "amax_o"): + spec(n, output=True) + return attn + + +def _support_reason(config, attn, version, capability): + """Preserve TE's hardware, recipe, and known cuDNN defect exclusions.""" + p = attn.params + ragged = attn.q_format == "thd" or attn.kv_format == "thd" + padding = "padding" in p["attn_mask_type"] + if ragged and not padding: + return "THD format requires PADDING / PADDING_CAUSAL / PADDING_CAUSAL_BOTTOM_RIGHT mask." + if (attn.q_format == "thd" and not attn.tq) or (attn.kv_format == "thd" and not attn.tk): + return "THD format requires num_tokens_q / num_tokens_kv to be set for the ragged inputs." + if p["page_table_k"] is not None: + if not padding: + return "Paged KV requires PADDING / PADDING_CAUSAL / PADDING_CAUSAL_BOTTOM_RIGHT mask." + if not all( + getattr(config, field + "_" + name) + for field in ("num_pages", "page_size", "max_pages_per_seq") + for name in ("k", "v") + ): + return "Paged KV requires num_pages, page_size and max_pages_per_seq for both K and V." + if not attn.fp8: + if attn.q_format == attn.kv_format == "thd" and capability < (9, 0): + return "F16/BF16 fused attention with THD format requires sm90 or later." + affected_cuda_graph = version <= 91500 and config.cuda_graph and config.max_seqlen_kv % 128 + if ( + affected_cuda_graph + and config.is_training + and attn.q_format == attn.kv_format + and attn.q_format in ("bshd", "sbhd") + and not padding + ): + return "Known cuDNN <= 9.15 issue with CUDA graph. Please upgrade cuDNN." + if ( + config.is_training + and attn.ragged_stats + and p["softmax_type"] == "learnable" + and version < 92600 + ): + return ( + "Known cuDNN < 9.26.0 issue with THD learnable softmax backward. Please upgrade" + " cuDNN." + ) + return "" + if config.return_max_logit: + return "FP8 fused attention does not support return_max_logit=True." + if attn.q_format != attn.kv_format: + return "FP8 fused attention supports BSHD/SBHD/BHSD/THD formats." + if config.is_training and config.dqkv_layout != tex.NVTE_QKV_Layout.NVTE_QKV_Layout_NOT_SET: + from .fused_attn import QKVLayout + + grad_layout = next(k for k, v in QKVLayout.items() if v == config.dqkv_layout) + qfmt, kvfmt = _formats(grad_layout) + if qfmt != kvfmt: + return "FP8 fused attention supports BSHD/SBHD/BHSD/THD gradient formats." + if attn.q_format == "thd": + if version < 92300: + return "FP8 fused attention with THD format requires cuDNN 9.23.0 or later!" + if config.is_training and capability < (10, 0): + return "FP8 fused attention with THD format supports backward on sm100+ only!" + if config.is_training and p["softmax_type"] != "vanilla" and version < 92600: + return ( + "FP8 THD attention with a sink token requires cuDNN 9.26.0 or later for backward!" + ) + if capability >= (10, 0) and max(attn.dq, attn.dv) > 128: + return "FP8 THD attention supports head dimensions up to 128 on sm100+ only!" + if p["attn_bias_type"] != "no_bias": + return "FP8 fused attention does not support attention bias." + half_types = (DType.kFloat16, DType.kBFloat16) + fp8_types = (DType.kFloat8E4M3, DType.kFloat8E5M2) + allowed_types = half_types if attn.mxfp8 else half_types + fp8_types + tensor_scaling = config.scaling_mode == tex.NVTEScalingMode.NVTE_DELAYED_TENSOR_SCALING + if ( + not (tensor_scaling or attn.mxfp8) + or DType.cast(config.o_dtype) not in allowed_types + or (config.is_training and DType.cast(config.dqkv_dtype) not in allowed_types) + ): + return ( + "FP8 fused attention only supports FP8DelayedScaling, FP8CurrentScaling, or MXFP8" + " recipes!" + ) + if attn.mxfp8 and version < 92100: + return "MXFP8 fused attention requires cuDNN 9.21.0 or later!" + return "" + + +def get_fused_attn_backend(config): + """Select an attention backend by validating the Python frontend graphs.""" + from .fused_attn import FusedAttnBackend + + cudnn = _import_cudnn_frontend() + if config.bias_type == tex.NVTE_Bias_Type.NVTE_PRE_SCALE_BIAS: + return FusedAttnBackend.No_Backend, "Fused attention does not support pre-scale bias." + if DType.cast(config.qkv_dtype) not in ( + DType.kFloat16, + DType.kBFloat16, + DType.kFloat8E4M3, + DType.kFloat8E5M2, + ): + return FusedAttnBackend.No_Backend, f"Unsupported QKV dtype qkv_dtype={config.qkv_dtype}." + try: + fwd = _query_graph(config, False) + reason = _support_reason( + config, fwd, cudnn.backend_version(), torch.cuda.get_device_capability() + ) + if reason: + return FusedAttnBackend.No_Backend, reason + fwd.entry() + if config.is_training: + _query_graph(config, True).entry() + return FusedAttnBackend.FP8 if fwd.fp8 else FusedAttnBackend.F16_arbitrary_seqlen, "" + except (cudnn.cudnnGraphNotSupportedError, ValueError) as exc: + return FusedAttnBackend.No_Backend, str(exc) + + +def fused_attn_fwd(params): + """Execute forward; preserve the extension's output and auxiliary tensor order.""" + p = _normalize(params) + with torch.cuda.device(p["q"].device): + p["rng_state"] = tex.fused_attn_rng_state(p["rng_gen"], p["rng_elts_per_thread"]) + attn = _AttentionGraph(p) + _prepare_inputs(attn) + out_shape = list(_output_shape(attn, "q", p["o_format"])) + out_shape[-1] = attn.dv + out = _allocate(out_shape, p["fake_dtype"], attn.device, p["o_quantizer"]) + if p["o_format"] == "thd": + _raw(out).zero_() + attn.matrix("o", out, p["o_format"], attn.h, attn.sq, p["cu_seqlens_q_padded"], output=True) + stats_shape = (attn.tq, attn.h, 1) if attn.ragged_stats else (attn.b, attn.h, attn.sq, 1) + stats = torch.empty(stats_shape, dtype=torch.float32, device=attn.device) + attn.stats("stats", stats, output=True) + aux_ctx_tensors = [stats, p["rng_state"]] + maximum = None + if p["return_max_logit"]: + maximum = torch.empty_like(stats) + attn.stats("max", maximum, output=True) + if p["attn_bias_type"] not in ("no_bias", "alibi"): + aux_ctx_tensors.append(p["attn_bias"]) + if p["softmax_type"] != "vanilla": + aux_ctx_tensors.append(p["softmax_offset"]) + if attn.fp8 and not attn.mxfp8: + sq = p["s_quantizer"] + attn.scalar("scale_s", sq.scale) + attn.scalar("descale_s", sq.scale.reciprocal()) + attn.scalar("amax_s", sq.amax, output=True) + if isinstance(p["o_quantizer"], Float8Quantizer): + scale_o, amax_o = p["o_quantizer"].scale, p["o_quantizer"].amax + else: + scale_o = torch.ones(1, dtype=torch.float32, device=attn.device) + amax_o = torch.empty_like(scale_o) + attn.scalar("scale_o", scale_o) + attn.scalar("amax_o", amax_o, output=True) + attn.execute() + return out, aux_ctx_tensors, maximum + + +def fused_attn_bwd(params): + """Execute backward, including FP8/MXFP8 and shared packed gradient storage.""" + p = _normalize(params) + aux = iter(p["aux_ctx_tensors"]) + stats, p["rng_state"] = next(aux), next(aux) + p["attn_bias"] = next(aux) if p["attn_bias_type"] not in ("no_bias", "alibi") else None + p["softmax_offset"] = next(aux) if p["softmax_type"] != "vanilla" else None + with torch.cuda.device(p["q"].device): + attn = _AttentionGraph(p, backward=True) + _prepare_inputs(attn) + attn.matrix("o", p["o"], p["o_format"], attn.h, attn.sq, p["cu_seqlens_q_padded"]) + attn.matrix("do", p["d_o"], p["do_format"], attn.h, attn.sq, p["cu_seqlens_q_padded"]) + attn.stats("stats", stats) + dq, dk, dv = _allocate_gradients(attn) + qfmt, kfmt = _formats(p["dqkv_layout"]) + for name, t, fmt, h, s, cu in ( + ("dq", dq, qfmt, attn.h, attn.sq, p["cu_seqlens_q_padded"]), + ("dk", dk, kfmt, attn.hg, attn.skv, p["cu_seqlens_kv_padded"]), + ("dv", dv, kfmt, attn.hg, attn.skv, p["cu_seqlens_kv_padded"]), + ): + if fmt == "thd": + _raw(t).zero_() + attn.matrix(name, t, fmt, h, s, cu, output=True) + dbias = torch.empty(0, device=attn.device, dtype=p["fake_dtype"]) + if p["attn_bias"] is not None: + dbias = torch.empty_like(p["attn_bias"]) + if p["attn_bias"].shape[-2] != 1 or p["attn_bias"].shape[:2] != (1, 1): + attn.add("dbias", dbias, output=True) + dsoftmax = torch.empty(0, device=attn.device, dtype=torch.float32) + if p["softmax_type"] != "vanilla": + dsoftmax = torch.empty_like(p["softmax_offset"]) + attn.add("dsoftmax_offset", dsoftmax, output=True) + if attn.fp8: + do_scale_fmt = p.get("do_scale_inv_format") or p["do_format"] + attn.scale("descale_do", _descale(p["d_o"]), "do", do_scale_fmt) + if attn.mxfp8: + do_f16 = next(aux) + attn.matrix( + "do_f16", do_f16, p["do_format"], attn.h, attn.sq, p["cu_seqlens_q_padded"] + ) + for name, t, fmt, h, s, cu in ( + ("q_t", p["q"], attn.q_format, attn.h, attn.sq, p["cu_seqlens_q_padded"]), + ("k_t", p["k"], attn.kv_format, attn.hg, attn.skv, p["cu_seqlens_kv_padded"]), + ("do_t", p["d_o"], p["do_format"], attn.h, attn.sq, p["cu_seqlens_q_padded"]), + ): + attn.matrix(name, t, fmt, h, s, cu, columnwise=True) + scale_fmt = ( + do_scale_fmt if name == "do_t" else p.get("qkv_scale_inv_format") or fmt + ) + attn.scale("descale_" + name, _descale(t, True), name, scale_fmt, True) + else: + attn.scalar( + "descale_o", + ( + _descale(p["o"]) + if hasattr(p["o"], "_fp8_dtype") + else torch.ones(1, dtype=torch.float32, device=attn.device) + ), + ) + for short, quantizer in (("s", p["s_quantizer"]), ("dp", p["dp_quantizer"])): + attn.scalar("scale_" + short, quantizer.scale) + attn.scalar("descale_" + short, quantizer.scale.reciprocal()) + attn.scalar("amax_dp", p["dp_quantizer"].amax, output=True) + amax_grads = [] + for name in ("dq", "dk", "dv"): + quantizer = p["dqkv_quantizer"] + scale = ( + quantizer.scale + if isinstance(quantizer, Float8Quantizer) + else (torch.ones(1, dtype=torch.float32, device=attn.device)) + ) + attn.scalar("scale_" + name, scale) + amax = torch.empty(1, dtype=torch.float32, device=attn.device) + attn.scalar("amax_" + name, amax, output=True) + amax_grads.append(amax) + attn.execute() + if attn.fp8 and not attn.mxfp8 and isinstance(p["dqkv_quantizer"], Float8Quantizer): + p["dqkv_quantizer"].amax.copy_(torch.stack(amax_grads).amax().reshape(1)) + return dq, dk, dv, dbias, dsoftmax diff --git a/transformer_engine/pytorch/cpp_extensions/fused_attn.py b/transformer_engine/pytorch/cpp_extensions/fused_attn.py index 9a33df7634b..98628205782 100644 --- a/transformer_engine/pytorch/cpp_extensions/fused_attn.py +++ b/transformer_engine/pytorch/cpp_extensions/fused_attn.py @@ -8,7 +8,6 @@ from enum import IntEnum from typing import Tuple, List, Union, Optional import torch -import transformer_engine_torch as tex from transformer_engine_torch import ( NVTE_QKV_Layout, NVTE_QKV_Format, @@ -363,52 +362,53 @@ def fused_attn_fwd( else: raise ValueError(f"Unsupported backend {fused_attention_backend}") - # execute kernel - output_tensors = tex.fused_attn_fwd( - max_seqlen_q, - max_seqlen_kv, - is_training, - attn_scale, - dropout, - fast_zero_fill, - QKVLayout[qkv_layout], - QKVFormat[o_format], - QKVFormat[qkv_scale_inv_format], - AttnBiasType[attn_bias_type], - AttnMaskType[attn_mask_type], - SoftmaxType[softmax_type], - window_size, - bottom_right_diagonal, - cu_seqlens_q, - cu_seqlens_kv, - q, - k, - v, - fake_dtype, - cu_seqlens_q_padded, - cu_seqlens_kv_padded, - page_table_k, - page_table_v, - s_quantizer, - o_quantizer, - attn_bias, - softmax_offset, - rng_gen, - rng_elts_per_thread, - return_max_logit, - cuda_graph, + # Import lazily: the frontend's quantized tensor classes also import this module. + from . import cudnn_fused_attn + + out, aux_ctx_tensors, max_tensor = cudnn_fused_attn.fused_attn_fwd( + { + "is_training": is_training, + "max_seqlen_q": max_seqlen_q, + "max_seqlen_kv": max_seqlen_kv, + "cu_seqlens_q": cu_seqlens_q, + "cu_seqlens_kv": cu_seqlens_kv, + "q": q, + "k": k, + "v": v, + "fake_dtype": fake_dtype, + "fused_attention_backend": fused_attention_backend, + "attn_bias": attn_bias, + "cu_seqlens_q_padded": cu_seqlens_q_padded, + "cu_seqlens_kv_padded": cu_seqlens_kv_padded, + "page_table_k": page_table_k, + "page_table_v": page_table_v, + "s_quantizer": s_quantizer, + "o_quantizer": o_quantizer, + "attn_scale": attn_scale, + "dropout": dropout, + "fast_zero_fill": fast_zero_fill, + "qkv_layout": qkv_layout, + "o_format": o_format, + "qkv_scale_inv_format": qkv_scale_inv_format, + "attn_bias_type": attn_bias_type, + "attn_mask_type": attn_mask_type, + "softmax_type": softmax_type, + "window_size": window_size, + "bottom_right_diagonal": bottom_right_diagonal, + "rng_gen": rng_gen, + "softmax_offset": softmax_offset, + "return_max_logit": return_max_logit, + "cuda_graph": cuda_graph, + "rng_elts_per_thread": rng_elts_per_thread, + } ) if return_max_logit: qkv_format = qkv_layout.replace("3", "").replace("2", "").split("_")[0] - # thd (newer cuDNN runtimes, non-sm120): output_tensors: out [tq, h, d], Stats [tq, h, 1], Max [tq, h, 1] + # thd (newer cuDNN runtimes, non-sm120): out [tq, h, d], Stats [tq, h, 1], Max [tq, h, 1] # thd (older cuDNN runtimes or sm120): output_tensors: out [tq, h, d], Stats [b, h, sq, 1], Max [b, h, sq, 1] # bshd: output_tensors: out [b, sq, h, d], Stats [b, h, sq, 1], Max [b, h, sq, 1] # sbhd: output_tensors: out [sq, b, h, d], Stats [b, h, sq, 1], Max [b, h, sq, 1] - aux_ctx_tensors = [output_tensors[1]] + list( - output_tensors[3:] - ) # Stats + rng_state + optional tensors - max_tensor = output_tensors[2] amax_dims = (0, 2) if max_tensor.ndim == 3 else (0, 2, 3) if qkv_format == "thd": @@ -438,11 +438,11 @@ def fused_attn_fwd( max_tensor = max_tensor.masked_fill(~valid.view(-1, 1, 1), float("-inf")) # Max -> max_logit [h] - max_logit = torch.amax(max_tensor, dim=amax_dims).to(dtype=output_tensors[0].dtype) - return output_tensors[0], aux_ctx_tensors, max_logit + max_logit = torch.amax(max_tensor, dim=amax_dims).to(dtype=out.dtype) + return out, aux_ctx_tensors, max_logit # out, aux_ctx_tensors - return output_tensors[0], output_tensors[1:] + return out, aux_ctx_tensors def fused_attn_bwd( @@ -615,39 +615,44 @@ def fused_attn_bwd( f" for backend={fused_attention_backend}." ) - output_tensors = tex.fused_attn_bwd( - max_seqlen_q, - max_seqlen_kv, - attn_scale, - dropout, - fast_zero_fill, - QKVLayout[qkv_layout], - QKVFormat[o_format], - QKVFormat[do_format], - QKVLayout[dqkv_layout], - QKVFormat[qkv_scale_inv_format], - QKVFormat[do_scale_inv_format], - AttnBiasType[attn_bias_type], - AttnMaskType[attn_mask_type], - SoftmaxType[softmax_type], - window_size, - bottom_right_diagonal, - deterministic, - cu_seqlens_q, - cu_seqlens_kv, - q, - k, - v, - o, - d_o, - fake_dtype, - aux_ctx_tensors, - cu_seqlens_q_padded, - cu_seqlens_kv_padded, - s_quantizer, - dp_quantizer, - dqkv_quantizer, - cuda_graph, + from . import cudnn_fused_attn + + output_tensors = cudnn_fused_attn.fused_attn_bwd( + { + "max_seqlen_q": max_seqlen_q, + "max_seqlen_kv": max_seqlen_kv, + "cu_seqlens_q": cu_seqlens_q, + "cu_seqlens_kv": cu_seqlens_kv, + "q": q, + "k": k, + "v": v, + "o": o, + "d_o": d_o, + "fake_dtype": fake_dtype, + "aux_ctx_tensors": aux_ctx_tensors, + "fused_attention_backend": fused_attention_backend, + "cu_seqlens_q_padded": cu_seqlens_q_padded, + "cu_seqlens_kv_padded": cu_seqlens_kv_padded, + "s_quantizer": s_quantizer, + "dp_quantizer": dp_quantizer, + "dqkv_quantizer": dqkv_quantizer, + "attn_scale": attn_scale, + "dropout": dropout, + "fast_zero_fill": fast_zero_fill, + "qkv_layout": qkv_layout, + "o_format": o_format, + "do_format": do_format, + "dqkv_layout": dqkv_layout, + "qkv_scale_inv_format": qkv_scale_inv_format, + "do_scale_inv_format": do_scale_inv_format, + "attn_bias_type": attn_bias_type, + "attn_mask_type": attn_mask_type, + "softmax_type": softmax_type, + "window_size": window_size, + "bottom_right_diagonal": bottom_right_diagonal, + "deterministic": deterministic, + "cuda_graph": cuda_graph, + } ) return output_tensors diff --git a/transformer_engine/pytorch/csrc/extensions.h b/transformer_engine/pytorch/csrc/extensions.h index 955763f43b1..4144953acc4 100644 --- a/transformer_engine/pytorch/csrc/extensions.h +++ b/transformer_engine/pytorch/csrc/extensions.h @@ -94,36 +94,7 @@ std::tuple moe_unpermute_bwd(at::Tensor input_bwd, at::T * Attention **************************************************************************************************/ -std::tuple get_fused_attn_backend( - const py::object &fused_attn_params); - -std::vector fused_attn_fwd( - size_t max_seqlen_q, size_t max_seqlen_kv, bool is_training, float attn_scale, float p_dropout, - bool set_zero, NVTE_QKV_Layout qkv_layout, NVTE_QKV_Format o_format, - NVTE_QKV_Format qkv_scale_inv_format, NVTE_Bias_Type bias_type, NVTE_Mask_Type attn_mask_type, - NVTE_Softmax_Type softmax_type, const std::vector window_size, - bool bottom_right_diagonal, const at::Tensor cu_seqlens_q, const at::Tensor cu_seqlens_kv, - const py::handle Q, const py::handle K, const py::handle V, const at::ScalarType fake_dtype, - const std::optional cu_seqlens_q_padded, - const std::optional cu_seqlens_kv_padded, - const std::optional page_table_k, const std::optional page_table_v, - py::handle s_quantizer, py::handle o_quantizer, const std::optional Bias, - const std::optional SoftmaxOffset, const std::optional rng_gen, - size_t rng_elts_per_thread, bool return_max_logit, bool cuda_graph); - -std::vector fused_attn_bwd( - size_t max_seqlen_q, size_t max_seqlen_kv, float attn_scale, float p_dropout, bool set_zero, - NVTE_QKV_Layout qkv_layout, NVTE_QKV_Format o_format, NVTE_QKV_Format do_format, - NVTE_QKV_Layout dqkv_layout, NVTE_QKV_Format qkv_scale_inv_format, - NVTE_QKV_Format do_scale_inv_format, NVTE_Bias_Type bias_type, NVTE_Mask_Type attn_mask_type, - NVTE_Softmax_Type softmax_type, const std::vector window_size, - bool bottom_right_diagonal, bool deterministic, const at::Tensor cu_seqlens_q, - const at::Tensor cu_seqlens_kv, const py::handle Q, const py::handle K, const py::handle V, - const py::handle O, const py::handle dO, const at::ScalarType fake_dtype, - const std::vector Aux_CTX_Tensors, - const std::optional cu_seqlens_q_padded, - const std::optional cu_seqlens_kv_padded, py::handle s_quantizer, - py::handle dp_quantizer, py::handle dqkv_quantizer, bool cuda_graph); +at::Tensor fused_attn_rng_state(const std::optional rng_gen, size_t elts_per_thread); at::Tensor fa_prepare_fwd(at::Tensor qkvi); at::Tensor fa_prepare_bwd(at::Tensor q, at::Tensor k, at::Tensor v); diff --git a/transformer_engine/pytorch/csrc/extensions/attention.cpp b/transformer_engine/pytorch/csrc/extensions/attention.cpp index 98d55d2b8b5..547cde9ee39 100644 --- a/transformer_engine/pytorch/csrc/extensions/attention.cpp +++ b/transformer_engine/pytorch/csrc/extensions/attention.cpp @@ -10,640 +10,16 @@ namespace transformer_engine::pytorch { -// get the fused attention backend -std::tuple get_fused_attn_backend(const py::object &p) { - FusedAttnConfigWrapper cfg; - cfg.set_is_training(p.attr("is_training").cast()) - .set_deterministic(p.attr("deterministic").cast()) - .set_cuda_graph(p.attr("cuda_graph").cast()) - .set_return_max_logit(p.attr("return_max_logit").cast()) - .set_attn_mask_type(p.attr("attn_mask_type").cast()) - .set_bias_type(p.attr("bias_type").cast()) - .set_window_size_left(p.attr("window_size_left").cast()) - .set_window_size_right(p.attr("window_size_right").cast()) - .set_bottom_right_diagonal(p.attr("bottom_right_diagonal").cast()) - .set_softmax_type(p.attr("softmax_type").cast()) - .set_scaling_mode(p.attr("scaling_mode").cast()) - .set_dropout(p.attr("dropout").cast()) - .set_attn_scale(p.attr("attn_scale").cast()) - .set_qkv_dtype(static_cast(p.attr("qkv_dtype").cast())) - .set_o_dtype(static_cast(p.attr("o_dtype").cast())) - .set_do_dtype(static_cast(p.attr("do_dtype").cast())) - .set_dqkv_dtype(static_cast(p.attr("dqkv_dtype").cast())) - .set_qkv_layout(p.attr("qkv_layout").cast()) - .set_o_format(p.attr("o_format").cast()) - .set_do_format(p.attr("do_format").cast()) - .set_dqkv_layout(p.attr("dqkv_layout").cast()) - .set_qkv_scale_inv_format(p.attr("qkv_scale_inv_format").cast()) - .set_do_scale_inv_format(p.attr("do_scale_inv_format").cast()) - .set_batch_size(p.attr("batch_size").cast()) - .set_num_attn_heads(p.attr("num_attn_heads").cast()) - .set_num_gqa_groups(p.attr("num_gqa_groups").cast()) - .set_head_dim_qk(p.attr("head_dim_qk").cast()) - .set_head_dim_v(p.attr("head_dim_v").cast()) - .set_max_seqlen_q(p.attr("max_seqlen_q").cast()) - .set_max_seqlen_kv(p.attr("max_seqlen_kv").cast()) - .set_num_tokens_q(p.attr("num_tokens_q").cast()) - .set_num_tokens_kv(p.attr("num_tokens_kv").cast()) - .set_num_pages_k(p.attr("num_pages_k").cast()) - .set_num_pages_v(p.attr("num_pages_v").cast()) - .set_page_size_k(p.attr("page_size_k").cast()) - .set_page_size_v(p.attr("page_size_v").cast()) - .set_max_pages_per_seq_k(p.attr("max_pages_per_seq_k").cast()) - .set_max_pages_per_seq_v(p.attr("max_pages_per_seq_v").cast()) - .set_bias_batch_size(p.attr("bias_batch_size").cast()) - .set_bias_num_heads(p.attr("bias_num_heads").cast()) - .set_bias_seqlen_q(p.attr("bias_seqlen_q").cast()) - .set_bias_seqlen_kv(p.attr("bias_seqlen_kv").cast()); - - py::gil_scoped_release nogil; - const char *message = nullptr; - NVTE_Fused_Attn_Backend fused_attention_backend = nvte_get_fused_attn_backend_v2(cfg, &message); - return {fused_attention_backend, std::string(message)}; -} - -// helper function for S and dP quantizers -std::tuple> quantizer_helper( - py::handle quantizer, const std::vector &shape, DType dtype, bool create_hp_tensor, - std::optional data) { - std::unique_ptr T_quantizer = convert_quantizer(quantizer); - TensorWrapper te_T; - py::object py_T; - std::optional amax_buf; - if (quantizer.is_none()) { - // high precision - auto *none_quantizer = dynamic_cast(T_quantizer.get()); - if (data.has_value()) { - std::tie(te_T, py_T) = none_quantizer->create_tensor(shape, dtype, data.value()); - } else { - std::tie(te_T, py_T) = none_quantizer->create_tensor(shape, dtype); - } - } else if (detail::IsFloat8Quantizers(quantizer.ptr())) { - // delayed scaling; this helps initialize scale_inv - auto *T_quantizer_fp8 = dynamic_cast(T_quantizer.get()); - std::tie(te_T, py_T) = - T_quantizer_fp8->create_tensor(shape, dtype, data, std::nullopt, std::nullopt); - } else if (detail::IsFloat8CurrentScalingQuantizers(quantizer.ptr())) { - // current scaling - auto *T_quantizer_fp8 = dynamic_cast(T_quantizer.get()); - if (create_hp_tensor) { - if (data.has_value()) { - std::tie(te_T, py_T, amax_buf) = - T_quantizer_fp8->create_unquantized_tensor_with_amax(shape, dtype, data.value()); - } else { - std::tie(te_T, py_T, amax_buf) = - T_quantizer_fp8->create_unquantized_tensor_with_amax(shape, dtype); - } - } else { - std::tie(te_T, py_T) = T_quantizer_fp8->create_tensor(shape, dtype); - NVTE_CHECK( - !data.has_value(), - "Float8CurrentScalingQuantizer::create_tensor() does not take data tensor as input!"); - } - } else if (detail::IsMXFP8Quantizers(quantizer.ptr())) { - // MXFP8 - if (create_hp_tensor) { - if (data.has_value()) { - std::tie(te_T, py_T) = NoneQuantizer(py::none()).create_tensor(shape, dtype, data.value()); - } else { - std::tie(te_T, py_T) = NoneQuantizer(py::none()).create_tensor(shape, dtype); - } - } else { - auto *T_quantizer_fp8 = dynamic_cast(T_quantizer.get()); - std::tie(te_T, py_T) = T_quantizer_fp8->create_tensor(shape, dtype); - NVTE_CHECK(!data.has_value(), - "MXFP8Quantizer::create_tensor() does not take data tensor as input!"); - } - } - return {std::move(te_T), std::move(py_T), std::move(amax_buf)}; -} - -// fused attention FWD with separate Q, K and V tensors -std::vector fused_attn_fwd( - size_t max_seqlen_q, size_t max_seqlen_kv, bool is_training, float attn_scale, float p_dropout, - bool set_zero, NVTE_QKV_Layout qkv_layout, NVTE_QKV_Format o_format, - NVTE_QKV_Format qkv_scale_inv_format, NVTE_Bias_Type bias_type, NVTE_Mask_Type attn_mask_type, - NVTE_Softmax_Type softmax_type, const std::vector window_size, - bool bottom_right_diagonal, const at::Tensor cu_seqlens_q, const at::Tensor cu_seqlens_kv, - const py::handle Q, const py::handle K, const py::handle V, const at::ScalarType fake_dtype, - const std::optional cu_seqlens_q_padded, - const std::optional cu_seqlens_kv_padded, - const std::optional page_table_k, const std::optional page_table_v, - py::handle s_quantizer, py::handle o_quantizer, const std::optional Bias, - const std::optional SoftmaxOffset, const std::optional rng_gen, - size_t rng_elts_per_thread, bool return_max_logit, bool cuda_graph) { - // Ensure that cuDNN handle is created on the correct device, - // overriding torch.cuda.set_device calls from user side. - // Assumes all tensors passed are on the same device. - at::cuda::CUDAGuard device_guard(cu_seqlens_q.device()); - - auto none = py::none(); - - // create QKV tensor wrappers - TensorWrapper te_Q, te_K, te_V; - te_Q = makeTransformerEngineTensor(Q, none); - te_K = makeTransformerEngineTensor(K, none); - te_V = makeTransformerEngineTensor(V, none); - const DType qkv_type = te_Q.dtype(); - - // create S tensor - auto [te_S, py_S, _] = quantizer_helper(s_quantizer, {0}, DType::kFloat32, false, std::nullopt); - - // create O tensor - std::unique_ptr O_quantizer = convert_quantizer(o_quantizer); - std::vector q_shape = convertShape(te_Q.shape()); - std::vector v_shape = convertShape(te_V.shape()); - auto o_shape_tmp = std::vector{q_shape.begin(), q_shape.end()}; - o_shape_tmp[o_shape_tmp.size() - 1] = v_shape[v_shape.size() - 1]; - auto o_shape = std::vector{o_shape_tmp.begin(), o_shape_tmp.end()}; - NVTE_QKV_Format q_format = nvte_get_q_format(qkv_layout); - AttentionShape o_parsed(q_format, o_shape_tmp.data()); - o_parsed.to_format(o_format, o_shape.data()); - const DType fake_dtype_te = GetTransformerEngineDType(fake_dtype); - auto [te_O, py_O, o_amax_buf] = - quantizer_helper(o_quantizer, o_shape, fake_dtype_te, true, std::nullopt); - - // construct NVTE tensors - TensorWrapper te_Bias; - TensorWrapper te_cu_seqlens_q, te_cu_seqlens_kv; - TensorWrapper te_cu_seqlens_q_padded, te_cu_seqlens_kv_padded; - TensorWrapper te_page_table_k, te_page_table_v; - if (qkv_type == DType::kFloat8E4M3 || qkv_type == DType::kFloat8E5M2) { - // FP8 - if (set_zero && (o_format == NVTE_QKV_Format::NVTE_THD)) { - // Initialize both the output data and its amax metadata. - te_O.zero_(at::cuda::getCurrentCUDAStream()); - } - } else if (qkv_type == DType::kBFloat16 || qkv_type == DType::kFloat16) { - if (o_format == NVTE_QKV_Format::NVTE_THD) { - te_O.zero_(at::cuda::getCurrentCUDAStream()); - } - } else { - NVTE_ERROR("Fused attention only supports FP8 and BF16/FP16 data types. \n"); - } - if ((bias_type != NVTE_NO_BIAS) && (bias_type != NVTE_ALIBI) && (Bias.has_value())) { - auto bias_sizes = Bias.value().sizes().vec(); - std::vector bias_shape{bias_sizes.begin(), bias_sizes.end()}; - te_Bias = makeTransformerEngineTensor(Bias.value().data_ptr(), bias_shape, DType::kFloat32); - } - auto cu_seqlens_q_sizes = cu_seqlens_q.sizes().vec(); - std::vector cu_seqlens_q_shape{cu_seqlens_q_sizes.begin(), cu_seqlens_q_sizes.end()}; - auto cu_seqlens_kv_sizes = cu_seqlens_kv.sizes().vec(); - std::vector cu_seqlens_kv_shape{cu_seqlens_kv_sizes.begin(), cu_seqlens_kv_sizes.end()}; - te_cu_seqlens_q = - makeTransformerEngineTensor(cu_seqlens_q.data_ptr(), cu_seqlens_q_shape, DType::kInt32); - te_cu_seqlens_kv = - makeTransformerEngineTensor(cu_seqlens_kv.data_ptr(), cu_seqlens_kv_shape, DType::kInt32); - - if ((cu_seqlens_q_padded.has_value()) && (cu_seqlens_kv_padded.has_value())) { - auto cu_seqlens_q_padded_sizes = cu_seqlens_q_padded.value().sizes().vec(); - std::vector cu_seqlens_q_padded_shape{cu_seqlens_q_padded_sizes.begin(), - cu_seqlens_q_padded_sizes.end()}; - auto cu_seqlens_kv_padded_sizes = cu_seqlens_kv_padded.value().sizes().vec(); - std::vector cu_seqlens_kv_padded_shape{cu_seqlens_kv_padded_sizes.begin(), - cu_seqlens_kv_padded_sizes.end()}; - te_cu_seqlens_q_padded = makeTransformerEngineTensor(cu_seqlens_q_padded.value().data_ptr(), - cu_seqlens_q_padded_shape, DType::kInt32); - te_cu_seqlens_kv_padded = makeTransformerEngineTensor( - cu_seqlens_kv_padded.value().data_ptr(), cu_seqlens_kv_padded_shape, DType::kInt32); - } - - if ((page_table_k.has_value()) && (page_table_v.has_value())) { - auto page_table_k_sizes = page_table_k.value().sizes().vec(); - std::vector page_table_k_shape{page_table_k_sizes.begin(), page_table_k_sizes.end()}; - auto page_table_v_sizes = page_table_v.value().sizes().vec(); - std::vector page_table_v_shape{page_table_v_sizes.begin(), page_table_v_sizes.end()}; - te_page_table_k = - makeTransformerEngineTensor(page_table_k.value().data_ptr(), page_table_k_shape, - DType::kInt32, nullptr, nullptr, nullptr); - te_page_table_v = - makeTransformerEngineTensor(page_table_v.value().data_ptr(), page_table_v_shape, - DType::kInt32, nullptr, nullptr, nullptr); - } - - // softmax offset - TensorWrapper te_SoftmaxOffset; - if ((softmax_type != NVTE_VANILLA_SOFTMAX) && (SoftmaxOffset.has_value())) { - auto SoftmaxOffset_sizes = SoftmaxOffset.value().sizes().vec(); - std::vector SoftmaxOffset_shape{SoftmaxOffset_sizes.begin(), SoftmaxOffset_sizes.end()}; - te_SoftmaxOffset = - makeTransformerEngineTensor(SoftmaxOffset.value().data_ptr(), SoftmaxOffset_shape, - DType::kFloat32, nullptr, nullptr, nullptr); - } - - // extract rng seed and offset +// Keep dropout on the framework's CUDA generator, including graph replay offsets. +// This helper does not construct or execute an attention graph. +at::Tensor fused_attn_rng_state(const std::optional rng_gen, + size_t elts_per_thread) { auto gen = at::get_generator_or_default( rng_gen, at::cuda::detail::getDefaultCUDAGenerator()); - at::PhiloxCudaState philox_args = init_philox_state(gen, rng_elts_per_thread); - auto options = torch::TensorOptions().dtype(torch::kInt64).device(torch::kCUDA); - auto rng_state = torch::empty({2}, options); - philox_unpack(philox_args, static_cast(rng_state.data_ptr())); - auto te_rng_state = makeTransformerEngineTensor(rng_state); - - // create auxiliary output tensors - NVTETensorPack nvte_aux_tensor_pack; - nvte_tensor_pack_create(&nvte_aux_tensor_pack); - - // create workspace - TensorWrapper workspace; - - // build the parameter object - FusedAttnFwdParamsWrapper params; - params.set_Q(te_Q.data()) - .set_K(te_K.data()) - .set_V(te_V.data()) - .set_Bias(te_Bias.data()) - .set_SoftmaxOffset(te_SoftmaxOffset.data()) - .set_S(te_S.data()) - .set_O(te_O.data()) - .set_Aux_CTX_Tensors(&nvte_aux_tensor_pack) - .set_cu_seqlens_q(te_cu_seqlens_q.data()) - .set_cu_seqlens_kv(te_cu_seqlens_kv.data()) - .set_cu_seqlens_q_padded(te_cu_seqlens_q_padded.data()) - .set_cu_seqlens_kv_padded(te_cu_seqlens_kv_padded.data()) - .set_page_table_k(te_page_table_k.data()) - .set_page_table_v(te_page_table_v.data()) - .set_rng_state(te_rng_state.data()) - .set_is_training(is_training) - .set_cuda_graph(cuda_graph) - .set_return_max_logit(return_max_logit) - .set_attn_mask_type(attn_mask_type) - .set_bias_type(bias_type) - .set_window_size_left(window_size[0]) - .set_window_size_right(window_size[1]) - .set_bottom_right_diagonal(bottom_right_diagonal) - .set_softmax_type(softmax_type) - .set_dropout(p_dropout) - .set_attn_scale(attn_scale) - .set_qkv_layout(qkv_layout) - .set_o_format(o_format) - .set_qkv_scale_inv_format(qkv_scale_inv_format) - .set_max_seqlen_q(max_seqlen_q) - .set_max_seqlen_kv(max_seqlen_kv) - .set_workspace(workspace.data()) - .set_stream(at::cuda::getCurrentCUDAStream()); - - // populate tensors with appropriate shapes and dtypes - NVTE_SCOPED_GIL_RELEASE({ nvte_fused_attn_fwd_v2(params); }); - - // allocate memory for workspace and auxiliary output tensors - auto workspace_data = allocateSpace(workspace.shape(), workspace.dtype()); - workspace = - makeTransformerEngineTensor(workspace_data.data_ptr(), workspace.shape(), workspace.dtype()); - params.set_workspace(workspace.data()); - - // output_tensors = [O, nvte_aux_tensor_pack.tensors] - std::vector output_tensors; - output_tensors.push_back(py_O); - auto set_tensor_param = [&](size_t i, const at::Tensor &output_tensor) { - output_tensors.push_back(py::cast(output_tensor)); - NVTEBasicTensor temp_data = {output_tensor.data_ptr(), - nvte_tensor_type(nvte_aux_tensor_pack.tensors[i]), - nvte_tensor_shape(nvte_aux_tensor_pack.tensors[i])}; - nvte_set_tensor_param(&nvte_aux_tensor_pack.tensors[i], kNVTERowwiseData, &temp_data); - }; - // allocate memory for nvte_aux_tensor_pack.tensors - // f16_arbitrary: S [b, h, sq, 1]/[tq, h, 1], (optional) Max [b, h, sq, 1]/[tq, h, 1], rng_state [2], (optional) Bias [1, h, sq, skv], (optional) SoftmaxOffset [1, h, 1, 1] - // fp8 : S [b, h, sq, 1], rng_state [2] - size_t i = 0; - at::Tensor output_tensor; - // intermediate softmax stats tensor S - output_tensor = - allocateSpace(nvte_shape_to_vector(nvte_tensor_shape(nvte_aux_tensor_pack.tensors[i])), - static_cast(nvte_tensor_type(nvte_aux_tensor_pack.tensors[i])), false); - set_tensor_param(i++, output_tensor); - // return_max_logit=true allocates Max after S - if (return_max_logit) { - output_tensor = - allocateSpace(nvte_shape_to_vector(nvte_tensor_shape(nvte_aux_tensor_pack.tensors[i])), - static_cast(nvte_tensor_type(nvte_aux_tensor_pack.tensors[i])), false); - set_tensor_param(i++, output_tensor); - } - // rng_state - if (i < nvte_aux_tensor_pack.size) { - set_tensor_param(i++, rng_state); - } - // bias (optional) - if ((bias_type != NVTE_NO_BIAS) && (bias_type != NVTE_ALIBI) && (Bias.has_value())) { - set_tensor_param(i++, Bias.value()); - } - // softmax_offset (optional) - if ((softmax_type != NVTE_VANILLA_SOFTMAX) && (SoftmaxOffset.has_value())) { - set_tensor_param(i++, SoftmaxOffset.value()); - } - - // execute the kernel - NVTE_SCOPED_GIL_RELEASE({ nvte_fused_attn_fwd_v2(params); }); - - // destroy tensor wrappers, but not allocated memory - nvte_tensor_pack_destroy(&nvte_aux_tensor_pack); - - // if training, [O, softmax-related tensors, rng_state]; if inference, [O] - return output_tensors; -} - -// fused attention BWD with separate Q, K and V -std::vector fused_attn_bwd( - size_t max_seqlen_q, size_t max_seqlen_kv, float attn_scale, float p_dropout, bool set_zero, - NVTE_QKV_Layout qkv_layout, NVTE_QKV_Format o_format, NVTE_QKV_Format do_format, - NVTE_QKV_Layout dqkv_layout, NVTE_QKV_Format qkv_scale_inv_format, - NVTE_QKV_Format do_scale_inv_format, NVTE_Bias_Type bias_type, NVTE_Mask_Type attn_mask_type, - NVTE_Softmax_Type softmax_type, const std::vector window_size, - bool bottom_right_diagonal, bool deterministic, const at::Tensor cu_seqlens_q, - const at::Tensor cu_seqlens_kv, const py::handle Q, const py::handle K, const py::handle V, - const py::handle O, const py::handle dO, const at::ScalarType fake_dtype, - const std::vector Aux_CTX_Tensors, - const std::optional cu_seqlens_q_padded, - const std::optional cu_seqlens_kv_padded, py::handle s_quantizer, - py::handle dp_quantizer, py::handle dqkv_quantizer, bool cuda_graph) { - auto none = py::none(); - - // create QKV, O, dO tensor wrappers - TensorWrapper te_Q, te_K, te_V, te_O, te_dO; - te_Q = makeTransformerEngineTensor(Q, none); - te_K = makeTransformerEngineTensor(K, none); - te_V = makeTransformerEngineTensor(V, none); - te_O = makeTransformerEngineTensor(O, none); - te_dO = makeTransformerEngineTensor(dO, none); - - // create S and dP tensors - auto [te_S, py_S, _s] = quantizer_helper(s_quantizer, {0}, DType::kFloat32, false, std::nullopt); - auto [te_dP, py_dP, _dp] = - quantizer_helper(dp_quantizer, {0}, DType::kFloat32, false, std::nullopt); - - // create dQ, dK, dV tensors - TensorWrapper te_dQ, te_dK, te_dV; - py::object py_dQ, py_dK, py_dV; - std::optional dq_amax_buf, dk_amax_buf, dv_amax_buf; - std::unique_ptr dQKV_quantizer = convert_quantizer(dqkv_quantizer); - std::vector q_shape = convertShape(te_Q.shape()); - std::vector k_shape = convertShape(te_K.shape()); - std::vector v_shape = convertShape(te_V.shape()); - const DType dqkv_fake_dtype = GetTransformerEngineDType(fake_dtype); - size_t ndim_q = q_shape.size(); - size_t ndim_kv = k_shape.size(); - std::vector dQ_shape(ndim_q), dK_shape(ndim_kv), dV_shape(ndim_kv); - NVTE_QKV_Format q_format = nvte_get_q_format(qkv_layout); - NVTE_QKV_Format kv_format = nvte_get_kv_format(qkv_layout); - NVTE_QKV_Format dq_format = nvte_get_q_format(dqkv_layout); - NVTE_QKV_Format dkv_format = nvte_get_kv_format(dqkv_layout); - AttentionShape q_parsed(q_format, q_shape.data()); - size_t h_q = q_parsed.h(); - q_parsed.to_format(dq_format, dQ_shape.data()); - AttentionShape k_parsed(kv_format, k_shape.data()); - k_parsed.to_format(dkv_format, dK_shape.data()); - AttentionShape v_parsed(kv_format, v_shape.data()); - v_parsed.to_format(dkv_format, dV_shape.data()); - at::Tensor dQ, dK, dV, dQKV, dKV; - // FP16/BF16: dqkv_fake_dtype = kFloat16/kBFloat16, dQ/dK/dV.dtype = torch.float16/torch.bfloat16 - // FP8DS: dqkv_fake_dtype = kFloat16/kBFloat16, dQ/dK/dV.dtype = torch.uint8 - // FP8CS/MXFP8: dqkv_fake_dtype = kFloat16/kBFloat16, dQ/dK/dV.dtype = torch.float16/torch.bfloat16 - auto options = torch::TensorOptions().dtype(fake_dtype).device(torch::kCUDA); - if (detail::IsFloat8Quantizers(dqkv_quantizer.ptr())) { - options = options.dtype(torch::kUInt8); - } - - NVTE_QKV_Layout_Group layout_group = nvte_get_qkv_layout_group(dqkv_layout); - std::vector tmp_shape; - switch (layout_group) { - case NVTE_QKV_Layout_Group::NVTE_3HD: - tmp_shape = std::vector{dQ_shape.begin(), dQ_shape.end()}; - tmp_shape.insert(tmp_shape.begin() + tmp_shape.size() - 2, int64_t(3)); - dQKV = torch::empty(c10::IntArrayRef(tmp_shape), options); - dQ = dQKV.index({"...", torch::indexing::Slice(0, 1, 1), - torch::indexing::Slice(0, torch::indexing::None, 1), - torch::indexing::Slice(0, torch::indexing::None, 1)}) - .squeeze(tmp_shape.size() - 3); - dK = dQKV.index({"...", torch::indexing::Slice(1, 2, 1), - torch::indexing::Slice(0, torch::indexing::None, 1), - torch::indexing::Slice(0, torch::indexing::None, 1)}) - .squeeze(tmp_shape.size() - 3); - dV = dQKV.index({"...", torch::indexing::Slice(2, torch::indexing::None, 1), - torch::indexing::Slice(0, torch::indexing::None, 1), - torch::indexing::Slice(0, torch::indexing::None, 1)}) - .squeeze(tmp_shape.size() - 3); - break; - case NVTE_QKV_Layout_Group::NVTE_H3D: - tmp_shape = std::vector{dQ_shape.begin(), dQ_shape.end()}; - tmp_shape.insert(tmp_shape.begin() + tmp_shape.size() - 1, int64_t(3)); - dQKV = torch::empty(c10::IntArrayRef(tmp_shape), options); - dQ = dQKV.index({"...", torch::indexing::Slice(0, 1, 1), - torch::indexing::Slice(0, torch::indexing::None, 1)}) - .squeeze(tmp_shape.size() - 2); - dK = dQKV.index({"...", torch::indexing::Slice(1, 2, 1), - torch::indexing::Slice(0, torch::indexing::None, 1)}) - .squeeze(tmp_shape.size() - 2); - dV = dQKV.index({"...", torch::indexing::Slice(2, torch::indexing::None, 1), - torch::indexing::Slice(0, torch::indexing::None, 1)}) - .squeeze(tmp_shape.size() - 2); - break; - case NVTE_QKV_Layout_Group::NVTE_HD_2HD: - tmp_shape = std::vector(dQ_shape.begin(), dQ_shape.end()); - dQ = torch::empty(tmp_shape, options); - tmp_shape = std::vector{dK_shape.begin(), dK_shape.end()}; - tmp_shape.insert(tmp_shape.begin() + tmp_shape.size() - 2, int64_t(2)); - dKV = torch::empty(c10::IntArrayRef(tmp_shape), options); - dK = dKV.index({"...", torch::indexing::Slice(0, 1, 1), - torch::indexing::Slice(0, torch::indexing::None, 1), - torch::indexing::Slice(0, torch::indexing::None, 1)}) - .squeeze(tmp_shape.size() - 3); - dV = dKV.index({"...", torch::indexing::Slice(1, torch::indexing::None, 1), - torch::indexing::Slice(0, torch::indexing::None, 1), - torch::indexing::Slice(0, torch::indexing::None, 1)}) - .squeeze(tmp_shape.size() - 3); - break; - case NVTE_QKV_Layout_Group::NVTE_HD_H2D: - tmp_shape = std::vector(dQ_shape.begin(), dQ_shape.end()); - dQ = torch::empty(tmp_shape, options); - tmp_shape = std::vector{dK_shape.begin(), dK_shape.end()}; - tmp_shape.insert(tmp_shape.begin() + tmp_shape.size() - 1, int64_t(2)); - dKV = torch::empty(c10::IntArrayRef(tmp_shape), options); - dK = dKV.index({"...", torch::indexing::Slice(0, 1, 1), - torch::indexing::Slice(0, torch::indexing::None, 1)}) - .squeeze(tmp_shape.size() - 2); - dV = dKV.index({"...", torch::indexing::Slice(1, torch::indexing::None, 1), - torch::indexing::Slice(0, torch::indexing::None, 1)}) - .squeeze(tmp_shape.size() - 2); - break; - case NVTE_QKV_Layout_Group::NVTE_HD_HD_HD: - case NVTE_QKV_Layout_Group::NVTE_SD_SD_SD: - tmp_shape = std::vector(dQ_shape.begin(), dQ_shape.end()); - dQ = torch::empty(tmp_shape, options); - tmp_shape = std::vector(dK_shape.begin(), dK_shape.end()); - dK = torch::empty(tmp_shape, options); - tmp_shape = std::vector(dV_shape.begin(), dV_shape.end()); - dV = torch::empty(tmp_shape, options); - break; - default: - NVTE_ERROR("QKV layout not supported!"); - } - - std::tie(te_dQ, py_dQ, dq_amax_buf) = - quantizer_helper(dqkv_quantizer, dQ_shape, dqkv_fake_dtype, true, dQ); - std::tie(te_dK, py_dK, dk_amax_buf) = - quantizer_helper(dqkv_quantizer, dK_shape, dqkv_fake_dtype, true, dK); - std::tie(te_dV, py_dV, dv_amax_buf) = - quantizer_helper(dqkv_quantizer, dV_shape, dqkv_fake_dtype, true, dV); - - // construct NVTE tensors - if (detail::IsFloat8Quantizers(dqkv_quantizer.ptr())) { - // FP8 - if (set_zero) { - // dQ/dK/dV may be strided views, so zero data through ATen and reset the shared amax - // separately. - if (dq_format == NVTE_QKV_Format::NVTE_THD) { - dQ.fill_(0); - } - if (dkv_format == NVTE_QKV_Format::NVTE_THD) { - dK.fill_(0); - dV.fill_(0); - } - if (dq_format == NVTE_QKV_Format::NVTE_THD || dkv_format == NVTE_QKV_Format::NVTE_THD) { - auto *fp8_quantizer = dynamic_cast(dQKV_quantizer.get()); - fp8_quantizer->amax.zero_(); - } - } - } else if (dqkv_quantizer.is_none() || - detail::IsFloat8CurrentScalingQuantizers(dqkv_quantizer.ptr()) || - detail::IsMXFP8Quantizers(dqkv_quantizer.ptr())) { - if (dq_format == NVTE_QKV_Format::NVTE_THD) { - dQ.fill_(0); - } - if (dkv_format == NVTE_QKV_Format::NVTE_THD) { - dK.fill_(0); - dV.fill_(0); - } - } else { - NVTE_ERROR("Fused attention only supports FP8 and BF16/FP16 data types. \n"); - } - - // create cu_seqlens tensorwrappers - auto cu_seqlens_q_sizes = cu_seqlens_q.sizes().vec(); - std::vector cu_seqlens_q_shape{cu_seqlens_q_sizes.begin(), cu_seqlens_q_sizes.end()}; - auto cu_seqlens_kv_sizes = cu_seqlens_kv.sizes().vec(); - std::vector cu_seqlens_kv_shape{cu_seqlens_kv_sizes.begin(), cu_seqlens_kv_sizes.end()}; - TensorWrapper te_cu_seqlens_q, te_cu_seqlens_kv; - te_cu_seqlens_q = makeTransformerEngineTensor(cu_seqlens_q.data_ptr(), cu_seqlens_q_shape, - DType::kInt32, nullptr, nullptr, nullptr); - te_cu_seqlens_kv = makeTransformerEngineTensor(cu_seqlens_kv.data_ptr(), cu_seqlens_kv_shape, - DType::kInt32, nullptr, nullptr, nullptr); - - TensorWrapper te_cu_seqlens_q_padded, te_cu_seqlens_kv_padded; - if ((cu_seqlens_q_padded.has_value()) && (cu_seqlens_kv_padded.has_value())) { - auto cu_seqlens_q_padded_sizes = cu_seqlens_q_padded.value().sizes().vec(); - std::vector cu_seqlens_q_padded_shape{cu_seqlens_q_padded_sizes.begin(), - cu_seqlens_q_padded_sizes.end()}; - auto cu_seqlens_kv_padded_sizes = cu_seqlens_kv_padded.value().sizes().vec(); - std::vector cu_seqlens_kv_padded_shape{cu_seqlens_kv_padded_sizes.begin(), - cu_seqlens_kv_padded_sizes.end()}; - te_cu_seqlens_q_padded = makeTransformerEngineTensor(cu_seqlens_q_padded.value().data_ptr(), - cu_seqlens_q_padded_shape, DType::kInt32); - te_cu_seqlens_kv_padded = makeTransformerEngineTensor( - cu_seqlens_kv_padded.value().data_ptr(), cu_seqlens_kv_padded_shape, DType::kInt32); - } - - // convert auxiliary tensors from forward to NVTETensors - NVTETensorPack nvte_aux_tensor_pack; - nvte_tensor_pack_create(&nvte_aux_tensor_pack); - nvte_aux_tensor_pack.size = Aux_CTX_Tensors.size(); - for (size_t i = 0; i < nvte_aux_tensor_pack.size; ++i) { - const std::vector &signed_shape = Aux_CTX_Tensors[i].sizes().vec(); - const std::vector tmp(signed_shape.begin(), signed_shape.end()); - - NVTEBasicTensor temp_data = { - Aux_CTX_Tensors[i].data_ptr(), - static_cast(GetTransformerEngineDType(Aux_CTX_Tensors[i].scalar_type())), - nvte_make_shape(tmp.data(), tmp.size())}; - nvte_set_tensor_param(&nvte_aux_tensor_pack.tensors[i], kNVTERowwiseData, &temp_data); - } - - // create dBias the same shape as Bias - at::Tensor dBias; - TensorWrapper te_dBias; - if ((bias_type != NVTE_NO_BIAS) && (bias_type != NVTE_ALIBI)) { - if (nvte_aux_tensor_pack.size >= 2) { - std::vector bias_shape(Aux_CTX_Tensors[nvte_aux_tensor_pack.size - 1].sizes().vec()); - dBias = torch::empty(bias_shape, options); - te_dBias = makeTransformerEngineTensor(dBias); - } else { - dBias = torch::empty({1, static_cast(h_q), static_cast(max_seqlen_q), - static_cast(max_seqlen_kv)}, - options); - te_dBias = makeTransformerEngineTensor(dBias); - } - if (nvte_get_qkv_format(qkv_layout) == NVTE_QKV_Format::NVTE_THD) { - dBias.fill_(0); - } - } - - // create dSoftmaxOffset in the same shape as SoftmaxOffset - at::Tensor dSoftmaxOffset; - TensorWrapper te_dSoftmaxOffset; - if (softmax_type != NVTE_VANILLA_SOFTMAX) { - options = torch::TensorOptions().dtype(at::kFloat).device(torch::kCUDA); - dSoftmaxOffset = torch::empty({1, static_cast(h_q), 1, 1}, options); - te_dSoftmaxOffset = makeTransformerEngineTensor(dSoftmaxOffset); - } - - // create workspace - TensorWrapper workspace; - - // build the parameter object - FusedAttnBwdParamsWrapper params; - params.set_Q(te_Q.data()) - .set_K(te_K.data()) - .set_V(te_V.data()) - .set_O(te_O.data()) - .set_dO(te_dO.data()) - .set_S(te_S.data()) - .set_dP(te_dP.data()) - .set_Aux_CTX_Tensors(&nvte_aux_tensor_pack) - .set_dQ(te_dQ.data()) - .set_dK(te_dK.data()) - .set_dV(te_dV.data()) - .set_dBias(te_dBias.data()) - .set_dSoftmaxOffset(te_dSoftmaxOffset.data()) - .set_cu_seqlens_q(te_cu_seqlens_q.data()) - .set_cu_seqlens_kv(te_cu_seqlens_kv.data()) - .set_cu_seqlens_q_padded(te_cu_seqlens_q_padded.data()) - .set_cu_seqlens_kv_padded(te_cu_seqlens_kv_padded.data()) - .set_deterministic(deterministic) - .set_cuda_graph(cuda_graph) - .set_attn_mask_type(attn_mask_type) - .set_bias_type(bias_type) - .set_window_size_left(window_size[0]) - .set_window_size_right(window_size[1]) - .set_bottom_right_diagonal(bottom_right_diagonal) - .set_softmax_type(softmax_type) - .set_dropout(p_dropout) - .set_attn_scale(attn_scale) - .set_qkv_layout(qkv_layout) - .set_o_format(o_format) - .set_do_format(do_format) - .set_dqkv_layout(dqkv_layout) - .set_qkv_scale_inv_format(qkv_scale_inv_format) - .set_do_scale_inv_format(do_scale_inv_format) - .set_max_seqlen_q(max_seqlen_q) - .set_max_seqlen_kv(max_seqlen_kv) - .set_workspace(workspace.data()) - .set_stream(at::cuda::getCurrentCUDAStream()); - - // populate tensors with appropriate shapes and dtypes - NVTE_SCOPED_GIL_RELEASE({ nvte_fused_attn_bwd_v2(params); }); - - // allocate memory for workspace - auto workspace_data = allocateSpace(workspace.shape(), workspace.dtype()); - workspace = - makeTransformerEngineTensor(workspace_data.data_ptr(), workspace.shape(), workspace.dtype()); - params.set_workspace(workspace.data()); - - // execute kernel - NVTE_SCOPED_GIL_RELEASE({ nvte_fused_attn_bwd_v2(params); }); - - // destroy tensor wrappers - nvte_tensor_pack_destroy(&nvte_aux_tensor_pack); - - return {py_dQ, py_dK, py_dV, py::cast(dBias), py::cast(dSoftmaxOffset)}; + auto state = init_philox_state(gen, elts_per_thread); + auto result = at::empty({2}, at::TensorOptions().dtype(at::kLong).device(at::kCUDA)); + philox_unpack(state, result.data_ptr()); + return result; } at::Tensor fa_prepare_fwd(at::Tensor qkvi) { diff --git a/transformer_engine/pytorch/csrc/extensions/pybind.cpp b/transformer_engine/pytorch/csrc/extensions/pybind.cpp index 2046408eba8..538d3cbaaf7 100644 --- a/transformer_engine/pytorch/csrc/extensions/pybind.cpp +++ b/transformer_engine/pytorch/csrc/extensions/pybind.cpp @@ -475,8 +475,6 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { m.def("swap_first_dims", &transformer_engine::pytorch::swap_first_dims, "Swap first two tensor dimensions", py::arg("tensor"), py::kw_only(), py::arg("out"), py::call_guard()); - m.def("get_fused_attn_backend", &transformer_engine::pytorch::get_fused_attn_backend, - "Get Fused Attention backend", py::arg("fused_attn_params")); m.def("compute_amax", &transformer_engine::pytorch::compute_amax, "Compute absolute max value in tensor", py::arg("input"), py::arg("amax"), py::call_guard()); @@ -569,10 +567,9 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { m.def("multi_tensor_pad_last_dim", &transformer_engine::pytorch::multi_tensor_pad_last_dim, "Pad multiple tensors' last dimension to a common alignment.", py::arg("inputs"), py::arg("alignment"), py::call_guard()); - m.def("fused_attn_fwd", &transformer_engine::pytorch::fused_attn_fwd, - "Fused Attention FP8/BF16/FP16 FWD with separate Q, K and V"); - m.def("fused_attn_bwd", &transformer_engine::pytorch::fused_attn_bwd, - "Fused Attention FP8/BF16/FP16 BWD with separate Q, K and V"); + m.def("fused_attn_rng_state", &transformer_engine::pytorch::fused_attn_rng_state, + "Reserve and unpack the CUDA generator's Philox state", py::arg("generator"), + py::arg("elts_per_thread")); m.def("copy_to_kv_cache", &transformer_engine::pytorch::copy_to_kv_cache, "Copy new KV tokens to KV cache", py::call_guard()); m.def("convert_thd_to_bshd", &transformer_engine::pytorch::convert_thd_to_bshd, From 78565fcc8fe0387e09f6099f8cf4a7eb1bdb4830 Mon Sep 17 00:00:00 2001 From: Vladimir Cherepanov Date: Sun, 11 Oct 2026 23:09:57 +0000 Subject: [PATCH 3/6] Share cuDNN Python attention graph builders across frontends Share tensor descriptors, SDPA construction, graph lifecycle stages, diagnostics, and score-modification cache utilities between PyTorch and JAX. Keep allocation, quantization, RNG, and execution in framework adapters. Separate support checks from compilation, fix JAX anonymous callback cache keys, preserve empty ragged capacities, and align frontend requirements and documentation. Validation: editable PyTorch build succeeded. Attention and Flex suites reported 2,877 passed and 2,426 skipped; both checkpoint-loading failures passed on retry with NVTE_ALLOW_UNSAFE_PICKLE_EXTRA_STATE=1. Pre-commit, Python lint, license checks, and JAX descriptor checks passed. JAX build and runtime validation remain for a separate session. Signed-off-by: Vladimir Cherepanov --- build_tools/jax.py | 2 +- build_tools/wheel_utils/build_wheels.sh | 2 +- docs/envvars.rst | 2 +- docs/installation.rst | 14 +- tests/jax/test_fused_attn.py | 75 +++ tests/jax/test_fused_attn_score_mod.py | 13 +- tests/pytorch/attention/test_attention.py | 106 ++++ transformer_engine/common/cudnn/__init__.py | 5 + transformer_engine/common/cudnn/attention.py | 219 +++++++ transformer_engine/common/cudnn/graph.py | 249 ++++++++ transformer_engine/common/cudnn/score_mod.py | 96 +++ .../jax/cpp_extensions/attention.py | 16 +- .../jax/cpp_extensions/cudnn_attention.py | 466 ++++++++------- .../jax/cpp_extensions/cudnn_graph.py | 30 +- .../jax/cpp_extensions/flex_attention.py | 355 ++++-------- .../{attention.cpp => cudnn_graph.cpp} | 2 +- .../dot_product_attention/flex_attention.py | 261 +++------ .../pytorch/cpp_extensions/cudnn.py | 34 +- .../cpp_extensions/cudnn_fused_attn.py | 547 +++++------------- 19 files changed, 1390 insertions(+), 1104 deletions(-) create mode 100644 transformer_engine/common/cudnn/__init__.py create mode 100644 transformer_engine/common/cudnn/attention.py create mode 100644 transformer_engine/common/cudnn/graph.py create mode 100644 transformer_engine/common/cudnn/score_mod.py rename transformer_engine/jax/csrc/extensions/{attention.cpp => cudnn_graph.cpp} (99%) diff --git a/build_tools/jax.py b/build_tools/jax.py index 891fb63eec8..9de778721ce 100644 --- a/build_tools/jax.py +++ b/build_tools/jax.py @@ -27,7 +27,7 @@ def install_requirements() -> List[str]: """Install dependencies for TE/JAX extensions.""" - return ["jax", "flax>=0.7.1", "nvidia-cudnn-frontend>=1.25.0", "triton"] + return ["jax", "flax>=0.7.1", "nvidia-cudnn-frontend>=1.29.0", "triton"] def test_requirements() -> List[str]: diff --git a/build_tools/wheel_utils/build_wheels.sh b/build_tools/wheel_utils/build_wheels.sh index 761f12390e1..0e49c2c2e7a 100644 --- a/build_tools/wheel_utils/build_wheels.sh +++ b/build_tools/wheel_utils/build_wheels.sh @@ -24,7 +24,7 @@ git checkout $TARGET_BRANCH git submodule update --init --recursive # Install deps -/opt/python/cp310-cp310/bin/pip install cmake pybind11[global] ninja setuptools wheel 'nvidia-cudnn-frontend>=1.25.0' 'apache-tvm-ffi>=0.1.12' +/opt/python/cp310-cp310/bin/pip install cmake pybind11[global] ninja setuptools wheel 'nvidia-cudnn-frontend>=1.29.0' 'apache-tvm-ffi>=0.1.12' if $BUILD_METAPACKAGE ; then cd /TransformerEngine diff --git a/docs/envvars.rst b/docs/envvars.rst index 511a59c7917..b2d9ffce1c5 100644 --- a/docs/envvars.rst +++ b/docs/envvars.rst @@ -229,7 +229,7 @@ backend-selection overview. :Type: ``int`` (0, 1 or 2), optionally followed by ``:`` :Default: ``0`` - :Description: Log FusedAttention graph cache activity to stderr, prefixed with ``[FUSED-ATTN-CACHE]``. ``1`` prints an end-of-run summary of the cache counters and the mean time of each cuDNN build stage. ``2`` additionally traces every event as it happens: each graph built, each graph cuDNN accepts and the cache keeps, each lookup and whether it hit or missed, each first execution that compiles kernels, and each execution. When the launcher exports a rank, only rank 0 logs; append ``:`` to override, as in ``1:all`` for level 1 on every rank or ``2:0,3`` for level 2 on ranks 0 and 3. + :Description: Log FusedAttention graph cache activity to stderr, prefixed with ``[FUSED-ATTN-CACHE]``. ``1`` prints an end-of-run summary of the cache counters and the mean time of each cuDNN build stage. ``2`` additionally traces every event as it happens: each graph built, each graph cuDNN accepts and the cache keeps, each lookup and whether it hit or missed, each plan compilation, and each execution. JAX reports Python graph-cache and build events during tracing; execution events apply to the PyTorch and common C++ paths. When the launcher exports a rank, only rank 0 logs; append ``:`` to override, as in ``1:all`` for level 1 on every rank or ``2:0,3`` for level 2 on ranks 0 and 3. .. envvar:: NVTE_ALLOW_NONDETERMINISTIC_ALGO diff --git a/docs/installation.rst b/docs/installation.rst index cb9ddccc16f..9da51c18bd2 100644 --- a/docs/installation.rst +++ b/docs/installation.rst @@ -16,10 +16,16 @@ Prerequisites 3. |driver link|_ supporting CUDA 12.1 or later. 4. `cuDNN 9.12 `__ or later. -PyTorch's cuDNN attention backend uses the cuDNN frontend Python API. The PyTorch -installation includes ``nvidia-cudnn-frontend>=1.29.0``. Attention graphs and plans -are cached by device and tensor metadata; dropout uses PyTorch's CUDA generator, -including CUDA graph capture and replay. +The PyTorch and JAX cuDNN attention backends use shared graph builders through +the cuDNN frontend Python API. Both installations include +``nvidia-cudnn-frontend>=1.29.0``. Graph descriptions are cached by tensor metadata +and target device or architecture. Support queries select supported plans without +compiling their kernels. PyTorch compiles on first execution and uses its CUDA +generator for dropout, including CUDA graph capture and replay. JAX compiles +during tracing when workspace sizes are needed, then serializes plans for its +native FFI executor. The JAX extension and the frontend Python package must use +matching cuDNN frontend versions; rebuild the extension after upgrading that +package. If the CUDA Toolkit headers are not available at runtime in a standard installation path, e.g. within `CUDA_HOME`, set diff --git a/tests/jax/test_fused_attn.py b/tests/jax/test_fused_attn.py index d8d6c088195..6da031c353f 100644 --- a/tests/jax/test_fused_attn.py +++ b/tests/jax/test_fused_attn.py @@ -428,6 +428,81 @@ def test_fused_attn_backend_message(): assert message != "" +@pytest.mark.parametrize("setup_error", (None, RuntimeError("frontend version mismatch"))) +def test_fused_attn_support_query_defers_compilation(setup_error, monkeypatch): + """Support queries must avoid compilation and propagate unexpected setup failures.""" + from transformer_engine.jax.cpp_extensions import cudnn_attention + + def forbidden(*_args, **_kwargs): + raise AssertionError("Support query compiled an execution plan") + + monkeypatch.setattr(cudnn_attention, "compile_graph", forbidden) + cudnn_attention._supported_graph.cache_clear() + if setup_error is not None: + + def fail(*_args, **_kwargs): + raise setup_error + + monkeypatch.setattr(cudnn_attention, "check_attention_support", fail) + helper = FusedAttnHelper( + is_training=True, + batch_size=2, + q_dtype=jnp.bfloat16, + kv_dtype=jnp.bfloat16, + qkv_layout=QKVLayout.BSHD_BSHD_BSHD, + attn_bias_type=AttnBiasType.NO_BIAS, + attn_mask_type=AttnMaskType.NO_MASK, + softmax_type=AttnSoftmaxType.VANILLA_SOFTMAX, + dropout_probability=0.0, + q_num_heads=4, + kv_num_heads=4, + q_max_seqlen=128, + kv_max_seqlen=128, + head_dim_qk=64, + head_dim_v=64, + window_size=(-1, -1), + attn_scale=0.125, + ) + if setup_error is not None: + with pytest.raises(RuntimeError, match=str(setup_error)): + helper.get_fused_attn_backend() + else: + backend, message = helper.get_fused_attn_backend() + if backend == NVTE_Fused_Attn_Backend.NVTE_No_Backend: + pytest.skip(f"Baseline attention is unsupported: {message}") + + +@pytest.mark.parametrize("batch,seqlen", ((0, 128), (2, 0))) +def test_cudnn_attention_empty_ragged_capacities(batch, seqlen, monkeypatch): + """Empty storage must retain positive graph capacities for ragged metadata.""" + from types import SimpleNamespace + from transformer_engine.jax.cpp_extensions import cudnn_attention + + monkeypatch.setattr(cudnn_attention, "cudnn_version", lambda: 92700) + config = SimpleNamespace( + qkv_layout=QKVLayout.THD_THD_THD, + max_segments_per_seq=1, + fp8=False, + attn_mask_type=AttnMaskType.PADDING_MASK, + attn_bias_type=AttnBiasType.NO_BIAS, + softmax_type=AttnSoftmaxType.VANILLA_SOFTMAX, + effective_window_size=(-1, -1), + scaling_factor=0.125, + bottom_right_diagonal=False, + dropout_probability=0.0, + is_training=False, + return_max_logit=False, + ) + shapes = tuple( + (shape, jnp.dtype(jnp.bfloat16)) + for shape in ((batch, seqlen, 4, 64),) * 3 + ((0,), (1, 4, 1, 1)) + ) + description = cudnn_attention._describe_graph("fwd", shapes, config, (100,), True, True) + assert description.metadata_batch == 32 + for name in ("q", "k", "v"): + assert dict(description.spec.tensors)[name].dim == (32, 4, 1024, 64) + + class BiasShape(Enum): """ Enum class to represent the different bias shapes used in the fused attention. diff --git a/tests/jax/test_fused_attn_score_mod.py b/tests/jax/test_fused_attn_score_mod.py index 5d9b1508589..4be2c2599a7 100644 --- a/tests/jax/test_fused_attn_score_mod.py +++ b/tests/jax/test_fused_attn_score_mod.py @@ -2,6 +2,7 @@ # # See LICENSE for license information. """Tests for cuDNN frontend score_mod fused attention.""" + from math import sqrt import jax @@ -22,7 +23,6 @@ from transformer_engine_jax import get_device_compute_capability, NVTE_Fused_Attn_Backend from test_fused_attn import FusedAttnRunner, SeqDescFormat - _CONFIG_TEST_HEAD_DIM = 128 _CONFIG_TEST_SCALING_FACTOR = 1.0 / sqrt(_CONFIG_TEST_HEAD_DIM) _SCORE_MOD_MIN_CUDNN_VERSION = (9, 23) @@ -675,6 +675,17 @@ class FakeCudnn: tex_attention._check_cudnn_frontend_version_match(FakeCudnn) +def test_fused_attn_score_mod_module_lambda_keys_do_not_collide(): + """Different anonymous functions cannot share a serialized attention graph.""" + first = lambda graph, score, tensors: score # noqa: E731 + second = lambda graph, score, tensors: graph.tanh(score) # noqa: E731 + first.__qualname__ = second.__qualname__ = "" + assert first.__qualname__ == second.__qualname__ == "" + assert tex_attention._score_mod_callback_cache_key(first) != ( + tex_attention._score_mod_callback_cache_key(second) + ) + + def test_fused_attn_score_mod_config_stabilizes_bound_method_cache_keys(): """Bound score_mod methods with graph keys share stable cache keys.""" softcap_score_mod = _ScoreModSoftcap() diff --git a/tests/pytorch/attention/test_attention.py b/tests/pytorch/attention/test_attention.py index 47350b7a6bb..a3b6a96c8c1 100644 --- a/tests/pytorch/attention/test_attention.py +++ b/tests/pytorch/attention/test_attention.py @@ -251,6 +251,112 @@ def common_attention(*_args, **_kwargs): _attention_backends["backend_selection_requires_update"] = True +@pytest.mark.parametrize("layout", ("bs3hd", "bsh3d", "sb3hd", "sbh3d", "t3hd", "th3d")) +def test_fused_attn_query_reuses_runtime_graph(layout, monkeypatch): + """Packed storage and runtime scalar values must preserve support-query graph identity.""" + from transformer_engine.pytorch.cpp_extensions import cudnn_fused_attn as implementation + + batch, seqlen, heads, dim = 2, 128, 4, 64 + fmt = implementation._formats(layout)[0] + base_shape = implementation._shape((batch, heads, seqlen, dim), fmt, batch * seqlen) + axis = -2 if "h3d" in layout else -3 + packed_shape = list(base_shape) + packed_shape.insert(len(packed_shape) + axis + 1, 3) + q, k, v = torch.empty(packed_shape, device="cuda", dtype=torch.bfloat16).unbind(axis) + cu = torch.arange(batch + 1, device="cuda", dtype=torch.int32) * seqlen + params = dict( + q=q, + k=k, + v=v, + cu_seqlens_q=cu, + cu_seqlens_kv=cu, + cu_seqlens_q_padded=cu, + cu_seqlens_kv_padded=cu, + qkv_layout=layout, + max_seqlen_q=seqlen, + max_seqlen_kv=seqlen, + attn_mask_type="padding" if fmt == "thd" else "no_mask", + attn_bias_type="no_bias", + softmax_type="vanilla", + dropout=0.0, + attn_scale=0.125, + window_size=(-1, -1), + bottom_right_diagonal=True, + ) + config = FusedAttentionParams( + qkv_layout=QKVLayout[layout], + o_format=QKVFormat[fmt], + batch_size=batch, + num_attn_heads=heads, + num_gqa_groups=heads, + head_dim_qk=dim, + head_dim_v=dim, + max_seqlen_q=seqlen, + max_seqlen_kv=seqlen, + num_tokens_q=batch * seqlen if fmt == "thd" else 0, + num_tokens_kv=batch * seqlen if fmt == "thd" else 0, + attn_mask_type=AttnMaskType[params["attn_mask_type"]], + ) + query = implementation._query_graph(config, False) + runtime = implementation._AttentionGraph(params) + implementation._prepare_inputs(runtime) + out = torch.empty(base_shape, device="cuda", dtype=torch.bfloat16) + runtime.matrix("o", out, fmt, heads, runtime.sq, cu, output=True) + stats_shape = (batch * seqlen, heads, 1) if runtime.ragged_stats else (batch, heads, seqlen, 1) + runtime.stats("stats", torch.empty(stats_shape, device="cuda"), output=True) + assert query.description() == runtime.description() + + # Isolate the cache to demonstrate that a query never compiles, and that a + # later execution with another scalar value compiles exactly the cached plan. + monkeypatch.setattr(implementation, "_plans", {}) + builds, compiles = [], [] + + def build(_self): + entry = implementation._PlanEntry(object(), {name: name for name in runtime.specs}) + builds.append(entry) + return entry + + def compile_plan(_cudnn, graph, _run): + compiles.append(graph) + return 17 + + monkeypatch.setattr(implementation._AttentionGraph, "build", build) + monkeypatch.setattr(implementation, "compile_graph", compile_plan) + monkeypatch.setattr(implementation, "_execute_cudnn_graph", lambda *_args: None) + queried = query.entry() + assert queried.workspace_size is None and not compiles + runtime.scalar("attn_scale", torch.tensor([0.25]), by_value=True) + runtime.execute() + runtime.execute() + assert builds == [queried] + assert compiles == [queried.graph] + assert queried.workspace_size == 17 + + +@pytest.mark.parametrize("error", (RuntimeError("device failure"), ImportError("frontend missing"))) +def test_fused_attn_query_propagates_setup_errors(error, monkeypatch): + """A failed setup must surface as an error instead of selecting a fallback backend.""" + from transformer_engine.pytorch.cpp_extensions import cudnn_fused_attn as implementation + + def fail(_self): + raise error + + monkeypatch.setattr(implementation._AttentionGraph, "build", fail) + monkeypatch.setattr(implementation, "_plans", {}) + config = FusedAttentionParams( + qkv_layout=QKVLayout["bshd_bshd_bshd"], + batch_size=2, + num_attn_heads=4, + num_gqa_groups=4, + head_dim_qk=64, + head_dim_v=64, + max_seqlen_q=128, + max_seqlen_kv=128, + ) + with pytest.raises(type(error), match=str(error)): + implementation.get_fused_attn_backend(config) + + @pytest.mark.parametrize( "changes,version,capability,reason", ( diff --git a/transformer_engine/common/cudnn/__init__.py b/transformer_engine/common/cudnn/__init__.py new file mode 100644 index 00000000000..7501e1ed649 --- /dev/null +++ b/transformer_engine/common/cudnn/__init__.py @@ -0,0 +1,5 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""Framework-independent helpers for the cuDNN Python graph API.""" diff --git a/transformer_engine/common/cudnn/attention.py b/transformer_engine/common/cudnn/attention.py new file mode 100644 index 00000000000..e5f872667e7 --- /dev/null +++ b/transformer_engine/common/cudnn/attention.py @@ -0,0 +1,219 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""Construct SDPA graphs from immutable descriptors, without framework tensors.""" + +from dataclasses import dataclass + +from .graph import TensorSpec, apply_tensor_spec, create_tensor + + +@dataclass(frozen=True) +class AttentionOptions: + """Normalized attributes that affect graph construction and support.""" + + backward: bool = False + fp8: bool = False + mxfp8: bool = False + attn_scale: str | float = "attn_scale" + bottom_right: bool = False + left_bound: int | None = None + right_bound: int | None = None + padding: bool = False + alibi: bool = False + dropout: float = 0.0 + deterministic: bool = False + generate_stats: bool = True + paged_max_seqlen_kv: int = 0 + max_total_seqlen_q: int = 0 + max_total_seqlen_kv: int = 0 + + +@dataclass(frozen=True) +class AttentionGraphSpec: + """A cacheable graph description; bindings and scalar values live in adapters.""" + + options: AttentionOptions + tensors: tuple + + @classmethod + def from_dict(cls, options, tensors): + """Canonicalize tensor order for both UIDs and graph cache keys.""" + return cls(options, tuple(sorted(tensors.items()))) + + +@dataclass +class BuiltAttentionGraph: + """A live graph and its tensor handles, without any runtime buffers.""" + + graph: object + tensors: dict + uids: dict + + +def build_attention_graph(cudnn, graph, spec, *, score_mod=None, score_mod_bprop=None): + """Build ordinary, FP8, MXFP8, or score-modified SDPA using logical bindings.""" + specs = dict(spec.tensors) + tensors = {} + uids = {name: index + 1 for index, (name, _) in enumerate(spec.tensors)} + + def tensor(name): + if name not in tensors: + desc = create_tensor(cudnn, graph, name, specs[name], uids[name]) + if specs[name].ragged_offset: + desc.set_ragged_offset(tensor(specs[name].ragged_offset)) + tensors[name] = desc + return tensors[name] + + for name, desc in spec.tensors: + if not desc.output: + tensor(name) + + options = spec.options + kwargs = { + "attn_scale": ( + tensor(options.attn_scale) + if isinstance(options.attn_scale, str) + else options.attn_scale + ), + "diagonal_alignment": ( + cudnn.diagonal_alignment.BOTTOM_RIGHT + if options.bottom_right + else cudnn.diagonal_alignment.TOP_LEFT + ), + } + bound_names = ( + ("left_bound", "right_bound") + if options.fp8 and (not options.mxfp8 or options.backward) + else ("diagonal_band_left_bound", "diagonal_band_right_bound") + ) + for name, bound in zip(bound_names, (options.left_bound, options.right_bound)): + if bound is not None: + kwargs[name] = bound + if options.padding: + kwargs.update(use_padding_mask=True, seq_len_q=tensor("seq_q"), seq_len_kv=tensor("seq_kv")) + if options.alibi: + kwargs["use_alibi_mask"] = True + if "bias" in specs: + kwargs["bias"] = tensor("bias") + if options.backward and "dbias" in specs: + kwargs["dBias"] = tensor("dbias") + if options.dropout: + kwargs["dropout"] = (options.dropout, tensor("seed"), tensor("offset")) + if "sink" in specs: + kwargs["sink_token"] = tensor("sink") + if options.backward: + kwargs["dSink_token"] = tensor("dsink") + if score_mod is not None: + kwargs["score_mod"] = score_mod + if score_mod_bprop is not None: + kwargs["score_mod_bprop"] = score_mod_bprop + if options.backward: + kwargs["use_deterministic_algorithm"] = options.deterministic + if options.max_total_seqlen_q: + kwargs["max_total_seq_len_q"] = options.max_total_seqlen_q + kwargs["max_total_seq_len_kv"] = options.max_total_seqlen_kv + else: + kwargs["generate_stats"] = options.generate_stats + if "max" in specs: + kwargs["score_max"] = tensor("max") + if options.paged_max_seqlen_kv: + kwargs.update( + paged_attention_k_table=tensor("page_table_k"), + paged_attention_v_table=tensor("page_table_v"), + paged_attention_max_seq_len_kv=options.paged_max_seqlen_kv, + ) + + q, k, v = (tensor(name) for name in ("q", "k", "v")) + if not options.fp8: + if options.backward: + outputs = graph.sdpa_backward( + q, k, v, tensor("o"), tensor("do"), tensor("stats"), **kwargs + ) + names = ("dq", "dk", "dv") + else: + outputs = graph.sdpa(q, k, v, **kwargs) + names = ("o", "stats") + elif not options.backward: + scales = [tensor(name) for name in ("descale_q", "descale_k", "descale_v")] + if options.mxfp8: + outputs = graph.sdpa_mxfp8(q, k, v, *scales, **kwargs) + names = ("o", "stats", "amax_o") + else: + outputs = graph.sdpa_fp8( + q, + k, + v, + *scales, + *(tensor(name) for name in ("descale_s", "scale_s", "scale_o")), + **kwargs + ) + names = ("o", "stats", "amax_s", "amax_o") + elif options.mxfp8: + outputs = graph.sdpa_mxfp8_backward( + q, + tensor("q_t"), + k, + tensor("k_t"), + v, + tensor("o"), + tensor("do_f16"), + tensor("do"), + tensor("do_t"), + tensor("stats"), + *( + tensor(name) + for name in ( + "descale_q", + "descale_q_t", + "descale_k", + "descale_k_t", + "descale_v", + "descale_do", + "descale_do_t", + ) + ), + **kwargs + ) + names = ("dq", "dk", "dv", "amax_dq", "amax_dk", "amax_dv") + else: + outputs = graph.sdpa_fp8_backward( + q, + k, + v, + tensor("o"), + tensor("do"), + tensor("stats"), + *( + tensor(name) + for name in ( + "descale_q", + "descale_k", + "descale_v", + "descale_o", + "descale_do", + "descale_s", + "descale_dp", + "scale_s", + "scale_dq", + "scale_dk", + "scale_dv", + "scale_dp", + ) + ), + **kwargs + ) + names = ("dq", "dk", "dv", "amax_dq", "amax_dk", "amax_dv", "amax_dp") + + for name, desc in zip(names, outputs): + if name not in specs: + # Unbound MXFP8 amax outputs stay virtual. Their descriptors are still required. + if desc is not None: + apply_tensor_spec(cudnn, desc, TensorSpec((1, 1, 1, 1), (1, 1, 1, 1), "FLOAT")) + continue + apply_tensor_spec(cudnn, desc, specs[name], uids[name]) + if specs[name].ragged_offset: + desc.set_ragged_offset(tensor(specs[name].ragged_offset)) + tensors[name] = desc + return BuiltAttentionGraph(graph, tensors, uids) diff --git a/transformer_engine/common/cudnn/graph.py b/transformer_engine/common/cudnn/graph.py new file mode 100644 index 00000000000..a60ea826f06 --- /dev/null +++ b/transformer_engine/common/cudnn/graph.py @@ -0,0 +1,249 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""Data-free tensor descriptions and explicit cuDNN graph build stages.""" + +import atexit +from collections import Counter +from dataclasses import dataclass +import importlib +import math +import os +import sys +import threading +import time + + +class UnsupportedGraph(ValueError): + """A configuration rejected by cuDNN, rather than an execution/setup failure.""" + + +def import_cudnn(): + """Import the optional graph frontend without importing either framework.""" + try: + return importlib.import_module("cudnn") + except ImportError as exc: + raise ImportError( + "cuDNN frontend Python package not found. " + "Install it with: pip install nvidia-cudnn-frontend" + ) from exc + + +def create_graph(cudnn, dtype, handle=None): + """Create a graph with FP32 intermediate values and accumulation.""" + kwargs = {} if handle is None else {"handle": handle} + return cudnn.pygraph( + io_data_type=getattr(cudnn.data_type, dtype) if isinstance(dtype, str) else dtype, + intermediate_data_type=cudnn.data_type.FLOAT, + compute_data_type=cudnn.data_type.FLOAT, + **kwargs, + ) + + +def check_graph_support(cudnn, graph, run_stage=None): + """Validate and select supported plans, without compiling their kernels.""" + run = run_stage or (lambda _name, function: function()) + try: + run("validate", graph.validate) + run("build_operation_graph", graph.build_operation_graph) + run( + "create_execution_plans", + lambda: graph.create_execution_plans([cudnn.heur_mode.A, cudnn.heur_mode.FALLBACK]), + ) + run("check_support", graph.check_support) + except cudnn.cudnnGraphNotSupportedError as exc: + raise UnsupportedGraph(f"cuDNN SDPA graph is not supported: {exc}") from exc + + +def compile_graph(cudnn, graph, run_stage=None): + """Compile the selected plan and return its nonempty workspace size.""" + run = run_stage or (lambda _name, function: function()) + run("build_plans", lambda: graph.build_plans(cudnn.build_plan_policy.HEURISTICS_CHOICE)) + return max(int(graph.get_workspace_size()), 1) + + +@dataclass(frozen=True) +class TensorSpec: + """A logical graph tensor, independent of its runtime buffer and framework.""" + + dim: tuple + stride: tuple + dtype: str + output: bool = False + ragged_offset: str = "" + reordered: bool = False + by_value: bool = False + + +def apply_tensor_spec(cudnn, tensor, spec, uid=None): + """Apply the same descriptor to an input or an operation-produced output.""" + tensor.set_dim(spec.dim).set_stride(spec.stride).set_data_type( + getattr(cudnn.data_type, spec.dtype) + ) + if uid is not None: + tensor.set_uid(uid) + if spec.output: + tensor.set_output(True) + return tensor + + +def create_tensor(cudnn, graph, name, spec, uid=None): + """Create an input or an explicitly supplied output with a stable UID.""" + kwargs = {} if uid is None else {"uid": uid} + tensor = graph.tensor( + name=name, + dim=spec.dim, + stride=spec.stride, + data_type=getattr(cudnn.data_type, spec.dtype), + is_pass_by_value=spec.by_value, + **kwargs, + ) + if spec.reordered: + tensor.set_reordering_type(cudnn.tensor_reordering.F8_128x4) + return apply_tensor_spec(cudnn, tensor, spec, uid) + + +def row_major_stride(shape): + """Return contiguous storage strides without consulting framework objects.""" + return tuple(math.prod(shape[index + 1 :]) for index in range(len(shape))) + + +def round_up(value, multiple): + """Round an integer up to a multiple.""" + return (value + multiple - 1) // multiple * multiple + + +def bucket(value, minimum, maximum): + """Bucket positive capacities by powers of two, then multiples of maximum.""" + if value == 0: + return 0 + if value <= maximum: + return max(minimum, 1 << (value - 1).bit_length()) + return round_up(value, maximum) + + +def attention_stride(dim, fmt): + """Describe contiguous attention storage in logical BHSD order.""" + b, h, s, d = dim + if fmt in ("bshd", "thd"): + return (s * h * d, d, h * d, 1) + if fmt == "sbhd": + return (h * d, d, b * h * d, 1) + if fmt == "bhsd": + return (h * s * d, s * d, d, 1) + raise ValueError(f"Unsupported attention format: {fmt}") + + +def attention_matrix_spec( + dim, fmt, dtype, *, storage_stride=None, packed="", output=False, ragged_offset="" +): + """Describe either actual storage strides or a hypothetical support-query tensor.""" + if storage_stride is not None and fmt == "thd": + token, head, last = storage_stride + stride = (dim[2] * token, head, token, last) + elif storage_stride is not None and len(storage_stride) == 4: + order = {"bshd": (0, 2, 1, 3), "sbhd": (1, 2, 0, 3), "bhsd": (0, 1, 2, 3)}[fmt] + stride = tuple(storage_stride[index] for index in order) + else: + stride = list(attention_stride(dim, fmt)) + if "3" in packed or "2" in packed: + count = 3 if "3" in packed else 2 + stride[0] *= count + stride[2] *= count + if "h3d" in packed or "h2d" in packed: + stride[1] *= count + stride = tuple(stride) + return TensorSpec(tuple(dim), stride, dtype, output=output, ragged_offset=ragged_offset) + + +def attention_stats_spec(batch, heads, seqlen, *, ragged=False, output=False): + """Describe dense BHS1 or packed TH1 softmax statistics.""" + stride = (heads * seqlen, 1, heads, 1) if ragged else (heads * seqlen, seqlen, 1, 1) + return TensorSpec( + (batch, heads, seqlen, 1), + stride, + "FLOAT", + output=output, + ragged_offset="stats_offset" if ragged else "", + ) + + +def mxfp8_scale_spec(dim, fmt, columnwise=False): + """Describe the reordered block-scale storage used by MXFP8 SDPA.""" + b, h, s, d = dim + shape = ( + (b, h, round_up((s + 31) // 32, 4), round_up(d, 128)) + if columnwise + else (b, h, round_up(s, 128), round_up((d + 31) // 32, 4)) + ) + return TensorSpec(shape, attention_stride(shape, fmt), "FP8_E8M0", reordered=True) + + +class CacheDiagnostics: + """Optional cache counters and build timings shared by Python frontends.""" + + def __init__(self, framework): + self.framework = framework + self.events = Counter() + self.times = Counter() + atexit.register(self.summary) + + @staticmethod + def level(): + """Resolve the existing debug level and rank filter.""" + setting = os.getenv("NVTE_FUSED_ATTN_CACHE_DEBUG", "0") + level, _, ranks = setting.partition(":") + rank = next( + ( + os.environ[n] + for n in ("RANK", "LOCAL_RANK", "OMPI_COMM_WORLD_RANK", "SLURM_PROCID") + if n in os.environ + ), + None, + ) + enabled = rank is None or ranks == "all" or rank in (ranks or "0").split(",") + return int(level or 0) if enabled else 0 + + def event(self, backend, direction, device, event, key=None): + """Record an event without retaining runtime tensors.""" + level = self.level() + if not level: + return + self.events[(backend, direction, event)] += 1 + if level >= 2: + print( + f"[FUSED-ATTN-CACHE] tid={threading.get_ident()} dev={device} | " + f"{backend} {direction} {event}" + + (f" | {key!r}" if key is not None else ""), + file=sys.stderr, + flush=True, + ) + + def run(self, backend, direction, stage, function): + """Time a build stage only when diagnostics are enabled.""" + if not self.level(): + return function() + start = time.perf_counter() + try: + return function() + finally: + self.times[(backend, direction, stage, "seconds")] += time.perf_counter() - start + self.times[(backend, direction, stage, "calls")] += 1 + + def summary(self): + """Print this frontend's collected counters and mean build times.""" + if not self.events: + return + print(f"[FUSED-ATTN-CACHE] {self.framework} Python frontend summary", file=sys.stderr) + for (backend, direction, event), count in sorted(self.events.items()): + print(f" {backend} {direction} {event}={count}", file=sys.stderr) + for backend, direction, stage, kind in sorted(self.times): + if kind == "calls": + calls = self.times[(backend, direction, stage, kind)] + seconds = self.times[(backend, direction, stage, "seconds")] + print( + f" {backend} {direction} {stage}: {calls} calls, " + f"{1000 * seconds / calls:.3f} ms average", + file=sys.stderr, + ) diff --git a/transformer_engine/common/cudnn/score_mod.py b/transformer_engine/common/cudnn/score_mod.py new file mode 100644 index 00000000000..f385a569d6a --- /dev/null +++ b/transformer_engine/common/cudnn/score_mod.py @@ -0,0 +1,96 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""Structural callback keys shared by score-modified attention frontends.""" + +from collections.abc import Mapping +import inspect + + +class UncacheableKey: + """Unique static identity for callbacks whose topology may change.""" + + +def freeze_cache_key(value, is_tensor): + """Freeze nested structural keys while rejecting framework runtime tensors.""" + if is_tensor(value): + raise TypeError( + "score_mod_graph_cache_key() must not include tensors. Pass runtime tensors " + "through score_mod_tensors or score_mod_bprop_tensors instead." + ) + if isinstance(value, Mapping): + return tuple( + sorted( + ( + (freeze_cache_key(key, is_tensor), freeze_cache_key(val, is_tensor)) + for key, val in value.items() + ), + key=repr, + ) + ) + if isinstance(value, (list, tuple)): + return tuple(freeze_cache_key(item, is_tensor) for item in value) + if isinstance(value, (set, frozenset)): + return tuple(sorted((freeze_cache_key(item, is_tensor) for item in value), key=repr)) + try: + hash(value) + except TypeError as exc: + raise TypeError( + "score_mod_graph_cache_key() must return a hashable value or a nested " + "combination of mapping/list/tuple/set values." + ) from exc + return value + + +def explicit_cache_key(owner, is_tensor): + """Read an optional structural key from a callback's owner.""" + key = getattr(owner, "score_mod_graph_cache_key", None) + if key is None: + return None + return freeze_cache_key(key() if callable(key) else key, is_tensor) + + +def callback_cache_key(callback, is_tensor, uncacheable): + """Cache stable functions and explicitly keyed owners; distinguish anonymous code.""" + if callback is None: + return None + owner = getattr(callback, "__self__", None) + function = getattr(callback, "__func__", None) + if owner is not None and function is not None: + key = explicit_cache_key(owner, is_tensor) + if key is None: + return uncacheable() + return ("bound_method", type(owner), function.__module__, function.__qualname__, key) + key = explicit_cache_key(callback, is_tensor) + if key is not None: + return ( + "callable", + type(callback), + getattr(callback, "__module__", None), + getattr(callback, "__qualname__", None), + key, + ) + if ( + inspect.isfunction(callback) + and callback.__closure__ is None + and "" not in callback.__qualname__ + ): + identity = ( + callback.__code__ + if callback.__name__ == "" or not callback.__qualname__ + else callback.__qualname__ + ) + return ("function", callback.__module__, identity) + return uncacheable() + + +def wrap_score_mod(callback, graph_tensors): + """Adapt TE's three-argument callback to cuDNN's two-argument callback.""" + if callback is None: + return None + + def wrapped(graph, scores): + return callback(graph, scores, graph_tensors) + + return wrapped diff --git a/transformer_engine/jax/cpp_extensions/attention.py b/transformer_engine/jax/cpp_extensions/attention.py index 702b9457017..b5f9d37fe23 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 @@ -49,7 +50,6 @@ with_sharding_constraint, ) - __all__ = [ "FusedAttnHelper", "fused_attn_fwd", @@ -244,15 +244,15 @@ def get_fused_attn_backend(self): ) avals += (jax.ShapeDtypeStruct((1, self.q_num_heads, 1, 1), jnp.float32),) try: - cudnn_attention.get_graph("fwd", avals, config) + cudnn_attention.check_attention_support("fwd", avals, config) if self.is_training: - cudnn_attention.get_graph("bwd", avals, config) + cudnn_attention.check_attention_support("bwd", avals, config) backend = ( NVTE_Fused_Attn_Backend.NVTE_FP8 if self.fp8 else NVTE_Fused_Attn_Backend.NVTE_F16_arbitrary_seqlen ) - except (ValueError, RuntimeError, ImportError) as exc: + except ValueError as exc: message = str(exc) AttentionLogging.setup_logging() @@ -2892,7 +2892,7 @@ def correction(output, softmax_aux, output_per_step, softmax_aux_per_step): else: for i in range(0, cp_size): carry = scan_kv_block(i, carry) - (kv, output, softmax_aux, max_logit) = carry + kv, output, softmax_aux, max_logit = carry output = output.astype(q.dtype) # Globalize the rank-local running [H] max across DP/CP. @@ -3117,7 +3117,7 @@ def jax_cond_wrap(): else: for i in range(0, cp_size): carry = scan_kv_block(i, carry) - (kv, dq, dk_dv, dbias) = carry + kv, dq, dk_dv, dbias = carry # Final permute to put gradients back to their final resting place. dk_dv = helper.permute_kv(dk_dv, cp_perm) @@ -3356,7 +3356,7 @@ def correction(output, softmax_aux, output_per_step, softmax_aux_per_step): else: for i in range(0, cp_size): carry = scan_kv_block(i, carry) - (_, _, _, output, softmax_aux, max_logit) = carry + _, _, _, output, softmax_aux, max_logit = carry output = output.astype(q.dtype) # Globalize the rank-local running [H] max across DP/CP. @@ -3488,7 +3488,7 @@ def compute(config): else: for idx in range(cp_size): carry = scan_kv_block(idx, carry) - (_, _, _, dq, dkv, dbias) = carry + _, _, _, dq, dkv, dbias = carry # Final permute to put gradients back to their final resting place. dkv = helper.permute_kv(dkv, cp_perm) diff --git a/transformer_engine/jax/cpp_extensions/cudnn_attention.py b/transformer_engine/jax/cpp_extensions/cudnn_attention.py index 1701d7c0d68..ed770b28554 100644 --- a/transformer_engine/jax/cpp_extensions/cudnn_attention.py +++ b/transformer_engine/jax/cpp_extensions/cudnn_attention.py @@ -5,6 +5,7 @@ import math import os +import threading from dataclasses import dataclass from functools import lru_cache @@ -13,7 +14,23 @@ import numpy as np from jax import ffi -from .cudnn_graph import import_cudnn, finalize_graph, serialize_graph, graph_attributes +from transformer_engine.common.cudnn.attention import ( + AttentionOptions, + AttentionGraphSpec, + BuiltAttentionGraph, + build_attention_graph, +) +from transformer_engine.common.cudnn.graph import ( + TensorSpec, + CacheDiagnostics, + bucket as _bucket, + row_major_stride as _stride, + create_graph, + check_graph_support, + compile_graph, +) + +from .cudnn_graph import import_cudnn, check_frontend_version, serialize_graph, graph_attributes def device_capabilities(): @@ -26,26 +43,16 @@ def cudnn_version(): return int(import_cudnn().backend_version()) -def _dtype(cudnn, dtype): +def _dtype(dtype): return { - jnp.dtype(jnp.float16): cudnn.data_type.HALF, - jnp.dtype(jnp.bfloat16): cudnn.data_type.BFLOAT16, - jnp.dtype(jnp.float32): cudnn.data_type.FLOAT, - jnp.dtype(jnp.float8_e4m3fn): cudnn.data_type.FP8_E4M3, - jnp.dtype(jnp.float8_e5m2): cudnn.data_type.FP8_E5M2, + jnp.dtype(jnp.float16): "HALF", + jnp.dtype(jnp.bfloat16): "BFLOAT16", + jnp.dtype(jnp.float32): "FLOAT", + jnp.dtype(jnp.float8_e4m3fn): "FP8_E4M3", + jnp.dtype(jnp.float8_e5m2): "FP8_E5M2", }[jnp.dtype(dtype)] -def _bucket(value, minimum, maximum): - if value <= maximum: - return max(minimum, 1 << (value - 1).bit_length()) - return (value + maximum - 1) // maximum * maximum - - -def _stride(shape): - return tuple(math.prod(shape[index + 1 :]) for index in range(len(shape))) - - def _dims(avals, config): q, k, v = avals[:3] layout = config.qkv_layout @@ -77,10 +84,58 @@ class AttentionGraph: fp8: bool -def get_graph(direction, avals, config): - """Cache by configuration, operand metadata and GPU architecture.""" +@dataclass(frozen=True) +class _AttentionDescription: + """Static graph metadata and JAX's packed-storage buffer bindings.""" + + spec: AttentionGraphSpec + inputs: tuple + outputs: tuple + scalars: tuple + metadata_batch: int + ragged_stats: bool + multipliers: tuple + fp8: bool + + +@dataclass +class _SupportedGraph: + """A supported live graph and its lazily compiled serialization.""" + + built: BuiltAttentionGraph + workspace_size: int | None = None + serialized_graph: bytes | None = None + + +_diagnostics = CacheDiagnostics("JAX") +_plan_lock = threading.RLock() + + +def _site(spec): + return "fp8" if spec.options.fp8 else "f16", "bwd" if spec.options.backward else "fwd" + + +@lru_cache(maxsize=256) +def _supported_graph(spec, capabilities, frontend_version, backend_version): + # Versions and target architecture are part of the plan identity, even when + # two targets happen to accept identical logical tensor descriptors. + del frontend_version, backend_version + cudnn = import_cudnn() + backend, direction = _site(spec) + _diagnostics.event(backend, direction, capabilities, "MISS", spec) + graph = create_graph(cudnn, dict(spec.tensors)["o"].dtype) + built = build_attention_graph(cudnn, graph, spec) + _diagnostics.event(backend, direction, capabilities, "CREATE_GRAPH") + check_graph_support( + cudnn, graph, lambda stage, function: _diagnostics.run(backend, direction, stage, function) + ) + _diagnostics.event(backend, direction, capabilities, "CACHE_GRAPH") + return _SupportedGraph(built) + + +def _request(direction, avals, config): shapes = tuple((tuple(aval.shape), jnp.dtype(aval.dtype)) for aval in avals[:5]) - return _build_graph( + return ( direction, shapes, config, @@ -90,20 +145,96 @@ def get_graph(direction, avals, config): ) +def _get_supported_graph(spec, capabilities, frontend_version, backend_version): + """Share support-query graphs with compilation, and report cache reuse.""" + # pylint infers the wrapped function's arguments for lru_cache.cache_info. + # pylint: disable=no-value-for-parameter + with _plan_lock: + hits = _supported_graph.cache_info().hits + entry = _supported_graph(spec, capabilities, frontend_version, backend_version) + if _supported_graph.cache_info().hits > hits: + _diagnostics.event(*_site(spec), capabilities, "HIT") + return entry + + +def check_attention_support(direction, avals, config): + """Check the common graph without compiling kernels or serializing a plan.""" + description = _describe_graph(*_request(direction, avals, config)) + cudnn = import_cudnn() + _get_supported_graph( + description.spec, + device_capabilities(), + check_frontend_version(cudnn), + int(cudnn.backend_version()), + ) + + +def get_graph(direction, avals, config): + """Compile when tracing requires workspace sizes and serialized FFI attributes.""" + return _build_graph(*_request(direction, avals, config)) + + @lru_cache(maxsize=256) def _build_graph(direction, shapes, config, capabilities, nondeterministic, fp8_backward): + description = _describe_graph( + direction, shapes, config, capabilities, nondeterministic, fp8_backward + ) cudnn = import_cudnn() + version = check_frontend_version(cudnn) + entry = _get_supported_graph( + description.spec, capabilities, version, int(cudnn.backend_version()) + ) + backend, direction = _site(description.spec) + with _plan_lock: + if entry.workspace_size is None: + workspace_size = compile_graph( + cudnn, + entry.built.graph, + lambda stage, function: _diagnostics.run(backend, direction, stage, function), + ) + serialized_graph = bytes(entry.built.graph.serialize()) + entry.workspace_size = workspace_size + entry.serialized_graph = serialized_graph + _diagnostics.event(backend, direction, capabilities, "BUILD_PLANS") + uids = entry.built.uids + serialized = serialize_graph( + serialized_graph=entry.serialized_graph, + cudnn_frontend_version=version, + workspace_size=entry.workspace_size, + input_uids=[uids[name] for name, _, _ in description.inputs], + output_uids=[uids[name] for name, _, _ in description.outputs], + scalar_uids=[uids[name] for name, _ in description.scalars], + scalar_values=[value for _, value in description.scalars], + ) + return AttentionGraph( + serialized, + tuple(source for _, source, _ in description.inputs), + tuple(range(len(description.inputs))), + tuple(offset for _, _, offset in description.inputs), + tuple(slot for _, slot, _ in description.outputs), + tuple(offset for _, _, offset in description.outputs), + description.metadata_batch, + description.ragged_stats, + description.multipliers, + description.fp8, + ) + + +def _describe_graph(direction, shapes, config, capabilities, nondeterministic, fp8_backward): + """Translate JAX's API and storage conventions into a common graph description.""" avals = tuple(jax.ShapeDtypeStruct(*shape) for shape in shapes) batch, sq, sk, h, hg, dq, dv = _dims(avals, config) backward = direction == "bwd" ragged = config.qkv_layout.is_thd() ragged_stats = ragged and capabilities[0] >= 90 and capabilities[0] != 120 segments = batch * config.max_segments_per_seq if ragged else batch - graph_batch = _bucket(segments, 32, 512) if ragged_stats else segments - graph_sq = _bucket(batch * sq, 1024, 32768) if ragged_stats else sq - graph_sk = _bucket(batch * sk, 1024, 32768) if ragged_stats else sk + # JAX uses fixed metadata capacities, including for empty ragged storage. + # Keep positive graph capacities while runtime segment lengths may be zero. + graph_batch = _bucket(max(segments, 1), 32, 512) if ragged_stats else segments + graph_sq = _bucket(max(batch * sq, 1), 1024, 32768) if ragged_stats else sq + graph_sk = _bucket(max(batch * sk, 1), 1024, 32768) if ragged_stats else sk fp8 = config.fp8 and (not backward or fp8_backward) - version = int(cudnn.backend_version()) + version = cudnn_version() if ragged and capabilities[0] < 90: raise ValueError("F16/BF16 fused attention with THD format requires sm90 or later.") if ( @@ -118,179 +249,95 @@ def _build_graph(direction, shapes, config, capabilities, nondeterministic, fp8_ raise ValueError("FP8 THD attention requires cuDNN 9.23.0 or later.") if backward and capabilities[0] < 100: raise ValueError("FP8 THD backward requires sm100 or later.") - io_dtype = _dtype(cudnn, avals[0].dtype) - qkv_dtype = cudnn.data_type.FP8_E4M3 if fp8 else io_dtype - graph = cudnn.pygraph( - io_data_type=io_dtype, - intermediate_data_type=cudnn.data_type.FLOAT, - compute_data_type=cudnn.data_type.FLOAT, - ) - names, input_uids, input_indices, input_offsets = [], [], [], [] - scalar_uids, scalar_values = [], [] - uid = 0 - - def tensor(name, dim, stride, dtype=None, offset=0): - nonlocal uid - uid += 1 - names.append(name) - input_uids.append(uid) - input_indices.append(len(names) - 1) - input_offsets.append(offset) - return graph.tensor( - name=name, - dim=dim, - stride=stride, - data_type=io_dtype if dtype is None else dtype, - uid=uid, - ) + if fp8 and backward and config.attn_bias_type.name != "NO_BIAS": + raise ValueError("cuDNN FP8 backward does not support attention bias or ALiBi.") + io_dtype = _dtype(avals[0].dtype) + qkv_dtype = "FP8_E4M3" if fp8 else io_dtype + specs, inputs, outputs, scalars = {}, [], [], [] + + def tensor(name, dim, stride, dtype=io_dtype, source=None, offset=0, ragged_offset=""): + specs[name] = TensorSpec(tuple(dim), tuple(stride), dtype, ragged_offset=ragged_offset) + inputs.append((name, name if source is None else source, offset)) def constant(name, value): - nonlocal uid - uid += 1 - scalar_uids.append(uid) - scalar_values.append(np.asarray(value, np.float32).tobytes()) - return graph.tensor( - name=name, - dim=(1, 1, 1, 1), - stride=(1, 1, 1, 1), - data_type=cudnn.data_type.FLOAT, - uid=uid, - is_pass_by_value=True, - ) + specs[name] = TensorSpec((1, 1, 1, 1), (1, 1, 1, 1), "FLOAT", by_value=True) + scalars.append((name, np.asarray(value, np.float32).tobytes())) - def output(value, dim, stride, slot, offset=0, dtype=None): - nonlocal uid - uid += 1 - value.set_output(True).set_uid(uid).set_dim(dim).set_stride(stride) - value.set_data_type(io_dtype if dtype is None else dtype) - outputs.append((uid, slot, offset)) - return value + def output(name, dim, stride, slot, offset=0, dtype=io_dtype, ragged_offset=""): + specs[name] = TensorSpec( + tuple(dim), tuple(stride), dtype, output=True, ragged_offset=ragged_offset + ) + outputs.append((name, slot, offset)) - outputs = [] qpack = 3 if config.qkv_layout.is_qkvpacked() else 1 kvpack = 3 if qpack == 3 else (2 if config.qkv_layout.is_kvpacked() else 1) - # FP8 conversion preserves the packed storage layout. itemsize = 1 if fp8 else avals[0].dtype.itemsize qstride = (graph_sq * qpack * h * dq, dq, qpack * h * dq, 1) kstride = (graph_sk * kvpack * hg * dq, dq, kvpack * hg * dq, 1) vstride = (graph_sk * kvpack * hg * dv, dv, kvpack * hg * dv, 1) ostride = (graph_sq * h * dv, dv, h * dv, 1) - qdim, kdim, vdim = ( - (graph_batch, h, graph_sq, dq), - (graph_batch, hg, graph_sk, dq), - (graph_batch, hg, graph_sk, dv), - ) + qdim = (graph_batch, h, graph_sq, dq) + kdim = (graph_batch, hg, graph_sk, dq) + vdim = (graph_batch, hg, graph_sk, dv) odim = (graph_batch, h, graph_sq, dv) - q = tensor("q", qdim, qstride, qkv_dtype) - k = tensor( - "q" if qpack == 3 else "k", kdim, kstride, qkv_dtype, h * dq * itemsize if qpack == 3 else 0 + tensor("q", qdim, qstride, qkv_dtype, ragged_offset="offset_q" if ragged else "") + tensor( + "k", + kdim, + kstride, + qkv_dtype, + source="q" if qpack == 3 else "k", + offset=h * dq * itemsize if qpack == 3 else 0, + ragged_offset="offset_k" if ragged else "", ) - v = tensor( - "q" if qpack == 3 else ("k" if kvpack == 2 else "v"), + tensor( + "v", vdim, vstride, qkv_dtype, - 2 * h * dq * itemsize if qpack == 3 else (hg * dq * itemsize if kvpack == 2 else 0), + source="q" if qpack == 3 else ("k" if kvpack == 2 else "v"), + offset=2 * h * dq * itemsize if qpack == 3 else (hg * dq * itemsize if kvpack == 2 else 0), + ragged_offset="offset_v" if ragged else "", ) - multipliers = (qpack * h * dq, kvpack * hg * dq, kvpack * hg * dv, h * dv, h) - offsets = [] if ragged: for name in ("offset_q", "offset_k", "offset_v", "offset_o", "offset_stats"): - offsets.append( - tensor(name, (graph_batch + 1, 1, 1, 1), (1, 1, 1, 1), cudnn.data_type.INT64) - ) - q.set_ragged_offset(offsets[0]) - k.set_ragged_offset(offsets[1]) - v.set_ragged_offset(offsets[2]) + tensor(name, (graph_batch + 1, 1, 1, 1), (1, 1, 1, 1), "INT64") stats_dim = (graph_batch, h, graph_sq, 1) stats_stride = (h * graph_sq, 1, h, 1) if ragged_stats else (h * graph_sq, graph_sq, 1, 1) - options = { - "attn_scale": float(config.scaling_factor), - "diagonal_alignment": ( - cudnn.diagonal_alignment.BOTTOM_RIGHT - if config.bottom_right_diagonal - else cudnn.diagonal_alignment.TOP_LEFT - ), - } + stats_offset = "offset_stats" if ragged_stats else "" + o_offset = "offset_o" if ragged else "" left, right = config.effective_window_size - if left >= 0: - options["left_bound" if fp8 else "diagonal_band_left_bound"] = left + 1 if config.attn_mask_type.is_causal(): right = 0 - if right >= 0: - options["right_bound" if fp8 else "diagonal_band_right_bound"] = right - if config.attn_mask_type.is_padding() or ragged: - options.update( - use_padding_mask=True, - seq_len_q=tensor("seq_q", (graph_batch, 1, 1, 1), (1, 1, 1, 1), cudnn.data_type.INT32), - seq_len_kv=tensor( - "seq_kv", (graph_batch, 1, 1, 1), (1, 1, 1, 1), cudnn.data_type.INT32 - ), - ) - if config.dropout_probability and config.is_training: - options["dropout"] = ( - float(config.dropout_probability), - tensor("rng", (1, 1, 1, 1), (1, 1, 1, 1), cudnn.data_type.INT64), - tensor("rng", (1, 1, 1, 1), (1, 1, 1, 1), cudnn.data_type.INT64, 8), - ) + padding = config.attn_mask_type.is_padding() or ragged + if padding: + for name in ("seq_q", "seq_kv"): + tensor(name, (graph_batch, 1, 1, 1), (1, 1, 1, 1), "INT32") + dropout = float(config.dropout_probability) if config.is_training else 0.0 + if dropout: + tensor("seed", (1, 1, 1, 1), (1, 1, 1, 1), "INT64", source="rng") + tensor("offset", (1, 1, 1, 1), (1, 1, 1, 1), "INT64", source="rng", offset=8) bias_type = config.attn_bias_type.name - if bias_type == "ALIBI": - if fp8 and backward: - raise ValueError("cuDNN FP8 backward does not support ALiBi.") - options["use_alibi_mask"] = True - elif bias_type != "NO_BIAS": - if fp8 and backward: - raise ValueError("cuDNN FP8 backward does not support attention bias.") - bias_shape = (*avals[3].shape[-3:],) - bias_shape = (math.prod(avals[3].shape[:-3]), *bias_shape) - options["bias"] = tensor("bias", bias_shape, _stride(bias_shape)) + if bias_type not in ("ALIBI", "NO_BIAS"): + bias_shape = (math.prod(avals[3].shape[:-3]), *avals[3].shape[-3:]) + tensor("bias", bias_shape, _stride(bias_shape)) if backward and bias_shape[:3] != (1, 1, 1): - dbias = graph.tensor( - name="dbias", dim=bias_shape, stride=_stride(bias_shape), data_type=io_dtype - ) - options["dBias"] = output(dbias, bias_shape, _stride(bias_shape), 3) + output("dbias", bias_shape, _stride(bias_shape), 3) if config.softmax_type.name != "VANILLA_SOFTMAX": - options["sink_token"] = tensor("sink", (1, h, 1, 1), (h, 1, 1, 1), cudnn.data_type.FLOAT) + tensor("sink", (1, h, 1, 1), (h, 1, 1, 1), "FLOAT") if backward: - dsink = graph.tensor( - name="dsink", dim=(1, h, 1, 1), stride=(h, 1, 1, 1), data_type=cudnn.data_type.FLOAT - ) - options["dSink_token"] = output( - dsink, (1, h, 1, 1), (h, 1, 1, 1), 4, dtype=cudnn.data_type.FLOAT - ) - + output("dsink", (1, h, 1, 1), (h, 1, 1, 1), 4, dtype="FLOAT") + constant("attn_scale", config.scaling_factor) if fp8: for name in ("descale_q", "descale_k", "descale_v"): - options[name] = tensor(name, (1, 1, 1, 1), (1, 1, 1, 1), cudnn.data_type.FLOAT) - # Softmax probabilities are in [0, 1]; use the full E4M3 range. - options["scale_s"] = constant("scale_s", 448.0) - options["descale_s"] = constant("descale_s", 1.0 / 448.0) + tensor(name, (1, 1, 1, 1), (1, 1, 1, 1), "FLOAT") + constant("scale_s", 448.0) + constant("descale_s", 1.0 / 448.0) if backward: - options["use_deterministic_algorithm"] = not nondeterministic - o = tensor("o", odim, ostride) - do = tensor("do", odim, ostride, cudnn.data_type.FP8_E5M2 if fp8 else io_dtype) - stats = tensor("stats", stats_dim, stats_stride, cudnn.data_type.FLOAT) - if ragged: - o.set_ragged_offset(offsets[3]) - do.set_ragged_offset(offsets[3]) - if ragged_stats: - stats.set_ragged_offset(offsets[4]) - if fp8: - options["descale_o"] = constant("descale_o", 1.0) - options["descale_dO"] = tensor( - "descale_do", (1, 1, 1, 1), (1, 1, 1, 1), cudnn.data_type.FLOAT - ) - options["scale_dP"] = constant("scale_dp", 1.0) - options["descale_dP"] = constant("descale_dp", 1.0) - for name in ("scale_dQ", "scale_dK", "scale_dV"): - options[name] = constant(name, 1.0) - gradients = graph.sdpa_fp8_backward(q=q, k=k, v=v, o=o, dO=do, stats=stats, **options) - else: - if ragged_stats: - options["max_total_seq_len_q"] = graph_sq - options["max_total_seq_len_kv"] = graph_sk - gradients = graph.sdpa_backward(q=q, k=k, v=v, o=o, dO=do, stats=stats, **options) + tensor("o", odim, ostride, ragged_offset=o_offset) + tensor("do", odim, ostride, "FP8_E5M2" if fp8 else io_dtype, ragged_offset=o_offset) + tensor("stats", stats_dim, stats_stride, "FLOAT", ragged_offset=stats_offset) grad_slots = (0, 0 if qpack == 3 else 1, 0 if qpack == 3 else (1 if kvpack == 2 else 2)) grad_offsets = ( 0, @@ -301,65 +348,48 @@ def output(value, dim, stride, slot, offset=0, dtype=None): else (hg * dq * avals[0].dtype.itemsize if kvpack == 2 else 0) ), ) - for index, (grad, dim, stride, slot, offset) in enumerate( - zip( - gradients[:3], - (qdim, kdim, vdim), - (qstride, kstride, vstride), - grad_slots, - grad_offsets, - ) + for name, dim, stride, slot, offset, ragged_offset in zip( + ("dq", "dk", "dv"), + (qdim, kdim, vdim), + (qstride, kstride, vstride), + grad_slots, + grad_offsets, + ("offset_q", "offset_k", "offset_v"), ): - output(grad, dim, stride, slot, offset) - if ragged: - grad.set_ragged_offset(offsets[index]) + output(name, dim, stride, slot, offset, ragged_offset=ragged_offset if ragged else "") if fp8: - for slot, amax in enumerate(gradients[3:7], 5): - output(amax, (1, 1, 1, 1), (1, 1, 1, 1), slot, dtype=cudnn.data_type.FLOAT) + tensor("descale_do", (1, 1, 1, 1), (1, 1, 1, 1), "FLOAT") + for name in ("descale_o", "scale_dp", "descale_dp", "scale_dq", "scale_dk", "scale_dv"): + constant(name, 1.0) + for slot, name in enumerate(("amax_dq", "amax_dk", "amax_dv", "amax_dp"), 5): + output(name, (1, 1, 1, 1), (1, 1, 1, 1), slot, dtype="FLOAT") else: + output("o", odim, ostride, 0, ragged_offset=o_offset) + output("stats", stats_dim, stats_stride, 1, dtype="FLOAT", ragged_offset=stats_offset) if config.return_max_logit: - score_max = graph.tensor( - name="max", dim=stats_dim, stride=stats_stride, data_type=cudnn.data_type.FLOAT - ) - if ragged_stats: - score_max.set_ragged_offset(offsets[4]) - options["score_max"] = output( - score_max, stats_dim, stats_stride, 2, dtype=cudnn.data_type.FLOAT - ) - # The public JAX operation remains differentiable with dropout disabled - # by is_training=False, just as the previous common-backed path did. - options["generate_stats"] = True + output("max", stats_dim, stats_stride, 2, dtype="FLOAT", ragged_offset=stats_offset) if fp8: - options["scale_o"] = constant("scale_o", 1.0) - o, stats, amax_s, amax_o = graph.sdpa_fp8(q=q, k=k, v=v, **options) - output(amax_s, (1, 1, 1, 1), (1, 1, 1, 1), 3, dtype=cudnn.data_type.FLOAT) - output(amax_o, (1, 1, 1, 1), (1, 1, 1, 1), 4, dtype=cudnn.data_type.FLOAT) - else: - o, stats = graph.sdpa(q=q, k=k, v=v, **options) - output(o, odim, ostride, 0) - if ragged: - o.set_ragged_offset(offsets[3]) - output(stats, stats_dim, stats_stride, 1, dtype=cudnn.data_type.FLOAT) - if ragged_stats: - stats.set_ragged_offset(offsets[4]) - - workspace, data, version = finalize_graph(cudnn, graph) - serialized = serialize_graph( - serialized_graph=data, - cudnn_frontend_version=version, - workspace_size=workspace, - input_uids=input_uids, - output_uids=[value[0] for value in outputs], - scalar_uids=scalar_uids, - scalar_values=scalar_values, + constant("scale_o", 1.0) + output("amax_s", (1, 1, 1, 1), (1, 1, 1, 1), 3, dtype="FLOAT") + output("amax_o", (1, 1, 1, 1), (1, 1, 1, 1), 4, dtype="FLOAT") + options = AttentionOptions( + backward=backward, + fp8=fp8, + bottom_right=config.bottom_right_diagonal, + left_bound=left + 1 if left >= 0 else None, + right_bound=right if right >= 0 else None, + padding=padding, + alibi=bias_type == "ALIBI", + dropout=dropout, + deterministic=backward and not nondeterministic, + max_total_seqlen_q=graph_sq if backward and ragged_stats and not fp8 else 0, + max_total_seqlen_kv=graph_sk if backward and ragged_stats and not fp8 else 0, ) - return AttentionGraph( - serialized, - tuple(names), - tuple(input_indices), - tuple(input_offsets), - tuple(value[1] for value in outputs), - tuple(value[2] for value in outputs), + return _AttentionDescription( + AttentionGraphSpec.from_dict(options, specs), + tuple(inputs), + tuple(outputs), + tuple(scalars), graph_batch, ragged_stats, multipliers, diff --git a/transformer_engine/jax/cpp_extensions/cudnn_graph.py b/transformer_engine/jax/cpp_extensions/cudnn_graph.py index e5acc167c59..b5cc35b7579 100644 --- a/transformer_engine/jax/cpp_extensions/cudnn_graph.py +++ b/transformer_engine/jax/cpp_extensions/cudnn_graph.py @@ -4,13 +4,18 @@ """Serialize Python-built cuDNN graphs for the JAX-owned FFI executor.""" import hashlib -import importlib from dataclasses import dataclass from typing import Sequence, Tuple import numpy as np import transformer_engine_jax +from transformer_engine.common.cudnn.graph import ( + import_cudnn as _import_cudnn, + check_graph_support, + compile_graph, +) + @dataclass(frozen=True) class SerializedGraph: @@ -55,12 +60,7 @@ def check_frontend_version(cudnn) -> int: def import_cudnn(): """Load the required Python frontend lazily.""" - try: - cudnn = importlib.import_module("cudnn") - except ImportError as exc: - raise ImportError( - "JAX fused_attn requires the cuDNN frontend Python package (`cudnn`)." - ) from exc + cudnn = _import_cudnn() check_frontend_version(cudnn) return cudnn @@ -109,19 +109,9 @@ def serialize_graph( def finalize_graph(cudnn, graph): """Validate, select and compile a supported execution plan.""" - try: - graph.validate() - graph.build_operation_graph() - graph.create_execution_plans([cudnn.heur_mode.A, cudnn.heur_mode.FALLBACK]) - graph.check_support() - graph.build_plans(cudnn.build_plan_policy.HEURISTICS_CHOICE) - except cudnn.cudnnGraphNotSupportedError as exc: - raise RuntimeError(f"cuDNN SDPA graph is not supported: {exc}") from exc - return ( - max(int(graph.get_workspace_size()), 1), - bytes(graph.serialize()), - check_frontend_version(cudnn), - ) + check_graph_support(cudnn, graph) + workspace_size = compile_graph(cudnn, graph) + return workspace_size, bytes(graph.serialize()), check_frontend_version(cudnn) def graph_attributes(graph): diff --git a/transformer_engine/jax/cpp_extensions/flex_attention.py b/transformer_engine/jax/cpp_extensions/flex_attention.py index d3789f39f20..86da72c80bf 100644 --- a/transformer_engine/jax/cpp_extensions/flex_attention.py +++ b/transformer_engine/jax/cpp_extensions/flex_attention.py @@ -3,7 +3,6 @@ # See LICENSE for license information. """cuDNN frontend score_mod fused attention helpers.""" -import inspect import os from dataclasses import dataclass from typing import Any, Callable, Dict, Mapping, Optional, Sequence, Tuple @@ -13,6 +12,26 @@ import numpy as np from jax import ffi +from transformer_engine.common.cudnn.attention import ( + AttentionOptions, + AttentionGraphSpec, + build_attention_graph, +) +from transformer_engine.common.cudnn.graph import ( + TensorSpec, + create_graph, + attention_matrix_spec, + row_major_stride as _row_major_stride, +) + +from transformer_engine.common.cudnn.score_mod import ( + UncacheableKey, + callback_cache_key, + explicit_cache_key, + freeze_cache_key, + wrap_score_mod as _wrap_score_mod, +) + from .cudnn_graph import ( SerializedGraph as _SerializedScoreModGraph, encode_frontend_version as _encode_cudnn_frontend_version, @@ -117,14 +136,7 @@ class _ScoreModScalarSpec: stride: Tuple[int, ...] = (1, 1, 1, 1) -class _UncacheableScoreModKey: - """Unique static key for callbacks that must not share compiled score_mod graphs.""" - - def __hash__(self): - return id(self) - - def __eq__(self, other): - return self is other +_UncacheableScoreModKey = UncacheableKey def _score_mod_key_is_uncacheable(key: Any) -> bool: @@ -132,86 +144,18 @@ def _score_mod_key_is_uncacheable(key: Any) -> bool: def _freeze_score_mod_cache_key(value: Any) -> Any: - """Convert a user-provided score_mod graph key into a hashable structure.""" - if _is_array_operand(value): - raise TypeError( - "score_mod_graph_cache_key() must not include tensors. Pass runtime tensors " - "through score_mod_tensors or score_mod_bprop_tensors instead." - ) - if isinstance(value, Mapping): - items = ( - ( - _freeze_score_mod_cache_key(key), - _freeze_score_mod_cache_key(val), - ) - for key, val in value.items() - ) - return tuple(sorted(items, key=repr)) - if isinstance(value, (list, tuple)): - return tuple(_freeze_score_mod_cache_key(item) for item in value) - if isinstance(value, (set, frozenset)): - items = (_freeze_score_mod_cache_key(item) for item in value) - return tuple(sorted(items, key=repr)) - try: - hash(value) - except TypeError as exc: - raise TypeError( - "score_mod_graph_cache_key() must return a hashable value or a nested " - "combination of mapping/list/tuple/set values." - ) from exc - return value + """Freeze a structural callback key using this framework's tensor predicate.""" + return freeze_cache_key(value, _is_array_operand) def _score_mod_explicit_cache_key(callback_owner: Any) -> Optional[Any]: - """Return a user-provided structural graph key for a score_mod callback.""" - explicit_key = getattr(callback_owner, "score_mod_graph_cache_key", None) - if explicit_key is None: - return None - explicit_key = explicit_key() if callable(explicit_key) else explicit_key - return _freeze_score_mod_cache_key(explicit_key) + """Read a callback owner's optional structural key.""" + return explicit_cache_key(callback_owner, _is_array_operand) def _score_mod_callback_cache_key(callback: Optional[Callable]) -> Any: - """Create a stable graph cache key for a score_mod callable. - - Module-level functions are assumed to have stable topology. Stateful bound methods and - callable instances need an explicit score_mod_graph_cache_key(); otherwise their graphs - are left uncached to avoid reusing stale graphs after Python object address reuse. - """ - if callback is None: - return None - self_obj = getattr(callback, "__self__", None) - func_obj = getattr(callback, "__func__", None) - if self_obj is not None and func_obj is not None: - explicit_key = _score_mod_explicit_cache_key(self_obj) - if explicit_key is None: - return _UncacheableScoreModKey() - return ( - "bound_method", - type(self_obj), - func_obj.__module__, - func_obj.__qualname__, - explicit_key, - ) - - explicit_key = _score_mod_explicit_cache_key(callback) - if explicit_key is not None: - return ( - "callable", - type(callback), - getattr(callback, "__module__", None), - getattr(callback, "__qualname__", None), - explicit_key, - ) - - if ( - inspect.isfunction(callback) - and callback.__closure__ is None - and "" not in callback.__qualname__ - ): - return ("function", callback.__module__, callback.__qualname__) - - return _UncacheableScoreModKey() + """Share stable callbacks and distinguish module-level anonymous functions.""" + return callback_cache_key(callback, _is_array_operand, _UncacheableScoreModKey) @dataclass(frozen=True) @@ -261,17 +205,7 @@ def __eq__(self, other): ) -# cuDNN frontend tensor UIDs are arbitrary, but assigning stable values makes serialized -# graphs deterministic and simplifies debugging across the Python graph builder and C++ executor. -_SCORE_MOD_UID_Q = 1 -_SCORE_MOD_UID_K = 2 -_SCORE_MOD_UID_V = 3 -_SCORE_MOD_UID_O = 4 -_SCORE_MOD_UID_STATS = 5 -_SCORE_MOD_UID_DO = 6 -_SCORE_MOD_UID_DQ = 7 -_SCORE_MOD_UID_DK = 8 -_SCORE_MOD_UID_DV = 9 +# Callback operands use a separate UID range from the shared SDPA builder. _SCORE_MOD_FWD_TENSOR_UID_BASE = 1000 _SCORE_MOD_BPROP_TENSOR_UID_BASE = 2000 _SCORE_MOD_FWD_SCALAR_UID_BASE = 3000 @@ -280,22 +214,14 @@ def __eq__(self, other): _score_mod_graph_cache: Dict[Tuple[Any, ...], _SerializedScoreModGraph] = {} -def _row_major_stride(shape: Sequence[int]) -> Tuple[int, ...]: - stride = [] - running = 1 - for dim in reversed(tuple(shape)): - stride.append(running) - running *= dim - return tuple(reversed(stride)) - - -def _bshd_as_bhsd_dim_stride(shape: Sequence[int]) -> Tuple[Tuple[int, ...], Tuple[int, ...]]: - if len(shape) != 4: - raise ValueError(f"score_mod requires rank-4 BSHD tensors, got shape={shape}.") - batch, seqlen, heads, head_dim = tuple(shape) - return ( - (batch, heads, seqlen, head_dim), - (seqlen * heads * head_dim, head_dim, heads * head_dim, 1), +def _matrix_spec(cudnn, aval, output=False): + """Translate JAX's contiguous BSHD storage into a shared logical descriptor.""" + batch, seqlen, heads, dim = aval.shape + return attention_matrix_spec( + (batch, heads, seqlen, dim), + "bshd", + _cudnn_data_type(cudnn, aval.dtype).name, + output=output, ) @@ -469,16 +395,6 @@ def _score_mod_graph_tensors( return graph_tensors, tuple(tensor_uids), tuple(scalar_uids), tuple(scalar_values) -def _wrap_score_mod(score_mod: Optional[Callable], graph_tensors: Dict[str, Any]): - if score_mod is None: - return None - - def wrapped_score_mod(sdpa_graph, score_tensor): - return score_mod(sdpa_graph, score_tensor, graph_tensors) - - return wrapped_score_mod - - def _graph_cache_key( direction: str, config: _FusedAttnScoreModConfig, @@ -501,28 +417,8 @@ def _shape_dtype(value) -> jax.ShapeDtypeStruct: def _build_score_mod_fwd_graph(q_aval, k_aval, v_aval, score_mod_avals, config): cudnn = _import_cudnn_for_score_mod() - - io_data_type = _cudnn_data_type(cudnn, q_aval.dtype) - graph = cudnn.pygraph( - io_data_type=io_data_type, - intermediate_data_type=cudnn.data_type.FLOAT, - compute_data_type=cudnn.data_type.FLOAT, - ) - - q_dim, q_stride = _bshd_as_bhsd_dim_stride(q_aval.shape) - k_dim, k_stride = _bshd_as_bhsd_dim_stride(k_aval.shape) - v_dim, v_stride = _bshd_as_bhsd_dim_stride(v_aval.shape) - q = graph.tensor( - name="q", dim=q_dim, stride=q_stride, data_type=io_data_type, uid=_SCORE_MOD_UID_Q - ) - k = graph.tensor( - name="k", dim=k_dim, stride=k_stride, data_type=io_data_type, uid=_SCORE_MOD_UID_K - ) - v = graph.tensor( - name="v", dim=v_dim, stride=v_stride, data_type=io_data_type, uid=_SCORE_MOD_UID_V - ) - - score_mod_graph_tensors, tensor_uids, scalar_uids, scalar_values = _score_mod_graph_tensors( + graph = create_graph(cudnn, _cudnn_data_type(cudnn, q_aval.dtype)) + extra, tensor_uids, scalar_uids, scalar_values = _score_mod_graph_tensors( cudnn, graph, config.score_mod_tensor_names, @@ -531,40 +427,35 @@ def _build_score_mod_fwd_graph(q_aval, k_aval, v_aval, score_mod_avals, config): _SCORE_MOD_FWD_TENSOR_UID_BASE, _SCORE_MOD_FWD_SCALAR_UID_BASE, ) - - output, stats = graph.sdpa( - name="te_score_mod_sdpa", - q=q, - k=k, - v=v, - generate_stats=config.is_training, - attn_scale=config.scaling_factor, - use_causal_mask=False, - score_mod=_wrap_score_mod(config.score_mod, score_mod_graph_tensors), - ) - - batch, q_seqlen, q_heads, _ = q_aval.shape - _, _, _, v_head_dim = v_aval.shape - output_dim, output_stride = _bshd_as_bhsd_dim_stride((batch, q_seqlen, q_heads, v_head_dim)) - output.set_output(True).set_uid(_SCORE_MOD_UID_O).set_dim(output_dim).set_stride(output_stride) - output.set_data_type(io_data_type) - - output_uids = [_SCORE_MOD_UID_O] + output_shape = (*q_aval.shape[:-1], v_aval.shape[-1]) + specs = { + "q": _matrix_spec(cudnn, q_aval), + "k": _matrix_spec(cudnn, k_aval), + "v": _matrix_spec(cudnn, v_aval), + "o": _matrix_spec(cudnn, jax.ShapeDtypeStruct(output_shape, q_aval.dtype), output=True), + } + output_names = ["o"] if config.is_training: - stats_shape = (batch, q_heads, q_seqlen, 1) - stats.set_output(True).set_uid(_SCORE_MOD_UID_STATS).set_dim(stats_shape).set_stride( - _row_major_stride(stats_shape) - ) - stats.set_data_type(cudnn.data_type.FLOAT) - output_uids.append(_SCORE_MOD_UID_STATS) - - workspace_size, serialized_graph, frontend_version = _finalize_score_mod_graph(cudnn, graph) + batch, seqlen, heads, _ = q_aval.shape + shape = (batch, heads, seqlen, 1) + specs["stats"] = TensorSpec(shape, _row_major_stride(shape), "FLOAT", output=True) + output_names.append("stats") + built = build_attention_graph( + cudnn, + graph, + AttentionGraphSpec.from_dict( + AttentionOptions(attn_scale=config.scaling_factor, generate_stats=config.is_training), + specs, + ), + score_mod=_wrap_score_mod(config.score_mod, extra), + ) + workspace_size, data, frontend_version = _finalize_score_mod_graph(cudnn, graph) return _serialized_score_mod_graph( - serialized_graph=serialized_graph, + serialized_graph=data, cudnn_frontend_version=frontend_version, workspace_size=workspace_size, - input_uids=[_SCORE_MOD_UID_Q, _SCORE_MOD_UID_K, _SCORE_MOD_UID_V, *tensor_uids], - output_uids=output_uids, + input_uids=[built.uids[name] for name in ("q", "k", "v")] + list(tensor_uids), + output_uids=[built.uids[name] for name in output_names], scalar_uids=scalar_uids, scalar_values=scalar_values, ) @@ -582,43 +473,8 @@ def _build_score_mod_bwd_graph( config, ): cudnn = _import_cudnn_for_score_mod() - - io_data_type = _cudnn_data_type(cudnn, q_aval.dtype) - graph = cudnn.pygraph( - io_data_type=io_data_type, - intermediate_data_type=cudnn.data_type.FLOAT, - compute_data_type=cudnn.data_type.FLOAT, - ) - - q_dim, q_stride = _bshd_as_bhsd_dim_stride(q_aval.shape) - k_dim, k_stride = _bshd_as_bhsd_dim_stride(k_aval.shape) - v_dim, v_stride = _bshd_as_bhsd_dim_stride(v_aval.shape) - o_dim, o_stride = _bshd_as_bhsd_dim_stride(output_aval.shape) - do_dim, do_stride = _bshd_as_bhsd_dim_stride(doutput_aval.shape) - q = graph.tensor( - name="q", dim=q_dim, stride=q_stride, data_type=io_data_type, uid=_SCORE_MOD_UID_Q - ) - k = graph.tensor( - name="k", dim=k_dim, stride=k_stride, data_type=io_data_type, uid=_SCORE_MOD_UID_K - ) - v = graph.tensor( - name="v", dim=v_dim, stride=v_stride, data_type=io_data_type, uid=_SCORE_MOD_UID_V - ) - output = graph.tensor( - name="o", dim=o_dim, stride=o_stride, data_type=io_data_type, uid=_SCORE_MOD_UID_O - ) - doutput = graph.tensor( - name="dO", dim=do_dim, stride=do_stride, data_type=io_data_type, uid=_SCORE_MOD_UID_DO - ) - stats = graph.tensor( - name="stats", - dim=tuple(int(dim) for dim in stats_aval.shape), - stride=_row_major_stride(stats_aval.shape), - data_type=cudnn.data_type.FLOAT, - uid=_SCORE_MOD_UID_STATS, - ) - - score_mod_graph_tensors, tensor_uids, scalar_uids, scalar_values = _score_mod_graph_tensors( + graph = create_graph(cudnn, _cudnn_data_type(cudnn, q_aval.dtype)) + extra, tensor_uids, scalar_uids, scalar_values = _score_mod_graph_tensors( cudnn, graph, config.score_mod_tensor_names, @@ -627,56 +483,49 @@ def _build_score_mod_bwd_graph( _SCORE_MOD_FWD_TENSOR_UID_BASE, _SCORE_MOD_FWD_SCALAR_UID_BASE, ) - ( - score_mod_bprop_graph_tensors, - bprop_tensor_uids, - bprop_scalar_uids, - bprop_scalar_values, - ) = _score_mod_graph_tensors( + bprop_extra, bprop_tensor_uids, bprop_scalar_uids, bprop_scalar_values = ( + _score_mod_graph_tensors( + cudnn, + graph, + config.score_mod_bprop_tensor_names, + score_mod_bprop_avals, + config.score_mod_bprop_scalars, + _SCORE_MOD_BPROP_TENSOR_UID_BASE, + _SCORE_MOD_BPROP_SCALAR_UID_BASE, + ) + ) + specs = { + "q": _matrix_spec(cudnn, q_aval), + "k": _matrix_spec(cudnn, k_aval), + "v": _matrix_spec(cudnn, v_aval), + "o": _matrix_spec(cudnn, output_aval), + "do": _matrix_spec(cudnn, doutput_aval), + "stats": TensorSpec(tuple(stats_aval.shape), _row_major_stride(stats_aval.shape), "FLOAT"), + "dq": _matrix_spec(cudnn, q_aval, output=True), + "dk": _matrix_spec(cudnn, k_aval, output=True), + "dv": _matrix_spec(cudnn, v_aval, output=True), + } + built = build_attention_graph( cudnn, graph, - config.score_mod_bprop_tensor_names, - score_mod_bprop_avals, - config.score_mod_bprop_scalars, - _SCORE_MOD_BPROP_TENSOR_UID_BASE, - _SCORE_MOD_BPROP_SCALAR_UID_BASE, + AttentionGraphSpec.from_dict( + AttentionOptions( + backward=True, attn_scale=config.scaling_factor, deterministic=config.deterministic + ), + specs, + ), + score_mod=_wrap_score_mod(config.score_mod, extra), + score_mod_bprop=_wrap_score_mod(config.score_mod_bprop, bprop_extra), ) - - dq, dk, dv = graph.sdpa_backward( - name="te_score_mod_sdpa_backward", - q=q, - k=k, - v=v, - o=output, - dO=doutput, - stats=stats, - attn_scale=config.scaling_factor, - use_causal_mask=False, - score_mod=_wrap_score_mod(config.score_mod, score_mod_graph_tensors), - score_mod_bprop=_wrap_score_mod(config.score_mod_bprop, score_mod_bprop_graph_tensors), - use_deterministic_algorithm=config.deterministic, - ) - - dq.set_output(True).set_uid(_SCORE_MOD_UID_DQ).set_dim(q_dim).set_stride(q_stride) - dk.set_output(True).set_uid(_SCORE_MOD_UID_DK).set_dim(k_dim).set_stride(k_stride) - dv.set_output(True).set_uid(_SCORE_MOD_UID_DV).set_dim(v_dim).set_stride(v_stride) - - workspace_size, serialized_graph, frontend_version = _finalize_score_mod_graph(cudnn, graph) + workspace_size, data, frontend_version = _finalize_score_mod_graph(cudnn, graph) return _serialized_score_mod_graph( - serialized_graph=serialized_graph, + serialized_graph=data, cudnn_frontend_version=frontend_version, workspace_size=workspace_size, - input_uids=[ - _SCORE_MOD_UID_Q, - _SCORE_MOD_UID_K, - _SCORE_MOD_UID_V, - _SCORE_MOD_UID_O, - _SCORE_MOD_UID_DO, - _SCORE_MOD_UID_STATS, - *tensor_uids, - *bprop_tensor_uids, - ], - output_uids=[_SCORE_MOD_UID_DQ, _SCORE_MOD_UID_DK, _SCORE_MOD_UID_DV], + input_uids=[built.uids[name] for name in ("q", "k", "v", "o", "do", "stats")] + + list(tensor_uids) + + list(bprop_tensor_uids), + output_uids=[built.uids[name] for name in ("dq", "dk", "dv")], scalar_uids=[*scalar_uids, *bprop_scalar_uids], scalar_values=[*scalar_values, *bprop_scalar_values], ) diff --git a/transformer_engine/jax/csrc/extensions/attention.cpp b/transformer_engine/jax/csrc/extensions/cudnn_graph.cpp similarity index 99% rename from transformer_engine/jax/csrc/extensions/attention.cpp rename to transformer_engine/jax/csrc/extensions/cudnn_graph.cpp index 4d4c2c522b3..40b9fa7cdee 100644 --- a/transformer_engine/jax/csrc/extensions/attention.cpp +++ b/transformer_engine/jax/csrc/extensions/cudnn_graph.cpp @@ -4,7 +4,7 @@ * See LICENSE for license information. ************************************************************************/ -// JAX-owned execution of plans built with the cuDNN Python frontend. +// JAX-owned execution of generic graphs built with the cuDNN Python frontend. // This translation unit deliberately has no TE common dependencies. #include #include diff --git a/transformer_engine/pytorch/attention/dot_product_attention/flex_attention.py b/transformer_engine/pytorch/attention/dot_product_attention/flex_attention.py index 22f388dc4b4..54a72a31aa8 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/flex_attention.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/flex_attention.py @@ -5,11 +5,24 @@ """cuDNN-backed Flex Attention helpers.""" from dataclasses import dataclass -import inspect from typing import Any, Callable, Dict, Optional, Tuple import torch +from transformer_engine.common.cudnn.attention import ( + AttentionOptions, + AttentionGraphSpec, + build_attention_graph, +) +from transformer_engine.common.cudnn.graph import TensorSpec + +from transformer_engine.common.cudnn.score_mod import ( + callback_cache_key, + explicit_cache_key, + freeze_cache_key, + wrap_score_mod as _wrap_score_mod, +) + from transformer_engine.pytorch.cpp_extensions.cudnn import ( _import_cudnn_frontend, _build_cudnn_pygraph, @@ -38,98 +51,35 @@ def _bhsd_dim_stride( raise ValueError(f"Flex Attention only supports SBHD/BSHD tensor formats, got {tensor_format}.") -def _bhsd_graph_tensor(graph, tensor: torch.Tensor, tensor_format: str): - """Create a cuDNN graph tensor with BHSD dims and TE-layout strides.""" - dim, stride = _bhsd_dim_stride(tensor, tensor_format) - return graph.tensor(dim=dim, stride=stride, data_type=tensor.dtype) +def _graph_tensor_spec(tensor, tensor_format=None, output=False): + """Translate framework metadata into the shared graph descriptor.""" + dim, stride = ( + _bhsd_dim_stride(tensor, tensor_format) + if tensor_format is not None + else (tuple(tensor.shape), tuple(tensor.stride())) + ) + dtype = {torch.float16: "HALF", torch.bfloat16: "BFLOAT16", torch.float32: "FLOAT"}[ + tensor.dtype + ] + return TensorSpec(dim, stride, dtype, output=output) # score_mod graph cache helpers. def _freeze_score_mod_cache_key(value: Any) -> Any: - """Convert a user-provided score_mod graph key into a hashable structure.""" - if isinstance(value, torch.Tensor): - raise TypeError( - "score_mod_graph_cache_key() must not include tensors. Pass runtime tensors " - "through score_mod_tensors or score_mod_bprop_tensors instead." - ) - if isinstance(value, dict): - items = ( - ( - _freeze_score_mod_cache_key(key), - _freeze_score_mod_cache_key(val), - ) - for key, val in value.items() - ) - return tuple(sorted(items, key=repr)) - if isinstance(value, (list, tuple)): - return tuple(_freeze_score_mod_cache_key(item) for item in value) - if isinstance(value, (set, frozenset)): - items = (_freeze_score_mod_cache_key(item) for item in value) - return tuple(sorted(items, key=repr)) - try: - hash(value) - except TypeError as exc: - raise TypeError( - "score_mod_graph_cache_key() must return a hashable value or a nested " - "combination of dict/list/tuple/set values." - ) from exc - return value + """Freeze a structural callback key using this framework's tensor predicate.""" + return freeze_cache_key(value, lambda value: isinstance(value, torch.Tensor)) def _score_mod_explicit_cache_key(callback_owner: Any) -> Optional[Any]: - """Return a user-provided structural graph key for a score_mod callback.""" - explicit_key = getattr(callback_owner, "score_mod_graph_cache_key", None) - if explicit_key is None: - return None - explicit_key = explicit_key() if callable(explicit_key) else explicit_key - return _freeze_score_mod_cache_key(explicit_key) + """Read a callback owner's optional structural key.""" + return explicit_cache_key(callback_owner, lambda value: isinstance(value, torch.Tensor)) def _score_mod_callback_cache_key(callback: Optional[Callable]) -> Any: - """Create a stable graph cache key for a score_mod callable. - - Module-level named functions are assumed to have stable topology. Anonymous functions - are keyed by code object because lambdas in the same module can share the same - qualname. Stateful bound methods and callable instances need an explicit - score_mod_graph_cache_key(); otherwise their graphs are left uncached to avoid reusing - stale graphs after Python object address reuse. - """ - if callback is None: - return None - self_obj = getattr(callback, "__self__", None) - func_obj = getattr(callback, "__func__", None) - if self_obj is not None and func_obj is not None: - explicit_key = _score_mod_explicit_cache_key(self_obj) - if explicit_key is None: - return _SCORE_MOD_UNCACHEABLE - return ( - "bound_method", - type(self_obj), - func_obj.__module__, - func_obj.__qualname__, - explicit_key, - ) - - explicit_key = _score_mod_explicit_cache_key(callback) - if explicit_key is not None: - return ( - "callable", - type(callback), - getattr(callback, "__module__", None), - getattr(callback, "__qualname__", None), - explicit_key, - ) - - if ( - inspect.isfunction(callback) - and callback.__closure__ is None - and "" not in callback.__qualname__ - ): - if callback.__name__ == "" or not callback.__qualname__: - return ("function", callback.__module__, callback.__code__) - return ("function", callback.__module__, callback.__qualname__) - - return _SCORE_MOD_UNCACHEABLE + """Share stable callbacks and distinguish module-level anonymous functions.""" + return callback_cache_key( + callback, lambda value: isinstance(value, torch.Tensor), lambda: _SCORE_MOD_UNCACHEABLE + ) def _score_mod_device_key(device: torch.device) -> Tuple[Any, ...]: @@ -175,17 +125,6 @@ def _make_cudnn_graph_tensor_dict(graph, tensors: Optional[Dict[str, torch.Tenso # cuDNN frontend score_mod graph helpers. -def _wrap_score_mod(score_mod: Optional[Callable], graph_tensors: Dict[str, Any]): - """Adapt TE's score_mod signature to cuDNN frontend's two-argument callback.""" - if score_mod is None: - return None - - def _wrapped_score_mod(sdpa_graph, score_tensor): - return score_mod(sdpa_graph, score_tensor, graph_tensors) - - return _wrapped_score_mod - - @dataclass class _CudnnScoreModFwdGraphEntry: """Cached cuDNN frontend graph and graph tensor handles for score_mod fprop.""" @@ -313,43 +252,33 @@ def _build_cudnn_score_mod_fwd_graph( cudnn = _import_cudnn_frontend() graph = _build_cudnn_pygraph(query_layer.dtype, query_layer.device) - q = _bhsd_graph_tensor(graph, query_layer, q_format) - k = _bhsd_graph_tensor(graph, key_layer, kv_format) - v = _bhsd_graph_tensor(graph, value_layer, kv_format) - - score_mod_graph_tensors = _make_cudnn_graph_tensor_dict(graph, score_mod_tensors) - wrapped_score_mod = _wrap_score_mod(score_mod, score_mod_graph_tensors) - - output_dim, output_stride = _bhsd_dim_stride(output_layer, q_format) - output, stats_tensor = graph.sdpa( - name="te_score_mod_sdpa", - q=q, - k=k, - v=v, - generate_stats=is_training, - attn_scale=attn_scale, - use_causal_mask=False, - score_mod=wrapped_score_mod, - ) - output.set_output(True).set_dim(output_dim).set_stride(output_stride) - + extra = _make_cudnn_graph_tensor_dict(graph, score_mod_tensors) + specs = { + "q": _graph_tensor_spec(query_layer, q_format), + "k": _graph_tensor_spec(key_layer, kv_format), + "v": _graph_tensor_spec(value_layer, kv_format), + "o": _graph_tensor_spec(output_layer, q_format, output=True), + } if is_training: assert stats is not None - stats_tensor.set_output(True).set_dim(stats.size()).set_stride( - stats.stride() - ).set_data_type(cudnn.data_type.FLOAT) - else: - stats_tensor = None - + specs["stats"] = _graph_tensor_spec(stats, output=True) + built = build_attention_graph( + cudnn, + graph, + AttentionGraphSpec.from_dict( + AttentionOptions(attn_scale=attn_scale, generate_stats=is_training), specs + ), + score_mod=_wrap_score_mod(score_mod, extra), + ) workspace_size = _finalize_cudnn_graph(graph) return _CudnnScoreModFwdGraphEntry( graph=graph, - q=q, - k=k, - v=v, - output=output, - stats=stats_tensor, - score_mod_graph_tensors=score_mod_graph_tensors, + q=built.tensors["q"], + k=built.tensors["k"], + v=built.tensors["v"], + output=built.tensors["o"], + stats=built.tensors.get("stats"), + score_mod_graph_tensors=extra, workspace_size=workspace_size, ) @@ -408,61 +337,49 @@ def _build_cudnn_score_mod_bwd_graph( deterministic: bool, ) -> _CudnnScoreModBwdGraphEntry: """Build a cached cuDNN frontend graph for score_mod bprop.""" + cudnn = _import_cudnn_frontend() graph = _build_cudnn_pygraph(query_layer.dtype, query_layer.device) - q = _bhsd_graph_tensor(graph, query_layer, q_format) - k = _bhsd_graph_tensor(graph, key_layer, kv_format) - v = _bhsd_graph_tensor(graph, value_layer, kv_format) - output = _bhsd_graph_tensor(graph, output_layer, q_format) - d_output = _bhsd_graph_tensor(graph, d_out, q_format) - stats_tensor = graph.tensor_like(stats) - - score_mod_graph_tensors = _make_cudnn_graph_tensor_dict(graph, score_mod_tensors) - score_mod_bprop_graph_tensors = ( + extra = _make_cudnn_graph_tensor_dict(graph, score_mod_tensors) + bprop_extra = ( _make_cudnn_graph_tensor_dict(graph, score_mod_bprop_tensors) if score_mod_bprop is not None else {} ) - wrapped_score_mod = _wrap_score_mod(score_mod, score_mod_graph_tensors) - wrapped_score_mod_bprop = _wrap_score_mod(score_mod_bprop, score_mod_bprop_graph_tensors) - - dq_layer = torch.empty_like(query_layer) - dk_layer = torch.empty_like(key_layer) - dv_layer = torch.empty_like(value_layer) - dq_dim, dq_stride = _bhsd_dim_stride(dq_layer, q_format) - dk_dim, dk_stride = _bhsd_dim_stride(dk_layer, kv_format) - dv_dim, dv_stride = _bhsd_dim_stride(dv_layer, kv_format) - dq, dk, dv = graph.sdpa_backward( - name="te_score_mod_sdpa_backward", - q=q, - k=k, - v=v, - o=output, - dO=d_output, - stats=stats_tensor, - attn_scale=attn_scale, - use_causal_mask=False, - score_mod=wrapped_score_mod, - score_mod_bprop=wrapped_score_mod_bprop, - use_deterministic_algorithm=deterministic, + specs = { + "q": _graph_tensor_spec(query_layer, q_format), + "k": _graph_tensor_spec(key_layer, kv_format), + "v": _graph_tensor_spec(value_layer, kv_format), + "o": _graph_tensor_spec(output_layer, q_format), + "do": _graph_tensor_spec(d_out, q_format), + "stats": _graph_tensor_spec(stats), + "dq": _graph_tensor_spec(torch.empty_like(query_layer), q_format, output=True), + "dk": _graph_tensor_spec(torch.empty_like(key_layer), kv_format, output=True), + "dv": _graph_tensor_spec(torch.empty_like(value_layer), kv_format, output=True), + } + built = build_attention_graph( + cudnn, + graph, + AttentionGraphSpec.from_dict( + AttentionOptions(backward=True, attn_scale=attn_scale, deterministic=deterministic), + specs, + ), + score_mod=_wrap_score_mod(score_mod, extra), + score_mod_bprop=_wrap_score_mod(score_mod_bprop, bprop_extra), ) - dq.set_output(True).set_dim(dq_dim).set_stride(dq_stride) - dk.set_output(True).set_dim(dk_dim).set_stride(dk_stride) - dv.set_output(True).set_dim(dv_dim).set_stride(dv_stride) - workspace_size = _finalize_cudnn_graph(graph) return _CudnnScoreModBwdGraphEntry( graph=graph, - q=q, - k=k, - v=v, - output=output, - d_output=d_output, - stats=stats_tensor, - dq=dq, - dk=dk, - dv=dv, - score_mod_graph_tensors=score_mod_graph_tensors, - score_mod_bprop_graph_tensors=score_mod_bprop_graph_tensors, + q=built.tensors["q"], + k=built.tensors["k"], + v=built.tensors["v"], + output=built.tensors["o"], + d_output=built.tensors["do"], + stats=built.tensors["stats"], + dq=built.tensors["dq"], + dk=built.tensors["dk"], + dv=built.tensors["dv"], + score_mod_graph_tensors=extra, + score_mod_bprop_graph_tensors=bprop_extra, workspace_size=workspace_size, ) diff --git a/transformer_engine/pytorch/cpp_extensions/cudnn.py b/transformer_engine/pytorch/cpp_extensions/cudnn.py index fee40d0e98d..10b36b2d53a 100644 --- a/transformer_engine/pytorch/cpp_extensions/cudnn.py +++ b/transformer_engine/pytorch/cpp_extensions/cudnn.py @@ -4,23 +4,18 @@ """cuDNN Python frontend graph lifecycle shared by attention implementations.""" -import importlib import threading import torch -_handles = threading.local() - +from transformer_engine.common.cudnn.graph import ( + import_cudnn as _import_cudnn_frontend, + create_graph, + check_graph_support, + compile_graph, +) -def _import_cudnn_frontend(): - """Import the cuDNN frontend Python package.""" - try: - return importlib.import_module("cudnn") - except ImportError as exc: - raise ImportError( - "cuDNN frontend Python package not found. " - "Install it with: pip install nvidia-cudnn-frontend" - ) from exc +_handles = threading.local() def _get_cudnn_current_stream_handle(cudnn, device): @@ -46,23 +41,14 @@ def _build_cudnn_pygraph(dtype, device): dtype = {torch.float16: cudnn.data_type.HALF, torch.bfloat16: cudnn.data_type.BFLOAT16}.get( dtype, dtype ) - return cudnn.pygraph( - io_data_type=dtype, - intermediate_data_type=cudnn.data_type.FLOAT, - compute_data_type=cudnn.data_type.FLOAT, - handle=_get_cudnn_current_stream_handle(cudnn, device), - ) + return create_graph(cudnn, dtype, _get_cudnn_current_stream_handle(cudnn, device)) def _finalize_cudnn_graph(graph): """Check support, build an execution plan, and return workspace bytes.""" cudnn = _import_cudnn_frontend() - graph.validate() - graph.build_operation_graph() - graph.create_execution_plans([cudnn.heur_mode.A, cudnn.heur_mode.FALLBACK]) - graph.check_support() - graph.build_plans(cudnn.build_plan_policy.HEURISTICS_CHOICE) - return max(graph.get_workspace_size(), 1) + check_graph_support(cudnn, graph) + return compile_graph(cudnn, graph) def _execute_cudnn_graph(graph, variant_pack, workspace_size, device): diff --git a/transformer_engine/pytorch/cpp_extensions/cudnn_fused_attn.py b/transformer_engine/pytorch/cpp_extensions/cudnn_fused_attn.py index 647c2667392..a05297763e1 100644 --- a/transformer_engine/pytorch/cpp_extensions/cudnn_fused_attn.py +++ b/transformer_engine/pytorch/cpp_extensions/cudnn_fused_attn.py @@ -9,16 +9,27 @@ """ from dataclasses import dataclass -import atexit -from collections import Counter -import os -import sys import threading -import time import torch import transformer_engine_torch as tex +from transformer_engine.common.cudnn.attention import ( + AttentionOptions, + AttentionGraphSpec, + build_attention_graph, +) +from transformer_engine.common.cudnn.graph import ( + TensorSpec as _TensorSpec, + CacheDiagnostics, + attention_matrix_spec, + attention_stats_spec, + mxfp8_scale_spec, + check_graph_support, + compile_graph, + bucket, +) + from .cudnn import ( _import_cudnn_frontend, _build_cudnn_pygraph, @@ -30,71 +41,26 @@ _plans = {} _cache_lock = threading.RLock() -_diagnostics = Counter() -_build_times = Counter() - - -def _debug_enabled(): - setting = os.getenv("NVTE_FUSED_ATTN_CACHE_DEBUG", "0") - level, _, ranks = setting.partition(":") - rank = next( - ( - os.environ[n] - for n in ("RANK", "LOCAL_RANK", "OMPI_COMM_WORLD_RANK", "SLURM_PROCID") - if n in os.environ - ), - None, - ) - enabled = rank is None or ranks == "all" or rank in (ranks or "0").split(",") - return int(level or 0) if enabled else 0 +_diagnostics = CacheDiagnostics("PyTorch") def _event(attn, event, key=None): - level = _debug_enabled() - if level == 0: - return - site = ("fp8" if attn.fp8 else "f16", "bwd" if attn.backward else "fwd") - _diagnostics[(*site, event)] += 1 - if level >= 2: - print( - f"[FUSED-ATTN-CACHE] tid={threading.get_ident()} dev={attn.device.index} | " - f"{site[0]} {site[1]} {event}" - + (f" | {key!r}" if key is not None else ""), - file=sys.stderr, - flush=True, - ) + _diagnostics.event( + "fp8" if attn.fp8 else "f16", + "bwd" if attn.backward else "fwd", + attn.device.index, + event, + key, + ) def _build_stage(attn, name, function): - if not _debug_enabled(): - return function() - start = time.perf_counter() - try: - return function() - finally: - site = ("fp8" if attn.fp8 else "f16", "bwd" if attn.backward else "fwd", name) - _build_times[(*site, "seconds")] += time.perf_counter() - start - _build_times[(*site, "calls")] += 1 - - -def _print_diagnostics(): - if not _diagnostics: - return - print("[FUSED-ATTN-CACHE] Python frontend summary", file=sys.stderr) - for (backend, direction, event), count in sorted(_diagnostics.items()): - print(f" {backend} {direction} {event}={count}", file=sys.stderr) - for backend, direction, stage, kind in sorted(_build_times): - if kind == "calls": - calls = _build_times[(backend, direction, stage, kind)] - seconds = _build_times[(backend, direction, stage, "seconds")] - print( - f" {backend} {direction} {stage}: {calls} calls, " - f"{1000 * seconds / calls:.3f} ms average", - file=sys.stderr, - ) - - -atexit.register(_print_diagnostics) + return _diagnostics.run( + "fp8" if attn.fp8 else "f16", + "bwd" if attn.backward else "fwd", + name, + function, + ) def _formats(layout): @@ -102,17 +68,6 @@ def _formats(layout): return parts[0], parts[1] if len(parts) > 1 else parts[0] -def _strides(dim, fmt): - b, h, s, d = dim - if fmt in ("bshd", "thd"): - return (s * h * d, d, h * d, 1) - if fmt == "sbhd": - return (h * d, d, b * h * d, 1) - if fmt == "bhsd": - return (h * s * d, s * d, d, 1) - raise ValueError(f"Unsupported attention format: {fmt}") - - def _shape(dim, fmt, tokens=0): b, h, s, d = dim if fmt == "bshd": @@ -151,57 +106,47 @@ def _descale(tensor, columnwise=False): return tensor._scale_inv -def _round_up(value, multiple): - return (value + multiple - 1) // multiple * multiple - - -def _token_bucket(tokens): - if tokens == 0: - return 0 - if tokens <= 1024: - return 1024 - if tokens <= 32768: - return 1 << (tokens - 1).bit_length() - return _round_up(tokens, 32768) +@dataclass +class _PlanEntry: + """A supported graph, compiled lazily under the cache lock.""" - -@dataclass(frozen=True) -class _TensorSpec: - dim: tuple - stride: tuple - dtype: str - output: bool = False - ragged_offset: str = "" - reordered: bool = False - by_value: bool = False + graph: object + tensors: dict + workspace_size: int | None = None class _AttentionGraph: """Runtime tensors and data-free descriptors for one attention invocation.""" - def __init__(self, params, backward=False): + def __init__( + self, params, backward=False, *, dimensions=None, device=None, fp8=None, mxfp8=None + ): self.params = params self.backward = backward - self.device = params["q"].device + self.device = params["q"].device if device is None else device self.q_format, self.kv_format = _formats(params["qkv_layout"]) - self.b = params["cu_seqlens_q"].numel() - 1 - self.h = params["q"].shape[1] if self.q_format == "bhsd" else params["q"].shape[-2] - self.hg = params["k"].shape[1] if self.kv_format == "bhsd" else params["k"].shape[-2] - self.dq = params["q"].shape[-1] - self.dv = params["v"].shape[-1] - self.tq = params["q"].shape[0] if self.q_format == "thd" else 0 - self.tk = params["k"].shape[0] if self.kv_format == "thd" else 0 + if dimensions is None: + dimensions = ( + params["cu_seqlens_q"].numel() - 1, + params["q"].shape[1] if self.q_format == "bhsd" else params["q"].shape[-2], + params["k"].shape[1] if self.kv_format == "bhsd" else params["k"].shape[-2], + params["q"].shape[-1], + params["v"].shape[-1], + params["q"].shape[0] if self.q_format == "thd" else 0, + params["k"].shape[0] if self.kv_format == "thd" else 0, + ) + self.b, self.h, self.hg, self.dq, self.dv, self.tq, self.tk = dimensions capability = torch.cuda.get_device_capability(self.device) ragged_graph = capability >= (9, 0) and capability != (12, 0) self.ragged_stats = self.q_format == "thd" and ragged_graph - self.sq = _token_bucket(self.tq) if self.ragged_stats else params["max_seqlen_q"] + self.sq = bucket(self.tq, 1024, 32768) if self.ragged_stats else params["max_seqlen_q"] self.skv = ( - _token_bucket(self.tk) + bucket(self.tk, 1024, 32768) if self.kv_format == "thd" and ragged_graph else params["max_seqlen_kv"] ) - self.fp8 = hasattr(params["q"], "_fp8_dtype") - self.mxfp8 = isinstance(params["q"], MXFP8TensorStorage) + self.fp8 = hasattr(params["q"], "_fp8_dtype") if fp8 is None else fp8 + self.mxfp8 = isinstance(params["q"], MXFP8TensorStorage) if mxfp8 is None else mxfp8 self.specs = {} self.tensors = {} @@ -223,59 +168,45 @@ def scalar(self, name, tensor, output=False, by_value=False): def matrix( self, name, tensor, fmt, heads, seqlen, cu_padded=None, columnwise=False, output=False ): - """Describe an attention matrix in logical BHSD order, retaining storage strides.""" + """Bind storage using the same matrix descriptors as support queries.""" data = _raw(tensor, columnwise) dim = (self.b, heads, seqlen, tensor.shape[-1]) - if fmt == "bshd": - order = (0, 2, 1, 3) - elif fmt == "sbhd": - order = (1, 2, 0, 3) - else: - order = (0, 1, 2, 3) - offset_name = "" - if fmt == "thd": - token_stride, head_stride, last_stride = data.stride() - stride = (seqlen * token_stride, head_stride, token_stride, last_stride) - offset_name = name + "_offset" - offsets = cu_padded.to(torch.int64) * token_stride + offset_name = name + "_offset" if fmt == "thd" else "" + if offset_name: + offsets = cu_padded.to(torch.int64) * data.stride(0) self.add(offset_name, offsets, (self.b + 1, 1, 1, 1), (1, 1, 1, 1)) - elif data.ndim == 4: - stride = tuple(data.stride(i) for i in order) - # Paged K/V have a page axis rather than a batch axis. - if name in ("k", "v") and self.params.get("page_table_k") is not None: - dim = tuple(data.shape[i] for i in order) - else: - stride = _strides(dim, fmt) - return self.add( - name, data, dim, stride, _dtype(tensor), output=output, ragged_offset=offset_name + elif name in ("k", "v") and self.params.get("page_table_k") is not None: + order = (0, 2, 1, 3) if fmt == "bshd" else (0, 1, 2, 3) + dim = tuple(data.shape[index] for index in order) + self.tensors[name] = data + self.specs[name] = attention_matrix_spec( + dim, + fmt, + _dtype(tensor), + storage_stride=data.stride(), + output=output, + ragged_offset=offset_name, ) + return name def stats(self, name, tensor, output=False): - """Describe softmax statistics, including offsets for packed THD storage.""" - dim = (self.b, self.h, self.sq, 1) + """Bind dense or packed softmax statistics.""" if self.ragged_stats: offsets = self.params["cu_seqlens_q_padded"].to(torch.int64) * self.h self.add("stats_offset", offsets, (self.b + 1, 1, 1, 1), (1, 1, 1, 1)) - return self.add( - name, - tensor, - dim, - (self.h * self.sq, 1, self.h, 1), - output=output, - ragged_offset="stats_offset", - ) - return self.add(name, tensor, dim, (self.h * self.sq, self.sq, 1, 1), output=output) + self.tensors[name] = tensor + self.specs[name] = attention_stats_spec( + self.b, self.h, self.sq, ragged=self.ragged_stats, output=output + ) + return name def scale(self, name, tensor, matrix_name, fmt, columnwise=False): - """Describe per-tensor FP8 scales or reordered MXFP8 block scales.""" + """Bind tensor scales or reordered MXFP8 block scales.""" if not self.mxfp8: return self.scalar(name, tensor) - b, h, s, d = self.specs[matrix_name].dim - if columnwise: - dim = (b, h, _round_up((s + 31) // 32, 4), _round_up(d, 128)) - else: - dim = (b, h, _round_up(s, 128), _round_up((d + 31) // 32, 4)) - return self.add(name, tensor, dim, _strides(dim, fmt), "FP8_E8M0", reordered=True) + self.tensors[name] = tensor + self.specs[name] = mxfp8_scale_spec(self.specs[matrix_name].dim, fmt, columnwise) + return name def options(self): """Return graph attributes that affect support and plan caching.""" @@ -288,24 +219,27 @@ def options(self): left, right = p["window_size"] or (-1, -1) if not causal and (left, right) == (-1, -1): bottom_right = False - return ( - mask, - p["attn_bias_type"], - p["softmax_type"], - left + 1 if left != -1 else None, - 0 if causal else (right if right != -1 else None), - bottom_right, - p["dropout"], - p.get("deterministic", False), - p.get("return_max_logit", False), - self.fp8, - self.mxfp8, - self.params["max_seqlen_kv"] if self.params.get("page_table_k") is not None else 0, + return AttentionOptions( + backward=self.backward, + fp8=self.fp8, + mxfp8=self.mxfp8, + bottom_right=bottom_right, + left_bound=left + 1 if left != -1 else None, + right_bound=0 if causal else (right if right != -1 else None), + padding="padding" in mask, + alibi=p["attn_bias_type"] == "alibi", + dropout=p["dropout"], + deterministic=p.get("deterministic", False), + paged_max_seqlen_kv=p["max_seqlen_kv"] if p.get("page_table_k") is not None else 0, ) + def description(self): + """Describe the graph independently of buffer addresses and scalar values.""" + return AttentionGraphSpec.from_dict(self.options(), self.specs) + def entry(self): """Share support-query graphs with execution without compiling on queries.""" - key = (self.device, self.backward, self.options(), tuple(sorted(self.specs.items()))) + key = (self.device, self.description()) # Autograd may run backward on another thread. Share device-specific plans; # each execution supplies its thread-local handle on the current stream. with _cache_lock: @@ -320,204 +254,32 @@ def entry(self): return entry def execute(self): - """Look up the graph by metadata, then bind this call's tensor addresses.""" + """Compile once, then bind this invocation's runtime tensor addresses.""" entry = self.entry() with _cache_lock: - if entry[2] is None: - cudnn = _import_cudnn_frontend() - _build_stage( - self, - "build_plans", - lambda: entry[0].build_plans(cudnn.build_plan_policy.HEURISTICS_CHOICE), + if entry.workspace_size is None: + entry.workspace_size = compile_graph( + _import_cudnn_frontend(), + entry.graph, + lambda name, function: _build_stage(self, name, function), ) - entry[2] = max(entry[0].get_workspace_size(), 1) _event(self, "BUILD_PLANS") - graph, tensors, workspace = entry _execute_cudnn_graph( - graph, {tensors[n]: t for n, t in self.tensors.items()}, workspace, self.device + entry.graph, + {entry.tensors[n]: t for n, t in self.tensors.items()}, + entry.workspace_size, + self.device, ) _event(self, "EXECUTE") def build(self): - """Construct and check a frontend graph without compiling its execution plan.""" + """Construct and check the common graph description without compiling it.""" cudnn = _import_cudnn_frontend() graph = _build_cudnn_pygraph(getattr(cudnn.data_type, self.specs["q"].dtype), self.device) - tensors = {} - - def tensor(name): - if name not in tensors: - spec = self.specs[name] - desc = graph.tensor( - name=name, - dim=spec.dim, - stride=spec.stride, - data_type=getattr(cudnn.data_type, spec.dtype), - is_pass_by_value=spec.by_value, - ) - if spec.ragged_offset: - desc.set_ragged_offset(tensor(spec.ragged_offset)) - if spec.reordered: - desc.set_reordering_type(cudnn.tensor_reordering.F8_128x4) - if spec.output: - desc.set_output(True) - tensors[name] = desc - return tensors[name] - - for name, spec in self.specs.items(): - if not spec.output: - tensor(name) - - mask, bias, softmax, left, right, bottom_right, dropout, deterministic, _, _, _, pages = ( - self.options() - ) - kwargs = { - "attn_scale": tensor("attn_scale"), - "diagonal_alignment": ( - cudnn.diagonal_alignment.BOTTOM_RIGHT - if bottom_right - else cudnn.diagonal_alignment.TOP_LEFT - ), - } - bound_names = ( - ("left_bound", "right_bound") - if (self.fp8 and (not self.mxfp8 or self.backward)) - else ("diagonal_band_left_bound", "diagonal_band_right_bound") - ) - kwargs.update(zip(bound_names, (left, right))) - if "padding" in mask: - kwargs.update( - use_padding_mask=True, seq_len_q=tensor("seq_q"), seq_len_kv=tensor("seq_kv") - ) - if bias == "alibi": - kwargs["use_alibi_mask"] = True - elif bias != "no_bias": - kwargs["bias"] = tensor("bias") - if self.backward and "dbias" in self.specs: - kwargs["dBias"] = tensor("dbias") - if dropout: - kwargs["dropout"] = (dropout, tensor("seed"), tensor("offset")) - if softmax != "vanilla": - kwargs["sink_token"] = tensor("softmax_offset") - if self.backward: - kwargs["dSink_token"] = tensor("dsoftmax_offset") - if self.backward: - kwargs["use_deterministic_algorithm"] = deterministic - else: - kwargs["generate_stats"] = True - if "max" in self.specs: - kwargs["score_max"] = tensor("max") - if pages: - kwargs.update( - paged_attention_k_table=tensor("page_table_k"), - paged_attention_v_table=tensor("page_table_v"), - paged_attention_max_seq_len_kv=pages, - ) - - q, k, v = (tensor(n) for n in ("q", "k", "v")) - if not self.fp8: - if self.backward: - outputs = graph.sdpa_backward( - q, k, v, tensor("o"), tensor("do"), tensor("stats"), **kwargs - ) - names = ("dq", "dk", "dv") - else: - outputs = graph.sdpa(q, k, v, **kwargs) - names = ("o", "stats") - elif not self.backward: - scale_args = [tensor(n) for n in ("descale_q", "descale_k", "descale_v")] - if self.mxfp8: - outputs = graph.sdpa_mxfp8(q, k, v, *scale_args, **kwargs) - names = ("o", "stats", "amax_o") - else: - outputs = graph.sdpa_fp8( - q, - k, - v, - *scale_args, - *(tensor(n) for n in ("descale_s", "scale_s", "scale_o")), - **kwargs, - ) - names = ("o", "stats", "amax_s", "amax_o") - elif self.mxfp8: - outputs = graph.sdpa_mxfp8_backward( - q, - tensor("q_t"), - k, - tensor("k_t"), - v, - tensor("o"), - tensor("do_f16"), - tensor("do"), - tensor("do_t"), - tensor("stats"), - *( - tensor(n) - for n in ( - "descale_q", - "descale_q_t", - "descale_k", - "descale_k_t", - "descale_v", - "descale_do", - "descale_do_t", - ) - ), - **kwargs, - ) - names = ("dq", "dk", "dv", "amax_dq", "amax_dk", "amax_dv") - else: - outputs = graph.sdpa_fp8_backward( - q, - k, - v, - tensor("o"), - tensor("do"), - tensor("stats"), - *( - tensor(n) - for n in ( - "descale_q", - "descale_k", - "descale_v", - "descale_o", - "descale_do", - "descale_s", - "descale_dp", - "scale_s", - "scale_dq", - "scale_dk", - "scale_dv", - "scale_dp", - ) - ), - **kwargs, - ) - names = ("dq", "dk", "dv", "amax_dq", "amax_dk", "amax_dv", "amax_dp") - - for name, desc in zip(names, outputs): - if name not in self.specs: - # MXFP8 amax outputs stay virtual, matching the C++ graph. - desc.set_dim((1, 1, 1, 1)).set_stride((1, 1, 1, 1)).set_data_type( - cudnn.data_type.FLOAT - ) - continue - spec = self.specs[name] - desc.set_output(True).set_dim(spec.dim).set_stride(spec.stride).set_data_type( - getattr(cudnn.data_type, spec.dtype) - ) - if spec.ragged_offset: - desc.set_ragged_offset(tensor(spec.ragged_offset)) - tensors[name] = desc + built = build_attention_graph(cudnn, graph, self.description()) _event(self, "CREATE_GRAPH") - _build_stage(self, "validate", graph.validate) - _build_stage(self, "build_operation_graph", graph.build_operation_graph) - _build_stage( - self, - "create_execution_plans", - lambda: graph.create_execution_plans([cudnn.heur_mode.A, cudnn.heur_mode.FALLBACK]), - ) - _build_stage(self, "check_support", graph.check_support) - return [graph, tensors, None] + check_graph_support(cudnn, graph, lambda name, function: _build_stage(self, name, function)) + return _PlanEntry(graph, built.tensors) def _prepare_inputs(attn): @@ -540,7 +302,7 @@ def _prepare_inputs(attn): if p["attn_bias_type"] not in ("no_bias", "alibi"): attn.add("bias", p["attn_bias"]) if p["softmax_type"] != "vanilla": - attn.add("softmax_offset", p["softmax_offset"]) + attn.add("sink", p["softmax_offset"]) if p["dropout"]: rng = p["rng_state"] attn.scalar("seed", rng[:1]) @@ -632,21 +394,6 @@ def name(mapping, value): if (config.dqkv_layout != tex.NVTE_QKV_Layout.NVTE_QKV_Layout_NOT_SET) else layout ) - attn = _AttentionGraph.__new__(_AttentionGraph) - attn.backward = backward - attn.device = torch.device("cuda", torch.cuda.current_device()) - attn.q_format, attn.kv_format = _formats(layout) - attn.b, attn.h, attn.hg = config.batch_size, config.num_attn_heads, config.num_gqa_groups - attn.dq, attn.dv = config.head_dim_qk, config.head_dim_v - attn.tq, attn.tk = config.num_tokens_q, config.num_tokens_kv - ragged_graph = torch.cuda.get_device_capability() >= (9, 0) and ( - torch.cuda.get_device_capability() != (12, 0) - ) - attn.ragged_stats = attn.q_format == "thd" and ragged_graph - attn.sq = _token_bucket(attn.tq) if attn.ragged_stats else config.max_seqlen_q - attn.skv = ( - _token_bucket(attn.tk) if attn.kv_format == "thd" and ragged_graph else config.max_seqlen_kv - ) types = { DType.kFloat16: "HALF", DType.kBFloat16: "BFLOAT16", @@ -654,9 +401,9 @@ def name(mapping, value): DType.kFloat8E4M3: "FP8_E4M3", DType.kFloat8E5M2: "FP8_E5M2", } - attn.fp8 = types[DType.cast(config.qkv_dtype)].startswith("FP8") - attn.mxfp8 = attn.fp8 and config.scaling_mode == tex.NVTEScalingMode.NVTE_MXFP8_1D_SCALING - attn.params = { + fp8 = types[DType.cast(config.qkv_dtype)].startswith("FP8") + mxfp8 = fp8 and config.scaling_mode == tex.NVTEScalingMode.NVTE_MXFP8_1D_SCALING + params = { "attn_mask_type": name(AttnMaskType, config.attn_mask_type), "attn_bias_type": name(AttnBiasType, config.bias_type), "softmax_type": name(SoftmaxType, config.softmax_type), @@ -668,7 +415,23 @@ def name(mapping, value): "max_seqlen_kv": config.max_seqlen_kv, "page_table_k": True if layout.startswith("paged_kv_") else None, } - attn.specs, attn.tensors = {}, {} + params.update(qkv_layout=layout, max_seqlen_q=config.max_seqlen_q) + attn = _AttentionGraph( + params, + backward, + dimensions=( + config.batch_size, + config.num_attn_heads, + config.num_gqa_groups, + config.head_dim_qk, + config.head_dim_v, + config.num_tokens_q, + config.num_tokens_kv, + ), + device=torch.device("cuda", torch.cuda.current_device()), + fp8=fp8, + mxfp8=mxfp8, + ) def spec(n, dim=(1, 1, 1, 1), stride=(1, 1, 1, 1), dtype="FLOAT", **kwargs): attn.specs[n] = _TensorSpec(tuple(dim), tuple(stride), dtype, **kwargs) @@ -677,41 +440,25 @@ def matrix(n, fmt, h, s, d, dtype, packed="", output=False): dim = (attn.b, h, s, d) if n in ("k", "v") and layout.startswith("paged_kv_"): dim = (getattr(config, "num_pages_" + n), h, getattr(config, "page_size_" + n), d) - stride = list(_strides(dim, fmt)) - if "3" in packed or "2" in packed: - count = 3 if "3" in packed else 2 - stride[0] *= count - stride[2] *= count - if "h3d" in packed or "h2d" in packed: - stride[1] *= count - offset = "" - if fmt == "thd": - offset = n + "_offset" + offset = n + "_offset" if fmt == "thd" else "" + if offset: spec(offset, (attn.b + 1, 1, 1, 1), dtype="INT64") - spec(n, dim, stride, dtype, output=output, ragged_offset=offset) + attn.specs[n] = attention_matrix_spec( + dim, fmt, dtype, packed=packed, output=output, ragged_offset=offset + ) def stats(n, output=False): - stride = ( - (attn.h * attn.sq, 1, attn.h, 1) - if attn.ragged_stats - else (attn.h * attn.sq, attn.sq, 1, 1) + if attn.ragged_stats: + spec("stats_offset", (attn.b + 1, 1, 1, 1), dtype="INT64") + attn.specs[n] = attention_stats_spec( + attn.b, attn.h, attn.sq, ragged=attn.ragged_stats, output=output ) - offset = "stats_offset" if attn.ragged_stats else "" - if offset: - spec(offset, (attn.b + 1, 1, 1, 1), dtype="INT64") - spec(n, (attn.b, attn.h, attn.sq, 1), stride, output=output, ragged_offset=offset) def scale(n, matrix_name, fmt, columnwise=False): if not attn.mxfp8: spec(n) - return - b, h, s, d = attn.specs[matrix_name].dim - dim = ( - (b, h, _round_up((s + 31) // 32, 4), _round_up(d, 128)) - if columnwise - else (b, h, _round_up(s, 128), _round_up((d + 31) // 32, 4)) - ) - spec(n, dim, _strides(dim, fmt), "FP8_E8M0", reordered=True) + else: + attn.specs[n] = mxfp8_scale_spec(attn.specs[matrix_name].dim, fmt, columnwise) parts = layout.removeprefix("paged_kv_").split("_") for n, fmt, h, s, d, part in ( @@ -743,9 +490,9 @@ def scale(n, matrix_name, fmt, columnwise=False): if backward and (b, h, sq) != (1, 1, 1): spec("dbias", dim, stride, types[DType.cast(config.qkv_dtype)], output=True) if attn.params["softmax_type"] != "vanilla": - spec("softmax_offset", (1, attn.h, 1, 1), (attn.h, 1, 1, 1)) + spec("sink", (1, attn.h, 1, 1), (attn.h, 1, 1, 1)) if backward: - spec("dsoftmax_offset", (1, attn.h, 1, 1), (attn.h, 1, 1, 1), output=True) + spec("dsink", (1, attn.h, 1, 1), (attn.h, 1, 1, 1), output=True) if layout.startswith("paged_kv_"): for n in ("k", "v"): pages = getattr(config, "max_pages_per_seq_" + n) @@ -991,7 +738,7 @@ def fused_attn_bwd(params): dsoftmax = torch.empty(0, device=attn.device, dtype=torch.float32) if p["softmax_type"] != "vanilla": dsoftmax = torch.empty_like(p["softmax_offset"]) - attn.add("dsoftmax_offset", dsoftmax, output=True) + attn.add("dsink", dsoftmax, output=True) if attn.fp8: do_scale_fmt = p.get("do_scale_inv_format") or p["do_format"] attn.scale("descale_do", _descale(p["d_o"]), "do", do_scale_fmt) From f564a84385739fb55c34fc70a3ba16de286eb8ce Mon Sep 17 00:00:00 2001 From: Vladimir Cherepanov Date: Thu, 15 Oct 2026 18:55:53 +0000 Subject: [PATCH 4/6] Fix cuDNN attention workspace, compile probe, and JAX dependencies Restore bucketed total-token bounds for FP16/BF16 packed backward graphs and forward Q and KV bounds independently to preserve mixed layouts. Extend graph reuse coverage to backward and compare packed and mixed layout outputs and gradients against FP32 SDPA. Update the torch.compile backend-selection regression to patch the Python support probe after removal of the C++ binding. Pin each JAX build's runtime cuDNN Frontend dependency to the distribution supplying its build headers while retaining the minimum supported version. Document the exact-version requirement and test generated dependency metadata for frontend versions 1.29.0 and 1.30.0. Validation: 382 compile tests passed, 24 skipped, and 1 xpassed; 34 focused attention tests and 3 dependency-metadata tests passed. Pre-commit, JAX Python lint, and changed-file license checks passed. JAX runtime validation was unavailable because JAX is not installed. Full PyTorch lint and repository license checks report existing issues in unrelated files and untracked nccl_ep headers. Signed-off-by: Vladimir Cherepanov --- build_tools/jax.py | 8 +- docs/installation.rst | 6 +- tests/jax/test_build_tools.py | 49 +++++++ tests/pytorch/attention/test_attention.py | 128 +++++++++++++++++- tests/pytorch/test_torch_compile.py | 7 +- transformer_engine/common/cudnn/attention.py | 1 + .../cpp_extensions/cudnn_fused_attn.py | 11 +- 7 files changed, 196 insertions(+), 14 deletions(-) create mode 100644 tests/jax/test_build_tools.py diff --git a/build_tools/jax.py b/build_tools/jax.py index 9de778721ce..56843e63c12 100644 --- a/build_tools/jax.py +++ b/build_tools/jax.py @@ -6,6 +6,7 @@ import os import warnings +from importlib.metadata import version as distribution_version from pathlib import Path from packaging import version @@ -27,7 +28,12 @@ def install_requirements() -> List[str]: """Install dependencies for TE/JAX extensions.""" - return ["jax", "flax>=0.7.1", "nvidia-cudnn-frontend>=1.29.0", "triton"] + # Python serializes graphs for C++ built with headers from this distribution. + # Keep the runtime serializer at the exact version used by this build. + frontend_version = distribution_version("nvidia-cudnn-frontend") + if version.parse(frontend_version) < version.parse("1.29.0"): + raise RuntimeError("Building TE/JAX requires nvidia-cudnn-frontend>=1.29.0.") + return ["jax", "flax>=0.7.1", f"nvidia-cudnn-frontend=={frontend_version}", "triton"] def test_requirements() -> List[str]: diff --git a/docs/installation.rst b/docs/installation.rst index 9da51c18bd2..e19aad3a427 100644 --- a/docs/installation.rst +++ b/docs/installation.rst @@ -17,8 +17,10 @@ Prerequisites 4. `cuDNN 9.12 `__ or later. The PyTorch and JAX cuDNN attention backends use shared graph builders through -the cuDNN frontend Python API. Both installations include -``nvidia-cudnn-frontend>=1.29.0``. Graph descriptions are cached by tensor metadata +the cuDNN frontend Python API. Both installations require +``nvidia-cudnn-frontend`` version 1.29.0 or later. JAX builds pin their runtime +dependency to the exact frontend version that supplies the build headers. +Graph descriptions are cached by tensor metadata and target device or architecture. Support queries select supported plans without compiling their kernels. PyTorch compiles on first execution and uses its CUDA generator for dropout, including CUDA graph capture and replay. JAX compiles diff --git a/tests/jax/test_build_tools.py b/tests/jax/test_build_tools.py new file mode 100644 index 00000000000..714aafec106 --- /dev/null +++ b/tests/jax/test_build_tools.py @@ -0,0 +1,49 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""Dependency metadata for builds that deserialize cuDNN frontend graphs.""" + +from email.parser import Parser +from io import StringIO + +import pytest +from packaging.requirements import Requirement +from setuptools import Distribution + +from build_tools import jax as jax_build + + +@pytest.mark.parametrize("frontend_version", ("1.29.0", "1.30.0")) +def test_jax_metadata_pins_build_frontend(monkeypatch, frontend_version): + """Built metadata must reject runtime serializers other than the build version.""" + + def distribution_version(package): + assert package == "nvidia-cudnn-frontend" + return frontend_version + + monkeypatch.setattr(jax_build, "distribution_version", distribution_version) + dist = Distribution( + { + "name": "transformer_engine_jax", + "version": "0.0.0", + "install_requires": jax_build.install_requirements(), + } + ) + metadata = StringIO() + dist.metadata.write_pkg_file(metadata) + requirements = [ + Requirement(value) + for value in Parser().parsestr(metadata.getvalue()).get_all("Requires-Dist") + ] + frontend = next(req for req in requirements if req.name == "nvidia-cudnn-frontend") + assert frontend.specifier.contains(frontend_version) + for other_version in ("1.28.0", "1.29.0", "1.29.1", "1.30.0", "1.30.1", "1.31.0"): + assert frontend.specifier.contains(other_version) == (other_version == frontend_version) + + +def test_jax_build_requires_minimum_frontend(monkeypatch): + """Pinning the installed version must preserve the minimum build requirement.""" + monkeypatch.setattr(jax_build, "distribution_version", lambda _package: "1.28.0") + with pytest.raises(RuntimeError, match="nvidia-cudnn-frontend>=1.29.0"): + jax_build.install_requirements() diff --git a/tests/pytorch/attention/test_attention.py b/tests/pytorch/attention/test_attention.py index a3b6a96c8c1..3b719dce0d3 100644 --- a/tests/pytorch/attention/test_attention.py +++ b/tests/pytorch/attention/test_attention.py @@ -11,6 +11,7 @@ import copy from dataclasses import replace from typing import Any, Dict, List, NamedTuple, Tuple, Union +from unittest.mock import Mock from packaging.version import Version as PkgVersion import pytest @@ -252,7 +253,8 @@ def common_attention(*_args, **_kwargs): @pytest.mark.parametrize("layout", ("bs3hd", "bsh3d", "sb3hd", "sbh3d", "t3hd", "th3d")) -def test_fused_attn_query_reuses_runtime_graph(layout, monkeypatch): +@pytest.mark.parametrize("backward", (False, True)) +def test_fused_attn_query_reuses_runtime_graph(layout, backward, monkeypatch): """Packed storage and runtime scalar values must preserve support-query graph identity.""" from transformer_engine.pytorch.cpp_extensions import cudnn_fused_attn as implementation @@ -297,13 +299,18 @@ def test_fused_attn_query_reuses_runtime_graph(layout, monkeypatch): num_tokens_kv=batch * seqlen if fmt == "thd" else 0, attn_mask_type=AttnMaskType[params["attn_mask_type"]], ) - query = implementation._query_graph(config, False) - runtime = implementation._AttentionGraph(params) + query = implementation._query_graph(config, backward) + runtime = implementation._AttentionGraph(params, backward) implementation._prepare_inputs(runtime) out = torch.empty(base_shape, device="cuda", dtype=torch.bfloat16) - runtime.matrix("o", out, fmt, heads, runtime.sq, cu, output=True) + runtime.matrix("o", out, fmt, heads, runtime.sq, cu, output=not backward) stats_shape = (batch * seqlen, heads, 1) if runtime.ragged_stats else (batch, heads, seqlen, 1) - runtime.stats("stats", torch.empty(stats_shape, device="cuda"), output=True) + runtime.stats("stats", torch.empty(stats_shape, device="cuda"), output=not backward) + if backward: + runtime.matrix("do", torch.empty_like(out), fmt, heads, runtime.sq, cu) + grads = torch.empty(packed_shape, device="cuda", dtype=torch.bfloat16).unbind(axis) + for name, grad in zip(("dq", "dk", "dv"), grads): + runtime.matrix(name, grad, fmt, heads, runtime.sq, cu, output=True) assert query.description() == runtime.description() # Isolate the cache to demonstrate that a query never compiles, and that a @@ -333,6 +340,117 @@ def compile_plan(_cudnn, graph, _run): assert queried.workspace_size == 17 +@pytest.mark.parametrize("dtype", (DType.kFloat16, DType.kBFloat16)) +@pytest.mark.parametrize("capability", ((10, 0), (12, 0))) +@pytest.mark.parametrize("layout", ("bshd_bshd_bshd", "thd_bshd_bshd", "thd_thd_thd")) +def test_fused_attn_backward_token_bounds(monkeypatch, dtype, capability, layout): + """Bound backward workspace by each ragged input's bucketed token capacity.""" + import cudnn + from transformer_engine.pytorch.cpp_extensions import cudnn_fused_attn as implementation + + monkeypatch.setattr(torch.cuda, "get_device_capability", lambda *_args: capability) + config = FusedAttentionParams( + qkv_layout=QKVLayout[layout], + qkv_dtype=dtype, + o_dtype=dtype, + do_dtype=dtype, + dqkv_dtype=dtype, + batch_size=4, + num_attn_heads=4, + num_gqa_groups=4, + head_dim_qk=64, + head_dim_v=64, + max_seqlen_q=512, + max_seqlen_kv=1024, + num_tokens_q=1153, + num_tokens_kv=2305, + attn_mask_type=AttnMaskType["padding"], + ) + query = implementation._query_graph(config, True) + graph = Mock() + graph.sdpa_backward.return_value = (Mock(), Mock(), Mock()) + implementation.build_attention_graph(cudnn, graph, query.description()) + bounds = { + name: value + for name, value in graph.sdpa_backward.call_args.kwargs.items() + if name.startswith("max_total_seq_len_") + } + expected = {} + if capability == (10, 0) and layout.startswith("thd"): + expected["max_total_seq_len_q"] = 2048 + if layout == "thd_thd_thd": + expected["max_total_seq_len_kv"] = 4096 + assert bounds == expected + + +@pytest.mark.parametrize("dtype", (torch.float16, torch.bfloat16)) +@pytest.mark.parametrize("kv_format", ("thd", "bshd", "sbhd")) +def test_fused_attn_packed_backward(dtype, kv_format): + """Packed backward must match FP32 SDPA with unequal lengths and padded capacities.""" + if device_compute_capability < (9, 0): + pytest.skip("FP16/BF16 THD fused attention requires SM90 or later") + q_lengths, kv_lengths = (64, 128, 32), (96, 160, 64) + q_offsets = (0, 128, 384, 512) + kv_offsets = (0, 384, 768, 1152) if kv_format == "thd" else (0, 256, 512, 768) + cu_q = torch.tensor((0, 64, 192, 224), device="cuda", dtype=torch.int32) + cu_kv = torch.tensor((0, 96, 256, 320), device="cuda", dtype=torch.int32) + q = torch.randn(512, 4, 64, device="cuda", dtype=dtype) + kv_shape = {"thd": (1152, 4, 64), "bshd": (3, 256, 4, 64), "sbhd": (256, 3, 4, 64)} + k, v = [torch.randn(kv_shape[kv_format], device="cuda", dtype=dtype) for _ in range(2)] + layout = f"thd_{kv_format}_{kv_format}" + kwargs = dict( + max_seqlen_q=128, + max_seqlen_kv=256, + cu_seqlens_q=cu_q, + cu_seqlens_kv=cu_kv, + cu_seqlens_q_padded=torch.tensor(q_offsets, device="cuda", dtype=torch.int32), + cu_seqlens_kv_padded=torch.tensor(kv_offsets, device="cuda", dtype=torch.int32), + q=q, + k=k, + v=v, + fake_dtype=dtype, + fused_attention_backend=FusedAttnBackend.F16_arbitrary_seqlen, + qkv_layout=layout, + o_format="thd", + attn_mask_type="padding", + ) + out, aux = fused_attn_fwd(is_training=True, **kwargs) + grad = torch.randn_like(out) + dq, dk, dv, *_ = fused_attn_bwd( + o=out, d_o=grad, aux_ctx_tensors=aux, do_format="thd", dqkv_layout=layout, **kwargs + ) + refs = [tensor.float().requires_grad_() for tensor in (q, k, v)] + reference = torch.zeros_like(refs[0]) + for index, (q_len, kv_len) in enumerate(zip(q_lengths, kv_lengths)): + q_slice = slice(q_offsets[index], q_offsets[index] + q_len) + if kv_format == "thd": + kv_slice = slice(kv_offsets[index], kv_offsets[index] + kv_len) + kr, vr = (tensor[kv_slice] for tensor in refs[1:]) + elif kv_format == "bshd": + kr, vr = (tensor[index, :kv_len] for tensor in refs[1:]) + else: + kr, vr = (tensor[:kv_len, index] for tensor in refs[1:]) + result = torch.nn.functional.scaled_dot_product_attention( + refs[0][q_slice].transpose(0, 1), kr.transpose(0, 1), vr.transpose(0, 1) + ) + reference[q_slice] = result.transpose(0, 1) + reference.backward(grad.float()) + rtol, atol = (2e-2, 1e-2) if dtype == torch.bfloat16 else (3e-3, 2e-3) + torch.testing.assert_close(out.float(), reference, rtol=rtol, atol=atol) + for actual, ref in zip((dq, dk, dv), refs): + if actual.ndim == 3: + torch.testing.assert_close(actual.float(), ref.grad, rtol=rtol, atol=atol) + else: + # cuDNN may leave dense KV padding gradients unwritten. + for index, length in enumerate(kv_lengths): + active = ( + (index, slice(0, length)) if kv_format == "bshd" else (slice(0, length), index) + ) + torch.testing.assert_close( + actual[active].float(), ref.grad[active], rtol=rtol, atol=atol + ) + + @pytest.mark.parametrize("error", (RuntimeError("device failure"), ImportError("frontend missing"))) def test_fused_attn_query_propagates_setup_errors(error, monkeypatch): """A failed setup must surface as an error instead of selecting a fallback backend.""" diff --git a/tests/pytorch/test_torch_compile.py b/tests/pytorch/test_torch_compile.py index 472917d56f0..cf0195e910d 100644 --- a/tests/pytorch/test_torch_compile.py +++ b/tests/pytorch/test_torch_compile.py @@ -1316,9 +1316,10 @@ def test_get_attention_backend_traceable(monkeypatch): """get_attention_backend must trace under torch.compile(fullgraph=True) without graph breaks. The compiled selection must stay consistent with eager when NVTE_* env vars flip (dynamo guards on os.environ) and when - attention params change, and the baked tex.get_fused_attn_backend result + attention params change, and the baked Python get_fused_attn_backend result must drive the selection.""" from transformer_engine.pytorch.attention.dot_product_attention import utils as dpa_utils + from transformer_engine.pytorch.cpp_extensions import cudnn_fused_attn def fn(x, params): ( @@ -1396,9 +1397,9 @@ def fn(x, params): # guard on the wrapped function). monkeypatch.setenv("NVTE_FLASH_ATTN", "0") monkeypatch.setattr( - dpa_utils.tex, + cudnn_fused_attn, "get_fused_attn_backend", - lambda *args: (tex.NVTE_Fused_Attn_Backend.NVTE_No_Backend, "disabled by test"), + lambda _config: (dpa_utils.FusedAttnBackend.No_Backend, "disabled by test"), ) def fn_no_backend(x, params): diff --git a/transformer_engine/common/cudnn/attention.py b/transformer_engine/common/cudnn/attention.py index e5f872667e7..e14f8e377cf 100644 --- a/transformer_engine/common/cudnn/attention.py +++ b/transformer_engine/common/cudnn/attention.py @@ -113,6 +113,7 @@ def tensor(name): kwargs["use_deterministic_algorithm"] = options.deterministic if options.max_total_seqlen_q: kwargs["max_total_seq_len_q"] = options.max_total_seqlen_q + if options.max_total_seqlen_kv: kwargs["max_total_seq_len_kv"] = options.max_total_seqlen_kv else: kwargs["generate_stats"] = options.generate_stats diff --git a/transformer_engine/pytorch/cpp_extensions/cudnn_fused_attn.py b/transformer_engine/pytorch/cpp_extensions/cudnn_fused_attn.py index a05297763e1..1812344b094 100644 --- a/transformer_engine/pytorch/cpp_extensions/cudnn_fused_attn.py +++ b/transformer_engine/pytorch/cpp_extensions/cudnn_fused_attn.py @@ -137,12 +137,12 @@ def __init__( ) self.b, self.h, self.hg, self.dq, self.dv, self.tq, self.tk = dimensions capability = torch.cuda.get_device_capability(self.device) - ragged_graph = capability >= (9, 0) and capability != (12, 0) - self.ragged_stats = self.q_format == "thd" and ragged_graph + self.ragged_graph = capability >= (9, 0) and capability != (12, 0) + self.ragged_stats = self.q_format == "thd" and self.ragged_graph self.sq = bucket(self.tq, 1024, 32768) if self.ragged_stats else params["max_seqlen_q"] self.skv = ( bucket(self.tk, 1024, 32768) - if self.kv_format == "thd" and ragged_graph + if self.kv_format == "thd" and self.ragged_graph else params["max_seqlen_kv"] ) self.fp8 = hasattr(params["q"], "_fp8_dtype") if fp8 is None else fp8 @@ -219,6 +219,7 @@ def options(self): left, right = p["window_size"] or (-1, -1) if not causal and (left, right) == (-1, -1): bottom_right = False + bounded_backward = self.backward and not self.fp8 and self.ragged_graph return AttentionOptions( backward=self.backward, fp8=self.fp8, @@ -231,6 +232,10 @@ def options(self): dropout=p["dropout"], deterministic=p.get("deterministic", False), paged_max_seqlen_kv=p["max_seqlen_kv"] if p.get("page_table_k") is not None else 0, + # Ragged sequence dimensions are bucketed total-token capacities. + # Bound cuDNN's backward scratch buffers by those same capacities. + max_total_seqlen_q=self.sq if bounded_backward and self.q_format == "thd" else 0, + max_total_seqlen_kv=self.skv if bounded_backward and self.kv_format == "thd" else 0, ) def description(self): From e50b3cfe9ced13a6b550491bcd907b6b39bd2596 Mon Sep 17 00:00:00 2001 From: Vladimir Cherepanov Date: Thu, 15 Oct 2026 19:43:27 +0000 Subject: [PATCH 5/6] Fix score-mod callback caching and SM12x attention layouts Address the remaining issues from PR #3289 in the shared cuDNN graph paths. Key plain callbacks on their code and typed snapshots of positional and keyword defaults so callbacks with different graph structure cannot reuse the same cached plan. Disable implicit caching for tensor, opaque, or recursive defaults while retaining the explicit structural-key protocol. Apply dense per-sequence statistics and dimensions throughout the SM12x family in common C++, PyTorch, and JAX. Share the Python layout predicate with context-parallel LSE handling and extend the related backend gates. Add forward/backward cache-lookup regressions for code and default changes, SM121/SM129 layout and eligibility coverage, an SM130 control, and matching JAX configuration tests. Document the caching and statistics conventions. Validation: - Rebuild the editable PyTorch installation with its cached configuration. - Compile suite: 382 passed, 24 skipped, 1 xpassed. - Focused attention regressions: 93 passed. - Common C API: four FP16/BF16 dense/packed support queries passed. - Two-GB200 THD context-parallel forward/backward regression passed. - Pre-commit, C++ lint, JAX Python lint, and changed-file license checks passed. JAX and physical SM121 runtime validation were unavailable. PyTorch Python lint still reports the pre-existing import-order issue in fused_mla_q_uproj.py. Signed-off-by: Vladimir Cherepanov --- docs/installation.rst | 11 ++ tests/jax/test_fused_attn.py | 42 +++++++ tests/jax/test_fused_attn_score_mod.py | 29 +++++ tests/pytorch/attention/test_attention.py | 12 +- .../test_attention_backend_selection.py | 41 +++++++ .../pytorch/attention/test_flex_attention.py | 105 ++++++++++++++++++ transformer_engine/common/cudnn/graph.py | 5 + transformer_engine/common/cudnn/score_mod.py | 41 ++++++- .../common/fused_attn/config_and_params.cpp | 3 +- .../jax/cpp_extensions/cudnn_attention.py | 5 +- .../dot_product_attention/context_parallel.py | 7 +- .../attention/dot_product_attention/utils.py | 33 +++--- .../cpp_extensions/cudnn_fused_attn.py | 3 +- 13 files changed, 306 insertions(+), 31 deletions(-) diff --git a/docs/installation.rst b/docs/installation.rst index e19aad3a427..9fac9ee56a6 100644 --- a/docs/installation.rst +++ b/docs/installation.rst @@ -29,6 +29,17 @@ native FFI executor. The JAX extension and the frontend Python package must use matching cuDNN frontend versions; rebuild the extension after upgrading that package. +Automatic ``score_mod`` graph caching includes the callback's code and snapshots +of builtin scalar and container defaults, including keyword-only defaults. +Callbacks with tensor or opaque stateful defaults bypass automatic caching. +Stateful callbacks can provide ``score_mod_graph_cache_key`` to explicitly +describe their graph structure; pass runtime data through ``score_mod_tensors`` +and ``score_mod_bprop_tensors``. + +On SM12x GPUs, including SM120 and SM121, packed attention uses dense +per-sequence softmax statistics internally. Graph construction and context +parallelism use the same statistics layout. + If the CUDA Toolkit headers are not available at runtime in a standard installation path, e.g. within `CUDA_HOME`, set `NVTE_CUDA_INCLUDE_PATH` in the environment. diff --git a/tests/jax/test_fused_attn.py b/tests/jax/test_fused_attn.py index 6da031c353f..14d7bd972cb 100644 --- a/tests/jax/test_fused_attn.py +++ b/tests/jax/test_fused_attn.py @@ -503,6 +503,48 @@ def test_cudnn_attention_empty_ragged_capacities(batch, seqlen, monkeypatch): assert dict(description.spec.tensors)[name].dim == (32, 4, 1024, 64) +@pytest.mark.parametrize("sm_arch", (100, 120, 121, 129, 130)) +@pytest.mark.parametrize("direction", ("fwd", "bwd")) +def test_cudnn_attention_sm12x_dense_statistics(sm_arch, direction, monkeypatch): + """All SM12x variants use dense per-segment statistics and per-sequence dimensions.""" + from types import SimpleNamespace + from transformer_engine.jax.cpp_extensions import cudnn_attention + + monkeypatch.setattr(cudnn_attention, "cudnn_version", lambda: 92700) + config = SimpleNamespace( + qkv_layout=QKVLayout.THD_THD_THD, + max_segments_per_seq=2, + fp8=False, + attn_mask_type=AttnMaskType.PADDING_MASK, + attn_bias_type=AttnBiasType.NO_BIAS, + softmax_type=AttnSoftmaxType.VANILLA_SOFTMAX, + effective_window_size=(-1, -1), + scaling_factor=0.125, + bottom_right_diagonal=False, + dropout_probability=0.0, + is_training=True, + return_max_logit=False, + ) + shapes = tuple( + (shape, jnp.dtype(jnp.bfloat16)) for shape in ((2, 128, 4, 64),) * 3 + ((0,), (1, 4, 1, 1)) + ) + description = cudnn_attention._describe_graph(direction, shapes, config, (sm_arch,), True, True) + tensors = dict(description.spec.tensors) + if 120 <= sm_arch < 130: + assert not description.ragged_stats + assert description.metadata_batch == 4 + assert tensors["q"].dim == (4, 4, 128, 64) + assert tensors["stats"].dim == (4, 4, 128, 1) + assert tensors["stats"].ragged_offset == "" + assert description.spec.options.max_total_seqlen_q == 0 + assert description.spec.options.max_total_seqlen_kv == 0 + else: + assert description.ragged_stats + assert description.metadata_batch == 32 + assert tensors["q"].dim == (32, 4, 1024, 64) + assert tensors["stats"].ragged_offset == "offset_stats" + + class BiasShape(Enum): """ Enum class to represent the different bias shapes used in the fused attention. diff --git a/tests/jax/test_fused_attn_score_mod.py b/tests/jax/test_fused_attn_score_mod.py index 4be2c2599a7..c4050f3b0fd 100644 --- a/tests/jax/test_fused_attn_score_mod.py +++ b/tests/jax/test_fused_attn_score_mod.py @@ -4,6 +4,7 @@ """Tests for cuDNN frontend score_mod fused attention.""" from math import sqrt +from types import FunctionType import jax import jax.numpy as jnp @@ -686,6 +687,34 @@ def test_fused_attn_score_mod_module_lambda_keys_do_not_collide(): ) +@pytest.mark.parametrize("keyword", (False, True)) +def test_fused_attn_score_mod_config_separates_function_defaults(keyword): + """Defaults must participate in both the static config and serialized-graph cache key.""" + first = FunctionType(_default_mode_score_mod.__code__, globals(), argdefs=(False,)) + second = FunctionType(_default_mode_score_mod.__code__, globals(), argdefs=(False,)) + first.__kwdefaults__ = second.__kwdefaults__ = {"keyword_mode": False} + if keyword: + second.__kwdefaults__ = {"keyword_mode": True} + else: + first.__defaults__ = (False,) + second.__defaults__ = (True,) + configs = [ + make_fused_attn_score_mod_config( + callback, None, None, None, _CONFIG_TEST_SCALING_FACTOR, True + )[0] + for callback in (first, first, second) + ] + assert configs[0] == configs[1] + assert configs[0] != configs[2] + assert tex_attention._graph_cache_key("fwd", configs[0], ()) != ( + tex_attention._graph_cache_key("fwd", configs[2], ()) + ) + + +def _default_mode_score_mod(graph, score, _tensors, mode=False, *, keyword_mode=False): + return graph.tanh(input=score) if mode or keyword_mode else score + + def test_fused_attn_score_mod_config_stabilizes_bound_method_cache_keys(): """Bound score_mod methods with graph keys share stable cache keys.""" softcap_score_mod = _ScoreModSoftcap() diff --git a/tests/pytorch/attention/test_attention.py b/tests/pytorch/attention/test_attention.py index 3b719dce0d3..4d4aa24d96d 100644 --- a/tests/pytorch/attention/test_attention.py +++ b/tests/pytorch/attention/test_attention.py @@ -341,7 +341,7 @@ def compile_plan(_cudnn, graph, _run): @pytest.mark.parametrize("dtype", (DType.kFloat16, DType.kBFloat16)) -@pytest.mark.parametrize("capability", ((10, 0), (12, 0))) +@pytest.mark.parametrize("capability", ((10, 0), (12, 0), (12, 1), (12, 9), (13, 0))) @pytest.mark.parametrize("layout", ("bshd_bshd_bshd", "thd_bshd_bshd", "thd_thd_thd")) def test_fused_attn_backward_token_bounds(monkeypatch, dtype, capability, layout): """Bound backward workspace by each ragged input's bucketed token capacity.""" @@ -376,11 +376,19 @@ def test_fused_attn_backward_token_bounds(monkeypatch, dtype, capability, layout if name.startswith("max_total_seq_len_") } expected = {} - if capability == (10, 0) and layout.startswith("thd"): + if capability[0] != 12 and layout.startswith("thd"): expected["max_total_seq_len_q"] = 2048 if layout == "thd_thd_thd": expected["max_total_seq_len_kv"] = 4096 assert bounds == expected + # cuDNN's graph layout and context-parallel LSE interpretation must agree. + from transformer_engine.pytorch.attention.dot_product_attention import context_parallel + + packed = context_parallel.uses_ragged_attention_graph(10 * capability[0] + capability[1]) + assert bool(query.specs["stats"].ragged_offset) == (packed and layout.startswith("thd")) + if capability[0] == 12: + assert query.sq == 512 + assert query.skv == 1024 @pytest.mark.parametrize("dtype", (torch.float16, torch.bfloat16)) diff --git a/tests/pytorch/attention/test_attention_backend_selection.py b/tests/pytorch/attention/test_attention_backend_selection.py index c2944eece3b..943152f79cf 100644 --- a/tests/pytorch/attention/test_attention_backend_selection.py +++ b/tests/pytorch/attention/test_attention_backend_selection.py @@ -5,10 +5,12 @@ """Focused FA4 eligibility and backend-cache regressions.""" from unittest.mock import Mock +from types import SimpleNamespace import pytest import torch from packaging.version import Version as PkgVersion +from transformer_engine.common import recipe from transformer_engine.pytorch import DotProductAttention from transformer_engine.pytorch.attention.dot_product_attention import ( @@ -82,6 +84,45 @@ def test_fa4_deterministic_backend_selection( assert available == [expect_flash, False, unfused] +@pytest.mark.parametrize("capability", ((12, 0), (12, 1), (12, 9))) +@pytest.mark.parametrize("case", ("fp8", "kv_cache", "mla", "thd_version", "thd_packed")) +def test_sm12x_fused_backend_restrictions(monkeypatch, capability, case): + """SM121 and other SM12x variants must retain the SM120 backend restrictions.""" + monkeypatch.setenv("NVTE_FUSED_ATTN", "1") + monkeypatch.setenv("NVTE_FLASH_ATTN", "0") + monkeypatch.setattr(dpa_utils, "get_device_compute_capability", lambda _device: capability) + version = (9, 18, 0) if case == "thd_version" else (9, 27, 0) + monkeypatch.setattr(dpa_utils, "get_cudnn_version", lambda: version) + probe = Mock(return_value=(dpa_utils.FusedAttnBackend.F16_arbitrary_seqlen.value, "")) + monkeypatch.setattr(dpa_utils, "_get_fused_attn_backend", probe) + options = dict( + device=torch.device("cuda:0"), + qkv_dtype=torch.bfloat16, + qkv_layout="bshd_bshd_bshd", + num_heads=4, + num_gqa_groups=4, + head_dim_qk=64, + head_dim_v=64, + deterministic=False, + ) + if case == "fp8": + options.update(fp8=True, fp8_meta={"recipe": recipe.DelayedScaling(fp8_dpa=True)}) + elif case == "kv_cache": + options.update(inference_params=SimpleNamespace(is_paged=False), is_training=False) + elif case == "mla": + options.update(head_dim_qk=192, head_dim_v=128) + else: + options.update( + qkv_layout="t3hd" if case == "thd_packed" else "thd_thd_thd", + attn_mask_type="padding", + num_tokens_q=128, + num_tokens_kv=128, + ) + result = dpa_utils.get_attention_backend(dpa_utils.AttentionParams(**options)) + assert not result[2] + probe.assert_not_called() + + @pytest.mark.parametrize("changed_field", ["device", "requires_backward"]) def test_thd_backend_cache_tracks_execution_context(monkeypatch, changed_field): selector = Mock(return_value=(False, None, False, None, True, [False, False, True])) diff --git a/tests/pytorch/attention/test_flex_attention.py b/tests/pytorch/attention/test_flex_attention.py index beed4069917..c40aaf2d0af 100644 --- a/tests/pytorch/attention/test_flex_attention.py +++ b/tests/pytorch/attention/test_flex_attention.py @@ -5,6 +5,8 @@ import os import sys import pathlib +from types import FunctionType +from unittest.mock import Mock import pytest import torch @@ -156,6 +158,109 @@ def _score_mod_cache_cpu_inputs(): return q, k, v, o, stats +def _score_mod_with_defaults(graph, score, _tensors, mode=False, *, options=None): + if options is not None: + mode = options["modes"][0] + return graph.tanh(input=score) if mode else score + + +def _clone_score_mod(default=False, options=None): + callback = FunctionType(_score_mod_with_defaults.__code__, globals(), argdefs=(default,)) + callback.__kwdefaults__ = {"options": options} + return callback + + +@pytest.mark.parametrize("direction", ("fwd", "bwd")) +@pytest.mark.parametrize("difference", ("positional", "keyword", "code", "type")) +def test_score_mod_cache_separates_function_defaults(monkeypatch, direction, difference): + """Different code or defaults must build separate forward and backward graphs.""" + first = _clone_score_mod() + second = _clone_score_mod(True) + if difference == "keyword": + second = _clone_score_mod(options={"modes": [True]}) + elif difference == "code": + second = FunctionType(_score_mod_identity_bprop.__code__, globals()) + second.__qualname__ = first.__qualname__ + elif difference == "type": + second = _clone_score_mod(0) + q, k, v, out, stats = _score_mod_cache_cpu_inputs() + monkeypatch.setattr(flex_attention, "_cudnn_score_mod_graph_cache", {}) + build = Mock(side_effect=lambda *_args: object()) + monkeypatch.setattr(flex_attention, f"_build_cudnn_score_mod_{direction}_graph", build) + + def get(callback): + if direction == "fwd": + return flex_attention._get_cudnn_score_mod_fwd_graph( + True, q, k, v, "bshd", "bshd", 1.0, callback, {}, out, stats + ) + return flex_attention._get_cudnn_score_mod_bwd_graph( + q, + k, + v, + out, + out, + stats, + "bshd", + "bshd", + 1.0, + _score_mod_identity_bprop, + callback, + {}, + {}, + False, + ) + + entry = get(first) + assert get(first) is entry + other = get(second) + assert other is not entry + assert get(second) is other + assert build.call_count == 2 + + # Keys snapshot nested defaults rather than retaining mutable references. + first.__kwdefaults__["options"] = {"modes": [False]} + changed = get(first) + assert changed is not entry + first.__kwdefaults__["options"]["modes"].append(True) + assert get(first) is not changed + assert build.call_count == 4 + + +@pytest.mark.parametrize("kind", ("tensor", "opaque", "cycle")) +def test_score_mod_cache_unsafe_defaults_are_uncached(kind): + """Opaque state, runtime tensors and recursive defaults cannot imply a stable graph.""" + if kind == "tensor": + default = torch.tensor(1.0) + elif kind == "opaque": + default = object() + else: + default = [] + default.append(default) + callback = _clone_score_mod(default) + key = flex_attention._score_mod_callback_cache_key(callback) + assert key is flex_attention._SCORE_MOD_UNCACHEABLE + + # An explicit structural key remains the caller's opt-in caching contract. + callback.score_mod_graph_cache_key = lambda: ("explicit-mode", True) + assert flex_attention._score_mod_callback_cache_key(callback) is not key + + +def test_score_mod_cache_snapshots_positional_defaults(): + """Changing positional defaults invalidates the old key while unchanged NaNs stay stable.""" + callback = _clone_score_mod(False) + before = flex_attention._score_mod_callback_cache_key(callback) + callback.__defaults__ = (True,) + assert flex_attention._score_mod_callback_cache_key(callback) != before + callback.__defaults__ = (float("nan"),) + assert flex_attention._score_mod_callback_cache_key(callback) == ( + flex_attention._score_mod_callback_cache_key(callback) + ) + callback.__defaults__ = (0.0,) + positive_zero = flex_attention._score_mod_callback_cache_key(callback) + callback.__defaults__ = (-0.0,) + assert flex_attention._score_mod_callback_cache_key(callback) != positive_zero + + def test_score_mod_cache_bound_method_requires_explicit_key(): """Unkeyed bound methods should be uncached instead of keyed by object id.""" diff --git a/transformer_engine/common/cudnn/graph.py b/transformer_engine/common/cudnn/graph.py index a60ea826f06..6393a476924 100644 --- a/transformer_engine/common/cudnn/graph.py +++ b/transformer_engine/common/cudnn/graph.py @@ -19,6 +19,11 @@ class UnsupportedGraph(ValueError): """A configuration rejected by cuDNN, rather than an execution/setup failure.""" +def uses_ragged_attention_graph(sm_arch): + """Use packed statistics on SM90+ except the SM12x family (encoded as 120-129).""" + return sm_arch >= 90 and not 120 <= sm_arch < 130 + + def import_cudnn(): """Import the optional graph frontend without importing either framework.""" try: diff --git a/transformer_engine/common/cudnn/score_mod.py b/transformer_engine/common/cudnn/score_mod.py index f385a569d6a..e48e35626cc 100644 --- a/transformer_engine/common/cudnn/score_mod.py +++ b/transformer_engine/common/cudnn/score_mod.py @@ -6,6 +6,7 @@ from collections.abc import Mapping import inspect +import struct class UncacheableKey: @@ -51,8 +52,29 @@ def explicit_cache_key(owner, is_tensor): return freeze_cache_key(key() if callable(key) else key, is_tensor) +def _freeze_function_default(value, is_tensor): + """Snapshot builtin defaults without trusting opaque objects' identity hashes.""" + value_type = type(value) + if is_tensor(value): + raise TypeError("Runtime tensors cannot be implicit callback cache keys.") + if value_type in (type(None), bool, int, str, bytes): + return value_type, value + if value_type is float: + return value_type, struct.pack("!d", value) + if value_type is complex: + return value_type, struct.pack("!dd", value.real, value.imag) + if value_type is dict: + return value_type, tuple( + (_freeze_function_default(key, is_tensor), _freeze_function_default(item, is_tensor)) + for key, item in value.items() + ) + if value_type in (list, tuple, set, frozenset): + return value_type, tuple(_freeze_function_default(item, is_tensor) for item in value) + raise TypeError("Opaque callback defaults require an explicit structural cache key.") + + def callback_cache_key(callback, is_tensor, uncacheable): - """Cache stable functions and explicitly keyed owners; distinguish anonymous code.""" + """Cache stateless functions by code/defaults and explicitly keyed callback owners.""" if callback is None: return None owner = getattr(callback, "__self__", None) @@ -76,12 +98,19 @@ def callback_cache_key(callback, is_tensor, uncacheable): and callback.__closure__ is None and "" not in callback.__qualname__ ): - identity = ( - callback.__code__ - if callback.__name__ == "" or not callback.__qualname__ - else callback.__qualname__ + try: + defaults = _freeze_function_default(callback.__defaults__, is_tensor) + kwdefaults = _freeze_function_default(callback.__kwdefaults__, is_tensor) + except (TypeError, RecursionError): + return uncacheable() + return ( + "function", + callback.__module__, + callback.__qualname__, + callback.__code__, + defaults, + kwdefaults, ) - return ("function", callback.__module__, identity) return uncacheable() diff --git a/transformer_engine/common/fused_attn/config_and_params.cpp b/transformer_engine/common/fused_attn/config_and_params.cpp index b12dbf30f61..f27859791e8 100644 --- a/transformer_engine/common/fused_attn/config_and_params.cpp +++ b/transformer_engine/common/fused_attn/config_and_params.cpp @@ -106,7 +106,8 @@ void FusedAttnConfig::derive() { // Use ragged (TH1) or dense (BHS1) graphs and stats const int sm_arch = cuda::sm_arch(cuda::current_device()); - uses_ragged_graph = sm_arch >= 90 && sm_arch != 120; + const bool is_sm12x = sm_arch >= 120 && sm_arch < 130; + uses_ragged_graph = sm_arch >= 90 && !is_sm12x; uses_ragged_stats = is_ragged_q && uses_ragged_graph; const bool buckets_the_batch = (is_ragged_q || is_ragged_kv) && uses_ragged_graph; graph_batch_size_fwd = diff --git a/transformer_engine/jax/cpp_extensions/cudnn_attention.py b/transformer_engine/jax/cpp_extensions/cudnn_attention.py index ed770b28554..054649e8585 100644 --- a/transformer_engine/jax/cpp_extensions/cudnn_attention.py +++ b/transformer_engine/jax/cpp_extensions/cudnn_attention.py @@ -28,6 +28,7 @@ create_graph, check_graph_support, compile_graph, + uses_ragged_attention_graph, ) from .cudnn_graph import import_cudnn, check_frontend_version, serialize_graph, graph_attributes @@ -226,7 +227,7 @@ def _describe_graph(direction, shapes, config, capabilities, nondeterministic, f batch, sq, sk, h, hg, dq, dv = _dims(avals, config) backward = direction == "bwd" ragged = config.qkv_layout.is_thd() - ragged_stats = ragged and capabilities[0] >= 90 and capabilities[0] != 120 + ragged_stats = ragged and uses_ragged_attention_graph(capabilities[0]) segments = batch * config.max_segments_per_seq if ragged else batch # JAX uses fixed metadata capacities, including for empty ragged storage. # Keep positive graph capacities while runtime segment lengths may be zero. @@ -425,7 +426,7 @@ def _metadata(value, size, fill=0): def _dense_stats_to_packed(value, offsets, lengths, shape): - """Convert SM120's per-segment BHS1 statistics to JAX's packed TH1 layout.""" + """Convert SM12x's per-segment BHS1 statistics to JAX's packed TH1 layout.""" tokens, heads = math.prod(shape[:-2]), shape[-2] rows = value.shape[-2] indices = offsets[: lengths.size, None] + jnp.arange(rows)[None, :] diff --git a/transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py b/transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py index 3a153a3f994..bf924de722f 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py @@ -8,6 +8,8 @@ import torch import transformer_engine_torch as tex +from transformer_engine.common.cudnn.graph import uses_ragged_attention_graph + from transformer_engine.pytorch.utils import ( get_cudnn_version, nvtx_range_pop, @@ -1837,7 +1839,10 @@ def forward( softmax_lse_in_packed_format = False if qkv_format == "thd": if use_fused_attention: - softmax_lse_in_packed_format = get_device_compute_capability() != (12, 0) + capability = get_device_compute_capability() + softmax_lse_in_packed_format = uses_ragged_attention_graph( + 10 * capability[0] + capability[1] + ) else: softmax_lse_in_packed_format = ( fa_utils.v2_6_0_plus or use_flash_attn_3 or use_flash_attn_4 diff --git a/transformer_engine/pytorch/attention/dot_product_attention/utils.py b/transformer_engine/pytorch/attention/dot_product_attention/utils.py index 06e016bfa3d..f4e56680881 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/utils.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/utils.py @@ -553,6 +553,7 @@ def get_attention_backend( if not logger.hasHandlers(): logger.addHandler(AttentionLogging._stream_handler) device_compute_capability = get_device_compute_capability(attention_params.device) + is_sm12x = device_compute_capability[0] == 12 cudnn_version = get_cudnn_version() run_config = { "transformer_engine_version": te.__version__, @@ -802,16 +803,16 @@ def _disable_all_flash_attention() -> None: logger.debug("Disabling FusedAttention for %s", fp8_recipe.__class__.__name__) use_fused_attention = False - if device_compute_capability == (12, 0): + if is_sm12x: if use_flash_attention: logger.debug( "Disabling FlashAttention as FP8 is not supported" - " for compute capability = sm120" + " for compute capability = sm12x" ) if use_fused_attention: logger.debug( "Disabling FusedAttention as FP8 is not supported" - " for compute capability = sm120" + " for compute capability = sm12x" ) use_flash_attention = False use_fused_attention = False @@ -950,14 +951,14 @@ def _disable_all_flash_attention() -> None: # Flash v4 | FP16/BF16 | TODO | sm80+ | bshd,sbhd,thd | TODO # Unfused | FP32/FP16/BF16 | non-paged/paged | all | bshd,sbhd,thd | >= 1 if inference_params is not None: - # Temporarily disabling fused attention for kv caching for sm89/sm120 irrespective of + # Temporarily disabling fused attention for kv caching for sm89/sm12x irrespective of # cuDNN version until the cuDNN bug is resolved. - if device_compute_capability in ((8, 9), (12, 0)): - logger.debug("Disabling FusedAttention for KV caching for sm89/sm120") + if device_compute_capability == (8, 9) or is_sm12x: + logger.debug("Disabling FusedAttention for KV caching for sm89/sm12x") use_fused_attention = False - # Temporarily disable FlashAttention for KV caching on sm120 - if device_compute_capability == (12, 0): - logger.debug("Disabling FlashAttention for KV caching for sm120") + # Temporarily disable FlashAttention for KV caching on sm12x + if is_sm12x: + logger.debug("Disabling FlashAttention for KV caching for sm12x") use_flash_attention = False if context_parallel: logger.debug("Disabling all backends for KV caching with context parallelism") @@ -1019,15 +1020,11 @@ def _disable_all_flash_attention() -> None: qkv_layout, ) use_fused_attention = False - if ( - device_compute_capability == (12, 0) - and (head_dim_qk > 128 or head_dim_qk % 8 != 0) - and is_training - ): + if is_sm12x and (head_dim_qk > 128 or head_dim_qk % 8 != 0) and is_training: if use_fused_attention: logger.debug( "Disabling FusedAttention as MLA for backward pass is not supported for compute" - " capability = sm120 for a head_dim_qk > 128 or head_dim_qk %%8 != 0. Found:" + " capability = sm12x for a head_dim_qk > 128 or head_dim_qk %%8 != 0. Found:" " head_dim_qk = %s", head_dim_qk, ) @@ -1193,19 +1190,19 @@ def _is_fa3_supported(num_heads, num_gqa_groups, head_dim_qk, head_dim_v, qkv_dt if use_unfused_attention: logger.debug("Disabling UnfusedDotProductAttention for pad_between_seqs = True") use_unfused_attention = False - if device_compute_capability == (12, 0): + if is_sm12x: if cudnn_version < (9, 18, 1): if use_fused_attention: logger.debug( "Disabling FusedAttention when Q or KV uses THD because it is" - " not supported for compute capability = sm120 and cuDNN version < 9.18.1" + " not supported for compute capability = sm12x and cuDNN version < 9.18.1" ) use_fused_attention = False elif qkv_layout in {"t3hd", "th3d"}: if use_fused_attention: logger.debug( "Disabling FusedAttention as qkv_layout = %s is not supported for" - " compute capability = sm120", + " compute capability = sm12x", qkv_layout, ) use_fused_attention = False diff --git a/transformer_engine/pytorch/cpp_extensions/cudnn_fused_attn.py b/transformer_engine/pytorch/cpp_extensions/cudnn_fused_attn.py index 1812344b094..da5a0f75c18 100644 --- a/transformer_engine/pytorch/cpp_extensions/cudnn_fused_attn.py +++ b/transformer_engine/pytorch/cpp_extensions/cudnn_fused_attn.py @@ -28,6 +28,7 @@ check_graph_support, compile_graph, bucket, + uses_ragged_attention_graph, ) from .cudnn import ( @@ -137,7 +138,7 @@ def __init__( ) self.b, self.h, self.hg, self.dq, self.dv, self.tq, self.tk = dimensions capability = torch.cuda.get_device_capability(self.device) - self.ragged_graph = capability >= (9, 0) and capability != (12, 0) + self.ragged_graph = uses_ragged_attention_graph(10 * capability[0] + capability[1]) self.ragged_stats = self.q_format == "thd" and self.ragged_graph self.sq = bucket(self.tq, 1024, 32768) if self.ragged_stats else params["max_seqlen_q"] self.skv = ( From 00ac7cf33a48696a9a671c5794a8f107f8837ff1 Mon Sep 17 00:00:00 2001 From: Vladimir Cherepanov Date: Thu, 15 Oct 2026 21:25:47 +0000 Subject: [PATCH 6/6] Align attention statistics metadata across SM12x variants Use the shared architecture predicate for JAX Max tensors and PyTorch fake saved statistics. Add regression coverage for SM121 and SM129 alongside neighboring architecture families. Validation: 19 focused JAX tests, pre-commit, and both Python lint suites passed. PyTorch tests require the unavailable torch package. The license check reports existing untracked nccl_ep headers. Signed-off-by: Vladimir Cherepanov --- tests/jax/test_fused_attn.py | 24 +++++++------ tests/pytorch/attention/test_attention.py | 36 +++++++++++++++++++ .../jax/cpp_extensions/attention.py | 13 ++++--- .../dot_product_attention/backends.py | 3 +- 4 files changed, 59 insertions(+), 17 deletions(-) diff --git a/tests/jax/test_fused_attn.py b/tests/jax/test_fused_attn.py index 14d7bd972cb..33157f2b183 100644 --- a/tests/jax/test_fused_attn.py +++ b/tests/jax/test_fused_attn.py @@ -2344,24 +2344,26 @@ def forbidden(*_args, **_kwargs): @pytest.mark.parametrize( - "kind", + "kind,sm_arch", [ - "dense", - "causal", - "padding", - "window", - "ragged", - "ragged_dense_stats", - "ragged_dense_stats_f16", + ("dense", None), + ("causal", None), + ("padding", None), + ("window", None), + ("ragged", None), + ("ragged_dense_stats", 120), + ("ragged_dense_stats_f16", 120), + ("ragged_dense_stats_f16", 121), + ("ragged_dense_stats_f16", 129), ], ) -def test_cudnn_python_fp8_gqa_masks(kind, monkeypatch, request): +def test_cudnn_python_fp8_gqa_masks(kind, sm_arch, monkeypatch, request): """Compare non-uniform FP8 attention and gradients against native FP32 GQA.""" ragged = kind.startswith("ragged") fp8 = not kind.endswith("_f16") if get_device_compute_capability(0) < (100 if ragged else 90): pytest.skip("FP8 ragged backward requires Blackwell; dense FP8 requires Hopper.") - if kind.startswith("ragged_dense_stats"): + if sm_arch is not None: from transformer_engine.jax.cpp_extensions import cudnn_attention # Test the dense statistics conversion on the available GPU. @@ -2369,7 +2371,7 @@ def test_cudnn_python_fp8_gqa_masks(kind, monkeypatch, request): # cache when mocking a different architecture's Max tensor layout. jax.clear_caches() request.addfinalizer(jax.clear_caches) - monkeypatch.setattr(cudnn_attention, "device_capabilities", lambda: (120,)) + monkeypatch.setattr(cudnn_attention, "device_capabilities", lambda: (sm_arch,)) keys = jax.random.split(jax.random.PRNGKey(123), 4) q = jax.random.choice(keys[0], jnp.array([-0.5, 0.5], jnp.bfloat16), (2, 128, 4, 64)) k = jax.random.choice(keys[1], jnp.array([-0.5, 0.5], jnp.bfloat16), (2, 128, 2, 64)) diff --git a/tests/pytorch/attention/test_attention.py b/tests/pytorch/attention/test_attention.py index 4d4aa24d96d..952873e8415 100644 --- a/tests/pytorch/attention/test_attention.py +++ b/tests/pytorch/attention/test_attention.py @@ -10,6 +10,7 @@ import pathlib import copy from dataclasses import replace +from types import SimpleNamespace from typing import Any, Dict, List, NamedTuple, Tuple, Union from unittest.mock import Mock @@ -391,6 +392,41 @@ def test_fused_attn_backward_token_bounds(monkeypatch, dtype, capability, layout assert query.skv == 1024 +@pytest.mark.parametrize("capability", ((10, 0), (12, 0), (12, 1), (12, 9), (13, 0))) +@pytest.mark.parametrize("layout", ("bshd_bshd_bshd", "thd_bshd_bshd", "thd_thd_thd")) +def test_fused_attn_fake_stats_shape(capability, layout, monkeypatch): + """Fake saved statistics must match the runtime layout across the SM12x family.""" + from transformer_engine.pytorch.attention.dot_product_attention import backends + + monkeypatch.setattr(backends, "get_device_compute_capability", lambda: capability) + batch, seqlen, heads, dim, tokens = 3, 128, 4, 64, 224 + ragged = layout.startswith("thd") + shape = (tokens, heads, dim) if ragged else (batch, seqlen, heads, dim) + kv_shape = (tokens, heads, dim) if layout.endswith("thd") else (batch, seqlen, heads, dim) + q = torch.empty(shape, device="meta", dtype=torch.bfloat16) + kv = torch.empty(kv_shape, device="meta", dtype=torch.bfloat16) + args = SimpleNamespace( + q=q, + k=kv, + v=kv, + packed_qkv=None, + packed_kv=None, + qkv_layout=layout, + cu_seqlens_q=torch.empty(batch + 1, device="meta", dtype=torch.int32), + max_seqlen_q=seqlen, + return_max_logit=False, + attn_bias_type="no_bias", + softmax_type="vanilla", + fused_attention_backend=FusedAttnBackend.F16_arbitrary_seqlen, + ) + _, _, saved, _ = backends._fused_attn_forward_fake(args) + expected = (batch, heads, seqlen, 1) + if ragged and capability[0] != 12: + expected = (tokens, heads, 1) + assert saved[10].shape == expected + assert saved[10].dtype == torch.float32 + + @pytest.mark.parametrize("dtype", (torch.float16, torch.bfloat16)) @pytest.mark.parametrize("kv_format", ("thd", "bshd", "sbhd")) def test_fused_attn_packed_backward(dtype, kv_format): diff --git a/transformer_engine/jax/cpp_extensions/attention.py b/transformer_engine/jax/cpp_extensions/attention.py index b5f9d37fe23..a5b1a747725 100644 --- a/transformer_engine/jax/cpp_extensions/attention.py +++ b/transformer_engine/jax/cpp_extensions/attention.py @@ -21,6 +21,7 @@ from transformer_engine_jax import ( NVTE_Fused_Attn_Backend, ) +from transformer_engine.common.cudnn.graph import uses_ragged_attention_graph from transformer_engine.jax.attention import ( AttnBiasType, AttnMaskType, @@ -503,8 +504,8 @@ def abstract( raise ValueError(f"Unsupported backend: {message}") softmax_aux_aval = q_aval.update(shape=softmax_shape, dtype=softmax_dtype) if config.return_max_logit: - # cuDNN Max is row-wise over S_kv. Dense and SM120 THD use - # [..., H, S_q, 1]; non-SM120 THD uses [..., S_q, H, 1]. + # cuDNN Max is row-wise over S_kv. Dense and SM12x THD use + # [..., H, S_q, 1]; other THD layouts use [..., S_q, H, 1]. # Both raw layouts are reduced to the public per-head [H] result below. if FusedAttnFwdPrimitive._uses_thd_ragged_max_tensor(config): max_tensor_shape = (*batch_shape, q_max_seqlen, attn_heads, 1) @@ -724,8 +725,8 @@ def convert_to_2d(offsets, batch, max_seqlen): def _reduce_max_logit(max_tensor, output, q_seqlen, q_seq_offsets, config): """Reduce cuDNN's row-wise Max tensor to the public per-head max_logit. - Dense and SM120 THD use ``[..., H, S_q, 1]``; non-SM120 - THD uses ``[..., S_q, H, 1]``. A rank-3 THD result is ``[T_q, H, 1]``. + Dense and SM12x THD use ``[..., H, S_q, 1]``; THD with packed statistics + uses ``[..., S_q, H, 1]``. A rank-3 THD result is ``[T_q, H, 1]``. All layouts reduce to ``[H]``. Static THD buffers can contain invalid query rows, so those rows are masked before reduction. """ @@ -764,7 +765,9 @@ def _reduce_max_logit(max_tensor, output, q_seqlen, q_seq_offsets, config): @staticmethod def _uses_thd_ragged_max_tensor(config): """Return whether cuDNN writes THD Max with BSH-like ragged-stats layout.""" - return config.qkv_layout.is_thd() and 120 not in cudnn_attention.device_capabilities() + return config.qkv_layout.is_thd() and uses_ragged_attention_graph( + cudnn_attention.device_capabilities()[0] + ) @staticmethod def _empty_or_neg_inf_max_logit(head, dtype, config): diff --git a/transformer_engine/pytorch/attention/dot_product_attention/backends.py b/transformer_engine/pytorch/attention/dot_product_attention/backends.py index 697e6b7b501..2b32df0e6ae 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/backends.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/backends.py @@ -17,6 +17,7 @@ import torch import torch.nn.functional as F +from transformer_engine.common.cudnn.graph import uses_ragged_attention_graph from transformer_engine.pytorch.utils import ( get_device_compute_capability, split_tensor_along_dim, @@ -2418,7 +2419,7 @@ def _fused_attn_stats_shape( if q_format == "thd": num_heads = q_shape[1] sm = get_device_compute_capability() - if sm >= (9, 0) and sm != (12, 0): + if uses_ragged_attention_graph(10 * sm[0] + sm[1]): return (q_shape[0], num_heads, 1) batch_size = args.cu_seqlens_q.shape[0] - 1 return (batch_size, num_heads, args.max_seqlen_q, 1)