diff --git a/cmake/QoLA.cmake b/cmake/QoLA.cmake new file mode 100644 index 0000000..e5e0ee6 --- /dev/null +++ b/cmake/QoLA.cmake @@ -0,0 +1,232 @@ +# SPDX-License-Identifier: MIT +# Copyright (C) 2026, Advanced Micro Devices, Inc. All rights reserved. +# +# CMake integration for QoLA's ahead-of-time AITER kernel builder. +# +# Consumers include this file and call qola_add_modules() once per module +# group. It owns everything between "here is a manifest" and "here are the +# headers and libraries to link": +# +# * parsing the manifest's pinned AITER commit, +# * syncing the AITER source tree to it (or honouring an override), +# * invoking `qola build` for the requested group and architectures, +# * locating the resulting public headers and shared objects. +# +# Usage: +# +# include(${QOLA_DIR}/cmake/QoLA.cmake) +# qola_add_modules( +# GROUP aiter_gemm +# MANIFEST ${CMAKE_CURRENT_LIST_DIR}/qola_manifest.toml +# BUILD_DIR ${CMAKE_CURRENT_BINARY_DIR}/qola +# ARCHS ${MY_ARCHS} +# OUT_INCLUDE_DIR QOLA_GEMM_INCLUDE_DIR +# OUT_LIB_DIR QOLA_GEMM_LIB_DIR +# OUT_LIBS QOLA_GEMM_LIBS +# OUT_AITER_DIR QOLA_AITER_SOURCE_DIR) +# +# Environment overrides honoured by this module: +# QOLA_AITER_SOURCE_DIR / NVTE_AITER_SOURCE_DIR +# Build against an existing AITER tree and skip checkout entirely. +# QOLA_PREBUILT_DIR_ +# Skip the build and consume prebuilt /lib and /include. + +include_guard(GLOBAL) + +set(QOLA_CMAKE_DIR "${CMAKE_CURRENT_LIST_DIR}") +get_filename_component(QOLA_ROOT_DIR "${QOLA_CMAKE_DIR}/.." ABSOLUTE) + +# Extract the single `aiter_commit = "..."` line from a manifest. +function(qola_read_manifest_commit manifest out_var) + if(NOT EXISTS "${manifest}") + message(FATAL_ERROR "[QoLA] Manifest not found: ${manifest}") + endif() + file(STRINGS "${manifest}" _lines + REGEX "^[ \t]*aiter_commit[ \t]*=[ \t]*\"[^\"]+\"") + list(LENGTH _lines _count) + if(NOT _count EQUAL 1) + message(FATAL_ERROR + "[QoLA] Expected exactly one 'aiter_commit = \"...\"' line in " + "${manifest}, found ${_count}.") + endif() + list(GET _lines 0 _line) + string(REGEX MATCH "\"([^\"]+)\"" _unused "${_line}") + if("${CMAKE_MATCH_1}" STREQUAL "") + message(FATAL_ERROR + "[QoLA] Failed to parse 'aiter_commit' from ${manifest}.") + endif() + set(${out_var} "${CMAKE_MATCH_1}" PARENT_SCOPE) +endfunction() + +# Locate a Python interpreter once, reusing the caller's if already found. +function(qola_find_python out_var) + if(Python_EXECUTABLE) + set(${out_var} "${Python_EXECUTABLE}" PARENT_SCOPE) + return() + endif() + find_package(Python COMPONENTS Interpreter QUIET) + if(NOT Python_EXECUTABLE) + message(FATAL_ERROR + "[QoLA] Python interpreter not found; it is required to check out " + "and build AITER kernels.") + endif() + set(${out_var} "${Python_EXECUTABLE}" PARENT_SCOPE) +endfunction() + +# Run `python -m qola.cli `, failing the configure step on error. +function(qola_run_cli what) + qola_find_python(_py) + execute_process( + COMMAND ${CMAKE_COMMAND} -E env "PYTHONPATH=${QOLA_ROOT_DIR}:$ENV{PYTHONPATH}" + "${_py}" -m qola.cli ${ARGN} + RESULT_VARIABLE _rc + OUTPUT_VARIABLE _out + ERROR_VARIABLE _err + OUTPUT_STRIP_TRAILING_WHITESPACE + ERROR_STRIP_TRAILING_WHITESPACE) + if(NOT _rc EQUAL 0) + message(FATAL_ERROR "[QoLA] ${what} failed.\n${_out}\n${_err}") + endif() +endfunction() + +# Sync the AITER source tree named by a manifest, unless overridden. +# +# Sets to the tree to build against. Idempotent across +# consumers: the first caller performs the checkout and later callers with the +# same manifest + destination short-circuit. +function(qola_checkout_aiter) + set(_opts) + set(_one MANIFEST DEFAULT_DIR OUT_DIR) + set(_multi) + cmake_parse_arguments(QCA "${_opts}" "${_one}" "${_multi}" ${ARGN}) + + set(_aiter_dir "${QCA_DEFAULT_DIR}") + set(_skip FALSE) + foreach(_env QOLA_AITER_SOURCE_DIR NVTE_AITER_SOURCE_DIR) + if(DEFINED ENV{${_env}} AND NOT "$ENV{${_env}}" STREQUAL "") + set(_aiter_dir "$ENV{${_env}}") + set(_skip TRUE) + message(STATUS "[QoLA] Using AITER source from ${_env}=${_aiter_dir}; skipping checkout.") + break() + endif() + endforeach() + + qola_read_manifest_commit("${QCA_MANIFEST}" _sha) + set_property(DIRECTORY APPEND PROPERTY CMAKE_CONFIGURE_DEPENDS "${QCA_MANIFEST}") + + if(NOT _skip) + # Two consumers sharing a manifest must not each re-run the checkout: + # it is idempotent but not free, and it re-touches patched headers. + get_property(_done GLOBAL PROPERTY QOLA_CHECKOUT_DONE_${_sha}_${_aiter_dir}) + if(_done) + message(STATUS "[QoLA] AITER already synced to ${_sha} at ${_aiter_dir}.") + else() + qola_run_cli("AITER checkout to ${_sha}" + checkout + --manifest "${QCA_MANIFEST}" + --aiter-root "${_aiter_dir}") + set_property(GLOBAL PROPERTY QOLA_CHECKOUT_DONE_${_sha}_${_aiter_dir} TRUE) + message(STATUS "[QoLA] Synced ${_aiter_dir} to ${_sha}") + endif() + endif() + + if(NOT EXISTS "${_aiter_dir}/csrc/include") + message(FATAL_ERROR + "[QoLA] Could not find AITER sources at ${_aiter_dir}/csrc/include.") + endif() + set(${QCA_OUT_DIR} "${_aiter_dir}" PARENT_SCOPE) +endfunction() + +# Build (or locate prebuilt) kernel libraries for one manifest module group. +function(qola_add_modules) + set(_opts) + set(_one GROUP MANIFEST BUILD_DIR AITER_DIR + OUT_INCLUDE_DIR OUT_LIB_DIR OUT_LIBS OUT_AITER_DIR OUT_CONFIG_DIR) + set(_multi ARCHS LIBS) + cmake_parse_arguments(QAM "${_opts}" "${_one}" "${_multi}" ${ARGN}) + + if(NOT QAM_GROUP) + message(FATAL_ERROR "[QoLA] qola_add_modules: GROUP is required.") + endif() + if(NOT QAM_MANIFEST) + message(FATAL_ERROR "[QoLA] qola_add_modules: MANIFEST is required.") + endif() + if(NOT QAM_BUILD_DIR) + message(FATAL_ERROR "[QoLA] qola_add_modules: BUILD_DIR is required.") + endif() + + # Prebuilt bypass: consume an existing lib/ + include/ pair. + string(TOUPPER "${QAM_GROUP}" _group_uc) + set(_prebuilt_env "QOLA_PREBUILT_DIR_${_group_uc}") + if(DEFINED ENV{${_prebuilt_env}} AND NOT "$ENV{${_prebuilt_env}}" STREQUAL "") + set(_prebuilt "$ENV{${_prebuilt_env}}") + message(STATUS "[QoLA] ${QAM_GROUP}: using prebuilt libraries from ${_prebuilt}") + set(_include_dir "${_prebuilt}/include") + set(_lib_dir "${_prebuilt}/lib") + set(_config_dir "${_prebuilt}/configs") + else() + if(QAM_AITER_DIR) + set(_aiter_dir "${QAM_AITER_DIR}") + else() + qola_checkout_aiter( + MANIFEST "${QAM_MANIFEST}" + DEFAULT_DIR "${QAM_BUILD_DIR}/third_party/aiter" + OUT_DIR _aiter_dir) + endif() + + string(REPLACE ";" ";" _archs "${QAM_ARCHS}") + list(JOIN _archs ";" _archs_str) + message(STATUS "[QoLA] ${QAM_GROUP}: building kernels for ${_archs_str}") + qola_run_cli("build of group '${QAM_GROUP}'" + build + --manifest "${QAM_MANIFEST}" + --aiter-root "${_aiter_dir}" + --output-dir "${QAM_BUILD_DIR}" + --group "${QAM_GROUP}" + --arch "${_archs_str}" + --skip-checkout) + set(_include_dir "${QAM_BUILD_DIR}/include") + set(_lib_dir "${QAM_BUILD_DIR}/lib") + set(_config_dir "${QAM_BUILD_DIR}/configs") + if(QAM_OUT_AITER_DIR) + set(${QAM_OUT_AITER_DIR} "${_aiter_dir}" PARENT_SCOPE) + endif() + endif() + + if(NOT EXISTS "${_include_dir}/qola_config.h") + message(FATAL_ERROR + "[QoLA] ${QAM_GROUP}: public headers missing at ${_include_dir}.") + endif() + + # Resolve the built shared objects. Explicit LIBS win; otherwise take + # whatever the group produced in lib/. + set(_libs) + if(QAM_LIBS) + foreach(_lib ${QAM_LIBS}) + if(NOT EXISTS "${_lib_dir}/${_lib}") + message(FATAL_ERROR + "[QoLA] ${QAM_GROUP}: expected library ${_lib} not found in ${_lib_dir}.") + endif() + list(APPEND _libs "${_lib}") + endforeach() + else() + file(GLOB _found RELATIVE "${_lib_dir}" "${_lib_dir}/*.so") + if(NOT _found) + message(FATAL_ERROR "[QoLA] ${QAM_GROUP}: no shared objects in ${_lib_dir}.") + endif() + set(_libs ${_found}) + endif() + + if(QAM_OUT_INCLUDE_DIR) + set(${QAM_OUT_INCLUDE_DIR} "${_include_dir}" PARENT_SCOPE) + endif() + if(QAM_OUT_LIB_DIR) + set(${QAM_OUT_LIB_DIR} "${_lib_dir}" PARENT_SCOPE) + endif() + if(QAM_OUT_LIBS) + set(${QAM_OUT_LIBS} "${_libs}" PARENT_SCOPE) + endif() + if(QAM_OUT_CONFIG_DIR) + set(${QAM_OUT_CONFIG_DIR} "${_config_dir}" PARENT_SCOPE) + endif() +endfunction() diff --git a/qola/build_tools/builder.py b/qola/build_tools/builder.py index 7a117e3..21f070e 100644 --- a/qola/build_tools/builder.py +++ b/qola/build_tools/builder.py @@ -34,6 +34,7 @@ def build_kernels( aiter_commit: Optional[str] = None, patches_dir: Optional[str] = None, skip_checkout: bool = False, + groups: Optional[List[str]] = None, ) -> dict[str, Any]: """Build AITER kernel modules from a consumer manifest. @@ -77,6 +78,16 @@ def build_kernels( ``aiter_commit`` and ``patches_dir`` are ignored in this mode; the only requirement is that *aiter_root* points at an existing git checkout. Defaults to ``False``. + groups + Restrict the build to ``[[modules]]`` entries carrying one of these + ``group`` values. Lets one manifest -- a single AITER commit and + patch set -- serve several consumers that each build their own + subset. When ``None``, falls back to the ``QOLA_BUILD_GROUPS`` + environment variable (a ``;``-separated list); when that is unset + too, every module in the manifest is built. + + The environment fallback exists for callers that invoke ``qola + build`` through an intermediary which does not forward ``--group``. Returns ------- @@ -86,6 +97,13 @@ def build_kernels( output_dir = str(Path(output_dir).resolve()) manifest_path = str(Path(manifest_path).resolve()) + if groups is None: + env_groups = os.environ.get("QOLA_BUILD_GROUPS", "") + parsed = [g.strip() for g in env_groups.replace(",", ";").split(";") if g.strip()] + if parsed: + groups = parsed + print(f"[QoLA] Restricting build to group(s) {groups} (QOLA_BUILD_GROUPS)") + # Save env vars we'll override so we can restore them on exit. prev_gpu_archs = os.environ.get("GPU_ARCHS") prev_jit_dir = os.environ.get("AITER_JIT_DIR") @@ -127,7 +145,7 @@ def build_kernels( try: return _build_kernels_inner( aiter_root, output_dir, manifest_path, archs, - jit_build_dir, verbose, build_mode, + jit_build_dir, verbose, build_mode, groups, ) finally: _restore_env("GPU_ARCHS", prev_gpu_archs) @@ -150,12 +168,13 @@ def _build_kernels_inner( jit_build_dir: str, verbose: bool, build_mode: str, + groups: Optional[List[str]] = None, ) -> dict[str, Any]: # 1. Resolve namespace ns = build_namespace(aiter_root) # 2. Parse manifest - specs = load_manifest(manifest_path, ns, build_mode=build_mode) + specs = load_manifest(manifest_path, ns, build_mode=build_mode, groups=groups) # 3. Load build_module from AITER build_module = load_build_module_fn(aiter_root) @@ -291,7 +310,12 @@ def _invoke_build(build_module_fn, spec: BuildSpec, verbose: bool) -> None: os.environ.pop("HIP_CLANG_PATH", None) -_PUBLIC_HEADERS = ("qola_common.h", "qola_mha_fwd.h", "qola_mha_bwd.h") +_PUBLIC_HEADERS = ( + "qola_common.h", + "qola_mha_fwd.h", + "qola_mha_bwd.h", + "qola_gemm_a4w4.h", +) def _export_public_headers(output_dir: str, namespace: str) -> None: diff --git a/qola/build_tools/config.py b/qola/build_tools/config.py index 13f1c11..f492e7c 100644 --- a/qola/build_tools/config.py +++ b/qola/build_tools/config.py @@ -38,6 +38,7 @@ class BuildSpec: hipify: bool = False hip_clang_path: Optional[str] = None hsa_subdirs: List[str] = field(default_factory=list) + third_party: List[str] = field(default_factory=list) # Defaults matching core.py's d_opt_build_args (line 712, commit 33f2e6af) @@ -77,6 +78,7 @@ def load_manifest( manifest_path: str, ns: AiterNamespace, build_mode: Optional[str] = None, + groups: Optional[List[str]] = None, ) -> List[BuildSpec]: """Parse a TOML manifest and return resolved :class:`BuildSpec` instances. @@ -93,6 +95,12 @@ def load_manifest( Per-module ``mode`` entries in ``[[modules]]`` still take final precedence (most specific scope). When ``None`` and unset in the manifest, defaults to ``"pybind"``. + groups + When provided, restricts the build to ``[[modules]]`` entries whose + ``group`` is in this list. This lets a single manifest -- one AITER + commit, one patch set, one checkout -- serve several independent + consumers that each build only their own subset of kernels. When + ``None``, every module in the manifest is built. Manifest schema:: @@ -106,6 +114,7 @@ def load_manifest( [[modules]] name = "libmha_fwd" + group = "ck_fused_attn" # optional; selectable via --group mode = "cpp_itfs" # optional per-module override receipt = 700 # optional CK codegen filter (default: whatever # optCompilerConfig.json specifies, typically 600) @@ -140,15 +149,45 @@ def load_manifest( gpu_archs_env = os.getenv("GPU_ARCHS", "") resolved_archs = [a.strip() for a in gpu_archs_env.split(";") if a.strip()] + # Restrict to the requested module groups, if any. Ungrouped modules are + # excluded when filtering is active: a manifest shared by several + # consumers should say explicitly who owns each module. + all_modules = manifest.get("modules", []) + if groups is not None: + wanted = set(groups) + declared = {m["group"] for m in all_modules if "group" in m} + unknown = wanted - declared + if unknown: + raise ValueError( + f"Unknown module group(s) {sorted(unknown)} for manifest " + f"{manifest_path}. Declared groups: {sorted(declared) or '(none)'}." + ) + selected_modules = [m for m in all_modules if m.get("group") in wanted] + if not selected_modules: + raise ValueError( + f"No modules selected for group(s) {sorted(wanted)} in " + f"manifest {manifest_path}." + ) + else: + selected_modules = all_modules + specs: List[BuildSpec] = [] fwd_section = manifest.get("mha_fwd_variants", []) - module_names = {m["name"] for m in manifest.get("modules", [])} + module_names = {m["name"] for m in selected_modules} has_fwd_variants = bool(fwd_section) has_static_fwd = "libmha_fwd" in module_names # Keys consumed by load_manifest before passing to _resolve_static_module. - _MANIFEST_KEYS = {"name", "mode", "drop_srcs", "drop_directions", "hsa_subdirs", "receipt"} + _MANIFEST_KEYS = { + "name", + "group", + "mode", + "drop_srcs", + "drop_directions", + "hsa_subdirs", + "receipt", + } # --- static modules --- # NOTE: Variant filtering is NOT applied to static libmha_fwd / @@ -156,7 +195,7 @@ def load_manifest( # files and the dispatch API file (fmha_*_api.cpp) on every call. # Running it N times with different --filter patterns overwrites the # API dispatch, leaving only the last filter's branches. - for mod_entry in manifest.get("modules", []): + for mod_entry in selected_modules: name = mod_entry["name"] mod_mode = mod_entry.get("mode", global_mode) drop_srcs = set(mod_entry.get("drop_srcs", [])) diff --git a/qola/build_tools/resolver.py b/qola/build_tools/resolver.py index 6c45485..f8fb38f 100644 --- a/qola/build_tools/resolver.py +++ b/qola/build_tools/resolver.py @@ -81,37 +81,38 @@ def build_namespace(aiter_root: str) -> AiterNamespace: # AITER_CONFIGS extraction # ------------------------------------------------------------------ -_CONFIG_START = "# config_env start here" -_CONFIG_END = "# config_env end here" - - def _build_aiter_configs(core_path: str, aiter_root_dir: str) -> Any: - """Extract the AITER_CONFIG class from core.py and instantiate it. - - Executes only the marked ``# config_env start/end`` block (lines 68-279) - in a namespace where ``AITER_ROOT_DIR`` is bound. All properties of - ``AITER_CONFIG`` reference ``AITER_ROOT_DIR`` through the module-level - ``AITER_CONFIG_*`` string constants defined in the same block. + """Return the ``AITER_CONFIGS`` singleton from AITER's core.py. + + Loads ``core.py`` as an isolated module via importlib (same mechanism as + ``load_build_module_fn``) and reads its module-level ``AITER_CONFIGS``. + + Historically this exec'd only the marked ``# config_env start/end`` block + with a hand-built namespace, to avoid running ``aiter/__init__.py``. But + newer AITER lineages grew ``get_config_file`` to reference module-level + helpers defined outside that block (``logger``, ``re``, ``mp_lock``, + ``FileBaton`` ...), which made the block non-self-contained. Loading the + full module (without importing the ``aiter`` package) is both simpler and + robust to where inside core.py these helpers live. """ - with open(core_path, "r") as f: - source = f.read() + jit_dir = os.path.dirname(os.path.abspath(core_path)) + utils_dir = os.path.join(jit_dir, "utils") + for d in (utils_dir, jit_dir): + if d not in sys.path: + sys.path.insert(0, d) - start = source.find(_CONFIG_START) - end = source.find(_CONFIG_END) - if start == -1 or end == -1: + spec = importlib.util.spec_from_file_location("_qola_jit_core_cfg", core_path) + if spec is None or spec.loader is None: + raise RuntimeError(f"Cannot load core.py from {core_path}") + mod = importlib.util.module_from_spec(spec) + spec.loader.exec_module(mod) + try: + return mod.AITER_CONFIGS # type: ignore[attr-defined] + except AttributeError as e: raise RuntimeError( - f"Could not find config_env markers in {core_path}. " - "AITER core.py structure may have changed." - ) - block = source[start : end + len(_CONFIG_END)] - - exec_ns: dict[str, Any] = { - "os": os, - "functools": __import__("functools"), - "AITER_ROOT_DIR": aiter_root_dir, - } - exec(compile(block, core_path, "exec"), exec_ns) # noqa: S102 - return exec_ns["AITER_CONFIGS"] + f"AITER core.py at {core_path} does not define AITER_CONFIGS; " + "its structure may have changed." + ) from e # ------------------------------------------------------------------ diff --git a/qola/cli.py b/qola/cli.py index b1d3a56..a9a6d6f 100644 --- a/qola/cli.py +++ b/qola/cli.py @@ -76,6 +76,16 @@ def main(argv: list[str] | None = None) -> int: "AITER source tree. --aiter-commit and --patches-dir are ignored " "when this is set.", ) + build_p.add_argument( + "--group", + action="append", + dest="groups", + help="Build only the [[modules]] entries whose 'group' matches. " + "Repeatable, and accepts a ';'-separated list. Lets one manifest " + "(one AITER commit, one patch set, one checkout) serve several " + "consumers that each build their own subset of kernels. When " + "omitted, every module in the manifest is built.", + ) build_p.add_argument( "--verbose", "-v", @@ -129,6 +139,10 @@ def main(argv: list[str] | None = None) -> int: if args.archs: archs = [a for entry in args.archs for a in entry.split(";") if a] + groups: list[str] | None = None + if args.groups: + groups = [g for entry in args.groups for g in entry.split(";") if g] + result = build_kernels( manifest_path=args.manifest, aiter_root=args.aiter_root, @@ -139,6 +153,7 @@ def main(argv: list[str] | None = None) -> int: aiter_commit=args.aiter_commit, patches_dir=args.patches_dir, skip_checkout=args.skip_checkout, + groups=groups, ) s = result["summary"] print( diff --git a/qola/cpp_itfs/qola_common.h b/qola/cpp_itfs/qola_common.h index 34e9352..9fb2554 100644 --- a/qola/cpp_itfs/qola_common.h +++ b/qola/cpp_itfs/qola_common.h @@ -1,5 +1,7 @@ // SPDX-License-Identifier: MIT // Copyright (C) 2026, Advanced Micro Devices, Inc. All rights reserved. +#pragma once + #include #include "qola_config.h" @@ -16,3 +18,27 @@ #define QOLA_NS_END } #define QOLA_NS(sym) qola::sym #endif + +// C-linkage exports cannot live in a namespace, so the same collision +// avoidance is applied as a symbol prefix instead. +// QOLA_NAMESPACE=te -> QOLA_C(foo) == qola_te_foo +// (unset) -> QOLA_C(foo) == qola_foo +// Consumers get the correct spelling for free by including the generated +// qola_config.h, so they never hardcode the namespace. +#define QOLA_C_CAT_(a, b) a##b +#define QOLA_C_CAT(a, b) QOLA_C_CAT_(a, b) +#ifdef QOLA_NAMESPACE +#define QOLA_C(sym) QOLA_C_CAT(QOLA_C_CAT(qola_, QOLA_NAMESPACE), QOLA_C_CAT(_, sym)) +#else +#define QOLA_C(sym) QOLA_C_CAT(qola_, sym) +#endif + +#ifdef __cplusplus +#define QOLA_C_BEGIN extern "C" { +#define QOLA_C_END } +#else +#define QOLA_C_BEGIN +#define QOLA_C_END +#endif + +#define QOLA_EXPORT __attribute__((visibility("default"))) diff --git a/qola/cpp_itfs/qola_exports.lds b/qola/cpp_itfs/qola_exports.lds index 1dad154..a09ec9e 100644 --- a/qola/cpp_itfs/qola_exports.lds +++ b/qola/cpp_itfs/qola_exports.lds @@ -12,6 +12,9 @@ */ { global: + /* C-linkage exports use a qola_[_] symbol prefix instead of a + * C++ namespace (see QOLA_C in qola_common.h). */ + qola_*; extern "C++" { qola::*; "typeinfo for qola::*"; diff --git a/qola/cpp_itfs/qola_gemm_a4w4.h b/qola/cpp_itfs/qola_gemm_a4w4.h new file mode 100644 index 0000000..7ae8df8 --- /dev/null +++ b/qola/cpp_itfs/qola_gemm_a4w4.h @@ -0,0 +1,83 @@ +// SPDX-License-Identifier: MIT +// Copyright (C) 2026, Advanced Micro Devices, Inc. All rights reserved. +// +// QoLA cpp_itfs C API for AITER's a4w4 (FP4 x FP4) GEMM kernels. +// +// This header deliberately pulls in nothing from AITER. Consumers see only +// a QoLA-owned POD descriptor and C-linkage entry points, so they can link +// the kernel libraries without exposing AITER headers (or AITER's enum +// ordering) to their own translation units. +// +// Error contract: every entry point returns 0 on success and non-zero on +// failure. A human-readable message is written to the caller-supplied +// `err_buf` when one is provided. The message is an out-parameter rather +// than a thread-local accessor on purpose -- the a4w4 backends ship as +// separate shared objects, so a shared `last_error()` symbol would resolve +// to whichever library the dynamic linker bound first and silently report +// the wrong (or an empty) message. +#pragma once + +#include +#include + +#include "qola_common.h" + +// Element type of an operand. Values are QoLA's own and are translated to +// AITER's AiterDtype inside the implementation, so this ABI is unaffected by +// reordering of AITER's enum. +typedef enum { + QOLA_DTYPE_FP4X2 = 0, /* two packed FP4 (E2M1) values per byte */ + QOLA_DTYPE_E8M0 = 1, /* 8-bit exponent-only microscale (1 byte) */ + QOLA_DTYPE_BF16 = 2, + QOLA_DTYPE_FP16 = 3, + QOLA_DTYPE_FP32 = 4, + QOLA_DTYPE_U8 = 5, + QOLA_DTYPE_I8 = 6, +} qola_dtype_t; + +// Lightweight device-tensor descriptor. The caller owns the storage; the +// descriptor must outlive the call but the storage need not. +// +// Field order is chosen for natural alignment and the explicit padding keeps +// the layout identical across compilers, so consumers may safely define a +// structurally identical type instead of including this header. +typedef struct { + void *ptr; + int32_t ndim; + int32_t dtype; /* one of qola_dtype_t */ + int32_t device_id; + int32_t reserved; + int64_t shape[8]; + int64_t strides[8]; +} qola_tensor_t; + +QOLA_C_BEGIN + +/* CK blockscale a4w4 GEMM: Y = XQ @ WQ^T with per-1x32 microscaling. + * XQ [M, K/2] fp4x2 + * WQ [N, K/2] fp4x2 + * x_scale [M, K/32] e8m0 + * w_scale [N, K/32] e8m0 + * Y [M, N] bf16 / fp16 (output, pre-allocated) + * + * `kernel_name` may be NULL or empty to request the default heuristic; + * a non-empty name must exist in the compiled registry. Kernel selection + * and weight/scale pre-shuffling are the caller's responsibility. + */ +QOLA_EXPORT int QOLA_C(gemm_a4w4_blockscale)(const qola_tensor_t *XQ, const qola_tensor_t *WQ, + const qola_tensor_t *x_scale, + const qola_tensor_t *w_scale, const qola_tensor_t *Y, + int split_k, const char *kernel_name, + hipStream_t stream, char *err_buf, size_t err_buf_size); + +/* ASM (f4gemm) a4w4 GEMM: D = alpha*A*B + beta*C. + * `bias` may be NULL; `kernel_name` may be NULL or empty for the heuristic. + */ +QOLA_EXPORT int QOLA_C(gemm_a4w4_asm)(const qola_tensor_t *A, const qola_tensor_t *B, + const qola_tensor_t *a_scale, const qola_tensor_t *b_scale, + const qola_tensor_t *out, const qola_tensor_t *bias, + const char *kernel_name, float alpha, float beta, + int bpreshuffle, int log2_k_split, hipStream_t stream, + char *err_buf, size_t err_buf_size); + +QOLA_C_END diff --git a/qola/cpp_itfs/qola_gemm_a4w4_asm.cu b/qola/cpp_itfs/qola_gemm_a4w4_asm.cu new file mode 100644 index 0000000..0e96f11 --- /dev/null +++ b/qola/cpp_itfs/qola_gemm_a4w4_asm.cu @@ -0,0 +1,83 @@ +// SPDX-License-Identifier: MIT +// Copyright (C) 2026, Advanced Micro Devices, Inc. All rights reserved. +// +// Thin cpp_itfs entry point for AITER's ASM a4w4 GEMM (f4gemm). + +#include "qola_gemm_a4w4_internal.h" + +// The torch-free ASM entrypoint is a C-ABI symbol defined in AITER's +// csrc/py_itfs_cu/asm_gemm_a4w4.cu via AITER_CTYPES_DEFINE_ENTRYPOINT_VOID. +// It returns 0 on success or -1 on failure, with the message stored in a +// thread-local retrievable through aiter_get_last_error(). We declare both +// here rather than pull in a torch-tainted AITER header. +extern "C" int gemm_a4w4_asm(aiter_tensor_t* A, + aiter_tensor_t* B, + aiter_tensor_t* A_scale, + aiter_tensor_t* B_scale, + aiter_tensor_t* out, + const char* kernelName, + aiter_tensor_t* bias, + float alpha, + float beta, + int bpreshuffle, + int log2_k_split, + hipStream_t stream); +extern "C" const char* aiter_get_last_error(); + +extern "C" int QOLA_C(gemm_a4w4_asm)(const qola_tensor_t* A, + const qola_tensor_t* B, + const qola_tensor_t* a_scale, + const qola_tensor_t* b_scale, + const qola_tensor_t* out, + const qola_tensor_t* bias, + const char* kernel_name, + float alpha, + float beta, + int bpreshuffle, + int log2_k_split, + hipStream_t stream, + char* err_buf, + size_t err_buf_size) +{ + if(A == nullptr || B == nullptr || a_scale == nullptr || b_scale == nullptr || out == nullptr) + { + qola_detail::set_error(err_buf, err_buf_size, "null tensor descriptor"); + return 1; + } + // AITER's entrypoint already reports failures as a status code, so the + // only thing that can throw here is descriptor translation. + int rc = 0; + int guard_rc = qola_detail::guarded(err_buf, err_buf_size, [&] { + aiter_tensor_t a_a = qola_detail::to_aiter_tensor(*A); + aiter_tensor_t a_b = qola_detail::to_aiter_tensor(*B); + aiter_tensor_t a_as = qola_detail::to_aiter_tensor(*a_scale); + aiter_tensor_t a_bs = qola_detail::to_aiter_tensor(*b_scale); + aiter_tensor_t a_o = qola_detail::to_aiter_tensor(*out); + aiter_tensor_t a_bias; + aiter_tensor_t* a_bias_ptr = nullptr; + if(bias != nullptr && bias->ptr != nullptr) + { + a_bias = qola_detail::to_aiter_tensor(*bias); + a_bias_ptr = &a_bias; + } + rc = ::gemm_a4w4_asm(&a_a, + &a_b, + &a_as, + &a_bs, + &a_o, + kernel_name ? kernel_name : "", + a_bias_ptr, + alpha, + beta, + bpreshuffle, + log2_k_split, + stream); + if(rc != 0) + { + const char* msg = ::aiter_get_last_error(); + qola_detail::set_error( + err_buf, err_buf_size, msg ? msg : "aiter gemm_a4w4_asm failed"); + } + }); + return guard_rc != 0 ? guard_rc : rc; +} diff --git a/qola/cpp_itfs/qola_gemm_a4w4_blockscale.cu b/qola/cpp_itfs/qola_gemm_a4w4_blockscale.cu new file mode 100644 index 0000000..d3f0a0a --- /dev/null +++ b/qola/cpp_itfs/qola_gemm_a4w4_blockscale.cu @@ -0,0 +1,41 @@ +// SPDX-License-Identifier: MIT +// Copyright (C) 2026, Advanced Micro Devices, Inc. All rights reserved. +// +// Thin cpp_itfs entry point for AITER's CK a4w4 blockscale GEMM. + +#include "qola_gemm_a4w4_internal.h" + +#include "gemm_a4w4_blockscale.h" // aiter::gemm_a4w4_blockscale + +extern "C" int QOLA_C(gemm_a4w4_blockscale)(const qola_tensor_t* XQ, + const qola_tensor_t* WQ, + const qola_tensor_t* x_scale, + const qola_tensor_t* w_scale, + const qola_tensor_t* Y, + int split_k, + const char* kernel_name, + hipStream_t stream, + char* err_buf, + size_t err_buf_size) +{ + if(XQ == nullptr || WQ == nullptr || x_scale == nullptr || w_scale == nullptr || Y == nullptr) + { + qola_detail::set_error(err_buf, err_buf_size, "null tensor descriptor"); + return 1; + } + return qola_detail::guarded(err_buf, err_buf_size, [&] { + aiter_tensor_t a_xq = qola_detail::to_aiter_tensor(*XQ); + aiter_tensor_t a_wq = qola_detail::to_aiter_tensor(*WQ); + aiter_tensor_t a_xs = qola_detail::to_aiter_tensor(*x_scale); + aiter_tensor_t a_ws = qola_detail::to_aiter_tensor(*w_scale); + aiter_tensor_t a_y = qola_detail::to_aiter_tensor(*Y); + ::aiter::gemm_a4w4_blockscale(a_xq, + a_wq, + a_xs, + a_ws, + a_y, + split_k, + stream, + kernel_name ? std::string(kernel_name) : std::string()); + }); +} diff --git a/qola/cpp_itfs/qola_gemm_a4w4_internal.h b/qola/cpp_itfs/qola_gemm_a4w4_internal.h new file mode 100644 index 0000000..2d475e3 --- /dev/null +++ b/qola/cpp_itfs/qola_gemm_a4w4_internal.h @@ -0,0 +1,113 @@ +// SPDX-License-Identifier: MIT +// Copyright (C) 2026, Advanced Micro Devices, Inc. All rights reserved. +// +// Internal helpers shared by the a4w4 cpp_itfs entry points. Not exported to +// consumers -- it includes AITER headers, which is exactly what the public +// qola_gemm_a4w4.h exists to keep out of downstream translation units. +#pragma once + +#include +#include +#include + +#include "aiter_tensor.h" // aiter_tensor_t, AiterDtype +#include "qola_gemm_a4w4.h" + +namespace qola_detail { + +inline AiterDtype to_aiter_dtype(int dtype) +{ + switch(dtype) + { + case QOLA_DTYPE_FP4X2: return AITER_DTYPE_fp4x2; + case QOLA_DTYPE_E8M0: return AITER_DTYPE_fp8_e8m0; + case QOLA_DTYPE_BF16: return AITER_DTYPE_bf16; + case QOLA_DTYPE_FP16: return AITER_DTYPE_fp16; + case QOLA_DTYPE_FP32: return AITER_DTYPE_fp32; + case QOLA_DTYPE_U8: return AITER_DTYPE_u8; + case QOLA_DTYPE_I8: return AITER_DTYPE_i8; + default: return AITER_DTYPE_u8; + } +} + +// Build an aiter_tensor_t from the public descriptor. Shares the caller's +// device pointer; no ownership is transferred. +inline aiter_tensor_t to_aiter_tensor(const qola_tensor_t& d) +{ + aiter_tensor_t t{}; + t.ptr = d.ptr; + t.ndim = d.ndim; + size_t numel = (d.ndim > 0) ? 1 : 0; + for(int i = 0; i < d.ndim && i < 8; ++i) + { + t.shape[i] = d.shape[i]; + t.strides[i] = d.strides[i]; + numel *= static_cast(d.shape[i]); + } + t.numel_ = numel; + t.dtype_ = to_aiter_dtype(d.dtype); + t.device_id = d.device_id; + return t; +} + +// AITER's AITER_CHECK routes through aiter_detail::check_fail, which calls +// std::abort() unless the thread-local g_aiter_can_throw is set. AITER's own +// C entry points flip it for the duration of a call (see +// csrc/include/aiter_ctypes_error.h) which is why the ASM path already reports +// failures as a status code; the CK blockscale path is a plain C++ function +// and has no such wrapper. QoLA's C ABI promises status codes rather than +// process death, so it must establish the same guarantee itself. +// +// Save/restore rather than unconditionally clearing, so nesting inside an +// AITER entry point that already set the flag is harmless. +class CanThrowGuard +{ + public: + CanThrowGuard() + : prev_(aiter_detail::g_aiter_can_throw) + { + aiter_detail::g_aiter_can_throw = true; + } + ~CanThrowGuard() { aiter_detail::g_aiter_can_throw = prev_; } + + CanThrowGuard(const CanThrowGuard&) = delete; + CanThrowGuard& operator=(const CanThrowGuard&) = delete; + + private: + bool prev_; +}; + +inline void set_error(char* err_buf, size_t err_buf_size, const char* msg) +{ + if(err_buf == nullptr || err_buf_size == 0) + return; + if(msg == nullptr) + msg = "unknown error"; + std::strncpy(err_buf, msg, err_buf_size - 1); + err_buf[err_buf_size - 1] = '\0'; +} + +// Runs `fn`, translating any escaping exception into the C status/message +// contract documented in qola_gemm_a4w4.h. +template +inline int guarded(char* err_buf, size_t err_buf_size, Fn&& fn) +{ + CanThrowGuard can_throw; + try + { + fn(); + } + catch(const std::exception& e) + { + set_error(err_buf, err_buf_size, e.what()); + return 1; + } + catch(...) + { + set_error(err_buf, err_buf_size, "unknown non-std exception"); + return 1; + } + return 0; +} + +} // namespace qola_detail diff --git a/qola/cpp_itfs/registry.toml b/qola/cpp_itfs/registry.toml index f055bfa..0fa4efe 100644 --- a/qola/cpp_itfs/registry.toml +++ b/qola/cpp_itfs/registry.toml @@ -62,3 +62,26 @@ add_blob_gen_cmd = [ # the CK `-d bwd` codegen via the manifest and sets ENABLE_CK=0 in the env, so # no separate registry entry is needed. +# The two a4w4 entries below are keyed on AITER's *pybind* module names rather +# than on dedicated `libgemm_a4w4_*` entries. cpp_itfs mode already forces +# torch_exclude=True / is_python_module=False, which is the only thing those +# dedicated entries added, so keying on `module_*` lets AITER carry a single +# recipe per kernel. Consumers pin the artifact name back via the manifest's +# `md_name` override. + +[module_gemm_a4w4_blockscale] +# Strip the pybind-only sources the `module_*` entry carries; the kernel TU +# (gemm_a4w4_blockscale.cu) is already torch-free. +drop_srcs = ["gemm_a4w4_blockscale_pybind.cu", "gemm_common.cu"] +add_srcs = ["qola/cpp_itfs/qola_gemm_a4w4_blockscale.cu"] +add_includes = ["qola/cpp_itfs"] +# CK blockscale path: kernels are generated (gen_instances.py), no HSA blobs. + +[module_gemm_a4w4_asm] +# Strip the pybind TU; the kernel TU (asm_gemm_a4w4.cu) is already torch-free, +# taking aiter_tensor_t* rather than torch::Tensor. +drop_srcs = ["gemm_a4w4_asm_pybind.cu"] +add_srcs = ["qola/cpp_itfs/qola_gemm_a4w4_asm.cu"] +add_includes = ["qola/cpp_itfs"] +hsa_subdirs = ["f4gemm"] +