Repository navigation
Conversation
The CK GEMM supported-architecture allowlist currently lives hardcoded in three
places: `torch_filter_ck_gemm_arches` (cmake/public/utils.cmake), the runtime
check `at::Context::ckGemmSupported()`, and the `gemm_internal_ck` dispatch in
`ck_gemm_{half,bfloat16}.hip`. Keeping them in sync is manual, and the
gfx11/gfx12 WMMA path (pytorch#143971) can only be re-enabled by editing source in
several spots.
This PR makes the list a single source of truth:
- A new CMake cache variable `TORCH_CK_GEMM_ARCH_LIST` (default
`gfx90a;gfx942;gfx950`) drives the `ck_gemm` build arch filter and is surfaced
to C++ through a generated `ATen/ROCmCKGEMMConfig.h` header, mirroring the
existing `AT_ROCM_CK_SDPA_ARCHS` mechanism. It is forwarded from the
environment via `cmake/EnvVarForwarding.cmake`.
- `Context.cpp` parses the list once through a shared `parseArchList` helper
used by both `ckSDPASupported()` and `ckGemmSupported()`.
- The `.hip` dispatch calls `at::globalContext().ckGemmSupported()` instead of
re-parsing the list, so the header include, the parser, and the fallback are
removed from both `.hip` files.
A WMMA arch can now be opted back in at both build and run time without patching
source, e.g. `TORCH_CK_GEMM_ARCH_LIST="gfx90a;gfx942;gfx950;gfx1200"`.
No. Default behavior is unchanged -- the default list is the same three archs
that were previously hardcoded.
@jeffdaily (requested this refactor in review).
CMake wiring, standalone project including `cmake/public/utils.cmake`:
- default: `torch_filter_ck_gemm_arches` filters
`gfx942;gfx1250-strict;gfx950;gfx1200` to `gfx942;gfx950`; generated header
`#define TORCH_CK_GEMM_ARCH_LIST "gfx90a,gfx942,gfx950"`.
- override `TORCH_CK_GEMM_ARCH_LIST="gfx90a;gfx942;gfx950;gfx1200"`: filter keeps
`gfx942;gfx950;gfx1200`; header `"gfx90a,gfx942,gfx950,gfx1200"`.
Full build:
```
PYTORCH_ROCM_ARCH="gfx90a gfx942 gfx950" pip install -e . --no-build-isolation -v
```
rc=0, 0 failed. `ck_gemm_half.hip.o` offload bundle = `gfx90a,gfx942,gfx950`;
`libtorch_hip.so` links (the `ck_gemm` -> `Context::ckGemmSupported()` cross-TU
call resolves).
Pull Request resolved: pytorch#201113
Approved by: https://github.com/jeffdaily
Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com>
Co-authored-by: Jeff Daily <jeff.daily@amd.com>
(cherry picked from commit 5182982)
(cherry picked from commit a117134)
This branch has not been deployed
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Related: ROCm/TheRock#9413
Backport of pytorch/pytorch@5182982d.
Problem
The release branch builds CK GEMM for every requested architecture except gfx1250. Adding gfx1170/gfx1171 therefore compiles unsupported CK kernels and breaks TheRock wheel builds.
Changes
TORCH_CK_GEMM_ARCH_LISTthe build-time and runtime source of truth.Tracked in AIPYTORCH-1538.
Validation
gfx942;gfx1170;gfx1250-strict;gfx950togfx942;gfx950.git diff --checkpassed.gfx94X-dcgpu+gfx117X-allvalidation build compiled PyTorch successfully.