feat(attention): cuDNN FROST attention backend for head_dim in (256, 512], with context parallelism by nvegesna-netizen · Pull Request #3527 · NVIDIA/TransformerEngine

nvegesna-netizen · 2026-09-16T04:32:15Z

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, which rules out
long-context training for them.

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 through the ordinary cuDNN graph API but had no
TE backend, so attn_forward_func_with_cp was unreachable for them. This adds one.

What this adds:

frost_attention.py: kernel wrapper with a plan cache
selection in get_attention_backend, plus FrostAttention / FrostAttnFunc in backends.py
CP dispatch for p2p, all_gather, a2a and a2a+p2p
CP test coverage at head_dim 512

Verification (B200)

Correctness uses the criterion FlashAttention applies to itself,
err(kernel, fp64) <= 2 * err(naive_bf16, fp64); observed 0.21x to 0.94x across square and
rectangular, causal and non-causal, GQA and MHA shapes.

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

a2a+p2p runs on 4 ranks (2 a2a x 2 p2p) and was added after the rest: it dispatches to the same
AttnFuncWithCPAndKVP2P as plain p2p, so it needed no new attention code, but
FrostAttention.forward computed context_parallel in a form that assumed a single process
group and raised TypeError: unhashable type: 'list' on the list a2a+p2p passes. It now uses
the same form as FlashAttention and FusedAttention.

Plus 2 nodes x 2 ranks for p2p, all_gather and a2a, and a FusedAttention regression
control passing throughout.

test_cp_with_frost_attention covers the same matrix from pytest: 24 tests, all passing on a
B200 with the dependencies installed.

Re-verified on 8x B200 at the branch tip (74b31a9a), in one clean run. Every section below is
from that single run, on the commit it describes:

check	result
non-distributed numerics vs a float64 reference	46 passed
context parallelism, 24 configurations	24 passed
`FusedAttention` CP regression control	pass
mixed-THD suite	33 passed, 1 skipped
`test_attention.py` sweep	388 passed, 237 skipped, 0 failed

and the specific behaviours the review changes introduced:

Stream binding. A forward issued inside torch.cuda.stream(side) is bitwise identical to
one on the default stream, over 50 repeats. This is the case the p2p ring actually exercises
on alternating steps, and no earlier probe could have detected the missing binding because
they all ran on the default stream.
Determinism. cuDNN refuses a deterministic FROST backward outright --
cudnnGraphNotSupportedError, no engine proposes a plan -- which is why the selector declines
rather than passing the flag. Confirmed both halves: the kernel still refuses, and
get_attention_backend returns FROST for deterministic=False and not for True.
Regression controls. FusedAttention CP passes, test_mixed_thd_attention.py is 33 passed
/ 1 skipped, and the broader test_attention.py sweep is 388 passed / 237 skipped / 0 failed.
The mixed-THD suite matters because get_attention_backend gained a return value, and the
THD policy path unpacks it.

test_frost_attention.py is new and non-distributed: 46 tests checking 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 error in the kernel appears on both sides and cancels. Note the reference must be
float64 rather than float32: 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 an fp16
kernel and the error floor collapses.

These tests will 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.

Two specifics that decide where coverage could come from, if you want it:

qa/L3_pytorch_FA_versions_test/test.sh is the lane that targets sm100+, but it pins
nvidia-cutlass-dsl[cu13]==4.4.2 for the FA4 path, so it would skip even on a B200.
nvidia-cutlass-dsl is not a declared TE dependency at all; it arrives transitively via
flash-attn-4. So no current lane guarantees the floor.

The arch gate itself follows existing practice — test_attention_with_cp.py already has an
SM100/SM103-only skip for the D=256 CP fused path. Happy to wire this into whichever lane you
consider the right home.

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

Performance

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

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) and runs out of
memory at seqlen 32768 on a 178 GiB device. 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

A separate use_frost_attention flag rather than a FusedAttnBackend value. That enum
mirrors NVTE_Fused_Attn_Backend value-for-value and is consumed by fused_attn_fwd, which
dispatches into C++ that caps at 256, so routing through it would feed a value into a path that
cannot honour it. The cost is widening the get_attention_backend return tuple.

One assertion relaxed, and it deserves scrutiny. context_parallel.py asserts
qkv_format != "sbhd" or use_fused_attention, inferring "not fused means flash". FROST is
neither, and builds its graphs from each tensor's actual strides, so sbhd is served directly.
This matters because Megatron uses sbhd internally. Other instances of the same inference are
safe by construction here (thd is declined by the selector, and the sliding-window asserts now
name FrostAttention), and the
softmax_type one is now declined in the selector rather than reaching the assert.

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, so the code checks the
selected plan by name rather than assuming it was used. Stock 26.08 ships exactly 4.6.2, which is
the version that satisfies the declared floor while declining every engine, so anyone upgrading
only cudnn-frontend lands in that trap. Separately, 4.7.1 is incompatible with flash-attn-4
4.0.0b11, which CI currently pins.

Naming. FrostAttention follows the wording already used for these engines in
tests/pytorch/attention/test_gdn_attention.py, but unlike FlashAttention / FusedAttention /
UnfusedDotProductAttention it names an engine family rather than an approach. Happy to rename.

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.

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, and
deterministic execution. thd is expressible with these kernels (ragged offsets plus a padding
mask do offer a FROST plan at d512) but is not implemented here. Determinism is declined because
cuDNN has no deterministic backward for them at all -- measured, not assumed.

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 -- the same rule FusedAttention carries. That matters for the motivating model:
Gemma 4's sliding layers are the other half of it.

Two items above would specifically benefit from the attention owners' view: the naming, and
the relaxed sbhd assertion in context_parallel.py. Happy to change either.

Known limitations and open decisions

Each of these was investigated against the tree rather than left to be found in review.

deterministic is now honoured. It was accepted and ignored: cuDNN's SDPA backward is
non-deterministic unless the graph asks otherwise, which is why the C++ path calls
set_deterministic_algorithm and flex_attention passes use_deterministic_algorithm to the
same sdpa_backward this module builds. The flag is threaded through and is part of the
plan-cache key. One thing to confirm on hardware: if cuDNN offers no deterministic FROST
backward plan on SM100, plan selection will raise where it previously ran silently
non-deterministically, and the right answer becomes declining FROST under determinism instead.

