From bd661f637e7b307b7850b60764fee987cebbcf8b Mon Sep 17 00:00:00 2001 From: David Langbehn <47395968+dlangbe@users.noreply.github.com> Date: Tue, 3 Nov 2026 10:53:13 -0600 Subject: [PATCH 1/8] perf(rocke): make gfx942 tiled attention heads, window and batch runtime (#14530) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit JIRA ID : AICK-2676 The gfx942 tiled attention kernels (2D, 4-warp GQA, 3D split-KV and reduce) now read the head counts and the sliding-window size as kernargs and derive the `cu_q` search trip count from `num_seqs`. One compiled binary therefore serves every head pair at a given GQA ratio, every window size and every batch. This is the kernel-side readiness hipDNN AOT packaging needs: a packaged catalog ships one binary per kernel symbol and cannot enumerate head counts, window sizes or batch sizes. Follow-up to #13372 and #13889, which removed raw shape fields from the tiled and scalar cache keys. Builds on #14216 and #14349, merged from develop. The gfx950 counterpart (AICK-2675) reuses the plumbing added here. ## Risk Assessment Risk 4. Every gfx942 tiled kernel changes its argument list and emitted IR, the tiled cache keys change, and the paged K/V descriptor gains runtime strides in both engines. That is a performance-critical path. Mitigations: byte identity is GREEN at llvm20, llvm22 and llvm23 for every family; gfx950 and gfx1250 tiled IR is unchanged (golden); and the gfx942 numerics and an A/B against develop were run on hardware (see Testing). ## ASIC Coverage - **gfx942:** all kernel changes. Verified manually on gfx942: numeric tests, the existing gfx942 GPU suites, and a develop-vs-branch A/B. Superbuild and TheRock CI do not run these GPU numerics on gfx942. - **gfx950 / gfx1250:** host-only changes (signature helpers, cache keys, launch-cache keys). Their tiled spec classes declare no runtime kernargs, so their IR and goldens do not move. PR CI is sufficient. ## Testing Summary - **Byte identity:** full `check_byte_identity.py` GREEN at llvm20, llvm22 and llvm23. The gfx942 2D, 2D stress, 3D and reduce emit configs cover windowed, i64 and sinks specs. - **Goldens:** the representative-IR golden moves only for the gfx942 tiled cases (2D, 3D, reduce and both 4-warp kernels); the strided-KV library golden only for its gfx942 configs. Every gfx950 and gfx1250 entry is unchanged. - **Cache keys (CPU):** - shapes that differ only in head pair (same GQA ratio), window size or batch bucket share one key and lower to byte-identical IR at every flavor, on 2D, 4-warp D256, the 4-warp D128 fold and 3D; - the invariant sweep holds; - gfx950 and gfx1250 still split on the batch bucket. - **Launch signature (CPU):** every production builder route and three gfx942 tuning candidates are checked: the kernel's params equal its launch signature by name and type. - **One binary, many shapes (gfx942 GPU):** `test_attention_gfx942_runtime_shape_numeric.py` runs pairs of shapes that share a key back to back through `run_unified_attention_torch`. The second launch must compile nothing, and both outputs must match an fp32 paged reference (max_abs ≤ 4e-2; bf16/fp16 inputs, randn) with unused block-table slots on a poison block. - Pairs cover 2D default, windowed 2D, D256 4-warp, the D128 window fold, 3D decode and two-token 3D. One pair changes only the head counts at identical token counts. - Checked against mutants: baking the window back in fails the windowed pair (max_abs 0.88); baking the kv head count into the K/V stride faults on the device. - **Existing GPU suites (gfx942):** 69 attention, MLA and strided-KV library test files. Results match develop; the remaining failures fail identically on develop (MLA comgr environment errors, strided-KV and selector-parity cases). **A/B, develop vs this branch (gfx942, 304 CUs)** - **Workload:** 116 problems through `run_unified_attention_torch` with production routing (`backend="auto"`): 98 on the 2D path, 18 on 3D. Cohorts cover d64/d128/d256, bf16/fp16 and fp8 K/V, GQA ratios 1, 4 and 8, sinks, sliding window, the 4-warp D256 and fold paths, and a 64/16 head pair. - **Method:** HIP-event mean over 50 launches after 10 warmup, median over 3 rounds. Arm order alternates per round, and each arm runs in a fresh process each round. - **Toolchain:** torch 2.13.0+rocm10.0.0, LLVM 23 from the ROCm 10 SDK wheels, Python lowering. - **Baseline:** develop at `f37a7056c791`. - **Numeric gate:** every row of both arms is within 4e-2 of an fp32 reference; worst max_abs is 1.24e-2 in both arms. - **Results (branch / develop):** | Measure | Result | |---|---| | Compiled binaries per process | 18 vs 54 | | Total first-launch time | 0.37× | | Steady state, geometric mean over 116 problems | 0.994× (range 0.87–1.06×) | | Per cohort | within 0.98–1.01×, except the two below | | d64 GQA-8 block-32 **sinks** (2 problems, 2D) | 1.05–1.06× | | d64 GQA-8 **fp8 K/V** decode (4 problems, 3D) | 1.01–1.03× | - **Sinks caveat, isolated:** - Timing each commit puts the whole sinks cost in the head-count commit, and only on the narrow 2D kernel. - A timing-only variant with the K/V token stride restored to a constant removes it; restoring the query head count changes nothing. - The per-copy K/V offset now multiplies the token index by the runtime stride inside the KV loop, where the constant stride used to compile to a shift. - Accepted here as a known cost. Hoisting the in-page offsets out of the loop belongs with the planned per-page i32 K/V loader rework. **AOT dry run:** a draft engine manifest pinned to one gfx942 tuning candidate (`transposed_x8`, 4 warps, tile 4× block) expands with `registry_enumerator.py --plan-space`: 24 cells (6 head configurations × bf16/fp16 × block 16/32), production sweep, 20 varying knobs. Result: - 324 instances (within the default per-arch budget of 1000); - 162 distinct kernels on this branch against 324 on develop, because head pairs at one GQA ratio now build the same kernel. ## Testing Checklist - [x] Byte identity - `ROCKE_LLVM_FLAVOR=llvm{20,22,23} python tools/check_byte_identity.py` - Status: Passed - [x] Representative-IR golden - `python tests/instances/rocke_ir_parity_harness.py --check tests/golden/rocke_representative_ir_sha256.json` - Status: Passed - [x] Library host tests (attention, MLA, strided KV, dispatch, layering), GPU hidden - `python -m pytest ` per file - Status: Passed (3 files fail identically on develop) - [x] Platform dynamic-descriptor tests - `python -m pytest tests/core/test_transforms_dynamic.py` - Status: Passed - [x] gfx942 GPU: new runtime-shape numeric test - `python -m pytest tests/test_attention_gfx942_runtime_shape_numeric.py` - Status: Passed (14/14) - [x] gfx942 GPU: existing attention/MLA/strided-KV library suites (69 files) - Status: Passed, matching develop - [x] gfx942 GPU: develop-vs-branch A/B, 116 problems, 3 rounds - Status: Passed (numeric gate), caveats above ## Technical Changes - **Runtime-kernarg plumbing.** Tiled spec classes declare `RUNTIME_KERNARG_FIELDS`. `_tiled_2d_signature` / `_tiled_3d_signatures` append those fields to the launch signature, and `_runtime_kernarg_values` packs them. Every gfx942 launch site uses the arch-aware signature. The 4-warp GQA kernels now declare `qq_bias_stride_0`, which the shared 2D signature already passed. - **Search trip count.** `binary_search_seq_idx_trip_count` computes `bit_length(num_seqs)` from the `num_seqs` kernarg. It's exact below 2^24 sequences; past that it adds at most one no-op iteration. Specs declare `RUNTIME_SEARCH_ITERS` and the keys drop the batch bucket. The existing helper and its other callers emit the same IR. - **Sliding window.** The size is a kernarg. Whether a window applies still selects structure, so the name token is `sw` and the keys carry `window > 0`. - **Head counts.** `num_query_heads` and `num_kv_heads` are kernargs. The Q/O, segment-output and paged K/V strides are runtime values through `DynamicTensorDescriptor`, which now takes mixed int/Value strides plus runtime strides in `offset_i64_split` (both engines) and `offset_i64` (Python). Names carry `g`, and the keys carry the GQA ratio. Every selector that reads absolute head counts feeds an output that's already keyed. - **Launch caches.** `_shape_key` (2D grid, 3D grids and workspace) now carries both head counts, since one binary now serves several head pairs. - **i64 name token.** The gfx942 2D name carries `i64` when `use_i64_kv_addr` is set; i32 and i64 builds used to share a symbol. - **C++ twins.** The 2D, 3D and reduce twins mirror all of the above. The gfx942 2D C++ entry also builds the gfx950 fastKV register-P twin, so its runtime changes are gated to gfx942. - **CI enrollment.** `test_attention_tiled_launch_signature.py` joins the installed attention CPU cohort; `test_attention_gfx942_runtime_shape_numeric.py` joins the library GPU cohort. **ABI note:** in this repo, the gfx942 tiled kernels are launched only through `run_unified_attention_torch`, the dispatch tuning candidates and the harnesses updated here. Anything that launches them by name must append `num_query_heads`, `num_kv_heads` and `sliding_window` (2D and 3D) or `num_query_heads` (reduce). **Notes for the packaging work:** - The gfx942 fast paged-KV descriptor path (`use_fast_paged_kv_desc`, reachable only through the `fastkv` tuning knob) still hard-codes 64/8 heads, pinned by its spec validator. A packaged catalog must serve it only at 64/8 or leave that knob out. - gfx942 routes windowed decode to 2D, so production never takes the windowed 3D path; byte identity covers it. - `AttentionTuningSpec.cache_key()` remains field-complete, so it still splits on the runtime fields. That over-keys but is never wrong. ## Submission Checklist - [ ] Look over the contributing guidelines at https://github.com/ROCm/TheRock/blob/main/GOVERNANCE.md#pull-requests. --------- Co-authored-by: Claude --- .../prefill/benchmark_prefill2d_live.py | 6 +- .../attention/prefill/bf16_wide_check.py | 6 +- .../prefill/parity_unified_attention.py | 6 +- .../library/dispatch/attention/candidate.py | 17 +- .../dispatch/attention/strided_decode.py | 6 +- .../kernels/common/attention_unified.py | 283 +++++-- .../kernels/gfx942/attention_tiled_2d.py | 196 +++-- .../kernels/gfx942/attention_tiled_3d.py | 131 ++-- .../attention_strided_kv_ir_sha256.json | 42 +- .../library/tests/test_attention_builds.py | 2 +- ..._attention_gfx942_runtime_shape_numeric.py | 299 ++++++++ .../test_attention_tiled_3d_running_max.py | 2 + .../tests/test_attention_tiled_cache_key.py | 325 +++++++- .../test_attention_tiled_launch_signature.py | 232 ++++++ .../library/tests/test_attn_bf16_d128_ring.py | 6 +- .../library/tests/test_strided_kv_decode.py | 4 +- .../rocke/platform/CMakeLists.txt | 16 +- .../platform/cpp/helpers/attention_ext.cpp | 55 +- .../rocke/platform/cpp/helpers/transforms.cpp | 77 +- .../helper_helper_rocke.helpers.attention.h | 17 + .../rocke/helper_rocke.helpers.transforms.h | 17 +- ...tance_gfx942_attention_tiled_2d_internal.h | 22 +- .../instance_gfx942_attention_tiled_3d.h | 9 +- ...tance_gfx942_attention_tiled_3d_internal.h | 14 +- .../attention_tiled_2d_ctx_prologue.cpp | 42 +- .../gfx942/attention_tiled_2d_epilogue.cpp | 8 +- .../attention_tiled_2d_kv_body_qk_softmax.cpp | 6 +- ...ttention_tiled_2d_kv_descriptors_loads.cpp | 38 +- .../attention_tiled_2d_loop_scaffold.cpp | 3 +- .../attention_tiled_2d_public_entry_glue.cpp | 16 +- .../attention_tiled_2d_q_load_gather.cpp | 22 +- ...ed_3d_gfx942_attn_tiled_3d_public_glue.cpp | 36 +- ...n_tiled_3d_gfx942_attn_tiled_3d_reduce.cpp | 35 +- ...d_gfx942_attn_tiled_3d_segment_compute.cpp | 7 +- ..._3d_gfx942_attn_tiled_3d_segment_setup.cpp | 102 +-- .../platform/dsl_docs/reference/file_index.md | 2 +- .../platform/python/rocke/helpers/__init__.py | 2 + .../python/rocke/helpers/attention.py | 20 +- .../python/rocke/helpers/transforms.py | 143 +++- .../tests/core/test_transforms_dynamic.py | 78 ++ .../rocke_representative_ir_sha256.json | 198 ++--- .../differential/golden/llvm_gfx_all.json | 696 +++++++++--------- 42 files changed, 2355 insertions(+), 889 deletions(-) create mode 100644 dnn-providers/hip-kernel-provider/rocke/library/tests/test_attention_gfx942_runtime_shape_numeric.py create mode 100644 dnn-providers/hip-kernel-provider/rocke/library/tests/test_attention_tiled_launch_signature.py diff --git a/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/gfx942/attention/prefill/benchmark_prefill2d_live.py b/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/gfx942/attention/prefill/benchmark_prefill2d_live.py index 3fd87cdb6f50..4f335715f6c4 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/gfx942/attention/prefill/benchmark_prefill2d_live.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/gfx942/attention/prefill/benchmark_prefill2d_live.py @@ -464,9 +464,9 @@ def build( supports_tiled_2d, ) from kernels.common.attention_unified import ( - _attn_signature, _enable_i64_kv_addr, _tiled_2d_impl, + _tiled_2d_signature, ) from rocke.runtime import KernelLauncher @@ -540,9 +540,7 @@ def build( KernelLauncher( hsaco=artifact.hsaco, kernel_name=artifact.kernel_name, - signature=_attn_signature( - dtype, include_bt_stride=True, include_qq_bias_stride=True - ), + signature=_tiled_2d_signature(ARCH, dtype), cache_key=("prefill2d_live_gfx942", key), ), spec, diff --git a/dnn-providers/hip-kernel-provider/rocke/library/builders/gfx942/attention/prefill/bf16_wide_check.py b/dnn-providers/hip-kernel-provider/rocke/library/builders/gfx942/attention/prefill/bf16_wide_check.py index 33ab7d493f21..949562d40ba9 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/builders/gfx942/attention/prefill/bf16_wide_check.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/builders/gfx942/attention/prefill/bf16_wide_check.py @@ -39,7 +39,7 @@ from kernels.gfx942.attention_tiled_2d import ( # noqa: E402 UnifiedAttention2DTiledSpec, ) -from kernels.common.attention_unified import _attn_signature # noqa: E402 +from kernels.common.attention_unified import _tiled_2d_signature # noqa: E402 from rocke.runtime import KernelLauncher # noqa: E402 BS = H.BLOCK_SIZE @@ -136,9 +136,7 @@ def build_launcher(s, spec, kind): launcher = KernelLauncher( hsaco=artifact.hsaco, kernel_name=artifact.kernel_name, - signature=_attn_signature( - s.dtype, include_bt_stride=True, include_qq_bias_stride=True - ), + signature=_tiled_2d_signature("gfx942", s.dtype), cache_key=("bf16wide", spec.kernel_name(), kind), ) return launcher diff --git a/dnn-providers/hip-kernel-provider/rocke/library/builders/gfx942/attention/prefill/parity_unified_attention.py b/dnn-providers/hip-kernel-provider/rocke/library/builders/gfx942/attention/prefill/parity_unified_attention.py index 0c5c3f1eadb0..0f957fda5697 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/builders/gfx942/attention/prefill/parity_unified_attention.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/builders/gfx942/attention/prefill/parity_unified_attention.py @@ -774,7 +774,7 @@ def _build_kernel(s: Shape): """ from rocke import compile_kernel from kernels import build_unified_attention_2d_tiled, supports_tiled_2d - from kernels.common.attention_unified import _attn_signature + from kernels.common.attention_unified import _tiled_2d_signature from rocke.runtime import KernelLauncher spec, config = _build_spec(s) @@ -810,9 +810,7 @@ def _build_kernel(s: Shape): launcher = KernelLauncher( hsaco=artifact.hsaco, kernel_name=artifact.kernel_name, - signature=_attn_signature( - s.dtype, include_bt_stride=True, include_qq_bias_stride=True - ), + signature=_tiled_2d_signature("gfx942", s.dtype), cache_key=("gfx942_attn_example", spec.kernel_name(), config), ) return launcher, spec, config diff --git a/dnn-providers/hip-kernel-provider/rocke/library/dispatch/attention/candidate.py b/dnn-providers/hip-kernel-provider/rocke/library/dispatch/attention/candidate.py index b95dc6a745ad..2ff45f51c115 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/dispatch/attention/candidate.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/dispatch/attention/candidate.py @@ -77,16 +77,21 @@ def build(spec, arch): return spec.build(arch) def signature(spec): - from kernels.common.attention_unified import _3d_signature, _attn_signature + from kernels.common.attention_unified import ( + _tiled_2d_signature, + _tiled_3d_signatures, + ) if spec.path == "3d": - return _3d_signature( - spec.kernel_spec.dtype, kv_dtype=spec.kernel_spec.kv_storage_dtype + segment, _ = _tiled_3d_signatures( + spec.arch, + spec.kernel_spec.dtype, + kv_dtype=spec.kernel_spec.kv_storage_dtype, ) - return _attn_signature( + return segment + return _tiled_2d_signature( + spec.arch, spec.kernel_spec.dtype, - include_bt_stride=True, - include_qq_bias_stride=True, kv_dtype=spec.kernel_spec.kv_storage_dtype, ) diff --git a/dnn-providers/hip-kernel-provider/rocke/library/dispatch/attention/strided_decode.py b/dnn-providers/hip-kernel-provider/rocke/library/dispatch/attention/strided_decode.py index 340bada8c9ec..d9da74e85bb1 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/dispatch/attention/strided_decode.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/dispatch/attention/strided_decode.py @@ -9,8 +9,8 @@ from kernels.common.attention_unified import ( UnifiedAttentionProblem, - _3d_signature, _strided_3d_specs_from_problem, + _tiled_3d_signatures, _validate_strided_3d_spec, ) from rocke.dispatch.core import Capability, KernelCandidate, ShapeRange, TorchBinding @@ -166,7 +166,9 @@ def launch(*, softmax_scale=softmax_scale, stream=stream): select_spec=select, sweep_space=lambda req: (select(req),) if candidate.admits(req)[0] else (), build=lambda spec, arch: spec.build(arch), - signature=lambda spec: _3d_signature(spec.kernel_spec.dtype, strided_kv=True), + signature=lambda spec: _tiled_3d_signatures( + spec.arch, spec.kernel_spec.dtype, strided_kv=True + )[0], grid=lambda spec, req: spec.launch_grid(_problem(req)), block=lambda spec: spec.launch_block(), bind_torch=bind, 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 641143464c89..043c682701c0 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 @@ -13,7 +13,7 @@ from collections import OrderedDict from dataclasses import dataclass, replace from threading import Lock, RLock -from typing import Any, Callable, Dict, Optional, Tuple +from typing import Any, Callable, Dict, Optional, Sequence, Tuple from rocke.core.arch import validate_arch from rocke.core.ir import ( @@ -1653,10 +1653,11 @@ def _kv_storage_dtype(problem: UnifiedAttentionProblem) -> Optional[str]: def _binary_search_iters(problem: UnifiedAttentionProblem, *, floor: int = 1) -> int: """Trip count of the seq_idx binary search the tiled kernels bake into IR. - Every tiled spec class derives ``binary_search_iters`` from ``spec.num_seqs`` - through the same ``binary_search_seq_idx_iters`` (gfx1250 3D with a floor of - 32). It is the only way ``num_seqs`` reaches emitted IR, so the tiled cache - keys carry this bucket instead of the raw batch count. + Every tiled spec class that bakes the search derives ``binary_search_iters`` + from ``spec.num_seqs`` through the same ``binary_search_seq_idx_iters`` + (gfx1250 3D with a floor of 32). It is the only way ``num_seqs`` reaches + emitted IR, so the tiled cache keys carry this bucket instead of the raw batch + count. A body that searches at runtime (``_searches_at_runtime``) bakes neither. """ return binary_search_seq_idx_iters(problem.num_seqs, floor=floor) @@ -1699,7 +1700,8 @@ def _tiled_cache_key(problem: UnifiedAttentionProblem, arch: str) -> Tuple: ``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 + ``_binary_search_iters``, and not at all in a body that searches at runtime + (``_searches_at_runtime``). Launch values that scale with the shape (the 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 @@ -1715,13 +1717,12 @@ def _tiled_cache_key(problem: UnifiedAttentionProblem, arch: str) -> Tuple: return ( "tiled", arch, - _binary_search_iters(problem), - problem.num_query_heads, - problem.num_kv_heads, + None if _searches_at_runtime(arch, "2d") else _binary_search_iters(problem), + _heads_key(problem, arch, "2d"), problem.head_size, problem.block_size, problem.dtype, - problem.sliding_window, + _window_key(problem, arch, "2d"), bool(problem.use_sinks), bool(problem.softcap > 0), bool(problem.use_alibi), @@ -4060,9 +4061,10 @@ def _tiled_3d_cache_key(problem: UnifiedAttentionProblem, arch: str) -> Tuple: with the shape are keyed separately through ``_shape_key``. """ decode_grid = _use_decode_grid(problem, arch) - if decode_grid: - # One CTA per sequence: the segment kernel indexes sequences directly - # and emits no search, so batch sizes in different buckets share it. + if decode_grid or _searches_at_runtime(arch, "3d"): + # The decode grid indexes sequences directly and emits no search; a + # runtime search reads the batch from num_seqs. Either way batch sizes + # in different buckets share the kernel. search_iters = None elif arch == "gfx1250": search_iters = _binary_search_iters(problem, floor=BINARY_SEARCH_FALLBACK_ITERS) @@ -4076,12 +4078,11 @@ def _tiled_3d_cache_key(problem: UnifiedAttentionProblem, arch: str) -> Tuple: # no way to check the key against the arch it compiles for. arch, search_iters, - problem.num_query_heads, - problem.num_kv_heads, + _heads_key(problem, arch, "3d"), problem.head_size, problem.block_size, problem.dtype, - problem.sliding_window, + _window_key(problem, arch, "3d"), bool(problem.use_sinks), bool(problem.softcap > 0), bool(problem.use_alibi), @@ -4119,12 +4120,113 @@ def _tiled_3d_cache_key(problem: UnifiedAttentionProblem, arch: str) -> Tuple: return base +def _runtime_kernarg_fields(spec_cls) -> Tuple[str, ...]: + """The problem fields a tiled spec class reads as i32 kernargs, in ABI order. + + The spec class declares them (``RUNTIME_KERNARG_FIELDS``) so its body's + params, the launch signature and the cache key read one tuple. A class + declaring none bakes those fields into its IR. + """ + return tuple(getattr(spec_cls, "RUNTIME_KERNARG_FIELDS", ())) + + +_TILED_SPEC_CLS: Dict[Tuple[str, str], type] = {} + + +def _tiled_spec_cls(arch: str, path: str) -> type: + """The tiled ``path`` ("2d" or "3d" segment) spec class for ``arch``. Memoized: + the cache keys read its class attributes on every launch.""" + key = (arch, path) + if key not in _TILED_SPEC_CLS: + impl = _tiled_2d_impl(arch) if path == "2d" else _tiled_3d_impl(arch) + _TILED_SPEC_CLS[key] = impl[0] + return _TILED_SPEC_CLS[key] + + +def _searches_at_runtime(arch: str, path: str) -> bool: + """True when the tiled ``path`` body for ``arch`` derives its cu_q search trip + count from the num_seqs kernarg (``RUNTIME_SEARCH_ITERS``), so num_seqs never + reaches its IR.""" + return bool(getattr(_tiled_spec_cls(arch, path), "RUNTIME_SEARCH_ITERS", False)) + + +def _heads_key(problem: UnifiedAttentionProblem, arch: str, path: str) -> Tuple: + """The head counts as the tiled ``path`` body for ``arch`` bakes them: both + counts, or only the GQA ratio when the counts are kernargs.""" + fields = _runtime_kernarg_fields(_tiled_spec_cls(arch, path)) + if "num_query_heads" in fields and "num_kv_heads" in fields: + return ("gqa", problem.num_queries_per_kv) + return (problem.num_query_heads, problem.num_kv_heads) + + +def _window_key(problem: UnifiedAttentionProblem, arch: str, path: str): + """The sliding window as the tiled ``path`` body for ``arch`` bakes it: the size + itself, or only whether there is one when the size is a kernarg.""" + if "sliding_window" in _runtime_kernarg_fields(_tiled_spec_cls(arch, path)): + return problem.sliding_window > 0 + return problem.sliding_window + + +def _runtime_kernarg_values(problem: UnifiedAttentionProblem) -> dict: + """Every runtime shape kernarg a tiled body can declare. The launcher packs + only the ones its signature names.""" + return { + "num_query_heads": int(problem.num_query_heads), + "num_kv_heads": int(problem.num_kv_heads), + "sliding_window": int(problem.sliding_window), + } + + +def _tiled_2d_signature( + arch: str, + dtype: str, + *, + kv_dtype: Optional[str] = None, + include_additive_bias: bool = False, +): + """Launch signature of the tiled 2D kernels built for ``arch``.""" + spec_cls, _, _ = _tiled_2d_impl(arch) + return _attn_signature( + dtype, + include_bt_stride=True, + include_qq_bias_stride=True, + include_additive_bias=include_additive_bias, + kv_dtype=kv_dtype, + runtime_fields=_runtime_kernarg_fields(spec_cls), + ) + + +def _tiled_3d_signatures( + arch: str, + dtype: str, + *, + kv_dtype: Optional[str] = None, + strided_kv: bool = False, + include_additive_bias: bool = False, + emit_lse: bool = False, +): + """Launch signatures of the tiled 3D segment and reduce kernels for ``arch``.""" + segment_cls, reduce_cls, _, _, _ = _tiled_3d_impl(arch) + segment = _3d_signature( + dtype, + kv_dtype=kv_dtype, + strided_kv=strided_kv, + include_additive_bias=include_additive_bias, + runtime_fields=_runtime_kernarg_fields(segment_cls), + ) + reduce = _reduce_signature( + dtype, emit_lse=emit_lse, runtime_fields=_runtime_kernarg_fields(reduce_cls) + ) + return segment, reduce + + def _3d_signature( dtype: str, *, kv_dtype: Optional[str] = None, strided_kv: bool = False, include_additive_bias: bool = False, + runtime_fields: Sequence[str] = (), ): from rocke.helpers.spec import SignatureBuilder @@ -4152,6 +4254,8 @@ def _3d_signature( .scalar("block_table_stride", "i32") .scalar("qq_bias_stride_0", "i32") ) + for name in runtime_fields: + sb.scalar(name, "i32") # Conditional params, in the kernel's declaration order (bias, then strided KV). if include_additive_bias: append_additive_bias_signature(sb) @@ -4164,7 +4268,9 @@ def _3d_signature( return sb.build() -def _reduce_signature(dtype: str, *, emit_lse: bool = False): +def _reduce_signature( + dtype: str, *, emit_lse: bool = False, runtime_fields: Sequence[str] = () +): from rocke.helpers.spec import SignatureBuilder io_dtype = "f16" if dtype == "fp16" else "bf16" @@ -4176,6 +4282,8 @@ def _reduce_signature(dtype: str, *, emit_lse: bool = False): .ptr("segm_expsum_ptr", "f32") .ptr("seq_lens_ptr", "i32") ) + for name in runtime_fields: + sig = sig.scalar(name, "i32") if emit_lse: sig = sig.ptr("lse_ptr", "f32") return sig.build() @@ -4188,6 +4296,7 @@ def _attn_signature( include_qq_bias_stride: bool = False, include_additive_bias: bool = False, kv_dtype: Optional[str] = None, + runtime_fields: Sequence[str] = (), ): from rocke.helpers.spec import SignatureBuilder @@ -4219,6 +4328,8 @@ def _attn_signature( sb.scalar("block_table_stride", "i32") if include_qq_bias_stride: sb.scalar("qq_bias_stride_0", "i32") + for name in runtime_fields: + sb.scalar(name, "i32") # Appended last, matching the kernel's conditional param declarations. if include_additive_bias: append_additive_bias_signature(sb) @@ -4279,6 +4390,7 @@ def _attn_values( "out_scale": float(out_scale), "softcap": float(softcap), "num_seqs": int(problem.num_seqs), + **_runtime_kernarg_values(problem), } if include_bt_stride: vals["block_table_stride"] = int(bt_stride) @@ -4625,6 +4737,7 @@ def _run_3d_tiled_on_stream( "additive_bias_batch_stride": int(additive_bias_batch_stride), "additive_bias_head_stride": int(additive_bias_head_stride), "additive_bias_sq_stride": int(additive_bias_sq_stride), + **_runtime_kernarg_values(problem), } if kv_cache_layout is not None: seg_vals.update(kv_cache_layout.arguments()) @@ -4634,6 +4747,7 @@ def _run_3d_tiled_on_stream( "segm_max_ptr": segm_max, "segm_expsum_ptr": segm_expsum, "seq_lens_ptr": seqused_k, + **_runtime_kernarg_values(problem), } if problem.emit_lse: red_vals["lse_ptr"] = lse @@ -5297,15 +5411,54 @@ 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. + The tiled binaries are shared across batch counts, token counts and, where + the body reads them as kernargs, head counts, but the grids (``total_q // + block_q + num_seqs`` by the kv or query heads), 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), + "heads", + int(problem.num_query_heads), + int(problem.num_kv_heads), + ) + + +def _tiled_3d_kernels( + problem: UnifiedAttentionProblem, + arch: str, + num_segments: int, + *, + strided_kv: bool = False, +): + """The tiled 3D segment and reduce ``KernelDef``s the production path builds.""" + ( + _, + UnifiedAttentionReduceTiledSpec, + build_unified_attention_3d_tiled, + build_unified_attention_reduce_tiled, + _, + ) = _tiled_3d_impl(arch) + if strided_kv: + segment_spec, reduce_spec = _strided_3d_specs_from_problem(problem, arch=arch) + else: + segment_spec = _tiled_3d_spec_from_problem(problem, arch) + reduce_spec = UnifiedAttentionReduceTiledSpec( + head_size=problem.head_size, + num_query_heads=problem.num_query_heads, + num_kv_heads=problem.num_kv_heads, + dtype=problem.dtype, + num_segments=num_segments, + waves_per_eu=_select_3d_waves_per_eu(problem, arch), + # Only when on: gfx1250's reduce spec has no such field. + **({"emit_lse": True} if problem.emit_lse else {}), + ) + return ( + build_unified_attention_3d_tiled(segment_spec, arch=arch), + build_unified_attention_reduce_tiled(reduce_spec, arch=arch), ) @@ -5333,31 +5486,9 @@ def _get_3d_pipeline( if tuning_spec is not None: seg_kernel, red_kernel = tuning_spec.build(arch) else: - ( - _, - UnifiedAttentionReduceTiledSpec, - build_unified_attention_3d_tiled, - build_unified_attention_reduce_tiled, - _, - ) = _tiled_3d_impl(arch) - if strided_kv: - segment_spec, reduce_spec = _strided_3d_specs_from_problem( - problem, arch=arch - ) - else: - segment_spec = _tiled_3d_spec_from_problem(problem, arch) - reduce_spec = UnifiedAttentionReduceTiledSpec( - head_size=problem.head_size, - num_query_heads=problem.num_query_heads, - num_kv_heads=problem.num_kv_heads, - dtype=problem.dtype, - num_segments=num_segments, - waves_per_eu=_select_3d_waves_per_eu(problem, arch), - # Only when on: gfx1250's reduce spec has no such field. - **({"emit_lse": True} if problem.emit_lse else {}), - ) - seg_kernel = build_unified_attention_3d_tiled(segment_spec, arch=arch) - red_kernel = build_unified_attention_reduce_tiled(reduce_spec, arch=arch) + seg_kernel, red_kernel = _tiled_3d_kernels( + problem, arch, num_segments, strided_kv=strided_kv + ) seg_art = compile_kernel(seg_kernel, arch=arch, capture_ir_text=False) red_art = compile_kernel(red_kernel, arch=arch, capture_ir_text=False) _ATTN_3D_TILED_CACHE[cache_key] = ( @@ -5367,28 +5498,31 @@ def _get_3d_pipeline( red_art.kernel_name, ) seg_hsaco, seg_kname, red_hsaco, red_kname = _ATTN_3D_TILED_CACHE[cache_key] + # problem.emit_lse matches the built reduce: an explicit tuning spec must + # agree with it (paged: run_unified_attention_torch; strided: + # _validate_strided_3d_spec refuses an LSE reduce). + seg_signature, red_signature = _tiled_3d_signatures( + arch, + problem.dtype, + strided_kv=strided_kv, + kv_dtype=( + tuning_spec.kernel_spec.kv_storage_dtype + if tuning_spec is not None + else _kv_storage_dtype(problem) + ), + include_additive_bias=problem.use_additive_bias, + emit_lse=problem.emit_lse, + ) seg_launcher = KernelLauncher( hsaco=seg_hsaco, kernel_name=seg_kname, - signature=_3d_signature( - problem.dtype, - strided_kv=strided_kv, - kv_dtype=( - tuning_spec.kernel_spec.kv_storage_dtype - if tuning_spec is not None - else _kv_storage_dtype(problem) - ), - include_additive_bias=problem.use_additive_bias, - ), + signature=seg_signature, cache_key=("3d_seg",) + cache_key, ) red_launcher = KernelLauncher( hsaco=red_hsaco, kernel_name=red_kname, - # problem.emit_lse matches the built reduce: an explicit tuning spec - # must agree with it (paged: run_unified_attention_torch; strided: - # _validate_strided_3d_spec refuses an LSE reduce). - signature=_reduce_signature(problem.dtype, emit_lse=problem.emit_lse), + signature=red_signature, cache_key=("3d_red",) + cache_key, ) pipeline = PipelineLauncher([seg_launcher, red_launcher]) @@ -5514,6 +5648,20 @@ def _select_2d_compile_backend(problem: UnifiedAttentionProblem, arch: str) -> s return "llvm" +def _tiled_2d_kernel(problem: UnifiedAttentionProblem, arch: str): + """The tiled 2D ``KernelDef`` the production path builds for ``problem``.""" + _, build_unified_attention_2d_tiled, _ = _tiled_2d_impl(arch) + spec = _tiled_spec_from_problem(problem, arch) + route = _gfx942_4warp_route(problem, arch) + if route is not None: + # Distinct 4-warp GQA paged builder (keyed separately in + # `_tiled_cache_key`; grid in `_get_2d_launch_meta`). Same paged + # ABI as the default builder. Parity+ with AITER @Sq4096/8192 + # (vs the 1-warp std-QK's 0.55x). + return route.builder(spec, arch=arch) + return build_unified_attention_2d_tiled(spec, arch=arch) + + def _get_2d_launcher( problem: UnifiedAttentionProblem, cache_key: Tuple, @@ -5536,17 +5684,7 @@ def _get_2d_launcher( kernel = tuning_spec.build(arch) backend = tuning_spec.compile_backend else: - _, build_unified_attention_2d_tiled, _ = _tiled_2d_impl(arch) - spec = _tiled_spec_from_problem(problem, arch) - route = _gfx942_4warp_route(problem, arch) - if route is not None: - # Distinct 4-warp GQA paged builder (keyed separately in - # `_tiled_cache_key`; grid in `_get_2d_launch_meta`). Same paged - # ABI as the default builder. Parity+ with AITER @Sq4096/8192 - # (vs the 1-warp std-QK's 0.55x). - kernel = route.builder(spec, arch=arch) - else: - kernel = build_unified_attention_2d_tiled(spec, arch=arch) + kernel = _tiled_2d_kernel(problem, arch) backend = _select_2d_compile_backend(problem, arch) if backend == "hipcc": from rocke.helpers.compile import compile_kernel_via_hipcc @@ -5559,10 +5697,9 @@ def _get_2d_launcher( launcher = KernelLauncher( hsaco=hsaco, kernel_name=kname, - signature=_attn_signature( + signature=_tiled_2d_signature( + arch, problem.dtype, - include_bt_stride=True, - include_qq_bias_stride=True, include_additive_bias=problem.use_additive_bias, kv_dtype=( tuning_spec.kernel_spec.kv_storage_dtype diff --git a/dnn-providers/hip-kernel-provider/rocke/library/kernels/gfx942/attention_tiled_2d.py b/dnn-providers/hip-kernel-provider/rocke/library/kernels/gfx942/attention_tiled_2d.py index 7608481a0641..ca4f7ad91b66 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/kernels/gfx942/attention_tiled_2d.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/kernels/gfx942/attention_tiled_2d.py @@ -36,7 +36,7 @@ from dataclasses import dataclass import os -from typing import Optional, Tuple +from typing import ClassVar, Optional, Tuple from rocke.core.ir import ( BF16, @@ -58,7 +58,7 @@ from rocke.helpers.attention import ( apply_softcap_log2, binary_search_seq_idx, - binary_search_seq_idx_iters, + binary_search_seq_idx_trip_count, dequant_fp8x8_to_dtype, mfma_16x16x16_for_dtype, mfma_16x16x32_for_dtype, @@ -73,7 +73,13 @@ from rocke.helpers.distribution import make_static_tile_distribution from rocke.helpers.layouts import TransposeLdsReader from rocke.helpers.mfma_gemm_inner import decode_mfma_lanes -from rocke.helpers.transforms import TensorDescriptor, embed, indirect, unmerge +from rocke.helpers.transforms import ( + DynamicTensorDescriptor, + TensorDescriptor, + embed, + indirect, + unmerge, +) MFMA_M = 16 @@ -161,6 +167,18 @@ def _sched_group_pin_mfma_step(b, ds_reads: int, mfma_count: int, sgid: int) -> @dataclass(frozen=True) class UnifiedAttention2DTiledSpec: + # Problem fields this body reads as trailing i32 kernargs instead of baking + # them, in ABI order. The launch signature, the argument pack and the cache + # key read the same tuple (``_runtime_kernarg_fields``). + RUNTIME_KERNARG_FIELDS: ClassVar[Tuple[str, ...]] = ( + "num_query_heads", + "num_kv_heads", + "sliding_window", + ) + # The cu_q search derives its trip count from the num_seqs kernarg, so + # num_seqs never reaches the IR and one binary serves every batch. + RUNTIME_SEARCH_ITERS: ClassVar[bool] = True + head_size: int block_size: int num_query_heads: int @@ -960,15 +978,6 @@ def n_blocks_per_tile(self) -> int: def dtype_ir(self) -> Type: return F16 if self.dtype == "fp16" else BF16 - @property - def binary_search_iters(self) -> int: - # AITER/Triton uses a true while-loop binary search. Our IR currently - # lowers this as a fixed-trip scf.for, so specialize the trip count to - # the known problem batch size instead of always paying 32 iterations. - # Falls back to 32 for direct unit-test specs that do not provide - # `num_seqs`. - return binary_search_seq_idx_iters(self.num_seqs) - def kernel_name(self) -> str: from rocke.helpers.spec import kernel_name_join @@ -981,12 +990,22 @@ def kernel_name(self) -> str: f"d{self.head_size}", f"b{self.block_size}", f"t{self.tile_size_eff}" if self.n_blocks_per_tile != 1 else "", - f"h{self.num_query_heads}kv{self.num_kv_heads}", + # Head counts are kernargs; only the GQA ratio is baked. + f"g{self.num_queries_per_kv}", + # Except on the fast paged-KV descriptor path, whose shifts bake them. + ( + f"h{self.num_query_heads}kv{self.num_kv_heads}" + if self.use_fast_paged_kv_desc + else "" + ), self.dtype, f"kv{self.kv_storage_dtype}" if self.kv_storage_dtype else "", "fnuz" if self.kv_storage_dtype else "", + # i64 K/V addressing changes the loaders, so it is a distinct kernel. + "i64" if self.use_i64_kv_addr else "", "" if not self.use_sinks else "sinks", - f"sw{self.sliding_window}" if self.sliding_window > 0 else "", + # The window size is a kernarg; only whether the body masks one is baked. + "sw" if self.sliding_window > 0 else "", "softcap" if self.has_softcap else "", "alibi" if self.use_alibi else "", "qqb" if self.use_qq_bias else "", @@ -1377,7 +1396,6 @@ def build_unified_attention_2d_tiled( # rows would duplicate the first tokens of the NEXT CTA and must be masked. VALID_ROWS = BLOCK_Q * NQK NUM_KV = spec.num_kv_heads - NUM_QH = spec.num_query_heads SLIDING_WINDOW = spec.sliding_window USE_SOFTCAP = spec.has_softcap USE_SINKS = spec.use_sinks @@ -1569,6 +1587,11 @@ def build_unified_attention_2d_tiled( num_seqs_p = b.param("num_seqs", I32) bt_stride_p = b.param("block_table_stride", I32) qq_bias_stride0_p = b.param("qq_bias_stride_0", I32) + # Shape values this body reads as kernargs (``RUNTIME_KERNARG_FIELDS``), + # after every fixed param so the shared launch signature can append them. + rt = {f: b.param(f, I32) for f in spec.RUNTIME_KERNARG_FIELDS} + num_qh_p = rt["num_query_heads"] + num_kv_p = rt["num_kv_heads"] if spec.use_q_major_grid: q_block_global_idx = b.block_id_x() @@ -1596,7 +1619,7 @@ def build_unified_attention_2d_tiled( q_block_global_idx, num_seqs_p, block_q=BLOCK_Q, - iterations=spec.binary_search_iters, + iterations=binary_search_seq_idx_trip_count(b, num_seqs_p), ) cu_q_start = b.global_load_i32(cu_q, seq_idx) cu_q_stop = b.global_load_i32(cu_q, b.add(seq_idx, b.const_i32(1))) @@ -2143,7 +2166,7 @@ def _v_load1(v_buf: Value, v_row: Value, v_n_col: Value) -> Value: # stays the same as the default bf16 path -- just softmax_scale / # ln(2) -- and we do NOT fold k_scale into qk_scale here. pv_fp8_scale = b.fdiv(v_scale_p, b.const_f32(240.0)) if FP8_MFMA_PV else None - sw_const = b.const_i32(int(SLIDING_WINDOW)) + sw_const = rt["sliding_window"] z8 = b.zero_vec(dtype, 8) # ---------------- Q -> LDS (cooperative vec8 chunks) ---------------- @@ -2156,15 +2179,11 @@ def _v_load1(v_buf: Value, v_row: Value, v_n_col: Value) -> Value: # Coordinate transform for Q (and the symmetric output buffer): # ``(token, head, dim)`` packed contiguously. The element-unit # descriptor is reused below for the output store too. - q_desc = TensorDescriptor.naive( + # The head count is a kernarg, so the token stride is too. + q_desc = DynamicTensorDescriptor.create( "Q", - # The runtime extents (total_q, num_query_heads, head_size) are - # only used by the validity predicate (which we don't request - # here). Use generous compile-time bounds so the descriptor's - # row-major stride product matches the kernel's layout - # assumptions exactly. - lengths=[1 << 30, NUM_QH, HD], coord_names=("token", "head", "dim"), + strides=(b.mul(num_qh_p, b.const_i32(HD)), HD, 1), ) if not Q_DIRECT_GLOBAL: for li in range(Q_VECS_PER_THREAD): @@ -2177,7 +2196,7 @@ def _v_load1(v_buf: Value, v_row: Value, v_n_col: Value) -> Value: ) _qmask_inner = b.land( b.cmp_lt(q_pos_t, cur_batch_q_len), - b.cmp_lt(qh_t, b.const_i32(NUM_QH)), + b.cmp_lt(qh_t, num_qh_p), ) qmask_t = ( b.land(_qmask_inner, b.cmp_lt(Q_row, b.const_i32(VALID_ROWS))) @@ -2359,7 +2378,7 @@ def _pack_p_a32(p_regs0_f32: list[Value], p_regs1_f32: list[Value]) -> Value: qh = b.add( b.mul(kv_head_idx, b.const_i32(NQK)), b.mod(row, b.const_i32(NQK)) ) - qh_in = b.cmp_lt(qh, b.const_i32(NUM_QH)) + qh_in = b.cmp_lt(qh, num_qh_p) sink_h = b.global_load(sinks, qh, dtype, align=2) sink_f = b.fmul(b.cast_to_f32(sink_h), rcp_ln2) m_inits.append(b.select(qh_in, sink_f, neg_inf)) @@ -2369,7 +2388,7 @@ def _pack_p_a32(p_regs0_f32: list[Value], p_regs1_f32: list[Value]) -> Value: qh = b.add( b.mul(kv_head_idx, b.const_i32(NQK)), b.mod(row, b.const_i32(NQK)) ) - qh_in = b.cmp_lt(qh, b.const_i32(NUM_QH)) + qh_in = b.cmp_lt(qh, num_qh_p) sink_h = b.global_load(sinks, qh, dtype, align=2) sink_f = b.fmul(b.cast_to_f32(sink_h), rcp_ln2) m_inits.append(b.select(qh_in, sink_f, neg_inf)) @@ -2463,10 +2482,11 @@ def _acc_idx(n: int, atom: int) -> int: # so for the FP8 K/V path the loader switches to a sync per-thread # dequant chain below; the byte-stride math here is shared so the # paged_kv_desc compiles consistently in both branches. - kv_stride_blk_b = BS * NUM_KV * HD * KV_BYTES - kv_stride_tok_b = NUM_KV * HD * KV_BYTES + # The token and block strides scale with the runtime kv head count. + kv_stride_tok_b = b.mul(num_kv_p, b.const_i32(HD * KV_BYTES)) + kv_stride_blk_b = b.mul(kv_stride_tok_b, b.const_i32(BS)) kv_stride_h_b = HD * KV_BYTES - kv_block_bytes_c = b.const_i32(kv_stride_blk_b) # one-block buffer bound + kv_block_bytes_c = kv_stride_blk_b # one-block buffer bound # Per-lane starting half index. With the gfx942 1-dword DMA each lane # contributes ``KV_HALVES_PER_LANE`` (= 2) contiguous halves per call, @@ -2636,13 +2656,10 @@ def _fast_paged_kv_voff(call: int, block0: Value, block1: Value): return None, b.add(block_b, within) else: - _kv_base = TensorDescriptor.naive( + _kv_base = DynamicTensorDescriptor.create( "paged_kv_bytes", - # ``lengths`` here is just informational (validity propagation - # is driven by the transforms above, not by these bounds). - lengths=[1 << 24, BS, NUM_KV, HD], - strides=[kv_stride_blk_b, kv_stride_tok_b, kv_stride_h_b, KV_BYTES], coord_names=("physical_block", "token", "kv_head", "dim"), + strides=(kv_stride_blk_b, kv_stride_tok_b, kv_stride_h_b, KV_BYTES), ) if N_BLOCKS_PER_TILE == 1: paged_kv_desc = _kv_base.transform( @@ -2711,7 +2728,7 @@ def _tile_pages( align=4, ) base = b.mul( - b.zext(b.to_sgpr_u32(physical), I64), b.const_i64(kv_stride_blk_b) + b.zext(b.to_sgpr_u32(physical), I64), b.zext(kv_stride_blk_b, I64) ) ptrs.append(b.global_ptr_add(src, base)) rsrcs = [b.buffer_rsrc(ptr, kv_block_bytes_c) for ptr in ptrs] @@ -2727,7 +2744,7 @@ def _page_voff(call: int, kv_head_b: Value) -> Value: token = b.div(half, b.const_i32(HD)) dim = b.mod(half, b.const_i32(HD)) return b.add( - b.add(b.mul(token, b.const_i32(kv_stride_tok_b)), kv_head_b), + b.add(b.mul(token, kv_stride_tok_b), kv_head_b), b.mul(dim, b.const_i32(KV_BYTES)), ) @@ -3843,7 +3860,7 @@ def _read_k8_mfma_operand( ) q32_mask = b.land( b.cmp_lt(q32_pos, cur_batch_q_len), - b.cmp_lt(q32_h, b.const_i32(NUM_QH)), + b.cmp_lt(q32_h, num_qh_p), ) q32_pos_safe = b.select(q32_mask, q32_pos, b.const_i32(0)) q32_h_safe = b.select(q32_mask, q32_h, b.const_i32(0)) @@ -3948,7 +3965,7 @@ def _read_k8_mfma_operand( qh_r = b.add(b.mul(kv_head_idx, b.const_i32(NQK)), b.mod(row, b.const_i32(NQK))) _row_ok_inner = b.land( b.cmp_lt(qp_r, cur_batch_q_len), - b.cmp_lt(qh_r, b.const_i32(NUM_QH)), + b.cmp_lt(qh_r, num_qh_p), ) row_ok = ( b.land(_row_ok_inner, b.cmp_lt(row, b.const_i32(VALID_ROWS))) @@ -3966,7 +3983,7 @@ def _read_k8_mfma_operand( hoist_alibi = [] for reg in range(REGS_PER_LANE): qh_r = hoist_qh_r[reg] - qh_ok = b.cmp_lt(qh_r, b.const_i32(NUM_QH)) + qh_ok = b.cmp_lt(qh_r, num_qh_p) slope = b.masked_global_load( alibi_slopes_ptr, qh_r, qh_ok, b.const_f32(0.0), dtype=F32, align=4 ) @@ -3982,7 +3999,7 @@ def _read_k8_mfma_operand( ) _st_row_ok_inner = b.land( b.cmp_lt(st_qp, cur_batch_q_len), - b.cmp_lt(st_qh, b.const_i32(NUM_QH)), + b.cmp_lt(st_qh, num_qh_p), ) st_row_ok = ( b.land(_st_row_ok_inner, b.cmp_lt(st_q_row, b.const_i32(VALID_ROWS))) @@ -4009,7 +4026,7 @@ def _read_k8_mfma_operand( and TRANSPOSED_QK_32X32 and TRANSPOSED_INVARIANT_HOIST ): - st_qh_ok = b.cmp_lt(st_qh, b.const_i32(NUM_QH)) + st_qh_ok = b.cmp_lt(st_qh, num_qh_p) st_alibi_slope = b.masked_global_load( alibi_slopes_ptr, st_qh, st_qh_ok, b.const_f32(0.0), dtype=F32, align=4 ) @@ -4071,7 +4088,7 @@ def _acc_get(n: int, atom: int) -> Value: ) _st_row_ok_iter_inner = b.land( b.cmp_lt(st_qp_iter, cur_batch_q_len), - b.cmp_lt(st_qh_iter, b.const_i32(NUM_QH)), + b.cmp_lt(st_qh_iter, num_qh_p), ) st_row_ok_iter = ( b.land( @@ -4083,7 +4100,7 @@ def _acc_get(n: int, atom: int) -> Value: ) st_causal_lim_iter = b.add(context_len, st_qp_iter) if USE_ALIBI: - st_qh_iter_ok = b.cmp_lt(st_qh_iter, b.const_i32(NUM_QH)) + st_qh_iter_ok = b.cmp_lt(st_qh_iter, num_qh_p) st_alibi_slope_iter = b.masked_global_load( alibi_slopes_ptr, st_qh_iter, @@ -4406,7 +4423,7 @@ def _kslot(group_idx: int) -> int: slope_for_lane = b.masked_global_load( alibi_slopes_ptr, qh_t_alibi, - b.cmp_lt(qh_t_alibi, b.const_i32(NUM_QH)), + b.cmp_lt(qh_t_alibi, num_qh_p), b.const_f32(0.0), dtype=F32, align=4, @@ -4461,7 +4478,7 @@ def _kslot(group_idx: int) -> int: ) _row_ok_inner_t = b.land( b.cmp_lt(qp_r, cur_batch_q_len), - b.cmp_lt(qh_r, b.const_i32(NUM_QH)), + b.cmp_lt(qh_r, num_qh_p), ) row_ok = ( b.land( @@ -5528,7 +5545,7 @@ def _acc_final_get(n: int, atom: int) -> Value: ) _op_mask_t_inner = b.land( b.cmp_lt(op_pos_t, cur_batch_q_len), - b.cmp_lt(op_qh_t, b.const_i32(NUM_QH)), + b.cmp_lt(op_qh_t, num_qh_p), ) op_mask_t = ( b.land(_op_mask_t_inner, b.cmp_lt(q_row_t, b.const_i32(VALID_ROWS))) @@ -5581,7 +5598,7 @@ def _acc_final_get(n: int, atom: int) -> Value: ) _op_mask32_inner = b.land( b.cmp_lt(op_pos32_base, cur_batch_q_len), - b.cmp_lt(op_qh32_base, b.const_i32(NUM_QH)), + b.cmp_lt(op_qh32_base, num_qh_p), ) op_mask32_base = ( b.land(_op_mask32_inner, b.cmp_lt(OUT_ROW_BASE32, b.const_i32(VALID_ROWS))) @@ -5680,7 +5697,7 @@ def _acc_final_get(n: int, atom: int) -> Value: ) _op_mask_inner = b.land( b.cmp_lt(op_pos, cur_batch_q_len), - b.cmp_lt(op_qh, b.const_i32(NUM_QH)), + b.cmp_lt(op_qh, num_qh_p), ) op_mask = ( b.land(_op_mask_inner, b.cmp_lt(OUT_ROW_BASE, b.const_i32(VALID_ROWS))) @@ -5817,8 +5834,6 @@ def _build_gfx942_4warp_gqa_lean( raise NotImplementedError( "4-warp GQA lean body is causal-only (no sliding-window mask)" ) - H = spec.num_query_heads - HKV = spec.num_kv_heads GQAG = spec.num_queries_per_kv BS = spec.block_size BN = 64 # 4-warp: 64-key tiles @@ -5829,7 +5844,6 @@ def _build_gfx942_4warp_gqa_lean( NK = HD // K NDdim = HD // 32 NKpv = BN // K - ITERS = spec.binary_search_iters b = IRBuilder(spec.kernel_name() + "_4wgqa") b.kernel.attrs["max_workgroup_size"] = 256 # 4 wave64 warps @@ -5868,6 +5882,14 @@ def _build_gfx942_4warp_gqa_lean( b.param("softcap", F32) num_seqs_p = b.param("num_seqs", I32) bt_stride_p = b.param("block_table_stride", I32) + # Unused here; declared so every later param sits where the shared 2D + # launch signature packs it. + b.param("qq_bias_stride_0", I32) + # Shape values this body reads as kernargs (``RUNTIME_KERNARG_FIELDS``), + # after every fixed param so the shared launch signature can append them. + rt = {f: b.param(f, I32) for f in spec.RUNTIME_KERNARG_FIELDS} + hq_p = rt["num_query_heads"] + hkv_p = rt["num_kv_heads"] tid = b.thread_id_x() wid = b.div(tid, b.const_i32(64)) @@ -5879,7 +5901,14 @@ def _build_gfx942_4warp_gqa_lean( gqb = b.block_id_y() # grid.y = total_q_blocks (BLOCK_M=128) # in-kernel CTA->seq via binary_search on cu_seqlens_q (mirrors :5590) - sid = binary_search_seq_idx(b, CUQ, gqb, num_seqs_p, block_q=128, iterations=ITERS) + sid = binary_search_seq_idx( + b, + CUQ, + gqb, + num_seqs_p, + block_q=128, + iterations=binary_search_seq_idx_trip_count(b, num_seqs_p), + ) cu_q_start = b.global_load_i32(CUQ, sid) cu_q_stop = b.global_load_i32(CUQ, b.add(sid, b.const_i32(1))) qlen = b.sub(cu_q_stop, cu_q_start) @@ -5894,8 +5923,10 @@ def _build_gfx942_4warp_gqa_lean( b.ret() V_lds = b.smem_alloc(dtype, [BN, HD], name_hint="Vlds") - q_desc = TensorDescriptor.naive( - "query_ptr", lengths=[1 << 30, H, HD], coord_names=("token", "head", "dim") + q_desc = DynamicTensorDescriptor.create( + "query_ptr", + coord_names=("token", "head", "dim"), + strides=(b.mul(hq_p, b.const_i32(HD)), HD, 1), ) def phys_key(kv, keytile): @@ -5934,9 +5965,7 @@ def bperm(v): key = b.div(lin, b.const_i32(HD)) hd = b.mod(lin, b.const_i32(HD)) pk = phys_key(kv, key) - velem = b.add( - b.mul(b.add(b.mul(pk, b.const_i32(HKV)), kvh), b.const_i32(HD)), hd - ) + velem = b.add(b.mul(b.add(b.mul(pk, hkv_p), kvh), b.const_i32(HD)), hd) b.smem_store_vN( V_lds, [key, hd], b.global_load_vN(Vp, velem, dtype, 8, align=16), 8 ) @@ -5955,9 +5984,7 @@ def bperm(v): q = b.global_load_vN(Q, q_off, dtype, BPL, align=BPL * 2) for kt in range(NKEYT): kelem = b.add( - b.mul( - b.add(b.mul(pk_kt[kt], b.const_i32(HKV)), kvh), b.const_i32(HD) - ), + b.mul(b.add(b.mul(pk_kt[kt], hkv_p), kvh), b.const_i32(HD)), koff, ) kf = b.global_load_vN(Kp, kelem, dtype, APL, align=APL * 2) @@ -6039,7 +6066,7 @@ def bperm(v): r, c = at.lane_to_output(b, lane, i) dim = b.add(b.mul(b.const_i32(nt), b.const_i32(32)), r) q_inseq = b.add(qbase, b.add(wq, c)) - oi = b.add(b.mul(b.add(qstart, b.add(wq, c)), b.const_i32(H)), qhead) + oi = b.add(b.mul(b.add(qstart, b.add(wq, c)), hq_p), qhead) val = b.cast_f32_to(b.fmul(b.vec_extract(accs_f[nt], i), recip), dtype) with b.scf_if(b.cmp_lt(q_inseq, qlen)): b.global_store(C, b.add(b.mul(oi, b.const_i32(HD)), dim), val, align=2) @@ -6072,8 +6099,6 @@ def build_gfx942_4warp_gqa( raise NotImplementedError("4-warp GQA kernel supports head_size in {128, 256}") if HD == _4WGQA_LEAN_HEAD_SIZE: return _build_gfx942_4warp_gqa_lean(spec, arch=arch) - H = spec.num_query_heads - HKV = spec.num_kv_heads GQAG = spec.num_queries_per_kv BS = spec.block_size HD128 = HD == 128 @@ -6093,7 +6118,6 @@ def build_gfx942_4warp_gqa( BN * HD ) // 256 # V-fill: elements per thread (256 = 4 wave64 threads) v_fill_nloads = v_fill_epw // 8 # 8-wide (dwordx4) loads per thread - ITERS = spec.binary_search_iters from ..common.attention_unified import gfx942_gqa_fold_eligible # GQA head-fold is the DEFAULT for the qualifying cohort (D128, GQA 4:1, sliding- @@ -6157,6 +6181,14 @@ def build_gfx942_4warp_gqa( b.param("softcap", F32) num_seqs_p = b.param("num_seqs", I32) bt_stride_p = b.param("block_table_stride", I32) + # Unused here; declared so every later param sits where the shared 2D + # launch signature packs it. + b.param("qq_bias_stride_0", I32) + # Shape values this body reads as kernargs (``RUNTIME_KERNARG_FIELDS``), + # after every fixed param so the shared launch signature can append them. + rt = {f: b.param(f, I32) for f in spec.RUNTIME_KERNARG_FIELDS} + hq_p = rt["num_query_heads"] + hkv_p = rt["num_kv_heads"] tid = b.thread_id_x() wid = b.div(tid, b.const_i32(64)) @@ -6175,7 +6207,12 @@ def build_gfx942_4warp_gqa( # in-kernel CTA->seq via binary_search on cu_seqlens_q (mirrors :5590) sid = binary_search_seq_idx( - b, CUQ, gqb, num_seqs_p, block_q=TOKBLK, iterations=ITERS + b, + CUQ, + gqb, + num_seqs_p, + block_q=TOKBLK, + iterations=binary_search_seq_idx_trip_count(b, num_seqs_p), ) cu_q_start = b.global_load_i32(CUQ, sid) cu_q_stop = b.global_load_i32(CUQ, b.add(sid, b.const_i32(1))) @@ -6221,8 +6258,10 @@ def build_gfx942_4warp_gqa( f"(V_lds {_v_lds_bytes} + K_lds {_k_lds_bytes}, HD={HD}, " f"HD128_PIPE={HD128_PIPE}) exceeds the {arch} {_lds_cap} B LDS capacity" ) - q_desc = TensorDescriptor.naive( - "query_ptr", lengths=[1 << 30, H, HD], coord_names=("token", "head", "dim") + q_desc = DynamicTensorDescriptor.create( + "query_ptr", + coord_names=("token", "head", "dim"), + strides=(b.mul(hq_p, b.const_i32(HD)), HD, 1), ) def phys_key(kv, keytile): @@ -6246,6 +6285,7 @@ def bperm(v): ] context_off = b.sub(klen, qlen) # prefix in KV cache (qlen!=klen: chunked/decode) window = int(spec.sliding_window) # 0 = causal; >0 = SWA (keep dist < window) + window_p = rt["sliding_window"] # the size itself is a kernarg causal_t = b.div( b.add(b.add(context_off, qbase), b.const_i32(TOKBLK + BN - 1)), b.const_i32(BN) ) @@ -6253,7 +6293,7 @@ def bperm(v): kvend = b.select(b.cmp_lt(causal_t, klen_t), causal_t, klen_t) # E6: windowed KV-loop start - skip tiles fully before the window (bounded work) if window > 0: - _ks = b.sub(b.add(context_off, qbase), b.const_i32(window)) + _ks = b.sub(b.add(context_off, qbase), window_p) _ks = b.select(b.cmp_gt(_ks, b.const_i32(0)), _ks, b.const_i32(0)) kvstart = b.div(_ks, b.const_i32(BN)) else: @@ -6275,7 +6315,7 @@ def bperm(v): int_end = b.select(b.cmp_lt(int_end_c, int_end_r), int_end_c, int_end_r) if window > 0: _sn = b.sub( - b.add(q_blk_hi, b.const_i32(1)), b.const_i32(window) + b.add(q_blk_hi, b.const_i32(1)), window_p ) # window-interior: kv*BN>=q_blk_hi-window+1 _cs = b.div(b.add(_sn, b.const_i32(BN - 1)), b.const_i32(BN)) # ceil int_start = b.select(b.cmp_lt(_sn, b.const_i32(0)), kvstart, _cs) @@ -6334,9 +6374,7 @@ def fill_tile(tkv, buf_off): key = b.div(lin, b.const_i32(HD)) hd = b.mod(lin, b.const_i32(HD)) pk = phys_key(tkv, key) - velem = b.add( - b.mul(b.add(b.mul(pk, b.const_i32(HKV)), kvh), b.const_i32(HD)), hd - ) + velem = b.add(b.mul(b.add(b.mul(pk, hkv_p), kvh), b.const_i32(HD)), hd) row = b.add(buf_off, key) b.smem_store_vN( V_lds, @@ -6361,9 +6399,7 @@ def fill_load( key = b.div(lin, b.const_i32(HD)) hd = b.mod(lin, b.const_i32(HD)) pk = phys_key(tkv, key) - velem = b.add( - b.mul(b.add(b.mul(pk, b.const_i32(HKV)), kvh), b.const_i32(HD)), hd - ) + velem = b.add(b.mul(b.add(b.mul(pk, hkv_p), kvh), b.const_i32(HD)), hd) vr = b.global_load_vN(Vp, velem, dtype, 8, align=16) kr = b.global_load_vN(Kp, velem, dtype, 8, align=16) pf.append((vr, kr, key, hd)) @@ -6434,7 +6470,7 @@ def body(kv, carry, masked): else: kelem = b.add( b.mul( - b.add(b.mul(pk_kt[kt], b.const_i32(HKV)), kvh), + b.add(b.mul(pk_kt[kt], hkv_p), kvh), b.const_i32(HD), ), koff, @@ -6465,7 +6501,7 @@ def body(kv, carry, masked): mask_cond = b.lor(b.cmp_gt(key_g, q_g), b.cmp_ge(key_g, klen)) if window > 0: mask_cond = b.lor( - mask_cond, b.cmp_ge(b.sub(q_g, key_g), b.const_i32(window)) + mask_cond, b.cmp_ge(b.sub(q_g, key_g), window_p) ) Sm[kt][i] = b.select(mask_cond, ninf, raw) else: @@ -6589,12 +6625,12 @@ def run_loop(start, end, init_iters, masked, iv): _hl = b.mod(_m, b.const_i32(FOLD_HEADS)) q_inseq = b.add(qbase, _tl) oi = b.add( - b.mul(b.add(qstart, _tl), b.const_i32(H)), + b.mul(b.add(qstart, _tl), hq_p), b.add(b.mul(kvh, b.const_i32(GQAG)), _hl), ) else: q_inseq = b.add(qbase, b.add(wq, c)) - oi = b.add(b.mul(b.add(qstart, b.add(wq, c)), b.const_i32(H)), qhead) + oi = b.add(b.mul(b.add(qstart, b.add(wq, c)), hq_p), qhead) val = b.cast_f32_to(b.fmul(b.vec_extract(accs_f[nt], i), recip), dtype) with b.scf_if(b.cmp_lt(q_inseq, qlen)): b.global_store(C, b.add(b.mul(oi, b.const_i32(HD)), dim), val, align=2) diff --git a/dnn-providers/hip-kernel-provider/rocke/library/kernels/gfx942/attention_tiled_3d.py b/dnn-providers/hip-kernel-provider/rocke/library/kernels/gfx942/attention_tiled_3d.py index a3632b712d29..5b3312e4a82f 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/kernels/gfx942/attention_tiled_3d.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/kernels/gfx942/attention_tiled_3d.py @@ -23,7 +23,8 @@ gfx942 2D kernel uses (gfx942-legal), via ``dequant_fp8x8_to_dtype``. The reduce kernel is pure f32 load / exp2 / store (no MFMA, no transpose read), -so it is arch-neutral and is a byte-for-byte port of the gfx950 reduce kernel. +so it is arch-neutral. It is a port of the gfx950 reduce kernel that reads the +query head count as a kernarg instead of baking it. See the gfx950 module docstring for the split-KV grid / workspace layout. """ @@ -31,7 +32,7 @@ from __future__ import annotations from dataclasses import dataclass, field -from typing import Optional, Tuple +from typing import ClassVar, Optional, Tuple from kernels.common._attention_strided_kv import ( declare_strided_kv_params, @@ -60,14 +61,19 @@ from rocke.helpers.attention import ( apply_softcap_log2 as _apply_softcap, binary_search_seq_idx as _binary_search_seq_idx_helper, - binary_search_seq_idx_iters, + binary_search_seq_idx_trip_count, dequant_fp8x8_to_dtype, mfma_16x16x16_for_dtype as _mfma_16x16x16, wave64_reduce_max as _wave64_reduce_max, wave64_reduce_sum as _wave64_reduce_sum, ) from rocke.helpers.distribution import make_static_tile_distribution -from rocke.helpers.transforms import TensorDescriptor, embed, indirect, unmerge +from rocke.helpers.transforms import ( + DynamicTensorDescriptor, + embed, + indirect, + unmerge, +) MFMA_M = 16 MFMA_N = 16 @@ -103,6 +109,18 @@ class UnifiedAttention3DTiledSpec: dispatcher and harnesses can construct it identically. """ + # Problem fields this body reads as trailing i32 kernargs instead of baking + # them, in ABI order. The launch signature, the argument pack and the cache + # key read the same tuple (``_runtime_kernarg_fields``). + RUNTIME_KERNARG_FIELDS: ClassVar[Tuple[str, ...]] = ( + "num_query_heads", + "num_kv_heads", + "sliding_window", + ) + # The cu_q search derives its trip count from the num_seqs kernarg, so + # num_seqs never reaches the IR and one binary serves every batch. + RUNTIME_SEARCH_ITERS: ClassVar[bool] = True + head_size: int block_size: int num_query_heads: int @@ -229,10 +247,6 @@ def tile_size(self) -> int: def dtype_ir(self) -> Type: return F16 if self.dtype == "fp16" else BF16 - @property - def binary_search_iters(self) -> int: - return binary_search_seq_idx_iters(self.num_seqs) - def kernel_name(self) -> str: from rocke.helpers.spec import kernel_name_join @@ -240,13 +254,15 @@ def kernel_name(self) -> str: "rocke_uattn3d_tiled_gfx942", f"d{self.head_size}", f"b{self.block_size}", - f"h{self.num_query_heads}kv{self.num_kv_heads}", + # Head counts are kernargs; only the GQA ratio is baked. + f"g{self.num_queries_per_kv}", f"seg{self.num_segments}", self.dtype, f"kv{self.kv_storage_dtype}" if self.kv_storage_dtype else "", "fnuz" if self.kv_storage_dtype else "", "sinks" if self.use_sinks else "", - f"sw{self.sliding_window}" if self.sliding_window > 0 else "", + # The window size is a kernarg; only whether the body masks one is baked. + "sw" if self.sliding_window > 0 else "", "softcap" if self.has_softcap else "", "alibi" if self.use_alibi else "", "qqb" if self.use_qq_bias else "", @@ -360,8 +376,6 @@ def build_unified_attention_3d_tiled( BLOCK_M = spec.block_m BLOCK_Q = spec.block_q NQK = spec.num_queries_per_kv - NUM_KV = spec.num_kv_heads - NUM_QH = spec.num_query_heads NUM_SEG = spec.num_segments SLIDING_WINDOW = spec.sliding_window USE_SOFTCAP = spec.has_softcap @@ -445,6 +459,11 @@ def build_unified_attention_3d_tiled( num_seqs_p = b.param("num_seqs", I32) bt_stride_p = b.param("block_table_stride", I32) qq_bias_stride0_p = b.param("qq_bias_stride_0", I32) + # Shape values this body reads as kernargs (``RUNTIME_KERNARG_FIELDS``), + # after every fixed param so the shared launch signature can append them. + rt = {f: b.param(f, I32) for f in spec.RUNTIME_KERNARG_FIELDS} + num_qh_p = rt["num_query_heads"] + num_kv_p = rt["num_kv_heads"] if spec.kv_layout == "strided": k_strides = declare_strided_kv_params(b, "k") v_strides = declare_strided_kv_params(b, "v") @@ -464,7 +483,7 @@ def build_unified_attention_3d_tiled( q_block_global_idx, num_seqs_p, block_q=BLOCK_Q, - iterations=spec.binary_search_iters, + iterations=binary_search_seq_idx_trip_count(b, num_seqs_p), ) cu_q_start = b.global_load_i32(cu_q, seq_idx) cu_q_stop = b.global_load_i32(cu_q, b.add(seq_idx, b.const_i32(1))) @@ -484,20 +503,21 @@ def build_unified_attention_3d_tiled( # tiles_per_segment = cdiv(seq_len, NUM_SEG * T) tps = b.div(b.add(seq_len, b.const_i32(NUM_SEG * T - 1)), b.const_i32(NUM_SEG * T)) - ml_desc = TensorDescriptor.naive( + # The head count is a kernarg, so each token stride is too. + ml_desc = DynamicTensorDescriptor.create( "segm_ml", - lengths=[1 << 30, NUM_QH, NUM_SEG], coord_names=("token", "head", "seg"), + strides=(b.mul(num_qh_p, b.const_i32(NUM_SEG)), NUM_SEG, 1), ) - seg_acc_desc = TensorDescriptor.naive( + seg_acc_desc = DynamicTensorDescriptor.create( "segm_output", - lengths=[1 << 30, NUM_QH, NUM_SEG, HD], coord_names=("token", "head", "seg", "dim"), + strides=(b.mul(num_qh_p, b.const_i32(NUM_SEG * HD)), NUM_SEG * HD, HD, 1), ) - q_desc = TensorDescriptor.naive( + q_desc = DynamicTensorDescriptor.create( "Q", - lengths=[1 << 30, NUM_QH, HD], coord_names=("token", "head", "dim"), + strides=(b.mul(num_qh_p, b.const_i32(HD)), HD, 1), ) seg_start_tile_pos = b.mul(b.mul(seg_idx, tps), b.const_i32(T)) @@ -511,9 +531,7 @@ def build_unified_attention_3d_tiled( qh_r = b.add( b.mul(kv_head_idx, b.const_i32(NQK)), b.mod(row, b.const_i32(NQK)) ) - row_ok = b.land( - b.cmp_lt(qp_r, cur_batch_q_len), b.cmp_lt(qh_r, b.const_i32(NUM_QH)) - ) + row_ok = b.land(b.cmp_lt(qp_r, cur_batch_q_len), b.cmp_lt(qh_r, num_qh_p)) qp_r_safe = b.select(row_ok, qp_r, b.const_i32(0)) qh_r_safe = b.select(row_ok, qh_r, b.const_i32(0)) qtoken = b.add(cu_q_start, qp_r_safe) @@ -531,7 +549,7 @@ def build_unified_attention_3d_tiled( b.mul(kv_head_idx, b.const_i32(NQK)), b.mod(row, b.const_i32(NQK)) ) row_ok = b.land( - b.cmp_lt(qp_r, cur_batch_q_len), b.cmp_lt(qh_r, b.const_i32(NUM_QH)) + b.cmp_lt(qp_r, cur_batch_q_len), b.cmp_lt(qh_r, num_qh_p) ) qp_r_safe = b.select(row_ok, qp_r, b.const_i32(0)) qh_r_safe = b.select(row_ok, qh_r, b.const_i32(0)) @@ -568,7 +586,7 @@ def _row_reduce(x: Value, op) -> Value: one_f = b.const_f32(1.0) rcp_ln2 = b.const_f32(1.4426950408889634) qk_scale = b.fmul(scale_p, rcp_ln2) - sw_const = b.const_i32(int(SLIDING_WINDOW)) + sw_const = rt["sliding_window"] z8 = b.zero_vec(dtype, 8) # ---------------- Q -> LDS ---------------- @@ -583,9 +601,7 @@ def _row_reduce(x: Value, op) -> Value: b.mul(kv_head_idx, b.const_i32(NQK)), b.mod(Q_row, b.const_i32(NQK)), ) - qmask_t = b.land( - b.cmp_lt(q_pos_t, cur_batch_q_len), b.cmp_lt(qh_t, b.const_i32(NUM_QH)) - ) + qmask_t = b.land(b.cmp_lt(q_pos_t, cur_batch_q_len), b.cmp_lt(qh_t, num_qh_p)) q_pos_safe = b.select(qmask_t, q_pos_t, b.const_i32(0)) qh_safe = b.select(qmask_t, qh_t, b.const_i32(0)) q_off_base, _ = q_desc.offset( @@ -628,7 +644,7 @@ def _row_reduce(x: Value, op) -> Value: qh = b.add( b.mul(kv_head_idx, b.const_i32(NQK)), b.mod(lane_col, b.const_i32(NQK)) ) - qh_in = b.cmp_lt(qh, b.const_i32(NUM_QH)) + qh_in = b.cmp_lt(qh, num_qh_p) sink_h = b.global_load(sinks, qh, dtype, align=2) sink_f = b.fmul(b.cast_to_f32(sink_h), rcp_ln2) sink_with_mask = b.select(qh_in, sink_f, neg_inf) @@ -667,8 +683,9 @@ def _row_reduce(x: Value, op) -> Value: assert (T * HD) % KV_HALVES_PER_CALL == 0 kv_calls_per_tile = (T * HD) // KV_HALVES_PER_CALL bytes_per_call = KV_HALVES_PER_CALL * 2 - kv_stride_blk_b = BS * NUM_KV * HD * KV_BYTES - kv_stride_tok_b = NUM_KV * HD * KV_BYTES + # The token and page strides scale with the runtime kv head count. + kv_stride_tok_b = b.mul(num_kv_p, b.const_i32(HD * KV_BYTES)) + kv_stride_blk_b = b.mul(kv_stride_tok_b, b.const_i32(BS)) kv_stride_h_b = HD * KV_BYTES bytes_per_buf = T * HD * 2 @@ -678,11 +695,10 @@ def _row_reduce(x: Value, op) -> Value: zero_soff = b.const_i32(0) seq_base = b.mul(seq_idx, bt_stride_p) if spec.kv_layout == "paged" else None - _kv_base = TensorDescriptor.naive( + _kv_base = DynamicTensorDescriptor.create( "paged_kv_bytes", - lengths=[1 << 24, BS, NUM_KV, HD], - strides=[kv_stride_blk_b, kv_stride_tok_b, kv_stride_h_b, KV_BYTES], coord_names=("physical_block", "token", "kv_head", "dim"), + strides=(kv_stride_blk_b, kv_stride_tok_b, kv_stride_h_b, KV_BYTES), ) if T == BS: paged_kv_desc = _kv_base.transform( @@ -742,7 +758,7 @@ def _issue_k_load( b.add( b.mul( b.div(src_half, b.const_i32(HD)), - b.const_i32(kv_stride_tok_b), + kv_stride_tok_b, ), b.mul(kv_head_idx, b.const_i32(kv_stride_h_b)), ), @@ -793,7 +809,7 @@ def _issue_v_load( b.add( b.mul( b.div(src_half, b.const_i32(HD)), - b.const_i32(kv_stride_tok_b), + kv_stride_tok_b, ), b.mul(kv_head_idx, b.const_i32(kv_stride_h_b)), ), @@ -861,15 +877,15 @@ def _tile_page(kv_tile_idx: Value) -> Value: return b.div(kv_tile_idx, b.const_i32(TILES_PER_PAGE)) def _tile_base( - page_id: Value, kv_tile_idx: Value, page_stride: int, token_stride: int + page_id: Value, kv_tile_idx: Value, page_stride: Value, token_stride: Value ) -> Value: # Offset of a tile's first token: the page's, plus the tile's rows # before it in the page (sub-page tiles). - base = b.mul(page_id, b.const_i32(page_stride)) + base = b.mul(page_id, page_stride) if TILES_PER_PAGE == 1: return base in_page = b.mod(kv_tile_idx, b.const_i32(TILES_PER_PAGE)) - return b.add(base, b.mul(in_page, b.const_i32(T * token_stride))) + return b.add(base, b.mul(in_page, b.mul(token_stride, b.const_i32(T)))) def _issue_wide_load( src, lds, kv_tile_idx: Value, buf_idx: Value, page_id: Value | None @@ -882,11 +898,14 @@ def _issue_wide_load( if page_id is not None: # The tile base plus the chunk's offset within the page, which # is loop-invariant: no page-table lookup per copy. + # Element strides (the wide loader indexes halves). + tok_e = b.mul(num_kv_p, b.const_i32(HD)) + blk_e = b.mul(tok_e, b.const_i32(BS)) elem_off = b.add( - _tile_base(page_id, kv_tile_idx, BS * NUM_KV * HD, NUM_KV * HD), + _tile_base(page_id, kv_tile_idx, blk_e, tok_e), b.add( b.add( - b.mul(row, b.const_i32(NUM_KV * HD)), + b.mul(row, tok_e), b.mul(kv_head_idx, b.const_i32(HD)), ), col, @@ -1052,13 +1071,11 @@ def _refill_page_id_window(page_id_window: Value, kv_tile_idx: Value) -> Value: qh_r = b.add( b.mul(kv_head_idx, b.const_i32(NQK)), b.mod(lane_col, b.const_i32(NQK)) ) - row_ok = b.land( - b.cmp_lt(qp_r, cur_batch_q_len), b.cmp_lt(qh_r, b.const_i32(NUM_QH)) - ) + row_ok = b.land(b.cmp_lt(qp_r, cur_batch_q_len), b.cmp_lt(qh_r, num_qh_p)) causal_lim = b.add(context_len, qp_r) if USE_ALIBI: - qh_ok = b.cmp_lt(qh_r, b.const_i32(NUM_QH)) + qh_ok = b.cmp_lt(qh_r, num_qh_p) alibi_slope = b.masked_global_load( alibi_slopes_ptr, qh_r, qh_ok, b.const_f32(0.0), dtype=F32, align=4 ) @@ -1261,12 +1278,17 @@ def _refill_page_id_window(page_id_window: Value, kv_tile_idx: Value) -> Value: # --------------------------------------------------------------------------- -# Reduce kernel (arch-neutral; byte-for-byte port of the gfx950 reduce kernel) +# Reduce kernel (arch-neutral; the gfx950 reduce with a runtime head count) # --------------------------------------------------------------------------- @dataclass(frozen=True) class UnifiedAttentionReduceTiledSpec: + # Problem fields this body reads as trailing i32 kernargs instead of baking + # them, in ABI order. The launch signature, the argument pack and the cache + # key read the same tuple (``_runtime_kernarg_fields``). + RUNTIME_KERNARG_FIELDS: ClassVar[Tuple[str, ...]] = ("num_query_heads",) + head_size: int num_query_heads: int num_kv_heads: int @@ -1297,7 +1319,6 @@ def kernel_name(self) -> str: parts = [ f"d{self.head_size}", - f"h{self.num_query_heads}", f"seg{self.num_segments}", self.dtype, ] @@ -1314,7 +1335,6 @@ def build_unified_attention_reduce_tiled( """Combine the per-segment partial state into the final fp16/bf16 output.""" HD = spec.head_size NUM_SEG = spec.num_segments - NUM_QH = spec.num_query_heads dtype = spec.dtype_ir THREADS = 64 @@ -1335,6 +1355,8 @@ def build_unified_attention_reduce_tiled( seg_max = b.param("segm_max_ptr", PtrType(F32, "global"), readonly=True, align=4) seg_l = b.param("segm_expsum_ptr", PtrType(F32, "global"), readonly=True, align=4) _seq_lens = b.param("seq_lens_ptr", PtrType(I32, "global"), readonly=True, align=4) + rt = {f: b.param(f, I32) for f in spec.RUNTIME_KERNARG_FIELDS} + num_qh_p = rt["num_query_heads"] lse = ( b.param( "lse_ptr", @@ -1354,20 +1376,21 @@ def build_unified_attention_reduce_tiled( neg_inf = b.const_f32(float("-inf")) zero_f = b.const_f32(0.0) - ml_desc_red = TensorDescriptor.naive( + # The head count is a kernarg, so each token stride is too. + ml_desc_red = DynamicTensorDescriptor.create( "segm_ml", - lengths=[1 << 30, NUM_QH, NUM_SEG], coord_names=("token", "head", "seg"), + strides=(b.mul(num_qh_p, b.const_i32(NUM_SEG)), NUM_SEG, 1), ) - seg_acc_desc_red = TensorDescriptor.naive( + seg_acc_desc_red = DynamicTensorDescriptor.create( "segm_output", - lengths=[1 << 30, NUM_QH, NUM_SEG, HD], coord_names=("token", "head", "seg", "dim"), + strides=(b.mul(num_qh_p, b.const_i32(NUM_SEG * HD)), NUM_SEG * HD, HD, 1), ) - out_desc_red = TensorDescriptor.naive( + out_desc_red = DynamicTensorDescriptor.create( "out", - lengths=[1 << 30, NUM_QH, HD], coord_names=("token", "head", "dim"), + strides=(b.mul(num_qh_p, b.const_i32(HD)), HD, 1), ) base_ml, _ = ml_desc_red.offset(b, token=q_token, head=q_head, seg=b.const_i32(0)) @@ -1459,7 +1482,7 @@ def build_unified_attention_reduce_tiled( # scale*log2e); lse_store converts to natural log and writes -inf for a # row with no mass (every segment empty), whose O is already 0. with b.scf_if(b.cmp_eq(tid, b.const_i32(0))): - lse_index = b.add(b.mul(q_token, b.const_i32(NUM_QH)), q_head) + lse_index = b.add(b.mul(q_token, num_qh_p), q_head) lse_store( b, lse, diff --git a/dnn-providers/hip-kernel-provider/rocke/library/tests/golden/attention_strided_kv_ir_sha256.json b/dnn-providers/hip-kernel-provider/rocke/library/tests/golden/attention_strided_kv_ir_sha256.json index a896280896bf..995dc1c6b744 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/tests/golden/attention_strided_kv_ir_sha256.json +++ b/dnn-providers/hip-kernel-provider/rocke/library/tests/golden/attention_strided_kv_ir_sha256.json @@ -1,50 +1,50 @@ { "llvm20": { - "0": "1626048eeaface4bf6eb73c5bc0c81a653da20fc18ca502898a55e6111d13fb2", - "1": "a9d06552869ca1485cb3a7b8a362d2a448fb8d55220f222d50d4a54ded0a2488", - "2": "8c0cb312ca581170ce8e387a94a1b77ee52b4c7fc233aa5b15f65f0c707eea8b", + "0": "8cfb54512441a845a2ae325065291e39e5367be96f320596e8622d7e308ec610", + "1": "b1cee01d468e92ecb5b8571323ffd897bff01c7d2e42dc8ba50814387f908a63", + "2": "a80d03691e079aacfb8653f7054a836d2ecd57256285274fa51c52a8a398c727", "3": "272f3227c873b525dc8e48c4f024f7885de375178e20f4c4d6686b4724d82c0b", "4": "b25f6bfc84f5db8903883a85671f13b109256697e33c67877686b4d90ac07606", "5": "bde0dfb3994fb2a89081fa05809323059111a5e7d4df55ed8859ec2f6b0cd930", - "6": "db68f7a59b6a018e310c3d86d065432fbc4a7201f3d914114dec2d58d3200a19", - "7": "e15ef2766af1f7f39f1e2e1e6f1602e17ed832ed960fee0fc87c54447742c45f", - "8": "15c5cc154e40e4e5f099d1bf6e476e283a70c450824db3ebb9df927b2f7aa2e0", + "6": "23c2881b225837148b997823910e598ad844d781441a1f122a5d370b4b97900c", + "7": "966f73aef9172bbe60348f970553576d609299131e029e6842e1f4dc49c96da1", + "8": "c5f4baa4b7f0237ce60bf63f6694fcca24bdd440c0fbc34a1dc0a6c8947b791c", "9": "7866c9eb6d414e0edf3678a477c387e2742057eed6a96616544f29057d609d0a", "10": "05d6f2c670f9c66f59974f7f2715eed844799ac869e2765d3d1f691bc417d2da", "11": "68530d1c8c9c9a12300743809054235d88e2a1698f811bbd86a531eb47113d40", - "12": "fa6ae13282f4e5748eae40227d1355529e793d2306f210136a498b6621fcff96", + "12": "6adc12acb6dd75f05927ef1e4f6683b85c780abae25d76588bf57ae25f8dfd88", "13": "d7982e56673824863076878c3461209f8557cfa8507b102071eae22d712b6414" }, "llvm22": { - "0": "bcef03e3b5555c8bbaaa03188901c1dfc8bee0cdfa80da32026367e82e7c727b", - "1": "9becf0d4f1b8eedc388e818afb01a40607a27efe2d415772ba1f6352c2ba35d7", - "2": "7fc95348ccfd43e6afef1b5bc83b0e17b55c0879d5d8c752d40e6f631d0b9c77", + "0": "14ccc0e86cb9a96163a8c32812cd718dc454f726dacced64fa3cf07db5c70c90", + "1": "62086a1cb544ba32772dc112e252813ffff1961bc2db23401bf0003bce718b73", + "2": "f2eb0fd45eb3693671f6d11bdce0eac92295e784e74efffe21351150e2807039", "3": "c88f5e1109c70e0c4f132b05578c6abf440d04bccb52bd774b492d67d3acc9bb", "4": "3e5522113635114b994bb40cf19bc51a417765ac827cbfa879acb45cb0e62796", "5": "3c03d60f49863bc856426e1431cb98e8e682da4b6818f1df320536fab12b1314", - "6": "2437aa43c1700bf2af3731b08d30869d68cf9110b2fbed6114fae27b78a4e0f0", - "7": "b1503cc02eed262c4d08b02d3802a16f3ee74ff0700f756e79e9a1c17580ff71", - "8": "8942a331f8dd8285230d4a53c02bca59ee5158284edaf12d7cbbd50900a5bfe5", + "6": "b4720021fe24f67cfdd7b6854a25f4b371cd5c086d7d804e4a1a451d8f885c82", + "7": "281cf49847fec9bf2f0133d0b0f83c9ff72e99e3f9bfd846d2d056289604609b", + "8": "8e1697deca33346efb263397db405b14b835a617e0fb8a9d6e5dd19958d767e4", "9": "43cbf9fc03696e3fcf0a132e31a769d92cd48267eaea3526e2cc7c46c35225da", "10": "35992a4d15f09b54bd668d4b21556a2c871fa8c6f8aa39a89169736e303476ea", "11": "ec52f6ee37cff049c8b81eb579b71c32be965bb4730fc7b30bd477196499727f", - "12": "002dfd5d5ce7d67b6992178da77d629eb26e80c5343c95a16a0e2f89462cf752", + "12": "b543ef77d6fd6bcdeab9b00f9c7dd8d9823777987f2e9e09b79195f3d230be37", "13": "f71e2da33c81aacfd3a6690c0fa2c540aa3a83586e0aa8ed0f1ed638023677db" }, "llvm23": { - "0": "67e7f5bd0d8861dc16e6cff53316d780d40f6232716a5a3ea25f06055ac48442", - "1": "a0b5229967c3892d492cbacb9db2d15f06ce919b1f15c8de8d298e0d74306e46", - "2": "4cded6921b9951b733d1dbfdfe93015aa7b78c5bd5def3dfbe503175949af968", + "0": "7ce6e33c41e2752f1863c4ef79eef519c4748cc9f1c3035cab5cab030f386d55", + "1": "38fa323beea8da252380810dbbe76686f3cff17bc15aca23c32a73c11f6a70e4", + "2": "5e39245a8752ce0229c9229aba2b11fc8cb07241e93ebaafa051cae24aab4514", "3": "393df7f6bb892e49b3c3031ff696d11addb1da00ad2088aa881ea37074c7db77", "4": "db62e65f3aa74588eb95d2bf69935aa6331d0a57d42996bbe81c7510e146f119", "5": "ee8248f21f5631e005750601f275daf2fd57d87f102a241c2fee76caa4567d33", - "6": "92c9560cfd6ae383a059972d0f4863e3699b56beba16aab1736d06bb41946143", - "7": "1b3cb6e2c817fce85c1bd189f7a7a01eb7e089c37ef1824cabdfbe5bd6915b47", - "8": "4b0b553b289d0de60dd9401f61f97d2f36cb201b6343437e28e5fbdc370cf529", + "6": "b3add7e87f7a74a750718e45156d80bc9a5725c33a8a25fadc55b71bb600af89", + "7": "94ea80d639a380a3cf0069ff8679e6f197fa1c40c9a4885413d747341964ef9c", + "8": "d8a78f5237a1c466494c27808c41590153f9ae47d899a5d6ad240510e680d89f", "9": "cb2426e5ff5ce00a3094ea4217df160031006f22cfe058568e4c959f1ecff086", "10": "400f75d835f4bb4d37dfe0039c609deb5e172c77e5a83b87b89071ed3b0017bf", "11": "0e441b0ae432165d00da483191b9470dcfc043c2bacd386b86fe6bef8f05266f", - "12": "03eb993466369d2ca4aabaff7880f23138bf08a0f71154159a0c9e4627e17df3", + "12": "0ec51a3f92220ca3dad8979ef3dbe976f89ab95e28de5a178eaee7477661b8d4", "13": "74d14799531a8e079e822a46af3a22f1875a85a46c6c561ef26427fd44ffd748" } } diff --git a/dnn-providers/hip-kernel-provider/rocke/library/tests/test_attention_builds.py b/dnn-providers/hip-kernel-provider/rocke/library/tests/test_attention_builds.py index 26c7d1d35c7b..a87bda7814fc 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/tests/test_attention_builds.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/tests/test_attention_builds.py @@ -140,7 +140,7 @@ def _compile_or_skip(kernel, *, arch: str): }, # gfx942 fp16 D128 tiled-2d (shipped default geometry): 2 waves/EU caps VGPR at # 256; the kernel wants 354, so it spills ~40 B to hold 2-wave occupancy. - "rocke_uattn2d_tiled_d128_b64_h32kv8_fp16_w4_wpe2_mw32_mfma32x8_stqk_s1_mask1_hoist_mlim_kvcpall_cfvst_ksring_rd2": { + "rocke_uattn2d_tiled_d128_b64_g4_fp16_w4_wpe2_mw32_mfma32x8_stqk_s1_mask1_hoist_mlim_kvcpall_cfvst_ksring_rd2": { "max_scratch_bytes": 64, "min_waves_per_simd": 2, }, diff --git a/dnn-providers/hip-kernel-provider/rocke/library/tests/test_attention_gfx942_runtime_shape_numeric.py b/dnn-providers/hip-kernel-provider/rocke/library/tests/test_attention_gfx942_runtime_shape_numeric.py new file mode 100644 index 000000000000..275464465829 --- /dev/null +++ b/dnn-providers/hip-kernel-provider/rocke/library/tests/test_attention_gfx942_runtime_shape_numeric.py @@ -0,0 +1,299 @@ +# Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +# SPDX-License-Identifier: MIT + +"""One gfx942 tiled binary serves every head pair, window size and batch. + +The gfx942 tiled 2D, 4-warp GQA and 3D kernels read the head counts and the +sliding window as kernargs and derive the cu_q search trip count from +``num_seqs``, so shapes that differ only in those share one cache key and one +compiled binary. Each case below runs a pair of such shapes back to back +through ``run_unified_attention_torch``: the second launch must reuse the +first binary (no new compile) and both must match an FP32 paged reference. + +The second shape of each pair changes at least one of the head counts (at the +same GQA ratio), the window size, or the batch (by several log2 search +buckets), so a value that is still baked into the binary corrupts it. Only the +2D pairs change the window; the 3D pairs keep it at 0. Block-table rows are +wider than the batch needs and unused slots point at a poison block, so a wrong +stride cannot pass by luck. + +The routing and key checks run on any host; the launches need a gfx942 GPU. +""" + +from __future__ import annotations + +import itertools +import math + +import pytest + +from kernels.common import attention_unified as au + +_ARCH = "gfx942" +_CUS = 304 + + +def _problem(*, q_lens, kv_lens, hq, hkv, d, bs, dtype, window, num_blocks): + return 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=dtype, + sliding_window=window, + num_cus=_CUS, + num_kv_blocks=num_blocks, + ) + + +def _shape(q_lens, kv_lens, hq, hkv, window): + return dict(q_lens=q_lens, kv_lens=kv_lens, hq=hq, hkv=hkv, window=window) + + +# (id, path, 4-warp cohort or None, (d, bs, dtype), first shape, second shape) +_CASES = ( + ( + "2d-bf16-d128", + "2d", + None, + (128, 16, "bf16"), + _shape([300, 200, 250], [300, 450, 250], 32, 8, 0), + _shape( + [90, 130, 70, 110, 40, 150, 60, 80, 120], + [90, 300, 70, 210, 40, 150, 60, 380, 120], + 16, + 4, + 0, + ), + ), + # Head counts alone: the token and sequence counts match, so only the head + # counts tell the two launches apart. + ( + "2d-bf16-d128-heads-only", + "2d", + None, + (128, 16, "bf16"), + _shape([1025], [1025], 32, 8, 0), + _shape([1025], [1025], 64, 16, 0), + ), + ( + "2d-fp16-d64-window", + "2d", + None, + (64, 32, "fp16"), + _shape([260, 180], [400, 180], 8, 8, 128), + _shape( + [70, 90, 50, 100, 60, 80, 40, 30, 110], + [200, 90, 300, 100, 60, 250, 40, 30, 110], + 16, + 16, + 512, + ), + ), + ( + "4warp-d256", + "2d", + "gfx942_d256_causal", + (256, 16, "bf16"), + _shape([1024, 1024], [1024, 1024], 16, 2, 0), + _shape( + [300, 500, 200, 400, 250, 350, 150, 450, 280], + [300, 700, 200, 400, 250, 600, 150, 450, 280], + 32, + 4, + 0, + ), + ), + ( + "4warp-fold-window", + "2d", + "gfx942_d128_swa", + (128, 16, "bf16"), + _shape([1024], [1024], 32, 8, 512), + _shape( + [200, 300, 150, 250, 180, 220, 260, 140, 310], + [400, 300, 350, 250, 180, 520, 260, 140, 310], + 16, + 4, + 256, + ), + ), + # The decode split sizes segments by num_seqs * num_kv_heads, so the 3D pairs + # hold that product fixed; gfx942 routes windowed decode to 2D. + ( + "3d-decode", + "3d", + None, + (128, 16, "bf16"), + _shape([1, 1], [4096, 2500], 64, 16, 0), + _shape([1] * 8, [3000, 900, 4096, 1700, 2200, 3500, 1200, 2800], 16, 4, 0), + ), + ( + "3d-two-token", + "3d", + None, + (128, 16, "bf16"), + _shape([2, 2], [4096, 3100], 64, 16, 0), + _shape([2] * 8, [2500, 4096, 1900, 2600, 3100, 2200, 2400, 2700], 16, 4, 0), + ), +) + + +def _problems(geometry, shape): + d, bs, dtype = geometry + needed = sum(math.ceil(n / bs) for n in shape["kv_lens"]) + return _problem( + q_lens=shape["q_lens"], + kv_lens=shape["kv_lens"], + hq=shape["hq"], + hkv=shape["hkv"], + d=d, + bs=bs, + dtype=dtype, + window=shape["window"], + num_blocks=needed + 1, + ) + + +def _key(problem, path): + if path == "2d": + return au._tiled_cache_key(problem, _ARCH) + return au._tiled_3d_cache_key(problem, _ARCH) + + +@pytest.mark.parametrize( + "path, cohort, geometry, first, second", + [pytest.param(*c[1:], id=c[0]) for c in _CASES], +) +def test_pair_shares_one_binary(path, cohort, geometry, first, second): + """Both shapes take the intended path and route and share one key, so the + GPU test below exercises one binary on two shapes.""" + a, b = _problems(geometry, first), _problems(geometry, second) + for p in (a, b): + assert p.select_path() == path + route = au._gfx942_4warp_route(p, _ARCH) if path == "2d" else None + assert (route.cohort if route else None) == cohort + assert a.num_queries_per_kv == b.num_queries_per_kv + batch_changes = a.num_seqs != b.num_seqs + if batch_changes and not (path == "3d" and au._use_decode_grid(a, _ARCH)): + # The decode grid indexes sequences directly and emits no search. + assert au._binary_search_iters(a) != au._binary_search_iters(b) + assert _key(a, path) == _key(b, path) + + +@pytest.fixture(scope="module") +def torch(): + torch = pytest.importorskip("torch") + if not torch.cuda.is_available(): + pytest.skip("requires a gfx942 GPU") + name = torch.cuda.get_device_properties(0).gcnArchName + if not name.startswith(_ARCH): + pytest.skip(f"requires a gfx942 GPU, found {name}") + return torch + + +def _case(torch, geometry, shape, seed): + d, bs, dtype = geometry + tdt = torch.bfloat16 if dtype == "bf16" else torch.float16 + q_lens, kv_lens, hq, hkv = ( + shape["q_lens"], + shape["kv_lens"], + shape["hq"], + shape["hkv"], + ) + 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() + table = torch.full((len(kv_lens), max(needed) + 3), poison, dtype=torch.int32) + for s, n in enumerate(needed): + table[s, :n] = torch.tensor(order[:n], dtype=torch.int32) + order = order[n:] + q = torch.randn(sum(q_lens), hq, d, generator=gen).to(tdt) + k = torch.randn(num_blocks, bs, hkv, d, generator=gen).to(tdt) + v = torch.randn(num_blocks, bs, hkv, d, generator=gen).to(tdt) + k[poison] = 64.0 + v[poison] = 64.0 + cu = torch.tensor([0, *itertools.accumulate(q_lens)], dtype=torch.int32) + return q.cuda(), k.cuda(), v.cuda(), cu.cuda(), table.cuda() + + +def _reference(torch, q, k, v, cu, table, shape, d, bs): + q_lens, kv_lens, hq, hkv = ( + shape["q_lens"], + shape["kv_lens"], + shape["hq"], + shape["hkv"], + ) + window = shape["window"] + out = torch.zeros(q.shape, dtype=torch.float32, device=q.device) + 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() + keys = keys.repeat_interleave(hq // hkv, dim=1) + values = values.repeat_interleave(hq // hkv, dim=1) + start = int(cu[s]) + queries = q[start : start + q_len].float() + scores = torch.einsum("qhd,khd->hqk", queries, keys) / math.sqrt(d) + q_pos = (kv_len - q_len) + torch.arange(q_len, device=q.device)[:, None] + k_pos = torch.arange(kv_len, device=q.device)[None, :] + visible = k_pos <= q_pos + if window > 0: + visible = visible & (q_pos - k_pos < window) + scores = scores.masked_fill(~visible, float("-inf")) + out[start : start + q_len] = torch.einsum( + "hqk,khd->qhd", scores.softmax(-1), values + ) + return out + + +def _cache(path): + return au._ATTN_3D_TILED_CACHE if path == "3d" else au._ATTN_TILED_CACHE + + +def _compiled(): + return len(au._ATTN_TILED_CACHE) + len(au._ATTN_3D_TILED_CACHE) + + +@pytest.mark.parametrize( + "path, geometry, first, second", + [pytest.param(c[1], *c[3:], id=c[0]) for c in _CASES], +) +def test_pair_numeric_on_one_binary(torch, path, geometry, first, second): + d, bs, _ = geometry + compiled = None + for i, shape in enumerate((first, second)): + problem = _problems(geometry, shape) + q, k, v, cu, table = _case(torch, geometry, shape, seed=i) + out = torch.empty_like(q) + au.run_unified_attention_torch( + problem=problem, + q=q, + k=k, + v=v, + out=out, + cu_seqlens_q=cu, + seqused_k=torch.tensor(shape["kv_lens"], dtype=torch.int32, device="cuda"), + softmax_scale=1.0 / math.sqrt(d), + block_table=table, + softcap=0.0, + backend="3d" if path == "3d" else "tiled", + use_graph=False, + ) + torch.cuda.synchronize() + assert _key(problem, path) in _cache(path), f"shape {i} missed the {path} cache" + if compiled is None: + compiled = _compiled() + else: + assert _compiled() == compiled, "the second shape compiled a new binary" + want = _reference(torch, q, k, v, cu, table, shape, d, bs) + max_abs = (out.float() - want).abs().max().item() + assert ( + math.isfinite(max_abs) and max_abs <= 4e-2 + ), f"shape {i}: max_abs {max_abs:.3e} vs fp32 reference" diff --git a/dnn-providers/hip-kernel-provider/rocke/library/tests/test_attention_tiled_3d_running_max.py b/dnn-providers/hip-kernel-provider/rocke/library/tests/test_attention_tiled_3d_running_max.py index 3f030181d3cc..34d566a0bbff 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/tests/test_attention_tiled_3d_running_max.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/tests/test_attention_tiled_3d_running_max.py @@ -206,6 +206,7 @@ def pages(x): "num_seqs": 1, "block_table_stride": blocks, "qq_bias_stride_0": 0, + **au._runtime_kernarg_values(problem), } reduce_values = { "output_ptr": out, @@ -213,6 +214,7 @@ def pages(x): "segm_max_ptr": segm_max, "segm_expsum_ptr": segm_expsum, "seq_lens_ptr": seq_lens, + **au._runtime_kernarg_values(problem), } # The launch fences the stream, so ``out`` is complete when it returns. au._launch_3d_pipeline( 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 203ee1d18fff..8e5da31b28d5 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 @@ -19,7 +19,8 @@ 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``. + * ``_binary_search_iters`` mirrors every spec class that bakes the search; + gfx942 searches at runtime, so its batch sizes share a binary. No GPU required. """ @@ -82,11 +83,21 @@ def routing_env(monkeypatch): def _norm_spec(spec): """Spec identity as the emitter sees it: ``num_seqs`` only reaches IR - through ``binary_search_iters``, which the 3D decode grid does not emit.""" + through ``binary_search_iters``, which the 3D decode grid does not emit and + a runtime-search body (``RUNTIME_SEARCH_ITERS``) does not bake.""" fields = dataclasses.asdict(spec) fields.pop("num_seqs") - if not getattr(spec, "use_decode_grid", False): + bakes_search = not getattr(type(spec), "RUNTIME_SEARCH_ITERS", False) + if bakes_search and not getattr(spec, "use_decode_grid", False): fields["binary_search_iters"] = spec.binary_search_iters + # A kernarg field reaches IR only through the structure it selects. + runtime = getattr(type(spec), "RUNTIME_KERNARG_FIELDS", ()) + if "sliding_window" in runtime: + fields["sliding_window"] = spec.sliding_window > 0 + if "num_query_heads" in runtime: + fields.pop("num_query_heads") + fields.pop("num_kv_heads") + fields["num_queries_per_kv"] = spec.num_queries_per_kv return fields @@ -100,12 +111,17 @@ def _built_2d(problem, arch): def _built_3d(problem, arch): + reduce_cls = au._tiled_3d_impl(arch)[1] + heads = ( + (problem.num_queries_per_kv,) + if "num_query_heads" in getattr(reduce_cls, "RUNTIME_KERNARG_FIELDS", ()) + else (problem.num_query_heads, problem.num_kv_heads) + ) return dict( spec=_norm_spec(au._tiled_3d_spec_from_problem(problem, arch)), reduce=( problem.head_size, - problem.num_query_heads, - problem.num_kv_heads, + *heads, problem.dtype, au._num_segments(problem, arch), au._select_3d_waves_per_eu(problem, arch), @@ -272,7 +288,8 @@ def _assert_split(key_fn, built_fn, a, b, field, arch): assert key_fn(a, arch) != key_fn(b, arch) -@pytest.mark.parametrize("arch", ("gfx942", "gfx950", "gfx1250")) +# gfx942 searches at runtime, so its buckets share a binary (_SHARED_PAIRS). +@pytest.mark.parametrize("arch", ("gfx950", "gfx1250")) def test_search_bucket_edge_splits_2d_key(routing_env, 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. @@ -293,7 +310,7 @@ def test_search_bucket_edge_splits_2d_key(routing_env, arch): assert _lower_2d(a, arch, "llvm22") != _lower_2d(b, arch, "llvm22") -@pytest.mark.parametrize("arch", ("gfx942", "gfx950")) +@pytest.mark.parametrize("arch", ("gfx950",)) def test_search_bucket_edge_splits_3d_key(routing_env, arch): # Two query tokens per sequence, so the launch takes the searching prefill # grid rather than the decode grid. @@ -424,6 +441,36 @@ def _lower_3d(problem, arch, flavor): dtype="bf16", ) +_D256_4WARP = dict( + num_query_heads=16, + num_kv_heads=2, + head_size=256, + block_size=16, + max_seqlen_q=4096, + max_seqlen_k=4096, + dtype="bf16", +) +_D128_FOLD = dict( + num_query_heads=32, + num_kv_heads=8, + head_size=128, + block_size=16, + max_seqlen_q=8192, + max_seqlen_k=8192, + sliding_window=4096, + dtype="bf16", +) +# Two query tokens per sequence: the 3D searching grid, not the decode grid. +_D128_TWO_TOKENS = dict( + num_query_heads=32, + num_kv_heads=8, + head_size=128, + block_size=16, + max_seqlen_q=2, + 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 = [ @@ -473,6 +520,132 @@ def _lower_3d(problem, arch, flavor): ) for arch in ("gfx942", "gfx950") ), + # gfx942 reads the search trip count from num_seqs at runtime, so batch + # sizes several log2 buckets apart share a binary on every builder route. + ( + "gfx942-2d-ns3-vs-ns9", + "gfx942", + "2d", + P( + total_q=3 * 1024, + num_seqs=3, + max_seqlen_q=1024, + max_seqlen_k=1024, + **_D64_GQA8, + ), + P( + total_q=9 * 1024, + num_seqs=9, + max_seqlen_q=1024, + max_seqlen_k=1024, + **_D64_GQA8, + ), + ), + ( + "gfx942-2d-4warp-ns2-vs-ns9", + "gfx942", + "2d", + P(total_q=2 * 4096, num_seqs=2, **_D256_4WARP), + P(total_q=9 * 4096, num_seqs=9, **_D256_4WARP), + ), + ( + "gfx942-2d-4warp-fold-ns2-vs-ns9", + "gfx942", + "2d", + P(total_q=2 * 8192, num_seqs=2, **_D128_FOLD), + P(total_q=9 * 8192, num_seqs=9, **_D128_FOLD), + ), + ( + "gfx942-3d-prefill-ns3-vs-ns9", + "gfx942", + "3d", + P(total_q=6, num_seqs=3, **_D128_TWO_TOKENS), + P(total_q=18, num_seqs=9, **_D128_TWO_TOKENS), + ), + # gfx942 reads the window size as a kernarg, so window sizes share a binary. + ( + "gfx942-2d-sw256-vs-sw1024", + "gfx942", + "2d", + P( + total_q=2048, + num_seqs=2, + max_seqlen_q=1024, + max_seqlen_k=1024, + sliding_window=256, + **_D64_GQA8, + ), + P( + total_q=2048, + num_seqs=2, + max_seqlen_q=1024, + max_seqlen_k=1024, + sliding_window=1024, + **_D64_GQA8, + ), + ), + ( + "gfx942-2d-4warp-fold-sw1024-vs-sw4096", + "gfx942", + "2d", + P(total_q=2 * 8192, num_seqs=2, **dict(_D128_FOLD, sliding_window=1024)), + P(total_q=2 * 8192, num_seqs=2, **_D128_FOLD), + ), + ( + "gfx942-3d-decode-sw256-vs-sw4096", + "gfx942", + "3d", + P(total_q=4, num_seqs=4, sliding_window=256, **_D128_DECODE), + P(total_q=4, num_seqs=4, sliding_window=4096, **_D128_DECODE), + ), + # gfx942 reads the head counts as kernargs, so head pairs at one GQA ratio + # share a binary. + ( + "gfx942-2d-h64kv8-vs-h32kv4", + "gfx942", + "2d", + P(total_q=2048, num_seqs=2, max_seqlen_q=1024, max_seqlen_k=1024, **_D64_GQA8), + P( + total_q=2048, + num_seqs=2, + max_seqlen_q=1024, + max_seqlen_k=1024, + **dict(_D64_GQA8, num_query_heads=32, num_kv_heads=4), + ), + ), + ( + "gfx942-2d-4warp-h16kv2-vs-h32kv4", + "gfx942", + "2d", + P(total_q=2 * 4096, num_seqs=2, **_D256_4WARP), + P( + total_q=2 * 4096, + num_seqs=2, + **dict(_D256_4WARP, num_query_heads=32, num_kv_heads=4), + ), + ), + ( + "gfx942-2d-4warp-fold-h32kv8-vs-h16kv4", + "gfx942", + "2d", + P(total_q=2 * 8192, num_seqs=2, **_D128_FOLD), + P( + total_q=2 * 8192, + num_seqs=2, + **dict(_D128_FOLD, num_query_heads=16, num_kv_heads=4), + ), + ), + ( + "gfx942-3d-prefill-h32kv8-vs-h64kv16", + "gfx942", + "3d", + P(total_q=6, num_seqs=3, **_D128_TWO_TOKENS), + P( + total_q=6, + num_seqs=3, + **dict(_D128_TWO_TOKENS, num_query_heads=64, num_kv_heads=16), + ), + ), # gfx1250 3D floors the search at 32 iterations, so ns 3 and 4 share a # binary even though they sit in different log2 buckets. ( @@ -495,7 +668,7 @@ def test_same_outputs_share_key_and_ir(routing_env, arch, path, a, b, flavor): 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 a != b assert key_fn(a, arch) == key_fn(b, arch) assert lower(a, arch, flavor) == lower(b, arch, flavor) @@ -530,15 +703,41 @@ def test_launch_meta_separates_num_seqs_under_one_key(routing_env): assert grid_b[1] - grid_a[1] == b.num_seqs - a.num_seqs +def test_launch_geometry_separates_head_counts_under_one_key(routing_env): + # gfx942 reads the head counts as kernargs, so one binary serves 32/8 and + # 64/16. The 2D grid spans the kv heads and the 3D workspace the query + # heads, so with equal token and sequence counts the launch caches must + # still tell them apart. + arch = "gfx942" + base = dict( + total_q=1025, + num_seqs=1, + head_size=128, + block_size=16, + max_seqlen_q=1025, + max_seqlen_k=1025, + dtype="bf16", + num_cus=304, + ) + a = P(num_query_heads=32, num_kv_heads=8, **base) + b = P(num_query_heads=64, num_kv_heads=16, **base) + key = au._tiled_cache_key(a, arch) + assert key == au._tiled_cache_key(b, arch) + assert au._shape_key(key, a) != au._shape_key(key, b) + au._2D_LAUNCH_META.clear() + grid_a = au._get_2d_launch_meta(a, key, arch=arch).grid + grid_b = au._get_2d_launch_meta(b, key, arch=arch).grid + assert (grid_a[0], grid_b[0]) == (a.num_kv_heads, b.num_kv_heads) + + # --------------------------------------------------------------------------- # _binary_search_iters mirrors the spec classes # --------------------------------------------------------------------------- -# (spec class, floor). gfx1250 3D floors the search at 32 iterations. +# (spec class, floor) for every class that bakes the search. gfx1250 3D floors +# it at 32 iterations. _SPEC_CLASSES = ( - (Gfx942Spec2D, 1), - (Gfx942Spec3D, 1), (Gfx950Spec2D, 1), (Gfx950Spec3D, 1), (Gfx1250Spec2D, 1), @@ -567,3 +766,107 @@ def test_binary_search_iters_mirrors_spec(spec_cls, floor, num_seqs): dtype="bf16", ) assert au._binary_search_iters(problem, floor=floor) == spec.binary_search_iters + + +@pytest.mark.parametrize("spec_cls", (Gfx942Spec2D, Gfx942Spec3D)) +def test_gfx942_specs_search_at_runtime(spec_cls): + assert spec_cls.RUNTIME_SEARCH_ITERS + assert not hasattr(spec_cls, "binary_search_iters") + arch, path = "gfx942", "2d" if spec_cls is Gfx942Spec2D else "3d" + assert au._searches_at_runtime(arch, path) + assert not au._searches_at_runtime("gfx950", path) + + +def _trip_count(num_seqs: int) -> int: + """``binary_search_seq_idx_trip_count`` on the host: the exponent of the + round-to-nearest f32 conversion LLVM's ``sitofp`` performs.""" + np = pytest.importorskip("numpy") + bits = int( + np.array([num_seqs], dtype=np.int32).astype(np.float32).view(np.int32)[0] + ) + return max((bits >> 23) - 126, 1) + + +_BATCHES = sorted( + {n for n in range(1, 4097)} + | {(1 << k) + d for k in range(12, 31) for d in (-1, 0, 1)} + | {(1 << 24) + d for d in range(-8, 9)} + | {(1 << 31) - 1} +) + + +def test_runtime_trip_count_covers_every_batch(): + # The search needs bit_length(num_seqs) iterations. Exact below 2**24; past + # it the conversion may round up to the next power of two, adding one + # iteration the done guard makes a no-op. + for n in _BATCHES: + trip = _trip_count(n) + if n < 1 << 24: + assert ( + trip + == n.bit_length() + == au._binary_search_iters( + P( + total_q=n, + num_seqs=n, + num_query_heads=8, + num_kv_heads=8, + head_size=128, + block_size=16, + max_seqlen_q=1, + max_seqlen_k=1, + dtype="bf16", + ) + ) + ), n + else: + assert n.bit_length() <= trip <= n.bit_length() + 1, n + + +def test_gfx942_i64_kv_addressing_is_in_the_kernel_name(routing_env): + # i64 K/V addressing changes the 2D loaders, so a binary is identified by + # it: the name, not only the in-process key, must tell the two apart. + base = dict( + total_q=2048, + num_seqs=2, + num_query_heads=32, + num_kv_heads=8, + head_size=128, + block_size=16, + max_seqlen_q=1024, + max_seqlen_k=1024, + dtype="bf16", + num_cus=304, + ) + block_bytes = 16 * 8 * 128 * 2 + small = P(num_kv_blocks=64, **base) + large = P(num_kv_blocks=(1 << 31) // block_bytes + 1, **base) + assert not au._enable_i64_kv_addr(small) and au._enable_i64_kv_addr(large) + names = [ + au._tiled_spec_from_problem(p, "gfx942").kernel_name() for p in (small, large) + ] + assert "i64" not in names[0].split("_") + assert "i64" in names[1].split("_") + assert names[0].replace("_bf16_", "_bf16_i64_", 1) == names[1] + + +def test_gfx942_fast_paged_kv_desc_names_its_head_counts(): + # The fast paged-KV descriptor body shifts by a 64/8 K/V stride instead of + # reading the head-count kernargs, so its name must not share the plain + # g8 symbol that serves every head pair at that ratio. + base = dict( + head_size=64, + block_size=32, + tile_size=64, + num_query_heads=64, + num_kv_heads=8, + dtype="bf16", + num_warps=4, + use_sinks=False, + sliding_window=0, + has_softcap=False, + ) + plain = Gfx942Spec2D(**base).kernel_name() + fast = Gfx942Spec2D(**base, use_fast_paged_kv_desc=True).kernel_name() + assert "h64kv8" not in plain.split("_") + assert fast == plain.replace("_g8_", "_g8_h64kv8_", 1) + "_fastkvdesc" diff --git a/dnn-providers/hip-kernel-provider/rocke/library/tests/test_attention_tiled_launch_signature.py b/dnn-providers/hip-kernel-provider/rocke/library/tests/test_attention_tiled_launch_signature.py new file mode 100644 index 000000000000..66f80973fb47 --- /dev/null +++ b/dnn-providers/hip-kernel-provider/rocke/library/tests/test_attention_tiled_launch_signature.py @@ -0,0 +1,232 @@ +# Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +# SPDX-License-Identifier: MIT + +"""Launch-signature contract for the tiled 2D / 3D attention kernels. + +``KernelLauncher`` packs kernargs by the signature's names and order, so a kernel +whose params drift from its launch signature reads the wrong argument without +faulting. Every builder route the production path takes is built here and its +params are compared, by name and type in order, with the signature its launcher +uses (``_tiled_2d_signature`` / ``_tiled_3d_signatures``). That includes the +runtime shape kernargs a spec class declares (``RUNTIME_KERNARG_FIELDS``). + +No GPU required. +""" + +from __future__ import annotations + +import pytest + +from kernels.common import attention_unified as au + +P = au.UnifiedAttentionProblem + +_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(autouse=True) +def routing_env(monkeypatch): + """Production routing: no env knob steers a selector.""" + for var in _ROUTING_ENV: + monkeypatch.delenv(var, raising=False) + + +def _prefill(*, dtype="bf16", head_size=128, hq=32, hkv=8, sq=1024, **kw): + return P( + total_q=2 * sq, + num_seqs=2, + num_query_heads=hq, + num_kv_heads=hkv, + head_size=head_size, + block_size=kw.pop("block_size", 16), + max_seqlen_q=sq, + max_seqlen_k=sq, + dtype=dtype, + **kw, + ) + + +def _decode(*, dtype="bf16", head_size=128, hq=32, hkv=8, sq=1, **kw): + return P( + total_q=4 * sq, + num_seqs=4, + num_query_heads=hq, + num_kv_heads=hkv, + head_size=head_size, + block_size=16, + max_seqlen_q=sq, + max_seqlen_k=4096, + dtype=dtype, + **kw, + ) + + +# (arch, case id, problem, kernel-name suffix the route must emit or None) +_CASES_2D = ( + ("gfx942", "narrow_d64", _prefill(head_size=64, hq=8, hkv=8, sq=128), None), + ("gfx942", "fp16_flash_d128", _prefill(dtype="fp16", sq=2048), None), + ("gfx942", "bf16_flash_d128", _prefill(sq=2048), None), + ("gfx942", "window", _prefill(sliding_window=256, hq=8, hkv=8), None), + ("gfx942", "sinks", _prefill(use_sinks=True), None), + ( + "gfx942", + "4warp_d256", + _prefill(head_size=256, hq=16, hkv=2, sq=4096), + "_4wgqa", + ), + ( + "gfx942", + "4warp_fold", + _prefill(sliding_window=4096, sq=8192), + "_4wgqa_fold", + ), + ( + "gfx942", + "4warp_fp16_window", + _prefill(dtype="fp16", sliding_window=4096, sq=8192), + "_4wgqa", + ), + ("gfx950", "bf16_d128", _prefill(), None), + ( + "gfx950", + "fp16_d64_window", + _prefill(dtype="fp16", head_size=64, sliding_window=256), + None, + ), +) + + +def _assert_matches(kernel, signature): + names = [p.name for p in kernel.params] + assert names == [a["name"] for a in signature], kernel.name + types = [p.type.name.replace(" ", "") for p in kernel.params] + assert types == [a["type"].replace(" ", "") for a in signature], kernel.name + + +@pytest.mark.parametrize( + "arch, problem, suffix", + [pytest.param(a, p, s, id=f"{a}-{i}") for a, i, p, s in _CASES_2D], +) +def test_2d_kernel_params_match_launch_signature(arch, problem, suffix): + kernel = au._tiled_2d_kernel(problem, arch) + if suffix is not None: + assert kernel.name.endswith(suffix), kernel.name + signature = au._tiled_2d_signature( + arch, + problem.dtype, + kv_dtype=au._kv_storage_dtype(problem), + include_additive_bias=problem.use_additive_bias, + ) + _assert_matches(kernel, signature) + + +_CASES_3D = ( + ("gfx942", "decode", _decode(), False), + ("gfx942", "decode_window", _decode(sliding_window=256), False), + ("gfx942", "decode_lse", _decode(emit_lse=True), False), + # Two queries per sequence leave the decode grid for the runtime search. + ("gfx942", "two_token", _decode(sq=2), False), + ("gfx942", "strided", _decode(hq=4, hkv=2), True), + ("gfx950", "decode", _decode(), False), + ("gfx950", "strided", _decode(hq=4, hkv=2), True), +) + + +@pytest.mark.parametrize( + "arch, problem, strided_kv", + [pytest.param(a, p, s, id=f"{a}-{i}") for a, i, p, s in _CASES_3D], +) +def test_3d_kernel_params_match_launch_signatures(arch, problem, strided_kv): + if problem.max_seqlen_q > 1: + assert not au._use_decode_grid(problem, arch) + num_segments = au._num_segments(problem, arch) + segment, reduce = au._tiled_3d_kernels( + problem, arch, num_segments, strided_kv=strided_kv + ) + seg_sig, red_sig = au._tiled_3d_signatures( + arch, + problem.dtype, + kv_dtype=au._kv_storage_dtype(problem), + strided_kv=strided_kv, + include_additive_bias=problem.use_additive_bias, + emit_lse=problem.emit_lse, + ) + _assert_matches(segment, seg_sig) + _assert_matches(reduce, red_sig) + + +@pytest.mark.parametrize( + "codepath", ["narrow", "transposed_x8", "gfx942_4warp", "splitkv"] +) +def test_tuning_candidate_signature_matches_its_kernel(codepath): + """The registry's tuning candidates launch through ``candidate.signature``.""" + from dispatch.attention import attention_candidates + from dispatch.attention.common import AttentionMaskType, AttentionRequest + from dispatch.attention.gfx942_unified import GFX942_TUNING_VARIANTS + + variants = [v for v in GFX942_TUNING_VARIANTS if v.codepath == codepath] + candidates = {c.spec_id: c for c in attention_candidates()} + head_size = 256 if codepath == "gfx942_4warp" else 128 + # Two queries per sequence keep the split-KV segment off the decode grid. + seqlen_q = 2 if codepath == "splitkv" else 1024 + checked = 0 + for variant in variants: + request = AttentionRequest( + batch=1, + nhead_q=16, + nhead_k=2, + seqlen_q=seqlen_q, + seqlen_k=1024, + hdim_q=head_size, + hdim_v=head_size, + arch="gfx942", + dtype="bf16", + mask_type=AttentionMaskType.BOTTOM_RIGHT_CAUSAL, + num_cus=304, + algorithm="unified_tuning", + spec_id=variant.spec_id, + ) + candidate = candidates[variant.spec_id] + if not candidate.admits(request)[0]: + continue + spec = candidate.select_spec(request) + kernel = candidate.build(spec, "gfx942") + if spec.path == "3d": + # The signature is the segment's; _get_3d_pipeline launches the + # reduce with the production 3D pair. + kernel, reduce = kernel + assert "_dg" not in kernel.name.split("_"), kernel.name + _, red_sig = au._tiled_3d_signatures( + "gfx942", + spec.kernel_spec.dtype, + kv_dtype=spec.kernel_spec.kv_storage_dtype, + ) + _assert_matches(reduce, red_sig) + _assert_matches(kernel, candidate.signature(spec)) + checked += 1 + break + assert checked, f"no {codepath} tuning candidate admits the request" + + +@pytest.mark.parametrize("arch", ["gfx942", "gfx950"]) +def test_runtime_values_cover_every_declared_field(arch): + """The shared argument pack supplies every field a spec class can declare, + so a launcher never raises on a runtime kernarg.""" + segment_cls, reduce_cls, _, _, _ = au._tiled_3d_impl(arch) + spec_2d_cls, _, _ = au._tiled_2d_impl(arch) + values = au._runtime_kernarg_values(_prefill()) + for cls in (spec_2d_cls, segment_cls, reduce_cls): + missing = set(au._runtime_kernarg_fields(cls)) - set(values) + assert not missing, (cls.__name__, missing) diff --git a/dnn-providers/hip-kernel-provider/rocke/library/tests/test_attn_bf16_d128_ring.py b/dnn-providers/hip-kernel-provider/rocke/library/tests/test_attn_bf16_d128_ring.py index 2f9fdb58d5b9..6736f225c5cd 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/tests/test_attn_bf16_d128_ring.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/tests/test_attn_bf16_d128_ring.py @@ -267,7 +267,7 @@ def test_d128_numeric_vs_fp32_oracle_at_magnitude(monkeypatch, dtype, tol, hq, h from rocke import compile_kernel from kernels import build_unified_attention_2d_tiled - from kernels.common.attention_unified import _attn_signature + from kernels.common.attention_unified import _tiled_2d_signature from rocke.runtime import KernelLauncher kernel = build_unified_attention_2d_tiled(spec, arch="gfx942") @@ -275,9 +275,7 @@ def test_d128_numeric_vs_fp32_oracle_at_magnitude(monkeypatch, dtype, tol, hq, h launcher = KernelLauncher( hsaco=artifact.hsaco, kernel_name=artifact.kernel_name, - signature=_attn_signature( - dtype, include_bt_stride=True, include_qq_bias_stride=True - ), + signature=_tiled_2d_signature("gfx942", dtype), cache_key=("d128guard", spec.kernel_name(), dtype, sq, hq, hk), ) out, _ = H._run_rocke(s, data, launcher, spec, warmup=3, attempts=3) diff --git a/dnn-providers/hip-kernel-provider/rocke/library/tests/test_strided_kv_decode.py b/dnn-providers/hip-kernel-provider/rocke/library/tests/test_strided_kv_decode.py index 40c3147389e6..dff5a094f4aa 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/tests/test_strided_kv_decode.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/tests/test_strided_kv_decode.py @@ -16,7 +16,7 @@ import pytest from kernels.common.attention_kv_cache import KvTensorLayout, StridedKvCacheLayout -from kernels.common.attention_unified import _3d_signature +from kernels.common.attention_unified import _tiled_3d_signatures from rocke.core.verify import verify @@ -105,7 +105,7 @@ def test_strided_segment_signature_and_identity(arch, dtype): assert strided.kernel_name() == paged.kernel_name() + "_stridedkv" kernel = module.build_unified_attention_3d_tiled(strided, arch=arch) assert not verify(kernel) - signature = _3d_signature(dtype, strided_kv=True) + signature, _ = _tiled_3d_signatures(arch, dtype, strided_kv=True) assert [p.name for p in kernel.params] == [p["name"] for p in signature] assert [p.type.name.replace(" ", "") for p in kernel.params] == [ p["type"].replace(" ", "") for p in signature diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/CMakeLists.txt b/dnn-providers/hip-kernel-provider/rocke/platform/CMakeLists.txt index 119785823c8d..d5c5efa2d482 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/CMakeLists.txt +++ b/dnn-providers/hip-kernel-provider/rocke/platform/CMakeLists.txt @@ -633,7 +633,8 @@ if(ROCKE_INSTALL_TESTS) "tests/library/tests/test_attention_dense_gfx950_lds.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") + "tests/library/tests/test_attention_tiled_3d_kv_cache_policy.py" + "tests/library/tests/test_attention_tiled_launch_signature.py") set(_ROCKE_LIBRARY_ATTENTION_CPU_TESTS "") foreach(_t ${_ROCKE_LIBRARY_ATTENTION_CPU_TEST_RELPATHS}) string(APPEND _ROCKE_LIBRARY_ATTENTION_CPU_TESTS @@ -647,6 +648,18 @@ if(ROCKE_INSTALL_TESTS) " ENVIRONMENT \"PYTHONPATH=${ROCKE_TEST_CTEST_REL}tests/library\"\n" " TIMEOUT 180)\n") + # Tiled attention cache-key contract (no GPU). Its invariant sweep compares + # key equality with build equality across every selector threshold for each + # tiled arch, so each sweep case runs for tens of seconds: it gets its own + # entry and budget rather than the 30 s per-test limit of the list above. + string(APPEND _ROCKE_INSTALLED_CTEST_ENTRIES + "add_test(rocke_library_attention_cache_key_pytest python -m pytest \"${ROCKE_TEST_CTEST_REL}tests/library/tests/test_attention_tiled_cache_key.py\" -v -rs --timeout=240)\n") + string(APPEND _ROCKE_INSTALLED_CTEST_ENTRIES + "set_tests_properties(rocke_library_attention_cache_key_pytest PROPERTIES\n" + " LABELS \"ckc\\;python\\;attention\\;host\\;quick\"\n" + " ENVIRONMENT \"PYTHONPATH=${ROCKE_TEST_CTEST_REL}tests/library\"\n" + " TIMEOUT 900)\n") + # DSA lightning-indexer CPU tests (no GPU: validator, dispatch wiring, and the # Python-lowered IR-sha golden). Same staged-library PYTHONPATH as the conv # lane. The on-GPU numeric tests are enrolled in the GPU lane below. @@ -840,6 +853,7 @@ if(ROCKE_INSTALL_TESTS) "tests/library/tests/test_gfx1250_attention.py" "tests/library/tests/test_gfx942_gqa_head_fold_numeric.py" "tests/library/tests/test_attention_gfx942_kv_pages.py" + "tests/library/tests/test_attention_gfx942_runtime_shape_numeric.py" "tests/library/tests/test_attention_tiled_3d_kv_lds_swizzle.py" "tests/library/tests/test_attention_tiled_3d_running_max.py" "tests/library/tests/test_attention_tiled_3d_wave_page_id_load.py" diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/helpers/attention_ext.cpp b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/helpers/attention_ext.cpp index c4fa74ed12a9..30451ad8e598 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/helpers/attention_ext.cpp +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/helpers/attention_ext.cpp @@ -129,13 +129,28 @@ rocke_value_t* rocke_wave64_reduce_sum(rocke_ir_builder_t* b, rocke_value_t* v) /* ----------------------------------------------- binary search on cu_q */ -rocke_value_t* rocke_binary_search_seq_idx(rocke_ir_builder_t* b, - rocke_value_t* cu_q, - rocke_value_t* q_block_global_idx, - rocke_value_t* num_seqs, - int block_q, - int iterations, - bool per_token) +rocke_value_t* rocke_binary_search_seq_idx_trip_count(rocke_ir_builder_t* b, + rocke_value_t* num_seqs) +{ + /* bits = b.bitcast(b.sitofp_f32(num_seqs), I32) + * exponent = b.lshr(bits, b.const_i32(23)) + * return b.smax(b.sub(exponent, b.const_i32(126)), b.const_i32(1)) */ + rocke_value_t* as_f32 = rocke_b_sitofp_f32(b, num_seqs); + rocke_value_t* bits = rocke_b_bitcast(b, as_f32, rocke_i32()); + rocke_value_t* exponent = rocke_b_lshr(b, bits, rocke_b_const_i32(b, 23)); + rocke_value_t* unbiased = rocke_b_sub(b, exponent, rocke_b_const_i32(b, 126)); + return rocke_b_smax(b, unbiased, rocke_b_const_i32(b, 1)); +} + +/* The search with a constant trip count (`upper == NULL`) or a runtime one. */ +static rocke_value_t* binary_search_seq_idx_impl(rocke_ir_builder_t* b, + rocke_value_t* cu_q, + rocke_value_t* q_block_global_idx, + rocke_value_t* num_seqs, + int block_q, + int iterations, + rocke_value_t* upper, + bool per_token) { rocke_value_t* bq; rocke_iter_arg_t iter_args[2]; @@ -162,7 +177,7 @@ rocke_value_t* rocke_binary_search_seq_idx(rocke_ir_builder_t* b, * hoist lb/ub/step first, then build the iter inits, to pin IR order. */ { rocke_value_t* lb = rocke_b_const_i32(b, 0); - rocke_value_t* ub = rocke_b_const_i32(b, (int64_t)iterations); + rocke_value_t* ub = upper != NULL ? upper : rocke_b_const_i32(b, (int64_t)iterations); rocke_value_t* step = rocke_b_const_i32(b, 1); iter_args[0].name = "left"; iter_args[0].init = rocke_b_const_i32(b, 0); @@ -216,3 +231,27 @@ rocke_value_t* rocke_binary_search_seq_idx(rocke_ir_builder_t* b, res0 = (loop.op != NULL) ? loop.op->results[0] : NULL; return rocke_b_sub(b, res0, rocke_b_const_i32(b, 1)); } + +rocke_value_t* rocke_binary_search_seq_idx(rocke_ir_builder_t* b, + rocke_value_t* cu_q, + rocke_value_t* q_block_global_idx, + rocke_value_t* num_seqs, + int block_q, + int iterations, + bool per_token) +{ + return binary_search_seq_idx_impl( + b, cu_q, q_block_global_idx, num_seqs, block_q, iterations, NULL, per_token); +} + +rocke_value_t* rocke_binary_search_seq_idx_dyn(rocke_ir_builder_t* b, + rocke_value_t* cu_q, + rocke_value_t* q_block_global_idx, + rocke_value_t* num_seqs, + int block_q, + rocke_value_t* iterations, + bool per_token) +{ + return binary_search_seq_idx_impl( + b, cu_q, q_block_global_idx, num_seqs, block_q, 0, iterations, per_token); +} diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/helpers/transforms.cpp b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/helpers/transforms.cpp index b960a6ff4706..3a2037849a41 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/helpers/transforms.cpp +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/helpers/transforms.cpp @@ -1646,6 +1646,10 @@ bool rocke_transforms_descriptor_offset_i64_split(rocke_ir_builder_t* b, { const char* name = desc->base_names[i]; int stride = desc->base_strides[i]; + rocke_value_t* stride_v + = desc->is_dynamic + ? ((const rocke_dynamic_tensor_descriptor_t*)desc)->dynamic_strides[i] + : NULL; const rocke_coord_var_t* c = rocke_i_map_get(&coords, name); if(c == NULL) { @@ -1653,7 +1657,19 @@ bool rocke_transforms_descriptor_offset_i64_split(rocke_ir_builder_t* b, return false; } valid = rocke_i_and(b, valid, c->valid); - if(strcmp(name, base_coord) == 0) + if(stride_v != NULL && strcmp(name, base_coord) == 0) + { + rocke_value_t* base_val = rocke_b_to_sgpr_u32(b, c->value); + rocke_value_t* base_w = rocke_b_zext(b, base_val, rocke_i64()); + rocke_value_t* stride_w = rocke_b_zext(b, stride_v, rocke_i64()); + base_i64 = rocke_b_mul(b, base_w, stride_w); + } + else if(stride_v != NULL) + { + rocke_value_t* term = rocke_b_mul(b, c->value, stride_v); + within = (within == NULL) ? term : rocke_b_add(b, within, term); + } + else if(strcmp(name, base_coord) == 0) { /* i64 term: pin the wave-uniform block id to an SGPR before widening * (Python b.mul(b.zext(b.to_sgpr_u32(c.value), I64), const_i64(stride))). @@ -1959,6 +1975,52 @@ rocke_dynamic_tensor_descriptor_t* return ddesc; } +rocke_dynamic_tensor_descriptor_t* + rocke_tensor_descriptor_naive_dynamic_mixed(rocke_ir_builder_t* b, + const char* name, + const char* const* coord_names, + int n_coords, + rocke_value_t* const* strides, + const int* static_strides) +{ + rocke_dynamic_tensor_descriptor_t* ddesc; + rocke_tensor_descriptor_t* base; + int fake_lengths[8]; + int base_strides[8]; + int i; + + if(!rocke_i_live(b) || n_coords < 1 || n_coords > 8) + { + return NULL; + } + ddesc = (rocke_dynamic_tensor_descriptor_t*)rocke_arena_calloc( + &b->arena, sizeof(rocke_dynamic_tensor_descriptor_t)); + if(ddesc == NULL) + { + return NULL; + } + /* Python create(): base_lengths all 0, base_strides the int stride or 1. */ + for(i = 0; i < n_coords; ++i) + { + fake_lengths[i] = 0; + base_strides[i] = strides[i] != NULL ? 1 : static_strides[i]; + } + base = rocke_tensor_descriptor_naive( + b, name, fake_lengths, n_coords, base_strides, coord_names, n_coords); + if(base == NULL) + { + return NULL; + } + ddesc->base = *base; + ddesc->n_dynamic = n_coords; + for(i = 0; i < n_coords; ++i) + { + ddesc->dynamic_strides[i] = strides[i]; + } + ddesc->base.is_dynamic = 1; + return ddesc; +} + bool rocke_dynamic_tensor_descriptor_offset(rocke_ir_builder_t* b, const rocke_dynamic_tensor_descriptor_t* desc, const char* const* in_names, @@ -2013,7 +2075,18 @@ bool rocke_dynamic_tensor_descriptor_offset(rocke_ir_builder_t* b, } valid = rocke_i_and(b, valid, cv->valid); stride_v = desc->dynamic_strides[i]; - term = rocke_b_mul(b, cv->value, stride_v); + if(stride_v != NULL) + { + term = rocke_b_mul(b, cv->value, stride_v); + } + else if(base->base_strides[i] == 1) + { + term = cv->value; + } + else + { + term = rocke_b_mul(b, cv->value, rocke_b_const_i32(b, (int64_t)base->base_strides[i])); + } off = (off == NULL) ? term : rocke_b_add(b, off, term); } if(off == NULL) diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/helper_helper_rocke.helpers.attention.h b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/helper_helper_rocke.helpers.attention.h index c68833fcc9a9..6f1d9886df89 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/helper_helper_rocke.helpers.attention.h +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/helper_helper_rocke.helpers.attention.h @@ -9,6 +9,8 @@ * PORTED SYMBOLS (this phase): * - rocke_apply_softcap_log2 (apply_softcap_log2) * - rocke_binary_search_seq_idx (binary_search_seq_idx) + * - rocke_binary_search_seq_idx_dyn (binary_search_seq_idx, runtime trip count) + * - rocke_binary_search_seq_idx_trip_count (binary_search_seq_idx_trip_count) * - rocke_mfma_16x16x16_for_dtype (mfma_16x16x16_for_dtype) * - rocke_wave64_reduce_max (wave64_reduce_max) * - rocke_wave64_reduce_sum (wave64_reduce_sum) @@ -71,6 +73,21 @@ rocke_value_t* rocke_binary_search_seq_idx(rocke_ir_builder_t* b, int iterations, bool per_token); +/* The same search with a runtime trip count, for kernels that serve any batch. */ +rocke_value_t* rocke_binary_search_seq_idx_dyn(rocke_ir_builder_t* b, + rocke_value_t* cu_q, + rocke_value_t* q_block_global_idx, + rocke_value_t* num_seqs, + int block_q, + rocke_value_t* iterations, + bool per_token); + +/* binary_search_seq_idx_trip_count: max(1, bit_length(num_seqs)) from the exponent + * of (float)num_seqs. Exact below 2^24 sequences; past that it may add one no-op + * iteration. */ +rocke_value_t* rocke_binary_search_seq_idx_trip_count(rocke_ir_builder_t* b, + rocke_value_t* num_seqs); + #ifdef __cplusplus } /* extern "C" */ #endif diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/helper_rocke.helpers.transforms.h b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/helper_rocke.helpers.transforms.h index 4798299c65e4..8183b1fb39da 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/helper_rocke.helpers.transforms.h +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/helper_rocke.helpers.transforms.h @@ -397,7 +397,9 @@ typedef struct rocke_dynamic_tensor_descriptor { rocke_tensor_descriptor_t base; /* must be first for safe cast */ int n_dynamic; /* == base.n_base */ - rocke_value_t* dynamic_strides[8]; /* parallel to base.base_names */ + /* Parallel to base.base_names. NULL marks a compile-time stride, read from + * base.base_strides with the static rule (no multiply for 1). */ + rocke_value_t* dynamic_strides[8]; } rocke_dynamic_tensor_descriptor_t; rocke_dynamic_tensor_descriptor_t* @@ -407,6 +409,16 @@ rocke_dynamic_tensor_descriptor_t* int n_coords, rocke_value_t* const* strides); +/* Python: DynamicTensorDescriptor.create with a mix of Value and int strides. + * strides[i] == NULL takes static_strides[i] as a compile-time stride. */ +rocke_dynamic_tensor_descriptor_t* + rocke_tensor_descriptor_naive_dynamic_mixed(rocke_ir_builder_t* b, + const char* name, + const char* const* coord_names, + int n_coords, + rocke_value_t* const* strides, + const int* static_strides); + bool rocke_dynamic_tensor_descriptor_offset(rocke_ir_builder_t* b, const rocke_dynamic_tensor_descriptor_t* desc, const char* const* in_names, @@ -495,7 +507,8 @@ bool rocke_transforms_descriptor_offset(rocke_ir_builder_t* b, * terms are summed into a small i32 within-block offset. *out_base_i64 and * *out_within are written on success (within is const_i32(0) when no other * terms); *out_valid is NULL == Python None. Returns false (builder error set) - * if base_coord is not among the descriptor's base coords. */ + * if base_coord is not among the descriptor's base coords. A dynamic descriptor + * multiplies by its runtime strides, widening the base_coord one to i64. */ bool rocke_transforms_descriptor_offset_i64_split(rocke_ir_builder_t* b, const rocke_tensor_descriptor_t* desc, const char* base_coord, diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/instance_gfx942_attention_tiled_2d_internal.h b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/instance_gfx942_attention_tiled_2d_internal.h index 9f195278430e..0b28fcf6f08e 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/instance_gfx942_attention_tiled_2d_internal.h +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/instance_gfx942_attention_tiled_2d_internal.h @@ -177,6 +177,14 @@ typedef struct rocke_gfx942_attn2d_build_ctx rocke_value_t* num_seqs_p; rocke_value_t* bt_stride_p; rocke_value_t* qq_bias_stride0_p; + /* RUNTIME_KERNARG_FIELDS params, in declaration order; set only when + * runtime_shape. */ + rocke_value_t* num_qh_p; + rocke_value_t* num_kv_p; + rocke_value_t* sliding_window_p; + /* The gfx942 Python body: it reads its runtime kernargs and searches at + * runtime. False when this body stands in for the gfx950 builder. */ + bool runtime_shape; /* ---- grid ids + wave decomposition (lines 1371-1388) ---- */ rocke_value_t* q_block_global_idx; @@ -268,20 +276,21 @@ typedef struct rocke_gfx942_attn2d_build_ctx rocke_value_t* lane_rg_is2_v; /* ---- SSA constants (lines 1891-1912) ---- */ - rocke_value_t* c0; rocke_value_t* zero_f; rocke_value_t* neg_inf_v; /* b.const_f32(-inf) (1892) */ rocke_value_t* one_f_v; /* b.const_f32(1.0) (1894) */ rocke_value_t* rcp_ln2_v; /* b.const_f32(1.4426950408889634) (1895) */ rocke_value_t* qk_scale_v; /* derived from scale_p */ - rocke_value_t* sw_const_v; /* b.const_i32(SLIDING_WINDOW) (1910) */ + rocke_value_t* sw_const_v; /* sliding_window_p, else const_i32(SLIDING_WINDOW) */ /* ---- paged-KV byte descriptor (full transform DAG, lines 2163-2351) ---- */ rocke_tensor_descriptor_t* q_desc; /* output/query [token,head,dim] */ rocke_tensor_descriptor_t* kv_desc; /* paged K/V byte descriptor */ rocke_value_t* seq_base; /* to_sgpr_u32(seq_idx*bt_stride) */ rocke_value_t* block_table_max_idx; /* to_sgpr_u32(num_seqs*bt_stride) */ - rocke_value_t* kv_block_bytes_c_v; /* const BS*NUM_KV*HD*KV_BYTES (2227) */ + rocke_value_t* kv_block_bytes_c_v; /* one-block bound: kv_stride_blk_b_v, or its const */ + rocke_value_t* kv_stride_tok_b_v; /* num_kv_p*HD*KV_BYTES (runtime_shape only) */ + rocke_value_t* kv_stride_blk_b_v; /* kv_stride_tok_b_v*BS (runtime_shape only) */ rocke_value_t* lane_half_base_v; /* tid * KV_HALVES_PER_LANE (2232) */ rocke_value_t* zero_soff_v; /* const_i32(0) soffset (2238) */ rocke_value_t* wave_lds_off_i64_v; /* wave_lds_offset_i64 (2254/2263) */ @@ -361,6 +370,13 @@ typedef struct rocke_gfx942_attn2d_build_ctx rocke_value_t* l_nonzero[ROCKE_GFX942_ATTN2D_MAX_REGS_PER_LANE]; } rocke_gfx942_attn2d_build_ctx_t; +/* The query head count as a row bound: the kernarg under runtime_shape, else + * the constant the gfx950 builder bakes, created at each use as Python does. */ +static inline rocke_value_t* rocke_gfx942_attn2d_num_qh(rocke_gfx942_attn2d_build_ctx_t* ctx) +{ + return ctx->runtime_shape ? ctx->num_qh_p : rocke_b_const_i32(ctx->b, ctx->NUM_QH); +} + /* ===================================================================== * * DRIVER-INTERNAL ctx population (the build prologue, lines 1104-1912). * diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/instance_gfx942_attention_tiled_3d.h b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/instance_gfx942_attention_tiled_3d.h index 6b689a020685..d0537f21aaf2 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/instance_gfx942_attention_tiled_3d.h +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/instance_gfx942_attention_tiled_3d.h @@ -34,9 +34,9 @@ * with a 4-reg-per-lane accumulator decoded via _mfma_16x16_c_row. * * REDUCE KERNEL (build_unified_attention_reduce_tiled, lines 1002-1133). Arch- - * neutral, pure f32 load / exp2 / store (byte-for-byte port of the gfx950 - * reduce): 64 threads, single active per (q_token, q_head), three passes (max, - * combine, normalize). + * neutral, pure f32 load / exp2 / store (the gfx950 reduce with a runtime + * query head count): 64 threads, single active per (q_token, q_head), three + * passes (max, combine, normalize). * * REUSED HELPER PORTS (no IR re-emitted here): * rocke/helper_helper_rocke.helpers.attention.h -- apply_softcap_log2, @@ -202,9 +202,6 @@ int rocke_unified_attention_3d_tiled_spec_tile_size( /* dtype_ir: F16 for "fp16", BF16 otherwise (returns a rocke_type_t*). */ const rocke_type_t* rocke_unified_attention_3d_tiled_spec_dtype_ir( const rocke_unified_attention_3d_tiled_spec_t* s); -/* binary_search_iters: 32 if num_seqs<=0, else max(1, ceil(log2(num_seqs+1))) */ -int rocke_unified_attention_3d_tiled_spec_binary_search_iters( - const rocke_unified_attention_3d_tiled_spec_t* s); /* kernel_name(): kernel_name_join(...) into caller's buf (NUL-terminated, up to * cap bytes incl. NUL). Returns the length that would be written (snprintf * convention), or -1 on a dead builder/encode error. */ diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/instance_gfx942_attention_tiled_3d_internal.h b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/instance_gfx942_attention_tiled_3d_internal.h index 881a70b64b20..848394ebb58d 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/instance_gfx942_attention_tiled_3d_internal.h +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/instance_gfx942_attention_tiled_3d_internal.h @@ -121,7 +121,6 @@ typedef struct rocke_gfx942_attn_tiled_3d_config int PV_N_TILES; /* HD // MFMA_N */ int THREADS; /* 64 */ - int binary_search_iters; /* spec.binary_search_iters */ /* ---- async / wide / fp8 KV feed geometry (lines 549-667) ---- */ int ASYNC_LDS_DWORDS; /* 1 */ @@ -225,6 +224,10 @@ typedef struct rocke_gfx942_attention_tiled_3d_build_ctx rocke_value_t* num_seqs_p; /* I32 num_seqs */ rocke_value_t* bt_stride_p; /* I32 block_table_stride */ rocke_value_t* qq_bias_stride0_p; /* I32 qq_bias_stride_0 */ + /* RUNTIME_KERNARG_FIELDS params, in declaration order. */ + rocke_value_t* num_qh_p; /* I32 num_query_heads */ + rocke_value_t* num_kv_p; /* I32 num_kv_heads */ + rocke_value_t* sliding_window_p; /* I32 sliding_window */ /* ---------- grid ids + thread (lines 331-334) ---------- */ rocke_value_t* q_block_global_idx; /* block_id_x() */ @@ -257,7 +260,7 @@ typedef struct rocke_gfx942_attention_tiled_3d_build_ctx rocke_value_t* one_f; /* const_f32(1.0) */ rocke_value_t* rcp_ln2; /* const_f32(1.4426950408889634) */ rocke_value_t* qk_scale; /* scale_p * rcp_ln2 */ - rocke_value_t* sw_const; /* const_i32(SLIDING_WINDOW) */ + rocke_value_t* sw_const; /* sliding_window_p */ rocke_value_t* z8; /* zero_vec(dtype, 8) */ /* ---------- lane decode (lines 482-483) ---------- */ @@ -280,6 +283,8 @@ typedef struct rocke_gfx942_attention_tiled_3d_build_ctx rocke_value_t* big_bytes; /* const_i32(0x7FFF0000) */ rocke_value_t* key_rsrc; /* buffer_rsrc(key, big_bytes) */ rocke_value_t* value_rsrc; /* buffer_rsrc(value, big_bytes) */ + rocke_value_t* kv_stride_tok_b; /* num_kv_p * HD * KV_BYTES */ + rocke_value_t* kv_stride_blk_b; /* kv_stride_tok_b * BS */ rocke_value_t* lane_half_base; /* tid * HALVES_PER_LANE */ rocke_value_t* K_lds_addr; /* smem_addr_of(K_lds) */ rocke_value_t* V_lds_addr; /* smem_addr_of(V_lds) */ @@ -320,6 +325,7 @@ typedef struct rocke_gfx942_attention_tiled_3d_build_ctx rocke_value_t* seg_max; /* F32* readonly */ rocke_value_t* seg_l; /* F32* readonly (segm_expsum) */ rocke_value_t* red_seq_lens; /* I32* readonly (_seq_lens, unused body) */ + rocke_value_t* red_num_qh_p; /* I32 num_query_heads (RUNTIME_KERNARG_FIELDS) */ /* ---------- grid ids (lines 1032-1034) ---------- */ rocke_value_t* q_token; /* block_id_x() */ @@ -401,8 +407,8 @@ rocke_value_t* rocke_gfx942_attention_tiled_3d_tile_base(rocke_gfx942_attention_tiled_3d_build_ctx_t* ctx, rocke_value_t* page_id, rocke_value_t* kv_tile_idx, - int page_stride, - int token_stride); + rocke_value_t* page_stride, + rocke_value_t* token_stride); /* _load_page_id_window(rel_page): lane l loads the id of the segment's page * rel_page + l (0 past the sequence's last page). */ diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/gfx942/attention_tiled_2d_ctx_prologue.cpp b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/gfx942/attention_tiled_2d_ctx_prologue.cpp index 6e9f5ca52c09..89a5016cdabb 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/gfx942/attention_tiled_2d_ctx_prologue.cpp +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/gfx942/attention_tiled_2d_ctx_prologue.cpp @@ -117,6 +117,9 @@ bool rocke_gfx942_attn2d_build_ctx_init(rocke_gfx942_attn2d_build_ctx_t* ctx, ctx->b = b; ctx->spec = spec; ctx->arch = arch; + /* This body also stands in for the gfx950 Python builder (the fastKV + * register-P twin), which bakes the shape values gfx942 reads at runtime. */ + ctx->runtime_shape = rocke_a2d_streq(arch, "gfx942"); /* ---- require_tiled_attention_arch(arch) (lines 1139-1141) ---- */ if(!rocke_a2d_require_tiled_attention_arch(b, arch)) @@ -398,6 +401,12 @@ bool rocke_gfx942_attn2d_build_ctx_init(rocke_gfx942_attn2d_build_ctx_t* ctx, ctx->num_seqs_p = rocke_b_param(b, "num_seqs", rocke_i32(), NULL); ctx->bt_stride_p = rocke_b_param(b, "block_table_stride", rocke_i32(), NULL); ctx->qq_bias_stride0_p = rocke_b_param(b, "qq_bias_stride_0", rocke_i32(), NULL); + if(ctx->runtime_shape) + { + ctx->num_qh_p = rocke_b_param(b, "num_query_heads", rocke_i32(), NULL); + ctx->num_kv_p = rocke_b_param(b, "num_kv_heads", rocke_i32(), NULL); + ctx->sliding_window_p = rocke_b_param(b, "sliding_window", rocke_i32(), NULL); + } } /* ---- grid ids + wave decomposition (lines 1371-1388) ---- */ @@ -427,14 +436,23 @@ bool rocke_gfx942_attn2d_build_ctx_init(rocke_gfx942_attn2d_build_ctx_t* ctx, } /* ---- seq lookup + Q-block geometry (lines 1390-1409) ---- */ - ctx->seq_idx - = rocke_binary_search_seq_idx(b, - ctx->cu_q, - ctx->q_block_global_idx, - ctx->num_seqs_p, - BLOCK_Q, - rocke_attention_tiled_2d_spec_binary_search_iters(spec), - false); + if(ctx->runtime_shape) + { + rocke_value_t* trip = rocke_binary_search_seq_idx_trip_count(b, ctx->num_seqs_p); + ctx->seq_idx = rocke_binary_search_seq_idx_dyn( + b, ctx->cu_q, ctx->q_block_global_idx, ctx->num_seqs_p, BLOCK_Q, trip, false); + } + else + { + ctx->seq_idx + = rocke_binary_search_seq_idx(b, + ctx->cu_q, + ctx->q_block_global_idx, + ctx->num_seqs_p, + BLOCK_Q, + rocke_attention_tiled_2d_spec_binary_search_iters(spec), + false); + } ctx->cu_q_start = rocke_b_global_load_i32(b, ctx->cu_q, ctx->seq_idx, 0); ctx->cu_q_stop = rocke_b_global_load_i32( b, ctx->cu_q, rocke_b_add(b, ctx->seq_idx, rocke_b_const_i32(b, 1)), 0); @@ -735,8 +753,8 @@ bool rocke_gfx942_attn2d_build_ctx_init(rocke_gfx942_attn2d_build_ctx_t* ctx, * (line 1994) and mask (lines 4111/4328/4403) regions. The previous port * created an unused i32 0 here and a FRESH const in the loop scaffold, * allocating an extra SSA value that shifted every later %N. Cache it. */ - ctx->sw_const_v = rocke_b_const_i32(b, ctx->SLIDING_WINDOW); - ctx->c0 = ctx->sw_const_v; + ctx->sw_const_v + = ctx->runtime_shape ? ctx->sliding_window_p : rocke_b_const_i32(b, ctx->SLIDING_WINDOW); /* ---- acc geometry the epilogue aliases (lines 2142-2143) ---- */ ctx->ACC_N_TILES = USE_MFMA_32X32 ? (HD / 32) : PV_N_TILES; @@ -994,7 +1012,7 @@ void rocke_gfx942_attn2d_emit_licm_hoist(rocke_gfx942_attn2d_build_ctx_t* ctx) /* row_ok = (qp_r < q_len) && (qh_r < NUM_QH) [&& (row < VALID_ROWS)]. * Skip the VALID_ROWS guard when VALID_ROWS == BLOCK_M (NQK divides BLOCK_M). */ rocke_value_t* row_ok_pos = rocke_b_cmp_lt(b, qp_r, ctx->cur_batch_q_len); - rocke_value_t* row_ok_qh = rocke_b_cmp_lt(b, qh_r, rocke_b_const_i32(b, ctx->NUM_QH)); + rocke_value_t* row_ok_qh = rocke_b_cmp_lt(b, qh_r, rocke_gfx942_attn2d_num_qh(ctx)); rocke_value_t* row_ok_inner = rocke_b_land(b, row_ok_pos, row_ok_qh); rocke_value_t* row_ok; if(ctx->VALID_ROWS < ctx->BLOCK_M) @@ -1022,7 +1040,7 @@ void rocke_gfx942_attn2d_emit_licm_hoist(rocke_gfx942_attn2d_build_ctx_t* ctx) for(reg = 0; reg < ctx->REGS_PER_LANE; ++reg) { rocke_value_t* qh_r = ctx->hoist_q_head[reg]; - rocke_value_t* qh_ok = rocke_b_cmp_lt(b, qh_r, rocke_b_const_i32(b, ctx->NUM_QH)); + rocke_value_t* qh_ok = rocke_b_cmp_lt(b, qh_r, rocke_gfx942_attn2d_num_qh(ctx)); rocke_value_t* slope = rocke_b_masked_global_load( b, ctx->alibi_slopes_ptr, qh_r, qh_ok, rocke_b_const_f32(b, 0.0), f32, 4); /* The softmax bucket reads the per-reg slope from hoist_q_head. */ diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/gfx942/attention_tiled_2d_epilogue.cpp b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/gfx942/attention_tiled_2d_epilogue.cpp index e84c88f2b596..56c5891f1389 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/gfx942/attention_tiled_2d_epilogue.cpp +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/gfx942/attention_tiled_2d_epilogue.cpp @@ -61,7 +61,6 @@ rocke_kernel_def_t* rocke_gfx942_attn2d_emit_epilogue(rocke_gfx942_attn2d_build_ const int ACC_M_ATOMS = ctx->ACC_M_ATOMS; const int REGS_PER_LANE = ctx->REGS_PER_LANE; const int NQK = ctx->NQK; - const int NUM_QH = ctx->NUM_QH; const int BLOCK_M = ctx->BLOCK_M; const int THREADS = ctx->THREADS; const int MFMA_N = 16; /* module-level constant MFMA_N (attention_tiled_2d.py:81) */ @@ -129,7 +128,7 @@ rocke_kernel_def_t* rocke_gfx942_attn2d_emit_epilogue(rocke_gfx942_attn2d_build_ rocke_value_t* op_qh_mod = rocke_b_mod(b, q_row_t, rocke_b_const_i32(b, NQK)); rocke_value_t* op_qh_t = rocke_b_add(b, op_qh_mul, op_qh_mod); rocke_value_t* op_ok_pos = rocke_b_cmp_lt(b, op_pos_t, cur_batch_q_len); - rocke_value_t* op_ok_qh = rocke_b_cmp_lt(b, op_qh_t, rocke_b_const_i32(b, NUM_QH)); + rocke_value_t* op_ok_qh = rocke_b_cmp_lt(b, op_qh_t, rocke_gfx942_attn2d_num_qh(ctx)); rocke_value_t* op_ok_inner = rocke_b_land(b, op_ok_pos, op_ok_qh); /* Skip the VALID_ROWS guard when VALID_ROWS == BLOCK_M (NQK divides BLOCK_M). */ rocke_value_t* op_mask_t; @@ -205,7 +204,8 @@ rocke_kernel_def_t* rocke_gfx942_attn2d_emit_epilogue(rocke_gfx942_attn2d_build_ rocke_value_t* op_qh32_base = rocke_b_add(b, op_qh32_mul, op_qh32_mod); /* Python emits the pos cmp BEFORE the qh cmp, then AND's VALID_ROWS; bind in order. */ rocke_value_t* op_ok_pos32 = rocke_b_cmp_lt(b, op_pos32_base, cur_batch_q_len); - rocke_value_t* op_ok_qh32 = rocke_b_cmp_lt(b, op_qh32_base, rocke_b_const_i32(b, NUM_QH)); + rocke_value_t* op_ok_qh32 + = rocke_b_cmp_lt(b, op_qh32_base, rocke_gfx942_attn2d_num_qh(ctx)); rocke_value_t* op_ok_inner32 = rocke_b_land(b, op_ok_pos32, op_ok_qh32); /* Skip the VALID_ROWS guard when VALID_ROWS == BLOCK_M (NQK divides BLOCK_M). */ rocke_value_t* op_mask32_base; @@ -303,7 +303,7 @@ rocke_kernel_def_t* rocke_gfx942_attn2d_emit_epilogue(rocke_gfx942_attn2d_build_ rocke_value_t* op_qh_mod = rocke_b_mod(b, OUT_ROW_BASE, rocke_b_const_i32(b, NQK)); rocke_value_t* op_qh = rocke_b_add(b, op_qh_mul, op_qh_mod); rocke_value_t* op_mask_pos = rocke_b_cmp_lt(b, op_pos, cur_batch_q_len); - rocke_value_t* op_mask_qh = rocke_b_cmp_lt(b, op_qh, rocke_b_const_i32(b, NUM_QH)); + rocke_value_t* op_mask_qh = rocke_b_cmp_lt(b, op_qh, rocke_gfx942_attn2d_num_qh(ctx)); rocke_value_t* op_mask_inner = rocke_b_land(b, op_mask_pos, op_mask_qh); /* Skip the VALID_ROWS guard when VALID_ROWS == BLOCK_M (NQK divides BLOCK_M). */ rocke_value_t* op_mask; diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/gfx942/attention_tiled_2d_kv_body_qk_softmax.cpp b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/gfx942/attention_tiled_2d_kv_body_qk_softmax.cpp index a680545387db..fd2a488a75af 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/gfx942/attention_tiled_2d_kv_body_qk_softmax.cpp +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/gfx942/attention_tiled_2d_kv_body_qk_softmax.cpp @@ -378,7 +378,7 @@ void rocke_gfx942_attn2d_emit_kv_body(rocke_gfx942_attn2d_build_ctx_t* ctx) /* row_ok = land(qpcur_batch_q_len); - rocke_value_t* st_ok_qh = rocke_b_cmp_lt(b, st_qh_iter, rocke_b_const_i32(b, ctx->NUM_QH)); + rocke_value_t* st_ok_qh = rocke_b_cmp_lt(b, st_qh_iter, rocke_gfx942_attn2d_num_qh(ctx)); rocke_value_t* st_ok_inner = rocke_b_land(b, st_ok_qp, st_ok_qh); if(ctx->VALID_ROWS < ctx->BLOCK_M) { @@ -394,7 +394,7 @@ void rocke_gfx942_attn2d_emit_kv_body(rocke_gfx942_attn2d_build_ctx_t* ctx) if(ctx->USE_ALIBI) { rocke_value_t* st_qh_iter_ok - = rocke_b_cmp_lt(b, st_qh_iter, rocke_b_const_i32(b, ctx->NUM_QH)); + = rocke_b_cmp_lt(b, st_qh_iter, rocke_gfx942_attn2d_num_qh(ctx)); st_alibi_slope_iter = rocke_b_masked_global_load(b, ctx->alibi_slopes_ptr, st_qh_iter, @@ -589,7 +589,7 @@ void rocke_gfx942_attn2d_emit_kv_body(rocke_gfx942_attn2d_build_ctx_t* ctx) * Skip the VALID_ROWS guard when VALID_ROWS == BLOCK_M. */ rocke_value_t* ok_pos = rocke_b_cmp_lt(b, qp_r, ctx->cur_batch_q_len); rocke_value_t* ok_qh - = rocke_b_cmp_lt(b, qh_r, rocke_b_const_i32(b, ctx->NUM_QH)); + = rocke_b_cmp_lt(b, qh_r, rocke_gfx942_attn2d_num_qh(ctx)); rocke_value_t* ok_inner = rocke_b_land(b, ok_pos, ok_qh); if(ctx->VALID_ROWS < ctx->BLOCK_M) { diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/gfx942/attention_tiled_2d_kv_descriptors_loads.cpp b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/gfx942/attention_tiled_2d_kv_descriptors_loads.cpp index 269d8cab0c4c..12775a5a4118 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/gfx942/attention_tiled_2d_kv_descriptors_loads.cpp +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/gfx942/attention_tiled_2d_kv_descriptors_loads.cpp @@ -104,7 +104,16 @@ static int rocke__v_bytes_per_call_swz(const rocke_gfx942_attn2d_build_ctx_t* ct * would shift downstream %values). */ static rocke_value_t* rocke__kv_block_bytes_c(rocke_gfx942_attn2d_build_ctx_t* ctx) { - if(ctx->kv_block_bytes_c_v == NULL) + if(ctx->kv_block_bytes_c_v == NULL && ctx->runtime_shape) + { + /* The token and block strides scale with the runtime kv head count. */ + ctx->kv_stride_tok_b_v = rocke_b_mul( + ctx->b, ctx->num_kv_p, rocke_b_const_i32(ctx->b, ctx->HD * ctx->KV_BYTES)); + ctx->kv_stride_blk_b_v + = rocke_b_mul(ctx->b, ctx->kv_stride_tok_b_v, rocke_b_const_i32(ctx->b, ctx->BS)); + ctx->kv_block_bytes_c_v = ctx->kv_stride_blk_b_v; + } + else if(ctx->kv_block_bytes_c_v == NULL) { int kv_stride_blk_b = ctx->BS * ctx->NUM_KV * ctx->HD * ctx->KV_BYTES; ctx->kv_block_bytes_c_v = rocke_b_const_i32(ctx->b, kv_stride_blk_b); @@ -394,8 +403,19 @@ void rocke_gfx942_attn2d_emit_preloop(rocke_gfx942_attn2d_build_ctx_t* ctx) int kv_lengths[4] = {1 << 24, ctx->BS, ctx->NUM_KV, ctx->HD}; int kv_strides[4] = {blk_b, tok_b, h_b, ctx->KV_BYTES}; static const char* const kv_coords[4] = {"physical_block", "token", "kv_head", "dim"}; - rocke_tensor_descriptor_t* kv_base = rocke_tensor_descriptor_naive( - b, "paged_kv_bytes", kv_lengths, 4, kv_strides, kv_coords, 4); + rocke_tensor_descriptor_t* kv_base; + if(ctx->runtime_shape) + { + rocke_value_t* dyn[4] = {ctx->kv_stride_blk_b_v, ctx->kv_stride_tok_b_v, NULL, NULL}; + kv_base = &rocke_tensor_descriptor_naive_dynamic_mixed( + b, "paged_kv_bytes", kv_coords, 4, dyn, kv_strides) + ->base; + } + else + { + kv_base = rocke_tensor_descriptor_naive( + b, "paged_kv_bytes", kv_lengths, 4, kv_strides, kv_coords, 4); + } /* Single-block tile (N_BLOCKS_PER_TILE == 1, T == BS): * indirect(tile_idx -> physical_block) ; unmerge(linear_half -> token,dim). @@ -616,7 +636,10 @@ static bool rocke__tile_pages(rocke_gfx942_attn2d_build_ctx_t* ctx, rocke_value_t* physical = rocke_b_masked_global_load(b, ctx->block_tables, idx, mask, def, rocke_i32(), 4); rocke_value_t* page_i64 = rocke_b_zext(b, rocke_b_to_sgpr_u32(b, physical), rocke_i64()); - rocke_value_t* base = rocke_b_mul(b, page_i64, rocke_b_const_i64(b, kv_stride_blk_b)); + rocke_value_t* stride64 = ctx->runtime_shape + ? rocke_b_zext(b, ctx->kv_stride_blk_b_v, rocke_i64()) + : rocke_b_const_i64(b, kv_stride_blk_b); + rocke_value_t* base = rocke_b_mul(b, page_i64, stride64); ptrs[page] = rocke_b_global_ptr_add(b, src, base); } rocke_value_t* kv_block_bytes_c = rocke__kv_block_bytes_c(ctx); @@ -639,8 +662,11 @@ static rocke_value_t* rocke_value_t* half = rocke_b_add(b, call_half, rocke__lane_half_base(ctx)); rocke_value_t* token = rocke_b_div(b, half, rocke_b_const_i32(b, ctx->HD)); rocke_value_t* dim = rocke_b_mod(b, half, rocke_b_const_i32(b, ctx->HD)); - rocke_value_t* token_b - = rocke_b_mul(b, token, rocke_b_const_i32(b, ctx->NUM_KV * ctx->HD * ctx->KV_BYTES)); + rocke_value_t* token_b = rocke_b_mul( + b, + token, + ctx->runtime_shape ? ctx->kv_stride_tok_b_v + : rocke_b_const_i32(b, ctx->NUM_KV * ctx->HD * ctx->KV_BYTES)); rocke_value_t* token_head_b = rocke_b_add(b, token_b, kv_head_b); rocke_value_t* dim_b = rocke_b_mul(b, dim, rocke_b_const_i32(b, ctx->KV_BYTES)); return rocke_b_add(b, token_head_b, dim_b); diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/gfx942/attention_tiled_2d_loop_scaffold.cpp b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/gfx942/attention_tiled_2d_loop_scaffold.cpp index 87e4c6755943..0ba9755606b9 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/gfx942/attention_tiled_2d_loop_scaffold.cpp +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/gfx942/attention_tiled_2d_loop_scaffold.cpp @@ -90,7 +90,6 @@ void rocke_gfx942_attn2d_emit_loop_bounds_and_inits(rocke_gfx942_attn2d_build_ct const int BLOCK_M = ctx->BLOCK_M; const int NQK = ctx->NQK; const int SLIDING_WINDOW = ctx->SLIDING_WINDOW; - const int NUM_QH = ctx->NUM_QH; const int SOFTMAX_STATE_SLOTS = ctx->SOFTMAX_STATE_SLOTS; /* ---- max_seq_prefix_len (1982-1984) ---- */ @@ -184,7 +183,7 @@ void rocke_gfx942_attn2d_emit_loop_bounds_and_inits(rocke_gfx942_attn2d_build_ct rocke_value_t* qh_mul = rocke_b_mul(B, ctx->kv_head_idx, rocke_b_const_i32(B, NQK)); rocke_value_t* qh_mod = rocke_b_mod(B, row, rocke_b_const_i32(B, NQK)); rocke_value_t* qh = rocke_b_add(B, qh_mul, qh_mod); - rocke_value_t* qh_in = rocke_b_cmp_lt(B, qh, rocke_b_const_i32(B, NUM_QH)); + rocke_value_t* qh_in = rocke_b_cmp_lt(B, qh, rocke_gfx942_attn2d_num_qh(ctx)); rocke_value_t* sink_h = rocke_b_global_load(B, ctx->sinks, qh, ctx->dtype, 2); rocke_value_t* sink_f = rocke_b_fmul(B, rocke_b_cast_to_f32(B, sink_h), rcp_ln2); m_inits[r] = rocke_b_select(B, qh_in, sink_f, neg_inf); diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/gfx942/attention_tiled_2d_public_entry_glue.cpp b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/gfx942/attention_tiled_2d_public_entry_glue.cpp index db62946fdd8c..ed757fc12c8e 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/gfx942/attention_tiled_2d_public_entry_glue.cpp +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/gfx942/attention_tiled_2d_public_entry_glue.cpp @@ -513,7 +513,7 @@ static rocke_status_t /* Value-carrying optionals become conditional strings; value-less flags * are passed as positional parts (kernel_name_join drops empties), exactly * as the Python kernel_name does. */ - char d_buf[32], b_buf[32], t_buf[32], hkv_buf[64], kv_buf[64]; + char d_buf[32], b_buf[32], t_buf[32], hkv_buf[64], fhkv_buf[64], kv_buf[64]; char sw_buf[32], w_buf[32], wpe_buf[32], mw_buf[32], kvcp_buf[64]; char rd_buf[32], ks_buf[32]; int nqh, nkv; @@ -530,13 +530,20 @@ static rocke_status_t snprintf(t_buf, sizeof(t_buf), "t%d", rocke_attention_tiled_2d_spec_tile_size_eff(s)); else t_buf[0] = '\0'; - snprintf(hkv_buf, sizeof(hkv_buf), "h%dkv%d", nqh, nkv); + /* Head counts are kernargs; only the GQA ratio is baked. */ + snprintf(hkv_buf, sizeof(hkv_buf), "g%d", nqh / nkv); + /* Except on the fast paged-KV descriptor path, whose shifts bake them. */ + if(s->use_fast_paged_kv_desc) + snprintf(fhkv_buf, sizeof(fhkv_buf), "h%dkv%d", nqh, nkv); + else + fhkv_buf[0] = '\0'; if(s->kv_storage_dtype) snprintf(kv_buf, sizeof(kv_buf), "kv%s", s->kv_storage_dtype); else kv_buf[0] = '\0'; + /* The window size is a kernarg; only whether the body masks one is baked. */ if(s->sliding_window > 0) - snprintf(sw_buf, sizeof(sw_buf), "sw%d", s->sliding_window); + snprintf(sw_buf, sizeof(sw_buf), "sw"); else sw_buf[0] = '\0'; if(s->num_warps != 1) @@ -578,11 +585,14 @@ static rocke_status_t b_buf, t_buf, hkv_buf, + fhkv_buf, s->dtype ? s->dtype : "", kv_buf, /* Python emits "fnuz" under the *same* condition as the kv part, * not under a format test, so it keys off kv_buf being non-empty. */ kv_buf[0] ? "fnuz" : "", + /* i64 K/V addressing changes the loaders, so it is a distinct kernel. */ + s->use_i64_kv_addr ? "i64" : "", s->use_sinks ? "sinks" : "", sw_buf, s->has_softcap ? "softcap" : "", diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/gfx942/attention_tiled_2d_q_load_gather.cpp b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/gfx942/attention_tiled_2d_q_load_gather.cpp index 47e7b7be0afc..6895f477491a 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/gfx942/attention_tiled_2d_q_load_gather.cpp +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/gfx942/attention_tiled_2d_q_load_gather.cpp @@ -87,6 +87,18 @@ void rocke_gfx942_attn2d_emit_q_load(rocke_gfx942_attn2d_build_ctx_t* ctx) /* q_desc = TensorDescriptor.naive("Q", lengths=[1<<30, NUM_QH, HD], * coord_names=("token","head","dim")). Stored on ctx for the * downstream Q gather (32x32 direct-global) + the epilogue store. */ + if(ctx->runtime_shape) + { + /* The head count is a kernarg, so the token stride is too. */ + static const char* const coord_names[3] = {"token", "head", "dim"}; + rocke_value_t* strides[3] + = {rocke_b_mul(b, ctx->num_qh_p, rocke_b_const_i32(b, ctx->HD)), NULL, NULL}; + const int static_strides[3] = {1, ctx->HD, 1}; + ctx->q_desc = &rocke_tensor_descriptor_naive_dynamic_mixed( + b, "Q", coord_names, 3, strides, static_strides) + ->base; + } + else { int lengths[3]; static const char* const coord_names[3] = {"token", "head", "dim"}; @@ -151,8 +163,7 @@ void rocke_gfx942_attn2d_emit_q_load(rocke_gfx942_attn2d_build_ctx_t* ctx) * guard is always true — omitting it avoids a redundant cmp+and on the hot path. */ { rocke_value_t* lt_pos = rocke_b_cmp_lt(b, q_pos_t, ctx->cur_batch_q_len); - rocke_value_t* lt_qh - = rocke_b_cmp_lt(b, qh_t, rocke_b_const_i32(b, (int64_t)ctx->NUM_QH)); + rocke_value_t* lt_qh = rocke_b_cmp_lt(b, qh_t, rocke_gfx942_attn2d_num_qh(ctx)); rocke_value_t* and_pos_qh = rocke_b_land(b, lt_pos, lt_qh); if(ctx->VALID_ROWS < ctx->BLOCK_M) { @@ -424,10 +435,9 @@ void rocke_gfx942_attn2d_emit_q_gather(rocke_gfx942_attn2d_build_ctx_t* ctx) rocke_b_mul(b, ctx->kv_head_idx, rocke_b_const_i32(b, (int64_t)ctx->NQK)), rocke_b_mod(b, q32_row, rocke_b_const_i32(b, (int64_t)ctx->NQK))); /* q32_mask = (q32_poscur_batch_q_len), - rocke_b_cmp_lt(b, q32_h, rocke_b_const_i32(b, (int64_t)ctx->NUM_QH))); + q32_mask = rocke_b_land(b, + rocke_b_cmp_lt(b, q32_pos, ctx->cur_batch_q_len), + rocke_b_cmp_lt(b, q32_h, rocke_gfx942_attn2d_num_qh(ctx))); q32_pos_safe = rocke_b_select(b, q32_mask, q32_pos, rocke_b_const_i32(b, 0)); q32_h_safe = rocke_b_select(b, q32_mask, q32_h, rocke_b_const_i32(b, 0)); diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/gfx942/attention_tiled_3d_gfx942_attn_tiled_3d_public_glue.cpp b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/gfx942/attention_tiled_3d_gfx942_attn_tiled_3d_public_glue.cpp index 394957d68e75..22a1c6de9cb3 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/gfx942/attention_tiled_3d_gfx942_attn_tiled_3d_public_glue.cpp +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/gfx942/attention_tiled_3d_gfx942_attn_tiled_3d_public_glue.cpp @@ -8,7 +8,7 @@ * SCOPE (this TU only): * - The two by-value spec structs' _default() / _validate() (__post_init__) * and every @property accessor (num_queries_per_kv, block_m=16, block_q, - * tile_size, dtype_ir, binary_search_iters, kernel_name via + * tile_size, dtype_ir, kernel_name via * kernel_name_join). (Python lines 91-172, 977-999.) * - supports_tiled_3d coverage gate (pure host predicate, static reasons). * (Python lines 175-231.) @@ -40,7 +40,6 @@ * the 16x16x16 f16+bf16 atoms present), which IS fully reproducible here. The * fallback reason strings mirror the Python text byte-for-byte. */ -#include #include #include #include @@ -243,25 +242,6 @@ const rocke_type_t* return rocke_bf16(); } -int rocke_unified_attention_3d_tiled_spec_binary_search_iters( - const rocke_unified_attention_3d_tiled_spec_t* s) -{ - double v; - int r; - if(s == NULL) - { - return -1; - } - if(s->num_seqs <= 0) - { - return 32; - } - /* max(1, ceil(log2(num_seqs + 1))) */ - v = ceil(log2((double)(s->num_seqs + 1))); - r = (int)v; - return r < 1 ? 1 : r; -} - int rocke_unified_attention_3d_tiled_spec_kernel_name( const rocke_unified_attention_3d_tiled_spec_t* s, char* buf, size_t cap) { @@ -298,7 +278,8 @@ int rocke_unified_attention_3d_tiled_spec_kernel_name( snprintf(d_part, sizeof(d_part), "d%d", s->head_size); snprintf(b_part, sizeof(b_part), "b%d", s->block_size); - snprintf(h_part, sizeof(h_part), "h%dkv%d", s->num_query_heads, s->num_kv_heads); + /* Head counts are kernargs; only the GQA ratio is baked. */ + snprintf(h_part, sizeof(h_part), "g%d", s->num_query_heads / s->num_kv_heads); snprintf(seg_part, sizeof(seg_part), "seg%d", s->num_segments); parts[np++] = "rocke_uattn3d_tiled_gfx942"; @@ -319,9 +300,10 @@ int rocke_unified_attention_3d_tiled_spec_kernel_name( { parts[np++] = "sinks"; } + /* The window size is a kernarg; only whether the body masks one is baked. */ if(s->sliding_window > 0) { - snprintf(sw_part, sizeof(sw_part), "sw%d", s->sliding_window); + snprintf(sw_part, sizeof(sw_part), "sw"); parts[np++] = sw_part; } if(s->has_softcap) @@ -394,10 +376,9 @@ const rocke_type_t* rocke_unified_attention_reduce_tiled_spec_dtype_ir( int rocke_unified_attention_reduce_tiled_spec_kernel_name( const rocke_unified_attention_reduce_tiled_spec_t* s, char* buf, size_t cap) { - /* kernel_name_join("rocke_uattn_reduce_tiled_gfx942", d, h, - * seg, dtype) */ + /* kernel_name_join("rocke_uattn_reduce_tiled_gfx942", d, seg, + * dtype); the head count is a kernarg. */ char d_part[32]; - char h_part[32]; char seg_part[32]; const char* parts[8]; size_t np = 0; @@ -410,12 +391,10 @@ int rocke_unified_attention_reduce_tiled_spec_kernel_name( } snprintf(d_part, sizeof(d_part), "d%d", s->head_size); - snprintf(h_part, sizeof(h_part), "h%d", s->num_query_heads); snprintf(seg_part, sizeof(seg_part), "seg%d", s->num_segments); parts[np++] = "rocke_uattn_reduce_tiled_gfx942"; parts[np++] = d_part; - parts[np++] = h_part; parts[np++] = seg_part; parts[np++] = (s->dtype != NULL) ? s->dtype : ""; @@ -669,7 +648,6 @@ bool rocke_gfx942_attn_tiled_3d_config_from_spec( out->PV_N_TILES = out->HD / ROCKE_ATTN3D_MFMA_N; out->THREADS = 64; - out->binary_search_iters = rocke_unified_attention_3d_tiled_spec_binary_search_iters(spec); if(out->PV_N_TILES < 0 || out->PV_N_TILES > 16) { diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/gfx942/attention_tiled_3d_gfx942_attn_tiled_3d_reduce.cpp b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/gfx942/attention_tiled_3d_gfx942_attn_tiled_3d_reduce.cpp index b810032b2ec0..5646bd143cb7 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/gfx942/attention_tiled_3d_gfx942_attn_tiled_3d_reduce.cpp +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/gfx942/attention_tiled_3d_gfx942_attn_tiled_3d_reduce.cpp @@ -17,7 +17,7 @@ * Pure f32 load / exp2 / store; no MFMA, no async DMA. wave64_reduce_max / * wave64_reduce_sum come from helper_helper_rocke.helpers.attention.h. The * builder-call sequence is a byte-identical translation of the Python span - * (which is itself a byte-for-byte port of the gfx950 reduce). + * (itself the gfx950 reduce with a runtime query head count). * * Peers (segment-kernel phases, ctx init, config_from_spec) live in sibling * translation units and are called only via the internal header. This file @@ -134,7 +134,6 @@ void rocke_gfx942_attention_tiled_3d_reduce_declare_and_prologue( const rocke_gfx942_attn_tiled_3d_config_t* cfg = &ctx->cfg; rocke_param_opts_t opts; - int NUM_QH = cfg->NUM_QH; int NUM_SEG = cfg->NUM_SEG; int HD = cfg->HD; @@ -187,6 +186,7 @@ void rocke_gfx942_attention_tiled_3d_reduce_declare_and_prologue( /* seq_lens_ptr : ptr, readonly, align=4 (unused in body) */ ctx->red_seq_lens = rocke_b_param(b, "seq_lens_ptr", rocke_ptr_type(b, rocke_i32(), "global"), &opts); + ctx->red_num_qh_p = rocke_b_param(b, "num_query_heads", rocke_i32(), NULL); /* ---- grid ids + thread (lines 1032-1034) ---- */ ctx->q_token = rocke_b_block_id_x(b); @@ -198,23 +198,36 @@ void rocke_gfx942_attention_tiled_3d_reduce_declare_and_prologue( ctx->zero_f = rocke_b_const_f32(b, 0.0); /* ---- descriptors (lines 1039-1053) ---- */ + /* The head count is a kernarg, so each token stride is too. */ { - int ml_lengths[3] = {1 << 30, NUM_QH, NUM_SEG}; + rocke_value_t* strides[3] + = {rocke_b_mul(b, ctx->red_num_qh_p, rocke_b_const_i32(b, NUM_SEG)), NULL, NULL}; + const int static_strides[3] = {1, NUM_SEG, 1}; const char* ml_coords[3] = {"token", "head", "seg"}; - ctx->ml_desc_red - = rocke_tensor_descriptor_naive(b, "segm_ml", ml_lengths, 3, NULL, ml_coords, 3); + ctx->ml_desc_red = &rocke_tensor_descriptor_naive_dynamic_mixed( + b, "segm_ml", ml_coords, 3, strides, static_strides) + ->base; } { - int sa_lengths[4] = {1 << 30, NUM_QH, NUM_SEG, HD}; + rocke_value_t* strides[4] + = {rocke_b_mul(b, ctx->red_num_qh_p, rocke_b_const_i32(b, NUM_SEG * HD)), + NULL, + NULL, + NULL}; + const int static_strides[4] = {1, NUM_SEG * HD, HD, 1}; const char* sa_coords[4] = {"token", "head", "seg", "dim"}; - ctx->seg_acc_desc_red - = rocke_tensor_descriptor_naive(b, "segm_output", sa_lengths, 4, NULL, sa_coords, 4); + ctx->seg_acc_desc_red = &rocke_tensor_descriptor_naive_dynamic_mixed( + b, "segm_output", sa_coords, 4, strides, static_strides) + ->base; } { - int out_lengths[3] = {1 << 30, NUM_QH, HD}; + rocke_value_t* strides[3] + = {rocke_b_mul(b, ctx->red_num_qh_p, rocke_b_const_i32(b, HD)), NULL, NULL}; + const int static_strides[3] = {1, HD, 1}; const char* out_coords[3] = {"token", "head", "dim"}; - ctx->out_desc_red - = rocke_tensor_descriptor_naive(b, "out", out_lengths, 3, NULL, out_coords, 3); + ctx->out_desc_red = &rocke_tensor_descriptor_naive_dynamic_mixed( + b, "out", out_coords, 3, strides, static_strides) + ->base; } /* ---- base_ml = ml_desc_red.offset(token=q_token, head=q_head, seg=0) ---- diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/gfx942/attention_tiled_3d_gfx942_attn_tiled_3d_segment_compute.cpp b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/gfx942/attention_tiled_3d_gfx942_attn_tiled_3d_segment_compute.cpp index 2510716e7b20..96a013265d71 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/gfx942/attention_tiled_3d_gfx942_attn_tiled_3d_segment_compute.cpp +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/gfx942/attention_tiled_3d_gfx942_attn_tiled_3d_segment_compute.cpp @@ -101,7 +101,7 @@ void rocke_gfx942_attention_tiled_3d_emit_loop_init( rocke_value_t* qh_mul = rocke_b_mul(b, ctx->kv_head_idx, rocke_b_const_i32(b, cfg->NQK)); rocke_value_t* qh_mod = rocke_b_mod(b, ctx->lane_col, rocke_b_const_i32(b, cfg->NQK)); rocke_value_t* qh = rocke_b_add(b, qh_mul, qh_mod); - rocke_value_t* qh_in = rocke_b_cmp_lt(b, qh, rocke_b_const_i32(b, cfg->NUM_QH)); + rocke_value_t* qh_in = rocke_b_cmp_lt(b, qh, ctx->num_qh_p); rocke_value_t* sink_h = rocke_b_global_load(b, ctx->sinks, qh, dtype, 2); rocke_value_t* sink_f = rocke_b_fmul(b, rocke_b_cast_to_f32(b, sink_h), ctx->rcp_ln2); rocke_value_t* sink_with_mask = rocke_b_select(b, qh_in, sink_f, ctx->neg_inf); @@ -223,14 +223,13 @@ void rocke_gfx942_attention_tiled_3d_emit_softmax_loop( qh_mod = rocke_b_mod(b, ctx->lane_col, rocke_b_const_i32(b, cfg->NQK)); ctx->row_qh = rocke_b_add(b, qh_mul, qh_mod); ok_a = rocke_b_cmp_lt(b, ctx->row_qp, ctx->cur_batch_q_len); - ok_b = rocke_b_cmp_lt(b, ctx->row_qh, rocke_b_const_i32(b, cfg->NUM_QH)); + ok_b = rocke_b_cmp_lt(b, ctx->row_qh, ctx->num_qh_p); ctx->row_ok = rocke_b_land(b, ok_a, ok_b); ctx->row_causal_lim = rocke_b_add(b, ctx->context_len, ctx->row_qp); ctx->row_alibi = NULL; if(cfg->USE_ALIBI) { - rocke_value_t* qh_ok - = rocke_b_cmp_lt(b, ctx->row_qh, rocke_b_const_i32(b, cfg->NUM_QH)); + rocke_value_t* qh_ok = rocke_b_cmp_lt(b, ctx->row_qh, ctx->num_qh_p); ctx->row_alibi = rocke_b_masked_global_load( b, ctx->alibi_slopes_ptr, ctx->row_qh, qh_ok, rocke_b_const_f32(b, 0.0), f32, 4); } diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/gfx942/attention_tiled_3d_gfx942_attn_tiled_3d_segment_setup.cpp b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/gfx942/attention_tiled_3d_gfx942_attn_tiled_3d_segment_setup.cpp index 386850a61551..d927f063900e 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/gfx942/attention_tiled_3d_gfx942_attn_tiled_3d_segment_setup.cpp +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/gfx942/attention_tiled_3d_gfx942_attn_tiled_3d_segment_setup.cpp @@ -318,6 +318,9 @@ void rocke_gfx942_attention_tiled_3d_declare_params( ctx->num_seqs_p = rocke_b_param(B, "num_seqs", rocke_i32(), NULL); ctx->bt_stride_p = rocke_b_param(B, "block_table_stride", rocke_i32(), NULL); ctx->qq_bias_stride0_p = rocke_b_param(B, "qq_bias_stride_0", rocke_i32(), NULL); + ctx->num_qh_p = rocke_b_param(B, "num_query_heads", rocke_i32(), NULL); + ctx->num_kv_p = rocke_b_param(B, "num_kv_heads", rocke_i32(), NULL); + ctx->sliding_window_p = rocke_b_param(B, "sliding_window", rocke_i32(), NULL); if(ctx->strided_kv) { ctx->k_strides = rocke_declare_strided_kv_params(B, "k"); @@ -337,7 +340,6 @@ static void rocke__build_paged_kv_desc(rocke_gfx942_attention_tiled_3d_build_ctx const int HD = CFG.HD; const int T = CFG.T; const int BS = CFG.BS; - const int NUM_KV = CFG.NUM_KV; /* seq_base = seq_idx * bt_stride_p (line 569) */ if(ctx->strided_kv) @@ -351,12 +353,12 @@ static void rocke__build_paged_kv_desc(rocke_gfx942_attention_tiled_3d_build_ctx * strides=[kv_stride_blk_b, kv_stride_tok_b, kv_stride_h_b, KV_BYTES], * coord_names=("physical_block","token","kv_head","dim")) (lines 570-575) */ { - const int lengths[4] = {1 << 24, BS, NUM_KV, HD}; - const int strides[4] - = {CFG.kv_stride_blk_b, CFG.kv_stride_tok_b, CFG.kv_stride_h_b, CFG.KV_BYTES}; + rocke_value_t* strides[4] = {ctx->kv_stride_blk_b, ctx->kv_stride_tok_b, NULL, NULL}; + const int static_strides[4] = {1, 1, CFG.kv_stride_h_b, CFG.KV_BYTES}; const char* coords[4] = {"physical_block", "token", "kv_head", "dim"}; - ctx->kv_base_desc - = rocke_tensor_descriptor_naive(B, "paged_kv_bytes", lengths, 4, strides, coords, 4); + ctx->kv_base_desc = &rocke_tensor_descriptor_naive_dynamic_mixed( + B, "paged_kv_bytes", coords, 4, strides, static_strides) + ->base; } /* Transform chain (lines 576-602). */ @@ -414,7 +416,6 @@ void rocke_gfx942_attention_tiled_3d_emit_prologue(rocke_gfx942_attention_tiled_ const int T = CFG.T; const int BLOCK_M = CFG.BLOCK_M; const int BLOCK_Q = CFG.BLOCK_Q; - const int NUM_QH = CFG.NUM_QH; const int NUM_SEG = CFG.NUM_SEG; const rocke_type_t* dtype = CFG.dtype; @@ -433,13 +434,9 @@ void rocke_gfx942_attention_tiled_3d_emit_prologue(rocke_gfx942_attention_tiled_ } else { - ctx->seq_idx = rocke_binary_search_seq_idx(B, - ctx->cu_q, - ctx->q_block_global_idx, - ctx->num_seqs_p, - BLOCK_Q, - CFG.binary_search_iters, - false); + rocke_value_t* trip = rocke_binary_search_seq_idx_trip_count(B, ctx->num_seqs_p); + ctx->seq_idx = rocke_binary_search_seq_idx_dyn( + B, ctx->cu_q, ctx->q_block_global_idx, ctx->num_seqs_p, BLOCK_Q, trip, false); } /* ---- cu_q bounds / per-sequence geometry (lines 344-350) ---- */ @@ -481,22 +478,33 @@ void rocke_gfx942_attention_tiled_3d_emit_prologue(rocke_gfx942_attention_tiled_ } /* ---- descriptors (lines 359-373) ---- */ + /* The head count is a kernarg, so each token stride is too. */ { - const int ml_lengths[3] = {1 << 30, NUM_QH, NUM_SEG}; + rocke_value_t* strides[3] + = {rocke_b_mul(B, ctx->num_qh_p, rocke_b_const_i32(B, NUM_SEG)), NULL, NULL}; + const int static_strides[3] = {1, NUM_SEG, 1}; const char* ml_coords[3] = {"token", "head", "seg"}; - ctx->ml_desc - = rocke_tensor_descriptor_naive(B, "segm_ml", ml_lengths, 3, NULL, ml_coords, 3); + ctx->ml_desc = &rocke_tensor_descriptor_naive_dynamic_mixed( + B, "segm_ml", ml_coords, 3, strides, static_strides) + ->base; } { - const int sa_lengths[4] = {1 << 30, NUM_QH, NUM_SEG, HD}; + rocke_value_t* strides[4] + = {rocke_b_mul(B, ctx->num_qh_p, rocke_b_const_i32(B, NUM_SEG * HD)), NULL, NULL, NULL}; + const int static_strides[4] = {1, NUM_SEG * HD, HD, 1}; const char* sa_coords[4] = {"token", "head", "seg", "dim"}; - ctx->seg_acc_desc - = rocke_tensor_descriptor_naive(B, "segm_output", sa_lengths, 4, NULL, sa_coords, 4); + ctx->seg_acc_desc = &rocke_tensor_descriptor_naive_dynamic_mixed( + B, "segm_output", sa_coords, 4, strides, static_strides) + ->base; } { - const int q_lengths[3] = {1 << 30, NUM_QH, HD}; + rocke_value_t* strides[3] + = {rocke_b_mul(B, ctx->num_qh_p, rocke_b_const_i32(B, HD)), NULL, NULL}; + const int static_strides[3] = {1, HD, 1}; const char* q_coords[3] = {"token", "head", "dim"}; - ctx->q_desc = rocke_tensor_descriptor_naive(B, "Q", q_lengths, 3, NULL, q_coords, 3); + ctx->q_desc = &rocke_tensor_descriptor_naive_dynamic_mixed( + B, "Q", q_coords, 3, strides, static_strides) + ->base; } /* seg_start_tile_pos = seg_idx*tps*T (line 375) */ @@ -525,7 +533,7 @@ void rocke_gfx942_attention_tiled_3d_emit_prologue(rocke_gfx942_attention_tiled_ ctx->one_f = rocke_b_const_f32(B, 1.0); ctx->rcp_ln2 = rocke_b_const_f32(B, 1.4426950408889634); ctx->qk_scale = rocke_b_fmul(B, ctx->scale_p, ctx->rcp_ln2); - ctx->sw_const = rocke_b_const_i32(B, CFG.SLIDING_WINDOW); + ctx->sw_const = ctx->sliding_window_p; ctx->z8 = rocke_b_zero_vec(B, dtype, 8); /* ---- Q -> LDS (lines 437-467) ---- */ @@ -583,6 +591,10 @@ void rocke_gfx942_attention_tiled_3d_emit_async_infra( ctx->key_rsrc = rocke_b_buffer_rsrc(B, ctx->key, ctx->big_bytes); ctx->value_rsrc = rocke_b_buffer_rsrc(B, ctx->value, ctx->big_bytes); } + /* The token and page strides scale with the runtime kv head count. */ + ctx->kv_stride_tok_b + = rocke_b_mul(B, ctx->num_kv_p, rocke_b_const_i32(B, CFG.HD * CFG.KV_BYTES)); + ctx->kv_stride_blk_b = rocke_b_mul(B, ctx->kv_stride_tok_b, rocke_b_const_i32(B, CFG.BS)); ctx->lane_half_base = rocke_b_mul(B, ctx->tid, rocke_b_const_i32(B, CFG.HALVES_PER_LANE)); ctx->K_lds_addr = rocke_b_smem_addr_of(B, ctx->K_lds); ctx->V_lds_addr = rocke_b_smem_addr_of(B, ctx->V_lds); @@ -599,7 +611,6 @@ void rocke_gfx942_attention_tiled_3d_emit_early_zero_fill( rocke_gfx942_attention_tiled_3d_build_ctx_t* ctx) { const int NQK = CFG.NQK; - const int NUM_QH = CFG.NUM_QH; const int PV_N_TILES = CFG.PV_N_TILES; int n, reg; @@ -622,7 +633,7 @@ void rocke_gfx942_attention_tiled_3d_emit_early_zero_fill( rocke_value_t* qh_r_mod = rocke_b_mod(B, row, rocke_b_const_i32(B, NQK)); rocke_value_t* qh_r = rocke_b_add(B, qh_r_mul, qh_r_mod); rocke_value_t* row_ok_a = rocke_b_cmp_lt(B, qp_r, ctx->cur_batch_q_len); - rocke_value_t* row_ok_b = rocke_b_cmp_lt(B, qh_r, rocke_b_const_i32(B, NUM_QH)); + rocke_value_t* row_ok_b = rocke_b_cmp_lt(B, qh_r, ctx->num_qh_p); rocke_value_t* row_ok = rocke_b_land(B, row_ok_a, row_ok_b); rocke_value_t* qp_r_safe = rocke_b_select(B, row_ok, qp_r, rocke_b_const_i32(B, 0)); rocke_value_t* qh_r_safe = rocke_b_select(B, row_ok, qh_r, rocke_b_const_i32(B, 0)); @@ -655,7 +666,7 @@ void rocke_gfx942_attention_tiled_3d_emit_early_zero_fill( rocke_value_t* qh_r_mod = rocke_b_mod(B, row, rocke_b_const_i32(B, NQK)); rocke_value_t* qh_r = rocke_b_add(B, qh_r_mul, qh_r_mod); rocke_value_t* row_ok_a = rocke_b_cmp_lt(B, qp_r, ctx->cur_batch_q_len); - rocke_value_t* row_ok_b = rocke_b_cmp_lt(B, qh_r, rocke_b_const_i32(B, NUM_QH)); + rocke_value_t* row_ok_b = rocke_b_cmp_lt(B, qh_r, ctx->num_qh_p); rocke_value_t* row_ok = rocke_b_land(B, row_ok_a, row_ok_b); rocke_value_t* qp_r_safe = rocke_b_select(B, row_ok, qp_r, rocke_b_const_i32(B, 0)); @@ -679,7 +690,6 @@ void rocke_gfx942_attention_tiled_3d_emit_early_zero_fill( void rocke_gfx942_attention_tiled_3d_emit_q_to_lds(rocke_gfx942_attention_tiled_3d_build_ctx_t* ctx) { const int NQK = CFG.NQK; - const int NUM_QH = CFG.NUM_QH; const int THREADS = CFG.THREADS; const int Q_VECS_PER_ROW = CFG.Q_VECS_PER_ROW; const int Q_VECS_PER_THREAD = CFG.Q_VECS_PER_THREAD; @@ -701,7 +711,7 @@ void rocke_gfx942_attention_tiled_3d_emit_q_to_lds(rocke_gfx942_attention_tiled_ rocke_value_t* qh_t_mod = rocke_b_mod(B, Q_row, rocke_b_const_i32(B, NQK)); rocke_value_t* qh_t = rocke_b_add(B, qh_t_mul, qh_t_mod); rocke_value_t* qmask_a = rocke_b_cmp_lt(B, q_pos_t, ctx->cur_batch_q_len); - rocke_value_t* qmask_b = rocke_b_cmp_lt(B, qh_t, rocke_b_const_i32(B, NUM_QH)); + rocke_value_t* qmask_b = rocke_b_cmp_lt(B, qh_t, ctx->num_qh_p); rocke_value_t* qmask_t = rocke_b_land(B, qmask_a, qmask_b); rocke_value_t* q_pos_safe = rocke_b_select(B, qmask_t, q_pos_t, rocke_b_const_i32(B, 0)); rocke_value_t* qh_safe = rocke_b_select(B, qmask_t, qh_t, rocke_b_const_i32(B, 0)); @@ -737,9 +747,10 @@ void rocke_gfx942_attention_tiled_3d_issue_k_load(rocke_gfx942_attention_tiled_3 /* The tile base goes into soffset; each copy adds its lane's offset within * the page. */ rocke_value_t* soff - = page_id != NULL ? rocke_gfx942_attention_tiled_3d_tile_base( - ctx, page_id, kv_tile_idx, CFG.kv_stride_blk_b, CFG.kv_stride_tok_b) - : ctx->zero_soff; + = page_id != NULL + ? rocke_gfx942_attention_tiled_3d_tile_base( + ctx, page_id, kv_tile_idx, ctx->kv_stride_blk_b, ctx->kv_stride_tok_b) + : ctx->zero_soff; rocke_value_t* buf_off_i32 = rocke_b_mul(B, buf_idx, rocke_b_const_i32(B, bytes_per_buf)); rocke_value_t* buf_off_i64 = rocke_b_zext(B, buf_off_i32, rocke_i64()); @@ -766,7 +777,7 @@ void rocke_gfx942_attention_tiled_3d_issue_k_load(rocke_gfx942_attention_tiled_3 rocke_value_t* src_half = rocke_gfx942_attention_tiled_3d_kv_src_linear(ctx, linear_half); rocke_value_t* tok = rocke_b_div(B, src_half, rocke_b_const_i32(B, CFG.HD)); - rocke_value_t* tok_off = rocke_b_mul(B, tok, rocke_b_const_i32(B, CFG.kv_stride_tok_b)); + rocke_value_t* tok_off = rocke_b_mul(B, tok, ctx->kv_stride_tok_b); rocke_value_t* head_off = rocke_b_mul(B, ctx->kv_head_idx, rocke_b_const_i32(B, CFG.kv_stride_h_b)); rocke_value_t* tok_head = rocke_b_add(B, tok_off, head_off); @@ -804,9 +815,10 @@ void rocke_gfx942_attention_tiled_3d_issue_v_load(rocke_gfx942_attention_tiled_3 /* The tile base goes into soffset; each copy adds its lane's offset within * the page. */ rocke_value_t* soff - = page_id != NULL ? rocke_gfx942_attention_tiled_3d_tile_base( - ctx, page_id, kv_tile_idx, CFG.kv_stride_blk_b, CFG.kv_stride_tok_b) - : ctx->zero_soff; + = page_id != NULL + ? rocke_gfx942_attention_tiled_3d_tile_base( + ctx, page_id, kv_tile_idx, ctx->kv_stride_blk_b, ctx->kv_stride_tok_b) + : ctx->zero_soff; rocke_value_t* buf_off_i32 = rocke_b_mul(B, buf_idx, rocke_b_const_i32(B, bytes_per_buf)); rocke_value_t* buf_off_i64 = rocke_b_zext(B, buf_off_i32, rocke_i64()); @@ -833,7 +845,7 @@ void rocke_gfx942_attention_tiled_3d_issue_v_load(rocke_gfx942_attention_tiled_3 rocke_value_t* src_half = rocke_gfx942_attention_tiled_3d_kv_src_linear(ctx, linear_half); rocke_value_t* tok = rocke_b_div(B, src_half, rocke_b_const_i32(B, CFG.HD)); - rocke_value_t* tok_off = rocke_b_mul(B, tok, rocke_b_const_i32(B, CFG.kv_stride_tok_b)); + rocke_value_t* tok_off = rocke_b_mul(B, tok, ctx->kv_stride_tok_b); rocke_value_t* head_off = rocke_b_mul(B, ctx->kv_head_idx, rocke_b_const_i32(B, CFG.kv_stride_h_b)); rocke_value_t* tok_head = rocke_b_add(B, tok_off, head_off); @@ -918,18 +930,19 @@ rocke_value_t* rocke_gfx942_attention_tiled_3d_tile_base(rocke_gfx942_attention_tiled_3d_build_ctx_t* ctx, rocke_value_t* page_id, rocke_value_t* kv_tile_idx, - int page_stride, - int token_stride) + rocke_value_t* page_stride, + rocke_value_t* token_stride) { - rocke_value_t* base = rocke_b_mul(B, page_id, rocke_b_const_i32(B, page_stride)); + rocke_value_t* base = rocke_b_mul(B, page_id, page_stride); rocke_value_t* in_page; + rocke_value_t* tile_stride; if(CFG.TILES_PER_PAGE == 1) { return base; } in_page = rocke_b_mod(B, kv_tile_idx, rocke_b_const_i32(B, CFG.TILES_PER_PAGE)); - return rocke_b_add( - B, base, rocke_b_mul(B, in_page, rocke_b_const_i32(B, CFG.T * token_stride))); + tile_stride = rocke_b_mul(B, token_stride, rocke_b_const_i32(B, CFG.T)); + return rocke_b_add(B, base, rocke_b_mul(B, in_page, tile_stride)); } /* _load_page_id_window(rel_page): lane l loads the id of the segment's page @@ -1034,9 +1047,12 @@ void rocke_gfx942_attention_tiled_3d_issue_wide_load( rocke_mem_opts_t opts = ROCKE_MEM_OPTS_INIT; if(page_id != NULL) { + /* Element strides (the wide loader indexes halves). */ + rocke_value_t* tok_e = rocke_b_mul(B, ctx->num_kv_p, rocke_b_const_i32(B, HD)); + rocke_value_t* blk_e = rocke_b_mul(B, tok_e, rocke_b_const_i32(B, CFG.BS)); rocke_value_t* page_base = rocke_gfx942_attention_tiled_3d_tile_base( - ctx, page_id, kv_tile_idx, CFG.BS * CFG.NUM_KV * HD, CFG.NUM_KV * HD); - rocke_value_t* row_off = rocke_b_mul(B, row, rocke_b_const_i32(B, CFG.NUM_KV * HD)); + ctx, page_id, kv_tile_idx, blk_e, tok_e); + rocke_value_t* row_off = rocke_b_mul(B, row, tok_e); rocke_value_t* head_off = rocke_b_mul(B, ctx->kv_head_idx, rocke_b_const_i32(B, HD)); rocke_value_t* in_page = rocke_b_add(B, rocke_b_add(B, row_off, head_off), col); elem_off = rocke_b_add(B, page_base, in_page); diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/dsl_docs/reference/file_index.md b/dnn-providers/hip-kernel-provider/rocke/platform/dsl_docs/reference/file_index.md index a5612cca5b1f..26bb95cb9544 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/dsl_docs/reference/file_index.md +++ b/dnn-providers/hip-kernel-provider/rocke/platform/dsl_docs/reference/file_index.md @@ -55,7 +55,7 @@ A by-file map of the `rocke` package. Symbols listed are the primary contents (p | `helpers/schedule.py` | `SchedulePolicy`. `sched_group_barrier`, `s_setprio`, named schedule flavors. | | `helpers/pipeline.py` | `SoftwarePipeline.run_ping_pong(...)`. | | `helpers/epilogues.py` | `DirectEpilogue`, `CShuffleEpilogue`. | -| `helpers/attention.py` | `Attention2DConfig`, `Attention3DConfig`, `OnlineSoftmaxState`, `PagedKvDescriptor`, `apply_softcap_log2`, `apply_softcap_scalar`, `binary_search_seq_idx`, `binary_search_seq_idx_iters`, `causal_mask`, `sliding_window_mask`, `mfma_16x16x16_for_dtype`, `mfma_16x16x32_for_dtype`, `select_2d_config`, `select_3d_config`, `use_2d_kernel`, `warp_xor_reduce_max`, `warp_xor_reduce_sum`. | +| `helpers/attention.py` | `Attention2DConfig`, `Attention3DConfig`, `OnlineSoftmaxState`, `PagedKvDescriptor`, `apply_softcap_log2`, `apply_softcap_scalar`, `binary_search_seq_idx`, `binary_search_seq_idx_iters`, `binary_search_seq_idx_trip_count`, `causal_mask`, `sliding_window_mask`, `mfma_16x16x16_for_dtype`, `mfma_16x16x32_for_dtype`, `select_2d_config`, `select_3d_config`, `use_2d_kernel`, `warp_xor_reduce_max`, `warp_xor_reduce_sum`. | | `helpers/compile.py` | `compile_kernel`, `KernelArtifact`. | | `helpers/manifest.py` | `MANIFEST_SCHEMA = "ck.dsl.example.manifest/v1"`, `attention_args_signature`, `gemm_args_signature`, `make_attention_manifest`, `make_conv_manifest`, `make_gemm_manifest`, `make_simple_op_manifest`, `write_artifact`. | | `helpers/tensor_view.py` | `TensorDescriptor`, `TensorView`, `TileWindow`, `TensorCoordinate`, `BufferResource`; constructors `make_global_view`, `make_lds_view`, `make_buffer_resource`, `make_buffer_view`, `make_naive_tensor_descriptor_packed`, `make_naive_tensor_view_packed`, `make_tile_window`, `make_tensor_coordinate`, `move_tensor_coordinate`, `view_from_transforms_descriptor`. | diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/__init__.py b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/__init__.py index d21338f7ef50..f4af461eb299 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/__init__.py +++ b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/__init__.py @@ -111,6 +111,7 @@ apply_softcap_scalar, binary_search_seq_idx, binary_search_seq_idx_iters, + binary_search_seq_idx_trip_count, apply_attention_mask, causal_mask, mfma_16x16x16_for_dtype, @@ -434,6 +435,7 @@ "BINARY_SEARCH_FALLBACK_ITERS", "binary_search_seq_idx", "binary_search_seq_idx_iters", + "binary_search_seq_idx_trip_count", "apply_attention_mask", "causal_mask", # MFMA GEMM inner helper diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/attention.py b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/attention.py index 8695eb4537fd..50776ac8f618 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/attention.py +++ b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/attention.py @@ -699,6 +699,19 @@ def binary_search_seq_idx_iters(num_seqs: int, *, floor: int = 1) -> int: return max(floor, int(math.ceil(math.log2(num_seqs + 1)))) +def binary_search_seq_idx_trip_count(b: IRBuilder, num_seqs: Value) -> Value: + """Runtime trip count for :func:`binary_search_seq_idx`: ``bit_length(num_seqs)``, + at least 1, read off the exponent of ``float(num_seqs)``. + + Equals :func:`binary_search_seq_idx_iters` below 2**24 sequences. Past that the + conversion may round up a power of two and add one iteration, which the + search's done guard turns into a no-op. Lets one kernel serve every batch size. + """ + bits = b.bitcast(b.sitofp_f32(num_seqs), I32) + exponent = b.lshr(bits, b.const_i32(23)) + return b.smax(b.sub(exponent, b.const_i32(126)), b.const_i32(1)) + + def binary_search_seq_idx( b: IRBuilder, cu_q: Value, @@ -706,7 +719,7 @@ def binary_search_seq_idx( num_seqs: Value, *, block_q: int, - iterations: int, + iterations: int | Value, per_token: bool = False, ) -> Value: """Triton-style binary search for the seq_idx for this q_block. @@ -716,7 +729,8 @@ def binary_search_seq_idx( invariant is ``cu_q[i] // BLOCK_Q + i <= target`` (i.e. the cumulative Q-block count up to sequence ``i``). The caller specializes ``iterations`` from the known problem batch size with - :func:`binary_search_seq_idx_iters`. + :func:`binary_search_seq_idx_iters`, or passes the runtime bound from + :func:`binary_search_seq_idx_trip_count`. P49: ``per_token=True`` switches to per-token mode (``cu_q[s] <= q_token < cu_q[s+1]``) used by ``fmha_paged_prefill``, @@ -729,7 +743,7 @@ def binary_search_seq_idx( bq = b.const_i32(block_q) loop = b.scf_for_iter( b.const_i32(0), - b.const_i32(iterations), + iterations if isinstance(iterations, Value) else b.const_i32(iterations), b.const_i32(1), [("left", b.const_i32(0)), ("right", num_seqs)], iv_name="bs_i", diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/transforms.py b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/transforms.py index bf4d34487075..4bd703678cec 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/transforms.py +++ b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/transforms.py @@ -70,7 +70,7 @@ from __future__ import annotations from dataclasses import dataclass, replace -from typing import Any, Dict, List, Optional, Sequence, Tuple +from typing import Any, Dict, List, Optional, Sequence, Tuple, Union from ..core.ir import F16, I32, I64, IRBuilder, Type, Value from . import magic_division @@ -1752,19 +1752,26 @@ def move( @dataclass class DynamicTensorDescriptor(TensorDescriptor): - """A :class:`TensorDescriptor` whose base strides are runtime SSA Values. + """A :class:`TensorDescriptor` whose base strides may be runtime SSA Values. Constructed via :meth:`DynamicTensorDescriptor.create`. The transform - chain and ``offset()`` work identically to the static descriptor except - the final stride-multiply in ``offset()`` uses runtime Values instead of - ``b.const_i32(stride)``. + chain works identically to the static descriptor; the final stride + reduction multiplies a base coord by its runtime ``Value`` stride, and an + ``int`` stride follows the static rule (no multiply for 1). - ``dynamic_strides`` holds one ``Value`` per base coord, in the same - order as ``base_names``. ``base_strides`` is set to all-ones (unused - sentinel) so the parent's ``offset()`` is not called directly. + ``dynamic_strides`` holds one entry per base coord, in the same order as + ``base_names``: the runtime ``Value``, or ``None`` where ``base_strides`` + holds the compile-time stride (1 at every runtime position). """ - dynamic_strides: Tuple[Value, ...] = () + dynamic_strides: Tuple[Optional[Value], ...] = () + + def _term(self, b: IRBuilder, i: int, value: Value) -> Value: + stride_v = self.dynamic_strides[i] + if stride_v is not None: + return b.mul(value, stride_v) + stride = self.base_strides[i] + return value if stride == 1 else b.mul(value, b.const_i32(stride)) def offset( self, @@ -1774,49 +1781,139 @@ def offset( coords = self._run_chain(b, upper_values) off: Optional[Value] = None valid: Optional[Value] = None - for name, stride_v in zip(self.base_names, self.dynamic_strides): + for i, name in enumerate(self.base_names): if name not in coords: raise ValueError( f"after chain, base coord {name!r} not in {sorted(coords.keys())}" ) c = coords[name] valid = _and(b, valid, c.valid) - term = b.mul(c.value, stride_v) + term = self._term(b, i, c.value) off = term if off is None else b.add(off, term) if off is None: off = b.const_i32(0) return off, valid + def offset_i64_split( + self, + b: IRBuilder, + base_coord: str, + **upper_values: Value, + ) -> Tuple[Value, Value, Optional[Value]]: + """:meth:`TensorDescriptor.offset_i64_split` with runtime strides; the + ``base_coord`` term widens a runtime stride to i64.""" + coords = self._run_chain(b, upper_values) + base_i64: Optional[Value] = None + within: Optional[Value] = None + valid: Optional[Value] = None + for i, name in enumerate(self.base_names): + if name not in coords: + raise ValueError( + f"after chain, base coord {name!r} not in {sorted(coords.keys())}" + ) + c = coords[name] + valid = _and(b, valid, c.valid) + if name == base_coord: + base_w = b.zext(b.to_sgpr_u32(c.value), I64) + stride_v = self.dynamic_strides[i] + if stride_v is not None: + base_i64 = b.mul(base_w, b.zext(stride_v, I64)) + else: + base_i64 = b.mul(base_w, b.const_i64(int(self.base_strides[i]))) + else: + term = self._term(b, i, c.value) + within = term if within is None else b.add(within, term) + if base_i64 is None: + raise ValueError( + f"offset_i64_split: base_coord {base_coord!r} not among " + f"base coords {list(self.base_names)}" + ) + if within is None: + within = b.const_i32(0) + return base_i64, within, valid + + def offset_i64( + self, + b: IRBuilder, + i64_coord: str, + **upper_values: Value, + ) -> Tuple[Value, Optional[Value]]: + """:meth:`TensorDescriptor.offset_i64` with runtime strides.""" + coords = self._run_chain(b, upper_values) + off: Optional[Value] = None + valid: Optional[Value] = None + for i, name in enumerate(self.base_names): + if name not in coords: + raise ValueError( + f"after chain, base coord {name!r} not in {sorted(coords.keys())}" + ) + c = coords[name] + valid = _and(b, valid, c.valid) + if name == i64_coord: + wide = b.zext(c.value, I64) + stride_v = self.dynamic_strides[i] + if stride_v is not None: + term = b.mul(wide, b.zext(stride_v, I64)) + else: + term = b.mul(wide, b.const_i64(int(self.base_strides[i]))) + else: + term = b.zext(self._term(b, i, c.value), I64) + off = term if off is None else b.add(off, term) + if off is None: + off = b.const_i64(0) + return off, valid + + def move( + self, + b: IRBuilder, + prev_offset: Value, + old_values: Dict[str, Value], + delta_values: Dict[str, Value], + ) -> Tuple[Value, Optional[Value]]: + """:meth:`TensorDescriptor.move` with runtime strides.""" + delta, new = self._run_chain_delta(b, old_values, delta_values) + off_delta: Optional[Value] = None + valid: Optional[Value] = None + for i, name in enumerate(self.base_names): + if name not in delta: + raise ValueError( + f"after chain, base coord {name!r} not in {sorted(delta.keys())}" + ) + valid = _and(b, valid, new[name].valid) + term = self._term(b, i, delta[name].value) + off_delta = term if off_delta is None else b.add(off_delta, term) + if off_delta is None: + return prev_offset, valid + return b.add(prev_offset, off_delta), valid + @classmethod def create( cls, name: str, *, coord_names: Sequence[str], - strides: Sequence[Value], + strides: Sequence[Union[Value, int]], ) -> "DynamicTensorDescriptor": - """Build a dynamic descriptor from runtime SSA strides. - - ``coord_names`` and ``strides`` must have the same length. - ``strides`` are i32 SSA Values representing the row-major strides - (in elements) for each coordinate, computed host-side and passed - as kernel arguments. + """Build a descriptor from per-coord strides, each a runtime i32 SSA + ``Value`` (in elements, usually from kernel arguments) or an ``int``. - The descriptor can then be extended with ``.transform(...)`` just - like a static descriptor. + ``coord_names`` and ``strides`` must have the same length. The + descriptor can then be extended with ``.transform(...)`` just like a + static descriptor. """ if len(coord_names) != len(strides): raise ValueError("coord_names and strides length mismatch") coord_names = tuple(coord_names) - strides = tuple(strides) + dynamic = tuple(s if isinstance(s, Value) else None for s in strides) + static = tuple(1 if isinstance(s, Value) else int(s) for s in strides) return cls( name=name, base_names=coord_names, base_lengths=tuple(0 for _ in coord_names), - base_strides=tuple(1 for _ in coord_names), + base_strides=static, chain=(), upper_names=coord_names, - dynamic_strides=strides, + dynamic_strides=dynamic, ) diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_transforms_dynamic.py b/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_transforms_dynamic.py index d0284324d7e7..95aa5468a120 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_transforms_dynamic.py +++ b/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_transforms_dynamic.py @@ -505,6 +505,18 @@ def test_no_bounds_no_predicate(self): # --------------------------------------------------------------------------- +def _depends_on(value, source) -> bool: + """True when ``source`` feeds ``value`` through some chain of ops.""" + stack = [value] + while stack: + v = stack.pop() + if v is source: + return True + if getattr(v, "op", None) is not None: + stack.extend(v.op.operands) + return False + + class TestDynamicTensorDescriptor: def test_offset_multiplies_by_runtime_strides(self): b = _builder() @@ -618,3 +630,69 @@ def test_strides_are_values(self): "t", coord_names=("a", "b"), strides=(s0, s1) ) assert all(isinstance(s, Value) for s in desc.dynamic_strides) + + def test_int_strides_follow_the_static_rule(self): + b = _builder() + a, c, d, s = _params(b, "a", "c", "d", "s") + desc = DynamicTensorDescriptor.create( + "t", coord_names=("a", "c", "d"), strides=(s, 3, 1) + ) + assert desc.dynamic_strides == (s, None, None) + start = len(b.kernel.body.ops) + off, valid = desc.offset(b, a=a, c=c, d=d) + ops = _new_ops(b, start) + muls = [op for op in ops if op.name == "arith.mul"] + # The Value stride and int stride 3 multiply; int stride 1 does not. + assert len(muls) == 2 + assert muls[0].operands == [a, s] + assert muls[1].operands[0] is c + assert muls[1].operands[1].op.attrs["value"] == 3 + env = {id(a): 5, id(c): 7, id(d): 11, id(s): 13} + _evaluate(ops, env) + assert env[id(off)] == 5 * 13 + 7 * 3 + 11 + assert valid is None + + def test_offset_i64_split_widens_a_runtime_base_stride(self): + b = _builder() + blk, tok, d, s_blk, s_tok = _params(b, "blk", "tok", "d", "s_blk", "s_tok") + desc = DynamicTensorDescriptor.create( + "kv", coord_names=("blk", "tok", "d"), strides=(s_blk, s_tok, 2) + ) + base, within, _ = desc.offset_i64_split(b, "blk", blk=blk, tok=tok, d=d) + # base = zext(to_sgpr(blk)) * zext(s_blk), in i64. + wide_blk, wide_stride = base.op.operands + assert wide_stride.op.operands == [s_blk] + assert _depends_on(wide_blk, blk) + # within = tok * s_tok + d * 2, in i32. + tok_term, d_term = within.op.operands + assert tok_term.op.operands == [tok, s_tok] + assert d_term.op.operands[0] is d + assert d_term.op.operands[1].op.attrs["value"] == 2 + + def test_offset_i64_widens_a_runtime_stride(self): + b = _builder() + blk, d, s_blk = _params(b, "blk", "d", "s_blk") + desc = DynamicTensorDescriptor.create( + "kv", coord_names=("blk", "d"), strides=(s_blk, 1) + ) + off, _ = desc.offset_i64(b, "blk", blk=blk, d=d) + blk_term, d_term = off.op.operands + assert blk_term.op.operands[0].op.operands == [blk] + assert blk_term.op.operands[1].op.operands == [s_blk] + # The int stride-1 term is the coord itself, widened. + assert d_term.op.operands == [d] + + def test_move_multiplies_deltas_by_runtime_strides(self): + b = _builder() + a, c, da, dc, s = _params(b, "a", "c", "da", "dc", "s") + desc = DynamicTensorDescriptor.create( + "t", coord_names=("a", "c"), strides=(s, 3) + ) + start = len(b.kernel.body.ops) + prev, _ = desc.offset(b, a=a, c=c) + off, valid = desc.move(b, prev, {"a": a, "c": c}, {"a": da, "c": dc}) + assert valid is None + env = {id(a): 5, id(c): 7, id(da): 2, id(dc): 4, id(s): 13} + _evaluate(_new_ops(b, start), env) + # The runtime stride scales the delta, not the 1 held in base_strides. + assert env[id(off)] == (5 + 2) * 13 + (7 + 4) * 3 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 b3afb40ec9e3..744e41dd6e7e 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 @@ -4,74 +4,74 @@ "cases": { "attention/gfx942/2d_bf16_d128_b16_decode_i64": { "arch": "gfx942", - "bytes": 452637, + "bytes": 454204, "family": "attention", - "kernel_name": "rocke_uattn2d_tiled_d128_b16_t64_h32kv8_bf16_wpe2_regpv", - "sha256": "2ca52b461a729bfaad2594dc8e6e5ad609c562aa11898f3de0a889d400fa0d7b", + "kernel_name": "rocke_uattn2d_tiled_d128_b16_t64_g4_bf16_i64_wpe2_regpv", + "sha256": "e4df972977ebf80ab8a8e70a5aa2354fe0aa1235fd11128c1a144e76dbce8012", "status": "ok" }, "attention/gfx942/2d_bf16_d128_b64": { "arch": "gfx942", - "bytes": 393929, + "bytes": 395436, "family": "attention", - "kernel_name": "rocke_uattn2d_tiled_d128_b64_h4kv2_bf16", - "sha256": "f2e657c99be9a646228a2ab2a8a04610437eab7fc46b15e8345bb6b6a1984f5f", + "kernel_name": "rocke_uattn2d_tiled_d128_b64_g2_bf16", + "sha256": "ea26b2967c5ec49452dbcf90307593dd83f12518aded03c7b5ba9ad9e81057f6", "status": "ok" }, "attention/gfx942/2d_fp16_d128_b64": { "arch": "gfx942", - "bytes": 382259, + "bytes": 383766, "family": "attention", - "kernel_name": "rocke_uattn2d_tiled_d128_b64_h4kv2_fp16", - "sha256": "09324ad9a9c0526bc91ece6ea5886946d97dd6b41d9e634652222a60d6889d74", + "kernel_name": "rocke_uattn2d_tiled_d128_b64_g2_fp16", + "sha256": "6b2f24ff38a2bd28e50c2a73e42571fc02c7034b6b8b1a72a92c752fa19eb59b", "status": "ok" }, "attention/gfx942/3d_bf16_d128_b64": { "arch": "gfx942", - "bytes": 355976, + "bytes": 358681, "family": "attention", - "kernel_name": "rocke_uattn3d_tiled_gfx942_d128_b64_h4kv2_seg4_bf16", - "sha256": "e4128a6e834957e3eaefa368f1409d69125dfca44ffa4633a684ab378f0f5bb2", + "kernel_name": "rocke_uattn3d_tiled_gfx942_d128_b64_g2_seg4_bf16", + "sha256": "9c14d2a299f5bf54629246deed221e8aec0198bfb80804c8ca21fce67f63d5c4", "status": "ok" }, "attention/gfx942/3d_fp16_d128_b64": { "arch": "gfx942", - "bytes": 343932, + "bytes": 346637, "family": "attention", - "kernel_name": "rocke_uattn3d_tiled_gfx942_d128_b64_h4kv2_seg4_fp16", - "sha256": "02a8532f212341b8ac080bf4f838859e782d18423c6edb799d5048157970414c", + "kernel_name": "rocke_uattn3d_tiled_gfx942_d128_b64_g2_seg4_fp16", + "sha256": "8891d9a78aad2540771b5b1b3253683be81401088ce9f6ca7b7e01e5840013fd", "status": "ok" }, "attention/gfx942/reduce_bf16_d128": { "arch": "gfx942", - "bytes": 9617, + "bytes": 9773, "family": "attention", - "kernel_name": "rocke_uattn_reduce_tiled_gfx942_d128_h4_seg4_bf16", - "sha256": "5ecc6f410ad0e6b674e4ea99868463538a75fe5c7c81e294f3db697e6f134eba", + "kernel_name": "rocke_uattn_reduce_tiled_gfx942_d128_seg4_bf16", + "sha256": "f978e1889a9b57c2c544a31e393a75ffd729837a967d44c4dd6361582459870e", "status": "ok" }, "attention/gfx942/reduce_bf16_d128_own_module": { "arch": "gfx942", - "bytes": 9617, + "bytes": 9773, "family": "attention", - "kernel_name": "rocke_uattn_reduce_tiled_gfx942_d128_h4_seg4_bf16", - "sha256": "5ecc6f410ad0e6b674e4ea99868463538a75fe5c7c81e294f3db697e6f134eba", + "kernel_name": "rocke_uattn_reduce_tiled_gfx942_d128_seg4_bf16", + "sha256": "f978e1889a9b57c2c544a31e393a75ffd729837a967d44c4dd6361582459870e", "status": "ok" }, "attention/gfx942/reduce_fp16_d128": { "arch": "gfx942", - "bytes": 9605, + "bytes": 9761, "family": "attention", - "kernel_name": "rocke_uattn_reduce_tiled_gfx942_d128_h4_seg4_fp16", - "sha256": "9793be461b2096fe71e56defe36858dd58257d0f8243d9dece3e53f4d6123207", + "kernel_name": "rocke_uattn_reduce_tiled_gfx942_d128_seg4_fp16", + "sha256": "afa3c4ede02acabade859cd09cc57a94fc27a7c6e893976954489f68dd932a28", "status": "ok" }, "attention/gfx942/reduce_lse_bf16_d128": { "arch": "gfx942", - "bytes": 10497, + "bytes": 10668, "family": "attention", - "kernel_name": "rocke_uattn_reduce_tiled_gfx942_d128_h4_seg4_bf16_lse", - "sha256": "09854df6ad25883f6bf6c43eecb394c9c210b244924394f6976c308dcdad13b0", + "kernel_name": "rocke_uattn_reduce_tiled_gfx942_d128_seg4_bf16_lse", + "sha256": "d8db2fd0d06898c27ac98f05951fd25bfa27f797f7e44eab109470c6e9ec6a39", "status": "ok" }, "attention/gfx950/2d_bf16_d128_b64": { @@ -316,18 +316,18 @@ }, "attention_d128_swa/gfx942/4warp_gqa_fold": { "arch": "gfx942", - "bytes": 408354, + "bytes": 408566, "family": "attention_d128_swa", - "kernel_name": "rocke_uattn2d_tiled_d128_b16_t32_h32kv8_bf16_sw4096_wpe2_mw32_4wgqa_fold", - "sha256": "f8147a55dfedf5d55283f17fdb43bd6e43acfaec80fb795466788d058071c2bb", + "kernel_name": "rocke_uattn2d_tiled_d128_b16_t32_g4_bf16_sw_wpe2_mw32_4wgqa_fold", + "sha256": "0c3c01d6023e7fece924282414af5eadda2b46bd9265a4711cea71b373b9439c", "status": "ok" }, "attention_d256/gfx942/4warp_gqa": { "arch": "gfx942", - "bytes": 467380, + "bytes": 469322, "family": "attention_d256", - "kernel_name": "rocke_uattn2d_tiled_d256_b16_t32_h16kv2_bf16_wpe2_mw32_regpv_4wgqa", - "sha256": "563eb76369b9ea9f8acba900077111767a9a7f43b4a4a130f99d161c01d3d7f7", + "kernel_name": "rocke_uattn2d_tiled_d256_b16_t32_g8_bf16_wpe2_mw32_regpv_4wgqa", + "sha256": "ecc9ec09432cc43c26dae760adc973cb0376ad87ebbb186f459dff0264abf9a2", "status": "ok" }, "attention_d256/gfx950/pad_interleave": { @@ -2393,74 +2393,74 @@ "cases": { "attention/gfx942/2d_bf16_d128_b16_decode_i64": { "arch": "gfx942", - "bytes": 453249, + "bytes": 454816, "family": "attention", - "kernel_name": "rocke_uattn2d_tiled_d128_b16_t64_h32kv8_bf16_wpe2_regpv", - "sha256": "5a6d4b54fb70d62786ec4c7fa3c6ce0da2db3f907aa20103244378dd5a413318", + "kernel_name": "rocke_uattn2d_tiled_d128_b16_t64_g4_bf16_i64_wpe2_regpv", + "sha256": "f370a0bf3c16dcc6a3e9e9553cfd619e3fadd52f8c88c2566d72830dff70c636", "status": "ok" }, "attention/gfx942/2d_bf16_d128_b64": { "arch": "gfx942", - "bytes": 394025, + "bytes": 395532, "family": "attention", - "kernel_name": "rocke_uattn2d_tiled_d128_b64_h4kv2_bf16", - "sha256": "864c0ee45f501c76c8a167a2cb3e36d16d54ce09cf7858e11d2795384d8b0140", + "kernel_name": "rocke_uattn2d_tiled_d128_b64_g2_bf16", + "sha256": "eed9f53bac44445aaaec5171e04c6d75efd045451b9a7342405fe6215c6d4522", "status": "ok" }, "attention/gfx942/2d_fp16_d128_b64": { "arch": "gfx942", - "bytes": 382355, + "bytes": 383862, "family": "attention", - "kernel_name": "rocke_uattn2d_tiled_d128_b64_h4kv2_fp16", - "sha256": "72dc6aad6e3a1db8786de0cf7edb1dfaa3637b3d0efd154253cf92c8d2d837ae", + "kernel_name": "rocke_uattn2d_tiled_d128_b64_g2_fp16", + "sha256": "3f87ed04f831ac8b2a75218e9b6990a45a5cd31e8a4bb517723b86f25d6d359d", "status": "ok" }, "attention/gfx942/3d_bf16_d128_b64": { "arch": "gfx942", - "bytes": 356076, + "bytes": 358781, "family": "attention", - "kernel_name": "rocke_uattn3d_tiled_gfx942_d128_b64_h4kv2_seg4_bf16", - "sha256": "38f97d98aa14759dda281c4952e6571d4e1c6ff21ac08274473091c0e0bb0146", + "kernel_name": "rocke_uattn3d_tiled_gfx942_d128_b64_g2_seg4_bf16", + "sha256": "fad6de65bd50402a69c278452eaf61951e7fe4db3a6eba8dafece7d108a6d1ab", "status": "ok" }, "attention/gfx942/3d_fp16_d128_b64": { "arch": "gfx942", - "bytes": 344028, + "bytes": 346733, "family": "attention", - "kernel_name": "rocke_uattn3d_tiled_gfx942_d128_b64_h4kv2_seg4_fp16", - "sha256": "2264bce69fae37533413945b5a0fa9778737b197ee8f3791043dbd13865736e4", + "kernel_name": "rocke_uattn3d_tiled_gfx942_d128_b64_g2_seg4_fp16", + "sha256": "902f7f7dc7eb0b2a5a21292899e44e816fb7112c844b6e199e9a087125df11b1", "status": "ok" }, "attention/gfx942/reduce_bf16_d128": { "arch": "gfx942", - "bytes": 9624, + "bytes": 9780, "family": "attention", - "kernel_name": "rocke_uattn_reduce_tiled_gfx942_d128_h4_seg4_bf16", - "sha256": "7f3bea52ba6bedf52e5219615a0988da1f6870faba69d9d02cfc81cffca477f5", + "kernel_name": "rocke_uattn_reduce_tiled_gfx942_d128_seg4_bf16", + "sha256": "be513de59d3fde4a6d0a372effed327b37d2d5de4b48161a6de5186dc692d721", "status": "ok" }, "attention/gfx942/reduce_bf16_d128_own_module": { "arch": "gfx942", - "bytes": 9624, + "bytes": 9780, "family": "attention", - "kernel_name": "rocke_uattn_reduce_tiled_gfx942_d128_h4_seg4_bf16", - "sha256": "7f3bea52ba6bedf52e5219615a0988da1f6870faba69d9d02cfc81cffca477f5", + "kernel_name": "rocke_uattn_reduce_tiled_gfx942_d128_seg4_bf16", + "sha256": "be513de59d3fde4a6d0a372effed327b37d2d5de4b48161a6de5186dc692d721", "status": "ok" }, "attention/gfx942/reduce_fp16_d128": { "arch": "gfx942", - "bytes": 9612, + "bytes": 9768, "family": "attention", - "kernel_name": "rocke_uattn_reduce_tiled_gfx942_d128_h4_seg4_fp16", - "sha256": "d6db7088aa420dd3e8f81576f4ccf2d5a35f9bb211695ec3b9c2c38ff33139f1", + "kernel_name": "rocke_uattn_reduce_tiled_gfx942_d128_seg4_fp16", + "sha256": "3d7c19d0b0422eeb6a1e558bf0c90ba458dfc659b5c49712b706960d39998a1e", "status": "ok" }, "attention/gfx942/reduce_lse_bf16_d128": { "arch": "gfx942", - "bytes": 10504, + "bytes": 10675, "family": "attention", - "kernel_name": "rocke_uattn_reduce_tiled_gfx942_d128_h4_seg4_bf16_lse", - "sha256": "edbfaa7b91c851b499b0cc4fac7ceaa9e3c4922e60dd6ff464bc18f251f37843", + "kernel_name": "rocke_uattn_reduce_tiled_gfx942_d128_seg4_bf16_lse", + "sha256": "1e43a204d3a86bdc3609fa571e5da182752e9fbbb8504d1e9a02413957975580", "status": "ok" }, "attention/gfx950/2d_bf16_d128_b64": { @@ -2705,18 +2705,18 @@ }, "attention_d128_swa/gfx942/4warp_gqa_fold": { "arch": "gfx942", - "bytes": 408361, + "bytes": 408573, "family": "attention_d128_swa", - "kernel_name": "rocke_uattn2d_tiled_d128_b16_t32_h32kv8_bf16_sw4096_wpe2_mw32_4wgqa_fold", - "sha256": "416175deefad07ef1f044052f5b5a7b54b97e6b08c05d4ce3923c958ea28565f", + "kernel_name": "rocke_uattn2d_tiled_d128_b16_t32_g4_bf16_sw_wpe2_mw32_4wgqa_fold", + "sha256": "c4b61f172d05ecb10e38d9c6cd7a29fcb5518f366f7a6d77a6e0899fe0a8d03b", "status": "ok" }, "attention_d256/gfx942/4warp_gqa": { "arch": "gfx942", - "bytes": 467387, + "bytes": 469329, "family": "attention_d256", - "kernel_name": "rocke_uattn2d_tiled_d256_b16_t32_h16kv2_bf16_wpe2_mw32_regpv_4wgqa", - "sha256": "299f3d565d7e70059d7b2988646e8e79cf36e5a752add300965b2c1f7fd69156", + "kernel_name": "rocke_uattn2d_tiled_d256_b16_t32_g8_bf16_wpe2_mw32_regpv_4wgqa", + "sha256": "9bec03ceb6570b738edb3aa9b443a89e714288595de4fff5f619e865046f9032", "status": "ok" }, "attention_d256/gfx950/pad_interleave": { @@ -4782,74 +4782,74 @@ "cases": { "attention/gfx942/2d_bf16_d128_b16_decode_i64": { "arch": "gfx942", - "bytes": 453313, + "bytes": 454880, "family": "attention", - "kernel_name": "rocke_uattn2d_tiled_d128_b16_t64_h32kv8_bf16_wpe2_regpv", - "sha256": "253477c2582804478fbe3995ebaf936f68515fcab3427271f3bcf2aabefcd76f", + "kernel_name": "rocke_uattn2d_tiled_d128_b16_t64_g4_bf16_i64_wpe2_regpv", + "sha256": "9f63b6295ff616606d52b5cd17630bd65552f3078b3d03c49f3e76f7cabade91", "status": "ok" }, "attention/gfx942/2d_bf16_d128_b64": { "arch": "gfx942", - "bytes": 394089, + "bytes": 395596, "family": "attention", - "kernel_name": "rocke_uattn2d_tiled_d128_b64_h4kv2_bf16", - "sha256": "3f6e283d958c5a27fbe2417f78e4081cfa67961cee68ba59b76ffb5e4dc48359", + "kernel_name": "rocke_uattn2d_tiled_d128_b64_g2_bf16", + "sha256": "c0110a4660e096a1ba69ceb5e849f274e0074ee9f191fc7177635db17bba3195", "status": "ok" }, "attention/gfx942/2d_fp16_d128_b64": { "arch": "gfx942", - "bytes": 382419, + "bytes": 383926, "family": "attention", - "kernel_name": "rocke_uattn2d_tiled_d128_b64_h4kv2_fp16", - "sha256": "9981b85dd33ec8a0fa9468bb2baa6e62f17c4ad2ac93cf4cec4d4cbf7b5f56dc", + "kernel_name": "rocke_uattn2d_tiled_d128_b64_g2_fp16", + "sha256": "ae18425059b9ef1e58d60dd4fba07237f9d2ff09e6593aab4509c33c5bc77f74", "status": "ok" }, "attention/gfx942/3d_bf16_d128_b64": { "arch": "gfx942", - "bytes": 356140, + "bytes": 358845, "family": "attention", - "kernel_name": "rocke_uattn3d_tiled_gfx942_d128_b64_h4kv2_seg4_bf16", - "sha256": "4b4572ad4abf3b71171d6bab1411d5f635c51a78806f0f4e9dd20c391be98336", + "kernel_name": "rocke_uattn3d_tiled_gfx942_d128_b64_g2_seg4_bf16", + "sha256": "e6cc4efe309c5291bf751ebc86e8beb998c17a980fbdb8a115aa871ef624f176", "status": "ok" }, "attention/gfx942/3d_fp16_d128_b64": { "arch": "gfx942", - "bytes": 344092, + "bytes": 346797, "family": "attention", - "kernel_name": "rocke_uattn3d_tiled_gfx942_d128_b64_h4kv2_seg4_fp16", - "sha256": "c8668c7491fad1da8a7437d92c0236b0991bc9583db525ac0a7d97c1718f8da0", + "kernel_name": "rocke_uattn3d_tiled_gfx942_d128_b64_g2_seg4_fp16", + "sha256": "1bc69ccd17c9d37f88503cfc4cad479c7ad5ee76e8e861600cc04e79b00c8668", "status": "ok" }, "attention/gfx942/reduce_bf16_d128": { "arch": "gfx942", - "bytes": 9688, + "bytes": 9844, "family": "attention", - "kernel_name": "rocke_uattn_reduce_tiled_gfx942_d128_h4_seg4_bf16", - "sha256": "71c374c060cef19fa32166265a6699405767e9ab966d032ce0a1645f4ecf441f", + "kernel_name": "rocke_uattn_reduce_tiled_gfx942_d128_seg4_bf16", + "sha256": "ef2e7b04539ca47744a03071852d51c58bc0042bb41962b7c11507fd7afc84b5", "status": "ok" }, "attention/gfx942/reduce_bf16_d128_own_module": { "arch": "gfx942", - "bytes": 9688, + "bytes": 9844, "family": "attention", - "kernel_name": "rocke_uattn_reduce_tiled_gfx942_d128_h4_seg4_bf16", - "sha256": "71c374c060cef19fa32166265a6699405767e9ab966d032ce0a1645f4ecf441f", + "kernel_name": "rocke_uattn_reduce_tiled_gfx942_d128_seg4_bf16", + "sha256": "ef2e7b04539ca47744a03071852d51c58bc0042bb41962b7c11507fd7afc84b5", "status": "ok" }, "attention/gfx942/reduce_fp16_d128": { "arch": "gfx942", - "bytes": 9676, + "bytes": 9832, "family": "attention", - "kernel_name": "rocke_uattn_reduce_tiled_gfx942_d128_h4_seg4_fp16", - "sha256": "c2dd8745dd4a4b54012c6ec13283b1813074e4eb527dc0c2bd5aea07ba836064", + "kernel_name": "rocke_uattn_reduce_tiled_gfx942_d128_seg4_fp16", + "sha256": "ad96a8fb6a051e71143a4d5f1309fc437ee7c730d3d4e09fb76a43cdab5c8735", "status": "ok" }, "attention/gfx942/reduce_lse_bf16_d128": { "arch": "gfx942", - "bytes": 10568, + "bytes": 10739, "family": "attention", - "kernel_name": "rocke_uattn_reduce_tiled_gfx942_d128_h4_seg4_bf16_lse", - "sha256": "38547bbde96c3f7fea0656216d3cda7814e8076218c51d756bb1d2f4000dc296", + "kernel_name": "rocke_uattn_reduce_tiled_gfx942_d128_seg4_bf16_lse", + "sha256": "8a49e4762779789e90c7d1066323d19c6ba8c629315a853e162fd9b43658b903", "status": "ok" }, "attention/gfx950/2d_bf16_d128_b64": { @@ -5094,18 +5094,18 @@ }, "attention_d128_swa/gfx942/4warp_gqa_fold": { "arch": "gfx942", - "bytes": 408425, + "bytes": 408637, "family": "attention_d128_swa", - "kernel_name": "rocke_uattn2d_tiled_d128_b16_t32_h32kv8_bf16_sw4096_wpe2_mw32_4wgqa_fold", - "sha256": "0cd18465bf7c11e4b85f2b9d27be69822f5b5a90e2ac2c63bc9db13beed9ff68", + "kernel_name": "rocke_uattn2d_tiled_d128_b16_t32_g4_bf16_sw_wpe2_mw32_4wgqa_fold", + "sha256": "41a0c77c12071ad34f67e2349290b9838aba509d44e7bb9973dcdb67393f6a7e", "status": "ok" }, "attention_d256/gfx942/4warp_gqa": { "arch": "gfx942", - "bytes": 467451, + "bytes": 469393, "family": "attention_d256", - "kernel_name": "rocke_uattn2d_tiled_d256_b16_t32_h16kv2_bf16_wpe2_mw32_regpv_4wgqa", - "sha256": "de4462883801cd10aab5035e8288669abfa60bdabc8731fbd1d36940251d2b62", + "kernel_name": "rocke_uattn2d_tiled_d256_b16_t32_g8_bf16_wpe2_mw32_regpv_4wgqa", + "sha256": "d69ccad50b8c1734e8d5c86fdf55727966c353d606fa3d974d3129a78386b262", "status": "ok" }, "attention_d256/gfx950/pad_interleave": { 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..871c330058c2 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 @@ -17,19 +17,19 @@ "5": "643cb437a319d7be09ef59750c9379e6de4f7458f3f4a6d580bb237009245d78" }, "attention_strided_kv": { - "0": "3db5edf30be9ce97cb3a092984fe922a77809c5e762026249928f658dcfb8f7e", - "1": "6d9013a4914a2815d92778aaca640fe6bc747db8209117e781fb67951abc5f07", + "0": "c99fe0ffaf340903e7c913700011f4e4e6aea5630ddeca60c24a66916228d9cb", + "1": "7a4040765b46c08f298f74dd3952e8a830aa02d4e397d6d3d37fa38f7cc76e02", "10": "8ec3555ed2e64c59f0a32847175a1dfe0416d83036b645c29da331e32f640f57", "11": "4bd0dd8a84612f0e25fe82e1188eea5a23e7951179a5ed266d969a391e994382", - "12": "86accec45af1c586b7594f090cb4750c55a65039a014c4f508e808d9c4226cbb", + "12": "8b9fdf9d81f108bb733518efe3f16ff3d3a512b5a735c2038666b5bc98c0a51f", "13": "6a0c63a1574445581eb6767c40600b6bffedde472c9704bbc5511855db030775", - "2": "76d7feb82487213ba920e107950d09b999982f3695467c56df38d7e9737dfa28", + "2": "5b2da99305995476f23c4aa3e274c249fb3dbab830d984fc641a57644d0de8b0", "3": "1caedf0196d6b6cb889496cefa61d844373d4db06dc9b98327f6aa44de2f8b48", "4": "d93e5587c7f3b2432d8248a8639f19ab75416920cb6f2139198bbcddc0603ba8", "5": "f7576a337c2599cb03991ef9a8e82541a50208c9ef1704b4b3cb63d3933ac7c6", - "6": "c53c3dfe719ab0c2c978eed42b9846f35a947d71f91b2d91cea008f6cf65c90e", - "7": "4d157a7f700be997c772d8eace3d48a1bd98e5a6575297d5ce7547fb672699f6", - "8": "668c6985a3d815d88eeeac3ea798755f50efb34859ced595e81fe6c273fe479a", + "6": "382bd4a5175627bbd9090a22346b587ddc3d86cc2b461f6d98b20ea7b9953260", + "7": "136adbde306c73fbbde41916fb93ac4b596cb116b2730727ea0802826580941d", + "8": "e5f2a61ca1e3b775cbac546a6cc531c22a5c978edeae0b0c082b9b05d083b87c", "9": "b932c7ec5c473d26f9c5d3364d503eba33b0d916d517114842dff4352274241e" }, "attention_unified": { @@ -1205,90 +1205,90 @@ "99": "ed953c5c5fe0bcae664c039ec8edf2016474863875f151b5e946b9e37fab8044" }, "gfx942_attention_tiled_2d": { - "0": "081b9ea44d7e6267b7adb233fd662024b71813fc22a7bcea78c64391632ab62b", - "1": "8b8e0a10881a1d49dd9557950a753483247653b331e3396ad44bc4f895b9456f", - "10": "50e97de5e5307362f008bae1b20c27902c0d4b9676a548453d835a8414de09f6", - "2": "7db61c7206590861ad98da1bcf3bb68ae42fbcdb7e87a634886428d66829f194", - "3": "497515c5e732e45b80ace2250b23e1e23923783a94ad872bdc8b55a55cc10157", - "4": "9cdc054059a9e15de4960c6f8decdd05d4e97cbbff5a6faa39cba6c10d8d2f76", - "5": "613ffa1352f5d13aa8d05a6f328ff541cd9703ad820bced3f3ef9d695925fdf3", - "6": "17ddfa772168f2006ac4696f07f9dbc02ec923bf51b6d1e1857ce19222da980b", - "7": "3084d1d243906da59e8e6bb346d445202cd45a1242481d1d5b9615d19d545218", - "8": "46ed703c45984abf61ee32d3b8629c73f95c521f580916d51b1d72d07ed9962a", - "9": "b52b187abb3912f7c2e7fac389309131c6c9d096cd5b7c8555ad0ccf1e9310b4" + "0": "243ce6e29f52f3e70df569b64fded97f7e2fb2dcfb7007ca6fc153165189a8a7", + "1": "e280cb555c592f55ca84a49f9e52a1da3244921b650dd7fc5c1e92f75b4aab52", + "10": "5c5130eeb4d69e97904d59f1c2f2d020ac2ad555241ceaf46ec80f041c1795e9", + "2": "7e75a335f62737a2a9d3e09cd98009b43f0d7eb0e9c0645ba73b9e324fb02cb2", + "3": "d33d637adf4c7f373ffd524c2748059bbb3a084215c6792ad2d41a0f39d5f639", + "4": "4994f64be200f59f96f9920b497ebfad8c141ed560761a0f403a2fcff5c0c1f5", + "5": "4ccc00edbb25c9bd254b788bc6a011c16f799abe26f267edf340dd64c74ca0f4", + "6": "ec7a9a4ecba3edcb1115711d4211500435a6c3b0ae2bf8de06f21c1393e2d2cf", + "7": "e8044cdeef8fbc9c1cabc12459cc1e0ee2bb521068b74f9668ecb23f16963c71", + "8": "e6b9b2d9dddd312502148246e0191f0529c26f04838ec3d0e82387ef3d344cb9", + "9": "2f79638f5bbc5b125ca1ee3501c428e7062d0354ba1ed1bc6eea4c45cc077227" }, "gfx942_attention_tiled_2d_stress": { - "0": "b41bf01c57fec2c9794ac89a3b26c30703141690d02b7511c882b0a4a228fca8", - "1": "ec96cbf156ce3c965ab73b2e98fcde8d45744c8cfc9028f461ebcf87a5748ab0", - "10": "73a51d59a5fef0c22d11495861ad12f141053609029917d5ecd717e3f9989c48", - "11": "36b79784082e885218f1c39dcc1d32d1a26bec6b58963210e4324a5bc743dead", - "12": "c24576cb2b7a2d457d053a816acbcb90f98599a7b38535355d8966bbf1df8f40", - "13": "2ef78bf52bfcb4af52690ee098a5c65483f12335ee618f32d49544bacbc1df97", - "14": "4867eb089e47bb1ca82e74699ee4cc9792dc0f89c6cee8b0acb9feb22941312d", - "15": "df68954fca97d47df24abba2adddf895f7b6f7ec18cefeaab6351c66752ae19f", - "16": "e6c8e96df74c882b3b8e6c15d140e35cc9166919f54c5e7093c552b532cec0bc", - "17": "ccc967143bedb28e0d475e7e237e7e418931c1026c662451422ea27102c509cf", - "18": "4c92ddee2c3f037a7d3062568b2ab94db45111261242c5123628064954f3b31a", - "19": "f1a68d02f90afd4fa092d47b90a7684a7b6c8b62c7d92485c3ec357e54134062", - "2": "34e7a008563f08d6c8fda39df04b58f23e37e4106d4ad3a4dda5345fbe890ef3", - "20": "2b3cf9c06055f2f55912b61676b5854aa9420bdba1c428b833cba1e127e2c591", - "21": "f811d4aeac6ff137ccbdaa01a81bd274bad47f531be8d9fd2cce37c504882b53", - "22": "34313e35c904c45a90c8816ea690c585399cbd8772a950216d69bc70f5e0af9a", - "23": "6251340438f17f16187a90ae97fc72be16845302adc356d715be3694b68dbbe5", - "24": "5dfdcf3b6647d8c4bd5ba28f6dbe1e460a16b3ebf91337b94eebdecb2452fc4d", - "25": "75004a08bd66dce13adb41248f7c15509d3ec5969916d5c787ec16414b2ef16e", - "26": "2bec92682a8568ad864cc103aeaae5eb745e48abe5079602113c8bed92d83b9b", - "27": "16fdf0f44630260e52edbe72bc0d1865ac6654ccae2f213cb0f2d75e6616604a", - "28": "4cdfb23b44ee904d4c5f26f64012dd3081478a7da3c1c4d49e75255cfc6bef13", - "29": "534bd15e9124447929d0e11386aacdc2ca3f2d4bb2a20b195f613b26cce74ebb", - "3": "fb4aad5cb8f24a2970eac4e837255819cf244aa6c2913aaa2d8a9109bd836b8c", - "30": "ec23f4c9254f609e18d58faf2f8894c2ae20eb6814d9d8092892710fe81266bb", - "31": "3b49f047b693ddbe704b4157df2f549a76f98c06f602a08c67ae191a7b1bf6d0", - "32": "f820fb95d3933340ba898f10d57ac8c12d40dce3aaeee38b5b16fefa9ecdc081", - "33": "0d4ac93c55201f0bf9385e3a268bc1b73120731f2cd33a237980b627bb5cb8a7", - "34": "8956e5de8a34ccd05b97821da447d2602be7e7927b1347764ec1416e7ef1b86d", - "35": "46728283bd432cef6631ab740d687ea72beb3237d1c6d864d97f38ab33a74d56", - "36": "64dee702d38c54c7d75ef1d07c213aecc648250a80f89c51c754ae31f4e3bc5b", - "37": "2470bba20d2e031a5acc89d8256459b54f231db6ecd0e69b974deb1d48628057", - "38": "b8cb68bb19f2ed952d2f99a3f043a8a03159e2f8df7793b3f98c554ddb3bf659", - "39": "8d284c9a4d63ad8272c926960d989f317044c581e11079951a0d45e09ae7f54f", - "4": "99c8943c23e13e149fb9a06ba65e29b2c97b9ae585efaf279b17a7377d3fafce", - "40": "17f454f15d1cdc9fec20fc7e570e93cddbed6bc7be699cb0997172f72868047e", - "41": "f3be3c4547587792ea2b9d362788c9001e94b467a612fc587a658419c603b55d", - "42": "ef132b0d611a5edc4fe75616c0147ffa6356619d32c07659655ade8097625698", - "43": "bdf8007a6e7409068023b8429bd7cdab30e8634846e1714558df1fdaa9eabb2d", - "44": "b88caee445d3c024b305f8ef08fbe49bc0f493ba344c050e4876002f37bfcd56", - "45": "2171aae193cab984d6001169ceb0b537183aba3cd2caf4e1e80fc07dce9818ab", - "46": "66c0d88f8fac33d7fb6451fe661e95dbdcbd7c71290b18cc60c5825e3c599dfa", - "47": "00276f07299a8675c570edd55491b61ca765195c6efaabbeee79c0319790393c", - "48": "d7a0ec9af2e31136da3f1e4a4667fe5013cb945e4d25fa4166ffa91dbcd0e32b", - "49": "ada093e8de1e26966b2e35995b2c2f5fe50fd78240e53405314f69b4e05ee2db", - "5": "5a068da599e5fbc032ffa581c133f52fb0dedde59dcaa1142890339eea3f9728", - "50": "ae956a3c1f9cf86e71709abf4ee4cbc5f31aaddc64bc83faa29f880fd270056a", - "6": "c972d836f2fcba3fb6c9051771f824a3c55a9e08867be6b4bd582c81cb43300b", - "7": "8b8e0a10881a1d49dd9557950a753483247653b331e3396ad44bc4f895b9456f", - "8": "3b87174c637e80bf221d571ea409e7fd54766a3d0f8a9f14f7516fff6f70a28f", - "9": "9cdc054059a9e15de4960c6f8decdd05d4e97cbbff5a6faa39cba6c10d8d2f76" + "0": "83585a59e8e35d5ba68d78308ed780cf58d805074052d09ce60e6c909b51bcb0", + "1": "db0e5daee6974938674506ae08f8a3e99872d92a62bd99911099415d69864b9e", + "10": "6b105d8f69e05eff092a163adce5867833575e175ee3889cb1d691cd83a80776", + "11": "2bb482c2ceb3604e3894405f6168f762df2183ef04c66869e9d21d22fd16e7ff", + "12": "4d50ee754cdcc81dfe277c51ee36cc07a4730556fdace2bcb5733faab8a82714", + "13": "a49723d729e5a25f90771acec641484861449ee48832c78839bab59d10b726b9", + "14": "215d0f48367f37848a077633f77735eb6c11024ae082cda72184995b4baaf921", + "15": "f04faf55611e1c9d69a1474eee901c1ea9301c81ddc88f0cf3fbbbdfdcfc33e5", + "16": "83585a59e8e35d5ba68d78308ed780cf58d805074052d09ce60e6c909b51bcb0", + "17": "83585a59e8e35d5ba68d78308ed780cf58d805074052d09ce60e6c909b51bcb0", + "18": "83585a59e8e35d5ba68d78308ed780cf58d805074052d09ce60e6c909b51bcb0", + "19": "759e4bbac9704e19b6ccf57a5026941056a5bf5c8dd31d3288b21da647678d7c", + "2": "243ce6e29f52f3e70df569b64fded97f7e2fb2dcfb7007ca6fc153165189a8a7", + "20": "f24fc4f6e5362a0feea5995c6e3b623410c894cc7f06a640d7f8722ee675e616", + "21": "d66573654a4c21bc2d33a28b5bf05fe9284be600b352fa27c3c5a41e5b55b14c", + "22": "cef847c49af44460efac117f1cfd1f19c4ed65dc5bf875a54db44db7828c8f58", + "23": "05bd3b7385482bffd421e55d356d34769a79026576e1152d3e810679439c13d3", + "24": "16906cdc1d6a3202650406ef0f0ca5836baa9ed1ec47fcf4d52b1f08192fb650", + "25": "a8579d5f404d226939b43b2aaff3f5a2614fe83c5dd48816a66426f4b4ed4f9c", + "26": "1a94a2bf533212234d9d713ede2d9438af35624637bbe066a93089c3110a7a08", + "27": "81d94db75cf0c253ab91a90fbf8730fc6999948eb487b5d0de314b0d0af0665d", + "28": "9bbc00d01882ab9e8d7003710e3fd259777963f46a36a6abeee79a0fb22ee00d", + "29": "a35819adb3ec8938b83aeb6db868721ec4b3c58d4c1c997a39aa39e3f1d1e302", + "3": "2893dbb65c3c6900ae2d803f9403e65fa78617c14ae02f8bc5fd43782b4c9429", + "30": "8622b96f0a2cc43881e8cb41c7ea4406e896c2a7e9d2f37bac26ed102fac8a73", + "31": "e17ab876e877af58abb87af63e71c66e1a4b383ccb2bebfdc3b5b53a764b1383", + "32": "83585a59e8e35d5ba68d78308ed780cf58d805074052d09ce60e6c909b51bcb0", + "33": "83585a59e8e35d5ba68d78308ed780cf58d805074052d09ce60e6c909b51bcb0", + "34": "83585a59e8e35d5ba68d78308ed780cf58d805074052d09ce60e6c909b51bcb0", + "35": "bedb63d5d6916aabf9389cdb5411027b03c4bf19e8a6bd65e09f8e508cb9d765", + "36": "77d5f19eec4203f30f8e114bb1bfe6048316df2ddca149068a03070f855712bf", + "37": "4d6d3cc972736804cafe2b64c5bf7271934f7d39d881edfa329db6a54befabc7", + "38": "e1d5afe9b3c75ca945f0e4abc76c65c1c4f2bdaaabb946bfc049d01860b27a6d", + "39": "d59c01539b723c778343764baff02e4c61edd81b04e2c37342b8cfb1658b8856", + "4": "b9093604e1ab3e7f680a90e420f4c02db3abbb19a27d20b0e8cffa8c5a30f844", + "40": "2a4ba52d4ae52f681d98ab4dd28c3824f893b17ddb66d49b07cc60f2ec6bdbc4", + "41": "da7f4787a0130c06f0c90bcdff9ed74a6ed7b68c6499a06a05919d473e9f9cf0", + "42": "fbf848a7d1dda3d1515fb953f898f77ac9b21b501181c285a6f572760fd77f3a", + "43": "f9c47b85d9e602ac2bb983fb9fa0b56b67c7fe193fdbfe7c72c05bc379dc6c74", + "44": "39671ec26a04bb905a45a9d2e1c93740095858c334646d0c9ed7fd8e2e0c3ebf", + "45": "5c7653515fef66d77bacf896727d7dbc82b95294eb612450f470c9708fc0e226", + "46": "838af06e86e9ab7241bfe7a231be81b80989e24c4146e943b184d32c7081b849", + "47": "69af5a7bf696af29dbfa708f06a32a975cc9aeec00a65b10afdf0b0cf28cb08c", + "48": "76c9cf2877b0cf92d2f9b8de9db917a98061687ab017e829580055501091fd6c", + "49": "d79fa2e4bcba2211f5f6295db55002a444d14a73ece4d4c25ad5dca508279398", + "5": "09078a9997dfb33458aaeb03fc0b94dcaa414105c932f39617ee0afd4140f48b", + "50": "9b050c0c07101a6926a02bc7995117fc4a3604ecf8d26e9068f7752c4979a1d4", + "6": "0edd8ea3ceff93475f6c3f1fb2e98128e834722054627e706d98b4e5d0a44039", + "7": "e280cb555c592f55ca84a49f9e52a1da3244921b650dd7fc5c1e92f75b4aab52", + "8": "41bd6c2b38735683337c3e14b40f1bf484ff10651ef1a789b04609cd96801533", + "9": "4994f64be200f59f96f9920b497ebfad8c141ed560761a0f403a2fcff5c0c1f5" }, "gfx942_attention_tiled_3d": { - "0": "82c3c72998cd823e2637f88b0a1c72f348eb967f0614731fb61c69957d00c4dd", - "1": "7e58c5e2e6bc2abf17c4b8c7beeb2493f06f659921daca2b03ad7ec7cc789b70", - "10": "e68e591dddb5a3597a3a94a8dd5198c7a26656ae7ec0f7b85930a5d290932709", - "11": "db36573107c98e859fb58951a102dfe2637e9cb91780768ef987e313d13f6fca", - "12": "16e0eac8fe2dc90647b16810453eaec2448c15d5d35d87d56eea7958b4300a95", - "13": "497672e27a6f1151ea11f56b69d770b8f7fad4d3ad497c30c2a9d4bb36efa5c4", - "14": "7017aad68c2f9bfbb7bc8b4155486d3beb26e708a69008ce85529b1d302b7c85", - "15": "74af62c3d1e130aa3be256ccd052038cf4979b0bfaebacb9a68203f92ba1d3aa", - "16": "13a61b43fa518268e0a66960b23eab2419f85915dcf852cdb4bace1bccf16e81", - "17": "77d0e1ab4cabb014c931b54c7bdc0aedadf53cb99174979686b4018e056a5a9a", - "2": "82cdca8f08ee9612cf5ea4f3ac729e826d283615ce0038eda2ce983c87bafd93", - "3": "0bc5cf7cbe950277ca54549cb078916678e75b4db7a3880e22328478d89ffb3d", - "4": "168b4be8598882c4beb38e84029c71293ab2f86538c4f8290c8c25b76a9634ef", - "5": "a7494ff8169f6e0c7788dfd9321e6878b87226448c14d838d1fd7a7135d05b6c", - "6": "333e7c93b7dbbb711e73bd6fd503e0c7d6bda0615221aaf2bf863ffc41b6765d", - "7": "579555ecebbee98e80d4b346cb21be3107eef939c2bd9edac5d934094f8605c7", - "8": "edcdbfb38f33f419b6cb5e35f30de6b6ba5cb0c15a163330966781f66784c342", - "9": "8101c7633b59e2202adbb913a3f9544eed1fe0003f67aa32de54f42642137c5d" + "0": "aa350c023bc7703af40f6f8b7ad7c5357b52be804fd2da3cb799809a5f528f71", + "1": "c2dd79a1b240714016fc2c91f953d7084e5ae3d6669780d0adbc52cb95578b30", + "10": "f90b66ee3e51ea2b98c85f439efe8a8ee907480d3f7a6743333a9bf8bc2d4f49", + "11": "c4c2bb8b050f8f793e5528d219fdc38d4d12f6da150499a84b304b472cf90ba3", + "12": "a1de2dcfce3a7b5b1e7ee71e17a71374cc6d83eb00c558af4b039e2fd193aeee", + "13": "72e27731f1fe5c334069b1a9e78784f7836c4e2298ad3b38b785de6033c01317", + "14": "d84c36ee99f9137188264361427cfdf85f5b602a2ad7914399ce4e19aff04a2b", + "15": "9d24d520c674c8e127276e0ef1dd37343beb6ba5534d3759d47b3a1b0a3151f9", + "16": "ef3452234835a8386dae0a4ca8d60ab160f0489d6d2748ba950990ebaf6d02b6", + "17": "236863bc3a96bb609c1c441894e8f208e1318ff5ee231d30929e9b611008c89d", + "2": "0aafc6fb66c1c4b9349a4ecc265ae0d70cef0697176488f9e365f3ad9ad72a1d", + "3": "25b66dc34eaf1122e985b4cdbe27e6ca17640a7714e7a75f1cd85264a31fa56c", + "4": "0b84bb6a3a6de5450589c36cc549ff01d325b5cee13e65b96498418a7cafc057", + "5": "0efd5e6e4f1590492f54ccdca10017498c64d0a7afccd573fda12cff8e1b038f", + "6": "3712baf35493e3d35d23b0510604b79599841682213c6b0f820e77ad0132032c", + "7": "5a489ec10027f3e4f4895b71c603be723e09cf753c70013370297cab462a2217", + "8": "1bddc9a616459227476f58036f7c33b38d2cda247776a193e37dc1ddcc20b62b", + "9": "d1dc34cb8dfca4f7fadafcd65d226b540ab47e634f24931abd0a802f02042dd5" }, "gfx950_attention_tiled_2d": { "0": "ca49dc00b5be88b8de75c90a1cc134b377df4c4e0af3d2bb2a7c33639cce51ae", @@ -1750,19 +1750,19 @@ "5": "9197a3464b35817737bfcf6bbcfd45d3dbbd6b3b441271df691b969126adeff6" }, "attention_strided_kv": { - "0": "1626048eeaface4bf6eb73c5bc0c81a653da20fc18ca502898a55e6111d13fb2", - "1": "a9d06552869ca1485cb3a7b8a362d2a448fb8d55220f222d50d4a54ded0a2488", + "0": "8cfb54512441a845a2ae325065291e39e5367be96f320596e8622d7e308ec610", + "1": "b1cee01d468e92ecb5b8571323ffd897bff01c7d2e42dc8ba50814387f908a63", "10": "05d6f2c670f9c66f59974f7f2715eed844799ac869e2765d3d1f691bc417d2da", "11": "68530d1c8c9c9a12300743809054235d88e2a1698f811bbd86a531eb47113d40", - "12": "fa6ae13282f4e5748eae40227d1355529e793d2306f210136a498b6621fcff96", + "12": "6adc12acb6dd75f05927ef1e4f6683b85c780abae25d76588bf57ae25f8dfd88", "13": "d7982e56673824863076878c3461209f8557cfa8507b102071eae22d712b6414", - "2": "8c0cb312ca581170ce8e387a94a1b77ee52b4c7fc233aa5b15f65f0c707eea8b", + "2": "a80d03691e079aacfb8653f7054a836d2ecd57256285274fa51c52a8a398c727", "3": "272f3227c873b525dc8e48c4f024f7885de375178e20f4c4d6686b4724d82c0b", "4": "b25f6bfc84f5db8903883a85671f13b109256697e33c67877686b4d90ac07606", "5": "bde0dfb3994fb2a89081fa05809323059111a5e7d4df55ed8859ec2f6b0cd930", - "6": "db68f7a59b6a018e310c3d86d065432fbc4a7201f3d914114dec2d58d3200a19", - "7": "e15ef2766af1f7f39f1e2e1e6f1602e17ed832ed960fee0fc87c54447742c45f", - "8": "15c5cc154e40e4e5f099d1bf6e476e283a70c450824db3ebb9df927b2f7aa2e0", + "6": "23c2881b225837148b997823910e598ad844d781441a1f122a5d370b4b97900c", + "7": "966f73aef9172bbe60348f970553576d609299131e029e6842e1f4dc49c96da1", + "8": "c5f4baa4b7f0237ce60bf63f6694fcca24bdd440c0fbc34a1dc0a6c8947b791c", "9": "7866c9eb6d414e0edf3678a477c387e2742057eed6a96616544f29057d609d0a" }, "attention_unified": { @@ -2851,90 +2851,90 @@ "5": "1e92940d1bd0df8fcc31d5e3ee7e763c1e6ee96f525d36dc9f3b9d28e4e1f2ac" }, "gfx942_attention_tiled_2d": { - "0": "a956b120d39272206978f3c417c935604cf613f080028a6bbbc5102f63fe8542", - "1": "868f36aeb9d5d997f371f31ea71e1a4ec63d4be09e95ca78260e3b5ea71a1708", - "10": "4985be273f5a50ba0518a3cf3a2aaf91a9477db0daa1d9143a79027166e415c7", - "2": "d4aa5087d93f60363bc7c0f6f9168e3c159f512afdab94d65505b17a490168d3", - "3": "1f28c9626f50ed87e159c02bb68d064c182496eaf514da5302788366dd818e4f", - "4": "901e017f616d1db39526eaf01ac905f9159a01a0843880862266880477d99071", - "5": "2b3ab76b5b37b656a848dd276eaa4e7499839707b4ae7fca0084d47381ce0341", - "6": "dbf46f3da55e2812bc3a20620c75919fd3c3257fef4aff88de8096b2e833d323", - "7": "2ca52b461a729bfaad2594dc8e6e5ad609c562aa11898f3de0a889d400fa0d7b", - "8": "2ae6b4a12ad4c0606d21fc510d40c3c40dc5f03a634411b59f8226e0feb2667c", - "9": "a81abb86b218dc327098774f8ef347f6ddee0a1b8636d4e394086fcc70f8182d" + "0": "658c0b07bfcccd49361d6891a8e50c51918fec590facde450c3d242dffa2ac69", + "1": "4c0a8ad2f2514ba61ca9e11be73f5025ab155b8f90c9376bc2c71b46a6f31ef7", + "10": "64dba803a7774a42e66278ff928f4e5fec9e0947ae0e7afbb2e065185fd1f4cf", + "2": "f42d95dc8002f6b6996b380c1095d9ecb7499329c379df17060f6ad761b50a9e", + "3": "d376de4ead91f6716c2489c85c57da42d89a16a8a9a15a5865f743285414ee16", + "4": "a4802916d6cc5dc4e705809b41ed9353082ee89e8c426b87befcabe04b031747", + "5": "50a7a2b08021facf58289e2cd9ac8905c1d242a4280c18004541160c67d9beb6", + "6": "36304b3bf020f6aa564ff3450f49d9c6a959ebe5eb18c7561b102fe921d01002", + "7": "e4df972977ebf80ab8a8e70a5aa2354fe0aa1235fd11128c1a144e76dbce8012", + "8": "04758a1b520bed624ae1a7777b1030556717761dcc5f787fb47f9ba1cde8c48d", + "9": "1f60a7dee7cbecdac0460b3eadb176c19733bd5a889a3753d595ec2bf37657f8" }, "gfx942_attention_tiled_2d_stress": { - "0": "0dfbeed6b1c98049d80a15bff63b27af0e7786ad93497aa180e2ad8bf9f09d47", - "1": "1dc828e18e02ed1b0398beec1c5b13860deb752ae6d044e285cc63ea7ec53b99", - "10": "dee6e805f3e8dbfa58cce4c8f3d27fd5f4daedb27cb92ba44448d16d701e7a1b", - "11": "e7f0028ee284931eb76f99aab3dbabc90aebd8d16d6c0db6e81829a418167f55", - "12": "3306a5b773b2100c95aa21d4bcb00d7141f12e7d9e4101a7be2a0d9e68355f7b", - "13": "52d34848d6ff71c84f6a17f26e72a361b146bd8083475a19d5b54797bdcbcb40", - "14": "3eae9322bef8247146d503605b0fb16604a40cb54d1a64cf98fd5c830a518254", - "15": "b2985e2817c50ac19af35a9e0a159ad3f9e4add74dee300eff984ae44fccdcf9", - "16": "97e6f19c78e02f6922b3666ab73901df1448c87df7fcb8c26a93c9729dd32538", - "17": "a78421230eb3069c5dcc3309429d79437c5df04082b617eea1e36f0633e7f0dd", - "18": "4ea948e4e3b4b903df1d4188d9b858aa6c52f2753d794bc58f57c56ef418dea4", - "19": "35794421435ba2148089109db22f5e869d9f7fc58aa99e6f31370fa166781d14", - "2": "32684e75876db002da89b7ce96e9477da26b3e0520ea81ac6683915a651da60d", - "20": "ab0b38aaa5ec96be2102c8f770437ccd98941889fd192a06332b36d057f60c00", - "21": "7a45cd29baa2e79015ae5bfa5318d15972997b0a79f969b28d4b764e87a63a10", - "22": "37b84efaec1a0e4e4c108afb30e986f521f86982243d10df17fa4028d58a93e2", - "23": "cf6702dd43174921dcb4460e9bc9eaf7628ddf58c256c5ddbd1e470e67863500", - "24": "01d8cccd08110ace4106a09201029b05b4751666d737511748d9b84ec03e91ab", - "25": "9d876d495c3e609ceb65b1b7fabce46c75cf63032c5ea1387e6883162c9d5918", - "26": "406d515cc68e55fbf9fbf74d936b43ef4a99b2e9b8c1fa720c4fa002a07059d8", - "27": "e4629709b2525ed2c9bd9e892eda60a65b397af5a470bfaf64ffc47b9ad0530c", - "28": "33b4dc3ebfd5e09e1d584d8ab30dc5f22642cab7ebcd606d3c4b6141dcab3f86", - "29": "42b851d0ede850db263a251cb851ce14eb32b26f8a0a63824460f293d5de2267", - "3": "f77a5b0c6e63bf9bd5a05dceb927fc8225f8604b1c917103f5117829efd130d9", - "30": "e7f8a147de31dab5dc2984b09fb6acb39527d55d0906f98b3830637cb3b20799", - "31": "4c5492305b5ce73b0bd7f43e114b00408f9b8461aa223d32b8324d46f31be0a6", - "32": "8caff259b398b3968e600616200387b4b5152ac4f7d09cb36c1387ef456ca4de", - "33": "a8fb4d06f79a4da815602cc8f428832234662e96e39d1aa2dad27adbee895382", - "34": "c02c45d7d053c6aba72ab78c677c7ad3cd2ab2fce9309155c071a709aea5b9db", - "35": "c0620cecad33eebdbdb592d8aa99ed1d9ffc2fdef57048e52c6a9f2f8e72a467", - "36": "8fa57da01daec0f76608174d405dcc74b1386aebc34e21423ea78d5d806b87fe", - "37": "69ef9a0dd80ace723f26105d8f91c149986b4264ff5a5c8276487c7347580486", - "38": "76038ced0278d4f4c72a96b325ba787e68ff2f9f386e8852f72ec2fb2dc90bce", - "39": "a71ac88c716fd289e0fd94b934bd01d54a5438e432dcb1a6a57d6faa55b10b1a", - "4": "181c8431daf57989840c118cdc7478bba2ea5f8693c585943e2bc4dd47f92fff", - "40": "1cdf18a2dcaa80f0ebc9e82fe80ce8e2649312483ee4c15488ee2534b9607790", - "41": "17b4254a1df1d379de245628032047b97b72481f9fe8910ee705bfb254aad119", - "42": "36c875af338a84f52d9f963ada96bd9197bc3b0e6c593db3efadd35307233213", - "43": "97ebf3af6ed9d9ae08f8cd4c7b5cfb2ad66451cb400c1ce031cc6bbc05680c22", - "44": "33cee6c91d48b4c9a9a649d3d1218979b7a2ff9bc6f3b414ec5fe321ca32896e", - "45": "7e9a0805ce307140f1fe3b2debc40e2f325beadd36eb5bd2b74e6b0c54467404", - "46": "33803554e820c81307c5bf8141cfcd4f7ed56de00e4feb0273ed1a2917378d1d", - "47": "a2b9ec545a7ebfe784be26bc9ff739723795b6b08322c8a55d3b66b61f5a7140", - "48": "db204a860e3990035a16928271212d191592fbc7581213b6e10de1fb2ebbad85", - "49": "262acd90c463dc85748c7f31d9ad2f76e6e1032e58c8bccefd20943c58eef538", - "5": "84e290226d1a6e61302664dd88475166a02b3679f5662ab1544802143ec8e7d9", - "50": "9df5103a18b0c1f5576c7d329fd8829ccaeffb3babc8d4bb6a984de018cf6866", - "6": "cf19945ddad3295637e50131047886ef8170e490208ad22f56bc943c47f9de52", - "7": "868f36aeb9d5d997f371f31ea71e1a4ec63d4be09e95ca78260e3b5ea71a1708", - "8": "e038d9c94900b29ba37a27a7d12a7f7e73882d63e1d1643a93486c4bd2bcf55a", - "9": "901e017f616d1db39526eaf01ac905f9159a01a0843880862266880477d99071" + "0": "ae7223b5a1c268f580248d875fde5d2b28b14626d25d619866ed61567f807c1c", + "1": "d2817dfa0dc44caa3260f0e019c0a02272c929328eb0647b64ff40b8f0a1da69", + "10": "5ad702da97fb5bd46a9347f9384afc4c73f1567276b4962df722de33ec3ac4e5", + "11": "598091decc07968ba9e25936e55d08bae5cb2311991a93c26e42ce69c8dfd5f1", + "12": "bc7512e7f0a3290f3b8301d50632e199e9ad90f6a72fb0744d035958f0c36a9f", + "13": "307f51d0f63492cf879fe9c5c84fec7c4fe7c68517aab4cf17123063accdf6e8", + "14": "ef61b035ab04324af128322535ac63943e4e9b0ce36d74306bfba949972edc46", + "15": "205d6f62c8f99147ffca3830ddc0d91b87bf78dd74cfceae867294701eb81a07", + "16": "ae7223b5a1c268f580248d875fde5d2b28b14626d25d619866ed61567f807c1c", + "17": "ae7223b5a1c268f580248d875fde5d2b28b14626d25d619866ed61567f807c1c", + "18": "ae7223b5a1c268f580248d875fde5d2b28b14626d25d619866ed61567f807c1c", + "19": "56e038bfdf91c9d94b9a355caf6e3d015dd62aaf66ce8197f677822c500a62e8", + "2": "658c0b07bfcccd49361d6891a8e50c51918fec590facde450c3d242dffa2ac69", + "20": "f09211ba13a58dfdfbf23f7c2eea4fb61a81c699f478a506173a14cc2cf778cf", + "21": "67e60634074e4bec021035fb3febee4b72f99c6643a50957d6363a1303a3d935", + "22": "2c5db7e41484c0ca84481b45d954b5de482c4305c28b8dd9113df697bf034827", + "23": "1bbeda22a9f03d1df57e733b5796fc80956560c5806be1a6f5b635af22125fec", + "24": "d11361dc7481249531a87cfa32cf66610df6ea56ad07849bf0a3d0c2289079eb", + "25": "886fcf55e47ef728a207db301139f5de5771721de53f44af2d1a78c0a6cecd94", + "26": "0db3d7acbf2f18c50729f2db54a16c50f004717bd4fcd7322f32944ed1791144", + "27": "1a7f13475cb8fd7db2436d037e64173873af87a28382ab9b12bd77b8ea325f5a", + "28": "efabc82e6dc1c610cf28673837ac72e8fb7233fa184a10e5e2b19b2f6c59b93e", + "29": "985f67d70128984ee8bbe59b2563397ab06569bbf6532ae0de20df2ff09622d5", + "3": "02fceea03690a12aa552fe608bf9bb836335fb40ddd1a9be719ba94773cc16f1", + "30": "2d4111b0b8e3352d00803ff793258f3258fe323266fa79db827928e571ff6a2f", + "31": "34ca88cf9e6deb205c42e860157d94148afbaa3a7f20a74aa0a8033870e1cf31", + "32": "ae7223b5a1c268f580248d875fde5d2b28b14626d25d619866ed61567f807c1c", + "33": "ae7223b5a1c268f580248d875fde5d2b28b14626d25d619866ed61567f807c1c", + "34": "ae7223b5a1c268f580248d875fde5d2b28b14626d25d619866ed61567f807c1c", + "35": "3e98338ca8d5b78a75cab39ec08c359cd215cb991b85c6fc1e43e5e5c9cabe6e", + "36": "1ee885d5d9aa146de5c9c245ca3d59566bbc23a11a7764adacece79a399c5f41", + "37": "e501cdfde807f0d3b06d5366ff1f44f610b13b46af48f98f4d520f23e423be63", + "38": "34c753aef239e05cf276aa81024c9c087724780f90cd0e54c3830510c7433310", + "39": "3576716291bb30bd5e0e68d15bf8d6df38eb65e8aae9c6883790a16173a74e11", + "4": "9130fd505353a9462508605212f6cf66d567c1b6ad4b7b32900f530851ae0166", + "40": "34392c8627b2627dffceb4574da58173470ba93fe3d0c8f4c3920d162a42cc37", + "41": "cf7d77264e94cb9db46c0c13855a70583cbd76f27a6fa220e3c355f67a4760a2", + "42": "f4534436a22c4879fdbf67550bc5f4cab9718fd4fdd998064ffc313683d47db3", + "43": "6d2f154b5c6328387f74ab2a7569681a7e1cd3aeff4a405c454c1c51c1f251cd", + "44": "63ab0924746d295ed004dfa1118b2c1bc535a7e5d468a0382b13af6fb4a111ce", + "45": "5bbc602de7ef123178c310ef937f268bdf30e1982e3f8de9cb4c23a9a1663982", + "46": "766323fd9b47accbcf084711126070c007185dd386054dc5a64d1c93c049222f", + "47": "7b52f054af405d1d37a0c8bf6f432bdc9b6481af8185d8bd48047bf9f3bd04d9", + "48": "a6f519406dc951e30bbfab3de2a53a269a2f73ff78113252babb350c238ceb2d", + "49": "86fa712ec430a5545e3c219826b3e1cb9f935066edfb84187effeb46c916ee4a", + "5": "91e2617c2ccbb65d8f11a6e439363fd708b2c8fc672b32a91370411890872eee", + "50": "54a53afd5e086c5f173cf8f8e47e332a733be329108363c0c0018c361af4ea19", + "6": "480fa44a8cd71f0d019c9c1b948f54d1cd7e18e33d79de323cafa3009ba9bb75", + "7": "4c0a8ad2f2514ba61ca9e11be73f5025ab155b8f90c9376bc2c71b46a6f31ef7", + "8": "cdbd49fa53c1c8265ad35229a657d69e1d12d7a002f7c0a9b4b57cbf760b010f", + "9": "a4802916d6cc5dc4e705809b41ed9353082ee89e8c426b87befcabe04b031747" }, "gfx942_attention_tiled_3d": { - "0": "f61a3620614b51a54fccba438176c0e362dc492e3c3b6ce406d6c06dcf02083e", - "1": "6d7582d1419b89d70494ee63e2e05bde6c60f299f73440179d60e4afbf5ae8e4", - "10": "ba19d4dd3407ad9d67333362a27140567a70468654dc174667ed52fcabd35bb7", - "11": "0df4aa34224377bf944bb3062b546d5eee768e775e5e8f1f07a94ad4f1c8d4c6", - "12": "3ff0d802cd6023feac84bb1a07bc48a8689cbea24046138a956b4e66d5988ef4", - "13": "4db7db604d1a8f19325aaaecd11cc674aa0c2fb6d035caa33467e50ad5156c73", - "14": "02aea5df8e61523f8acb97d7e72aea32a2ac847544b55dfdffef362f21461ef0", - "15": "e9d7c0a309f76242a4ec4d1199cf8d05c7a7fc80d88a398e68808fd158a8976b", - "16": "d44a86588284102e502dd7f329604faae765fabd32c134fca6a83a37681031dc", - "17": "3f7ceb0edd63bb22f911e8f14279a8b5be14b1c4f951a9ad42afa7c9f7c38acf", - "2": "d776e66b764f75793fc25b17e2639b306c289592fff73953fcede2e76b189322", - "3": "41f971cb9da04b95343d06bd7c15143ff7cf3e82eb9baee619b563bc672d3838", - "4": "5eff7adc1f3424caab3164b362e805f6561e65955071ba010deb46f6a8aa33a2", - "5": "26bb8e65bf6a98aa64a294fc00756b12910f3f2a2c4a3c7e12a7a05d5f0c0244", - "6": "806338a1f5ecaaef9aa176e6a8ba08cae73fd17ab9b65ca0cf742df738f73bf7", - "7": "5d12ced5cd4697a9a82c85407b2becb0c54e0385e0968b6de06924f4624d5a8c", - "8": "ce51f9cb5d66e24e641647e9fdc7490bff294a358b9258ff2a86b9512c710740", - "9": "124c151ef4a8be0963321ca6f6e49e4efde00256aba1414da772d5ac3d074690" + "0": "d94d49bbb05d2df758da221ad23ee5f680233b40bf7408903d8e3625fa955f39", + "1": "09752826b00ca337e7315807d33fd8aa0dc51b6ded27bd4a472f77a972b6864d", + "10": "1a15b979a525eea328760b7cbb2a346a7d9060bce4a0cd8f2f6058e6674e76fd", + "11": "ca406f9c288b7f671ef4ade88db8aee72c7f46cba9594ae271f83b60c5874c46", + "12": "1db13e53d9e4d8bc04f935b4705b11646c10b39ffd833b7e8e7b9c7ecd298b4f", + "13": "c46889d1715ab6a2f89f632b904aba71f970b8d638616cac468ab797ca595d2c", + "14": "d0068ca273a3714dbc9920f8aa11e131e647617eb775c48720502eaa3cf57406", + "15": "3ac273b458497f3df8ef5fc73202e88316cc92fa74c680e5ffae51ab7c312718", + "16": "70c3cf16e608cff4c036354a7c5a88fe88a12376938d7399bababc587078d3fa", + "17": "a478df9215919e6ee5f5bf7ab2efdafeb77becd737bce7a78493baeec0b3ef9b", + "2": "9976606b42a736348651dd1f709a3b74a2d1306d5c1cf4a89921f85023b9bc13", + "3": "e977cee9184c8bd95675b33a9b34e7b11218471ced98a709f711f2821deaabf7", + "4": "0cd0c07d6a2c87caa5bba1c2720ea9e9a703299f077a18668a9e971a1fc8dda8", + "5": "33804a39dcf82487f3ed9b1d166c47be6ac4f9e768184a14b26d1fb089185440", + "6": "5973c29007a027cf1bfb84d6d6c2e4b94b03d7e3dec3bde6275e7ec71947bc6a", + "7": "18b7f6c6136b8cf4dd23bdef33013ae10aba43beab09282c5b03a5a3f756cef9", + "8": "a840ade391b97e553602f3cd32dabd909de586c92b2a9f44dd51db7975e74a01", + "9": "f9794098f04a1cfad57dad0a169e28e5915f03806cfe4f573c50602f04132a6b" }, "gfx950_attention_tiled_2d": { "0": "f8d4b006fc2a64a141ad3099defa8f2a1ca8bbe6307700292dc4c72a951acee1", @@ -3395,19 +3395,19 @@ "5": "c1623e50052e670b53bc35ceaa5c2e1375c0a49a40c351cee09f67f2f72e6471" }, "attention_strided_kv": { - "0": "bcef03e3b5555c8bbaaa03188901c1dfc8bee0cdfa80da32026367e82e7c727b", - "1": "9becf0d4f1b8eedc388e818afb01a40607a27efe2d415772ba1f6352c2ba35d7", + "0": "14ccc0e86cb9a96163a8c32812cd718dc454f726dacced64fa3cf07db5c70c90", + "1": "62086a1cb544ba32772dc112e252813ffff1961bc2db23401bf0003bce718b73", "10": "35992a4d15f09b54bd668d4b21556a2c871fa8c6f8aa39a89169736e303476ea", "11": "ec52f6ee37cff049c8b81eb579b71c32be965bb4730fc7b30bd477196499727f", - "12": "002dfd5d5ce7d67b6992178da77d629eb26e80c5343c95a16a0e2f89462cf752", + "12": "b543ef77d6fd6bcdeab9b00f9c7dd8d9823777987f2e9e09b79195f3d230be37", "13": "f71e2da33c81aacfd3a6690c0fa2c540aa3a83586e0aa8ed0f1ed638023677db", - "2": "7fc95348ccfd43e6afef1b5bc83b0e17b55c0879d5d8c752d40e6f631d0b9c77", + "2": "f2eb0fd45eb3693671f6d11bdce0eac92295e784e74efffe21351150e2807039", "3": "c88f5e1109c70e0c4f132b05578c6abf440d04bccb52bd774b492d67d3acc9bb", "4": "3e5522113635114b994bb40cf19bc51a417765ac827cbfa879acb45cb0e62796", "5": "3c03d60f49863bc856426e1431cb98e8e682da4b6818f1df320536fab12b1314", - "6": "2437aa43c1700bf2af3731b08d30869d68cf9110b2fbed6114fae27b78a4e0f0", - "7": "b1503cc02eed262c4d08b02d3802a16f3ee74ff0700f756e79e9a1c17580ff71", - "8": "8942a331f8dd8285230d4a53c02bca59ee5158284edaf12d7cbbd50900a5bfe5", + "6": "b4720021fe24f67cfdd7b6854a25f4b371cd5c086d7d804e4a1a451d8f885c82", + "7": "281cf49847fec9bf2f0133d0b0f83c9ff72e99e3f9bfd846d2d056289604609b", + "8": "8e1697deca33346efb263397db405b14b835a617e0fb8a9d6e5dd19958d767e4", "9": "43cbf9fc03696e3fcf0a132e31a769d92cd48267eaea3526e2cc7c46c35225da" }, "attention_unified": { @@ -4496,90 +4496,90 @@ "5": "0ca6d5f7108bf00a9acbccd47bafbe6a067bbf7cb7d4c64438010f406175cd7f" }, "gfx942_attention_tiled_2d": { - "0": "ce207c4ae527e57be870cc96a0e54ad218477af39635338648f223c469594b16", - "1": "213ee124983091d878f7caa81893be93d406e92f6d1938e5b96e123201226c54", - "10": "7f4a92cdf39b50015aca158467798c3c982d32b144bd7e5a4fb323534165dd9f", - "2": "1f7c7636818c44b41e7c62064fd9bd3149f337a79aa0612f6073399021f42f80", - "3": "13d57ca5fe2d62d8b7bcfd92315edce782db6081ecb8fd4a6dad65eaa962d488", - "4": "b1ac91fddeda3792f2e46f4f37a0956f9e145e04c230b9167dfd1bd95e1e57e2", - "5": "9c6ccd210d6e9c80b6964d68f6d606873c6dac7e7ef7cc94d4d41f9e9f85acf6", - "6": "82e3a89767957d40237d2b71a8cdc8dafa9be11c3d1918161a8e56ebb05f76f4", - "7": "5a6d4b54fb70d62786ec4c7fa3c6ce0da2db3f907aa20103244378dd5a413318", - "8": "3837bb7103b23810d73be4191f11649059b4412de17e6599b8e2bc8661e226e5", - "9": "402e33ef0ff9269296148c2b029f44662466d6d4c162ae30d56fb9f674ad5b97" + "0": "98fe2fab85d720915e27214e187367f528849aa6f86c1a8d4c26ce7929963cc1", + "1": "6c7de2b5477f939d037a53a07db441ad7dc633f61a7e35c6e6f1951e348901f4", + "10": "ad68db14b411e90f6af00b0927823071ae86720ed7560ba33778463a08724c7b", + "2": "8b62d0f96b9954d475090a432ee6ea0705b35bfc9e1eff1f0d5764e9b16bc7f5", + "3": "827d0772f480a6aa6f985a6b75bb22c082172660bb005cb920f36aeb4811509b", + "4": "0bee8cedf0247f1d06728b30ccc5d03b5b2bd23f4049a14a123d11d2d67708f7", + "5": "e08bd214aec13735cf2e8823cd2844e177b2109762b6f87d7ce3574494476a03", + "6": "e9e71eb2ce686ad554cef1417f19ab01bd7c3a1f5d4849520603f2143020cf61", + "7": "f370a0bf3c16dcc6a3e9e9553cfd619e3fadd52f8c88c2566d72830dff70c636", + "8": "e1b2782ea4affcd15aa6109f329374cefe528aec5e740c14bc4961e124364f20", + "9": "b6d28c03b8f1b3d4369f2eba932743299e0f8d7a4decd0792d044884a63d751c" }, "gfx942_attention_tiled_2d_stress": { - "0": "9f84d35d8d40423921150bb2748c35c4caba0a128b6b3ba570e73977ed2e009f", - "1": "c0768666702a5ea7293ceca4ad759b897928938859d59a65fe84686d47282221", - "10": "68852c54f4b22ec7b694ae826507666cbdc132d535c7e7a239a7573c578ad384", - "11": "1cee802257f5c623c203383afac26bdaab173e16c4d16ce4502bc8d539c417df", - "12": "ff3a35a4093972a266bb04e21053b420d3b93428c88b6c1d6af6cebebcef0fc9", - "13": "d4144178e1b64192fd352e4f45c792e154930bc009d2473b2269f7545576ce9b", - "14": "04eed9253cfac668b188780221d0a5fbbbf12cef5c9792d0768e8514ee6733dc", - "15": "debdedf7dedd9151dad4b472b8cddb9234d834d49bf72b9cd2590372956f1006", - "16": "c453af6947552d3c4367dc0d60a32e64a1f2a3d0a4d72e931d62fbb2cdfe34ea", - "17": "9024b767cf923bdbff2a5bacfebd3498a40c1767ba94f456abb180d79759081d", - "18": "938b3e247ef22cfe674e4a2fb3d528ebd16d42e9bb0c3e508a22b76fbfe4582f", - "19": "ac1e914ef5647819516aedaee18c90c674bf240c6f0b0368abf6ecdd02e13068", - "2": "affcff591aad2cc7a9bc7315f6a3dba605c71f9f38c922e8c101becbdadd1603", - "20": "be8954f0f2de26dd2a05ad04bb4a305ba10814098c12f30c5b44ea46e3904302", - "21": "d201b1ecd26578276a7c48a169c7e9bb2656f33d6ee200b01e2d7930690b9b3d", - "22": "7d72de7d963d4ecc27e4ae1facbe2f2cf9eba2521490b563801889791da5b27e", - "23": "71bc2948e67e0c1dfbda30d7699280d2600d9fc3e6d66f6f9223f51c08db585b", - "24": "dbe56afe12707b4eaf06d6b7cfb41dd7232dd8275e14a95a7ea42bc04d178f4b", - "25": "93f03ff8f910ba9cf8bdfeeed14a5fd607e4a180fc0d9486a9eba62a93567ce3", - "26": "86dc4d94312c39514f01fee6d3f6fb475805297d08af0cf78ae61ce6db3495d2", - "27": "903e94e8b3d9fa6b20ca451ba07f1993c63040b65094d8885d2fefb7442a1432", - "28": "5038356247d3c89972b69f4081f7ff50be0cecc5d3b55c535dd5c69a1059419a", - "29": "e695ba35ffce37b0b5a1149d268bccbb7adabc9bde628582dcfb0bc07dc1afd3", - "3": "ba68073cd010c68ff5cd5658e7b1ed625d5129c41918966923028f8dfcb82a08", - "30": "b38282c34d1b550927257af51bab55c3590c11440daa51657e32d95b3a57e51a", - "31": "fb66fc56dd4f63e5286f30312140751c0a7684211c19370d30e80be437d4fe6a", - "32": "1aa3154a0db07407f18bf5d6c9a0948b2cd51a687a1aa6c4ced0672c7df966e2", - "33": "34d07424595b75d4288fef8cdb12c24c01f57812e6fb0e3a46eee0b7aaf81307", - "34": "944cd092df508ad92936342964c90509b406f3b2fba515d97967fba41092ca1c", - "35": "bbc144a79af8ff93e40eb0c854eb83ae6a730b75a93b72edfd0edac0f5f2ac16", - "36": "5f33b7267514e25c3f4bfe33332c976b6be249b6142ee46916dfd33fb3ec043f", - "37": "0ffa90112362e07f0c5191463f40cbf0927197b2576ad7be099b4ddb6aee495c", - "38": "b09f3359b45fe6043715d3f6b2897a47f4b6767b918ea5a170ab56c9dbf90d8d", - "39": "89da105b57d2527088ec2a52d8c2512150c1d37ec654c2d221fbaddfc5408231", - "4": "02d965b16086e833d3ba97430034908ff2ebe3583a46318595a0f67c4468314a", - "40": "57b177ce073974f53ab5821f79e2d3a8c5f622bbe2cb717b1f49b8f89693f413", - "41": "da37097a6394a3f11491d51dc1891a67d9855825c70de565baaf9142766e0bf4", - "42": "e7086edb0d64c362e9d9d1118037f29829beb2bd13f981a2c76201e22bd1e7d1", - "43": "ee472deeb10b15c19df4367f8b7ed7f7d406fa61e657a0f7e342253db336c12d", - "44": "9ad000d5f39233eafd57634d3ca25d62f87a5d774fb58b3d7f23277d903eaa62", - "45": "9ff3524f1c115974cd7e77d057d8cbf42a61b04b24412f167a704f0eab6e1696", - "46": "6fbc1c8740d573a35ae4398d09b68d17147ab94ce6d0290880cb27da017d08fd", - "47": "75525f7e4696f7506d1b915a6e47e237d406262ff163dd4a3fa25923d070a613", - "48": "174e097b4023a02b59dbed329cfd8aea8047056b6c71e3066c001305ade01620", - "49": "63799ce0e5ae46f96bbbbc5ae6523364190759d2012ff0eaa7b865ee8c72385b", - "5": "915603bd7c78baef390f3d4f17c494ed8f11a2ec124ce75cc86b3a4cf2121484", - "50": "2c2f883f554a2f831597d9e2e9453a878012ad9ee1329459dd16c20da831f297", - "6": "3eff31c2eebf5b40810ab171ada43bdbc83abddb4a0493122d182f85219d1734", - "7": "213ee124983091d878f7caa81893be93d406e92f6d1938e5b96e123201226c54", - "8": "1ad6118c6396a70ae2cc55aa49de5ffab6ed80719a2f4175ff6a3542c2773563", - "9": "b1ac91fddeda3792f2e46f4f37a0956f9e145e04c230b9167dfd1bd95e1e57e2" + "0": "ebf988e14dd9bafc4ed04b3bfdd4a24f2104dd7919107284e5323121cc0b84f8", + "1": "d06203651f8ea522eb20610db9baad58fac91dffba0616d2a67e08517bea151b", + "10": "a4446a4e630b32797c27b1bb292ec78c59b790108476a597e366a3217a10d6c8", + "11": "aaa214baf0d070ed859a39f235ad10dab2a5490cfeaf2b98f718bfd556cd0c78", + "12": "4aa0cc0c8633b6643932b93db37ced480aab130fe96a39b4266607b3c4cb545f", + "13": "4c4c5f6284427b473be4dce7f3b180f08723a2c7ee6245e29ccf2c2161a8cc01", + "14": "a187b6152c2e26835ab1a5934434d50487ba8fb407b9c4a14e0b53f7f64a2ff9", + "15": "0804519025f8961e0a9ecd7fd673b82e5aa99c27432b34fd971162fb41a5a5e1", + "16": "ebf988e14dd9bafc4ed04b3bfdd4a24f2104dd7919107284e5323121cc0b84f8", + "17": "ebf988e14dd9bafc4ed04b3bfdd4a24f2104dd7919107284e5323121cc0b84f8", + "18": "ebf988e14dd9bafc4ed04b3bfdd4a24f2104dd7919107284e5323121cc0b84f8", + "19": "ed4f756a9bf01b9604d033ee1f1e2b003461026ac5c93c38be38411554168752", + "2": "98fe2fab85d720915e27214e187367f528849aa6f86c1a8d4c26ce7929963cc1", + "20": "f8d10cbb7227e4f7a490aa4377da3f18ed8bfa1ae7572127dfbfcef889473a38", + "21": "a42f78186d5871fc983bae1d491febdc9fcdb7fd170b013074cb2cd655cb926b", + "22": "746191a9e8c92f035564eb0c5f80690da26cfc576d81a2deb24dda4d612b75f0", + "23": "e43befef39d27146ae41d7a286eb39af3f76284e101a2dc4bd6d18ad169e9ad1", + "24": "f046bd084317ff25cc2382287e26aa6db17873068f811a8c7bab2231abac0f20", + "25": "0062a09fed5c700347ab67dd5cd920dfd425832650dcb8e12f7c0cf0223236cd", + "26": "97e507bdb52a069a232330065f79d6ce97c4797875d89401df939062096b8849", + "27": "ed5134d8da808bb37b701111300921e9994acd70b3c47a1041b8b2038158fee5", + "28": "198064654041e7340d5cb499c3ffd7114a93a39c1b43d5f76a6ad8e98a4b6594", + "29": "e8253a6fcd2a86753602d1bbeb6e269fe5b76ba476d8a92b44ce5752e4f03d8b", + "3": "e1cb4f55d2d41f49e0e2bede52387a60509c695bba99b866c6765bf997d06d9b", + "30": "7314113ff6bb106ca80de8a2404aeb3365064a5919aa376fe63ef0d1507eb2d8", + "31": "c0f2e8913776fb6f32027ee4e76645a45b851bb4cfc04efc884d6da6975aec1c", + "32": "ebf988e14dd9bafc4ed04b3bfdd4a24f2104dd7919107284e5323121cc0b84f8", + "33": "ebf988e14dd9bafc4ed04b3bfdd4a24f2104dd7919107284e5323121cc0b84f8", + "34": "ebf988e14dd9bafc4ed04b3bfdd4a24f2104dd7919107284e5323121cc0b84f8", + "35": "81ef1ab782e2780eddb3fe87f1da8834ad4388be864e2e1a72f98634e1f12a9a", + "36": "9e0961b2751a35e420c707a1fa1d836c2d48e766882fd1a46b71f3536dc9c91d", + "37": "e0294f81f2d451891065fa9afe6a1cd33e011d365213f87c4435f3b82c49d9dd", + "38": "0f71c60ab87b32d7c20de4cd21e7132889512f4f81661b2410a304c5e4aed99e", + "39": "5aa8ecb43f05a4f9d71e31f60ff4089faa1af7b9d9e55b373b746ea16ffce7af", + "4": "76cad0354f786ac97568c791762580cbdadfe1e1f7d042fe995ab8a924ffaf1c", + "40": "b6e0df9b6fffc6f88460e84581be61aa354476c8f2aa7c2fa6a6fdfc0940d78a", + "41": "c89d27e2b1ccec03ce438c4be523503b2f3702e2a271b4bbee0f5071626a50c8", + "42": "07c104955c0a2bf3695b0852ef3942bed1ac1c60be9c506a7f0314630d6af4ae", + "43": "3f1fac66f6bf8f4e6ab26bf911aa4b0329025c8d4dc104041369edbe3e05a173", + "44": "bfebd0a89451dc2671cca07f2579132ca83cf4126b950bd130fe088dff265ee5", + "45": "04f4f79143d08f9fcd7db3960aa731f4a7bdfb9348e0b6ae134070dce595d073", + "46": "6a3e1a9d72563b212484296656e97dd061e9e4be01a12166f2d931f70d5315d5", + "47": "2e0106c085b10ddf7eb0132739c65bf7aa2ca7a0829045fafb675fd99835aa7b", + "48": "f3beee7e3684d0f5f9ca53014377ef9705b4dab9e0403398bebadf093ee6574c", + "49": "f74457ac84783928a9ec3f01d254b3cf1e6bf4de51ce4e79eecf85bccc16c7bc", + "5": "9f72d4197ca6afc0d81250c9410bb0e5880d26297683138c04d9dabe565b4514", + "50": "db81d8461b690fa4e66171a9d866023df5fc1203c98ed14be438eefa6cc17e36", + "6": "d606354576d9165ad84f497ca39ed638ff1c96591d0836f46b53452e9a257a17", + "7": "6c7de2b5477f939d037a53a07db441ad7dc633f61a7e35c6e6f1951e348901f4", + "8": "567553ee2dce530cec844d0ac9d24d7ad15a8bbfb0a811fee8d8116f868eb784", + "9": "0bee8cedf0247f1d06728b30ccc5d03b5b2bd23f4049a14a123d11d2d67708f7" }, "gfx942_attention_tiled_3d": { - "0": "cc3f2041d379a84dc3bf7b6cce098b0e09de51e2a53da24b8b005b76a88ef2a3", - "1": "0d08e69fd75f45e2213afabb749a795ef854d57fbcfa282185eefc167c6ecfdd", - "10": "b156fd210ce4b5c5dd7fc1acd1713dc8938fd1032387a061178d9e0ef7d165cd", - "11": "bbc6de3c202c39f2b5188698176f94638bc5489527b30e13484234a72e1ad6fd", - "12": "31888b625d263d42d838c7193cfbb8fd9a795aa565a8f5616cbad167158fe60f", - "13": "4dad26133f26869a0d87133e35a2d3a9f0d9732d763852474a30851882b9e453", - "14": "61a575de4d2bf56076920f7a9892b8809b2d1951290e229f560b5cb25bc0234f", - "15": "cd9141e0c9f06a5b405d0ed1f19625eed22c8a594a9c117aa97a6a20c66e8f69", - "16": "9d7f10dd845a8b7399a4dfd6ecaee077cc9c3831cc9e618be75a711cc9acfc4a", - "17": "6302c517f41d7cfcc32f3b1cf3843a121b25ac90e6451c0c4a71d9402f138b95", - "2": "56153f0bef9d8262404f7d0887de981b0bd919840bc2310b6473b98314df26aa", - "3": "9398d2999e4caa2d1fa3ed7be10683b884fb813f22eb0404bb9cf8d472575904", - "4": "e5eddf3ecfd8fae5d62b08af12677b6d15fd9a1111225b4a08ea828cb12d0905", - "5": "523279f91929d6188a955c2490f810830fc15abf00df25dc055c093af8db6da5", - "6": "99e272ef0dff025b3e4b1afb9d3ec8de09fbd6b66ea2d2050755a3c854bbfabd", - "7": "ed5f5549086d0f47a9cd6623d625babf5fb2e6dfe6b5beef8818894049aab80a", - "8": "6e2e0a93bc06db1b2d9c5cdc6af832c212f04af161a6b2efd1197179b42bd5e5", - "9": "cd075c6ad036b6ca70f01126387f4bfc25baa6dc4cfb02752afe038ebb143850" + "0": "ad844ff28e6848c399f7147d29c5620e0546354621463ed36e76487f4cb982f4", + "1": "63e3953c0178ff760382c45c2fe3bb979b2dec3d440b5c143fc4ae45b8aa3683", + "10": "04bb65a8b60b9eb14429a482e53646ee23e16660b21b32870c802975b0c83867", + "11": "6af9bf0abdecfc56386a4de0eeced09ec2d6b52410aed5f5575455a869e01627", + "12": "a74a6d52055c3e08763f5f504ec70fbb84f5ad0bd790da4ddc7f6879933ec4d8", + "13": "05166ad2f3fcb437afed4f6e00553696f72e347b18df0df107e3aa2a5ceda845", + "14": "6b69a670ae8eaefd49e4e05d666be63d3bd4df8cf92f810f6dcf2d29c8950a8f", + "15": "41845128ddaff2006010b88d387299f3fc19b3b65312d77c07ee86953732c1ca", + "16": "c6ccbad3a8d6f43c723fcad24952032adcdbe227176a3f9b8ec9c7b3e25d6ce1", + "17": "a60e6a0999fb68e771ee6be08163ea410e4af178f7c8597382638f8e9194f61b", + "2": "835fcee389e09b34d06562c485a598d812de675b4b61df933d9e925725668227", + "3": "c8bc9de015a452c5e154a9ffe4346b93a9fd90e6093b31e371efc6439c4f4749", + "4": "9dff59a8af03107f77a726caf069ace2298dd58cfef97779e1445dc92e14e0b5", + "5": "b32a6ed27c303eff3d023e6a1872cc7447f9eeb2d576bb22d2e78c4104f59073", + "6": "0763109a0f0ecb853339e0027645a591019a6c45b3ec4d8624e2ea1070c32830", + "7": "b4fd3d0ceebaf0a3d709b162472eaec7e02b8e8b896274cbe302faa3bce10595", + "8": "f0f518ad50fd102173d52d1f7f4096f5bf91a9e103837bc4e69a62fb7a49b4a0", + "9": "de3ca04e9008824143d5e906a0542a32367d1402fbbf49ffc33dc9ccaa0d333d" }, "gfx950_attention_tiled_2d": { "0": "0a23fb0cacc8dae6e12f8c82403e1fcbddb2445715fac624d8519dcd79d5e2ce", @@ -5040,19 +5040,19 @@ "5": "c8eff38fc6979cc723b3bd8e276884108a0655d1284a4e548bab6ecce69116ce" }, "attention_strided_kv": { - "0": "67e7f5bd0d8861dc16e6cff53316d780d40f6232716a5a3ea25f06055ac48442", - "1": "a0b5229967c3892d492cbacb9db2d15f06ce919b1f15c8de8d298e0d74306e46", + "0": "7ce6e33c41e2752f1863c4ef79eef519c4748cc9f1c3035cab5cab030f386d55", + "1": "38fa323beea8da252380810dbbe76686f3cff17bc15aca23c32a73c11f6a70e4", "10": "400f75d835f4bb4d37dfe0039c609deb5e172c77e5a83b87b89071ed3b0017bf", "11": "0e441b0ae432165d00da483191b9470dcfc043c2bacd386b86fe6bef8f05266f", - "12": "03eb993466369d2ca4aabaff7880f23138bf08a0f71154159a0c9e4627e17df3", + "12": "0ec51a3f92220ca3dad8979ef3dbe976f89ab95e28de5a178eaee7477661b8d4", "13": "74d14799531a8e079e822a46af3a22f1875a85a46c6c561ef26427fd44ffd748", - "2": "4cded6921b9951b733d1dbfdfe93015aa7b78c5bd5def3dfbe503175949af968", + "2": "5e39245a8752ce0229c9229aba2b11fc8cb07241e93ebaafa051cae24aab4514", "3": "393df7f6bb892e49b3c3031ff696d11addb1da00ad2088aa881ea37074c7db77", "4": "db62e65f3aa74588eb95d2bf69935aa6331d0a57d42996bbe81c7510e146f119", "5": "ee8248f21f5631e005750601f275daf2fd57d87f102a241c2fee76caa4567d33", - "6": "92c9560cfd6ae383a059972d0f4863e3699b56beba16aab1736d06bb41946143", - "7": "1b3cb6e2c817fce85c1bd189f7a7a01eb7e089c37ef1824cabdfbe5bd6915b47", - "8": "4b0b553b289d0de60dd9401f61f97d2f36cb201b6343437e28e5fbdc370cf529", + "6": "b3add7e87f7a74a750718e45156d80bc9a5725c33a8a25fadc55b71bb600af89", + "7": "94ea80d639a380a3cf0069ff8679e6f197fa1c40c9a4885413d747341964ef9c", + "8": "d8a78f5237a1c466494c27808c41590153f9ae47d899a5d6ad240510e680d89f", "9": "cb2426e5ff5ce00a3094ea4217df160031006f22cfe058568e4c959f1ecff086" }, "attention_unified": { @@ -6247,90 +6247,90 @@ "99": "056fa9ad97d50744d41853632710588920df00099f28006244673116483888d2" }, "gfx942_attention_tiled_2d": { - "0": "683c08c7dffc3cfc7cd0c0853efae565ca4d01acdc02e0149990831e28e791b1", - "1": "f60cb522d2cd4ee1c14b9057e1be4f665d15e090503a2b97909f85a68bcb6301", - "10": "880bdc958260b111bca83758a86805ba077b18452d766a7a17171ece7a8057cd", - "2": "b0edace3965c60e1f2a68dde86a11fcd2ea016fc711bcad083622d6464f3ba53", - "3": "855f84545da23f579b52932f83f9f8c70c40b00bacc3c51b2a28d4adf19964a6", - "4": "c6124f604a4e6ef8e53042aa51534532e369ce7a67fdf51214ba7dbfd2e1ce86", - "5": "737d97c4090c7145c9c8f7de9945e092736a702c83cc81aaca143abe4a46d4ce", - "6": "be28e420e181d8f16f3d7d762eec6b93f33ae94db860c29f31cfa4f2867d5412", - "7": "253477c2582804478fbe3995ebaf936f68515fcab3427271f3bcf2aabefcd76f", - "8": "f84c3a6c9ead604be246102e5e9d8ca35c2208c00dee3b6e8ad74c695a83d3d7", - "9": "038b636e4692b8e73cfdd2da8f7e8baa40fa78a773541afe2a6f340e8bb09f54" + "0": "4ffb68341b74c824a28dbd06e5809df057633e7ddd3bedc0a5b8a9e874e411bc", + "1": "f90746ff0848486572dce468bbe5ac6e80b85573a8d525f87fee60c955cc5b0d", + "10": "77b32296c327d6170b10f8ac8c9ce5d5cb8bef2f462d3d8b5ce8bc95287b4660", + "2": "1ecd8f1b39aab6dd3d1849e3275c3534f989fbd1a877b6a775311b32b2a4ebd5", + "3": "7516ac66acb86878afc9b685d58b7ae9b2d4fc226cb32350df2e9a1223263fa6", + "4": "7e09cc0af6308d4ff0091b53dbff6d80d4d5ea76858ae79f93eb5ba81f46acae", + "5": "4665587b411e3ba1f274c7357f47a514c49b4668c05306e00e0d3f40d6211ce6", + "6": "ef62d2b564e1bdabbf319c07c607f3cd50a17d0df7e95cb59a22debee6dd812c", + "7": "9f63b6295ff616606d52b5cd17630bd65552f3078b3d03c49f3e76f7cabade91", + "8": "668338ed2b1726e12130c0a3e6adbf19aa94e99417c3d56e569eab93eb55846f", + "9": "507fbba3cfe4013106b475479db1c230fce54bdef528d76a4a13a40e94a6778d" }, "gfx942_attention_tiled_2d_stress": { - "0": "3c257f633839e4a04a0f6ca8bcaafcb115889b70376a2acd6cbdfc1b8b33b969", - "1": "9994772d3475461ba2b9d4cefdbef460d2b14d4b5f8bc3498b6e9edc9036dc1e", - "10": "48f8ba43cfac3a2e6da48c0c15790e2b8ed9ec7d1352997ceb5d0be1fb5d71c8", - "11": "12e49591c7c38efb6fff030f380813b42f217b8e2b3a4005358b0f3a40f52e6b", - "12": "68b23f96a31e3262ff5fd1f0f08a0332bd8d4615cf65f6693656ef70957436c1", - "13": "72349d0c343d64d1bd5ce82260e34bb87b95fff05d9b857549bf32d5550c3c02", - "14": "113161bcfbf0788a94c526ac60a4a4f19b7e727715df0c22c86cecc8bed34938", - "15": "7db2356ccfff211545897ee898708a6551b362296ccf1c1065d0c45a78a6be10", - "16": "4e2ae5983ecc1b1a1561406118a9b7e70cbacefe15cb7a64a1ee4b7d12690f4d", - "17": "093a168f00dbd8a88129df3640a5f934a4fbd41fd3dcd35d6793bc1b98e1dcd5", - "18": "ecab656a1e9ec7997164a274c88c6b598c6f2e90f0388660daa63461ab99e1ce", - "19": "df700cb58d63f3a785f8fd977fa7764e15cb9aef62ff4a4b0cbc6fd26faa92c4", - "2": "cae0835ca60926686b9d866dd827db4b5d26228c37532e9ecb3ac5f3e3021b09", - "20": "94894ff5fb06fd2c3babca42e79db1d7d290e67345d72b64af753ae7df8ee201", - "21": "9dff89360ab462fa19206b5f2121e5d42dbd82fa89babb800229fec5c48131a0", - "22": "bbc287df4ae583798b9c18a58b273743304428efa8212b0acf970f360d9a2375", - "23": "cc15a612474b649fd4bec44bcbe9eb69df4387d834f91e9c61ce6fa17420cc76", - "24": "90ab9a31e7543716eb69fac0d30eba43bd9236d07d9cd457d059d200dd674165", - "25": "444c9549d6271944d5aa2cb5579dd7aeeffa0bcbfff03cfe81ff066a4242646b", - "26": "b63acc95bcdfb55c5cd50f33e0a49cfa401f1142cde1148da63457a19fcdf735", - "27": "3bf72339ab3850ed764135ebca922372c7ec994a1bc4694177742ee8cd295efe", - "28": "e8815c584989b0a283182d2b5f103499512485b347927b6530010a2aa3f7f5d4", - "29": "f5b5cdc9f45a4d588344ab1bada174f3bfdecf2e0a2cc8e296eaa60db8046967", - "3": "6c9fa4af8b158d63b06b48aaa7268ee763048c235b5547d8166130dc5a36be11", - "30": "77cb707949de494b6606ca090e20e692e74de58f4d18736c4697f75a5e66adf1", - "31": "949ec4860e313f8a1fb741fb59d1ff8e854a7251f05bb2ed6e1f7a78cf18f774", - "32": "8de2279af7e6d8481187974a1c1cd5f210c4abf2039eb3cea8ffe50a428493ef", - "33": "5a4ff05b311be746655d69c1396d1f086b16b216053a2baaa56e36eb56eed133", - "34": "d35d8a02af195d14051ee3b33238e0a6419457eb4717db37eea2e23b644d6f57", - "35": "e4ae43386d8ec48cf480c0c08b04fafeb38d78e7799f0c98988b7ba8471b3250", - "36": "d0a1e8d743c0896d36873cbdb3356626c604607a290b454cb73b151ea831f178", - "37": "d8f8893533bb07905ae93e5fb120d1d1d7e3f7fc39e564896caf5524f3789337", - "38": "4f79fd30f43b9d6dfbc73038d564b2dd993ca7529b4780d2e22a9f4fd8e1ef5a", - "39": "6bfb00574e61382094beecfc7e23203a99a149f0e7495de48964d82effc3ee8d", - "4": "f5cd7b4764f0c56167dd29ba57773599fdb177f06f55b2581e16fdbbd70aee92", - "40": "64167e911d818ebee3905791215e914983cacc475bf1ea06a40506964a71483f", - "41": "1056e6751129cdf36786ba777abb9588be0be548b328878923b7717e9d2201f2", - "42": "7bea4c58a8c7b2fac886942df6306aea2fb3218bf6d8c1cc2f0ad932259d0bd3", - "43": "2e40a5ca1d56a413b4fa7c69e300a7f4b95e3cc6251e1eaac347acd5b0b178e8", - "44": "7f1d6b3ba2ea51aa057d3725b9feb46f0e685e4336b447c9abd9d488917654e1", - "45": "c9df1f0d7021b66591eec33e3ce826e9604c95fc9eef2aa59471d93388d76b0f", - "46": "c5c46c6ff9b3016447eac716f4c4be43103f1896c2a50985e82515378925b4d6", - "47": "039508f2b259b79c39df687eb61ce0c784976a239f4a89cf0de1a6307648cd68", - "48": "8552d394a743f002fcff91c97ba9c5af5c52de4cc3184c64ccf2dd0330ba5124", - "49": "4110bd2e34aeda9a2304e45b075452629a1f736762669dbc706a316c725103e1", - "5": "9cda0884fd25671f525b3aa61966528e3a9bd638aa8456922911349e6e6d5c94", - "50": "d3e933b40976b8b1a2b04f049ad321c7b38249dca1490db48032bf6ec148c8ae", - "6": "46cf9f774b755cd75522c9a2248506cead8f1171aa7681c3b33e5839882bf4a1", - "7": "f60cb522d2cd4ee1c14b9057e1be4f665d15e090503a2b97909f85a68bcb6301", - "8": "0a0e7e36c71075a84ec51a2e69b733a642645589338385f24e007c5c020e52f0", - "9": "c6124f604a4e6ef8e53042aa51534532e369ce7a67fdf51214ba7dbfd2e1ce86" + "0": "86789f4387016e3deeaf4e529748ec27ef5c8fb1ad32a597b7fa97ec521544bb", + "1": "7ab4604ff2aa1cd71effecbc8db6da5c3f070278b4311034b698e7b4af81af56", + "10": "1749380b9e7e78cc4552ae8233937eac2477ae22cef73488dfd4f85c95376dbe", + "11": "b3ecc5505842ddfa15e5921cc076adcd260e2ce0d6c7fb6dc8a20aa03aef26ca", + "12": "88025cc7e0dcfbf0846e60206edc8855c2cdf4dfe0d287e2f18e26eac93786fa", + "13": "c07232ff8b2f88bc4515dc8cf8bc6c50163a593827c2e0d343aceb898e763b49", + "14": "8391a70cb6408c34515e6c1cd2696cbc0de2a23cd83dfd7dccf7d063be341961", + "15": "2edc5e85548261810f3087b0d1ac5f8fc200770a2d9b11dc131b22448c358d4d", + "16": "86789f4387016e3deeaf4e529748ec27ef5c8fb1ad32a597b7fa97ec521544bb", + "17": "86789f4387016e3deeaf4e529748ec27ef5c8fb1ad32a597b7fa97ec521544bb", + "18": "86789f4387016e3deeaf4e529748ec27ef5c8fb1ad32a597b7fa97ec521544bb", + "19": "f02ca6a5aabafe06558df856660a5d05825fc76c5948c539b19dfef32d703637", + "2": "4ffb68341b74c824a28dbd06e5809df057633e7ddd3bedc0a5b8a9e874e411bc", + "20": "7bd6e1d04d483c7a0644a12abbd43858515e31ed764a1e9e62ed767169516300", + "21": "e89ff280e975cafdfaf5c2756f1b8a9487b5361df41f72b2d7d4464d463d7376", + "22": "c75aa5841567b0320b62be3ba4dc943d35d25bb462a2f414af7353e10771fb79", + "23": "0e28c4f45ef58a9b186f0c75a64a998e5c52b29c809e41df4d8af6f32aa78a57", + "24": "15d96cc7b4952c4a2fadf32f9cdd74a9c7bd7bbe3b29c37204b09e9fd574312f", + "25": "29bd94808f00e9e78ad546d850a30b83f09d62cf3d8a860e9dee649aacbaf84e", + "26": "21eaaba3d60624c49a56fe97637f8ea5b82cb062279a645dc62368ebf784aeed", + "27": "7e25b606affc8ca5dd3245c5bf847f6172c3ecee66139d752c9f64f1c364232a", + "28": "13840460e68c708f652735d9fb09a405321c3e27d34c601c23d2ead5f77cd6d6", + "29": "2ce680139cddf578de6822216d3681173f0b9d0f0786c102034b0ba11ec63314", + "3": "5c39f0189ccaa8e5ce0e7572c567ba3e8ce006b583ea709f4947c91f81e43885", + "30": "5793da387e336bad853f8c20294267383126fa719938a5de147a1ae1f12876bc", + "31": "78a1f674796fef7666ebd7e59a9c23f6387bd9aa870f2c8c559b0516dd838554", + "32": "86789f4387016e3deeaf4e529748ec27ef5c8fb1ad32a597b7fa97ec521544bb", + "33": "86789f4387016e3deeaf4e529748ec27ef5c8fb1ad32a597b7fa97ec521544bb", + "34": "86789f4387016e3deeaf4e529748ec27ef5c8fb1ad32a597b7fa97ec521544bb", + "35": "9c8fc0e32cf741bf2561324c7ced51676ab7e6b10d9d6ca2f6d0c21570c266da", + "36": "df373202e9db27337dacece58ece45e8f51e70b30d7f0c949736b30b85bd0d33", + "37": "4d0d6d559027b2fb6f510b7c4398bf91b5c08f24a53f670082b4c8e42090cad3", + "38": "b1ec4a85fd8f18339dd1012ea88d019a8379bdfd4a6150862b6dbfe8fcaa7bb5", + "39": "76f37a22beff460b1035a57941564e7ec50c1dbcaa896bf2da5e2e41e5e388c4", + "4": "5574fa3bd96351d39394b32773e9df4d53b200c656d882ac491d5b49e1e55d1d", + "40": "78dd706334a2598b8606e4e44fa7f28acf7e223fe0a542e6f52e95269152237f", + "41": "fc3ea47719568291c6ea9300b5a25499c031e9b9f38bbce3df353d9673f16845", + "42": "4f4d18bee80c33bc721e4fbbff6f1f5cb85304ccbd4e02c4224c69ab699c96cb", + "43": "5cb0ee4f37e61557d926cd6b91d37aa33dd8c484c3afcf342b94d4c3c1a2d472", + "44": "c52bdbfc84684ee0c70ef6257d1adf687462dd56f40236a0247521643e5fc980", + "45": "c4eb8afab9997fe7a1e6782177c67fb49337251d40e2b42b93249f48ec4a7f6c", + "46": "e07adc0e5cc5e7f108146808277dee2bc81f6766eb7bbdd8960c020361a23cf4", + "47": "a84a16836ed2eb51efc9e9b123e3d727ebeac364bf9819959f1b6c5c0d8c8def", + "48": "c1dff6bf77261f0c2d44265d27b190a735dd30ba66579578b0a8d1f203ad2750", + "49": "92cb8c2c2b79148804947d05eb3904642908231cf80660f6fe3de8052ffac522", + "5": "6ecaf40f3c945da47ded7e53abb57c677762ecd83c8120892e8687599eb3e488", + "50": "f668a0e5626f9d40ebdd44a54103c3ad40f18ee5ba683dd5a7aeabd020d79a16", + "6": "3e77361fe5aeb076821337a1b0dc929849de765ee04183d1c5c6fe56cf636620", + "7": "f90746ff0848486572dce468bbe5ac6e80b85573a8d525f87fee60c955cc5b0d", + "8": "9e0b9df4efe3dc0d8761b6c52c4f743624ee1e28d345aa746a369ce0e6b6c130", + "9": "7e09cc0af6308d4ff0091b53dbff6d80d4d5ea76858ae79f93eb5ba81f46acae" }, "gfx942_attention_tiled_3d": { - "0": "57ca5ee38d86fc8b4e651192c70ba44ac718c9b2fe30e999ec9e71f91eed1135", - "1": "226af192115e35b584bfbd1b7b8b9ffd8ebdeef9dbded9a5a09e2c9057477e26", - "10": "f1183de138bcf10016206d28a24a162e62f131f56f17c20833daa7c91e55ae30", - "11": "bfef7228e0bdd07aa26811572456b8f7cd9ce439b6c3b43495d2444f6f8eb693", - "12": "41d82ee2d11ab4d157092deb75442f377119feebd48d2a2a03f33b9620c6aae3", - "13": "0a8e2378f921147722c87e3157dfbb13ca751664547b11a9ff80ae0e0670552c", - "14": "70f152fdc8e454549e258c41687d6d9a674bab828876496506af2d88544a6ec8", - "15": "aa50f643e437fa9ed1cf32e4691c7fbdf3a0fe69ecbe1075695ff40b62103f57", - "16": "48a1e5d30c4d66a9cddcdeb9a701854ae76a3e7f00019ab02c22c42bc1535ad3", - "17": "baad169f072e48314087b69b424e32a40a09e96ffe17d4fb0e5779a1ef7c6d3a", - "2": "6f4de2c5041fdca3d24abc5d74cb661146e716ad06b113600cefd059875f26c1", - "3": "e8153af8740777bab3d5613e33d52b40b4616cf4b1890a0559b76c0eb31e2d85", - "4": "3e17cdf6ffb2daa7a7c457c1741f3a6350a5939897d5d7b6e9bc21fa64ca60e0", - "5": "dff48e7a6319aae6c1f0a67bc8eea2e79686a4caaf34b9cb1704cc62d2574f67", - "6": "aa8ab4e1c38d7a1f31eed51bcf6ae1ce78daac1fcb309e60106f96190f36103f", - "7": "d1233449ba68ce5d12e96a6201f487536db220eb32f9290b15940b2d53a56195", - "8": "c38439219df6ba61583c59378ee8842bfb5b62fab6576b3592570c8b52cee17b", - "9": "20855a3cdd3d37414ef6e184050ec341cec92f0e332e47543ec2c3e8310ecc05" + "0": "05f9208d14e20625c2355d08fead09242c01f3cfb6e692c305e4b53e06864f09", + "1": "57c6d9ec3cab9ad101daa297ed465cd9ec8369dad8d838e4fa6b7b38abb72bf1", + "10": "c9b6044bd6ad5613cac0e1271686587923ed527944e2bf551653ca45a4dc2193", + "11": "a95831cc76d51ac3a373bb21eb52c3f2d70f96b7e916e663e569a0498843a38d", + "12": "c23a71e2dfd74511e3b90c3db3a32cb712012e08950c836c5fbead37f67fb40f", + "13": "9743c4e3857b3e4d0ca7fcea728942f3cb4b5921e3f5c1171def03bf5c7cf490", + "14": "44eb3c2c08af758651044a93fe35e868724141edbd74166500744bb12ec207ab", + "15": "ab732e91ba542852f9b8d8d69c99b2cdf743bb9d5391d4ada0f31fe5d261af36", + "16": "6b4ed78e7af2fa4eaa0a2bf72c035401116e480af9e57bc257149a50b577a559", + "17": "3c93a4c7f94f8d4850c5a2c49a7d0bcc11a121dd635f1c687c1060f73742f429", + "2": "640e6b3eb547b74221c2c3d4d2c1dc5b91da62f92be5e6e570e99978b4961ce3", + "3": "14931a184dcdd75b8c70d2fa7fcb7486a215f10cccd37867e0b95f9ae45b6566", + "4": "47309bd8554215f56236b38fb0b9c37d3253eb0d54b4437b94244fb27437a3b7", + "5": "1946bc1f67069ffda2dc161d2a8d11daededa922c1458942b266f481c97bec08", + "6": "3403e4686fcba6f26cb2ab502b028205f18c651d9ca0a84b961ac08add09f1f1", + "7": "6f45ac9918b237fee5ef9f51734569ff2eff1e6871f4500965a0a0ab01de94d1", + "8": "ccd6c331cc69c39faab3dae1759c3fc90aa63d5383f68f0493bb055e612edccd", + "9": "0780bfda4117929b13f56941e11205d93fc75afe332d5ef49b91c55d76d4bf34" }, "gfx950_attention_tiled_2d": { "0": "8fcdbf0c36481459f4ead47cc0d3fdd8dc2bbe1aef001e292c9135b38d06ca16", From bdb78d8e123ea7d94d3a47fdadddffa2066b17c0 Mon Sep 17 00:00:00 2001 From: David Langbehn <47395968+dlangbe@users.noreply.github.com> Date: Tue, 3 Nov 2026 18:03:18 -0600 Subject: [PATCH 2/8] perf(rocke): make gfx950 tiled attention heads, window and batch runtime (#14727) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit JIRA ID : AICK-2675 ## Motivation hipDNN AOT packaging ships one compiled binary per kernel symbol, and a packaged catalog cannot enumerate head counts, sliding-window sizes or batch sizes. The gfx950 tiled attention kernels (2D, fastKV register-P 2D, 3D split-KV and reduce) baked all three into their IR. So every head pair, every window and every batch bucket compiled its own binary, and the fastKV register-P builder could not be packed at all because it had no spec-first support predicate. This PR applies to gfx950 what #14530 (AICK-2676) did for gfx942. The kernels read the head counts and the window as kernargs and derive the `cu_q` search trip count from `num_seqs`. One gfx950 binary now serves every head pair at a given GQA ratio, every window size and every batch, and kernel names and cache keys carry exactly what stays baked. Stacked on #14530, which adds the shared plumbing this reuses: `RUNTIME_KERNARG_FIELDS` / `RUNTIME_SEARCH_ITERS`, the arch-aware launch signatures, `binary_search_seq_idx_trip_count`, and mixed/i64 runtime strides in `DynamicTensorDescriptor`. ## Technical Details One commit per change; every kernel change is mirrored in the C++ engine in the same commit. - **Batch** (`3e02491`): the gfx950 2D, fastKV register-P and 3D segment kernels take the search trip count from `num_seqs` at runtime, and the tiled cache keys drop the batch bucket. The single-batch schedule flags in the 2D key now collapse to `None` when neither is set, so `num_seqs == 1` and `> 1` no longer split identical kernels. - **Sliding window** (`ddc344e`): the window size is a trailing i32 kernarg. Whether a window applies still selects structure, so the name token becomes `sw` and the keys carry `sliding_window > 0`. The gfx950 benchmark, parity and build-test launch sites switch to `_tiled_2d_signature` / `_tiled_3d_signatures`. - **Head counts** (`74e8aa5`): - The kernels read `num_query_heads` / `num_kv_heads` as kernargs and derive the Q/O, segment-output and paged K/V strides at runtime. Only the GQA ratio stays baked: names carry `g`, the reduce name drops `h`, and the keys carry the ratio. - The fast paged-KV descriptor keeps its pinned 64/8 geometry. `_enable_fast_paged_kv_desc` now resolves it for both the spec builder and the 2D key, so a 32/4 problem at ratio 8 never reuses a 64/8 fast binary. - Every caller of the shared gfx942 2D C++ prologue (which also builds the gfx950 fastKV register-P twin) now reads its shape at runtime, so its `runtime_shape` / `runtime_search` / `runtime_window` gates and the unused 2D `binary_search_iters` helper are gone. - **K/V copy offsets and page-id prefetch** (`25bac46`): with a runtime kv head count, every bf16/fp16 K/V copy in the 2D loop multiplied its token index by the runtime stride. - When a copy's tokens sit in one page, its offset is now built from three parts: - the page base, passed in `soffset`; - the copy's first-token offset, built once before the loop; - the lane's own part, also built once before the loop. - The KV loop carries the page ids of the current and next tile and reads the ids of the tile after next right after the iter-start drain, as the 3D page-id loader does. - Copies that span pages, the fp8 loaders, and the K-ring / grouped-KV2 schedules keep their existing addressing. - New emit config 56 (the shipping bf16 d64/b16 multi-batch spec) keeps byte-identity coverage on the page-spanning path. - **Packaging predicate** (`12d8242`): `supports_fastkv_register_p_2d_spec` reads its parameters off the spec and also refuses what the builder raises on; it shares those checks with the builder. `hkp_pack` maps the builder to it in `_PREDICATE_ALIASES`. - **i64 name token** (`6569f53`): the gfx950 2D and fastKV register-P names carry `i64` when `use_i64_kv_addr` is set, matching gfx942 2D. gfx950 3D already carries `i64kv`. - **Tests** (`5d10262`, `802b732`): - new `test_attention_gfx950_runtime_shape_numeric.py`, enrolled in the library GPU test list in `platform/CMakeLists.txt`; - the cache-key invariant sweep now has two head pairs per GQA ratio and a second window size, so it catches an under-key between them; - new targeted key and launch-geometry tests for gfx950. ## Risk Assessment Risk 4. Every gfx950 tiled kernel changes its argument list and emitted IR, the tiled cache keys change, and the 2D K/V loaders and KV-loop carry change in both engines. That is a performance-critical path. Mitigations: - byte identity is GREEN at llvm20, llvm22 and llvm23 for every family; - only gfx950 golden entries move; - gfx950 numerics and a baseline-vs-branch A/B were run on hardware (below). ## ASIC Coverage - **gfx950:** all kernel changes. Verified manually on gfx950: the new numeric test, the existing gfx950 GPU suites, and an A/B against #14530's head. Superbuild and TheRock CI do not run these GPU numerics on gfx950. - **gfx942:** C++-only cleanup of the shared 2D prologue (constant branches nobody calls any more); the gfx942 Python IR and goldens do not move. - **gfx1250:** test parametrization only. ## Test Plan - **Byte identity:** `ROCKE_LLVM_FLAVOR=llvm{20,22,23} python tools/check_byte_identity.py`, every family. - **Goldens:** - `python tests/instances/rocke_ir_parity_harness.py --check tests/golden/rocke_representative_ir_sha256.json`; - the strided-KV library golden; - both re-blessed only where a commit moves gfx950 IR. - **Library host tests (GPU hidden):** `python -m pytest ` per file over the attention, MLA, strided-KV, dispatch and layering set (69 files), plus `descriptor-packaging/tests/rocke/test_hkp_pack_producer_guards.py`. - **Cache keys (CPU):** - the invariant sweep (key equality ⇔ build equality over every selector threshold); - targeted share/split pairs; - an IR-mode A/B over 311,040 gfx950 problems, checking routing and that no problem is served a binary built for a different spec. - **gfx950 GPU numerics:** - `test_attention_gfx950_runtime_shape_numeric.py`: pairs of shapes that share one key run back to back through `run_unified_attention_torch`. The second launch must compile nothing, and both must match an fp32 paged reference (max_abs ≤ 4e-2; bf16/fp16 randn inputs) with unused block-table slots on a poison block. - Pairs cover 2D default, a window, the single-batch combo with and without a window, the 64/8 fast descriptor (batch only) and its 32/4 neighbour, the D256 fast route, 3D decode and two-token 3D. - **Existing gfx950 GPU suites:** unified-mask, additive-bias, LSE, 3D running-max, KV-LDS swizzle, wave page-id load, strided-KV decode and dense. - **A/B:** #14530's head vs this branch, run end to end through `run_unified_attention_torch` with production routing (`backend="auto"`). - **AOT dry run:** a draft gfx950 2D engine manifest, pinned to one tuning candidate, run through `registry_enumerator.py --plan-space`. ## Test Result - **Byte identity:** GREEN at llvm20, llvm22 and llvm23. - **Goldens:** only gfx950 entries move, at every flavor. Representative-IR cases per commit: - batch: 21; - window: 21; - head counts: 24 (adds the 3 reduce cases); - K/V loader: 15 (2D only). - The 7 gfx950 strided-KV configs move with the batch, window and head-count commits. - No case is added or removed, and no gfx942 or gfx1250 entry changes. - **Host tests:** results match develop. The 26 remaining failures in `test_strided_kv_lds.py`, `test_strided_kv_decode.py` and `test_selector_parity.py` fail identically on develop. - **hkp_pack:** `test_hkp_pack_producer_guards.py` passes (24/24), including the fastKV register-P accept/refuse case. The rest of `descriptor-packaging/tests/rocke/` needs the hipDNN plugin SDK data, so it is left to PR CI. - **IR-mode A/B (CPU):** - all 311,040 problems route identically; - no problem is served another spec's binary; - distinct binaries for the sweep drop from 10,662 to 2,094 (2D) and from 5,040 to 3,216 (3D). - **gfx950 GPU numerics:** - runtime-shape test 18/18; - unified-mask 315 passed (16 skipped, other arches); - additive-bias 185/185, LSE 57/57, dense 322/322; - 3D running-max, KV-LDS swizzle, wave page-id load and strided-KV decode pass (other-arch cases skipped). - The 3D-only suites ran before the 2D loader commit, which does not touch 3D. **A/B, #14530 head vs this branch (gfx950, 256 CUs)** - **Workload:** 132 problems across 11 cohorts: - d64/d128/d256, bf16/fp16 and fp8 K/V; - GQA ratios 1, 4 and 8, including 64/8 vs 32/4 and a 64/16 head pair; - sinks, sliding windows (128, 256, 4096 fold), the D256 fast route and the single-batch combo; - 2D and 3D paths. - **Method:** HIP-event mean over 50 launches after 10 warmup, median over 3 rounds. Arm order alternates per round, and each arm runs in a fresh process each round. - **Toolchain:** torch 2.10.0+rocm7.0 wheels with the wheel's comgr (LLVM 20 flavor), Python lowering, the same for both arms. - **Baseline:** #14530's head, `c620a3ee6add`. Candidate: this branch at `25bac467c15a`. - **Numeric gate:** every row of both arms is within 4e-2 of an fp32 reference. Worst max_abs is 1.54e-2 in both arms, and max_abs is identical between arms on every row. - **Results (branch / baseline):** | Measure | Result | |---|---| | Compiled binaries per process | 35 vs 73 | | Total first-launch time | 0.32× | | Steady state, geometric mean over 132 problems | 0.980× (range 0.82–1.06×) | | D256 fast route (6 problems) | 0.947× | | bf16 d128 GQA-4 (40 problems) | 0.961× | | bf16 d128, 64/16 heads (12 problems) | 0.978× | | Sliding-window cohorts (3 cohorts, 7 problems) | cohort geomeans 0.887–0.955× | | d64 GQA-8 block-32, incl. the 64/8 fast descriptor (35 problems) | 1.007× | | Other cohorts | within 0.99–1.02× | - **Outliers:** one problem is above 1.05×: bf16 d128 GQA-4, 4 sequences, q 257, k 4096, at 1.064×. - **Without the loader commit:** per-commit timing on the same harness put the head-count change at 1.49–1.82× on the D256 fast route and 1.05–1.10× on long multi-batch d128 prefill, all from the per-copy runtime stride multiply. The `25bac46` loader change recovers both. **AOT dry run (`--plan-space`)** - **Manifest:** a draft gfx950 2D engine manifest pinned to `gfx950_u2d_transposed32_nw4_mw32_t2xb_llvm`, 24 cells (6 head configurations × bf16/fp16 × block 16/32), with the 21 varying gfx950 knobs classified as performance fields. - **Result:** 445 instances, within the default per-arch budget of 1000. They build 211 distinct kernels on this branch against 421 on the baseline, because head pairs at one GQA ratio now share a kernel. ## Testing Checklist - [x] Byte identity - `ROCKE_LLVM_FLAVOR=llvm{20,22,23} python tools/check_byte_identity.py` - Status: Passed - [x] Representative-IR golden - `python tests/instances/rocke_ir_parity_harness.py --check tests/golden/rocke_representative_ir_sha256.json` - Status: Passed - [x] Library host tests, GPU hidden - `python -m pytest ` per file - Status: Passed (3 files fail identically on develop) - [x] hkp_pack producer guards - `python -m pytest descriptor-packaging/tests/rocke/test_hkp_pack_producer_guards.py` - Status: Passed - [x] gfx950 GPU: new runtime-shape numeric test - Status: Passed (18/18) - [x] gfx950 GPU: existing attention suites - Status: Passed - [x] gfx950 GPU: baseline-vs-branch A/B, 132 problems, 3 rounds - Status: Passed (numeric gate) ## Submission Checklist - [ ] Look over the contributing guidelines at https://github.com/ROCm/TheRock/blob/main/GOVERNANCE.md#pull-requests. --------- Co-authored-by: Claude --- ...ROCKE_ENGINE_test_categories_external.yaml | 3 + .../python/hkp_pack/rocke_compile.py | 10 +- .../rocke/test_hkp_pack_producer_guards.py | 31 + .../prefill/benchmark_prefill2d_live.py | 4 +- .../benchmark_prefill2d_fastkv_regp.py | 24 +- .../prefill/benchmark_prefill2d_live.py | 10 +- .../prefill/benchmark_prefill2d_traces.py | 16 +- .../builders/common/attention_spec_builder.py | 11 +- .../prefill/parity_unified_attention.py | 10 +- .../library/dispatch/attention/common.py | 4 +- .../dispatch/attention/strided_decode.py | 3 +- .../kernels/common/attention_unified.py | 60 +- .../kernels/gfx950/attention_tiled_2d.py | 255 +++++-- .../gfx950/attention_tiled_2d_fastkv_regp.py | 64 +- .../kernels/gfx950/attention_tiled_3d.py | 134 ++-- .../attention/test_unified_lse_dispatch.py | 8 +- .../attention_strided_kv_ir_sha256.json | 42 +- .../parity/gfx950_attention_tiled_2d_emit.c | 62 ++ .../parity/gfx950_attention_tiled_2d_emit.py | 57 ++ ...x950_attention_tiled_2d_fastkv_regp_emit.c | 10 + ...950_attention_tiled_2d_fastkv_regp_emit.py | 4 + .../library/tests/test_attention_builds.py | 16 +- ..._attention_gfx950_runtime_shape_numeric.py | 329 ++++++++ .../tests/test_attention_tiled_cache_key.py | 233 ++++-- .../rocke/platform/CMakeLists.txt | 17 +- ...ocke.instances.gfx942.attention_tiled_2d.h | 3 +- ...tance_gfx942_attention_tiled_2d_internal.h | 17 +- .../instance_gfx942_attention_tiled_3d.h | 4 + ...tance_gfx950_attention_tiled_2d_internal.h | 41 +- .../instance_gfx950_attention_tiled_3d.h | 3 +- ...tance_gfx950_attention_tiled_3d_internal.h | 21 +- .../attention_tiled_2d_ctx_prologue.cpp | 38 +- .../gfx942/attention_tiled_2d_epilogue.cpp | 7 +- .../attention_tiled_2d_kv_body_qk_softmax.cpp | 8 +- ...ttention_tiled_2d_kv_descriptors_loads.cpp | 46 +- .../attention_tiled_2d_loop_scaffold.cpp | 21 +- .../attention_tiled_2d_public_entry_glue.cpp | 3 + .../attention_tiled_2d_q_load_gather.cpp | 21 +- ...ed_3d_gfx942_attn_tiled_3d_public_glue.cpp | 24 + .../attention_tiled_2d_ctx_prologue.cpp | 25 +- ...attention_tiled_2d_fastkv_regp_helpers.cpp | 19 +- ...attention_tiled_2d_kv_body_pv_epilogue.cpp | 24 +- .../attention_tiled_2d_kv_body_qk_softmax.cpp | 30 +- ...ttention_tiled_2d_kv_descriptors_loads.cpp | 221 +++++- .../attention_tiled_2d_public_entry_glue.cpp | 40 +- .../attention_tiled_2d_q_load_gather.cpp | 71 +- .../gfx950/attention_tiled_3d_public_glue.cpp | 36 +- .../gfx950/attention_tiled_3d_reduce.cpp | 40 +- ...tention_tiled_3d_segment_loop_epilogue.cpp | 7 +- .../attention_tiled_3d_segment_prologue.cpp | 116 +-- .../python/rocke/helpers/transforms.py | 6 +- .../tests/core/test_transforms_dynamic.py | 9 + .../rocke_representative_ir_sha256.json | 432 +++++------ .../differential/golden/llvm_gfx_all.json | 704 +++++++++--------- 54 files changed, 2344 insertions(+), 1110 deletions(-) create mode 100644 dnn-providers/hip-kernel-provider/rocke/library/tests/test_attention_gfx950_runtime_shape_numeric.py diff --git a/dnn-providers/hip-kernel-provider/ROCKE_ENGINE_test_categories_external.yaml b/dnn-providers/hip-kernel-provider/ROCKE_ENGINE_test_categories_external.yaml index e692d3e0cf2c..f9c41c9667c9 100644 --- a/dnn-providers/hip-kernel-provider/ROCKE_ENGINE_test_categories_external.yaml +++ b/dnn-providers/hip-kernel-provider/ROCKE_ENGINE_test_categories_external.yaml @@ -67,6 +67,7 @@ test_categories: - "rocke_differential_golden_canonical" - "rocke_library_conv_pytest" - "rocke_library_attention_cpu_pytest" + - "rocke_library_attention_cache_key_pytest" - "rocke_library_attention_pytest" - "rocke_library_gdn_pytest" - "rocke_library_attention_gpu_pytest" @@ -108,6 +109,7 @@ test_categories: - "rocke_differential_golden_canonical" - "rocke_library_conv_pytest" - "rocke_library_attention_cpu_pytest" + - "rocke_library_attention_cache_key_pytest" - "rocke_library_attention_pytest" - "rocke_library_gdn_pytest" - "rocke_library_attention_gpu_pytest" @@ -125,6 +127,7 @@ test_categories: description: "rocKE generated-artifact reproduction checks" test_patterns: - "rocke_arch_domain_reproducibility" + - "rocke_library_attention_cache_key_sweep_pytest" labels: - "host" - "toolchain" diff --git a/dnn-providers/hip-kernel-provider/descriptor-packaging/python/hkp_pack/rocke_compile.py b/dnn-providers/hip-kernel-provider/descriptor-packaging/python/hkp_pack/rocke_compile.py index a1898df7bc7c..a4af0dba3d17 100644 --- a/dnn-providers/hip-kernel-provider/descriptor-packaging/python/hkp_pack/rocke_compile.py +++ b/dnn-providers/hip-kernel-provider/descriptor-packaging/python/hkp_pack/rocke_compile.py @@ -323,14 +323,16 @@ def _module_from_source(source): # Builders whose validation predicate is not derivable from the builder name. -# The tiled ``supports_tiled_2d``/``3d`` are keyword-only, so these builders use the -# spec-first ``supports_tiled_*_spec`` wrappers in the gfx942/gfx950 modules. The -# reduce and fastkv-register-P tiled builders and all gfx1250 tiled builders have no -# spec-first predicate, so ``_check_support_predicate`` refuses to pack them. +# The tiled ``supports_tiled_2d``/``3d`` and ``supports_fastkv_register_p_2d`` are +# keyword-only, so these builders use the spec-first ``supports_*_spec`` wrappers in +# the gfx942/gfx950 modules. The reduce tiled builders and all gfx1250 tiled +# builders have no spec-first predicate, so ``_check_support_predicate`` refuses to +# pack them. _PREDICATE_ALIASES = { "build_unified_attention_2d_tiled": "supports_tiled_2d_spec", "build_gfx942_4warp_gqa": "supports_gfx942_4warp_gqa_spec", "build_unified_attention_3d_tiled": "supports_tiled_3d_spec", + "build_unified_attention_2d_fastkv_register_p": "supports_fastkv_register_p_2d_spec", } diff --git a/dnn-providers/hip-kernel-provider/descriptor-packaging/tests/rocke/test_hkp_pack_producer_guards.py b/dnn-providers/hip-kernel-provider/descriptor-packaging/tests/rocke/test_hkp_pack_producer_guards.py index 18411294c2df..2ae1f72e8fc5 100644 --- a/dnn-providers/hip-kernel-provider/descriptor-packaging/tests/rocke/test_hkp_pack_producer_guards.py +++ b/dnn-providers/hip-kernel-provider/descriptor-packaging/tests/rocke/test_hkp_pack_producer_guards.py @@ -22,6 +22,7 @@ _module_from_source, _require_spec_arch_signature, _resolve_spec_class, + _resolve_support_predicate, build_spec, rocke_compile_arch, ) @@ -340,6 +341,36 @@ def test_every_shipped_rocke_builder_passes_the_pack_time_predicate_check(): pytest.fail(f"{source}:{builder} @ {arch} ({target}): {exc}") +@pytest.mark.skipif(not _HAS_ROCKE_LIBRARY, reason="rocke/library not in this tree") +def test_fastkv_register_p_builder_packs_through_its_spec_predicate(): + """The gfx950 fastKV register-P builder resolves its spec-first predicate by + alias; the predicate accepts the experiment's spec and refuses one the + builder would raise on.""" + module = import_module("kernels.gfx950.attention_tiled_2d_fastkv_regp") + builder = "build_unified_attention_2d_fastkv_register_p" + name, _ = _resolve_support_predicate(module, builder) + assert name == "supports_fastkv_register_p_2d_spec" + spec = module.make_fastkv_register_p_spec( + module.UnifiedAttention2DTiledSpec( + head_size=64, + block_size=32, + num_query_heads=64, + num_kv_heads=8, + dtype="bf16", + use_sinks=False, + sliding_window=0, + has_softcap=False, + ) + ) + _check_support_predicate(module, builder, spec, "gfx950") + no_fast_desc = dataclasses.replace(spec, use_fast_paged_kv_desc=False) + with pytest.raises(HkpPackError, match="use_fast_paged_kv_desc"): + _check_support_predicate(module, builder, no_fast_desc, "gfx950") + # The shared tiled gate admits gfx942 for its narrow builder; this one must not. + with pytest.raises(HkpPackError, match="gfx950 kernel"): + _check_support_predicate(module, builder, spec, "gfx942") + + # --- C. Literal validation -------------------------------------------------- @dataclasses.dataclass class _LiteralSpec: diff --git a/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/gfx942/attention/prefill/benchmark_prefill2d_live.py b/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/gfx942/attention/prefill/benchmark_prefill2d_live.py index 4f335715f6c4..f03ef6c5f7c9 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/gfx942/attention/prefill/benchmark_prefill2d_live.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/gfx942/attention/prefill/benchmark_prefill2d_live.py @@ -525,12 +525,12 @@ def build( use_i64_kv_addr=flags["use_i64_kv_addr"] or _enable_i64_kv_addr(problem), use_register_pv=flags["use_register_pv"], ) - # kernel_name() does not encode the KV addressing width, so key on it too. + # The name carries i64 addressing; the signature keys the per-shape + # launch values the cached spec feeds. key = ( shape.signature, variant, spec.kernel_name(), - spec.use_i64_kv_addr, self.compile_backend, ) if key not in self._launchers: diff --git a/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/gfx950/attention/prefill/benchmark_prefill2d_fastkv_regp.py b/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/gfx950/attention/prefill/benchmark_prefill2d_fastkv_regp.py index 4c6e990e648e..646bff265ff5 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/gfx950/attention/prefill/benchmark_prefill2d_fastkv_regp.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/gfx950/attention/prefill/benchmark_prefill2d_fastkv_regp.py @@ -8,11 +8,10 @@ ``attention_tiled_2d_fastkv_regp.py`` and optionally joins against a Triton CSV by ``shape_signature``. -The emitted 2D kernel name omits specialization constants such as ``num_seqs`` -and the derived binary-search trip count. This harness appends a short shape -hash to each built ``KernelDef.name`` before compiling, so the ROCm module/HSACO -cache never sees two different shape-specialized code objects with the same -kernel symbol. +Head counts, window and batch are kernargs, so shapes that share a kernel name +share a code object. This harness still appends a short shape hash to each built +``KernelDef.name`` before compiling, so every trace record loads its own module +and the timings stay independent of the order records run in. Example: @@ -308,7 +307,7 @@ def _variant_spec_and_builder( def _launcher(self, shape, problem, variant: str, sliding_window: int): from rocke import compile_kernel from kernels.common.attention_unified import ( - _attn_signature, + _tiled_2d_signature, _enable_i64_kv_addr, ) from rocke.runtime import KernelLauncher @@ -317,28 +316,25 @@ def _launcher(self, shape, problem, variant: str, sliding_window: int): shape, problem, variant, sliding_window ) spec = replace(spec, use_i64_kv_addr=_enable_i64_kv_addr(problem)) - # kernel_name() does not encode the KV addressing width, so key on it too. + # The name carries i64 addressing; the signature keeps one module per + # trace record (see the module docstring). key = ( variant, shape.signature, spec.kernel_name(), - spec.use_i64_kv_addr, self.compile_backend, ) if key not in self._launchers: kernel = builder(spec) - # The display name omits compile-time shape constants. Keep the - # human-readable spec name in CSV, but compile a shape-unique symbol - # to avoid ROCm HSACO/module-cache aliasing across trace records. + # Keep the human-readable spec name in CSV, but compile a + # shape-unique symbol so each trace record loads its own module. kernel.name = _shape_kernel_name(kernel.name, variant, shape.signature) artifact = compile_kernel(kernel, capture_ir_text=False) self._launchers[key] = ( KernelLauncher( hsaco=artifact.hsaco, kernel_name=artifact.kernel_name, - signature=_attn_signature( - spec.dtype, include_bt_stride=True, include_qq_bias_stride=True - ), + signature=_tiled_2d_signature("gfx950", spec.dtype), cache_key=("prefill2d_fastkv_regp", key), ), spec, diff --git a/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/gfx950/attention/prefill/benchmark_prefill2d_live.py b/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/gfx950/attention/prefill/benchmark_prefill2d_live.py index 934731c21c58..5373adc5f35e 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/gfx950/attention/prefill/benchmark_prefill2d_live.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/gfx950/attention/prefill/benchmark_prefill2d_live.py @@ -436,7 +436,7 @@ def build( supports_tiled_2d, ) from kernels.common.attention_unified import ( - _attn_signature, + _tiled_2d_signature, _enable_i64_kv_addr, _select_2d_compile_backend, ) @@ -498,12 +498,12 @@ def build( # A cache past the i32 buffer range needs i64 whatever the variant asks. use_i64_kv_addr=flags["use_i64_kv_addr"] or _enable_i64_kv_addr(problem), ) - # kernel_name() does not encode the KV addressing width, so key on it too. + # The name carries i64 addressing; the signature keys the per-shape + # launch values the cached spec feeds. key = ( shape.signature, variant, spec.kernel_name(), - spec.use_i64_kv_addr, self.compile_backend, ) if key not in self._launchers: @@ -521,9 +521,7 @@ def build( KernelLauncher( hsaco=artifact.hsaco, kernel_name=artifact.kernel_name, - signature=_attn_signature( - dtype, include_bt_stride=True, include_qq_bias_stride=True - ), + signature=_tiled_2d_signature("gfx950", dtype), cache_key=("prefill2d_live", key), ), spec, diff --git a/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/gfx950/attention/prefill/benchmark_prefill2d_traces.py b/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/gfx950/attention/prefill/benchmark_prefill2d_traces.py index 3af6de7a0664..dffe20e7b64f 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/gfx950/attention/prefill/benchmark_prefill2d_traces.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/gfx950/attention/prefill/benchmark_prefill2d_traces.py @@ -128,7 +128,7 @@ def _launcher(self, shape, problem, sliding_window: int): supports_tiled_2d, ) from kernels.common.attention_unified import ( - _attn_signature, + _tiled_2d_signature, _enable_i64_kv_addr, ) from rocke.runtime import KernelLauncher @@ -181,15 +181,13 @@ def _launcher(self, shape, problem, sliding_window: int): use_fast_paged_kv_desc=use_combo, use_i64_kv_addr=_enable_i64_kv_addr(problem), ) - # The display name intentionally omits several compile-time constants - # (notably num_seqs/binary_search_iters). Cache on the trace signature - # as well so shapes with the same kernel_name never reuse an HSACO that - # was specialized for a different batch geometry. - # kernel_name() does not encode the KV addressing width either. + # The kernel name identifies the binary (head counts, window and batch + # are kernargs; i64 addressing is in the name). Key on the trace + # signature as well because the cached spec also feeds per-shape launch + # values such as the K/V address limit. key = ( shape.signature, spec.kernel_name(), - spec.use_i64_kv_addr, self.compile_backend, ) if key not in self._launchers: @@ -199,9 +197,7 @@ def _launcher(self, shape, problem, sliding_window: int): KernelLauncher( hsaco=artifact.hsaco, kernel_name=artifact.kernel_name, - signature=_attn_signature( - dtype, include_bt_stride=True, include_qq_bias_stride=True - ), + signature=_tiled_2d_signature("gfx950", dtype), cache_key=("prefill2d_trace_combo", key), ), spec, diff --git a/dnn-providers/hip-kernel-provider/rocke/library/builders/common/attention_spec_builder.py b/dnn-providers/hip-kernel-provider/rocke/library/builders/common/attention_spec_builder.py index ef9b99611332..4ae910048889 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/builders/common/attention_spec_builder.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/builders/common/attention_spec_builder.py @@ -21,6 +21,7 @@ _has_score_bias, _enable_combo_2d, _enable_early_v_schedule, + _enable_fast_paged_kv_desc, _enable_fp8_mfma_qk, _enable_gfx942_3d_wide_kv_load, _enable_gfx942_bf16_flash, @@ -234,15 +235,7 @@ def _base_2d_generic_fields(problem: UnifiedAttentionProblem, arch: str) -> dict use_transposed_mask_limit=mask_opts, use_mfma32_skip_legacy_qreg=skip_legacy_qreg, use_early_v_schedule=_enable_early_v_schedule(problem, arch), - use_fast_paged_kv_desc=( - combo_no_sw - and problem.dtype - == "bf16" # fast descriptor is bf16-only (spec validator enforces) - and not problem.use_fp8 - and problem.num_query_heads == 64 - and problem.num_kv_heads == 8 - and _select_2d_tile_size(problem, arch) == 64 - ), + use_fast_paged_kv_desc=_enable_fast_paged_kv_desc(problem, arch), use_register_pv=_enable_register_pv(problem, arch), use_fp8_mfma_qk=_enable_fp8_mfma_qk(problem, arch), use_i64_kv_addr=_enable_i64_kv_addr(problem), diff --git a/dnn-providers/hip-kernel-provider/rocke/library/builders/gfx950/attention/prefill/parity_unified_attention.py b/dnn-providers/hip-kernel-provider/rocke/library/builders/gfx950/attention/prefill/parity_unified_attention.py index 36c0d147abe7..1f104f41c8a5 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/builders/gfx950/attention/prefill/parity_unified_attention.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/builders/gfx950/attention/prefill/parity_unified_attention.py @@ -1720,7 +1720,7 @@ def _run_rocke( supports_tiled_2d, ) from kernels.common.attention_unified import ( - _attn_signature, + _tiled_2d_signature, _attn_values, _spec_kv_addr_limit, _select_2d_compile_backend, @@ -1760,9 +1760,7 @@ def _run_rocke( launcher = KernelLauncher( hsaco=artifact.hsaco, kernel_name=artifact.kernel_name, - signature=_attn_signature( - dtype_str, include_bt_stride=True, include_qq_bias_stride=True - ), + signature=_tiled_2d_signature(_ARCH, dtype_str), cache_key=("d256_fast_direct", spec.kernel_name()), ) if probe_occupancy: @@ -1833,9 +1831,7 @@ def call_once(): launcher = KernelLauncher( hsaco=artifact.hsaco, kernel_name=artifact.kernel_name, - signature=_attn_signature( - dtype_str, include_bt_stride=True, include_qq_bias_stride=True - ), + signature=_tiled_2d_signature(_ARCH, dtype_str), cache_key=("prod_2d_direct", spec.kernel_name()), ) vals = _attn_values( diff --git a/dnn-providers/hip-kernel-provider/rocke/library/dispatch/attention/common.py b/dnn-providers/hip-kernel-provider/rocke/library/dispatch/attention/common.py index f2ddac05f769..3ac56419d90e 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/dispatch/attention/common.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/dispatch/attention/common.py @@ -72,7 +72,9 @@ from rocke.dispatch.tuning.identity import jsonable, normalize_knobs FAMILY = "attention_unified" -ATTENTION_ABI_VERSION = "hipkg-attention-unified/v1" +# v2: the gfx942 and gfx950 tiled kernels take the head counts and the sliding +# window as kernargs. +ATTENTION_ABI_VERSION = "hipkg-attention-unified/v2" class AttentionMaskType(IntEnum): diff --git a/dnn-providers/hip-kernel-provider/rocke/library/dispatch/attention/strided_decode.py b/dnn-providers/hip-kernel-provider/rocke/library/dispatch/attention/strided_decode.py index d9da74e85bb1..e31f92739476 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/dispatch/attention/strided_decode.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/dispatch/attention/strided_decode.py @@ -148,7 +148,8 @@ def launch(*, softmax_scale=softmax_scale, stream=stream): family=FAMILY, algorithm="strided_decode", spec_id="strided_decode", - abi_version="rocke-attention-strided-kv/v1", + # v2: the head counts and the sliding window became kernargs. + abi_version="rocke-attention-strided-kv/v2", priority=20, capability=Capability( arches=("gfx942", "gfx950"), 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 043c682701c0..7ddef488cf9f 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 @@ -1714,6 +1714,23 @@ def _tiled_cache_key(problem: UnifiedAttentionProblem, arch: str) -> Tuple: num_warps_2d = None if fp16_flash else _select_2d_num_warps(problem, arch) tile_size = _select_2d_tile_size(problem, arch) block_m_per_warp = _select_2d_block_m_per_warp(problem, arch) + # gfx950 schedule tail of `_spec_gfx950_generic`. Both flags are false + # outside the single-batch cohort (num_seqs == 1), so skip them there. Keyed + # as None when both are off, because the cohort itself does not change the + # kernel and a runtime batch search must not split ns=1 from ns>1 on it. + single_batch_sched = None + if _enable_single_batch_combo(problem, arch): + flags = ( + _enable_softmax_mfma_interleave(problem, arch), + _enable_k_single_buffer( + problem, + arch, + num_warps=num_warps_2d, + block_m_per_warp=block_m_per_warp, + tile_size=tile_size, + ), + ) + single_batch_sched = flags if any(flags) else None return ( "tiled", arch, @@ -1758,24 +1775,12 @@ def _tiled_cache_key(problem: UnifiedAttentionProblem, arch: str) -> Tuple: # 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, arch), - # gfx950 schedule tail of `_spec_gfx950_generic`. Both are false - # outside the single-batch cohort (num_seqs == 1), so gate on it and - # spare every other launch the selector walks. sched_barrier is implied - # by _enable_v_double_buffer, which is already in the key above. - ( - ( - _enable_softmax_mfma_interleave(problem, arch), - _enable_k_single_buffer( - problem, - arch, - num_warps=num_warps_2d, - block_m_per_warp=block_m_per_warp, - tile_size=tile_size, - ), - ) - if _enable_single_batch_combo(problem, arch) - else None - ), + # Bakes the 64/8 head geometry, which the heads entry above does not + # carry when the counts are kernargs. + _enable_fast_paged_kv_desc(problem, arch), + # sched_barrier is implied by _enable_v_double_buffer, which is already + # in the key above. + single_batch_sched, # 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. @@ -2899,6 +2904,25 @@ def _enable_transposed_subflags(problem: UnifiedAttentionProblem, arch: str) -> return _enable_transposed_qk_32x32(problem, arch) +def _enable_fast_paged_kv_desc(problem: UnifiedAttentionProblem, arch: str) -> bool: + """The fast multi-block paged-KV descriptor of the no-window bf16 2D combo. + + Its loader bakes the h64kv8 bf16 T=64 geometry (the spec validator pins it), + so it is on only for exactly that head pair. The head counts are kernargs + and the cache key carries only their ratio, so the key carries this flag + too: a 32/4 problem at the same ratio must not reuse a 64/8 fast binary. + """ + return ( + _enable_combo_2d(problem, arch) + and problem.sliding_window == 0 + and problem.dtype == "bf16" + and not problem.use_fp8 + and problem.num_query_heads == 64 + and problem.num_kv_heads == 8 + and _select_2d_tile_size(problem, arch) == 64 + ) + + # Largest paged KV cache, in bytes, that i32 addressing reads correctly on the # tiled builders. They bound K/V buffer loads at num_records=0x7FFF0000, so a # load at or past that offset returns zero. diff --git a/dnn-providers/hip-kernel-provider/rocke/library/kernels/gfx950/attention_tiled_2d.py b/dnn-providers/hip-kernel-provider/rocke/library/kernels/gfx950/attention_tiled_2d.py index 2fb7454b2555..8e7106307e08 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/kernels/gfx950/attention_tiled_2d.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/kernels/gfx950/attention_tiled_2d.py @@ -40,7 +40,7 @@ from __future__ import annotations from dataclasses import dataclass -from typing import Optional, Tuple +from typing import ClassVar, Optional, Tuple from kernels.common._attention_additive_bias import ( add_additive_bias, @@ -65,7 +65,7 @@ from rocke.helpers.attention import ( apply_softcap_log2, binary_search_seq_idx, - binary_search_seq_idx_iters, + binary_search_seq_idx_trip_count, dequant_fp8x8_to_dtype, mfma_16x16x16_for_dtype, mfma_16x16x32_for_dtype, @@ -78,7 +78,13 @@ ) from rocke.helpers.distribution import make_static_tile_distribution from rocke.helpers.layouts import TransposeLdsReader -from rocke.helpers.transforms import TensorDescriptor, embed, indirect, unmerge +from rocke.helpers.transforms import ( + DynamicTensorDescriptor, + TensorDescriptor, + embed, + indirect, + unmerge, +) MFMA_M = 16 @@ -154,6 +160,18 @@ def _mfma_32x32_c_col(b, lane, n_tile32: int = 0): @dataclass(frozen=True) class UnifiedAttention2DTiledSpec: + # Problem fields this body reads as trailing i32 kernargs instead of baking + # them, in ABI order. The launch signature, the argument pack and the cache + # key read the same tuple (``_runtime_kernarg_fields``). + RUNTIME_KERNARG_FIELDS: ClassVar[Tuple[str, ...]] = ( + "num_query_heads", + "num_kv_heads", + "sliding_window", + ) + # The cu_q search derives its trip count from the num_seqs kernarg, so + # num_seqs never reaches the IR and one binary serves every batch. + RUNTIME_SEARCH_ITERS: ClassVar[bool] = True + head_size: int block_size: int num_query_heads: int @@ -884,15 +902,6 @@ def n_blocks_per_tile(self) -> int: def dtype_ir(self) -> Type: return F16 if self.dtype == "fp16" else BF16 - @property - def binary_search_iters(self) -> int: - # AITER/Triton uses a true while-loop binary search. Our IR currently - # lowers this as a fixed-trip scf.for, so specialize the trip count to - # the known problem batch size instead of always paying 32 iterations. - # Falls back to 32 for direct unit-test specs that do not provide - # `num_seqs`. - return binary_search_seq_idx_iters(self.num_seqs) - def kernel_name(self) -> str: from rocke.helpers.spec import kernel_name_join @@ -905,11 +914,21 @@ def kernel_name(self) -> str: f"d{self.head_size}", f"b{self.block_size}", f"t{self.tile_size_eff}" if self.n_blocks_per_tile != 1 else "", - f"h{self.num_query_heads}kv{self.num_kv_heads}", + # Head counts are kernargs; only the GQA ratio is baked. + f"g{self.num_queries_per_kv}", + # Except on the fast paged-KV descriptor path, whose shifts bake them. + ( + f"h{self.num_query_heads}kv{self.num_kv_heads}" + if self.use_fast_paged_kv_desc + else "" + ), self.dtype, f"kv{self.kv_storage_dtype}" if self.kv_storage_dtype else "", + # i64 K/V addressing changes the loaders, so it is a distinct kernel. + "i64" if self.use_i64_kv_addr else "", "" if not self.use_sinks else "sinks", - f"sw{self.sliding_window}" if self.sliding_window > 0 else "", + # The window size is a kernarg; only whether the body masks one is baked. + "sw" if self.sliding_window > 0 else "", "softcap" if self.has_softcap else "", "alibi" if self.use_alibi else "", "qqb" if self.use_qq_bias else "", @@ -1144,7 +1163,6 @@ def build_unified_attention_2d_tiled( NQK = spec.num_queries_per_kv VALID_ROWS = BLOCK_Q * NQK NUM_KV = spec.num_kv_heads - NUM_QH = spec.num_query_heads SLIDING_WINDOW = spec.sliding_window USE_SOFTCAP = spec.has_softcap USE_SINKS = spec.use_sinks @@ -1296,6 +1314,11 @@ def build_unified_attention_2d_tiled( num_seqs_p = b.param("num_seqs", I32) bt_stride_p = b.param("block_table_stride", I32) qq_bias_stride0_p = b.param("qq_bias_stride_0", I32) + # Shape values this body reads as kernargs (``RUNTIME_KERNARG_FIELDS``), + # after every fixed param so the shared launch signature can append them. + rt = {f: b.param(f, I32) for f in spec.RUNTIME_KERNARG_FIELDS} + num_qh_p = rt["num_query_heads"] + num_kv_p = rt["num_kv_heads"] # Additive-bias params are appended only when the feature is on, so the # non-bias kernel ABI (and its IR) is unchanged. additive_bias = declare_additive_bias_params(b) if USE_ADDITIVE_BIAS else None @@ -1322,7 +1345,7 @@ def build_unified_attention_2d_tiled( q_block_global_idx, num_seqs_p, block_q=BLOCK_Q, - iterations=spec.binary_search_iters, + iterations=binary_search_seq_idx_trip_count(b, num_seqs_p), ) cu_q_start = b.global_load_i32(cu_q, seq_idx) cu_q_stop = b.global_load_i32(cu_q, b.add(seq_idx, b.const_i32(1))) @@ -1660,7 +1683,7 @@ def _band_row_mask(row_ok, upper_ok, in_prefix): # stays the same as the default bf16 path -- just softmax_scale / # ln(2) -- and we do NOT fold k_scale into qk_scale here. pv_fp8_scale = b.fdiv(v_scale_p, b.const_f32(240.0)) if FP8_MFMA_PV else None - sw_const = b.const_i32(int(SLIDING_WINDOW)) + sw_const = rt["sliding_window"] z8 = b.zero_vec(dtype, 8) # ---------------- Q -> LDS (cooperative vec8 chunks) ---------------- @@ -1673,15 +1696,11 @@ def _band_row_mask(row_ok, upper_ok, in_prefix): # Coordinate transform for Q (and the symmetric output buffer): # ``(token, head, dim)`` packed contiguously. The element-unit # descriptor is reused below for the output store too. - q_desc = TensorDescriptor.naive( + # The head count is a kernarg, so the token stride is too. + q_desc = DynamicTensorDescriptor.create( "Q", - # The runtime extents (total_q, num_query_heads, head_size) are - # only used by the validity predicate (which we don't request - # here). Use generous compile-time bounds so the descriptor's - # row-major stride product matches the kernel's layout - # assumptions exactly. - lengths=[1 << 30, NUM_QH, HD], coord_names=("token", "head", "dim"), + strides=(b.mul(num_qh_p, b.const_i32(HD)), HD, 1), ) # direct-reg: skip the cooperative Q->LDS staging store entirely. # The per-lane Q32 MFMA operand is gathered straight from global in the @@ -1696,7 +1715,7 @@ def _band_row_mask(row_ok, upper_ok, in_prefix): ) _qmask_inner = b.land( b.cmp_lt(q_pos_t, cur_batch_q_len), - b.cmp_lt(qh_t, b.const_i32(NUM_QH)), + b.cmp_lt(qh_t, num_qh_p), ) qmask_t = ( b.land(_qmask_inner, b.cmp_lt(Q_row, b.const_i32(VALID_ROWS))) @@ -1886,7 +1905,7 @@ def _pack_p_a32(p_regs0_f32: list[Value], p_regs1_f32: list[Value]) -> Value: qh = b.add( b.mul(kv_head_idx, b.const_i32(NQK)), b.mod(row, b.const_i32(NQK)) ) - qh_in = b.cmp_lt(qh, b.const_i32(NUM_QH)) + qh_in = b.cmp_lt(qh, num_qh_p) sink_h = b.global_load(sinks, qh, dtype, align=2) sink_f = b.fmul(b.cast_to_f32(sink_h), rcp_ln2) m_inits.append(b.select(qh_in, sink_f, neg_inf)) @@ -1896,7 +1915,7 @@ def _pack_p_a32(p_regs0_f32: list[Value], p_regs1_f32: list[Value]) -> Value: qh = b.add( b.mul(kv_head_idx, b.const_i32(NQK)), b.mod(row, b.const_i32(NQK)) ) - qh_in = b.cmp_lt(qh, b.const_i32(NUM_QH)) + qh_in = b.cmp_lt(qh, num_qh_p) sink_h = b.global_load(sinks, qh, dtype, align=2) sink_f = b.fmul(b.cast_to_f32(sink_h), rcp_ln2) m_inits.append(b.select(qh_in, sink_f, neg_inf)) @@ -1982,10 +2001,11 @@ def _acc_idx(n: int, atom: int) -> int: # so for the FP8 K/V path the loader switches to a sync per-thread # dequant chain below; the byte-stride math here is shared so the # paged_kv_desc compiles consistently in both branches. - kv_stride_blk_b = BS * NUM_KV * HD * KV_BYTES - kv_stride_tok_b = NUM_KV * HD * KV_BYTES + # The token and block strides scale with the runtime kv head count. + kv_stride_tok_b = b.mul(num_kv_p, b.const_i32(HD * KV_BYTES)) + kv_stride_blk_b = b.mul(kv_stride_tok_b, b.const_i32(BS)) kv_stride_h_b = HD * KV_BYTES - kv_block_bytes_c = b.const_i32(kv_stride_blk_b) # one-block buffer bound + kv_block_bytes_c = kv_stride_blk_b # one-block buffer bound lane_half_base = b.mul(tid, b.const_i32(8)) @@ -2086,6 +2106,23 @@ def _acc_idx(n: int, atom: int) -> int: # return ``0`` (a guaranteed-valid block id). Downstream causal / # in_prefix masks already discard the bogus tokens. block_table_max_idx = b.to_sgpr_u32(b.mul(num_seqs_p, bt_stride_p)) + # A copy covers COPY_TOKENS whole tokens (KV_HALVES_PER_CALL is a multiple + # of HD). When they fit in one page the bf16/fp16 loaders build the page + # base, the in-page token offset and the lane part separately (see + # ``_paged_copy``); otherwise each wave's share still sits in one page and + # the loaders keep the paged-KV descriptor. + COPY_TOKENS = KV_HALVES_PER_CALL // HD + PAGE_COPIES = not FAST_PAGED_KV_DESC and not KV_FP8 and BS % COPY_TOKENS == 0 + # With page copies, the loop carries the page ids of the current and the + # next tile and reads the ids of the tile after next one iteration ahead, + # so no copy waits on a table read issued just before it. The schedules + # that load two tiles ahead (the K ring, grouped KV2) read in place, and so + # do tiles of more than 16 pages, which the C++ twin's carry arrays bound. + PAGE_PREFETCH = ( + PAGE_COPIES and not RING3 and not GROUPED_KV2 and N_BLOCKS_PER_TILE <= 16 + ) + # Page ids already read for a tile, by the tile index value. + tile_pages = {} if FAST_PAGED_KV_DESC: assert (BS, T, HD, NUM_KV, KV_BYTES, NUM_WARPS, kv_calls_per_tile) == ( 32, @@ -2139,13 +2176,10 @@ def _fast_paged_kv_voff(call: int, block0: Value, block1: Value): return None, b.add(block_b, within) else: - _kv_base = TensorDescriptor.naive( + _kv_base = DynamicTensorDescriptor.create( "paged_kv_bytes", - # ``lengths`` here is just informational (validity propagation - # is driven by the transforms above, not by these bounds). - lengths=[1 << 24, BS, NUM_KV, HD], - strides=[kv_stride_blk_b, kv_stride_tok_b, kv_stride_h_b, KV_BYTES], coord_names=("physical_block", "token", "kv_head", "dim"), + strides=(kv_stride_blk_b, kv_stride_tok_b, kv_stride_h_b, KV_BYTES), ) if N_BLOCKS_PER_TILE == 1: paged_kv_desc = _kv_base.transform( @@ -2178,6 +2212,77 @@ def _fast_paged_kv_voff(call: int, block0: Value, block1: Value): max_idx=block_table_max_idx, ), ) + if PAGE_COPIES: + # When a copy's tokens sit in one page, the bf16/fp16 loaders address + # it as its page's base, its first token's offset in the page and + # the lane's own part. The token stride is a kernarg, so build the + # lane part and the per-copy token offsets once here rather than + # multiplying the token index in every copy of every tile. + lane_tok = b.div(lane_half_base, b.const_i32(HD)) + lane_dim = b.mod(lane_half_base, b.const_i32(HD)) + lane_within_b = b.add( + b.add( + b.mul(lane_tok, kv_stride_tok_b), + b.mul(kv_head_idx, b.const_i32(kv_stride_h_b)), + ), + b.mul(lane_dim, b.const_i32(KV_BYTES)), + ) + copy_tok_b = {} + for call in range(kv_calls_per_tile): + tok0 = (call * COPY_TOKENS) % BS + if tok0 and tok0 not in copy_tok_b: + copy_tok_b[tok0] = b.mul(kv_stride_tok_b, b.const_i32(tok0)) + if I64_KV_ADDR: + kv_stride_blk_i64 = b.zext(kv_stride_blk_b, I64) + + def _tile_page_id(kv_tile_idx: Value, blk: int) -> Value: + """The paged-KV descriptor's indirect lookup for page ``blk`` of a tile: + one guarded block-table read, wave-uniform.""" + idx = b.add( + seq_base, + ( + kv_tile_idx + if N_BLOCKS_PER_TILE == 1 + else b.add( + b.mul(kv_tile_idx, b.const_i32(N_BLOCKS_PER_TILE)), + b.const_i32(blk), + ) + ), + ) + page = b.masked_global_load( + block_tables, + idx, + b.cmp_lt(idx, block_table_max_idx), + b.const_i32(0), + dtype=I32, + align=4, + ) + return b.to_sgpr_u32(page) + + def _tile_page_ids(kv_tile_idx: Value) -> list: + return [_tile_page_id(kv_tile_idx, j) for j in range(N_BLOCKS_PER_TILE)] + + if PAGE_COPIES: + + def _paged_copy(src: Value, src_rsrc: Value, kv_tile_idx: Value, call, pages): + """(rsrc, voffset, soffset) of one bf16/fp16 K/V copy. ``pages`` + caches this tile's page ids by block within the tile when the + loop does not carry them.""" + blk = (call * COPY_TOKENS) // BS + tok0 = (call * COPY_TOKENS) % BS + carried = tile_pages.get(id(kv_tile_idx)) + if carried is not None: + pages = carried + elif blk not in pages: + pages[blk] = _tile_page_id(kv_tile_idx, blk) + voff = lane_within_b + if tok0: + voff = b.add(voff, copy_tok_b[tok0]) + if I64_KV_ADDR: + base_i64 = b.mul(b.zext(pages[blk], I64), kv_stride_blk_i64) + rsrc = b.buffer_rsrc(b.global_ptr_add(src, base_i64), kv_block_bytes_c) + return rsrc, voff, zero_soff + return src_rsrc, voff, b.mul(pages[blk], kv_stride_blk_b) def _issue_k_load_runtime(kv_tile_idx: Value, buf_idx: Value) -> None: """Issue async K loads for one tile into K_lds[buf_idx]. @@ -2200,14 +2305,20 @@ def _issue_k_load_runtime(kv_tile_idx: Value, buf_idx: Value) -> None: K_wave_base = b.smem_ptr_add(K_buf_base, wave_lds_offset_K_i64) if FAST_PAGED_KV_DESC: fast_block0, fast_block1 = _fast_paged_kv_blocks(kv_tile_idx) + pages = {} for call in range(kv_calls_per_tile): k_rsrc = key_rsrc + soff = zero_soff if FAST_PAGED_KV_DESC: base_i64, voff = _fast_paged_kv_voff(call, fast_block0, fast_block1) if base_i64 is not None: # I64_KV_ADDR: per-block 64-bit base k_rsrc = b.buffer_rsrc( b.global_ptr_add(key, base_i64), kv_block_bytes_c ) + elif PAGE_COPIES: + k_rsrc, voff, soff = _paged_copy( + key, key_rsrc, kv_tile_idx, call, pages + ) else: linear_half = b.add( b.const_i32(call * KV_HALVES_PER_CALL), lane_half_base @@ -2236,7 +2347,7 @@ def _issue_k_load_runtime(kv_tile_idx: Value, buf_idx: Value) -> None: # ``dsl_docs/primitives/intrinsics_and_primitives.md`` as the # right hint for K-loop streaming tile loads. b.async_buffer_load_lds_addr( - k_rsrc, k_dst, voff, zero_soff, 4, coherency=CACHE_STREAM + k_rsrc, k_dst, voff, soff, 4, coherency=CACHE_STREAM ) def _issue_v_load_runtime(kv_tile_idx: Value, buf_idx: Value) -> None: @@ -2258,14 +2369,20 @@ def _issue_v_load_runtime(kv_tile_idx: Value, buf_idx: Value) -> None: V_wave_base = b.smem_ptr_add(V_lds_addr, wave_lds_offset_i64) if FAST_PAGED_KV_DESC: fast_block0, fast_block1 = _fast_paged_kv_blocks(kv_tile_idx) + pages = {} for call in range(kv_calls_per_tile): v_rsrc = value_rsrc + soff = zero_soff if FAST_PAGED_KV_DESC: base_i64, voff = _fast_paged_kv_voff(call, fast_block0, fast_block1) if base_i64 is not None: v_rsrc = b.buffer_rsrc( b.global_ptr_add(value, base_i64), kv_block_bytes_c ) + elif PAGE_COPIES: + v_rsrc, voff, soff = _paged_copy( + value, value_rsrc, kv_tile_idx, call, pages + ) else: linear_half = b.add( b.const_i32(call * KV_HALVES_PER_CALL), lane_half_base @@ -2293,7 +2410,7 @@ def _issue_v_load_runtime(kv_tile_idx: Value, buf_idx: Value) -> None: # re-read within this kernel; see _issue_k_load_runtime for # the rationale. b.async_buffer_load_lds_addr( - v_rsrc, v_dst, voff, zero_soff, 4, coherency=CACHE_STREAM + v_rsrc, v_dst, voff, soff, 4, coherency=CACHE_STREAM ) # ---------------- FP8 K/V cache: async DMA loader (round 2) ---------------- @@ -2929,7 +3046,7 @@ def _read_q32(k: int) -> Value: # so the <8 x dtype> operand is bit-identical to the staged path. # row -> q_pos = qb_start_pos + row // NQK ; qh = kv_head*NQK + row%NQK # col -> dim - # Out-of-range (token >= q_len or head >= NUM_QH) lanes read 0 (masked), + # Out-of-range (token >= q_len or head >= num_query_heads) lanes read 0 (masked), # matching the staging store's zero-fill. def _read_q32_global(k: int) -> Value: q32_col = b.add(b.const_i32(k * 16), b.mul(lane_half, b.const_i32(8))) @@ -2940,7 +3057,7 @@ def _read_q32_global(k: int) -> Value: ) qg_mask = b.land( b.cmp_lt(qg_pos, cur_batch_q_len), - b.cmp_lt(qg_h, b.const_i32(NUM_QH)), + b.cmp_lt(qg_h, num_qh_p), ) qg_pos_safe = b.select(qg_mask, qg_pos, b.const_i32(0)) qg_h_safe = b.select(qg_mask, qg_h, b.const_i32(0)) @@ -3002,6 +3119,16 @@ def _read_q32_global(k: int) -> Value: # the next thing scheduled, so this ordering is naturally # correct in the IR. + if PAGE_PREFETCH: + # Page ids of the first tile and of the one after it, carried into the + # loop; iteration i then reads the ids of tile i + 2. + start_next_raw = b.add(tile_start, b.const_i32(1)) + start_next = b.select( + b.cmp_lt(start_next_raw, tile_end), start_next_raw, tile_start + ) + start_pages = _tile_page_ids(tile_start) + start_next_pages = _tile_page_ids(start_next) + tile_pages[id(tile_start)] = start_pages # Prefetch tile_start's K into buffer 0 BEFORE the loop. _issue_k(tile_start, b.const_i32(0)) if RING3: @@ -3037,12 +3164,17 @@ def _read_q32_global(k: int) -> Value: # 6. Yield ``(m, l, acc, nxt_buf)`` so the next iter consumes nxt_buf. cur_buf_init = b.const_i32(0) iter_args.append(("cur_buf", cur_buf_init)) + if PAGE_PREFETCH: + for j in range(N_BLOCKS_PER_TILE): + iter_args.append((f"page_cur{j}", start_pages[j])) + for j in range(N_BLOCKS_PER_TILE): + iter_args.append((f"page_next{j}", start_next_pages[j])) # ---------------- LICM hoist: per-reg invariants ---------------- # ``qp_r``, ``qh_r``, ``row_ok``, ``causal_lim``, and ``alibi_per_row`` # depend only on CTA-level constants (``wave_row_base``, ``lane_rg``, # ``qb_start_pos``, ``kv_head_idx``, ``cur_batch_q_len``, ``context_len``, - # ``NUM_QH``, ``NQK``). Computing them BEFORE the kvloop avoids paying + # ``num_qh_p``, ``NQK``). Computing them BEFORE the kvloop avoids paying # them per kv tile. LLVM LICM should hoist them automatically, but # explicit hoisting makes the IR's hot path leaner and eliminates a # source of compiler-scheduling variability. The per-reg arrays @@ -3062,7 +3194,7 @@ def _read_q32_global(k: int) -> Value: qh_r = b.add(b.mul(kv_head_idx, b.const_i32(NQK)), b.mod(row, b.const_i32(NQK))) _row_ok_inner = b.land( b.cmp_lt(qp_r, cur_batch_q_len), - b.cmp_lt(qh_r, b.const_i32(NUM_QH)), + b.cmp_lt(qh_r, num_qh_p), ) row_ok = ( b.land(_row_ok_inner, b.cmp_lt(row, b.const_i32(VALID_ROWS))) @@ -3080,7 +3212,7 @@ def _read_q32_global(k: int) -> Value: hoist_alibi = [] for reg in range(REGS_PER_LANE): qh_r = hoist_qh_r[reg] - qh_ok = b.cmp_lt(qh_r, b.const_i32(NUM_QH)) + qh_ok = b.cmp_lt(qh_r, num_qh_p) slope = b.masked_global_load( alibi_slopes_ptr, qh_r, qh_ok, b.const_f32(0.0), dtype=F32, align=4 ) @@ -3096,7 +3228,7 @@ def _read_q32_global(k: int) -> Value: ) _st_row_ok_inner = b.land( b.cmp_lt(st_qp, cur_batch_q_len), - b.cmp_lt(st_qh, b.const_i32(NUM_QH)), + b.cmp_lt(st_qh, num_qh_p), ) st_row_ok = ( b.land(_st_row_ok_inner, b.cmp_lt(st_q_row, b.const_i32(VALID_ROWS))) @@ -3123,7 +3255,7 @@ def _read_q32_global(k: int) -> Value: and TRANSPOSED_QK_32X32 and TRANSPOSED_INVARIANT_HOIST ): - st_qh_ok = b.cmp_lt(st_qh, b.const_i32(NUM_QH)) + st_qh_ok = b.cmp_lt(st_qh, num_qh_p) st_alibi_slope = b.masked_global_load( alibi_slopes_ptr, st_qh, st_qh_ok, b.const_f32(0.0), dtype=F32, align=4 ) @@ -3157,6 +3289,13 @@ def _acc_get(n: int, atom: int) -> Value: return acc_vals[n * ACC_M_ATOMS + atom] cur_buf = carry[ml_count + ACC_N_TILES * ACC_M_ATOMS] + if PAGE_PREFETCH: + pages_at = ml_count + ACC_N_TILES * ACC_M_ATOMS + 1 + cur_pages = list(carry[pages_at : pages_at + N_BLOCKS_PER_TILE]) + next_pages = list( + carry[pages_at + N_BLOCKS_PER_TILE : pages_at + 2 * N_BLOCKS_PER_TILE] + ) + tile_pages[id(kv_tile_iv)] = cur_pages if RING3: # 3-slot ring: cur_buf holds K[i]; (cur_buf+1)%3 holds K[i+1] # (prefetched last iter / prologue); the prefetch target for K[i+2] @@ -3202,7 +3341,7 @@ def _mod3_add(x, c): ) _st_row_ok_iter_inner = b.land( b.cmp_lt(st_qp_iter, cur_batch_q_len), - b.cmp_lt(st_qh_iter, b.const_i32(NUM_QH)), + b.cmp_lt(st_qh_iter, num_qh_p), ) st_row_ok_iter = ( b.land( @@ -3214,7 +3353,7 @@ def _mod3_add(x, c): ) st_causal_lim_iter = b.add(context_len, st_qp_iter) if USE_ALIBI: - st_qh_iter_ok = b.cmp_lt(st_qh_iter, b.const_i32(NUM_QH)) + st_qh_iter_ok = b.cmp_lt(st_qh_iter, num_qh_p) st_alibi_slope_iter = b.masked_global_load( alibi_slopes_ptr, st_qh_iter, @@ -3240,6 +3379,8 @@ def _mod3_add(x, c): next_tile_iv_raw = b.add(kv_tile_iv, b.const_i32(1)) in_range_next = b.cmp_lt(next_tile_iv_raw, tile_end) safe_next_tile = b.select(in_range_next, next_tile_iv_raw, kv_tile_iv) + if PAGE_PREFETCH: + tile_pages[id(safe_next_tile)] = next_pages if RING3: # Ring depth 3 prefetches K[i+2] (two tiles ahead) into ring_pf_buf. ring_pf_tile_raw = b.add(kv_tile_iv, b.const_i32(2)) @@ -3282,6 +3423,15 @@ def _mod3_add(x, c): # GEMM MFMA-interleave, level 1 = attention-style interleave. if USE_SOFTMAX_INTERLEAVE and SOFTMAX_INTERLEAVE_MODE in (0, 1): b.iglp_opt(SOFTMAX_INTERLEAVE_MODE) + if PAGE_PREFETCH: + # Read the page ids of the tile after next now, right after the + # iter-start drain, so they arrive long before next iteration's + # copies need them. That tile is next iteration's safe_next_tile. + after_next_raw = b.add(kv_tile_iv, b.const_i32(2)) + after_next = b.select( + b.cmp_lt(after_next_raw, tile_end), after_next_raw, next_tile_iv_raw + ) + after_next_pages = _tile_page_ids(after_next) if EARLY_V_SCHEDULE: if V_DOUBLE_BUF: # V[i] (cur_buf) was prefetched in the previous iter / prologue @@ -3485,7 +3635,7 @@ def _mod3_add(x, c): slope_for_lane = b.masked_global_load( alibi_slopes_ptr, qh_t_alibi, - b.cmp_lt(qh_t_alibi, b.const_i32(NUM_QH)), + b.cmp_lt(qh_t_alibi, num_qh_p), b.const_f32(0.0), dtype=F32, align=4, @@ -3540,7 +3690,7 @@ def _mod3_add(x, c): ) _row_ok_inner_t = b.land( b.cmp_lt(qp_r, cur_batch_q_len), - b.cmp_lt(qh_r, b.const_i32(NUM_QH)), + b.cmp_lt(qh_r, num_qh_p), ) row_ok = ( b.land( @@ -4436,6 +4586,9 @@ def _apply_transposed_pv_regs(acc32: Value, n: int, p_regs) -> Value: yields.append(ring_next_cur) else: yields.append(cur_buf if GROUPED_KV2 else nxt_buf) + if PAGE_PREFETCH: + yields.extend(next_pages) + yields.extend(after_next_pages) b.scf_yield(*yields) # ---- drive the (one or two) KV-loop phases ---- @@ -4545,7 +4698,7 @@ def _acc_final_get(n: int, atom: int) -> Value: ) _op_mask_t_inner = b.land( b.cmp_lt(op_pos_t, cur_batch_q_len), - b.cmp_lt(op_qh_t, b.const_i32(NUM_QH)), + b.cmp_lt(op_qh_t, num_qh_p), ) op_mask_t = ( b.land(_op_mask_t_inner, b.cmp_lt(q_row_t, b.const_i32(VALID_ROWS))) @@ -4598,7 +4751,7 @@ def _acc_final_get(n: int, atom: int) -> Value: ) _op_mask32_inner = b.land( b.cmp_lt(op_pos32_base, cur_batch_q_len), - b.cmp_lt(op_qh32_base, b.const_i32(NUM_QH)), + b.cmp_lt(op_qh32_base, num_qh_p), ) op_mask32_base = ( b.land(_op_mask32_inner, b.cmp_lt(OUT_ROW_BASE32, b.const_i32(VALID_ROWS))) @@ -4697,7 +4850,7 @@ def _acc_final_get(n: int, atom: int) -> Value: ) _op_mask_inner = b.land( b.cmp_lt(op_pos, cur_batch_q_len), - b.cmp_lt(op_qh, b.const_i32(NUM_QH)), + b.cmp_lt(op_qh, num_qh_p), ) op_mask = ( b.land(_op_mask_inner, b.cmp_lt(OUT_ROW_BASE, b.const_i32(VALID_ROWS))) diff --git a/dnn-providers/hip-kernel-provider/rocke/library/kernels/gfx950/attention_tiled_2d_fastkv_regp.py b/dnn-providers/hip-kernel-provider/rocke/library/kernels/gfx950/attention_tiled_2d_fastkv_regp.py index dfabc19a5e89..0323ea05d48e 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/kernels/gfx950/attention_tiled_2d_fastkv_regp.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/kernels/gfx950/attention_tiled_2d_fastkv_regp.py @@ -141,6 +141,53 @@ def supports_fastkv_register_p_2d( return True, "supported" +# The body uses gfx950-only wide MFMA atoms and transposed LDS reads. +_NOT_GFX950 = "fastKV register-P is a gfx950 kernel; {arch} is not supported" + + +def _fastkv_register_p_spec_error(spec: UnifiedAttention2DTiledSpec) -> str | None: + """Why the fastKV register-P builder refuses ``spec``, or None if it builds.""" + if not spec.use_fast_paged_kv_desc: + return "fastKV register-P experiment requires use_fast_paged_kv_desc" + if not (spec.use_mfma_32x32 and spec.use_transposed_qk_32x32): + return "fastKV register-P experiment requires transposed R4" + if spec.kv_storage_dtype is not None: + return "fastKV register-P experiment does not support FP8 KV cache" + return None + + +def supports_fastkv_register_p_2d_spec( + spec: UnifiedAttention2DTiledSpec, *, arch: str = "gfx950" +) -> tuple[bool, str]: + """Spec-shaped wrapper for ``supports_fastkv_register_p_2d``. + + ``supports_fastkv_register_p_2d`` takes its parameters as keyword-only + arguments, so ``hkp_pack``'s ``_check_support_predicate`` cannot call it with a + spec instance. This wrapper reads them off the spec and delegates, after + refusing what :func:`build_unified_attention_2d_fastkv_register_p` raises on, + so the builder is validatable at pack time. + """ + if arch != "gfx950": + return False, _NOT_GFX950.format(arch=arch) + error = _fastkv_register_p_spec_error(spec) + if error is not None: + return False, error + return supports_fastkv_register_p_2d( + head_size=spec.head_size, + block_size=spec.block_size, + dtype=spec.dtype, + num_queries_per_kv=spec.num_queries_per_kv, + use_alibi=spec.use_alibi, + use_qq_bias=spec.use_qq_bias, + use_fp8=False, + q_dtype=spec.dtype, + num_query_heads=spec.num_query_heads, + num_kv_heads=spec.num_kv_heads, + tile_size=spec.tile_size, + arch=arch, + ) + + def build_unified_attention_2d_fastkv_register_p( spec: UnifiedAttention2DTiledSpec, *, @@ -148,15 +195,14 @@ def build_unified_attention_2d_fastkv_register_p( ) -> KernelDef: """Build the experimental fastKV + register-P 2D attention kernel. - Thin wrapper over :func:`build_unified_attention_2d_tiled`; the ``arch`` - target (gfx950-only, see that builder) is threaded straight through, so - a gfx942 request fails with the same clean structured error. + Thin wrapper over :func:`build_unified_attention_2d_tiled`. That builder + also admits gfx942 for its narrow-atom variant, so this one refuses every + target but gfx950 itself. """ - if not spec.use_fast_paged_kv_desc: - raise ValueError("fastKV register-P experiment requires use_fast_paged_kv_desc") - if not (spec.use_mfma_32x32 and spec.use_transposed_qk_32x32): - raise ValueError("fastKV register-P experiment requires transposed R4") - if spec.kv_storage_dtype is not None: - raise ValueError("fastKV register-P experiment does not support FP8 KV cache") + if arch != "gfx950": + raise ValueError(_NOT_GFX950.format(arch=arch)) + error = _fastkv_register_p_spec_error(spec) + if error is not None: + raise ValueError(error) return build_unified_attention_2d_tiled(_FastKvRegisterPProxy(spec), arch=arch) diff --git a/dnn-providers/hip-kernel-provider/rocke/library/kernels/gfx950/attention_tiled_3d.py b/dnn-providers/hip-kernel-provider/rocke/library/kernels/gfx950/attention_tiled_3d.py index 59ac897a9aad..47d29a0a0a54 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/kernels/gfx950/attention_tiled_3d.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/kernels/gfx950/attention_tiled_3d.py @@ -28,7 +28,7 @@ from __future__ import annotations from dataclasses import dataclass, field -from typing import Optional, Tuple +from typing import ClassVar, Optional, Tuple from kernels.common._attention_additive_bias import ( add_additive_bias, @@ -60,7 +60,7 @@ from rocke.helpers.attention import ( apply_softcap_log2 as _apply_softcap, binary_search_seq_idx as _binary_search_seq_idx_helper, - binary_search_seq_idx_iters, + binary_search_seq_idx_trip_count, dequant_fp8x8_to_dtype, mfma_16x16x16_for_dtype as _mfma_16x16x16, mfma_16x16x32_for_dtype as _mfma_16x16x32, @@ -69,7 +69,7 @@ ) from rocke.helpers.distribution import make_static_tile_distribution from rocke.helpers.layouts import TransposeLdsReader -from rocke.helpers.transforms import TensorDescriptor, indirect, unmerge +from rocke.helpers.transforms import DynamicTensorDescriptor, indirect, unmerge MFMA_M = 16 MFMA_N = 16 @@ -125,6 +125,18 @@ class UnifiedAttention3DTiledSpec: knob exactly as AITER's ``select_3d_config`` derives it. """ + # Problem fields this body reads as trailing i32 kernargs instead of baking + # them, in ABI order. The launch signature, the argument pack and the cache + # key read the same tuple (``_runtime_kernarg_fields``). + RUNTIME_KERNARG_FIELDS: ClassVar[Tuple[str, ...]] = ( + "num_query_heads", + "num_kv_heads", + "sliding_window", + ) + # The cu_q search derives its trip count from the num_seqs kernarg, so + # num_seqs never reaches the IR and one binary serves every batch. + RUNTIME_SEARCH_ITERS: ClassVar[bool] = True + head_size: int block_size: int num_query_heads: int @@ -298,10 +310,6 @@ def n_blocks_per_tile(self) -> int: def dtype_ir(self) -> Type: return F16 if self.dtype == "fp16" else BF16 - @property - def binary_search_iters(self) -> int: - return binary_search_seq_idx_iters(self.num_seqs) - def kernel_name(self) -> str: from rocke.helpers.spec import kernel_name_join @@ -310,13 +318,15 @@ def kernel_name(self) -> str: f"d{self.head_size}", f"b{self.block_size}", f"t{self.tile_size}" if self.n_blocks_per_tile != 1 else "", - f"h{self.num_query_heads}kv{self.num_kv_heads}", + # Head counts are kernargs; only the GQA ratio is baked. + f"g{self.num_queries_per_kv}", f"seg{self.num_segments}", self.dtype, f"kv{self.kv_storage_dtype}" if self.kv_storage_dtype else "", "i64kv" if self.use_i64_kv_addr else "", "sinks" if self.use_sinks else "", - f"sw{self.sliding_window}" if self.sliding_window > 0 else "", + # The window size is a kernarg; only whether the body masks one is baked. + "sw" if self.sliding_window > 0 else "", "softcap" if self.has_softcap else "", "alibi" if self.use_alibi else "", "qqb" if self.use_qq_bias else "", @@ -428,8 +438,6 @@ def build_unified_attention_3d_tiled( BLOCK_M = spec.block_m BLOCK_Q = spec.block_q NQK = spec.num_queries_per_kv - NUM_KV = spec.num_kv_heads - NUM_QH = spec.num_query_heads NUM_SEG = spec.num_segments SLIDING_WINDOW = spec.sliding_window USE_SOFTCAP = spec.has_softcap @@ -511,6 +519,11 @@ def build_unified_attention_3d_tiled( num_seqs_p = b.param("num_seqs", I32) bt_stride_p = b.param("block_table_stride", I32) qq_bias_stride0_p = b.param("qq_bias_stride_0", I32) + # Shape values this body reads as kernargs (``RUNTIME_KERNARG_FIELDS``), + # after every fixed param so the shared launch signature can append them. + rt = {f: b.param(f, I32) for f in spec.RUNTIME_KERNARG_FIELDS} + num_qh_p = rt["num_query_heads"] + num_kv_p = rt["num_kv_heads"] # Additive-bias params are appended only when the feature is on, so the # non-bias kernel ABI (and its IR) is unchanged. additive_bias = declare_additive_bias_params(b) if USE_ADDITIVE_BIAS else None @@ -533,7 +546,7 @@ def build_unified_attention_3d_tiled( q_block_global_idx, num_seqs_p, block_q=BLOCK_Q, - iterations=spec.binary_search_iters, + iterations=binary_search_seq_idx_trip_count(b, num_seqs_p), ) cu_q_start = b.global_load_i32(cu_q, seq_idx) cu_q_stop = b.global_load_i32(cu_q, b.add(seq_idx, b.const_i32(1))) @@ -578,23 +591,22 @@ def _band_row_mask(row_ok, upper_ok, in_prefix): # segm_max [total_q, num_qh, num_segments] # segm_expsum [total_q, num_qh, num_segments] # segm_output [total_q, num_qh, num_segments, head_size] - # plus Q / output which both share ``(token, head, dim)``. ``1<<30`` - # is a compile-time upper bound on ``total_q`` so the row-major - # stride product matches the kernel's layout assumptions exactly. - ml_desc = TensorDescriptor.naive( + # plus Q / output which both share ``(token, head, dim)``. The head count + # is a kernarg, so each token stride is too. + ml_desc = DynamicTensorDescriptor.create( "segm_ml", - lengths=[1 << 30, NUM_QH, NUM_SEG], coord_names=("token", "head", "seg"), + strides=(b.mul(num_qh_p, b.const_i32(NUM_SEG)), NUM_SEG, 1), ) - seg_acc_desc = TensorDescriptor.naive( + seg_acc_desc = DynamicTensorDescriptor.create( "segm_output", - lengths=[1 << 30, NUM_QH, NUM_SEG, HD], coord_names=("token", "head", "seg", "dim"), + strides=(b.mul(num_qh_p, b.const_i32(NUM_SEG * HD)), NUM_SEG * HD, HD, 1), ) - q_desc = TensorDescriptor.naive( + q_desc = DynamicTensorDescriptor.create( "Q", - lengths=[1 << 30, NUM_QH, HD], coord_names=("token", "head", "dim"), + strides=(b.mul(num_qh_p, b.const_i32(HD)), HD, 1), ) seg_start_tile_pos = b.mul(b.mul(seg_idx, tps), b.const_i32(T)) @@ -608,9 +620,7 @@ def _band_row_mask(row_ok, upper_ok, in_prefix): qh_r = b.add( b.mul(kv_head_idx, b.const_i32(NQK)), b.mod(row, b.const_i32(NQK)) ) - row_ok = b.land( - b.cmp_lt(qp_r, cur_batch_q_len), b.cmp_lt(qh_r, b.const_i32(NUM_QH)) - ) + row_ok = b.land(b.cmp_lt(qp_r, cur_batch_q_len), b.cmp_lt(qh_r, num_qh_p)) qp_r_safe = b.select(row_ok, qp_r, b.const_i32(0)) qh_r_safe = b.select(row_ok, qh_r, b.const_i32(0)) qtoken = b.add(cu_q_start, qp_r_safe) @@ -630,7 +640,7 @@ def _band_row_mask(row_ok, upper_ok, in_prefix): b.mul(kv_head_idx, b.const_i32(NQK)), b.mod(row, b.const_i32(NQK)) ) row_ok = b.land( - b.cmp_lt(qp_r, cur_batch_q_len), b.cmp_lt(qh_r, b.const_i32(NUM_QH)) + b.cmp_lt(qp_r, cur_batch_q_len), b.cmp_lt(qh_r, num_qh_p) ) qp_r_safe = b.select(row_ok, qp_r, b.const_i32(0)) qh_r_safe = b.select(row_ok, qh_r, b.const_i32(0)) @@ -672,7 +682,7 @@ def _row_reduce(x: Value, op) -> Value: one_f = b.const_f32(1.0) rcp_ln2 = b.const_f32(1.4426950408889634) qk_scale = b.fmul(scale_p, rcp_ln2) - sw_const = b.const_i32(int(SLIDING_WINDOW)) + sw_const = rt["sliding_window"] z8 = b.zero_vec(dtype, 8) # ---------------- Q -> LDS ---------------- @@ -687,9 +697,7 @@ def _row_reduce(x: Value, op) -> Value: b.mul(kv_head_idx, b.const_i32(NQK)), b.mod(Q_row, b.const_i32(NQK)), ) - qmask_t = b.land( - b.cmp_lt(q_pos_t, cur_batch_q_len), b.cmp_lt(qh_t, b.const_i32(NUM_QH)) - ) + qmask_t = b.land(b.cmp_lt(q_pos_t, cur_batch_q_len), b.cmp_lt(qh_t, num_qh_p)) q_pos_safe = b.select(qmask_t, q_pos_t, b.const_i32(0)) qh_safe = b.select(qmask_t, qh_t, b.const_i32(0)) q_off_base, _ = q_desc.offset( @@ -743,7 +751,7 @@ def _row_reduce(x: Value, op) -> Value: qh = b.add( b.mul(kv_head_idx, b.const_i32(NQK)), b.mod(lane_col, b.const_i32(NQK)) ) - qh_in = b.cmp_lt(qh, b.const_i32(NUM_QH)) + qh_in = b.cmp_lt(qh, num_qh_p) sink_h = b.global_load(sinks, qh, dtype, align=2) sink_f = b.fmul(b.cast_to_f32(sink_h), rcp_ln2) sink_with_mask = b.select(qh_in, sink_f, neg_inf) @@ -774,14 +782,15 @@ def _row_reduce(x: Value, op) -> Value: # fp8e4m3. The FP8 sync loader path below dequantises and stores to # LDS in the working dtype, so ``bytes_per_buf`` stays at the working # dtype slab size. - kv_stride_blk_b = BS * NUM_KV * HD * KV_BYTES - kv_stride_tok_b = NUM_KV * HD * KV_BYTES + # The token and page strides scale with the runtime kv head count. + kv_stride_tok_b = b.mul(num_kv_p, b.const_i32(HD * KV_BYTES)) + kv_stride_blk_b = b.mul(kv_stride_tok_b, b.const_i32(BS)) kv_stride_h_b = HD * KV_BYTES bytes_per_buf = T * HD * 2 # One-block buffer bound for the i64-addressing path: each per-block # buffer descriptor only spans a single physical KV block, so the # within-block i32 voffset can never overflow. - kv_block_bytes_c = b.const_i32(kv_stride_blk_b) + kv_block_bytes_c = kv_stride_blk_b lane_half_base = b.mul(tid, b.const_i32(8)) K_lds_addr = b.smem_addr_of(K_lds) @@ -806,11 +815,10 @@ def _row_reduce(x: Value, op) -> Value: PAGE_IDS = not KV_FP8 and spec.kv_layout == "paged" WAVE_PAGE_ID_LOAD = spec.use_wave_page_id_load if not PAGE_IDS: - paged_kv_desc = TensorDescriptor.naive( + paged_kv_desc = DynamicTensorDescriptor.create( "paged_kv_bytes", - lengths=[1 << 24, T, NUM_KV, HD], - strides=[kv_stride_blk_b, kv_stride_tok_b, kv_stride_h_b, KV_BYTES], coord_names=("physical_block", "token", "kv_head", "dim"), + strides=(kv_stride_blk_b, kv_stride_tok_b, kv_stride_h_b, KV_BYTES), ).transform( indirect( "tile_idx", into="physical_block", table=block_tables, base=seq_base @@ -832,7 +840,7 @@ def _row_reduce(x: Value, op) -> Value: lane_tok = b.div(lane_half_base, b.const_i32(HD)) lane_dim = b.mod(lane_half_base, b.const_i32(HD)) lane_tok_head = b.add( - b.mul(lane_tok, b.const_i32(kv_stride_tok_b)), + b.mul(lane_tok, kv_stride_tok_b), b.mul(kv_head_idx, b.const_i32(kv_stride_h_b)), ) # Per-lane byte offset within a page, keyed by the call's swizzle @@ -856,6 +864,16 @@ def _row_reduce(x: Value, op) -> Value: lane_within_b[0] = b.add( lane_tok_head, b.mul(lane_dim, b.const_i32(KV_BYTES)) ) + # The token stride is a kernarg, so each call's token offset within the + # page is a multiply. They are wave-uniform and loop-invariant: build + # them once here rather than in every tile's copies. + call_tok_b = {} + for call in range(kv_calls_per_tile): + tok0 = (call % calls_per_block) * tokens_per_call + if tok0 and tok0 not in call_tok_b: + call_tok_b[tok0] = b.mul(kv_stride_tok_b, b.const_i32(tok0)) + if I64_KV_ADDR: + kv_stride_blk_i64 = b.zext(kv_stride_blk_b, I64) def _kv_col(row: Value, col: Value) -> Value: # LDS column of K/V element (row, col): its 16-byte chunk is permuted @@ -906,18 +924,18 @@ def _issue_paged_tile( blk = call // calls_per_block tok0 = (call % calls_per_block) * tokens_per_call ph = (call * tokens_per_call) % KV_CHUNKS_PER_ROW if KV_SWZ else 0 - voff = b.add(lane_within_b[ph], b.const_i32(tok0 * kv_stride_tok_b)) + voff = lane_within_b[ph] + if tok0: + voff = b.add(voff, call_tok_b[tok0]) if I64_KV_ADDR: - base_i64 = b.mul( - b.zext(page_ids[blk], I64), b.const_i64(kv_stride_blk_b) - ) + base_i64 = b.mul(b.zext(page_ids[blk], I64), kv_stride_blk_i64) call_rsrc = b.buffer_rsrc( b.global_ptr_add(src, base_i64), kv_block_bytes_c ) soff = zero_soff else: call_rsrc = src_rsrc - soff = b.mul(page_ids[blk], b.const_i32(kv_stride_blk_b)) + soff = b.mul(page_ids[blk], kv_stride_blk_b) dst = b.smem_ptr_add(buf_base, b.const_i64(call * bytes_per_call)) b.async_buffer_load_lds_addr( call_rsrc, dst, voff, soff, 4, coherency=KV_CACHE_BITS @@ -1118,13 +1136,11 @@ def _refill_page_id_window(page_id_window: Value, kv_tile_idx: Value) -> Value: qh_r = b.add( b.mul(kv_head_idx, b.const_i32(NQK)), b.mod(lane_col, b.const_i32(NQK)) ) - row_ok = b.land( - b.cmp_lt(qp_r, cur_batch_q_len), b.cmp_lt(qh_r, b.const_i32(NUM_QH)) - ) + row_ok = b.land(b.cmp_lt(qp_r, cur_batch_q_len), b.cmp_lt(qh_r, num_qh_p)) causal_lim = b.add(context_len, qp_r) if USE_ALIBI: - qh_ok = b.cmp_lt(qh_r, b.const_i32(NUM_QH)) + qh_ok = b.cmp_lt(qh_r, num_qh_p) alibi_slope = b.masked_global_load( alibi_slopes_ptr, qh_r, qh_ok, b.const_f32(0.0), dtype=F32, align=4 ) @@ -1369,6 +1385,11 @@ def _refill_page_id_window(page_id_window: Value, kv_tile_idx: Value) -> Value: @dataclass(frozen=True) class UnifiedAttentionReduceTiledSpec: + # Problem fields this body reads as trailing i32 kernargs instead of baking + # them, in ABI order. The launch signature, the argument pack and the cache + # key read the same tuple (``_runtime_kernarg_fields``). + RUNTIME_KERNARG_FIELDS: ClassVar[Tuple[str, ...]] = ("num_query_heads",) + head_size: int num_query_heads: int num_kv_heads: int @@ -1401,7 +1422,6 @@ def kernel_name(self) -> str: parts = [ f"d{self.head_size}", - f"h{self.num_query_heads}", f"seg{self.num_segments}", self.dtype, ] @@ -1433,7 +1453,6 @@ def build_unified_attention_reduce_tiled( """ HD = spec.head_size NUM_SEG = spec.num_segments - NUM_QH = spec.num_query_heads dtype = spec.dtype_ir THREADS = 64 @@ -1454,6 +1473,8 @@ def build_unified_attention_reduce_tiled( seg_max = b.param("segm_max_ptr", PtrType(F32, "global"), readonly=True, align=4) seg_l = b.param("segm_expsum_ptr", PtrType(F32, "global"), readonly=True, align=4) _seq_lens = b.param("seq_lens_ptr", PtrType(I32, "global"), readonly=True, align=4) + rt = {f: b.param(f, I32) for f in spec.RUNTIME_KERNARG_FIELDS} + num_qh_p = rt["num_query_heads"] lse = ( b.param( "lse_ptr", @@ -1478,21 +1499,22 @@ def build_unified_attention_reduce_tiled( # segm_output : [total_q, num_qh, num_segments, head_size] # And the final output uses ``(token, head, dim)``. Encoding them # as descriptors makes the per-pass offset math one call instead of - # an `add(mul(qt,...), mul(qh,...))` ladder. - ml_desc_red = TensorDescriptor.naive( + # an `add(mul(qt,...), mul(qh,...))` ladder. The head count is a kernarg, + # so each token stride is too. + ml_desc_red = DynamicTensorDescriptor.create( "segm_ml", - lengths=[1 << 30, NUM_QH, NUM_SEG], coord_names=("token", "head", "seg"), + strides=(b.mul(num_qh_p, b.const_i32(NUM_SEG)), NUM_SEG, 1), ) - seg_acc_desc_red = TensorDescriptor.naive( + seg_acc_desc_red = DynamicTensorDescriptor.create( "segm_output", - lengths=[1 << 30, NUM_QH, NUM_SEG, HD], coord_names=("token", "head", "seg", "dim"), + strides=(b.mul(num_qh_p, b.const_i32(NUM_SEG * HD)), NUM_SEG * HD, HD, 1), ) - out_desc_red = TensorDescriptor.naive( + out_desc_red = DynamicTensorDescriptor.create( "out", - lengths=[1 << 30, NUM_QH, HD], coord_names=("token", "head", "dim"), + strides=(b.mul(num_qh_p, b.const_i32(HD)), HD, 1), ) # Per-segment ml base for this (q_token, q_head); the @@ -1612,7 +1634,7 @@ def build_unified_attention_reduce_tiled( # scale*log2e); lse_store converts to natural log and writes -inf for a # row with no mass (every segment empty), whose O is already 0. with b.scf_if(b.cmp_eq(tid, b.const_i32(0))): - lse_index = b.add(b.mul(q_token, b.const_i32(NUM_QH)), q_head) + lse_index = b.add(b.mul(q_token, num_qh_p), q_head) lse_store( b, lse, diff --git a/dnn-providers/hip-kernel-provider/rocke/library/tests/dispatch/attention/test_unified_lse_dispatch.py b/dnn-providers/hip-kernel-provider/rocke/library/tests/dispatch/attention/test_unified_lse_dispatch.py index 7967bb6e1b0d..eab9215099af 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/tests/dispatch/attention/test_unified_lse_dispatch.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/tests/dispatch/attention/test_unified_lse_dispatch.py @@ -88,11 +88,13 @@ def test_pinned_split_kv_spec_follows_the_request(arch): # What a long-lived cache stored for ``_decode_request`` before emit_lse existed: # the unpinned spec_hash, and the split-KV pin with its spec_hash. These are # develop's own values (gfx950's split-KV config_key moved with #13120's -# additive-bias field); emit_lse must leave them exactly as develop computes them. +# additive-bias field, and both pins with the attention ABI's move to v2 when +# the head counts and window became kernargs); emit_lse must leave them exactly +# as develop computes them. _PRE_LSE_UNPINNED_SPEC_HASH = "65416b2a8b9ffa90" _PRE_LSE_PIN = { - "gfx942": ("splitkv_seg16_t1xb_wpenone@d50574be81e67c73", "8dc6060ed42707f6"), - "gfx950": ("splitkv_seg16_t1xb_wpenone@79ee1d118140d024", "b9c537be7298900c"), + "gfx942": ("splitkv_seg16_t1xb_wpenone@cc45239d2c3577c2", "e62e88afa2989d49"), + "gfx950": ("splitkv_seg16_t1xb_wpenone@6328ef0bf325716b", "512dd0bf3b909fb8"), } diff --git a/dnn-providers/hip-kernel-provider/rocke/library/tests/golden/attention_strided_kv_ir_sha256.json b/dnn-providers/hip-kernel-provider/rocke/library/tests/golden/attention_strided_kv_ir_sha256.json index 995dc1c6b744..e46608f6623b 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/tests/golden/attention_strided_kv_ir_sha256.json +++ b/dnn-providers/hip-kernel-provider/rocke/library/tests/golden/attention_strided_kv_ir_sha256.json @@ -3,48 +3,48 @@ "0": "8cfb54512441a845a2ae325065291e39e5367be96f320596e8622d7e308ec610", "1": "b1cee01d468e92ecb5b8571323ffd897bff01c7d2e42dc8ba50814387f908a63", "2": "a80d03691e079aacfb8653f7054a836d2ecd57256285274fa51c52a8a398c727", - "3": "272f3227c873b525dc8e48c4f024f7885de375178e20f4c4d6686b4724d82c0b", - "4": "b25f6bfc84f5db8903883a85671f13b109256697e33c67877686b4d90ac07606", - "5": "bde0dfb3994fb2a89081fa05809323059111a5e7d4df55ed8859ec2f6b0cd930", + "3": "e9f5684da874f4c1bedfccc3a38750263a8d9364803d660871cba0b1176be7f4", + "4": "67c89f7a4550323dd055e30f2fdf5c12a7708b41f2af7e9cfdf1dc5faaa090e7", + "5": "77d12b3c9210688cfa94e005484498b7c285e7c173f48b46d2b87cbf9a3c3ff2", "6": "23c2881b225837148b997823910e598ad844d781441a1f122a5d370b4b97900c", "7": "966f73aef9172bbe60348f970553576d609299131e029e6842e1f4dc49c96da1", "8": "c5f4baa4b7f0237ce60bf63f6694fcca24bdd440c0fbc34a1dc0a6c8947b791c", - "9": "7866c9eb6d414e0edf3678a477c387e2742057eed6a96616544f29057d609d0a", - "10": "05d6f2c670f9c66f59974f7f2715eed844799ac869e2765d3d1f691bc417d2da", - "11": "68530d1c8c9c9a12300743809054235d88e2a1698f811bbd86a531eb47113d40", + "9": "977a23a9b768a42645ef5dd393278c4d577de5f3cb0c5b7d847260c1bf0b4462", + "10": "1238b48be61bb4a6c83e3581ed269a5de59b717901772d81d17577988a788e9e", + "11": "c996154b2f38e896037f4e14e8574c7c85d938d297ec93068407f0cca78cbde8", "12": "6adc12acb6dd75f05927ef1e4f6683b85c780abae25d76588bf57ae25f8dfd88", - "13": "d7982e56673824863076878c3461209f8557cfa8507b102071eae22d712b6414" + "13": "5266e99d4e82df6558795288cbdaef86099a84229f002945d227977834b89ffe" }, "llvm22": { "0": "14ccc0e86cb9a96163a8c32812cd718dc454f726dacced64fa3cf07db5c70c90", "1": "62086a1cb544ba32772dc112e252813ffff1961bc2db23401bf0003bce718b73", "2": "f2eb0fd45eb3693671f6d11bdce0eac92295e784e74efffe21351150e2807039", - "3": "c88f5e1109c70e0c4f132b05578c6abf440d04bccb52bd774b492d67d3acc9bb", - "4": "3e5522113635114b994bb40cf19bc51a417765ac827cbfa879acb45cb0e62796", - "5": "3c03d60f49863bc856426e1431cb98e8e682da4b6818f1df320536fab12b1314", + "3": "3a8fe959ee0737f990adf184e812dfe983b24e0f50cfdbcad667aeda89eade72", + "4": "2f21d041985082346f34105cef0ed4b58df655eed4c600a176deaaec0350cfc9", + "5": "d2770b39913bb089531c1db4d9fc31db4114053922f0bdbca3b34cdcad122a69", "6": "b4720021fe24f67cfdd7b6854a25f4b371cd5c086d7d804e4a1a451d8f885c82", "7": "281cf49847fec9bf2f0133d0b0f83c9ff72e99e3f9bfd846d2d056289604609b", "8": "8e1697deca33346efb263397db405b14b835a617e0fb8a9d6e5dd19958d767e4", - "9": "43cbf9fc03696e3fcf0a132e31a769d92cd48267eaea3526e2cc7c46c35225da", - "10": "35992a4d15f09b54bd668d4b21556a2c871fa8c6f8aa39a89169736e303476ea", - "11": "ec52f6ee37cff049c8b81eb579b71c32be965bb4730fc7b30bd477196499727f", + "9": "5f423f4b85911405a7548c7641e955fcdb485e040872744970c16c92326ae3e1", + "10": "0c4719c51595a01254f8a7c7c26bcf7cf10f7164fe54dc0c697fe36b938bc1b9", + "11": "c8bdc92a7fafee7c1187375671ea8caab00edceaf7b5f82c844d8343767f22b3", "12": "b543ef77d6fd6bcdeab9b00f9c7dd8d9823777987f2e9e09b79195f3d230be37", - "13": "f71e2da33c81aacfd3a6690c0fa2c540aa3a83586e0aa8ed0f1ed638023677db" + "13": "a2a213edfcbe00373fe8d3962f4b20a7216c9489a0c5c039536d63483f5bae57" }, "llvm23": { "0": "7ce6e33c41e2752f1863c4ef79eef519c4748cc9f1c3035cab5cab030f386d55", "1": "38fa323beea8da252380810dbbe76686f3cff17bc15aca23c32a73c11f6a70e4", "2": "5e39245a8752ce0229c9229aba2b11fc8cb07241e93ebaafa051cae24aab4514", - "3": "393df7f6bb892e49b3c3031ff696d11addb1da00ad2088aa881ea37074c7db77", - "4": "db62e65f3aa74588eb95d2bf69935aa6331d0a57d42996bbe81c7510e146f119", - "5": "ee8248f21f5631e005750601f275daf2fd57d87f102a241c2fee76caa4567d33", + "3": "76bca1175f651389730a2856cf8bb9cf9d79998445425cde57ac37700dd1d2ef", + "4": "0ea423ae0c4cdd6d4b72ed754b1a18dab37fb3f6b9fbad492b1bcd036529c823", + "5": "29304ac08dddbce980440041340d863e3996b72fa786fecd5f739ca8c14bd150", "6": "b3add7e87f7a74a750718e45156d80bc9a5725c33a8a25fadc55b71bb600af89", "7": "94ea80d639a380a3cf0069ff8679e6f197fa1c40c9a4885413d747341964ef9c", "8": "d8a78f5237a1c466494c27808c41590153f9ae47d899a5d6ad240510e680d89f", - "9": "cb2426e5ff5ce00a3094ea4217df160031006f22cfe058568e4c959f1ecff086", - "10": "400f75d835f4bb4d37dfe0039c609deb5e172c77e5a83b87b89071ed3b0017bf", - "11": "0e441b0ae432165d00da483191b9470dcfc043c2bacd386b86fe6bef8f05266f", + "9": "67bfeb3341425741eb76f57e12f56c47a87bbdb73b1e5641c4a4f256f2647591", + "10": "59ba76e31fd186ba00c06d841b8a038c8941a8419e56b8fa46daf2ea9eb61028", + "11": "a30d1148dddca335071fee50a5cedd26bb5323693c278b14d22580d6cc5df811", "12": "0ec51a3f92220ca3dad8979ef3dbe976f89ab95e28de5a178eaee7477661b8d4", - "13": "74d14799531a8e079e822a46af3a22f1875a85a46c6c561ef26427fd44ffd748" + "13": "9bb3761cb7e928a0efc7789cda47a011260003ba29f7db4cdde89af0a0bda125" } } diff --git a/dnn-providers/hip-kernel-provider/rocke/library/tests/parity/gfx950_attention_tiled_2d_emit.c b/dnn-providers/hip-kernel-provider/rocke/library/tests/parity/gfx950_attention_tiled_2d_emit.c index a15978519ef9..acc324ba9202 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/tests/parity/gfx950_attention_tiled_2d_emit.c +++ b/dnn-providers/hip-kernel-provider/rocke/library/tests/parity/gfx950_attention_tiled_2d_emit.c @@ -914,6 +914,68 @@ static int make_spec(int idx, rocke_attention_tiled_2d_spec_t* s) s->tile_size = 64; break; + /* bf16 D64/b16 GQA-4 multi-batch: the shipping spec, whose one K/V copy per + * tile (32 tokens) spans two pages, so the bf16 loaders keep the paged-KV + * descriptor. Every other bf16/fp16 case takes the per-page copy offsets. */ + case 56: + s->head_size = 64; + s->block_size = 16; + s->num_query_heads = 32; + s->num_kv_heads = 8; + s->dtype = "bf16"; + s->use_sinks = false; + s->sliding_window = 0; + s->has_softcap = false; + s->num_seqs = 2; + s->num_warps = 4; + s->has_waves_per_eu = true; + s->waves_per_eu = 2; + s->has_tile_size = true; + s->tile_size = 32; + s->block_m_per_warp = 32; + s->use_mfma_32x32 = true; + s->use_transposed_qk_32x32 = true; + s->use_transposed_scalar_state = true; + s->use_transposed_mask_once = true; + s->use_transposed_half_local_pv = true; + s->use_mfma32_skip_legacy_qreg = true; + s->use_transposed_mask_limit = true; + break; + /* cfg57/58: page copies at the 16-page prefetch bound (page ids carried) + * and past it (page ids read in place). */ + case 57: + s->head_size = 64; + s->block_size = 16; + s->num_query_heads = 32; + s->num_kv_heads = 8; + s->dtype = "bf16"; + s->use_sinks = false; + s->sliding_window = 0; + s->has_softcap = false; + s->num_seqs = 2; + s->has_tile_size = true; + s->tile_size = 256; + s->block_m_per_warp = 32; + s->use_mfma_32x32 = true; + s->use_transposed_qk_32x32 = true; + break; + case 58: + s->head_size = 64; + s->block_size = 16; + s->num_query_heads = 32; + s->num_kv_heads = 8; + s->dtype = "bf16"; + s->use_sinks = false; + s->sliding_window = 0; + s->has_softcap = false; + s->num_seqs = 2; + s->has_tile_size = true; + s->tile_size = 288; + s->block_m_per_warp = 32; + s->use_mfma_32x32 = true; + s->use_transposed_qk_32x32 = true; + break; + default: return -1; } diff --git a/dnn-providers/hip-kernel-provider/rocke/library/tests/parity/gfx950_attention_tiled_2d_emit.py b/dnn-providers/hip-kernel-provider/rocke/library/tests/parity/gfx950_attention_tiled_2d_emit.py index e1be2aa15915..2d8b5f534b71 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/tests/parity/gfx950_attention_tiled_2d_emit.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/tests/parity/gfx950_attention_tiled_2d_emit.py @@ -898,6 +898,63 @@ block_m_per_warp=32, tile_size=64, ), + # bf16 D64/b16 GQA-4 multi-batch: the shipping spec, whose one K/V copy per + # tile (32 tokens) spans two pages, so the bf16 loaders keep the paged-KV + # descriptor. Every other bf16/fp16 case takes the per-page copy offsets. + 56: dict( + head_size=64, + block_size=16, + num_query_heads=32, + num_kv_heads=8, + dtype="bf16", + use_sinks=False, + sliding_window=0, + has_softcap=False, + num_seqs=2, + num_warps=4, + waves_per_eu=2, + tile_size=32, + block_m_per_warp=32, + use_mfma_32x32=True, + use_transposed_qk_32x32=True, + use_transposed_scalar_state=True, + use_transposed_mask_once=True, + use_transposed_half_local_pv=True, + use_mfma32_skip_legacy_qreg=True, + use_transposed_mask_limit=True, + ), + # cfg57/58: page copies at the 16-page prefetch bound (page ids carried) + # and past it (page ids read in place). + 57: dict( + head_size=64, + block_size=16, + num_query_heads=32, + num_kv_heads=8, + dtype="bf16", + use_sinks=False, + sliding_window=0, + has_softcap=False, + num_seqs=2, + tile_size=256, + block_m_per_warp=32, + use_mfma_32x32=True, + use_transposed_qk_32x32=True, + ), + 58: dict( + head_size=64, + block_size=16, + num_query_heads=32, + num_kv_heads=8, + dtype="bf16", + use_sinks=False, + sliding_window=0, + has_softcap=False, + num_seqs=2, + tile_size=288, + block_m_per_warp=32, + use_mfma_32x32=True, + use_transposed_qk_32x32=True, + ), } diff --git a/dnn-providers/hip-kernel-provider/rocke/library/tests/parity/gfx950_attention_tiled_2d_fastkv_regp_emit.c b/dnn-providers/hip-kernel-provider/rocke/library/tests/parity/gfx950_attention_tiled_2d_fastkv_regp_emit.c index 7f3d4efa1579..0cd599e0d53d 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/tests/parity/gfx950_attention_tiled_2d_fastkv_regp_emit.c +++ b/dnn-providers/hip-kernel-provider/rocke/library/tests/parity/gfx950_attention_tiled_2d_fastkv_regp_emit.c @@ -75,6 +75,16 @@ static int make_spec(int idx, rocke_attention_tiled_2d_spec_t* s) case 5: s->use_grouped_kv2_softmax = true; break; + /* The pack predicate admits windowed, sinks and i64 specs; pin each body. */ + case 6: + s->sliding_window = 256; + break; + case 7: + s->use_sinks = true; + break; + case 8: + s->use_i64_kv_addr = true; + break; default: return -1; } diff --git a/dnn-providers/hip-kernel-provider/rocke/library/tests/parity/gfx950_attention_tiled_2d_fastkv_regp_emit.py b/dnn-providers/hip-kernel-provider/rocke/library/tests/parity/gfx950_attention_tiled_2d_fastkv_regp_emit.py index d17b1cad2b35..517151d181d8 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/tests/parity/gfx950_attention_tiled_2d_fastkv_regp_emit.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/tests/parity/gfx950_attention_tiled_2d_fastkv_regp_emit.py @@ -47,6 +47,10 @@ 3: dict(_BASE, use_transposed_half_local_pv=True, use_mfma32_skip_legacy_qreg=True), 4: dict(_BASE, use_agpr_alloc_zero=True, use_transposed_half_local_pv=True), 5: dict(_BASE, use_grouped_kv2_softmax=True), + # The pack predicate admits windowed, sinks and i64 specs; pin each body. + 6: dict(_BASE, sliding_window=256), + 7: dict(_BASE, use_sinks=True), + 8: dict(_BASE, use_i64_kv_addr=True), } diff --git a/dnn-providers/hip-kernel-provider/rocke/library/tests/test_attention_builds.py b/dnn-providers/hip-kernel-provider/rocke/library/tests/test_attention_builds.py index a87bda7814fc..d5316acf0580 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/tests/test_attention_builds.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/tests/test_attention_builds.py @@ -1135,7 +1135,10 @@ def test_unified_attention_additive_bias_abi_matches_signature(self): build_unified_attention_2d_tiled, build_unified_attention_3d_tiled, ) - from kernels.common.attention_unified import _3d_signature, _attn_signature + from kernels.common.attention_unified import ( + _tiled_2d_signature, + _tiled_3d_signatures, + ) common = dict( head_size=64, @@ -1152,11 +1155,8 @@ def test_unified_attention_additive_bias_abi_matches_signature(self): k = build_unified_attention_2d_tiled( UnifiedAttention2DTiledSpec(**common, use_additive_bias=use_ab) ) - sig = _attn_signature( - "bf16", - include_bt_stride=True, - include_qq_bias_stride=True, - include_additive_bias=use_ab, + sig = _tiled_2d_signature( + "gfx950", "bf16", include_additive_bias=use_ab ) self.assertEqual([p.name for p in k.params], [a["name"] for a in sig]) self.assertEqual( @@ -1168,7 +1168,9 @@ def test_unified_attention_additive_bias_abi_matches_signature(self): **common, num_segments=4, use_additive_bias=use_ab ) ) - sig = _3d_signature("bf16", include_additive_bias=use_ab) + sig, _ = _tiled_3d_signatures( + "gfx950", "bf16", include_additive_bias=use_ab + ) self.assertEqual([p.name for p in seg.params], [a["name"] for a in sig]) def test_unified_attention_additive_bias_must_be_bool(self): diff --git a/dnn-providers/hip-kernel-provider/rocke/library/tests/test_attention_gfx950_runtime_shape_numeric.py b/dnn-providers/hip-kernel-provider/rocke/library/tests/test_attention_gfx950_runtime_shape_numeric.py new file mode 100644 index 000000000000..db5687342b18 --- /dev/null +++ b/dnn-providers/hip-kernel-provider/rocke/library/tests/test_attention_gfx950_runtime_shape_numeric.py @@ -0,0 +1,329 @@ +# Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +# SPDX-License-Identifier: MIT + +"""One gfx950 tiled binary serves every head pair, window size and batch. + +The gfx950 tiled 2D and 3D kernels read the head counts and the sliding window +as kernargs and derive the cu_q search trip count from ``num_seqs``, so shapes +that differ only in those share one cache key and one compiled binary. Each +case below runs a pair of such shapes back to back through +``run_unified_attention_torch``: the second launch must reuse the first binary +(no new compile) and both must match an FP32 paged reference. + +The second shape of each pair changes the head counts at the same GQA ratio, +the window size, or the batch by several log2 search buckets, so a value that +is still baked into the binary corrupts it. Block-table rows are wider than the +batch needs and unused slots point at a poison block, so a wrong stride cannot +pass by luck. + +The fast paged-KV descriptor bakes the 64/8 head geometry, so its pair keeps +64/8 and changes the batch; its 32/4 and 16/2 neighbours at the same ratio take +the plain combo binary. + +The routing and key checks run on any host; the launches need a gfx950 GPU. +""" + +from __future__ import annotations + +import itertools +import math + +import pytest + +from kernels.common import attention_unified as au + +_ARCH = "gfx950" +_CUS = 256 + + +def _problem(*, q_lens, kv_lens, hq, hkv, d, bs, dtype, window, num_blocks): + return 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=dtype, + sliding_window=window, + num_cus=_CUS, + num_kv_blocks=num_blocks, + ) + + +def _shape(q_lens, kv_lens, hq, hkv, window): + return dict(q_lens=q_lens, kv_lens=kv_lens, hq=hq, hkv=hkv, window=window) + + +# (id, path, route predicate or None, (d, bs, dtype), first shape, second shape). +# The route predicate is an attention_unified selector both shapes must satisfy. +# Both shapes of a pair sit on one side of every seqlen threshold the selectors +# read (combo above 256 query tokens, hipcc above 512), so only the head counts, +# the window and the batch differ. +_CASES = ( + ( + "2d-bf16-d128", + "2d", + None, + (128, 16, "bf16"), + _shape([300, 200, 250], [300, 450, 250], 32, 8, 0), + _shape( + [90, 130, 70, 110, 40, 300, 60, 80, 120], + [90, 300, 70, 210, 40, 300, 60, 380, 120], + 16, + 4, + 0, + ), + ), + # Head counts alone: the token and sequence counts match, so only the head + # counts tell the two launches apart. + ( + "2d-bf16-d128-single-batch-heads-only", + "2d", + "_enable_single_batch_combo", + (128, 16, "bf16"), + _shape([1025], [1025], 32, 8, 0), + _shape([1025], [1025], 64, 16, 0), + ), + ( + "2d-fp16-d64-window", + "2d", + None, + (64, 32, "fp16"), + _shape([260, 180], [400, 180], 8, 8, 128), + _shape( + [70, 90, 50, 100, 60, 80, 40, 30, 110], + [200, 90, 300, 100, 60, 250, 40, 30, 110], + 16, + 16, + 512, + ), + ), + # The single-batch fp16 d128 window cohort: heads and window change. + ( + "2d-fp16-d128-single-batch-window", + "2d", + "_enable_single_batch_combo", + (128, 16, "fp16"), + _shape([2048], [2048], 32, 8, 256), + _shape([2048], [2048], 64, 16, 1024), + ), + ( + "2d-bf16-d64-fast-desc-batch", + "2d", + "_enable_fast_paged_kv_desc", + (64, 32, "bf16"), + _shape([500, 400], [700, 900], 64, 8, 0), + _shape( + [300, 280, 400, 260, 350, 310, 270, 290, 330], + [300, 600, 400, 800, 350, 310, 500, 290, 330], + 64, + 8, + 0, + ), + ), + ( + "2d-bf16-d64-combo-heads", + "2d", + "_enable_combo_2d", + (64, 32, "bf16"), + _shape([500, 400], [700, 900], 32, 4, 0), + _shape( + [300, 280, 400, 260, 350, 310, 270, 290, 330], + [300, 600, 400, 800, 350, 310, 500, 290, 330], + 16, + 2, + 0, + ), + ), + ( + "2d-bf16-d256", + "2d", + "_d256_gfx950_fast", + (256, 16, "bf16"), + _shape([512, 512], [1024, 1024], 16, 2, 0), + _shape( + [300, 500, 200, 400, 250, 350, 150, 450, 280], + [300, 700, 200, 400, 250, 600, 150, 450, 280], + 32, + 4, + 0, + ), + ), + # The decode split sizes segments by num_seqs * num_kv_heads, so the 3D pairs + # hold that product fixed; windowed decode routes to 2D. + ( + "3d-decode", + "3d", + None, + (128, 16, "bf16"), + _shape([1, 1], [4096, 2500], 64, 16, 0), + _shape([1] * 8, [3000, 900, 4096, 1700, 2200, 3500, 1200, 2800], 16, 4, 0), + ), + ( + "3d-two-token", + "3d", + None, + (128, 16, "bf16"), + _shape([2, 2], [4096, 3100], 64, 16, 0), + _shape([2] * 8, [2500, 4096, 1900, 2600, 3100, 2200, 2400, 2700], 16, 4, 0), + ), +) + + +def _problems(geometry, shape): + d, bs, dtype = geometry + needed = sum(math.ceil(n / bs) for n in shape["kv_lens"]) + return _problem( + q_lens=shape["q_lens"], + kv_lens=shape["kv_lens"], + hq=shape["hq"], + hkv=shape["hkv"], + d=d, + bs=bs, + dtype=dtype, + window=shape["window"], + num_blocks=needed + 1, + ) + + +def _key(problem, path): + if path == "2d": + return au._tiled_cache_key(problem, _ARCH) + return au._tiled_3d_cache_key(problem, _ARCH) + + +@pytest.mark.parametrize( + "path, route, geometry, first, second", + [pytest.param(*c[1:], id=c[0]) for c in _CASES], +) +def test_pair_shares_one_binary(path, route, geometry, first, second): + """Both shapes take the intended path and route and share one key, so the + GPU test below exercises one binary on two shapes.""" + a, b = _problems(geometry, first), _problems(geometry, second) + for p in (a, b): + assert p.select_path() == path + if route is not None: + assert getattr(au, route)(p, _ARCH) + assert a.num_queries_per_kv == b.num_queries_per_kv + batch_changes = a.num_seqs != b.num_seqs + if batch_changes and not (path == "3d" and au._use_decode_grid(a, _ARCH)): + # The decode grid indexes sequences directly and emits no search. + assert au._binary_search_iters(a) != au._binary_search_iters(b) + assert _key(a, path) == _key(b, path) + + +@pytest.fixture(scope="module") +def torch(): + torch = pytest.importorskip("torch") + if not torch.cuda.is_available(): + pytest.skip("requires a gfx950 GPU") + name = torch.cuda.get_device_properties(0).gcnArchName + if not name.startswith(_ARCH): + pytest.skip(f"requires a gfx950 GPU, found {name}") + return torch + + +def _case(torch, geometry, shape, seed): + d, bs, dtype = geometry + tdt = torch.bfloat16 if dtype == "bf16" else torch.float16 + q_lens, kv_lens, hq, hkv = ( + shape["q_lens"], + shape["kv_lens"], + shape["hq"], + shape["hkv"], + ) + 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() + table = torch.full((len(kv_lens), max(needed) + 3), poison, dtype=torch.int32) + for s, n in enumerate(needed): + table[s, :n] = torch.tensor(order[:n], dtype=torch.int32) + order = order[n:] + q = torch.randn(sum(q_lens), hq, d, generator=gen).to(tdt) + k = torch.randn(num_blocks, bs, hkv, d, generator=gen).to(tdt) + v = torch.randn(num_blocks, bs, hkv, d, generator=gen).to(tdt) + k[poison] = 64.0 + v[poison] = 64.0 + cu = torch.tensor([0, *itertools.accumulate(q_lens)], dtype=torch.int32) + return q.cuda(), k.cuda(), v.cuda(), cu.cuda(), table.cuda() + + +def _reference(torch, q, k, v, cu, table, shape, d, bs): + q_lens, kv_lens, hq, hkv = ( + shape["q_lens"], + shape["kv_lens"], + shape["hq"], + shape["hkv"], + ) + window = shape["window"] + out = torch.zeros(q.shape, dtype=torch.float32, device=q.device) + 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() + keys = keys.repeat_interleave(hq // hkv, dim=1) + values = values.repeat_interleave(hq // hkv, dim=1) + start = int(cu[s]) + queries = q[start : start + q_len].float() + scores = torch.einsum("qhd,khd->hqk", queries, keys) / math.sqrt(d) + q_pos = (kv_len - q_len) + torch.arange(q_len, device=q.device)[:, None] + k_pos = torch.arange(kv_len, device=q.device)[None, :] + visible = k_pos <= q_pos + if window > 0: + visible = visible & (q_pos - k_pos < window) + scores = scores.masked_fill(~visible, float("-inf")) + out[start : start + q_len] = torch.einsum( + "hqk,khd->qhd", scores.softmax(-1), values + ) + return out + + +def _cache(path): + return au._ATTN_3D_TILED_CACHE if path == "3d" else au._ATTN_TILED_CACHE + + +def _compiled(): + return len(au._ATTN_TILED_CACHE) + len(au._ATTN_3D_TILED_CACHE) + + +@pytest.mark.parametrize( + "path, geometry, first, second", + [pytest.param(c[1], *c[3:], id=c[0]) for c in _CASES], +) +def test_pair_numeric_on_one_binary(torch, path, geometry, first, second): + d, bs, _ = geometry + compiled = None + for i, shape in enumerate((first, second)): + problem = _problems(geometry, shape) + q, k, v, cu, table = _case(torch, geometry, shape, seed=i) + out = torch.empty_like(q) + au.run_unified_attention_torch( + problem=problem, + q=q, + k=k, + v=v, + out=out, + cu_seqlens_q=cu, + seqused_k=torch.tensor(shape["kv_lens"], dtype=torch.int32, device="cuda"), + softmax_scale=1.0 / math.sqrt(d), + block_table=table, + softcap=0.0, + backend="3d" if path == "3d" else "tiled", + use_graph=False, + ) + torch.cuda.synchronize() + assert _key(problem, path) in _cache(path), f"shape {i} missed the {path} cache" + if compiled is None: + compiled = _compiled() + else: + assert _compiled() == compiled, "the second shape compiled a new binary" + want = _reference(torch, q, k, v, cu, table, shape, d, bs) + max_abs = (out.float() - want).abs().max().item() + assert ( + math.isfinite(max_abs) and max_abs <= 4e-2 + ), f"shape {i}: max_abs {max_abs:.3e} vs fp32 reference" 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 8e5da31b28d5..7cba980d1254 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 @@ -20,7 +20,7 @@ lower to byte-identical IR. * Launch geometry: the grid caches still separate batch counts. * ``_binary_search_iters`` mirrors every spec class that bakes the search; - gfx942 searches at runtime, so its batch sizes share a binary. + gfx942 and gfx950 search at runtime, so their batch sizes share a binary. No GPU required. """ @@ -144,6 +144,7 @@ def _built_3d(problem, arch): _FEATURES = ( (0, False, None), (128, False, None), + (512, False, None), (0, True, None), (0, False, "softcap"), (0, False, "alibi"), @@ -156,7 +157,9 @@ def _static_grid(arch): (False, True), (64, 128, 256), (16, 32, 64), - ((8, 8), (32, 8), (64, 8)), + # Two pairs share each of ratios 1 and 8, so a body that reads the + # head counts as kernargs is checked against an under-key at one ratio. + ((8, 8), (16, 16), (32, 8), (32, 4), (64, 8)), _FEATURES, ): yield dict( @@ -288,8 +291,9 @@ def _assert_split(key_fn, built_fn, a, b, field, arch): assert key_fn(a, arch) != key_fn(b, arch) -# gfx942 searches at runtime, so its buckets share a binary (_SHARED_PAIRS). -@pytest.mark.parametrize("arch", ("gfx950", "gfx1250")) +# gfx942 and gfx950 search at runtime, so their buckets share a binary +# (_SHARED_PAIRS). +@pytest.mark.parametrize("arch", ("gfx1250",)) def test_search_bucket_edge_splits_2d_key(routing_env, 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. @@ -310,25 +314,6 @@ def test_search_bucket_edge_splits_2d_key(routing_env, arch): assert _lower_2d(a, arch, "llvm22") != _lower_2d(b, arch, "llvm22") -@pytest.mark.parametrize("arch", ("gfx950",)) -def test_search_bucket_edge_splits_3d_key(routing_env, arch): - # Two query tokens per sequence, so the launch takes the searching prefill - # grid rather than the decode grid. - base = dict( - num_query_heads=32, - num_kv_heads=8, - head_size=128, - block_size=16, - max_seqlen_q=2, - max_seqlen_k=8192, - dtype="bf16", - ) - a = P(total_q=6, num_seqs=3, **base) - b = P(total_q=8, num_seqs=4, **base) - assert not au._use_decode_grid(a, arch) - _assert_split(au._tiled_3d_cache_key, _built_3d, a, b, "binary_search_iters", arch) - - def test_fp8_mfma_qk_kv_threshold_splits_key(routing_env): # 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. @@ -460,6 +445,15 @@ def _lower_3d(problem, arch, flavor): sliding_window=4096, dtype="bf16", ) +_D128_SB_FP16 = dict( + num_query_heads=32, + num_kv_heads=8, + head_size=128, + block_size=16, + max_seqlen_q=4096, + max_seqlen_k=4096, + dtype="fp16", +) # Two query tokens per sequence: the 3D searching grid, not the decode grid. _D128_TWO_TOKENS = dict( num_query_heads=32, @@ -520,6 +514,135 @@ def _lower_3d(problem, arch, flavor): ) for arch in ("gfx942", "gfx950") ), + # gfx950 reads the search trip count from num_seqs at runtime too. + ( + "gfx950-2d-ns3-vs-ns9", + "gfx950", + "2d", + P( + total_q=3 * 1024, + num_seqs=3, + max_seqlen_q=1024, + max_seqlen_k=1024, + **_D64_GQA8, + ), + P( + total_q=9 * 1024, + num_seqs=9, + max_seqlen_q=1024, + max_seqlen_k=1024, + **_D64_GQA8, + ), + ), + ( + "gfx950-3d-prefill-ns3-vs-ns9", + "gfx950", + "3d", + P(total_q=6, num_seqs=3, **_D128_TWO_TOKENS), + P(total_q=18, num_seqs=9, **_D128_TWO_TOKENS), + ), + # gfx950 reads the window size as a kernarg, so window sizes share a binary. + ( + "gfx950-2d-sw256-vs-sw1024", + "gfx950", + "2d", + P( + total_q=2048, + num_seqs=2, + max_seqlen_q=1024, + max_seqlen_k=1024, + sliding_window=256, + **_D64_GQA8, + ), + P( + total_q=2048, + num_seqs=2, + max_seqlen_q=1024, + max_seqlen_k=1024, + sliding_window=1024, + **_D64_GQA8, + ), + ), + # The single-batch fp16 d128 window cohort picks its own KV tile. + ( + "gfx950-2d-single-batch-sw256-vs-sw1024", + "gfx950", + "2d", + P(total_q=4096, num_seqs=1, sliding_window=256, **_D128_SB_FP16), + P(total_q=4096, num_seqs=1, sliding_window=1024, **_D128_SB_FP16), + ), + ( + "gfx950-3d-decode-sw256-vs-sw4096", + "gfx950", + "3d", + P(total_q=4, num_seqs=4, sliding_window=256, **_D128_DECODE), + P(total_q=4, num_seqs=4, sliding_window=4096, **_D128_DECODE), + ), + # gfx950 reads the head counts as kernargs, so head pairs at one GQA ratio + # share a binary. 32/4 and 16/2 stay off the 64/8 fast paged-KV descriptor. + ( + "gfx950-2d-h32kv4-vs-h16kv2", + "gfx950", + "2d", + P( + total_q=2048, + num_seqs=2, + max_seqlen_q=1024, + max_seqlen_k=1024, + **dict(_D64_GQA8, num_query_heads=32, num_kv_heads=4), + ), + P( + total_q=2048, + num_seqs=2, + max_seqlen_q=1024, + max_seqlen_k=1024, + **dict(_D64_GQA8, num_query_heads=16, num_kv_heads=2), + ), + ), + ( + "gfx950-2d-single-batch-h32kv8-vs-h64kv16", + "gfx950", + "2d", + P(total_q=4096, num_seqs=1, **_D128_SB_FP16), + P( + total_q=4096, + num_seqs=1, + **dict(_D128_SB_FP16, num_query_heads=64, num_kv_heads=16), + ), + ), + ( + "gfx950-2d-d256-h16kv2-vs-h32kv4", + "gfx950", + "2d", + P(total_q=2 * 4096, num_seqs=2, **_D256_4WARP), + P( + total_q=2 * 4096, + num_seqs=2, + **dict(_D256_4WARP, num_query_heads=32, num_kv_heads=4), + ), + ), + ( + "gfx950-3d-prefill-h32kv8-vs-h64kv16", + "gfx950", + "3d", + P(total_q=6, num_seqs=3, **_D128_TWO_TOKENS), + P( + total_q=6, + num_seqs=3, + **dict(_D128_TWO_TOKENS, num_query_heads=64, num_kv_heads=16), + ), + ), + ( + "gfx950-3d-decode-h32kv8-vs-h64kv16", + "gfx950", + "3d", + P(total_q=31, num_seqs=31, **_D128_DECODE), + P( + total_q=31, + num_seqs=31, + **dict(_D128_DECODE, num_query_heads=64, num_kv_heads=16), + ), + ), # gfx942 reads the search trip count from num_seqs at runtime, so batch # sizes several log2 buckets apart share a binary on every builder route. ( @@ -673,6 +796,19 @@ def test_same_outputs_share_key_and_ir(routing_env, arch, path, a, b, flavor): assert lower(a, arch, flavor) == lower(b, arch, flavor) +def test_fast_paged_kv_desc_splits_head_pairs_at_one_ratio(routing_env): + # The gfx950 fast paged-KV descriptor bakes the 64/8 head geometry. The key + # carries only the GQA ratio, so it must carry the fast-desc flag too: a + # 32/4 problem at ratio 8 must not reuse the 64/8 fast binary. + arch = "gfx950" + base = dict(total_q=2048, num_seqs=2, max_seqlen_q=1024, max_seqlen_k=1024) + a = P(**base, **_D64_GQA8) + b = P(**base, **dict(_D64_GQA8, num_query_heads=32, num_kv_heads=4)) + assert au._tiled_spec_from_problem(a, arch).use_fast_paged_kv_desc + assert not au._tiled_spec_from_problem(b, arch).use_fast_paged_kv_desc + assert au._tiled_cache_key(a, arch) != au._tiled_cache_key(b, arch) + + # --------------------------------------------------------------------------- # Launch geometry still separates batch counts # --------------------------------------------------------------------------- @@ -703,12 +839,14 @@ def test_launch_meta_separates_num_seqs_under_one_key(routing_env): assert grid_b[1] - grid_a[1] == b.num_seqs - a.num_seqs -def test_launch_geometry_separates_head_counts_under_one_key(routing_env): - # gfx942 reads the head counts as kernargs, so one binary serves 32/8 and - # 64/16. The 2D grid spans the kv heads and the 3D workspace the query - # heads, so with equal token and sequence counts the launch caches must - # still tell them apart. - arch = "gfx942" +@pytest.mark.parametrize("arch,num_cus", [("gfx942", 304), ("gfx950", 256)]) +def test_launch_geometry_separates_head_counts_under_one_key( + routing_env, arch, num_cus +): + # gfx942 and gfx950 read the head counts as kernargs, so one binary serves + # 32/8 and 64/16. The 2D grid spans the kv heads and the 3D workspace the + # query heads, so with equal token and sequence counts the launch caches + # must still tell them apart. base = dict( total_q=1025, num_seqs=1, @@ -717,7 +855,7 @@ def test_launch_geometry_separates_head_counts_under_one_key(routing_env): max_seqlen_q=1025, max_seqlen_k=1025, dtype="bf16", - num_cus=304, + num_cus=num_cus, ) a = P(num_query_heads=32, num_kv_heads=8, **base) b = P(num_query_heads=64, num_kv_heads=16, **base) @@ -738,8 +876,6 @@ def test_launch_geometry_separates_head_counts_under_one_key(routing_env): # (spec class, floor) for every class that bakes the search. gfx1250 3D floors # it at 32 iterations. _SPEC_CLASSES = ( - (Gfx950Spec2D, 1), - (Gfx950Spec3D, 1), (Gfx1250Spec2D, 1), (Gfx1250Spec3D, 32), ) @@ -768,13 +904,20 @@ def test_binary_search_iters_mirrors_spec(spec_cls, floor, num_seqs): assert au._binary_search_iters(problem, floor=floor) == spec.binary_search_iters -@pytest.mark.parametrize("spec_cls", (Gfx942Spec2D, Gfx942Spec3D)) -def test_gfx942_specs_search_at_runtime(spec_cls): +@pytest.mark.parametrize( + "arch, path, spec_cls", + ( + ("gfx942", "2d", Gfx942Spec2D), + ("gfx942", "3d", Gfx942Spec3D), + ("gfx950", "2d", Gfx950Spec2D), + ("gfx950", "3d", Gfx950Spec3D), + ), +) +def test_specs_search_at_runtime(arch, path, spec_cls): assert spec_cls.RUNTIME_SEARCH_ITERS assert not hasattr(spec_cls, "binary_search_iters") - arch, path = "gfx942", "2d" if spec_cls is Gfx942Spec2D else "3d" assert au._searches_at_runtime(arch, path) - assert not au._searches_at_runtime("gfx950", path) + assert not au._searches_at_runtime("gfx1250", path) def _trip_count(num_seqs: int) -> int: @@ -823,7 +966,8 @@ def test_runtime_trip_count_covers_every_batch(): assert n.bit_length() <= trip <= n.bit_length() + 1, n -def test_gfx942_i64_kv_addressing_is_in_the_kernel_name(routing_env): +@pytest.mark.parametrize("arch,num_cus", [("gfx942", 304), ("gfx950", 256)]) +def test_i64_kv_addressing_is_in_the_2d_kernel_name(routing_env, arch, num_cus): # i64 K/V addressing changes the 2D loaders, so a binary is identified by # it: the name, not only the in-process key, must tell the two apart. base = dict( @@ -836,21 +980,22 @@ def test_gfx942_i64_kv_addressing_is_in_the_kernel_name(routing_env): max_seqlen_q=1024, max_seqlen_k=1024, dtype="bf16", - num_cus=304, + num_cus=num_cus, ) block_bytes = 16 * 8 * 128 * 2 small = P(num_kv_blocks=64, **base) large = P(num_kv_blocks=(1 << 31) // block_bytes + 1, **base) assert not au._enable_i64_kv_addr(small) and au._enable_i64_kv_addr(large) - names = [ - au._tiled_spec_from_problem(p, "gfx942").kernel_name() for p in (small, large) - ] + names = [au._tiled_spec_from_problem(p, arch).kernel_name() for p in (small, large)] assert "i64" not in names[0].split("_") assert "i64" in names[1].split("_") assert names[0].replace("_bf16_", "_bf16_i64_", 1) == names[1] -def test_gfx942_fast_paged_kv_desc_names_its_head_counts(): +@pytest.mark.parametrize( + "spec_cls", [Gfx942Spec2D, Gfx950Spec2D], ids=["gfx942", "gfx950"] +) +def test_fast_paged_kv_desc_names_its_head_counts(spec_cls): # The fast paged-KV descriptor body shifts by a 64/8 K/V stride instead of # reading the head-count kernargs, so its name must not share the plain # g8 symbol that serves every head pair at that ratio. @@ -866,7 +1011,7 @@ def test_gfx942_fast_paged_kv_desc_names_its_head_counts(): sliding_window=0, has_softcap=False, ) - plain = Gfx942Spec2D(**base).kernel_name() - fast = Gfx942Spec2D(**base, use_fast_paged_kv_desc=True).kernel_name() + plain = spec_cls(**base).kernel_name() + fast = spec_cls(**base, use_fast_paged_kv_desc=True).kernel_name() assert "h64kv8" not in plain.split("_") assert fast == plain.replace("_g8_", "_g8_h64kv8_", 1) + "_fastkvdesc" diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/CMakeLists.txt b/dnn-providers/hip-kernel-provider/rocke/platform/CMakeLists.txt index d5c5efa2d482..21ed12deb892 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/CMakeLists.txt +++ b/dnn-providers/hip-kernel-provider/rocke/platform/CMakeLists.txt @@ -648,16 +648,24 @@ if(ROCKE_INSTALL_TESTS) " ENVIRONMENT \"PYTHONPATH=${ROCKE_TEST_CTEST_REL}tests/library\"\n" " TIMEOUT 180)\n") - # Tiled attention cache-key contract (no GPU). Its invariant sweep compares + # Tiled attention cache-key contract (no GPU). The targeted share/split + # checks take seconds and run in the quick tier. The invariant sweep compares # key equality with build equality across every selector threshold for each - # tiled arch, so each sweep case runs for tens of seconds: it gets its own - # entry and budget rather than the 30 s per-test limit of the list above. + # tiled arch, so each sweep case runs for tens of seconds: it is a separate + # standard-tier entry with its own budget. string(APPEND _ROCKE_INSTALLED_CTEST_ENTRIES - "add_test(rocke_library_attention_cache_key_pytest python -m pytest \"${ROCKE_TEST_CTEST_REL}tests/library/tests/test_attention_tiled_cache_key.py\" -v -rs --timeout=240)\n") + "add_test(rocke_library_attention_cache_key_pytest python -m pytest \"${ROCKE_TEST_CTEST_REL}tests/library/tests/test_attention_tiled_cache_key.py\" -k \"not test_key_matches_build_identity\" -v -rs --timeout=30)\n") string(APPEND _ROCKE_INSTALLED_CTEST_ENTRIES "set_tests_properties(rocke_library_attention_cache_key_pytest PROPERTIES\n" " LABELS \"ckc\\;python\\;attention\\;host\\;quick\"\n" " ENVIRONMENT \"PYTHONPATH=${ROCKE_TEST_CTEST_REL}tests/library\"\n" + " TIMEOUT 180)\n") + string(APPEND _ROCKE_INSTALLED_CTEST_ENTRIES + "add_test(rocke_library_attention_cache_key_sweep_pytest python -m pytest \"${ROCKE_TEST_CTEST_REL}tests/library/tests/test_attention_tiled_cache_key.py\" -k test_key_matches_build_identity -v -rs --timeout=240)\n") + string(APPEND _ROCKE_INSTALLED_CTEST_ENTRIES + "set_tests_properties(rocke_library_attention_cache_key_sweep_pytest PROPERTIES\n" + " LABELS \"ckc\\;python\\;attention\\;host\\;standard\"\n" + " ENVIRONMENT \"PYTHONPATH=${ROCKE_TEST_CTEST_REL}tests/library\"\n" " TIMEOUT 900)\n") # DSA lightning-indexer CPU tests (no GPU: validator, dispatch wiring, and the @@ -854,6 +862,7 @@ if(ROCKE_INSTALL_TESTS) "tests/library/tests/test_gfx942_gqa_head_fold_numeric.py" "tests/library/tests/test_attention_gfx942_kv_pages.py" "tests/library/tests/test_attention_gfx942_runtime_shape_numeric.py" + "tests/library/tests/test_attention_gfx950_runtime_shape_numeric.py" "tests/library/tests/test_attention_tiled_3d_kv_lds_swizzle.py" "tests/library/tests/test_attention_tiled_3d_running_max.py" "tests/library/tests/test_attention_tiled_3d_wave_page_id_load.py" diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/helper_helper_rocke.instances.gfx942.attention_tiled_2d.h b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/helper_helper_rocke.instances.gfx942.attention_tiled_2d.h index de8da22bdbe9..85e10caefd4d 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/helper_helper_rocke.instances.gfx942.attention_tiled_2d.h +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/helper_helper_rocke.instances.gfx942.attention_tiled_2d.h @@ -203,7 +203,8 @@ int rocke_attention_tiled_2d_spec_n_blocks_per_tile(const rocke_attention_tiled_ /* dtype_ir: F16 for "fp16", BF16 otherwise (returns a rocke_type_t*). */ const rocke_type_t* rocke_attention_tiled_2d_spec_dtype_ir(const rocke_attention_tiled_2d_spec_t* s); -/* binary_search_iters: 32 if num_seqs<=0, else max(1, ceil(log2(num_seqs+1))) */ +/* binary_search_iters: 32 if num_seqs<=0, else max(1, ceil(log2(num_seqs+1))). + * Exported for existing callers; the kernels derive the search at runtime. */ int rocke_attention_tiled_2d_spec_binary_search_iters(const rocke_attention_tiled_2d_spec_t* s); /* __post_init__ validator (builder variant). Reproduces the Python validation diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/instance_gfx942_attention_tiled_2d_internal.h b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/instance_gfx942_attention_tiled_2d_internal.h index 0b28fcf6f08e..c2d36242dee6 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/instance_gfx942_attention_tiled_2d_internal.h +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/instance_gfx942_attention_tiled_2d_internal.h @@ -177,14 +177,10 @@ typedef struct rocke_gfx942_attn2d_build_ctx rocke_value_t* num_seqs_p; rocke_value_t* bt_stride_p; rocke_value_t* qq_bias_stride0_p; - /* RUNTIME_KERNARG_FIELDS params, in declaration order; set only when - * runtime_shape. */ + /* RUNTIME_KERNARG_FIELDS params, in declaration order. */ rocke_value_t* num_qh_p; rocke_value_t* num_kv_p; rocke_value_t* sliding_window_p; - /* The gfx942 Python body: it reads its runtime kernargs and searches at - * runtime. False when this body stands in for the gfx950 builder. */ - bool runtime_shape; /* ---- grid ids + wave decomposition (lines 1371-1388) ---- */ rocke_value_t* q_block_global_idx; @@ -289,8 +285,8 @@ typedef struct rocke_gfx942_attn2d_build_ctx rocke_value_t* seq_base; /* to_sgpr_u32(seq_idx*bt_stride) */ rocke_value_t* block_table_max_idx; /* to_sgpr_u32(num_seqs*bt_stride) */ rocke_value_t* kv_block_bytes_c_v; /* one-block bound: kv_stride_blk_b_v, or its const */ - rocke_value_t* kv_stride_tok_b_v; /* num_kv_p*HD*KV_BYTES (runtime_shape only) */ - rocke_value_t* kv_stride_blk_b_v; /* kv_stride_tok_b_v*BS (runtime_shape only) */ + rocke_value_t* kv_stride_tok_b_v; /* num_kv_p*HD*KV_BYTES */ + rocke_value_t* kv_stride_blk_b_v; /* kv_stride_tok_b_v*BS */ rocke_value_t* lane_half_base_v; /* tid * KV_HALVES_PER_LANE (2232) */ rocke_value_t* zero_soff_v; /* const_i32(0) soffset (2238) */ rocke_value_t* wave_lds_off_i64_v; /* wave_lds_offset_i64 (2254/2263) */ @@ -370,13 +366,6 @@ typedef struct rocke_gfx942_attn2d_build_ctx rocke_value_t* l_nonzero[ROCKE_GFX942_ATTN2D_MAX_REGS_PER_LANE]; } rocke_gfx942_attn2d_build_ctx_t; -/* The query head count as a row bound: the kernarg under runtime_shape, else - * the constant the gfx950 builder bakes, created at each use as Python does. */ -static inline rocke_value_t* rocke_gfx942_attn2d_num_qh(rocke_gfx942_attn2d_build_ctx_t* ctx) -{ - return ctx->runtime_shape ? ctx->num_qh_p : rocke_b_const_i32(ctx->b, ctx->NUM_QH); -} - /* ===================================================================== * * DRIVER-INTERNAL ctx population (the build prologue, lines 1104-1912). * diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/instance_gfx942_attention_tiled_3d.h b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/instance_gfx942_attention_tiled_3d.h index d0537f21aaf2..2e49daf862da 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/instance_gfx942_attention_tiled_3d.h +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/instance_gfx942_attention_tiled_3d.h @@ -202,6 +202,10 @@ int rocke_unified_attention_3d_tiled_spec_tile_size( /* dtype_ir: F16 for "fp16", BF16 otherwise (returns a rocke_type_t*). */ const rocke_type_t* rocke_unified_attention_3d_tiled_spec_dtype_ir( const rocke_unified_attention_3d_tiled_spec_t* s); +/* binary_search_iters: 32 if num_seqs<=0, else max(1, ceil(log2(num_seqs+1))). + * Exported for existing callers; the kernels derive the search at runtime. */ +int rocke_unified_attention_3d_tiled_spec_binary_search_iters( + const rocke_unified_attention_3d_tiled_spec_t* s); /* kernel_name(): kernel_name_join(...) into caller's buf (NUL-terminated, up to * cap bytes incl. NUL). Returns the length that would be written (snprintf * convention), or -1 on a dead builder/encode error. */ diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/instance_gfx950_attention_tiled_2d_internal.h b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/instance_gfx950_attention_tiled_2d_internal.h index 3e5902944206..9f48be1dae7a 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/instance_gfx950_attention_tiled_2d_internal.h +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/instance_gfx950_attention_tiled_2d_internal.h @@ -87,6 +87,9 @@ extern "C" { #define ROCKE_GFX950_ATTN2D_MAX_ACCS 64 /* Online-softmax iter-args: 2*SOFTMAX_STATE_SLOTS (m,l per slot) + ACCS + cur_buf. */ #define ROCKE_GFX950_ATTN2D_MAX_ITER_ARGS 96 +/* Most pages per tile whose ids the KV loop carries; larger tiles read them in + * place (Python's PAGE_PREFETCH applies the same bound). */ +#define ROCKE_GFX950_ATTN2D_MAX_PREFETCH_PAGES 16 /* Per-lane register slots (REGS_PER_LANE): 4 (M=16), 8 (M=32), 16 (32x32). */ #define ROCKE_GFX950_ATTN2D_MAX_REGS_PER_LANE 16 /* QK / PV K-iters (HD/QK_K_STEP or T/PV_K_STEP); HD<=256, T<=128 -> <=16. */ @@ -117,7 +120,7 @@ typedef struct rocke_gfx950_attn2d_build_ctx /* ---- ALL-CAPS geometry constants (lines 754-847) ---- */ int HD, T, BS, N_BLOCKS_PER_TILE; - int BLOCK_M, BLOCK_Q, NQK, NUM_KV, NUM_QH; + int BLOCK_M, BLOCK_Q, NQK, NUM_KV; /* NUM_KV only for the fast-desc 64/8 pin */ int VALID_ROWS; /* BLOCK_Q * NQK -- rows that map to complete GQA groups */ int SLIDING_WINDOW; bool USE_SOFTCAP, USE_SINKS, USE_ALIBI, USE_QQ_BIAS; @@ -175,6 +178,10 @@ typedef struct rocke_gfx950_attn2d_build_ctx rocke_value_t* num_seqs_p; rocke_value_t* bt_stride_p; rocke_value_t* qq_bias_stride0_p; + /* RUNTIME_KERNARG_FIELDS params, in declaration order. */ + rocke_value_t* num_qh_p; + rocke_value_t* num_kv_p; + rocke_value_t* sliding_window_p; /* ---- grid ids + wave decomposition ---- */ rocke_value_t* q_block_global_idx; @@ -246,21 +253,40 @@ typedef struct rocke_gfx950_attn2d_build_ctx rocke_value_t* lane_rg_is2_v; /* ---- SSA constants ---- */ - rocke_value_t* c0; rocke_value_t* zero_f; rocke_value_t* neg_inf_v; /* b.const_f32(-inf) */ rocke_value_t* one_f_v; /* b.const_f32(1.0) */ rocke_value_t* rcp_ln2_v; /* b.const_f32(1.4426950408889634) */ rocke_value_t* qk_scale_v; /* derived from scale_p */ rocke_value_t* pv_fp8_scale_v; /* fdiv(v_scale, 240) when FP8_MFMA_PV (line 1149) */ - rocke_value_t* sw_const_v; /* b.const_i32(SLIDING_WINDOW) */ + rocke_value_t* sw_const_v; /* the sliding_window kernarg */ /* ---- paged-KV byte descriptor (full transform DAG, lines 1163-1630) ---- */ rocke_tensor_descriptor_t* q_desc; /* output/query [token,head,dim] */ rocke_tensor_descriptor_t* kv_desc; /* paged K/V byte descriptor */ rocke_value_t* seq_base; /* to_sgpr_u32(seq_idx*bt_stride) */ rocke_value_t* block_table_max_idx; /* to_sgpr_u32(num_seqs*bt_stride) */ - rocke_value_t* kv_block_bytes_c_v; /* const BS*NUM_KV*HD*KV_BYTES */ + rocke_value_t* kv_stride_tok_b_v; /* num_kv_p*HD*KV_BYTES */ + rocke_value_t* kv_stride_blk_b_v; /* kv_stride_tok_b_v*BS */ + rocke_value_t* kv_block_bytes_c_v; /* kv_stride_blk_b_v (one block) */ + /* bf16/fp16 loaders (!KV_FP8, !FAST_PAGED_KV_DESC): the lane's in-page byte + * offset, each copy's first-token offset in its page (indexed by that + * token; [0] stays NULL), and the i64 page stride (I64_KV_ADDR). */ + rocke_value_t* lane_within_b_v; + rocke_value_t* copy_tok_b[64]; + rocke_value_t* kv_stride_blk_i64_v; + /* Page ids the KV loop carries when page copies prefetch them (one per page + * of a tile), with the tile index value each set belongs to: the first tile + * and the one after it (prologue), the current and next tile, and the tile + * after next read one iteration ahead. */ + rocke_value_t* page_start_tile; + rocke_value_t* page_start[ROCKE_GFX950_ATTN2D_MAX_PREFETCH_PAGES]; + rocke_value_t* page_start_next[ROCKE_GFX950_ATTN2D_MAX_PREFETCH_PAGES]; + rocke_value_t* page_cur_tile; + rocke_value_t* page_cur[ROCKE_GFX950_ATTN2D_MAX_PREFETCH_PAGES]; + rocke_value_t* page_next_tile; + rocke_value_t* page_next[ROCKE_GFX950_ATTN2D_MAX_PREFETCH_PAGES]; + rocke_value_t* page_after_next[ROCKE_GFX950_ATTN2D_MAX_PREFETCH_PAGES]; rocke_value_t* lane_half_base_v; /* tid * KV_HALVES_PER_LANE */ rocke_value_t* zero_soff_v; /* const_i32(0) soffset */ rocke_value_t* wave_lds_off_i64_v; /* wave_lds_offset_i64 */ @@ -434,6 +460,13 @@ void rocke_gfx950_attn2d_issue_k_load_runtime(rocke_gfx950_attn2d_build_ctx_t* c void rocke_gfx950_attn2d_issue_v_load_runtime(rocke_gfx950_attn2d_build_ctx_t* ctx, rocke_value_t* kv_tile_idx, rocke_value_t* buf_idx); +/* True when the bf16/fp16 loaders use per-page copy offsets and the KV loop + * carries the page ids, reading each tile's ids one iteration ahead. */ +bool rocke_gfx950_attn2d_page_prefetch(const rocke_gfx950_attn2d_build_ctx_t* ctx); +/* The page ids of every page of a tile: one guarded block-table read each. */ +void rocke_gfx950_attn2d_tile_page_ids(rocke_gfx950_attn2d_build_ctx_t* ctx, + rocke_value_t* kv_tile_idx, + rocke_value_t** out); void rocke_gfx950_attn2d_issue_k(rocke_gfx950_attn2d_build_ctx_t* ctx, rocke_value_t* tile_idx, rocke_value_t* buf_idx); diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/instance_gfx950_attention_tiled_3d.h b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/instance_gfx950_attention_tiled_3d.h index 6506e70467d6..9f22f935f38e 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/instance_gfx950_attention_tiled_3d.h +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/instance_gfx950_attention_tiled_3d.h @@ -280,7 +280,8 @@ int rocke_gfx950_unified_attention_3d_tiled_spec_tile_size( /* dtype_ir: F16 for "fp16", BF16 otherwise (returns a rocke_type_t*). */ const rocke_type_t* rocke_gfx950_unified_attention_3d_tiled_spec_dtype_ir( const rocke_unified_attention_3d_tiled_spec_t* s); -/* binary_search_iters: 32 if num_seqs<=0, else max(1, ceil(log2(num_seqs+1))) */ +/* binary_search_iters: 32 if num_seqs<=0, else max(1, ceil(log2(num_seqs+1))). + * Exported for existing callers; the kernels derive the search at runtime. */ int rocke_gfx950_unified_attention_3d_tiled_spec_binary_search_iters( const rocke_unified_attention_3d_tiled_spec_t* s); /* kernel_name(): kernel_name_join(...) into caller's buf (NUL-terminated, up to diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/instance_gfx950_attention_tiled_3d_internal.h b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/instance_gfx950_attention_tiled_3d_internal.h index 4363fedd6162..c768f6c013b5 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/instance_gfx950_attention_tiled_3d_internal.h +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/instance_gfx950_attention_tiled_3d_internal.h @@ -129,8 +129,6 @@ typedef struct rocke_gfx950_attn_tiled_3d_config int BLOCK_M; /* spec.block_m (== 16) */ int BLOCK_Q; /* spec.block_q */ int NQK; /* spec.num_queries_per_kv */ - int NUM_KV; /* spec.num_kv_heads */ - int NUM_QH; /* spec.num_query_heads */ int NUM_SEG; /* spec.num_segments */ int SLIDING_WINDOW; /* spec.sliding_window */ bool USE_SOFTCAP; /* spec.has_softcap */ @@ -153,7 +151,6 @@ typedef struct rocke_gfx950_attn_tiled_3d_config int PV_N_TILES; /* HD // MFMA_N */ int THREADS; /* 64 */ - int binary_search_iters; /* spec.binary_search_iters */ /* ---- async / fp8 KV feed geometry (lines 565-643) ---- */ int ASYNC_LDS_DWORDS; /* 4 (CDNA4) */ @@ -161,8 +158,6 @@ typedef struct rocke_gfx950_attn_tiled_3d_config int KV_HALVES_PER_CALL; /* THREADS * 8 */ int kv_calls_per_tile; /* (T*HD) // KV_HALVES_PER_CALL */ int bytes_per_call; /* KV_HALVES_PER_CALL * 2 */ - int kv_stride_blk_b; /* BS*NUM_KV*HD*KV_BYTES */ - int kv_stride_tok_b; /* NUM_KV*HD*KV_BYTES */ int kv_stride_h_b; /* HD*KV_BYTES */ int bytes_per_buf; /* T*HD*2 */ int N_BLOCKS_PER_TILE; /* T // BS (multi-page tile when > 1) */ @@ -252,6 +247,10 @@ typedef struct rocke_gfx950_attention_tiled_3d_build_ctx rocke_value_t* num_seqs_p; /* I32 num_seqs */ rocke_value_t* bt_stride_p; /* I32 block_table_stride */ rocke_value_t* qq_bias_stride0_p; /* I32 qq_bias_stride_0 */ + /* RUNTIME_KERNARG_FIELDS params, in declaration order. */ + rocke_value_t* num_qh_p; /* I32 num_query_heads */ + rocke_value_t* num_kv_p; /* I32 num_kv_heads */ + rocke_value_t* sliding_window_p; /* I32 sliding_window */ /* ---------- grid ids + thread (lines 360-363) ---------- */ rocke_value_t* q_block_global_idx; /* block_id_x() */ @@ -284,7 +283,7 @@ typedef struct rocke_gfx950_attention_tiled_3d_build_ctx rocke_value_t* one_f; /* const_f32(1.0) */ rocke_value_t* rcp_ln2; /* const_f32(1.4426950408889634) */ rocke_value_t* qk_scale; /* scale_p * rcp_ln2 */ - rocke_value_t* sw_const; /* const_i32(SLIDING_WINDOW) */ + rocke_value_t* sw_const; /* the sliding_window kernarg */ rocke_value_t* z8; /* zero_vec(dtype, 8) */ /* ---------- lane decode (lines 533-534) ---------- */ @@ -326,8 +325,14 @@ typedef struct rocke_gfx950_attention_tiled_3d_build_ctx rocke_value_t* first_page_ids[8]; /* tile_start's page ids (PAGE_IDS; N <= 8) */ rocke_value_t* seg_first_page; /* tile_start * N_BLOCKS_PER_TILE (WAVE_PAGE_ID_LOAD) */ rocke_value_t* first_page_id_window; /* the first 64-page id window (WAVE_PAGE_ID_LOAD) */ - rocke_value_t* - kv_block_bytes_c; /* const_i32(kv_stride_blk_b): 1-block buffer bound (i64 path) */ + /* K/V byte strides; they scale with the runtime kv head count. */ + rocke_value_t* kv_stride_tok_b; /* num_kv_p * HD * KV_BYTES */ + rocke_value_t* kv_stride_blk_b; /* kv_stride_tok_b * BS */ + rocke_value_t* kv_block_bytes_c; /* kv_stride_blk_b: 1-block buffer bound (i64 path) */ + /* Each call's token offset within a page (PAGE_IDS), indexed by + * tok0 / tokens_per_call; [0] stays NULL (no offset). */ + rocke_value_t* call_tok_b[64]; + rocke_value_t* kv_stride_blk_i64; /* zext(kv_stride_blk_b) (PAGE_IDS && I64_KV_ADDR) */ /* ---------- online-softmax loop init carry (lines 536-562, 717-718) ---------- * * m_init: the lane's query row (l starts at one_f). acc_inits: PV_N_TILES diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/gfx942/attention_tiled_2d_ctx_prologue.cpp b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/gfx942/attention_tiled_2d_ctx_prologue.cpp index 89a5016cdabb..28ae05284e3f 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/gfx942/attention_tiled_2d_ctx_prologue.cpp +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/gfx942/attention_tiled_2d_ctx_prologue.cpp @@ -117,9 +117,6 @@ bool rocke_gfx942_attn2d_build_ctx_init(rocke_gfx942_attn2d_build_ctx_t* ctx, ctx->b = b; ctx->spec = spec; ctx->arch = arch; - /* This body also stands in for the gfx950 Python builder (the fastKV - * register-P twin), which bakes the shape values gfx942 reads at runtime. */ - ctx->runtime_shape = rocke_a2d_streq(arch, "gfx942"); /* ---- require_tiled_attention_arch(arch) (lines 1139-1141) ---- */ if(!rocke_a2d_require_tiled_attention_arch(b, arch)) @@ -401,12 +398,9 @@ bool rocke_gfx942_attn2d_build_ctx_init(rocke_gfx942_attn2d_build_ctx_t* ctx, ctx->num_seqs_p = rocke_b_param(b, "num_seqs", rocke_i32(), NULL); ctx->bt_stride_p = rocke_b_param(b, "block_table_stride", rocke_i32(), NULL); ctx->qq_bias_stride0_p = rocke_b_param(b, "qq_bias_stride_0", rocke_i32(), NULL); - if(ctx->runtime_shape) - { - ctx->num_qh_p = rocke_b_param(b, "num_query_heads", rocke_i32(), NULL); - ctx->num_kv_p = rocke_b_param(b, "num_kv_heads", rocke_i32(), NULL); - ctx->sliding_window_p = rocke_b_param(b, "sliding_window", rocke_i32(), NULL); - } + ctx->num_qh_p = rocke_b_param(b, "num_query_heads", rocke_i32(), NULL); + ctx->num_kv_p = rocke_b_param(b, "num_kv_heads", rocke_i32(), NULL); + ctx->sliding_window_p = rocke_b_param(b, "sliding_window", rocke_i32(), NULL); } /* ---- grid ids + wave decomposition (lines 1371-1388) ---- */ @@ -436,23 +430,11 @@ bool rocke_gfx942_attn2d_build_ctx_init(rocke_gfx942_attn2d_build_ctx_t* ctx, } /* ---- seq lookup + Q-block geometry (lines 1390-1409) ---- */ - if(ctx->runtime_shape) { rocke_value_t* trip = rocke_binary_search_seq_idx_trip_count(b, ctx->num_seqs_p); ctx->seq_idx = rocke_binary_search_seq_idx_dyn( b, ctx->cu_q, ctx->q_block_global_idx, ctx->num_seqs_p, BLOCK_Q, trip, false); } - else - { - ctx->seq_idx - = rocke_binary_search_seq_idx(b, - ctx->cu_q, - ctx->q_block_global_idx, - ctx->num_seqs_p, - BLOCK_Q, - rocke_attention_tiled_2d_spec_binary_search_iters(spec), - false); - } ctx->cu_q_start = rocke_b_global_load_i32(b, ctx->cu_q, ctx->seq_idx, 0); ctx->cu_q_stop = rocke_b_global_load_i32( b, ctx->cu_q, rocke_b_add(b, ctx->seq_idx, rocke_b_const_i32(b, 1)), 0); @@ -748,13 +730,9 @@ bool rocke_gfx942_attn2d_build_ctx_init(rocke_gfx942_attn2d_build_ctx_t* ctx, qk_scale = rocke_b_fmul(b, qk_scale, ctx->k_scale_p); ctx->qk_scale_v = qk_scale; } - /* sw_const = b.const_i32(int(SLIDING_WINDOW)) (line 1910). Python creates - * this once here in the constants block and REUSES it in the tile-bound - * (line 1994) and mask (lines 4111/4328/4403) regions. The previous port - * created an unused i32 0 here and a FRESH const in the loop scaffold, - * allocating an extra SSA value that shifted every later %N. Cache it. */ - ctx->sw_const_v - = ctx->runtime_shape ? ctx->sliding_window_p : rocke_b_const_i32(b, ctx->SLIDING_WINDOW); + /* sw_const = rt["sliding_window"]. Python binds it once here in the + * constants block and REUSES it in the tile-bound and mask regions. */ + ctx->sw_const_v = ctx->sliding_window_p; /* ---- acc geometry the epilogue aliases (lines 2142-2143) ---- */ ctx->ACC_N_TILES = USE_MFMA_32X32 ? (HD / 32) : PV_N_TILES; @@ -1012,7 +990,7 @@ void rocke_gfx942_attn2d_emit_licm_hoist(rocke_gfx942_attn2d_build_ctx_t* ctx) /* row_ok = (qp_r < q_len) && (qh_r < NUM_QH) [&& (row < VALID_ROWS)]. * Skip the VALID_ROWS guard when VALID_ROWS == BLOCK_M (NQK divides BLOCK_M). */ rocke_value_t* row_ok_pos = rocke_b_cmp_lt(b, qp_r, ctx->cur_batch_q_len); - rocke_value_t* row_ok_qh = rocke_b_cmp_lt(b, qh_r, rocke_gfx942_attn2d_num_qh(ctx)); + rocke_value_t* row_ok_qh = rocke_b_cmp_lt(b, qh_r, ctx->num_qh_p); rocke_value_t* row_ok_inner = rocke_b_land(b, row_ok_pos, row_ok_qh); rocke_value_t* row_ok; if(ctx->VALID_ROWS < ctx->BLOCK_M) @@ -1040,7 +1018,7 @@ void rocke_gfx942_attn2d_emit_licm_hoist(rocke_gfx942_attn2d_build_ctx_t* ctx) for(reg = 0; reg < ctx->REGS_PER_LANE; ++reg) { rocke_value_t* qh_r = ctx->hoist_q_head[reg]; - rocke_value_t* qh_ok = rocke_b_cmp_lt(b, qh_r, rocke_gfx942_attn2d_num_qh(ctx)); + rocke_value_t* qh_ok = rocke_b_cmp_lt(b, qh_r, ctx->num_qh_p); rocke_value_t* slope = rocke_b_masked_global_load( b, ctx->alibi_slopes_ptr, qh_r, qh_ok, rocke_b_const_f32(b, 0.0), f32, 4); /* The softmax bucket reads the per-reg slope from hoist_q_head. */ diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/gfx942/attention_tiled_2d_epilogue.cpp b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/gfx942/attention_tiled_2d_epilogue.cpp index 56c5891f1389..3cf0ce4417bf 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/gfx942/attention_tiled_2d_epilogue.cpp +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/gfx942/attention_tiled_2d_epilogue.cpp @@ -128,7 +128,7 @@ rocke_kernel_def_t* rocke_gfx942_attn2d_emit_epilogue(rocke_gfx942_attn2d_build_ rocke_value_t* op_qh_mod = rocke_b_mod(b, q_row_t, rocke_b_const_i32(b, NQK)); rocke_value_t* op_qh_t = rocke_b_add(b, op_qh_mul, op_qh_mod); rocke_value_t* op_ok_pos = rocke_b_cmp_lt(b, op_pos_t, cur_batch_q_len); - rocke_value_t* op_ok_qh = rocke_b_cmp_lt(b, op_qh_t, rocke_gfx942_attn2d_num_qh(ctx)); + rocke_value_t* op_ok_qh = rocke_b_cmp_lt(b, op_qh_t, ctx->num_qh_p); rocke_value_t* op_ok_inner = rocke_b_land(b, op_ok_pos, op_ok_qh); /* Skip the VALID_ROWS guard when VALID_ROWS == BLOCK_M (NQK divides BLOCK_M). */ rocke_value_t* op_mask_t; @@ -204,8 +204,7 @@ rocke_kernel_def_t* rocke_gfx942_attn2d_emit_epilogue(rocke_gfx942_attn2d_build_ rocke_value_t* op_qh32_base = rocke_b_add(b, op_qh32_mul, op_qh32_mod); /* Python emits the pos cmp BEFORE the qh cmp, then AND's VALID_ROWS; bind in order. */ rocke_value_t* op_ok_pos32 = rocke_b_cmp_lt(b, op_pos32_base, cur_batch_q_len); - rocke_value_t* op_ok_qh32 - = rocke_b_cmp_lt(b, op_qh32_base, rocke_gfx942_attn2d_num_qh(ctx)); + rocke_value_t* op_ok_qh32 = rocke_b_cmp_lt(b, op_qh32_base, ctx->num_qh_p); rocke_value_t* op_ok_inner32 = rocke_b_land(b, op_ok_pos32, op_ok_qh32); /* Skip the VALID_ROWS guard when VALID_ROWS == BLOCK_M (NQK divides BLOCK_M). */ rocke_value_t* op_mask32_base; @@ -303,7 +302,7 @@ rocke_kernel_def_t* rocke_gfx942_attn2d_emit_epilogue(rocke_gfx942_attn2d_build_ rocke_value_t* op_qh_mod = rocke_b_mod(b, OUT_ROW_BASE, rocke_b_const_i32(b, NQK)); rocke_value_t* op_qh = rocke_b_add(b, op_qh_mul, op_qh_mod); rocke_value_t* op_mask_pos = rocke_b_cmp_lt(b, op_pos, cur_batch_q_len); - rocke_value_t* op_mask_qh = rocke_b_cmp_lt(b, op_qh, rocke_gfx942_attn2d_num_qh(ctx)); + rocke_value_t* op_mask_qh = rocke_b_cmp_lt(b, op_qh, ctx->num_qh_p); rocke_value_t* op_mask_inner = rocke_b_land(b, op_mask_pos, op_mask_qh); /* Skip the VALID_ROWS guard when VALID_ROWS == BLOCK_M (NQK divides BLOCK_M). */ rocke_value_t* op_mask; diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/gfx942/attention_tiled_2d_kv_body_qk_softmax.cpp b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/gfx942/attention_tiled_2d_kv_body_qk_softmax.cpp index fd2a488a75af..f3da0c3c74ee 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/gfx942/attention_tiled_2d_kv_body_qk_softmax.cpp +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/gfx942/attention_tiled_2d_kv_body_qk_softmax.cpp @@ -378,7 +378,7 @@ void rocke_gfx942_attn2d_emit_kv_body(rocke_gfx942_attn2d_build_ctx_t* ctx) /* row_ok = land(qpcur_batch_q_len); - rocke_value_t* st_ok_qh = rocke_b_cmp_lt(b, st_qh_iter, rocke_gfx942_attn2d_num_qh(ctx)); + rocke_value_t* st_ok_qh = rocke_b_cmp_lt(b, st_qh_iter, ctx->num_qh_p); rocke_value_t* st_ok_inner = rocke_b_land(b, st_ok_qp, st_ok_qh); if(ctx->VALID_ROWS < ctx->BLOCK_M) { @@ -393,8 +393,7 @@ void rocke_gfx942_attn2d_emit_kv_body(rocke_gfx942_attn2d_build_ctx_t* ctx) st_causal_lim_iter = rocke_b_add(b, ctx->context_len, st_qp_iter); if(ctx->USE_ALIBI) { - rocke_value_t* st_qh_iter_ok - = rocke_b_cmp_lt(b, st_qh_iter, rocke_gfx942_attn2d_num_qh(ctx)); + rocke_value_t* st_qh_iter_ok = rocke_b_cmp_lt(b, st_qh_iter, ctx->num_qh_p); st_alibi_slope_iter = rocke_b_masked_global_load(b, ctx->alibi_slopes_ptr, st_qh_iter, @@ -588,8 +587,7 @@ void rocke_gfx942_attn2d_emit_kv_body(rocke_gfx942_attn2d_build_ctx_t* ctx) /* row_ok = land(qpcur_batch_q_len); - rocke_value_t* ok_qh - = rocke_b_cmp_lt(b, qh_r, rocke_gfx942_attn2d_num_qh(ctx)); + rocke_value_t* ok_qh = rocke_b_cmp_lt(b, qh_r, ctx->num_qh_p); rocke_value_t* ok_inner = rocke_b_land(b, ok_pos, ok_qh); if(ctx->VALID_ROWS < ctx->BLOCK_M) { diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/gfx942/attention_tiled_2d_kv_descriptors_loads.cpp b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/gfx942/attention_tiled_2d_kv_descriptors_loads.cpp index 12775a5a4118..277b707ce3cc 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/gfx942/attention_tiled_2d_kv_descriptors_loads.cpp +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/gfx942/attention_tiled_2d_kv_descriptors_loads.cpp @@ -104,7 +104,7 @@ static int rocke__v_bytes_per_call_swz(const rocke_gfx942_attn2d_build_ctx_t* ct * would shift downstream %values). */ static rocke_value_t* rocke__kv_block_bytes_c(rocke_gfx942_attn2d_build_ctx_t* ctx) { - if(ctx->kv_block_bytes_c_v == NULL && ctx->runtime_shape) + if(ctx->kv_block_bytes_c_v == NULL) { /* The token and block strides scale with the runtime kv head count. */ ctx->kv_stride_tok_b_v = rocke_b_mul( @@ -113,11 +113,6 @@ static rocke_value_t* rocke__kv_block_bytes_c(rocke_gfx942_attn2d_build_ctx_t* c = rocke_b_mul(ctx->b, ctx->kv_stride_tok_b_v, rocke_b_const_i32(ctx->b, ctx->BS)); ctx->kv_block_bytes_c_v = ctx->kv_stride_blk_b_v; } - else if(ctx->kv_block_bytes_c_v == NULL) - { - int kv_stride_blk_b = ctx->BS * ctx->NUM_KV * ctx->HD * ctx->KV_BYTES; - ctx->kv_block_bytes_c_v = rocke_b_const_i32(ctx->b, kv_stride_blk_b); - } return ctx->kv_block_bytes_c_v; } @@ -397,25 +392,12 @@ void rocke_gfx942_attn2d_emit_preloop(rocke_gfx942_attn2d_build_ctx_t* ctx) * lengths=[1<<24, BS, NUM_KV, HD], * strides=[kv_stride_blk_b, kv_stride_tok_b, kv_stride_h_b, KV_BYTES], * coord_names=("physical_block","token","kv_head","dim")). */ - int blk_b = ctx->BS * ctx->NUM_KV * ctx->HD * ctx->KV_BYTES; - int tok_b = ctx->NUM_KV * ctx->HD * ctx->KV_BYTES; - int h_b = ctx->HD * ctx->KV_BYTES; - int kv_lengths[4] = {1 << 24, ctx->BS, ctx->NUM_KV, ctx->HD}; - int kv_strides[4] = {blk_b, tok_b, h_b, ctx->KV_BYTES}; + int kv_strides[4] = {1, 1, ctx->HD * ctx->KV_BYTES, ctx->KV_BYTES}; static const char* const kv_coords[4] = {"physical_block", "token", "kv_head", "dim"}; - rocke_tensor_descriptor_t* kv_base; - if(ctx->runtime_shape) - { - rocke_value_t* dyn[4] = {ctx->kv_stride_blk_b_v, ctx->kv_stride_tok_b_v, NULL, NULL}; - kv_base = &rocke_tensor_descriptor_naive_dynamic_mixed( - b, "paged_kv_bytes", kv_coords, 4, dyn, kv_strides) - ->base; - } - else - { - kv_base = rocke_tensor_descriptor_naive( - b, "paged_kv_bytes", kv_lengths, 4, kv_strides, kv_coords, 4); - } + rocke_value_t* dyn[4] = {ctx->kv_stride_blk_b_v, ctx->kv_stride_tok_b_v, NULL, NULL}; + rocke_tensor_descriptor_t* kv_base = &rocke_tensor_descriptor_naive_dynamic_mixed( + b, "paged_kv_bytes", kv_coords, 4, dyn, kv_strides) + ->base; /* Single-block tile (N_BLOCKS_PER_TILE == 1, T == BS): * indirect(tile_idx -> physical_block) ; unmerge(linear_half -> token,dim). @@ -555,8 +537,9 @@ static void rocke__fast_paged_kv_voff_split(rocke_gfx942_attn2d_build_ctx_t* ctx (void)base_ptr_src; if(ctx->I64_KV_ADDR) { - rocke_value_t* base_i64 - = rocke_b_shl(b, rocke_b_zext(b, physical, rocke_i64()), rocke_b_const_i64(b, 15)); + /* zext before the shift constant, as Python evaluates them. */ + rocke_value_t* physical_i64 = rocke_b_zext(b, physical, rocke_i64()); + rocke_value_t* base_i64 = rocke_b_shl(b, physical_i64, rocke_b_const_i64(b, 15)); if(out_base_i64) *out_base_i64 = base_i64; if(out_voff) @@ -623,7 +606,6 @@ static bool rocke__tile_pages(rocke_gfx942_attn2d_build_ctx_t* ctx, snprintf(b->err, ROCKE_ERR_MSG_CAP, "out of memory for %d tile pages", n_pages); return false; } - int64_t kv_stride_blk_b = (int64_t)ctx->BS * ctx->NUM_KV * ctx->HD * ctx->KV_BYTES; rocke_value_t* seq_base = rocke__seq_base(ctx); rocke_value_t* max_idx = rocke__block_table_max_idx(ctx); rocke_value_t* tile_page0 = rocke_b_mul(b, kv_tile_idx, rocke_b_const_i32(b, n_pages)); @@ -636,9 +618,7 @@ static bool rocke__tile_pages(rocke_gfx942_attn2d_build_ctx_t* ctx, rocke_value_t* physical = rocke_b_masked_global_load(b, ctx->block_tables, idx, mask, def, rocke_i32(), 4); rocke_value_t* page_i64 = rocke_b_zext(b, rocke_b_to_sgpr_u32(b, physical), rocke_i64()); - rocke_value_t* stride64 = ctx->runtime_shape - ? rocke_b_zext(b, ctx->kv_stride_blk_b_v, rocke_i64()) - : rocke_b_const_i64(b, kv_stride_blk_b); + rocke_value_t* stride64 = rocke_b_zext(b, ctx->kv_stride_blk_b_v, rocke_i64()); rocke_value_t* base = rocke_b_mul(b, page_i64, stride64); ptrs[page] = rocke_b_global_ptr_add(b, src, base); } @@ -662,11 +642,7 @@ static rocke_value_t* rocke_value_t* half = rocke_b_add(b, call_half, rocke__lane_half_base(ctx)); rocke_value_t* token = rocke_b_div(b, half, rocke_b_const_i32(b, ctx->HD)); rocke_value_t* dim = rocke_b_mod(b, half, rocke_b_const_i32(b, ctx->HD)); - rocke_value_t* token_b = rocke_b_mul( - b, - token, - ctx->runtime_shape ? ctx->kv_stride_tok_b_v - : rocke_b_const_i32(b, ctx->NUM_KV * ctx->HD * ctx->KV_BYTES)); + rocke_value_t* token_b = rocke_b_mul(b, token, ctx->kv_stride_tok_b_v); rocke_value_t* token_head_b = rocke_b_add(b, token_b, kv_head_b); rocke_value_t* dim_b = rocke_b_mul(b, dim, rocke_b_const_i32(b, ctx->KV_BYTES)); return rocke_b_add(b, token_head_b, dim_b); diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/gfx942/attention_tiled_2d_loop_scaffold.cpp b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/gfx942/attention_tiled_2d_loop_scaffold.cpp index 0ba9755606b9..c3edf6420ac4 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/gfx942/attention_tiled_2d_loop_scaffold.cpp +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/gfx942/attention_tiled_2d_loop_scaffold.cpp @@ -172,9 +172,24 @@ void rocke_gfx942_attn2d_emit_loop_bounds_and_inits(rocke_gfx942_attn2d_build_ct rocke_value_t* m_inits[ROCKE_GFX942_ATTN2D_MAX_REGS_PER_LANE]; rocke_value_t* l_inits[ROCKE_GFX942_ATTN2D_MAX_REGS_PER_LANE]; - if(ctx->USE_SINKS) + if(ctx->USE_SINKS && ctx->USE_MFMA_32X32 && ctx->TRANSPOSED_QK_32X32 + && ctx->TRANSPOSED_SCALAR_STATE) { - /* Narrow gfx942 path (the transposed-scalar 32x32 branch is gfx950). */ + /* Transposed-scalar-state branch, reached by the gfx950 fastKV + * register-P twin: one state slot whose row is wave_row_base + + * lane_col32. */ + rocke_value_t* row = rocke_b_add(B, ctx->wave_row_base, ctx->lane_col32_v); + rocke_value_t* qh_mul = rocke_b_mul(B, ctx->kv_head_idx, rocke_b_const_i32(B, NQK)); + rocke_value_t* qh_mod = rocke_b_mod(B, row, rocke_b_const_i32(B, NQK)); + rocke_value_t* qh = rocke_b_add(B, qh_mul, qh_mod); + rocke_value_t* qh_in = rocke_b_cmp_lt(B, qh, ctx->num_qh_p); + rocke_value_t* sink_h = rocke_b_global_load(B, ctx->sinks, qh, ctx->dtype, 2); + rocke_value_t* sink_f = rocke_b_fmul(B, rocke_b_cast_to_f32(B, sink_h), rcp_ln2); + m_inits[0] = rocke_b_select(B, qh_in, sink_f, neg_inf); + } + else if(ctx->USE_SINKS) + { + /* Narrow gfx942 path. */ for(int r = 0; r < SOFTMAX_STATE_SLOTS; ++r) { rocke_value_t* row = rocke_b_add(B, ctx->wave_row_base, sc_state_row(ctx, lane_rg, r)); @@ -183,7 +198,7 @@ void rocke_gfx942_attn2d_emit_loop_bounds_and_inits(rocke_gfx942_attn2d_build_ct rocke_value_t* qh_mul = rocke_b_mul(B, ctx->kv_head_idx, rocke_b_const_i32(B, NQK)); rocke_value_t* qh_mod = rocke_b_mod(B, row, rocke_b_const_i32(B, NQK)); rocke_value_t* qh = rocke_b_add(B, qh_mul, qh_mod); - rocke_value_t* qh_in = rocke_b_cmp_lt(B, qh, rocke_gfx942_attn2d_num_qh(ctx)); + rocke_value_t* qh_in = rocke_b_cmp_lt(B, qh, ctx->num_qh_p); rocke_value_t* sink_h = rocke_b_global_load(B, ctx->sinks, qh, ctx->dtype, 2); rocke_value_t* sink_f = rocke_b_fmul(B, rocke_b_cast_to_f32(B, sink_h), rcp_ln2); m_inits[r] = rocke_b_select(B, qh_in, sink_f, neg_inf); diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/gfx942/attention_tiled_2d_public_entry_glue.cpp b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/gfx942/attention_tiled_2d_public_entry_glue.cpp index ed757fc12c8e..9cffd3422cb9 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/gfx942/attention_tiled_2d_public_entry_glue.cpp +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/gfx942/attention_tiled_2d_public_entry_glue.cpp @@ -523,6 +523,9 @@ static rocke_status_t nqh = s->num_query_heads; nkv = s->num_kv_heads; + /* The name carries the GQA ratio; reject a spec it cannot be formed for. */ + if(nkv <= 0) + return ROCKE_ERR_VALUE; snprintf(d_buf, sizeof(d_buf), "d%d", s->head_size); snprintf(b_buf, sizeof(b_buf), "b%d", s->block_size); diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/gfx942/attention_tiled_2d_q_load_gather.cpp b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/gfx942/attention_tiled_2d_q_load_gather.cpp index 6895f477491a..e1cc34c0d767 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/gfx942/attention_tiled_2d_q_load_gather.cpp +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/gfx942/attention_tiled_2d_q_load_gather.cpp @@ -84,12 +84,10 @@ void rocke_gfx942_attn2d_emit_q_load(rocke_gfx942_attn2d_build_ctx_t* ctx) * recompute here for the staged-store fallback path. */ z8 = rocke_b_zero_vec(b, dtype, 8); - /* q_desc = TensorDescriptor.naive("Q", lengths=[1<<30, NUM_QH, HD], - * coord_names=("token","head","dim")). Stored on ctx for the - * downstream Q gather (32x32 direct-global) + the epilogue store. */ - if(ctx->runtime_shape) + /* q_desc over (token, head, dim), stored on ctx for the downstream Q gather + * (32x32 direct-global) + the epilogue store. The head count is a kernarg, + * so the token stride is too. */ { - /* The head count is a kernarg, so the token stride is too. */ static const char* const coord_names[3] = {"token", "head", "dim"}; rocke_value_t* strides[3] = {rocke_b_mul(b, ctx->num_qh_p, rocke_b_const_i32(b, ctx->HD)), NULL, NULL}; @@ -98,15 +96,6 @@ void rocke_gfx942_attn2d_emit_q_load(rocke_gfx942_attn2d_build_ctx_t* ctx) b, "Q", coord_names, 3, strides, static_strides) ->base; } - else - { - int lengths[3]; - static const char* const coord_names[3] = {"token", "head", "dim"}; - lengths[0] = 1 << 30; - lengths[1] = ctx->NUM_QH; - lengths[2] = ctx->HD; - ctx->q_desc = rocke_tensor_descriptor_naive(b, "Q", lengths, 3, NULL, coord_names, 3); - } if(ctx->Q_DIRECT_GLOBAL) { @@ -163,7 +152,7 @@ void rocke_gfx942_attn2d_emit_q_load(rocke_gfx942_attn2d_build_ctx_t* ctx) * guard is always true — omitting it avoids a redundant cmp+and on the hot path. */ { rocke_value_t* lt_pos = rocke_b_cmp_lt(b, q_pos_t, ctx->cur_batch_q_len); - rocke_value_t* lt_qh = rocke_b_cmp_lt(b, qh_t, rocke_gfx942_attn2d_num_qh(ctx)); + rocke_value_t* lt_qh = rocke_b_cmp_lt(b, qh_t, ctx->num_qh_p); rocke_value_t* and_pos_qh = rocke_b_land(b, lt_pos, lt_qh); if(ctx->VALID_ROWS < ctx->BLOCK_M) { @@ -437,7 +426,7 @@ void rocke_gfx942_attn2d_emit_q_gather(rocke_gfx942_attn2d_build_ctx_t* ctx) /* q32_mask = (q32_poscur_batch_q_len), - rocke_b_cmp_lt(b, q32_h, rocke_gfx942_attn2d_num_qh(ctx))); + rocke_b_cmp_lt(b, q32_h, ctx->num_qh_p)); q32_pos_safe = rocke_b_select(b, q32_mask, q32_pos, rocke_b_const_i32(b, 0)); q32_h_safe = rocke_b_select(b, q32_mask, q32_h, rocke_b_const_i32(b, 0)); diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/gfx942/attention_tiled_3d_gfx942_attn_tiled_3d_public_glue.cpp b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/gfx942/attention_tiled_3d_gfx942_attn_tiled_3d_public_glue.cpp index 22a1c6de9cb3..dd46f7a54812 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/gfx942/attention_tiled_3d_gfx942_attn_tiled_3d_public_glue.cpp +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/gfx942/attention_tiled_3d_gfx942_attn_tiled_3d_public_glue.cpp @@ -40,6 +40,7 @@ * the 16x16x16 f16+bf16 atoms present), which IS fully reproducible here. The * fallback reason strings mirror the Python text byte-for-byte. */ +#include #include #include #include @@ -242,6 +243,25 @@ const rocke_type_t* return rocke_bf16(); } +int rocke_unified_attention_3d_tiled_spec_binary_search_iters( + const rocke_unified_attention_3d_tiled_spec_t* s) +{ + double v; + int r; + if(s == NULL) + { + return -1; + } + if(s->num_seqs <= 0) + { + return 32; + } + /* max(1, ceil(log2(num_seqs + 1))) */ + v = ceil(log2((double)(s->num_seqs + 1))); + r = (int)v; + return r < 1 ? 1 : r; +} + int rocke_unified_attention_3d_tiled_spec_kernel_name( const rocke_unified_attention_3d_tiled_spec_t* s, char* buf, size_t cap) { @@ -279,6 +299,10 @@ int rocke_unified_attention_3d_tiled_spec_kernel_name( snprintf(d_part, sizeof(d_part), "d%d", s->head_size); snprintf(b_part, sizeof(b_part), "b%d", s->block_size); /* Head counts are kernargs; only the GQA ratio is baked. */ + if(s->num_kv_heads <= 0) + { + return -1; + } snprintf(h_part, sizeof(h_part), "g%d", s->num_query_heads / s->num_kv_heads); snprintf(seg_part, sizeof(seg_part), "seg%d", s->num_segments); diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/gfx950/attention_tiled_2d_ctx_prologue.cpp b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/gfx950/attention_tiled_2d_ctx_prologue.cpp index 25c952141626..6cfc97e5c964 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/gfx950/attention_tiled_2d_ctx_prologue.cpp +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/gfx950/attention_tiled_2d_ctx_prologue.cpp @@ -139,7 +139,6 @@ bool rocke_gfx950_attn2d_build_ctx_init(rocke_gfx950_attn2d_build_ctx_t* ctx, const int BLOCK_Q = rocke_attention_tiled_2d_spec_block_q(spec); const int NQK = rocke_attention_tiled_2d_spec_num_queries_per_kv(spec); const int NUM_KV = spec->num_kv_heads; - const int NUM_QH = spec->num_query_heads; const int SLIDING_WINDOW = spec->sliding_window; ctx->HD = HD; @@ -150,7 +149,6 @@ bool rocke_gfx950_attn2d_build_ctx_init(rocke_gfx950_attn2d_build_ctx_t* ctx, ctx->BLOCK_Q = BLOCK_Q; ctx->NQK = NQK; ctx->NUM_KV = NUM_KV; - ctx->NUM_QH = NUM_QH; ctx->SLIDING_WINDOW = SLIDING_WINDOW; ctx->USE_SOFTCAP = spec->has_softcap; ctx->USE_SINKS = spec->use_sinks; @@ -292,6 +290,9 @@ bool rocke_gfx950_attn2d_build_ctx_init(rocke_gfx950_attn2d_build_ctx_t* ctx, ctx->num_seqs_p = rocke_b_param(b, "num_seqs", rocke_i32(), NULL); ctx->bt_stride_p = rocke_b_param(b, "block_table_stride", rocke_i32(), NULL); ctx->qq_bias_stride0_p = rocke_b_param(b, "qq_bias_stride_0", rocke_i32(), NULL); + ctx->num_qh_p = rocke_b_param(b, "num_query_heads", rocke_i32(), NULL); + ctx->num_kv_p = rocke_b_param(b, "num_kv_heads", rocke_i32(), NULL); + ctx->sliding_window_p = rocke_b_param(b, "sliding_window", rocke_i32(), NULL); } /* ---- grid ids + wave decomposition (lines 891-904) ---- * @@ -315,14 +316,11 @@ bool rocke_gfx950_attn2d_build_ctx_init(rocke_gfx950_attn2d_build_ctx_t* ctx, } /* ---- seq lookup + Q-block geometry (lines 906-925) ---- */ - ctx->seq_idx - = rocke_binary_search_seq_idx(b, - ctx->cu_q, - ctx->q_block_global_idx, - ctx->num_seqs_p, - BLOCK_Q, - rocke_attention_tiled_2d_spec_binary_search_iters(spec), - false); + { + rocke_value_t* trip = rocke_binary_search_seq_idx_trip_count(b, ctx->num_seqs_p); + ctx->seq_idx = rocke_binary_search_seq_idx_dyn( + b, ctx->cu_q, ctx->q_block_global_idx, ctx->num_seqs_p, BLOCK_Q, trip, false); + } ctx->cu_q_start = rocke_b_global_load_i32(b, ctx->cu_q, ctx->seq_idx, 0); ctx->cu_q_stop = rocke_b_global_load_i32( b, ctx->cu_q, rocke_b_add(b, ctx->seq_idx, rocke_b_const_i32(b, 1)), 0); @@ -470,10 +468,9 @@ bool rocke_gfx950_attn2d_build_ctx_init(rocke_gfx950_attn2d_build_ctx_t* ctx, ctx->pv_fp8_scale_v = ctx->FP8_MFMA_PV ? rocke_b_fdiv(b, ctx->v_scale_p, rocke_b_const_f32(b, 240.0)) : NULL; - /* sw_const = const_i32(SLIDING_WINDOW). Created once in the constants block - * and reused by the tile-bound + mask regions (peer phases). */ - ctx->sw_const_v = rocke_b_const_i32(b, SLIDING_WINDOW); - ctx->c0 = ctx->sw_const_v; + /* sw_const = rt["sliding_window"], reused by the tile-bound + mask regions + * (peer phases). */ + ctx->sw_const_v = ctx->sliding_window_p; /* ---- acc geometry the loop-bounds / epilogue phases alias (lines 1380-1387) ---- */ ctx->ACC_N_TILES = USE_MFMA_32X32 ? (HD / 32) : PV_N_TILES; diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/gfx950/attention_tiled_2d_fastkv_regp_helpers.cpp b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/gfx950/attention_tiled_2d_fastkv_regp_helpers.cpp index dabd67225c4a..cda218408421 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/gfx950/attention_tiled_2d_fastkv_regp_helpers.cpp +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/gfx950/attention_tiled_2d_fastkv_regp_helpers.cpp @@ -38,10 +38,11 @@ static rocke_status_t rocke__fastkv_gfx950_base_name(const rocke_attention_tiled char* out, size_t out_cap) { - char t_buf[32], hkv_buf[64], kv_buf[64], sw_buf[32], w_buf[32], mw_buf[32]; + char t_buf[32], hkv_buf[64], fhkv_buf[64], kv_buf[64], sw_buf[32], w_buf[32], mw_buf[32]; char d_buf[16], b_buf[16]; - if(s == NULL || out == NULL) + /* The name carries the GQA ratio; reject a spec it cannot be formed for. */ + if(s == NULL || out == NULL || s->num_kv_heads <= 0) return ROCKE_ERR_VALUE; snprintf(d_buf, sizeof(d_buf), "d%d", s->head_size); @@ -50,13 +51,20 @@ static rocke_status_t rocke__fastkv_gfx950_base_name(const rocke_attention_tiled snprintf(t_buf, sizeof(t_buf), "t%d", rocke_attention_tiled_2d_spec_tile_size_eff(s)); else t_buf[0] = '\0'; - snprintf(hkv_buf, sizeof(hkv_buf), "h%dkv%d", s->num_query_heads, s->num_kv_heads); + /* Head counts are kernargs; only the GQA ratio is baked. */ + snprintf(hkv_buf, sizeof(hkv_buf), "g%d", s->num_query_heads / s->num_kv_heads); + /* Except on the fast paged-KV descriptor path, whose shifts bake them. */ + if(s->use_fast_paged_kv_desc) + snprintf(fhkv_buf, sizeof(fhkv_buf), "h%dkv%d", s->num_query_heads, s->num_kv_heads); + else + fhkv_buf[0] = '\0'; if(s->kv_storage_dtype != NULL) snprintf(kv_buf, sizeof(kv_buf), "kv%s", s->kv_storage_dtype); else kv_buf[0] = '\0'; + /* The window size is a kernarg; only whether the body masks one is baked. */ if(s->sliding_window > 0) - snprintf(sw_buf, sizeof(sw_buf), "sw%d", s->sliding_window); + snprintf(sw_buf, sizeof(sw_buf), "sw"); else sw_buf[0] = '\0'; if(s->num_warps != 1) @@ -74,8 +82,11 @@ static rocke_status_t rocke__fastkv_gfx950_base_name(const rocke_attention_tiled b_buf, t_buf, hkv_buf, + fhkv_buf, s->dtype ? s->dtype : "", kv_buf, + /* i64 K/V addressing changes the loaders, so it is a distinct kernel. */ + s->use_i64_kv_addr ? "i64" : "", s->use_sinks ? "sinks" : "", sw_buf, s->has_softcap ? "softcap" : "", diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/gfx950/attention_tiled_2d_kv_body_pv_epilogue.cpp b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/gfx950/attention_tiled_2d_kv_body_pv_epilogue.cpp index 0a8b47281437..486f323cde00 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/gfx950/attention_tiled_2d_kv_body_pv_epilogue.cpp +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/gfx950/attention_tiled_2d_kv_body_pv_epilogue.cpp @@ -656,6 +656,13 @@ void rocke_gfx950_attn2d_emit_pv_bucket(rocke_gfx950_attn2d_build_ctx_t* ctx, for(atom = 0; atom < ctx->ACC_M_ATOMS; ++atom) yields[yc++] = new_acc[n * ctx->ACC_M_ATOMS + atom]; yields[yc++] = ctx->GROUPED_KV2 ? in->cur_buf : ctx->nxt_buf_v; + if(rocke_gfx950_attn2d_page_prefetch(ctx)) + { + for(i = 0; i < ctx->N_BLOCKS_PER_TILE; ++i) + yields[yc++] = ctx->page_next[i]; + for(i = 0; i < ctx->N_BLOCKS_PER_TILE; ++i) + yields[yc++] = ctx->page_after_next[i]; + } /* Record the rewritten carry for callers that thread phases (the driver's * single-phase path reads this back via the loop results). */ @@ -711,6 +718,15 @@ rocke_for_t rocke_gfx950_attn2d_drive_kv_loop(rocke_gfx950_attn2d_build_ctx_t* c for(i = 0; i < num_accs; ++i) ctx->acc_cur[i] = kvloop.iter_vars[ml + i]; ctx->cur_buf = kvloop.iter_vars[ml + num_accs]; + if(rocke_gfx950_attn2d_page_prefetch(ctx)) + { + for(i = 0; i < ctx->N_BLOCKS_PER_TILE; ++i) + { + ctx->page_cur[i] = kvloop.iter_vars[ml + num_accs + 1 + i]; + ctx->page_next[i] = kvloop.iter_vars[ml + num_accs + 1 + ctx->N_BLOCKS_PER_TILE + i]; + } + ctx->page_cur_tile = kvloop.iv; + } ctx->skip_mask = false; rocke_gfx950_attn2d_emit_kv_body(ctx); @@ -778,8 +794,7 @@ rocke_kernel_def_t* rocke_gfx950_attn2d_emit_epilogue(rocke_gfx950_attn2d_build_ rocke_value_t* op_qh_mod = rocke_b_mod(b, q_row_t, rocke_b_const_i32(b, ctx->NQK)); rocke_value_t* op_qh_t = rocke_b_add(b, op_qh_mul, op_qh_mod); rocke_value_t* op_mask_pos = rocke_b_cmp_lt(b, op_pos_t, ctx->cur_batch_q_len); - rocke_value_t* op_mask_qh - = rocke_b_cmp_lt(b, op_qh_t, rocke_b_const_i32(b, ctx->NUM_QH)); + rocke_value_t* op_mask_qh = rocke_b_cmp_lt(b, op_qh_t, ctx->num_qh_p); rocke_value_t* op_mask_inner = rocke_b_land(b, op_mask_pos, op_mask_qh); rocke_value_t* op_mask_t; if(ctx->VALID_ROWS < ctx->BLOCK_M) @@ -875,8 +890,7 @@ rocke_kernel_def_t* rocke_gfx950_attn2d_emit_epilogue(rocke_gfx950_attn2d_build_ } { rocke_value_t* mask_pos = rocke_b_cmp_lt(b, op_pos32_base, ctx->cur_batch_q_len); - rocke_value_t* mask_qh - = rocke_b_cmp_lt(b, op_qh32_base, rocke_b_const_i32(b, ctx->NUM_QH)); + rocke_value_t* mask_qh = rocke_b_cmp_lt(b, op_qh32_base, ctx->num_qh_p); rocke_value_t* mask_inner = rocke_b_land(b, mask_pos, mask_qh); if(ctx->VALID_ROWS < ctx->BLOCK_M) { @@ -994,7 +1008,7 @@ rocke_kernel_def_t* rocke_gfx950_attn2d_emit_epilogue(rocke_gfx950_attn2d_build_ } { rocke_value_t* mask_pos = rocke_b_cmp_lt(b, op_pos, ctx->cur_batch_q_len); - rocke_value_t* mask_qh = rocke_b_cmp_lt(b, op_qh, rocke_b_const_i32(b, ctx->NUM_QH)); + rocke_value_t* mask_qh = rocke_b_cmp_lt(b, op_qh, ctx->num_qh_p); rocke_value_t* mask_inner = rocke_b_land(b, mask_pos, mask_qh); if(ctx->VALID_ROWS < ctx->BLOCK_M) { diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/gfx950/attention_tiled_2d_kv_body_qk_softmax.cpp b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/gfx950/attention_tiled_2d_kv_body_qk_softmax.cpp index e4ad658b356d..d44eac2d1f2a 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/gfx950/attention_tiled_2d_kv_body_qk_softmax.cpp +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/gfx950/attention_tiled_2d_kv_body_qk_softmax.cpp @@ -243,7 +243,7 @@ void rocke_gfx950_attn2d_emit_licm_hoist(rocke_gfx950_attn2d_build_ctx_t* ctx) rocke_value_t* qh_mod = rocke_b_mod(b, row, rocke_b_const_i32(b, ctx->NQK)); rocke_value_t* qh_r = rocke_b_add(b, qh_mul, qh_mod); rocke_value_t* row_ok_pos = rocke_b_cmp_lt(b, qp_r, ctx->cur_batch_q_len); - rocke_value_t* row_ok_qh = rocke_b_cmp_lt(b, qh_r, rocke_b_const_i32(b, ctx->NUM_QH)); + rocke_value_t* row_ok_qh = rocke_b_cmp_lt(b, qh_r, ctx->num_qh_p); rocke_value_t* row_ok_inner = rocke_b_land(b, row_ok_pos, row_ok_qh); rocke_value_t* row_ok; if(ctx->VALID_ROWS < ctx->BLOCK_M) @@ -271,7 +271,7 @@ void rocke_gfx950_attn2d_emit_licm_hoist(rocke_gfx950_attn2d_build_ctx_t* ctx) for(reg = 0; reg < ctx->REGS_PER_LANE; ++reg) { rocke_value_t* qh_r = ctx->hoist_q_head[reg]; - rocke_value_t* qh_ok = rocke_b_cmp_lt(b, qh_r, rocke_b_const_i32(b, ctx->NUM_QH)); + rocke_value_t* qh_ok = rocke_b_cmp_lt(b, qh_r, ctx->num_qh_p); rocke_value_t* slope = rocke_b_masked_global_load( b, ctx->alibi_slopes_ptr, qh_r, qh_ok, rocke_b_const_f32(b, 0.0), f32, 4); ctx->hoist_q_head[reg] = slope; @@ -294,7 +294,7 @@ void rocke_gfx950_attn2d_emit_licm_hoist(rocke_gfx950_attn2d_build_ctx_t* ctx) rocke_value_t* st_qh_mod = rocke_b_mod(b, st_q_row, rocke_b_const_i32(b, ctx->NQK)); rocke_value_t* st_qh = rocke_b_add(b, st_qh_mul, st_qh_mod); rocke_value_t* st_ok_qp = rocke_b_cmp_lt(b, st_qp, ctx->cur_batch_q_len); - rocke_value_t* st_ok_qh = rocke_b_cmp_lt(b, st_qh, rocke_b_const_i32(b, ctx->NUM_QH)); + rocke_value_t* st_ok_qh = rocke_b_cmp_lt(b, st_qh, ctx->num_qh_p); rocke_value_t* st_ok_inner = rocke_b_land(b, st_ok_qp, st_ok_qh); ctx->st_qp_hoist = st_qp; ctx->st_qh_hoist = st_qh; @@ -311,7 +311,7 @@ void rocke_gfx950_attn2d_emit_licm_hoist(rocke_gfx950_attn2d_build_ctx_t* ctx) ctx->st_causal_lim_hoist = rocke_b_add(b, ctx->context_len, st_qp); if(ctx->USE_ALIBI) { - rocke_value_t* st_qh_ok = rocke_b_cmp_lt(b, st_qh, rocke_b_const_i32(b, ctx->NUM_QH)); + rocke_value_t* st_qh_ok = rocke_b_cmp_lt(b, st_qh, ctx->num_qh_p); ctx->st_alibi_slope_hoist = rocke_b_masked_global_load(b, ctx->alibi_slopes_ptr, st_qh, @@ -528,7 +528,7 @@ void rocke_gfx950_attn2d_emit_kv_body(rocke_gfx950_attn2d_build_ctx_t* ctx) rocke_value_t* st_qh_mod = rocke_b_mod(b, st_q_row_iter, rocke_b_const_i32(b, ctx->NQK)); rocke_value_t* st_qh_iter = rocke_b_add(b, st_qh_mul, st_qh_mod); rocke_value_t* st_ok_qp = rocke_b_cmp_lt(b, st_qp_iter, ctx->cur_batch_q_len); - rocke_value_t* st_ok_qh = rocke_b_cmp_lt(b, st_qh_iter, rocke_b_const_i32(b, ctx->NUM_QH)); + rocke_value_t* st_ok_qh = rocke_b_cmp_lt(b, st_qh_iter, ctx->num_qh_p); rocke_value_t* st_ok_inner = rocke_b_land(b, st_ok_qp, st_ok_qh); if(ctx->VALID_ROWS < ctx->BLOCK_M) { @@ -543,8 +543,7 @@ void rocke_gfx950_attn2d_emit_kv_body(rocke_gfx950_attn2d_build_ctx_t* ctx) st_causal_lim_iter = rocke_b_add(b, ctx->context_len, st_qp_iter); if(ctx->USE_ALIBI) { - rocke_value_t* st_qh_iter_ok - = rocke_b_cmp_lt(b, st_qh_iter, rocke_b_const_i32(b, ctx->NUM_QH)); + rocke_value_t* st_qh_iter_ok = rocke_b_cmp_lt(b, st_qh_iter, ctx->num_qh_p); st_alibi_slope_iter = rocke_b_masked_global_load(b, ctx->alibi_slopes_ptr, st_qh_iter, @@ -561,6 +560,8 @@ void rocke_gfx950_attn2d_emit_kv_body(rocke_gfx950_attn2d_build_ctx_t* ctx) rocke_value_t* in_range_next = rocke_b_cmp_lt(b, next_tile_iv_raw, ctx->tile_end); rocke_value_t* safe_next_tile = rocke_b_select(b, in_range_next, next_tile_iv_raw, ctx->kv_tile_iv); + if(rocke_gfx950_attn2d_page_prefetch(ctx)) + ctx->page_next_tile = safe_next_tile; /* softmax-derived state forwarded to the PV bucket. */ rocke_value_t* alpha_regs[ROCKE_GFX950_ATTN2D_MAX_REGS_PER_LANE]; @@ -595,6 +596,19 @@ void rocke_gfx950_attn2d_emit_kv_body(rocke_gfx950_attn2d_build_ctx_t* ctx) && (ctx->SOFTMAX_INTERLEAVE_MODE == 0 || ctx->SOFTMAX_INTERLEAVE_MODE == 1)) rocke_b_iglp_opt(b, ctx->SOFTMAX_INTERLEAVE_MODE); + /* Read the page ids of the tile after next now, right after the iter-start + * drain, so they arrive long before next iteration's copies need them. That + * tile is next iteration's safe_next_tile. */ + if(rocke_gfx950_attn2d_page_prefetch(ctx)) + { + rocke_value_t* c_two = rocke_b_const_i32(b, 2); + rocke_value_t* after_next_raw = rocke_b_add(b, ctx->kv_tile_iv, c_two); + rocke_value_t* in_range_after = rocke_b_cmp_lt(b, after_next_raw, ctx->tile_end); + rocke_value_t* after_next + = rocke_b_select(b, in_range_after, after_next_raw, next_tile_iv_raw); + rocke_gfx950_attn2d_tile_page_ids(ctx, after_next, ctx->page_after_next); + } + /* Early-V schedule: V_lds is single-buffered and the iter-start full drain * guarantees the previous PV's V reads retired, so the current V copy can be * issued before QK to overlap with QK + softmax (Python 2566-2570). */ @@ -760,7 +774,7 @@ void rocke_gfx950_attn2d_emit_kv_body(rocke_gfx950_attn2d_build_ctx_t* ctx) qh_mod = rocke_b_mod(b, q_row_t, rocke_b_const_i32(b, ctx->NQK)); qh_r = rocke_b_add(b, qh_mul, qh_mod); row_ok_pos = rocke_b_cmp_lt(b, qp_r, ctx->cur_batch_q_len); - row_ok_qh = rocke_b_cmp_lt(b, qh_r, rocke_b_const_i32(b, ctx->NUM_QH)); + row_ok_qh = rocke_b_cmp_lt(b, qh_r, ctx->num_qh_p); rocke_value_t* row_ok_inner = rocke_b_land(b, row_ok_pos, row_ok_qh); if(ctx->VALID_ROWS < ctx->BLOCK_M) { diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/gfx950/attention_tiled_2d_kv_descriptors_loads.cpp b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/gfx950/attention_tiled_2d_kv_descriptors_loads.cpp index 82129cb1b3bd..4854c9b805b7 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/gfx950/attention_tiled_2d_kv_descriptors_loads.cpp +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/gfx950/attention_tiled_2d_kv_descriptors_loads.cpp @@ -47,6 +47,9 @@ * Lifetime: every emitted node is arena-owned (rocke_ir_builder_t.arena). Nothing * is freed individually; the arena bulk-frees the whole graph. */ +#include /* snprintf */ +#include /* memset */ + #include "rocke/instance_gfx950_attention_tiled_2d_internal.h" /* ========================================================================== @@ -88,16 +91,20 @@ static int rocke__wave_bytes(const rocke_gfx950_attn2d_build_ctx_t* ctx) return ctx->WAVE * 16; } -/* kv_block_bytes_c value = BS * NUM_KV * HD * KV_BYTES (one-block buffer bound, - * Python 1458). Python creates this const ONCE and reuses it in every loader; - * cache it on ctx so we do not allocate a duplicate const per loader (which would - * shift downstream %values). */ +/* kv_block_bytes_c = kv_stride_blk_b (one-block buffer bound). Python derives + * the token and block strides ONCE and reuses them in every loader; cache them + * on ctx so we do not emit a duplicate multiply per loader (which would shift + * downstream %values). */ static rocke_value_t* rocke__kv_block_bytes_c(rocke_gfx950_attn2d_build_ctx_t* ctx) { if(ctx->kv_block_bytes_c_v == NULL) { - int kv_stride_blk_b = ctx->BS * ctx->NUM_KV * ctx->HD * ctx->KV_BYTES; - ctx->kv_block_bytes_c_v = rocke_b_const_i32(ctx->b, kv_stride_blk_b); + /* The token and block strides scale with the runtime kv head count. */ + ctx->kv_stride_tok_b_v = rocke_b_mul( + ctx->b, ctx->num_kv_p, rocke_b_const_i32(ctx->b, ctx->HD * ctx->KV_BYTES)); + ctx->kv_stride_blk_b_v + = rocke_b_mul(ctx->b, ctx->kv_stride_tok_b_v, rocke_b_const_i32(ctx->b, ctx->BS)); + ctx->kv_block_bytes_c_v = ctx->kv_stride_blk_b_v; } return ctx->kv_block_bytes_c_v; } @@ -175,6 +182,155 @@ static rocke_value_t* rocke__wave_lds_offset_i64(rocke_gfx950_attn2d_build_ctx_t return ctx->wave_lds_off_i64_v; } +/* A copy covers copy_tokens whole tokens (KV_HALVES_PER_CALL is a multiple of + * HD). When they fit in one page the bf16/fp16 loaders address a copy as its + * page's base, its first token's offset in the page and the lane's own part; + * otherwise each wave's share still sits in one page and the loaders keep the + * paged-KV descriptor. The token stride is a kernarg, so the lane part and the + * per-copy token offsets are built once, after the paged-KV descriptor, rather + * than multiplying the token index in every copy. */ +static int rocke__copy_tokens(const rocke_gfx950_attn2d_build_ctx_t* ctx) +{ + return rocke__kv_halves_per_call(ctx) / ctx->HD; +} + +static bool rocke__page_copies(const rocke_gfx950_attn2d_build_ctx_t* ctx) +{ + return !ctx->FAST_PAGED_KV_DESC && !ctx->KV_FP8 && ctx->BS % rocke__copy_tokens(ctx) == 0; +} + +static void rocke__emit_paged_copy_offsets(rocke_gfx950_attn2d_build_ctx_t* ctx) +{ + rocke_ir_builder_t* b = ctx->b; + const int copy_tokens = rocke__copy_tokens(ctx); + const int kv_calls_per_tile = rocke__kv_calls_per_tile(ctx); + rocke_value_t* lane_half_base = rocke__lane_half_base(ctx); + + memset(ctx->copy_tok_b, 0, sizeof(ctx->copy_tok_b)); + ctx->kv_stride_blk_i64_v = NULL; + if(ctx->BS / copy_tokens > 64) + { + b->status = ROCKE_ERR_VALUE; + snprintf(b->err, + ROCKE_ERR_MSG_CAP, + "%d copies per page exceed the 64-entry offset table", + ctx->BS / copy_tokens); + return; + } + rocke_value_t* c_hd0 = rocke_b_const_i32(b, ctx->HD); + rocke_value_t* lane_tok = rocke_b_div(b, lane_half_base, c_hd0); + rocke_value_t* c_hd1 = rocke_b_const_i32(b, ctx->HD); + rocke_value_t* lane_dim = rocke_b_mod(b, lane_half_base, c_hd1); + rocke_value_t* tok_b = rocke_b_mul(b, lane_tok, ctx->kv_stride_tok_b_v); + rocke_value_t* c_h = rocke_b_const_i32(b, ctx->HD * ctx->KV_BYTES); + rocke_value_t* head_b = rocke_b_mul(b, ctx->kv_head_idx, c_h); + rocke_value_t* tok_head = rocke_b_add(b, tok_b, head_b); + rocke_value_t* c_bytes = rocke_b_const_i32(b, ctx->KV_BYTES); + rocke_value_t* dim_b = rocke_b_mul(b, lane_dim, c_bytes); + ctx->lane_within_b_v = rocke_b_add(b, tok_head, dim_b); + for(int call = 0; call < kv_calls_per_tile; ++call) + { + const int tok_i = call % (ctx->BS / copy_tokens); + if(tok_i == 0 || ctx->copy_tok_b[tok_i] != NULL) + continue; + rocke_value_t* c_tok = rocke_b_const_i32(b, tok_i * copy_tokens); + ctx->copy_tok_b[tok_i] = rocke_b_mul(b, ctx->kv_stride_tok_b_v, c_tok); + } + if(ctx->I64_KV_ADDR) + ctx->kv_stride_blk_i64_v = rocke_b_zext(b, ctx->kv_stride_blk_b_v, rocke_i64()); +} + +/* The paged-KV descriptor's indirect lookup for page `blk` of a tile: one + * guarded block-table read, wave-uniform. */ +static rocke_value_t* + rocke__tile_page_id(rocke_gfx950_attn2d_build_ctx_t* ctx, rocke_value_t* kv_tile_idx, int blk) +{ + rocke_ir_builder_t* b = ctx->b; + rocke_value_t* tile_blk = kv_tile_idx; + if(ctx->N_BLOCKS_PER_TILE != 1) + { + rocke_value_t* c_nb = rocke_b_const_i32(b, ctx->N_BLOCKS_PER_TILE); + rocke_value_t* first = rocke_b_mul(b, kv_tile_idx, c_nb); + rocke_value_t* c_blk = rocke_b_const_i32(b, blk); + tile_blk = rocke_b_add(b, first, c_blk); + } + rocke_value_t* idx = rocke_b_add(b, rocke__seq_base(ctx), tile_blk); + rocke_value_t* mask = rocke_b_cmp_lt(b, idx, rocke__block_table_max_idx(ctx)); + rocke_value_t* zero = rocke_b_const_i32(b, 0); + rocke_value_t* page + = rocke_b_masked_global_load(b, ctx->block_tables, idx, mask, zero, rocke_i32(), 4); + return rocke_b_to_sgpr_u32(b, page); +} + +bool rocke_gfx950_attn2d_page_prefetch(const rocke_gfx950_attn2d_build_ctx_t* ctx) +{ + /* The C++ port has no deep K ring (kv_ring_depth != 2 is refused). */ + return rocke__page_copies(ctx) && !ctx->GROUPED_KV2 + && ctx->N_BLOCKS_PER_TILE <= ROCKE_GFX950_ATTN2D_MAX_PREFETCH_PAGES; +} + +void rocke_gfx950_attn2d_tile_page_ids(rocke_gfx950_attn2d_build_ctx_t* ctx, + rocke_value_t* kv_tile_idx, + rocke_value_t** out) +{ + for(int j = 0; j < ctx->N_BLOCKS_PER_TILE; ++j) + out[j] = rocke__tile_page_id(ctx, kv_tile_idx, j); +} + +/* The page ids the loop carries for `kv_tile_idx`, or NULL when it reads them + * in place. */ +static rocke_value_t* const* rocke__carried_pages(const rocke_gfx950_attn2d_build_ctx_t* ctx, + const rocke_value_t* kv_tile_idx) +{ + if(!rocke_gfx950_attn2d_page_prefetch(ctx)) + return NULL; + if(kv_tile_idx == ctx->page_cur_tile) + return ctx->page_cur; + if(kv_tile_idx == ctx->page_next_tile) + return ctx->page_next; + if(kv_tile_idx == ctx->page_start_tile) + return ctx->page_start; + return NULL; +} + +/* (rsrc, voffset, soffset) of one bf16/fp16 K/V copy. `pages` caches this + * tile's wave-uniform page ids by block within the tile. */ +static void rocke__paged_copy(rocke_gfx950_attn2d_build_ctx_t* ctx, + rocke_value_t* src, + rocke_value_t* src_rsrc, + rocke_value_t* kv_tile_idx, + int call, + rocke_value_t** pages, + rocke_value_t** out_rsrc, + rocke_value_t** out_voff, + rocke_value_t** out_soff) +{ + rocke_ir_builder_t* b = ctx->b; + const int copy_tokens = rocke__copy_tokens(ctx); + const int blk = (call * copy_tokens) / ctx->BS; + const int tok_i = call % (ctx->BS / copy_tokens); + rocke_value_t* const* carried = rocke__carried_pages(ctx, kv_tile_idx); + if(carried != NULL) + pages = (rocke_value_t**)carried; + else if(pages[blk] == NULL) + pages[blk] = rocke__tile_page_id(ctx, kv_tile_idx, blk); + rocke_value_t* voff = ctx->lane_within_b_v; + if(tok_i != 0) + voff = rocke_b_add(b, voff, ctx->copy_tok_b[tok_i]); + *out_voff = voff; + if(ctx->I64_KV_ADDR) + { + rocke_value_t* page_i64 = rocke_b_zext(b, pages[blk], rocke_i64()); + rocke_value_t* base_i64 = rocke_b_mul(b, page_i64, ctx->kv_stride_blk_i64_v); + rocke_value_t* ptr = rocke_b_global_ptr_add(b, src, base_i64); + *out_rsrc = rocke_b_buffer_rsrc(b, ptr, rocke__kv_block_bytes_c(ctx)); + *out_soff = rocke__zero_soff(ctx); + return; + } + *out_rsrc = src_rsrc; + *out_soff = rocke_b_mul(b, pages[blk], ctx->kv_stride_blk_b_v); +} + /* ========================================================================== * Pre-loop: K/V buffer resources, byte derivation, paged-KV descriptor * (Python lines 1163-1630, the K/V-descriptor parts). @@ -198,9 +354,9 @@ void rocke_gfx950_attn2d_emit_preloop(rocke_gfx950_attn2d_build_ctx_t* ctx) (void)((ctx->T * ctx->HD) % kv_halves_per_call); } - /* kv_block_bytes_c (1458): Python creates this one-block-bound const here, - * after the byte-stride derivation and BEFORE lane_half_base, and reuses it - * in every loader. Emit it now so its SSA value lands in source order. */ + /* kv_block_bytes_c (1458): Python derives the byte strides and this + * one-block bound here, BEFORE lane_half_base, and reuses them in every + * loader. Emit them now so their SSA values land in source order. */ (void)rocke__kv_block_bytes_c(ctx); /* Per-lane starting half index (1460), then K/V LDS base ptrtoints (1462- @@ -243,18 +399,16 @@ void rocke_gfx950_attn2d_emit_preloop(rocke_gfx950_attn2d_build_ctx_t* ctx) } /* ---- Paged KV byte descriptor (full transform DAG, 1591-1629) ---- - * _kv_base = naive("paged_kv_bytes", - * lengths=[1<<24, BS, NUM_KV, HD], - * strides=[kv_stride_blk_b, kv_stride_tok_b, kv_stride_h_b, KV_BYTES], - * coord_names=("physical_block","token","kv_head","dim")). */ - int blk_b = ctx->BS * ctx->NUM_KV * ctx->HD * ctx->KV_BYTES; - int tok_b = ctx->NUM_KV * ctx->HD * ctx->KV_BYTES; - int h_b = ctx->HD * ctx->KV_BYTES; - int kv_lengths[4] = {1 << 24, ctx->BS, ctx->NUM_KV, ctx->HD}; - int kv_strides[4] = {blk_b, tok_b, h_b, ctx->KV_BYTES}; + * _kv_base = DynamicTensorDescriptor.create("paged_kv_bytes", + * coord_names=("physical_block","token","kv_head","dim"), + * strides=(kv_stride_blk_b, kv_stride_tok_b, kv_stride_h_b, KV_BYTES)), + * the first two the runtime strides derived above. */ + int kv_strides[4] = {1, 1, ctx->HD * ctx->KV_BYTES, ctx->KV_BYTES}; + rocke_value_t* kv_dyn[4] = {ctx->kv_stride_blk_b_v, ctx->kv_stride_tok_b_v, NULL, NULL}; static const char* const kv_coords[4] = {"physical_block", "token", "kv_head", "dim"}; - rocke_tensor_descriptor_t* kv_base = rocke_tensor_descriptor_naive( - b, "paged_kv_bytes", kv_lengths, 4, kv_strides, kv_coords, 4); + rocke_tensor_descriptor_t* kv_base = &rocke_tensor_descriptor_naive_dynamic_mixed( + b, "paged_kv_bytes", kv_coords, 4, kv_dyn, kv_strides) + ->base; /* Single-block tile (N_BLOCKS_PER_TILE == 1, T == BS): * indirect(tile_idx -> physical_block) ; unmerge(linear_half -> token,dim). @@ -285,6 +439,8 @@ void rocke_gfx950_attn2d_emit_preloop(rocke_gfx950_attn2d_build_ctx_t* ctx) b, "linear_block_idx", "physical_block", ctx->block_tables, seq_base, max_idx, 0); ctx->kv_desc = rocke_tensor_descriptor_transform(b, kv_base, chain, 3); } + if(rocke__page_copies(ctx)) + rocke__emit_paged_copy_offsets(ctx); /* fp8 async-loader setup (Python 1785-1796), after the paged-KV * descriptor DAG. */ @@ -350,8 +506,9 @@ static void rocke__fast_paged_kv_voff_split(rocke_gfx950_attn2d_build_ctx_t* ctx if(ctx->I64_KV_ADDR) { /* base_i64 = zext(physical, i64) << 15 ; return (base_i64, within). */ - rocke_value_t* base_i64 - = rocke_b_shl(b, rocke_b_zext(b, physical, rocke_i64()), rocke_b_const_i64(b, 15)); + /* zext before the shift constant, as Python evaluates them. */ + rocke_value_t* physical_i64 = rocke_b_zext(b, physical, rocke_i64()); + rocke_value_t* base_i64 = rocke_b_shl(b, physical_i64, rocke_b_const_i64(b, 15)); if(out_base_i64) *out_base_i64 = base_i64; if(out_voff) @@ -410,10 +567,12 @@ void rocke_gfx950_attn2d_issue_k_load_runtime(rocke_gfx950_attn2d_build_ctx_t* c if(ctx->FAST_PAGED_KV_DESC) rocke_gfx950_attn2d_fast_paged_kv_blocks(ctx, kv_tile_idx, &fast_block0, &fast_block1); + rocke_value_t* pages[64] = {NULL}; /* this tile's page ids by block */ for(int call = 0; call < kv_calls_per_tile; ++call) { rocke_value_t* k_rsrc = ctx->k_rsrc; rocke_value_t* voff = NULL; + rocke_value_t* soff = zero_soff; if(ctx->FAST_PAGED_KV_DESC) { rocke_value_t* base_i64 = NULL; @@ -424,6 +583,11 @@ void rocke_gfx950_attn2d_issue_k_load_runtime(rocke_gfx950_attn2d_build_ctx_t* c b, rocke_b_global_ptr_add(b, ctx->key, base_i64), kv_block_bytes_c); } } + else if(rocke__page_copies(ctx)) + { + rocke__paged_copy( + ctx, ctx->key, ctx->k_rsrc, kv_tile_idx, call, pages, &k_rsrc, &voff, &soff); + } else { /* Non-fast paged-KV (N_BLOCKS_PER_TILE==1 / multi-block): @@ -460,8 +624,7 @@ void rocke_gfx950_attn2d_issue_k_load_runtime(rocke_gfx950_attn2d_build_ctx_t* c rocke_value_t* k_dst = rocke_b_smem_ptr_add( b, K_wave_base, rocke_b_const_i64(b, (int64_t)call * bytes_per_call)); /* CACHE_STREAM (SLC): one-shot streaming load; dwords=4 (1687-1689). */ - rocke_b_async_buffer_load_lds_addr( - b, k_rsrc, k_dst, voff, zero_soff, 4, ROCKE_CACHE_STREAM); + rocke_b_async_buffer_load_lds_addr(b, k_rsrc, k_dst, voff, soff, 4, ROCKE_CACHE_STREAM); } } @@ -489,10 +652,12 @@ void rocke_gfx950_attn2d_issue_v_load_runtime(rocke_gfx950_attn2d_build_ctx_t* c if(ctx->FAST_PAGED_KV_DESC) rocke_gfx950_attn2d_fast_paged_kv_blocks(ctx, kv_tile_idx, &fast_block0, &fast_block1); + rocke_value_t* pages[64] = {NULL}; /* this tile's page ids by block */ for(int call = 0; call < kv_calls_per_tile; ++call) { rocke_value_t* v_rsrc = ctx->v_rsrc; rocke_value_t* voff = NULL; + rocke_value_t* soff = zero_soff; if(ctx->FAST_PAGED_KV_DESC) { rocke_value_t* base_i64 = NULL; @@ -503,6 +668,11 @@ void rocke_gfx950_attn2d_issue_v_load_runtime(rocke_gfx950_attn2d_build_ctx_t* c b, rocke_b_global_ptr_add(b, ctx->value, base_i64), kv_block_bytes_c); } } + else if(rocke__page_copies(ctx)) + { + rocke__paged_copy( + ctx, ctx->value, ctx->v_rsrc, kv_tile_idx, call, pages, &v_rsrc, &voff, &soff); + } else { /* Non-fast paged-KV: voff = paged_kv_desc.offset(tile_idx=, @@ -538,8 +708,7 @@ void rocke_gfx950_attn2d_issue_v_load_runtime(rocke_gfx950_attn2d_build_ctx_t* c rocke_value_t* v_dst = rocke_b_smem_ptr_add( b, V_wave_base, rocke_b_const_i64(b, (int64_t)call * bytes_per_call)); /* CACHE_STREAM (SLC): V consumed once per iter; dwords=4 (1737-1739). */ - rocke_b_async_buffer_load_lds_addr( - b, v_rsrc, v_dst, voff, zero_soff, 4, ROCKE_CACHE_STREAM); + rocke_b_async_buffer_load_lds_addr(b, v_rsrc, v_dst, voff, soff, 4, ROCKE_CACHE_STREAM); } } diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/gfx950/attention_tiled_2d_public_entry_glue.cpp b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/gfx950/attention_tiled_2d_public_entry_glue.cpp index 859ece99e9d4..4336a1774cb4 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/gfx950/attention_tiled_2d_public_entry_glue.cpp +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/gfx950/attention_tiled_2d_public_entry_glue.cpp @@ -566,7 +566,6 @@ static bool rocke_g950_build_ctx_init_local(rocke_gfx950_attn2d_build_ctx_t* ctx const int BLOCK_Q = rocke_attention_tiled_2d_spec_block_q(spec); const int NQK = rocke_attention_tiled_2d_spec_num_queries_per_kv(spec); const int NUM_KV = spec->num_kv_heads; - const int NUM_QH = spec->num_query_heads; const int SLIDING_WINDOW = spec->sliding_window; ctx->HD = HD; @@ -578,7 +577,6 @@ static bool rocke_g950_build_ctx_init_local(rocke_gfx950_attn2d_build_ctx_t* ctx ctx->VALID_ROWS = BLOCK_Q * NQK; ctx->NQK = NQK; ctx->NUM_KV = NUM_KV; - ctx->NUM_QH = NUM_QH; ctx->SLIDING_WINDOW = SLIDING_WINDOW; ctx->USE_SOFTCAP = spec->has_softcap; ctx->USE_SINKS = spec->use_sinks; @@ -742,6 +740,9 @@ static bool rocke_g950_build_ctx_init_local(rocke_gfx950_attn2d_build_ctx_t* ctx ctx->num_seqs_p = rocke_b_param(b, "num_seqs", rocke_i32(), NULL); ctx->bt_stride_p = rocke_b_param(b, "block_table_stride", rocke_i32(), NULL); ctx->qq_bias_stride0_p = rocke_b_param(b, "qq_bias_stride_0", rocke_i32(), NULL); + ctx->num_qh_p = rocke_b_param(b, "num_query_heads", rocke_i32(), NULL); + ctx->num_kv_p = rocke_b_param(b, "num_kv_heads", rocke_i32(), NULL); + ctx->sliding_window_p = rocke_b_param(b, "sliding_window", rocke_i32(), NULL); } /* ---- grid ids + wave decomposition (Py891-904) ---- */ @@ -763,14 +764,11 @@ static bool rocke_g950_build_ctx_init_local(rocke_gfx950_attn2d_build_ctx_t* ctx } /* ---- seq lookup + Q-block geometry (Py906-925) ---- */ - ctx->seq_idx - = rocke_binary_search_seq_idx(b, - ctx->cu_q, - ctx->q_block_global_idx, - ctx->num_seqs_p, - BLOCK_Q, - rocke_attention_tiled_2d_spec_binary_search_iters(spec), - false); + { + rocke_value_t* trip = rocke_binary_search_seq_idx_trip_count(b, ctx->num_seqs_p); + ctx->seq_idx = rocke_binary_search_seq_idx_dyn( + b, ctx->cu_q, ctx->q_block_global_idx, ctx->num_seqs_p, BLOCK_Q, trip, false); + } ctx->cu_q_start = rocke_b_global_load_i32(b, ctx->cu_q, ctx->seq_idx, 0); ctx->cu_q_stop = rocke_b_global_load_i32( b, ctx->cu_q, rocke_b_add(b, ctx->seq_idx, rocke_b_const_i32(b, 1)), 0); @@ -914,8 +912,7 @@ static bool rocke_g950_build_ctx_init_local(rocke_gfx950_attn2d_build_ctx_t* ctx * here so the SSA id matches Python; the PV epilogue reuses this value. */ ctx->pv_fp8_scale_v = FP8_MFMA_PV ? rocke_b_fdiv(b, ctx->v_scale_p, rocke_b_const_f32(b, 240.0)) : NULL; - ctx->sw_const_v = rocke_b_const_i32(b, SLIDING_WINDOW); /* sw_const (1150) */ - ctx->c0 = ctx->sw_const_v; + ctx->sw_const_v = ctx->sliding_window_p; /* sw_const = rt["sliding_window"] */ /* ---- acc geometry the epilogue aliases (Py1381-1382) ---- */ ctx->ACC_N_TILES = USE_MFMA_32X32 ? (HD / 32) : PV_N_TILES; @@ -1025,7 +1022,7 @@ static rocke_status_t rocke_g950_attn2d_kernel_name(const rocke_attention_tiled_ char* out, size_t out_cap) { - char d_buf[32], b_buf[32], t_buf[32], hkv_buf[64], kv_buf[64]; + char d_buf[32], b_buf[32], t_buf[32], hkv_buf[64], fhkv_buf[64], kv_buf[64]; char sw_buf[32], w_buf[32], mw_buf[32], schedb_buf[32], smxil_buf[32]; int nqh, nkv; @@ -1034,6 +1031,9 @@ static rocke_status_t rocke_g950_attn2d_kernel_name(const rocke_attention_tiled_ nqh = s->num_query_heads; nkv = s->num_kv_heads; + /* The name carries the GQA ratio; reject a spec it cannot be formed for. */ + if(nkv <= 0) + return ROCKE_ERR_VALUE; snprintf(d_buf, sizeof(d_buf), "d%d", s->head_size); snprintf(b_buf, sizeof(b_buf), "b%d", s->block_size); @@ -1041,13 +1041,20 @@ static rocke_status_t rocke_g950_attn2d_kernel_name(const rocke_attention_tiled_ snprintf(t_buf, sizeof(t_buf), "t%d", rocke_attention_tiled_2d_spec_tile_size_eff(s)); else t_buf[0] = '\0'; - snprintf(hkv_buf, sizeof(hkv_buf), "h%dkv%d", nqh, nkv); + /* Head counts are kernargs; only the GQA ratio is baked. */ + snprintf(hkv_buf, sizeof(hkv_buf), "g%d", nqh / nkv); + /* Except on the fast paged-KV descriptor path, whose shifts bake them. */ + if(s->use_fast_paged_kv_desc) + snprintf(fhkv_buf, sizeof(fhkv_buf), "h%dkv%d", nqh, nkv); + else + fhkv_buf[0] = '\0'; if(s->kv_storage_dtype) snprintf(kv_buf, sizeof(kv_buf), "kv%s", s->kv_storage_dtype); else kv_buf[0] = '\0'; + /* The window size is a kernarg; only whether the body masks one is baked. */ if(s->sliding_window > 0) - snprintf(sw_buf, sizeof(sw_buf), "sw%d", s->sliding_window); + snprintf(sw_buf, sizeof(sw_buf), "sw"); else sw_buf[0] = '\0'; if(s->num_warps != 1) @@ -1083,8 +1090,11 @@ static rocke_status_t rocke_g950_attn2d_kernel_name(const rocke_attention_tiled_ b_buf, t_buf, hkv_buf, + fhkv_buf, s->dtype ? s->dtype : "", kv_buf, + /* i64 K/V addressing changes the loaders, so it is a distinct kernel. */ + s->use_i64_kv_addr ? "i64" : "", s->use_sinks ? "sinks" : "", sw_buf, s->has_softcap ? "softcap" : "", diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/gfx950/attention_tiled_2d_q_load_gather.cpp b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/gfx950/attention_tiled_2d_q_load_gather.cpp index 81b1e09a4da5..b1b880569d9a 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/gfx950/attention_tiled_2d_q_load_gather.cpp +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/gfx950/attention_tiled_2d_q_load_gather.cpp @@ -110,16 +110,17 @@ void rocke_gfx950_attn2d_emit_q_load(rocke_gfx950_attn2d_build_ctx_t* ctx) Q_VECS_PER_ROW = ctx->HD / 8; Q_VECS_PER_THREAD = (ctx->BLOCK_M * Q_VECS_PER_ROW) / ctx->THREADS; - /* q_desc = TensorDescriptor.naive("Q", lengths=[1<<30, NUM_QH, HD], - * coord_names=("token","head","dim")). Stored on ctx for the - * downstream Q gather + the epilogue output store. */ + /* q_desc over (token, head, dim), stored on ctx for the downstream Q gather + * + the epilogue output store. The head count is a kernarg, so the token + * stride is too. */ { - int lengths[3]; static const char* const coord_names[3] = {"token", "head", "dim"}; - lengths[0] = 1 << 30; - lengths[1] = ctx->NUM_QH; - lengths[2] = ctx->HD; - ctx->q_desc = rocke_tensor_descriptor_naive(b, "Q", lengths, 3, NULL, coord_names, 3); + rocke_value_t* strides[3] + = {rocke_b_mul(b, ctx->num_qh_p, rocke_b_const_i32(b, ctx->HD)), NULL, NULL}; + const int static_strides[3] = {1, ctx->HD, 1}; + ctx->q_desc = &rocke_tensor_descriptor_naive_dynamic_mixed( + b, "Q", coord_names, 3, strides, static_strides) + ->base; } for(li = 0; li < Q_VECS_PER_THREAD; ++li) @@ -164,12 +165,11 @@ void rocke_gfx950_attn2d_emit_q_load(rocke_gfx950_attn2d_build_ctx_t* ctx) rocke_value_t* qh_mod = rocke_b_mod(b, Q_row, rocke_b_const_i32(b, (int64_t)ctx->NQK)); qh_t = rocke_b_add(b, qh_mul, qh_mod); } - /* qmask_t = (q_pos_t < cur_batch_q_len) && (qh_t < NUM_QH) [&& (Q_row < VALID_ROWS)]. + /* qmask_t = (q_pos_t < cur_batch_q_len) && (qh_t < num_query_heads) [&& (Q_row < VALID_ROWS)]. * Skip the VALID_ROWS guard when VALID_ROWS == BLOCK_M (NQK divides BLOCK_M). */ { rocke_value_t* lt_pos = rocke_b_cmp_lt(b, q_pos_t, ctx->cur_batch_q_len); - rocke_value_t* lt_qh - = rocke_b_cmp_lt(b, qh_t, rocke_b_const_i32(b, (int64_t)ctx->NUM_QH)); + rocke_value_t* lt_qh = rocke_b_cmp_lt(b, qh_t, ctx->num_qh_p); rocke_value_t* and_pos_qh = rocke_b_land(b, lt_pos, lt_qh); if(ctx->VALID_ROWS < ctx->BLOCK_M) { @@ -248,7 +248,6 @@ void rocke_gfx950_attn2d_emit_loop_bounds_and_inits(rocke_gfx950_attn2d_build_ct const int BLOCK_M = ctx->BLOCK_M; const int NQK = ctx->NQK; const int SLIDING_WINDOW = ctx->SLIDING_WINDOW; - const int NUM_QH = ctx->NUM_QH; const int SOFTMAX_STATE_SLOTS = ctx->SOFTMAX_STATE_SLOTS; /* ---- max_seq_prefix_len (1221-1223) ---- */ @@ -334,7 +333,7 @@ void rocke_gfx950_attn2d_emit_loop_bounds_and_inits(rocke_gfx950_attn2d_build_ct rocke_value_t* qh_mul = rocke_b_mul(B, ctx->kv_head_idx, rocke_b_const_i32(B, NQK)); rocke_value_t* qh_mod = rocke_b_mod(B, row, rocke_b_const_i32(B, NQK)); rocke_value_t* qh = rocke_b_add(B, qh_mul, qh_mod); - rocke_value_t* qh_in = rocke_b_cmp_lt(B, qh, rocke_b_const_i32(B, NUM_QH)); + rocke_value_t* qh_in = rocke_b_cmp_lt(B, qh, ctx->num_qh_p); rocke_value_t* sink_h = rocke_b_global_load(B, ctx->sinks, qh, ctx->dtype, 2); rocke_value_t* sink_f = rocke_b_fmul(B, rocke_b_cast_to_f32(B, sink_h), rcp_ln2); m_inits[0] = rocke_b_select(B, qh_in, sink_f, neg_inf); @@ -349,7 +348,7 @@ void rocke_gfx950_attn2d_emit_loop_bounds_and_inits(rocke_gfx950_attn2d_build_ct rocke_value_t* qh_mul = rocke_b_mul(B, ctx->kv_head_idx, rocke_b_const_i32(B, NQK)); rocke_value_t* qh_mod = rocke_b_mod(B, row, rocke_b_const_i32(B, NQK)); rocke_value_t* qh = rocke_b_add(B, qh_mul, qh_mod); - rocke_value_t* qh_in = rocke_b_cmp_lt(B, qh, rocke_b_const_i32(B, NUM_QH)); + rocke_value_t* qh_in = rocke_b_cmp_lt(B, qh, ctx->num_qh_p); rocke_value_t* sink_h = rocke_b_global_load(B, ctx->sinks, qh, ctx->dtype, 2); rocke_value_t* sink_f = rocke_b_fmul(B, rocke_b_cast_to_f32(B, sink_h), rcp_ln2); m_inits[r] = rocke_b_select(B, qh_in, sink_f, neg_inf); @@ -626,6 +625,20 @@ void rocke_gfx950_attn2d_emit_preloop_prefetch(rocke_gfx950_attn2d_build_ctx_t* if(ctx == NULL) return; + /* Page ids of the first tile and of the one after it, carried into the + * loop; iteration i then reads the ids of tile i + 2. */ + const bool page_prefetch = rocke_gfx950_attn2d_page_prefetch(ctx); + if(page_prefetch) + { + rocke_value_t* c_one = rocke_b_const_i32(B, 1); + rocke_value_t* start_next_raw = rocke_b_add(B, ctx->tile_start, c_one); + rocke_value_t* in_range = rocke_b_cmp_lt(B, start_next_raw, ctx->tile_end); + rocke_value_t* start_next = rocke_b_select(B, in_range, start_next_raw, ctx->tile_start); + rocke_gfx950_attn2d_tile_page_ids(ctx, ctx->tile_start, ctx->page_start); + rocke_gfx950_attn2d_tile_page_ids(ctx, start_next, ctx->page_start_next); + ctx->page_start_tile = ctx->tile_start; + } + /* Prefetch tile_start's K into buffer 0 BEFORE the loop (2384). */ rocke_gfx950_attn2d_issue_k(ctx, ctx->tile_start, rocke_b_const_i32(B, 0)); @@ -636,6 +649,36 @@ void rocke_gfx950_attn2d_emit_preloop_prefetch(rocke_gfx950_attn2d_build_ctx_t* ctx->iter_args_names[idx] = ctx->iter_args_name_buf[idx]; ctx->iter_args[idx] = rocke_b_const_i32(B, 0); ctx->iter_args_count = idx + 1; + if(page_prefetch + && ctx->iter_args_count + 2 * ctx->N_BLOCKS_PER_TILE > ROCKE_GFX950_ATTN2D_MAX_ITER_ARGS) + { + B->status = ROCKE_ERR_VALUE; + snprintf(B->err, + ROCKE_ERR_MSG_CAP, + "%d loop carries exceed the %d-entry iter_args table", + ctx->iter_args_count + 2 * ctx->N_BLOCKS_PER_TILE, + ROCKE_GFX950_ATTN2D_MAX_ITER_ARGS); + return; + } + if(page_prefetch) + { + for(int j = 0; j < ctx->N_BLOCKS_PER_TILE; ++j) + { + idx = ctx->iter_args_count++; + snprintf( + ctx->iter_args_name_buf[idx], sizeof(ctx->iter_args_name_buf[0]), "page_cur%d", j); + ctx->iter_args_names[idx] = ctx->iter_args_name_buf[idx]; + ctx->iter_args[idx] = ctx->page_start[j]; + } + for(int j = 0; j < ctx->N_BLOCKS_PER_TILE; ++j) + { + idx = ctx->iter_args_count++; + snprintf( + ctx->iter_args_name_buf[idx], sizeof(ctx->iter_args_name_buf[0]), "page_next%d", j); + ctx->iter_args_names[idx] = ctx->iter_args_name_buf[idx]; + ctx->iter_args[idx] = ctx->page_start_next[j]; + } + } /* kv_step (2482) is emitted by the caller AFTER emit_licm_hoist to keep the * const in Python emission order; see rocke_gfx950_attn2d_emit_kv_step. */ diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/gfx950/attention_tiled_3d_public_glue.cpp b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/gfx950/attention_tiled_3d_public_glue.cpp index 65f68f628c46..d9986d95ab64 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/gfx950/attention_tiled_3d_public_glue.cpp +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/gfx950/attention_tiled_3d_public_glue.cpp @@ -9,7 +9,7 @@ * SCOPE (this TU only): * - The shared spec PODs' _default() / _validate() (__post_init__) and every * @property accessor (num_queries_per_kv, block_m=16, block_q, tile_size, - * dtype_ir, binary_search_iters, kernel_name via kernel_name_join). These + * dtype_ir, kernel_name via kernel_name_join). These * spec struct/default definitions are GUARDED against the gfx942 3D header's * duplicate definitions (the two arch ports share the identical frozen * dataclass shape -- see the SHARED SPEC POD note in @@ -412,7 +412,12 @@ int rocke_gfx950_unified_attention_3d_tiled_spec_kernel_name( snprintf(d_part, sizeof(d_part), "d%d", s->head_size); snprintf(b_part, sizeof(b_part), "b%d", s->block_size); - snprintf(h_part, sizeof(h_part), "h%dkv%d", s->num_query_heads, s->num_kv_heads); + /* Head counts are kernargs; only the GQA ratio is baked. */ + if(s->num_kv_heads <= 0) + { + return -1; + } + snprintf(h_part, sizeof(h_part), "g%d", s->num_query_heads / s->num_kv_heads); snprintf(seg_part, sizeof(seg_part), "seg%d", s->num_segments); parts[np++] = "rocke_uattn3d_tiled"; @@ -442,9 +447,10 @@ int rocke_gfx950_unified_attention_3d_tiled_spec_kernel_name( { parts[np++] = "sinks"; } + /* The window size is a kernarg; only whether the body masks one is baked. */ if(s->sliding_window > 0) { - snprintf(sw_part, sizeof(sw_part), "sw%d", s->sliding_window); + snprintf(sw_part, sizeof(sw_part), "sw"); parts[np++] = sw_part; } if(s->has_softcap) @@ -502,10 +508,9 @@ const rocke_type_t* rocke_gfx950_unified_attention_reduce_tiled_spec_dtype_ir( int rocke_gfx950_unified_attention_reduce_tiled_spec_kernel_name( const rocke_unified_attention_reduce_tiled_spec_t* s, char* buf, size_t cap) { - /* kernel_name_join("rocke_uattn_reduce_tiled", d, h, - * seg, dtype) */ + /* kernel_name_join("rocke_uattn_reduce_tiled", d, seg, dtype); + * the head count is a kernarg. */ char d_part[32]; - char h_part[32]; char seg_part[32]; const char* parts[8]; size_t np = 0; @@ -518,12 +523,10 @@ int rocke_gfx950_unified_attention_reduce_tiled_spec_kernel_name( } snprintf(d_part, sizeof(d_part), "d%d", s->head_size); - snprintf(h_part, sizeof(h_part), "h%d", s->num_query_heads); snprintf(seg_part, sizeof(seg_part), "seg%d", s->num_segments); parts[np++] = "rocke_uattn_reduce_tiled"; parts[np++] = d_part; - parts[np++] = h_part; parts[np++] = seg_part; parts[np++] = (s->dtype != NULL) ? s->dtype : ""; @@ -707,8 +710,6 @@ bool rocke_gfx950_attn_tiled_3d_config_from_spec( out->BLOCK_M = rocke_gfx950_unified_attention_3d_tiled_spec_block_m(spec); out->BLOCK_Q = rocke_gfx950_unified_attention_3d_tiled_spec_block_q(spec); out->NQK = rocke_gfx950_unified_attention_3d_tiled_spec_num_queries_per_kv(spec); - out->NUM_KV = spec->num_kv_heads; - out->NUM_QH = spec->num_query_heads; out->NUM_SEG = spec->num_segments; out->SLIDING_WINDOW = spec->sliding_window; out->USE_SOFTCAP = spec->has_softcap; @@ -742,8 +743,6 @@ bool rocke_gfx950_attn_tiled_3d_config_from_spec( out->PV_N_TILES = out->HD / ROCKE_ATTN3D950_MFMA_N; out->THREADS = ROCKE_ATTN3D950_SEGMENT_THREADS; - out->binary_search_iters - = rocke_gfx950_unified_attention_3d_tiled_spec_binary_search_iters(spec); if(out->PV_N_TILES < 0 || out->PV_N_TILES > 16) { @@ -770,8 +769,6 @@ bool rocke_gfx950_attn_tiled_3d_config_from_spec( } out->kv_calls_per_tile = (out->T * out->HD) / out->KV_HALVES_PER_CALL; out->bytes_per_call = out->KV_HALVES_PER_CALL * 2; - out->kv_stride_blk_b = out->BS * out->NUM_KV * out->HD * out->KV_BYTES; - out->kv_stride_tok_b = out->NUM_KV * out->HD * out->KV_BYTES; out->kv_stride_h_b = out->HD * out->KV_BYTES; out->bytes_per_buf = out->T * out->HD * 2; out->N_BLOCKS_PER_TILE = out->BS > 0 ? out->T / out->BS : 0; @@ -796,6 +793,15 @@ bool rocke_gfx950_attn_tiled_3d_config_from_spec( return false; } out->calls_per_block = out->BS / out->tokens_per_call; + /* The prologue caches each call's in-page token offset in call_tok_b. */ + if(out->calls_per_block > 64) + { + rocke_i_set_err(b, + ROCKE_ERR_VALUE, + "%d calls per page exceed the 64-entry offset table", + out->calls_per_block); + return false; + } } out->fp8_elems_per_chunk = 8; @@ -856,8 +862,6 @@ bool rocke_gfx950_attn_tiled_3d_reduce_config_from_spec( out->HD = spec->head_size; out->NUM_SEG = spec->num_segments; - out->NUM_QH = spec->num_query_heads; - out->NUM_KV = spec->num_kv_heads; out->dtype = rocke_gfx950_unified_attention_reduce_tiled_spec_dtype_ir(spec); out->THREADS = 64; diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/gfx950/attention_tiled_3d_reduce.cpp b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/gfx950/attention_tiled_3d_reduce.cpp index 0ea23c18ea85..339d6c44fc89 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/gfx950/attention_tiled_3d_reduce.cpp +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/gfx950/attention_tiled_3d_reduce.cpp @@ -135,7 +135,6 @@ void rocke_gfx950_attention_tiled_3d_reduce_declare_and_prologue( const rocke_gfx950_attn_tiled_3d_config_t* cfg = &ctx->cfg; rocke_param_opts_t opts; - int NUM_QH = cfg->NUM_QH; int NUM_SEG = cfg->NUM_SEG; int HD = cfg->HD; @@ -189,6 +188,9 @@ void rocke_gfx950_attention_tiled_3d_reduce_declare_and_prologue( ctx->red_seq_lens = rocke_b_param(b, "seq_lens_ptr", rocke_ptr_type(b, rocke_i32(), "global"), &opts); + /* RUNTIME_KERNARG_FIELDS : the query head count */ + ctx->num_qh_p = rocke_b_param(b, "num_query_heads", rocke_i32(), NULL); + /* ---- grid ids + thread (lines 1029-1031) ---- */ ctx->q_token = rocke_b_block_id_x(b); ctx->q_head = rocke_b_block_id_y(b); @@ -198,24 +200,34 @@ void rocke_gfx950_attention_tiled_3d_reduce_declare_and_prologue( ctx->neg_inf = rocke_b_const_f32(b, -INFINITY); ctx->zero_f = rocke_b_const_f32(b, 0.0); - /* ---- descriptors (lines 1042-1056) ---- */ + /* ---- descriptors (lines 1042-1056) ---- + * The head count is a kernarg, so each token stride is too. */ { - int ml_lengths[3] = {1 << 30, NUM_QH, NUM_SEG}; - const char* ml_coords[3] = {"token", "head", "seg"}; - ctx->ml_desc_red - = rocke_tensor_descriptor_naive(b, "segm_ml", ml_lengths, 3, NULL, ml_coords, 3); + static const char* const ml_coords[3] = {"token", "head", "seg"}; + rocke_value_t* ml_dyn[3] + = {rocke_b_mul(b, ctx->num_qh_p, rocke_b_const_i32(b, NUM_SEG)), NULL, NULL}; + const int ml_static[3] = {1, NUM_SEG, 1}; + ctx->ml_desc_red = &rocke_tensor_descriptor_naive_dynamic_mixed( + b, "segm_ml", ml_coords, 3, ml_dyn, ml_static) + ->base; } { - int sa_lengths[4] = {1 << 30, NUM_QH, NUM_SEG, HD}; - const char* sa_coords[4] = {"token", "head", "seg", "dim"}; - ctx->seg_acc_desc_red - = rocke_tensor_descriptor_naive(b, "segm_output", sa_lengths, 4, NULL, sa_coords, 4); + static const char* const sa_coords[4] = {"token", "head", "seg", "dim"}; + rocke_value_t* sa_dyn[4] + = {rocke_b_mul(b, ctx->num_qh_p, rocke_b_const_i32(b, NUM_SEG * HD)), NULL, NULL, NULL}; + const int sa_static[4] = {1, NUM_SEG * HD, HD, 1}; + ctx->seg_acc_desc_red = &rocke_tensor_descriptor_naive_dynamic_mixed( + b, "segm_output", sa_coords, 4, sa_dyn, sa_static) + ->base; } { - int out_lengths[3] = {1 << 30, NUM_QH, HD}; - const char* out_coords[3] = {"token", "head", "dim"}; - ctx->out_desc_red - = rocke_tensor_descriptor_naive(b, "out", out_lengths, 3, NULL, out_coords, 3); + static const char* const out_coords[3] = {"token", "head", "dim"}; + rocke_value_t* out_dyn[3] + = {rocke_b_mul(b, ctx->num_qh_p, rocke_b_const_i32(b, HD)), NULL, NULL}; + const int out_static[3] = {1, HD, 1}; + ctx->out_desc_red = &rocke_tensor_descriptor_naive_dynamic_mixed( + b, "out", out_coords, 3, out_dyn, out_static) + ->base; } /* ---- base_ml = ml_desc_red.offset(token=q_token, head=q_head, seg=0) ---- diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/gfx950/attention_tiled_3d_segment_loop_epilogue.cpp b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/gfx950/attention_tiled_3d_segment_loop_epilogue.cpp index c3d0d8f85991..9df6fbda3c37 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/gfx950/attention_tiled_3d_segment_loop_epilogue.cpp +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/gfx950/attention_tiled_3d_segment_loop_epilogue.cpp @@ -116,7 +116,7 @@ void rocke_gfx950_attention_tiled_3d_emit_loop_init( rocke_value_t* qh_mul = rocke_b_mul(b, ctx->kv_head_idx, rocke_b_const_i32(b, cfg->NQK)); rocke_value_t* qh_mod = rocke_b_mod(b, ctx->lane_col, rocke_b_const_i32(b, cfg->NQK)); rocke_value_t* qh = rocke_b_add(b, qh_mul, qh_mod); - rocke_value_t* qh_in = rocke_b_cmp_lt(b, qh, rocke_b_const_i32(b, cfg->NUM_QH)); + rocke_value_t* qh_in = rocke_b_cmp_lt(b, qh, ctx->num_qh_p); rocke_value_t* sink_h = rocke_b_global_load(b, ctx->sinks, qh, dtype, 2); rocke_value_t* sink_f = rocke_b_fmul(b, rocke_b_cast_to_f32(b, sink_h), ctx->rcp_ln2); rocke_value_t* sink_with_mask = rocke_b_select(b, qh_in, sink_f, ctx->neg_inf); @@ -248,14 +248,13 @@ void rocke_gfx950_attention_tiled_3d_emit_softmax_loop( qh_mod = rocke_b_mod(b, ctx->lane_col, rocke_b_const_i32(b, cfg->NQK)); ctx->row_qh = rocke_b_add(b, qh_mul, qh_mod); ok_a = rocke_b_cmp_lt(b, ctx->row_qp, ctx->cur_batch_q_len); - ok_b = rocke_b_cmp_lt(b, ctx->row_qh, rocke_b_const_i32(b, cfg->NUM_QH)); + ok_b = rocke_b_cmp_lt(b, ctx->row_qh, ctx->num_qh_p); ctx->row_ok = rocke_b_land(b, ok_a, ok_b); ctx->row_causal_lim = rocke_b_add(b, ctx->context_len, ctx->row_qp); ctx->row_alibi = NULL; if(cfg->USE_ALIBI) { - rocke_value_t* qh_ok - = rocke_b_cmp_lt(b, ctx->row_qh, rocke_b_const_i32(b, cfg->NUM_QH)); + rocke_value_t* qh_ok = rocke_b_cmp_lt(b, ctx->row_qh, ctx->num_qh_p); ctx->row_alibi = rocke_b_masked_global_load( b, ctx->alibi_slopes_ptr, ctx->row_qh, qh_ok, rocke_b_const_f32(b, 0.0), f32, 4); } diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/gfx950/attention_tiled_3d_segment_prologue.cpp b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/gfx950/attention_tiled_3d_segment_prologue.cpp index b911ba358e04..50b7607d7e1e 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/gfx950/attention_tiled_3d_segment_prologue.cpp +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/gfx950/attention_tiled_3d_segment_prologue.cpp @@ -328,6 +328,9 @@ void rocke_gfx950_attention_tiled_3d_declare_params( ctx->num_seqs_p = rocke_b_param(B, "num_seqs", rocke_i32(), NULL); ctx->bt_stride_p = rocke_b_param(B, "block_table_stride", rocke_i32(), NULL); ctx->qq_bias_stride0_p = rocke_b_param(B, "qq_bias_stride_0", rocke_i32(), NULL); + ctx->num_qh_p = rocke_b_param(B, "num_query_heads", rocke_i32(), NULL); + ctx->num_kv_p = rocke_b_param(B, "num_kv_heads", rocke_i32(), NULL); + ctx->sliding_window_p = rocke_b_param(B, "sliding_window", rocke_i32(), NULL); if(ctx->strided_kv) { ctx->k_strides = rocke_declare_strided_kv_params(B, "k"); @@ -344,7 +347,6 @@ void rocke_gfx950_attention_tiled_3d_emit_prologue(rocke_gfx950_attention_tiled_ const int T = CFG.T; const int BLOCK_M = CFG.BLOCK_M; const int BLOCK_Q = CFG.BLOCK_Q; - const int NUM_QH = CFG.NUM_QH; const int NUM_SEG = CFG.NUM_SEG; const rocke_type_t* dtype = CFG.dtype; @@ -363,13 +365,9 @@ void rocke_gfx950_attention_tiled_3d_emit_prologue(rocke_gfx950_attention_tiled_ } else { - ctx->seq_idx = rocke_binary_search_seq_idx(B, - ctx->cu_q, - ctx->q_block_global_idx, - ctx->num_seqs_p, - BLOCK_Q, - CFG.binary_search_iters, - false); + rocke_value_t* trip = rocke_binary_search_seq_idx_trip_count(B, ctx->num_seqs_p); + ctx->seq_idx = rocke_binary_search_seq_idx_dyn( + B, ctx->cu_q, ctx->q_block_global_idx, ctx->num_seqs_p, BLOCK_Q, trip, false); } /* ---- cu_q bounds / per-sequence geometry (lines 373-379) ---- */ @@ -410,23 +408,34 @@ void rocke_gfx950_attention_tiled_3d_emit_prologue(rocke_gfx950_attention_tiled_ ctx->tps = rocke_b_div(B, tps_num, rocke_b_const_i32(B, NUM_SEG * T)); } - /* ---- descriptors (lines 399-413) ---- */ + /* ---- descriptors (lines 399-413) ---- + * The head count is a kernarg, so each token stride is too. */ { - const int ml_lengths[3] = {1 << 30, NUM_QH, NUM_SEG}; - const char* ml_coords[3] = {"token", "head", "seg"}; - ctx->ml_desc - = rocke_tensor_descriptor_naive(B, "segm_ml", ml_lengths, 3, NULL, ml_coords, 3); + static const char* const ml_coords[3] = {"token", "head", "seg"}; + rocke_value_t* ml_dyn[3] + = {rocke_b_mul(B, ctx->num_qh_p, rocke_b_const_i32(B, NUM_SEG)), NULL, NULL}; + const int ml_static[3] = {1, NUM_SEG, 1}; + ctx->ml_desc = &rocke_tensor_descriptor_naive_dynamic_mixed( + B, "segm_ml", ml_coords, 3, ml_dyn, ml_static) + ->base; } { - const int sa_lengths[4] = {1 << 30, NUM_QH, NUM_SEG, HD}; - const char* sa_coords[4] = {"token", "head", "seg", "dim"}; - ctx->seg_acc_desc - = rocke_tensor_descriptor_naive(B, "segm_output", sa_lengths, 4, NULL, sa_coords, 4); + static const char* const sa_coords[4] = {"token", "head", "seg", "dim"}; + rocke_value_t* sa_dyn[4] + = {rocke_b_mul(B, ctx->num_qh_p, rocke_b_const_i32(B, NUM_SEG * HD)), NULL, NULL, NULL}; + const int sa_static[4] = {1, NUM_SEG * HD, HD, 1}; + ctx->seg_acc_desc = &rocke_tensor_descriptor_naive_dynamic_mixed( + B, "segm_output", sa_coords, 4, sa_dyn, sa_static) + ->base; } { - const int q_lengths[3] = {1 << 30, NUM_QH, HD}; - const char* q_coords[3] = {"token", "head", "dim"}; - ctx->q_desc = rocke_tensor_descriptor_naive(B, "Q", q_lengths, 3, NULL, q_coords, 3); + static const char* const q_coords[3] = {"token", "head", "dim"}; + rocke_value_t* q_dyn[3] + = {rocke_b_mul(B, ctx->num_qh_p, rocke_b_const_i32(B, HD)), NULL, NULL}; + const int q_static[3] = {1, HD, 1}; + ctx->q_desc + = &rocke_tensor_descriptor_naive_dynamic_mixed(B, "Q", q_coords, 3, q_dyn, q_static) + ->base; } /* seg_start_tile_pos = seg_idx*tps*T (line 415) */ @@ -466,7 +475,7 @@ void rocke_gfx950_attention_tiled_3d_emit_prologue(rocke_gfx950_attention_tiled_ ctx->one_f = rocke_b_const_f32(B, 1.0); ctx->rcp_ln2 = rocke_b_const_f32(B, 1.4426950408889634); ctx->qk_scale = rocke_b_fmul(B, ctx->scale_p, ctx->rcp_ln2); - ctx->sw_const = rocke_b_const_i32(B, CFG.SLIDING_WINDOW); + ctx->sw_const = ctx->sliding_window_p; ctx->z8 = rocke_b_zero_vec(B, dtype, 8); /* ---- Q -> LDS (lines 482-512) ---- */ @@ -519,7 +528,6 @@ void rocke_gfx950_attention_tiled_3d_emit_async_infra( { const int HD = CFG.HD; const int T = CFG.T; - const int NUM_KV = CFG.NUM_KV; /* ---- async DMA infra (lines 565-585) ---- */ ctx->big_bytes = rocke_b_const_i32(B, 0x7FFF0000); @@ -535,10 +543,12 @@ void rocke_gfx950_attention_tiled_3d_emit_async_infra( ctx->key_rsrc = rocke_b_buffer_rsrc(B, ctx->key, ctx->big_bytes); ctx->value_rsrc = rocke_b_buffer_rsrc(B, ctx->value, ctx->big_bytes); } - /* kv_block_bytes_c = const_i32(kv_stride_blk_b): one-block buffer bound for - * the i64-addressing path (Python creates it unconditionally right after the - * byte strides; unused in the i32 path, where it is DCE'd). */ - ctx->kv_block_bytes_c = rocke_b_const_i32(B, CFG.kv_stride_blk_b); + /* The token and page byte strides scale with the runtime kv head count. + * kv_block_bytes_c is the page stride: the one-block buffer bound for the + * i64-addressing path (unused in the i32 path, where it is DCE'd). */ + ctx->kv_stride_tok_b = rocke_b_mul(B, ctx->num_kv_p, rocke_b_const_i32(B, HD * CFG.KV_BYTES)); + ctx->kv_stride_blk_b = rocke_b_mul(B, ctx->kv_stride_tok_b, rocke_b_const_i32(B, CFG.BS)); + ctx->kv_block_bytes_c = ctx->kv_stride_blk_b; ctx->lane_half_base = rocke_b_mul(B, ctx->tid, rocke_b_const_i32(B, CFG.HALVES_PER_LANE)); ctx->K_lds_addr = rocke_b_smem_addr_of(B, ctx->K_lds); ctx->V_lds_addr = rocke_b_smem_addr_of(B, ctx->V_lds); @@ -546,6 +556,8 @@ void rocke_gfx950_attention_tiled_3d_emit_async_infra( ctx->seq_pages = NULL; memset(ctx->lane_within_b, 0, sizeof(ctx->lane_within_b)); + memset(ctx->call_tok_b, 0, sizeof(ctx->call_tok_b)); + ctx->kv_stride_blk_i64 = NULL; ctx->paged_kv_desc = NULL; /* seq_base = seq_idx * bt_stride_p (line 592) */ if(ctx->strided_kv) @@ -555,15 +567,15 @@ void rocke_gfx950_attention_tiled_3d_emit_async_infra( ctx->seq_base = rocke_b_mul(B, ctx->seq_idx, ctx->bt_stride_p); if(!CFG.PAGE_IDS) { - /* paged_kv_desc = naive("paged_kv_bytes", [1<<24, T, NUM_KV, HD], byte + /* paged_kv_desc = DynamicTensorDescriptor("paged_kv_bytes", byte * strides).transform(indirect("tile_idx" -> "physical_block", * base=seq_base), unmerge("linear_half" -> ("token","dim"), (T, HD))) */ - const int lengths[4] = {1 << 24, T, NUM_KV, HD}; - const int strides[4] - = {CFG.kv_stride_blk_b, CFG.kv_stride_tok_b, CFG.kv_stride_h_b, CFG.KV_BYTES}; - const char* coords[4] = {"physical_block", "token", "kv_head", "dim"}; - rocke_tensor_descriptor_t* base - = rocke_tensor_descriptor_naive(B, "paged_kv_bytes", lengths, 4, strides, coords, 4); + const int strides[4] = {1, 1, CFG.kv_stride_h_b, CFG.KV_BYTES}; + rocke_value_t* dyn[4] = {ctx->kv_stride_blk_b, ctx->kv_stride_tok_b, NULL, NULL}; + static const char* const coords[4] = {"physical_block", "token", "kv_head", "dim"}; + rocke_tensor_descriptor_t* base = &rocke_tensor_descriptor_naive_dynamic_mixed( + B, "paged_kv_bytes", coords, 4, dyn, strides) + ->base; const rocke_transform_t* chain[2]; const char* into_td[2] = {"token", "dim"}; const int dims_td[2] = {T, HD}; @@ -587,7 +599,7 @@ void rocke_gfx950_attention_tiled_3d_emit_async_infra( { rocke_value_t* lane_tok = rocke_b_div(B, ctx->lane_half_base, rocke_b_const_i32(B, HD)); rocke_value_t* lane_dim = rocke_b_mod(B, ctx->lane_half_base, rocke_b_const_i32(B, HD)); - rocke_value_t* tok_b = rocke_b_mul(B, lane_tok, rocke_b_const_i32(B, CFG.kv_stride_tok_b)); + rocke_value_t* tok_b = rocke_b_mul(B, lane_tok, ctx->kv_stride_tok_b); rocke_value_t* head_b = rocke_b_mul(B, ctx->kv_head_idx, rocke_b_const_i32(B, CFG.kv_stride_h_b)); rocke_value_t* tok_head = rocke_b_add(B, tok_b, head_b); @@ -616,6 +628,22 @@ void rocke_gfx950_attention_tiled_3d_emit_async_infra( rocke_value_t* dim_b = rocke_b_mul(B, lane_dim, rocke_b_const_i32(B, CFG.KV_BYTES)); ctx->lane_within_b[0] = rocke_b_add(B, tok_head, dim_b); } + /* The token stride is a kernarg, so each call's token offset within the + * page is a multiply. They are wave-uniform and loop-invariant: build + * them once here rather than in every tile's copies. */ + { + int call; + for(call = 0; call < CFG.kv_calls_per_tile; ++call) + { + const int tok_i = call % CFG.calls_per_block; + if(tok_i == 0 || ctx->call_tok_b[tok_i] != NULL) + continue; + ctx->call_tok_b[tok_i] = rocke_b_mul( + B, ctx->kv_stride_tok_b, rocke_b_const_i32(B, tok_i * CFG.tokens_per_call)); + } + } + if(CFG.I64_KV_ADDR) + ctx->kv_stride_blk_i64 = rocke_b_zext(B, ctx->kv_stride_blk_b, rocke_i64()); } } @@ -775,17 +803,17 @@ static void rocke__issue_paged_tile(rocke_gfx950_attention_tiled_3d_build_ctx_t* for(call = 0; call < CFG.kv_calls_per_tile; ++call) { const int blk = call / CFG.calls_per_block; - const int tok0 = (call % CFG.calls_per_block) * CFG.tokens_per_call; + const int tok_i = call % CFG.calls_per_block; const int ph_i = CFG.KV_SWZ ? (call % (CFG.KV_CHUNKS_PER_ROW / CFG.tokens_per_call)) : 0; - rocke_value_t* voff = rocke_b_add( - B, ctx->lane_within_b[ph_i], rocke_b_const_i32(B, tok0 * CFG.kv_stride_tok_b)); + rocke_value_t* voff = ctx->lane_within_b[ph_i]; + if(tok_i != 0) + voff = rocke_b_add(B, voff, ctx->call_tok_b[tok_i]); rocke_value_t* call_rsrc; rocke_value_t* soff; if(CFG.I64_KV_ADDR) { rocke_value_t* pid_i64 = rocke_b_zext(B, page_ids[blk], rocke_i64()); - rocke_value_t* base_i64 - = rocke_b_mul(B, pid_i64, rocke_b_const_i64(B, (int64_t)CFG.kv_stride_blk_b)); + rocke_value_t* base_i64 = rocke_b_mul(B, pid_i64, ctx->kv_stride_blk_i64); call_rsrc = rocke_b_buffer_rsrc( B, rocke_b_global_ptr_add(B, src, base_i64), ctx->kv_block_bytes_c); soff = ctx->zero_soff; @@ -793,7 +821,7 @@ static void rocke__issue_paged_tile(rocke_gfx950_attention_tiled_3d_build_ctx_t* else { call_rsrc = src_rsrc; - soff = rocke_b_mul(B, page_ids[blk], rocke_b_const_i32(B, CFG.kv_stride_blk_b)); + soff = rocke_b_mul(B, page_ids[blk], ctx->kv_stride_blk_b); } rocke_value_t* dst = rocke_b_smem_ptr_add( B, buf_base, rocke_b_const_i64(B, (int64_t)call * CFG.bytes_per_call)); @@ -809,7 +837,6 @@ void rocke_gfx950_attention_tiled_3d_emit_early_zero_fill( rocke_gfx950_attention_tiled_3d_build_ctx_t* ctx) { const int NQK = CFG.NQK; - const int NUM_QH = CFG.NUM_QH; const int PV_N_TILES = CFG.PV_N_TILES; int n, reg; @@ -832,7 +859,7 @@ void rocke_gfx950_attention_tiled_3d_emit_early_zero_fill( rocke_value_t* qh_r_mod = rocke_b_mod(B, row, rocke_b_const_i32(B, NQK)); rocke_value_t* qh_r = rocke_b_add(B, qh_r_mul, qh_r_mod); rocke_value_t* row_ok_a = rocke_b_cmp_lt(B, qp_r, ctx->cur_batch_q_len); - rocke_value_t* row_ok_b = rocke_b_cmp_lt(B, qh_r, rocke_b_const_i32(B, NUM_QH)); + rocke_value_t* row_ok_b = rocke_b_cmp_lt(B, qh_r, ctx->num_qh_p); rocke_value_t* row_ok = rocke_b_land(B, row_ok_a, row_ok_b); rocke_value_t* qp_r_safe = rocke_b_select(B, row_ok, qp_r, rocke_b_const_i32(B, 0)); rocke_value_t* qh_r_safe = rocke_b_select(B, row_ok, qh_r, rocke_b_const_i32(B, 0)); @@ -865,7 +892,7 @@ void rocke_gfx950_attention_tiled_3d_emit_early_zero_fill( rocke_value_t* qh_r_mod = rocke_b_mod(B, row, rocke_b_const_i32(B, NQK)); rocke_value_t* qh_r = rocke_b_add(B, qh_r_mul, qh_r_mod); rocke_value_t* row_ok_a = rocke_b_cmp_lt(B, qp_r, ctx->cur_batch_q_len); - rocke_value_t* row_ok_b = rocke_b_cmp_lt(B, qh_r, rocke_b_const_i32(B, NUM_QH)); + rocke_value_t* row_ok_b = rocke_b_cmp_lt(B, qh_r, ctx->num_qh_p); rocke_value_t* row_ok = rocke_b_land(B, row_ok_a, row_ok_b); rocke_value_t* qp_r_safe = rocke_b_select(B, row_ok, qp_r, rocke_b_const_i32(B, 0)); @@ -889,7 +916,6 @@ void rocke_gfx950_attention_tiled_3d_emit_early_zero_fill( void rocke_gfx950_attention_tiled_3d_emit_q_to_lds(rocke_gfx950_attention_tiled_3d_build_ctx_t* ctx) { const int NQK = CFG.NQK; - const int NUM_QH = CFG.NUM_QH; const int THREADS = CFG.THREADS; const int Q_VECS_PER_ROW = CFG.Q_VECS_PER_ROW; const int Q_VECS_PER_THREAD = CFG.Q_VECS_PER_THREAD; @@ -911,7 +937,7 @@ void rocke_gfx950_attention_tiled_3d_emit_q_to_lds(rocke_gfx950_attention_tiled_ rocke_value_t* qh_t_mod = rocke_b_mod(B, Q_row, rocke_b_const_i32(B, NQK)); rocke_value_t* qh_t = rocke_b_add(B, qh_t_mul, qh_t_mod); rocke_value_t* qmask_a = rocke_b_cmp_lt(B, q_pos_t, ctx->cur_batch_q_len); - rocke_value_t* qmask_b = rocke_b_cmp_lt(B, qh_t, rocke_b_const_i32(B, NUM_QH)); + rocke_value_t* qmask_b = rocke_b_cmp_lt(B, qh_t, ctx->num_qh_p); rocke_value_t* qmask_t = rocke_b_land(B, qmask_a, qmask_b); rocke_value_t* q_pos_safe = rocke_b_select(B, qmask_t, q_pos_t, rocke_b_const_i32(B, 0)); rocke_value_t* qh_safe = rocke_b_select(B, qmask_t, qh_t, rocke_b_const_i32(B, 0)); diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/transforms.py b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/transforms.py index 4bd703678cec..570b57323e48 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/transforms.py +++ b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/transforms.py @@ -69,6 +69,7 @@ from __future__ import annotations +import operator from dataclasses import dataclass, replace from typing import Any, Dict, List, Optional, Sequence, Tuple, Union @@ -1905,7 +1906,10 @@ def create( raise ValueError("coord_names and strides length mismatch") coord_names = tuple(coord_names) dynamic = tuple(s if isinstance(s, Value) else None for s in strides) - static = tuple(1 if isinstance(s, Value) else int(s) for s in strides) + # operator.index rejects a float or string stride instead of truncating it. + static = tuple( + 1 if isinstance(s, Value) else operator.index(s) for s in strides + ) return cls( name=name, base_names=coord_names, diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_transforms_dynamic.py b/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_transforms_dynamic.py index 95aa5468a120..2bb0094e1330 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_transforms_dynamic.py +++ b/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_transforms_dynamic.py @@ -617,6 +617,15 @@ def test_missing_chain_input_rejected(self): with pytest.raises(ValueError, match="missing upper coords"): desc.offset(b, x=x) + @pytest.mark.parametrize("bad", [3.75, "3"]) + def test_create_rejects_a_non_integer_stride(self, bad): + b = _builder() + (s0,) = _params(b, "s0") + with pytest.raises(TypeError): + DynamicTensorDescriptor.create( + "t", coord_names=("a", "b"), strides=(s0, bad) + ) + def test_create_length_mismatch_rejected(self): b = _builder() (s0,) = _params(b, "s0") 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 744e41dd6e7e..eca00575a5e9 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 @@ -76,186 +76,186 @@ }, "attention/gfx950/2d_bf16_d128_b64": { "arch": "gfx950", - "bytes": 167466, + "bytes": 139954, "family": "attention", - "kernel_name": "rocke_uattn2d_tiled_d128_b64_h4kv2_bf16", - "sha256": "0bbe7c757b4aa3d732775c27d14da275156e644f103de624011ccd6ed42f68b9", + "kernel_name": "rocke_uattn2d_tiled_d128_b64_g2_bf16", + "sha256": "1cd621d8b243a40e34ce7965e0a9e9021c7356ad105fad57e74b424e17b0c0f6", "status": "ok" }, "attention/gfx950/2d_bf16_d64_b32_gqa8_sinks_combo": { "arch": "gfx950", - "bytes": 116995, + "bytes": 117790, "family": "attention", - "kernel_name": "rocke_uattn2d_tiled_d64_b32_t64_h64kv8_bf16_sinks_w4_mw32_mfma32_stqk_s1_mask1_hlpv_skipqreg_mlim_fastkvdesc", - "sha256": "94e09ed1ce5b5ad97cd73408730b6c16a1bc0a0065c836ffe0d849cd1b7ee7af", + "kernel_name": "rocke_uattn2d_tiled_d64_b32_t64_g8_h64kv8_bf16_sinks_w4_mw32_mfma32_stqk_s1_mask1_hlpv_skipqreg_mlim_fastkvdesc", + "sha256": "5e6ac1abe5f83dbb53ee0d0dbb0e796204a229b5f9a7d08c64ace181a06e8806", "status": "ok" }, "attention/gfx950/2d_fp16_d128_b64": { "arch": "gfx950", - "bytes": 165889, + "bytes": 138377, "family": "attention", - "kernel_name": "rocke_uattn2d_tiled_d128_b64_h4kv2_fp16", - "sha256": "c5b0131c0a7446097c2ff53299a254ae6f52047a9b3977ffedca0823fefcdf21", + "kernel_name": "rocke_uattn2d_tiled_d128_b64_g2_fp16", + "sha256": "82f78e923d03767891be1810bbd3926b4c617db871c4321667fdf8fd12280d42", "status": "ok" }, "attention/gfx950/2d_fp16_d128_b64_alibi": { "arch": "gfx950", - "bytes": 171064, + "bytes": 143607, "family": "attention", - "kernel_name": "rocke_uattn2d_tiled_d128_b64_h4kv2_fp16_alibi", - "sha256": "f7664dfabbdcef3e5d5e501e6b058ee4872ee11e0354204c5506013f77453dcf", + "kernel_name": "rocke_uattn2d_tiled_d128_b64_g2_fp16_alibi", + "sha256": "96d3f747c3691b7cb8f44c4ddc9bbe787cd6a3be8a02078abf378d8f9b07a4e9", "status": "ok" }, "attention/gfx950/2d_fp16_d128_b64_left_only_sw64": { "arch": "gfx950", - "bytes": 167358, + "bytes": 139934, "family": "attention", - "kernel_name": "rocke_uattn2d_tiled_d128_b64_h4kv2_fp16_sw64_tl_rbu", - "sha256": "7c1664c4fc235ada95079912aa4a80f0a6a97ad63add9b693288e050007c2eae", + "kernel_name": "rocke_uattn2d_tiled_d128_b64_g2_fp16_sw_tl_rbu", + "sha256": "663a3e8ab2d1abc8c8fd9052be1daa6a1d478c767230a6b178bfd5089b49e98a", "status": "ok" }, "attention/gfx950/2d_fp16_d128_b64_no_mask": { "arch": "gfx950", - "bytes": 164521, + "bytes": 137005, "family": "attention", - "kernel_name": "rocke_uattn2d_tiled_d128_b64_h4kv2_fp16_rbu", - "sha256": "2f3938213b1d5d05b4243527cc2de470920721436cbeb032ecf366314b23af85", + "kernel_name": "rocke_uattn2d_tiled_d128_b64_g2_fp16_rbu", + "sha256": "13fb8518225982fccfcf7739a5f40ca6ac0b389467068f9841d467c1404810b8", "status": "ok" }, "attention/gfx950/2d_fp16_d128_b64_rb16": { "arch": "gfx950", - "bytes": 166668, + "bytes": 139187, "family": "attention", - "kernel_name": "rocke_uattn2d_tiled_d128_b64_h4kv2_fp16_rb16", - "sha256": "5f79198fe6c537d6a00242f4ef19e1b6c1ec44efc430584cbab2da12c54355ae", + "kernel_name": "rocke_uattn2d_tiled_d128_b64_g2_fp16_rb16", + "sha256": "9eaf6d69014ea46b249b3f47b2c705ef275892644b3777b124f37d80433adf98", "status": "ok" }, "attention/gfx950/2d_fp16_d128_b64_softcap": { "arch": "gfx950", - "bytes": 173921, + "bytes": 146373, "family": "attention", - "kernel_name": "rocke_uattn2d_tiled_d128_b64_h4kv2_fp16_softcap", - "sha256": "12dcee5c39cf7f0486d3a7418b2511ac3c78046d9e5dba42c4e35d649158b807", + "kernel_name": "rocke_uattn2d_tiled_d128_b64_g2_fp16_softcap", + "sha256": "b877ea971a20a69d48d96eda7a6d571ef5ebff0459f36a559a23ccacc1976c35", "status": "ok" }, "attention/gfx950/2d_fp16_d128_b64_top_left": { "arch": "gfx950", - "bytes": 165929, + "bytes": 138417, "family": "attention", - "kernel_name": "rocke_uattn2d_tiled_d128_b64_h4kv2_fp16_tl", - "sha256": "bfee42e8173c5a1d555048935e0dcc483cd1e506622e59acb71cfb5ace2895d7", + "kernel_name": "rocke_uattn2d_tiled_d128_b64_g2_fp16_tl", + "sha256": "949f1641ed1707e557629d58dd465e1b6bdc93f75a3ac520728a73238fcf2a37", "status": "ok" }, "attention/gfx950/2d_fp16_d128_b64_top_left_sw64": { "arch": "gfx950", - "bytes": 168719, + "bytes": 141281, "family": "attention", - "kernel_name": "rocke_uattn2d_tiled_d128_b64_h4kv2_fp16_sw64_tl", - "sha256": "0e93b29f935eb1fca8cb92b727c408f51c41bda825b15927498ff064763d2233", + "kernel_name": "rocke_uattn2d_tiled_d128_b64_g2_fp16_sw_tl", + "sha256": "4e087145bff9db19ae47b0b08f24394c83823a4076249026e0262e373d6c05f7", "status": "ok" }, "attention/gfx950/2d_fp16_d128_b64_two_sided_sw64_rb16": { "arch": "gfx950", - "bytes": 169536, + "bytes": 142120, "family": "attention", - "kernel_name": "rocke_uattn2d_tiled_d128_b64_h4kv2_fp16_sw64_tl_rb16", - "sha256": "6e7a95e3894e2fab38cf0ed15943ac01a5f0e448091253e56f7a92590a42b1ce", + "kernel_name": "rocke_uattn2d_tiled_d128_b64_g2_fp16_sw_tl_rb16", + "sha256": "20539f1b40b16d3c759bd1c747627d0753e4ebe598f39fc39b9da2777e1d5605", "status": "ok" }, "attention/gfx950/2d_fp16_d64_b32_gqa8_combo_no_mask": { "arch": "gfx950", - "bytes": 115746, + "bytes": 115142, "family": "attention", - "kernel_name": "rocke_uattn2d_tiled_d64_b32_t64_h64kv8_fp16_rbu_w4_mw32_mfma32_stqk_s1_mask1_hlpv_skipqreg_mlim", - "sha256": "32766ba93ebcba66bc4afb411a6991959521b7447502946ffe49bf8c955f2ab0", + "kernel_name": "rocke_uattn2d_tiled_d64_b32_t64_g8_fp16_rbu_w4_mw32_mfma32_stqk_s1_mask1_hlpv_skipqreg_mlim", + "sha256": "fa721ea3bde849fda8048d7594e18d547f6faadeab9a864cbf40371dbaf89600", "status": "ok" }, "attention/gfx950/2d_fp16_d64_b32_gqa8_combo_rb16": { "arch": "gfx950", - "bytes": 116082, + "bytes": 115495, "family": "attention", - "kernel_name": "rocke_uattn2d_tiled_d64_b32_t64_h64kv8_fp16_rb16_w4_mw32_mfma32_stqk_s1_mask1_hlpv_skipqreg_mlim", - "sha256": "ccaa9788e7d4bb9b295494f79c343f5163c9376ce67f110db9c515e57430aa80", + "kernel_name": "rocke_uattn2d_tiled_d64_b32_t64_g8_fp16_rb16_w4_mw32_mfma32_stqk_s1_mask1_hlpv_skipqreg_mlim", + "sha256": "d41de88c7ff3d56c58710c58770cea39d21072e27b6fd9739da660b0969a8635", "status": "ok" }, "attention/gfx950/2d_fp16_d64_b32_gqa8_combo_top_left": { "arch": "gfx950", - "bytes": 115874, + "bytes": 115273, "family": "attention", - "kernel_name": "rocke_uattn2d_tiled_d64_b32_t64_h64kv8_fp16_tl_w4_mw32_mfma32_stqk_s1_mask1_hlpv_skipqreg_mlim", - "sha256": "1a1f6efc9c29e49bfbe80414965bb2e1a937dee5e98f4646721fabca781a5171", + "kernel_name": "rocke_uattn2d_tiled_d64_b32_t64_g8_fp16_tl_w4_mw32_mfma32_stqk_s1_mask1_hlpv_skipqreg_mlim", + "sha256": "e5f80bfe55013aeb78c30b572571adcf3b7db48e811674e495d4b166038f0ab7", "status": "ok" }, "attention/gfx950/2d_fp16_d64_b32_gqa8_sinks_combo": { "arch": "gfx950", - "bytes": 116553, + "bytes": 115994, "family": "attention", - "kernel_name": "rocke_uattn2d_tiled_d64_b32_t64_h64kv8_fp16_sinks_w4_mw32_mfma32_stqk_s1_mask1_hlpv_skipqreg_mlim", - "sha256": "59487b61670b4a62dcea91ee659fcdabd595d06bdec14fd4db94c6c92d807eb2", + "kernel_name": "rocke_uattn2d_tiled_d64_b32_t64_g8_fp16_sinks_w4_mw32_mfma32_stqk_s1_mask1_hlpv_skipqreg_mlim", + "sha256": "c7d22d271ef6c66e1c02a40b540aee15cd560cddb08f05f3623b702f348b67dd", "status": "ok" }, "attention/gfx950/3d_bf16_d128_b64": { "arch": "gfx950", - "bytes": 161417, + "bytes": 163351, "family": "attention", - "kernel_name": "rocke_uattn3d_tiled_d128_b64_h4kv2_seg4_bf16", - "sha256": "68dcecb4f369b8bcab80ed8cc7d7c664fbd12ca2dd4c8a21b1f1d2533dd0f52b", + "kernel_name": "rocke_uattn3d_tiled_d128_b64_g2_seg4_bf16", + "sha256": "ecf4ac152e550ad5aaeeb6ac76bf69c7846fea33e71e18a1ffdf98e790fa9c01", "status": "ok" }, "attention/gfx950/3d_bf16_d128_b64_no_mask": { "arch": "gfx950", - "bytes": 159936, + "bytes": 161870, "family": "attention", - "kernel_name": "rocke_uattn3d_tiled_d128_b64_h4kv2_seg4_bf16_rbu", - "sha256": "29722169fae8ba2f9cf874a8aebd5eef91bd5323e23a6d63f8be393c9d16f8b3", + "kernel_name": "rocke_uattn3d_tiled_d128_b64_g2_seg4_bf16_rbu", + "sha256": "ae54fba61b58bfac4ba0b3785085990844912143a73aa4c2a47d1472c17d8126", "status": "ok" }, "attention/gfx950/3d_bf16_d128_b64_rb16": { "arch": "gfx950", - "bytes": 162096, + "bytes": 164030, "family": "attention", - "kernel_name": "rocke_uattn3d_tiled_d128_b64_h4kv2_seg4_bf16_rb16", - "sha256": "60c24d2b8e9a7c62c626e9b6e253dd618ee942d0ee612cf34e162a43f2200421", + "kernel_name": "rocke_uattn3d_tiled_d128_b64_g2_seg4_bf16_rb16", + "sha256": "6f44b77f8b70a7dce8bd1169f1c955080a8de140c34cb7e5d34659e57e36d1a0", "status": "ok" }, "attention/gfx950/3d_bf16_d128_b64_top_left": { "arch": "gfx950", - "bytes": 161412, + "bytes": 163346, "family": "attention", - "kernel_name": "rocke_uattn3d_tiled_d128_b64_h4kv2_seg4_bf16_tl", - "sha256": "55f1c89711ab6956d3c1730aa0f680af09f02f72f114546d37ae4cefec583911", + "kernel_name": "rocke_uattn3d_tiled_d128_b64_g2_seg4_bf16_tl", + "sha256": "a22800436a8303278d75a142e84c6ab26e60d436e24a0494f5dfc2d4f56249d4", "status": "ok" }, "attention/gfx950/3d_fp16_d128_b64": { "arch": "gfx950", - "bytes": 159456, + "bytes": 161390, "family": "attention", - "kernel_name": "rocke_uattn3d_tiled_d128_b64_h4kv2_seg4_fp16", - "sha256": "a9de23680e469af7b7c240b5ceba497e9ab4a1d096ac522bb3de5f03b8debd02", + "kernel_name": "rocke_uattn3d_tiled_d128_b64_g2_seg4_fp16", + "sha256": "a72ec1595bf9f2688997ed5dba8b2aef5ad42c4f9f5e7e5ccc8f9eaae50904f8", "status": "ok" }, "attention/gfx950/reduce_bf16_d128": { "arch": "gfx950", - "bytes": 9582, + "bytes": 9738, "family": "attention", - "kernel_name": "rocke_uattn_reduce_tiled_d128_h4_seg4_bf16", - "sha256": "ffd2175efd452e81fae6d9a0e673054b8a983b39655e2799cc62c45ea00f500c", + "kernel_name": "rocke_uattn_reduce_tiled_d128_seg4_bf16", + "sha256": "a3e43426dbbeb00409c2c52791998a12a05aa4664c42c26ef77649ce7d65e4ff", "status": "ok" }, "attention/gfx950/reduce_fp16_d128": { "arch": "gfx950", - "bytes": 9570, + "bytes": 9726, "family": "attention", - "kernel_name": "rocke_uattn_reduce_tiled_d128_h4_seg4_fp16", - "sha256": "beabc28143983d8b7435934ad4ac4c8f9b21ee351d15132153f37b0a15272a97", + "kernel_name": "rocke_uattn_reduce_tiled_d128_seg4_fp16", + "sha256": "2f121f51da067142e50d33d27ae2d7a82d58792f119ad4a3d9bd7248718d2db7", "status": "ok" }, "attention/gfx950/reduce_lse_bf16_d128": { "arch": "gfx950", - "bytes": 10462, + "bytes": 10633, "family": "attention", - "kernel_name": "rocke_uattn_reduce_tiled_d128_h4_seg4_bf16_lse", - "sha256": "aeb2bb788bd933c14b9ef4b6098ce2b14da748e9e501cbe3a87454b1f8a2cc35", + "kernel_name": "rocke_uattn_reduce_tiled_d128_seg4_bf16_lse", + "sha256": "88d4bda6d396c24fd4030bdb63be810f47d5cf973c6f11f2d830ea1208de6c0b", "status": "ok" }, "attention/gfx950/scalar2d_fp16_d128_no_mask": { @@ -332,10 +332,10 @@ }, "attention_d256/gfx950/pad_interleave": { "arch": "gfx950", - "bytes": 597130, + "bytes": 533839, "family": "attention_d256", - "kernel_name": "rocke_uattn2d_tiled_d256_b16_t64_h16kv2_bf16_w2_mw32_mfma32_stqk_s1_mask1_hlpv_mlim_ksb_mphsplit_smxil2g4_qdreg", - "sha256": "cd4307bce99bb62a7f06701560e1b8d398295b7597a4b6701e27380694469e84", + "kernel_name": "rocke_uattn2d_tiled_d256_b16_t64_g8_bf16_w2_mw32_mfma32_stqk_s1_mask1_hlpv_mlim_ksb_mphsplit_smxil2g4_qdreg", + "sha256": "7a8d8eb868ad222426a35257023fe53b8b2625ad6d9e45e5d4d6319400d3e7f5", "status": "ok" }, "attention_dense/gfx942/causal_sq512": { @@ -2465,186 +2465,186 @@ }, "attention/gfx950/2d_bf16_d128_b64": { "arch": "gfx950", - "bytes": 167562, + "bytes": 140050, "family": "attention", - "kernel_name": "rocke_uattn2d_tiled_d128_b64_h4kv2_bf16", - "sha256": "27fa0afc2d4df5e0549e55804c28fe525769dff8cbe4e1a59129df1327c3fc8f", + "kernel_name": "rocke_uattn2d_tiled_d128_b64_g2_bf16", + "sha256": "f352909660a1002ee8c0c74650f12e6dc902541df4873d80183527cae96b9edc", "status": "ok" }, "attention/gfx950/2d_bf16_d64_b32_gqa8_sinks_combo": { "arch": "gfx950", - "bytes": 117091, + "bytes": 117886, "family": "attention", - "kernel_name": "rocke_uattn2d_tiled_d64_b32_t64_h64kv8_bf16_sinks_w4_mw32_mfma32_stqk_s1_mask1_hlpv_skipqreg_mlim_fastkvdesc", - "sha256": "1231b94b52ee265b1f17afcd6967955c00b483fc6308dac27c0dbe25ec4d6e23", + "kernel_name": "rocke_uattn2d_tiled_d64_b32_t64_g8_h64kv8_bf16_sinks_w4_mw32_mfma32_stqk_s1_mask1_hlpv_skipqreg_mlim_fastkvdesc", + "sha256": "86fa538e17e4ff0549251cfe3c1781d0386cfbf9b1a3f1bf6fc7e01edc9aa6dc", "status": "ok" }, "attention/gfx950/2d_fp16_d128_b64": { "arch": "gfx950", - "bytes": 165985, + "bytes": 138473, "family": "attention", - "kernel_name": "rocke_uattn2d_tiled_d128_b64_h4kv2_fp16", - "sha256": "b0f98884afccd99351b73d031557f5c699cb8401fb4e980b32089818dcfc53e0", + "kernel_name": "rocke_uattn2d_tiled_d128_b64_g2_fp16", + "sha256": "7486da1b589d2715398373e68a993c4edb95351fa27fcebbfafbddd7091bec37", "status": "ok" }, "attention/gfx950/2d_fp16_d128_b64_alibi": { "arch": "gfx950", - "bytes": 171160, + "bytes": 143703, "family": "attention", - "kernel_name": "rocke_uattn2d_tiled_d128_b64_h4kv2_fp16_alibi", - "sha256": "33a3260ff1479a0453fd6fd06ac15531096330da39fc926b47d22b4c51c12d4d", + "kernel_name": "rocke_uattn2d_tiled_d128_b64_g2_fp16_alibi", + "sha256": "0e266bda7c779840e6b21cf0ddd13211a3dda55a4de306829cdb8582531cb1fa", "status": "ok" }, "attention/gfx950/2d_fp16_d128_b64_left_only_sw64": { "arch": "gfx950", - "bytes": 167454, + "bytes": 140030, "family": "attention", - "kernel_name": "rocke_uattn2d_tiled_d128_b64_h4kv2_fp16_sw64_tl_rbu", - "sha256": "d19efce4834f0aaff0701c6bf25507b2ad4eb379fefc0bbdb2e9749eeb48a8c0", + "kernel_name": "rocke_uattn2d_tiled_d128_b64_g2_fp16_sw_tl_rbu", + "sha256": "28df84cf88041c2ab9bfd8d1b44f4e900056da3e17691bb56bb3b7be26b5e744", "status": "ok" }, "attention/gfx950/2d_fp16_d128_b64_no_mask": { "arch": "gfx950", - "bytes": 164617, + "bytes": 137101, "family": "attention", - "kernel_name": "rocke_uattn2d_tiled_d128_b64_h4kv2_fp16_rbu", - "sha256": "d95f12a14d004edb17c85cddca4c04791dcbed265e8eeaeba7ba749f4c038325", + "kernel_name": "rocke_uattn2d_tiled_d128_b64_g2_fp16_rbu", + "sha256": "c771a4a8d84a7b770d1a7e2598200bb85b1df83e0e1cca2e10d866c20bc2b9dc", "status": "ok" }, "attention/gfx950/2d_fp16_d128_b64_rb16": { "arch": "gfx950", - "bytes": 166764, + "bytes": 139283, "family": "attention", - "kernel_name": "rocke_uattn2d_tiled_d128_b64_h4kv2_fp16_rb16", - "sha256": "74bbd4536ebdd8dd28246a9275a54b96eef54d8a87f2bed76cc26da8f71b275c", + "kernel_name": "rocke_uattn2d_tiled_d128_b64_g2_fp16_rb16", + "sha256": "4a4301f8b791e8f3bc6f41aa0ac09852115af604ebf8cf6fe586a76a5b35cd04", "status": "ok" }, "attention/gfx950/2d_fp16_d128_b64_softcap": { "arch": "gfx950", - "bytes": 174017, + "bytes": 146469, "family": "attention", - "kernel_name": "rocke_uattn2d_tiled_d128_b64_h4kv2_fp16_softcap", - "sha256": "820d4d097c0edbba5b494aa134eabbd04e79809f1beba6ee3f98bc719ae1ec08", + "kernel_name": "rocke_uattn2d_tiled_d128_b64_g2_fp16_softcap", + "sha256": "8d39d5e1ed34cf1d86c1e30d5c0bbf314c6179cabdd711c9cd799a5c725b86e1", "status": "ok" }, "attention/gfx950/2d_fp16_d128_b64_top_left": { "arch": "gfx950", - "bytes": 166025, + "bytes": 138513, "family": "attention", - "kernel_name": "rocke_uattn2d_tiled_d128_b64_h4kv2_fp16_tl", - "sha256": "e76d393f5aab2deca97d1816445d4a40480dba4761dc05e288ee171da0a4d22d", + "kernel_name": "rocke_uattn2d_tiled_d128_b64_g2_fp16_tl", + "sha256": "e1369a130b356735177e7f708b61f583ba7c91793d5405bc1b8db5e72c06e728", "status": "ok" }, "attention/gfx950/2d_fp16_d128_b64_top_left_sw64": { "arch": "gfx950", - "bytes": 168815, + "bytes": 141377, "family": "attention", - "kernel_name": "rocke_uattn2d_tiled_d128_b64_h4kv2_fp16_sw64_tl", - "sha256": "ae166ef5b28bc668cba7c0dcdb84a90a70f140c385b7b50fea30dd7c607ddb5a", + "kernel_name": "rocke_uattn2d_tiled_d128_b64_g2_fp16_sw_tl", + "sha256": "487b7d80e1f993561fccaa5c85950b96dd6943538d3aa69654f3e43e3ec2b82a", "status": "ok" }, "attention/gfx950/2d_fp16_d128_b64_two_sided_sw64_rb16": { "arch": "gfx950", - "bytes": 169632, + "bytes": 142216, "family": "attention", - "kernel_name": "rocke_uattn2d_tiled_d128_b64_h4kv2_fp16_sw64_tl_rb16", - "sha256": "8c66a2894e7cf463ac5e5f568089b0970418c57dccad4e91a213437a3acf87de", + "kernel_name": "rocke_uattn2d_tiled_d128_b64_g2_fp16_sw_tl_rb16", + "sha256": "f02165ada6fd40c116e314ad81281b02356c8d418969ec30a6efb7bc45c049ad", "status": "ok" }, "attention/gfx950/2d_fp16_d64_b32_gqa8_combo_no_mask": { "arch": "gfx950", - "bytes": 115842, + "bytes": 115238, "family": "attention", - "kernel_name": "rocke_uattn2d_tiled_d64_b32_t64_h64kv8_fp16_rbu_w4_mw32_mfma32_stqk_s1_mask1_hlpv_skipqreg_mlim", - "sha256": "b3d3f4e0ca7185d8e5fcba0b5ddf3d4535ac639a3dda1a217ad8b903efba0987", + "kernel_name": "rocke_uattn2d_tiled_d64_b32_t64_g8_fp16_rbu_w4_mw32_mfma32_stqk_s1_mask1_hlpv_skipqreg_mlim", + "sha256": "ca0ac3d11d058dba04fa4606827854aa3f357791a4058b1cce8ee00b22aa7515", "status": "ok" }, "attention/gfx950/2d_fp16_d64_b32_gqa8_combo_rb16": { "arch": "gfx950", - "bytes": 116178, + "bytes": 115591, "family": "attention", - "kernel_name": "rocke_uattn2d_tiled_d64_b32_t64_h64kv8_fp16_rb16_w4_mw32_mfma32_stqk_s1_mask1_hlpv_skipqreg_mlim", - "sha256": "f720d49b695181cfe6c356df3a033c161c187cb736f396b9366e023b1d3c144a", + "kernel_name": "rocke_uattn2d_tiled_d64_b32_t64_g8_fp16_rb16_w4_mw32_mfma32_stqk_s1_mask1_hlpv_skipqreg_mlim", + "sha256": "a8321999e6772ce51750ab7e56942f23d88d307a3d7f8b5b821f47e80c6366f2", "status": "ok" }, "attention/gfx950/2d_fp16_d64_b32_gqa8_combo_top_left": { "arch": "gfx950", - "bytes": 115970, + "bytes": 115369, "family": "attention", - "kernel_name": "rocke_uattn2d_tiled_d64_b32_t64_h64kv8_fp16_tl_w4_mw32_mfma32_stqk_s1_mask1_hlpv_skipqreg_mlim", - "sha256": "8dfc866984e944487ec1bb7feae2f87e76e5d643741b25a89ceb9ad4397bcfb1", + "kernel_name": "rocke_uattn2d_tiled_d64_b32_t64_g8_fp16_tl_w4_mw32_mfma32_stqk_s1_mask1_hlpv_skipqreg_mlim", + "sha256": "131cd77438c920a26d79f613f3e9e0006a7a21ffbce6d054ce4aedb749f3ba0e", "status": "ok" }, "attention/gfx950/2d_fp16_d64_b32_gqa8_sinks_combo": { "arch": "gfx950", - "bytes": 116649, + "bytes": 116090, "family": "attention", - "kernel_name": "rocke_uattn2d_tiled_d64_b32_t64_h64kv8_fp16_sinks_w4_mw32_mfma32_stqk_s1_mask1_hlpv_skipqreg_mlim", - "sha256": "be9dd416f73bd072385cde4ca0be82c7f5c34405c5a1d51910f35b396f2a7ab6", + "kernel_name": "rocke_uattn2d_tiled_d64_b32_t64_g8_fp16_sinks_w4_mw32_mfma32_stqk_s1_mask1_hlpv_skipqreg_mlim", + "sha256": "ffb1faa63471ca481a39cf9ba8d89e4e18e6777ac8425c125baa150356a38107", "status": "ok" }, "attention/gfx950/3d_bf16_d128_b64": { "arch": "gfx950", - "bytes": 161513, + "bytes": 163447, "family": "attention", - "kernel_name": "rocke_uattn3d_tiled_d128_b64_h4kv2_seg4_bf16", - "sha256": "fc2317de951975a2b4e8cb60cc366ba9c5a247c8ca095615e936d14a1acd4f56", + "kernel_name": "rocke_uattn3d_tiled_d128_b64_g2_seg4_bf16", + "sha256": "050a52b77e77931513af83a3ab5c69f088a3d7483a21cfcdba82a1a9bbb5fe8f", "status": "ok" }, "attention/gfx950/3d_bf16_d128_b64_no_mask": { "arch": "gfx950", - "bytes": 160032, + "bytes": 161966, "family": "attention", - "kernel_name": "rocke_uattn3d_tiled_d128_b64_h4kv2_seg4_bf16_rbu", - "sha256": "0995eea05db416b61c1e1cb41687c37c263710a0f67c3976b836bf17d9af66a3", + "kernel_name": "rocke_uattn3d_tiled_d128_b64_g2_seg4_bf16_rbu", + "sha256": "c4d124595b1970e86378eb29358bfd18f1813bad78226e81e25b52c2dbdda5b7", "status": "ok" }, "attention/gfx950/3d_bf16_d128_b64_rb16": { "arch": "gfx950", - "bytes": 162192, + "bytes": 164126, "family": "attention", - "kernel_name": "rocke_uattn3d_tiled_d128_b64_h4kv2_seg4_bf16_rb16", - "sha256": "2fc3b55e07054ec1ba3fedfd0daa9e85ea56d4df60424550814481b8f9319f7d", + "kernel_name": "rocke_uattn3d_tiled_d128_b64_g2_seg4_bf16_rb16", + "sha256": "bc7f23ff8c8c30262bd7e6725c2bbc3e021bbf6413eb5ae4378e861864364cc6", "status": "ok" }, "attention/gfx950/3d_bf16_d128_b64_top_left": { "arch": "gfx950", - "bytes": 161508, + "bytes": 163442, "family": "attention", - "kernel_name": "rocke_uattn3d_tiled_d128_b64_h4kv2_seg4_bf16_tl", - "sha256": "caf565da35f0a6574e0d17853f07ab7b8474232254fe61160fff6a5b89fde405", + "kernel_name": "rocke_uattn3d_tiled_d128_b64_g2_seg4_bf16_tl", + "sha256": "0032025571b4726db2a16f0baedb4b4f8676a11fcad8902d7bd7ee3874022b9a", "status": "ok" }, "attention/gfx950/3d_fp16_d128_b64": { "arch": "gfx950", - "bytes": 159552, + "bytes": 161486, "family": "attention", - "kernel_name": "rocke_uattn3d_tiled_d128_b64_h4kv2_seg4_fp16", - "sha256": "a886f8bfde1bc76812a5a27a2ad13a7b9353efdf2950b182025aa7dd056b133f", + "kernel_name": "rocke_uattn3d_tiled_d128_b64_g2_seg4_fp16", + "sha256": "fb3929f3c7de873b1860b4d64ad068dfea612b2f6bcc001faf5baa30f2d44fb7", "status": "ok" }, "attention/gfx950/reduce_bf16_d128": { "arch": "gfx950", - "bytes": 9589, + "bytes": 9745, "family": "attention", - "kernel_name": "rocke_uattn_reduce_tiled_d128_h4_seg4_bf16", - "sha256": "98e0ba897d0a0ed2d7c2d809bdfb39f6c49c5f1c33b854e73ec5583124f6283d", + "kernel_name": "rocke_uattn_reduce_tiled_d128_seg4_bf16", + "sha256": "098580ab1c7803d72d1b2bd362d07aef3376e502038da6fea1f9802852ae109b", "status": "ok" }, "attention/gfx950/reduce_fp16_d128": { "arch": "gfx950", - "bytes": 9577, + "bytes": 9733, "family": "attention", - "kernel_name": "rocke_uattn_reduce_tiled_d128_h4_seg4_fp16", - "sha256": "e02b0ee7cfad94216ed2f945b38f047c5ed4e7d1e53d9a10a5ad090ccc41f3be", + "kernel_name": "rocke_uattn_reduce_tiled_d128_seg4_fp16", + "sha256": "d5203691a7de9c415f0b9a13d7449e59001d688dc12ee0e3fae1398990911638", "status": "ok" }, "attention/gfx950/reduce_lse_bf16_d128": { "arch": "gfx950", - "bytes": 10469, + "bytes": 10640, "family": "attention", - "kernel_name": "rocke_uattn_reduce_tiled_d128_h4_seg4_bf16_lse", - "sha256": "ac383a27abbc3ce0195b48e843e12175f598ad1b1a4a5fe1f51641442c202750", + "kernel_name": "rocke_uattn_reduce_tiled_d128_seg4_bf16_lse", + "sha256": "0da89e4bd78242188bf43d5b2b1978e8d7692ae80bb5fd397e52ec7f55ea7cff", "status": "ok" }, "attention/gfx950/scalar2d_fp16_d128_no_mask": { @@ -2721,10 +2721,10 @@ }, "attention_d256/gfx950/pad_interleave": { "arch": "gfx950", - "bytes": 597230, + "bytes": 533939, "family": "attention_d256", - "kernel_name": "rocke_uattn2d_tiled_d256_b16_t64_h16kv2_bf16_w2_mw32_mfma32_stqk_s1_mask1_hlpv_mlim_ksb_mphsplit_smxil2g4_qdreg", - "sha256": "cf3163aba7b796909e402cf81ca203af1f865ba4343aec58d9bdeb228df731ab", + "kernel_name": "rocke_uattn2d_tiled_d256_b16_t64_g8_bf16_w2_mw32_mfma32_stqk_s1_mask1_hlpv_mlim_ksb_mphsplit_smxil2g4_qdreg", + "sha256": "18573d232e5d665e89532cf7fe5c84b0af59fe7373144b93e7d8d5f46989c784", "status": "ok" }, "attention_dense/gfx942/causal_sq512": { @@ -4854,186 +4854,186 @@ }, "attention/gfx950/2d_bf16_d128_b64": { "arch": "gfx950", - "bytes": 167626, + "bytes": 140114, "family": "attention", - "kernel_name": "rocke_uattn2d_tiled_d128_b64_h4kv2_bf16", - "sha256": "23f1c0b559bbadb15d9179ecda2b1efcb58fc1bc742b51d3f1d09f2811a759af", + "kernel_name": "rocke_uattn2d_tiled_d128_b64_g2_bf16", + "sha256": "268943d5f0007844d17c946fee623eca5c60be9ec428091131624efa646ea537", "status": "ok" }, "attention/gfx950/2d_bf16_d64_b32_gqa8_sinks_combo": { "arch": "gfx950", - "bytes": 117155, + "bytes": 117950, "family": "attention", - "kernel_name": "rocke_uattn2d_tiled_d64_b32_t64_h64kv8_bf16_sinks_w4_mw32_mfma32_stqk_s1_mask1_hlpv_skipqreg_mlim_fastkvdesc", - "sha256": "eb41c11c96b2c6e28734961cc896b0c003df7e05722786ce9db0540edf4bd997", + "kernel_name": "rocke_uattn2d_tiled_d64_b32_t64_g8_h64kv8_bf16_sinks_w4_mw32_mfma32_stqk_s1_mask1_hlpv_skipqreg_mlim_fastkvdesc", + "sha256": "8fd094b334a261b8eeb0e8ca4a7c35b354b859279cb85164569cf34e3694e8ab", "status": "ok" }, "attention/gfx950/2d_fp16_d128_b64": { "arch": "gfx950", - "bytes": 166049, + "bytes": 138537, "family": "attention", - "kernel_name": "rocke_uattn2d_tiled_d128_b64_h4kv2_fp16", - "sha256": "b0de24b68cd58ec2d8cdd8aa616c15feeb20861ac77fb6d0d8eba7bde98abf66", + "kernel_name": "rocke_uattn2d_tiled_d128_b64_g2_fp16", + "sha256": "d4f2f43e054bfe905f479129541916d8686e63cdcd89fe380153990fd59e5ada", "status": "ok" }, "attention/gfx950/2d_fp16_d128_b64_alibi": { "arch": "gfx950", - "bytes": 171224, + "bytes": 143767, "family": "attention", - "kernel_name": "rocke_uattn2d_tiled_d128_b64_h4kv2_fp16_alibi", - "sha256": "b52b0a4b9e8302e2fa7e3ed50271705a951b54fec94ab0ed84e344f04babd9e5", + "kernel_name": "rocke_uattn2d_tiled_d128_b64_g2_fp16_alibi", + "sha256": "567b3e2bf54c921100a0bc3823af70eecb739fefa93afd74e33e61bc53a04c51", "status": "ok" }, "attention/gfx950/2d_fp16_d128_b64_left_only_sw64": { "arch": "gfx950", - "bytes": 167518, + "bytes": 140094, "family": "attention", - "kernel_name": "rocke_uattn2d_tiled_d128_b64_h4kv2_fp16_sw64_tl_rbu", - "sha256": "186cb956204eb9da83fc33f6b32112e0370605fcae1d41714b3e6088431f4de6", + "kernel_name": "rocke_uattn2d_tiled_d128_b64_g2_fp16_sw_tl_rbu", + "sha256": "6cde8de59ff728d387701f5ad1964b082dd895c7c106809b9770c691008899d5", "status": "ok" }, "attention/gfx950/2d_fp16_d128_b64_no_mask": { "arch": "gfx950", - "bytes": 164681, + "bytes": 137165, "family": "attention", - "kernel_name": "rocke_uattn2d_tiled_d128_b64_h4kv2_fp16_rbu", - "sha256": "8c308c7a4917bc28e5be4fe385d285580d83a30d17bc49e86b0e4790bb463728", + "kernel_name": "rocke_uattn2d_tiled_d128_b64_g2_fp16_rbu", + "sha256": "ab2c5a05db4a35db4eae69ced70ef7fe47ce043bc9c5f222d2544de85aedb434", "status": "ok" }, "attention/gfx950/2d_fp16_d128_b64_rb16": { "arch": "gfx950", - "bytes": 166828, + "bytes": 139347, "family": "attention", - "kernel_name": "rocke_uattn2d_tiled_d128_b64_h4kv2_fp16_rb16", - "sha256": "d85e184618cd02e4a37d142369149cabb1fba773792cdaab0094747f3d5430e9", + "kernel_name": "rocke_uattn2d_tiled_d128_b64_g2_fp16_rb16", + "sha256": "c7c8f460f6940be8bb1858764893b065315054171c2e0e5cc2a3a9deeb18eb56", "status": "ok" }, "attention/gfx950/2d_fp16_d128_b64_softcap": { "arch": "gfx950", - "bytes": 174081, + "bytes": 146533, "family": "attention", - "kernel_name": "rocke_uattn2d_tiled_d128_b64_h4kv2_fp16_softcap", - "sha256": "61db0ee642edee415fb63f3523104a3e681489f40f5e5873a2edcad69894f917", + "kernel_name": "rocke_uattn2d_tiled_d128_b64_g2_fp16_softcap", + "sha256": "00997452c34eefbf48b8be50b17866466588d49bdabc54caa0c7693235d5cd50", "status": "ok" }, "attention/gfx950/2d_fp16_d128_b64_top_left": { "arch": "gfx950", - "bytes": 166089, + "bytes": 138577, "family": "attention", - "kernel_name": "rocke_uattn2d_tiled_d128_b64_h4kv2_fp16_tl", - "sha256": "9280e49718784dd6670bc29a85adb6e72548b6c6c5390d21bee26f58c1ce3577", + "kernel_name": "rocke_uattn2d_tiled_d128_b64_g2_fp16_tl", + "sha256": "95654a3a3d334deb799c0326461ed975af8bfe951f67a2148330a6a04c31a1ae", "status": "ok" }, "attention/gfx950/2d_fp16_d128_b64_top_left_sw64": { "arch": "gfx950", - "bytes": 168879, + "bytes": 141441, "family": "attention", - "kernel_name": "rocke_uattn2d_tiled_d128_b64_h4kv2_fp16_sw64_tl", - "sha256": "d9630540ab733a9d2e053387250b3b6529fb36c7837d0d74127ff2d924afbd1a", + "kernel_name": "rocke_uattn2d_tiled_d128_b64_g2_fp16_sw_tl", + "sha256": "ff870db94fc12e71db0489cd3836606b0a524980416ed77ac34eafa17800fe88", "status": "ok" }, "attention/gfx950/2d_fp16_d128_b64_two_sided_sw64_rb16": { "arch": "gfx950", - "bytes": 169696, + "bytes": 142280, "family": "attention", - "kernel_name": "rocke_uattn2d_tiled_d128_b64_h4kv2_fp16_sw64_tl_rb16", - "sha256": "b963e9550fe3929bc9597e4aa605a64d96183ed2643bfa53833a8df3cbefad50", + "kernel_name": "rocke_uattn2d_tiled_d128_b64_g2_fp16_sw_tl_rb16", + "sha256": "32a11090b0f5c871d13a95bc75977167cc4b8aded226a96662934fe09acecf2d", "status": "ok" }, "attention/gfx950/2d_fp16_d64_b32_gqa8_combo_no_mask": { "arch": "gfx950", - "bytes": 115906, + "bytes": 115302, "family": "attention", - "kernel_name": "rocke_uattn2d_tiled_d64_b32_t64_h64kv8_fp16_rbu_w4_mw32_mfma32_stqk_s1_mask1_hlpv_skipqreg_mlim", - "sha256": "c24e991ae22f1ad2e30ba506904ec47ec40565b3895c5b871fa2e9bb8a761b7e", + "kernel_name": "rocke_uattn2d_tiled_d64_b32_t64_g8_fp16_rbu_w4_mw32_mfma32_stqk_s1_mask1_hlpv_skipqreg_mlim", + "sha256": "a11f4b52e7e0405348e2428f5346ded517961a7766135a37f32da5abdeb1edaa", "status": "ok" }, "attention/gfx950/2d_fp16_d64_b32_gqa8_combo_rb16": { "arch": "gfx950", - "bytes": 116242, + "bytes": 115655, "family": "attention", - "kernel_name": "rocke_uattn2d_tiled_d64_b32_t64_h64kv8_fp16_rb16_w4_mw32_mfma32_stqk_s1_mask1_hlpv_skipqreg_mlim", - "sha256": "6ad687ecf488ca19a78703424cec5d3bee5e84911218d381d0ddb40aa67b349d", + "kernel_name": "rocke_uattn2d_tiled_d64_b32_t64_g8_fp16_rb16_w4_mw32_mfma32_stqk_s1_mask1_hlpv_skipqreg_mlim", + "sha256": "5affdb04bad2b7b65d4c8acb6de4aca0fa83b4e34edb98ca8e44507fd985be47", "status": "ok" }, "attention/gfx950/2d_fp16_d64_b32_gqa8_combo_top_left": { "arch": "gfx950", - "bytes": 116034, + "bytes": 115433, "family": "attention", - "kernel_name": "rocke_uattn2d_tiled_d64_b32_t64_h64kv8_fp16_tl_w4_mw32_mfma32_stqk_s1_mask1_hlpv_skipqreg_mlim", - "sha256": "84021273881d949e6cae97d37193713add60a70ba389d04fa5df2088463c15a6", + "kernel_name": "rocke_uattn2d_tiled_d64_b32_t64_g8_fp16_tl_w4_mw32_mfma32_stqk_s1_mask1_hlpv_skipqreg_mlim", + "sha256": "facfc044b68c2d867d207afc874c2171dea6b43dcc39722bdf548db24ae77697", "status": "ok" }, "attention/gfx950/2d_fp16_d64_b32_gqa8_sinks_combo": { "arch": "gfx950", - "bytes": 116713, + "bytes": 116154, "family": "attention", - "kernel_name": "rocke_uattn2d_tiled_d64_b32_t64_h64kv8_fp16_sinks_w4_mw32_mfma32_stqk_s1_mask1_hlpv_skipqreg_mlim", - "sha256": "b66dad2f8db97c367c43f8b628347b265f33a98be11ab353d19917277a87fcc6", + "kernel_name": "rocke_uattn2d_tiled_d64_b32_t64_g8_fp16_sinks_w4_mw32_mfma32_stqk_s1_mask1_hlpv_skipqreg_mlim", + "sha256": "e991ccdd7d485073fa24d51b8e0506e11e9d6514dafde686e0e1216e450be56c", "status": "ok" }, "attention/gfx950/3d_bf16_d128_b64": { "arch": "gfx950", - "bytes": 161577, + "bytes": 163511, "family": "attention", - "kernel_name": "rocke_uattn3d_tiled_d128_b64_h4kv2_seg4_bf16", - "sha256": "8bb93d90bb6b918c9e5ca741713b8f527330c08f0ec64fb8399c5607f95f93e4", + "kernel_name": "rocke_uattn3d_tiled_d128_b64_g2_seg4_bf16", + "sha256": "ede4f268f6987f25f6502f1567e583439e01b7406e41f3102022458b7ce874c5", "status": "ok" }, "attention/gfx950/3d_bf16_d128_b64_no_mask": { "arch": "gfx950", - "bytes": 160096, + "bytes": 162030, "family": "attention", - "kernel_name": "rocke_uattn3d_tiled_d128_b64_h4kv2_seg4_bf16_rbu", - "sha256": "94a385dc5f7a749e5b3286996fc9c1891c353bcfc4caac8ebf07bbd40475ba36", + "kernel_name": "rocke_uattn3d_tiled_d128_b64_g2_seg4_bf16_rbu", + "sha256": "291ebec1509a9aa66b0e4b08e1ef847eab247af77ba43d05f8979fc5b37d77c3", "status": "ok" }, "attention/gfx950/3d_bf16_d128_b64_rb16": { "arch": "gfx950", - "bytes": 162256, + "bytes": 164190, "family": "attention", - "kernel_name": "rocke_uattn3d_tiled_d128_b64_h4kv2_seg4_bf16_rb16", - "sha256": "9ba252e2a18e12bb50e49d9a53de0d44519fdc53445e893b966ed6510c3272de", + "kernel_name": "rocke_uattn3d_tiled_d128_b64_g2_seg4_bf16_rb16", + "sha256": "90038c0ae2fdf637d30f59b02bf1545c4e3414b1fa7ab0a1e2cee3a12ad5c33a", "status": "ok" }, "attention/gfx950/3d_bf16_d128_b64_top_left": { "arch": "gfx950", - "bytes": 161572, + "bytes": 163506, "family": "attention", - "kernel_name": "rocke_uattn3d_tiled_d128_b64_h4kv2_seg4_bf16_tl", - "sha256": "044a112eaa821c7cdb13e0e501b0019d12f13189eb6eedee6118696b9063ab86", + "kernel_name": "rocke_uattn3d_tiled_d128_b64_g2_seg4_bf16_tl", + "sha256": "489d7be5499ac03f654285e6532d867f87a64d85d3d7fb2f37c60b17bca591d0", "status": "ok" }, "attention/gfx950/3d_fp16_d128_b64": { "arch": "gfx950", - "bytes": 159616, + "bytes": 161550, "family": "attention", - "kernel_name": "rocke_uattn3d_tiled_d128_b64_h4kv2_seg4_fp16", - "sha256": "6ec6b7296f4019b340866cec1f71786d02f696256b3195ffae0d1387aa52f46c", + "kernel_name": "rocke_uattn3d_tiled_d128_b64_g2_seg4_fp16", + "sha256": "43f2721a7212e095b5805d2bdd5b875a09b2050be6cb68b5405feb6292e18d6b", "status": "ok" }, "attention/gfx950/reduce_bf16_d128": { "arch": "gfx950", - "bytes": 9653, + "bytes": 9809, "family": "attention", - "kernel_name": "rocke_uattn_reduce_tiled_d128_h4_seg4_bf16", - "sha256": "19d7d61c6f782d75843f9f8feb9b266a793b186da1d817b477612cc47e4a32ca", + "kernel_name": "rocke_uattn_reduce_tiled_d128_seg4_bf16", + "sha256": "622b5c24eb92d4b7fcf99a4506b63d6793cf92bc683efe3302418559aa8c3cf7", "status": "ok" }, "attention/gfx950/reduce_fp16_d128": { "arch": "gfx950", - "bytes": 9641, + "bytes": 9797, "family": "attention", - "kernel_name": "rocke_uattn_reduce_tiled_d128_h4_seg4_fp16", - "sha256": "86a97d40e5b330935ede77f65469c2b6be4842406951f2d75192d713a07a0010", + "kernel_name": "rocke_uattn_reduce_tiled_d128_seg4_fp16", + "sha256": "86bfc8e5818512ff10660abe064f7614d71cab37f33447828062c802f0cf36b7", "status": "ok" }, "attention/gfx950/reduce_lse_bf16_d128": { "arch": "gfx950", - "bytes": 10533, + "bytes": 10704, "family": "attention", - "kernel_name": "rocke_uattn_reduce_tiled_d128_h4_seg4_bf16_lse", - "sha256": "7518fc3a768b5f4598c9827648481455b523eb21bcf1e1f8b8c6134c30b195ae", + "kernel_name": "rocke_uattn_reduce_tiled_d128_seg4_bf16_lse", + "sha256": "f687a90d6d466326ae5d6fd95c4037ce647e22601377181c203e9b836ad408e7", "status": "ok" }, "attention/gfx950/scalar2d_fp16_d128_no_mask": { @@ -5110,10 +5110,10 @@ }, "attention_d256/gfx950/pad_interleave": { "arch": "gfx950", - "bytes": 597294, + "bytes": 534003, "family": "attention_d256", - "kernel_name": "rocke_uattn2d_tiled_d256_b16_t64_h16kv2_bf16_w2_mw32_mfma32_stqk_s1_mask1_hlpv_mlim_ksb_mphsplit_smxil2g4_qdreg", - "sha256": "e20ccb61ab8ca4ba8413c69ba4dec0f50f9ec7025db28b15e608a7db0672729f", + "kernel_name": "rocke_uattn2d_tiled_d256_b16_t64_g8_bf16_w2_mw32_mfma32_stqk_s1_mask1_hlpv_mlim_ksb_mphsplit_smxil2g4_qdreg", + "sha256": "8b4eba316195303225c042ce46ef15ee2fbc0eb25762a627a1f145b0b8ea6888", "status": "ok" }, "attention_dense/gfx942/causal_sq512": { 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 871c330058c2..6029d5259eab 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 @@ -19,18 +19,18 @@ "attention_strided_kv": { "0": "c99fe0ffaf340903e7c913700011f4e4e6aea5630ddeca60c24a66916228d9cb", "1": "7a4040765b46c08f298f74dd3952e8a830aa02d4e397d6d3d37fa38f7cc76e02", - "10": "8ec3555ed2e64c59f0a32847175a1dfe0416d83036b645c29da331e32f640f57", - "11": "4bd0dd8a84612f0e25fe82e1188eea5a23e7951179a5ed266d969a391e994382", + "10": "e15390610f1eda44ef679b865738e82fd9ecd1646ff82f7da5dcd8f158e6cfff", + "11": "88f4619bb13a46d71165b57e6af170f437760aa740d07fa8730e6d097f034d0d", "12": "8b9fdf9d81f108bb733518efe3f16ff3d3a512b5a735c2038666b5bc98c0a51f", - "13": "6a0c63a1574445581eb6767c40600b6bffedde472c9704bbc5511855db030775", + "13": "e60f0ad39c4f7a93deb98820eb0ff822d9bd687c7c7b6489e22c1a956f863700", "2": "5b2da99305995476f23c4aa3e274c249fb3dbab830d984fc641a57644d0de8b0", - "3": "1caedf0196d6b6cb889496cefa61d844373d4db06dc9b98327f6aa44de2f8b48", - "4": "d93e5587c7f3b2432d8248a8639f19ab75416920cb6f2139198bbcddc0603ba8", - "5": "f7576a337c2599cb03991ef9a8e82541a50208c9ef1704b4b3cb63d3933ac7c6", + "3": "2e17e193a646496fb2307f2bb4fe061c18310622743072ead233902b2f2591af", + "4": "9a78878c3b5e7b90231fd8ad907fc629e164c1bfa7265a8af16019ffaa957297", + "5": "b4b7ee5c63360b9231f6f6301326ad35b3877263050c43631a60704339bc77ca", "6": "382bd4a5175627bbd9090a22346b587ddc3d86cc2b461f6d98b20ea7b9953260", "7": "136adbde306c73fbbde41916fb93ac4b596cb116b2730727ea0802826580941d", "8": "e5f2a61ca1e3b775cbac546a6cc531c22a5c978edeae0b0c082b9b05d083b87c", - "9": "b932c7ec5c473d26f9c5d3364d503eba33b0d916d517114842dff4352274241e" + "9": "ec8c7102236fbae47258f43dcbda0cbc0f77c2c9fe516184ae302ab27a6e59e5" }, "attention_unified": { "0": "2dd9c60d1dba62d1dba0402468528886c520441e3447900b690b17e0f6a84a3c", @@ -1291,88 +1291,94 @@ "9": "d1dc34cb8dfca4f7fadafcd65d226b540ab47e634f24931abd0a802f02042dd5" }, "gfx950_attention_tiled_2d": { - "0": "ca49dc00b5be88b8de75c90a1cc134b377df4c4e0af3d2bb2a7c33639cce51ae", - "1": "db4f40bb92bfee28ec1dc70ed0453a4b09efe417571402b87117658f2ce9df7d", - "10": "ac27f33fdce4e8b065eb54c82fe6a2bc8a3fcb4edba590032a1ecf37c9c0ac7d", - "11": "638a8ceb5ebc0bb7575dc721b81dda689d1c7a9b5ab9973308a75d49632f20ed", - "12": "b4ed5b3c793ac4d5b4fb20e09743ac4f55ac8f5c1ba0c642a909e564d52ffcfd", - "13": "4068ce911395ea155a2d4d5e62b5c66441c0a21dd45dba2f63c5be88a61d25e8", - "14": "4a3556ec2b6b9631032fb0ebef3647cb74e02dffed88e9ef4b15dc0dfac53bf1", - "15": "81b84031b5ad1e13268077c440110963f895cdc154881974c72a99b5cbe1b61f", - "16": "0dd78c1f28099feb7aa00875c5dee6d551df60a9e5f83d2a92cb5967b3fc1ed5", - "17": "4930cb63381a923cf92afac2210dc19e7c68b315faa0e4bd264e6de179bad9fe", - "18": "1f59bc34daa84e4200c563a35223d871a0bae241bb7c5a9452695a7524fd6196", - "19": "6f0040da04e2aa1dd18e15e285d7f8a5b8eeaed5c5f3041ac0069984d9a8b1e5", - "2": "717a1f45d8152a793c8652a3a56fb6a923626be8cb890ffe0462823fa13c755b", - "20": "b49b4d828e01473bb3bf6a5b68f35d0a5c388f3b0f0a271df019f982c4eaf987", - "21": "c108fb11bd5b4456cece927e07f49bc69c5686277b439f0c4370ab139b12df1b", - "22": "f89c7eeb4d0b424ad8edc79330891b78642a5fff9dd26d779a8cc905aaf96258", - "23": "6633e44e6aff0780d13ce09522c78db297477bb7e27394c19b11ea9cf147662b", - "24": "1fd5bfcd6d6ff5b74293424998e6097f5506249e9127e8158e6cbf3e872f707a", - "25": "9ca066433e8d2bd3e12ca72a38a50ac6680f2f0f038b87f96654cc7c283d41ed", - "26": "7dd920b867ac40d2b9e141d4b4659ccc69ba2202eef8bafc0c6a0f44383915ad", - "27": "b24018e49cc0e1a586ba328064c394127ce85aadda29e5727baa94db26feff29", - "28": "6ac82a8165b3d6a941c157f4817e48cb70247ea54b4605b27ea7e297423fb1d6", - "29": "3747f4d7b0e61b886d05b3205c869aee188993634217561f903f223d71c06c41", - "3": "828e9f041e7f23de6b73243df119c738945deac20ba8865dd5a9a98e985d4b05", - "30": "cb461dd028ffdb9cfe5c2480d21f2b138adf6a6e5b7c49c49e0c37528fb2bead", - "31": "d90c34dfdb3426c06999a45a66e1cb4d791ac3b80a6216066752b72fb4de17e8", - "32": "0d90c05cd049dc492a6143af3dcd4cab1006bca2f1d9c091c4f33fbdd5ab63af", - "33": "61a075e7a8518161afd92322b24e1bf6520fc1782a3743bd69a31e966f93fd73", - "34": "aa6d8786f71d6a72ca1a0d3e5a9467b17bbd187ab741ba07210053c2acbe8b5f", - "35": "f86b5de9a0d800de5c37c8309bae053ec8816ffcebdfdb55726aed4930362f39", - "36": "cd94fbfd9e8ec84ffffd2e86a6d9c175bad508911715afd1597cebf216e77940", - "37": "3e0273ca319c8f594d2809fff80ae4764d26c84762d450d9d3d5a06ee2daf3aa", - "38": "0fa70ef9fdd6154224d1ac035f0f759187d994e11c7ec82df8b6a99ef36a421a", - "39": "1c1fd39333532d3bd4e45ce96ff9b191b6c7cdb753e7c470a42535134b1a280c", - "4": "2ea80c30737b05259c15cd6a1a485fde67231ab36ff8e8e4b4e189a9dd617770", - "40": "80a03006d8056a77ebe19763ac3b7c1904517481920484864d7dc1d951268736", - "41": "5cd307b7295de60593430aa4fa1ee97187de1fe66c4400f66de069b100d17931", - "42": "2aa9b65105b53ad6abcf1ac735b0ac2f1153cbd79dc8bb08cf0abd3bee6b3d01", - "43": "83c1f65e9bd3adea3cd71d06158d3b4e8f14afd242349fa3cb33baf55db4e2b3", - "44": "81421e63d31f74612b54c1de06afc53df0a182d38f0c843df0b3763f65e71a52", - "45": "b47031495fee7dcf3087a9564d60fba772f867e7f98cfb324bee59263c4e8167", - "46": "3e8e5eabac9ce291e7ecf04456fbc7e62b55b33fd9ebe2e7e0f3b75928e21e65", - "47": "021f16ec8e4ecd8244cfa14516f7f7ecf2259bebe4cad7436a8ad434e96286b0", - "48": "8bf4e2cfac82bb4c5532e6c51342f6225b0f17941754eaa66381d39357865e89", - "49": "bbef9803cec91734b012585d1f1774fa09caf5bd87fe4869f5b59a276cb56c1c", - "5": "5d2b06c569b12221ea8888aa9af516fa33e7c8e21820f5171ab74c1240f17a04", - "50": "73adb0791bebe6e03aca1bc85e266ebf736e3a45a43978f7f20a28dcaabb260b", - "51": "eaedcc80f928aa2ec0ebadd2bbe888db960df4710b9e19821004f843ba3feb22", - "52": "717c6765996c6b9f2ddcafa25d583072236bc38966c1001bfa53f53a9538f9bb", - "53": "df4d128af016237edebf828d8e24dd5404d3473b28b509764a85cbfabb1e3bfe", - "54": "677abed3136298df3a2ecc8aaa13a9aed2447f17bb38f3a4b5c1dc08e2bbd094", - "55": "520e2482bf88fde065bcd238f0f2d7674c49111f726c07b1f7fbdf4b2ad8d54b", - "6": "555adce2cbd9bb832d640131769903e5d639db70bc07461eea5c8a59c2358718", - "7": "bbcd9758625b990d8d78cbc25e11756950e85d7db44c1b5050724f0efd721b1f", - "8": "cea4c8a1c2e10ab7d794c144103859fb72be23ab8cd875339df07d9fb6965aa4", - "9": "7180bac09e314b919c16242218d43935eaf463660488b5f7986fdefab664dca5" + "0": "93094984410adace374ae2120509bba148ed7fa3971bb5ec1f7893a6e71a86f5", + "1": "25da386ddc67492e7d6a0f2df62c8451f1422398b1d27e9a3689a1e9a90b350c", + "10": "f0fdf2a47235a06298ca13002c4f35f436664251af56d1eddc026661db230e5e", + "11": "77a8cb2c5b0d44018373403cec28ecc606aec07a9116993a39eb243bba5c4efe", + "12": "544aab4cff9017a05f0c3c84598f4aaf04a0a97472f364d07efa8a7a0adbbb7b", + "13": "2bec2ee7624c840a85375821928911af931a7f8f8026c04213e14ae05a49ac99", + "14": "695c7b899b46cef3cc241b45efa0940baa73745fad710cfdb4a31fcdc5f8e2a9", + "15": "d840bfac8b20548cdfd3bc60059bc5b4cc87c15019c0ac4f29ce8c178ffc29a2", + "16": "b695c0545119e926a529eacaadbc51af021391b57ab92e7cfb5e0b62cb9b373e", + "17": "835e96f0c3d2f3d9dac2a87fc76df715cafc6c2cd3e107f5f9440d60114c1024", + "18": "1a8347304eeb7d77ba3d02eb27cb87887cbebecf9a64e9edaed3ca4e98cdfb1b", + "19": "4464542d2ddbbd8530730e1c3e319b0313abae90ac0a85bba37003bf86d5c7c8", + "2": "04cdf86241fb2a0eeca86836ac9a6c901463b377800be8502c530cf0609b0ab6", + "20": "f66343671400fc1501968503b41dd09cdce34cc62c21a74eab779dc783e65aee", + "21": "8ab787255d932fcd1fb59ef1c2a95efed02c1e3ed2e59a171ca2c575be1dc896", + "22": "7fceb1a3fc2b4dd3c02823bbdf55ad86bf46bef4d0195571ce8d99f4fd3e20ec", + "23": "eab9a53a4e086ff638bd14e17d4dea119b281ebc478a4b85bc8d774bb2ce6a0f", + "24": "101a034ed4d1cb4c2e31137f90a8bd8ef1a920ade65e805f48f47b74546c68a3", + "25": "77a8cb2c5b0d44018373403cec28ecc606aec07a9116993a39eb243bba5c4efe", + "26": "77a8cb2c5b0d44018373403cec28ecc606aec07a9116993a39eb243bba5c4efe", + "27": "10bcf52cfab732148496ab043966a6c8118ed142c5fa76afc23d0342613383f3", + "28": "7994cf603bcdf4cd51b94bb5e084b8ee47aebf9a031ab13be2eead9c5441ba00", + "29": "df9bc830adbb2973972a0ba8ef874c629191420292d393d12cf37e904d48573d", + "3": "547ef63bacccf56f7062da4f099f33ed69f77f1051323c47458164654440243a", + "30": "c5050869dd16607d88acd47bc78706bf49e5eb3c44c3c0916109c51cc67190fa", + "31": "559229d57bb4a1f5dc4af47d6467106928f72d7c278c88ad49c8eb28bd54078c", + "32": "e640bf5024946d357b9631ede24c7443e0933a89da27079dda1be88ba13e7d8b", + "33": "1b14019ed4b2747e450181c29b01b9d1b05823dfb65f9425a01ef4d1b92d2a08", + "34": "db33b45dcbd971b9f5b9900d75f85b79e3be03efe92c497cd1d82df5373c7ed5", + "35": "ee24361cdb7979418b9c51f0c1c919e26b3ca63ba8facf21270322c59fb1296a", + "36": "9ab13a2f7b56804156ae0c3aa3abec50de5a1eb3cb6b4eedd34c746735aac0a8", + "37": "53b8f4482e62900960732c42cf4c2ad7e49d0dad891c0d576f7197c4402ac469", + "38": "e1241b6ccc8a38513719063c4044c94179f14c1a364a688d6ffb72ceb84802ae", + "39": "571dddd5f4b6138845ad79aed3ffc92aed7ae13b663390678f3fea9e1641698f", + "4": "952269e365f54bcd32703e5b07562b727ccd7c51f27c5ba570cde183e5c3a91f", + "40": "a8053c9a70f95e67c21bb2e6830d79d4af523bd5ff2d75e3abcb7bbe0f0558e7", + "41": "7b2fe7266c257a39ee4870d94edecd09b082dc1e88f3c0b0d58372fbe6b417b4", + "42": "dbc692de12e89eaa846d393f39a30c5c3f4a819cd71e9a6fff66dcac83ea54d3", + "43": "69d0bbc4357bd5f282d4f8c38a73652cba0db6b0bd2c67c22cc450e5b1aa7a5e", + "44": "c4045f34a441662d05cbedace9f567eeb27196325c98be7eaaf3710fafbd93db", + "45": "a933247995e58e4ee07ffd7a973be719db87a12c1610b0a80d944a1931f40f70", + "46": "248a30bcabee03057f549678bdd3ce844ef72133d4f8c183d9cd615c80ab27c3", + "47": "634d95a7ca62e9e3ddcf3114f543bbc4e6e83c8c1652c0c80b927f5747381ccf", + "48": "ca718ef8a2ca1208caab2e2d413ec2849da3d3e18bb3a968c7573bafe167c6ac", + "49": "1da57ee34387d35a3f7f2ffd1e451da77b6953fe88afaefe3d735ffe781d7126", + "5": "77a8cb2c5b0d44018373403cec28ecc606aec07a9116993a39eb243bba5c4efe", + "50": "fc7f60e06fc5b1e8cc6e55006155e36ef0f6a4dfbfd6767cccb55455d9146df7", + "51": "e603ec95c1858767b8b3b5653dd771bda6af39ed949680f23592f0e7e3b02357", + "52": "bb86cb63a7fed6a92fca3edd76bde36130ad1398c768a084b4a6f8593b3cac09", + "53": "b4ffa3fcf32e5c36edb3ca1bc5991492116d14beaa1aa1e0e3ecd4d7a083c02a", + "54": "1ab0964abc00fd9f8027e605692987477add15dd452e384aa43ed265813e1e18", + "55": "40d3462127f00a1e29809de4e5f190372fdedda88ac4ef2953c18d97da0dad25", + "56": "286ad9bfa77214e95d883171ea724b15bfb5d87163f62b1ff012200df96d41ef", + "57": "eee5907eb445fa9e393c3aac1cd418dcd976831fbec0d81016b377c1569a0b5b", + "58": "bec04fa7cb0ad1bb9c97397b8fdc9e003d7bb86f54b8790b76e0c6abfc3b10d9", + "6": "3e29a53bc91fa717c99e3b3aa1ad6ba152131af99de69d01c3c43e24973e4319", + "7": "0c9809831df1e01c317980d40d64acdeac96ae67b086beef58fb6f32cd390eba", + "8": "e2f87644b2be0cd03a9a62de2d6c0755212cf5ddd176e9d24174417d122967be", + "9": "759a4e30b2a0d4d9d34df454da62344d83f87503da8dd87ca79b55e2d0045e1e" }, "gfx950_attention_tiled_2d_fastkv_regp": { - "0": "a4c8150a95b082f8c1b7c30ddb66f98e6f7cbeaddac33e44145c26d0ce9b2c31", - "1": "ad26425a30d4b82e07cd386c13bbaf7be295f1b9344080fc3d61fd8c71ce029e", - "2": "f9b1c2fd068f8ad98baa1c8c530ab43f05c71db565b3c497571077f73d31e9f3", - "3": "fb39ace41275b01fe8dde795cad062df4e8c798e80ae7d364f61825941fe4089", - "4": "57297f16d0de1afcf3a1675e8ee573a4132c55ec52abd80ce14810d607949821", - "5": "a063e54aa2522eb9b9fce2f21def9d2c8419a6d5e336e659df94c4b801e0b9f3" + "0": "e3de6d4dd7a0e66b00ebb61eab7ba40cf3f9dacd910e834a67e0faeeb1d3921e", + "1": "532530ccecefe7b4c1cf085a5987955c5ed5e223cf1261f0b0a78a6a20a25032", + "2": "7f26dd7e52095be42734082dc59b818a5f75ac3130c540b15f244ea1ebdc7e7c", + "3": "15dae46a70e94f2bc332510904d724f846d0e857c21775e7dc8eecc25f74be9e", + "4": "233ba1c8f69997c272b74e9e7cd8dab50fab47f2d11f2382c387bfed9b28f1c7", + "5": "66a9e70454fc68630eddba35d6229705503685c8f95091c0c13ba7a4046f8f45", + "6": "faa49a0047ec908249a68cf85fcc471580126875015d0a852ae8388170ab449b", + "7": "0495650db5ea7510079340da73f0eba7916fab3f5c405b80d9e1a04d3a5ba5d5", + "8": "587bde17a07e9f9fd61c48819a8e7de5a82ba00267924213d05bd97f5d679a51" }, "gfx950_attention_tiled_3d": { - "0": "63c8cefe4601631d737cc00b2a88c787f723f2ba33cab0223da78d809c6c428c", - "1": "8a239b390bcd3c552c2e7f7bb9012c5b0ca292f86468d7bfda242bda1eb391fa", - "10": "27609aa7e76bf0f6c4ca2c54ca743d097b1f673e214c6569564021e8b96f12ba", - "11": "76edcfea215c922cb47cd15263b4400b4615ebfc9f8a06137c34c23ae409f675", - "12": "39ad47d2eca33acfecbc757d8a64de44c32bfbb1bcc5f8972bcd77baff135461", - "13": "7fdb55c05de4dccfeb8f1b15ae41c25ebf399b98f7b8b1c3896a50478d8cadff", - "14": "a86b7bf57b351c35c60e17f70a593e3036bc35bae025454bdc1589840a1645ac", - "15": "77cbc7e611c9046832a959c50392299f8f007f65024a66d6fa6fd9d5ce3bef36", - "2": "7ca8ac489d8db2ad713d0c94d5e2b5814a38e1df54d82105e8ba1232ed6d1115", - "3": "030fa0f8de882d8ecf675814aeed194d744c6642345c2e1f87a736f386a5ff7b", - "4": "206a5d7824acf278b093b00d7b26723a9b2e6815b66e097d64296c4df3518ec9", - "5": "2f141c4cb4050ea9408f0597541b1427dc73c4a000b08a03e8c9013633fc85d9", - "6": "e9059831e394df815a98157611853cce5f53641781c5da8202dd05550956ac52", - "7": "96e48c0a6a6cc1464a40bfb01aa61a205a37e5dc446125addc3737f3e1e71eaf", - "8": "ed364897692af264f57616e281ced0602de18fccc5b0c9c788f590a8f7e3e5e8", - "9": "a2c55c19220b6d907582f436785f1495c1ab5dab4549c3f3ff6e45f34ac7f6c9" + "0": "794ef52dabc1565190db106078194ea67db68870c0f19e229148f6418f4c2c8c", + "1": "b93369f85eab79d0dd554e8d9c3296b51c2cc460a4aca45b17141a6657793a11", + "10": "7f2a38dcc59a132fe20a9770b20e806659f6a63ee2794edc1712f8ecbc6c9581", + "11": "39cfa97d6cfb65524952e5c56bffb65e970d71ed464d71c21082c8bf8b3cd151", + "12": "2d30096a2f0f39c5556acd2a8a9ec014369ea9e9fba758afa2b0e9ae70e2d4c6", + "13": "c1a45fe713393b43e90642c66c4c6ed045ffd9c854b1c006954f6724613f356b", + "14": "1bb245d5a9c95809a84a423b4f2502479c0ac60f7f2f496cfbdcba06d4f18c69", + "15": "f9d045dd07f039f52374cb96c4e67eee5d640dbbf5debce34b2d56d7adc9f060", + "2": "5e8da8b6d8d6a96d6f5eb1e74b05f449a69d35bca1dabe1ad3f848839e433085", + "3": "e68b749c715ce0e26d17b3dc8b356639764eaffcd64e7a0b93c5878c1134928f", + "4": "8e6255d95fca99ac69a29158860ddff1b38c405be43ae8cc1b7979aca4736ba0", + "5": "8641b5c679a1b04b9aa07bec3ef984c3c80db2702cb87806c6f0ce74dfc5cc98", + "6": "f17b42418b62a82c894431dff094caf6da7e03e57f21c8762b22e474a220d099", + "7": "a6a9a2c2cec53389bbb153f2518c44f230344157b6f5d6d4023995a36900a16e", + "8": "77862d101a9fb155ce10e340ff9de1755e4607c6d471a5fc8b0e87fdabd0abc3", + "9": "40d658bd39ef0ecdb80772d43aacfb64b402173ca07d2b56989e6c9613f4ea3b" }, "gfx950_block_scaled_gemm": { "0": "2a9ffaa380e90d74c3ae5e89b280ee5cbec54917e940b4ebc05541adcdd7ac9c", @@ -1752,18 +1758,18 @@ "attention_strided_kv": { "0": "8cfb54512441a845a2ae325065291e39e5367be96f320596e8622d7e308ec610", "1": "b1cee01d468e92ecb5b8571323ffd897bff01c7d2e42dc8ba50814387f908a63", - "10": "05d6f2c670f9c66f59974f7f2715eed844799ac869e2765d3d1f691bc417d2da", - "11": "68530d1c8c9c9a12300743809054235d88e2a1698f811bbd86a531eb47113d40", + "10": "1238b48be61bb4a6c83e3581ed269a5de59b717901772d81d17577988a788e9e", + "11": "c996154b2f38e896037f4e14e8574c7c85d938d297ec93068407f0cca78cbde8", "12": "6adc12acb6dd75f05927ef1e4f6683b85c780abae25d76588bf57ae25f8dfd88", - "13": "d7982e56673824863076878c3461209f8557cfa8507b102071eae22d712b6414", + "13": "5266e99d4e82df6558795288cbdaef86099a84229f002945d227977834b89ffe", "2": "a80d03691e079aacfb8653f7054a836d2ecd57256285274fa51c52a8a398c727", - "3": "272f3227c873b525dc8e48c4f024f7885de375178e20f4c4d6686b4724d82c0b", - "4": "b25f6bfc84f5db8903883a85671f13b109256697e33c67877686b4d90ac07606", - "5": "bde0dfb3994fb2a89081fa05809323059111a5e7d4df55ed8859ec2f6b0cd930", + "3": "e9f5684da874f4c1bedfccc3a38750263a8d9364803d660871cba0b1176be7f4", + "4": "67c89f7a4550323dd055e30f2fdf5c12a7708b41f2af7e9cfdf1dc5faaa090e7", + "5": "77d12b3c9210688cfa94e005484498b7c285e7c173f48b46d2b87cbf9a3c3ff2", "6": "23c2881b225837148b997823910e598ad844d781441a1f122a5d370b4b97900c", "7": "966f73aef9172bbe60348f970553576d609299131e029e6842e1f4dc49c96da1", "8": "c5f4baa4b7f0237ce60bf63f6694fcca24bdd440c0fbc34a1dc0a6c8947b791c", - "9": "7866c9eb6d414e0edf3678a477c387e2742057eed6a96616544f29057d609d0a" + "9": "977a23a9b768a42645ef5dd393278c4d577de5f3cb0c5b7d847260c1bf0b4462" }, "attention_unified": { "0": "64362d51c5c69a238b6ef0315e46e670e8fbe98f2aab1695d4bbe43bd1365014", @@ -2937,88 +2943,94 @@ "9": "f9794098f04a1cfad57dad0a169e28e5915f03806cfe4f573c50602f04132a6b" }, "gfx950_attention_tiled_2d": { - "0": "f8d4b006fc2a64a141ad3099defa8f2a1ca8bbe6307700292dc4c72a951acee1", - "1": "3fd1b2cb0b3030c1642fe5d8d7736bc1dc39935de5c229d6e4fc64348b6c6cc1", - "10": "3be6ab529bf1b83a8c79acb73a4d31091ae69f1d51f543ea3961d95ac250d228", - "11": "c699a7c2c5ff85ef053df3f1d8ffde08d07ff2b4812c4461a2ccc62601b605ad", - "12": "5329dc933a77f35a2007fad3a88950a5a706589c26eea780842abd315bfda434", - "13": "34ccf508be395083d937a5932a7e7991641ce78e7cfc1b6d83f5bde3511ebe98", - "14": "f13fbc763699d0453ccde62e1ee805186be481649202b965a706f02c21dc96a1", - "15": "28952d5c90d0f0ef6a344ed98739710f037a78287d91c4ea6167760573d3c55e", - "16": "e070a5ae005cfdbcdc2f19697a43901ee8112ed496dad04132e79bb25d43644f", - "17": "eda873b82842ca378d2776131bc34ce9f88b6bb2902942b3538eb7cdda416560", - "18": "204450366ce40405c3c9f4adde05ac8dfb61363743c52fd88e631e93dba2845b", - "19": "153041d81bef99e10d2a09ae3f7aa68ce2d1e34592f2998dc36ade4e4dc76dca", - "2": "3e52ea6ad02ca2a7800e6d89572973e4094d700a793df3678bbf9ee84e0dfbf0", - "20": "64b3c1e810bf1543004443df45f156ae89a95684ff696c971d0cda6fe1818140", - "21": "a81f19a8a4e84711772faf257fdee60d053add8c12d89cc085018241f23be905", - "22": "2a45eae61b462853b9d7ff4854fec2c6396c991c719d875f48adcd09d5727392", - "23": "93bbebf02c4b3cc284c0dbaf67695ea8f422e84b16faef2264e392b601e6cfeb", - "24": "0e205be3008102649aa2a1d197fcc3d2813bb11b267280f9925e52712c25a3f5", - "25": "8dd9a33e53173c5e682f625836dc1b2bb13feffa426f1d3fcf54e80c9c1d7597", - "26": "fa4996112f26a831bbee7a71687912c0a2cc5b3d14a1dd485bf7487018c25f9c", - "27": "ecc250c6b37ddc3d087f9c43ef1cfbcd856ed012351cf82d7f54f453ced8b1ec", - "28": "2dfcc80e8b66bcfcb15ae9a032d8b62ca1a2e8fc1d74373b1521592f3f54ecc4", - "29": "2390c228a644d0fe305d99e0dd2f606df52a361d406f5ba16a8bd95300dd23ec", - "3": "a2d9ac98063b7c95d03eda9777ec83a2e58077155b36eb1b9e1e2486c67ba6ed", - "30": "3efdd6ae314696b4c738df0ec739e185cc6359734aeb6524aadaa83b047d3bc7", - "31": "73fb20fc03e38648da94d1e11da2c4de60045246ef29054e3dc8746452083201", - "32": "7f892189170be9ce77db471b13e49b6a7769709cf748ec7aab349019768b411d", - "33": "9fdb1ff71d73e5490cd15b6f5567eb35985b9df08cd5973a9dff792095c0a6cf", - "34": "9618014a69973256c2af3e0602f1f6a8b489e4452eb8615fbab49928768d5ddd", - "35": "baad20333c440d0d82e273855df1b91d1febc29ddf3002c2943ab1aea59866e6", - "36": "52145a337031a2081d82745d40783d1dc280601133ad03cf6a334a7fdd359f7a", - "37": "d4845ee7ce04dd527c88982f6d13e4bf2c96727b936d8718c047ea3111a3cd5a", - "38": "4e534de061ebdcdda0c354b1c690644b0b1cb684515ee8493494de420a9d40fc", - "39": "bffabe074fe35697b952b43821d6ece576271da577590fbfd36d1a2b297dc1f7", - "4": "7d3c7454603e7bc78fc430b4300af791955e398b479ce53c72802f663548352f", - "40": "9342eba3cbc54373b2ce8ed5d2c6143a499d86afecf77de6f228dcb4c6b4f911", - "41": "8001fa88a41e98f713906df139865ad70734658a4c2132c6ebe6b9cabe0d12e0", - "42": "1a125a0daad95609c3b8cc3f6f3916e8e9c5f583be4fc9ee1bbd422bc08d6b97", - "43": "4376d9e26f3a97599aa55d4600dc65a954ce8361e7b7057ea2a4f8807379e8a7", - "44": "2beeefdce488c6ef2db1c1408fd142370affb2b9c1018d8366c7d438f666296b", - "45": "ffcc24514feed181298eca60e41be5c85b15e6ebb04165b825f02fc8d4aad1e9", - "46": "7d6b496b79928b6ef33fa82d56df4839c0aa16dffca6a0efd2b44000a198c2da", - "47": "477a2c76b1652047d67fd995a37ad07a924b2db83c9b836244b8a0f79f845bad", - "48": "27ccd36f5172739443afcd092f75a06677c150273f9e00be361ce8f893eeb987", - "49": "417b8561f36c9024fbd2f1c66b5a4c688fdce1579263ff8255404f22d318f156", - "5": "27454ea01193a02cb95747f60d6683f08584bf0007355cd264ea5ca9fae00726", - "50": "42cb73b9efec995f2d2510fd3f4821fde4af119d109f6860f240ef1fb84d0a96", - "51": "37811e034ddf56a2cbcdd802684ab9390952ba4ff3c4702f904e14dfe8007b03", - "52": "d25925c4082ddeca75c72d05aa79b5e0b73bc11745afeee959efb9712cf10dd8", - "53": "fbb3913dda415398fe593588b79842d74230fb7c5c918950ca87fc7625ca1e75", - "54": "59487b61670b4a62dcea91ee659fcdabd595d06bdec14fd4db94c6c92d807eb2", - "55": "94e09ed1ce5b5ad97cd73408730b6c16a1bc0a0065c836ffe0d849cd1b7ee7af", - "6": "a8e770f914300d621e67fe470600bf63f5368f0c298b0390fd439cd96d0cc247", - "7": "758aec903672f77f8a5ddc109e3571302baa37470d709972d829547d7eba4783", - "8": "43299e6f6d1b5c05ea7ab4b81ddaf3693b0e87ce51bedaf0f5d6d70855707566", - "9": "49bf3ded9d268ab03c2a8d6d371b2e6dca582177e26dfd89a58167626e276594" + "0": "8418f07d1282dbf4c02ccee4545f1d1cc21d2150a025eae50ed4ae31d2edf5ee", + "1": "ad58ac14b02cd66f1174761c626f309f531d5fc4e5b3a00830103d019e3ad3fe", + "10": "8428165fedee9fbece8a185e5ed2922c790d88fb6caf9428623345be4232df2f", + "11": "c90d00f49adbaeb70d6c14489bf05f17a2cf3e6e4011eb724ee93a0d15b9de47", + "12": "3e6c0e1a14494fde091d449b5783c3915af8500ebff2aa190dc2c9cfe208c4e1", + "13": "063e418d032a3244514052f6cfa003ea9ba9d68c59fc4f0880cffe3ed4cfbb09", + "14": "c51a79503e1b99f80012a7859f49082f8b3955cd3d5d3af939ff647df65014ea", + "15": "a4de49aa21870c56980a1fb09cff324a47cb35b16c186ec43e814ef1ab31508c", + "16": "87869d8ec5b5bcab9488026fe73ebcbe1230448b9a760f8eff87a1e770971798", + "17": "1b74e403eb2f00cabfde3b750ecdc74e30763a996c0f0e9b44dcd25e2cca6609", + "18": "f3a8cc1271a9f446ca2734c25a8b9d0d2d9765f27cd4b9a6d4965564b4f92d01", + "19": "61bc43d9621c9f1bd0b34b778a97af985efb7567a0b80f0da5f24980aa4a7164", + "2": "68f48e6292f8350d38fdbbb8458540bce28a0f146c1e3d80f54819beef0aba10", + "20": "cff91be2d44012e501d22eb232b8705994e08a6b48f1c241ac6c3ee4f9ffe536", + "21": "c25a3b0c91aaf8ca628f75264b72e56362e9861fb47a282b8d7a56034d83d9fb", + "22": "a2f103544de99d1bef71d2cc9a993e28c65e1f3dba19e9b30d38c981dc6c9446", + "23": "ee8fdf9a3cc7fb49c8598cac44a7ae1ff51a1b2cc15a5f4f79f8ecb4ffcffd6c", + "24": "6e023513d597926e7c41ec4f4f4523ad55d29cabf082b719471b03a0188eb08a", + "25": "c90d00f49adbaeb70d6c14489bf05f17a2cf3e6e4011eb724ee93a0d15b9de47", + "26": "c90d00f49adbaeb70d6c14489bf05f17a2cf3e6e4011eb724ee93a0d15b9de47", + "27": "75e025f02c84be176ba504c84629bb106e5afd1776854157bd603db855103223", + "28": "b9c850b5611d60ffe7b40b6563ad8ce24c193bf24befa35f09b39d0de31fa231", + "29": "ac86511613bc25857cb299a10425784b44b1bc874604b1bf9a30b52255a15221", + "3": "76338cf03716801e4e27068fc578892e4fe47caab060fb580e64c426a3ef0836", + "30": "64a14898a78a578de868b84fee84051cdc0999bb00b83ad17f486b6266906510", + "31": "fa74be3d36ccc07d840c6043fd390699e94e3e44c58b5936f34fd3f437fada96", + "32": "59bb465333f70ec7b3a04a90c23b6112832f29de14cee3fdcebc9620bc75f4db", + "33": "0064b4ce915a9e514c972a2fa2668a6481562b796597b73db06ef4c28dc146e3", + "34": "4aaabcde1b7c9f5eb7f8bc8f5775935a0beb1216edf28f5be1b63fc8efc7a866", + "35": "31df4cd1a814ea6847b1fc43ff2619eded86df4b5e402e71ee53f5a75ff9c809", + "36": "d6b5b4b63bb19f070b83fb55dfa1cc38bf48b810551c53ecdd42123799840a26", + "37": "c7b48476c019af8627c88c259e275ae2a5cdbf667ff0e584ffeca848e104e1d3", + "38": "52df06a1d2ee3b0055256b5c1930321e0a6e88e426a6af8ace252ac63110c71d", + "39": "e98a91a24339f5f8a936536910421092ade2ac313f7aa371a974ab130cb7c284", + "4": "0c60ad5ae99b8e29e9f21fc8caea004607d8004740df1a0ed03a4a640c56d007", + "40": "a0e27a89d0a7e1e10247569ae29b0106024f7c6bc4cdc9f71f508eeed33883c1", + "41": "45199c39bc3a0d5d138268fbd069caf56494c896247cb3a260fbfa82c498ef6b", + "42": "16bf6a1f3edd2db9bf58866a4b17bce589005bd6403a8ce20904fd1badcc8ad0", + "43": "897007ab1bdd4c3ce0f88058529c2f0398e1d26511edee68cb3701220b852d9a", + "44": "f55f18a9a75310290e18a43d0d348a8a3e85fc88ef6b23e07ef14e7d73ce4c1d", + "45": "26ea61efe5cea6b8785ac4be000573d04671777bb71c9783ceb598beca9284be", + "46": "d8054da7c13cad2ac7a8e2ff60418740880f5a6ebeb1e6a99d629bbaf72b61c9", + "47": "b2f49a3e667216992b2b4003fd112a6b525e43e3aae037831d097c95930d6bfa", + "48": "d5ca836f0e25e0aa3ec65b2ac7748fcde5b36a2b16214dad966386a4cb93b2b7", + "49": "9b272697484920ec9704cabae391d3d7888176086778e2a423f241733e2b80db", + "5": "c90d00f49adbaeb70d6c14489bf05f17a2cf3e6e4011eb724ee93a0d15b9de47", + "50": "0ef77b4937e20911b63a35d7921bb38efd431959c5b57696af9c2ae35f60e98f", + "51": "23376d55e779f56abb8a16125cde2ce5e309b75947321f3f386fd2dba42dd28b", + "52": "4357209ea9871d1e150e5a9d31ae6497a7cc70ea71ee76a4092950a641846672", + "53": "a6f3716314400f2d5b033e39d76aceb5643eb31ce3cef93ad5a6c04e70a1ff55", + "54": "c7d22d271ef6c66e1c02a40b540aee15cd560cddb08f05f3623b702f348b67dd", + "55": "5e6ac1abe5f83dbb53ee0d0dbb0e796204a229b5f9a7d08c64ace181a06e8806", + "56": "940b5ae04863b17ec33fc6148e02b0a5156cee1ae0291488137e9308e4311a9b", + "57": "f4b181612b3e8ed9a84ffc8be3b9e2e5c87019cef5abcfee3e2f34f685f6fa25", + "58": "e1992defd14591924e4d351b76e96a1ef175b4add0daa76b816b0891d967dc9a", + "6": "56f34291d823fe3e6a3a9a9053990c71fc97b27fd58d942b4f0362f541c49051", + "7": "02c8b639d85c38ac2e372bbc0d01cb5fd15c9b2e7f651474903516ff0713db2c", + "8": "3439187b9c8410bc01047b6887587e517b716a269f7dc87479b12bd77a5e32a2", + "9": "25f85b4ef668051a224e218000c466492f2cb8223e59525c310c31c2b718804a" }, "gfx950_attention_tiled_2d_fastkv_regp": { - "0": "a0b3959ec02993bed82b8a6f9bbda8b8f04b6a32532cdcc98b40c02d7d18fcef", - "1": "baedc71fdfb1039fa0a8fc3f046eca8edc17e7e500a9e0b081b404a9e50e605d", - "2": "0feec4b8b01ede6a4bea69de371c1c389a7187f880a2676c9b214739a3478e71", - "3": "b70baf0a5756a8153d3bda95dbe72efdadfa6285ae6148f914b4e817bfe9c6d1", - "4": "c0f8151997ed98c54ca51d3bff50761e5f4f629997edc66e3ac789e684e2b0cb", - "5": "071a4a0bd27aa461838b0242e2ddd9e0f97b67d2753ba9be0ace3b81699a67fd" + "0": "d233c1b6836cb4efb4944e467eca48734192718a5cef19ebac0b02260c8128fd", + "1": "fc2ab8924bc64fb42ee0e4c76445ffbdd8f7dfbb2d97a141dddccba61b09890c", + "2": "a7324eaae142d606de910aeb0f0cd90eac0cac87473e62c35ae2f988a68e5add", + "3": "f6fc0254a041ec5037c397a65eadb433fe4203d1409926d95b90ce842c52f278", + "4": "3bc26b865ca95ed979c3397e4d803b0bdf5a404dfc259fb413ae2b1a3a0bda5a", + "5": "c493b4db35adca98da972d7e2f22a0885f36e3fc5192f944674e08369b456a07", + "6": "22406abd8e1df7d5622935fd1bf471f53b7fbe72837fe022de7134901de9aa52", + "7": "6a613ef20711a54842af2c1031cb27880dc0c3068ec147ab4b33a7568b9f55f2", + "8": "cd973c18da740ecb042d487ffb5262157122e21166d13203caa78d22fb782933" }, "gfx950_attention_tiled_3d": { - "0": "211ac134586448e278d5662ad89a605338b755e91bf63d1584e267dfc79254bf", - "1": "d4edaafd12d2efec90c73cefb53184fcf0d82bec148db3cd292db16a342867a8", - "10": "8cd481553082c11d6302bde9d47605eb0aedc7f081c24d212e5c09c808fa5895", - "11": "da1378094b7054d0afd8811eaf12733e7aa899e4a675c25555fa93823f250269", - "12": "9460de8acdaea4eb6feb3489f338fc4ab9b0c8f56abc16221f4d68fe23f01943", - "13": "f9a6211224a477eae84fa9a316ad84dc0674756e7fff8dd663c536438b57d6b1", - "14": "d95a4ee4691a5d4498aca014e680cc4bb9dad85731e6f7b913aa005b7542e0ea", - "15": "5922ae9e30f10d84ba06293cf1b9f530872f26ac4b31d264ee0e7e41d96f85f2", - "2": "43225f4bde00693107c9eba8bb29ede7589c65b60cee13b6e310bce7260bdce2", - "3": "011317d12e0cd324dfdf21438343e02c900930fb3c5ccd02ec2752f73188734b", - "4": "68468fbf2b4c26089c675cc417e8af43b308a6a7838e9604853845fa154e6dd4", - "5": "fe17ac0fd5f0ba6ee6ceda19ee333e679d2734a21d083aac0a84e306b2a66d49", - "6": "5a40a5a12f9bb3b9e0d83625990829bdbf4b8165dc92b8c50dcd3f8321907766", - "7": "6a4c575380ba8cc9c2ed2bf2b7c645cbc37bec336a11eb35d2c6c2b6d8f7eabe", - "8": "79ff0b255d637b2b899291ca7a8ec600d0cde17280932b2e0279fcc0b37f408e", - "9": "e125f3349708cf7b6e6e4066957438905b2a9ff7e93da177de8d0771d83e0658" + "0": "fe675929a9634bd88d72b450c85337132c928137776c495e32b5f28ca255d380", + "1": "57ab70bab54350aa49d3baba87fa087105836e6842425649c314ebaefe8c7732", + "10": "aff41145919896c345f3f878958e01c97e0813ad7d5002c5d0b8751f130ad8a0", + "11": "9ff1a50a1380881e4e300ff8dffa21b83509eb9719688096c2d236d9a00820c1", + "12": "fe0123e9b13702deccaf6953b12d74929850aa602555419e8e6385eed184b272", + "13": "cfcc2c5f687396f119a6dd8a1e252d7f73be717d230c0a2591210ff0274dc1ba", + "14": "b0bf7e57e30193170e54bde16a00c4bb2db9a395d3d7d479520f8934918ce9ba", + "15": "ebfbf262cbf7c31026eb26a20825d580de27864ae6a01c46de1c5fb47f34f00f", + "2": "d894f9af92f6e93f9c3c31ea1eae12bb5e41a0a48d553e8c4d4f78897522161b", + "3": "a23f992ad7619d51697ff9fe98320985bbd24dde1dc0119cc4e86e508561f3bb", + "4": "106ae0d2ec8b80c3ea740f952f09de97420489008902dd0d8dd99d26d48ad098", + "5": "ce56b406998b94386798028bd17a420415cc79f5e36536db08766d6ea2815434", + "6": "37d8ec6157885dab20d3ada0a61745c9741e5d0f26c475359b48321e6611d0d9", + "7": "3ea23963d10f6d71a867a5aab0dc65813f5f6f38728c9f575dfc4b8a5834fd01", + "8": "9046803e1f3e588a2e41b45f7d885fced25771097d95f0598ea2f7b4856ab544", + "9": "135d33061aeb2d7494e23c685581fce8caf78a539420fac27cf3936aca4fcac5" }, "gfx950_block_scaled_gemm": { "0": "66bb0b6afc125672530de28db30992a4984e6aaf449efbfa50712adae21d59b2", @@ -3397,18 +3409,18 @@ "attention_strided_kv": { "0": "14ccc0e86cb9a96163a8c32812cd718dc454f726dacced64fa3cf07db5c70c90", "1": "62086a1cb544ba32772dc112e252813ffff1961bc2db23401bf0003bce718b73", - "10": "35992a4d15f09b54bd668d4b21556a2c871fa8c6f8aa39a89169736e303476ea", - "11": "ec52f6ee37cff049c8b81eb579b71c32be965bb4730fc7b30bd477196499727f", + "10": "0c4719c51595a01254f8a7c7c26bcf7cf10f7164fe54dc0c697fe36b938bc1b9", + "11": "c8bdc92a7fafee7c1187375671ea8caab00edceaf7b5f82c844d8343767f22b3", "12": "b543ef77d6fd6bcdeab9b00f9c7dd8d9823777987f2e9e09b79195f3d230be37", - "13": "f71e2da33c81aacfd3a6690c0fa2c540aa3a83586e0aa8ed0f1ed638023677db", + "13": "a2a213edfcbe00373fe8d3962f4b20a7216c9489a0c5c039536d63483f5bae57", "2": "f2eb0fd45eb3693671f6d11bdce0eac92295e784e74efffe21351150e2807039", - "3": "c88f5e1109c70e0c4f132b05578c6abf440d04bccb52bd774b492d67d3acc9bb", - "4": "3e5522113635114b994bb40cf19bc51a417765ac827cbfa879acb45cb0e62796", - "5": "3c03d60f49863bc856426e1431cb98e8e682da4b6818f1df320536fab12b1314", + "3": "3a8fe959ee0737f990adf184e812dfe983b24e0f50cfdbcad667aeda89eade72", + "4": "2f21d041985082346f34105cef0ed4b58df655eed4c600a176deaaec0350cfc9", + "5": "d2770b39913bb089531c1db4d9fc31db4114053922f0bdbca3b34cdcad122a69", "6": "b4720021fe24f67cfdd7b6854a25f4b371cd5c086d7d804e4a1a451d8f885c82", "7": "281cf49847fec9bf2f0133d0b0f83c9ff72e99e3f9bfd846d2d056289604609b", "8": "8e1697deca33346efb263397db405b14b835a617e0fb8a9d6e5dd19958d767e4", - "9": "43cbf9fc03696e3fcf0a132e31a769d92cd48267eaea3526e2cc7c46c35225da" + "9": "5f423f4b85911405a7548c7641e955fcdb485e040872744970c16c92326ae3e1" }, "attention_unified": { "0": "63b0e18b1115541aa4761eefe2dffc2e7ce9f00563e55079ddaffe5c7f61121f", @@ -4582,88 +4594,94 @@ "9": "de3ca04e9008824143d5e906a0542a32367d1402fbbf49ffc33dc9ccaa0d333d" }, "gfx950_attention_tiled_2d": { - "0": "0a23fb0cacc8dae6e12f8c82403e1fcbddb2445715fac624d8519dcd79d5e2ce", - "1": "eb3b0b047c578a087b1b7cf87b73fd78cd317ea5360445f355ef1209b647c18f", - "10": "b94995dc503be7d4bd42e3b2d3d5ddfb5b6038bab6ec5f3f73002cc1d051ab35", - "11": "14128d04529d2eee6cae4b42997920617481ef487066b1929c707d0c5ee7bf95", - "12": "4563234ba35a529fa6ff0d9b5cf4b291ac6d0454af67a5fb68ba2061ded93b30", - "13": "50f304a5f0fbf4fdcfa3ac0f56aed3719bbc84fe8f8efb66a9c4ad51a4faa1e5", - "14": "1f8feaba93da23316ed6adb7f8342f9f52841b043196d93ab8073118be73fca0", - "15": "d1280c506b2059d12b42e2ade24bb924a0dfe9a354702a6881327785a98ed306", - "16": "71870be41bf69919450417b5084528bde114b7fdf472dee211938b30e02c177a", - "17": "68b3e487128ff136f75ea16fa9d3bf9088036defe7f6146adfae933ef2cb0027", - "18": "062f56f5da3047d0f73c25ac0115a65df45a6a7320fcd41c8752be0e82b14604", - "19": "145bb32e2e8cdbde9f30fdfe309b4f85688bdd88ed9eae6b4fb279887a797b22", - "2": "af13ad6004494575067ad487c80fc88b39fb9d1a54669bc76bf1d7aa2d7d19fa", - "20": "f7ecb54b12ffc809a9c75445fea2735b82738ad97183e95c9e29f5a249696ecc", - "21": "06a0f58c5ad46954af4869ae7d1ae994cd63753a70f7baf6bb79e2e0e7c9320d", - "22": "729f82ab7e2037a3f132d5c10038bae697c08ca392fec411134331269768e7bd", - "23": "e3ceef9bfb0b05dd484864f2426310f7995b7cd2aed34c6972d3c563eb1c0a0a", - "24": "5b7fd415bf9882e0562853ca965eb9dcb86041a10da4882f962321b5b96e4641", - "25": "a0dde86c3ee04a8693f330c021b4d446276d8168d0f589f34d3e5d1741784327", - "26": "cb9c8fdd3f3205a3146a77c708c5ba127e1191ad4dafd2ec2b1594cf44736c18", - "27": "67352ade3d1110d321cb01f0569eeca1c43fc124b0c3088e4cf802dfdcf5816a", - "28": "3f395d2aa78799f1b09482b39cbf5ba09b645435f9bd4b99ab134199d4b0b13f", - "29": "8200282e4e8b9ee44c05cb1e48d1bb1839860599a73a4c6cf9228e36fdc2082a", - "3": "061f7c500c1e4d801e8d35aa0a1d899f02db0e1e92a3ab3ada9460080b719815", - "30": "78e3cd080fa4cffc891dec6618b9d11d0f3ec543843cadcb9f180c157641ffb5", - "31": "9c9dd61b7b279b6b310e9ef312db015881092905fc6c11b8a799a147062853d6", - "32": "3ae9bc57dda3195d679fff57c757a66b73b660aef1e76e7109c5db70abb49e51", - "33": "447bb16da5473e15803e19597672ed613ebb5514f9f1688ff01e17ca48dd6f9b", - "34": "94ff011215964d95efbe5aaaf9f09aaed9073aac73ebab189f9be4814d07c298", - "35": "507d70baa755c0933bd06ff2046364247bac737153dc576f465333a9816a20c7", - "36": "60152f214bcbed0c2f77772814423da6f1bb378cceb7e78e9e421d92c5248b50", - "37": "3ad82567405e0b7def5c742f7fdd1f08e033cabbf1090242f54b325c892766c7", - "38": "60d33194d2677453a497e54d0d1dbf048885819605cbcf7af1605eabfdda96e0", - "39": "4e02711efd3101b554d8071ff31c50b99d2dbb271a8a09962a04644ae0e92da5", - "4": "5a4820726133c386135da36d962e6f935dae94eb1b421d5d8d0efdda8b232cd9", - "40": "60a55ca1764928a06058b05818c274662ebe616bd71d6f818a51c5c76d307bf2", - "41": "2a96cb4c4bf58d1df223e65719644573d21541f51235ffcdda4d59493ff227ef", - "42": "a13ff696925abbc6c461cdd3ecb08a601bbbd5753907f5bbdbee900dfa1ebff0", - "43": "73e0cdb40d32ab7cd21493e8d285d4444ad6bd9ec94e4be5edd6aba811537efc", - "44": "97a934ca73f4ff1cfd9f310064d38d6541daa1e1dfedfb15840d4bb20c53624d", - "45": "0bacb8b3eb61b42e9e95d16ab918aa3bded2b49f6564d7c7919a15c441b3f1b0", - "46": "f4aebddb56e2fa7da78e1964970d2221d6bcae8cc3634bf3f294cd81ff99064b", - "47": "f4b819dbc4c06dfa74c3356a96a73fbdbb3338c0440914bfcb9f235f6d7aa27d", - "48": "ab2c00358daa7f1455b0cd44cc11289b89191c567616afbff393ef54c8966ef3", - "49": "eabe56ec4f189d98b199e25f468f944ab099a2e7222e8eb90dbfb4d30fef0d0f", - "5": "9096f0a92f1e04a20d7ab1d0a5d8024fb53b2c4f5cb9905147ca6812ccff65e2", - "50": "a5a9e77b94d74b00cddaf3e70cf6438ccaadfca168c4b2f172558cf13321695f", - "51": "eeda727b98f24882aba1ce68ceb39b856281310dd770506c2c504eeaad7a58e3", - "52": "8132e501ba57380b65eb1bf3957c2d92774f62dc1421de516464a4950552f283", - "53": "b7221e797dd9eb2a5678a126fc0143f933ced071904055eb2461d0ba011db7fb", - "54": "be9dd416f73bd072385cde4ca0be82c7f5c34405c5a1d51910f35b396f2a7ab6", - "55": "1231b94b52ee265b1f17afcd6967955c00b483fc6308dac27c0dbe25ec4d6e23", - "6": "9cfce58150ae462a04619701d8d43d3f2c31e3befadb1203c65d69a639dd948e", - "7": "286f9016a8a3230e2e5d927897e81708b099862c61afabc890f50076e4ae0462", - "8": "3fe933c10995d0eaa679ddb53b935f1f44eb9bb23f7f0fe68ff84a598c557d8c", - "9": "44dc12f2553925df374df58403abbd0ee6e754c8987323d490199d742334b026" + "0": "69afd30813461604e1fd18434cb72c1562a254210eff4874f4923d7825db5646", + "1": "97e111cb68bc08ca8a6b6d7d9ee85507d0a19a0a4968903c59e19300cb1cb037", + "10": "e2238915829fe82b1802f884e7fc1d39a4ef2edecdab72c72cdbe403b755a387", + "11": "5bbfe2c8fa835bf493ad649f084f2e17eeb32262fd128f80af633f6cc6a6aa38", + "12": "48e907a2906be8d2a013d7c111cb8da4590b7991d8ae9cd0b1db590e3f415684", + "13": "92f10e18645cd74b066334ece11f885dc5544770b32c6fd46fc8ce8b971fbcc2", + "14": "69940c9cca1e12209832a16c9bcfa28957b9a8a257637bb2418415697d687c0d", + "15": "5fda4a0ecd5ff4be90c0783e730902f9632dc841dd408dbdf74d518dd4fa41c5", + "16": "e7fb03b1568c2a32e90dca6318a6d4035127da0ba0f5495a448153d521a6f49c", + "17": "0d991047d8f92ab749c7eaadbd1d2cb79e02725e71487be75ffd306b392f6ee6", + "18": "43c4eff974f6bae0cfd02fe0e07f758e1b5969123a4520bcfd005c80ef1d3900", + "19": "b115d44793217ba408d95795766dd0e3c73cb76cbe1bae7e7954d83102c86dbd", + "2": "ca3f65a5836bf3eb1d6f424cd6448e927d4a0a52c383a998d28025bb25f31d65", + "20": "02b97fe48ee19c070f9a109edbac35ced69ff253bd11ee41ad865e1b7cda368b", + "21": "c396077ca11123155a0f43beb8ebb9bdbf4e03c0a6dd5faf31cbfdae2d92cd11", + "22": "c466b5575586898e8a44891fbe660681eab81ebf36a46ae5a3be7f7ee867fa77", + "23": "a4c9151bcc11733a91fd71b5fc529ff9ed25e29725182e7f999076936a3a7004", + "24": "e70bf9089204593cec414ec1169a128b7593e9207a3e27588cb6775daf95ddf2", + "25": "5bbfe2c8fa835bf493ad649f084f2e17eeb32262fd128f80af633f6cc6a6aa38", + "26": "5bbfe2c8fa835bf493ad649f084f2e17eeb32262fd128f80af633f6cc6a6aa38", + "27": "932a5dcd3f077fbaf1acd8f7f6e161fd4ea558f2f8fc32aaa0f2f39e6b65152b", + "28": "eabe2798f89d47ebffc8b6b2c93df8b4b2c2c7e768f3e0504898ec903d01429a", + "29": "76c08ad4b6cfe7088a0e6ed9eeb9edac5d43f2880c327791b814452aada768fa", + "3": "7e738a6c24ab5f682581f132e8d5a9459c504d0137edd8e5e917cf778268eae0", + "30": "6b45e64cd1c6ab6cc80d6ee3f88cf518d1cf8c61be5bd1a63407cbdacb4902fc", + "31": "d03e155604eae6c093e51e0895a350a2c1d97d3d071c22a34ec58cacb7a411a7", + "32": "2fd23e33cf6f0a3497f696c2adcb9a33b28d2fe2fbb62c33b7e2918ef2266b77", + "33": "0cab19b10754fc4552ed30576b4c560307d6ff59a5a4acb898547ca65d1e8a33", + "34": "cfb0a56a84657fb71d9244b96b2e182dd82fd08ddd57e6cf7b52f110b13cd7da", + "35": "8d4c0305fe45f0683f5d77831f7e8fa0b388be92c4e4a0102dcb4288e2887fdc", + "36": "f12075855414f1cd378b6806b231ba9e12449c1942c0f3015f3fae113977d1e6", + "37": "1c25a4337fe9bfa18651be1665eb87a85f52315e1d0d089e240e2583c2fbd530", + "38": "6240443479c0adf218f2729d0ec626b7be6a99134ef11e74aa3cd468c842dc46", + "39": "05d243a1bb32436c4dae1f4e4d4236193958afc2f88bd9d587d8d65a143a4553", + "4": "118ad5644bfda0b2b192108f8ade1c27aeea9c2fc46c1478aa9ed74b989c13ce", + "40": "9feddd169852030580c080e6c8bdb3c94cb77b11e95231fa6b096de9c1993f25", + "41": "1c5d21566de30cd21d6e20d40ccd5d8c28a580a79e5d048fab010408d3f00c05", + "42": "f82fb958809c402f0b2aa705fdda31bd4a492898e83c45515f3ba6eb07cd1a36", + "43": "c261724442a58dfe8c0ed6e9d84dc85a6330f2d175feabea811d0244c93d776a", + "44": "6f38de382e45e047e174f3cd69e9d69e1ac1ad17594c3a504cf03b1b7ec0f3bb", + "45": "9db906b718877da455f31f21ae644a3729e9b49e5036273d7785f25fb287bf73", + "46": "a988ee9fe02977c2fc8fa3f1b5f8a4c3baba467432de4a471a8a6d5c08d100f6", + "47": "06ae74c9db90c48a62c1460431281f827fa6cbaf5e3ecb9bdbff9b3ac2cecb61", + "48": "ec8cdf33665f1131f9c28c617a294c98d9dd775866231486ad29812c0cac9728", + "49": "6f5df7591785bf63ebc4c186da4c543027563f7b99061649d3f8352e98302eb5", + "5": "5bbfe2c8fa835bf493ad649f084f2e17eeb32262fd128f80af633f6cc6a6aa38", + "50": "d5df2534c417dd89f782547a4d0c42e5e3eb22f56d085b6e77dd8120fcd23300", + "51": "a4350ecebabc1e2c0009853ae9ba2b2d3dfe725d326b9b1665b7065f4b465996", + "52": "cb4c568774351899f29a0d75d6dc1699d53d3a61fac8d90bcc29065d640334b1", + "53": "61dc2622386455ad11a1798156297aec1efa021eeb76a1478711e8361125fc40", + "54": "ffb1faa63471ca481a39cf9ba8d89e4e18e6777ac8425c125baa150356a38107", + "55": "86fa538e17e4ff0549251cfe3c1781d0386cfbf9b1a3f1bf6fc7e01edc9aa6dc", + "56": "1eb67bb32b57e2294ccab9edf95ec4177a0bbcb761b0f94eca574f6a1dd9d385", + "57": "4c0e20289761374384f9f6626e1be46fb637a25cd1a7f5e31b9b3182c1cbe8e3", + "58": "cd1b1b74c25c4dc0c884fe84efffd674b8e86f4f572d98b18feda7c65de067ce", + "6": "601699de7121888b804cd44d98ee0ba2aea05d0d422f1c0f4c2bf8f52b1312d5", + "7": "4698cb337db88884bc1e5dc91b0b02ebc3692e5d52f619fdbb4f8fc1d1043d69", + "8": "43562e62ab343738adb4c6c0c8d05be2ef7cabc1f8deab97b79195481ebcb6ac", + "9": "b4230e125d0b11d010d2cb4d54591d1421a3022066b0a0bf08634aa9f17f5ee4" }, "gfx950_attention_tiled_2d_fastkv_regp": { - "0": "6eeeed91e3d8db4564a81e23182fc30b78d31d65ccd1116cce5bd311716da86c", - "1": "4b577379eeb73d38a07e0fae9697bd401bdabd8e43880494006da1aed2aa2c8e", - "2": "8adc74c946292f860230227abaa7b105e238c4df4afc470b11e7fba725bb1577", - "3": "6ac78192a59bd823582807179719d756a2b11af0f37d43a02603a484a9e701ef", - "4": "1c8772e203ba3782d2b50096a40676a6bf0dbfc46537de905e918ec6887d3b8f", - "5": "d83efc28766894ba4973bca0bd2493db8181530b3460dd51fc3cde7ca00db3cf" + "0": "41bbb3d3201a871550f75502990ce9121ac47df093dc19d3efafd15c579f0c85", + "1": "e577acb27a35a97c94824c0847be635d6433de4b919c473c5a8833b4da95040e", + "2": "e9346206978098447a4de0f62101498a45d33aa903c546a1d375fdad929f9b53", + "3": "d0aba205fa7dbac5de19a7b8a2673838c84aec30ebdb17494384d39bea3a84a7", + "4": "2003dbc1b30d23af8796b89cd3d6fe5a90b7bc2d2094930de2562a36ff07073a", + "5": "4a6d8ed4d0629421615579b64cb2d3618c2d0beea77fe4d8ae8942ba617e573a", + "6": "df1670a4164f394ccc9e32aeef651b3264b156b09d8df4ed2afaf0207f75423a", + "7": "d1b3387b27962a4b3babfe31e4c65228eae437c03daae3650af70e1fd9e91e31", + "8": "85610b89daf488f7b18663dc5a1c7ce0813fc8f4c7c118f85a65661170d59133" }, "gfx950_attention_tiled_3d": { - "0": "22c3c50dcdeecfdd1706c2c3806786e493f4576336bd5960ed9c9029bb8ee919", - "1": "af5e9cd856f18112da7820b9d8a59caa361fce8ae3204658071db605101e60cd", - "10": "621a818700d6c87a14554b10f1d99367a4f3506b56b7dd4fcb35153f17f1e60d", - "11": "7f03d2d86ed48a324e9f40cf0da2a32e00d71bb1b1bcc6e563befd5014e63e39", - "12": "05b11a1897269a47d19a68662c4484557a920b1ddb4d8f8e482a3b2c78f9bf7d", - "13": "468671325ffc7a4e9d4d713fea1ba7373be4239e875b97cde89c95c8192d970a", - "14": "2429340adb1e6297cc6e3b7b8ad5f599459bf976d4f0e0973b7f4596709d9fb1", - "15": "af9726d24b47db9c198912f2d3f6ecfd2701cb3a720da6f3619d0949abd9f6bf", - "2": "f0a777ee162213e55e259419d46d1791a975efb8fb0bf272d7536cd9d9efaeca", - "3": "f92162681e4e1b1c3bd8ff591baefe4dcfa3a840ade95e46c4b3ae47d3498880", - "4": "41bf0207a6966ba320af8c251bfc51a0791e82f3782b204b3cc1e07e2c265a57", - "5": "01ebc828573b3f80c1477b713a26688d6be326c0471338acc4e1f4425bec1bd2", - "6": "7ea68a213e1fd528c99eb137b420e4a638913c1dc16271bbe5d7a0f1aad4c6a4", - "7": "f208f171bb95f08b72962196681ca36bccf1a43a1d201d92dfcb1d5df79d1d73", - "8": "4f26d93dba745e47c56d09b5e28e04a486a07bb98dcc6bf133efa042003aca8d", - "9": "8cb90ae44cf07b9dd7f14d5c5bc39c7f82a9d1786b010dbf9be5ed204a957968" + "0": "9c2eb707c5e6f816cb8a8cde2043486839be7744ba14fab9ec1c2a9040423f3a", + "1": "d2d475c2f77f51a66e7e89e3dd8ff1382b24287a8e1f7cc67916972f71a2bcf0", + "10": "5b2639db7625c46c6f3674c99405dba02cc4edddd01de074a429ebb655bca80d", + "11": "4980053626a2642194170dfda1e0100b42233c3d326581255ab0c91a51c69228", + "12": "3bae7978951e52b791ce2e2717b6b14ffbc627e64bed3772c6bf3fbfa411d29c", + "13": "87a148b6f9549fdd82fa70d06df5b8bb10d9f02cd5f7926738ffd9558c8a6e9d", + "14": "9c30ffc9aa7961ea59c6131e97fefb8b2e2a1e65bd154cc17f3e194f41eb102c", + "15": "797161b0cc40f517d50369b8dd509fe3111e7e8985e96814cc84ff2dad60fe2c", + "2": "ed7166fa1dda465ceff4f41aae8d4f83234a4c267b63377a681ece2a09169118", + "3": "8160d94c1f921b2464387efb27f2697ca3b9a812cad6d9625bc6b68917f8c897", + "4": "f3a1153ea9c1d6e847538c6deccbada4f65ba44841ff72ed5712c50581465806", + "5": "88c518ff4d4c3490ac01bf456a8e801c1ba3fa3f399ac84361084b0e0bf82e4e", + "6": "0176fe74a827d2f881405f0ce2f7175979eccfac3e82e6f7989d4ed93e9bbfeb", + "7": "fae6dd54d2a3430fe83677da1710c609b951684908c3391e892bb342a3314f28", + "8": "33d3cbde3a27b309661b5605df4acc0829d7277d10771981961373d641a5d9d4", + "9": "9181d129aac71a08956752065d4746be5dfa024168ea37b59dbbc473939b2690" }, "gfx950_block_scaled_gemm": { "0": "e5e987f27ee79ef21e95d40eb956f71f784be46d0284512d94da6a08abf0b1d1", @@ -5042,18 +5060,18 @@ "attention_strided_kv": { "0": "7ce6e33c41e2752f1863c4ef79eef519c4748cc9f1c3035cab5cab030f386d55", "1": "38fa323beea8da252380810dbbe76686f3cff17bc15aca23c32a73c11f6a70e4", - "10": "400f75d835f4bb4d37dfe0039c609deb5e172c77e5a83b87b89071ed3b0017bf", - "11": "0e441b0ae432165d00da483191b9470dcfc043c2bacd386b86fe6bef8f05266f", + "10": "59ba76e31fd186ba00c06d841b8a038c8941a8419e56b8fa46daf2ea9eb61028", + "11": "a30d1148dddca335071fee50a5cedd26bb5323693c278b14d22580d6cc5df811", "12": "0ec51a3f92220ca3dad8979ef3dbe976f89ab95e28de5a178eaee7477661b8d4", - "13": "74d14799531a8e079e822a46af3a22f1875a85a46c6c561ef26427fd44ffd748", + "13": "9bb3761cb7e928a0efc7789cda47a011260003ba29f7db4cdde89af0a0bda125", "2": "5e39245a8752ce0229c9229aba2b11fc8cb07241e93ebaafa051cae24aab4514", - "3": "393df7f6bb892e49b3c3031ff696d11addb1da00ad2088aa881ea37074c7db77", - "4": "db62e65f3aa74588eb95d2bf69935aa6331d0a57d42996bbe81c7510e146f119", - "5": "ee8248f21f5631e005750601f275daf2fd57d87f102a241c2fee76caa4567d33", + "3": "76bca1175f651389730a2856cf8bb9cf9d79998445425cde57ac37700dd1d2ef", + "4": "0ea423ae0c4cdd6d4b72ed754b1a18dab37fb3f6b9fbad492b1bcd036529c823", + "5": "29304ac08dddbce980440041340d863e3996b72fa786fecd5f739ca8c14bd150", "6": "b3add7e87f7a74a750718e45156d80bc9a5725c33a8a25fadc55b71bb600af89", "7": "94ea80d639a380a3cf0069ff8679e6f197fa1c40c9a4885413d747341964ef9c", "8": "d8a78f5237a1c466494c27808c41590153f9ae47d899a5d6ad240510e680d89f", - "9": "cb2426e5ff5ce00a3094ea4217df160031006f22cfe058568e4c959f1ecff086" + "9": "67bfeb3341425741eb76f57e12f56c47a87bbdb73b1e5641c4a4f256f2647591" }, "attention_unified": { "0": "de5cde9088dd20c8d6c5e2c7c792d6a9d4a8b42cc423a6cdf26f5779c1b12ba6", @@ -6333,88 +6351,94 @@ "9": "0780bfda4117929b13f56941e11205d93fc75afe332d5ef49b91c55d76d4bf34" }, "gfx950_attention_tiled_2d": { - "0": "8fcdbf0c36481459f4ead47cc0d3fdd8dc2bbe1aef001e292c9135b38d06ca16", - "1": "00ad5abe7f8df2a03f1d5ecf57e291fdfa1fe767359e7c25fb47a93ae603656a", - "10": "9568c9deb479a1beba8c080793235fea51ad6210acaf6ece5d331c31760187cb", - "11": "66caf6b211607e42a2a4783e68981f12e4192078532d3db83c3cb94104a47174", - "12": "22ea40707cff17e16bc1f2bbe24841d962a88c9e478012ca077f6d9be77d7f25", - "13": "bc1946d33f9945953fcc7cb7d900bce7ce1bb7bb0343e617ada53869c19988f0", - "14": "de5e5a92ba78fe6939e1a6593f1016552b698485d92fe21b09d4008dd2a9e308", - "15": "e42cda4e52e845acf4073c2f7afef76992a9076efc494082ee2a5cd79712fddf", - "16": "7cc948c9f527928a1cc3494417c6e98b07853c3859ff880fa05f8d43f1e8f60f", - "17": "df9570d3106dda1e7424b0dfc85c74ba5042adc091a5a0b53aacb23ed10cc5e3", - "18": "a38a3c482c6b4db56a7b0268033d3880d6dd09aea4915f99d991139a8272594e", - "19": "ca5cf64ba9a99184756bb575decb3cfa3bfa801af3e01279e7a8b65811b0eff1", - "2": "8e96a42d2b394d6e6cdbcac5b75476b6708922f695bcf0e607146bd071f0007e", - "20": "c1b131e3db54cfc229f4905769bb3fbf02c151f62f28c5468b2e8eb86ca338e0", - "21": "f74451f4e28b0924a3c68d330a44016184e6e037f92f97fe84de4226b6fcdf5f", - "22": "365e2e07f4821ba28cf561a7d02ff0437b7f8c58e060d812bd10c71922caf6fe", - "23": "ab0e4ada1453e9349ae9a993e141839a43ae249e18b25a888ea991878b587f87", - "24": "20e780d43a5aff908d810c1f1b362315b2f7ca74eab7c2343cdc9bc035d151c3", - "25": "6790e8837095e85027863075b9e8c741415b48097a4067b990cebdea3edf810f", - "26": "4d462351d6c4faf4d296d7ed3a11317b157db18a0355a4f266901f2499f61bff", - "27": "d2d61bbb588568d436bdc372b1fd34a3a678d7fe1ed8deabf9b9814a16217732", - "28": "ec276705e1effa1e66bee8f9e3cce4395afc7bc2b6dac220ab61b14115acbd64", - "29": "3ea70fadb80471727a3315327f3468b72e9f089717af1f9c1347cc9c1a1a401f", - "3": "bdfaabd03c3108a248a902099acdd49a25a09f1a52ddbedb86f4da0d66caefe3", - "30": "e118f95f390629c618633b4a653e91c235d627961baa65e21f905b05e588c6c4", - "31": "2d843f1f562525bbf30aa70a6e69c55fdb1330e3c8a315088a5313a09b8705cc", - "32": "cca9ae248ba2729e3bde67b99e0d78b76835151db330d416a21fc408be978a5f", - "33": "c6396f78b3fc5e485ec6d97e628b4150e57723ce80893b2e8d1a67e1a9e4f876", - "34": "9ad85ba3f10300fcdb02b46d84c0f247c8592ffeec9127c934c27aceaed8653d", - "35": "af34f34671c389b186219df77f803c6fecd6d29e5fdb58e254629dc34901c511", - "36": "2e17e230a9b469b52449428ca6e65520e059d9347a5f15209d2ac7c6fababe8c", - "37": "16e7b1cd901bc581154a9e34ba2e6faa12ce49167a8af51ffa56736234af22b1", - "38": "792b8e0883eaac9b48055f3b2cf5f31eb00dc59967600b56b07715728014c126", - "39": "2b24bc12eee9b94c6680cd1746a2d6ef84839f13968c8012be027115b2cc5dcd", - "4": "8b88be7277ebf0ef024a88a9636a09c2e071e3d3bfa15404253952673090d711", - "40": "a528d3e6aa69881d97086386a28248d3d78a0845df21f8bd5a208117b4f354a9", - "41": "091d207e652e2633552d81b057e8442f009e24a96f9e167591d114e61601cb53", - "42": "a4acf8e88e53626569831a7674c7c757d3936ce3e6e36696d66a19668fcce717", - "43": "b45482e1a40873c9c1dc30fb976fc8d7d10e22f518d9534d65810790bbaa9175", - "44": "fe57578d36b5cf75d44e43e47fa8c72fa4484e5c127d4f735be815edb2c9ec69", - "45": "e3e4603775d7cbf05e4f54b512dd56539e650a4feeca15c278735c09fb209eea", - "46": "7ffa5dd09960f5e3ad9dd3a4df1482488bcf965468ea76be09b2cc4c34fe1ee2", - "47": "c52cd0212b029786067bb53ae852af6539d80f28d9226bca02809d640db016c3", - "48": "c9994518966952418ddeec5f176913a01d8b1e19cfecceb4ad3c4328cf0a097e", - "49": "d3d9c46f038c7be692c1faeb6426312c9e11775dcbf5591e65e3008fe7f161c6", - "5": "e0c0cd11a3d48dbf99b9e56f6ab71eab219dd2daac465dd9ab3a119eeb60a062", - "50": "13283007200c7f037b39a3f8277d0cd3ff67d7812876c37cc0b28e73c548cbb1", - "51": "f5277e42aa0d7544d6bdd56c90f9821e393f7d071d7c7a52f6f8f36f8b6dab63", - "52": "d778e3582bbf4921da262552e477a9cbbe287882b9d1cd9dba89f8358597a785", - "53": "c8b03a311dd9da61aee18b40645cba0e0bcc7b0649659a6ed20c327586ba43e8", - "54": "b66dad2f8db97c367c43f8b628347b265f33a98be11ab353d19917277a87fcc6", - "55": "eb41c11c96b2c6e28734961cc896b0c003df7e05722786ce9db0540edf4bd997", - "6": "cb2c0b480840f421ebb818e886e506b333f18969d09d722e30e2168d9910cbeb", - "7": "79d7633c087a7eb374095998f7c72ecd22f707930db6343f54bc64e866f722ed", - "8": "1b7d6305a39a83570adb2c90ba34ea90c090e5c7f9492fd1799886dfd4c6649c", - "9": "11eac5a6a6b571b6bdc6e4eab286782607a2358ee8e26ba7d8a10f35464d109a" + "0": "d11a068c8d19d89e053f8370cb5abd27fe3736ee6b9bd0a5db97103cbe6ad7b1", + "1": "9a839106808b01ce48c8d6ab6e6e1d988a5774674731a4419096a45ecdf2063a", + "10": "9a98947bb71b46d9e4ba29599e9f32592d16b70a546c1b973a6c16ff7ed2fcfe", + "11": "4cd41914106c24f9082decab56baeaf3e00ac91a9e01a31a49e8b45753b14267", + "12": "564430eb52a97e5af212c8cd531bc05d687a1a58be84e1149d50d47d39978ab6", + "13": "881342b2cf683852b0992ccc3e67b05398bd39d365fcad0d9357fd8dcc278584", + "14": "ecd11f20a00cdb6f1b9c7b94b14837fca8275b9f54f69efb975b67070103a13a", + "15": "fa2be77e7cc0f6602e9e89d4015a2dccbef5a905def81187de9d15bf0c87f070", + "16": "57213541b7420076efb0bb778ee4da186ecd5b0bb070885c3b42991ea2d1c4ae", + "17": "f2be22e871ca785b1f50645e1d9ea8c9fc7e9d370e0b03152611e3bd81ef71bf", + "18": "db8cf615fc261db8c7c32115e6b00bad7c51e22978f29519c9ba4a7eed65ae52", + "19": "31d8fc126b2a9c9f4126fcd1c3670459441cca3ec0dbb077460840b3dec5f47e", + "2": "47ce61dc0c5aecfbd09ee22a44eb3b3a84d46bfbc8c91ab49dce170175370819", + "20": "0ad1bb3a25a1af6a9263d17dafe0f0b4e52d435db345a0f687c166ef7bc7703f", + "21": "82e88f302a107a70f85fd2d35f2e1dcac26697d96253b73f1ae3bbfc45b631ae", + "22": "7b4e1e6f632c0f67a544ae174b66c77b5a9a88c99563d5d79f1f67f65de3e564", + "23": "38ae63d87a307be1bd99df037fab1a80f4d15c611652851860448433308ccc82", + "24": "cfc0d62508e01e6fa47f5fd3012c3e3a19dd9ed6adeafc931a9fe4390df106a5", + "25": "4cd41914106c24f9082decab56baeaf3e00ac91a9e01a31a49e8b45753b14267", + "26": "4cd41914106c24f9082decab56baeaf3e00ac91a9e01a31a49e8b45753b14267", + "27": "38030182311b81bcbcd7fb4361c0a103e63778da58d9da2f63f81bff76797655", + "28": "a98e68354de25d573b015e1e879cc4543f428db34909c3bb8f809f51bd3eb708", + "29": "526722237e5b4fd92cc3ac85019176942703b7073e42cf7b54628ded6fdcf0f7", + "3": "70c9b7c0055d3b9725ba6009d6fd8339aae2b2dfe9919cf5a325a97920fd211b", + "30": "bd7fe6ea2867cef249abb8134363be56afb241eb1ead73e84adb0528128ec51d", + "31": "bab51c59b4bc6e2086fc6f7c4cbacf6d3d50d029777a636e70b698a5e5b0fb3d", + "32": "f4ddea46c32e6b28d816b3cc2bdea4936e9f8402f194c969c03118d9630f87ec", + "33": "d6b278ec63bbe590323fb19a7cef974e7b7a63fbf84b8f6cd4c3938ab2a59d0f", + "34": "9db698a3e555dc0627a4a65b732551703cfc077f7ad5ea4164a8a16b8c013e37", + "35": "2b3c32b1e196a360fbb24f3f673d06d2894ed8c05dfa3abe5980b06004d06fab", + "36": "c37e06816538a2022bf1d225a9afa3f83343b51619ce55b38277eb34548f5605", + "37": "81817a5685d7d22f869d93f2f8d280f92c1900310c5ab7a28a92b27fcdf3a503", + "38": "7ab97c50302239432f777c8bfa8c99ffd2748e5a9c4cab17de728f65c3dc5017", + "39": "f1ed8d5cdba8dfe24e043ce57adb01baf17451c706d3c5eb35d5b3c596617011", + "4": "07c34a6af29757e08f08febbca6678f10a8747495a60e1e4fcca606fe3e921d6", + "40": "72b1acd763265d3c074b2096ccd547b13c99abb70f0ef083262e6c7811a6b867", + "41": "35c2d919ea7c37be4dbf11a79499540f9abe89f14d649af210cdf4ceabaad7a8", + "42": "8c88eed76fb9794f90cd17e37e211fb45c1f1ad46c8583e50221279f77aecd7b", + "43": "d5f71c4f01e4d06ecf6462019ce2d3f34c32d39a76f92865a62324ecd85eee03", + "44": "4a4f02c1a950e97a3d8599ba26aaa49dcd80eb5609cc82755dc918742d4a4b40", + "45": "639c656428c2bbeef15291da74351ca3bf73c5e3e403268d4f69918f8b9752b3", + "46": "bae0bc4e1725c32777e5c41f2b339958e970a993b867529f18b80e6a28fbbb87", + "47": "7f19a9c3f716e2b976f2cb9da062574d0dac4e849cd7e5ef7bf535c3f35a290a", + "48": "8cc828326bfeca2cf2ebf53e10346953d6c190dab0c7fe1a39cd8cd5edde906d", + "49": "a697c9aa0fca2ff754aedf973164234573188fe48e32dce3bd93a89137bb2eee", + "5": "4cd41914106c24f9082decab56baeaf3e00ac91a9e01a31a49e8b45753b14267", + "50": "7eb5f81d6983052f02d0bdeb83b3ed20bff952ee5ebf293b4b1d2d38f08a841e", + "51": "0431cb8a30c0c3fdef8418ad326810dfcb0b3902309a0b8270b231a4c60bbddb", + "52": "4a765e55deb0ae5c5420be117e7015f48a3eb60d19e5d445110ccafae28ce304", + "53": "f06109220185e28a131a498235f5dbff4a278ee259cb115b8459d06971ed36b4", + "54": "e991ccdd7d485073fa24d51b8e0506e11e9d6514dafde686e0e1216e450be56c", + "55": "8fd094b334a261b8eeb0e8ca4a7c35b354b859279cb85164569cf34e3694e8ab", + "56": "efc6a08c973764bc16b5f506145043755aa9a8548ca30bcd93b25bb3ffab5a7d", + "57": "6b9774131c3fad814abf6cd04e7f2952b3bb9480cee0ca363689dfd23db044a3", + "58": "f9f7151e1665fa216cfe0c4fdbfeea4830dc72b126230449cce96c5cd6081934", + "6": "b21885e2357685a1dd19687462448ee0ade77735a90166a7df5d72804517780a", + "7": "ac76698f1ae9e2a1fc56c57202af960d6e3702d27152df8c7a1b62a933657e31", + "8": "51278c135bdf9f150386d7f34007caed2cb2798cbfbec858b86f01a2fd968776", + "9": "515d78c1a99eca9f1ab90d4e8e90a339b1cdb66b5b945af1c3c3e4a40bd25577" }, "gfx950_attention_tiled_2d_fastkv_regp": { - "0": "c89843f84db4565f7bd93364c173ee310386fa248d44a62400a49245affbaea1", - "1": "724e561a527e529a689d45b602ffecb891e4f849f623ae81f1a071edbf8cc3e3", - "2": "d5b38b492bf3b7fe961a02ae29d9593cf8d4c27ce8c17293d7742ac25d01668f", - "3": "7b11b81762944f9cb12718422ed46b70f7999fb057d0eb4e65bb375a89deeaaf", - "4": "3602d357be0a34988c98c27b58f7add55aa5d46e5b32479a5ae544282febabd6", - "5": "5ae5b81340ee7836e4e3950e535f024e6eb80e4f7b6ea6f6fd60ef1a6af0093d" + "0": "a6eaeb9879aa5b723d141578e9a76c62f213894cf1b79584704057cca967e963", + "1": "eca0eb2d5db59e3af816ba32baaaf5f6f6b0ebe3ad530e7516ff50b4981f971c", + "2": "1cb99b70e2c8d01dcab577e907fecc38ccb97bf7a3907a0a8722b4dfba9b4bc7", + "3": "0f1ec29b41d8eb406b3fd4dbaf086930de65f86bac82d3bb182d941e9f829882", + "4": "d72e6d8604f69997b0af4478aba4deed68af36d99755aadbf702d2192dff4968", + "5": "f368ec1319f8e2f3aaeca7205d2bd1d7902e9e13edb4bafaf6993838f2d5d04a", + "6": "2a3a26aef4692ab99cab30d92b4837dd9fcc634e4a6c7aa9db9fc57c62677e0d", + "7": "06f4cd1bf03d1c01b66a1e70dec7bf7fb0e850e5d84c2b5a42752d386a526a0d", + "8": "c4b62e99e3fa746d5f89d010e46022f5d61dfded176bed8f58ae83e09b4e6679" }, "gfx950_attention_tiled_3d": { - "0": "2d15cfdb485bc112c4364d6b000586b3da7cfa0403b170a7f4c4662a3003ab84", - "1": "0476ad71ba35fe47953f586a272750b88fba23a5f66aecabcec34062a1a5a7b7", - "10": "b0d08ee5431786ce546dce7a44e221a76b0ac9cdcb5afa2c8313fb2c2e0e1b44", - "11": "4870f2163281ea108cb8ae6383de76d6fb427ec80e558f6b6c076a607da18d6f", - "12": "9680fd6292de1079d9387a0e9947acbf0831cb64de2858ce4f19f879e5379d8e", - "13": "d34f0c7de3fd5d0467d139803dfe5bd17dad4c26628bd868f4245c05e292d0f4", - "14": "7a984f9f07137c77ea2f0704ad7a65a350e9fc96ab759497cba626b3fb766f52", - "15": "942c3c72cc504c539d4abc6201db2bab548871427f35a335258eee0614c1491e", - "2": "1680787ea992fe51bc10f0d663dce835be3c0dc742b4a36023964a9c1e671ebf", - "3": "d543f94c41159c5af6822c0656497b16bcada09f475aa2205c2ed937becaa526", - "4": "e3c11f0bea0406fd826b858b25c0273db7685e0fc1d04cc8f90c4789fa474a9f", - "5": "d43dcd71e7156cfd4f15bec3022fd26ecbe83352068d1313fc824a1c17475de5", - "6": "f44fa8e200198874f65206294b327a10f1fcc62ec4b6f36e78efd52654299b44", - "7": "efaa5a76dca2b7fafcdd6a775d13559f38cbcabfea66b05e2cb94e86da519a7a", - "8": "2c72ff6b89e1d4a1883e7d3e82ca01a7d4de9dd98d54c1142bd40de8012fcf1a", - "9": "62a5bae4eef3aff5335a242803f79ef319a5e1a6ff80e984709c0b51e241be11" + "0": "185fe233b422855899789edbb842c03f215fad00c1695917e21cad755a16de46", + "1": "a13ef2bebf3c8d27861c800653beb21b0efdec7e1865146ad540b5e4333521cb", + "10": "bc305699e45ba27ddc4d9c747ef9f3557dcc34346e558e8c07405467ef11535a", + "11": "9753247ce408ad7f445143ca04a20c3ea3a6edd86a823386109750a3b38f44f1", + "12": "06e34a54b23416e64f2b9b995f2cfceb7e18e6bac98bfa5fa6245df6ee8b70ed", + "13": "4c9f99009b4997667cd8f5d378869782f35ef55664924a42094859a696df3138", + "14": "d657f73cff7d7c5ca2f8ff5c6143e5471564608be25a8e55bb0a39c7638d67fb", + "15": "23ecb293b7e44b601d6db8bc37f89de539e9089940561c836d8b832c81bb0d39", + "2": "fc583225c3e1a96067496e8ab7afa37fa6c6eb993b7bbc3d2f237fd57607f063", + "3": "7104dd650ee025a86f1830a1805541c315425bbac5c24b25c4e115874f50feff", + "4": "3bb308c0df37f7f375ebb7cfa77188a3f847cdd40627db54591a3a64f51ea694", + "5": "5a6f8e3e40767268e22d73f6dbe7a1db0d682130f795e49bc46fc8e3e18836eb", + "6": "dbc33c5d0880dd8a476682f38346ab2305dfc211c9511f0f252a8e978419cc6b", + "7": "f95f155e633137d51def2007f23ed9b0f5123e0521f43fd712d4ce02b3d0850a", + "8": "048a04ffe10fda0e1a274e78ec9ac2a06c216e6415d38e9e309478ae57d6d7ae", + "9": "51c6643e3d5d31cf6ae87b54e8c40b66ac2fd0d6f6651a6117f09505ad0ae60f" }, "gfx950_block_scaled_gemm": { "0": "a550fa8769391a2bf246ee73ad70769826d1f902db0a49bb60a80e7c0efda41e", From a2f4d565e3dd72cbef87352bac68f7439aefac64 Mon Sep 17 00:00:00 2001 From: dlangbe Date: Wed, 4 Nov 2026 13:31:30 -0500 Subject: [PATCH 3/8] test(hip-kernel-provider): count the 2D unified attention scalar pairs with runtime kernargs The gfx950 unified 2D kernel now takes eleven scalars, so the like-slot exchange test sees C(11, 2) scalar pairs instead of C(8, 2). Co-Authored-By: Claude --- .../packs/TestGfx950AttentionUnified2DSignature.cpp | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/dnn-providers/hip-kernel-provider/src/tests/engines/kernel_ingestor_engine/packs/TestGfx950AttentionUnified2DSignature.cpp b/dnn-providers/hip-kernel-provider/src/tests/engines/kernel_ingestor_engine/packs/TestGfx950AttentionUnified2DSignature.cpp index ae0eeeefe6b6..aa59ee202dbc 100644 --- a/dnn-providers/hip-kernel-provider/src/tests/engines/kernel_ingestor_engine/packs/TestGfx950AttentionUnified2DSignature.cpp +++ b/dnn-providers/hip-kernel-provider/src/tests/engines/kernel_ingestor_engine/packs/TestGfx950AttentionUnified2DSignature.cpp @@ -222,10 +222,10 @@ TEST(TestGfx950AttentionUnified2DSignature, RefusesEachExchangeOfTwoLikeSlots) ++exchanged; } } - // Ten pointers and eight scalars, every one of each group the same kind and size: - // C(10, 2) + C(8, 2) pairs. A signature that grew or lost a slot changes this count. + // Ten pointers and eleven scalars, every one of each group the same kind and size: + // C(10, 2) + C(11, 2) pairs. A signature that grew or lost a slot changes this count. constexpr std::size_t POINTER_PAIRS = 45; - constexpr std::size_t SCALAR_PAIRS = 28; + constexpr std::size_t SCALAR_PAIRS = 55; EXPECT_EQ(exchanged, POINTER_PAIRS + SCALAR_PAIRS); } From ca119e0be1617e2d0bd814afcea1584cddc61c74 Mon Sep 17 00:00:00 2001 From: AnaghaRaoAMD Date: Wed, 4 Nov 2026 17:24:57 -0500 Subject: [PATCH 4/8] fix(rocke): address #15024 review on tiled attention runtime shapes - Head counts: one check_head_counts rule (Python _fmha_common, C++ attention_unified_selectors) runs first in every gfx942/gfx950 tiled validator; spec.num_queries_per_kv raises ValueError instead of dividing by zero. Both 2D C++ _new entries init the builder before naming, so the lower no longer frees an uninitialised builder on 0 KV heads. - gfx950 2D K/V loaders: page-id and copy-offset tables are sized from the spec on the arena instead of fixed 64-entry stack/ctx arrays. - fastKV register-P: the kwargs predicate, C++ predicate and C++ builder refuse non-gfx950 targets with the Python text. - Descriptors: TensorDescriptor gets _term/_term_i64 stride hooks so DynamicTensorDescriptor stops copying the four offset/move loops; C++ _naive_dynamic wraps _mixed, which returns the base view and sets the builder error on failure. New dynamic_helpers parity case (config 10). - Runtime-shape numerics: one torch-free, arch-parametrized module (DeviceArray + _launch_3d_pipeline) replaces the two torch-only files. - gfx950 2D hipDNN pack: drop the 511-sequence cap left over from the baked binary search, and its dead capacity helpers. Byte identity GREEN at llvm20/22/23; no golden moved. Co-Authored-By: Claude --- .../library/kernels/common/_fmha_common.py | 47 +- .../kernels/gfx942/attention_tiled_2d.py | 6 +- .../kernels/gfx942/attention_tiled_3d.py | 6 +- .../kernels/gfx950/attention_tiled_2d.py | 6 +- .../gfx950/attention_tiled_2d_fastkv_regp.py | 19 +- .../kernels/gfx950/attention_tiled_3d.py | 19 +- .../test_tiled_support_spec_wrappers.py | 60 +++ ..._attention_gfx942_runtime_shape_numeric.py | 299 ----------- ..._attention_gfx950_runtime_shape_numeric.py | 329 ------------- .../test_attention_spec_predicate_parity.py | 2 +- .../test_attention_tiled_cpp_entry_guards.py | 75 +++ ...t_attention_tiled_runtime_shape_numeric.py | 464 ++++++++++++++++++ .../rocke/platform/CMakeLists.txt | 8 +- .../rocke/platform/cpp/helpers/transforms.cpp | 68 +-- ...ances.common.attention_unified_selectors.h | 10 + .../rocke/helper_rocke.helpers.transforms.h | 9 +- ...tance_gfx950_attention_tiled_2d_internal.h | 5 +- .../common/attention_unified_selectors.cpp | 17 + ...ttention_tiled_2d_kv_descriptors_loads.cpp | 5 +- .../attention_tiled_2d_public_entry_glue.cpp | 16 +- .../attention_tiled_2d_q_load_gather.cpp | 5 +- ...ed_3d_gfx942_attn_tiled_3d_public_glue.cpp | 5 + ...n_tiled_3d_gfx942_attn_tiled_3d_reduce.cpp | 15 +- ..._3d_gfx942_attn_tiled_3d_segment_setup.cpp | 20 +- ...attention_tiled_2d_fastkv_regp_helpers.cpp | 28 ++ ...ttention_tiled_2d_kv_descriptors_loads.cpp | 47 +- .../attention_tiled_2d_public_entry_glue.cpp | 29 +- .../attention_tiled_2d_q_load_gather.cpp | 5 +- .../gfx950/attention_tiled_3d_public_glue.cpp | 5 + .../gfx950/attention_tiled_3d_reduce.cpp | 15 +- .../attention_tiled_3d_segment_prologue.cpp | 18 +- .../python/rocke/helpers/transforms.py | 175 ++----- .../tests/core/test_transforms_dynamic.py | 2 +- .../instances/parity/dynamic_helpers_emit.c | 32 ++ .../instances/parity/dynamic_helpers_emit.py | 26 +- .../gfx950_attention_unified_2d/README.md | 5 +- .../Gfx950AttentionUnified2DGeometry.hpp | 22 +- .../packs/Gfx950AttentionUnified2DNative.cpp | 41 +- .../packs/UnifiedAttentionGeometry.hpp | 41 +- ...egrationGpuGfx950AttentionUnifiedKnobs.cpp | 18 +- .../src/tests/CMakeLists.txt | 2 +- .../Gfx950AttentionUnified2DTestGraphs.hpp | 3 +- ...stGfx950AttentionUnified2DCatalogTiles.cpp | 35 +- .../TestGfx950AttentionUnified2DDispatch.cpp | 41 +- .../TestGfx950AttentionUnified2DGeometry.cpp | 49 +- .../TestGfx950AttentionUnified2DMatchers.cpp | 42 +- .../packs/TestUnifiedAttentionGeometry.cpp | 1 - 47 files changed, 1039 insertions(+), 1158 deletions(-) delete mode 100644 dnn-providers/hip-kernel-provider/rocke/library/tests/test_attention_gfx942_runtime_shape_numeric.py delete mode 100644 dnn-providers/hip-kernel-provider/rocke/library/tests/test_attention_gfx950_runtime_shape_numeric.py create mode 100644 dnn-providers/hip-kernel-provider/rocke/library/tests/test_attention_tiled_cpp_entry_guards.py create mode 100644 dnn-providers/hip-kernel-provider/rocke/library/tests/test_attention_tiled_runtime_shape_numeric.py diff --git a/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/_fmha_common.py b/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/_fmha_common.py index 4dcf33643bf6..70fa500cab53 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/_fmha_common.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/_fmha_common.py @@ -53,6 +53,33 @@ FmhaMaskMode = Literal["none", "causal", "sliding_window", "alibi", "custom"] +def check_head_counts(num_query_heads: int, num_kv_heads: int) -> Tuple[bool, str]: + """Admit a positive head pair whose query heads split evenly over the KV heads. + + The one head-count rule every attention spec validator runs first: the + kernels index heads by the GQA ratio ``num_query_heads // num_kv_heads``. + """ + if num_query_heads <= 0 or num_kv_heads <= 0: + return False, ( + f"num_query_heads / num_kv_heads must be > 0 " + f"(got {num_query_heads}, {num_kv_heads})" + ) + if num_query_heads % num_kv_heads != 0: + return False, ( + f"num_query_heads ({num_query_heads}) must be divisible by " + f"num_kv_heads ({num_kv_heads}) for GQA / MQA" + ) + return True, "" + + +def gqa_ratio(num_query_heads: int, num_kv_heads: int) -> int: + """``num_query_heads // num_kv_heads``; ``ValueError`` if :func:`check_head_counts` rejects the pair.""" + ok, reason = check_head_counts(num_query_heads, num_kv_heads) + if not ok: + raise ValueError(reason) + return num_query_heads // num_kv_heads + + @dataclass(frozen=True) class FmhaShape: """Shape bundle used by every FMHA flavour.""" @@ -65,12 +92,7 @@ class FmhaShape: @property def num_queries_per_kv(self) -> int: - if self.num_query_heads % self.num_kv_heads: - raise ValueError( - f"num_query_heads ({self.num_query_heads}) must be divisible by " - f"num_kv_heads ({self.num_kv_heads})" - ) - return self.num_query_heads // self.num_kv_heads + return gqa_ratio(self.num_query_heads, self.num_kv_heads) @dataclass(frozen=True) @@ -123,16 +145,9 @@ def validate_common_spec(spec: FmhaCommonSpec) -> Tuple[bool, str]: f"head_size {s.head_size} not in the supported set " "{32, 64, 128, 192, 256} (CK Tile's standard FMHA shapes)" ) - if s.num_query_heads <= 0 or s.num_kv_heads <= 0: - return False, ( - f"num_query_heads / num_kv_heads must be > 0 " - f"(got {s.num_query_heads}, {s.num_kv_heads})" - ) - if s.num_query_heads % s.num_kv_heads != 0: - return False, ( - f"num_query_heads ({s.num_query_heads}) must be divisible by " - f"num_kv_heads ({s.num_kv_heads}) for GQA / MQA" - ) + ok, reason = check_head_counts(s.num_query_heads, s.num_kv_heads) + if not ok: + return False, reason if s.block_size_q not in (16, 32, 64, 128): return False, f"block_size_q {s.block_size_q} not in {{16, 32, 64, 128}}" if s.block_size_k not in (16, 32, 64, 128, 256): diff --git a/dnn-providers/hip-kernel-provider/rocke/library/kernels/gfx942/attention_tiled_2d.py b/dnn-providers/hip-kernel-provider/rocke/library/kernels/gfx942/attention_tiled_2d.py index acb635fe2df4..2a56e4368862 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/kernels/gfx942/attention_tiled_2d.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/kernels/gfx942/attention_tiled_2d.py @@ -38,6 +38,7 @@ import os from typing import ClassVar, Optional, Tuple +from kernels.common._fmha_common import check_head_counts, gqa_ratio from kernels.common._lse_store import lse_row_has_mass, lse_store from rocke.core.ir import ( @@ -946,7 +947,7 @@ def __post_init__(self): @property def num_queries_per_kv(self) -> int: - return self.num_query_heads // self.num_kv_heads + return gqa_ratio(self.num_query_heads, self.num_kv_heads) @property def block_m(self) -> int: @@ -6781,6 +6782,9 @@ def supports_gfx942_4warp_gqa_spec( def _supports_tiled_2d_spec( spec: UnifiedAttention2DTiledSpec, *, arch: str, use_d256_fast: bool ) -> Tuple[bool, str]: + ok, reason = check_head_counts(spec.num_query_heads, spec.num_kv_heads) + if not ok: + return False, reason use_fp8 = spec.kv_storage_dtype is not None return supports_tiled_2d( head_size=spec.head_size, diff --git a/dnn-providers/hip-kernel-provider/rocke/library/kernels/gfx942/attention_tiled_3d.py b/dnn-providers/hip-kernel-provider/rocke/library/kernels/gfx942/attention_tiled_3d.py index aeac46b8354d..9a905742f54d 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/kernels/gfx942/attention_tiled_3d.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/kernels/gfx942/attention_tiled_3d.py @@ -34,6 +34,7 @@ from dataclasses import dataclass, field from typing import ClassVar, Optional, Tuple +from kernels.common._fmha_common import check_head_counts, gqa_ratio from kernels.common._attention_strided_kv import ( declare_strided_kv_params, strided_kv_offset, @@ -235,7 +236,7 @@ def __post_init__(self): @property def num_queries_per_kv(self) -> int: - return self.num_query_heads // self.num_kv_heads + return gqa_ratio(self.num_query_heads, self.num_kv_heads) @property def block_m(self) -> int: @@ -1540,6 +1541,9 @@ def supports_tiled_3d_spec( This wrapper reads each of them off the spec and delegates, making the 3D unified builder validatable at pack time. """ + ok, reason = check_head_counts(spec.num_query_heads, spec.num_kv_heads) + if not ok: + return False, reason use_fp8 = spec.kv_storage_dtype is not None return supports_tiled_3d( head_size=spec.head_size, diff --git a/dnn-providers/hip-kernel-provider/rocke/library/kernels/gfx950/attention_tiled_2d.py b/dnn-providers/hip-kernel-provider/rocke/library/kernels/gfx950/attention_tiled_2d.py index 333f971ddb77..5333b0bc5784 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/kernels/gfx950/attention_tiled_2d.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/kernels/gfx950/attention_tiled_2d.py @@ -42,6 +42,7 @@ from dataclasses import dataclass, field from typing import ClassVar, Optional, Tuple +from kernels.common._fmha_common import check_head_counts, gqa_ratio from kernels.common._attention_additive_bias import ( add_additive_bias, declare_additive_bias_params, @@ -876,7 +877,7 @@ def __post_init__(self): @property def num_queries_per_kv(self) -> int: - return self.num_query_heads // self.num_kv_heads + return gqa_ratio(self.num_query_heads, self.num_kv_heads) @property def block_m(self) -> int: @@ -5045,6 +5046,9 @@ def supports_tiled_2d_spec( ``use_conflict_free_v_store``, ``use_k_sliced_ring`` and ``use_d256_fast`` are not gfx950 spec fields. """ + ok, reason = check_head_counts(spec.num_query_heads, spec.num_kv_heads) + if not ok: + return False, reason use_fp8 = spec.kv_storage_dtype is not None return supports_tiled_2d( head_size=spec.head_size, diff --git a/dnn-providers/hip-kernel-provider/rocke/library/kernels/gfx950/attention_tiled_2d_fastkv_regp.py b/dnn-providers/hip-kernel-provider/rocke/library/kernels/gfx950/attention_tiled_2d_fastkv_regp.py index 0323ea05d48e..4a1b73eb18e5 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/kernels/gfx950/attention_tiled_2d_fastkv_regp.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/kernels/gfx950/attention_tiled_2d_fastkv_regp.py @@ -18,6 +18,7 @@ from dataclasses import replace from typing import Any +from kernels.common._fmha_common import check_head_counts from rocke.core.ir import KernelDef from .attention_tiled_2d import ( UnifiedAttention2DTiledSpec, @@ -94,6 +95,10 @@ def make_fastkv_register_p_spec( ) +# The body uses gfx950-only wide MFMA atoms and transposed LDS reads. +_NOT_GFX950 = "fastKV register-P is a gfx950 kernel; {arch} is not supported" + + def supports_fastkv_register_p_2d( *, head_size: int, @@ -109,8 +114,13 @@ def supports_fastkv_register_p_2d( tile_size: int | None = None, arch: str = "gfx950", ) -> tuple[bool, str]: - """Check support for the experimental fastKV + register-P kernel.""" + """Check support for the experimental fastKV + register-P kernel. + ``supports_tiled_2d`` also admits gfx942 for its narrow-atom variant; this + kernel's body is gfx950-only, so every other target is refused first. + """ + if arch != "gfx950": + return False, _NOT_GFX950.format(arch=arch) ok, reason = supports_tiled_2d( head_size=head_size, block_size=block_size, @@ -141,10 +151,6 @@ def supports_fastkv_register_p_2d( return True, "supported" -# The body uses gfx950-only wide MFMA atoms and transposed LDS reads. -_NOT_GFX950 = "fastKV register-P is a gfx950 kernel; {arch} is not supported" - - def _fastkv_register_p_spec_error(spec: UnifiedAttention2DTiledSpec) -> str | None: """Why the fastKV register-P builder refuses ``spec``, or None if it builds.""" if not spec.use_fast_paged_kv_desc: @@ -167,6 +173,9 @@ def supports_fastkv_register_p_2d_spec( refusing what :func:`build_unified_attention_2d_fastkv_register_p` raises on, so the builder is validatable at pack time. """ + ok, reason = check_head_counts(spec.num_query_heads, spec.num_kv_heads) + if not ok: + return False, reason if arch != "gfx950": return False, _NOT_GFX950.format(arch=arch) error = _fastkv_register_p_spec_error(spec) diff --git a/dnn-providers/hip-kernel-provider/rocke/library/kernels/gfx950/attention_tiled_3d.py b/dnn-providers/hip-kernel-provider/rocke/library/kernels/gfx950/attention_tiled_3d.py index 67b96cb0969f..7552e6ba8437 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/kernels/gfx950/attention_tiled_3d.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/kernels/gfx950/attention_tiled_3d.py @@ -30,6 +30,7 @@ from dataclasses import dataclass, field from typing import ClassVar, Optional, Tuple +from kernels.common._fmha_common import check_head_counts, gqa_ratio from kernels.common._attention_additive_bias import ( add_additive_bias, declare_additive_bias_params, @@ -315,7 +316,7 @@ def __post_init__(self): @property def num_queries_per_kv(self) -> int: - return self.num_query_heads // self.num_kv_heads + return gqa_ratio(self.num_query_heads, self.num_kv_heads) @property def block_m(self) -> int: @@ -1725,6 +1726,9 @@ def supports_tiled_3d_spec( This wrapper reads each of them off the spec and delegates, making the 3D unified builder validatable at pack time. """ + ok, reason = check_head_counts(spec.num_query_heads, spec.num_kv_heads) + if not ok: + return False, reason use_fp8 = spec.kv_storage_dtype is not None return supports_tiled_3d( head_size=spec.head_size, @@ -1785,16 +1789,9 @@ def supports_tiled_reduce_spec( "tiled reduce kernel needs head_size divisible by " f"{_REDUCE_THREADS} (got {spec.head_size})", ) - if ( - spec.num_kv_heads < 1 - or spec.num_query_heads < 1 - or spec.num_query_heads % spec.num_kv_heads != 0 - ): - return ( - False, - "tiled reduce kernel needs a positive num_query_heads divisible by " - f"num_kv_heads (got {spec.num_query_heads} / {spec.num_kv_heads})", - ) + ok, reason = check_head_counts(spec.num_query_heads, spec.num_kv_heads) + if not ok: + return False, f"tiled reduce kernel: {reason}" # ``arch`` passed validate_tiled_attention_arch, so it has a catalog row. target = ArchTarget.from_gfx(arch) lds_bytes = split_kv_reduce_lds_bytes(spec.num_segments) diff --git a/dnn-providers/hip-kernel-provider/rocke/library/tests/dispatch/attention/test_tiled_support_spec_wrappers.py b/dnn-providers/hip-kernel-provider/rocke/library/tests/dispatch/attention/test_tiled_support_spec_wrappers.py index 2f950bb88f52..0d628eea54fb 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/tests/dispatch/attention/test_tiled_support_spec_wrappers.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/tests/dispatch/attention/test_tiled_support_spec_wrappers.py @@ -10,6 +10,7 @@ import kernels.gfx942.attention_tiled_2d as gfx942_2d import kernels.gfx942.attention_tiled_3d as gfx942_3d import kernels.gfx950.attention_tiled_2d as gfx950_2d +import kernels.gfx950.attention_tiled_2d_fastkv_regp as fastkv import kernels.gfx950.attention_tiled_3d as gfx950_3d _PROBLEM = dict( @@ -89,6 +90,65 @@ def test_every_wrapper_judges_head_size(self): self.assertFalse(accepted) self.assertIn("head_size", reason) + def test_every_wrapper_refuses_bad_head_counts_first(self): + """The name carries the GQA ratio, so a zero or non-dividing head pair is + refused by the wrapper (not a ZeroDivisionError) and by kernel_name().""" + seg = {"num_segments": 16} + cases = ( + (gfx942_2d, "2d", {}), + (gfx950_2d, "2d", {}), + (gfx942_3d, "3d", seg), + (gfx950_3d, "3d", seg), + ) + heads = ( + ((8, 0), "must be > 0"), + ((0, 1), "must be > 0"), + ((8, 3), "divisible"), + ) + for module, dim, extra in cases: + spec_cls = getattr(module, f"UnifiedAttention{dim.upper()}TiledSpec") + wrapper = getattr(module, f"supports_tiled_{dim}_spec") + for (hq, hkv), reason in heads: + with self.subTest(module=module.__name__, hq=hq, hkv=hkv): + problem = {**_PROBLEM, "num_query_heads": hq, "num_kv_heads": hkv} + spec = spec_cls(head_size=128, **problem, **extra) + accepted, why = wrapper(spec) + self.assertFalse(accepted) + self.assertIn(reason, why) + with self.assertRaisesRegex(ValueError, reason): + spec.kernel_name() + + def test_fastkv_register_p_refuses_bad_head_counts(self): + spec = gfx950_2d.UnifiedAttention2DTiledSpec( + head_size=64, **{**_PROBLEM, "num_query_heads": 64, "num_kv_heads": 0} + ) + accepted, why = fastkv.supports_fastkv_register_p_2d_spec(spec) + self.assertFalse(accepted) + self.assertIn("must be > 0", why) + + def test_fastkv_register_p_kwargs_predicate_refuses_gfx942(self): + """supports_tiled_2d admits gfx942, but the fastKV register-P body is + gfx950-only; the kwargs predicate must agree with the spec wrapper.""" + kwargs = { + "head_size": 64, + "block_size": 32, + "dtype": "bf16", + "num_queries_per_kv": 8, + "use_alibi": False, + "use_qq_bias": False, + "use_fp8": False, + "q_dtype": "bf16", + "num_query_heads": 64, + "num_kv_heads": 8, + } + self.assertEqual( + fastkv.supports_fastkv_register_p_2d(**kwargs, arch="gfx950"), + (True, "supported"), + ) + accepted, why = fastkv.supports_fastkv_register_p_2d(**kwargs, arch="gfx942") + self.assertFalse(accepted) + self.assertIn("gfx950 kernel", why) + if __name__ == "__main__": unittest.main() diff --git a/dnn-providers/hip-kernel-provider/rocke/library/tests/test_attention_gfx942_runtime_shape_numeric.py b/dnn-providers/hip-kernel-provider/rocke/library/tests/test_attention_gfx942_runtime_shape_numeric.py deleted file mode 100644 index 275464465829..000000000000 --- a/dnn-providers/hip-kernel-provider/rocke/library/tests/test_attention_gfx942_runtime_shape_numeric.py +++ /dev/null @@ -1,299 +0,0 @@ -# Copyright (c) Advanced Micro Devices, Inc., or its affiliates. -# SPDX-License-Identifier: MIT - -"""One gfx942 tiled binary serves every head pair, window size and batch. - -The gfx942 tiled 2D, 4-warp GQA and 3D kernels read the head counts and the -sliding window as kernargs and derive the cu_q search trip count from -``num_seqs``, so shapes that differ only in those share one cache key and one -compiled binary. Each case below runs a pair of such shapes back to back -through ``run_unified_attention_torch``: the second launch must reuse the -first binary (no new compile) and both must match an FP32 paged reference. - -The second shape of each pair changes at least one of the head counts (at the -same GQA ratio), the window size, or the batch (by several log2 search -buckets), so a value that is still baked into the binary corrupts it. Only the -2D pairs change the window; the 3D pairs keep it at 0. Block-table rows are -wider than the batch needs and unused slots point at a poison block, so a wrong -stride cannot pass by luck. - -The routing and key checks run on any host; the launches need a gfx942 GPU. -""" - -from __future__ import annotations - -import itertools -import math - -import pytest - -from kernels.common import attention_unified as au - -_ARCH = "gfx942" -_CUS = 304 - - -def _problem(*, q_lens, kv_lens, hq, hkv, d, bs, dtype, window, num_blocks): - return 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=dtype, - sliding_window=window, - num_cus=_CUS, - num_kv_blocks=num_blocks, - ) - - -def _shape(q_lens, kv_lens, hq, hkv, window): - return dict(q_lens=q_lens, kv_lens=kv_lens, hq=hq, hkv=hkv, window=window) - - -# (id, path, 4-warp cohort or None, (d, bs, dtype), first shape, second shape) -_CASES = ( - ( - "2d-bf16-d128", - "2d", - None, - (128, 16, "bf16"), - _shape([300, 200, 250], [300, 450, 250], 32, 8, 0), - _shape( - [90, 130, 70, 110, 40, 150, 60, 80, 120], - [90, 300, 70, 210, 40, 150, 60, 380, 120], - 16, - 4, - 0, - ), - ), - # Head counts alone: the token and sequence counts match, so only the head - # counts tell the two launches apart. - ( - "2d-bf16-d128-heads-only", - "2d", - None, - (128, 16, "bf16"), - _shape([1025], [1025], 32, 8, 0), - _shape([1025], [1025], 64, 16, 0), - ), - ( - "2d-fp16-d64-window", - "2d", - None, - (64, 32, "fp16"), - _shape([260, 180], [400, 180], 8, 8, 128), - _shape( - [70, 90, 50, 100, 60, 80, 40, 30, 110], - [200, 90, 300, 100, 60, 250, 40, 30, 110], - 16, - 16, - 512, - ), - ), - ( - "4warp-d256", - "2d", - "gfx942_d256_causal", - (256, 16, "bf16"), - _shape([1024, 1024], [1024, 1024], 16, 2, 0), - _shape( - [300, 500, 200, 400, 250, 350, 150, 450, 280], - [300, 700, 200, 400, 250, 600, 150, 450, 280], - 32, - 4, - 0, - ), - ), - ( - "4warp-fold-window", - "2d", - "gfx942_d128_swa", - (128, 16, "bf16"), - _shape([1024], [1024], 32, 8, 512), - _shape( - [200, 300, 150, 250, 180, 220, 260, 140, 310], - [400, 300, 350, 250, 180, 520, 260, 140, 310], - 16, - 4, - 256, - ), - ), - # The decode split sizes segments by num_seqs * num_kv_heads, so the 3D pairs - # hold that product fixed; gfx942 routes windowed decode to 2D. - ( - "3d-decode", - "3d", - None, - (128, 16, "bf16"), - _shape([1, 1], [4096, 2500], 64, 16, 0), - _shape([1] * 8, [3000, 900, 4096, 1700, 2200, 3500, 1200, 2800], 16, 4, 0), - ), - ( - "3d-two-token", - "3d", - None, - (128, 16, "bf16"), - _shape([2, 2], [4096, 3100], 64, 16, 0), - _shape([2] * 8, [2500, 4096, 1900, 2600, 3100, 2200, 2400, 2700], 16, 4, 0), - ), -) - - -def _problems(geometry, shape): - d, bs, dtype = geometry - needed = sum(math.ceil(n / bs) for n in shape["kv_lens"]) - return _problem( - q_lens=shape["q_lens"], - kv_lens=shape["kv_lens"], - hq=shape["hq"], - hkv=shape["hkv"], - d=d, - bs=bs, - dtype=dtype, - window=shape["window"], - num_blocks=needed + 1, - ) - - -def _key(problem, path): - if path == "2d": - return au._tiled_cache_key(problem, _ARCH) - return au._tiled_3d_cache_key(problem, _ARCH) - - -@pytest.mark.parametrize( - "path, cohort, geometry, first, second", - [pytest.param(*c[1:], id=c[0]) for c in _CASES], -) -def test_pair_shares_one_binary(path, cohort, geometry, first, second): - """Both shapes take the intended path and route and share one key, so the - GPU test below exercises one binary on two shapes.""" - a, b = _problems(geometry, first), _problems(geometry, second) - for p in (a, b): - assert p.select_path() == path - route = au._gfx942_4warp_route(p, _ARCH) if path == "2d" else None - assert (route.cohort if route else None) == cohort - assert a.num_queries_per_kv == b.num_queries_per_kv - batch_changes = a.num_seqs != b.num_seqs - if batch_changes and not (path == "3d" and au._use_decode_grid(a, _ARCH)): - # The decode grid indexes sequences directly and emits no search. - assert au._binary_search_iters(a) != au._binary_search_iters(b) - assert _key(a, path) == _key(b, path) - - -@pytest.fixture(scope="module") -def torch(): - torch = pytest.importorskip("torch") - if not torch.cuda.is_available(): - pytest.skip("requires a gfx942 GPU") - name = torch.cuda.get_device_properties(0).gcnArchName - if not name.startswith(_ARCH): - pytest.skip(f"requires a gfx942 GPU, found {name}") - return torch - - -def _case(torch, geometry, shape, seed): - d, bs, dtype = geometry - tdt = torch.bfloat16 if dtype == "bf16" else torch.float16 - q_lens, kv_lens, hq, hkv = ( - shape["q_lens"], - shape["kv_lens"], - shape["hq"], - shape["hkv"], - ) - 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() - table = torch.full((len(kv_lens), max(needed) + 3), poison, dtype=torch.int32) - for s, n in enumerate(needed): - table[s, :n] = torch.tensor(order[:n], dtype=torch.int32) - order = order[n:] - q = torch.randn(sum(q_lens), hq, d, generator=gen).to(tdt) - k = torch.randn(num_blocks, bs, hkv, d, generator=gen).to(tdt) - v = torch.randn(num_blocks, bs, hkv, d, generator=gen).to(tdt) - k[poison] = 64.0 - v[poison] = 64.0 - cu = torch.tensor([0, *itertools.accumulate(q_lens)], dtype=torch.int32) - return q.cuda(), k.cuda(), v.cuda(), cu.cuda(), table.cuda() - - -def _reference(torch, q, k, v, cu, table, shape, d, bs): - q_lens, kv_lens, hq, hkv = ( - shape["q_lens"], - shape["kv_lens"], - shape["hq"], - shape["hkv"], - ) - window = shape["window"] - out = torch.zeros(q.shape, dtype=torch.float32, device=q.device) - 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() - keys = keys.repeat_interleave(hq // hkv, dim=1) - values = values.repeat_interleave(hq // hkv, dim=1) - start = int(cu[s]) - queries = q[start : start + q_len].float() - scores = torch.einsum("qhd,khd->hqk", queries, keys) / math.sqrt(d) - q_pos = (kv_len - q_len) + torch.arange(q_len, device=q.device)[:, None] - k_pos = torch.arange(kv_len, device=q.device)[None, :] - visible = k_pos <= q_pos - if window > 0: - visible = visible & (q_pos - k_pos < window) - scores = scores.masked_fill(~visible, float("-inf")) - out[start : start + q_len] = torch.einsum( - "hqk,khd->qhd", scores.softmax(-1), values - ) - return out - - -def _cache(path): - return au._ATTN_3D_TILED_CACHE if path == "3d" else au._ATTN_TILED_CACHE - - -def _compiled(): - return len(au._ATTN_TILED_CACHE) + len(au._ATTN_3D_TILED_CACHE) - - -@pytest.mark.parametrize( - "path, geometry, first, second", - [pytest.param(c[1], *c[3:], id=c[0]) for c in _CASES], -) -def test_pair_numeric_on_one_binary(torch, path, geometry, first, second): - d, bs, _ = geometry - compiled = None - for i, shape in enumerate((first, second)): - problem = _problems(geometry, shape) - q, k, v, cu, table = _case(torch, geometry, shape, seed=i) - out = torch.empty_like(q) - au.run_unified_attention_torch( - problem=problem, - q=q, - k=k, - v=v, - out=out, - cu_seqlens_q=cu, - seqused_k=torch.tensor(shape["kv_lens"], dtype=torch.int32, device="cuda"), - softmax_scale=1.0 / math.sqrt(d), - block_table=table, - softcap=0.0, - backend="3d" if path == "3d" else "tiled", - use_graph=False, - ) - torch.cuda.synchronize() - assert _key(problem, path) in _cache(path), f"shape {i} missed the {path} cache" - if compiled is None: - compiled = _compiled() - else: - assert _compiled() == compiled, "the second shape compiled a new binary" - want = _reference(torch, q, k, v, cu, table, shape, d, bs) - max_abs = (out.float() - want).abs().max().item() - assert ( - math.isfinite(max_abs) and max_abs <= 4e-2 - ), f"shape {i}: max_abs {max_abs:.3e} vs fp32 reference" diff --git a/dnn-providers/hip-kernel-provider/rocke/library/tests/test_attention_gfx950_runtime_shape_numeric.py b/dnn-providers/hip-kernel-provider/rocke/library/tests/test_attention_gfx950_runtime_shape_numeric.py deleted file mode 100644 index db5687342b18..000000000000 --- a/dnn-providers/hip-kernel-provider/rocke/library/tests/test_attention_gfx950_runtime_shape_numeric.py +++ /dev/null @@ -1,329 +0,0 @@ -# Copyright (c) Advanced Micro Devices, Inc., or its affiliates. -# SPDX-License-Identifier: MIT - -"""One gfx950 tiled binary serves every head pair, window size and batch. - -The gfx950 tiled 2D and 3D kernels read the head counts and the sliding window -as kernargs and derive the cu_q search trip count from ``num_seqs``, so shapes -that differ only in those share one cache key and one compiled binary. Each -case below runs a pair of such shapes back to back through -``run_unified_attention_torch``: the second launch must reuse the first binary -(no new compile) and both must match an FP32 paged reference. - -The second shape of each pair changes the head counts at the same GQA ratio, -the window size, or the batch by several log2 search buckets, so a value that -is still baked into the binary corrupts it. Block-table rows are wider than the -batch needs and unused slots point at a poison block, so a wrong stride cannot -pass by luck. - -The fast paged-KV descriptor bakes the 64/8 head geometry, so its pair keeps -64/8 and changes the batch; its 32/4 and 16/2 neighbours at the same ratio take -the plain combo binary. - -The routing and key checks run on any host; the launches need a gfx950 GPU. -""" - -from __future__ import annotations - -import itertools -import math - -import pytest - -from kernels.common import attention_unified as au - -_ARCH = "gfx950" -_CUS = 256 - - -def _problem(*, q_lens, kv_lens, hq, hkv, d, bs, dtype, window, num_blocks): - return 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=dtype, - sliding_window=window, - num_cus=_CUS, - num_kv_blocks=num_blocks, - ) - - -def _shape(q_lens, kv_lens, hq, hkv, window): - return dict(q_lens=q_lens, kv_lens=kv_lens, hq=hq, hkv=hkv, window=window) - - -# (id, path, route predicate or None, (d, bs, dtype), first shape, second shape). -# The route predicate is an attention_unified selector both shapes must satisfy. -# Both shapes of a pair sit on one side of every seqlen threshold the selectors -# read (combo above 256 query tokens, hipcc above 512), so only the head counts, -# the window and the batch differ. -_CASES = ( - ( - "2d-bf16-d128", - "2d", - None, - (128, 16, "bf16"), - _shape([300, 200, 250], [300, 450, 250], 32, 8, 0), - _shape( - [90, 130, 70, 110, 40, 300, 60, 80, 120], - [90, 300, 70, 210, 40, 300, 60, 380, 120], - 16, - 4, - 0, - ), - ), - # Head counts alone: the token and sequence counts match, so only the head - # counts tell the two launches apart. - ( - "2d-bf16-d128-single-batch-heads-only", - "2d", - "_enable_single_batch_combo", - (128, 16, "bf16"), - _shape([1025], [1025], 32, 8, 0), - _shape([1025], [1025], 64, 16, 0), - ), - ( - "2d-fp16-d64-window", - "2d", - None, - (64, 32, "fp16"), - _shape([260, 180], [400, 180], 8, 8, 128), - _shape( - [70, 90, 50, 100, 60, 80, 40, 30, 110], - [200, 90, 300, 100, 60, 250, 40, 30, 110], - 16, - 16, - 512, - ), - ), - # The single-batch fp16 d128 window cohort: heads and window change. - ( - "2d-fp16-d128-single-batch-window", - "2d", - "_enable_single_batch_combo", - (128, 16, "fp16"), - _shape([2048], [2048], 32, 8, 256), - _shape([2048], [2048], 64, 16, 1024), - ), - ( - "2d-bf16-d64-fast-desc-batch", - "2d", - "_enable_fast_paged_kv_desc", - (64, 32, "bf16"), - _shape([500, 400], [700, 900], 64, 8, 0), - _shape( - [300, 280, 400, 260, 350, 310, 270, 290, 330], - [300, 600, 400, 800, 350, 310, 500, 290, 330], - 64, - 8, - 0, - ), - ), - ( - "2d-bf16-d64-combo-heads", - "2d", - "_enable_combo_2d", - (64, 32, "bf16"), - _shape([500, 400], [700, 900], 32, 4, 0), - _shape( - [300, 280, 400, 260, 350, 310, 270, 290, 330], - [300, 600, 400, 800, 350, 310, 500, 290, 330], - 16, - 2, - 0, - ), - ), - ( - "2d-bf16-d256", - "2d", - "_d256_gfx950_fast", - (256, 16, "bf16"), - _shape([512, 512], [1024, 1024], 16, 2, 0), - _shape( - [300, 500, 200, 400, 250, 350, 150, 450, 280], - [300, 700, 200, 400, 250, 600, 150, 450, 280], - 32, - 4, - 0, - ), - ), - # The decode split sizes segments by num_seqs * num_kv_heads, so the 3D pairs - # hold that product fixed; windowed decode routes to 2D. - ( - "3d-decode", - "3d", - None, - (128, 16, "bf16"), - _shape([1, 1], [4096, 2500], 64, 16, 0), - _shape([1] * 8, [3000, 900, 4096, 1700, 2200, 3500, 1200, 2800], 16, 4, 0), - ), - ( - "3d-two-token", - "3d", - None, - (128, 16, "bf16"), - _shape([2, 2], [4096, 3100], 64, 16, 0), - _shape([2] * 8, [2500, 4096, 1900, 2600, 3100, 2200, 2400, 2700], 16, 4, 0), - ), -) - - -def _problems(geometry, shape): - d, bs, dtype = geometry - needed = sum(math.ceil(n / bs) for n in shape["kv_lens"]) - return _problem( - q_lens=shape["q_lens"], - kv_lens=shape["kv_lens"], - hq=shape["hq"], - hkv=shape["hkv"], - d=d, - bs=bs, - dtype=dtype, - window=shape["window"], - num_blocks=needed + 1, - ) - - -def _key(problem, path): - if path == "2d": - return au._tiled_cache_key(problem, _ARCH) - return au._tiled_3d_cache_key(problem, _ARCH) - - -@pytest.mark.parametrize( - "path, route, geometry, first, second", - [pytest.param(*c[1:], id=c[0]) for c in _CASES], -) -def test_pair_shares_one_binary(path, route, geometry, first, second): - """Both shapes take the intended path and route and share one key, so the - GPU test below exercises one binary on two shapes.""" - a, b = _problems(geometry, first), _problems(geometry, second) - for p in (a, b): - assert p.select_path() == path - if route is not None: - assert getattr(au, route)(p, _ARCH) - assert a.num_queries_per_kv == b.num_queries_per_kv - batch_changes = a.num_seqs != b.num_seqs - if batch_changes and not (path == "3d" and au._use_decode_grid(a, _ARCH)): - # The decode grid indexes sequences directly and emits no search. - assert au._binary_search_iters(a) != au._binary_search_iters(b) - assert _key(a, path) == _key(b, path) - - -@pytest.fixture(scope="module") -def torch(): - torch = pytest.importorskip("torch") - if not torch.cuda.is_available(): - pytest.skip("requires a gfx950 GPU") - name = torch.cuda.get_device_properties(0).gcnArchName - if not name.startswith(_ARCH): - pytest.skip(f"requires a gfx950 GPU, found {name}") - return torch - - -def _case(torch, geometry, shape, seed): - d, bs, dtype = geometry - tdt = torch.bfloat16 if dtype == "bf16" else torch.float16 - q_lens, kv_lens, hq, hkv = ( - shape["q_lens"], - shape["kv_lens"], - shape["hq"], - shape["hkv"], - ) - 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() - table = torch.full((len(kv_lens), max(needed) + 3), poison, dtype=torch.int32) - for s, n in enumerate(needed): - table[s, :n] = torch.tensor(order[:n], dtype=torch.int32) - order = order[n:] - q = torch.randn(sum(q_lens), hq, d, generator=gen).to(tdt) - k = torch.randn(num_blocks, bs, hkv, d, generator=gen).to(tdt) - v = torch.randn(num_blocks, bs, hkv, d, generator=gen).to(tdt) - k[poison] = 64.0 - v[poison] = 64.0 - cu = torch.tensor([0, *itertools.accumulate(q_lens)], dtype=torch.int32) - return q.cuda(), k.cuda(), v.cuda(), cu.cuda(), table.cuda() - - -def _reference(torch, q, k, v, cu, table, shape, d, bs): - q_lens, kv_lens, hq, hkv = ( - shape["q_lens"], - shape["kv_lens"], - shape["hq"], - shape["hkv"], - ) - window = shape["window"] - out = torch.zeros(q.shape, dtype=torch.float32, device=q.device) - 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() - keys = keys.repeat_interleave(hq // hkv, dim=1) - values = values.repeat_interleave(hq // hkv, dim=1) - start = int(cu[s]) - queries = q[start : start + q_len].float() - scores = torch.einsum("qhd,khd->hqk", queries, keys) / math.sqrt(d) - q_pos = (kv_len - q_len) + torch.arange(q_len, device=q.device)[:, None] - k_pos = torch.arange(kv_len, device=q.device)[None, :] - visible = k_pos <= q_pos - if window > 0: - visible = visible & (q_pos - k_pos < window) - scores = scores.masked_fill(~visible, float("-inf")) - out[start : start + q_len] = torch.einsum( - "hqk,khd->qhd", scores.softmax(-1), values - ) - return out - - -def _cache(path): - return au._ATTN_3D_TILED_CACHE if path == "3d" else au._ATTN_TILED_CACHE - - -def _compiled(): - return len(au._ATTN_TILED_CACHE) + len(au._ATTN_3D_TILED_CACHE) - - -@pytest.mark.parametrize( - "path, geometry, first, second", - [pytest.param(c[1], *c[3:], id=c[0]) for c in _CASES], -) -def test_pair_numeric_on_one_binary(torch, path, geometry, first, second): - d, bs, _ = geometry - compiled = None - for i, shape in enumerate((first, second)): - problem = _problems(geometry, shape) - q, k, v, cu, table = _case(torch, geometry, shape, seed=i) - out = torch.empty_like(q) - au.run_unified_attention_torch( - problem=problem, - q=q, - k=k, - v=v, - out=out, - cu_seqlens_q=cu, - seqused_k=torch.tensor(shape["kv_lens"], dtype=torch.int32, device="cuda"), - softmax_scale=1.0 / math.sqrt(d), - block_table=table, - softcap=0.0, - backend="3d" if path == "3d" else "tiled", - use_graph=False, - ) - torch.cuda.synchronize() - assert _key(problem, path) in _cache(path), f"shape {i} missed the {path} cache" - if compiled is None: - compiled = _compiled() - else: - assert _compiled() == compiled, "the second shape compiled a new binary" - want = _reference(torch, q, k, v, cu, table, shape, d, bs) - max_abs = (out.float() - want).abs().max().item() - assert ( - math.isfinite(max_abs) and max_abs <= 4e-2 - ), f"shape {i}: max_abs {max_abs:.3e} vs fp32 reference" diff --git a/dnn-providers/hip-kernel-provider/rocke/library/tests/test_attention_spec_predicate_parity.py b/dnn-providers/hip-kernel-provider/rocke/library/tests/test_attention_spec_predicate_parity.py index 01b218e209d5..675f2c82b899 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/tests/test_attention_spec_predicate_parity.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/tests/test_attention_spec_predicate_parity.py @@ -262,7 +262,7 @@ def test_reduce_predicate_admits_the_reduce_spec_of_every_admitted_3d_problem( (dict(dtype="fp32"), "fp16/bf16"), (dict(head_size=96), "divisible by 64"), (dict(num_query_heads=30, num_kv_heads=8), "divisible by num_kv_heads"), - (dict(num_query_heads=-8, num_kv_heads=8), "positive num_query_heads"), + (dict(num_query_heads=-8, num_kv_heads=8), "must be > 0"), (dict(num_segments=40961), "over gfx950's 163840 B LDS (at most 40960)"), ], ) diff --git a/dnn-providers/hip-kernel-provider/rocke/library/tests/test_attention_tiled_cpp_entry_guards.py b/dnn-providers/hip-kernel-provider/rocke/library/tests/test_attention_tiled_cpp_entry_guards.py new file mode 100644 index 000000000000..76ccc968bce0 --- /dev/null +++ b/dnn-providers/hip-kernel-provider/rocke/library/tests/test_attention_tiled_cpp_entry_guards.py @@ -0,0 +1,75 @@ +# Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +# SPDX-License-Identifier: MIT + +"""The C++ tiled-attention entries refuse what the Python builders refuse, with +the same reason, instead of crashing. + +* A zero or non-dividing head pair: the 2D ``_new`` entries used to fail at + naming before initialising the builder, and the caller then freed that + uninitialised builder (signal 11). +* fastKV register-P on gfx942: the body is gfx950-only, so the C++ builder must + refuse it as Python does rather than emit gfx950 MFMA for gfx942. +""" + +from __future__ import annotations + +import dataclasses + +import pytest +from kernels.gfx942.attention_tiled_2d import ( + UnifiedAttention2DTiledSpec as Gfx942Spec2D, +) +from kernels.gfx950.attention_tiled_2d import ( + UnifiedAttention2DTiledSpec as Gfx950Spec2D, +) +from kernels.gfx950.attention_tiled_2d_fastkv_regp import ( + build_unified_attention_2d_fastkv_register_p, + make_fastkv_register_p_spec, +) + +rocke_engine = pytest.importorskip("rocke_engine") + +_PROBLEM = { + "head_size": 128, + "block_size": 16, + "dtype": "bf16", + "use_sinks": False, + "sliding_window": 0, + "has_softcap": False, +} + + +@pytest.mark.parametrize( + "arch, spec_cls", + [("gfx942", Gfx942Spec2D), ("gfx950", Gfx950Spec2D)], +) +@pytest.mark.parametrize( + "heads, reason", + [((32, 0), "must be > 0"), ((32, 6), "must be divisible by num_kv_heads")], +) +def test_2d_entry_refuses_bad_head_counts_like_python(arch, spec_cls, heads, reason): + hq, hkv = heads + spec = spec_cls(num_query_heads=hq, num_kv_heads=hkv, **_PROBLEM) + with pytest.raises(ValueError, match=reason) as py_err: + spec.kernel_name() + lower = getattr(rocke_engine, f"{arch}_attention_tiled_2d_lower_llvm") + with pytest.raises(Exception) as cpp_err: + lower(dataclasses.asdict(spec), arch) + assert str(py_err.value) in str(cpp_err.value) + + +def test_fastkv_register_p_entry_refuses_gfx942_like_python(): + spec = make_fastkv_register_p_spec( + Gfx950Spec2D( + num_query_heads=64, + num_kv_heads=8, + **{**_PROBLEM, "head_size": 64, "block_size": 32}, + ) + ) + with pytest.raises(ValueError, match="gfx950 kernel") as py_err: + build_unified_attention_2d_fastkv_register_p(spec, arch="gfx942") + with pytest.raises(Exception) as cpp_err: + rocke_engine.gfx950_attention_tiled_2d_fastkv_regp_lower_llvm( + dataclasses.asdict(spec), "gfx942" + ) + assert str(py_err.value) in str(cpp_err.value) diff --git a/dnn-providers/hip-kernel-provider/rocke/library/tests/test_attention_tiled_runtime_shape_numeric.py b/dnn-providers/hip-kernel-provider/rocke/library/tests/test_attention_tiled_runtime_shape_numeric.py new file mode 100644 index 000000000000..867e9b3e022f --- /dev/null +++ b/dnn-providers/hip-kernel-provider/rocke/library/tests/test_attention_tiled_runtime_shape_numeric.py @@ -0,0 +1,464 @@ +# Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +# SPDX-License-Identifier: MIT + +"""One gfx942 or gfx950 tiled binary serves every head pair, window size and batch. + +The tiled 2D (gfx942 also 4-warp GQA) and 3D kernels read the head counts and +the sliding window as kernargs and derive the cu_q search trip count from +``num_seqs``, so shapes that differ only in those share one cache key and one +compiled binary. Each case below runs a pair of such shapes back to back: the +second launch must reuse the first binary (no new compile) and both must match +an FP32 paged reference. + +The second shape of each pair changes at least one of the head counts (at the +same GQA ratio), the window size, or the batch (by several log2 search +buckets), so a value that is still baked into the binary corrupts it. Block-table +rows are wider than the batch needs and unused slots point at a poison block, +so a wrong stride cannot pass by luck. + +gfx942 routes windowed decode to 2D, so its 3D pairs keep the window at 0. The +gfx950 fast paged-KV descriptor bakes the 64/8 head geometry, so its pair keeps +64/8 and changes the batch; its 32/4 and 16/2 neighbours at the same ratio take +the plain combo binary. + +The kernels run on ``rocke.runtime`` device buffers filled from numpy, and the +reference is numpy over the fp16/bf16-rounded inputs, so the launches run in GPU +lanes without torch. 2D launches go through ``run_unified_attention_torch``; +3D launches build and launch the segment + reduce pair as its 3D path does. +The routing and key checks run on any host; the launches need the case's GPU. +""" + +from __future__ import annotations + +import itertools +import math + +import numpy as np +import pytest + +from kernels.common import attention_unified as au +from rocke.numeric.bf16 import bf16_to_f32, f32_to_bf16 +from rocke.runtime.hip_module import get_device_arch + +from ._attention_device_array import DeviceArray + +_CUS = {"gfx942": 304, "gfx950": 256} + + +def _shape(q_lens, kv_lens, hq, hkv, window): + return { + "q_lens": q_lens, + "kv_lens": kv_lens, + "hq": hq, + "hkv": hkv, + "window": window, + } + + +# Pairs both arches run unchanged. +_HEADS_ONLY = (_shape([1025], [1025], 32, 8, 0), _shape([1025], [1025], 64, 16, 0)) +_D64_WINDOW = ( + _shape([260, 180], [400, 180], 8, 8, 128), + _shape( + [70, 90, 50, 100, 60, 80, 40, 30, 110], + [200, 90, 300, 100, 60, 250, 40, 30, 110], + 16, + 16, + 512, + ), +) +# The decode split sizes segments by num_seqs * num_kv_heads, so the 3D pairs +# hold that product fixed; windowed decode routes to 2D. +_DECODE = ( + _shape([1, 1], [4096, 2500], 64, 16, 0), + _shape([1] * 8, [3000, 900, 4096, 1700, 2200, 3500, 1200, 2800], 16, 4, 0), +) +_TWO_TOKEN = ( + _shape([2, 2], [4096, 3100], 64, 16, 0), + _shape([2] * 8, [2500, 4096, 1900, 2600, 3100, 2200, 2400, 2700], 16, 4, 0), +) +_D64_BATCH_KV = [300, 600, 400, 800, 350, 310, 500, 290, 330] +_D64_BATCH_Q = [300, 280, 400, 260, 350, 310, 270, 290, 330] + +# (arch, id, path, route, (d, bs, dtype), first shape, second shape). +# gfx942: ``route`` is the 4-warp GQA cohort both 2D shapes take, or None for +# the plain tiled kernel. gfx950: ``route`` names an attention_unified selector +# both shapes must satisfy, or None for no check. gfx950 pairs sit on one side +# of every seqlen threshold the selectors read (combo above 256 query tokens, +# hipcc above 512), so only the head counts, the window and the batch differ. +_CASES = ( + ( + "gfx942", + "2d-bf16-d128", + "2d", + None, + (128, 16, "bf16"), + _shape([300, 200, 250], [300, 450, 250], 32, 8, 0), + _shape( + [90, 130, 70, 110, 40, 150, 60, 80, 120], + [90, 300, 70, 210, 40, 150, 60, 380, 120], + 16, + 4, + 0, + ), + ), + # Head counts alone: the token and sequence counts match, so only the head + # counts tell the two launches apart. + ("gfx942", "2d-bf16-d128-heads-only", "2d", None, (128, 16, "bf16"), *_HEADS_ONLY), + ("gfx942", "2d-fp16-d64-window", "2d", None, (64, 32, "fp16"), *_D64_WINDOW), + ( + "gfx942", + "4warp-d256", + "2d", + "gfx942_d256_causal", + (256, 16, "bf16"), + _shape([1024, 1024], [1024, 1024], 16, 2, 0), + _shape( + [300, 500, 200, 400, 250, 350, 150, 450, 280], + [300, 700, 200, 400, 250, 600, 150, 450, 280], + 32, + 4, + 0, + ), + ), + ( + "gfx942", + "4warp-fold-window", + "2d", + "gfx942_d128_swa", + (128, 16, "bf16"), + _shape([1024], [1024], 32, 8, 512), + _shape( + [200, 300, 150, 250, 180, 220, 260, 140, 310], + [400, 300, 350, 250, 180, 520, 260, 140, 310], + 16, + 4, + 256, + ), + ), + ("gfx942", "3d-decode", "3d", None, (128, 16, "bf16"), *_DECODE), + ("gfx942", "3d-two-token", "3d", None, (128, 16, "bf16"), *_TWO_TOKEN), + ( + "gfx950", + "2d-bf16-d128", + "2d", + None, + (128, 16, "bf16"), + _shape([300, 200, 250], [300, 450, 250], 32, 8, 0), + _shape( + [90, 130, 70, 110, 40, 300, 60, 80, 120], + [90, 300, 70, 210, 40, 300, 60, 380, 120], + 16, + 4, + 0, + ), + ), + ( + "gfx950", + "2d-bf16-d128-single-batch-heads-only", + "2d", + "_enable_single_batch_combo", + (128, 16, "bf16"), + *_HEADS_ONLY, + ), + ("gfx950", "2d-fp16-d64-window", "2d", None, (64, 32, "fp16"), *_D64_WINDOW), + # The single-batch fp16 d128 window cohort: heads and window change. + ( + "gfx950", + "2d-fp16-d128-single-batch-window", + "2d", + "_enable_single_batch_combo", + (128, 16, "fp16"), + _shape([2048], [2048], 32, 8, 256), + _shape([2048], [2048], 64, 16, 1024), + ), + ( + "gfx950", + "2d-bf16-d64-fast-desc-batch", + "2d", + "_enable_fast_paged_kv_desc", + (64, 32, "bf16"), + _shape([500, 400], [700, 900], 64, 8, 0), + _shape(_D64_BATCH_Q, _D64_BATCH_KV, 64, 8, 0), + ), + ( + "gfx950", + "2d-bf16-d64-combo-heads", + "2d", + "_enable_combo_2d", + (64, 32, "bf16"), + _shape([500, 400], [700, 900], 32, 4, 0), + _shape(_D64_BATCH_Q, _D64_BATCH_KV, 16, 2, 0), + ), + ( + "gfx950", + "2d-bf16-d256", + "2d", + "_d256_gfx950_fast", + (256, 16, "bf16"), + _shape([512, 512], [1024, 1024], 16, 2, 0), + _shape( + [300, 500, 200, 400, 250, 350, 150, 450, 280], + [300, 700, 200, 400, 250, 600, 150, 450, 280], + 32, + 4, + 0, + ), + ), + ("gfx950", "3d-decode", "3d", None, (128, 16, "bf16"), *_DECODE), + ("gfx950", "3d-two-token", "3d", None, (128, 16, "bf16"), *_TWO_TOKEN), +) + + +def _params(*fields): + return [pytest.param(*(c[i] for i in fields), id=f"{c[0]}-{c[1]}") for c in _CASES] + + +def _problem(arch, geometry, shape): + d, bs, dtype = geometry + q_lens, kv_lens = shape["q_lens"], shape["kv_lens"] + return au.UnifiedAttentionProblem( + total_q=sum(q_lens), + num_seqs=len(q_lens), + num_query_heads=shape["hq"], + num_kv_heads=shape["hkv"], + head_size=d, + block_size=bs, + max_seqlen_q=max(q_lens), + max_seqlen_k=max(kv_lens), + dtype=dtype, + sliding_window=shape["window"], + num_cus=_CUS[arch], + num_kv_blocks=sum(math.ceil(n / bs) for n in kv_lens) + 1, + ) + + +def _key(problem, path, arch): + if path == "2d": + return au._tiled_cache_key(problem, arch) + return au._tiled_3d_cache_key(problem, arch) + + +def _check_route(arch, path, route, problem): + if arch == "gfx942": + found = au._gfx942_4warp_route(problem, arch) if path == "2d" else None + assert (found.cohort if found else None) == route + elif route is not None: + assert getattr(au, route)(problem, arch) + + +@pytest.mark.parametrize( + "arch, path, route, geometry, first, second", _params(0, 2, 3, 4, 5, 6) +) +def test_pair_shares_one_binary(arch, path, route, geometry, first, second): + """Both shapes take the intended path and route and share one key, so the + GPU test below exercises one binary on two shapes.""" + a, b = _problem(arch, geometry, first), _problem(arch, geometry, second) + for p in (a, b): + assert p.select_path() == path + _check_route(arch, path, route, p) + assert a.num_queries_per_kv == b.num_queries_per_kv + batch_changes = a.num_seqs != b.num_seqs + if batch_changes and not (path == "3d" and au._use_decode_grid(a, arch)): + # The decode grid indexes sequences directly and emits no search. + assert au._binary_search_iters(a) != au._binary_search_iters(b) + assert _key(a, path, arch) == _key(b, path, arch) + + +def _encode(x, dtype): + """fp32 values as the kernel's fp16 or bf16 (raw uint16) host array.""" + return x.astype(np.float16) if dtype == "fp16" else f32_to_bf16(x) + + +def _decode(x, dtype): + return x.astype(np.float32) if dtype == "fp16" else bf16_to_f32(x) + + +def _case(geometry, shape, seed): + """fp32 q, k, v exact in the kernel dtype, cu_seqlens_q and a block table + wider than needed whose unused slots point at a poison block.""" + d, bs, dtype = geometry + q_lens, kv_lens, hq, hkv = (shape[f] for f in ("q_lens", "kv_lens", "hq", "hkv")) + 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 = iter(rng.permutation(poison)) + table = np.full((len(kv_lens), max(needed) + 3), poison, dtype=np.int32) + for s, n in enumerate(needed): + table[s, :n] = [next(order) for _ in range(n)] + + def randn(*dims): + return _decode(_encode(rng.standard_normal(dims, np.float32), dtype), dtype) + + 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 = np.array([0, *itertools.accumulate(q_lens)], dtype=np.int32) + return q, k, v, cu, table + + +def _reference(q, k, v, cu, table, shape, bs): + """FP32 causal (optionally windowed) paged attention, one KV head group at + a time.""" + q_lens, kv_lens, hkv, window = ( + shape[f] for f in ("q_lens", "kv_lens", "hkv", "window") + ) + d = q.shape[-1] + ratio = shape["hq"] // hkv + out = np.zeros(q.shape, dtype=np.float32) + for s, (q_len, kv_len) in enumerate(zip(q_lens, kv_lens)): + blocks = table[s, : math.ceil(kv_len / bs)] + keys = k[blocks].reshape(-1, hkv, d)[:kv_len] + values = v[blocks].reshape(-1, hkv, d)[:kv_len] + start = int(cu[s]) + q_pos = (kv_len - q_len) + np.arange(q_len)[:, None] + k_pos = np.arange(kv_len)[None, :] + visible = k_pos <= q_pos + if window > 0: + visible &= q_pos - k_pos < window + for g in range(hkv): + heads = slice(g * ratio, (g + 1) * ratio) + queries = q[start : start + q_len, heads].transpose(1, 0, 2) + scores = (queries @ keys[:, g].T) / math.sqrt(d) + scores = np.where(visible, scores, -np.inf) + probs = np.exp(scores - scores.max(-1, keepdims=True)) + probs /= probs.sum(-1, keepdims=True) + out[start : start + q_len, heads] = (probs @ values[:, g]).transpose( + 1, 0, 2 + ) + return out + + +def _launch_3d(rt, arch, problem, q, k, v, out, cu, seq_lens, table, scale): + """The 3D segment + reduce pair, launched as ``run_unified_attention_torch``'s + 3D path does (same support check, cache key, segment count and grids).""" + ok, why = au._explicit_path_supported(problem, None, "3d", arch) + assert ok, why + num_segments = au._num_segments(problem, arch) + prepared = au._get_3d_pipeline( + problem, au._tiled_3d_cache_key(problem, arch), num_segments, arch=arch + ) + # NaN until the kernels write it, so a skipped workspace slot fails. + segm_output, segm_max, segm_expsum = ( + DeviceArray(rt, np.full(spec.shape, np.nan, np.float32)) + for spec in au._attention_3d_workspace_specs(problem, num_segments, None) + ) + seg_vals = { + "segm_output_ptr": segm_output, + "segm_max_ptr": segm_max, + "segm_expsum_ptr": segm_expsum, + "query_ptr": q, + "key_cache_ptr": k, + "value_cache_ptr": v, + "sink_ptr": None, + "block_tables_ptr": table, + "seq_lens_ptr": seq_lens, + "alibi_slopes_ptr": 0, + "qq_bias_ptr": 0, + "query_start_len_ptr": cu, + "scale": scale, + "k_scale": 1.0, + "v_scale": 1.0, + "softcap": 0.0, + "num_seqs": problem.num_seqs, + "block_table_stride": table.stride(0), + "qq_bias_stride_0": 0, + **au._runtime_kernarg_values(problem), + } + red_vals = { + "output_ptr": out, + "segm_output_ptr": segm_output, + "segm_max_ptr": segm_max, + "segm_expsum_ptr": segm_expsum, + "seq_lens_ptr": seq_lens, + **au._runtime_kernarg_values(problem), + } + au._launch_3d_pipeline(prepared, seg_vals, red_vals, stream=0, capturing=False) + + +def _launch(rt, arch, path, problem, geometry, q, k, v, cu, table, kv_lens): + """Run one shape on ``path``; returns O as float32.""" + d, _, dtype = geometry + host_q = _encode(q, dtype) + # Every byte 0xFF is NaN in fp16 and bf16: an element the kernel skips + # fails the check instead of passing on stale memory. + out_host = np.empty_like(host_q) + out_host.view(np.uint8)[...] = 0xFF + out = DeviceArray(rt, out_host) + args = { + "q": DeviceArray(rt, host_q), + "k": DeviceArray(rt, _encode(k, dtype)), + "v": DeviceArray(rt, _encode(v, dtype)), + "out": out, + "cu": DeviceArray(rt, cu), + "seq_lens": DeviceArray(rt, np.array(kv_lens, dtype=np.int32)), + "table": DeviceArray(rt, table), + "scale": 1.0 / math.sqrt(d), + } + if path == "3d": + _launch_3d(rt, arch, problem, **args) + else: + au.run_unified_attention_torch( + problem=problem, + q=args["q"], + k=args["k"], + v=args["v"], + out=out, + cu_seqlens_q=args["cu"], + seqused_k=args["seq_lens"], + softmax_scale=args["scale"], + block_table=args["table"], + softcap=0.0, + backend="tiled", + use_graph=False, + ) + # The read is a blocking copy, so the launch has finished. + return _decode(out.read(), dtype) + + +def _cache(path): + return au._ATTN_3D_TILED_CACHE if path == "3d" else au._ATTN_TILED_CACHE + + +def _compiled(): + return len(au._ATTN_TILED_CACHE) + len(au._ATTN_3D_TILED_CACHE) + + +@pytest.fixture(scope="module") +def device_arch(): + return get_device_arch() + + +@pytest.mark.gpu +@pytest.mark.parametrize("arch, path, geometry, first, second", _params(0, 2, 4, 5, 6)) +def test_pair_numeric_on_one_binary( + device_arch, monkeypatch, arch, path, geometry, first, second +): + if device_arch != arch: + pytest.skip(f"requires a {arch} GPU, found {device_arch or 'no GPU'}") + from rocke.runtime.hip_module import Runtime + + # The gfx942 2D graph replay captures through torch; launch eagerly. + monkeypatch.setenv("HIPDNN_GFX942_2D_GRAPH", "0") + rt = Runtime() + _, bs, _ = geometry + compiled = None + for i, shape in enumerate((first, second)): + problem = _problem(arch, geometry, shape) + q, k, v, cu, table = _case(geometry, shape, seed=i) + got = _launch( + rt, arch, path, problem, geometry, q, k, v, cu, table, shape["kv_lens"] + ) + missed = f"shape {i} missed the {path} cache" + assert _key(problem, path, arch) in _cache(path), missed + if compiled is None: + compiled = _compiled() + else: + assert _compiled() == compiled, "the second shape compiled a new binary" + want = _reference(q, k, v, cu, table, shape, bs) + max_abs = float(np.abs(got - want).max()) + error = f"shape {i}: max_abs {max_abs:.3e} vs fp32 reference" + assert math.isfinite(max_abs) and max_abs <= 4e-2, error diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/CMakeLists.txt b/dnn-providers/hip-kernel-provider/rocke/platform/CMakeLists.txt index 8c8669bded02..5d04d72609ef 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/CMakeLists.txt +++ b/dnn-providers/hip-kernel-provider/rocke/platform/CMakeLists.txt @@ -731,8 +731,8 @@ if(ROCKE_INSTALL_TESTS) # host-side SDPA sequence-length and padding-mask adapters/specs, the # fmha_mfma free-stride ABI and layout contracts, the 2D/3D mask_type spec # fields, the 3D segment+reduce module's pairing, the reduce's LDS segment - # bound against the compiled code object, and the backward-test memo bound - # in tests/conftest.py. + # bound against the compiled code object, the tiled C++ entries' head-count + # and arch refusals, and the backward-test memo bound in tests/conftest.py. # File-list style (same pattern as conv_pytest) so PYTHONPATH resolves rocke/ # correctly from the installed tree. set(_ROCKE_LIBRARY_ATTENTION_CPU_TEST_RELPATHS @@ -744,6 +744,7 @@ if(ROCKE_INSTALL_TESTS) "tests/library/tests/test_attention_tiled_3d_decode_split.py" "tests/library/tests/test_attention_tiled_3d_kv_cache_policy.py" "tests/library/tests/test_attention_tiled_launch_signature.py" + "tests/library/tests/test_attention_tiled_cpp_entry_guards.py" "tests/library/tests/test_attention_tiled_3d_mask_type.py" "tests/library/tests/test_attention_tiled_3d_module.py" "tests/library/tests/test_attention_tiled_3d_reduce_lds.py" @@ -1031,8 +1032,7 @@ if(ROCKE_INSTALL_TESTS) "tests/library/tests/test_gfx1250_attention.py" "tests/library/tests/test_gfx942_gqa_head_fold_numeric.py" "tests/library/tests/test_attention_gfx942_kv_pages.py" - "tests/library/tests/test_attention_gfx942_runtime_shape_numeric.py" - "tests/library/tests/test_attention_gfx950_runtime_shape_numeric.py" + "tests/library/tests/test_attention_tiled_runtime_shape_numeric.py" "tests/library/tests/test_attention_tiled_3d_kv_lds_swizzle.py" "tests/library/tests/test_attention_tiled_3d_running_max.py" "tests/library/tests/test_attention_unified_2d_masked_tile_numeric.py" diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/helpers/transforms.cpp b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/helpers/transforms.cpp index 3a2037849a41..939731e179ce 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/helpers/transforms.cpp +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/helpers/transforms.cpp @@ -1926,56 +1926,13 @@ rocke_dynamic_tensor_descriptor_t* int n_coords, rocke_value_t* const* strides) { - rocke_dynamic_tensor_descriptor_t* ddesc; - int i; - /* Fake lengths (all 0) — never used for offset computation in dynamic path */ - int fake_lengths[8]; - - if(!rocke_i_live(b) || n_coords < 1 || n_coords > 8) - { - return NULL; - } - ddesc = (rocke_dynamic_tensor_descriptor_t*)rocke_arena_calloc( - &b->arena, sizeof(rocke_dynamic_tensor_descriptor_t)); - if(ddesc == NULL) - { - return NULL; - } - - for(i = 0; i < n_coords; ++i) - { - fake_lengths[i] = 0; - } - - /* Build the static base using fake lengths (never used for dynamic offset). */ - { - /* Use row-major int strides of all-1 as placeholder. */ - int fake_strides[8]; - for(i = 0; i < n_coords; ++i) - { - fake_strides[i] = 1; - } - /* We construct the base descriptor fields directly. */ - rocke_tensor_descriptor_t* base = rocke_tensor_descriptor_naive( - b, name, fake_lengths, n_coords, fake_strides, coord_names, n_coords); - if(base == NULL) - { - return NULL; - } - ddesc->base = *base; - } - - /* Override with runtime strides. */ - ddesc->n_dynamic = n_coords; - for(i = 0; i < n_coords; ++i) - { - ddesc->dynamic_strides[i] = strides[i]; - } - ddesc->base.is_dynamic = 1; - return ddesc; + /* All-Value create(): a NULL stride takes the static rule's 1. The base is + * the first field, so the cast recovers the dynamic descriptor. */ + return (rocke_dynamic_tensor_descriptor_t*)rocke_tensor_descriptor_naive_dynamic_mixed( + b, name, coord_names, n_coords, strides, NULL); } -rocke_dynamic_tensor_descriptor_t* +rocke_tensor_descriptor_t* rocke_tensor_descriptor_naive_dynamic_mixed(rocke_ir_builder_t* b, const char* name, const char* const* coord_names, @@ -1989,21 +1946,30 @@ rocke_dynamic_tensor_descriptor_t* int base_strides[8]; int i; - if(!rocke_i_live(b) || n_coords < 1 || n_coords > 8) + if(!rocke_i_live(b)) { return NULL; } + if(n_coords < 1 || n_coords > 8) + { + rocke_i_set_err(b, + ROCKE_ERR_VALUE, + "DynamicTensorDescriptor.create: %d coords (expected 1..8)", + n_coords); + return NULL; + } ddesc = (rocke_dynamic_tensor_descriptor_t*)rocke_arena_calloc( &b->arena, sizeof(rocke_dynamic_tensor_descriptor_t)); if(ddesc == NULL) { + rocke_i_set_err(b, ROCKE_ERR_OOM, "DynamicTensorDescriptor.create: OOM"); return NULL; } /* Python create(): base_lengths all 0, base_strides the int stride or 1. */ for(i = 0; i < n_coords; ++i) { fake_lengths[i] = 0; - base_strides[i] = strides[i] != NULL ? 1 : static_strides[i]; + base_strides[i] = (strides[i] != NULL || static_strides == NULL) ? 1 : static_strides[i]; } base = rocke_tensor_descriptor_naive( b, name, fake_lengths, n_coords, base_strides, coord_names, n_coords); @@ -2018,7 +1984,7 @@ rocke_dynamic_tensor_descriptor_t* ddesc->dynamic_strides[i] = strides[i]; } ddesc->base.is_dynamic = 1; - return ddesc; + return &ddesc->base; } bool rocke_dynamic_tensor_descriptor_offset(rocke_ir_builder_t* b, 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 39843a3938fb..31b6f8445ff5 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 @@ -91,6 +91,16 @@ typedef struct rocke_unified_attn_problem int rocke_unified_attn_num_queries_per_kv(rocke_ir_builder_t* b, const rocke_unified_attn_problem_t* p); +/* Python: kernels.common._fmha_common.check_head_counts. True when both head + * counts are positive and num_query_heads splits evenly over num_kv_heads; + * otherwise false, with the Python reason text written to `reason`. Every + * tiled attention entry runs it before naming the kernel (the name carries + * the GQA ratio). */ +bool rocke_attention_check_head_counts(int num_query_heads, + int num_kv_heads, + char* reason, + size_t reason_cap); + /* Python: _resolve_attention_arch() override hook. The Python memoizes the * device arch (falling back to "gfx950"); this lets a host pin the arch the * selectors resolve to (NULL restores the "gfx950" fallback). Pass a string diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/helper_rocke.helpers.transforms.h b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/helper_rocke.helpers.transforms.h index 8183b1fb39da..4f2f7a912f7e 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/helper_rocke.helpers.transforms.h +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/helper_rocke.helpers.transforms.h @@ -410,8 +410,13 @@ rocke_dynamic_tensor_descriptor_t* rocke_value_t* const* strides); /* Python: DynamicTensorDescriptor.create with a mix of Value and int strides. - * strides[i] == NULL takes static_strides[i] as a compile-time stride. */ -rocke_dynamic_tensor_descriptor_t* + * strides[i] == NULL takes static_strides[i] as a compile-time stride (1 when + * static_strides is NULL). Returns the descriptor as its base view (is_dynamic + * set; pass it to rocke_tensor_descriptor_transform / _offset directly), or NULL + * with the builder's error set -- the descriptor calls then no-op on the dead + * builder, so callers need not check. rocke_tensor_descriptor_naive_dynamic is + * this with every stride a Value. */ +rocke_tensor_descriptor_t* rocke_tensor_descriptor_naive_dynamic_mixed(rocke_ir_builder_t* b, const char* name, const char* const* coord_names, diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/instance_gfx950_attention_tiled_2d_internal.h b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/instance_gfx950_attention_tiled_2d_internal.h index b4201005505b..8665f4c206c4 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/instance_gfx950_attention_tiled_2d_internal.h +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/instance_gfx950_attention_tiled_2d_internal.h @@ -272,9 +272,10 @@ typedef struct rocke_gfx950_attn2d_build_ctx rocke_value_t* kv_block_bytes_c_v; /* kv_stride_blk_b_v (one block) */ /* bf16/fp16 loaders (!KV_FP8, !FAST_PAGED_KV_DESC): the lane's in-page byte * offset, each copy's first-token offset in its page (indexed by that - * token; [0] stays NULL), and the i64 page stride (I64_KV_ADDR). */ + * copy within the page, BS / copy tokens entries on the arena; [0] stays + * NULL), and the i64 page stride (I64_KV_ADDR). */ rocke_value_t* lane_within_b_v; - rocke_value_t* copy_tok_b[64]; + rocke_value_t** copy_tok_b; rocke_value_t* kv_stride_blk_i64_v; /* Page ids the KV loop carries when page copies prefetch them (one per page * of a tile), with the tile index value each set belongs to: the first tile 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 720b4dae8355..01f3f2ae1ab6 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 @@ -115,6 +115,23 @@ int rocke_unified_attn_num_queries_per_kv(rocke_ir_builder_t* b, return p->num_query_heads / p->num_kv_heads; } +bool rocke_attention_check_head_counts(int num_query_heads, + int num_kv_heads, + char* reason, + size_t reason_cap) +{ + const char* fmt = NULL; + if(num_query_heads <= 0 || num_kv_heads <= 0) + fmt = "num_query_heads / num_kv_heads must be > 0 (got %d, %d)"; + else if(num_query_heads % num_kv_heads != 0) + fmt = "num_query_heads (%d) must be divisible by num_kv_heads (%d) for GQA / MQA"; + if(fmt == NULL) + return true; + if(reason != NULL && reason_cap > 0) + (void)snprintf(reason, reason_cap, fmt, num_query_heads, num_kv_heads); + return false; +} + /* Internal convenience: same value without a builder (the selectors call the * property freely; the divisibility precondition is guaranteed for any problem * that reached the spec stage). */ diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/gfx942/attention_tiled_2d_kv_descriptors_loads.cpp b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/gfx942/attention_tiled_2d_kv_descriptors_loads.cpp index 277b707ce3cc..086f4d9a9047 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/gfx942/attention_tiled_2d_kv_descriptors_loads.cpp +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/gfx942/attention_tiled_2d_kv_descriptors_loads.cpp @@ -395,9 +395,8 @@ void rocke_gfx942_attn2d_emit_preloop(rocke_gfx942_attn2d_build_ctx_t* ctx) int kv_strides[4] = {1, 1, ctx->HD * ctx->KV_BYTES, ctx->KV_BYTES}; static const char* const kv_coords[4] = {"physical_block", "token", "kv_head", "dim"}; rocke_value_t* dyn[4] = {ctx->kv_stride_blk_b_v, ctx->kv_stride_tok_b_v, NULL, NULL}; - rocke_tensor_descriptor_t* kv_base = &rocke_tensor_descriptor_naive_dynamic_mixed( - b, "paged_kv_bytes", kv_coords, 4, dyn, kv_strides) - ->base; + rocke_tensor_descriptor_t* kv_base = rocke_tensor_descriptor_naive_dynamic_mixed( + b, "paged_kv_bytes", kv_coords, 4, dyn, kv_strides); /* Single-block tile (N_BLOCKS_PER_TILE == 1, T == BS): * indirect(tile_idx -> physical_block) ; unmerge(linear_half -> token,dim). diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/gfx942/attention_tiled_2d_public_entry_glue.cpp b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/gfx942/attention_tiled_2d_public_entry_glue.cpp index 5396c333fd01..40adb961e96d 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/gfx942/attention_tiled_2d_public_entry_glue.cpp +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/gfx942/attention_tiled_2d_public_entry_glue.cpp @@ -41,6 +41,7 @@ #include "rocke/helper_helper_rocke.instances.common.attention_unified_selectors.h" /* rocke_unified_attn_tiled_2d_lds_bytes */ #include "rocke/helper_rocke.helpers.spec.h" /* rocke_kernel_name_join */ +#include "rocke/arena.h" /* rocke_arena_strdup */ #include "rocke/error_boundary.hpp" /* ckc::guard_builder boundary shim */ #include #include @@ -652,12 +653,21 @@ rocke_kernel_def_t* rocke_build_unified_attention_2d_tiled_scalar_new( { return ckc::guard_builder(b, [&]() -> rocke_kernel_def_t* { char kname[1024]; + char why[160]; if(b == NULL || spec == NULL) return NULL; - if(rocke__attn2d_kernel_name(spec, kname, sizeof(kname)) != ROCKE_OK) - return NULL; - if(rocke_ir_builder_init(b, kname) != ROCKE_OK) + /* Init before anything can fail: the lower convenience reads and frees + * `b` whatever this returns. */ + if(rocke_ir_builder_init(b, "") != ROCKE_OK) return NULL; + if(!rocke_attention_check_head_counts( + spec->num_query_heads, spec->num_kv_heads, why, sizeof(why))) + ckc::raise_status(ROCKE_ERR_VALUE, why); + if(rocke__attn2d_kernel_name(spec, kname, sizeof(kname)) != ROCKE_OK) + ckc::raise_status(ROCKE_ERR_VALUE, "tiled 2D kernel_name encode failed"); + b->kernel->name = rocke_arena_strdup(&b->arena, kname); + if(b->kernel->name == NULL) + ckc::raise_status(ROCKE_ERR_OOM, "tiled 2D kernel name OOM"); return rocke_build_unified_attention_2d_tiled_scalar(b, spec, arch); }); } diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/gfx942/attention_tiled_2d_q_load_gather.cpp b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/gfx942/attention_tiled_2d_q_load_gather.cpp index e1cc34c0d767..81135a4eb790 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/gfx942/attention_tiled_2d_q_load_gather.cpp +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/gfx942/attention_tiled_2d_q_load_gather.cpp @@ -92,9 +92,8 @@ void rocke_gfx942_attn2d_emit_q_load(rocke_gfx942_attn2d_build_ctx_t* ctx) rocke_value_t* strides[3] = {rocke_b_mul(b, ctx->num_qh_p, rocke_b_const_i32(b, ctx->HD)), NULL, NULL}; const int static_strides[3] = {1, ctx->HD, 1}; - ctx->q_desc = &rocke_tensor_descriptor_naive_dynamic_mixed( - b, "Q", coord_names, 3, strides, static_strides) - ->base; + ctx->q_desc = rocke_tensor_descriptor_naive_dynamic_mixed( + b, "Q", coord_names, 3, strides, static_strides); } if(ctx->Q_DIRECT_GLOBAL) diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/gfx942/attention_tiled_3d_gfx942_attn_tiled_3d_public_glue.cpp b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/gfx942/attention_tiled_3d_gfx942_attn_tiled_3d_public_glue.cpp index 0a5de6937f5d..6933a1049246 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/gfx942/attention_tiled_3d_gfx942_attn_tiled_3d_public_glue.cpp +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/gfx942/attention_tiled_3d_gfx942_attn_tiled_3d_public_glue.cpp @@ -47,6 +47,7 @@ #include "rocke/arena.h" /* rocke_arena_strdup */ #include "rocke/error_boundary.hpp" /* ckc::guard_builder boundary shim */ +#include "rocke/helper_helper_rocke.instances.common.attention_unified_selectors.h" /* rocke_attention_check_head_counts */ #include "rocke/helper_rocke.core.arch.h" /* rocke_archtarget_from_gfx, op_for_shape */ #include "rocke/helper_rocke.helpers.atoms.h" /* rocke_mfma_atom, make_c_warp_dstr_encoding */ #include "rocke/helper_rocke.helpers.distribution.h" /* make_static_tile_distribution */ @@ -1157,6 +1158,10 @@ rocke_status_t rocke_build_unified_attention_3d_tiled_gfx942_lower_to_llvm( return ROCKE_ERR_VALUE; } + /* The name carries the GQA ratio: refuse a head pair it cannot be formed + * for with the validator's reason, not the encoder's. */ + if(!rocke_attention_check_head_counts(spec->num_query_heads, spec->num_kv_heads, err, err_cap)) + return ROCKE_ERR_VALUE; if(rocke_unified_attention_3d_tiled_spec_kernel_name(spec, name, sizeof(name)) < 0) { rocke_attn3d_set_err_buf(err, err_cap, "lower_to_llvm: kernel_name encode failed"); diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/gfx942/attention_tiled_3d_gfx942_attn_tiled_3d_reduce.cpp b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/gfx942/attention_tiled_3d_gfx942_attn_tiled_3d_reduce.cpp index 59d68142d0a0..56e1fd522f61 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/gfx942/attention_tiled_3d_gfx942_attn_tiled_3d_reduce.cpp +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/gfx942/attention_tiled_3d_gfx942_attn_tiled_3d_reduce.cpp @@ -219,9 +219,8 @@ void rocke_gfx942_attention_tiled_3d_reduce_declare_and_prologue( = {rocke_b_mul(b, ctx->red_num_qh_p, rocke_b_const_i32(b, NUM_SEG)), NULL, NULL}; const int static_strides[3] = {1, NUM_SEG, 1}; const char* ml_coords[3] = {"token", "head", "seg"}; - ctx->ml_desc_red = &rocke_tensor_descriptor_naive_dynamic_mixed( - b, "segm_ml", ml_coords, 3, strides, static_strides) - ->base; + ctx->ml_desc_red = rocke_tensor_descriptor_naive_dynamic_mixed( + b, "segm_ml", ml_coords, 3, strides, static_strides); } { rocke_value_t* strides[4] @@ -231,18 +230,16 @@ void rocke_gfx942_attention_tiled_3d_reduce_declare_and_prologue( NULL}; const int static_strides[4] = {1, NUM_SEG * HD, HD, 1}; const char* sa_coords[4] = {"token", "head", "seg", "dim"}; - ctx->seg_acc_desc_red = &rocke_tensor_descriptor_naive_dynamic_mixed( - b, "segm_output", sa_coords, 4, strides, static_strides) - ->base; + ctx->seg_acc_desc_red = rocke_tensor_descriptor_naive_dynamic_mixed( + b, "segm_output", sa_coords, 4, strides, static_strides); } { rocke_value_t* strides[3] = {rocke_b_mul(b, ctx->red_num_qh_p, rocke_b_const_i32(b, HD)), NULL, NULL}; const int static_strides[3] = {1, HD, 1}; const char* out_coords[3] = {"token", "head", "dim"}; - ctx->out_desc_red = &rocke_tensor_descriptor_naive_dynamic_mixed( - b, "out", out_coords, 3, strides, static_strides) - ->base; + ctx->out_desc_red = rocke_tensor_descriptor_naive_dynamic_mixed( + b, "out", out_coords, 3, strides, static_strides); } /* ---- base_ml = ml_desc_red.offset(token=q_token, head=q_head, seg=0) ---- diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/gfx942/attention_tiled_3d_gfx942_attn_tiled_3d_segment_setup.cpp b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/gfx942/attention_tiled_3d_gfx942_attn_tiled_3d_segment_setup.cpp index f25fa21192c5..fc30292abe1a 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/gfx942/attention_tiled_3d_gfx942_attn_tiled_3d_segment_setup.cpp +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/gfx942/attention_tiled_3d_gfx942_attn_tiled_3d_segment_setup.cpp @@ -356,9 +356,8 @@ static void rocke__build_paged_kv_desc(rocke_gfx942_attention_tiled_3d_build_ctx rocke_value_t* strides[4] = {ctx->kv_stride_blk_b, ctx->kv_stride_tok_b, NULL, NULL}; const int static_strides[4] = {1, 1, CFG.kv_stride_h_b, CFG.KV_BYTES}; const char* coords[4] = {"physical_block", "token", "kv_head", "dim"}; - ctx->kv_base_desc = &rocke_tensor_descriptor_naive_dynamic_mixed( - B, "paged_kv_bytes", coords, 4, strides, static_strides) - ->base; + ctx->kv_base_desc = rocke_tensor_descriptor_naive_dynamic_mixed( + B, "paged_kv_bytes", coords, 4, strides, static_strides); } /* Transform chain (lines 576-602). */ @@ -484,27 +483,24 @@ void rocke_gfx942_attention_tiled_3d_emit_prologue(rocke_gfx942_attention_tiled_ = {rocke_b_mul(B, ctx->num_qh_p, rocke_b_const_i32(B, NUM_SEG)), NULL, NULL}; const int static_strides[3] = {1, NUM_SEG, 1}; const char* ml_coords[3] = {"token", "head", "seg"}; - ctx->ml_desc = &rocke_tensor_descriptor_naive_dynamic_mixed( - B, "segm_ml", ml_coords, 3, strides, static_strides) - ->base; + ctx->ml_desc = rocke_tensor_descriptor_naive_dynamic_mixed( + B, "segm_ml", ml_coords, 3, strides, static_strides); } { rocke_value_t* strides[4] = {rocke_b_mul(B, ctx->num_qh_p, rocke_b_const_i32(B, NUM_SEG * HD)), NULL, NULL, NULL}; const int static_strides[4] = {1, NUM_SEG * HD, HD, 1}; const char* sa_coords[4] = {"token", "head", "seg", "dim"}; - ctx->seg_acc_desc = &rocke_tensor_descriptor_naive_dynamic_mixed( - B, "segm_output", sa_coords, 4, strides, static_strides) - ->base; + ctx->seg_acc_desc = rocke_tensor_descriptor_naive_dynamic_mixed( + B, "segm_output", sa_coords, 4, strides, static_strides); } { rocke_value_t* strides[3] = {rocke_b_mul(B, ctx->num_qh_p, rocke_b_const_i32(B, HD)), NULL, NULL}; const int static_strides[3] = {1, HD, 1}; const char* q_coords[3] = {"token", "head", "dim"}; - ctx->q_desc = &rocke_tensor_descriptor_naive_dynamic_mixed( - B, "Q", q_coords, 3, strides, static_strides) - ->base; + ctx->q_desc = rocke_tensor_descriptor_naive_dynamic_mixed( + B, "Q", q_coords, 3, strides, static_strides); } /* seg_start_tile_pos = seg_idx*tps*T (line 375) */ diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/gfx950/attention_tiled_2d_fastkv_regp_helpers.cpp b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/gfx950/attention_tiled_2d_fastkv_regp_helpers.cpp index e4709d5d5e77..3717be914642 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/gfx950/attention_tiled_2d_fastkv_regp_helpers.cpp +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/gfx950/attention_tiled_2d_fastkv_regp_helpers.cpp @@ -29,6 +29,19 @@ ckc::raise_status(st, msg ? msg : ""); } +/* Python _NOT_GFX950: the body uses gfx950-only wide MFMA atoms and transposed + * LDS reads. True (with the reason written) for any target but gfx950; NULL arch + * is the Python default "gfx950". */ +static bool rocke__fastkv_regp_not_gfx950(const char* arch, char* reason, size_t reason_cap) +{ + if(arch == NULL || strcmp(arch, "gfx950") == 0) + return false; + if(reason != NULL && reason_cap > 0) + snprintf( + reason, reason_cap, "fastKV register-P is a gfx950 kernel; %s is not supported", arch); + return true; +} + /* gfx950 UnifiedAttention2DTiledSpec.kernel_name() (rocke/instances/gfx950/ * attention_tiled_2d.py kernel_name). The gfx950 part list is SHORTER than the * gfx942 one (no wpe, no mfma32x8/qdir/kvcp/gldlds/qgrid/cfv/... parts), so the @@ -241,6 +254,11 @@ bool rocke_gfx950_supports_fastkv_register_p_2d( if(args == NULL) return false; + /* if arch != "gfx950": return False, _NOT_GFX950 -- supports_tiled_2d also + * admits gfx942, but this body is gfx950-only. NULL arch == Python "gfx950". */ + if(rocke__fastkv_regp_not_gfx950(args->arch, reason, reason_cap)) + return false; + /* ok, reason = supports_tiled_2d(... num_warps=4, kv_storage_dtype=None, * tile_size = tile_size if tile_size is not None else 2*block_size, * arch=arch) */ @@ -296,6 +314,16 @@ rocke_kernel_def_t* rocke_build_unified_attention_2d_fastkv_register_p( if(b->status != ROCKE_OK) /* dead builder: no-op */ return NULL; + /* if arch != "gfx950": raise ValueError(_NOT_GFX950) */ + { + char why[128]; + if(rocke__fastkv_regp_not_gfx950(arch, why, sizeof(why))) + { + rocke__fastkv_regp_set_err(b, ROCKE_ERR_VALUE, why); + return NULL; + } + } + /* if not spec.use_fast_paged_kv_desc: raise ValueError(...) */ if(!spec->use_fast_paged_kv_desc) { diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/gfx950/attention_tiled_2d_kv_descriptors_loads.cpp b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/gfx950/attention_tiled_2d_kv_descriptors_loads.cpp index 4854c9b805b7..057e3195d343 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/gfx950/attention_tiled_2d_kv_descriptors_loads.cpp +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/gfx950/attention_tiled_2d_kv_descriptors_loads.cpp @@ -50,6 +50,7 @@ #include /* snprintf */ #include /* memset */ +#include "rocke/arena.h" /* rocke_arena_calloc */ #include "rocke/instance_gfx950_attention_tiled_2d_internal.h" /* ========================================================================== @@ -199,6 +200,22 @@ static bool rocke__page_copies(const rocke_gfx950_attn2d_build_ctx_t* ctx) return !ctx->FAST_PAGED_KV_DESC && !ctx->KV_FP8 && ctx->BS % rocke__copy_tokens(ctx) == 0; } +/* A zeroed table of `n` IR values on the builder arena, sized from the spec + * (the Python twins are dicts, so no fixed bound applies). NULL with the + * builder's error set on OOM. */ +static rocke_value_t** rocke__value_table(rocke_gfx950_attn2d_build_ctx_t* ctx, int n) +{ + rocke_ir_builder_t* b = ctx->b; + rocke_value_t** t = (rocke_value_t**)rocke_arena_calloc( + &b->arena, (size_t)(n > 0 ? n : 1) * sizeof(rocke_value_t*)); + if(t == NULL) + { + b->status = ROCKE_ERR_OOM; + snprintf(b->err, ROCKE_ERR_MSG_CAP, "out of memory for a %d-entry K/V table", n); + } + return t; +} + static void rocke__emit_paged_copy_offsets(rocke_gfx950_attn2d_build_ctx_t* ctx) { rocke_ir_builder_t* b = ctx->b; @@ -206,17 +223,12 @@ static void rocke__emit_paged_copy_offsets(rocke_gfx950_attn2d_build_ctx_t* ctx) const int kv_calls_per_tile = rocke__kv_calls_per_tile(ctx); rocke_value_t* lane_half_base = rocke__lane_half_base(ctx); - memset(ctx->copy_tok_b, 0, sizeof(ctx->copy_tok_b)); + const int copies_per_page = ctx->BS / copy_tokens; + ctx->kv_stride_blk_i64_v = NULL; - if(ctx->BS / copy_tokens > 64) - { - b->status = ROCKE_ERR_VALUE; - snprintf(b->err, - ROCKE_ERR_MSG_CAP, - "%d copies per page exceed the 64-entry offset table", - ctx->BS / copy_tokens); + ctx->copy_tok_b = rocke__value_table(ctx, copies_per_page); + if(ctx->copy_tok_b == NULL) return; - } rocke_value_t* c_hd0 = rocke_b_const_i32(b, ctx->HD); rocke_value_t* lane_tok = rocke_b_div(b, lane_half_base, c_hd0); rocke_value_t* c_hd1 = rocke_b_const_i32(b, ctx->HD); @@ -230,7 +242,7 @@ static void rocke__emit_paged_copy_offsets(rocke_gfx950_attn2d_build_ctx_t* ctx) ctx->lane_within_b_v = rocke_b_add(b, tok_head, dim_b); for(int call = 0; call < kv_calls_per_tile; ++call) { - const int tok_i = call % (ctx->BS / copy_tokens); + const int tok_i = call % copies_per_page; if(tok_i == 0 || ctx->copy_tok_b[tok_i] != NULL) continue; rocke_value_t* c_tok = rocke_b_const_i32(b, tok_i * copy_tokens); @@ -406,9 +418,8 @@ void rocke_gfx950_attn2d_emit_preloop(rocke_gfx950_attn2d_build_ctx_t* ctx) int kv_strides[4] = {1, 1, ctx->HD * ctx->KV_BYTES, ctx->KV_BYTES}; rocke_value_t* kv_dyn[4] = {ctx->kv_stride_blk_b_v, ctx->kv_stride_tok_b_v, NULL, NULL}; static const char* const kv_coords[4] = {"physical_block", "token", "kv_head", "dim"}; - rocke_tensor_descriptor_t* kv_base = &rocke_tensor_descriptor_naive_dynamic_mixed( - b, "paged_kv_bytes", kv_coords, 4, kv_dyn, kv_strides) - ->base; + rocke_tensor_descriptor_t* kv_base = rocke_tensor_descriptor_naive_dynamic_mixed( + b, "paged_kv_bytes", kv_coords, 4, kv_dyn, kv_strides); /* Single-block tile (N_BLOCKS_PER_TILE == 1, T == BS): * indirect(tile_idx -> physical_block) ; unmerge(linear_half -> token,dim). @@ -567,7 +578,10 @@ void rocke_gfx950_attn2d_issue_k_load_runtime(rocke_gfx950_attn2d_build_ctx_t* c if(ctx->FAST_PAGED_KV_DESC) rocke_gfx950_attn2d_fast_paged_kv_blocks(ctx, kv_tile_idx, &fast_block0, &fast_block1); - rocke_value_t* pages[64] = {NULL}; /* this tile's page ids by block */ + /* This tile's page ids by block, read on first use. */ + rocke_value_t** pages = rocke__value_table(ctx, ctx->N_BLOCKS_PER_TILE); + if(pages == NULL) + return; for(int call = 0; call < kv_calls_per_tile; ++call) { rocke_value_t* k_rsrc = ctx->k_rsrc; @@ -652,7 +666,10 @@ void rocke_gfx950_attn2d_issue_v_load_runtime(rocke_gfx950_attn2d_build_ctx_t* c if(ctx->FAST_PAGED_KV_DESC) rocke_gfx950_attn2d_fast_paged_kv_blocks(ctx, kv_tile_idx, &fast_block0, &fast_block1); - rocke_value_t* pages[64] = {NULL}; /* this tile's page ids by block */ + /* This tile's page ids by block, read on first use. */ + rocke_value_t** pages = rocke__value_table(ctx, ctx->N_BLOCKS_PER_TILE); + if(pages == NULL) + return; for(int call = 0; call < kv_calls_per_tile; ++call) { rocke_value_t* v_rsrc = ctx->v_rsrc; diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/gfx950/attention_tiled_2d_public_entry_glue.cpp b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/gfx950/attention_tiled_2d_public_entry_glue.cpp index a3a56e7a6ed0..acf6d1d308ac 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/gfx950/attention_tiled_2d_public_entry_glue.cpp +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/gfx950/attention_tiled_2d_public_entry_glue.cpp @@ -42,8 +42,10 @@ #include "rocke/instance_gfx950_attention_tiled_2d.h" #include "rocke/instance_gfx950_attention_tiled_2d_internal.h" +#include "rocke/arena.h" /* rocke_arena_strdup */ #include "rocke/error_boundary.hpp" /* ckc::guard_builder boundary shim */ #include "rocke/helper_helper_rocke.helpers.attention.h" /* binary_search_seq_idx */ +#include "rocke/helper_helper_rocke.instances.common.attention_unified_selectors.h" /* rocke_attention_check_head_counts */ #include "rocke/helper_helper_rocke.instances.gfx942.attention_tiled_2d.h" /* rocke__mfma_32x32_c_* */ #include "rocke/helper_rocke.helpers.spec.h" /* rocke_kernel_name_join */ @@ -1160,14 +1162,25 @@ static rocke_status_t rocke_g950_attn2d_kernel_name(const rocke_attention_tiled_ rocke_kernel_def_t* rocke_gfx950_build_unified_attention_2d_tiled_new( rocke_ir_builder_t* b, const rocke_attention_tiled_2d_spec_t* spec, const char* arch) { - char kname[1024]; - if(b == NULL || spec == NULL) - return NULL; - if(rocke_g950_attn2d_kernel_name(spec, kname, sizeof(kname)) != ROCKE_OK) - return NULL; - if(rocke_ir_builder_init(b, kname) != ROCKE_OK) - return NULL; - return rocke_gfx950_build_unified_attention_2d_tiled(b, spec, arch); + return ckc::guard_builder(b, [&]() -> rocke_kernel_def_t* { + char kname[1024]; + char why[160]; + if(b == NULL || spec == NULL) + return NULL; + /* Init before anything can fail: the lower convenience reads and frees + * `b` whatever this returns. */ + if(rocke_ir_builder_init(b, "") != ROCKE_OK) + return NULL; + if(!rocke_attention_check_head_counts( + spec->num_query_heads, spec->num_kv_heads, why, sizeof(why))) + ckc::raise_status(ROCKE_ERR_VALUE, why); + if(rocke_g950_attn2d_kernel_name(spec, kname, sizeof(kname)) != ROCKE_OK) + ckc::raise_status(ROCKE_ERR_VALUE, "tiled 2D kernel_name encode failed"); + b->kernel->name = rocke_arena_strdup(&b->arena, kname); + if(b->kernel->name == NULL) + ckc::raise_status(ROCKE_ERR_OOM, "tiled 2D kernel name OOM"); + return rocke_gfx950_build_unified_attention_2d_tiled(b, spec, arch); + }); } /* ===================================================================== * diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/gfx950/attention_tiled_2d_q_load_gather.cpp b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/gfx950/attention_tiled_2d_q_load_gather.cpp index b1b880569d9a..ae3d0609cf0f 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/gfx950/attention_tiled_2d_q_load_gather.cpp +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/gfx950/attention_tiled_2d_q_load_gather.cpp @@ -118,9 +118,8 @@ void rocke_gfx950_attn2d_emit_q_load(rocke_gfx950_attn2d_build_ctx_t* ctx) rocke_value_t* strides[3] = {rocke_b_mul(b, ctx->num_qh_p, rocke_b_const_i32(b, ctx->HD)), NULL, NULL}; const int static_strides[3] = {1, ctx->HD, 1}; - ctx->q_desc = &rocke_tensor_descriptor_naive_dynamic_mixed( - b, "Q", coord_names, 3, strides, static_strides) - ->base; + ctx->q_desc = rocke_tensor_descriptor_naive_dynamic_mixed( + b, "Q", coord_names, 3, strides, static_strides); } for(li = 0; li < Q_VECS_PER_THREAD; ++li) diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/gfx950/attention_tiled_3d_public_glue.cpp b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/gfx950/attention_tiled_3d_public_glue.cpp index bbfee08284f1..257cf526a7d7 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/gfx950/attention_tiled_3d_public_glue.cpp +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/gfx950/attention_tiled_3d_public_glue.cpp @@ -72,6 +72,7 @@ * _default() definitions below -- the gfx942 TU owns them. The structs are * byte-identical so the single shared definition is correct for both. */ #include "rocke/error_boundary.hpp" /* ckc::guard_builder boundary shim */ +#include "rocke/helper_helper_rocke.instances.common.attention_unified_selectors.h" /* rocke_attention_check_head_counts */ #include "rocke/instance_gfx942_attention_tiled_3d.h" #include "rocke/instance_gfx950_attention_tiled_3d.h" #include "rocke/instance_gfx950_attention_tiled_3d_internal.h" @@ -1245,6 +1246,10 @@ rocke_status_t rocke_build_unified_attention_3d_tiled_gfx950_lower_to_llvm( return ROCKE_ERR_VALUE; } + /* The name carries the GQA ratio: refuse a head pair it cannot be formed + * for with the validator's reason, not the encoder's. */ + if(!rocke_attention_check_head_counts(spec->num_query_heads, spec->num_kv_heads, err, err_cap)) + return ROCKE_ERR_VALUE; if(rocke_gfx950_unified_attention_3d_tiled_spec_kernel_name(spec, name, sizeof(name)) < 0) { rocke_attn3d950_set_err_buf(err, err_cap, "lower_to_llvm: kernel_name encode failed"); diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/gfx950/attention_tiled_3d_reduce.cpp b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/gfx950/attention_tiled_3d_reduce.cpp index 199c070ebeaf..14a506cf91aa 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/gfx950/attention_tiled_3d_reduce.cpp +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/gfx950/attention_tiled_3d_reduce.cpp @@ -222,27 +222,24 @@ void rocke_gfx950_attention_tiled_3d_reduce_declare_and_prologue( rocke_value_t* ml_dyn[3] = {rocke_b_mul(b, ctx->num_qh_p, rocke_b_const_i32(b, NUM_SEG)), NULL, NULL}; const int ml_static[3] = {1, NUM_SEG, 1}; - ctx->ml_desc_red = &rocke_tensor_descriptor_naive_dynamic_mixed( - b, "segm_ml", ml_coords, 3, ml_dyn, ml_static) - ->base; + ctx->ml_desc_red = rocke_tensor_descriptor_naive_dynamic_mixed( + b, "segm_ml", ml_coords, 3, ml_dyn, ml_static); } { static const char* const sa_coords[4] = {"token", "head", "seg", "dim"}; rocke_value_t* sa_dyn[4] = {rocke_b_mul(b, ctx->num_qh_p, rocke_b_const_i32(b, NUM_SEG * HD)), NULL, NULL, NULL}; const int sa_static[4] = {1, NUM_SEG * HD, HD, 1}; - ctx->seg_acc_desc_red = &rocke_tensor_descriptor_naive_dynamic_mixed( - b, "segm_output", sa_coords, 4, sa_dyn, sa_static) - ->base; + ctx->seg_acc_desc_red = rocke_tensor_descriptor_naive_dynamic_mixed( + b, "segm_output", sa_coords, 4, sa_dyn, sa_static); } { static const char* const out_coords[3] = {"token", "head", "dim"}; rocke_value_t* out_dyn[3] = {rocke_b_mul(b, ctx->num_qh_p, rocke_b_const_i32(b, HD)), NULL, NULL}; const int out_static[3] = {1, HD, 1}; - ctx->out_desc_red = &rocke_tensor_descriptor_naive_dynamic_mixed( - b, "out", out_coords, 3, out_dyn, out_static) - ->base; + ctx->out_desc_red = rocke_tensor_descriptor_naive_dynamic_mixed( + b, "out", out_coords, 3, out_dyn, out_static); } /* ---- base_ml = ml_desc_red.offset(token=q_token, head=q_head, seg=0) ---- diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/gfx950/attention_tiled_3d_segment_prologue.cpp b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/gfx950/attention_tiled_3d_segment_prologue.cpp index 79a3da17a344..691050c530e7 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/gfx950/attention_tiled_3d_segment_prologue.cpp +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/gfx950/attention_tiled_3d_segment_prologue.cpp @@ -415,18 +415,16 @@ void rocke_gfx950_attention_tiled_3d_emit_prologue(rocke_gfx950_attention_tiled_ rocke_value_t* ml_dyn[3] = {rocke_b_mul(B, ctx->num_qh_p, rocke_b_const_i32(B, NUM_SEG)), NULL, NULL}; const int ml_static[3] = {1, NUM_SEG, 1}; - ctx->ml_desc = &rocke_tensor_descriptor_naive_dynamic_mixed( - B, "segm_ml", ml_coords, 3, ml_dyn, ml_static) - ->base; + ctx->ml_desc = rocke_tensor_descriptor_naive_dynamic_mixed( + B, "segm_ml", ml_coords, 3, ml_dyn, ml_static); } { static const char* const sa_coords[4] = {"token", "head", "seg", "dim"}; rocke_value_t* sa_dyn[4] = {rocke_b_mul(B, ctx->num_qh_p, rocke_b_const_i32(B, NUM_SEG * HD)), NULL, NULL, NULL}; const int sa_static[4] = {1, NUM_SEG * HD, HD, 1}; - ctx->seg_acc_desc = &rocke_tensor_descriptor_naive_dynamic_mixed( - B, "segm_output", sa_coords, 4, sa_dyn, sa_static) - ->base; + ctx->seg_acc_desc = rocke_tensor_descriptor_naive_dynamic_mixed( + B, "segm_output", sa_coords, 4, sa_dyn, sa_static); } { static const char* const q_coords[3] = {"token", "head", "dim"}; @@ -434,8 +432,7 @@ void rocke_gfx950_attention_tiled_3d_emit_prologue(rocke_gfx950_attention_tiled_ = {rocke_b_mul(B, ctx->num_qh_p, rocke_b_const_i32(B, HD)), NULL, NULL}; const int q_static[3] = {1, HD, 1}; ctx->q_desc - = &rocke_tensor_descriptor_naive_dynamic_mixed(B, "Q", q_coords, 3, q_dyn, q_static) - ->base; + = rocke_tensor_descriptor_naive_dynamic_mixed(B, "Q", q_coords, 3, q_dyn, q_static); } /* seg_start_tile_pos = seg_idx*tps*T (line 415) */ @@ -573,9 +570,8 @@ void rocke_gfx950_attention_tiled_3d_emit_async_infra( const int strides[4] = {1, 1, CFG.kv_stride_h_b, CFG.KV_BYTES}; rocke_value_t* dyn[4] = {ctx->kv_stride_blk_b, ctx->kv_stride_tok_b, NULL, NULL}; static const char* const coords[4] = {"physical_block", "token", "kv_head", "dim"}; - rocke_tensor_descriptor_t* base = &rocke_tensor_descriptor_naive_dynamic_mixed( - B, "paged_kv_bytes", coords, 4, dyn, strides) - ->base; + rocke_tensor_descriptor_t* base = rocke_tensor_descriptor_naive_dynamic_mixed( + B, "paged_kv_bytes", coords, 4, dyn, strides); const rocke_transform_t* chain[2]; const char* into_td[2] = {"token", "dim"}; const int dims_td[2] = {T, HD}; diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/transforms.py b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/transforms.py index 570b57323e48..95589cc8406f 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/transforms.py +++ b/dnn-providers/hip-kernel-provider/rocke/platform/python/rocke/helpers/transforms.py @@ -1543,6 +1543,18 @@ def _run_chain( remaining = next_remaining return coords + def _term(self, b: IRBuilder, i: int, value: Value) -> Value: + """Base coord ``i``'s i32 offset term: ``value`` times its stride, with no + multiply for a stride of 1. :class:`DynamicTensorDescriptor` overrides + the stride source; every offset/move loop goes through this hook.""" + stride = self.base_strides[i] + return value if stride == 1 else b.mul(value, b.const_i32(stride)) + + def _term_i64(self, b: IRBuilder, i: int, wide: Value) -> Value: + """Base coord ``i``'s i64 term: the already-widened ``wide`` times its + stride (always multiplied, as the i64 paths do).""" + return b.mul(wide, b.const_i64(int(self.base_strides[i]))) + def offset( self, b: IRBuilder, @@ -1558,17 +1570,14 @@ def offset( coords = self._run_chain(b, upper_values) offset: Optional[Value] = None valid: Optional[Value] = None - for name, stride in zip(self.base_names, self.base_strides): + for i, name in enumerate(self.base_names): if name not in coords: raise ValueError( f"after chain, base coord {name!r} not in {sorted(coords.keys())}" ) c = coords[name] valid = _and(b, valid, c.valid) - if stride == 1: - term = c.value - else: - term = b.mul(c.value, b.const_i32(stride)) + term = self._term(b, i, c.value) offset = term if offset is None else b.add(offset, term) if offset is None: offset = b.const_i32(0) @@ -1593,7 +1602,7 @@ def offset_i64_split( base_i64: Optional[Value] = None within: Optional[Value] = None valid: Optional[Value] = None - for name, stride in zip(self.base_names, self.base_strides): + for i, name in enumerate(self.base_names): if name not in coords: raise ValueError( f"after chain, base coord {name!r} not in {sorted(coords.keys())}" @@ -1605,9 +1614,9 @@ def offset_i64_split( # descriptor (uniform per wave), so pin the (wave-uniform) # block id to an SGPR before widening. base_val = b.to_sgpr_u32(c.value) - base_i64 = b.mul(b.zext(base_val, I64), b.const_i64(int(stride))) + base_i64 = self._term_i64(b, i, b.zext(base_val, I64)) else: - term = c.value if stride == 1 else b.mul(c.value, b.const_i32(stride)) + term = self._term(b, i, c.value) within = term if within is None else b.add(within, term) if base_i64 is None: raise ValueError( @@ -1635,7 +1644,7 @@ def offset_i64( coords = self._run_chain(b, upper_values) off: Optional[Value] = None valid: Optional[Value] = None - for name, stride in zip(self.base_names, self.base_strides): + for i, name in enumerate(self.base_names): if name not in coords: raise ValueError( f"after chain, base coord {name!r} not in {sorted(coords.keys())}" @@ -1643,10 +1652,9 @@ def offset_i64( c = coords[name] valid = _and(b, valid, c.valid) if name == i64_coord: - term = b.mul(b.zext(c.value, I64), b.const_i64(int(stride))) + term = self._term_i64(b, i, b.zext(c.value, I64)) else: - t32 = c.value if stride == 1 else b.mul(c.value, b.const_i32(stride)) - term = b.zext(t32, I64) + term = b.zext(self._term(b, i, c.value), I64) off = term if off is None else b.add(off, term) if off is None: off = b.const_i64(0) @@ -1737,14 +1745,13 @@ def move( delta, new = self._run_chain_delta(b, old_values, delta_values) off_delta: Optional[Value] = None valid: Optional[Value] = None - for name, stride in zip(self.base_names, self.base_strides): + for i, name in enumerate(self.base_names): if name not in delta: raise ValueError( f"after chain, base coord {name!r} not in {sorted(delta.keys())}" ) valid = _and(b, valid, new[name].valid) - dv = delta[name].value - term = dv if stride == 1 else b.mul(dv, b.const_i32(stride)) + term = self._term(b, i, delta[name].value) off_delta = term if off_delta is None else b.add(off_delta, term) if off_delta is None: return prev_offset, valid @@ -1771,121 +1778,13 @@ def _term(self, b: IRBuilder, i: int, value: Value) -> Value: stride_v = self.dynamic_strides[i] if stride_v is not None: return b.mul(value, stride_v) - stride = self.base_strides[i] - return value if stride == 1 else b.mul(value, b.const_i32(stride)) - - def offset( - self, - b: IRBuilder, - **upper_values: Value, - ) -> Tuple[Value, Optional[Value]]: - coords = self._run_chain(b, upper_values) - off: Optional[Value] = None - valid: Optional[Value] = None - for i, name in enumerate(self.base_names): - if name not in coords: - raise ValueError( - f"after chain, base coord {name!r} not in {sorted(coords.keys())}" - ) - c = coords[name] - valid = _and(b, valid, c.valid) - term = self._term(b, i, c.value) - off = term if off is None else b.add(off, term) - if off is None: - off = b.const_i32(0) - return off, valid + return super()._term(b, i, value) - def offset_i64_split( - self, - b: IRBuilder, - base_coord: str, - **upper_values: Value, - ) -> Tuple[Value, Value, Optional[Value]]: - """:meth:`TensorDescriptor.offset_i64_split` with runtime strides; the - ``base_coord`` term widens a runtime stride to i64.""" - coords = self._run_chain(b, upper_values) - base_i64: Optional[Value] = None - within: Optional[Value] = None - valid: Optional[Value] = None - for i, name in enumerate(self.base_names): - if name not in coords: - raise ValueError( - f"after chain, base coord {name!r} not in {sorted(coords.keys())}" - ) - c = coords[name] - valid = _and(b, valid, c.valid) - if name == base_coord: - base_w = b.zext(b.to_sgpr_u32(c.value), I64) - stride_v = self.dynamic_strides[i] - if stride_v is not None: - base_i64 = b.mul(base_w, b.zext(stride_v, I64)) - else: - base_i64 = b.mul(base_w, b.const_i64(int(self.base_strides[i]))) - else: - term = self._term(b, i, c.value) - within = term if within is None else b.add(within, term) - if base_i64 is None: - raise ValueError( - f"offset_i64_split: base_coord {base_coord!r} not among " - f"base coords {list(self.base_names)}" - ) - if within is None: - within = b.const_i32(0) - return base_i64, within, valid - - def offset_i64( - self, - b: IRBuilder, - i64_coord: str, - **upper_values: Value, - ) -> Tuple[Value, Optional[Value]]: - """:meth:`TensorDescriptor.offset_i64` with runtime strides.""" - coords = self._run_chain(b, upper_values) - off: Optional[Value] = None - valid: Optional[Value] = None - for i, name in enumerate(self.base_names): - if name not in coords: - raise ValueError( - f"after chain, base coord {name!r} not in {sorted(coords.keys())}" - ) - c = coords[name] - valid = _and(b, valid, c.valid) - if name == i64_coord: - wide = b.zext(c.value, I64) - stride_v = self.dynamic_strides[i] - if stride_v is not None: - term = b.mul(wide, b.zext(stride_v, I64)) - else: - term = b.mul(wide, b.const_i64(int(self.base_strides[i]))) - else: - term = b.zext(self._term(b, i, c.value), I64) - off = term if off is None else b.add(off, term) - if off is None: - off = b.const_i64(0) - return off, valid - - def move( - self, - b: IRBuilder, - prev_offset: Value, - old_values: Dict[str, Value], - delta_values: Dict[str, Value], - ) -> Tuple[Value, Optional[Value]]: - """:meth:`TensorDescriptor.move` with runtime strides.""" - delta, new = self._run_chain_delta(b, old_values, delta_values) - off_delta: Optional[Value] = None - valid: Optional[Value] = None - for i, name in enumerate(self.base_names): - if name not in delta: - raise ValueError( - f"after chain, base coord {name!r} not in {sorted(delta.keys())}" - ) - valid = _and(b, valid, new[name].valid) - term = self._term(b, i, delta[name].value) - off_delta = term if off_delta is None else b.add(off_delta, term) - if off_delta is None: - return prev_offset, valid - return b.add(prev_offset, off_delta), valid + def _term_i64(self, b: IRBuilder, i: int, wide: Value) -> Value: + stride_v = self.dynamic_strides[i] + if stride_v is not None: + return b.mul(wide, b.zext(stride_v, I64)) + return super()._term_i64(b, i, wide) @classmethod def create( @@ -1906,15 +1805,23 @@ def create( raise ValueError("coord_names and strides length mismatch") coord_names = tuple(coord_names) dynamic = tuple(s if isinstance(s, Value) else None for s in strides) - # operator.index rejects a float or string stride instead of truncating it. - static = tuple( - 1 if isinstance(s, Value) else operator.index(s) for s in strides - ) + static = [] + for i, s in enumerate(strides): + if isinstance(s, Value): + static.append(1) + continue + # operator.index rejects a float or string stride instead of truncating it. + try: + static.append(operator.index(s)) + except TypeError: + raise TypeError( + f"strides[{i}] must be a Value or an int, got {type(s).__name__}" + ) from None return cls( name=name, base_names=coord_names, base_lengths=tuple(0 for _ in coord_names), - base_strides=static, + base_strides=tuple(static), chain=(), upper_names=coord_names, dynamic_strides=dynamic, diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_transforms_dynamic.py b/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_transforms_dynamic.py index 2bb0094e1330..5b90ba8edfe8 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_transforms_dynamic.py +++ b/dnn-providers/hip-kernel-provider/rocke/platform/tests/core/test_transforms_dynamic.py @@ -621,7 +621,7 @@ def test_missing_chain_input_rejected(self): def test_create_rejects_a_non_integer_stride(self, bad): b = _builder() (s0,) = _params(b, "s0") - with pytest.raises(TypeError): + with pytest.raises(TypeError, match=r"strides\[1\]"): DynamicTensorDescriptor.create( "t", coord_names=("a", "b"), strides=(s0, bad) ) diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/tests/instances/parity/dynamic_helpers_emit.c b/dnn-providers/hip-kernel-provider/rocke/platform/tests/instances/parity/dynamic_helpers_emit.c index 0d5cbfdf1f15..c35a45c88620 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/tests/instances/parity/dynamic_helpers_emit.c +++ b/dnn-providers/hip-kernel-provider/rocke/platform/tests/instances/parity/dynamic_helpers_emit.c @@ -337,6 +337,37 @@ static void build_unmerge_mixed(rocke_ir_builder_t* b) rocke_b_ret(b); } +/* Python build_mixed_strides: create() with a Value, an int and an int-1 + * stride, through offset and offset_i64_split. */ +static void build_mixed_strides(rocke_ir_builder_t* b) +{ + rocke_value_t* o = ptr_param(b, "O", rocke_i32(), false); + rocke_value_t* o64 = ptr_param(b, "O64", rocke_i64(), false); + rocke_value_t* p_sb = rocke_b_param(b, "p_sB", rocke_i32(), NULL); + rocke_value_t* tid = rocke_b_thread_id_x(b); + + const char* base_names[3] = {"blk", "tok", "d"}; + rocke_value_t* dyn[3] = {p_sb, NULL, NULL}; + const int stat[3] = {0, 64, 1}; + rocke_tensor_descriptor_t* desc + = rocke_tensor_descriptor_naive_dynamic_mixed(b, "kv_mixed", base_names, 3, dyn, stat); + + rocke_value_t* in_values[3] = {tid, tid, tid}; + rocke_value_t* off = NULL; + rocke_value_t* valid = NULL; + if(!rocke_transforms_descriptor_offset(b, desc, base_names, in_values, 3, &off, &valid)) + return; + rocke_b_global_store(b, o, tid, off, 1); + rocke_value_t* base_i64 = NULL; + rocke_value_t* within = NULL; + if(!rocke_transforms_descriptor_offset_i64_split( + b, desc, "blk", base_names, in_values, 3, &base_i64, &within, &valid)) + return; + rocke_b_global_store(b, o64, tid, base_i64, 1); + rocke_b_global_store(b, o, within, tid, 1); + rocke_b_ret(b); +} + static void build_pad_dynamic(rocke_ir_builder_t* b) { rocke_value_t* x = ptr_param(b, "X", rocke_f16(), true); @@ -431,6 +462,7 @@ static const build_fn_t CONFIGS[] = { build_unmerge_mixed, build_chiplet_runtime_default, build_chiplet_runtime_chunk, + build_mixed_strides, }; static const int NUM_CONFIGS = (int)(sizeof(CONFIGS) / sizeof(CONFIGS[0])); diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/tests/instances/parity/dynamic_helpers_emit.py b/dnn-providers/hip-kernel-provider/rocke/platform/tests/instances/parity/dynamic_helpers_emit.py index c2dc9aa0f6b4..292b7e4136a2 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/tests/instances/parity/dynamic_helpers_emit.py +++ b/dnn-providers/hip-kernel-provider/rocke/platform/tests/instances/parity/dynamic_helpers_emit.py @@ -18,7 +18,7 @@ # under different names. # # arch = gfx950, llvm_flavor = AUTO, matching the C side. -from rocke.core.ir import F16, F32, I32, IRBuilder, KernelDef, PtrType +from rocke.core.ir import F16, F32, I32, I64, IRBuilder, KernelDef, PtrType from rocke.helpers.grid import chiplet_transform_chunked_runtime from rocke.helpers.pipeline import SoftwarePipeline from rocke.helpers.schedule import SchedulePolicy @@ -207,6 +207,28 @@ def build_unmerge_mixed(b: IRBuilder) -> None: b.ret() +def build_mixed_strides(b: IRBuilder) -> None: + """DynamicTensorDescriptor.create with a Value, an int and an int-1 stride. + + The tiled attention K/V descriptor's shape: the page stride is a kernarg, + the token stride a constant, the head dim contiguous. Covers both stride + hooks: offset (i32 terms) and offset_i64_split (the i64 page term). + """ + o = _ptr(b, "O", I32) + o64 = _ptr(b, "O64", I64) + p_sb = b.param("p_sB", I32) + tid = b.thread_id_x() + desc = DynamicTensorDescriptor.create( + "kv_mixed", coord_names=("blk", "tok", "d"), strides=(p_sb, 64, 1) + ) + off, _ = desc.offset(b, blk=tid, tok=tid, d=tid) + b.global_store(o, tid, off) + base_i64, within, _ = desc.offset_i64_split(b, "blk", blk=tid, tok=tid, d=tid) + b.global_store(o64, tid, base_i64) + b.global_store(o, within, tid) + b.ret() + + def build_pad_dynamic(b: IRBuilder) -> None: """embed_dynamic with every argument a Value, plus pad_dynamic variants. @@ -318,6 +340,8 @@ def build(b: IRBuilder) -> None: _chiplet_runtime(explicit_chunk=False), # 9: chiplet_transform_chunked_runtime with an explicit runtime chunk. _chiplet_runtime(explicit_chunk=True), + # 10: DynamicTensorDescriptor with Value / int / int-1 strides, i32 and i64 paths. + build_mixed_strides, ] diff --git a/dnn-providers/hip-kernel-provider/src/engines/kernel_ingestor_engine/descriptors/rocKE/gfx950_attention_unified_2d/README.md b/dnn-providers/hip-kernel-provider/src/engines/kernel_ingestor_engine/descriptors/rocKE/gfx950_attention_unified_2d/README.md index f9d7280c376c..a10ae6fc94cc 100644 --- a/dnn-providers/hip-kernel-provider/src/engines/kernel_ingestor_engine/descriptors/rocKE/gfx950_attention_unified_2d/README.md +++ b/dnn-providers/hip-kernel-provider/src/engines/kernel_ingestor_engine/descriptors/rocKE/gfx950_attention_unified_2d/README.md @@ -26,8 +26,9 @@ plan's workspace (`packs/Gfx950AttentionUnified2DGeometry.hpp`). The full space's wide 32x32 geometry (4 warps x 32 rows, KV tile 64) won no graph and is not shipped. -- Batch capacity 511: `num_seqs` 511 bakes 9 binary-search iterations, which resolve at most - 2^9 - 1 sequences. The engine declines a graph with a larger batch. +- Batch: any. The kernel takes the batch as the `num_seqs` kernarg and derives its `cu_q` + binary-search trip count from it at runtime, so no batch is baked. The KMD still declares + `num_seqs` (511); the engine does not read it. - Decode (seqlen_q 1) whose GQA group divides 16 is left to `hipkernel:Gfx950AttentionUnified3D`. - Fixed in every variant: no sliding window, no sinks, no ALiBi, no softcap, no fp8, 32-bit K/V addressing (K/V larger than 0x7FFF0000 bytes is declined). diff --git a/dnn-providers/hip-kernel-provider/src/engines/kernel_ingestor_engine/packs/Gfx950AttentionUnified2DGeometry.hpp b/dnn-providers/hip-kernel-provider/src/engines/kernel_ingestor_engine/packs/Gfx950AttentionUnified2DGeometry.hpp index a8405b5968d2..590f6541882e 100644 --- a/dnn-providers/hip-kernel-provider/src/engines/kernel_ingestor_engine/packs/Gfx950AttentionUnified2DGeometry.hpp +++ b/dnn-providers/hip-kernel-provider/src/engines/kernel_ingestor_engine/packs/Gfx950AttentionUnified2DGeometry.hpp @@ -14,17 +14,17 @@ /** * @file Gfx950AttentionUnified2DGeometry.hpp * @brief What gfx950 unified-attention 2D restates from its Python builder: the geometry - * rules, and the launch grid, batch capacity and paged-KV metadata of - * UnifiedAttentionGeometry.hpp bound to the gfx950 traits. + * rules, and the launch grid and paged-KV metadata of UnifiedAttentionGeometry.hpp + * bound to the gfx950 traits. * * A rocKE kernel is launched from Python; an ingestor engine relaunches the same binary * from C++. Nothing compares the two halves, and a mismatch does not fault: the kernel * computes something else. Everything here is a pure function of a candidate's KMD fields * and the graph's dimensions, so it is testable with no HIP context. * - * The arch-neutral contract (identity block table over a contiguous BSHD K/V, the batch - * capacity a baked `num_seqs` gives, the grid) lives in UnifiedAttentionGeometry.hpp; this - * header adds only what gfx950's builder fixes: its traits and its geometry legality. + * The arch-neutral contract (identity block table over a contiguous BSHD K/V, the grid) + * lives in UnifiedAttentionGeometry.hpp; this header adds only what gfx950's builder fixes: + * its traits and its geometry legality. */ namespace hip_kernel_provider::kernel_ingestor_engine { @@ -137,18 +137,6 @@ inline bool isSupportedGfx950Unified2DGeometry(const Gfx950Unified2DGeometryFiel return !g.useTransposedQk32x32 || g.useMfma32x32; } -/// `unifiedAttentionBinarySearchIters`, named for this pack. -inline int64_t gfx950Unified2DBinarySearchIters(int64_t numSeqs) -{ - return unifiedAttentionBinarySearchIters(numSeqs); -} - -/// `unifiedAttentionBatchCapacity`: 511 -> 9 iterations -> 511 sequences. -inline int64_t gfx950Unified2DBatchCapacity(int64_t numSeqs) -{ - return unifiedAttentionBatchCapacity(numSeqs); -} - using Gfx950Unified2DLaunchGeometry = UnifiedAttentionLaunchGeometry; /// `unifiedAttention2DLaunchGeometry` with the gfx950 traits: kv-major grid diff --git a/dnn-providers/hip-kernel-provider/src/engines/kernel_ingestor_engine/packs/Gfx950AttentionUnified2DNative.cpp b/dnn-providers/hip-kernel-provider/src/engines/kernel_ingestor_engine/packs/Gfx950AttentionUnified2DNative.cpp index 3eeea72d3a69..2062a3deec9c 100644 --- a/dnn-providers/hip-kernel-provider/src/engines/kernel_ingestor_engine/packs/Gfx950AttentionUnified2DNative.cpp +++ b/dnn-providers/hip-kernel-provider/src/engines/kernel_ingestor_engine/packs/Gfx950AttentionUnified2DNative.cpp @@ -77,7 +77,9 @@ constexpr std::string_view DISPATCH_SYMBOL = "hipkernel.gfx950_attention_unified // KMD fields (14-field schema). Fields a shipped variant holds at one value across the whole // catalog (waves_per_eu, kv_storage_dtype, the sink/alibi/qq-bias/softcap switches) are absent: // a variant that moved one of them would have to add it here first, or the candidates collide -// on the catalog key and the loader drops one. +// on the catalog key and the loader drops one. `num_seqs` is declared but not read: the kernel +// takes the batch as a kernarg and derives its cu_q search trip count from it, so no batch is +// baked. constexpr std::string_view DTYPE_FIELD = "dtype"; constexpr std::string_view HEAD_SIZE_FIELD = "head_size"; constexpr std::string_view NUM_QUERY_HEADS_FIELD = "num_query_heads"; @@ -85,7 +87,6 @@ constexpr std::string_view NUM_KV_HEADS_FIELD = "num_kv_heads"; constexpr std::string_view BLOCK_SIZE_FIELD = "block_size"; constexpr std::string_view MASK_TYPE_FIELD = "mask_type"; constexpr std::string_view SLIDING_WINDOW_FIELD = "sliding_window"; -constexpr std::string_view NUM_SEQS_FIELD = "num_seqs"; constexpr std::string_view TILE_SIZE_FIELD = "tile_size"; constexpr std::string_view NUM_WARPS_FIELD = "num_warps"; constexpr std::string_view BLOCK_M_PER_WARP_FIELD = "block_m_per_warp"; @@ -241,7 +242,6 @@ struct AttentionUnified2DCandidate int64_t blockSize = 0; int64_t maskType = 0; int64_t slidingWindow = 0; - int64_t numSeqs = 0; int64_t tileSize = 0; int64_t numWarps = 0; int64_t blockMPerWarp = 0; @@ -266,15 +266,13 @@ std::optional candidateFor(const KernelDefinition& const auto blockSize = integerMetadata(kernel, BLOCK_SIZE_FIELD); const auto maskType = integerMetadata(kernel, MASK_TYPE_FIELD); const auto slidingWindow = integerMetadata(kernel, SLIDING_WINDOW_FIELD); - const auto numSeqs = integerMetadata(kernel, NUM_SEQS_FIELD); const auto tileSize = integerMetadata(kernel, TILE_SIZE_FIELD); const auto numWarps = integerMetadata(kernel, NUM_WARPS_FIELD); const auto blockMPerWarp = integerMetadata(kernel, BLOCK_M_PER_WARP_FIELD); const auto useMfma32x32 = integerMetadata(kernel, USE_MFMA_32X32_FIELD); const auto useTransposedQk32x32 = integerMetadata(kernel, USE_TRANSPOSED_QK_32X32_FIELD); if(!headSize || !numQueryHeads || !numKvHeads || !blockSize || !maskType || !slidingWindow - || !numSeqs || !tileSize || !numWarps || !blockMPerWarp || !useMfma32x32 - || !useTransposedQk32x32) + || !tileSize || !numWarps || !blockMPerWarp || !useMfma32x32 || !useTransposedQk32x32) { return std::nullopt; } @@ -293,9 +291,7 @@ std::optional candidateFor(const KernelDefinition& { return std::nullopt; } - // A capacity of 0 or less would make the kernel run its unspecialized 32-iteration - // search; the catalog always states one, so a record without it is malformed. - if(*numSeqs <= 0 || *useMfma32x32 < 0 || *useMfma32x32 > 1 || *useTransposedQk32x32 < 0 + if(*useMfma32x32 < 0 || *useMfma32x32 > 1 || *useTransposedQk32x32 < 0 || *useTransposedQk32x32 > 1) { return std::nullopt; @@ -318,7 +314,6 @@ std::optional candidateFor(const KernelDefinition& *blockSize, *maskType, *slidingWindow, - *numSeqs, *tileSize, *numWarps, *blockMPerWarp, @@ -335,10 +330,10 @@ std::optional candidateFor(const KernelDefinition& * what this checks against the graph: * * - dtype, head_size, num_query_heads, num_kv_heads, the mask kind and sliding window; - * - block_size, which must divide seqlen_kv (the identity cache is whole blocks); - * - the batch CAPACITY (`num_seqs`): the kernel's binary search resolves at most - * 2^iters - 1 sequences, so a larger batch reads another sequence's lengths without a - * fault. + * - block_size, which must divide seqlen_kv (the identity cache is whole blocks). + * + * The batch is not bounded here: the kernel derives its cu_q search trip count from the + * `num_seqs` kernarg, so one binary serves every batch the identity table can index. */ bool kernelMatches(const MatchContext& context, const BoundTokens& bound, @@ -396,13 +391,8 @@ bool kernelMatches(const MatchContext& context, // The identity cache must be whole blocks of THIS candidate's size, and small enough for // the table the kernel indexes in 32 bits. - if(!gfx950Unified2DWorkspaceLayout(problem.batch, problem.seqLenKv, candidate->blockSize) - .has_value()) - { - return false; - } - - return problem.batch <= gfx950Unified2DBatchCapacity(candidate->numSeqs); + return gfx950Unified2DWorkspaceLayout(problem.batch, problem.seqLenKv, candidate->blockSize) + .has_value(); } /// Scores for the cold ranking. The baseline geometry scores `BASELINE_SCORE`; every @@ -599,15 +589,6 @@ Gfx950Unified2DWorkspaceLayout requireLayout(const AttentionUnified2DCandidate& const KernelDefinition& kernel) { const std::string name = toString(kernel.kernelId); - const auto capacity = gfx950Unified2DBatchCapacity(candidate.numSeqs); - if(problem.batch > capacity) - { - throw hipdnn_plugin_sdk::HipdnnPluginException( - HIPDNN_PLUGIN_STATUS_BAD_PARAM, - "gfx950 attention_unified_2d: kernel '" + name + "' serves at most " - + std::to_string(capacity) + " sequences and the graph has " - + std::to_string(problem.batch)); - } const auto layout = gfx950Unified2DWorkspaceLayout(problem.batch, problem.seqLenKv, candidate.blockSize); if(!layout.has_value()) diff --git a/dnn-providers/hip-kernel-provider/src/engines/kernel_ingestor_engine/packs/UnifiedAttentionGeometry.hpp b/dnn-providers/hip-kernel-provider/src/engines/kernel_ingestor_engine/packs/UnifiedAttentionGeometry.hpp index 55abeebd5681..dc79c07a5c4e 100644 --- a/dnn-providers/hip-kernel-provider/src/engines/kernel_ingestor_engine/packs/UnifiedAttentionGeometry.hpp +++ b/dnn-providers/hip-kernel-provider/src/engines/kernel_ingestor_engine/packs/UnifiedAttentionGeometry.hpp @@ -16,8 +16,8 @@ /** * @file UnifiedAttentionGeometry.hpp * @brief The arch-neutral half of the unified-attention (vLLM `unified_attention`) packs: - * launch grid, batch capacity and the identity paged-KV metadata a contiguous K/V - * graph is served with, parameterized by the target's UnifiedAttentionArchTraits. + * launch grid and the identity paged-KV metadata a contiguous K/V graph is served + * with, parameterized by the target's UnifiedAttentionArchTraits. * * rocKE builds the unified-attention kernels per architecture, but the host contract around * them is shared: the kernel reads a paged K/V cache through `block_tables`, per-sequence @@ -95,41 +95,6 @@ inline bool isSupportedUnifiedAttentionBlockSize(const UnifiedAttentionArchTrait return detail::isOneOf(traits.blockSizes, blockSize); } -/** - * @brief Binary-search trip count a kernel was compiled with for a baked `num_seqs`. - * - * `binary_search_seq_idx_iters`: `max(1, ceil(log2(num_seqs + 1)))`, and 32 for a spec that - * declares no `num_seqs`. The same on every architecture. - */ -inline int64_t unifiedAttentionBinarySearchIters(int64_t numSeqs) -{ - constexpr int64_t UNSPECIALIZED_ITERS = 32; - if(numSeqs <= 0) - { - return UNSPECIALIZED_ITERS; - } - int64_t iters = 1; - // Smallest iters with 2^iters >= numSeqs + 1; bounded by the shift width anyway. - while(iters < UNSPECIALIZED_ITERS && (int64_t{1} << iters) < numSeqs + 1) - { - ++iters; - } - return iters; -} - -/** - * @brief The most sequences one launch of a kernel baked with @p numSeqs can address. - * - * The kernel bisects `[0, num_seqs_runtime]` in a fixed number of iterations (an unrolled - * scf.for), which resolves at most `2^iters - 1` sequences. A larger batch leaves the search - * unfinished: the last blocks read the wrong sequence's lengths and the launch still - * succeeds. 511 -> 9 iterations -> 511. - */ -inline int64_t unifiedAttentionBatchCapacity(int64_t numSeqs) -{ - return (int64_t{1} << unifiedAttentionBinarySearchIters(numSeqs)) - 1; -} - /// The grid and block a candidate must launch with. struct UnifiedAttentionLaunchGeometry { @@ -158,7 +123,7 @@ struct UnifiedAttentionLaunchGeometry * = (qblocks, num_kv_heads, 1) q-major (traits.qMajorGrid) * * and `launch_block`, `(waveLanes * num_warps, 1, 1)`. `total_q` is `batch * seqLenQ` and - * `num_seqs` the RUNTIME batch, not the baked capacity. + * `num_seqs` the runtime batch. * * Throws instead of returning a degenerate grid: an empty launch returns cleanly having * written nothing, and prepare() is the last place a named failure is cheap. diff --git a/dnn-providers/hip-kernel-provider/src/integration_tests/kernel_ingestor_engine/IntegrationGpuGfx950AttentionUnifiedKnobs.cpp b/dnn-providers/hip-kernel-provider/src/integration_tests/kernel_ingestor_engine/IntegrationGpuGfx950AttentionUnifiedKnobs.cpp index 5d390f64fe57..46ecef48abcd 100644 --- a/dnn-providers/hip-kernel-provider/src/integration_tests/kernel_ingestor_engine/IntegrationGpuGfx950AttentionUnifiedKnobs.cpp +++ b/dnn-providers/hip-kernel-provider/src/integration_tests/kernel_ingestor_engine/IntegrationGpuGfx950AttentionUnifiedKnobs.cpp @@ -249,8 +249,8 @@ std::string maskName(Mask mask) /// Every Unified2D kernel class -- geometry x dtype x mask -- forced once, on a prefill shape /// with B > 1 and Sq != Skv: the default geometry cycling through the head configurations -/// (every cell ships it), the others on a cell that ships them; then one cold case, which -/// the engine's ranking serves with the default geometry. +/// (every cell ships it), the others on a cell that ships them; then two cold cases, one at +/// B 512, which the engine's ranking serves with the default geometry. std::vector knobCases2D() { std::vector cases; @@ -290,14 +290,20 @@ std::vector knobCases2D() {DataType::BFLOAT16, 128, 8, 2, Mask::TOP_LEFT, 3, 512, 512}, {}, knobsOf(NARROW_4X16)}); + // B 512 is one past what the old baked 9-iteration cu_q search resolved; the kernel now + // derives the trip count from the num_seqs kernarg, so every batch is served. + cases.push_back({"Unified2D_Cold_BatchPastTheOldBakedSearch", + ENGINE_2D, + {DataType::BFLOAT16, 128, 8, 2, Mask::TOP_LEFT, 512, 16, 64}, + {}, + knobsOf(NARROW_4X16)}); return cases; } /// Every Unified3D kernel class -- num_segments x dtype x causal mask -- forced once on a /// decode shape (Sq 1): the default num_segments cycling through the head configurations, the /// others on a cell that ships them; a NO_MASK decode graph (served by the bottom-right -/// kernel); and a cold B = 1024 decode case, past any 2D batch capacity: the decode grid has -/// none. +/// kernel); and a cold B = 1024 decode case. std::vector knobCases3D() { std::vector cases; @@ -347,10 +353,6 @@ std::vector knobCases() std::vector declineCases() { return { - // 2D bakes a batch capacity of 511 into its binary search. - {"Unified2D_BatchBeyondTheBakedCapacity", - ENGINE_2D, - {DataType::BFLOAT16, 128, 8, 2, Mask::NO_MASK, 512, 16, 64}}, // The identity cache is cut into whole 16-token blocks. {"Unified2D_SeqKvNotAWholeNumberOfBlocks", ENGINE_2D, diff --git a/dnn-providers/hip-kernel-provider/src/tests/CMakeLists.txt b/dnn-providers/hip-kernel-provider/src/tests/CMakeLists.txt index ea88df505de3..ba7a81c952d3 100644 --- a/dnn-providers/hip-kernel-provider/src/tests/CMakeLists.txt +++ b/dnn-providers/hip-kernel-provider/src/tests/CMakeLists.txt @@ -441,7 +441,7 @@ if(HIPDNN_ENABLE_KERNEL_INGESTOR) # suite below. foreach(_gfx950_engine IN ITEMS "dense|Gfx950AttentionDense|EveryShippedKernelsTileIsOneTheEngineLaunches" - "unified_2d|Gfx950AttentionUnified2D|EveryShippedKernelsGeometryIsOneTheEngineLaunches,EveryShippedKernelServesTheManifestsBatchCapacity" + "unified_2d|Gfx950AttentionUnified2D|EveryShippedKernelsGeometryIsOneTheEngineLaunches" "unified_3d|Gfx950AttentionUnified3D|EveryShippedKernelsGeometryIsOneTheEngineLaunches,EveryEntryLaunchesItsSegmentAndReduceKernels") string(REPLACE "|" ";" _gfx950_engine "${_gfx950_engine}") list(GET _gfx950_engine 0 _snake) diff --git a/dnn-providers/hip-kernel-provider/src/tests/engines/kernel_ingestor_engine/packs/Gfx950AttentionUnified2DTestGraphs.hpp b/dnn-providers/hip-kernel-provider/src/tests/engines/kernel_ingestor_engine/packs/Gfx950AttentionUnified2DTestGraphs.hpp index 91819e46d7e1..85a59adc262a 100644 --- a/dnn-providers/hip-kernel-provider/src/tests/engines/kernel_ingestor_engine/packs/Gfx950AttentionUnified2DTestGraphs.hpp +++ b/dnn-providers/hip-kernel-provider/src/tests/engines/kernel_ingestor_engine/packs/Gfx950AttentionUnified2DTestGraphs.hpp @@ -87,7 +87,8 @@ inline std::optional matchGraph(const GraphSpec& spec) /// The fields the engine's KMD declares, as a completed catalog record carries them. The /// defaults are the shipped baseline kernel for the default graph (bf16, D128, Hq 8 / Hkv 2, -/// top-left causal, block 16, T64, 4 warps x 16 rows, capacity 511). +/// top-left causal, block 16, T64, 4 warps x 16 rows, num_seqs 511). The engine does not read +/// num_seqs -- the batch is a kernarg -- but the KMD still declares it. struct KernelSpec { std::string dtype = "BF16"; diff --git a/dnn-providers/hip-kernel-provider/src/tests/engines/kernel_ingestor_engine/packs/TestGfx950AttentionUnified2DCatalogTiles.cpp b/dnn-providers/hip-kernel-provider/src/tests/engines/kernel_ingestor_engine/packs/TestGfx950AttentionUnified2DCatalogTiles.cpp index acdd98a544e8..6b7f2e5ed03f 100644 --- a/dnn-providers/hip-kernel-provider/src/tests/engines/kernel_ingestor_engine/packs/TestGfx950AttentionUnified2DCatalogTiles.cpp +++ b/dnn-providers/hip-kernel-provider/src/tests/engines/kernel_ingestor_engine/packs/TestGfx950AttentionUnified2DCatalogTiles.cpp @@ -22,9 +22,8 @@ /** * @file TestGfx950AttentionUnified2DCatalogTiles.cpp - * @brief Census of the geometry and batch capacity the shipped - * hipkernel:Gfx950AttentionUnified2D catalog declares, judged by the checks the engine - * itself applies. + * @brief Census of the geometry the shipped hipkernel:Gfx950AttentionUnified2D catalog + * declares, judged by the checks the engine itself applies. * * The engine launches a loaded kernel only when its candidateFor() accepts the record's * geometry through isSupportedGfx950Unified2DGeometry(); a kernel it declines ships but is @@ -33,11 +32,6 @@ * unbuildable warp/row/tile combination, or a field missing from the record) fails here * instead of silently never running. * - * It also pins the batch capacity: the kernel's binary search resolves at most 2^iters - 1 - * sequences, and a catalog regenerated with a smaller `enumerate.fixed.batch` would ship - * binaries that quietly serve fewer sequences than the shipped ones. Raise the floor below - * deliberately, with the manifest. - * * TestGfx950AttentionUnified2DDecodeHandoff checks the decode boundary between the two * unified engines: Unified2D declines every decode graph whose GQA group divides * GFX950_UNIFIED_3D_BLOCK_M and leaves it to Unified3D, so every such cell the 2D catalog @@ -56,9 +50,6 @@ using namespace hip_kernel_provider::kernel_ingestor_engine; using namespace hip_kernel_provider::kernel_ingestor_engine::testing; namespace ingestor = hipdnn_plugin_sdk::ingestor; -/// The batch every shipped kernel must serve at least: the manifest's `enumerate.fixed.batch`. -constexpr int64_t MINIMUM_BATCH_CAPACITY = 511; - /// A kernel's integer metadata field, or nullopt when it is absent or another type: the cases /// candidateFor() declines. std::optional integerField(const ingestor::KernelDescriptor& kernel, const char* field) @@ -148,28 +139,6 @@ TEST(TestGfx950AttentionUnified2DCatalogTiles, EveryShippedKernelsGeometryIsOneT EXPECT_GT(checked, 0U) << "the loaded set holds no kernels to check"; } -TEST(TestGfx950AttentionUnified2DCatalogTiles, EveryShippedKernelServesTheManifestsBatchCapacity) -{ - const auto& set = loadedSet(GFX950_ATTENTION_UNIFIED_2D_ENGINE_NAME); - - std::size_t checked = 0; - for(const auto& pack : set.packs) - { - for(const auto& kernel : pack.kernels) - { - ++checked; - const auto numSeqs = integerField(kernel, "num_seqs"); - ASSERT_TRUE(numSeqs.has_value()) - << "kernel '" << kernel.name << "' does not carry an integer num_seqs"; - EXPECT_GT(*numSeqs, 0) << "kernel '" << kernel.name << "'"; - EXPECT_GE(gfx950Unified2DBatchCapacity(*numSeqs), MINIMUM_BATCH_CAPACITY) - << "kernel '" << kernel.name << "' bakes num_seqs " << *numSeqs - << ", which resolves fewer than " << MINIMUM_BATCH_CAPACITY << " sequences"; - } - } - EXPECT_GT(checked, 0U) << "the loaded set holds no kernels to check"; -} - TEST(TestGfx950AttentionUnified2DDecodeHandoff, EveryDecodeCellHandedToUnified3DHasAUnified3DKernel) { std::set unified3dCells; diff --git a/dnn-providers/hip-kernel-provider/src/tests/engines/kernel_ingestor_engine/packs/TestGfx950AttentionUnified2DDispatch.cpp b/dnn-providers/hip-kernel-provider/src/tests/engines/kernel_ingestor_engine/packs/TestGfx950AttentionUnified2DDispatch.cpp index d6f5096ced42..49b3ea1b5cda 100644 --- a/dnn-providers/hip-kernel-provider/src/tests/engines/kernel_ingestor_engine/packs/TestGfx950AttentionUnified2DDispatch.cpp +++ b/dnn-providers/hip-kernel-provider/src/tests/engines/kernel_ingestor_engine/packs/TestGfx950AttentionUnified2DDispatch.cpp @@ -98,7 +98,6 @@ std::string prepareFailureFor(const KernelSpec& spec, const GraphSpec& graphSpec constexpr const char* MISSING_TOKEN_MARKER = "missing bound token"; constexpr const char* OUTPUT_LAYOUT_MARKER = "not dense BSHD"; constexpr const char* GEOMETRY_MARKER = "declares no supported geometry"; -constexpr const char* CAPACITY_MARKER = "serves at most"; constexpr const char* BLOCK_MARKER = "does not divide seqlen_kv"; } // namespace @@ -250,19 +249,16 @@ TEST(TestGfx950AttentionUnified2DDispatch, BindsExactlyTheTokensTheDispatchReads /// The positive control the refusals below depend on: a prepare() that threw for every input /// would satisfy all of them and dispatch nothing. A failure is still expected -- the archive -/// is absent -- but it must be the loader's, which means the bindings, the geometry, the -/// capacity and the identity tables were all accepted first. +/// is absent -- but it must be the loader's, which means the bindings, the geometry and the +/// identity tables were all accepted first. TEST(TestGfx950AttentionUnified2DDispatch, ReachesTheLoaderOnceEverythingBeforeItIsGood) { const std::string failure = prepareFailureFor(KernelSpec{}); ASSERT_FALSE(failure.empty()) << "the archive named by the fixture does not exist, so " "prepare() cannot have succeeded"; - for(const char* marker : {MISSING_TOKEN_MARKER, - OUTPUT_LAYOUT_MARKER, - GEOMETRY_MARKER, - CAPACITY_MARKER, - BLOCK_MARKER}) + for(const char* marker : + {MISSING_TOKEN_MARKER, OUTPUT_LAYOUT_MARKER, GEOMETRY_MARKER, BLOCK_MARKER}) { EXPECT_EQ(failure.find(marker), std::string::npos) << failure; } @@ -363,7 +359,6 @@ TEST(TestGfx950AttentionUnified2DDispatch, RefusesToPrepareACandidateWithoutASup with([](KernelSpec& s) { s.blockSize = 24; }); with([](KernelSpec& s) { s.useMfma32x32 = 1; }); // needs 32 rows per warp with([](KernelSpec& s) { s.useTransposedQk32x32 = 1; }); // needs the 32x32 path - with([](KernelSpec& s) { s.numSeqs = 0; }); with([](KernelSpec& s) { s.maskType = 3; }); with([](KernelSpec& s) { s.headSize = 96; }); with([](KernelSpec& s) { s.numQueryHeads = 9; }); @@ -390,6 +385,10 @@ TEST(TestGfx950AttentionUnified2DDispatch, RefusesToPrepareACandidateMissingAFie { continue; // read by the matcher, not by prepare() } + if(field == "num_seqs") + { + continue; // declared by the KMD, not read: the batch is a kernarg + } SCOPED_TRACE(field); auto kernel = makeKpackKernel(KernelSpec{}, pythonAbi()); kernel.metadata.erase(field); @@ -404,25 +403,21 @@ TEST(TestGfx950AttentionUnified2DDispatch, RefusesToPrepareACandidateMissingAFie EXPECT_NE(failure.find(GEOMETRY_MARKER), std::string::npos) << failure; } -/// One past the capacity the kernel was built for. The search would leave the last blocks -/// reading another sequence's lengths without a fault, so prepare() refuses by name. -TEST(TestGfx950AttentionUnified2DDispatch, RefusesToPrepareABatchAboveTheCandidatesCapacity) +/// The kernel derives its cu_q search trip count from the `num_seqs` kernarg, so a batch past +/// what the old baked 9-iteration search resolved (2^9 - 1) reaches the loader like any other. +TEST(TestGfx950AttentionUnified2DDispatch, PreparesABatchPastTheOldBakedSearch) { GraphSpec spec; - spec.batch = 4; + spec.batch = 512; spec.seqLenQ = 16; spec.seqLenKv = 16; - KernelSpec capacityThree; - capacityThree.numSeqs = 3; // 2 iterations -> 3 sequences - const std::string tooMany = prepareFailureFor(capacityThree, spec); - ASSERT_FALSE(tooMany.empty()); - EXPECT_NE(tooMany.find(CAPACITY_MARKER), std::string::npos) << tooMany; - - spec.batch = 3; - const std::string exactlyFull = prepareFailureFor(capacityThree, spec); - ASSERT_FALSE(exactlyFull.empty()); - EXPECT_EQ(exactlyFull.find(CAPACITY_MARKER), std::string::npos) << exactlyFull; + const std::string failure = prepareFailureFor(KernelSpec{}, spec); + ASSERT_FALSE(failure.empty()) << "the archive named by the fixture does not exist, so " + "prepare() cannot have succeeded"; + EXPECT_EQ(failure.find(GEOMETRY_MARKER), std::string::npos) << failure; + EXPECT_EQ(failure.find(BLOCK_MARKER), std::string::npos) << failure; + EXPECT_NE(failure.find("does not exist"), std::string::npos) << failure; } /// A block size that does not divide the graph's Skv is not a whole-block cache. diff --git a/dnn-providers/hip-kernel-provider/src/tests/engines/kernel_ingestor_engine/packs/TestGfx950AttentionUnified2DGeometry.cpp b/dnn-providers/hip-kernel-provider/src/tests/engines/kernel_ingestor_engine/packs/TestGfx950AttentionUnified2DGeometry.cpp index 0a8df84f6886..088e3efcf37c 100644 --- a/dnn-providers/hip-kernel-provider/src/tests/engines/kernel_ingestor_engine/packs/TestGfx950AttentionUnified2DGeometry.cpp +++ b/dnn-providers/hip-kernel-provider/src/tests/engines/kernel_ingestor_engine/packs/TestGfx950AttentionUnified2DGeometry.cpp @@ -21,7 +21,7 @@ /** * @file TestGfx950AttentionUnified2DGeometry.cpp * @brief Pins what hipkernel:Gfx950AttentionUnified2D restates from its Python builder: the - * geometry rules, the launch grid, the batch capacity and the paged-KV metadata. + * geometry rules, the launch grid and the paged-KV metadata. * * Every expectation is a literal worked out from the Python (attention_tiled_2d.py, * `AttentionTuningSpec.launch_grid`), not computed the way the header computes it: a helper @@ -172,53 +172,6 @@ TEST(TestGfx950AttentionUnified2DGeometry, TheThirtyTwoByThirtyTwoPathsKeepTheir EXPECT_TRUE(isSupportedGfx950Unified2DGeometry(transposed)); } -// ============================================================================= -// Batch capacity -// ============================================================================= - -/// The kernel bisects [0, num_seqs] in `iters` unrolled steps, which resolves 2^iters - 1 -/// sequences. Table values are ceil(log2(num_seqs + 1)) worked out by hand. -TEST(TestGfx950AttentionUnified2DGeometry, BinarySearchItersIsCeilLogTwoOfNumSeqsPlusOne) -{ - const std::vector> table{{1, 1}, - {2, 2}, - {3, 2}, - {4, 3}, - {7, 3}, - {8, 4}, - {15, 4}, - {16, 5}, - {255, 8}, - {256, 9}, - {511, 9}, - {512, 10}, - {1023, 10}, - {1024, 11}}; - for(const auto& [numSeqs, iters] : table) - { - EXPECT_EQ(gfx950Unified2DBinarySearchIters(numSeqs), iters) << "num_seqs " << numSeqs; - } -} - -/// A spec that declares no num_seqs gets the kernel's conservative 32 iterations. -TEST(TestGfx950AttentionUnified2DGeometry, AnUndeclaredNumSeqsIsTheThirtyTwoIterationFallback) -{ - EXPECT_EQ(gfx950Unified2DBinarySearchIters(0), 32); - EXPECT_EQ(gfx950Unified2DBinarySearchIters(-7), 32); -} - -/// Capacity is 2^iters - 1, so the baked 511 serves exactly 511 and not 512. The boundary is -/// the whole point: one past it reads another sequence's seq_lens with no fault. -TEST(TestGfx950AttentionUnified2DGeometry, BatchCapacityIsTheLargestBatchTheSearchResolves) -{ - EXPECT_EQ(gfx950Unified2DBatchCapacity(511), 511); - EXPECT_EQ(gfx950Unified2DBatchCapacity(512), 1023); - EXPECT_EQ(gfx950Unified2DBatchCapacity(1), 1); - EXPECT_EQ(gfx950Unified2DBatchCapacity(2), 3); - EXPECT_EQ(gfx950Unified2DBatchCapacity(3), 3); - EXPECT_EQ(gfx950Unified2DBatchCapacity(4), 7); -} - // ============================================================================= // Launch geometry // ============================================================================= diff --git a/dnn-providers/hip-kernel-provider/src/tests/engines/kernel_ingestor_engine/packs/TestGfx950AttentionUnified2DMatchers.cpp b/dnn-providers/hip-kernel-provider/src/tests/engines/kernel_ingestor_engine/packs/TestGfx950AttentionUnified2DMatchers.cpp index 100513660ac4..08372dbb62dd 100644 --- a/dnn-providers/hip-kernel-provider/src/tests/engines/kernel_ingestor_engine/packs/TestGfx950AttentionUnified2DMatchers.cpp +++ b/dnn-providers/hip-kernel-provider/src/tests/engines/kernel_ingestor_engine/packs/TestGfx950AttentionUnified2DMatchers.cpp @@ -23,7 +23,7 @@ * survivors rank. * * One case per rule, in severity order: the silent-wrong-answer rules first (layout, shape, - * mask, batch capacity), then faults (addressing limits), then declined features. A graph + * mask), then faults (addressing limits), then declined features. A graph * match that returns nullopt empties this engine's WHOLE catalog, so every case that expects * a decline is paired with a positive control that shows the same graph one field away is * served. @@ -550,11 +550,10 @@ TEST(TestGfx950AttentionUnified2DMatchers, KernelMatchRequiresTheCandidatesBlock EXPECT_FALSE(matchesKernel(sixteenOnly, block64)); } -/// The kernel bisects the sequences in a fixed number of unrolled steps, which resolves at -/// most 2^iters - 1 of them. Past that the last blocks read another sequence's lengths and the -/// launch still succeeds, so the boundary is exact: capacity 3 (2 iterations) serves a batch -/// of 3 and not 4, and the shipped 511 (9 iterations) serves 511 and not 512. -TEST(TestGfx950AttentionUnified2DMatchers, KernelMatchBoundsTheBatchByTheBakedCapacity) +/// The kernel derives its cu_q search trip count from the `num_seqs` kernarg, so no batch is +/// baked: a record declaring the shipped `num_seqs` 511 (or a smaller one, or none) serves a +/// batch past 2^9 - 1, which the old baked 9-iteration search could not resolve. +TEST(TestGfx950AttentionUnified2DMatchers, KernelMatchServesABatchPastTheOldBakedSearch) { const auto graphOfBatch = [](int64_t batch) { GraphSpec spec; @@ -563,20 +562,19 @@ TEST(TestGfx950AttentionUnified2DMatchers, KernelMatchBoundsTheBatchByTheBakedCa spec.seqLenKv = 16; return spec; }; - const auto kernelOfCapacity = [](int64_t numSeqs) { - KernelSpec spec; - spec.numSeqs = numSeqs; - return spec; - }; + KernelSpec numSeqsThree; + numSeqsThree.numSeqs = 3; + + for(const int64_t batch : {3, 4, 511, 512, 4096}) + { + EXPECT_TRUE(matchesKernel(graphOfBatch(batch), KernelSpec{})) << "batch " << batch; + EXPECT_TRUE(matchesKernel(graphOfBatch(batch), numSeqsThree)) << "batch " << batch; + } - EXPECT_TRUE(matchesKernel(graphOfBatch(3), kernelOfCapacity(3))); - EXPECT_FALSE(matchesKernel(graphOfBatch(4), kernelOfCapacity(3))); - EXPECT_TRUE(matchesKernel(graphOfBatch(511), kernelOfCapacity(511))); - EXPECT_FALSE(matchesKernel(graphOfBatch(512), kernelOfCapacity(511))); - // A capacity that is not 2^k - 1 still rounds UP to the iterations it needs: num_seqs 4 - // takes 3 iterations, so it serves 7. - EXPECT_TRUE(matchesKernel(graphOfBatch(7), kernelOfCapacity(4))); - EXPECT_FALSE(matchesKernel(graphOfBatch(8), kernelOfCapacity(4))); + auto withoutNumSeqs = makeKernel(KernelSpec{}); + withoutNumSeqs.metadata.erase(std::string("num_seqs")); + const Graph batch512(graphOfBatch(512)); + EXPECT_TRUE(matchesKernel(batch512, withoutNumSeqs)); } /// The metadata carries no shape: one binary serves any Sq and Skv for its configuration. @@ -613,8 +611,6 @@ TEST(TestGfx950AttentionUnified2DMatchers, KernelMatchDeclinesACandidateWithAnUn with([](KernelSpec& s) { s.blockSize = 24; }); with([](KernelSpec& s) { s.useMfma32x32 = 1; }); with([](KernelSpec& s) { s.useTransposedQk32x32 = 1; }); - with([](KernelSpec& s) { s.numSeqs = 0; }); - with([](KernelSpec& s) { s.numSeqs = -3; }); with([](KernelSpec& s) { s.maskType = 7; }); for(std::size_t i = 0; i < unbuildable.size(); ++i) @@ -630,6 +626,10 @@ TEST(TestGfx950AttentionUnified2DMatchers, KernelMatchDeclinesAnIncompleteOrMist const Graph graph; for(const auto& entry : metadataOf(KernelSpec{})) { + if(entry.first == "num_seqs") + { + continue; // declared by the KMD, not read: the batch is a kernarg + } SCOPED_TRACE(entry.first); auto kernel = makeKernel(KernelSpec{}); kernel.metadata.erase(entry.first); diff --git a/dnn-providers/hip-kernel-provider/src/tests/engines/kernel_ingestor_engine/packs/TestUnifiedAttentionGeometry.cpp b/dnn-providers/hip-kernel-provider/src/tests/engines/kernel_ingestor_engine/packs/TestUnifiedAttentionGeometry.cpp index d1f0f22984a1..be371db8969a 100644 --- a/dnn-providers/hip-kernel-provider/src/tests/engines/kernel_ingestor_engine/packs/TestUnifiedAttentionGeometry.cpp +++ b/dnn-providers/hip-kernel-provider/src/tests/engines/kernel_ingestor_engine/packs/TestUnifiedAttentionGeometry.cpp @@ -77,7 +77,6 @@ TEST(TestUnifiedAttentionGeometry, Gfx950WrappersAreTheNeutralFunctions) EXPECT_EQ(gfx950Unified2DLaunchGeometry(2, 32, 16, 4, 300, 3, "k"), unifiedAttention2DLaunchGeometry( GFX950_UNIFIED_ATTENTION_TRAITS, 2, 32, 16, 4, 300, 3, "k")); - EXPECT_EQ(gfx950Unified2DBatchCapacity(511), unifiedAttentionBatchCapacity(511)); EXPECT_EQ(gfx950Unified2DWorkspaceLayout(3, 512, 16), unifiedAttentionWorkspaceLayout(3, 512, 16)); EXPECT_EQ(gfx950Unified2DWorkspaceImage(3, 7, 512, 16), From f9dd8b6086accf82ea1702beb680037d47938e52 Mon Sep 17 00:00:00 2001 From: AnaghaRaoAMD Date: Wed, 4 Nov 2026 17:59:57 -0500 Subject: [PATCH 5/8] fix(rocke): address #15024 Copilot review on gfx950 loaders and reduce launches - gfx950 2D page copies: require KV_HALVES_PER_CALL % head_size == 0 in both engines. A one-warp head_size 192 copy (512 halves, 2 2/3 rows) truncated to 2 tokens and addressed the wrong row and dim; those shapes keep the paged-KV descriptor. New cfg71 (the D192 one-warp spec dispatch picks) in both emitters and a D192 one-warp numeric pair. - Dense segment combine: launch the tiled reduce with its runtime kernargs (num_query_heads); it was launched with the pointer-only signature. - Binding parity mirror: add cfg68-71 so prove_parity_binding's completeness check matches the emitters. - Launch-signature test: check the "dg" token (split on "_" drops the underscore, so "_dg" never matched). - Differential golden: re-bless gfx942_attention_tiled_2d_stress (cfg21, 32/5 heads, is now refused by both engines under the head-count rule) and gfx950_attention_tiled_2d (cfg71), all flavors + canonical. Byte identity GREEN at llvm20/22/23; representative-IR and strided-KV goldens unchanged. Co-Authored-By: Claude --- .../common/attention_dense_kv_split.py | 13 +++- .../kernels/gfx950/attention_tiled_2d.py | 17 ++-- .../gfx942_attention_tiled_2d_stress_emit.c | 2 + .../gfx942_attention_tiled_2d_stress_emit.py | 2 + .../parity/gfx950_attention_tiled_2d_emit.c | 19 +++++ .../parity/gfx950_attention_tiled_2d_emit.py | 17 ++++ ...st_attention_dense_kv_split_combine_abi.py | 77 +++++++++++++++++++ .../test_attention_tiled_launch_signature.py | 2 +- ...t_attention_tiled_runtime_shape_numeric.py | 20 +++++ .../rocke/platform/CMakeLists.txt | 4 +- .../attention_tiled_parity_configs.py | 66 ++++++++++++++++ ...ttention_tiled_2d_kv_descriptors_loads.cpp | 13 ++-- .../differential/golden/llvm_gfx_all.json | 10 +-- 13 files changed, 240 insertions(+), 22 deletions(-) create mode 100644 dnn-providers/hip-kernel-provider/rocke/library/tests/test_attention_dense_kv_split_combine_abi.py diff --git a/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/attention_dense_kv_split.py b/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/attention_dense_kv_split.py index 8da79a427125..736458e6738b 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/attention_dense_kv_split.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/attention_dense_kv_split.py @@ -190,27 +190,34 @@ def combine_spec(spec, arch: str): def launch_segment_combine(spec, arch: str, out, workspace, stream: int) -> None: """Merge the segments in ``workspace`` into ``out`` ([B, Sq, Hq, D]).""" - from kernels.common.attention_unified import _reduce_signature + from kernels.common.attention_unified import ( + _reduce_signature, + _runtime_kernarg_fields, + ) from rocke.helpers.compile import compile_kernel from rocke.runtime import KernelLauncher, LaunchConfig red = combine_spec(spec, arch) + spec_type, build = _combine_impl(arch) + # The reduce reads its runtime shape fields (num_query_heads) as kernargs. + runtime_fields = _runtime_kernarg_fields(spec_type) key = (arch, red) launcher = _COMBINE_LAUNCHER_CACHE.get(key) if launcher is None: - _, build = _combine_impl(arch) art = compile_kernel( build(red, arch=arch), arch=arch, backend="python", capture_ir_text=False ) launcher = KernelLauncher( hsaco=art.hsaco, kernel_name=art.kernel_name, - signature=_reduce_signature(spec.dtype), + signature=_reduce_signature(spec.dtype, runtime_fields=runtime_fields), ) _COMBINE_LAUNCHER_CACHE[key] = launcher import torch vals = dict(workspace) + for name in runtime_fields: + vals[name] = int(getattr(red, name)) vals["output_ptr"] = out # Unused by the combine body; any int32 buffer satisfies the ABI. vals["seq_lens_ptr"] = workspace["segm_max_ptr"].view(torch.int32) diff --git a/dnn-providers/hip-kernel-provider/rocke/library/kernels/gfx950/attention_tiled_2d.py b/dnn-providers/hip-kernel-provider/rocke/library/kernels/gfx950/attention_tiled_2d.py index 5333b0bc5784..262f0ce4c0a0 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/kernels/gfx950/attention_tiled_2d.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/kernels/gfx950/attention_tiled_2d.py @@ -2159,13 +2159,18 @@ def _acc_idx(n: int, atom: int) -> int: # return ``0`` (a guaranteed-valid block id). Downstream causal / # in_prefix masks already discard the bogus tokens. block_table_max_idx = b.to_sgpr_u32(b.mul(num_seqs_p, bt_stride_p)) - # A copy covers COPY_TOKENS whole tokens (KV_HALVES_PER_CALL is a multiple - # of HD). When they fit in one page the bf16/fp16 loaders build the page - # base, the in-page token offset and the lane part separately (see - # ``_paged_copy``); otherwise each wave's share still sits in one page and - # the loaders keep the paged-KV descriptor. + # When a copy covers COPY_TOKENS whole tokens and they fit in one page, the + # bf16/fp16 loaders build the page base, the in-page token offset and the + # lane part separately (see ``_paged_copy``). A copy that ends mid-row + # (head_size 192 with one warp: 512 halves are 2 2/3 rows) or spans pages + # keeps the paged-KV descriptor, which addresses each lane's own row. COPY_TOKENS = KV_HALVES_PER_CALL // HD - PAGE_COPIES = not FAST_PAGED_KV_DESC and not KV_FP8 and BS % COPY_TOKENS == 0 + PAGE_COPIES = ( + not FAST_PAGED_KV_DESC + and not KV_FP8 + and KV_HALVES_PER_CALL % HD == 0 + and BS % COPY_TOKENS == 0 + ) # With page copies, the loop carries the page ids of the current and the # next tile and reads the ids of the tile after next one iteration ahead, # so no copy waits on a table read issued just before it. The schedules diff --git a/dnn-providers/hip-kernel-provider/rocke/library/tests/parity/gfx942_attention_tiled_2d_stress_emit.c b/dnn-providers/hip-kernel-provider/rocke/library/tests/parity/gfx942_attention_tiled_2d_stress_emit.c index 50212bbfbd5e..58d24f5e9670 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/tests/parity/gfx942_attention_tiled_2d_stress_emit.c +++ b/dnn-providers/hip-kernel-provider/rocke/library/tests/parity/gfx942_attention_tiled_2d_stress_emit.c @@ -231,6 +231,8 @@ static int make_spec(int idx, rocke_attention_tiled_2d_spec_t* s) s->sliding_window = 0; s->has_softcap = false; break; + /* 32/5 does not divide: both engines refuse it (check_head_counts), so it + * is a both-reject parity case rather than an emitted kernel. */ case 21: s->head_size = 64; s->block_size = 32; diff --git a/dnn-providers/hip-kernel-provider/rocke/library/tests/parity/gfx942_attention_tiled_2d_stress_emit.py b/dnn-providers/hip-kernel-provider/rocke/library/tests/parity/gfx942_attention_tiled_2d_stress_emit.py index ae69cac3b1a7..bb6948fa4521 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/tests/parity/gfx942_attention_tiled_2d_stress_emit.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/tests/parity/gfx942_attention_tiled_2d_stress_emit.py @@ -228,6 +228,8 @@ sliding_window=0, has_softcap=False, ), + # 32/5 does not divide: both engines refuse it (check_head_counts), so it + # is a both-reject parity case rather than an emitted kernel. 21: dict( head_size=64, block_size=32, diff --git a/dnn-providers/hip-kernel-provider/rocke/library/tests/parity/gfx950_attention_tiled_2d_emit.c b/dnn-providers/hip-kernel-provider/rocke/library/tests/parity/gfx950_attention_tiled_2d_emit.c index ceab4bf460a3..82705f94e692 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/tests/parity/gfx950_attention_tiled_2d_emit.c +++ b/dnn-providers/hip-kernel-provider/rocke/library/tests/parity/gfx950_attention_tiled_2d_emit.c @@ -1184,6 +1184,25 @@ static int make_spec(int idx, rocke_attention_tiled_2d_spec_t* s) s->use_mfma_32x32 = true; s->use_transposed_qk_32x32 = true; break; + /* cfg71: bf16 D192/b16 one warp, the spec dispatch picks for short D192 + * prefill. One copy is 512 halves, 2 2/3 rows, so it is not a whole number + * of tokens and the bf16 loaders must keep the paged-KV descriptor. */ + case 71: + s->head_size = 192; + s->block_size = 16; + s->num_query_heads = 16; + s->num_kv_heads = 2; + s->dtype = "bf16"; + s->use_sinks = false; + s->sliding_window = 0; + s->has_softcap = false; + s->num_seqs = 4; + s->has_waves_per_eu = true; + s->waves_per_eu = 2; + s->use_register_pv = true; + s->has_tile_size = true; + s->tile_size = 32; + break; default: return -1; diff --git a/dnn-providers/hip-kernel-provider/rocke/library/tests/parity/gfx950_attention_tiled_2d_emit.py b/dnn-providers/hip-kernel-provider/rocke/library/tests/parity/gfx950_attention_tiled_2d_emit.py index e927ce605c00..da8a3ea2608e 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/tests/parity/gfx950_attention_tiled_2d_emit.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/tests/parity/gfx950_attention_tiled_2d_emit.py @@ -1128,6 +1128,23 @@ use_mfma_32x32=True, use_transposed_qk_32x32=True, ), + # cfg71: bf16 D192/b16 one warp, the spec dispatch picks for short D192 + # prefill. One copy is 512 halves, 2 2/3 rows, so it is not a whole number + # of tokens and the bf16 loaders must keep the paged-KV descriptor. + 71: dict( + head_size=192, + block_size=16, + num_query_heads=16, + num_kv_heads=2, + dtype="bf16", + use_sinks=False, + sliding_window=0, + has_softcap=False, + num_seqs=4, + waves_per_eu=2, + use_register_pv=True, + tile_size=32, + ), } # --- idx64-67: emit_lse (trailing lse_ptr + FP32 LSE store), one per epilogue diff --git a/dnn-providers/hip-kernel-provider/rocke/library/tests/test_attention_dense_kv_split_combine_abi.py b/dnn-providers/hip-kernel-provider/rocke/library/tests/test_attention_dense_kv_split_combine_abi.py new file mode 100644 index 000000000000..49d1e3e9d6fa --- /dev/null +++ b/dnn-providers/hip-kernel-provider/rocke/library/tests/test_attention_dense_kv_split_combine_abi.py @@ -0,0 +1,77 @@ +# Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +# SPDX-License-Identifier: MIT + +"""The dense segment-combine launch packs every argument the tiled reduce declares. + +``launch_segment_combine`` launches the tiled 3D reduce kernel. That kernel reads +``num_query_heads`` as a kernarg (``RUNTIME_KERNARG_FIELDS``), so a launcher built +from the pointer-only reduce signature leaves the head count, and with it the +reduce's token strides, undefined. This captures the launcher's signature and +values on the host and checks them against the reduce kernel's params. +""" + +from __future__ import annotations + +import pytest +from kernels.common import attention_dense_kv_split as kv_split +from kernels.gfx942.attention_dense import Gfx942AttentionDenseSpec +from kernels.gfx950.attention_dense import Gfx950AttentionDenseSpec + +torch = pytest.importorskip("torch") + + +class _Art: + hsaco = b"" + + def __init__(self, kernel): + self.kernel_name = kernel.name + + +@pytest.mark.parametrize( + "arch, spec_cls", + [("gfx942", Gfx942AttentionDenseSpec), ("gfx950", Gfx950AttentionDenseSpec)], +) +def test_combine_launch_matches_reduce_params(monkeypatch, arch, spec_cls): + import rocke.helpers.compile as compile_mod + import rocke.runtime as runtime_mod + + built = {} + + def fake_compile(kernel, **_): + built["kernel"] = kernel + return _Art(kernel) + + class FakeLauncher: + def __init__(self, *, hsaco, kernel_name, signature): + built["signature"] = signature + + def __call__(self, vals, *, config): + built["vals"] = vals + + monkeypatch.setattr(compile_mod, "compile_kernel", fake_compile) + monkeypatch.setattr(runtime_mod, "KernelLauncher", FakeLauncher) + monkeypatch.setattr(kv_split, "_COMBINE_LAUNCHER_CACHE", {}) + + spec = spec_cls( + batch=2, + seqlen_q=256, + seqlen_kv=2048, + num_query_heads=16, + num_kv_heads=2, + head_size=128, + causal=True, + dtype="bf16", + kv_segments=4, + ) + workspace = kv_split.allocate_segment_workspace(spec, "cpu") + out = torch.empty( + (spec.batch, spec.seqlen_q, spec.num_query_heads, spec.head_size), + dtype=torch.bfloat16, + ) + kv_split.launch_segment_combine(spec, arch, out, workspace, stream=0) + + params = [p.name for p in built["kernel"].params] + assert params == [a["name"] for a in built["signature"]] + assert "num_query_heads" in params + assert set(params) <= set(built["vals"]) + assert built["vals"]["num_query_heads"] == spec.num_query_heads diff --git a/dnn-providers/hip-kernel-provider/rocke/library/tests/test_attention_tiled_launch_signature.py b/dnn-providers/hip-kernel-provider/rocke/library/tests/test_attention_tiled_launch_signature.py index 66f80973fb47..8e900974a432 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/tests/test_attention_tiled_launch_signature.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/tests/test_attention_tiled_launch_signature.py @@ -207,7 +207,7 @@ def test_tuning_candidate_signature_matches_its_kernel(codepath): # The signature is the segment's; _get_3d_pipeline launches the # reduce with the production 3D pair. kernel, reduce = kernel - assert "_dg" not in kernel.name.split("_"), kernel.name + assert "dg" not in kernel.name.split("_"), kernel.name _, red_sig = au._tiled_3d_signatures( "gfx942", spec.kernel_spec.dtype, diff --git a/dnn-providers/hip-kernel-provider/rocke/library/tests/test_attention_tiled_runtime_shape_numeric.py b/dnn-providers/hip-kernel-provider/rocke/library/tests/test_attention_tiled_runtime_shape_numeric.py index 867e9b3e022f..0bfba7f0d7a4 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/tests/test_attention_tiled_runtime_shape_numeric.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/tests/test_attention_tiled_runtime_shape_numeric.py @@ -205,6 +205,17 @@ def _shape(q_lens, kv_lens, hq, hkv, window): 0, ), ), + # D192 short prefill takes the one-warp 2D kernel, whose 512-half K/V copy + # ends mid-row: the loaders must keep the paged-KV descriptor. + ( + "gfx950", + "2d-bf16-d192-one-warp", + "2d", + None, + (192, 16, "bf16"), + _shape([64, 48], [256, 192], 16, 2, 0), + _shape([64, 60, 50, 40, 56], [256, 192, 320, 200, 280], 32, 4, 0), + ), ("gfx950", "3d-decode", "3d", None, (128, 16, "bf16"), *_DECODE), ("gfx950", "3d-two-token", "3d", None, (128, 16, "bf16"), *_TWO_TOKEN), ) @@ -265,6 +276,15 @@ def test_pair_shares_one_binary(arch, path, route, geometry, first, second): assert _key(a, path, arch) == _key(b, path, arch) +def test_d192_pair_takes_the_one_warp_kernel(): + """The D192 pair must stay on the one-warp spec: that is the geometry whose + K/V copy ends mid-row, which the numeric pair is there to cover.""" + case = next(c for c in _CASES if c[1] == "2d-bf16-d192-one-warp") + for shape in case[5:]: + spec = au._tiled_spec_from_problem(_problem("gfx950", case[4], shape), "gfx950") + assert spec.head_size == 192 and spec.num_warps == 1, spec + + def _encode(x, dtype): """fp32 values as the kernel's fp16 or bf16 (raw uint16) host array.""" return x.astype(np.float16) if dtype == "fp16" else f32_to_bf16(x) diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/CMakeLists.txt b/dnn-providers/hip-kernel-provider/rocke/platform/CMakeLists.txt index 5d04d72609ef..b44e0265bc93 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/CMakeLists.txt +++ b/dnn-providers/hip-kernel-provider/rocke/platform/CMakeLists.txt @@ -732,7 +732,8 @@ if(ROCKE_INSTALL_TESTS) # fmha_mfma free-stride ABI and layout contracts, the 2D/3D mask_type spec # fields, the 3D segment+reduce module's pairing, the reduce's LDS segment # bound against the compiled code object, the tiled C++ entries' head-count - # and arch refusals, and the backward-test memo bound in tests/conftest.py. + # and arch refusals, the dense segment-combine launch ABI, and the + # backward-test memo bound in tests/conftest.py. # File-list style (same pattern as conv_pytest) so PYTHONPATH resolves rocke/ # correctly from the installed tree. set(_ROCKE_LIBRARY_ATTENTION_CPU_TEST_RELPATHS @@ -745,6 +746,7 @@ if(ROCKE_INSTALL_TESTS) "tests/library/tests/test_attention_tiled_3d_kv_cache_policy.py" "tests/library/tests/test_attention_tiled_launch_signature.py" "tests/library/tests/test_attention_tiled_cpp_entry_guards.py" + "tests/library/tests/test_attention_dense_kv_split_combine_abi.py" "tests/library/tests/test_attention_tiled_3d_mask_type.py" "tests/library/tests/test_attention_tiled_3d_module.py" "tests/library/tests/test_attention_tiled_3d_reduce_lds.py" diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/bindings/attention_tiled_parity_configs.py b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/bindings/attention_tiled_parity_configs.py index b15211806f7f..f69c46110918 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/bindings/attention_tiled_parity_configs.py +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/bindings/attention_tiled_parity_configs.py @@ -1280,6 +1280,72 @@ "use_alibi": True, "emit_lse": True, }, + { + "head_size": 64, + "block_size": 16, + "num_query_heads": 32, + "num_kv_heads": 8, + "dtype": "bf16", + "use_sinks": False, + "sliding_window": 0, + "has_softcap": False, + "num_seqs": 2, + "num_warps": 4, + "waves_per_eu": 2, + "tile_size": 32, + "block_m_per_warp": 32, + "use_mfma_32x32": True, + "use_transposed_qk_32x32": True, + "use_transposed_scalar_state": True, + "use_transposed_mask_once": True, + "use_transposed_half_local_pv": True, + "use_mfma32_skip_legacy_qreg": True, + "use_transposed_mask_limit": True, + }, + { + "head_size": 64, + "block_size": 16, + "num_query_heads": 32, + "num_kv_heads": 8, + "dtype": "bf16", + "use_sinks": False, + "sliding_window": 0, + "has_softcap": False, + "num_seqs": 2, + "tile_size": 256, + "block_m_per_warp": 32, + "use_mfma_32x32": True, + "use_transposed_qk_32x32": True, + }, + { + "head_size": 64, + "block_size": 16, + "num_query_heads": 32, + "num_kv_heads": 8, + "dtype": "bf16", + "use_sinks": False, + "sliding_window": 0, + "has_softcap": False, + "num_seqs": 2, + "tile_size": 288, + "block_m_per_warp": 32, + "use_mfma_32x32": True, + "use_transposed_qk_32x32": True, + }, + { + "head_size": 192, + "block_size": 16, + "num_query_heads": 16, + "num_kv_heads": 2, + "dtype": "bf16", + "use_sinks": False, + "sliding_window": 0, + "has_softcap": False, + "num_seqs": 4, + "waves_per_eu": 2, + "use_register_pv": True, + "tile_size": 32, + }, ] GFX942_ATTENTION_TILED_3D = [ diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/gfx950/attention_tiled_2d_kv_descriptors_loads.cpp b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/gfx950/attention_tiled_2d_kv_descriptors_loads.cpp index 057e3195d343..b613250ea548 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/gfx950/attention_tiled_2d_kv_descriptors_loads.cpp +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/gfx950/attention_tiled_2d_kv_descriptors_loads.cpp @@ -183,11 +183,11 @@ static rocke_value_t* rocke__wave_lds_offset_i64(rocke_gfx950_attn2d_build_ctx_t return ctx->wave_lds_off_i64_v; } -/* A copy covers copy_tokens whole tokens (KV_HALVES_PER_CALL is a multiple of - * HD). When they fit in one page the bf16/fp16 loaders address a copy as its - * page's base, its first token's offset in the page and the lane's own part; - * otherwise each wave's share still sits in one page and the loaders keep the - * paged-KV descriptor. The token stride is a kernarg, so the lane part and the +/* When a copy covers copy_tokens whole tokens and they fit in one page, the + * bf16/fp16 loaders address a copy as its page's base, its first token's + * offset in the page and the lane's own part. A copy that ends mid-row + * (head_size 192 with one warp: 512 halves are 2 2/3 rows) or spans pages + * keeps the paged-KV descriptor, which addresses each lane's own row. The token stride is a kernarg, so the lane part and the * per-copy token offsets are built once, after the paged-KV descriptor, rather * than multiplying the token index in every copy. */ static int rocke__copy_tokens(const rocke_gfx950_attn2d_build_ctx_t* ctx) @@ -197,7 +197,8 @@ static int rocke__copy_tokens(const rocke_gfx950_attn2d_build_ctx_t* ctx) static bool rocke__page_copies(const rocke_gfx950_attn2d_build_ctx_t* ctx) { - return !ctx->FAST_PAGED_KV_DESC && !ctx->KV_FP8 && ctx->BS % rocke__copy_tokens(ctx) == 0; + return !ctx->FAST_PAGED_KV_DESC && !ctx->KV_FP8 && rocke__kv_halves_per_call(ctx) % ctx->HD == 0 + && ctx->BS % rocke__copy_tokens(ctx) == 0; } /* A zeroed table of `n` IR values on the builder arena, sized from the spec 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 82e14b6674c7..aaecbd0e8e0c 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 @@ -1331,7 +1331,6 @@ "19": "8dff4368d397f6a9e209fb74ab18c6ecb183d1b83d628099ba313247509ee385", "2": "7f76dbe0d5e41e601888811064883de1a13d712d9ce1f5feae1f1e4c3eb18de0", "20": "74d73e46de452e86c8feaaa1be020dd556fce4b54e4b57419036a440c960e936", - "21": "40c87f693fa6ef5ca428df0d1b74539759eadea8381bf65465a0687be7661186", "22": "4f009d277496d367a65b276f7a4756a514bc555ac1380275c2edc1aefcd072ad", "23": "6545559600bd6b1f66e0f7847b8a57ec453407892c650f10b5cbc5ce61d4aac3", "24": "e786f0a81576ff0dd372ed0ff4f44532441fd4fb1fdb3bc174e48e313b4464ae", @@ -1464,6 +1463,7 @@ "69": "e2c2040c8e161845722093b62e9d173e79947e4f4a22e40d212149fc9e126bc7", "7": "1067ddfadbdb481e9623ec4389070eca0a385d2968de4d5218cdaf62f68481f0", "70": "6b98a1c3fe47f3794abb2eac2c2be60dddd0030b9247848df12867450260a76b", + "71": "f072b9b61cc857838d5cddd76a6b9796a00c613392c172c6e837a20f76969019", "8": "f4f0267af6f85a06591533f53a3cc74339723fa2e33c5f68e5fcd1b1ac4a6520", "9": "3c0bbb4435538d794f348fb385377e07b8274d96162257c6117e15e995afdf25" }, @@ -3116,7 +3116,6 @@ "19": "0cb1d827d6464648426db3ccf0e91f9f40f307a5a0fc5b7c91003ce9a42a90f1", "2": "c48d213c51a2cb49f42b8e2e483bae3d60cabea50358d34cac100a023754b4c3", "20": "508c96b15f0e4feebbea38f924341ba6d6678439635d8f74f648defffe495731", - "21": "8ecbb98a18c4e531ce7acf9782c6c2b2e0145dd6fb0e00463f4ddbbdf3549b94", "22": "acc7d51ac2c6c950d0c2f086d14ebfdddfc5c9cad8101b00f2739d90ee9090d1", "23": "93bd007f170e9fee09ac8d8cc421d7596c8a3698ba34d144c0c007b4555c6294", "24": "01e6a1de28c7824d8089216ef73eb70713b01fccf81e5a7d3a0b1730b5035edd", @@ -3249,6 +3248,7 @@ "69": "63b458ff53037edd27126db2f8028159d48239199c9494522bfedc95fa54727f", "7": "763832662414ae026dad9670f8f064ceea6145dd42ec2159941d37f3927ed2f7", "70": "2a3870347ab70cbcb276925da1726e5c51dcb737771d9d2af5e742a8c22680fe", + "71": "2e6df5bd9a6d6430bdd7cec8fb061975b840cc885565eb058307102a26e52b8b", "8": "436a62f0670b97e464439f466724c5d911a96b10e4bb035429b37d5be8d85243", "9": "07eadd3cdcb25d33e4162b75d3543681a9f260e6a972c64ce4bfe3a28dee2bf9" }, @@ -4900,7 +4900,6 @@ "19": "0bb698f1e7d115fc6e1b441348c9bc56f4a23a77fdc70b0f4f6dced8d77e5ecb", "2": "6f3d7370e78f1dbdae597abc516b6d19f43c4f608f9188a9710462d37b1747bc", "20": "84e4242ce5816d3dfa10bfa1ffab5095427764f228bfce69ffa758c875e397d9", - "21": "eac4ebb36eb69acec64f300b320fd959a3f52adfc444f8874a5271b2c3967795", "22": "2da29508514cb9e860c70addc2e17ed2f45e4d805524c181679cb16c68e427a5", "23": "fcf1ba8c03dcd55dbfdc62779924908131a15b5723cc43efc1ef73b5eab6cb4f", "24": "9e680483e5d92516aeaa918728a6fa56cb0e642f6165dd0b1bbc22f6f1c40b12", @@ -5033,6 +5032,7 @@ "69": "2045c5ff62259cd0fa0883a039ef937c7245952c9fd75520b8125f1cdd335ef7", "7": "8212f1103d5177844f6a8f6ae2e4beb73c28486bb07d975da6aa94cbac407c41", "70": "6a08f23a301f5f3025c4f0aa673f1848a6a91af57b81f853f3150e347b6fd164", + "71": "b8eff71e24ee9632e4c37ab8c589975411b851bc1abaaf0e47e4a2dada05ce10", "8": "7d827babcb8192445726b06e4250deb674158b5a77d87ecdd9e873118dc0ec93", "9": "5454f3658b3f3e436f29eb8a9f7184ed12eb0209155508c15d6d0b09eae9e2b0" }, @@ -6832,7 +6832,6 @@ "19": "9270dc04118c3c54288bb0f26595dee9627c0193a553f23d6ae0e583f063ed0f", "2": "d2287389d440d609f01fcde8f6a5eafd6a52175118b1722d175a92044cd859ce", "20": "d31da8b6037b65e794d03ec8763e392787cdc6e453bd1a005261c42ff29d0fc0", - "21": "27718afc2eda45622dae1760aedf641a3a71e1b6917de22388951f9817eb92f1", "22": "8b0386f0372e9380c7878b110e94c69e95b63acbcc9de6819982d331a1e6a5b1", "23": "0de5ae39bff1832101f673c9e89bc157e6ecc82be7bc67cd97b7897e43d99ce4", "24": "98b7af693c24567312660088a8846a5bcb52e8f235e9cfefc302a19ce5951b7f", @@ -6965,6 +6964,7 @@ "69": "1b2746a0d11d9d22993c5f58314b7eaf913926a7b80f3c28ac96c4df8824b22f", "7": "9f8857f7a94b9b8524bd631b0b1fcd3c89d8733d317d2ccfd9e26af3ff72e3ea", "70": "2d702df8314f65d553d8f0e1d72d0d202625361d26ec2a361f7ff3dd012ea947", + "71": "a1e8c2259d9baa180cd31a07c3974ff01d000af87bc24a21a7e4b3a157e939ab", "8": "9a4431d34fde7e376552fc757ff75be86b1c4c6194b3f95c93c0ab9fbbf9f2a4", "9": "6267b3845952589a6f3594a1e07902bfd52ec3ecfaa85d398c4b002570f37e07" }, @@ -9074,7 +9074,6 @@ "19": "9270dc04118c3c54288bb0f26595dee9627c0193a553f23d6ae0e583f063ed0f", "2": "d2287389d440d609f01fcde8f6a5eafd6a52175118b1722d175a92044cd859ce", "20": "d31da8b6037b65e794d03ec8763e392787cdc6e453bd1a005261c42ff29d0fc0", - "21": "27718afc2eda45622dae1760aedf641a3a71e1b6917de22388951f9817eb92f1", "22": "8b0386f0372e9380c7878b110e94c69e95b63acbcc9de6819982d331a1e6a5b1", "23": "0de5ae39bff1832101f673c9e89bc157e6ecc82be7bc67cd97b7897e43d99ce4", "24": "98b7af693c24567312660088a8846a5bcb52e8f235e9cfefc302a19ce5951b7f", @@ -9207,6 +9206,7 @@ "69": "1b2746a0d11d9d22993c5f58314b7eaf913926a7b80f3c28ac96c4df8824b22f", "7": "9f8857f7a94b9b8524bd631b0b1fcd3c89d8733d317d2ccfd9e26af3ff72e3ea", "70": "2d702df8314f65d553d8f0e1d72d0d202625361d26ec2a361f7ff3dd012ea947", + "71": "a1e8c2259d9baa180cd31a07c3974ff01d000af87bc24a21a7e4b3a157e939ab", "8": "9a4431d34fde7e376552fc757ff75be86b1c4c6194b3f95c93c0ab9fbbf9f2a4", "9": "6267b3845952589a6f3594a1e07902bfd52ec3ecfaa85d398c4b002570f37e07" }, From 32019169558886a57371c123b065c70d71b01cac Mon Sep 17 00:00:00 2001 From: ThruptiRajLakshmanaGowda Date: Wed, 4 Nov 2026 23:38:48 -0500 Subject: [PATCH 6/8] test(hip-kernel-provider): Unified3D catalog census follows the runtime kernargs The segment and reduce kernels gained runtime kernargs (num_query_heads, num_kv_heads and sliding_window on the segment, num_query_heads on the reduce), so the packed signatures are 22 and 6 arguments and the reduce's last is num_query_heads. TestGfx950AttentionUnified3DSignature already says so; the catalog census still expected 19 and 5. Co-Authored-By: Claude --- .../packs/TestGfx950AttentionUnified3DCatalogTiles.cpp | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/dnn-providers/hip-kernel-provider/src/tests/engines/kernel_ingestor_engine/packs/TestGfx950AttentionUnified3DCatalogTiles.cpp b/dnn-providers/hip-kernel-provider/src/tests/engines/kernel_ingestor_engine/packs/TestGfx950AttentionUnified3DCatalogTiles.cpp index c5b2c41ab46d..98193ff466d4 100644 --- a/dnn-providers/hip-kernel-provider/src/tests/engines/kernel_ingestor_engine/packs/TestGfx950AttentionUnified3DCatalogTiles.cpp +++ b/dnn-providers/hip-kernel-provider/src/tests/engines/kernel_ingestor_engine/packs/TestGfx950AttentionUnified3DCatalogTiles.cpp @@ -40,8 +40,8 @@ using namespace hip_kernel_provider::kernel_ingestor_engine; using namespace hip_kernel_provider::kernel_ingestor_engine::testing; namespace ingestor = hipdnn_plugin_sdk::ingestor; -constexpr std::size_t SEGMENT_ARGUMENTS = 19; -constexpr std::size_t REDUCE_ARGUMENTS = 5; +constexpr std::size_t SEGMENT_ARGUMENTS = 22; +constexpr std::size_t REDUCE_ARGUMENTS = 6; std::optional integerField(const ingestor::KernelDescriptor& kernel, const char* field) { @@ -124,7 +124,7 @@ TEST(TestGfx950AttentionUnified3DCatalogTiles, EveryEntryLaunchesItsSegmentAndRe ASSERT_EQ(reduce.signature.size(), REDUCE_ARGUMENTS); EXPECT_EQ(segment.signature.front().name, "segm_output_ptr"); EXPECT_EQ(reduce.signature.front().name, "output_ptr"); - EXPECT_EQ(reduce.signature.back().name, "seq_lens_ptr"); + EXPECT_EQ(reduce.signature.back().name, "num_query_heads"); } } EXPECT_GT(checked, 0U) << "the loaded set holds no kernels to check"; From 77ef4f77897164c1b6556d464ccbd9ce37628def Mon Sep 17 00:00:00 2001 From: dlangbe Date: Wed, 4 Nov 2026 23:39:29 -0500 Subject: [PATCH 7/8] fix(rocke): run the tiled attention entry guards in CI and share one C++ head-count rule - The C++ head-count check moves into _fmha_common, like Python's check_head_counts: rocke_fmha_validate_common_spec calls it instead of repeating both texts, and the tiled entries include it from there. - The C++ fastKV register-P builder runs the head-count check before the shared scaffold divides by num_kv_heads, as the Python build raises. - New host CTest rocke_tiled_attention_entry_guards: the gfx942/gfx950 2D and 3D lowers and the fastKV register-P lower refuse 32/0 and 32/6 heads, and the fastKV lower and predicate refuse gfx942, with the Python text. The pytest twin needs the rocke_engine binding, which no CI lane builds. - Dense segment combine: combine_launch_abi() gives the reduce launch signature and runtime values; a torch-free test checks it against the reduce kernel's params, so the CPU lanes run it. Co-Authored-By: Claude --- ...ROCKE_ENGINE_test_categories_external.yaml | 2 + .../common/attention_dense_kv_split.py | 31 ++- ...st_attention_dense_kv_split_combine_abi.py | 51 +++-- ...ances.common.attention_unified_selectors.h | 10 - ...lper_rocke.instances.common._fmha_common.h | 11 + .../cpp/instances/common/_fmha_common.cpp | 55 ++--- .../common/attention_unified_selectors.cpp | 17 -- .../attention_tiled_2d_public_entry_glue.cpp | 1 + ...ed_3d_gfx942_attn_tiled_3d_public_glue.cpp | 2 +- ...attention_tiled_2d_fastkv_regp_helpers.cpp | 13 ++ .../attention_tiled_2d_public_entry_glue.cpp | 2 +- .../gfx950/attention_tiled_3d_public_glue.cpp | 2 +- .../platform/tests/cmake/InstanceTests.cmake | 22 ++ .../tiled_attention_entry_guards.cpp | 205 ++++++++++++++++++ 14 files changed, 343 insertions(+), 81 deletions(-) create mode 100644 dnn-providers/hip-kernel-provider/rocke/platform/tests/instances/tiled_attention_entry_guards.cpp diff --git a/dnn-providers/hip-kernel-provider/ROCKE_ENGINE_test_categories_external.yaml b/dnn-providers/hip-kernel-provider/ROCKE_ENGINE_test_categories_external.yaml index 885624add2f7..2b3fbe4c3602 100644 --- a/dnn-providers/hip-kernel-provider/ROCKE_ENGINE_test_categories_external.yaml +++ b/dnn-providers/hip-kernel-provider/ROCKE_ENGINE_test_categories_external.yaml @@ -74,6 +74,7 @@ test_categories: - "rocke_mx_scale_decode" - "rocke_buffer_rsrc_hip" - "rocke_tiled_attention_2d_reentrancy" + - "rocke_tiled_attention_entry_guards" - "rocke_unified_attention_scalar_lower_to_llvm" - "rocke_python_import_smoke" - "rocke_golden_static" @@ -133,6 +134,7 @@ test_categories: - "rocke_mx_scale_decode" - "rocke_buffer_rsrc_hip" - "rocke_tiled_attention_2d_reentrancy" + - "rocke_tiled_attention_entry_guards" - "rocke_unified_attention_scalar_lower_to_llvm" - "rocke_python_import_smoke" - "rocke_golden_static" diff --git a/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/attention_dense_kv_split.py b/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/attention_dense_kv_split.py index 736458e6738b..b1e6f2360416 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/attention_dense_kv_split.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/attention_dense_kv_split.py @@ -188,19 +188,33 @@ def combine_spec(spec, arch: str): ) -def launch_segment_combine(spec, arch: str, out, workspace, stream: int) -> None: - """Merge the segments in ``workspace`` into ``out`` ([B, Sq, Hq, D]).""" +def combine_launch_abi(spec, arch: str): + """The combine's reduce spec, launch signature and runtime scalar values. + + The reduce reads its runtime shape fields (``num_query_heads``) as kernargs, + so the signature carries them after the pointers and the values supply them. + Host only: no torch and no compile. + """ from kernels.common.attention_unified import ( _reduce_signature, _runtime_kernarg_fields, ) - from rocke.helpers.compile import compile_kernel - from rocke.runtime import KernelLauncher, LaunchConfig red = combine_spec(spec, arch) - spec_type, build = _combine_impl(arch) - # The reduce reads its runtime shape fields (num_query_heads) as kernargs. + spec_type, _ = _combine_impl(arch) runtime_fields = _runtime_kernarg_fields(spec_type) + signature = _reduce_signature(spec.dtype, runtime_fields=runtime_fields) + scalars = {name: int(getattr(red, name)) for name in runtime_fields} + return red, signature, scalars + + +def launch_segment_combine(spec, arch: str, out, workspace, stream: int) -> None: + """Merge the segments in ``workspace`` into ``out`` ([B, Sq, Hq, D]).""" + from rocke.helpers.compile import compile_kernel + from rocke.runtime import KernelLauncher, LaunchConfig + + red, signature, scalars = combine_launch_abi(spec, arch) + _, build = _combine_impl(arch) key = (arch, red) launcher = _COMBINE_LAUNCHER_CACHE.get(key) if launcher is None: @@ -210,14 +224,13 @@ def launch_segment_combine(spec, arch: str, out, workspace, stream: int) -> None launcher = KernelLauncher( hsaco=art.hsaco, kernel_name=art.kernel_name, - signature=_reduce_signature(spec.dtype, runtime_fields=runtime_fields), + signature=signature, ) _COMBINE_LAUNCHER_CACHE[key] = launcher import torch vals = dict(workspace) - for name in runtime_fields: - vals[name] = int(getattr(red, name)) + vals.update(scalars) vals["output_ptr"] = out # Unused by the combine body; any int32 buffer satisfies the ABI. vals["seq_lens_ptr"] = workspace["segm_max_ptr"].view(torch.int32) diff --git a/dnn-providers/hip-kernel-provider/rocke/library/tests/test_attention_dense_kv_split_combine_abi.py b/dnn-providers/hip-kernel-provider/rocke/library/tests/test_attention_dense_kv_split_combine_abi.py index 49d1e3e9d6fa..5f3bae4d121b 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/tests/test_attention_dense_kv_split_combine_abi.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/tests/test_attention_dense_kv_split_combine_abi.py @@ -8,6 +8,9 @@ from the pointer-only reduce signature leaves the head count, and with it the reduce's token strides, undefined. This captures the launcher's signature and values on the host and checks them against the reduce kernel's params. + +``test_combine_abi_matches_reduce_params`` needs neither torch nor a GPU, so it +runs in every CPU lane; the launch test runs where torch is installed. """ from __future__ import annotations @@ -17,7 +20,35 @@ from kernels.gfx942.attention_dense import Gfx942AttentionDenseSpec from kernels.gfx950.attention_dense import Gfx950AttentionDenseSpec -torch = pytest.importorskip("torch") +_ARCHES = [("gfx942", Gfx942AttentionDenseSpec), ("gfx950", Gfx950AttentionDenseSpec)] + + +def _dense_spec(spec_cls): + return spec_cls( + batch=2, + seqlen_q=256, + seqlen_kv=2048, + num_query_heads=16, + num_kv_heads=2, + head_size=128, + causal=True, + dtype="bf16", + kv_segments=4, + ) + + +@pytest.mark.parametrize("arch, spec_cls", _ARCHES) +def test_combine_abi_matches_reduce_params(arch, spec_cls): + """The signature the launcher is built from names the reduce kernel's params + in order, and every runtime scalar among them gets the combine's value.""" + spec = _dense_spec(spec_cls) + red, signature, scalars = kv_split.combine_launch_abi(spec, arch) + _, build = kv_split._combine_impl(arch) + params = [p.name for p in build(red, arch=arch).params] + assert params == [a["name"] for a in signature] + assert scalars == {"num_query_heads": spec.num_query_heads} + pointers = {"output_ptr", "seq_lens_ptr"} | set(kv_split.SEGMENT_PARAMS) + assert set(params) == pointers | set(scalars) class _Art: @@ -27,11 +58,9 @@ def __init__(self, kernel): self.kernel_name = kernel.name -@pytest.mark.parametrize( - "arch, spec_cls", - [("gfx942", Gfx942AttentionDenseSpec), ("gfx950", Gfx950AttentionDenseSpec)], -) +@pytest.mark.parametrize("arch, spec_cls", _ARCHES) def test_combine_launch_matches_reduce_params(monkeypatch, arch, spec_cls): + torch = pytest.importorskip("torch") import rocke.helpers.compile as compile_mod import rocke.runtime as runtime_mod @@ -52,17 +81,7 @@ def __call__(self, vals, *, config): monkeypatch.setattr(runtime_mod, "KernelLauncher", FakeLauncher) monkeypatch.setattr(kv_split, "_COMBINE_LAUNCHER_CACHE", {}) - spec = spec_cls( - batch=2, - seqlen_q=256, - seqlen_kv=2048, - num_query_heads=16, - num_kv_heads=2, - head_size=128, - causal=True, - dtype="bf16", - kv_segments=4, - ) + spec = _dense_spec(spec_cls) workspace = kv_split.allocate_segment_workspace(spec, "cpu") out = torch.empty( (spec.batch, spec.seqlen_q, spec.num_query_heads, spec.head_size), 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 31b6f8445ff5..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 @@ -91,16 +91,6 @@ typedef struct rocke_unified_attn_problem int rocke_unified_attn_num_queries_per_kv(rocke_ir_builder_t* b, const rocke_unified_attn_problem_t* p); -/* Python: kernels.common._fmha_common.check_head_counts. True when both head - * counts are positive and num_query_heads splits evenly over num_kv_heads; - * otherwise false, with the Python reason text written to `reason`. Every - * tiled attention entry runs it before naming the kernel (the name carries - * the GQA ratio). */ -bool rocke_attention_check_head_counts(int num_query_heads, - int num_kv_heads, - char* reason, - size_t reason_cap); - /* Python: _resolve_attention_arch() override hook. The Python memoizes the * device arch (falling back to "gfx950"); this lets a host pin the arch the * selectors resolve to (NULL restores the "gfx950" fallback). Pass a string diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/helper_rocke.instances.common._fmha_common.h b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/helper_rocke.instances.common._fmha_common.h index d0e0dfd9922e..ee5ae7fa4c6b 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/helper_rocke.instances.common._fmha_common.h +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/helper_rocke.instances.common._fmha_common.h @@ -175,6 +175,17 @@ static inline bool rocke_fmha_common_spec_use_custom_mask(const rocke_fmha_commo * ("ok" on the accept path); pass a live `arena`. The reason strings reproduce * the Python messages byte-for-byte. On a formatting OOM the predicate result is * still returned faithfully and the reason is a best-effort static fallback. */ +/* Python: kernels.common._fmha_common.check_head_counts. True when both head + * counts are positive and num_query_heads splits evenly over num_kv_heads; + * otherwise false, with the Python reason text written to `reason`. The one + * head-count rule: rocke_fmha_validate_common_spec and every tiled attention + * entry run it (the tiled entries before naming the kernel, whose name + * carries the GQA ratio). */ +bool rocke_attention_check_head_counts(int num_query_heads, + int num_kv_heads, + char* reason, + size_t reason_cap); + bool rocke_fmha_validate_common_spec(rocke_arena_t* arena, const rocke_fmha_common_spec_t* spec, const char** out_reason); diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/common/_fmha_common.cpp b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/common/_fmha_common.cpp index df4baf078c9e..1d9364bb7e64 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/common/_fmha_common.cpp +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/common/_fmha_common.cpp @@ -8,6 +8,7 @@ */ #include "rocke/helper_rocke.instances.common._fmha_common.h" +#include #include #include "rocke/helper_rocke.helpers.io.h" /* rocke_b_io_ir_type */ @@ -100,6 +101,23 @@ rocke_fmha_common_spec_t rocke_fmha_common_spec_default(rocke_fmha_shape_t shape /* Set *out_reason (if non-NULL) to an arena-printf'd message, falling back to * `fallback` on OOM. Returns `ret` so callers can `return reject(...)`. */ +bool rocke_attention_check_head_counts(int num_query_heads, + int num_kv_heads, + char* reason, + size_t reason_cap) +{ + const char* fmt = NULL; + if(num_query_heads <= 0 || num_kv_heads <= 0) + fmt = "num_query_heads / num_kv_heads must be > 0 (got %d, %d)"; + else if(num_query_heads % num_kv_heads != 0) + fmt = "num_query_heads (%d) must be divisible by num_kv_heads (%d) for GQA / MQA"; + if(fmt == NULL) + return true; + if(reason != NULL && reason_cap > 0) + (void)snprintf(reason, reason_cap, fmt, num_query_heads, num_kv_heads); + return false; +} + static bool rocke_i_fmha_set_reason(rocke_arena_t* arena, const char** out_reason, bool ret, @@ -157,33 +175,18 @@ bool rocke_fmha_validate_common_spec(rocke_arena_t* arena, s->head_size), "head_size not in the supported set"); } - /* if s.num_query_heads <= 0 or s.num_kv_heads <= 0: */ - if(s->num_query_heads <= 0 || s->num_kv_heads <= 0) + /* ok, reason = check_head_counts(s.num_query_heads, s.num_kv_heads) */ { - return rocke_i_fmha_set_reason( - arena, - out_reason, - false, - rocke_arena_printf(arena, - "num_query_heads / num_kv_heads must be > 0 " - "(got %d, %d)", - s->num_query_heads, - s->num_kv_heads), - "num_query_heads / num_kv_heads must be > 0"); - } - /* if s.num_query_heads % s.num_kv_heads != 0: */ - if((s->num_query_heads % s->num_kv_heads) != 0) - { - return rocke_i_fmha_set_reason( - arena, - out_reason, - false, - rocke_arena_printf(arena, - "num_query_heads (%d) must be divisible by " - "num_kv_heads (%d) for GQA / MQA", - s->num_query_heads, - s->num_kv_heads), - "num_query_heads must be divisible by num_kv_heads"); + char why[160]; + if(!rocke_attention_check_head_counts( + s->num_query_heads, s->num_kv_heads, why, sizeof(why))) + { + return rocke_i_fmha_set_reason(arena, + out_reason, + false, + rocke_arena_printf(arena, "%s", why), + "num_query_heads / num_kv_heads not a valid GQA pair"); + } } /* if s.block_size_q not in (16, 32, 64, 128): */ if(s->block_size_q != 16 && s->block_size_q != 32 && s->block_size_q != 64 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 01f3f2ae1ab6..720b4dae8355 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 @@ -115,23 +115,6 @@ int rocke_unified_attn_num_queries_per_kv(rocke_ir_builder_t* b, return p->num_query_heads / p->num_kv_heads; } -bool rocke_attention_check_head_counts(int num_query_heads, - int num_kv_heads, - char* reason, - size_t reason_cap) -{ - const char* fmt = NULL; - if(num_query_heads <= 0 || num_kv_heads <= 0) - fmt = "num_query_heads / num_kv_heads must be > 0 (got %d, %d)"; - else if(num_query_heads % num_kv_heads != 0) - fmt = "num_query_heads (%d) must be divisible by num_kv_heads (%d) for GQA / MQA"; - if(fmt == NULL) - return true; - if(reason != NULL && reason_cap > 0) - (void)snprintf(reason, reason_cap, fmt, num_query_heads, num_kv_heads); - return false; -} - /* Internal convenience: same value without a builder (the selectors call the * property freely; the divisibility precondition is guaranteed for any problem * that reached the spec stage). */ diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/gfx942/attention_tiled_2d_public_entry_glue.cpp b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/gfx942/attention_tiled_2d_public_entry_glue.cpp index 40adb961e96d..79e3128f5fb6 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/gfx942/attention_tiled_2d_public_entry_glue.cpp +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/gfx942/attention_tiled_2d_public_entry_glue.cpp @@ -40,6 +40,7 @@ #include "rocke/helper_helper_rocke.instances.common.attention_unified_selectors.h" /* rocke_unified_attn_tiled_2d_lds_bytes */ #include "rocke/helper_rocke.helpers.spec.h" /* rocke_kernel_name_join */ +#include "rocke/helper_rocke.instances.common._fmha_common.h" /* rocke_attention_check_head_counts */ #include "rocke/arena.h" /* rocke_arena_strdup */ #include "rocke/error_boundary.hpp" /* ckc::guard_builder boundary shim */ diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/gfx942/attention_tiled_3d_gfx942_attn_tiled_3d_public_glue.cpp b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/gfx942/attention_tiled_3d_gfx942_attn_tiled_3d_public_glue.cpp index 6933a1049246..5b6f049e7f4f 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/gfx942/attention_tiled_3d_gfx942_attn_tiled_3d_public_glue.cpp +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/gfx942/attention_tiled_3d_gfx942_attn_tiled_3d_public_glue.cpp @@ -47,11 +47,11 @@ #include "rocke/arena.h" /* rocke_arena_strdup */ #include "rocke/error_boundary.hpp" /* ckc::guard_builder boundary shim */ -#include "rocke/helper_helper_rocke.instances.common.attention_unified_selectors.h" /* rocke_attention_check_head_counts */ #include "rocke/helper_rocke.core.arch.h" /* rocke_archtarget_from_gfx, op_for_shape */ #include "rocke/helper_rocke.helpers.atoms.h" /* rocke_mfma_atom, make_c_warp_dstr_encoding */ #include "rocke/helper_rocke.helpers.distribution.h" /* make_static_tile_distribution */ #include "rocke/helper_rocke.helpers.spec.h" /* rocke_kernel_name_join */ +#include "rocke/helper_rocke.instances.common._fmha_common.h" /* rocke_attention_check_head_counts */ #include "rocke/instance_gfx942_attention_tiled_3d.h" #include "rocke/instance_gfx942_attention_tiled_3d_internal.h" #include "rocke/ir.h" diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/gfx950/attention_tiled_2d_fastkv_regp_helpers.cpp b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/gfx950/attention_tiled_2d_fastkv_regp_helpers.cpp index 3717be914642..ad988823cddc 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/gfx950/attention_tiled_2d_fastkv_regp_helpers.cpp +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/gfx950/attention_tiled_2d_fastkv_regp_helpers.cpp @@ -9,6 +9,7 @@ #include "rocke/error.hpp" #include "rocke/helper_instance_gfx950_attention_tiled_2d_fastkv_regp.h" #include "rocke/helper_rocke.helpers.spec.h" /* rocke_kernel_name_join */ +#include "rocke/helper_rocke.instances.common._fmha_common.h" /* rocke_attention_check_head_counts */ #include "rocke/error_boundary.hpp" /* ckc::guard_builder boundary shim */ #include @@ -345,6 +346,18 @@ rocke_kernel_def_t* rocke_build_unified_attention_2d_fastkv_register_p( b, ROCKE_ERR_VALUE, "fastKV register-P experiment does not support FP8 KV cache"); return NULL; } + /* The tiled build reads spec.num_queries_per_kv, which raises ValueError + * with the head-count rule's text; check it here, before the shared + * scaffold divides by num_kv_heads. */ + { + char why[160]; + if(!rocke_attention_check_head_counts( + spec->num_query_heads, spec->num_kv_heads, why, sizeof(why))) + { + rocke__fastkv_regp_set_err(b, ROCKE_ERR_VALUE, why); + return NULL; + } + } /* return build_unified_attention_2d_tiled(_FastKvRegisterPProxy(spec), arch) * diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/gfx950/attention_tiled_2d_public_entry_glue.cpp b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/gfx950/attention_tiled_2d_public_entry_glue.cpp index acf6d1d308ac..3fd542998938 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/gfx950/attention_tiled_2d_public_entry_glue.cpp +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/gfx950/attention_tiled_2d_public_entry_glue.cpp @@ -45,9 +45,9 @@ #include "rocke/arena.h" /* rocke_arena_strdup */ #include "rocke/error_boundary.hpp" /* ckc::guard_builder boundary shim */ #include "rocke/helper_helper_rocke.helpers.attention.h" /* binary_search_seq_idx */ -#include "rocke/helper_helper_rocke.instances.common.attention_unified_selectors.h" /* rocke_attention_check_head_counts */ #include "rocke/helper_helper_rocke.instances.gfx942.attention_tiled_2d.h" /* rocke__mfma_32x32_c_* */ #include "rocke/helper_rocke.helpers.spec.h" /* rocke_kernel_name_join */ +#include "rocke/helper_rocke.instances.common._fmha_common.h" /* rocke_attention_check_head_counts */ #include #include diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/gfx950/attention_tiled_3d_public_glue.cpp b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/gfx950/attention_tiled_3d_public_glue.cpp index 257cf526a7d7..80fe704bcddc 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/gfx950/attention_tiled_3d_public_glue.cpp +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/gfx950/attention_tiled_3d_public_glue.cpp @@ -72,7 +72,7 @@ * _default() definitions below -- the gfx942 TU owns them. The structs are * byte-identical so the single shared definition is correct for both. */ #include "rocke/error_boundary.hpp" /* ckc::guard_builder boundary shim */ -#include "rocke/helper_helper_rocke.instances.common.attention_unified_selectors.h" /* rocke_attention_check_head_counts */ +#include "rocke/helper_rocke.instances.common._fmha_common.h" /* rocke_attention_check_head_counts */ #include "rocke/instance_gfx942_attention_tiled_3d.h" #include "rocke/instance_gfx950_attention_tiled_3d.h" #include "rocke/instance_gfx950_attention_tiled_3d_internal.h" diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/tests/cmake/InstanceTests.cmake b/dnn-providers/hip-kernel-provider/rocke/platform/tests/cmake/InstanceTests.cmake index b1278a892ef4..47210d22e3f5 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/tests/cmake/InstanceTests.cmake +++ b/dnn-providers/hip-kernel-provider/rocke/platform/tests/cmake/InstanceTests.cmake @@ -24,6 +24,28 @@ if(EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/instances/tiled_attention_2d_reentrancy.c rocke_tiled_attention_2d_reentrancy) endif() +# The C tiled-attention entries refuse bad head counts, and fastKV register-P +# on gfx942, with the Python reason text (host only, no GPU). +if(EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/instances/tiled_attention_entry_guards.cpp") + add_executable(rocke_tiled_attention_entry_guards + instances/tiled_attention_entry_guards.cpp) + target_link_libraries(rocke_tiled_attention_entry_guards PRIVATE rocke_core) + if(UNIX) + target_link_libraries(rocke_tiled_attention_entry_guards PRIVATE m) + endif() + add_test(NAME rocke_tiled_attention_entry_guards + COMMAND rocke_tiled_attention_entry_guards) + set_tests_properties(rocke_tiled_attention_entry_guards + PROPERTIES LABELS "ckc;unit;host") + rocke_cxx_test_installed_name(_rocke_installed_name rocke_tiled_attention_entry_guards) + set_target_properties(rocke_tiled_attention_entry_guards + PROPERTIES OUTPUT_NAME "${_rocke_installed_name}") + install(TARGETS rocke_tiled_attention_entry_guards + RUNTIME DESTINATION ${ROCKE_CXX_TEST_INSTALL_DIR}) + set_property(GLOBAL APPEND PROPERTY ROCKE_INSTALL_TEST_TARGETS + rocke_tiled_attention_entry_guards) +endif() + # The three scalar-attention _lower_to_llvm wrappers take no target, so they # lower for a fixed one while the lowerer itself rejects a NULL arch. Nothing # else calls them, so this is the only thing that notices when they stop diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/tests/instances/tiled_attention_entry_guards.cpp b/dnn-providers/hip-kernel-provider/rocke/platform/tests/instances/tiled_attention_entry_guards.cpp new file mode 100644 index 000000000000..62ada3bf0c58 --- /dev/null +++ b/dnn-providers/hip-kernel-provider/rocke/platform/tests/instances/tiled_attention_entry_guards.cpp @@ -0,0 +1,205 @@ +// Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +// SPDX-License-Identifier: MIT +/* + * Host-only test that the C tiled-attention entries refuse what the Python + * builders refuse, with the Python reason text, instead of crashing: + * + * - a zero or non-dividing head pair, on the gfx942/gfx950 2D and 3D lowers + * and the gfx950 fastKV register-P lower (the kernel names carry the GQA + * ratio, and the shared 2D scaffold divides by num_kv_heads); + * - fastKV register-P on gfx942, on its lower and its support predicate (the + * body is gfx950-only). + * + * The expected texts are the Python ones: kernels.common._fmha_common + * check_head_counts and the fastKV module's _NOT_GFX950. The pytest twin, + * test_attention_tiled_cpp_entry_guards.py, compares the two engines directly + * where the rocke_engine binding is built. + */ +#include +#include +#include + +#include "rocke/instance_gfx942_attention_tiled_2d.h" +#include "rocke/instance_gfx942_attention_tiled_3d.h" +#include "rocke/instance_gfx950_attention_tiled_2d.h" +#include "rocke/instance_gfx950_attention_tiled_2d_fastkv_regp.h" +#include "rocke/instance_gfx950_attention_tiled_3d.h" +#include "rocke/lower_llvm.h" + +static const char* const NOT_POSITIVE = "num_query_heads / num_kv_heads must be > 0 (got 32, 0)"; +static const char* const NOT_DIVIDING + = "num_query_heads (32) must be divisible by num_kv_heads (6) for GQA / MQA"; +static const char* const NOT_GFX950 + = "fastKV register-P is a gfx950 kernel; gfx942 is not supported"; + +typedef struct head_pair +{ + int num_query_heads; + int num_kv_heads; + const char* reason; +} head_pair_t; + +static const head_pair_t BAD_PAIRS[] = {{32, 0, NOT_POSITIVE}, {32, 6, NOT_DIVIDING}}; + +/* Fails unless the call was refused (non-OK, no output) with `reason` in its + * diagnostic. Frees any output it was handed. */ +static int expect_refused( + const char* what, rocke_status_t st, char* ll, const char* err, const char* reason) +{ + const bool refused = st != ROCKE_OK && ll == NULL; + const bool says_why = strstr(err, reason) != NULL; + free(ll); + if(refused && says_why) + { + printf("%-40s refused: %s\n", what, err); + return 0; + } + fprintf(stderr, + "%-40s expected a refusal with \"%s\", got status=%d ll=%s err=\"%s\"\n", + what, + reason, + (int)st, + ll != NULL ? "" : "NULL", + err); + return 1; +} + +static rocke_attention_tiled_2d_spec_t spec_2d(int num_query_heads, int num_kv_heads) +{ + rocke_attention_tiled_2d_spec_t s = rocke_attention_tiled_2d_spec_default(); + s.head_size = 128; + s.block_size = 16; + s.num_query_heads = num_query_heads; + s.num_kv_heads = num_kv_heads; + s.dtype = "bf16"; + s.use_sinks = false; + s.sliding_window = 0; + s.has_softcap = false; + return s; +} + +/* make_fastkv_register_p_spec over the experiment's bf16 d64_b32 shape. */ +static rocke_attention_tiled_2d_spec_t spec_fastkv(int num_query_heads, int num_kv_heads) +{ + rocke_attention_tiled_2d_spec_t s = spec_2d(num_query_heads, num_kv_heads); + s.head_size = 64; + s.block_size = 32; + s.use_fast_paged_kv_desc = true; + s.use_mfma_32x32 = true; + s.use_transposed_qk_32x32 = true; + return s; +} + +static rocke_unified_attention_3d_tiled_spec_t spec_3d(int num_query_heads, int num_kv_heads) +{ + rocke_unified_attention_3d_tiled_spec_t s = rocke_unified_attention_3d_tiled_spec_default(); + s.head_size = 128; + s.block_size = 16; + s.num_query_heads = num_query_heads; + s.num_kv_heads = num_kv_heads; + s.dtype = "bf16"; + s.use_sinks = false; + s.sliding_window = 0; + s.has_softcap = false; + s.num_segments = 16; + return s; +} + +static int check_head_counts(void) +{ + int failures = 0; + for(size_t i = 0; i < sizeof(BAD_PAIRS) / sizeof(BAD_PAIRS[0]); ++i) + { + const head_pair_t* p = &BAD_PAIRS[i]; + char what[64]; + char err[512]; + char* ll; + rocke_status_t st; + + const rocke_attention_tiled_2d_spec_t s2 = spec_2d(p->num_query_heads, p->num_kv_heads); + const rocke_unified_attention_3d_tiled_spec_t s3 + = spec_3d(p->num_query_heads, p->num_kv_heads); + const rocke_attention_tiled_2d_spec_t sf = spec_fastkv(p->num_query_heads, p->num_kv_heads); + + ll = NULL; + err[0] = '\0'; + st = rocke_gfx942_attention_tiled_2d_lower_to_llvm( + &s2, "gfx942", ROCKE_LLVM_FLAVOR_AUTO, &ll, err, sizeof(err)); + snprintf(what, sizeof(what), "gfx942 2D %d/%d", p->num_query_heads, p->num_kv_heads); + failures += expect_refused(what, st, ll, err, p->reason); + + ll = NULL; + err[0] = '\0'; + st = rocke_gfx950_attention_tiled_2d_lower_to_llvm( + &s2, "gfx950", ROCKE_LLVM_FLAVOR_AUTO, &ll, err, sizeof(err)); + snprintf(what, sizeof(what), "gfx950 2D %d/%d", p->num_query_heads, p->num_kv_heads); + failures += expect_refused(what, st, ll, err, p->reason); + + ll = NULL; + err[0] = '\0'; + st = rocke_build_unified_attention_3d_tiled_gfx942_lower_to_llvm( + &s3, "gfx942", ROCKE_LLVM_FLAVOR_AUTO, &ll, err, sizeof(err)); + snprintf(what, sizeof(what), "gfx942 3D %d/%d", p->num_query_heads, p->num_kv_heads); + failures += expect_refused(what, st, ll, err, p->reason); + + ll = NULL; + err[0] = '\0'; + st = rocke_build_unified_attention_3d_tiled_gfx950_lower_to_llvm( + &s3, "gfx950", ROCKE_LLVM_FLAVOR_AUTO, &ll, err, sizeof(err)); + snprintf(what, sizeof(what), "gfx950 3D %d/%d", p->num_query_heads, p->num_kv_heads); + failures += expect_refused(what, st, ll, err, p->reason); + + ll = NULL; + err[0] = '\0'; + st = rocke_gfx950_attention_tiled_2d_fastkv_regp_lower_to_llvm( + &sf, "gfx950", ROCKE_LLVM_FLAVOR_AUTO, &ll, err, sizeof(err)); + snprintf( + what, sizeof(what), "gfx950 fastKV regp %d/%d", p->num_query_heads, p->num_kv_heads); + failures += expect_refused(what, st, ll, err, p->reason); + } + return failures; +} + +static int check_fastkv_arch(void) +{ + int failures = 0; + char err[512]; + char* ll = NULL; + const rocke_attention_tiled_2d_spec_t s = spec_fastkv(64, 8); + + err[0] = '\0'; + const rocke_status_t st = rocke_gfx950_attention_tiled_2d_fastkv_regp_lower_to_llvm( + &s, "gfx942", ROCKE_LLVM_FLAVOR_AUTO, &ll, err, sizeof(err)); + failures += expect_refused("gfx950 fastKV regp lower on gfx942", st, ll, err, NOT_GFX950); + + /* supports_fastkv_register_p_2d(arch="gfx942") over a shape it otherwise admits. */ + rocke_gfx950_supports_fastkv_register_p_2d_args_t a; + memset(&a, 0, sizeof(a)); + a.head_size = 64; + a.block_size = 32; + a.dtype = "bf16"; + a.num_queries_per_kv = 8; + a.num_query_heads = 64; + a.num_kv_heads = 8; + a.arch = "gfx942"; + err[0] = '\0'; + const bool ok = rocke_gfx950_supports_fastkv_register_p_2d(&a, err, sizeof(err)); + failures += expect_refused("gfx950 fastKV regp predicate on gfx942", + ok ? ROCKE_OK : ROCKE_ERR_VALUE, + NULL, + err, + NOT_GFX950); + return failures; +} + +int main(void) +{ + const int failures = check_head_counts() + check_fastkv_arch(); + if(failures != 0) + { + fprintf(stderr, "%d TILED ATTENTION ENTRY GUARD CHECK(S) FAILED\n", failures); + return 1; + } + printf("ALL TILED ATTENTION ENTRY GUARDS REFUSED WITH THE PYTHON REASON\n"); + return 0; +} From ffcfd98dd6485ddc23b1b7ec32fdc06c9b31d802 Mon Sep 17 00:00:00 2001 From: dlangbe Date: Wed, 4 Nov 2026 23:40:09 -0500 Subject: [PATCH 8/8] perf(rocke): cut the tiled cache-key sweep from 109 s to 39 s - 2D on gfx942/gfx950 sweeps num_seqs 1 and 2: a body that searches cu_q at runtime reads num_seqs only as the single-batch cohort. 3D (segment sizing reads the batch) and gfx1250 (baked search) keep 1 to 4. - The routing-env states pool on 2D only; on 3D every flipped state reproduced the default sweep's keys and kernels. - _norm_spec reads the spec fields directly instead of dataclasses.asdict. - validate_tiled_attention_arch is cached per arch string; every tiled support gate calls it and it depends only on the arch. Dropping each of the 44 2D key entries in turn, the trimmed sweep catches the same 12 (gfx942) and 18 (gfx950) as the full one. Co-Authored-By: Claude --- .../library/kernels/common/attention_arch.py | 4 +++ .../tests/test_attention_tiled_cache_key.py | 31 ++++++++++++++----- 2 files changed, 27 insertions(+), 8 deletions(-) diff --git a/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/attention_arch.py b/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/attention_arch.py index 394e67bcd6a8..28aa51aaa4c0 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/attention_arch.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/attention_arch.py @@ -37,6 +37,7 @@ from __future__ import annotations +import functools from typing import Tuple from rocke.core.arch import ArchTarget @@ -83,6 +84,9 @@ def _narrow_k_mfma_available(target: ArchTarget) -> bool: ) +# A function of the arch string alone (the catalog is static), and every tiled +# support gate calls it, so it is computed once per arch. +@functools.lru_cache(maxsize=None) def validate_tiled_attention_arch(arch: str) -> Tuple[bool, str]: """Return ``(ok, reason)`` for running a tiled attention kernel on ``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 fbfa072837d8..749cc99b737c 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 @@ -85,7 +85,9 @@ def _norm_spec(spec): """Spec identity as the emitter sees it: ``num_seqs`` only reaches IR through ``binary_search_iters``, which the 3D decode grid does not emit and a runtime-search body (``RUNTIME_SEARCH_ITERS``) does not bake.""" - fields = dataclasses.asdict(spec) + # The spec fields are scalars and strings, so a shallow read is + # ``dataclasses.asdict`` without its recursive copy. + fields = {f.name: getattr(spec, f.name) for f in dataclasses.fields(spec)} fields.pop("num_seqs") bakes_search = not getattr(type(spec), "RUNTIME_SEARCH_ITERS", False) if bakes_search and not getattr(spec, "use_decode_grid", False): @@ -138,6 +140,11 @@ def _built_3d(problem, arch): # 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) +# A 2D body that searches cu_q at runtime reads num_seqs only as the +# single-batch cohort (1 vs >= 2), so its sweep stops at 2. The 3D split-KV +# selectors size segments from the batch, and gfx1250 bakes the search, so +# both keep every value. +_NUM_SEQS_RUNTIME_SEARCH_2D = (1, 2) _SEQLEN_Q = (1, 2, 64, 65, 256, 257, 512, 513, 768, 769, 1024, 1025, 1536, 2048) _SEQLEN_K = ("q", 512, 513, 1024, 1025) @@ -180,9 +187,14 @@ def _static_grid(arch): ) -def _sweep(arch): +def _sweep(arch, path): + num_seqs = ( + _NUM_SEQS_RUNTIME_SEARCH_2D + if path == "2d" and au._searches_at_runtime(arch, "2d") + else _NUM_SEQS + ) for static in _static_grid(arch): - for ns, q, k in itertools.product(_NUM_SEQS, _SEQLEN_Q, _SEQLEN_K): + for ns, q, k in itertools.product(num_seqs, _SEQLEN_Q, _SEQLEN_K): k = q if k == "q" else k if k < q: continue @@ -191,9 +203,11 @@ def _sweep(arch): ) -# Env states to pool per arch. Routing knobs are read on every launch and can -# flip mid-process; pooling them catches a selector whose output only changes -# when a knob flips. +# Env states to pool per arch on the 2D path. Routing knobs are read on every +# launch and can flip mid-process; pooling them catches a selector whose output +# only changes when a knob flips. These knobs steer 2D routing only: on the 3D +# path every flipped state reproduces the default sweep's keys and kernels, so +# 3D sweeps once. _ENV_STATES = { "gfx942": ( {}, @@ -235,10 +249,11 @@ def test_key_matches_build_identity(routing_env, monkeypatch, arch, path): first = {} key_of_build = {} checked = 0 - for env in _ENV_STATES.get(arch, ({},)): + env_states = _ENV_STATES.get(arch, ({},)) if path == "2d" else ({},) + for env in env_states: for var, val in env.items(): monkeypatch.setenv(var, val) - for problem in _sweep(arch): + for problem in _sweep(arch, path): if not supported(problem, arch)[0]: continue try: