blockwise fp8 gemm integration for gfx942 and gfx950#658
Conversation
…dant HIP guards, revert unnecessary common.h change
# Conflicts: # tests/cpp/operator/CMakeLists.txt
…hipkittens_optim # Conflicts: # transformer_engine/common/CMakeLists.txt # transformer_engine/common/gemm/kittens/CMakeLists.txt # transformer_engine/common/gemm/rocm_gemm.cu # transformer_engine/pytorch/quantization.py
…gemm_hipkittens_optim # Conflicts: # transformer_engine/common/gemm/kittens/cdna4/mxfp8_gemm.cpp # transformer_engine/pytorch/quantization.py
![github-actions[bot]](https://avatars.githubusercontent.com/in/15368?s=60&v=4){kind=small}
| [submodule "3rdparty/hipkittens_cdna3"] | ||
| path = 3rdparty/hipkittens_cdna3 | ||
| url = https://github.com/asdfvg123/HipKittens.git | ||
| branch = yeonsoo/cdna3_fp8 |
There was a problem hiding this comment.
Submodule points to a personal user fork on an unreleased branch — must be relocated before merge.
url = https://github.com/asdfvg123/HipKittens.gitis a personal repo, not a canonical one. The sibling3rdparty/hipkittenscorrectly points atHazyResearch/HipKittens.git. Anyone cloning ROCm/TransformerEngine will failgit submodule update --initthe moment this repo goes private or gets renamed. It also blocks reproducible ROCm CI wheel builds (this submodule is now in the init list in.github/workflows/rocm-wheels-build.yml).- The pinned branch
yeonsoo/cdna3_fp8is likewise a personal in-flight branch; the CDNA3 kittens sources need to land in a persistent upstream branch (HazyResearch or a ROCm-owned mirror) and this URL/branch should be updated before merge.
| NVTE_CHECK((m % 8) == 0, | ||
| "Blockwise FP8 GEMM requires N divisible by 8"); | ||
| NVTE_CHECK(inputB->scaling_mode != NVTE_BLOCK_SCALING_1D || (n % 16) == 0, | ||
| "Blockwise FP8 GEMM requires M divisible by 16 for 1D scaling"); |
There was a problem hiding this comment.
Error messages have M/N swapped relative to the values being checked.
The check is on m but the message says "N divisible by 8", and the check on n says "M divisible by 16". A user hitting either of these gets a message that names the wrong dimension, which will be very confusing to debug (especially since the arch impls internally swap kM=N/kN=M — a reader seeing the message will not know which side of the swap is being enforced).
Fix: name the dimension after the value actually being checked, e.g.
| NVTE_CHECK((m % 8) == 0, | |
| "Blockwise FP8 GEMM requires N divisible by 8"); | |
| NVTE_CHECK(inputB->scaling_mode != NVTE_BLOCK_SCALING_1D || (n % 16) == 0, | |
| "Blockwise FP8 GEMM requires M divisible by 16 for 1D scaling"); | |
| NVTE_CHECK((k % 16) == 0, | |
| "Blockwise FP8 GEMM requires K divisible by 16"); | |
| NVTE_CHECK((m % 8) == 0, | |
| "Blockwise FP8 GEMM requires M divisible by 8"); | |
| NVTE_CHECK(inputB->scaling_mode != NVTE_BLOCK_SCALING_1D || (n % 16) == 0, | |
| "Blockwise FP8 GEMM requires N divisible by 16 for 1D scaling"); |
(If the intent was to speak in the kernel's post-swap convention, then the two messages need to be swapped with each other — either way the current pairing is wrong.)
Also note: since the preceding NVTE_CHECK(inputB->scaling_mode == NVTE_BLOCK_SCALING_1D, ...) already forces B to be 1D, the inputB->scaling_mode != NVTE_BLOCK_SCALING_1D || guard on the last check is always false-branch — you can drop it and just check (n % 16) == 0 unconditionally.
| if IS_HIP_EXTENSION and get_device_compute_capability() in ((9, 4), (9, 5)): | ||
| expected_err_msg = None | ||
| expected_err_cls = RuntimeError |
There was a problem hiding this comment.
This weakens the "constraint enforced" tests to a near-tautology on ROCm.
Setting expected_err_msg = None here means the test now accepts any RuntimeError from the ROCm code path, regardless of whether it comes from the constraint we're supposedly enforcing or from an unrelated failure (setup error, shape mismatch, OOM, etc.). Every call site of cublas_gemm_test_constraint_enforced is paired with a specific constraint (illegal dtype, illegal 2Dx2D, split-accumulator, bgrad, gelu-unsupported…), and none of that specificity is preserved on ROCm.
Prefer either:
- Mapping the cuBLAS-flavored
expected_err_msgto the corresponding ROCm one (theNVTE_CHECKstrings you added inrocm_gemm.cuabove, e.g."Blockwise FP8 GEMM does not support ..."), so the test still asserts the right thing, or - Passing a per-callsite ROCm expected message via the caller instead of blanket-nulling it here.
| if key in _FP8_KEYS: | ||
| value = ( | ||
| get_torch_float8_e4m3_type() if key is tex.DType.kFloat8E4M3 | ||
| get_torch_float8_e4m3_type() if key == tex.DType.kFloat8E4M3 |
There was a problem hiding this comment.
This change also affects the CUDA path — please classify and document.
Custom_DType_Dict is shared code that both CUDA and ROCm execute. Changing is to == on the pybind11 enum is likely a real bug fix (enum identity via is is brittle across module reloads / repeated pybind imports), not a ROCm-specific tweak. Per the ROCm-fork upstream-compat rule, generic bug fixes to shared code are fine but should be called out explicitly (PR description or a short comment) so it can be upstreamed and future IFU merges know why the line diverges from NVIDIA's TE.
If instead this change is only motivated by a ROCm-observed failure, please add a one-line comment explaining that, or guard it with IS_HIP_EXTENSION.
| gpu_arch = get_device_compute_capability() | ||
| if gpu_arch in ((9, 4), (9, 5)): # TODO: enabled for gfx1250 when ready | ||
| return True, "" | ||
| return False, "Device arch gfx94x or newer is required for FP8 block scaling execution." |
There was a problem hiding this comment.
This drops is_fp8_block_scaling_quantization_available — that is a public API change.
is_fp8_block_scaling_quantization_available and the underlying check_fp8_block_scaling_quantization_support were being removed together with the corresponding re-export from transformer_engine/pytorch/__init__.py and the __all__ entry above. That symbol was exposed on transformer_engine.pytorch, so any downstream code (or ROCm-side tests / benchmarks / megablocks / megatron patches) that imported is_fp8_block_scaling_quantization_available(...) will now break with ImportError.
Two options:
- If block scaling GEMM being available implies block scaling quantization is available (the new
_compute_fp8_block_scaling_supportgating on(9,4)/(9,5)suggests yes), keep the old API as a thin alias tois_fp8_block_scaling_availableand mark it deprecated, so we don't break importers. - If we're intentionally taking the API away, please call this out in the PR description as a breaking change.
Claude review — PR 658Reviewed the ROCm-only 3-dot diff ( Verdict: solid direction, but there are a few things worth fixing before merge. Findings posted inline:
Copyright headers: OK — every touched file carries an up-to-date AMD 2026 header in the correct format; NVIDIA lines preserved on modified upstream files. |
|
|
Is CDNA3 support expected to be ported to HK? |
Currently HK supports CDNA3, CDNA4, and UDNA1. There is currently effort with the external HK team to move away from having separate git branches for each arch to having directory level separation. |
In this case can we avoid having 2 copies of HK submodule? |
Yes, looks like they merged it in late last week. I think that will require some integration/cmake changes so that the right headers are imported for each kernel. We'll probably want to move our kittens kernels into cdna3/cdna4/udna1 folders as well, so might be better as a follow up ticket? |
| emulated = get_device_compute_capability() >= (10, 0) | ||
| return supported and not emulated | ||
|
|
||
| def rocm_blockwise_unsupported_reason( |
There was a problem hiding this comment.
Here generally I think we should return a tuple -- True/False, followed by an optional reason. Maybe rename to rocm_blockwise_is_supported too.
| ): | ||
| if not fp8_blockwise_gemm_supported(): | ||
| pytest.skip("CUDA version does not support blockwise FP8 gemm.") | ||
| if IS_HIP_EXTENSION: |
There was a problem hiding this comment.
We could also potentially use the rocm function above here too?
| is_w_1d_scaled, | ||
| ) -> None: | ||
| # NOTE: BGRAD epilogue is not supported for fp8. | ||
| if IS_HIP_EXTENSION: |
There was a problem hiding this comment.
One of the great things about hipKittens is how easy it is to fuse epilogues into kernels. If you have time it might be worth looking into bgrad/accumulator support.
| ) -> None: | ||
| # e5m2 by e5m2 not supported. | ||
| if IS_HIP_EXTENSION: | ||
| expected_err_msg = "does not support e5m2 by e5m2 inputs" |
There was a problem hiding this comment.
e5m2 should be supported by hipkittens, have a look at the mxfp8 kernels to see how we template for that
| use_split_accumulator, | ||
| is_x_1d_scaled, | ||
| is_w_1d_scaled, | ||
| expected_err_msg="dimension requirement", |
There was a problem hiding this comment.
This change should be hip guarded
|
|
||
| #ifndef KITTENS_DTYPE_ENUM_DEFINED | ||
| #define KITTENS_DTYPE_ENUM_DEFINED | ||
| enum KittensDType { |
There was a problem hiding this comment.
I'd rather not have to redefine these. I think we can have a parent folder header file that contains these enums and any other shared values.
| #endif // KITTENS_SCALING_MODE_DEFINED | ||
|
|
||
| namespace blockwise_gfx942 { | ||
| void kittens_blockwise_fp8_gemm_impl_cdna3( |
There was a problem hiding this comment.
Rather than a namespace for gfx942 and gfx950, etc, I think we would be better off with a purely virtual class at the parent folder level, with child class definitions for functions within each architectural folder. That way we can track more easily what is available and what is not for each architecture.
| @@ -0,0 +1,1459 @@ | |||
| /************************************************************************* | |||
There was a problem hiding this comment.
Most of my comments regarding the cdna3 kernel applies here as well -- otherwise the kernel looks good.
| option(NVTE_KITTENS_USE_POWER_OF_2_SCALE | ||
| "Use HW E8M0 MFMA scaling for power-of-2 blockwise FP8 scales (gfx950)" ON) | ||
|
|
||
| function(try_enable_hipkittens_gemm) |
There was a problem hiding this comment.
Not sure this needs to bea function if we are only calling it once?
| return; | ||
| } | ||
| } | ||
| #endif |
There was a problem hiding this comment.
if hipkittens is not being used, we should add an nvte_error here for blockwise usage
4 participants
Description
Integrating blockwise FP8 gemm for gfx942 and gfx950 using hipkittens. Using 2 submodules for now. The hipkittens submodule uses separate branch for CDNA3 and CDNA4.
Fixes # (issue)
Type of change
Changes
Please list the changes introduced in this PR:
separate the kernel for gfx942 and gfx950. The host selects the kernel in runtime.
The kernel supports 1d2d, 1d1d GEMM with output bf16, fp32, fp16 / input fp8e5m2, fp8e4m3 / and different epilogues.
Checklist: