Skip to content

fix(rocke): centralize dtype names and shared MXFP contracts - #12287

Merged
tenpercent merged 17 commits into
developfrom
users/mpodkory/rocke/mxfp-common
Sep 23, 2026
Merged

tenpercent merged 17 commits into
developfrom
users/mpodkory/rocke/mxfp-common

Conversation

@tenpercent

@tenpercent tenpercent commented Sep 18, 2026 •

Copy link
Copy Markdown
Contributor

Native FP8/BF8 lookup must use canonical dtype keys, and scaled WMMA selection must distinguish matrix dtypes from scale types and their shared K-group size. This PR centralizes dtype normalization and adds a compact scaled-WMMA contract shared by the Python and C++ engines.

  • Keep the existing A/B/C MMA fields, fragments, and layouts, with C describing both accumulator input and result. Add independent a_scale_dtype / b_scale_dtype fields and one shared scale_block_k: Python MmaScaleBlockK.K16 / K32, mirrored by a C enum. All-absent metadata denotes an unscaled atom; partial contracts and invalid sizes reject.
  • Use one scale-aware matching path per language. Python queries accept scales=(a_type, b_type, block_k); C queries accept an optional rocke_mma_scale_filter_t pointer. An omitted filter is unconstrained; an all-absent filter explicitly selects unscaled atoms. Exact selection and largest-K ties reject ambiguity.
  • Select four supported homogeneous FP8/BF8 + E8M0 K16/K32 contracts from one wmma_scaled family. Builders emit generic tile.mma. IDs extend the existing underscore convention, for example wmma_gfx1250_f32_16x16x128_fp8_fp8_scale_e8m0_e8m0_k32; existing unscaled IDs retain their spelling.
  • Derive A/B selectors, packed carriers and intrinsic declarations from catalog metadata in the shared LLVM/HIP backend contract. Scale types remain distinct from integer carrier types. Declarations are shared where physical signatures agree. Generic MMA result naming uses a target-independent family lookup in both engines. Native lookup builds an immutable index once and uses binary search thereafter; Python caches its projection.
  • Expose per-source scale fragment counts and optional lane/slot layouts. The four supported gfx1250 contracts populate these maps, and the GEMM loader consumes them while applying tensor offsets and strides separately. K32 carries four E8M0 elements per lane in i32; K16 carries eight in i64. Their matrix A/B maps remain unavailable.
  • Preserve unscaled WMMA and dense MFMA lowering. Remove unrelated emitter edits and cover malformed zero/two-result WMMA with native regression tests; invalid result counts return an error instead of crashing.
  • Preserve the target-independent dtype API, shared MXFP8 example arguments/packing/verifier, and representative golden coverage. Recognition of FP4/FP6 names or other scale types does not enable them in this backend: executable scaled-WMMA catalog coverage remains the four E8M0 contracts. Independent destination metadata is outside this change.
  • Correct four existing unscaled gfx1250 FP8/BF8 C catalog entries to report eight packed i32 source elements.

Compatibility: the old scaled-WMMA IDs and dedicated builder wrappers are removed without aliases; serialized IR using them must be regenerated. Native consumers must rebuild for the added fields and query argument. C++ catalog query calls may omit the trailing filter; C callers pass a filter pointer or NULL. The ArchTarget helper requires the trailing filter argument in both languages; existing unscaled helper calls pass nullptr.

Validation

The scale-helper follow-up d6606f7d03883db907b0d5b0011bdb5b89aa2782 passes the scale filter through rocke_archtarget_op_for_shape, updates all 17 existing unscaled callers to pass nullptr, and extends the catalog-row test to cover helper selection, including the real FP8/BF8 K16/K32 contracts. A fresh native Release build, 11/11 CTests, C99 header/call compilation, and full LLVM20/22/23 parity passed: 70 families and 785 configurations per flavor, with zero drift. LLVM20/22 each have 732 identical outputs and 53 matching rejections; LLVM23 has 742 identical outputs and 43 matching rejections. The regression test fails with the expected ambiguity error when the filter is deliberately dropped. Full-PR hooks passed on a byte-identical local snapshot after an NFS stall; source hashes were verified before committing. No GPU rerun was performed for this host-side forwarding change.

The earlier host lookup follow-up is 0026b45992d0df58184dda254f515e6b93c2bf38: 11/11 native CTests, 52 Python tests plus 26 subtests, full LLVM20/22/23 parity, shared goldens and full-PR hooks passed. All 2,340 parity verdicts/output hashes match the parent; 34 captured native LLVM/HIP artifacts are byte-identical before/after. All 1,179 platform file hashes match the validated source. No GPU rerun was performed for this host-only change.

