From ae00949cbb7e1fc1d2c43cdf3bf2d90bcb2dc924 Mon Sep 17 00:00:00 2001 From: Yang Xu Date: Fri, 16 Oct 2026 11:39:02 -0700 Subject: [PATCH 1/5] 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/5] 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/5] 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/5] 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/5] 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