Log when thd with dropout falls to the composite cuDNN engine

[bzantium]

What does this PR do?

Adds a logger.debug line for the case where qkv_format="thd" and attention_dropout > 0, which is served by the composite cuDNN engine rather than the unified one.

Related to #3312.

Why

fused_attn_f16_arbitrary_seqlen.cu already notes that dropout and stats generation cannot be combined on the unified engine, so a thd request with dropout is routed to the composite engine, which does not support cu_seqlens:

      // This extra restriction is needed because cuDNN frontend doesn't yet allow
      // the combination of dropout and stats generation for the fprop unified engine,
      // so any such request would always get routed to the old composite SDPA engine
      // (which doesn't support cu_seqlens). Remove this restriction when possible.
      !is_dropout;

The routing is correct, but it is expensive and silent. Forward + backward through one DotProductAttention, 4096 tokens, 16 heads, head_dim 128, bf16, median of 50 iterations after 20 warmup:

GPU	layout	p=0.0	p=0.1	dropout cost
B300 (sm103, TE 2.14.1)	thd	0.585 ms	3.060 ms	+2.476 ms (5.24x)
B300 (sm103, TE 2.14.1)	sbhd	0.480 ms	0.525 ms	+0.045 ms (1.09x)
H200 (sm90, TE 2.15.0)	sbhd	0.594 ms	0.614 ms	+0.020 ms (1.03x)

A profile attributes it to cudnn::fusion::gen_dropout_mask_4bit and its transpose variant, which together take 43% of CUDA time in an 8-layer training step on B300.

Backend selection already logs every case where a backend is disabled, so a user reading NVTE_DEBUG_LEVEL=2 output sees why a backend was not chosen. This case is different: the backend is chosen and quietly costs several times more. Most configurations set attention_dropout to 0 and never see it, but Megatron-Core's TransformerConfig defaults it to 0.1, so a packed run that does not set it explicitly inherits the slow path with nothing in the log to suggest it.

This does not change behaviour — it only makes the situation visible. The underlying fix belongs to the cuDNN frontend restriction quoted above.

Testing

black (repo settings) and pylint --rcfile=pylintrc clean on the changed file. No behavioural change, so no new tests; the line appears in existing NVTE_DEBUG_LEVEL=2 output when the condition holds.

[greptile-apps]

Greptile Summary

This PR adds a debug advisory for thd FusedAttention requests with dropout, intended to expose their slower composite-cuDNN execution path.

• Logs the composite-engine performance warning during dropout filtering.
• Does not alter backend selection or attention behavior.

Confidence Score: 4/5

The PR is safe to merge, though the new advisory should be moved or gated on final FusedAttention selection to avoid misleading diagnostics.

Later backend filters can reject FusedAttention after the new message claims its composite cuDNN engine is in use, while final execution proceeds through another backend.

Files Needing Attention: transformer_engine/pytorch/attention/dot_product_attention/utils.py

Important Files Changed

Filename	Overview
transformer_engine/pytorch/attention/dot_product_attention/utils.py	Adds the intended diagnostic, but emits it before later filters establish that FusedAttention will actually be selected.

_{Reviews (1): Last reviewed commit: "Log when thd with dropout falls to the c..." | Re-trigger Greptile}

[greptile-apps]

Advisory precedes backend selection

The advisory runs before later filters can reject FusedAttention, so configurations such as thd with dropout and an arbitrary mask log that the composite cuDNN engine is in use even though another backend is ultimately selected. This sends users toward the wrong performance diagnosis; emit the advisory only after final FusedAttention selection.

Knowledge Base Used: PyTorch Attention Stack

Note: If this suggestion doesn't match your team's coding style, reply to this and let me know. I'll remember it for next time!

[cyanguwa]

Thanks for the detailed reporting in #3312!

I agree that this affects users' performance and without warning, but I don't think it's TE's responsibility to cache these default-value issues coming from Megatron. Given that dropout is going out of fashion, Megatron can probably change their default to 0 but in a gradual deprecation code cycle. This way, it wouldn't incorrectly and silently call FlashAttention for dropout either.

With cuDNN, we can check with them to see if they can speed up the dropout implementation on Blackwell for THD, to match FlashAttention's performance. For now, we can disable FusedAttention for THD + dropout + Blackwell, and allow users to use FlashAttention instead. @KshitijLakhani, could you please lead the discussion with cuDNN, file a bug if necessary, and guide @bzantium to disable FusedAttention in this PR? Thanks!