[PyTorch] Allow F16 GDN inside FP8 autocast

[layalir]

Description

GatedDeltaNetAttention currently rejects any enclosing FP8 autocast or calibration context, preventing FP16/BF16 GDN from being composed with MXFP8 linear layers.

Add NVTE_LA_FP8_RECIPE=F16 to run GDN in its input FP16/BF16 dtype while retaining the enclosing linear-layer recipe. Set the variable before importing Transformer Engine. Linear attention has its own override, independent of NVTE_DPA_FP8_RECIPE.

Type of change

• Bug fix (non-breaking change which fixes an issue)

Changes

• Disable quantization and calibration around the complete GDN forward lifecycle, including metadata preparation and checkpoint dispatch; restore the enclosing context on normal and exceptional exits.
• Preserve FP32 gate/state requirements and the existing rejection when the override is absent or unsupported.
• Document the variable and add regression coverage for calibration, context restoration, and FP16/BF16 GDN between MXFP8 linear layers with eager and checkpointed backward.

Validation

On GB300, the following selection passed 16 tests with zero failures or skips, including forward comparison with the reference recurrence and finite gradients:

NVTE_GDN_TEST_REQUIRED=1 CUDA_DEVICE_MAX_CONNECTIONS=1 \
python -m pytest -q tests/pytorch/attention/test_gdn_attention.py \
  -k 'f16_override or rejects_fp8_autocast'

Also completed ten synthetic policy-training steps on eight GB300 GPUs. Runtime observation confirmed BF16 GDN/GQA and MXFP8 expert grouped-GEMM forward, input-gradient, and weight-gradient operations on every rank. This verifies mixed-precision integration, not convergence or performance.

Checklist

• Read the contributing guidelines; both commits are signed off.
• Functionality, explanatory comments, documentation, and regression tests are included.
• Targeted GPU regression tests pass.
• Full repository L0 test and lint suites (not run; targeted validation is described above).

[greptile-apps]

Confidence Score: 5/5

The PR appears safe to merge, with no actionable correctness, security, or repository-rule issues identified.

Summary

This PR permits FP16/BF16 GatedDeltaNetAttention execution within an enclosing FP8 autocast or calibration region when explicitly configured.

• Adds the import-time NVTE_LA_FP8_RECIPE=F16 override.
• Temporarily disables FP8 quantization and calibration across GDN metadata preparation, execution, and checkpoint dispatch while restoring the enclosing context afterward.
• Preserves rejection of unsupported overrides and existing FP32 gate/state requirements.
• Documents the override and adds eager, checkpointed, calibration, exception-restoration, BF16, FP16, and MXFP8 integration coverage.

Diagram

%%{init: {'theme': 'neutral'}}%%
flowchart LR
    A[Enter enclosing FP8 autocast] --> B[MXFP8 input Linear]
    B --> C[Enter GDN-local autocast disabled]
    C --> D[Prepare GDN forward metadata]
    D --> E{Checkpointed?}
    E -->|No| F[Run GDN in FP16/BF16]
    E -->|Yes| G[Dispatch checkpoint with FP8 disabled state]
    G --> H[Recompute GDN in FP16/BF16]
    F --> I[End GDN lifecycle]
    H --> I
    I --> J[Restore enclosing FP8 state and recipe]
    J --> K[MXFP8 output Linear]

Loading

_{Reviews (1) · Last reviewed commit: "[PyTorch] Give linear attention its own ..."}

[ksivaman]

This can be closed in favor of #3521 which removes this restriction and achieves the same outcome. I would avoid introducing environment variables here, especially as multiple recipes are not supported for LA as of yet.

[layalir]

Makes sense. Closing.