feat(attention): cuDNN FROST attention backend for head_dim in (256, 512], with context parallelism

docs/envvars.rst

@@ -149,10 +149,13 @@ Then it applies a performance-based preference order among the remaining eligibl
 In PyTorch, the broad preference order is ``FlashAttention > FusedAttention >
 UnfusedDotProductAttention`` on supported pre-Hopper GPUs such as Ampere/Ada, and
 ``FusedAttention > FlashAttention > UnfusedDotProductAttention`` on Hopper and newer GPUs,
-including Blackwell. In JAX, Transformer Engine uses cuDNN fused attention when
-``NVTE_FUSED_ATTN=1`` and an eligible cuDNN kernel is available; otherwise it falls back to the
-JAX-native implementation. See :doc:`examples/attention/attention` for a longer
-backend-selection overview.
+including Blackwell. On Blackwell SM100/SM103 the order is ``FusedAttention > FlashAttention >
+FrostAttention > UnfusedDotProductAttention``; FrostAttention only becomes eligible for
+symmetric ``head_dim`` in (256, 512], which flash and fused attention do not serve, so the
+backend it can displace is UnfusedDotProductAttention. In JAX, Transformer Engine uses cuDNN
+fused attention when ``NVTE_FUSED_ATTN=1`` and an eligible cuDNN kernel is available; otherwise
+it falls back to the JAX-native implementation. See :doc:`examples/attention/attention` for a
+longer backend-selection overview.
 
 .. envvar:: NVTE_FLASH_ATTN
 
@@ -184,6 +187,12 @@ backend-selection overview.
    :Default: ``1``
    :Description: Enable or disable FusedAttention backend (cuDNN-based) for DotProductAttention. When set to ``0``, FusedAttention will not be used.
 
+.. envvar:: NVTE_FROST_ATTN
+
+   :Type: ``int`` (0 or 1)
+   :Default: ``1``
+   :Description: Enable or disable FrostAttention backend (the cuDNN FROST CuTe-DSL SDPA kernels in cuDNN Frontend) for DotProductAttention. **This backend is experimental and subject to change**, including the possibility of being folded into FusedAttention; the underlying cuDNN FROST engines are themselves experimental. When set to ``0``, FrostAttention will not be used. From released components it is the only backend serving symmetric ``head_dim`` in (256, 512] together with context parallelism; without context parallelism UnfusedDotProductAttention also covers that range, and FrostAttention is preferred over it where both are eligible. It is limited to SM100/SM103 with BF16/FP16 inputs, a ``head_dim`` that is a multiple of 8, and ``nvidia-cudnn-frontend>=1.29.0`` and ``nvidia-cutlass-dsl>=4.7.0`` installed. It supports context parallelism with ``cp_comm_type`` of ``p2p``, ``all_gather``, ``a2a`` or ``a2a+p2p``, and sliding-window attention with ``all_gather`` or ``a2a`` (declined with ``p2p`` and ``a2a+p2p``, whose ring shards KV across steps). It declines FP8, ``thd`` layouts, dropout, attention bias, softcap, KV caching, ``max_logit``, and deterministic execution, the last because cuDNN offers no deterministic backward for these kernels.
+
 .. envvar:: NVTE_UNFUSED_ATTN
 
    :Type: ``int`` (0 or 1)

qa/L0_pytorch_unittest/test.sh

@@ -63,6 +63,7 @@ python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_hybrid_quantizat
 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"
 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"
+python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_frost_attention.xml $TE_PATH/tests/pytorch/attention/test_frost_attention.py || test_fail "test_frost_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"
 NVTE_GDN2_TEST_REQUIRED=1 python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_gdn2_attention.xml $TE_PATH/tests/pytorch/attention/test_gdn2_attention.py || test_fail "test_gdn2_attention.py"
 NVTE_GDP_TEST_REQUIRED=1 python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_gdp_attention.xml $TE_PATH/tests/pytorch/attention/test_gdp_attention.py || test_fail "test_gdp_attention.py"

