[PyTorch] ONNX test fix + export for FP8 attention

[pggPL]

Description

1. Fixes incorrect implementation of no_torch_dynamo decorator, which results in errors for newest PyTorch. The decorator was not correctly disabled during export to onnx.
2. Adds support for FP8 attention export.

Fixes #2588

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

Please list the changes introduced in this PR:

• Change A
• Change B

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

[greptile-apps]

Greptile Overview

Greptile Summary

This PR fixes ONNX export issues and adds FP8 attention support for ONNX export.

Key Changes:

• Fixed no_torch_dynamo decorator bug: The previous lambda-based implementation incorrectly evaluated is_in_onnx_export_mode() at decoration time instead of at runtime, causing errors with newer PyTorch versions during ONNX export. The new implementation uses a proper wrapper function that checks the export mode at each call.

• Added FP8 attention ONNX export: Implemented onnx_forward method in FP8EmulationFunc that uses ONNX-compatible operations (flatten, concatenate, quantize, split) to emulate FP8 quantization during export.

• Enabled FP8 emulation during ONNX export: Modified attention backend selection logic to allow FP8 emulation when is_in_onnx_export_mode() returns true, even without the environment variable.

• Updated tests: Added parameterization for FP8 recipes (DelayedScaling and Float8CurrentScaling with fp8_dpa=True), removed attention_dropout=0.5 to avoid non-deterministic outputs, and adjusted tolerance to 5e-1 for FP8 tests.

• Updated CI: Set NVTE_UnfusedDPA_Emulate_FP8=1 environment variable in test script to enable FP8 emulation in CI environments without native FP8 hardware.

The implementation follows best practices for ONNX export by using operations with defined ONNX translations and properly handling the export mode detection.

Confidence Score: 4/5

• This PR is safe to merge with minor considerations
• The changes correctly address the no_torch_dynamo decorator bug that was causing issues with newer PyTorch versions, and properly implement FP8 attention ONNX export. The implementation follows best practices by using ONNX-compatible operations and runtime mode detection. Minor consideration: assert statement in onnx_forward could be improved with better error handling.
• Pay attention to transformer_engine/pytorch/attention/dot_product_attention/backends.py - the assert statement could cause issues if non-Float8 quantizers are passed during ONNX export

Important Files Changed

Filename	Overview
transformer_engine/pytorch/jit.py	Fixed no_torch_dynamo decorator to properly check ONNX export mode at runtime instead of at decoration time, preventing errors with newer PyTorch versions
transformer_engine/pytorch/attention/dot_product_attention/backends.py	Added onnx_forward method to FP8EmulationFunc for ONNX-compatible FP8 quantization/dequantization using flatten+concat+quantize+split operations
tests/pytorch/test_onnx_export.py	Added FP8 recipe parameterization to core attention tests, removed attention_dropout=0.5 parameter, and adjusted tolerance for FP8 tests

Sequence Diagram

sequenceDiagram
    participant User
    participant DotProductAttention
    participant AttentionBackend as get_attention_backend
    participant UnfusedDPA as UnfusedDotProductAttention
    participant FP8Emulation as FP8EmulationFunc
    participant Quantizer as Float8Quantizer
    participant ONNX as ONNX Export

    User->>DotProductAttention: forward(query, key, value)
    DotProductAttention->>AttentionBackend: get_attention_backend()
    
    alt ONNX Export Mode
        AttentionBackend->>AttentionBackend: is_in_onnx_export_mode() == True
        AttentionBackend->>AttentionBackend: allow_emulation = True
        AttentionBackend-->>DotProductAttention: use UnfusedDotProductAttention
        
        DotProductAttention->>UnfusedDPA: forward(Q, K, V)
        UnfusedDPA->>FP8Emulation: apply(Q, K, V, quantizer, "QKV_quantizer")
        FP8Emulation->>FP8Emulation: is_in_onnx_export_mode() == True
        FP8Emulation->>FP8Emulation: onnx_forward()
        
        FP8Emulation->>FP8Emulation: flatten Q, K, V
        FP8Emulation->>FP8Emulation: concatenate tensors
        FP8Emulation->>Quantizer: onnx_quantize(combined)
        Quantizer-->>FP8Emulation: FP8 tensor
        FP8Emulation->>Quantizer: onnx_dequantize(fp8_tensor)
        Quantizer-->>FP8Emulation: dequantized tensor
        FP8Emulation->>FP8Emulation: split and reshape Q, K, V
        FP8Emulation-->>UnfusedDPA: emulated FP8 Q, K, V
        
        UnfusedDPA->>UnfusedDPA: compute attention(Q, K, V)
        UnfusedDPA-->>DotProductAttention: attention output
        DotProductAttention->>ONNX: export to ONNX graph
        ONNX-->>User: ONNX model with FP8 attention
    else Normal Training/Inference
        AttentionBackend->>AttentionBackend: check env var NVTE_UnfusedDPA_Emulate_FP8
        AttentionBackend-->>DotProductAttention: select appropriate backend
        DotProductAttention-->>User: attention output
    end

Loading

[greptile-apps]

_{No files reviewed, no comments}

_{Edit Code Review Agent Settings | Greptile}

[greptile-apps]

_{1 file reviewed, 1 comment}

_{Edit Code Review Agent Settings | Greptile}

[pggPL]

/te-ci pytorch L1

[timmoon10]

LGTM

[timmoon10]

This is fine for FP8 attention, although we'll need to revisit whenever we support MXFP8 or NVFP4. Why can't we concatenate the 2D tensors?

[pggPL]

I'm not sure if this will work for all layouts and different max_q_length and max_kv_length. Added asserions that's it not mxfp8, because I want to merge it fast. I will rethink it when adding support for mxfp8.

[greptile-apps]

_{1 file reviewed, 1 comment}

_{Edit Code Review Agent Settings | Greptile}

[greptile-apps]

Assert statement will cause ONNX export to fail if non-Float8 quantizers are used. Consider replacing with a runtime check that raises a more descriptive error or logging a warning.

[pggPL]

/te-ci pytorch L1