Repository navigation
[HIP] [CI] [JIT] [Feature] Add gfx:cu_num build targets so one build serves multiple SKUs of an arch - #5340
[HIP] [CI] [JIT] [Feature] Add gfx:cu_num build targets so one build serves multiple SKUs of an arch#5340RElbers wants to merge 20 commits into
Conversation
🏷️ CI GuideRuns automatically on every PR:
Extended tests (opt-in via labels):
PR title tags & labels: |
ad04192 to
9c06df2
Compare
Review (SKU expression / non-comm enablement)This is the build-time half of SILOTIGER-971 / SILOTIGER-978: The exact Request changes on one P1 before merge:
Also worth fixing in this PR (P2):
|
- Adds AITER_GPU_TARGETS, a ';'- or ','-separated list of gfx or gfx:cu_num entries, parsed into (gfx, cu_num) build targets with arch names case-folded and validated and CU counts required positive. - Case folds GPU_ARCHS in the same parser, so every consumer of the arch set agrees on its spelling rather than only the one that builds targets. - Adds gpu_archs_env_names as the single definition of when GPU_ARCHS names an explicit arch set, so its readers cannot disagree over native or its case. - Adds get_build_archs_env so a caller can recover the arch set a build compiles. - Adds unmatched_targets, which names the build targets a tuning frame holds no rows for, leaving filter_tune_df a filter. - Covers parsing, validation, dedup, unmatched targets, and the GPU_ARCHS fallback in the tests. filter_tune_df compares the arch string exactly, so an unvalidated typo matched no row and detuned the whole build with no error; validation names the offending variable instead. An arch named only in AITER_GPU_TARGETS must still reach the offload-arch flags, or a module ships dispatch entries for device code it does not contain.
- Resolves the runtime dispatch arch from the live device when it is one of the named targets, else an order-independent max, replacing the last-entry rule. - Reads the live device CU count for a GPU_ARCHS build when CU_NUM is unset, so a binned or partitioned part is not resolved to the full-SKU count, and logs the substitution. - Makes AITER_GPU_TARGETS authoritative for the offload-arch flags, warning when GPU_ARCHS disagrees, and keys the gfx1250 cluster-launch flag off the resolved arch list, the same membership test the compiler flags use. - Inspects GPU_ARCHS through gpu_archs_env_names everywhere it is read, so the disagreement warning cannot fire on a spelling the rest of the build treats as native. - Validates the JIT arch set against KNOWN_GFX, so a target the parser accepts is not rejected afterwards by a second hand maintained list. - Selects and reports tuned rows in one place for both tune-dict builders, after the libtype filter so the report covers the rows a module would have used. - Covers runtime arch resolution and order independence in the tests. Reading the live CU count makes a bare GPU_ARCHS build depend on the builder's device, so release and CI builds should pin AITER_GPU_TARGETS or CU_NUM. The runtime max is lexicographic, chosen so reordering a target list cannot change what a cross-build host selects, and one helper applies it to a named target list and a multi-arch GPU_ARCHS alike.
- Keys the template-library build directory on the resolved arch set, so a specialization built for one arch cannot satisfy the built-check for another. - Widens the cpp_itfs arch allowlist to KNOWN_GFX and returns a sorted, deduplicated list, so both allowlists now derive from one table. - Resolves hsaco directories and the prebuilt MLA stage-1 object from the live device rather than the build-target string, keeping the compile_hsaco, check_hsaco, get_hsaco_launcher, and run_hsaco signatures unchanged. - Caches the resolved arch key and keeps the lookup cache on run_lib, so the per-launch path repeats no string work. - Scopes the gluon AOT temporary outputs and their lock by architecture. - Covers both cache paths in the codegen tests. Template libraries and hsaco files hold device code, but their paths were keyed on the build-target string alone, so a multi-arch build looked up a composite directory that no single-arch build ever wrote. The arch key is fixed for the life of a process, so caching it is safe.
- Adds an autotuning-pipeline section on tuning a second SKU that shares a gfx arch, keyed on cu_num so its catalog is additive rather than colliding. - Describes building with AITER_GPU_TARGETS: the gfx or gfx:cu_num list, its authority over GPU_ARCHS for the arch set, and the CU_NUM override. - Notes that build-time filtering makes adding rows require a rebuild, and that CU_NUM must not be used to impersonate hardware in a serving environment.
- Adds a gpu_targets input, blank by default so an unset value leaves the gpu_archs path unchanged. - Names gfx942:304 and gfx950:256 from the release and nightly callers, the same pair the architecture map already resolves. - Drops the nightly default to blank when gpu_archs is overridden, so a dispatch naming its own architectures is not silently re-pinned. Naming only architectures takes the CU count from whichever host ran the job, so one commit could resolve different targets on different builders.
- Moves the guidance out of the autotuning pipeline document, whose subject is the two CI workflows, since almost none of this material concerns CI. - Restructures it as four steps, tune, check for collisions, validate, and build, with the mechanism behind each one moved to footnotes. - Corrects two commands that could not work as printed, one naming a build variable the validation test never reads and one prebuilding no kernels under the default mode. - Adds the collision guard, keeping the gfx column through a deduplication rewrite, the build warning that names an unmatched target, and the JIT path for exercising a SKU without building a package. - States that CU_NUM is a build-time flag, and that some operators still read it during dispatch, so it is not inert in a serving environment. Config rows are filtered before code generation, so a package built without a target named contains no tuned kernels for it at all rather than merely leaving them unselected. That and the default prebuild mode compiling nothing are both easy to miss, and either one leaves a reader with a package that silently falls back to the default kernel.
…10 inputs and the form was full, so gpu_targets is passed by workflow_call only.
Adds the environment variable row and examples covering the default CU count, a single non-default one, and both in one build.
…r in the catalog untouched with nothing to signal the gap, and the build variable the page existed to explain is covered by the installation guide.
- Prefers the active HIP device in _detect_native, moving the rocminfo scan to _detect_native_rocminfo as the fallback for hosts with no HIP context. - Leaves _detect_native uncached, so a call made during a build cannot pin the rocminfo answer for a later caller that has a device to ask. - Replaces the lru_cache on get_gfx_runtime with a memo that stores only an answer the HIP context backed, keeping cache_clear for callers that reset it. - Reads the CU count from the active device in get_cu_num_custom_op before falling back to rocminfo, which also picks up the current partition's count. - Defaults _get_pci_chip_id to the active ordinal rather than device 0, so the MI308 chip id probe describes the launching device. - Bounds get_gfx_custom_op_core at one cache entry, which is all a function taking no arguments can hold. rocminfo names the first GPU agent on the host, which stops being the launching device once HIP_VISIBLE_DEVICES or torch.cuda.set_device has selected another. Tuned configs are keyed on (gfx, cu_num), so a host carrying two architectures dispatched another SKU's kernels with nothing to signal it, while the CU probe asserted that every agent reports the same count.
- Adds op_tests/test_gpu_target_resolution.py, covering arch and CU resolution from the active HIP device on a faked host whose launching device is not the first one enumerated. - Moves test_runtime_arch_resolution and test_cpp_itfs_cache_identity there from test_gemm_codegen.py, whose subject is GEMM build targeting and dispatch, which neither one exercises. - Leaves test_unmatched_targets behind, since matching tuned config rows to build targets is what that file is for. - Replaces the _check counter with plain asserts, so a failing case stops the test under pytest as well as under the file runner. - Drops the mock and SimpleNamespace imports the move left unused in test_gemm_codegen.py. The _check helper records a failure and returns, so pytest saw no exception and reported every test in test_gemm_codegen.py as passing whatever its checks did; only the __main__ block, which CI invokes, propagated the result. Plain asserts fail under both runners.
- Tells a host with no GPU apart from one whose arch and CU probe failed, and warns only for the second, where the default count can be wrong. - Names the targets that took a default count, and the two variables that pin one instead. Both hosts fall back to the per-arch default, and until now both did so in silence. On a host that has a GPU, that default can be a full SKU count baked onto a binned or partitioned part, which is the failure this branch exists to remove.
- Warns when CU_NUM names a count no AITER_GPU_TARGETS entry was built for, naming the count to add or telling the reader to drop the variable. - Documents CU_NUM in the installation guide, which described the variable AITER_GPU_TARGETS overrides without describing the variable itself. - Covers the substitution that gives a bare GPU_ARCHS the live device's count rather than the arch default, which nothing exercised. The two variables decide different things: the targets decide which kernels a build contains, and CU_NUM still decides the count the runtime looks them up by. Setting a count no target was built for therefore compiles one SKU's kernels and then asks for another's, and every tuned shape falls back.
- MainFunc now tests gfx1250 membership through get_gfx_list, which the test never provided, so every case reaching it raised NameError.
a44faa5 to
c1736a3
Compare
Resolves the chip_info.py conflict against the shared rocminfo output cache added upstream. The rocminfo arch detector keeps its new name and now reads through _rocminfo_output(), while the HIP-preferring _detect_native wrapper, the get_gfx_runtime memo, and upstream's ASIC revision helpers are all kept.
- Removes _active_device_index, _active_device_props and _active_device_arch, so runtime GPU identity comes from rocminfo alone again. - Folds the renamed rocminfo detector back into _detect_native and restores the lru_cache on get_gfx_runtime and get_asic_revision, dropping the hand rolled memo and the cache_clear shim that stood in for it. - Restores the rocminfo compute unit read in get_cu_num_custom_op and the fixed device argument in _get_pci_chip_id. - Collapses the missing GPU and failed probe warnings into one message, since rocminfo cannot tell those two cases apart. - Drops the active device test and reworks the binned part test to fake rocminfo, which leaves this layer with no torch dependency. The probes only changed the answer on a host whose devices differ in architecture or CU count, which is rare, and aiter does not special case partition mode anywhere else. They also made the answer depend on which device was current when the first caller ran, and callers reached during import fixed it there for the rest of the process.
Adds a table to the installation guide mapping AITER_GPU_TARGETS, GPU_ARCHS and CU_NUM to the targets a build compiles for. The existing entries describe each variable on its own, so the precedence between them, what a bare architecture entry defaults to, and what happens with no GPU present were not written down anywhere.
Separates runtime resolution, OPUS flags, and cache behavior into focused tests. Checks runtime architecture selection across both target orderings with clearly named inputs and expected results. Consolidates environment and cache cleanup so test cases remain isolated. Simplifies GPU mocks and temporary HSACO setup.
…u-target-resolution
|
@RElbers is this PR still valid? It has seen no activity in a month. |
…u-target-resolution - Keeps this branch's AITER_GPU_TARGETS arch resolution in validate_and_update_archs and carries main's Windows native-detection branch into it. - Keeps allowed_archs derived from KNOWN_GFX, which is a superset of the hardcoded list main still carries, so no target is lost. - Loads and probes template libs under this branch's arch-keyed build dir while adopting main's LIB_BASENAME and its Windows-safe CDLL mode.
Yes, it's still valid. I just synced with latest main again. |
There was a problem hiding this comment.
🟡 Changes recommended
Architecture normalization is incomplete, causing uppercase or whitespace-only GPU_ARCHS configurations to fail in affected build paths.
2 open findings
What changed in this PR
Adds explicit (gfx, CU count) build targets so one wheel can include tuned kernels for multiple GPU SKUs sharing an architecture.
Changes:
- Adds target parsing, validation, dispatch resolution, and tuning-row filtering.
- Architecture-scopes template caches and runtime HSACO/MLA paths.
- Documents and wires targets into release/nightly workflows with expanded tests.
| File | Description |
|---|---|
aiter/jit/utils/build_targets.py |
Parses and filters explicit build targets. |
aiter/jit/utils/chip_info.py |
Resolves build/runtime architecture and CU selection. |
aiter/jit/core.py |
Applies target architecture sets to JIT compilation. |
aiter/jit/optCompilerConfig.json |
Uses architecture membership for gfx1250 flags. |
csrc/cpp_itfs/utils.py |
Architecture-scopes template and HSACO caches. |
csrc/cpp_itfs/mla/asm_mla_decode_fwd.py |
Selects MLA assembly by runtime architecture. |
csrc/cpp_itfs/pa_gluon_aot/pa_decode_gluon_aot.py |
Scopes AOT files and locks by architecture. |
docs/installation.rst |
Documents target resolution and examples. |
op_tests/test_gpu_target_resolution.py |
Tests resolution, flags, caches, and HSACO lookup. |
op_tests/test_gemm_codegen.py |
Tests target parsing and unmatched tuning rows. |
op_tests/test_jit_cache_transaction.py |
Updates the mocked JIT environment. |
.github/workflows/aiter-release.yaml |
Passes targets into build containers. |
.github/workflows/nightly.yaml |
Configures default nightly targets. |
.github/workflows/release-automation.yaml |
Pins release architecture/CU pairs. |
🧠 Review effort: Balanced
Give feedback about Copilot approvals in this survey to enter a drawing for a $150 gift card.
| def gpu_archs_env_names() -> list[str]: | ||
| """Arch names GPU_ARCHS explicitly requests, empty when unset or native.""" | ||
| archs = _parse_gpu_archs_env(os.getenv("GPU_ARCHS") or "native") | ||
| return [a for a in archs if a != "native"] |
| if GPU_ARCH is None: | ||
| # AITER_GPU_TARGETS outranks GPU_ARCHS for the arch set, as in aiter/jit/core.py. | ||
| # HSACO paths are keyed separately by the live GPU. | ||
| GPU_ARCH = (os.environ.get("GPU_ARCHS") or "").strip() |
|
I've put the PR back in draft state, considering what I'm changing now based on latest merge with main. I expect to re-open within a couple of days. |