tests/pytorch/attention/run_attention_with_cp.py

@@ -17,6 +17,7 @@
 from transformer_engine.pytorch import DType
 from test_attention_with_cp import (
     model_configs_flash_attn,
+    model_configs_frost_attn,
     model_configs_fused_attn,
 )
 from transformer_engine.pytorch import (
@@ -273,6 +274,14 @@ def run_dpa_with_cp(
             config = copy.deepcopy(model_configs_fused_attn[model])
         else:
             assert False, f"{model=} is not a known FusedAttention CP config!"
+    if kernel_backend == "FrostAttention":
+        # Leave NVTE_FLASH_ATTN and NVTE_FUSED_ATTN at 0: FROST is the only backend that serves
+        # head_dim > 256, so get_attention_backend selects it on its own.
+        os.environ["NVTE_FROST_ATTN"] = "1"
+        if model in model_configs_frost_attn:
+            config = copy.deepcopy(model_configs_frost_attn[model])
+        else:
+            assert False, f"{model=} is not a known FrostAttention CP config!"
     assert config.attn_mask_type in [
         "causal",
         "no_mask",
@@ -593,6 +602,19 @@ def run_dpa_with_cp(
             pad_between_seqs=pad_between_seqs,
             fp8_output=fp8_mha,
         )
+        if kernel_backend == "FrostAttention":
+            # Assert the backend actually used, not just the one requested. FROST is currently
+            # the only selectable backend for these configs -- flash and fused are env-gated off
+            # and CP disables unfused -- so a silent substitution is impossible today and this
+            # would pass by construction. It is here so it stops passing if that stops being
+            # true, rather than quietly testing some other kernel.
+            from transformer_engine.pytorch.attention.dot_product_attention.dot_product_attention import (  # pylint: disable=import-outside-toplevel
+                _attention_backends,
+            )
+
+            assert _attention_backends[
+                "use_frost_attention"
+            ], "expected FrostAttention to be selected, got %s" % (_attention_backends,)
         if config.return_max_logit:
             out_, max_logit_ = out_
         if is_training:

tests/pytorch/attention/test_attention_with_cp.py

@@ -459,6 +459,16 @@ def test_cp_with_flash_attention_softcap(cp_pool, cp_comm_type):
     )
 
 
+# cuDNN FROST: symmetric head_dim in (256, 512] on SM100/SM103, the range no other backend
+# serves together with context parallelism. Shapes are Gemma-4 global layers, which is what
+# motivated the backend. seqlen must stay divisible by cp_size * 2 for causal load balancing.
+model_configs_frost_attn = {
+    #   test:         ModelConfig(b, sq, hq, dqk)
+    "cp_hd512_0": ModelConfig(2, 4096, 8, 512, num_gqa_groups=4, attn_mask_type="causal"),
+    "cp_hd512_1": ModelConfig(2, 4096, 8, 512, num_gqa_groups=4, attn_mask_type="no_mask"),
+    "cp_hd512_2": ModelConfig(2, 2048, 8, 512, num_gqa_groups=8, attn_mask_type="causal"),
+}
+
 model_configs_fused_attn = {
     # test: ModelConfig(b, sq, hq, dqk)
     "cp_1_0": ModelConfig(2, 4096, 12, 128, attn_mask_type="causal", return_max_logit=True),  # MHA
@@ -748,6 +758,99 @@ def test_cp_with_fused_attention(
     )
 
 
+def _frost_availability():
+    """Why FrostAttention cannot run here, or None if it can.
+
+    The backend needs cuDNN Frontend >= 1.29.0 and, less obviously,
+    nvidia-cutlass-dsl >= 4.7.0: cudnn-frontend only declares >= 4.6.2, and below the FROST floor
+    every FROST engine silently declines and ordinary backend plans are returned with no error.
+    Reporting the reason as a skip keeps that distinguishable from a real failure.
+    """
+    if get_device_compute_capability() not in ((10, 0), (10, 3)):
+        return "FrostAttention requires SM100/SM103 (the cuDNN d512 backward is Blackwell-only)."
+    from transformer_engine.pytorch.attention.dot_product_attention.frost_attention import (
+        is_frost_attention_available,
+    )
+
+    ok, reason = is_frost_attention_available()
+    return None if ok else reason
+
+
+@pytest.mark.parametrize("model", model_configs_frost_attn.keys())
+@pytest.mark.parametrize("qkv_format", ["bshd", "sbhd"])
+@pytest.mark.parametrize("cp_comm_type", ["p2p", "all_gather", "a2a", "a2a+p2p"])
+def test_cp_with_frost_attention(cp_pool, model, qkv_format, cp_comm_type):
+    """Context parallelism at head_dim 512, which no other backend serves.
+
+    thd is excluded because the backend declines it: it needs varlen support that is not
+    implemented.
+
+    a2a+p2p needs four ranks rather than two -- an a2a subgroup crossed with a p2p subgroup -- and
+    exercises no new attention code: it dispatches to the same AttnFuncWithCPAndKVP2P as plain p2p,
+    with an a2a communication stage on either side of the ring. It is covered here so that claim is
+    measured rather than assumed.
+    """
+    reason = _frost_availability()
+    if reason is not None:
+        pytest.skip(reason)
+
+    config = model_configs_frost_attn[model]
+    config.context_parallel = True
+    config.cp_comm_type = cp_comm_type
+
+    # a2a requires num_heads and num_gqa_groups divisible by the a2a subgroup size; every config
+    # here satisfies that, but assert rather than rely on it staying true.
+    if cp_comm_type == "a2a+p2p":
+        assert config.num_heads % 2 == 0 and config.num_gqa_groups % 2 == 0, (
+            f"cp_comm_type=a2a+p2p needs num_heads ({config.num_heads}) and num_gqa_groups"
+            f" ({config.num_gqa_groups}) divisible by the a2a subgroup size"
+        )
+
+    pool = cp_pool(4 if cp_comm_type == "a2a+p2p" else 2)
+
+    _submit(
+        pool,
+        dtype="bf16",
+        model=model,
+        qkv_format=qkv_format,
+        kernel_backend="FrostAttention",
+        cp_comm_type=cp_comm_type,
+        is_training=True,
+        log_level=pytest_logging_level,
+    )
+
+
+@pytest.mark.parametrize("cp_comm_type", ["p2p", "all_gather", "a2a"])
+def test_cp_with_frost_attention_fp16(cp_pool, cp_comm_type):
+    """One fp16 arm per comm type, since the matrix above is bf16 throughout.
+
+    The backend serves BF16 and FP16, but every context-parallel configuration was covered in bf16
+    only. fp16 has a far narrower exponent range, and the ring correction exponentiates a difference
+    of log-sum-exp values across steps, so a range problem would surface here rather than in the
+    non-CP numerics. One model and one layout keeps the cost to three cases rather than doubling
+    the matrix; a2a+p2p is omitted because it would need a second four-rank pool for a dtype that
+    exercises no additional code path.
+    """
+    reason = _frost_availability()
+    if reason is not None:
+        pytest.skip(reason)
+
+    config = model_configs_frost_attn["cp_hd512_0"]
+    config.context_parallel = True
+    config.cp_comm_type = cp_comm_type
+
+    _submit(
+        cp_pool(2),
+        dtype="fp16",
+        model="cp_hd512_0",
+        qkv_format="bshd",
+        kernel_backend="FrostAttention",
+        cp_comm_type=cp_comm_type,
+        is_training=True,
+        log_level=pytest_logging_level,
+    )
+
+
 @pytest.mark.skipif(get_cudnn_version() < (8, 9, 7), reason="cuDNN 8.9.7+ is required.")
 @pytest.mark.skipif(
     get_device_compute_capability() < (9, 0), reason="FusedAttention THD requires sm90+."