From e6210349a19bf4c0d4bab47cc831dfabcb267432 Mon Sep 17 00:00:00 2001 From: Minho Ryu Date: Tue, 4 Aug 2026 19:39:46 +0900 Subject: [PATCH] Log when thd with dropout falls to the composite cuDNN engine Dropout keeps a thd request off cuDNN's unified engine, and the composite engine it lands on instead generates the dropout mask in separate kernels. On sm103 that is 5x the no-dropout cost for the same attention, and nothing in the backend selection log says so. Measured at 4096 tokens, 16 heads, head_dim 128, bf16, forward+backward: thd 0.585 ms at p=0 against 3.060 ms at p=0.1, while sbhd goes 0.480 ms to 0.525 ms for the same dropout. Signed-off-by: Minho Ryu --- .../pytorch/attention/dot_product_attention/utils.py | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/transformer_engine/pytorch/attention/dot_product_attention/utils.py b/transformer_engine/pytorch/attention/dot_product_attention/utils.py index 9ee6ad0101..c31d27cc53 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/utils.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/utils.py @@ -1082,6 +1082,13 @@ def _is_fa3_supported(num_heads, num_gqa_groups, head_dim_qk, head_dim_v, qkv_dt if use_flash_attention_4 and FlashAttentionUtils.v4_is_installed: logger.debug("Disabling FlashAttention 4 for dropout") use_flash_attention_4 = False + if use_fused_attention and qkv_format == "thd": + # Dropout keeps thd off cuDNN's unified engine, so it falls to the much slower + # composite one. Nothing else reports this. + logger.debug( + "FusedAttention with dropout and qkv_format = thd uses the composite cuDNN" + " engine, which is much slower than the unified engine" + ) # Filter: Softmax type # context_parallel | softmax_type | supported backends