[xmiksay]

ROCm/HIP backend for AMD GPUs

Continues @airpods69's #3424, which established the foundation this builds on:
rocm-rs for the bindings, the rocm feature and its wiring through Device,
Storage and the backend traits, rocm_backend shaped as a mirror of
cuda_backend, and a working device, allocator, RNG and matmul. Its open
checklist — convolutions, training a model, and "update the API to make more
sense" — is what this finishes. Thanks also to #3186 for opening the topic.

Tested on a Radeon RX 7800 XT (gfx1101, RDNA3, wave32), ROCm 7.2.4 — the same
architecture as #3424's RX 7700 XT, one ROCm major version newer. Everything
below is measured on that card; nothing is extrapolated from another vendor's
numbers or another architecture's.

The main design decision: share the CUDA kernel sources

#3424 hand-wrote .hip kernels. This compiles candle-kernels/src/*.cu — the
very files the CUDA backend uses — with hipcc, bridged by a ~360-line shim in
candle-rocm-kernels/src/hip_shim/ that shadows the CUDA header names the
sources #include (cuda_fp16.h, cuda_bf16.h, cuda_fp8.h, cuda.h,
cuda/std/limits) from its own include directory.

All eleven modules compile unmodified, including the 4,845-line quantized.cu.
The shim supplies only what HIP genuinely lacks: a __nv_bfloat16 typedef,
__hmax_nan/__hmin_nan for bf16, 16-bit atomicAdd via CAS, __dp4a,
__vsubss4, and *_sync shuffle wrappers.

Hand-written kernels drift, and the ones in #3424 already had: utan where
candle calls utanh (so tanh was broken on every dtype), a bminimum that
dropped NaN handling, a UNARY_OP missing its in-place path, and a BINARY_OP
missing the out-type parameter comparison operators need — which is why cmp
had no kernels to call. Sharing the sources makes that class of bug
structurally impossible.

What this costs the CUDA path

An earlier revision of this description claimed candle-kernels was untouched.
That is no longer true, so here is the exact diff — +155/-43 across two
files, none of which changes an existing kernel's behaviour under nvcc:

• quantized.cu, +90/-43. The mul_mat_q* launchers had their tile
  geometry hard-coded to the Ampere set. They now select through an MMQ_ARCH
  macro that resolves to AMPERE unless RDNA2/RDNA3 is defined, and only
  candle-rocm-kernels ever defines those — nvcc resolves to the same
  constants the file used before. The RDNA arms additionally carry
  __launch_bounds__, which expands to nothing off RDNA. Without it the AMD
  backend budgets registers for a 1024-thread block, eight times what these
  launch, and spills 1748 bytes per lane on mul_mat_q3_K. The rest of that
  diff is two genuine bug fixes that apply to CUDA as well — see 5. below.
• unary.cu, +65/-0. One new kernel family, transpose2d_*. Purely
  additive; no CUDA-side code calls it yet, though cuda_backend's
  copy_strided_src has the same slow transpose and could.

I have no CUDA toolchain here, so none of this is nvcc-compiled. The
MMQ_ARCH indirection is macro paste and the transpose is plain CUDA C++, but
that is reasoning, not evidence, and it is the one part of this PR that CI has
to check rather than me. git diff origin/main -- candle-core/src/cuda_backend candle-core/src/metal_backend candle-core/src/cpu_backend is still empty.

Non-invasive by design

Device::Rocm, Storage::Rocm, QStorage::Rocm, rocm_fwd and
pub mod rocm_backend all sit behind #[cfg(feature = "rocm")]. There is
deliberately no dummy backend, unlike cuda/metal — so a default or
--features cuda build gains no enum variant, no trait method and no public
type. The cost is ~40 extra #[cfg] lines inside candle-core, which is what
storage.rs and tensor.rs already do.

Features: rocm (default path, im2col + rocBLAS GEMM) and opt-in
miopen = ["rocm", "rocm-rs/miopen"], which swaps convolutions onto MIOpen the
way cudnn layers over cuda. A plain rocm build never links libMIOpen.

Evidence it works

Every candle test suite, on the GPU, as of the head of this branch:

