TheRock/rocm-libraries/dnn-providers/hip-kernel-provider/rocke/library/tests/test_gfx1250_attention.py

# Copyright (c) Advanced Micro Devices, Inc., or its affiliates.
# SPDX-License-Identifier: MIT
"""gfx1250 unified-attention tests (no-GPU lowering + GPU probes).

Covers the wave32 WMMA attention surface:
  * the standalone dense WMMA attention forward;
  * the tiled-2D prefill kernel (gate / lower / dispatch);
  * the split-KV tiled-3D decode kernel (gate / lower / dispatch /
    multi-wave regression);
  * the scalar fp8-KV unified-attention leg;
  * Qwen3-30B-A3B decode/prefill routing through the unified instance and the
    KV-cache dequant / append-RoPE kernels.

The GPU probe class is skipped unless run on a gfx1250 device.

    PYTHONPATH=python python3 -m pytest rocke/tests/test_gfx1250_attention.py -k "not Gpu"
"""

from __future__ import annotations

import os
import subprocess
import sys
import unittest

from rocke.runtime.hip_module import get_device_arch

# Subprocess `-m` runs (rocke.*, kernels.*, builders.*) resolve through the
# editable-installed rocke + library packages; no PYTHONPATH wiring needed
# (see rocke/BUILDING.md).

# Running device's gfx arch via the rocke HIP runtime (no torch dependency).
_ARCH = get_device_arch(0)


class TestGfx1250DenseAttention(unittest.TestCase):
    def test_attention_fwd_lowers(self):
        # The standalone gfx1250 dense WMMA attention forward (BLOCK_K=32) must
        # build and lower to the K=32 WMMA intrinsic. H=64 -> 2 QK d-tiles x 2
        # N-sub-tiles + 4 PV d-tiles = 8 WMMA calls per K-loop iteration.
        from rocke.core.lower_llvm import lower_kernel_to_llvm
        from kernels.gfx1250.wmma_attention_fwd import (
            WmmaAttentionFwdSpec,
            build_wmma_attention_fwd,
        )

        spec = WmmaAttentionFwdSpec(head_size=64, num_query_heads=4)
        ll = lower_kernel_to_llvm(
            build_wmma_attention_fwd(spec, arch="gfx1250"), arch="gfx1250"
        )
        self.assertIn("llvm.amdgcn.wmma.f32.16x16x32.f16.v8f32.v16f16", ll)


