GEMM reference computation offload

[matthiasdiener]

Description

Introduce a HIP implementation of the GEMM reference computation to speed up these computations.

Partly addresses https://github.com/ROCm/frameworks-internal/issues/14746

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:

• Introduce a HIP implementation of the GEMM reference computation

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

[Copilot]

Pull request overview

This PR introduces a GPU-accelerated implementation of the GEMM reference computation using HIP/CUDA kernels to improve performance over the previous CPU-based implementation. The reference computation is critical for validating GEMM operations, and offloading it to the GPU significantly speeds up testing.

Key Changes

• Replaced CPU OpenMP-based reference GEMM computation with GPU kernel implementation
• Introduced compute_ref_kernel to perform matrix multiplication, bias addition, GELU activation, and scaling on GPU
• Refactored both tensor-wise and MXFP8 code paths to use a common compute_ref_impl function that manages device memory and kernel execution

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[alextmagro]

Hi Matthias! Looks great, just have a couple performance questions

[wangye805]

We can adapt existing test tensor classes (

TransformerEngine/tests/cpp/test_common.cu

Line 226 in 669b556

Tensor::Tensor(const std::string& name,

) and their space allocation functions (

TransformerEngine/tests/cpp/operator/test_cast_mxfp8.cu

Line 237 in 669b556

Tensor output_c("output_c", shape, otype, rowwise, colwise, NVTE_MXFP8_1D_SCALING);

) defined in tests/cpp/test_common.cu instead of reinventing.

In fact, we can change the api of reference computing by taking directly const tensor& therefore we don't need to re-allocate the input and do one extra copy

[matthiasdiener]

What do you think of 3ecea7f? This also merges the mxfp8/non-mxfp8 paths.

[wangye805]

Thanks for consolidating with existing apis in test_common.cu.

In fact, I still see some cudaMalloc and cudaFree, which can be replaced by using existing test tensor class apis.
For example, the device pointer for scale (

TransformerEngine/tests/cpp/operator/test_cublaslt_gemm.cu

Line 246 in e32fb3d

NVTE_CHECK_CUDA(cudaMalloc(&d_a_scale_packed, a_scale_packed.size() * sizeof(fp8e8m0)));

), its corresponding test tensor allocation can be found here:

TransformerEngine/tests/cpp/test_common.cu

Lines 321 to 335 in 2bc74c8

	if (rowwise) {
	(void)cudaMalloc((void**)&rowwise_scale_inv, rowwise_scale_size); // NOLINT(*)
	(void)cudaMemset(rowwise_scale_inv, 0, rowwise_scale_size);
	rowwise_scale_inv_cpu_data_ = std::make_unique<unsigned char[]>(rowwise_scale_size);
	std::fill_n(rowwise_scale_inv_cpu_data_.get(), rowwise_scale_size, 0);
	auto scale_dtype = rowwise_scale_meta.type;
	tensor_.set_rowwise_scale_inv(rowwise_scale_inv, scale_dtype, scale_shape);
	}
	if (columnwise) {
	(void)cudaMalloc((void**)&columnwise_scale_inv, columnwise_scale_size); // NOLINT(*)
	(void)cudaMemset(columnwise_scale_inv, 0, columnwise_scale_size);
	columnwise_scale_inv_cpu_data_ = std::make_unique<unsigned char[]>(columnwise_scale_size);
	std::fill_n(columnwise_scale_inv_cpu_data_.get(), columnwise_scale_size, 0);
	auto scale_dtype = colwise_scale_meta.type;
	tensor_.set_columnwise_scale_inv(columnwise_scale_inv, scale_dtype, columnwise_scale_shape);

[matthiasdiener]

I replaced the remaining raw allocations in the reference path with test::Tensor for the temporary device buffers (RefD/RefGelu/RefAmax) in e11e400.

[wangye805]

I see. Yeah, it indeed saved some cudaMalloc/cudaFrees.

How about we put the RefD instantiation inside PerformTest, and pass the Tensor RefD (including its RefAmax D) and RefPreGeluOut to run_reference directly (instead of std::unique_ptr<D_Type[]>& ref_D, float* ref_amax_d, std::unique_ptr<Gelu_Type[]>& ref_pre_gelu_out). Then this can save some ref cpu ptr allocation.

[matthiasdiener]

What do you think of 325ece6?

[wangye805]

LGTM