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

[nvegesna-netizen]

Problem

get_attention_backend selects no backend at all for symmetric head_dim=512 with context
parallelism, so that configuration raises rather than running:

• FlashAttention 2 and 3 instantiate kernels up to head_dim 256
• FA4 is disabled at symmetric 512 on SM100/SM110 (TMEM budget)
• the C++ cuDNN fused path caps at 256
• UnfusedDotProductAttention serves 512 but is disabled under CP

Hybrid models such as Gemma 4, which interleave sliding-window layers at head_dim 256 with
global layers at 512, therefore cannot use context parallelism at all.

To be precise about the claim: this makes context parallelism available at symmetric 512 from
released components. It is not a claim to unbounded context. Dao-AILab/flash-attention#2877 adds
symmetric D512 kernels to FA4 and, with #3532, that path also works and scales further, but it is
unmerged and unreviewed. Measured capacity here, and the workspace behaviour that bounds it, are
in Known limitations.

Approach

cuDNN Frontend >= 1.29.0 ships CuTe-DSL ("FROST") SDPA kernels that serve symmetric 512 forward
and backward on SM100/SM103. They are reachable only through the cuDNN graph Python API, since
the FROST engines register at Python import time and need nvidia-cutlass-dsl, while TE's C++
builds against cuDNN Frontend headers only.

FROST is a FusedAttention sub-backend, FusedAttnBackend["FROST"], not a fourth top-level
backend:

• frost_attention.py holds the kernel wrapper and fused_attn_fwd/bwd behind the
  cpp_extensions.fused_attn signatures
• cpp_extensions/fused_attn.py routes to those two functions when the sub-backend is FROST,
  before any C++ call
• utils.py consults is_frost_attention_supported inside _get_fused_attn_backend, where the
  C++ selector returns No_Backend, and checks availability once at the end of
  get_attention_backend, the way flash-attn's version is checked

So FusedAttnFunc and attn_forward_func_with_cp reach these kernels without knowing which
sub-backend they got, and dot_product_attention.py needs no change at all.

Several declines now come from the existing fused filters rather than their own copies: the
context-parallel mask, window and a2a restrictions; the score_mod and softcap filters;
and FP8, since FusedAttnFunc forces the FP8 sub-backend on that path. The all-gather path's
causal -> causal_bottom_right rewrite applies to FROST for free, which is what gives it the
bottom-right band it needs on a trimmed KV range.

Verification

Full suite on B200 against this head, in a container with nvidia-cudnn-frontend 1.29.0 and
nvidia-cutlass-dsl 4.8.0:

check	result
FROST is reached as a FusedAttention sub-backend	8 gates pass
flex_attention.py regression	27 passed
non-distributed numerics vs a float64 reference	43 passed
context parallelism	22 passed
broader test_attention.py suite	2652 passed, 2605 skipped

The broad suite is identical to the pre-change baseline and flex_attention.py's own suite is
clean, which is what says the shared-module extraction below preserved behaviour rather than merely
compiling.

One of the eight gates covers the shared module itself: that cudnn_pygraph owns the only plan
cache and the only cache key, that neither backend kept its own, and that execute_graph cannot
have its handle name defaulted away. A plan built for one graph being served to another is the
failure that would not surface as a crash.

Correctness uses the criterion FlashAttention applies to itself,
err(kernel, fp64) <= 2 * err(inputs rounded to dtype, fp64), with that floor measured per case so
the bar tracks shape and dtype instead of encoding a number that rots. Coverage is square and
rectangular, causal / bottom-right / non-causal, GQA and MHA, forward and backward in bf16 and
fp16, plus fp16 arms under CP.

Context parallelism, each compared against the non-CP reference via run_attention_with_cp.py:

	bshd	sbhd
p2p	CP=2, CP=4	CP=2, CP=4
all_gather	CP=2, CP=4	CP=2
a2a	CP=2, CP=4	CP=2

Plus 2 nodes x 2 ranks for every comm type including a2a+p2p, with a FusedAttention control
passing throughout. a2a+p2p needs the multi-node arm specifically: at world_size 4 the a2a
level takes consecutive ranks (0,1),(2,3) and the p2p level takes strided ranks (0,2),(1,3),
so block assignment across two nodes puts the all-to-all inside each node and the ring across the
fabric. On one node both levels sit on NVLink and the arrangement is never exercised. The run
establishes the rank-to-host mapping first and declines to report unless it is block ordered,
since round-robin would invert the two levels and pass while testing the opposite topology.

