Repository navigation
[rocjitsu] Offload independent MMA work to VM-owned helper threads - #11767
Merged
Merged
Conversation
github-actions
Bot
added
documentation
❌ PR Check — Action Required
📖 Need help? See the Policy FAQ for details on every check and how to fix failures. |
|
🚫 Please fix the failed policies before requesting reviews. The following policy checks failed:
The |
Contributor
There was a problem hiding this comment.
🔵 Needs a closer look
One or more issues must be addressed before approval.
Pull request overview
Adds VM-owned helper threads for asynchronous MMA execution, with dependency tracking, plugin opt-in, configuration/checkpoint support, and extensive tests and documentation.
Changes:
- Adds async MMA scheduling, admission caching, helper pools, and plugin notifications.
- Extends thread allocation and Mirage configuration with helper counts.
- Adds coverage for concurrency, hazards, lifetimes, checkpoints, plugins, and policies.
File summaries
| File | Description |
|---|---|
| emulation/rocjitsu/tools/rocjitsu/main.cpp | Updated as part of this pull request. |
| emulation/rocjitsu/tests/partitioning_test.cpp | Updated as part of this pull request. |
| emulation/rocjitsu/tests/mma_test_util.h | Updated as part of this pull request. |
| emulation/rocjitsu/tests/matrix_coexecution_test.cpp | Updated as part of this pull request. |
| emulation/rocjitsu/tests/matrix_coexecution_lifetime_test.cpp | Updated as part of this pull request. |
| emulation/rocjitsu/tests/instruction_cache_test.cpp | Updated as part of this pull request. |
| emulation/rocjitsu/tests/execution_plugin_test.cpp | Updated as part of this pull request. |
| emulation/rocjitsu/tests/config_test.cpp | Updated as part of this pull request. |
| emulation/rocjitsu/tests/CMakeLists.txt | Updated as part of this pull request. |
| emulation/rocjitsu/tests/async_mma_test_util.h | Updated as part of this pull request. |
| emulation/rocjitsu/tests/async_mma_policy_test.cpp | Updated as part of this pull request. |
| emulation/rocjitsu/schemas/simulation_config.fbs | Updated as part of this pull request. |
| emulation/rocjitsu/lib/rocjitsu/src/rocjitsu/vm/rj_vm.cpp | Updated as part of this pull request. |
| emulation/rocjitsu/lib/rocjitsu/src/rocjitsu/vm/plugins/throughput/plugin.h | Updated as part of this pull request. |
| emulation/rocjitsu/lib/rocjitsu/src/rocjitsu/vm/plugins/throughput/plugin.cpp | Updated as part of this pull request. |
| emulation/rocjitsu/lib/rocjitsu/src/rocjitsu/vm/plugins/logging/plugin.h | Updated as part of this pull request. |
| emulation/rocjitsu/lib/rocjitsu/src/rocjitsu/vm/plugins/logging/plugin.cpp | Updated as part of this pull request. |
| emulation/rocjitsu/lib/rocjitsu/src/rocjitsu/vm/plugins/execution_plugin.h | Updated as part of this pull request. |
| emulation/rocjitsu/lib/rocjitsu/src/rocjitsu/vm/plugins/execution_plugin_group.h | Updated as part of this pull request. |
| emulation/rocjitsu/lib/rocjitsu/src/rocjitsu/vm/amdgpu/mma_admission.h | Updated as part of this pull request. |
| emulation/rocjitsu/lib/rocjitsu/src/rocjitsu/vm/amdgpu/matrix_coexecution.h | Updated as part of this pull request. |
| emulation/rocjitsu/lib/rocjitsu/src/rocjitsu/vm/amdgpu/instruction_cache.h | Updated as part of this pull request. |
| emulation/rocjitsu/lib/rocjitsu/src/rocjitsu/vm/amdgpu/compute_unit.h | Updated as part of this pull request. |
| emulation/rocjitsu/lib/rocjitsu/src/rocjitsu/vm/amdgpu/compute_unit.cpp | Updated as part of this pull request. |
| emulation/rocjitsu/lib/rocjitsu/src/rocjitsu/vm/amdgpu/async_scoreboard.h | Updated as part of this pull request. |
| emulation/rocjitsu/lib/rocjitsu/src/rocjitsu/isa/isa_traits.h | Updated as part of this pull request. |
| emulation/rocjitsu/lib/rocjitsu/src/rocjitsu/isa/arch/amdgpu/shared/rdna_isa_base.h | Updated as part of this pull request. |
| emulation/rocjitsu/lib/rocjitsu/src/rocjitsu/isa/arch/amdgpu/shared/cdna_isa_base.h | Updated as part of this pull request. |
| emulation/rocjitsu/lib/rocjitsu/src/rocjitsu/isa/arch/amdgpu/rdna4/isa.h | Updated as part of this pull request. |
| emulation/rocjitsu/lib/rocjitsu/src/rocjitsu/isa/arch/amdgpu/cdna5/isa.h | Updated as part of this pull request. |
| emulation/rocjitsu/lib/rocjitsu/src/rocjitsu/isa/arch/amdgpu/cdna4/isa.h | Updated as part of this pull request. |
| emulation/rocjitsu/lib/rocjitsu/src/rocjitsu/isa/arch/amdgpu/cdna3/isa.h | Updated as part of this pull request. |
| emulation/rocjitsu/lib/rocjitsu/src/rocjitsu/isa/arch/amdgpu/async_mma_policy.h | Updated as part of this pull request. |
| emulation/rocjitsu/lib/rocjitsu/src/rocjitsu/config/config_loader.h | Updated as part of this pull request. |
| emulation/rocjitsu/lib/rocjitsu/src/rocjitsu/config/config_loader.cpp | Updated as part of this pull request. |
| emulation/rocjitsu/lib/rocjitsu/src/rocjitsu/config/checkpoint.h | Updated as part of this pull request. |
| emulation/rocjitsu/lib/rocjitsu/src/rocjitsu/config/checkpoint.cpp | Updated as part of this pull request. |
| emulation/rocjitsu/docs/sphinx/how-to/configure-topology.md | Updated as part of this pull request. |
| emulation/rocjitsu/docs/sphinx/conceptual/json-configuration.md | Updated as part of this pull request. |
| emulation/rocjitsu/docs/plugins.md | Updated as part of this pull request. |
| emulation/rocjitsu/docs/mma-admission.md | Updated as part of this pull request. |
| emulation/rocjitsu/docs/configuration.md | Updated as part of this pull request. |
| emulation/rocjitsu/docs/concurrent-dispatch.md | Updated as part of this pull request. |
| emulation/rocjitsu/docs/async-instructions.md | Updated as part of this pull request. |
| emulation/rocjitsu/configs/gfx950_mi355x.json | Updated as part of this pull request. |
| emulation/rocjitsu/configs/gfx950_mi355x_kmd.json | Updated as part of this pull request. |
| emulation/rocjitsu/configs/gfx950_mi355x_kmd_2gpu.json | Updated as part of this pull request. |
| emulation/rocjitsu/configs/gfx942_cdna3.json | Updated as part of this pull request. |
| emulation/rocjitsu/configs/gfx942_cdna3_kmd.json | Updated as part of this pull request. |
| emulation/rocjitsu/configs/gfx90a_mi210_kmd.json | Updated as part of this pull request. |
| emulation/rocjitsu/configs/gfx1250_mi455x.json | Updated as part of this pull request. |
| emulation/rocjitsu/configs/gfx1250_mi455x_kmd_4gpu.json | Updated as part of this pull request. |
| emulation/rocjitsu/configs/gfx1201_r9700.json | Updated as part of this pull request. |
| emulation/rocjitsu/configs/gfx1151.json | Updated as part of this pull request. |
| emulation/rocjitsu/configs/gfx1100_w7900.json | Updated as part of this pull request. |
| emulation/mirage/tests/e2e.rs | Updated as part of this pull request. |
| emulation/mirage/rocjitsu/src/lib.rs | Updated as part of this pull request. |
| emulation/mirage/rocjitsu/build.rs | Updated as part of this pull request. |
Review details
Suppressed comments (2)
emulation/rocjitsu/lib/rocjitsu/src/rocjitsu/config/config_loader.cpp:86
- This condition still applies the automatic budget to every candidate, so explicit execution-thread requests are rejected whenever their effective cost exceeds the host budget. For example, the updated
ExplicitOverridesAndSafeFallbacktest requests E=8/D=33/H=8 with host_threads=4, giving cost 48, and this code falls back to the default 1/1/0 plan instead of the expected explicit plan. Allow an explicit E, D, or helper request to bypass thecost <= budgetcheck (while retaining the automatic CPU-cap check for unspecified E/D).
const uint64_t cost = cpu_cost + candidate.helpers;
if (cost <= budget && cost >= best_cost) {
emulation/rocjitsu/lib/rocjitsu/src/rocjitsu/isa/arch/amdgpu/cdna3/isa.h:30
HasAsyncMmais the compile-time capability used byIsaExecComputeUnit::supports_async_execution, but this marks CDNA3 as supported even thoughasync_mma_policy::supported()explicitly excludes CDNA3 and the production factory therefore never enables the adapter for it. Any caller relying on the public trait can select async execution for CDNA3 while the runtime silently uses the synchronous path. Keep the trait at its base value for CDNA3 (or update the policy/factory so the capability and runtime selection agree).
static constexpr uint32_t ASYNC_MMA_WAVE_SIZE = 64; ///< Async MMA adapter wave size.
- Files reviewed: 58/58 changed files
- Comments generated: 0
- Review effort level: Lite
💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
kuhar
force-pushed
the
users/kuhar/async-mma
branch
from
September 17, 2026 16:01
7e20e76 to
838c90d
Compare
Adding 16 async helpers to the existing 32 CPU execution threads improves IREE MXFP4 dispatch throughput by 108% (73% by process wall time), AITER MXFP4 by 32%, and small-grid Gluon MXFP4 by 33%. Helpers overlap independent MMA work within active waves when few workgroups limit CU parallelism. ## Plugin and architectural impact - The ExecutionPlugin API/ABI changes: rebuild plugins with the host. Async execution requires every plugin to opt in via supports_async_instructions(). Architectural observers, including ConSan, keep synchronous execution by default. - Offloaded instructions replace before/after callbacks with an issuer-side, metadata-only issue notification. Register hooks may run concurrently within the same wave and precede that notification. Plugins cannot inspect architectural snapshots or retain instruction/wave references; serializing callbacks alone does not satisfy this contract. - Throughput and kernel logging opt in. Instruction counts and dispatch wall time remain valid; throughput reports overlapping execution-duration metrics as unavailable. - This changes functional host execution only. Helpers access the wave's live registers under a dependency scoreboard and cached independence lookahead. Jobs finish before CU rescheduling or error propagation; precise async exception rollback is unsupported. Debugging, clocked mode and unsupported wave states retain synchronous execution. ## Performance Final implementation: 8/25/0 versus 8/25/16 (32 versus 48 execution threads). Speedup = (baseline time / helper time - 1) * 100%; 0% is parity. | Target / workload | Dispatch speedup | Process-wall speedup | |---|---:|---:| | gfx950 / IREE MXFP4 (16 workgroups) | +107.8% | +73.0% | | gfx950 / AITER MXFP4 GEMM | +31.7% | +1.9% | | gfx1250 / Gluon MXFP4 (16 workgroups) | +33.2% | +1.1% | | gfx950 / TensileLite FP8 GEMM | +13.9% | +1.2% | | gfx1250 / Gluon FP16 K32 | +3.4% | +0.9% | | gfx950 / AITER FP8 attention | +1.7% | -0.1% | | gfx1250 / AITER FP8 attention | +0.7% | -0.8% | Medians of four fresh-process samples per configuration on 64 reserved Threadripper PRO 9995WX cores, Clang 23 -O2, with warmed compiler caches. The two budgets were measured in separate blocks; small differences are inconclusive. Process wall includes startup and teardown. An earlier equal-48-thread comparison used 8/41/0 versus 8/25/16, with dispatch limits lifted only in the local control build. Helper dispatch gains were 106% for IREE MXFP4, 26% for AITER MXFP4 and 31% for small-grid Gluon MXFP4. TensileLite FP8 tied; large Gluon FP16 and gfx1250 FP8 attention favored extra dispatch workers (helper gains -6% and -11%). These results predate the final simplification and show the workload-dependent tradeoff. ## Thread policy and ownership Expose async_helper_threads in JSON and Mirage. Preserve engine/dispatch allocations through 32 threads; single-GPU gfx950/gfx1250 presets add 2/4/8/16 helpers at affinity budgets 34/36/40/48. Their triples are 8/25/H: E/D/H means engines / inclusive dispatch width / helpers, with total execution threads = E + D - 1 + H. Automatic selection stops at 48; other targets and multi-GPU presets keep zero helpers. Explicit JSON settings can override the policy. No async environment controls. Own the lazy helper pool in the VM, with nonblocking inline fallback and bounded spinning followed by atomic wait/notify. Cache accepted and rejected lookahead plans, revalidate code bytes after I-cache invalidation, and destroy temporary decodes on the issuer. Preserve floating-point state, fork fallback and checkpoint compatibility. Performance probes remain local.
kuhar
force-pushed
the
users/kuhar/async-mma
branch
from
September 18, 2026 00:58
838c90d to
37001b8
Compare
amd-arosa
added a commit
that referenced
this pull request
Sep 18, 2026
One conflict, in tests/execution_plugin_test.cpp, in two places: * The plugin include block -- this branch adds instruction_mix/plugin.h where develop's MMA helper-thread work (#11767) adds logging/plugin.h. Both kept. * A ~650-line region where both sides append a block of tests after the throughput cases: this branch's instruction-mix helpers and tests, develop's AsyncEventPlugin coverage for VM-owned helper threads. The two blocks are independent -- the collision is purely positional -- so both are kept, in that order, with the brace that previously closed whichever block came last now closing develop's. Verified on the merged tree: 4731/4731 ctest, clean build with no new diagnostics, and 58/58 across InstructionMixPluginTest, ThroughputPluginTest and develop's new async/coexecution cases. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Contributor
TheRock Submodule Bump ActivityNewest first
|
IanNod
added a commit
that referenced
this pull request
Sep 24, 2026
Develop landed #11767 on emulation/mirage while this was in flight, and the rebase merged its three files into cli/ cleanly. What it did not do is reconcile what those files now mean. Helper threads are a third dimension of an allocation, not a detail: ExecutionThreadAllocation gained `helpers`, and this PR's own C entry point -- which develop does not have, so nothing warned -- reported only engines and dispatch. The table would have gone on rendering, quietly short by a column. rj_config_resolve_execution_threads returns the count now, ThreadAllocation carries it, the total includes it, and the row list gained the additive helper budgets 34, 36 and 40 that only exist because engine and dispatch cost is capped at 32. The e2e case that pins those rows to configuration.md is re-pinned to develop's rewritten table rather than relaxed. It exists to stop the allocator and the documentation drifting apart, and the two agreeing on new numbers is the test doing its job. Develop also grew two VFIO capabilities on the C++ CLI this PR deletes, and `scripts/run-vfio-guest.py` uses both: `--check-vfio-user` to decide whether the build can serve at all, and `--vfio-ready-fd` to wait for the socket instead of polling for the file. Both are carried over -- check-vfio-user as a subcommand with the bare flag routed to it, the way --thread-budget-table already is, and the ready descriptor threaded through rj_run_vfio_server. The ctest that probed the old binary runs the CLI instead. The schema moved too: seven KfdDeviceInfo fields are optional now, and `= null` is not a value core/build.rs can emit. An optional scalar has no default, which is the one case this generator has nothing to say about, so it publishes zero. Develop's own ServerSignalMask replaces the guard round three added. Same fix, better placed -- it owns the signal set and supports an early restore the server actually uses. Co-authored-by: Cursor <cursoragent@cursor.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Adding 16 async helpers to the existing 32 CPU execution threads improves IREE MXFP4 dispatch throughput by 108% (73% by process wall time), AITER MXFP4 by 32%, and small-grid Gluon MXFP4 by 33%. Helpers overlap independent MMA work within active waves when few workgroups limit CU parallelism.
Original RFC: https://gist.github.com/kuhar/dff788b1ef24a9584cf98a35a9890fff
Plugin and architectural impact
Performance
Final implementation: 8/25/0 versus 8/25/16 (32 versus 48 execution threads). Speedup = (baseline time / helper time - 1) * 100%; 0% is parity.
Medians of four fresh-process samples per configuration on 64 reserved Threadripper PRO 9995WX cores, Clang 23 -O2, with warmed compiler caches. The two budgets were measured in separate blocks; small differences are inconclusive. Process wall includes startup and teardown.
An earlier equal-48-thread comparison used 8/41/0 versus 8/25/16, with dispatch limits lifted only in the local control build. Helper dispatch gains were 106% for IREE MXFP4, 26% for AITER MXFP4 and 31% for small-grid Gluon MXFP4. TensileLite FP8 tied; large Gluon FP16 and gfx1250 FP8 attention favored extra dispatch workers (helper gains -6% and -11%). These results predate the final simplification and show the workload-dependent tradeoff.
Thread policy and ownership
Expose async_helper_threads in JSON and Mirage. Preserve engine/dispatch allocations through 32 threads; single-GPU gfx950/gfx1250 presets add 2/4/8/16 helpers at affinity budgets 34/36/40/48. Their triples are 8/25/H: E/D/H means engines / inclusive dispatch width / helpers, with total execution threads = E + D - 1 + H. Automatic selection stops at 48; other targets and multi-GPU presets keep zero helpers. Explicit JSON settings can override the policy. No async environment controls.
Own the lazy helper pool in the VM, with nonblocking inline fallback and bounded spinning followed by atomic wait/notify. Cache accepted and rejected lookahead plans, revalidate code bytes after I-cache invalidation, and destroy temporary decodes on the issuer. Preserve floating-point state, fork fallback and checkpoint compatibility. Performance probes remain local.