class TestGfx1250TiledAttention2D(unittest.TestCase):
    def test_gate_accepts_fp8_rejects_bf16_kv(self):
        from kernels.gfx1250.attention_tiled_2d import supports_tiled_2d

        ok, why = supports_tiled_2d(
            head_size=64,
            block_size=32,
            dtype="bf16",
            num_queries_per_kv=8,
            use_alibi=False,
            use_qq_bias=False,
            use_fp8=True,
            q_dtype=None,
            num_warps=1,
            block_m_per_warp=16,
            kv_storage_dtype="fp8e4m3",
            tile_size=32,
            arch="gfx1250",
        )
        self.assertTrue(ok, why)

        ok, why = supports_tiled_2d(
            head_size=64,
            block_size=32,
            dtype="bf16",
            num_queries_per_kv=8,
            use_alibi=False,
            use_qq_bias=False,
            use_fp8=False,
            q_dtype=None,
            num_warps=1,
            block_m_per_warp=16,
            kv_storage_dtype=None,
            tile_size=32,
            arch="gfx1250",
        )
        self.assertFalse(ok)
        self.assertIn("fp8e4m3", why)

    def test_lowers_to_wmma_and_fp8_dequant(self):
        from rocke.core.lower_llvm import lower_kernel_to_llvm
        from kernels.gfx1250.attention_tiled_2d import (
            UnifiedAttention2DTiledSpec,
            build_unified_attention_2d_tiled,
        )

        spec = UnifiedAttention2DTiledSpec(
            head_size=64,
            block_size=32,
            num_query_heads=64,
            num_kv_heads=8,
            dtype="bf16",
            use_sinks=True,
            sliding_window=128,
            has_softcap=False,
            num_seqs=4,
            kv_storage_dtype="fp8e4m3",
            tile_size=32,
        )
        ll = lower_kernel_to_llvm(
            build_unified_attention_2d_tiled(spec, arch="gfx1250"), arch="gfx1250"
        )
        self.assertIn("llvm.amdgcn.wmma.f32.16x16x32.bf16.v8f32.v16bf16", ll)
        self.assertIn("llvm.amdgcn.cvt.pk.f32.fp8", ll)

    def test_common_dispatch_selects_wave32_tiled_meta(self):
        from kernels.common import attention_unified as au

        old_arch = au._RESOLVED_ATTENTION_ARCH
        au._RESOLVED_ATTENTION_ARCH = "gfx1250"
        try:
            problem = au.UnifiedAttentionProblem(
                total_q=4,
                num_seqs=4,
                num_query_heads=64,
                num_kv_heads=8,
                head_size=64,
                block_size=32,
                max_seqlen_q=1,
                max_seqlen_k=256,
                dtype="bf16",
                use_sinks=True,
                sliding_window=128,
                use_fp8=True,
            )
            ok, why = au.supports_native_unified_attention_tiled(problem)
            self.assertTrue(ok, why)
            spec = au._tiled_spec_from_problem(problem)
            self.assertEqual(spec.kernel_name().split("_")[2], "tiled")
            self.assertEqual(spec.num_warps, 1)
            self.assertEqual(spec.block_q, 2)
            meta = au._get_2d_launch_meta(problem, au._tiled_cache_key(problem))
            self.assertEqual(meta.block, (32, 1, 1))
            self.assertEqual(meta.grid, (8, 6, 1))
        finally:
            au._RESOLVED_ATTENTION_ARCH = old_arch