test_frost_attention.py is non-distributed and checks the forward, the LSE convention and the
backward gradients against an independent float64 reference, across both causal alignments, bf16
and fp16, GQA and MHA, and a rectangular shape where top-left and bottom-right masking differ.
The CP tests cannot do this: run_attention_with_cp.py compares a CP run against a non-CP run of
the same backend, so a systematic kernel error appears on both sides and cancels. The reference
must be float64 rather than float32, because torch computes fp32 matmuls in TF32 on Ampere and
newer and TF32's significand is 11 bits, the same as fp16, so an fp32 reference is no more
accurate than the kernel under test.

End to end, a Gemma 4 dense model (sliding layers on FlashAttention, global layers on FROST)
matched its CP=1 result to 9.7e-05 on the loss.

These tests skip in CI as it stands, and the blocker is the dependency floors rather than the
Blackwell requirement. Stock nvcr.io/nvidia/pytorch:26.08-py3 ships nvidia-cudnn-frontend
1.26.0 and nvidia-cutlass-dsl 4.6.2, both below what FROST needs, so
is_frost_attention_available() declines and the tests skip with that reason rather than failing.
qa/L3_pytorch_FA_versions_test/test.sh targets sm100+ but pins nvidia-cutlass-dsl[cu13]==4.4.2
for the FA4 path, so it would skip even on a B200; and nvidia-cutlass-dsl is not a declared TE
dependency at all, arriving transitively via flash-attn-4. Happy to wire this into whichever
lane you consider the right home.

Performance

Against UnfusedDotProductAttention at the same shape (b2 hq8 hkv4 d512 causal bf16, CP=1,
single layer), the only other backend serving this head dim from released components:

seqlen	fwd speedup	fwd+bwd speedup	unfused peak (fwd)	FROST peak (fwd)
4096	5.8x	2.4x	1520 MiB	224 MiB
16384	15.2x	6.4x	51232 MiB	1057 MiB
32768	n/a (OOM)	n/a (OOM)	OOM on 178 GiB	2850 MiB

The unfused path materialises the full s x s score matrix, so it grows O(s^2). Note that in
forward+backward FROST uses more memory than unfused below roughly 8k, where the saved tensors
and workspace exceed the small score matrix; the crossover is between 8192 and 16384.

Notes for reviewers

Two places the "no changes outside frost_attention.py" goal does not quite hold, both the
same root cause: the selected sub-backend was being discarded and re-derived as
F16_arbitrary_seqlen in seven places, so those places now keep it instead.

• context_parallel.py (+29/-7): each of the three autograd classes recomputed the sub-backend
  locally rather than receiving it. attn_forward_func_with_cp gains one argument, threaded
  through, with the existing derivation kept as the fallback. One further line: the p2p forward
  step returned max_logit with a starred tail, which is only unambiguous while no backend has a
  statically known return length. The FROST branch gives pylint one, so the tail resolved to empty
  and the five-label unpacks tripped unbalanced-tuple-unpacking. Indexing says what the function
  actually returns.
• backends.py (+5/-5): the non-FP8 backward hardcoded F16_arbitrary_seqlen, discarding the
  value the forward had already saved, so it now simply reads
  ctx_attrs["fused_attention_backend"]. The CP assert admitted only that sub-backend.

FusedAttnBackend.FROST has no pybind counterpart. It is 3, past the C++ enum, and the
import-time sync assert now exempts python-only members. Nothing in transformer_engine/pytorch
constructs NVTE_Fused_Attn_Backend(int), and the dispatch happens before any C++ call, so the
value never reaches pybind. Every post-selection check in utils.py is guarded on == FP8 or
== F16_arbitrary_seqlen, so a FROST value falls through them.

is_frost_attention_supported deliberately does not probe availability. Probing imports cuDNN
Frontend with the FROST engines enabled, which is process-wide and reorders plan selection for
every other cuDNN consumer. That must not happen for the overwhelming majority of configs, which
are nowhere near head_dim 512. Availability is checked once at the end, beside the flash-attn
version checks.

Plan selection is deliberately strict. heur_mode.A with an explicit name-checked
select_plan, rather than A|FALLBACK with HEURISTICS_CHOICE. Without a pin, build_plans
walks the ranked list and finalizes the first plan that builds, logging declines at INFO. At d512
that matters in the forward, where an ordinary engine may build and compute a different function;
the backward is self-limiting, since no non-FROST d512 backward exists. Pinning is also what makes
check_support fatal rather than advisory, which is why it follows the selection.

