[Proposal] ROCm (AMD Instinct) support for the core attention and distributed init #739

Summary

We run FlashDreams' LingBot-World path on AMD Instinct MI355X (gfx950) on 1, 2, 4 and 8 cards. On a stock ROCm PyTorch build, two places in flashdreams.core stop it from starting. We work around both with a small runtime shim today. Both are still present on main (59cf6a0).

Before sending PRs, we would like to check three things:

1. Is ROCm support something you would accept upstream?
2. Do you agree with the shape proposed below?
3. Is anyone already working on this?

Neither change alters behaviour on CUDA.

Where it breaks on ROCm

1. The SDPA backend is hardcoded to cuDNN.

• NativeAttention defaults to backend="cudnn" (native.py#L24).
• ContextParallelAttention also defaults to "cudnn", and it accepts only cudnn or flash (cp.py#L57, cp.py#L71).
• Several recipes pass backend="cudnn" explicitly:
  • the Wan blocks (modules.py#L178, #L402);
  • Cosmos (modules.py#L256);
  • the template recipe (network.py#L119);
  • OmniDreams (modules.py#L297).

ROCm builds of PyTorch have no cuDNN attention. LingBot-World therefore fails at its first attention call, and so would anything else built on the Wan blocks, since they share this path. We have only run LingBot. In our shim, we change only that one constructor argument to "flash". That sends the same call to aten._scaled_dot_product_flash_attention, which runs AOTriton on gfx950. Nothing else in the attention path needed to change.

2. flashdreams.core.distributed.init() calls NVIDIA-only libraries. See distributed/init.py#L185 and #L218-L221.

• pynvml.nvmlInit() runs before the try that already handles NVMLError. On a machine without NVML it raises, and the process never reaches init_process_group.
• The function ends with ctypes.CDLL("libcudart.so") and cudaDeviceSetLimit(0x05, 128), an L2 fetch-granularity tweak. ROCm has no libcudart, so this raises OSError.

Only multi-rank starts reach either call. In our shim, nvmlInit becomes a no-op, so the existing except skips CPU affinity, and libcudart resolves to a stub. init_process_group(backend="nccl") itself works unchanged, because ROCm builds of PyTorch route it to RCCL. RCCL also asks for HSA_NO_SCRATCH_RECLAIM=1, which we set in the environment.

Proposed changes (two small PRs)

PR 1: resolve the attention backend by platform.

• Change the default of NativeAttention and ContextParallelAttention from "cudnn" to "auto".
• "auto" resolves once, at construction:
  • CUDA builds: "cudnn", exactly as today;
  • ROCm builds (torch.version.hip set): "flash".
• Remove the explicit backend="cudnn" from the recipes listed above, so they inherit the default.
• An explicit backend=... still wins everywhere, so anyone who pins a backend sees no change.
• Add CPU tests (ci_cpu) for the resolution, with torch.version.hip monkeypatched.

PR 2: make distributed.init() tolerate a machine without NVML or libcudart.

• Move nvmlInit() inside the existing try. If NVML is missing, init logs the same warning it already logs and carries on without setting CPU affinity.
• Run the L2 fetch-granularity tweak only on CUDA builds, and keep a guard for OSError.
• Mention the RCCL environment variable in the docs. No code change for it.

Possible follow-up, separate and optional: allow backend="efficient" in ContextParallelAttention.

aten._scaled_dot_product_efficient_attention also returns the log-sum-exp when compute_log_sumexp=True, so the ring path could use it too. On gfx950 it measured 7% faster than flash on our widest attention shape (see the kernel table below). This one is only worth doing if you want it.

Not in scope. These are CUDA-only kernels:

• sage, fp8_tma, scaled_fp8 self-attention;
• rowwise_fp8 linears.

The new lingbot-world-v2-1p3b-causal-fast-perf preset uses sage with TMA. We have not tried it on ROCm. A clear error when a CUDA-only backend is requested on ROCm seems better than a deep failure, if you agree.

What we measured

Setup. Measured on FlashDreams 0f39773 (2026-07-17), the commit our integration is pinned to. We have not yet re-run on 0.2.0.

• Software: the public rocm/pytorch image, with torch 2.12.0+rocm7.14.1, HIP 7.14 and Triton 3.8.
• Hardware: AMD Instinct MI355X (gfx950, 288 GiB).
• Model: LingBot-World v2, bf16, compile_mode=default, TAEHV decoder, 12-frame chunks, UMT5 on the GPU.
• Method: 30–40 measured chunks after 8–10 warm-up chunks, one seed, camera held forward. p95 was within 2% of p50 in every run.
• The H100 SXM column uses the same driver script, seed and presets.
• The fps figures are native generation rate, not delivered or interpolated fps.

One card

• On GEMM-bound shapes, MI355X runs about 1.5–2× faster than H100.
• On attention-bound shapes it runs at about the same speed. An example is the 1.3B at window 30 / sink 6.
• hipGraph capture through FlashDreams' own CUDAGraphWrapper works unchanged:
  • 1.3B: 53.3 → 53.6 fps;
  • 14B: 16.3 → 16.35 fps.

Context parallel

These runs use the stock Ulysses path. The model is the 14B at 832×480. The H100 ladder was measured on 2026-08-30 with the same driver script. A re-run of its CP=2 and CP=8 points on 2026-09-13 came within 1% of it.

SDPA kernels on gfx950

This is a raw SDPA call on the window-30 shape: query 4680 × key 51480 tokens, 12 heads, d128, bf16, warm.

The CK preference is process-wide. It also breaks the Wan VAE encoder, which uses a head dimension above 256. That is why we do not propose it.

FP8 (FYI, not a FlashDreams change). We also run torchao float8 rowwise quantization on the DiT. torchao 0.17 refused it on gfx950. torchao 0.18.0 accepts MI350-class GPUs, so nothing is needed here.

Related: Ulysses KV traffic (we can file this separately)

On a two-card MI355X pair, the stock 14B lost speed at CP=2: 10.6 → 9.06 fps at 832×480.

The cause is in _impl_ulysses. On every attention call, it all-to-alls the rank's whole cached K/V window, not just the new chunk. On the 1.3B at 720p with window 15 / sink 3, we profiled about 265 MB per call. A two-card XGMI pair moves it at the link limit.

We keep a head-sharded mirror of the KV cache and exchange only the new chunk's K/V on each step. Results:

• In eager mode, the output matches the stock path exactly: max abs diff 0 over 757 attention calls, identical frames (1.3B, 720p, CP=2).
• Traffic on the same run: 34.8 → 9.95 GB per rank per chunk.
• 14B at the 704×384 frame on MI355X: 2 / 4 / 8 cards go from 12.7 / 28.8 / – to 17.3 / 35.1 / 40.4 fps.
• On H100 with NVLink, the gain is 4–7%.

We have not checked this against the tensor- and context-parallel runtime added in #619. If it is of interest, we will rebase it on main and open a separate design issue.

Testing we can provide

• CPU tests (ci_cpu) for both PRs.
• For each PR, before and after numbers on MI355X (1, 2 and 4 cards), plus a CUDA check on H100 that frames are unchanged in eager mode.

We cannot offer an AMD runner for your CI.