Suite	Result
candle-core --lib	136 (140 with --features miopen)
tensor_tests	35
quantized_tests	26
bilinear_tests	12
conv_tests	9
grad_tests	6
matmul_tests	5
pool_tests	4
layout_tests	1
candle-nn ops	10
candle-nn --lib	9
candle-rocm-kernels --lib	19
ug micro-kernels (--features "rocm ug")	6

test_device! gained a fifth ident, so every existing suite generates its
*_rocm case rather than ROCm getting a bespoke harness.

Numerical agreement is checked against the other backends, not just
self-consistency. In quantized_tests, the ROCm expectations are
bit-identical to CUDA's — ROCm runs the same MMVQ kernels from the same
.cu, so it lands on the same numbers. The DMMV decode path is verified
against a dequantize-then-matmul reference across all ten dtypes at k=256 and
k=4096, and MMQ against the same reference for all ten.

End to end:

$ cargo run --release --example bert --features rocm
score: 0.89 'The new movie is awesome' 'The new movie is so great'
score: 0.68 'The cat sits outside' 'The cat plays in the garden'

$ cargo run --release --example quantized --features rocm -- \
    --which SmoLM2-360M-Instruct --sample-len 100
loaded 290 tensors (384.62MB) in 0.33s
  6 prompt tokens processed: 151.35 token/s
 99 tokens generated: 143.96 token/s

CI parity, run locally: cargo fmt --all -- --check, cargo clippy --workspace --tests --examples --benches -- -D warnings, cargo check --workspace (also
with +avx2), and cargo test --workspace all pass. The arm-neon and macOS
metal jobs are unverified here for want of a cross toolchain.

Included

• Ops: the full BackendStorage/BackendDevice surface — elementwise,
  cast, affine/powf/elu, reduce (with u32 argmin/argmax), cmp, where_cond,
  matmul via rocBLAS, conv1d/2d and transposes via im2col, pooling, upsample,
  the gather/scatter/index_add family, arg_sort, and RNG including f16/bf16.
• candle-nn: rocm_fwd for softmax, rms_norm, layer_norm, sigmoid and all
  three rotary-embedding variants, across BF16/F16/F32/F64.
• Quantized/GGUF: four matmul paths — MMVQ (batches 1–8, f32/f16/bf16
  activations), the fused DMMV decode kernel, MMQ (mul_mat_q*) for larger
  batches, and dequantize-then-GEMM. Plus quantize_imatrix and a raw weight
  pointer for downstream kernels.
• Quantized MoE: moe_gemm_gguf at every batch size, via the MMVQ kernels
  with the expert index folded into the weight pointer.
• F8E4M3 arithmetic — compiled and reachable now. gfx1101 has no fp8
  hardware, so this is conversion-based, but it is exercised.
• ug micro-kernels (--features "rocm ug"), compiled through the same
  runtime hipcc path, and KernelCache::custom_function so downstream crates
  can compile kernels of their own in the same dialect and disk cache.
• A GPU-free CI job (.github/workflows/ci_rocm.yaml). rocm-rs needs
  ROCm headers, not a device, so the job extracts seven -dev debs (~14 MB)
  rather than apt-installing the ~9.5 GB runtime. A self-hosted GPU job is
  present but workflow_dispatch-only, inert until a runner exists.
• rocm feature for candle-pyo3; docs in the README and the book.

Quantized matmul: how the dispatch was chosen

Every threshold in quantized/rocm/ is a measurement, not a shape heuristic,
and quantized::rocm::bench* (--ignored) reprints the sweeps. Two results
are worth surfacing because they are counter-intuitive:

MMQ is not near either roofline at prefill. At Qwen3.5-2B's shapes it
reaches 20–22 TOP/s against gfx1101's ~75 TOPS dp4a ceiling, and generates
165–185 GB/s of weight traffic against 624 GB/s of bandwidth — bound by
neither. Raising mmq_x from 64 to 128 made Q8_0 and Q4_0 worse, so the 28%
is the kernel design, not the tile geometry.

rocBLAS on gfx1101 is erratic with a transposed operand. The dequantize
kernels write weights as n rows of k; describing that buffer as (k, n)
with swapped strides is free on the host but reaches rocBLAS as
Operation::Transpose, where the same f16 GEMM ranges over 6–45 TFLOP/s with
no monotonicity in the batch — against 25–57 for the plain form. So the dense
path now materialises the other orientation first, through a new tiled
transpose kernel, and fwd prefers it over MMQ for f16 from a batch of 256 up.
Worst cell over three dtypes and four shapes, mmq_ms over the dense path
(above 1.0, dense wins):

