From 892ef5c63ba1b503a18b359a6303361b336787c3 Mon Sep 17 00:00:00 2001 From: dlangbe Date: Wed, 7 Oct 2026 21:59:39 -0400 Subject: [PATCH 01/16] fix(rocke/attention): key every resolved selector in the tiled cache keys The tiled 2D and 3D cache keys are computed per launch without building the spec, so they must carry the resolved output of every selector the spec builder consults. Three did not reach the key: - binary_search_iters, the seq_idx search trip count every tiled spec bakes into IR from num_seqs. Only the raw num_seqs key field separated it. - use_softmax_mfma_interleave and use_k_single_buffer, gated on the gfx950 single-batch cohort (num_seqs == 1). Also separated only by raw num_seqs. - use_fp8_mfma_qk, which reads max_seqlen_k. Nothing in the key reads that field, so two fp8 problems straddling max_seqlen_k = 16 * T shared one key and the cache served whichever kernel compiled first. Add those outputs, plus the combo predicate and sched_barrier, so the key covers every spec field the builder resolves. The interleave / K single-buffer / sched_barrier entries are gated on the single-batch cohort so other launches skip those selector walks. _enable_fp8_mfma_qk now checks max_seqlen_q before the loader-fit walk; both checks are pure, so the result is unchanged. The fp8 entry is skipped on gfx1250, whose builder never reads the field. Raw num_seqs stays in the key here. The follow-up drops it. test_attention_tiled_cache_key sweeps every num_seqs / max_seqlen_q / max_seqlen_k selector threshold on gfx942, gfx950 and gfx1250, 2D and 3D, and asserts an equal key implies an equal built spec, builder route and compile backend. Co-Authored-By: Claude --- .../kernels/common/attention_unified.py | 70 +++- .../tests/test_attention_tiled_cache_key.py | 330 ++++++++++++++++++ 2 files changed, 393 insertions(+), 7 deletions(-) create mode 100644 dnn-providers/hip-kernel-provider/rocke/library/tests/test_attention_tiled_cache_key.py diff --git a/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/attention_unified.py b/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/attention_unified.py index 495309bed9aa..0a345e37f1ed 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/attention_unified.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/attention_unified.py @@ -10,6 +10,7 @@ from __future__ import annotations +import math from dataclasses import dataclass, replace from typing import Any, Callable, Dict, Optional, Tuple @@ -1436,6 +1437,19 @@ def _kv_storage_dtype(problem: UnifiedAttentionProblem) -> Optional[str]: return "fp8e4m3" if problem.use_fp8 else None +def _binary_search_iters(problem: UnifiedAttentionProblem) -> int: + """Trip count of the seq_idx binary search the tiled kernels bake into IR. + + Mirrors ``binary_search_iters`` on the gfx942/gfx950/gfx1250 tiled 2D and + gfx942/gfx950 tiled 3D spec classes, which derive it from ``spec.num_seqs``. + It is the only way ``num_seqs`` reaches emitted IR, so the tiled cache keys + carry this bucket instead of the raw batch count. + """ + if problem.num_seqs <= 0: + return 32 + return max(1, int(math.ceil(math.log2(problem.num_seqs + 1)))) + + def _tiled_cache_key(problem: UnifiedAttentionProblem) -> Tuple: """Compute the tiled-2D cache key WITHOUT building the spec dataclass. @@ -1445,14 +1459,23 @@ def _tiled_cache_key(problem: UnifiedAttentionProblem) -> Tuple: which is material on decode kernels (~20µs). The spec is only built on cache miss (first launch per shape) inside ``_get_2d_launcher``. - Note: this MUST match the spec-derived key — every selector knob - that affects the kernel build is included. If a selector changes - behaviour, update both this function and ``_tiled_spec_from_problem``. + Note: this MUST match the spec-derived key. The key carries static problem + fields plus the RESOLVED output of every selector the builder consults, and + no raw dynamic shape field (``num_seqs``, ``total_q``, ``max_seqlen_q``, + ``max_seqlen_k``). A raw field compiles one binary per value; a selector + missing from the key lets two shapes that build different kernels share + one cache slot. ``num_seqs`` reaches IR only through + ``_binary_search_iters``. Launch values that scale with the shape (the + grid) are keyed separately in ``_get_2d_launch_meta``. When adding a + selector to ``_tiled_spec_from_problem``, add its resolved output here; + ``test_attention_tiled_cache_key.py`` sweeps the selector thresholds and + fails if two problems share a key but build different specs. """ return ( "tiled", _resolve_attention_arch(), problem.num_seqs, + _binary_search_iters(problem), problem.num_query_heads, problem.num_kv_heads, problem.head_size, @@ -1494,6 +1517,29 @@ def _tiled_cache_key(problem: UnifiedAttentionProblem) -> Tuple: _enable_transposed_subflags(problem), _enable_v_double_buffer(problem), _enable_early_v_schedule(problem), + # The builder derives the scalar-state / mask / skip-qreg / fast-paged-KV + # spec fields from the combo predicate, which reads max_seqlen_q. + _enable_combo_2d(problem), + # gfx950 schedule tail of `_spec_gfx950_generic`. All three are false + # outside the single-batch cohort (num_seqs == 1), so gate on it and + # spare every other launch the selector walks. + ( + ( + _enable_sched_barrier(problem), + _enable_softmax_mfma_interleave(problem), + _enable_k_single_buffer(problem), + ) + if _enable_single_batch_combo(problem) + else None + ), + # Reads max_seqlen_q and max_seqlen_k (the 16 * T threshold). The + # gfx1250 builder never sets this field, so keying it there would only + # split identical kernels. + ( + _enable_fp8_mfma_qk(problem) + if _resolve_attention_arch() != "gfx1250" + else None + ), _enable_register_pv(problem), _enable_i64_kv_addr(problem), _select_2d_compile_backend(problem), @@ -1651,13 +1697,14 @@ def _enable_fp8_mfma_qk(problem: UnifiedAttentionProblem) -> bool: # path. Never combine them. if _enable_combo_2d(problem): return False - if not _fp8_qk_loader_fits(problem): - return False # Large prefill (long per-CTA KV loop) prefers the sync-dequant path # regardless of sliding window; the K-in-LDS loader only wins when the - # KV loop is short (decode / short prefill). + # KV loop is short (decode / short prefill). Checked before the loader fit + # because the cache key evaluates this on every fp8 launch. if problem.max_seqlen_q > 256: return False + if not _fp8_qk_loader_fits(problem): + return False T_eff = _select_2d_tile_size(problem) return problem.sliding_window > 0 or problem.max_seqlen_k <= 16 * T_eff @@ -3169,6 +3216,12 @@ def _tiled_3d_spec_from_problem( def _tiled_3d_cache_key(problem: UnifiedAttentionProblem) -> Tuple: + """HSACO cache key for the 3D split-KV segment + reduce pair. + + Same contract as ``_tiled_cache_key``: static fields plus resolved selector + outputs, no raw dynamic shape field. Grid and workspace values that scale + with the shape are keyed separately in ``_get_3d_pipeline``. + """ base = ( "tiled3d", problem.num_seqs, @@ -3193,6 +3246,9 @@ def _tiled_3d_cache_key(problem: UnifiedAttentionProblem) -> Tuple: if _resolve_attention_arch() == "gfx1250": sp = _tiled_3d_spec_from_problem(problem) return base + ( + # gfx1250 floors the search at 32 iterations, so read it from the + # spec rather than _binary_search_iters. + sp.binary_search_iters, sp.use_invariant_hoist, sp.use_wide_kv_load, sp.use_register_p, @@ -3204,7 +3260,7 @@ def _tiled_3d_cache_key(problem: UnifiedAttentionProblem) -> Tuple: getattr(sp, "use_fused_reduce", False), getattr(sp, "use_dpp_softmax", False), ) - return base + return base + (_binary_search_iters(problem),) def _3d_signature(dtype: str, *, kv_dtype: Optional[str] = None): diff --git a/dnn-providers/hip-kernel-provider/rocke/library/tests/test_attention_tiled_cache_key.py b/dnn-providers/hip-kernel-provider/rocke/library/tests/test_attention_tiled_cache_key.py new file mode 100644 index 000000000000..0fafb70fc16e --- /dev/null +++ b/dnn-providers/hip-kernel-provider/rocke/library/tests/test_attention_tiled_cache_key.py @@ -0,0 +1,330 @@ +# Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +# SPDX-License-Identifier: MIT + +"""Cache-key contract for the tiled 2D / 3D attention launchers. + +``_tiled_cache_key`` and ``_tiled_3d_cache_key`` are computed on every launch +without building the spec, so they must carry the resolved output of every +selector the spec builder consults. Two failure directions: + + * Under-keying: two problems share a key but build different kernels, and + the cache serves whichever compiled first. + * Over-keying: a raw shape field splits problems that build the same kernel, + so every batch count or seqlen compiles its own copy. + +Layers: + * Invariant sweep: across every selector threshold, an equal key implies an + equal spec (plus builder route and compile backend). + * Targeted pairs: shapes that resolve to different selector outputs get + different keys. + * ``_binary_search_iters`` mirrors every spec class's ``binary_search_iters``. + +No GPU required. +""" + +from __future__ import annotations + +import dataclasses +import itertools + +import pytest + +from kernels.common import attention_unified as au + +P = au.UnifiedAttentionProblem + +ARCHES = ("gfx942", "gfx950", "gfx1250") + +# Env knobs that steer selectors. Cleared so the sweep sees production routing. +_ROUTING_ENV = ( + "HIPDNN_GFX942_K_SLICE_HD", + "HIPDNN_GFX942_K_SLICED_RING", + "HIPDNN_GFX942_K_LDSSEQ", + "HIPDNN_GFX942_FLASH_WIDE", + "HIPDNN_GFX942_FLASH_MLIM", + "HIPDNN_GFX942_BF16_CFVST", + "HIPDNN_GFX942_D128_SMALLTILE_DK", + "HIPDNN_GFX942_3D_HOIST", + "HIPDNN_GFX942_3D_WKV", + "HIPDNN_GFX950_D128_SMALL_TILE", + "HIPDNN_GFX950_D128_SOFTMAX_INTERLEAVE", +) + + +@pytest.fixture +def pin_arch(monkeypatch): + """Pin the memoized arch (conftest restores it) and clear routing env.""" + for var in _ROUTING_ENV: + monkeypatch.delenv(var, raising=False) + + def _pin(arch): + au._RESOLVED_ATTENTION_ARCH = arch + + return _pin + + +def _norm_spec(spec): + """Spec identity as the emitter sees it: ``num_seqs`` only reaches IR + through ``binary_search_iters``.""" + fields = dataclasses.asdict(spec) + fields.pop("num_seqs") + fields["binary_search_iters"] = spec.binary_search_iters + return fields + + +def _built_2d(problem): + route = au._gfx942_4warp_route(problem) + return dict( + spec=_norm_spec(au._tiled_spec_from_problem(problem)), + builder=route.cohort if route is not None else "default", + backend=au._select_2d_compile_backend(problem), + ) + + +def _built_3d(problem): + return dict( + spec=_norm_spec(au._tiled_3d_spec_from_problem(problem)), + reduce=( + problem.head_size, + problem.num_query_heads, + problem.num_kv_heads, + problem.dtype, + au._num_segments(problem), + au._select_3d_waves_per_eu(problem), + ), + ) + + +# --------------------------------------------------------------------------- +# Invariant sweep +# --------------------------------------------------------------------------- + +# Dynamic axes straddle every num_seqs / max_seqlen_q / max_seqlen_k threshold +# the selectors read: ns 1 vs >=2 and the log2 search buckets; q at 1, 64, 128, +# 256, 512, 768, 1024, 1536, 2048; k across 16 * T for T in {32, 64}. +_NUM_SEQS = (1, 2, 3, 4) +_SEQLEN_Q = (1, 2, 64, 65, 256, 257, 512, 513, 768, 769, 1024, 1025, 1536, 2048) +_SEQLEN_K = ("q", 512, 513, 1024, 1025) + +# (sliding_window, use_sinks, bias) combinations; bias is one feature at a time. +_FEATURES = ( + (0, False, None), + (128, False, None), + (0, True, None), + (0, False, "softcap"), + (0, False, "alibi"), +) + + +def _static_grid(arch): + for dtype, fp8, hd, bs, heads, (sw, sinks, bias) in itertools.product( + ("bf16", "fp16"), + (False, True), + (64, 128, 256), + (16, 32, 64), + ((8, 8), (32, 8), (64, 8)), + _FEATURES, + ): + yield dict( + dtype=dtype, + use_fp8=fp8, + fp8_fnuz=fp8 and arch == "gfx942", + head_size=hd, + block_size=bs, + num_query_heads=heads[0], + num_kv_heads=heads[1], + sliding_window=sw, + use_sinks=sinks, + softcap=30.0 if bias == "softcap" else 0.0, + use_alibi=bias == "alibi", + ) + + +def _sweep(arch): + for static in _static_grid(arch): + for ns, q, k in itertools.product(_NUM_SEQS, _SEQLEN_Q, _SEQLEN_K): + k = q if k == "q" else k + if k < q: + continue + yield P( + total_q=ns * q, num_seqs=ns, max_seqlen_q=q, max_seqlen_k=k, **static + ) + + +@pytest.mark.parametrize("arch", ARCHES) +@pytest.mark.parametrize("path", ("2d", "3d")) +def test_equal_key_implies_equal_build(pin_arch, arch, path): + """No two problems may share a key while building different kernels. + + This is the guard for the next selector: one added to the spec builder but + not to the key fails here as soon as the sweep straddles its threshold. + """ + pin_arch(arch) + if path == "2d": + supported = au.supports_native_unified_attention_tiled + key_fn, built_fn = au._tiled_cache_key, _built_2d + else: + supported = au.supports_native_unified_attention_3d_tiled + key_fn, built_fn = au._tiled_3d_cache_key, _built_3d + + first = {} + checked = 0 + for problem in _sweep(arch): + if not supported(problem)[0]: + continue + try: + # gfx1250's 3D key builds the spec itself, so it raises too. + key = key_fn(problem) + built = built_fn(problem) + except (ValueError, NotImplementedError): + # Admitted by the gate but rejected by the spec; never launches. + continue + checked += 1 + seen = first.setdefault(key, (problem, built)) + if seen[1] != built: + other, other_built = seen + diff = { + name: (other_built["spec"].get(name), built["spec"].get(name)) + for name in built["spec"] + if other_built["spec"].get(name) != built["spec"].get(name) + } + for name in built: + if name != "spec" and other_built[name] != built[name]: + diff[name] = (other_built[name], built[name]) + pytest.fail( + f"{arch} {path}: one key, two kernels\n {other}\n {problem}\n" + f" differing: {diff}" + ) + assert checked > 0 + + +# --------------------------------------------------------------------------- +# Different selector outputs -> different keys +# --------------------------------------------------------------------------- + + +def _assert_split(key_fn, built_fn, a, b, field): + assert built_fn(a)["spec"][field] != built_fn(b)["spec"][field] + assert key_fn(a) != key_fn(b) + + +@pytest.mark.parametrize("arch", ("gfx942", "gfx950", "gfx1250")) +def test_search_bucket_edge_splits_2d_key(pin_arch, arch): + # num_seqs 3 and 4 fall in different log2 search buckets (2 vs 3 iterations). + # d64 GQA-8 (+ fp8 K/V on gfx1250): a geometry every arch's tiled 2D supports. + pin_arch(arch) + base = dict( + num_query_heads=64, + num_kv_heads=8, + head_size=64, + block_size=32, + max_seqlen_q=1024, + max_seqlen_k=1024, + dtype="bf16", + use_fp8=arch == "gfx1250", + ) + a = P(total_q=3 * 1024, num_seqs=3, **base) + b = P(total_q=4 * 1024, num_seqs=4, **base) + _assert_split(au._tiled_cache_key, _built_2d, a, b, "binary_search_iters") + + +@pytest.mark.parametrize("arch", ("gfx942", "gfx950")) +def test_search_bucket_edge_splits_3d_key(pin_arch, arch): + pin_arch(arch) + base = dict( + num_query_heads=32, + num_kv_heads=8, + head_size=128, + block_size=16, + max_seqlen_q=1, + max_seqlen_k=8192, + dtype="bf16", + ) + a = P(total_q=3, num_seqs=3, **base) + b = P(total_q=4, num_seqs=4, **base) + _assert_split(au._tiled_3d_cache_key, _built_3d, a, b, "binary_search_iters") + + +def test_fp8_mfma_qk_kv_threshold_splits_key(pin_arch): + # T = 2 * block_size = 32 -> the K-in-LDS fp8 path turns off above + # max_seqlen_k = 16 * T = 512. Nothing else in the key reads max_seqlen_k. + pin_arch("gfx950") + base = dict( + total_q=1, + num_seqs=1, + num_query_heads=8, + num_kv_heads=8, + head_size=64, + block_size=16, + max_seqlen_q=1, + dtype="bf16", + use_fp8=True, + ) + a = P(max_seqlen_k=512, **base) + b = P(max_seqlen_k=513, **base) + assert au._enable_fp8_mfma_qk(a) and not au._enable_fp8_mfma_qk(b) + _assert_split(au._tiled_cache_key, _built_2d, a, b, "use_fp8_mfma_qk") + + +def test_single_batch_interleave_splits_key(pin_arch): + # The d128 single-batch cohort turns on the softmax/MFMA interleave; the + # same shape at num_seqs=2 does not. + pin_arch("gfx950") + base = dict( + num_query_heads=8, + num_kv_heads=8, + head_size=128, + block_size=16, + max_seqlen_q=257, + max_seqlen_k=257, + dtype="bf16", + ) + a = P(total_q=257, num_seqs=1, **base) + b = P(total_q=514, num_seqs=2, **base) + _assert_split(au._tiled_cache_key, _built_2d, a, b, "use_softmax_mfma_interleave") + + +# --------------------------------------------------------------------------- +# _binary_search_iters mirrors the spec classes +# --------------------------------------------------------------------------- + + +def _spec_classes(): + from kernels.gfx942.attention_tiled_2d import ( + UnifiedAttention2DTiledSpec as Gfx942Spec2D, + ) + from kernels.gfx942.attention_tiled_3d import ( + UnifiedAttention3DTiledSpec as Gfx942Spec3D, + ) + from kernels.gfx950.attention_tiled_2d import ( + UnifiedAttention2DTiledSpec as Gfx950Spec2D, + ) + from kernels.gfx950.attention_tiled_3d import ( + UnifiedAttention3DTiledSpec as Gfx950Spec3D, + ) + from kernels.gfx1250.attention_tiled_2d import ( + UnifiedAttention2DTiledSpec as Gfx1250Spec2D, + ) + + return (Gfx942Spec2D, Gfx942Spec3D, Gfx950Spec2D, Gfx950Spec3D, Gfx1250Spec2D) + + +@pytest.mark.parametrize("num_seqs", (0, 1, 2, 3, 4, 7, 8, 15, 16, 450, 1 << 20)) +@pytest.mark.parametrize("spec_cls", _spec_classes(), ids=lambda c: c.__module__) +def test_binary_search_iters_mirrors_spec(spec_cls, num_seqs): + # binary_search_iters is a pure function of num_seqs; bypass __init__ + # validation so the comparison needs no full geometry. + spec = object.__new__(spec_cls) + object.__setattr__(spec, "num_seqs", num_seqs) + problem = P( + total_q=1, + num_seqs=num_seqs, + num_query_heads=8, + num_kv_heads=8, + head_size=128, + block_size=16, + max_seqlen_q=1, + max_seqlen_k=1, + dtype="bf16", + ) + assert au._binary_search_iters(problem) == spec.binary_search_iters From 8352908d452a0333022f757c6a6ca72f01c70434 Mon Sep 17 00:00:00 2001 From: dlangbe Date: Wed, 7 Oct 2026 23:29:31 -0400 Subject: [PATCH 02/16] perf(rocke/attention): drop raw num_seqs from the tiled cache keys With every resolved selector output in the key, raw num_seqs only split problems that build the same kernel. Remove it from _tiled_cache_key and _tiled_3d_cache_key. Batch count now reaches the key only through _binary_search_iters, so 2D and 3D binaries are shared within a log2 batch bucket (and across all batch counts on gfx1250 3D, which floors the search at 32 iterations). The launch-geometry caches still need the batch count: both grids add one padding q-block per sequence. The 2D launch-meta key and the 3D prepared-pipeline key now go through _shape_key, which appends total_q and num_seqs; the 3D prepared pipeline covers both grids and the workspace, not the workspace alone. The 3D CUDA-graph key gains num_seqs too, since the captured launch bakes in the num_seqs kernarg and the segment grid. test_attention_tiled_cache_key now asserts both directions: an equal key implies an equal build, and an equal build implies an equal key, so a raw shape field or a selector the builder never reads fails as over-keying. New pair tests check that same-output shapes share a key and lower to byte-identical IR at llvm20 and llvm22, and that launch meta still separates batch counts under one key. Co-Authored-By: Claude --- .../kernels/common/attention_unified.py | 32 ++- .../tests/test_attention_tiled_cache_key.py | 184 +++++++++++++++++- 2 files changed, 202 insertions(+), 14 deletions(-) diff --git a/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/attention_unified.py b/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/attention_unified.py index 0a345e37f1ed..49ea1347a48f 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/attention_unified.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/attention_unified.py @@ -1466,15 +1466,14 @@ def _tiled_cache_key(problem: UnifiedAttentionProblem) -> Tuple: missing from the key lets two shapes that build different kernels share one cache slot. ``num_seqs`` reaches IR only through ``_binary_search_iters``. Launch values that scale with the shape (the - grid) are keyed separately in ``_get_2d_launch_meta``. When adding a + grid) are keyed separately through ``_shape_key``. When adding a selector to ``_tiled_spec_from_problem``, add its resolved output here; ``test_attention_tiled_cache_key.py`` sweeps the selector thresholds and - fails if two problems share a key but build different specs. + fails wherever key equality and built-spec equality disagree. """ return ( "tiled", _resolve_attention_arch(), - problem.num_seqs, _binary_search_iters(problem), problem.num_query_heads, problem.num_kv_heads, @@ -3220,11 +3219,10 @@ def _tiled_3d_cache_key(problem: UnifiedAttentionProblem) -> Tuple: Same contract as ``_tiled_cache_key``: static fields plus resolved selector outputs, no raw dynamic shape field. Grid and workspace values that scale - with the shape are keyed separately in ``_get_3d_pipeline``. + with the shape are keyed separately through ``_shape_key``. """ base = ( "tiled3d", - problem.num_seqs, problem.num_query_heads, problem.num_kv_heads, problem.head_size, @@ -3445,6 +3443,9 @@ def _run_3d_tiled( graph_key = ( cache_key, int(problem.total_q), + # The captured launch bakes in the num_seqs kernarg and the segment + # grid, and cache_key no longer separates batch counts. + int(problem.num_seqs), int(stream), id(q), id(k), @@ -4025,6 +4026,21 @@ def attention_3d_workspace_nbytes( ) +def _shape_key(cache_key: Tuple, problem: UnifiedAttentionProblem) -> Tuple: + """Extend a binary cache key with the shape values the launch derives from. + + The tiled binaries are shared across batch counts and token counts, but the + grids (``total_q // block_q + num_seqs``), the 3D reduce grid and the 3D + workspace are not. Everything cached per launch geometry keys on this. + """ + return cache_key + ( + "total_q", + int(problem.total_q), + "num_seqs", + int(problem.num_seqs), + ) + + def _get_3d_pipeline( problem: UnifiedAttentionProblem, cache_key: Tuple, @@ -4032,7 +4048,9 @@ def _get_3d_pipeline( *, tuning_spec=None, ) -> _Attention3DPrepared: - prepared_key = cache_key + ("total_q", int(problem.total_q)) + # The prepared pipeline owns both launch grids and the workspace, so it is + # per shape; the HSACO pair under cache_key is shared. + prepared_key = _shape_key(cache_key, problem) if prepared_key in _3D_PIPELINES: return _3D_PIPELINES[prepared_key] if cache_key not in _ATTN_3D_TILED_CACHE: @@ -4319,7 +4337,7 @@ def _get_2d_launch_meta( *, tuning_spec=None, ) -> _Attention2DLaunchMeta: - meta_key = cache_key + ("total_q", int(problem.total_q)) + meta_key = _shape_key(cache_key, problem) if meta_key in _2D_LAUNCH_META: return _2D_LAUNCH_META[meta_key] arch = _resolve_attention_arch() diff --git a/dnn-providers/hip-kernel-provider/rocke/library/tests/test_attention_tiled_cache_key.py b/dnn-providers/hip-kernel-provider/rocke/library/tests/test_attention_tiled_cache_key.py index 0fafb70fc16e..081d8419d4c9 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/tests/test_attention_tiled_cache_key.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/tests/test_attention_tiled_cache_key.py @@ -13,10 +13,12 @@ so every batch count or seqlen compiles its own copy. Layers: - * Invariant sweep: across every selector threshold, an equal key implies an - equal spec (plus builder route and compile backend). + * Invariant sweep: across every selector threshold, two problems share a key + exactly when they build the same spec, builder route and compile backend. * Targeted pairs: shapes that resolve to different selector outputs get - different keys. + different keys; shapes that resolve to the same outputs share a key and + lower to byte-identical IR. + * Launch geometry: the grid caches still separate batch counts. * ``_binary_search_iters`` mirrors every spec class's ``binary_search_iters``. No GPU required. @@ -153,11 +155,15 @@ def _sweep(arch): @pytest.mark.parametrize("arch", ARCHES) @pytest.mark.parametrize("path", ("2d", "3d")) -def test_equal_key_implies_equal_build(pin_arch, arch, path): - """No two problems may share a key while building different kernels. +def test_key_matches_build_identity(pin_arch, arch, path): + """Key equality and build equality must coincide across the sweep. - This is the guard for the next selector: one added to the spec builder but - not to the key fails here as soon as the sweep straddles its threshold. + Equal key, different build: under-keying. This is the guard for the next + selector; one added to the spec builder but not to the key fails here as + soon as the sweep straddles its threshold. + + Equal build, different key: over-keying. A raw shape field, or a selector + the arch's builder never reads, compiles duplicate binaries. """ pin_arch(arch) if path == "2d": @@ -168,6 +174,7 @@ def test_equal_key_implies_equal_build(pin_arch, arch, path): key_fn, built_fn = au._tiled_3d_cache_key, _built_3d first = {} + key_of_build = {} checked = 0 for problem in _sweep(arch): if not supported(problem)[0]: @@ -195,6 +202,14 @@ def test_equal_key_implies_equal_build(pin_arch, arch, path): f"{arch} {path}: one key, two kernels\n {other}\n {problem}\n" f" differing: {diff}" ) + build_id = repr(sorted((name, repr(value)) for name, value in built.items())) + other_key, other = key_of_build.setdefault(build_id, (key, problem)) + if other_key != key: + diff = [(i, a, b) for i, (a, b) in enumerate(zip(other_key, key)) if a != b] + pytest.fail( + f"{arch} {path}: one kernel, two keys\n {other}\n {problem}\n" + f" differing key entries (index, a, b): {diff}" + ) assert checked > 0 @@ -226,6 +241,8 @@ def test_search_bucket_edge_splits_2d_key(pin_arch, arch): a = P(total_q=3 * 1024, num_seqs=3, **base) b = P(total_q=4 * 1024, num_seqs=4, **base) _assert_split(au._tiled_cache_key, _built_2d, a, b, "binary_search_iters") + # The trip count is a literal in the IR, so the binaries really differ. + assert _lower_2d(a, arch, "llvm22") != _lower_2d(b, arch, "llvm22") @pytest.mark.parametrize("arch", ("gfx942", "gfx950")) @@ -284,6 +301,159 @@ def test_single_batch_interleave_splits_key(pin_arch): _assert_split(au._tiled_cache_key, _built_2d, a, b, "use_softmax_mfma_interleave") +# --------------------------------------------------------------------------- +# Same selector outputs -> same key and byte-identical IR +# --------------------------------------------------------------------------- + + +def _lower_2d(problem, arch, flavor): + from rocke.core.lower_llvm import _lower_kernel_to_llvm_python + + _, build, _ = au._tiled_2d_impl(arch) + spec = au._tiled_spec_from_problem(problem) + route = au._gfx942_4warp_route(problem) + kernel = route.builder(spec, arch=arch) if route else build(spec, arch=arch) + return _lower_kernel_to_llvm_python(kernel, arch=arch, llvm_flavor=flavor) + + +def _lower_3d(problem, arch, flavor): + from rocke.core.lower_llvm import _lower_kernel_to_llvm_python + + _, reduce_spec_cls, build_seg, build_red, _ = au._tiled_3d_impl(arch) + seg = build_seg(au._tiled_3d_spec_from_problem(problem), arch=arch) + red = build_red( + reduce_spec_cls( + head_size=problem.head_size, + num_query_heads=problem.num_query_heads, + num_kv_heads=problem.num_kv_heads, + dtype=problem.dtype, + num_segments=au._num_segments(problem), + waves_per_eu=au._select_3d_waves_per_eu(problem), + ), + arch=arch, + ) + return tuple( + _lower_kernel_to_llvm_python(k, arch=arch, llvm_flavor=flavor) + for k in (seg, red) + ) + + +_D64_GQA8 = dict( + num_query_heads=64, num_kv_heads=8, head_size=64, block_size=32, dtype="bf16" +) +_D128_DECODE = dict( + num_query_heads=32, + num_kv_heads=8, + head_size=128, + block_size=16, + max_seqlen_q=1, + max_seqlen_k=8192, + dtype="bf16", +) + +# (id, arch, path, problem_a, problem_b). Each pair resolves to the same +# selector outputs, so the pair must share one binary. +_SHARED_PAIRS = [ + *( + ( + f"{arch}-2d-ns4-vs-ns7", + arch, + "2d", + P( + total_q=4 * 1024, + num_seqs=4, + max_seqlen_q=1024, + max_seqlen_k=1024, + use_fp8=arch == "gfx1250", + **_D64_GQA8, + ), + P( + total_q=7 * 1024, + num_seqs=7, + max_seqlen_q=1024, + max_seqlen_k=1024, + use_fp8=arch == "gfx1250", + **_D64_GQA8, + ), + ) + for arch in ARCHES + ), + # Inside one max_seqlen_q band of the gfx950 combo cohort (512 < q <= 768). + ( + "gfx950-2d-q600-vs-q700", + "gfx950", + "2d", + P(total_q=1200, num_seqs=2, max_seqlen_q=600, max_seqlen_k=600, **_D64_GQA8), + P(total_q=1400, num_seqs=2, max_seqlen_q=700, max_seqlen_k=700, **_D64_GQA8), + ), + *( + ( + f"{arch}-3d-ns4-vs-ns7", + arch, + "3d", + P(total_q=4, num_seqs=4, **_D128_DECODE), + P(total_q=7, num_seqs=7, **_D128_DECODE), + ) + for arch in ("gfx942", "gfx950") + ), + # gfx1250 3D floors the search at 32 iterations, so ns 3 and 4 share a + # binary even though they sit in different log2 buckets. + ( + "gfx1250-3d-ns3-vs-ns4", + "gfx1250", + "3d", + P(total_q=3, num_seqs=3, max_seqlen_q=1, max_seqlen_k=8192, **_D64_GQA8), + P(total_q=4, num_seqs=4, max_seqlen_q=1, max_seqlen_k=8192, **_D64_GQA8), + ), +] + + +@pytest.mark.parametrize("flavor", ("llvm20", "llvm22")) +@pytest.mark.parametrize( + "arch,path,a,b", [p[1:] for p in _SHARED_PAIRS], ids=[p[0] for p in _SHARED_PAIRS] +) +def test_same_outputs_share_key_and_ir(pin_arch, arch, path, a, b, flavor): + pin_arch(arch) + key_fn, lower = ( + (au._tiled_cache_key, _lower_2d) + if path == "2d" + else (au._tiled_3d_cache_key, _lower_3d) + ) + assert a.num_seqs != b.num_seqs or a.max_seqlen_q != b.max_seqlen_q + assert key_fn(a) == key_fn(b) + assert lower(a, arch, flavor) == lower(b, arch, flavor) + + +# --------------------------------------------------------------------------- +# Launch geometry still separates batch counts +# --------------------------------------------------------------------------- + + +def test_shape_key_carries_batch_and_token_counts(): + base = dict(max_seqlen_q=100, max_seqlen_k=100, **_D64_GQA8) + a = P(total_q=400, num_seqs=4, **base) + b = P(total_q=400, num_seqs=7, **base) + c = P(total_q=700, num_seqs=7, **base) + key = ("k",) + assert au._shape_key(key, a) != au._shape_key(key, b) + assert au._shape_key(key, b) != au._shape_key(key, c) + + +def test_launch_meta_separates_num_seqs_under_one_key(pin_arch): + # Same binary (one cache key) and same total_q, but the grid adds one + # padding q-block per sequence, so the launch geometry must differ. + pin_arch("gfx950") + base = dict(total_q=400, max_seqlen_q=100, max_seqlen_k=100, **_D64_GQA8) + a = P(num_seqs=4, **base) + b = P(num_seqs=7, **base) + key = au._tiled_cache_key(a) + assert key == au._tiled_cache_key(b) + au._2D_LAUNCH_META.clear() + grid_a = au._get_2d_launch_meta(a, key).grid + grid_b = au._get_2d_launch_meta(b, key).grid + assert grid_b[1] - grid_a[1] == b.num_seqs - a.num_seqs + + # --------------------------------------------------------------------------- # _binary_search_iters mirrors the spec classes # --------------------------------------------------------------------------- From dc2a8f3946d9a81c1d122351fea63d2de6d0ba5a Mon Sep 17 00:00:00 2001 From: dlangbe Date: Wed, 7 Oct 2026 23:30:35 -0400 Subject: [PATCH 03/16] feat(rocke/attention): add an A/B harness for the tiled cache keys benchmarks/common/tiled_cache_key_ab.py compares two rocke trees, each imported in its own process, on the tiled 2D/3D attention launch path. --mode ir (CPU) routes a sweep that straddles every num_seqs / max_seqlen_q / max_seqlen_k selector threshold like backend="auto", lowers each distinct full spec to LLVM IR, and checks per problem that the binary the launcher cache serves is the one the problem builds from its own spec, and that this IR matches the baseline tree's. It also reports compiled-binary counts per tree. --mode ab (GPU) runs a smaller sweep end to end with arm order alternated across rounds and a fresh process per arm per round, and records first-launch time, steady-state latency, max_abs against an fp32 reference, cache-key host cost and compiled-binary counts. Measured numbers stay out of the repo. Co-Authored-By: Claude --- .../benchmarks/common/tiled_cache_key_ab.py | 827 ++++++++++++++++++ 1 file changed, 827 insertions(+) create mode 100644 dnn-providers/hip-kernel-provider/rocke/library/benchmarks/common/tiled_cache_key_ab.py diff --git a/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/common/tiled_cache_key_ab.py b/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/common/tiled_cache_key_ab.py new file mode 100644 index 000000000000..2ddd1a8e3217 --- /dev/null +++ b/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/common/tiled_cache_key_ab.py @@ -0,0 +1,827 @@ +# Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +# SPDX-License-Identifier: MIT +"""A/B harness for the tiled attention cache keys. + +Compares two rocke trees ("arms"), typically a develop worktree and a branch +that changes ``_tiled_cache_key`` / ``_tiled_3d_cache_key``. Each arm runs in +its own process and imports ``kernels`` and ``rocke`` from its own root, so the +arms never share a module or a launcher cache. The problem sweeps live in this +file, so both arms walk the same problems in the same order. + +``--mode ir`` (CPU only) + Walks a sweep that straddles every num_seqs / max_seqlen_q / max_seqlen_k + selector threshold, routes each problem the way + ``run_unified_attention_torch(backend="auto")`` does, and lowers its + kernel(s) to LLVM IR with the Python lowerer. Per problem it records the + hash of the IR the problem builds from its own spec, and of the IR the + launcher cache serves it: the binary built for the first problem in sweep + order with the same key. Each distinct full spec (raw ``num_seqs`` + included) is lowered once. The candidate arm passes when it serves every + problem its own IR and that IR equals the baseline arm's own IR for the + same problem. The summary also counts the baseline's existing mismatches + and the compiled binaries (distinct keys) per arm. + +``--mode ab`` (GPU) + Runs a smaller sweep end to end through ``run_unified_attention_torch`` + (auto routing) on real tensors. Every arm and round is a fresh process and + the arm order alternates across rounds, so first-launch time includes the + compile on a cache miss. Per problem and arm: first-launch wall time, + steady-state latency (HIP events, median over rounds), max_abs against an + fp32 reference, and host cost of the cache-key call. Per arm: compiled + binary count. + +Outputs go to ``--out``: per-arm raw data plus ``summary.md`` and a CSV. +Measured numbers belong on the ticket or the access-controlled record, never +in the repo or a PR body (see platform/AGENTS.md, Compliance). + +Example, develop worktree vs this checkout: + + python library/benchmarks/common/tiled_cache_key_ab.py --mode ir \\ + --arm base=/dnn-providers/hip-kernel-provider/rocke \\ + --arm branch=/dnn-providers/hip-kernel-provider/rocke \\ + --out + + python library/benchmarks/common/tiled_cache_key_ab.py --mode ab --rounds 3 \\ + --arm base=... --arm branch=... --out +""" + +from __future__ import annotations + +import argparse +import csv +import hashlib +import itertools +import json +import math +import os +import statistics +import subprocess +import sys +import tempfile +import time +from pathlib import Path +from typing import Dict, Iterator, List, Optional, Tuple + +ARCHES = ("gfx942", "gfx950", "gfx1250") + +# Selector thresholds the sweeps straddle. num_seqs: 1 vs >= 2 and the log2 +# search buckets. max_seqlen_q: decode, 64, 128, 256, 512, 768, 1024, 1536. +# max_seqlen_k: 16 * T for T in {32, 64}. +_IR_NUM_SEQS = (1, 2, 3, 4, 7, 8, 16, 64) +_IR_SEQLEN_Q = (1, 2, 64, 65, 128, 129, 256, 257, 512, 513, 768, 769, 1024, 1025, 1536) +_IR_SEQLEN_K = ("q", 512, 513, 1024, 1025, 8192) +_IR_FEATURES = ( + (0, False, None), + (128, False, None), + (0, True, None), + (0, False, "softcap"), + (0, False, "alibi"), +) + + +def _shape_grid(num_seqs, seqlens_q, seqlens_k) -> Iterator[Tuple[int, int, int]]: + """(num_seqs, max_seqlen_q, max_seqlen_k) with k >= q and no duplicates; + "q" in ``seqlens_k`` means k == q.""" + for ns, q in itertools.product(num_seqs, seqlens_q): + for k in sorted({q if k == "q" else k for k in seqlens_k}): + if k >= q: + yield ns, q, k + + +# --------------------------------------------------------------------------- +# Arm bootstrap +# --------------------------------------------------------------------------- + + +def _bootstrap(root: str, arch: Optional[str]): + """Import ``kernels`` from ``root`` and pin the attention arch.""" + root_path = Path(root).resolve() + sys.path[:0] = [str(root_path / "library"), str(root_path / "platform" / "python")] + from kernels.common import attention_unified as au + + loaded = Path(au.__file__).resolve() + if root_path not in loaded.parents: + raise SystemExit(f"kernels imported from {loaded}, not from {root_path}") + if arch is not None: + au._RESOLVED_ATTENTION_ARCH = arch + return au + + +def _route(au, problem) -> Optional[str]: + """Path ``run_unified_attention_torch(backend="auto")`` takes, or None for + the scalar fallback (out of scope here).""" + if problem.select_path() != "2d": + if au.supports_native_unified_attention_3d_tiled(problem)[0]: + return "3d" + if au.supports_native_unified_attention_tiled(problem)[0]: + return "2d" + return None + + +def _key(au, problem, path: str): + if path == "2d": + return au._tiled_cache_key(problem) + return au._tiled_3d_cache_key(problem) + + +def _short_hash(text: str) -> str: + return hashlib.sha256(text.encode()).hexdigest()[:16] + + +# --------------------------------------------------------------------------- +# IR mode +# --------------------------------------------------------------------------- + + +def _ir_sweep(problem_cls, arch: str, num_cus: int) -> Iterator: + for dtype, fp8, hd, bs, heads, (sw, sinks, bias) in itertools.product( + ("bf16", "fp16"), + (False, True), + (64, 128, 256), + (16, 32, 64), + ((8, 8), (32, 8), (64, 8)), + _IR_FEATURES, + ): + for ns, q, k in _shape_grid(_IR_NUM_SEQS, _IR_SEQLEN_Q, _IR_SEQLEN_K): + yield problem_cls( + total_q=ns * q, + num_seqs=ns, + num_query_heads=heads[0], + num_kv_heads=heads[1], + head_size=hd, + block_size=bs, + max_seqlen_q=q, + max_seqlen_k=k, + dtype=dtype, + sliding_window=sw, + softcap=30.0 if bias == "softcap" else 0.0, + use_sinks=sinks, + use_alibi=bias == "alibi", + use_fp8=fp8, + fp8_fnuz=fp8 and arch == "gfx942", + num_cus=num_cus, + ) + + +def _build_identity(au, problem, path: str) -> str: + """Full input of the binary: spec (raw num_seqs kept), builder, backend.""" + if path == "2d": + route = au._gfx942_4warp_route(problem) + return repr( + ( + au._tiled_spec_from_problem(problem), + route.cohort if route else "default", + au._select_2d_compile_backend(problem), + ) + ) + return repr( + ( + au._tiled_3d_spec_from_problem(problem), + au._num_segments(problem), + au._select_3d_waves_per_eu(problem), + ) + ) + + +_POOL_AU = None +_POOL_ARCH = None +_POOL_FLAVOR = None + + +def _pool_init(root: str, arch: str, flavor: str) -> None: + global _POOL_AU, _POOL_ARCH, _POOL_FLAVOR + _POOL_AU = _bootstrap(root, arch) + _POOL_ARCH, _POOL_FLAVOR = arch, flavor + + +def _pool_lower(item) -> str: + """Lower one problem's kernel(s) to LLVM IR; return the IR hash.""" + problem, path = item + au, arch = _POOL_AU, _POOL_ARCH + from rocke.core.lower_llvm import _lower_kernel_to_llvm_python + + if path == "2d": + _, build, _ = au._tiled_2d_impl(arch) + spec = au._tiled_spec_from_problem(problem) + route = au._gfx942_4warp_route(problem) + kernels = [route.builder(spec, arch=arch) if route else build(spec, arch=arch)] + tag = au._select_2d_compile_backend(problem) + else: + _, red_cls, build_seg, build_red, _ = au._tiled_3d_impl(arch) + kernels = [ + build_seg(au._tiled_3d_spec_from_problem(problem), arch=arch), + build_red( + red_cls( + head_size=problem.head_size, + num_query_heads=problem.num_query_heads, + num_kv_heads=problem.num_kv_heads, + dtype=problem.dtype, + num_segments=au._num_segments(problem), + waves_per_eu=au._select_3d_waves_per_eu(problem), + ), + arch=arch, + ), + ] + tag = "3d" + text = "\n".join( + _lower_kernel_to_llvm_python(k, arch=arch, llvm_flavor=_POOL_FLAVOR) + for k in kernels + ) + # The backend picks the compiler for the same IR, so it is part of identity. + return _short_hash(tag + "\n" + text) + + +def _worker_ir(args) -> None: + from concurrent.futures import ProcessPoolExecutor + + au = _bootstrap(args.root, args.arch) + P = au.UnifiedAttentionProblem + paths, keys, builds = [], [], [] + rep_of_build: Dict[str, Tuple] = {} + for problem in _ir_sweep(P, args.arch, args.num_cus): + try: + path = _route(au, problem) + if path is not None: + key = _key(au, problem, path) + build = _build_identity(au, problem, path) + except (ValueError, NotImplementedError): + # Admitted by a gate but rejected by the spec builder; never launches. + path = None + if path is None: + paths.append(None) + keys.append(None) + builds.append(None) + continue + paths.append(path) + keys.append(_short_hash(repr(key))) + builds.append(build) + rep_of_build.setdefault(build, (problem, path)) + + order = list(rep_of_build) + jobs = args.jobs or max(1, (os.cpu_count() or 2) // 2) + t0 = time.time() + with ProcessPoolExecutor( + max_workers=jobs, + initializer=_pool_init, + initargs=(args.root, args.arch, args.flavor), + ) as pool: + hashes = list( + pool.map(_pool_lower, (rep_of_build[b] for b in order), chunksize=8) + ) + ir_of_build = dict(zip(order, hashes)) + + own = [ir_of_build[b] if b is not None else None for b in builds] + first_of_key: Dict[str, str] = {} + served = [] + for k, h in zip(keys, own): + served.append(None if k is None else first_of_key.setdefault(k, h)) + Path(args.result).write_text( + json.dumps( + dict( + arch=args.arch, + root=str(Path(args.root).resolve()), + git=_git_head(args.root), + flavor=args.flavor, + lowered=len(order), + lower_seconds=round(time.time() - t0, 1), + path=paths, + key=keys, + own=own, + served=served, + ) + ) + ) + + +def _summarize_ir(arms: List[Tuple[str, str]], out: Path, arches, num_cus: int) -> int: + base_name, cand_name = arms[0][0], arms[1][0] + lines = [ + "# Tiled cache-key A/B: per-problem IR identity", + "", + f"Baseline arm `{base_name}`, candidate arm `{cand_name}`.", + "", + ] + failed = False + au = None + for arch in arches: + data = { + name: json.loads((out / f"ir_{name}_{arch}.json").read_text()) + for name, _ in arms + } + b, c = data[base_name], data[cand_name] + if au is None: + au = _bootstrap(arms[1][1], None) + P = au.UnifiedAttentionProblem + problems = list(_ir_sweep(P, arch, num_cus)) + n = len(problems) + route_diff = [i for i in range(n) if b["path"][i] != c["path"][i]] + routed = [i for i in range(n) if c["path"][i] is not None] + own_diff = [i for i in routed if b["own"][i] != c["own"][i]] + cand_bad = [i for i in routed if c["served"][i] != c["own"][i]] + base_bad = [i for i in routed if b["served"][i] != b["own"][i]] + vs_base_own = [i for i in routed if c["served"][i] != b["own"][i]] + failed |= bool(route_diff or own_diff or cand_bad or vs_base_own) + lines += [ + f"## {arch}", + "", + f"- git: {base_name} `{b['git']}`, {cand_name} `{c['git']}`;" + f" flavor {c['flavor']}; Python lowerer", + f"- problems routed to a tiled path: {len(routed)} of {n}" + f" (2D {sum(c['path'][i] == '2d' for i in routed)}," + f" 3D {sum(c['path'][i] == '3d' for i in routed)})", + f"- distinct specs lowered: {base_name} {b['lowered']}," + f" {cand_name} {c['lowered']}", + "", + "| check | count | expected |", + "|---|---|---|", + f"| route differs between arms | {len(route_diff)} | 0 |", + f"| own IR differs between arms | {len(own_diff)} | 0 |", + f"| {cand_name} serves IR != own | {len(cand_bad)} | 0 |", + f"| {cand_name} serves IR != {base_name} own | {len(vs_base_own)} | 0 |", + f"| {base_name} serves IR != own | {len(base_bad)} | existing gaps |", + "", + "| path | binaries " + base_name + " | binaries " + cand_name + " |", + "|---|---|---|", + ] + for path in ("2d", "3d"): + nb = len({b["key"][i] for i in routed if b["path"][i] == path}) + nc = len({c["key"][i] for i in routed if c["path"][i] == path}) + lines.append(f"| {path} | {nb} | {nc} |") + lines.append("") + for label, idx in ( + ("route differs", route_diff), + ("own IR differs", own_diff), + (f"{cand_name} serves != own", cand_bad), + (f"{base_name} serves != own", base_bad), + ): + if idx: + lines.append(f"First {label} examples:") + lines += [f"- `{problems[i]}`" for i in idx[:5]] + lines.append("") + lines.append("RESULT: " + ("FAIL" if failed else "PASS")) + (out / "summary.md").write_text("\n".join(lines) + "\n") + print("\n".join(lines)) + return 1 if failed else 0 + + +# --------------------------------------------------------------------------- +# A/B mode (GPU) +# --------------------------------------------------------------------------- + +# (label, static fields, num_seqs values, max_seqlen_q values, max_seqlen_k +# values). "q" means max_seqlen_k == max_seqlen_q. +_AB_COHORTS = ( + ( + "d64-gqa8-bs32-bf16", + dict(num_query_heads=64, num_kv_heads=8, head_size=64, block_size=32), + (1, 2, 3, 4, 7, 8, 16), + (64, 257, 600, 1025, 2048), + ("q",), + ), + ( + "d128-gqa4-bs16-bf16", + dict(num_query_heads=32, num_kv_heads=8, head_size=128, block_size=16), + (1, 2, 3, 4, 8), + (1, 257, 1025, 2048), + ("q", 4096), + ), + ( + "d64-mha-bs16-fp16", + dict( + num_query_heads=8, num_kv_heads=8, head_size=64, block_size=16, dtype="fp16" + ), + (1, 2, 4, 7), + (129, 513, 1536), + ("q",), + ), + ( + "d64-gqa8-bs16-fp8kv", + dict( + num_query_heads=64, + num_kv_heads=8, + head_size=64, + block_size=16, + use_fp8=True, + ), + (16, 64), + (1, 64), + (512, 513, 4096), + ), + ( + "d128-gqa4-bs32-bf16-sw128", + dict( + num_query_heads=32, + num_kv_heads=8, + head_size=128, + block_size=32, + sliding_window=128, + ), + (1, 2, 4), + (1025,), + ("q",), + ), + ( + "d64-gqa8-bs32-bf16-sinks", + dict( + num_query_heads=64, + num_kv_heads=8, + head_size=64, + block_size=32, + use_sinks=True, + ), + (1, 2), + (1025,), + ("q",), + ), +) + + +def _ab_sweep(problem_cls, arch: str, num_cus: int) -> Iterator[Tuple[str, object]]: + for label, static, nss, qs, ks in _AB_COHORTS: + static = {"dtype": "bf16", **static} + if static.get("use_fp8"): + static["fp8_fnuz"] = arch == "gfx942" + for ns, q, k in _shape_grid(nss, qs, ks): + yield label, problem_cls( + total_q=ns * q, + num_seqs=ns, + max_seqlen_q=q, + max_seqlen_k=k, + num_cus=num_cus, + **static, + ) + + +def _make_inputs(problem, seed: int, device: str = "cuda"): + """Paged inputs with every sequence at (max_seqlen_q, max_seqlen_k), plus + their fp32 upcasts for the reference.""" + import torch + + torch.manual_seed(seed) + io = torch.float16 if problem.dtype == "fp16" else torch.bfloat16 + ns, q_len, kv_len, bs = ( + problem.num_seqs, + problem.max_seqlen_q, + problem.max_seqlen_k, + problem.block_size, + ) + blocks_per_seq = (kv_len + bs - 1) // bs + num_blocks = ns * blocks_per_seq + query = torch.randn( + problem.total_q, + problem.num_query_heads, + problem.head_size, + dtype=io, + device=device, + ) + cache_shape = (num_blocks, bs, problem.num_kv_heads, problem.head_size) + key_cache = torch.randn(cache_shape, dtype=io, device=device) + value_cache = torch.randn(cache_shape, dtype=io, device=device) + if problem.use_fp8: + fp8 = torch.float8_e4m3fnuz if problem.fp8_fnuz else torch.float8_e4m3fn + key_cache, value_cache = key_cache.to(fp8), value_cache.to(fp8) + block_tables = ( + torch.randperm(num_blocks, device=device, dtype=torch.int64) + .to(torch.int32) + .view(ns, blocks_per_seq) + ) + cu_q = torch.arange(0, ns + 1, dtype=torch.int32, device=device) * q_len + kv_lens = torch.full((ns,), kv_len, dtype=torch.int32, device=device) + sinks = ( + torch.randn(problem.num_query_heads, dtype=io, device=device) + if problem.use_sinks + else None + ) + return dict( + query=query, + key_cache=key_cache, + value_cache=value_cache, + block_tables=block_tables, + cu_q=cu_q, + kv_lens=kv_lens, + sinks=sinks, + scale=problem.head_size**-0.5, + out=torch.empty_like(query), + ) + + +def _reference_fp32(problem, data): + from builders.gfx950.attention.prefill.parity_unified_attention import ( + ref_paged_attn, + ) + + return ref_paged_attn( + query=data["query"].float(), + key_cache=data["key_cache"].float(), + value_cache=data["value_cache"].float(), + query_lens=[problem.max_seqlen_q] * problem.num_seqs, + kv_lens=[problem.max_seqlen_k] * problem.num_seqs, + block_tables=data["block_tables"], + scale=data["scale"], + sliding_window=problem.sliding_window or None, + sinks=data["sinks"].float() if data["sinks"] is not None else None, + ) + + +def _binaries(au) -> int: + return len(au._ATTN_TILED_CACHE) + len(au._ATTN_3D_TILED_CACHE) + + +def _worker_ab(args) -> None: + import torch + + if not torch.cuda.is_available(): + raise SystemExit("--mode ab needs a HIP-visible GPU") + from rocke.runtime.hip_module import get_device_arch + + arch = get_device_arch() + au = _bootstrap(args.root, None) + from kernels import run_unified_attention_torch + from rocke.runtime import synchronize_and_release, time_launches + + num_cus = args.num_cus or torch.cuda.get_device_properties(0).multi_processor_count + stream = int(torch.cuda.current_stream().cuda_stream) + rows = [] + for idx, (label, problem) in enumerate( + _ab_sweep(au.UnifiedAttentionProblem, arch, num_cus) + ): + path = _route(au, problem) + if path is None: + continue + data = _make_inputs(problem, seed=idx) + + def call(): + run_unified_attention_torch( + problem=problem, + q=data["query"], + k=data["key_cache"], + v=data["value_cache"], + out=data["out"], + cu_seqlens_q=data["cu_q"], + seqused_k=data["kv_lens"], + softmax_scale=data["scale"], + block_table=data["block_tables"], + softcap=0.0, + sinks=data["sinks"], + backend="auto", + stream=stream, + ) + + before = _binaries(au) + torch.cuda.synchronize() + t0 = time.perf_counter() + call() + torch.cuda.synchronize() + first_ms = (time.perf_counter() - t0) * 1e3 + compiled = _binaries(au) - before + + ref = _reference_fp32(problem, data) + max_abs = float((data["out"].float() - ref).abs().max().item()) + + steady_ms = time_launches( + call, warmup=args.warmup, iters=args.iters, stream=stream + ) + synchronize_and_release(stream) + + key_us = [] + for _ in range(5): + t0 = time.perf_counter() + for _ in range(1000): + _key(au, problem, path) + key_us.append((time.perf_counter() - t0) * 1e3) + rows.append( + dict( + idx=idx, + cohort=label, + num_seqs=problem.num_seqs, + max_seqlen_q=problem.max_seqlen_q, + max_seqlen_k=problem.max_seqlen_k, + path=path, + first_ms=first_ms, + compiled=compiled, + steady_ms=steady_ms, + max_abs=max_abs, + key_us=statistics.median(key_us), + ) + ) + del data, ref + from rocke.core.backend import resolve_backend + + Path(args.result).write_text( + json.dumps( + dict( + arch=arch, + root=str(Path(args.root).resolve()), + git=_git_head(args.root), + backend=resolve_backend(), + num_cus=num_cus, + binaries=_binaries(au), + rows=rows, + ) + ) + ) + + +def _summarize_ab(arms: List[Tuple[str, str]], out: Path, rounds: int) -> int: + base_name, cand_name = arms[0][0], arms[1][0] + runs = { + name: [ + json.loads((out / f"ab_{name}_r{r}.json").read_text()) + for r in range(rounds) + ] + for name, _ in arms + } + meta = runs[cand_name][0] + if not meta["rows"]: + print(f"no problem routed to a tiled path on {meta['arch']}") + return 1 + table = [] + for row in meta["rows"]: + idx = row["idx"] + rec = dict( + idx=idx, + cohort=row["cohort"], + num_seqs=row["num_seqs"], + max_seqlen_q=row["max_seqlen_q"], + max_seqlen_k=row["max_seqlen_k"], + path=row["path"], + ) + for name in (base_name, cand_name): + hits = [ + next(x for x in run["rows"] if x["idx"] == idx) for run in runs[name] + ] + rec[f"{name}_steady_ms"] = statistics.median(h["steady_ms"] for h in hits) + rec[f"{name}_first_ms"] = statistics.median(h["first_ms"] for h in hits) + rec[f"{name}_compiled"] = hits[0]["compiled"] + rec[f"{name}_max_abs"] = max(h["max_abs"] for h in hits) + rec[f"{name}_key_us"] = statistics.median(h["key_us"] for h in hits) + rec["ratio"] = rec[f"{cand_name}_steady_ms"] / rec[f"{base_name}_steady_ms"] + table.append(rec) + + with (out / "ab.csv").open("w", newline="") as fh: + writer = csv.DictWriter(fh, fieldnames=list(table[0])) + writer.writeheader() + writer.writerows(table) + + def total(name, field): + return sum(r[f"{name}_{field}"] for r in table) + + geomean = math.exp(statistics.fmean(math.log(r["ratio"]) for r in table)) + lines = [ + "# Tiled cache-key A/B: latency, compiles, correctness", + "", + f"- arch {meta['arch']}, num_cus {meta['num_cus']}, {rounds} rounds, arm order" + " alternated per round, one fresh process per arm per round", + f"- {base_name}: `{runs[base_name][0]['git']}`" + f" (lowering {runs[base_name][0]['backend']});" + f" {cand_name}: `{runs[cand_name][0]['git']}`" + f" (lowering {runs[cand_name][0]['backend']})", + f"- problems: {len(table)} (steady-state = median over rounds of the HIP-event" + " mean per call)", + "", + f"| metric | {base_name} | {cand_name} |", + "|---|---|---|", + "| compiled binaries (per process) | " + + " / ".join(str(run["binaries"]) for run in runs[base_name]) + + " | " + + " / ".join(str(run["binaries"]) for run in runs[cand_name]) + + " |", + f"| sum first-launch ms | {total(base_name, 'first_ms'):.1f} |" + f" {total(cand_name, 'first_ms'):.1f} |", + f"| max max_abs vs fp32 ref | {max(r[f'{base_name}_max_abs'] for r in table):.3e} |" + f" {max(r[f'{cand_name}_max_abs'] for r in table):.3e} |", + f"| median key us/call | " + f"{statistics.median(r[f'{base_name}_key_us'] for r in table):.2f} | " + f"{statistics.median(r[f'{cand_name}_key_us'] for r in table):.2f} |", + "", + f"Steady-state geomean {cand_name}/{base_name}: {geomean:.4f};" + f" problems > 1.05x slower: {sum(r['ratio'] > 1.05 for r in table)};" + f" < 0.95x: {sum(r['ratio'] < 0.95 for r in table)}.", + "", + "| cohort | ns | q | k | path | " + f"{base_name} ms | {cand_name} ms | ratio | " + f"{base_name} max_abs | {cand_name} max_abs | " + f"{base_name} compiled | {cand_name} compiled |", + "|---|---|---|---|---|---|---|---|---|---|---|---|", + ] + for r in table: + lines.append( + f"| {r['cohort']} | {r['num_seqs']} | {r['max_seqlen_q']} |" + f" {r['max_seqlen_k']} | {r['path']} |" + f" {r[f'{base_name}_steady_ms']:.4f} | {r[f'{cand_name}_steady_ms']:.4f} |" + f" {r['ratio']:.3f} | {r[f'{base_name}_max_abs']:.2e} |" + f" {r[f'{cand_name}_max_abs']:.2e} | {r[f'{base_name}_compiled']} |" + f" {r[f'{cand_name}_compiled']} |" + ) + (out / "summary.md").write_text("\n".join(lines) + "\n") + print("\n".join(lines[:16])) + return 0 + + +# --------------------------------------------------------------------------- +# Driver +# --------------------------------------------------------------------------- + + +def _git_head(root: str) -> str: + try: + return subprocess.run( + ["git", "-C", root, "rev-parse", "--short=12", "HEAD"], + capture_output=True, + text=True, + check=True, + ).stdout.strip() + ( + "+dirty" + if subprocess.run( + ["git", "-C", root, "status", "--porcelain", "--", "."], + capture_output=True, + text=True, + ).stdout.strip() + else "" + ) + except (OSError, subprocess.CalledProcessError): + return "unknown" + + +def _spawn(worker: str, root: str, result: Path, extra: List[str]) -> None: + cmd = [ + sys.executable, + str(Path(__file__).resolve()), + "--worker", + worker, + "--root", + root, + "--result", + str(result), + *extra, + ] + subprocess.run(cmd, check=True) + + +def main(argv: Optional[List[str]] = None) -> int: + ap = argparse.ArgumentParser(description=__doc__.split("\n\n")[0]) + ap.add_argument("--mode", choices=("ir", "ab")) + ap.add_argument( + "--arm", + action="append", + default=[], + metavar="NAME=ROOT", + help="rocke root of an arm; give exactly two, baseline first", + ) + ap.add_argument("--out", type=Path, default=None) + ap.add_argument("--arch", action="append", choices=ARCHES, help="ir mode arches") + ap.add_argument("--flavor", default="llvm22", choices=("llvm20", "llvm22")) + ap.add_argument("--num-cus", type=int, default=0, help="0: device (ab), 120 (ir)") + ap.add_argument("--jobs", type=int, default=0, help="ir lowering processes") + ap.add_argument("--rounds", type=int, default=3) + ap.add_argument("--iters", type=int, default=50) + ap.add_argument("--warmup", type=int, default=10) + # Internal: one arm in one process. + ap.add_argument("--worker", choices=("ir", "ab"), help=argparse.SUPPRESS) + ap.add_argument("--root", help=argparse.SUPPRESS) + ap.add_argument("--result", help=argparse.SUPPRESS) + args = ap.parse_args(argv) + + if args.worker == "ir": + (args.arch,) = args.arch + args.num_cus = args.num_cus or 120 + _worker_ir(args) + return 0 + if args.worker == "ab": + _worker_ab(args) + return 0 + + if args.mode is None or len(args.arm) != 2: + ap.error("--mode and exactly two --arm NAME=ROOT are required") + arms = [tuple(a.split("=", 1)) for a in args.arm] + out = args.out or Path(tempfile.mkdtemp(prefix="tiled_cache_key_ab_")) + out.mkdir(parents=True, exist_ok=True) + print(f"writing to {out}") + + if args.mode == "ir": + arches = args.arch or list(ARCHES) + common = ["--flavor", args.flavor, "--jobs", str(args.jobs)] + common += ["--num-cus", str(args.num_cus or 120)] + for arch in arches: + for name, root in arms: + print(f"[ir] {name} {arch}", flush=True) + _spawn( + "ir", + root, + out / f"ir_{name}_{arch}.json", + ["--arch", arch, *common], + ) + return _summarize_ir(arms, out, arches, args.num_cus or 120) + + common = ["--iters", str(args.iters), "--warmup", str(args.warmup)] + common += ["--num-cus", str(args.num_cus)] + for r in range(args.rounds): + order = arms if r % 2 == 0 else arms[::-1] + for name, root in order: + print(f"[ab] round {r} arm {name}", flush=True) + _spawn("ab", root, out / f"ab_{name}_r{r}.json", common) + return _summarize_ab(arms, out, args.rounds) + + +if __name__ == "__main__": + sys.exit(main()) From 9544316699ff164adbea126cb7ccec5b2c8b1d82 Mon Sep 17 00:00:00 2001 From: dlangbe Date: Thu, 8 Oct 2026 02:38:17 -0400 Subject: [PATCH 04/16] fix(rocke/attention): keep garbage collection out of the tiled cache-key A/B timings --- .../benchmarks/common/tiled_cache_key_ab.py | 28 ++++++++++++++++--- 1 file changed, 24 insertions(+), 4 deletions(-) diff --git a/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/common/tiled_cache_key_ab.py b/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/common/tiled_cache_key_ab.py index 2ddd1a8e3217..e2c8e45b2b19 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/common/tiled_cache_key_ab.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/common/tiled_cache_key_ab.py @@ -28,7 +28,8 @@ compile on a cache miss. Per problem and arm: first-launch wall time, steady-state latency (HIP events, median over rounds), max_abs against an fp32 reference, and host cost of the cache-key call. Per arm: compiled - binary count. + binary count. Python's garbage collector is off inside each worker and + runs between problems, so a collection pass never lands in a timed region. Outputs go to ``--out``: per-arm raw data plus ``summary.md`` and a CSV. Measured numbers belong on the ticket or the access-controlled record, never @@ -49,6 +50,7 @@ import argparse import csv +import gc import hashlib import itertools import json @@ -542,12 +544,18 @@ def _worker_ab(args) -> None: num_cus = args.num_cus or torch.cuda.get_device_properties(0).multi_processor_count stream = int(torch.cuda.current_stream().cuda_stream) rows = [] + # A collection pass inside a timed loop stalls the host while the GPU idles, + # which reads as a per-launch slowdown, and where the collector fires + # depends on each arm's allocation history. Keep it off and collect between + # problems instead, as timeit does. + gc.disable() for idx, (label, problem) in enumerate( _ab_sweep(au.UnifiedAttentionProblem, arch, num_cus) ): path = _route(au, problem) if path is None: continue + gc.collect() data = _make_inputs(problem, seed=idx) def call(): @@ -605,15 +613,20 @@ def call(): ) ) del data, ref - from rocke.core.backend import resolve_backend + from rocke.core.backend import cpp_fallbacks, resolve_backend + # Report the lowering that ran, not the configured one: without the + # rocke_engine extension every cpp lowering falls back to Python. + backend = resolve_backend() + if backend == "cpp" and cpp_fallbacks(): + backend = f"python (cpp fell back on {len(cpp_fallbacks())} kernels)" Path(args.result).write_text( json.dumps( dict( arch=arch, root=str(Path(args.root).resolve()), git=_git_head(args.root), - backend=resolve_backend(), + backend=backend, num_cus=num_cus, binaries=_binaries(au), rows=rows, @@ -740,7 +753,14 @@ def _git_head(root: str) -> str: else "" ) except (OSError, subprocess.CalledProcessError): - return "unknown" + # No git metadata (an exported tree, or a container that mounts only + # part of the checkout): fingerprint the dispatcher source instead so + # the arms stay identifiable. + src = Path(root) / "library" / "kernels" / "common" / "attention_unified.py" + try: + return "src:" + hashlib.sha256(src.read_bytes()).hexdigest()[:12] + except OSError: + return "unknown" def _spawn(worker: str, root: str, result: Path, extra: List[str]) -> None: From 09297f93c7d3d2f91f053a7beca8ee5461b4ccb7 Mon Sep 17 00:00:00 2001 From: dlangbe Date: Thu, 8 Oct 2026 12:42:32 -0400 Subject: [PATCH 05/16] docs(rocke/attention): require new attention selectors in the tiled cache key --- .../hip-kernel-provider/rocke/KERNEL_AUTHORING.md | 8 ++++++++ .../rocke/library/builders/common/README.md | 9 +++++++++ 2 files changed, 17 insertions(+) diff --git a/dnn-providers/hip-kernel-provider/rocke/KERNEL_AUTHORING.md b/dnn-providers/hip-kernel-provider/rocke/KERNEL_AUTHORING.md index 9af9ee5c4cc7..7ba4ef4551c4 100644 --- a/dnn-providers/hip-kernel-provider/rocke/KERNEL_AUTHORING.md +++ b/dnn-providers/hip-kernel-provider/rocke/KERNEL_AUTHORING.md @@ -281,6 +281,14 @@ documented in [library/builders/common/README.md](library/builders/common/README `_enable_()` predicate, following the branch order documented in [library/builders/common/README.md](library/builders/common/README.md). Default-OFF opt-in knobs may stay unwired (env/flag-only) until a sweep proves a win. + **Add the predicate's resolved output to the launcher cache key** + (`_tiled_cache_key` / `_tiled_3d_cache_key` in + [attention_unified.py](library/kernels/common/attention_unified.py)). The key is + computed per launch without building the spec, so a selector missing from it lets + two shapes that build different kernels share one cached binary. + [test_attention_tiled_cache_key.py](library/tests/test_attention_tiled_cache_key.py) + sweeps the selector thresholds and fails when key equality and built-spec equality + disagree. 5. **Document the knob** in the "Knob reference" section of [library/builders/common/README.md](library/builders/common/README.md): what it rewrites, its correctness guarantee, its measured effect, and its enable condition. diff --git a/dnn-providers/hip-kernel-provider/rocke/library/builders/common/README.md b/dnn-providers/hip-kernel-provider/rocke/library/builders/common/README.md index c07034ad7d0b..1269a0a52fd8 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/builders/common/README.md +++ b/dnn-providers/hip-kernel-provider/rocke/library/builders/common/README.md @@ -81,6 +81,15 @@ All remaining shapes. Key knob decisions: ### Knob reference +Every knob a selector resolves must also reach the launcher cache key +(`_tiled_cache_key` / `_tiled_3d_cache_key` in +[`attention_unified.py`](../../kernels/common/attention_unified.py)): add the +selector's resolved output there when you wire it. The key is computed per launch +without building the spec, so a missing entry lets two shapes that build different +kernels share one cached binary. +[`test_attention_tiled_cache_key.py`](../../tests/test_attention_tiled_cache_key.py) +fails when key equality and built-spec equality disagree. + **`use_transposed_half_local_pv`** — Rewrites the PV phase of the transposed 32x32 kernel so each 32-lane warp-half reads only the P rows it owns, eliminating the cross-half `lane^32` shuffle fetch. Bit-identical to the From 4d980283bdd0f833900b69a917e50612408c7d27 Mon Sep 17 00:00:00 2001 From: dlangbe Date: Thu, 8 Oct 2026 12:43:28 -0400 Subject: [PATCH 06/16] refactor(rocke/attention): modernize typing and imports in the tiled cache-key test and harness --- .../benchmarks/common/tiled_cache_key_ab.py | 30 ++++++++----- .../kernels/common/attention_unified.py | 2 +- .../tests/test_attention_tiled_cache_key.py | 45 +++++++++---------- 3 files changed, 40 insertions(+), 37 deletions(-) diff --git a/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/common/tiled_cache_key_ab.py b/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/common/tiled_cache_key_ab.py index e2c8e45b2b19..7c8aec566574 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/common/tiled_cache_key_ab.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/common/tiled_cache_key_ab.py @@ -61,8 +61,8 @@ import sys import tempfile import time +from collections.abc import Iterator from pathlib import Path -from typing import Dict, Iterator, List, Optional, Tuple ARCHES = ("gfx942", "gfx950", "gfx1250") @@ -81,7 +81,7 @@ ) -def _shape_grid(num_seqs, seqlens_q, seqlens_k) -> Iterator[Tuple[int, int, int]]: +def _shape_grid(num_seqs, seqlens_q, seqlens_k) -> Iterator[tuple[int, int, int]]: """(num_seqs, max_seqlen_q, max_seqlen_k) with k >= q and no duplicates; "q" in ``seqlens_k`` means k == q.""" for ns, q in itertools.product(num_seqs, seqlens_q): @@ -95,8 +95,14 @@ def _shape_grid(num_seqs, seqlens_q, seqlens_k) -> Iterator[Tuple[int, int, int] # --------------------------------------------------------------------------- -def _bootstrap(root: str, arch: Optional[str]): - """Import ``kernels`` from ``root`` and pin the attention arch.""" +def _bootstrap(root: str, arch: str | None): + """Import ``kernels`` from ``root`` and pin the attention arch. + + Prepending the arm's tree to ``sys.path`` is the point of the harness: each + arm is its own process importing its own checkout, so the comparison never + mixes modules or caches. The root comes from ``--arm``, never a fixed path, + and the import is verified to resolve under it. + """ root_path = Path(root).resolve() sys.path[:0] = [str(root_path / "library"), str(root_path / "platform" / "python")] from kernels.common import attention_unified as au @@ -109,7 +115,7 @@ def _bootstrap(root: str, arch: Optional[str]): return au -def _route(au, problem) -> Optional[str]: +def _route(au, problem) -> str | None: """Path ``run_unified_attention_torch(backend="auto")`` takes, or None for the scalar fallback (out of scope here).""" if problem.select_path() != "2d": @@ -239,7 +245,7 @@ def _worker_ir(args) -> None: au = _bootstrap(args.root, args.arch) P = au.UnifiedAttentionProblem paths, keys, builds = [], [], [] - rep_of_build: Dict[str, Tuple] = {} + rep_of_build: dict[str, tuple] = {} for problem in _ir_sweep(P, args.arch, args.num_cus): try: path = _route(au, problem) @@ -273,7 +279,7 @@ def _worker_ir(args) -> None: ir_of_build = dict(zip(order, hashes)) own = [ir_of_build[b] if b is not None else None for b in builds] - first_of_key: Dict[str, str] = {} + first_of_key: dict[str, str] = {} served = [] for k, h in zip(keys, own): served.append(None if k is None else first_of_key.setdefault(k, h)) @@ -295,7 +301,7 @@ def _worker_ir(args) -> None: ) -def _summarize_ir(arms: List[Tuple[str, str]], out: Path, arches, num_cus: int) -> int: +def _summarize_ir(arms: list[tuple[str, str]], out: Path, arches, num_cus: int) -> int: base_name, cand_name = arms[0][0], arms[1][0] lines = [ "# Tiled cache-key A/B: per-problem IR identity", @@ -438,7 +444,7 @@ def _summarize_ir(arms: List[Tuple[str, str]], out: Path, arches, num_cus: int) ) -def _ab_sweep(problem_cls, arch: str, num_cus: int) -> Iterator[Tuple[str, object]]: +def _ab_sweep(problem_cls, arch: str, num_cus: int) -> Iterator[tuple[str, object]]: for label, static, nss, qs, ks in _AB_COHORTS: static = {"dtype": "bf16", **static} if static.get("use_fp8"): @@ -635,7 +641,7 @@ def call(): ) -def _summarize_ab(arms: List[Tuple[str, str]], out: Path, rounds: int) -> int: +def _summarize_ab(arms: list[tuple[str, str]], out: Path, rounds: int) -> int: base_name, cand_name = arms[0][0], arms[1][0] runs = { name: [ @@ -763,7 +769,7 @@ def _git_head(root: str) -> str: return "unknown" -def _spawn(worker: str, root: str, result: Path, extra: List[str]) -> None: +def _spawn(worker: str, root: str, result: Path, extra: list[str]) -> None: cmd = [ sys.executable, str(Path(__file__).resolve()), @@ -778,7 +784,7 @@ def _spawn(worker: str, root: str, result: Path, extra: List[str]) -> None: subprocess.run(cmd, check=True) -def main(argv: Optional[List[str]] = None) -> int: +def main(argv: list[str] | None = None) -> int: ap = argparse.ArgumentParser(description=__doc__.split("\n\n")[0]) ap.add_argument("--mode", choices=("ir", "ab")) ap.add_argument( diff --git a/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/attention_unified.py b/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/attention_unified.py index 49ea1347a48f..b45096520448 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/attention_unified.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/attention_unified.py @@ -4026,7 +4026,7 @@ def attention_3d_workspace_nbytes( ) -def _shape_key(cache_key: Tuple, problem: UnifiedAttentionProblem) -> Tuple: +def _shape_key(cache_key: tuple, problem: UnifiedAttentionProblem) -> tuple: """Extend a binary cache key with the shape values the launch derives from. The tiled binaries are shared across batch counts and token counts, but the diff --git a/dnn-providers/hip-kernel-provider/rocke/library/tests/test_attention_tiled_cache_key.py b/dnn-providers/hip-kernel-provider/rocke/library/tests/test_attention_tiled_cache_key.py index 081d8419d4c9..ba878fc15a1b 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/tests/test_attention_tiled_cache_key.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/tests/test_attention_tiled_cache_key.py @@ -32,6 +32,22 @@ import pytest from kernels.common import attention_unified as au +from kernels.gfx1250.attention_tiled_2d import ( + UnifiedAttention2DTiledSpec as Gfx1250Spec2D, +) +from kernels.gfx942.attention_tiled_2d import ( + UnifiedAttention2DTiledSpec as Gfx942Spec2D, +) +from kernels.gfx942.attention_tiled_3d import ( + UnifiedAttention3DTiledSpec as Gfx942Spec3D, +) +from kernels.gfx950.attention_tiled_2d import ( + UnifiedAttention2DTiledSpec as Gfx950Spec2D, +) +from kernels.gfx950.attention_tiled_3d import ( + UnifiedAttention3DTiledSpec as Gfx950Spec3D, +) +from rocke.core.lower_llvm import LLVM_FLAVORS, _lower_kernel_to_llvm_python P = au.UnifiedAttentionProblem @@ -307,8 +323,6 @@ def test_single_batch_interleave_splits_key(pin_arch): def _lower_2d(problem, arch, flavor): - from rocke.core.lower_llvm import _lower_kernel_to_llvm_python - _, build, _ = au._tiled_2d_impl(arch) spec = au._tiled_spec_from_problem(problem) route = au._gfx942_4warp_route(problem) @@ -317,8 +331,6 @@ def _lower_2d(problem, arch, flavor): def _lower_3d(problem, arch, flavor): - from rocke.core.lower_llvm import _lower_kernel_to_llvm_python - _, reduce_spec_cls, build_seg, build_red, _ = au._tiled_3d_impl(arch) seg = build_seg(au._tiled_3d_spec_from_problem(problem), arch=arch) red = build_red( @@ -408,7 +420,7 @@ def _lower_3d(problem, arch, flavor): ] -@pytest.mark.parametrize("flavor", ("llvm20", "llvm22")) +@pytest.mark.parametrize("flavor", LLVM_FLAVORS) @pytest.mark.parametrize( "arch,path,a,b", [p[1:] for p in _SHARED_PAIRS], ids=[p[0] for p in _SHARED_PAIRS] ) @@ -459,28 +471,13 @@ def test_launch_meta_separates_num_seqs_under_one_key(pin_arch): # --------------------------------------------------------------------------- -def _spec_classes(): - from kernels.gfx942.attention_tiled_2d import ( - UnifiedAttention2DTiledSpec as Gfx942Spec2D, - ) - from kernels.gfx942.attention_tiled_3d import ( - UnifiedAttention3DTiledSpec as Gfx942Spec3D, - ) - from kernels.gfx950.attention_tiled_2d import ( - UnifiedAttention2DTiledSpec as Gfx950Spec2D, - ) - from kernels.gfx950.attention_tiled_3d import ( - UnifiedAttention3DTiledSpec as Gfx950Spec3D, - ) - from kernels.gfx1250.attention_tiled_2d import ( - UnifiedAttention2DTiledSpec as Gfx1250Spec2D, - ) - - return (Gfx942Spec2D, Gfx942Spec3D, Gfx950Spec2D, Gfx950Spec3D, Gfx1250Spec2D) +# gfx1250 3D floors the search at 32 iterations, so it is checked through the key +# (test_same_outputs_share_key_and_ir) rather than mirrored here. +_SPEC_CLASSES = (Gfx942Spec2D, Gfx942Spec3D, Gfx950Spec2D, Gfx950Spec3D, Gfx1250Spec2D) @pytest.mark.parametrize("num_seqs", (0, 1, 2, 3, 4, 7, 8, 15, 16, 450, 1 << 20)) -@pytest.mark.parametrize("spec_cls", _spec_classes(), ids=lambda c: c.__module__) +@pytest.mark.parametrize("spec_cls", _SPEC_CLASSES, ids=lambda c: c.__module__) def test_binary_search_iters_mirrors_spec(spec_cls, num_seqs): # binary_search_iters is a pure function of num_seqs; bypass __init__ # validation so the comparison needs no full geometry. From 8c66efcf0423d3b403de0955a7a1d3d51fb43a14 Mon Sep 17 00:00:00 2001 From: dlangbe Date: Thu, 8 Oct 2026 15:24:56 -0400 Subject: [PATCH 07/16] perf(rocke/attention): reuse resolved geometry in the tiled cache key The tiled 2D cache key is computed on every launch. Two of its entries re-ran selectors whose outputs the key already holds: _enable_k_single_buffer re-walked the num-warps, rows-per-warp and tile selectors for its block_m <= tile_size check, and _enable_fp8_mfma_qk re-walked the tile and num-warps selectors for its loader-fit check. On short, host-bound launches that recomputation was visible per call. _tiled_cache_key now resolves num-warps, tile size and rows-per-warp once and passes them to both predicates (and to _fp8_qk_loader_fits) through new keyword-only parameters. The parameters default to None, in which case each helper resolves the value itself, so the spec builder and every other caller are unchanged. The generic 2D warp count is only passed when it is the count the key records; on the gfx942 fp16 flash path the predicates still resolve it themselves. This is an exact rewrite: the key tuple is identical for every problem in the cache-key test sweep and its fp8 / single-batch extensions on gfx942, gfx950 and gfx1250, and test_attention_tiled_cache_key still enforces key equality <=> build equality in both directions. Co-Authored-By: Claude --- .../kernels/common/attention_unified.py | 78 ++++++++++++++----- 1 file changed, 57 insertions(+), 21 deletions(-) diff --git a/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/attention_unified.py b/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/attention_unified.py index b45096520448..335ccbb5dc49 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/attention_unified.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/attention_unified.py @@ -676,7 +676,13 @@ def _enable_d128_small_tile(problem: UnifiedAttentionProblem) -> bool: ) -def _enable_k_single_buffer(problem: UnifiedAttentionProblem) -> bool: +def _enable_k_single_buffer( + problem: UnifiedAttentionProblem, + *, + num_warps: int | None = None, + block_m_per_warp: int | None = None, + tile_size: int | None = None, +) -> bool: """d128 long-context lever: K single-buffer at T=64 (== 2*block_size). For the same single-batch d128 cohort that ``_enable_d128_small_tile`` gates, @@ -704,9 +710,15 @@ def _enable_k_single_buffer(problem: UnifiedAttentionProblem) -> bool: if not _enable_d128_small_tile(problem): return False # block_m <= tile_size, asked of the selectors that actually set the geometry - # (no recursion: none of these consult _enable_k_single_buffer). - block_m = _select_2d_num_warps(problem) * _select_2d_block_m_per_warp(problem) - return block_m <= _select_2d_tile_size(problem) + # (no recursion: none of these consult _enable_k_single_buffer). The cache key + # passes the values it already resolved instead of walking them again. + if num_warps is None: + num_warps = _select_2d_num_warps(problem) + if block_m_per_warp is None: + block_m_per_warp = _select_2d_block_m_per_warp(problem) + if tile_size is None: + tile_size = _select_2d_tile_size(problem) + return num_warps * block_m_per_warp <= tile_size def _d256_gfx950_cohort(problem: "UnifiedAttentionProblem") -> bool: @@ -1471,6 +1483,13 @@ def _tiled_cache_key(problem: UnifiedAttentionProblem) -> Tuple: ``test_attention_tiled_cache_key.py`` sweeps the selector thresholds and fails wherever key equality and built-spec equality disagree. """ + # Resolved once and reused by the predicates below that would otherwise walk + # the same selectors again. Those predicates size against the generic 2D + # warp count, so they only take it when that is the count keyed here. + fp16_flash = _enable_gfx942_fp16_flash(problem) + num_warps_2d = None if fp16_flash else _select_2d_num_warps(problem) + tile_size = _select_2d_tile_size(problem) + block_m_per_warp = _select_2d_block_m_per_warp(problem) return ( "tiled", _resolve_attention_arch(), @@ -1485,15 +1504,11 @@ def _tiled_cache_key(problem: UnifiedAttentionProblem) -> Tuple: bool(problem.softcap > 0), bool(problem.use_alibi), bool(problem.use_qq_bias), - ( - _select_gfx942_flash_num_warps(problem) - if _enable_gfx942_fp16_flash(problem) - else _select_2d_num_warps(problem) - ), + (_select_gfx942_flash_num_warps(problem) if fp16_flash else num_warps_2d), _kv_storage_dtype(problem), - _select_2d_tile_size(problem), + tile_size, _select_2d_waves_per_eu(problem), - _select_2d_block_m_per_warp(problem), + block_m_per_warp, # Layer-2 override marker: flips the key exactly when the D256 gfx950 # codegen override applies in `_tiled_spec_from_problem`, so the key stays # faithful to the built kernel (key === built kernel) without duplicating @@ -1526,7 +1541,12 @@ def _tiled_cache_key(problem: UnifiedAttentionProblem) -> Tuple: ( _enable_sched_barrier(problem), _enable_softmax_mfma_interleave(problem), - _enable_k_single_buffer(problem), + _enable_k_single_buffer( + problem, + num_warps=num_warps_2d, + block_m_per_warp=block_m_per_warp, + tile_size=tile_size, + ), ) if _enable_single_batch_combo(problem) else None @@ -1535,7 +1555,7 @@ def _tiled_cache_key(problem: UnifiedAttentionProblem) -> Tuple: # gfx1250 builder never sets this field, so keying it there would only # split identical kernels. ( - _enable_fp8_mfma_qk(problem) + _enable_fp8_mfma_qk(problem, tile_size=tile_size, num_warps=num_warps_2d) if _resolve_attention_arch() != "gfx1250" else None ), @@ -1656,16 +1676,26 @@ def _select_2d_waves_per_eu(problem: UnifiedAttentionProblem) -> Optional[int]: return 2 -def _fp8_qk_loader_fits(problem: UnifiedAttentionProblem) -> bool: +def _fp8_qk_loader_fits( + problem: UnifiedAttentionProblem, + *, + tile_size: int | None = None, + num_warps: int | None = None, +) -> bool: """True iff the async-fp8 K loader can tile the per-iter K bytes. ``raw.ptr.buffer.load.lds`` accepts dwords ∈ {1, 3, 4} = {4, 12, 16} bytes per lane (a hardware quirk -- dwords=2 is rejected). We need one of those per-call payloads to evenly tile ``T*HD`` bytes given - ``num_warps * 64`` lanes per CTA. + ``num_warps * 64`` lanes per CTA. ``tile_size`` / ``num_warps`` take the + already-resolved selector outputs when the caller has them. """ - tile_bytes = _select_2d_tile_size(problem) * problem.head_size - threads = _select_2d_num_warps(problem) * 64 + if tile_size is None: + tile_size = _select_2d_tile_size(problem) + if num_warps is None: + num_warps = _select_2d_num_warps(problem) + tile_bytes = tile_size * problem.head_size + threads = num_warps * 64 for bpl in (16, 12, 4): payload = threads * bpl if tile_bytes >= payload and tile_bytes % payload == 0: @@ -1673,7 +1703,12 @@ def _fp8_qk_loader_fits(problem: UnifiedAttentionProblem) -> bool: return False -def _enable_fp8_mfma_qk(problem: UnifiedAttentionProblem) -> bool: +def _enable_fp8_mfma_qk( + problem: UnifiedAttentionProblem, + *, + tile_size: int | None = None, + num_warps: int | None = None, +) -> bool: """Heuristic: enable the ULP-correct fp8-K-LDS path when it helps. The path is bit-identical to the sync-dequant default. The win @@ -1702,10 +1737,11 @@ def _enable_fp8_mfma_qk(problem: UnifiedAttentionProblem) -> bool: # because the cache key evaluates this on every fp8 launch. if problem.max_seqlen_q > 256: return False - if not _fp8_qk_loader_fits(problem): + if tile_size is None: + tile_size = _select_2d_tile_size(problem) + if not _fp8_qk_loader_fits(problem, tile_size=tile_size, num_warps=num_warps): return False - T_eff = _select_2d_tile_size(problem) - return problem.sliding_window > 0 or problem.max_seqlen_k <= 16 * T_eff + return problem.sliding_window > 0 or problem.max_seqlen_k <= 16 * tile_size def _enable_single_batch_combo(problem: UnifiedAttentionProblem) -> bool: From 4f75b57128ebff80c2de9b1d393fa29fe18aec2f Mon Sep 17 00:00:00 2001 From: dlangbe Date: Thu, 8 Oct 2026 20:21:45 -0400 Subject: [PATCH 08/16] perf(rocke/attention): check shape before env knobs in single-batch predicates The tiled cache key evaluates the gfx950 single-batch schedule predicates on every launch in that cohort. _enable_softmax_mfma_interleave and _enable_d128_small_tile read their environment knob before the head-size and fp8 checks, and _enable_sched_barrier walked the V-schedule predicates before checking head size, so every d64 single-batch launch paid for predicates that can only fire at head size 128. Check the problem fields first. The result is unchanged for every input; the cache key is identical across the cache-key test sweep and its fp8 / single-batch extensions on gfx942, gfx950 and gfx1250, with the knobs both unset and disabled. --- .../kernels/common/attention_unified.py | 24 +++++++++---------- 1 file changed, 12 insertions(+), 12 deletions(-) diff --git a/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/attention_unified.py b/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/attention_unified.py index 335ccbb5dc49..7d37d0514c1e 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/attention_unified.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/attention_unified.py @@ -661,6 +661,10 @@ def _enable_d128_small_tile(problem: UnifiedAttentionProblem) -> bool: DISABLES the lever, restoring the prior routing (T=64, nw=2/4). Any other value (or unset) -> default-ON for the cohort. """ + # Shape checks first: the cache key asks on every launch, and most problems + # fail them before the env knob or the cohort walk would matter. + if problem.head_size != 128 or problem.use_fp8: + return False env = ( __import__("os") .environ.get("HIPDNN_GFX950_D128_SMALL_TILE", "") @@ -669,11 +673,7 @@ def _enable_d128_small_tile(problem: UnifiedAttentionProblem) -> bool: ) if env in ("0", "false", "no", "off"): return False - return ( - problem.head_size == 128 - and not problem.use_fp8 - and _enable_single_batch_combo(problem) - ) + return _enable_single_batch_combo(problem) def _enable_k_single_buffer( @@ -853,6 +853,9 @@ def _enable_softmax_mfma_interleave(problem: UnifiedAttentionProblem) -> bool: ESCAPE HATCH: ``HIPDNN_GFX950_D128_SOFTMAX_INTERLEAVE=0`` (or off/no/false) force-DISABLES the lever, restoring the prior nw=2 / no-interleave routing. """ + # Shape checks first, as in _enable_d128_small_tile. + if problem.head_size != 128 or problem.use_fp8: + return False env = ( __import__("os") .environ.get("HIPDNN_GFX950_D128_SOFTMAX_INTERLEAVE", "") @@ -861,11 +864,7 @@ def _enable_softmax_mfma_interleave(problem: UnifiedAttentionProblem) -> bool: ) if env in ("0", "false", "no", "off"): return False - return ( - problem.head_size == 128 - and not problem.use_fp8 - and _enable_single_batch_combo(problem) - ) + return _enable_single_batch_combo(problem) @dataclass(frozen=True) @@ -1878,9 +1877,10 @@ def _enable_sched_barrier(problem: UnifiedAttentionProblem) -> bool: short single-batch combo uses num_warps==4 (BLOCK_M=128, ample ILP), where the fence is the over-constrained num_warps>=2 regime, so leave it OFF. """ - if not _enable_v_double_buffer(problem): + # Head size first: it is a field read, the V-schedule walk is not. + if problem.head_size != 128: return False - return problem.head_size == 128 + return _enable_v_double_buffer(problem) def _enable_early_v_schedule(problem: UnifiedAttentionProblem) -> bool: From 08a4a275020fb98864acc1c4b91c6aa6e9e9c87b Mon Sep 17 00:00:00 2001 From: dlangbe Date: Tue, 13 Oct 2026 21:08:38 -0400 Subject: [PATCH 09/16] test(rocke/attention): pin the single-batch tail of the tiled cache key --- .../tests/test_attention_tiled_cache_key.py | 27 ++++++++++++++++++- 1 file changed, 26 insertions(+), 1 deletion(-) diff --git a/dnn-providers/hip-kernel-provider/rocke/library/tests/test_attention_tiled_cache_key.py b/dnn-providers/hip-kernel-provider/rocke/library/tests/test_attention_tiled_cache_key.py index ba878fc15a1b..339fce81db30 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/tests/test_attention_tiled_cache_key.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/tests/test_attention_tiled_cache_key.py @@ -301,7 +301,8 @@ def test_fp8_mfma_qk_kv_threshold_splits_key(pin_arch): def test_single_batch_interleave_splits_key(pin_arch): # The d128 single-batch cohort turns on the softmax/MFMA interleave; the - # same shape at num_seqs=2 does not. + # same shape at num_seqs=2 does not. The search bucket also separates this + # pair, so the tail entry itself is pinned by the env-flip test below. pin_arch("gfx950") base = dict( num_query_heads=8, @@ -317,6 +318,30 @@ def test_single_batch_interleave_splits_key(pin_arch): _assert_split(au._tiled_cache_key, _built_2d, a, b, "use_softmax_mfma_interleave") +def test_env_knob_flip_splits_single_batch_key(pin_arch, monkeypatch): + # Routing env knobs are read on every launch and can change mid-process. + # q=257 under default env and q=1025 with both d128 knobs off resolve the + # same geometry; only the single-batch tail separates their keys. + pin_arch("gfx950") + base = dict( + num_seqs=1, + num_query_heads=8, + num_kv_heads=8, + head_size=128, + block_size=16, + dtype="bf16", + ) + a = P(total_q=257, max_seqlen_q=257, max_seqlen_k=257, **base) + key_a = au._tiled_cache_key(a) + spec_a = _built_2d(a)["spec"] + monkeypatch.setenv("HIPDNN_GFX950_D128_SMALL_TILE", "0") + monkeypatch.setenv("HIPDNN_GFX950_D128_SOFTMAX_INTERLEAVE", "off") + b = P(total_q=1025, max_seqlen_q=1025, max_seqlen_k=1025, **base) + field = "use_softmax_mfma_interleave" + assert spec_a[field] != _built_2d(b)["spec"][field] + assert key_a != au._tiled_cache_key(b) + + # --------------------------------------------------------------------------- # Same selector outputs -> same key and byte-identical IR # --------------------------------------------------------------------------- From 0009508ef757de397c22e4126da4258773bf7103 Mon Sep 17 00:00:00 2001 From: dlangbe Date: Fri, 16 Oct 2026 14:29:18 -0400 Subject: [PATCH 10/16] fix(rocke/attention): pass the scalar block-table stride as a kernarg The scalar fallback compiled the block-table row stride as ceil(max_seqlen_k / block_size) and ignored the caller's table stride. A table with rows wider than the batch needs made every sequence after the first read the wrong blocks. Because the stride was baked into the IR, the cache key also had to carry raw max_seqlen_k, and the kernel name carried total_q, so each batch shape compiled its own binary. - Add a trailing block_table_stride kernarg to the scalar 2D and 3D kernels and index block-table rows with it, in both engines. - Launch the scalar kernel with the table's real row stride. - Drop q from the 2D, 3D and reduce scalar kernel names. - Drop total_q, num_seqs, max_seqlen_q and max_seqlen_k from the scalar _cache_key; none of them reach the IR any more. - Add CPU tests for key sharing, byte-identical IR across shapes at every LLVM flavor, and the launch signature, plus a GPU test with wide and sliced block tables against an fp32 reference. Emitted IR changes for the scalar kernels only (name, new kernarg, row stride multiply). The representative golden has no scalar cases. Co-Authored-By: Claude --- .../kernels/common/attention_unified.py | 47 +++--- ...st_attention_scalar_block_table_numeric.py | 141 ++++++++++++++++++ .../tests/test_attention_scalar_cache_key.py | 129 ++++++++++++++++ ...rocke.instances.common.attention_unified.h | 4 +- ...ances.common.attention_unified_selectors.h | 18 ++- .../rocke/instance_attention_unified.h | 2 +- .../instance_attention_unified_internal.h | 6 +- .../instances/common/attention_unified_2d.cpp | 7 +- .../common/attention_unified_3d_reduce.cpp | 2 + .../common/attention_unified_entry.cpp | 40 ++--- .../common/attention_unified_prologue.cpp | 12 -- .../common/attention_unified_selectors.cpp | 19 +-- 12 files changed, 346 insertions(+), 81 deletions(-) create mode 100644 dnn-providers/hip-kernel-provider/rocke/library/tests/test_attention_scalar_block_table_numeric.py create mode 100644 dnn-providers/hip-kernel-provider/rocke/library/tests/test_attention_scalar_cache_key.py diff --git a/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/attention_unified.py b/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/attention_unified.py index f3eb282048f0..873cdcf1bb14 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/attention_unified.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/attention_unified.py @@ -652,18 +652,19 @@ def _cache_key(problem: UnifiedAttentionProblem, arch: str) -> Tuple: ``arch`` sits at index 1 like every other key family in this module, so ``_assert_cache_key_arch`` can check key-vs-compile agreement uniformly and two arches can never alias to one cached kernel. + + Like the tiled keys, it holds only fields the scalar 2D build reads. The + dynamic shape (``total_q``, ``num_seqs``, ``max_seqlen_q``, ``max_seqlen_k``) + reaches the kernel as kernargs or the launch grid, never as IR, so one binary + serves every batch. ``test_attention_scalar_cache_key.py`` enforces this. """ return ( "scalar", arch, - problem.total_q, - problem.num_seqs, problem.num_query_heads, problem.num_kv_heads, problem.head_size, problem.block_size, - problem.max_seqlen_q, - problem.max_seqlen_k, problem.dtype, problem.q_dtype, problem.sliding_window, @@ -4998,7 +4999,7 @@ def _get_scalar_launcher( kernel_name=kname, signature=_attn_signature( problem.dtype, - include_bt_stride=False, + include_bt_stride=True, kv_dtype=_kv_storage_dtype(problem), ), cache_key=("scalar",) + cache_key, @@ -5396,7 +5397,7 @@ def run_unified_attention_torch( softcap=softcap, sinks=sinks, bt_stride=bt_stride, - include_bt_stride=False, + include_bt_stride=True, k_scale=k_scale, v_scale=v_scale, out_scale=out_scale, @@ -5436,7 +5437,6 @@ def kernel_name(self) -> str: p = self.problem return kernel_name_join( self.name, - f"q{p.total_q}", f"h{p.num_query_heads}", f"kv{p.num_kv_heads}", f"d{p.head_size}", @@ -5495,6 +5495,7 @@ def build_unified_attention_2d( _out_scale = b.param("out_scale", F32) softcap = b.param("softcap", F32) num_seqs = b.param("num_seqs", I32) + bt_stride = b.param("block_table_stride", I32) q_tok = b.block_id_x() q_head = b.block_id_y() @@ -5553,14 +5554,7 @@ def build_unified_attention_2d( with loop as (kpos, (m_val, l_val, acc_val)): block_idx, token_in_block = _magic_div_mod(b, kpos, p.block_size) physical = b.global_load_i32( - block_tables, - b.add( - b.mul( - seq_idx, - b.const_i32((p.max_seqlen_k + p.block_size - 1) // p.block_size), - ), - block_idx, - ), + block_tables, b.add(b.mul(seq_idx, bt_stride), block_idx) ) # Vectorised QK dot product. Both Q[q_tok, q_head, :] and @@ -5672,7 +5666,6 @@ def kernel_name(self) -> str: p = self.problem return kernel_name_join( self.name, - f"q{p.total_q}", f"h{p.num_query_heads}", f"kv{p.num_kv_heads}", f"d{p.head_size}", @@ -5723,6 +5716,7 @@ def build_unified_attention_3d( scale = abi["scale"] _softcap = b.param("softcap", F32) num_seqs = b.param("num_seqs", I32) + bt_stride = b.param("block_table_stride", I32) q_tok = b.block_id_x() q_head = b.block_id_y() @@ -5773,6 +5767,7 @@ def build_unified_attention_3d( query, key, block_tables, + bt_stride, seq_idx, q_tok, q_head, @@ -5787,7 +5782,9 @@ def build_unified_attention_3d( alpha = b.exp2(b.fsub(m_val, new_m)) prob = b.exp2(b.fsub(score, new_m)) new_l = b.fadd(b.fmul(l_val, alpha), prob) - vv = _emit_v_load(b, p, dtype, value, block_tables, seq_idx, kv_head, kpos, dim) + vv = _emit_v_load( + b, p, dtype, value, block_tables, bt_stride, seq_idx, kv_head, kpos, dim + ) new_acc = b.fadd(b.fmul(acc_val, alpha), b.fmul(prob, vv)) b.scf_yield(new_m, new_l, new_acc) @@ -5824,7 +5821,6 @@ def kernel_name(self) -> str: p = self.problem return kernel_name_join( self.name, - f"q{p.total_q}", f"h{p.num_query_heads}", f"d{p.head_size}", f"seg{self.num_segments}", @@ -6063,6 +6059,7 @@ def _emit_qk_score( query: Value, key: Value, block_tables: Value, + bt_stride: Value, seq_idx: Value, q_tok: Value, q_head: Value, @@ -6091,7 +6088,7 @@ def _emit_qk_score( """ score = b.const_f32(0.0) physical, token_in_block = _physical_block_and_token( - b, p, block_tables, seq_idx, kpos + b, p, block_tables, bt_stride, seq_idx, kpos ) q_desc = _q_descriptor(p) kv_desc = _paged_kv_descriptor(p) @@ -6142,13 +6139,14 @@ def _emit_v_load( dtype: Type, value: Value, block_tables: Value, + bt_stride: Value, seq_idx: Value, kv_head: Value, kpos: Value, dim: Value, ) -> Value: physical, token_in_block = _physical_block_and_token( - b, p, block_tables, seq_idx, kpos + b, p, block_tables, bt_stride, seq_idx, kpos ) v_off = _paged_kv_descriptor(p).offset( b, @@ -6164,12 +6162,17 @@ def _physical_block_and_token( b: IRBuilder, p: UnifiedAttentionProblem, block_tables: Value, + bt_stride: Value, seq_idx: Value, kpos: Value, ) -> Tuple[Value, Value]: + """Look up ``kpos``'s physical KV block in row ``seq_idx`` of the block table. + + ``bt_stride`` is the table's row stride kernarg, not a value derived from + ``max_seqlen_k``: callers may allocate rows wider than the longest sequence. + """ block_idx, token_in_block = _magic_div_mod(b, kpos, p.block_size) - max_blocks = (p.max_seqlen_k + p.block_size - 1) // p.block_size physical = b.global_load_i32( - block_tables, b.add(b.mul(seq_idx, b.const_i32(max_blocks)), block_idx) + block_tables, b.add(b.mul(seq_idx, bt_stride), block_idx) ) return physical, token_in_block diff --git a/dnn-providers/hip-kernel-provider/rocke/library/tests/test_attention_scalar_block_table_numeric.py b/dnn-providers/hip-kernel-provider/rocke/library/tests/test_attention_scalar_block_table_numeric.py new file mode 100644 index 000000000000..6db01be49d0a --- /dev/null +++ b/dnn-providers/hip-kernel-provider/rocke/library/tests/test_attention_scalar_block_table_numeric.py @@ -0,0 +1,141 @@ +# Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +# SPDX-License-Identifier: MIT + +"""Scalar fallback launches with block-table rows wider than the batch needs. + +Callers usually size block-table rows for the longest sequence they admit, not +the longest one in the current batch. The scalar kernels index rows by the +table's real row stride (the ``block_table_stride`` kernarg), so every sequence +after the first must still find its own blocks. Unused table slots point at a +poison block, so a wrong row stride corrupts the output instead of passing by +luck. + +Reference attention is computed independently in FP32 from the logical K/V +after rounding inputs to bf16. Requires a ROCm GPU; skips otherwise. +""" + +import itertools +import math + +import pytest + +from kernels.common import attention_unified as au + +_HQ, _HKV, _D, _BS = 8, 2, 64, 16 +_ATOL = 1e-2 + + +@pytest.fixture(scope="module") +def torch(): + torch = pytest.importorskip("torch") + if not torch.cuda.is_available(): + pytest.skip("requires a ROCm GPU") + return torch + + +def _case(torch, q_lens, kv_lens, row_width, *, view_of_wider=False, seed=0): + """Build a paged case whose block-table rows hold ``row_width`` slots.""" + gen = torch.Generator().manual_seed(seed) + needed = [math.ceil(n / _BS) for n in kv_lens] + num_blocks = sum(needed) + 1 + poison = num_blocks - 1 + order = torch.randperm(num_blocks - 1, generator=gen).tolist() + alloc_width = row_width + 5 if view_of_wider else row_width + table = torch.full((len(kv_lens), alloc_width), poison, dtype=torch.int32) + for s, n in enumerate(needed): + table[s, :n] = torch.tensor(order[:n], dtype=torch.int32) + order = order[n:] + + def randn(*shape): + return torch.randn(*shape, generator=gen).to(torch.bfloat16) + + q = randn(sum(q_lens), _HQ, _D) + k = randn(num_blocks, _BS, _HKV, _D) + v = randn(num_blocks, _BS, _HKV, _D) + k[poison] = 64.0 + v[poison] = 64.0 + cu = torch.tensor([0, *itertools.accumulate(q_lens)], dtype=torch.int32) + problem = au.UnifiedAttentionProblem( + total_q=sum(q_lens), + num_seqs=len(q_lens), + num_query_heads=_HQ, + num_kv_heads=_HKV, + head_size=_D, + block_size=_BS, + max_seqlen_q=max(q_lens), + max_seqlen_k=max(kv_lens), + dtype="bf16", + ) + dev = table.cuda() + if view_of_wider: + dev = dev[:, :row_width] + return problem, q, k, v, cu, table, dev + + +def _reference(torch, q, k, v, cu, table, q_lens, kv_lens): + out = torch.zeros(q.shape, dtype=torch.float32) + ratio = _HQ // _HKV + for s, (q_len, kv_len) in enumerate(zip(q_lens, kv_lens)): + blocks = table[s, : math.ceil(kv_len / _BS)].long() + keys = k[blocks].reshape(-1, _HKV, _D)[:kv_len].float() + values = v[blocks].reshape(-1, _HKV, _D)[:kv_len].float() + context = kv_len - q_len + for i in range(q_len): + tok = int(cu[s]) + i + visible = context + i + 1 + for h in range(_HQ): + scores = keys[:visible, h // ratio] @ q[tok, h].float() + probs = (scores / math.sqrt(_D)).softmax(0) + out[tok, h] = probs @ values[:visible, h // ratio] + return out + + +def _launch(torch, problem, q, k, v, cu, kv_lens, block_table): + out = torch.empty(q.shape, dtype=q.dtype, device="cuda") + au.run_unified_attention_torch( + problem=problem, + q=q.cuda(), + k=k.cuda(), + v=v.cuda(), + out=out, + cu_seqlens_q=cu.cuda(), + seqused_k=torch.tensor(kv_lens, dtype=torch.int32, device="cuda"), + softmax_scale=1.0 / math.sqrt(_D), + block_table=block_table, + softcap=0.0, + backend="scalar", + ) + torch.cuda.synchronize() + return out.float().cpu() + + +@pytest.mark.parametrize("view_of_wider", (False, True), ids=("wide", "view")) +def test_wide_block_table_rows(torch, view_of_wider): + q_lens, kv_lens = [5, 1, 9], [37, 20, 61] + row_width = math.ceil(max(kv_lens) / _BS) + 6 + problem, q, k, v, cu, table, dev = _case( + torch, q_lens, kv_lens, row_width, view_of_wider=view_of_wider + ) + assert dev.stride(0) > math.ceil(problem.max_seqlen_k / _BS) + got = _launch(torch, problem, q, k, v, cu, kv_lens, dev) + want = _reference(torch, q, k, v, cu, table, q_lens, kv_lens) + max_abs = (got - want).abs().max().item() + assert max_abs <= _ATOL, f"max_abs {max_abs:.3e} vs fp32 reference" + + +def test_one_binary_serves_different_batches(torch): + """Shapes that differ only in batch geometry share a key; the second + launch must reuse the first binary and still match the reference.""" + first = ([4, 2], [30, 17], 4) + second = ([1, 7, 3, 2], [90, 64, 15, 33], 12) + launchers = None + for i, (q_lens, kv_lens, width) in enumerate((first, second)): + problem, q, k, v, cu, table, dev = _case(torch, q_lens, kv_lens, width, seed=i) + got = _launch(torch, problem, q, k, v, cu, kv_lens, dev) + if launchers is None: + launchers = len(au._SCALAR_LAUNCHERS) + else: + assert len(au._SCALAR_LAUNCHERS) == launchers + want = _reference(torch, q, k, v, cu, table, q_lens, kv_lens) + max_abs = (got - want).abs().max().item() + assert max_abs <= _ATOL, f"case {i}: max_abs {max_abs:.3e}" diff --git a/dnn-providers/hip-kernel-provider/rocke/library/tests/test_attention_scalar_cache_key.py b/dnn-providers/hip-kernel-provider/rocke/library/tests/test_attention_scalar_cache_key.py new file mode 100644 index 000000000000..7591be089ed7 --- /dev/null +++ b/dnn-providers/hip-kernel-provider/rocke/library/tests/test_attention_scalar_cache_key.py @@ -0,0 +1,129 @@ +# Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +# SPDX-License-Identifier: MIT + +"""Cache-key contract for the scalar unified-attention fallback. + +The scalar 2D / 3D / reduce kernels read the dynamic shape (``total_q``, +``num_seqs``, ``max_seqlen_q``, ``max_seqlen_k``) only through kernargs and the +launch grid. The block-table row stride is the ``block_table_stride`` kernarg, +so a caller's table may be wider than ``ceil(max_seqlen_k / block_size)``. +``_cache_key`` therefore omits the dynamic shape, and one binary serves every +batch. + +Layers: + * Shapes that differ only in the dynamic shape share a key and lower to + byte-identical IR at every LLVM flavor. + * Every static field the build reads splits the key. + * The launch signature matches the kernel's params, including the stride. + +No GPU required. ``test_attention_scalar_block_table_numeric.py`` runs the +launch path on hardware. +""" + +from __future__ import annotations + +import dataclasses +import re + +import pytest + +from kernels.common import attention_unified as au +from rocke.core.lower_llvm import LLVM_FLAVORS, _lower_kernel_to_llvm_python + +P = au.UnifiedAttentionProblem + +_BASE = P( + total_q=64, + num_seqs=4, + num_query_heads=8, + num_kv_heads=2, + head_size=128, + block_size=16, + max_seqlen_q=16, + max_seqlen_k=256, + dtype="bf16", +) + +# Each entry moves only the dynamic shape. max_seqlen_k=1000 changes the block +# count, which used to be compiled into the block-table row stride. +_DYNAMIC_VARIANTS = ( + dict(total_q=200), + dict(num_seqs=1, total_q=16), + dict(num_seqs=37, total_q=37, max_seqlen_q=1), + dict(max_seqlen_q=50), + dict(max_seqlen_k=250), + dict(max_seqlen_k=1000), + dict(total_q=1, num_seqs=1, max_seqlen_q=1, max_seqlen_k=65536), +) + +# Each entry moves one static field the scalar build reads. +_STATIC_VARIANTS = ( + dict(num_query_heads=16), + dict(num_kv_heads=8), + dict(head_size=64), + dict(block_size=32), + dict(dtype="fp16"), + dict(sliding_window=128), + dict(use_sinks=True), + dict(softcap=30.0), +) + + +def _lower_all(problem: P, flavor: str) -> tuple[str, str, str]: + kernels = ( + au.build_unified_attention_2d(au.UnifiedAttention2DSpec(problem)), + au.build_unified_attention_3d(au.UnifiedAttention3DSpec(problem)), + au.build_unified_attention_reduce(au.UnifiedAttentionReduceSpec(problem, 8)), + ) + return tuple( + _lower_kernel_to_llvm_python(k, arch="gfx950", llvm_flavor=flavor) + for k in kernels + ) + + +@pytest.mark.parametrize("change", _DYNAMIC_VARIANTS, ids=repr) +def test_dynamic_shape_shares_key(change): + other = dataclasses.replace(_BASE, **change) + assert au._cache_key(other, "gfx950") == au._cache_key(_BASE, "gfx950") + + +@pytest.mark.parametrize("flavor", LLVM_FLAVORS) +def test_dynamic_shape_lowers_identical_ir(flavor): + """Same key means same binary: every variant must lower byte-identically + for the 2D, 3D and reduce scalar kernels, kernel name included.""" + base = _lower_all(_BASE, flavor) + for change in _DYNAMIC_VARIANTS: + other = _lower_all(dataclasses.replace(_BASE, **change), flavor) + for kind, a, b in zip(("2d", "3d", "reduce"), base, other): + assert a == b, f"{kind} IR differs for {change} at {flavor}" + + +@pytest.mark.parametrize("change", _STATIC_VARIANTS, ids=repr) +def test_static_field_splits_key_and_ir(change): + other = dataclasses.replace(_BASE, **change) + assert au._cache_key(other, "gfx950") != au._cache_key(_BASE, "gfx950") + a = _lower_all(_BASE, "llvm20")[0] + b = _lower_all(other, "llvm20")[0] + assert a != b, f"{change} splits the key but not the 2D IR" + + +def test_arch_splits_key(): + assert au._cache_key(_BASE, "gfx942") != au._cache_key(_BASE, "gfx950") + + +def test_block_table_stride_is_a_kernarg(): + """The row stride must come from the launch, not ``max_seqlen_k``.""" + ll = _lower_all(_BASE, "llvm20") + for kind, text in zip(("2d", "3d"), ll[:2]): + assert "i32 %block_table_stride)" in text, kind + assert re.search(r"mul nsw i32 %\w+, %block_table_stride$", text, re.M), kind + + +def test_launch_signature_matches_kernel_params(): + """``_get_scalar_launcher`` packs kernargs by this signature; a name or + order mismatch with the kernel's params silently shifts every argument.""" + kernel = au.build_unified_attention_2d(au.UnifiedAttention2DSpec(_BASE)) + signature = au._attn_signature( + _BASE.dtype, include_bt_stride=True, kv_dtype=au._kv_storage_dtype(_BASE) + ) + assert [item["name"] for item in signature] == [p.name for p in kernel.params] diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/helper_helper_rocke.instances.common.attention_unified.h b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/helper_helper_rocke.instances.common.attention_unified.h index b795a85c52e0..d9bcb2a32967 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/helper_helper_rocke.instances.common.attention_unified.h +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/helper_helper_rocke.instances.common.attention_unified.h @@ -13,7 +13,7 @@ * build_unified_attention_2d rocke_build_unified_attention_2d_scalar * UnifiedAttention2DSpec.kernel_name rocke_unified_attention_2d_scalar_kernel_name * (block_id grid of build_..._2d) rocke_unified_attention_2d_scalar_grid - * _attn_signature(include_bt_stride=False) rocke_unified_attention_2d_scalar_signature + * _attn_signature(include_bt_stride=True) rocke_unified_attention_2d_scalar_signature * build_unified_attention_3d rocke_build_unified_attention_3d_scalar * UnifiedAttention3DSpec.kernel_name rocke_unified_attention_3d_scalar_kernel_name * (block_id grid of build_..._3d) rocke_unified_attention_3d_scalar_grid @@ -123,7 +123,7 @@ rocke_kernel_def_t* rocke_build_unified_attention_2d_scalar( * extents = (total_q, num_query_heads, head_size). Writes out[0..2]. */ void rocke_unified_attention_2d_scalar_grid(const rocke_unified_attention_problem_t* p, int out[3]); -/* _attn_signature(dtype, include_bt_stride=False): the 2D scalar kernel ABI. +/* _attn_signature(dtype, include_bt_stride=True): the 2D scalar kernel ABI. * On ROCKE_OK *out_items / *out_count hold the arena-owned array. */ rocke_status_t rocke_unified_attention_2d_scalar_signature(const rocke_unified_attention_problem_t* p, diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/helper_helper_rocke.instances.common.attention_unified_selectors.h b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/helper_helper_rocke.instances.common.attention_unified_selectors.h index d59fd5bb26e8..39843a3938fb 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/helper_helper_rocke.instances.common.attention_unified_selectors.h +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/helper_helper_rocke.instances.common.attention_unified_selectors.h @@ -267,22 +267,24 @@ bool rocke_unified_attn_segm_descriptors(rocke_ir_builder_t* b, /* ------------------------------------------------------- IR emit helpers */ -/* Python: _physical_block_and_token(b, p, block_tables, seq_idx, kpos) +/* Python: _physical_block_and_token(b, p, block_tables, bt_stride, seq_idx, kpos) * -> (physical, token_in_block). * (block_idx, token_in_block) = magic_div_mod(kpos, block_size); - * max_blocks = ceil(max_seqlen_k / block_size); - * physical = global_load_i32(block_tables, seq_idx*max_blocks + block_idx). + * physical = global_load_i32(block_tables, seq_idx*bt_stride + block_idx). + * bt_stride is the block_table_stride kernarg (the table's row stride). * Writes *out_physical and *out_token_in_block; returns true on success. */ bool rocke_unified_attn_physical_block_and_token(rocke_ir_builder_t* b, const rocke_unified_attn_problem_t* p, rocke_value_t* block_tables, + rocke_value_t* bt_stride, rocke_value_t* seq_idx, rocke_value_t* kpos, rocke_value_t** out_physical, rocke_value_t** out_token_in_block); -/* Python: _emit_qk_score(b, p, dtype, query, key, block_tables, seq_idx, q_tok, - * q_head, kv_head, kpos, scale, rcp_ln2) -> Value. +/* Python: _emit_qk_score(b, p, dtype, query, key, block_tables, bt_stride, + * seq_idx, q_tok, q_head, kv_head, kpos, scale, + * rcp_ln2) -> Value. * Per-(query_token, query_head) QK dot product for the scalar kernels (vec8 main * fold + scalar tail), scaled by scale*rcp_ln2. `dtype` is the IR element type * (the Python passes a Type). */ @@ -292,6 +294,7 @@ rocke_value_t* rocke_unified_attn_emit_qk_score(rocke_ir_builder_t* b, rocke_value_t* query, rocke_value_t* key, rocke_value_t* block_tables, + rocke_value_t* bt_stride, rocke_value_t* seq_idx, rocke_value_t* q_tok, rocke_value_t* q_head, @@ -300,14 +303,15 @@ rocke_value_t* rocke_unified_attn_emit_qk_score(rocke_ir_builder_t* b, rocke_value_t* scale, rocke_value_t* rcp_ln2); -/* Python: _emit_v_load(b, p, dtype, value, block_tables, seq_idx, kv_head, kpos, - * dim) -> Value. +/* Python: _emit_v_load(b, p, dtype, value, block_tables, bt_stride, seq_idx, + * kv_head, kpos, dim) -> Value. * Loads V[physical, token_in_block, kv_head, dim] and casts to f32. */ rocke_value_t* rocke_unified_attn_emit_v_load(rocke_ir_builder_t* b, const rocke_unified_attn_problem_t* p, const rocke_type_t* dtype, rocke_value_t* value, rocke_value_t* block_tables, + rocke_value_t* bt_stride, rocke_value_t* seq_idx, rocke_value_t* kv_head, rocke_value_t* kpos, diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/instance_attention_unified.h b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/instance_attention_unified.h index 1a3dae8dfaa1..2bfb844f1770 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/instance_attention_unified.h +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/instance_attention_unified.h @@ -36,7 +36,7 @@ * build_unified_attention_reduce(spec) rocke_build_unified_attention_reduce_scalar * *.kernel_name() rocke_unified_attention_*_kernel_name * (block_id grid extents) rocke_unified_attention_*_grid - * _attn_signature(include_bt_stride=False) rocke_unified_attention_2d_scalar_signature + * _attn_signature(include_bt_stride=True) rocke_unified_attention_2d_scalar_signature * supports_native_unified_attention(problem) rocke_unified_attention_supports_scalar * (build -> lower .ll convenience, NEW) rocke_unified_attention_*_lower_to_llvm * diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/instance_attention_unified_internal.h b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/instance_attention_unified_internal.h index 67d0128605c5..5e3605c426f0 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/instance_attention_unified_internal.h +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/instance_attention_unified_internal.h @@ -108,7 +108,6 @@ typedef struct rocke_attn_unified_build_ctx const rocke_type_t* dtype; /* spec.dtype_ir (F16 / BF16) */ int num_segments; /* 3D/reduce; 0 for 2D */ int num_queries_per_kv; /* p->num_query_heads / kv */ - int max_blocks; /* ceil(max_seqlen_k/block_size) */ /* ---------- kernel result + extra outputs ---------- */ rocke_kernel_def_t* kernel; /* == b->kernel; returned by the driver */ @@ -139,11 +138,12 @@ typedef struct rocke_attn_unified_build_ctx rocke_value_t* abi_v_scale; /* "v_scale" (F32) */ /* ---------- kernel-specific tail params (after the ABI prefix) ---------- - * 2D appends out_scale, softcap, num_seqs. 3D appends softcap (unused) + - * num_seqs. reduce appends none. */ + * 2D appends out_scale, softcap, num_seqs, block_table_stride. 3D appends + * softcap (unused), num_seqs, block_table_stride. reduce appends none. */ rocke_value_t* out_scale; /* 2D "out_scale" (F32; unused by body) */ rocke_value_t* softcap; /* 2D/3D "softcap" (F32; 3D unused) */ rocke_value_t* num_seqs; /* 2D/3D "num_seqs" (I32) */ + rocke_value_t* block_table_stride; /* 2D/3D "block_table_stride" (I32) */ /* ---------- grid ids + thread predicate (prologue) ---------- */ rocke_value_t* q_tok; /* block_id_x() */ diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/common/attention_unified_2d.cpp b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/common/attention_unified_2d.cpp index 4d5ac35da835..66cf4b20e881 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/common/attention_unified_2d.cpp +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/common/attention_unified_2d.cpp @@ -120,14 +120,11 @@ void rocke_attn_unified_emit_2d_softmax_loop(rocke_attn_unified_build_ctx_t* ctx rocke_value_t* token_in_block = NULL; (void)rocke_unified_attn_magic_div_mod(b, kpos, p->block_size, &block_idx, &token_in_block); - /* physical = global_load_i32(block_tables, - * seq_idx * ceil(max_seqlen_k/block_size) + block_idx) */ - int max_blocks_2d = (p->max_seqlen_k + p->block_size - 1) / p->block_size; + /* physical = global_load_i32(block_tables, seq_idx * bt_stride + block_idx) */ rocke_value_t* physical = rocke_b_global_load_i32( b, ctx->abi_block_tables, - rocke_b_add( - b, rocke_b_mul(b, ctx->seq_idx, rocke_b_const_i32(b, max_blocks_2d)), block_idx), + rocke_b_add(b, rocke_b_mul(b, ctx->seq_idx, ctx->block_table_stride), block_idx), /*align=*/4); /* Vectorised QK dot product (vec8 main fold + scalar tail). */ diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/common/attention_unified_3d_reduce.cpp b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/common/attention_unified_3d_reduce.cpp index df051acff079..569c16653ef9 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/common/attention_unified_3d_reduce.cpp +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/common/attention_unified_3d_reduce.cpp @@ -152,6 +152,7 @@ void rocke_attn_unified_emit_3d_segment_loop(rocke_attn_unified_build_ctx_t* ctx ctx->abi_query, ctx->abi_key, ctx->abi_block_tables, + ctx->block_table_stride, ctx->seq_idx, ctx->q_tok, ctx->q_head, @@ -174,6 +175,7 @@ void rocke_attn_unified_emit_3d_segment_loop(rocke_attn_unified_build_ctx_t* ctx ctx->dtype, ctx->abi_value, ctx->abi_block_tables, + ctx->block_table_stride, ctx->seq_idx, ctx->kv_head, kpos, diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/common/attention_unified_entry.cpp b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/common/attention_unified_entry.cpp index f0a081d2bf56..eee5fa81bdea 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/common/attention_unified_entry.cpp +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/common/attention_unified_entry.cpp @@ -17,7 +17,7 @@ * - the three launch grid fns * rocke_unified_attention_{2d,3d,reduce}_scalar_grid * - the 2D signature - * rocke_unified_attention_2d_scalar_signature (_attn_signature, no bt_stride) + * rocke_unified_attention_2d_scalar_signature (_attn_signature, with bt_stride) * - the three build drivers, each ctx_init -> phase functions in Python * execution order (see instance_attention_unified_internal.h) * rocke_build_unified_attention_{2d,3d,reduce}_scalar @@ -188,25 +188,24 @@ bool rocke_unified_attention_supports_scalar(const rocke_unified_attention_probl rocke_status_t rocke_unified_attention_2d_scalar_kernel_name( const rocke_unified_attention_problem_t* p, const char* name, char* out, size_t out_cap) { - char qbuf[32], hbuf[32], kvbuf[32], dbuf[32], bbuf[32]; + char hbuf[32], kvbuf[32], dbuf[32], bbuf[32]; const char* prefix; if(p == NULL || out == NULL) return ROCKE_ERR_VALUE; prefix = name ? name : "rocke_unified_attention_2d_scalar"; - snprintf(qbuf, sizeof(qbuf), "q%d", p->total_q); snprintf(hbuf, sizeof(hbuf), "h%d", p->num_query_heads); snprintf(kvbuf, sizeof(kvbuf), "kv%d", p->num_kv_heads); snprintf(dbuf, sizeof(dbuf), "d%d", p->head_size); snprintf(bbuf, sizeof(bbuf), "b%d", p->block_size); { - const char* parts[6] = {qbuf, hbuf, kvbuf, dbuf, bbuf, p->dtype ? p->dtype : ""}; + const char* parts[5] = {hbuf, kvbuf, dbuf, bbuf, p->dtype ? p->dtype : ""}; /* flags={"sink":use_sinks,"sw":sliding_window>0,"softcap":softcap>0} */ const char* flag_names[3] = {"sink", "sw", "softcap"}; int flag_on[3] = {p->use_sinks ? 1 : 0, p->sliding_window > 0 ? 1 : 0, p->softcap > 0.0 ? 1 : 0}; - return rocke_kernel_name_join(prefix, parts, 6, flag_names, flag_on, 3, out, out_cap, NULL); + return rocke_kernel_name_join(prefix, parts, 5, flag_names, flag_on, 3, out, out_cap, NULL); } } @@ -217,13 +216,12 @@ rocke_status_t char* out, size_t out_cap) { - char qbuf[32], hbuf[32], kvbuf[32], dbuf[32], bbuf[32], segbuf[32]; + char hbuf[32], kvbuf[32], dbuf[32], bbuf[32], segbuf[32]; const char* prefix; if(p == NULL || out == NULL) return ROCKE_ERR_VALUE; prefix = name ? name : "rocke_unified_attention_3d_scalar"; - snprintf(qbuf, sizeof(qbuf), "q%d", p->total_q); snprintf(hbuf, sizeof(hbuf), "h%d", p->num_query_heads); snprintf(kvbuf, sizeof(kvbuf), "kv%d", p->num_kv_heads); snprintf(dbuf, sizeof(dbuf), "d%d", p->head_size); @@ -231,8 +229,8 @@ rocke_status_t snprintf(segbuf, sizeof(segbuf), "seg%d", num_segments); { - const char* parts[7] = {qbuf, hbuf, kvbuf, dbuf, bbuf, segbuf, p->dtype ? p->dtype : ""}; - return rocke_kernel_name_join(prefix, parts, 7, NULL, NULL, 0, out, out_cap, NULL); + const char* parts[6] = {hbuf, kvbuf, dbuf, bbuf, segbuf, p->dtype ? p->dtype : ""}; + return rocke_kernel_name_join(prefix, parts, 6, NULL, NULL, 0, out, out_cap, NULL); } } @@ -243,20 +241,19 @@ rocke_status_t char* out, size_t out_cap) { - char qbuf[32], hbuf[32], dbuf[32], segbuf[32]; + char hbuf[32], dbuf[32], segbuf[32]; const char* prefix; if(p == NULL || out == NULL) return ROCKE_ERR_VALUE; prefix = name ? name : "rocke_unified_attention_reduce_scalar"; - snprintf(qbuf, sizeof(qbuf), "q%d", p->total_q); snprintf(hbuf, sizeof(hbuf), "h%d", p->num_query_heads); snprintf(dbuf, sizeof(dbuf), "d%d", p->head_size); snprintf(segbuf, sizeof(segbuf), "seg%d", num_segments); { - const char* parts[5] = {qbuf, hbuf, dbuf, segbuf, p->dtype ? p->dtype : ""}; - return rocke_kernel_name_join(prefix, parts, 5, NULL, NULL, 0, out, out_cap, NULL); + const char* parts[4] = {hbuf, dbuf, segbuf, p->dtype ? p->dtype : ""}; + return rocke_kernel_name_join(prefix, parts, 4, NULL, NULL, 0, out, out_cap, NULL); } } @@ -295,7 +292,7 @@ void rocke_unified_attention_reduce_scalar_grid(const rocke_unified_attention_pr } /* ===================================================================== * - * 2D signature -- _attn_signature(dtype, include_bt_stride=False) + * 2D signature -- _attn_signature(dtype, include_bt_stride=True) * ===================================================================== */ rocke_status_t @@ -331,7 +328,8 @@ rocke_status_t rocke_signature_builder_scalar(&sb, "out_scale", "f32"); rocke_signature_builder_scalar(&sb, "softcap", "f32"); rocke_signature_builder_scalar(&sb, "num_seqs", "i32"); - /* include_bt_stride=False, include_qq_bias_stride=False => no tail scalars */ + rocke_signature_builder_scalar(&sb, "block_table_stride", "i32"); + /* include_qq_bias_stride=False => no qq_bias_stride_0 tail scalar */ return rocke_signature_builder_build(&sb, out_items, out_count); } @@ -341,9 +339,9 @@ rocke_status_t * * ctx_init -> declare_scalar_params (+ 2D tail params) -> emit_find_seq_idx -> * emit_prologue -> emit_2d_softmax_loop -> emit_2d_epilogue -> ctx.kernel. - * The 2D tail params (out_scale, softcap, num_seqs) follow the shared ABI - * prefix; declare_scalar_params declares only the prefix, so the tail is - * appended here -- matching the Python build body's param order. + * The 2D tail params (out_scale, softcap, num_seqs, block_table_stride) follow + * the shared ABI prefix; declare_scalar_params declares only the prefix, so + * the tail is appended here, matching the Python build body's param order. * ===================================================================== */ rocke_kernel_def_t* rocke_build_unified_attention_2d_scalar( rocke_ir_builder_t* b, const rocke_unified_attention_problem_t* p, const char* name) @@ -382,10 +380,11 @@ rocke_kernel_def_t* rocke_build_unified_attention_2d_scalar( /* abi = _declare_scalar_attn_params(b, dtype) */ rocke_attn_unified_declare_scalar_params(&ctx); - /* 2D tail params: out_scale, softcap, num_seqs. */ + /* 2D tail params: out_scale, softcap, num_seqs, block_table_stride. */ ctx.out_scale = rocke_b_param(b, "out_scale", rocke_f32(), NULL); ctx.softcap = rocke_b_param(b, "softcap", rocke_f32(), NULL); ctx.num_seqs = rocke_b_param(b, "num_seqs", rocke_i32(), NULL); + ctx.block_table_stride = rocke_b_param(b, "block_table_stride", rocke_i32(), NULL); /* prologue: grid ids + seq-idx scan + per-seq geometry + SSA constants. * The seq-idx scan reads q_tok, which emit_prologue computes first; the @@ -450,9 +449,10 @@ rocke_kernel_def_t* /* abi = _declare_scalar_attn_params(b, dtype) */ rocke_attn_unified_declare_scalar_params(&ctx); - /* 3D tail params: softcap (unused by body) + num_seqs. */ + /* 3D tail params: softcap (unused by body), num_seqs, block_table_stride. */ ctx.softcap = rocke_b_param(b, "softcap", rocke_f32(), NULL); ctx.num_seqs = rocke_b_param(b, "num_seqs", rocke_i32(), NULL); + ctx.block_table_stride = rocke_b_param(b, "block_table_stride", rocke_i32(), NULL); /* emit_prologue runs the seq-idx scan internally after computing q_tok; * do not call emit_find_seq_idx here (q_tok would still be NULL). */ diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/common/attention_unified_prologue.cpp b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/common/attention_unified_prologue.cpp index cabf3e734094..6bfdefac03de 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/common/attention_unified_prologue.cpp +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/common/attention_unified_prologue.cpp @@ -122,18 +122,6 @@ bool rocke_attn_unified_ctx_init(rocke_attn_unified_build_ctx_t* ctx, return false; } - /* max_blocks = ceil(max_seqlen_k / block_size). Python folds this inline as - * (p.max_seqlen_k + p.block_size - 1) // p.block_size when computing the - * block_tables row stride. */ - if(p->block_size > 0) - { - ctx->max_blocks = (p->max_seqlen_k + p->block_size - 1) / p->block_size; - } - else - { - ctx->max_blocks = 0; - } - return true; } diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/common/attention_unified_selectors.cpp b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/common/attention_unified_selectors.cpp index 838cf19fdd52..a57d41a47754 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/common/attention_unified_selectors.cpp +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/common/attention_unified_selectors.cpp @@ -1313,6 +1313,7 @@ bool rocke_unified_attn_segm_descriptors(rocke_ir_builder_t* b, bool rocke_unified_attn_physical_block_and_token(rocke_ir_builder_t* b, const rocke_unified_attn_problem_t* p, rocke_value_t* block_tables, + rocke_value_t* bt_stride, rocke_value_t* seq_idx, rocke_value_t* kpos, rocke_value_t** out_physical, @@ -1320,19 +1321,17 @@ bool rocke_unified_attn_physical_block_and_token(rocke_ir_builder_t* b, { rocke_value_t* block_idx = NULL; rocke_value_t* token_in_block = NULL; - int max_blocks; rocke_value_t* physical; if(!rocke_unified_attn_magic_div_mod(b, kpos, p->block_size, &block_idx, &token_in_block)) { return false; } - max_blocks = (p->max_seqlen_k + p->block_size - 1) / p->block_size; - physical = rocke_b_global_load_i32( - b, - block_tables, - rocke_b_add(b, rocke_b_mul(b, seq_idx, rocke_b_const_i32(b, max_blocks)), block_idx), - 0 /* align default */); + physical + = rocke_b_global_load_i32(b, + block_tables, + rocke_b_add(b, rocke_b_mul(b, seq_idx, bt_stride), block_idx), + 0 /* align default */); if(out_physical != NULL) { @@ -1351,6 +1350,7 @@ rocke_value_t* rocke_unified_attn_emit_qk_score(rocke_ir_builder_t* b, rocke_value_t* query, rocke_value_t* key, rocke_value_t* block_tables, + rocke_value_t* bt_stride, rocke_value_t* seq_idx, rocke_value_t* q_tok, rocke_value_t* q_head, @@ -1372,7 +1372,7 @@ rocke_value_t* rocke_unified_attn_emit_qk_score(rocke_ir_builder_t* b, int d; if(!rocke_unified_attn_physical_block_and_token( - b, p, block_tables, seq_idx, kpos, &physical, &token_in_block)) + b, p, block_tables, bt_stride, seq_idx, kpos, &physical, &token_in_block)) { return NULL; } @@ -1451,6 +1451,7 @@ rocke_value_t* rocke_unified_attn_emit_v_load(rocke_ir_builder_t* b, const rocke_type_t* dtype, rocke_value_t* value, rocke_value_t* block_tables, + rocke_value_t* bt_stride, rocke_value_t* seq_idx, rocke_value_t* kv_head, rocke_value_t* kpos, @@ -1462,7 +1463,7 @@ rocke_value_t* rocke_unified_attn_emit_v_load(rocke_ir_builder_t* b, rocke_value_t* v_off; if(!rocke_unified_attn_physical_block_and_token( - b, p, block_tables, seq_idx, kpos, &physical, &token_in_block)) + b, p, block_tables, bt_stride, seq_idx, kpos, &physical, &token_in_block)) { return NULL; } From 6b4a95d2bb6c9afacd76df21acad39da69dd8cee Mon Sep 17 00:00:00 2001 From: dlangbe Date: Fri, 16 Oct 2026 14:30:12 -0400 Subject: [PATCH 11/16] test(rocke): add a scalar family to the cache-key A/B harness The harness compared only the tiled 2D/3D launcher keys. The scalar fallback key change needs the same evidence: per-problem IR identity, compiled binary count, and A/B latency with max_abs against fp32. - Add --family {tiled,scalar}; scalar routes problems through the scalar gate, keys them with _cache_key, lowers the 2D scalar kernel the launcher compiles, and launches with backend="scalar". - Use smaller scalar sweeps, since every problem is lowered and the kernel is slow. - In the scalar family, report but do not fail on own-IR differences between arms: a key change that also changes the kernel ABI makes them differ by design. Co-Authored-By: Claude --- .../benchmarks/common/tiled_cache_key_ab.py | 174 ++++++++++++++---- 1 file changed, 136 insertions(+), 38 deletions(-) diff --git a/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/common/tiled_cache_key_ab.py b/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/common/tiled_cache_key_ab.py index 15db625a4873..c7fe1e9390ef 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/common/tiled_cache_key_ab.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/common/tiled_cache_key_ab.py @@ -1,9 +1,10 @@ # Copyright (c) Advanced Micro Devices, Inc., or its affiliates. # SPDX-License-Identifier: MIT -"""A/B harness for the tiled attention cache keys. +"""A/B harness for the attention launcher cache keys. Compares two rocke trees ("arms"), typically a develop worktree and a branch -that changes ``_tiled_cache_key`` / ``_tiled_3d_cache_key``. Each arm runs in +that changes ``_tiled_cache_key`` / ``_tiled_3d_cache_key`` or, with +``--family scalar``, the scalar fallback's ``_cache_key``. Each arm runs in its own process and imports ``kernels`` and ``rocke`` from its own root, so the arms never share a module or a launcher cache. The problem sweeps live in this file, so both arms walk the same problems in the same order. @@ -31,6 +32,13 @@ binary count. Python's garbage collector is off inside each worker and runs between problems, so a collection pass never lands in a timed region. +``--family scalar`` runs both modes on the scalar fallback instead: problems +the scalar gate admits, the 2D scalar kernel (the one the launcher compiles), +and ``backend="scalar"``. Its sweeps are smaller because every problem is +lowered and the kernel is slow. A key change that also changes the kernel ABI +makes the arms' own IR differ by design, so in this family that count is +reported but does not fail the run. + Outputs go to ``--out``: per-arm raw data plus ``summary.md`` and a CSV. Measured numbers belong on the ticket or the access-controlled record, never in the repo or a PR body (see platform/AGENTS.md, Compliance). @@ -81,6 +89,13 @@ ) +# Scalar sweep: the dynamic shape the old key carried raw, with block counts +# on both sides of every block size. +_SCALAR_IR_NUM_SEQS = (1, 2, 7, 64) +_SCALAR_IR_SEQLEN_Q = (1, 16, 129) +_SCALAR_IR_SEQLEN_K = ("q", 100, 1000, 4096) + + def _shape_grid(num_seqs, seqlens_q, seqlens_k) -> Iterator[tuple[int, int, int]]: """(num_seqs, max_seqlen_q, max_seqlen_k) with k >= q and no duplicates; "q" in ``seqlens_k`` means k == q.""" @@ -115,9 +130,14 @@ def _bootstrap(root: str, arch: str | None): return au -def _route(au, problem, arch: str) -> str | None: - """Path ``run_unified_attention_torch(backend="auto")`` takes, or None for - the scalar fallback (out of scope here).""" +def _route(au, problem, arch: str, family: str) -> str | None: + """Path ``run_unified_attention_torch`` takes for ``family``, or None. + + Tiled: the ``backend="auto"`` route, None for the scalar fallback. Scalar: + ``"scalar"`` when the scalar gate admits the problem.""" + if family == "scalar": + ok = au.supports_native_unified_attention(problem, arch)[0] + return "scalar" if ok else None if problem.select_path() != "2d": if au.supports_native_unified_attention_3d_tiled(problem, arch)[0]: return "3d" @@ -127,6 +147,8 @@ def _route(au, problem, arch: str) -> str | None: def _key(au, problem, path: str, arch: str): + if path == "scalar": + return au._cache_key(problem, arch) if path == "2d": return au._tiled_cache_key(problem, arch) return au._tiled_3d_cache_key(problem, arch) @@ -141,7 +163,12 @@ def _short_hash(text: str) -> str: # --------------------------------------------------------------------------- -def _ir_sweep(problem_cls, arch: str, num_cus: int) -> Iterator: +def _ir_sweep(problem_cls, arch: str, num_cus: int, family: str) -> Iterator: + grid = ( + (_SCALAR_IR_NUM_SEQS, _SCALAR_IR_SEQLEN_Q, _SCALAR_IR_SEQLEN_K) + if family == "scalar" + else (_IR_NUM_SEQS, _IR_SEQLEN_Q, _IR_SEQLEN_K) + ) for dtype, fp8, hd, bs, heads, (sw, sinks, bias) in itertools.product( ("bf16", "fp16"), (False, True), @@ -150,7 +177,7 @@ def _ir_sweep(problem_cls, arch: str, num_cus: int) -> Iterator: ((8, 8), (32, 8), (64, 8)), _IR_FEATURES, ): - for ns, q, k in _shape_grid(_IR_NUM_SEQS, _IR_SEQLEN_Q, _IR_SEQLEN_K): + for ns, q, k in _shape_grid(*grid): yield problem_cls( total_q=ns * q, num_seqs=ns, @@ -172,7 +199,10 @@ def _ir_sweep(problem_cls, arch: str, num_cus: int) -> Iterator: def _build_identity(au, problem, path: str, arch: str) -> str: - """Full input of the binary: spec (raw num_seqs kept), builder, backend.""" + """Full input of the binary: spec (raw num_seqs kept), builder, backend. + The scalar spec is the problem itself.""" + if path == "scalar": + return repr(problem) if path == "2d": route = au._gfx942_4warp_route(problem, arch) return repr( @@ -208,7 +238,11 @@ def _pool_lower(item) -> str: au, arch = _POOL_AU, _POOL_ARCH from rocke.core.lower_llvm import _lower_kernel_to_llvm_python - if path == "2d": + if path == "scalar": + spec = au.UnifiedAttention2DSpec(problem=problem) + kernels = [au.build_unified_attention_2d(spec, arch=arch)] + tag = "scalar" + elif path == "2d": _, build, _ = au._tiled_2d_impl(arch) spec = au._tiled_spec_from_problem(problem, arch) route = au._gfx942_4warp_route(problem, arch) @@ -246,9 +280,9 @@ def _worker_ir(args) -> None: P = au.UnifiedAttentionProblem paths, keys, builds = [], [], [] rep_of_build: dict[str, tuple] = {} - for problem in _ir_sweep(P, args.arch, args.num_cus): + for problem in _ir_sweep(P, args.arch, args.num_cus, args.family): try: - path = _route(au, problem, args.arch) + path = _route(au, problem, args.arch, args.family) if path is not None: key = _key(au, problem, path, args.arch) build = _build_identity(au, problem, path, args.arch) @@ -301,10 +335,16 @@ def _worker_ir(args) -> None: ) -def _summarize_ir(arms: list[tuple[str, str]], out: Path, arches, num_cus: int) -> int: +def _summarize_ir( + arms: list[tuple[str, str]], out: Path, arches, num_cus: int, family: str +) -> int: base_name, cand_name = arms[0][0], arms[1][0] + paths = ("scalar",) if family == "scalar" else ("2d", "3d") + # Scalar: the arms' own IR differs by design when the key change also changes + # the kernel ABI, so only same-arm identity gates the run. + own_expected = "any (ABI change)" if family == "scalar" else "0" lines = [ - "# Tiled cache-key A/B: per-problem IR identity", + f"# {family.capitalize()} cache-key A/B: per-problem IR identity", "", f"Baseline arm `{base_name}`, candidate arm `{cand_name}`.", "", @@ -320,7 +360,7 @@ def _summarize_ir(arms: list[tuple[str, str]], out: Path, arches, num_cus: int) if au is None: au = _bootstrap(arms[1][1], None) P = au.UnifiedAttentionProblem - problems = list(_ir_sweep(P, arch, num_cus)) + problems = list(_ir_sweep(P, arch, num_cus, family)) n = len(problems) route_diff = [i for i in range(n) if b["path"][i] != c["path"][i]] routed = [i for i in range(n) if c["path"][i] is not None] @@ -328,37 +368,40 @@ def _summarize_ir(arms: list[tuple[str, str]], out: Path, arches, num_cus: int) cand_bad = [i for i in routed if c["served"][i] != c["own"][i]] base_bad = [i for i in routed if b["served"][i] != b["own"][i]] vs_base_own = [i for i in routed if c["served"][i] != b["own"][i]] - failed |= bool(route_diff or own_diff or cand_bad or vs_base_own) + failed |= bool(route_diff or cand_bad) + if family != "scalar": + failed |= bool(own_diff or vs_base_own) lines += [ f"## {arch}", "", f"- git: {base_name} `{b['git']}`, {cand_name} `{c['git']}`;" f" flavor {c['flavor']}; Python lowerer", - f"- problems routed to a tiled path: {len(routed)} of {n}" - f" (2D {sum(c['path'][i] == '2d' for i in routed)}," - f" 3D {sum(c['path'][i] == '3d' for i in routed)})", + f"- problems routed to the {family} path: {len(routed)} of {n} (" + + ", ".join(f"{p} {sum(c['path'][i] == p for i in routed)}" for p in paths) + + ")", f"- distinct specs lowered: {base_name} {b['lowered']}," f" {cand_name} {c['lowered']}", "", "| check | count | expected |", "|---|---|---|", f"| route differs between arms | {len(route_diff)} | 0 |", - f"| own IR differs between arms | {len(own_diff)} | 0 |", + f"| own IR differs between arms | {len(own_diff)} | {own_expected} |", f"| {cand_name} serves IR != own | {len(cand_bad)} | 0 |", - f"| {cand_name} serves IR != {base_name} own | {len(vs_base_own)} | 0 |", + f"| {cand_name} serves IR != {base_name} own | {len(vs_base_own)}" + f" | {own_expected} |", f"| {base_name} serves IR != own | {len(base_bad)} | existing gaps |", "", "| path | binaries " + base_name + " | binaries " + cand_name + " |", "|---|---|---|", ] - for path in ("2d", "3d"): + for path in paths: nb = len({b["key"][i] for i in routed if b["path"][i] == path}) nc = len({c["key"][i] for i in routed if c["path"][i] == path}) lines.append(f"| {path} | {nb} | {nc} |") lines.append("") for label, idx in ( ("route differs", route_diff), - ("own IR differs", own_diff), + ("own IR differs", own_diff if family != "scalar" else []), (f"{cand_name} serves != own", cand_bad), (f"{base_name} serves != own", base_bad), ): @@ -444,8 +487,51 @@ def _summarize_ir(arms: list[tuple[str, str]], out: Path, arches, num_cus: int) ) -def _ab_sweep(problem_cls, arch: str, num_cus: int) -> Iterator[tuple[str, object]]: - for label, static, nss, qs, ks in _AB_COHORTS: +# The scalar kernel runs one workgroup per output element with one active +# lane, so its cohorts stay small. +_SCALAR_AB_COHORTS = ( + ( + "d64-gqa4-bs16-bf16", + dict(num_query_heads=8, num_kv_heads=2, head_size=64, block_size=16), + (1, 2, 4, 7), + (1, 16, 64), + ("q", 256, 1000), + ), + ( + "d128-mha-bs32-fp16", + dict( + num_query_heads=8, + num_kv_heads=8, + head_size=128, + block_size=32, + dtype="fp16", + ), + (1, 3), + (1, 32), + (128, 512), + ), + ( + "d64-gqa4-bs16-bf16-sw64-sinks", + dict( + num_query_heads=8, + num_kv_heads=2, + head_size=64, + block_size=16, + sliding_window=64, + use_sinks=True, + ), + (2, 5), + (16,), + (300,), + ), +) + + +def _ab_sweep( + problem_cls, arch: str, num_cus: int, family: str +) -> Iterator[tuple[str, object]]: + cohorts = _SCALAR_AB_COHORTS if family == "scalar" else _AB_COHORTS + for label, static, nss, qs, ks in cohorts: static = {"dtype": "bf16", **static} if static.get("use_fp8"): static["fp8_fnuz"] = arch == "gfx942" @@ -531,7 +617,9 @@ def _reference_fp32(problem, data): ) -def _binaries(au) -> int: +def _binaries(au, family: str) -> int: + if family == "scalar": + return len(au._ATTN_CACHE) return len(au._ATTN_TILED_CACHE) + len(au._ATTN_3D_TILED_CACHE) @@ -556,9 +644,9 @@ def _worker_ab(args) -> None: # problems instead, as timeit does. gc.disable() for idx, (label, problem) in enumerate( - _ab_sweep(au.UnifiedAttentionProblem, arch, num_cus) + _ab_sweep(au.UnifiedAttentionProblem, arch, num_cus, args.family) ): - path = _route(au, problem, arch) + path = _route(au, problem, arch, args.family) if path is None: continue gc.collect() @@ -577,17 +665,17 @@ def call(): block_table=data["block_tables"], softcap=0.0, sinks=data["sinks"], - backend="auto", + backend="scalar" if args.family == "scalar" else "auto", stream=stream, ) - before = _binaries(au) + before = _binaries(au, args.family) torch.cuda.synchronize() t0 = time.perf_counter() call() torch.cuda.synchronize() first_ms = (time.perf_counter() - t0) * 1e3 - compiled = _binaries(au) - before + compiled = _binaries(au, args.family) - before ref = _reference_fp32(problem, data) max_abs = float((data["out"].float() - ref).abs().max().item()) @@ -634,14 +722,16 @@ def call(): git=_git_head(args.root), backend=backend, num_cus=num_cus, - binaries=_binaries(au), + binaries=_binaries(au, args.family), rows=rows, ) ) ) -def _summarize_ab(arms: list[tuple[str, str]], out: Path, rounds: int) -> int: +def _summarize_ab( + arms: list[tuple[str, str]], out: Path, rounds: int, family: str +) -> int: base_name, cand_name = arms[0][0], arms[1][0] runs = { name: [ @@ -652,7 +742,7 @@ def _summarize_ab(arms: list[tuple[str, str]], out: Path, rounds: int) -> int: } meta = runs[cand_name][0] if not meta["rows"]: - print(f"no problem routed to a tiled path on {meta['arch']}") + print(f"no problem routed to the {family} path on {meta['arch']}") return 1 table = [] for row in meta["rows"]: @@ -687,7 +777,7 @@ def total(name, field): geomean = math.exp(statistics.fmean(math.log(r["ratio"]) for r in table)) lines = [ - "# Tiled cache-key A/B: latency, compiles, correctness", + f"# {family.capitalize()} cache-key A/B: latency, compiles, correctness", "", f"- arch {meta['arch']}, num_cus {meta['num_cus']}, {rounds} rounds, arm order" " alternated per round, one fresh process per arm per round", @@ -787,6 +877,12 @@ def _spawn(worker: str, root: str, result: Path, extra: list[str]) -> None: def main(argv: list[str] | None = None) -> int: ap = argparse.ArgumentParser(description=__doc__.split("\n\n")[0]) ap.add_argument("--mode", choices=("ir", "ab")) + ap.add_argument( + "--family", + default="tiled", + choices=("tiled", "scalar"), + help="launcher family whose key is under test", + ) ap.add_argument( "--arm", action="append", @@ -826,7 +922,8 @@ def main(argv: list[str] | None = None) -> int: if args.mode == "ir": arches = args.arch or list(ARCHES) - common = ["--flavor", args.flavor, "--jobs", str(args.jobs)] + common = ["--family", args.family, "--flavor", args.flavor] + common += ["--jobs", str(args.jobs)] common += ["--num-cus", str(args.num_cus or 120)] for arch in arches: for name, root in arms: @@ -837,16 +934,17 @@ def main(argv: list[str] | None = None) -> int: out / f"ir_{name}_{arch}.json", ["--arch", arch, *common], ) - return _summarize_ir(arms, out, arches, args.num_cus or 120) + return _summarize_ir(arms, out, arches, args.num_cus or 120, args.family) - common = ["--iters", str(args.iters), "--warmup", str(args.warmup)] + common = ["--family", args.family] + common += ["--iters", str(args.iters), "--warmup", str(args.warmup)] common += ["--num-cus", str(args.num_cus)] for r in range(args.rounds): order = arms if r % 2 == 0 else arms[::-1] for name, root in order: print(f"[ab] round {r} arm {name}", flush=True) _spawn("ab", root, out / f"ab_{name}_r{r}.json", common) - return _summarize_ab(arms, out, args.rounds) + return _summarize_ab(arms, out, args.rounds, args.family) if __name__ == "__main__": From ec41895c34120c88aa040c71b8e91409c1e20e4f Mon Sep 17 00:00:00 2001 From: dlangbe Date: Thu, 22 Oct 2026 11:57:59 -0400 Subject: [PATCH 12/16] test(rocke): re-bless scalar attention IR goldens for the stride kernarg The scalar 2D and 3D kernels now take block_table_stride as a kernarg and drop q from their names, which moves the representative IR of the scalar cases develop added since this branch was cut. - Update sha256, bytes and kernel_name for 7 scalar cases at llvm20, llvm22 and llvm23. No other entry changes; the stale summary block is left as develop has it. Co-Authored-By: Claude --- .../rocke_representative_ir_sha256.json | 126 +++++++++--------- 1 file changed, 63 insertions(+), 63 deletions(-) diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/tests/golden/rocke_representative_ir_sha256.json b/dnn-providers/hip-kernel-provider/rocke/platform/tests/golden/rocke_representative_ir_sha256.json index cd2dfb08b99a..01408b616709 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/tests/golden/rocke_representative_ir_sha256.json +++ b/dnn-providers/hip-kernel-provider/rocke/platform/tests/golden/rocke_representative_ir_sha256.json @@ -228,58 +228,58 @@ }, "attention/gfx950/scalar2d_fp16_d128_no_mask": { "arch": "gfx950", - "bytes": 46260, + "bytes": 46297, "family": "attention", - "kernel_name": "rocke_unified_attention_2d_scalar_q64_h4_kv2_d128_b16_fp16_rbu", - "sha256": "534453293b4d4318daf0a9fad16ddf298a7ffaf8dbbfeb9cc901f269d46f4765", + "kernel_name": "rocke_unified_attention_2d_scalar_h4_kv2_d128_b16_fp16_rbu", + "sha256": "3f9a03e5ce27d736bfd18e922729a750d674374d64efba79c997e91f3057a323", "status": "ok" }, "attention/gfx950/scalar2d_fp16_d128_q40_kv8": { "arch": "gfx950", - "bytes": 46293, + "bytes": 46330, "family": "attention", - "kernel_name": "rocke_unified_attention_2d_scalar_q64_h40_kv8_d128_b16_fp16", - "sha256": "f542181c3d28ec098300697fe0c61323706a6ce2f788b3728bce99fad908acd9", + "kernel_name": "rocke_unified_attention_2d_scalar_h40_kv8_d128_b16_fp16", + "sha256": "6bb8ac9d58bce9f3aacc6dafbebeb93b9312d860cdbfbe99c0f181bb5f6fb8e6", "status": "ok" }, "attention/gfx950/scalar2d_fp16_d128_sw16_rb8": { "arch": "gfx950", - "bytes": 46594, + "bytes": 46631, "family": "attention", - "kernel_name": "rocke_unified_attention_2d_scalar_q64_h4_kv2_d128_b16_fp16_sw_rb8", - "sha256": "d81575919d468f5c88b915fac219132c9808b5fa2ccdab6031f7d868c0489c03", + "kernel_name": "rocke_unified_attention_2d_scalar_h4_kv2_d128_b16_fp16_sw_rb8", + "sha256": "1b0fd719499c69cbb6672949b471bf52ff9f5f14fc6402ef3e97a485643c8111", "status": "ok" }, "attention/gfx950/scalar2d_fp16_d128_top_left": { "arch": "gfx950", - "bytes": 46364, + "bytes": 46401, "family": "attention", - "kernel_name": "rocke_unified_attention_2d_scalar_q64_h4_kv2_d128_b16_fp16_tl", - "sha256": "e8104b7de52d1e84b12f28a0e836dcee13477021397d3b6dbea935c1a296a335", + "kernel_name": "rocke_unified_attention_2d_scalar_h4_kv2_d128_b16_fp16_tl", + "sha256": "4e9da57dfd9249086c6aeeeb9baefd731f50941b94a029ab26967ea4eea364ae", "status": "ok" }, "attention/gfx950/scalar3d_fp16_d128_q40_kv8": { "arch": "gfx950", - "bytes": 48783, + "bytes": 48834, "family": "attention", - "kernel_name": "rocke_unified_attention_3d_scalar_q64_h40_kv8_d128_b16_seg8_fp16", - "sha256": "9804fc433bba54dffb78e7f36f8f303d185a9401dc93b05d1da75e9b304077d1", + "kernel_name": "rocke_unified_attention_3d_scalar_h40_kv8_d128_b16_seg8_fp16", + "sha256": "41dcdeed861eb19399f523ec4715c28d48eb8fbc1622c0e4bd1e58e9f51135aa", "status": "ok" }, "attention/gfx950/scalar3d_fp16_d128_sw16_rb8": { "arch": "gfx950", - "bytes": 49095, + "bytes": 49146, "family": "attention", - "kernel_name": "rocke_unified_attention_3d_scalar_q64_h4_kv2_d128_b16_seg8_fp16_rb8_sw", - "sha256": "5d8bf24697392fca54352be7717cfeac566b5d6dd1ed8df23f7b652cdf8a7357", + "kernel_name": "rocke_unified_attention_3d_scalar_h4_kv2_d128_b16_seg8_fp16_rb8_sw", + "sha256": "7e27fcad07b5406d73e4ca69e37cecd70fb4694e2ba7d0b0c421a705e50d9575", "status": "ok" }, "attention/gfx950/scalar3d_fp16_d128_top_left": { "arch": "gfx950", - "bytes": 48855, + "bytes": 48908, "family": "attention", - "kernel_name": "rocke_unified_attention_3d_scalar_q64_h4_kv2_d128_b16_seg8_fp16_tl", - "sha256": "49a72d16ffa227df568072c50f77647261cb8b8023669fcec8321c76fd9a6429", + "kernel_name": "rocke_unified_attention_3d_scalar_h4_kv2_d128_b16_seg8_fp16_tl", + "sha256": "5427e7f6a7e14f59bd67d240872d1ca8d05b7818b3443edda3b57c9949874cab", "status": "ok" }, "attention_d128_swa/gfx942/4warp_gqa_fold": { @@ -2223,58 +2223,58 @@ }, "attention/gfx950/scalar2d_fp16_d128_no_mask": { "arch": "gfx950", - "bytes": 46267, + "bytes": 46304, "family": "attention", - "kernel_name": "rocke_unified_attention_2d_scalar_q64_h4_kv2_d128_b16_fp16_rbu", - "sha256": "51fe41aaaa3623be88d35d880007c86f723d1821ae0d32a511a3bbe530c682c4", + "kernel_name": "rocke_unified_attention_2d_scalar_h4_kv2_d128_b16_fp16_rbu", + "sha256": "7bfb21eb2170edf8d4256e5cba6726a66fbcafa8915278bd59ad710de5b3e33e", "status": "ok" }, "attention/gfx950/scalar2d_fp16_d128_q40_kv8": { "arch": "gfx950", - "bytes": 46300, + "bytes": 46337, "family": "attention", - "kernel_name": "rocke_unified_attention_2d_scalar_q64_h40_kv8_d128_b16_fp16", - "sha256": "946bac7a0332cf806cbc23efe6128d7e8a147a7de5542ce27bd3d239af24d4e9", + "kernel_name": "rocke_unified_attention_2d_scalar_h40_kv8_d128_b16_fp16", + "sha256": "e44040c1a0ca888c9265a7179849e9817bb4c86453d1ac4f17ffd6d878f6e00c", "status": "ok" }, "attention/gfx950/scalar2d_fp16_d128_sw16_rb8": { "arch": "gfx950", - "bytes": 46601, + "bytes": 46638, "family": "attention", - "kernel_name": "rocke_unified_attention_2d_scalar_q64_h4_kv2_d128_b16_fp16_sw_rb8", - "sha256": "1acd98d1834d52cec79c2a70bd09572d96d168423fdc6a11a1deac4e95901916", + "kernel_name": "rocke_unified_attention_2d_scalar_h4_kv2_d128_b16_fp16_sw_rb8", + "sha256": "9d2cd3e99bbb6496c4b592fe0b207d57d8010a5519492c22be25f7103d035187", "status": "ok" }, "attention/gfx950/scalar2d_fp16_d128_top_left": { "arch": "gfx950", - "bytes": 46371, + "bytes": 46408, "family": "attention", - "kernel_name": "rocke_unified_attention_2d_scalar_q64_h4_kv2_d128_b16_fp16_tl", - "sha256": "56f3f2d5411ee663955e06e799c0be8770248af65df1d5dc7fede2ea8ff048be", + "kernel_name": "rocke_unified_attention_2d_scalar_h4_kv2_d128_b16_fp16_tl", + "sha256": "4ee7f6b39f7d0f87e5d06dff0041f97c72a0179398047310bc29eeb70b734a73", "status": "ok" }, "attention/gfx950/scalar3d_fp16_d128_q40_kv8": { "arch": "gfx950", - "bytes": 48790, + "bytes": 48841, "family": "attention", - "kernel_name": "rocke_unified_attention_3d_scalar_q64_h40_kv8_d128_b16_seg8_fp16", - "sha256": "188f9394d8d01d0f00aea78d2172ea2836bac7612e7f6a338d6ccd04f73a0151", + "kernel_name": "rocke_unified_attention_3d_scalar_h40_kv8_d128_b16_seg8_fp16", + "sha256": "1d905ab5c0eba1cb417cfa4cf3b6c43dfb0316a4fa9c7e23aab2a2dfbff2089d", "status": "ok" }, "attention/gfx950/scalar3d_fp16_d128_sw16_rb8": { "arch": "gfx950", - "bytes": 49102, + "bytes": 49153, "family": "attention", - "kernel_name": "rocke_unified_attention_3d_scalar_q64_h4_kv2_d128_b16_seg8_fp16_rb8_sw", - "sha256": "1e75fb5a2c80f393ec5687d943261ba4ee4afa9304e2fb6f143a1e1accc3fd07", + "kernel_name": "rocke_unified_attention_3d_scalar_h4_kv2_d128_b16_seg8_fp16_rb8_sw", + "sha256": "dec69e6c8561b0b8143243e254531dd0224320d1a63123b2f54208c001d66415", "status": "ok" }, "attention/gfx950/scalar3d_fp16_d128_top_left": { "arch": "gfx950", - "bytes": 48862, + "bytes": 48915, "family": "attention", - "kernel_name": "rocke_unified_attention_3d_scalar_q64_h4_kv2_d128_b16_seg8_fp16_tl", - "sha256": "99353476cc9677d60f117ba680c708c996663844c0bb4b51f2a16308948a65ba", + "kernel_name": "rocke_unified_attention_3d_scalar_h4_kv2_d128_b16_seg8_fp16_tl", + "sha256": "a3ec8c6fbc92410df68aeb62edc5589c354f898242d2e3334c55b9b2df8ab7f3", "status": "ok" }, "attention_d128_swa/gfx942/4warp_gqa_fold": { @@ -4218,58 +4218,58 @@ }, "attention/gfx950/scalar2d_fp16_d128_no_mask": { "arch": "gfx950", - "bytes": 46331, + "bytes": 46368, "family": "attention", - "kernel_name": "rocke_unified_attention_2d_scalar_q64_h4_kv2_d128_b16_fp16_rbu", - "sha256": "1439195339b0bf240984da6a0e706798211e6266b0effecd21c730e056221d71", + "kernel_name": "rocke_unified_attention_2d_scalar_h4_kv2_d128_b16_fp16_rbu", + "sha256": "75a15fc5580b38822ceb1c20b63d62f4e807e8e7e2bfff87c0268a19f55e7b4d", "status": "ok" }, "attention/gfx950/scalar2d_fp16_d128_q40_kv8": { "arch": "gfx950", - "bytes": 46364, + "bytes": 46401, "family": "attention", - "kernel_name": "rocke_unified_attention_2d_scalar_q64_h40_kv8_d128_b16_fp16", - "sha256": "7a80904e488e71aa45b5a983b3ca207e411f4797b21e8b68bc1a5ef0e5a2772f", + "kernel_name": "rocke_unified_attention_2d_scalar_h40_kv8_d128_b16_fp16", + "sha256": "e58b5d9f581c300df4066d6baa7bfd0b3a6180c50a18f4fe34baaff875e79dc3", "status": "ok" }, "attention/gfx950/scalar2d_fp16_d128_sw16_rb8": { "arch": "gfx950", - "bytes": 46665, + "bytes": 46702, "family": "attention", - "kernel_name": "rocke_unified_attention_2d_scalar_q64_h4_kv2_d128_b16_fp16_sw_rb8", - "sha256": "06b90e3c577d968c02ac6171c6f0b7bda0d798b36c2059818111baa0b8994965", + "kernel_name": "rocke_unified_attention_2d_scalar_h4_kv2_d128_b16_fp16_sw_rb8", + "sha256": "3507ed5f852025708ebc67b95e6831bca37d2937c336e030c05793a17755e8d7", "status": "ok" }, "attention/gfx950/scalar2d_fp16_d128_top_left": { "arch": "gfx950", - "bytes": 46435, + "bytes": 46472, "family": "attention", - "kernel_name": "rocke_unified_attention_2d_scalar_q64_h4_kv2_d128_b16_fp16_tl", - "sha256": "076fcd4b4e1f618f5c75534cba8333098b03497e5393d157543433476bb3d158", + "kernel_name": "rocke_unified_attention_2d_scalar_h4_kv2_d128_b16_fp16_tl", + "sha256": "af239152736fb1333cf0b27337db79f7389af7b0eace59060696a629a3a1f75d", "status": "ok" }, "attention/gfx950/scalar3d_fp16_d128_q40_kv8": { "arch": "gfx950", - "bytes": 48854, + "bytes": 48905, "family": "attention", - "kernel_name": "rocke_unified_attention_3d_scalar_q64_h40_kv8_d128_b16_seg8_fp16", - "sha256": "f2041465d5d9e8128b2ab20a34dc1baaaad76b9b818b6c69c3f64c590b1e0a37", + "kernel_name": "rocke_unified_attention_3d_scalar_h40_kv8_d128_b16_seg8_fp16", + "sha256": "888055e6c6a30f6d67cd73d46d375e28e1ab48eae8d3c71a1106e6a971d8de89", "status": "ok" }, "attention/gfx950/scalar3d_fp16_d128_sw16_rb8": { "arch": "gfx950", - "bytes": 49166, + "bytes": 49217, "family": "attention", - "kernel_name": "rocke_unified_attention_3d_scalar_q64_h4_kv2_d128_b16_seg8_fp16_rb8_sw", - "sha256": "0b5f59549c2dba463219d00aea1c016c547479fe1c5a3ea51daccdf3c8973ee1", + "kernel_name": "rocke_unified_attention_3d_scalar_h4_kv2_d128_b16_seg8_fp16_rb8_sw", + "sha256": "68805f1e75cd1e88e20e82aeb8727e8fb12c873041d76b3f4c10423a79716cc4", "status": "ok" }, "attention/gfx950/scalar3d_fp16_d128_top_left": { "arch": "gfx950", - "bytes": 48926, + "bytes": 48979, "family": "attention", - "kernel_name": "rocke_unified_attention_3d_scalar_q64_h4_kv2_d128_b16_seg8_fp16_tl", - "sha256": "1b55556afefa9291e98a2a07cea73a144f8f9be96fb1790e0d3d3ab2c690d7ad", + "kernel_name": "rocke_unified_attention_3d_scalar_h4_kv2_d128_b16_seg8_fp16_tl", + "sha256": "6f3dc2facda457fc1f91e1c50f6edd2ecdd08f03b7a767c85be21a66eb132589", "status": "ok" }, "attention_d128_swa/gfx942/4warp_gqa_fold": { From 6a6ca5a72e4fa6f47d64f59c86d21a8d82532569 Mon Sep 17 00:00:00 2001 From: dlangbe Date: Fri, 23 Oct 2026 11:01:58 -0400 Subject: [PATCH 13/16] refactor(rocke): share the scalar block-table lookup between 2D and 3D The scalar 2D body repeated the block-table lookup that the 3D path already takes from a helper, so the two could drift apart. - Call _physical_block_and_token from build_unified_attention_2d. - Call rocke_unified_attn_physical_block_and_token from the C++ 2D softmax loop. Its default alignment resolves to 4, the same as the call it replaces. No emitted IR changes. Co-Authored-By: Claude --- .../kernels/common/attention_unified.py | 5 ++--- .../instances/common/attention_unified_2d.cpp | 21 ++++++++++--------- 2 files changed, 13 insertions(+), 13 deletions(-) diff --git a/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/attention_unified.py b/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/attention_unified.py index e13344bbb6d7..7de67d567c89 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/attention_unified.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/attention_unified.py @@ -6333,9 +6333,8 @@ def build_unified_attention_2d( iv_name="kpos", ) with loop as (kpos, (m_val, l_val, acc_val)): - block_idx, token_in_block = _magic_div_mod(b, kpos, p.block_size) - physical = b.global_load_i32( - block_tables, b.add(b.mul(seq_idx, bt_stride), block_idx) + physical, token_in_block = _physical_block_and_token( + b, p, block_tables, bt_stride, seq_idx, kpos ) # Vectorised QK dot product. Both Q[q_tok, q_head, :] and diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/common/attention_unified_2d.cpp b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/common/attention_unified_2d.cpp index 66cf4b20e881..6389cbb01ad2 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/common/attention_unified_2d.cpp +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/common/attention_unified_2d.cpp @@ -115,17 +115,18 @@ void rocke_attn_unified_emit_2d_softmax_loop(rocke_attn_unified_build_ctx_t* ctx rocke_value_t* l_val = loop.iter_vars[1]; rocke_value_t* acc_val = loop.iter_vars[2]; - /* block_idx, token_in_block = _magic_div_mod(b, kpos, p.block_size) */ - rocke_value_t* block_idx = NULL; + /* physical, token_in_block = _physical_block_and_token( + * b, p, block_tables, bt_stride, seq_idx, kpos) */ + rocke_value_t* physical = NULL; rocke_value_t* token_in_block = NULL; - (void)rocke_unified_attn_magic_div_mod(b, kpos, p->block_size, &block_idx, &token_in_block); - - /* physical = global_load_i32(block_tables, seq_idx * bt_stride + block_idx) */ - rocke_value_t* physical = rocke_b_global_load_i32( - b, - ctx->abi_block_tables, - rocke_b_add(b, rocke_b_mul(b, ctx->seq_idx, ctx->block_table_stride), block_idx), - /*align=*/4); + (void)rocke_unified_attn_physical_block_and_token(b, + &ctx->sel_p, + ctx->abi_block_tables, + ctx->block_table_stride, + ctx->seq_idx, + kpos, + &physical, + &token_in_block); /* Vectorised QK dot product (vec8 main fold + scalar tail). */ rocke_value_t* score = ctx->zero_f; From 405e46c7aa597e1b9eab47eac013cf55923e3dbd Mon Sep 17 00:00:00 2001 From: dlangbe Date: Fri, 23 Oct 2026 11:02:27 -0400 Subject: [PATCH 14/16] test(rocke): run the scalar attention stride tests in CI Neither scalar test file was in a CTest list, and the GPU test needed torch, which the CI GPU lane does not have. - Add test_attention_scalar_cache_key.py to the attention CPU lane and test_attention_scalar_block_table_numeric.py to the library GPU lane. - Launch the GPU test through run_unified_attention_torch with rocke.runtime device buffers and a numpy fp32 reference, so it runs without torch. It skips on devices whose target id the attention builders do not know. - Add a CPU test that drives run_unified_attention_torch with a recording launcher and checks the table's real row stride is packed. Co-Authored-By: Claude --- ...st_attention_scalar_block_table_numeric.py | 146 ++++++++++++------ .../tests/test_attention_scalar_cache_key.py | 53 +++++++ .../rocke/platform/CMakeLists.txt | 5 +- 3 files changed, 152 insertions(+), 52 deletions(-) diff --git a/dnn-providers/hip-kernel-provider/rocke/library/tests/test_attention_scalar_block_table_numeric.py b/dnn-providers/hip-kernel-provider/rocke/library/tests/test_attention_scalar_block_table_numeric.py index 6db01be49d0a..cdf3703a6c30 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/tests/test_attention_scalar_block_table_numeric.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/tests/test_attention_scalar_block_table_numeric.py @@ -10,51 +10,98 @@ poison block, so a wrong row stride corrupts the output instead of passing by luck. -Reference attention is computed independently in FP32 from the logical K/V -after rounding inputs to bf16. Requires a ROCm GPU; skips otherwise. +The launch goes through ``run_unified_attention_torch(backend="scalar")`` with +``rocke.runtime`` device buffers instead of torch tensors, so the test runs in +GPU lanes without torch. Reference attention is computed independently in FP32 +with numpy from the logical K/V, after rounding inputs to fp16. Skips without a +GPU, and on devices whose target id the attention builders do not know. """ +import ctypes import itertools import math +import numpy as np import pytest from kernels.common import attention_unified as au +from rocke.core.arch.target import known_arches +from rocke.runtime.hip_module import get_device_arch _HQ, _HKV, _D, _BS = 8, 2, 64, 16 _ATOL = 1e-2 +# The launch compiles for the device's own target id, which the builders must +# know (gfx11 devices are only known through their generic target). +_ARCH = get_device_arch() +pytestmark = [ + pytest.mark.gpu, + pytest.mark.skipif( + _ARCH not in known_arches(), + reason=f"needs a GPU with a known attention target, found {_ARCH}", + ), +] + @pytest.fixture(scope="module") -def torch(): - torch = pytest.importorskip("torch") - if not torch.cuda.is_available(): - pytest.skip("requires a ROCm GPU") - return torch +def rt(): + from rocke.runtime.hip_module import Runtime + + return Runtime() + + +def _u8(array): + return (ctypes.c_uint8 * array.nbytes).from_buffer(array) + + +class _DeviceArray: + """A device copy of a numpy array, exposing the tensor attributes + ``run_unified_attention_torch`` reads. ``shape`` may be narrower than the + host array along the last axis, which models a column view of it.""" + + device = None + + def __init__(self, rt, host, *, shape=None): + from rocke.runtime.launcher import DeviceMem + self._rt = rt + self._host = np.ascontiguousarray(host) + self._mem = DeviceMem(self._host.nbytes) + rt.memcpy_h2d(self._mem.ptr(), _u8(self._host), self._host.nbytes) + self.shape = self._host.shape if shape is None else shape -def _case(torch, q_lens, kv_lens, row_width, *, view_of_wider=False, seed=0): - """Build a paged case whose block-table rows hold ``row_width`` slots.""" - gen = torch.Generator().manual_seed(seed) + def stride(self, dim): + return self._host.strides[dim] // self._host.itemsize + + def data_ptr(self): + return self._mem.ptr() + + def read(self): + out = np.empty_like(self._host) + self._rt.memcpy_d2h(_u8(out), self._mem.ptr(), out.nbytes) + return out + + +def _case(q_lens, kv_lens, row_width, *, seed): + """A paged case whose block-table rows hold ``row_width`` slots.""" + rng = np.random.default_rng(seed) needed = [math.ceil(n / _BS) for n in kv_lens] num_blocks = sum(needed) + 1 poison = num_blocks - 1 - order = torch.randperm(num_blocks - 1, generator=gen).tolist() - alloc_width = row_width + 5 if view_of_wider else row_width - table = torch.full((len(kv_lens), alloc_width), poison, dtype=torch.int32) + order = iter(rng.permutation(num_blocks - 1)) + table = np.full((len(kv_lens), row_width), poison, dtype=np.int32) for s, n in enumerate(needed): - table[s, :n] = torch.tensor(order[:n], dtype=torch.int32) - order = order[n:] + table[s, :n] = [next(order) for _ in range(n)] def randn(*shape): - return torch.randn(*shape, generator=gen).to(torch.bfloat16) + return rng.standard_normal(shape).astype(np.float16) q = randn(sum(q_lens), _HQ, _D) k = randn(num_blocks, _BS, _HKV, _D) v = randn(num_blocks, _BS, _HKV, _D) k[poison] = 64.0 v[poison] = 64.0 - cu = torch.tensor([0, *itertools.accumulate(q_lens)], dtype=torch.int32) + cu = np.array([0, *itertools.accumulate(q_lens)], dtype=np.int32) problem = au.UnifiedAttentionProblem( total_q=sum(q_lens), num_seqs=len(q_lens), @@ -64,78 +111,75 @@ def randn(*shape): block_size=_BS, max_seqlen_q=max(q_lens), max_seqlen_k=max(kv_lens), - dtype="bf16", + dtype="fp16", ) - dev = table.cuda() - if view_of_wider: - dev = dev[:, :row_width] - return problem, q, k, v, cu, table, dev + return problem, q, k, v, cu, table -def _reference(torch, q, k, v, cu, table, q_lens, kv_lens): - out = torch.zeros(q.shape, dtype=torch.float32) +def _reference(q, k, v, cu, table, q_lens, kv_lens): + out = np.zeros(q.shape, dtype=np.float32) ratio = _HQ // _HKV for s, (q_len, kv_len) in enumerate(zip(q_lens, kv_lens)): - blocks = table[s, : math.ceil(kv_len / _BS)].long() - keys = k[blocks].reshape(-1, _HKV, _D)[:kv_len].float() - values = v[blocks].reshape(-1, _HKV, _D)[:kv_len].float() + blocks = table[s, : math.ceil(kv_len / _BS)] + keys = k[blocks].reshape(-1, _HKV, _D)[:kv_len].astype(np.float32) + values = v[blocks].reshape(-1, _HKV, _D)[:kv_len].astype(np.float32) context = kv_len - q_len for i in range(q_len): tok = int(cu[s]) + i visible = context + i + 1 for h in range(_HQ): - scores = keys[:visible, h // ratio] @ q[tok, h].float() - probs = (scores / math.sqrt(_D)).softmax(0) - out[tok, h] = probs @ values[:visible, h // ratio] + scores = keys[:visible, h // ratio] @ q[tok, h].astype(np.float32) + scores /= math.sqrt(_D) + probs = np.exp(scores - scores.max()) + out[tok, h] = (probs / probs.sum()) @ values[:visible, h // ratio] return out -def _launch(torch, problem, q, k, v, cu, kv_lens, block_table): - out = torch.empty(q.shape, dtype=q.dtype, device="cuda") +def _launch(rt, problem, q, k, v, cu, kv_lens, block_table): + out = _DeviceArray(rt, np.zeros_like(q)) au.run_unified_attention_torch( problem=problem, - q=q.cuda(), - k=k.cuda(), - v=v.cuda(), + q=_DeviceArray(rt, q), + k=_DeviceArray(rt, k), + v=_DeviceArray(rt, v), out=out, - cu_seqlens_q=cu.cuda(), - seqused_k=torch.tensor(kv_lens, dtype=torch.int32, device="cuda"), + cu_seqlens_q=_DeviceArray(rt, cu), + seqused_k=_DeviceArray(rt, np.array(kv_lens, dtype=np.int32)), softmax_scale=1.0 / math.sqrt(_D), block_table=block_table, softcap=0.0, backend="scalar", ) - torch.cuda.synchronize() - return out.float().cpu() + return out.read().astype(np.float32) @pytest.mark.parametrize("view_of_wider", (False, True), ids=("wide", "view")) -def test_wide_block_table_rows(torch, view_of_wider): +def test_wide_block_table_rows(rt, view_of_wider): q_lens, kv_lens = [5, 1, 9], [37, 20, 61] row_width = math.ceil(max(kv_lens) / _BS) + 6 - problem, q, k, v, cu, table, dev = _case( - torch, q_lens, kv_lens, row_width, view_of_wider=view_of_wider - ) + alloc_width = row_width + 5 if view_of_wider else row_width + problem, q, k, v, cu, table = _case(q_lens, kv_lens, alloc_width, seed=0) + dev = _DeviceArray(rt, table, shape=(len(q_lens), row_width)) assert dev.stride(0) > math.ceil(problem.max_seqlen_k / _BS) - got = _launch(torch, problem, q, k, v, cu, kv_lens, dev) - want = _reference(torch, q, k, v, cu, table, q_lens, kv_lens) - max_abs = (got - want).abs().max().item() + got = _launch(rt, problem, q, k, v, cu, kv_lens, dev) + want = _reference(q, k, v, cu, table, q_lens, kv_lens) + max_abs = float(np.abs(got - want).max()) assert max_abs <= _ATOL, f"max_abs {max_abs:.3e} vs fp32 reference" -def test_one_binary_serves_different_batches(torch): +def test_one_binary_serves_different_batches(rt): """Shapes that differ only in batch geometry share a key; the second launch must reuse the first binary and still match the reference.""" first = ([4, 2], [30, 17], 4) second = ([1, 7, 3, 2], [90, 64, 15, 33], 12) launchers = None for i, (q_lens, kv_lens, width) in enumerate((first, second)): - problem, q, k, v, cu, table, dev = _case(torch, q_lens, kv_lens, width, seed=i) - got = _launch(torch, problem, q, k, v, cu, kv_lens, dev) + problem, q, k, v, cu, table = _case(q_lens, kv_lens, width, seed=i) + got = _launch(rt, problem, q, k, v, cu, kv_lens, _DeviceArray(rt, table)) if launchers is None: launchers = len(au._SCALAR_LAUNCHERS) else: assert len(au._SCALAR_LAUNCHERS) == launchers - want = _reference(torch, q, k, v, cu, table, q_lens, kv_lens) - max_abs = (got - want).abs().max().item() + want = _reference(q, k, v, cu, table, q_lens, kv_lens) + max_abs = float(np.abs(got - want).max()) assert max_abs <= _ATOL, f"case {i}: max_abs {max_abs:.3e}" diff --git a/dnn-providers/hip-kernel-provider/rocke/library/tests/test_attention_scalar_cache_key.py b/dnn-providers/hip-kernel-provider/rocke/library/tests/test_attention_scalar_cache_key.py index 7591be089ed7..518e72289017 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/tests/test_attention_scalar_cache_key.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/tests/test_attention_scalar_cache_key.py @@ -15,6 +15,7 @@ byte-identical IR at every LLVM flavor. * Every static field the build reads splits the key. * The launch signature matches the kernel's params, including the stride. + * ``run_unified_attention_torch`` packs the table's real row stride. No GPU required. ``test_attention_scalar_block_table_numeric.py`` runs the launch path on hardware. @@ -127,3 +128,55 @@ def test_launch_signature_matches_kernel_params(): _BASE.dtype, include_bt_stride=True, kv_dtype=au._kv_storage_dtype(_BASE) ) assert [item["name"] for item in signature] == [p.name for p in kernel.params] + + +class _FakeArray: + """The tensor attributes ``run_unified_attention_torch`` reads on the + scalar path; nothing is launched, so the pointer is never dereferenced.""" + + device = None + + def __init__(self, shape, *, row_stride=None): + self.shape = shape + self._row_stride = shape[-1] if row_stride is None else row_stride + + def stride(self, dim): + assert dim == 0 + return self._row_stride + + def data_ptr(self): + return 0x1000 + + +@pytest.mark.parametrize("row_stride", (16, 24), ids=("exact", "wider")) +def test_scalar_launch_packs_the_table_row_stride(monkeypatch, row_stride): + """The scalar launch packs ``block_table.stride(0)``, not a stride derived + from ``max_seqlen_k``. 24 models a column view of a wider table: its rows + hold 16 slots, but the next row starts 24 slots on.""" + launches = [] + monkeypatch.setattr(au, "_RESOLVED_ATTENTION_ARCH", "gfx950") + monkeypatch.setattr( + au, + "_get_scalar_launcher", + lambda problem, key, *, arch: lambda vals, *, config: launches.append( + (vals, config) + ), + ) + p = _BASE + blocks_per_seq = -(-p.max_seqlen_k // p.block_size) + au.run_unified_attention_torch( + problem=p, + q=_FakeArray((p.total_q, p.num_query_heads, p.head_size)), + k=_FakeArray((64, p.block_size, p.num_kv_heads, p.head_size)), + v=_FakeArray((64, p.block_size, p.num_kv_heads, p.head_size)), + out=_FakeArray((p.total_q, p.num_query_heads, p.head_size)), + cu_seqlens_q=_FakeArray((p.num_seqs + 1,)), + seqused_k=_FakeArray((p.num_seqs,)), + softmax_scale=0.125, + block_table=_FakeArray((p.num_seqs, blocks_per_seq), row_stride=row_stride), + softcap=0.0, + backend="scalar", + ) + [(vals, config)] = launches + assert vals["block_table_stride"] == row_stride + assert config.grid == (p.total_q, p.num_query_heads, p.head_size) diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/CMakeLists.txt b/dnn-providers/hip-kernel-provider/rocke/platform/CMakeLists.txt index d5a2229d1367..81d17eebfba9 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/CMakeLists.txt +++ b/dnn-providers/hip-kernel-provider/rocke/platform/CMakeLists.txt @@ -497,11 +497,13 @@ if(ROCKE_INSTALL_TESTS) " TIMEOUT 180)\n") # Library attention CPU tests (no GPU required): gfx950 LDS budget model, - # spec-wrapper parity gate, and the 3D decode split rule and K/V cache policy. + # spec-wrapper parity gate, the 3D decode split rule and K/V cache policy, and + # the scalar fallback's cache key and block-table stride kernarg. # File-list style (same pattern as conv_pytest) so PYTHONPATH resolves rocke/ # correctly from the installed tree. set(_ROCKE_LIBRARY_ATTENTION_CPU_TEST_RELPATHS "tests/library/tests/test_attention_dense_gfx950_lds.py" + "tests/library/tests/test_attention_scalar_cache_key.py" "tests/library/tests/test_attention_spec_predicate_parity.py" "tests/library/tests/test_attention_tiled_3d_decode_split.py" "tests/library/tests/test_attention_tiled_3d_kv_cache_policy.py") @@ -695,6 +697,7 @@ if(ROCKE_INSTALL_TESTS) "tests/library/tests/test_conv_dgrad_correctness.py" "tests/library/tests/test_direct_conv_correctness.py" "tests/library/tests/test_attention_dense_gfx942_numeric.py" + "tests/library/tests/test_attention_scalar_block_table_numeric.py" "tests/library/tests/test_attn_bf16_d128_ring.py" "tests/library/tests/test_dtype_dispatch.py" "tests/library/tests/test_gdn_decode_gfx950_numeric.py" From 7189edc583ff998eaea95c0158368ad79005c023 Mon Sep 17 00:00:00 2001 From: dlangbe Date: Sun, 25 Oct 2026 11:29:20 -0400 Subject: [PATCH 15/16] test(rocke): cover the remaining scalar key fields and drop stale max_blocks docs Review follow-ups for the scalar cache-key change. - Add use_fp8, causal_top_left and right_bound to _STATIC_VARIANTS. _cache_key keeps all three and the 2D build reads them, so the test now fails if a later change drops one from the key. - Remove the max_blocks mentions from the ctx_init docs in the C++ prologue and internal header; the field went away with the stride kernarg. Co-Authored-By: Claude --- .../rocke/library/tests/test_attention_scalar_cache_key.py | 3 +++ .../cpp/include/rocke/instance_attention_unified_internal.h | 5 ++--- .../cpp/instances/common/attention_unified_prologue.cpp | 3 +-- 3 files changed, 6 insertions(+), 5 deletions(-) diff --git a/dnn-providers/hip-kernel-provider/rocke/library/tests/test_attention_scalar_cache_key.py b/dnn-providers/hip-kernel-provider/rocke/library/tests/test_attention_scalar_cache_key.py index 518e72289017..819cc9217006 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/tests/test_attention_scalar_cache_key.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/tests/test_attention_scalar_cache_key.py @@ -67,6 +67,9 @@ dict(sliding_window=128), dict(use_sinks=True), dict(softcap=30.0), + dict(use_fp8=True), + dict(causal_top_left=True), + dict(right_bound=8), ) diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/instance_attention_unified_internal.h b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/instance_attention_unified_internal.h index 5e3605c426f0..13f7aaf61705 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/instance_attention_unified_internal.h +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/instance_attention_unified_internal.h @@ -209,9 +209,8 @@ typedef struct rocke_attn_unified_build_ctx * ============================================================ */ /* Zero-init the ctx and copy in the problem/spec slice (b, kind, p, sel_p, - * dtype, num_segments) + derive num_queries_per_kv / max_blocks. On a non- - * divisible head ratio or unsupported dtype, sets b's sticky error and returns - * false. */ + * dtype, num_segments) + derive num_queries_per_kv. On a non-divisible head + * ratio or unsupported dtype, sets b's sticky error and returns false. */ bool rocke_attn_unified_ctx_init(rocke_attn_unified_build_ctx_t* ctx, rocke_ir_builder_t* b, rocke_attn_unified_kind_t kind, diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/common/attention_unified_prologue.cpp b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/common/attention_unified_prologue.cpp index 5e1c10f828e0..c6f6ad3f9796 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/common/attention_unified_prologue.cpp +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/common/attention_unified_prologue.cpp @@ -40,8 +40,7 @@ * if p.dtype not in ("fp16", "bf16"): raise ValueError(...) (2D only) * dtype = spec.dtype_ir * plus deriving num_queries_per_kv (= p.num_query_heads // num_kv_heads, which - * raises on a non-divisible ratio) and max_blocks - * (= ceil(max_seqlen_k / block_size)). + * raises on a non-divisible ratio). * * Zero-inits the ctx, copies the problem/spec slice into it, mirrors `p` into * the selector-helper problem shape (ctx->sel_p) so the ported emit/descriptor From 56ae95258ca18b7fba6a79314093c2dd048d9ea9 Mon Sep 17 00:00:00 2001 From: dlangbe Date: Wed, 28 Oct 2026 13:40:35 -0400 Subject: [PATCH 16/16] test(rocke): re-bless the scalar attention differential golden The differential golden now runs in CI and anchors the scalar attention_unified and attention_unified_stress parity families. Their IR changed with the block_table_stride kernarg and the q name drop, so their hashes are stale. - Re-record both families with --only attention_unified: 38 configs at llvm20, llvm22, llvm23 and in the canonical section. No other family's hashes change. --- .../differential/golden/llvm_gfx_all.json | 304 +++++++++--------- 1 file changed, 152 insertions(+), 152 deletions(-) diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/tests/instances/differential/golden/llvm_gfx_all.json b/dnn-providers/hip-kernel-provider/rocke/platform/tests/instances/differential/golden/llvm_gfx_all.json index 168a95b7ce71..7215500602ed 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/tests/instances/differential/golden/llvm_gfx_all.json +++ b/dnn-providers/hip-kernel-provider/rocke/platform/tests/instances/differential/golden/llvm_gfx_all.json @@ -33,46 +33,46 @@ "9": "b932c7ec5c473d26f9c5d3364d503eba33b0d916d517114842dff4352274241e" }, "attention_unified": { - "0": "2dd9c60d1dba62d1dba0402468528886c520441e3447900b690b17e0f6a84a3c", - "1": "500a2736bc4eba1f612490447600591e4d7002b05b292346b7787fe779d4b67d", - "2": "ec98172e062d413a143e0e987421bb8cda6bf4d91e5cec5a201859f85284155f", - "3": "f8888fdc0c659524e857dfd909bd9897f5c1cbbec38c6807f96489730a7bebbe", - "4": "2bcac227c33864e46a72a5609b6f7f9107580f3ebbfd2232525499a23a710599", - "5": "507240ae9e57b775882908f4e0fb21725cdc2dec4ab29013271ec1a3878e0191" + "0": "31365624504bcda43d0d0cfafcfb6c259ca957a82f4f2d85ba284cd682695eb9", + "1": "92576fd0ba5f17da40f9bf1d5b0663ed7f4dbe74df8ad76c56aa9f20004442fa", + "2": "30e88d793bf93069b76f4918ed5ada53447ade3f2f280bcaeb0ed9b23ef22cc7", + "3": "7f59dca66aa6119cc19a8268751e54d86c60ae8554bb91eb474760ba34f4e2f6", + "4": "c230ddb11e58111e9f7fdb6d75024350207d9ed77a19baac2ba23b8a6e51ecb6", + "5": "43254a2c4eed8a972f2b17ffba1087931728145c610318b644570c3e6a44c512" }, "attention_unified_stress": { - "0": "2dd9c60d1dba62d1dba0402468528886c520441e3447900b690b17e0f6a84a3c", - "1": "500a2736bc4eba1f612490447600591e4d7002b05b292346b7787fe779d4b67d", - "10": "0252808c628de2641859b01fb8c905dae101fd1c89c0614cd6d9e2302864de92", - "11": "344929d8ba914413685616a57b9ef7d85c0f6354a07e86c7ad4c9e76b683fb2f", - "12": "32cb5f9bf7b8388f44ec41287e0762de15c5bdcc0f5f5596e72bbf56cc7450cb", - "13": "3af0ee7e88a948504f3b177d4145ccb4691fa8e76d0f09fadcb3666b903cb488", - "14": "8459bd20674e91bce53c82f6dcfdc47ea3cf5f58cdc3ab1b6b6d0ce1b2ec7275", - "15": "4af6694657669d7a6e50f3e90ac388d88bfecf9748a1f52a2bf58e179639ec8c", - "16": "d8bc53b746652a81b02d0e02bc198899b991fa10ef652079faed79774cf4197e", - "17": "78eb0bcaa63935b69928b7f266ba9a33c650e57021ec1b33fb8b8b9e2e00ea9c", - "18": "b46c13b3b8a18de871fd1a47e27814910503bd0a4adbf800dd4adf33a6e055ca", - "19": "466014afe6266009a2aeff7f58235f6312e64c239108e67f7548340f38a35da1", - "2": "ec98172e062d413a143e0e987421bb8cda6bf4d91e5cec5a201859f85284155f", - "20": "1311efa57f82300b0b8184ae10063d74342c961cfc45ad40f00283ccd3256759", - "21": "a288655c459c268557dc7dbe148d6105589486ef3feed146a5a01a39f13c8a0e", - "22": "53463bdbb2771fcc16a65fff886241248ef9209cd33933de75cebeaf88620edc", - "23": "d1c913a15235ff21085201b4dabcc96233dc192ea68a36d218fdcd25927b1b3f", - "24": "b608fa983c1db132a217aba4b905586aa4fe5ebd71fe33f84250289183b2359f", - "25": "c0a5be7d39c1472564838b8a1f102d9d12a71e9f4833ae9e09afd73d012f480a", - "26": "bb4faebc78eb45065c7a7320d4630d542b330b1f2bb4b48956b5ca2d738ee3a0", - "27": "2dd9c60d1dba62d1dba0402468528886c520441e3447900b690b17e0f6a84a3c", - "28": "5ba1652b701d867cae30d118eba0ecd46521c2447139cb036884822b3dde6ce7", - "29": "b3174eade54bf0e48739e71988d1f632b9db51fe11d78980601a5861cdfb2763", - "3": "f8888fdc0c659524e857dfd909bd9897f5c1cbbec38c6807f96489730a7bebbe", - "30": "f4aabbdeab2628b77a7d204f76c9b316913ae91568bd30cb264251cbb4144c70", - "31": "06f02bf44a1c2069f141a3da1ef7314041100f782a8fd25fef1dfad03dbbe5a3", - "4": "2bcac227c33864e46a72a5609b6f7f9107580f3ebbfd2232525499a23a710599", - "5": "507240ae9e57b775882908f4e0fb21725cdc2dec4ab29013271ec1a3878e0191", - "6": "d9cfd6ca31deca3da37b87080c06f3328803f7143076d74a356148e7921772c1", - "7": "70882681c478c0bc142ca64c5abd7f3af95ba503717531aa952f52485686f03d", - "8": "9c5a8b2eca7fc04b7f164579e512d4ce9ce5a5a764eac9c32e29c54af85125c7", - "9": "072e5956a22a8a34a3291cecad17e55c3e8247f8be7d18a4f4312f0273e32ee1" + "0": "31365624504bcda43d0d0cfafcfb6c259ca957a82f4f2d85ba284cd682695eb9", + "1": "92576fd0ba5f17da40f9bf1d5b0663ed7f4dbe74df8ad76c56aa9f20004442fa", + "10": "8b47e9d9475935d18a43207e4c63c3b7567b6c3acbc5e95f552e82b109d7b871", + "11": "cc7a4dafe5ffa895450d075efe0121e1322443e93d7a9ae4e628555efa821651", + "12": "15fad9a8a6822d7b89b3ffe2d9476bbcf26197768dddac2ed24795cb764245f5", + "13": "f7e5e0c800235f9d4b3e4a51002ac6e18738803cfa9d47199b21f2783513d7e5", + "14": "38fe6f4974281696d20fe972c2d08a4fdeed72d281467298f6a55050d95986c7", + "15": "b702e5d22882daf200562b83874c53582c9e368bd602d1285c5c57aa8270cbee", + "16": "c57fcdfd6725811330f121e4f1bb353d2e17ab6ecde278259d65069c5635533a", + "17": "e7c73cd1f14af9ab6cb47e2a794a70075f2a04d03f3c55dd0fca0913053e60be", + "18": "982db10de1ecd90ce2955f5d5770f25cc1ad65ebba9c6d3e6db57d904d138b5a", + "19": "976b7ec8661f6be82a927f4ff5a5e6568b607d6ef3a70829110fb60122fe1bcc", + "2": "30e88d793bf93069b76f4918ed5ada53447ade3f2f280bcaeb0ed9b23ef22cc7", + "20": "dce7bab41ea08f6bc03fadef0a72bc4a0b275eb83017abac3d770f129311c513", + "21": "b8f928a4c87449119ee5b6b41571d88c15bc4a905337f16cf85c1366d28fc336", + "22": "b3c0fc075f8d696ae0c5f38866036b120d6e8ed2dffcdb92040f64b200a4775d", + "23": "b25a2f9e0c18ef0d21a075b7f461ff311eb4f78cc735e9bd3ab0f41c7c2a5d0b", + "24": "da4cbca6de5d4f50624e59cee13f8dabc8d12399e94249171c1e55bac2e32fc7", + "25": "cb699355b1d3217ff7455b741ae2466bc844a4d0318b4ed8d295b4eeba82ab3a", + "26": "145b26679e5dec3ecef20e7296a0d16e13089a705018fa8d7c334c1bb243523b", + "27": "31365624504bcda43d0d0cfafcfb6c259ca957a82f4f2d85ba284cd682695eb9", + "28": "b4487fddfc3cfd2584da70879eb6ae5258f3681f95285b3e29ce07bf7b0dda62", + "29": "0533f9cee5dfa32f658a3039a2c5ba69238eaead0e1107fedba4e4fe92e23bef", + "3": "7f59dca66aa6119cc19a8268751e54d86c60ae8554bb91eb474760ba34f4e2f6", + "30": "f6fb24276751812e4db8f7ba938ecfd73219ba8ba995176a4285888cfeef25da", + "31": "0ec993d53c1a5ca42c81b5ca2889ef9f8f51f90e0a211151eef8b6911437663e", + "4": "c230ddb11e58111e9f7fdb6d75024350207d9ed77a19baac2ba23b8a6e51ecb6", + "5": "43254a2c4eed8a972f2b17ffba1087931728145c610318b644570c3e6a44c512", + "6": "ae1644dc03ec4578e8a45faba66d95d5fe893652c497a50df3fa6bb6799bfcee", + "7": "0fbfd451acd5479fd7507b150f773fd05412bd53b1c830c7afda0205808053f1", + "8": "03115c4463d615e408c9635d4055112f158771425ee64edfa03632f3cd8b5c3d", + "9": "359aae0c9f49985c77b777b0d5afe766c3b836e671fd525ae3efa00e80d84da5" }, "barrier_fence": { "0": "3d4810bee983ee04bc7b34ac2f2bd9767cc94ee2b217c0017821f9cccaf5ec66", @@ -1766,46 +1766,46 @@ "9": "7866c9eb6d414e0edf3678a477c387e2742057eed6a96616544f29057d609d0a" }, "attention_unified": { - "0": "64362d51c5c69a238b6ef0315e46e670e8fbe98f2aab1695d4bbe43bd1365014", - "1": "27447752f7aa56093807880e18639d7f52c888a64c4916b831183925b6a035e2", - "2": "627ef0db51fa9ca96ab05dd5394fbb038e91ea5f30f9392fb04874bc87ecb346", - "3": "170c01ebeb641dca63b1f5fea4fabc8baf65f36e747ad244438596a345de2847", - "4": "aa793bb3d984f1268d321d8920b84336c2ac0de18e0928943bfc37b22a5df8b2", - "5": "4c4d2131287fadfd01b2e555cf28dc1272e2a943d5bfcd0fd35f6525392faa75" + "0": "1c3cfc68c551423819d493f80b861e0073a74f6b50d34a0818963691d0ffa420", + "1": "77fe5a7e29ce4f4d6f9b91c43f18e1eebcae586134664298015045f739fb7a52", + "2": "66b75c98be92b50bf1f4dde89157c2557ddcd29902bb94fb6a3b4e4ade77e8a8", + "3": "d3e453afa76ad7d2198c1edbe0d7be0fc681f7a59ef0a9c61080b8f52be56b16", + "4": "026059224bd3587a5c90c94aa2f4ebbc19410fb8e1c0c15911d1259a86a7ca38", + "5": "a9af175dd5bfbb8ee5b7d0c6c4ffdee063a90846cbcf2d2b2e05da552aca0deb" }, "attention_unified_stress": { - "0": "64362d51c5c69a238b6ef0315e46e670e8fbe98f2aab1695d4bbe43bd1365014", - "1": "27447752f7aa56093807880e18639d7f52c888a64c4916b831183925b6a035e2", - "10": "b05b137dfff261c4eee5eeb851a595b52e98cc9e798009b77f56abe889edc5a4", - "11": "982af8227c8e1882d3df85b4445eda5ea40e8c7af591d69a7343da74483d236a", - "12": "28d12263ccf69764495db43c29c0d3007530f2b7d7c3f1df9169b31691397ced", - "13": "9e794d6fed08e3c08d6601fdbb0b55e863509127967b3ec68ce5bfd20029f86d", - "14": "41d04a277cfbd2b14b4fa0f54590d611fa23a88e0a763f39bb010a77166742cf", - "15": "a48c6d5c83c5a2c79b5a731cc74beac4938537a6f82e119d0da8dc54ddd57de3", - "16": "a7029785098f56ce632d070b7454f92ca024435507a5aa17fd31468b230977f2", - "17": "bdefa081d85e49cbbc892685dd3d967bfe6d0f1c5ecc6fad3577cf00e291e0db", - "18": "ee866a91ab399aa9fee011bc14d423d1cb44341f2e71e498f633c2f7586919b2", - "19": "05ba90538ed03e1358ff92c321a1145c5b864c56dbb2b291e2759c12f99957b7", - "2": "627ef0db51fa9ca96ab05dd5394fbb038e91ea5f30f9392fb04874bc87ecb346", - "20": "b9b6662dce40d70a0dd5d70432f012b74b3a5bd172ed11685211a1097f9cdac7", - "21": "657c285a690bb173062a27abab45f734877f6635e53df357875323d2110cc9a2", - "22": "cd7476bf9e693d7bdaad4aa476ba7aabdda404143ca161e038b287ea8e53cef5", - "23": "f8f1b9df83ef485df6e4c082cd0af552d57ae6b7cb3646aaa2559605cb1ad010", - "24": "1b36b8951c3db1d1a4fa5c6c6858d1f319e4fd1b42a9c3c158e2d86128f15496", - "25": "c0b5900afe2b0512a7959fe775cc3301dd3d41cd6be6565a1b3ebc1a6d9e1c95", - "26": "ca52d4910739e1178dbf9999ce6a2b356c49fb9e3c314b92a97688290a4b4900", - "27": "64362d51c5c69a238b6ef0315e46e670e8fbe98f2aab1695d4bbe43bd1365014", - "28": "cc45552cfbb86dab6b8761c19f547683fd3b3d68dc3a4bfb977c87783ac10adc", - "29": "8dc49b2b40620742b65df33124897a82238ad4fcb46ca5298a7b364277133ee8", - "3": "170c01ebeb641dca63b1f5fea4fabc8baf65f36e747ad244438596a345de2847", - "30": "e82e39f0ba7f9cd4bb3083442967a8e9e985f88d2135617888644c7bcb5dd34b", - "31": "a54317710e58234ad790f887194187eb525d1dcc48ec610cfb381b70e9946d72", - "4": "aa793bb3d984f1268d321d8920b84336c2ac0de18e0928943bfc37b22a5df8b2", - "5": "4c4d2131287fadfd01b2e555cf28dc1272e2a943d5bfcd0fd35f6525392faa75", - "6": "e4898af527d70f177992e2f65e9abbf3300f6153083d352d8f7947952ccd47cd", - "7": "85deb6f4462b83d1b554223d09ef06463d8e1997618884ff7eae22da7207c875", - "8": "41ece1db50fb5ab6467f99d0a12304b3c3cb248afbe55a25b88b609f7eb1980c", - "9": "20adb829c57d108c13a3177a66d1d8242431fbffb6472ca1c6309856cae089b9" + "0": "1c3cfc68c551423819d493f80b861e0073a74f6b50d34a0818963691d0ffa420", + "1": "77fe5a7e29ce4f4d6f9b91c43f18e1eebcae586134664298015045f739fb7a52", + "10": "057dda7dcf2370a9180dfb2f24c501c89fe43637db1b633dc89a1210645104e9", + "11": "0ab771335f7c25bc1eaa80bd0d5f7bc423b20c6f9a5aeab53794a6bf007c5775", + "12": "bfc29c5dfcb8e1d0d6ce667b9d21c8eb1d0e401a152853bb8b8c832279d876b2", + "13": "34c5b7611e78cdd9061047c651e56fdd7f5b35583272539df164005dc5393441", + "14": "fe64b6559a72913e18242de5fee9dc4a2644f6873186a0d3c048dae0893ec272", + "15": "620b805dc704b9096d605ccad38d69464ac673dff44f243af879585bf7f84e03", + "16": "2f88ed7b12f45cee48e7658d4fc8d9110fa1ce18eb2642d768d479e361e81a03", + "17": "497d259cdd4e3f5d8ce95d13528e61323cea17ba57215bd3af3add509ae3d13d", + "18": "1b0fcd6e3fd01c78e93ea4b2954d085bbc3c9c0af2bb127af52aed1182c2cc95", + "19": "9540eae9743e166c86718b759bf9b93ebb3e635a362b2fe9fe5ff14acffeba26", + "2": "66b75c98be92b50bf1f4dde89157c2557ddcd29902bb94fb6a3b4e4ade77e8a8", + "20": "14c76d8f655f5c1cafe58e6d1f7f842ef9e6ec18cbfd7cda02ab98228a94907b", + "21": "f1d6ed718c431ce725cc087945a1992751aec072965ef530724a14a29889e1d4", + "22": "75d345bf65119baab803528474b3553407fd806a630bd4a7253a81eb2a354f5f", + "23": "daed7d9b41303a26e8893c803a7b4c9b0100126a21924b27fa71d6629baeba42", + "24": "99364e51e92abc2644e7276f5b538054cd8c593af7b23a15ddc1de3e337b7046", + "25": "a3ab049f136e18484f87fe79ff069f73aae17bb2b2c9a6e6e175df74ba7fb81a", + "26": "60736a8eaa3ec285141feeb30c195f22a82a32bb3226797e9e4d9b45b8c094d1", + "27": "1c3cfc68c551423819d493f80b861e0073a74f6b50d34a0818963691d0ffa420", + "28": "ddb0f88cb28a05e320f2c34e99e5da694240d30b0a6f9dc73e08b26a7880bf20", + "29": "23f8babd384e973c944434b683b295a1e44ef62a292499dbe41376f5a8bcec12", + "3": "d3e453afa76ad7d2198c1edbe0d7be0fc681f7a59ef0a9c61080b8f52be56b16", + "30": "3cab5855ff4bdde156554e3dc57c15d1f1639405ac20401b5b03273126c298d7", + "31": "8593c2303ce1e802d6c362ee100e48d45901f9b04ce636e8bf56ae011705ccde", + "4": "026059224bd3587a5c90c94aa2f4ebbc19410fb8e1c0c15911d1259a86a7ca38", + "5": "a9af175dd5bfbb8ee5b7d0c6c4ffdee063a90846cbcf2d2b2e05da552aca0deb", + "6": "3015b1b8a562946706637c5df13e6706c7f97246539979810cd153ea0daff69e", + "7": "789007c3a14934ac5533dedcf6b617d5575054a0121a159ed26b54fae9b4e975", + "8": "79af89461585dab768b203c8801c850428a576623b4a01935957375950e02ecd", + "9": "55718bd1f65c19b033aba009d38104999d6ed6c4f9a2caaa3b42e0466f3f0e66" }, "barrier_fence": { "0": "84cfa4eaf14b6c0c8a1861d4bbffc69958c4710f97eaf2b94a9237d3e375d2ed", @@ -3411,46 +3411,46 @@ "9": "43cbf9fc03696e3fcf0a132e31a769d92cd48267eaea3526e2cc7c46c35225da" }, "attention_unified": { - "0": "63b0e18b1115541aa4761eefe2dffc2e7ce9f00563e55079ddaffe5c7f61121f", - "1": "6c7f2c3f393ed7aa55ef217842463d5ea82a3f13e568dd1dc9aebef4415ce29a", - "2": "a597a879541f23a43f1c5545b64437fe69e369cfbd8d84e065ed618175d62e53", - "3": "99c2955fda52f59ec0b563cd1df6e5b6ad02f0198a19b30765fda4f04a783cbc", - "4": "46fc238f5534033fbd09e7b9daf0b399f885d9eeeb7ec5c39081454487127bc7", - "5": "a1840a581b8124abdfa7fe996cc6b637405939604a45e4b7725e17361014e66c" + "0": "621b87ad838637c25d93b776cbf09f8353951fc953c1445e32e1ff40cf4bf37c", + "1": "9d250722f0adb3180fe9fb61b1e1813555f596a015b747ab79b88b6f2aad36a8", + "2": "ca34c7c9daf99525f98f5f92de56a2f5be745441e85c7edad1d3378c9c06668a", + "3": "db8560425d4c6243491b352371468b4e65385deecbb48ec56bf978ef95f6b7f6", + "4": "88ff1c2b6c354ca1df1968ae24e51e0b96af797010d0ac66dccf13161e64f84e", + "5": "7944b498b84e783ebe6595554ea7477ef54a0d056a77e284150d880e8b3d7379" }, "attention_unified_stress": { - "0": "63b0e18b1115541aa4761eefe2dffc2e7ce9f00563e55079ddaffe5c7f61121f", - "1": "6c7f2c3f393ed7aa55ef217842463d5ea82a3f13e568dd1dc9aebef4415ce29a", - "10": "1c1adcf989f3f67f7f8e43f04287548e4a455981da11c0ec00aaeb669e5268d6", - "11": "720f3f304d96b386ca8de6edd54d46fbe3f1255b75519c4ed84a29efa7b18a7d", - "12": "15d150b6adc02132f2f2643fd50cd9d6b5404faf46f57bc26ce32a1509d96879", - "13": "41c065a6b5d93bf5fea38729c414e6b89f62833982b0bd07c32a839f7c1f2ff6", - "14": "b0fded81e70015374613eb01764d124a8ce0f30646476940ea8dad781115696c", - "15": "49bdfc467a6c49e9155475772bfb90ea74f69e9651e356b0ac852d87aaa67d4c", - "16": "c231377d6682233dea711f323a19b945a7be54e8bb80d4d132e51aef3bac9f5f", - "17": "61e864d477fc0fe7bf8ef8292e27097f39658c598ea39f2c62a9a0a045245a2a", - "18": "1cb2f0333fb0f0a7a6c8a90198bccf8d5dc6c75723af8f36344babc7667cc4f8", - "19": "799b048d3043ab1a26b564c3cf82a2b3a63f2c3dd15c61e997abb7dfa84d0107", - "2": "a597a879541f23a43f1c5545b64437fe69e369cfbd8d84e065ed618175d62e53", - "20": "c2ae484227fa7bb8dfae26dd5014f30a024caed64eff4f84d90ef5ef18c5709b", - "21": "e0ecb473a332344288b6fd94977c093cfc189735bd20ed3d0f52132e0424898a", - "22": "df158209ce45647964ad114a22eff05fa3a6ebe8331fc587d0ea7571408ebbab", - "23": "ad7eac5e70ae41231e760bcf52d3899f1275bc573fa15e46bff9496c5cfc1191", - "24": "a459a8268208f3e20bd0783a7eb9a0bb84c8add8fd26876c58433e21896dc1ac", - "25": "26fe9ded321476d62d0588a31d65f26f5ea2e916fd2127a401bd2d7a798c4803", - "26": "1d2c0212b12236f9e64a022877b96fb2db376d0fe6f715e35efc5fde3c1875be", - "27": "63b0e18b1115541aa4761eefe2dffc2e7ce9f00563e55079ddaffe5c7f61121f", - "28": "4846529cdb6cb1de94d4a03d5d61de4bf8e76167d9883afc36de15bdde8543a8", - "29": "da53673153ebf1f5c3b5a88e2a388f40f03b1c07777a1ff1fcec9ae003ed60fa", - "3": "99c2955fda52f59ec0b563cd1df6e5b6ad02f0198a19b30765fda4f04a783cbc", - "30": "c4b85baa5ad7804b0c06237a9dc87d4e05e6e9e2802274637630de9bb8bf24f6", - "31": "70b18e643a5b0a33fd4b42f0b0cb402c2b7d63a88e653779b88410cd8d2072d4", - "4": "46fc238f5534033fbd09e7b9daf0b399f885d9eeeb7ec5c39081454487127bc7", - "5": "a1840a581b8124abdfa7fe996cc6b637405939604a45e4b7725e17361014e66c", - "6": "d662718e9c3e7569f96269d423b4b308781e5ae38bbcbd35d0203c70e5dbc6d2", - "7": "6174c69e7692759da8be5930262f2eca74827475d341e498a8d1eac0a00ce1f2", - "8": "193c8ce9e0c72e0ce1ce46a64aeaf2d6b83fb70e68d4684f9623d47be30666ad", - "9": "37bd385bd44e1c9178ff2198e0ccedec5e0e7975ce17df77f8936374139cf040" + "0": "621b87ad838637c25d93b776cbf09f8353951fc953c1445e32e1ff40cf4bf37c", + "1": "9d250722f0adb3180fe9fb61b1e1813555f596a015b747ab79b88b6f2aad36a8", + "10": "249b587a9ae2da36ce1a5055892c8e15be09fa965d6059962093e97af6701776", + "11": "17369247188acc3d770d98af31fc14254263e92b59df8065a5f70091cb771a07", + "12": "fadfa9559417fd7297328d5e1db64543d5547d059207a1ba1bf0519564a6e405", + "13": "6ad089df8f09baa56dc8b849b7180a13014d2cc62ffb06a597c6c734bef17058", + "14": "7ab12e8cad3f47e77f43bd1c7bfcb1a71c33a2c5d9c1810a956004b0d84aa429", + "15": "8aab65afb70f5bc449767ead8eddb02c2282a2aa2225d865f042a7ae7a2a0701", + "16": "82ed45992370e55c4201d02626832c0c2a3d852ef7e81c72c0e92bebccf08ec0", + "17": "8aa1f4d358b0918613e7e3913ad7c25608e7b451870551fa3ddad8d2fbacb130", + "18": "5207d667a03eaaf1ae83423c0c3ab03eb960b129c97194825a5f4de51d5775f3", + "19": "7feffe317a47453b4b07485df1c99dc247ff2d0434d4fd1d76e86aa90caf9666", + "2": "ca34c7c9daf99525f98f5f92de56a2f5be745441e85c7edad1d3378c9c06668a", + "20": "098b13bd2e5efff60f6f1d17430727bf718cea40363a12cb14a1d0fbe7ff9d71", + "21": "3244c914c1181b7749ce9149dc851a8e2ebad4c1c693accc0d5337eb5fd074d3", + "22": "4fba3049b3f96bbc74a867d87687012aa74d3bf0e32973dc0cc5d4c34eb5a4b9", + "23": "feb91d3f21e09701f044d9b4c9e7b8340d44e91b44354395bc7113cce36407d0", + "24": "457cf85374d32c4531f6bd8d41de73e994b99b2c18c5f043c917025a6d81e5e2", + "25": "b2d9e41107b1b88fcf90fb56dfd5f9656b6618a014bcff1af43d34e301a132db", + "26": "0582ba3534cabd9fa29df3cf014cce41b82d878fe35915a2cef4f5bb10418182", + "27": "621b87ad838637c25d93b776cbf09f8353951fc953c1445e32e1ff40cf4bf37c", + "28": "ba84f869b0ac00db17d66eac97ce329ba1d351fbf565a7ef9b9a76d34d2705d1", + "29": "878ae0b876381442d4a0e10cddbad8f35e6f9db07a0054648eea897ff2c68063", + "3": "db8560425d4c6243491b352371468b4e65385deecbb48ec56bf978ef95f6b7f6", + "30": "1645b9c2ffe51a5fd57894abeccf7135d47d42e1002c2752e1ac8618b1ffdb1c", + "31": "e9bb038e50b8338a83b4311d6b81e1f66ebdc5673bdd1253608b3bd4a6e5fc1f", + "4": "88ff1c2b6c354ca1df1968ae24e51e0b96af797010d0ac66dccf13161e64f84e", + "5": "7944b498b84e783ebe6595554ea7477ef54a0d056a77e284150d880e8b3d7379", + "6": "b0cc6fef08547f21207faee8db31a3e0c2a520af46e18ad8dfebf7ab489b9b98", + "7": "cfe74d23ce67785935ce2682f35f8451c8ef6c7e672970a0280b86e523a2d94f", + "8": "755dcc1f76e4e9c92110731810875b07a9bd8631d823d91ca530167ec3ed3050", + "9": "8c7db97abda6d315164e84c312ab5118719f628377c1e46c3155bead674bbf0a" }, "barrier_fence": { "0": "0681ccf5b97842f61b0ff9fed898b4a696cd0f7df2b2bc3cd5eb41d13e9cc044", @@ -5056,46 +5056,46 @@ "9": "cb2426e5ff5ce00a3094ea4217df160031006f22cfe058568e4c959f1ecff086" }, "attention_unified": { - "0": "de5cde9088dd20c8d6c5e2c7c792d6a9d4a8b42cc423a6cdf26f5779c1b12ba6", - "1": "f99195661a2ff3b53c46ab521e7ac39f562faa18822cfa390cd0384f3ebcfda3", - "2": "f0ee3ebf4e8a939851d5ebc43ab69ccad8077a134820b472ab9e867fbf546721", - "3": "3874bf5a465c5a34b8bcd8a6d0938367541d42b25e4beec6b6e49f765b909a60", - "4": "516da7e699092a229c993d3853dd3f96c73c9856c086ef86c89a5f5da2d339ca", - "5": "18bc3d432dc491521d17acb3c8755d1888a94141cb3e11d842a0bd48c1917770" + "0": "24e760857f19125cb3f1231179d12ccaf3cd4147c9dd11b43839a25362bd7589", + "1": "f563ab8ccb4778a72e05506b3b12b32d4142c7e7fba1733128c50c139cbaace4", + "2": "4ade93e9211772dd24489a231046d4c81e3efb1a7af35e8a01ddde469d1636bb", + "3": "eb0d167d88528507471a33788130052506306f39b47f432b72671db087591e2c", + "4": "9bb87b1079352a6e2f34fe2c0d95767a93dd597b579868483a5b71cd424a477a", + "5": "1f5b9dec2bee8b2ef90f91e1c58de84a561eaa724b280b0b359fe34116aabcfe" }, "attention_unified_stress": { - "0": "de5cde9088dd20c8d6c5e2c7c792d6a9d4a8b42cc423a6cdf26f5779c1b12ba6", - "1": "f99195661a2ff3b53c46ab521e7ac39f562faa18822cfa390cd0384f3ebcfda3", - "10": "7fc12a187cdf0586c4e6ba2160d90009cedc2c6c3b534babaf14af25998852f8", - "11": "60e391c8c42c4112ae41b67104ac9595b923492081f9f555cb358a1bdc48c50c", - "12": "b7b11c5ef9153c71bd3d55de065f399e393e7e0f59cdeac35f6f37d55624ead4", - "13": "0987f3ba80d0938b994b3fe52960a3ad63ca24133a78edbc374277e066c747ab", - "14": "b4f61f8b093486dc28e9f839102d5684a0562e136f7ddd8526f8d81d9541ed0e", - "15": "759ed646471f703eadd51c65ac535829bea52e1132260247fffa5839a9a796ca", - "16": "7eec1407f76792a3f7e91e626c5938d9685e228ed3ab961f4a3453140a6a443c", - "17": "a5cf7f1b67898ddec192e39114af82d28137f15403c3d80bc6e29ebfff805b0e", - "18": "e8d32e69908a40423e4b98bf24380ea7f67f1328fdaac6c46ca7877ecceee7da", - "19": "f8c26ab84b87576cd4ec91537889207afb0cf45513707bf09e0f0dede5109e4a", - "2": "f0ee3ebf4e8a939851d5ebc43ab69ccad8077a134820b472ab9e867fbf546721", - "20": "2140a93a660f467d4116ef3b3b83d6804b7739c1667d63f3b2446a90496fe93d", - "21": "7dc48535c5d3fc65d7e0403ed96232675bf7bc0f454155308bb4f7ba97f43499", - "22": "14a427161eedb72e80ae6fd8f320368e851fc0a75b31857e7381d3f1aa8d250f", - "23": "a2484ecd0c43fe6932ad5bc0b31bb19f872c3020247fc7ea0d27c89a9b4489b9", - "24": "1a7a79fd8fe13d9990119777c228e11bfea46269d9c039e125898fe70c35a7d0", - "25": "cdcca635d8bd968dc740566e748c7c5345c70ab6b66d3cca6b967d071b5fbb8c", - "26": "dfa69ab3eda7acd0cc589a6f533f574ac1c306913b9feaad2de0a7369b8667ec", - "27": "de5cde9088dd20c8d6c5e2c7c792d6a9d4a8b42cc423a6cdf26f5779c1b12ba6", - "28": "ef7707f598907038871634b1eedd15b32918404fda6f5caf90d0da6795551e8a", - "29": "05af36dfd05cc9c62fe3b1a20e93e952a3096f1e1c9563295c5b396182adf0e8", - "3": "3874bf5a465c5a34b8bcd8a6d0938367541d42b25e4beec6b6e49f765b909a60", - "30": "93cc1be104c5fcf9662eaaf92718f6c8bbf88dd5c52335baa2930dc42961b96f", - "31": "a7c46650acc8639d1eed8fa4972614f385374500928610989e518196f05c35af", - "4": "516da7e699092a229c993d3853dd3f96c73c9856c086ef86c89a5f5da2d339ca", - "5": "18bc3d432dc491521d17acb3c8755d1888a94141cb3e11d842a0bd48c1917770", - "6": "abdb5d661d6d01af4a8f15bf3679a91ec1d7a718007e820935c39fea50163e7a", - "7": "c095e489d273198c1a59545333e07a629f06af5968c645b7c3d934af46ba0c20", - "8": "d50ef4453d39a454a67d6173852e693d5f8b1a92ec81f75e0f1b2806c5877791", - "9": "54e63bdbf0a3e34496b94278e052ad7dd69cc78e6836220f924ec65843b5b806" + "0": "24e760857f19125cb3f1231179d12ccaf3cd4147c9dd11b43839a25362bd7589", + "1": "f563ab8ccb4778a72e05506b3b12b32d4142c7e7fba1733128c50c139cbaace4", + "10": "365f5450456182cb954c328911d0d9bb04352016362453ee59485315b026e036", + "11": "67abfd4e59bba9f9e81c571e4b109bd318e8053402c18421d90a1c0a4462fb7d", + "12": "ebb85bfee8ae7aecff844f304fc8e3427d1916b6f03d9cc05b29421c724da114", + "13": "863f4366816cfa2bd96b851313c3063249c236fe1d78984f5615768b8998b763", + "14": "ee2867678a52a2c0266f969fa5626339cacfdbf6905075a41a782ab3577ed471", + "15": "f05beb49523f5d7c97548c033eff996153ae03533782ce45a53f0f2e32611583", + "16": "6c4881eee07b859ebf45fc0f5a275bc73d270cd96fc247ae221c205a21cc08a0", + "17": "4dea3a909cfd93b638cbd7a9bd2bdf27286abe90ca04a9a4d6b4cc6755c0541c", + "18": "a3df309d5ebf90b95f0cbd50cdd40a6e49c3a509fdc90e665b9ac6d5dc1722b4", + "19": "1e1cfb449e715befe0236f9d2c34c99d9caaa5c84826e04395d646e74522dca5", + "2": "4ade93e9211772dd24489a231046d4c81e3efb1a7af35e8a01ddde469d1636bb", + "20": "aaad442b85ed576e22a8e6b8f77d27e626b61c2166425f3ba6bb58e4cff2057d", + "21": "c346403c584fd99cc9d1d3b28cab373ca19a0eff238a24f7c9ad9e2679d61c61", + "22": "239e6c04a2a53d46e9586b332f7064644404099dac068eb80e64fbb13a92ec4d", + "23": "b2d6f355b4930417fe1e4c384cdfc7b36cabe5088d0b26ad8ee6f34b69a73983", + "24": "ea46244f7f0623a632a2243fff81f9129c1d6ec7ebf9a93a8318db03b2f9b907", + "25": "615e7a2b4724db9443b3c1bf14d7ded0e8db366006437f93f3dbae92e07f53e2", + "26": "1f3e729f3d74efe8c2b3f99709e427e7c7c3097dffb441926a27681389fc6331", + "27": "24e760857f19125cb3f1231179d12ccaf3cd4147c9dd11b43839a25362bd7589", + "28": "34409f552ad8071902ebfe84d2f1898725c2cffed7b70013b770797ad037e243", + "29": "ac3a5e2f897a8a74ac5a9fe87004f14607a2a8b1aff0ae4dbae226dbb8cdf850", + "3": "eb0d167d88528507471a33788130052506306f39b47f432b72671db087591e2c", + "30": "2f8bb9edc15c71cf628284be3296b34f68a80c90d5471cdec76f3322047e856c", + "31": "6585b4b31bea3542cb8f34eead1733c09563b76264794dacab9ea63b588ba88d", + "4": "9bb87b1079352a6e2f34fe2c0d95767a93dd597b579868483a5b71cd424a477a", + "5": "1f5b9dec2bee8b2ef90f91e1c58de84a561eaa724b280b0b359fe34116aabcfe", + "6": "e2762b5fa7a1aff2b071072c9c3c122d841c9c6e2c54c4d6abcc8e451713c09f", + "7": "fa2ddafef2cd6763da14bf4ad3d0f06b40c5a4927ff44f66e4907c3d8f3a1d7f", + "8": "4b83508bd5f0f93fdcb5d05433c5d3edb6c97f61783b580a3303688b82df9334", + "9": "6965c02a06c21b47903cccd41ee46de9780d24890ff716f25e67fb8d51ac660c" }, "barrier_fence": { "0": "0af35781f2467b4cfa91824a3ea0fce54703b0af5264db42a3c27d8ea1cfd4d1",