fix(rocke/attention): exclude gfx942 D128 from sliced-K ring (correctness)

dnn-providers/hip-kernel-provider/rocke/library/kernels/common/attention_unified.py

@@ -1763,18 +1763,24 @@ def _enable_gfx942_flash_mask_limit(problem: UnifiedAttentionProblem) -> bool:
 
 
 def _enable_gfx942_flash_k_sliced_ring(problem: UnifiedAttentionProblem) -> bool:
-    # The sliced-K ring (32-wide K slices -> k_groups = HD/32) wins on BOTH head
-    # sizes: D128 (k_groups=4) and D64 (k_groups=2). Measured T=64+ring+cfvst+
-    # mask-limit (nw4) vs the prior per-head bests: D64 13-17% faster (beats Torch
-    # at S2048, ~parity elsewhere); D128 beats Torch S2048/S4096. So D64 and D128
-    # prefill now share the ring path.
-    # bf16 D128 now shares the ring path: the earlier exclusion attached the ring
-    # to the non-ring bf16-WIDE geometry (nw=2, no cfvst), which the ring cannot
-    # use -- the ring requires the conflict-free-V store (cfvst) and the wide nw=4
-    # flash geometry. The spec builder's ring branch instead uses the fp16-flash
-    # geometry (nw=4, tile=64, cfvst) -- verified numerically correct on gfx942
-    # (max_abs 0.00049, no NaN/Inf, GQA + MHA) on both the Python and C++ engines,
-    # with the byte-identity gate GREEN. So D64 and D128 bf16 prefill share the ring.
+    # The sliced-K ring (32-wide K slices -> k_groups = HD/32) is correctness-
+    # verified only for D64 (k_groups=2): measured T=64+ring+cfvst+mask-limit (nw4)
+    # is 13-17% faster than the prior D64 best (beats Torch at S2048, ~parity
+    # elsewhere).
+    #
+    # D128 (k_groups=4) is EXCLUDED for BOTH dtypes. The k_groups=4 ring produces
+    # numerically wrong results at realistic input magnitude -- max_abs ~0.5-1.3
+    # vs the fp32 oracle (bf16 AND fp16, GQA), even at S512 when the ring is forced
+    # -- while the non-ring D128 path (nw2/nw4, T=64, no ring) passes at max_abs
+    # ~0.0156 (bf16) / ~0.002 (fp16) on every S. #9198's "max_abs 0.00049 verified"
+    # claim was measured against a uniform_(-0.1,0.1) oracle whose near-uniform
+    # softmax never exercises the online-softmax running-max rescale across the four
+    # K slices, so the defect was masked; it reproduces immediately with randn
+    # (unit-variance) inputs (study s34_..._ring_correctness_regression). Until the
+    # k_groups=4 sliced-K accumulation is fixed in attention_tiled_2d.py, D128 stays
+    # on the non-ring flash geometry (correct-but-slower).
+    if _resolve_attention_arch() == "gfx942" and problem.head_size == 128:
+        return False
     if not (
         (_enable_gfx942_fp16_flash(problem) or _enable_gfx942_bf16_flash(problem))
         and problem.head_size in (64, 128)

dnn-providers/hip-kernel-provider/rocke/library/tests/test_attn_bf16_d128_ring.py

@@ -1,28 +1,36 @@
 # Copyright (c) Advanced Micro Devices, Inc., or its affiliates.
 # SPDX-License-Identifier: MIT
 
-"""Regression tests for gfx942 bf16 D128 sharing the sliced-K ring geometry.
-
-Pure Python (no GPU / compilation). #9057 excluded D128 bf16 from the ring
-because the ring was attached to the non-ring bf16-wide geometry (nw=2, no
-cfvst), which the ring cannot use -- it requires the conflict-free-V store and
-the wide nw=4 flash geometry. The fix routes D128 bf16 through the fp16-flash
-ring geometry (nw=4, tile=64, cfvst) instead -- verified numerically correct
-(max_abs 0.00049, GQA + MHA) on both the Python and C++ engines with the
-byte-identity gate GREEN.
-
-These lock in:
-  1. D128 bf16 prefill now enables the sliced-K ring (no longer excluded).
-  2. The resulting spec is the fp16-flash geometry (nw=4, tile=64, ring, cfvst,
-     mask-limit) -- i.e. identical to the fp16 D128 spec modulo dtype.
-  3. _get_2d_launch_meta computes the grid/block for the ring geometry (nw=4),
-     matching the spec the launcher builds (not the non-ring bf16-wide nw=2).
-  4. D64 bf16 and fp16 D128 are unchanged (no regression).
+"""Regression tests for the gfx942 D128 sliced-K ring exclusion.
+
+The sliced-K ring (32-wide K slices -> k_groups = HD/32) is correctness-verified
+only for D64 (k_groups=2). The D128 ring (k_groups=4) is numerically WRONG at
+realistic input magnitude -- max_abs ~0.5-1.3 vs the fp32 oracle for BOTH bf16
+and fp16, GQA and MHA, even at S512 when the ring is forced -- while the non-ring
+D128 flash path (nw2/nw4, T=64) passes at max_abs ~0.0156 (bf16) / ~0.002 (fp16).
+
+#9198 briefly enabled the D128 bf16 ring on the strength of a "max_abs 0.00049"
+check, but that check used a uniform_(-0.1, 0.1) oracle whose near-uniform softmax
+never exercises the online-softmax running-max rescale across the four K slices,
+so the defect was masked. It reproduces immediately with randn (unit-variance)
+inputs (study s34_gfx942_bf16_d128_prefill_ring_correctness_regression). The fix
+excludes ALL D128 (both dtypes) from the ring until the k_groups=4 sliced-K
+accumulation is fixed in attention_tiled_2d.py.
+
+Two layers:
+  * Pure-Python spec assertions (always run): D128 ring OFF, geometry sane,
+    D64 ring unchanged.
+  * On-GPU numeric guardrail (skipped without a gfx942 device + torch): launch
+    the production kernel at randn magnitude and assert max_abs vs fp32 SDPA.
+    A spec-only test cannot catch a numerics regression in the emitted kernel --
+    that is exactly what let #9198 through.
 """
 
 from __future__ import annotations
 
-import dataclasses
+import importlib.util
+import os
+from pathlib import Path
 
 import pytest
 
@@ -39,9 +47,6 @@ def gfx942(monkeypatch):
     old_arch = au._RESOLVED_ATTENTION_ARCH
     au._RESOLVED_ATTENTION_ARCH = "gfx942"
     # Neutralize any inherited env overrides so we test the default policy.
-    # HIPDNN_GFX942_FLASH_WIDE and _K_LDSSEQ are included because they change the
-    # ring num_warps / LDS-sequence geometry these tests assert on -- a set value
-    # in a dev/CI shell would otherwise make the assertions non-deterministic.
     for var in (
         "HIPDNN_GFX942_K_SLICED_RING",
         "HIPDNN_GFX942_BF16_WIDE",
@@ -72,61 +77,184 @@ def _problem(dtype, sq=4096, hq=32, hk=8, d=128, bs=64):
     )
 
 
+# ---------------------------------------------------------------------------
+# Pure-Python spec assertions (no GPU)
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.parametrize("dtype", ["bf16", "fp16"])
 @pytest.mark.parametrize("hq,hk", [(32, 8), (16, 16), (64, 8), (128, 8), (64, 4)])
-def test_d128_bf16_enables_ring_by_default(gfx942, hq, hk):
-    p = _problem("bf16", hq=hq, hk=hk)
-    assert au._enable_gfx942_bf16_flash(p), "bf16 D128 prefill should be flash-eligible"
-    assert au._enable_gfx942_flash_k_sliced_ring(
+def test_d128_excluded_from_ring(gfx942, dtype, hq, hk):
+    # The k_groups=4 D128 ring is numerically wrong; it must stay OFF for both
+    # dtypes and every geometry.
+    p = _problem(dtype, hq=hq, hk=hk)
+    assert not au._enable_gfx942_flash_k_sliced_ring(
         p
-    ), "D128 bf16 should now share the sliced-K ring (exclusion removed)"
+    ), "D128 ring is numerically wrong (k_groups=4) and must be excluded"
+    spec = au._tiled_spec_from_problem(p)
+    assert not spec.use_k_sliced_ring
+    assert spec.tile_size == 64
 
 
-@pytest.mark.parametrize("hq,hk", [(32, 8), (16, 16)])
-def test_d128_bf16_spec_matches_fp16_modulo_dtype(gfx942, hq, hk):
-    sb = au._tiled_spec_from_problem(_problem("bf16", hq=hq, hk=hk))
-    sf = au._tiled_spec_from_problem(_problem("fp16", hq=hq, hk=hk))
-    fb = {f.name: getattr(sb, f.name) for f in dataclasses.fields(sb)}
-    ff = {f.name: getattr(sf, f.name) for f in dataclasses.fields(sf)}
-    differing = {k for k in fb if fb[k] != ff[k]}
-    assert differing == {"dtype"}, f"unexpected spec diffs: {differing}"
-    # Sanity: it really is the ring geometry.
-    assert sb.num_warps == 4
-    assert sb.tile_size == 64
-    assert sb.use_k_sliced_ring
-    assert sb.use_conflict_free_v_store
-    assert sb.block_m_per_warp == 32
+@pytest.mark.parametrize("dtype", ["bf16", "fp16"])
+def test_d128_ring_force_on_still_off(gfx942, monkeypatch, dtype):
+    # Even the explicit force-on env must NOT re-enable the D128 ring: the
+    # head_size==128 guard short-circuits before the env is consulted.
+    monkeypatch.setenv("HIPDNN_GFX942_K_SLICED_RING", "1")
+    p = _problem(dtype)
+    assert not au._enable_gfx942_flash_k_sliced_ring(p)
 
 
 @pytest.mark.parametrize("hq,hk", [(32, 8), (16, 16), (64, 8), (128, 8), (64, 4)])
-def test_d128_bf16_launch_meta_matches_ring_spec(gfx942, hq, hk):
+def test_d128_launch_meta_matches_nonring_spec(gfx942, hq, hk):
+    # The launch grid/block must match the (non-ring) geometry the launcher
+    # actually builds -- a mismatch silently corrupts output.
     au._2D_LAUNCH_META.clear()
     p = _problem("bf16", hq=hq, hk=hk)
     spec = au._tiled_spec_from_problem(p)
     meta = au._get_2d_launch_meta(p, au._tiled_cache_key(p))
-    block_m = spec.num_warps * spec.block_m_per_warp
-    nqk = hq // hk
-    block_q = block_m // nqk if nqk <= block_m else 1
-    expected_blocks = p.total_q // block_q + p.num_seqs
-    assert meta.grid[1] == expected_blocks
-    assert meta.block[0] == 64 * spec.num_warps  # 64*4 == 256, not the wide 64*2
+    assert meta.block[0] == 64 * spec.num_warps
 
 
-def test_d64_bf16_unchanged(gfx942):
-    # D64 already used the ring pre-fix; must stay nw=4 ring.
-    s = au._tiled_spec_from_problem(_problem("bf16", d=64, bs=16))
-    assert s.use_k_sliced_ring
-    assert s.num_warps == 4
+def test_d64_ring_unchanged(gfx942):
+    # D64 (k_groups=2) is correctness-verified and keeps the ring.
+    for dtype in ("bf16", "fp16"):
+        s = au._tiled_spec_from_problem(_problem(dtype, d=64, bs=16))
+        assert s.use_k_sliced_ring, f"D64 {dtype} must keep the ring"
+        assert s.num_warps == 4
 
 
-def test_fp16_d128_unchanged(gfx942):
-    s = au._tiled_spec_from_problem(_problem("fp16"))
-    assert s.num_warps == 4
-    assert s.tile_size == 64
-    assert s.use_k_sliced_ring
+# ---------------------------------------------------------------------------
+# On-GPU numeric guardrail -- the check #9198 lacked
+# ---------------------------------------------------------------------------
 
 
-def test_short_context_bf16_still_narrow(gfx942):
-    # small_q_narrow (q<=768) must NOT take the ring (per _enable_gfx942_bf16_flash).
-    p = _problem("bf16", sq=512)
-    assert au._enable_gfx942_small_q_narrow(p)
-    assert not au._enable_gfx942_flash_k_sliced_ring(p)
+def _load_parity_harness():
+    import sys
+
+    here = Path(__file__).resolve()
+    # rocke/library/tests -> rocke/library/builders/gfx942/attention/prefill
+    harness = (
+        here.parents[1]
+        / "builders/gfx942/attention/prefill/parity_unified_attention.py"
+    )
+    if not harness.exists():
+        return None
+    spec = importlib.util.spec_from_file_location("parity_unified_attention", harness)
+    mod = importlib.util.module_from_spec(spec)
+    # Register before exec so dataclasses in the harness (e.g. Shape) can resolve
+    # cls.__module__ during field introspection.
+    sys.modules["parity_unified_attention"] = mod
+    spec.loader.exec_module(mod)
+    return mod
+
+
+def _gpu_ready():
+    try:
+        import torch
+    except Exception:  # noqa: BLE001
+        return False
+    if not torch.cuda.is_available():
+        return False
+    name = torch.cuda.get_device_name(0).lower()
+    return "mi300" in name or "gfx942" in name
+
+
+requires_gfx942_gpu = pytest.mark.skipif(
+    not _gpu_ready(), reason="needs a gfx942 (MI300X) GPU with ROCm torch"
+)
+
+
+@requires_gfx942_gpu
+@pytest.mark.parametrize("dtype,tol", [("bf16", 5e-2), ("fp16", 5e-2)])
+@pytest.mark.parametrize("hq,hk", [(32, 8), (16, 16)])
+@pytest.mark.parametrize("sq", [512, 1024, 4096])
+def test_d128_numeric_vs_fp32_oracle_at_magnitude(dtype, tol, hq, hk, sq):
+    """Launch the production D128 kernel at realistic (unit-variance) magnitude
+    and assert max_abs vs the fp32 paged-attn oracle.
+
+    This is the guardrail #9198 lacked: with the default uniform_(-0.1, 0.1)
+    inputs the ring bug is invisible (softmax ~uniform); randn inputs expose it
+    immediately. If a future change re-enables the broken D128 ring, this fails.
+    """
+    import torch
+
+    H = _load_parity_harness()
+    if H is None:
+        pytest.skip("parity harness not present in this checkout")
+
+    # Force the arch resolver so the production selector picks the gfx942 path.
+    old_arch = au._RESOLVED_ATTENTION_ARCH
+    au._RESOLVED_ATTENTION_ARCH = "gfx942"
+    for var in (
+        "HIPDNN_GFX942_K_SLICED_RING",
+        "HIPDNN_GFX942_FLASH_WIDE",
+        "HIPDNN_GFX942_FLASH_MLIM",
+    ):
+        os.environ.pop(var, None)
+    try:
+        from builders.common.attention_spec_builder import _tiled_spec_from_problem
+
+        s = H.Shape(
+            f"d128_{hq}_{hk}_S{sq}", dtype, sq, sq, 128, hq, hk, 1, True, "perf"
+        )
+
+        # randn (unit-variance) inputs -- NOT the uniform_(-0.1,0.1) default that
+        # masked the ring bug.
+        torch.manual_seed(0)
+        td = torch.float16 if dtype == "fp16" else torch.bfloat16
+        BS = H.BLOCK_SIZE
+        max_blk = (sq + BS - 1) // BS
+        query = torch.randn(sq, hq, 128, device="cuda").to(td)
+        key_cache = torch.randn(max_blk, BS, hk, 128, device="cuda").to(td)
+        value_cache = torch.randn(max_blk, BS, hk, 128, device="cuda").to(td)
+        cu_q = torch.tensor([0, sq], dtype=torch.int32, device="cuda")
+        kv_lens = torch.tensor([sq], dtype=torch.int32, device="cuda")
+        block_tables = torch.arange(max_blk, dtype=torch.int32, device="cuda").view(
+            1, max_blk
+        )
+        scale = 128**-0.5
+        data = {
+            "query": query,
+            "key_cache": key_cache,
+            "value_cache": value_cache,
+            "cu_q": cu_q,
+            "kv_lens": kv_lens,
+            "query_lens": [sq],
+            "kv_lens_list": [sq],
+            "block_tables": block_tables,
+            "scale": scale,
+            "max_query_len": sq,
+            "max_kv_len": sq,
+        }
+        ref = H.ref_paged_attn(
+            query, key_cache, value_cache, [sq], [sq], block_tables, scale
+        )
+
+        p = _problem(dtype, sq=sq, hq=hq, hk=hk)
+        spec = _tiled_spec_from_problem(p)
+        assert not spec.use_k_sliced_ring, "production D128 spec must not use the ring"
+
+        from rocke import compile_kernel
+        from kernels import build_unified_attention_2d_tiled
+        from kernels.common.attention_unified import _attn_signature
+        from rocke.runtime import KernelLauncher
+
+        kernel = build_unified_attention_2d_tiled(spec, arch="gfx942")
+        artifact = compile_kernel(kernel, arch="gfx942", capture_ir_text=False)
+        launcher = KernelLauncher(
+            hsaco=artifact.hsaco,
+            kernel_name=artifact.kernel_name,
+            signature=_attn_signature(
+                dtype, include_bt_stride=True, include_qq_bias_stride=True
+            ),
+            cache_key=("d128guard", spec.kernel_name(), dtype, sq, hq, hk),
+        )
+        out, _ = H._run_rocke(s, data, launcher, spec, warmup=3, attempts=3)
+        m = H.compare(ref, out)
+        assert m["max_abs"] <= tol, (
+            f"D128 {dtype} {hq}/{hk} S{sq}: max_abs {m['max_abs']:.4f} > {tol} "
+            f"(ring regression?)"
+        )
+    finally:
+        au._RESOLVED_ATTENTION_ARCH = old_arch