Version constraint worth knowing. FROST enforces nvidia-cutlass-dsl >= 4.7.0 at plan-build
time while cudnn-frontend only declares >= 4.6.2. Below that floor every FROST engine silently
declines and ordinary backend plans are returned with no error, which is why the code checks the
selected plan by name rather than assuming the engine was used. Stock 26.08 ships exactly 4.6.2.
Separately, 4.7.1 is incompatible with flash-attn-4 4.0.0b11, which CI currently pins.

Scope

SM100/SM103 only (the cuDNN d512 backward is Blackwell-only), bf16/fp16, symmetric head_dim in
(256, 512], bshd and sbhd, mask types no_mask / causal / causal_bottom_right.

Symmetric specifically, and the constraint is the backward engine's. v carries its own graph
node and its own cache-key entry, so an asymmetric pair is expressible here. It is declined because
sdpa_bwd_sm100, the only f16 FROST backward reachable on SM100/SM103, leaves dqk_ge_dv unset,
which cuDNN reads as requiring d_qk == d_v at every head dim rather than only above 256. The
forward is the opposite: sdpa_fwd_prefill_sm100 lists (192, 128) among its native d_shapes
and has no dqk_ge_dv concept at all. So the decline covers both directions rather than training
alone, since is_training is module.training and eval() does not disable autograd, which makes
a served forward no guarantee that no backward follows.

The head-dim range is that same engine's declared envelope: d_envelope_floor=256 (exclusive),
d={512} and d_pad_multiple=8, which is _MIN_HEAD_DIM 257, _MAX_HEAD_DIM 512 and
_HEAD_DIM_MULTIPLE 8 here.

FP8 forward with a FROST backward is not supported, and cannot arise: FP8 serves head_dim
up to 256 and FROST only above it, so their ranges are disjoint. _fused_attn_setup_ctx therefore
hands the backward whatever the forward recorded, which is FP8 for an FP8 forward and FROST for a
FROST one. Supporting a mixed pair would need get_attention_backend to return a backend per
direction, which is out of scope here.

Declined by the selector, each an explicit decline rather than a silent fallback: FP8, attention
bias, dropout, softcap, non-vanilla softmax, thd, max_logit, KV caching, CUDA graph
capture, and deterministic execution. thd is expressible with these kernels but is not implemented here.
Determinism is declined because cuDNN has no deterministic backward for them at all; measured,
not assumed.

Also declined: a right-bounded window on a non-causal mask when max_seqlen_q != max_seqlen_kv.
That combination takes its alignment only from bottom_right_diagonal, which defaults to
top-left, while the all-gather ring measures its window against the bottom-right diagonal. The two
differ exactly when the lengths do, so FROST declines rather than guessing.

Sliding window is supported, with all_gather or a2a. It is declined with p2p and
a2a+p2p, whose ring shards KV across steps so a bound measured against the full sequence does
not survive the per-step tiles, which is the same rule FusedAttention carries. That matters for
the motivating model: Gemma 4's sliding layers are the other half of it.

Known limitations and open decisions

Backward memory grows super-linearly with sequence length. Measured on B200,
b2 hq8 hkv4 d512 causal bf16, single GPU, peak allocated for forward + backward:

seqlen	4,096	8,192	16,384	32,768	65,536
peak	2.0 GiB	8.0 GiB	8.0 GiB	32.0 GiB	112.0 GiB

The forward is not the source: its workspace measures zero at every length above. The growth is in
the backward, where the cuDNN graph's get_workspace_size() request rises from 1.6 GiB at 4k to
74 GiB at 64k, and this module allocates what the graph asks for.

The consequence is a ceiling on the per-rank shard rather than on the model: shards beyond roughly
32k tokens are impractical on a 180 GB device, so a long-context configuration needs a
correspondingly higher cp_size. Largest global sequence that fits:

cp_size	1	2	4	8
longest context	65,536	131,072	131,072	262,144

The request is also layout-dependent, which matters if you try to reproduce these numbers.
Same logical problem, same plan, differing only in the strides the graph is built against:

seqlen	4,096	8,192	16,384	32,768	65,536
bshd strided view	1.13 GiB	4.25 GiB	4.50 GiB	17.00 GiB	66.00 GiB
bhsd contiguous	1.63 GiB	5.25 GiB	6.50 GiB	21.00 GiB	74.00 GiB

