[pggPL]

Description

Some native operations return compact MXFP8 or NVFP4 scales even when their quantizer has optimize_for_gemm=True. Callers and fake implementations must then predict the selected kernel's output layout.

Honor the requested GEMM scale layout at the native operation boundary. When a kernel can only emit compact scales, it writes directly into the output with the output's scale pointers temporarily redirected to scratch buffers, and the scales are then swizzled into the output's own buffers. Output buffers are never reallocated, so their addresses stay stable. Kernels that already emit GEMM-swizzled scales skip the conversion.

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

• Add redirect_scales_to_scratch / restore_and_swizzle_scales (single- and multi-tensor) for kernels that only emit compact scales. They touch only the scale pointers and the swizzle flag, so the output tensor itself is used by the kernel.
• NVFP4 quantization: use the redirect only when the kernel cannot emit GEMM-swizzled scales (e.g. 1D quantization without RHT cast fusion, row-scaled, 4over6, unaligned shapes). With a noop_flag, the previous scales are unswizzled into the scratch buffers first, so a GPU no-op leaves the output unchanged.
• Multi/split quantization: decide per tensor whether compact scratch scales are needed. Those tensors share one scratch allocation and one swizzle kernel per usage; tensors whose kernel emits GEMM-swizzled scales, or that do not request them, are quantized directly.
• Normalization: add Impl::FUSED_NORM_QUANT_UNSWIZZLED for cuDNN MXFP8 LayerNorm/RMSNorm with optimize_for_gemm, using the redirect within the existing intermediate-output switch. Keep 2D MXFP8 normalization on the compatible unfused path.
• GEMM: convert a provided quantized output to the quantizer's layout where it is constructed when unfused quantization is needed.
• Remove post-quantize inplace_swizzle_scale_for_gemm calls, which replaced the scale buffers of the output.
• Simplify fake_bgrad to derive the layout directly from the quantizer and remove the unused layout-query binding.
• Cover native producers, multi/split quantization layouts (including mixed usages, formats and empty tensors) and compiled bias gradients.

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

[timmoon10]

Good catch, it definitely is hacky that we sometimes ignore optimize_for_gemm. I'm skeptical about the changes to the C++ infrastructure since QuantizationWorkspace's design is rather ad hoc.

[pggPL]

My agent says:

Scale corruption on main with noop_flag (fixed by this PR)

On main, an in-place NVFP4 weight update with optimize_for_gemm=True corrupts the scales when noop_flag skips quantization.

• It happens when the quantization kernel can't emit GEMM-swizzled scales itself, e.g. 1D quantization without RHT cast fusion, row-scaled NVFP4, 4over6, or shapes that aren't 128-aligned.
• The first quantization writes compact scales, and the post-quantize inplace_swizzle_scale_for_gemm swizzles them.
• On the next update, convert_and_update_tensor resets _with_gemm_swizzled_scales to False, but the buffer still holds swizzled scales.
• If noop_flag skips the kernel (CUDA graphs with weight caching), the fallback swizzles these already-swizzled scales a second time. The FP4 data stays intact, but ~90% of the scale bytes are wrong and GEMM output is off by ~33% relative norm (fprop and dgrad).

In this PR the output's scales are redirected to scratch buffers during quantization. With noop_flag, the current scales are first unswizzled into the scratch, so a skipped update reproduces them exactly. The cost is one small scale-only kernel per update, and only when noop_flag is passed. MXFP8 is not affected, because its kernel always emits swizzled scales.

[greptile-apps]

Confidence Score: 4/5

[Medium impact] The change reaches native quantization, normalization, and GEMM output workflows in PyTorch.

The PR appears safe to merge, with a non-blocking test gap around no-op replay.

Findings

1.  Swizzled no-op replay lacks coverage ▶

Summary

Native quantized outputs now honor optimize_for_gemm without replacing their scale buffers.

• Native quantization returns the scale layout the quantizer requests.
• Normalization keeps GEMM-ready scales on cuDNN MXFP8 outputs.
• GEMM converts supplied outputs when it must quantize them afterward.
• Compiled bias gradients use the quantizer’s requested scale layout.

Diagram

%%{init: {'theme': 'neutral'}}%%
flowchart TD
  A[Quantized output] --> B{Kernel writes requested layout?}
  B -->|Yes| C[Write directly into output]
  B -->|No| D[Redirect scale pointers to scratch]
  D --> E[Kernel writes data and compact scales]
  E --> F[Restore output scale pointers]
  F --> G[Swizzle scratch scales into output buffers]

Loading

_{Reviews (1) · Last reviewed commit: "Drop NVFP4 shape-gate tests for GEMM-rea..." · Reviewed by Greptile}

[greptile-apps]

Swizzled no-op replay lacks coverage

The new tests never pass noop_flag, and test_skip_quantization_with_noop_flag leaves optimize_for_gemm disabled. That leaves the new preserve_scales branch untested, so a regression could change cached scales during a step that should leave them untouched.

Extend the no-op replay test with optimize_for_gemm=True for NVFP4 outputs. Check that data, scales, and their addresses stay unchanged when the captured graph skips quantization.