The implementation and GPU validation below apply to parent 082651964df5de0090e65135dcb2067b2ddd38ce. All 1,178 platform source-file hashes were verified against the tested archive when committing; the worktree is clean. GPU archive SHA-256: 1e6ff5ec74677a61658dbfcd7a18f363dcc5f5d3e5974dd99034cdfe72f459a9.

  • All applicable repository hooks over the full PR file set plus the new test, and whitespace/relative-path checks.
  • 104 focused Python tests and 26 subtests; zero skips.
  • Fresh native Release build and all 11 CTests, including malformed float/integer WMMA error handling.
  • Existing public metadata and records preserved across all seven catalogs, including all sixteen gfx950 records; native family/layout/default checks cover every target.
  • Full Python/C++ LLVM parity: 70 families and 780 configurations per flavor, zero drift. LLVM20/22 each have 727 identical outputs and 53 matching rejections; LLVM23 has 737 identical outputs and 43 matching rejections.
  • Shared representative golden passes at all three LLVM flavors. Only four MXFP8 LLVM23 hashes/byte counts changed for map-generated address arithmetic; existing cases and expected failures are preserved.
  • Independent scale-coordinate and packed-word tests cover both operands, every lane/slot, duplication, invalid slots, missing layouts, multiple tiles and K steps.
  • Final-source gfx1250 GPU validation on ROCm 10.1.0: 13 tests / 49 numerical cases, zero mismatches/errors/skips. Coverage includes COMGR/HIP compilation, scale-group isolation, multi-tile cases and an unscaled FP8 regression.
  • Eight stable ROCm 10.0.0 HIP/COMGR compilation cases passed with Python sources byte-identical to the final commit. This was compilation evidence only; the latest numerical run used ROCm 10.1.0.

GPU coverage is the bounded FP8/BF8 SCALE/SCALE16 suite through the Python backend against a CPU NumPy reference. Native C++ is covered by build/tests/emission parity; native C++ GPU runtime and arbitrary format combinations are outside this validation.

ISSUE ID : AICK-2056

@therock-pr-bot

therock-pr-bot Bot commented Sep 18, 2026 •

Copy link
Copy Markdown

✅ All Checks Passed — Ready for Review

Check Status Details
📝 PR Description ✅ Pass —
⛔ Forbidden Files ✅ Pass —
🧪 Unit Test ✅ Pass —
🔎 pre-commit ✅ Pass —
🚫 Draft PR 🔜 To Be Enabled —
🚩 Feature Flag 🔜 To Be Enabled —
📊 Code Coverage 🔜 To Be Enabled —
🤖 therock-pr-bot ✅ Pass —

🎉 All checks passed! This PR is ready for review.

📖 Need help? See the Policy FAQ for details on every check and how to fix failures.

🙋 Wish to Override Policy?

@tenpercent tenpercent changed the title fix(rocke): share MXFP contracts and canonical FP8/BF8 catalog lookups fix(rocke): centralize dtype names and shared MXFP contracts Sep 18, 2026
@therock-pr-bot

Copy link
Copy Markdown

🎉 All checks passed! This PR is ready for review.

@tenpercent
tenpercent marked this pull request as ready for review September 18, 2026 17:25
@tenpercent
tenpercent requested review from a team as code owners September 18, 2026 17:25
@tenpercent
tenpercent marked this pull request as draft September 18, 2026 23:30
@tenpercent
tenpercent marked this pull request as ready for review September 19, 2026 03:19
@tenpercent
tenpercent marked this pull request as draft September 21, 2026 17:15
@yraparti

yraparti commented Sep 21, 2026 •

Copy link
Copy Markdown
Contributor

Review: Software Architect Lens

Strengths:

  • The dtypes.cpp extraction is the right call. Having normalize_dtype in arch/data.cpp was an architectural layering violation — dtype normalization has no business depending on arch target data. The new rocke/dtypes.h header with a standalone rocke_normalize_dtype is
    a proper separation of concerns.
  • The shift from builder-per-variant (2 explicit wrappers) to catalog-driven dispatch is the correct architectural direction. Adding a third scaled variant should now be a catalog edit, not a code change in 4 different files.
  • Ambiguity detection in select_largest_k and op_for_shape is a correctness improvement — silent first-match behavior is a latent correctness hazard.