class TestGfx1250TiledAttention3D(unittest.TestCase):
    def test_gate_and_lowers(self):
        from rocke.core.lower_llvm import lower_kernel_to_llvm
        from kernels.gfx1250.attention_tiled_3d import (
            UnifiedAttention3DTiledSpec,
            UnifiedAttentionReduceTiledSpec,
            build_unified_attention_3d_tiled,
            build_unified_attention_reduce_tiled,
            supports_tiled_3d,
        )

        ok, why = supports_tiled_3d(
            head_size=64,
            block_size=16,
            dtype="bf16",
            num_queries_per_kv=8,
            use_alibi=False,
            use_qq_bias=False,
            use_fp8=True,
            q_dtype=None,
            kv_storage_dtype="fp8e4m3",
            arch="gfx1250",
        )
        self.assertTrue(ok, why)
        bad, why_bad = supports_tiled_3d(
            head_size=64,
            block_size=64,
            dtype="bf16",
            num_queries_per_kv=8,
            use_alibi=False,
            use_qq_bias=False,
            use_fp8=False,
            q_dtype=None,
            kv_storage_dtype=None,
            arch="gfx1250",
        )
        self.assertFalse(bad)
        self.assertIn("block_size", why_bad)

        seg = UnifiedAttention3DTiledSpec(
            head_size=64,
            block_size=16,
            num_query_heads=32,
            num_kv_heads=4,
            dtype="bf16",
            use_sinks=True,
            sliding_window=0,
            has_softcap=False,
            num_segments=16,
            num_seqs=2,
            kv_storage_dtype="fp8e4m3",
        )
        ll_seg = lower_kernel_to_llvm(
            build_unified_attention_3d_tiled(seg, arch="gfx1250"), arch="gfx1250"
        )
        self.assertIn("llvm.amdgcn.wmma.f32.16x16x32.bf16.v8f32.v16bf16", ll_seg)
        self.assertIn("llvm.amdgcn.cvt.pk.f32.fp8", ll_seg)
        red = UnifiedAttentionReduceTiledSpec(
            head_size=64,
            num_query_heads=32,
            num_kv_heads=4,
            dtype="bf16",
            num_segments=16,
        )
        ll_red = lower_kernel_to_llvm(
            build_unified_attention_reduce_tiled(red, arch="gfx1250"), arch="gfx1250"
        )
        self.assertIn("llvm.amdgcn.s.barrier", ll_red)

    def test_large_batch_regression_binary_search_floor(self):
        # Regression: large-batch decode binary search uses 32+ iterations
        # (the historically flaky 256-seq fp8 shape).
        from rocke.core.lower_llvm import lower_kernel_to_llvm
        from kernels.gfx1250.attention_tiled_3d import (
            UnifiedAttention3DTiledSpec,
            build_unified_attention_3d_tiled,
        )

        seg256 = UnifiedAttention3DTiledSpec(
            head_size=64,
            block_size=16,
            num_query_heads=32,
            num_kv_heads=4,
            dtype="bf16",
            use_sinks=False,
            sliding_window=0,
            has_softcap=False,
            num_segments=16,
            num_seqs=256,
            kv_storage_dtype="fp8e4m3",
        )
        self.assertGreaterEqual(seg256.binary_search_iters, 32)
        ll_256 = lower_kernel_to_llvm(
            build_unified_attention_3d_tiled(seg256, arch="gfx1250"), arch="gfx1250"
        )
        self.assertIn("llvm.amdgcn.wmma.f32.16x16x32.bf16.v8f32.v16bf16", ll_256)

    def test_multi_wave_widths_lower(self):
        # Cooperative multi-wave32 CTA (LDS inter-wave reduction) lowers for
        # every supported width.
        from rocke.core.lower_llvm import lower_kernel_to_llvm
        from kernels.gfx1250.attention_tiled_3d import (
            UnifiedAttention3DTiledSpec,
            build_unified_attention_3d_tiled,
        )

        for waves in (1, 2, 4, 8):
            spec = UnifiedAttention3DTiledSpec(
                head_size=64,
                block_size=16,
                num_query_heads=32,
                num_kv_heads=4,
                dtype="bf16",
                use_sinks=True,
                sliding_window=0,
                has_softcap=False,
                num_segments=16,
                num_seqs=256,
                kv_storage_dtype="fp8e4m3",
                num_waves=waves,
                use_wide_lds_reads=(waves == 1),
            )
            ll = lower_kernel_to_llvm(
                build_unified_attention_3d_tiled(spec, arch="gfx1250"), arch="gfx1250"
            )
            self.assertIn("llvm.amdgcn.wmma.f32.16x16x32.bf16.v8f32.v16bf16", ll)

    def test_dtla_prefetch_bf16_lowers(self):
        # DTLA bf16: async global->LDS V staging double-buffer + prefetch.
        # gfx1250 uses global_load_async_to_lds_b128 + the dedicated ASYNC
        # counter (s_wait_asynccnt), NOT the gfx9 buffer/global load-to-LDS
        # intrinsics (which do not select on gfx1250). The 32-token tile spans
        # 2 paged blocks at block_size=16, 1 at 32; both must lower (and
        # compile -- see test_dtla_prefetch_bf16_compiles).
        from rocke.core.lower_llvm import lower_kernel_to_llvm
        from kernels.gfx1250.attention_tiled_3d import (
            UnifiedAttention3DTiledSpec,
            build_unified_attention_3d_tiled,
        )

        for bs in (16, 32):
            spec = UnifiedAttention3DTiledSpec(
                head_size=64,
                block_size=bs,
                num_query_heads=32,
                num_kv_heads=4,
                dtype="bf16",
                use_sinks=True,
                sliding_window=0,
                has_softcap=False,
                num_segments=16,
                num_seqs=2,
                kv_storage_dtype="bf16",
                use_dtla_prefetch=True,
                use_wide_lds_reads=False,
            )
            self.assertIn("dtla", spec.kernel_name())
            ll = lower_kernel_to_llvm(
                build_unified_attention_3d_tiled(spec, arch="gfx1250"), arch="gfx1250"
            )
            self.assertIn("llvm.amdgcn.wmma.f32.16x16x32.bf16.v8f32.v16bf16", ll)
            self.assertIn("llvm.amdgcn.global.load.async.to.lds.b128", ll)
            self.assertIn("llvm.amdgcn.s.wait.asynccnt", ll)
            # The gfx9 DirectToLDS intrinsics must NOT appear (not selectable).
            self.assertNotIn("raw.ptr.buffer.load.lds", ll)
            self.assertNotIn("global.load.lds", ll)

    def test_dtla_prefetch_bf16_compiles(self):
        # Regression guard: lowering to LLVM IR is NOT enough -- the gfx9
        # buffer/global load-to-LDS intrinsics lower cleanly but FAIL to select
        # on gfx1250. Compile to a real code object so a future regression to a
        # non-selectable async path is caught here, not on the GPU box.
        from kernels.gfx1250.attention_tiled_3d import (
            UnifiedAttention3DTiledSpec,
            build_unified_attention_3d_tiled,
        )

        try:
            from rocke.helpers.compile import compile_kernel
        except Exception as e:  # pragma: no cover - env-dependent
            self.skipTest(f"compile toolchain unavailable: {e}")

        for bs in (16, 32):
            spec = UnifiedAttention3DTiledSpec(
                head_size=64,
                block_size=bs,
                num_query_heads=32,
                num_kv_heads=4,
                dtype="bf16",
                use_sinks=True,
                sliding_window=0,
                has_softcap=False,
                num_segments=16,
                num_seqs=2,
                kv_storage_dtype="bf16",
                use_dtla_prefetch=True,
                use_wide_lds_reads=False,
            )
            try:
                art = compile_kernel(
                    build_unified_attention_3d_tiled(spec, arch="gfx1250"),
                    arch="gfx1250",
                )
            except Exception as e:  # pragma: no cover - env-dependent
                self.skipTest(f"gfx1250 comgr compile unavailable: {e}")
            self.assertGreater(art.hsaco_bytes, 0)

    def test_dtla_prefetch_rejects_incompatible_levers(self):
        # DTLA owns the V_lds layout + pipeline; guard the mutually exclusive
        # combinations and the not-yet-implemented fp8 path.
        from kernels.gfx1250.attention_tiled_3d import (
            UnifiedAttention3DTiledSpec,
        )

        base = dict(
            head_size=64,
            block_size=16,
            num_query_heads=32,
            num_kv_heads=4,
            dtype="bf16",
            use_sinks=False,
            sliding_window=0,
            has_softcap=False,
            num_segments=16,
            num_seqs=2,
            kv_storage_dtype="bf16",
            use_dtla_prefetch=True,
            use_wide_lds_reads=False,
        )
        with self.assertRaises(ValueError):
            UnifiedAttention3DTiledSpec(**{**base, "use_wide_lds_reads": True})
        with self.assertRaises(ValueError):
            UnifiedAttention3DTiledSpec(**{**base, "num_waves": 2})
        with self.assertRaises(ValueError):
            UnifiedAttention3DTiledSpec(**{**base, "use_register_p": True})
        with self.assertRaises(ValueError):
            UnifiedAttention3DTiledSpec(**{**base, "kv_storage_dtype": "fp8e4m3"})

    def test_common_dispatch_selects_3d(self):
        from kernels.common import attention_unified as au

        old_arch = au._RESOLVED_ATTENTION_ARCH
        au._RESOLVED_ATTENTION_ARCH = "gfx1250"
        try:
            problem = au.UnifiedAttentionProblem(
                total_q=2,
                num_seqs=2,
                num_query_heads=32,
                num_kv_heads=4,
                head_size=64,
                block_size=16,
                max_seqlen_q=1,
                max_seqlen_k=1024,
                dtype="bf16",
                q_dtype="bf16",
                sliding_window=0,
                use_sinks=True,
                use_fp8=True,
            )
            ok, why = au.supports_native_unified_attention_3d_tiled(problem)
            self.assertTrue(ok, why)
            spec = au._tiled_3d_spec_from_problem(problem)
            self.assertEqual(spec.block_q, 2)
            self.assertGreaterEqual(spec.num_segments, 1)
            Spec, _, _, _, _ = au._tiled_3d_impl("gfx1250")
            self.assertEqual(Spec.__module__, "kernels.gfx1250.attention_tiled_3d")
        finally:
            au._RESOLVED_ATTENTION_ARCH = old_arch


