[Common] row-scaled nvfp4 path: fuse row/col amax into a single TMA-tiled kernel

[cael-ling]

Description

The NVFP4 row-scaled path that was originally proposed in #2931 computes per-row and per-column amax with two separate kernels. This PR fuses both directions into a single kernel that streams 128x128 chunks through shared memory via TMA (coalesced loads) and does the column reduction from SMEM. Amax is an exact max reduction, so results are byte-identical to the two-kernel path. The fused path is used only when the quantize call needs both directions (rowwise + columnwise amax) on a BF16 input with 128-aligned dims; the kernel then produces both amaxes in one pass. Any other case keeps the original two kernels. NVTE_NVFP4_FUSED_AMAX=0 forces the fallback at runtime.

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 compute_fused_amax_kernel (rowwise + columnwise) and its host wrappers
  fused_amax_supported / compute_fused_amax in quantize_transpose_nvfp4.cuh.
• Dispatch to the fused kernel in dispatch/quantize.cuh (fwd and bwd) when supported, else fall back to the standalone amax kernels.
• Add NVTE_NVFP4_FUSED_AMAX kill switch (default on).

Performance

Full-quantize median latency, fused (fused amax + cast) vs (row-wise & columnwise amax + cast), the cast kernel is byte-identical so the delta is the amax step:

shape	fused (ms)	fallback (ms)	speedup
4096x4096	0.098	0.138	1.41x
8192x8192	0.148	0.360	2.44x
8192x16384	0.210	0.648	3.09x
32768x8192	0.346	1.222	3.53x
16384x16384	0.348	1.263	3.63x

Reproduce

Single Blackwell (SM100) GPU.

• pytest tests/pytorch/nvfp4/test_nvfp4_quantize_exact.py -k "row_scaled and both_directions"
  passes; rowwise/columnwise amax and qx/qx_t match the reference exactly.
• Fixed-input repeat runs give byte-identical amax (deterministic) and match the two-kernel path.
• Add NVTE_NVFP4_FUSED_AMAX env var (default enabled): set to 0 to disable the fused path at runtime and fall back to the two standalone amax kernels, no rebuild needed.

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]

Confidence Score: 5/5

The PR appears safe to merge.

Summary

The PR replaces separate rowwise and columnwise NVFP4 amax dispatch with a fused TMA-tiled reduction for eligible BF16 inputs while preserving the existing fallback.

• Adds fused row/column amax zeroing, reduction, support checks, and host launch logic.
• Routes forward and backward row-scaled quantization through a shared amax dispatcher.
• Preserves amax state for no-op calls by guarding both fused kernels before writes.

Diagram

%%{init: {'theme': 'neutral'}}%%
flowchart TD
    A[Row-scaled NVFP4 quantization] --> B{BF16, both directions,
128-aligned dimensions}
    B -->|Yes| C[Fused zero and TMA amax kernels]
    B -->|No| D[Standalone rowwise amax]
    D --> E{Columnwise output requested}
    E -->|Yes| F[Standalone columnwise amax]
    E -->|No| G[Continue quantization]
    C --> G
    F --> G

Loading

_{Reviews (5) · Last reviewed commit: "[pre-commit.ci] auto fixes from pre-comm..."}

[timmoon10]

Do we need to compile for the row-only or col-only cases? If the data or amax pointers are missing, it seems like we should call the unfused kernels. Also, we check amax pointers in fused_amax_supported, so will never hit the DO_ROW=false or DO_COL=false cases.

Suggested change

	TRANSFORMER_ENGINE_SWITCH_CONDITION(
	do_row, DO_ROW, TRANSFORMER_ENGINE_SWITCH_CONDITION(do_col, DO_COL, {
	auto kernel = compute_fused_amax_kernel<DO_ROW, DO_COL>;
	cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, dshmem_size);
	kernel<<<grid, block, dshmem_size, stream>>>(
	tensor_map_input, do_row ? row_amax_ptr : nullptr, do_col ? col_amax_ptr : nullptr,
	noop_ptr, rows, cols);
	}));
	auto kernel = compute_fused_amax_kernel<true, true>;
	cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, dshmem_size);
	kernel<<<grid, block, dshmem_size, stream>>>(
	tensor_map_input, row_amax_ptr, col_amax_ptr, noop_ptr, rows, cols);

[timmoon10]

LGTM, pending CI. I made a change to avoid unnecessary kernel compilations (#3454 (comment)), but can you sanity-check?

[timmoon10]

/te-ci