From c105f9d44875d70d6a936fed46b8090261b72ee9 Mon Sep 17 00:00:00 2001 From: Vladimir Cherepanov Date: Sun, 11 Oct 2026 21:09:30 +0000 Subject: [PATCH 01/49] 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 02/49] 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 03/49] 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 04/49] 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 05/49] 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 06/49] 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) From 13141e302f835acd7692073695bdab67587ce02c Mon Sep 17 00:00:00 2001 From: Vladimir Cherepanov Date: Mon, 19 Oct 2026 22:32:55 +0000 Subject: [PATCH 07/49] Fix cuDNN attention support guards and output initialization Restore specialized FP16/BF16 THD learnable-softmax kernels before cuDNN 9.26 and reject all affected FP8 THD sink-token backward paths in JAX. Preserve the FP16/BF16 backward fallback. Skip fills for fully written JAX outputs while retaining initialization for ragged buffers and unwritten optional results. Add support-policy, buffer-binding, and numerical regression coverage, and document the version restrictions. Signed-off-by: Vladimir Cherepanov --- docs/examples/jax/attention.rst | 8 + tests/jax/test_fused_attn.py | 224 ++++++++++++++++++ tests/pytorch/attention/test_attention.py | 77 ++++++ .../jax/cpp_extensions/cudnn_attention.py | 39 +-- .../jax/csrc/extensions/cudnn_graph.cpp | 6 +- .../cpp_extensions/cudnn_fused_attn.py | 1 + 6 files changed, 338 insertions(+), 17 deletions(-) diff --git a/docs/examples/jax/attention.rst b/docs/examples/jax/attention.rst index 23f52b105d2..7ed82edff5d 100644 --- a/docs/examples/jax/attention.rst +++ b/docs/examples/jax/attention.rst @@ -39,6 +39,14 @@ does not maintain delayed-scaling history. 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 THD attention requires cuDNN 9.23 or later. FP8 THD backward with off-by-one +or learnable softmax requires cuDNN 9.26 or later; select +``NVTE_FP8_DPA_BWD=0`` to use FP16/BF16 backward instead. For FP16/BF16 THD +learnable-softmax backward with packed statistics, cuDNN versions before 9.26 +remain eligible for value-head dimensions 64, 128, and 256. Other value-head +dimensions require cuDNN 9.26 or later. Execution-plan support still depends on +the complete attention configuration. + 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 diff --git a/tests/jax/test_fused_attn.py b/tests/jax/test_fused_attn.py index 33157f2b183..568c526b974 100644 --- a/tests/jax/test_fused_attn.py +++ b/tests/jax/test_fused_attn.py @@ -545,6 +545,230 @@ def test_cudnn_attention_sm12x_dense_statistics(sm_arch, direction, monkeypatch) assert tensors["stats"].ragged_offset == "offset_stats" +@pytest.mark.parametrize("dtype", (jnp.float16, jnp.bfloat16)) +@pytest.mark.parametrize("head_dim_v", (64, 128, 192, 256)) +@pytest.mark.parametrize("version", (92500, 92600)) +@pytest.mark.parametrize("direction", ("fwd", "bwd")) +def test_cudnn_attention_thd_learnable_support(dtype, head_dim_v, version, direction, monkeypatch): + """Preserve specialized FP16/BF16 THD learnable-softmax kernels before cuDNN 9.26.""" + from types import SimpleNamespace + from transformer_engine.jax.cpp_extensions import cudnn_attention + + monkeypatch.setattr(cudnn_attention, "cudnn_version", lambda: version) + 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.LEARNABLE_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(dtype)) + for shape in ((2, 128, 4, 64),) * 2 + ((2, 128, 4, head_dim_v), (0,), (1, 4, 1, 1)) + ) + if direction == "bwd" and version < 92600 and head_dim_v == 192: + with pytest.raises(ValueError, match="THD learnable softmax backward requires cuDNN 9.26"): + cudnn_attention._describe_graph(direction, shapes, config, (100,), True, True) + else: + description = cudnn_attention._describe_graph(direction, shapes, config, (100,), True, True) + assert not description.fp8 + assert dict(description.spec.tensors)["v"].dim[-1] == head_dim_v + + +@pytest.mark.parametrize("softmax", tuple(AttnSoftmaxType)) +@pytest.mark.parametrize("version", (92300, 92500, 92600)) +@pytest.mark.parametrize("direction", ("fwd", "bwd")) +@pytest.mark.parametrize("fp8_backward", (False, True)) +def test_cudnn_attention_fp8_thd_sink_support( + softmax, version, direction, fp8_backward, monkeypatch +): + """All sink-token FP8 backward paths require 9.26; F16 backward remains available.""" + from types import SimpleNamespace + from transformer_engine.jax.cpp_extensions import cudnn_attention + + monkeypatch.setattr(cudnn_attention, "cudnn_version", lambda: version) + config = SimpleNamespace( + qkv_layout=QKVLayout.THD_THD_THD, + max_segments_per_seq=1, + fp8=True, + attn_mask_type=AttnMaskType.PADDING_MASK, + attn_bias_type=AttnBiasType.NO_BIAS, + softmax_type=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)) + ) + if ( + direction == "bwd" + and fp8_backward + and version < 92600 + and softmax != AttnSoftmaxType.VANILLA_SOFTMAX + ): + with pytest.raises( + ValueError, match="FP8 THD backward with a sink token requires cuDNN 9.26" + ): + cudnn_attention._describe_graph(direction, shapes, config, (100,), True, fp8_backward) + else: + description = cudnn_attention._describe_graph( + direction, shapes, config, (100,), True, fp8_backward + ) + assert description.fp8 == (direction == "fwd" or fp8_backward) + + +@pytest.mark.parametrize( + "layout", + ( + QKVLayout.BS3HD, + QKVLayout.BSHD_BS2HD, + QKVLayout.BSHD_BSHD_BSHD, + QKVLayout.T3HD, + QKVLayout.THD_T2HD, + QKVLayout.THD_THD_THD, + ), +) +@pytest.mark.parametrize("direction", ("fwd", "bwd")) +@pytest.mark.parametrize("fp8", (False, True)) +def test_cudnn_attention_output_initialization(layout, direction, fp8, monkeypatch): + """Only ragged storage and unwritten optional results request an FFI fill.""" + from types import SimpleNamespace + from transformer_engine.jax.cpp_extensions import cudnn_attention + + q = jnp.ones((1, 128, 2, 64), jnp.bfloat16) + empty = jnp.empty((0,), q.dtype) + if layout.is_qkvpacked(): + operands = (jnp.stack((q, q, q), axis=-3), empty, empty) + grad_slots = (0, 0, 0) + elif layout.is_kvpacked(): + operands = (q, jnp.stack((q, q), axis=-3), empty) + grad_slots = (0, 1, 1) + else: + operands = (q, q, q) + grad_slots = (0, 1, 2) + backward = direction == "bwd" + output_slots = (*grad_slots, 4) if backward else (0, 1) + if fp8: + output_slots += (5, 6, 7, 8) if backward else (3, 4) + elif not backward: + output_slots += (2,) + plan = SimpleNamespace( + graph=SimpleNamespace(workspace_size=0), + names=(), + input_indices=(), + input_offsets=(), + output_indices=output_slots, + output_offsets=(0,) * len(output_slots), + metadata_batch=1, + ragged_stats=layout.is_thd(), + multipliers=(128, 128, 128, 128, 2), + fp8=fp8, + ) + monkeypatch.setattr(cudnn_attention, "get_graph", lambda *_args: plan) + monkeypatch.setattr(cudnn_attention, "graph_attributes", lambda _graph: {}) + captured = {} + + def ffi_call(_name, specs): + def execute(*_args, **attributes): + captured.update(attributes) + return tuple(jnp.empty(spec.shape, spec.dtype) for spec in specs) + + return execute + + monkeypatch.setattr(cudnn_attention.ffi, "ffi_call", ffi_call) + config = SimpleNamespace(qkv_layout=layout, fp8=fp8, return_max_logit=not fp8) + stats_shape = (1, 128, 2, 1) if layout.is_thd() else (1, 2, 128, 1) + stats = jnp.empty(stats_shape, jnp.float32) + rng = jnp.zeros((4,), jnp.uint32) + cu_seqlens = jnp.array([0, 128], jnp.int32) + bias = jnp.zeros((1, 1, 1, 128), q.dtype) + sink = jnp.zeros((1, 2, 1, 1), jnp.float32) + args = (*operands, bias, sink) + if backward: + args += (stats, rng, q, q, cu_seqlens, cu_seqlens, cu_seqlens, cu_seqlens) + results = (*operands, bias, sink, empty) + # dBias is deliberately unbound; packed gradient slots also have dummy results. + expected = [ + -1, + 0 if layout.is_qkvpacked() else -1, + 0 if not layout.is_separate() else -1, + 0, + -1, + ] + if layout.is_thd(): + expected[:3] = [0, 0, 0] + if fp8: + expected += [-1] * 4 + else: + args += (rng, cu_seqlens, cu_seqlens, cu_seqlens, cu_seqlens) + results = (q, stats, empty if fp8 else stats, rng, empty) + expected = [0, 0xF0, 0xF0] if layout.is_thd() else [-1, -1, 0 if fp8 else -1] + if fp8: + expected += [-1, -1] + cudnn_attention.execute(direction, args, config, results) + np.testing.assert_array_equal(captured["output_fill_values"], expected) + + +@pytest.mark.parametrize("dtype", (jnp.float16, jnp.bfloat16)) +@pytest.mark.parametrize("layout", (QKVLayout.T3HD, QKVLayout.THD_T2HD, QKVLayout.THD_THD_THD)) +def test_cudnn_attention_ragged_padding_outputs(dtype, layout): + """Real FFI execution must keep unused ragged outputs and packed gradients zero.""" + if get_device_compute_capability(0) < 90: + pytest.skip("FP16/BF16 THD attention requires Hopper or newer.") + tensor = jnp.ones((1, 128, 2, 64), dtype) + if layout.is_qkvpacked(): + operands = (jnp.stack((tensor, tensor, tensor), axis=-3),) + elif layout.is_kvpacked(): + operands = (tensor, jnp.stack((tensor, tensor), axis=-3)) + else: + operands = (tensor, tensor, tensor) + segment_ids = jnp.where(jnp.arange(128)[None, :] < 64, 1, 0).astype(jnp.int32) + positions = jnp.where(segment_ids > 0, jnp.arange(128)[None, :], 0).astype(jnp.int32) + descriptor = SequenceDescriptor.from_segment_ids_and_pos(segment_ids, positions) + + def attention(arrays): + return fused_attn( + arrays, + None, + descriptor, + None, + attn_bias_type=AttnBiasType.NO_BIAS, + attn_mask_type=AttnMaskType.PADDING_MASK, + qkv_layout=layout, + softmax_type=AttnSoftmaxType.VANILLA_SOFTMAX, + scaling_factor=0.125, + dropout_probability=0.0, + is_training=True, + return_max_logit=True, + ) + + output, maximum = jax.jit(attention)(operands) + expected = jnp.zeros_like(tensor).at[:, :64].set(1) + np.testing.assert_array_equal(output, expected) + np.testing.assert_array_equal(maximum, jnp.full((2,), 8.0, jnp.float32)) + gradients = jax.jit(jax.grad(lambda arrays: jnp.sum(attention(arrays)[0])))(operands) + zero = jnp.zeros_like(tensor) + if layout.is_qkvpacked(): + expected_gradients = (jnp.stack((zero, zero, expected), axis=-3),) + elif layout.is_kvpacked(): + expected_gradients = (zero, jnp.stack((zero, expected), axis=-3)) + else: + expected_gradients = (zero, zero, expected) + for actual, reference in zip(gradients, expected_gradients): + np.testing.assert_allclose(actual, reference, rtol=1e-3, atol=1e-3) + + class BiasShape(Enum): """ Enum class to represent the different bias shapes used in the fused attention. diff --git a/tests/pytorch/attention/test_attention.py b/tests/pytorch/attention/test_attention.py index 6f8eae3b5d6..efe2859dc83 100644 --- a/tests/pytorch/attention/test_attention.py +++ b/tests/pytorch/attention/test_attention.py @@ -53,6 +53,7 @@ QKVLayout, fused_attn_bwd, fused_attn_fwd, + SoftmaxType, ) from transformer_engine.pytorch.cpp_extensions.cudnn_fused_attn import get_fused_attn_backend from transformer_engine.pytorch.distributed import CudaRNGStatesTracker @@ -553,6 +554,7 @@ def fail(_self): "qkv_layout": QKVLayout["thd_thd_thd"], "attn_mask_type": AttnMaskType["padding"], "softmax_type": tex.NVTE_Softmax_Type.NVTE_LEARNABLE_SOFTMAX, + "head_dim_v": 192, "num_tokens_q": 256, "num_tokens_kv": 256, }, @@ -617,6 +619,81 @@ def build(_self): assert message.startswith(reason) +@pytest.mark.parametrize("dtype", (DType.kFloat16, DType.kBFloat16)) +@pytest.mark.parametrize("head_dim_v", (64, 128, 192, 256)) +@pytest.mark.parametrize("version", (92500, 92600)) +@pytest.mark.parametrize("is_training", (False, True)) +def test_fused_attn_python_thd_learnable_support( + dtype, head_dim_v, version, is_training, monkeypatch +): + """Only generic THD learnable-softmax backward kernels need cuDNN 9.26.""" + 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: (10, 0)) + monkeypatch.setattr(_AttentionGraph, "entry", lambda _self: None) + config = FusedAttentionParams( + is_training=is_training, + qkv_dtype=dtype, + qkv_layout=QKVLayout["thd_thd_thd"], + attn_mask_type=AttnMaskType["padding"], + softmax_type=SoftmaxType["learnable"], + batch_size=2, + num_attn_heads=4, + num_gqa_groups=4, + head_dim_qk=64, + head_dim_v=head_dim_v, + max_seqlen_q=128, + max_seqlen_kv=128, + num_tokens_q=256, + num_tokens_kv=256, + ) + backend, message = get_fused_attn_backend(config) + if is_training and version < 92600 and head_dim_v == 192: + assert backend == FusedAttnBackend.No_Backend + assert "THD learnable softmax backward" in message + else: + assert backend == FusedAttnBackend.F16_arbitrary_seqlen + assert message == "" + + +@pytest.mark.parametrize("softmax", ("vanilla", "off-by-one", "learnable")) +@pytest.mark.parametrize("version", (92300, 92500, 92600)) +@pytest.mark.parametrize("is_training", (False, True)) +def test_fused_attn_python_fp8_thd_sink_support(softmax, version, is_training, monkeypatch): + """Every FP8 THD sink-token backward kernel needs cuDNN 9.26, even at D=64.""" + 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: (10, 0)) + monkeypatch.setattr(_AttentionGraph, "entry", lambda _self: None) + config = FusedAttentionParams( + is_training=is_training, + qkv_dtype=DType.kFloat8E4M3, + qkv_layout=QKVLayout["thd_thd_thd"], + attn_mask_type=AttnMaskType["padding"], + softmax_type=SoftmaxType[softmax], + 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, + num_tokens_q=256, + num_tokens_kv=256, + ) + backend, message = get_fused_attn_backend(config) + if is_training and version < 92600 and softmax != "vanilla": + assert backend == FusedAttnBackend.No_Backend + assert "sink token requires cuDNN 9.26" in message + else: + assert backend == FusedAttnBackend.FP8 + assert message == "" + + @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.""" diff --git a/transformer_engine/jax/cpp_extensions/cudnn_attention.py b/transformer_engine/jax/cpp_extensions/cudnn_attention.py index 054649e8585..f96f74f5b94 100644 --- a/transformer_engine/jax/cpp_extensions/cudnn_attention.py +++ b/transformer_engine/jax/cpp_extensions/cudnn_attention.py @@ -238,18 +238,20 @@ def _describe_graph(direction, shapes, config, capabilities, nondeterministic, f version = cudnn_version() if ragged and capabilities[0] < 90: raise ValueError("F16/BF16 fused attention with THD format requires sm90 or later.") - if ( - backward - and ragged_stats - and config.softmax_type.name == "LEARNABLE_SOFTMAX" - and version < 92600 - ): - raise ValueError("THD learnable softmax backward requires cuDNN 9.26.0 or later.") + if backward and not fp8 and ragged_stats: + if ( + config.softmax_type.name == "LEARNABLE_SOFTMAX" + and version < 92600 + and dv not in (64, 128, 256) + ): + raise ValueError("THD learnable softmax backward requires cuDNN 9.26.0 or later.") if fp8 and ragged: if version < 92300: raise ValueError("FP8 THD attention requires cuDNN 9.23.0 or later.") if backward and capabilities[0] < 100: raise ValueError("FP8 THD backward requires sm100 or later.") + if backward and config.softmax_type.name != "VANILLA_SOFTMAX" and version < 92600: + raise ValueError("FP8 THD backward with a sink token requires cuDNN 9.26.0 or later.") if fp8 and backward and config.attn_bias_type.name != "NO_BIAS": raise ValueError("cuDNN FP8 backward does not support attention bias or ALiBi.") io_dtype = _dtype(avals[0].dtype) @@ -520,13 +522,22 @@ def execute(direction, args, config, result_avals): 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 + # Fully written outputs need no initialization. Unbound optional results + # (including unsupported bias gradients) must still return zeros. + fills = np.full(len(specs) - 1, -1, np.int64) + for slot in set(range(len(fills))) - set(plan.output_indices): + fills[slot] = 0 + if config.qkv_layout.is_thd(): + if backward: + # Packed gradients share result slots: clear each backing buffer once. + fills[:3] = 0 + else: + fills[0] = 0 + # 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[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) diff --git a/transformer_engine/jax/csrc/extensions/cudnn_graph.cpp b/transformer_engine/jax/csrc/extensions/cudnn_graph.cpp index 40b9fa7cdee..2499899ca3f 100644 --- a/transformer_engine/jax/csrc/extensions/cudnn_graph.cpp +++ b/transformer_engine/jax/csrc/extensions/cudnn_graph.cpp @@ -133,9 +133,9 @@ ffi::Error ExecuteGraph(cudaStream_t stream, ffi::RemainingArgs args, ffi::Remai 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; + // Negative or absent fill values leave fully written outputs uninitialized. + const int64_t fill = output_fill_values.has_value() ? output_fill_values.value()[i] : -1; + if (fill >= 0 && buf.value()->size_bytes()) { CheckCuda( cudaMemsetAsync(buf.value()->untyped_data(), fill, buf.value()->size_bytes(), stream)); } diff --git a/transformer_engine/pytorch/cpp_extensions/cudnn_fused_attn.py b/transformer_engine/pytorch/cpp_extensions/cudnn_fused_attn.py index da5a0f75c18..fcdf2363571 100644 --- a/transformer_engine/pytorch/cpp_extensions/cudnn_fused_attn.py +++ b/transformer_engine/pytorch/cpp_extensions/cudnn_fused_attn.py @@ -595,6 +595,7 @@ def _support_reason(config, attn, version, capability): and attn.ragged_stats and p["softmax_type"] == "learnable" and version < 92600 + and attn.dv not in (64, 128, 256) ): return ( "Known cuDNN < 9.26.0 issue with THD learnable softmax backward. Please upgrade" From e9e9b8c6f7553b4360da69c090bc6f00825088ff Mon Sep 17 00:00:00 2001 From: Vladimir Cherepanov Date: Tue, 20 Oct 2026 03:40:21 +0000 Subject: [PATCH 08/49] Fix build tools imports in JAX dependency tests Temporarily expose the checkout root while importing build_tools so the new tests collect from the CI working directory. Restore the import path immediately to preserve installed-package resolution. Validation: all three tests pass with direct pytest and from the JAX CI working directory; import-path restoration and git diff --check pass. Signed-off-by: Vladimir Cherepanov --- tests/jax/test_build_tools.py | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/tests/jax/test_build_tools.py b/tests/jax/test_build_tools.py index 714aafec106..5f39fb4dbb5 100644 --- a/tests/jax/test_build_tools.py +++ b/tests/jax/test_build_tools.py @@ -6,12 +6,17 @@ from email.parser import Parser from io import StringIO +from pathlib import Path import pytest from packaging.requirements import Requirement from setuptools import Distribution -from build_tools import jax as jax_build +# CI runs outside the checkout root and tests the installed TE package. +# Expose the source-only build tools just for this import. +with pytest.MonkeyPatch.context() as _build_import_path: + _build_import_path.syspath_prepend(str(Path(__file__).resolve().parents[2])) + from build_tools import jax as jax_build @pytest.mark.parametrize("frontend_version", ("1.29.0", "1.30.0")) From 9e0f6e7a27b9799f3ea2d3435a53860a9b8b29f8 Mon Sep 17 00:00:00 2001 From: Vladimir Cherepanov Date: Tue, 20 Oct 2026 04:24:34 +0000 Subject: [PATCH 09/49] Skip deterministic cuDNN execution tests below SM90 Match the existing backend policy when the Python execution test forces cuDNN. Skip deterministic training below SM90 before disabling the fallback backends. Signed-off-by: Vladimir Cherepanov --- tests/pytorch/attention/test_attention.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/tests/pytorch/attention/test_attention.py b/tests/pytorch/attention/test_attention.py index efe2859dc83..85e53fb59b9 100644 --- a/tests/pytorch/attention/test_attention.py +++ b/tests/pytorch/attention/test_attention.py @@ -224,6 +224,8 @@ def test_fused_attn_backend_message(): @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.""" + if _deterministic and device_compute_capability < (9, 0): + pytest.skip("TE disables deterministic cuDNN attention training below SM90") def common_attention(*_args, **_kwargs): raise AssertionError("PyTorch attention called the TE common attention API") From 01ffce20069889d2a571616c0a188ed538028eb4 Mon Sep 17 00:00:00 2001 From: Vladimir Cherepanov Date: Tue, 20 Oct 2026 04:59:13 +0000 Subject: [PATCH 10/49] Share MXFP8 attention preparation rules across frontends Centralize operand directions, scale padding and reordering, alignment checks, and known backend exclusions. Adapt the existing PyTorch preparation without changing tensor ownership or GPU execution. Validation: 14 JAX metadata tests and changed-file pre-commit checks pass. PyTorch runtime validation is unavailable in this JAX environment. Signed-off-by: Vladimir Cherepanov --- tests/jax/test_mxfp8_attention.py | 55 +++++++++++ transformer_engine/common/cudnn/mxfp8.py | 57 +++++++++++ .../attention/dot_product_attention/utils.py | 95 ++++++++----------- 3 files changed, 149 insertions(+), 58 deletions(-) create mode 100644 tests/jax/test_mxfp8_attention.py create mode 100644 transformer_engine/common/cudnn/mxfp8.py diff --git a/tests/jax/test_mxfp8_attention.py b/tests/jax/test_mxfp8_attention.py new file mode 100644 index 00000000000..adf614b3d18 --- /dev/null +++ b/tests/jax/test_mxfp8_attention.py @@ -0,0 +1,55 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""MXFP8 attention preparation, numerical, and integration tests.""" + +import math + +import pytest + +from transformer_engine.common.cudnn import mxfp8 +from transformer_engine.common.cudnn.graph import mxfp8_scale_spec + + +@pytest.mark.parametrize("columnwise", (False, True)) +@pytest.mark.parametrize("dim", ((2, 3, 128, 64), (1, 4, 256, 96), (2, 1, 384, 160))) +def test_scale_reordering_matches_graph(dim, columnwise): + """Preparation and cuDNN agree on padding each batch/head scale matrix.""" + batch, heads, seqlen, head_dim = dim + scale_shape = ( + (batch, heads, seqlen // 32, head_dim) + if columnwise + else (batch, heads, seqlen, head_dim // 32) + ) + padded, tiled, permutation = mxfp8.scale_reordering(scale_shape, columnwise) + assert padded == mxfp8_scale_spec(dim, "bhsd", columnwise).dim + assert math.prod(tiled) == math.prod(padded) + assert sorted(permutation) == list(range(len(tiled))) + + +@pytest.mark.parametrize( + "sm,version,fmt,error", + ( + (90, 92600, "bshd", "SM100"), + (100, 92000, "bshd", "9.21"), + (100, 92300, "bshd", "defects"), + (100, 92301, "bshd", "defects"), + (100, 92600, "thd", "THD"), + (100, 92100, "bshd", ""), + (100, 92600, "sbhd", ""), + ), +) +def test_shared_mxfp8_support(sm, version, fmt, error): + """Both frontends retain the same hardware and known-defect exclusions.""" + reason = mxfp8.support_error(sm, version, fmt) + assert error in reason if error else not reason + + +def test_quantization_directions(): + """Forward V contracts over sequence; backward Q/K need both directions.""" + assert mxfp8.quantization_layouts() == ((True, False), (True, False), (False, True)) + assert mxfp8.quantization_layouts(True) == ((True, True), (True, True), (True, False)) + assert not mxfp8.shape_error(128, 256, 64, 96) + assert mxfp8.shape_error(129, 256, 64, 96) + assert mxfp8.shape_error(128, 256, 64, 80) diff --git a/transformer_engine/common/cudnn/mxfp8.py b/transformer_engine/common/cudnn/mxfp8.py new file mode 100644 index 00000000000..80d5626523f --- /dev/null +++ b/transformer_engine/common/cudnn/mxfp8.py @@ -0,0 +1,57 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""Framework-independent MXFP8 attention operand preparation rules.""" + +from .graph import round_up + + +def quantization_layouts(backward=False): + """Return (rowwise, columnwise) requirements for Q, K, and V. + + Rowwise quantization groups 32 values along the head dimension; columnwise + quantization groups 32 sequence positions. The two representations must be + quantized independently from the higher-precision operand. + """ + return ((True, backward), (True, backward), (backward, not backward)) + + +def support_error(sm_arch, backend_version, qkv_format): + """Return shared hardware, backend-defect, and layout exclusions.""" + if sm_arch < 100: + return "MXFP8 attention requires SM100 or later." + if backend_version < 92100: + return "MXFP8 attention requires cuDNN 9.21.0 or later." + if backend_version in (92300, 92301): + return "MXFP8 attention is disabled for known cuDNN 9.23.0/9.23.1 defects." + if qkv_format == "thd": + return "MXFP8 attention does not support THD format." + return "" + + +def shape_error(seqlen_q, seqlen_kv, head_dim_qk, head_dim_v): + """Return the alignment requirements of the existing preparation kernels.""" + if seqlen_q % 128 or seqlen_kv % 128 or head_dim_qk % 32 or head_dim_v % 32: + return ( + "MXFP8 attention requires sequence lengths divisible by 128 and head " + "dimensions divisible by 32. " + f"Received {seqlen_q=}, {seqlen_kv=}, {head_dim_qk=}, {head_dim_v=}." + ) + return "" + + +def scale_reordering(shape, columnwise=False): + """Describe padding and the F8_128x4 permutation of logical BHSD scales. + + Return the padded BHSD shape, the tiled reshape, and its permutation. The + frontend applies these operations to device arrays. Each batch/head matrix + is reordered independently, including columnwise scales. + """ + batch, heads, rows, cols = shape + if columnwise: + rows, cols = cols, rows + rows, cols = round_up(rows, 128), round_up(cols, 4) + padded_shape = (batch, heads, cols, rows) if columnwise else (batch, heads, rows, cols) + tiled_shape = (batch, heads, rows // 128, 4, 32, cols // 4, 4) + return padded_shape, tiled_shape, (0, 1, 2, 5, 4, 3, 6) diff --git a/transformer_engine/pytorch/attention/dot_product_attention/utils.py b/transformer_engine/pytorch/attention/dot_product_attention/utils.py index 9ed286b44a6..ca582956cf5 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/utils.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/utils.py @@ -36,6 +36,12 @@ META_S, META_DP, ) +from transformer_engine.common.cudnn.mxfp8 import ( + quantization_layouts as mxfp8_quantization_layouts, + scale_reordering as mxfp8_scale_reordering, + shape_error as mxfp8_shape_error, + support_error as mxfp8_support_error, +) from transformer_engine.pytorch.attention.inference import InferenceParams from transformer_engine.pytorch.cpu_offload import is_cpu_offload_enabled from transformer_engine.pytorch.quantized_tensor import QuantizedTensorStorage @@ -784,23 +790,16 @@ def _disable_all_flash_attention() -> None: ) use_fused_attention = False if use_fused_attention and fp8_recipe.mxfp8(): - if device_compute_capability < (10, 0): - logger.debug("Disabling FusedAttention for MXFP8 on arch < sm100") - use_fused_attention = False - elif fp8_recipe.fp8_mha: - logger.debug("Disabling FusedAttention for MXFP8 with fp8_mha=True") + reason = mxfp8_support_error( + 10 * device_compute_capability[0] + device_compute_capability[1], + 10000 * cudnn_version[0] + 100 * cudnn_version[1] + cudnn_version[2], + qkv_format, + ) + if fp8_recipe.fp8_mha: + reason = "MXFP8 attention does not support fp8_mha=True." + if reason: + logger.debug("Disabling FusedAttention: %s", reason) use_fused_attention = False - else: - if cudnn_version < (9, 21, 0): - logger.debug("Disabling FusedAttention for MXFP8 with cuDNN < 9.21.0") - use_fused_attention = False - elif cudnn_version in ((9, 23, 0), (9, 23, 1)): - # 9.23.0/9.23.1: known bugs with MXFP8 SDPA - logger.debug("Disabling FusedAttention for MXFP8 with cuDNN 9.23.0/9.23.1") - use_fused_attention = False - elif qkv_format == "thd": - logger.debug("Disabling FusedAttention for MXFP8 with qkv_format = thd") - use_fused_attention = False if use_fused_attention and (fp8_recipe.float8_block_scaling() or fp8_recipe.nvfp4()): logger.debug("Disabling FusedAttention for %s", fp8_recipe.__class__.__name__) use_fused_attention = False @@ -3457,25 +3456,21 @@ def mxfp8_transpose_swizzle(fp8_tensors, src_format): rs_list = [t._rowwise_scale_inv for t in fp8_tensors] cs_list = [t._columnwise_scale_inv for t in fp8_tensors] - def _align_up(x, a): - return ((x + a - 1) // a) * a - - def _bhsd_shape(src_4d, d_pad): + def _bhsd_shape(src_4d): if src_format == "sbhd": S, B, H, _ = src_4d.shape else: B, S, H, _ = src_4d.shape - return (B, H, S, d_pad) + return (B, H, S, src_4d.shape[-1]) - def _build_outputs(scale_list, alignment): + def _build_outputs(scale_list, columnwise): entries = [] total = 0 for s in scale_list: if s is None: entries.append(None) continue - d_pad = _align_up(s.shape[-1], alignment) - shape = _bhsd_shape(s, d_pad) + shape, _, _ = mxfp8_scale_reordering(_bhsd_shape(s), columnwise) numel = 1 for dim in shape: numel *= dim @@ -3487,8 +3482,8 @@ def _build_outputs(scale_list, alignment): buf = torch.empty(total, dtype=torch.uint8, device=device) return [buf[e[0] : e[0] + e[1]].view(e[2]) if e is not None else None for e in entries] - rs_outs = _build_outputs(rs_list, 4) - cs_outs = _build_outputs(cs_list, 128) + rs_outs = _build_outputs(rs_list, False) + cs_outs = _build_outputs(cs_list, True) rs_permuted = tex.multi_tensor_transpose_to_bhsd( rs_list, original_format=src_format, outputs=rs_outs @@ -3541,10 +3536,8 @@ def combine_and_quantize( d_v = v.shape[_d_dim[qkv_format]] s_q = q.shape[_s_dim[q_format]] s_kv = v.shape[_s_dim[kv_format]] - assert s_q % 128 == 0 and s_kv % 128 == 0 and d_qk % 32 == 0 and d_v % 32 == 0, ( - "MXFP8 quantization requires s_q % 128 == 0, s_kv % 128 == 0, d_qk % 32 == 0, d_v % 32" - f" == 0. Found {s_q=}, {s_kv=}, {d_qk=}, {d_v=}." - ) + alignment_error = mxfp8_shape_error(s_q, s_kv, d_qk, d_v) + assert not alignment_error, alignment_error if qkv_layout not in ("bshd_bshd_bshd", "sbhd_sbhd_sbhd"): keep_same_data_and_scale_inv_format = True @@ -3552,20 +3545,13 @@ def combine_and_quantize( # ---- Fast path: quantize in original layout, permute scale_inv to BHSD, then swizzle ---- if not keep_same_data_and_scale_inv_format: q_quantizer, k_quantizer, v_quantizer = [qkv_quantizer.copy() for _ in range(3)] - if used_in_forward and not used_in_backward: - q_quantizer.rowwise_usage = True - q_quantizer.columnwise_usage = False - k_quantizer.rowwise_usage = True - k_quantizer.columnwise_usage = False - v_quantizer.rowwise_usage = False - v_quantizer.columnwise_usage = True - elif (not used_in_forward) and used_in_backward: - q_quantizer.rowwise_usage = True - q_quantizer.columnwise_usage = True - k_quantizer.rowwise_usage = True - k_quantizer.columnwise_usage = True - v_quantizer.rowwise_usage = True - v_quantizer.columnwise_usage = False + if used_in_forward != used_in_backward: + for quantizer, (rowwise, columnwise) in zip( + (q_quantizer, k_quantizer, v_quantizer), + mxfp8_quantization_layouts(backward=used_in_backward), + ): + quantizer.rowwise_usage = rowwise + quantizer.columnwise_usage = columnwise (q_fp8, k_fp8, v_fp8), qkv_scale_inv_format = mxfp8_quantize_fast_path( [(q, q_quantizer), (k, k_quantizer), (v, v_quantizer)], qkv_format ) @@ -3588,20 +3574,13 @@ def combine_and_quantize( q, k, v = [x.view(-1, x.shape[-1]) for x in [q, k, v]] q_quantizer, k_quantizer, v_quantizer = [qkv_quantizer.copy() for _ in range(3)] - if used_in_forward and not used_in_backward: - q_quantizer.rowwise_usage = True - q_quantizer.columnwise_usage = False - k_quantizer.rowwise_usage = True - k_quantizer.columnwise_usage = False - v_quantizer.rowwise_usage = False - v_quantizer.columnwise_usage = True - elif (not used_in_forward) and used_in_backward: - q_quantizer.rowwise_usage = True - q_quantizer.columnwise_usage = True - k_quantizer.rowwise_usage = True - k_quantizer.columnwise_usage = True - v_quantizer.rowwise_usage = True - v_quantizer.columnwise_usage = False + if used_in_forward != used_in_backward: + for quantizer, (rowwise, columnwise) in zip( + (q_quantizer, k_quantizer, v_quantizer), + mxfp8_quantization_layouts(backward=used_in_backward), + ): + quantizer.rowwise_usage = rowwise + quantizer.columnwise_usage = columnwise q_fp8, k_fp8, v_fp8 = [ quant(x) for quant, x in zip([q_quantizer, k_quantizer, v_quantizer], [q, k, v]) ] From da5c01127f1b9747fc159c66d51c99ed5002022f Mon Sep 17 00:00:00 2001 From: Vladimir Cherepanov Date: Tue, 20 Oct 2026 05:09:00 +0000 Subject: [PATCH 11/49] Add dense MXFP8 attention to JAX Expose fp8="mxfp8" through fused_attn and Flax modules. Prepare independent contraction-axis quantizations and per-head reordered E8M0 scales using existing common quantization kernels, then execute shared MXFP8 graphs through the existing JAX FFI executor. Validation on four GB200 GPUs with cuDNN 9.26 / Frontend 1.29: 30 focused MXFP8 tests and 15 existing attention regressions pass. Changed-file pre-commit and focused JAX pylint pass. Full license checking encounters existing failures including untracked nccl_ep headers. Signed-off-by: Vladimir Cherepanov --- docs/examples/jax/attention.rst | 11 +- tests/jax/test_mxfp8_attention.py | 127 ++++++++++++++++++ transformer_engine/jax/attention.py | 19 +-- .../jax/cpp_extensions/attention.py | 10 +- .../jax/cpp_extensions/cudnn_attention.py | 70 ++++++++-- .../jax/cpp_extensions/mxfp8_attention.py | 69 ++++++++++ transformer_engine/jax/flax/transformer.py | 21 +-- 7 files changed, 292 insertions(+), 35 deletions(-) create mode 100644 transformer_engine/jax/cpp_extensions/mxfp8_attention.py diff --git a/docs/examples/jax/attention.rst b/docs/examples/jax/attention.rst index 7ed82edff5d..88b2a43f2b2 100644 --- a/docs/examples/jax/attention.rst +++ b/docs/examples/jax/attention.rst @@ -51,8 +51,15 @@ 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 +Pass ``fp8="mxfp8"`` to select MXFP8 attention on SM100 or later. MXFP8 uses +E4M3 operands with E8M0 block scales along the attention contraction axes in +both passes. Inputs, outputs and gradients remain FP16/BF16. Sequence lengths +must be divisible by 128 and head dimensions by 32. This mode currently supports +separate BSHD inputs; THD and cuDNN 9.23.0/9.23.1 are excluded. It requires +cuDNN 9.21 or later. ``NVTE_FP8_DPA_BWD=0`` keeps the original FP16/BF16 operands +for MXFP8 backward, matching the PyTorch MXFP8 path. + +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 diff --git a/tests/jax/test_mxfp8_attention.py b/tests/jax/test_mxfp8_attention.py index adf614b3d18..da1790326b1 100644 --- a/tests/jax/test_mxfp8_attention.py +++ b/tests/jax/test_mxfp8_attention.py @@ -6,10 +6,23 @@ import math +import jax +import jax.numpy as jnp +import numpy as np import pytest from transformer_engine.common.cudnn import mxfp8 from transformer_engine.common.cudnn.graph import mxfp8_scale_spec +from transformer_engine.jax.attention import ( + AttnBiasType, + AttnMaskType, + AttnSoftmaxType, + QKVLayout, + SequenceDescriptor, + fused_attn, +) +from transformer_engine.jax.cpp_extensions import cudnn_attention +from transformer_engine.jax.cpp_extensions import mxfp8_attention @pytest.mark.parametrize("columnwise", (False, True)) @@ -53,3 +66,117 @@ def test_quantization_directions(): assert not mxfp8.shape_error(128, 256, 64, 96) assert mxfp8.shape_error(129, 256, 64, 96) assert mxfp8.shape_error(128, 256, 64, 80) + + +def _require_mxfp8(): + reason = mxfp8.support_error( + cudnn_attention.device_capabilities()[0], cudnn_attention.cudnn_version(), "bshd" + ) + if reason: + pytest.skip(reason) + + +def _attention(qkv, **kwargs): + options = dict( + attn_bias_type=AttnBiasType.NO_BIAS, + attn_mask_type=AttnMaskType.NO_MASK, + softmax_type=AttnSoftmaxType.VANILLA_SOFTMAX, + qkv_layout=QKVLayout.BSHD_BSHD_BSHD, + scaling_factor=0.125, + dropout_probability=0.0, + is_training=True, + fp8="mxfp8", + ) + options.update(kwargs) + return fused_attn(qkv, None, SequenceDescriptor(), None, **options) + + +def _reference(qkv): + q, k, v = (x.astype(jnp.float32) for x in qkv) + repeats = q.shape[-2] // k.shape[-2] + k, v = (jnp.repeat(x, repeats, axis=-2) for x in (k, v)) + scores = jnp.einsum("bqhd,bkhd->bhqk", q, k) * 0.125 + return jnp.einsum("bhqk,bkhd->bqhd", jax.nn.softmax(scores, -1), v) + + +@pytest.mark.parametrize("dtype", (jnp.float16, jnp.bfloat16)) +@pytest.mark.parametrize("fp8_backward", (False, True)) +@pytest.mark.parametrize("is_training", (False, True)) +def test_dense_forward_backward(dtype, fp8_backward, is_training, monkeypatch): + """Exactly representable inputs/probabilities retain model-precision accuracy.""" + _require_mxfp8() + monkeypatch.setenv("NVTE_FP8_DPA_BWD", "1" if fp8_backward else "0") + q = jnp.full((1, 128, 2, 64), 0.5, dtype) + k = jnp.full((1, 128, 1, 64), 0.5, dtype) + v = jnp.zeros_like(k).at[:, ::2, :, 0].set(1.0) + inputs = (q, k, v) + attention = lambda xs: _attention(xs, is_training=is_training) + actual = jax.jit(attention)(inputs) + expected = jax.jit(_reference)(inputs) + np.testing.assert_allclose(actual, expected, rtol=5e-3, atol=5e-4) + assert actual.dtype == dtype + grads = jax.jit(jax.grad(lambda xs: attention(xs).sum()))(inputs) + reference_grads = jax.jit(jax.grad(lambda xs: _reference(xs).sum()))(inputs) + for actual_grad, expected_grad in zip(grads, reference_grads): + assert actual_grad.dtype == dtype + np.testing.assert_allclose(actual_grad, expected_grad, rtol=1e-2, atol=1e-3) + + +@pytest.mark.parametrize("columnwise", (False, True)) +@pytest.mark.parametrize("head_dim", (64, 96, 160)) +def test_quantization_and_scale_bytes(columnwise, head_dim): + """Native quantization and cuDNN scale packing agree with an independent reference.""" + _require_mxfp8() + rng = np.random.default_rng(17) + value = jnp.asarray(rng.normal(size=(2, 3, 128, head_dim)), jnp.bfloat16) + value *= jnp.asarray([0.01, 1.0, 100.0], jnp.bfloat16)[None, :, None, None] + row, col = jax.jit(lambda x: mxfp8_attention.quantize_operand(x, not columnwise, columnwise))( + value + ) + tensor = col if columnwise else row + source = np.asarray(value, np.float32) + oriented = source.swapaxes(-1, -2) if columnwise else source + blocks = oriented.reshape(*oriented.shape[:-1], -1, 32) + maximum = np.max(np.abs(blocks), axis=-1) + exponent = np.maximum(np.ceil(np.log2(np.maximum(maximum / 448.0, 2.0**-127))), -127) + scales = np.exp2(exponent) + encoded = (blocks / scales[..., None]).astype(jnp.float8_e4m3fn).reshape(oriented.shape) + if columnwise: + encoded = encoded.swapaxes(-1, -2) + np.testing.assert_array_equal(np.asarray(tensor.data).view(np.uint8), encoded.view(np.uint8)) + scales = (exponent + 127).astype(np.uint8) + rows, cols = scales.shape[-2:] + padded_rows = ((rows + 127) // 128) * 128 + padded_cols = ((cols + 3) // 4) * 4 + expected = np.zeros((*scales.shape[:2], padded_rows * padded_cols), np.uint8) + # Physical offset within each 128x4 tile, derived independently of the + # reshape/transpose used by the implementation. + for i in range(rows): + for j in range(cols): + index = ((i // 128) * (padded_cols // 4) + j // 4) * 512 + index += (i % 32) * 16 + ((i % 128) // 32) * 4 + j % 4 + expected[..., index] = scales[..., i, j] + actual = jax.jit(lambda s: mxfp8_attention.reorder_scale(s, columnwise))(tensor.scale_inv) + np.testing.assert_array_equal(np.asarray(actual).reshape(expected.shape), expected) + + +@pytest.mark.parametrize("deterministic", (False, True)) +def test_flax_mxfp8_mode(deterministic): + """Flax forwards the requested block-scaling mode without coercing it to bool.""" + from transformer_engine.jax.flax import DotProductAttention + + _require_mxfp8() + module = DotProductAttention( + head_dim=64, + num_attention_heads=2, + attn_mask_type="no_mask", + transpose_batch_sequence=False, + fp8="mxfp8", + ) + q = jnp.full((1, 128, 2, 64), 0.5, jnp.bfloat16) + value = jnp.zeros_like(q).at[:, ::2, :, 0].set(1.0) + variables = module.init(jax.random.key(0), q, q, value, deterministic=deterministic) + forward = lambda v: module.apply(variables, q, q, v, deterministic=deterministic) + actual = jax.jit(forward)(value) + np.testing.assert_allclose(actual[..., 0], 0.5, rtol=5e-3, atol=5e-4) + np.testing.assert_allclose(jax.jit(jax.grad(lambda v: forward(v).sum()))(value), 1.0) diff --git a/transformer_engine/jax/attention.py b/transformer_engine/jax/attention.py index 2765e1fbb5b..02c23d7f531 100644 --- a/transformer_engine/jax/attention.py +++ b/transformer_engine/jax/attention.py @@ -348,7 +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, + fp8: bool | str = False, ): """ To check whether the fused attention kernel is supported. @@ -1144,7 +1144,7 @@ def _legacy_fused_attn( context_parallel_axis: str = "", softmax_offset: Optional[jnp.ndarray] = None, return_max_logit: bool = False, - fp8: bool = False, + fp8: bool | str = False, ): """ Perform non-THD (non-packed) cuDNN fused attention. @@ -1266,7 +1266,7 @@ def fused_attn_thd( context_parallel_axis: str = "", softmax_offset: Optional[jnp.ndarray] = None, return_max_logit: bool = False, - fp8: bool = False, + fp8: bool | str = False, ): """ Deprecated THD fused attn, please use fusd_attn with SequenceDescriptor @@ -1353,7 +1353,7 @@ def _fused_attn( context_checkpoint_name: str = "context", stripe_size: int | None = None, return_max_logit: bool = False, - fp8: bool = False, + fp8: bool | str = False, ): output, _ = _fused_attn_fwd_rule( qkv, @@ -1592,7 +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, + fp8: bool | str = False, ): """ Perform cuDNN fused attention. @@ -1649,9 +1649,10 @@ 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. + fp8 (bool | str): True enables tensor-scaled cuDNN FP8 SDPA; "mxfp8" + enables MXFP8 block-scaled attention on SM100 or later. Outputs and + gradients retain the input FP16/BF16 dtype. MXFP8 uses E4M3 operands + in both passes, with independent scales along the contraction axes. return_max_logit (bool): If True, also return per-head maximum attention logits with shape ``[h]``. Returns: @@ -1689,6 +1690,8 @@ def fused_attn( AttnBiasType.NO_BIAS, AttnMaskType.PADDING_CAUSAL_MASK, QKVLayout.T3HD, 0.125, 0, True, 3) """ + if fp8 not in (False, True, "mxfp8"): + raise ValueError('fp8 must be False, True, or "mxfp8".') if score_mod is None: score_mod_only_args = [ name diff --git a/transformer_engine/jax/cpp_extensions/attention.py b/transformer_engine/jax/cpp_extensions/attention.py index a5b1a747725..22cf802ea99 100644 --- a/transformer_engine/jax/cpp_extensions/attention.py +++ b/transformer_engine/jax/cpp_extensions/attention.py @@ -136,7 +136,7 @@ class _FusedAttnConfig: ) # Only for CP + Striped. For Ring P2P, stripe_size=1 only.For AG, stripe_size>=1. return_max_logit: bool = False allow_fast_causal_path: bool = True - fp8: bool = False + fp8: bool | str = False @property def effective_window_size(self) -> Tuple[int, int]: @@ -176,7 +176,7 @@ class FusedAttnHelper: bias_seqlen_q: Optional[int] = None bias_seqlen_kv: Optional[int] = None max_segments_per_seq: int = 1 - fp8: bool = False + fp8: bool | str = False def is_fused_attn_kernel_available(self): """Check if there is available fused attention kernel""" @@ -193,6 +193,8 @@ 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. """ + if self.fp8 not in (False, True, "mxfp8"): + raise ValueError('fp8 must be False, True, or "mxfp8".') if self.q_dtype != self.kv_dtype: raise ValueError("Q and KV must have the same data type.") backend = NVTE_Fused_Attn_Backend.NVTE_No_Backend @@ -3541,7 +3543,7 @@ def fused_attn_fwd( context_parallel_axis: str = "", stripe_size: int | None = None, return_max_logit: bool = False, - fp8: bool = False, + fp8: bool | str = False, ) -> jnp.ndarray: """ Perform the forward pass of with cuDNN fused attention implementations. @@ -3710,7 +3712,7 @@ def fused_attn_bwd( context_parallel_causal_load_balanced: bool = False, context_parallel_axis: str = "", stripe_size: int | None = None, - fp8: bool = False, + fp8: bool | str = False, ): """ Perform the backward pass of the cuDNN fused attention implementations. diff --git a/transformer_engine/jax/cpp_extensions/cudnn_attention.py b/transformer_engine/jax/cpp_extensions/cudnn_attention.py index f96f74f5b94..6f4c17ed622 100644 --- a/transformer_engine/jax/cpp_extensions/cudnn_attention.py +++ b/transformer_engine/jax/cpp_extensions/cudnn_attention.py @@ -29,9 +29,12 @@ check_graph_support, compile_graph, uses_ragged_attention_graph, + mxfp8_scale_spec, ) +from transformer_engine.common.cudnn import mxfp8 as mxfp8_rules from .cudnn_graph import import_cudnn, check_frontend_version, serialize_graph, graph_attributes +from . import mxfp8_attention def device_capabilities(): @@ -235,7 +238,20 @@ def _describe_graph(direction, shapes, config, capabilities, nondeterministic, f 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) + mxfp8 = fp8 and config.fp8 == "mxfp8" version = cudnn_version() + if config.fp8 == "mxfp8": + error = mxfp8_rules.support_error( + capabilities[0], version, "thd" if ragged else "bshd" + ) or mxfp8_rules.shape_error(sq, sk, dq, dv) + if error: + raise ValueError(error) + if config.qkv_layout.name != "BSHD_BSHD_BSHD": + raise ValueError("MXFP8 attention currently requires separate BSHD query/key/value.") + if config.return_max_logit: + raise ValueError("MXFP8 attention does not support return_max_logit=True.") + if config.attn_bias_type.name != "NO_BIAS": + raise ValueError("MXFP8 attention does not support attention bias or ALiBi.") if ragged and capabilities[0] < 90: raise ValueError("F16/BF16 fused attention with THD format requires sm90 or later.") if backward and not fp8 and ragged_stats: @@ -283,6 +299,9 @@ def output(name, dim, stride, slot, offset=0, dtype=io_dtype, ragged_offset=""): kdim = (graph_batch, hg, graph_sk, dq) vdim = (graph_batch, hg, graph_sk, dv) odim = (graph_batch, h, graph_sq, dv) + grad_strides = (qstride, kstride, vstride) + if mxfp8: + qstride, kstride, vstride = (_stride(dim) for dim in (qdim, kdim, vdim)) tensor("q", qdim, qstride, qkv_dtype, ragged_offset="offset_q" if ragged else "") tensor( "k", @@ -333,13 +352,26 @@ def output(name, dim, stride, slot, offset=0, dtype=io_dtype, ragged_offset=""): output("dsink", (1, h, 1, 1), (h, 1, 1, 1), 4, dtype="FLOAT") constant("attn_scale", config.scaling_factor) if fp8: - for name in ("descale_q", "descale_k", "descale_v"): - tensor(name, (1, 1, 1, 1), (1, 1, 1, 1), "FLOAT") - constant("scale_s", 448.0) - constant("descale_s", 1.0 / 448.0) + for name, dim in zip(("q", "k", "v"), (qdim, kdim, vdim)): + if mxfp8: + specs["descale_" + name] = mxfp8_scale_spec( + dim, "bhsd", columnwise=name == "v" and not backward + ) + inputs.append(("descale_" + name, "descale_" + name, 0)) + else: + tensor("descale_" + name, (1, 1, 1, 1), (1, 1, 1, 1), "FLOAT") + if not mxfp8: + constant("scale_s", 448.0) + constant("descale_s", 1.0 / 448.0) if backward: tensor("o", odim, ostride, ragged_offset=o_offset) - tensor("do", odim, ostride, "FP8_E5M2" if fp8 else io_dtype, ragged_offset=o_offset) + tensor( + "do", + odim, + _stride(odim) if mxfp8 else ostride, + "FP8_E4M3" if mxfp8 else "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 = ( @@ -354,13 +386,21 @@ def output(name, dim, stride, slot, offset=0, dtype=io_dtype, ragged_offset=""): for name, dim, stride, slot, offset, ragged_offset in zip( ("dq", "dk", "dv"), (qdim, kdim, vdim), - (qstride, kstride, vstride), + grad_strides, grad_slots, grad_offsets, ("offset_q", "offset_k", "offset_v"), ): output(name, dim, stride, slot, offset, ragged_offset=ragged_offset if ragged else "") - if fp8: + if mxfp8: + tensor("do_f16", odim, _stride(odim)) + for name, dim in (("q_t", qdim), ("k_t", kdim), ("do_t", odim)): + tensor(name, dim, _stride(dim), "FP8_E4M3") + specs["descale_" + name] = mxfp8_scale_spec(dim, "bhsd", columnwise=True) + inputs.append(("descale_" + name, "descale_" + name, 0)) + specs["descale_do"] = mxfp8_scale_spec(odim, "bhsd") + inputs.append(("descale_do", "descale_do", 0)) + elif fp8: tensor("descale_do", (1, 1, 1, 1), (1, 1, 1, 1), "FLOAT") for name in ("descale_o", "scale_dp", "descale_dp", "scale_dq", "scale_dk", "scale_dv"): constant(name, 1.0) @@ -371,13 +411,14 @@ def output(name, dim, stride, slot, offset=0, dtype=io_dtype, ragged_offset=""): output("stats", stats_dim, stats_stride, 1, dtype="FLOAT", ragged_offset=stats_offset) if config.return_max_logit: output("max", stats_dim, stats_stride, 2, dtype="FLOAT", ragged_offset=stats_offset) - if fp8: + if fp8 and not mxfp8: 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, + mxfp8=mxfp8, bottom_right=config.bottom_right_diagonal, left_bound=left + 1 if left >= 0 else None, right_bound=right if right >= 0 else None, @@ -475,7 +516,10 @@ def execute(direction, args, config, result_avals): plan.multipliers, ): values[name] = _offset64(offsets, multiplier) - if config.fp8: + if config.fp8 == "mxfp8": + if plan.fp8: + values.update(mxfp8_attention.prepare_qkv((q, k, v), backward)) + elif config.fp8: for name in ("q", "k", "v"): # Empty dummy operands are not quantized. if values[name].size: @@ -499,10 +543,12 @@ def execute(direction, args, config, result_avals): values["seq_q"], q.shape[-4] if config.qkv_layout.is_qkvpacked() else q.shape[-3], ) - if plan.fp8: + if plan.fp8 and config.fp8 == "mxfp8": + values.update(mxfp8_attention.prepare_doutput(do)) + elif plan.fp8: values["do"], values["descale_do"] = _quantize(do, jnp.float8_e5m2) specs = list(result_avals[:-1]) - if plan.fp8: + if plan.fp8 and config.fp8 != "mxfp8": specs += [jax.ShapeDtypeStruct((1,), jnp.float32)] * 4 else: specs = list(result_avals[:3]) @@ -512,7 +558,7 @@ def execute(direction, args, config, result_avals): specs[1] = jax.ShapeDtypeStruct(dense_shape, jnp.float32) if config.return_max_logit: specs[2] = jax.ShapeDtypeStruct(dense_shape, jnp.float32) - if plan.fp8: + if plan.fp8 and config.fp8 != "mxfp8": specs += [jax.ShapeDtypeStruct((1,), jnp.float32)] * 2 specs.append(jax.ShapeDtypeStruct((plan.graph.workspace_size,), jnp.uint8)) attributes = graph_attributes(plan.graph) diff --git a/transformer_engine/jax/cpp_extensions/mxfp8_attention.py b/transformer_engine/jax/cpp_extensions/mxfp8_attention.py new file mode 100644 index 00000000000..237ed580982 --- /dev/null +++ b/transformer_engine/jax/cpp_extensions/mxfp8_attention.py @@ -0,0 +1,69 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""Adapt JAX MXFP8 quantization buffers to the shared attention graph contract.""" + +import jax.numpy as jnp + +from transformer_engine.common.cudnn.mxfp8 import quantization_layouts, scale_reordering +from ..quantize import QuantizerFactory, QuantizeLayout, ScalingMode, ScaledTensor2x +from .quantization import quantize + + +def reorder_scale(scale, columnwise=False): + """Pad and reorder each logical BHSD scale matrix for cuDNN.""" + scale = scale.view(jnp.uint8) + padded, tiled, permutation = scale_reordering(scale.shape, columnwise) + scale = jnp.pad(scale, tuple((0, target - size) for size, target in zip(scale.shape, padded))) + if columnwise: + scale = scale.swapaxes(-2, -1) + return scale.reshape(tiled).transpose(permutation).reshape(padded) + + +def quantize_operand(value, rowwise=True, columnwise=False): + """Quantize a logical BHSD operand along its attention contraction axes.""" + layout = ( + QuantizeLayout.ROWWISE_COLWISE + if rowwise and columnwise + else QuantizeLayout.ROWWISE if rowwise else QuantizeLayout.COLWISE + ) + quantizer = QuantizerFactory.create( + scaling_mode=ScalingMode.MXFP8_1D_SCALING, + q_dtype=jnp.float8_e4m3fn, + q_layout=layout, + ) + result = quantize(value, quantizer) + if isinstance(result, ScaledTensor2x): + return result.rowwise_tensor, result.colwise_tensor + return (result, None) if rowwise else (None, result) + + +def prepare_qkv(qkv, backward=False): + """Return data and reordered scales named by the common graph builder.""" + values = {} + for name, value, (rowwise, columnwise) in zip( + ("q", "k", "v"), qkv, quantization_layouts(backward) + ): + value = value.reshape(-1, *value.shape[-3:]).transpose(0, 2, 1, 3) + row, col = quantize_operand(value, rowwise, columnwise) + primary = row if rowwise else col + values[name] = primary.data + values["descale_" + name] = reorder_scale(primary.scale_inv, not rowwise) + if rowwise and columnwise: + values[name + "_t"] = col.data + values["descale_" + name + "_t"] = reorder_scale(col.scale_inv, True) + return values + + +def prepare_doutput(value): + """Keep the original gradient and prepare both MXFP8 representations.""" + value = value.reshape(-1, *value.shape[-3:]).transpose(0, 2, 1, 3) + row, col = quantize_operand(value, True, True) + return { + "do": row.data, + "do_t": col.data, + "do_f16": value, + "descale_do": reorder_scale(row.scale_inv), + "descale_do_t": reorder_scale(col.scale_inv, True), + } diff --git a/transformer_engine/jax/flax/transformer.py b/transformer_engine/jax/flax/transformer.py index 453ceab49ed..13460b535d6 100644 --- a/transformer_engine/jax/flax/transformer.py +++ b/transformer_engine/jax/flax/transformer.py @@ -310,7 +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 + fp8: bool | str = False score_mod_requested: bool = False return_max_logit: bool = False @@ -371,7 +371,7 @@ def __call__( "return_max_logit": self.return_max_logit, } if self.fp8: - fused_attn_kwargs["fp8"] = True + fused_attn_kwargs["fp8"] = self.fp8 if self.qkv_layout.is_qkvpacked(): """qkvpacked format, treat @@ -635,10 +635,11 @@ 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 + fp8: bool | str, 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. + independent of the GEMM quantization recipe. Set to ``"mxfp8"`` for + block-scaled MXFP8 attention on SM100 or later. Optimization parameters ----------------------- @@ -668,7 +669,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 + fp8: bool | str = False def __post_init__(self): # TODO(KshitijLakhani): Remove warning in TransformerEngine v2.12 @@ -1290,9 +1291,10 @@ 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 + fp8: bool | str, 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. + Set to ``"mxfp8"`` for block-scaled MXFP8 attention. """ head_dim: int @@ -1335,7 +1337,7 @@ 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 + fp8: bool | str = False def __post_init__(self): # Deal with changed defaults in API @@ -2085,10 +2087,11 @@ 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 + fp8: bool | str, 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. + Set to ``"mxfp8"`` for block-scaled MXFP8 attention. Optimization parameters ----------------------- @@ -2157,7 +2160,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 + fp8: bool | str = False def __post_init__(self): if self.mha_kernel_init is None: From a32c9fa21132737927751623560e9ad6f40dcd96 Mon Sep 17 00:00:00 2001 From: Vladimir Cherepanov Date: Tue, 20 Oct 2026 05:26:32 +0000 Subject: [PATCH 12/49] Extend JAX MXFP8 attention layouts and feature coverage Support packed QKV/KV storage, GQA masks, sliding windows, sink softmax, unequal head dimensions, and deterministic execution. Keep original model-precision dO layout and use the common fused quantizer for columnwise operands to avoid constant-zero NaNs in the pure-JAX fallback. Validate 60 focused cases across these capabilities, FP16/BF16 backward, and checkpointing on GB200 with cuDNN 9.26. Add explicit SM12x and dropout exclusions. Changed-file pre-commit and focused JAX pylint pass. Signed-off-by: Vladimir Cherepanov --- docs/examples/jax/attention.rst | 4 +- tests/jax/test_mxfp8_attention.py | 173 +++++++++++++++++- transformer_engine/common/cudnn/mxfp8.py | 2 + .../jax/cpp_extensions/cudnn_attention.py | 26 ++- .../jax/cpp_extensions/mxfp8_attention.py | 18 +- 5 files changed, 202 insertions(+), 21 deletions(-) diff --git a/docs/examples/jax/attention.rst b/docs/examples/jax/attention.rst index 88b2a43f2b2..13a5f2e6b3a 100644 --- a/docs/examples/jax/attention.rst +++ b/docs/examples/jax/attention.rst @@ -55,9 +55,11 @@ Pass ``fp8="mxfp8"`` to select MXFP8 attention on SM100 or later. MXFP8 uses E4M3 operands with E8M0 block scales along the attention contraction axes in both passes. Inputs, outputs and gradients remain FP16/BF16. Sequence lengths must be divisible by 128 and head dimensions by 32. This mode currently supports -separate BSHD inputs; THD and cuDNN 9.23.0/9.23.1 are excluded. It requires +separate, QKV-packed and KV-packed BSHD inputs. THD, SM12x, and cuDNN +9.23.0/9.23.1 are excluded. It requires cuDNN 9.21 or later. ``NVTE_FP8_DPA_BWD=0`` keeps the original FP16/BF16 operands for MXFP8 backward, matching the PyTorch MXFP8 path. +Attention bias, ALiBi, dropout and ``return_max_logit`` are excluded from this mode. Paged KV-cache APIs and quantized module inputs/outputs require further JAX API and frontend-binding work. ``score_mod`` currently supports FP16/BF16 only. diff --git a/tests/jax/test_mxfp8_attention.py b/tests/jax/test_mxfp8_attention.py index da1790326b1..d5c810d9ec0 100644 --- a/tests/jax/test_mxfp8_attention.py +++ b/tests/jax/test_mxfp8_attention.py @@ -45,6 +45,8 @@ def test_scale_reordering_matches_graph(dim, columnwise): "sm,version,fmt,error", ( (90, 92600, "bshd", "SM100"), + (120, 92600, "bshd", "SM12x"), + (121, 92600, "bshd", "SM12x"), (100, 92000, "bshd", "9.21"), (100, 92300, "bshd", "defects"), (100, 92301, "bshd", "defects"), @@ -88,15 +90,27 @@ def _attention(qkv, **kwargs): fp8="mxfp8", ) options.update(kwargs) - return fused_attn(qkv, None, SequenceDescriptor(), None, **options) + descriptor = options.pop("sequence_descriptor", SequenceDescriptor()) + seed = options.pop("seed", None) + return fused_attn(qkv, None, descriptor, seed, **options) -def _reference(qkv): +def _reference(qkv, mask=None, sink=None): q, k, v = (x.astype(jnp.float32) for x in qkv) repeats = q.shape[-2] // k.shape[-2] k, v = (jnp.repeat(x, repeats, axis=-2) for x in (k, v)) scores = jnp.einsum("bqhd,bkhd->bhqk", q, k) * 0.125 - return jnp.einsum("bhqk,bkhd->bqhd", jax.nn.softmax(scores, -1), v) + if mask is not None: + scores = jnp.where(mask[:, None], scores, -1e30) + if sink is None: + probability = jax.nn.softmax(scores, -1) + else: + sink = jnp.broadcast_to(sink, (*scores.shape[:-1], 1)) + probability = jax.nn.softmax(jnp.concatenate((scores, sink), axis=-1), -1)[..., :-1] + output = jnp.einsum("bhqk,bkhd->bqhd", probability, v) + if mask is not None: + output = jnp.where(jnp.any(mask, axis=-1)[..., None, None], output, 0.0) + return output @pytest.mark.parametrize("dtype", (jnp.float16, jnp.bfloat16)) @@ -180,3 +194,156 @@ def test_flax_mxfp8_mode(deterministic): actual = jax.jit(forward)(value) np.testing.assert_allclose(actual[..., 0], 0.5, rtol=5e-3, atol=5e-4) np.testing.assert_allclose(jax.jit(jax.grad(lambda v: forward(v).sum()))(value), 1.0) + + +def _pack(qkv, layout): + q, k, v = qkv + if layout.is_qkvpacked(): + return (jnp.stack((q, k, v), axis=-3),) + if layout.is_kvpacked(): + return q, jnp.stack((k, v), axis=-3) + return qkv + + +def _unpack(qkv, layout): + if layout.is_qkvpacked(): + return tuple(qkv[0][..., i, :, :] for i in range(3)) + if layout.is_kvpacked(): + return qkv[0], qkv[1][..., 0, :, :], qkv[1][..., 1, :, :] + return qkv + + +@pytest.mark.parametrize("layout", (QKVLayout.BS3HD, QKVLayout.BSHD_BS2HD)) +@pytest.mark.parametrize("dtype", (jnp.float16, jnp.bfloat16)) +@pytest.mark.parametrize("fp8_backward", (False, True)) +def test_packed_forward_backward(layout, dtype, fp8_backward, monkeypatch): + """Packed result storage preserves every gradient slice and original layout.""" + _require_mxfp8() + monkeypatch.setenv("NVTE_FP8_DPA_BWD", "1" if fp8_backward else "0") + q = jnp.full((2, 128, 2, 64), 0.5, dtype) + kv = q if layout.is_qkvpacked() else q[:, :, :1] + value = jnp.zeros_like(kv).at[:, ::2, :, 0].set(1.0) + inputs = _pack((q, kv, value), layout) + attention = lambda xs: _attention(xs, qkv_layout=layout) + reference = lambda xs: _reference(_unpack(xs, layout)) + np.testing.assert_allclose(jax.jit(attention)(inputs), reference(inputs), rtol=5e-3, atol=5e-4) + actual = jax.jit(jax.grad(lambda xs: attention(xs).sum()))(inputs) + expected = jax.jit(jax.grad(lambda xs: reference(xs).sum()))(inputs) + for grad, ref_grad, operand in zip(actual, expected, inputs): + assert grad.shape == operand.shape + assert grad.dtype == dtype + np.testing.assert_allclose(grad, ref_grad, rtol=1e-2, atol=1e-3) + + +@pytest.mark.parametrize("kind", ("dense", "causal", "bottom_right", "padding", "window")) +@pytest.mark.parametrize("fp8_backward", (False, True)) +def test_gqa_masks(kind, fp8_backward, monkeypatch): + """Bound the error from cuDNN's internal FP8 probabilities and dP on nonuniform data.""" + _require_mxfp8() + monkeypatch.setenv("NVTE_FP8_DPA_BWD", "1" if fp8_backward else "0") + keys = jax.random.split(jax.random.key(123), 4) + sq, sk = (128, 256) if kind == "bottom_right" else (128, 128) + q = jax.random.choice(keys[0], jnp.array([-0.5, 0.5], jnp.bfloat16), (2, sq, 4, 64)) + k = jax.random.choice(keys[1], jnp.array([-0.5, 0.5], jnp.bfloat16), (2, sk, 2, 64)) + v = jax.random.choice(keys[2], jnp.array([-1.0, 1.0], jnp.bfloat16), k.shape) + weight = jax.random.choice(keys[3], jnp.array([-1.0, 1.0], jnp.bfloat16), q.shape) + qpos, kpos = jnp.arange(sq)[:, None], jnp.arange(sk)[None, :] + mask = jnp.ones((2, sq, sk), jnp.bool_) + options = {} + if kind in ("causal", "bottom_right"): + diagonal = sk - sq if kind == "bottom_right" else 0 + mask &= kpos <= qpos + diagonal + options["attn_mask_type"] = ( + AttnMaskType.CAUSAL_BOTTOM_RIGHT_MASK if diagonal else AttnMaskType.CAUSAL_MASK + ) + elif kind == "padding": + qlength, klength = jnp.array([112, 96]), jnp.array([96, 80]) + mask &= (qpos[None] < qlength[:, None, None]) & (kpos[None] < klength[:, None, None]) + options.update( + attn_mask_type=AttnMaskType.PADDING_MASK, + sequence_descriptor=SequenceDescriptor.from_seqlens((qlength, klength)), + ) + elif kind == "window": + mask &= (kpos >= qpos - 31) & (kpos <= qpos + 15) + options["window_size"] = (31, 15) + inputs = (q, k, v) + attention = lambda xs: _attention(xs, **options) + reference = lambda xs: _reference(xs, mask) + actual = jax.jit(attention)(inputs) + expected = jax.jit(reference)(inputs) + relative_error = lambda a, b: float(jnp.linalg.norm(a - b) / jnp.linalg.norm(b)) + assert relative_error(actual, expected) < 0.04 + actual_grads = jax.jit(jax.grad(lambda xs: (attention(xs) * weight).sum()))(inputs) + expected_grads = jax.jit(jax.grad(lambda xs: (reference(xs) * weight).sum()))(inputs) + for actual_grad, expected_grad in zip(actual_grads, expected_grads): + assert relative_error(actual_grad, expected_grad) < 0.08 + + +@pytest.mark.parametrize("softmax", tuple(AttnSoftmaxType)) +@pytest.mark.parametrize("deterministic", (False, True)) +def test_softmax_variants(softmax, deterministic, monkeypatch): + """MXFP8 attention preserves attention-sink outputs and learnable-sink gradients.""" + _require_mxfp8() + monkeypatch.setenv("NVTE_ALLOW_NONDETERMINISTIC_ALGO", "0" if deterministic else "1") + q = jnp.zeros((1, 128, 2, 64), jnp.bfloat16) + v = jnp.zeros_like(q).at[:, ::2, :, 0].set(1.0) + sink = jnp.zeros((1, 2, 1, 1), jnp.float32) + options = dict(softmax_type=softmax) + if softmax == AttnSoftmaxType.LEARNABLE_SOFTMAX: + options["softmax_offset"] = sink + actual = jax.jit(lambda xs: _attention(xs, **options))((q, q, v)) + denominator = 128 if softmax == AttnSoftmaxType.VANILLA_SOFTMAX else 129 + expected = jnp.zeros_like(actual).at[..., 0].set(64 / denominator) + np.testing.assert_allclose(actual, expected, rtol=1e-2, atol=5e-4) + gradients = jax.jit(jax.grad(lambda xs: _attention(xs, **options).sum()))((q, q, v)) + assert all(bool(jnp.all(jnp.isfinite(grad))) for grad in gradients) + np.testing.assert_allclose(gradients[2], 128 / denominator, rtol=1e-2, atol=5e-4) + if softmax == AttnSoftmaxType.LEARNABLE_SOFTMAX: + dsink = jax.jit( + jax.grad(lambda s: _attention((q, q, v), softmax_type=softmax, softmax_offset=s).sum()) + )(sink) + np.testing.assert_allclose(dsink, -128 * 64 / denominator**2, rtol=2e-2, atol=1e-3) + + +@pytest.mark.parametrize("head_dims", ((64, 128), (128, 64), (128, 128))) +def test_different_head_dimensions(head_dims): + """Per-head scales retain their direction when QK and V dimensions differ.""" + _require_mxfp8() + dq, dv = head_dims + q = jnp.full((1, 128, 2, dq), 0.5, jnp.bfloat16) + k = jnp.full((1, 256, 1, dq), 0.5, jnp.bfloat16) + v = jnp.zeros((1, 256, 1, dv), jnp.bfloat16).at[:, ::2, :, 0].set(1.0) + inputs = (q, k, v) + np.testing.assert_allclose(jax.jit(_attention)(inputs), _reference(inputs), atol=5e-4) + actual = jax.jit(jax.grad(lambda xs: _attention(xs).sum()))(inputs) + expected = jax.jit(jax.grad(lambda xs: _reference(xs).sum()))(inputs) + for grad, ref_grad in zip(actual, expected): + np.testing.assert_allclose(grad, ref_grad, rtol=1e-2, atol=1e-3) + + +def test_checkpoint_and_dropout_guard(monkeypatch): + """Recomputation preserves gradients and unsupported dropout has an actionable error.""" + _require_mxfp8() + monkeypatch.setenv("NVTE_FP8_DPA_BWD", "0") + q = jnp.zeros((1, 128, 2, 64), jnp.bfloat16) + v = jnp.zeros_like(q).at[..., 0].set(1.0) + inputs = (q, q, v) + forward = lambda xs, key: _attention(xs, seed=key) + key = jax.random.PRNGKey(41) + actual = jax.jit(forward)(inputs, key) + np.testing.assert_array_equal(actual, jax.jit(forward)(inputs, key)) + np.testing.assert_allclose(actual[..., 0], 1.0) + plain = jax.jit(jax.value_and_grad(lambda xs: forward(xs, key).sum()))(inputs) + recomputed = jax.jit(jax.value_and_grad(lambda xs: jax.checkpoint(forward)(xs, key).sum()))( + inputs + ) + for value, reference in zip(jax.tree.leaves(plain), jax.tree.leaves(recomputed)): + assert bool(jnp.all(jnp.isfinite(value))) + np.testing.assert_array_equal(value, reference) + with pytest.raises(ValueError, match="MXFP8 attention does not support dropout"): + jax.jit(lambda xs: _attention(xs, seed=key, dropout_probability=0.25))(inputs) + # Inference ignores the configured training dropout probability. + np.testing.assert_array_equal( + actual, + jax.jit(lambda xs: _attention(xs, dropout_probability=0.25, is_training=False))(inputs), + ) diff --git a/transformer_engine/common/cudnn/mxfp8.py b/transformer_engine/common/cudnn/mxfp8.py index 80d5626523f..7418c4f629e 100644 --- a/transformer_engine/common/cudnn/mxfp8.py +++ b/transformer_engine/common/cudnn/mxfp8.py @@ -21,6 +21,8 @@ def support_error(sm_arch, backend_version, qkv_format): """Return shared hardware, backend-defect, and layout exclusions.""" if sm_arch < 100: return "MXFP8 attention requires SM100 or later." + if 120 <= sm_arch < 130: + return "cuDNN MXFP8 attention is not supported on SM12x." if backend_version < 92100: return "MXFP8 attention requires cuDNN 9.21.0 or later." if backend_version in (92300, 92301): diff --git a/transformer_engine/jax/cpp_extensions/cudnn_attention.py b/transformer_engine/jax/cpp_extensions/cudnn_attention.py index 6f4c17ed622..b36b790a70e 100644 --- a/transformer_engine/jax/cpp_extensions/cudnn_attention.py +++ b/transformer_engine/jax/cpp_extensions/cudnn_attention.py @@ -246,12 +246,12 @@ def _describe_graph(direction, shapes, config, capabilities, nondeterministic, f ) or mxfp8_rules.shape_error(sq, sk, dq, dv) if error: raise ValueError(error) - if config.qkv_layout.name != "BSHD_BSHD_BSHD": - raise ValueError("MXFP8 attention currently requires separate BSHD query/key/value.") if config.return_max_logit: raise ValueError("MXFP8 attention does not support return_max_logit=True.") if config.attn_bias_type.name != "NO_BIAS": raise ValueError("MXFP8 attention does not support attention bias or ALiBi.") + if config.is_training and config.dropout_probability: + raise ValueError("cuDNN MXFP8 attention does not support dropout.") if ragged and capabilities[0] < 90: raise ValueError("F16/BF16 fused attention with THD format requires sm90 or later.") if backward and not fp8 and ragged_stats: @@ -308,8 +308,8 @@ def output(name, dim, stride, slot, offset=0, dtype=io_dtype, ragged_offset=""): kdim, kstride, qkv_dtype, - source="q" if qpack == 3 else "k", - offset=h * dq * itemsize if qpack == 3 else 0, + source="q" if qpack == 3 and not mxfp8 else "k", + offset=h * dq * itemsize if qpack == 3 and not mxfp8 else 0, ragged_offset="offset_k" if ragged else "", ) tensor( @@ -317,8 +317,12 @@ def output(name, dim, stride, slot, offset=0, dtype=io_dtype, ragged_offset=""): vdim, vstride, qkv_dtype, - source="q" if qpack == 3 else ("k" if kvpack == 2 else "v"), - offset=2 * h * dq * itemsize if qpack == 3 else (hg * dq * itemsize if kvpack == 2 else 0), + source="v" if mxfp8 else "q" if qpack == 3 else ("k" if kvpack == 2 else "v"), + offset=( + 0 + if mxfp8 + else 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) @@ -393,7 +397,7 @@ def output(name, dim, stride, slot, offset=0, dtype=io_dtype, ragged_offset=""): ): output(name, dim, stride, slot, offset, ragged_offset=ragged_offset if ragged else "") if mxfp8: - tensor("do_f16", odim, _stride(odim)) + tensor("do_f16", odim, ostride) for name, dim in (("q_t", qdim), ("k_t", kdim), ("do_t", odim)): tensor(name, dim, _stride(dim), "FP8_E4M3") specs["descale_" + name] = mxfp8_scale_spec(dim, "bhsd", columnwise=True) @@ -518,7 +522,13 @@ def execute(direction, args, config, result_avals): values[name] = _offset64(offsets, multiplier) if config.fp8 == "mxfp8": if plan.fp8: - values.update(mxfp8_attention.prepare_qkv((q, k, v), backward)) + if config.qkv_layout.is_qkvpacked(): + unpacked = tuple(q[..., index, :, :] for index in range(3)) + elif config.qkv_layout.is_kvpacked(): + unpacked = (q, k[..., 0, :, :], k[..., 1, :, :]) + else: + unpacked = (q, k, v) + values.update(mxfp8_attention.prepare_qkv(unpacked, backward)) elif config.fp8: for name in ("q", "k", "v"): # Empty dummy operands are not quantized. diff --git a/transformer_engine/jax/cpp_extensions/mxfp8_attention.py b/transformer_engine/jax/cpp_extensions/mxfp8_attention.py index 237ed580982..267c95270dc 100644 --- a/transformer_engine/jax/cpp_extensions/mxfp8_attention.py +++ b/transformer_engine/jax/cpp_extensions/mxfp8_attention.py @@ -23,11 +23,11 @@ def reorder_scale(scale, columnwise=False): def quantize_operand(value, rowwise=True, columnwise=False): """Quantize a logical BHSD operand along its attention contraction axes.""" - layout = ( - QuantizeLayout.ROWWISE_COLWISE - if rowwise and columnwise - else QuantizeLayout.ROWWISE if rowwise else QuantizeLayout.COLWISE - ) + # The public colwise-only wrapper uses a pure-JAX fallback. Its division + # by the minimum E8M0 scale can turn constant zero blocks into NaNs during + # XLA constant folding. Request the fused common kernel and discard its + # unused rowwise result when only columnwise data is needed. + layout = QuantizeLayout.ROWWISE_COLWISE if columnwise else QuantizeLayout.ROWWISE quantizer = QuantizerFactory.create( scaling_mode=ScalingMode.MXFP8_1D_SCALING, q_dtype=jnp.float8_e4m3fn, @@ -35,8 +35,8 @@ def quantize_operand(value, rowwise=True, columnwise=False): ) result = quantize(value, quantizer) if isinstance(result, ScaledTensor2x): - return result.rowwise_tensor, result.colwise_tensor - return (result, None) if rowwise else (None, result) + return result.rowwise_tensor if rowwise else None, result.colwise_tensor + return result, None def prepare_qkv(qkv, backward=False): @@ -58,8 +58,8 @@ def prepare_qkv(qkv, backward=False): def prepare_doutput(value): """Keep the original gradient and prepare both MXFP8 representations.""" - value = value.reshape(-1, *value.shape[-3:]).transpose(0, 2, 1, 3) - row, col = quantize_operand(value, True, True) + value = value.reshape(-1, *value.shape[-3:]) + row, col = quantize_operand(value.transpose(0, 2, 1, 3), True, True) return { "do": row.data, "do_t": col.data, From 950cf8d76e9151b3e7738ecbe3ce05b265fe9967 Mon Sep 17 00:00:00 2001 From: Vladimir Cherepanov Date: Tue, 20 Oct 2026 05:33:14 +0000 Subject: [PATCH 13/49] Reuse prequantized MXFP8 operands in JAX attention Accept independent rowwise/columnwise ScaledTensor2x operands alongside model-precision gradient targets. Consume supplied FP8 buffers directly in the shared cuDNN graph, reorder scales after partitioning, and preserve original operands for model-precision backward. Add quantize_attention_qkv, API documentation, and nine passing GB200 tests for packed layouts, both backward modes, supplied values, validation, and two-GPU data/tensor sharding. Focused JAX pylint and changed-file pre-commit pass. Signed-off-by: Vladimir Cherepanov --- docs/api/jax.rst | 1 + docs/examples/jax/attention.rst | 9 + tests/jax/test_mxfp8_attention.py | 77 ++++++++ transformer_engine/jax/attention.py | 43 +++++ .../jax/cpp_extensions/attention.py | 23 ++- .../jax/cpp_extensions/cudnn_attention.py | 8 +- .../jax/cpp_extensions/mxfp8_attention.py | 171 ++++++++++++++++++ 7 files changed, 326 insertions(+), 6 deletions(-) diff --git a/docs/api/jax.rst b/docs/api/jax.rst index 130e66bd5c1..08c99bd3b49 100644 --- a/docs/api/jax.rst +++ b/docs/api/jax.rst @@ -64,4 +64,5 @@ Attention functions ------------------- .. autoapifunction:: transformer_engine.jax.attention.fused_attn +.. autoapifunction:: transformer_engine.jax.attention.quantize_attention_qkv .. autoapifunction:: transformer_engine.jax.attention.is_fused_attn_kernel_available diff --git a/docs/examples/jax/attention.rst b/docs/examples/jax/attention.rst index 13a5f2e6b3a..12466cb334f 100644 --- a/docs/examples/jax/attention.rst +++ b/docs/examples/jax/attention.rst @@ -61,6 +61,15 @@ cuDNN 9.21 or later. ``NVTE_FP8_DPA_BWD=0`` keeps the original FP16/BF16 operand for MXFP8 backward, matching the PyTorch MXFP8 path. Attention bias, ALiBi, dropout and ``return_max_logit`` are excluded from this mode. +To reuse quantizations, call ``quantize_attention_qkv((q, k, v))`` and pass its +result to ``fused_attn(..., fp8="mxfp8", quantized_qkv=quantized)``. Each operand +is a JAX ``ScaledTensor2x`` with independent rowwise/columnwise E4M3 data and +unswizzled E8M0 scales in BHSD order. Equivalent operands from a producer may +be supplied directly. The original BSHD arrays supply the gradient targets and +the original FP16/BF16 operands for ``NVTE_FP8_DPA_BWD=0``. These auxiliary +buffers are non-differentiable. Scale reordering happens after partitioning, +so data and tensor parallel sharding preserve per-head scale boundaries. + Paged KV-cache APIs and quantized module inputs/outputs require further JAX API and frontend-binding work. ``score_mod`` currently supports FP16/BF16 only. diff --git a/tests/jax/test_mxfp8_attention.py b/tests/jax/test_mxfp8_attention.py index d5c810d9ec0..0d2ac95a6b0 100644 --- a/tests/jax/test_mxfp8_attention.py +++ b/tests/jax/test_mxfp8_attention.py @@ -20,9 +20,11 @@ QKVLayout, SequenceDescriptor, fused_attn, + quantize_attention_qkv, ) from transformer_engine.jax.cpp_extensions import cudnn_attention from transformer_engine.jax.cpp_extensions import mxfp8_attention +from transformer_engine.jax.sharding import MeshResource, global_shard_guard @pytest.mark.parametrize("columnwise", (False, True)) @@ -347,3 +349,78 @@ def test_checkpoint_and_dropout_guard(monkeypatch): actual, jax.jit(lambda xs: _attention(xs, dropout_probability=0.25, is_training=False))(inputs), ) + + +@pytest.mark.parametrize( + "layout", (QKVLayout.BS3HD, QKVLayout.BSHD_BS2HD, QKVLayout.BSHD_BSHD_BSHD) +) +@pytest.mark.parametrize("fp8_backward", (False, True)) +def test_prequantized_inputs(layout, fp8_backward, monkeypatch): + """Reuse supplied quantizations and return gradients to the model-precision arrays.""" + _require_mxfp8() + monkeypatch.setenv("NVTE_FP8_DPA_BWD", "1" if fp8_backward else "0") + q = jnp.full((1, 128, 2, 64), 0.5, jnp.bfloat16) + v = jnp.zeros_like(q).at[:, ::2, :, 0].set(1.0) + inputs = _pack((q, q, v), layout) + prepared = jax.jit(quantize_attention_qkv)((q, q, v)) + reference = lambda xs: _attention(xs, qkv_layout=layout) + expected = jax.jit(jax.value_and_grad(lambda xs: reference(xs).sum()))(inputs) + + def fail_requantization(*args, **kwargs): + raise AssertionError("Supplied QKV must be consumed without requantization") + + monkeypatch.setattr(mxfp8_attention, "prepare_qkv", fail_requantization) + forward = lambda xs, buffers: _attention(xs, qkv_layout=layout, quantized_qkv=buffers) + actual = jax.jit(jax.value_and_grad(lambda xs, ys: forward(xs, ys).sum()))(inputs, prepared) + for value, ref in zip(jax.tree.leaves(actual), jax.tree.leaves(expected)): + assert value.dtype == ref.dtype + np.testing.assert_array_equal(value, ref) + # In one compiled training step, buffer preparation is non-differentiable + # while the attention VJP still supplies model-precision QKV gradients. + integrated = jax.jit( + jax.value_and_grad( + lambda xs: forward(xs, quantize_attention_qkv(_unpack(xs, layout))).sum() + ) + )(inputs) + for value, ref in zip(jax.tree.leaves(integrated), jax.tree.leaves(expected)): + np.testing.assert_array_equal(value, ref) + + +@pytest.mark.parametrize("axis", ("dp", "tp")) +def test_prequantized_sharding(axis): + """Per-head scale packing happens after data/tensor parallel partitioning.""" + _require_mxfp8() + if jax.local_device_count() < 2: + pytest.skip("Requires two GPUs for attention sharding") + mesh = jax.sharding.Mesh(np.array(jax.local_devices()[:2]), (axis,)) + spec = (axis, None, None, None) if axis == "dp" else (None, None, axis, None) + sharding = jax.sharding.NamedSharding(mesh, jax.sharding.PartitionSpec(*spec)) + q = jnp.full((2, 128, 4, 64), 0.5, jnp.bfloat16) + kv = q[:, :, :2] + v = jnp.zeros_like(kv).at[:, ::2, :, 0].set(1.0) + inputs = (q, kv, v) + loss = lambda xs: _attention(xs, quantized_qkv=quantize_attention_qkv(xs)).sum() + expected = jax.jit(jax.value_and_grad(loss))(inputs) + with jax.set_mesh(mesh), global_shard_guard(MeshResource(**{axis + "_resource": axis})): + actual = jax.jit(jax.value_and_grad(loss), in_shardings=((sharding,) * 3,))(inputs) + for value, ref in zip(jax.tree.leaves(actual), jax.tree.leaves(expected)): + np.testing.assert_array_equal(value, ref) + + +def test_prequantized_values_and_validation(): + """Supplied FP8 values determine the output; incompatible modes fail at the API boundary.""" + _require_mxfp8() + q = jnp.zeros((1, 128, 2, 64), jnp.bfloat16) + v = q.at[..., 0].set(1.0) + prepared = jax.jit(quantize_attention_qkv)((q, q, v)) + output = jax.jit(lambda xs, ys: _attention(xs, quantized_qkv=ys))((q, q, q), prepared) + np.testing.assert_allclose(output, v) + with pytest.raises(ValueError, match="requires fp8"): + _attention((q, q, v), quantized_qkv=prepared, fp8=True) + with pytest.raises(ValueError, match="three separate BSHD"): + quantize_attention_qkv((q,)) + with pytest.raises(ValueError, match="FP16/BF16"): + quantize_attention_qkv((q.astype(jnp.float32), q, v)) + with pytest.raises(ValueError, match="unswizzled E4M3"): + swapped = prepared[0].__class__(prepared[2].colwise_tensor, prepared[2].rowwise_tensor) + _attention((q, q, v), quantized_qkv=(prepared[0], prepared[0], swapped)) diff --git a/transformer_engine/jax/attention.py b/transformer_engine/jax/attention.py index 02c23d7f531..932c2cdacda 100644 --- a/transformer_engine/jax/attention.py +++ b/transformer_engine/jax/attention.py @@ -1354,6 +1354,7 @@ def _fused_attn( stripe_size: int | None = None, return_max_logit: bool = False, fp8: bool | str = False, + quantized_qkv=None, ): output, _ = _fused_attn_fwd_rule( qkv, @@ -1377,6 +1378,7 @@ def _fused_attn( stripe_size=stripe_size, return_max_logit=return_max_logit, fp8=fp8, + quantized_qkv=quantized_qkv, ) return output @@ -1403,6 +1405,7 @@ def _fused_attn_fwd_rule( stripe_size, return_max_logit, fp8, + quantized_qkv=None, ): output, softmax_aux, rng_state, max_logit = tex.fused_attn_fwd( qkv, @@ -1425,6 +1428,7 @@ def _fused_attn_fwd_rule( stripe_size=stripe_size, return_max_logit=return_max_logit, fp8=fp8, + quantized_qkv=quantized_qkv, ) output = checkpoint_name(output, context_checkpoint_name) softmax_aux = checkpoint_name(softmax_aux, context_checkpoint_name) @@ -1439,6 +1443,7 @@ def _fused_attn_fwd_rule( rng_state, softmax_offset, output, + quantized_qkv, ) @@ -1473,6 +1478,7 @@ def _fused_attn_bwd_rule( rng_state, softmax_offset, output, + quantized_qkv, ) = ctx grad_qkv, grad_bias, grad_softmax_offset = tex.fused_attn_bwd( qkv, @@ -1497,6 +1503,7 @@ def _fused_attn_bwd_rule( context_parallel_axis=context_parallel_axis, stripe_size=stripe_size, fp8=fp8, + quantized_qkv=quantized_qkv, ) if attn_bias_type in (AttnBiasType.NO_BIAS, AttnBiasType.ALIBI): grad_bias = None @@ -1508,6 +1515,7 @@ def _fused_attn_bwd_rule( grad_softmax_offset, None, None, + None, ) @@ -1567,6 +1575,23 @@ def _fused_attn_score_mod_bwd_rule(config, context_checkpoint_name, ctx, dz): _fused_attn_score_mod.defvjp(_fused_attn_score_mod_fwd_rule, _fused_attn_score_mod_bwd_rule) +def quantize_attention_qkv(qkv): + """Quantize separate BSHD Q/K/V into attention-ready BHSD MXFP8 tensor pairs. + + The returned ``ScaledTensor2x`` operands can be passed as ``quantized_qkv`` + to :func:`fused_attn`. Both directions are quantized independently from the + original operands. Quantized buffers are non-differentiable auxiliary data; + attention gradients flow to the original ``qkv`` arrays. + """ + from .cpp_extensions.mxfp8_attention import quantize_qkv + + if len(qkv) != 3 or any(x.ndim != 4 for x in qkv): + raise ValueError("quantize_attention_qkv requires three separate BSHD arrays.") + if any(x.dtype not in (jnp.float16, jnp.bfloat16) for x in qkv): + raise ValueError("quantize_attention_qkv requires FP16/BF16 operands.") + return quantize_qkv(qkv) + + def fused_attn( qkv: Tuple[jnp.ndarray, ...], bias: Optional[jnp.ndarray], @@ -1593,6 +1618,7 @@ def fused_attn( score_mod_bprop_tensors: Optional[Mapping[str, Any]] = None, return_max_logit: bool = False, fp8: bool | str = False, + quantized_qkv=None, ): """ Perform cuDNN fused attention. @@ -1653,6 +1679,11 @@ def fused_attn( enables MXFP8 block-scaled attention on SM100 or later. Outputs and gradients retain the input FP16/BF16 dtype. MXFP8 uses E4M3 operands in both passes, with independent scales along the contraction axes. + quantized_qkv (Optional[Tuple[ScaledTensor2x, ...]]): Prequantized MXFP8 + Q/K/V, each with independent unswizzled BHSD rowwise/columnwise data + and E8M0 scales. Requires ``fp8="mxfp8"``. The original ``qkv`` + supplies shapes, model dtype, gradient targets and FP16/BF16 backward + operands. See :func:`quantize_attention_qkv`. return_max_logit (bool): If True, also return per-head maximum attention logits with shape ``[h]``. Returns: @@ -1692,6 +1723,17 @@ def fused_attn( """ if fp8 not in (False, True, "mxfp8"): raise ValueError('fp8 must be False, True, or "mxfp8".') + if quantized_qkv is not None: + if fp8 != "mxfp8": + raise ValueError('quantized_qkv requires fp8="mxfp8".') + if not isinstance(sequence_descriptor, SequenceDescriptor): + raise ValueError("quantized_qkv requires a SequenceDescriptor.") + from .cpp_extensions.mxfp8_attention import prequantized_buffers + + # Keep auxiliary buffers as plain arrays across the custom VJP. The + # quantized tensor constructors validate real arrays, whereas JAX may + # reconstruct cotangent pytrees with opaque sentinel leaves. + quantized_qkv = prequantized_buffers(qkv, quantized_qkv, qkv_layout) if score_mod is None: score_mod_only_args = [ name @@ -1802,5 +1844,6 @@ def fused_attn( stripe_size=stripe_size, return_max_logit=return_max_logit, fp8=fp8, + quantized_qkv=quantized_qkv, ) return output diff --git a/transformer_engine/jax/cpp_extensions/attention.py b/transformer_engine/jax/cpp_extensions/attention.py index 22cf802ea99..06ad5e0c0dc 100644 --- a/transformer_engine/jax/cpp_extensions/attention.py +++ b/transformer_engine/jax/cpp_extensions/attention.py @@ -3544,6 +3544,7 @@ def fused_attn_fwd( stripe_size: int | None = None, return_max_logit: bool = False, fp8: bool | str = False, + quantized_qkv=None, ) -> jnp.ndarray: """ Perform the forward pass of with cuDNN fused attention implementations. @@ -3678,14 +3679,21 @@ def fused_attn_fwd( primitive = FusedRingAttnFwdPrimitive.outer_primitive seq_desc_flatten, _ = jax.tree.flatten(sequence_descriptor) - output, softmax_aux, rng_state, max_logit = primitive.bind( + operands = ( *qkv_for_primitive, bias, softmax_offset, seed, *seq_desc_flatten, - config=fused_config, ) + if quantized_qkv is None: + output, softmax_aux, rng_state, max_logit = primitive.bind(*operands, config=fused_config) + else: + from .mxfp8_attention import execute_prequantized + + output, softmax_aux, rng_state, max_logit = execute_prequantized( + operands, quantized_qkv, fused_config, False + ) rng_state = with_sharding_constraint(rng_state, PartitionSpec(get_all_mesh_axes(), None)) return (output, softmax_aux, rng_state, max_logit) @@ -3713,6 +3721,7 @@ def fused_attn_bwd( context_parallel_axis: str = "", stripe_size: int | None = None, fp8: bool | str = False, + quantized_qkv=None, ): """ Perform the backward pass of the cuDNN fused attention implementations. @@ -3856,7 +3865,7 @@ def fused_attn_bwd( primitive = FusedRingAttnBwdPrimitive.outer_primitive seq_desc_flatten, _ = jax.tree.flatten(sequence_descriptor) - *qkv_grads, bias_grad, softmax_offset_grad = primitive.bind( + operands = ( *qkv_for_primitive, bias, softmax_offset, @@ -3865,6 +3874,12 @@ def fused_attn_bwd( output, doutput, *seq_desc_flatten, - config=fused_config, ) + if quantized_qkv is None: + results = primitive.bind(*operands, config=fused_config) + else: + from .mxfp8_attention import execute_prequantized + + results = execute_prequantized(operands, quantized_qkv, fused_config, True) + *qkv_grads, bias_grad, softmax_offset_grad = results return tuple(qkv_grads[: len(qkv)]), bias_grad, softmax_offset_grad diff --git a/transformer_engine/jax/cpp_extensions/cudnn_attention.py b/transformer_engine/jax/cpp_extensions/cudnn_attention.py index b36b790a70e..af11ec9066a 100644 --- a/transformer_engine/jax/cpp_extensions/cudnn_attention.py +++ b/transformer_engine/jax/cpp_extensions/cudnn_attention.py @@ -493,7 +493,7 @@ def _packed_stats_to_dense(value, offsets, lengths, rows): return dense.transpose(0, 2, 1, 3) -def execute(direction, args, config, result_avals): +def execute(direction, args, config, result_avals, prepared_qkv=None): """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] @@ -528,7 +528,11 @@ def execute(direction, args, config, result_avals): unpacked = (q, k[..., 0, :, :], k[..., 1, :, :]) else: unpacked = (q, k, v) - values.update(mxfp8_attention.prepare_qkv(unpacked, backward)) + values.update( + mxfp8_attention.prepare_qkv(unpacked, backward) + if prepared_qkv is None + else prepared_qkv + ) elif config.fp8: for name in ("q", "k", "v"): # Empty dummy operands are not quantized. diff --git a/transformer_engine/jax/cpp_extensions/mxfp8_attention.py b/transformer_engine/jax/cpp_extensions/mxfp8_attention.py index 267c95270dc..fd772954d40 100644 --- a/transformer_engine/jax/cpp_extensions/mxfp8_attention.py +++ b/transformer_engine/jax/cpp_extensions/mxfp8_attention.py @@ -4,11 +4,17 @@ """Adapt JAX MXFP8 quantization buffers to the shared attention graph contract.""" +from functools import partial + +import jax import jax.numpy as jnp +from jax.experimental.custom_partitioning import custom_partitioning +from jax.sharding import NamedSharding, PartitionSpec from transformer_engine.common.cudnn.mxfp8 import quantization_layouts, scale_reordering from ..quantize import QuantizerFactory, QuantizeLayout, ScalingMode, ScaledTensor2x from .quantization import quantize +from .base import _JAX_GSPMD_SUPPORTED def reorder_scale(scale, columnwise=False): @@ -67,3 +73,168 @@ def prepare_doutput(value): "descale_do": reorder_scale(row.scale_inv), "descale_do_t": reorder_scale(col.scale_inv, True), } + + +def quantize_qkv(qkv): + """Prepare separate BSHD operands with independent BHSD row/column quantizations.""" + return tuple( + # Quantized buffers are auxiliary inputs, not differentiation targets. + ScaledTensor2x( + *quantize_operand(jax.lax.stop_gradient(x).transpose(0, 2, 1, 3), True, True) + ) + for x in qkv + ) + + +def prequantized_buffers(qkv, quantized_qkv, layout): + """Validate and extract unswizzled BHSD MXFP8 operands without dequantizing them.""" + if layout.is_thd(): + raise ValueError("MXFP8 attention does not support THD format.") + if layout.is_qkvpacked(): + operands = tuple(qkv[0][..., i, :, :] for i in range(3)) + elif layout.is_kvpacked(): + operands = qkv[0], qkv[1][..., 0, :, :], qkv[1][..., 1, :, :] + else: + operands = qkv + if len(quantized_qkv) != 3: + raise ValueError("quantized_qkv must contain separate Q, K, and V ScaledTensor2x operands.") + buffers = [] + for value, tensor in zip(operands, quantized_qkv): + if value.ndim != 4 or not isinstance(tensor, ScaledTensor2x): + raise ValueError("Prequantized attention requires BSHD arrays and BHSD ScaledTensor2x.") + shape = (value.shape[0], value.shape[2], value.shape[1], value.shape[3]) + row, col = tensor.rowwise_tensor, tensor.colwise_tensor + for part, columnwise in ((row, False), (col, True)): + valid = ( + part.scaling_mode == ScalingMode.MXFP8_1D_SCALING, + part.data.dtype == jnp.float8_e4m3fn, + part.data.shape == shape, + part.scale_inv.dtype == jnp.float8_e8m0fnu, + part.flatten_axis == 3, + part.is_colwise == columnwise, + part.data_layout == "N", + not part.has_rht_applied, + ) + if not all(valid): + raise ValueError( + "quantized_qkv requires unswizzled E4M3 MXFP8 BHSD data with independent " + "rowwise/columnwise scales and no Hadamard transform." + ) + buffers.extend((row.data, col.data, row.scale_inv, col.scale_inv)) + return tuple(jax.lax.stop_gradient(x) for x in buffers) + + +def _prepared_values(buffers, backward): + values = {} + for i, name in enumerate(("q", "k", "v")): + row, col, row_scale, col_scale = buffers[4 * i : 4 * i + 4] + columnwise = name == "v" and not backward + values[name] = col if columnwise else row + values["descale_" + name] = reorder_scale( + col_scale if columnwise else row_scale, columnwise + ) + if backward and name != "v": + values[name + "_t"] = col + values["descale_" + name + "_t"] = reorder_scale(col_scale, True) + return values + + +def _execute_prequantized(args, buffers, config, backward): + # Local imports keep the quantization adapter independent of attention's + # primitive registration order. + from .attention import FusedAttnFwdPrimitive, FusedAttnBwdPrimitive, generate_cu_seqlen + from . import cudnn_attention + from ..attention import SequenceDescriptor + + base = FusedAttnBwdPrimitive if backward else FusedAttnFwdPrimitive + result_avals = base.abstract( + *(jax.core.ShapedArray(x.shape, x.dtype) for x in args), config=config + ) + start = 9 if backward else 6 + descriptor = SequenceDescriptor( + seqlens=args[start : start + 2], + seq_offsets=args[start + 2 : start + 4], + segment_ids=args[start + 4 : start + 6], + segment_pos=args[start + 6 : start + 8], + ) + lengths, _ = descriptor._get_seqlens_and_offsets( + config.attn_mask_type, + config.qkv_layout, + config.window_size, + config.max_segments_per_seq, + allow_fast_causal_path=config.allow_fast_causal_path, + ) + args = (*args[:start], *(generate_cu_seqlen(x.flatten()) for x in lengths), *args[start + 2 :]) + results = cudnn_attention.execute( + "bwd" if backward else "fwd", + args, + config, + result_avals, + prepared_qkv=_prepared_values(buffers, backward), + ) + if backward: + return results[:5] + output, stats, max_logit, rng, _ = results + return output, stats, rng, max_logit + + +@partial(custom_partitioning, static_argnums=(2, 3)) +def execute_prequantized(args, buffers, config, backward): + """Execute supplied quantizations with the standard attention gradient contract.""" + return _execute_prequantized(args, buffers, config, backward) + + +def _prequantized_partition(config, backward, mesh, arg_infos, result_infos): + from .attention import FusedAttnFwdPrimitive, FusedAttnBwdPrimitive + from ..sharding import get_mesh_axis_size + + if get_mesh_axis_size(config.cp_axis, mesh) > 1: + raise ValueError("Prequantized MXFP8 context parallelism requires local requantization.") + base = FusedAttnBwdPrimitive if backward else FusedAttnFwdPrimitive + _, _, out_shardings, arg_shardings = base.partition(config, mesh, arg_infos[0], result_infos) + data_shardings = [] + for index in range(3): + source = ( + 0 + if config.qkv_layout.is_qkvpacked() + else min(index, 1) if config.qkv_layout.is_kvpacked() else index + ) + spec = arg_shardings[source].spec + # Model inputs are BSHD (possibly packed); supplied buffers are BHSD. + prepared = NamedSharding(mesh, PartitionSpec(spec[0], spec[-2], spec[1], spec[-1])) + data_shardings.extend((prepared,) * 4) + + def impl(local_args, local_buffers): + result = _execute_prequantized(local_args, local_buffers, config, backward) + if backward and config.softmax_type.name == "LEARNABLE_SOFTMAX": + from ..sharding import all_reduce_sum_along_dp_fsdp + + result = (*result[:4], all_reduce_sum_along_dp_fsdp(result[4], mesh)) + return result + + return mesh, impl, out_shardings, (arg_shardings, tuple(data_shardings)) + + +def _prequantized_rule(config, backward, mesh, value_types, result_types): + from .attention import FusedAttnFwdPrimitive, FusedAttnBwdPrimitive + from jax.experimental.custom_partitioning import SdyShardingRule + + base = FusedAttnBwdPrimitive if backward else FusedAttnFwdPrimitive + count = 17 if backward else 14 + rule = base.shardy_sharding_rule(config, mesh, value_types[:count], result_types) + inputs = (*rule.operand_mappings, *((f"…{100 + i}",) for i in range(12))) + return SdyShardingRule(inputs, rule.result_mappings) + + +def _prequantized_infer(config, backward, mesh, arg_infos, result_infos): + from .attention import FusedAttnFwdPrimitive, FusedAttnBwdPrimitive + + base = FusedAttnBwdPrimitive if backward else FusedAttnFwdPrimitive + return base.infer_sharding_from_operands(config, mesh, arg_infos[0], result_infos) + + +execute_prequantized.def_partition( + partition=_prequantized_partition, + sharding_rule=_prequantized_rule, + **({"infer_sharding_from_operands": _prequantized_infer} if _JAX_GSPMD_SUPPORTED else {}), +) From f5efdfc2c6f0cdc8fc1cf5520d1ba783ed399476 Mon Sep 17 00:00:00 2001 From: Vladimir Cherepanov Date: Tue, 20 Oct 2026 05:50:21 +0000 Subject: [PATCH 14/49] Enable packed and prequantized MXFP8 context parallelism Decode supplied MXFP8 operands before model-precision Ring/AllGather communication and quantize each local attention step. Support packed storage, causal load balancing, prequantized Flax BSHD/SBHD calls, checkpointing, and vmapped training. Preserve scale boundaries for replicated, data-parallel, and tensor-parallel execution. Reject the confirmed cuDNN Frontend 1.29.0 backward scale-stride defect for Q/K heads wider than 128 in the shared graph builder. Model-precision backward remains available via NVTE_FP8_DPA_BWD=0. Validation on four GB200 GPUs (SM100), cuDNN 9.26 / Frontend 1.29.0: full focused suite passes 109 tests, followed by 15 final checks including the added replicated-mesh case (110 distinct MXFP8 cases total). Existing attention, score-modification, and build-tools regressions: 58 passed, one skipped. JAX Python lint, changed-file pre-commit, and changed-file license checks pass. Tree-wide license checking encounters existing failures; PyTorch runtime validation is unavailable in this environment. Signed-off-by: Vladimir Cherepanov --- docs/examples/jax/attention.rst | 18 +- tests/jax/test_mxfp8_attention.py | 193 +++++++++++++++++- transformer_engine/common/cudnn/attention.py | 14 ++ transformer_engine/jax/attention.py | 13 ++ .../jax/cpp_extensions/mxfp8_attention.py | 86 +++++++- transformer_engine/jax/flax/transformer.py | 11 + 6 files changed, 317 insertions(+), 18 deletions(-) diff --git a/docs/examples/jax/attention.rst b/docs/examples/jax/attention.rst index 12466cb334f..13489601830 100644 --- a/docs/examples/jax/attention.rst +++ b/docs/examples/jax/attention.rst @@ -60,6 +60,9 @@ separate, QKV-packed and KV-packed BSHD inputs. THD, SM12x, and cuDNN cuDNN 9.21 or later. ``NVTE_FP8_DPA_BWD=0`` keeps the original FP16/BF16 operands for MXFP8 backward, matching the PyTorch MXFP8 path. Attention bias, ALiBi, dropout and ``return_max_logit`` are excluded from this mode. +Exact head-dimension availability is checked by cuDNN. Frontend 1.29.0 has an +incorrect MXFP8 backward scale stride for Q/K heads wider than 128; these +configurations require ``NVTE_FP8_DPA_BWD=0``. This check protects both frontends. To reuse quantizations, call ``quantize_attention_qkv((q, k, v))`` and pass its result to ``fused_attn(..., fp8="mxfp8", quantized_qkv=quantized)``. Each operand @@ -69,9 +72,18 @@ be supplied directly. The original BSHD arrays supply the gradient targets and the original FP16/BF16 operands for ``NVTE_FP8_DPA_BWD=0``. These auxiliary buffers are non-differentiable. Scale reordering happens after partitioning, so data and tensor parallel sharding preserve per-head scale boundaries. - -Paged KV-cache APIs and quantized module inputs/outputs require further JAX API -and frontend-binding work. ``score_mod`` currently supports FP16/BF16 only. +Ring and AllGather context parallelism communicate FP16/BF16 operands and +quantize each local attention step. Prequantized inputs are decoded before CP +communication, as in the PyTorch MXFP8 implementation. Separate and packed +BSHD storage, GQA, and causal dual-chunk load balancing are supported. Every +local attention step must meet the 128-token alignment requirement; 256 tokens +per CP rank accommodates the half-sequence steps used by causal load balancing +and AllGather. The existing JAX CP restrictions on sliding windows, padding, +dropout and non-vanilla softmax still apply. + +Paged KV-cache APIs, quantized module outputs, and automatic quantized +projection-to-attention integration 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 diff --git a/tests/jax/test_mxfp8_attention.py b/tests/jax/test_mxfp8_attention.py index 0d2ac95a6b0..35eee619e9f 100644 --- a/tests/jax/test_mxfp8_attention.py +++ b/tests/jax/test_mxfp8_attention.py @@ -17,6 +17,7 @@ AttnBiasType, AttnMaskType, AttnSoftmaxType, + CPStrategy, QKVLayout, SequenceDescriptor, fused_attn, @@ -191,8 +192,12 @@ def test_flax_mxfp8_mode(deterministic): ) q = jnp.full((1, 128, 2, 64), 0.5, jnp.bfloat16) value = jnp.zeros_like(q).at[:, ::2, :, 0].set(1.0) - variables = module.init(jax.random.key(0), q, q, value, deterministic=deterministic) - forward = lambda v: module.apply(variables, q, q, v, deterministic=deterministic) + variables = module.init( + jax.random.key(0), q, q, value, SequenceDescriptor(), deterministic=deterministic + ) + forward = lambda v: module.apply( + variables, q, q, v, SequenceDescriptor(), deterministic=deterministic + ) actual = jax.jit(forward)(value) np.testing.assert_allclose(actual[..., 0], 0.5, rtol=5e-3, atol=5e-4) np.testing.assert_allclose(jax.jit(jax.grad(lambda v: forward(v).sum()))(value), 1.0) @@ -323,6 +328,26 @@ def test_different_head_dimensions(head_dims): np.testing.assert_allclose(grad, ref_grad, rtol=1e-2, atol=1e-3) +def test_large_head_backend_limits(monkeypatch): + """Reject the identified Frontend stride defect; retain valid model-precision backward.""" + _require_mxfp8() + cudnn = cudnn_attention.import_cudnn() + q = jnp.full((1, 128, 2, 192), 0.5, jnp.bfloat16) + k = jnp.full((1, 256, 1, 192), 0.5, q.dtype) + v = jnp.zeros((1, 256, 1, 128), q.dtype).at[:, ::2, :, 0].set(1.0) + inputs = (q, k, v) + if cudnn.__version__.split("+", 1)[0].split("-", 1)[0] == "1.29.0": + monkeypatch.setenv("NVTE_FP8_DPA_BWD", "1") + with pytest.raises(ValueError, match="incorrect scale strides"): + jax.jit(_attention)(inputs) + monkeypatch.setenv("NVTE_FP8_DPA_BWD", "0") + np.testing.assert_allclose(jax.jit(_attention)(inputs), _reference(inputs), atol=5e-4) + grads = jax.jit(jax.grad(lambda xs: _attention(xs).sum()))(inputs) + reference = jax.jit(jax.grad(lambda xs: _reference(xs).sum()))(inputs) + for grad, ref in zip(grads, reference): + np.testing.assert_allclose(grad, ref, rtol=1e-2, atol=1e-3) + + def test_checkpoint_and_dropout_guard(monkeypatch): """Recomputation preserves gradients and unsupported dropout has an actionable error.""" _require_mxfp8() @@ -386,14 +411,18 @@ def fail_requantization(*args, **kwargs): np.testing.assert_array_equal(value, ref) -@pytest.mark.parametrize("axis", ("dp", "tp")) +@pytest.mark.parametrize("axis", ("dp", "tp", "replica")) def test_prequantized_sharding(axis): """Per-head scale packing happens after data/tensor parallel partitioning.""" _require_mxfp8() if jax.local_device_count() < 2: pytest.skip("Requires two GPUs for attention sharding") mesh = jax.sharding.Mesh(np.array(jax.local_devices()[:2]), (axis,)) - spec = (axis, None, None, None) if axis == "dp" else (None, None, axis, None) + spec = ( + (axis, None, None, None) + if axis == "dp" + else (None, None, axis, None) if axis == "tp" else () + ) sharding = jax.sharding.NamedSharding(mesh, jax.sharding.PartitionSpec(*spec)) q = jnp.full((2, 128, 4, 64), 0.5, jnp.bfloat16) kv = q[:, :, :2] @@ -401,7 +430,8 @@ def test_prequantized_sharding(axis): inputs = (q, kv, v) loss = lambda xs: _attention(xs, quantized_qkv=quantize_attention_qkv(xs)).sum() expected = jax.jit(jax.value_and_grad(loss))(inputs) - with jax.set_mesh(mesh), global_shard_guard(MeshResource(**{axis + "_resource": axis})): + resource = MeshResource() if axis == "replica" else MeshResource(**{axis + "_resource": axis}) + with jax.set_mesh(mesh), global_shard_guard(resource): actual = jax.jit(jax.value_and_grad(loss), in_shardings=((sharding,) * 3,))(inputs) for value, ref in zip(jax.tree.leaves(actual), jax.tree.leaves(expected)): np.testing.assert_array_equal(value, ref) @@ -424,3 +454,156 @@ def test_prequantized_values_and_validation(): with pytest.raises(ValueError, match="unswizzled E4M3"): swapped = prepared[0].__class__(prepared[2].colwise_tensor, prepared[2].rowwise_tensor) _attention((q, q, v), quantized_qkv=(prepared[0], prepared[0], swapped)) + mesh = jax.sharding.Mesh(np.asarray(jax.local_devices()[:2]), ("cp",)) + with jax.set_mesh(mesh), global_shard_guard(MeshResource(cp_resource="cp")): + with pytest.raises(ValueError, match="does not support THD"): + _attention((q, q, v), qkv_layout=QKVLayout.THD_THD_THD, context_parallel_axis="cp") + + +@pytest.mark.parametrize("strategy", (CPStrategy.ALL_GATHER, CPStrategy.RING)) +@pytest.mark.parametrize("cp_size", (2, 4)) +@pytest.mark.parametrize("kind", ("dense", "causal", "balanced_causal")) +@pytest.mark.parametrize("fp8_backward", (False, True)) +def test_context_parallel(strategy, cp_size, kind, fp8_backward, monkeypatch): + """Context-parallel outputs and gradients agree with an independent global reference.""" + from transformer_engine.jax.cpp_extensions.attention import reorder_causal_dual_chunk_swap + + _require_mxfp8() + if jax.local_device_count() < cp_size: + pytest.skip(f"Requires {cp_size} GPUs for context parallelism") + monkeypatch.setenv("NVTE_FP8_DPA_BWD", "1" if fp8_backward else "0") + mesh = jax.sharding.Mesh(np.asarray(jax.local_devices()[:cp_size]), ("cp",)) + sharding = jax.sharding.NamedSharding(mesh, jax.sharding.PartitionSpec(None, "cp", None, None)) + length = 256 * cp_size + q = jnp.full((1, length, 4, 64), 0.5, jnp.bfloat16) + k = q[:, :, :2] + v = jnp.zeros_like(k).at[:, ::2, :, 0].set(1.0) + # Distinct rank-local values make missing communication or incorrect + # dual-chunk ordering observable, while retaining exact MXFP8 operands. + amplitudes = jnp.exp2((jnp.arange(length) // 128) % 4 - 2).astype(v.dtype) + v *= amplitudes[None, :, None, None] + inputs = (q, k, v) + balanced = kind == "balanced_causal" + causal = kind != "dense" + mask = ( + jnp.arange(length)[None, :, None] >= jnp.arange(length)[None, None, :] if causal else None + ) + reference = lambda xs: _reference(xs, mask) + expected = jax.jit(reference)(inputs) + ref_grads = jax.jit(jax.grad(lambda xs: reference(xs).sum()))(inputs) + reorder = lambda x, reverse: reorder_causal_dual_chunk_swap(x, cp_size, 1, reverse) + if balanced: + inputs = tuple(reorder(x, False) for x in inputs) + attention = lambda xs: _attention( + xs, + attn_mask_type=AttnMaskType.CAUSAL_MASK if causal else AttnMaskType.NO_MASK, + context_parallel_axis="cp", + context_parallel_strategy=strategy, + context_parallel_causal_load_balanced=balanced, + ) + with jax.set_mesh(mesh), global_shard_guard(MeshResource(cp_resource="cp")): + actual = jax.jit(attention, in_shardings=((sharding,) * 3,))(inputs) + grads = jax.jit(jax.grad(lambda xs: attention(xs).sum()), in_shardings=((sharding,) * 3,))( + inputs + ) + if balanced: + actual = reorder(actual, True) + grads = tuple(reorder(x, True) for x in grads) + np.testing.assert_allclose(actual, expected, rtol=1e-2, atol=1e-3) + for grad, ref in zip(grads, ref_grads): + np.testing.assert_allclose(grad, ref, rtol=3e-2, atol=2e-2) + + +@pytest.mark.parametrize("strategy", (CPStrategy.ALL_GATHER, CPStrategy.RING)) +@pytest.mark.parametrize( + "layout", (QKVLayout.BS3HD, QKVLayout.BSHD_BS2HD, QKVLayout.BSHD_BSHD_BSHD) +) +@pytest.mark.parametrize("prequantized", (False, True)) +def test_context_parallel_storage(strategy, layout, prequantized): + """Packed and supplied FP8 inputs preserve nonuniform cross-rank values and gradients.""" + _require_mxfp8() + if jax.local_device_count() < 2: + pytest.skip("Requires two GPUs for context parallelism") + mesh = jax.sharding.Mesh(np.asarray(jax.local_devices()[:2]), ("cp",)) + sharding = jax.sharding.NamedSharding(mesh, jax.sharding.PartitionSpec(None, "cp", None, None)) + rng = np.random.default_rng(42) + q = jnp.asarray(rng.choice([-0.5, 0.5], (1, 512, 2, 64)), jnp.bfloat16) + k = jnp.asarray(rng.choice([-0.5, 0.5], q.shape), q.dtype) + v = jnp.asarray(rng.choice([-1.0, 1.0], q.shape), q.dtype) + amplitude = jnp.exp2(jnp.arange(512) // 128 - 2).astype(q.dtype) + k *= amplitude[None, :, None, None] + v *= amplitude[None, :, None, None] + weight = jnp.asarray(rng.choice([-1.0, 1.0], q.shape), q.dtype) + inputs = _pack((q, k, v), layout) + reference = lambda xs: _reference(_unpack(xs, layout)) + expected = jax.jit(reference)(inputs) + ref_grads = jax.jit(jax.grad(lambda xs: (reference(xs) * weight).sum()))(inputs) + + def attention(xs): + supplied = quantize_attention_qkv(_unpack(xs, layout)) if prequantized else None + return _attention( + xs, + qkv_layout=layout, + quantized_qkv=supplied, + context_parallel_axis="cp", + context_parallel_strategy=strategy, + ) + + with jax.set_mesh(mesh), global_shard_guard(MeshResource(cp_resource="cp")): + actual = jax.jit(attention, in_shardings=((sharding,) * len(inputs),))(inputs) + grads = jax.jit( + jax.grad(lambda xs: (attention(xs) * weight).sum()), + in_shardings=((sharding,) * len(inputs),), + )(inputs) + relative_error = lambda a, b: float(jnp.linalg.norm(a - b) / jnp.linalg.norm(b)) + assert relative_error(actual, expected) < 0.04 + for grad, ref in zip(grads, ref_grads): + assert relative_error(grad, ref) < 0.08 + + +def test_prequantized_vmap_and_checkpoint(): + """Independent mapped attention calls retain scale boundaries and model gradients.""" + _require_mxfp8() + q = jnp.full((2, 1, 128, 2, 64), 0.5, jnp.bfloat16) + v = jnp.zeros_like(q).at[:, :, ::2, :, 0].set(1.0) + inputs = q, q, v + forward = lambda xs: _attention(xs, quantized_qkv=quantize_attention_qkv(xs)) + plain = jax.jit(jax.vmap(jax.value_and_grad(lambda xs: forward(xs).sum())))(inputs) + remat = jax.jit(jax.vmap(jax.value_and_grad(lambda xs: jax.checkpoint(forward)(xs).sum())))( + inputs + ) + reference = jax.jit(jax.vmap(jax.value_and_grad(lambda xs: _reference(xs).sum())))(inputs) + for value, recomputed, ref in zip( + jax.tree.leaves(plain), jax.tree.leaves(remat), jax.tree.leaves(reference) + ): + np.testing.assert_array_equal(value, recomputed) + np.testing.assert_allclose(value, ref, rtol=1e-2, atol=1e-3) + + +@pytest.mark.parametrize("transpose", (False, True)) +def test_flax_prequantized_inputs(transpose): + """BSHD and SBHD Flax calls consume the same prequantized BHSD operands.""" + from transformer_engine.jax.flax import DotProductAttention + + _require_mxfp8() + q = jnp.full((1, 128, 2, 64), 0.5, jnp.bfloat16) + v = jnp.zeros_like(q).at[:, ::2, :, 0].set(1.0) + prepared = jax.jit(quantize_attention_qkv)((q, q, v)) + if transpose: + q, v = q.swapaxes(0, 1), v.swapaxes(0, 1) + module = DotProductAttention( + head_dim=64, + num_attention_heads=2, + attn_mask_type="no_mask", + transpose_batch_sequence=transpose, + fp8="mxfp8", + ) + variables = module.init( + jax.random.key(0), q, q, v, SequenceDescriptor(), quantized_qkv=prepared + ) + forward = lambda xs: module.apply(variables, *xs, SequenceDescriptor(), quantized_qkv=prepared) + actual = jax.jit(forward)((q, q, v)) + expected = jnp.zeros_like(actual).at[..., 0].set(0.5) + np.testing.assert_allclose(actual, expected, atol=5e-4) + grads = jax.jit(jax.grad(lambda xs: forward(xs).sum()))((q, q, v)) + np.testing.assert_allclose(grads[2], 1.0, atol=1e-3) diff --git a/transformer_engine/common/cudnn/attention.py b/transformer_engine/common/cudnn/attention.py index e14f8e377cf..311b4156311 100644 --- a/transformer_engine/common/cudnn/attention.py +++ b/transformer_engine/common/cudnn/attention.py @@ -55,6 +55,20 @@ class BuiltAttentionGraph: 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) + if ( + spec.options.mxfp8 + and spec.options.backward + and specs["q"].dim[-1] > 128 + and str(getattr(cudnn, "__version__", "")).split("+", 1)[0].split("-", 1)[0] == "1.29.0" + ): + # Frontend 1.29.0 sdpa_fp8_bwd.h overwrites Q_T/K_T scale strides with + # a hard-coded width of 128, even for wider Q/K heads. Do not permit + # silently incorrect gradients through either framework adapter. + raise ValueError( + "cuDNN Frontend 1.29.0 MXFP8 backward has incorrect scale strides for " + "Q/K head dimensions greater than 128. Use model-precision backward " + "with NVTE_FP8_DPA_BWD=0." + ) tensors = {} uids = {name: index + 1 for index, (name, _) in enumerate(spec.tensors)} diff --git a/transformer_engine/jax/attention.py b/transformer_engine/jax/attention.py index 932c2cdacda..b12942ce0d7 100644 --- a/transformer_engine/jax/attention.py +++ b/transformer_engine/jax/attention.py @@ -1734,6 +1734,19 @@ def fused_attn( # quantized tensor constructors validate real arrays, whereas JAX may # reconstruct cotangent pytrees with opaque sentinel leaves. quantized_qkv = prequantized_buffers(qkv, quantized_qkv, qkv_layout) + if fp8 == "mxfp8": + if qkv_layout.is_thd(): + raise ValueError("MXFP8 attention does not support THD format.") + from .sharding import is_mesh_available, global_mesh_resource, get_mesh_axis_size + + if is_mesh_available(): + cp_axis = context_parallel_axis or global_mesh_resource().cp_resource + if get_mesh_axis_size(cp_axis) > 1: + from .cpp_extensions.mxfp8_attention import context_parallel_qkv + + qkv = context_parallel_qkv(qkv, quantized_qkv, qkv_layout) + qkv_layout = QKVLayout.BSHD_BSHD_BSHD + quantized_qkv = None if score_mod is None: score_mod_only_args = [ name diff --git a/transformer_engine/jax/cpp_extensions/mxfp8_attention.py b/transformer_engine/jax/cpp_extensions/mxfp8_attention.py index fd772954d40..b948ef32b86 100644 --- a/transformer_engine/jax/cpp_extensions/mxfp8_attention.py +++ b/transformer_engine/jax/cpp_extensions/mxfp8_attention.py @@ -4,15 +4,17 @@ """Adapt JAX MXFP8 quantization buffers to the shared attention graph contract.""" -from functools import partial +from functools import lru_cache, partial import jax import jax.numpy as jnp from jax.experimental.custom_partitioning import custom_partitioning from jax.sharding import NamedSharding, PartitionSpec +from jax.custom_batching import sequential_vmap from transformer_engine.common.cudnn.mxfp8 import quantization_layouts, scale_reordering from ..quantize import QuantizerFactory, QuantizeLayout, ScalingMode, ScaledTensor2x +from ..quantize.tensor import ScaledTensorFactory from .quantization import quantize from .base import _JAX_GSPMD_SUPPORTED @@ -77,19 +79,44 @@ def prepare_doutput(value): def quantize_qkv(qkv): """Prepare separate BSHD operands with independent BHSD row/column quantizations.""" - return tuple( - # Quantized buffers are auxiliary inputs, not differentiation targets. - ScaledTensor2x( - *quantize_operand(jax.lax.stop_gradient(x).transpose(0, 2, 1, 3), True, True) + buffers = _quantize_qkv(tuple(jax.lax.stop_gradient(x) for x in qkv)) + results = [] + for value, parts in zip(qkv, buffers): + tensors = [] + for i, columnwise in enumerate((False, True)): + data, scale, amax = parts[3 * i : 3 * i + 3] + tensors.append( + ScaledTensorFactory.create_1x( + data, + scale, + amax, + scaling_mode=ScalingMode.MXFP8_1D_SCALING, + dq_dtype=value.dtype, + is_colwise=columnwise, + ) + ) + results.append(ScaledTensor2x(tensors[0], tensors[1])) + return tuple(results) + + +@sequential_vmap +def _quantize_qkv(qkv): + results = [] + for value in qkv: + pair = quantize_operand(value.transpose(0, 2, 1, 3), True, True) + results.append( + tuple(x for tensor in pair for x in (tensor.data, tensor.scale_inv, tensor.amax)) ) - for x in qkv - ) + return tuple(results) def prequantized_buffers(qkv, quantized_qkv, layout): """Validate and extract unswizzled BHSD MXFP8 operands without dequantizing them.""" if layout.is_thd(): raise ValueError("MXFP8 attention does not support THD format.") + expected = 1 if layout.is_qkvpacked() else 2 if layout.is_kvpacked() else 3 + if len(qkv) != expected: + raise ValueError(f"Expected {expected} model operands for {layout}, received {len(qkv)}.") if layout.is_qkvpacked(): operands = tuple(qkv[0][..., i, :, :] for i in range(3)) elif layout.is_kvpacked(): @@ -124,6 +151,29 @@ def prequantized_buffers(qkv, quantized_qkv, layout): return tuple(jax.lax.stop_gradient(x) for x in buffers) +def context_parallel_qkv(qkv, buffers, layout): + """Unpack model storage and decode supplied inputs before CP communication. + + MXFP8 CP communicates model-precision values and quantizes the local Q/K/V + slices independently for each contraction. As in PyTorch, prequantized + inputs are decoded before communication. Gradients target the original + model arrays, treating the quantization as a straight-through operation. + """ + if layout.is_qkvpacked(): + qkv = tuple(qkv[0][..., i, :, :] for i in range(3)) + elif layout.is_kvpacked(): + qkv = qkv[0], qkv[1][..., 0, :, :], qkv[1][..., 1, :, :] + if buffers is None: + return qkv + decoded = [] + for i, source in enumerate(qkv): + data, _, scale, _ = buffers[4 * i : 4 * i + 4] + value = data.astype(jnp.float32) * jnp.repeat(scale.astype(jnp.float32), 32, axis=-1) + value = value.transpose(0, 2, 1, 3).astype(source.dtype) + decoded.append(jax.lax.stop_gradient(value) + (source - jax.lax.stop_gradient(source))) + return tuple(decoded) + + def _prepared_values(buffers, backward): values = {} for i, name in enumerate(("q", "k", "v")): @@ -178,15 +228,31 @@ def _execute_prequantized(args, buffers, config, backward): return output, stats, rng, max_logit -@partial(custom_partitioning, static_argnums=(2, 3)) def execute_prequantized(args, buffers, config, backward): """Execute supplied quantizations with the standard attention gradient contract.""" + return _batched_executor(config, backward)(args, buffers) + + +@lru_cache(maxsize=256) +def _batched_executor(config, backward): + # Independent vmap elements retain their per-head block boundaries. Use a + # loop rather than the generic quantizer's flattening of mapped dimensions. + @sequential_vmap + def execute(args, buffers): + return _partitioned_prequantized(args, buffers, config, backward) + + return execute + + +@partial(custom_partitioning, static_argnums=(2, 3)) +def _partitioned_prequantized(args, buffers, config, backward): return _execute_prequantized(args, buffers, config, backward) def _prequantized_partition(config, backward, mesh, arg_infos, result_infos): from .attention import FusedAttnFwdPrimitive, FusedAttnBwdPrimitive from ..sharding import get_mesh_axis_size + from .misc import get_padded_spec if get_mesh_axis_size(config.cp_axis, mesh) > 1: raise ValueError("Prequantized MXFP8 context parallelism requires local requantization.") @@ -199,7 +265,7 @@ def _prequantized_partition(config, backward, mesh, arg_infos, result_infos): if config.qkv_layout.is_qkvpacked() else min(index, 1) if config.qkv_layout.is_kvpacked() else index ) - spec = arg_shardings[source].spec + spec = get_padded_spec(arg_infos[0][source]) # Model inputs are BSHD (possibly packed); supplied buffers are BHSD. prepared = NamedSharding(mesh, PartitionSpec(spec[0], spec[-2], spec[1], spec[-1])) data_shardings.extend((prepared,) * 4) @@ -233,7 +299,7 @@ def _prequantized_infer(config, backward, mesh, arg_infos, result_infos): return base.infer_sharding_from_operands(config, mesh, arg_infos[0], result_infos) -execute_prequantized.def_partition( +_partitioned_prequantized.def_partition( partition=_prequantized_partition, sharding_rule=_prequantized_rule, **({"infer_sharding_from_operands": _prequantized_infer} if _JAX_GSPMD_SUPPORTED else {}), diff --git a/transformer_engine/jax/flax/transformer.py b/transformer_engine/jax/flax/transformer.py index 13460b535d6..9c7faf7a997 100644 --- a/transformer_engine/jax/flax/transformer.py +++ b/transformer_engine/jax/flax/transformer.py @@ -327,6 +327,7 @@ def __call__( deterministic: bool = False, score_mod_tensors: Optional[Mapping[str, Any]] = None, score_mod_bprop_tensors: Optional[Mapping[str, Any]] = None, + quantized_qkv=None, ) -> Array: seed = None @@ -372,6 +373,8 @@ def __call__( } if self.fp8: fused_attn_kwargs["fp8"] = self.fp8 + if quantized_qkv is not None: + fused_attn_kwargs["quantized_qkv"] = quantized_qkv if self.qkv_layout.is_qkvpacked(): """qkvpacked format, treat @@ -713,6 +716,7 @@ def __call__( mask: Optional[Union[SequenceDescriptor, Array]] = None, score_mod_tensors: Optional[Mapping[str, Any]] = None, score_mod_bprop_tensors: Optional[Mapping[str, Any]] = None, + quantized_qkv=None, ) -> Array: """ Parameters @@ -737,6 +741,10 @@ def __call__( Additional tensor/scalar operands passed to ``score_mod``. score_mod_bprop_tensors: Optional[Mapping[str, Any]], default = None Additional tensor/scalar operands passed to ``score_mod_bprop``. + quantized_qkv: Optional[Tuple[ScaledTensor2x, ...]], default = None + Separate prequantized MXFP8 Q/K/V in BHSD order, independent of the + model input layout. Requires ``fp8="mxfp8"`` and a SequenceDescriptor. + See :func:`transformer_engine.jax.attention.quantize_attention_qkv`. Returns ------- @@ -744,6 +752,8 @@ def __call__( Output tensor, or ``(output, max_logit)`` when ``return_max_logit`` is enabled. """ input_dtype = query.dtype + if quantized_qkv is not None and self.fp8 != "mxfp8": + raise ValueError('quantized_qkv requires fp8="mxfp8".') if mask is not None: if sequence_descriptor is not None: @@ -980,6 +990,7 @@ def __call__( deterministic=deterministic, score_mod_tensors=score_mod_tensors, score_mod_bprop_tensors=score_mod_bprop_tensors, + quantized_qkv=quantized_qkv, ) output = x[0] if self.return_max_logit else x assert ( From 88fd9d209efac4b25f2f6e4b38a1728011877822 Mon Sep 17 00:00:00 2001 From: Vladimir Cherepanov Date: Wed, 21 Oct 2026 18:51:33 +0000 Subject: [PATCH 15/49] Skip unsupported JAX cuDNN attention test configurations Query fused attention backend support before running FP8 tests instead of assuming Hopper supports FP16/BF16 outputs. Preserve cuDNN rejection reasons and match context-parallel queries to their per-device shapes and layouts. Signed-off-by: Vladimir Cherepanov --- tests/jax/test_fused_attn.py | 104 ++++++++++++++++++++++++----------- 1 file changed, 71 insertions(+), 33 deletions(-) diff --git a/tests/jax/test_fused_attn.py b/tests/jax/test_fused_attn.py index 568c526b974..8e14fe4805c 100644 --- a/tests/jax/test_fused_attn.py +++ b/tests/jax/test_fused_attn.py @@ -2449,6 +2449,33 @@ def test_backward( ) +def _require_cudnn_attention(q, k, v, **overrides): + """Skip unsupported graphs with cuDNN's reason, including FP8 output restrictions.""" + helper = FusedAttnHelper( + is_training=True, + batch_size=q.shape[0], + q_dtype=q.dtype, + kv_dtype=k.dtype, + 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=q.shape[-2], + kv_num_heads=k.shape[-2], + q_max_seqlen=q.shape[1], + kv_max_seqlen=k.shape[1], + head_dim_qk=q.shape[-1], + head_dim_v=v.shape[-1], + window_size=(-1, -1), + bottom_right_diagonal=False, + attn_scale=0.125, + ) + backend, reason = replace(helper, **overrides).get_fused_attn_backend() + if backend == NVTE_Fused_Attn_Backend.NVTE_No_Backend: + pytest.skip(reason) + + @pytest.mark.parametrize("dtype", [jnp.float16, jnp.bfloat16]) @pytest.mark.parametrize("fp8_backward", [False, True]) @pytest.mark.parametrize( @@ -2500,8 +2527,7 @@ def attention(operands): fp8=True, ) - if get_device_compute_capability(0) < 90: - pytest.skip("cuDNN FP8 SDPA requires Hopper or newer.") + _require_cudnn_attention(q, k, v, qkv_layout=layout, fp8=True) actual = jax.jit(attention)(operands) expected = jax.jit(reference)(operands) np.testing.assert_allclose(actual, expected, rtol=5e-3, atol=5e-3) @@ -2585,8 +2611,8 @@ 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 ragged and get_device_compute_capability(0) < 100: + pytest.skip("Ragged backward and emulated SM12x statistics require Blackwell.") if sm_arch is not None: from transformer_engine.jax.cpp_extensions import cudnn_attention @@ -2627,6 +2653,18 @@ def test_cudnn_python_fp8_gqa_masks(kind, sm_arch, monkeypatch, request): attn_mask = AttnMaskType.PADDING_CAUSAL_MASK layout = QKVLayout.THD_THD_THD + _require_cudnn_attention( + q, + k, + v, + qkv_layout=layout, + attn_mask_type=attn_mask, + max_segments_per_seq=2 if ragged else 1, + window_size=window if window is not None else (-1, -1), + fp8=fp8, + return_max_logit=True, + ) + 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)) @@ -2755,8 +2793,6 @@ def test_cudnn_python_flax_forward_backward(softmax, deterministic, fp8, monkeyp """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, @@ -2768,6 +2804,18 @@ def test_cudnn_python_flax_forward_backward(softmax, deterministic, fp8, monkeyp ) q = jnp.zeros((1, 128, 2, 64), jnp.bfloat16) v = jnp.zeros_like(q).at[:, ::2, :, 0].set(1.0) + if fp8: + _require_cudnn_attention( + q, + q, + v, + fp8=True, + softmax_type=( + AttnSoftmaxType.LEARNABLE_SOFTMAX + if softmax == "learnable" + else AttnSoftmaxType.VANILLA_SOFTMAX + ), + ) variables = module.init(jax.random.PRNGKey(0), q, q, v, deterministic=deterministic) def forward(variables, value): @@ -2795,33 +2843,9 @@ def forward(variables, value): @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)) + _require_cudnn_attention(*arrays, dropout_probability=0.2, fp8=fp8) def attention(arrays, seed): return fused_attn( @@ -2856,12 +2880,26 @@ 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) + if fp8: + # Query the per-device graph: all-gather retains the global KV sequence, + # while ring attention uses local Q and KV blocks. + _require_cudnn_attention( + q, + q, + v, + fp8=True, + qkv_layout=( + QKVLayout.BSHD_BSHD_BSHD + if strategy == CPStrategy.ALL_GATHER + else QKVLayout.BSHD_BS2HD + ), + q_max_seqlen=q.shape[1] // 2, + kv_max_seqlen=q.shape[1] if strategy == CPStrategy.ALL_GATHER else q.shape[1] // 2, + ) arrays = tuple(jax.device_put(value, sharding) for value in (q, q, v)) def attention(arrays): From 6bb64820b0c392224fb46bcbf3c7dbfa8985b6ae Mon Sep 17 00:00:00 2001 From: Vladimir Cherepanov Date: Tue, 27 Oct 2026 19:15:38 +0000 Subject: [PATCH 16/49] Require precompiled cuDNN plans during JAX graph deserialization Signed-off-by: Vladimir Cherepanov --- docs/installation.rst | 7 ++++--- transformer_engine/jax/csrc/extensions/cudnn_graph.cpp | 3 ++- 2 files changed, 6 insertions(+), 4 deletions(-) diff --git a/docs/installation.rst b/docs/installation.rst index 9fac9ee56a6..36f52bce1bb 100644 --- a/docs/installation.rst +++ b/docs/installation.rst @@ -25,9 +25,10 @@ and target device or architecture. Support queries select supported plans withou 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. +native FFI executor. The executor requires a serialized precompiled execution +plan and reports an error if it cannot be restored. 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. diff --git a/transformer_engine/jax/csrc/extensions/cudnn_graph.cpp b/transformer_engine/jax/csrc/extensions/cudnn_graph.cpp index 2499899ca3f..157a7d17cfe 100644 --- a/transformer_engine/jax/csrc/extensions/cudnn_graph.cpp +++ b/transformer_engine/jax/csrc/extensions/cudnn_graph.cpp @@ -81,7 +81,8 @@ std::shared_ptr GetGraph(cudnnHandle_t handle, ffi::Dictionary &attrs) { 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())); + auto status = candidate->deserialize(handle, std::vector(data.begin(), data.end()), + /*enforce_precompiled=*/true); if (!status.is_good()) throw std::runtime_error(status.get_message()); graph = std::move(candidate); } From 710e419bf832f8f6eb9479d4f88c05843db42235 Mon Sep 17 00:00:00 2001 From: Vladimir Cherepanov Date: Tue, 27 Oct 2026 19:36:24 +0000 Subject: [PATCH 17/49] Temporarily disable FROST plans in JAX cuDNN attention Signed-off-by: Vladimir Cherepanov --- .../jax/cpp_extensions/cudnn_attention.py | 9 ++++++++- .../jax/cpp_extensions/cudnn_graph.py | 13 +++++++++++++ 2 files changed, 21 insertions(+), 1 deletion(-) diff --git a/transformer_engine/jax/cpp_extensions/cudnn_attention.py b/transformer_engine/jax/cpp_extensions/cudnn_attention.py index af11ec9066a..4e38f29440b 100644 --- a/transformer_engine/jax/cpp_extensions/cudnn_attention.py +++ b/transformer_engine/jax/cpp_extensions/cudnn_attention.py @@ -33,7 +33,13 @@ ) from transformer_engine.common.cudnn import mxfp8 as mxfp8_rules -from .cudnn_graph import import_cudnn, check_frontend_version, serialize_graph, graph_attributes +from .cudnn_graph import ( + import_cudnn, + check_frontend_version, + deselect_frost_plans, + serialize_graph, + graph_attributes, +) from . import mxfp8_attention @@ -191,6 +197,7 @@ def _build_graph(direction, shapes, config, capabilities, nondeterministic, fp8_ backend, direction = _site(description.spec) with _plan_lock: if entry.workspace_size is None: + deselect_frost_plans(entry.built.graph) workspace_size = compile_graph( cudnn, entry.built.graph, diff --git a/transformer_engine/jax/cpp_extensions/cudnn_graph.py b/transformer_engine/jax/cpp_extensions/cudnn_graph.py index b5cc35b7579..ddd891f2910 100644 --- a/transformer_engine/jax/cpp_extensions/cudnn_graph.py +++ b/transformer_engine/jax/cpp_extensions/cudnn_graph.py @@ -107,9 +107,22 @@ def serialize_graph( ) +def deselect_frost_plans(graph): + """Temporarily exclude FROST SDPA plans from the JAX C++ serialization path.""" + # cuDNN cannot currently serialize FROST SDPA plans for C++ execution. + # Remove this workaround once cuDNN fixes its serialization path. + plan_names = [ + graph.get_plan_name_at_index(index) for index in range(graph.get_execution_plan_count()) + ] + graph.deselect_engines( + [name for name in plan_names if name.startswith(("sdpa_fwd_", "sdpa_bwd_"))] + ) + + def finalize_graph(cudnn, graph): """Validate, select and compile a supported execution plan.""" check_graph_support(cudnn, graph) + deselect_frost_plans(graph) workspace_size = compile_graph(cudnn, graph) return workspace_size, bytes(graph.serialize()), check_frontend_version(cudnn) From 070c008e44771e49ac359d6b56411c19434c5655 Mon Sep 17 00:00:00 2001 From: Vladimir Cherepanov Date: Thu, 29 Oct 2026 23:37:37 +0000 Subject: [PATCH 18/49] Fix PyTorch fused attention precision and inference dropout Validate model-precision backward graphs when NVTE_FP8_DPA_BWD=0, retaining the original QKV layout and nominal dtype. Normalize inference dropout before preparing runtime graph inputs. Add regression coverage for FP8 backward selection and inference output against zero-dropout execution and FP32 SDPA. Signed-off-by: Vladimir Cherepanov --- tests/pytorch/attention/test_attention.py | 132 ++++++++++++++++++ .../attention/dot_product_attention/utils.py | 6 + .../cpp_extensions/cudnn_fused_attn.py | 35 ++++- .../pytorch/cpp_extensions/fused_attn.py | 2 +- 4 files changed, 173 insertions(+), 2 deletions(-) diff --git a/tests/pytorch/attention/test_attention.py b/tests/pytorch/attention/test_attention.py index 85e53fb59b9..6eac0d7de8b 100644 --- a/tests/pytorch/attention/test_attention.py +++ b/tests/pytorch/attention/test_attention.py @@ -220,6 +220,138 @@ def test_fused_attn_backend_message(): assert message != "" +@pytest.mark.parametrize("dtype", (DType.kFloat16, DType.kBFloat16)) +@pytest.mark.parametrize("fp8_recipe", ("delayed", "current", "mxfp8")) +@pytest.mark.parametrize("fp8_bwd", (None, "0", "1")) +@pytest.mark.parametrize("layout", ("bshd_bshd_bshd", "bs3hd")) +@pytest.mark.parametrize("head_dim", (128, 192)) +def test_fused_attn_query_backward_precision( + monkeypatch, dtype, fp8_recipe, fp8_bwd, layout, head_dim +): + """Validate the executed backward, including the Frontend 1.29 MXFP8 workaround.""" + import cudnn + from transformer_engine.pytorch.cpp_extensions import cudnn_fused_attn as implementation + + monkeypatch.delenv("NVTE_FP8_DPA_BWD", raising=False) + if fp8_bwd is not None: + monkeypatch.setenv("NVTE_FP8_DPA_BWD", fp8_bwd) + monkeypatch.setattr(cudnn, "__version__", "1.29.0") + monkeypatch.setattr(torch.cuda, "get_device_capability", lambda *_args: (10, 0)) + is_mxfp8 = fp8_recipe == "mxfp8" + delayed = fp8_recipe == "delayed" + effective_layout = "bhsd_bhsd_bhsd" if is_mxfp8 and layout == "bs3hd" else layout + config = FusedAttentionParams( + qkv_layout=QKVLayout[effective_layout], + dqkv_layout=QKVLayout[layout], + o_format=QKVFormat["bshd"], + do_format=QKVFormat["bshd"], + qkv_dtype=DType.kFloat8E4M3, + o_dtype=DType.kFloat8E4M3 if delayed else dtype, + do_dtype=DType.kFloat8E5M2, + dqkv_dtype=DType.kFloat8E5M2 if delayed else dtype, + nominal_dtype=dtype if delayed else None, + scaling_mode=( + tex.NVTEScalingMode.NVTE_MXFP8_1D_SCALING + if is_mxfp8 + else tex.NVTEScalingMode.NVTE_DELAYED_TENSOR_SCALING + ), + batch_size=2, + num_attn_heads=4, + num_gqa_groups=4, + head_dim_qk=head_dim, + head_dim_v=128, + max_seqlen_q=128, + max_seqlen_kv=128, + ) + queries = [] + + def entry(attn): + queries.append(attn) + graph = Mock() + for name, count in ( + ("sdpa_mxfp8", 3), + ("sdpa_fp8", 4), + ("sdpa_mxfp8_backward", 6), + ("sdpa_fp8_backward", 7), + ("sdpa_backward", 3), + ): + getattr(graph, name).return_value = tuple(Mock() for _ in range(count)) + implementation.build_attention_graph(cudnn, graph, attn.description()) + + monkeypatch.setattr(implementation._AttentionGraph, "entry", entry) + backend, message = implementation.get_fused_attn_backend(config) + assert len(queries) == 2 + assert queries[0].fp8 + if is_mxfp8 and head_dim > 128 and fp8_bwd != "0": + assert backend == FusedAttnBackend.No_Backend + assert "incorrect scale strides" in message + return + assert backend == FusedAttnBackend.FP8 + assert message == "" + backward = queries[1] + if fp8_bwd == "0": + expected = implementation._query_graph( + replace( + config, + qkv_dtype=dtype, + o_dtype=dtype, + do_dtype=dtype, + dqkv_dtype=dtype, + qkv_layout=QKVLayout[layout], + ), + True, + ) + assert backward.description() == expected.description() + assert not backward.fp8 and not backward.mxfp8 + assert not any(name.startswith("descale_") for name in backward.specs) + + # Falling back to model precision must still validate backward support. + def reject_backward(attn): + if attn.backward: + raise ValueError("unsupported model-precision backward") + + monkeypatch.setattr(implementation._AttentionGraph, "entry", reject_backward) + backend, message = implementation.get_fused_attn_backend(config) + assert backend == FusedAttnBackend.No_Backend + assert message == "unsupported model-precision backward" + else: + assert backward.fp8 + assert backward.mxfp8 == is_mxfp8 + + +@pytest.mark.parametrize("dtype", (torch.float16, torch.bfloat16)) +@pytest.mark.parametrize("qkv_format", ("bshd", "sbhd")) +def test_fused_attn_inference_ignores_dropout(dtype, qkv_format): + """Configured inference dropout must match zero-dropout execution and FP32 SDPA.""" + shape = (2, 128, 4, 64) if qkv_format == "bshd" else (128, 2, 4, 64) + q, k, v = [torch.randn(shape, device="cuda", dtype=dtype) for _ in range(3)] + cu = torch.arange(3, device="cuda", dtype=torch.int32) * 128 + kwargs = dict( + is_training=False, + max_seqlen_q=128, + max_seqlen_kv=128, + cu_seqlens_q=cu, + cu_seqlens_kv=cu, + q=q, + k=k, + v=v, + fake_dtype=dtype, + fused_attention_backend=FusedAttnBackend.F16_arbitrary_seqlen, + qkv_layout="_".join([qkv_format] * 3), + o_format=qkv_format, + attn_mask_type="causal", + ) + actual, _ = fused_attn_fwd(dropout=0.25, **kwargs) + expected, _ = fused_attn_fwd(dropout=0.0, **kwargs) + torch.testing.assert_close(actual, expected, rtol=0, atol=0) + order = (0, 2, 1, 3) if qkv_format == "bshd" else (1, 2, 0, 3) + reference = torch.nn.functional.scaled_dot_product_attention( + *(tensor.float().permute(order) for tensor in (q, k, v)), is_causal=True + ) + rtol, atol = (2e-2, 1e-2) if dtype == torch.bfloat16 else (3e-3, 2e-3) + torch.testing.assert_close(actual.float().permute(order), reference, rtol=rtol, atol=atol) + + @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): diff --git a/transformer_engine/pytorch/attention/dot_product_attention/utils.py b/transformer_engine/pytorch/attention/dot_product_attention/utils.py index ca582956cf5..25a4fceda7e 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/utils.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/utils.py @@ -395,6 +395,8 @@ class FusedAttentionParams: o_dtype: DType = DType.kBFloat16 do_dtype: DType = DType.kBFloat16 dqkv_dtype: DType = DType.kBFloat16 + # Model precision for backward when NVTE_FP8_DPA_BWD=0. + nominal_dtype: Optional[DType] = None # tensor layouts qkv_layout: tex.NVTE_QKV_Layout = tex.NVTE_QKV_Layout.NVTE_QKV_Layout_NOT_SET @@ -1640,6 +1642,9 @@ def _is_fa3_supported(num_heads, num_gqa_groups, head_dim_qk, head_dim_v, qkv_dt bias_batch_size, bias_num_heads, bias_seqlen_q, bias_seqlen_kv = ( fu_core_attention_bias_shape ) + model_dtype = qkv_dtype + if model_dtype in (torch.float8_e4m3fn, torch.float8_e5m2): + model_dtype = nominal_dtype if nominal_dtype is not None else torch.bfloat16 base_fused_attn_kwargs = { "is_training": is_training, "deterministic": deterministic, @@ -1658,6 +1663,7 @@ def _is_fa3_supported(num_heads, num_gqa_groups, head_dim_qk, head_dim_v, qkv_dt "o_dtype": o_type, "do_dtype": do_type, "dqkv_dtype": dqkv_type, + "nominal_dtype": TE_DType[model_dtype], "qkv_layout": QKVLayout[spec.qkv_layout], "o_format": QKVFormat[o_format], "do_format": QKVFormat[do_format], diff --git a/transformer_engine/pytorch/cpp_extensions/cudnn_fused_attn.py b/transformer_engine/pytorch/cpp_extensions/cudnn_fused_attn.py index fcdf2363571..279571e1fe3 100644 --- a/transformer_engine/pytorch/cpp_extensions/cudnn_fused_attn.py +++ b/transformer_engine/pytorch/cpp_extensions/cudnn_fused_attn.py @@ -8,7 +8,8 @@ graph construction, plan caching, and execution do not use the common attention API. """ -from dataclasses import dataclass +from dataclasses import dataclass, replace +import os import threading import torch @@ -391,6 +392,36 @@ 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 + if ( + backward + and DType.cast(config.qkv_dtype) in (DType.kFloat8E4M3, DType.kFloat8E5M2) + and not int(os.getenv("NVTE_FP8_DPA_BWD", "1")) + ): + # FusedAttnFunc saves model-precision Q/K/V and O for this backward. + # The gradient layout records the inputs' layout before FP8 quantization. + dtype = config.nominal_dtype + if dtype is None: + dtype = next( + ( + t + for t in (config.dqkv_dtype, config.o_dtype) + if t in (DType.kFloat16, DType.kBFloat16) + ), + DType.kBFloat16, + ) + config = replace( + config, + qkv_dtype=dtype, + o_dtype=dtype, + do_dtype=dtype, + dqkv_dtype=dtype, + qkv_layout=( + config.dqkv_layout + if config.dqkv_layout != tex.NVTE_QKV_Layout.NVTE_QKV_Layout_NOT_SET + else config.qkv_layout + ), + ) + def name(mapping, value): return next(k for k, v in mapping.items() if v == value) @@ -676,6 +707,8 @@ def get_fused_attn_backend(config): def fused_attn_fwd(params): """Execute forward; preserve the extension's output and auxiliary tensor order.""" p = _normalize(params) + if not p["is_training"]: + p["dropout"] = 0.0 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) diff --git a/transformer_engine/pytorch/cpp_extensions/fused_attn.py b/transformer_engine/pytorch/cpp_extensions/fused_attn.py index 98628205782..533b8a830a8 100644 --- a/transformer_engine/pytorch/cpp_extensions/fused_attn.py +++ b/transformer_engine/pytorch/cpp_extensions/fused_attn.py @@ -255,7 +255,7 @@ def fused_attn_fwd( if None, use 1.0/sqrt(head_dim_qk) as the default dropout : float, default = 0.0 dropout probability, 0.0 means no dropout, 1.0 means no output; - dropout must be 0.0 if is_training is False + ignored if is_training is False fast_zero_fill : bool, default = True if True, initializes the output tensor O to zero using the fast filling method; if False, uses PyTorch's .fill_() method From 5ab71e0d628523ea004e11b24ebddeb8ddbf6a90 Mon Sep 17 00:00:00 2001 From: Vladimir Cherepanov Date: Thu, 29 Oct 2026 23:54:53 +0000 Subject: [PATCH 19/49] Fix SBHD paged-cache dimensions in PyTorch attention Map paged K/V dimensions to BHSD consistently with their strides. Add FP16/BF16 regression coverage for all six paged layouts against FP32 SDPA. Signed-off-by: Vladimir Cherepanov --- tests/pytorch/attention/test_attention.py | 90 +++++++++++++++++++ .../cpp_extensions/cudnn_fused_attn.py | 2 +- 2 files changed, 91 insertions(+), 1 deletion(-) diff --git a/tests/pytorch/attention/test_attention.py b/tests/pytorch/attention/test_attention.py index 6eac0d7de8b..f72b0e0699f 100644 --- a/tests/pytorch/attention/test_attention.py +++ b/tests/pytorch/attention/test_attention.py @@ -478,6 +478,96 @@ def compile_plan(_cudnn, graph, _run): assert queried.workspace_size == 17 +@pytest.mark.parametrize("dtype", (torch.float16, torch.bfloat16)) +@pytest.mark.parametrize("q_format", ("bshd", "sbhd", "thd")) +@pytest.mark.parametrize("kv_format", ("bshd", "sbhd")) +def test_fused_attn_paged_cache_layout(dtype, q_format, kv_format): + """Paged cache dimensions must match support queries and independent FP32 SDPA.""" + from torch.nn.attention import SDPBackend, sdpa_kernel + + batch, seqlen, heads, groups, dim = 2, 4, 4, 2, 64 + num_pages, page_size = 6, 128 + lengths = (203, 137) + q_dense = torch.randn(batch, seqlen, heads, dim, device="cuda", dtype=dtype) + q = q_dense.transpose(0, 1).contiguous() if q_format == "sbhd" else q_dense + if q_format == "thd": + q = q.reshape(batch * seqlen, heads, dim) + pages = [ + torch.randn(num_pages, page_size, groups, dim, device="cuda", dtype=dtype) for _ in range(2) + ] + k, v = [ + tensor.transpose(0, 1).contiguous() if kv_format == "sbhd" else tensor for tensor in pages + ] + # Nonconsecutive, distinct K/V mappings also exercise page-table addressing. + table_ids = ((3, 0), (5, 2)), ((4, 1), (0, 3)) + tables = [torch.tensor(ids, device="cuda", dtype=torch.int32) for ids in table_ids] + cu_q = torch.tensor((0, seqlen, batch * seqlen), device="cuda", dtype=torch.int32) + cu_kv = torch.tensor((0, lengths[0], sum(lengths)), device="cuda", dtype=torch.int32) + layout = f"paged_kv_{q_format}_{kv_format}_{kv_format}" + te_dtype = DType.kFloat16 if dtype == torch.float16 else DType.kBFloat16 + config = FusedAttentionParams( + is_training=False, + qkv_layout=QKVLayout[layout], + o_format=QKVFormat[q_format], + qkv_dtype=te_dtype, + o_dtype=te_dtype, + batch_size=batch, + num_attn_heads=heads, + num_gqa_groups=groups, + head_dim_qk=dim, + head_dim_v=dim, + max_seqlen_q=seqlen, + max_seqlen_kv=2 * page_size, + num_tokens_q=batch * seqlen if q_format == "thd" else 0, + num_pages_k=num_pages, + num_pages_v=num_pages, + page_size_k=page_size, + page_size_v=page_size, + max_pages_per_seq_k=2, + max_pages_per_seq_v=2, + attn_mask_type=AttnMaskType["padding"], + ) + backend, message = get_fused_attn_backend(config) + if backend == FusedAttnBackend.No_Backend: + pytest.skip(f"Paged fused attention is unsupported for {layout}: {message}") + out, _ = fused_attn_fwd( + is_training=False, + max_seqlen_q=seqlen, + max_seqlen_kv=2 * page_size, + cu_seqlens_q=cu_q, + cu_seqlens_kv=cu_kv, + q=q, + k=k, + v=v, + fake_dtype=dtype, + fused_attention_backend=backend, + page_table_k=tables[0], + page_table_v=tables[1], + qkv_layout=layout, + o_format=q_format, + attn_mask_type="padding", + ) + expected = [] + with sdpa_kernel(SDPBackend.MATH): + for index, length in enumerate(lengths): + kr, vr = [ + tensor[list(ids[index])] + .reshape(-1, groups, dim)[:length] + .float() + .transpose(0, 1) + .repeat_interleave(heads // groups, dim=0) + for tensor, ids in zip(pages, table_ids) + ] + expected.append( + torch.nn.functional.scaled_dot_product_attention( + q_dense[index].float().transpose(0, 1), kr, vr + ).transpose(0, 1) + ) + actual = out.transpose(0, 1) if q_format == "sbhd" else out.reshape(q_dense.shape) + rtol, atol = (2e-2, 1e-2) if dtype == torch.bfloat16 else (3e-3, 2e-3) + torch.testing.assert_close(actual.float(), torch.stack(expected), rtol=rtol, atol=atol) + + @pytest.mark.parametrize("dtype", (DType.kFloat16, DType.kBFloat16)) @pytest.mark.parametrize("capability", ((10, 0), (12, 0), (12, 1), (12, 9), (13, 0))) @pytest.mark.parametrize("layout", ("bshd_bshd_bshd", "thd_bshd_bshd", "thd_thd_thd")) diff --git a/transformer_engine/pytorch/cpp_extensions/cudnn_fused_attn.py b/transformer_engine/pytorch/cpp_extensions/cudnn_fused_attn.py index 279571e1fe3..82ae790bf9c 100644 --- a/transformer_engine/pytorch/cpp_extensions/cudnn_fused_attn.py +++ b/transformer_engine/pytorch/cpp_extensions/cudnn_fused_attn.py @@ -178,7 +178,7 @@ def matrix( offsets = cu_padded.to(torch.int64) * data.stride(0) self.add(offset_name, offsets, (self.b + 1, 1, 1, 1), (1, 1, 1, 1)) elif name in ("k", "v") and self.params.get("page_table_k") is not None: - order = (0, 2, 1, 3) if fmt == "bshd" else (0, 1, 2, 3) + order = {"bshd": (0, 2, 1, 3), "sbhd": (1, 2, 0, 3), "bhsd": (0, 1, 2, 3)}[fmt] dim = tuple(data.shape[index] for index in order) self.tensors[name] = data self.specs[name] = attention_matrix_spec( From cad318825a21e3b57133b06c1d2e4a63bf299e3b Mon Sep 17 00:00:00 2001 From: Vladimir Cherepanov Date: Fri, 30 Oct 2026 00:30:54 +0000 Subject: [PATCH 20/49] Fall back from fused attention when the cuDNN frontend is unavailable Return No_Backend instead of raising from PyTorch and JAX backend selection when the cuDNN frontend package is missing, its version does not match the C++ build, or it rejects the graph with a frontend error. Report a clear error when building TE/JAX without nvidia-cudnn-frontend. Co-Authored-By: Claude Opus 5.5 Signed-off-by: Vladimir Cherepanov --- build_tools/jax.py | 9 +++++++-- transformer_engine/jax/cpp_extensions/attention.py | 4 +++- .../pytorch/cpp_extensions/cudnn_fused_attn.py | 9 +++++++-- 3 files changed, 17 insertions(+), 5 deletions(-) diff --git a/build_tools/jax.py b/build_tools/jax.py index 56843e63c12..65c40fc020e 100644 --- a/build_tools/jax.py +++ b/build_tools/jax.py @@ -6,7 +6,7 @@ import os import warnings -from importlib.metadata import version as distribution_version +from importlib.metadata import PackageNotFoundError, version as distribution_version from pathlib import Path from packaging import version @@ -30,7 +30,12 @@ def install_requirements() -> List[str]: """Install dependencies for TE/JAX extensions.""" # 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") + try: + frontend_version = distribution_version("nvidia-cudnn-frontend") + except PackageNotFoundError as exc: + raise RuntimeError( + "Building TE/JAX requires nvidia-cudnn-frontend>=1.29.0 to be installed." + ) from exc 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"] diff --git a/transformer_engine/jax/cpp_extensions/attention.py b/transformer_engine/jax/cpp_extensions/attention.py index 06ad5e0c0dc..12cda12d7bf 100644 --- a/transformer_engine/jax/cpp_extensions/attention.py +++ b/transformer_engine/jax/cpp_extensions/attention.py @@ -255,7 +255,9 @@ def get_fused_attn_backend(self): if self.fp8 else NVTE_Fused_Attn_Backend.NVTE_F16_arbitrary_seqlen ) - except ValueError as exc: + except (ValueError, ImportError, RuntimeError) as exc: + # ImportError: frontend package missing; RuntimeError: Python/C++ frontend + # version mismatch. Either way, fall back rather than fail backend selection. message = str(exc) AttentionLogging.setup_logging() diff --git a/transformer_engine/pytorch/cpp_extensions/cudnn_fused_attn.py b/transformer_engine/pytorch/cpp_extensions/cudnn_fused_attn.py index 82ae790bf9c..cf1ab3ea9a5 100644 --- a/transformer_engine/pytorch/cpp_extensions/cudnn_fused_attn.py +++ b/transformer_engine/pytorch/cpp_extensions/cudnn_fused_attn.py @@ -679,7 +679,6 @@ 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 ( @@ -690,6 +689,7 @@ def get_fused_attn_backend(config): ): return FusedAttnBackend.No_Backend, f"Unsupported QKV dtype qkv_dtype={config.qkv_dtype}." try: + cudnn = _import_cudnn_frontend() fwd = _query_graph(config, False) reason = _support_reason( config, fwd, cudnn.backend_version(), torch.cuda.get_device_capability() @@ -700,8 +700,13 @@ def get_fused_attn_backend(config): 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: + except (ImportError, ValueError, RuntimeError) as exc: return FusedAttnBackend.No_Backend, str(exc) + except Exception as exc: # pylint: disable=broad-exception-caught + # cudnnGraphNotSupportedError and other frontend errors derive directly from Exception. + if type(exc).__module__.startswith("cudnn"): + return FusedAttnBackend.No_Backend, str(exc) + raise def fused_attn_fwd(params): From ff6d4a03c3a9b90bb9ce9b9a0d76567454c7e513 Mon Sep 17 00:00:00 2001 From: Vladimir Cherepanov Date: Fri, 30 Oct 2026 00:35:51 +0000 Subject: [PATCH 21/49] Cache the fused attention cache-diagnostics level Resolve NVTE_FUSED_ATTN_CACHE_DEBUG and the rank filter once, on first use, instead of re-reading the environment on every cache event in the attention hot path. Co-Authored-By: Claude Opus 5.5 Signed-off-by: Vladimir Cherepanov --- transformer_engine/common/cudnn/graph.py | 9 ++++++++- 1 file changed, 8 insertions(+), 1 deletion(-) diff --git a/transformer_engine/common/cudnn/graph.py b/transformer_engine/common/cudnn/graph.py index 6393a476924..845f86330de 100644 --- a/transformer_engine/common/cudnn/graph.py +++ b/transformer_engine/common/cudnn/graph.py @@ -192,10 +192,17 @@ def __init__(self, framework): self.framework = framework self.events = Counter() self.times = Counter() + self._level = None atexit.register(self.summary) + def level(self): + """Return the debug level, resolved from the environment on first use.""" + if self._level is None: + self._level = self._resolve_level() + return self._level + @staticmethod - def level(): + def _resolve_level(): """Resolve the existing debug level and rank filter.""" setting = os.getenv("NVTE_FUSED_ATTN_CACHE_DEBUG", "0") level, _, ranks = setting.partition(":") From 8ee7f42a78c8327dc66a61fafe2467c908277261 Mon Sep 17 00:00:00 2001 From: Vladimir Cherepanov Date: Fri, 30 Oct 2026 01:23:40 +0000 Subject: [PATCH 22/49] Propagate fused attention setup errors again Restore the original handling in PyTorch and JAX backend selection: only a configuration rejected by cuDNN falls back to another backend, while a missing frontend, a frontend version mismatch, or a device failure is raised. Falling back hid broken installs behind a silent switch to unfused attention and broke test_fused_attn_query_propagates_setup_errors. This reverts the attention changes from cad31882 and keeps its build_tools/jax.py error message. Co-Authored-By: Claude Opus 5.5 Signed-off-by: Vladimir Cherepanov --- transformer_engine/jax/cpp_extensions/attention.py | 4 +--- .../pytorch/cpp_extensions/cudnn_fused_attn.py | 9 ++------- 2 files changed, 3 insertions(+), 10 deletions(-) diff --git a/transformer_engine/jax/cpp_extensions/attention.py b/transformer_engine/jax/cpp_extensions/attention.py index 12cda12d7bf..06ad5e0c0dc 100644 --- a/transformer_engine/jax/cpp_extensions/attention.py +++ b/transformer_engine/jax/cpp_extensions/attention.py @@ -255,9 +255,7 @@ def get_fused_attn_backend(self): if self.fp8 else NVTE_Fused_Attn_Backend.NVTE_F16_arbitrary_seqlen ) - except (ValueError, ImportError, RuntimeError) as exc: - # ImportError: frontend package missing; RuntimeError: Python/C++ frontend - # version mismatch. Either way, fall back rather than fail backend selection. + except ValueError as exc: message = str(exc) AttentionLogging.setup_logging() diff --git a/transformer_engine/pytorch/cpp_extensions/cudnn_fused_attn.py b/transformer_engine/pytorch/cpp_extensions/cudnn_fused_attn.py index cf1ab3ea9a5..82ae790bf9c 100644 --- a/transformer_engine/pytorch/cpp_extensions/cudnn_fused_attn.py +++ b/transformer_engine/pytorch/cpp_extensions/cudnn_fused_attn.py @@ -679,6 +679,7 @@ 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 ( @@ -689,7 +690,6 @@ def get_fused_attn_backend(config): ): return FusedAttnBackend.No_Backend, f"Unsupported QKV dtype qkv_dtype={config.qkv_dtype}." try: - cudnn = _import_cudnn_frontend() fwd = _query_graph(config, False) reason = _support_reason( config, fwd, cudnn.backend_version(), torch.cuda.get_device_capability() @@ -700,13 +700,8 @@ def get_fused_attn_backend(config): if config.is_training: _query_graph(config, True).entry() return FusedAttnBackend.FP8 if fwd.fp8 else FusedAttnBackend.F16_arbitrary_seqlen, "" - except (ImportError, ValueError, RuntimeError) as exc: + except (cudnn.cudnnGraphNotSupportedError, ValueError) as exc: return FusedAttnBackend.No_Backend, str(exc) - except Exception as exc: # pylint: disable=broad-exception-caught - # cudnnGraphNotSupportedError and other frontend errors derive directly from Exception. - if type(exc).__module__.startswith("cudnn"): - return FusedAttnBackend.No_Backend, str(exc) - raise def fused_attn_fwd(params): From 5b5cf2d744a0b8916ecbcab2d6f189fc20b91b46 Mon Sep 17 00:00:00 2001 From: Vladimir Cherepanov Date: Thu, 5 Nov 2026 20:10:18 +0000 Subject: [PATCH 23/49] Fix FP8 attention test backward mode isolation Explicitly enable FP8 backward before the custom MHA support query so it matches the direct FP8 backward call and skips unsupported H100 configurations. Restore the parameterized tests' FP8 environment settings with monkeypatch to avoid test-order dependencies. Validation: pre-commit and L0_license pass. Focused pytest collection is blocked by the installed native extension missing DType.kFloat8E8M0. Signed-off-by: Vladimir Cherepanov --- tests/pytorch/attention/test_attention.py | 16 +++++++++++----- 1 file changed, 11 insertions(+), 5 deletions(-) diff --git a/tests/pytorch/attention/test_attention.py b/tests/pytorch/attention/test_attention.py index f72b0e0699f..fced977a7bd 100644 --- a/tests/pytorch/attention/test_attention.py +++ b/tests/pytorch/attention/test_attention.py @@ -3893,12 +3893,13 @@ def test_mha_fp8_vs_f16( RoPE, is_training, scaling_mode, + monkeypatch, ): """Test MultiHeadAttention module in FP8""" if not is_training and fp8_dpa_bwd: pytest.skip("fp8_dpa_bwd=True not applicable for inference") - os.environ["NVTE_FP8_DPA_BWD"] = "1" if fp8_dpa_bwd else "0" + monkeypatch.setenv("NVTE_FP8_DPA_BWD", "1" if fp8_dpa_bwd else "0") config = _get_fp8_vs_f16_config(model, qkv_format) # Test backend availability @@ -4136,7 +4137,9 @@ def get_dummy_cuda_rng_tracker() -> CudaRNGStatesTracker: @pytest.mark.parametrize("fp8_dpa_bwd", [True, False]) @pytest.mark.parametrize("is_training", [True, False]) @pytest.mark.parametrize("scaling_mode", ["delayed", "current", "mxfp8"]) -def test_dpa_fp8_vs_f16(dtype, model, qkv_layout, fp8_dpa_bwd, is_training, scaling_mode): +def test_dpa_fp8_vs_f16( + dtype, model, qkv_layout, fp8_dpa_bwd, is_training, scaling_mode, monkeypatch +): """Test DotProductAttention module in FP8""" config = _get_fp8_vs_f16_config(model, qkv_layout) if config.num_heads != config.num_gqa_groups and "3" in qkv_layout: @@ -4154,8 +4157,8 @@ def test_dpa_fp8_vs_f16(dtype, model, qkv_layout, fp8_dpa_bwd, is_training, scal # if get_device_compute_capability() >= (10, 0): # config.dropout_p = 0.1 - os.environ["NVTE_FP8_DPA_BWD"] = "1" if fp8_dpa_bwd else "0" - os.environ["NVTE_UnfusedDPA_Emulate_FP8"] = "1" + monkeypatch.setenv("NVTE_FP8_DPA_BWD", "1" if fp8_dpa_bwd else "0") + monkeypatch.setenv("NVTE_UnfusedDPA_Emulate_FP8", "1") # Test backend availability if scaling_mode == "delayed": @@ -4469,7 +4472,7 @@ def get_dummy_cuda_rng_tracker() -> CudaRNGStatesTracker: @pytest.mark.skipif(not fp8_attn_available, reason=reason_for_no_fp8_attn) @pytest.mark.parametrize("dtype", param_types_fp8) @pytest.mark.parametrize("model", model_configs_fp8) -def test_custom_mha_fp8_vs_f16(dtype, model): +def test_custom_mha_fp8_vs_f16(dtype, model, monkeypatch): """Test FP8 dot product attention implementations based on cuDNN frontend v0.9 and v1.0+. Each test compares results from a custom implementation of an FP8 MHA module, i.e. Custom_MHA_FP8(), to results from an F16 MHA @@ -4478,6 +4481,9 @@ def test_custom_mha_fp8_vs_f16(dtype, model): config = model_configs_fp8[model] + # The custom implementation always executes FP8 backward. + monkeypatch.setenv("NVTE_FP8_DPA_BWD", "1") + # Test backend availability is_training = True fp8_meta = {} From 6737567a1c789ea8d6a66000f8f86cbfd023edd3 Mon Sep 17 00:00:00 2001 From: Vladimir Cherepanov Date: Thu, 5 Nov 2026 14:43:28 -0800 Subject: [PATCH 24/49] Fix host scalar lifetime in JAX cuDNN graph execution Retain immutable aligned scalar buffers across asynchronous cuDNN execution and CUDA graph replay to prevent corrupted FP8 GQA key and value gradients. Extend numerical coverage to different attention scales sharing compiled plans. Signed-off-by: Vladimir Cherepanov --- tests/jax/test_fused_attn.py | 19 ++++++++------- .../jax/csrc/extensions/cudnn_graph.cpp | 23 ++++++++++++++++--- 2 files changed, 31 insertions(+), 11 deletions(-) diff --git a/tests/jax/test_fused_attn.py b/tests/jax/test_fused_attn.py index 8e14fe4805c..b658f56f92c 100644 --- a/tests/jax/test_fused_attn.py +++ b/tests/jax/test_fused_attn.py @@ -2593,6 +2593,7 @@ def forbidden(*_args, **_kwargs): np.testing.assert_array_equal(gradients, tensor) +@pytest.mark.parametrize("scaling_factor", [0.125, 0.25]) @pytest.mark.parametrize( "kind,sm_arch", [ @@ -2607,8 +2608,8 @@ def forbidden(*_args, **_kwargs): ("ragged_dense_stats_f16", 129), ], ) -def test_cudnn_python_fp8_gqa_masks(kind, sm_arch, monkeypatch, request): - """Compare non-uniform FP8 attention and gradients against native FP32 GQA.""" +def test_cudnn_python_fp8_gqa_masks(kind, sm_arch, scaling_factor, monkeypatch, request): + """Compare FP8 GQA against FP32, including cached plans with different scalar values.""" ragged = kind.startswith("ragged") fp8 = not kind.endswith("_f16") if ragged and get_device_compute_capability(0) < 100: @@ -2662,13 +2663,14 @@ def test_cudnn_python_fp8_gqa_masks(kind, sm_arch, monkeypatch, request): max_segments_per_seq=2 if ragged else 1, window_size=window if window is not None else (-1, -1), fp8=fp8, + attn_scale=scaling_factor, return_max_logit=True, ) 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.einsum("bqhd,bkhd->bhqk", query, key) * scaling_factor scores = jnp.where(mask[:, None], scores, -jnp.inf) result = jnp.einsum("bhqk,bkhd->bqhd", jax.nn.softmax(scores, axis=-1), value) if kind == "padding": @@ -2687,7 +2689,7 @@ def attention(arrays): attn_mask_type=attn_mask, qkv_layout=layout, softmax_type=AttnSoftmaxType.VANILLA_SOFTMAX, - scaling_factor=0.125, + scaling_factor=scaling_factor, dropout_probability=0.0, is_training=True, max_segments_per_seq=2 if ragged else 1, @@ -2711,8 +2713,9 @@ def relative_error(actual, expected): 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) + for name, actual_grad, expected_grad in zip(("dQ", "dK", "dV"), actual_grads, expected_grads): + error = relative_error(actual_grad, expected_grad) + assert error < (0.1 if fp8 else 0.01), f"{name} relative error: {error}" _, max_logit = jax.jit( lambda xs: fused_attn( xs, @@ -2723,7 +2726,7 @@ def relative_error(actual, expected): attn_mask_type=attn_mask, qkv_layout=layout, softmax_type=AttnSoftmaxType.VANILLA_SOFTMAX, - scaling_factor=0.125, + scaling_factor=scaling_factor, dropout_probability=0.0, is_training=True, max_segments_per_seq=2 if ragged else 1, @@ -2738,7 +2741,7 @@ def relative_error(actual, expected): q.astype(jnp.float32), jnp.repeat(k.astype(jnp.float32), 2, axis=-2), ) - * 0.125 + * scaling_factor ) max_mask = mask if kind == "padding": diff --git a/transformer_engine/jax/csrc/extensions/cudnn_graph.cpp b/transformer_engine/jax/csrc/extensions/cudnn_graph.cpp index 157a7d17cfe..b00bde6d77f 100644 --- a/transformer_engine/jax/csrc/extensions/cudnn_graph.cpp +++ b/transformer_engine/jax/csrc/extensions/cudnn_graph.cpp @@ -93,6 +93,24 @@ struct Scalar { alignas(16) std::array value{}; }; +const std::vector &GetScalars(ffi::Span values) { + // cuDNN's FP8 GQA reductions can read host scalars after execute returns. + // Immutable cached storage must also outlive CUDA graph capture and replay. + const std::string key(values.begin(), values.end()); + static std::mutex mutex; + static std::unordered_map> cache; + std::lock_guard lock(mutex); + auto [entry, inserted] = cache.try_emplace(key); + if (inserted) { + auto &scalars = entry->second; + scalars.resize(values.size() / 16); + for (size_t i = 0; i < scalars.size(); ++i) { + std::copy_n(values.begin() + i * 16, 16, scalars[i].value.begin()); + } + } + return entry->second; +} + ffi::Error InitializeGraph(cudaStream_t stream, ffi::RemainingArgs, ffi::RemainingRets, ffi::Dictionary attrs) { try { @@ -157,13 +175,12 @@ ffi::Error ExecuteGraph(cudaStream_t stream, ffi::RemainingArgs args, ffi::Remai scalar_values.size() != scalar_uids.size() * 16) { throw std::runtime_error("Invalid cuDNN pass-by-value scalar bindings."); } - std::vector scalars(scalar_uids.size()); + const auto &scalars = GetScalars(scalar_values); 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()); + pointers.emplace(scalar_uids[i], const_cast(scalars[i].value.data())); } auto workspace = rets.get(rets.size() - 1); if (workspace.has_error()) throw std::runtime_error("Missing cuDNN graph workspace."); From a2596cec3cd56e77d7ceeadd04eaa8bd2614afdb Mon Sep 17 00:00:00 2001 From: Vladimir Cherepanov Date: Fri, 6 Nov 2026 20:22:09 +0000 Subject: [PATCH 25/49] Fix JAX attention documentation formatting Signed-off-by: Vladimir Cherepanov --- docs/examples/jax/attention.rst | 2 +- transformer_engine/jax/attention.py | 13 ++++++++----- 2 files changed, 9 insertions(+), 6 deletions(-) diff --git a/docs/examples/jax/attention.rst b/docs/examples/jax/attention.rst index 13489601830..54906d368e0 100644 --- a/docs/examples/jax/attention.rst +++ b/docs/examples/jax/attention.rst @@ -17,7 +17,7 @@ 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 diff --git a/transformer_engine/jax/attention.py b/transformer_engine/jax/attention.py index b12942ce0d7..de85641afa0 100644 --- a/transformer_engine/jax/attention.py +++ b/transformer_engine/jax/attention.py @@ -1623,11 +1623,15 @@ def fused_attn( """ Perform cuDNN fused attention. - This function implements the following formula: + This function implements the following formula:: + BMM1 -> (PreBias) -> ScaleMaskSoftmax -> (PostBias) -> (Dropout) -> BMM2 + Args: qkv (Tuple[jnp.ndarray, ...]): A tuple containing query, key, and value tensors. - It supports three formats: + + It supports three formats: + - `(qkv_packed,)`: For interleaved QKV packed format, typically used when query, key, and value have the same shape (e.g., self-attention). - `(query, kv_packed)`: For separate query and KV packed format, typically used when @@ -1687,9 +1691,8 @@ def fused_attn( return_max_logit (bool): If True, also return per-head maximum attention logits with shape ``[h]``. Returns: - jnp.ndarray: - Attention output when ``return_max_logit`` is False. - tuple[jnp.ndarray, jnp.ndarray]: + jnp.ndarray or tuple[jnp.ndarray, jnp.ndarray]: + Attention output when ``return_max_logit`` is False, or ``(output, max_logit)`` when ``return_max_logit`` is True. Examples (non-THD, also known as non-packed): From 8598bcd39698251b67602a1a60d00a876271ad50 Mon Sep 17 00:00:00 2001 From: Vladimir Cherepanov Date: Tue, 10 Nov 2026 21:34:04 +0000 Subject: [PATCH 26/49] Simplify cuDNN attention documentation Keep the JAX attention overview focused on supported module options and remove version-specific support tables, workarounds, and implementation details. Retain installation requirements, clarify cache diagnostics and backward precision, move callback cache-key guidance to API docstrings, and remove internal JAX functions from the public API reference. Validation: pre-commit, both framework Python lint suites, changed-file license checks, and docstring-only AST comparisons passed. The full Sphinx build could not proceed without generated Doxygen XML in this container. Signed-off-by: Vladimir Cherepanov --- docs/api/jax.rst | 7 -- docs/envvars.rst | 4 +- docs/examples/jax/attention.rst | 80 +++---------------- docs/installation.rst | 28 +------ transformer_engine/jax/attention.py | 8 +- transformer_engine/jax/flax/transformer.py | 6 +- .../dot_product_attention.py | 4 + 7 files changed, 31 insertions(+), 106 deletions(-) diff --git a/docs/api/jax.rst b/docs/api/jax.rst index 5bc498d7d42..9157ff16d9a 100644 --- a/docs/api/jax.rst +++ b/docs/api/jax.rst @@ -60,13 +60,6 @@ Modules .. autoapifunction:: transformer_engine.jax.flax.extend_logical_axis_rules -Attention functions -------------------- - -.. autoapifunction:: transformer_engine.jax.attention.fused_attn -.. autoapifunction:: transformer_engine.jax.attention.quantize_attention_qkv -.. autoapifunction:: transformer_engine.jax.attention.is_fused_attn_kernel_available - Mixture of Experts ------------------ Router, routing kernels, and grouped dense for MoE layers. See diff --git a/docs/envvars.rst b/docs/envvars.rst index 64dae198800..04108f8b773 100644 --- a/docs/envvars.rst +++ b/docs/envvars.rst @@ -240,7 +240,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 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. + :Description: Log FusedAttention graph cache activity to stderr, prefixed with ``[FUSED-ATTN-CACHE]``. ``1`` prints a summary of cache counters and build timings at process exit. ``2`` also prints individual events. At both levels, JAX diagnostics describe graph construction and caching during tracing, rather than each execution of a compiled function. 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 @@ -282,7 +282,7 @@ FP8 Configuration :Type: ``int`` (0 or 1) :Default: ``1`` - :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``. + :Description: When FP8 attention is enabled in PyTorch or JAX, control whether the backward pass also uses FP8. ``1`` = FP8 forward and backward, ``0`` = FP8 forward and FP16/BF16 backward. .. envvar:: NVTE_DPA_FP8CS_O_in_F16 diff --git a/docs/examples/jax/attention.rst b/docs/examples/jax/attention.rst index 54906d368e0..17aed37f453 100644 --- a/docs/examples/jax/attention.rst +++ b/docs/examples/jax/attention.rst @@ -16,74 +16,18 @@ 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 THD attention requires cuDNN 9.23 or later. FP8 THD backward with off-by-one -or learnable softmax requires cuDNN 9.26 or later; select -``NVTE_FP8_DPA_BWD=0`` to use FP16/BF16 backward instead. For FP16/BF16 THD -learnable-softmax backward with packed statistics, cuDNN versions before 9.26 -remain eligible for value-head dimensions 64, 128, and 256. Other value-head -dimensions require cuDNN 9.26 or later. Execution-plan support still depends on -the complete attention configuration. - -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. -Pass ``fp8="mxfp8"`` to select MXFP8 attention on SM100 or later. MXFP8 uses -E4M3 operands with E8M0 block scales along the attention contraction axes in -both passes. Inputs, outputs and gradients remain FP16/BF16. Sequence lengths -must be divisible by 128 and head dimensions by 32. This mode currently supports -separate, QKV-packed and KV-packed BSHD inputs. THD, SM12x, and cuDNN -9.23.0/9.23.1 are excluded. It requires -cuDNN 9.21 or later. ``NVTE_FP8_DPA_BWD=0`` keeps the original FP16/BF16 operands -for MXFP8 backward, matching the PyTorch MXFP8 path. -Attention bias, ALiBi, dropout and ``return_max_logit`` are excluded from this mode. -Exact head-dimension availability is checked by cuDNN. Frontend 1.29.0 has an -incorrect MXFP8 backward scale stride for Q/K heads wider than 128; these -configurations require ``NVTE_FP8_DPA_BWD=0``. This check protects both frontends. - -To reuse quantizations, call ``quantize_attention_qkv((q, k, v))`` and pass its -result to ``fused_attn(..., fp8="mxfp8", quantized_qkv=quantized)``. Each operand -is a JAX ``ScaledTensor2x`` with independent rowwise/columnwise E4M3 data and -unswizzled E8M0 scales in BHSD order. Equivalent operands from a producer may -be supplied directly. The original BSHD arrays supply the gradient targets and -the original FP16/BF16 operands for ``NVTE_FP8_DPA_BWD=0``. These auxiliary -buffers are non-differentiable. Scale reordering happens after partitioning, -so data and tensor parallel sharding preserve per-head scale boundaries. -Ring and AllGather context parallelism communicate FP16/BF16 operands and -quantize each local attention step. Prequantized inputs are decoded before CP -communication, as in the PyTorch MXFP8 implementation. Separate and packed -BSHD storage, GQA, and causal dual-chunk load balancing are supported. Every -local attention step must meet the 128-token alignment requirement; 256 tokens -per CP rank accommodates the half-sequence steps used by causal load balancing -and AllGather. The existing JAX CP restrictions on sliding windows, padding, -dropout and non-vanilla softmax still apply. - -Paged KV-cache APIs, quantized module outputs, and automatic quantized -projection-to-attention integration require further JAX API and frontend-binding -work. ``score_mod`` currently supports FP16/BF16 only. +FP8 attention +-------------- + +Set ``fp8=True`` on the Flax ``DotProductAttention``, ``MultiHeadAttention`` +or ``TransformerLayer`` module to enable FP8 attention with current tensor +scaling. Set ``fp8="mxfp8"`` for block-scaled MXFP8 attention. These options +are independent of the GEMM quantization recipe. Inputs, outputs and input +gradients retain their FP16/BF16 dtype. + +:envvar:`NVTE_FP8_DPA_BWD` controls whether the backward pass also uses FP8. +Availability depends on the GPU, installed cuDNN version, and attention +configuration. ``AttnBiasType.ALIBI`` (``attn_bias_type="alibi"`` in Flax) enables cuDNN's ALiBi mask without an explicit bias tensor. Explicit dropout keys have JAX's diff --git a/docs/installation.rst b/docs/installation.rst index 36f52bce1bb..1314910c7bd 100644 --- a/docs/installation.rst +++ b/docs/installation.rst @@ -16,30 +16,10 @@ Prerequisites 3. |driver link|_ supporting CUDA 12.1 or later. 4. `cuDNN 9.12 `__ or later. -The PyTorch and JAX cuDNN attention backends use shared graph builders through -the cuDNN frontend Python API. Both installations require -``nvidia-cudnn-frontend`` version 1.29.0 or later. JAX builds pin their runtime -dependency to the exact frontend version that supplies the build headers. -Graph descriptions are cached by tensor metadata -and target device or architecture. Support queries select supported plans without -compiling their kernels. PyTorch compiles on first execution and uses its CUDA -generator for dropout, including CUDA graph capture and replay. JAX compiles -during tracing when workspace sizes are needed, then serializes plans for its -native FFI executor. The executor requires a serialized precompiled execution -plan and reports an error if it cannot be restored. 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. +PyTorch and JAX cuDNN attention require ``nvidia-cudnn-frontend`` version +1.29.0 or later. JAX builds pin the runtime dependency to the frontend version +used for the build headers. The JAX extension and frontend Python package must +use matching 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/transformer_engine/jax/attention.py b/transformer_engine/jax/attention.py index de85641afa0..91437b7a9a9 100644 --- a/transformer_engine/jax/attention.py +++ b/transformer_engine/jax/attention.py @@ -1679,10 +1679,10 @@ 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 | str): True enables tensor-scaled cuDNN FP8 SDPA; "mxfp8" - enables MXFP8 block-scaled attention on SM100 or later. Outputs and - gradients retain the input FP16/BF16 dtype. MXFP8 uses E4M3 operands - in both passes, with independent scales along the contraction axes. + fp8 (bool | str): False disables FP8 attention; True enables FP8 attention + with current tensor scaling; "mxfp8" enables block-scaled MXFP8 + attention. This option is independent of the GEMM quantization recipe. + Outputs and gradients retain the input FP16/BF16 dtype. quantized_qkv (Optional[Tuple[ScaledTensor2x, ...]]): Prequantized MXFP8 Q/K/V, each with independent unswizzled BHSD rowwise/columnwise data and E8M0 scales. Requires ``fp8="mxfp8"``. The original ``qkv`` diff --git a/transformer_engine/jax/flax/transformer.py b/transformer_engine/jax/flax/transformer.py index 9c7faf7a997..10389518778 100644 --- a/transformer_engine/jax/flax/transformer.py +++ b/transformer_engine/jax/flax/transformer.py @@ -628,6 +628,10 @@ class DotProductAttention(nn.Module): # pylint: disable=too-few-public-methods score_mod: Optional[Callable], default = None Experimental cuDNN frontend score modification callback. When set, fused attention is required and the same score_mod limitations as :func:`fused_attn` apply. + For stateful callbacks, provide ``score_mod_graph_cache_key`` on the callable + or bound-method owner to enable graph caching. The key must describe the graph + structure and change when that structure changes; it must not contain tensors. + Pass runtime data through ``score_mod_tensors`` and ``score_mod_bprop_tensors``. score_mod_bprop: Optional[Callable], default = None Optional score modification backward callback. score_mod_tensors: Optional[Mapping[str, Any]], default = None @@ -642,7 +646,7 @@ class DotProductAttention(nn.Module): # pylint: disable=too-few-public-methods 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. Set to ``"mxfp8"`` for - block-scaled MXFP8 attention on SM100 or later. + block-scaled MXFP8 attention. Optimization parameters ----------------------- diff --git a/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py b/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py index 9a3c8a44dbf..4e52f27209f 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py @@ -2190,6 +2190,10 @@ def forward( and is mutually exclusive with masks, bias, ALiBi, sink attention, dropout, FP8, context parallelism, THD format, KV caching, and return_max_logit. The callback signature is ``score_mod(graph, score, tensors) -> score``. + For stateful callbacks, provide ``score_mod_graph_cache_key`` on the callable + or bound-method owner to enable graph caching. The key must describe the graph + structure and change when that structure changes; it must not contain tensors. + Pass runtime data through ``score_mod_tensors`` and ``score_mod_bprop_tensors``. score_mod_bprop: Optional[Callable], default = None Optional cuDNN frontend callback for the backward pass of score_mod. The callback signature is ``score_mod_bprop(graph, dP, tensors) -> dP``. From a93a55f8f01cbf2363b4da82d981104951ccd4e2 Mon Sep 17 00:00:00 2001 From: Vladimir Cherepanov Date: Wed, 11 Nov 2026 02:15:00 +0000 Subject: [PATCH 27/49] Match common attention batch bucketing Separate logical batch sizes from graph capacities in the PyTorch cuDNN attention path. Bucket eligible ragged backward graphs and converted-length forward graphs, padding sequence metadata and terminal offsets. Use cumulative lengths and ragged-offset multipliers for eligible dropout-free forward graphs, matching the common path's version gates. Extend cache-reuse and numerical coverage across batch buckets, including packed padding, mixed layouts, FP8 backward precision, and dropout checked against an independent FP32 reference. Validation: 394 selected PyTorch cases passed and 281 were skipped. Pre-commit, both framework Python lint suites, and edited-file license checks passed. Validation used GB200 with PyTorch 2.15.0a0, CUDA 13.4, cuDNN 9.27, and cuDNN Frontend 1.31. JAX runtime validation was unavailable because JAX is not installed. Signed-off-by: Vladimir Cherepanov --- tests/pytorch/attention/test_attention.py | 170 ++++++++++++++++-- transformer_engine/common/cudnn/attention.py | 11 +- transformer_engine/common/cudnn/graph.py | 27 ++- .../cpp_extensions/cudnn_fused_attn.py | 98 ++++++++-- 4 files changed, 277 insertions(+), 29 deletions(-) diff --git a/tests/pytorch/attention/test_attention.py b/tests/pytorch/attention/test_attention.py index fced977a7bd..1975fddc61b 100644 --- a/tests/pytorch/attention/test_attention.py +++ b/tests/pytorch/attention/test_attention.py @@ -392,7 +392,8 @@ def common_attention(*_args, **_kwargs): @pytest.mark.parametrize("layout", ("bs3hd", "bsh3d", "sb3hd", "sbh3d", "t3hd", "th3d")) @pytest.mark.parametrize("backward", (False, True)) -def test_fused_attn_query_reuses_runtime_graph(layout, backward, monkeypatch): +@pytest.mark.parametrize("dropout", (0.0, 0.125)) +def test_fused_attn_query_reuses_runtime_graph(layout, backward, dropout, 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 @@ -418,7 +419,8 @@ def test_fused_attn_query_reuses_runtime_graph(layout, backward, monkeypatch): attn_mask_type="padding" if fmt == "thd" else "no_mask", attn_bias_type="no_bias", softmax_type="vanilla", - dropout=0.0, + dropout=dropout, + rng_state=torch.zeros(2, device="cuda", dtype=torch.int64), attn_scale=0.125, window_size=(-1, -1), bottom_right_diagonal=True, @@ -436,6 +438,7 @@ def test_fused_attn_query_reuses_runtime_graph(layout, backward, monkeypatch): 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"]], + dropout=dropout, ) query = implementation._query_graph(config, backward) runtime = implementation._AttentionGraph(params, backward) @@ -470,6 +473,10 @@ def compile_plan(_cudnn, graph, _run): monkeypatch.setattr(implementation, "_execute_cudnn_graph", lambda *_args: None) queried = query.entry() assert queried.workspace_size is None and not compiles + if fmt == "thd" and not query.use_cu_seqlens: + other_query = implementation._query_graph(replace(config, batch_size=batch + 1), backward) + assert other_query.description() == query.description() + assert other_query.entry() is queried runtime.scalar("attn_scale", torch.tensor([0.25]), by_value=True) runtime.execute() runtime.execute() @@ -478,6 +485,74 @@ def compile_plan(_cudnn, graph, _run): assert queried.workspace_size == 17 +@pytest.mark.parametrize( + "backward,dropout,version,fp8,layout,capability,expected_batch,direct", + ( + (False, 0.0, 92300, False, "thd_thd_thd", (10, 0), 64, False), + (False, 0.0, 92400, False, "thd_thd_thd", (10, 0), 33, True), + (False, 0.0, 92400, True, "thd_thd_thd", (10, 0), 64, False), + (False, 0.0, 92500, True, "thd_thd_thd", (10, 0), 33, True), + (False, 0.125, 92700, False, "thd_thd_thd", (10, 0), 64, False), + (True, 0.0, 92700, False, "thd_thd_thd", (10, 0), 64, False), + (True, 0.0, 92700, True, "thd_thd_thd", (10, 0), 64, False), + (True, 0.0, 92700, False, "thd_bshd_bshd", (10, 0), 64, False), + (True, 0.0, 92700, False, "thd_sbhd_sbhd", (10, 0), 64, False), + (True, 0.0, 92700, False, "thd_thd_thd", (12, 0), 33, False), + (True, 0.0, 92700, False, "thd_thd_thd", (12, 1), 33, False), + (True, 0.0, 92700, False, "bshd_bshd_bshd", (10, 0), 33, False), + ), +) +def test_fused_attn_ragged_batch_capacity( + backward, dropout, version, fp8, layout, capability, expected_batch, direct, monkeypatch +): + """Match common forward/backward batch capacities and cumulative-length inputs.""" + import cudnn + from transformer_engine.pytorch.cpp_extensions import cudnn_fused_attn as implementation + + monkeypatch.setattr(cudnn, "backend_version", lambda: version) + monkeypatch.setattr(torch.cuda, "get_device_capability", lambda *_args: capability) + dtype = DType.kFloat8E4M3 if fp8 else DType.kBFloat16 + config = FusedAttentionParams( + qkv_layout=QKVLayout[layout], + qkv_dtype=dtype, + batch_size=33, + num_attn_heads=4, + num_gqa_groups=4, + head_dim_qk=64, + head_dim_v=64, + max_seqlen_q=128, + max_seqlen_kv=128, + num_tokens_q=4096, + num_tokens_kv=4096, + attn_mask_type=AttnMaskType["padding"], + dropout=dropout, + ) + query = implementation._query_graph(config, backward) + assert query.b == 33 + assert query.graph_b == expected_batch + assert query.use_cu_seqlens == direct + assert query.specs["q"].dim[0] == expected_batch + assert query.specs["seq_q"].dim == (expected_batch + direct, 1, 1, 1) + if layout.startswith("thd"): + assert query.specs["q_offset"].dim == (expected_batch + 1, 1, 1, 1) + assert query.specs["q_offset"].dtype == ("INT32" if direct else "INT64") + assert query.specs["q"].ragged_offset_multiplier == (256 if direct else 1) + + graph = Mock() + if fp8: + operation = graph.sdpa_fp8_backward if backward else graph.sdpa_fp8 + operation.return_value = tuple(Mock() for _ in range(7 if backward else 4)) + else: + operation = graph.sdpa_backward if backward else graph.sdpa + operation.return_value = tuple(Mock() for _ in range(3 if backward else 2)) + implementation.build_attention_graph(cudnn, graph, query.description()) + kwargs = operation.call_args.kwargs + assert ("cu_seq_len_q" in kwargs) == direct + assert ("seq_len_q" in kwargs) == (not direct) + if direct: + assert kwargs["implementation"] == cudnn.attention_implementation.UNIFIED + + @pytest.mark.parametrize("dtype", (torch.float16, torch.bfloat16)) @pytest.mark.parametrize("q_format", ("bshd", "sbhd", "thd")) @pytest.mark.parametrize("kv_format", ("bshd", "sbhd")) @@ -656,17 +731,29 @@ def test_fused_attn_fake_stats_shape(capability, layout, monkeypatch): @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): +@pytest.mark.parametrize("batch", (3, 33)) +def test_fused_attn_packed_backward(dtype, kv_format, batch): """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)} + q_lengths, kv_lengths = (64, 128, 32) * (batch // 3), (96, 160, 64) * (batch // 3) + q_offsets = [0] + for capacity in (128, 256, 128) * (batch // 3): + q_offsets.append(q_offsets[-1] + capacity) + kv_capacity = 384 if kv_format == "thd" else 256 + kv_offsets = tuple(index * kv_capacity for index in range(batch + 1)) + cu_q = torch.tensor((0, *q_lengths), device="cuda", dtype=torch.int32).cumsum( + 0, dtype=torch.int32 + ) + cu_kv = torch.tensor((0, *kv_lengths), device="cuda", dtype=torch.int32).cumsum( + 0, dtype=torch.int32 + ) + q = torch.randn(q_offsets[-1], 4, 64, device="cuda", dtype=dtype) + kv_shape = { + "thd": (kv_offsets[-1], 4, 64), + "bshd": (batch, 256, 4, 64), + "sbhd": (256, batch, 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( @@ -722,6 +809,69 @@ def test_fused_attn_packed_backward(dtype, kv_format): ) +@pytest.mark.parametrize("dtype", (torch.float16, torch.bfloat16)) +@pytest.mark.parametrize("batch", (3, 33)) +def test_fused_attn_ragged_dropout(dtype, batch): + """Bucketed forward/backward graphs preserve dropout numerics and packed padding.""" + if device_compute_capability < (9, 0): + pytest.skip("FP16/BF16 THD fused attention requires SM90 or later") + heads, dim, capacity, dropout = 2, 64, 64, 0.25 + lengths = (16, 32, 64) * (batch // 3) + cu = torch.tensor((0, *lengths), device="cuda", dtype=torch.int32).cumsum(0, dtype=torch.int32) + offsets = torch.arange(batch + 1, device="cuda", dtype=torch.int32) * capacity + q, k = [ + 0.1 * torch.randn(batch * capacity, heads, dim, device="cuda", dtype=dtype) + for _ in range(2) + ] + v = torch.zeros_like(q) + for index, length in enumerate(lengths): + start = index * capacity + v[start : start + length] = torch.eye(length, dim, device="cuda", dtype=dtype)[:, None, :] + kwargs = dict( + max_seqlen_q=capacity, + max_seqlen_kv=capacity, + cu_seqlens_q=cu, + cu_seqlens_kv=cu, + cu_seqlens_q_padded=offsets, + cu_seqlens_kv_padded=offsets, + q=q, + k=k, + v=v, + fake_dtype=dtype, + fused_attention_backend=FusedAttnBackend.F16_arbitrary_seqlen, + qkv_layout="thd_thd_thd", + o_format="thd", + attn_mask_type="padding", + dropout=dropout, + ) + 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="thd_thd_thd", + **kwargs, + ) + refs = [tensor.float().requires_grad_() for tensor in (q, k, v)] + reference = torch.zeros_like(refs[0]) + for index, length in enumerate(lengths): + active = slice(index * capacity, index * capacity + length) + qr, kr, vr = [tensor[active].transpose(0, 1) for tensor in refs] + # Identity V exposes the sampled dropout mask in O. Recompute attention + # independently in FP32 using that mask, without depending on cuDNN's RNG algorithm. + mask = out[active, :, :length].transpose(0, 1) != 0 + assert mask.any() and (~mask).any() + probabilities = torch.softmax(qr @ kr.transpose(-1, -2) / dim**0.5, dim=-1) + reference[active] = ((probabilities * mask / (1.0 - dropout)) @ vr).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): + torch.testing.assert_close(actual.float(), ref.grad, 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/transformer_engine/common/cudnn/attention.py b/transformer_engine/common/cudnn/attention.py index 311b4156311..5ee27c2fc88 100644 --- a/transformer_engine/common/cudnn/attention.py +++ b/transformer_engine/common/cudnn/attention.py @@ -21,6 +21,7 @@ class AttentionOptions: left_bound: int | None = None right_bound: int | None = None padding: bool = False + use_cu_seqlens: bool = False alibi: bool = False dropout: float = 0.0 deterministic: bool = False @@ -106,7 +107,15 @@ def tensor(name): 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")) + kwargs["use_padding_mask"] = True + if options.use_cu_seqlens: + kwargs.update( + cu_seq_len_q=tensor("seq_q"), + cu_seq_len_kv=tensor("seq_kv"), + implementation=cudnn.attention_implementation.UNIFIED, + ) + else: + kwargs.update(seq_len_q=tensor("seq_q"), seq_len_kv=tensor("seq_kv")) if options.alibi: kwargs["use_alibi_mask"] = True if "bias" in specs: diff --git a/transformer_engine/common/cudnn/graph.py b/transformer_engine/common/cudnn/graph.py index cbb2337a128..8651042a3d6 100644 --- a/transformer_engine/common/cudnn/graph.py +++ b/transformer_engine/common/cudnn/graph.py @@ -88,6 +88,7 @@ class TensorSpec: ragged_offset: str = "" reordered: bool = False by_value: bool = False + ragged_offset_multiplier: int = 1 def apply_tensor_spec(cudnn, tensor, spec, uid=None): @@ -99,6 +100,8 @@ def apply_tensor_spec(cudnn, tensor, spec, uid=None): tensor.set_uid(uid) if spec.output: tensor.set_output(True) + if spec.ragged_offset_multiplier != 1: + tensor.set_ragged_offset_multiplier(spec.ragged_offset_multiplier) return tensor @@ -150,7 +153,15 @@ def attention_stride(dim, fmt): def attention_matrix_spec( - dim, fmt, dtype, *, storage_stride=None, packed="", output=False, ragged_offset="" + dim, + fmt, + dtype, + *, + storage_stride=None, + packed="", + output=False, + ragged_offset="", + ragged_offset_multiplier=1, ): """Describe either actual storage strides or a hypothetical support-query tensor.""" if storage_stride is not None and fmt == "thd": @@ -168,10 +179,19 @@ def attention_matrix_spec( 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) + return TensorSpec( + tuple(dim), + stride, + dtype, + output=output, + ragged_offset=ragged_offset, + ragged_offset_multiplier=ragged_offset_multiplier, + ) -def attention_stats_spec(batch, heads, seqlen, *, ragged=False, output=False): +def attention_stats_spec( + batch, heads, seqlen, *, ragged=False, output=False, ragged_offset_multiplier=1 +): """Describe dense BHS1 or packed TH1 softmax statistics.""" stride = (heads * seqlen, 1, heads, 1) if ragged else (heads * seqlen, seqlen, 1, 1) return TensorSpec( @@ -180,6 +200,7 @@ def attention_stats_spec(batch, heads, seqlen, *, ragged=False, output=False): "FLOAT", output=output, ragged_offset="stats_offset" if ragged else "", + ragged_offset_multiplier=ragged_offset_multiplier, ) diff --git a/transformer_engine/pytorch/cpp_extensions/cudnn_fused_attn.py b/transformer_engine/pytorch/cpp_extensions/cudnn_fused_attn.py index 82ae790bf9c..f6be18829a3 100644 --- a/transformer_engine/pytorch/cpp_extensions/cudnn_fused_attn.py +++ b/transformer_engine/pytorch/cpp_extensions/cudnn_fused_attn.py @@ -138,17 +138,31 @@ def __init__( 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 + 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 capability = torch.cuda.get_device_capability(self.device) self.ragged_graph = uses_ragged_attention_graph(10 * capability[0] + capability[1]) self.ragged_stats = self.q_format == "thd" and self.ragged_graph + # The supported Python frontend versions include cumulative-length inputs. + # Like the common path, use them for dropout-free forward graphs only. + self.use_cu_seqlens = ( + not backward + and "padding" in params["attn_mask_type"] + and not params["dropout"] + and _import_cudnn_frontend().backend_version() >= (92500 if self.fp8 else 92400) + ) + ragged = self.q_format == "thd" or self.kv_format == "thd" + self.graph_b = ( + bucket(self.b, 32, 512) + if ragged and self.ragged_graph and not self.use_cu_seqlens + else self.b + ) self.sq = bucket(self.tq, 1024, 32768) if self.ragged_stats else params["max_seqlen_q"] self.skv = ( bucket(self.tk, 1024, 32768) if self.kv_format == "thd" and self.ragged_graph else params["max_seqlen_kv"] ) - self.fp8 = hasattr(params["q"], "_fp8_dtype") if fp8 is None else fp8 - self.mxfp8 = isinstance(params["q"], MXFP8TensorStorage) if mxfp8 is None else mxfp8 self.specs = {} self.tensors = {} @@ -172,11 +186,15 @@ def matrix( ): """Bind storage using the same matrix descriptors as support queries.""" data = _raw(tensor, columnwise) - dim = (self.b, heads, seqlen, tensor.shape[-1]) + dim = (self.graph_b, heads, seqlen, tensor.shape[-1]) offset_name = name + "_offset" if fmt == "thd" else "" if offset_name: - offsets = cu_padded.to(torch.int64) * data.stride(0) - self.add(offset_name, offsets, (self.b + 1, 1, 1, 1), (1, 1, 1, 1)) + self.add( + offset_name, + self.offsets(cu_padded, data.stride(0)), + (self.graph_b + 1, 1, 1, 1), + (1, 1, 1, 1), + ) elif name in ("k", "v") and self.params.get("page_table_k") is not None: order = {"bshd": (0, 2, 1, 3), "sbhd": (1, 2, 0, 3), "bhsd": (0, 1, 2, 3)}[fmt] dim = tuple(data.shape[index] for index in order) @@ -188,17 +206,36 @@ def matrix( storage_stride=data.stride(), output=output, ragged_offset=offset_name, + ragged_offset_multiplier=data.stride(0) if offset_name and self.use_cu_seqlens else 1, ) return name + def offsets(self, cu_padded, multiplier): + """Bind token offsets directly or pad converted offsets with the terminal value.""" + if self.use_cu_seqlens: + return cu_padded + offsets = cu_padded.to(torch.int64) * multiplier + if self.graph_b > self.b: + offsets = torch.cat((offsets, offsets[-1:].expand(self.graph_b - self.b))) + return offsets + def stats(self, name, tensor, output=False): """Bind dense or packed softmax statistics.""" if self.ragged_stats: - offsets = self.params["cu_seqlens_q_padded"].to(torch.int64) * self.h - self.add("stats_offset", offsets, (self.b + 1, 1, 1, 1), (1, 1, 1, 1)) + offsets = self.offsets(self.params["cu_seqlens_q_padded"], self.h) + self.add("stats_offset", offsets, (self.graph_b + 1, 1, 1, 1), (1, 1, 1, 1)) + elif tensor.shape[0] < self.graph_b: + tensor = torch.nn.functional.pad( + tensor, (0, 0, 0, 0, 0, 0, 0, self.graph_b - tensor.shape[0]) + ) self.tensors[name] = tensor self.specs[name] = attention_stats_spec( - self.b, self.h, self.sq, ragged=self.ragged_stats, output=output + self.graph_b, + self.h, + self.sq, + ragged=self.ragged_stats, + output=output, + ragged_offset_multiplier=self.h if self.ragged_stats and self.use_cu_seqlens else 1, ) return name @@ -230,6 +267,7 @@ def options(self): 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, + use_cu_seqlens=self.use_cu_seqlens, alibi=p["attn_bias_type"] == "alibi", dropout=p["dropout"], deterministic=p.get("deterministic", False), @@ -305,7 +343,15 @@ def _prepare_inputs(attn): 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 attn.use_cu_seqlens: + lengths = cu + else: + lengths = cu[1:] - cu[:-1] + if attn.graph_b > attn.b: + lengths = torch.nn.functional.pad(lengths, (0, attn.graph_b - attn.b)) + attn.add( + name, lengths, (attn.graph_b + int(attn.use_cu_seqlens), 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": @@ -474,21 +520,37 @@ 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) + dim = (attn.graph_b, h, s, d) if n in ("k", "v") and layout.startswith("paged_kv_"): dim = (getattr(config, "num_pages_" + n), h, getattr(config, "page_size_" + n), d) offset = n + "_offset" if fmt == "thd" else "" if offset: - spec(offset, (attn.b + 1, 1, 1, 1), dtype="INT64") + spec( + offset, + (attn.graph_b + 1, 1, 1, 1), + dtype="INT32" if attn.use_cu_seqlens else "INT64", + ) attn.specs[n] = attention_matrix_spec( dim, fmt, dtype, packed=packed, output=output, ragged_offset=offset ) + if offset and attn.use_cu_seqlens: + desc = attn.specs[n] + attn.specs[n] = replace(desc, ragged_offset_multiplier=desc.stride[2]) def stats(n, output=False): if attn.ragged_stats: - spec("stats_offset", (attn.b + 1, 1, 1, 1), dtype="INT64") + spec( + "stats_offset", + (attn.graph_b + 1, 1, 1, 1), + dtype="INT32" if attn.use_cu_seqlens else "INT64", + ) attn.specs[n] = attention_stats_spec( - attn.b, attn.h, attn.sq, ragged=attn.ragged_stats, output=output + attn.graph_b, + attn.h, + attn.sq, + ragged=attn.ragged_stats, + output=output, + ragged_offset_multiplier=attn.h if attn.ragged_stats and attn.use_cu_seqlens else 1, ) def scale(n, matrix_name, fmt, columnwise=False): @@ -510,7 +572,7 @@ def scale(n, matrix_name, fmt, columnwise=False): 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") + spec(n, (attn.graph_b + int(attn.use_cu_seqlens), 1, 1, 1), dtype="INT32") if attn.params["dropout"]: spec("seed", dtype="INT64") spec("offset", dtype="INT64") @@ -719,7 +781,9 @@ def fused_attn_fwd(params): 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_shape = ( + (attn.tq, attn.h, 1) if attn.ragged_stats else (attn.graph_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"]] @@ -744,6 +808,10 @@ def fused_attn_fwd(params): attn.scalar("scale_o", scale_o) attn.scalar("amax_o", amax_o, output=True) attn.execute() + if not attn.ragged_stats and attn.graph_b > attn.b: + aux_ctx_tensors[0] = stats[: attn.b] + if maximum is not None: + maximum = maximum[: attn.b] return out, aux_ctx_tensors, maximum From dd4457a45bcd663f1474ed7a53af189f5a0a661f Mon Sep 17 00:00:00 2001 From: Vladimir Cherepanov Date: Wed, 11 Nov 2026 02:40:57 +0000 Subject: [PATCH 28/49] Remove JAX attention cache size limits Retain supported graphs and serialized bindings without the arbitrary 256-entry eviction limit, matching the existing PyTorch cache policy and avoiding unnecessary rebuilds. Validation: focused JAX attention checks (12 passed, 1 skipped), pre-commit, changed-file pylint and license checks passed. FP8-backward dropout is unsupported by cuDNN on this Blackwell configuration. Full JAX lint reports existing tuple-unpacking warnings in unchanged files; the full license check fails only on pre-existing untracked nccl_ep headers. Signed-off-by: Vladimir Cherepanov --- transformer_engine/jax/cpp_extensions/cudnn_attention.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/transformer_engine/jax/cpp_extensions/cudnn_attention.py b/transformer_engine/jax/cpp_extensions/cudnn_attention.py index 4e38f29440b..9dc4ed1b516 100644 --- a/transformer_engine/jax/cpp_extensions/cudnn_attention.py +++ b/transformer_engine/jax/cpp_extensions/cudnn_attention.py @@ -125,7 +125,7 @@ def _site(spec): return "fp8" if spec.options.fp8 else "f16", "bwd" if spec.options.backward else "fwd" -@lru_cache(maxsize=256) +@lru_cache(maxsize=None) 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. @@ -184,7 +184,7 @@ def get_graph(direction, avals, config): return _build_graph(*_request(direction, avals, config)) -@lru_cache(maxsize=256) +@lru_cache(maxsize=None) def _build_graph(direction, shapes, config, capabilities, nondeterministic, fp8_backward): description = _describe_graph( direction, shapes, config, capabilities, nondeterministic, fp8_backward From c0890e27927802dd67e33edb86a195abb8bf1e2b Mon Sep 17 00:00:00 2001 From: Vladimir Cherepanov Date: Wed, 11 Nov 2026 02:47:04 +0000 Subject: [PATCH 29/49] Address PyTorch cuDNN attention review feedback Check hardware and recipe exclusions before constructing descriptors and reuse the shared MXFP8 support policy. Validate backward independently of training mode, including model-precision fallback and trainable bias or sink parameters with frozen Q/K/V. Keep dropout tied to training mode and order graph build before execute. Expose the existing common sequence-length conversion kernel through a utility C API and private PyTorch binding. Convert both cumulative-length arrays and zero-fill batch padding in one launch, preserving direct cumulative inputs where supported. Validation: rebuilt the editable PyTorch/common installation; 520 focused PyTorch tests and 32 common C++ tests passed on GB200 with CUDA 13.4, cuDNN 9.27, and cuDNN Frontend 1.31. Both framework lint suites, pre-commit, and changed-file license checks passed. Coverage includes independent numerical references, backward selection, nondefault CUDA streams, and CUDA graph replay. Signed-off-by: Vladimir Cherepanov --- docs/release_notes.md | 3 +- tests/cpp/operator/test_splits_to_offsets.cu | 49 +++++ tests/pytorch/attention/test_attention.py | 195 +++++++++++++++++- .../test_attention_backend_selection.py | 94 +++++++++ transformer_engine/common/fused_attn/utils.cu | 33 +++ .../common/include/transformer_engine/utils.h | 16 ++ .../dot_product_attention.py | 11 +- .../attention/dot_product_attention/utils.py | 14 +- .../cpp_extensions/cudnn_fused_attn.py | 154 ++++++++------ transformer_engine/pytorch/csrc/extensions.h | 4 + .../pytorch/csrc/extensions/attention.cpp | 25 +++ .../pytorch/csrc/extensions/pybind.cpp | 3 + 12 files changed, 515 insertions(+), 86 deletions(-) diff --git a/docs/release_notes.md b/docs/release_notes.md index 37f112a10b1..1504eff9a76 100644 --- a/docs/release_notes.md +++ b/docs/release_notes.md @@ -6,7 +6,8 @@ ## Fixed Issues -- +- PyTorch fused-attention backend selection checks backward support when gradients are needed, + including in evaluation mode. ## Breaking Changes in This Release diff --git a/tests/cpp/operator/test_splits_to_offsets.cu b/tests/cpp/operator/test_splits_to_offsets.cu index f0b72d9a6b8..5ff780865c3 100644 --- a/tests/cpp/operator/test_splits_to_offsets.cu +++ b/tests/cpp/operator/test_splits_to_offsets.cu @@ -13,10 +13,59 @@ #include #include +#include #include "../test_common.h" class SplitsToOffsetsTestSuite : public ::testing::TestWithParam> {}; +TEST(SequenceLengthsTest, ConvertAndPad) { + using namespace transformer_engine; + for (const size_t batch : {0, 3, 129, 513}) { + for (const size_t padding : {0, 7}) { + const size_t capacity = batch + padding; + SCOPED_TRACE(::testing::Message() << "batch=" << batch << ", capacity=" << capacity); + test::Tensor cu_q("cu_q", std::vector{batch + 1}, DType::kInt32); + test::Tensor cu_kv("cu_kv", std::vector{batch + 1}, DType::kInt32); + test::Tensor seq_q("seq_q", std::vector{capacity}, DType::kInt32); + test::Tensor seq_kv("seq_kv", std::vector{capacity}, DType::kInt32); + auto *q = cu_q.rowwise_cpu_dptr(); + auto *kv = cu_kv.rowwise_cpu_dptr(); + q[0] = kv[0] = 0; + for (size_t i = 0; i < batch; ++i) { + q[i + 1] = q[i] + i % 7; + kv[i + 1] = kv[i] + i % 11; + } + cu_q.from_cpu(); + cu_kv.from_cpu(); + nvte_cu_seqlens_to_actual_seqlens(cu_q.data(), cu_kv.data(), seq_q.data(), seq_kv.data(), + nullptr); + seq_q.to_cpu(); + seq_kv.to_cpu(); + for (size_t i = 0; i < capacity; ++i) { + EXPECT_EQ(seq_q.rowwise_cpu_dptr()[i], i < batch ? i % 7 : 0); + EXPECT_EQ(seq_kv.rowwise_cpu_dptr()[i], i < batch ? i % 11 : 0); + } + } + } +} + +TEST(SequenceLengthsTest, RejectInvalidMetadata) { + using namespace transformer_engine; + test::Tensor cu("cu", std::vector{3}, DType::kInt32); + test::Tensor lengths("lengths", std::vector{2}, DType::kInt32); + test::Tensor small("small", std::vector{1}, DType::kInt32); + test::Tensor wrong_type("wrong_type", std::vector{2}, DType::kInt64); + test::Tensor wrong_rank("wrong_rank", std::vector{1, 2}, DType::kInt32); + EXPECT_ANY_THROW(nvte_cu_seqlens_to_actual_seqlens(cu.data(), small.data(), lengths.data(), + lengths.data(), nullptr)); + EXPECT_ANY_THROW(nvte_cu_seqlens_to_actual_seqlens(cu.data(), cu.data(), small.data(), + small.data(), nullptr)); + EXPECT_ANY_THROW(nvte_cu_seqlens_to_actual_seqlens(cu.data(), cu.data(), wrong_type.data(), + wrong_type.data(), nullptr)); + EXPECT_ANY_THROW(nvte_cu_seqlens_to_actual_seqlens(cu.data(), cu.data(), wrong_rank.data(), + wrong_rank.data(), nullptr)); +} + TEST_P(SplitsToOffsetsTestSuite, TestSplitsToOffsets) { const size_t num_tensors = std::get<0>(GetParam()); const int64_t logical_last_dim = std::get<1>(GetParam()); diff --git a/tests/pytorch/attention/test_attention.py b/tests/pytorch/attention/test_attention.py index 1975fddc61b..0e3a09c32d7 100644 --- a/tests/pytorch/attention/test_attention.py +++ b/tests/pytorch/attention/test_attention.py @@ -956,7 +956,7 @@ def fail(_self): }, 92000, (10, 0), - "MXFP8 fused attention requires cuDNN 9.21.0", + "MXFP8 attention requires cuDNN 9.21.0", ), ( {"qkv_layout": QKVLayout["paged_kv_bshd_bshd_bshd"]}, @@ -969,15 +969,15 @@ def fail(_self): 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 + from transformer_engine.pytorch.cpp_extensions import cudnn_fused_attn as implementation 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") + def query(*_args): + raise AssertionError("Excluded configuration reached descriptor construction") - monkeypatch.setattr(_AttentionGraph, "build", build) + monkeypatch.setattr(implementation, "_query_graph", query) baseline = FusedAttentionParams( qkv_layout=QKVLayout["bshd_bshd_bshd"], batch_size=2, @@ -993,12 +993,138 @@ def build(_self): assert message.startswith(reason) +@pytest.mark.parametrize( + "capability,version,layout,reason", + ( + ((9, 0), 92700, "bshd_bshd_bshd", "requires SM100"), + ((12, 1), 92700, "bshd_bshd_bshd", "not supported on SM12x"), + ((10, 0), 92300, "bshd_bshd_bshd", "known cuDNN 9.23.0/9.23.1 defects"), + ((10, 0), 92301, "bshd_bshd_bshd", "known cuDNN 9.23.0/9.23.1 defects"), + ((10, 0), 92700, "thd_thd_thd", "does not support THD"), + ), +) +def test_fused_attn_python_mxfp8_exclusions(monkeypatch, capability, version, layout, reason): + """Apply the shared MXFP8 policy before constructing any descriptors.""" + import cudnn + from transformer_engine.pytorch.cpp_extensions import cudnn_fused_attn as implementation + + monkeypatch.setattr(cudnn, "backend_version", lambda: version) + monkeypatch.setattr(torch.cuda, "get_device_capability", lambda *_args: capability) + query = Mock( + side_effect=AssertionError("Excluded MXFP8 config reached descriptor construction") + ) + monkeypatch.setattr(implementation, "_query_graph", query) + config = FusedAttentionParams( + qkv_dtype=DType.kFloat8E4M3, + scaling_mode=tex.NVTEScalingMode.NVTE_MXFP8_1D_SCALING, + qkv_layout=QKVLayout[layout], + attn_mask_type=AttnMaskType["padding"], + num_tokens_q=128, + num_tokens_kv=128, + ) + backend, message = get_fused_attn_backend(config) + assert backend == FusedAttnBackend.No_Backend + assert reason in message + query.assert_not_called() + + +@pytest.mark.parametrize("is_training", (False, True)) +@pytest.mark.parametrize("requires_backward", (False, True, None)) +def test_fused_attn_python_backward_selection(monkeypatch, is_training, requires_backward): + """Backward support and dropout must follow independent flags.""" + from transformer_engine.pytorch.cpp_extensions import cudnn_fused_attn as implementation + + queried = [] + + def entry(attn): + queried.append(attn.backward) + assert attn.params["dropout"] == (0.125 if is_training else 0.0) + if attn.backward: + raise ValueError("unsupported backward") + + monkeypatch.setattr(implementation._AttentionGraph, "entry", entry) + config = FusedAttentionParams( + is_training=is_training, + requires_backward=requires_backward, + dropout=0.125, + 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(config) + backward = is_training if requires_backward is None else requires_backward + assert queried == ([False, True] if backward else [False]) + assert backend == ( + FusedAttnBackend.No_Backend if backward else FusedAttnBackend.F16_arbitrary_seqlen + ) + assert message == ("unsupported backward" if backward else "") + + +@pytest.mark.parametrize( + "batch,capacity", ((0, 0), (0, 32), (3, 3), (3, 32), (129, 256), (513, 1024)) +) +def test_cu_seqlens_to_actual_seqlens(batch, capacity): + """Convert both length arrays on the current stream, including CUDA graph replay.""" + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(stream): + q_lengths = torch.arange(batch, dtype=torch.int32, device="cuda") % 7 + kv_lengths = torch.arange(batch, dtype=torch.int32, device="cuda") % 11 + zero = torch.zeros(1, dtype=torch.int32, device="cuda") + cu_q = torch.cat((zero, q_lengths.cumsum(0, dtype=torch.int32))) + cu_kv = torch.cat((zero, kv_lengths.cumsum(0, dtype=torch.int32))) + expected = [ + torch.nn.functional.pad(t, (0, capacity - batch)) for t in (q_lengths, kv_lengths) + ] + actual = tex.cu_seqlens_to_actual_seqlens(cu_q, cu_kv, capacity) + for output, reference in zip(actual, expected): + torch.testing.assert_close(output, reference) + if capacity: + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph, stream=stream): + captured = tex.cu_seqlens_to_actual_seqlens(cu_q, cu_kv, capacity) + cu_q.mul_(2) + cu_kv.mul_(3) + graph.replay() + for output, reference, factor in zip(captured, expected, (2, 3)): + torch.testing.assert_close(output, reference * factor) + torch.cuda.current_stream().wait_stream(stream) + + +@pytest.mark.parametrize("case", ("dtype", "rank", "stride", "batch", "capacity", "cpu")) +def test_cu_seqlens_to_actual_seqlens_invalid(case): + """Reject invalid metadata before launching the common kernel.""" + q = torch.tensor([0, 2, 5], dtype=torch.int32, device="cuda") + kv, capacity = q, 3 + if case == "dtype": + q = q.long() + elif case == "rank": + q = q.unsqueeze(0) + elif case == "stride": + q = torch.empty(6, dtype=torch.int32, device="cuda")[::2] + elif case == "batch": + kv = kv[:2] + elif case == "capacity": + capacity = 1 + else: + q = q.cpu() + with pytest.raises(RuntimeError): + tex.cu_seqlens_to_actual_seqlens(q, kv, capacity) + + @pytest.mark.parametrize("dtype", (DType.kFloat16, DType.kBFloat16)) @pytest.mark.parametrize("head_dim_v", (64, 128, 192, 256)) @pytest.mark.parametrize("version", (92500, 92600)) -@pytest.mark.parametrize("is_training", (False, True)) +@pytest.mark.parametrize( + "is_training,requires_backward", ((False, None), (True, None), (False, True), (True, False)) +) def test_fused_attn_python_thd_learnable_support( - dtype, head_dim_v, version, is_training, monkeypatch + dtype, head_dim_v, version, is_training, requires_backward, monkeypatch ): """Only generic THD learnable-softmax backward kernels need cuDNN 9.26.""" import cudnn @@ -1009,6 +1135,7 @@ def test_fused_attn_python_thd_learnable_support( monkeypatch.setattr(_AttentionGraph, "entry", lambda _self: None) config = FusedAttentionParams( is_training=is_training, + requires_backward=requires_backward, qkv_dtype=dtype, qkv_layout=QKVLayout["thd_thd_thd"], attn_mask_type=AttnMaskType["padding"], @@ -1024,7 +1151,8 @@ def test_fused_attn_python_thd_learnable_support( num_tokens_kv=256, ) backend, message = get_fused_attn_backend(config) - if is_training and version < 92600 and head_dim_v == 192: + backward = is_training if requires_backward is None else requires_backward + if backward and version < 92600 and head_dim_v == 192: assert backend == FusedAttnBackend.No_Backend assert "THD learnable softmax backward" in message else: @@ -1032,10 +1160,53 @@ def test_fused_attn_python_thd_learnable_support( assert message == "" +@pytest.mark.parametrize("fp8_bwd", ("0", "1")) +def test_fused_attn_python_eval_backward_precision(monkeypatch, fp8_bwd): + """FP8 THD backward exclusions must allow a model-precision backward on SM90.""" + import cudnn + from transformer_engine.pytorch.cpp_extensions import cudnn_fused_attn as implementation + + monkeypatch.setenv("NVTE_FP8_DPA_BWD", fp8_bwd) + monkeypatch.setattr(cudnn, "backend_version", lambda: 92500) + monkeypatch.setattr(torch.cuda, "get_device_capability", lambda *_args: (9, 0)) + queried = [] + monkeypatch.setattr(implementation._AttentionGraph, "entry", lambda attn: queried.append(attn)) + config = FusedAttentionParams( + is_training=False, + requires_backward=True, + qkv_dtype=DType.kFloat8E4M3, + nominal_dtype=DType.kBFloat16, + qkv_layout=QKVLayout["thd_thd_thd"], + attn_mask_type=AttnMaskType["padding"], + 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, + num_tokens_q=256, + num_tokens_kv=256, + ) + backend, message = get_fused_attn_backend(config) + if fp8_bwd == "0": + assert backend == FusedAttnBackend.FP8 + assert message == "" + assert [attn.fp8 for attn in queried] == [True, False] + else: + assert backend == FusedAttnBackend.No_Backend + assert "supports backward on sm100+" in message + assert not queried + + @pytest.mark.parametrize("softmax", ("vanilla", "off-by-one", "learnable")) @pytest.mark.parametrize("version", (92300, 92500, 92600)) -@pytest.mark.parametrize("is_training", (False, True)) -def test_fused_attn_python_fp8_thd_sink_support(softmax, version, is_training, monkeypatch): +@pytest.mark.parametrize( + "is_training,requires_backward", ((False, None), (True, None), (False, True), (True, False)) +) +def test_fused_attn_python_fp8_thd_sink_support( + softmax, version, is_training, requires_backward, monkeypatch +): """Every FP8 THD sink-token backward kernel needs cuDNN 9.26, even at D=64.""" import cudnn from transformer_engine.pytorch.cpp_extensions.cudnn_fused_attn import _AttentionGraph @@ -1045,6 +1216,7 @@ def test_fused_attn_python_fp8_thd_sink_support(softmax, version, is_training, m monkeypatch.setattr(_AttentionGraph, "entry", lambda _self: None) config = FusedAttentionParams( is_training=is_training, + requires_backward=requires_backward, qkv_dtype=DType.kFloat8E4M3, qkv_layout=QKVLayout["thd_thd_thd"], attn_mask_type=AttnMaskType["padding"], @@ -1060,7 +1232,8 @@ def test_fused_attn_python_fp8_thd_sink_support(softmax, version, is_training, m num_tokens_kv=256, ) backend, message = get_fused_attn_backend(config) - if is_training and version < 92600 and softmax != "vanilla": + backward = is_training if requires_backward is None else requires_backward + if backward and version < 92600 and softmax != "vanilla": assert backend == FusedAttnBackend.No_Backend assert "sink token requires cuDNN 9.26" in message else: diff --git a/tests/pytorch/attention/test_attention_backend_selection.py b/tests/pytorch/attention/test_attention_backend_selection.py index 97b1f510f6c..73e79118aba 100644 --- a/tests/pytorch/attention/test_attention_backend_selection.py +++ b/tests/pytorch/attention/test_attention_backend_selection.py @@ -86,6 +86,100 @@ def test_fa4_deterministic_backend_selection( assert available == [expect_flash, False, unfused] +@pytest.mark.parametrize("is_training", (False, True)) +@pytest.mark.parametrize("requires_backward", (False, True)) +def test_fused_backend_backward_requirement(monkeypatch, is_training, requires_backward): + """Pass the autograd requirement to the fused backend independently of training.""" + monkeypatch.setenv("NVTE_FUSED_ATTN", "1") + monkeypatch.setenv("NVTE_FLASH_ATTN", "0") + probe = Mock(return_value=(dpa_utils.FusedAttnBackend.F16_arbitrary_seqlen.value, "")) + monkeypatch.setattr(dpa_utils, "_get_fused_attn_backend", probe) + params = dpa_utils.AttentionParams( + 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, + is_training=is_training, + requires_backward=requires_backward, + ) + _, _, fused, _, _, _ = dpa_utils.get_attention_backend(params) + assert fused + probe.assert_called_once() + assert probe.call_args.kwargs["is_training"] is is_training + assert probe.call_args.kwargs["requires_backward"] is requires_backward + + +@pytest.mark.parametrize("grad_target", ("bias", "softmax_offset")) +def test_backward_requirement_with_frozen_qkv(monkeypatch, grad_target): + """Bias and sink parameters still need backward support when Q/K/V are frozen.""" + selector = Mock(return_value=(False, None, False, None, True, [False, False, True])) + monkeypatch.setattr(dpa_utils, "get_attention_backend", selector) + attention = ( + DotProductAttention( + num_attention_heads=4, + kv_channels=64, + qkv_format="bshd", + attn_mask_type="no_mask", + softmax_type="learnable" if grad_target == "softmax_offset" else "vanilla", + ) + .cuda() + .eval() + ) + qkv = [torch.randn(2, 128, 4, 64, device="cuda", dtype=torch.bfloat16) for _ in range(3)] + kwargs = {} + if grad_target == "bias": + kwargs = { + "core_attention_bias_type": "post_scale_bias", + "core_attention_bias": torch.zeros( + 1, 4, 128, 128, device="cuda", dtype=torch.bfloat16, requires_grad=True + ), + } + output = attention(*qkv, **kwargs) + assert output.requires_grad + assert selector.call_args.args[0].requires_backward + + +@pytest.mark.parametrize("dtype", (torch.float16, torch.bfloat16)) +def test_fused_eval_backward_after_forward_only(monkeypatch, dtype): + """Evaluation with autograd must retain numerical parity after a cached inference call.""" + monkeypatch.setenv("NVTE_FUSED_ATTN", "1") + monkeypatch.setenv("NVTE_FLASH_ATTN", "0") + qkv = [ + torch.randn(2, 128, 4, 64, device="cuda", dtype=dtype, requires_grad=True) for _ in range(3) + ] + attention = ( + DotProductAttention( + num_attention_heads=4, + kv_channels=64, + qkv_format="bshd", + attn_mask_type="no_mask", + attention_dropout=0.125, + ) + .cuda() + .eval() + ) + with torch.no_grad(): + inference = attention(*qkv) + output = attention(*qkv) + assert dpa_module._attention_backends["use_fused_attention"] + torch.testing.assert_close(output, inference) + refs = [tensor.detach().float().requires_grad_() for tensor in qkv] + qr, kr, vr = [tensor.transpose(1, 2) for tensor in refs] + reference = (torch.softmax(qr @ kr.transpose(-1, -2) / 8, dim=-1) @ vr).transpose(1, 2) + reference = reference.reshape_as(output) + gradient = torch.randn_like(output) + output.backward(gradient) + reference.backward(gradient.float()) + rtol, atol = (2e-2, 1e-2) if dtype == torch.bfloat16 else (3e-3, 2e-3) + torch.testing.assert_close(output.float(), reference, rtol=rtol, atol=atol) + for actual, ref in zip(qkv, refs): + torch.testing.assert_close(actual.grad.float(), ref.grad, rtol=rtol, atol=atol) + + @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): diff --git a/transformer_engine/common/fused_attn/utils.cu b/transformer_engine/common/fused_attn/utils.cu index b74dcfe4344..4cee59c87d0 100644 --- a/transformer_engine/common/fused_attn/utils.cu +++ b/transformer_engine/common/fused_attn/utils.cu @@ -11,6 +11,7 @@ #include "../common.h" #include "../util/cuda_runtime.h" #include "transformer_engine/fused_attn.h" +#include "transformer_engine/utils.h" #include "utils.h" namespace transformer_engine { @@ -523,6 +524,38 @@ __global__ void extract_seed_and_offset(int64_t *rng_state_ptr, bool captured, i } // namespace fused_attn } // namespace transformer_engine +void nvte_cu_seqlens_to_actual_seqlens(const NVTETensor cu_seqlens_q, + const NVTETensor cu_seqlens_kv, NVTETensor seqlens_q, + NVTETensor seqlens_kv, cudaStream_t stream) { + NVTE_API_CALL(nvte_cu_seqlens_to_actual_seqlens); + using namespace transformer_engine; + const auto &q_in = *convertNVTETensorCheck(cu_seqlens_q); + const auto &kv_in = *convertNVTETensorCheck(cu_seqlens_kv); + const auto &q_out = *convertNVTETensorCheck(seqlens_q); + const auto &kv_out = *convertNVTETensorCheck(seqlens_kv); + for (const auto *tensor : {&q_in, &kv_in, &q_out, &kv_out}) { + NVTE_CHECK(tensor->data.dtype == DType::kInt32 && tensor->data.shape.size() == 1, + "Sequence lengths must be one-dimensional INT32 tensors."); + } + NVTE_CHECK(q_in.data.shape == kv_in.data.shape && q_in.data.shape[0] >= 1, + "Cumulative sequence lengths must have matching nonempty shapes."); + NVTE_CHECK(q_out.data.shape == kv_out.data.shape && q_out.data.shape[0] >= q_in.data.shape[0] - 1, + "Output sequence lengths must have matching shapes covering the input batch."); + CheckInputTensor(q_in, "cu_seqlens_q"); + CheckInputTensor(kv_in, "cu_seqlens_kv"); + CheckOutputTensor(q_out, "seqlens_q", true); + CheckOutputTensor(kv_out, "seqlens_kv", true); + const auto max_batch = q_out.data.shape[0]; + if (max_batch == 0) return; + constexpr int threads = 128; + const auto blocks = (max_batch + threads - 1) / threads; + fused_attn::cu_seqlens_to_actual_seqlens<<>>( + q_in.data.shape[0] - 1, max_batch, static_cast(q_in.data.dptr), + static_cast(kv_in.data.dptr), static_cast(q_out.data.dptr), + static_cast(kv_out.data.dptr)); + NVTE_CHECK_CUDA(cudaGetLastError()); +} + void nvte_extract_seed_and_offset(int64_t *rng_state_ptr, int captured, int64_t *seed_ptr, uint64_t seed_val, int64_t *offset_ptr, uint64_t offset_val, uint32_t offset_intragraph, cudaStream_t stream) { diff --git a/transformer_engine/common/include/transformer_engine/utils.h b/transformer_engine/common/include/transformer_engine/utils.h index fda49dd5499..2488fadc74c 100644 --- a/transformer_engine/common/include/transformer_engine/utils.h +++ b/transformer_engine/common/include/transformer_engine/utils.h @@ -40,6 +40,22 @@ void nvte_copy_host_to_device_via_kernel(const void *host_ptr, void *device_ptr, void nvte_convert_pointers_to_tensor(const uint64_t *host_ptrs, NVTETensor output, int64_t count, cudaStream_t stream); +/*! \brief Convert cumulative Q and KV sequence lengths to individual lengths. + * + * All tensors are contiguous, one-dimensional INT32 device tensors. The inputs + * each contain batch_size + 1 cumulative lengths. The outputs have the same + * size, at least batch_size, and any extra entries are set to zero. + * + * \param[in] cu_seqlens_q Cumulative Q sequence lengths. + * \param[in] cu_seqlens_kv Cumulative KV sequence lengths. + * \param[out] seqlens_q Individual Q sequence lengths. + * \param[out] seqlens_kv Individual KV sequence lengths. + * \param[in] stream CUDA stream for the operation. + */ +void nvte_cu_seqlens_to_actual_seqlens(const NVTETensor cu_seqlens_q, + const NVTETensor cu_seqlens_kv, NVTETensor seqlens_q, + NVTETensor seqlens_kv, cudaStream_t stream); + #ifdef __cplusplus } // extern "C" #endif diff --git a/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py b/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py index 4e52f27209f..f15d6fcd563 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py @@ -2816,7 +2816,16 @@ def forward( "is_training": self.training, "device": query_layer.device, "requires_backward": torch.is_grad_enabled() - and any(t.requires_grad for t in (query_layer, key_layer, value_layer)), + and any( + t is not None and t.requires_grad + for t in ( + query_layer, + key_layer, + value_layer, + core_attention_bias, + self.softmax_offset, + ) + ), "fp8": self.fp8, "fp8_meta": self.fp8_meta, "inference_params": inference_params, diff --git a/transformer_engine/pytorch/attention/dot_product_attention/utils.py b/transformer_engine/pytorch/attention/dot_product_attention/utils.py index 056bcd6b041..00fe6f056ab 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/utils.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/utils.py @@ -436,6 +436,9 @@ class FusedAttentionParams: bias_seqlen_q: int = 0 bias_seqlen_kv: int = 0 + # None preserves the training-mode default for direct support queries. + requires_backward: Optional[bool] = None + class _NoOpLogger: """ @@ -1652,6 +1655,7 @@ def _is_fa3_supported(num_heads, num_gqa_groups, head_dim_qk, head_dim_v, qkv_dt model_dtype = nominal_dtype if nominal_dtype is not None else torch.bfloat16 base_fused_attn_kwargs = { "is_training": is_training, + "requires_backward": attention_params.requires_backward, "deterministic": deterministic, "cuda_graph": cuda_graph, "return_max_logit": return_max_logit, @@ -1825,7 +1829,7 @@ def _is_fa3_supported(num_heads, num_gqa_groups, head_dim_qk, head_dim_v, qkv_dt fused_attention_backend = None if ( fused_attention_backend == FusedAttnBackend.FP8.value - and is_training + and attention_params.requires_backward and (device_compute_capability < (9, 0) or cudnn_version < (9, 19, 0)) ): logger.debug( @@ -1836,7 +1840,7 @@ def _is_fa3_supported(num_heads, num_gqa_groups, head_dim_qk, head_dim_v, qkv_dt fused_attention_backend = None if ( fused_attention_backend == FusedAttnBackend.F16_arbitrary_seqlen.value - and is_training + and attention_params.requires_backward and (device_compute_capability < (9, 0) or core_attention_bias_requires_grad) ): logger.debug("Disabling FusedAttention for determinism reasons with post_scale_bias") @@ -1848,15 +1852,15 @@ def _is_fa3_supported(num_heads, num_gqa_groups, head_dim_qk, head_dim_v, qkv_dt # was observed to grow super-linearly (B=4 took ~4x the B=2 amount, not 2x) — # revisit if a config uses B>2. SM90_DET_FUSED_THD_BWD_MAX_BHSS = 1 << 30 - is_sm90_f16_thd_training = ( + is_sm90_f16_thd_backward = ( fused_attention_backend == FusedAttnBackend.F16_arbitrary_seqlen.value - and is_training + and attention_params.requires_backward and qkv_format == "thd" and device_compute_capability == (9, 0) ) if ( use_fused_attention - and is_sm90_f16_thd_training + and is_sm90_f16_thd_backward and batch_size * num_heads * max_seqlen_q * max_seqlen_kv >= SM90_DET_FUSED_THD_BWD_MAX_BHSS ): diff --git a/transformer_engine/pytorch/cpp_extensions/cudnn_fused_attn.py b/transformer_engine/pytorch/cpp_extensions/cudnn_fused_attn.py index f6be18829a3..ffcf36186a8 100644 --- a/transformer_engine/pytorch/cpp_extensions/cudnn_fused_attn.py +++ b/transformer_engine/pytorch/cpp_extensions/cudnn_fused_attn.py @@ -4,8 +4,8 @@ """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. +The extension supplies capture-safe Philox state and sequence-length conversion. +Graph construction, plan caching, and execution use the Python frontend. """ from dataclasses import dataclass, replace @@ -31,6 +31,7 @@ bucket, uses_ragged_attention_graph, ) +from transformer_engine.common.cudnn.mxfp8 import support_error as mxfp8_support_error from .cudnn import ( _import_cudnn_frontend, @@ -298,6 +299,15 @@ def entry(self): _event(self, "HIT", key) return entry + def build(self): + """Construct and check the common graph description without compiling it.""" + cudnn = _import_cudnn_frontend() + graph = _build_cudnn_pygraph(getattr(cudnn.data_type, self.specs["q"].dtype), self.device) + built = build_attention_graph(cudnn, graph, self.description()) + _event(self, "CREATE_GRAPH") + check_graph_support(cudnn, graph, lambda name, function: _build_stage(self, name, function)) + return _PlanEntry(graph, built.tensors) + def execute(self): """Compile once, then bind this invocation's runtime tensor addresses.""" entry = self.entry() @@ -317,15 +327,6 @@ def execute(self): ) _event(self, "EXECUTE") - def build(self): - """Construct and check the common graph description without compiling it.""" - cudnn = _import_cudnn_frontend() - graph = _build_cudnn_pygraph(getattr(cudnn.data_type, self.specs["q"].dtype), self.device) - built = build_attention_graph(cudnn, graph, self.description()) - _event(self, "CREATE_GRAPH") - check_graph_support(cudnn, graph, lambda name, function: _build_stage(self, name, function)) - return _PlanEntry(graph, built.tensors) - def _prepare_inputs(attn): p = attn.params @@ -342,16 +343,11 @@ def _prepare_inputs(attn): 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"])): - if attn.use_cu_seqlens: - lengths = cu - else: - lengths = cu[1:] - cu[:-1] - if attn.graph_b > attn.b: - lengths = torch.nn.functional.pad(lengths, (0, attn.graph_b - attn.b)) - attn.add( - name, lengths, (attn.graph_b + int(attn.use_cu_seqlens), 1, 1, 1), (1, 1, 1, 1) - ) + lengths = (p["cu_seqlens_q"], p["cu_seqlens_kv"]) + if not attn.use_cu_seqlens: + lengths = tex.cu_seqlens_to_actual_seqlens(*lengths, attn.graph_b) + for name, tensor in zip(("seq_q", "seq_kv"), lengths): + attn.add(name, tensor, (attn.graph_b + int(attn.use_cu_seqlens), 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": @@ -434,14 +430,10 @@ def _normalize(params): 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 - - if ( - backward - and DType.cast(config.qkv_dtype) in (DType.kFloat8E4M3, DType.kFloat8E5M2) - and not int(os.getenv("NVTE_FP8_DPA_BWD", "1")) +def _backward_config(config): + """Match the precision and layout saved by FusedAttnFunc for backward.""" + if DType.cast(config.qkv_dtype) in (DType.kFloat8E4M3, DType.kFloat8E5M2) and not int( + os.getenv("NVTE_FP8_DPA_BWD", "1") ): # FusedAttnFunc saves model-precision Q/K/V and O for this backward. # The gradient layout records the inputs' layout before FP8 quantization. @@ -467,6 +459,15 @@ def _query_graph(config, backward): else config.qkv_layout ), ) + return config + + +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 + + if backward: + config = _backward_config(config) def name(mapping, value): return next(k for k, v in mapping.items() if v == value) @@ -653,16 +654,31 @@ def scale(n, matrix_name, fmt, columnwise=False): 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"] +def _support_reason(config, version, capability, backward=False): + """Check TE's hardware, recipe, and known cuDNN defect exclusions.""" + from .fused_attn import QKVLayout, AttnMaskType, SoftmaxType + + layout = next(k for k, v in QKVLayout.items() if v == config.qkv_layout) + q_format, kv_format = _formats(layout) + fp8 = DType.cast(config.qkv_dtype) in (DType.kFloat8E4M3, DType.kFloat8E5M2) + mxfp8 = fp8 and config.scaling_mode == tex.NVTEScalingMode.NVTE_MXFP8_1D_SCALING + sm_arch = 10 * capability[0] + capability[1] + if mxfp8: + reason = mxfp8_support_error(sm_arch, version, q_format) + if reason: + return reason + mask = next(k for k, v in AttnMaskType.items() if v == config.attn_mask_type) + softmax = next(k for k, v in SoftmaxType.items() if v == config.softmax_type) + ragged = q_format == "thd" or kv_format == "thd" + ragged_stats = q_format == "thd" and uses_ragged_attention_graph(sm_arch) + padding = "padding" in mask 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): + if (q_format == "thd" and not config.num_tokens_q) or ( + kv_format == "thd" and not config.num_tokens_kv + ): 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 layout.startswith("paged_kv_"): if not padding: return "Paged KV requires PADDING / PADDING_CAUSAL / PADDING_CAUSAL_BOTTOM_RIGHT mask." if not all( @@ -671,24 +687,24 @@ def _support_reason(config, attn, version, capability): 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): + if not fp8: + if q_format == 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 q_format == kv_format + and 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" + backward + and ragged_stats + and softmax == "learnable" and version < 92600 - and attn.dv not in (64, 128, 256) + and config.head_dim_v not in (64, 128, 256) ): return ( "Known cuDNN < 9.26.0 issue with THD learnable softmax backward. Please upgrade" @@ -697,43 +713,39 @@ def _support_reason(config, attn, version, capability): return "" if config.return_max_logit: return "FP8 fused attention does not support return_max_logit=True." - if attn.q_format != attn.kv_format: + if q_format != 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 - + if backward and config.dqkv_layout != tex.NVTE_QKV_Layout.NVTE_QKV_Layout_NOT_SET: 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 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): + if backward 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: + if backward and softmax != "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: + if capability >= (10, 0) and max(config.head_dim_qk, config.head_dim_v) > 128: return "FP8 THD attention supports head dimensions up to 128 on sm100+ only!" - if p["attn_bias_type"] != "no_bias": + if config.bias_type != tex.NVTE_Bias_Type.NVTE_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 + allowed_types = half_types if mxfp8 else half_types + fp8_types tensor_scaling = config.scaling_mode == tex.NVTEScalingMode.NVTE_DELAYED_TENSOR_SCALING if ( - not (tensor_scaling or attn.mxfp8) + not (tensor_scaling or 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) + or (backward 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 "" @@ -751,17 +763,23 @@ def get_fused_attn_backend(config): DType.kFloat8E5M2, ): return FusedAttnBackend.No_Backend, f"Unsupported QKV dtype qkv_dtype={config.qkv_dtype}." + requires_backward = config.requires_backward + if requires_backward is None: + requires_backward = config.is_training 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, "" + configs = [(config, False)] + if requires_backward: + configs.append((_backward_config(config), True)) + version = cudnn.backend_version() + capability = torch.cuda.get_device_capability() + for direction_config, backward in configs: + reason = _support_reason(direction_config, version, capability, backward) + if reason: + return FusedAttnBackend.No_Backend, reason + for direction_config, backward in configs: + _query_graph(direction_config, backward).entry() + fp8 = DType.cast(config.qkv_dtype) in (DType.kFloat8E4M3, DType.kFloat8E5M2) + return FusedAttnBackend.FP8 if fp8 else FusedAttnBackend.F16_arbitrary_seqlen, "" except (cudnn.cudnnGraphNotSupportedError, ValueError) as exc: return FusedAttnBackend.No_Backend, str(exc) diff --git a/transformer_engine/pytorch/csrc/extensions.h b/transformer_engine/pytorch/csrc/extensions.h index 55133bb2932..96fc944c0eb 100644 --- a/transformer_engine/pytorch/csrc/extensions.h +++ b/transformer_engine/pytorch/csrc/extensions.h @@ -97,6 +97,10 @@ std::tuple moe_unpermute_bwd(at::Tensor input_bwd, at::T at::Tensor fused_attn_rng_state(const std::optional rng_gen, size_t elts_per_thread); +std::tuple cu_seqlens_to_actual_seqlens(const at::Tensor &cu_seqlens_q, + const at::Tensor &cu_seqlens_kv, + int64_t max_batch); + 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 547cde9ee39..c52a303af0f 100644 --- a/transformer_engine/pytorch/csrc/extensions/attention.cpp +++ b/transformer_engine/pytorch/csrc/extensions/attention.cpp @@ -10,6 +10,31 @@ namespace transformer_engine::pytorch { +std::tuple cu_seqlens_to_actual_seqlens(const at::Tensor &cu_seqlens_q, + const at::Tensor &cu_seqlens_kv, + int64_t max_batch) { + NVTE_CHECK(cu_seqlens_q.is_cuda() && cu_seqlens_kv.device() == cu_seqlens_q.device(), + "Cumulative sequence lengths must be on the same CUDA device."); + for (const auto &tensor : {cu_seqlens_q, cu_seqlens_kv}) { + NVTE_CHECK(tensor.scalar_type() == at::kInt && tensor.dim() == 1 && tensor.is_contiguous(), + "Cumulative sequence lengths must be contiguous one-dimensional INT32 tensors."); + } + NVTE_CHECK(cu_seqlens_q.numel() >= 1 && cu_seqlens_q.numel() == cu_seqlens_kv.numel(), + "Cumulative sequence lengths must have matching nonempty shapes."); + NVTE_CHECK(max_batch >= cu_seqlens_q.numel() - 1, "Output capacity must cover the input batch."); + at::cuda::CUDAGuard device_guard(cu_seqlens_q.device()); + auto lengths = at::empty({2, max_batch}, cu_seqlens_q.options()); + auto seqlens_q = lengths[0]; + auto seqlens_kv = lengths[1]; + auto te_cu_q = makeTransformerEngineTensor(cu_seqlens_q); + auto te_cu_kv = makeTransformerEngineTensor(cu_seqlens_kv); + auto te_q = makeTransformerEngineTensor(seqlens_q); + auto te_kv = makeTransformerEngineTensor(seqlens_kv); + nvte_cu_seqlens_to_actual_seqlens(te_cu_q.data(), te_cu_kv.data(), te_q.data(), te_kv.data(), + at::cuda::getCurrentCUDAStream()); + return {seqlens_q, seqlens_kv}; +} + // 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, diff --git a/transformer_engine/pytorch/csrc/extensions/pybind.cpp b/transformer_engine/pytorch/csrc/extensions/pybind.cpp index acae02c6a81..c7cc5039787 100644 --- a/transformer_engine/pytorch/csrc/extensions/pybind.cpp +++ b/transformer_engine/pytorch/csrc/extensions/pybind.cpp @@ -590,6 +590,9 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { 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("cu_seqlens_to_actual_seqlens", &transformer_engine::pytorch::cu_seqlens_to_actual_seqlens, + "Convert and pad cumulative Q and KV sequence lengths", py::arg("cu_seqlens_q"), + py::arg("cu_seqlens_kv"), py::arg("max_batch"), py::call_guard()); 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 5213d192da84aa0b94b1285eafd721f1e154be57 Mon Sep 17 00:00:00 2001 From: Vladimir Cherepanov Date: Wed, 11 Nov 2026 03:03:55 +0000 Subject: [PATCH 30/49] Retire the legacy PyTorch fused attention wrapper Merge validation, RNG sizing, and max-logit reduction into the cuDNN Python forward and backward entry points. Move shared attention enums, layout mappings, and FP8 metadata constants to pytorch.constants and update callers. Preserve package-level exports, function signatures, and guarded graph execution. Validation: 479 focused PyTorch tests passed on GB200, including packed FP8 attention, max-logit reduction, CUDA graph replay, and torch.compile backend selection. Fresh import and signature checks, PyTorch/common Python lint, pre-commit, and changed-file license checks passed. Signed-off-by: Vladimir Cherepanov --- tests/pytorch/attention/test_attention.py | 10 +- tests/pytorch/test_identity_quantizer.py | 2 +- tests/pytorch/utils.py | 2 +- .../pytorch/attention/custom_ops.py | 2 +- .../dot_product_attention/backends.py | 8 +- .../dot_product_attention/context_parallel.py | 4 +- .../attention/dot_product_attention/utils.py | 6 +- transformer_engine/pytorch/attention/rope.py | 2 +- transformer_engine/pytorch/constants.py | 136 ++++ .../pytorch/cpp_extensions/__init__.py | 2 +- .../cpp_extensions/cudnn_fused_attn.py | 276 +++++++- .../pytorch/cpp_extensions/fused_attn.py | 658 ------------------ 12 files changed, 419 insertions(+), 689 deletions(-) delete mode 100644 transformer_engine/pytorch/cpp_extensions/fused_attn.py diff --git a/tests/pytorch/attention/test_attention.py b/tests/pytorch/attention/test_attention.py index 0e3a09c32d7..c01379dbc10 100644 --- a/tests/pytorch/attention/test_attention.py +++ b/tests/pytorch/attention/test_attention.py @@ -46,16 +46,18 @@ ) from transformer_engine.pytorch.attention import RotaryPositionEmbedding import transformer_engine.pytorch.cpp_extensions as ext -from transformer_engine.pytorch.cpp_extensions.fused_attn import ( +from transformer_engine.pytorch.constants import ( FusedAttnBackend, AttnMaskType, QKVFormat, QKVLayout, - fused_attn_bwd, - fused_attn_fwd, SoftmaxType, ) -from transformer_engine.pytorch.cpp_extensions.cudnn_fused_attn import get_fused_attn_backend +from transformer_engine.pytorch.cpp_extensions.cudnn_fused_attn import ( + fused_attn_fwd, + fused_attn_bwd, + 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 ( diff --git a/tests/pytorch/test_identity_quantizer.py b/tests/pytorch/test_identity_quantizer.py index 2c120af0fe3..e213aed5b98 100644 --- a/tests/pytorch/test_identity_quantizer.py +++ b/tests/pytorch/test_identity_quantizer.py @@ -629,7 +629,7 @@ def test_hybrid_quantizer_rejects_nested_quantizer_requests(self): def test_fp8_dpa_rejects_identity_quantizer_with_type_error(self): from transformer_engine.pytorch.attention.dot_product_attention import utils as dpa_utils - from transformer_engine.pytorch.cpp_extensions.fused_attn import ( + from transformer_engine.pytorch.constants import ( META_DO, META_DP, META_DQKV, diff --git a/tests/pytorch/utils.py b/tests/pytorch/utils.py index 2a848842b52..b3a4c624363 100644 --- a/tests/pytorch/utils.py +++ b/tests/pytorch/utils.py @@ -33,7 +33,7 @@ AttentionLogging, check_set_window_size, ) -from transformer_engine.pytorch.cpp_extensions.fused_attn import FusedAttnBackend +from transformer_engine.pytorch.constants import FusedAttnBackend from transformer_engine.pytorch.module.base import get_dummy_wgrad diff --git a/transformer_engine/pytorch/attention/custom_ops.py b/transformer_engine/pytorch/attention/custom_ops.py index f364633607e..234b09ee122 100644 --- a/transformer_engine/pytorch/attention/custom_ops.py +++ b/transformer_engine/pytorch/attention/custom_ops.py @@ -7,7 +7,7 @@ import torch import transformer_engine_torch as tex -from transformer_engine.pytorch.cpp_extensions.fused_attn import QKVFormat +from transformer_engine.pytorch.constants import QKVFormat # The ops take the format's value rather than the pybind enum: converting the # enum inside a traced region makes dynamo recurse until it gives up. diff --git a/transformer_engine/pytorch/attention/dot_product_attention/backends.py b/transformer_engine/pytorch/attention/dot_product_attention/backends.py index 0ccc2ed85ba..dc4340faa18 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/backends.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/backends.py @@ -46,14 +46,14 @@ CPLoadBalancingStrategy, QKVLayouts, dist_group_type, -) -from transformer_engine.pytorch.cpp_extensions.fused_attn import ( - fused_attn_fwd, - fused_attn_bwd, FusedAttnBackend, META_O, META_QKV, ) +from transformer_engine.pytorch.cpp_extensions.cudnn_fused_attn import ( + fused_attn_fwd, + fused_attn_bwd, +) from transformer_engine.pytorch.quantization import get_fp8_torch_dtype, FP8GlobalStateManager from transformer_engine.pytorch.distributed import get_distributed_world_size from transformer_engine.pytorch.attention.custom_ops import ( 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 8f7bf952a9d..59dff95ad31 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py @@ -17,10 +17,9 @@ nvtx_range_push, get_device_compute_capability, ) -from transformer_engine.pytorch.cpp_extensions.fused_attn import ( +from transformer_engine.pytorch.cpp_extensions.cudnn_fused_attn import ( fused_attn_fwd, fused_attn_bwd, - FusedAttnBackend, ) from transformer_engine.pytorch.quantization import FP8GlobalStateManager from transformer_engine.pytorch.tensor.float8_tensor import Float8Tensor @@ -31,6 +30,7 @@ from transformer_engine.pytorch.constants import ( CPLoadBalancingStrategy, dist_group_type, + FusedAttnBackend, ) from transformer_engine.pytorch.distributed import ( get_distributed_world_size, diff --git a/transformer_engine/pytorch/attention/dot_product_attention/utils.py b/transformer_engine/pytorch/attention/dot_product_attention/utils.py index 00fe6f056ab..6b9fe808181 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/utils.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/utils.py @@ -23,7 +23,7 @@ import transformer_engine_torch as tex import transformer_engine as te from transformer_engine.pytorch.device import te_device_type -from transformer_engine.pytorch.cpp_extensions.fused_attn import ( +from transformer_engine.pytorch.constants import ( QKVLayout, QKVFormat, AttnBiasType, @@ -36,6 +36,9 @@ META_DO, META_S, META_DP, + TE_DType, + DType, + MXFP8_BLOCK_SCALING_SIZE, ) from transformer_engine.common.cudnn.mxfp8 import ( quantization_layouts as mxfp8_quantization_layouts, @@ -56,7 +59,6 @@ from transformer_engine.pytorch.tensor.storage.mxfp8_tensor_storage import MXFP8TensorStorage from transformer_engine.pytorch.quantization import get_fp8_te_dtype -from transformer_engine.pytorch.constants import TE_DType, DType, MXFP8_BLOCK_SCALING_SIZE from transformer_engine.pytorch.utils import ( diff --git a/transformer_engine/pytorch/attention/rope.py b/transformer_engine/pytorch/attention/rope.py index 83369547f1f..3a21e92a059 100644 --- a/transformer_engine/pytorch/attention/rope.py +++ b/transformer_engine/pytorch/attention/rope.py @@ -9,7 +9,7 @@ import torch import transformer_engine_torch as tex -from transformer_engine.pytorch.cpp_extensions.fused_attn import QKVFormat +from transformer_engine.pytorch.constants import QKVFormat from transformer_engine.pytorch.device import te_device_type diff --git a/transformer_engine/pytorch/constants.py b/transformer_engine/pytorch/constants.py index da4a9038747..b55f4f9deb8 100644 --- a/transformer_engine/pytorch/constants.py +++ b/transformer_engine/pytorch/constants.py @@ -192,3 +192,139 @@ def __hash__(self) -> int: MXFP8_BLOCK_SCALING_SIZE = 32 NVFP4_BLOCK_SCALING_SIZE = 16 + + +QKVFormat = { + None: tex.NVTE_QKV_Format.NVTE_QKV_Format_NOT_SET, + "bshd": tex.NVTE_QKV_Format.NVTE_BSHD, + "sbhd": tex.NVTE_QKV_Format.NVTE_SBHD, + "thd": tex.NVTE_QKV_Format.NVTE_THD, + "sbhd_2bshd": tex.NVTE_QKV_Format.NVTE_SBHD_2BSHD, + "bshd_2sbhd": tex.NVTE_QKV_Format.NVTE_BSHD_2SBHD, + "thd_2bshd": tex.NVTE_QKV_Format.NVTE_THD_2BSHD, + "thd_2sbhd": tex.NVTE_QKV_Format.NVTE_THD_2SBHD, + "bhsd": tex.NVTE_QKV_Format.NVTE_BHSD, +} + +QKVLayout = { + "sb3hd": tex.NVTE_QKV_Layout.NVTE_SB3HD, + "sbh3d": tex.NVTE_QKV_Layout.NVTE_SBH3D, + "sbhd_sb2hd": tex.NVTE_QKV_Layout.NVTE_SBHD_SB2HD, + "sbhd_sbh2d": tex.NVTE_QKV_Layout.NVTE_SBHD_SBH2D, + "sbhd_sbhd_sbhd": tex.NVTE_QKV_Layout.NVTE_SBHD_SBHD_SBHD, + "bs3hd": tex.NVTE_QKV_Layout.NVTE_BS3HD, + "bsh3d": tex.NVTE_QKV_Layout.NVTE_BSH3D, + "bshd_bs2hd": tex.NVTE_QKV_Layout.NVTE_BSHD_BS2HD, + "bshd_bsh2d": tex.NVTE_QKV_Layout.NVTE_BSHD_BSH2D, + "bshd_bshd_bshd": tex.NVTE_QKV_Layout.NVTE_BSHD_BSHD_BSHD, + "t3hd": tex.NVTE_QKV_Layout.NVTE_T3HD, + "th3d": tex.NVTE_QKV_Layout.NVTE_TH3D, + "thd_t2hd": tex.NVTE_QKV_Layout.NVTE_THD_T2HD, + "thd_th2d": tex.NVTE_QKV_Layout.NVTE_THD_TH2D, + "thd_thd_thd": tex.NVTE_QKV_Layout.NVTE_THD_THD_THD, + "sbhd_bshd_bshd": tex.NVTE_QKV_Layout.NVTE_SBHD_BSHD_BSHD, + "bshd_sbhd_sbhd": tex.NVTE_QKV_Layout.NVTE_BSHD_SBHD_SBHD, + "thd_bshd_bshd": tex.NVTE_QKV_Layout.NVTE_THD_BSHD_BSHD, + "thd_sbhd_sbhd": tex.NVTE_QKV_Layout.NVTE_THD_SBHD_SBHD, + "paged_kv_bshd_bshd_bshd": tex.NVTE_QKV_Layout.NVTE_Paged_KV_BSHD_BSHD_BSHD, + "paged_kv_bshd_sbhd_sbhd": tex.NVTE_QKV_Layout.NVTE_Paged_KV_BSHD_SBHD_SBHD, + "paged_kv_sbhd_bshd_bshd": tex.NVTE_QKV_Layout.NVTE_Paged_KV_SBHD_BSHD_BSHD, + "paged_kv_sbhd_sbhd_sbhd": tex.NVTE_QKV_Layout.NVTE_Paged_KV_SBHD_SBHD_SBHD, + "paged_kv_thd_bshd_bshd": tex.NVTE_QKV_Layout.NVTE_Paged_KV_THD_BSHD_BSHD, + "paged_kv_thd_sbhd_sbhd": tex.NVTE_QKV_Layout.NVTE_Paged_KV_THD_SBHD_SBHD, + "bhsd_bhsd_bhsd": tex.NVTE_QKV_Layout.NVTE_BHSD_BHSD_BHSD, +} + +AttnBiasType = { + "no_bias": tex.NVTE_Bias_Type.NVTE_NO_BIAS, + "pre_scale_bias": tex.NVTE_Bias_Type.NVTE_PRE_SCALE_BIAS, + "post_scale_bias": tex.NVTE_Bias_Type.NVTE_POST_SCALE_BIAS, + "alibi": tex.NVTE_Bias_Type.NVTE_ALIBI, +} + +AttnMaskType = { + "no_mask": tex.NVTE_Mask_Type.NVTE_NO_MASK, + "padding": tex.NVTE_Mask_Type.NVTE_PADDING_MASK, + "causal": tex.NVTE_Mask_Type.NVTE_CAUSAL_MASK, + "padding_causal": tex.NVTE_Mask_Type.NVTE_PADDING_CAUSAL_MASK, + "causal_bottom_right": tex.NVTE_Mask_Type.NVTE_CAUSAL_BOTTOM_RIGHT_MASK, + "padding_causal_bottom_right": tex.NVTE_Mask_Type.NVTE_PADDING_CAUSAL_BOTTOM_RIGHT_MASK, +} + +SoftmaxType = { + "vanilla": tex.NVTE_Softmax_Type.NVTE_VANILLA_SOFTMAX, + "off-by-one": tex.NVTE_Softmax_Type.NVTE_OFF_BY_ONE_SOFTMAX, + "learnable": tex.NVTE_Softmax_Type.NVTE_LEARNABLE_SOFTMAX, +} + + +class FusedAttnBackend(enum.IntEnum): + """Fused attention sub-backends. + + This is the canonical fused-attention backend enum for + ``transformer_engine.pytorch``. It mirrors the backend + ``transformer_engine_torch.NVTE_Fused_Attn_Backend`` (pybind11) enum + value-for-value, and instances of the two enums compare equal when they + share the same integer value. Unlike the pybind enum, a plain-python + ``IntEnum`` is traceable by ``torch.compile``: comparisons against a member + constant-fold cleanly. Lookup by name (``FusedAttnBackend["FP8"]``) works + the same way as with the dict this used to be. + + Members do not survive a graph break, though, so ``get_attention_backend`` + returns the sub-backend as a plain int and ``cast`` turns it back into a + member. + """ + + No_Backend = int(tex.NVTE_Fused_Attn_Backend.NVTE_No_Backend) + F16_arbitrary_seqlen = int(tex.NVTE_Fused_Attn_Backend.NVTE_F16_arbitrary_seqlen) + FP8 = int(tex.NVTE_Fused_Attn_Backend.NVTE_FP8) + + @classmethod + def cast( + cls, backend: "Union[FusedAttnBackend, tex.NVTE_Fused_Attn_Backend]" + ) -> "FusedAttnBackend": + """Normalize a backend value to the canonical ``FusedAttnBackend`` member. + + The pybind ``transformer_engine_torch.NVTE_Fused_Attn_Backend`` enum is + accepted as input for backward compatibility and mapped to the matching + ``FusedAttnBackend`` member. + """ + if isinstance(backend, cls): + return backend + return cls(int(backend)) + + def __eq__(self, other: object) -> bool: + # ``FusedAttnBackend`` is an ``IntEnum`` while ``tex.NVTE_Fused_Attn_Backend`` + # is a pybind11 enum. Compare by integer value so the two enums stay + # equivalent regardless of the pybind11 version (the pybind ``__eq__`` + # handles the reverse order). + if isinstance(other, tex.NVTE_Fused_Attn_Backend): + return int(self) == int(other) + return int.__eq__(self, other) + + def __ne__(self, other: object) -> bool: + result = self.__eq__(other) + if result is NotImplemented: + return result + return not result + + def __hash__(self) -> int: + return int.__hash__(self) + + +# Fail fast at import time if a new enumerator is added on the C++ side +# without being mirrored above. +assert {f"NVTE_{m.name}" for m in FusedAttnBackend} == set( + tex.NVTE_Fused_Attn_Backend.__members__ +), ( + "FusedAttnBackend in python is out of sync with" + " transformer_engine_torch.NVTE_Fused_Attn_Backend defined on the C++ side." + " Please make sure TE C++ and python are in sync." +) + +META_QKV = FP8FwdTensorIdx.GEMM1_OUTPUT +META_DQKV = FP8BwdTensorIdx.GRAD_OUTPUT1 +META_O = FP8FwdTensorIdx.GEMM2_INPUT +META_DO = FP8BwdTensorIdx.GRAD_INPUT2 +META_S = FP8FwdTensorIdx.GEMM3_OUTPUT +META_DP = FP8BwdTensorIdx.GRAD_INPUT3 diff --git a/transformer_engine/pytorch/cpp_extensions/__init__.py b/transformer_engine/pytorch/cpp_extensions/__init__.py index bb6e9211326..e8b8ed2d2ad 100644 --- a/transformer_engine/pytorch/cpp_extensions/__init__.py +++ b/transformer_engine/pytorch/cpp_extensions/__init__.py @@ -5,5 +5,5 @@ """Python interface for c++ extensions""" from transformer_engine_torch import * -from .fused_attn import * +from .cudnn_fused_attn import fused_attn_fwd, fused_attn_bwd from .gemm import * diff --git a/transformer_engine/pytorch/cpp_extensions/cudnn_fused_attn.py b/transformer_engine/pytorch/cpp_extensions/cudnn_fused_attn.py index ffcf36186a8..81652f9b728 100644 --- a/transformer_engine/pytorch/cpp_extensions/cudnn_fused_attn.py +++ b/transformer_engine/pytorch/cpp_extensions/cudnn_fused_attn.py @@ -9,8 +9,10 @@ """ from dataclasses import dataclass, replace +import math import os import threading +from typing import List, Optional, Tuple, Union import torch import transformer_engine_torch as tex @@ -40,7 +42,21 @@ ) from ..tensor.float8_tensor import Float8Quantizer from ..tensor.storage.mxfp8_tensor_storage import MXFP8TensorStorage -from ..constants import DType +from ..quantized_tensor import Quantizer +from ..constants import ( + DType, + QKVLayout, + QKVFormat, + AttnMaskType, + AttnBiasType, + SoftmaxType, + FusedAttnBackend, +) + +__all__ = ["fused_attn_fwd", "fused_attn_bwd"] + +_FP8_THREADS_PER_CTA = 128 +_F16_RNG_ELTS_PER_THREAD = 16 _plans = {} _cache_lock = threading.RLock() @@ -464,8 +480,6 @@ def _backward_config(config): 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 - if backward: config = _backward_config(config) @@ -656,8 +670,6 @@ def scale(n, matrix_name, fmt, columnwise=False): def _support_reason(config, version, capability, backward=False): """Check TE's hardware, recipe, and known cuDNN defect exclusions.""" - from .fused_attn import QKVLayout, AttnMaskType, SoftmaxType - layout = next(k for k, v in QKVLayout.items() if v == config.qkv_layout) q_format, kv_format = _formats(layout) fp8 = DType.cast(config.qkv_dtype) in (DType.kFloat8E4M3, DType.kFloat8E5M2) @@ -751,8 +763,6 @@ def _support_reason(config, version, capability, backward=False): 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." @@ -784,9 +794,118 @@ def get_fused_attn_backend(config): 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) +def fused_attn_fwd( + is_training: bool, + max_seqlen_q: int, + max_seqlen_kv: int, + cu_seqlens_q: torch.Tensor, + cu_seqlens_kv: torch.Tensor, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + fake_dtype: torch.dtype, + fused_attention_backend: FusedAttnBackend, + attn_bias: torch.Tensor = None, + cu_seqlens_q_padded: torch.Tensor = None, + cu_seqlens_kv_padded: torch.Tensor = None, + page_table_k: torch.Tensor = None, + page_table_v: torch.Tensor = None, + s_quantizer: Quantizer = None, + o_quantizer: Quantizer = None, + attn_scale: float = None, + dropout: float = 0.0, + fast_zero_fill: bool = True, + qkv_layout: str = "sbh3d", + o_format: str = "sbhd", + qkv_scale_inv_format: str = None, + attn_bias_type: str = "no_bias", + attn_mask_type: str = "padding", + softmax_type: str = "vanilla", + window_size: Tuple[int, int] = (-1, -1), + bottom_right_diagonal: bool = None, + rng_gen: torch.Generator = None, + softmax_offset: torch.Tensor = None, + return_max_logit: bool = False, + cuda_graph: bool = False, +) -> Tuple[Union[torch.Tensor, None], ...]: + """Execute forward, returning output, backward auxiliaries, and optional max logits.""" + if bottom_right_diagonal is None: + bottom_right_diagonal = attn_mask_type in { + "causal_bottom_right", + "padding_causal_bottom_right", + } + + if attn_scale is None: + d = q.size(-1) + attn_scale = 1.0 / math.sqrt(d) + + if attn_bias_type not in ["no_bias", "alibi"]: + if attn_bias is None: + raise ValueError( + f"attn_bias tensor cannot be None when attn_bias_type={attn_bias_type!r}." + ) + if attn_bias.dtype != q.dtype: + raise ValueError( + "attn_bias tensor must have the same dtype as q and kv: " + f"attn_bias.dtype={attn_bias.dtype} but q.dtype={q.dtype}." + ) + + # Accept the pybind enum for backward compatibility. + fused_attention_backend = FusedAttnBackend.cast(fused_attention_backend) + if fused_attention_backend == FusedAttnBackend["No_Backend"]: + raise ValueError( + "Fused attention does not support this input combination:" + f" qkv_layout={qkv_layout!r}, attn_bias_type={attn_bias_type!r}," + f" attn_mask_type={attn_mask_type!r}, q.shape={list(q.shape)}," + f" q.dtype={q.dtype}, backend={fused_attention_backend}." + ) + + if fused_attention_backend == FusedAttnBackend["F16_arbitrary_seqlen"]: + rng_elts_per_thread = _F16_RNG_ELTS_PER_THREAD + elif fused_attention_backend == FusedAttnBackend["FP8"]: + rng_elts_per_thread = ( + max_seqlen_q * max_seqlen_q + _FP8_THREADS_PER_CTA - 1 + ) // _FP8_THREADS_PER_CTA + else: + raise ValueError(f"Unsupported backend {fused_attention_backend}") + + p = _normalize( + { + "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 not p["is_training"]: p["dropout"] = 0.0 with torch.cuda.device(p["q"].device): @@ -830,12 +949,141 @@ def fused_attn_fwd(params): aux_ctx_tensors[0] = stats[: attn.b] if maximum is not None: maximum = maximum[: attn.b] - return out, aux_ctx_tensors, maximum + if return_max_logit: + qkv_format = qkv_layout.replace("3", "").replace("2", "").split("_")[0] + # Reduce over batch and sequence positions, excluding THD storage padding. + amax_dims = (0, 2) if maximum.ndim == 3 else (0, 2, 3) + + if qkv_format == "thd": + if maximum.ndim == 4: + seqlens_q = (cu_seqlens_q[1:] - cu_seqlens_q[:-1]).to(device=maximum.device) + sq_idx = torch.arange(maximum.shape[2], device=maximum.device).view(1, 1, -1, 1) + valid = sq_idx < seqlens_q.view(-1, 1, 1, 1) + maximum = maximum.masked_fill(~valid, float("-inf")) + elif maximum.ndim == 3: + if cu_seqlens_q_padded is not None: + # Exclude padding; CP may pass nonzero padded offsets. + actual_seqlens = (cu_seqlens_q[1:] - cu_seqlens_q[:-1]).to( + device=maximum.device + ) + tq = maximum.shape[0] + starts = cu_seqlens_q_padded[:-1].to(device=maximum.device) + ends = (starts + actual_seqlens).clamp(max=tq) + delta = torch.zeros(tq + 1, dtype=torch.int32, device=maximum.device) + updates = torch.ones_like(starts, dtype=torch.int32) + delta.scatter_add_(0, starts.clamp(max=tq), updates) + delta.scatter_add_(0, ends, -updates) + valid = delta[:-1].cumsum(0) > 0 + maximum = maximum.masked_fill(~valid.view(-1, 1, 1), float("-inf")) + + # Max -> max_logit [h] + max_logit = torch.amax(maximum, dim=amax_dims).to(dtype=out.dtype) + return out, aux_ctx_tensors, max_logit + + return out, aux_ctx_tensors + + +def fused_attn_bwd( + max_seqlen_q: int, + max_seqlen_kv: int, + cu_seqlens_q: torch.Tensor, + cu_seqlens_kv: torch.Tensor, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + o: torch.Tensor, + d_o: torch.Tensor, + fake_dtype: torch.dtype, + aux_ctx_tensors: List[torch.Tensor], + fused_attention_backend: FusedAttnBackend, + cu_seqlens_q_padded: torch.Tensor = None, + cu_seqlens_kv_padded: torch.Tensor = None, + s_quantizer: Quantizer = None, + dp_quantizer: Quantizer = None, + dqkv_quantizer: Quantizer = None, + attn_scale: Optional[float] = None, + dropout: float = 0.0, + fast_zero_fill: bool = True, + qkv_layout: str = "sbh3d", + o_format: str = "sbhd", + do_format: str = "sbhd", + dqkv_layout: str = "sbh3d", + qkv_scale_inv_format: str = None, + do_scale_inv_format: str = None, + attn_bias_type: str = "no_bias", + attn_mask_type: str = "padding", + softmax_type: str = "vanilla", + window_size: Tuple[int, int] = (-1, -1), + bottom_right_diagonal: bool = None, + deterministic: bool = False, + cuda_graph: bool = False, +) -> Tuple[Union[torch.Tensor, None], ...]: + """Execute backward and return Q/K/V, bias, and softmax-offset gradients.""" + if bottom_right_diagonal is None: + bottom_right_diagonal = attn_mask_type in { + "causal_bottom_right", + "padding_causal_bottom_right", + } + + if attn_scale is None: + d = q.size(-1) + attn_scale = 1.0 / math.sqrt(d) + + # Accept the pybind enum for backward compatibility. + fused_attention_backend = FusedAttnBackend.cast(fused_attention_backend) + if fused_attention_backend == FusedAttnBackend["No_Backend"]: + raise ValueError( + "Fused attention backward does not support this input combination:" + f" qkv_layout={qkv_layout!r}, attn_bias_type={attn_bias_type!r}," + f" attn_mask_type={attn_mask_type!r}, q.shape={list(q.shape)}," + f" q.dtype={q.dtype}, backend={fused_attention_backend}." + ) + + if len(aux_ctx_tensors) < 1: + raise ValueError( + "aux_ctx_tensors must contain rng_state as its last element," + f" but got len(aux_ctx_tensors)={len(aux_ctx_tensors)}" + f" for backend={fused_attention_backend}." + ) -def fused_attn_bwd(params): - """Execute backward, including FP8/MXFP8 and shared packed gradient storage.""" - p = _normalize(params) + p = _normalize( + { + "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, + } + ) 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 diff --git a/transformer_engine/pytorch/cpp_extensions/fused_attn.py b/transformer_engine/pytorch/cpp_extensions/fused_attn.py deleted file mode 100644 index 533b8a830a8..00000000000 --- a/transformer_engine/pytorch/cpp_extensions/fused_attn.py +++ /dev/null @@ -1,658 +0,0 @@ -# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# -# See LICENSE for license information. - -"""Python interface for fused attention extensions""" - -import math -from enum import IntEnum -from typing import Tuple, List, Union, Optional -import torch -from transformer_engine_torch import ( - NVTE_QKV_Layout, - NVTE_QKV_Format, - NVTE_Bias_Type, - NVTE_Mask_Type, - NVTE_Softmax_Type, - NVTE_Fused_Attn_Backend, -) -from ..quantized_tensor import Quantizer -from ..constants import FP8BwdTensorIdx, FP8FwdTensorIdx, DType - -__all__ = [ - "fused_attn_fwd", - "fused_attn_bwd", -] - - -TORCH_DType = { - DType.kFloat8E4M3: torch.uint8, - DType.kFloat8E5M2: torch.uint8, - DType.kFloat16: torch.half, - DType.kBFloat16: torch.bfloat16, - DType.kFloat32: torch.float32, - DType.kInt32: torch.int32, -} - -QKVFormat = { - None: NVTE_QKV_Format.NVTE_QKV_Format_NOT_SET, - "bshd": NVTE_QKV_Format.NVTE_BSHD, - "sbhd": NVTE_QKV_Format.NVTE_SBHD, - "thd": NVTE_QKV_Format.NVTE_THD, - "sbhd_2bshd": NVTE_QKV_Format.NVTE_SBHD_2BSHD, - "bshd_2sbhd": NVTE_QKV_Format.NVTE_BSHD_2SBHD, - "thd_2bshd": NVTE_QKV_Format.NVTE_THD_2BSHD, - "thd_2sbhd": NVTE_QKV_Format.NVTE_THD_2SBHD, - "bhsd": NVTE_QKV_Format.NVTE_BHSD, -} - -QKVLayout = { - "sb3hd": NVTE_QKV_Layout.NVTE_SB3HD, - "sbh3d": NVTE_QKV_Layout.NVTE_SBH3D, - "sbhd_sb2hd": NVTE_QKV_Layout.NVTE_SBHD_SB2HD, - "sbhd_sbh2d": NVTE_QKV_Layout.NVTE_SBHD_SBH2D, - "sbhd_sbhd_sbhd": NVTE_QKV_Layout.NVTE_SBHD_SBHD_SBHD, - "bs3hd": NVTE_QKV_Layout.NVTE_BS3HD, - "bsh3d": NVTE_QKV_Layout.NVTE_BSH3D, - "bshd_bs2hd": NVTE_QKV_Layout.NVTE_BSHD_BS2HD, - "bshd_bsh2d": NVTE_QKV_Layout.NVTE_BSHD_BSH2D, - "bshd_bshd_bshd": NVTE_QKV_Layout.NVTE_BSHD_BSHD_BSHD, - "t3hd": NVTE_QKV_Layout.NVTE_T3HD, - "th3d": NVTE_QKV_Layout.NVTE_TH3D, - "thd_t2hd": NVTE_QKV_Layout.NVTE_THD_T2HD, - "thd_th2d": NVTE_QKV_Layout.NVTE_THD_TH2D, - "thd_thd_thd": NVTE_QKV_Layout.NVTE_THD_THD_THD, - "sbhd_bshd_bshd": NVTE_QKV_Layout.NVTE_SBHD_BSHD_BSHD, - "bshd_sbhd_sbhd": NVTE_QKV_Layout.NVTE_BSHD_SBHD_SBHD, - "thd_bshd_bshd": NVTE_QKV_Layout.NVTE_THD_BSHD_BSHD, - "thd_sbhd_sbhd": NVTE_QKV_Layout.NVTE_THD_SBHD_SBHD, - "paged_kv_bshd_bshd_bshd": NVTE_QKV_Layout.NVTE_Paged_KV_BSHD_BSHD_BSHD, - "paged_kv_bshd_sbhd_sbhd": NVTE_QKV_Layout.NVTE_Paged_KV_BSHD_SBHD_SBHD, - "paged_kv_sbhd_bshd_bshd": NVTE_QKV_Layout.NVTE_Paged_KV_SBHD_BSHD_BSHD, - "paged_kv_sbhd_sbhd_sbhd": NVTE_QKV_Layout.NVTE_Paged_KV_SBHD_SBHD_SBHD, - "paged_kv_thd_bshd_bshd": NVTE_QKV_Layout.NVTE_Paged_KV_THD_BSHD_BSHD, - "paged_kv_thd_sbhd_sbhd": NVTE_QKV_Layout.NVTE_Paged_KV_THD_SBHD_SBHD, - "bhsd_bhsd_bhsd": NVTE_QKV_Layout.NVTE_BHSD_BHSD_BHSD, -} - -AttnBiasType = { - "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, -} - -AttnMaskType = { - "no_mask": NVTE_Mask_Type.NVTE_NO_MASK, - "padding": NVTE_Mask_Type.NVTE_PADDING_MASK, - "causal": NVTE_Mask_Type.NVTE_CAUSAL_MASK, - "padding_causal": NVTE_Mask_Type.NVTE_PADDING_CAUSAL_MASK, - "causal_bottom_right": NVTE_Mask_Type.NVTE_CAUSAL_BOTTOM_RIGHT_MASK, - "padding_causal_bottom_right": NVTE_Mask_Type.NVTE_PADDING_CAUSAL_BOTTOM_RIGHT_MASK, -} - -SoftmaxType = { - "vanilla": NVTE_Softmax_Type.NVTE_VANILLA_SOFTMAX, - "off-by-one": NVTE_Softmax_Type.NVTE_OFF_BY_ONE_SOFTMAX, - "learnable": NVTE_Softmax_Type.NVTE_LEARNABLE_SOFTMAX, -} - - -class FusedAttnBackend(IntEnum): - """Fused attention sub-backends. - - This is the canonical fused-attention backend enum for - ``transformer_engine.pytorch``. It mirrors the backend - ``transformer_engine_torch.NVTE_Fused_Attn_Backend`` (pybind11) enum - value-for-value, and instances of the two enums compare equal when they - share the same integer value. Unlike the pybind enum, a plain-python - ``IntEnum`` is traceable by ``torch.compile``: comparisons against a member - constant-fold cleanly. Lookup by name (``FusedAttnBackend["FP8"]``) works - the same way as with the dict this used to be. - - Members do not survive a graph break, though, so ``get_attention_backend`` - returns the sub-backend as a plain int and ``cast`` turns it back into a - member. - """ - - No_Backend = int(NVTE_Fused_Attn_Backend.NVTE_No_Backend) - F16_arbitrary_seqlen = int(NVTE_Fused_Attn_Backend.NVTE_F16_arbitrary_seqlen) - FP8 = int(NVTE_Fused_Attn_Backend.NVTE_FP8) - - @classmethod - def cast( - cls, backend: "Union[FusedAttnBackend, NVTE_Fused_Attn_Backend]" - ) -> "FusedAttnBackend": - """Normalize a backend value to the canonical ``FusedAttnBackend`` member. - - The pybind ``transformer_engine_torch.NVTE_Fused_Attn_Backend`` enum is - accepted as input for backward compatibility and mapped to the matching - ``FusedAttnBackend`` member. - """ - if isinstance(backend, cls): - return backend - return cls(int(backend)) - - def __eq__(self, other: object) -> bool: - # ``FusedAttnBackend`` is an ``IntEnum`` while ``NVTE_Fused_Attn_Backend`` - # is a pybind11 enum. Compare by integer value so the two enums stay - # equivalent regardless of the pybind11 version (the pybind ``__eq__`` - # handles the reverse order). - if isinstance(other, NVTE_Fused_Attn_Backend): - return int(self) == int(other) - return int.__eq__(self, other) - - def __ne__(self, other: object) -> bool: - result = self.__eq__(other) - if result is NotImplemented: - return result - return not result - - def __hash__(self) -> int: - return int.__hash__(self) - - -# Fail fast at import time if a new enumerator is added on the C++ side -# without being mirrored above. -assert {f"NVTE_{m.name}" for m in FusedAttnBackend} == set(NVTE_Fused_Attn_Backend.__members__), ( - "FusedAttnBackend in python is out of sync with" - " transformer_engine_torch.NVTE_Fused_Attn_Backend defined on the C++ side." - " Please make sure TE C++ and python are in sync." -) - -BACKEND_FP8_THREADS_PER_CTA = 128 -BACKEND_F16arb_ELTS_PER_THREADS = 16 - -META_QKV = FP8FwdTensorIdx.GEMM1_OUTPUT -META_DQKV = FP8BwdTensorIdx.GRAD_OUTPUT1 -META_O = FP8FwdTensorIdx.GEMM2_INPUT -META_DO = FP8BwdTensorIdx.GRAD_INPUT2 -META_S = FP8FwdTensorIdx.GEMM3_OUTPUT -META_DP = FP8BwdTensorIdx.GRAD_INPUT3 - - -def fused_attn_fwd( - is_training: bool, - max_seqlen_q: int, - max_seqlen_kv: int, - cu_seqlens_q: torch.Tensor, - cu_seqlens_kv: torch.Tensor, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - fake_dtype: torch.dtype, - fused_attention_backend: FusedAttnBackend, - attn_bias: torch.Tensor = None, - cu_seqlens_q_padded: torch.Tensor = None, - cu_seqlens_kv_padded: torch.Tensor = None, - page_table_k: torch.Tensor = None, - page_table_v: torch.Tensor = None, - s_quantizer: Quantizer = None, - o_quantizer: Quantizer = None, - attn_scale: float = None, - dropout: float = 0.0, - fast_zero_fill: bool = True, - qkv_layout: str = "sbh3d", - o_format: str = "sbhd", - qkv_scale_inv_format: str = None, - attn_bias_type: str = "no_bias", - attn_mask_type: str = "padding", - softmax_type: str = "vanilla", - window_size: Tuple[int, int] = (-1, -1), - bottom_right_diagonal: bool = None, - rng_gen: torch.Generator = None, - softmax_offset: torch.Tensor = None, - return_max_logit: bool = False, - cuda_graph: bool = False, -) -> Tuple[Union[torch.Tensor, None], ...]: - """Fused Attention FWD for separate QKV input. - - Parameters - ---------- - is_training : bool - if True, runs training and produces auxiliary tensors aux_ctx_tensors - for the backward; if False, runs inference and doesn't produce aux_ctx_tensors - max_seqlen_q : int - max sequence length for Q, used for padding; - may be larger than max(seqlens_q), - seqlens_q = cu_seqlens_q[1:] - cu_seqlens_q[:-1] - max_seqlen_kv : int - max sequence length for K and V, used for padding; - may be larger than max(seqlens_kv), - seqlens_kv = cu_seqlens_kv[1:] - cu_seqlens_kv[:-1] - cu_seqlens_q : torch.Tensor - cumulative sequence lengths for Q; shape [batch_size + 1] - cu_seqlens_kv : torch.Tensor - cumulative sequence lengths for K and V; shape [batch_size + 1] - q : torch.Tensor - input tensor Q; shape sbhd, bshd or thd (see `qkv_layout` for details) - k : torch.Tensor - input tensor K; shape sbhd, bshd or thd (see `qkv_layout` for details) - v : torch.Tensor - input tensor V; shape sbhd, bshd or thd (see `qkv_layout` for details) - fake_dtype : DType - data type of Q, K and V - in case of high precision, fake dtype in case of FP8; - in torch.dtype - fused_attention_backend : FusedAttnBackend - please see FusedAttention module for details on supported backends. - attn_bias : torch.Tensor, default = None - input tensor Bias when attn_bias_type is "pre_scale_bias" or "post_scale_bias"; - shape [1, num_heads, max_seqlen_q, max_seqlen_kv], same data type as q, k and v - cu_seqlens_q_padded : torch.Tensor, default = None - cumulative sequence offsets for Q; shape [batch_size + 1] - cu_seqlens_kv_padded : torch.Tensor, default = None - cumulative sequence offsets for KV; shape [batch_size + 1] - page_table_k : torch.Tensor, default = None - page table for K cache; shape [batch_size, max_pages_per_seq_k] - page_table_v : torch.Tensor, default = None - page table for V cache; shape [batch_size, max_pages_per_seq_v] - s_quantizer : Quantizer, default = None - Quantizer object for the intermediate value S. - o_quantizer : Quantizer, default = None - Quantizer object for the output of the attention. - attn_scale : float, default = None - if not None, use attn_scale as the attention scale for Q*K.T BMM; - if None, use 1.0/sqrt(head_dim_qk) as the default - dropout : float, default = 0.0 - dropout probability, 0.0 means no dropout, 1.0 means no output; - ignored if is_training is False - fast_zero_fill : bool, default = True - if True, initializes the output tensor O to zero using the fast filling method; - if False, uses PyTorch's .fill_() method - qkv_layout : str, default = "sbh3d" - layout of Q, K and V; - {"sb3hd", "sbh3d", "sbhd_sb2hd", "sbhd_sbh2d", "sbhd_sbhd_sbhd", - "bs3hd", "bsh3d", "bshd_bs2hd", "bshd_bsh2d", "bshd_bshd_bshd", - "t3hd", "th3d", "thd_t2hd", "thd_th2d", "thd_thd_thd"} - o_format : str, default = "sbhd" - format of O; {"sbhd", "bshd", "thd"} - qkv_scale_inv_format : str, default = None - format of the scale-inverse tensors for QKV; {"sbhd", "bshd", "thd", "bhsd"}; - if None, defaults to the format inferred from qkv_layout. - attn_bias_type : str, default = "no_bias" - type of the bias; {"no_bias", "pre_scale_bias", "post_scale_bias", "alibi"} - attn_mask_type : str, default = "padding" - type of the attention mask; {"padding", "causal", "padding_causal", "no_mask"} - softmax_type : str, default = "vanilla" - type of the attention softmax; {"vanilla", "off-by-one", "learnable"} - window_size : Tuple[int, int], default = (-1, -1) - sliding window size for local attention, where query at position i attends to keys - in [i + seqlen_k - seqlen_q - window_size[0], i + seqlen_k - seqlen_q - + window_size[1]] inclusive. Special cases (-1, -1) and (-1, 0) mean no sliding - window and causal mask specifically. - bottom_right_diagonal: bool, default = None - whether to align sliding window and ALiBi diagonal to the top left (False) or - bottom right (True) corner of the softmax matrix. - rng_gen : torch.Generator, default = None - random number generator; - if None, uses the default CUDA generator from PyTorch; otherwise, uses rng_gen - softmax_offset : torch.Tensor, default = None - softmax offset tensor of shape [1, h_q, 1, 1]. - See softmax_type in DotProductAttention for details. - return_max_logit : bool, default = False - whether to return the maximum attention score - cuda_graph : bool, default = False - whether or not cuda graph capture is enabled. - - Returns - ---------- - o : torch.Tensor - output tensor O, of the attention calculation; same data type as Q, K and V; - same shape as Q - aux_ctx_tensors : List[torch.Tensor] - auxiliary output tensors used for the backward; - if is_training is True, aux_ctx_tensors = [softmax-related tensors, rng_state] - if is_training is False, aux_ctx_tensors = None - - softmax-related tensors: - 1. if fused_attention_backend == FusedAttnBackend["F16_arbitrary_seqlen"] - softmaxStats: torch.Tensor - log(sum(e^(x - max(x)))), where x=Q*K.T - shape [batch_size, num_heads, max_seqlen_q, 1], dtype float32 - Max: torch.Tensor, only when return_max_logit is True - shape [batch_size, num_heads, max_seqlen_q, 1], dtype float32 - 2. if fused_attention_backend == FusedAttnBackend["FP8"] - softmaxStats: torch.Tensor - log(sum(e^(x - max(x)))), where x=Q*K.T - shape [batch_size, num_heads, max_seqlen_q, 1], dtype float32 - rng_state: torch.Tensor - state of the random number generator; - [seed, offset], dtype uint64 - max_logit : if return_max_logit = True, shape [h] and same data type as O; otherwise None - """ - - if bottom_right_diagonal is None: - bottom_right_diagonal = attn_mask_type in { - "causal_bottom_right", - "padding_causal_bottom_right", - } - - if attn_scale is None: - d = q.size(-1) - attn_scale = 1.0 / math.sqrt(d) - - if attn_bias_type not in ["no_bias", "alibi"]: - if attn_bias is None: - raise ValueError( - f"attn_bias tensor cannot be None when attn_bias_type={attn_bias_type!r}." - ) - if attn_bias.dtype != q.dtype: - raise ValueError( - "attn_bias tensor must have the same dtype as q and kv: " - f"attn_bias.dtype={attn_bias.dtype} but q.dtype={q.dtype}." - ) - - # Accept the pybind enum for backward compatibility. - fused_attention_backend = FusedAttnBackend.cast(fused_attention_backend) - if fused_attention_backend == FusedAttnBackend["No_Backend"]: - raise ValueError( - "Fused attention does not support this input combination:" - f" qkv_layout={qkv_layout!r}, attn_bias_type={attn_bias_type!r}," - f" attn_mask_type={attn_mask_type!r}, q.shape={list(q.shape)}," - f" q.dtype={q.dtype}, backend={fused_attention_backend}." - ) - - if fused_attention_backend == FusedAttnBackend["F16_arbitrary_seqlen"]: - rng_elts_per_thread = BACKEND_F16arb_ELTS_PER_THREADS - # FP8 fused attention API from fmha_v2 - elif fused_attention_backend == FusedAttnBackend["FP8"]: - rng_elts_per_thread = ( - max_seqlen_q * max_seqlen_q + BACKEND_FP8_THREADS_PER_CTA - 1 - ) // BACKEND_FP8_THREADS_PER_CTA - else: - raise ValueError(f"Unsupported backend {fused_attention_backend}") - - # 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): 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] - amax_dims = (0, 2) if max_tensor.ndim == 3 else (0, 2, 3) - - if qkv_format == "thd": - if max_tensor.ndim == 4: - # For THD on cuDNN <= 9.6 or THD on sm120, Max tensor can be [b, h, sq, 1] - # with padded sequence positions. Exclude those padded positions when computing max_logit. - seqlens_q = (cu_seqlens_q[1:] - cu_seqlens_q[:-1]).to(device=max_tensor.device) - sq_idx = torch.arange(max_tensor.shape[2], device=max_tensor.device).view( - 1, 1, -1, 1 - ) - valid = sq_idx < seqlens_q.view(-1, 1, 1, 1) - max_tensor = max_tensor.masked_fill(~valid, float("-inf")) - elif max_tensor.ndim == 3: - if cu_seqlens_q_padded is not None: - # Exclude padding; CP may pass nonzero padded offsets. - actual_seqlens = (cu_seqlens_q[1:] - cu_seqlens_q[:-1]).to( - device=max_tensor.device - ) - tq = max_tensor.shape[0] - starts = cu_seqlens_q_padded[:-1].to(device=max_tensor.device) - ends = (starts + actual_seqlens).clamp(max=tq) - delta = torch.zeros(tq + 1, dtype=torch.int32, device=max_tensor.device) - updates = torch.ones_like(starts, dtype=torch.int32) - delta.scatter_add_(0, starts.clamp(max=tq), updates) - delta.scatter_add_(0, ends, -updates) - valid = delta[:-1].cumsum(0) > 0 - 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=out.dtype) - return out, aux_ctx_tensors, max_logit - - # out, aux_ctx_tensors - return out, aux_ctx_tensors - - -def fused_attn_bwd( - max_seqlen_q: int, - max_seqlen_kv: int, - cu_seqlens_q: torch.Tensor, - cu_seqlens_kv: torch.Tensor, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - o: torch.Tensor, - d_o: torch.Tensor, - fake_dtype: torch.dtype, - aux_ctx_tensors: List[torch.Tensor], - fused_attention_backend: FusedAttnBackend, - cu_seqlens_q_padded: torch.Tensor = None, - cu_seqlens_kv_padded: torch.Tensor = None, - s_quantizer: Quantizer = None, - dp_quantizer: Quantizer = None, - dqkv_quantizer: Quantizer = None, - attn_scale: Optional[float] = None, - dropout: float = 0.0, - fast_zero_fill: bool = True, - qkv_layout: str = "sbh3d", - o_format: str = "sbhd", - do_format: str = "sbhd", - dqkv_layout: str = "sbh3d", - qkv_scale_inv_format: str = None, - do_scale_inv_format: str = None, - attn_bias_type: str = "no_bias", - attn_mask_type: str = "padding", - softmax_type: str = "vanilla", - window_size: Tuple[int, int] = (-1, -1), - bottom_right_diagonal: bool = None, - deterministic: bool = False, - cuda_graph: bool = False, -) -> Tuple[Union[torch.Tensor, None], ...]: - """Fused Attention BWD for packed KV input. - - Parameters - ---------- - max_seqlen_q : int - max sequence length for Q, used for padding; may be larger than max(seqlens_q), - seqlens_q = cu_seqlens_q[1:] - cu_seqlens_q[:-1] - max_seqlen_kv : int - max sequence length for K and V, used for padding; - may be larger than max(seqlens_kv), - seqlens_kv = cu_seqlens_kv[1:] - cu_seqlens_kv[:-1] - cu_seqlens_q : torch.Tensor - cumulative sequence lengths for Q; shape [batch_size + 1] - cu_seqlens_kv : torch.Tensor - cumulative sequence lengths for K and V; shape [batch_size + 1] - q : torch.Tensor - input tensor Q; shape sbhd, bshd or thd (see `qkv_layout` for details) - k : torch.Tensor - input tensor K; shape sbhd, bshd or thd (see `qkv_layout` for details) - v : torch.Tensor - input tensor V; shape sbhd, bshd or thd (see `qkv_layout` for details) - o : torch.Tensor - input tensor O (output of forward); same data type as Q, K and V; - same shape as Q - d_o : torch.Tensor - input tensor dO (gradient of O); same data type as Q, K and V; - same shape as Q - fake_dtype : DType - data type of Q, K and V - in case of high precision, fake dtype in case of FP8; - in torch.dtype - aux_ctx_tensors : List[torch.Tensor] - auxiliary output tensors of the forward pass when its is_training is True, - e.g. aux_ctx_tensors = [S, Max, rng_state] - fused_attention_backend : FusedAttnBackend - please see FusedAttention module for details on supported backends. - cu_seqlens_q_padded : torch.Tensor, default = None - cumulative sequence offsets for Q; shape [batch_size + 1] - cu_seqlens_kv_padded : torch.Tensor, default = None - cumulative sequence offsets for KV; shape [batch_size + 1] - s_quantizer : Quantizer, default = None - Quantizer object for the intermediate value S. - dp_quantizer : Quantizer, default = None - Quantizer object for the intermediate value dP. - dqkv_quantizer : Quantizer, default = None - Quantizer object for the output values of the fused_attn_bwd. - attn_scale : float, default = None - if not None, use attn_scale as the attention scale for Q*K.T BMM; - if None, use 1.0/sqrt(head_dim_qk) as the default - dropout : float, default = 0.0 - dropout probability, 0.0 means no dropout, 1.0 means no output; - dropout must be 0.0 if is_training is False - fast_zero_fill : bool, default = True - if True, initializes the output tensor O to zero using the fast filling method; - if False, uses PyTorch's .fill_() method - qkv_layout : str, default = "sbh3d" - layout of Q, K and V; - {"sb3hd", "sbh3d", "sbhd_sb2hd", "sbhd_sbh2d", "sbhd_sbhd_sbhd", - "bs3hd", "bsh3d", "bshd_bs2hd", "bshd_bsh2d", "bshd_bshd_bshd", - "t3hd", "th3d", "thd_t2hd", "thd_th2d", "thd_thd_thd"} - o_format : str, default = "sbhd" - format of O; {"sbhd", "bshd", "thd"} - do_format : str, default = "sbhd" - format of dO; {"sbhd", "bshd", "thd"} - dqkv_layout : str, default = "sbh3d" - layout of dQ, dK and dV; - {"sb3hd", "sbh3d", "sbhd_sb2hd", "sbhd_sbh2d", "sbhd_sbhd_sbhd", - "bs3hd", "bsh3d", "bshd_bs2hd", "bshd_bsh2d", "bshd_bshd_bshd", - "t3hd", "th3d", "thd_t2hd", "thd_th2d", "thd_thd_thd"} - qkv_scale_inv_format : str, default = None - format of the scale-inverse tensors for QKV; {"sbhd", "bshd", "thd", "bhsd"}; - if None, defaults to the format inferred from qkv_layout. - do_scale_inv_format : str, default = None - format of the scale-inverse tensors for dO; {"sbhd", "bshd", "thd", "bhsd"}; - if None, defaults to the format inferred from the output layout. - attn_bias_type : str, default = "no_bias" - type of the bias; {"no_bias", "pre_scale_bias", "post_scale_bias", "alibi"} - attn_mask_type : str, default = "padding" - type of the attention mask; {"padding", "causal", "padding_causal", "no_mask"} - softmax_type : str, default = "vanilla" - type of the attention softmax; {"vanilla", "off-by-one", "learnable"} - window_size : Tuple[int, int], default = (-1, -1) - sliding window size for local attention, where query at position i attends to keys - in [i + seqlen_k - seqlen_q - window_size[0], i + seqlen_k - seqlen_q - + window_size[1]] inclusive. Special cases (-1, -1) and (-1, 0) mean no sliding - window and causal mask specifically. - bottom_right_diagonal: bool, default = None - whether to align sliding window and ALiBi diagonal to the top left (False) or - bottom right (True) corner of the softmax matrix. - deterministic : bool, default = False - whether to execute the backward pass with deterministic behaviours. - cuda_graph : bool, default = False - whether or not cuda graph capture is enabled. - - Returns - ---------- - d_q : torch.Tensor - gradient tensor of Q; same data type and shape as Q - d_k : torch.Tensor - gradient tensor of K; same data type and shape as K - d_v : torch.Tensor - gradient tensor of V; same data type and shape as V - d_bias : torch.Tensor, optional - gradient tensor of Bias when attn_bias_type is "pre_scale_bias" - or "post_scale_bias"; same data type and shape as Bias - d_softmax_offset : torch.Tensor, optional - gradient tensor of softmax offset of shape [1, h_q, 1, 1]. - See softmax_type in DotProductAttention for details. - """ - if bottom_right_diagonal is None: - bottom_right_diagonal = attn_mask_type in { - "causal_bottom_right", - "padding_causal_bottom_right", - } - - if attn_scale is None: - d = q.size(-1) - attn_scale = 1.0 / math.sqrt(d) - - # Accept the pybind enum for backward compatibility. - fused_attention_backend = FusedAttnBackend.cast(fused_attention_backend) - if fused_attention_backend == FusedAttnBackend["No_Backend"]: - raise ValueError( - "Fused attention backward does not support this input combination:" - f" qkv_layout={qkv_layout!r}, attn_bias_type={attn_bias_type!r}," - f" attn_mask_type={attn_mask_type!r}, q.shape={list(q.shape)}," - f" q.dtype={q.dtype}, backend={fused_attention_backend}." - ) - - if len(aux_ctx_tensors) < 1: - raise ValueError( - "aux_ctx_tensors must contain rng_state as its last element," - f" but got len(aux_ctx_tensors)={len(aux_ctx_tensors)}" - f" for backend={fused_attention_backend}." - ) - - 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 From 5c2100250eff6fbdf318d5f6b7a3f024ea0c20ae Mon Sep 17 00:00:00 2001 From: Vladimir Cherepanov Date: Wed, 11 Nov 2026 03:55:07 +0000 Subject: [PATCH 31/49] Address straightforward PyTorch attention review feedback Consolidate attention tests, retire the custom FP8 harness, and retain compact metadata coverage. Use the TE unfused reference and move paged-layout cases into the KV-cache suite. Document C attention API deprecation, clarify max-logit naming, and avoid redundant device guards while retaining safe cross-device stream binding. Signed-off-by: Vladimir Cherepanov --- docs/api/c/fused_attn.rst | 11 + tests/pytorch/attention/test_attention.py | 691 ++---------------- tests/pytorch/attention/test_kv_cache.py | 102 +++ .../include/transformer_engine/fused_attn.h | 99 ++- transformer_engine/pytorch/constants.py | 4 +- .../pytorch/cpp_extensions/cudnn.py | 7 +- .../cpp_extensions/cudnn_fused_attn.py | 36 +- 7 files changed, 277 insertions(+), 673 deletions(-) diff --git a/docs/api/c/fused_attn.rst b/docs/api/c/fused_attn.rst index 00333716b8c..eb3e48efda4 100644 --- a/docs/api/c/fused_attn.rst +++ b/docs/api/c/fused_attn.rst @@ -6,4 +6,15 @@ fused_attn.h ============ +.. deprecated:: 2.22 + + The C fused-attention execution and support-query APIs, including their v2 + variants, configuration and parameter accessors, and C++ parameter wrappers, + will be removed in the next major release. Use + :class:`transformer_engine.pytorch.DotProductAttention` or + :func:`transformer_engine.jax.attention.fused_attn` instead. + + Layout enums and helpers, RNG utilities, KV-cache operations, and + context-parallel tensor utilities are not deprecated by this change. + .. doxygenfile:: fused_attn.h diff --git a/tests/pytorch/attention/test_attention.py b/tests/pytorch/attention/test_attention.py index c01379dbc10..de7195bbf08 100644 --- a/tests/pytorch/attention/test_attention.py +++ b/tests/pytorch/attention/test_attention.py @@ -27,7 +27,6 @@ DotProductAttention, MultiheadAttention, get_device_compute_capability, - Quantizer, is_fp8_available, is_bf16_available, ) @@ -45,7 +44,9 @@ get_thd_padding_mask, ) from transformer_engine.pytorch.attention import RotaryPositionEmbedding -import transformer_engine.pytorch.cpp_extensions as ext +from transformer_engine.pytorch.attention.dot_product_attention.backends import ( + UnfusedDotProductAttention, +) from transformer_engine.pytorch.constants import ( FusedAttnBackend, AttnMaskType, @@ -59,19 +60,13 @@ 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 ( init_method_normal, scaled_init_method_normal, ) from transformer_engine.pytorch.utils import get_cudnn_version -from transformer_engine.pytorch.constants import DType, FP8BwdTensorIdx, FP8FwdTensorIdx +from transformer_engine.pytorch.constants import DType import transformer_engine_torch as tex -from transformer_engine.pytorch.quantized_tensor import ( - Quantizer, - prepare_for_saving, - restore_from_func_ctx, -) _current_file = pathlib.Path(__file__).resolve() sys.path = [str(_current_file.parent.parent)] + sys.path @@ -324,7 +319,7 @@ def reject_backward(attn): @pytest.mark.parametrize("dtype", (torch.float16, torch.bfloat16)) @pytest.mark.parametrize("qkv_format", ("bshd", "sbhd")) def test_fused_attn_inference_ignores_dropout(dtype, qkv_format): - """Configured inference dropout must match zero-dropout execution and FP32 SDPA.""" + """Configured inference dropout must match zero dropout and the unfused backend.""" shape = (2, 128, 4, 64) if qkv_format == "bshd" else (128, 2, 4, 64) q, k, v = [torch.randn(shape, device="cuda", dtype=dtype) for _ in range(3)] cu = torch.arange(3, device="cuda", dtype=torch.int32) * 128 @@ -346,56 +341,34 @@ def test_fused_attn_inference_ignores_dropout(dtype, qkv_format): actual, _ = fused_attn_fwd(dropout=0.25, **kwargs) expected, _ = fused_attn_fwd(dropout=0.0, **kwargs) torch.testing.assert_close(actual, expected, rtol=0, atol=0) - order = (0, 2, 1, 3) if qkv_format == "bshd" else (1, 2, 0, 3) - reference = torch.nn.functional.scaled_dot_product_attention( - *(tensor.float().permute(order) for tensor in (q, k, v)), is_causal=True - ) + unfused = UnfusedDotProductAttention(softmax_scale=64**-0.5).eval() + reference = unfused({}, q.float(), k.float(), v.float(), qkv_layout=kwargs["qkv_layout"]) rtol, atol = (2e-2, 1e-2) if dtype == torch.bfloat16 else (3e-3, 2e-3) - torch.testing.assert_close(actual.float().permute(order), reference, rtol=rtol, atol=atol) + torch.testing.assert_close(actual.float(), reference.reshape(shape), rtol=rtol, atol=atol) -@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.""" - if _deterministic and device_compute_capability < (9, 0): - pytest.skip("TE disables deterministic cuDNN attention training below SM90") - - 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 +def test_cudnn_attention_handle_streams(): + """Cached handles follow the target device's stream without changing the caller's device.""" + import cudnn + from transformer_engine.pytorch.cpp_extensions.cudnn import _get_cudnn_current_stream_handle + + current_device = torch.cuda.current_device() + device = torch.device("cuda", (current_device + 1) % torch.cuda.device_count()) + handle = _get_cudnn_current_stream_handle(cudnn, device) + assert torch.cuda.current_device() == current_device + for stream in (torch.cuda.Stream(device=device), torch.cuda.default_stream(device)): + with torch.cuda.stream(stream): + for caller_device in (device.index, current_device): + with torch.cuda.device(caller_device): + assert _get_cudnn_current_stream_handle(cudnn, device) is handle + assert cudnn.get_stream(handle) == stream.cuda_stream + assert torch.cuda.current_device() == caller_device @pytest.mark.parametrize("layout", ("bs3hd", "bsh3d", "sb3hd", "sbh3d", "t3hd", "th3d")) @pytest.mark.parametrize("backward", (False, True)) @pytest.mark.parametrize("dropout", (0.0, 0.125)) -def test_fused_attn_query_reuses_runtime_graph(layout, backward, dropout, monkeypatch): +def test_fused_attn_query_matches_runtime_graph(layout, backward, dropout): """Packed storage and runtime scalar values must preserve support-query graph identity.""" from transformer_engine.pytorch.cpp_extensions import cudnn_fused_attn as implementation @@ -456,35 +429,11 @@ def test_fused_attn_query_reuses_runtime_graph(layout, backward, dropout, monkey runtime.matrix(name, grad, fmt, heads, runtime.sq, cu, output=True) assert query.description() == runtime.description() - # Isolate the cache to demonstrate that a query never compiles, and that a - # later execution with another scalar value compiles exactly the cached plan. - monkeypatch.setattr(implementation, "_plans", {}) - builds, compiles = [], [] - - def build(_self): - entry = implementation._PlanEntry(object(), {name: name for name in runtime.specs}) - builds.append(entry) - return entry - - def compile_plan(_cudnn, graph, _run): - compiles.append(graph) - return 17 - - monkeypatch.setattr(implementation._AttentionGraph, "build", build) - monkeypatch.setattr(implementation, "compile_graph", compile_plan) - monkeypatch.setattr(implementation, "_execute_cudnn_graph", lambda *_args: None) - queried = query.entry() - assert queried.workspace_size is None and not compiles if fmt == "thd" and not query.use_cu_seqlens: other_query = implementation._query_graph(replace(config, batch_size=batch + 1), backward) assert other_query.description() == query.description() - assert other_query.entry() is queried 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 + assert query.description() == runtime.description() @pytest.mark.parametrize( @@ -555,100 +504,18 @@ def test_fused_attn_ragged_batch_capacity( assert kwargs["implementation"] == cudnn.attention_implementation.UNIFIED -@pytest.mark.parametrize("dtype", (torch.float16, torch.bfloat16)) -@pytest.mark.parametrize("q_format", ("bshd", "sbhd", "thd")) -@pytest.mark.parametrize("kv_format", ("bshd", "sbhd")) -def test_fused_attn_paged_cache_layout(dtype, q_format, kv_format): - """Paged cache dimensions must match support queries and independent FP32 SDPA.""" - from torch.nn.attention import SDPBackend, sdpa_kernel - - batch, seqlen, heads, groups, dim = 2, 4, 4, 2, 64 - num_pages, page_size = 6, 128 - lengths = (203, 137) - q_dense = torch.randn(batch, seqlen, heads, dim, device="cuda", dtype=dtype) - q = q_dense.transpose(0, 1).contiguous() if q_format == "sbhd" else q_dense - if q_format == "thd": - q = q.reshape(batch * seqlen, heads, dim) - pages = [ - torch.randn(num_pages, page_size, groups, dim, device="cuda", dtype=dtype) for _ in range(2) - ] - k, v = [ - tensor.transpose(0, 1).contiguous() if kv_format == "sbhd" else tensor for tensor in pages - ] - # Nonconsecutive, distinct K/V mappings also exercise page-table addressing. - table_ids = ((3, 0), (5, 2)), ((4, 1), (0, 3)) - tables = [torch.tensor(ids, device="cuda", dtype=torch.int32) for ids in table_ids] - cu_q = torch.tensor((0, seqlen, batch * seqlen), device="cuda", dtype=torch.int32) - cu_kv = torch.tensor((0, lengths[0], sum(lengths)), device="cuda", dtype=torch.int32) - layout = f"paged_kv_{q_format}_{kv_format}_{kv_format}" - te_dtype = DType.kFloat16 if dtype == torch.float16 else DType.kBFloat16 - config = FusedAttentionParams( - is_training=False, - qkv_layout=QKVLayout[layout], - o_format=QKVFormat[q_format], - qkv_dtype=te_dtype, - o_dtype=te_dtype, - batch_size=batch, - num_attn_heads=heads, - num_gqa_groups=groups, - head_dim_qk=dim, - head_dim_v=dim, - max_seqlen_q=seqlen, - max_seqlen_kv=2 * page_size, - num_tokens_q=batch * seqlen if q_format == "thd" else 0, - num_pages_k=num_pages, - num_pages_v=num_pages, - page_size_k=page_size, - page_size_v=page_size, - max_pages_per_seq_k=2, - max_pages_per_seq_v=2, - attn_mask_type=AttnMaskType["padding"], - ) - backend, message = get_fused_attn_backend(config) - if backend == FusedAttnBackend.No_Backend: - pytest.skip(f"Paged fused attention is unsupported for {layout}: {message}") - out, _ = fused_attn_fwd( - is_training=False, - max_seqlen_q=seqlen, - max_seqlen_kv=2 * page_size, - cu_seqlens_q=cu_q, - cu_seqlens_kv=cu_kv, - q=q, - k=k, - v=v, - fake_dtype=dtype, - fused_attention_backend=backend, - page_table_k=tables[0], - page_table_v=tables[1], - qkv_layout=layout, - o_format=q_format, - attn_mask_type="padding", - ) - expected = [] - with sdpa_kernel(SDPBackend.MATH): - for index, length in enumerate(lengths): - kr, vr = [ - tensor[list(ids[index])] - .reshape(-1, groups, dim)[:length] - .float() - .transpose(0, 1) - .repeat_interleave(heads // groups, dim=0) - for tensor, ids in zip(pages, table_ids) - ] - expected.append( - torch.nn.functional.scaled_dot_product_attention( - q_dense[index].float().transpose(0, 1), kr, vr - ).transpose(0, 1) - ) - actual = out.transpose(0, 1) if q_format == "sbhd" else out.reshape(q_dense.shape) - rtol, atol = (2e-2, 1e-2) if dtype == torch.bfloat16 else (3e-3, 2e-3) - torch.testing.assert_close(actual.float(), torch.stack(expected), rtol=rtol, atol=atol) - - -@pytest.mark.parametrize("dtype", (DType.kFloat16, DType.kBFloat16)) -@pytest.mark.parametrize("capability", ((10, 0), (12, 0), (12, 1), (12, 9), (13, 0))) -@pytest.mark.parametrize("layout", ("bshd_bshd_bshd", "thd_bshd_bshd", "thd_thd_thd")) -def test_fused_attn_backward_token_bounds(monkeypatch, dtype, capability, layout): +@pytest.mark.parametrize( + "capability,layout", + ( + ((10, 0), "bshd_bshd_bshd"), + ((10, 0), "thd_bshd_bshd"), + ((10, 0), "thd_thd_thd"), + ((12, 0), "thd_thd_thd"), + ((12, 1), "thd_bshd_bshd"), + ((13, 0), "thd_thd_thd"), + ), +) +def test_fused_attn_backward_token_bounds(monkeypatch, 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 @@ -656,10 +523,10 @@ def test_fused_attn_backward_token_bounds(monkeypatch, dtype, capability, layout 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, + qkv_dtype=DType.kBFloat16, + o_dtype=DType.kBFloat16, + do_dtype=DType.kBFloat16, + dqkv_dtype=DType.kBFloat16, batch_size=4, num_attn_heads=4, num_gqa_groups=4, @@ -696,8 +563,16 @@ 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")) +@pytest.mark.parametrize( + "capability,layout", + ( + ((10, 0), "bshd_bshd_bshd"), + ((10, 0), "thd_bshd_bshd"), + ((10, 0), "thd_thd_thd"), + ((12, 0), "thd_thd_thd"), + ((12, 1), "thd_bshd_bshd"), + ), +) 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 @@ -4778,465 +4653,3 @@ def get_dummy_cuda_rng_tracker() -> CudaRNGStatesTracker: if is_training: return out, (inp[0].grad, inp[1].grad, inp[2].grad, d_softmax_offset) return out, (None, None, None, d_softmax_offset) - - -model_configs_fp8 = { - # test: ModelConfig(b, sq, hq, dqk) - "fp8_1": ModelConfig(1, 512, 1, 64), - "fp8_2": ModelConfig(4, 512, 16, 64), - "fp8_3": ModelConfig(1, 2048, 1, 128), - "fp8_4": ModelConfig(2, 2048, 24, 128), - "fp8_5": ModelConfig(1, 512, 1, 64, attn_mask_type="causal"), - "fp8_6": ModelConfig(4, 512, 16, 64, attn_mask_type="causal"), - "fp8_7": ModelConfig(1, 2048, 1, 128, attn_mask_type="causal"), - "fp8_8": ModelConfig(2, 2048, 24, 128, attn_mask_type="causal"), -} -param_types_fp8 = [torch.float16, torch.bfloat16] - - -@pytest.mark.skipif(not fp8_attn_available, reason=reason_for_no_fp8_attn) -@pytest.mark.parametrize("dtype", param_types_fp8) -@pytest.mark.parametrize("model", model_configs_fp8) -def test_custom_mha_fp8_vs_f16(dtype, model, monkeypatch): - """Test FP8 dot product attention implementations based on cuDNN frontend - v0.9 and v1.0+. Each test compares results from a custom implementation of - an FP8 MHA module, i.e. Custom_MHA_FP8(), to results from an F16 MHA - implementation, i.e. transformer_engine.pytorch.attention.MultiHeadAttention. - Both paths take F16 input and output. QKV layout is bs3hd""" - - config = model_configs_fp8[model] - - # The custom implementation always executes FP8 backward. - monkeypatch.setenv("NVTE_FP8_DPA_BWD", "1") - - # Test backend availability - is_training = True - fp8_meta = {} - fp8_recipe = recipe.DelayedScaling( - margin=0, - fp8_format=recipe.Format.HYBRID, - amax_history_len=1, - amax_compute_algo="most_recent", - fp8_dpa=True, - ) - fp8_meta["recipe"] = fp8_recipe - available_backends, _, fused_attn_backends = get_available_attention_backends( - config, - qkv_dtype=torch.float8_e4m3fn, - nominal_dtype=dtype, - qkv_layout="bs3hd", - fp8=True, - fp8_meta=fp8_meta, - is_training=is_training, - deterministic=_deterministic, - ) - flash_attn_supported, fused_attn_supported, unfused_attn_supported = available_backends - if not (fused_attn_backends and unfused_attn_supported): - pytest.skip("Not enough backends to run this test with.") - - fused_attn_fwd_fp8, fused_attn_bwd_fp8 = _run_custom_mha_fp8(dtype, config, "FusedAttention") - unfused_attn_fwd_f16, unfused_attn_bwd_f16 = _run_ref_mha_f16( - dtype, config, "UnfusedDotProductAttention" - ) - - atol = 5e-1 - rtol = 5e-1 - rmse_tol = 0.13 - compare_and_assert( - fused_attn_fwd_fp8, - unfused_attn_fwd_f16, - "fused_attn_fwd_fp8", - "unfused_attn_fwd_f16", - atol, - rtol, - rmse_tol, - True, - ) - compare_and_assert( - fused_attn_bwd_fp8, - unfused_attn_bwd_f16, - "fused_attn_bwd_fp8", - "unfused_attn_bwd_f16", - atol, - rtol, - rmse_tol, - True, - ) - - -def _run_custom_mha_fp8(dtype, config, backend): - """Run Custom_MHA_FP8 with FP8 FusedAttention backend. Both input and output - are in F16. QKV GEMM, DPA, and projection GEMM are calculated in FP8.""" - reset_rng_states() - os.environ["NVTE_FLASH_ATTN"] = "0" - os.environ["NVTE_FUSED_ATTN"] = "0" - os.environ["NVTE_UNFUSED_ATTN"] = "0" - if backend == "FlashAttention": - os.environ["NVTE_FLASH_ATTN"] = "1" - if backend == "FusedAttention": - os.environ["NVTE_FUSED_ATTN"] = "1" - if backend == "UnfusedDotProductAttention": - os.environ["NVTE_UNFUSED_ATTN"] = "1" - _attention_backends["backend_selection_requires_update"] = True - - inp = 0.0001 * torch.randint( - -100, - 100, - (config.batch_size * config.max_seqlen_q, config.num_heads * config.head_dim_qk), - dtype=dtype, - device="cuda", - requires_grad=True, - ) - seqlens = torch.full([config.batch_size], config.max_seqlen_q, dtype=torch.int32, device="cuda") - cu_seqlens = torch.zeros(config.batch_size + 1, device="cuda", dtype=torch.int32) - cu_seqlens[1:] = torch.cumsum(seqlens, dim=0) - - out_grad = 0.01 * torch.randn( - config.batch_size * config.max_seqlen_q, - config.num_heads * config.head_dim_qk, - dtype=dtype, - device="cuda", - ) - torch.save(out_grad, "out_grad.pt") - - fp8_recipe = recipe.DelayedScaling( - margin=0, - fp8_format=recipe.Format.HYBRID, - amax_history_len=1, - amax_compute_algo="most_recent", - fp8_dpa=True, - ) - - mha = Custom_MHA_FP8(config).to(dtype=dtype, device="cuda") - with autocast(enabled=True, recipe=fp8_recipe): - out = mha(inp, cu_seqlens, config.max_seqlen_q) - out.backward(out_grad) - - out = torch.load("out.pt") - dqkv = torch.load("dqkv.pt") - return ( - out.view(config.batch_size, config.max_seqlen_q, -1), - dqkv.view( - config.batch_size, config.max_seqlen_q, 3, config.num_heads, config.head_dim_qk - ).contiguous(), - ) - - -def _run_ref_mha_f16(dtype, config, backend): - """Run reference F16 FusedAttention. Both input and output - are in F16. QKV GEMM, DPA, and projection GEMM are also in F16.""" - - os.environ["NVTE_FLASH_ATTN"] = "0" - os.environ["NVTE_FUSED_ATTN"] = "0" - os.environ["NVTE_UNFUSED_ATTN"] = "0" - if backend == "FlashAttention": - os.environ["NVTE_FLASH_ATTN"] = "1" - if backend == "FusedAttention": - os.environ["NVTE_FUSED_ATTN"] = "1" - if backend == "UnfusedDotProductAttention": - os.environ["NVTE_UNFUSED_ATTN"] = "1" - _attention_backends["backend_selection_requires_update"] = True - - inp = torch.load("qkv.pt").to(device="cuda") - inp.requires_grad = True - seqlens = torch.full([config.batch_size], config.max_seqlen_q, dtype=torch.int32, device="cuda") - cu_seqlens = torch.zeros(config.batch_size + 1, device="cuda", dtype=torch.int32) - cu_seqlens[1:] = torch.cumsum(seqlens, dim=0) - out_grad = ( - torch.load("out_grad.pt").to(device="cuda").view(config.batch_size, config.max_seqlen_q, -1) - ) - - _DUMMY_CUDA_RNG_STATE_TRACKER = CudaRNGStatesTracker() - _DUMMY_CUDA_RNG_STATE_TRACKER.add("model-parallel-rng", seed) - - def get_dummy_cuda_rng_tracker() -> CudaRNGStatesTracker: - """Get cuda rng tracker.""" - return _DUMMY_CUDA_RNG_STATE_TRACKER - - block = DotProductAttention( - config.num_heads, - config.head_dim_qk, - attention_dropout=config.dropout_p, - sequence_parallel=False, - tp_size=1, - get_rng_state_tracker=get_dummy_cuda_rng_tracker, - tp_group=None, - layer_number=1, - attention_type="self", - qkv_format="bshd", - ).to(dtype=dtype, device="cuda") - - q = inp[:, :, 0, :, :] - k = inp[:, :, 1, :, :] - v = inp[:, :, 2, :, :] - out = block(q, k, v, attn_mask_type=config.attn_mask_type) - out.backward(out_grad) - - return out, inp.grad - - -_CUBLASLT_WORKSPACE_SIZE_BYTES = 33_554_432 # 32MiB -_2X_ACC_FPROP = False -_2X_ACC_DGRAD = False -_2X_ACC_WGRAD = False - -META_QKV = FP8FwdTensorIdx.GEMM1_OUTPUT -META_DQKV = FP8BwdTensorIdx.GRAD_OUTPUT1 -META_O = FP8FwdTensorIdx.GEMM2_INPUT -META_DO = FP8BwdTensorIdx.GRAD_INPUT2 -META_S = FP8FwdTensorIdx.GEMM3_OUTPUT -META_DP = FP8BwdTensorIdx.GRAD_INPUT3 - - -class _custom_mha_fp8(torch.autograd.Function): - @staticmethod - def forward( - ctx, - inp: torch.Tensor, - qkv_weight: torch.Tensor, - qkv_bias: torch.Tensor, - cu_seqlens: torch.Tensor, - num_heads: int, - p_dropout: float, - max_s: int, - fast_zero_fill: bool, - fp8_meta: Dict[str, Any], - is_training: bool, - mask_type: str, - quantizers: list[Quantizer], - ) -> torch.Tensor: - qkv_dtype = inp.dtype - - assert inp.dim() == 2 - in_features = qkv_weight.shape[-1] - h = num_heads - d = in_features // h - b = cu_seqlens.numel() - 1 - - input_quantizer = quantizers["scaling_fwd"][FP8FwdTensorIdx.GEMM1_INPUT] - qkv_quantizer = quantizers["scaling_fwd"][FP8FwdTensorIdx.GEMM2_INPUT] - qkv_weight_quantizer = quantizers["scaling_fwd"][FP8FwdTensorIdx.GEMM1_WEIGHT] - o_quantizer = quantizers["scaling_fwd"][FP8FwdTensorIdx.GEMM1_OUTPUT] - dO_quantizer = quantizers["scaling_bwd"][FP8BwdTensorIdx.GRAD_OUTPUT1] - dQKV_quantizer = quantizers["scaling_bwd"][FP8BwdTensorIdx.GRAD_INPUT1] - s_quantizer = quantizers["scaling_bwd"][FP8BwdTensorIdx.GRAD_OUTPUT2] - dP_quantizer = quantizers["scaling_bwd"][FP8BwdTensorIdx.GRAD_OUTPUT3] - - inp_fp8 = input_quantizer(inp) - - qkv_weight_fp8 = qkv_weight_quantizer(qkv_weight) - - qkv, *_ = ext.general_gemm( - qkv_weight_fp8, - inp_fp8, - bias=qkv_bias, - out_dtype=qkv_weight_fp8.dtype, - quantization_params=qkv_quantizer, - use_split_accumulator=_2X_ACC_FPROP, - ) - qkv_layout = "bs3hd" - o_format = "bshd" - qkv = qkv.view(-1, 3, h, d) - qkv_fp16 = qkv.dequantize().view(b, max_s, 3, h, d).contiguous() - torch.save(qkv_fp16, "qkv.pt") - qkv = qkv.view(b, max_s, 3, h, d) # bs3hd - - # FMHA - q_data = qkv._data[:, :, 0, :, :] - k_data = qkv._data[:, :, 1, :, :] - v_data = qkv._data[:, :, 2, :, :] - q = qkv.make_like(tensor=qkv, data=q_data, shape=q_data.shape) - k = qkv.make_like(tensor=qkv, data=k_data, shape=k_data.shape) - v = qkv.make_like(tensor=qkv, data=v_data, shape=v_data.shape) - - out, aux_ctx_tensors = fused_attn_fwd( - is_training, - max_s, - max_s, - cu_seqlens, - cu_seqlens, - q, - k, - v, - qkv_dtype, - FusedAttnBackend["FP8"], - attn_scale=None, - dropout=p_dropout, - fast_zero_fill=fast_zero_fill, - qkv_layout=qkv_layout, - o_format=o_format, - attn_bias_type="no_bias", - attn_mask_type=mask_type, - rng_gen=None, - o_quantizer=o_quantizer, - s_quantizer=s_quantizer, - ) - - tensors_to_save, tensor_objects = prepare_for_saving(q, k, v, inp_fp8, qkv_weight_fp8, out) - - ctx.save_for_backward(*tensors_to_save) - ctx.tensor_objects = tensor_objects - ctx.aux_ctx_tensors = aux_ctx_tensors - ctx.qkv_dtype = qkv_dtype - ctx.fp8_meta = fp8_meta - ctx.cu_seqlens = cu_seqlens - ctx.p_dropout = p_dropout - ctx.max_s = max_s - ctx.fast_zero_fill = fast_zero_fill - ctx.hidden_size = in_features - ctx.num_heads = num_heads - ctx.mask_type = mask_type - ctx.dtype = inp.dtype - ctx.qkv_layout = qkv_layout - ctx.o_format = o_format - - ctx.dQKV_quantizer = dQKV_quantizer - ctx.dO_quantizer = dO_quantizer - ctx.dP_quantizer = dP_quantizer - ctx.S_quantizer = s_quantizer - - out = out.view(-1, in_features) # (bs)(hd) - out_fp16 = out.dequantize() - torch.save(out_fp16, "out.pt") # (bs)(hd) - return out_fp16 - - @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) - - proj_dgrad = ctx.dO_quantizer(grad_output) - - dq, dk, dv, *rest = fused_attn_bwd( - ctx.max_s, - ctx.max_s, - ctx.cu_seqlens, - ctx.cu_seqlens, - q, - k, - v, - out, - proj_dgrad.view_as(out), - ctx.qkv_dtype, - ctx.aux_ctx_tensors, - FusedAttnBackend["FP8"], - None, - None, - ctx.S_quantizer, - ctx.dP_quantizer, - ctx.dQKV_quantizer, - attn_scale=None, - dropout=ctx.p_dropout, - fast_zero_fill=ctx.fast_zero_fill, - qkv_layout=ctx.qkv_layout, - o_format=ctx.o_format, - do_format=ctx.o_format, - dqkv_layout=ctx.qkv_layout, - attn_bias_type="no_bias", - attn_mask_type=ctx.mask_type, - ) - dim = 2 - dqkv = torch.Tensor().to(device=dq._data.device, dtype=dq._data.dtype) - dqkv_shape = list(dq._data.shape) - dqkv_shape.insert(dim, 3) - dqkv_stride = list(dq._data.stride()) - dqkv_stride.insert(dim, int(dqkv_stride[-3] / 3)) - dqkv.set_( - dq._data.untyped_storage(), dq._data.storage_offset(), dqkv_shape, dqkv_stride - ) # bs3hd - - dqkv_c = dqkv.view(-1, 3 * ctx.hidden_size) - dqkv_c = dq.make_like(tensor=dq, data=dqkv_c, shape=dqkv_c.shape) - dqkv_c_fp16 = dqkv_c.dequantize() - torch.save(dqkv_c_fp16, "dqkv.pt") - - qkv_bgrad, dqkv = ext.bgrad_quantize(dqkv_c_fp16, ctx.dQKV_quantizer) - dqkv_c._transpose = None - dqkv_c._create_transpose() - - # QKV DGRAD - qkv_dgrad, *_ = ext.general_gemm( - qkv_weight_fp8, - dqkv_c, - ctx.dtype, - use_split_accumulator=_2X_ACC_DGRAD, - layout="NN", - ) - - # QKV WGRAD - qkv_wgrad, *_ = ext.general_gemm( - inp_fp8, - dqkv, - ctx.dtype, - use_split_accumulator=_2X_ACC_WGRAD, - layout="NT", - ) - - return ( - qkv_dgrad, - qkv_wgrad, - qkv_bgrad, - None, - None, - None, - None, - None, - None, - None, - None, - None, - None, - ) - - -class Custom_MHA_FP8(TransformerEngineBaseModule): - def __init__(self, config, params_dtype: torch.dtype = torch.float32): - super().__init__() - self.p_dropout = config.dropout_p - self.h = config.num_heads - self.hidden_size = config.hidden_size - self.head_dim = config.head_dim_qk - self.fast_zero_fill = True - self.mask_type = config.attn_mask_type - - self.qkv_weight = torch.nn.Parameter( - torch.empty( - self.hidden_size * 3, - self.hidden_size, - device=torch.cuda.current_device(), - dtype=params_dtype, - ) - ) - self.qkv_bias = torch.nn.Parameter( - torch.empty( - self.hidden_size * 3, - device=torch.cuda.current_device(), - dtype=params_dtype, - ) - ) - with torch.no_grad(): - self.qkv_bias.zero_() - self.qkv_weight.fill_(1.0) - - def forward( - self, - inp: torch.Tensor, - cu_seqlens, - max_s, - ) -> torch.Tensor: - with self.prepare_forward_ctx(inp, num_gemms=3) as inp: - out = _custom_mha_fp8.apply( - inp, - self.qkv_weight, - self.qkv_bias, - cu_seqlens, - self.h, - self.p_dropout, - max_s, - self.fast_zero_fill, - self.fp8_meta, - self.training, - self.mask_type, - self.quantizers, - ) - return out diff --git a/tests/pytorch/attention/test_kv_cache.py b/tests/pytorch/attention/test_kv_cache.py index 46cd2eb58cf..03dba8369a1 100644 --- a/tests/pytorch/attention/test_kv_cache.py +++ b/tests/pytorch/attention/test_kv_cache.py @@ -25,8 +25,20 @@ is_bf16_available, ) from transformer_engine.common import recipe +from transformer_engine.pytorch.constants import ( + DType, + FusedAttnBackend, + QKVLayout, + QKVFormat, + AttnMaskType, +) +from transformer_engine.pytorch.cpp_extensions.cudnn_fused_attn import ( + fused_attn_fwd, + get_fused_attn_backend, +) from transformer_engine.pytorch.attention.dot_product_attention.utils import ( FlashAttentionUtils as fa_utils, + FusedAttentionParams, ) from transformer_engine.pytorch.utils import ( init_method_normal, @@ -58,6 +70,96 @@ qkv_formats = ["bshd", "sbhd", "thd"] +@pytest.mark.parametrize("dtype", (torch.float16, torch.bfloat16)) +@pytest.mark.parametrize("q_format", ("bshd", "sbhd", "thd")) +@pytest.mark.parametrize("kv_format", ("bshd", "sbhd")) +def test_fused_attn_paged_cache_layout(dtype, q_format, kv_format): + """Paged cache dimensions must match support queries and independent FP32 SDPA.""" + from torch.nn.attention import SDPBackend, sdpa_kernel + + batch, seqlen, heads, groups, dim = 2, 4, 4, 2, 64 + num_pages, page_size = 6, 128 + lengths = (203, 137) + q_dense = torch.randn(batch, seqlen, heads, dim, device="cuda", dtype=dtype) + q = q_dense.transpose(0, 1).contiguous() if q_format == "sbhd" else q_dense + if q_format == "thd": + q = q.reshape(batch * seqlen, heads, dim) + pages = [ + torch.randn(num_pages, page_size, groups, dim, device="cuda", dtype=dtype) for _ in range(2) + ] + k, v = [ + tensor.transpose(0, 1).contiguous() if kv_format == "sbhd" else tensor for tensor in pages + ] + # Nonconsecutive, distinct K/V mappings also exercise page-table addressing. + table_ids = ((3, 0), (5, 2)), ((4, 1), (0, 3)) + tables = [torch.tensor(ids, device="cuda", dtype=torch.int32) for ids in table_ids] + cu_q = torch.tensor((0, seqlen, batch * seqlen), device="cuda", dtype=torch.int32) + cu_kv = torch.tensor((0, lengths[0], sum(lengths)), device="cuda", dtype=torch.int32) + layout = f"paged_kv_{q_format}_{kv_format}_{kv_format}" + te_dtype = DType.kFloat16 if dtype == torch.float16 else DType.kBFloat16 + config = FusedAttentionParams( + is_training=False, + qkv_layout=QKVLayout[layout], + o_format=QKVFormat[q_format], + qkv_dtype=te_dtype, + o_dtype=te_dtype, + batch_size=batch, + num_attn_heads=heads, + num_gqa_groups=groups, + head_dim_qk=dim, + head_dim_v=dim, + max_seqlen_q=seqlen, + max_seqlen_kv=2 * page_size, + num_tokens_q=batch * seqlen if q_format == "thd" else 0, + num_pages_k=num_pages, + num_pages_v=num_pages, + page_size_k=page_size, + page_size_v=page_size, + max_pages_per_seq_k=2, + max_pages_per_seq_v=2, + attn_mask_type=AttnMaskType["padding"], + ) + backend, message = get_fused_attn_backend(config) + if backend == FusedAttnBackend.No_Backend: + pytest.skip(f"Paged fused attention is unsupported for {layout}: {message}") + out, _ = fused_attn_fwd( + is_training=False, + max_seqlen_q=seqlen, + max_seqlen_kv=2 * page_size, + cu_seqlens_q=cu_q, + cu_seqlens_kv=cu_kv, + q=q, + k=k, + v=v, + fake_dtype=dtype, + fused_attention_backend=backend, + page_table_k=tables[0], + page_table_v=tables[1], + qkv_layout=layout, + o_format=q_format, + attn_mask_type="padding", + ) + expected = [] + with sdpa_kernel(SDPBackend.MATH): + for index, length in enumerate(lengths): + kr, vr = [ + tensor[list(ids[index])] + .reshape(-1, groups, dim)[:length] + .float() + .transpose(0, 1) + .repeat_interleave(heads // groups, dim=0) + for tensor, ids in zip(pages, table_ids) + ] + expected.append( + torch.nn.functional.scaled_dot_product_attention( + q_dense[index].float().transpose(0, 1), kr, vr + ).transpose(0, 1) + ) + actual = out.transpose(0, 1) if q_format == "sbhd" else out.reshape(q_dense.shape) + rtol, atol = (2e-2, 1e-2) if dtype == torch.bfloat16 else (3e-3, 2e-3) + torch.testing.assert_close(actual.float(), torch.stack(expected), rtol=rtol, atol=atol) + + def to_pretty_string(x: torch.Tensor): return "[" + ",".join(["{:>3s}".format(str(i)) for i in x.tolist()]) + "]" diff --git a/transformer_engine/common/include/transformer_engine/fused_attn.h b/transformer_engine/common/include/transformer_engine/fused_attn.h index 7d1ffe13198..57f31f737e0 100644 --- a/transformer_engine/common/include/transformer_engine/fused_attn.h +++ b/transformer_engine/common/include/transformer_engine/fused_attn.h @@ -257,18 +257,34 @@ enum NVTEFusedAttnConfigAttribute { kNVTEFusedAttnConfigNumAttributes }; -/*! \brief Create a fused-attention configuration. */ +/*! \brief Create a fused-attention configuration. + * + * \deprecated The C fused-attention API will be removed in the next major release. + * Use Transformer Engine Python attention APIs instead. + */ NVTEFusedAttnConfig nvte_create_fused_attn_config(void); -/*! \brief Destroy a fused-attention configuration. */ +/*! \brief Destroy a fused-attention configuration. + * + * \deprecated The C fused-attention API will be removed in the next major release. + * Use Transformer Engine Python attention APIs instead. + */ void nvte_destroy_fused_attn_config(NVTEFusedAttnConfig config); -/*! \brief Query an attribute in a fused-attention configuration. */ +/*! \brief Query an attribute in a fused-attention configuration. + * + * \deprecated The C fused-attention API will be removed in the next major release. + * Use Transformer Engine Python attention APIs instead. + */ void nvte_get_fused_attn_config_attribute(NVTEFusedAttnConfig config, NVTEFusedAttnConfigAttribute attr, void *buf, size_t size_in_bytes, size_t *size_written); -/*! \brief Set an attribute in a fused-attention configuration. */ +/*! \brief Set an attribute in a fused-attention configuration. + * + * \deprecated The C fused-attention API will be removed in the next major release. + * Use Transformer Engine Python attention APIs instead. + */ void nvte_set_fused_attn_config_attribute(NVTEFusedAttnConfig config, NVTEFusedAttnConfigAttribute attr, const void *buf, size_t size_in_bytes); @@ -324,18 +340,34 @@ enum NVTEFusedAttnFwdParamsAttribute { kNVTEFusedAttnFwdParamsNumAttributes }; -/*! \brief Create a fused-attention forward-parameter object. */ +/*! \brief Create a fused-attention forward-parameter object. + * + * \deprecated The C fused-attention API will be removed in the next major release. + * Use Transformer Engine Python attention APIs instead. + */ NVTEFusedAttnFwdParams nvte_create_fused_attn_fwd_params(void); -/*! \brief Destroy a fused-attention forward-parameter object. */ +/*! \brief Destroy a fused-attention forward-parameter object. + * + * \deprecated The C fused-attention API will be removed in the next major release. + * Use Transformer Engine Python attention APIs instead. + */ void nvte_destroy_fused_attn_fwd_params(NVTEFusedAttnFwdParams params); -/*! \brief Query an attribute in a fused-attention forward-parameter object. */ +/*! \brief Query an attribute in a fused-attention forward-parameter object. + * + * \deprecated The C fused-attention API will be removed in the next major release. + * Use Transformer Engine Python attention APIs instead. + */ void nvte_get_fused_attn_fwd_params_attribute(NVTEFusedAttnFwdParams params, NVTEFusedAttnFwdParamsAttribute attr, void *buf, size_t size_in_bytes, size_t *size_written); -/*! \brief Set an attribute in a fused-attention forward-parameter object. */ +/*! \brief Set an attribute in a fused-attention forward-parameter object. + * + * \deprecated The C fused-attention API will be removed in the next major release. + * Use Transformer Engine Python attention APIs instead. + */ void nvte_set_fused_attn_fwd_params_attribute(NVTEFusedAttnFwdParams params, NVTEFusedAttnFwdParamsAttribute attr, const void *buf, size_t size_in_bytes); @@ -395,18 +427,34 @@ enum NVTEFusedAttnBwdParamsAttribute { kNVTEFusedAttnBwdParamsNumAttributes }; -/*! \brief Create a fused-attention backward-parameter object. */ +/*! \brief Create a fused-attention backward-parameter object. + * + * \deprecated The C fused-attention API will be removed in the next major release. + * Use Transformer Engine Python attention APIs instead. + */ NVTEFusedAttnBwdParams nvte_create_fused_attn_bwd_params(void); -/*! \brief Destroy a fused-attention backward-parameter object. */ +/*! \brief Destroy a fused-attention backward-parameter object. + * + * \deprecated The C fused-attention API will be removed in the next major release. + * Use Transformer Engine Python attention APIs instead. + */ void nvte_destroy_fused_attn_bwd_params(NVTEFusedAttnBwdParams params); -/*! \brief Query an attribute in a fused-attention backward-parameter object. */ +/*! \brief Query an attribute in a fused-attention backward-parameter object. + * + * \deprecated The C fused-attention API will be removed in the next major release. + * Use Transformer Engine Python attention APIs instead. + */ void nvte_get_fused_attn_bwd_params_attribute(NVTEFusedAttnBwdParams params, NVTEFusedAttnBwdParamsAttribute attr, void *buf, size_t size_in_bytes, size_t *size_written); -/*! \brief Set an attribute in a fused-attention backward-parameter object. */ +/*! \brief Set an attribute in a fused-attention backward-parameter object. + * + * \deprecated The C fused-attention API will be removed in the next major release. + * Use Transformer Engine Python attention APIs instead. + */ void nvte_set_fused_attn_bwd_params_attribute(NVTEFusedAttnBwdParams params, NVTEFusedAttnBwdParamsAttribute attr, const void *buf, size_t size_in_bytes); @@ -423,6 +471,9 @@ void nvte_set_fused_attn_bwd_params_attribute(NVTEFusedAttnBwdParams params, * * \return Fused-attention backend, `NVTE_F16_arbitrary_seqlen` or `NVTE_FP8`, * if the configuration is supported; otherwise, `NVTE_No_Backend`. + * + * \deprecated The C fused-attention API will be removed in the next major release. + * Use Transformer Engine Python attention APIs instead. */ NVTE_Fused_Attn_Backend nvte_get_fused_attn_backend_v2(NVTEFusedAttnConfig cfg, const char **message); @@ -449,7 +500,8 @@ NVTE_Fused_Attn_Backend nvte_get_fused_attn_backend_v2(NVTEFusedAttnConfig cfg, * \param[in] cuda_graph Whether cuda graph capture is enabled or not. * \param[in] deterministic Whether determinism is required or not. * - * \deprecated This function has been deprecated in favor of `nvte_get_fused_attn_backend_v2`. + * \deprecated The C fused-attention API will be removed in the next major release. + * Use Transformer Engine Python attention APIs instead. * * \note `nvte_get_fused_attn_backend` has a narrower input signature compared to * `nvte_get_fused_attn_backend_v2`. For the fields it cannot express, it fills them with @@ -474,6 +526,9 @@ NVTE_Fused_Attn_Backend nvte_get_fused_attn_backend( * destroy it with `nvte_destroy_fused_attn_fwd_params()`. * * \param[in,out] params Fused-attention forward-parameter handle. + * + * \deprecated The C fused-attention API will be removed in the next major release. + * Use Transformer Engine Python attention APIs instead. */ void nvte_fused_attn_fwd_v2(NVTEFusedAttnFwdParams params); @@ -537,7 +592,8 @@ void nvte_fused_attn_fwd_v2(NVTEFusedAttnFwdParams params); * \param[in] workspace Workspace tensor. * \param[in] stream CUDA stream used for this operation. * - * \deprecated This function has been deprecated in favor of `nvte_fused_attn_fwd_v2`. + * \deprecated The C fused-attention API will be removed in the next major release. + * Use Transformer Engine Python attention APIs instead. */ void nvte_fused_attn_fwd(const NVTETensor Q, const NVTETensor K, const NVTETensor V, const NVTETensor Bias, const NVTETensor SoftmaxOffset, NVTETensor S, @@ -562,6 +618,9 @@ void nvte_fused_attn_fwd(const NVTETensor Q, const NVTETensor K, const NVTETenso * destroy it with `nvte_destroy_fused_attn_bwd_params()`. * * \param[in,out] params Fused-attention backward-parameter handle. + * + * \deprecated The C fused-attention API will be removed in the next major release. + * Use Transformer Engine Python attention APIs instead. */ void nvte_fused_attn_bwd_v2(NVTEFusedAttnBwdParams params); @@ -625,7 +684,8 @@ void nvte_fused_attn_bwd_v2(NVTEFusedAttnBwdParams params); * \param[in] workspace Workspace tensor. * \param[in] stream CUDA stream used for this operation. * - * \deprecated This function has been deprecated in favor of `nvte_fused_attn_bwd_v2`. + * \deprecated The C fused-attention API will be removed in the next major release. + * Use Transformer Engine Python attention APIs instead. */ void nvte_fused_attn_bwd(const NVTETensor Q, const NVTETensor K, const NVTETensor V, const NVTETensor O, const NVTETensor dO, const NVTETensor S, NVTETensor dP, @@ -997,6 +1057,9 @@ class AttentionShape { * It owns an opaque `NVTEFusedAttnConfig` handle created by * `nvte_create_fused_attn_config()`, and provides a convenient, * chainable interface for setting every field in `FusedAttnConfig`. + * + * \deprecated The C fused-attention API will be removed in the next major release. + * Use Transformer Engine Python attention APIs instead. */ class FusedAttnConfigWrapper { public: @@ -1172,6 +1235,9 @@ class FusedAttnConfigWrapper { * It owns an opaque `NVTEFusedAttnFwdParams` handle created by * `nvte_create_fused_attn_fwd_params()`, and provides a convenient, * chainable interface for setting every field in `FusedAttnFwdParams`. + * + * \deprecated The C fused-attention API will be removed in the next major release. + * Use Transformer Engine Python attention APIs instead. */ class FusedAttnFwdParamsWrapper { public: @@ -1320,6 +1386,9 @@ class FusedAttnFwdParamsWrapper { * It owns an opaque `NVTEFusedAttnBwdParams` handle created by * `nvte_create_fused_attn_bwd_params()`, and provides a convenient, * chainable interface for setting every field in `FusedAttnBwdParams`. + * + * \deprecated The C fused-attention API will be removed in the next major release. + * Use Transformer Engine Python attention APIs instead. */ class FusedAttnBwdParamsWrapper { public: diff --git a/transformer_engine/pytorch/constants.py b/transformer_engine/pytorch/constants.py index b55f4f9deb8..af619ddf8f9 100644 --- a/transformer_engine/pytorch/constants.py +++ b/transformer_engine/pytorch/constants.py @@ -262,7 +262,9 @@ class FusedAttnBackend(enum.IntEnum): """Fused attention sub-backends. This is the canonical fused-attention backend enum for - ``transformer_engine.pytorch``. It mirrors the backend + ``transformer_engine.pytorch``. The F16 and FP8 labels describe the precision + of the Python cuDNN attention path and do not select the deprecated C API. + It mirrors the backend ``transformer_engine_torch.NVTE_Fused_Attn_Backend`` (pybind11) enum value-for-value, and instances of the two enums compare equal when they share the same integer value. Unlike the pybind enum, a plain-python diff --git a/transformer_engine/pytorch/cpp_extensions/cudnn.py b/transformer_engine/pytorch/cpp_extensions/cudnn.py index 994fe43f354..911b0fcbd85 100644 --- a/transformer_engine/pytorch/cpp_extensions/cudnn.py +++ b/transformer_engine/pytorch/cpp_extensions/cudnn.py @@ -5,6 +5,7 @@ """cuDNN Python frontend graph lifecycle shared by attention implementations.""" import threading +from contextlib import nullcontext import torch @@ -27,7 +28,11 @@ def _get_cudnn_current_stream_handle(cudnn, device): device = torch.device(te_device_type(), torch.cuda.current_device()) if not hasattr(_handles, "by_device"): _handles.by_device = {} - with torch.cuda.device(device): + # cuDNN also needs the handle's device when rebinding its stream. + device_context = ( + torch.cuda.device(device) if device.index != torch.cuda.current_device() else nullcontext() + ) + with device_context: handle = _handles.by_device.get(device) if handle is None: handle = cudnn.create_handle() diff --git a/transformer_engine/pytorch/cpp_extensions/cudnn_fused_attn.py b/transformer_engine/pytorch/cpp_extensions/cudnn_fused_attn.py index 81652f9b728..403d31dfbad 100644 --- a/transformer_engine/pytorch/cpp_extensions/cudnn_fused_attn.py +++ b/transformer_engine/pytorch/cpp_extensions/cudnn_fused_attn.py @@ -924,10 +924,10 @@ def fused_attn_fwd( 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 + max_s_tensor = None if p["return_max_logit"]: - maximum = torch.empty_like(stats) - attn.stats("max", maximum, output=True) + max_s_tensor = torch.empty_like(stats) + attn.stats("max", max_s_tensor, output=True) if p["attn_bias_type"] not in ("no_bias", "alibi"): aux_ctx_tensors.append(p["attn_bias"]) if p["softmax_type"] != "vanilla": @@ -947,38 +947,40 @@ def fused_attn_fwd( attn.execute() if not attn.ragged_stats and attn.graph_b > attn.b: aux_ctx_tensors[0] = stats[: attn.b] - if maximum is not None: - maximum = maximum[: attn.b] + if max_s_tensor is not None: + max_s_tensor = max_s_tensor[: attn.b] if return_max_logit: qkv_format = qkv_layout.replace("3", "").replace("2", "").split("_")[0] # Reduce over batch and sequence positions, excluding THD storage padding. - amax_dims = (0, 2) if maximum.ndim == 3 else (0, 2, 3) + amax_dims = (0, 2) if max_s_tensor.ndim == 3 else (0, 2, 3) if qkv_format == "thd": - if maximum.ndim == 4: - seqlens_q = (cu_seqlens_q[1:] - cu_seqlens_q[:-1]).to(device=maximum.device) - sq_idx = torch.arange(maximum.shape[2], device=maximum.device).view(1, 1, -1, 1) + if max_s_tensor.ndim == 4: + seqlens_q = (cu_seqlens_q[1:] - cu_seqlens_q[:-1]).to(device=max_s_tensor.device) + sq_idx = torch.arange(max_s_tensor.shape[2], device=max_s_tensor.device).view( + 1, 1, -1, 1 + ) valid = sq_idx < seqlens_q.view(-1, 1, 1, 1) - maximum = maximum.masked_fill(~valid, float("-inf")) - elif maximum.ndim == 3: + max_s_tensor = max_s_tensor.masked_fill(~valid, float("-inf")) + elif max_s_tensor.ndim == 3: if cu_seqlens_q_padded is not None: # Exclude padding; CP may pass nonzero padded offsets. actual_seqlens = (cu_seqlens_q[1:] - cu_seqlens_q[:-1]).to( - device=maximum.device + device=max_s_tensor.device ) - tq = maximum.shape[0] - starts = cu_seqlens_q_padded[:-1].to(device=maximum.device) + tq = max_s_tensor.shape[0] + starts = cu_seqlens_q_padded[:-1].to(device=max_s_tensor.device) ends = (starts + actual_seqlens).clamp(max=tq) - delta = torch.zeros(tq + 1, dtype=torch.int32, device=maximum.device) + delta = torch.zeros(tq + 1, dtype=torch.int32, device=max_s_tensor.device) updates = torch.ones_like(starts, dtype=torch.int32) delta.scatter_add_(0, starts.clamp(max=tq), updates) delta.scatter_add_(0, ends, -updates) valid = delta[:-1].cumsum(0) > 0 - maximum = maximum.masked_fill(~valid.view(-1, 1, 1), float("-inf")) + max_s_tensor = max_s_tensor.masked_fill(~valid.view(-1, 1, 1), float("-inf")) # Max -> max_logit [h] - max_logit = torch.amax(maximum, dim=amax_dims).to(dtype=out.dtype) + max_logit = torch.amax(max_s_tensor, dim=amax_dims).to(dtype=out.dtype) return out, aux_ctx_tensors, max_logit return out, aux_ctx_tensors From 5344e71a567353a93f8fda2775dc7ac9dbfcd54e Mon Sep 17 00:00:00 2001 From: Vladimir Cherepanov Date: Wed, 11 Nov 2026 04:11:59 +0000 Subject: [PATCH 32/49] Remove redundant attention metadata tests Remove the backward token-bound and fake-statistics shape tests requested in review, along with the unused SimpleNamespace import. Signed-off-by: Vladimir Cherepanov --- tests/pytorch/attention/test_attention.py | 103 ---------------------- 1 file changed, 103 deletions(-) diff --git a/tests/pytorch/attention/test_attention.py b/tests/pytorch/attention/test_attention.py index de7195bbf08..bd8fae47312 100644 --- a/tests/pytorch/attention/test_attention.py +++ b/tests/pytorch/attention/test_attention.py @@ -10,7 +10,6 @@ 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 @@ -504,108 +503,6 @@ def test_fused_attn_ragged_batch_capacity( assert kwargs["implementation"] == cudnn.attention_implementation.UNIFIED -@pytest.mark.parametrize( - "capability,layout", - ( - ((10, 0), "bshd_bshd_bshd"), - ((10, 0), "thd_bshd_bshd"), - ((10, 0), "thd_thd_thd"), - ((12, 0), "thd_thd_thd"), - ((12, 1), "thd_bshd_bshd"), - ((13, 0), "thd_thd_thd"), - ), -) -def test_fused_attn_backward_token_bounds(monkeypatch, 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.kBFloat16, - o_dtype=DType.kBFloat16, - do_dtype=DType.kBFloat16, - dqkv_dtype=DType.kBFloat16, - batch_size=4, - num_attn_heads=4, - num_gqa_groups=4, - head_dim_qk=64, - head_dim_v=64, - max_seqlen_q=512, - max_seqlen_kv=1024, - num_tokens_q=1153, - num_tokens_kv=2305, - attn_mask_type=AttnMaskType["padding"], - ) - query = implementation._query_graph(config, True) - graph = Mock() - graph.sdpa_backward.return_value = (Mock(), Mock(), Mock()) - implementation.build_attention_graph(cudnn, graph, query.description()) - bounds = { - name: value - for name, value in graph.sdpa_backward.call_args.kwargs.items() - if name.startswith("max_total_seq_len_") - } - expected = {} - if capability[0] != 12 and layout.startswith("thd"): - expected["max_total_seq_len_q"] = 2048 - if layout == "thd_thd_thd": - expected["max_total_seq_len_kv"] = 4096 - assert bounds == expected - # cuDNN's graph layout and context-parallel LSE interpretation must agree. - from transformer_engine.pytorch.attention.dot_product_attention import context_parallel - - packed = context_parallel.uses_ragged_attention_graph(10 * capability[0] + capability[1]) - assert bool(query.specs["stats"].ragged_offset) == (packed and layout.startswith("thd")) - if capability[0] == 12: - assert query.sq == 512 - assert query.skv == 1024 - - -@pytest.mark.parametrize( - "capability,layout", - ( - ((10, 0), "bshd_bshd_bshd"), - ((10, 0), "thd_bshd_bshd"), - ((10, 0), "thd_thd_thd"), - ((12, 0), "thd_thd_thd"), - ((12, 1), "thd_bshd_bshd"), - ), -) -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")) @pytest.mark.parametrize("batch", (3, 33)) From 6727540d0917e6d5fcc35556e169453518f72323 Mon Sep 17 00:00:00 2001 From: Vladimir Cherepanov Date: Wed, 11 Nov 2026 20:36:16 +0000 Subject: [PATCH 33/49] Fix JAX cuDNN support checks and reuse common handles Filter unserializable SDPA plans before support checking for ordinary and score-modified attention. Pin queries to eligible plans and retain compilation fallback across backend candidates. Reuse the common cuDNN handle in FFI initialization and execution, group attention options before CP defaults, and remove unused bias metadata. Add regressions for excluded-plan rejection and compilation fallback. Signed-off-by: Vladimir Cherepanov --- tests/jax/test_fused_attn.py | 6 +- tests/jax/test_fused_attn_score_mod.py | 61 ++++++++++++++++++- transformer_engine/common/cudnn/graph.py | 6 +- .../jax/cpp_extensions/attention.py | 10 +-- .../jax/cpp_extensions/cudnn_attention.py | 6 +- .../jax/cpp_extensions/cudnn_graph.py | 44 ++++++++++--- .../jax/csrc/extensions/cudnn_graph.cpp | 23 ++----- 7 files changed, 110 insertions(+), 46 deletions(-) diff --git a/tests/jax/test_fused_attn.py b/tests/jax/test_fused_attn.py index b658f56f92c..3508dd3ef00 100644 --- a/tests/jax/test_fused_attn.py +++ b/tests/jax/test_fused_attn.py @@ -2538,13 +2538,13 @@ def attention(operands): 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.""" +def test_cudnn_python_attention_does_not_call_common_attention(monkeypatch): + """Backend selection, workspace sizing and execution bypass the common attention APIs.""" import transformer_engine_jax from transformer_engine.jax.cpp_extensions import cudnn_attention def forbidden(*_args, **_kwargs): - raise AssertionError("JAX attention called TE common") + raise AssertionError("JAX attention called a common attention API") for name in ( "get_fused_attn_backend", diff --git a/tests/jax/test_fused_attn_score_mod.py b/tests/jax/test_fused_attn_score_mod.py index df90eb83328..40ef9330e67 100644 --- a/tests/jax/test_fused_attn_score_mod.py +++ b/tests/jax/test_fused_attn_score_mod.py @@ -658,10 +658,10 @@ def test_fused_attn_score_mod_config_splits_tensors_and_pass_by_value_scalars(): @pytest.mark.parametrize("selected_engine", (None, "sdpa_fwd", "sdpa_bwd")) def test_fused_attn_score_mod_backend_guard(monkeypatch, selected_engine): """Score-mod graphs must use cuDNN while retaining JAX's serialization exclusions.""" - cudnn = Mock(__version__="1.31.0") + cudnn = Mock(__version__="1.31.0", cudnnGraphNotSupportedError=NotImplementedError) graph = Mock(selected_engine=selected_engine) graph.get_execution_plan_count.return_value = 2 - graph.get_plan_name_at_index.side_effect = ["sdpa_fwd_frost", "cudnn_plan"] + graph.get_plan_name_at_index.side_effect = lambda index: ("sdpa_fwd_frost", "cudnn_plan")[index] graph.get_workspace_size.return_value = 0 graph.serialize.return_value = b"graph" monkeypatch.setattr(tex_cudnn_graph, "check_frontend_version", lambda _cudnn: 13100) @@ -679,10 +679,65 @@ def test_fused_attn_score_mod_backend_guard(monkeypatch, selected_engine): assert graph.deselect_engines.call_args_list == [ call(["sdpa_fwd", "sdpa_bwd"]), - call(["sdpa_fwd_frost"]), + call(["sdpa_fwd_", "sdpa_bwd_"]), ] +@pytest.mark.parametrize("has_score_mod", (False, True)) +@pytest.mark.parametrize( + "plans,backend_supported", + ( + (("sdpa_fwd_frost",), True), + (("sdpa_bwd_frost", "cudnn_plan"), True), + (("sdpa_fwd_frost", "cudnn_plan"), False), + (("cudnn_plan", "sdpa_bwd_frost"), False), + (("cudnn_plan",), True), + ), +) +def test_cudnn_graph_support_requires_serializable_plan(plans, backend_supported, has_score_mod): + """Excluded Python plans must not hide missing or unsupported backend plans.""" + from transformer_engine.common.cudnn.graph import UnsupportedGraph + + cudnn = Mock(cudnnGraphNotSupportedError=NotImplementedError) + graph = Mock() + graph.get_execution_plan_count.return_value = len(plans) + graph.get_plan_name_at_index.side_effect = lambda index: plans[index] + selected = [] + graph.select_plan.side_effect = lambda index: selected.append(index) + + def check_support(): + # Match the frontend's behavior: an unpinned query may accept a Python + # plan even when the backend cannot execute the graph. + if not selected or plans[selected[-1]].startswith("sdpa_"): + return + if not backend_supported: + raise NotImplementedError("Backend rejected graph") + + graph.check_support.side_effect = check_support + if backend_supported and "cudnn_plan" in plans: + tex_cudnn_graph.check_graph_support(cudnn, graph, has_score_mod=has_score_mod) + assert plans[selected[-1]] == "cudnn_plan" + else: + with pytest.raises(UnsupportedGraph): + tex_cudnn_graph.check_graph_support(cudnn, graph, has_score_mod=has_score_mod) + graph.build_plans.assert_not_called() + graph.serialize.assert_not_called() + + +def test_cudnn_graph_compilation_retains_backend_fallback(): + """Pinning support to the backend must not discard its later execution plans.""" + cudnn = Mock(cudnnGraphNotSupportedError=NotImplementedError) + graph = Mock() + plans = ("sdpa_fwd_frost", "cudnn_first", "cudnn_fallback") + graph.get_execution_plan_count.return_value = len(plans) + graph.get_plan_name_at_index.side_effect = lambda index: plans[index] + graph.build_plans.side_effect = [NotImplementedError("First plan failed"), None] + graph.get_workspace_size.return_value = 64 + + assert tex_cudnn_graph.compile_graph(cudnn, graph) == 64 + assert graph.select_plan.call_args_list == [call(1), call(2)] + + def test_fused_attn_score_mod_cudnn_frontend_version_check(monkeypatch): """cuDNN frontend Python and C++ versions must match.""" diff --git a/transformer_engine/common/cudnn/graph.py b/transformer_engine/common/cudnn/graph.py index 8651042a3d6..d19c7a10995 100644 --- a/transformer_engine/common/cudnn/graph.py +++ b/transformer_engine/common/cudnn/graph.py @@ -46,8 +46,8 @@ def create_graph(cudnn, dtype, handle=None): ) -def check_graph_support(cudnn, graph, run_stage=None, *, has_score_mod=False): - """Validate and select supported plans, without compiling their kernels.""" +def check_graph_support(cudnn, graph, run_stage=None, *, has_score_mod=False, filter_plans=None): + """Validate support, applying an optional plan filter before checking candidates.""" run = run_stage or (lambda _name, function: function()) try: run("validate", graph.validate) @@ -59,6 +59,8 @@ def check_graph_support(cudnn, graph, run_stage=None, *, has_score_mod=False): "create_execution_plans", lambda: graph.create_execution_plans([cudnn.heur_mode.A, cudnn.heur_mode.FALLBACK]), ) + if filter_plans is not None: + run("filter_plans", lambda: filter_plans(graph)) run("check_support", graph.check_support) except cudnn.cudnnGraphNotSupportedError as exc: raise UnsupportedGraph(f"cuDNN SDPA graph is not supported: {exc}") from exc diff --git a/transformer_engine/jax/cpp_extensions/attention.py b/transformer_engine/jax/cpp_extensions/attention.py index 06ad5e0c0dc..985a639c71d 100644 --- a/transformer_engine/jax/cpp_extensions/attention.py +++ b/transformer_engine/jax/cpp_extensions/attention.py @@ -215,12 +215,12 @@ def get_fused_attn_backend(self): max_segments_per_seq=self.max_segments_per_seq, window_size=self.window_size, bottom_right_diagonal=self.bottom_right_diagonal, + return_max_logit=self.return_max_logit, + fp8=self.fp8, 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) @@ -527,12 +527,6 @@ 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 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) - graph = cudnn_attention.get_graph( "fwd", (q_aval, k_aval, v_aval, bias_aval, softmax_offset_aval), config ) diff --git a/transformer_engine/jax/cpp_extensions/cudnn_attention.py b/transformer_engine/jax/cpp_extensions/cudnn_attention.py index 9dc4ed1b516..98fdbc5ce1b 100644 --- a/transformer_engine/jax/cpp_extensions/cudnn_attention.py +++ b/transformer_engine/jax/cpp_extensions/cudnn_attention.py @@ -26,8 +26,6 @@ bucket as _bucket, row_major_stride as _stride, create_graph, - check_graph_support, - compile_graph, uses_ragged_attention_graph, mxfp8_scale_spec, ) @@ -36,7 +34,8 @@ from .cudnn_graph import ( import_cudnn, check_frontend_version, - deselect_frost_plans, + check_graph_support, + compile_graph, serialize_graph, graph_attributes, ) @@ -197,7 +196,6 @@ def _build_graph(direction, shapes, config, capabilities, nondeterministic, fp8_ backend, direction = _site(description.spec) with _plan_lock: if entry.workspace_size is None: - deselect_frost_plans(entry.built.graph) workspace_size = compile_graph( cudnn, entry.built.graph, diff --git a/transformer_engine/jax/cpp_extensions/cudnn_graph.py b/transformer_engine/jax/cpp_extensions/cudnn_graph.py index 81d74e75104..4da11f172ec 100644 --- a/transformer_engine/jax/cpp_extensions/cudnn_graph.py +++ b/transformer_engine/jax/cpp_extensions/cudnn_graph.py @@ -12,8 +12,8 @@ from transformer_engine.common.cudnn.graph import ( import_cudnn as _import_cudnn, - check_graph_support, - compile_graph, + check_graph_support as _check_graph_support, + compile_graph as _compile_graph, ) @@ -107,22 +107,52 @@ def serialize_graph( ) -def deselect_frost_plans(graph): - """Temporarily exclude FROST SDPA plans from the JAX C++ serialization path.""" +def _serializable_plan_indices(graph, *, has_score_mod=False): + """Return plan indices eligible for the JAX C++ serialization path.""" # cuDNN cannot currently serialize FROST SDPA plans for C++ execution. # Remove this workaround once cuDNN fixes its serialization path. plan_names = [ graph.get_plan_name_at_index(index) for index in range(graph.get_execution_plan_count()) ] - graph.deselect_engines( - [name for name in plan_names if name.startswith(("sdpa_fwd_", "sdpa_bwd_"))] + prefixes = ("sdpa_fwd", "sdpa_bwd") if has_score_mod else ("sdpa_fwd_", "sdpa_bwd_") + return tuple(index for index, name in enumerate(plan_names) if not name.startswith(prefixes)) + + +def check_graph_support(cudnn, graph, run_stage=None, *, has_score_mod=False): + """Require support from a plan that the JAX executor can deserialize.""" + + def filter_plans(candidate): + candidate.deselect_engines(["sdpa_fwd_", "sdpa_bwd_"]) + indices = _serializable_plan_indices(candidate, has_score_mod=has_score_mod) + if not indices: + raise cudnn.cudnnGraphNotSupportedError("No serializable cuDNN SDPA execution plan.") + # Frontend 1.29/1.31 check_support can accept an excluded Python plan, + # or suppress a backend rejection if Python plans remain. Pin an eligible + # plan so support comes from the same executor used at runtime. + candidate.select_plan(indices[0]) + + _check_graph_support( + cudnn, graph, run_stage, has_score_mod=has_score_mod, filter_plans=filter_plans ) +def compile_graph(cudnn, graph, run_stage=None, *, has_score_mod=False): + """Compile an eligible plan, retaining fallback after support-query pinning.""" + last_error = None + for index in _serializable_plan_indices(graph, has_score_mod=has_score_mod): + graph.select_plan(index) + try: + return _compile_graph(cudnn, graph, run_stage, has_score_mod=has_score_mod) + except cudnn.cudnnGraphNotSupportedError as exc: + last_error = exc + if last_error is not None: + raise last_error + raise cudnn.cudnnGraphNotSupportedError("No serializable cuDNN SDPA execution plan.") + + def finalize_graph(cudnn, graph, *, has_score_mod=False): """Validate, select and compile a supported execution plan.""" check_graph_support(cudnn, graph, has_score_mod=has_score_mod) - deselect_frost_plans(graph) workspace_size = compile_graph(cudnn, graph, has_score_mod=has_score_mod) return workspace_size, bytes(graph.serialize()), check_frontend_version(cudnn) diff --git a/transformer_engine/jax/csrc/extensions/cudnn_graph.cpp b/transformer_engine/jax/csrc/extensions/cudnn_graph.cpp index b00bde6d77f..4973708591e 100644 --- a/transformer_engine/jax/csrc/extensions/cudnn_graph.cpp +++ b/transformer_engine/jax/csrc/extensions/cudnn_graph.cpp @@ -5,7 +5,6 @@ ************************************************************************/ // JAX-owned execution of generic graphs built with the cuDNN Python frontend. -// This translation unit deliberately has no TE common dependencies. #include #include #include @@ -21,6 +20,8 @@ #include #include +#include "cudnn_utils.h" + namespace transformer_engine { namespace jax { namespace { @@ -44,22 +45,6 @@ T Attribute(ffi::Dictionary &attrs, const char *name) { return value.value(); } -struct Handle { - cudnnHandle_t value = nullptr; - Handle() { CheckCudnn(cudnnCreate(&value)); } - ~Handle() { cudnnDestroy(value); } -}; - -cudnnHandle_t GetHandle() { - // Handles are thread-local and device-local; assign the stream on each call. - thread_local std::unordered_map> handles; - int device; - CheckCuda(cudaGetDevice(&device)); - auto &handle = handles[device]; - if (!handle) handle = std::make_unique(); - return handle->value; -} - using Graph = cudnn_frontend::graph::Graph; std::shared_ptr GetGraph(cudnnHandle_t handle, ffi::Dictionary &attrs) { @@ -114,7 +99,7 @@ const std::vector &GetScalars(ffi::Span values) { ffi::Error InitializeGraph(cudaStream_t stream, ffi::RemainingArgs, ffi::RemainingRets, ffi::Dictionary attrs) { try { - auto handle = GetHandle(); + auto handle = nvte_get_cudnn_handle(); CheckCudnn(cudnnSetStream(handle, stream)); GetGraph(handle, attrs); return ffi::Error::Success(); @@ -126,7 +111,7 @@ ffi::Error InitializeGraph(cudaStream_t stream, ffi::RemainingArgs, ffi::Remaini ffi::Error ExecuteGraph(cudaStream_t stream, ffi::RemainingArgs args, ffi::RemainingRets rets, ffi::Dictionary attrs) { try { - auto handle = GetHandle(); + auto handle = nvte_get_cudnn_handle(); CheckCudnn(cudnnSetStream(handle, stream)); auto graph = GetGraph(handle, attrs); auto input_uids = Attribute>(attrs, "input_uids"); From 303d61818309d8daa174f6b188805001d97d1be9 Mon Sep 17 00:00:00 2001 From: Vladimir Cherepanov Date: Wed, 11 Nov 2026 21:57:34 +0000 Subject: [PATCH 34/49] Replace the JAX backend binding with a Python enum Keep the existing backend query and numeric results while moving its FusedAttnBackend enum into Python. Update JAX, Flax and test consumers, remove the unused pybind enum, and retain descriptive backend logging. Validation: rebuilt the JAX extension; 57 focused tests passed. Pre-commit, focused pylint, C++ lint and changed-file licenses passed. The full lint launcher reports no diagnostics beyond the existing ones. Signed-off-by: Vladimir Cherepanov --- tests/jax/test_fused_attn.py | 14 ++++++------- tests/jax/test_fused_attn_score_mod.py | 7 ++++--- transformer_engine/jax/attention.py | 10 +++++++++- .../jax/cpp_extensions/attention.py | 20 ++++++++----------- .../jax/csrc/extensions/pybind.cpp | 5 ----- transformer_engine/jax/flax/transformer.py | 4 ++-- 6 files changed, 30 insertions(+), 30 deletions(-) diff --git a/tests/jax/test_fused_attn.py b/tests/jax/test_fused_attn.py index 3508dd3ef00..82d1c2d91b8 100644 --- a/tests/jax/test_fused_attn.py +++ b/tests/jax/test_fused_attn.py @@ -26,6 +26,7 @@ from transformer_engine.jax import autocast from transformer_engine.jax.sharding import MeshResource from transformer_engine.jax.attention import ( + FusedAttnBackend, AttnBiasType, AttnMaskType, AttnSoftmaxType, @@ -42,7 +43,6 @@ ) from transformer_engine.jax.cpp_extensions import FusedAttnHelper from transformer_engine_jax import ( - NVTE_Fused_Attn_Backend, get_cudnn_version, get_device_compute_capability, ) @@ -412,19 +412,19 @@ def test_fused_attn_backend_message(): backend, message = replace( baseline, attn_bias_type=AttnBiasType.PRE_SCALE_BIAS ).get_fused_attn_backend() - assert backend == NVTE_Fused_Attn_Backend.NVTE_No_Backend + assert backend == FusedAttnBackend.No_Backend assert message == "Fused attention does not support pre-scale bias." # No error message if supported; otherwise skip the test backend, message = baseline.get_fused_attn_backend() - if backend == NVTE_Fused_Attn_Backend.NVTE_No_Backend: + if backend == FusedAttnBackend.No_Backend: pytest.skip(f"FusedAttention does not support the baseline config: {message}") assert message == "" # All TE rules have cleared; now gets rejected by cuDNN's support check # cuDNN's error message might change across cuDNN versions, so only verify the presence of the string backend, message = replace(baseline, head_dim_qk=1024, head_dim_v=1024).get_fused_attn_backend() - assert backend == NVTE_Fused_Attn_Backend.NVTE_No_Backend + assert backend == FusedAttnBackend.No_Backend assert message != "" @@ -468,7 +468,7 @@ def fail(*_args, **_kwargs): helper.get_fused_attn_backend() else: backend, message = helper.get_fused_attn_backend() - if backend == NVTE_Fused_Attn_Backend.NVTE_No_Backend: + if backend == FusedAttnBackend.No_Backend: pytest.skip(f"Baseline attention is unsupported: {message}") @@ -1021,7 +1021,7 @@ def _check_configs(self): bias_seqlen_kv=bias_seqlen_kv, max_segments_per_seq=self._get_max_segments_per_sequence(), ).get_fused_attn_backend() - if self.backend != NVTE_Fused_Attn_Backend.NVTE_F16_arbitrary_seqlen: + if self.backend != FusedAttnBackend.F16_arbitrary_seqlen: pytest.skip(message) if ( @@ -2472,7 +2472,7 @@ def _require_cudnn_attention(q, k, v, **overrides): attn_scale=0.125, ) backend, reason = replace(helper, **overrides).get_fused_attn_backend() - if backend == NVTE_Fused_Attn_Backend.NVTE_No_Backend: + if backend == FusedAttnBackend.No_Backend: pytest.skip(reason) diff --git a/tests/jax/test_fused_attn_score_mod.py b/tests/jax/test_fused_attn_score_mod.py index 40ef9330e67..644245c6065 100644 --- a/tests/jax/test_fused_attn_score_mod.py +++ b/tests/jax/test_fused_attn_score_mod.py @@ -15,6 +15,7 @@ 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 ( + FusedAttnBackend, AttnBiasType, AttnMaskType, AttnSoftmaxType, @@ -22,7 +23,7 @@ ) from transformer_engine.jax.cpp_extensions import make_fused_attn_score_mod_config from transformer_engine.jax.flax import transformer as flax_transformer -from transformer_engine_jax import get_device_compute_capability, NVTE_Fused_Attn_Backend +from transformer_engine_jax import get_device_compute_capability from test_fused_attn import FusedAttnRunner, SeqDescFormat _CONFIG_TEST_HEAD_DIM = 128 @@ -406,9 +407,9 @@ def __init__(self, *args, **kwargs): def get_fused_attn_backend(self): if kernel_available: - return NVTE_Fused_Attn_Backend.NVTE_F16_arbitrary_seqlen, "" + return FusedAttnBackend.F16_arbitrary_seqlen, "" return ( - NVTE_Fused_Attn_Backend.NVTE_No_Backend, + FusedAttnBackend.No_Backend, "fake FusedAttnHelper: no fused attention backend available for this configuration", ) diff --git a/transformer_engine/jax/attention.py b/transformer_engine/jax/attention.py index 91437b7a9a9..26f43637172 100644 --- a/transformer_engine/jax/attention.py +++ b/transformer_engine/jax/attention.py @@ -3,7 +3,7 @@ # See LICENSE for license information. """JAX multi-head attention modules""" from __future__ import annotations -from enum import Enum +from enum import Enum, IntEnum from functools import partial import math from typing import Any, Callable, Mapping, Optional, Tuple, Union @@ -22,6 +22,14 @@ from . import cpp_extensions as tex +class FusedAttnBackend(IntEnum): + """Fused attention support results, with values matching the common C API.""" + + No_Backend = -1 + F16_arbitrary_seqlen = 1 + FP8 = 2 + + class AttnBiasType(Enum): """ NO_BIAS: Softmax is performed as softmax(scale * qk) diff --git a/transformer_engine/jax/cpp_extensions/attention.py b/transformer_engine/jax/cpp_extensions/attention.py index 985a639c71d..2c04ff4ee2a 100644 --- a/transformer_engine/jax/cpp_extensions/attention.py +++ b/transformer_engine/jax/cpp_extensions/attention.py @@ -18,11 +18,9 @@ from jax.sharding import PartitionSpec, NamedSharding from jax.experimental.custom_partitioning import SdyShardingRule -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 ( + FusedAttnBackend, AttnBiasType, AttnMaskType, AttnSoftmaxType, @@ -181,7 +179,7 @@ class FusedAttnHelper: def is_fused_attn_kernel_available(self): """Check if there is available fused attention kernel""" backend, _ = self.get_fused_attn_backend() - return backend != NVTE_Fused_Attn_Backend.NVTE_No_Backend + return backend != FusedAttnBackend.No_Backend def get_fused_attn_backend(self): """Get the fused attention kernel backend. @@ -197,7 +195,7 @@ def get_fused_attn_backend(self): raise ValueError('fp8 must be False, True, or "mxfp8".') if self.q_dtype != self.kv_dtype: raise ValueError("Q and KV must have the same data type.") - backend = NVTE_Fused_Attn_Backend.NVTE_No_Backend + backend = FusedAttnBackend.No_Backend message = "" if self.attn_bias_type == AttnBiasType.PRE_SCALE_BIAS: message = "Fused attention does not support pre-scale bias." @@ -251,9 +249,7 @@ def get_fused_attn_backend(self): if self.is_training: cudnn_attention.check_attention_support("bwd", avals, config) backend = ( - NVTE_Fused_Attn_Backend.NVTE_FP8 - if self.fp8 - else NVTE_Fused_Attn_Backend.NVTE_F16_arbitrary_seqlen + FusedAttnBackend.FP8 if self.fp8 else FusedAttnBackend.F16_arbitrary_seqlen ) except ValueError as exc: message = str(exc) @@ -261,14 +257,14 @@ def get_fused_attn_backend(self): AttentionLogging.setup_logging() logger = AttentionLogging.logger logger.debug("Running fused attention backend selection with config=%s", self) - if backend == NVTE_Fused_Attn_Backend.NVTE_No_Backend: + if backend == FusedAttnBackend.No_Backend: logger.info("No fused attention backend available; falling back to unfused attention.") logger.debug( "Reason fused attention was rejected: %s", message or "(no diagnostic message available)", ) else: - logger.info("Selected fused attention backend: %s", backend) + logger.info("Selected fused attention backend: %s", backend.name) if message: logger.debug("Fused attention backend diagnostic message: %s", message) return backend, message @@ -494,8 +490,8 @@ def abstract( ).get_fused_attn_backend() if backend in ( - NVTE_Fused_Attn_Backend.NVTE_F16_arbitrary_seqlen, - NVTE_Fused_Attn_Backend.NVTE_FP8, + FusedAttnBackend.F16_arbitrary_seqlen, + FusedAttnBackend.FP8, ): if config.qkv_layout.is_thd(): softmax_shape = (*batch_shape, q_max_seqlen, attn_heads, 1) diff --git a/transformer_engine/jax/csrc/extensions/pybind.cpp b/transformer_engine/jax/csrc/extensions/pybind.cpp index e56a7214565..09b2df5d4b8 100644 --- a/transformer_engine/jax/csrc/extensions/pybind.cpp +++ b/transformer_engine/jax/csrc/extensions/pybind.cpp @@ -244,11 +244,6 @@ PYBIND11_MODULE(transformer_engine_jax, m) { .value("CLAMPED_SWIGLU", NVTE_Activation_Type::CLAMPED_SWIGLU) .export_values(); - pybind11::enum_(m, "NVTE_Fused_Attn_Backend", pybind11::module_local()) - .value("NVTE_No_Backend", NVTE_Fused_Attn_Backend::NVTE_No_Backend) - .value("NVTE_F16_arbitrary_seqlen", NVTE_Fused_Attn_Backend::NVTE_F16_arbitrary_seqlen) - .value("NVTE_FP8", NVTE_Fused_Attn_Backend::NVTE_FP8); - pybind11::enum_(m, "NVTE_Norm_Type", pybind11::module_local()) .value("LayerNorm", NVTE_Norm_Type::LayerNorm) .value("RMSNorm", NVTE_Norm_Type::RMSNorm) diff --git a/transformer_engine/jax/flax/transformer.py b/transformer_engine/jax/flax/transformer.py index 10389518778..c204bd29eea 100644 --- a/transformer_engine/jax/flax/transformer.py +++ b/transformer_engine/jax/flax/transformer.py @@ -21,11 +21,11 @@ from jax import random as jax_random from jax import lax, vmap from jax.ad_checkpoint import checkpoint_name -from transformer_engine_jax import NVTE_Fused_Attn_Backend from .module import DenseGeneral, LayerNormDenseGeneral, LayerNormMLP from .module import LayerNorm, Softmax from ..attention import ( + FusedAttnBackend, AttnBiasType, AttnMaskType, AttnSoftmaxType, @@ -873,7 +873,7 @@ def __call__( 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 + has_fused_attn_kernel = fused_attn_backend != FusedAttnBackend.No_Backend if attn_bias_type == AttnBiasType.ALIBI and ( not enable_fused_attn or not has_fused_attn_kernel ): From e39650069303d3c9ca72336c5d625ffda51fbd06 Mon Sep 17 00:00:00 2001 From: Vladimir Cherepanov Date: Wed, 11 Nov 2026 22:23:08 +0000 Subject: [PATCH 35/49] Support both O precisions in JAX FP8 backward Honor NVTE_DPA_FP8CS_O_in_F16 for current-scaled attention. Include the resolved O precision in graph requests, describe and quantize FP8 O with its inverse scale, and retain model-precision forward outputs. Ignore the option for MXFP8 and model-precision backward. Document the precision option and the absence of delayed-scaling JAX attention. Extend existing numerical, context-parallel and cache tests. Validation: 132 focused tests passed, one skipped, including eight two-GPU context-parallel cases. Existing numerical tolerances retained. Pre-commit, focused pylint and changed-file license checks passed. Full JAX lint reports only the existing diagnostics. Signed-off-by: Vladimir Cherepanov --- docs/envvars.rst | 6 +- tests/jax/test_fused_attn.py | 66 ++++++++++++++++++- transformer_engine/jax/attention.py | 9 ++- .../jax/cpp_extensions/cudnn_attention.py | 33 ++++++++-- 4 files changed, 101 insertions(+), 13 deletions(-) diff --git a/docs/envvars.rst b/docs/envvars.rst index 04108f8b773..8a55ff508a4 100644 --- a/docs/envvars.rst +++ b/docs/envvars.rst @@ -288,7 +288,11 @@ FP8 Configuration :Type: ``int`` (0 or 1) :Default: ``1`` - :Description: For Float8CurrentScaling in DotProductAttention, use FP16/BF16 for the output tensor in the backward pass. ``1`` = use F16/BF16 output in backward, ``0`` = use FP8 output in backward. + :Description: For current-scaled FP8 attention in PyTorch or JAX, control the precision + of the forward output tensor consumed by the backward pass. ``1`` = use FP16/BF16 + output in backward, ``0`` = quantize it to FP8. This option does not affect MXFP8 + attention or FP16/BF16 backward selected with ``NVTE_FP8_DPA_BWD=0``. In JAX, + public forward outputs are unchanged; set this option before tracing/compilation. .. envvar:: NVTE_DPA_FP8_RECIPE diff --git a/tests/jax/test_fused_attn.py b/tests/jax/test_fused_attn.py index 82d1c2d91b8..ed8f69ad615 100644 --- a/tests/jax/test_fused_attn.py +++ b/tests/jax/test_fused_attn.py @@ -472,6 +472,59 @@ def fail(*_args, **_kwargs): pytest.skip(f"Baseline attention is unsupported: {message}") +@pytest.mark.parametrize( + "fp8,fp8_backward", ((False, True), (True, False), (True, True), ("mxfp8", True)) +) +def test_cudnn_attention_backward_o_precision_cache(fp8, fp8_backward, monkeypatch): + """Support and execution share O precision, without duplicating unaffected graphs.""" + from transformer_engine.jax.cpp_extensions import cudnn_attention + from transformer_engine.jax.cpp_extensions.attention import _FusedAttnConfig + + monkeypatch.setenv("NVTE_FP8_DPA_BWD", "1" if fp8_backward else "0") + q = jax.ShapeDtypeStruct((1, 128, 2, 64), jnp.bfloat16) + empty = jax.ShapeDtypeStruct((0,), q.dtype) + avals = (q, q, q, empty, jax.ShapeDtypeStruct((0,), jnp.float32)) + _require_cudnn_attention(q, q, q, fp8=fp8) + config = _FusedAttnConfig( + attn_bias_type=AttnBiasType.NO_BIAS, + attn_mask_type=AttnMaskType.NO_MASK, + softmax_type=AttnSoftmaxType.VANILLA_SOFTMAX, + qkv_layout=QKVLayout.BSHD_BSHD_BSHD, + scaling_factor=0.125, + dropout_probability=0.0, + is_training=True, + max_segments_per_seq=1, + window_size=(-1, -1), + bottom_right_diagonal=False, + context_parallel_load_balanced=False, + cp_axis="", + cp_striped_window_size=None, + stripe_size=None, + fp8=fp8, + ) + backward, forward = [], [] + for o_in_f16 in (True, False): + monkeypatch.setenv("NVTE_DPA_FP8CS_O_in_F16", "1" if o_in_f16 else "0") + cudnn_attention.check_attention_support("bwd", avals, config) + plan = cudnn_attention.get_graph("bwd", avals, config) + expected_fp8_o = fp8 is True and fp8_backward and not o_in_f16 + assert plan.fp8_o == expected_fp8_o + assert ("descale_o" in plan.names) == expected_fp8_o + description = cudnn_attention._describe_graph( + *cudnn_attention._request("bwd", avals, config) + ) + assert dict(description.spec.tensors)["o"].dtype == ( + "FP8_E4M3" if expected_fp8_o else "BFLOAT16" + ) + backward.append(plan) + forward.append(cudnn_attention.get_graph("fwd", avals, config)) + assert forward[0] is forward[1] + if fp8 is True and fp8_backward: + assert backward[0].graph.graph_hash != backward[1].graph.graph_hash + else: + assert backward[0] is backward[1] + + @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.""" @@ -674,6 +727,7 @@ def test_cudnn_attention_output_initialization(layout, direction, fp8, monkeypat ragged_stats=layout.is_thd(), multipliers=(128, 128, 128, 128, 2), fp8=fp8, + fp8_o=False, ) monkeypatch.setattr(cudnn_attention, "get_graph", lambda *_args: plan) monkeypatch.setattr(cudnn_attention, "graph_attributes", lambda _graph: {}) @@ -2476,6 +2530,7 @@ def _require_cudnn_attention(q, k, v, **overrides): pytest.skip(reason) +@pytest.mark.parametrize("o_in_f16", [True, False]) @pytest.mark.parametrize("dtype", [jnp.float16, jnp.bfloat16]) @pytest.mark.parametrize("fp8_backward", [False, True]) @pytest.mark.parametrize( @@ -2486,9 +2541,10 @@ def _require_cudnn_attention(q, k, v, **overrides): QKVLayout.BSHD_BSHD_BSHD, ], ) -def test_cudnn_python_fp8_forward_backward(dtype, layout, fp8_backward, monkeypatch): +def test_cudnn_python_fp8_forward_backward(dtype, layout, fp8_backward, o_in_f16, monkeypatch): """Exercise real FP8 SDPA with exactly representable operands and probabilities.""" monkeypatch.setenv("NVTE_FP8_DPA_BWD", "1" if fp8_backward else "0") + monkeypatch.setenv("NVTE_DPA_FP8CS_O_in_F16", "1" if o_in_f16 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) @@ -2593,6 +2649,7 @@ def forbidden(*_args, **_kwargs): np.testing.assert_array_equal(gradients, tensor) +@pytest.mark.parametrize("o_in_f16", [True, False]) @pytest.mark.parametrize("scaling_factor", [0.125, 0.25]) @pytest.mark.parametrize( "kind,sm_arch", @@ -2608,8 +2665,9 @@ def forbidden(*_args, **_kwargs): ("ragged_dense_stats_f16", 129), ], ) -def test_cudnn_python_fp8_gqa_masks(kind, sm_arch, scaling_factor, monkeypatch, request): +def test_cudnn_python_fp8_gqa_masks(kind, sm_arch, scaling_factor, o_in_f16, monkeypatch, request): """Compare FP8 GQA against FP32, including cached plans with different scalar values.""" + monkeypatch.setenv("NVTE_DPA_FP8CS_O_in_F16", "1" if o_in_f16 else "0") ragged = kind.startswith("ragged") fp8 = not kind.endswith("_f16") if ragged and get_device_compute_capability(0) < 100: @@ -2877,10 +2935,12 @@ def attention(arrays, seed): np.testing.assert_array_equal(actual, expected) +@pytest.mark.parametrize("o_in_f16", [True, False]) @pytest.mark.parametrize("strategy", [CPStrategy.ALL_GATHER, CPStrategy.RING]) @pytest.mark.parametrize("fp8", [False, True]) -def test_cudnn_python_context_parallel(strategy, fp8): +def test_cudnn_python_context_parallel(strategy, fp8, o_in_f16, monkeypatch): """Two-device context parallel attention preserves FP8 forward and backward.""" + monkeypatch.setenv("NVTE_DPA_FP8CS_O_in_F16", "1" if o_in_f16 else "0") if jax.local_device_count() < 2: pytest.skip("Context parallel attention requires two GPUs.") mesh = Mesh(np.asarray(jax.local_devices()[:2]), ("cp",)) diff --git a/transformer_engine/jax/attention.py b/transformer_engine/jax/attention.py index 26f43637172..cd5d6626472 100644 --- a/transformer_engine/jax/attention.py +++ b/transformer_engine/jax/attention.py @@ -1688,9 +1688,12 @@ def fused_attn( score_mod_bprop_tensors (Optional[Mapping[str, Any]]): Additional tensors or Python/NumPy scalars made available to `score_mod_bprop`. fp8 (bool | str): False disables FP8 attention; True enables FP8 attention - with current tensor scaling; "mxfp8" enables block-scaled MXFP8 - attention. This option is independent of the GEMM quantization recipe. - Outputs and gradients retain the input FP16/BF16 dtype. + with current tensor scaling (CS); "mxfp8" enables block-scaled MXFP8 + attention. Delayed scaling (DS) is not supported for JAX attention. + This option is independent of the GEMM quantization recipe. + Outputs and gradients retain the input FP16/BF16 dtype. For CS FP8 + backward, ``NVTE_DPA_FP8CS_O_in_F16`` selects whether the saved output + is consumed in FP16/BF16 (default, ``1``) or quantized to FP8 (``0``). quantized_qkv (Optional[Tuple[ScaledTensor2x, ...]]): Prequantized MXFP8 Q/K/V, each with independent unswizzled BHSD rowwise/columnwise data and E8M0 scales. Requires ``fp8="mxfp8"``. The original ``qkv`` diff --git a/transformer_engine/jax/cpp_extensions/cudnn_attention.py b/transformer_engine/jax/cpp_extensions/cudnn_attention.py index 98fdbc5ce1b..e5a7eefa7f0 100644 --- a/transformer_engine/jax/cpp_extensions/cudnn_attention.py +++ b/transformer_engine/jax/cpp_extensions/cudnn_attention.py @@ -91,6 +91,7 @@ class AttentionGraph: ragged_stats: bool multipliers: tuple fp8: bool + fp8_o: bool = False @dataclass(frozen=True) @@ -144,13 +145,22 @@ def _supported_graph(spec, capabilities, frontend_version, backend_version): def _request(direction, avals, config): shapes = tuple((tuple(aval.shape), jnp.dtype(aval.dtype)) for aval in avals[:5]) + fp8_backward = bool(int(os.getenv("NVTE_FP8_DPA_BWD", "1"))) + fp8_o = ( + direction == "bwd" + and bool(config.fp8) + and config.fp8 != "mxfp8" + and fp8_backward + and not bool(int(os.getenv("NVTE_DPA_FP8CS_O_in_F16", "1"))) + ) return ( direction, shapes, config, device_capabilities(), bool(int(os.getenv("NVTE_ALLOW_NONDETERMINISTIC_ALGO", "1"))), - bool(int(os.getenv("NVTE_FP8_DPA_BWD", "1"))), + fp8_backward, + fp8_o, ) @@ -184,9 +194,11 @@ def get_graph(direction, avals, config): @lru_cache(maxsize=None) -def _build_graph(direction, shapes, config, capabilities, nondeterministic, fp8_backward): +def _build_graph( + direction, shapes, config, capabilities, nondeterministic, fp8_backward, fp8_o=False +): description = _describe_graph( - direction, shapes, config, capabilities, nondeterministic, fp8_backward + direction, shapes, config, capabilities, nondeterministic, fp8_backward, fp8_o ) cudnn = import_cudnn() version = check_frontend_version(cudnn) @@ -226,10 +238,13 @@ def _build_graph(direction, shapes, config, capabilities, nondeterministic, fp8_ description.ragged_stats, description.multipliers, description.fp8, + fp8_o, ) -def _describe_graph(direction, shapes, config, capabilities, nondeterministic, fp8_backward): +def _describe_graph( + direction, shapes, config, capabilities, nondeterministic, fp8_backward, fp8_o=False +): """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) @@ -373,7 +388,7 @@ def output(name, dim, stride, slot, offset=0, dtype=io_dtype, ragged_offset=""): constant("scale_s", 448.0) constant("descale_s", 1.0 / 448.0) if backward: - tensor("o", odim, ostride, ragged_offset=o_offset) + tensor("o", odim, ostride, "FP8_E4M3" if fp8_o else io_dtype, ragged_offset=o_offset) tensor( "do", odim, @@ -411,7 +426,11 @@ def output(name, dim, stride, slot, offset=0, dtype=io_dtype, ragged_offset=""): inputs.append(("descale_do", "descale_do", 0)) elif fp8: tensor("descale_do", (1, 1, 1, 1), (1, 1, 1, 1), "FLOAT") - for name in ("descale_o", "scale_dp", "descale_dp", "scale_dq", "scale_dk", "scale_dv"): + if fp8_o: + tensor("descale_o", (1, 1, 1, 1), (1, 1, 1, 1), "FLOAT") + else: + constant("descale_o", 1.0) + for name in ("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") @@ -566,6 +585,8 @@ def execute(direction, args, config, result_avals, prepared_qkv=None): values.update(mxfp8_attention.prepare_doutput(do)) elif plan.fp8: values["do"], values["descale_do"] = _quantize(do, jnp.float8_e5m2) + if plan.fp8_o: + values["o"], values["descale_o"] = _quantize(o, jnp.float8_e4m3fn) specs = list(result_avals[:-1]) if plan.fp8 and config.fp8 != "mxfp8": specs += [jax.ShapeDtypeStruct((1,), jnp.float32)] * 4 From d5fdde21aedb22012fa9c4212a487f7e020546a3 Mon Sep 17 00:00:00 2001 From: Vladimir Cherepanov Date: Wed, 11 Nov 2026 22:33:25 +0000 Subject: [PATCH 36/49] Match main cuDNN attention heuristic selection Request heuristic mode A without FALLBACK, matching the common fused attention path on main. Update the existing score-modified attention guard expectation. Validation: 158 focused PyTorch tests passed, including Flex Attention, cache reuse, ragged numerics, and delayed/current/MXFP8 attention. Black, whitespace, and changed-file license checks passed. Signed-off-by: Vladimir Cherepanov --- tests/pytorch/attention/test_flex_attention.py | 4 +--- transformer_engine/common/cudnn/graph.py | 2 +- 2 files changed, 2 insertions(+), 4 deletions(-) diff --git a/tests/pytorch/attention/test_flex_attention.py b/tests/pytorch/attention/test_flex_attention.py index 5a8452029ce..c6e8d7775fe 100644 --- a/tests/pytorch/attention/test_flex_attention.py +++ b/tests/pytorch/attention/test_flex_attention.py @@ -54,9 +54,7 @@ def test_cudnn_score_mod_backend_guard(monkeypatch, has_score_mod, selected_engi assert ( graph.mock_calls.index(call.build_operation_graph()) < graph.mock_calls.index(exclusion) - < graph.mock_calls.index( - call.create_execution_plans([cudnn.heur_mode.A, cudnn.heur_mode.FALLBACK]) - ) + < graph.mock_calls.index(call.create_execution_plans([cudnn.heur_mode.A])) ) else: graph.deselect_engines.assert_not_called() diff --git a/transformer_engine/common/cudnn/graph.py b/transformer_engine/common/cudnn/graph.py index d19c7a10995..ee0f111d3dd 100644 --- a/transformer_engine/common/cudnn/graph.py +++ b/transformer_engine/common/cudnn/graph.py @@ -57,7 +57,7 @@ def check_graph_support(cudnn, graph, run_stage=None, *, has_score_mod=False, fi run("deselect_engines", lambda: graph.deselect_engines(["sdpa_fwd", "sdpa_bwd"])) run( "create_execution_plans", - lambda: graph.create_execution_plans([cudnn.heur_mode.A, cudnn.heur_mode.FALLBACK]), + lambda: graph.create_execution_plans([cudnn.heur_mode.A]), ) if filter_plans is not None: run("filter_plans", lambda: filter_plans(graph)) From 509322ffc95d79aaa425e903436f9205628c9af6 Mon Sep 17 00:00:00 2001 From: Vladimir Cherepanov Date: Wed, 11 Nov 2026 22:39:28 +0000 Subject: [PATCH 37/49] Group cuDNN attention dispatch and rename MXFP8 helpers Keep F16, tensor-scaled FP8, and MXFP8 dispatch in separate sections, with forward and backward calls together. Clarify that fp8 describes the effective precision already resolved by the frontend. Rename the shared preparation module to mxfp8_utils and update both frontend imports and the existing JAX tests. Validation: 182 focused PyTorch tests passed, covering all three FP8 recipes and both backward precision settings. Black, syntax, whitespace, and changed-file license checks passed. JAX runtime validation was unavailable because JAX and its extension are not installed. Signed-off-by: Vladimir Cherepanov --- tests/jax/test_mxfp8_attention.py | 2 +- transformer_engine/common/cudnn/attention.py | 132 +++++++++--------- .../common/cudnn/{mxfp8.py => mxfp8_utils.py} | 0 .../jax/cpp_extensions/cudnn_attention.py | 2 +- .../jax/cpp_extensions/mxfp8_attention.py | 2 +- .../attention/dot_product_attention/utils.py | 2 +- .../cpp_extensions/cudnn_fused_attn.py | 2 +- 7 files changed, 72 insertions(+), 70 deletions(-) rename transformer_engine/common/cudnn/{mxfp8.py => mxfp8_utils.py} (100%) diff --git a/tests/jax/test_mxfp8_attention.py b/tests/jax/test_mxfp8_attention.py index 35eee619e9f..5a34081c258 100644 --- a/tests/jax/test_mxfp8_attention.py +++ b/tests/jax/test_mxfp8_attention.py @@ -11,7 +11,7 @@ import numpy as np import pytest -from transformer_engine.common.cudnn import mxfp8 +from transformer_engine.common.cudnn import mxfp8_utils as mxfp8 from transformer_engine.common.cudnn.graph import mxfp8_scale_spec from transformer_engine.jax.attention import ( AttnBiasType, diff --git a/transformer_engine/common/cudnn/attention.py b/transformer_engine/common/cudnn/attention.py index 5ee27c2fc88..44f45584d24 100644 --- a/transformer_engine/common/cudnn/attention.py +++ b/transformer_engine/common/cudnn/attention.py @@ -14,7 +14,7 @@ class AttentionOptions: """Normalized attributes that affect graph construction and support.""" backward: bool = False - fp8: bool = False + fp8: bool = False # Effective precision for this pass, resolved by the frontend. mxfp8: bool = False attn_scale: str | float = "attn_scale" bottom_right: bool = False @@ -151,20 +151,17 @@ def tensor(name): q, k, v = (tensor(name) for name in ("q", "k", "v")) if not options.fp8: - if options.backward: + if not options.backward: + outputs = graph.sdpa(q, k, v, **kwargs) + names = ("o", "stats") + else: 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: + elif not options.mxfp8: + if not options.backward: + scales = [tensor(name) for name in ("descale_q", "descale_k", "descale_v")] outputs = graph.sdpa_fp8( q, k, @@ -174,61 +171,66 @@ def tensor(name): **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") 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") + if not options.backward: + scales = [tensor(name) for name in ("descale_q", "descale_k", "descale_v")] + outputs = graph.sdpa_mxfp8(q, k, v, *scales, **kwargs) + names = ("o", "stats", "amax_o") + else: + 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") for name, desc in zip(names, outputs): if name not in specs: diff --git a/transformer_engine/common/cudnn/mxfp8.py b/transformer_engine/common/cudnn/mxfp8_utils.py similarity index 100% rename from transformer_engine/common/cudnn/mxfp8.py rename to transformer_engine/common/cudnn/mxfp8_utils.py diff --git a/transformer_engine/jax/cpp_extensions/cudnn_attention.py b/transformer_engine/jax/cpp_extensions/cudnn_attention.py index e5a7eefa7f0..27b681dd37e 100644 --- a/transformer_engine/jax/cpp_extensions/cudnn_attention.py +++ b/transformer_engine/jax/cpp_extensions/cudnn_attention.py @@ -29,7 +29,7 @@ uses_ragged_attention_graph, mxfp8_scale_spec, ) -from transformer_engine.common.cudnn import mxfp8 as mxfp8_rules +from transformer_engine.common.cudnn import mxfp8_utils as mxfp8_rules from .cudnn_graph import ( import_cudnn, diff --git a/transformer_engine/jax/cpp_extensions/mxfp8_attention.py b/transformer_engine/jax/cpp_extensions/mxfp8_attention.py index b948ef32b86..cb222c93e7b 100644 --- a/transformer_engine/jax/cpp_extensions/mxfp8_attention.py +++ b/transformer_engine/jax/cpp_extensions/mxfp8_attention.py @@ -12,7 +12,7 @@ from jax.sharding import NamedSharding, PartitionSpec from jax.custom_batching import sequential_vmap -from transformer_engine.common.cudnn.mxfp8 import quantization_layouts, scale_reordering +from transformer_engine.common.cudnn.mxfp8_utils import quantization_layouts, scale_reordering from ..quantize import QuantizerFactory, QuantizeLayout, ScalingMode, ScaledTensor2x from ..quantize.tensor import ScaledTensorFactory from .quantization import quantize diff --git a/transformer_engine/pytorch/attention/dot_product_attention/utils.py b/transformer_engine/pytorch/attention/dot_product_attention/utils.py index 6b9fe808181..c71cc78a46a 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/utils.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/utils.py @@ -40,7 +40,7 @@ DType, MXFP8_BLOCK_SCALING_SIZE, ) -from transformer_engine.common.cudnn.mxfp8 import ( +from transformer_engine.common.cudnn.mxfp8_utils import ( quantization_layouts as mxfp8_quantization_layouts, scale_reordering as mxfp8_scale_reordering, shape_error as mxfp8_shape_error, diff --git a/transformer_engine/pytorch/cpp_extensions/cudnn_fused_attn.py b/transformer_engine/pytorch/cpp_extensions/cudnn_fused_attn.py index 403d31dfbad..8dbae20a3d7 100644 --- a/transformer_engine/pytorch/cpp_extensions/cudnn_fused_attn.py +++ b/transformer_engine/pytorch/cpp_extensions/cudnn_fused_attn.py @@ -33,7 +33,7 @@ bucket, uses_ragged_attention_graph, ) -from transformer_engine.common.cudnn.mxfp8 import support_error as mxfp8_support_error +from transformer_engine.common.cudnn.mxfp8_utils import support_error as mxfp8_support_error from .cudnn import ( _import_cudnn_frontend, From c2904577678109a6f96db80a4e9240307ec32c3e Mon Sep 17 00:00:00 2001 From: Vladimir Cherepanov Date: Wed, 11 Nov 2026 23:10:27 +0000 Subject: [PATCH 38/49] Match fused attention bias query and runtime descriptors Carry actual bias strides into support queries so strided bias inputs reuse queried graphs. Allocate contiguous dBias to satisfy cuDNN for permuted bias inputs. Add graph identity, compilation count, and independent numerical checks for FP16/BF16 bias layouts. Extend optional-output descriptor coverage. Validation: 251 focused PyTorch tests passed, with 12 existing skips. Repository pre-commit, Python lint, syntax, whitespace, and changed-file license checks passed. Signed-off-by: Vladimir Cherepanov --- tests/pytorch/attention/test_attention.py | 94 ++++++++++++++++++- .../dot_product_attention.py | 5 + .../attention/dot_product_attention/utils.py | 10 ++ .../cpp_extensions/cudnn_fused_attn.py | 5 +- 4 files changed, 110 insertions(+), 4 deletions(-) diff --git a/tests/pytorch/attention/test_attention.py b/tests/pytorch/attention/test_attention.py index bd8fae47312..159c1d9cde1 100644 --- a/tests/pytorch/attention/test_attention.py +++ b/tests/pytorch/attention/test_attention.py @@ -368,7 +368,7 @@ def test_cudnn_attention_handle_streams(): @pytest.mark.parametrize("backward", (False, True)) @pytest.mark.parametrize("dropout", (0.0, 0.125)) def test_fused_attn_query_matches_runtime_graph(layout, backward, dropout): - """Packed storage and runtime scalar values must preserve support-query graph identity.""" + """Packed storage, optional outputs, and scalars preserve query/runtime graph identity.""" from transformer_engine.pytorch.cpp_extensions import cudnn_fused_attn as implementation batch, seqlen, heads, dim = 2, 128, 4, 64 @@ -392,7 +392,9 @@ def test_fused_attn_query_matches_runtime_graph(layout, backward, dropout): max_seqlen_kv=seqlen, attn_mask_type="padding" if fmt == "thd" else "no_mask", attn_bias_type="no_bias", - softmax_type="vanilla", + softmax_type="learnable", + softmax_offset=torch.zeros(1, heads, 1, 1, device="cuda", dtype=torch.float32), + return_max_logit=not backward, dropout=dropout, rng_state=torch.zeros(2, device="cuda", dtype=torch.int64), attn_scale=0.125, @@ -412,6 +414,8 @@ def test_fused_attn_query_matches_runtime_graph(layout, backward, dropout): 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"]], + softmax_type=SoftmaxType["learnable"], + return_max_logit=not backward, dropout=dropout, ) query = implementation._query_graph(config, backward) @@ -421,11 +425,14 @@ def test_fused_attn_query_matches_runtime_graph(layout, backward, dropout): runtime.matrix("o", out, fmt, heads, runtime.sq, cu, output=not backward) stats_shape = (batch * seqlen, heads, 1) if runtime.ragged_stats else (batch, heads, seqlen, 1) runtime.stats("stats", torch.empty(stats_shape, device="cuda"), output=not backward) + if not backward: + runtime.stats("max", torch.empty(stats_shape, device="cuda"), output=True) 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) + runtime.add("dsink", torch.empty_like(params["softmax_offset"]), output=True) assert query.description() == runtime.description() if fmt == "thd" and not query.use_cu_seqlens: @@ -435,6 +442,89 @@ def test_fused_attn_query_matches_runtime_graph(layout, backward, dropout): assert query.description() == runtime.description() +@pytest.mark.parametrize("dtype", (torch.float16, torch.bfloat16)) +@pytest.mark.parametrize("qkv_format", ("bshd", "sbhd")) +@pytest.mark.parametrize("bias_layout", ("contiguous", "gapped", "permuted")) +def test_fused_attn_bias_graph_cache(dtype, qkv_format, bias_layout, monkeypatch): + """Bias layouts reuse queried plans and preserve forward/backward numerics.""" + from transformer_engine.pytorch.cpp_extensions import cudnn_fused_attn as implementation + + monkeypatch.setenv("NVTE_FLASH_ATTN", "0") + monkeypatch.setenv("NVTE_FUSED_ATTN", "1") + monkeypatch.setenv("NVTE_UNFUSED_ATTN", "0") + monkeypatch.setattr(implementation, "_plans", {}) + compile_graph = Mock(wraps=implementation.compile_graph) + monkeypatch.setattr(implementation, "compile_graph", compile_graph) + original_entry = implementation._AttentionGraph.entry + queried = {} + executed = [] + + def entry(attn): + description = attn.description() + runtime = "q" in attn.params + if runtime: + assert description in queried, "Execution changed the queried tensor descriptors" + result = original_entry(attn) + if runtime: + assert result is queried[description] + executed.append(attn.backward) + else: + queried[description] = result + return result + + monkeypatch.setattr(implementation._AttentionGraph, "entry", entry) + _attention_backends["backend_selection_requires_update"] = True + batch, seqlen, heads, dim = 2, 128, 4, 64 + block = DotProductAttention( + heads, dim, qkv_format=qkv_format, attn_mask_type="no_mask", layer_number=1 + ).to(dtype=dtype, device="cuda") + rtol, atol = (2e-2, 1e-2) if dtype == torch.bfloat16 else (3e-3, 2e-3) + for _ in range(2): + qkv = [torch.randn(batch, seqlen, heads, dim, dtype=dtype, device="cuda") for _ in range(3)] + if bias_layout == "gapped": + bias = torch.randn(1, heads, 2 * seqlen, seqlen, dtype=dtype, device="cuda") + bias = bias[..., ::2, :] + elif bias_layout == "permuted": + bias = torch.randn(1, seqlen, heads, seqlen, dtype=dtype, device="cuda") + bias = bias.transpose(1, 2) + else: + bias = torch.randn(1, heads, seqlen, seqlen, dtype=dtype, device="cuda") + bias.requires_grad_() + inputs = [ + (x if qkv_format == "bshd" else x.transpose(0, 1).contiguous()).requires_grad_() + for x in qkv + ] + refs = [x.float().detach().requires_grad_() for x in qkv] + ref_bias = bias.detach().float().requires_grad_() + qr, kr, vr = [x.transpose(1, 2) for x in refs] + scores = qr @ kr.transpose(-1, -2) / dim**0.5 + ref_bias + reference = (scores.softmax(-1) @ vr).transpose(1, 2) + grad = torch.randn_like(reference, dtype=dtype) + reference.backward(grad.float()) + + output = block( + *inputs, + attn_mask_type="no_mask", + core_attention_bias_type="post_scale_bias", + core_attention_bias=bias, + ) + output_grad = grad if qkv_format == "bshd" else grad.transpose(0, 1).contiguous() + output.backward(output_grad.reshape(output.shape)) + output = output.reshape(*inputs[0].shape) + if qkv_format == "sbhd": + output = output.transpose(0, 1) + torch.testing.assert_close(output.float(), reference, rtol=rtol, atol=atol) + for actual, ref in zip(inputs, refs): + actual_grad = actual.grad if qkv_format == "bshd" else actual.grad.transpose(0, 1) + torch.testing.assert_close(actual_grad.float(), ref.grad, rtol=rtol, atol=atol) + # dBias subtracts dO.O using low-precision saved O, amplifying pointwise cancellation. + error = torch.linalg.vector_norm(bias.grad.float() - ref_bias.grad) + assert error <= torch.finfo(dtype).eps * torch.linalg.vector_norm(ref_bias.grad) + assert len(implementation._plans) == 2 + assert compile_graph.call_count == 2 + assert executed == [False, True, False, True] + + @pytest.mark.parametrize( "backward,dropout,version,fp8,layout,capability,expected_batch,direct", ( diff --git a/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py b/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py index f15d6fcd563..3ba6f069e1d 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py @@ -2802,6 +2802,11 @@ def forward( ), "core_attention_bias_type": core_attention_bias_type, "core_attention_bias_shape": core_attention_bias_shape, + "core_attention_bias_stride": ( + tuple(core_attention_bias.stride()) + if core_attention_bias_type != "no_bias" and core_attention_bias is not None + else None + ), "core_attention_bias_requires_grad": ( core_attention_bias.requires_grad if core_attention_bias_type != "no_bias" and core_attention_bias is not None diff --git a/transformer_engine/pytorch/attention/dot_product_attention/utils.py b/transformer_engine/pytorch/attention/dot_product_attention/utils.py index c71cc78a46a..52522f6ac2f 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/utils.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/utils.py @@ -263,6 +263,8 @@ class AttentionParams: Attention bias type, {`no_bias`, `pre_scale_bias`, `post_scale_bias`, `alibi`}. core_attention_bias_shape : Optional[Tuple[int, int, int, int]], default = None Attention bias shape, (b, h, sq, skv). + core_attention_bias_stride : Optional[Tuple[int, int, int, int]], default = None + Attention bias strides. None describes contiguous storage. core_attention_bias_requires_grad : bool, default = True Whether attention bias requires gradient. pad_between_seqs : bool, default = False @@ -356,6 +358,7 @@ class AttentionParams: has_score_mod_bprop: bool = False device: Optional[torch.device] = None requires_backward: bool = True + core_attention_bias_stride: Optional[Tuple[int, int, int, int]] = None def __eq__(self, other): """ @@ -440,6 +443,7 @@ class FusedAttentionParams: # None preserves the training-mode default for direct support queries. requires_backward: Optional[bool] = None + bias_stride: Optional[Tuple[int, int, int, int]] = None class _NoOpLogger: @@ -1700,6 +1704,12 @@ def _is_fa3_supported(num_heads, num_gqa_groups, head_dim_qk, head_dim_v, qkv_dt "bias_num_heads": bias_num_heads, "bias_seqlen_q": bias_seqlen_q, "bias_seqlen_kv": bias_seqlen_kv, + # CP steps materialize contiguous bias slices; ALiBi conversion does too. + "bias_stride": ( + attention_params.core_attention_bias_stride + if not context_parallel and fu_core_attention_bias_type == core_attention_bias_type + else None + ), } # Context-parallel per-step configs if context_parallel: diff --git a/transformer_engine/pytorch/cpp_extensions/cudnn_fused_attn.py b/transformer_engine/pytorch/cpp_extensions/cudnn_fused_attn.py index 8dbae20a3d7..a23bcaddf6f 100644 --- a/transformer_engine/pytorch/cpp_extensions/cudnn_fused_attn.py +++ b/transformer_engine/pytorch/cpp_extensions/cudnn_fused_attn.py @@ -600,7 +600,7 @@ def scale(n, matrix_name, fmt, columnwise=False): ) b, h, sq, skv = dim stride = (h * sq * skv, sq * skv, skv, 1) - spec("bias", dim, stride, types[DType.cast(config.qkv_dtype)]) + spec("bias", dim, config.bias_stride or 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": @@ -1108,7 +1108,8 @@ def fused_attn_bwd( 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"]) + # cuDNN requires row-major dBias even when the bias input is a strided view. + dbias = torch.empty_like(p["attn_bias"], memory_format=torch.contiguous_format) 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) From b65fd239833a336321acd6bc1198b51867f84443 Mon Sep 17 00:00:00 2001 From: Vladimir Cherepanov Date: Thu, 12 Nov 2026 00:03:16 +0000 Subject: [PATCH 39/49] Remove MXFP8 support-surface and scale-byte tests Remove the explicit architecture/version support matrix and the quantization/scale-byte test requested in review. Retain the remaining scale-layout, numerical, and integration coverage. Validation: 31 focused MXFP8 tests passed; pre-commit and the changed-file license check passed. Signed-off-by: Vladimir Cherepanov --- tests/jax/test_mxfp8_attention.py | 58 ------------------------------- 1 file changed, 58 deletions(-) diff --git a/tests/jax/test_mxfp8_attention.py b/tests/jax/test_mxfp8_attention.py index 5a34081c258..9b849b30153 100644 --- a/tests/jax/test_mxfp8_attention.py +++ b/tests/jax/test_mxfp8_attention.py @@ -44,26 +44,6 @@ def test_scale_reordering_matches_graph(dim, columnwise): assert sorted(permutation) == list(range(len(tiled))) -@pytest.mark.parametrize( - "sm,version,fmt,error", - ( - (90, 92600, "bshd", "SM100"), - (120, 92600, "bshd", "SM12x"), - (121, 92600, "bshd", "SM12x"), - (100, 92000, "bshd", "9.21"), - (100, 92300, "bshd", "defects"), - (100, 92301, "bshd", "defects"), - (100, 92600, "thd", "THD"), - (100, 92100, "bshd", ""), - (100, 92600, "sbhd", ""), - ), -) -def test_shared_mxfp8_support(sm, version, fmt, error): - """Both frontends retain the same hardware and known-defect exclusions.""" - reason = mxfp8.support_error(sm, version, fmt) - assert error in reason if error else not reason - - def test_quantization_directions(): """Forward V contracts over sequence; backward Q/K need both directions.""" assert mxfp8.quantization_layouts() == ((True, False), (True, False), (False, True)) @@ -139,44 +119,6 @@ def test_dense_forward_backward(dtype, fp8_backward, is_training, monkeypatch): np.testing.assert_allclose(actual_grad, expected_grad, rtol=1e-2, atol=1e-3) -@pytest.mark.parametrize("columnwise", (False, True)) -@pytest.mark.parametrize("head_dim", (64, 96, 160)) -def test_quantization_and_scale_bytes(columnwise, head_dim): - """Native quantization and cuDNN scale packing agree with an independent reference.""" - _require_mxfp8() - rng = np.random.default_rng(17) - value = jnp.asarray(rng.normal(size=(2, 3, 128, head_dim)), jnp.bfloat16) - value *= jnp.asarray([0.01, 1.0, 100.0], jnp.bfloat16)[None, :, None, None] - row, col = jax.jit(lambda x: mxfp8_attention.quantize_operand(x, not columnwise, columnwise))( - value - ) - tensor = col if columnwise else row - source = np.asarray(value, np.float32) - oriented = source.swapaxes(-1, -2) if columnwise else source - blocks = oriented.reshape(*oriented.shape[:-1], -1, 32) - maximum = np.max(np.abs(blocks), axis=-1) - exponent = np.maximum(np.ceil(np.log2(np.maximum(maximum / 448.0, 2.0**-127))), -127) - scales = np.exp2(exponent) - encoded = (blocks / scales[..., None]).astype(jnp.float8_e4m3fn).reshape(oriented.shape) - if columnwise: - encoded = encoded.swapaxes(-1, -2) - np.testing.assert_array_equal(np.asarray(tensor.data).view(np.uint8), encoded.view(np.uint8)) - scales = (exponent + 127).astype(np.uint8) - rows, cols = scales.shape[-2:] - padded_rows = ((rows + 127) // 128) * 128 - padded_cols = ((cols + 3) // 4) * 4 - expected = np.zeros((*scales.shape[:2], padded_rows * padded_cols), np.uint8) - # Physical offset within each 128x4 tile, derived independently of the - # reshape/transpose used by the implementation. - for i in range(rows): - for j in range(cols): - index = ((i // 128) * (padded_cols // 4) + j // 4) * 512 - index += (i % 32) * 16 + ((i % 128) // 32) * 4 + j % 4 - expected[..., index] = scales[..., i, j] - actual = jax.jit(lambda s: mxfp8_attention.reorder_scale(s, columnwise))(tensor.scale_inv) - np.testing.assert_array_equal(np.asarray(actual).reshape(expected.shape), expected) - - @pytest.mark.parametrize("deterministic", (False, True)) def test_flax_mxfp8_mode(deterministic): """Flax forwards the requested block-scaling mode without coercing it to bool.""" From 9231140bc952a38933687389f7f038b8b5415d46 Mon Sep 17 00:00:00 2001 From: Vladimir Cherepanov Date: Thu, 12 Nov 2026 00:10:07 +0000 Subject: [PATCH 40/49] Remove the empty ragged-capacity descriptor test Remove the private descriptor-bucketing test requested in review. Retain real FFI numerical coverage for ragged padding and gradients. Validation: six ragged-padding cases passed; pre-commit and the changed-file license check passed. Signed-off-by: Vladimir Cherepanov --- tests/jax/test_fused_attn.py | 31 ------------------------------- 1 file changed, 31 deletions(-) diff --git a/tests/jax/test_fused_attn.py b/tests/jax/test_fused_attn.py index ed8f69ad615..008825449f7 100644 --- a/tests/jax/test_fused_attn.py +++ b/tests/jax/test_fused_attn.py @@ -525,37 +525,6 @@ def test_cudnn_attention_backward_o_precision_cache(fp8, fp8_backward, monkeypat assert backward[0] is backward[1] -@pytest.mark.parametrize("batch,seqlen", ((0, 128), (2, 0))) -def test_cudnn_attention_empty_ragged_capacities(batch, seqlen, monkeypatch): - """Empty storage must retain positive graph capacities for ragged metadata.""" - from types import SimpleNamespace - from transformer_engine.jax.cpp_extensions import cudnn_attention - - monkeypatch.setattr(cudnn_attention, "cudnn_version", lambda: 92700) - config = SimpleNamespace( - qkv_layout=QKVLayout.THD_THD_THD, - max_segments_per_seq=1, - fp8=False, - attn_mask_type=AttnMaskType.PADDING_MASK, - attn_bias_type=AttnBiasType.NO_BIAS, - softmax_type=AttnSoftmaxType.VANILLA_SOFTMAX, - effective_window_size=(-1, -1), - scaling_factor=0.125, - bottom_right_diagonal=False, - dropout_probability=0.0, - is_training=False, - return_max_logit=False, - ) - shapes = tuple( - (shape, jnp.dtype(jnp.bfloat16)) - for shape in ((batch, seqlen, 4, 64),) * 3 + ((0,), (1, 4, 1, 1)) - ) - description = cudnn_attention._describe_graph("fwd", shapes, config, (100,), True, True) - assert description.metadata_batch == 32 - for name in ("q", "k", "v"): - assert dict(description.spec.tensors)[name].dim == (32, 4, 1024, 64) - - @pytest.mark.parametrize("sm_arch", (100, 120, 121, 129, 130)) @pytest.mark.parametrize("direction", ("fwd", "bwd")) def test_cudnn_attention_sm12x_dense_statistics(sm_arch, direction, monkeypatch): From b52a40c888f0c46f890d096ab0efd12f234b229d Mon Sep 17 00:00:00 2001 From: Vladimir Cherepanov Date: Thu, 12 Nov 2026 00:22:19 +0000 Subject: [PATCH 41/49] Group JAX distributed attention tests Move F16/CS context-parallel and MXFP8 sharding/communication tests into the existing distributed attention suite. L0's selection now excludes those cases and the L1 launcher collects them, with all numerical statements and parameter matrices preserved. Remove the legacy common-API routing and redundant Flax MXFP8 mode forwarding tests. Reuse the existing numerical helpers without adding another launcher or test infrastructure. Validation: 47 distributed and 92 single-GPU cases passed. Confirmed L0 excludes the relocated cases and L1 collects all 47. Pre-commit and changed-file license checks passed. Signed-off-by: Vladimir Cherepanov --- tests/jax/test_distributed_fused_attn.py | 204 +++++++++++++++++++++++ tests/jax/test_fused_attn.py | 118 ------------- tests/jax/test_mxfp8_attention.py | 154 ----------------- 3 files changed, 204 insertions(+), 272 deletions(-) diff --git a/tests/jax/test_distributed_fused_attn.py b/tests/jax/test_distributed_fused_attn.py index e9ae8cd3e4a..9bd3d819cb7 100644 --- a/tests/jax/test_distributed_fused_attn.py +++ b/tests/jax/test_distributed_fused_attn.py @@ -6,6 +6,8 @@ import pytest import jax import jax.numpy as jnp +import numpy as np +from jax.sharding import Mesh, NamedSharding, PartitionSpec from jax import random from distributed_test_base import ( generate_configs, @@ -14,6 +16,7 @@ ) from test_fused_attn import ( FusedAttnRunner, + _require_cudnn_attention, BiasShape, SeqDescFormat, ) @@ -21,10 +24,21 @@ ScoreModFusedAttnRunner, _has_cudnn_frontend_python, ) +from test_mxfp8_attention import ( + _attention as _mxfp8_attention, + _reference as _mxfp8_reference, + _require_mxfp8, + _pack as _pack_mxfp8, + _unpack as _unpack_mxfp8, +) from utils import pytest_parametrize_wrapper +from transformer_engine.jax import autocast +from transformer_engine.jax.sharding import MeshResource, global_shard_guard from transformer_engine_jax import get_cudnn_version, get_device_compute_capability from transformer_engine.jax.attention import ( is_fused_attn_kernel_available, + fused_attn, + quantize_attention_qkv, AttnBiasType, AttnMaskType, AttnSoftmaxType, @@ -919,3 +933,193 @@ def test(self, cp_size, shape, qkv_format, reorder_strategy, stripe_size): inversed = inverse(reordered, reorder_strategy, cp_size, seq_dim, stripe_size) assert jnp.array_equal(inversed, ref) + + +@pytest.mark.parametrize("o_in_f16", [True, False]) +@pytest.mark.parametrize("strategy", [CPStrategy.ALL_GATHER, CPStrategy.RING]) +@pytest.mark.parametrize("fp8", [False, True]) +def test_cudnn_python_context_parallel(strategy, fp8, o_in_f16, monkeypatch): + """Two-device context parallel attention preserves FP8 forward and backward.""" + monkeypatch.setenv("NVTE_DPA_FP8CS_O_in_F16", "1" if o_in_f16 else "0") + if jax.local_device_count() < 2: + pytest.skip("Context parallel attention requires two GPUs.") + 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) + if fp8: + # Query the per-device graph: all-gather retains the global KV sequence, + # while ring attention uses local Q and KV blocks. + _require_cudnn_attention( + q, + q, + v, + fp8=True, + qkv_layout=( + QKVLayout.BSHD_BSHD_BSHD + if strategy == CPStrategy.ALL_GATHER + else QKVLayout.BSHD_BS2HD + ), + q_max_seqlen=q.shape[1] // 2, + kv_max_seqlen=q.shape[1] if strategy == CPStrategy.ALL_GATHER else q.shape[1] // 2, + ) + 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) + + +@pytest.mark.parametrize("axis", ("dp", "tp", "replica")) +def test_mxfp8_prequantized_sharding(axis): + """Per-head scale packing happens after data/tensor parallel partitioning.""" + _require_mxfp8() + if jax.local_device_count() < 2: + pytest.skip("Requires two GPUs for attention sharding") + mesh = jax.sharding.Mesh(np.array(jax.local_devices()[:2]), (axis,)) + spec = ( + (axis, None, None, None) + if axis == "dp" + else (None, None, axis, None) if axis == "tp" else () + ) + sharding = jax.sharding.NamedSharding(mesh, jax.sharding.PartitionSpec(*spec)) + q = jnp.full((2, 128, 4, 64), 0.5, jnp.bfloat16) + kv = q[:, :, :2] + v = jnp.zeros_like(kv).at[:, ::2, :, 0].set(1.0) + inputs = (q, kv, v) + loss = lambda xs: _mxfp8_attention(xs, quantized_qkv=quantize_attention_qkv(xs)).sum() + expected = jax.jit(jax.value_and_grad(loss))(inputs) + resource = MeshResource() if axis == "replica" else MeshResource(**{axis + "_resource": axis}) + with jax.set_mesh(mesh), global_shard_guard(resource): + actual = jax.jit(jax.value_and_grad(loss), in_shardings=((sharding,) * 3,))(inputs) + for value, ref in zip(jax.tree.leaves(actual), jax.tree.leaves(expected)): + np.testing.assert_array_equal(value, ref) + + +@pytest.mark.parametrize("strategy", (CPStrategy.ALL_GATHER, CPStrategy.RING)) +@pytest.mark.parametrize("cp_size", (2, 4)) +@pytest.mark.parametrize("kind", ("dense", "causal", "balanced_causal")) +@pytest.mark.parametrize("fp8_backward", (False, True)) +def test_mxfp8_context_parallel(strategy, cp_size, kind, fp8_backward, monkeypatch): + """Context-parallel outputs and gradients agree with an independent global reference.""" + from transformer_engine.jax.cpp_extensions.attention import reorder_causal_dual_chunk_swap + + _require_mxfp8() + if jax.local_device_count() < cp_size: + pytest.skip(f"Requires {cp_size} GPUs for context parallelism") + monkeypatch.setenv("NVTE_FP8_DPA_BWD", "1" if fp8_backward else "0") + mesh = jax.sharding.Mesh(np.asarray(jax.local_devices()[:cp_size]), ("cp",)) + sharding = jax.sharding.NamedSharding(mesh, jax.sharding.PartitionSpec(None, "cp", None, None)) + length = 256 * cp_size + q = jnp.full((1, length, 4, 64), 0.5, jnp.bfloat16) + k = q[:, :, :2] + v = jnp.zeros_like(k).at[:, ::2, :, 0].set(1.0) + # Distinct rank-local values make missing communication or incorrect + # dual-chunk ordering observable, while retaining exact MXFP8 operands. + amplitudes = jnp.exp2((jnp.arange(length) // 128) % 4 - 2).astype(v.dtype) + v *= amplitudes[None, :, None, None] + inputs = (q, k, v) + balanced = kind == "balanced_causal" + causal = kind != "dense" + mask = ( + jnp.arange(length)[None, :, None] >= jnp.arange(length)[None, None, :] if causal else None + ) + reference = lambda xs: _mxfp8_reference(xs, mask) + expected = jax.jit(reference)(inputs) + ref_grads = jax.jit(jax.grad(lambda xs: reference(xs).sum()))(inputs) + reorder = lambda x, reverse: reorder_causal_dual_chunk_swap(x, cp_size, 1, reverse) + if balanced: + inputs = tuple(reorder(x, False) for x in inputs) + attention = lambda xs: _mxfp8_attention( + xs, + attn_mask_type=AttnMaskType.CAUSAL_MASK if causal else AttnMaskType.NO_MASK, + context_parallel_axis="cp", + context_parallel_strategy=strategy, + context_parallel_causal_load_balanced=balanced, + ) + with jax.set_mesh(mesh), global_shard_guard(MeshResource(cp_resource="cp")): + actual = jax.jit(attention, in_shardings=((sharding,) * 3,))(inputs) + grads = jax.jit(jax.grad(lambda xs: attention(xs).sum()), in_shardings=((sharding,) * 3,))( + inputs + ) + if balanced: + actual = reorder(actual, True) + grads = tuple(reorder(x, True) for x in grads) + np.testing.assert_allclose(actual, expected, rtol=1e-2, atol=1e-3) + for grad, ref in zip(grads, ref_grads): + np.testing.assert_allclose(grad, ref, rtol=3e-2, atol=2e-2) + + +@pytest.mark.parametrize("strategy", (CPStrategy.ALL_GATHER, CPStrategy.RING)) +@pytest.mark.parametrize( + "layout", (QKVLayout.BS3HD, QKVLayout.BSHD_BS2HD, QKVLayout.BSHD_BSHD_BSHD) +) +@pytest.mark.parametrize("prequantized", (False, True)) +def test_mxfp8_context_parallel_storage(strategy, layout, prequantized): + """Packed and supplied FP8 inputs preserve nonuniform cross-rank values and gradients.""" + _require_mxfp8() + if jax.local_device_count() < 2: + pytest.skip("Requires two GPUs for context parallelism") + mesh = jax.sharding.Mesh(np.asarray(jax.local_devices()[:2]), ("cp",)) + sharding = jax.sharding.NamedSharding(mesh, jax.sharding.PartitionSpec(None, "cp", None, None)) + rng = np.random.default_rng(42) + q = jnp.asarray(rng.choice([-0.5, 0.5], (1, 512, 2, 64)), jnp.bfloat16) + k = jnp.asarray(rng.choice([-0.5, 0.5], q.shape), q.dtype) + v = jnp.asarray(rng.choice([-1.0, 1.0], q.shape), q.dtype) + amplitude = jnp.exp2(jnp.arange(512) // 128 - 2).astype(q.dtype) + k *= amplitude[None, :, None, None] + v *= amplitude[None, :, None, None] + weight = jnp.asarray(rng.choice([-1.0, 1.0], q.shape), q.dtype) + inputs = _pack_mxfp8((q, k, v), layout) + reference = lambda xs: _mxfp8_reference(_unpack_mxfp8(xs, layout)) + expected = jax.jit(reference)(inputs) + ref_grads = jax.jit(jax.grad(lambda xs: (reference(xs) * weight).sum()))(inputs) + + def attention(xs): + supplied = quantize_attention_qkv(_unpack_mxfp8(xs, layout)) if prequantized else None + return _mxfp8_attention( + xs, + qkv_layout=layout, + quantized_qkv=supplied, + context_parallel_axis="cp", + context_parallel_strategy=strategy, + ) + + with jax.set_mesh(mesh), global_shard_guard(MeshResource(cp_resource="cp")): + actual = jax.jit(attention, in_shardings=((sharding,) * len(inputs),))(inputs) + grads = jax.jit( + jax.grad(lambda xs: (attention(xs) * weight).sum()), + in_shardings=((sharding,) * len(inputs),), + )(inputs) + relative_error = lambda a, b: float(jnp.linalg.norm(a - b) / jnp.linalg.norm(b)) + assert relative_error(actual, expected) < 0.04 + for grad, ref in zip(grads, ref_grads): + assert relative_error(grad, ref) < 0.08 diff --git a/tests/jax/test_fused_attn.py b/tests/jax/test_fused_attn.py index 008825449f7..69b77cc9cc5 100644 --- a/tests/jax/test_fused_attn.py +++ b/tests/jax/test_fused_attn.py @@ -2563,61 +2563,6 @@ def attention(operands): np.testing.assert_allclose(actual_grad, expected_grad, rtol=1e-2, atol=1e-2) -def test_cudnn_python_attention_does_not_call_common_attention(monkeypatch): - """Backend selection, workspace sizing and execution bypass the common attention APIs.""" - import transformer_engine_jax - from transformer_engine.jax.cpp_extensions import cudnn_attention - - def forbidden(*_args, **_kwargs): - raise AssertionError("JAX attention called a common attention API") - - 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("o_in_f16", [True, False]) @pytest.mark.parametrize("scaling_factor", [0.125, 0.25]) @pytest.mark.parametrize( @@ -2902,66 +2847,3 @@ def attention(arrays, seed): 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("o_in_f16", [True, False]) -@pytest.mark.parametrize("strategy", [CPStrategy.ALL_GATHER, CPStrategy.RING]) -@pytest.mark.parametrize("fp8", [False, True]) -def test_cudnn_python_context_parallel(strategy, fp8, o_in_f16, monkeypatch): - """Two-device context parallel attention preserves FP8 forward and backward.""" - monkeypatch.setenv("NVTE_DPA_FP8CS_O_in_F16", "1" if o_in_f16 else "0") - if jax.local_device_count() < 2: - pytest.skip("Context parallel attention requires two GPUs.") - 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) - if fp8: - # Query the per-device graph: all-gather retains the global KV sequence, - # while ring attention uses local Q and KV blocks. - _require_cudnn_attention( - q, - q, - v, - fp8=True, - qkv_layout=( - QKVLayout.BSHD_BSHD_BSHD - if strategy == CPStrategy.ALL_GATHER - else QKVLayout.BSHD_BS2HD - ), - q_max_seqlen=q.shape[1] // 2, - kv_max_seqlen=q.shape[1] if strategy == CPStrategy.ALL_GATHER else q.shape[1] // 2, - ) - 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_mxfp8_attention.py b/tests/jax/test_mxfp8_attention.py index 9b849b30153..bc37e2a5c01 100644 --- a/tests/jax/test_mxfp8_attention.py +++ b/tests/jax/test_mxfp8_attention.py @@ -17,7 +17,6 @@ AttnBiasType, AttnMaskType, AttnSoftmaxType, - CPStrategy, QKVLayout, SequenceDescriptor, fused_attn, @@ -119,32 +118,6 @@ def test_dense_forward_backward(dtype, fp8_backward, is_training, monkeypatch): np.testing.assert_allclose(actual_grad, expected_grad, rtol=1e-2, atol=1e-3) -@pytest.mark.parametrize("deterministic", (False, True)) -def test_flax_mxfp8_mode(deterministic): - """Flax forwards the requested block-scaling mode without coercing it to bool.""" - from transformer_engine.jax.flax import DotProductAttention - - _require_mxfp8() - module = DotProductAttention( - head_dim=64, - num_attention_heads=2, - attn_mask_type="no_mask", - transpose_batch_sequence=False, - fp8="mxfp8", - ) - q = jnp.full((1, 128, 2, 64), 0.5, jnp.bfloat16) - value = jnp.zeros_like(q).at[:, ::2, :, 0].set(1.0) - variables = module.init( - jax.random.key(0), q, q, value, SequenceDescriptor(), deterministic=deterministic - ) - forward = lambda v: module.apply( - variables, q, q, v, SequenceDescriptor(), deterministic=deterministic - ) - actual = jax.jit(forward)(value) - np.testing.assert_allclose(actual[..., 0], 0.5, rtol=5e-3, atol=5e-4) - np.testing.assert_allclose(jax.jit(jax.grad(lambda v: forward(v).sum()))(value), 1.0) - - def _pack(qkv, layout): q, k, v = qkv if layout.is_qkvpacked(): @@ -353,32 +326,6 @@ def fail_requantization(*args, **kwargs): np.testing.assert_array_equal(value, ref) -@pytest.mark.parametrize("axis", ("dp", "tp", "replica")) -def test_prequantized_sharding(axis): - """Per-head scale packing happens after data/tensor parallel partitioning.""" - _require_mxfp8() - if jax.local_device_count() < 2: - pytest.skip("Requires two GPUs for attention sharding") - mesh = jax.sharding.Mesh(np.array(jax.local_devices()[:2]), (axis,)) - spec = ( - (axis, None, None, None) - if axis == "dp" - else (None, None, axis, None) if axis == "tp" else () - ) - sharding = jax.sharding.NamedSharding(mesh, jax.sharding.PartitionSpec(*spec)) - q = jnp.full((2, 128, 4, 64), 0.5, jnp.bfloat16) - kv = q[:, :, :2] - v = jnp.zeros_like(kv).at[:, ::2, :, 0].set(1.0) - inputs = (q, kv, v) - loss = lambda xs: _attention(xs, quantized_qkv=quantize_attention_qkv(xs)).sum() - expected = jax.jit(jax.value_and_grad(loss))(inputs) - resource = MeshResource() if axis == "replica" else MeshResource(**{axis + "_resource": axis}) - with jax.set_mesh(mesh), global_shard_guard(resource): - actual = jax.jit(jax.value_and_grad(loss), in_shardings=((sharding,) * 3,))(inputs) - for value, ref in zip(jax.tree.leaves(actual), jax.tree.leaves(expected)): - np.testing.assert_array_equal(value, ref) - - def test_prequantized_values_and_validation(): """Supplied FP8 values determine the output; incompatible modes fail at the API boundary.""" _require_mxfp8() @@ -402,107 +349,6 @@ def test_prequantized_values_and_validation(): _attention((q, q, v), qkv_layout=QKVLayout.THD_THD_THD, context_parallel_axis="cp") -@pytest.mark.parametrize("strategy", (CPStrategy.ALL_GATHER, CPStrategy.RING)) -@pytest.mark.parametrize("cp_size", (2, 4)) -@pytest.mark.parametrize("kind", ("dense", "causal", "balanced_causal")) -@pytest.mark.parametrize("fp8_backward", (False, True)) -def test_context_parallel(strategy, cp_size, kind, fp8_backward, monkeypatch): - """Context-parallel outputs and gradients agree with an independent global reference.""" - from transformer_engine.jax.cpp_extensions.attention import reorder_causal_dual_chunk_swap - - _require_mxfp8() - if jax.local_device_count() < cp_size: - pytest.skip(f"Requires {cp_size} GPUs for context parallelism") - monkeypatch.setenv("NVTE_FP8_DPA_BWD", "1" if fp8_backward else "0") - mesh = jax.sharding.Mesh(np.asarray(jax.local_devices()[:cp_size]), ("cp",)) - sharding = jax.sharding.NamedSharding(mesh, jax.sharding.PartitionSpec(None, "cp", None, None)) - length = 256 * cp_size - q = jnp.full((1, length, 4, 64), 0.5, jnp.bfloat16) - k = q[:, :, :2] - v = jnp.zeros_like(k).at[:, ::2, :, 0].set(1.0) - # Distinct rank-local values make missing communication or incorrect - # dual-chunk ordering observable, while retaining exact MXFP8 operands. - amplitudes = jnp.exp2((jnp.arange(length) // 128) % 4 - 2).astype(v.dtype) - v *= amplitudes[None, :, None, None] - inputs = (q, k, v) - balanced = kind == "balanced_causal" - causal = kind != "dense" - mask = ( - jnp.arange(length)[None, :, None] >= jnp.arange(length)[None, None, :] if causal else None - ) - reference = lambda xs: _reference(xs, mask) - expected = jax.jit(reference)(inputs) - ref_grads = jax.jit(jax.grad(lambda xs: reference(xs).sum()))(inputs) - reorder = lambda x, reverse: reorder_causal_dual_chunk_swap(x, cp_size, 1, reverse) - if balanced: - inputs = tuple(reorder(x, False) for x in inputs) - attention = lambda xs: _attention( - xs, - attn_mask_type=AttnMaskType.CAUSAL_MASK if causal else AttnMaskType.NO_MASK, - context_parallel_axis="cp", - context_parallel_strategy=strategy, - context_parallel_causal_load_balanced=balanced, - ) - with jax.set_mesh(mesh), global_shard_guard(MeshResource(cp_resource="cp")): - actual = jax.jit(attention, in_shardings=((sharding,) * 3,))(inputs) - grads = jax.jit(jax.grad(lambda xs: attention(xs).sum()), in_shardings=((sharding,) * 3,))( - inputs - ) - if balanced: - actual = reorder(actual, True) - grads = tuple(reorder(x, True) for x in grads) - np.testing.assert_allclose(actual, expected, rtol=1e-2, atol=1e-3) - for grad, ref in zip(grads, ref_grads): - np.testing.assert_allclose(grad, ref, rtol=3e-2, atol=2e-2) - - -@pytest.mark.parametrize("strategy", (CPStrategy.ALL_GATHER, CPStrategy.RING)) -@pytest.mark.parametrize( - "layout", (QKVLayout.BS3HD, QKVLayout.BSHD_BS2HD, QKVLayout.BSHD_BSHD_BSHD) -) -@pytest.mark.parametrize("prequantized", (False, True)) -def test_context_parallel_storage(strategy, layout, prequantized): - """Packed and supplied FP8 inputs preserve nonuniform cross-rank values and gradients.""" - _require_mxfp8() - if jax.local_device_count() < 2: - pytest.skip("Requires two GPUs for context parallelism") - mesh = jax.sharding.Mesh(np.asarray(jax.local_devices()[:2]), ("cp",)) - sharding = jax.sharding.NamedSharding(mesh, jax.sharding.PartitionSpec(None, "cp", None, None)) - rng = np.random.default_rng(42) - q = jnp.asarray(rng.choice([-0.5, 0.5], (1, 512, 2, 64)), jnp.bfloat16) - k = jnp.asarray(rng.choice([-0.5, 0.5], q.shape), q.dtype) - v = jnp.asarray(rng.choice([-1.0, 1.0], q.shape), q.dtype) - amplitude = jnp.exp2(jnp.arange(512) // 128 - 2).astype(q.dtype) - k *= amplitude[None, :, None, None] - v *= amplitude[None, :, None, None] - weight = jnp.asarray(rng.choice([-1.0, 1.0], q.shape), q.dtype) - inputs = _pack((q, k, v), layout) - reference = lambda xs: _reference(_unpack(xs, layout)) - expected = jax.jit(reference)(inputs) - ref_grads = jax.jit(jax.grad(lambda xs: (reference(xs) * weight).sum()))(inputs) - - def attention(xs): - supplied = quantize_attention_qkv(_unpack(xs, layout)) if prequantized else None - return _attention( - xs, - qkv_layout=layout, - quantized_qkv=supplied, - context_parallel_axis="cp", - context_parallel_strategy=strategy, - ) - - with jax.set_mesh(mesh), global_shard_guard(MeshResource(cp_resource="cp")): - actual = jax.jit(attention, in_shardings=((sharding,) * len(inputs),))(inputs) - grads = jax.jit( - jax.grad(lambda xs: (attention(xs) * weight).sum()), - in_shardings=((sharding,) * len(inputs),), - )(inputs) - relative_error = lambda a, b: float(jnp.linalg.norm(a - b) / jnp.linalg.norm(b)) - assert relative_error(actual, expected) < 0.04 - for grad, ref in zip(grads, ref_grads): - assert relative_error(grad, ref) < 0.08 - - def test_prequantized_vmap_and_checkpoint(): """Independent mapped attention calls retain scale boundaries and model gradients.""" _require_mxfp8() From 06536d6376ca8b5b7e59aadc7b6ea7a1455f1471 Mon Sep 17 00:00:00 2001 From: Vladimir Cherepanov Date: Thu, 12 Nov 2026 00:30:51 +0000 Subject: [PATCH 42/49] Share cuDNN attention caches across SDPA paths Use a frontend-owned common cache for ordinary and score-modified attention. Share synchronized graph lookup, compile-once state, and optional serialization while retaining callback identity, unsafe-callback bypass, and engine eligibility rules. Consolidate PyTorch score-modified key, build, and lookup helpers around AttentionGraphSpec and a shared live entry type. Preserve JAX serialized results and request-specific binding memoization. Validation: 257 focused PyTorch checks passed with 12 existing skips. All 51 Flex Attention tests passed again after the final fixture edits. Repository pre-commit, both Python lint launchers, syntax, whitespace, and changed-file license checks passed. JAX runtime validation is unavailable in this container. Signed-off-by: Vladimir Cherepanov --- tests/jax/test_fused_attn.py | 2 +- tests/pytorch/attention/test_attention.py | 6 +- .../pytorch/attention/test_flex_attention.py | 408 +++++++--------- transformer_engine/common/cudnn/graph.py | 46 ++ .../jax/cpp_extensions/cudnn_attention.py | 58 +-- .../jax/cpp_extensions/cudnn_graph.py | 3 + .../jax/cpp_extensions/flex_attention.py | 34 +- .../dot_product_attention/flex_attention.py | 436 ++++-------------- .../pytorch/cpp_extensions/cudnn.py | 2 + .../cpp_extensions/cudnn_fused_attn.py | 54 +-- 10 files changed, 378 insertions(+), 671 deletions(-) diff --git a/tests/jax/test_fused_attn.py b/tests/jax/test_fused_attn.py index 69b77cc9cc5..ee0bcc3e336 100644 --- a/tests/jax/test_fused_attn.py +++ b/tests/jax/test_fused_attn.py @@ -437,7 +437,7 @@ 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() + monkeypatch.setattr(cudnn_attention._graph_cache, "entries", {}) if setup_error is not None: def fail(*_args, **_kwargs): diff --git a/tests/pytorch/attention/test_attention.py b/tests/pytorch/attention/test_attention.py index 159c1d9cde1..6d6dd31f040 100644 --- a/tests/pytorch/attention/test_attention.py +++ b/tests/pytorch/attention/test_attention.py @@ -452,7 +452,7 @@ def test_fused_attn_bias_graph_cache(dtype, qkv_format, bias_layout, monkeypatch monkeypatch.setenv("NVTE_FLASH_ATTN", "0") monkeypatch.setenv("NVTE_FUSED_ATTN", "1") monkeypatch.setenv("NVTE_UNFUSED_ATTN", "0") - monkeypatch.setattr(implementation, "_plans", {}) + monkeypatch.setattr(implementation._graph_cache, "entries", {}) compile_graph = Mock(wraps=implementation.compile_graph) monkeypatch.setattr(implementation, "compile_graph", compile_graph) original_entry = implementation._AttentionGraph.entry @@ -520,7 +520,7 @@ def entry(attn): # dBias subtracts dO.O using low-precision saved O, amplifying pointwise cancellation. error = torch.linalg.vector_norm(bias.grad.float() - ref_bias.grad) assert error <= torch.finfo(dtype).eps * torch.linalg.vector_norm(ref_bias.grad) - assert len(implementation._plans) == 2 + assert len(implementation._graph_cache.entries) == 2 assert compile_graph.call_count == 2 assert executed == [False, True, False, True] @@ -745,7 +745,7 @@ def fail(_self): raise error monkeypatch.setattr(implementation._AttentionGraph, "build", fail) - monkeypatch.setattr(implementation, "_plans", {}) + monkeypatch.setattr(implementation._graph_cache, "entries", {}) config = FusedAttentionParams( qkv_layout=QKVLayout["bshd_bshd_bshd"], batch_size=2, diff --git a/tests/pytorch/attention/test_flex_attention.py b/tests/pytorch/attention/test_flex_attention.py index c6e8d7775fe..68de5f86848 100644 --- a/tests/pytorch/attention/test_flex_attention.py +++ b/tests/pytorch/attention/test_flex_attention.py @@ -5,12 +5,20 @@ import os import sys import pathlib -from types import FunctionType +from types import FunctionType, SimpleNamespace from unittest.mock import Mock, call +from concurrent.futures import ThreadPoolExecutor +import time import pytest import torch +from transformer_engine.common.cudnn.attention import ( + AttentionOptions, + AttentionGraphSpec, + BuiltAttentionGraph, +) +from transformer_engine.common.cudnn.graph import GraphCache, GraphCacheEntry from transformer_engine.pytorch import DotProductAttention, is_bf16_available from transformer_engine.pytorch.attention.dot_product_attention import _attention_backends import transformer_engine.pytorch.attention.dot_product_attention.flex_attention as flex_attention @@ -60,6 +68,59 @@ def test_cudnn_score_mod_backend_guard(monkeypatch, has_score_mod, selected_engi graph.deselect_engines.assert_not_called() +@pytest.mark.parametrize("serialize", (False, True)) +def test_cudnn_graph_cache_compiles_once_across_threads(serialize): + """Concurrent requests share one supported graph and one completed compilation.""" + cache = GraphCache() + graph = Mock() + graph.serialize.return_value = b"graph" + + def build(): + time.sleep(0.01) + return GraphCacheEntry(BuiltAttentionGraph(graph, {}, {})) + + def compile_plan(): + time.sleep(0.01) + return 64 + + build = Mock(side_effect=build) + compile_plan = Mock(side_effect=compile_plan) + + def request(_index): + entry = cache.get("graph", build) + assert cache.compile(entry, compile_plan, serialize=serialize) == 64 + return entry + + with ThreadPoolExecutor(max_workers=4) as pool: + entries = list(pool.map(request, range(8))) + assert all(entry is entries[0] for entry in entries) + assert build.call_count == compile_plan.call_count == 1 + assert graph.serialize.call_count == int(serialize) + + +@pytest.mark.parametrize("failure", ("support", "compile", "serialize")) +def test_cudnn_graph_cache_retries_failed_build(failure): + """Failures must not publish an incomplete cache entry or compilation state.""" + cache = GraphCache() + graph = Mock() + graph.serialize.return_value = b"graph" + entry = GraphCacheEntry(BuiltAttentionGraph(graph, {}, {})) + build = Mock(side_effect=[RuntimeError("failure"), entry] if failure == "support" else [entry]) + compile_plan = Mock( + side_effect=[RuntimeError("failure"), 64] if failure == "compile" else [64, 64] + ) + if failure == "serialize": + graph.serialize.side_effect = [RuntimeError("failure"), b"graph"] + with pytest.raises(RuntimeError, match="failure"): + cached = cache.get("graph", build) + cache.compile(cached, compile_plan, serialize=True) + assert entry.workspace_size is entry.serialized_graph is None + assert bool(cache.entries) == (failure != "support") + cached = cache.get("graph", build) + assert cache.compile(cached, compile_plan, serialize=True) == 64 + assert cached.serialized_graph == b"graph" + + def _score_mod_causal(score_mod_graph, score_tensor, tensors): """cuDNN frontend score_mod implementing top-left causal masking.""" cudnn = flex_attention._import_cudnn_frontend() @@ -201,6 +262,41 @@ def _clone_score_mod(default=False, options=None): return callback +def _mock_score_mod_build(monkeypatch): + """Replace cuDNN construction while exercising the real cache and bindings.""" + monkeypatch.setattr(cudnn_graph._graph_cache, "entries", {}) + + def build(spec, _device, _callback, _bprop, extra, bprop_extra): + handles = flex_attention._score_mod_bindings( + {name: object() for name, _ in spec.tensors}, + {name: object() for name in extra}, + {name: object() for name in bprop_extra}, + ) + return GraphCacheEntry(BuiltAttentionGraph(object(), handles, {})) + + build = Mock(side_effect=build) + monkeypatch.setattr(flex_attention, "_build_cudnn_score_mod_graph", build) + monkeypatch.setattr(flex_attention, "compile_graph", Mock(return_value=1)) + return build + + +def _get_score_mod_test_graph(callback, *, backward=False, extra=None): + q, k, v, out, stats = _score_mod_cache_cpu_inputs() + tensors = {"q": q, "k": k, "v": v, "o": out, "stats": stats} + if backward: + tensors.update(do=out, dq=q, dk=k, dv=v) + entry, _ = flex_attention._get_cudnn_score_mod_graph( + AttentionOptions(backward=backward, attn_scale=1.0), + tensors, + "bshd", + "bshd", + _score_mod_identity_bprop if backward else callback, + callback if backward else None, + score_mod_tensors=extra, + ) + return entry + + @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): @@ -214,32 +310,10 @@ def test_score_mod_cache_separates_function_defaults(monkeypatch, direction, dif 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) + build = _mock_score_mod_build(monkeypatch) 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, - ) + return _get_score_mod_test_graph(callback, backward=direction == "bwd") entry = get(first) assert get(first) is entry @@ -355,137 +429,28 @@ def test_score_mod_cache_module_lambda_keys_do_not_collide(): def test_score_mod_cache_key_ignores_pass_by_value_values(): """Scalar CPU tensor values are runtime inputs, not execution-plan metadata.""" - q, k, v, o, stats = _score_mod_cache_cpu_inputs() - key_0 = flex_attention._cudnn_score_mod_fwd_cache_key( - True, - q, - k, - v, - "bshd", - "bshd", - 1.0, - _score_mod_causal, - {"softcap": torch.tensor(0.8, dtype=torch.float32)}, - o, - stats, - ) - key_1 = flex_attention._cudnn_score_mod_fwd_cache_key( - True, - q, - k, - v, - "bshd", - "bshd", - 1.0, - _score_mod_causal, - {"softcap": torch.tensor(1.2, dtype=torch.float32)}, - o, - stats, - ) - key_2 = flex_attention._cudnn_score_mod_fwd_cache_key( - True, - q, - k, - v, - "bshd", - "bshd", - 1.0, - _score_mod_causal, - {"softcap": torch.tensor([0.8], dtype=torch.float32)}, - o, - stats, - ) + spec = AttentionGraphSpec.from_dict(AttentionOptions(), {}) - assert key_0 == key_1 - assert key_0 != key_2 + def key(tensor): + return flex_attention._cudnn_score_mod_cache_key( + spec, torch.device("cpu"), _score_mod_causal, None, {"softcap": tensor}, {} + ) + + assert key(torch.tensor(0.8)) == key(torch.tensor(1.2)) + assert key(torch.tensor(0.8)) != key(torch.tensor([0.8])) def test_score_mod_cache_fwd_reuses_graph_for_pass_by_value_changes(monkeypatch): """Fprop graph cache should reuse entries when only scalar CPU tensor values change.""" - q, k, v, o, stats = _score_mod_cache_cpu_inputs() - cache = flex_attention._cudnn_score_mod_graph_cache - saved_cache = dict(cache) - build_entries = [] - - def fake_build( - is_training, - query_layer, - key_layer, - value_layer, - q_format, - kv_format, - attn_scale, - score_mod, - score_mod_tensors, - output_layer, - stats, - ): - del ( - is_training, - query_layer, - key_layer, - value_layer, - q_format, - kv_format, - attn_scale, - score_mod, - score_mod_tensors, - output_layer, - stats, - ) - entry = object() - build_entries.append(entry) - return entry - - monkeypatch.setattr(flex_attention, "_build_cudnn_score_mod_fwd_graph", fake_build) - try: - cache.clear() - entry_0 = flex_attention._get_cudnn_score_mod_fwd_graph( - True, - q, - k, - v, - "bshd", - "bshd", - 1.0, - _score_mod_causal, - {"softcap": torch.tensor(0.8, dtype=torch.float32)}, - o, - stats, - ) - entry_1 = flex_attention._get_cudnn_score_mod_fwd_graph( - True, - q, - k, - v, - "bshd", - "bshd", - 1.0, - _score_mod_causal, - {"softcap": torch.tensor(1.2, dtype=torch.float32)}, - o, - stats, - ) - entry_2 = flex_attention._get_cudnn_score_mod_fwd_graph( - True, - q, - k, - v, - "bshd", - "bshd", - 1.0, - _score_mod_causal, - {"softcap": torch.tensor([0.8], dtype=torch.float32)}, - o, - stats, - ) - finally: - cache.clear() - cache.update(saved_cache) - - assert entry_0 is entry_1 - assert entry_2 is not entry_0 - assert len(build_entries) == 2 + build = _mock_score_mod_build(monkeypatch) + entries = [ + _get_score_mod_test_graph(_score_mod_causal, extra={"softcap": torch.tensor(value)}) + for value in (0.8, 1.2, [0.8]) + ] + assert entries[0] is entries[1] + assert entries[2] is not entries[0] + assert build.call_count == 2 + assert flex_attention.compile_graph.call_count == 2 def test_score_mod_cache_fwd_skips_cache_for_unkeyed_bound_method(monkeypatch): @@ -495,91 +460,18 @@ class UnkeyedScoreMod: def forward(self, _score_mod_graph, score_tensor, _tensors): return score_tensor - q, k, v, o, stats = _score_mod_cache_cpu_inputs() - score_mod = UnkeyedScoreMod() - cache = flex_attention._cudnn_score_mod_graph_cache - saved_cache = dict(cache) - build_entries = [] - - def fake_build( - is_training, - query_layer, - key_layer, - value_layer, - q_format, - kv_format, - attn_scale, - score_mod, - score_mod_tensors, - output_layer, - stats, - ): - del ( - is_training, - query_layer, - key_layer, - value_layer, - q_format, - kv_format, - attn_scale, - score_mod, - score_mod_tensors, - output_layer, - stats, - ) - entry = object() - build_entries.append(entry) - return entry - - monkeypatch.setattr(flex_attention, "_build_cudnn_score_mod_fwd_graph", fake_build) - try: - cache.clear() - entry_0 = flex_attention._get_cudnn_score_mod_fwd_graph( - True, - q, - k, - v, - "bshd", - "bshd", - 1.0, - score_mod.forward, - None, - o, - stats, - ) - entry_1 = flex_attention._get_cudnn_score_mod_fwd_graph( - True, - q, - k, - v, - "bshd", - "bshd", - 1.0, - score_mod.forward, - None, - o, - stats, - ) - assert len(cache) == 0 - finally: - cache.clear() - cache.update(saved_cache) - - assert entry_0 is not entry_1 - assert len(build_entries) == 2 + build = _mock_score_mod_build(monkeypatch) + callback = UnkeyedScoreMod().forward + assert _get_score_mod_test_graph(callback) is not _get_score_mod_test_graph(callback) + assert not cudnn_graph._graph_cache.entries + assert build.call_count == 2 def test_score_mod_tensors_are_version_checked_for_backward(monkeypatch): """In-place score_mod tensor updates before backward should be rejected.""" class FakeEntry: - graph = object() - q = object() - k = object() - v = object() - output = object() - stats = object() - score_mod_graph_tensors = {"softcap": object()} + built = SimpleNamespace(graph=object()) workspace_size = 1 def fake_execute(graph, variant_pack, workspace_size, device): @@ -591,7 +483,9 @@ def fake_execute(graph, variant_pack, workspace_size, device): v = v.requires_grad_() softcap = torch.tensor(0.8, dtype=torch.float32) - monkeypatch.setattr(flex_attention, "_get_cudnn_score_mod_fwd_graph", lambda *args: FakeEntry()) + monkeypatch.setattr( + flex_attention, "_get_cudnn_score_mod_graph", lambda *_args, **_kwargs: (FakeEntry(), {}) + ) monkeypatch.setattr(flex_attention, "_execute_cudnn_graph", fake_execute) out = flex_attention.FusedAttentionWithScoreModFunc.apply( @@ -685,8 +579,13 @@ def _pytorch_softcap_attention(q, k, v, qkv_format, softmax_scale, softcap): ("post_scale_bias", False), ], ) -def test_dot_product_attention_score_mod(dtype, qkv_format, score_mod_case, scalar_loss): +def test_dot_product_attention_score_mod( + dtype, qkv_format, score_mod_case, scalar_loss, monkeypatch +): """Compare score_mod attention against equivalent reference implementations.""" + monkeypatch.setattr(cudnn_graph._graph_cache, "entries", {}) + compile_plan = Mock(wraps=flex_attention.compile_graph) + monkeypatch.setattr(flex_attention, "compile_graph", compile_plan) try: flex_attention._import_cudnn_frontend() except ImportError: @@ -841,3 +740,48 @@ def test_dot_product_attention_score_mod(dtype, qkv_format, score_mod_case, scal torch.testing.assert_close(q.grad, q_ref.grad, **tols) torch.testing.assert_close(k.grad, k_ref.grad, **tols) torch.testing.assert_close(v.grad, v_ref.grad, **tols) + + # New buffers with identical metadata must reuse both plans and rebind all operands. + repeated = [x.detach().clone().requires_grad_() for x in (q, k, v)] + repeated_kwargs = dict(score_mod_kwargs) + for name in ("score_mod_tensors", "score_mod_bprop_tensors"): + if name in repeated_kwargs: + repeated_kwargs[name] = { + key: value.clone() for key, value in repeated_kwargs[name].items() + } + repeated_out = flex_attn( + *repeated, qkv_format=qkv_format, attn_mask_type="no_mask", **repeated_kwargs + ) + repeated_out.backward(torch.ones_like(repeated_out) if scalar_loss else d_out) + torch.testing.assert_close(repeated_out, out_ref, **tols) + for actual, ref in zip(repeated, (q_ref, k_ref, v_ref)): + torch.testing.assert_close(actual.grad, ref.grad, **tols) + assert compile_plan.call_count == 2 + assert sum(key[0] == "score_mod" for key in cudnn_graph._graph_cache.entries) == 2 + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required.") +def test_score_mod_inference_reuses_graph(monkeypatch): + """Inference omits Stats and reuses its plan on new runtime buffers.""" + monkeypatch.setenv("NVTE_FLASH_ATTN", "0") + monkeypatch.setenv("NVTE_FUSED_ATTN", "1") + monkeypatch.setenv("NVTE_UNFUSED_ATTN", "0") + monkeypatch.setattr(cudnn_graph._graph_cache, "entries", {}) + compile_plan = Mock(wraps=flex_attention.compile_graph) + monkeypatch.setattr(flex_attention, "compile_graph", compile_plan) + _attention_backends["backend_selection_requires_update"] = True + attention = ( + DotProductAttention(4, 64, qkv_format="bshd", attn_mask_type="no_mask").cuda().eval() + ) + for _ in range(2): + q, k, v = (torch.randn(2, 128, 4, 64, device="cuda", dtype=torch.float16) for _ in range(3)) + qr, kr, vr = (x.transpose(1, 2).float() for x in (q, k, v)) + indices = torch.arange(128, device="cuda", dtype=torch.float32) + scores = qr @ kr.transpose(-1, -2) / 8 + indices[:, None] - indices[None, :] + reference = (scores.softmax(-1) @ vr).transpose(1, 2) + with torch.no_grad(): + out = attention(q, k, v, score_mod=_score_mod_post_scale_bias) + torch.testing.assert_close( + out.reshape(reference.shape).float(), reference, rtol=0.003, atol=0.002 + ) + assert compile_plan.call_count == 1 diff --git a/transformer_engine/common/cudnn/graph.py b/transformer_engine/common/cudnn/graph.py index ee0f111d3dd..33e3813b8bd 100644 --- a/transformer_engine/common/cudnn/graph.py +++ b/transformer_engine/common/cudnn/graph.py @@ -15,6 +15,52 @@ import time +@dataclass +class GraphCacheEntry: + """A supported graph and its optional compiled plan and serialization.""" + + built: object + workspace_size: int | None = None + serialized_graph: bytes | None = None + + +class GraphCache: + """Own graph entries and serialize lookup, construction, and compilation.""" + + def __init__(self): + self.entries = {} + self.lock = threading.RLock() + + def get(self, key, build, event=None): + """Build once per key; None explicitly bypasses storage for mutable callbacks.""" + if key is None: + return build() + with self.lock: + if key in self.entries: + if event is not None: + event("HIT", key) + return self.entries[key] + if event is not None: + event("MISS", key) + entry = build() + self.entries[key] = entry + if event is not None: + event("CACHE_GRAPH") + return entry + + def compile(self, entry, compile_plan, *, serialize=False, event=None): + """Publish compilation state only after all requested stages succeed.""" + with self.lock: + if entry.workspace_size is None: + workspace_size = compile_plan() + data = bytes(entry.built.graph.serialize()) if serialize else None + entry.workspace_size = workspace_size + entry.serialized_graph = data + if event is not None: + event("BUILD_PLANS") + return entry.workspace_size + + class UnsupportedGraph(ValueError): """A configuration rejected by cuDNN, rather than an execution/setup failure.""" diff --git a/transformer_engine/jax/cpp_extensions/cudnn_attention.py b/transformer_engine/jax/cpp_extensions/cudnn_attention.py index 27b681dd37e..1784f84fdf5 100644 --- a/transformer_engine/jax/cpp_extensions/cudnn_attention.py +++ b/transformer_engine/jax/cpp_extensions/cudnn_attention.py @@ -5,7 +5,6 @@ import math import os -import threading from dataclasses import dataclass from functools import lru_cache @@ -17,12 +16,12 @@ from transformer_engine.common.cudnn.attention import ( AttentionOptions, AttentionGraphSpec, - BuiltAttentionGraph, build_attention_graph, ) from transformer_engine.common.cudnn.graph import ( TensorSpec, CacheDiagnostics, + GraphCacheEntry, bucket as _bucket, row_major_stride as _stride, create_graph, @@ -38,6 +37,7 @@ compile_graph, serialize_graph, graph_attributes, + _graph_cache, ) from . import mxfp8_attention @@ -108,39 +108,23 @@ class _AttentionDescription: 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=None) -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 +def _build_supported_graph(spec, capabilities): 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) + return GraphCacheEntry(built) def _request(direction, avals, config): @@ -166,14 +150,11 @@ def _request(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 + return _graph_cache.get( + ("attention", spec, capabilities, frontend_version, backend_version), + lambda: _build_supported_graph(spec, capabilities), + event=lambda event, key=None: _diagnostics.event(*_site(spec), capabilities, event, key), + ) def check_attention_support(direction, avals, config): @@ -206,17 +187,16 @@ def _build_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") + _graph_cache.compile( + entry, + lambda: compile_graph( + cudnn, + entry.built.graph, + lambda stage, function: _diagnostics.run(backend, direction, stage, function), + ), + serialize=True, + event=lambda event: _diagnostics.event(backend, direction, capabilities, event), + ) uids = entry.built.uids serialized = serialize_graph( serialized_graph=entry.serialized_graph, diff --git a/transformer_engine/jax/cpp_extensions/cudnn_graph.py b/transformer_engine/jax/cpp_extensions/cudnn_graph.py index 4da11f172ec..0687522bae3 100644 --- a/transformer_engine/jax/cpp_extensions/cudnn_graph.py +++ b/transformer_engine/jax/cpp_extensions/cudnn_graph.py @@ -14,8 +14,11 @@ import_cudnn as _import_cudnn, check_graph_support as _check_graph_support, compile_graph as _compile_graph, + GraphCache, ) +_graph_cache = GraphCache() + @dataclass(frozen=True) class SerializedGraph: diff --git a/transformer_engine/jax/cpp_extensions/flex_attention.py b/transformer_engine/jax/cpp_extensions/flex_attention.py index 09cf381d691..8cb29c70cff 100644 --- a/transformer_engine/jax/cpp_extensions/flex_attention.py +++ b/transformer_engine/jax/cpp_extensions/flex_attention.py @@ -5,7 +5,7 @@ import os from dataclasses import dataclass -from typing import Any, Callable, Dict, Mapping, Optional, Sequence, Tuple +from typing import Any, Callable, Mapping, Optional, Sequence, Tuple import jax import jax.numpy as jnp @@ -41,6 +41,7 @@ serialize_graph as _serialized_score_mod_graph, finalize_graph as _finalize_score_mod_graph, import_cudnn as _import_cudnn_for_score_mod, + _graph_cache, ) __all__ = [ @@ -211,8 +212,6 @@ def __eq__(self, other): _SCORE_MOD_FWD_SCALAR_UID_BASE = 3000 _SCORE_MOD_BPROP_SCALAR_UID_BASE = 4000 -_score_mod_graph_cache: Dict[Tuple[Any, ...], _SerializedScoreModGraph] = {} - def _matrix_spec(cudnn, aval, output=False): """Translate JAX's contiguous BSHD storage into a shared logical descriptor.""" @@ -405,6 +404,7 @@ def _graph_cache_key( ): return None return ( + "score_mod", direction, config, tuple((tuple(aval.shape), _dtype_name(aval.dtype)) for aval in avals), @@ -545,14 +545,9 @@ def _fused_attn_score_mod_fwd( q_aval, k_aval, v_aval = map(_shape_dtype, (q, k, v)) score_mod_avals = tuple(_shape_dtype(arg) for arg in score_mod_tensors) key = _graph_cache_key("fwd", config, (q_aval, k_aval, v_aval, *score_mod_avals)) - if key is None: - graph = _build_score_mod_fwd_graph(q_aval, k_aval, v_aval, score_mod_avals, config) - else: - if key not in _score_mod_graph_cache: - _score_mod_graph_cache[key] = _build_score_mod_fwd_graph( - q_aval, k_aval, v_aval, score_mod_avals, config - ) - graph = _score_mod_graph_cache[key] + graph = _graph_cache.get( + key, lambda: _build_score_mod_fwd_graph(q_aval, k_aval, v_aval, score_mod_avals, config) + ) batch, q_seqlen, q_heads, _ = q.shape _, _, _, v_head_dim = v.shape @@ -599,22 +594,15 @@ def _fused_attn_score_mod_bwd( all_inputs = (*all_inputs, *score_mod_bprop_tensors) avals = tuple(_shape_dtype(arg) for arg in all_inputs) key = _graph_cache_key("bwd", config, avals) - if key is None: - graph = _build_score_mod_bwd_graph( + graph = _graph_cache.get( + key, + lambda: _build_score_mod_bwd_graph( *avals[:6], avals[6 : 6 + len(score_mod_tensors)], avals[6 + len(score_mod_tensors) :], config, - ) - else: - if key not in _score_mod_graph_cache: - _score_mod_graph_cache[key] = _build_score_mod_bwd_graph( - *avals[:6], - avals[6 : 6 + len(score_mod_tensors)], - avals[6 + len(score_mod_tensors) :], - config, - ) - graph = _score_mod_graph_cache[key] + ), + ) dq = jax.ShapeDtypeStruct(q.shape, q.dtype) dk = jax.ShapeDtypeStruct(k.shape, k.dtype) 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 e0972aedca3..08f50cfa341 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/flex_attention.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/flex_attention.py @@ -4,7 +4,6 @@ """cuDNN-backed Flex Attention helpers.""" -from dataclasses import dataclass from typing import Any, Callable, Dict, Optional, Tuple import torch @@ -15,7 +14,12 @@ AttentionGraphSpec, build_attention_graph, ) -from transformer_engine.common.cudnn.graph import TensorSpec +from transformer_engine.common.cudnn.graph import ( + TensorSpec, + GraphCacheEntry, + check_graph_support, + compile_graph, +) from transformer_engine.common.cudnn.score_mod import ( callback_cache_key, @@ -27,11 +31,10 @@ from transformer_engine.pytorch.cpp_extensions.cudnn import ( _import_cudnn_frontend, _build_cudnn_pygraph, - _finalize_cudnn_graph, + _graph_cache, _execute_cudnn_graph, ) -_cudnn_score_mod_graph_cache: Dict[Tuple[Any, ...], Any] = {} _SCORE_MOD_UNCACHEABLE = object() @@ -112,12 +115,6 @@ def _score_mod_tensor_dict_metadata( return tuple((name, _score_mod_tensor_metadata(tensor)) for name, tensor in tensors.items()) -def _score_mod_bhsd_tensor_metadata(tensor: torch.Tensor, tensor_format: str) -> Tuple[Any, ...]: - """Describe an SBHD/BSHD runtime tensor as a cuDNN BHSD graph tensor.""" - dim, stride = _bhsd_dim_stride(tensor, tensor_format) - return (dim, stride, tensor.dtype, _score_mod_device_key(tensor.device)) - - def _make_cudnn_graph_tensor_dict(graph, tensors: Optional[Dict[str, torch.Tensor]]): """Create cuDNN graph tensors matching runtime tensors.""" if tensors is None: @@ -125,307 +122,91 @@ def _make_cudnn_graph_tensor_dict(graph, tensors: Optional[Dict[str, torch.Tenso return {name: graph.tensor_like(tensor) for name, tensor in tensors.items()} -# cuDNN frontend score_mod graph helpers. -@dataclass -class _CudnnScoreModFwdGraphEntry: - """Cached cuDNN frontend graph and graph tensor handles for score_mod fprop.""" - - graph: Any - q: Any - k: Any - v: Any - output: Any - stats: Optional[Any] - score_mod_graph_tensors: Dict[str, Any] - workspace_size: int - - -@dataclass -class _CudnnScoreModBwdGraphEntry: - """Cached cuDNN frontend graph and graph tensor handles for score_mod bprop.""" - - graph: Any - q: Any - k: Any - v: Any - output: Any - d_output: Any - stats: Any - dq: Any - dk: Any - dv: Any - score_mod_graph_tensors: Dict[str, Any] - score_mod_bprop_graph_tensors: Dict[str, Any] - workspace_size: int - - -def _cudnn_score_mod_fwd_cache_key( - is_training: bool, - query_layer: torch.Tensor, - key_layer: torch.Tensor, - value_layer: torch.Tensor, - q_format: str, - kv_format: str, - attn_scale: float, - score_mod: Callable, - score_mod_tensors: Optional[Dict[str, torch.Tensor]], - output_layer: torch.Tensor, - stats: Optional[torch.Tensor], -) -> Optional[Tuple[Any, ...]]: - """Pre-build cache key for score_mod fprop execution plans. - - cuDNN exposes graph.key(), but only after graph construction has run the user callback. - This key avoids rebuilding the Python graph on cache hits. - """ - score_mod_key = _score_mod_callback_cache_key(score_mod) - if score_mod_key is _SCORE_MOD_UNCACHEABLE: - return None - return ( - "fwd", - is_training, - q_format, - kv_format, - attn_scale, - score_mod_key, - _score_mod_bhsd_tensor_metadata(query_layer, q_format), - _score_mod_bhsd_tensor_metadata(key_layer, kv_format), - _score_mod_bhsd_tensor_metadata(value_layer, kv_format), - _score_mod_bhsd_tensor_metadata(output_layer, q_format), - _score_mod_tensor_metadata(stats) if stats is not None else None, - _score_mod_tensor_dict_metadata(score_mod_tensors), - ) +def _score_mod_bindings(tensors, extra, bprop_extra): + """Keep callback operand names separate from SDPA and from each other.""" + return { + **tensors, + **{f"score_mod:{name}": tensor for name, tensor in extra.items()}, + **{f"score_mod_bprop:{name}": tensor for name, tensor in bprop_extra.items()}, + } -def _cudnn_score_mod_bwd_cache_key( - query_layer: torch.Tensor, - key_layer: torch.Tensor, - value_layer: torch.Tensor, - output_layer: torch.Tensor, - d_out: torch.Tensor, - stats: torch.Tensor, - q_format: str, - kv_format: str, - attn_scale: float, - score_mod: Callable, - score_mod_bprop: Optional[Callable], - score_mod_tensors: Optional[Dict[str, torch.Tensor]], - score_mod_bprop_tensors: Optional[Dict[str, torch.Tensor]], - deterministic: bool, -) -> Optional[Tuple[Any, ...]]: - """Pre-build cache key for score_mod bprop execution plans.""" - score_mod_key = _score_mod_callback_cache_key(score_mod) - score_mod_bprop_key = _score_mod_callback_cache_key(score_mod_bprop) - if score_mod_key is _SCORE_MOD_UNCACHEABLE or score_mod_bprop_key is _SCORE_MOD_UNCACHEABLE: +def _cudnn_score_mod_cache_key(spec, device, score_mod, score_mod_bprop, extra, bprop_extra): + """Combine the shared SDPA descriptor with callback topology and operand metadata.""" + keys = tuple( + _score_mod_callback_cache_key(callback) for callback in (score_mod, score_mod_bprop) + ) + if any(key is _SCORE_MOD_UNCACHEABLE for key in keys): return None return ( - "bwd", - q_format, - kv_format, - attn_scale, - deterministic, - score_mod_key, - score_mod_bprop_key, - _score_mod_bhsd_tensor_metadata(query_layer, q_format), - _score_mod_bhsd_tensor_metadata(key_layer, kv_format), - _score_mod_bhsd_tensor_metadata(value_layer, kv_format), - _score_mod_bhsd_tensor_metadata(output_layer, q_format), - _score_mod_bhsd_tensor_metadata(d_out, q_format), - _score_mod_tensor_metadata(stats), - _score_mod_tensor_dict_metadata(score_mod_tensors), - _score_mod_tensor_dict_metadata(score_mod_bprop_tensors), + "score_mod", + _score_mod_device_key(device), + spec, + keys, + _score_mod_tensor_dict_metadata(extra), + _score_mod_tensor_dict_metadata(bprop_extra), ) -def _build_cudnn_score_mod_fwd_graph( - is_training: bool, - query_layer: torch.Tensor, - key_layer: torch.Tensor, - value_layer: torch.Tensor, - q_format: str, - kv_format: str, - attn_scale: float, - score_mod: Callable, - score_mod_tensors: Optional[Dict[str, torch.Tensor]], - output_layer: torch.Tensor, - stats: Optional[torch.Tensor], -) -> _CudnnScoreModFwdGraphEntry: - """Build a cached cuDNN frontend graph for score_mod fprop.""" +def _build_cudnn_score_mod_graph(spec, device, score_mod, score_mod_bprop, extra, bprop_extra): + """Build and check either direction without retaining runtime tensors.""" cudnn = _import_cudnn_frontend() - - graph = _build_cudnn_pygraph(query_layer.dtype, query_layer.device) - 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 - specs["stats"] = _graph_tensor_spec(stats, output=True) + graph = _build_cudnn_pygraph(getattr(cudnn.data_type, dict(spec.tensors)["q"].dtype), device) + extra = _make_cudnn_graph_tensor_dict(graph, extra) + bprop_extra = _make_cudnn_graph_tensor_dict(graph, bprop_extra) built = build_attention_graph( cudnn, graph, - AttentionGraphSpec.from_dict( - AttentionOptions(attn_scale=attn_scale, generate_stats=is_training), specs - ), + spec, score_mod=_wrap_score_mod(score_mod, extra), + score_mod_bprop=_wrap_score_mod(score_mod_bprop, bprop_extra), ) - workspace_size = _finalize_cudnn_graph(graph, has_score_mod=True) - return _CudnnScoreModFwdGraphEntry( - graph=graph, - 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, - ) - - -def _get_cudnn_score_mod_fwd_graph( - is_training: bool, - query_layer: torch.Tensor, - key_layer: torch.Tensor, - value_layer: torch.Tensor, - q_format: str, - kv_format: str, - attn_scale: float, - score_mod: Callable, - score_mod_tensors: Optional[Dict[str, torch.Tensor]], - output_layer: torch.Tensor, - stats: Optional[torch.Tensor], -) -> _CudnnScoreModFwdGraphEntry: - """Return a cached cuDNN frontend graph for score_mod fprop.""" - build_args = ( - is_training, - query_layer, - key_layer, - value_layer, - q_format, - kv_format, - attn_scale, - score_mod, - score_mod_tensors, - output_layer, - stats, - ) - key = _cudnn_score_mod_fwd_cache_key(*build_args) - if key is None: - return _build_cudnn_score_mod_fwd_graph(*build_args) - entry = _cudnn_score_mod_graph_cache.get(key) - if entry is None: - entry = _build_cudnn_score_mod_fwd_graph(*build_args) - _cudnn_score_mod_graph_cache[key] = entry - return entry - - -def _build_cudnn_score_mod_bwd_graph( - query_layer: torch.Tensor, - key_layer: torch.Tensor, - value_layer: torch.Tensor, - output_layer: torch.Tensor, - d_out: torch.Tensor, - stats: torch.Tensor, - q_format: str, - kv_format: str, - attn_scale: float, - score_mod: Callable, - score_mod_bprop: Optional[Callable], - score_mod_tensors: Optional[Dict[str, torch.Tensor]], - score_mod_bprop_tensors: Optional[Dict[str, torch.Tensor]], - 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) - 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 {} - ) + built.tensors = _score_mod_bindings(built.tensors, extra, bprop_extra) + check_graph_support(cudnn, graph, has_score_mod=True) + return GraphCacheEntry(built) + + +def _get_cudnn_score_mod_graph( + options, + tensors, + q_format, + kv_format, + score_mod, + score_mod_bprop=None, + score_mod_tensors=None, + score_mod_bprop_tensors=None, +): + """Use the ordinary plan cache and return bindings for this invocation.""" + extra = score_mod_tensors or {} + bprop_extra = (score_mod_bprop_tensors or {}) if score_mod_bprop is not None else {} + outputs = ("dq", "dk", "dv") if options.backward else ("o", "stats") 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), + name: _graph_tensor_spec( + tensor, + ( + None + if name == "stats" + else (kv_format if name in ("k", "v", "dk", "dv") else q_format) + ), + output=name in outputs, + ) + for name, tensor in tensors.items() } - built = build_attention_graph( - cudnn, - graph, - AttentionGraphSpec.from_dict( - AttentionOptions(backward=True, attn_scale=attn_scale, deterministic=deterministic), - specs, + spec = AttentionGraphSpec.from_dict(options, specs) + device = tensors["q"].device + key = _cudnn_score_mod_cache_key(spec, device, score_mod, score_mod_bprop, extra, bprop_extra) + entry = _graph_cache.get( + key, + lambda: _build_cudnn_score_mod_graph( + spec, device, score_mod, score_mod_bprop, extra, bprop_extra ), - score_mod=_wrap_score_mod(score_mod, extra), - score_mod_bprop=_wrap_score_mod(score_mod_bprop, bprop_extra), ) - workspace_size = _finalize_cudnn_graph(graph, has_score_mod=True) - return _CudnnScoreModBwdGraphEntry( - graph=graph, - 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, + _graph_cache.compile( + entry, + lambda: compile_graph(_import_cudnn_frontend(), entry.built.graph, has_score_mod=True), ) - - -def _get_cudnn_score_mod_bwd_graph( - query_layer: torch.Tensor, - key_layer: torch.Tensor, - value_layer: torch.Tensor, - output_layer: torch.Tensor, - d_out: torch.Tensor, - stats: torch.Tensor, - q_format: str, - kv_format: str, - attn_scale: float, - score_mod: Callable, - score_mod_bprop: Optional[Callable], - score_mod_tensors: Optional[Dict[str, torch.Tensor]], - score_mod_bprop_tensors: Optional[Dict[str, torch.Tensor]], - deterministic: bool, -) -> _CudnnScoreModBwdGraphEntry: - """Return a cached cuDNN frontend graph for score_mod bprop.""" - build_args = ( - query_layer, - key_layer, - value_layer, - output_layer, - d_out, - stats, - q_format, - kv_format, - attn_scale, - score_mod, - score_mod_bprop, - score_mod_tensors, - score_mod_bprop_tensors, - deterministic, - ) - key = _cudnn_score_mod_bwd_cache_key(*build_args) - if key is None: - return _build_cudnn_score_mod_bwd_graph(*build_args) - entry = _cudnn_score_mod_graph_cache.get(key) - if entry is None: - entry = _build_cudnn_score_mod_bwd_graph(*build_args) - _cudnn_score_mod_graph_cache[key] = entry - return entry + bindings = _score_mod_bindings(tensors, extra, bprop_extra) + return entry, {entry.built.tensors[name]: tensor for name, tensor in bindings.items()} class FusedAttentionWithScoreModFunc(torch.autograd.Function): @@ -462,32 +243,20 @@ def forward( else: stats = None - entry = _get_cudnn_score_mod_fwd_graph( - is_training, - query_layer, - key_layer, - value_layer, + tensors = {"q": query_layer, "k": key_layer, "v": value_layer, "o": output_layer} + if stats is not None: + tensors["stats"] = stats + entry, variant_pack = _get_cudnn_score_mod_graph( + AttentionOptions(attn_scale=attn_scale, generate_stats=is_training), + tensors, q_format, kv_format, - attn_scale, score_mod, - score_mod_tensors, - output_layer, - stats, + score_mod_tensors=score_mod_tensors, ) - variant_pack = { - entry.q: query_layer, - entry.k: key_layer, - entry.v: value_layer, - entry.output: output_layer, - } - if is_training: - variant_pack[entry.stats] = stats - for name, graph_tensor in entry.score_mod_graph_tensors.items(): - variant_pack[graph_tensor] = score_mod_tensors[name] _execute_cudnn_graph( - entry.graph, + entry.built.graph, variant_pack, entry.workspace_size, query_layer.device, @@ -541,40 +310,31 @@ def backward(ctx, d_out: torch.Tensor): dq_layer = torch.empty_like(query_layer) dk_layer = torch.empty_like(key_layer) dv_layer = torch.empty_like(value_layer) - entry = _get_cudnn_score_mod_bwd_graph( - query_layer, - key_layer, - value_layer, - output_layer, - d_out, - stats, + entry, variant_pack = _get_cudnn_score_mod_graph( + AttentionOptions( + backward=True, attn_scale=ctx.attn_scale, deterministic=ctx.deterministic + ), + { + "q": query_layer, + "k": key_layer, + "v": value_layer, + "o": output_layer, + "do": d_out, + "stats": stats, + "dq": dq_layer, + "dk": dk_layer, + "dv": dv_layer, + }, ctx.q_format, ctx.kv_format, - ctx.attn_scale, ctx.score_mod, ctx.score_mod_bprop, score_mod_tensors, score_mod_bprop_tensors, - ctx.deterministic, ) - variant_pack = { - entry.q: query_layer, - entry.k: key_layer, - entry.v: value_layer, - entry.output: output_layer, - entry.d_output: d_out, - entry.stats: stats, - entry.dq: dq_layer, - entry.dk: dk_layer, - entry.dv: dv_layer, - } - for name, graph_tensor in entry.score_mod_graph_tensors.items(): - variant_pack[graph_tensor] = score_mod_tensors[name] - for name, graph_tensor in entry.score_mod_bprop_graph_tensors.items(): - variant_pack[graph_tensor] = score_mod_bprop_tensors[name] _execute_cudnn_graph( - entry.graph, + entry.built.graph, variant_pack, entry.workspace_size, query_layer.device, diff --git a/transformer_engine/pytorch/cpp_extensions/cudnn.py b/transformer_engine/pytorch/cpp_extensions/cudnn.py index 911b0fcbd85..e430be77eb1 100644 --- a/transformer_engine/pytorch/cpp_extensions/cudnn.py +++ b/transformer_engine/pytorch/cpp_extensions/cudnn.py @@ -14,10 +14,12 @@ create_graph, check_graph_support, compile_graph, + GraphCache, ) from transformer_engine.pytorch.device import te_device_type _handles = threading.local() +_graph_cache = GraphCache() def _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 index a23bcaddf6f..b4847f77f7f 100644 --- a/transformer_engine/pytorch/cpp_extensions/cudnn_fused_attn.py +++ b/transformer_engine/pytorch/cpp_extensions/cudnn_fused_attn.py @@ -8,10 +8,10 @@ Graph construction, plan caching, and execution use the Python frontend. """ -from dataclasses import dataclass, replace +from dataclasses import replace +from functools import partial import math import os -import threading from typing import List, Optional, Tuple, Union import torch @@ -25,6 +25,7 @@ from transformer_engine.common.cudnn.graph import ( TensorSpec as _TensorSpec, CacheDiagnostics, + GraphCacheEntry, attention_matrix_spec, attention_stats_spec, mxfp8_scale_spec, @@ -39,6 +40,7 @@ _import_cudnn_frontend, _build_cudnn_pygraph, _execute_cudnn_graph, + _graph_cache, ) from ..tensor.float8_tensor import Float8Quantizer from ..tensor.storage.mxfp8_tensor_storage import MXFP8TensorStorage @@ -58,8 +60,6 @@ _FP8_THREADS_PER_CTA = 128 _F16_RNG_ELTS_PER_THREAD = 16 -_plans = {} -_cache_lock = threading.RLock() _diagnostics = CacheDiagnostics("PyTorch") @@ -125,15 +125,6 @@ def _descale(tensor, columnwise=False): return tensor._scale_inv -@dataclass -class _PlanEntry: - """A supported graph, compiled lazily under the cache lock.""" - - graph: object - tensors: dict - workspace_size: int | None = None - - class _AttentionGraph: """Runtime tensors and data-free descriptors for one attention invocation.""" @@ -301,19 +292,11 @@ def description(self): def entry(self): """Share support-query graphs with execution without compiling on queries.""" - key = (self.device, self.description()) # Autograd may run backward on another thread. Share device-specific plans; # each execution supplies its thread-local handle on the current stream. - with _cache_lock: - entry = _plans.get(key) - if entry is None: - _event(self, "MISS", key) - entry = self.build() - _plans[key] = entry - _event(self, "CACHE_GRAPH") - else: - _event(self, "HIT", key) - return entry + return _graph_cache.get( + (self.device, self.description()), self.build, event=partial(_event, self) + ) def build(self): """Construct and check the common graph description without compiling it.""" @@ -322,22 +305,23 @@ def build(self): built = build_attention_graph(cudnn, graph, self.description()) _event(self, "CREATE_GRAPH") check_graph_support(cudnn, graph, lambda name, function: _build_stage(self, name, function)) - return _PlanEntry(graph, built.tensors) + return GraphCacheEntry(built) def execute(self): """Compile once, then bind this invocation's runtime tensor addresses.""" entry = self.entry() - with _cache_lock: - if entry.workspace_size is None: - entry.workspace_size = compile_graph( - _import_cudnn_frontend(), - entry.graph, - lambda name, function: _build_stage(self, name, function), - ) - _event(self, "BUILD_PLANS") + _graph_cache.compile( + entry, + lambda: compile_graph( + _import_cudnn_frontend(), + entry.built.graph, + lambda name, function: _build_stage(self, name, function), + ), + event=partial(_event, self), + ) _execute_cudnn_graph( - entry.graph, - {entry.tensors[n]: t for n, t in self.tensors.items()}, + entry.built.graph, + {entry.built.tensors[n]: t for n, t in self.tensors.items()}, entry.workspace_size, self.device, ) From d1249bf55c0ba5fe9e90b174f8b17a6ba5baec86 Mon Sep 17 00:00:00 2001 From: Vladimir Cherepanov Date: Thu, 12 Nov 2026 01:09:38 +0000 Subject: [PATCH 43/49] Consolidate JAX FP8 attention numerical tests Reuse the existing FP32 attention reference for CS and MXFP8, replace broad descriptor and initialization matrices with focused contracts, and cover realistic scaled MHA, GQA, and cross-attention cases. Keep Flax-specific sink and prequantized boundary checks, and merge context-parallel storage coverage into a curated distributed matrix. Signed-off-by: Vladimir Cherepanov --- tests/jax/test_distributed_fused_attn.py | 104 +++--- tests/jax/test_fused_attn.py | 384 +++++++++-------------- tests/jax/test_mxfp8_attention.py | 22 +- 3 files changed, 190 insertions(+), 320 deletions(-) diff --git a/tests/jax/test_distributed_fused_attn.py b/tests/jax/test_distributed_fused_attn.py index 9bd3d819cb7..1643197dd18 100644 --- a/tests/jax/test_distributed_fused_attn.py +++ b/tests/jax/test_distributed_fused_attn.py @@ -1025,11 +1025,26 @@ def test_mxfp8_prequantized_sharding(axis): @pytest.mark.parametrize("strategy", (CPStrategy.ALL_GATHER, CPStrategy.RING)) -@pytest.mark.parametrize("cp_size", (2, 4)) -@pytest.mark.parametrize("kind", ("dense", "causal", "balanced_causal")) -@pytest.mark.parametrize("fp8_backward", (False, True)) -def test_mxfp8_context_parallel(strategy, cp_size, kind, fp8_backward, monkeypatch): - """Context-parallel outputs and gradients agree with an independent global reference.""" +@pytest.mark.parametrize( + "kind,cp_size,fp8_backward,layout,prequantized", + [ + ("dense", 2, False, QKVLayout.BSHD_BSHD_BSHD, False), + ("dense", 4, True, QKVLayout.BSHD_BSHD_BSHD, False), + ("causal", 2, True, QKVLayout.BSHD_BSHD_BSHD, False), + ("causal", 4, False, QKVLayout.BSHD_BSHD_BSHD, False), + ("balanced_causal", 2, False, QKVLayout.BSHD_BSHD_BSHD, False), + ("balanced_causal", 4, True, QKVLayout.BSHD_BSHD_BSHD, False), + ("dense", 2, True, QKVLayout.BS3HD, False), + ("dense", 2, True, QKVLayout.BSHD_BS2HD, False), + ("dense", 2, True, QKVLayout.BS3HD, True), + ("dense", 2, True, QKVLayout.BSHD_BS2HD, True), + ("dense", 2, True, QKVLayout.BSHD_BSHD_BSHD, True), + ], +) +def test_mxfp8_context_parallel( + strategy, kind, cp_size, fp8_backward, layout, prequantized, monkeypatch +): + """CP masks, load balancing and packed/supplied inputs agree with global FP32 attention.""" from transformer_engine.jax.cpp_extensions.attention import reorder_causal_dual_chunk_swap _require_mxfp8() @@ -1039,69 +1054,29 @@ def test_mxfp8_context_parallel(strategy, cp_size, kind, fp8_backward, monkeypat mesh = jax.sharding.Mesh(np.asarray(jax.local_devices()[:cp_size]), ("cp",)) sharding = jax.sharding.NamedSharding(mesh, jax.sharding.PartitionSpec(None, "cp", None, None)) length = 256 * cp_size - q = jnp.full((1, length, 4, 64), 0.5, jnp.bfloat16) - k = q[:, :, :2] - v = jnp.zeros_like(k).at[:, ::2, :, 0].set(1.0) - # Distinct rank-local values make missing communication or incorrect - # dual-chunk ordering observable, while retaining exact MXFP8 operands. - amplitudes = jnp.exp2((jnp.arange(length) // 128) % 4 - 2).astype(v.dtype) - v *= amplitudes[None, :, None, None] - inputs = (q, k, v) + rng = np.random.default_rng(42) + q = jnp.asarray(rng.choice([-0.5, 0.5], (1, length, 4, 64)), jnp.bfloat16) + kv_shape = q.shape if layout.is_qkvpacked() else (1, length, 2, 64) + k = jnp.asarray(rng.choice([-0.5, 0.5], kv_shape), q.dtype) + v = jnp.asarray(rng.choice([-1.0, 1.0], kv_shape), q.dtype) + # Nonuniform rank-local keys and values expose missing communication and reordering. + amplitude = jnp.exp2((jnp.arange(length) // 128) % 4 - 2).astype(q.dtype) + k *= amplitude[None, :, None, None] + v *= amplitude[None, :, None, None] + weight = jnp.asarray(rng.choice([-1.0, 1.0], q.shape), q.dtype) + inputs = _pack_mxfp8((q, k, v), layout) balanced = kind == "balanced_causal" causal = kind != "dense" mask = ( jnp.arange(length)[None, :, None] >= jnp.arange(length)[None, None, :] if causal else None ) - reference = lambda xs: _mxfp8_reference(xs, mask) + reference = lambda xs: _mxfp8_reference(_unpack_mxfp8(xs, layout), mask) expected = jax.jit(reference)(inputs) - ref_grads = jax.jit(jax.grad(lambda xs: reference(xs).sum()))(inputs) + ref_grads = jax.jit(jax.grad(lambda xs: (reference(xs) * weight).sum()))(inputs) reorder = lambda x, reverse: reorder_causal_dual_chunk_swap(x, cp_size, 1, reverse) + local_weight = reorder(weight, False) if balanced else weight if balanced: inputs = tuple(reorder(x, False) for x in inputs) - attention = lambda xs: _mxfp8_attention( - xs, - attn_mask_type=AttnMaskType.CAUSAL_MASK if causal else AttnMaskType.NO_MASK, - context_parallel_axis="cp", - context_parallel_strategy=strategy, - context_parallel_causal_load_balanced=balanced, - ) - with jax.set_mesh(mesh), global_shard_guard(MeshResource(cp_resource="cp")): - actual = jax.jit(attention, in_shardings=((sharding,) * 3,))(inputs) - grads = jax.jit(jax.grad(lambda xs: attention(xs).sum()), in_shardings=((sharding,) * 3,))( - inputs - ) - if balanced: - actual = reorder(actual, True) - grads = tuple(reorder(x, True) for x in grads) - np.testing.assert_allclose(actual, expected, rtol=1e-2, atol=1e-3) - for grad, ref in zip(grads, ref_grads): - np.testing.assert_allclose(grad, ref, rtol=3e-2, atol=2e-2) - - -@pytest.mark.parametrize("strategy", (CPStrategy.ALL_GATHER, CPStrategy.RING)) -@pytest.mark.parametrize( - "layout", (QKVLayout.BS3HD, QKVLayout.BSHD_BS2HD, QKVLayout.BSHD_BSHD_BSHD) -) -@pytest.mark.parametrize("prequantized", (False, True)) -def test_mxfp8_context_parallel_storage(strategy, layout, prequantized): - """Packed and supplied FP8 inputs preserve nonuniform cross-rank values and gradients.""" - _require_mxfp8() - if jax.local_device_count() < 2: - pytest.skip("Requires two GPUs for context parallelism") - mesh = jax.sharding.Mesh(np.asarray(jax.local_devices()[:2]), ("cp",)) - sharding = jax.sharding.NamedSharding(mesh, jax.sharding.PartitionSpec(None, "cp", None, None)) - rng = np.random.default_rng(42) - q = jnp.asarray(rng.choice([-0.5, 0.5], (1, 512, 2, 64)), jnp.bfloat16) - k = jnp.asarray(rng.choice([-0.5, 0.5], q.shape), q.dtype) - v = jnp.asarray(rng.choice([-1.0, 1.0], q.shape), q.dtype) - amplitude = jnp.exp2(jnp.arange(512) // 128 - 2).astype(q.dtype) - k *= amplitude[None, :, None, None] - v *= amplitude[None, :, None, None] - weight = jnp.asarray(rng.choice([-1.0, 1.0], q.shape), q.dtype) - inputs = _pack_mxfp8((q, k, v), layout) - reference = lambda xs: _mxfp8_reference(_unpack_mxfp8(xs, layout)) - expected = jax.jit(reference)(inputs) - ref_grads = jax.jit(jax.grad(lambda xs: (reference(xs) * weight).sum()))(inputs) def attention(xs): supplied = quantize_attention_qkv(_unpack_mxfp8(xs, layout)) if prequantized else None @@ -1109,17 +1084,24 @@ def attention(xs): xs, qkv_layout=layout, quantized_qkv=supplied, + attn_mask_type=AttnMaskType.CAUSAL_MASK if causal else AttnMaskType.NO_MASK, context_parallel_axis="cp", context_parallel_strategy=strategy, + context_parallel_causal_load_balanced=balanced, ) with jax.set_mesh(mesh), global_shard_guard(MeshResource(cp_resource="cp")): actual = jax.jit(attention, in_shardings=((sharding,) * len(inputs),))(inputs) grads = jax.jit( - jax.grad(lambda xs: (attention(xs) * weight).sum()), + jax.grad(lambda xs: (attention(xs) * local_weight).sum()), in_shardings=((sharding,) * len(inputs),), )(inputs) - relative_error = lambda a, b: float(jnp.linalg.norm(a - b) / jnp.linalg.norm(b)) + if balanced: + actual = reorder(actual, True) + grads = tuple(reorder(x, True) for x in grads) + relative_error = lambda a, b: float( + jnp.linalg.norm(a.astype(jnp.float32) - b) / jnp.linalg.norm(b) + ) assert relative_error(actual, expected) < 0.04 for grad, ref in zip(grads, ref_grads): assert relative_error(grad, ref) < 0.08 diff --git a/tests/jax/test_fused_attn.py b/tests/jax/test_fused_attn.py index ee0bcc3e336..f05d9ed32f7 100644 --- a/tests/jax/test_fused_attn.py +++ b/tests/jax/test_fused_attn.py @@ -144,6 +144,31 @@ def general_dot_product_attention( return context +def attention_reference( + qkv, mask=None, sink=None, scaling_factor=0.125, return_max_logit=False, bias=None +): + """FP32 unfused attention; mask marks valid query/key pairs in BQK order.""" + result = general_dot_product_attention( + *qkv, + softmax_offset=sink, + bias=bias, + mask=None if mask is None else ~mask[:, None], + deterministic=True, + softmax_type=( + AttnSoftmaxType.VANILLA_SOFTMAX if sink is None else AttnSoftmaxType.LEARNABLE_SOFTMAX + ), + scale_factor=scaling_factor, + dropout_rate=0.0, + dropout_rng=None, + dtype=jnp.float32, + is_max_logit_enabled=return_max_logit, + ) + if mask is not None and not return_max_logit: + # cuDNN returns zero for fully masked queries, including their gradients. + result = jnp.where(jnp.any(mask, axis=-1)[..., None, None], result, 0.0) + return result + + @jax.jit def make_causal_mask( segment_ids_q: ArrayLike, @@ -525,177 +550,25 @@ def test_cudnn_attention_backward_o_precision_cache(fp8, fp8_backward, monkeypat assert backward[0] is backward[1] -@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" - - -@pytest.mark.parametrize("dtype", (jnp.float16, jnp.bfloat16)) -@pytest.mark.parametrize("head_dim_v", (64, 128, 192, 256)) -@pytest.mark.parametrize("version", (92500, 92600)) -@pytest.mark.parametrize("direction", ("fwd", "bwd")) -def test_cudnn_attention_thd_learnable_support(dtype, head_dim_v, version, direction, monkeypatch): - """Preserve specialized FP16/BF16 THD learnable-softmax kernels before cuDNN 9.26.""" - from types import SimpleNamespace - from transformer_engine.jax.cpp_extensions import cudnn_attention - - monkeypatch.setattr(cudnn_attention, "cudnn_version", lambda: version) - 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.LEARNABLE_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(dtype)) - for shape in ((2, 128, 4, 64),) * 2 + ((2, 128, 4, head_dim_v), (0,), (1, 4, 1, 1)) - ) - if direction == "bwd" and version < 92600 and head_dim_v == 192: - with pytest.raises(ValueError, match="THD learnable softmax backward requires cuDNN 9.26"): - cudnn_attention._describe_graph(direction, shapes, config, (100,), True, True) - else: - description = cudnn_attention._describe_graph(direction, shapes, config, (100,), True, True) - assert not description.fp8 - assert dict(description.spec.tensors)["v"].dim[-1] == head_dim_v - - -@pytest.mark.parametrize("softmax", tuple(AttnSoftmaxType)) -@pytest.mark.parametrize("version", (92300, 92500, 92600)) -@pytest.mark.parametrize("direction", ("fwd", "bwd")) -@pytest.mark.parametrize("fp8_backward", (False, True)) -def test_cudnn_attention_fp8_thd_sink_support( - softmax, version, direction, fp8_backward, monkeypatch -): - """All sink-token FP8 backward paths require 9.26; F16 backward remains available.""" - from types import SimpleNamespace - from transformer_engine.jax.cpp_extensions import cudnn_attention - - monkeypatch.setattr(cudnn_attention, "cudnn_version", lambda: version) - config = SimpleNamespace( - qkv_layout=QKVLayout.THD_THD_THD, - max_segments_per_seq=1, - fp8=True, - attn_mask_type=AttnMaskType.PADDING_MASK, - attn_bias_type=AttnBiasType.NO_BIAS, - softmax_type=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)) - ) - if ( - direction == "bwd" - and fp8_backward - and version < 92600 - and softmax != AttnSoftmaxType.VANILLA_SOFTMAX - ): - with pytest.raises( - ValueError, match="FP8 THD backward with a sink token requires cuDNN 9.26" - ): - cudnn_attention._describe_graph(direction, shapes, config, (100,), True, fp8_backward) - else: - description = cudnn_attention._describe_graph( - direction, shapes, config, (100,), True, fp8_backward - ) - assert description.fp8 == (direction == "fwd" or fp8_backward) - - -@pytest.mark.parametrize( - "layout", - ( - QKVLayout.BS3HD, - QKVLayout.BSHD_BS2HD, - QKVLayout.BSHD_BSHD_BSHD, - QKVLayout.T3HD, - QKVLayout.THD_T2HD, - QKVLayout.THD_THD_THD, - ), -) -@pytest.mark.parametrize("direction", ("fwd", "bwd")) -@pytest.mark.parametrize("fp8", (False, True)) -def test_cudnn_attention_output_initialization(layout, direction, fp8, monkeypatch): - """Only ragged storage and unwritten optional results request an FFI fill.""" +def test_cudnn_attention_unbound_backward_outputs(monkeypatch): + """Unwritten dBias and packed-gradient dummy slots request zero initialization.""" from types import SimpleNamespace from transformer_engine.jax.cpp_extensions import cudnn_attention q = jnp.ones((1, 128, 2, 64), jnp.bfloat16) empty = jnp.empty((0,), q.dtype) - if layout.is_qkvpacked(): - operands = (jnp.stack((q, q, q), axis=-3), empty, empty) - grad_slots = (0, 0, 0) - elif layout.is_kvpacked(): - operands = (q, jnp.stack((q, q), axis=-3), empty) - grad_slots = (0, 1, 1) - else: - operands = (q, q, q) - grad_slots = (0, 1, 2) - backward = direction == "bwd" - output_slots = (*grad_slots, 4) if backward else (0, 1) - if fp8: - output_slots += (5, 6, 7, 8) if backward else (3, 4) - elif not backward: - output_slots += (2,) + operands = (jnp.stack((q, q, q), axis=-3), empty, empty) plan = SimpleNamespace( graph=SimpleNamespace(workspace_size=0), names=(), input_indices=(), input_offsets=(), - output_indices=output_slots, - output_offsets=(0,) * len(output_slots), + output_indices=(0, 0, 0, 4), + output_offsets=(0, 0, 0, 0), metadata_batch=1, - ragged_stats=layout.is_thd(), + ragged_stats=False, multipliers=(128, 128, 128, 128, 2), - fp8=fp8, + fp8=False, fp8_o=False, ) monkeypatch.setattr(cudnn_attention, "get_graph", lambda *_args: plan) @@ -710,37 +583,15 @@ def execute(*_args, **attributes): return execute monkeypatch.setattr(cudnn_attention.ffi, "ffi_call", ffi_call) - config = SimpleNamespace(qkv_layout=layout, fp8=fp8, return_max_logit=not fp8) - stats_shape = (1, 128, 2, 1) if layout.is_thd() else (1, 2, 128, 1) - stats = jnp.empty(stats_shape, jnp.float32) + config = SimpleNamespace(qkv_layout=QKVLayout.BS3HD, fp8=False, return_max_logit=False) + stats = jnp.empty((1, 2, 128, 1), jnp.float32) rng = jnp.zeros((4,), jnp.uint32) cu_seqlens = jnp.array([0, 128], jnp.int32) bias = jnp.zeros((1, 1, 1, 128), q.dtype) sink = jnp.zeros((1, 2, 1, 1), jnp.float32) - args = (*operands, bias, sink) - if backward: - args += (stats, rng, q, q, cu_seqlens, cu_seqlens, cu_seqlens, cu_seqlens) - results = (*operands, bias, sink, empty) - # dBias is deliberately unbound; packed gradient slots also have dummy results. - expected = [ - -1, - 0 if layout.is_qkvpacked() else -1, - 0 if not layout.is_separate() else -1, - 0, - -1, - ] - if layout.is_thd(): - expected[:3] = [0, 0, 0] - if fp8: - expected += [-1] * 4 - else: - args += (rng, cu_seqlens, cu_seqlens, cu_seqlens, cu_seqlens) - results = (q, stats, empty if fp8 else stats, rng, empty) - expected = [0, 0xF0, 0xF0] if layout.is_thd() else [-1, -1, 0 if fp8 else -1] - if fp8: - expected += [-1, -1] - cudnn_attention.execute(direction, args, config, results) - np.testing.assert_array_equal(captured["output_fill_values"], expected) + args = (*operands, bias, sink, stats, rng, q, q, cu_seqlens, cu_seqlens, cu_seqlens, cu_seqlens) + cudnn_attention.execute("bwd", args, config, (*operands, bias, sink, empty)) + np.testing.assert_array_equal(captured["output_fill_values"], [-1, 0, 0, 0, -1]) @pytest.mark.parametrize("dtype", (jnp.float16, jnp.bfloat16)) @@ -2499,9 +2350,8 @@ def _require_cudnn_attention(q, k, v, **overrides): pytest.skip(reason) -@pytest.mark.parametrize("o_in_f16", [True, False]) +@pytest.mark.parametrize("fp8_backward,o_in_f16", [(False, True), (True, True), (True, False)]) @pytest.mark.parametrize("dtype", [jnp.float16, jnp.bfloat16]) -@pytest.mark.parametrize("fp8_backward", [False, True]) @pytest.mark.parametrize( "layout", [ @@ -2532,15 +2382,13 @@ def unpack(operands): 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) + return attention_reference(unpack(operands)) def attention(operands): return fused_attn( operands, None, - None, + SequenceDescriptor(), None, attn_bias_type=AttnBiasType.NO_BIAS, attn_mask_type=AttnMaskType.NO_MASK, @@ -2563,6 +2411,85 @@ def attention(operands): np.testing.assert_allclose(actual_grad, expected_grad, rtol=1e-2, atol=1e-2) +@pytest.mark.parametrize( + "fp8,fp8_backward,o_in_f16", + [ + pytest.param(True, False, True, id="cs-f16-bwd"), + pytest.param(True, True, True, id="cs-fp8-bwd-f16-o"), + pytest.param(True, True, False, id="cs-fp8-bwd-fp8-o"), + pytest.param("mxfp8", False, True, id="mxfp8-f16-bwd"), + pytest.param("mxfp8", True, True, id="mxfp8-fp8-bwd"), + ], +) +@pytest.mark.parametrize( + "shape,input_scales,causal", + [ + pytest.param((1, 512, 512, 16, 16, 64), (0.5, 1.0, 2.0), False, id="mha-512"), + pytest.param((2, 1024, 1024, 16, 4, 128), (1.0, 0.5, 4.0), True, id="gqa-1024"), + ] + + ( + [ + pytest.param((1, 2048, 2048, 32, 8, 128), (0.5, 2.0, 1.0), True, id="gqa-2048"), + pytest.param((2, 512, 2048, 16, 4, 64), (2.0, 0.25, 0.5), False, id="cross-2048"), + ] + if _TEST_LEVEL == "L2" + else [] + ), +) +def test_cudnn_python_fp8_model_configs( + shape, input_scales, causal, fp8, fp8_backward, o_in_f16, monkeypatch +): + """Representative CS/MXFP8 outputs and gradients agree with FP32 unfused attention.""" + monkeypatch.setenv("NVTE_FP8_DPA_BWD", "1" if fp8_backward else "0") + monkeypatch.setenv("NVTE_DPA_FP8CS_O_in_F16", "1" if o_in_f16 else "0") + batch, sq, sk, q_heads, kv_heads, dim = shape + keys = jax.random.split(jax.random.key(321), 4) + shapes = ((batch, sq, q_heads, dim), (batch, sk, kv_heads, dim), (batch, sk, kv_heads, dim)) + arrays = [] + for key, tensor_shape, scale in zip(keys, shapes, input_scales): + values = jnp.clip(jax.random.normal(key, tensor_shape), -7.0, 7.0) + values = values.astype(jnp.float8_e4m3fn).astype(jnp.bfloat16) + # Fix the maximum so CS scaling is a power of two. Both recipes then + # preserve these random operands exactly, isolating internal FP8 P/dP error. + arrays.append(values.at[0, 0, 0, 0].set(7.0) * scale) + arrays = tuple(arrays) + weight = jax.random.normal(keys[3], shapes[0]) + scaling_factor = dim**-0.5 + attn_mask = AttnMaskType.CAUSAL_MASK if causal else AttnMaskType.NO_MASK + mask = (jnp.arange(sq)[None, :, None] >= jnp.arange(sk)[None, None, :]) if causal else None + _require_cudnn_attention(*arrays, fp8=fp8, attn_mask_type=attn_mask, attn_scale=scaling_factor) + + def attention(xs): + return fused_attn( + xs, + None, + SequenceDescriptor(), + None, + attn_bias_type=AttnBiasType.NO_BIAS, + attn_mask_type=attn_mask, + qkv_layout=QKVLayout.BSHD_BSHD_BSHD, + softmax_type=AttnSoftmaxType.VANILLA_SOFTMAX, + scaling_factor=scaling_factor, + dropout_probability=0.0, + is_training=True, + fp8=fp8, + ) + + reference = lambda xs: attention_reference(xs, mask, scaling_factor=scaling_factor) + actual, expected = jax.jit(attention)(arrays), jax.jit(reference)(arrays) + relative_error = lambda a, b: float( + jnp.linalg.norm(a.astype(jnp.float32) - b) / jnp.linalg.norm(b) + ) + assert actual.dtype == jnp.bfloat16 + assert relative_error(actual, expected) < (0.04 if fp8 == "mxfp8" else 0.06) + actual_grads = jax.jit(jax.grad(lambda xs: (attention(xs) * weight).sum()))(arrays) + expected_grads = jax.jit(jax.grad(lambda xs: (reference(xs) * weight).sum()))(arrays) + for name, grad, ref in zip(("dQ", "dK", "dV"), actual_grads, expected_grads): + assert grad.dtype == jnp.bfloat16 + error = relative_error(grad, ref) + assert error < (0.08 if fp8 == "mxfp8" else 0.1), f"{name} relative error: {error}" + + @pytest.mark.parametrize("o_in_f16", [True, False]) @pytest.mark.parametrize("scaling_factor", [0.125, 0.25]) @pytest.mark.parametrize( @@ -2575,8 +2502,6 @@ def attention(operands): ("ragged", None), ("ragged_dense_stats", 120), ("ragged_dense_stats_f16", 120), - ("ragged_dense_stats_f16", 121), - ("ragged_dense_stats_f16", 129), ], ) def test_cudnn_python_fp8_gqa_masks(kind, sm_arch, scaling_factor, o_in_f16, monkeypatch, request): @@ -2602,7 +2527,7 @@ def test_cudnn_python_fp8_gqa_masks(kind, sm_arch, scaling_factor, o_in_f16, mon 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 + descriptor = SequenceDescriptor() attn_mask = AttnMaskType.NO_MASK layout = QKVLayout.BSHD_BSHD_BSHD window = None @@ -2615,7 +2540,9 @@ def test_cudnn_python_fp8_gqa_masks(kind, sm_arch, scaling_factor, o_in_f16, mon if kind == "padding": lengths = jnp.array([128, 91], jnp.int32) descriptor = SequenceDescriptor.from_seqlens((lengths, lengths)) - mask &= positions[None, None, :] < lengths[:, None, None] + mask &= (positions[None, None, :] < lengths[:, None, None]) & ( + positions[None, :, None] < lengths[:, None, None] + ) attn_mask = AttnMaskType.PADDING_MASK if ragged: ids = jnp.broadcast_to(positions // 64 + 1, (2, 128)) @@ -2640,16 +2567,7 @@ def test_cudnn_python_fp8_gqa_masks(kind, sm_arch, scaling_factor, o_in_f16, mon ) 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) * scaling_factor - 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 + return attention_reference(arrays, mask, scaling_factor=scaling_factor) def attention(arrays): result = fused_attn( @@ -2707,18 +2625,9 @@ def relative_error(actual, expected): fp8=fp8, ) )(arrays) - scores = ( - jnp.einsum( - "bqhd,bkhd->bhqk", - q.astype(jnp.float32), - jnp.repeat(k.astype(jnp.float32), 2, axis=-2), - ) - * scaling_factor + expected_max = attention_reference( + arrays, mask, scaling_factor=scaling_factor, return_max_logit=True ) - 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) @@ -2731,17 +2640,17 @@ def test_cudnn_python_alibi_forward_backward(): 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) + return attention_reference( + arrays, + causal[None], + bias=slopes[None, :, None, None] * positions[None, None, None, :], + ) def attention(arrays): return fused_attn( arrays, None, - None, + SequenceDescriptor(), None, attn_bias_type=AttnBiasType.ALIBI, attn_mask_type=AttnMaskType.CAUSAL_MASK, @@ -2761,11 +2670,9 @@ def attention(arrays): 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.""" +def test_cudnn_python_flax_sink_gradients(fp8, monkeypatch): + """The cuDNN Python backend returns gradients to Flax's learnable sink parameter.""" from transformer_engine.jax.flax import DotProductAttention monkeypatch.setenv("NVTE_FUSED_ATTN", "1") @@ -2774,7 +2681,7 @@ def test_cudnn_python_flax_forward_backward(softmax, deterministic, fp8, monkeyp num_attention_heads=2, attn_mask_type="no_mask", transpose_batch_sequence=False, - softmax_type=softmax, + softmax_type="learnable", fp8=fp8, ) q = jnp.zeros((1, 128, 2, 64), jnp.bfloat16) @@ -2785,19 +2692,15 @@ def test_cudnn_python_flax_forward_backward(softmax, deterministic, fp8, monkeyp q, v, fp8=True, - softmax_type=( - AttnSoftmaxType.LEARNABLE_SOFTMAX - if softmax == "learnable" - else AttnSoftmaxType.VANILLA_SOFTMAX - ), + softmax_type=AttnSoftmaxType.LEARNABLE_SOFTMAX, ) - variables = module.init(jax.random.PRNGKey(0), q, q, v, deterministic=deterministic) + variables = module.init(jax.random.PRNGKey(0), q, q, v, deterministic=True) def forward(variables, value): - return module.apply(variables, q, q, value, deterministic=deterministic) + return module.apply(variables, q, q, value, deterministic=True) actual = jax.jit(forward)(variables, v) - denominator = 129 if softmax == "learnable" else 128 + denominator = 129 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( @@ -2808,11 +2711,10 @@ def forward(variables, value): )(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) + 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)]) diff --git a/tests/jax/test_mxfp8_attention.py b/tests/jax/test_mxfp8_attention.py index bc37e2a5c01..f7024095155 100644 --- a/tests/jax/test_mxfp8_attention.py +++ b/tests/jax/test_mxfp8_attention.py @@ -11,6 +11,8 @@ import numpy as np import pytest +from test_fused_attn import attention_reference as _reference + from transformer_engine.common.cudnn import mxfp8_utils as mxfp8 from transformer_engine.common.cudnn.graph import mxfp8_scale_spec from transformer_engine.jax.attention import ( @@ -77,24 +79,6 @@ def _attention(qkv, **kwargs): return fused_attn(qkv, None, descriptor, seed, **options) -def _reference(qkv, mask=None, sink=None): - q, k, v = (x.astype(jnp.float32) for x in qkv) - repeats = q.shape[-2] // k.shape[-2] - k, v = (jnp.repeat(x, repeats, axis=-2) for x in (k, v)) - scores = jnp.einsum("bqhd,bkhd->bhqk", q, k) * 0.125 - if mask is not None: - scores = jnp.where(mask[:, None], scores, -1e30) - if sink is None: - probability = jax.nn.softmax(scores, -1) - else: - sink = jnp.broadcast_to(sink, (*scores.shape[:-1], 1)) - probability = jax.nn.softmax(jnp.concatenate((scores, sink), axis=-1), -1)[..., :-1] - output = jnp.einsum("bhqk,bkhd->bqhd", probability, v) - if mask is not None: - output = jnp.where(jnp.any(mask, axis=-1)[..., None, None], output, 0.0) - return output - - @pytest.mark.parametrize("dtype", (jnp.float16, jnp.bfloat16)) @pytest.mark.parametrize("fp8_backward", (False, True)) @pytest.mark.parametrize("is_training", (False, True)) @@ -377,6 +361,8 @@ def test_flax_prequantized_inputs(transpose): q = jnp.full((1, 128, 2, 64), 0.5, jnp.bfloat16) v = jnp.zeros_like(q).at[:, ::2, :, 0].set(1.0) prepared = jax.jit(quantize_attention_qkv)((q, q, v)) + # A dropped quantized_qkv argument would produce zeros from the nominal V. + v = jnp.zeros_like(v) if transpose: q, v = q.swapaxes(0, 1), v.swapaxes(0, 1) module = DotProductAttention( From 9080bddfb9ecf38f3c287b5ec7f8593edef2299d Mon Sep 17 00:00:00 2001 From: Vladimir Cherepanov Date: Thu, 12 Nov 2026 20:29:22 +0000 Subject: [PATCH 44/49] Skip the bias graph cache test in deterministic mode Fused attention backward with dBias does not support deterministic execution. Skip the bias graph cache test when deterministic algorithms are required, while retaining its nondeterministic coverage. Validation: all 12 deterministic cases skipped and all 12 nondeterministic cases passed on GB200. Pre-commit, PyTorch Python lint, and the changed-file license check passed. Signed-off-by: Vladimir Cherepanov --- tests/pytorch/attention/test_attention.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/tests/pytorch/attention/test_attention.py b/tests/pytorch/attention/test_attention.py index 6d6dd31f040..ed7f9a47e59 100644 --- a/tests/pytorch/attention/test_attention.py +++ b/tests/pytorch/attention/test_attention.py @@ -445,6 +445,10 @@ def test_fused_attn_query_matches_runtime_graph(layout, backward, dropout): @pytest.mark.parametrize("dtype", (torch.float16, torch.bfloat16)) @pytest.mark.parametrize("qkv_format", ("bshd", "sbhd")) @pytest.mark.parametrize("bias_layout", ("contiguous", "gapped", "permuted")) +@pytest.mark.skipif( + _deterministic, + reason="Fused attention backward with dBias does not support deterministic execution", +) def test_fused_attn_bias_graph_cache(dtype, qkv_format, bias_layout, monkeypatch): """Bias layouts reuse queried plans and preserve forward/backward numerics.""" from transformer_engine.pytorch.cpp_extensions import cudnn_fused_attn as implementation From 303ab1802c02664a12c8ff9e48230fd4e9ef8ceb Mon Sep 17 00:00:00 2001 From: Vladimir Cherepanov Date: Thu, 12 Nov 2026 20:40:19 +0000 Subject: [PATCH 45/49] Fix architecture assumptions in PyTorch attention tests Align THD FP8 MHA token counts to 16 for weight-gradient GEMMs and reserve enough space in the final sequence for alignment. Limit batch-bucketing identity assertions to ragged graphs. Skip the deterministic fused eval/backward numerical test below SM90. Validation on GB200: 243 focused tests passed with 50 expected skips. Checked 96 graph-metadata cases with simulated SM80/90/100/120 capabilities and token alignment/bounds for 384 configurations. Pre-commit, PyTorch Python lint, and changed-file license checks passed. Actual H100/A100 execution remains for CI. Signed-off-by: Vladimir Cherepanov --- tests/pytorch/attention/test_attention.py | 10 +++++----- .../attention/test_attention_backend_selection.py | 4 ++++ 2 files changed, 9 insertions(+), 5 deletions(-) diff --git a/tests/pytorch/attention/test_attention.py b/tests/pytorch/attention/test_attention.py index ed7f9a47e59..36c344c90d2 100644 --- a/tests/pytorch/attention/test_attention.py +++ b/tests/pytorch/attention/test_attention.py @@ -435,7 +435,7 @@ def test_fused_attn_query_matches_runtime_graph(layout, backward, dropout): runtime.add("dsink", torch.empty_like(params["softmax_offset"]), output=True) assert query.description() == runtime.description() - if fmt == "thd" and not query.use_cu_seqlens: + if fmt == "thd" and query.ragged_graph and not query.use_cu_seqlens: other_query = implementation._query_graph(replace(config, batch_size=batch + 1), backward) assert other_query.description() == query.description() runtime.scalar("attn_scale", torch.tensor([0.25]), by_value=True) @@ -4043,8 +4043,8 @@ def draw(high, count): def random_seqlens(max_seqlen): if qkv_format != "thd": return draw(max_seqlen, config.batch_size) - # Reserve seven positions so total-token alignment only increases the final length. - return torch.cat((draw(max_seqlen, config.batch_size - 1), draw(max_seqlen - 6, 1))) + # Reserve fifteen positions so total-token alignment only increases the final length. + return torch.cat((draw(max_seqlen, config.batch_size - 1), draw(max_seqlen - 14, 1))) if config.attn_type == "self": seqlens_q = random_seqlens(config.max_seqlen_q) @@ -4061,8 +4061,8 @@ def random_seqlens(max_seqlen): ) if qkv_format == "thd": # FP8 Linear flattens THD input to [t, h*d], so align total tokens for cuBLAS. - seqlens_q[-1] += -seqlens_q.sum() % 8 - seqlens_kv[-1] += -seqlens_kv.sum() % 8 + seqlens_q[-1] += -seqlens_q.sum() % 16 + seqlens_kv[-1] += -seqlens_kv.sum() % 16 return seqlens_q, seqlens_kv diff --git a/tests/pytorch/attention/test_attention_backend_selection.py b/tests/pytorch/attention/test_attention_backend_selection.py index 73e79118aba..1ce59ba919a 100644 --- a/tests/pytorch/attention/test_attention_backend_selection.py +++ b/tests/pytorch/attention/test_attention_backend_selection.py @@ -144,6 +144,10 @@ def test_backward_requirement_with_frozen_qkv(monkeypatch, grad_target): @pytest.mark.parametrize("dtype", (torch.float16, torch.bfloat16)) +@pytest.mark.skipif( + torch.cuda.get_device_capability() < (9, 0), + reason="Deterministic fused attention backward requires SM90 or later", +) def test_fused_eval_backward_after_forward_only(monkeypatch, dtype): """Evaluation with autograd must retain numerical parity after a cached inference call.""" monkeypatch.setenv("NVTE_FUSED_ATTN", "1") From b9315cf075d8c2868ff8597e9ac4eb17433f59b0 Mon Sep 17 00:00:00 2001 From: Vladimir Cherepanov Date: Thu, 12 Nov 2026 20:55:19 +0000 Subject: [PATCH 46/49] Remove JAX build dependency metadata tests Signed-off-by: Vladimir Cherepanov --- tests/jax/test_build_tools.py | 54 ----------------------------------- 1 file changed, 54 deletions(-) delete mode 100644 tests/jax/test_build_tools.py diff --git a/tests/jax/test_build_tools.py b/tests/jax/test_build_tools.py deleted file mode 100644 index 5f39fb4dbb5..00000000000 --- a/tests/jax/test_build_tools.py +++ /dev/null @@ -1,54 +0,0 @@ -# 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 -from pathlib import Path - -import pytest -from packaging.requirements import Requirement -from setuptools import Distribution - -# CI runs outside the checkout root and tests the installed TE package. -# Expose the source-only build tools just for this import. -with pytest.MonkeyPatch.context() as _build_import_path: - _build_import_path.syspath_prepend(str(Path(__file__).resolve().parents[2])) - 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() From 3bdbc8303e6ba17568136929d44a813e55bc15f2 Mon Sep 17 00:00:00 2001 From: Vladimir Cherepanov Date: Thu, 12 Nov 2026 21:28:07 +0000 Subject: [PATCH 47/49] Reuse shared graph attributes in JAX Flex Attention Use graph_attributes for forward and backward FFI calls instead of repeating the nine serialized-graph attributes. Validation: pre-commit, JAX Python lint, syntax, and changed-file license checks passed. Original-versus-new FFI argument parity passed for inference forward, training forward, and backward. JAX numerical tests were not run because JAX is not installed. Signed-off-by: Vladimir Cherepanov --- .../jax/cpp_extensions/flex_attention.py | 21 +++---------------- 1 file changed, 3 insertions(+), 18 deletions(-) diff --git a/transformer_engine/jax/cpp_extensions/flex_attention.py b/transformer_engine/jax/cpp_extensions/flex_attention.py index 8cb29c70cff..61254a2c22f 100644 --- a/transformer_engine/jax/cpp_extensions/flex_attention.py +++ b/transformer_engine/jax/cpp_extensions/flex_attention.py @@ -41,6 +41,7 @@ serialize_graph as _serialized_score_mod_graph, finalize_graph as _finalize_score_mod_graph, import_cudnn as _import_cudnn_for_score_mod, + graph_attributes, _graph_cache, ) @@ -563,15 +564,7 @@ def _fused_attn_score_mod_fwd( k, v, *score_mod_tensors, - 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, + **graph_attributes(graph), ) return output, softmax_stats @@ -620,15 +613,7 @@ def _fused_attn_score_mod_bwd( softmax_stats, *score_mod_tensors, *score_mod_bprop_tensors, - 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, + **graph_attributes(graph), ) return dq, dk, dv From b8b8b6beeb379de98c5fc1ed5c645172858d3889 Mon Sep 17 00:00:00 2001 From: Vladimir Cherepanov Date: Thu, 12 Nov 2026 21:38:48 +0000 Subject: [PATCH 48/49] Add Flex Attention cache diagnostics Share per-framework diagnostic collectors between ordinary SDPA and Flex Attention. Report cache events and graph-build timings, plus PyTorch executions, while retaining JAX tracing-only reporting and correct uncached-callback behavior. Document the shared summaries and add coverage for debug levels 0, 1, and 2 with cacheable and uncacheable callbacks. Validation: 70 PyTorch tests passed on GB200 and 21 isolated JAX unit cases passed with framework stubs and mocked cuDNN. Both framework Python lints, pre-commit, syntax, and changed-file license checks passed. Native JAX execution was not tested because JAX and its extension are not installed. Signed-off-by: Vladimir Cherepanov --- docs/envvars.rst | 2 +- tests/jax/test_fused_attn_score_mod.py | 67 ++++++++++++++++++- .../pytorch/attention/test_flex_attention.py | 62 +++++++++++++++++ .../jax/cpp_extensions/cudnn_attention.py | 5 +- .../jax/cpp_extensions/cudnn_graph.py | 8 ++- .../jax/cpp_extensions/flex_attention.py | 24 ++++++- .../dot_product_attention/flex_attention.py | 21 +++++- .../pytorch/cpp_extensions/cudnn.py | 2 + .../cpp_extensions/cudnn_fused_attn.py | 4 +- 9 files changed, 178 insertions(+), 17 deletions(-) diff --git a/docs/envvars.rst b/docs/envvars.rst index 8a55ff508a4..0d5b833f78f 100644 --- a/docs/envvars.rst +++ b/docs/envvars.rst @@ -240,7 +240,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 a summary of cache counters and build timings at process exit. ``2`` also prints individual events. At both levels, JAX diagnostics describe graph construction and caching during tracing, rather than each execution of a compiled function. 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 a summary of cache counters and build timings at process exit. ``2`` also prints individual events. Ordinary SDPA and Flex Attention contribute to the same per-framework summary; Flex Attention uses the ``f16`` forward/backward counters. Uncacheable callbacks report graph construction and compilation without cache hits, misses, or insertions. At both levels, JAX diagnostics describe graph construction and caching during tracing, rather than each execution of a compiled function. 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/tests/jax/test_fused_attn_score_mod.py b/tests/jax/test_fused_attn_score_mod.py index 644245c6065..d47f2cb64e8 100644 --- a/tests/jax/test_fused_attn_score_mod.py +++ b/tests/jax/test_fused_attn_score_mod.py @@ -4,7 +4,8 @@ """Tests for cuDNN frontend score_mod fused attention.""" from math import sqrt -from types import FunctionType +from collections import Counter +from types import FunctionType, SimpleNamespace from unittest.mock import Mock, call import jax @@ -739,6 +740,70 @@ def test_cudnn_graph_compilation_retains_backend_fallback(): assert graph.select_plan.call_args_list == [call(1), call(2)] +@pytest.mark.parametrize("level", (0, 1, 2)) +@pytest.mark.parametrize("cacheable", (False, True)) +def test_fused_attn_score_mod_cache_diagnostics(monkeypatch, capsys, level, cacheable): + """Tracing reports real graph/cache stages, never executions of the FFI call.""" + diagnostics = tex_cudnn_graph._diagnostics + assert diagnostics is tex_attention._diagnostics + monkeypatch.setattr(diagnostics, "_level", level) + monkeypatch.setattr(diagnostics, "events", Counter()) + monkeypatch.setattr(diagnostics, "times", Counter()) + monkeypatch.setattr(tex_cudnn_graph._graph_cache, "entries", {}) + devices = Mock(return_value=[SimpleNamespace(compute_capability="10.0")]) + monkeypatch.setattr(jax, "local_devices", devices) + + cudnn = Mock(__version__="1.31.0", cudnnGraphNotSupportedError=NotImplementedError) + cudnn.data_type = SimpleNamespace( + HALF=SimpleNamespace(name="HALF"), FLOAT=SimpleNamespace(name="FLOAT") + ) + graph = Mock(selected_engine=None) + graph.sdpa.return_value = (Mock(), Mock()) + graph.sdpa_backward.return_value = (Mock(), Mock(), Mock()) + graph.get_execution_plan_count.return_value = 1 + graph.get_plan_name_at_index.return_value = "cudnn_plan" + graph.get_workspace_size.return_value = 64 + graph.serialize.return_value = b"graph" + monkeypatch.setattr(tex_attention, "_import_cudnn_for_score_mod", lambda: cudnn) + monkeypatch.setattr(tex_attention, "create_graph", lambda *_args: graph) + monkeypatch.setattr(tex_cudnn_graph, "check_frontend_version", lambda _cudnn: 13100) + monkeypatch.setattr( + tex_attention.ffi, "ffi_call", lambda _target, outputs: lambda *_args, **_kwargs: outputs + ) + + class UnkeyedScoreMod: + def forward(self, _graph, score, _tensors): + return score + + callback = _identity_score_mod if cacheable else UnkeyedScoreMod().forward + config, _, _ = make_fused_attn_score_mod_config(callback, callback, None, None, 0.125, True) + qkv = (jax.ShapeDtypeStruct((2, 16, 4, 64), jnp.float16),) * 3 + for _ in range(2): + out, stats = tex_attention.fused_attn_score_mod_fwd(qkv, (), config) + tex_attention.fused_attn_score_mod_bwd(qkv, out, out, stats, (), (), config) + + assert len(tex_cudnn_graph._graph_cache.entries) == (2 if cacheable else 0) + for direction in ("fwd", "bwd"): + expected = {"CREATE_GRAPH": 1 if cacheable else 2, "BUILD_PLANS": 1 if cacheable else 2} + if cacheable: + expected.update(MISS=1, HIT=1, CACHE_GRAPH=1) + assert { + event: count + for (backend, event_direction, event), count in diagnostics.events.items() + if backend == "f16" and event_direction == direction + } == (expected if level else {}) + assert diagnostics.times[("f16", direction, "build_plans", "calls")] == ( + expected["BUILD_PLANS"] if level else 0 + ) + if not level: + devices.assert_not_called() + diagnostics.summary() + stderr = capsys.readouterr().err + assert ("JAX Python frontend summary" in stderr) == bool(level) + assert (" f16 fwd CREATE_GRAPH\n" in stderr) == (level == 2) + assert "EXECUTE" not in stderr + + def test_fused_attn_score_mod_cudnn_frontend_version_check(monkeypatch): """cuDNN frontend Python and C++ versions must match.""" diff --git a/tests/pytorch/attention/test_flex_attention.py b/tests/pytorch/attention/test_flex_attention.py index 68de5f86848..5dbedec4985 100644 --- a/tests/pytorch/attention/test_flex_attention.py +++ b/tests/pytorch/attention/test_flex_attention.py @@ -5,6 +5,7 @@ import os import sys import pathlib +from collections import Counter from types import FunctionType, SimpleNamespace from unittest.mock import Mock, call from concurrent.futures import ThreadPoolExecutor @@ -23,6 +24,7 @@ from transformer_engine.pytorch.attention.dot_product_attention import _attention_backends import transformer_engine.pytorch.attention.dot_product_attention.flex_attention as flex_attention from transformer_engine.pytorch.cpp_extensions import cudnn as cudnn_graph +from transformer_engine.pytorch.cpp_extensions import cudnn_fused_attn from transformer_engine.pytorch.utils import get_device_compute_capability _current_file = pathlib.Path(__file__).resolve() @@ -467,6 +469,66 @@ def forward(self, _score_mod_graph, score_tensor, _tensors): assert build.call_count == 2 +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required.") +@pytest.mark.parametrize("level", (0, 1, 2)) +@pytest.mark.parametrize("cacheable", (False, True)) +def test_score_mod_cache_diagnostics(monkeypatch, capsys, level, cacheable): + """Report real forward/backward builds and executions without caching unsafe callbacks.""" + diagnostics = cudnn_graph._diagnostics + assert diagnostics is flex_attention._diagnostics is cudnn_fused_attn._diagnostics + monkeypatch.setattr(diagnostics, "_level", level) + monkeypatch.setattr(diagnostics, "events", Counter()) + monkeypatch.setattr(diagnostics, "times", Counter()) + monkeypatch.setattr(cudnn_graph._graph_cache, "entries", {}) + + class UnkeyedScoreMod: + def forward(self, graph, score, tensors): + return _score_mod_post_scale_bias(graph, score, tensors) + + callback = _score_mod_post_scale_bias if cacheable else UnkeyedScoreMod().forward + for _ in range(2): + qkv = [ + torch.randn(2, 16, 4, 64, device="cuda", dtype=torch.float16, requires_grad=True) + for _ in range(3) + ] + out = flex_attention.FusedAttentionWithScoreModFunc.apply( + True, + *qkv, + "bshd", + "bshd", + 0.125, + callback, + _score_mod_identity_bprop, + None, + None, + False, + ) + out.backward(torch.ones_like(out)) + torch.cuda.synchronize() + + assert len(cudnn_graph._graph_cache.entries) == (2 if cacheable else 0) + for direction in ("fwd", "bwd"): + expected = { + "CREATE_GRAPH": 1 if cacheable else 2, + "BUILD_PLANS": 1 if cacheable else 2, + "EXECUTE": 2, + } + if cacheable: + expected.update(MISS=1, HIT=1, CACHE_GRAPH=1) + assert { + event: count + for (backend, event_direction, event), count in diagnostics.events.items() + if backend == "f16" and event_direction == direction + } == (expected if level else {}) + assert diagnostics.times[("f16", direction, "build_plans", "calls")] == ( + expected["BUILD_PLANS"] if level else 0 + ) + diagnostics.summary() + stderr = capsys.readouterr().err + assert ("PyTorch Python frontend summary" in stderr) == bool(level) + assert (" f16 fwd EXECUTE\n" in stderr) == (level == 2) + + def test_score_mod_tensors_are_version_checked_for_backward(monkeypatch): """In-place score_mod tensor updates before backward should be rejected.""" diff --git a/transformer_engine/jax/cpp_extensions/cudnn_attention.py b/transformer_engine/jax/cpp_extensions/cudnn_attention.py index 1784f84fdf5..3c0ef41a062 100644 --- a/transformer_engine/jax/cpp_extensions/cudnn_attention.py +++ b/transformer_engine/jax/cpp_extensions/cudnn_attention.py @@ -20,7 +20,6 @@ ) from transformer_engine.common.cudnn.graph import ( TensorSpec, - CacheDiagnostics, GraphCacheEntry, bucket as _bucket, row_major_stride as _stride, @@ -38,6 +37,7 @@ serialize_graph, graph_attributes, _graph_cache, + _diagnostics, ) from . import mxfp8_attention @@ -108,9 +108,6 @@ class _AttentionDescription: fp8: bool -_diagnostics = CacheDiagnostics("JAX") - - def _site(spec): return "fp8" if spec.options.fp8 else "f16", "bwd" if spec.options.backward else "fwd" diff --git a/transformer_engine/jax/cpp_extensions/cudnn_graph.py b/transformer_engine/jax/cpp_extensions/cudnn_graph.py index 0687522bae3..6295aa42e0c 100644 --- a/transformer_engine/jax/cpp_extensions/cudnn_graph.py +++ b/transformer_engine/jax/cpp_extensions/cudnn_graph.py @@ -15,9 +15,11 @@ check_graph_support as _check_graph_support, compile_graph as _compile_graph, GraphCache, + CacheDiagnostics, ) _graph_cache = GraphCache() +_diagnostics = CacheDiagnostics("JAX") @dataclass(frozen=True) @@ -153,10 +155,10 @@ def compile_graph(cudnn, graph, run_stage=None, *, has_score_mod=False): raise cudnn.cudnnGraphNotSupportedError("No serializable cuDNN SDPA execution plan.") -def finalize_graph(cudnn, graph, *, has_score_mod=False): +def finalize_graph(cudnn, graph, run_stage=None, *, has_score_mod=False): """Validate, select and compile a supported execution plan.""" - check_graph_support(cudnn, graph, has_score_mod=has_score_mod) - workspace_size = compile_graph(cudnn, graph, has_score_mod=has_score_mod) + check_graph_support(cudnn, graph, run_stage, has_score_mod=has_score_mod) + workspace_size = compile_graph(cudnn, graph, run_stage, has_score_mod=has_score_mod) return workspace_size, bytes(graph.serialize()), check_frontend_version(cudnn) diff --git a/transformer_engine/jax/cpp_extensions/flex_attention.py b/transformer_engine/jax/cpp_extensions/flex_attention.py index 61254a2c22f..495945dae38 100644 --- a/transformer_engine/jax/cpp_extensions/flex_attention.py +++ b/transformer_engine/jax/cpp_extensions/flex_attention.py @@ -5,6 +5,7 @@ import os from dataclasses import dataclass +from functools import partial from typing import Any, Callable, Mapping, Optional, Sequence, Tuple import jax @@ -43,6 +44,7 @@ import_cudnn as _import_cudnn_for_score_mod, graph_attributes, _graph_cache, + _diagnostics, ) __all__ = [ @@ -416,6 +418,15 @@ def _shape_dtype(value) -> jax.ShapeDtypeStruct: return jax.ShapeDtypeStruct(tuple(value.shape), value.dtype) +def _score_mod_cache_event(direction, event, key=None): + """Report tracing activity without querying devices when diagnostics are disabled.""" + if _diagnostics.level(): + capabilities = tuple( + int(float(device.compute_capability) * 10) for device in jax.local_devices() + ) + _diagnostics.event("f16", direction, capabilities, event, key) + + def _build_score_mod_fwd_graph(q_aval, k_aval, v_aval, score_mod_avals, config): cudnn = _import_cudnn_for_score_mod() graph = create_graph(cudnn, _cudnn_data_type(cudnn, q_aval.dtype)) @@ -450,9 +461,11 @@ def _build_score_mod_fwd_graph(q_aval, k_aval, v_aval, score_mod_avals, config): ), score_mod=_wrap_score_mod(config.score_mod, extra), ) + _score_mod_cache_event("fwd", "CREATE_GRAPH") workspace_size, data, frontend_version = _finalize_score_mod_graph( - cudnn, graph, has_score_mod=True + cudnn, graph, partial(_diagnostics.run, "f16", "fwd"), has_score_mod=True ) + _score_mod_cache_event("fwd", "BUILD_PLANS") return _serialized_score_mod_graph( serialized_graph=data, cudnn_frontend_version=frontend_version, @@ -520,9 +533,11 @@ def _build_score_mod_bwd_graph( score_mod=_wrap_score_mod(config.score_mod, extra), score_mod_bprop=_wrap_score_mod(config.score_mod_bprop, bprop_extra), ) + _score_mod_cache_event("bwd", "CREATE_GRAPH") workspace_size, data, frontend_version = _finalize_score_mod_graph( - cudnn, graph, has_score_mod=True + cudnn, graph, partial(_diagnostics.run, "f16", "bwd"), has_score_mod=True ) + _score_mod_cache_event("bwd", "BUILD_PLANS") return _serialized_score_mod_graph( serialized_graph=data, cudnn_frontend_version=frontend_version, @@ -547,7 +562,9 @@ def _fused_attn_score_mod_fwd( score_mod_avals = tuple(_shape_dtype(arg) for arg in score_mod_tensors) key = _graph_cache_key("fwd", config, (q_aval, k_aval, v_aval, *score_mod_avals)) graph = _graph_cache.get( - key, lambda: _build_score_mod_fwd_graph(q_aval, k_aval, v_aval, score_mod_avals, config) + key, + lambda: _build_score_mod_fwd_graph(q_aval, k_aval, v_aval, score_mod_avals, config), + event=partial(_score_mod_cache_event, "fwd"), ) batch, q_seqlen, q_heads, _ = q.shape @@ -595,6 +612,7 @@ def _fused_attn_score_mod_bwd( avals[6 + len(score_mod_tensors) :], config, ), + event=partial(_score_mod_cache_event, "bwd"), ) dq = jax.ShapeDtypeStruct(q.shape, q.dtype) 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 08f50cfa341..89af9f37fde 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/flex_attention.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/flex_attention.py @@ -4,6 +4,7 @@ """cuDNN-backed Flex Attention helpers.""" +from functools import partial from typing import Any, Callable, Dict, Optional, Tuple import torch @@ -33,6 +34,7 @@ _build_cudnn_pygraph, _graph_cache, _execute_cudnn_graph, + _diagnostics, ) _SCORE_MOD_UNCACHEABLE = object() @@ -162,7 +164,11 @@ def _build_cudnn_score_mod_graph(spec, device, score_mod, score_mod_bprop, extra score_mod_bprop=_wrap_score_mod(score_mod_bprop, bprop_extra), ) built.tensors = _score_mod_bindings(built.tensors, extra, bprop_extra) - check_graph_support(cudnn, graph, has_score_mod=True) + direction = "bwd" if spec.options.backward else "fwd" + _diagnostics.event("f16", direction, device.index, "CREATE_GRAPH") + check_graph_support( + cudnn, graph, partial(_diagnostics.run, "f16", direction), has_score_mod=True + ) return GraphCacheEntry(built) @@ -194,16 +200,25 @@ def _get_cudnn_score_mod_graph( } spec = AttentionGraphSpec.from_dict(options, specs) device = tensors["q"].device + direction = "bwd" if options.backward else "fwd" + event = partial(_diagnostics.event, "f16", direction, device.index) key = _cudnn_score_mod_cache_key(spec, device, score_mod, score_mod_bprop, extra, bprop_extra) entry = _graph_cache.get( key, lambda: _build_cudnn_score_mod_graph( spec, device, score_mod, score_mod_bprop, extra, bprop_extra ), + event=event, ) _graph_cache.compile( entry, - lambda: compile_graph(_import_cudnn_frontend(), entry.built.graph, has_score_mod=True), + lambda: compile_graph( + _import_cudnn_frontend(), + entry.built.graph, + partial(_diagnostics.run, "f16", direction), + has_score_mod=True, + ), + event=event, ) bindings = _score_mod_bindings(tensors, extra, bprop_extra) return entry, {entry.built.tensors[name]: tensor for name, tensor in bindings.items()} @@ -261,6 +276,7 @@ def forward( entry.workspace_size, query_layer.device, ) + _diagnostics.event("f16", "fwd", query_layer.device.index, "EXECUTE") ctx.is_training = is_training ctx.q_format = q_format @@ -339,6 +355,7 @@ def backward(ctx, d_out: torch.Tensor): entry.workspace_size, query_layer.device, ) + _diagnostics.event("f16", "bwd", query_layer.device.index, "EXECUTE") return ( None, diff --git a/transformer_engine/pytorch/cpp_extensions/cudnn.py b/transformer_engine/pytorch/cpp_extensions/cudnn.py index e430be77eb1..a843e02ac80 100644 --- a/transformer_engine/pytorch/cpp_extensions/cudnn.py +++ b/transformer_engine/pytorch/cpp_extensions/cudnn.py @@ -15,11 +15,13 @@ check_graph_support, compile_graph, GraphCache, + CacheDiagnostics, ) from transformer_engine.pytorch.device import te_device_type _handles = threading.local() _graph_cache = GraphCache() +_diagnostics = CacheDiagnostics("PyTorch") def _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 index b4847f77f7f..e24c47fcbf8 100644 --- a/transformer_engine/pytorch/cpp_extensions/cudnn_fused_attn.py +++ b/transformer_engine/pytorch/cpp_extensions/cudnn_fused_attn.py @@ -24,7 +24,6 @@ ) from transformer_engine.common.cudnn.graph import ( TensorSpec as _TensorSpec, - CacheDiagnostics, GraphCacheEntry, attention_matrix_spec, attention_stats_spec, @@ -41,6 +40,7 @@ _build_cudnn_pygraph, _execute_cudnn_graph, _graph_cache, + _diagnostics, ) from ..tensor.float8_tensor import Float8Quantizer from ..tensor.storage.mxfp8_tensor_storage import MXFP8TensorStorage @@ -60,8 +60,6 @@ _FP8_THREADS_PER_CTA = 128 _F16_RNG_ELTS_PER_THREAD = 16 -_diagnostics = CacheDiagnostics("PyTorch") - def _event(attn, event, key=None): _diagnostics.event( From 8ad483327d59ea1c33cb6997ffae60e51093aa82 Mon Sep 17 00:00:00 2001 From: Vladimir Cherepanov Date: Fri, 13 Nov 2026 20:47:13 +0000 Subject: [PATCH 49/49] Fix attention test backend availability checks Match backward support queries and autograd state to the test execution mode. Check the custom FP8 test's model-precision reference separately from its fused FP8 path to avoid depending on leaked emulation settings. Validate all 44 previously extra-skipped cases and the shared 5,259-case selection. Outcomes now match main: 2,734 passed and 2,525 skipped, with no failures. Signed-off-by: Vladimir Cherepanov --- tests/pytorch/attention/test_attention.py | 54 ++++++++++++++--------- tests/pytorch/utils.py | 8 +++- 2 files changed, 39 insertions(+), 23 deletions(-) diff --git a/tests/pytorch/attention/test_attention.py b/tests/pytorch/attention/test_attention.py index 9d5608056c4..fbc62328291 100644 --- a/tests/pytorch/attention/test_attention.py +++ b/tests/pytorch/attention/test_attention.py @@ -3329,7 +3329,8 @@ def get_dummy_cuda_rng_tracker() -> CudaRNGStatesTracker: ) if forward_kwargs is not None: dpa_forward_kwargs.update(forward_kwargs) - out = block(q, k, v, **packed_kwargs, **dpa_forward_kwargs) + with torch.set_grad_enabled(is_training): + out = block(q, k, v, **packed_kwargs, **dpa_forward_kwargs) max_logit = None if config.return_max_logit: out, max_logit = out @@ -3769,21 +3770,22 @@ def _run_transformer_layer( alibi_slopes = torch.randn(config.num_heads).abs().to(dtype=torch.float32, device="cuda") # Run a forward and backward pass - out = block( - inp, - self_attn_mask_type=config.attn_mask_type, - encoder_output=inp_enc if config.attn_type == "cross" else None, - enc_dec_attn_mask_type=config.attn_mask_type if config.attn_type == "cross" else None, - checkpoint_core_attention=False, - rotary_pos_emb=rotary_pos_emb, - core_attention_bias_type=config.attn_bias_type, - core_attention_bias=bias, - alibi_slopes=alibi_slopes, - max_seqlen_q=config.max_seqlen_q, - max_seqlen_kv=config.max_seqlen_kv, - cu_seqlens_q=cu_seqlens_q, - cu_seqlens_kv=cu_seqlens_kv, - ) + with torch.set_grad_enabled(is_training): + out = block( + inp, + self_attn_mask_type=config.attn_mask_type, + encoder_output=inp_enc if config.attn_type == "cross" else None, + enc_dec_attn_mask_type=config.attn_mask_type if config.attn_type == "cross" else None, + checkpoint_core_attention=False, + rotary_pos_emb=rotary_pos_emb, + core_attention_bias_type=config.attn_bias_type, + core_attention_bias=bias, + alibi_slopes=alibi_slopes, + max_seqlen_q=config.max_seqlen_q, + max_seqlen_kv=config.max_seqlen_kv, + cu_seqlens_q=cu_seqlens_q, + cu_seqlens_kv=cu_seqlens_kv, + ) if is_training: loss = out.sum() loss.backward() @@ -4300,7 +4302,7 @@ def get_dummy_cuda_rng_tracker() -> CudaRNGStatesTracker: tensor = 0.01 * torch.randn(tensor_shape, dtype=dtype, device="cuda") out_grad = tensor.view(*tensor.shape[:-2], -1) - with autocast(enabled=fp8_mha, recipe=fp8_recipe): + with torch.set_grad_enabled(is_training), autocast(enabled=fp8_mha, recipe=fp8_recipe): out = mha( hidden_states, attn_mask_type=config.attn_mask_type, @@ -4626,7 +4628,7 @@ def get_dummy_cuda_rng_tracker() -> CudaRNGStatesTracker: out_grad_shape_new = [*out_grad_shape[:-2], out_grad_shape[-2] * out_grad_shape[-1]] out_grad = torch.randn(out_grad_shape_new, dtype=dtype, device="cuda") - with autocast(enabled=fp8_dpa, recipe=fp8_recipe): + with torch.set_grad_enabled(is_training), autocast(enabled=fp8_dpa, recipe=fp8_recipe): out = dpa( inp[0], inp[1], @@ -4693,7 +4695,7 @@ def test_custom_mha_fp8_vs_f16(dtype, model, monkeypatch): fp8_dpa=True, ) fp8_meta["recipe"] = fp8_recipe - available_backends, _, fused_attn_backends = get_available_attention_backends( + _, _, fused_attn_backends = get_available_attention_backends( config, qkv_dtype=torch.float8_e4m3fn, nominal_dtype=dtype, @@ -4703,9 +4705,17 @@ def test_custom_mha_fp8_vs_f16(dtype, model, monkeypatch): is_training=is_training, deterministic=_deterministic, ) - flash_attn_supported, fused_attn_supported, unfused_attn_supported = available_backends - if not (fused_attn_backends and unfused_attn_supported): - pytest.skip("Not enough backends to run this test with.") + if not fused_attn_backends: + pytest.skip("FusedAttention is unavailable for this FP8 configuration.") + reference_backends, _, _ = get_available_attention_backends( + config, + qkv_dtype=dtype, + qkv_layout="bs3hd", + is_training=is_training, + deterministic=_deterministic, + ) + if not reference_backends[2]: + pytest.skip("UnfusedDotProductAttention is unavailable for the model-precision reference.") fused_attn_fwd_fp8, fused_attn_bwd_fp8 = _run_custom_mha_fp8(dtype, config, "FusedAttention") unfused_attn_fwd_f16, unfused_attn_bwd_f16 = _run_ref_mha_f16( diff --git a/tests/pytorch/utils.py b/tests/pytorch/utils.py index b3a4c624363..f39e2659a8b 100644 --- a/tests/pytorch/utils.py +++ b/tests/pytorch/utils.py @@ -399,8 +399,13 @@ def get_available_attention_backends( cp_size_a2a: int = 1, num_tokens_q: Optional[int] = None, num_tokens_kv: Optional[int] = None, + requires_backward: Optional[bool] = None, ) -> Tuple[List, List]: - """Check for all available attention backends that support a model configuration""" + """Check which backends support a model configuration. + + By default, only training requires backward support. Override requires_backward + to query evaluation with autograd or training without backward. + """ _, q_format, kv_format = get_qkv_format(qkv_layout, inference_params) if num_tokens_q is None: @@ -480,6 +485,7 @@ def test(): fp8=fp8, fp8_meta=fp8_meta, is_training=is_training, + requires_backward=is_training if requires_backward is None else requires_backward, inference_params=inference_params, softmax_type=config.softmax_type, return_max_logit=config.return_max_logit,