Skip to content

[HIP] [CI] [JIT] [Feature] Add gfx:cu_num build targets so one build serves multiple SKUs of an arch - #5340

Draft
RElbers wants to merge 20 commits into
mainfrom
cu-aware-targets/gpu-target-resolution
Draft

RElbers wants to merge 20 commits into
mainfrom
cu-aware-targets/gpu-target-resolution

Conversation

@RElbers

@RElbers RElbers commented Sep 8, 2026 •

Copy link
Copy Markdown
Contributor

Motivation

AITER selects tuned GEMM kernels at build time using (gfx, cu_num, ...). GPUs sharing an architecture can have different CU counts. For example, gfx950 with 128 CUs on MI350P and 256 CUs on MI350X/MI355X. GPU_ARCHS and a single global CU_NUM cannot request both configurations in one build. This PR adds AITER_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_TARGETS is optional. Builds that leave it unset continue to use GPU_ARCHS, CU_NUM, or native detection, but this PR also changes some of their behavior:

  • For explicitly named GPU_ARCHS targets, omitting CU_NUM previously 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.
  • For multiple named architectures, 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.
  • Explicit GPU_ARCHS builds previously accepted zero and negative CU_NUM overrides; now only positive integers are accepted.
  • _parse_gpu_archs_env previously preserved architecture-name case; it now normalizes names to lowercase, so entries such as GFX942 resolve as gfx942.

Build target resolution

AITER_GPU_TARGETS GPU_ARCHS CU_NUM Build targets
gfx950:128;gfx950:256 ignored ignored gfx950:128, gfx950:256
gfx942;gfx950:128 ignored ignored gfx942:304, gfx950:128
unset gfx942;gfx950 unset gfx942:304, gfx950:2561
unset gfx942;gfx950 80 gfx942:80, gfx950:80
unset gfx950 unset, no GPU gfx950:256
unset unset or native unset live architecture at its live CU count
unset unset or native 80 live architecture at 80
unset unset or native any, no GPU error

CU_NUM applies to every architecture named in GPU_ARCHS. A bare AITER_GPU_TARGETS entry takes its
architecture's default count. When AITER_GPU_TARGETS is set, a GPU_ARCHS naming a different architecture set warns. A CU_NUM matching none of the targets warns.

