Repository navigation
fix(rocke): carry the raw running max in unified attention - #14542
Merged
AviralGoelAMD merged 30 commits intoOct 31, 2026
Merged
Conversation
…2D C++; register the masked-tile tests
The empty-segment guards compared the segment sum and the final denominator with an ordered greater-than, which is false for NaN, so a NaN sum was dropped as an empty segment and a NaN denominator gave O = 0. Treat only an exact zero as empty so NaN propagates as before.
Run every epilogue branch at page sizes 16 and 64, using the 1x-block gfx942 variants at page size 64 where the 2x-block ones exceed LDS, assert each geometry still gives a row a fully masked leading tile, and correct the fp16/bf16 P underflow thresholds in the docstring.
…erage Scalar 3D: write NaN into one segment sum of a row that sees keys between the segment and reduce launches; that row's O must be NaN and every other row must match the reference. An ogt-style empty-segment guard would skip the NaN segment and leave the row finite. Tiled 2D: run the pinned epilogue branches at head sizes 64, 128 and 256 (256 skipped, with the canonicalizer's reason, where the variant exceeds the LDS budget), and add one auto-route D64 bf16 case per arch whose production kernel is the multi-warp default builder.
Build the scalar split-3D reduce in the representative IR golden
(scalar_reduce_{fp16,bf16}_d128_seg8) and in the attention_unified C/Python
parity emitter pair (configs 6-7), so check_byte_identity compares the
reduce between the C++ and Python engines. Goldens gain only the new keys.
Also fix a gfx950 comment that named m_new as the state handed to the PV
bucket; it is m_carry, the raw running max.
…ry mask The scalar 2D epilogue guarded acc * rcp(l) only for non-default masks, on the claim that the default mask never yields a row without keys. Bottom-right causal with q_len > kv_len does: its leading rows see no key, l is 0 and O was NaN. The C++ twin had no guard at all. Guard every mask, and test l == 0 (fcmp oeq) instead of l > 0 (fcmp ogt) so a NaN sum still reaches O instead of being stored as a keyless 0. Mirror the same ops in the C++ engine and re-bless the representative and differential goldens for the scalar 2D cases. Tests: a default-mask q_len > kv_len case whose keyless rows must be 0, and a NaN query-row probe that must give NaN O under the default and top-left masks.
The tiled 2D epilogue on gfx942 and gfx950 stored 0 for any row whose softmax sum failed 'l > 0'. A NaN sum fails that test too, so a corrupt row came out as a finite 0. Store 0 only for an exact-zero sum (fcmp oeq), the rule the scalar 2D epilogue and the scalar reduce use. Rows with a finite nonzero sum are unchanged. Python builders and C++ twins change in step; the fastKV register-P variant inherits the change. Also correct the scalar 2D and scalar 3D comments about keyless rows, and re-bless the representative and differential goldens of the tiled 2D families.
Add a NaN-query-row case per tiled 2D epilogue branch, state what each mask case of the scalar 2D NaN test pins, and drop the unused segm_max argument of the scalar 3D workspace hook.
github-actions
Bot
added
project: hip-kernel-provider
rocke-lib: attention
✅ All Checks Passed — Ready for Review
📖 Need help? See the Policy FAQ for details on every check and how to fix failures. |
|
🎉 All checks passed! This PR is ready for review. |
…arry Resolve the differential golden conflict with #14501 by regeneration: take develop's file, then re-record the attention_tiled_2d and attention_unified families (ll llvm20/22/23 + ir canonical, python-only). Other families keep develop's entries.
…arry Resolve the representative IR golden conflict with #14357 (fp64 GEMM on gfx942) by regeneration: take develop's file, then re-record it with rocke_ir_parity_harness.py --write on the merged tree. Every entry equals develop's or this branch's hash. The differential golden and the parity harness merged cleanly, and check-golden is OK at every flavor.
AviralGoelAMD
added a commit
that referenced
this pull request
Oct 29, 2026
…nified-2d-lse Bring in #14542's block_table_stride fix for the scalar 3D test launch and its merge of develop (#14357, fp64 GEMM on gfx942). Resolve the representative IR golden conflict by regeneration: take #14542's file, then re-record it with rocke_ir_parity_harness.py --write on the merged tree. Every entry equals one side's hash apart from the summary counters.
…heck The attention_unified emitter pair built only the scalar 2D kernel and the reduce, so nothing compared the Python and C++ scalar 3D segment kernel. Add configs 8-9: the split-3D segment kernel (8 segments) of two default-mask problems, built by build_unified_attention_3d and rocke_build_unified_attention_3d_scalar_new. The two engines did not produce the same IR for it. Reconcile the C++ builder with Python: - search the sequence index with rocke_binary_search_seq_idx (per token, 32 iterations), as _emit_find_seq_idx_scan does, instead of a linear scan; - emit nested builder calls in Python's left-to-right order (QK extracts, the segment end, and the accumulator update), because C leaves argument evaluation order unspecified. Record the new configs in the differential golden.
The segment kernels now store the raw running max, so an empty segment already has max = -inf and exp2(max - overall) is 0 whenever the overall max is finite. The per-segment "sum == 0" select only mattered when every segment is empty (overall = -inf), where exp2(-inf - -inf) is NaN. Clamp the exponent reference once instead, the same way the segment kernels do: ref = select(overall > -inf, overall, 0), and factor = exp2(max - ref). Results are unchanged (an all-empty row gets factor 0, den 0 and O = 0; a NaN sum still gives a NaN den), with one fcmp and one select less per segment, and no reliance on an exact-zero sum. Both engines emit the same ops. The IR test now pins one reference outside the loops, read by the combine pass, and no compare or select inside the loops. Re-bless the scalar reduce representative entries and the differential reduce configs (6-7).
The default mask is bottom-right causal with no right bound, and it can have a sliding window. The scalar 3D segment kernel applied only the causal bound under that mask, so a bottom-right windowed problem attended to every key before the diagonal, outside its window. Both engines had the same gap. Apply the window under every mask: Python drops the window=False special case (and the now-unused parameter of _scalar_band_ok), and the C++ segment loop emits the same left-bound check after the causal one when sliding_window > 0. The kernel name now carries the "sw" flag for any windowed problem, so the fixed kernel does not share a name with the old one. Tests: the scalar 3D numeric test gains the bottom-right windowed case (rows past the first window were wrong before this change), and the IR test checks that a window compare masks the score for every windowed scalar 2D and 3D problem. The byte-identity emitters gain a windowed segment config (10).
AviralGoelAMD
added a commit
that referenced
this pull request
Oct 29, 2026
…nified-2d-lse Bring in #14542's review fixes: scalar 3D segment configs in the byte-identity gate (with the C++ 3D seq-index search and op order matched to Python), one exponent reference in the scalar reduce combine pass, and the sliding window applied under every mask in the scalar 3D body. The merge is clean.
Contributor
|
Automated review: Model - qwen3.6:35b-a3b-coding-mtp-q4_K_M, effort: max PR #14542 — carry the raw running max in unified attention (AviralGoelAMD, AICK-2829) — 🟠 First-round: not merge-ready — rebase needed — both goldens conflict with #13889's scalar re-bless; the numerics fix checks out
Utility: 🟢 Positive — fixes masked-first-tile output collapse and hidden NaN rows; GPU numerics author-run. Legend: 🔴 blocking · 🟠 should-fix · 🟡 minor nit · 🟢 info/good · ⚪ retracted First-round review: a local qwen3.6 model made the independent first pass, and the orchestrator (Claude Opus 5.5) verified it. The PR conflicts with develop, so there is no merged tree yet. The build, byte identity, full golden checks and test runs wait for the rebase. The orchestrator simulated the merge with |
Contributor
|
Automated review: Model - qwen3.6:35b-a3b-coding-mtp-q4_K_M, effort: max PR #14542 — carry the raw running max in unified attention (AviralGoelAMD, AICK-2829) — 🟡 Follow-up: no closer look needed — every earlier item is fixed at head and the numerics now run in CI; arai713's approval predates the codegen changes, and CI at head is still building
Net-new: none. Utility: 🟢 Positive — fixes masked-tile collapse, hidden NaN rows and the 3D window; numerics CI-run. Legend: 🔴 blocking · 🟠 should-fix · 🟡 minor nit · 🟢 info/good · ⚪ retracted Follow-up review: a local qwen3.6 model made the independent first pass. The orchestrator (Claude Opus 5.5) verified it with an interdiff of the five new commits against |
1 task done
AnaghaRaoAMD
approved these changes
Oct 29, 2026
8 of 9 tasks
7 of 8 tasks
…arry Resolve the representative-IR golden conflict with the per-family shard split: drop the single-file golden develop deleted, then re-bless only the attention and attention_d256 shards. The re-bless changes the same 84 flavor/case hashes this branch changed in the old file; other shards keep develop's content.
…arry Resolve the golden conflicts with develop's attention changes by regeneration: take develop's files, then re-record this branch's six attention families in the differential golden (ll llvm20/22/23 + ir canonical, python-only) and re-bless the attention and attention_d256 representative-IR shards. The re-bless touches the same 84 flavor/case entries as before the merge; other families keep develop's entries.
AviralGoelAMD
deleted the
users/avirgoel/rocke/tiled-2d-running-max-carry
branch
October 31, 2026 01:21
19 of 20 tasks
10 tasks done
11 of 13 tasks
Contributor
TheRock Submodule Bump ActivityNewest first
|
cgmillette
added a commit
that referenced
this pull request
Nov 4, 2026
…r engines (#14242) ## Summary Packages rocKE's gfx950 unified attention as two hipDNN kernel-ingestor engines: `hipkernel:Gfx950AttentionUnified2D` (prefill and the decode shapes 3D does not cover) and `hipkernel:Gfx950AttentionUnified3D` (split-KV decode, a segment and a reduce launch per kernel through #14158's multi-launch `sources`). Both ship AOT, so an SDPA-forward graph on gfx950 can be served by unified attention with no JIT. Both catalogs are reduced to what a correctness-checked MI355X sweep measured as needed: Unified2D ships 222 of its 792-kernel full space and Unified3D 123 of 348, with no graph losing coverage or more than the measured noise (the final reduced-vs-full sweep is below). JIRA ID : AICK-2542 JIRA ID : AICK-2682 ## Risk Assessment Risk 2 of 5. The PR adds two engines and their descriptors. No existing engine's code, catalog or ranking changes: an unpinned SdpaFwd graph still ranks `hipkernel:Gfx950AttentionDense` first (device test `DenseStillRanksAheadWhenUnpinned`). The shared pieces it touches: - generator tooling: `registry_enumerator` full-space and multi-launch sections, the regeneration gates, provenance stamping, and two fixes in `variant_reduce`. Only `provenance.generated_from` is opt-in. The expanded registry output now always writes per-kernel `provenance.rocke_spec_id` / `rocke_tuning_id`, including Dense's 840 catalog kernels (its manifest pins a `spec_id`). The compact `variants:` selections drop it, so no shipped descriptor changes, but any other consumer of the expanded output sees the new field; - one rocKE builder, `build_unified_attention_3d_tiled_module`; - the ingestor save census's list of dispatch symbols. These are covered by their own tests and by the existing dense regeneration gates, which still pass. 3D is the first production user of `launchNames()` (#14158). ## ASIC Coverage gfx950 only: both engines ship only for gfx950, and other architectures' shards carry neither engine. PR CI has no gfx950 device lane (`hkp-probe-gfx950` packs and asserts but launches nothing), so the device evidence comes from MI355X (and, before the rebase, MI350X) runs. ## Testing Summary - Host: unit and census tests for both packs, the arch-neutral geometry, and the generator suite under `hipdnn_strict_skips`, including regeneration gates g1–g7 for all three gfx950 manifests; static descriptor gates on the installed shards (variant sets, coverage rungs 1–2, validator, launch surface). - Device, MI355X (gfx950), on the shipped (reduced) catalogs at this head's code: - engine-pinned SdpaFwd quick and standard suites against the CPU and GPU references, both engines plus a Dense regression run; - the knobs suite: every shipped kernel class forced and checked against the CPU reference, plus declines, precedence and a NaN-prefilled workspace; - a timing comparison of every engine that serves each corpus graph, and PyTorch; - RUNBOOK stage 7's final sweep: the full catalog against the reduced one in one session. - How the catalogs were reduced: a 3-round dnn-benchmark sweep with correctness on, over 188 prefill graphs (2D) and 88 decode graphs (3D), then `variant_reduce`. ## Packaging and performance statistics All numbers below come from tool output (packaging logs, device suite runs, dnn-benchmark sweeps). **Packaging** (gfx950 shard, shipped catalogs; pack time is the build's own `hkp_pack` command per engine, median of 3, relative to Dense's on the same job) | Engine | Kernels (full space) | AOT budget | Code objects, zstd | KDP | Pack time vs Dense, build default workers | Serial time per kernel vs Dense | |---|---:|---:|---:|---:|---:|---:| | Unified2D | 222 (792) | 1000 | 1.55 MiB | 0.29 MB | 0.11× | 1.58× | | Unified3D | 123 (348), 246 launches | 400 | 0.64 MiB | 0.26 MB | 0.41× | 1.06× | | Dense (catalog unchanged, for scale) | 840 (52 distinct binaries since #14059) | 1000 | 6.16 MiB | 1.26 MB | 1 | 1 | - The gfx950 kpack is 8.49 MiB at this head; the two engines are 2.19 MiB, 26% of it. Since #14059 (runtime head counts for Dense), Dense's 840 catalog kernels compile to 52 distinct binaries, but the kpack still stores one frame per kernel, so Dense grew from 5.79 to 6.16 MiB. Collapsing those entries is #14059's stated follow-up. The pack-time ratios and the full-space figure (14.45 MiB) were measured against Dense before #14059. - Provider load (`hipdnn_list_engines`, median of 5) is 1.24× its time without these descriptors. **Correctness and coverage** (MI355X, shipped catalogs) | Suite | Unified2D | Unified3D | Dense | |---|---|---|---| | SdpaFwd quick + standard, CPU and GPU reference (served / declined / failed) | 43 / 6 / 0 | 4 / 45 / 0 | 45 / 4 / 0 | | Knobs suite (forced kernel classes, declines, precedence) | 30 / 30 (both engines) | | 62 / 62 | | Corpus graphs served and validated against PyTorch | 188 / 188 | 82 / 88 | | Every decline is intentional: Unified2D leaves decode to Unified3D and declines Hkv > Hq; Unified3D serves decode only, and its 6 corpus declines are batch-256 decode graphs whose K/V is over the 32-bit addressing cap (ASM SDPA serves those). **Kernel time** (MI355X, autotuned, median of 3 rounds; ratio = Unified / other, below 1 = Unified faster; every engine's round-to-round drift ≤ 0.4%). The Unified2D rows were re-measured at this head, on the kernels develop's tiled-2D changes recompiled (see the checklist). | | Graphs both serve | Geomean ratio | Time-weighted (other / Unified) | Unified faster on | |---|---:|---:|---:|---:| | Unified3D vs ASM SDPA (decode) | 60 | 0.200 | 8.17 | 60 / 60 | | Unified3D vs PyTorch | 82 | 0.062 | 34.2 | 82 / 82 | | Unified2D vs Dense (prefill) | 179 | 2.37 | 0.38 | 0 / 179 | | Unified2D vs PyTorch | 188 | 1.38 | 0.55 | 39 / 188 | - Unified3D serves 22 decode graphs no existing hipDNN engine serves. - Unified2D is slower than Dense wherever Dense serves, and Dense keeps ranking first there; it adds coverage Dense lacks (the 9 bottom-right causal graphs in these corpora). **Variant reduction** | Engine | Kernels | Chosen eps | Measured drift | Final sweep, reduced / full kernel time (geomean, rounds 1–3) | |---|---|---:|---:|---| | Unified2D | 792 → 222 | 0.02 | 0.62% | 1.000, 0.999, 1.000 (u2d_attn); 1.000, 0.999, 1.000 (u2d_rocke) | | Unified3D | 348 → 123 | 0.03 | 1.08% | 1.004, 1.002, 0.999 (u3d_attn); 1.002, 1.007, 1.004 (u3d_rocke) | Both arms of the final sweep served the same graphs and passed correctness on every corpus. ## Testing Checklist - [x] Unit (host, no GPU) - `HIP_VISIBLE_DEVICES=-1 hip_kernel_provider_tests` - 1454 passed, including `TestKernelIngestorEngine.EveryHandlerReportsSupportsSavingExactlyWhenItSavesAndRestores`. Our build host also fails 11 Dense signature/dispatch tests: it has a `/nonexistent` directory, which those tests assume is absent. They pass in CI - Status: Passed - [x] Census - `ctest -R 'hkp-census|rocke_unified'` - 36/36 (the census inventories are the reduced catalogs') - Status: Passed - [x] Generator suite (strict skips) - `pytest -p hipdnn_strict_skips IngestorGenerator/tests` - 1646 passed, 0 skipped - Status: Passed - [x] Regeneration gates on the reduced selections - `tests/test_registry_regeneration.py` (g1–g7) 39 passed; AOT budget tests 12 passed - Status: Passed - [x] rocKE attention CPU tests and packaging guards - 2D/3D mask_type, 3D module, decode split, KV cache policy, spec-predicate parity: 145 passed, plus 6 develop KV-cache tests that need the compiled `rocke_engine`; `test_hkp_pack_producer_guards.py` and `test_gfx950_attention_unified_3d_bundle_coverage.py`: 26 passed - Status: Passed - [x] Golden IR, before the reduction - `ctest -R rocke_golden_static` - the reduction changes no rocKE code - Status: Passed - [x] Launch surface - `launch_surface.py configs/gfx950_attention_unified_{2d,3d}.engine.yaml --check` - Status: Passed - [x] Variant sets - `verify_variant_sets.py --mode structural` and `--mode full --arch gfx950 --manifest …` - 2D 222/222 distinct binaries, 3D 123/123 - Status: Passed - [x] Coverage gate rungs 1–2 - `coverage_gate.py --mode full --expect-engine …` - Status: Passed - [x] Validator - `hipdnn_validate_descriptors <shard> --expect-engine … --json` - success, no warnings - Status: Passed - [x] Packaging probe - gfx942 configure with `HIPKERNELPROVIDER_ENABLE_PACKAGING_PROBES=ON`, `ctest -R '^hkp-probe-gfx950'` - 12/12; `hkp-probe-gfx950` packs 51 UKDs of the current descriptors in a gfx942-only build - Status: Passed - [x] Device suites, MI355X - engine-pinned SdpaFwd quick + standard with GPU reference (`hip_kernel_provider_gfx950_attention_*_gpu_ref_integration_tests`) and CPU reference (`*-external-integration_{quick,standard}_suite`) - 0 failures: Unified2D 43 served / 6 declined, Unified3D 4 / 45, Dense 45 / 4 - Status: Passed - [x] Device knobs suite, MI355X - `hip_kernel_provider_integration_tests --gtest_filter='*Gfx950AttentionUnified*'` - 30/30; Dense knobs 62/62 - Status: Passed - [x] Fault re-check, MI355X - the whole `u2d_rocke` corpus in one autotuned process, twice and then with serialized launches - 93/93 each, 0 errors - Status: Passed - [x] Reduction sweep + `variant_reduce.py` - 3 rounds, correctness on; 2D `SWEEP_DONE`, 3D `SWEEP_DONE`; 2D 792 → 222, 3D 348 → 123 - Status: Passed - [x] Final sweep, reduced against full (RUNBOOK §7) - same coverage, correctness passed on both arms, ratios above - Status: Passed - [x] Unified3D catalog regenerated for #14367's LLC-aware K/V cache policy - all 123 kernels now `kv_cache_policy: "all"` (was `"stream"`), same kernel classes; MI355X re-run: suites and knobs unchanged, 246/246 decode rows correct, kernel time against the `"stream"` catalog geomean 0.998, time-weighted 0.999 - Status: Passed - [x] Device evidence at this head (rebased onto develop 2026-11-03, after the review-round refactors) - Dense's and Unified3D's code objects are byte-identical to the previous device runs; develop's #13652 recompiled 201 of the 222 Unified2D code objects (same kernels and ids), and the Unified2D/3D matchers moved onto a shared header, so both engines' device runs were repeated at this head: suites Unified2D 43 / 6 / 0, Unified3D 4 / 45 / 0, Dense 45 / 4 / 0 (passed / skipped / failed), knobs 30/30 and 62/62; compare 188/188 Unified2D graphs (564/564 rows) and 82 Unified3D graphs (246/246 rows) correct against PyTorch in all 3 rounds. Kernel time vs the previous runs: Unified2D geomean 0.985 (none of 188 graphs >2% slower), Unified3D 1.003 (ASM_SDPA control 1.006). A final rebase picked up #14059, which recompiled Dense's 840 code objects only (Unified2D/3D byte-identical to the run above); the suites and knobs were re-run at this head with the same results, `DenseStillRanksAheadWhenUnpinned` included - Status: Passed - [x] DVC - both reduced KDPs on the `ingestor` remote; `dvc status -c` in sync - Status: Passed - [ ] PR CI - GitHub PR checks - Status: Pending ## Technical Changes - **Unified2D pack** (`Gfx950AttentionUnified2DNative.cpp`): matches BSHD SDPA forward on the identity paged-KV layout and builds block_tables, seq_lens and query_start_len in the workspace. - It declines: batch > 511 (the baked 9-iteration binary search), Skv not a multiple of 16, i32 KV over 0x7FFF0000, any window, and the decode cells Unified3D covers (Sq = 1 with the GQA group dividing 16). - Engine knobs: tile_size, num_warps, block_m_per_warp, use_mfma_32x32 and use_transposed_qk_32x32. - **Unified3D pack** (`Gfx950AttentionUnified3DNative.cpp`): decode only. - `launchNames()` is `{segment, reduce}`, and both definitions are resolved through `launchKernelDefinition` with a signature check before any load. - The segment launch writes `[token, Hq, seg, D]` scratch and the reduce always runs, so the workspace needs no zeroing. - Knob: num_segments. No batch limit. - **Arch-neutral geometry** (`UnifiedAttentionGeometry.hpp`): the launch geometry and workspace layouts take a `UnifiedAttentionArchTraits`, and the gfx950 headers are thin wrappers, so a gfx942 pack (AICK-2544) only supplies traits. - **rocKE**: `build_unified_attention_3d_tiled_module` packs the segment and its derived reduce (`reduce_spec_for`) as one `KernelModule`. The module predicate gates both halves. The 3D pack compiles develop's running-max fix for split-KV (#14276). - **Generator**: - `enumerate.also` sections may use `space: full` and carry manifest `launches`. The enumerator refuses a registry that pairs a reduce when the manifest declares no launches, and keeps one default per cell. - The gates are launch-aware: g1 keys on launches, g3 uses `--subset` for full spaces, g4 compares launch sets, and g6 walks every launch and checks module roles. - g2 and g3 run per engine with the shipped ids reused (#14109), comparing every descriptor byte for byte. - `generate.py --stamp-provenance` stamps `provenance.generated_from` (config, manifest, generator and rocKE commits). It is opt-in, so a default regeneration stays byte-identical. - `variant_reduce` fixes: - it reads the format-2 winner lines the runtime has written since #13028; - it keys a graph on its content with `Graph.id` (minted fresh at each finalization) zeroed, so sweep rounds match. - **Catalogs** (the manifests enumerate the full spaces; the selections are `variant_reduce`'s): - Unified2D: 222 kernels over 198 cells (33 head configs × bf16/fp16 × 3 masks): the narrow nw4/mw16/t64 geometry in every cell, transposed 32x32 nw4/mw32/t64 in 21 and narrow nw2/mw16/t32 in 3. The wide 32x32 geometry won no graph and is not shipped. - Unified3D: 123 kernels: num_segments 16 in all 116 decode cells, 32 or 64 in 7 bf16 bottom-right D128 variants. The K/V cache policy follows #14367's LLC-aware rule, which picks `"all"` for every shipped decode cell. - **Descriptors**: both are generated with develop's generator (post-#13028 UED/UHD, UKD `sources` for 3D), every kernel id reused from the full catalog. A regeneration with the shipped ids reused reproduces them byte for byte. The KDPs live in DVC (`ingestor` remote). - **Shared code** (one home per fact): - `packs/UnifiedAttentionGraphMatch.hpp`: the graph-match helpers both unified packs use (operand/BSHD checks, SDPA node, mask classification, dtype, metadata readers); the tests' graph builder is `UnifiedAttentionTestGraphs.hpp`. - rocKE `kernels/common/attention_mask_type.py` (mask ordinals and band -> kind, used by both tiled specs and re-exported by dispatch) and `attention_split_kv_reduce.py` (the segment -> reduce derivation the dispatcher and the module export share). - CMake registers Dense, Unified2D and Unified3D by one loop per lane (external integration, gpu_ref, census). - **Tests**: - The device knobs suite (`IntegrationGpuGfx950AttentionUnifiedKnobs.cpp`) adds a harness `prepareWorkspace` hook. It forces the default geometry/segment count across the head configurations and each other shipped class on a cell that carries it, and `CoversEveryCatalogKernelClass` holds those tables to the catalog. - The engine-pinned `gpu_ref` targets and external-integration lanes cover both engines. - The ingestor save census names all seven pack dispatch symbols. - The Unified2D/3D integration TOMLs are reserved with the gfx950 census shard (`hkp_reserve_census_shard(... FILES ...)`, as #14573 requires of every shard entry). **Constraints and follow-ups** - #13372 and #13889 are both in develop; this branch no longer carries a copy of #13889. - Heads are catalog cells (baked), window = 0, 2D batch ≤ 511, i32 KV ≤ 0x7FFF0000. D256 is deferred for both engines. The kernarg work is tracked in AICK-2675. - There is no 2D/3D `select_path` crossover replay: the split is the static decode rule above. - Runtimes older than UKD 1.1 skip the 3D descriptors. Strided-KV 3D is a follow-up. - Follow-up: execution-plan save/restore (#13315) for both handlers. Today they report `supportsSaving() == false`, as the RDNA handler does. - One earlier MI355X sweep hit a single GPU fault in Unified2D on `u2d_rocke/rocke_sdpa_035` (bf16, D 128, Hq 64 / Hkv 8, causal, S 512). It has not reproduced since: not in the forced geometries on that cell, 15 single-graph runs, 4 full-corpus runs on MI350X, or the MI355X full-corpus re-check above. - Unified2D kernel time moved with develop's tiled-2D changes (#14542, #14616), not with this PR: at this head it is 1.035× (geomean) its pre-rebase time on the same 188 graphs, up to 1.36× on long causal bf16 D128 prefill (S 8192), while Dense in the same runs is flat (0.997). Dense still ranks first wherever it serves; the table above is at this head. - Measurement note: the pinned dnn-benchmarking (v1 result schema, which `sweep.py` reads) has no bottom-right causal PyTorch reference. The sweeps used upstream dnn-benchmarking #84's reference, with the 9 bottom-right corpus graphs written as `causal_mask: false, right_bound: 0, diagonal_alignment: BOTTOM_RIGHT` (the same band under hipDNN's `resolveDiagonalBand`): #84 reads `causal_mask: true` as top-left regardless of the alignment. 🤖 Generated with [Claude Code](https://claude.com/claude-code) --------- Co-authored-by: Claude <noreply@anthropic.com> Co-authored-by: Christopher Millette <63608002+cgmillette@users.noreply.github.com>
dlangbe
added a commit
that referenced
this pull request
Nov 5, 2026
… live set Carrying the raw running max (#14542) kept it live through the whole transposed softmax body alongside the clamped exponent reference. At fp16 D128 with i64 K/V addressing that one extra value pushed the kernel to the VGPR ceiling, and the compiler then serialized the per-tile compares through VCC. Rebuild the carry at the yield as new_l == 0 ? -inf : m_new instead. The raw max is -inf exactly while the row has seen no key, which is when l is 0, so the carried value is unchanged. Testing l rather than the max compare keeps LLVM from folding the raw max back in. The IR test accepts the rebuilt carry and still rejects a clamped one. Co-Authored-By: Claude <noreply@anthropic.com>
1 task
cycheng
pushed a commit
that referenced
this pull request
Nov 6, 2026
…r engines (#14242) ## Summary Packages rocKE's gfx950 unified attention as two hipDNN kernel-ingestor engines: `hipkernel:Gfx950AttentionUnified2D` (prefill and the decode shapes 3D does not cover) and `hipkernel:Gfx950AttentionUnified3D` (split-KV decode, a segment and a reduce launch per kernel through #14158's multi-launch `sources`). Both ship AOT, so an SDPA-forward graph on gfx950 can be served by unified attention with no JIT. Both catalogs are reduced to what a correctness-checked MI355X sweep measured as needed: Unified2D ships 222 of its 792-kernel full space and Unified3D 123 of 348, with no graph losing coverage or more than the measured noise (the final reduced-vs-full sweep is below). JIRA ID : AICK-2542 JIRA ID : AICK-2682 ## Risk Assessment Risk 2 of 5. The PR adds two engines and their descriptors. No existing engine's code, catalog or ranking changes: an unpinned SdpaFwd graph still ranks `hipkernel:Gfx950AttentionDense` first (device test `DenseStillRanksAheadWhenUnpinned`). The shared pieces it touches: - generator tooling: `registry_enumerator` full-space and multi-launch sections, the regeneration gates, provenance stamping, and two fixes in `variant_reduce`. Only `provenance.generated_from` is opt-in. The expanded registry output now always writes per-kernel `provenance.rocke_spec_id` / `rocke_tuning_id`, including Dense's 840 catalog kernels (its manifest pins a `spec_id`). The compact `variants:` selections drop it, so no shipped descriptor changes, but any other consumer of the expanded output sees the new field; - one rocKE builder, `build_unified_attention_3d_tiled_module`; - the ingestor save census's list of dispatch symbols. These are covered by their own tests and by the existing dense regeneration gates, which still pass. 3D is the first production user of `launchNames()` (#14158). ## ASIC Coverage gfx950 only: both engines ship only for gfx950, and other architectures' shards carry neither engine. PR CI has no gfx950 device lane (`hkp-probe-gfx950` packs and asserts but launches nothing), so the device evidence comes from MI355X (and, before the rebase, MI350X) runs. ## Testing Summary - Host: unit and census tests for both packs, the arch-neutral geometry, and the generator suite under `hipdnn_strict_skips`, including regeneration gates g1–g7 for all three gfx950 manifests; static descriptor gates on the installed shards (variant sets, coverage rungs 1–2, validator, launch surface). - Device, MI355X (gfx950), on the shipped (reduced) catalogs at this head's code: - engine-pinned SdpaFwd quick and standard suites against the CPU and GPU references, both engines plus a Dense regression run; - the knobs suite: every shipped kernel class forced and checked against the CPU reference, plus declines, precedence and a NaN-prefilled workspace; - a timing comparison of every engine that serves each corpus graph, and PyTorch; - RUNBOOK stage 7's final sweep: the full catalog against the reduced one in one session. - How the catalogs were reduced: a 3-round dnn-benchmark sweep with correctness on, over 188 prefill graphs (2D) and 88 decode graphs (3D), then `variant_reduce`. ## Packaging and performance statistics All numbers below come from tool output (packaging logs, device suite runs, dnn-benchmark sweeps). **Packaging** (gfx950 shard, shipped catalogs; pack time is the build's own `hkp_pack` command per engine, median of 3, relative to Dense's on the same job) | Engine | Kernels (full space) | AOT budget | Code objects, zstd | KDP | Pack time vs Dense, build default workers | Serial time per kernel vs Dense | |---|---:|---:|---:|---:|---:|---:| | Unified2D | 222 (792) | 1000 | 1.55 MiB | 0.29 MB | 0.11× | 1.58× | | Unified3D | 123 (348), 246 launches | 400 | 0.64 MiB | 0.26 MB | 0.41× | 1.06× | | Dense (catalog unchanged, for scale) | 840 (52 distinct binaries since #14059) | 1000 | 6.16 MiB | 1.26 MB | 1 | 1 | - The gfx950 kpack is 8.49 MiB at this head; the two engines are 2.19 MiB, 26% of it. Since #14059 (runtime head counts for Dense), Dense's 840 catalog kernels compile to 52 distinct binaries, but the kpack still stores one frame per kernel, so Dense grew from 5.79 to 6.16 MiB. Collapsing those entries is #14059's stated follow-up. The pack-time ratios and the full-space figure (14.45 MiB) were measured against Dense before #14059. - Provider load (`hipdnn_list_engines`, median of 5) is 1.24× its time without these descriptors. **Correctness and coverage** (MI355X, shipped catalogs) | Suite | Unified2D | Unified3D | Dense | |---|---|---|---| | SdpaFwd quick + standard, CPU and GPU reference (served / declined / failed) | 43 / 6 / 0 | 4 / 45 / 0 | 45 / 4 / 0 | | Knobs suite (forced kernel classes, declines, precedence) | 30 / 30 (both engines) | | 62 / 62 | | Corpus graphs served and validated against PyTorch | 188 / 188 | 82 / 88 | | Every decline is intentional: Unified2D leaves decode to Unified3D and declines Hkv > Hq; Unified3D serves decode only, and its 6 corpus declines are batch-256 decode graphs whose K/V is over the 32-bit addressing cap (ASM SDPA serves those). **Kernel time** (MI355X, autotuned, median of 3 rounds; ratio = Unified / other, below 1 = Unified faster; every engine's round-to-round drift ≤ 0.4%). The Unified2D rows were re-measured at this head, on the kernels develop's tiled-2D changes recompiled (see the checklist). | | Graphs both serve | Geomean ratio | Time-weighted (other / Unified) | Unified faster on | |---|---:|---:|---:|---:| | Unified3D vs ASM SDPA (decode) | 60 | 0.200 | 8.17 | 60 / 60 | | Unified3D vs PyTorch | 82 | 0.062 | 34.2 | 82 / 82 | | Unified2D vs Dense (prefill) | 179 | 2.37 | 0.38 | 0 / 179 | | Unified2D vs PyTorch | 188 | 1.38 | 0.55 | 39 / 188 | - Unified3D serves 22 decode graphs no existing hipDNN engine serves. - Unified2D is slower than Dense wherever Dense serves, and Dense keeps ranking first there; it adds coverage Dense lacks (the 9 bottom-right causal graphs in these corpora). **Variant reduction** | Engine | Kernels | Chosen eps | Measured drift | Final sweep, reduced / full kernel time (geomean, rounds 1–3) | |---|---|---:|---:|---| | Unified2D | 792 → 222 | 0.02 | 0.62% | 1.000, 0.999, 1.000 (u2d_attn); 1.000, 0.999, 1.000 (u2d_rocke) | | Unified3D | 348 → 123 | 0.03 | 1.08% | 1.004, 1.002, 0.999 (u3d_attn); 1.002, 1.007, 1.004 (u3d_rocke) | Both arms of the final sweep served the same graphs and passed correctness on every corpus. ## Testing Checklist - [x] Unit (host, no GPU) - `HIP_VISIBLE_DEVICES=-1 hip_kernel_provider_tests` - 1454 passed, including `TestKernelIngestorEngine.EveryHandlerReportsSupportsSavingExactlyWhenItSavesAndRestores`. Our build host also fails 11 Dense signature/dispatch tests: it has a `/nonexistent` directory, which those tests assume is absent. They pass in CI - Status: Passed - [x] Census - `ctest -R 'hkp-census|rocke_unified'` - 36/36 (the census inventories are the reduced catalogs') - Status: Passed - [x] Generator suite (strict skips) - `pytest -p hipdnn_strict_skips IngestorGenerator/tests` - 1646 passed, 0 skipped - Status: Passed - [x] Regeneration gates on the reduced selections - `tests/test_registry_regeneration.py` (g1–g7) 39 passed; AOT budget tests 12 passed - Status: Passed - [x] rocKE attention CPU tests and packaging guards - 2D/3D mask_type, 3D module, decode split, KV cache policy, spec-predicate parity: 145 passed, plus 6 develop KV-cache tests that need the compiled `rocke_engine`; `test_hkp_pack_producer_guards.py` and `test_gfx950_attention_unified_3d_bundle_coverage.py`: 26 passed - Status: Passed - [x] Golden IR, before the reduction - `ctest -R rocke_golden_static` - the reduction changes no rocKE code - Status: Passed - [x] Launch surface - `launch_surface.py configs/gfx950_attention_unified_{2d,3d}.engine.yaml --check` - Status: Passed - [x] Variant sets - `verify_variant_sets.py --mode structural` and `--mode full --arch gfx950 --manifest …` - 2D 222/222 distinct binaries, 3D 123/123 - Status: Passed - [x] Coverage gate rungs 1–2 - `coverage_gate.py --mode full --expect-engine …` - Status: Passed - [x] Validator - `hipdnn_validate_descriptors <shard> --expect-engine … --json` - success, no warnings - Status: Passed - [x] Packaging probe - gfx942 configure with `HIPKERNELPROVIDER_ENABLE_PACKAGING_PROBES=ON`, `ctest -R '^hkp-probe-gfx950'` - 12/12; `hkp-probe-gfx950` packs 51 UKDs of the current descriptors in a gfx942-only build - Status: Passed - [x] Device suites, MI355X - engine-pinned SdpaFwd quick + standard with GPU reference (`hip_kernel_provider_gfx950_attention_*_gpu_ref_integration_tests`) and CPU reference (`*-external-integration_{quick,standard}_suite`) - 0 failures: Unified2D 43 served / 6 declined, Unified3D 4 / 45, Dense 45 / 4 - Status: Passed - [x] Device knobs suite, MI355X - `hip_kernel_provider_integration_tests --gtest_filter='*Gfx950AttentionUnified*'` - 30/30; Dense knobs 62/62 - Status: Passed - [x] Fault re-check, MI355X - the whole `u2d_rocke` corpus in one autotuned process, twice and then with serialized launches - 93/93 each, 0 errors - Status: Passed - [x] Reduction sweep + `variant_reduce.py` - 3 rounds, correctness on; 2D `SWEEP_DONE`, 3D `SWEEP_DONE`; 2D 792 → 222, 3D 348 → 123 - Status: Passed - [x] Final sweep, reduced against full (RUNBOOK §7) - same coverage, correctness passed on both arms, ratios above - Status: Passed - [x] Unified3D catalog regenerated for #14367's LLC-aware K/V cache policy - all 123 kernels now `kv_cache_policy: "all"` (was `"stream"`), same kernel classes; MI355X re-run: suites and knobs unchanged, 246/246 decode rows correct, kernel time against the `"stream"` catalog geomean 0.998, time-weighted 0.999 - Status: Passed - [x] Device evidence at this head (rebased onto develop 2026-11-03, after the review-round refactors) - Dense's and Unified3D's code objects are byte-identical to the previous device runs; develop's #13652 recompiled 201 of the 222 Unified2D code objects (same kernels and ids), and the Unified2D/3D matchers moved onto a shared header, so both engines' device runs were repeated at this head: suites Unified2D 43 / 6 / 0, Unified3D 4 / 45 / 0, Dense 45 / 4 / 0 (passed / skipped / failed), knobs 30/30 and 62/62; compare 188/188 Unified2D graphs (564/564 rows) and 82 Unified3D graphs (246/246 rows) correct against PyTorch in all 3 rounds. Kernel time vs the previous runs: Unified2D geomean 0.985 (none of 188 graphs >2% slower), Unified3D 1.003 (ASM_SDPA control 1.006). A final rebase picked up #14059, which recompiled Dense's 840 code objects only (Unified2D/3D byte-identical to the run above); the suites and knobs were re-run at this head with the same results, `DenseStillRanksAheadWhenUnpinned` included - Status: Passed - [x] DVC - both reduced KDPs on the `ingestor` remote; `dvc status -c` in sync - Status: Passed - [ ] PR CI - GitHub PR checks - Status: Pending ## Technical Changes - **Unified2D pack** (`Gfx950AttentionUnified2DNative.cpp`): matches BSHD SDPA forward on the identity paged-KV layout and builds block_tables, seq_lens and query_start_len in the workspace. - It declines: batch > 511 (the baked 9-iteration binary search), Skv not a multiple of 16, i32 KV over 0x7FFF0000, any window, and the decode cells Unified3D covers (Sq = 1 with the GQA group dividing 16). - Engine knobs: tile_size, num_warps, block_m_per_warp, use_mfma_32x32 and use_transposed_qk_32x32. - **Unified3D pack** (`Gfx950AttentionUnified3DNative.cpp`): decode only. - `launchNames()` is `{segment, reduce}`, and both definitions are resolved through `launchKernelDefinition` with a signature check before any load. - The segment launch writes `[token, Hq, seg, D]` scratch and the reduce always runs, so the workspace needs no zeroing. - Knob: num_segments. No batch limit. - **Arch-neutral geometry** (`UnifiedAttentionGeometry.hpp`): the launch geometry and workspace layouts take a `UnifiedAttentionArchTraits`, and the gfx950 headers are thin wrappers, so a gfx942 pack (AICK-2544) only supplies traits. - **rocKE**: `build_unified_attention_3d_tiled_module` packs the segment and its derived reduce (`reduce_spec_for`) as one `KernelModule`. The module predicate gates both halves. The 3D pack compiles develop's running-max fix for split-KV (#14276). - **Generator**: - `enumerate.also` sections may use `space: full` and carry manifest `launches`. The enumerator refuses a registry that pairs a reduce when the manifest declares no launches, and keeps one default per cell. - The gates are launch-aware: g1 keys on launches, g3 uses `--subset` for full spaces, g4 compares launch sets, and g6 walks every launch and checks module roles. - g2 and g3 run per engine with the shipped ids reused (#14109), comparing every descriptor byte for byte. - `generate.py --stamp-provenance` stamps `provenance.generated_from` (config, manifest, generator and rocKE commits). It is opt-in, so a default regeneration stays byte-identical. - `variant_reduce` fixes: - it reads the format-2 winner lines the runtime has written since #13028; - it keys a graph on its content with `Graph.id` (minted fresh at each finalization) zeroed, so sweep rounds match. - **Catalogs** (the manifests enumerate the full spaces; the selections are `variant_reduce`'s): - Unified2D: 222 kernels over 198 cells (33 head configs × bf16/fp16 × 3 masks): the narrow nw4/mw16/t64 geometry in every cell, transposed 32x32 nw4/mw32/t64 in 21 and narrow nw2/mw16/t32 in 3. The wide 32x32 geometry won no graph and is not shipped. - Unified3D: 123 kernels: num_segments 16 in all 116 decode cells, 32 or 64 in 7 bf16 bottom-right D128 variants. The K/V cache policy follows #14367's LLC-aware rule, which picks `"all"` for every shipped decode cell. - **Descriptors**: both are generated with develop's generator (post-#13028 UED/UHD, UKD `sources` for 3D), every kernel id reused from the full catalog. A regeneration with the shipped ids reused reproduces them byte for byte. The KDPs live in DVC (`ingestor` remote). - **Shared code** (one home per fact): - `packs/UnifiedAttentionGraphMatch.hpp`: the graph-match helpers both unified packs use (operand/BSHD checks, SDPA node, mask classification, dtype, metadata readers); the tests' graph builder is `UnifiedAttentionTestGraphs.hpp`. - rocKE `kernels/common/attention_mask_type.py` (mask ordinals and band -> kind, used by both tiled specs and re-exported by dispatch) and `attention_split_kv_reduce.py` (the segment -> reduce derivation the dispatcher and the module export share). - CMake registers Dense, Unified2D and Unified3D by one loop per lane (external integration, gpu_ref, census). - **Tests**: - The device knobs suite (`IntegrationGpuGfx950AttentionUnifiedKnobs.cpp`) adds a harness `prepareWorkspace` hook. It forces the default geometry/segment count across the head configurations and each other shipped class on a cell that carries it, and `CoversEveryCatalogKernelClass` holds those tables to the catalog. - The engine-pinned `gpu_ref` targets and external-integration lanes cover both engines. - The ingestor save census names all seven pack dispatch symbols. - The Unified2D/3D integration TOMLs are reserved with the gfx950 census shard (`hkp_reserve_census_shard(... FILES ...)`, as #14573 requires of every shard entry). **Constraints and follow-ups** - #13372 and #13889 are both in develop; this branch no longer carries a copy of #13889. - Heads are catalog cells (baked), window = 0, 2D batch ≤ 511, i32 KV ≤ 0x7FFF0000. D256 is deferred for both engines. The kernarg work is tracked in AICK-2675. - There is no 2D/3D `select_path` crossover replay: the split is the static decode rule above. - Runtimes older than UKD 1.1 skip the 3D descriptors. Strided-KV 3D is a follow-up. - Follow-up: execution-plan save/restore (#13315) for both handlers. Today they report `supportsSaving() == false`, as the RDNA handler does. - One earlier MI355X sweep hit a single GPU fault in Unified2D on `u2d_rocke/rocke_sdpa_035` (bf16, D 128, Hq 64 / Hkv 8, causal, S 512). It has not reproduced since: not in the forced geometries on that cell, 15 single-graph runs, 4 full-corpus runs on MI350X, or the MI355X full-corpus re-check above. - Unified2D kernel time moved with develop's tiled-2D changes (#14542, #14616), not with this PR: at this head it is 1.035× (geomean) its pre-rebase time on the same 188 graphs, up to 1.36× on long causal bf16 D128 prefill (S 8192), while Dense in the same runs is flat (0.997). Dense still ranks first wherever it serves; the table above is at this head. - Measurement note: the pinned dnn-benchmarking (v1 result schema, which `sweep.py` reads) has no bottom-right causal PyTorch reference. The sweeps used upstream dnn-benchmarking #84's reference, with the 9 bottom-right corpus graphs written as `causal_mask: false, right_bound: 0, diagonal_alignment: BOTTOM_RIGHT` (the same band under hipDNN's `resolveDiagonalBand`): #84 reads `causal_mask: true` as top-left regardless of the alignment. 🤖 Generated with [Claude Code](https://claude.com/claude-code) --------- Co-authored-by: Claude <noreply@anthropic.com> Co-authored-by: Christopher Millette <63608002+cgmillette@users.noreply.github.com>
cgmillette
pushed a commit
that referenced
this pull request
Nov 6, 2026
JIRA ID : AICK-2606 Partially addresses #14651 (the `TestKvAddrLimit` classification failure on develop). ## Motivation #12446 added a launch-time guard that rejects a paged KV cache larger than the compiled kernel can address. Some gfx942 K/V loaders form an i32 byte offset from the cache base even when the spec is built with `use_i64_kv_addr=True`. The guard caps those specs through `_I32_ONLY_KV_LOADERS`. The conflict-free V store (cfvst) loader is the capped loader that matters in production. It is the default V loader for gfx942 bf16 D64 and fp16 D64/D128 prefill (no sliding window). A sweep of the gfx942 default-route specs one block past 2^31 bytes found that cfvst is the only capped loader any shipped i64 spec selects. Its V offsets wrap past 2^31 bytes. Before #12446 that read wrong data or faulted the GPU; since #12446 the guard raises `ValueError` before launch. Dispatch has no fallback route, so a deployment whose per-layer K cache is over 2^31 bytes fails every prefill on these configurations. This PR gives the cfvst loader an i64 path that costs no more than the i32 path, and removes the cap. That lifts the cap for every shipped gfx942 configuration. ## Technical Details **cfvst V loader** (`library/kernels/gfx942/attention_tiled_2d.py`, `_cfvst_load_v_regs`) - Each cfvst load call covers a fixed range of tokens: 16 at D64 and 8 at D128 with 4 warps. On every geometry the launcher and the tuning space build, that range lies inside one page, and which page is known at build time. - Under `I64_KV_ADDR` the loader loads the tile's page ids once, builds one 64-bit page pointer per page, and indexes each V load with an i32 offset within the page. The per-lane select and divide stay i32, as on the i32 path. The pointers come from `_tile_page_ptrs`, split out of #14349's `_tile_pages` with the same op order. - The in-page byte offset, the kv head's byte offset and the `PAGE_RSRC_I64` rule come from `_kv_page_off`, `_kv_head_page_b` and the shared predicate, which the per-page K/V loaders use too, so a change to the paged K/V layout reaches cfvst as well. They use the runtime token stride from #15024. - A geometry whose load calls would span pages raises `ValueError` at build time. This follows the existing rule for K/V copies that straddle pages. No such geometry can launch: the launcher accepts block sizes 16, 32 and 64, and the tuning space offers 1, 2 or 4 warps. - It covers every load form: the vector `global_load_vN` path and the `HIPDNN_GFX942_CFV_STORE_SEPOFF` / `_SCALAR_LOAD` diagnostic paths. - The i32 cfvst loader emits the same IR as before. - An earlier commit in this PR used full per-lane i64 offsets (`offset_i64`). That was correct, but at fp16 D128, which already sits at the VGPR limit, it added register spills and made the i64 kernel noticeably slower than the i32 one. The page-pointer version replaces it. **Running-max carry on the transposed softmax path** (`attention_tiled_2d.py`, the loop yield) - After merging develop, the fp16 D128 i64 cfvst kernel was 1.3x slower than its i32 form. A bisect over the merged develop commits found #14542, which carries the raw running max. On the transposed path the raw max then stayed live through the whole body beside the clamped exponent reference. This kernel already sits at the VGPR ceiling, so the one extra value made the compiler route the per-tile compares through VCC one at a time. - The carry is now rebuilt at the yield as `new_l == 0 ? -inf : m_new`. The raw max is -inf exactly while the row has seen no key, which is when l is 0 (a seen key adds P = 1 at the max), so the carried value is the same and #14542's behaviour is kept. Testing l rather than the max compare matters: LLVM folds a rebuild from the compare back into the raw max. - This changes the IR of every gfx942 tiled 2D kernel on the transposed path, i32 included. The non-transposed paths and gfx950 are unchanged. **Scalar `global_load` lowering** (`core/lower_llvm.py`, `cpp/core/lower_llvm/mem.cpp`) - `global_load_typed` always wrote its GEP index as `i32`, so an i64 element index gave invalid IR. It now uses the index operand's own type, as `global_load_vN` already does. This was found while writing the first i64 loader. The final cfvst loader no longer needs it, but the lowering was wrong for any i64 index. Every existing caller passes an i32 index, so no existing kernel changes. **C++ twin** - The page-pointer cfvst loader and the same build error are mirrored in `attention_tiled_2d_conflict_free_v.cpp`. `rocke_gfx942_attn2d_tile_page_ptrs`, `rocke_gfx942_attn2d_kv_page_off`, `rocke_gfx942_attn2d_kv_head_page_b` and `rocke_gfx942_attn2d_page_rsrc_i64` are shared with the per-page K/V loaders. - The rebuilt running-max carry is mirrored in `attention_tiled_2d_kv_body_qk_softmax.cpp` and `attention_tiled_2d_kv_body_pv.cpp`, only when that body builds for gfx942. gfx950 fastkv_regp shares the C++ body but its Python (`gfx950/attention_tiled_2d.py`) still carries the raw max, so it keeps that carry. - No emit config covered cfvst before, and the twin did not match Python. It is now byte-identical: - It computes the per-tile values once and reuses the cached `max_seq_prefix_len`. - It places the split load and store where the transposed softmax issues them. - It fixes C argument evaluation order (unspecified, right-to-left in practice) in `v_t_slot`, the transposed PV token base, and the cfvst block guard and prezero loop. **Lift the cap** - `use_conflict_free_v_store` is removed from `_I32_ONLY_KV_LOADERS` and classified as i64-safe in `test_every_tiled_flag_is_classified_for_kv_addressing`. - `test_gfx942_cfvst_meta_has_no_limit`, `test_limit_memo_is_keyed_per_kernel` and `TestGfx942RouteGuard.test_cfvst_route_launches_past_2gib` now expect no limit. The route test launches at 65535, 65536, 65537 and 69632 blocks and checks the spec is i64 cfvst. - `use_additive_bias` (#13120) is also classified (`_NO_KV_ADDRESS_FLAGS`). The classification test already failed on develop without it. **Byte-identity coverage** - `library/tests/parity/gfx942_attention_tiled_2d_emit.{py,c}`: cfg18 is i32 cfvst and cfg19 is i64 cfvst, both on the bf16 D64 GQA 32/8 prefill geometry. The shipped spec also sets the sliced-K ring, direct-Q-global and the mask-limit stack. The C twin does not port those yet, so the configs leave them off. They still cover the new V loader code. - `llvm_gfx_all.json`: cfg18/19 blessed at llvm20, llvm22, llvm23, llvm24 and in canonical IR. The running-max change re-blesses the 18 transposed-path configs of `gfx942_attention_tiled_2d` and `gfx942_attention_tiled_2d_stress`; no other family moves. - Representative-IR golden: only `attention/gfx942/2d_lse_transposed32_bf16_d128` moves, from the running-max change. **Tests and docs** - `platform/tests/core/test_global_load_index_width.{py,cpp}` check the GEP index width; the C++ one is a host CTest. - `library/tests/test_gfx942_cfvst_i64_kv_addr.py` (CPU): - every V load of the shipped i64 cfvst spec indexes a page pointer with an i32 offset, including both diagnostic forms; - the i32 spec keeps i32 offsets from the cache base; - a page-spanning geometry is refused; - every cfvst geometry in the gfx942 tuning space builds with i64. - `library/tests/test_gfx942_cfvst_i64_kv_addr_numeric.py` (gfx942 GPU lane): a CPU routing anchor per shipped shape, plus an on-GPU check for bf16 D64, fp16 D64 and fp16 D128. Each runs `run_unified_attention_torch` on a 2.125 GiB GQA 32/8 causal cache with every block-table entry above 2^31 bytes, against an fp32 reference. It gates on `get_device_arch() == "gfx942"`. - `library/tests/test_attention_unified_masked_tile_ir.py` accepts the rebuilt carry in its exact form and still rejects a clamped carry. - CI lanes: `test_gfx942_cfvst_i64_kv_addr.py` runs in `rocke_library_attention_cpu_pytest`, and the numeric file in the library GPU lane. `rocke_global_load_index_width` is registered in `tests/cmake/CoreTests.cmake` and listed in the quick and ffm-quick tiers of `ROCKE_ENGINE_test_categories_external.yaml`. - `builders/gfx950/attention/README.md`: describes the cfvst page-pointer addressing; the paged-cache K/V limit list no longer names cfvst. **Out of scope (follow-ups)** - The other `_I32_ONLY_KV_LOADERS` entries: cfv, vhbm and khbm on gfx942, and fp8-PV on gfx942 and gfx950. No shipped spec selects them under i64. - Under i64 the sliced-K ring K loader still builds a buffer resource per copy, the pattern #14349 replaced in the other K/V loaders. - `global_store_typed` still writes its GEP index as `i32` in both engines, so an i64 store index would give invalid IR. No caller passes one. - Running `test_selector_parity.py` in a CI lane (the rest of #14651). - Replacing the hardcoded `0x7FFF0000` K/V `num_records` with one named limit (changes IR and goldens). - Exposing the cache-size limit to `supports_native_unified_attention*`. - Porting the sliced-K ring, direct-Q-global and the mask-limit stack to the gfx942 C twin, so an emit config can match a full shipped spec. ## Test Plan CPU: - Byte-identity gate at every flavor: `ROCKE_LLVM_FLAVOR=llvm20|llvm22|llvm23|llvm24 python tools/check_byte_identity.py` - Differential golden, as CI runs it: `run_diff.py --mode ll --python-only --check-golden --flavor llvm20,llvm22,llvm23,llvm24` `run_diff.py --mode ir --canonical --python-only --check-golden` - Representative-IR golden: `rocke_ir_parity_harness.py --check tests/golden/rocke_representative_ir_sha256 --flavor <flavor>` at each flavor. - `test_selector_parity.py`, `test_gfx942_cfvst_i64_kv_addr.py`, `test_attention_unified_masked_tile_ir.py`, `test_attention_gfx942_kv_pages.py`, `test_global_load_index_width.py` and the attention spec/selector suites, with the GPU hidden. - `rocke_global_load_index_width` CTest. gfx942 GPU (torch 2.13.0+rocm10.0.0 wheels, Python lowering): - `test_gfx942_cfvst_i64_kv_addr_numeric.py`. - #14542's `test_attention_unified_2d_masked_tile_numeric.py` and `test_attention_unified_scalar_masked_tile_numeric.py` (fully masked first tile, NaN rows), plus `test_attention_tiled_runtime_shape_numeric.py` and `test_gfx942_4wgqa_bounds_numeric.py`. - `builders/gfx942/attention/prefill/parity_unified_attention.py --scenario default`. - Kernel resources (VGPR and SGPR counts, spills) from the code object notes of the shipped cfvst kernels, i32 and i64. - Bisect of the fp16 D128 i64 slowdown: each develop commit merged into the pre-merge head, compiled, and its SGPR count read; the boundary then timed. - Performance on the shipped prefill shapes, i32 vs i64 on this branch. The workload is 2 sequences of 4096 tokens, causal, GQA 32/8, 32-token pages. - Both arms run the same small cache. The i64 arm only claims a larger block count, which selects the i64 spec. - Each arm is first checked against an fp32 reference (declared gate max_abs ≤ 4e-2). - Time is the median of 50 launches after 10 warm-up launches, measured with HIP events, over two repeats. ## Test Result - Byte-identity (`check_byte_identity.py`, every family): GREEN at llvm20, llvm22, llvm23 and llvm24, including cfg18/19. - Differential golden: GOLDEN OK at every flavor and in canonical IR. Representative-IR golden: OK at all four flavors. - CPU tests pass. The new CPU tests fail on develop and pass here. - gfx942 numeric, at this head: all three shipped shapes pass on caches past 2^31 bytes (bf16 gate 4e-2, fp16 gate 2e-2). #14542's masked-tile tests pass, including the NaN rows on the transposed path. - gfx942 parity (default group): 22 PASS, 6 SKIP, 0 FAIL. The skips are D256 shapes over the LDS budget. - Kernel resources: the fp16 D128 i64 kernel uses 252 VGPRs and 101 SGPRs with no spills; before the running-max change it used 254 and 49. - Performance, at these measured configurations, as time ratios. Every arm passed the numeric gate (max_abs 7.9e-3 bf16, 1.1e-3 fp16). | shape | branch i64 / branch i32 | fp16 D128 i64, running-max change / without it | |---|---|---| | bf16 D64 | 1.01–1.02 | | | fp16 D64 | 0.99–1.00 | | | fp16 D128 | 0.95 | 0.72 | The running-max change leaves the i32 kernels' times unchanged within run-to-run noise. ## Submission Checklist - [ ] Look over the contributing guidelines at https://github.com/ROCm/TheRock/blob/main/GOVERNANCE.md#pull-requests. --------- Co-authored-by: Claude <noreply@anthropic.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
ISSUE ID : AICK-2829
Also fixes the empty-segment and keyless-row cases of AICK-2686.
What was wrong
Online softmax keeps a running max
mper query row. When a row's first KV tile was fully masked (sliding window, or a leading split-KV segment), the kernels replacedm = -infwith 0 and carried that 0 to later tiles. Real scores were then exponentiated against 0 instead of the row's max. If every visible score is very negative, P rounds to 0 and the row's output is wrong (0 or NaN) even though the row sees keys.The output guard had a second problem. Rows with no key were detected with
l > 0, which also turned a NaN row into 0 and hid it.This blocks LSE output on the 2D path (AICK-2685).
What changed
build_unified_attention_2d_tiledl > 0guard hid NaNl == 0build_unified_attention_2dl == 0guard under every maskbuild_unified_attention_3dbuild_unified_attention_reduce-inf) gets weight 0; zero total writes O = 0test_attention_unified_2d_masked_tile_numeric.pyandtest_attention_unified_scalar_masked_tile_numeric.py, both in the CTest GPU list.Why it works
The loop carries the raw max, which stays
-infuntil the row sees a key. The-inf → 0substitute is used only as the exponent reference for P and the rescale factor. At the first tile with real keys the reference becomes that tile's max, so the largest P is 1 and cannot underflow. The tiled 3D kernels already use this update (#14276). Thel == 0guard marks only truly keyless rows, so NaN still reaches the output.How we validated
Develop
56ad847e204vs this head, one GPU allocation per arch, Python backend, LLVM 23.l > 0tiled guard)After review: the scalar numeric and IR tests pass on gfx942 and gfx950. The new bottom-right windowed scalar 3D case fails before the window fix.
Performance
Cold memory, D128, GQA 32/8, KV page 16, Sq = Skv = 8192, 5 alternating rounds. Time relative to develop (lower is faster):
No scratch on any timed kernel. The gfx942 sliding-window shape was timed with the 4-warp route forced off, because production sends it to the 4-warp kernel, which this PR does not change.
Notes
> 0guard to the C++ scalar kernels. Whichever PR merges second must adopt the raw-max carry and thel == 0guard.-1e30sentinel, no clamp), and the tiled 3D reduce's> 0guard. A shared helper for these rules is AICK-2831.