The difference is exactly linear, 128 KiB per query token at both ends of the sweep, while the
super-linear term is identical in both layouts. Note the direction: the contiguous layout is the
more expensive one. This module keys its plan cache on each tensor's actual strides, so which
figure applies follows from the caller's layout. Both columns are reproducible without
TransformerEngine by building the same two graphs through the cuDNN graph API and reading
get_workspace_size(), which is how they were measured: nothing executed, no workspace allocated.

The cudnn_pygraph.py extraction is done. frost_attention.py and flex_attention.py now
drive cuDNN through one shared module: the frontend import and its process-wide engine switch, the
per-device handle with its per-call set_stream, the dtype mapping, the BHSD tensor description,
the SDPA forward and backward graph builders, the plan cache and the key that identifies a
plan, plan finalization and graph execution. flex_attention.py comes out at +126/-268 against main.

The reason to share these is ownership rather than line count, and the module docstring says so: the
state behind them is process-global, and giving it two owners is how this code has produced bugs
before. flex also picks up a fix on the way, since it was building graphs under whatever device
happened to be current while only the handle named the right one.

Worth stating plainly: this takes the duplication from three pygraph sites to two, not to one.
transformer_engine/jax/cpp_extensions/flex_attention.py carries its own copy, shares no code with
either, and cannot consume a torch-dependent module.

The graph builders are shared as cudnn_pygraph.build_fwd and build_bwd. Each backend passes
tensor descriptors, any auxiliary tensors its callback reads, and its own sdpa arguments, so a
mask, a score_mod and the choice of which engine to pin or bar stay with the backend that wants
them.

What stayed behind is the part that is not cuDNN mechanics: the selector, the TE mask vocabulary,
the runtime guards and the cpp_extensions shims.

flex is barred from the FROST engines, by default rather than on request. Those engines accept
a score_mod graph, pass check_support, build, and then compute without the callback. Measured on
B200: a FROST plan ranks at index 0 and an unpinned build selects it at head_dim 64, 128, 256 and
512, with the output tracking a float64 reference computed without the modifier. The switch that
offers them is process-wide, so declining to ask for them is not enough.

finalize_plans now bars them unless the caller pins a plan by name, which FROST does and flex does
not. The default direction is deliberate: forgetting to exclude gives silently wrong numbers, while
excluding wrongly gives a slower plan or a loud decline. The engine names live in one place, so the
pin and the bar cannot name different engines; previously they were separate string literals in two
files, and cuDNN has renamed that family once already.

It cannot be checked any cheaper than end to end. deselect_engines marks rather than removes,
measured on B200 at head_dim 64, 128 and 256: the plan count is unchanged and index 0 still names
the barred engine after build_plans, so a post-build assertion on the selected plan would fire
falsely. What verifies the bar is test_frost_switch_does_not_change_what_flex_computes, which runs
the score_mod path with the engines on and off and compares. NVTE_FLEX_TEST_REQUIRED exists so a
lane can make that test fail rather than skip.

The underlying cuDNN defect is filed upstream: those engines do have a score_mod capability gate,
but it reads a key the sdpa() path never writes. The bar carries that as a removal condition in
its docstring rather than becoming permanent by default.

One handle is shared across overlapped CP streams. The wait_stream serialization in
context_parallel.py exists for FA3/FA4's internal per-call workspace, and its comment says
FusedAttention keeps the per-step overlap. FROST's workspace is caller-owned and per-call, and
FusedAttention uses the identical one-handle-per-device plus cudnnSetStream-per-call model while
being deliberately left overlapped, so I have not added FROST to that guard.

Dependency floors are intentionally not enforced by the build. Bumping
nvidia-cudnn-frontend to 1.29.0 would force nvidia-cutlass-dsl and apache-tvm-ffi into every
TE-PyTorch install, because 1.29.0 drops the cutedsl extra marker, and it still would not
guarantee the 4.7.0 floor, since the transitive requirement is 4.6.2. Every other optional backend
here (flash-attn 2/3/4, GDN, quack) is undeclared and runtime-probed, which is what this does.
NVTE_FROST_TEST_REQUIRED and NVTE_FLEX_TEST_REQUIRED mirror NVTE_GDN_TEST_REQUIRED, so a lane
meant to cover either can fail loudly rather than skipping. Neither is set in qa: there is no
Blackwell L0 lane for the first, and I cannot verify the L0 container carries the frontend package
for the second. The flex one matters more than it used to, since that file's single numerical test
is now what says this refactor preserved its behaviour.

