Dev.to AI 🤖 Ai 👁 0 📖 7 min read

Running FlashVSR on an AMD GPU with ROCm (and what broke along the way)

FlashVSR is one of the more interesting video super-resolution models right now. It's a one-step diffusion model built on Wan 2.1, it streams long videos in chunks, and on an NVIDIA card it's fast enough to be practical.

FlashVSR is one of the more interesting video super-resolution models right now. It's a one-step diffusion model built on Wan 2.1, it streams long videos in chunks, and on an NVIDIA card it's fast enough to be practical. I used the FlashVSR-Pro fork, which wraps it in a usable infer.py.

It also has a hard dependency on Block-Sparse-Attention, a CUDA extension. No CUDA, no attention kernel, no model.

I build a video restoration app that has to work on AMD cards too, so I spent a good chunk of August getting FlashVSR running on a Radeon RX 7800 XT (16 GB, RDNA3, gfx1101). It works now. This post is everything I had to change, in the order I hit it, so you don't have to rediscover it.

Setup that finally worked:

GPU RX 7800 XT, 16 GB
Host Windows 11 + WSL2, Ubuntu 24.04
ROCm 7.2 system install inside WSL
PyTorch 2.10.0+rocm7.0
Precision fp16

1. Use WSL2, not native Windows

AMD now ships preview PyTorch wheels for Windows, and they import fine. torch.cuda.is_available() returns True, your card shows up, and you feel great for about six seconds, which is roughly how long a full FlashVSR inference ran before it segfaulted. AMD's own notes mention intermittent crashes during inference on Windows, and that matched what I saw.

The same card inside WSL2 was stable, and a lot faster: a plain fp16 matmul benchmark gave me about 50 TFLOP/s in WSL against about 16 on native Windows.

Getting the GPU visible inside WSL took a few environment variables:

export HSA_ENABLE_DXG_DETECTION=1
export ROCM_PATH=/opt/rocm-7.2.0
export LD_LIBRARY_PATH=/opt/rocm-7.2.0/lib:$LD_LIBRARY_PATH
export HSA_OVERRIDE_GFX_VERSION=11.0.1   # gfx1101; adjust for your card

One non-obvious fix: the PyTorch ROCm wheel bundles its own libhsa-runtime64.so, and that copy didn't detect the GPU through WSL's DXG layer for me. Replacing it with the system ROCm copy (keep a backup) fixed it. If rocminfo sees your card but PyTorch doesn't, check this first.

2. Make Block-Sparse-Attention optional

FlashVSR imports the kernel at the top of diffsynth/models/wan_video_dit.py. Step one is letting it fail:

try:
    from block_sparse_attn import block_sparse_attn_func
    BLOCK_SPARSE_AVAILABLE = True
except ModuleNotFoundError:
    block_sparse_attn_func = None
    BLOCK_SPARSE_AVAILABLE = False

Then, where the attention mask is used, fall back to PyTorch's scaled_dot_product_attention when the kernel isn't there:

if attention_mask is not None and not BLOCK_SPARSE_AVAILABLE:
    q = rearrange(q, "b s (n d) -> b n s d", n=num_heads)
    k = rearrange(k, "b s (n d) -> b n s d", n=num_heads)
    v = rearrange(v, "b s (n d) -> b n s d", n=num_heads)
    x = F.scaled_dot_product_attention(q, k, v)
    x = rearrange(x, "b n s d -> b s (n d)", n=num_heads)

That's dense attention. It ignores FlashVSR's sparse locality mask completely, which costs speed, and more on that below.

3. Turn on AOTriton, or SDPA will eat all your memory

On ROCm, scaled_dot_product_attention only uses a fused flash-attention kernel (AOTriton) on consumer RDNA3 cards if you ask for it:

export TORCH_ROCM_AOTRITON_ENABLE_EXPERIMENTAL=1

Without it, PyTorch quietly falls back to the math backend, which builds the full attention matrix. For FlashVSR's sequence lengths that's tens of gigabytes. You won't get a helpful error, just an out-of-memory crash, or a process that dies with no traceback at all.

If you only remember one thing from this post, make it this variable.

4. The bug that turned every frame into noise

Dense attention throws away information the model was trained with, so the obvious next step was to keep it. I expanded FlashVSR's block-level sparsity mask into a token-level boolean mask and passed it to SDPA as attn_mask. In theory that gives you the same attention pattern as the CUDA kernel, just slower.

In practice every output frame was pure latent noise. Not slightly wrong, not blurry: static.

