Skip to content

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

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
sudhu2k wants to merge 6 commits into
base: dev
Choose a base branch
from
Open

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

Changes from 2 commits
Commits
Show all changes
6 commits
Select commit Hold shift + click to select a range
a0a7884
Support FP32 accumulation in multi-tensor GEMM for 16-bit inputs
sudhu2k Aug 11, 2026
517e8ba
Clarify supported dtype combinations in multi-tensor GEMM comment
sudhu2k Aug 11, 2026
3b0eeba
Add grouped GEMM test suite for CK backend with fallback detection
sudhu2k Aug 12, 2026
1637b2b
Assert CK grouped GEMM does not fall back to hipBLASLt for supported …
sudhu2k Aug 12, 2026
6774018
Move CK grouped GEMM tests to dedicated test file
sudhu2k Aug 14, 2026
53561fd
Remove unwanted includes
sudhu2k Aug 21, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
8 changes: 7 additions & 1 deletion transformer_engine/common/gemm/cublaslt_gemm.cu
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1187,8 +1187,14 @@ void nvte_multi_tensor_gemm(const NVTETensor *A, const NVTETensor *B, NVTETensor
auto B_dt = effective_dtype(inputB);
auto D_dt = OutputD->data.dtype;

// Supported grouped-GEMM dtype combinations for the CK path:
// - FP8 inputs: A and B both FP8 (D dtype handled by the FP8 path).
// - 16-bit inputs: A and B the same 16-bit type (both BF16 or both FP16),
// with output D either that same 16-bit type or Float32
// (fused wgrad accumulation into an FP32 main_grad buffer).
return ((is_fp8_dtype(A_dt) && is_fp8_dtype(B_dt)) ||
((A_dt == B_dt) && (A_dt == D_dt) && is_fp16_dtype(A_dt)));
((A_dt == B_dt) && is_fp16_dtype(A_dt) &&
(A_dt == D_dt || D_dt == transformer_engine::DType::kFloat32)));

@ipanfilo ipanfilo Aug 11, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Corresponding GEMMTestSuite needs to be added to test_cublas_gem

@matthiasdiener matthiasdiener Aug 12, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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 sudhu2k Aug 12, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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 sudhu2k Aug 12, 2026 •
edited
Loading

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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 matthiasdiener Aug 12, 2026 •
edited
Loading

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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 ipanfilo Aug 14, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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 sudhu2k Aug 14, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Moved the test to test_ck_grouped_gemm.cu in 6774018

@sudhu2k sudhu2k Aug 18, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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

@ipanfilo ipanfilo Aug 21, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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


#else
auto A_type = get_cuda_dtype(inputA->data.dtype);
Expand Down
Loading