batch	128	256	384	512	1024	2048
transposed rhs	1.21	0.51	1.62	1.46	0.71	0.72
reoriented rhs	0.85	1.19	1.38	1.55	1.36	1.90

Reorienting is often the slower of the two, by up to 1.8x. What it buys is a
path that is predictable: the transposed row dips under 1.0 at three of six
batches with no pattern, so no threshold makes it safe to dispatch on. End to
end through QMatMul::forward, batch 255 (MMQ) against 256 (dense) on the same
shape: 1.2–2.1x, rising to 3.0x at larger batches.

The transpose kernel is not quantized-specific — copy_strided_src routes any
swap of the last two dimensions to it, and it moves 292–662 GB/s where the
generic strided copy managed about 20.

An earlier round of this was equally counter-intuitive in the other direction:
preferring MMVQ everywhere looked better in microbenchmarks and regressed real
generation by 28%.

Deliberately not included

candle-rocm-kernels/README.md has the full list with the reason for each, so
the gaps don't have to be found by grep. The substantial ones:

• fast_mmq / fast_mmvq. They bind to sources using inline PTX,
  nvcuda::wmma and __reduce_add_sync. The portable mul_mat_q* kernels are
  what this launches instead.
• The dense fused MoE GEMM (moe_gemm, and the prefill half of
  moe_gemm_gguf) — nvcuda::wmma compiled into a host-side libmoe.a. The
  quantized MoE is implemented, so FusedMoeGGUF runs at every batch; prefill
  goes through the mat-vec kernel, which is correct but not tile-optimised.
• Flash attention. candle-flash-attn{,-v3} depend on candle-core with
  features = ["cuda"] unconditionally, so there is nothing to gate.
• NCCL/RCCL, and so multi-GPU collectives.
• CUDA graphs. No hipGraph counterpart here.
• upsample_nearest1d (CUDA bails too), and i16/i32 for reduce/indexing/cast
  (the shared .cu never instantiates them).

Upstream bugs found

Building this surfaced five defects in existing backends, all documented in
candle-rocm-kernels/README.md under "Divergences from cuda_backend". The
first four are left alone — each needs its own fix and tests, and this PR is
large enough:

1. cuda_backend::index_add passes the source layout as the freshly
   allocated accumulator's, so a non-zero start_offset is applied twice.
   metal_backend doesn't do this and Tensor::scatter uses
   Layout::contiguous; ROCm follows those.
2. quantized/cuda.rs's dequantize fallback builds a typed slice over a
   Vec<u8> via from_raw_parts. Vec<u8> guarantees only byte alignment, so
   that's UB for any block type with stricter alignment. ROCm uses
   read_unaligned.
3. cpu_backend's MatMul::f collapses the batch loop when
   b_skip == 0 && a_skip == m*k without checking that the lhs rows are
   contiguous. conv_transpose1d's col2im path hits exactly that shape, so
   CPU conv_transpose1d with batch > 1 is wrong from batch 1 onward.
4. Every CUDA F8E4M3 launch names a kernel that does not exist.
   cuda_backend::kernel_name builds {root}_f8e4m3 from DType::as_str(),
   while the shared sources spell their entry points _f8_e4m3 or
   _fp8_e4m3 — no symbol in candle-kernels/src ends in _f8e4m3. Since
   cuda_backend/utils.rs does dispatch the F8E4M3 arm of every Map*
   trait, those ops resolve a missing function at run time rather than failing
   to compile. This backend keeps both spellings in FP8_SPELLED_ROOTS, with a
   test against the shipped .cu text.
5. Two bugs in the shared quantized.cu are fixed here rather than only
   reported, because they are in code this backend had to touch anyway and are
   wrong on CUDA too: the MMVQ store ran one row past the output for an odd row
   count (4f7eff3d), and indexed_moe_forward strode its weights by QK_K
   instead of the dtype's qk (6a35ea18).