Concerns:

  1. Breaking change is underspecified in the PR description. The PR states "serialized IR using old scaled-WMMA IDs must be regenerated" but doesn't state what the migration path is for existing compiled artifacts in any deployment. For a platform DSL used by multiple
    consumers, a deprecation alias or a deserializer shim would be standard — the PR explicitly says there are none. This is an acceptable tradeoff only if the team has confirmed no external serialized artifacts exist.
  2. rocke_mma_scale_filter_t pointer-or-NULL API is fragile. The C query API uses NULL to mean "unconstrained" and a struct with all-zero fields to mean "select unscaled only". This two-level null semantics is easy to misuse — a caller who zero-initializes a struct on
    the stack and passes &filter accidentally requests unscaled-only rather than unconstrained. A named sentinel or a separate enum field would be safer.
  3. wmma_scale_internal.h is a new shared internal header coupling lower_hip and lower_llvm through a side channel. Both backends now call rocke_gfx1250_scaled_wmma_from_op and rocke_scaled_wmma_contract. This is better than duplication but means a future third backend
    also has to know about this header. The right long-term home is probably in the catalog/query layer as a proper public query function rather than an internal header.

Review: rocke Architect Lens

Strengths:

  • ID naming convention wmma_gfx1250_f32_16x16x128_fp8_fp8_scale_e8m0_e8m0_k32 is unambiguous and self-documenting. Encoding target, result type, shape, operand dtype, scale dtype, and block_k into the ID means you can reconstruct the full contract from the ID string
    alone.
  • The MmaScaleBlockK.K16 / K32 Python enum mirrored by a C enum is correct — the block-K value is a first-class dimension of the contract, not a property you infer from the ID string.
  • Correcting the a_frag_len/b_frag_len from 32→8 on the unscaled gfx1250 WMMA atoms is a legitimate bug fix. For a 16x16x64 WMMA in FP8, A is packed as 8 i32 words (64 bytes = 16 FP8 elements per lane in 128-lane wave? Need to double-check the wave size). The PR says
    wave_size=32, so 64 FP8 elements / 8 elements per i32 = 8 words. This is correct.
  • The wmma_scaled family name in the catalog is distinct from wmma, which correctly prevents unscaled query paths from matching scaled atoms.

Concerns:

  1. rocke_mma_result_hint in ir_tile.cpp now does a catalog lookup at IR emission time. Previously it was a static table. For gfx1250, it calls rocke_arch_target_from_gfx("gfx1250") and rocke_mma_catalog_by_op_id on every call. This is a hot-path regression for kernels
    with many MMA ops. The hint should be cached on the atom, derived once at catalog build time, or the lookup should be gated by a family string check before the catalog scan.
  2. rocke_ati_scale_dtype in query.cpp performs a silent fp8e4m3→e4m3 aliasing without any documented rationale. This conversion maps "fp8e4m3" to "e4m3" but leaves "bf8e5m2" unmodified. The asymmetry is suspicious — either both need mapping or neither. If the scale
    dtype namespace is different from the matrix dtype namespace, that should be made explicit in the type system, not handled by a one-off function.
  3. Four rocke_ati_validate_scales calls (one per query function) with identical validation logic is repetitive and means adding a new constraint requires 4 edits. A single validation call at the entry point or a wrapper that normalizes the filter before passing it to
    the inner loop would be cleaner.

Review: GPU Programming Expert Lens

Strengths:

  • The distinction between K16 and K32 block-K group sizes maps directly to the hardware: __builtin_amdgcn_wmma_scale_f32_16x16x128_f8f6f4 vs __builtin_amdgcn_wmma_scale16_f32_16x16x128_f8f6f4 differ in their scale operand packing — i32 (4 bytes = 4 E8M0 values) vs i64
    (8 bytes = 8 E8M0 values). The PR correctly maps this to matrix_words and scale_formats fields, so the intrinsic arguments are now derived from the contract rather than hand-coded per-variant.
  • Encoding matrix_formats[0/1] as immarg integer format selectors (the i32 immarg leading arguments to the f8f6f4 intrinsic) is correct — these select which of the 6 FP8/FP6/FP4 interpretations to apply. Previously these were hardcoded to 0, which worked only because
    FP8 was index 0. Now BF8 will use a different selector, making mixed-format contracts possible.
  • The LLVM declaration key rename (wmma.scale.gfx1250.f32.16x16x128.fp8.fp8 → wmma.scale.block32.gfx1250.f32.16x16x128.f8f6f4.v8f32.v16i32.v16i32) aligns with the actual LLVM intrinsic mangling and is more correct — the old key name implied it was FP8-only when the
    intrinsic is generic.

