[liayan]

Description

On cuDNN 9.26, NVTE_CUDNN_MXFP8_NORM_OUTPUT_IN_INPUT_DTYPE=1 makes CudnnNormalizationPlan type the MXFP8 RMSNorm virtual output as the input dtype (BF16). Heuristics then die with SIGSEGV while planning that graph at exactly 2048×7168, forward training (row-wise and column-wise block-scale quantize). create_execution_plans never returns, so no norm kernel runs. There is no Python exception: the process exits -11 after the warning in common.cpp.

Single GPU, GB200, cuDNN 92600, Transformer Engine from the image that prints that warning at common.cpp:312. Each case is LayerNormLinear(hidden, 128, normalization="RMSNorm"), BF16, MXFP8, cuDNN norm on in both directions, in its own process.

Failing case:

=== BEGIN 2048x7168 flag=1 train ===
cudnn 92600 train True
/workspace/transformerengine/transformer_engine/common/normalization/common.cpp:312 in function CudnnNormalizationPlan: The cuDNN MXFP8 normalization intermediate output uses the input dtype (itype) instead of the compute dtype; otype still applies to the final quantized output.
=== END 2048x7168 flag=1 train exit=-11 ===

Same call at 4096×7168 prints the same warning and finishes:

=== BEGIN 4096x7168 flag=1 train ===
cudnn 92600 train True
/workspace/transformerengine/transformer_engine/common/normalization/common.cpp:312 in function CudnnNormalizationPlan: The cuDNN MXFP8 normalization intermediate output uses the input dtype (itype) instead of the compute dtype; otype still applies to the final quantized output.
forward-ok (4096, 128)
backward-ok
PASS
=== END 4096x7168 flag=1 train exit=0 ===

case	result
2048×7168, flag 1, train	SIGSEGV, exit=-11, no forward-ok
4096×7168, flag 1, train	forward and backward pass
2048×7168, flag 0, train	forward and backward pass, no common.cpp:312 warning
2048×6144, flag 1, train	forward and backward pass
2048×7168, flag 1, eval	forward pass

4096×7168 is the tensor-parallel-1 shape of a 4096-token, hidden-7168 model. 2048×7168 is that sequence split by tensor parallel 2. Hidden 6144 at the same token count plans, so this is one shape, not a width threshold. Eval plans because the column-wise quantize output is absent.

This keeps the FP32 intermediate for that one shape on cuDNN 9.26 (cudnnGetVersion() 92600–92699) even when the flag is 1. cuDNN norm stays on. Other shapes still take the BF16 intermediate. Drop the guard once a cuDNN release fixes the heuristic.

The overflow is in libcudnn_engines_runtime_compiled.so.9.26.0, reached from EngineHeuristics::finalize_internal. cudnn-frontend only forwards HeurMode_t::A into query_cudnn_heuristics_impl, and the process dies before an engine id comes back. A cuDNN bug report is being filed separately.

Type of change

• Documentation change (change only to the documentation, either a fix or a new content)
• Bug fix (non-breaking change which fixes an issue)
• New feature (non-breaking change which adds functionality)
• Breaking change (fix or feature that would cause existing functionality to not work as expected)
• Infra/Build change
• Code refactoring

Changes

• Skip the BF16 MXFP8 RMSNorm intermediate for 2048×7168 forward training on cuDNN 9.26, and warn once when the flag asked for it.
• Note that exception on NVTE_CUDNN_MXFP8_NORM_OUTPUT_IN_INPUT_DTYPE.

Checklist:

• I have read and followed the contributing guidelines
• The functionality is complete
• I have commented my code, particularly in hard-to-understand areas
• I have made corresponding changes to the documentation
• My changes generate no new warnings
• I have added tests that prove my fix is effective or that my feature works
• New and existing unit tests pass locally with my changes

No unit test is included. The failure is a SIGSEGV inside cuDNN during create_execution_plans, before execute(), and it needs cuDNN 9.26 and the exact shape. A test would kill the process on that library. The five cases above were run against the unmodified library in the image. This change was not compiled in the environment that opened the pull request.

[greptile-apps]

Confidence Score: 5/5

[Medium impact] The PR appears safe to merge; no actionable issues were found.

Summary

Keeps the FP32 MXFP8 RMSNorm intermediate for 2048×7168 forward training on cuDNN 9.26, even when NVTE_CUDNN_MXFP8_NORM_OUTPUT_IN_INPUT_DTYPE is enabled.

• RMSNorm keeps an FP32 intermediate for one cuDNN 9.26 shape.

_{Reviews (1) · Last reviewed commit: "[pre-commit.ci] auto fixes from pre-comm..." · Reviewed by Greptile}