Also worth flagging for review: rocm-rs is a third-party 0.5.x bindgen crate
rather than something with cudarc's track record. It's an optional dependency
reachable only via --features rocm, and it's used with
default-features = false because its default gpu-sort feature builds an
amdgcn kernel through a proc macro that shells out to a nested nightly
-Zbuild-std invocation. Happy to discuss alternatives.

[xmiksay]

Status update: CUDA/ROCm gap closure, plus four bugs that affect CUDA

Latest round of work on this branch closes most of the remaining functional gaps between the CUDA and ROCm backends, and fixes several defects found along the way. Four of those defects are in shared code and affect the CUDA backend too — those are called out first, since they matter independently of whether this PR lands.

Pre-existing bugs fixed in shared code (affect CUDA)

Commit	Bug
4f7eff3d	MMVQ wrote one row past the end of each output column. mul_mat_vec_q guarded its store on threadIdx.x < rows_per_cuda_block only, missing upstream llama.cpp's row0 + threadIdx.x < nrows_dst. rows_per_cuda_block is 2 for ncols_y >= 2 on both backends (the 1-row arm requires RDNA2/RDNA3, which neither build defines), the grid is nrows.div_ceil(2), and dst is allocated at exactly nrows * b_size floats. Every odd nrows with a batch of 2..8 therefore wrote into the next batch row's first output — silent corruption, not a fault. Measured on gfx1101 with the fix reverted: n=4097, m=2, k=512, Q4_0 returns 207.7 against a reference of -214.6.
6a35ea18	indexed_moe_forward strided between experts by QK_K in a template over qk. The five K-quant instantiations pass QK_K as qk and are unaffected, but the q8_0 instantiation passes QK8_0 = 32, so it strode by an eighth of the expert matrix and every routed pair not going to expert 0 read the wrong weights.
f4e6ef83	WithDType::to_f64 for F8E4M3 recursed into itself. The macro was passed `
678088f5	Custom kernel modules were cached in memory by name while the disk cache keyed on source, so a revised source silently resolved against the module it replaced. ROCm-only, but the same name-only keying exists in CudaDevice::get_or_load_custom_func.

Gaps closed

Gap	Result
fp8 (F8E4M3)	Unary, binary, comparison, affine/powf/elu, where_cond, gather/index_select/index_add/scatter_add, both cast directions, strided copy, const_set. Two causes: the compile flags pinned __CUDA_ARCH__=800 while the fp8 kernels sit behind 890, and RocmStorageSlice::F8E4M3 held SendSyncDeviceMemory<u8>, so generic dispatch resolved *_u8 kernels.
Quantized MoE	QRocmStorage::indexed_moe_forward and candle_nn::moe::moe_gemm_gguf, built on indexed_moe_forward_*_q8_1 in quantized.cu with no new device code. FusedMoeGGUF, and so quantized_qwen3_moe, runs on ROCm.
ug codegen	ug emits CUDA C, not PTX, so hipcc consumes it. HIP code generator added to candle-ug behind a rocm feature (no ug-hip exists upstream, and ug-cuda cannot be built without nvcc). RocmDevice::compile + UgIOp1::rocm_fwd.
imatrix quantization	quantize_imatrix/_onto, byte-identical to CPU for Q2K..Q6K. No device dependency — the CUDA path also round-trips through host k_quants.
QStorage::device_ptr	Plus QRocmPtrGuard, which the ROCm allocator requires: freed blocks go to a free list rather than back to the driver, so an escaped pointer aliases the next tensor rather than merely dangling.
rocBLAS handle, stream constructors	RocmBlas carries its stream binding and deliberately does not deref to rocm_rs::rocblas::Handle, whose set_stream takes &self. new_with_stream was a naming gap: hipStreamCreate already gives every ROCm device a private stream.

MMQ performance

quantized.cu carries one MMQ tile set per architecture, but all ten launchers named MMQ_*_AMPERE unconditionally, so the *_RDNA* constants were dead text. With no __launch_bounds__ the AMD backend budgets registers for the default 1024-thread block — 8x what these launch — capping every MMQ kernel at 192 VGPRs and spilling accumulators to scratch (mul_mat_q3_K: 557 spills, 1748 bytes per lane; now 243 VGPRs, zero spills).

After re-benchmarking (9 weight sizes x 4 batches x 10 dtypes, three independent runs, each cell a median of five repeats), all ten quantized dtypes now beat the dense dequantize+rocBLAS path; four had previously been excluded from MMQ dispatch entirely. Q3K went from 0.20-0.54x to 1.04-1.70x.

CUDA is unaffected: the RDNA define is set only on the hipcc command line, and preprocessing quantized.cu without it is byte-identical to preprocessing the original.

Tests

Environment: RX 7800 XT (gfx1101, RDNA3), HIP 7.2.53211, rustc 1.97.0.

Target	Result
make test-rocm-core	226 passed, 0 failed, 3 ignored
make test-rocm-nn	19 passed, 0 failed
make test-rocm-ug	6 passed, 0 failed
cargo test -p candle-rocm-kernels	19 passed, 0 failed
make rocm-shim-test	11 modules compiled under hipcc, 0 failures
make check-rocm / make check	pass
make clippy-rocm (-D warnings) / make fmt-check	pass

One caveat on earlier green reports from this branch: ROCM_FILTER defaulted to _rocm, which matches the integration suites' ..._rocm suffix but not the backend module tests, whose names carry rocm only in the module path (quantized::rocm::tests_mmvq::…). make test-rocm-core was running 5 tests and skipping 92 — the entire mmvq/mmq hardware suite had never executed under make. Fixed in a2af1863; the counts above are post-fix.

New GPU coverage this round: MMVQ odd-row regression, MMQ geometry/threshold policy, fp8 against CPU references, MoE against CPU references for all six dtypes, imatrix byte-equality with CPU, device_ptr guard ordering, ug kernels executing on device, and a downstream rocblas_saxpy call through the exposed handle.

Not implemented, with reasons

• Flash attention — candle-flash-attn, candle-flash-attn-v3 are CUDA-only crates.
• NCCL/RCCL — no multi-GPU collectives.
• CUDA graphs — no ROCm counterpart to enable_cuda_graph_htod_cache.
• External-stream device construction — declined. An audit found ~120 sites assuming one stream per device: the free-list allocator recycling blocks with no event query, every launcher dropping a scratch buffer while its kernel is in flight, rocBLAS/rocRAND/MIOpen handles bound once, and the MIOpen shared workspace. Admitting a second stream needs the allocator reworked first.
• moe_gemm (dense f16/bf16) and the WMMA prefill kernel — nvcuda::wmma, and compiled into a static library rather than a loadable module. ROCm takes the mat-vec path at every batch size: correct, not tile-optimised. No tensor-expression fallback was written on purpose; the options were a per-row weight gather (3.2 GB at prefill) or one matmul per expert.
• fast_mmq / fast_mmvq — inline PTX, nvcuda::wmma, __reduce_add_sync.
• set_gemm_reduced_precision_bf16 / _f32 — accepted and stored but inert; rocBLAS has no bf16 compute type and the xf32 analogue is CDNA-only. The f16 knob is live.

fp8 reduce/argmin/argmax/softmax/rmsnorm/rope, scatter-set, and sort/conv/pool/matmul are shared limitations, not ROCm gaps — those instantiations are commented out in the shared sources (reduce.cu:657-664, sort.cu:78-79, conv.cu:817-826) and indexing.cu has SA_OP_F8 with no S_OP, so CUDA lacks them too. The README now separates the two categories, which it previously conflated.

Two things worth a maintainer's eye

1. CUDA's fp8 dispatch appears to resolve names that do not exist. kernel_name builds {root}_{dtype} from DType::as_str, which is "f8e4m3" (dtype.rs:87), but every fp8 symbol in candle-kernels/src is spelled _f8_e4m3 or _fp8_e4m3. cuda_backend/utils.rs does dispatch the F8E4M3 arm of each Map* trait, and cuda_backend/mod.rs:2498 hardcodes "ucopy_f8e4m3". Not touched here — there is no NVIDIA hardware on this machine to confirm the runtime failure or validate a fix.

2. storage.rs performs no same_device check on the ROCm copy paths, so a cross-RocmDevice copy is reachable and would be ordered on the wrong stream. Not addressed in this round.

Verification limits

There is no CUDA toolkit or NVIDIA GPU available here — cudarc's build script hard-fails without nvcc — so make check-cuda could not be run. Where shared sources changed, CUDA safety rests on source review, and for the MMQ tile change on a demonstration that preprocessing without the RDNA define is byte-identical to the original. The CUDA-affecting fixes above are unverified on NVIDIA hardware and would benefit from a second pair of eyes.

[xmiksay]

Follow-up: moe_gemm_gguf now dispatches on the device rather than the feature set (7bcfb9d6)

A review of how candle_nn::moe reaches the ROCm backend turned up a wiring defect in the commit above (b59d346a). Worth flagging separately since it is a Cargo-feature footgun rather than a kernel issue.

The ROCm implementation was gated:

#[cfg(all(not(feature = "cuda"), feature = "rocm"))]
mod rocm;

Cargo features are additive and nothing in the tree forbids cuda + rocm together — one dependent asking for each is enough to unify them. In that build:

• mod rocm is not compiled, so rocm::moe_gemm_gguf and its tests disappear silently
• the CUDA variant becomes the only moe_gemm_gguf, and a ROCm tensor reaches input.device().as_cuda_device()? and errors at runtime

candle-core already answers this question the other way — matching the storage at runtime with both arms compiled in, e.g. QStorage::device_ptr at quantized/mod.rs:310-312. candle-nn/src/moe.rs was the only place selecting a backend at compile time.

It now matches on input.device() under any(cuda, rocm), and the CUDA entry points move to moe/cuda.rs mirroring moe/rocm.rs, so moe.rs is just the dispatcher. The move is verbatim — only the quantized::{self, QTensor} import narrows to QTensor, self having been unused in the moved code.

moe_gemm (dense f16/bf16) is untouched: it has no ROCm implementation either way, so its existing cuda/not-cuda split already covers every build.

Verified on gfx1101: make test-rocm-nn 9 passed / 0 failed with all five moe::rocm tests green, plus make check-rocm, make clippy-rocm (-D warnings), make fmt-check and the workspace make check. The cuda + rocm combination itself could not be built here — no nvcc — so that half rests on reading the cfg logic.

[cryptomilk]

Thank you very much for your work. As we have a day of learning in the company today. I though I would give this a try.

I'm working on Crane. So I thought lets try to wire it up to run Qwen 3.5 on my GPU (Radeon RX 9070 XT). However I couldn't run it:

JIT-compiling the quantized kernel module (via candle_core::rocm_backend's hipcc-at-runtime pipeline) fails on ROCm 7.14.0 with:

error: redefinition of 'atomicAdd'

for both __half and __hip_bfloat16.

candle-rocm-kernels/src/hip_shim/hip_compat.h's CANDLE_HIP_ATOMIC_ADD_16 macro is built on the documented assumption:

HIP has atomicAdd for float/double/int, but not for the 16-bit float types.

That was true against whatever ROCm version this crate was developed and tested against (ROCm 7.2.4 / gfx1101, per the crate's own README), but is no longer true on ROCm 7.14: /opt/rocm/include/hip/amd_detail/amd_hip_fp16.h
and amd_hip_bf16.h already define atomicAdd for __half and __hip_bfloat16 themselves.

https://github.com/cryptomilk/Crane/tree/asn-rocm has the code and a compose.yaml:

podman compose up

would reproduce it.

[xmiksay]

Thank you very much for your work. As we have a day of learning in the company today. I though I would give this a try.

I'm working on Crane. So I thought lets try to wire it up to run Qwen 3.5 on my GPU (Radeon RX 9070 XT). However I couldn't run it:

JIT-compiling the quantized kernel module (via candle_core::rocm_backend's hipcc-at-runtime pipeline) fails on ROCm 7.14.0 with:

error: redefinition of 'atomicAdd'

for both __half and __hip_bfloat16.

candle-rocm-kernels/src/hip_shim/hip_compat.h's CANDLE_HIP_ATOMIC_ADD_16 macro is built on the documented assumption:

HIP has atomicAdd for float/double/int, but not for the 16-bit float types.

That was true against whatever ROCm version this crate was developed and tested against (ROCm 7.2.4 / gfx1101, per the crate's own README), but is no longer true on ROCm 7.14: /opt/rocm/include/hip/amd_detail/amd_hip_fp16.h and amd_hip_bf16.h already define atomicAdd for __half and __hip_bfloat16 themselves.

https://github.com/cryptomilk/Crane/tree/asn-rocm has the code and a compose.yaml:

podman compose up

would reproduce it.

Hi, when I wrote it, I tested it with Crane as well, and I was able to run it (this was a reason why I wrote it :) )
Thx for your test; I will rebase it and adopt a new version of ROCm + test it with Crane as well one more time. The goal was to be as fast as Ollama...

I have prepared branche for Crane for ROCm as well.

[cryptomilk]

Oh this is awesome! We have quite some changes in the meantime :-)

[cryptomilk]

And if you are not aware of it yet, there is data/crane-model-download for downloading models now :-)

[cryptomilk]

Awesome, I was able to load the model now. But with my first request, I already ran out of memory :-)

model=qwen3.5-9b-gguf quant=Q5_K_M

[Hahihula]

It's great to see the broadening of HW support! Hope this gets to main as soon as possible.

[cryptomilk]

The last days I've tested Crane build against this ROCm branch with LLM (Qwen 3.5), TTS (Voxtral) and ASR (Qwen3 ASR) and it is running smooth.

[zachelnet]

Hi @xmiksay,

thanks for your PR. I found the following:

ROCm conv path: f32 im2col blows up at ~640p and is a blocker for Stable Diffusion / FLUX

I benchmarked it with a real conv-heavy model (Real-CUGAN 2x super-resolution, ~30 conv/conv_transpose layers) on my Radeon RX 9070 XT (gfx1201, RDNA4), ROCm 7.14.0. Correctness is fine — outputs match CPU to ~1e-5.

Impact

• f32 convs hit a scaling cliff at large spatial sizes. USE_IM2COL_CONV2D/1D = true (rocm_backend/ops_conv.rs) always materializes (b·h·w, c_in·kh·kw):
  • CUGAN: 480p→640p is 1.78× the pixels but 4.2× the time (167→700 ms); 720p = 1107 ms. f16 stays linear throughout (73→127→229→290 ms), so this is memory-bandwidth pathology, not fp32 throughput. burn/Vulkan runs the same model at 720p in 313 ms (f32) / 46 ms (fp16).
  • On a display-driven GPU the multi-GB allocations can starve the driver: amdgpu_dm ... Failed to pin framebuffer with error -12 → GNOME/GDM restarts.
  • SD/FLUX are affected, not just SR — on the SD-VAE-decoder shapes (3×3, 512ch): 40 ms at 256², 210 ms at 512² (≈1 s+ per VAE decode), and 19.3 GB im2col at 1024² → hipErrorOutOfMemory, so SDXL/FLUX are effectively unrunnable on 16 GB cards.

Cause

• im2col materializes h·w × c_in·kh·kw f32; past a few hundred MB it's memory-bound and eventually OOMs. A direct-conv backend (no im2col) I also tested scales linearly to 1080p and allocates only O(input).

Fix

• Dispatch threshold: route to a fallback non-im2col (direct) conv path when the estimated im2col size in bytes — h_out·w_out·c_in·kh·kw·sizeof(dtype), so the threshold naturally differs for f32 vs f16 — would exceed ~256–512 MB. Small change; removes the cliff, the desktop crash, and the SD blocker.
• Longer term: a shared-memory-tiled direct conv would eliminate the im2col round-trip entirely and get f32 toward burn/MIOpen numbers.

Happy to share the benchmark if useful as a conv-heavy smoke/perf test.

[zachelnet]

Follow-up: I tested the direct-conv fallback too, and it's not a free win. Flipping USE_IM2COL_CONV2D = false (the existing ops_conv_direct path) fixes the OOM/launch failures — 1024² now runs in both f32 and f16 with only ~O(input) memory. But the naive direct conv is ~18× slower than im2col+GEMM at 512² (3.8 s vs 210 ms) and 17 s per layer at 1024² — unusable in practice. So a dispatch threshold only makes sense together with a tiled / implicit-GEMM direct conv (or MIOpen); the naive direct path alone just trades OOM for unusable latency.

[cryptomilk]

This is also working on ROCm 10.

[cryptomilk]

@xmiksay could you add the four commits of https://github.com/cryptomilk/candle/tree/asn-feat-rocm-backend at the top?

[cryptomilk]

@Christof23 @ivarflakstad @EricLBuehler gentle ping