get_attention_backend now returns 7 values instead of 6. All in-tree callers are updated.
dot_product_attention.py treats it as a plugin override point, so a plugin still returning 6
will raise. I did not add a compatibility shim: the tuple's arity has changed before without one
(6 from 5 in #1704), and it was silently reordered in #1542, which is a strictly more dangerous
change; the function is not exported from any public namespace; and the plugin hook itself is
recent and has no in-tree implementation or test. Happy to add a normalizing wrapper at the
override site if you consider the hook public.

Duplication against flex_attention.py. Worth stating precisely, since a fourth copy of
cuDNN plumbing is a reasonable thing to object to: fused_mla_q_uproj.py and grouped_mlp.py
do not build pygraphs at all, so this is the second PyTorch pygraph site, not the fourth. Four
lines are verbatim-identical across the two files, and the one genuinely duplicated unit is the
11-line per-device handle helper. Plan selection is deliberately opposite -- heur_mode.A plus
an explicit name-checked select_plan, versus A|FALLBACK with HEURISTICS_CHOICE -- because a
fallback plan at d512 is exactly the silently-non-FROST plan this backend must reject. I would
rather extract the handle helper as a follow-up than refactor a shipped, tested backend here.

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 also 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 FROST needs, 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. Instead NVTE_FROST_TEST_REQUIRED mirrors NVTE_GDN_TEST_REQUIRED, so a lane
meant to cover FROST fails loudly rather than skipping. It is not set in qa yet, since there is
no Blackwell L0 lane to set it on.

Documentation. NVTE_FROST_ATTN is now in docs/envvars.rst, and the stated backend
preference order is corrected -- it omitted FROST entirely. The backend tables in
docs/examples/attention/attention.ipynb are not updated yet; I would rather settle the
backend's name and status with you first, and note that the recently added flex-attention and
GDN backends shipped with no docs at all, so this is raising the bar rather than following it.

greptile-apps · 2026-09-16T04:38:27Z

Confidence Score: 5/5

The PR appears safe to merge; no new actionable failures remain after the latest test-coverage changes.

Summary

This PR adds a cuDNN FROST attention backend for symmetric head dimensions in (256, 512] on SM100/SM103, including context-parallel execution.

Adds FROST graph construction, plan caching, availability checks, and autograd integration.
Extends backend selection and context-parallel dispatch for p2p, all-gather, a2a, and a2a+p2p.
Adds independent numerical tests, context-parallel coverage, backend-selection regressions, and environment-variable documentation.
The changes since the previous review add FP16 coverage for both context-parallel execution and standalone backward gradients.

Diagram

%%{init: {'theme': 'neutral'}}%%
flowchart TD
  A[DotProductAttention] --> B[get_attention_backend]
  B -->|Eligible SM100/SM103 configuration| C[FrostAttention]
  C --> D{Context parallel?}
  D -->|No| E[FrostAttnFunc]
  E --> F[cuDNN FROST forward/backward]
  D -->|Yes| G[attn_forward_func_with_cp]
  G --> H[p2p or a2a+p2p]
  G --> I[all_gather]
  G --> J[a2a]
  H --> F
  I --> F
  J --> F

_{Reviews (25) · Last reviewed commit: "test(attention): cover fp16 in the backw..."}

Adds frost_attention.py, a thin PyTorch wrapper around the CuTe-DSL ("FROST") SDPA kernels in cuDNN Frontend >= 1.29.0. These are the only kernels that serve symmetric head_dim 512 forward and backward on SM100/SM103, a range no other backend covers: FlashAttention 2 and 3 cap at 256, FA4 is disabled at symmetric 512, and the C++ cuDNN fused path caps at 256. Measured on B200 before writing the integration, and each result shaped the code: - Correctness against the criterion FlashAttention applies to itself, err(kernel, fp64) <= 2 * err(naive_bf16, fp64): 0.21x to 0.94x across square and rectangular, causal and non-causal, GQA and MHA shapes. An absolute error is uninterpretable without that floor. - The forward LSE is natural-log logsumexp in fp32, matching an fp64 reference to 1.8e-06. This is what makes a context-parallel ring merge valid at all. - Outputs are bitwise reproducible across runs, ruling out a racing split-KV or atomic reduction. - Plan building costs ~1972 ms cold and ~12 ms once cuDNN caches the JIT, against a ~0.129 ms execute. Hence _PLAN_CACHE: at ~15000x an execute, caching is required rather than an optimisation. Design notes: - The cache holds compiled plans only, never output buffers. Buffers are allocated per call so a reused plan cannot make one call overwrite another's result, and with torch.empty_strided rather than empty_like, which does not preserve an arbitrary permuted stride. - Graphs are built from each tensor's ACTUAL strides, so bshd and sbhd are both served without a transpose. sbhd matters because Megatron uses it internally and copying every tensor per call would be a real cost. - _MASK_MODES lists only spellings verified behaviourally. cudnn sdpa() takes **kwargs and silently ignores names it does not recognise, so a typo would apply no mask and still build and run; inspect.signature is no help either, reporting no mask parameters at all. Both top-left and bottom-right causal are needed: the p2p ring produces square diagonal tiles where the two coincide, while all_gather trims KV and relies on bottom-right, where they differ by three orders of magnitude. - Unsupported configurations are refused rather than approximated, because the failure mode of guessing is silent numerical corruption, not an exception. Scope: SM100/SM103 only (the cuDNN d512 backward is Blackwell-only), bf16/fp16, symmetric head_dim in (256, 512], bshd and sbhd. thd needs varlen support that is feasible but not implemented here. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Signed-off-by: Nitin Vegesna <nvegesna@nvidia.com>

…tention Makes the FROST kernels reachable. Before this, get_attention_backend selected NO backend for symmetric head_dim 512 with context parallelism: FlashAttention and FusedAttention decline the head dim, and UnfusedDotProductAttention is disabled under CP. That combination raised rather than running, which is the gap this series closes. - get_attention_backend admits head_dim in (256, 512] on SM100/SM103 and returns a new use_frost_attention flag. It is consulted only where the established backends cannot run the shape, so it never displaces a faster path, and it is preferred over the unfused path, which covers the same shapes but cannot do CP. - FrostAttention and FrostAttnFunc in backends.py. TE selects a module class per backend, and there was none for cuDNN's Python kernels, so attn_forward_func_with_cp was unreachable for them. - FrostAttention is deliberately narrow: no FP8, bias, dropout, softmax offset or paging. Threading a flag through FusedAttention instead would have pulled FROST into all of that machinery; the selector declines those configurations first, so anything reaching the module is already supported. - FrostAttnFunc covers the non-CP path only. The CP path does not go through it because the ring must interleave per-step kernel calls with KV exchange and LSE correction rather than treating attention as one opaque autograd node. use_frost_attention is a separate flag rather than a FusedAttnBackend value: that enum mirrors NVTE_Fused_Attn_Backend value-for-value and is consumed by fused_attn_fwd, which dispatches into C++ that caps at 256, so routing FROST through it would feed a value into a path that cannot honour it. Two contracts worth noting, both of which produce runtime errors rather than type errors when missed: TE attention modules return heads flattened into the last dimension ([b, s, h*d]), and both return paths need it; and the "no backend is available" guard must count the new flag, or selecting FROST alone raises the very error this change removes. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Signed-off-by: Nitin Vegesna <nvegesna@nvidia.com>

… and a2a Dispatches the FROST kernels per step in all three CP comm types, which is what makes head_dim 512 usable for long-context training rather than only at CP=1. All three are needed for a real model. Gemma-4 dense is hybrid: sliding-window layers at head_dim 256 alongside global layers at 512. TE asserts that sliding-window attention requires a2a or all_gather, never p2p, so a p2p-only backend passes every attention test and still cannot run the target model. (MCore accepts cp_comm_type as a per-layer list, so a mixed configuration also works: sliding layers on all_gather, global layers on the cheaper p2p ring.) - p2p adds cp_p2p_{fwd,bwd}_frost_attn beside the existing fused and flash helpers. A ring step is only ever given causal, no_mask or a padding variant, so the dense case needs just causal on or off. - all_gather is simpler: KV is already gathered and trimmed, so each step is one call with no LSE correction. It does require BOTTOM-RIGHT causal, because get_kv_seq_info_after_all_gather trims KV and returns a window that is causal relative to the trimmed range. Top-left and bottom-right coincide only when SQ == SKV, which all_gather never produces, so the wrong choice here would be silent corruption rather than an error. - a2a is simplest of all: after the all-to-all each rank holds the full sequence for a subset of heads, so there is no ring and no correction. Ring tiles are slices and do not carry the strides the cuDNN graphs are built for, so every to_frost_layout call site passes contiguous tiles. This can copy; correctness first, worth revisiting if it shows up in a profile. Also relaxes an sbhd guard that inferred "not fused means flash". FROST is neither, and its graphs are built from actual strides, so sbhd is served directly; this matters because Megatron uses sbhd internally. The remaining instances of that inference are safe by construction (sliding windows and thd are already declined by the selector). Note for future work here: the three CP autograd classes are similar enough to invite generalisation and different enough to punish it. They do not carry identical ctx state, their aux_ctx_tensors differ in shape, and the tensor saved for backward is not always the value the branch returns. Adding a branch means checking what the enclosing function initialises and later consumes, not what the neighbouring class does. Each class also requires its backward to return one gradient per forward input; adding a parameter without the matching None breaks every existing user of that comm type, not just the new path. Verified on B200: {p2p, all_gather, a2a} x {bshd, sbhd} x {CP=2, CP=4} against the non-CP reference, plus 2 nodes x 2 ranks, with a FusedAttention regression control passing throughout. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Signed-off-by: Nitin Vegesna <nvegesna@nvidia.com>

Adds model_configs_frost_attn with Gemma-4 global-layer shapes (head_dim 512, GQA and MHA, causal and no_mask) and a FrostAttention kernel_backend in the CP runner. TE's existing CP matrix stops at head_dim 192, so nothing covered the range this backend exists for. The runner leaves NVTE_FLASH_ATTN and NVTE_FUSED_ATTN at 0 and sets NVTE_FROST_ATTN=1 rather than relying on fallthrough. FROST is the only backend serving head_dim > 256, so the selector would pick it either way, but making it an explicit kernel_backend keeps the test honest about what it exercises. Also updates the three call sites that unpack get_attention_backend for the added return value. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Signed-off-by: Nitin Vegesna <nvegesna@nvidia.com>

for more information, see https://pre-commit.ci

model_configs_frost_attn was defined but no pytest function parametrized over it, so the configs were only reachable by invoking run_attention_with_cp.py directly and CI would never have executed them. test_cp_with_frost_attention covers p2p / all_gather / a2a across bshd and sbhd. It skips rather than fails where the backend cannot run, reporting the reason from is_frost_attention_available(). That matters for the less obvious dependency: cuDNN Frontend declares nvidia-cutlass-dsl >= 4.6.2 but FROST enforces >= 4.7.0 at plan-build time, and below that floor every FROST engine silently declines and ordinary backend plans come back with no error. An environment without the dependencies, or without an SM100/SM103 GPU, therefore reports a skip with a reason rather than a failure that looks like a bug. thd and a2a+p2p are excluded because the backend declines them. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Signed-off-by: Nitin Vegesna <nvegesna@nvidia.com>

…s by device Two defects found in review. The availability check verified nvidia-cutlass-dsl but never the cuDNN Frontend version, and 1.29.0 is the first release carrying the head_dim=512 backward. The repo pins nvidia-cudnn-frontend>=1.28.0, and a 1.28.0 wheel does ship the d512 forward along with an importable cudnn.sdpa, so the import guard passed, the forward plan built, and the first backward raised mid-step. The plan-build error also named only cutlass-dsl, pointing users at the wrong package; it now reports both versions and their floors. The plan cache key omitted the device, so in a single-process multi-GPU run the same shape on a second device would reuse a graph built under the first while allocating tensors and workspace on the second. Both of TE's other cuDNN caches already guard against this: the C++ fused-attn cache keys on device_id to "distinguish graphs on different GPUs in a single-process run", and flex_attention keys its Python cache on device. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Signed-off-by: Nitin Vegesna <nvegesna@nvidia.com>

…version gate The device element added to _key() in 63294e8 made the key 12 items while _build_fwd and _build_bwd still unpacked 11, so the first FROST call of any kind raised ValueError: too many values to unpack. Every path was affected, forward and backward, CP and non-CP. Nothing caught it because the only test that reaches _key is gated on SM100/SM103. The unpack is now starred, so further device components cannot reintroduce the same break. The key also lacked device.type, letting a CPU tensor alias cuda:0, and called torch.cuda.current_device() unguarded for a normalisation that a materialized CUDA tensor never needs -- which would raise a confusing CUDA-init error on a CPU-only host. It now keys on (type, index) directly, following _score_mod_device_key in flex_attention.py. Version handling moves to packaging.Version over distribution metadata, matching _cudnn_frontend_version_supported in fused_mla_q_uproj.py. The hand-rolled parser accepted 1.29.0rc1 as 1.29.0, and by replacing the old int() parse it had quietly relaxed the cutlass floor to admit 4.7.0rc1 as well; both are rejected again. An undeterminable version no longer reports as "0" and hard-declines a valid source install -- it defers to _select_frost_plan, which checks the plan by name. That error message now looks both versions up defensively, since it previously could raise PackageNotFoundError while formatting the very diagnostic explaining a failure. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Signed-off-by: Nitin Vegesna <nvegesna@nvidia.com>

…ion probe Both FROST graphs declare v with k's shape and stride, and the plan-cache key records only q's and k's, so a v laid out differently from k would hit a plan built for k's layout and read the wrong elements with no error -- and in the backward, dv is allocated from v's own stride, disagreeing with the stride the graph declared. The forward checked shapes but not strides; the backward checked neither. Both now share one guard. Callers in TE always split k and v from a single QKV tensor, so this costs nothing and only closes a silent wrong answer. The version probe now returns the raw string alongside the parsed version, so "not installed" is distinguishable from "installed but unparseable". Previously both collapsed to None, which made an odd version string report as not installed and hard-decline a valid install -- the failure this was meant to remove. Only absence declines now; an unparseable version defers to the plan-name check, which is what the accompanying comment already claimed. The module fallback applies to that case too, and the plan-build error prints the raw string rather than a tuple. Also corrects a comment in backends.py stating frost_attention raises on any non-BSHD-contiguous layout. It does not: the graphs are built from each tensor's actual strides, and .contiguous() is there to keep one plan per shape. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Signed-off-by: Nitin Vegesna <nvegesna@nvidia.com>

…n value get_attention_backend gained a seventh return value, but the mixed-THD mask-policy path in _get_thd_policy_attention_backend still unpacked six, so every caller of that path raised ValueError: too many values to unpack. This had nothing to do with FROST -- it broke existing users of mixed-THD attention. The same function also rebuilt _attention_backends without use_frost_attention, leaving a stale value for the read at the scalar forward. Both are fixed, and the fake selector in test_mixed_thd_attention.py is updated to the same arity so it keeps matching the real signature rather than masking a mismatch. The CP runner now asserts that FrostAttention was the backend actually selected, not merely the one requested. The guarantee was previously emergent -- flash and fused are env-gated off and CP disables unfused, leaving FROST the only candidate -- so the assert passes by construction today. It is there so the tests fail rather than silently exercise another kernel if that ever stops holding. Docstrings drop the measured timings and error magnitudes. They were accurate, but no comment in transformer_engine/pytorch or transformer_engine/common cites figures like these; the fused-attention graph cache states the same constraint qualitatively. Benchmark numbers from one machine and one shape rot silently, so the constraints stay and the measurements live in the pull request instead. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Signed-off-by: Nitin Vegesna <nvegesna@nvidia.com>

…ong-answer paths The graphs were built and executed without a cuDNN handle, so cuDNN ran on its default handle's stream while the tensors and workspace were allocated on PyTorch's current stream, with nothing ordering the two. That is live on the path this backend exists for: the p2p ring issues attention inside `with torch.cuda.stream(cp_stream)`, so on alternating ring steps the kernel and its buffers sat on different streams. flex_attention.py and the C++ fused path both bind the stream explicitly; this now does the same, per device, rebinding on every call because one cached plan is executed from different streams. Validation now covers what the builders assume. Every node but stats is declared from q's dtype and execute() binds raw pointers, so a tensor of another dtype had its bits reinterpreted silently -- dout mattered most, since it arrives from autograd. k's batch and head_dim, out and dout's shapes, and softmax_lse's dtype and shape were likewise assumed and unchecked, and the backward additionally skipped the GQA divisibility check the forward has, which the CP ring can reach by calling it directly. Three configurations were selectable but unsupported, each a wrong answer rather than an error: CP with causal cross-attention or bottom-right masking, which the ring's square-tile chunking cannot serve and which both other CP backends already decline; return_max_logit, where this returns a bare tensor while the unfused path it displaces returns a pair; and load_balancing_strategy, which was dropped on the way to the ring and silently reverted to DUAL_CHUNK_SWAP. The first two now decline, the third is threaded through. KV caching declines explicitly -- it was already unreachable via the padding-mask assert, but only indirectly. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Signed-off-by: Nitin Vegesna <nvegesna@nvidia.com>

… names _handle_for created the handle inside a `with torch.cuda.device(...)` block but returned before cudnn.pygraph() was called, so graph construction and the CuTe-DSL plan build ran under whatever device happened to be current, with only the handle carrying the intended one. A JIT compile path is more likely to read the ambient CUDA context than the handle, and the guard costs nothing, so the build now happens under the device the cache key names. Unreachable from TE's own callers, which always run on the rank's own device, and flex_attention.py has the same shape -- this is hardening, not a fix for a live bug. Also corrects the rationale on the new context-parallel mask declines. It read as though any unequal q/kv length is wrong under CP, which would indict no_mask too; the restriction is specifically about where the causal diagonal sits, and no_mask stays allowed when the lengths differ. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Signed-off-by: Nitin Vegesna <nvegesna@nvidia.com>

…tself The CP tests compare a context-parallel run against a non-CP run of the same backend, so they validate the ring plumbing and nothing about the kernel: a wrong softmax scale, a causal mask anchored to the wrong corner, or an LSE in the wrong log base appears identically on both sides and cancels. FrostAttnFunc, the path a single-GPU head_dim 512 user takes, had no coverage at all. test_frost_attention.py checks forward output, the LSE convention and the backward gradients against an fp32 reference computed independently of TE and of cuDNN, over both causal alignments, both dtypes, GQA and MHA, and a rectangular shape where top-left and bottom-right masking differ. The bar is the criterion FlashAttention applies to itself -- error within 2x what the reference itself incurs from reduced-precision inputs -- measured per case rather than hard-coded, so it tracks the shape instead of encoding a number that rots. Inputs are generated in fp32 and cast down, because rounding an already-rounded tensor would collapse that floor to zero. It also covers the decline paths and the k/v mismatch guards. Registered in qa/L0 alongside the sibling backends. Separately, the availability probe now runs after the shape and dtype checks rather than before. Probing imports cuDNN Frontend and sets CUDNN_FRONTEND_ENABLE_FROST_ENGINES, which registers engines process-wide and is therefore visible to flex_attention and the GDN path. That happened for every attention configuration on any Blackwell machine, at any head dim, including the overwhelming majority nowhere near 512. It now happens only for a configuration FROST could actually serve. An explicit CUDNN_FRONTEND_ENABLE_FROST_ENGINES=0 also declines cleanly instead of raising later from plan selection. Finally, the claim that TE's C++ fused path caps at 256 was wrong: that dispatch applies no head-dim test and simply asks cuDNN for a graph, so the ceiling is cuDNN's engine coverage. Stated correctly, along with the actual reason a Python backend is required -- FROST engines register at Python import time and need nvidia-cutlass-dsl, while TE's C++ builds against frontend headers only. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Signed-off-by: Nitin Vegesna <nvegesna@nvidia.com>

… the backend cuDNN's SDPA backward is non-deterministic unless the graph asks otherwise -- that is why the C++ fused path calls set_deterministic_algorithm and why flex_attention passes use_deterministic_algorithm to the same sdpa_backward this module builds. FROST passed neither, so NVTE_ALLOW_NONDETERMINISTIC_ALGO=0 was silently not honoured while every other backend either honours it or declines. The flag is now threaded from DotProductAttention through FrostAttnFunc and the three context-parallel backward wrappers into the graph, and it is part of the plan-cache key: the deterministic backward is a different algorithm, so a plan built one way must not serve a call that asked for the other. The availability probe gains an NVTE_FROST_TEST_REQUIRED escape hatch, mirroring NVTE_GDN_TEST_REQUIRED, so a lane intended to cover this backend fails loudly instead of skipping silently. It is deliberately not set in qa yet, since no Blackwell L0 lane exists to set it on. Documents NVTE_FROST_ATTN in docs/envvars.rst, placed by that file's backend-preference ordering rather than alphabetically, and corrects the stated preference order, which omitted FrostAttention entirely. FROST sits between FusedAttention and UnfusedDotProductAttention and is only ever eligible in the (256, 512] head_dim band that flash and fused do not serve, so it never displaces a backend that could otherwise have run. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Signed-off-by: Nitin Vegesna <nvegesna@nvidia.com>

…lism UnfusedDotProductAttention also serves symmetric head_dim in (256, 512] -- there is no head-dim filter against it anywhere -- so calling FrostAttention the only backend for that range was wrong, and would tell a user without context parallelism that they need a Blackwell-only dependency stack they do not. FrostAttention is the only backend for that range *with* context parallelism; without it, unfused covers the same shapes and FROST is merely preferred. The same paragraph also claimed FrostAttention never displaces a backend that could otherwise have run, which is wrong in the other direction: it suppresses UnfusedDotProductAttention when both are eligible. Says so now. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Signed-off-by: Nitin Vegesna <nvegesna@nvidia.com>

… p2p wrapper Threading deterministic into the FROST backward wrappers used a match on the trailing out_part/dout_part/section parameters, which is not unique to the FROST one: cp_p2p_bwd_fused_attn ends the same way and already took deterministic positionally. It therefore gained a second, keyword copy and the module stopped compiling, taking all of transformer_engine.pytorch down with it. Not caught before pushing because the syntax check used ast.parse, which parses a duplicate argument happily -- CPython only rejects it when building the symbol table in compile(). Verified now with compile() across every file the branch touches, plus an AST scan for any repeated parameter name. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Signed-off-by: Nitin Vegesna <nvegesna@nvidia.com>

Measured on B200 with cuDNN Frontend 1.29.0: asking sdpa_backward for a deterministic algorithm is refused outright -- cudnnGraphNotSupportedError, no engine proposes a plan for the graph. So unlike the C++ fused path, which opts in via set_deterministic_algorithm, there is nothing here to opt into, and passing the flag alone would turn NVTE_ALLOW_NONDETERMINISTIC_ALGO=0 from a silent violation into a hard failure at plan build. The selector now declines FROST when determinism is required during training, which is what the other backends do where they cannot honour it. The graph still passes use_deterministic_algorithm, so the decline lifts on its own if cuDNN ships a deterministic d512 backward. The context-parallel tests are unaffected: their runner sets NVTE_ALLOW_NONDETERMINISTIC_ALGO=1 explicitly. Also fixes the k/v layout guard test, which could never have failed: it built the mismatched v as a [b, h, s, d] contiguous tensor, whose strides are exactly those of a contiguous k, so there was nothing to reject. It is now built as sbhd and permuted, which keeps the shape and the contiguous head dimension while genuinely differing in stride order. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Signed-off-by: Nitin Vegesna <nvegesna@nvidia.com>

…not a reference The oracle compared the kernel against an fp32 reference, and on Ampere and newer torch computes fp32 matmuls in TF32. TF32's significand is 11 bits -- exactly fp16's -- so for fp16 inputs the "reference" was no more accurate than the kernel it was judging. Rounding the inputs to fp16 then changed the reference almost not at all, and the measured error floor collapsed from about 1e-03 to 3e-08, reducing the bound to the bare absolute slack. That is how it presented on B200: all seven failures were fp16 with a causal mask, where the kernel's error is a perfectly normal 1.55e-03 but the bound had become 1e-03. bf16 was unaffected because its 8-bit significand is far coarser than TF32, so its floor stayed honest -- which is exactly why the flaw looked like an fp16-specific kernel problem rather than a broken reference. The reference is now float64 throughout, immune to TF32 and to whatever the ambient precision flags are. With it the fp16 causal floor returns to 1.49e-03 and the bound to 3.99e-03, comfortably above the observed error, while a deliberate 1% scale error is still rejected in every dtype and mask combination. Tests renamed accordingly, since they no longer compare against fp32. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Signed-off-by: Nitin Vegesna <nvegesna@nvidia.com>

…ding window The backend allowed exactly three mask spellings from a hardcoded table and declined every sliding window. Neither restriction was necessary: cuDNN's own engine descriptor for sdpa_bwd_sm100 declares swa and right_band_widening, and the legacy spellings are not a separate mechanism at all -- pygraph/sdpa.cpp desugars use_causal_mask to (TOP_LEFT, right_bound=0) and use_causal_mask_bottom_right to (BOTTOM_RIGHT, right_bound=0), and refuses to combine either with an explicit right bound. Masking is therefore built the way the C++ fused path and the in-flight Python port both build it: a diagonal alignment plus a two-sided band. Causal, bottom-right and sliding window come from one mechanism instead of three spellings, the window travels in the plan-cache key, and the all-gather path no longer raises on a window it can now serve. The old justification for the allowlist was also wrong. It claimed sdpa() silently ignores unknown kwargs, so a misspelling would apply no mask and still run. sdpa is a pybind function with an explicit named-argument list and no kwargs catch-all; an unknown keyword raises TypeError. The error is deferred to plan creation rather than raised at validate, which is presumably where the belief came from, but it is loud, not silent. Separately, head_dim is now required to be a multiple of 8. The engine pads to that multiple, so 260 sat inside the advertised (256, 512] range, passed the gate, and then failed at plan selection complaining about missing engines instead of declining cleanly. The oracle test gains sliding-window cases against the float64 reference, including an assertion that a window changes the output -- a dropped bound would otherwise still produce finite, plausible numbers. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Signed-off-by: Nitin Vegesna <nvegesna@nvidia.com>

…for p2p Allowing sliding window opened two context-parallel paths that could not serve it. The a2a helpers took no window and so ran plain causal attention with the left bound silently dropped -- finite, plausible output and wrong gradients. The p2p ring cannot serve it at all, because a left bound measured against the full sequence does not survive the per-step KV tiles. a2a now carries the window, which matches what it can actually do: after the all-to-all each rank holds the full sequence for a subset of heads, so the user's window applies unchanged. p2p and a2a+p2p decline, which is the same rule FusedAttention already carries a few hundred lines above, for the same reason. all_gather was already correct. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Signed-off-by: Nitin Vegesna <nvegesna@nvidia.com>

…y, and per cp_comm_type The window was exercised only in the forward, yet it changes the backward graph's dK/dV accumulation rather than just a mask fill, so dq/dk/dv under a window were entirely unvalidated. The backward test now parametrizes over it. Adds window=(0,0), the boundary of cuDNN's convention: left_bound counts visible tokens including the diagonal and has a documented minimum of 1, so this is the value where an off-by-one stops producing wrong numbers and starts producing an error instead. Adds the window-validation cases to the decline test, which were unreachable from the suite even though is_frost_attention_supported accepts and routes the argument. Adds a selector test for the rules the previous commit introduced, which shipped untested: all_gather and a2a may serve a window, p2p and a2a+p2p decline it, and configurations without a real window must still select FROST under p2p -- the decline has to key on the window rather than on p2p itself. Also corrects docs/envvars.rst, which still said the backend declines sliding window, and which omitted both the multiple-of-8 head_dim constraint and the determinism decline. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Signed-off-by: Nitin Vegesna <nvegesna@nvidia.com>

Enabling sliding window made two pre-existing assertions reachable that had never heard of this backend. Both the all_gather and a2a forwards allow a window only if FusedAttention or some FlashAttention is in play, and when FROST is selected every one of those flags is False -- so the assert survived solely on fa_utils.v2_3_plus, which reports whether flash-attn happens to be installed rather than which backend is running. Sliding window with all_gather or a2a would therefore fail or pass on an unrelated package, on exactly the Blackwell d512 box this backend exists for. Both allowlists and both messages now include FROST. Also declines a windowed non-causal mask when the q and kv lengths differ. FROST anchors the band from the mask type, so that case always lands top-left, while TE's bottom_right_diagonal defaults to True and the C++ fused path picks the alignment from it -- the two would disagree silently. Declining is better than guessing the anchor. The remaining fixes are gate hygiene. window_size now rejects a left below -1, which would otherwise build diagonal_band_left_bound=-1 and fail at plan build, and a non-iterable window declines instead of raising TypeError out of backend selection, which is not an exception the selector catches. window_size moves after deterministic in frost_attn_bwd so an existing positional caller cannot silently reinterpret one as the other. Tests cover the new gates, including the head_dim multiple-of-8 rule, which had none. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Signed-off-by: Nitin Vegesna <nvegesna@nvidia.com>

The docstring asserted that FROST engines register at cudnn import time, so a setdefault running after another module had already imported cudnn would be too late and leave no FROST engine. That is what the documentation implies, and it was the basis for a concern about the in-flight port of cuDNN attention to the Python API, whose shared import helper sets no such variable. Measured on B200 with cuDNN Frontend 1.29.0 and it does not hold: importing cudnn and cudnn.sdpa first with the switch unset, then setting it and building a plan, still selects a FROST engine. The switch is still set before the import, because that is what the documentation asks for and it costs nothing, but nothing depends on winning the race and the plan-name check verifies the engine either way. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Signed-off-by: Nitin Vegesna <nvegesna@nvidia.com>

The float64 reference skipped masking entirely whenever attn_mask_type was "no_mask", so a window passed alongside it was ignored. That is not TE's rule: the SWA construction in utils.py applies the window to any mask type, treats -1 as unbounded on that side, and lets a causal mask type pin the right bound to the diagonal. no_mask with (w, 0) is therefore a causal band of width w, not an unmasked attention. Caught by the B200 run, which reported dq off by 2.9 against a reference maximum of 0.83 for exactly the two no_mask windowed backward cases, while every causal windowed case passed -- the signature of a reference that is masking differently rather than a kernel that is computing wrongly. The reference now derives blocking from (left, right) plus the diagonal offset, which reproduces TE's keep rule for every mask type and window combination that carries a window. The previous form is unchanged for unwindowed causal and bottom-right, so the cases already verified on hardware keep their meaning. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Signed-off-by: Nitin Vegesna <nvegesna@nvidia.com>

…, not precedent The comment said FusedAttention carries the same rule and left the reason as an assertion about per-step tiles. The actual reason is visible in the p2p path: it hardcodes the per-step window to (-1, 0) or (-1, -1) at every kernel call, so a user window is discarded there regardless of backend. all_gather by contrast computes window_size_per_step through get_kv_seq_info_after_all_gather and passes it down, and a2a sees the whole sequence after the all-to-all. Citing that is checkable; citing another backend's rule is not. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Signed-off-by: Nitin Vegesna <nvegesna@nvidia.com>

…s silently ignored DotProductAttention.forward takes a separate branch under ONNX export that skips get_attention_backend and binds the backend flags itself. Adding a seventh flag without binding it there made the availability check below read an unassigned local: UnboundLocalError on every ONNX export, on every GPU, at any head dim, and NVTE_FROST_ATTN=0 does not help because that branch never consults it. This is the one defect in the series that reaches users who will never touch head_dim 512. Five capabilities were silently ignored rather than declined, all reachable because at head_dim 512 the fused path is unavailable and FROST becomes the sole survivor of filters written to disable everything: score_mod (including the score_mod_bprop-without-score_mod case, which is meant to end in "no backend available" and was instead being rescued into a wrong answer), a quantized qkv_type carrying a nominal bf16 dtype outside an fp8 autocast, num_splits, checkpoint_core_attention, and CUDA graph capture. Each now declines, matching what the neighbouring filters do for the other backends. Lint: the new module scored 8.91 against the repo's own pylint gate, which does not disable consider-using-f-string, while the sibling flex_attention.py scores 10.00. All thirty percent-format sites are now f-strings, verified by comparing the rendered decline messages before and after. Also clears the regressions this branch introduced elsewhere -- an unused-argument pair, a used-before-assignment that three branches made unprovable, and a condition one clause over the limit. Every touched file is back to 10.00 under the pinned pylint and CI's Python. Removes a deterministic parameter accidentally added to cp_p2p_bwd_flash_attn, which never read it, and asserts in the all_gather window helper the invariant that the selector enforces in another file. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Signed-off-by: Nitin Vegesna <nvegesna@nvidia.com>

…local The new qkv_type decline read a local that get_attention_backend rebinds far earlier: the fused-attention dtype spec assigns qkv_type, o_type, do_type and dqkv_type from spec, so by the time the FROST guards run the name holds an NVTE dtype enum rather than the tensor class. Comparing that against torch.Tensor is unequal for every input, so FROST was declined unconditionally -- the selector reported "Disabling FrostAttention for qkv_type = 6" and no backend at all for head_dim 512. Reading attention_params.qkv_type is unambiguous and cannot be shadowed. Audited the other names these guards read for the same hazard: only window_size is also rebound, at the check_set_window_size normalisation, which is the canonical value every neighbouring filter uses and is the right one to read. Caught by test_frost_sliding_window_selection_by_cp_comm_type, which exists because a reviewer pointed out the selector rules had no coverage at all. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Signed-off-by: Nitin Vegesna <nvegesna@nvidia.com>

…n it The ONNX fix has never executed. The section meant to verify it ran test_onnx_export.py, which imports onnxruntime -- absent from the container -- so it failed at collection and proved nothing. The bug was an UnboundLocalError, not anything about ONNX serialization: the export branch skips get_attention_backend and binds the backend flags by hand, and the availability check below it reads all of them. Entering export mode and running one ordinary head_dim-64 attention reproduces it without onnxruntime. That test must not be Blackwell-gated -- the bug hit every user on every GPU -- so the module-level pytestmark becomes a named decorator applied to the six tests that genuinely need FROST, leaving the new one to run wherever there is a CUDA device. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>

nvegesna-netizen · 2026-09-17T01:31:30Z

Verified end-to-end on B200 (SM100) at e5071dec, cudnn-frontend 1.29.0 / cutlass-dsl 4.7.1:

Numerics — 59 cases against an independent float64 reference, forward and backward, across head_dim 320/512, bf16/fp16, and the mask and sliding-window matrix. The bar is the one FlashAttention applies to itself: error within 2x the floor the reference itself incurs from reduced-precision inputs, measured per case rather than hard-coded. float64 rather than float32 because TF32's significand is 11 bits, the same as fp16, so an fp32 reference cannot judge an fp16 kernel.
Context parallelism — 18 CP tests pass (all_gather, a2a), plus a FusedAttention CP control to confirm nothing regressed for existing users.
Stream safety — the cuDNN handle is per-device with the stream bound per call; a FROST execute issued on a non-default stream is bitwise identical to the default-stream result over 50 repeats.
Determinism — cuDNN refuses a deterministic backward, and the selector declines it rather than silently returning a non-deterministic one.
No collateral — broad attention suite 388 passed / 0 failed; mixed-THD 33 passed.

One gap worth naming: tests/pytorch/test_torch_compile.py cannot run in the 26.08 container — it dies at collection on MXFP8Quantizer is not registered as an opaque type, from the torch opaque-object API move (the registration in dynamo/quantizer_opaque.py is swallowed by design to keep TE importable). It fails identically on main, so it is unrelated to this PR, but it does mean the fullgraph=True tracing of get_attention_backend is covered here only by the traceable unit test, not by that suite. If there is a lane with a matching torch, that is the useful one to run.

The decline said a2a+p2p "is not wired up". It is. context_parallel.py dispatches `cp_comm_type in ["p2p", "a2a+p2p"]` to the same AttnFuncWithCPAndKVP2P and passes use_frost_attention into it, where FrostAttention is called at all four forward section sites and all four backward sites. The a2a stage is flash_attn_a2a_communicate: a redistribution between sequence- and head-sharding that invokes no attention kernel. Under a2a+p2p the per-step calls are therefore the ordinary p2p section calls with fewer heads per rank. What was actually true is that it was untested. a2a+p2p needs four ranks, an a2a subgroup crossed with a p2p subgroup, and every FrostAttention CP arm ran on a pool of two. Declining an untested path is defensible; describing it as unwired was not, and it would have misled anyone deciding whether to enable it. The sliding-window decline for a2a+p2p stays and is unrelated: it rings across sub-groups, so a window measured against the full sequence still does not survive the per-step tiles, exactly as with plain p2p. Test coverage extends to four ranks for this case only, and asserts the a2a divisibility requirement rather than relying on the current configs happening to satisfy it. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>

cp_comm_type="a2a+p2p" passes cp_group as [a2a_group, p2p_group]. FrostAttention computed context_parallel with a one-liner that assumed a single group, so get_distributed_world_size received a list and raised TypeError: unhashable type: 'list' at backends.py in FrostAttention.forward, before any attention ran. The signature already declared Optional[Union[dist_group_type, List[dist_group_type]]]; the body did not honour it. Now the same form FlashAttention and FusedAttention use a few hundred lines above and below: multiply the sub-group sizes when a list arrives. Found by enabling a2a+p2p and running it, after the previous commit claimed on code-reading grounds that the path was already complete. It was reachable, but it crashed on the first line of the forward. All six a2a+p2p arms failed deterministically while the eighteen existing arms passed. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>

…lism The backend serves BF16 and FP16, but coverage was uneven: only the forward numerics ran both dtypes. The backward and all 24 context-parallel configurations were bf16 only. That is the wrong way round. fp16 has a far narrower exponent range than bf16, and the two places it would show first are exactly the two that were untested: the gradient of a softmax subtracts similarly sized terms, and the ring correction exponentiates a difference of log-sum-exp values across steps. The backward test now runs both dtypes. Context parallelism gains one fp16 arm per comm type rather than a doubled matrix -- one model, one layout, three cases. a2a+p2p is omitted from the fp16 arm because it would need a second four-rank pool for a dtype that exercises no additional code path. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>

github-actions Bot added the community-contribution PRs from external contributor outside the core maintainers, representing community-driven work. label Sep 16, 2026

nvegesna-netizen mentioned this pull request Sep 16, 2026

feat(gemma4): allow CP > 1 on the dense model by building its layer spec with TE NVIDIA-NeMo/Megatron-Bridge#6120

Draft

nvegesna-netizen marked this pull request as ready for review September 16, 2026 04:35

nvegesna-netizen requested a review from cyanguwa as a code owner September 16, 2026 04:35

greptile-apps Bot reviewed Sep 16, 2026

View reviewed changes

Comment thread transformer_engine/pytorch/attention/dot_product_attention/frost_attention.py Outdated

Comment thread transformer_engine/pytorch/attention/dot_product_attention/frost_attention.py Outdated

Comment thread tests/pytorch/attention/test_attention_with_cp.py

nvegesna-netizen and others added 8 commits September 15, 2026 23:21

[pre-commit.ci] auto fixes from pre-commit.com hooks

bc15d3a

for more information, see https://pre-commit.ci

nvegesna-netizen force-pushed the nvegesna/te-frost-d512-cp branch from 7d453a4 to 4069bfa Compare September 16, 2026 06:22

nvegesna-netizen and others added 6 commits September 15, 2026 23:29

greptile-apps Bot reviewed Sep 16, 2026

View reviewed changes

Comment thread docs/envvars.rst Outdated

nvegesna-netizen and others added 5 commits September 16, 2026 00:47

greptile-apps Bot reviewed Sep 16, 2026

View reviewed changes

Comment thread transformer_engine/pytorch/attention/dot_product_attention/utils.py

nvegesna-netizen and others added 2 commits September 16, 2026 09:43

nvegesna-netizen and others added 7 commits September 16, 2026 09:52

nvegesna-netizen and others added 3 commits September 17, 2026 01:47

Navigation Menu

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

feat(attention): cuDNN FROST attention backend for head_dim in (256, 512], with context parallelism#3527
nvegesna-netizen wants to merge 31 commits into
NVIDIA:mainfrom
nvegesna-netizen:nvegesna/te-frost-d512-cp

nvegesna-netizen commented Sep 16, 2026 •

edited

Loading

Uh oh!

greptile-apps Bot commented Sep 16, 2026 •

edited

Loading

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

nvegesna-netizen commented Sep 17, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant

Conversation

nvegesna-netizen commented Sep 16, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Problem

Approach

Verification (B200)

Performance

Notes for reviewers

Scope

Known limitations and open decisions

Uh oh!

greptile-apps Bot commented Sep 16, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Confidence Score: 5/5

Summary

Diagram

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

nvegesna-netizen commented Sep 17, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant

nvegesna-netizen commented Sep 16, 2026 •

edited

Loading

greptile-apps Bot commented Sep 16, 2026 •

edited

Loading