BEAM=2 hangs the GPU on gfx1201 (RDNA4) with bf16 matmul, requiring driver reset

BEAM search on a bf16 matmul hangs the GPU on gfx1201. The kernel driver
detects a hung queue and performs a MODE1 reset, losing VRAM. BEAM=0 on the
identical workload completes in 1.4s.

Reproduced 4x from clean state, on both 5b05106 and current master 0319b1e.

Repro

import time
from tinygrad import Tensor, Device, dtypes
from tinygrad.helpers import Context

N = 2048
for beam in (0, 2):
    print(f"BEAM={beam} ...", end=" ", flush=True)
    t0 = time.perf_counter()
    try:
        with Context(TC=1, BEAM=beam):
            a = Tensor.rand(N, N).cast(dtypes.bfloat16).contiguous().realize()
            b = Tensor.rand(N, N).cast(dtypes.bfloat16).contiguous().realize()
            (a @ b).realize()
        print(f"OK in {time.perf_counter()-t0:.1f}s")
    except Exception as e:
        print(f"FAILED after {time.perf_counter()-t0:.1f}s: {type(e).__name__}: {e}")

$ AMD=1 python3 repro.py
BEAM=0 ... OK in 1.4s
BEAM=2 ... FAILED after 30.5s: RuntimeError: Wait timeout: 30000 ms! (the signal is not set to 27, but 26)

Raised from runtime/support/hcq.py:287 in HCQSignal.wait, via
HCQCompiled.synchronize at hcq.py:427.

dmesg

amdgpu 0000:23:00.0: GPU reset begin!. Source: 3
amdgpu 0000:23:00.0: MES(1) failed to respond to msg=REMOVE_QUEUE
amdgpu 0000:23:00.0: failed to unmap legacy queue
amdgpu 0000:23:00.0: MODE1 reset
amdgpu 0000:23:00.0: GPU reset succeeded, trying to resume
amdgpu 0000:23:00.0: VRAM is lost due to GPU reset!
amdgpu 0000:23:00.0: GPU reset(3) succeeded!
amdgpu 0000:23:00.0: [drm] device wedged, but recovered through reset

The GPU recovers each time; no reboot required. Any concurrent GPU work dies.

Environment

GPU        AMD Radeon RX 9070 XT (device 0x7550, gfx1201 / RDNA4)
renderer   AMDLLVMRenderer
tensor     WMMA_16_16_16_half_float, WMMA_16_16_16_half_half,
 cores     WMMA_16_16_16___bf16_float, WMMA_16_16_16___bf16___bf16
ROCm       7.2.4
kernel     7.0.0-30-generic
OS         Ubuntu 24.04.4 LTS
Python     3.12.3
tinygrad   master 0319b1e (also reproduces on 5b05106)

Secondary: TC=1 is slower than TC=0 for half precision

Possibly the same root cause, so noting rather than filing separately. Same
2048x2048 matmul, BEAM=0, TFLOP/s from wall clock and 2*n^3 (not
GlobalCounters), median of 5, fresh operands each iteration:

fp32 is unaffected, as expected. Both half precisions are meaningfully slower
with tensor cores enabled — and TC defaults to 1, so this is the out-of-box
path on this hardware. AMD documents ~191 TFLOP/s dense bf16 WMMA for RDNA4;
we measure 2.24.

Note

test/opt/test_tensor_cores.py:84 still carries
# TODO: don't skip bf16 for real device (METAL, AMD) on master, so this path
appears to be untested on real AMD hardware — consistent with it going
unnoticed.

Not tested

• Only AMDLLVMRenderer. AMDDevice is constructed with
  [HIPRenderer, AMDLLVMRenderer, HIPCCRenderer] and selects automatically;
  we found no env var to force a different one (AMD_LLVM=0 had no effect),
  so we could not check whether HIP/HIPCC are also affected.
• Only gfx1201. No other RDNA4 or RDNA3 part available to compare.
• Did not bisect which of the 44 commits between 5b05106 and 0319b1e matter,
  since it reproduces on both.
• Did not try TC_SELECT / TC_OPT to narrow which WMMA config hangs.

Run the supplied Python reproducer with BEAM=0 and BEAM=2 on the AMDLLVMRenderer, then inspect runtime/support/hcq.py at HCQSignal.wait and HCQCompiled.synchronize. Review test/opt/test_tensor_cores.py:84 and the BEAM and bf16 matmul paths to identify the failing configuration. Done means the workload no longer hangs or resets the GPU, with coverage for the affected real-device path.