diff --git a/benchmarks/attention/benchmark_mxfp8_causal_own_block.py b/benchmarks/attention/benchmark_mxfp8_causal_own_block.py new file mode 100644 index 00000000000..7ccb9dc14e0 --- /dev/null +++ b/benchmarks/attention/benchmark_mxfp8_causal_own_block.py @@ -0,0 +1,218 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""Benchmark causal MXFP8 attention with ``fp8_dpa_causal_own_block`` against BF16 and the +ordinary MXFP8 path. + +Every path runs ``DotProductAttention`` with FusedAttention (cuDNN) on identical BF16 inputs: + +* ``bf16``: BF16 cuDNN causal attention; +* ``mxfp8``: ``MXFP8BlockScaling(fp8_dpa=True)``, whose causal output can read future tokens + through shared V scales; +* ``mxfp8_own_block``: ``MXFP8BlockScaling(fp8_dpa=True, fp8_dpa_causal_own_block=True)``. + +For each shape the script reports the first-call time (plan building and JIT compilation), +the steady-state median of the training forward and of forward + backward (CUDA events +around the module call, after warm-up, including quantization and Python launch overhead), +and the slowest GPU kernels of one forward + backward from ``torch.profiler``. The output relative L2 error against BF16 +is printed alongside, so a timing is never read without its accuracy. + +Example:: + + python benchmarks/attention/benchmark_mxfp8_causal_own_block.py \\ + --batch 2 --heads 16 --kv-heads 16 --seqlens 2048 4096 8192 +""" + +import argparse +import os +import time + +import torch +from torch.profiler import ProfilerActivity, profile + +from transformer_engine.common import recipe +from transformer_engine.pytorch import DotProductAttention, autocast +from transformer_engine.pytorch.attention.dot_product_attention import _attention_backends + +HEAD_DIM = 128 +PATHS = ("bf16", "mxfp8", "mxfp8_own_block") + + +def recipe_for(path): + if path == "bf16": + return None + return recipe.MXFP8BlockScaling( + fp8_dpa=True, fp8_dpa_causal_own_block=path == "mxfp8_own_block" + ) + + +def make_inputs(args, seqlen): + generator = torch.Generator(device="cuda").manual_seed(1234) + if args.format == "bshd": + q_shape = (args.batch, seqlen, args.heads, HEAD_DIM) + else: + q_shape = (seqlen, args.batch, args.heads, HEAD_DIM) + kv_shape = (*q_shape[:2], args.kv_heads, HEAD_DIM) + q, k, v = ( + torch.randn( + shape, dtype=torch.bfloat16, device="cuda", generator=generator + ).requires_grad_() + for shape in (q_shape, kv_shape, kv_shape) + ) + out_grad = torch.randn( + (*q_shape[:2], args.heads * HEAD_DIM), + dtype=torch.bfloat16, + device="cuda", + generator=generator, + ) + return (q, k, v), out_grad + + +def make_step(args, path, inputs, out_grad): + """Return a (forward, forward_backward) pair of closures for one path.""" + module = DotProductAttention( + args.heads, + HEAD_DIM, + num_gqa_groups=args.kv_heads, + attention_type="self", + qkv_format=args.format, + ).cuda() + fp8_recipe = recipe_for(path) + + def forward(): + with autocast(enabled=fp8_recipe is not None, recipe=fp8_recipe): + return module(*inputs, qkv_format=args.format, attn_mask_type="causal") + + def forward_backward(): + for x in inputs: + x.grad = None # write the gradients instead of accumulating them + out = forward() + out.backward(out_grad) + return out + + return forward, forward_backward + + +def median_ms(fn, warmup, iters): + for _ in range(warmup): + fn() + starts = [torch.cuda.Event(enable_timing=True) for _ in range(iters)] + ends = [torch.cuda.Event(enable_timing=True) for _ in range(iters)] + torch.cuda.synchronize() + for start, end in zip(starts, ends): + start.record() + fn() + end.record() + torch.cuda.synchronize() + times = sorted(start.elapsed_time(end) for start, end in zip(starts, ends)) + return times[len(times) // 2] + + +def kernel_breakdown(fn, iters, top): + """The ``top`` GPU kernels of ``fn`` by time, as (name, ms per call), from torch.profiler.""" + torch.cuda.synchronize() + with profile(activities=[ProfilerActivity.CUDA]) as prof: + for _ in range(iters): + fn() + torch.cuda.synchronize() + kernels = [ + (event.key, event.device_time_total / 1e3 / iters) + for event in prof.key_averages() + if event.device_type == torch.autograd.DeviceType.CUDA + ] + kernels.sort(key=lambda kernel: -kernel[1]) + total = sum(ms for _, ms in kernels) + return kernels[:top], total + + +def causal_tflops(args, seqlen, ms, backward): + # Causal attention reads half of the score matrix; the backward costs 2.5x the forward. + flops = 4 * args.batch * args.heads * seqlen * seqlen * HEAD_DIM / 2 + if backward: + flops *= 3.5 + return flops / (ms * 1e-3) / 1e12 + + +def main(): + parser = argparse.ArgumentParser(description=__doc__.split("\n\n")[0]) + parser.add_argument("--batch", type=int, default=2) + parser.add_argument("--heads", type=int, default=16) + parser.add_argument("--kv-heads", type=int, default=16) + parser.add_argument("--seqlens", type=int, nargs="+", default=[1024, 2048, 4096, 8192]) + parser.add_argument("--format", choices=("bshd", "sbhd"), default="bshd") + parser.add_argument("--paths", nargs="+", choices=PATHS, default=list(PATHS)) + parser.add_argument("--warmup", type=int, default=10) + parser.add_argument("--iters", type=int, default=50) + parser.add_argument("--profile-iters", type=int, default=10) + parser.add_argument("--top-kernels", type=int, default=8) + args = parser.parse_args() + + os.environ["NVTE_FLASH_ATTN"] = "0" + os.environ["NVTE_UNFUSED_ATTN"] = "0" + props = torch.cuda.get_device_properties(torch.cuda.current_device()) + import cudnn # pylint: disable=import-outside-toplevel + import transformer_engine # pylint: disable=import-outside-toplevel + + print( + f"{props.name} (sm{props.major}{props.minor}), torch {torch.__version__}, " + f"CUDA {torch.version.cuda}, cuDNN {cudnn.backend_version()}, " + f"cuDNN Frontend {cudnn.__version__}, TE {transformer_engine.__version__}, " + f"NVTE_ALLOW_NONDETERMINISTIC_ALGO={os.getenv('NVTE_ALLOW_NONDETERMINISTIC_ALGO', '1')}" + ) + print( + f"batch={args.batch} heads={args.heads} kv_heads={args.kv_heads} d={HEAD_DIM} " + f"format={args.format} mask=causal; median of {args.iters} after {args.warmup} warm-up" + ) + + rows = [] + for seqlen in args.seqlens: + inputs, out_grad = make_inputs(args, seqlen) + reference = None + for path in args.paths: + _attention_backends["backend_selection_requires_update"] = True + forward, forward_backward = make_step(args, path, inputs, out_grad) + torch.cuda.synchronize() + start = time.perf_counter() + out = forward_backward() + torch.cuda.synchronize() + first_s = time.perf_counter() - start + if path == "bf16": + reference = out.detach().float() + error = ( + (out.detach().float() - reference).norm() / reference.norm() + if reference is not None + else float("nan") + ) + fwd_ms = median_ms(forward, args.warmup, args.iters) + fwd_bwd_ms = median_ms(forward_backward, args.warmup, args.iters) + kernels, kernel_ms = kernel_breakdown( + forward_backward, args.profile_iters, args.top_kernels + ) + rows.append((seqlen, path, first_s, fwd_ms, fwd_bwd_ms, error)) + print( + f"\ns={seqlen:<6} {path:<16} first call {first_s:7.2f} s | fwd {fwd_ms:7.3f} ms" + f" ({causal_tflops(args, seqlen, fwd_ms, False):6.1f} TFLOPS) | fwd+bwd" + f" {fwd_bwd_ms:7.3f} ms" + f" ({causal_tflops(args, seqlen, fwd_bwd_ms, True):6.1f} TFLOPS) | out rel L2 vs" + f" bf16 {error:.4f}" + ) + print(f" GPU kernel time per fwd+bwd: {kernel_ms:.3f} ms") + for name, ms in kernels: + print(f" {ms:8.3f} ms {name[:100]}") + + print("\nseqlen path fwd ms fwd+bwd ms fwd+bwd vs mxfp8 fwd+bwd vs bf16") + by_key = {(row[0], row[1]): row for row in rows} + for seqlen, path, _, fwd_ms, fwd_bwd_ms, _ in rows: + mxfp8 = by_key.get((seqlen, "mxfp8")) + bf16 = by_key.get((seqlen, "bf16")) + vs_mxfp8 = f"{fwd_bwd_ms / mxfp8[4]:.2f}x" if mxfp8 else "-" + vs_bf16 = f"{fwd_bwd_ms / bf16[4]:.2f}x" if bf16 else "-" + print( + f"{seqlen:<7} {path:<16} {fwd_ms:7.3f} {fwd_bwd_ms:10.3f} {vs_mxfp8:>16} " + f" {vs_bf16:>15}" + ) + + +if __name__ == "__main__": + main() diff --git a/build_tools/jax.py b/build_tools/jax.py index cd571d7da53..8f1f6eec2fd 100644 --- a/build_tools/jax.py +++ b/build_tools/jax.py @@ -6,7 +6,9 @@ import os import warnings +from importlib.metadata import PackageNotFoundError, version as distribution_version from pathlib import Path +from packaging import version import setuptools @@ -26,7 +28,17 @@ def install_requirements() -> List[str]: """Install dependencies for TE/JAX extensions.""" - return ["jax>=0.5.3", "flax>=0.7.1", "nvidia-cudnn-frontend>=1.31.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. + try: + frontend_version = distribution_version("nvidia-cudnn-frontend") + except PackageNotFoundError as exc: + raise RuntimeError( + "Building TE/JAX requires nvidia-cudnn-frontend>=1.31.0 to be installed." + ) from exc + if version.parse(frontend_version) < version.parse("1.31.0"): + raise RuntimeError("Building TE/JAX requires nvidia-cudnn-frontend>=1.31.0.") + return ["jax>=0.5.3", "flax>=0.7.1", f"nvidia-cudnn-frontend=={frontend_version}", "triton"] def test_requirements() -> List[str]: 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/docs/api/jax.rst b/docs/api/jax.rst index 99282be422c..9157ff16d9a 100644 --- a/docs/api/jax.rst +++ b/docs/api/jax.rst @@ -60,7 +60,6 @@ Modules .. autoapifunction:: transformer_engine.jax.flax.extend_logical_axis_rules - Mixture of Experts ------------------ Router, routing kernels, and grouped dense for MoE layers. See diff --git a/docs/envvars.rst b/docs/envvars.rst index a0a64b3f637..a2478588bd1 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 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 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 @@ -282,13 +282,23 @@ 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: 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_FP8_DPA_CAUSAL_OWN_BLOCK + + :Type: ``int`` (0 or 1) + :Default: ``0`` + :Description: Default of ``MXFP8BlockScaling.fp8_dpa_causal_own_block`` in PyTorch. ``1`` = causal MXFP8 DotProductAttention keeps every output and gradient independent of later tokens: cuDNN runs MXFP8 attention over completed 32-token blocks, a Triton kernel adds each query's own block with V scaled per token, and the backward runs in FP16/BF16 on the decoded operands. Overrides ``NVTE_FP8_DPA_BWD`` for these calls. ``0`` = ordinary causal MXFP8 attention, whose 32-token columnwise V scales can carry a later token into an earlier output. .. envvar:: NVTE_DPA_FP8CS_O_in_F16 :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/docs/examples/jax/attention.rst b/docs/examples/jax/attention.rst index f82a6ffe621..17aed37f453 100644 --- a/docs/examples/jax/attention.rst +++ b/docs/examples/jax/attention.rst @@ -16,6 +16,24 @@ 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. +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 +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/docs/installation.rst b/docs/installation.rst index f45063191d2..c2e3803ce36 100644 --- a/docs/installation.rst +++ b/docs/installation.rst @@ -17,6 +17,11 @@ Prerequisites 4. `cuDNN 9.12 `__ or later. 5. `cuDNN Frontend 1.31.0 `__ or later. +PyTorch and JAX cuDNN attention require ``nvidia-cudnn-frontend`` version +1.31.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 `NVTE_CUDA_INCLUDE_PATH` in the environment. 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/jax/test_distributed_fused_attn.py b/tests/jax/test_distributed_fused_attn.py index e9ae8cd3e4a..1643197dd18 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,175 @@ 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( + "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() + 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 + 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(_unpack_mxfp8(xs, layout), mask) + expected = jax.jit(reference)(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) + + 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, + 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) * local_weight).sum()), + in_shardings=((sharding,) * len(inputs),), + )(inputs) + 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 4e5c2948a21..f05d9ed32f7 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, ) @@ -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, @@ -412,22 +437,212 @@ 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 != "" +@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) + monkeypatch.setattr(cudnn_attention._graph_cache, "entries", {}) + 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 == FusedAttnBackend.No_Backend: + 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] + + +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) + operands = (jnp.stack((q, q, q), axis=-3), empty, empty) + plan = SimpleNamespace( + graph=SimpleNamespace(workspace_size=0), + names=(), + input_indices=(), + input_offsets=(), + output_indices=(0, 0, 0, 4), + output_offsets=(0, 0, 0, 0), + metadata_batch=1, + ragged_stats=False, + multipliers=(128, 128, 128, 128, 2), + fp8=False, + fp8_o=False, + ) + 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=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, 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)) +@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. @@ -680,7 +895,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 ( @@ -2106,3 +2321,431 @@ def test_backward( seq_desc_format, is_self_attention=None, ) + + +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 == FusedAttnBackend.No_Backend: + pytest.skip(reason) + + +@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( + "layout", + [ + QKVLayout.BS3HD, + QKVLayout.BSHD_BS2HD, + QKVLayout.BSHD_BSHD_BSHD, + ], +) +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) + 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): + return attention_reference(unpack(operands)) + + def attention(operands): + return fused_attn( + operands, + None, + SequenceDescriptor(), + 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, + ) + + _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) + 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) + + +@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( + "kind,sm_arch", + [ + ("dense", None), + ("causal", None), + ("padding", None), + ("window", None), + ("ragged", None), + ("ragged_dense_stats", 120), + ("ragged_dense_stats_f16", 120), + ], +) +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: + 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 + + # 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: (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)) + 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 = SequenceDescriptor() + 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]) & ( + 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 + + _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, + attn_scale=scaling_factor, + return_max_logit=True, + ) + + def reference(arrays): + return attention_reference(arrays, mask, scaling_factor=scaling_factor) + + 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=scaling_factor, + 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 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, + None, + descriptor, + None, + attn_bias_type=AttnBiasType.NO_BIAS, + attn_mask_type=attn_mask, + qkv_layout=layout, + softmax_type=AttnSoftmaxType.VANILLA_SOFTMAX, + scaling_factor=scaling_factor, + 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) + expected_max = attention_reference( + arrays, mask, scaling_factor=scaling_factor, return_max_logit=True + ) + 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): + return attention_reference( + arrays, + causal[None], + bias=slopes[None, :, None, None] * positions[None, None, None, :], + ) + + def attention(arrays): + return fused_attn( + arrays, + None, + SequenceDescriptor(), + 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("fp8", [False, True]) +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") + module = DotProductAttention( + head_dim=64, + num_attention_heads=2, + attn_mask_type="no_mask", + transpose_batch_sequence=False, + softmax_type="learnable", + fp8=fp8, + ) + 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, + ) + variables = module.init(jax.random.PRNGKey(0), q, q, v, deterministic=True) + + def forward(variables, value): + return module.apply(variables, q, q, value, deterministic=True) + + actual = jax.jit(forward)(variables, v) + 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( + 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) + 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.""" + monkeypatch.setenv("NVTE_FP8_DPA_BWD", "1" if fp8_backward else "0") + 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( + 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) diff --git a/tests/jax/test_fused_attn_score_mod.py b/tests/jax/test_fused_attn_score_mod.py index 5a613101585..d47f2cb64e8 100644 --- a/tests/jax/test_fused_attn_score_mod.py +++ b/tests/jax/test_fused_attn_score_mod.py @@ -2,7 +2,11 @@ # # See LICENSE for license information. """Tests for cuDNN frontend score_mod fused attention.""" + from math import sqrt +from collections import Counter +from types import FunctionType, SimpleNamespace +from unittest.mock import Mock, call import jax import jax.numpy as jnp @@ -10,7 +14,9 @@ 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 ( + FusedAttnBackend, AttnBiasType, AttnMaskType, AttnSoftmaxType, @@ -18,10 +24,9 @@ ) 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 _CONFIG_TEST_SCALING_FACTOR = 1.0 / sqrt(_CONFIG_TEST_HEAD_DIM) _SCORE_MOD_MIN_CUDNN_VERSION = (9, 23) @@ -403,9 +408,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", ) @@ -652,6 +657,153 @@ def test_fused_attn_score_mod_config_splits_tensors_and_pass_by_value_scalars(): assert len(config.score_mod_scalars[0].value) == np.dtype(np.float32).itemsize +@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", cudnnGraphNotSupportedError=NotImplementedError) + graph = Mock(selected_engine=selected_engine) + graph.get_execution_plan_count.return_value = 2 + 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) + + if selected_engine is not None: + with pytest.raises(RuntimeError, match="Only the cuDNN backend implements score_mod"): + tex_attention._finalize_score_mod_graph(cudnn, graph, has_score_mod=True) + graph.serialize.assert_not_called() + else: + assert tex_attention._finalize_score_mod_graph(cudnn, graph, has_score_mod=True) == ( + 1, + b"graph", + 13100, + ) + + assert graph.deselect_engines.call_args_list == [ + call(["sdpa_fwd", "sdpa_bwd"]), + 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)] + + +@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.""" @@ -659,14 +811,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, ) @@ -674,6 +826,45 @@ 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) + ) + + +@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/jax/test_mxfp8_attention.py b/tests/jax/test_mxfp8_attention.py new file mode 100644 index 00000000000..f7024095155 --- /dev/null +++ b/tests/jax/test_mxfp8_attention.py @@ -0,0 +1,383 @@ +# 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 jax +import jax.numpy as jnp +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 ( + AttnBiasType, + AttnMaskType, + AttnSoftmaxType, + 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)) +@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))) + + +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) + + +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) + descriptor = options.pop("sequence_descriptor", SequenceDescriptor()) + seed = options.pop("seed", None) + return fused_attn(qkv, None, descriptor, seed, **options) + + +@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) + + +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_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() + 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), + ) + + +@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) + + +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)) + 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") + + +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)) + # 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( + 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/tests/pytorch/attention/test_attention.py b/tests/pytorch/attention/test_attention.py index 8238031847d..e14da92bf4b 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 @@ -25,7 +26,6 @@ DotProductAttention, MultiheadAttention, get_device_compute_capability, - Quantizer, is_fp8_available, is_bf16_available, ) @@ -44,27 +44,35 @@ ) 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.attention.dot_product_attention.backends import ( + UnfusedDotProductAttention, +) +from transformer_engine.pytorch.constants import ( FusedAttnBackend, + AttnMaskType, QKVFormat, QKVLayout, - fused_attn_bwd, + SoftmaxType, +) +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.quantized_tensor import ( + Quantizer, + prepare_for_saving, + restore_from_func_ctx, +) 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 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, - prepare_for_saving, - restore_from_func_ctx, -) _current_file = pathlib.Path(__file__).resolve() sys.path = [str(_current_file.parent.parent)] + sys.path @@ -160,23 +168,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 != "" @@ -203,20 +209,1165 @@ 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", (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 + + +def _own_block_query_config(): + """FusedAttentionParams for MXFP8BlockScaling(fp8_dpa_causal_own_block=True) selection.""" + own_block = recipe.MXFP8BlockScaling(fp8_dpa=True, fp8_dpa_causal_own_block=True) + spec = get_fused_attn_spec(own_block, torch.bfloat16, "bshd_bshd_bshd", cs_o_in_f16=True) + return FusedAttentionParams( + attn_mask_type=AttnMaskType["causal"], + scaling_mode=spec.scaling_mode, + attn_scale=128**-0.5, + qkv_dtype=spec.qkv, + o_dtype=spec.o, + do_dtype=spec.do, + dqkv_dtype=spec.dqkv, + nominal_dtype=DType.kBFloat16, + qkv_layout=QKVLayout[spec.qkv_layout], + o_format=QKVFormat[spec.o_format], + do_format=QKVFormat[spec.do_format], + dqkv_layout=QKVLayout[spec.dqkv_layout], + qkv_scale_inv_format=QKVFormat[spec.scale_inv_format], + do_scale_inv_format=QKVFormat[spec.scale_inv_format], + batch_size=2, + num_attn_heads=4, + num_gqa_groups=4, + head_dim_qk=128, + head_dim_v=128, + max_seqlen_q=256, + max_seqlen_kv=256, + causal_exclude_own_block=32, + ) + + +def _skip_unless_frontend_has_own_block(): + """Skip when the installed cuDNN Frontend's sdpa_mxfp8 lacks causal_exclude_own_block.""" + from transformer_engine.pytorch.cpp_extensions.cudnn_fused_attn import get_fused_attn_backend + + _, message = get_fused_attn_backend(_own_block_query_config()) + if "causal_exclude_own_block" in message: + pytest.skip(message) + + +@pytest.mark.parametrize("frontend_has_option", (True, False), ids=("frontend", "old_frontend")) +def test_fused_attn_query_mxfp8_causal_own_block(monkeypatch, frontend_has_option): + """Selection for MXFP8BlockScaling(fp8_dpa_causal_own_block=True) checks the graphs that run: + the MXFP8 forward with causal_exclude_own_block and a constant scale, and the + model-precision backward. A cuDNN Frontend whose sdpa_mxfp8 lacks the option is rejected + at selection instead of failing in the first forward.""" + import cudnn + from transformer_engine.pytorch.cpp_extensions import cudnn_fused_attn as implementation + + monkeypatch.delenv("NVTE_FP8_DPA_BWD", raising=False) + monkeypatch.setattr(torch.cuda, "get_device_capability", lambda *_args: (10, 0)) + config = _own_block_query_config() + + def sdpa_mxfp8(*_args, **kwargs): + if not frontend_has_option and "causal_exclude_own_block" in kwargs: + raise TypeError( + "sdpa_mxfp8() got an unexpected keyword argument 'causal_exclude_own_block'" + ) + return Mock(), Mock(), Mock() + + queries, graphs = [], [] + + def entry(attn): + queries.append(attn) + graph = Mock() + graph.sdpa_mxfp8.side_effect = sdpa_mxfp8 + graph.sdpa_backward.return_value = (Mock(), Mock(), Mock()) + graphs.append(graph) + implementation.build_attention_graph(cudnn, graph, attn.description()) + + monkeypatch.setattr(implementation._AttentionGraph, "entry", entry) + backend, message = implementation.get_fused_attn_backend(config) + if not frontend_has_option: + assert backend == FusedAttnBackend.No_Backend + assert "causal_exclude_own_block" in message + assert len(queries) == 1 + return + assert backend == FusedAttnBackend.FP8 + assert message == "" + forward, backward = queries + assert forward.mxfp8 and "attn_scale" not in forward.specs + assert forward.options().attn_scale == config.attn_scale + kwargs = graphs[0].sdpa_mxfp8.call_args.kwargs + assert kwargs["causal_exclude_own_block"] == 32 and kwargs["attn_scale"] == config.attn_scale + expected = implementation._query_graph( + replace( + config, + qkv_dtype=DType.kBFloat16, + o_dtype=DType.kBFloat16, + do_dtype=DType.kBFloat16, + dqkv_dtype=DType.kBFloat16, + causal_exclude_own_block=0, + ), + True, + ) + assert backward.description() == expected.description() + assert not backward.fp8 and not any(name.startswith("descale_") for name in backward.specs) + graphs[1].sdpa_backward.assert_called_once() + graphs[1].sdpa_mxfp8_backward.assert_not_called() + + +@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 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 + 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) + 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(), reference.reshape(shape), rtol=rtol, atol=atol) + + +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_matches_runtime_graph(layout, backward, dropout): + """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 + 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="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, + 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"]], + softmax_type=SoftmaxType["learnable"], + return_max_logit=not backward, + dropout=dropout, + ) + 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=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 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) + 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")) +@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 + + monkeypatch.setenv("NVTE_FLASH_ATTN", "0") + monkeypatch.setenv("NVTE_FUSED_ATTN", "1") + monkeypatch.setenv("NVTE_UNFUSED_ATTN", "0") + 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 + 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._graph_cache.entries) == 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", + ( + (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("kv_format", ("thd", "bshd", "sbhd")) +@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) * (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( + 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("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.""" + 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._graph_cache, "entries", {}) + 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", + ( + ({"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, + "head_dim_v": 192, + "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 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 import cudnn_fused_attn as implementation + + monkeypatch.setattr(cudnn, "backend_version", lambda: version) + monkeypatch.setattr(torch.cuda, "get_device_capability", lambda *_args: capability) + + def query(*_args): + raise AssertionError("Excluded configuration reached descriptor construction") + + monkeypatch.setattr(implementation, "_query_graph", query) + 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.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,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, requires_backward, 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, + requires_backward=requires_backward, + 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) + 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: + assert backend == FusedAttnBackend.F16_arbitrary_seqlen + 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,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 + + 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, + requires_backward=requires_backward, + 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) + 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: + 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.""" + 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", [ @@ -2999,8 +4150,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) @@ -3017,8 +4168,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 @@ -3040,12 +4191,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 @@ -3283,7 +4435,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: @@ -3301,8 +4455,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": @@ -3599,6 +4753,128 @@ def get_dummy_cuda_rng_tracker() -> CudaRNGStatesTracker: return out, (None, None, None, d_softmax_offset) +@pytest.mark.parametrize( + "change,reason", + [ + ({}, ""), + ({"sm_arch": 90}, "SM100 or SM103"), + ({"sm_arch": 120}, "SM100 or SM103"), + ({"attn_mask_type": "padding_causal"}, "top-left causal"), + ({"attn_mask_type": "causal_bottom_right"}, "top-left causal"), + ({"window_size": (128, 0)}, "without window"), + ({"qkv_layout": "bs3hd"}, "separate BSHD or SBHD"), + ({"qkv_layout": "thd_thd_thd"}, "separate BSHD or SBHD"), + ({"head_dim_qk": 64, "head_dim_v": 64}, "head dimension 128"), + ({"seqlen_kv": 1024}, "equal query and key"), + ({"seqlen_q": 1000, "seqlen_kv": 1000}, "divisible by 128"), + ({"dropout": 0.1}, "dropout"), + ({"bias_type": "post_scale_bias"}, "bias"), + ({"softmax_type": "learnable"}, "softmax offset"), + ], +) +def test_mxfp8_causal_own_block_support_rules(change, reason): + """The configurations MXFP8BlockScaling(fp8_dpa_causal_own_block=True) serves.""" + from transformer_engine.common.cudnn.mxfp8_utils import causal_own_block_error + + config = dict( + sm_arch=100, + qkv_layout="bshd_bshd_bshd", + attn_mask_type="causal", + window_size=(-1, 0), + head_dim_qk=128, + head_dim_v=128, + seqlen_q=512, + seqlen_kv=512, + dropout=0.0, + bias_type="no_bias", + softmax_type="vanilla", + ) + error = causal_own_block_error(**{**config, **change}) + assert (reason in error) if reason else error == "" + + +def _run_dpa_causal(inputs, out_grad, qkv_format, num_heads, num_gqa_groups, fp8_recipe): + """Causal DotProductAttention on fixed BF16 inputs; returns the output and dQ/dK/dV.""" + q, k, v = (x.detach().clone().requires_grad_() for x in inputs) + _attention_backends["backend_selection_requires_update"] = True + dpa = DotProductAttention( + num_heads, 128, num_gqa_groups=num_gqa_groups, attention_type="self", qkv_format=qkv_format + ).cuda() + with autocast(enabled=fp8_recipe is not None, recipe=fp8_recipe): + out = dpa(q, k, v, qkv_format=qkv_format, attn_mask_type="causal") + out.backward(out_grad) + return out.detach(), (q.grad, k.grad, v.grad) + + +@pytest.mark.skipif(not fp8_attn_available, reason=reason_for_no_fp8_attn) +@pytest.mark.skipif( + get_device_compute_capability() not in ((10, 0), (10, 3)), + reason="MXFP8 causal own-block attention requires SM100 or SM103.", +) +@pytest.mark.parametrize("qkv_format", ["bshd", "sbhd"]) +@pytest.mark.parametrize("num_gqa_groups", [8, 2], ids=["mha", "gqa"]) +def test_dpa_mxfp8_causal_own_block(qkv_format, num_gqa_groups, monkeypatch): + """MXFP8BlockScaling(fp8_dpa_causal_own_block=True): causal MXFP8 attention in which no + output or gradient of a token depends on a later token. + + Checks the output and gradients against BF16 attention, then scales every token from + position ``cut`` on (mid-block, so the cut shares a 32-token V scale block with visible + tokens) by 2**10. Outputs and gradients before ``cut`` must stay bitwise identical, with + zero gradients after it under a prefix-only loss. The ordinary MXFP8 path must change + rows before ``cut`` under the same perturbation, which shows the check can see the leak. + """ + _skip_unless_frontend_has_own_block() + monkeypatch.setenv("NVTE_ALLOW_NONDETERMINISTIC_ALGO", "0") + monkeypatch.setenv("NVTE_FLASH_ATTN", "0") + monkeypatch.setenv("NVTE_UNFUSED_ATTN", "0") + b, s, h, cut = 2, 512, 8, 153 + own_block = recipe.MXFP8BlockScaling(fp8_dpa=True, fp8_dpa_causal_own_block=True) + ordinary = recipe.MXFP8BlockScaling(fp8_dpa=True) + shape = (b, s, h, 128) if qkv_format == "bshd" else (s, b, h, 128) + kv_shape = (*shape[:2], num_gqa_groups, 128) + seq = 1 if qkv_format == "bshd" else 0 + generator = torch.Generator(device="cuda").manual_seed(1234) + inputs = [ + torch.randn(x, dtype=torch.bfloat16, device="cuda", generator=generator) + for x in (shape, kv_shape, kv_shape) + ] + out_grad = 0.1 * torch.randn( + (*shape[:2], h * 128), dtype=torch.bfloat16, device="cuda", generator=generator + ) + run = lambda x, g, r: _run_dpa_causal(x, g, qkv_format, h, num_gqa_groups, r) # noqa: E731 + + out, grads = run(inputs, out_grad, own_block) + out_ref, grads_ref = run(inputs, out_grad, None) + for name, got, want in zip(("out", "dq", "dk", "dv"), (out, *grads), (out_ref, *grads_ref)): + compare_and_assert(got, want, f"own_block {name}", f"bf16 {name}", 5e-1, 5e-2, 0.11, True) + # B200 measured ~0.046 (out), ~0.052 (dq, dk), ~0.033 (dv); the ordinary MXFP8 path + # measures ~0.046 (out) and ~0.066 / ~0.049 for the gradients on the same inputs. + error = (got.float() - want.float()).norm() / want.float().norm() + assert error < 0.08, f"{name}: relative L2 error {error:.4f} against BF16" + + def prefix(x): + return x.narrow(seq, 0, cut) + + def future(x): + return x.narrow(seq, cut, s - cut) + + prefix_grad = out_grad.clone() + future(prefix_grad).zero_() + perturbed = [x.clone() for x in inputs] + for x in perturbed: + future(x).mul_(1024.0) + base_out, base_grads = run(inputs, prefix_grad, own_block) + new_out, new_grads = run(perturbed, prefix_grad, own_block) + assert torch.equal(prefix(base_out), prefix(new_out)), "an output read a future token" + for name, before, after in zip(("dq", "dk", "dv"), base_grads, new_grads): + assert torch.equal(prefix(before), prefix(after)), f"{name} read a future token" + assert future(after).count_nonzero() == 0, f"{name} is nonzero after the prefix" + + leaky_out, _ = run(inputs, prefix_grad, ordinary) + leaky_new, _ = run(perturbed, prefix_grad, ordinary) + assert not torch.equal(prefix(leaky_out), prefix(leaky_new)), "the control path did not leak" + + model_configs_fp8 = { # test: ModelConfig(b, sq, hq, dqk) "fp8_1": ModelConfig(1, 512, 1, 64), @@ -3616,7 +4892,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 @@ -3625,6 +4901,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 = {} @@ -3920,7 +5199,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/tests/pytorch/attention/test_attention_backend_selection.py b/tests/pytorch/attention/test_attention_backend_selection.py index 2e41e199f0f..1ce59ba919a 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 ( @@ -84,6 +86,143 @@ 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)) +@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") + 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): + """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( "version,backward,expect_flash", [ diff --git a/tests/pytorch/attention/test_flex_attention.py b/tests/pytorch/attention/test_flex_attention.py index beed4069917..5dbedec4985 100644 --- a/tests/pytorch/attention/test_flex_attention.py +++ b/tests/pytorch/attention/test_flex_attention.py @@ -5,13 +5,26 @@ 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 +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 +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() @@ -27,6 +40,89 @@ param_types.append(torch.bfloat16) +@pytest.mark.parametrize("has_score_mod", (False, True)) +@pytest.mark.parametrize("selected_engine", (None, "sdpa_fwd", "sdpa_bwd", "missing")) +def test_cudnn_score_mod_backend_guard(monkeypatch, has_score_mod, selected_engine): + """Exclude Python engines before planning and reject them if they are still selected.""" + cudnn = Mock(__version__="1.31.0") + graph = Mock(selected_engine=selected_engine) + if selected_engine == "missing": + del graph.selected_engine + graph.get_workspace_size.return_value = 0 + monkeypatch.setattr(cudnn_graph, "_import_cudnn_frontend", lambda: cudnn) + + if has_score_mod and selected_engine in ("sdpa_fwd", "sdpa_bwd"): + with pytest.raises(RuntimeError, match="Only the cuDNN backend implements score_mod"): + cudnn_graph._finalize_cudnn_graph(graph, has_score_mod=has_score_mod) + graph.get_workspace_size.assert_not_called() + else: + assert cudnn_graph._finalize_cudnn_graph(graph, has_score_mod=has_score_mod) == 1 + + if has_score_mod: + exclusion = call.deselect_engines(["sdpa_fwd", "sdpa_bwd"]) + graph.deselect_engines.assert_called_once_with(["sdpa_fwd", "sdpa_bwd"]) + 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])) + ) + else: + 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() @@ -156,6 +252,122 @@ 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 + + +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): + """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) + build = _mock_score_mod_build(monkeypatch) + + def get(callback): + return _get_score_mod_test_graph(callback, backward=direction == "bwd") + + 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.""" @@ -219,137 +431,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): @@ -359,91 +462,78 @@ 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 + 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 - 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( + +@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, - q, - k, - v, + *qkv, "bshd", "bshd", - 1.0, - score_mod.forward, + 0.125, + callback, + _score_mod_identity_bprop, 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, + False, ) - assert len(cache) == 0 - finally: - cache.clear() - cache.update(saved_cache) - - assert entry_0 is not entry_1 - assert len(build_entries) == 2 + 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.""" 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): @@ -455,7 +545,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( @@ -549,8 +641,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: @@ -705,3 +802,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/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/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/test_torch_compile.py b/tests/pytorch/test_torch_compile.py index 6dea6836187..9dc0b1dc7f3 100644 --- a/tests/pytorch/test_torch_compile.py +++ b/tests/pytorch/test_torch_compile.py @@ -1323,9 +1323,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): ( @@ -1403,9 +1404,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/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/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..5bad38383d7 --- /dev/null +++ b/transformer_engine/common/cudnn/attention.py @@ -0,0 +1,258 @@ +# 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 # Effective precision for this pass, resolved by the frontend. + 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 + use_cu_seqlens: 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 + causal_exclude_own_block: 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) + 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)} + + 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["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: + 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 + if options.max_total_seqlen_kv: + kwargs["max_total_seq_len_kv"] = options.max_total_seqlen_kv + else: + kwargs["generate_stats"] = options.generate_stats + if options.causal_exclude_own_block: + # MXFP8 forward only: the causal mask also drops each query's own aligned block of + # keys, so every columnwise V scale read covers visible tokens only. + kwargs["causal_exclude_own_block"] = options.causal_exclude_own_block + 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 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") + 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, + v, + *scales, + *(tensor(name) for name in ("descale_s", "scale_s", "scale_o")), + **kwargs + ) + names = ("o", "stats", "amax_s", "amax_o") + 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: + if not options.backward: + scales = [tensor(name) for name in ("descale_q", "descale_k", "descale_v")] + try: + outputs = graph.sdpa_mxfp8(q, k, v, *scales, **kwargs) + except TypeError as exc: + if options.causal_exclude_own_block and "causal_exclude_own_block" in str(exc): + raise ValueError( + "MXFP8 causal attention without future-token scale leakage requires a" + " cuDNN Frontend whose sdpa_mxfp8 accepts causal_exclude_own_block." + ) from exc + raise + 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: + # 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..33e3813b8bd --- /dev/null +++ b/transformer_engine/common/cudnn/graph.py @@ -0,0 +1,339 @@ +# 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 + + +@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.""" + + +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: + 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, *, 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) + run("build_operation_graph", graph.build_operation_graph) + if has_score_mod: + # Python SDPA engines can silently ignore score_mod; only cuDNN implements it. + run("deselect_engines", lambda: graph.deselect_engines(["sdpa_fwd", "sdpa_bwd"])) + run( + "create_execution_plans", + lambda: graph.create_execution_plans([cudnn.heur_mode.A]), + ) + 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 + + +def compile_graph(cudnn, graph, run_stage=None, *, has_score_mod=False): + """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)) + selected_engine = getattr(graph, "selected_engine", None) if has_score_mod else None + if selected_engine is not None: + raise RuntimeError( + f"cuDNN frontend {cudnn.__version__} selected {selected_engine} for a score_mod graph. " + "Only the cuDNN backend implements score_mod." + ) + 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 + ragged_offset_multiplier: int = 1 + + +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) + if spec.ragged_offset_multiplier != 1: + tensor.set_ragged_offset_multiplier(spec.ragged_offset_multiplier) + 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="", + ragged_offset_multiplier=1, +): + """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, + ragged_offset_multiplier=ragged_offset_multiplier, + ) + + +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( + (batch, heads, seqlen, 1), + stride, + "FLOAT", + output=output, + ragged_offset="stats_offset" if ragged else "", + ragged_offset_multiplier=ragged_offset_multiplier, + ) + + +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() + 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 _resolve_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/mxfp8_utils.py b/transformer_engine/common/cudnn/mxfp8_utils.py new file mode 100644 index 00000000000..f34b955297a --- /dev/null +++ b/transformer_engine/common/cudnn/mxfp8_utils.py @@ -0,0 +1,101 @@ +# 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 + +# Tokens that share a columnwise MXFP8 scale; MXFP8 causal own-block attention computes each +# query's own aligned block of this many keys separately. +OWN_BLOCK = 32 + + +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 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): + 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 causal_own_block_error( + sm_arch, + qkv_layout, + attn_mask_type, + window_size, + head_dim_qk, + head_dim_v, + seqlen_q, + seqlen_kv, + dropout, + bias_type, + softmax_type, +): + """Return why causal MXFP8 attention split at the 32-token scale block cannot serve a call. + + The completed-block half runs on cuDNN Frontend's SM100 MXFP8 forward engine + (``sdpa_mxfp8(causal_exclude_own_block=32)``); the own-block half and the backward + operand decode are Triton kernels for D = 128 and BSHD/SBHD payloads. + """ + if sm_arch not in (100, 103): + return "MXFP8 causal own-block attention requires SM100 or SM103." + if attn_mask_type != "causal" or tuple(window_size or (-1, 0)) not in ((-1, 0), (-1, -1)): + return "MXFP8 causal own-block attention requires a top-left causal mask without window." + if qkv_layout not in ("bshd_bshd_bshd", "sbhd_sbhd_sbhd"): + return "MXFP8 causal own-block attention requires separate BSHD or SBHD Q, K and V." + if head_dim_qk != 128 or head_dim_v != 128: + return "MXFP8 causal own-block attention requires head dimension 128." + if seqlen_q != seqlen_kv: + return "MXFP8 causal own-block attention requires equal query and key sequence lengths." + if seqlen_q % 128: + return "MXFP8 causal own-block attention requires sequence lengths divisible by 128." + if dropout: + return "MXFP8 causal own-block attention does not support dropout." + if bias_type != "no_bias" or softmax_type != "vanilla": + return "MXFP8 causal own-block attention does not support bias or a softmax offset." + 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/common/cudnn/score_mod.py b/transformer_engine/common/cudnn/score_mod.py new file mode 100644 index 00000000000..e48e35626cc --- /dev/null +++ b/transformer_engine/common/cudnn/score_mod.py @@ -0,0 +1,125 @@ +# 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 +import struct + + +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 _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 stateless functions by code/defaults and explicitly keyed callback owners.""" + 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__ + ): + 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 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/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/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/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/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/common/recipe/__init__.py b/transformer_engine/common/recipe/__init__.py index 258504f9443..a96089833f0 100644 --- a/transformer_engine/common/recipe/__init__.py +++ b/transformer_engine/common/recipe/__init__.py @@ -399,6 +399,22 @@ class MXFP8BlockScaling(Recipe): compute dtype (e.g. BF16/FP16/FP32) for backward. enable_2d_quantization : bool, default = False If set to `True`, 2D block scaling is used for weight tensors. + fp8_dpa_causal_own_block : bool, default = False + With `fp8_dpa=True` and a causal mask, keep each query's attention independent + of future tokens. A columnwise MXFP8 V scale is shared by 32 consecutive tokens, + so ordinary causal MXFP8 attention lets a later token in a query's own 32-token + block change how earlier, visible values are rounded. With this option, cuDNN + computes MXFP8 attention over completed 32-token blocks only, a Triton kernel + adds each query's own block using V scaled per token, and the backward runs in + high precision on the decoded MXFP8 Q/K and the per-token-scaled V (the + completed-block forward read V with its columnwise MXFP8 scales, so, as with + `NVTE_FP8_DPA_BWD=0`, the backward differentiates slightly different operands). + Requires SM100/SM103, a cuDNN Frontend whose `sdpa_mxfp8` accepts + `causal_exclude_own_block`, separate BSHD or SBHD Q/K/V with head dimension 128, + sequence lengths that are equal and divisible by 128, no dropout, bias, softmax + offset or context parallelism. Only FusedAttention runs it: when it cannot, no + other backend is used and DotProductAttention raises an error. Defaults to the + environment variable `NVTE_FP8_DPA_CAUSAL_OWN_BLOCK`. """ margin: int = 0 @@ -407,6 +423,7 @@ class MXFP8BlockScaling(Recipe): fp8_mha: bool = False backward_override: Optional[str] = os.getenv("NVTE_BACKWARD_OVERRIDE", None) enable_2d_quantization: bool = False + fp8_dpa_causal_own_block: bool = os.getenv("NVTE_FP8_DPA_CAUSAL_OWN_BLOCK", "0") == "1" def __post_init__(self) -> None: assert self.fp8_format in (Format.E4M3, Format.HYBRID), "Unsupported FP8 format." @@ -420,7 +437,8 @@ def _make_repr(self) -> str: f"margin={self.margin}, " f"format={str(self.fp8_format).split('.')[1]}, " f"backward_override={self.backward_override}, " - f"enable_2d_quantization={self.enable_2d_quantization}" + f"enable_2d_quantization={self.enable_2d_quantization}, " + f"fp8_dpa_causal_own_block={self.fp8_dpa_causal_own_block}" ) diff --git a/transformer_engine/common/triton/mxfp8_causal_attention.py b/transformer_engine/common/triton/mxfp8_causal_attention.py new file mode 100644 index 00000000000..d52584f9328 --- /dev/null +++ b/transformer_engine/common/triton/mxfp8_causal_attention.py @@ -0,0 +1,250 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""Triton kernels for causal MXFP8 attention without future-token scale leakage. + +A columnwise MXFP8 V scale is shared by 32 consecutive tokens, so under an ordinary causal +mask a query can read a V value whose scale was raised by a later token in its own 32-token +block (https://github.com/NVIDIA/TransformerEngine/issues/3248). The forward therefore +splits each query's keys in two: + +* keys in completed 32-token blocks, ``j < 32 * floor(i / 32)``: cuDNN's MXFP8 SDPA with + ``causal_exclude_own_block=32``, whose V scales cover visible tokens only; +* keys in the query's own block: :func:`own_block_attention_kernel`, which multiplies the + softmax weights by a token-normalized copy of V (one power-of-two amplitude per token, + from :func:`token_values_kernel`) and merges the result with the first partial through + the two log-sum-exps. + +:func:`decode_kernel` prepares the backward: BF16 Q and K decoded from their rowwise MXFP8 +bytes and V decoded from the token-normalized copy, so no backward operand depends on a +future token either. + +All kernels take D = 128, E4M3 payloads, and rowwise scales in the F8_128x4 layout of each +(batch, head) matrix: a 128-token tile holds 512 E8M0 bytes, token ``t`` and 32-feature +group ``g`` at ``(t % 32) * 16 + (t % 128 // 32) * 4 + g``. +""" + +import triton +import triton.language as tl + +ZERO_TOKEN = tl.constexpr(-1000) # amplitude exponent of an all-zero V row +LOG2_E = tl.constexpr(1.4426950408889634) +LOG2_E4M3_MAX = tl.constexpr(8.807354922057604) # log2(448) + + +@triton.jit +def _scale_offset(token, group, heads_index, num_tokens: tl.constexpr): + """Byte offset of a rowwise E8M0 scale in the F8_128x4 layout.""" + tile = heads_index * (num_tokens // 128) + token // 128 + return tile * 512 + (token % 32) * 16 + (token % 128 // 32) * 4 + group + + +@triton.jit +def _decode_scale(code): + """E8M0 code to its power-of-two value; code 0 is 2**-127 (an FP32 subnormal).""" + code = code.to(tl.int32) + return tl.where(code == 0, 0x00400000, code << 23).to(tl.float32, bitcast=True) + + +@triton.jit +def token_values_kernel( + V, + U, + E, + v_stride_b, + v_stride_s, + v_stride_h, + num_tokens: tl.constexpr, + num_kv_heads: tl.constexpr, +): + """Token-normalized V: ``U = V * 2**(8 - E)`` in E4M3, one exponent ``E`` per token. + + ``E = ceil(log2(max |V_t|))`` (bounded below by -126), so ``|U| <= 256`` and the cast to + E4M3 depends on the token's own values only. ``U`` is ``[B, H_kv, S, 128]`` and ``E`` is + ``[B, H_kv, S]``; an all-zero row stores ``U = 0`` and ``E = ZERO_TOKEN``. + """ + block, head, batch = tl.program_id(0), tl.program_id(1), tl.program_id(2) + t = block * 32 + tl.arange(0, 32) + d = tl.arange(0, 128) + v = tl.load( + V + batch * v_stride_b + t[:, None] * v_stride_s + head * v_stride_h + d[None, :] + ).to(tl.float32) + maximum = tl.max(tl.abs(v), 1) + bits = maximum.to(tl.int32, bitcast=True) + e = ((bits >> 23) & 255) - 127 + ((bits & 0x7FFFFF) != 0).to(tl.int32) + # Exact v * 2**(8 - e) as two normal power-of-two factors (8 - e lies in [-120, 135]). + shift = 8 - e + first = (((shift >> 1) + 127) << 23).to(tl.float32, bitcast=True) + second = ((shift - (shift >> 1) + 127) << 23).to(tl.float32, bitcast=True) + u = (v * first[:, None] * second[:, None]).to(tl.float8e4nv) + u = tl.where(maximum[:, None] == 0, 0, u) + row = (batch * num_kv_heads + head) * num_tokens + t + tl.store(U + row[:, None] * 128 + d[None, :], u) + tl.store(E + row, tl.where(maximum == 0, ZERO_TOKEN, e)) + + +@triton.jit +def own_block_attention_kernel( + Q, + K, + SQ, + SK, + U, + E, + O, + LSE, + q_stride_b, + q_stride_s, + q_stride_h, + k_stride_b, + k_stride_s, + k_stride_h, + o_stride_b, + o_stride_s, + o_stride_h, + attn_scale, + num_tokens: tl.constexpr, + num_heads: tl.constexpr, + num_kv_heads: tl.constexpr, + PACK: tl.constexpr, +): + """Attention of each query over its own 32-token block, merged into the completed-block + partial. + + ``O`` and ``LSE`` hold the completed-block output and natural-log softmax statistics on + entry (rows 0..31: O = 0, LSE = -inf) and the full causal result on exit; each program + reads its rows before overwriting them. One program covers 128 MMA rows: 128 / PACK + tokens for each of PACK query heads that share a KV head. + """ + tile, packed_head, batch = tl.program_id(0), tl.program_id(1), tl.program_id(2) + KEYS: tl.constexpr = 128 // PACK + GROUPS: tl.constexpr = KEYS // 32 + kv_head = packed_head // ((num_heads // num_kv_heads) // PACK) + row = tl.arange(0, 128) + head = packed_head * PACK + row // KEYS + token = tile * KEYS + row % KEYS + key = tile * KEYS + tl.arange(0, KEYS) + d = tl.arange(0, 128) + g = tl.arange(0, 4) + + # Issue the partial loads first; they overlap the two matrix products. + o_offset = batch * o_stride_b + token[:, None] * o_stride_s + head[:, None] * o_stride_h + partial = tl.load(O + o_offset + d[None, :]).to(tl.float32) + lse_offset = (batch * num_heads + head) * num_tokens + token + partial_lse = tl.load(LSE + lse_offset) + q = tl.load( + Q + batch * q_stride_b + token[:, None] * q_stride_s + head[:, None] * q_stride_h + d + ) + k = tl.load( + K + batch * k_stride_b + key[:, None] * k_stride_s + kv_head * k_stride_h + d[None, :] + ) + sq = tl.load( + SQ + + _scale_offset(token[:, None], g[None, :], batch * num_heads + head[:, None], num_tokens) + ) + sk = tl.load( + SK + _scale_offset(key[:, None], g[None, :], batch * num_kv_heads + kv_head, num_tokens) + ) + score = tl.dot_scaled(q, sq, "e4m3", tl.trans(k), sk, "e4m3") * attn_scale + + # Each row sees only its own 32-key group: select it so the softmax runs on 128 x 32. + j = tl.arange(0, 32) + pick = tl.arange(0, GROUPS)[None, :, None] == ((row % KEYS) // 32)[:, None, None] + if GROUPS == 1: + own_score = score + e = tl.load(E + (batch * num_kv_heads + kv_head) * num_tokens + key)[None, :] + else: + own_score = tl.sum(tl.where(pick, tl.reshape(score, (128, GROUPS, 32)), 0.0), 1) + e = tl.load( + E + + (batch * num_kv_heads + kv_head) * num_tokens + + (token[:, None] // 32) * 32 + + j[None, :] + ) + visible = j[None, :] <= (token % 32)[:, None] + own_score = tl.where(visible, own_score, -float("inf")) + m = tl.max(own_score, 1) + p = tl.exp(own_score - m[:, None]) + denominator = tl.sum(p, 1) + + # W = P * 2**E in E4M3 with one E8M0 scale per row: P V = (W * 2**-8) U exactly in the + # scales, with the only rounding in W and U themselves. + log_w = (own_score - m[:, None]) * LOG2_E + e + log_w = tl.where(visible & (e != ZERO_TOKEN), log_w, -float("inf")) + w_max = tl.max(log_w, 1) + w_exponent = tl.where(w_max == -float("inf"), 0, tl.ceil(w_max - LOG2_E4M3_MAX)) + w_exponent = tl.minimum(127, tl.maximum(-127, w_exponent)) + own_w = tl.exp2(log_w - w_exponent[:, None]) + if GROUPS == 1: + w = own_w.to(tl.float8e4nv) + else: + # Block-diagonal W: zero outside the row's own key group. + w = tl.reshape(tl.where(pick, own_w[:, None, :], 0.0), (128, KEYS)).to(tl.float8e4nv) + w_scale = tl.broadcast_to((w_exponent + 127).to(tl.uint8)[:, None], (128, GROUPS)) + u = tl.load(U + ((batch * num_kv_heads + kv_head) * num_tokens + key[:, None]) * 128 + d) + u_scale = tl.full((128, GROUPS), 127 - 8, tl.uint8) # 2**-8 for every 32 keys of U + numerator = tl.dot_scaled(w, w_scale, "e4m3", u, u_scale, "e4m3") + + # The partial is normalized, so exp(partial_lse) is its denominator. Rows 0..31 have no + # completed block; discard their partial before any arithmetic. + partial_lse = tl.where(token >= 32, partial_lse, -float("inf")) + partial = tl.where(token[:, None] >= 32, partial, 0.0) + common = tl.maximum(m, partial_lse) + a, z = tl.exp(partial_lse - common), tl.exp(m - common) + total = a + z * denominator + result = (a[:, None] * partial + z[:, None] * numerator) / total[:, None] + tl.store(O + o_offset + d[None, :], result.to(O.dtype.element_ty)) + tl.store(LSE + lse_offset, common + tl.log(total)) + + +@triton.jit +def decode_kernel( + Q8, + SQ, + K8, + SK, + U, + E, + Q, + K, + V, + q_stride_b, + q_stride_s, + q_stride_h, + k_stride_b, + k_stride_s, + k_stride_h, + num_batches: tl.constexpr, + num_tokens: tl.constexpr, + num_heads: tl.constexpr, + num_kv_heads: tl.constexpr, + BLOCK: tl.constexpr, +): + """Backward operands: Q and K decoded from rowwise MXFP8, V from the token-normalized copy. + + Programs ``[0, B * H)`` along axis 1 decode Q; the next ``B * H_kv`` decode K and V. The + decoded tensors use the strides of the MXFP8 payloads (Q like Q8; K and V like K8). Every + value is exact in BF16 or FP16 barring underflow: an E4M3 value times a power of two. + """ + t = tl.program_id(0) * BLOCK + tl.arange(0, BLOCK) + matrix = tl.program_id(1) + d = tl.arange(0, 128) + g = d // 32 + if matrix < num_batches * num_heads: + batch, head = matrix // num_heads, matrix % num_heads + offset = batch * q_stride_b + t[:, None] * q_stride_s + head * q_stride_h + d[None, :] + code = tl.load(SQ + _scale_offset(t[:, None], g[None, :], matrix, num_tokens)) + q = tl.load(Q8 + offset).to(tl.float32) * _decode_scale(code) + tl.store(Q + offset, q.to(Q.dtype.element_ty)) + else: + matrix -= num_batches * num_heads + batch, head = matrix // num_kv_heads, matrix % num_kv_heads + offset = batch * k_stride_b + t[:, None] * k_stride_s + head * k_stride_h + d[None, :] + code = tl.load(SK + _scale_offset(t[:, None], g[None, :], matrix, num_tokens)) + k = tl.load(K8 + offset).to(tl.float32) * _decode_scale(code) + tl.store(K + offset, k.to(K.dtype.element_ty)) + row = matrix * num_tokens + t + amplitude = tl.exp2((tl.load(E + row) - 8).to(tl.float32)) + v = tl.load(U + row[:, None] * 128 + d[None, :]).to(tl.float32) * amplitude[:, None] + tl.store(V + offset, v.to(V.dtype.element_ty)) diff --git a/transformer_engine/jax/attention.py b/transformer_engine/jax/attention.py index edf418f15ea..cd5d6626472 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 @@ -17,12 +17,19 @@ 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 +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) @@ -33,6 +40,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 +147,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 +356,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 | str = False, ): """ To check whether the fused attention kernel is supported. @@ -384,6 +393,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 +1152,7 @@ def _legacy_fused_attn( context_parallel_axis: str = "", softmax_offset: Optional[jnp.ndarray] = None, return_max_logit: bool = False, + fp8: bool | str = False, ): """ Perform non-THD (non-packed) cuDNN fused attention. @@ -1236,6 +1247,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 +1274,7 @@ def fused_attn_thd( context_parallel_axis: str = "", softmax_offset: Optional[jnp.ndarray] = None, return_max_logit: bool = False, + fp8: bool | str = False, ): """ Deprecated THD fused attn, please use fusd_attn with SequenceDescriptor @@ -1317,6 +1330,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 +1338,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 +1361,8 @@ def _fused_attn( context_checkpoint_name: str = "context", stripe_size: int | None = None, return_max_logit: bool = False, + fp8: bool | str = False, + quantized_qkv=None, ): output, _ = _fused_attn_fwd_rule( qkv, @@ -1369,6 +1385,8 @@ def _fused_attn( context_checkpoint_name=context_checkpoint_name, stripe_size=stripe_size, return_max_logit=return_max_logit, + fp8=fp8, + quantized_qkv=quantized_qkv, ) return output @@ -1394,6 +1412,8 @@ def _fused_attn_fwd_rule( context_checkpoint_name, stripe_size, return_max_logit, + fp8, + quantized_qkv=None, ): output, softmax_aux, rng_state, max_logit = tex.fused_attn_fwd( qkv, @@ -1415,6 +1435,8 @@ def _fused_attn_fwd_rule( context_parallel_axis=context_parallel_axis, 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) @@ -1429,6 +1451,7 @@ def _fused_attn_fwd_rule( rng_state, softmax_offset, output, + quantized_qkv, ) @@ -1448,6 +1471,7 @@ def _fused_attn_bwd_rule( context_checkpoint_name, stripe_size, return_max_logit, + fp8, ctx, dz, ): @@ -1462,6 +1486,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, @@ -1485,8 +1510,10 @@ 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, + quantized_qkv=quantized_qkv, ) - 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 @@ -1496,6 +1523,7 @@ def _fused_attn_bwd_rule( grad_softmax_offset, None, None, + None, ) @@ -1555,6 +1583,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], @@ -1580,15 +1625,21 @@ 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 | str = False, + quantized_qkv=None, ): """ 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 @@ -1636,12 +1687,23 @@ 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): False disables FP8 attention; True enables FP8 attention + 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`` + 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: - 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): @@ -1673,6 +1735,32 @@ 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 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 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 @@ -1686,6 +1774,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 +1838,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 +1870,7 @@ def fused_attn( context_checkpoint_name=context_checkpoint_name, 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 3a1de7bd677..2c04ff4ee2a 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 @@ -12,22 +13,14 @@ 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.common.cudnn.graph import uses_ragged_attention_graph from transformer_engine.jax.attention import ( + FusedAttnBackend, AttnBiasType, AttnMaskType, AttnSoftmaxType, @@ -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, @@ -59,7 +49,6 @@ with_sharding_constraint, ) - __all__ = [ "FusedAttnHelper", "fused_attn_fwd", @@ -118,6 +107,7 @@ def setup_logging(): "stripe_size", "return_max_logit", "allow_fast_causal_path", + "fp8", ], ) @dataclass(frozen=True) @@ -144,6 +134,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 | str = False @property def effective_window_size(self) -> Tuple[int, int]: @@ -153,64 +144,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,11 +174,12 @@ class FusedAttnHelper: bias_seqlen_q: Optional[int] = None bias_seqlen_kv: Optional[int] = None max_segments_per_seq: int = 1 + fp8: bool | str = False 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. @@ -257,67 +191,80 @@ 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.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.") - 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 = FusedAttnBackend.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, + 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, ) - ) + 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.check_attention_support("fwd", avals, config) + if self.is_training: + cudnn_attention.check_attention_support("bwd", avals, config) + backend = ( + FusedAttnBackend.FP8 if self.fp8 else FusedAttnBackend.F16_arbitrary_seqlen + ) + except ValueError as exc: + message = str(exc) 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 @@ -401,7 +348,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 +486,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 ( + FusedAttnBackend.F16_arbitrary_seqlen, + FusedAttnBackend.FP8, + ): if config.qkv_layout.is_thd(): softmax_shape = (*batch_shape, q_max_seqlen, attn_heads, 1) else: @@ -550,8 +502,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) @@ -571,42 +523,10 @@ 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: - 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 +576,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) - - 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("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 +596,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 @@ -839,8 +717,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. """ @@ -879,7 +757,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 get_all_device_compute_capability() + 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): @@ -1114,57 +994,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() - - 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, - ) + attn_heads = FusedAttnHelper.parse_qkv_aval(q_aval, k_aval, v_aval, config.qkv_layout)[3] - 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 +1059,12 @@ def lowering( *, config, ): - """ - Fused attention bwd lowering rules - """ - q_aval, k_aval, v_aval, bias_aval, *_ = ctx.avals_in + """Lower Python-built cuDNN attention through the JAX-owned graph executor.""" - ( - 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) + def execute(*args): + return cudnn_attention.execute("bwd", args, config, ctx.avals_out) - 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 - - return ffi.ffi_lowering(FusedAttnBwdPrimitive.name)( + return mlir.lower_fun(execute, multiple_results=True)( ctx, q, k, @@ -1266,28 +1082,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 +1509,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 +1531,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 +2604,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: @@ -3089,7 +2887,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. @@ -3314,7 +3112,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) @@ -3553,7 +3351,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. @@ -3685,7 +3483,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) @@ -3735,6 +3533,8 @@ def fused_attn_fwd( context_parallel_axis: str = "", 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. @@ -3798,7 +3598,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 +3651,7 @@ def fused_attn_fwd( cp_striped_window_size=None, stripe_size=stripe_size, return_max_logit=return_max_logit, + fp8=fp8, ) primitive = None @@ -3868,14 +3669,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) @@ -3902,6 +3710,8 @@ def fused_attn_bwd( context_parallel_causal_load_balanced: bool = False, 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. @@ -3967,7 +3777,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 +3805,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 +3825,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, @@ -4041,7 +3855,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, @@ -4050,6 +3864,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 new file mode 100644 index 00000000000..3c0ef41a062 --- /dev/null +++ b/transformer_engine/jax/cpp_extensions/cudnn_attention.py @@ -0,0 +1,616 @@ +# 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 transformer_engine.common.cudnn.attention import ( + AttentionOptions, + AttentionGraphSpec, + build_attention_graph, +) +from transformer_engine.common.cudnn.graph import ( + TensorSpec, + GraphCacheEntry, + bucket as _bucket, + row_major_stride as _stride, + create_graph, + uses_ragged_attention_graph, + mxfp8_scale_spec, +) +from transformer_engine.common.cudnn import mxfp8_utils as mxfp8_rules + +from .cudnn_graph import ( + import_cudnn, + check_frontend_version, + check_graph_support, + compile_graph, + serialize_graph, + graph_attributes, + _graph_cache, + _diagnostics, +) +from . import mxfp8_attention + + +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(dtype): + return { + 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 _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 + fp8_o: bool = False + + +@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 + + +def _site(spec): + return "fp8" if spec.options.fp8 else "f16", "bwd" if spec.options.backward else "fwd" + + +def _build_supported_graph(spec, capabilities): + cudnn = import_cudnn() + backend, direction = _site(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) + ) + return GraphCacheEntry(built) + + +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"))), + fp8_backward, + fp8_o, + ) + + +def _get_supported_graph(spec, capabilities, frontend_version, backend_version): + """Share support-query graphs with compilation, and report cache reuse.""" + 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): + """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=None) +def _build_graph( + direction, shapes, config, capabilities, nondeterministic, fp8_backward, fp8_o=False +): + description = _describe_graph( + direction, shapes, config, capabilities, nondeterministic, fp8_backward, fp8_o + ) + 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) + _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, + 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, + fp8_o, + ) + + +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) + backward = direction == "bwd" + ragged = config.qkv_layout.is_thd() + 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. + 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) + 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.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: + 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) + 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): + 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(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)) + + 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) + 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 = (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) + 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", + kdim, + kstride, + qkv_dtype, + 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( + "v", + vdim, + vstride, + qkv_dtype, + 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) + if ragged: + for name in ("offset_q", "offset_k", "offset_v", "offset_o", "offset_stats"): + 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) + stats_offset = "offset_stats" if ragged_stats else "" + o_offset = "offset_o" if ragged else "" + left, right = config.effective_window_size + if config.attn_mask_type.is_causal(): + right = 0 + 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 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): + output("dbias", bias_shape, _stride(bias_shape), 3) + if config.softmax_type.name != "VANILLA_SOFTMAX": + tensor("sink", (1, h, 1, 1), (h, 1, 1, 1), "FLOAT") + if backward: + output("dsink", (1, h, 1, 1), (h, 1, 1, 1), 4, dtype="FLOAT") + constant("attn_scale", config.scaling_factor) + if fp8: + 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, "FP8_E4M3" if fp8_o 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 = ( + 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 name, dim, stride, slot, offset, ragged_offset in zip( + ("dq", "dk", "dv"), + (qdim, kdim, vdim), + 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 mxfp8: + 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) + 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") + 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") + 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: + output("max", stats_dim, stats_stride, 2, dtype="FLOAT", ragged_offset=stats_offset) + 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, + 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 _AttentionDescription( + AttentionGraphSpec.from_dict(options, specs), + tuple(inputs), + tuple(outputs), + tuple(scalars), + 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 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, :] + 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, 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] + 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 == "mxfp8": + if plan.fp8: + 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) + if prepared_qkv is None + else prepared_qkv + ) + elif 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 and config.fp8 == "mxfp8": + 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 + 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 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) + 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), + ) + # 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) + 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..6295aa42e0c --- /dev/null +++ b/transformer_engine/jax/cpp_extensions/cudnn_graph.py @@ -0,0 +1,177 @@ +# 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 +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 as _check_graph_support, + compile_graph as _compile_graph, + GraphCache, + CacheDiagnostics, +) + +_graph_cache = GraphCache() +_diagnostics = CacheDiagnostics("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.""" + cudnn = _import_cudnn() + 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 _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()) + ] + 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, run_stage=None, *, has_score_mod=False): + """Validate, select and compile a supported execution plan.""" + 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) + + +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..495945dae38 100644 --- a/transformer_engine/jax/cpp_extensions/flex_attention.py +++ b/transformer_engine/jax/cpp_extensions/flex_attention.py @@ -3,19 +3,49 @@ # See LICENSE for license information. """cuDNN frontend score_mod fused attention helpers.""" -import hashlib -import importlib -import inspect import os from dataclasses import dataclass -from typing import Any, Callable, Dict, Mapping, Optional, Sequence, Tuple +from functools import partial +from typing import Any, Callable, Mapping, Optional, Sequence, Tuple import jax import jax.numpy as jnp import numpy as np from jax import ffi -import transformer_engine_jax +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, + 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, + graph_attributes, + _graph_cache, + _diagnostics, +) __all__ = [ "FusedAttnScoreModHelper", @@ -110,14 +140,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: @@ -125,86 +148,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) @@ -254,56 +209,21 @@ 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 -_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 _SCORE_MOD_BPROP_SCALAR_UID_BASE = 4000 -_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, ) @@ -477,102 +397,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 - - def wrapped_score_mod(sdpa_graph, score_tensor): - return score_mod(sdpa_graph, score_tensor, graph_tensors) - - 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, @@ -583,6 +407,7 @@ def _graph_cache_key( ): return None return ( + "score_mod", direction, config, tuple((tuple(aval.shape), _dtype_name(aval.dtype)) for aval in avals), @@ -593,41 +418,19 @@ 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 _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() - - 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, @@ -636,40 +439,39 @@ 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), + ) + _score_mod_cache_event("fwd", "CREATE_GRAPH") + workspace_size, data, frontend_version = _finalize_score_mod_graph( + 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=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, ) @@ -687,43 +489,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, @@ -732,56 +499,53 @@ 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, - ) - - 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) + 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), + ) + _score_mod_cache_event("bwd", "CREATE_GRAPH") + workspace_size, data, frontend_version = _finalize_score_mod_graph( + 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=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], ) @@ -797,14 +561,11 @@ 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), + event=partial(_score_mod_cache_event, "fwd"), + ) batch, q_seqlen, q_heads, _ = q.shape _, _, _, v_head_dim = v.shape @@ -813,22 +574,14 @@ 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, 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 @@ -851,29 +604,23 @@ 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] + ), + event=partial(_score_mod_cache_event, "bwd"), + ) dq = jax.ShapeDtypeStruct(q.shape, q.dtype) dk = jax.ShapeDtypeStruct(k.shape, k.dtype) 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, @@ -884,15 +631,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 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..cb222c93e7b --- /dev/null +++ b/transformer_engine/jax/cpp_extensions/mxfp8_attention.py @@ -0,0 +1,306 @@ +# 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.""" + +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_utils 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 + + +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.""" + # 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, + q_layout=layout, + ) + result = quantize(value, quantizer) + if isinstance(result, ScaledTensor2x): + return result.rowwise_tensor if rowwise else None, result.colwise_tensor + return result, None + + +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:]) + row, col = quantize_operand(value.transpose(0, 2, 1, 3), 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), + } + + +def quantize_qkv(qkv): + """Prepare separate BSHD operands with independent BHSD row/column quantizations.""" + 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)) + ) + 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(): + 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 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")): + 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 + + +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.") + 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 = 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) + + 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) + + +_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/csrc/extensions.h b/transformer_engine/jax/csrc/extensions.h index 893374b3de3..2924c909c0d 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 deleted file mode 100644 index def5e0149b8..00000000000 --- a/transformer_engine/jax/csrc/extensions/attention.cpp +++ /dev/null @@ -1,1115 +0,0 @@ -/************************************************************************* - * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. - * - * See LICENSE for license information. - ************************************************************************/ - -#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 { - -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); - - nvte_tensor_pack_destroy(&aux_output_tensors); -} - -#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(); -} - -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); - } - - 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()); -} - -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 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() { - 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; - } - } - - 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(); -} - -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()); - } -} - -} // 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(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 - .Attrs()); - -} // namespace jax -} // namespace transformer_engine diff --git a/transformer_engine/jax/csrc/extensions/cudnn_graph.cpp b/transformer_engine/jax/csrc/extensions/cudnn_graph.cpp new file mode 100644 index 00000000000..4973708591e --- /dev/null +++ b/transformer_engine/jax/csrc/extensions/cudnn_graph.cpp @@ -0,0 +1,199 @@ +/************************************************************************* + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * + * See LICENSE for license information. + ************************************************************************/ + +// JAX-owned execution of generic graphs built with the cuDNN Python frontend. +#include +#include +#include +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include + +#include "cudnn_utils.h" + +namespace transformer_engine { +namespace jax { +namespace { + +namespace ffi = xla::ffi; + +void CheckCuda(cudaError_t status) { + if (status != cudaSuccess) throw std::runtime_error(cudaGetErrorString(status)); +} + +void CheckCudnn(cudnnStatus_t status) { + if (status != CUDNN_STATUS_SUCCESS) throw std::runtime_error(cudnnGetErrorString(status)); +} + +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); + } + return value.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; + 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()), + /*enforce_precompiled=*/true); + if (!status.is_good()) throw std::runtime_error(status.get_message()); + graph = std::move(candidate); + } + return graph; +} + +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 { + auto handle = nvte_get_cudnn_handle(); + 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 = nvte_get_cudnn_handle(); + 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."); + // 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)); + } + } + 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."); + } + 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."); + } + 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."); + 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 + +XLA_FFI_DEFINE_HANDLER_SYMBOL(CudnnGraphHandler, ExecuteGraph, + xla::ffi::Ffi::Bind() + .Ctx>() + .RemainingArgs() + .RemainingRets() + .Attrs(), + {xla::ffi::Traits::kCmdBufferCompatible}); + +XLA_FFI_DEFINE_HANDLER_SYMBOL(CudnnGraphInitializeHandler, InitializeGraph, + xla::ffi::Ffi::Bind() + .Ctx>() + .RemainingArgs() + .RemainingRets() + .Attrs()); + +} // namespace jax +} // namespace transformer_engine diff --git a/transformer_engine/jax/csrc/extensions/pybind.cpp b/transformer_engine/jax/csrc/extensions/pybind.cpp index d82d109f650..09b2df5d4b8 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"] = @@ -160,8 +151,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); @@ -172,10 +161,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); @@ -211,7 +197,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) @@ -257,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 5a01572f1fa..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, @@ -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 | str = False score_mod_requested: bool = False return_max_logit: bool = False @@ -326,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 @@ -369,6 +371,10 @@ def __call__( "score_mod_bprop_tensors": score_mod_bprop_tensors, "return_max_logit": self.return_max_logit, } + 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 @@ -622,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 @@ -632,6 +642,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 | 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. Set to ``"mxfp8"`` for + block-scaled MXFP8 attention. Optimization parameters ----------------------- @@ -661,6 +676,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 | str = False def __post_init__(self): # TODO(KshitijLakhani): Remove warning in TransformerEngine v2.12 @@ -704,6 +720,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 @@ -728,6 +745,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 ------- @@ -735,6 +756,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: @@ -761,7 +784,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 +814,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 +857,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 +870,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 + 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 + ): + 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 +981,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, )( @@ -961,6 +994,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 ( @@ -1272,6 +1306,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 | 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 @@ -1314,6 +1352,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 | str = False + def __post_init__(self): # Deal with changed defaults in API # TODO(KshitijLakhani): Remove warning in TransformerEngine v2.12 @@ -1735,6 +1775,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 +2102,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 | 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 ----------------------- @@ -2129,6 +2175,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 | str = False def __post_init__(self): if self.mha_kernel_init is None: @@ -2300,6 +2347,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 +2443,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( 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 c3503d0f1bf..78da0f51dd1 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/backends.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/backends.py @@ -18,6 +18,8 @@ import torch import torch.nn.functional as F +from transformer_engine.common.cudnn.graph import uses_ragged_attention_graph +from transformer_engine.common.cudnn.mxfp8_utils import OWN_BLOCK as MXFP8_OWN_BLOCK from transformer_engine.pytorch.device import te_device_type from transformer_engine.pytorch.utils import ( @@ -45,14 +47,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 ( @@ -1634,6 +1636,9 @@ class FusedAttnBwdArgs: rng_state: Optional[torch.Tensor] = None aux_bias: Optional[torch.Tensor] = None aux_softmax_offset: Optional[torch.Tensor] = None + # MXFP8 causal own-block attention: per-token-scaled V (E4M3) and its exponents. + v_token_values: Optional[torch.Tensor] = None + v_token_exponents: Optional[torch.Tensor] = None # --- Attention config --- max_seqlen_q: Optional[int] = None @@ -1656,6 +1661,7 @@ class FusedAttnBwdArgs: # --- FP8 --- fp8: Optional[bool] = None + causal_own_block: Optional[bool] = None is_input_fp8: Optional[bool] = None bf16_backward: Optional[bool] = None qkv_type: Optional[str] = None @@ -1689,6 +1695,8 @@ def setup_saved_tensors(self, ctx: torch.autograd.function.FunctionCtx) -> None: self.aux_softmax_offset, self.packed_qkv, self.packed_kv, + self.v_token_values, + self.v_token_exponents, ) = restore_from_func_ctx( ctx ) # pylint: disable=unbalanced-tuple-unpacking @@ -1773,6 +1781,15 @@ def _fused_attn_forward_impl( # without depending on `fp8_recipe.()`. qkv_type = _qkv_quantizer_type(QKV_quantizer) if fp8 else None + # MXFP8BlockScaling.fp8_dpa_causal_own_block: completed 32-token blocks in cuDNN MXFP8, + # each query's own block in Triton with per-token V scales, backward in high precision + # on the decoded operands (see transformer_engine/common/triton/mxfp8_causal_attention.py). + own_block = qkv_type == "mxfp8" and getattr(fp8_recipe, "fp8_dpa_causal_own_block", False) + v_token_values = v_token_exponents = None + if own_block: + assert not is_input_fp8, "MXFP8 causal own-block attention requires high-precision Q/K/V." + is_bwd_fp8 = False + # get nominal data type for out # FP16/BF16 attention: torch.float16 or torch.bfloat16 # FP8 attention: torch.float16 or torch.bfloat16 @@ -1799,7 +1816,7 @@ def _fused_attn_forward_impl( k, v, QKV_quantizer, - used_in_backward=is_training, + used_in_backward=is_training and not own_block, combined_qkv=args.packed_qkv, combined_kv=args.packed_kv, ) @@ -1853,7 +1870,28 @@ def _fused_attn_forward_impl( args.rng_gen, args.softmax_offset, cuda_graph=is_graph_capturing(), + causal_exclude_own_block=MXFP8_OWN_BLOCK if own_block else 0, ) + if own_block: + # pylint: disable=import-outside-toplevel + from transformer_engine.pytorch.triton.mxfp8_causal_attention import ( + own_block_attention, + token_values, + ) + + _, q_format, kv_format = dpa_utils.get_qkv_format(qkv_layout) + v_token_values, v_token_exponents = token_values(v, kv_format) + own_block_attention( + q_fp8, + k_fp8, + v_token_values, + v_token_exponents, + out_, + aux_ctx_tensors[0], + args.attn_scale, + q_format, + o_format, + ) # out_fp8: Float8Tensor/MXFP8Tensor; dtype = torch.float16 or torch.bfloat16 # fp8_dtype = tex.DType.kFloat8E4M3 @@ -1897,7 +1935,11 @@ def _fused_attn_forward_impl( # save q, k, v, o tensors fp8_tensors = (None, None, None, None) f16_tensors = (None, None, None, None, None, None) - if is_bwd_fp8: + if own_block: + # Backward decodes Q/K from these and V from the per-token-scaled copy. + fp8_tensors = (q_fp8, k_fp8, None, None) + f16_tensors = (None, None, None, out_f16, None, None) + elif is_bwd_fp8: if (qkv_type == "current" and _dpa_fp8_cs_o_in_f16) or qkv_type == "mxfp8": fp8_tensors = (q_fp8, k_fp8, v_fp8, None) f16_tensors = (None, None, None, out_f16, None, None) @@ -1959,7 +2001,9 @@ def _fused_attn_forward_impl( nvtx_range_pop(f"{nvtx_label}") if is_cpu_offload_enabled() and NVTE_CPU_OFFLOAD_V1: - if is_bwd_fp8: + if own_block: + tensor_list = [q_fp8, k_fp8, v_token_values, v_token_exponents, out_f16] + elif is_bwd_fp8: tensor_list = fp8_tensors else: tensor_list = [q, k, v, out_f16] @@ -1987,6 +2031,8 @@ def _fused_attn_forward_impl( None, "attn_bias" if has_bias else None, "softmax_offset" if has_softmax_offset else None, + None, + None, ) tensors_to_save = ( *fp8_tensors, @@ -1995,6 +2041,8 @@ def _fused_attn_forward_impl( rng_state, None, None, + v_token_values, + v_token_exponents, ) tensors_to_save = tuple( None if src else t for t, src in zip(tensors_to_save, saved_from, strict=True) @@ -2006,6 +2054,7 @@ def _fused_attn_forward_impl( ctx_attrs = { "fp8": is_bwd_fp8, + "causal_own_block": own_block, "is_input_fp8": is_input_fp8, "saved_from": saved_from, "fused_attention_backend": fused_attention_backend, @@ -2056,6 +2105,7 @@ def _fused_attn_setup_ctx( fp8 = ctx_attrs["fp8"] bwd_args.fp8 = fp8 + bwd_args.causal_own_block = ctx_attrs["causal_own_block"] bwd_args.is_input_fp8 = ctx_attrs["is_input_fp8"] # assume fwd and bwd always use the same high precision, i.e. torch.float16 or torch.bfloat16 # used when some tensors are base tensors and loose the "dtype" attribute @@ -2131,6 +2181,8 @@ def _fused_attn_setup_ctx( rng_state, aux_bias, aux_softmax_offset, + v_token_values, + v_token_exponents, ) = ( sources[src] if src else saved for saved, src in zip(tensors_to_save_from_forward, ctx_attrs["saved_from"], strict=True) @@ -2154,6 +2206,8 @@ def _fused_attn_setup_ctx( aux_softmax_offset, packed_qkv, packed_kv, + v_token_values, + v_token_exponents, ) @@ -2181,6 +2235,18 @@ def _fused_attn_backward_impl( else: d_out_fp8 = args.dO_quantizer(d_out) q_fp8, k_fp8, v_fp8, out_fp8 = args.q_fp8, args.k_fp8, args.v_fp8, args.out_fp8 + if args.causal_own_block: + # pylint: disable=import-outside-toplevel + from transformer_engine.pytorch.triton.mxfp8_causal_attention import decode_operands + + args.q, args.k, args.v = decode_operands( + q_fp8, + k_fp8, + args.v_token_values, + args.v_token_exponents, + args.nominal_dtype, + dpa_utils.get_qkv_format(args.qkv_layout)[1], + ) q, k, v = _fused_attn_qkv(args) out = args.out cu_seqlens_q, cu_seqlens_kv = args.cu_seqlens_q, args.cu_seqlens_kv @@ -2455,7 +2521,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) @@ -2486,7 +2552,7 @@ def _fused_attn_forward_fake( rng_state = TensorSpec(shape=(2,), dtype=torch.int64, device=q.device) has_bias = args.attn_bias_type not in ["no_bias", "alibi"] and args.attn_bias is not None has_softmax_offset = args.softmax_type != "vanilla" and args.softmax_offset is not None - tensors_to_save = (*(None,) * 10, softmax_stats, rng_state, None, None) + tensors_to_save = (*(None,) * 10, softmax_stats, rng_state, None, None, None, None) saved_from = ( *(None,) * 4, "q", @@ -2499,9 +2565,12 @@ def _fused_attn_forward_fake( None, "attn_bias" if has_bias else None, "softmax_offset" if has_softmax_offset else None, + None, + None, ) ctx_attrs = { "fp8": False, + "causal_own_block": False, "is_input_fp8": False, "saved_from": saved_from, "fused_attention_backend": args.fused_attention_backend, @@ -2613,13 +2682,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/context_parallel.py b/transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py index 42fea51cecf..59dff95ad31 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.device import te_device_type from transformer_engine.pytorch.utils import ( get_cudnn_version, @@ -15,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 @@ -29,6 +30,7 @@ from transformer_engine.pytorch.constants import ( CPLoadBalancingStrategy, dist_group_type, + FusedAttnBackend, ) from transformer_engine.pytorch.distributed import ( get_distributed_world_size, @@ -1840,7 +1842,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/dot_product_attention.py b/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py index 80844adca43..0e2c1a05330 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 @@ -2189,6 +2189,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``. @@ -2794,6 +2798,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 @@ -2808,7 +2817,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, @@ -2911,10 +2929,22 @@ def forward( # raise exception if no backend is available if sum([use_flash_attention, use_fused_attention, use_unfused_attention]) == 0: + fp8_recipe = self.fp8_meta["recipe"] if self.fp8 else None + own_block = ( + fp8_recipe is not None + and fp8_recipe.fp8_dpa + and getattr(fp8_recipe, "fp8_dpa_causal_own_block", False) + ) raise ValueError( - "No dot product attention backend is available for the provided inputs. Please" - " run with NVTE_DEBUG=1 NVTE_DEBUG_LEVEL=2 to find out the reasons for" - " disabling all backends." + "No dot product attention backend is available for the provided inputs." + + ( + " MXFP8BlockScaling(fp8_dpa_causal_own_block=True) runs only on" + " FusedAttention and has no fallback backend." + if own_block + else "" + ) + + " Please run with NVTE_DEBUG=1 NVTE_DEBUG_LEVEL=2 to find out the reasons" + " for disabling all backends." ) # run attention 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 c122b8faa43..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,28 +4,40 @@ """cuDNN-backed Flex Attention helpers.""" -from dataclasses import dataclass -import importlib -import inspect +from functools import partial from typing import Any, Callable, Dict, Optional, Tuple import torch from transformer_engine.pytorch.device import te_device_type -_cudnn_score_mod_handles: Dict[torch.device, Any] = {} -_cudnn_score_mod_graph_cache: Dict[Tuple[Any, ...], Any] = {} -_SCORE_MOD_UNCACHEABLE = object() - +from transformer_engine.common.cudnn.attention import ( + AttentionOptions, + AttentionGraphSpec, + build_attention_graph, +) +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, + 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, + _graph_cache, + _execute_cudnn_graph, + _diagnostics, +) -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 +_SCORE_MOD_UNCACHEABLE = object() def _bhsd_dim_stride( @@ -45,98 +57,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, ...]: @@ -168,12 +117,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: @@ -181,425 +124,104 @@ 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. -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 +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 _get_cudnn_current_stream_handle(cudnn, device: torch.device): - """Return a cuDNN handle for device, bound to PyTorch's current stream.""" - if device.type != te_device_type(): - raise ValueError(f"Flex Attention only supports CUDA tensors, got device {device}.") - if device.index is None: - device = torch.device(te_device_type(), 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), +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) ) - return graph - - -@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 _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() - # Only the cuDNN backend implements score_mod. cuDNN frontend 1.31 can offer its Python SDPA - # engines for these graphs, and those engines run plain attention without the modifier. - graph.deselect_engines(["sdpa_fwd", "sdpa_bwd"]) - 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) - selected_engine = getattr(graph, "selected_engine", None) - if selected_engine is not None: - raise RuntimeError( - f"cuDNN frontend {cudnn.__version__} selected {selected_engine} for a score_mod graph. " - "Only the cuDNN backend implements score_mod." - ) - 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 == te_device_type() and device.index is None: - device = torch.device(te_device_type(), 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, - 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: + if any(key is _SCORE_MOD_UNCACHEABLE for key in keys): 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), + "score_mod", + _score_mod_device_key(device), + spec, + keys, + _score_mod_tensor_dict_metadata(extra), + _score_mod_tensor_dict_metadata(bprop_extra), ) -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: - 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), - ) - - -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) - 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) - - 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 - - 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, - 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.""" - 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 = ( - _make_cudnn_graph_tensor_dict(graph, score_mod_bprop_tensors) - if score_mod_bprop is not None - else {} + 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, + spec, + score_mod=_wrap_score_mod(score_mod, extra), + score_mod_bprop=_wrap_score_mod(score_mod_bprop, bprop_extra), ) - 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, + built.tensors = _score_mod_bindings(built.tensors, extra, bprop_extra) + 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 ) - 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, - workspace_size=workspace_size, + 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 = { + 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() + } + 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, ) - - -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, + _graph_cache.compile( + entry, + lambda: compile_graph( + _import_cudnn_frontend(), + entry.built.graph, + partial(_diagnostics.run, "f16", direction), + has_score_mod=True, + ), + event=event, ) - 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): @@ -636,36 +258,25 @@ 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, ) + _diagnostics.event("f16", "fwd", query_layer.device.index, "EXECUTE") ctx.is_training = is_training ctx.q_format = q_format @@ -715,44 +326,36 @@ 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, ) + _diagnostics.event("f16", "bwd", query_layer.device.index, "EXECUTE") return ( None, diff --git a/transformer_engine/pytorch/attention/dot_product_attention/utils.py b/transformer_engine/pytorch/attention/dot_product_attention/utils.py index 2ecf8a60090..6862bd6cb53 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,17 @@ META_DO, META_S, META_DP, + TE_DType, + DType, + MXFP8_BLOCK_SCALING_SIZE, +) +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, + support_error as mxfp8_support_error, + causal_own_block_error as mxfp8_causal_own_block_error, + OWN_BLOCK as MXFP8_OWN_BLOCK, ) from transformer_engine.pytorch.attention.inference import InferenceParams from transformer_engine.pytorch.cpu_offload import is_cpu_offload_enabled @@ -50,7 +61,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 ( @@ -255,6 +265,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 @@ -348,6 +360,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): """ @@ -388,12 +401,17 @@ class FusedAttentionParams: scaling_mode: tex.NVTEScalingMode = tex.NVTEScalingMode.NVTE_DELAYED_TENSOR_SCALING dropout: float = 0.0 attn_scale: float = 1.0 + # MXFP8 causal own-block attention (MXFP8BlockScaling.fp8_dpa_causal_own_block): the forward + # excludes each query's own block of this many keys and the backward runs in model precision. + causal_exclude_own_block: int = 0 # tensor types qkv_dtype: DType = DType.kBFloat16 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 @@ -428,6 +446,10 @@ 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 + bias_stride: Optional[Tuple[int, int, int, int]] = None + class _NoOpLogger: """ @@ -454,7 +476,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 @@ -463,7 +485,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 @@ -558,6 +582,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__, @@ -729,10 +754,14 @@ def _disable_all_flash_attention() -> None: use_fused_attention = False # Filter: Execution type + causal_own_block = False if fp8 and fp8_meta["recipe"].fp8_dpa: fp8_recipe = fp8_meta["recipe"] if fp8_meta.get("local_recipes", None) is not None: fp8_recipe = fp8_meta["local_recipes"][0] + causal_own_block = fp8_recipe.mxfp8() and getattr( + fp8_recipe, "fp8_dpa_causal_own_block", False + ) if use_flash_attention_2 and FlashAttentionUtils.is_installed: logger.debug("Disabling FlashAttention 2 for FP8 attention") use_flash_attention_2 = False @@ -757,6 +786,13 @@ def _disable_all_flash_attention() -> None: if not allow_emulation: logger.debug("Disabling UnfusedDotProductAttention for FP8 attention") use_unfused_attention = False + elif causal_own_block: + # Emulation would quantize V with shared columnwise scales and reintroduce the + # future-token dependence the option removes. + logger.debug( + "Disabling UnfusedDotProductAttention for MXFP8 causal own-block attention" + ) + use_unfused_attention = False if use_fused_attention and fp8_recipe.delayed(): if ( device_compute_capability >= (10, 0) @@ -786,37 +822,50 @@ 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 not reason and causal_own_block: + reason = mxfp8_causal_own_block_error( + 10 * device_compute_capability[0] + device_compute_capability[1], + qkv_layout, + attn_mask_type, + window_size, + head_dim_qk, + head_dim_v, + max_seqlen_q, + max_seqlen_kv, + attention_dropout, + core_attention_bias_type, + softmax_type, + ) + if not reason and issubclass(qkv_type, QuantizedTensorStorage): + reason = "MXFP8 causal own-block attention requires high-precision Q, K, V." + if not reason and context_parallel: + reason = ( + "MXFP8 causal own-block attention does not support context parallelism." + ) + 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 - 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 @@ -955,14 +1004,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") @@ -1024,15 +1073,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, ) @@ -1198,19 +1243,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 @@ -1647,8 +1692,12 @@ 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, + "requires_backward": attention_params.requires_backward, "deterministic": deterministic, "cuda_graph": cuda_graph, "return_max_logit": return_max_logit, @@ -1665,6 +1714,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], @@ -1690,6 +1740,13 @@ 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 + ), + "causal_exclude_own_block": MXFP8_OWN_BLOCK if causal_own_block else 0, } # Context-parallel per-step configs if context_parallel: @@ -1821,7 +1878,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( @@ -1832,7 +1889,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") @@ -1844,15 +1901,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 ): @@ -3467,25 +3524,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 @@ -3497,8 +3550,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 @@ -3551,10 +3604,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 @@ -3562,20 +3613,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 ) @@ -3598,20 +3642,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]) ] 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..af619ddf8f9 100644 --- a/transformer_engine/pytorch/constants.py +++ b/transformer_engine/pytorch/constants.py @@ -192,3 +192,141 @@ 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``. 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 + ``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.py b/transformer_engine/pytorch/cpp_extensions/cudnn.py new file mode 100644 index 00000000000..a843e02ac80 --- /dev/null +++ b/transformer_engine/pytorch/cpp_extensions/cudnn.py @@ -0,0 +1,74 @@ +# 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 threading +from contextlib import nullcontext + +import torch + +from transformer_engine.common.cudnn.graph import ( + import_cudnn as _import_cudnn_frontend, + create_graph, + 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): + """Get a thread-local, device-specific handle on PyTorch's current stream.""" + if device.type != te_device_type(): + raise ValueError(f"cuDNN attention requires {te_device_type()} tensors, got {device}.") + if device.index is None: + device = torch.device(te_device_type(), torch.cuda.current_device()) + if not hasattr(_handles, "by_device"): + _handles.by_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() + _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 create_graph(cudnn, dtype, _get_cudnn_current_stream_handle(cudnn, device)) + + +def _finalize_cudnn_graph(graph, *, has_score_mod=False): + """Check support, build an execution plan, and return workspace bytes.""" + cudnn = _import_cudnn_frontend() + check_graph_support(cudnn, graph, has_score_mod=has_score_mod) + return compile_graph(cudnn, graph, has_score_mod=has_score_mod) + + +def _execute_cudnn_graph(graph, variant_pack, workspace_size, device): + """Execute on the current stream using framework-owned workspace.""" + cudnn = _import_cudnn_frontend() + if device.type == te_device_type() and device.index is None: + device = torch.device(te_device_type(), 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), + ) 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..c1cffe10a7d --- /dev/null +++ b/transformer_engine/pytorch/cpp_extensions/cudnn_fused_attn.py @@ -0,0 +1,1168 @@ +# 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 supplies capture-safe Philox state and sequence-length conversion. +Graph construction, plan caching, and execution use the Python frontend. +""" + +from dataclasses import replace +from functools import partial +import math +import os +from typing import List, Optional, Tuple, Union + +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, + GraphCacheEntry, + attention_matrix_spec, + attention_stats_spec, + mxfp8_scale_spec, + check_graph_support, + compile_graph, + bucket, + uses_ragged_attention_graph, +) +from transformer_engine.common.cudnn.mxfp8_utils import support_error as mxfp8_support_error + +from .cudnn import ( + _import_cudnn_frontend, + _build_cudnn_pygraph, + _execute_cudnn_graph, + _graph_cache, + _diagnostics, +) +from ..tensor.float8_tensor import Float8Quantizer +from ..tensor.storage.mxfp8_tensor_storage import MXFP8TensorStorage +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 + + +def _event(attn, event, key=None): + _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): + return _diagnostics.run( + "fp8" if attn.fp8 else "f16", + "bwd" if attn.backward else "fwd", + name, + function, + ) + + +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 _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 + + +class _AttentionGraph: + """Runtime tensors and data-free descriptors for one attention invocation.""" + + 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 if device is None else device + self.q_format, self.kv_format = _formats(params["qkv_layout"]) + 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 + 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.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 + ): + """Bind storage using the same matrix descriptors as support queries.""" + data = _raw(tensor, columnwise) + dim = (self.graph_b, heads, seqlen, tensor.shape[-1]) + offset_name = name + "_offset" if fmt == "thd" else "" + if offset_name: + 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) + self.tensors[name] = data + self.specs[name] = attention_matrix_spec( + dim, + fmt, + _dtype(tensor), + 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.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.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 + + def scale(self, name, tensor, matrix_name, fmt, columnwise=False): + """Bind tensor scales or reordered MXFP8 block scales.""" + if not self.mxfp8: + return self.scalar(name, tensor) + 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.""" + 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 + bounded_backward = self.backward and not self.fp8 and self.ragged_graph + # The SM100 MXFP8 engine that serves causal_exclude_own_block takes a constant scale. + own_block = 0 if self.backward else p.get("causal_exclude_own_block", 0) + return AttentionOptions( + backward=self.backward, + attn_scale=float(p["attn_scale"]) if own_block else "attn_scale", + 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, + use_cu_seqlens=self.use_cu_seqlens, + 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, + # 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, + causal_exclude_own_block=own_block, + ) + + 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.""" + # Autograd may run backward on another thread. Share device-specific plans; + # each execution supplies its thread-local handle on the current stream. + 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.""" + 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 GraphCacheEntry(built) + + def execute(self): + """Compile once, then bind this invocation's runtime tensor addresses.""" + entry = self.entry() + _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.built.graph, + {entry.built.tensors[n]: t for n, t in self.tensors.items()}, + entry.workspace_size, + self.device, + ) + _event(self, "EXECUTE") + + +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) + if not p.get("causal_exclude_own_block"): + # The own-block graph takes the scale as a graph constant (AttentionOptions.attn_scale). + attn.scalar( + "attn_scale", torch.tensor([p["attn_scale"]], dtype=torch.float32), by_value=True + ) + if "padding" in p["attn_mask_type"]: + 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": + attn.add("sink", 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 _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 ( + config.causal_exclude_own_block or not int(os.getenv("NVTE_FP8_DPA_BWD", "1")) + ): + # FusedAttnFunc saves model-precision Q/K/V and O for this backward (MXFP8 causal + # own-block attention decodes them from its MXFP8 operands). + # 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 + ), + ) + return config + + +def _query_graph(config, backward): + """Describe a support query without allocating tensors or touching RNG state.""" + if backward: + config = _backward_config(config) + + 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 + ) + types = { + DType.kFloat16: "HALF", + DType.kBFloat16: "BFLOAT16", + DType.kFloat32: "FLOAT", + DType.kFloat8E4M3: "FP8_E4M3", + DType.kFloat8E5M2: "FP8_E5M2", + } + 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), + "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_scale": config.attn_scale, + "causal_exclude_own_block": 0 if backward else config.causal_exclude_own_block, + } + 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) + + def matrix(n, fmt, h, s, d, dtype, packed="", output=False): + 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.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.graph_b + 1, 1, 1, 1), + dtype="INT32" if attn.use_cu_seqlens else "INT64", + ) + attn.specs[n] = attention_stats_spec( + 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): + if not attn.mxfp8: + spec(n) + 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 ( + ("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) + if not attn.params["causal_exclude_own_block"]: + # The own-block graph takes the scale as a graph constant (AttentionOptions.attn_scale). + spec("attn_scale", by_value=True) + if "padding" in attn.params["attn_mask_type"]: + for n in ("seq_q", "seq_kv"): + 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") + 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, 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": + spec("sink", (1, attn.h, 1, 1), (attn.h, 1, 1, 1)) + if backward: + 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) + 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, version, capability, backward=False): + """Check TE's hardware, recipe, and known cuDNN defect exclusions.""" + 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 (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 layout.startswith("paged_kv_"): + 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 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 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 ( + backward + and ragged_stats + and softmax == "learnable" + and version < 92600 + and config.head_dim_v not in (64, 128, 256) + ): + 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 q_format != kv_format: + return "FP8 fused attention supports BSHD/SBHD/BHSD/THD formats." + 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 q_format == "thd": + if version < 92300: + return "FP8 fused attention with THD format requires cuDNN 9.23.0 or later!" + if backward and capability < (10, 0): + return "FP8 fused attention with THD format supports backward on sm100+ only!" + 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(config.head_dim_qk, config.head_dim_v) > 128: + return "FP8 THD attention supports head dimensions up to 128 on sm100+ only!" + 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 mxfp8 else half_types + fp8_types + tensor_scaling = config.scaling_mode == tex.NVTEScalingMode.NVTE_DELAYED_TENSOR_SCALING + if ( + not (tensor_scaling or mxfp8) + or DType.cast(config.o_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!" + ) + return "" + + +def get_fused_attn_backend(config): + """Select an attention backend by validating the Python frontend graphs.""" + 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}." + requires_backward = config.requires_backward + if requires_backward is None: + requires_backward = config.is_training + try: + 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) + + +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, + causal_exclude_own_block: int = 0, +) -> Tuple[Union[torch.Tensor, None], ...]: + """Execute forward, returning output, backward auxiliaries, and optional max logits. + + ``causal_exclude_own_block`` (MXFP8 only): if nonzero, the causal mask also excludes each + query's own aligned block of this many keys, so query ``i`` reads keys + ``j < B * floor(i / B)`` and rows ``0 .. B-1`` return zero output and ``-inf`` statistics. + With ``B = 32`` no columnwise V scale the kernel reads covers a token the query cannot see. + """ + 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, + "causal_exclude_own_block": causal_exclude_own_block, + "rng_elts_per_thread": rng_elts_per_thread, + } + ) + 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) + _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.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"]] + max_s_tensor = None + if p["return_max_logit"]: + 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": + 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() + if not attn.ragged_stats and attn.graph_b > attn.b: + aux_ctx_tensors[0] = stats[: 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 max_s_tensor.ndim == 3 else (0, 2, 3) + + if qkv_format == "thd": + 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) + 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=max_s_tensor.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=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 + max_s_tensor = max_s_tensor.masked_fill(~valid.view(-1, 1, 1), float("-inf")) + + # Max -> max_logit [h] + 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 + + +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}." + ) + + 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 + 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: + # 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) + if p["softmax_type"] != "vanilla": + dsoftmax = torch.empty_like(p["softmax_offset"]) + 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) + 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 deleted file mode 100644 index 9a33df7634b..00000000000 --- a/transformer_engine/pytorch/cpp_extensions/fused_attn.py +++ /dev/null @@ -1,653 +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 -import transformer_engine_torch as tex -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; - 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"} - 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}") - - # 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, - ) - - 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 (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": - 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=output_tensors[0].dtype) - return output_tensors[0], aux_ctx_tensors, max_logit - - # out, aux_ctx_tensors - return output_tensors[0], output_tensors[1:] - - -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}." - ) - - 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, - ) - - return output_tensors diff --git a/transformer_engine/pytorch/csrc/extensions.h b/transformer_engine/pytorch/csrc/extensions.h index 02bf6773442..81db70b454e 100644 --- a/transformer_engine/pytorch/csrc/extensions.h +++ b/transformer_engine/pytorch/csrc/extensions.h @@ -95,36 +95,11 @@ 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); + +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 98d55d2b8b5..c52a303af0f 100644 --- a/transformer_engine/pytorch/csrc/extensions/attention.cpp +++ b/transformer_engine/pytorch/csrc/extensions/attention.cpp @@ -10,640 +10,41 @@ 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!"); - } +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."); } - 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. + 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}; +} - 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 df6d16782d2..822862d10bb 100644 --- a/transformer_engine/pytorch/csrc/extensions/pybind.cpp +++ b/transformer_engine/pytorch/csrc/extensions/pybind.cpp @@ -494,8 +494,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()); @@ -588,10 +586,12 @@ 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("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, diff --git a/transformer_engine/pytorch/triton/mxfp8_causal_attention.py b/transformer_engine/pytorch/triton/mxfp8_causal_attention.py new file mode 100644 index 00000000000..655b716762a --- /dev/null +++ b/transformer_engine/pytorch/triton/mxfp8_causal_attention.py @@ -0,0 +1,128 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""PyTorch wrappers for causal MXFP8 attention without future-token scale leakage.""" + +import torch + +from transformer_engine.common.cudnn.mxfp8_utils import OWN_BLOCK +from transformer_engine.common.triton.mxfp8_causal_attention import ( + decode_kernel, + own_block_attention_kernel, + token_values_kernel, +) + +__all__ = ["OWN_BLOCK", "token_values", "own_block_attention", "decode_operands"] + + +def _bsh_strides(x: torch.Tensor, fmt: str): + """(batch, sequence, head) element strides of a bshd/sbhd tensor with unit feature stride.""" + assert x.stride(-1) == 1, "the feature dimension must be contiguous" + if fmt == "bshd": + return x.stride(0), x.stride(1), x.stride(2) + if fmt == "sbhd": + return x.stride(1), x.stride(0), x.stride(2) + if fmt == "bhsd": + return x.stride(0), x.stride(2), x.stride(1) + raise ValueError(f"Unsupported format {fmt!r}; expected bshd, sbhd or bhsd.") + + +def _bshd_dims(x: torch.Tensor, fmt: str): + """(batch, sequence, heads, features) of a bshd/sbhd/bhsd tensor.""" + order = {"bshd": (0, 1, 2, 3), "sbhd": (1, 0, 2, 3), "bhsd": (0, 2, 1, 3)}[fmt] + return tuple(x.shape[i] for i in order) + + +def _payload(x) -> torch.Tensor: + """Rowwise E4M3 payload of an MXFP8 tensor (stored as uint8) typed for Triton.""" + return x._rowwise_data.view(torch.float8_e4m3fn) + + +def _pack_heads(num_heads: int, num_kv_heads: int) -> int: + """Query heads sharing a KV head that fill the boundary kernel's 128 MMA rows.""" + return next(p for p in (4, 2, 1) if (num_heads // num_kv_heads) % p == 0) + + +def token_values(v: torch.Tensor, fmt: str): + """Token-normalized V for the own-block product and the backward. + + Returns ``u`` (E4M3, ``[B, H_kv, S, 128]``) and ``e`` (int32, ``[B, H_kv, S]``) with + ``V = u * 2**(e - 8)`` up to the E4M3 rounding of ``u``. + """ + b, s, h, d = _bshd_dims(v, fmt) + assert d == 128 and s % OWN_BLOCK == 0 + u = torch.empty((b, h, s, d), dtype=torch.float8_e4m3fn, device=v.device) + e = torch.empty((b, h, s), dtype=torch.int32, device=v.device) + token_values_kernel[(s // OWN_BLOCK, h, b)](v, u, e, *_bsh_strides(v, fmt), s, h, num_warps=4) + return u, e + + +def own_block_attention( + q_fp8, k_fp8, u, e, out, stats, attn_scale: float, qkv_format: str, o_format: str +): + """Merge each query's own-block attention into the completed-block result, in place. + + ``q_fp8`` and ``k_fp8`` are rowwise MXFP8 tensors whose payloads are in ``qkv_format`` and + whose scales are in the per-(batch, head) F8_128x4 layout; ``out`` (``o_format``) and + ``stats`` (``[B, H, S, 1]``, natural log) hold cuDNN's ``causal_exclude_own_block=32`` + result on entry and the causal result on exit. + """ + fmt = qkv_format + q8, k8 = _payload(q_fp8), _payload(k_fp8) + b, s, h, d = _bshd_dims(q8, fmt) + h_kv = _bshd_dims(k8, fmt)[2] + assert d == 128 and s % 128 == 0 and h % h_kv == 0 + assert stats.is_contiguous() and stats.numel() == b * h * s + pack = _pack_heads(h, h_kv) + own_block_attention_kernel[(s // (128 // pack), h // pack, b)]( + q8, + k8, + q_fp8._rowwise_scale_inv, + k_fp8._rowwise_scale_inv, + u, + e, + out, + stats, + *_bsh_strides(q8, fmt), + *_bsh_strides(k8, fmt), + *_bsh_strides(out, o_format), + attn_scale, + s, + h, + h_kv, + pack, + num_warps=4, + num_stages=1, + ) + + +def decode_operands(q_fp8, k_fp8, u, e, dtype: torch.dtype, fmt: str): + """BF16/FP16 Q, K and V for the backward, shaped like the MXFP8 payloads (``fmt``).""" + q8, k8 = _payload(q_fp8), _payload(k_fp8) + b, s, h, _ = _bshd_dims(q8, fmt) + h_kv = _bshd_dims(k8, fmt)[2] + q = torch.empty_strided(q8.shape, q8.stride(), dtype=dtype, device=q8.device) + k = torch.empty_strided(k8.shape, k8.stride(), dtype=dtype, device=k8.device) + v = torch.empty_like(k) + block = 64 + decode_kernel[(s // block, b * (h + h_kv))]( + q8, + q_fp8._rowwise_scale_inv, + k8, + k_fp8._rowwise_scale_inv, + u, + e, + q, + k, + v, + *_bsh_strides(q, fmt), + *_bsh_strides(k, fmt), + b, + s, + h, + h_kv, + block, + num_warps=4, + ) + return q, k, v