Motivation
AITER selects tuned GEMM kernels at build time using
(gfx, cu_num, ...). GPUs sharing an architecture can have different CU counts. For example,gfx950with 128 CUs on MI350P and 256 CUs on MI350X/MI355X.GPU_ARCHSand a single globalCU_NUMcannot request both configurations in one build. This PR addsAITER_GPU_TARGETS="gfx950:128;gfx950:256"so the existing code-generation path can include tuned kernels for both targets. It provides the target-selection mechanism; corresponding tuning data is still required.AITER_GPU_TARGETSis optional. Builds that leave it unset continue to useGPU_ARCHS,CU_NUM, or native detection, but this PR also changes some of their behavior:GPU_ARCHStargets, omittingCU_NUMpreviously selected each architecture's default CU count; it now uses the live device's count for the matching architecture, retaining defaults for other targets or when detection fails.get_gfx()previously selected the last entry; it now selects the detected architecture if included, otherwise the lexicographically greatest name, making the fallback independent of list order.GPU_ARCHSbuilds previously accepted zero and negativeCU_NUMoverrides; now only positive integers are accepted._parse_gpu_archs_envpreviously preserved architecture-name case; it now normalizes names to lowercase, so entries such asGFX942resolve asgfx942.Build target resolution
AITER_GPU_TARGETSGPU_ARCHSCU_NUMgfx950:128;gfx950:256gfx950:128,gfx950:256gfx942;gfx950:128gfx942:304,gfx950:128gfx942;gfx950gfx942:304,gfx950:2561gfx942;gfx95080gfx942:80,gfx950:80gfx950gfx950:256nativenative8080nativeCU_NUMapplies to every architecture named inGPU_ARCHS. A bareAITER_GPU_TARGETSentry takes itsarchitecture's default count. When
AITER_GPU_TARGETSis set, aGPU_ARCHSnaming a different architecture set warns. ACU_NUMmatching none of the targets warns.Implementation/Changes
Target parsing and selection.
build_targets.pyparsesAITER_GPU_TARGETSinto(gfx, cu_num)pairs, validates architecture names and positive CU counts, and removes exact duplicates. Entries without a CU count use the architecture default. These pairs replaceGPU_ARCHS/CU_NUMfor build-target selection and feed the existing CSV filtering, so two CU counts of one architecture can both contribute tuned rows.Architecture validation and gfx1250 flags.
jit/core.pyandcpp_itfs/utils.pyderive compiler architecture lists from the new targets when set. Both validators useKNOWN_GFXinstead of separate allowlists. The gfx1250 compiler flags and cluster-launch define check membership inget_gfx_list()instead of readingGPU_ARCHSdirectly; cluster launch still requires ROCm 7.0 or later.Template-library cache keys. C++ template-library cache paths now include the sorted, deduplicated architecture set. Different sets use separate caches;
gfx950;gfx942,gfx942;gfx950, and equivalent lists with duplicates share one. Paged-attention AOT temporary files and locks use the same key. Existing template libraries rebuild under the new paths on first use. The main JIT module cache is unchanged and does not distinguish CU counts.HSACO and MLA assembly paths. HSACO cache reads and default writes use
get_gfx_runtime(), as does MLA assembly lookup. This replaces paths based on the rawGPU_ARCHSvalue, which could contain multiple architecture names.Warnings and documentation. Build-target resolution reports when no live CU count can be read, and warns when
CU_NUMrequests a count absent fromAITER_GPU_TARGETS. CSV filtering warns about targets with no rows for the selected library type; those targets continue to use default kernels. The installation guide documents the new variable and thatCU_NUMstill overrides runtime lookup.Release and nightly workflows. The reusable release workflow accepts a
gpu_targetsinput and passes it into the build containers asAITER_GPU_TARGETS. Automated release and default nightly builds specifygfx942:304;gfx950:256, so their target CU counts do not depend on the build host.Out of scope
OPUS and FlyDSL AOT target-filtering integration and adding the gfx950:128 tuning rows is a separate effort and is out of scope here. This is the
AITER_GPU_TARGETSmechanism only. Both are addressed in the follow-ups below.Test Plan
Test Result
Verified on MI300X and MI355X :
python op_tests/test_gemm_codegen.py: 64 checks passed, 0 failed.python -m pytest -q op_tests/test_gpu_target_resolution.py op_tests/test_jit_cache_transaction.py: 56 tests passed, 0 failed.Follow-up PRs
Footnotes
The architecture matching the live device takes that device's CU count. ↩