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 introduces a Blackwell-targeted fused MLA Q up-projection, RoPE, and MXFP8 quantization path and integrates its pre-quantized output with fused attention.

• Exposes FusedMLAQUpProjRopeQuant through the PyTorch attention API.
• Adds MXFP8 layout preparation and optional BF16 attention-gradient handling.
• Adds forward numerical coverage and registers the test in the L0 PyTorch job.

Confidence Score: 5/5

The PR appears safe to merge based on the eligible follow-up findings available in this review.

No blocking failure remains.

Important Files Changed

Filename	Overview
transformer_engine/pytorch/attention/fused_mla_q_uproj.py	Adds feature detection and the cuDNN frontend wrapper for fused Q projection, RoPE, and MXFP8 output construction.
transformer_engine/pytorch/attention/dot_product_attention/utils.py	Splits MXFP8 preparation into quantization and scale transpose/swizzle phases to accommodate pre-quantized tensors.
transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py	Adds pre-quantized MXFP8 layout detection and forwards the requested backward precision to fused attention.
transformer_engine/pytorch/attention/dot_product_attention/backends.py	Integrates pre-quantized MXFP8 scale metadata and optional BF16 Q/K/V gradients into the fused-attention autograd path.
tests/pytorch/attention/test_fused_mla_q_uproj.py	Adds Blackwell-gated forward numerical and saved-activation checks for the new fused operation.
qa/L0_pytorch_unittest/test.sh	Registers the new fused MLA test in the L0 PyTorch unit-test job.

Sequence Diagram

sequenceDiagram
participant Model
participant Fusion as Fused MLA Q Up-Proj
participant Quant as MXFP8 Layout Preparation
participant DPA as DotProductAttention
participant cuDNN as cuDNN Fused Attention
Model->>Fusion: Q-LORA activation, weight, cos, sin
Fusion->>Fusion: GEMM + per-head RoPE + MXFP8 quantization
Fusion-->>Quant: Quantized Q and saved activation
Model->>Quant: K and V
Quant->>Quant: Quantize and transpose/swizzle scales
Quant->>DPA: Pre-quantized Q, K, V
DPA->>cuDNN: MXFP8 fused-attention request
cuDNN-->>DPA: Attention output
DPA-->>Model: Output and backward gradients

Loading

_{Reviews (11): Last reviewed commit: "Merge branch 'main' into qrope_fusion" | 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.

[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.

[sraman-rgb]

/te-ci pytorch

[sudhakarsingh27]

LGTM