Fix: fused QUproj + RoPE + Quant

[ghadiaravi13]

Description

Modify the FusedMLAQUpProjFunction to absorb norm inside the autgrad function, so that the fp32 output from the normalization is directly quantized to MXFP8 without intermediate bf16 rounding.

Fixes # (issue)

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 Summary

This PR moves RMSNorm inside the fused MLA Q up-projection autograd function so its FP32 accumulator can be quantized directly to MXFP8 before fused GEMM, RoPE, and output quantization.

• Adds an unfused custom backward covering inverse RoPE, projection gradients, and RMSNorm gradients.
• Reuses the normalization-produced MXFP8 activation for forward GEMM and weight-gradient computation.
• Exposes FusedMLAQUpProjFunction through the attention and root PyTorch namespaces.
• Restricts this temporary fused path to RMSNorm and single-rank tensor parallelism.

Confidence Score: 4/5

The PR should not merge until the RoPE backward launch validates that the rotary interval fits within each query head; the missing API documentation is non-blocking.

Caller-controlled head metadata can make the new Triton backward kernel cross head boundaries and access beyond the gradient allocation because its mask bounds only the starting offset.

Files Needing Attention: transformer_engine/pytorch/attention/fused_mla_q_uproj.py and transformer_engine/pytorch/init.py

Important Files Changed

Filename	Overview
transformer_engine/pytorch/attention/fused_mla_q_uproj.py	Adds normalization and complete custom autograd around fused MLA projection, but backward lacks a required head-dimension bound before launching the Triton RoPE kernel.
transformer_engine/pytorch/attention/init.py	Publicly exports the new autograd function; this makes its low-level caller-controlled dimension contract directly reachable.
transformer_engine/pytorch/init.py	Re-exports the new function at the root PyTorch API without adding the corresponding API documentation.

Flowchart

%%{init: {'theme': 'neutral'}}%%
flowchart LR
  X[Pre-normalization input] --> N[RMSNorm]
  N --> Q[Direct MXFP8 quantization]
  Q --> G[Fused Q up-projection GEMM]
  G --> R[Per-head RoPE]
  R --> O[MXFP8 query]
  O --> RB[Inverse RoPE backward]
  RB --> LB[Linear backward]
  LB --> NB[RMSNorm backward]

Loading

_{Reviews (1): Last reviewed commit: "refactor: simplify comments and document..." | Re-trigger Greptile}

[greptile-apps]

Unbounded per-head RoPE access

When q_head_dim matches the fused output width but qk_head_dim + qk_pos_emb_head_dim exceeds it, the backward kernel's mask bounds only each head's starting offset, so its loads and stores cross into adjacent heads and beyond the final head, causing corrupted gradients or a CUDA illegal-memory-access failure.

Knowledge Base Used: PyTorch attention execution

[greptile-apps]

Undocumented public autograd API

FusedMLAQUpProjFunction is now exported from the root and attention namespaces without a corresponding entry in docs/api/pytorch.rst, leaving its extensive argument and shape contract absent from the generated public API reference.

Knowledge Base Used: PyTorch runtime and public API

Note: If this suggestion doesn't match your team's coding style, reply to this and let me know. I'll remember it for next time!