Concerns:

  1. rocke_mma_result_hint does rocke_arch_target_from_gfx("gfx1250") unconditionally on every call for any op_id. This is a string-literal gfx lookup per MMA op emitted. On a kernel with 100+ tile.mma ops, this is 100+ catalog scans. The mxacc hint should be encoded in
    the catalog atom's metadata or checked via the family name before the catalog scan.
  2. The PR only covers homogeneous FP8/BF8 contracts (A and B same type). The gfx1250 scaled-WMMA intrinsic physically supports heterogeneous mixes (fp8×bf8, bf8×fp8) via the format selector arguments. The PR description acknowledges this is not covered but doesn't say
    whether a follow-on is planned. Given the format selectors are now in the contract struct, adding heterogeneous contracts would be a catalog addition only — which is correct — but the fact that the PR's Python wmma_scale.py may not enumerate these should be called
    out.
  3. The scale_formats encoding is implicit. The matrix format integer values (index into the f8f6f4 format enum: 0=fp8, 1=bf8, 2=fp6e2m3, 3=fp6e3m2, 4=fp4) are computed somewhere in wmma_scale_internal.h and stored in the contract struct, but the diff doesn't show that
    header. The correctness of BF8 getting format selector 1 rather than 0 is critical and should be visible in the review.

Summary Verdict

Approve with follow-up items. The core refactoring is architecturally sound and the bug fix on frag lengths is legitimate. The main risks are:

image

@yraparti yraparti left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@tenpercent

Copy link
Copy Markdown
Contributor Author

Follow-up to the review, which was submitted against 856e2452. The fixes are now in 082651964df5 and 0026b45992d0.

