From ae00949cbb7e1fc1d2c43cdf3bf2d90bcb2dc924 Mon Sep 17 00:00:00 2001 From: Yang Xu Date: Fri, 16 Oct 2026 11:39:02 -0700 Subject: [PATCH 1/8] Add cuDNN OSS backward dispatch for context parallel attention Signed-off-by: Yang Xu --- docs/envvars.rst | 55 ++++ tests/pytorch/attention/test_frost_attn.py | 183 ++++++++++++ transformer_engine/common/CMakeLists.txt | 3 +- transformer_engine/common/CuTeDSL/__init__.py | 6 + .../common/CuTeDSL/frost_attn.py | 196 +++++++++++++ .../common/fused_attn/frost_attn.cu | 263 ++++++++++++++++++ .../common/fused_attn/frost_attn.h | 30 ++ .../fused_attn_f16_arbitrary_seqlen.cu | 48 +++- 8 files changed, 776 insertions(+), 8 deletions(-) create mode 100644 tests/pytorch/attention/test_frost_attn.py create mode 100644 transformer_engine/common/CuTeDSL/frost_attn.py create mode 100644 transformer_engine/common/fused_attn/frost_attn.cu create mode 100644 transformer_engine/common/fused_attn/frost_attn.h diff --git a/docs/envvars.rst b/docs/envvars.rst index 09e2b462780..30fca64d248 100644 --- a/docs/envvars.rst +++ b/docs/envvars.rst @@ -216,6 +216,61 @@ 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. + 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. + +.. 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. 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 or gapped THD 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 :Type: ``int`` (0 or 1) diff --git a/tests/pytorch/attention/test_frost_attn.py b/tests/pytorch/attention/test_frost_attn.py new file mode 100644 index 00000000000..584707fcea2 --- /dev/null +++ b/tests/pytorch/attention/test_frost_attn.py @@ -0,0 +1,183 @@ +# 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 os + +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") + yield + setter(int(os.environ.get("NVTE_ENABLE_CUTEDSL_BACKEND", "0") != "0")) + + +def _case(dim, thd, unequal, causal, seed=31): + generator = torch.Generator(device="cpu").manual_seed(seed) + q_lengths, k_lengths = ((129, 63), (193, 97)) if unequal else ((128, 128), (256, 256)) + 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], -torch.inf + ) + result = scores.softmax(-1) @ 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): + 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") + ) + 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, + cu_seqlens_kv_padded=cu_kv, + 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", [(128, True), (256, True)]) +def test_frost_backward_replay(frost_enabled, dim, thd): + """Repeated replay overwrites poisoned gradients and reinitializes scratch.""" + del frost_enabled + backward, reference = _calls(_case(dim, thd, False, True), dim, thd, True) + 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() 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..7704a3d2fcf --- /dev/null +++ b/transformer_engine/common/CuTeDSL/frost_attn.py @@ -0,0 +1,196 @@ +# 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): + """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: + shape, stride = (batch, h, s, dim), (s * 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 +): + """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 + 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) + 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, + 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..1023239d10c --- /dev/null +++ b/transformer_engine/common/fused_attn/frost_attn.cu @@ -0,0 +1,263 @@ +/************************************************************************* + * 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" + +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; +}; + +bool eligible(const FusedAttnConfig &c) { + 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.window_size_left != -1 || + (c.window_size_right != -1 && c.window_size_right != 0) || + (c.window_size_right == 0 && !c.is_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. + return getenv("NVTE_FUSED_ATTN_FROST_UNIFORM") && + c.num_tokens_q == c.batch_size * c.max_seqlen_q && + c.num_tokens_kv == c.batch_size * c.max_seqlen_kv; +} + +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; + // 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, + static_cast(c.attn_scale), cuda::current_device()); + // 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) { + 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 * sq || + origins_kv[i] != i * sk)) { + 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) { + 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 * sequence + 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 uniform = cfg.is_ragged_q && cfg.head_dim_qk == 256; + 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 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, op.bytes, 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, op.bytes, reinterpret_cast(pointers[other.role]), + other.bytes), + "cuDNN OSS gradient aliases an operand"); + } + } + } + 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]), op.bytes), + "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); + 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); + NVTE_CHECK_CUDA(cudaGetLastError()); + } + 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/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) { From 827e7cfb5534ace796a89ffe08b3d684574033cc Mon Sep 17 00:00:00 2001 From: Yang Xu Date: Fri, 16 Oct 2026 13:56:06 -0700 Subject: [PATCH 2/8] Support bottom-right causal masks in cuDNN OSS backward Signed-off-by: Yang Xu --- docs/envvars.rst | 2 + tests/pytorch/attention/test_frost_attn.py | 49 ++++++++++++++++--- .../common/CuTeDSL/frost_attn.py | 7 ++- .../common/fused_attn/frost_attn.cu | 23 ++++----- 4 files changed, 62 insertions(+), 19 deletions(-) diff --git a/docs/envvars.rst b/docs/envvars.rst index 30fca64d248..74fd370c69e 100644 --- a/docs/envvars.rst +++ b/docs/envvars.rst @@ -235,6 +235,8 @@ backend-selection overview. 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, including rectangular calls. + 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 diff --git a/tests/pytorch/attention/test_frost_attn.py b/tests/pytorch/attention/test_frost_attn.py index 584707fcea2..729de3d0510 100644 --- a/tests/pytorch/attention/test_frost_attn.py +++ b/tests/pytorch/attention/test_frost_attn.py @@ -38,7 +38,7 @@ def frost_enabled(monkeypatch): setter(int(os.environ.get("NVTE_ENABLE_CUTEDSL_BACKEND", "0") != "0")) -def _case(dim, thd, unequal, causal, seed=31): +def _case(dim, thd, unequal, causal, seed=31, bottom_right=False): generator = torch.Generator(device="cpu").manual_seed(seed) q_lengths, k_lengths = ((129, 63), (193, 97)) if unequal else ((128, 128), (256, 256)) heads, kv_heads = 4, 1 @@ -53,7 +53,9 @@ def draw(s, h): 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], -torch.inf + torch.arange(sk)[None, :] + > torch.arange(sq)[:, None] + (sk - sq if bottom_right else 0), + -torch.inf, ) result = scores.softmax(-1) @ vr.repeat_interleave(heads // kv_heads, 0) grads = torch.autograd.grad(result, (qr, kr, vr), do.transpose(0, 1).double()) @@ -69,7 +71,7 @@ def draw(s, h): return q, k, v, do, prefixes, reference, max(q_lengths), max(k_lengths) -def _calls(case, dim, thd, causal): +def _calls(case, dim, thd, causal, bottom_right=False): from transformer_engine.pytorch.cpp_extensions.fused_attn import ( FusedAttnBackend, fused_attn_fwd, @@ -81,6 +83,9 @@ def _calls(case, dim, thd, causal): 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, @@ -140,11 +145,15 @@ def test_frost_backward_and_rebinding(frost_enabled, dim, thd, unequal, causal): _check(result, reference) -@pytest.mark.parametrize("dim,thd", [(128, True), (256, True)]) -def test_frost_backward_replay(frost_enabled, dim, thd): +@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), dim, thd, True) + 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): @@ -181,3 +190,31 @@ def test_frost_d256_ragged_fallback(frost_enabled, monkeypatch): 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_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() diff --git a/transformer_engine/common/CuTeDSL/frost_attn.py b/transformer_engine/common/CuTeDSL/frost_attn.py index 7704a3d2fcf..143f0a903ac 100644 --- a/transformer_engine/common/CuTeDSL/frost_attn.py +++ b/transformer_engine/common/CuTeDSL/frost_attn.py @@ -79,7 +79,7 @@ def _native_metadata(spec, dim, thd, expected_prefix): def get_frost_attention_backward( - key, batch, heads, kv_heads, sq, sk, dim, tq, tk, thd, causal, scale, device + key, batch, heads, kv_heads, sq, sk, dim, tq, tk, thd, causal, scale, device, bottom_right=False ): """Compile a graph declaration and register its native chain; allocate no data.""" try: @@ -102,6 +102,8 @@ def get_frost_attention_backward( 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 from cutlass.base_dsl.version_info import CUDA_VERSION if CUDA_VERSION.major < 13: @@ -158,7 +160,8 @@ def attach(role, tensor): dO=refs["do"], stats=refs["stats"], attn_scale=scale, - use_causal_mask=causal, + 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"), diff --git a/transformer_engine/common/fused_attn/frost_attn.cu b/transformer_engine/common/fused_attn/frost_attn.cu index 1023239d10c..5c1621d14e6 100644 --- a/transformer_engine/common/fused_attn/frost_attn.cu +++ b/transformer_engine/common/fused_attn/frost_attn.cu @@ -39,15 +39,16 @@ struct Plan { }; 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.window_size_left != -1 || - (c.window_size_right != -1 && c.window_size_right != 0) || - (c.window_size_right == 0 && !c.is_causal) || !std::isfinite(c.attn_scale) || - c.attn_scale <= 0 || c.batch_size == 0 || c.batch_size > 65535 || c.num_gqa_groups == 0 || + (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) @@ -91,13 +92,13 @@ const Plan *get_plan(const FusedAttnConfig &c, cudaStream_t stream) { const bool thd = c.is_ragged_q && c.head_dim_qk == 128; // 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, - static_cast(c.attn_scale), cuda::current_device()); + 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); // 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(); From 4152a201e9f4dcfc8c3cd3f80cd5743fa2b614c5 Mon Sep 17 00:00:00 2001 From: Yang Xu Date: Fri, 16 Oct 2026 14:48:58 -0700 Subject: [PATCH 3/8] Support FROST backward for context-parallel THD AllGather Signed-off-by: Yang Xu --- docs/envvars.rst | 23 ++++- tests/pytorch/attention/test_frost_attn.py | 97 ++++++++++++++++++- .../common/CuTeDSL/frost_attn.py | 27 +++++- .../common/fused_attn/frost_attn.cu | 81 +++++++++++++--- 4 files changed, 202 insertions(+), 26 deletions(-) diff --git a/docs/envvars.rst b/docs/envvars.rst index 74fd370c69e..1790114ee8d 100644 --- a/docs/envvars.rst +++ b/docs/envvars.rst @@ -257,18 +257,37 @@ backend-selection overview. 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. +.. 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. This allows + 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 or gapped THD must leave this option disabled. + 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. diff --git a/tests/pytorch/attention/test_frost_attn.py b/tests/pytorch/attention/test_frost_attn.py index 729de3d0510..5f727017ee0 100644 --- a/tests/pytorch/attention/test_frost_attn.py +++ b/tests/pytorch/attention/test_frost_attn.py @@ -38,9 +38,12 @@ def frost_enabled(monkeypatch): setter(int(os.environ.get("NVTE_ENABLE_CUTEDSL_BACKEND", "0") != "0")) -def _case(dim, thd, unequal, causal, seed=31, bottom_right=False): +def _case(dim, thd, unequal, causal, seed=31, bottom_right=False, kv_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 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): @@ -71,7 +74,7 @@ def draw(s, h): return q, k, v, do, prefixes, reference, max(q_lengths), max(k_lengths) -def _calls(case, dim, thd, causal, bottom_right=False): +def _calls(case, dim, thd, causal, bottom_right=False, origins=None): from transformer_engine.pytorch.cpp_extensions.fused_attn import ( FusedAttnBackend, fused_attn_fwd, @@ -96,8 +99,8 @@ def _calls(case, dim, thd, causal, bottom_right=False): v=v, fake_dtype=torch.bfloat16, fused_attention_backend=FusedAttnBackend.F16_arbitrary_seqlen, - cu_seqlens_q_padded=cu_q, - cu_seqlens_kv_padded=cu_kv, + 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, @@ -218,3 +221,89 @@ def test_frost_d128_bottom_right_fallback(frost_enabled, monkeypatch): 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) diff --git a/transformer_engine/common/CuTeDSL/frost_attn.py b/transformer_engine/common/CuTeDSL/frost_attn.py index 143f0a903ac..0b180395fa4 100644 --- a/transformer_engine/common/CuTeDSL/frost_attn.py +++ b/transformer_engine/common/CuTeDSL/frost_attn.py @@ -21,7 +21,7 @@ _ENGINES = {128: 20609, 256: 20607} -def _declarations(batch, heads, kv_heads, sq, sk, dim, tq, thd): +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: @@ -31,7 +31,8 @@ def _declarations(batch, heads, kv_heads, sq, sk, dim, tq, thd): shape = (batch, heads, sq, 1) stride = (tq * heads, 1, heads, 1) if thd else (heads * sq, sq, 1, 1) else: - shape, stride = (batch, h, s, dim), (s * h * dim, dim, h * dim, 1) + 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 @@ -79,7 +80,22 @@ def _native_metadata(spec, dim, thd, expected_prefix): def get_frost_attention_backward( - key, batch, heads, kv_heads, sq, sk, dim, tq, tk, thd, causal, scale, device, bottom_right=False + 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: @@ -104,6 +120,9 @@ def get_frost_attention_backward( 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: @@ -117,7 +136,7 @@ def get_frost_attention_backward( 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) + declarations = _declarations(batch, heads, kv_heads, sq, sk, dim, tq, thd, q_pitch, kv_pitch) refs = {} def attach(role, tensor): diff --git a/transformer_engine/common/fused_attn/frost_attn.cu b/transformer_engine/common/fused_attn/frost_attn.cu index 5c1621d14e6..6b011cabbf0 100644 --- a/transformer_engine/common/fused_attn/frost_attn.cu +++ b/transformer_engine/common/fused_attn/frost_attn.cu @@ -38,6 +38,31 @@ struct Plan { 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 || @@ -65,9 +90,8 @@ bool eligible(const FusedAttnConfig &c) { 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. - return getenv("NVTE_FUSED_ATTN_FROST_UNIFORM") && - c.num_tokens_q == c.batch_size * c.max_seqlen_q && - c.num_tokens_kv == c.batch_size * c.max_seqlen_kv; + UniformLayout layout; + return uniform_layout(c, &layout); } const Plan *get_plan(const FusedAttnConfig &c, cudaStream_t stream) { @@ -90,6 +114,10 @@ const Plan *get_plan(const FusedAttnConfig &c, cudaStream_t stream) { 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)( @@ -98,7 +126,8 @@ const Plan *get_plan(const FusedAttnConfig &c, cudaStream_t stream) { 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(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(); @@ -146,22 +175,23 @@ bool overlaps(uintptr_t a, size_t a_bytes, uintptr_t b, size_t b_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 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 * sq || - origins_kv[i] != i * sk)) { + 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 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 * sequence + s) * heads + h]; + dense[i] = packed[(b * pitch + origin + s) * heads + h]; } } @@ -185,6 +215,10 @@ bool frost_attention_backward(const FusedAttnConfig &cfg, const FrostPointers &p NVTE_CHECK(cache_cfg.device_id == cuda::current_device(), "cuDNN OSS plan belongs to another device"); 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"); @@ -198,19 +232,29 @@ bool frost_attention_backward(const FusedAttnConfig &cfg, const FrostPointers &p *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, op.bytes, reinterpret_cast(workspace), required), - "cuDNN OSS operand overlaps workspace"); + 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, op.bytes, reinterpret_cast(pointers[other.role]), - other.bytes), + NVTE_CHECK(!overlaps(address, bound_bytes(op), + reinterpret_cast(pointers[other.role]), bound_bytes(other)), "cuDNN OSS gradient aliases an operand"); } } @@ -228,7 +272,7 @@ bool frost_attention_backward(const FusedAttnConfig &cfg, const FrostPointers &p for (const auto &op : plan->operands) { if (op.role >= 6 && op.role <= 8) { NVTE_CHECK(!overlaps(address, metadata_bytes, - reinterpret_cast(pointers[op.role]), op.bytes), + reinterpret_cast(pointers[op.role]), bound_bytes(op)), "cuDNN OSS gradient aliases THD metadata"); } } @@ -236,14 +280,19 @@ bool frost_attention_backward(const FusedAttnConfig &cfg, const FrostPointers &p 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.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); + 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; From cf04d4e11821b3b93b76915a03637bc2b572b460 Mon Sep 17 00:00:00 2001 From: Yang Xu Date: Fri, 16 Oct 2026 18:37:07 -0700 Subject: [PATCH 4/8] Avoid redundant initialization for compact D128 FROST gradients Share the strict compact-layout predicate between PyTorch output allocation and common dispatch, and validate device metadata on every execution. Keep the policy default-off and retain native/fallback initialization. Add FP64, changed-input, poisoned-replay, actual-fill and isolated metadata-guard tests; all 61 directed tests pass on B200. Signed-off-by: Yang Xu --- docs/envvars.rst | 20 +++ tests/pytorch/attention/test_frost_attn.py | 166 ++++++++++++++++++ .../common/fused_attn/frost_attn.cu | 15 ++ .../common/fused_attn/frost_output_policy.h | 32 ++++ .../pytorch/csrc/extensions/attention.cpp | 13 +- 5 files changed, 244 insertions(+), 2 deletions(-) create mode 100644 transformer_engine/common/fused_attn/frost_output_policy.h diff --git a/docs/envvars.rst b/docs/envvars.rst index 1790114ee8d..6b40442b9ff 100644 --- a/docs/envvars.rst +++ b/docs/envvars.rst @@ -292,6 +292,26 @@ backend-selection overview. 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/tests/pytorch/attention/test_frost_attn.py b/tests/pytorch/attention/test_frost_attn.py index 5f727017ee0..9b7aff71c6d 100644 --- a/tests/pytorch/attention/test_frost_attn.py +++ b/tests/pytorch/attention/test_frost_attn.py @@ -5,7 +5,11 @@ """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 @@ -34,6 +38,7 @@ def frost_enabled(monkeypatch): 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")) @@ -307,3 +312,164 @@ def unexpected_python(*args, **kwargs): _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/fused_attn/frost_attn.cu b/transformer_engine/common/fused_attn/frost_attn.cu index 6b011cabbf0..12900fac441 100644 --- a/transformer_engine/common/fused_attn/frost_attn.cu +++ b/transformer_engine/common/fused_attn/frost_attn.cu @@ -19,6 +19,7 @@ #include "../util/cuda_runtime.h" #include "../util/system.h" #include "frost_attn.h" +#include "frost_output_policy.h" namespace transformer_engine { namespace fused_attn { @@ -214,6 +215,10 @@ bool frost_attention_backward(const FusedAttnConfig &cfg, const FrostPointers &p } 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) { @@ -259,6 +264,16 @@ bool frost_attention_backward(const FusedAttnConfig &cfg, const FrostPointers &p } } } + // 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); 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/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); } From a7e91a561021374493671690009566af38a957d6 Mon Sep 17 00:00:00 2001 From: Yang Xu Date: Fri, 16 Oct 2026 19:18:13 -0700 Subject: [PATCH 5/8] Add qualified QA coverage for context parallel FROST attention Signed-off-by: Yang Xu --- AGENTS.md | 9 ++++ docs/envvars.rst | 14 +++++- qa/L0_pytorch_unittest/test.sh | 3 ++ qa/L1_pytorch_frost_unittest/test.sh | 50 ++++++++++++++++++++++ tests/pytorch/attention/test_frost_attn.py | 28 +++++++++++- 5 files changed, 101 insertions(+), 3 deletions(-) create mode 100644 qa/L1_pytorch_frost_unittest/test.sh 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 6b40442b9ff..84248f32660 100644 --- a/docs/envvars.rst +++ b/docs/envvars.rst @@ -235,7 +235,10 @@ backend-selection overview. 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, including rectangular calls. + 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 @@ -257,6 +260,15 @@ backend-selection overview. 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`` 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 index 9b7aff71c6d..77787486866 100644 --- a/tests/pytorch/attention/test_frost_attn.py +++ b/tests/pytorch/attention/test_frost_attn.py @@ -43,9 +43,12 @@ def frost_enabled(monkeypatch): 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): +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) @@ -65,7 +68,10 @@ def draw(s, h): > torch.arange(sq)[:, None] + (sk - sq if bottom_right else 0), -torch.inf, ) - result = scores.softmax(-1) @ vr.repeat_interleave(heads // kv_heads, 0) + # 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)) @@ -216,6 +222,24 @@ def test_frost_d256_bottom_right(frost_enabled, thd): _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 From 020ed08be23a87948db80591c5e089056ea97097 Mon Sep 17 00:00:00 2001 From: Yang Xu Date: Fri, 16 Oct 2026 21:55:48 -0700 Subject: [PATCH 6/8] Add native FROST forward for compact context-parallel attention Signed-off-by: Yang Xu --- docs/envvars.rst | 46 ++- qa/L1_pytorch_frost_unittest/test.sh | 10 +- .../pytorch/attention/test_frost_attn_fwd.py | 290 ++++++++++++++++++ transformer_engine/common/CuTeDSL/__init__.py | 12 + .../common/CuTeDSL/frost_attn_fwd.py | 213 +++++++++++++ .../common/fused_attn/frost_attn.cu | 178 +++++++++++ .../common/fused_attn/frost_attn.h | 6 + .../fused_attn_f16_arbitrary_seqlen.cu | 27 +- 8 files changed, 771 insertions(+), 11 deletions(-) create mode 100644 tests/pytorch/attention/test_frost_attn_fwd.py create mode 100644 transformer_engine/common/CuTeDSL/frost_attn_fwd.py diff --git a/docs/envvars.rst b/docs/envvars.rst index 84248f32660..e9665ca2872 100644 --- a/docs/envvars.rst +++ b/docs/envvars.rst @@ -225,7 +225,8 @@ backend-selection overview. 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. + use it to verify benchmark routing. Forward remains cuDNN + unless separately enabled with ``NVTE_FUSED_ATTN_FROST_FWD``. Requires a TE build with ``NVTE_WITH_CUTEDSL=1`` (CMake: ``-DNVTE_WITH_CUTEDSL=ON``), ``NVTE_ENABLE_CUTEDSL_BACKEND=1``, @@ -269,6 +270,49 @@ backend-selection overview. ``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_FWD + + :Default: ``0`` + :Values: ``0``, ``1``, ``2`` + :Description: Experimental cuDNN OSS forward provider through TE's common + C++ entry. ``1`` attempts the provider and retains the backend + for unsupported declarations; ``2`` requires the provider. + Requires the same build/DSL setup as the backward provider, + plus ``NVTE_FUSED_ATTN_FROST_FWD_COMPACT=1``. + + Initial scope: SM100, BF16 training, separate compact THD Q/K/V and output, + D128, top-left causal padding mask, no bias, dropout, paging, sliding window, + or maximum-logit output. Stats retain TE's packed FP32 natural-log format. + Forward and backward routing are independent; enabling this option does not + enable the OSS backward provider. + + Cold preparation uses the ordinary FE graph contract and retains the complete + compiled host entry, including metadata/remap and descriptor construction. + Hot execution and warmed CUDA graph capture enter native C++/TVM-FFI without + Python callbacks. Current operands, workspace and stream are rebound each call; + cached tuple arguments contain immutable declaration geometry only. + + This adapter currently targets FE ``ThdLaunchSpec`` at development revision + ``2bcdd1bb487a65771d971e6d7380e956cea9ec22``. FE PR 1696 optimizes its THD work + lookup but does not change the ABI. The export is experimental: an unfamiliar + host signature or optional feature reports an error rather than dropping an + argument. Warm up each declaration and device before capture. + +.. envvar:: NVTE_FUSED_ATTN_FROST_FWD_COMPACT + + :Default: ``0`` + :Values: ``0``, ``1`` + :Description: Promise that THD physical origins equal the zero-based cumulative + lengths. Lengths may be unequal and may change during replay; + this option does not convert THD to a dense declaration. + + Every execution validates the promise on device, including monotonic lengths, + declared bounds and physical origins. Live queries require a nonempty KV + sequence. Invalid promised metadata traps the CUDA context, so this flag must + not be applied to storage with physical holes. Such layouts, including pitched + AllGather subcalls, must retain the existing forward route. Passing A2A tests + does not establish AllGather support or distributed performance. + .. envvar:: NVTE_FUSED_ATTN_FROST_ALLGATHER :Default: ``0`` diff --git a/qa/L1_pytorch_frost_unittest/test.sh b/qa/L1_pytorch_frost_unittest/test.sh index 2aa9b585df7..9bb936ffef1 100644 --- a/qa/L1_pytorch_frost_unittest/test.sh +++ b/qa/L1_pytorch_frost_unittest/test.sh @@ -32,7 +32,8 @@ 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" + "$TE_PATH/tests/pytorch/attention/test_frost_attn.py" \ + "$TE_PATH/tests/pytorch/attention/test_frost_attn_fwd.py" # Reject a green all-skipped run or selection that omitted native-route guards. python3 - "$XML_LOG_DIR/pytest_test_frost_attn.xml" <<'RESULTS' @@ -40,11 +41,14 @@ 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 len(cases) >= 79, "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") +for kind in ("q_origin", "kv_origin", "q_limit", "kv_limit", "q_decrease", "q_base", "graph_origin"): + assert f"test_frost_forward_metadata_guard[{kind}]" in names, "Missing isolated forward guard: " + kind +assert "test_frost_forward_compact_decline" in names +print("FROST QA passed:", len(cases), "tests, zero skips, all twelve isolated guards") RESULTS diff --git a/tests/pytorch/attention/test_frost_attn_fwd.py b/tests/pytorch/attention/test_frost_attn_fwd.py new file mode 100644 index 00000000000..57376d3eb82 --- /dev/null +++ b/tests/pytorch/attention/test_frost_attn_fwd.py @@ -0,0 +1,290 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""Actual TE common-library FROST forward, including training Stats and replay.""" + +import ctypes +import json +import os +from pathlib import Path +import subprocess +import sys + +import pytest +import torch + +from test_frost_attn import frost_enabled # noqa: F401 + + +@pytest.fixture +def forward_enabled(frost_enabled, monkeypatch): + del frost_enabled + monkeypatch.setenv("NVTE_FUSED_ATTN_FROST_FWD", "2") + monkeypatch.setenv("NVTE_FUSED_ATTN_FROST_FWD_COMPACT", "1") + + +def _case(heads=4, kv_heads=1, seed=51): + gen = torch.Generator(device="cpu").manual_seed(seed) + q, k, v, do = [ + (torch.randn(512, h, 128, generator=gen) * 0.5).bfloat16().cuda() + for h in (heads, kv_heads, kv_heads, heads) + ] + cuq = torch.tensor([0, 129, 192], device="cuda", dtype=torch.int32) + cuk = torch.tensor([0, 193, 290], device="cuda", dtype=torch.int32) + return q, k, v, do, cuq, cuk, cuq.clone(), cuk.clone() + + +def _forward(case): + from transformer_engine.pytorch.cpp_extensions.fused_attn import ( + FusedAttnBackend, + fused_attn_fwd, + ) + + q, k, v, _, cuq, cuk, oq, ok = case + return fused_attn_fwd( + is_training=True, + max_seqlen_q=256, + max_seqlen_kv=256, + cu_seqlens_q=cuq, + cu_seqlens_kv=cuk, + q=q, + k=k, + v=v, + fake_dtype=torch.bfloat16, + fused_attention_backend=FusedAttnBackend.F16_arbitrary_seqlen, + cu_seqlens_q_padded=oq, + cu_seqlens_kv_padded=ok, + attn_scale=128**-0.5, + qkv_layout="thd_thd_thd", + o_format="thd", + attn_mask_type="padding_causal", + dropout=0.0, + ) + + +def _reference(case): + q, k, v, do, cuq, cuk, _, _ = [x.cpu() for x in case] + outputs, stats, gradients = [], [], [] + for i in range(2): + qs = slice(int(cuq[i]), int(cuq[i + 1])) + ks = slice(int(cuk[i]), int(cuk[i + 1])) + qr, kr, vr = [x.double().transpose(0, 1).requires_grad_() for x in (q[qs], k[ks], v[ks])] + repeat = q.shape[1] // k.shape[1] + score = qr @ kr.repeat_interleave(repeat, 0).transpose(-1, -2) * 128**-0.5 + mask = torch.arange(kr.shape[1])[None, :] > torch.arange(qr.shape[1])[:, None] + score = score.masked_fill(mask, -torch.inf) + o = score.softmax(-1) @ vr.repeat_interleave(repeat, 0) + grads = torch.autograd.grad(o, (qr, kr, vr), do[qs].transpose(0, 1).double()) + outputs.append(o.transpose(0, 1).detach()) + stats.append(score.logsumexp(-1).transpose(0, 1).detach()) + gradients.append(tuple(g.transpose(0, 1).detach() for g in grads)) + return ( + torch.cat(outputs), + torch.cat(stats), + tuple(torch.cat([g[i] for g in gradients]) for i in range(3)), + ) + + +def _check_forward(result, expected): + o, aux = result + ref, lse, _ = expected + torch.testing.assert_close(o[: len(ref)].double().cpu(), ref, atol=0.015, rtol=0.015) + torch.testing.assert_close(aux[0][: len(lse), :, 0].double().cpu(), lse, atol=1e-4, rtol=1e-4) + + +@pytest.mark.parametrize("heads,kv_heads", [(4, 1), (8, 2), (4, 4)]) +@pytest.mark.parametrize("seed", [51, 73]) +def test_frost_forward_training(forward_enabled, heads, kv_heads, seed): + """Check O, natural-log packed Stats and gradients against CPU FP64.""" + del forward_enabled + from transformer_engine.pytorch.cpp_extensions.fused_attn import ( + FusedAttnBackend, + fused_attn_bwd, + ) + + case = _case(heads, kv_heads, seed) + expected = _reference(case) + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(stream): + result = _forward(case) + torch.cuda.current_stream().wait_stream(stream) + _check_forward(result, expected) + q, k, v, do, cuq, cuk, oq, ok = case + o, aux = result + actual = fused_attn_bwd( + max_seqlen_q=256, + max_seqlen_kv=256, + cu_seqlens_q=cuq, + cu_seqlens_kv=cuk, + q=q, + k=k, + v=v, + o=o, + d_o=do, + fake_dtype=torch.bfloat16, + fused_attention_backend=FusedAttnBackend.F16_arbitrary_seqlen, + aux_ctx_tensors=aux, + cu_seqlens_q_padded=oq, + cu_seqlens_kv_padded=ok, + attn_scale=128**-0.5, + qkv_layout="thd_thd_thd", + dqkv_layout="thd_thd_thd", + o_format="thd", + do_format="thd", + attn_mask_type="padding_causal", + dropout=0.0, + deterministic=False, + )[:3] + for got, want in zip(actual, expected[2]): + torch.testing.assert_close(got[: len(want)].double().cpu(), want, atol=0.025, rtol=0.025) + relative = ( + got[: len(want)].double().cpu() - want + ).square().mean().sqrt() / want.square().mean().sqrt().clamp_min(1e-8) + assert relative < 0.02 + + +@pytest.mark.parametrize("heads,kv_heads", [(4, 1), (8, 2), (4, 4)]) +def test_frost_forward_native_replay(forward_enabled, heads, kv_heads): + """No Python hot dispatch; graph reads changed lengths, origins and inputs.""" + del forward_enabled + import tvm_ffi + + case = _case(heads, kv_heads) + expected = _reference(case) + _check_forward(_forward(case), expected) + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(stream): + _forward(case) + torch.cuda.current_stream().wait_stream(stream) + names = ("get_frost_attention_forward", "get_frost_attention_forward_constant") + originals = {name: tvm_ffi.get_global_func(name) for name in names} + + def unexpected_python(*args): + raise AssertionError("Steady-state forward entered Python preparation") + + try: + for name in names: + tvm_ffi.register_global_func(name, unexpected_python, override=True) + # A second allocation uses the same declaration with new pointers. + other = _case(heads, kv_heads, 97) + _check_forward(_forward(other), _reference(other)) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph, stream=stream): + result = _forward(case) + for q_lengths, k_lengths in ( + ((128, 128), (128, 256)), + ((0, 0), (0, 0)), + ((63, 129), (97, 193)), + ): + for cu, origin, lengths in ( + (case[4], case[6], q_lengths), + (case[5], case[7], k_lengths), + ): + cu.copy_( + torch.tensor([0, lengths[0], sum(lengths)], dtype=torch.int32, device="cuda") + ) + origin.copy_(cu) + for tensor in case[:3]: + tensor.neg_() + result[0].fill_(float("nan")) + result[1][0].fill_(float("nan")) + graph.replay() + if sum(q_lengths): + _check_forward(result, _reference(case)) + else: + torch.cuda.synchronize() + finally: + for name, original in originals.items(): + tvm_ffi.register_global_func(name, original, override=True) + + +def test_frost_forward_compact_decline(forward_enabled, monkeypatch): + """The compact-layout promise is required even after a plan has been cached.""" + del forward_enabled + case = _case() + expected = _reference(case) + _check_forward(_forward(case), expected) + monkeypatch.setenv("NVTE_FUSED_ATTN_FROST_FWD_COMPACT", "0") + monkeypatch.setenv("NVTE_FUSED_ATTN_FROST_FWD", "1") + _check_forward(_forward(case), expected) + monkeypatch.setenv("NVTE_FUSED_ATTN_FROST_FWD", "2") + with pytest.raises(RuntimeError, match="Requested cuDNN OSS forward unavailable"): + _forward(case) + + +def _invalid_metadata(kind): + 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, core.nvte_set_cutedsl_backend.restype = [ + ctypes.c_int + ], None + core.nvte_set_cutedsl_backend(1) + os.environ["NVTE_FUSED_ATTN_FROST_FWD"] = "2" + os.environ["NVTE_FUSED_ATTN_FROST_FWD_COMPACT"] = "1" + case = _case() + _check_forward(_forward(case), _reference(case)) + graph = None + if kind == "graph_origin": + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(stream): + _forward(case) + torch.cuda.current_stream().wait_stream(stream) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph, stream=stream): + _forward(case) + case[7][-1] += 1 + elif kind in ("q_origin", "kv_origin"): + case[6 if kind == "q_origin" else 7][1] += 1 + else: + cu, origin = ( + (case[4], case[6]) + if kind in ("q_limit", "q_decrease", "q_base") + else (case[5], case[7]) + ) + if kind == "q_decrease": + cu[-1] = 100 + elif kind == "q_base": + cu[0] = 1 + else: + cu[-1] = 513 + origin.copy_(cu) + print("EXPECT_FORWARD_GUARD_TRAP", kind, flush=True) + if graph is None: + _forward(case) + else: + graph.replay() + torch.cuda.synchronize() + raise AssertionError("Invalid forward metadata escaped the device guard") + + +@pytest.mark.parametrize( + "kind", ["q_origin", "kv_origin", "q_limit", "kv_limit", "q_decrease", "q_base", "graph_origin"] +) +def test_frost_forward_metadata_guard(forward_enabled, kind): + del forward_enabled + command = [sys.executable] + (["-S"] if sys.flags.no_site else []) + code = ( + "import json, runpy, sys; sys.path[:] = json.loads(sys.argv[1]);" + " runpy.run_path(sys.argv[2])['_invalid_metadata'](sys.argv[3])" + ) + result = subprocess.run( + command + ["-c", code, json.dumps(sys.path), str(Path(__file__).resolve()), kind], + text=True, + capture_output=True, + timeout=240, + check=False, + ) + output = result.stdout + result.stderr + assert result.returncode != 0, output + assert "EXPECT_FORWARD_GUARD_TRAP " + kind in output, output + assert "NVTE_FUSED_ATTN_FROST_FWD_COMPACT: invalid sequence/offset" in output, output diff --git a/transformer_engine/common/CuTeDSL/__init__.py b/transformer_engine/common/CuTeDSL/__init__.py index 3745001d7db..6bdb28bd5e9 100644 --- a/transformer_engine/common/CuTeDSL/__init__.py +++ b/transformer_engine/common/CuTeDSL/__init__.py @@ -28,3 +28,15 @@ def register_cutedsl_backends(): tvm_ffi.register_global_func( "get_frost_attention_backward", get_frost_attention_backward, override=True ) + + from transformer_engine.common.CuTeDSL.frost_attn_fwd import ( + get_frost_attention_forward, + get_frost_attention_forward_constant, + ) + + tvm_ffi.register_global_func( + "get_frost_attention_forward", get_frost_attention_forward, override=True + ) + tvm_ffi.register_global_func( + "get_frost_attention_forward_constant", get_frost_attention_forward_constant, override=True + ) diff --git a/transformer_engine/common/CuTeDSL/frost_attn_fwd.py b/transformer_engine/common/CuTeDSL/frost_attn_fwd.py new file mode 100644 index 00000000000..5fc60acc20f --- /dev/null +++ b/transformer_engine/common/CuTeDSL/frost_attn_fwd.py @@ -0,0 +1,213 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""Cold preparation for TE's native cuDNN OSS THD forward dispatcher. + +The complete compiled host entry launches metadata/remap, descriptors and SDPA. +Only immutable declaration arguments cross Python during preparation; C++ binds +the current pointers, workspace and stream on every execution. +""" + +from importlib.metadata import PackageNotFoundError, version +import logging +import math + +_LOG = logging.getLogger("transformer_engine.frost") +_OWNERS = {} +_ORDER = tuple( + "q_ptr k_ptr v_ptr o_ptr lse_ptr sinks_ptr meta_ptr o_desc_ptr problem_size " + "q_strides k_strides v_strides o_strides lse_strides lse_ext scale_softmax_log2 " + "n_thd_units seq_q_lens_addr thd_q_lens_ptr thd_kv_lens_ptr thd_lens_form " + "o_partial_ptr block_table_ptr block_table_v_ptr table_strides n_pages stream".split() +) +_POINTERS = dict(q_ptr=0, k_ptr=1, v_ptr=2, o_ptr=3, lse_ptr=4, thd_q_lens_ptr=5, thd_kv_lens_ptr=6) + + +def get_frost_attention_forward_constant(key, index): + """Return one retained declaration value to the C++ cache, only at plan time.""" + return _OWNERS[key][1][index] + + +def _export(graph, batch, heads, kv_heads, sq, sk, tq, tk, scale, device): + import tvm_ffi + + spec = graph._compiled_plans[graph._plan_index]._prepared.spec + if ( + tuple(spec.order) != _ORDER + or not isinstance(spec.fn, tvm_ffi.Function) + or spec.quant is not None + or spec.paged + or spec.has_sink + or spec.gate_expect is not None + or not spec.has_lse + or spec.lse_padded + or spec.lse_head_major + or spec.lse_stride_override + or spec.lse_fill_plan is not None + or spec.split_workspace is not None + or spec.lens_form != 3 + or spec.n_q_lens != batch + 1 + or spec.n_kv_lens != batch + 1 + or spec.device_index != device + or (spec.b, spec.qh, spec.kh, spec.d_qk, spec.d_v) != (batch, heads, kv_heads, 128, 128) + or spec.total_q != tq + or spec.total_kv != tk + or spec.s_q_max != sq + ): + raise RuntimeError("Unsupported cuDNN OSS prepared forward contract") + args = dict(zip(spec.order, spec.template)) + for role, h in (("q", heads), ("k", kv_heads), ("v", kv_heads), ("o", heads)): + if spec.expect[role] not in ("bf16", "bfloat16"): + raise RuntimeError("cuDNN OSS forward requires BF16 operands") + if tuple(args[role + "_strides"]) != (h * 128, h * 128, 128): + raise RuntimeError("cuDNN OSS forward strides differ from the TE declaration") + if ( + args["lse_strides"] != (0, 0, 0) + or args["lse_ext"] != 0 + or args["seq_q_lens_addr"] != 0 + or args["thd_lens_form"] != 3 + or args["table_strides"] != (0, 0) + or args["n_pages"] != 0 + or not math.isclose(args["scale_softmax_log2"], scale * math.log2(math.e), rel_tol=1e-12) + or any( + args[name] is not None + for name in ("sinks_ptr", "o_partial_ptr", "block_table_ptr", "block_table_v_ptr") + ) + ): + raise RuntimeError("Unsupported cuDNN OSS forward constant arguments") + args["problem_size"] = (batch, heads, kv_heads, tq, tk, 0) + args["sinks_ptr"] = 0 # Required pointer slot; FE binds absent sinks as address zero. + # Match FE's native binder: capacity is fixed by the TE declaration; lengths + # remain device data and may change on a retained graph. + args["n_thd_units"] = min(args["n_thd_units"], ((tq - 1) // spec.cga_tile_m + batch) * heads) + workspace = int(graph.get_workspace_size()) + alignment = int(spec.workspace_alignment) + if not (0 <= spec.off_o_desc < spec.scratch_bytes <= workspace < 2**63): + raise RuntimeError("Invalid cuDNN OSS forward workspace contract") + if alignment < 16 or alignment & (alignment - 1): + raise RuntimeError("Invalid cuDNN OSS forward workspace alignment") + spans = ( + tq * heads * 128 * 2, + tk * kv_heads * 128 * 2, + tk * kv_heads * 128 * 2, + tq * heads * 128 * 2, + tq * heads * 4, + *((batch + 1) * 4 for _ in range(4)), + ) + if any(n <= 0 or n >= 2**63 for n in spans): + raise RuntimeError("cuDNN OSS forward operand span is not representable") + metadata = [1, workspace, alignment, len(_ORDER)] + for i, span in enumerate(spans): + metadata.extend((span, 16 if i < 4 else 4)) + constants = [] + for i, name in enumerate(_ORDER): + # kind 0 = immutable value; 1 = current operand; 2 = workspace offset; + # kind 3 = current stream. The C++ bridge never constructs FFI arrays. + kind, value = 0, i + if name in _POINTERS: + kind, value = 1, _POINTERS[name] + elif name in ("meta_ptr", "o_desc_ptr"): + kind, value = 2, 0 if name == "meta_ptr" else int(spec.off_o_desc) + elif name == "stream": + kind, value = 3, 0 + metadata.extend((kind, value)) + constants.append(tvm_ffi.convert(args[name]) if kind == 0 else None) + return spec.fn, tuple(constants), metadata + + +def get_frost_attention_forward(key, batch, heads, kv_heads, sq, sk, tq, tk, scale, device): + """Compile a compact BF16 THD training declaration without allocating data.""" + try: + import cudnn + import tvm_ffi + from cuda.bindings import runtime + from packaging.version import Version + + if Version(version("nvidia-cutlass-dsl")) < Version("4.8.0"): + raise ImportError("cuDNN OSS forward requires CuTeDSL >= 4.8.0") + from cutlass.base_dsl.version_info import CUDA_VERSION + + if CUDA_VERSION.major < 13: + raise ImportError("cuDNN OSS forward requires nvidia-cutlass-dsl[cu13]") + except (ImportError, PackageNotFoundError) as error: + _LOG.warning("cuDNN OSS forward dependency unavailable: %s", error) + 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") + io = cudnn.data_type.BFLOAT16 + graph = cudnn.pygraph( + io_data_type=io, + intermediate_data_type=cudnn.data_type.FLOAT, + compute_data_type=cudnn.data_type.FLOAT, + ) + + def attach(tensor, name, factor): + origin = graph.tensor( + name=name + "_origin", + dim=[batch + 1, 1, 1, 1], + stride=[1, 1, 1, 1], + data_type=cudnn.data_type.INT32, + ) + tensor.set_ragged_offset(origin).set_ragged_offset_multiplier(factor) + + refs = {} + for role, h, s in (("q", heads, sq), ("k", kv_heads, sk), ("v", kv_heads, sk)): + refs[role] = graph.tensor( + name=role, dim=[batch, h, s, 128], stride=[s * h * 128, 128, h * 128, 1], data_type=io + ) + attach(refs[role], role, h * 128) + lengths = [ + graph.tensor( + name="cu_" + role, + dim=[batch + 1, 1, 1, 1], + stride=[1, 1, 1, 1], + data_type=cudnn.data_type.INT32, + ) + for role in ("q", "kv") + ] + o, stats = graph.sdpa( + **refs, + generate_stats=True, + attn_scale=scale, + use_padding_mask=True, + use_causal_mask=True, + cu_seq_len_q=lengths[0], + cu_seq_len_kv=lengths[1], + max_total_seq_len_q=tq, + max_total_seq_len_kv=tk + ) + o.set_output(True).set_dim([batch, heads, sq, 128]).set_stride( + [sq * heads * 128, 128, heads * 128, 1] + ) + stats.set_output(True).set_data_type(cudnn.data_type.FLOAT).set_dim( + [batch, heads, sq, 1] + ).set_stride([sq * heads, 1, heads, 1]) + attach(o, "o", heads * 128) + attach(stats, "stats", heads) + try: + graph.validate() + graph.build_operation_graph() + graph.create_execution_plans([cudnn.heur_mode.A]) + index = next( + ( + i + for i in range(len(graph.plans)) + if graph.get_engine_and_knobs_at_index(i)[0] == 20511 + ), + None, + ) + if index is None: + return False + graph.select_plan(index) + graph.build_plans() + except (NotImplementedError, cudnn.cudnnGraphNotSupportedError) as error: + _LOG.warning("cuDNN OSS forward declined this declaration: %s", error) + return False + native, constants, metadata = _export( + graph, batch, heads, kv_heads, sq, sk, tq, tk, scale, device + ) + _OWNERS[key] = (graph, constants) + 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 index 12900fac441..cfaf42282b8 100644 --- a/transformer_engine/common/fused_attn/frost_attn.cu +++ b/transformer_engine/common/fused_attn/frost_attn.cu @@ -324,5 +324,183 @@ bool frost_attention_backward(const FusedAttnConfig &cfg, const FrostPointers &p return true; } +namespace { + +struct ForwardArgument { + int kind; + int64_t value; + tvm::ffi::Any constant; +}; + +struct ForwardPlan { + tvm::ffi::Function function; + size_t workspace_bytes; + size_t workspace_alignment; + std::array operands; + std::vector arguments; +}; + +bool forward_eligible(const FusedAttnConfig &c) { + // The current FE forward host derives compact physical positions from lengths. + // Separate padded origins therefore require an explicit, device-checked promise. + return getenv("NVTE_FUSED_ATTN_FROST_FWD_COMPACT") && + cuda::sm_arch(cuda::current_device()) == 100 && c.is_training && !c.return_max_logit && + c.qkv_dtype == kNVTEBFloat16 && c.o_dtype == kNVTEBFloat16 && c.head_dim_qk == 128 && + c.head_dim_v == 128 && c.bias_type == NVTE_NO_BIAS && + c.softmax_type == NVTE_VANILLA_SOFTMAX && c.dropout == 0 && !c.is_paged_kv && + c.qkv_layout == NVTE_THD_THD_THD && c.o_format == NVTE_THD && c.is_padding && + c.is_ragged_q && c.uses_ragged_stats && c.is_causal && !c.is_causal_bottom_right && + c.window_size_left == -1 && (c.window_size_right == -1 || c.window_size_right == 0) && + 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 <= 65535 && + c.num_attn_heads >= c.num_gqa_groups && c.num_attn_heads % c.num_gqa_groups == 0 && + c.max_seqlen_q > 0 && c.max_seqlen_kv > 0 && c.max_seqlen_q <= INT32_MAX - 256 && + c.max_seqlen_kv <= INT32_MAX - 256 && c.num_tokens_q > 0 && c.num_tokens_kv > 0 && + c.num_tokens_q <= INT32_MAX && c.num_tokens_kv <= INT32_MAX; +} + +const ForwardPlan *get_forward_plan(const FusedAttnConfig &c, cudaStream_t stream) { + static auto *cache = new std::map>; + static auto *mutex = new std::mutex; + static std::atomic serial{0}; + { + std::lock_guard lock(*mutex); + auto it = cache->find(c); + if (it != cache->end()) return it->second.get(); + } + cudaStreamCaptureStatus capture; + NVTE_CHECK_CUDA(cudaStreamIsCapturing(stream, &capture)); + NVTE_CHECK(capture == cudaStreamCaptureStatusNone, + "Warm up cuDNN OSS forward 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_forward"); + auto constant = tvm::ffi::Function::GetGlobal("get_frost_attention_forward_constant"); + if (!prepare || !constant) return nullptr; + const std::string key = "transformer_engine.frost.forward." + std::to_string(serial++); + // Python is entered only on a cache miss, without holding a C++ cache 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.num_tokens_q), static_cast(c.num_tokens_kv), + static_cast(c.attn_scale), c.device_id); + auto encoded = result.try_cast(); + auto declined = result.try_cast(); + NVTE_CHECK(encoded || (declined && !*declined), "Invalid cuDNN OSS forward preparation response"); + std::unique_ptr plan; + if (encoded) { + NVTE_CHECK(encoded->size() < 4096, "Oversized cuDNN OSS forward metadata"); + std::istringstream input{std::string(*encoded)}; + std::vector v; + int64_t value; + while (input >> value) v.push_back(value); + NVTE_CHECK(input.eof() && v.size() >= 22 && v[0] == 1 && v[1] > 0 && v[2] >= 16 && + (v[2] & (v[2] - 1)) == 0 && v[3] == 27 && v.size() == 22 + 2 * v[3], + "Malformed cuDNN OSS forward metadata"); + auto function = tvm::ffi::Function::GetGlobal(key); + NVTE_CHECK(function.has_value(), "cuDNN OSS did not register its native forward entry"); + plan = std::make_unique( + ForwardPlan{*function, static_cast(v[1]), static_cast(v[2]), {}, {}}); + size_t cursor = 4; + for (int role = 0; role < 9; ++role) { + const auto bytes = v[cursor++], alignment = v[cursor++]; + NVTE_CHECK(bytes > 0 && alignment > 0 && (alignment & (alignment - 1)) == 0, + "Invalid cuDNN OSS forward operand metadata"); + plan->operands[role] = {role, static_cast(bytes), static_cast(alignment)}; + } + for (int64_t i = 0; i < v[3]; ++i) { + const auto kind = v[cursor++], index = v[cursor++]; + NVTE_CHECK((kind == 0 && index == i) || (kind == 1 && index >= 0 && index < 7) || + (kind == 2 && index >= 0 && index < v[1]) || (kind == 3 && index == 0), + "Invalid cuDNN OSS forward argument metadata"); + tvm::ffi::Any immutable; + if (kind == 0) immutable = (*constant)(tvm::ffi::String(key), index); + plan->arguments.push_back({static_cast(kind), index, std::move(immutable)}); + } + } + std::lock_guard lock(*mutex); + return cache->emplace(c, std::move(plan)).first->second.get(); +} + +__global__ void validate_compact_forward(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 tq, + int64_t tk) { + const int64_t i = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + if (i > batch) return; + const int64_t q = cu_q[i], k = cu_kv[i]; + bool valid = q >= 0 && q <= tq && k >= 0 && k <= tk && origins_q[i] == q && origins_kv[i] == k; + if (i == 0) valid = valid && q == 0 && k == 0; + if (i < batch) { + const int64_t lq = static_cast(cu_q[i + 1]) - q; + const int64_t lk = static_cast(cu_kv[i + 1]) - k; + valid = valid && lq >= 0 && lq <= sq && lk >= 0 && lk <= sk && (lq == 0 || lk > 0); + } + if (!valid) { + printf("NVTE_FUSED_ATTN_FROST_FWD_COMPACT: invalid sequence/offset at %d\n", + static_cast(i)); + asm volatile("trap;"); + } +} + +} // namespace + +bool frost_attention_forward(const FusedAttnConfig &cfg, const FrostForwardPointers &pointers, + void *workspace, size_t *workspace_size, cudaStream_t stream) { + const int mode = getenv("NVTE_FUSED_ATTN_FROST_FWD"); + if (mode == 0) return false; + NVTE_CHECK(mode == 1 || mode == 2, "NVTE_FUSED_ATTN_FROST_FWD must be 0, 1 or 2"); + cfg.check_derived(); + FusedAttnConfig cached = cfg; + cached.device_id = cuda::current_device(); + const ForwardPlan *plan = forward_eligible(cfg) ? get_forward_plan(cached, stream) : nullptr; + if (!plan) { + NVTE_CHECK(mode != 2, + "Requested cuDNN OSS forward unavailable; check compact THD D128 causal scope, " + "NVTE_FUSED_ATTN_FROST_FWD_COMPACT=1 and compatible FE/CuTeDSL installation"); + return false; + } + if (workspace == nullptr) { + *workspace_size = plan->workspace_bytes; + return true; + } + NVTE_CHECK(*workspace_size >= plan->workspace_bytes, "cuDNN OSS forward workspace is too small"); + *workspace_size = plan->workspace_bytes; + const auto address = reinterpret_cast(workspace); + NVTE_CHECK(address % plan->workspace_alignment == 0, "cuDNN OSS forward workspace misaligned"); + for (const auto &op : plan->operands) { + const auto ptr = reinterpret_cast(pointers[op.role]); + NVTE_CHECK(ptr != 0 && ptr % op.alignment == 0, + "cuDNN OSS forward operand is null or misaligned, role=", op.role); + NVTE_CHECK(!overlaps(ptr, op.bytes, address, plan->workspace_bytes), + "cuDNN OSS forward operand overlaps workspace"); + if (op.role == 3 || op.role == 4) { + for (const auto &other : plan->operands) { + if (other.role == op.role) continue; + NVTE_CHECK(!overlaps(ptr, op.bytes, reinterpret_cast(pointers[other.role]), + other.bytes), + "cuDNN OSS forward output aliases an operand"); + } + } + } + validate_compact_forward<<<(cfg.batch_size + 256) / 256, 256, 0, stream>>>( + static_cast(pointers[5]), static_cast(pointers[6]), + static_cast(pointers[7]), static_cast(pointers[8]), + cfg.batch_size, cfg.max_seqlen_q, cfg.max_seqlen_kv, cfg.num_tokens_q, cfg.num_tokens_kv); + NVTE_CHECK_CUDA(cudaGetLastError()); + std::array frame; + for (size_t i = 0; i < plan->arguments.size(); ++i) { + const auto &arg = plan->arguments[i]; + if (arg.kind == 0) frame[i] = static_cast(arg.constant); + if (arg.kind == 1) frame[i] = reinterpret_cast(pointers[arg.value]); + if (arg.kind == 2) frame[i] = static_cast(address + arg.value); + if (arg.kind == 3) frame[i] = reinterpret_cast(stream); + } + tvm::ffi::Any result; + plan->function.CallPacked(frame.data(), static_cast(frame.size()), &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 index 360e378d7e0..3fc58b7b4ea 100644 --- a/transformer_engine/common/fused_attn/frost_attn.h +++ b/transformer_engine/common/fused_attn/frost_attn.h @@ -20,6 +20,12 @@ 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; +// q,k,v,o,stats,cu_q,cu_kv,physical_origin_q,physical_origin_kv. +using FrostForwardPointers = std::array; + +bool frost_attention_forward(const FusedAttnConfig &cfg, const FrostForwardPointers &pointers, + void *workspace, size_t *workspace_size, cudaStream_t stream); + // 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, 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 5794a30ad9a..f26e88825d8 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 @@ -1087,13 +1087,26 @@ void fused_attn_arbitrary_seqlen_fwd(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_fwd_impl( - cfg, devPtrQ, devPtrK, devPtrV, devPtrBias, devPtrSoftmaxOffset, devPtrS1, devPtrS2, devPtrO, - devPtrDropoutSeed, devPtrDropoutOffset, devPtrCuSeqlensQ, devPtrCuSeqlensKV, devPtrPageTableK, - devPtrPageTableV, 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::FrostForwardPointers frost_pointers{ + devPtrQ, devPtrK, devPtrV, devPtrO, devPtrS1, + devPtrCuSeqlensQ, devPtrCuSeqlensKV, devPtrSeqOffsetsQ, devPtrSeqOffsetsKV}; + used_frost = fused_attn::frost_attention_forward(cfg, frost_pointers, workspace->data.dptr, + &workspace_size, stream); +#else + NVTE_CHECK(transformer_engine::getenv("NVTE_FUSED_ATTN_FROST_FWD") != 2, + "cuDNN OSS forward requires a TE build with NVTE_WITH_CUTEDSL=ON"); +#endif + if (!used_frost) { + workspace_size = 0; + fused_attn_arbitrary_seqlen_fwd_impl( + cfg, devPtrQ, devPtrK, devPtrV, devPtrBias, devPtrSoftmaxOffset, devPtrS1, devPtrS2, + devPtrO, devPtrDropoutSeed, devPtrDropoutOffset, devPtrCuSeqlensQ, devPtrCuSeqlensKV, + devPtrPageTableK, devPtrPageTableV, devPtrSeqOffsetsQ, devPtrSeqOffsetsKV, + workspace->data.dptr, &workspace_size, stream, handle); + } if (workspace_size > 0) { if (workspace->data.dptr == nullptr) { From 11aab95c9e7a120349a9f04477a40dc13142818d Mon Sep 17 00:00:00 2001 From: Yang Xu Date: Fri, 16 Oct 2026 22:30:46 -0700 Subject: [PATCH 7/8] Fix lint in native OSS forward preparation Signed-off-by: Yang Xu --- .../common/CuTeDSL/frost_attn_fwd.py | 18 ++++++++++++++---- 1 file changed, 14 insertions(+), 4 deletions(-) diff --git a/transformer_engine/common/CuTeDSL/frost_attn_fwd.py b/transformer_engine/common/CuTeDSL/frost_attn_fwd.py index 5fc60acc20f..76c5e9abe08 100644 --- a/transformer_engine/common/CuTeDSL/frost_attn_fwd.py +++ b/transformer_engine/common/CuTeDSL/frost_attn_fwd.py @@ -21,7 +21,15 @@ "n_thd_units seq_q_lens_addr thd_q_lens_ptr thd_kv_lens_ptr thd_lens_form " "o_partial_ptr block_table_ptr block_table_v_ptr table_strides n_pages stream".split() ) -_POINTERS = dict(q_ptr=0, k_ptr=1, v_ptr=2, o_ptr=3, lse_ptr=4, thd_q_lens_ptr=5, thd_kv_lens_ptr=6) +_POINTERS = { + "q_ptr": 0, + "k_ptr": 1, + "v_ptr": 2, + "o_ptr": 3, + "lse_ptr": 4, + "thd_q_lens_ptr": 5, + "thd_kv_lens_ptr": 6, +} def get_frost_attention_forward_constant(key, index): @@ -29,7 +37,9 @@ def get_frost_attention_forward_constant(key, index): return _OWNERS[key][1][index] -def _export(graph, batch, heads, kv_heads, sq, sk, tq, tk, scale, device): +def _export(graph, batch, heads, kv_heads, sq, _sk, tq, tk, scale, device): + # Keep the fail-closed prepared-contract checks explicit and short-circuiting. + # pylint: disable=too-many-boolean-expressions import tvm_ffi spec = graph._compiled_plans[graph._plan_index]._prepared.spec @@ -83,7 +93,7 @@ def _export(graph, batch, heads, kv_heads, sq, sk, tq, tk, scale, device): args["n_thd_units"] = min(args["n_thd_units"], ((tq - 1) // spec.cga_tile_m + batch) * heads) workspace = int(graph.get_workspace_size()) alignment = int(spec.workspace_alignment) - if not (0 <= spec.off_o_desc < spec.scratch_bytes <= workspace < 2**63): + if not 0 <= spec.off_o_desc < spec.scratch_bytes <= workspace < 2**63: raise RuntimeError("Invalid cuDNN OSS forward workspace contract") if alignment < 16 or alignment & (alignment - 1): raise RuntimeError("Invalid cuDNN OSS forward workspace alignment") @@ -133,7 +143,7 @@ def get_frost_attention_forward(key, batch, heads, kv_heads, sq, sk, tq, tk, sca except (ImportError, PackageNotFoundError) as error: _LOG.warning("cuDNN OSS forward dependency unavailable: %s", error) return False - error, current = runtime.cudaGetDevice() + error, current = runtime.cudaGetDevice() # pylint: disable=c-extension-no-member if int(error) != 0 or int(current) != device: raise RuntimeError("cuDNN OSS preparation must use the calling CUDA device") io = cudnn.data_type.BFLOAT16 From a708a5a83598053d49afd86e505413025f086007 Mon Sep 17 00:00:00 2001 From: Yang Xu Date: Mon, 19 Oct 2026 06:47:00 -0700 Subject: [PATCH 8/8] Support known cuDNN OSS forward launch signatures Signed-off-by: Yang Xu --- AGENTS.md | 9 + qa/L1_pytorch_frost_unittest/test.sh | 10 +- .../attention/test_frost_attn_fwd_contract.py | 163 ++++++++++++++++++ .../common/CuTeDSL/frost_attn_fwd.py | 37 +++- .../common/fused_attn/frost_attn.cu | 8 +- 5 files changed, 217 insertions(+), 10 deletions(-) create mode 100644 tests/pytorch/attention/test_frost_attn_fwd_contract.py diff --git a/AGENTS.md b/AGENTS.md index 0f80257a3bd..d6b488c5134 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -87,6 +87,15 @@ but steady-state dispatch should use the cached C++/TVM FFI function. Framework-specific tensor handling, autograd or transformation semantics, state, and user-facing composition belong in the corresponding frontend. +Adapters to a private prepared/FFI host ABI must allowlist exact argument orders +and validate neutral values for unsupported optional features. When FE adds host +slots, update both the Python exporter and the bounded native argument frame; +pass the actual argument count to FFI. Do not accept unknown slots merely because +they fit the frame. For the FROST forward bridge, run +`tests/pytorch/attention/test_frost_attn_fwd_contract.py` for legacy/current ABI and +rejection coverage, then the actual common-library forward/replay tests. CPU +contract checks alone do not validate a compiled GPU entry. + ### Working architecture rules These rules describe the intended architecture. Check them against the diff --git a/qa/L1_pytorch_frost_unittest/test.sh b/qa/L1_pytorch_frost_unittest/test.sh index 9bb936ffef1..18d6413065d 100644 --- a/qa/L1_pytorch_frost_unittest/test.sh +++ b/qa/L1_pytorch_frost_unittest/test.sh @@ -33,7 +33,8 @@ 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" \ - "$TE_PATH/tests/pytorch/attention/test_frost_attn_fwd.py" + "$TE_PATH/tests/pytorch/attention/test_frost_attn_fwd.py" \ + "$TE_PATH/tests/pytorch/attention/test_frost_attn_fwd_contract.py" # Reject a green all-skipped run or selection that omitted native-route guards. python3 - "$XML_LOG_DIR/pytest_test_frost_attn.xml" <<'RESULTS' @@ -41,7 +42,7 @@ import sys import xml.etree.ElementTree as ET root = ET.parse(sys.argv[1]).getroot() cases = list(root.iter("testcase")) -assert len(cases) >= 79, "FROST QA did not execute the complete regression suite" +assert len(cases) >= 104, "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"): @@ -50,5 +51,10 @@ assert "test_frost_d256_bottom_right_q_longer_fallback" in names for kind in ("q_origin", "kv_origin", "q_limit", "kv_limit", "q_decrease", "q_base", "graph_origin"): assert f"test_frost_forward_metadata_guard[{kind}]" in names, "Missing isolated forward guard: " + kind assert "test_frost_forward_compact_decline" in names +for prefix, count in (("test_frost_forward_known_contract", 3), + ("test_frost_forward_declines_optional_features", 10), + ("test_frost_forward_declines_unknown_contract", 5), + ("test_frost_forward_retains_feature_guards", 7)): + assert sum(name.startswith(prefix + "[") for name in names) == count, "Missing forward contract tests: " + prefix print("FROST QA passed:", len(cases), "tests, zero skips, all twelve isolated guards") RESULTS diff --git a/tests/pytorch/attention/test_frost_attn_fwd_contract.py b/tests/pytorch/attention/test_frost_attn_fwd_contract.py new file mode 100644 index 00000000000..1ea08b6a0f6 --- /dev/null +++ b/tests/pytorch/attention/test_frost_attn_fwd_contract.py @@ -0,0 +1,163 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""Known FE host signatures and fail-closed forward export; no GPU required.""" + +import math +from types import SimpleNamespace + +import pytest +import tvm_ffi + +from transformer_engine.common.CuTeDSL import frost_attn_fwd + +_BASE_ORDER = tuple( + "q_ptr k_ptr v_ptr o_ptr lse_ptr sinks_ptr meta_ptr o_desc_ptr problem_size " + "q_strides k_strides v_strides o_strides lse_strides lse_ext scale_softmax_log2 " + "n_thd_units seq_q_lens_addr thd_q_lens_ptr thd_kv_lens_ptr thd_lens_form " + "o_partial_ptr block_table_ptr block_table_v_ptr table_strides n_pages stream".split() +) +_SCORE = dict( + score_max_ptr=None, + score_sum_exp_ptr=None, + score_max_strides=(0, 0, 0), + score_sum_exp_strides=(0, 0, 0), +) +_BIAS = dict( + bias_ptr=None, + bias_strides=(0, 0, 0), + bias_extents=(1, 1), + kv_mask_ptr=None, + kv_mask_batch_stride=0, + inv_scale_softmax=1.0, +) + + +def _graph(extra): + order = _BASE_ORDER[:-1] + tuple(extra) + ("stream",) + args = dict.fromkeys(order) + for role, h in [("q", 4), ("k", 1), ("v", 1), ("o", 4)]: + args[role + "_strides"] = (h * 128, h * 128, 128) + args.update( + lse_strides=(0, 0, 0), + lse_ext=0, + seq_q_lens_addr=0, + thd_lens_form=3, + table_strides=(0, 0), + n_pages=0, + scale_softmax_log2=128**-0.5 * math.log2(math.e), + n_thd_units=64, + **extra, + ) + spec = SimpleNamespace( + order=order, + template=[args[n] for n in order], + fn=tvm_ffi.convert(lambda: None), + quant=None, + paged=False, + has_sink=False, + gate_expect=None, + has_lse=True, + lse_padded=False, + lse_head_major=False, + lse_stride_override=False, + lse_fill_plan=None, + split_workspace=None, + lens_form=3, + n_q_lens=3, + n_kv_lens=3, + device_index=0, + b=2, + qh=4, + kh=1, + d_qk=128, + d_v=128, + total_q=512, + total_kv=512, + s_q_max=256, + expect={r: "bf16" for r in ("q", "k", "v", "o")}, + cga_tile_m=128, + workspace_alignment=256, + off_o_desc=256, + scratch_bytes=512, + ) + return SimpleNamespace( + _compiled_plans=[SimpleNamespace(_prepared=SimpleNamespace(spec=spec))], + _plan_index=0, + get_workspace_size=lambda: 1024, + ) + + +def _export(module, g): + return module._export(g, 2, 4, 1, 256, 256, 512, 512, 128**-0.5, 0) + + +@pytest.mark.parametrize("extra", [{}, _SCORE, {**_SCORE, **_BIAS}]) +def test_frost_forward_known_contract(extra): + """Legacy and current FE declarations keep the stream and operands bound.""" + g = _graph(extra) + native, constants, metadata = _export(frost_attn_fwd, g) + spec = g._compiled_plans[0]._prepared.spec + count = len(spec.order) + assert native is spec.fn + assert len(constants) == metadata[3] == count + assert len(metadata) == 22 + 2 * count + assert metadata[-2:] == [3, 0] + for index in range(26, count - 1): + assert metadata[22 + 2 * index : 24 + 2 * index] == [0, index] + + +@pytest.mark.parametrize("name,value", list({**_SCORE, **_BIAS}.items())) +def test_frost_forward_declines_optional_features(name, value): + """New slots never silently enable score outputs, bias or KV masks.""" + g = _graph({**_SCORE, **_BIAS}) + spec = g._compiled_plans[0]._prepared.spec + changed = 8 if value is None else ((8,) * len(value) if isinstance(value, tuple) else 8) + spec.template[spec.order.index(name)] = changed + with pytest.raises(RuntimeError, match="Unsupported cuDNN OSS forward optional argument"): + _export(frost_attn_fwd, g) + + +@pytest.mark.parametrize( + "kind", ["unknown", "missing", "reordered", "duplicate", "truncated_template"] +) +def test_frost_forward_declines_unknown_contract(kind): + """A changed positional ABI must fail before any FFI invocation.""" + g = _graph({**_SCORE, **_BIAS}) + spec = g._compiled_plans[0]._prepared.spec + order = list(spec.order) + if kind == "unknown": + order[-2] = "new_slot" + elif kind == "missing": + order.pop(-2) + spec.template.pop(-2) + elif kind == "reordered": + order[0], order[1] = order[1], order[0] + elif kind == "duplicate": + order[1] = order[0] + else: + spec.template.pop() + spec.order = order + with pytest.raises(RuntimeError, match="Unsupported cuDNN OSS prepared forward contract"): + _export(frost_attn_fwd, g) + + +@pytest.mark.parametrize( + "name,value", + [ + ("gate_expect", "bf16"), + ("has_sink", True), + ("paged", True), + ("lse_padded", True), + ("quant", object()), + ("split_workspace", object()), + ("lens_form", 2), + ], +) +def test_frost_forward_retains_feature_guards(name, value): + """ABI compatibility does not broaden the supported operation contract.""" + g = _graph({**_SCORE, **_BIAS}) + setattr(g._compiled_plans[0]._prepared.spec, name, value) + with pytest.raises(RuntimeError, match="Unsupported cuDNN OSS prepared forward contract"): + _export(frost_attn_fwd, g) diff --git a/transformer_engine/common/CuTeDSL/frost_attn_fwd.py b/transformer_engine/common/CuTeDSL/frost_attn_fwd.py index 76c5e9abe08..5a8f87502bc 100644 --- a/transformer_engine/common/CuTeDSL/frost_attn_fwd.py +++ b/transformer_engine/common/CuTeDSL/frost_attn_fwd.py @@ -21,6 +21,28 @@ "n_thd_units seq_q_lens_addr thd_q_lens_ptr thd_kv_lens_ptr thd_lens_form " "o_partial_ptr block_table_ptr block_table_v_ptr table_strides n_pages stream".split() ) +# Known optional, dense-only host slots. THD must keep every one disabled. +_SCORE_DEFAULTS = { + "score_max_ptr": None, + "score_sum_exp_ptr": None, + "score_max_strides": (0, 0, 0), + "score_sum_exp_strides": (0, 0, 0), +} +_BIAS_DEFAULTS = { + "bias_ptr": None, + "bias_strides": (0, 0, 0), + "bias_extents": (1, 1), + "kv_mask_ptr": None, + "kv_mask_batch_stride": 0, + "inv_scale_softmax": 1.0, +} +_ORDERS = ( + _ORDER, + _ORDER[:-1] + tuple(_SCORE_DEFAULTS) + _ORDER[-1:], + _ORDER[:-1] + tuple(_SCORE_DEFAULTS) + tuple(_BIAS_DEFAULTS) + _ORDER[-1:], +) + + _POINTERS = { "q_ptr": 0, "k_ptr": 1, @@ -43,8 +65,10 @@ def _export(graph, batch, heads, kv_heads, sq, _sk, tq, tk, scale, device): import tvm_ffi spec = graph._compiled_plans[graph._plan_index]._prepared.spec + order = tuple(spec.order) if ( - tuple(spec.order) != _ORDER + order not in _ORDERS + or len(spec.template) != len(order) or not isinstance(spec.fn, tvm_ffi.Function) or spec.quant is not None or spec.paged @@ -66,7 +90,10 @@ def _export(graph, batch, heads, kv_heads, sq, _sk, tq, tk, scale, device): or spec.s_q_max != sq ): raise RuntimeError("Unsupported cuDNN OSS prepared forward contract") - args = dict(zip(spec.order, spec.template)) + args = dict(zip(order, spec.template)) + for name, expected in {**_SCORE_DEFAULTS, **_BIAS_DEFAULTS}.items(): + if name in args and args[name] != expected: + raise RuntimeError("Unsupported cuDNN OSS forward optional argument: " + name) for role, h in (("q", heads), ("k", kv_heads), ("v", kv_heads), ("o", heads)): if spec.expect[role] not in ("bf16", "bfloat16"): raise RuntimeError("cuDNN OSS forward requires BF16 operands") @@ -107,11 +134,11 @@ def _export(graph, batch, heads, kv_heads, sq, _sk, tq, tk, scale, device): ) if any(n <= 0 or n >= 2**63 for n in spans): raise RuntimeError("cuDNN OSS forward operand span is not representable") - metadata = [1, workspace, alignment, len(_ORDER)] + metadata = [1, workspace, alignment, len(order)] for i, span in enumerate(spans): metadata.extend((span, 16 if i < 4 else 4)) constants = [] - for i, name in enumerate(_ORDER): + for i, name in enumerate(order): # kind 0 = immutable value; 1 = current operand; 2 = workspace offset; # kind 3 = current stream. The C++ bridge never constructs FFI arrays. kind, value = 0, i @@ -186,7 +213,7 @@ def attach(tensor, name, factor): cu_seq_len_q=lengths[0], cu_seq_len_kv=lengths[1], max_total_seq_len_q=tq, - max_total_seq_len_kv=tk + max_total_seq_len_kv=tk, ) o.set_output(True).set_dim([batch, heads, sq, 128]).set_stride( [sq * heads * 128, 128, heads * 128, 1] diff --git a/transformer_engine/common/fused_attn/frost_attn.cu b/transformer_engine/common/fused_attn/frost_attn.cu index cfaf42282b8..f71c90da51d 100644 --- a/transformer_engine/common/fused_attn/frost_attn.cu +++ b/transformer_engine/common/fused_attn/frost_attn.cu @@ -396,7 +396,8 @@ const ForwardPlan *get_forward_plan(const FusedAttnConfig &c, cudaStream_t strea int64_t value; while (input >> value) v.push_back(value); NVTE_CHECK(input.eof() && v.size() >= 22 && v[0] == 1 && v[1] > 0 && v[2] >= 16 && - (v[2] & (v[2] - 1)) == 0 && v[3] == 27 && v.size() == 22 + 2 * v[3], + (v[2] & (v[2] - 1)) == 0 && (v[3] == 27 || v[3] == 31 || v[3] == 37) && + v.size() == 22 + 2 * v[3], "Malformed cuDNN OSS forward metadata"); auto function = tvm::ffi::Function::GetGlobal(key); NVTE_CHECK(function.has_value(), "cuDNN OSS did not register its native forward entry"); @@ -489,7 +490,8 @@ bool frost_attention_forward(const FusedAttnConfig &cfg, const FrostForwardPoint static_cast(pointers[7]), static_cast(pointers[8]), cfg.batch_size, cfg.max_seqlen_q, cfg.max_seqlen_kv, cfg.num_tokens_q, cfg.num_tokens_kv); NVTE_CHECK_CUDA(cudaGetLastError()); - std::array frame; + // Largest explicitly supported host ABI; retain stack storage on every call. + std::array frame; for (size_t i = 0; i < plan->arguments.size(); ++i) { const auto &arg = plan->arguments[i]; if (arg.kind == 0) frame[i] = static_cast(arg.constant); @@ -498,7 +500,7 @@ bool frost_attention_forward(const FusedAttnConfig &cfg, const FrostForwardPoint if (arg.kind == 3) frame[i] = reinterpret_cast(stream); } tvm::ffi::Any result; - plan->function.CallPacked(frame.data(), static_cast(frame.size()), &result); + plan->function.CallPacked(frame.data(), static_cast(plan->arguments.size()), &result); return true; }