Log when thd with dropout falls to the composite cuDNN engine - #3313
Log when thd with dropout falls to the composite cuDNN engine#3313bzantium wants to merge 1 commit into
Conversation
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 <ryumin93@gmail.com>
github-actions
Bot
added
the
community-contribution
Greptile SummaryThis PR adds a debug advisory for
Confidence Score: 4/5The 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
Reviews (1): Last reviewed commit: "Log when thd with dropout falls to the c..." | Re-trigger Greptile |
![greptile-apps[bot]](https://avatars.githubusercontent.com/in/867647?s=60&v=4){kind=small}
| 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" | ||
| ) |
There was a problem hiding this comment.
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!
There was a problem hiding this comment.
Remove the current early debug advisory from the dropout-filtering section. It runs before final backend selection and can therefore report FusedAttention even when a later filter selects another backend.
| 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" | |
| ) |
|
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! |
@bzantium thanks for reporting this. Please set NVTE_FUSED_ATTN=0 NVTE_FLASH_ATTN=1 (hopefully this unblocks you temporarily)
For THD with dropout 0.1, switching to FA2 reduced time from 4.407 ms to approximately 0.728 ms, about a 6.1× speedup.
To verify that FA2 is selected, temporarily use: The log should report something like this: NOTE: I am suggesting NVTE_UNFUSED_ATTN=0 only as a validation guard to prevent silent fallback. |
| 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" | ||
| ) |
There was a problem hiding this comment.
Remove the current early debug advisory from the dropout-filtering section. It runs before final backend selection and can therefore report FusedAttention even when a later filter selects another backend.
| 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" | |
| ) |
| # Select FusedAttention for performance | ||
| if use_flash_attention and use_fused_attention and device_compute_capability >= (9, 0): | ||
| logger.debug( | ||
| "Disabling FlashAttention to give FusedAttention preference on Hopper+ " | ||
| "for performance reasons" | ||
| ) | ||
| use_flash_attention = False |
There was a problem hiding this comment.
Replace the existing Hopper+ FusedAttention preference block with the following narrow selection rule. This safely prefers FA2 only when it is confirmed installed, enabled, and eligible. Otherwise, FusedAttention remains enabled.
# Prefer FA2 for THD training with dropout on SM100/103, where FusedAttention has a known
# performance issue. At this point use_flash_attention_2 confirms a usable installation.
if (
is_training
and qkv_format == "thd"
and attention_dropout != 0.0
and device_compute_capability in ((10, 0), (10, 3))
and use_flash_attention_2
and use_fused_attention
):
logger.debug(
"Disabling FusedAttention to give FlashAttention 2 preference for THD with dropout on SM100/103"
)
use_fused_attention = False
fused_attention_backend = None
# Select FusedAttention for performance in all other Hopper+ configurations.
elif use_flash_attention and use_fused_attention and device_compute_capability >= (9, 0):
logger.debug(
"Disabling FlashAttention to give FusedAttention preference on Hopper+ "
"for performance reasons"
)
use_flash_attention = False
Please place this after unavailable FlashAttention installations have been filtered and use_flash_attention_2 has been finalized. This placement is important because it ensures that FusedAttention is disabled only when FA2 is actually usable.
|
@bzantium Here's a quick update: I discussed with the cuDNN team and it's a known shortcoming. The issue is not unified v/s composite but it is the slow dropout kernels used for this very specific combination being discussed here. The fact that the slow dropout kernels are being triggered in a composite setup worsens the timing. cuDNN does not have this on their roadmap for the foreseeable future as it is low priority and the suggested "fix" is for TE users to use FA2 instead if performance is of importance. If you'd like to contribute to TE, which we highly encourage, I've reviewed your PR and added suggested changes. Please review these changes and commit them. Once you've done this I can help you launch CI on this PR and review/approve the PR to have it merged in to main. On the contrary, if you'd rather have me add in these changes, I'm happy to do so in a separate PR. Let me know. Notes:
|
|
Lastly, here's a small mock up test I came up with using my agent. I just asked it to write a test to confirm the backend choice. This is juts a guidance but it does the trick. Feel free to be creative :) |
3 participants
What does this PR do?
Adds a
logger.debugline for the case whereqkv_format="thd"andattention_dropout > 0, which is served by the composite cuDNN engine rather than the unified one.Related to #3312.
Why
fused_attn_f16_arbitrary_seqlen.cualready 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 supportcu_seqlens:The routing is correct, but it is expensive and silent. Forward + backward through one
DotProductAttention, 4096 tokens, 16 heads,head_dim128, bf16, median of 50 iterations after 20 warmup:thdsbhdsbhdA profile attributes it to
cudnn::fusion::gen_dropout_mask_4bitand 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=2output sees why a backend was not chosen. This case is different: the backend is chosen and quietly costs several times more. Most configurations setattention_dropoutto 0 and never see it, but Megatron-Core'sTransformerConfigdefaults 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) andpylint --rcfile=pylintrcclean on the changed file. No behavioural change, so no new tests; the line appears in existingNVTE_DEBUG_LEVEL=2output when the condition holds.