Add support for fused Q Up-Proj GEMM/RoPE/Quant.

[chaseblock]

This PR adds support for fusing the GEMM in the Q Up Proj step of DeepseekV3 training with the following RoPE and MXFP8 quantization operations. This uses a custom kernel from cudnn_frontend, and supports both 16-bit projection and mxfp8 projection.

Type of change

• New feature (non-breaking change which adds functionality)

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 Summary

The PR adds a public fused MLA Q up-projection wrapper and extends fused attention to consume MXFP8 Q/K/V while optionally returning BF16 gradients.

• Exports FusedMLAQUpProjRopeQuant from the PyTorch attention APIs.
• Adds BF16-backward control to fused attention.
• Adds two-phase MXFP8 quantization, transpose, and scale-swizzle helpers.
• Adds MXFP8 storage-aware QKV layout detection.

Confidence Score: 5/5

The PR appears safe to merge within the scope of this follow-up review.

No blocking failure remains in the eligible review scope.

Important Files Changed

Filename	Overview
transformer_engine/pytorch/attention/fused_mla_q_uproj.py	Adds the SM100 fused Q up-projection, RoPE, and MXFP8 quantization wrapper with BF16 and MXFP8 weight paths.
transformer_engine/pytorch/attention/dot_product_attention/utils.py	Adds separable MXFP8 quantization and BHSD scale-transpose/swizzle helpers.
transformer_engine/pytorch/attention/dot_product_attention/backends.py	Adds an option for fused-attention backward to return BF16 gradients for FP8 inputs while preserving the autograd argument contract.
transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py	Adds MXFP8 storage-aware layout detection and propagates the BF16-backward option to fused attention.
transformer_engine/pytorch/attention/init.py	Exports the fused MLA Q up-projection wrapper from the attention package.
transformer_engine/pytorch/init.py	Exports the new wrapper from the top-level PyTorch API.

Flowchart

%%{init: {'theme': 'neutral'}}%%
flowchart LR
  X["Input activation"] --> U["Fused Q up-projection"]
  W["BF16 or MXFP8 weight"] --> U
  R["RoPE cos/sin"] --> U
  U --> Q["MXFP8 query"]
  K["Key"] --> KQ["MXFP8 quantization"]
  V["Value"] --> VQ["MXFP8 quantization"]
  Q --> S["BHSD scale transpose and GEMM swizzle"]
  KQ --> S
  VQ --> S
  S --> A["cuDNN fused attention"]
  A --> O["Attention output"]
  A --> B["Optional BF16 Q/K/V gradients"]

Loading

_{Reviews (3): Last reviewed commit: "Adjust comment SBHD/BSHD" | Re-trigger Greptile}

[sudhakarsingh27]

1. Commented feedback inline
2. Pls fix DCA here and in the other PR
3. Also, I wonder if you've measured perf for mxfp8_quantize_only and its swizzle counterpart. I feel they could be further optimized as I see a lot of for loops.

[sudhakarsingh27]

do we need a dict for this const

[chaseblock]

1. Commented feedback inline
2. Pls fix DCA here and in the other PR
3. Also, I wonder if you've measured perf for mxfp8_quantize_only and its swizzle counterpart. I feel they could be further optimized as I see a lot of for loops.

Addressed inline feedback.

DCA fixed.

I haven't explicitly modeled the perf of these two functions, but they were essentially extracted from the existing combine_and_quantize, which we needed to break apart since one of the tensors is now being quantized beforehand.