class TestGfx1250ScalarFp8Attention(unittest.TestCase):
    @staticmethod
    def _small_fp8_problem(**overrides):
        from kernels.common.attention_unified import UnifiedAttentionProblem

        kwargs = dict(
            total_q=8,
            num_seqs=1,
            num_query_heads=4,
            num_kv_heads=1,
            head_size=64,
            block_size=16,
            max_seqlen_q=8,
            max_seqlen_k=16,
            dtype="bf16",
            q_dtype="bf16",
            use_fp8=True,
        )
        kwargs.update(overrides)
        return UnifiedAttentionProblem(**kwargs)

    def test_scalar_gate_accepts_fp8_kv_bf16_query(self):
        from kernels.common.attention_unified import (
            supports_native_unified_attention,
        )

        ok, reason = supports_native_unified_attention(self._small_fp8_problem())
        self.assertTrue(ok, reason)
        self.assertIn("supported", reason)

    def test_scalar_gate_keeps_unsupported_biases_rejected(self):
        from kernels.common.attention_unified import (
            supports_native_unified_attention,
        )

        ok_alibi, reason_alibi = supports_native_unified_attention(
            self._small_fp8_problem(use_alibi=True)
        )
        ok_qq, reason_qq = supports_native_unified_attention(
            self._small_fp8_problem(use_qq_bias=True)
        )
        self.assertFalse(ok_alibi)
        self.assertIn("ALiBi", reason_alibi)
        self.assertFalse(ok_qq)
        self.assertIn("QQ bias", reason_qq)

    def test_scalar_fp8_kv_lowers_dequant_and_scale(self):
        from rocke.core.lower_llvm import lower_kernel_to_llvm
        from kernels.common.attention_unified import (
            UnifiedAttention2DSpec,
            build_unified_attention_2d,
        )

        kernel = build_unified_attention_2d(
            UnifiedAttention2DSpec(self._small_fp8_problem())
        )
        ll = lower_kernel_to_llvm(kernel, arch="gfx1250")
        self.assertIn("@llvm.amdgcn.cvt.f32.fp8", ll)
        self.assertIn("%k_scale", ll)
        self.assertIn("%v_scale", ll)