Implementation/Changes

  • Target parsing and selection. build_targets.py parses AITER_GPU_TARGETS into (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 replace GPU_ARCHS/CU_NUM for 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.py and cpp_itfs/utils.py derive compiler architecture lists from the new targets when set. Both validators use KNOWN_GFX instead of separate allowlists. The gfx1250 compiler flags and cluster-launch define check membership in get_gfx_list() instead of reading GPU_ARCHS directly; 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 raw GPU_ARCHS value, 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_NUM requests a count absent from AITER_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 that CU_NUM still overrides runtime lookup.

  • Release and nightly workflows. The reusable release workflow accepts a gpu_targets input and passes it into the build containers as AITER_GPU_TARGETS. Automated release and default nightly builds specify gfx942: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_TARGETS mechanism only. Both are addressed in the follow-ups below.

Test Plan

  • Extend GEMM code-generation tests to cover multiple CU counts for one architecture, target deduplication, default counts, environment precedence, and reporting targets without tuned rows.
  • Add GPU target-resolution tests with the live GPU probes mocked. Cover live CU detection for a binned part, order-independent architecture selection, gfx1250 compiler flags, and cache separation.
  • Run the existing GEMM dispatch/code-generation and JIT-cache transaction regression suites.

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

  1. The architecture matching the live device takes that device's CU count. ↩

@github-actions

github-actions Bot commented Sep 8, 2026

Copy link
Copy Markdown
Contributor

🏷️ CI Guide

Runs automatically on every PR:

  • ✅ Pre-checks (submodule verification, code formatting)
  • ✅ Aiter op tests (gfx942 + gfx950)
  • ✅ Triton tests on MI35X (only when aiter/ops/triton/** or related paths are changed)

Extended tests (opt-in via labels):

Label Tests
ci:gfx1250-ffm-triton Run the five-shard gfx1250 FFM Triton test suite
ci:triton-300x Run an additional Triton test job on MI300X in PRs; main branch always runs both MI35X and MI300X
multigpu Aiter multi-GPU tests on the 8-GPU runner
ci:sglang SGLang integration tests: DeepSeek-R1-MXFP4 accuracy, Qwen 3.5 accuracy
ci:atom ATOM benchmark: DeepSeek-R1-0528, GPT-OSS-120B
ci:atom_full ATOM accuracy suite for PR and main models from ATOM models_accuracy.json
ci:vllm vLLM benchmark: GPT-OSS-120B, DeepSeek-R1-0528, Kimi-K2.5
ci:all All standard extended tests (excludes ci:atom_full)

Only add ci:atom_full for FlyDSL or Triton upgrades.
Add labels via the sidebar or gh pr edit 5340 --add-label <label>

PR title tags & labels:
Component tags ([Triton/Gluon], [HIP], [CK], [ASM], ...) are added to the PR title and as PR labels automatically from the changed files and re-synced on every push — change-type tags like [fix]/[Perf], op tags like [MLA], and human labels (ci:*) are left untouched. Add the no-auto-title label to opt this PR out.

@RElbers
RElbers marked this pull request as ready for review September 8, 2026 13:48
@RElbers
RElbers requested a review from a team September 8, 2026 13:48
@github-actions github-actions Bot changed the title [Feature] Add gfx:cu_num build targets so one build serves both SKUs of an arch [HIP] [JIT] [Feature] Add gfx:cu_num build targets so one build serves both SKUs of an arch Sep 8, 2026
@github-actions github-actions Bot changed the title [HIP] [JIT] [Feature] Add gfx:cu_num build targets so one build serves both SKUs of an arch [HIP] [CI] [JIT] [Feature] Add gfx:cu_num build targets so one build serves both SKUs of an arch Sep 8, 2026
@github-actions github-actions Bot added the CI label Sep 8, 2026
@RElbers
RElbers force-pushed the cu-aware-targets/gpu-target-resolution branch from ad04192 to 9c06df2 Compare September 8, 2026 15:37
@zufayu
zufayu requested a review from gyohuangxin September 9, 2026 01:24
@RElbers RElbers changed the title [HIP] [CI] [JIT] [Feature] Add gfx:cu_num build targets so one build serves both SKUs of an arch [HIP] [CI] [JIT] [Feature] Add gfx:cu_num build targets so one build serves multiple SKUs of an arch Sep 9, 2026
@samremes

samremes commented Sep 9, 2026 •

Copy link
Copy Markdown
Contributor

Review (SKU expression / non-comm enablement)

This is the build-time half of SILOTIGER-971 / SILOTIGER-978: AITER_GPU_TARGETS=gfx950:128;gfx950:256 is the right knob. Not adding gfx950:128 rows here is correct (SILOTIGER-972 comes after dispatch is safe). Runtime XCC identity stays in #5182; this PR should key tuning on CU count.

The exact (gfx, cu_num) filter and the release pin gfx942:304;gfx950:256 match the plan. The get_gfx() last-entry-wins fix on multi-arch wheels is necessary.

Request changes on one P1 before merge:

  1. Runtime arch still is not the active HIP device. _detect_native takes the first gfx token from rocminfo. For a multi-arch wheel this replaces last-entry-wins with first-rocminfo-wins. On a mixed host that is not the device that will launch. Resolve from the current HIP device (same direction as [JIT] Identify MI350P, and stop hardcoding the XCD count #5182's per-device XCC/PCI queries). Add a test where the first enumerated arch differs from the active device.

Also worth fixing in this PR (P2):

  1. Bare GPU_ARCHS live-CU probe: a failed get_cu_num() / get_gfx_runtime() falls through a broad except to the full-SKU table default with no warning. That is indistinguishable from a GPU-less cross-build and can bake 256-CU identity on a 128-CU box.

  2. Precedence vs CU_NUM. Implementation lets AITER_GPU_TARGETS win over CU_NUM; docs/installation.rst only says it overrides GPU_ARCHS. Same-arch conflict AITER_GPU_TARGETS=gfx950:256 + CU_NUM=128 does not warn.

  3. Tests miss the new validation surface. Unknown-arch still goes through GPU_ARCHS; no case for malformed AITER_GPU_TARGETS, live-CU substitution, probe failure, or the unmatched-target warning after libtype filtering.

Merge order: land this (after the P1) before #5341 / #5342.

@RElbers
RElbers added this pull request to stack #5449 September 11, 2026 11:13
@RElbers
RElbers removed this pull request from stack #5449 September 11, 2026 15:19
@RElbers
RElbers added this pull request to stack #5466 September 11, 2026 15:20
- 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.
@RElbers
RElbers force-pushed the cu-aware-targets/gpu-target-resolution branch from a44faa5 to c1736a3 Compare September 11, 2026 15:50
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.
@vgokhale

vgokhale commented Oct 7, 2026

Copy link
Copy Markdown
Contributor

@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.
Copilot AI balanced review requested due to automatic review settings October 8, 2026 12:25
@RElbers

RElbers commented Oct 8, 2026

Copy link
Copy Markdown
Contributor Author

@RElbers is this PR still valid? It has seen no activity in a month.

Yes, it's still valid. I just synced with latest main again.

Copilot AI left a comment

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.

🟡 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.

Comment on lines +137 to +140
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"]
Comment thread csrc/cpp_itfs/utils.py
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()
@RElbers
RElbers marked this pull request as draft October 8, 2026 20:09
@RElbers

RElbers commented Oct 8, 2026

Copy link
Copy Markdown
Contributor Author

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.

@zufayu
zufayu removed the request for review from gyohuangxin October 9, 2026 00:29

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

4 participants