Addressed in this PR

  • Hardcoded gfx1250 lookup and repeated family scans (rocKE architect Initial workflows #1 / GPU Initial workflows #1): generic MMA naming now resolves families across targets. Native code builds an immutable index once and uses binary search thereafter; Python already caches its projection. Concurrent first use, copied IDs, duplicate/conflicting IDs, unknown IDs and existing result hints are tested. See native query and private index. This removes repeated family scans; no measured speedup is claimed.
  • Unrelated comment/formatting and mma.cpp changes: restored the unrelated unscaled WMMA/dense MFMA implementation while retaining scaled lowering. Malformed zero/two-result WMMA cases now have native regression coverage.
  • Scale layout metadata: both sources now expose scale fragment counts and optional lane/slot maps, populated for the four supported gfx1250 contracts. The loader uses those maps and applies tensor offsets/strides separately. Layout, address and packing tests cover both operands.

Deferred to follow-up

  • Consumer migration: inventory supported consumers of retired IDs/wrappers and persisted IR, then determine whether regeneration/rebuild is sufficient or a bounded compatibility bridge is needed. The current migration is documented; this PR does not establish that external persisted IR is absent. Raw HSACO loading does not interpret MMA IDs, while native API consumers need to rebuild/update their calls.
  • C filter usability: add named unconstrained/unscaled helpers or initializers while retaining the current semantics. This is a discoverability improvement, not a change to matching behavior.
  • Heterogeneous FP8/BF8 contracts: handle as a separate feature with coordinated Python/native catalog and fragment metadata and numerical validation. Independent selectors support that direction, but adding production support is not merely a single catalog-row edit. This PR remains scoped to the four homogeneous E8M0 K16/K32 contracts.

These follow-up items remain outside this PR.

Intentional architectural choices / no further change planned here

  • Shared lowering contract stays internal. The helper derives backend intrinsic names, register carriers and selector values from catalog metadata. Keeping those backend details out of the public hardware-query API avoids exposing LLVM-specific declarations there.
  • NULL versus explicit unscaled remains intentional. An omitted/NULL filter is unconstrained; an all-absent filter selects unscaled atoms. This is documented and tested. Follow-up convenience helpers can make the distinction clearer without redefining it.
  • The scale alias is correct as written. fp8e4m3 -> e4m3 is documented. bf8e5m2 denotes E5M2, not E5M3, so mapping it to e5m3 for symmetry would change the format incorrectly.
  • Validation is already shared. Four query entry points call one validator; a new constraint belongs in that one implementation.
  • Selector values are explicit and reviewed. The shared header is included in the PR diff and spells out FP8=0, BF8=1 and E8M0=0. Both engine tests cover the selectors. No missing-header or selector fix is required.

Validation at 0026b45992d0

Fresh native build and 11/11 CTests; 52 focused Python tests plus 26 subtests; full parity for 70 families / 780 configurations per LLVM flavor. LLVM20/22 each have 727 identical outputs and 53 matching rejections; LLVM23 has 737 identical outputs and 43 matching rejections. All 2,340 verdicts/output hashes match the parent. Also passed: 34 before/after native LLVM/HIP artifact comparisons, shared goldens for all three flavors, full-PR hooks and the relative-path guard. All 1,179 platform file hashes match the validated source.

The prior 082651964df5 gfx1250 numerical run passed 13 tests / 49 cases on ROCm 10.1.0. This host-only cache change was not GPU-rerun; its emitted code is unchanged. Stable ROCm 10.0.0 evidence remains compilation only.

Assistance: codex / astra / medium.

@tenpercent
tenpercent marked this pull request as ready for review September 21, 2026 21:17
@tenpercent
tenpercent requested a review from CongMa13 September 21, 2026 23:49

@CongMa13 CongMa13 left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Static review of 76b60c7c7bbd80d059e353487a4fdcaad260b901 found one P2 issue in the C/C++ rocke.core.arch helper mirror: the canonical catalog lookup gained a scale filter, but the helper-level rocke_archtarget_op_for_shape still drops that dimension. The latest MmaScaleDType change appears to address the existing E5M3/E5M2 discussion, so I have not duplicated that thread.

Static review only; no build, tests, or project-code execution performed.

@yraparti
yraparti self-requested a review September 22, 2026 19:03

@CongMa13 CongMa13 left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Great. LGTM.

@tenpercent
tenpercent enabled auto-merge (squash) September 22, 2026 19:38
@tenpercent
tenpercent merged commit d874683 into develop Sep 23, 2026
92 checks passed
@tenpercent
tenpercent deleted the users/mpodkory/rocke/mxfp-common branch September 23, 2026 08:21
tenpercent added a commit that referenced this pull request Sep 24, 2026
`platform/tests/run_all.py` fails its existing architecture-isolation
test after #12287 because `core/arch/wmma_scale.py` constructs LLVM
intrinsic names. Move the LLVM carrier types, intrinsic names, and
declaration keys into `core/isa/wmma_scale.py`, and make the ISA emitter
consume that representation. Scale packing and catalog metadata remain
in `core/arch`.

The existing isolation guard is unchanged. Update the scaled-WMMA tests
to check LLVM spelling through the ISA representation while keeping
catalog assertions about fragment metadata. Native code and golden files
are unchanged.

ISSUE ID : AICK-2417

Validation:

- Focused architecture, scaled-WMMA, catalog, and debug object tests: 78
passed plus 26 subtests. The four debug object tests use ROCm 7.14's
LLVM tools.
- Fresh native build and 11/11 CTests passed.
- Python/C++ LLVM-IR parity covers all 70 families / 785 configurations
for LLVM20, LLVM22, and LLVM23. LLVM20/22: 732 identical outputs and 53
matching rejections each; LLVM23: 742 identical outputs and 43 matching
rejections. The initial LLVM22 pass hit three Python emitter timeouts;
rerunning all 20 configurations in those three families passed.
- Repository pre-commit hooks and relative-path guard passed.
- Host `tests/run_all.py --no-gate --build-root <fresh-build>` is GREEN:
1,036 passed, 92 skipped, and 11/11 CTests. ROCm 7.14 is selected
specifically for the debug object tools via `ROCKE_TEST_LLVM_BIN`. The
optional `ROCKE_BACKEND=both` pytest pass is skipped because the pybind
extension is not installed; native/Python comparison is covered
separately by the byte-identity gates.
- Full suite with `module load ubuntu-24 rocm/7.14.0`, compared with
unchanged base `3066840e5a0e`: 9 failures / 1,034 passes before, 8
failures / 1,035 passes after; 92 skips in each. The removed failure is
the architecture guard. Both revisions retain seven datalayout subtest
failures from reference-toolchain version detection and one
portable-recipe HSACO parity failure.

No GPU numerical validation was performed for this Python organization
change.
@assistant-librarian

Copy link
Copy Markdown
Contributor

TheRock Submodule Bump Activity

Newest first

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants