[common] [NVFP4] Per-token (row-scaled transpose) cast + CUTLASS GEMM kernels (1/3, split from #3045)

[cael-ling]

Description

This is the first of three PRs splitting #3045 (NVFP4 per-token quantization recipe) into reviewable pieces. It lands only the common (C/CUDA) layer: the NVFP4 row-scaled transpose ("per-token") quantization kernels and CUTLASS GEMM.

Following how #2931 integrates row-scaled NVFP4, the kernels are reached only through TransformerEngine's existing nvte_quantize_v2 / nvte_cublas_gemm_v2 entry points — no new top-level C-API. Selection is driven by tensor metadata (the existing row_scaled_nvfp4 flag from #2931) rather than a dedicated dispatch: a row-scaled NVFP4 tensor that also allocates a columnwise buffer is produced/consumed by the transpose kernels. Naming follows the merged row-scaled NVFP4 convention (cutlass_nvfp4_*, row_scaled_transpose_*) so the new kernels sit alongside the existing ones.

What's included

Cast kernels

• cast/nvfp4/quantize_transpose_row_scaled_nvfp4.{cu,h} — single-tensor row-scaled transpose cast: per-row (rowwise) and per-col (columnwise) NVFP4 scales produced in one pass.
• cast/nvfp4/group_quantize_transpose_row_scaled_nvfp4.{cu,h} — grouped (multi-tensor) variant.

GEMM kernels

• gemm/cutlass_nvfp4_gemm.{cu,h} + gemm/cutlass_nvfp4_grouped_gemm.cu — CUTLASS NVFP4 GEMM (dense + grouped) with per-row*per-col rescale fused into the EVT epilogue.
• gemm/cutlass_nvfp4_post_scale.cu — standalone post-scale kernel (internal helper).

API integration (no new public C-API)

• Cast: cast/dispatch/quantize.cuh selects the row-scaled transpose kernels inside nvte_quantize_v2 when a row-scaled NVFP4 output also allocates a columnwise buffer (output->row_scaled_nvfp4 && has_columnwise_data()) — single + grouped. No dedicated config attribute.
• GEMM: cublas_gemm() (nvte_cublas_gemm_v2) dispatches to the CUTLASS row-scaled kernel when both operands carry the row_scaled_nvfp4 flag, replacing the previous hard reject. The dense entry point is the internal cutlass_nvfp4::gemm_row_scaled; the grouped GEMM and post-scale kernels are internalized (no C-API), to be wired into their dispatch paths by follow-up PRs.
• New QuantizationConfig attributes for the RHT / random-sign-mask options (transformer_engine.h / common.h / transformer_engine.cpp), plumbed through nvte_quantize_v2 (default-off in this PR).

Tests (C++ gtest, no Python dependency)

• Extended tests/cpp/operator/test_cast_nvfp4_transpose.cu with a row-scaled-transpose suite, cross-validated against the merged row-scaled 1D kernel (Implement row-scaled NVFP4 fprop recipe #2931, RowScaled1D): rowwise vs RowScaled1D(input), columnwise vs RowScaled1D(input^T). Amaxes and FP8 block scales match bitwise; FP4 codes may differ only at rounding-midpoint ties (< 0.1%).
• New tests/cpp/operator/test_cutlass_nvfp4_gemm.cu validating the row-scaled GEMM through the public nvte_cublas_gemm_v2 path (row-scaled NVFP4 operands with per-row amax vectors) against a dequant-matmul-rescale fp32 reference.

What's NOT included (deferred)

To keep this PR focused on the kernels, the following from #3045 land in follow-up PRs:

• PyTorch bindings, recipe, quantizer, and module wiring
• Grouped-GEMM and post-scale dispatch wiring (kernels are internalized here; their callers land later)
• Backward pass, stochastic rounding, RHT, and 2D-weight paths
• Docs / examples

Note: the SR / RHT / random-sign-mask code paths are present in the cast kernels and plumbed through the nvte_quantize_v2 config in this PR, but are exercised only in their default-off state. Their non-default behavior is validated together with the recipe that turns them on (follow-up PR), since SR is non-deterministic and RHT needs a dedicated reference (the #2931 RowScaled1D kernel used for cross-validation here has neither).

Testing

C++ operator gtests on Blackwell (SM100), all passing:

• CastNVFP4RowScaledTransposeTestSuite
• CutlassNVFP4RowScaledGemmTestSuite (exercises the nvte_cublas_gemm_v2 fold)

Related

• Splits Add NVFP4 per-token quantization recipe #3045
• Follows the row-scaled NVFP4 integration pattern from Implement row-scaled NVFP4 fprop recipe #2931

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 adds the common C/CUDA layer for row-scaled NVFP4 transpose quantization and CUTLASS GEMM. The main changes are:

• New single and grouped row-scaled NVFP4 transpose cast kernels.
• New dense and grouped CUTLASS NVFP4 GEMM kernels.
• Dispatch through the existing quantize and GEMM entry points.
• Quantization config plumbing for stochastic rounding, RHT, and random sign masks.
• C++ tests for row-scaled transpose casts and dense CUTLASS NVFP4 GEMM.

Confidence Score: 5/5

This looks safe to merge.

• No blocking issues found in the changed code.

Important Files Changed

Filename	Overview
transformer_engine/common/cast/dispatch/quantize.cuh	Routes row-scaled NVFP4 tensors with columnwise output into the new transpose quantization kernels.
transformer_engine/common/gemm/cutlass_nvfp4_gemm.cu	Adds the dense CUTLASS row-scaled NVFP4 GEMM implementation.
transformer_engine/common/gemm/cutlass_nvfp4_grouped_gemm.cu	Adds the grouped CUTLASS row-scaled NVFP4 GEMM implementation.
transformer_engine/common/cast/nvfp4/quantize_transpose_row_scaled_nvfp4.cu	Adds the single-tensor row-scaled transpose cast kernel.
transformer_engine/common/cast/nvfp4/group_quantize_transpose_row_scaled_nvfp4.cu	Adds the grouped row-scaled transpose cast kernel.

_{Reviews (3): Last reviewed commit: "[common] Add NVFP4 row-scaled transpose ..." | Re-trigger Greptile}

[greptile-apps]

Shared Scratch Cross-Stream Race

When two grouped NVFP4 GEMMs run on different streams, both calls reuse the same static metadata and workspace buffers. One launch can overwrite or free the scratch while another stream's CUTLASS kernel is still reading it, which can produce wrong grouped GEMM outputs or allocator failures.

[greptile-apps]

Hardware Info Uses Device Zero

This launch always reports device 0 to CUTLASS, but the API runs on the caller's current CUDA device. On a multi-GPU job where the active device is not 0, the grouped scheduler can use the wrong SM count/device id and size persistent work incorrectly, causing poor scheduling or launch failures.

Suggested change

	cutlass::KernelHardwareInfo hw_info;
	hw_info.device_id = 0;
	hw_info.sm_count = cached_sm_count();
	int device_id = 0;
	NVTE_CHECK_CUDA(cudaGetDevice(&device_id));
	cutlass::KernelHardwareInfo hw_info;
	hw_info.device_id = device_id;
	hw_info.sm_count = cutlass::KernelHardwareInfo::query_device_multiprocessor_count(device_id);

[greptile-apps]

Grouped Scale Dtype Unchecked

The grouped API validates A/B/D and alpha tensors but never checks a_sf[g] or b_sf[g] before passing their bytes to CUTLASS as NVFP4 scale factors. A caller that passes an FP32 or otherwise wrong scale tensor gets silently corrupted GEMM results instead of the same clear error the dense APIs return.

Suggested change

	NVTE_CHECK(a_t->data.dtype == DType::kFloat4E2M1 && b_t->data.dtype == DType::kFloat4E2M1,
	"group ", g, ": A/B must be FP4 e2m1");
	NVTE_CHECK((d_t->data.dtype == DType::kFloat32) == d_is_fp32, "group ", g,
	": D dtype must be uniform across groups");
	NVTE_CHECK(a_t->data.dtype == DType::kFloat4E2M1 && b_t->data.dtype == DType::kFloat4E2M1,
	"group ", g, ": A/B must be FP4 e2m1");
	NVTE_CHECK(sa_t->data.dtype == DType::kFloat8E4M3 || sa_t->data.dtype == DType::kByte,
	"group ", g, ": A scale must be FP8 e4m3 (or raw uint8 byte)");
	NVTE_CHECK(sb_t->data.dtype == DType::kFloat8E4M3 || sb_t->data.dtype == DType::kByte,
	"group ", g, ": B scale must be FP8 e4m3 (or raw uint8 byte)");
	NVTE_CHECK((d_t->data.dtype == DType::kFloat32) == d_is_fp32, "group ", g,
	": D dtype must be uniform across groups");

[greptile-apps]

RNG State Length Unchecked

The stochastic-rounding path only verifies that rng_state is int64 before treating it as the seed/offset pair used by the kernel. If a caller passes a one-element int64 tensor with row-scaled transpose enabled, the kernel reads the missing offset past the allocation and can produce invalid randomness or an illegal memory access.