diff --git a/AGENTS.md b/AGENTS.md index 7dceca65f63..0f80257a3bd 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -280,6 +280,15 @@ particular container, GPU architecture, GPU count, or external repository. Use the higher-level launcher for distributed, Megatron Core, FSDP, ONNX, and attention-backend coverage rather than reconstructing its environment. +Adding a test file does not automatically add QA coverage: for example, +`qa/L0_pytorch_unittest/test.sh` explicitly enumerates attention test files. +Connect new specialized suites to a checked-in qualified target. The SM100 +FROST target is `qa/L1_pytorch_frost_unittest/test.sh`; a provisioned L0 lane +can enable its hook with `NVTE_FROST_TEST_REQUIRED=1`. A required target must +fail for missing capabilities or skipped tests and verify that isolated +negative-test children actually execute, rather than accepting an all-skipped +pytest result. + ### Validation by changed area - Shared Python: focused behavior through each affected frontend. diff --git a/docs/envvars.rst b/docs/envvars.rst index 09e2b462780..84248f32660 100644 --- a/docs/envvars.rst +++ b/docs/envvars.rst @@ -216,6 +216,114 @@ backend-selection overview. :Default: ``1`` :Description: Enable or disable FlashAttention 4 (the ``flash-attn-4`` package) for DotProductAttention, without affecting FlashAttention 2 or 3. When set to ``0``, FlashAttention 4 will not be used even if it is installed. + +.. envvar:: NVTE_FUSED_ATTN_FROST + + :Default: ``0`` + :Values: ``0``, ``1``, ``2`` + :Description: Experimental cuDNN OSS (FROST) backward provider for fused + attention on SM100. ``1`` attempts the OSS provider and retains + the cuDNN backend for unsupported declarations. ``2`` requires + the OSS provider and reports an error instead of falling back; + use it to verify benchmark routing. Forward remains cuDNN. + Requires a TE build with ``NVTE_WITH_CUTEDSL=1`` + (CMake: ``-DNVTE_WITH_CUTEDSL=ON``), + ``NVTE_ENABLE_CUTEDSL_BACKEND=1``, + ``CUDNN_FRONTEND_ENABLE_FROST_ENGINES=1``, CuTeDSL 4.8 or newer + with its CUDA 13 compiler flavor (``nvidia-cutlass-dsl[cu13]``), + and a compatible ``nvidia-cudnn-frontend`` Python installation. + + The initial contract covers BF16, separate Q/K/V, D128/D256, nondeterministic + backward, no dropout, no bias/sinks, and unmasked or top-left causal attention. + D256 also supports bottom-right causal attention when ``max_seqlen_q <= + max_seqlen_kv``. Q-longer-than-KV uniform THD declarations use the backend + in mode ``1`` and report unsupported OSS backward in mode ``2``. Dense BSHD + Q-longer-than-KV calls are rejected by the existing forward path. + D128 bottom-right declarations retain the backend fallback (or fail in strict mode). + BSHD lengths must be multiples of 128. D128 also supports native packed THD + with device lengths and physical origins. D256 THD requires the separate + uniform-layout promise below. Other layouts, masks and dtypes retain their + existing route in mode ``1``. Kernel support does not establish performance + for every context-parallel communication type. + + Python builds the cuDNN graph and registers its compiled native function on + the first call for a declaration and device. Subsequent C++ calls pass the + current pointers, stream and framework-owned workspace directly through + TVM-FFI. The complete backward chain retains all required scratch + initialization and gradient reductions. Warm up before CUDA graph capture. + + This draft depends on the prepared-plan contract in cuDNN Frontend + `PR 1688 `_ + (development dependency revision ``333278627eb2ccffc5f47da26d733f87a2d0a8aa``). + That contract is not yet a stable public C++ export API. The adapter checks + its operand order and metadata before publishing a native entry; an ABI + mismatch reports an error. Install the compatible FE source in the same + Python environment as TE. An older wheel is not equivalent to this dependency. + + A dedicated SM100 regression target is available as + ``TE_PATH=/path/to/TransformerEngine XML_LOG_DIR=/logs bash + qa/L1_pytorch_frost_unittest/test.sh``. Provision the compatible FE source, + CUDA 13 CuTeDSL dependencies, pytest, and a TE build with the bridge enabled + before invoking it. The target fails for missing capabilities, skipped tests, + or missing isolated metadata-guard cases. Qualified L0 lanes can also set + ``NVTE_FROST_TEST_REQUIRED=1`` to invoke it from the PyTorch L0 launcher; + ordinary L0 lanes do not acquire these experimental dependencies. + +.. envvar:: NVTE_FUSED_ATTN_FROST_ALLGATHER + + :Default: ``0`` + :Values: ``0``, ``1`` + :Description: Experimental D256 BF16 THD backward adaptation for causal + AllGather with equal-length sequences and CP sizes 2, 4 or 8. + Requires ``NVTE_FUSED_ATTN_FROST_UNIFORM=1`` and a cuDNN OSS + D256 engine with `FE PR 1690 + `_. The graph + retains full Q/K/V buffers and their batch pitch; Q/O/dO/dQ + pointers select the current half-sequence, and packed Stats + are read at the corresponding physical origin. Sequence + lengths and origins are validated on the GPU on every call. + This adds no Q/K/V layout-copy kernel. Existing THD gradient + initialization remains necessary for unused physical rows. + Other gapped or unequal THD layouts remain unsupported. + +.. envvar:: NVTE_FUSED_ATTN_FROST_UNIFORM + + :Default: ``0`` + :Values: ``0``, ``1`` + :Description: For D256 THD only, promise that every actual Q/KV sequence fills + its declared maximum length and that physical token origins + are the corresponding compact cumulative offsets (or the + supported per-step pitches and origins when + ``NVTE_FUSED_ATTN_FROST_ALLGATHER=1``). This allows + an equivalent dense OSS backward declaration. TE checks the + promise on the GPU on every invocation and converts the packed + FP32 Stats into the dense layout on the same stream. A violated + promise raises a device error; it is not an unchecked hint. + Unequal THD and other gapped layouts must leave this option disabled. + + D128 uses native THD and does not need this promise. The option does not alter + attention precision, masks, communication or the public fused-attention API. + +.. envvar:: NVTE_FUSED_ATTN_FROST_D128_NOZERO + + :Default: ``0`` + :Values: ``0``, ``1`` + :Description: Omit the three PyTorch dQ/dK/dV output fills when the D128 + FROST backward chain overwrites every physical output row. + Requires strict routing (``NVTE_FUSED_ATTN_FROST=2``), BF16, + separate Q/K/V, grouped-query attention, and square compact + equal-length THD with lengths divisible by 128. Native and + fallback modes, rectangular shapes, and unequal layouts retain + the existing initialization. Internal FROST scratch + initialization and grouped-gradient reductions are unchanged. + + Enabling this option promises that device sequence lengths and physical + origins match the compact declaration. TE checks that promise on the GPU + on every execution, including CUDA graph replay. Invalid metadata raises a + device error before output may be used. D128 retains its native THD graph; + this option does not require ``NVTE_FUSED_ATTN_FROST_UNIFORM`` or change + attention precision, masks, communication, or output allocation sizes. + .. envvar:: NVTE_FUSED_ATTN :Type: ``int`` (0 or 1) diff --git a/qa/L0_pytorch_unittest/test.sh b/qa/L0_pytorch_unittest/test.sh index f45e9369d10..ea1773ea839 100644 --- a/qa/L0_pytorch_unittest/test.sh +++ b/qa/L0_pytorch_unittest/test.sh @@ -63,6 +63,9 @@ NVTE_FLASH_ATTN=0 NVTE_CPU_OFFLOAD_V1=1 python3 -m pytest --tb=auto --junitxml=$ python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_hybrid_quantization.xml $TE_PATH/tests/pytorch/test_hybrid_quantization.py || test_fail "test_hybrid_quantization.py" python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_identity_quantizer.xml $TE_PATH/tests/pytorch/test_identity_quantizer.py || test_fail "test_identity_quantizer.py" NVTE_ALLOW_UNSAFE_PICKLE_EXTRA_STATE=1 python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_attention.xml $TE_PATH/tests/pytorch/attention/test_attention.py || test_fail "test_attention.py" +if [ "${NVTE_FROST_TEST_REQUIRED:-0}" = "1" ]; then + bash "$TE_PATH/qa/L1_pytorch_frost_unittest/test.sh" || test_fail "test_frost_attn.py" +fi python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_attention_backend_selection.xml $TE_PATH/tests/pytorch/attention/test_attention_backend_selection.py || test_fail "test_attention_backend_selection.py" python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_flex_attention.xml $TE_PATH/tests/pytorch/attention/test_flex_attention.py || test_fail "test_flex_attention.py" NVTE_GDN_TEST_REQUIRED=1 python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_gdn_attention.xml $TE_PATH/tests/pytorch/attention/test_gdn_attention.py || test_fail "test_gdn_attention.py" diff --git a/qa/L1_pytorch_frost_unittest/test.sh b/qa/L1_pytorch_frost_unittest/test.sh new file mode 100644 index 00000000000..2aa9b585df7 --- /dev/null +++ b/qa/L1_pytorch_frost_unittest/test.sh @@ -0,0 +1,50 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +# Qualified SM100 target: provision compatible FE/DSL and build TE with CuTeDSL. +set -euo pipefail +ulimit -c 0 +: "${TE_PATH:=/opt/transformerengine}" +: "${XML_LOG_DIR:=/logs}" +mkdir -p "$XML_LOG_DIR" +export NVTE_ENABLE_CUTEDSL_BACKEND=1 +export CUDNN_FRONTEND_ENABLE_FROST_ENGINES=1 +export NVTE_FUSED_ATTN_FROST=2 +export NVTE_ALLOW_NONDETERMINISTIC_ALGO=1 +export NVTE_FLASH_ATTN=0 +export NVTE_FUSED_ATTN=1 + +python3 - <<'PREFLIGHT' +from importlib.metadata import version +from packaging.version import Version +import torch +import cudnn +import tvm_ffi +from cutlass.base_dsl.version_info import CUDA_VERSION +import transformer_engine +import transformer_engine.pytorch +assert torch.cuda.is_available() and torch.cuda.get_device_capability() == (10, 0), "FROST QA requires SM100" +assert Version(version("nvidia-cutlass-dsl")) >= Version("4.8.0"), "FROST QA requires CuTeDSL >= 4.8" +assert CUDA_VERSION.major >= 13, "FROST QA requires the CUDA 13 CuTeDSL compiler flavor" +print("FROST QA imports:", transformer_engine.__file__, cudnn.__file__, tvm_ffi.__file__) +PREFLIGHT + +python3 -m pytest -v --tb=auto \ + --junitxml="$XML_LOG_DIR/pytest_test_frost_attn.xml" \ + "$TE_PATH/tests/pytorch/attention/test_frost_attn.py" + +# Reject a green all-skipped run or selection that omitted native-route guards. +python3 - "$XML_LOG_DIR/pytest_test_frost_attn.xml" <<'RESULTS' +import sys +import xml.etree.ElementTree as ET +root = ET.parse(sys.argv[1]).getroot() +cases = list(root.iter("testcase")) +assert len(cases) >= 62, "FROST QA did not execute the complete regression suite" +assert not any(child.tag in ("failure", "error", "skipped") for case in cases for child in case), "FROST QA requires zero failures, errors or skips" +names = [case.get("name", "") for case in cases] +for kind in ("q_length", "kv_length", "q_origin", "kv_origin", "graph_terminal"): + assert f"test_frost_d128_nozero_metadata_guard[{kind}]" in names, "Missing isolated metadata-guard test: " + kind +assert "test_frost_d256_bottom_right_q_longer_fallback" in names +print("FROST QA passed:", len(cases), "tests, zero skips, all five isolated guards") +RESULTS diff --git a/tests/pytorch/attention/test_frost_attn.py b/tests/pytorch/attention/test_frost_attn.py new file mode 100644 index 00000000000..77787486866 --- /dev/null +++ b/tests/pytorch/attention/test_frost_attn.py @@ -0,0 +1,499 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""cuDNN OSS through the real TE common-library backward dispatcher.""" + +import ctypes +import json +import os +from pathlib import Path +import subprocess +import sys + +import pytest +import torch + + +@pytest.fixture +def frost_enabled(monkeypatch): + """Require native execution instead of accepting a silent backend fallback.""" + if not torch.cuda.is_available() or torch.cuda.get_device_capability() != (10, 0): + pytest.skip("cuDNN OSS integration currently requires SM100") + monkeypatch.setenv("CUDNN_FRONTEND_ENABLE_FROST_ENGINES", "1") + pytest.importorskip("cudnn") + pytest.importorskip("tvm_ffi") + from transformer_engine.common import ( + _get_shared_object_file, + _load_tvm_ffi_library, + _register_cutedsl_backends, + ) + + assert _load_tvm_ffi_library() + assert _register_cutedsl_backends() + + core = ctypes.CDLL(str(_get_shared_object_file("core"))) + setter = core.nvte_set_cutedsl_backend + setter.argtypes, setter.restype = [ctypes.c_int], None + setter(1) + monkeypatch.setenv("NVTE_FUSED_ATTN_FROST", "2") + monkeypatch.setenv("NVTE_FUSED_ATTN_FROST_UNIFORM", "1") + monkeypatch.setenv("NVTE_FUSED_ATTN_FROST_D128_NOZERO", "0") + yield + setter(int(os.environ.get("NVTE_ENABLE_CUTEDSL_BACKEND", "0") != "0")) + + +def _case(dim, thd, unequal, causal, seed=31, bottom_right=False, kv_length=None, q_length=None): + generator = torch.Generator(device="cpu").manual_seed(seed) + q_lengths, k_lengths = ((129, 63), (193, 97)) if unequal else ((128, 128), (256, 256)) + if q_length is not None: + assert not unequal + q_lengths = (q_length, q_length) + if kv_length is not None: + assert not unequal + k_lengths = (kv_length, kv_length) + heads, kv_heads = 4, 1 + inputs, expected = [], [] + for sq, sk in zip(q_lengths, k_lengths): + + def draw(s, h): + return (torch.randn(s, h, dim, generator=generator) * 0.5).bfloat16() + + q, k, v, do = draw(sq, heads), draw(sk, kv_heads), draw(sk, kv_heads), draw(sq, heads) + qr, kr, vr = (x.double().transpose(0, 1).requires_grad_() for x in (q, k, v)) + scores = (qr @ kr.repeat_interleave(heads // kv_heads, 0).transpose(-1, -2)) * dim**-0.5 + if causal: + scores = scores.masked_fill( + torch.arange(sk)[None, :] + > torch.arange(sq)[:, None] + (sk - sq if bottom_right else 0), + -torch.inf, + ) + # Bottom-right Q > KV has keyless leading rows: their output/gradient is zero. + keyless = torch.isneginf(scores).all(dim=-1, keepdim=True) + probabilities = scores.masked_fill(keyless, 0).softmax(-1).masked_fill(keyless, 0) + result = probabilities @ vr.repeat_interleave(heads // kv_heads, 0) + grads = torch.autograd.grad(result, (qr, kr, vr), do.transpose(0, 1).double()) + inputs.append((q, k, v, do)) + expected.append(tuple(x.transpose(0, 1).detach() for x in grads)) + combine = torch.cat if thd else torch.stack + q, k, v, do = (combine([row[i] for row in inputs]).cuda() for i in range(4)) + reference = tuple(combine([row[i] for row in expected]) for i in range(3)) + prefixes = [ + torch.tensor([0, lengths[0], sum(lengths)], device="cuda", dtype=torch.int32) + for lengths in (q_lengths, k_lengths) + ] + return q, k, v, do, prefixes, reference, max(q_lengths), max(k_lengths) + + +def _calls(case, dim, thd, causal, bottom_right=False, origins=None): + from transformer_engine.pytorch.cpp_extensions.fused_attn import ( + FusedAttnBackend, + fused_attn_fwd, + fused_attn_bwd, + ) + + q, k, v, do, (cu_q, cu_kv), reference, sq, sk = case + fmt = "thd" if thd else "bshd" + mask = ( + ("padding_causal" if causal else "padding") if thd else ("causal" if causal else "no_mask") + ) + if bottom_right: + assert causal + mask += "_bottom_right" + options = dict( + max_seqlen_q=sq, + max_seqlen_kv=sk, + cu_seqlens_q=cu_q, + cu_seqlens_kv=cu_kv, + q=q, + k=k, + v=v, + fake_dtype=torch.bfloat16, + fused_attention_backend=FusedAttnBackend.F16_arbitrary_seqlen, + cu_seqlens_q_padded=cu_q if origins is None else origins[0], + cu_seqlens_kv_padded=cu_kv if origins is None else origins[1], + attn_scale=dim**-0.5, + qkv_layout="_".join([fmt] * 3), + o_format=fmt, + attn_mask_type=mask, + dropout=0.0, + ) + o, aux = fused_attn_fwd(is_training=True, **options) + + def backward(): + return fused_attn_bwd( + o=o, + d_o=do, + aux_ctx_tensors=aux, + do_format=fmt, + dqkv_layout=options["qkv_layout"], + deterministic=False, + **options + )[:3] + + return backward, reference + + +def _check(actual, reference): + for got, want in zip(actual, reference): + torch.testing.assert_close(got.double().cpu(), want, atol=0.025, rtol=0.025) + error = (got.double().cpu() - want).square().mean().sqrt() + assert error / want.square().mean().sqrt().clamp_min(1e-8) < 0.02 + + +@pytest.mark.parametrize( + "dim,thd,unequal", + [(128, False, False), (128, True, True), (256, False, False), (256, True, False)], +) +@pytest.mark.parametrize("causal", [False, True]) +def test_frost_backward_and_rebinding(frost_enabled, dim, thd, unequal, causal): + """New pointers and changed inputs reuse the cached native plan.""" + del frost_enabled + for seed in (31, 47): + backward, reference = _calls(_case(dim, thd, unequal, causal, seed), dim, thd, causal) + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(stream): + result = backward() + torch.cuda.current_stream().wait_stream(stream) + _check(result, reference) + + +@pytest.mark.parametrize( + "dim,thd,bottom_right", [(128, True, False), (256, True, False), (256, True, True)] +) +def test_frost_backward_replay(frost_enabled, dim, thd, bottom_right): + """Repeated replay overwrites poisoned gradients and reinitializes scratch.""" + del frost_enabled + backward, reference = _calls( + _case(dim, thd, False, True, bottom_right=bottom_right), dim, thd, True, bottom_right + ) + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(stream): + _ = backward() + torch.cuda.current_stream().wait_stream(stream) + # Removing preparation must not affect cached C++ dispatch or graph capture. + import tvm_ffi + + def unexpected_python(*args, **kwargs): + raise AssertionError("Steady-state backward re-entered Python preparation") + + original = tvm_ffi.get_global_func("get_frost_attention_backward") + tvm_ffi.register_global_func("get_frost_attention_backward", unexpected_python, override=True) + try: + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph, stream=stream): + result = backward() + for _ in range(3): + for value in result: + value.fill_(float("nan")) + graph.replay() + _check(result, reference) + finally: + tvm_ffi.register_global_func("get_frost_attention_backward", original, override=True) + + +def test_frost_d256_ragged_fallback(frost_enabled, monkeypatch): + """D256 unequal THD uses the existing backend unless strict mode is requested.""" + del frost_enabled + monkeypatch.setenv("NVTE_FUSED_ATTN_FROST_UNIFORM", "0") + monkeypatch.setenv("NVTE_FUSED_ATTN_FROST", "1") + backward, reference = _calls(_case(256, True, True, True), 256, True, True) + _check(backward(), reference) + monkeypatch.setenv("NVTE_FUSED_ATTN_FROST", "2") + with pytest.raises(RuntimeError, match="Requested cuDNN OSS backward unavailable"): + backward() + + +@pytest.mark.parametrize("thd", [False, True]) +def test_frost_d256_bottom_right(frost_enabled, thd): + """Rectangular masks preserve their diagonal and keep distinct cached plans.""" + del frost_enabled + for seed in (31, 47): + for bottom_right in (True, False, True): + case = _case(256, thd, False, True, seed, bottom_right) + backward, reference = _calls(case, 256, thd, True, bottom_right) + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(stream): + result = backward() + torch.cuda.current_stream().wait_stream(stream) + _check(result, reference) + + +def test_frost_d256_bottom_right_q_longer_fallback(frost_enabled, monkeypatch): + """Q > KV declines the dense OSS declaration and retains correct THD fallback.""" + del frost_enabled + case = _case(256, True, False, True, bottom_right=True, kv_length=128, q_length=256) + monkeypatch.setenv("NVTE_FUSED_ATTN_FROST", "1") + backward, reference = _calls(case, 256, True, True, True) + for mode in (1, 2, 1): + monkeypatch.setenv("NVTE_FUSED_ATTN_FROST", str(mode)) + if mode == 2: + with pytest.raises(RuntimeError, match="Requested cuDNN OSS backward unavailable"): + backward() + else: + result = backward() + _check(result, reference) + dq = result[0].view(2, 256, 4, 256) + assert torch.count_nonzero(dq[:, :128]).item() == 0 + + +def test_frost_d128_bottom_right_fallback(frost_enabled, monkeypatch): + """The D128 engine does not declare bottom-right causal support.""" + del frost_enabled + monkeypatch.setenv("NVTE_FUSED_ATTN_FROST", "1") + case = _case(128, True, False, True, bottom_right=True) + backward, reference = _calls(case, 128, True, True, True) + _check(backward(), reference) + monkeypatch.setenv("NVTE_FUSED_ATTN_FROST", "2") + with pytest.raises(RuntimeError, match="Requested cuDNN OSS backward unavailable"): + backward() + + +def _pitched_case(cp, sk, seed=31): + """Full THD buffers with the physical origins of a causal AllGather step.""" + case = _case(256, True, False, True, seed, bottom_right=True, kv_length=sk) + q, k, v, do, prefixes, reference, sq, _ = case + pq, pk = 2 * sq, 2 * cp * sq + origin = sq if sk > pk // 2 else 0 + + def expand(value, kv, fill): + pitch, length, start = (pk, sk, 0) if kv else (pq, sq, origin) + heads, dim = value.shape[1:] + full = torch.full((2, pitch, heads, dim), fill, device=value.device, dtype=value.dtype) + full[:, start : start + length].copy_(value.view(2, length, heads, dim)) + return full.view(2 * pitch, heads, dim) + + q, k, v, do = [expand(value, i in (1, 2), 11) for i, value in enumerate((q, k, v, do))] + reference = tuple(expand(value, i != 0, 0) for i, value in enumerate(reference)) + origins = ( + torch.tensor([origin, pq + origin, 2 * pq], device="cuda", dtype=torch.int32), + torch.tensor([0, pk, 2 * pk], device="cuda", dtype=torch.int32), + ) + return (q, k, v, do, prefixes, reference, sq, sk), origins + + +@pytest.mark.parametrize( + "cp,chunk", [(cp, chunk) for cp in (2, 4, 8) for chunk in range(1, 2 * cp + 1)] +) +def test_frost_d256_pitched_allgather(frost_enabled, monkeypatch, cp, chunk): + """Every rank/step geometry, including the zero gradients in physical holes.""" + del frost_enabled + monkeypatch.setenv("NVTE_FUSED_ATTN_FROST_ALLGATHER", "1") + for seed in (31, 47): + case, origins = _pitched_case(cp, chunk * 128, seed) + backward, reference = _calls(case, 256, True, True, True, origins) + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(stream): + result = backward() + torch.cuda.current_stream().wait_stream(stream) + _check(result, reference) + + +def test_frost_d256_pitched_opt_in(frost_enabled, monkeypatch): + """The extra layout contract requires its own opt-in and preserves fallback.""" + del frost_enabled + monkeypatch.setenv("NVTE_FUSED_ATTN_FROST_ALLGATHER", "0") + monkeypatch.setenv("NVTE_FUSED_ATTN_FROST", "1") + case, origins = _pitched_case(2, 384) + backward, reference = _calls(case, 256, True, True, True, origins) + _check(backward(), reference) + monkeypatch.setenv("NVTE_FUSED_ATTN_FROST", "2") + with pytest.raises(RuntimeError, match="Requested cuDNN OSS backward unavailable"): + backward() + + +@pytest.mark.parametrize("sk", [128, 384]) +def test_frost_d256_pitched_replay(frost_enabled, monkeypatch, sk): + """Graph replay reuses native dispatch, preserving physical origins and zero holes.""" + del frost_enabled + monkeypatch.setenv("NVTE_FUSED_ATTN_FROST_ALLGATHER", "1") + case, origins = _pitched_case(2, sk) + backward, reference = _calls(case, 256, True, True, True, origins) + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(stream): + _ = backward() + torch.cuda.current_stream().wait_stream(stream) + import tvm_ffi + + def unexpected_python(*args, **kwargs): + raise AssertionError("Steady-state backward re-entered Python preparation") + + original = tvm_ffi.get_global_func("get_frost_attention_backward") + tvm_ffi.register_global_func("get_frost_attention_backward", unexpected_python, override=True) + try: + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph, stream=stream): + result = backward() + for _ in range(3): + for value in result: + value.fill_(float("nan")) + graph.replay() + _check(result, reference) + finally: + tvm_ffi.register_global_func("get_frost_attention_backward", original, override=True) + + +@pytest.mark.parametrize("causal", [False, True]) +def test_frost_d128_nozero_changed_inputs(frost_enabled, monkeypatch, causal): + """Changed inputs and fresh bindings agree with FP64 with either output policy.""" + del frost_enabled + for seed in (31, 47, 61): + case = _case(128, True, False, causal, seed, kv_length=128) + backward, reference = _calls(case, 128, True, causal) + outputs = [] + for flag in ("0", "1"): + monkeypatch.setenv("NVTE_FUSED_ATTN_FROST_D128_NOZERO", flag) + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(stream): + actual = backward() + torch.cuda.current_stream().wait_stream(stream) + _check(actual, reference) + outputs.append(actual) + for a, b in zip(*outputs): + torch.testing.assert_close(a, b, atol=0, rtol=0) + + +@pytest.mark.parametrize("causal", [False, True]) +def test_frost_d128_nozero_poisoned_graph(frost_enabled, monkeypatch, causal): + """The native chain overwrites every poisoned gradient without output fills.""" + del frost_enabled + monkeypatch.setenv("NVTE_FUSED_ATTN_FROST_D128_NOZERO", "1") + backward, reference = _calls(_case(128, True, False, causal, kv_length=128), 128, True, causal) + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(stream): + _ = backward() + torch.cuda.current_stream().wait_stream(stream) + import tvm_ffi + + def unexpected_python(*args, **kwargs): + raise AssertionError("Steady-state backward re-entered Python preparation") + + original = tvm_ffi.get_global_func("get_frost_attention_backward") + tvm_ffi.register_global_func("get_frost_attention_backward", unexpected_python, override=True) + try: + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph, stream=stream): + result = backward() + for _ in range(4): + for value in result: + value.fill_(float("nan")) + graph.replay() + _check(result, reference) + finally: + tvm_ffi.register_global_func("get_frost_attention_backward", original, override=True) + + +@pytest.mark.parametrize( + "mode,flag,kv_length,unequal,expected", + [ + (2, 0, 128, False, 3), + (2, 1, 128, False, 0), + (1, 1, 128, False, 3), + (0, 1, 128, False, 3), + (2, 1, 256, False, 3), + (2, 1, None, True, 3), + ], +) +def test_frost_d128_nozero_fill_calls( + frost_enabled, monkeypatch, mode, flag, kv_length, unequal, expected +): + """Only strict compact square THD may omit fully overwritten output fills.""" + del frost_enabled + monkeypatch.setenv("NVTE_FUSED_ATTN_FROST", str(mode)) + monkeypatch.setenv("NVTE_FUSED_ATTN_FROST_D128_NOZERO", str(flag)) + case = _case(128, True, unequal, True, kv_length=kv_length) + backward, reference = _calls(case, 128, True, True) + _check(backward(), reference) + shapes = [list(t.shape) for t in case[:3]] + with torch.profiler.profile( + activities=[torch.profiler.ProfilerActivity.CPU], record_shapes=True + ) as prof: + result = backward() + _check(result, reference) + fills = [ + e.input_shapes + for e in prof.events() + if e.name == "aten::fill_" and e.input_shapes and e.input_shapes[0] in shapes + ] + assert len(fills) == expected, (mode, flag, kv_length, fills) + print("OUTPUT_FILL_CALLS", mode, flag, kv_length, len(fills), flush=True) + + +def _run_invalid_compact_metadata(kind): + """Each intentional device trap runs in its own child process.""" + from transformer_engine.common import ( + _get_shared_object_file, + _load_tvm_ffi_library, + _register_cutedsl_backends, + ) + + assert _load_tvm_ffi_library() and _register_cutedsl_backends() + core = ctypes.CDLL(str(_get_shared_object_file("core"))) + core.nvte_set_cutedsl_backend.argtypes = [ctypes.c_int] + core.nvte_set_cutedsl_backend.restype = None + core.nvte_set_cutedsl_backend(1) + os.environ["NVTE_FUSED_ATTN_FROST"] = "2" + os.environ["NVTE_FUSED_ATTN_FROST_D128_NOZERO"] = "1" + case = _case(128, True, False, True, kv_length=128) + origins = [x.clone() for x in case[4]] + backward, reference = _calls(case, 128, True, True, origins=origins) + _check(backward(), reference) + if kind == "graph_terminal": + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(stream): + _ = backward() + torch.cuda.current_stream().wait_stream(stream) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph, stream=stream): + _ = backward() + origins[1][-1] -= 1 + print("EXPECT_DEVICE_GUARD_TRAP", kind, flush=True) + graph.replay() + else: + target = { + "q_length": case[4][0], + "kv_length": case[4][1], + "q_origin": origins[0], + "kv_origin": origins[1], + }[kind] + target[1] -= 1 + print("EXPECT_DEVICE_GUARD_TRAP", kind, flush=True) + _ = backward() + torch.cuda.synchronize() + raise AssertionError("Invalid metadata escaped the per-execution GPU guard") + + +@pytest.mark.parametrize( + "kind", ["q_length", "kv_length", "q_origin", "kv_origin", "graph_terminal"] +) +def test_frost_d128_nozero_metadata_guard(frost_enabled, kind): + """Invalid live metadata must fail before uninitialized output can be returned.""" + del frost_enabled + command = [sys.executable] + if sys.flags.no_site: + command.append("-S") + # Retain the test interpreter's import paths, including isolated build roots. + # CUDA device traps run in a fresh process so the parent context stays usable. + code = ( + "import json, runpy, sys; sys.path[:] = json.loads(sys.argv[1]); " + "runpy.run_path(sys.argv[2])['_run_invalid_compact_metadata'](sys.argv[3])" + ) + result = subprocess.run( + command + ["-c", code, json.dumps(sys.path), str(Path(__file__).resolve()), kind], + text=True, + capture_output=True, + check=False, + timeout=240, + ) + output = result.stdout + result.stderr + assert result.returncode != 0, output + assert "EXPECT_DEVICE_GUARD_TRAP " + kind in output, output + assert "NVTE_FUSED_ATTN_FROST_UNIFORM: invalid sequence/offset" in output, output diff --git a/transformer_engine/common/CMakeLists.txt b/transformer_engine/common/CMakeLists.txt index 077a9a2b3a1..35676293092 100644 --- a/transformer_engine/common/CMakeLists.txt +++ b/transformer_engine/common/CMakeLists.txt @@ -168,7 +168,7 @@ set(CUTLASS_TOOLS_INCLUDE_DIR find_package(Python COMPONENTS Interpreter Development.Module REQUIRED) option(NVTE_WITH_CUTEDSL - "Build the CuTeDSL quantize backend bridge (requires the apache-tvm-ffi Python package)" OFF) + "Build the CuTeDSL kernel bridge (requires the apache-tvm-ffi Python package)" OFF) if(NVTE_WITH_CUTEDSL) # Building libtransformer_engine.so with CuTeDSL support requires the TVM-FFI headers. if(NOT EXISTS "${TVM_FFI_INCLUDE_DIR}/tvm/ffi/any.h") @@ -283,6 +283,7 @@ list(APPEND transformer_engine_cpp_sources if(NVTE_WITH_CUTEDSL) list(APPEND transformer_engine_cpp_sources init_python.cpp) + list(APPEND transformer_engine_cuda_sources fused_attn/frost_attn.cu) endif() list(APPEND transformer_engine_cuda_sources diff --git a/transformer_engine/common/CuTeDSL/__init__.py b/transformer_engine/common/CuTeDSL/__init__.py index 50d536852c2..3745001d7db 100644 --- a/transformer_engine/common/CuTeDSL/__init__.py +++ b/transformer_engine/common/CuTeDSL/__init__.py @@ -22,3 +22,9 @@ def register_cutedsl_backends(): tvm_ffi.register_global_func( "get_mxfp8_quantization_function", get_mxfp8_quantization_function, override=True ) + + from transformer_engine.common.CuTeDSL.frost_attn import get_frost_attention_backward + + tvm_ffi.register_global_func( + "get_frost_attention_backward", get_frost_attention_backward, override=True + ) diff --git a/transformer_engine/common/CuTeDSL/frost_attn.py b/transformer_engine/common/CuTeDSL/frost_attn.py new file mode 100644 index 00000000000..0b180395fa4 --- /dev/null +++ b/transformer_engine/common/CuTeDSL/frost_attn.py @@ -0,0 +1,218 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""Prepare cuDNN OSS backward plans for TE's native dispatcher. + +Only preparation enters Python. The registry entry is the compiled TVM-FFI +function for the complete backward chain, including its scratch initialization. +The private FE prepared-plan adapter is deliberately confined to this module. +""" + +from functools import partial +from importlib.metadata import PackageNotFoundError, version +import logging + +_LOG = logging.getLogger("transformer_engine.frost") +_OWNERS = {} +_ROLES = ("q", "k", "v", "o", "do", "stats", "dq", "dk", "dv") +_ORIGIN_ROLES = ("q", "k", "v", "o", "do", "dq", "dk", "dv", "stats") +_POINTER_ROLES = _ROLES + ("seq_q", "seq_kv") + tuple("ragged_" + r for r in _ORIGIN_ROLES) +_ENGINES = {128: 20609, 256: 20607} + + +def _declarations(batch, heads, kv_heads, sq, sk, dim, tq, thd, q_pitch=0, kv_pitch=0): + """Logical BHSD declarations over TE's BSHD or packed token-major storage.""" + result = {} + for role in _ROLES: + kv = role in ("k", "v", "dk", "dv") + h, s = (kv_heads, sk) if kv else (heads, sq) + if role == "stats": + shape = (batch, heads, sq, 1) + stride = (tq * heads, 1, heads, 1) if thd else (heads * sq, sq, 1, 1) + else: + pitch = (kv_pitch if kv else q_pitch) or s + shape, stride = (batch, h, s, dim), (pitch * h * dim, dim, h * dim, 1) + result[role] = (shape, stride) + return result + + +def _native_metadata(spec, dim, thd, expected_prefix): + """Refuse an unfamiliar prepared ABI instead of dropping optional operands.""" + from cudnn.frost.compiled_cache import positional_entry + import tvm_ffi + + if ( + spec.name != f"sdpa_bwd_sm100_d{dim}" + or spec.scale_log2 + or spec.packed_tile_groups + or bool(spec.length_form) != thd + ): + raise RuntimeError("Unsupported cuDNN OSS prepared backward contract") + roles = tuple(spec.roles) + expected = _POINTER_ROLES if thd else _ROLES + (("seq_kv", "delta") if dim == 256 else ()) + if roles != expected or len(spec.operands) != len(roles): + raise RuntimeError(f"Unsupported cuDNN OSS operand order: {roles}") + native = positional_entry(spec.artifact.entry) + if not isinstance(native, tvm_ffi.Function): + raise RuntimeError("cuDNN OSS must expose a compiled native TVM-FFI entry") + if thd: + bound = getattr(spec.fn, "entry", None) + if not isinstance(bound, partial) or tuple(bound.args) != expected_prefix: + raise RuntimeError("Unsupported cuDNN OSS packed backward prefix") + elif expected_prefix: + raise RuntimeError("Dense backward cannot have a packed prefix") + metadata = [1, int(spec.workspace_bytes), int(thd), len(expected_prefix), len(roles)] + metadata.extend(expected_prefix) + for role, operand in zip(roles, spec.operands): + if operand is None: + metadata.extend((-1, 0, 1)) + continue + if role not in _POINTER_ROLES or operand.opaque_bytes or operand.packed_tile_bytes: + raise RuntimeError(f"Unsupported cuDNN OSS operand: {role}") + span = int(operand.span) * int(operand.itemsize) + if role in ("seq_q", "seq_kv"): + span += 4 # TE supplies B+1 cumulative lengths, lens_form=3. + if span <= 0 or span > 2**63 - 1: + raise RuntimeError("cuDNN OSS operand byte span is not representable") + metadata.extend((_POINTER_ROLES.index(role), span, int(operand.alignment))) + return native, metadata + + +def get_frost_attention_backward( + key, + batch, + heads, + kv_heads, + sq, + sk, + dim, + tq, + tk, + thd, + causal, + scale, + device, + bottom_right=False, + q_pitch=0, + kv_pitch=0, +): + """Compile a graph declaration and register its native chain; allocate no data.""" + try: + import cudnn + from cuda.bindings import runtime + import tvm_ffi + except ImportError as error: + _LOG.warning("cuDNN OSS attention dependency unavailable: %s", error) + return False + + from packaging.version import Version + + try: + dsl_version = Version(version("nvidia-cutlass-dsl")) + except PackageNotFoundError: + _LOG.warning("cuDNN OSS attention requires nvidia-cutlass-dsl >= 4.8.0") + return False + if dsl_version < Version("4.8.0"): + _LOG.warning("cuDNN OSS attention requires nvidia-cutlass-dsl >= 4.8.0") + return False + if dim not in _ENGINES or (thd and dim != 128): + return False + if bottom_right and (dim != 256 or not causal): + return False + if q_pitch or kv_pitch: + if thd or dim != 256 or q_pitch < sq or kv_pitch < sk: + return False + from cutlass.base_dsl.version_info import CUDA_VERSION + + if CUDA_VERSION.major < 13: + _LOG.warning("cuDNN OSS backward requires the CUDA 13 flavor: nvidia-cutlass-dsl[cu13]") + return False + error, current = runtime.cudaGetDevice() + if int(error) != 0 or int(current) != device: + raise RuntimeError("cuDNN OSS preparation must use the calling CUDA device") + graph = cudnn.pygraph( + io_data_type=cudnn.data_type.BFLOAT16, + intermediate_data_type=cudnn.data_type.FLOAT, + compute_data_type=cudnn.data_type.FLOAT, + ) + declarations = _declarations(batch, heads, kv_heads, sq, sk, dim, tq, thd, q_pitch, kv_pitch) + refs = {} + + def attach(role, tensor): + if not thd: + return + kv = role in ("k", "v", "dk", "dv") + h = kv_heads if kv else heads + offset = graph.tensor( + name=role + "_origin", + dim=[batch + 1, 1, 1, 1], + stride=[1, 1, 1, 1], + data_type=cudnn.data_type.INT32, + ) + tensor.set_ragged_offset(offset) + tensor.set_ragged_offset_multiplier(h if role == "stats" else h * dim) + + for role in _ROLES[:6]: + shape, stride = declarations[role] + refs[role] = graph.tensor( + name=role, + dim=shape, + stride=stride, + data_type=cudnn.data_type.FLOAT if role == "stats" else cudnn.data_type.BFLOAT16, + ) + attach(role, refs[role]) + lengths = {} + if thd: + for role in ("q", "kv"): + lengths[role] = graph.tensor( + name="length_" + role, + dim=[batch, 1, 1, 1], + stride=[1, 1, 1, 1], + data_type=cudnn.data_type.INT32, + ) + outputs = graph.sdpa_backward( + q=refs["q"], + k=refs["k"], + v=refs["v"], + o=refs["o"], + dO=refs["do"], + stats=refs["stats"], + attn_scale=scale, + use_causal_mask=causal and not bottom_right, + use_causal_mask_bottom_right=causal and bottom_right, + use_padding_mask=thd, + seq_len_q=lengths.get("q"), + seq_len_kv=lengths.get("kv"), + max_total_seq_len_q=tq if thd else None, + max_total_seq_len_kv=tk if thd else None, + use_deterministic_algorithm=False, + ) + for role, tensor in zip(_ROLES[6:], outputs): + shape, stride = declarations[role] + tensor.set_output(True).set_data_type(cudnn.data_type.BFLOAT16) + tensor.set_dim(shape).set_stride(stride) + attach(role, tensor) + try: + graph.validate() + graph.build_operation_graph() + try: + graph.create_execution_plan(_ENGINES[dim], {}) + except ValueError as error: + if "no python engine on this graph owns engine_id" not in str(error): + raise + _LOG.warning("Installed cuDNN Frontend lacks the requested OSS engine: %s", error) + return False + graph.build_plans() + except (NotImplementedError, cudnn.cudnnGraphNotSupportedError) as error: + _LOG.warning("cuDNN OSS attention declined this declaration: %s", error) + return False + spec = graph._compiled_plans[graph._plan_index]._prepared.spec + prefix = (min(batch * sq, tq), min(batch * sk, tk), tq, tk, sq, sk, 0) if thd else () + native, metadata = _native_metadata(spec, dim, thd, prefix) + if spec.device_index != device or float(spec.scale) != float(scale): + raise RuntimeError("cuDNN OSS plan device/scale differs from its TE declaration") + # Keep the graph, compiled owner and shared library alive through graph replay. + _OWNERS[key] = graph + tvm_ffi.register_global_func(key, native, override=False) + return " ".join(str(value) for value in metadata) diff --git a/transformer_engine/common/fused_attn/frost_attn.cu b/transformer_engine/common/fused_attn/frost_attn.cu new file mode 100644 index 00000000000..12900fac441 --- /dev/null +++ b/transformer_engine/common/fused_attn/frost_attn.cu @@ -0,0 +1,328 @@ +/************************************************************************* + * 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 "../tvm_ffi_bridge.h" +#include "../util/cuda_runtime.h" +#include "../util/system.h" +#include "frost_attn.h" +#include "frost_output_policy.h" + +namespace transformer_engine { +namespace fused_attn { +namespace { + +struct Operand { + int role; + size_t bytes; + size_t alignment; +}; + +struct Plan { + tvm::ffi::Function function; + size_t workspace_bytes; + bool length_form; + std::vector prefix; + std::vector operands; +}; + +struct UniformLayout { + size_t q_pitch; + size_t kv_pitch; + size_t q_origin; +}; + +bool uniform_layout(const FusedAttnConfig &c, UniformLayout *layout) { + if (!getenv("NVTE_FUSED_ATTN_FROST_UNIFORM")) return false; + *layout = {c.max_seqlen_q, c.max_seqlen_kv, 0}; + if (c.num_tokens_q == c.batch_size * c.max_seqlen_q && + c.num_tokens_kv == c.batch_size * c.max_seqlen_kv) + return true; + if (!getenv("NVTE_FUSED_ATTN_FROST_ALLGATHER") || !c.is_causal_bottom_right || + c.max_seqlen_q % 128 != 0 || c.num_tokens_kv % c.batch_size != 0) + return false; + const size_t pq = 2 * c.max_seqlen_q; + const size_t pk = c.num_tokens_kv / c.batch_size; + const size_t cp = pk / pq; + if ((cp != 2 && cp != 4 && cp != 8) || pk % pq != 0 || c.num_tokens_q != c.batch_size * pq || + c.max_seqlen_kv > pk || c.max_seqlen_kv % c.max_seqlen_q != 0) + return false; + *layout = {pq, pk, c.max_seqlen_kv > pk / 2 ? c.max_seqlen_q : 0}; + return true; +} + +bool eligible(const FusedAttnConfig &c) { + const bool causal = c.is_causal || c.is_causal_bottom_right; + if (cuda::sm_arch(cuda::current_device()) != 100 || c.deterministic || !c.is_training || + c.qkv_dtype != kNVTEBFloat16 || c.o_dtype != kNVTEBFloat16 || c.do_dtype != kNVTEBFloat16 || + c.dqkv_dtype != kNVTEBFloat16 || c.head_dim_qk != c.head_dim_v || + (c.head_dim_qk != 128 && c.head_dim_qk != 256) || c.bias_type != NVTE_NO_BIAS || + c.softmax_type != NVTE_VANILLA_SOFTMAX || c.dropout != 0 || c.is_paged_kv || + (c.bottom_right_diagonal && (c.head_dim_qk != 256 || !c.is_causal_bottom_right)) || + c.window_size_left != -1 || (c.window_size_right != -1 && c.window_size_right != 0) || + (c.window_size_right == 0 && !causal) || !std::isfinite(c.attn_scale) || c.attn_scale <= 0 || + c.batch_size == 0 || c.batch_size > 65535 || c.num_gqa_groups == 0 || + c.num_attn_heads < c.num_gqa_groups || c.num_attn_heads % c.num_gqa_groups != 0 || + c.num_attn_heads > 65535 || c.max_seqlen_q == 0 || c.max_seqlen_kv == 0 || + c.max_seqlen_q > INT32_MAX - 256 || c.max_seqlen_kv > INT32_MAX - 256) + return false; + const auto format = c.is_ragged_q ? NVTE_THD : NVTE_BSHD; + if (c.o_format != format || c.do_format != format || c.dqkv_layout != c.qkv_layout) return false; + if (c.qkv_layout == NVTE_BSHD_BSHD_BSHD) { + return !c.is_padding && c.max_seqlen_q % 128 == 0 && c.max_seqlen_kv % 128 == 0; + } + if (c.qkv_layout != NVTE_THD_THD_THD || !c.is_padding || !c.uses_ragged_stats || + c.num_tokens_q == 0 || c.num_tokens_kv == 0 || c.num_tokens_q > INT32_MAX || + c.num_tokens_kv > INT32_MAX) + return false; + if (c.head_dim_qk == 128) return true; + // D256's current OSS engine declines ragged attention. This separate promise + // enables an equivalent dense declaration; validate its values on EVERY call. + UniformLayout layout; + return uniform_layout(c, &layout); +} + +const Plan *get_plan(const FusedAttnConfig &c, cudaStream_t stream) { + // Process lifetime, like TVMFFIConfigCache: graph executables can outlive Python. + static auto *cache = new std::map>; + static auto *mutex = new std::mutex; + static std::atomic serial{0}; + { + std::lock_guard lock(*mutex); + auto found = cache->find(c); + if (found != cache->end()) return found->second.get(); + } + cudaStreamCaptureStatus capture; + NVTE_CHECK_CUDA(cudaStreamIsCapturing(stream, &capture)); + NVTE_CHECK(capture == cudaStreamCaptureStatusNone, + "Warm up cuDNN OSS attention before CUDA graph capture"); + auto ¢ral = tvm_ffi_bridge::TVMFFICentral::getInstance(); + if (!central.get_cutedsl_backend_enabled() || !central.get_tvm_ffi_available()) return nullptr; + auto prepare = tvm::ffi::Function::GetGlobal("get_frost_attention_backward"); + if (!prepare) return nullptr; + const std::string key = "transformer_engine.frost.backward." + std::to_string(serial++); + const bool thd = c.is_ragged_q && c.head_dim_qk == 128; + UniformLayout layout{0, 0, 0}; + if (c.is_ragged_q && c.head_dim_qk == 256) { + NVTE_CHECK(uniform_layout(c, &layout), "Invalid cuDNN OSS uniform layout"); + } + // Never hold a C++ cache lock while entering Python: another Python thread + // might hold the GIL and be waiting for that lock. + auto result = (*prepare)( + tvm::ffi::String(key), static_cast(c.batch_size), + static_cast(c.num_attn_heads), static_cast(c.num_gqa_groups), + static_cast(c.max_seqlen_q), static_cast(c.max_seqlen_kv), + static_cast(c.head_dim_qk), static_cast(c.num_tokens_q), + static_cast(c.num_tokens_kv), thd, c.is_causal || c.is_causal_bottom_right, + static_cast(c.attn_scale), cuda::current_device(), c.bottom_right_diagonal, + static_cast(layout.q_pitch), static_cast(layout.kv_pitch)); + // A bounded scalar string keeps the optional bridge free of TVM object-type + // registration symbols, so loading TE with the DSL backend disabled still works. + auto encoded = result.try_cast(); + auto declined = result.try_cast(); + NVTE_CHECK(encoded || (declined && !*declined), "Invalid cuDNN OSS preparation response"); + std::unique_ptr plan; + if (encoded) { + NVTE_CHECK(encoded->size() < 4096, "Oversized cuDNN OSS native plan metadata"); + std::istringstream input{std::string(*encoded)}; + std::vector v; + int64_t value; + while (input >> value) v.push_back(value); + NVTE_CHECK(input.eof(), "Malformed cuDNN OSS native plan metadata"); + auto at = [&v](size_t i) { return v[i]; }; + NVTE_CHECK(v.size() >= 5 && at(0) == 1, "Unknown cuDNN OSS native plan metadata"); + const auto prefix_count = at(3), operand_count = at(4); + NVTE_CHECK(at(1) > 0 && at(2) == static_cast(thd) && prefix_count == (thd ? 7 : 0) && + operand_count >= 9 && operand_count <= 20 && + v.size() == 5 + prefix_count + operand_count * 3, + "Malformed cuDNN OSS native plan metadata"); + auto function = tvm::ffi::Function::GetGlobal(key); + NVTE_CHECK(function.has_value(), "cuDNN OSS did not register its native function"); + plan = std::make_unique(Plan{*function, static_cast(at(1)), thd, {}, {}}); + size_t cursor = 5; + for (int64_t i = 0; i < prefix_count; ++i) plan->prefix.push_back(at(cursor++)); + for (int64_t i = 0; i < operand_count; ++i) { + const auto role = at(cursor++), bytes = at(cursor++), alignment = at(cursor++); + NVTE_CHECK(role >= -1 && role < 20 && bytes >= 0 && alignment > 0 && + (alignment & (alignment - 1)) == 0, + "Invalid cuDNN OSS operand metadata"); + plan->operands.push_back( + {static_cast(role), static_cast(bytes), static_cast(alignment)}); + } + } + std::lock_guard lock(*mutex); + auto inserted = cache->emplace(c, std::move(plan)); + return inserted.first->second.get(); +} + +bool overlaps(uintptr_t a, size_t a_bytes, uintptr_t b, size_t b_bytes) { + NVTE_CHECK(a_bytes <= UINTPTR_MAX - a && b_bytes <= UINTPTR_MAX - b, + "cuDNN OSS pointer range overflows"); + return a < b + b_bytes && b < a + a_bytes; +} + +__global__ void validate_uniform(const int32_t *cu_q, const int32_t *cu_kv, + const int32_t *origins_q, const int32_t *origins_kv, int64_t batch, + int64_t sq, int64_t sk, int64_t pq, int64_t pk, int64_t origin) { + const int64_t i = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + if (i <= batch && + (cu_q[i] != i * sq || cu_kv[i] != i * sk || + origins_q[i] != i * pq + (i < batch ? origin : 0) || origins_kv[i] != i * pk)) { + printf("NVTE_FUSED_ATTN_FROST_UNIFORM: invalid sequence/offset at %d\n", static_cast(i)); + asm volatile("trap;"); + } +} + +__global__ void unpack_stats(const float *packed, float *dense, int64_t size, int64_t heads, + int64_t sequence, int64_t pitch, int64_t origin) { + const int64_t i = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + if (i < size) { + const int64_t s = i % sequence, h = (i / sequence) % heads; + const int64_t b = i / (sequence * heads); + dense[i] = packed[(b * pitch + origin + s) * heads + h]; + } +} + +} // namespace + +bool frost_attention_backward(const FusedAttnConfig &cfg, const FrostPointers &pointers, + void *workspace, size_t *workspace_size, cudaStream_t stream) { + const int mode = getenv("NVTE_FUSED_ATTN_FROST"); + if (mode == 0) return false; + NVTE_CHECK(mode == 1 || mode == 2, "NVTE_FUSED_ATTN_FROST must be 0, 1 or 2"); + cfg.check_derived(); + FusedAttnConfig cache_cfg = cfg; + cache_cfg.device_id = cuda::current_device(); + const Plan *plan = eligible(cfg) ? get_plan(cache_cfg, stream) : nullptr; + if (!plan) { + NVTE_CHECK(mode != 2, + "Requested cuDNN OSS backward unavailable; check shape, mask, " + "cuDNN Frontend installation and NVTE_ENABLE_CUTEDSL_BACKEND=1"); + return false; + } + NVTE_CHECK(cache_cfg.device_id == cuda::current_device(), + "cuDNN OSS plan belongs to another device"); + const bool compact_d128 = frost_d128_compact_outputs( + cfg.qkv_layout, cfg.dqkv_layout, cfg.dqkv_dtype, cfg.head_dim_qk, cfg.head_dim_v, + cfg.batch_size, cfg.num_tokens_q, cfg.num_tokens_kv, cfg.max_seqlen_q, cfg.max_seqlen_kv, + cfg.num_attn_heads, cfg.num_gqa_groups); + const bool uniform = cfg.is_ragged_q && cfg.head_dim_qk == 256; + UniformLayout layout{0, 0, 0}; + if (uniform) { + NVTE_CHECK(uniform_layout(cfg, &layout), "Invalid cuDNN OSS uniform layout"); + } + const size_t stats_bytes = + uniform ? cfg.batch_size * cfg.num_attn_heads * cfg.max_seqlen_q * 4 : 0; + NVTE_CHECK(plan->workspace_bytes <= SIZE_MAX - 127 - stats_bytes, "cuDNN OSS workspace overflow"); + const size_t kernel_bytes = (plan->workspace_bytes + 127) & ~size_t(127); + const size_t required = kernel_bytes + stats_bytes; + if (workspace == nullptr) { + *workspace_size = required; + return true; + } + NVTE_CHECK(*workspace_size >= required, "cuDNN OSS workspace is too small"); + *workspace_size = required; + NVTE_CHECK(reinterpret_cast(workspace) % 128 == 0, + "cuDNN OSS workspace must be 128-byte aligned"); + // Validate the full declared physical buffers, including holes, before any + // launch or pointer adjustment. Packed Stats is larger than the dense scratch. + auto bound_bytes = [&](const Operand &op) -> size_t { + if (!uniform || op.role < 0 || op.role > 8) return op.bytes; + if (op.role == 5) return cfg.num_tokens_q * cfg.num_attn_heads * sizeof(float); + const bool kv = op.role == 1 || op.role == 2 || op.role == 7 || op.role == 8; + return (kv ? cfg.num_tokens_kv * cfg.num_gqa_groups : cfg.num_tokens_q * cfg.num_attn_heads) * + cfg.head_dim_qk * 2; + }; + // Validate caller pointers before any launch or Stats conversion. + for (const auto &op : plan->operands) { + if (op.role < 0) continue; + const auto address = reinterpret_cast(pointers[op.role]); + NVTE_CHECK(address != 0 && address % op.alignment == 0, + "cuDNN OSS operand is null or misaligned, role=", op.role); + NVTE_CHECK( + !overlaps(address, bound_bytes(op), reinterpret_cast(workspace), required), + "cuDNN OSS operand overlaps workspace"); + if (op.role >= 6 && op.role <= 8) { + for (const auto &other : plan->operands) { + if (other.role < 0 || other.role == op.role) continue; + NVTE_CHECK(!overlaps(address, bound_bytes(op), + reinterpret_cast(pointers[other.role]), bound_bytes(other)), + "cuDNN OSS gradient aliases an operand"); + } + } + } + // D128 keeps native THD Stats and descriptors. This guard proves that every + // physical output row is live; it does not lower THD to a dense engine. + if (compact_d128) { + validate_uniform<<<(cfg.batch_size + 256) / 256, 256, 0, stream>>>( + static_cast(pointers[9]), static_cast(pointers[10]), + static_cast(pointers[11]), static_cast(pointers[12]), + cfg.batch_size, cfg.max_seqlen_q, cfg.max_seqlen_kv, cfg.max_seqlen_q, cfg.max_seqlen_kv, + 0); + NVTE_CHECK_CUDA(cudaGetLastError()); + } + auto actual = pointers; + if (uniform) { + const size_t metadata_bytes = (cfg.batch_size + 1) * sizeof(int32_t); + for (int i : {9, 10, 11, 12}) { + const auto address = reinterpret_cast(pointers[i]); + NVTE_CHECK(address != 0 && address % alignof(int32_t) == 0, + "THD metadata is null or misaligned"); + NVTE_CHECK( + !overlaps(address, metadata_bytes, reinterpret_cast(workspace), required), + "THD metadata overlaps workspace"); + for (const auto &op : plan->operands) { + if (op.role >= 6 && op.role <= 8) { + NVTE_CHECK(!overlaps(address, metadata_bytes, + reinterpret_cast(pointers[op.role]), bound_bytes(op)), + "cuDNN OSS gradient aliases THD metadata"); + } + } + } + validate_uniform<<<(cfg.batch_size + 256) / 256, 256, 0, stream>>>( + static_cast(pointers[9]), static_cast(pointers[10]), + static_cast(pointers[11]), static_cast(pointers[12]), + cfg.batch_size, cfg.max_seqlen_q, cfg.max_seqlen_kv, layout.q_pitch, layout.kv_pitch, + layout.q_origin); + NVTE_CHECK_CUDA(cudaGetLastError()); + actual[5] = static_cast(workspace) + kernel_bytes; + unpack_stats<<<(stats_bytes / 4 + 255) / 256, 256, 0, stream>>>( + static_cast(pointers[5]), static_cast(actual[5]), stats_bytes / 4, + cfg.num_attn_heads, cfg.max_seqlen_q, layout.q_pitch, layout.q_origin); + NVTE_CHECK_CUDA(cudaGetLastError()); + } + if (uniform && layout.q_origin != 0) { + const size_t offset = layout.q_origin * cfg.num_attn_heads * cfg.head_dim_qk * 2; + for (int role : {0, 3, 4, 6}) actual[role] = static_cast(actual[role]) + offset; + } + std::array frame; + size_t cursor = 0; + for (auto prefix : plan->prefix) frame[cursor++] = prefix; + for (const auto &op : plan->operands) { + if (op.role >= 0) frame[cursor] = reinterpret_cast(actual[op.role]); + ++cursor; + } + frame[cursor++] = reinterpret_cast(workspace); + frame[cursor++] = static_cast(cfg.attn_scale); + if (plan->length_form) frame[cursor++] = int32_t(3); + frame[cursor++] = reinterpret_cast(stream); + tvm::ffi::Any result; + plan->function.CallPacked(frame.data(), static_cast(cursor), &result); + return true; +} + +} // namespace fused_attn +} // namespace transformer_engine diff --git a/transformer_engine/common/fused_attn/frost_attn.h b/transformer_engine/common/fused_attn/frost_attn.h new file mode 100644 index 00000000000..360e378d7e0 --- /dev/null +++ b/transformer_engine/common/fused_attn/frost_attn.h @@ -0,0 +1,30 @@ +/************************************************************************* + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * + * See LICENSE for license information. + ************************************************************************/ + +#ifndef TRANSFORMER_ENGINE_COMMON_FUSED_ATTN_FROST_ATTN_H_ +#define TRANSFORMER_ENGINE_COMMON_FUSED_ATTN_FROST_ATTN_H_ + +#include + +#include +#include + +#include "config_and_params.h" + +namespace transformer_engine { +namespace fused_attn { + +// q,k,v,o,do,stats,dq,dk,dv,cu_q,cu_kv followed by the nine physical origins. +using FrostPointers = std::array; + +// A null workspace prepares the plan and returns the framework-owned byte count. +// False selects the existing backend. Once execution starts errors propagate. +bool frost_attention_backward(const FusedAttnConfig &cfg, const FrostPointers &pointers, + void *workspace, size_t *workspace_size, cudaStream_t stream); + +} // namespace fused_attn +} // namespace transformer_engine +#endif // TRANSFORMER_ENGINE_COMMON_FUSED_ATTN_FROST_ATTN_H_ diff --git a/transformer_engine/common/fused_attn/frost_output_policy.h b/transformer_engine/common/fused_attn/frost_output_policy.h new file mode 100644 index 00000000000..1e7c627860b --- /dev/null +++ b/transformer_engine/common/fused_attn/frost_output_policy.h @@ -0,0 +1,32 @@ +/************************************************************************* + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * + * See LICENSE for license information. + ************************************************************************/ + +#ifndef TRANSFORMER_ENGINE_COMMON_FUSED_ATTN_FROST_OUTPUT_POLICY_H_ +#define TRANSFORMER_ENGINE_COMMON_FUSED_ATTN_FROST_OUTPUT_POLICY_H_ + +#include + +#include "../util/system.h" +#include "transformer_engine/fused_attn.h" + +namespace transformer_engine::fused_attn { + +// Default-off output initialization policy. The framework may omit output initialization +// only under this same predicate used by common dispatch. Strict mode prevents +// native fallback; common dispatch validates the device lengths/origins on every +// execution before FROST writes all rows. Pitched and unequal layouts keep zeros. +inline bool frost_d128_compact_outputs(NVTE_QKV_Layout qkv, NVTE_QKV_Layout dqkv, NVTEDType dtype, + size_t dq, size_t dv, size_t batch, size_t tq, size_t tk, + size_t sq, size_t sk, size_t hq, size_t hkv) { + return getenv("NVTE_FUSED_ATTN_FROST") == 2 && + getenv("NVTE_FUSED_ATTN_FROST_D128_NOZERO") && qkv == NVTE_THD_THD_THD && + dqkv == qkv && dtype == kNVTEBFloat16 && dq == 128 && dv == 128 && batch > 0 && sq > 0 && + sq == sk && sq % 128 == 0 && tq % batch == 0 && tk % batch == 0 && tq / batch == sq && + tk / batch == sk && hkv > 0 && hq > hkv && hq % hkv == 0; +} + +} // namespace transformer_engine::fused_attn +#endif // TRANSFORMER_ENGINE_COMMON_FUSED_ATTN_FROST_OUTPUT_POLICY_H_ diff --git a/transformer_engine/common/fused_attn/fused_attn_f16_arbitrary_seqlen.cu b/transformer_engine/common/fused_attn/fused_attn_f16_arbitrary_seqlen.cu index 723e1991db9..5794a30ad9a 100644 --- a/transformer_engine/common/fused_attn/fused_attn_f16_arbitrary_seqlen.cu +++ b/transformer_engine/common/fused_attn/fused_attn_f16_arbitrary_seqlen.cu @@ -17,6 +17,9 @@ #include "../util/system.h" #include "fused_attn_f16_arbitrary_seqlen.h" #include "graph_cache.h" +#ifdef NVTE_WITH_CUTEDSL +#include "frost_attn.h" +#endif #include "graph_cache_debug.h" #include "utils.h" @@ -1155,13 +1158,44 @@ void fused_attn_arbitrary_seqlen_bwd(const FusedAttnConfig &cfg, const Tensor *i void *devPtrDropoutOffset = reinterpret_cast(reinterpret_cast(rng_state->data.dptr) + 1); - size_t workspace_size = 0; - - fused_attn_arbitrary_seqlen_bwd_impl( - cfg, devPtrQ, devPtrK, devPtrV, devPtrO, devPtrSoftmaxStats, devPtrBias, devPtrSoftmaxOffset, - devPtrdQ, devPtrdK, devPtrdV, devPtrdO, devPtrdBias, devPtrdSoftmaxOffset, devPtrDropoutSeed, - devPtrDropoutOffset, devPtrCuSeqlensQ, devPtrCuSeqlensKV, devPtrSeqOffsetsQ, - devPtrSeqOffsetsKV, workspace->data.dptr, &workspace_size, stream, handle); + size_t workspace_size = workspace->data.buffer_size_bytes(); + bool used_frost = false; +#ifdef NVTE_WITH_CUTEDSL + const fused_attn::FrostPointers frost_pointers{devPtrQ, + devPtrK, + devPtrV, + devPtrO, + devPtrdO, + devPtrSoftmaxStats, + devPtrdQ, + devPtrdK, + devPtrdV, + devPtrCuSeqlensQ, + devPtrCuSeqlensKV, + devPtrSeqOffsetsQ, + devPtrSeqOffsetsKV, + devPtrSeqOffsetsKV, + devPtrSeqOffsetsQ, + devPtrSeqOffsetsQ, + devPtrSeqOffsetsQ, + devPtrSeqOffsetsKV, + devPtrSeqOffsetsKV, + devPtrSeqOffsetsQ}; + used_frost = fused_attn::frost_attention_backward(cfg, frost_pointers, workspace->data.dptr, + &workspace_size, stream); +#else + NVTE_CHECK(transformer_engine::getenv("NVTE_FUSED_ATTN_FROST") != 2, + "cuDNN OSS attention requires a TE build with NVTE_WITH_CUTEDSL=ON"); +#endif + if (!used_frost) { + workspace_size = 0; + fused_attn_arbitrary_seqlen_bwd_impl( + cfg, devPtrQ, devPtrK, devPtrV, devPtrO, devPtrSoftmaxStats, devPtrBias, + devPtrSoftmaxOffset, devPtrdQ, devPtrdK, devPtrdV, devPtrdO, devPtrdBias, + devPtrdSoftmaxOffset, devPtrDropoutSeed, devPtrDropoutOffset, devPtrCuSeqlensQ, + devPtrCuSeqlensKV, devPtrSeqOffsetsQ, devPtrSeqOffsetsKV, workspace->data.dptr, + &workspace_size, stream, handle); + } if (workspace_size > 0) { if (workspace->data.dptr == nullptr) { diff --git a/transformer_engine/pytorch/csrc/extensions/attention.cpp b/transformer_engine/pytorch/csrc/extensions/attention.cpp index 98d55d2b8b5..f38789d4c1f 100644 --- a/transformer_engine/pytorch/csrc/extensions/attention.cpp +++ b/transformer_engine/pytorch/csrc/extensions/attention.cpp @@ -6,6 +6,7 @@ #include "../extensions.h" #include "common.h" +#include "common/fused_attn/frost_output_policy.h" #include "pybind.h" namespace transformer_engine::pytorch { @@ -506,10 +507,18 @@ std::vector fused_attn_bwd( } else if (dqkv_quantizer.is_none() || detail::IsFloat8CurrentScalingQuantizers(dqkv_quantizer.ptr()) || detail::IsMXFP8Quantizers(dqkv_quantizer.ptr())) { - if (dq_format == NVTE_QKV_Format::NVTE_THD) { + const bool compact_d128 = dqkv_quantizer.is_none() && cu_seqlens_q.dim() == 1 && + cu_seqlens_q.numel() > 1 && + cu_seqlens_kv.sizes() == cu_seqlens_q.sizes() && + q_shape.size() == 3 && k_shape.size() == 3 && v_shape == k_shape && + fused_attn::frost_d128_compact_outputs( + qkv_layout, dqkv_layout, static_cast(dqkv_fake_dtype), + q_shape[2], v_shape[2], cu_seqlens_q.numel() - 1, q_shape[0], + k_shape[0], max_seqlen_q, max_seqlen_kv, q_shape[1], k_shape[1]); + if (dq_format == NVTE_QKV_Format::NVTE_THD && !compact_d128) { dQ.fill_(0); } - if (dkv_format == NVTE_QKV_Format::NVTE_THD) { + if (dkv_format == NVTE_QKV_Format::NVTE_THD && !compact_d128) { dK.fill_(0); dV.fill_(0); }