class TestGfx1250Qwen3AttentionRouting(unittest.TestCase):
    """Qwen3 decode/prefill route through the unified attention instance."""

    @staticmethod
    def _decode_problem(shape):
        from kernels.common.attention_unified import UnifiedAttentionProblem

        return UnifiedAttentionProblem(
            total_q=shape.total_q,
            num_seqs=shape.num_seqs,
            num_query_heads=shape.num_query_heads,
            num_kv_heads=shape.num_kv_heads,
            head_size=shape.head_size,
            block_size=shape.block_size,
            max_seqlen_q=shape.max_seqlen_q,
            max_seqlen_k=shape.max_seqlen_k,
            dtype=shape.dtype,
            q_dtype=shape.dtype,
            use_sinks=False,
            use_fp8=shape.kv_storage_dtype in ("fp8e4m3", "bf8e5m2"),
        )

    def test_decode_3d_routes_through_unified_split_kv(self):
        from rocke.core.lower_llvm import lower_kernel_to_llvm
        from rocke.examples.gfx1250.qwen3_30b_a3b.qwen3_30b_a3b_shapes import (
            qwen3_decode_attention_shapes,
        )
        from kernels.common import attention_unified as au

        old_arch = au._RESOLVED_ATTENTION_ARCH
        au._RESOLVED_ATTENTION_ARCH = "gfx1250"
        try:
            for kv_storage in ("bf16", "fp8e4m3"):
                shapes = qwen3_decode_attention_shapes(kv_storage_dtype=kv_storage)
                for shape in shapes:
                    problem = self._decode_problem(shape)
                    ok, why = au.supports_native_unified_attention_3d_tiled(problem)
                    self.assertTrue(ok, f"{kv_storage} kv{shape.kv_len}: {why}")
                    (_SegSpec, ReduceSpec, build_seg, build_red, _) = au._tiled_3d_impl(
                        "gfx1250"
                    )
                    seg_spec = au._tiled_3d_spec_from_problem(problem)
                    ll = lower_kernel_to_llvm(
                        build_seg(seg_spec, arch="gfx1250"), arch="gfx1250"
                    )
                    self.assertIn(
                        "llvm.amdgcn.wmma.f32.16x16x32.bf16.v8f32.v16bf16", ll
                    )
                    if kv_storage == "fp8e4m3":
                        self.assertIn("llvm.amdgcn.cvt.pk.f32.fp8", ll)
                    red_spec = ReduceSpec(
                        head_size=shape.head_size,
                        num_query_heads=shape.num_query_heads,
                        num_kv_heads=shape.num_kv_heads,
                        dtype=shape.dtype,
                        num_segments=seg_spec.num_segments,
                    )
                    lower_kernel_to_llvm(
                        build_red(red_spec, arch="gfx1250"), arch="gfx1250"
                    )
        finally:
            au._RESOLVED_ATTENTION_ARCH = old_arch

    def test_prefill_2d_routes_through_unified_scalar(self):
        from rocke.core.lower_llvm import lower_kernel_to_llvm
        from rocke.examples.gfx1250.qwen3_30b_a3b.qwen3_30b_a3b_shapes import (
            qwen3_prefill_attention_shapes,
        )
        from kernels.common.attention_unified import (
            UnifiedAttention2DSpec,
            UnifiedAttentionProblem,
            build_unified_attention_2d,
            supports_native_unified_attention,
        )

        for shape in qwen3_prefill_attention_shapes():
            problem = UnifiedAttentionProblem(
                total_q=shape.total_q,
                num_seqs=shape.num_seqs,
                num_query_heads=shape.num_query_heads,
                num_kv_heads=shape.num_kv_heads,
                head_size=shape.head_size,
                block_size=shape.block_size,
                max_seqlen_q=shape.max_seqlen_q,
                max_seqlen_k=shape.max_seqlen_k,
                dtype=shape.dtype,
                use_sinks=False,
            )
            ok, why = supports_native_unified_attention(problem)
            self.assertTrue(ok, f"prefill q{shape.q_len}: {why}")
            ll = lower_kernel_to_llvm(
                build_unified_attention_2d(
                    UnifiedAttention2DSpec(
                        problem=problem, name="rocke_gfx1250_qwen3_prefill2d_scalar"
                    )
                ),
                arch="gfx1250",
            )
            self.assertIn("rocke_gfx1250_qwen3_prefill2d_scalar", ll)
            self.assertIn("define amdgpu_kernel void", ll)
            self.assertIn("addrspace(1)", ll)


