[PyTorch] Fix NVFP4 metadata in compiled Linear operations

[pggPL]

Fix NVFP4 torch.compile failures for non-matrix Linear inputs and incorrect gradients from Linear operations with bias. Native Bias backward can produce compact scales while its fake output describes GEMM-swizzled scales, causing compiled GEMM to interpret the reconstructed tensor incorrectly.

Native execution and fake metadata now share alignment requirements through NVFP4Quantizer. C++ reads the RHT-fusion environment flag only when needed; Python reads it explicitly so Dynamo can guard changes. TensorSpec preserves the actual scale layout, and NVFP4 metadata matches native packed columnwise buffers and row-scaled amax shapes. Columnwise-only shape warnings are skipped during tracing.

Bias backward uses a small shared fake helper. Follow-up normalization and activation operations should use the same fused-output metadata helper.

Regression coverage stays in test_torch_compile.py, including dynamic shapes, environment changes, quantizer configurations and CUDA Graphs. Focused checks on SM89 passed: 98 tests passed and 44 skipped, including hardware-dependent NVFP4 cases. The C++ extension build and binding checks, pre-commit and changed-source pylint also passed. Validation of the final NVFP4 test matrix on SM100 remains pending; PyTorch L0 CI is requested separately.

Fixes #3626.

[greptile-apps]

Confidence Score: 5/5

[Medium impact] The changes appear safe to merge; no actionable new defects were found.

Summary

This PR aligns compiled NVFP4 metadata with native outputs for Linear operations with bias and non-matrix inputs.

• Compiled NVFP4 outputs describe the scale layout native code produces.
• NVFP4 fake tensors use the packed shapes and amax sizes native code uses.

The latest changes flatten the columnwise shape helper and test a size-one middle dimension. No new actionable issues were found. Tests were not run during this review.

Diagram

%%{init: {'theme': 'neutral'}}%%
flowchart LR
  A[NVFP4 quantizer settings] --> B[Shared alignment rules]
  B --> C[Native output scale layout]
  B --> D[Fake output scale layout]
  D --> E[TensorSpec metadata]
  C --> F[Compiled GEMM]
  E --> F

Loading

_{Reviews (12) · Last reviewed commit: "Fix NVFP4 columnwise shapes in the share..." · Reviewed by Greptile}

[pggPL]

/te-ci pytorch L0

[pggPL]

/te-ci pytorch