The likely culprit is mask polarity (SDPA's boolean attn_mask means "True = attend", and it's very easy to invert the block mask's meaning when expanding it), possibly combined with how the expanded mask interacts with the fused kernel. I didn't fully chase it down, because the plain dense path produced correct, clean output and the masked one wasn't faster on this card anyway: once you hand SDPA an arbitrary mask, you lose most of the fused-kernel advantage.

So the shipped path is dense, and mask expansion stays behind an off-by-default experiment flag. If you try this yourself and get noise, turn the mask off before you start debugging anything else.

5. Memory: ROCm can't defragment, so don't let it fragment

The NVIDIA advice for long diffusion runs is PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True. That wasn't available to me on ROCm, and FlashVSR's streaming loop fragments VRAM badly without it. Eventually even a 176 MiB allocation fails on a card with gigabytes "free".

Three patches fixed it.

Offload the DiT before decoding. FlashVSR's TCDecoder runs after the diffusion step in each chunk, with the DiT still resident. Moving it off first gives the decoder room:

self.load_models_to_device([])
torch.cuda.empty_cache()
cur_frames = self.TCDecoder.decode_video(...)
self.load_models_to_device(['dit'])

Stream decoded frames to the CPU. The decoder accumulates finished frames on the GPU and stacks them at the end. At 1280×960 that's around 3 GB of extra residency, scattered across the heap. Moving each finished frame to host memory as it's produced keeps the GPU heap compact:

out.append(xt.to("cpu", non_blocking=False))

Clean up after every streaming segment. Delete the per-segment tensors and empty the cache before the next DiT step, so the allocator can hand back contiguous blocks:

del cur_frames, cur_latents, noise_pred_posi, cur_LQ_frame
gc.collect()
torch.cuda.empty_cache()

None of these change the output. They just stop the run from dying halfway through.

6. Pick a resolution the card can actually do

With all of the above, 16 GB comfortably handles 2× upscaling from an input around 320 to 384 pixels wide. A 512-wide input ran one streaming step and then died without a traceback. Some other knobs I tried:

  • fp16 over bf16. bf16 was less reliable on RDNA3 in my runs.
  • DiT tiling (tile_dit) fits larger inputs into memory, but was about 8× slower. I only use it as a last resort.
  • For low-res sources, keep the native width. For a 320×240 clip, I found it better to shorten the chunks (2 s down to 1 s) than to downscale to 256 wide first. Downscaling a source that small throws away detail you can't get back.

7. Don't trust WSL's idea of how much RAM you have

Inside WSL, psutil reports the VM's memory cap, not the machine's. On a 32 GB host, WSL showed about 10 GB, and my memory-tuning code treated it as a small machine. It reloaded the model every two chunks to be safe, which roughly doubled the run time.

Two fixes:

  • Raise the VM limit in %UserProfile%\.wslconfig (memory=12GB on a 32 GB host), then run wsl --shutdown once.
  • Pass the real host RAM in from the Windows side as an environment variable, and size things from that.

Keeping one model process warm across chunks, and restarting it every 12 or so instead of every 2, was the single biggest speedup after getting attention right.

How fast is it?

My first working run took about 230 seconds for 2 seconds of 384×384 video upscaled to 768×768. After the memory and warm-worker fixes, a realistic number on the 7800 XT is about a minute of processing per second of 320 to 360p source.

That's fine for clips and highlights. A full 2-hour tape is not a reasonable job for this card, and dense attention is the reason. NVIDIA cards with the real sparse kernel are much faster.

The quality catch: diffusion warps on motion

Getting FlashVSR to run is half the story. On real footage, especially fast motion like a kick or a head turn, pure FlashVSR output tends to warp: edges drift a few pixels from where they are in the source, and faces can subtly change shape between frames.

What worked better for me was using it as a detail donor rather than the final picture:

  1. Run a conventional super-resolution model (Real-ESRGAN class) on the source. It's less detailed, but its geometry stays put.
  2. Run FlashVSR and line its output up with that base.
  3. Take only the high-frequency detail from FlashVSR, and only where the two agree at low frequency. Where they disagree (the warp), fall back to the base.
  4. Smooth that detail layer over time a little, so it doesn't shimmer.

The result keeps FlashVSR's texture and the source's shape. It also gets rid of the bright "sticker outline" that shows up around moving people when you sharpen a slightly misaligned diffusion output.

Summary

If you want FlashVSR on an AMD card:

  1. Use WSL2 with a system ROCm install, and swap in the system libhsa-runtime64.so if PyTorch can't see the GPU.
  2. Make Block-Sparse-Attention optional and fall back to dense SDPA.
  3. Set TORCH_ROCM_AOTRITON_ENABLE_EXPERIMENTAL=1.
  4. Don't pass an expanded sparsity mask to SDPA unless you enjoy static.
  5. Offload the DiT before decoding, move decoded frames to the CPU, and clean up after every segment.
  6. Stay around 320 to 384 px input width on 16 GB, in fp16.
  7. Tell your code the real host RAM, not WSL's.

I packaged all of this into Phoenix Upscaler, a desktop app for restoring old video (tapes, DVDs, early phone footage), where the AMD generative mode sets up WSL and ROCm for you. The setup guide covers what it installs, and there's a longer write-up on restoring old video in general. But everything above works on a plain FlashVSR-Pro checkout too.

If you get it running on a different RDNA card (a 7900 XTX, or anything RDNA2), I'd like to hear what resolution it handles.

📰 Read the original article on Dev.to AI

Originally published by Dev.to AI. Aggregated on AIWithGhost for educational purposes — full credit and traffic to the original publisher.