class TestGfx1250Qwen3KvCache(unittest.TestCase):
    def test_fp8_kv_dequant_lowers_with_explicit_scale_multiply(self):
        from rocke.core.lower_llvm import lower_kernel_to_llvm
        from rocke.instances.gfx1250.qwen3_kv_cache import (
            Qwen3KvDequantSpec,
            build_qwen3_kv_dequant_smoke,
        )

        ll = lower_kernel_to_llvm(
            build_qwen3_kv_dequant_smoke(
                Qwen3KvDequantSpec(kv_storage_dtype="fp8e4m3"),
                arch="gfx1250",
            ),
            arch="gfx1250",
        )
        self.assertIn("llvm.amdgcn.cvt.pk.f32.fp8", ll)
        self.assertIn("fmul float", ll)
        self.assertNotIn("cvt.scalef32.pk.f32.fp8", ll)

    def test_bf8_kv_dequant_lowers_with_explicit_scale_multiply(self):
        from rocke.core.lower_llvm import lower_kernel_to_llvm
        from rocke.instances.gfx1250.qwen3_kv_cache import (
            Qwen3KvDequantSpec,
            build_qwen3_kv_dequant_smoke,
        )

        ll = lower_kernel_to_llvm(
            build_qwen3_kv_dequant_smoke(
                Qwen3KvDequantSpec(kv_storage_dtype="bf8e5m2"),
                arch="gfx1250",
            ),
            arch="gfx1250",
        )
        self.assertIn("llvm.amdgcn.cvt.pk.f32.bf8", ll)
        self.assertIn("fmul float", ll)
        self.assertNotIn("cvt.scalef32.pk.f32.bf8", ll)

    def test_kv_append_rope_lowers_for_bf16_fp8_bf8_storage(self):
        from rocke.core.lower_llvm import lower_kernel_to_llvm
        from rocke.instances.gfx1250.qwen3_kv_cache import (
            Qwen3KvAppendRopeSpec,
            build_qwen3_kv_append_rope,
        )

        expectations = {
            "bf16": "store bfloat",
            "fp8e4m3": "llvm.amdgcn.cvt.pk.fp8.f32",
            "bf8e5m2": "llvm.amdgcn.cvt.pk.bf8.f32",
        }
        for kv_storage_dtype, needle in expectations.items():
            ll = lower_kernel_to_llvm(
                build_qwen3_kv_append_rope(
                    Qwen3KvAppendRopeSpec(kv_storage_dtype=kv_storage_dtype),
                    arch="gfx1250",
                ),
                arch="gfx1250",
            )
            self.assertIn("rocke_gfx1250_qwen3_kv_append_rope", ll)
            self.assertIn("fsub float", ll)
            self.assertIn("fadd float", ll)
            self.assertIn(needle, ll)


