MXFP4 Cast Transpose Triton [WIP]

[sarthak-amd]

Description

Implements the MXFP4 rowwise and columnwise FP32/BF16 -> MXFP4 fused quantization + cast kernel

• Verify Tolerances and functional Unit Tests

• The triton te_cast_transpose_mxfp4_triton currently outputs FP4 data in linear layout [M, N/2] with contiguous byte packing. AITER's gemm_a4w4 requires the B matrix in MFMA shuffle layout for tensor cores. This layout shuffle can be fused into the triton kernel in future.

[wangye805]

You

[wangye805]

We previously already defined env NVTE_USE_CAST_TRANSPOSE_TRITON.

[wangye805]

If FP4 data shuffle is not yet supported in Triton kernel, why do we need to add it here?

[sarthak-amd]

This is kept to ensure API consistency between Triton and the upcoming hip kernel for which I'll create a separate PR. In the hip kernel we were able to fuse the shuffle.

[sarthak-amd]

hip vs triton flow
Input: BF16 [M, N]
↓
MXFP4Quantizer.update_quantized()
↓
tex.cast_transpose_mxfp4_fused_shuffle() [Single HIP kernel]
↓
├─→ Rowwise FP4 [M, K/2] (MFMA shuffled)
├─→ Rowwise Scale [M_pad, K/32_pad] (shuffled)
├─→ Colwise FP4 [N, M/2] (MFMA shuffled)
└─→ Colwise Scale [N_pad, M/32_pad] (shuffled)
↓
AITER gemm_a4w4 (zero-copy)

vs

Input: BF16 [M, N]
↓
MXFP4Quantizer.update_quantized()
↓
te_cast_transpose_mxfp4_triton() [Triton JIT kernel]
↓
├─→ Rowwise FP4 [M, K/2] (linear layout)
├─→ Rowwise Scale [M_pad, K/32_pad] (shuffled)
├─→ Colwise FP4 [N, M/2] (linear layout)
└─→ Colwise Scale [N_pad, M/32_pad] (shuffled)
↓
aiter.ops.shuffle.shuffle_weight() [External call]
↓
FP4 data → MFMA layout
↓
AITER gemm_a4w4

[wangye805]

Can we add some prime numbers like

TransformerEngine/tests/cpp/operator/test_cast_transpose.cu

Lines 90 to 92 in 9d6b0e5

	{1, 3221}, // Prime 456
	{2333, 1}, // Prime 345
	{1481, 677}}; // Primes 234, 123

[wangye805]

Data tol seems to be quite large. You can follow our mxfp8 scale and data adjustment scheme:

TransformerEngine/tests/cpp/test_common.cu

Line 730 in 9d6b0e5

void adjust_ref_for_e8m0_scale_error(const std::string &name,

[wangye805]

What is fp4 shuffle?

[sarthak-amd]

fp4 shuffle basically rearranges [M, K/2] linear layout → MFMA instruction layout (16×16).

The currently flow training workflow if TE MXFP4 Quantization Kernel is used is as follows
TE Triton Kernel → Linear FP4 [N, K/2] → aiter.ops.shuffle_weight() → MFMA FP4 → aiter.gemm_a4w4()

You can find the shuffle code in aiter/aiter/ops/shuffle.py

[wangye805]

If we are going to enable kFloat4E2M1, there are other related changes needed. Search for https://github.com/search?q=repo%3AROCm%2FTransformerEngine%20kFloat4E2M1&type=code for more details:

[wangye805]

Is there alignment/padding requirements for M and K?

[wangye805]

TE currently supported 2D matrices from flatted high-dimensional tensors:

TransformerEngine/transformer_engine/common/common.h

Lines 238 to 262 in 9d6b0e5

	size_t flat_first_dim() const {
	const auto &full_shape = shape();
	size_t ret = 1;
	if (!full_shape.empty()) {
	for (size_t i = 0; i < full_shape.size() - 1; i++) {
	ret *= full_shape[i];
	}
	}
	return ret;
	}
	/*! Matrix width after tensor is flattened to 2D
	*
	* If a tensor has dimensions (D1, D2, ..., Dn), it is reinterpreted
	* as a (D1*D2*...*D(n-1), Dn) matrix.
	*/
	size_t flat_last_dim() const {
	const auto &full_shape = shape();
	if (full_shape.empty()) {
	return 1;
	} else {
	return full_shape.back();
	}
	}
	};

[wangye805]

I presume this 256 is from some alignment/padding requirement?

[sarthak-amd]

The 256 alignment is required by AITER's CK-based MXFP4 GEMM kernels for scale tensor swizzle/shuffle layout.
Required for scale swizzle layout: 256 = ScaleBlockSize(32) × 8 waves.
See aiter/aiter/utility/fp4_utils.py:398 and gemm_a4w4_blockscale_common.cuh:66

[wangye805]

You will need to add this pytest into our ci script (somewhere near

TransformerEngine/ci/pytorch.sh

Line 74 in 9d6b0e5

run_default_fa 1 triton_kernels/test_norms.py

) otherwise it won't be tested