Fix CK grouped-GEMM fallback for fused wgrad accumulation (BF16 in / FP32 out)

[sudhu2k]

Description

On ROCm, the grouped-GEMM backend selector gated the CK fast path behind is_supported_dtype(), which for 16-bit inputs required A_dt == B_dt == D_dt. This excluded the very common fused wgrad-accumulation case, where GroupedLinear/MoE backward accumulates the weight gradient into an FP32 main_grad buffer, i.e. BF16 inputs (A/B) with an FP32 output (D). As a result every wgrad grouped GEMM fell back to the multi-stream hipBLASLt path.

This PR relaxes the gate to allow D to be either the matching 16-bit type or Float32 for 16-bit inputs, routing the fused-accumulation wgrad GEMM through CK instead of hipBLASLt.

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:

• Relax the ROCm is_supported_dtype() check in nvte_multi_tensor_gemm (transformer_engine/common/gemm/cublaslt_gemm.cu) so 16-bit (BF16/FP16) inputs with matching A/B are accepted when the output D is either the same 16-bit type or Float32.
• This lets the fused wgrad-accumulation path (BF16 in / FP32 main_grad out) use the CK grouped-GEMM backend instead of falling back to hipBLASLt with reason=unsupported_dtype.

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

[ipanfilo]

Corresponding GEMMTestSuite needs to be added to test_cublas_gem

[matthiasdiener]

There are no grouped GEMM tests in test_cublaslt_gemm.cu currently, it's the single GEMM test suite. Would it be more natural to put these with the existing grouped GEMM tests? Two options:

• C++: add the tests to test_grouped_gemm.cu, or following the pattern in test_ck_grouped_mxfp8.cu (which already has a fallback detection); or
• pytest: the scenario this PR fixes exists already as a test (but is disabled on ROCm) in

  TransformerEngine/tests/pytorch/test_grouped_linear.py

  Lines 301 to 302 in ebbd623

  	if IS_HIP_EXTENSION and dtype not in (torch.float32,) and fuse_wgrad_accumulation and not fp8:
  	pytest.skip(f"ROCm does not support fused wgrad accumulation for {dtype}.")

  
  For the fallback-detection pattern, the restored test in restore test_grouped_gemm_unaligned pytest #703 could be reused

[sudhu2k]

I was originally thinking of adding the tests to test_grouped_gemm.cu but I noticed that we currently remove that file from testing in

TransformerEngine/tests/cpp/operator/CMakeLists.txt

Lines 38 to 47 in ebbd623

	test_multi_swizzle.cu
	test_swap_first_dims.cu
	test_grouped_gemm.cu #CUDA-only test
	../test_common.cu)
	if(USE_ROCM)
	get_target_property(test_cuda_sources test_operator SOURCES)
	# Remove CUDA-only tests and add ROCm specific ones
	list(REMOVE_ITEM test_cuda_sources
	test_grouped_gemm.cu)

So I've added the tests to test_cublas_gemm.cu itself, since the groupedgemm implementation is part of the cublaslt_gemm.cu file.
I can create a test_ck_grouped_gemm.cu file, and add it there, but let me know!
The .cu tests that I've added already checks and fails if the test falls back to hipblasLT multi-stream grouped gemm.

[sudhu2k]

I think test_grouped_linear_accuracy_rocm_backends pytest is more suited to test this PR's scenario on the module level. I've added a similar fallback detection pattern that one of your test uses in 3b0eeba

[matthiasdiener]

Thanks! Not sure if we need to keep the new C++ tests in addition to the pytest ones, but I'll defer to @ipanfilo and you on that.

[ipanfilo]

test_grouped_gemm.cu is for grouped_gemm API. while this PR is for multi_tensor_gemm. I think it is more appropriate to make test_ck_grouped_gemm based on test_ck_grouped_mxfp8

[sudhu2k]

Moved the test to test_ck_grouped_gemm.cu in 6774018

[sudhu2k]

Hi @ipanfilo, could you please take a quick look at this PR again? Thanks!

[ipanfilo]

Well, I mean repurpose test_ck_grouped_mxfp8 to be more generic. Sorry for confusion, it can be separate PR

[ipanfilo]

Well, I mean repurpose test_ck_grouped_mxfp8 to be more generic. Sorry for confusion, it can be separate PR