Add MIGraphX backend for AMD GPUs (ROCm) — 2.94x over OpenCL on MI300X - #1235
Add MIGraphX backend for AMD GPUs (ROCm) — 2.94x over OpenCL on MI300X#1235zhihuidu-amd wants to merge 1 commit into
Conversation
Adds a third GPU backend targeting AMD via MIGraphX, ROCm's graph compiler, alongside the existing CUDA/TensorRT and OpenCL paths. The backend reuses the ONNX ModelProto that OnnxModelBuilder already emits for the TensorRT path and hands the identical bytes to MIGraphX's parse_onnx_buffer, so network construction is shared and onnxmodelbuilder.cpp is untouched. Measured on MI300X (gfx942), ROCm 7.2.0, b18c384nbt, 19x19, FP16, with both backends built and benchmarked in a single job on one node: visits=3200 OpenCL (tuned) 1564.38 MIGraphX 4599.00 2.94x visits=800 OpenCL (tuned) 1548.94 MIGraphX 4239.44 2.74x KataGo's OpenCL tuner reports canUseFP16TensorCores=0 on gfx942, so the OpenCL path never issues MFMA; MIGraphX routes convolutions through rocMLIR/MIOpen, which do. Validated with runnnonmanyposestest over 254 positions across all 5 nets in cpp/tests/models: FP32 agrees with OpenCL to 2.6e-11..5.6e-10 policyProbSquerr, and MIGraphX's FP16 is 2.2x-51x closer to the FP32 reference than OpenCL's FP16. Notes: - MIGraphX compiles one static shape, so the program is compiled at maxBatchSize and short batches are zero-padded. Padding rows get an all-ones mask, since the graph divides by maskSum for masked means and a zero mask row is a division by zero that propagates NaN into real rows. - Graph outputs are exposed as positional main:#output_N parameters while inputs keep their ONNX names; the mapping is asserted against declared shapes so an emitter reordering fails loudly rather than silently swapping tensors. - migraphxTransformerNHWC defaults to false, unlike TensorRT's trtTransformerNHWC. The channel-last trunk produces wrong policy output on transformer nets under MIGraphX (policySqErr 136 vs 6e-10) while value heads stay correct; root cause is still open, so the safe NCHW default ships. - Protobuf must be linked statically with -Wl,--exclude-libs,ALL, because libmigraphx_onnx exports its bundled protobuf as weak symbols that a shared libprotobuf would preempt. Documented in Compiling.md.
|
Can you test #1234 and see how it compares? |
|
Thanks for pointing me at #1234 — I built it and benchmarked it against mine on the same SetupBoth backends built and benchmarked inside a single job on one node, so nothing A note on the hardware: these absolutes are MI325X, while my PR description quotes ResultsFive interleaved trials per point (yours, mine, yours, mine, ...) so node drift hits both
Both CIs exclude 1.0, so this is a real difference rather than noise: yours is 8.8% One secondary observation, offered as a property rather than a consolation: my run-to-run Measurement note: warm-up matters for the MIOpen pathWorth flagging in case you benchmark this yourself. Your
A single-pass sweep therefore understates the MIOpen path badly in whichever cell runs That asymmetry is also a small real difference in its own right: MIGraphX compiles the Why the three are so closeWorth spelling out, since three AMD paths sounds like three different engines and it
I dumped my own compiled graph ( Both end up on the same MFMA units, but not via the same kernels. Profiling mine shows So I can't explain the 9-11% yet. It isn't layout, and it isn't an algorithm choice I can For completeness on the AMD ceiling: I instrumented my eval loop and it's GPU-bound — How the two approaches differThroughput favours yours, as above. The structural difference is the same split KataGo
#1234 adds the AMD analogue of Concretely, what the graph-compiler path buys:
The flip side, stated plainly: the per-op path has more headroom. A human can always The two findings below apply to the ONNX path as well, so they're worth having on record Two findings worth keeping either way1. Channel-last trunk is wrong for transformer nets under MIGraphX. With
Every position is affected and the logits collapse toward flat, so it's a wrong 2. Zero-padding short batches needs an all-ones mask, not zeros. The emitted graph Validation
Happy to run anything else you'd find useful — I have access to MI300X, MI325X and |
@zhihuidu-amd Thanks for the extremely careful benchmark - interleaved trials with CIs and the warm-up analysis is exactly how this should be measured. A few responses: On the cause of the gap. My "multiple static .mxr models" explanation in #1188 came from my own older MIGraphX branch, where I compiled one model per shape and paid for it. Your single-static-shape + all-ones-mask padding design avoids that problem entirely, so I agree the remaining 9-11% has a different cause. Your im2col suspect looks plausible to me: MIOpen's Winograd/implicit-GEMM kernels consume the tensor in place, and a ~9.4% Im2d2Col share matches the size of the gap almost suspiciously well. If you ever want to confirm it, profiling one conv where rocMLIR picks im2col+GEMM against MIOpen's Winograd on the same shape should settle it - but I agree it's not fixable from outside MIGraphX. On the transformer NHWC policy bug. A useful data point for triangulation: TensorRT's NHWC transformer path works, and my ROCm backend also computes transformers channel-last (BSHD attention) with correct policy output. Two independent NHWC consumers producing correct policy strongly suggests the bug is inside MIGraphX's NHWC lowering rather than in the emitted ONNX graph. On keep-vs-drop. Your structural point is fair, and I'll soften what I said in #1188: a ~900-line backend that reuses onnxmodelbuilder and inherits compiler improvements for free is a much smaller maintenance burden than I assumed. Whether KataGo carries both is @lightvector's call. My remaining concerns about the graph path are practical rather than ideological:
Either way, your two findings (the NHWC policy bug and the all-ones mask padding) are worth having on record - thanks for the measurement effort. |
|
@Looong01 Four things: a dead hypothesis I should have retired sooner, the low-visits data The im2col suspect is deadYou called it plausible, but I'd already tested it by the time you replied and it doesn't
Nothing moves the im2col share — it sits at 9.26–9.39% in every configuration. NHWC is You suggested profiling one conv where rocMLIR picks im2col+GEMM against MIOpen's Winograd on Also, you were right about the warm-up point: KataGo front-loads MIOpen compilation into Low visits: I was wrong about this, and the correction favours my backendI need to retract something before I show the data. In an earlier sweep I saw my backend at
reallocateNNEvalWithEnoughBatchSize(maxThreads); // batchSizeLimit = maxNumThreadsSo I measured it rather than just asserting it — same point, same job, only the compiled shape
Your per-op path resizes descriptors per call and is essentially immune (4%); mine pays Rerun properly — each thread count in its own process, so the compiled shape matches the
Two things fall out. First, the same Second, the trend runs opposite to the one you predicted. At 100 visits my backend is 39–47% That looked at the time like "padding waste is real but outweighed by something the graph Your concern was well-posed and worth testing properly, and I'd have gone on believing the Two things I'll state rather than let someone find them. The harness artefact isn't specific to my backend. It cuts against fixed-shape backends The 9–11% deficit was self-inflicted, and it's fixedI said above that the oversized compiled shape might explain the gap at high thread counts. Two changes, measured separately. 1. The batch cap. Sweeping
Your backend varies 1.14x across the whole sweep; mine varies 1.68x. That asymmetry is the 2. Batch bucketing in the backend. Since users shouldn't have to tune a config value to
To be careful about which comparison is which: 1.369x is my own before/after at the default Bucketing is numerically neutral — ON-vs-OFF output differences (mean 1.29e-07) fall inside The two overlap rather than compose: once the cap is sensible, bucketing's marginal gain Worth stating plainly: the 9–11% you were ahead by was my bug, not a property of the On the NHWC triangulation — it's my bug, not MIGraphX'sYour triangulation was useful, and it turned out to point the other way. Rather than pass it Both pass. MIGraphX computes the channel-last graph correctly. So "MIGraphX's NHWC lowering is at fault" is disproved, and the policySqErr 136 has to come An NVIDIA data point, for calibrationSince the question underneath all of this is "does the graph-compiler approach cost
On NVIDIA the graph-compiler backend is 27% faster than the hand-written one — the same (Hardware isn't matched — H100 80GB vs MI325X, different bandwidth and FP16 peak. We don't Where that leaves thingsYour Windows-support concern stands and I have no answer to it. Performance has changed since my last comment, so to be explicit about what I'm now
Caveat I'd rather state than have found: the low-visits grid predates the bucketing change, Thank you for pushing on the low-visits case. Without it I'd have published a number that |
|
@Looong01 Two follow-ups: the low-visits grid you asked for, re-measured now that bucketing The low-visits grid, redone — and the old one was wrong in my favourLast time I flagged that the 100–400 visit numbers predated the bucketing fix and deserved a This time I ran three arms interleaved in one job on one node, each thread count in its
The
Please treat this table as replacing the earlier one, not extending it. Your backend Bucketing's benefit tracks thread count, not visits — about 1.5% at t=16, 16% at t=32, So the answer to what you actually asked is yes, static-shape padding was hurting these Caveat: single trial per cell. The direction is consistent across all 12 and the bucketing Correction: the im2col hypothesis is dead, and I said so too weaklyIn my last comment I wrote that the im2col suspect "doesn't hold". I want to state the
So the hypothesis is disproved rather than merely unsupported. What I can show is that the Separately, while chasing the padding problem I did find a genuine MIGraphX bug — |
@zhihuidu-amd Thanks for running this properly - the three-arm interleaved design with the ours-off control arm is exactly right, and I accept the new table as replacing the old one. For the record: the drift you flagged (1481 vs 1291 etc.) is real and worth its own caveat on any cross-job comparison; your handling of it is more honest than most benchmark reports I've seen. I'll also plainly retract what I said in #1188 about dropping MIGraphX. The bucketing data is convincing: with bucketing on, you win or tie every cell, and the 2x at 100v/64t is not noise-sized. My earlier "padding hurts at low visits" prediction was right about the mechanism existing but wrong about the axis (thread count, not visits) and wrong about the conclusion once bucketing exists. And on the im2col question you were right to close it - I called it plausible, your knob sweep disproved it. Noted and accepted. The coherent picture as I now read it: MIOpen's per-op path wins when GPU-bound (big batches, high threads - your first 160/192-thread finding), the graph-compiler path wins when dispatch-bound (small batches, low-to-mid threads). That's exactly the CUDA vs TensorRT split on the NVIDIA side, and I no longer object to carrying both. One important scope note, though: everything we've both measured so far is on b18c384nbt, a pure convnet. That model form is legacy as of v1.17 - the main training run is switching to transformers, and the three new strong nets (b10c384h6nbttflrs, b10c512h8nbt3tflrs-fson-silu-rsnh, b11c768h12nbt3tflrs-fson-silu) are all transformers. The transformer comparison between our two backends is still a blank page, and right now it isn't a level comparison:
So if you're willing, the most informative next measurement would be one of the new transformer nets (say b10c384h6nbttflrs for speed, or b11c768h12nbt3tflrs-fson-silu for the flagship) - same interleaved setup, plus a correctness spot-check of policy output against the Eigen or CUDA reference, since on transformers a backend can be fast and quietly wrong (we just caught exactly that class of bug on our own older ROCm branch the hard way). If the NHWC policy bug turns out to be in MIGraphX's lowering, it may be worth filing alongside your Softplus/Softsign find. And thanks again for the rigor here - this is how backend comparisons should be done. |
Hii, thank u and good to see u. Do u have any suggestions in detail? Like what backends? |
|
@Looong01 Ran the transformer comparison you asked for. Correctness first, since your warning MI325X (gfx942), ROCm 7.2.0, Caveat up front, because it shapes everything below: my arm is on the NCHW fallback. Correctness
Two independent implementations — a graph compiled by MIGraphX versus your hand-written One result goes against me and I'd rather flag it than have you find it: in FP16 my error Comparing both GPU backends against my FP32 only proves we agree, not that we're both right,
Both FP32 paths land within ~1e-9 of a CPU implementation that shares no code with either. The FP16 gap survives the independent reference unchanged, so it is real and mine, not an ThroughputSame protocol as the low-visits grid: each thread count in its own process, warm-up discarded
What I read from itYour dispatch-bound/GPU-bound synthesis holds, and the transformer shifts the whole curve. I'll say the part that's awkward for my own framing: I expected your fused-attention path to The flagship net, where it goes the other wayI also ran
Side by side at the corners: 2.31x → 1.68x at 400/16, and 1.29x → 0.97x at 3200/128. Two caveats, one in each direction. At 0.97x/1.00x the gap is inside the run-to-run noise Correctness on the flagship matches the small net: #1234 FP32 vs mine at 1.58e-10, and my Limits worth stating: single trial per cell; one machine, one ROCm version, one net per size On the NHWC bug: |
|
@Looong01 Prefer to merge this PR now and put follow-ups in smaller PRs rather than expanding this one: https://testing.googleblog.com/2024/07/in-praise-of-small-pull-requests.html OK to merge? |
|
@lightvector Could you review this PR? Thank you. |
|
We can run KataGo on AMD GPUs in a few different ways, similar to how NVIDIA supports OpenCL, cuDNN, and TensorRT. Supporting multiple options keeps the door open while we keep improving performance. Users can pick the one that fits their workload and GPU version. |
|
I'm reluctant to merge yet another backend. KataGo already has seven backends. This is already too much maintenance burden. Do you have suggestions on how to deal with this? For example, does this backend 100% obsolete the rocm backend that we just merged, such that we can delete it? @Looong01 - curious for your thoughts as well, given that it was your rocm backend that was just merged, and earlier you had worked on a migraphx backend that I had reviewed that you had claimed was significantly worse which was why we went with rocm instead. |
|
@lightvector @Looong01 Fair question on maintenance — we would not frame this as “replace ROCm.” The benchmark thread shows MIGraphX and ROCm trade off. The two paths look complementary, similar to TensorRT vs CUDA. At AMD, we support multiple paths on ROCm depending on customer requests.
We’d suggest merging this PR, keeping ROCm, and putting performance improvement follow-ups in small PRs with clear docs on when to use each recipe. Happy to help with validation and maintenance from AMD’s side. |
|
@yaoliu13 Thanks. Can you explain more about the difference between this backend and #1222 which was merged just recently? If this backend is using ONNX anyways, and #1222 has migraphx as a provider for onnx's runtime, what is it doing differently than this PR? Is this PR more efficient somehow and what is the reason for the difference? |
|
I built the MIGraphX backend branch from https://github.com/zhihuidu-amd/KataGo/tree/migraphx-backend locally and ran the same four release models through both the MIGraphX backend and the ROCm backend on my RX 7900 XTX (gfx1100 ROCm 7.14, MIOpen 3.5.2. The GTP protocol works on MIGraphX (first move generated correctly), and the benchmark numbers below are the measured end-to-end visits/s from Comparison table
What the numbers show
Important caveats about this MIGraphX buildThis is not the possible MIGraphX configuration. My local ROCm 7.14 apt repository does not provide the
I also had to apply small local patches to make MIGraphX compile at all in this configuration: stubs for So the MIGraphX arm here is effectively IOpen for convolutions + rocBLAS for GEMMs, no rocMLIR fusion, no CK. That is a lower-bound performance for MIGraphX; the published numbers from your setup may be noticeably better if you have rocMLIR working. Even with that caveat, the ROCm backend is still faster on the same hardware, and the thread-scaling pattern is striking. Bottom lineThe MIGraphX backend compiles and runs correctly, but my RDNA3 machine it does not match the ROCm backend's throughput or scaling. The difference is modest on the heaviest network (b40) and large on the smaller transformer networks where overhead dominates. I think the ROCm backend should remain the primary AMD GPU path in KataGo, and MIGraphX is better kept as a separate experimental branch until it can match this performance or provide a clear feature advantage. |
Follow-up: testing the "dispatch-bound" hypothesis with hipGraph captureI want to refine one point from my earlier numbers, because your transformer comment included a specific hypothesis:
I ran a direct experiment to test this: capture the whole neural-net forward pass into a single HIP graph and replay it, collapsing the 100+ individual HIP dispatches down to one graph launch. This is the cleanest way to isolate dispatch overhead from kernel execution time. On gfx1100 (RX 7900 XTX) the result is neutral within ~3% run-to-run noise. In other words, on this RDNA3 card, per-op dispatch overhead is not the bottleneck. If CK FMHA were a real per-block win but being hidden by dispatch, hipGraph would have exposed it; it did not. So the MIGraphX 2× advantage on gfx942 (MI325X) is probably not coming from "fewer dispatches" in the graph-compiler sense. A sharper explanation is that MIGraphX's op fusion reduces the number of kernels and the amount of device-memory traffic, and at small batch sizes those kernels are tiny and memory-latency-dominated. That is a different bottleneck than dispatch count. Architecture asymmetry caveatThis experiment only holds for RDNA3. CDNA (gfx942) has a different driver/firmware dispatch path, and dispatch overhead may genuinely be more expensive there. I do not have a CDNA machine, so I cannot run the same hipGraph check on MI325X. Someone with CDNA access would need to repeat the capture test to settle whether the two architectures have different bottlenecks. Scorecard summary after the transformer run
My reading is unchanged: the two backends win in different regions, and neither is universally better. The ROCm backend should stay as the default AMD path because it is already faster on the same hardware in most KataGo use cases and matches or exceeds MIGraphX on the largest models at high batch sizes. |
|
On my side, as my result: Maybe MIGraphX backend runs faster than ROCm backend with CK on CDNA machine, but I only have RDNA machine so I can NOT prove @zhihuidu-amd's result on my side. Maybe anyone who have CDNA resources can help. In fact, ROCm runs better on my RDNA machine than MIGraphX. |
|
@lightvector #1222 mentioned that other providers (cuda / migraphx / coreml) are wired in and should work but are unverified, and the code stated that MIGraphX from-source build slots in here once validated (needs ROCm; deferred). Some users may run ONNX with MIGraphX but some may prefer running MIGraphX directly. @Looong01 MIGraphX officially supports MI300X on Ubuntu [1] but doesn't officially support RX 7900 XTX (gfx1100) on Windows [2]. 1: https://rocm.docs.amd.com/en/latest/compatibility/compatibility-matrix.html?fam=instinct&w=compute&gpu=mi300x&gfx=gfx942&os=ubuntu |
|
Thanks, so if migraphx is already available under the onnx backend that was merged, then my request would be to test and make whatever adjustments are necessary to get migraphx working well there instead of merging this PR, which seems to also a wrapper around onnx. Is that possible or am I missing something? Dealing with all the all the different accelerator providers each requiring custom implementations is a big burden, and I'd like to be moving towards unifying under a smaller number of common frameworks. If TensorRT weren't already grandfathered in due to the historical development path, I'd also be pushing for TensorRT to be available via only the onnx backend if possible rather than being its own backend. |
|
@lightvector This PR does not use ONNX Runtime, and it does not load a pre-built .onnx file. We use ONNX only as an in-memory intermediate graph format — built by onnxmodelbuilder.cpp / OnnxModelBuilder::build(), the same emitter the TensorRT backend uses — and then pass that buffer directly to MIGraphX via parse_onnx_buffer. That mirrors how the TensorRT backend works today (ONNX emitter → TensorRT via nvonnxparser), without going through ONNX Runtime. We’re happy to investigate getting MIGraphX working under the ONNX backend (#1222) as you suggested. For context, https://rocm.blogs.amd.com/artificial-intelligence/triton-inference-server/README.html benchmarks ONNX with MIGraphX vs. ONNX with TensorRT. For fair comparisons, we should align on the serving stack. Otherwise, comparing peak-performance native TensorRT on one side against ONNX + MIGraphX (not native) on the other wouldn’t be apples-to-apples. We should also allow users to use native MIGraphX for flexibility. |
|
@lightvector On the #1222 question — I built the comparison arm rather than argue from the code, ONNX Runtime + migraphx EP, measuredORT built from source at v1.20.1 with
The EP works — at threads=1 it is ~380x the CPU provider, so the graph really is running on the The cause is in ORT's own source // input shapes are different, needs to re-parse onnx and
// re-compile the program
if (!input_shape_match) {MIGraphX compiles static shapes. MCTS submits a different batch size on nearly every eval — that Is there a config fix? Not in the version a #1222 user would build today. ORT 1.27 added a (I nearly got this wrong in your favour: I first read the recompile code in a 1.27 checkout that What "get migraphx working well under #1222" would actually takeYour instinct to unify is right, and I am not going to argue against it on weak grounds. But it is
Option 2 is worth dwelling on, because it is the same problem this PR already solves. The reason So the two paths are not really "new backend vs. reuse the ONNX one" — they are "this logic lives Two other differences worth knowing, since both are ONNX-backend gaps rather than design limits:
@Looong01 Thank you for actually building and running it — that is more than I had any right to One thing in your setup needs flagging, and it is what your own caveat already names. You built The second line is an entire attention block — Q·Kᵀ, scale, softmax, ·V — fused into one Rather than assert that matters, I measured it — same machine, same graphs, isolated compiled
That is MLIR alone on gfx942; your build also had CK and hipBLASLt off, so the real handicap is Your writeup already calls this "a lower-bound performance for MIGraphX" — I would just ask that the To be explicit about what I am not claiming: I am not saying MIGraphX would win on RDNA3 with a My own numbers, including the ones against meHead-to-head vs merged master ROCM (not the #1234 branch), one tree, one node, one job, 5
Two corrections to my own published numbers, both against me:
@lightvector So the direct answer to your original question is no — this does not 100% obsolete the What I would ask instead: the small-batch region is not a corner case, it is analysis, GTP play, and |
All my tests are on Ubuntu, not Windows. Btw, there is another reason to support ROCm, that ROCm support both Linux and Windows, but MIGraphX only support Linux. |
|
I have no ill intentions towards adding a new backend; I just want to discuss the issue itself. @lightvector mentioned that his current maintenance workload is quite heavy. I want to step back and ask what user population a merged MIGraphX backend would actually serve. The use case is CDNA training, not RDNA inferenceYour small-batch wins (b10/t16 2.05×, b10/t32 1.69×, b11/t32 1.07×) are measured on MI300X/MI325X (gfx942). At large batch and on larger models the same table shows ROCm winning. That means the MIGraphX advantage region is specifically:
In KataGo, the only workload that fits all three is large-scale selfplay / training data generation. Individual users running GTP or analysis on a single workstation do not drive enough small-batch evals to make that 2× region decisive; they are more likely to be memory- or large-batch-limited, where the ROCm backend is already faster or equal. But training is not on the tableThe public training pipeline at https://katagotraining.org/ is CUDA-based. As far as I know, @lightvector has not announced plans to move KataGo training to AMD GPUs, and neither ROCm nor MIGraphX backends are part of that infrastructure today. If the MIGraphX backend is not going to be used for KataGo's own distributed training, then its natural use case disappears. Most end users are on RDNAThe people who actually download KataGo binaries and run them on AMD hardware are overwhelmingly on consumer RDNA cards (RX 6800/6900/7800/7900 series). On that hardware:
So for the user base that actually exists today, the ROCm backend is the right path. A merged MIGraphX backend would primarily serve a hypothetical future where KataGo training runs on CDNA clusters, and that future is not currently planned. If, after considering all the factors I've mentioned above, @lightvector still believes the MIGraphX backend deserves to be merged, then I have no further comments. I fully respect our repo owner's ideas, plans, and arrangements. |
|
@zhihuidu-amd thanks for the new report and benchmarks! These are extremely helpful to understand the direction... which I think does actually suggest not merging this backend. @Looong01 thanks for the input and nice argument, but if your post is AI-written, also consider extra-double-checking the AI-written posts or try to have them be less overconfident about places they might be making claims that are inaccurate.
Selfplay data generation is large-batch-size, not small batch size, to benefit from the greater throughput of large batch by running a large number of games in parallel.
Not quite. The data generation pipeline (which is the vast majority of the compute) is whatever hardware all the contributors running KataGo are using to generate and upload data. There is nothing that ties it to CUDA, although it is true that in practice NVIDIA gpus are more common among users at the moment. The minority of the compute, the actual training of the neural net itself, is pytorch, which is also technically cuda in practice on the particular GPU machines we use for it but not relevant to a discussion of the C++ backend. Those are the two major inaccuracies I noticed. But in any case, @Looong01 correcting the first inaccuracy above actually improves your argument. Since even self-play data generation does not fit the above advantage profile (being large-batch), and if indeed RDNA is what the vast majority of consumers have, that would leave the advantage region of MIGraphX without any use case. I'll note that also b11c768h12nbt3tflrs-fson-silu is unambiguously the best model even considering compute cost, and there the benchmark showed rocm was better. The only reason for releasing the small transformers (on which MIGraphX appeared to be better) was to offer an alternative for very weak hardware, but the small transformers are unambiguously weaker per compute cost, the only reason you would run them is if your hardware was so weak that you simply could not run the larger models to a minimal number of visits at all, or for research purposes into suboptimally small models. If indeed CDNA is the "datacenter" generation of devices, then I would be surprised if users would be going out of their way to rent datacenter GPUs only to run a weak small model. Thoughts? Is there some case where MIGraphX would still add a lot of value? |
|
@lightvector One important finding before those benchmarks are used to draw a conclusion: we The configuration errorMIGraphX compiles a static shape. This backend compiles one program at
MCTS fills almost exactly half the batch, at every thread count. So the benchmarks I posted had The fix is batch bucketing: compile a small geometric ladder of shapes and dispatch each eval to The results I have so farNot projections — measured. But partial, and I am labelling exactly how partial. 1. Bucketing on vs off vs ROCm, on merged master (
A 0.83x loss becomes a 1.15x win; bucketing alone is worth 1.36x to this backend. 2. Bucketing on vs off at larger batch (
The gain does not fade at high thread counts, which is the relevant question for your point about 3. Correctness, checked before any timing — policy squerr 1.8e-10 against the What I am not claiming yetI do not know that this backend now wins everywhere, and I would rather say so than find out in Results in a few hours. This bears directly on your remark that Everything else in my previous comment stands unchanged — the ONNX-EP recompile finding, and @Looong01 Your point about Windows support on RDNA is fair and I cannot answer it with hardware I |
4 participants
Add a MIGraphX backend (AMD ROCm)
KataGo has a fast TensorRT backend for NVIDIA and a portable OpenCL backend that is
substantially slower on AMD datacenter GPUs. This adds a third GPU backend targeting AMD
via MIGraphX, ROCm's graph compiler.
Measured on MI300X (gfx942), ROCm 7.2.0,
b18c384nbt, 19x19, FP16. Both backends werebuilt and benchmarked inside a single job on one node with one ROCm install, so the ratio
cannot be contaminated by a toolchain difference.
opencltunerautotune)2.94x at 3200 visits, 2.74x at 800. The ratio is stable across search depth because both
backends benefit similarly from NN-cache reuse in longer searches, and
avgBatchSizematchesclosely between them, so this is like-for-like rather than a batching artifact. The 800-visit
pair was reproduced across three independent builds (4239.44 / 4254.13 / 4237.89).
Why it is faster
Not because the port is clever. KataGo's own OpenCL tuner reports
canUseFP16TensorCores=0on gfx942 -- the OpenCL path uses FP16 storage but never issuesMFMA. MIGraphX routes convolutions through rocMLIR/MIOpen, which do, and fuses
conv+bias+activation. The OpenCL ceiling here is structural, not a tuning gap: it
saturates flat at ~1549 nnEvals/s from 160 threads onward while MIGraphX keeps scaling.
Implementation
The backend reuses
onnxmodelbuilder.cpp-- the same self-contained ONNX ModelProto theTensorRT backend feeds to
nvonnxparseris handed to MIGraphX'sparse_onnx_bufferunmodified. Network construction is therefore shared, and the output decode is the
TensorRT backend's decode. New code is one ~900-line backend file plus CMake wiring.
Notable points:
program is compiled at
maxBatchSizeand short batches are zero-padded. Padding rowsget an all-ones mask, not zeros: the graph divides by
maskSumfor masked means, and azero mask row is a division by zero that propagates NaN into real rows.
set_offload_copy(false)) and allocated once.run_async, noteval()--eval()runs onMIGraphX's internal stream and is not ordered against caller-issued copies.
Correctness
Validated with KataGo's own
runnnonmanyposestest(254 positions from a pro game) acrossall 5 nets in
cpp/tests/models/, against the OpenCL backend on the same node:(policyProbSquerr; win/score errors are of the same order.) FP32 agreement is essentially
exact, and MIGraphX's FP16 is 2.2x-51x closer to the FP32 reference than OpenCL's FP16
is. The harness includes a self-check -- OpenCL re-run against its own reference must give
~0 -- which caught two harness bugs during development.
Known limitation:
migraphxTransformerNHWCdefaults to falseThe TensorRT backend defaults
trtTransformerNHWCto true. The MIGraphX equivalentdefaults to false, because the channel-last trunk produces wrong policy output on
transformer nets under MIGraphX while the value heads stay correct:
Every board position on every test position is affected and the logits collapse toward a
near-flat distribution, so this is a wrong computation, not a layout permutation. Root
cause is still open -- either MIGraphX's lowering of an op the channel-last path emits, or
an emitter assumption that only holds under TensorRT. Convnets never take this path (the
emitter only goes channel-last when the model has transformer blocks). The flag is
retained so the NHWC path can be re-enabled once fixed.
Build
Requires ROCm with MIGraphX runtime and headers (
migraphx,migraphx-dev), plus thestatic protobuf library
libprotobuf.a.Compiling.mdgains a short sectionexplaining why the static link is mandatory:
libmigraphx_onnxexports its bundledprotobuf symbols as weak, so an application linking its own shared libprotobuf preempts
them and the ONNX parse aborts inside protobuf's
repeated_field.h-- at model load, notat link time, which makes it confusing to diagnose. Linking statically with
-Wl,--exclude-libs,ALLkeeps the two copies apart;nm -D --defined-only ./katago | grep -c protobufprints 0 on a correct build (verified on the tree in this PR).Relationship to #1188
#1188 also adds AMD support, including a MIGraphX backend, and is the larger change
(ROCm/MIOpen + MIGraphX + Windows, +16k lines). This PR is independent of it -- it branches
from the v1.17.2 tag and shares no files with that branch -- and takes a different approach
to MIGraphX specifically: it reuses the existing ONNX emitter rather than building the
network op-by-op, which is why it is ~900 lines instead of ~1900. The two are
complementary and either can be taken without the other; I'd defer to whichever the
maintainer prefers.
Not included