Documentation. NVTE_FROST_ATTN is in docs/envvars.rst, described as a FusedAttention
sub-backend. The backend tables in docs/examples/attention/attention.ipynb are not updated yet.

[greptile-apps]

Confidence Score: 5/5

[High impact] The changes since the last review appear safe to merge.

Summary

This PR adds FROST as a Python-only sub-backend of FusedAttention for supported Blackwell attention configurations above head dimension 256.

• FusedAttention can use FROST for supported head dimensions through 512.
• FROST and Flex Attention share cuDNN graph tools.

The supplied previous threads are unnumbered, so they receive no numbered ledger entries. Their fixes and withdrawn findings remain consistent with the current diff.

Diagram

%%{init: {'theme': 'neutral'}}%%
flowchart TD
    A[Attention configuration] --> B[Existing fused selector]
    B -->|Existing backend available| C[C++ fused attention]
    B -->|No backend| D[FROST support checks]
    D --> E[Architecture and dependency checks]
    E -->|Supported| F[FusedAttention with FROST]
    F --> G[Ordinary or context-parallel calls]
    G --> H[Python FROST dispatch]
    H --> I[Shared cuDNN graph cache and execution]

Loading

_{Reviews (95) · Last reviewed commit: "test(attention): drop the FROST k/v mism..." · Reviewed by Greptile}

[cyanguwa]

@nvegesna-netizen, I made a few changes in 7c56bcb. I think these files should be good now: backends.py, context_parallel.py, utils.py, fused_attn.py. But these files will need another look: cudnn_pygraph.py, flex_attention.py, frost_attention.py. There's still some duplication happening in them, for example, _bhsd_dim_stride and _bhsd_tensor_spec. Could you please take a closer look at these files, consolidate their functions, and make their code structure as lean as safely possible?

If we have taken a function out of Flex to cudnn_pygraph, we can remove their shims as well, for example, _bhsd_dim_stride -> bhsd_dim_stride, device_key -> _score_mod_device_key, finalize_plans -> finalize_cudnn_graph. Don't be afraid to touch the Flex code :) FusedAttentionWithScoreModFunc should be left alone but the rest of the file should be fine for some modification.

I think at the end of the day, we want to just have one set of these functions in cudnn_pygraph: build_graph_fwd/bwd, cache_graph (including the key generation, the global cache dict, etc), finalize_plans (which is already good I think), and so on. Flex and FROST can then call these shared functions, with a couple of changes on the input args (e.g. Flex with score_mod), and it should all just work!

Some of the opts, such as bottom_right_diagonal, are specific to FROST, but they can be passed to cudnn_pygraph.build_graph_fwd/bwd as well. They just need to be False or some other specific values, for a specific BACKEND_NAME to work. FROST has is_frost_attention_supported to disable some of the configs it doesn't support. Flex is the same, and some of its condition checks are in utils.py (can be moved to flex_attention technically).

Meanwhile, please keep the comments/docstrings short and sweet - I think so far TE has adopted a mostly defensive (sparse and empirical) documentation style.

Thanks.

[cyanguwa]

/te-ci pytorch L0

[nvegesna-netizen]

Done in 6d65eb6, 42eaaab, a66f88f, e72c5cd and e88dc70.

Removed: both _import_cudnn_frontend, _bhsd_dim_stride, _bhsd_tensor_spec, _score_mod_device_key, _score_mod_bhsd_tensor_metadata, _finalize_cudnn_graph, _execute_cudnn_graph, and frost's _bhsd and its described closure. _finalize_cudnn_graph had no production caller.

cudnn_pygraph now owns the cache and the key. cached_graph takes backend_name and keys one module-level cache; sdpa_cache_key states what identifies a plan, with anything backend-specific passed through an opaque extras, so flex's score_mod protocol stays in flex; execute_graph requires backend_name, since it picks the handle.

On comments, I kept the cuDNN and engine facts and moved decline reasons into the messages, since the subtree explains declines through logger.debug rather than comments.

One thing to flag: FusedAttentionWithScoreModFunc was touched three times, all call-site updates, with no change to its logic, signature, ctx protocol or return. One follows from removing _bhsd_dim_stride, which you asked for; the other two from removing _execute_cudnn_graph, which was my call.