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

[sudhu2k]

Hi @sarthak-amd
Can't we merge the mxfp8 triton kernel with the mxfp4 triton kernel?

TransformerEngine/transformer_engine/pytorch/triton_kernels/cast_transpose.py

Line 335 in ef83316

def _cast_transpose_triton_mxfp8(

Kernel wise it should be exactly similar except how the block is being casted.
in mxfp8 we use a separate function to cast the values.

TransformerEngine/transformer_engine/pytorch/triton_kernels/cast_transpose.py

Line 301 in ef83316

def float_to_e8m0_triton(val: tl.float32) -> tl.uint8:

So we can just put this part in a separate function:

TransformerEngine/transformer_engine/pytorch/triton_kernels/cast_transpose.py

Lines 562 to 571 in ef83316

	# Nearest-neighbor quantization to E2M1 values
	# E2M1 representable values: {0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0}
	idx_row = tl.zeros([MXFP4_BLOCK_SIZE, MXFP4_BLOCK_SIZE], dtype=tl.uint8)
	idx_row = tl.where(abs_qx_row >= 0.25, 1, idx_row) # → 0.5
	idx_row = tl.where(abs_qx_row >= 0.75, 2, idx_row) # → 1.0
	idx_row = tl.where(abs_qx_row >= 1.25, 3, idx_row) # → 1.5
	idx_row = tl.where(abs_qx_row >= 1.75, 4, idx_row) # → 2.0
	idx_row = tl.where(abs_qx_row >= 2.5, 5, idx_row) # → 3.0
	idx_row = tl.where(abs_qx_row >= 3.5, 6, idx_row) # → 4.0
	idx_row = tl.where(abs_qx_row >= 5.0, 7, idx_row) # → 6.0

and alternate the calls based on whether it's a MXFP8 kernel or MXFP4 kernel.

The following operation can be replaced with the exp2f_rcp_triton function

TransformerEngine/transformer_engine/pytorch/triton_kernels/cast_transpose.py

Lines 549 to 551 in ef83316

	scale_unbiased_row = tl.log2(tl.maximum(amax_rounded, 1e-45)).floor() - 2
	scale_unbiased_row = tl.clamp(scale_unbiased_row, min=-127.0, max=127.0)
	quant_scale_row = tl.exp2(-scale_unbiased_row)

Another difference seems to be that MXFP8 kernels doesn't have the shuffle op fused in it. Do we want that mxfp8 as well? or is that specific to fp4 data?

Also @wangye805 noticed
For the column-wise quantization of an input tensor MxN, the mxfp8 seems to still have columnwise output as shape MxN but mxfp4 implementation is NxM. Technically we can add something like STORE_TRANSPOSE, to the kernel and invert the strides in the mxfp8 kernel itself.

Let me know what you think @sarthak-amd

[sudhu2k]

New changes

1. Refactor of kernel wrappers and mxfp4 tensors to be consistent with other recipes.
2. Decomposition of repeated work into a separate function for the mxfp4 quantize kernel.
3. Added shuffle tests.
4. Bug fixes.

[wangye805]

Okay, currently the target vs ref comparison is not coupled: we validated quantized data content and scale independently and allow for certain mismatch rate.

A better way is to do the validation jointly by adjusting the quanized value if the scale is mismatched:

TransformerEngine/tests/cpp/operator/test_cast_mxfp8.cu

Lines 314 to 353 in 9b31283

	#ifdef __HIP_PLATFORM_AMD__
	const double abs_tolerable_mismatches_limit = 1.0;
	const double rel_tolerable_mismatches_limit = 1.0e-4;
	#else
	const double abs_tolerable_mismatches_limit = 0.0;
	const double rel_tolerable_mismatches_limit = 0.0;
	#endif
	std::vector<size_t> mismatches_scales_indices;
	size_t mismatches_scales = 0;
	compare_e8m0_scaling_factors("scales", gpu_scales_ptr, ref_output_scales.get(),
	unpadded_blocks_Y, unpadded_blocks_X, scales_stride,
	mismatches_scales_indices, mismatches_scales,
	scale_diff_abs_tolerance,
	abs_tolerable_mismatches_limit,
	rel_tolerable_mismatches_limit);
	#ifdef __HIP_PLATFORM_AMD__
	if (::testing::Test::HasFatalFailure()) return;
	adjust_ref_for_e8m0_scale_error("scales", mismatches_scales_indices, gpu_scales_ptr,
	ref_output_scales.get(), scales_stride, rows, cols, rowwise,
	ref_output_c.get(), otype);
	mismatches_scales = 0;
	#endif
	const size_t mismatches_elts = 32 * mismatches_scales;
	auto [atol, rtol] = getTolerances(otype);
	compareResults("output_c", output_c, ref_output_c.get(), rowwise, atol, rtol, true, mismatches_elts);
	if (processing_method == ProcessingMethod::CAST_DBIAS
	|| processing_method == ProcessingMethod::CAST_DBIAS_DACT)
	{
	auto [atol_dbias, rtol_dbias] = getTolerances(itype);
	if (itype == DType::kFloat32) {
	atol_dbias = 1e-4;
	rtol_dbias *= sqrt(static_cast<double>(rows)) ;
	} else {
	rtol_dbias *= 4;
	}
	compareResults("output_dbias", output_dbias, ref_output_dbias.get(), true, atol_dbias, rtol_dbias);