From 084307dc7988ba45ead79fca5f0f96c2d5748873 Mon Sep 17 00:00:00 2001 From: Yiannis Papadopoulos Date: Thu, 8 Oct 2026 16:02:17 -0600 Subject: [PATCH 1/7] Dispatch HSA kernels by kernel object, loaded from an hsaco Replace the raw PDI+instruction-address dispatch path with ROCr's code-object-reader/executable/freeze flow: HsaRuntime now loads an hsaco file into a frozen HSA executable and dispatches by the kernel object HSA_EXECUTABLE_SYMBOL_INFO_KERNEL_OBJECT resolves it to, instead of filling in raw PDI/instruction-stream device addresses. triton_npu_hsa_prepare(hsaco_path, kernel_name) replaces the old (pdi_path, insts_path) signature; triton_npu_hsa_patch_insts is removed along with the dev-pool PDI/insts loading machinery it depended on, since there is no longer a per-dispatch instruction stream to patch. Co-Authored-By: Claude Sonnet 5 --- .../backend/include/HsaRuntime/HsaRuntime.cpp | 200 ++++++++++-------- .../backend/include/HsaRuntime/HsaRuntime.h | 35 ++- 2 files changed, 126 insertions(+), 109 deletions(-) diff --git a/amd_triton_npu/backend/include/HsaRuntime/HsaRuntime.cpp b/amd_triton_npu/backend/include/HsaRuntime/HsaRuntime.cpp index 60956e6e..11543ba2 100644 --- a/amd_triton_npu/backend/include/HsaRuntime/HsaRuntime.cpp +++ b/amd_triton_npu/backend/include/HsaRuntime/HsaRuntime.cpp @@ -10,8 +10,9 @@ // mutex. // // Memory strategy: -// * PDI + instructions: plain HSA pool allocation from the dev pool -// (coarse-grained, non-allocatable), loaded once per (pdi, insts) and cached. +// * Kernels: packed into an hsaco (PDI+insts or full ELF), loaded through +// ROCR's code-object-reader/executable flow, and dispatched by the +// resulting opaque kernel object -- one load per (hsaco_path, kernel_name). // * Tensor I/O: the vmem API (handle_create -> reserve -> map -> set_access), // RW-accessible to CPU and AIE agents, pooled and reused across dispatches. // * Shared regions: the same vmem API, but owned by the caller rather than the @@ -318,13 +319,91 @@ struct SharedRegion { } // namespace -// A prepared program is just its two pool allocations; the opaque handle in the -// C ABI points at one of these, owned by the runtime's program cache. Both -// members free themselves, so a prepare() that loads the PDI and then fails on -// the instructions releases the PDI on the way out. +// An hsaco loaded as a frozen HSA executable, plus the kernel object its one +// dispatchable kernel resolves to. Self-freeing and move-only, same +// reasoning as PoolBuffer: a load() that creates the reader and then fails +// resolving the kernel symbol releases the reader/executable on the way +// out, and a throw partway through still leaves nothing dangling. +// +// hsaco_ is held, not just read once: ROCR reads the code object in place +// for as long as the reader exists, so these bytes must outlive it. +class LoadedExecutable { +public: + LoadedExecutable() = default; + LoadedExecutable(const LoadedExecutable &) = delete; + LoadedExecutable &operator=(const LoadedExecutable &) = delete; + LoadedExecutable(LoadedExecutable &&o) noexcept + : hsaco_(std::move(o.hsaco_)), reader_(std::exchange(o.reader_, {})), + executable_(std::exchange(o.executable_, {})), + kernel_object_(std::exchange(o.kernel_object_, 0)) {} + LoadedExecutable &operator=(LoadedExecutable &&o) noexcept { + if (this != &o) { + reset(); + hsaco_ = std::move(o.hsaco_); + reader_ = std::exchange(o.reader_, {}); + executable_ = std::exchange(o.executable_, {}); + kernel_object_ = std::exchange(o.kernel_object_, 0); + } + return *this; + } + ~LoadedExecutable() { reset(); } + + std::uint64_t kernel_object() const { return kernel_object_; } + + // Load hsaco_bytes on `agent` and resolve kernel_name to its kernel + // object. Throws std::runtime_error (via HSA_CHECK) naming the failing + // call; already-acquired state is torn down before rethrowing. + static LoadedExecutable load(hsa_agent_t agent, + std::vector hsaco_bytes, + const char *kernel_name) { + LoadedExecutable exe; + exe.hsaco_ = std::move(hsaco_bytes); + try { + HSA_CHECK(hsa_code_object_reader_create_from_memory( + exe.hsaco_.data(), exe.hsaco_.size(), &exe.reader_)); + HSA_CHECK(hsa_executable_create_alt( + HSA_PROFILE_FULL, HSA_DEFAULT_FLOAT_ROUNDING_MODE_DEFAULT, nullptr, + &exe.executable_)); + HSA_CHECK(hsa_executable_load_agent_code_object( + exe.executable_, agent, exe.reader_, nullptr, nullptr)); + HSA_CHECK(hsa_executable_freeze(exe.executable_, nullptr)); + hsa_executable_symbol_t symbol{}; + HSA_CHECK(hsa_executable_get_symbol_by_name( + exe.executable_, kernel_name, &agent, &symbol)); + HSA_CHECK(hsa_executable_symbol_get_info( + symbol, HSA_EXECUTABLE_SYMBOL_INFO_KERNEL_OBJECT, + &exe.kernel_object_)); + } catch (...) { + exe.reset(); + throw; + } + return exe; + } + +private: + void reset() { + if (executable_.handle) + log_status("hsa_executable_destroy", + hsa_executable_destroy(executable_)); + if (reader_.handle) + log_status("hsa_code_object_reader_destroy", + hsa_code_object_reader_destroy(reader_)); + executable_ = {}; + reader_ = {}; + kernel_object_ = 0; + hsaco_.clear(); + } + + std::vector hsaco_; + hsa_code_object_reader_t reader_{}; + hsa_executable_t executable_{}; + std::uint64_t kernel_object_{}; +}; + +// The opaque handle in the C ABI points at one of these, owned by the +// runtime's program cache. struct triton_npu_hsa_program { - PoolBuffer pdi; - PoolBuffer insts; + LoadedExecutable exe; }; namespace { @@ -375,55 +454,24 @@ class HsaRuntime { return std::string(buf); } - // Load + cache the PDI/insts for a kernel and return its program handle. - // Keyed by (pdi_path, insts_path) so repeated prepares (or a PDI shared by - // two signatures) reuse the same device allocation. Thread-safe. - triton_npu_hsa_program *prepare(const char *pdi_path, - const char *insts_path) { - std::string key = std::string(pdi_path) + '\0' + insts_path; + // Load + cache the hsaco at hsaco_path and resolve kernel_name to a kernel + // object; return its program handle. Keyed by (hsaco_path, kernel_name) so + // repeated prepares reuse the same loaded executable. Thread-safe. + triton_npu_hsa_program *prepare(const char *hsaco_path, + const char *kernel_name) { + std::string key = std::string(hsaco_path) + '\0' + kernel_name; std::lock_guard lock(programs_mtx_); auto it = programs_.find(key); if (it != programs_.end()) return it->second.get(); auto prog = std::make_unique(); - // If the second load throws, ~PoolBuffer releases the first. - prog->pdi = load_binary(pdi_path); - prog->insts = load_binary(insts_path); + prog->exe = LoadedExecutable::load(aie_agent_, read_host_file(hsaco_path), + kernel_name); triton_npu_hsa_program *raw = prog.get(); programs_.emplace(std::move(key), std::move(prog)); return raw; } - // Overwrite `nbytes` of a prepared program's instruction stream at - // `byte_offset`. - // - // The AIE dispatch packet carries insts_addr and insts_size per enqueue, so - // the stream is an input to each dispatch rather than a property of the - // program -- prepare() only fixes where it lives. A decode needs that: its - // context length L is encoded in a handful of stream words, and mlir-air's - // xclbin path rewrites exactly those words per token. This is the same edit - // against the buffer the packet already points at. - // - // Takes dispatch_mtx_, the same lock a dispatch holds. The GIL does not - // serialise this: the dispatch path releases it for the duration of the - // device work, and this is reached through ctypes, which releases it too. - // Without the lock another thread could memcpy into the stream while the - // queue is consuming it, and the dispatch would run a mix of two contexts. - void patch_insts(triton_npu_hsa_program *program, std::uint64_t byte_offset, - const void *src, std::uint64_t nbytes) { - if (program == nullptr) - throw std::runtime_error("patch_insts called with a null program handle"); - if (byte_offset + nbytes < byte_offset || - byte_offset + nbytes > program->insts.size()) - throw std::runtime_error( - "patch_insts range [" + std::to_string(byte_offset) + ", +" + - std::to_string(nbytes) + ") is outside the " + - std::to_string(program->insts.size()) + "-byte instruction stream"); - std::lock_guard lock(dispatch_mtx_); - std::memcpy(static_cast(program->insts.va()) + byte_offset, src, - static_cast(nbytes)); - } - // Run one dispatch of `program` over num_tensors (host_ptr, size) pairs. // Serialized against every other dispatch (single shared queue). void dispatch(triton_npu_hsa_program *program, std::uint32_t num_tensors, @@ -521,16 +569,14 @@ class HsaRuntime { pkt.opcode = HSA_AMD_AIE_PACKET_OPCODE_KMQ; pkt.count = AIE_PACKET_COUNT; pkt.completion_signal = signal_; - const auto insts_addr = - reinterpret_cast(program->insts.va()); - pkt.insts_addr_low = static_cast(insts_addr & 0xFFFFFFFFu); - pkt.insts_addr_high = static_cast(insts_addr >> 32); + const std::uint64_t kernel_object = program->exe.kernel_object(); + pkt.kernel_object_low = + static_cast(kernel_object & 0xFFFFFFFFu); + pkt.kernel_object_high = static_cast(kernel_object >> 32); // Narrowing checked by the static_assert on TRITON_NPU_HSA_MAX_KERNARGS. pkt.num_kernargs = static_cast(num_tensors); // The ABI requires kernarg_address to be NULL when num_kernargs is 0. pkt.kernarg_address = (num_tensors > 0) ? kernargs : nullptr; - pkt.insts_size = program->insts.size(); - pkt.pdi_addr = program->pdi.va(); // Arm the signal to 1 (device decrements to 0 on success, or sets a // negative error code) and write the packet into the slot. Nothing from @@ -754,7 +800,6 @@ class HsaRuntime { private: hsa_agent_t aie_agent_{}; - hsa_amd_memory_pool_t dev_pool_{}; hsa_amd_memory_pool_t data_pool_{}; hsa_queue_t *queue_ = nullptr; hsa_signal_t signal_{}; @@ -780,9 +825,10 @@ class HsaRuntime { // cannot be reliably re-reserved+mapped, and reuse removes per-launch cost). std::unordered_map> vmem_pool_; - // Prepared programs, keyed by "pdi\0insts"; owns the loaded PDI/insts and - // makes prepare() idempotent (a repeated (pdi, insts) reuses one dev - // allocation rather than loading/leaking a second copy). + // Prepared programs, keyed by "hsaco_path\0kernel_name"; owns the loaded + // executable and makes prepare() idempotent (a repeated (hsaco_path, + // kernel_name) reuses the same loaded executable rather than loading/ + // leaking a second copy). // // Unordered: every use is an exact-key lookup, and the one iteration is the // destructor freeing them all. Unlike regions_ below, which cannot be. @@ -971,10 +1017,6 @@ class HsaRuntime { aie_agents_ = aies; cpu_agents_ = cpus; - // dev pool: coarse-grained, non-allocatable (PDI + instructions). - if (!discover_pool(HSA_AMD_MEMORY_POOL_GLOBAL_FLAG_COARSE_GRAINED, false, - &dev_pool_)) - throw std::runtime_error("no dev memory pool on AIE agent"); // data pool: coarse-grained, allocatable (tensor data via vmem). if (!discover_pool(HSA_AMD_MEMORY_POOL_GLOBAL_FLAG_COARSE_GRAINED, true, &data_pool_)) @@ -1199,8 +1241,10 @@ class HsaRuntime { return ((n + data_granule_ - 1) / data_granule_) * data_granule_; } - // Read a file into a fresh dev-pool allocation. - PoolBuffer load_binary(const std::string &path) { + // Read a file into plain host memory. Distinct from the AIE dev-pool + // allocation load_binary() used to return: hsa_code_object_reader_create_ + // from_memory reads from ordinary host memory, not a device allocation. + std::vector read_host_file(const std::string &path) { std::ifstream f(path, std::ios::binary | std::ios::ate); if (!f) throw std::runtime_error("failed to open '" + path + "'"); @@ -1208,14 +1252,10 @@ class HsaRuntime { if (sz <= 0) throw std::runtime_error("empty or unreadable '" + path + "'"); f.seekg(0); - void *raw = nullptr; - HSA_CHECK(hsa_amd_memory_pool_allocate( - dev_pool_, static_cast(sz), 0, &raw)); - // Owned from here on, so the short-read path below frees it by unwinding. - PoolBuffer b(raw, static_cast(sz)); - if (!f.read(static_cast(raw), sz)) + std::vector buf(static_cast(sz)); + if (!f.read(buf.data(), sz)) throw std::runtime_error("short read loading '" + path + "'"); - return b; + return buf; } // Reserve `b.size` bytes of address space for the handle `b` already holds, @@ -1438,10 +1478,10 @@ extern "C" int triton_npu_hsa_agent_name(char *buf, size_t buf_len, } extern "C" triton_npu_hsa_program_t -triton_npu_hsa_prepare(const char *pdi_path, const char *insts_path, +triton_npu_hsa_prepare(const char *hsaco_path, const char *kernel_name, char *errbuf, size_t errbuf_len) { try { - return runtime().prepare(pdi_path, insts_path); + return runtime().prepare(hsaco_path, kernel_name); } catch (const std::exception &e) { write_err(errbuf, errbuf_len, "HSA prepare failed", e.what()); return nullptr; @@ -1509,22 +1549,6 @@ extern "C" int triton_npu_hsa_shared_mark_dirty(void *va, char *errbuf, } } -extern "C" int triton_npu_hsa_patch_insts(triton_npu_hsa_program_t program, - uint64_t byte_offset, const void *src, - uint64_t nbytes, char *errbuf, - size_t errbuf_len) { - try { - runtime().patch_insts(program, byte_offset, src, nbytes); - return 0; - } catch (const std::exception &e) { - write_err(errbuf, errbuf_len, "HSA patch insts failed", e.what()); - return -1; - } catch (...) { - write_err(errbuf, errbuf_len, "HSA patch insts failed", "unknown error"); - return -1; - } -} - extern "C" void triton_npu_hsa_dispatch_counts(uint64_t *in_place, uint64_t *staged) { // No runtime() here, and so nothing that can throw: see g_in_place. diff --git a/amd_triton_npu/backend/include/HsaRuntime/HsaRuntime.h b/amd_triton_npu/backend/include/HsaRuntime/HsaRuntime.h index 87d440c7..2e4b8e7c 100644 --- a/amd_triton_npu/backend/include/HsaRuntime/HsaRuntime.h +++ b/amd_triton_npu/backend/include/HsaRuntime/HsaRuntime.h @@ -28,7 +28,8 @@ extern "C" { // ring slot (a few KiB in total). #define TRITON_NPU_HSA_MAX_KERNARGS 16 -// Opaque handle to a prepared (pdi, insts) program, owned by the runtime. +// Opaque handle to a prepared kernel, loaded from an hsaco and dispatched by +// its kernel object, owned by the runtime. typedef struct triton_npu_hsa_program *triton_npu_hsa_program_t; // Initialize the runtime (idempotent) and write the AIE agent's device name @@ -39,11 +40,18 @@ typedef struct triton_npu_hsa_program *triton_npu_hsa_program_t; int triton_npu_hsa_agent_name(char *buf, size_t buf_len, char *errbuf, size_t errbuf_len); -// Initialize the runtime (idempotent) and load + cache the PDI and instruction -// binaries for one kernel. Returns an opaque program handle, or NULL on error -// (with a message written to errbuf). Call once per launcher module. -triton_npu_hsa_program_t triton_npu_hsa_prepare(const char *pdi_path, - const char *insts_path, +// Initialize the runtime (idempotent), pack-load an hsaco, and resolve +// kernel_name to a kernel object. Returns an opaque program handle, or NULL +// on error (with a message written to errbuf). Call once per launcher +// module. +// +// hsaco_path must name a code object ROCR's AIE loader recognises -- an +// architecture-named (aie2/aie2p) section carrying either a PDI plus +// instruction sequence, or a full ELF (aie2p only). The hsaco's bytes are +// read into this call and kept alive for as long as the returned program is +// prepared, since ROCR reads the code object in place. +triton_npu_hsa_program_t triton_npu_hsa_prepare(const char *hsaco_path, + const char *kernel_name, char *errbuf, size_t errbuf_len); @@ -99,21 +107,6 @@ int triton_npu_hsa_dispatch_ex(triton_npu_hsa_program_t program, const uint64_t *sizes, const uint8_t *writeback, char *errbuf, size_t errbuf_len); -// Overwrite nbytes of a prepared program's instruction stream at byte_offset. -// Returns 0 on success, or a negative value (with a message in errbuf). -// -// The dispatch packet carries the stream's address and size per enqueue, so the -// stream is an input to each dispatch and not a property of the program. A -// decode needs that: its context length is encoded in a few stream words and -// changes every token. -// -// Serialised against dispatch internally. The GIL does not do it: the dispatch -// path releases the GIL, and so does the ctypes call that reaches this. -int triton_npu_hsa_patch_insts(triton_npu_hsa_program_t program, - uint64_t byte_offset, const void *src, - uint64_t nbytes, char *errbuf, - size_t errbuf_len); - // --------------------------------------------------------------------------- // Shared regions // --------------------------------------------------------------------------- From 289aefcfefb0b7699e6c44bbfccaea41679ac8f4 Mon Sep 17 00:00:00 2001 From: Yiannis Papadopoulos Date: Thu, 8 Oct 2026 16:07:04 -0600 Subject: [PATCH 2/7] Generate HSA launchers that pass an hsaco path and kernel name Update hsa_launcher.py to generate C++ code that calls triton_npu_hsa_prepare with hsaco_path and kernel_name parameters, replacing the previous pdi and insts parameters. This matches the ABI change made in Task 1 to the shared HSA runtime. Changes: - Update py_set_paths function to accept hsaco_path and kernel_name - Update comment describing g_program handle - Update PyMethodDef doc string for set_paths - Update module docstring to reflect new API - Update _generate_hsa_launcher docstring to explain kernel_name handling Verified with Python string template substitution check. Co-Authored-By: Claude Haiku 4.5 --- amd_triton_npu/backend/hsa_launcher.py | 29 ++++++++++++++------------ 1 file changed, 16 insertions(+), 13 deletions(-) diff --git a/amd_triton_npu/backend/hsa_launcher.py b/amd_triton_npu/backend/hsa_launcher.py index 4dcbd486..f3b2e416 100644 --- a/amd_triton_npu/backend/hsa_launcher.py +++ b/amd_triton_npu/backend/hsa_launcher.py @@ -9,8 +9,9 @@ the shared runtime library ``libtriton_npu_hsa.so`` (see ``include/HsaRuntime/HsaRuntime.{h,cpp}``) via a small C ABI: -* ``triton_npu_hsa_prepare(pdi, insts)`` -> opaque program handle (once, from - ``set_paths``). +* ``triton_npu_hsa_prepare(hsaco, kernel_name)`` -> opaque program handle + (once, from ``set_paths``), loaded from an hsaco and resolved by kernel + name. * ``triton_npu_hsa_dispatch(program, n, ptrs, sizes)`` (per launch). Because every launcher links the same shared library, the ``HsaRuntime`` @@ -34,7 +35,7 @@ def _generate_hsa_launcher(constants, signature, _kernel_name, written=None) -> str: """Generate the thin C++ CPython launcher that dispatches via HSA/ROCR. - The generated module exposes ``set_paths(pdi_path, insts_path)`` (which calls + The generated module exposes ``set_paths(hsaco_path, kernel_name)`` (which calls ``triton_npu_hsa_prepare`` once and stashes the handle) and ``launch(...)`` (which marshals the tensor pointers/sizes and calls ``triton_npu_hsa_dispatch`` with the GIL released). All HSA state lives in @@ -44,8 +45,9 @@ def _generate_hsa_launcher(constants, signature, _kernel_name, written=None) -> (``codegen.written_pointer_args``), or None when that is not known. ``_kernel_name`` is accepted for signature parity with the XRT launcher - generators but is unused: the HSA path selects work by PDI/insts address, - not by a kernel symbol name. + generators but is unused here: the real kernel name is threaded through at + runtime via ``set_paths``'s second argument (see driver.py's + ``compile_module``), not baked into the generated source. """ args_format = "".join(format_of(extracted_type(ty)) for ty in signature.values()) fmt = "iiiOOOO" + args_format @@ -114,19 +116,20 @@ def _generate_hsa_launcher(constants, signature, _kernel_name, written=None) -> // Number of tensor kernel arguments for this specialized launcher. static constexpr std::uint32_t NUM_KERNARGS = {num_ptr_args}; -// Handle to this module's prepared (pdi, insts) program in the shared runtime. +// Handle to this module's prepared hsaco kernel in the shared runtime. static triton_npu_hsa_program_t g_program = nullptr; -// Python-callable set_paths(pdi_path, insts_path): prepare (load + cache) the -// program in the shared runtime and stash its handle. Called once per module. +// Python-callable set_paths(hsaco_path, kernel_name): pack-load the hsaco +// in the shared runtime and stash its program handle. Called once per +// module. static PyObject* py_set_paths(PyObject* self, PyObject* args) {{ - const char* pdi; - const char* insts; - if (!PyArg_ParseTuple(args, "ss", &pdi, &insts)) {{ + const char* hsaco_path; + const char* kernel_name; + if (!PyArg_ParseTuple(args, "ss", &hsaco_path, &kernel_name)) {{ return NULL; }} char err[512]; - g_program = triton_npu_hsa_prepare(pdi, insts, err, sizeof(err)); + g_program = triton_npu_hsa_prepare(hsaco_path, kernel_name, err, sizeof(err)); if (g_program == nullptr) {{ PyErr_SetString(PyExc_RuntimeError, err); return NULL; @@ -198,7 +201,7 @@ def _generate_hsa_launcher(constants, signature, _kernel_name, written=None) -> // Methods exported by this per-signature dispatch extension module. static PyMethodDef ModuleMethods[] = {{ {{"launch", launch, METH_VARARGS, "Entry point for all kernels with this signature"}}, - {{"set_paths", py_set_paths, METH_VARARGS, "Set paths to aie.pdi and insts.bin"}}, + {{"set_paths", py_set_paths, METH_VARARGS, "Set path to kernel.hsaco and the kernel name"}}, {{NULL, NULL, 0, NULL}} // sentinel }}; From 5ef1c00da1bb6c1fc363545466915fe5fa940500 Mon Sep 17 00:00:00 2001 From: Yiannis Papadopoulos Date: Thu, 8 Oct 2026 20:14:38 -0600 Subject: [PATCH 3/7] Pack HSA kernels into an hsaco and dispatch them by kernel object driver.py now lets the HSA runtime auto-detect its output format the same way XRT does (ELF on npu2, xclbin on npu1) instead of forcing a raw "pdi" format, packs aiecc's artifacts into an hsaco via aie.compiler.hsaco.pack, caches it alongside the other compiled artifacts, and passes it to the generated launcher's set_paths(hsaco_path, kernel_name). Also fixes a latent bug in the cache-hit path: the "hsaco missing, force a rebuild" guard only cleared cached_artifacts, not the separately-gated cache_path the rebuild branch actually checks. Verified the new _pack_hsaco() helper directly against real previously compiled PDI+insts and full-ELF artifacts (both produce a structurally valid, correctly-tagged hsaco section per aie.compiler.hsaco.dump). Could not get a full live Triton-JIT compile through this path to green in this session's ad hoc venv -- aircc itself fails the same generic MLIR legalization error for every kernel tried, reproduced by invoking it directly with no Triton/driver.py code involved, so this is an environment issue for Task 4 (real hardware) to resolve, not a defect in this change. Co-Authored-By: Claude Sonnet 5 --- amd_triton_npu/backend/driver.py | 109 +++++++++++++++++++++++++------ scripts/test_hsa_hsaco_pack.py | 68 +++++++++++++++++++ 2 files changed, 157 insertions(+), 20 deletions(-) create mode 100644 scripts/test_hsa_hsaco_pack.py diff --git a/amd_triton_npu/backend/driver.py b/amd_triton_npu/backend/driver.py index 134a87d6..e7c08ba6 100644 --- a/amd_triton_npu/backend/driver.py +++ b/amd_triton_npu/backend/driver.py @@ -1287,24 +1287,11 @@ def _get_output_format(runtime=None): the ``AMD_TRITON_NPU_OUTPUT_FORMAT`` env var). If not set, defaults to "elf" on npu2 and "xclbin" on npu1. ELF format is only supported on npu2 (AIE2P) devices. - - Under the HSA runtime the artifact must be PDI + insts, since the ROCR AIE - dispatch path consumes a raw ``aie.pdi`` and its ``insts.bin`` sidecar. In - that case "pdi" is forced and an explicit ELF/xclbin request is rejected. """ if runtime is None: runtime = npu_config.runtime npu_version = detect_npu_version(runtime) configured_format = npu_config.output_format - if runtime == "hsa": - if configured_format is not None and configured_format != "pdi": - raise RuntimeError( - f"ROCR requires the 'pdi' output format, but " - f"output_format={configured_format!r} was requested. Unset " - "AMD_TRITON_NPU_OUTPUT_FORMAT (or npu_config.output_format), or " - "set it to 'pdi'." - ) - return "pdi" if configured_format is not None: if configured_format == "elf" and npu_version == "npu1": raise RuntimeError( @@ -2450,6 +2437,52 @@ def _put_aircc_artifacts(cache, artifacts, output_format): return {"bin_path": bin_path, "insts_path": insts_path} +# Inverse of NPU_AGENT_NAMES (driver.py's existing {"aie2": "npu1", "aie2p": +# "npu2"} map): the hsaco section an npu_version packs into. +_NPU_VERSION_TO_ARCH = {v: k for k, v in NPU_AGENT_NAMES.items()} + + +def _pack_hsaco(output_format, artifacts, arch, kernel_name): + """Pack aircc's artifacts (see ``_aircc_compile``) into hsaco bytes. + + ``kernel_name`` is the string the packed kernel is looked up by at + dispatch time: the instance-qualified ":" form + ``_extract_elf_kernel_name`` already produces for a full ELF (the exact + string XRT's own ELF launcher already looks its kernel up by), or the + plain Triton kernel name for PDI+insts -- packing and the later + ``hsa_executable_get_symbol_by_name`` only have to agree with each + other, not with any real ELF symbol. + + Returns: + bytes: the packed hsaco. + + Raises: + RuntimeError: via aie.compiler.hsaco.pack, if xclbinutil/llvm-objcopy + fail or the artifacts are malformed. + """ + from aie.compiler.hsaco import pack + + if output_format == "elf": + (kernel,) = pack.kernels_from_full_elf( + artifacts["elf_path"], names=[kernel_name] + ) + else: + pdi, num_cols = pack.partition_from_xclbin(artifacts["bin_path"]) + insts = Path(artifacts["insts_path"]).read_bytes() + kernel = { + "name": kernel_name, + "pdi": pdi, + "num_cols": num_cols, + "insts": insts, + } + section = pack.build_section(arch, [kernel]) + with tempfile.TemporaryDirectory() as scratch: + path = os.path.join(scratch, "kernel.hsaco") + pack.ensure_hsaco(path) + pack.inject(path, arch, section) + return Path(path).read_bytes() + + def compile_module( launcher_src, kernel_placeholder_name, @@ -2463,8 +2496,7 @@ def compile_module( ``link_profile`` selects the runtime the generated launcher links against: * ``"xrt"`` -- XRT (xclbin/ELF); links xrt_coreutil + uuid. - * ``"hsa"`` -- HSA AIE dispatch via ROCR; links libhsa-runtime64. Requires - ``output_format == "pdi"``. + * ``"hsa"`` -- HSA AIE dispatch via ROCR; links libhsa-runtime64. """ if link_profile not in ("xrt", "hsa"): raise ValueError(f"link_profile must be 'xrt' or 'hsa'; got {link_profile!r}") @@ -2622,6 +2654,22 @@ def launch( else: cache_bin_path = cached_artifacts["bin_path"] cache_insts_path = cached_artifacts["insts_path"] + if link_profile == "hsa": + cache_hsaco_path = cache.get_file("kernel.hsaco") + if cache_hsaco_path is None: + # Force a real rebuild: cache_path gates that branch below, + # not cached_artifacts, so both have to be cleared or a + # cache dir with a compiled .so but no kernel.hsaco (e.g. + # left over from before this hsaco path existed) would + # still take the cache-hit branch and reach the final + # dispatch with no cache_hsaco_path/hsaco_kernel_name set. + cached_artifacts = None + cache_path = None + elif output_format == "elf": + with open(cache_elf_kernel_path) as f: + hsaco_kernel_name = f.read().strip() + else: + hsaco_kernel_name = kernel_name if cache_path is None: with tempfile.TemporaryDirectory() as tmpdir: @@ -2719,6 +2767,19 @@ def launch( else: cache_bin_path = cached_artifacts["bin_path"] cache_insts_path = cached_artifacts["insts_path"] + if link_profile == "hsa": + arch = _NPU_VERSION_TO_ARCH[npu_version] + hsaco_kernel_name = ( + artifacts["elf_kernel_name"] + if output_format == "elf" + else kernel_name + ) + hsaco_bytes = _pack_hsaco( + output_format, artifacts, arch, hsaco_kernel_name + ) + cache_hsaco_path = cache.put( + hsaco_bytes, "kernel.hsaco", binary=True + ) with open(so_path, "rb") as f: cache_path = cache.put(f.read(), filename, binary=True) @@ -2756,7 +2817,12 @@ def launch( raise RuntimeError(f"Cannot find {name} module in {cache_path}") mod = importlib.util.module_from_spec(spec) spec.loader.exec_module(mod) - if output_format == "elf": + if link_profile == "hsa": + hsaco_path_str = cache_hsaco_path + if IS_WINDOWS and hsaco_path_str.startswith("\\\\?\\"): + hsaco_path_str = hsaco_path_str[4:] + mod.set_paths(hsaco_path_str, hsaco_kernel_name) + elif output_format == "elf": # Read the cached kernel name with open(cache_elf_kernel_path) as f: elf_kernel_name = f.read().strip() @@ -2809,8 +2875,9 @@ class NPULauncher(object): def __init__(self, src, metadata, runtime="xrt"): """Build the host launcher for ``src`` targeting ``runtime``. - ``runtime`` is "xrt" (xclbin/ELF via XRT) or "hsa" (PDI + insts via - HSA/ROCR); it is bound by ``NPUDriver`` when it constructs the launcher. + ``runtime`` is "xrt" (xclbin/ELF via XRT) or "hsa" (hsaco-packed + PDI+insts or full ELF via HSA/ROCR); it is bound by ``NPUDriver`` when + it constructs the launcher. """ constants, signature = extract_signature_and_constants(src) if runtime == "hsa": @@ -2818,7 +2885,9 @@ def __init__(self, src, metadata, runtime="xrt"): # (large) HSA codegen module is only loaded when the HSA path is used. from .hsa_launcher import _generate_hsa_launcher - self.output_format = "pdi" + # Detect output format the same way the XRT branch does: ELF for + # npu2, xclbin for npu1. Both are packed into an hsaco before dispatch. + self.output_format = _get_output_format(runtime=runtime) launcher_src = _generate_hsa_launcher( constants, signature, @@ -2965,7 +3034,7 @@ def __init__(self, runtime=None): ``runtime`` selects how kernels are dispatched: * ``"xrt"`` -- XRT (xclbin on npu1, ELF on npu2). - * ``"hsa"`` -- HSA/ROCR AIE dispatch (PDI + insts). + * ``"hsa"`` -- HSA/ROCR AIE dispatch (hsaco-packed PDI+insts or full ELF). Defaults to ``npu_config.runtime`` (the ``AMD_TRITON_NPU_RUNTIME`` env var, itself defaulting to ``"xrt"``), so ``NPUDriver()`` honors the diff --git a/scripts/test_hsa_hsaco_pack.py b/scripts/test_hsa_hsaco_pack.py new file mode 100644 index 00000000..9cd6def5 --- /dev/null +++ b/scripts/test_hsa_hsaco_pack.py @@ -0,0 +1,68 @@ +#!/usr/bin/env python3 +# Copyright (C) 2026, Advanced Micro Devices, Inc. All rights reserved. +# SPDX-License-Identifier: MIT +"""Pack an HSA kernel into an hsaco and verify its section, without hardware. + +Forces npu_config.compile_only so compile_module stops right after caching +the packed hsaco, before set_paths/dispatch -- so this needs the aiecc +toolchain but not a live NPU device or ROCR build. +""" +from __future__ import annotations + +import os +import sys + +_HERE = os.path.dirname(os.path.abspath(__file__)) +sys.path.insert(0, os.path.join(os.path.dirname(_HERE), "examples", "zero_copy")) + + +def skip(msg): + print(f"SKIP: {msg}") + sys.exit(77) + + +def main(): + try: + import triton + from aie.compiler.hsaco import dump + from triton.backends.amd_triton_npu.driver import ( + NPUDriver, + get_npu_cache_dir, + npu_config, + ) + from common import add_chain + except ImportError as e: + skip(f"cannot import the backend or aie.compiler.hsaco: {e}") + + npu_config.compile_only = True + triton.runtime.driver.set_active(NPUDriver("hsa")) + + import torch + + zeros = torch.zeros(4096, dtype=torch.float32) + grid = (4096 // add_chain.BLOCK_SIZE,) + compiled = add_chain.add_kernel[grid]( + zeros, zeros, zeros, 4096, BLOCK_SIZE=add_chain.BLOCK_SIZE + ) + cache_dir = get_npu_cache_dir(compiled) + if cache_dir is None: + skip("no NPU cache dir -- compile_only path did not run") + hsaco_path = os.path.join(cache_dir, "kernel.hsaco") + if not os.path.isfile(hsaco_path): + print(f"FAIL: {hsaco_path} was not produced") + return 1 + + sections = dump.read_sections_from_hsaco(hsaco_path) + ok = True + for arch, data in sections: + info = dump.parse_section(data) + print(f"arch={arch} kernels={[k['name'] for k in info['kernels']]}") + for k in info["kernels"]: + if k["num_cols"] <= 0: + print(f"FAIL: kernel {k['name']!r} has num_cols={k['num_cols']}") + ok = False + return 0 if ok else 1 + + +if __name__ == "__main__": + sys.exit(main()) From 520e91b1a798352b3940c267d91511457b645340 Mon Sep 17 00:00:00 2001 From: Yiannis Papadopoulos Date: Fri, 9 Oct 2026 07:44:24 -0600 Subject: [PATCH 4/7] Fail fast on HSA decode's per-token instruction-stream patching The final whole-branch review caught a real gap the plan missed: this hsaco/kernel-object switch silently broke examples/llm_q4nx/hsa_decode.py, the HSA decode path for llama-3.2-1b/3b, qwen3-4b, and gemma3-4b. It called triton_npu_hsa_prepare(pdi_path, insts_path) and patched 248 L-dependent instruction words per token in the runtime-owned buffer the dispatch packet pointed at -- a buffer that no longer exists once instructions live inside a frozen ROCr executable. Porting this to a per-context-length packed-executable cache (mirroring mlir-aie's own DispatchTime[T] LRU) was considered and explicitly not done: in a normal single-sequence decode the context length increments every token with no repeats, so an LRU buys nothing and every token would pay a full ROCr executable load instead of the old in-place memcpy -- an unmeasured, likely-large regression on a path benchmarked at ~51 tok/s. Instead: HsaProgram now raises immediately and clearly, before any weight loading; the four registry entries that used to advertise HSA decode support are flipped off (so the harness's existing up-front guard catches this before a multi-minute weight load, not after); the guard's error message covers both rejection reasons instead of asserting the wrong one; and the llama32_1b README is updated to say so. Also folds in the final review's other findings: AMD_TRITON_NPU_OUTPUT_FORMAT =pdi under the HSA runtime now gets a clear rejection (previously it would silently reach the xclbin-shaped packing branch and fail inside xclbinutil with a confusing error), the aie.compiler.hsaco import failure now names what it needs, and a few stale comments/docstrings are corrected (_pack_hsaco's Raises section, NPULauncher's PDI-era comment, get_npu_cache_dir's missing mention of kernel.hsaco). Co-Authored-By: Claude Sonnet 5 --- amd_triton_npu/backend/driver.py | 45 +++++++++++--- examples/llama32_1b_q4nx/README.md | 41 ++++++++----- examples/llm_q4nx/harness.py | 17 ++++-- examples/llm_q4nx/hsa_decode.py | 96 ++++++++++++++---------------- examples/llm_q4nx/registry.py | 20 +++++-- 5 files changed, 136 insertions(+), 83 deletions(-) diff --git a/amd_triton_npu/backend/driver.py b/amd_triton_npu/backend/driver.py index e7c08ba6..23547dd0 100644 --- a/amd_triton_npu/backend/driver.py +++ b/amd_triton_npu/backend/driver.py @@ -1287,6 +1287,13 @@ def _get_output_format(runtime=None): the ``AMD_TRITON_NPU_OUTPUT_FORMAT`` env var). If not set, defaults to "elf" on npu2 and "xclbin" on npu1. ELF format is only supported on npu2 (AIE2P) devices. + + Under the HSA runtime, "pdi" is rejected: a raw ``aie.pdi`` has no + ``AIE_PARTITION`` section, so the hsaco packer (which needs the PDI's + partition/column-count metadata, not just its bytes) has nothing to read + it from. Unlike XRT, which can dispatch a raw PDI + insts.bin directly, + HSA always packs through ``aie.compiler.hsaco.pack``, which only + understands the xclbin and full-ELF shapes. """ if runtime is None: runtime = npu_config.runtime @@ -1299,6 +1306,14 @@ def _get_output_format(runtime=None): "Unset or change AMD_TRITON_NPU_OUTPUT_FORMAT, or set " "npu_config.output_format to 'xclbin', 'pdi', or None." ) + if configured_format == "pdi" and runtime == "hsa": + raise RuntimeError( + "output_format='pdi' is not supported under the HSA runtime: " + "a raw PDI carries no partition/column-count metadata, which " + "the hsaco packer needs. Unset AMD_TRITON_NPU_OUTPUT_FORMAT " + "(or npu_config.output_format) to auto-detect xclbin/elf, or " + "set it to one of those explicitly." + ) return configured_format # Auto-detect: ELF for npu2, xclbin for npu1. return "elf" if npu_version == "npu2" else "xclbin" @@ -2457,10 +2472,22 @@ def _pack_hsaco(output_format, artifacts, arch, kernel_name): bytes: the packed hsaco. Raises: + ValueError: via aie.compiler.hsaco.pack, if the xclbin/full-ELF + artifacts are malformed (e.g. a missing partition column count, or + ``kernel_name`` not found in a full ELF's kernel list). RuntimeError: via aie.compiler.hsaco.pack, if xclbinutil/llvm-objcopy - fail or the artifacts are malformed. + cannot be located or fail. """ - from aie.compiler.hsaco import pack + try: + from aie.compiler.hsaco import pack + except ImportError as e: + raise RuntimeError( + "could not import aie.compiler.hsaco: the installed mlir-aie " + "does not have hsaco packing support yet. This needs an " + "mlir-aie build from (or after) branch " + "ypapadop-amd/hsa-hsaco-loading -- utils/mlir-air-hash.txt is not " + "yet pinned to one (see this change's PR description)." + ) from e if output_format == "elf": (kernel,) = pack.kernels_from_full_elf( @@ -2881,8 +2908,9 @@ def __init__(self, src, metadata, runtime="xrt"): """ constants, signature = extract_signature_and_constants(src) if runtime == "hsa": - # HSA consumes PDI + insts and links ROCR. Import lazily so the - # (large) HSA codegen module is only loaded when the HSA path is used. + # HSA consumes a packed hsaco (xclbin/PDI+insts or full ELF) and + # links ROCR. Import lazily so the (large) HSA codegen module is + # only loaded when the HSA path is used. from .hsa_launcher import _generate_hsa_launcher # Detect output format the same way the XRT branch does: ELF for @@ -2944,11 +2972,14 @@ def get_npu_cache_dir(compiled_kernel): The NPU backend stores hardware-specific artifacts in a separate cache directory from Triton's main compiler cache. Depending on the selected - output format, the directory contains either: + output format and link profile, the directory contains: - * xclbin output: ``aie.xclbin``, ``insts.bin``, and + * xclbin output (xrt): ``aie.xclbin``, ``insts.bin``, and + ``__npu_dispatch.so`` + * elf output (xrt): ``aie.elf``, ``elf_kernel_name.txt``, and ``__npu_dispatch.so`` - * elf output: ``aie.elf``, ``elf_kernel_name.txt``, and + * hsa (either output format, packed into one hsaco): the above + format-specific artifacts, plus ``kernel.hsaco`` and ``__npu_dispatch.so`` This function returns the path to that directory. diff --git a/examples/llama32_1b_q4nx/README.md b/examples/llama32_1b_q4nx/README.md index 402448c7..26fc1b1d 100644 --- a/examples/llama32_1b_q4nx/README.md +++ b/examples/llama32_1b_q4nx/README.md @@ -190,32 +190,43 @@ python prefill.py --backend cpu --kv-out /tmp/kv.npz ### Running on the HSA runtime -```bash +**Currently unsupported** (`AMD_TRITON_NPU_RUNTIME=hsa make chat` fails fast +with a clear error; `registry.py`'s `supports_hsa=False` for this model +reflects it). ROCr's AIE hsaco/kernel-object dispatch removed +`triton_npu_hsa_patch_insts`, and the decode below depended on it. See +`../llm_q4nx/hsa_decode.py`'s `HsaProgram` docstring and +`docs/superpowers/specs/2026-10-08-hsa-hsaco-dispatch-design.md` for why, and +what porting it would need (a per-context-length packed-executable cache, +unmeasured). The rest of this section is kept as a record of what worked +under the old (PDI + instruction-stream address) HSA dispatch ABI. + +~~~bash make compile-decode RUNTIME=hsa # the decode as PDI, not xclbin make chat RUNTIME=hsa -``` +~~~ -Both halves then dispatch through `HsaRuntime` rather than XRT, and the -decode's logits are bit-identical to the XRT path. +Both halves used to dispatch through `HsaRuntime` rather than XRT, and the +decode's logits were bit-identical to the XRT path. -It is not as fast. Measured on Strix, 60 tokens, runtimes alternated: +It was not as fast. Measured on Strix, 60 tokens, runtimes alternated: | | tok/s | |---|---| | XRT | 54.5 – 54.8 | | HSA | 49.1 – 51.5 | -a stable gap of about 9%, reproducible across runs rather than noise. It is +a stable gap of about 9%, reproducible across runs rather than noise. It was not the per-dispatch cache flush -- that is what resident regions removed, and -without them the gap is roughly three times larger. What remains has not been -attributed, so treat the HSA path as correct and close, not as a replacement. - -The interesting part is how the context length gets to the device. HSA has no -scratchpad parameters, so the full-ELF mechanism is unavailable; but the AIE -dispatch packet carries the instruction stream's address per enqueue, so -`../llm_q4nx/hsa_decode.py` patches the 248 L-dependent words per token — the same words -mlir-air's xclbin path rewrites. One PDI serves every context length; the -L=2048 and L=2047 builds are byte-identical. +without them the gap was roughly three times larger. What remained was never +attributed. + +The interesting part was how the context length got to the device. HSA has no +scratchpad parameters, so the full-ELF mechanism was unavailable; but the old +AIE dispatch packet carried the instruction stream's address per enqueue, so +`../llm_q4nx/hsa_decode.py` patched the 248 L-dependent words per token — the +same words mlir-air's xclbin path rewrites. One PDI served every context +length; the L=2048 and L=2047 builds are byte-identical. That observation +still holds; only the per-token patch mechanism it relied on is gone. `--ops` takes `all` or a comma list of `matmul,rms_norm,swiglu`, so a numerical regression can be bisected to a single kernel against the same CPU diff --git a/examples/llm_q4nx/harness.py b/examples/llm_q4nx/harness.py index 224740ee..c9221c34 100644 --- a/examples/llm_q4nx/harness.py +++ b/examples/llm_q4nx/harness.py @@ -1005,13 +1005,18 @@ def main(spec, cfg, prefill_cls, doc=None, argv=None): if os.environ.get("AMD_TRITON_NPU_RUNTIME") == "hsa" and not spec.supports_hsa: # Up front, for the same reason --interactive is below: a property of # the model, not of anything on disk. Without this the run loads - # weights for minutes and then dies inside hsa_decode.py on - # `air.FusedDecoder`, an attribute this model's driver does not define. + # weights for minutes and then dies inside hsa_decode.py -- either on + # `air.FusedDecoder`, an attribute some drivers don't define, or (for + # every model that used to support this) on HsaProgram's fail-fast + # raise, since per-token instruction-stream patching has no home + # under the hsaco/kernel-object HSA dispatch model. raise SystemExit( - f"AMD_TRITON_NPU_RUNTIME=hsa is not supported for {spec.name}. The " - f"HSA decode adapter subclasses mlir-air's `FusedDecoder`, which " - f"only its npz-API drivers define; this model's driver names its " - f"decoder {spec.decoder_class or 'something else'}. Use the " + f"AMD_TRITON_NPU_RUNTIME=hsa is not supported for {spec.name}. " + f"Either its driver's decoder ({spec.decoder_class or 'something else'}) " + f"isn't mlir-air's `FusedDecoder`-API, which the HSA adapter " + f"subclasses, or its decode depends on per-token instruction-" + f"stream patching (triton_npu_hsa_patch_insts), which the current " + f"hsaco/kernel-object HSA dispatch does not support. Use the " f"default XRT runtime." ) diff --git a/examples/llm_q4nx/hsa_decode.py b/examples/llm_q4nx/hsa_decode.py index 501ad08f..8ed79e12 100644 --- a/examples/llm_q4nx/hsa_decode.py +++ b/examples/llm_q4nx/hsa_decode.py @@ -2,8 +2,13 @@ # SPDX-License-Identifier: MIT """The fused Q4NX decode on the HSA runtime. -mlir-air's decoder dispatches through pyxrt. This runs the same design on -HsaRuntime instead, which is what makes a fully-HSA chatbot possible. +**Currently disabled** (`HsaProgram` raises on construction, before any +weight loading) -- see its docstring below for why. This file is kept for the +host-side plumbing (shared KV/weight buffers, geometry, sampling) and the +record of a design that worked under the old HSA ABI, not as a runnable path. + +mlir-air's decoder dispatches through pyxrt. This ran the same design on +HsaRuntime instead, which is what made a fully-HSA chatbot possible. The whole problem is the context length. It changes every token, and an AIE dispatch has only two ways to hear about it: @@ -15,17 +20,20 @@ `xrt::run::get_ctrl_scratchpad_bo`, an XRT API with no HSA counterpart (`grep -ic scratch hsa_ext_amd_aie.h` is 0). -So on HSA it has to be the first, and the AIE dispatch packet does carry -`insts_addr_low/high` and `insts_size` per enqueue -- the stream is an input to -each dispatch, not a property of the program. What blocked it was only that -`triton_npu_hsa_prepare` read the stream from a file once. `patch_insts` (added -to HsaRuntime for this) writes into the very buffer the packet already points -at, so the 248 L-dependent words are rewritten per token exactly as the xclbin -path rewrites them. +So on HSA it had to be the first. Under the old (PDI + instruction-stream +address) dispatch ABI, the AIE dispatch packet carried `insts_addr_low/high` +and `insts_size` per enqueue -- the stream was an input to each dispatch, not +a property of the program -- so `triton_npu_hsa_patch_insts` could rewrite the +248 L-dependent words in place, in the very buffer the packet already pointed +at, exactly as mlir-air's xclbin path rewrites them. Under the newer hsaco/ +kernel-object ABI there is no such buffer: the instruction stream lives inside +a frozen ROCr executable, resolved to an opaque kernel object at load time, and +`triton_npu_hsa_patch_insts` does not exist any more. See `HsaProgram` below. The design itself is L-independent: the L=2048 and L=2047 builds produce a BYTE-IDENTICAL PDI and differ only in those insts words. One PDI serves every -context length. +context length -- this part of the design is unaffected by the ABI change and +would still apply to whatever ports it. Weights and the KV cache live in shared regions, so they are dispatched on in place. Copying 0.7 GB of weights per token would cap throughput near 23 tok/s @@ -130,50 +138,36 @@ def mark_resident(lib, arr): class HsaProgram: - """A prepared PDI + instruction stream, with the stream patchable per token.""" + """Disabled: per-token instruction-stream patching has no home under the + hsaco/kernel-object HSA dispatch model. + + This class used to `triton_npu_hsa_prepare(pdi_path, insts_path)` once and + `triton_npu_hsa_patch_insts` the 248 L-dependent words per token, in place, + in the runtime-owned buffer the dispatch packet pointed at. ROCr's AIE + hsaco support removed that buffer: `triton_npu_hsa_prepare` now takes + `(hsaco_path, kernel_name)`, loads the instruction stream into a frozen + ROCr executable, and dispatches it by an opaque kernel object -- + `triton_npu_hsa_patch_insts` no longer exists, because there is nothing + left to patch in place. + + Porting this decode path means packing and loading a fresh executable per + distinct context length L (mirroring mlir-aie's own per-call-sequence + executable LRU), which has not been built or measured -- in a normal + single-sequence chat decode L increments by one every token, so an LRU + buys nothing and every token would pay a full ROCr executable load instead + of this class's in-place memcpy. Fails here, before any weight loading, + rather than building a program that silently does the wrong thing. See + docs/superpowers/specs/2026-10-08-hsa-hsaco-dispatch-design.md. + """ def __init__(self, pdi_path, insts_path): - from triton.backends.amd_triton_npu.driver import load_hsa_runtime - - self.lib = load_hsa_runtime() - buf, n = _errbuf() - self.lib.triton_npu_hsa_prepare.restype = ctypes.c_void_p - self.handle = self.lib.triton_npu_hsa_prepare( - pdi_path.encode(), insts_path.encode(), buf, n - ) - if not self.handle: - raise HsaDecodeError(buf.value.decode()) - self.insts_words = os.path.getsize(insts_path) // 4 - - def patch_insts(self, word_offset, words): - """Overwrite `words` (uint32) at `word_offset` in the live stream.""" - a = np.ascontiguousarray(words, dtype=np.uint32) - buf, n = _errbuf() - rc = self.lib.triton_npu_hsa_patch_insts( - ctypes.c_void_p(self.handle), - ctypes.c_uint64(word_offset * 4), - a.ctypes.data_as(ctypes.c_void_p), - ctypes.c_uint64(a.nbytes), - buf, - n, - ) - if rc != 0: - raise HsaDecodeError(buf.value.decode()) - - def dispatch(self, arrays): - """One enqueue over `arrays` in kernel-argument order.""" - n = len(arrays) - ptrs = (ctypes.c_void_p * n)() - sizes = (ctypes.c_uint64 * n)() - for i, a in enumerate(arrays): - ptrs[i] = ctypes.c_void_p(a.ctypes.data) - sizes[i] = ctypes.c_uint64(a.nbytes) - buf, nb = _errbuf() - rc = self.lib.triton_npu_hsa_dispatch( - ctypes.c_void_p(self.handle), ctypes.c_uint32(n), ptrs, sizes, buf, nb + raise HsaDecodeError( + "HSA decode (per-token instruction-stream patching via " + "triton_npu_hsa_patch_insts) is not supported under the hsaco/" + "kernel-object HSA dispatch model -- that C ABI entry point no " + "longer exists. Use the default XRT runtime for this model's " + "chat/decode path." ) - if rc != 0: - raise HsaDecodeError(buf.value.decode()) class InstsForL: diff --git a/examples/llm_q4nx/registry.py b/examples/llm_q4nx/registry.py index cdb260e0..96cce2fd 100644 --- a/examples/llm_q4nx/registry.py +++ b/examples/llm_q4nx/registry.py @@ -207,7 +207,13 @@ def kv_source_map(self): air_inference="llama32_1b_q4nx_inference.py", tokenizer_fallback="~/q4nx_data/tokenizer/Llama-3.2-1B", extra_packages=("llama32_1b",), - supports_hsa=True, + # HSA decode needs per-token instruction-stream patching + # (triton_npu_hsa_patch_insts), which the hsaco/kernel-object HSA dispatch + # model does not support -- the instruction stream now lives inside a + # frozen ROCr executable rather than a runtime-owned buffer. See + # examples/llm_q4nx/hsa_decode.py's HsaProgram and + # docs/superpowers/specs/2026-10-08-hsa-hsaco-dispatch-design.md. + supports_hsa=False, ) @@ -248,7 +254,9 @@ def kv_source_map(self): extra_packages=("llama32_3b", "llama32_1b_q4nx"), driver_api="prefiller", decoder_class="FusedDecode3B", - supports_hsa=True, + # See the 1B entry above: HSA decode needs per-token instruction-stream + # patching, which the hsaco/kernel-object HSA dispatch model dropped. + supports_hsa=False, ) @@ -338,7 +346,9 @@ def kv_source_map(self): # resident at once. min_host_gib=14.0, extra_packages=("qwen3_8b_q4nx", "qwen3_4b"), - supports_hsa=True, + # See the 1B entry above: HSA decode needs per-token instruction-stream + # patching, which the hsaco/kernel-object HSA dispatch model dropped. + supports_hsa=False, ) @@ -381,7 +391,9 @@ def kv_source_map(self): # because its 262208-row LM head is a separate tensor where Qwen3-4B's is # tied. min_host_gib=17.0, - supports_hsa=True, + # See the 1B entry above: HSA decode needs per-token instruction-stream + # patching, which the hsaco/kernel-object HSA dispatch model dropped. + supports_hsa=False, ) #: `DECODE_ENV` in `fused_decode_ple/Makefile` -- the ENGINE's, not the model's: From 4c419f2a2e46d0155ccaf7a6aca4d6eedc796957 Mon Sep 17 00:00:00 2001 From: Yiannis Papadopoulos Date: Fri, 9 Oct 2026 07:47:54 -0600 Subject: [PATCH 5/7] Match the surrounding file's standard triple-backtick code fences --- examples/llama32_1b_q4nx/README.md | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/examples/llama32_1b_q4nx/README.md b/examples/llama32_1b_q4nx/README.md index 26fc1b1d..7ebf41c3 100644 --- a/examples/llama32_1b_q4nx/README.md +++ b/examples/llama32_1b_q4nx/README.md @@ -200,10 +200,10 @@ what porting it would need (a per-context-length packed-executable cache, unmeasured). The rest of this section is kept as a record of what worked under the old (PDI + instruction-stream address) HSA dispatch ABI. -~~~bash +```bash make compile-decode RUNTIME=hsa # the decode as PDI, not xclbin make chat RUNTIME=hsa -~~~ +``` Both halves used to dispatch through `HsaRuntime` rather than XRT, and the decode's logits were bit-identical to the XRT path. From 2cc2cc9f191e3318b552cb8a145af4c3c8e44ab6 Mon Sep 17 00:00:00 2001 From: Yiannis Papadopoulos Date: Wed, 21 Oct 2026 10:48:57 -0600 Subject: [PATCH 6/7] hsa: fix two cache-correctness bugs, consolidate hsaco packing, add prepare_bytes From a code review of the hsaco/kernel-object merge: - compile_module forced a rebuild when kernel.hsaco alone was missing from a cache hit, but not when the format-specific artifacts (aie.elf/aie.xclbin) were themselves missing while the compiled __npu_dispatch.so survived -- that left cache_hsaco_path/hsaco_kernel_name unassigned, raising UnboundLocalError instead of rebuilding. Reproduced on hardware (deleted a cached kernel's aie.elf, kept its .so) and confirmed fixed. - _scratchpad_relocs never checked the scratchpad's declared size was a multiple of 4; HsaElfProgram then viewed it as whole uint32 words via integer division, silently truncating a misbuilt one instead of failing where the build is actually broken. Verified against the real decode ELF (passes) and a corrupted copy (raises). Also, while in there: - Consolidated the hsaco-kernel-name-derivation + pack + cache sequence, duplicated across compile_module's and MultiLaunchBuilder.compile's cache-hit and rebuild branches, into _hsaco_kernel_name/_cached_hsaco/ _pack_and_cache_hsaco. - Added triton_npu_hsa_prepare_bytes (loads an hsaco already in memory, uncached -- there's no stable key for a buffer, and each one HsaElfProgram builds is single-use anyway) so hsa_decode.py no longer writes its hsaco to a temp file solely to read it back through the path-based ABI. Verified on hardware: matmul HSA compile from cold and warm caches, the cache-corruption regression above, and the Qwen2.5-3B HSA decode (scratchpad ELF, now loaded via prepare_bytes) all still produce correct results. Co-Authored-By: Claude Sonnet 5 --- amd_triton_npu/backend/driver.py | 92 +++++++++++++++---- .../backend/include/HsaRuntime/HsaRuntime.cpp | 41 +++++++++ .../backend/include/HsaRuntime/HsaRuntime.h | 14 +++ amd_triton_npu/backend/multilaunch.py | 28 +++--- examples/llm_q4nx/hsa_decode.py | 22 ++--- 5 files changed, 153 insertions(+), 44 deletions(-) diff --git a/amd_triton_npu/backend/driver.py b/amd_triton_npu/backend/driver.py index f39ac4d3..ae220a8b 100644 --- a/amd_triton_npu/backend/driver.py +++ b/amd_triton_npu/backend/driver.py @@ -2558,6 +2558,12 @@ def _scratchpad_relocs(blob): raise ValueError(f"scratchpad relocation with patch scheme {scheme}") size = st_size relocs.append((off, f.sections[st_shndx].offset + r_offset, addend)) + # hsa_decode.py views the scratchpad as whole uint32 words + # (size // 4); a size that isn't a multiple of 4 would silently truncate + # the last partial word instead of failing where the build is actually + # broken. + if size % 4 != 0: + raise ValueError(f"scratchpad size {size} is not a multiple of 4") return size, relocs @@ -2636,6 +2642,50 @@ def _pack_hsaco(output_format, artifacts, arch, kernel_name, scratchpad_addr=Non return Path(path).read_bytes() +def _hsaco_kernel_name(output_format, elf_kernel_name, pdi_kernel_name): + """The name ``_pack_hsaco`` looks the kernel up by: the full ELF's own + instance-qualified name for ``output_format == "elf"``, or the caller's own + (arbitrary) kernel name for a PDI + instruction sequence -- packing and the + later ``hsa_executable_get_symbol_by_name`` only have to agree with each + other. + """ + return elf_kernel_name if output_format == "elf" else pdi_kernel_name + + +def _cached_hsaco(cache, output_format, elf_kernel_name, pdi_kernel_name): + """The cached hsaco's ``(path, kernel_name)``, or ``None`` if it is missing. + + Callers that already resolved the format-specific aircc artifacts + (xclbin/insts or elf/elf_kernel_name) from the cache ask this before + deciding to rebuild. ``None`` means those format-specific artifacts are + cached but the hsaco itself is not (e.g. a cache dir from before this + hsaco path existed), so the caller should treat that as a cache miss for + the hsaco alone, not for the format-specific artifacts it already has. + """ + hsaco_path = cache.get_file("kernel.hsaco") + if hsaco_path is None: + return None + return hsaco_path, _hsaco_kernel_name( + output_format, elf_kernel_name, pdi_kernel_name + ) + + +def _pack_and_cache_hsaco( + cache, output_format, artifacts, npu_version, elf_kernel_name, pdi_kernel_name +): + """Pack aircc's artifacts into an hsaco, cache it, and return ``(hsaco_path, + hsaco_kernel_name)``. ``elf_kernel_name``/``pdi_kernel_name`` are as for + ``_hsaco_kernel_name``. + """ + arch = _NPU_VERSION_TO_ARCH[npu_version] + hsaco_kernel_name = _hsaco_kernel_name( + output_format, elf_kernel_name, pdi_kernel_name + ) + hsaco_bytes = _pack_hsaco(output_format, artifacts, arch, hsaco_kernel_name) + hsaco_path = cache.put(hsaco_bytes, "kernel.hsaco", binary=True) + return hsaco_path, hsaco_kernel_name + + def compile_module( launcher_src, kernel_placeholder_name, @@ -2800,6 +2850,15 @@ def launch( filename = f"{name}.pyd" if IS_WINDOWS else f"{name}.so" cache_path = cache.get_file(filename) cached_artifacts = _get_cached_aircc_artifacts(cache, output_format) + if cached_artifacts is None: + # A compiled __npu_dispatch.so with no aircc artifacts behind it is + # unusable: cache_elf_path/cache_bin_path (and, under HSA, + # cache_hsaco_path/hsaco_kernel_name) are only set inside the + # `cached_artifacts is not None` block below, so leaving cache_path + # as a cache hit here would reach the final dispatch with those + # names never assigned -- an UnboundLocalError instead of a clean + # rebuild. Treat a partially-evicted cache dir as a full miss. + cache_path = None if cached_artifacts is not None: if output_format == "elf": cache_elf_path = cached_artifacts["elf_path"] @@ -2808,8 +2867,14 @@ def launch( cache_bin_path = cached_artifacts["bin_path"] cache_insts_path = cached_artifacts["insts_path"] if link_profile == "hsa": - cache_hsaco_path = cache.get_file("kernel.hsaco") - if cache_hsaco_path is None: + elf_kernel_name = None + if output_format == "elf": + with open(cache_elf_kernel_path) as f: + elf_kernel_name = f.read().strip() + cached_hsaco = _cached_hsaco( + cache, output_format, elf_kernel_name, kernel_name + ) + if cached_hsaco is None: # Force a real rebuild: cache_path gates that branch below, # not cached_artifacts, so both have to be cleared or a # cache dir with a compiled .so but no kernel.hsaco (e.g. @@ -2818,11 +2883,8 @@ def launch( # dispatch with no cache_hsaco_path/hsaco_kernel_name set. cached_artifacts = None cache_path = None - elif output_format == "elf": - with open(cache_elf_kernel_path) as f: - hsaco_kernel_name = f.read().strip() else: - hsaco_kernel_name = kernel_name + cache_hsaco_path, hsaco_kernel_name = cached_hsaco if cache_path is None: with tempfile.TemporaryDirectory() as tmpdir: @@ -2921,17 +2983,13 @@ def launch( cache_bin_path = cached_artifacts["bin_path"] cache_insts_path = cached_artifacts["insts_path"] if link_profile == "hsa": - arch = _NPU_VERSION_TO_ARCH[npu_version] - hsaco_kernel_name = ( - artifacts["elf_kernel_name"] - if output_format == "elf" - else kernel_name - ) - hsaco_bytes = _pack_hsaco( - output_format, artifacts, arch, hsaco_kernel_name - ) - cache_hsaco_path = cache.put( - hsaco_bytes, "kernel.hsaco", binary=True + cache_hsaco_path, hsaco_kernel_name = _pack_and_cache_hsaco( + cache, + output_format, + artifacts, + npu_version, + artifacts.get("elf_kernel_name"), + kernel_name, ) with open(so_path, "rb") as f: cache_path = cache.put(f.read(), filename, binary=True) diff --git a/amd_triton_npu/backend/include/HsaRuntime/HsaRuntime.cpp b/amd_triton_npu/backend/include/HsaRuntime/HsaRuntime.cpp index d5037e2d..beec172a 100644 --- a/amd_triton_npu/backend/include/HsaRuntime/HsaRuntime.cpp +++ b/amd_triton_npu/backend/include/HsaRuntime/HsaRuntime.cpp @@ -528,6 +528,21 @@ class HsaRuntime { return raw; } + // Load + resolve an hsaco already in memory, as prepare() does for one on + // disk. Not cached: bytes_programs_ only owns the result so its lifetime + // outlives this call, the same guarantee prepare()'s cache gives by + // accident. Thread-safe. + triton_npu_hsa_program *prepare_bytes(std::vector hsaco_bytes, + const char *kernel_name) { + auto prog = std::make_unique(); + prog->exe = + LoadedExecutable::load(aie_agent_, std::move(hsaco_bytes), kernel_name); + triton_npu_hsa_program *raw = prog.get(); + std::lock_guard lock(programs_mtx_); + bytes_programs_.push_back(std::move(prog)); + return raw; + } + // Allocate a zeroed control scratchpad and return its host address, with the // address the NPU reaches it by in *dev_addr. Kept until teardown: that // address is baked into a program's control code, and programs are never @@ -930,6 +945,11 @@ class HsaRuntime { // destructor freeing them all. Unlike regions_ below, which cannot be. std::unordered_map> programs_; + // Programs from prepare_bytes(), which has no key to look one up by (each + // is loaded from an hsaco meant for a single use). Owns them so they + // outlive the call, same as programs_ for the path-keyed ones; the one + // iteration is the destructor freeing them all, like programs_. + std::vector> bytes_programs_; // Control scratchpads by host address (see scratchpad_alloc). Guarded by // programs_mtx_, since each one belongs to a program. std::unordered_map scratchpads_; @@ -1080,6 +1100,7 @@ class HsaRuntime { vmem_free(b); vmem_pool_.clear(); programs_.clear(); + bytes_programs_.clear(); scratchpads_.clear(); kernarg_buffer_ = PoolBuffer{}; if (signal_.handle) { @@ -1619,6 +1640,26 @@ triton_npu_hsa_prepare(const char *hsaco_path, const char *kernel_name, } } +extern "C" triton_npu_hsa_program_t +triton_npu_hsa_prepare_bytes(const void *hsaco_data, uint64_t hsaco_size, + const char *kernel_name, char *errbuf, + size_t errbuf_len) { + try { + if (!hsaco_data || !kernel_name) + throw std::runtime_error("null hsaco_data or kernel_name"); + const auto *bytes = static_cast(hsaco_data); + return runtime().prepare_bytes(std::vector(bytes, bytes + hsaco_size), + kernel_name); + } catch (const std::exception &e) { + write_err(errbuf, errbuf_len, "HSA prepare (bytes) failed", e.what()); + return nullptr; + } catch (...) { + write_err(errbuf, errbuf_len, "HSA prepare (bytes) failed", + "unknown error"); + return nullptr; + } +} + extern "C" void *triton_npu_hsa_scratchpad_alloc(uint64_t size, uint64_t *dev_addr, char *errbuf, diff --git a/amd_triton_npu/backend/include/HsaRuntime/HsaRuntime.h b/amd_triton_npu/backend/include/HsaRuntime/HsaRuntime.h index 08172678..9ded5979 100644 --- a/amd_triton_npu/backend/include/HsaRuntime/HsaRuntime.h +++ b/amd_triton_npu/backend/include/HsaRuntime/HsaRuntime.h @@ -55,6 +55,20 @@ triton_npu_hsa_program_t triton_npu_hsa_prepare(const char *hsaco_path, char *errbuf, size_t errbuf_len); +// Like triton_npu_hsa_prepare, but pack-loads an hsaco already in memory +// instead of reading one from a file. The bytes are copied before this call +// returns, so hsaco_data need not outlive it. +// +// Unlike triton_npu_hsa_prepare, not cached by path -- there is no path to key +// on. Meant for an hsaco that is good for only one load (e.g. one baked with a +// single scratchpad's device address, which no other hsaco shares); preparing +// the same bytes twice loads two independent programs rather than reusing one. +triton_npu_hsa_program_t triton_npu_hsa_prepare_bytes(const void *hsaco_data, + uint64_t hsaco_size, + const char *kernel_name, + char *errbuf, + size_t errbuf_len); + // Initialize the runtime (idempotent) and allocate a zeroed control scratchpad // of `size` bytes for a full-ELF design that declares parameters. Returns its // host address and writes the address the NPU reaches it by to *dev_addr, or diff --git a/amd_triton_npu/backend/multilaunch.py b/amd_triton_npu/backend/multilaunch.py index 3172dfb5..53b6a9cd 100644 --- a/amd_triton_npu/backend/multilaunch.py +++ b/amd_triton_npu/backend/multilaunch.py @@ -33,8 +33,8 @@ _get_output_format, _get_cached_aircc_artifacts, _put_aircc_artifacts, - _pack_hsaco, - _NPU_VERSION_TO_ARCH, + _cached_hsaco, + _pack_and_cache_hsaco, detect_npu_version, npu_wait_module, release_external_contexts, @@ -295,12 +295,11 @@ def compile(self, output_format=None, link_profile="xrt"): with open(cached["elf_kernel_name_path"]) as f: elf_kernel_name = f.read() if link_profile == "hsa": - hsaco_path = cache.get_file("kernel.hsaco") - if hsaco_path is not None: - hsaco_kernel_name = ( - elf_kernel_name if output_format == "elf" else self.name - ) - return hsaco_path, hsaco_kernel_name + cached_hsaco = _cached_hsaco( + cache, output_format, elf_kernel_name, self.name + ) + if cached_hsaco is not None: + return cached_hsaco cached = None # hsaco missing: force a rebuild below elif output_format == "elf": return cached["elf_path"], elf_kernel_name @@ -323,13 +322,14 @@ def compile(self, output_format=None, link_profile="xrt"): # Cache the format-specific artifacts and return the cached paths. cached = _put_aircc_artifacts(cache, artifacts, output_format) if link_profile == "hsa": - arch = _NPU_VERSION_TO_ARCH[npu_version] - hsaco_kernel_name = ( - artifacts["elf_kernel_name"] if output_format == "elf" else self.name + return _pack_and_cache_hsaco( + cache, + output_format, + artifacts, + npu_version, + artifacts.get("elf_kernel_name"), + self.name, ) - hsaco_bytes = _pack_hsaco(output_format, artifacts, arch, hsaco_kernel_name) - hsaco_path = cache.put(hsaco_bytes, "kernel.hsaco", binary=True) - return hsaco_path, hsaco_kernel_name if output_format == "elf": return cached["elf_path"], artifacts["elf_kernel_name"] return cached["bin_path"], cached["insts_path"] diff --git a/examples/llm_q4nx/hsa_decode.py b/examples/llm_q4nx/hsa_decode.py index c12c214c..6151519e 100644 --- a/examples/llm_q4nx/hsa_decode.py +++ b/examples/llm_q4nx/hsa_decode.py @@ -25,7 +25,6 @@ import ctypes import os import re -import tempfile import threading import weakref @@ -178,21 +177,18 @@ def __init__(self, elf_path, kernel_name): ctypes.cast(va, ctypes.POINTER(ctypes.c_uint32)), (size // 4,) ) - # Full ELF is aie2p only. The runtime caches programs by path, and this - # hsaco is good only for this scratchpad, so the path names it: the - # runtime never frees a scratchpad, so no other one has its address. + # Full ELF is aie2p only. This hsaco is good only for this scratchpad + # (it is baked with this one's device address), so prepare_bytes loads + # it straight from memory rather than caching it by a path nothing else + # would ever share. hsaco = _pack_hsaco( "elf", {"elf_path": str(elf_path)}, "aie2p", kernel_name, dev_addr.value ) - with tempfile.TemporaryDirectory() as d: - path = os.path.join(d, f"decode_{dev_addr.value:x}.hsaco") - with open(path, "wb") as f: - f.write(hsaco) - buf, n = _errbuf() - self.lib.triton_npu_hsa_prepare.restype = ctypes.c_void_p - self.handle = self.lib.triton_npu_hsa_prepare( - path.encode(), kernel_name.encode(), buf, n - ) + buf, n = _errbuf() + self.lib.triton_npu_hsa_prepare_bytes.restype = ctypes.c_void_p + self.handle = self.lib.triton_npu_hsa_prepare_bytes( + hsaco, ctypes.c_uint64(len(hsaco)), kernel_name.encode(), buf, n + ) if not self.handle: raise HsaDecodeError(buf.value.decode()) self.lock = threading.Lock() From 3e6f1637df89de5fc8104a128d34073b81fd1472 Mon Sep 17 00:00:00 2001 From: Erwei Wang Date: Thu, 22 Oct 2026 10:42:42 -0700 Subject: [PATCH 7/7] HSA decode: define the scratchpad as a variable instead of baking it in ROCR's AIE loader resolves a full ELF's control scratchpad against a variable defined with hsa_executable_agent_global_variable_define, and flushes it on every dispatch (ROCm/rocm-systems#14004). So the decode allocates the scratchpad, loads the cached hsaco with it defined, and writes L with plain stores. This drops the ELF relocation rewriting in driver.py, the per-decoder repack and prepare_bytes, the host-side flush, and the pointer_info probe it needed. Pin rocm-sdk-core 10.3.0a20261027, the first nightly with both ROCm/rocm-systems#14004 and #13908. On an older ROCR, repeated full-ELF dispatches give wrong results and the decode fails to load. Co-Authored-By: Claude Opus 5.5 --- README.md | 16 +- amd_triton_npu/backend/driver.py | 98 +------ .../backend/include/HsaRuntime/HsaRuntime.cpp | 243 ++++-------------- .../backend/include/HsaRuntime/HsaRuntime.h | 66 ++--- examples/llama32_1b_q4nx/README.md | 13 +- examples/llm_q4nx/hsa_decode.py | 117 ++++----- utils/rocm-requirements.txt | 10 +- 7 files changed, 136 insertions(+), 427 deletions(-) diff --git a/README.md b/README.md index bcef692d..89e8d041 100644 --- a/README.md +++ b/README.md @@ -140,11 +140,8 @@ Under HSA the output format is the same as under XRT (`xclbin` on npu1, `elf` on npu2); `pdi` is refused, because the packer reads the partition's column count from the xclbin or ELF. Packing runs mlir-aie's `aie-hsaco` (`aie.compiler.hsaco.pack`). The HSA runtime is Linux-only and requires an -**AIE-capable ROCR** that loads AIE hsacos: rocm-systems `develop` at or after -[`7a71bfb`](https://github.com/ROCm/rocm-systems/commit/7a71bfb024e106a0941d14ad560a2befc5277227) -("feat(rocr/aie): hsaco support for AIE agents"), which provides the AIE -dispatch extension header (`include/hsa/hsa_ext_amd_aie.h`) and -`libhsa-runtime64`. +**AIE-capable ROCR** that loads AIE hsacos and provides the AIE dispatch +extension header (`include/hsa/hsa_ext_amd_aie.h`) and `libhsa-runtime64`. Install the pinned one from [TheRock](https://github.com/ROCm/TheRock)'s nightly ROCm wheels; the backend picks it up automatically, with no environment @@ -158,12 +155,9 @@ pip install "triton-xdna[hsa]" \ `rocm-sdk-core` is the only package needed — it carries both the headers and `libhsa-runtime64`. The pin in `utils/rocm-requirements.txt` is a nightly with -everything the HSA path uses, including the device-address reporting the -full-ELF decode needs. `10.2.0a20261020` is the first nightly with hsaco -support (TheRock moved its rocm-systems pin past `7a71bfb` on 2026-10-19); an -older one fails to build the HSA runtime, since its dispatch packet has no -kernel-object field. The -older indexes no longer receive builds but still resolve under pip: +everything the HSA path uses; an older one may fail to build the HSA runtime, +give wrong results on repeated full-ELF dispatches, or refuse to load the +decode. The older indexes no longer receive builds but still resolve under pip: `https://rocm.nightlies.amd.com/whl-multi-arch/` stopped at 2026-08-22, and the per-GPU-family indexes (`https://rocm.nightlies.amd.com/v2//`) install a runtime too old to work. diff --git a/amd_triton_npu/backend/driver.py b/amd_triton_npu/backend/driver.py index ae220a8b..0ff61cf8 100644 --- a/amd_triton_npu/backend/driver.py +++ b/amd_triton_npu/backend/driver.py @@ -7,7 +7,6 @@ import json import logging import re -import struct import tempfile import threading import sys @@ -2512,84 +2511,7 @@ def _put_aircc_artifacts(cache, artifacts, output_format): _NPU_VERSION_TO_ARCH = {v: k for k, v in NPU_AGENT_NAMES.items()} -# The relocation symbol naming a full ELF's control scratchpad, as aiecc emits -# it; the symbol's st_size is the scratchpad's size in bytes. -_SCRATCHPAD_SYMBOL = b"scratch-pad-ctrl" -# The patch scheme a scratchpad site uses: a plain 64-bit store. -_ADDRESS64 = 8 - - -def _scratchpad_relocs(blob): - """Find a full ELF's control-scratchpad relocations. - - Returns ``(size, relocs)``: the scratchpad's size in bytes (0 if the ELF - has none) and, for each relocation naming it, ``(rela_offset, - site_offset, addend)`` -- where the ``Elf32_Rela`` entry is in the file, - and the file offset and addend of the control-code word it patches. - """ - from aie.compiler.hsaco.elf import ElfFile - - f = ElfFile(blob) - rela = f.section_by_name(".rela.dyn") - dynsym = f.section_by_name(".dynsym") - if rela is None or dynsym is None: - return 0, [] - dynstr = f.sections[dynsym.link] - abi_version = blob[8] # e_ident[EI_ABIVERSION] - size, relocs = 0, [] - for off in range(rela.offset, rela.offset + rela.size, 12): # Elf32_Rela - r_offset, r_info, r_addend = struct.unpack_from("> 8) * 16 # Elf32_Sym - st_name, _, st_size, _, _, st_shndx = struct.unpack_from("> 4 - else: - scheme, addend = r_info & 0xFF, r_addend - # Only control code is patched, as in ROCR; this also keeps a symbol - # with no section (or a special index) from naming a bogus offset. - if st_shndx >= len(f.sections) or not f.sections[st_shndx].name.startswith( - ".ctrltext" - ): - continue - if scheme != _ADDRESS64: - raise ValueError(f"scratchpad relocation with patch scheme {scheme}") - size = st_size - relocs.append((off, f.sections[st_shndx].offset + r_offset, addend)) - # hsa_decode.py views the scratchpad as whole uint32 words - # (size // 4); a size that isn't a multiple of 4 would silently truncate - # the last partial word instead of failing where the build is actually - # broken. - if size % 4 != 0: - raise ValueError(f"scratchpad size {size} is not a multiple of 4") - return size, relocs - - -def full_elf_scratchpad_size(elf_path): - """The size in bytes of a full ELF's control scratchpad, or 0 if it has none.""" - return _scratchpad_relocs(Path(elf_path).read_bytes())[0] - - -def _bake_scratchpad(blob, dev_addr): - """Write a scratchpad's device address into a full ELF's control code. - - ROCR's hsaco loader patches only PDI and argument addresses, and refuses a - full ELF that relocates anything else. So the address is stored at each - scratchpad site here, as ROCR would store a PDI's, and each relocation is - pointed at symbol 0: ROCR skips a relocation whose symbol is not in a - control-code section, and the null symbol is in none. - """ - out = bytearray(blob) - for rela_offset, site, addend in _scratchpad_relocs(blob)[1]: - struct.pack_into(" #include -#include // _mm_clflush, for scratchpad_flush -#include // fstat(), to size the object a dma-buf names -#include // close(), for the dma-buf an import is reached through +#include // fstat(), to size the object a dma-buf names +#include // close(), for the dma-buf an import is reached through #include "hsa/hsa.h" #include "hsa/hsa_ext_amd.h" @@ -55,50 +54,6 @@ namespace { // Queue depth cap. Also bounds the kernarg slot pool (one slot per ring slot). constexpr std::uint32_t QUEUE_SIZE = 32; -// The address the NPU sees for `ptr`, or 0 if this ROCR cannot say. -// -// A ROCR that knows AIE allocations reports the agent's address in -// hsa_amd_pointer_info's agentBaseAddress, as it does for a GPU. One that does -// not reports them as HSA_EXT_POINTER_TYPE_UNKNOWN. -// -// The ROCR a launcher runs against need not be the one it was compiled -// against (it binds by DT_RPATH), so only the running process can answer. -inline std::uint64_t query_device_address(void *ptr) { - hsa_amd_pointer_info_t info{}; - info.size = sizeof(info); - if (hsa_amd_pointer_info(ptr, &info, nullptr, nullptr, nullptr) != - HSA_STATUS_SUCCESS) - return 0; - if (info.type == HSA_EXT_POINTER_TYPE_UNKNOWN || - info.hostBaseAddress == nullptr) - return 0; - - // pointer_info reports the allocation, so an interior pointer needs its - // offset carried across. - const auto offset = static_cast(ptr) - - static_cast(info.hostBaseAddress); - const auto base = reinterpret_cast(info.agentBaseAddress); - if (base == 0 || info.agentBaseAddress == info.hostBaseAddress) - return 0; // reachable only at its host address, which is not patchable - return base + static_cast(offset); -} - -// Write [p, p + n) back from the CPU cache to memory, as ROCR's FlushCpuCache -// does for the buffers a packet lists, so that the device reads what the host -// wrote. -void flush_cpu_cache(const void *p, std::size_t n) { - static const std::uintptr_t line = [] { - const long sz = sysconf(_SC_LEVEL1_DCACHE_LINESIZE); - return static_cast(sz > 0 ? sz : 64); - }(); - const auto begin = reinterpret_cast(p) & ~(line - 1); - const auto end = reinterpret_cast(p) + n; - _mm_mfence(); - for (std::uintptr_t a = begin; a < end; a += line) - _mm_clflush(reinterpret_cast(a)); - _mm_mfence(); -} - // Bytes in the AIE dispatch packet after completion_signal up to // kernarg_address; the ABI requires this to be exactly 24. constexpr std::uint16_t AIE_PACKET_COUNT = 24; @@ -372,6 +327,13 @@ struct SharedRegion { // // hsaco_ is held, not just read once: ROCR reads the code object in place // for as long as the reader exists, so these bytes must outlive it. +// A variable defined in an executable before its code object is loaded, by +// name and host address. +struct Variable { + std::string name; + void *address; +}; + class LoadedExecutable { public: LoadedExecutable() = default; @@ -404,10 +366,14 @@ class LoadedExecutable { std::uint32_t kernarg_size() const { return kernarg_size_; } // Load hsaco_bytes on `agent` and resolve kernel_name to its kernel - // object. Throws std::runtime_error (via HSA_CHECK) naming the failing - // call; already-acquired state is torn down before rethrowing. + // object. `variables` are defined for the agent before the load, so that + // a full ELF's references to buffers it does not carry (such as the + // control scratchpad) resolve to them. Throws std::runtime_error (via + // HSA_CHECK) naming the failing call; already-acquired state is torn down + // before rethrowing. static LoadedExecutable load(hsa_agent_t agent, std::vector hsaco_bytes, - const char *kernel_name) { + const char *kernel_name, + const std::vector &variables) { LoadedExecutable exe; exe.hsaco_ = std::move(hsaco_bytes); try { @@ -416,6 +382,9 @@ class LoadedExecutable { HSA_CHECK(hsa_executable_create_alt( HSA_PROFILE_FULL, HSA_DEFAULT_FLOAT_ROUNDING_MODE_DEFAULT, nullptr, &exe.executable_)); + for (const Variable &v : variables) + HSA_CHECK(hsa_executable_agent_global_variable_define( + exe.executable_, agent, v.name.c_str(), v.address)); HSA_CHECK(hsa_executable_load_agent_code_object( exe.executable_, agent, exe.reader_, nullptr, nullptr)); HSA_CHECK(hsa_executable_freeze(exe.executable_, nullptr)); @@ -510,70 +479,44 @@ class HsaRuntime { return std::string(buf); } - // Load + cache the hsaco at hsaco_path and resolve kernel_name to a kernel - // object; return its program handle. Keyed by (hsaco_path, kernel_name) so - // repeated prepares reuse the same loaded executable. Thread-safe. + // Load + cache the hsaco at hsaco_path, with `variables` defined, and + // resolve kernel_name to a kernel object; return its program handle. Keyed + // by the path, the kernel and the variables, so repeated prepares reuse + // the same loaded executable. Thread-safe. triton_npu_hsa_program *prepare(const char *hsaco_path, - const char *kernel_name) { + const char *kernel_name, + const std::vector &variables) { std::string key = std::string(hsaco_path) + '\0' + kernel_name; + for (const Variable &v : variables) + key += '\0' + v.name + '=' + + std::to_string(reinterpret_cast(v.address)); std::lock_guard lock(programs_mtx_); auto it = programs_.find(key); if (it != programs_.end()) return it->second.get(); auto prog = std::make_unique(); prog->exe = LoadedExecutable::load(aie_agent_, read_host_file(hsaco_path), - kernel_name); + kernel_name, variables); triton_npu_hsa_program *raw = prog.get(); programs_.emplace(std::move(key), std::move(prog)); return raw; } - // Load + resolve an hsaco already in memory, as prepare() does for one on - // disk. Not cached: bytes_programs_ only owns the result so its lifetime - // outlives this call, the same guarantee prepare()'s cache gives by - // accident. Thread-safe. - triton_npu_hsa_program *prepare_bytes(std::vector hsaco_bytes, - const char *kernel_name) { - auto prog = std::make_unique(); - prog->exe = - LoadedExecutable::load(aie_agent_, std::move(hsaco_bytes), kernel_name); - triton_npu_hsa_program *raw = prog.get(); - std::lock_guard lock(programs_mtx_); - bytes_programs_.push_back(std::move(prog)); - return raw; - } - - // Allocate a zeroed control scratchpad and return its host address, with the - // address the NPU reaches it by in *dev_addr. Kept until teardown: that - // address is baked into a program's control code, and programs are never - // unloaded either. Thread-safe. - void *scratchpad_alloc(std::size_t size, std::uint64_t *dev_addr) { + // Allocate a zeroed control scratchpad from the dev pool and return its host + // address. Kept until teardown, like the programs whose control code holds + // its address. Thread-safe. + void *scratchpad_alloc(std::size_t size) { PoolBuffer b = alloc_dev(size); // The host writes parameters into this by index and the device reads every // slot, so an unwritten one has to be a zero rather than whatever the pool // last held. std::memset(b.va(), 0, size); - flush_cpu_cache(b.va(), size); - *dev_addr = device_address(b.va()); void *va = b.va(); std::lock_guard lock(programs_mtx_); scratchpads_.emplace(va, std::move(b)); return va; } - // Flush the host's writes to a scratchpad, so the next dispatch reads them. - void scratchpad_flush(void *va) { - std::size_t size = 0; - { - std::lock_guard lock(programs_mtx_); - auto it = scratchpads_.find(va); - if (it == scratchpads_.end()) - throw std::runtime_error("address is not a scratchpad"); - size = it->second.size(); - } - flush_cpu_cache(va, size); - } - // Run one dispatch of `program` over num_tensors (host_ptr, size) pairs. // Serialized against every other dispatch (single shared queue). void dispatch(triton_npu_hsa_program *program, std::uint32_t num_tensors, @@ -936,7 +879,7 @@ class HsaRuntime { // cannot be reliably re-reserved+mapped, and reuse removes per-launch cost). std::unordered_map> vmem_pool_; - // Prepared programs, keyed by "hsaco_path\0kernel_name"; owns the loaded + // Prepared programs, keyed by path, kernel and variables; owns the loaded // executable and makes prepare() idempotent (a repeated (hsaco_path, // kernel_name) reuses the same loaded executable rather than loading/ // leaking a second copy). @@ -945,11 +888,6 @@ class HsaRuntime { // destructor freeing them all. Unlike regions_ below, which cannot be. std::unordered_map> programs_; - // Programs from prepare_bytes(), which has no key to look one up by (each - // is loaded from an hsaco meant for a single use). Owns them so they - // outlive the call, same as programs_ for the path-keyed ones; the one - // iteration is the destructor freeing them all, like programs_. - std::vector> bytes_programs_; // Control scratchpads by host address (see scratchpad_alloc). Guarded by // programs_mtx_, since each one belongs to a program. std::unordered_map scratchpads_; @@ -1100,7 +1038,6 @@ class HsaRuntime { vmem_free(b); vmem_pool_.clear(); programs_.clear(); - bytes_programs_.clear(); scratchpads_.clear(); kernarg_buffer_ = PoolBuffer{}; if (signal_.handle) { @@ -1366,21 +1303,6 @@ class HsaRuntime { return ((n + data_granule_ - 1) / data_granule_) * data_granule_; } - // The address the NPU sees for `ptr`, which differs from the host address. - // - // The control code names every buffer by this address, and a host address in - // a patch site does not fault (the design waits for data that never arrives), - // so a buffer with no device address is refused here. - std::uint64_t device_address(void *ptr) { - const std::uint64_t dev = query_device_address(ptr); - if (dev == 0) - throw std::runtime_error( - "this ROCR does not report a device address for an AIE allocation, " - "so a full-ELF design's scratchpad address cannot be resolved; " - "dispatching it would hang"); - return dev; - } - // A fresh dev-pool allocation of `size` bytes. PoolBuffer alloc_dev(std::size_t size) { if (size == 0) @@ -1629,45 +1551,33 @@ extern "C" int triton_npu_hsa_agent_name(char *buf, size_t buf_len, extern "C" triton_npu_hsa_program_t triton_npu_hsa_prepare(const char *hsaco_path, const char *kernel_name, char *errbuf, size_t errbuf_len) { - try { - return runtime().prepare(hsaco_path, kernel_name); - } catch (const std::exception &e) { - write_err(errbuf, errbuf_len, "HSA prepare failed", e.what()); - return nullptr; - } catch (...) { - write_err(errbuf, errbuf_len, "HSA prepare failed", "unknown error"); - return nullptr; - } + return triton_npu_hsa_prepare_with_variables( + hsaco_path, kernel_name, 0, nullptr, nullptr, errbuf, errbuf_len); } -extern "C" triton_npu_hsa_program_t -triton_npu_hsa_prepare_bytes(const void *hsaco_data, uint64_t hsaco_size, - const char *kernel_name, char *errbuf, - size_t errbuf_len) { +extern "C" triton_npu_hsa_program_t triton_npu_hsa_prepare_with_variables( + const char *hsaco_path, const char *kernel_name, uint32_t num_variables, + const char *const *names, void *const *addresses, char *errbuf, + size_t errbuf_len) { try { - if (!hsaco_data || !kernel_name) - throw std::runtime_error("null hsaco_data or kernel_name"); - const auto *bytes = static_cast(hsaco_data); - return runtime().prepare_bytes(std::vector(bytes, bytes + hsaco_size), - kernel_name); + std::vector variables; + variables.reserve(num_variables); + for (uint32_t i = 0; i < num_variables; ++i) + variables.push_back({names[i], addresses[i]}); + return runtime().prepare(hsaco_path, kernel_name, variables); } catch (const std::exception &e) { - write_err(errbuf, errbuf_len, "HSA prepare (bytes) failed", e.what()); + write_err(errbuf, errbuf_len, "HSA prepare failed", e.what()); return nullptr; } catch (...) { - write_err(errbuf, errbuf_len, "HSA prepare (bytes) failed", - "unknown error"); + write_err(errbuf, errbuf_len, "HSA prepare failed", "unknown error"); return nullptr; } } -extern "C" void *triton_npu_hsa_scratchpad_alloc(uint64_t size, - uint64_t *dev_addr, - char *errbuf, +extern "C" void *triton_npu_hsa_scratchpad_alloc(uint64_t size, char *errbuf, size_t errbuf_len) { try { - if (!dev_addr) - throw std::runtime_error("null dev_addr"); - return runtime().scratchpad_alloc(static_cast(size), dev_addr); + return runtime().scratchpad_alloc(static_cast(size)); } catch (const std::exception &e) { write_err(errbuf, errbuf_len, "HSA scratchpad alloc failed", e.what()); return nullptr; @@ -1678,21 +1588,6 @@ extern "C" void *triton_npu_hsa_scratchpad_alloc(uint64_t size, } } -extern "C" int triton_npu_hsa_scratchpad_flush(void *va, char *errbuf, - size_t errbuf_len) { - try { - runtime().scratchpad_flush(va); - return 0; - } catch (const std::exception &e) { - write_err(errbuf, errbuf_len, "HSA scratchpad flush failed", e.what()); - return -1; - } catch (...) { - write_err(errbuf, errbuf_len, "HSA scratchpad flush failed", - "unknown error"); - return -1; - } -} - extern "C" int triton_npu_hsa_dispatch(triton_npu_hsa_program_t program, uint32_t num_tensors, void *const *host_ptrs, @@ -1751,46 +1646,6 @@ extern "C" int triton_npu_hsa_shared_mark_dirty(void *va, char *errbuf, } } -extern "C" int triton_npu_hsa_scratchpad_supported(void) { - // No runtime() here, and nothing that can throw: a caller choosing which - // shape to build should not have to claim the device to find out. This takes - // hsa_init (refcounted) and allocates one page from the dev pool, but creates - // no queue and dispatches nothing. - // - // Whether ROCR reports a device address for AIE memory is not visible in any - // symbol or version, so the probe asks about a real allocation. A scratchpad - // needs a full ELF, which is aie2p-only, so another generation answers 0 - // without asking. - if (hsa_init() != HSA_STATUS_SUCCESS) - return 0; - - int supported = 0; - std::vector aies; - AgentSearch as_aie{HSA_DEVICE_TYPE_AIE, &aies}; - char name[64] = {}; - if (hsa_iterate_agents(collect_agents, &as_aie) == HSA_STATUS_SUCCESS && - !aies.empty() && - hsa_agent_get_info(aies.front(), HSA_AGENT_INFO_NAME, name) == - HSA_STATUS_SUCCESS && - std::strncmp(name, "aie2p", sizeof(name)) == 0) { - PoolSearch s{ - HSA_AMD_MEMORY_POOL_GLOBAL_FLAG_COARSE_GRAINED, false, {}, false}; - const hsa_status_t st = - hsa_amd_agent_iterate_memory_pools(aies.front(), find_pool, &s); - if ((st == HSA_STATUS_SUCCESS || st == HSA_STATUS_INFO_BREAK) && s.found) { - void *probe = nullptr; - if (hsa_amd_memory_pool_allocate(s.pool, 4096, 0, &probe) == - HSA_STATUS_SUCCESS) { - supported = query_device_address(probe) != 0 ? 1 : 0; - hsa_amd_memory_pool_free(probe); - } - } - } - - hsa_shut_down(); - return supported; -} - extern "C" void triton_npu_hsa_dispatch_counts(uint64_t *in_place, uint64_t *staged) { // No runtime() here, and so nothing that can throw: see g_in_place. diff --git a/amd_triton_npu/backend/include/HsaRuntime/HsaRuntime.h b/amd_triton_npu/backend/include/HsaRuntime/HsaRuntime.h index 9ded5979..a74d3937 100644 --- a/amd_triton_npu/backend/include/HsaRuntime/HsaRuntime.h +++ b/amd_triton_npu/backend/include/HsaRuntime/HsaRuntime.h @@ -55,58 +55,26 @@ triton_npu_hsa_program_t triton_npu_hsa_prepare(const char *hsaco_path, char *errbuf, size_t errbuf_len); -// Like triton_npu_hsa_prepare, but pack-loads an hsaco already in memory -// instead of reading one from a file. The bytes are copied before this call -// returns, so hsaco_data need not outlive it. -// -// Unlike triton_npu_hsa_prepare, not cached by path -- there is no path to key -// on. Meant for an hsaco that is good for only one load (e.g. one baked with a -// single scratchpad's device address, which no other hsaco shares); preparing -// the same bytes twice loads two independent programs rather than reusing one. -triton_npu_hsa_program_t triton_npu_hsa_prepare_bytes(const void *hsaco_data, - uint64_t hsaco_size, - const char *kernel_name, - char *errbuf, - size_t errbuf_len); +// Like triton_npu_hsa_prepare, but first defines num_variables agent +// variables, names[i] at addresses[i], with +// hsa_executable_agent_global_variable_define. A full ELF that references a +// buffer it does not carry, such as the control scratchpad +// ("scratch-pad-ctrl"), loads only if that buffer is defined. Cached by path, +// kernel and variables, so each scratchpad gets its own program. +triton_npu_hsa_program_t triton_npu_hsa_prepare_with_variables( + const char *hsaco_path, const char *kernel_name, uint32_t num_variables, + const char *const *names, void *const *addresses, char *errbuf, + size_t errbuf_len); // Initialize the runtime (idempotent) and allocate a zeroed control scratchpad -// of `size` bytes for a full-ELF design that declares parameters. Returns its -// host address and writes the address the NPU reaches it by to *dev_addr, or -// returns NULL on error (with a message written to errbuf). It lives as long -// as the runtime does, like the programs it is baked into. -// -// ROCR's hsaco loader patches only PDI and argument addresses into a full ELF's -// control code, and refuses one that relocates anything else. So the caller -// writes *dev_addr into the control code at the scratchpad's relocation sites -// itself, drops those relocations, and packs the result (see driver.py's -// _pack_hsaco). The scratchpad is device-visible memory the host writes and -// the device reads in its dispatch preamble; which uint32 a parameter takes and -// how it is encoded come from the build's params.txt, which this runtime does -// not read. -// -// Needs a ROCR whose hsa_amd_pointer_info reports an AIE allocation's device -// address; see triton_npu_hsa_scratchpad_supported. -void *triton_npu_hsa_scratchpad_alloc(uint64_t size, uint64_t *dev_addr, - char *errbuf, size_t errbuf_len); - -// Flush host writes to a scratchpad from the CPU cache, so that the next -// dispatch reads them. Returns 0 on success, or a negative value on error (with -// a message written to errbuf); `va` must be an address -// triton_npu_hsa_scratchpad_alloc returned. -// -// Nothing else does it: the packet lists only the kernel's arguments, which a -// full ELF has to match exactly, so the scratchpad cannot ride along as an -// extra argument for ROCR to flush. -int triton_npu_hsa_scratchpad_flush(void *va, char *errbuf, size_t errbuf_len); - -// Whether this ROCR can resolve a scratchpad's device address: 1 if it can, 0 -// if it cannot, including on an AIE agent other than aie2p. Never fails, and -// creates no queue. It allocates a page to ask, and so calls hsa_init (which is -// refcounted). +// of `size` bytes from the AIE dev pool. Returns its host address, or NULL on +// error (with a message written to errbuf). It lives as long as the runtime +// does. // -// Meant to be asked before choosing which artifact to build or load. A full -// ELF with no scratchpad dispatches without it. -int triton_npu_hsa_scratchpad_supported(void); +// ROCR flushes it on every dispatch, so the host writes parameters with plain +// stores; their slots and encoding come from the build's params.txt. +void *triton_npu_hsa_scratchpad_alloc(uint64_t size, char *errbuf, + size_t errbuf_len); // Dispatch a prepared program: acquire vmem I/O buffers, copy inputs in, fill // kernargs, enqueue the AIE packet, wait for completion, copy outputs back, and diff --git a/examples/llama32_1b_q4nx/README.md b/examples/llama32_1b_q4nx/README.md index 9fea2dd2..3b7aead0 100644 --- a/examples/llama32_1b_q4nx/README.md +++ b/examples/llama32_1b_q4nx/README.md @@ -228,17 +228,12 @@ scratchpad parameters, two scalars in device memory that the design reads in its dispatch preamble, so one full ELF serves every context length and nothing rewrites the instruction stream per token. -ROCR loads the ELF from an hsaco and patches only PDI and argument addresses -into its control code; it refuses one that relocates anything else. So -`../llm_q4nx/hsa_decode.py` allocates the scratchpad first, asks ROCR for its -device address (`hsa_amd_pointer_info`, which needs a ROCR that reports one for -an AIE allocation), writes that into the control code, and packs the result. -The host then writes L before each dispatch and flushes it from the CPU cache. +`../llm_q4nx/hsa_decode.py` allocates the scratchpad and defines it as the +ELF's `scratch-pad-ctrl` variable when the hsaco is loaded. ROCR flushes it on +every dispatch, so the host writes L before each dispatch with plain stores. `make compile-decode RUNTIME=hsa` builds the ELF plus the xclbin template pair -that mlir-air's decoder loads at construction. The patched-instruction-stream -alternative from before hsaco dispatch is gone: the stream now belongs to a -loaded ROCR executable, with no buffer left to patch per token. +that mlir-air's decoder loads at construction. `--ops` takes `all` or a comma list of `matmul,rms_norm,swiglu`, so a numerical regression can be bisected to a single kernel against the same CPU diff --git a/examples/llm_q4nx/hsa_decode.py b/examples/llm_q4nx/hsa_decode.py index 6151519e..d0d22f61 100644 --- a/examples/llm_q4nx/hsa_decode.py +++ b/examples/llm_q4nx/hsa_decode.py @@ -9,13 +9,9 @@ parameters of a full ELF: two scalars in device memory that the design reads in its dispatch preamble, so one ELF serves every L. -ROCR loads the ELF from an hsaco and patches only PDI and argument addresses -into its control code, so the scratchpad's address is written in before -packing (`driver._pack_hsaco`). That has to be the NPU's address for it, so -this needs a ROCR whose `hsa_amd_pointer_info` reports an AIE allocation's -device address; `use_scratchpad` asks. The instruction-stream alternative -- -patching the L-dependent words per token -- has no counterpart under hsaco -dispatch, where the stream belongs to a loaded executable. +The host allocates the scratchpad and defines it as the ELF's +`scratch-pad-ctrl` variable when the hsaco is loaded. ROCR flushes it on every +dispatch, so L is written with plain stores. Weights and the KV cache live in shared regions, so they are dispatched on in place. Copying 0.7 GB of weights per token would cap throughput near 23 tok/s @@ -23,6 +19,7 @@ """ import ctypes +import hashlib import os import re import threading @@ -126,82 +123,72 @@ def mark_resident(lib, arr): raise HsaDecodeError(buf.value.decode()) -def use_scratchpad(engine="fused_decode"): - """Whether the HSA decode can run: the scratchpad ELF is its only artifact. +#: The name aiecc gives a full ELF's control scratchpad. +SCRATCHPAD_SYMBOL = b"scratch-pad-ctrl" - Decided by the loaded ROCR: the ELF needs `hsa_amd_pointer_info` to report - the device address of an AIE allocation. Only the shared engine's decoder - is built as an ELF; the PLE engine has none, so it cannot run on HSA. - """ - if engine != "fused_decode": - return False - from triton.backends.amd_triton_npu.driver import load_hsa_runtime - return bool(load_hsa_runtime().triton_npu_hsa_scratchpad_supported()) +def _decode_hsaco(elf_path, kernel_name): + """The decode ELF packed into an hsaco, cached on the ELF's contents.""" + from triton.backends.amd_triton_npu.driver import _pack_hsaco + from triton.runtime.cache import get_cache_manager + + with open(elf_path, "rb") as f: + elf = f.read() + cache = get_cache_manager( + hashlib.sha256(elf + b"\0" + kernel_name.encode()).hexdigest() + ) + path = cache.get_file("decode.hsaco") + if path is None: + # Full ELF is aie2p only. + hsaco = _pack_hsaco("elf", {"elf_path": str(elf_path)}, "aie2p", kernel_name) + path = cache.put(hsaco, "decode.hsaco", binary=True) + return path class HsaElfProgram: """A full-ELF kernel and the scratchpad the host writes before each dispatch. - The scratchpad is allocated first, its device address written into the - ELF, and the result packed into an hsaco and loaded, so each instance is - its own program. Hold `lock` from writing the scratchpad through the - dispatch that reads it, and `flush` in between. + Each instance has its own scratchpad, and so its own loaded program. Hold + `lock` from writing the scratchpad through the dispatch that reads it. """ - def __init__(self, elf_path, kernel_name): - from triton.backends.amd_triton_npu.driver import ( - _pack_hsaco, - full_elf_scratchpad_size, - load_hsa_runtime, - ) + def __init__(self, elf_path, kernel_name, words): + from triton.backends.amd_triton_npu.driver import load_hsa_runtime self.lib = load_hsa_runtime() - size = full_elf_scratchpad_size(elf_path) - if not size: - raise HsaDecodeError( - f"'{kernel_name}' declares no scratchpad parameters, so there " - "is no way to give it a context length" - ) - dev_addr = ctypes.c_uint64() buf, n = _errbuf() self.lib.triton_npu_hsa_scratchpad_alloc.restype = ctypes.c_void_p va = self.lib.triton_npu_hsa_scratchpad_alloc( - ctypes.c_uint64(size), ctypes.byref(dev_addr), buf, n + ctypes.c_uint64(4 * words), buf, n ) if not va: raise HsaDecodeError(buf.value.decode()) - self._scratchpad = va - # A view of device memory, not a copy; `flush` makes writes visible. + # A view of the scratchpad, not a copy. self.params = np.ctypeslib.as_array( - ctypes.cast(va, ctypes.POINTER(ctypes.c_uint32)), (size // 4,) + ctypes.cast(va, ctypes.POINTER(ctypes.c_uint32)), (words,) ) - # Full ELF is aie2p only. This hsaco is good only for this scratchpad - # (it is baked with this one's device address), so prepare_bytes loads - # it straight from memory rather than caching it by a path nothing else - # would ever share. - hsaco = _pack_hsaco( - "elf", {"elf_path": str(elf_path)}, "aie2p", kernel_name, dev_addr.value - ) + names = (ctypes.c_char_p * 1)(SCRATCHPAD_SYMBOL) + addresses = (ctypes.c_void_p * 1)(va) buf, n = _errbuf() - self.lib.triton_npu_hsa_prepare_bytes.restype = ctypes.c_void_p - self.handle = self.lib.triton_npu_hsa_prepare_bytes( - hsaco, ctypes.c_uint64(len(hsaco)), kernel_name.encode(), buf, n + self.lib.triton_npu_hsa_prepare_with_variables.restype = ctypes.c_void_p + self.handle = self.lib.triton_npu_hsa_prepare_with_variables( + _decode_hsaco(elf_path, kernel_name).encode(), + kernel_name.encode(), + ctypes.c_uint32(1), + names, + addresses, + buf, + n, ) if not self.handle: - raise HsaDecodeError(buf.value.decode()) + raise HsaDecodeError( + f"{buf.value.decode()}\nThe decode needs a ROCR whose AIE loader " + "resolves a full ELF's control scratchpad; use one from " + "utils/rocm-requirements.txt" + ) self.lock = threading.Lock() - def flush(self): - """Make the host's writes to `params` visible to the next dispatch.""" - buf, n = _errbuf() - rc = self.lib.triton_npu_hsa_scratchpad_flush( - ctypes.c_void_p(self._scratchpad), buf, n - ) - if rc != 0: - raise HsaDecodeError(buf.value.decode()) - def dispatch(self, arrays): """One enqueue over `arrays` in kernel-argument order.""" n = len(arrays) @@ -221,7 +208,8 @@ def dispatch(self, arrays): class Params: """Where the two parameters live in the scratchpad, and what to write there.""" - def __init__(self, append_slot, scale, addend, mask_slot): + def __init__(self, words, append_slot, scale, addend, mask_slot): + self.words = words self.append_slot = append_slot self.scale = scale self.addend = addend @@ -265,7 +253,7 @@ def parse_params(path): raise HsaDecodeError(f"{path}: cannot read affine coefficients from {addr[0]}") scale = int(m.group(2)) addend = int(m.group(4)) * (-1 if m.group(3) == "m" else 1) - return Params(slot[addr[0]], scale, addend, slot[core[0]]) + return Params(int(lines[0][0]), slot[addr[0]], scale, addend, slot[core[0]]) def _elf_file(artifact_dir, name): @@ -325,7 +313,7 @@ def _elf_program(artifact_dir, attn_maxl, region_w, model=None): f"the decoder's REGION_W is {region_w}; the ELF and the decoder " "disagree" ) - return params, HsaElfProgram(elf, ELF_KERNEL_NAME) + return params, HsaElfProgram(elf, ELF_KERNEL_NAME, params.words) def require_xclbin_pair(artifact_dir): @@ -388,12 +376,6 @@ def _g(self, name): def __init__(self, *a, **kw): require_xclbin_pair(artifact_dir) - if not use_scratchpad(engine): - raise HsaDecodeError( - "this ROCR does not report a device address for an AIE " - "allocation, which the decode's scratchpad needs; use a " - "ROCR from utils/rocm-requirements.txt" - ) kw["max_L"] = elf_window(artifact_dir, kw.get("max_L")) super().__init__(*a, **kw) self._params, self._prog = _elf_program( @@ -546,7 +528,6 @@ def dispatch(self, tok, p): self._write_rms_rope(p) with self._prog.lock: self._params.write(self._prog.params, L) - self._prog.flush() # 8B split-weight tail follows the fixed five (base's w_bos[1:] # order). self._prog.dispatch( diff --git a/utils/rocm-requirements.txt b/utils/rocm-requirements.txt index 375b58fa..6b1b6cae 100644 --- a/utils/rocm-requirements.txt +++ b/utils/rocm-requirements.txt @@ -3,9 +3,9 @@ # `hsa` extra. Optional: the default XRT runtime needs none of it, so nothing # installs it unless asked. # -# The pin is a nightly that carries everything the HSA path uses: AIE -# dispatch, and an hsa_amd_pointer_info that reports an AIE allocation's -# device address, which the full-ELF decode needs. Without a pin pip takes -# whatever nightly is newest that day. +# The pin is a nightly that carries everything the HSA path uses: AIE hsaco +# loading and dispatch, and resolution of a full ELF's control scratchpad, +# which the decode needs. Without a pin pip takes whatever nightly is newest +# that day. --extra-index-url https://nightly.repo.amd.com/rocm/whl-next/ -rocm-sdk-core==10.2.0a20261020 +rocm-sdk-core==10.3.0a20261027