@unittest.skipUnless(
    _ARCH == "gfx1250",
    f"gfx1250 WMMA GPU probe; running on {_ARCH!r} (needs a gfx1250 device)",
)
class TestGfx1250Gpu(unittest.TestCase):
    """GPU probe: build + numerically verify WMMA GEMM / attention on hardware."""

    def _run(self, module, *args, timeout=300):
        env = dict(os.environ)
        env["PYTHONDONTWRITEBYTECODE"] = "1"
        r = subprocess.run(
            [sys.executable, "-m", module, *args],
            env=env,
            capture_output=True,
            text=True,
            timeout=timeout,
        )
        return r.returncode, (r.stdout + r.stderr)

    def test_wmma_probe_16x16x32(self):
        rc, out = self._run(
            "rocke.examples.gfx1250.wmma_probe", "--m", "16", "--n", "16", "--k", "32"
        )
        self.assertEqual(rc, 0, f"gfx1250 WMMA probe 16x16x32 failed:\n{out[-2500:]}")
        self.assertIn("PASS", out)

    def test_wmma_probe_64x64x64(self):
        rc, out = self._run(
            "rocke.examples.gfx1250.wmma_probe", "--m", "64", "--n", "64", "--k", "64"
        )
        self.assertEqual(rc, 0, f"gfx1250 WMMA probe 64x64x64 failed:\n{out[-2500:]}")
        self.assertIn("PASS", out)

    def test_attention_fwd_verify(self):
        rc, out = self._run(
            "builders.gfx1250.attention.wmma_attention_fwd_verify",
            "--seqlen-q",
            "64",
            "--seqlen-k",
            "64",
            "--head-size",
            "64",
            "--heads",
            "4",
        )
        self.assertEqual(rc, 0, f"gfx1250 attention fwd failed:\n{out[-2500:]}")
        self.assertIn("PASS", out)

    def test_attention_fwd_causal(self):
        # Causal currently verifies through the HIP/clang path. The direct LLVM
        # path still has a WMMA scheduling/codegen issue that produces
        # deterministic NaNs; hipcc emits a correct code object for the same IR.
        rc, out = self._run(
            "builders.gfx1250.attention.wmma_attention_fwd_verify",
            "--seqlen-q",
            "64",
            "--seqlen-k",
            "64",
            "--head-size",
            "64",
            "--heads",
            "4",
            "--causal",
        )
        self.assertEqual(rc, 0, f"gfx1250 causal attention fwd failed:\n{out[-2500:]}")
        self.assertIn("PASS", out)


if __name__ == "__main__":
    unittest.main(verbosity=2)