From b6567c4be5f40d23c4de19418856cefed596ae15 Mon Sep 17 00:00:00 2001 From: Nripesh Niketan Date: Thu, 13 Aug 2026 12:33:56 +0100 Subject: [PATCH] Integrate test-backed RMSNorm dispatch frontier Add a pinned host-dispatch contract for the RMSNorm forward and gradient workloads exercised by the MLX fast-operation tests. Translate the resulting 12 DirectX artifacts independently, retain their dispatch provenance, and compile every artifact in the Windows project-porting gate. Update the repository frontier accounting, evidence, documentation, and focused tests while preserving the explicit boundary between compiler validation and runtime or numerical parity. --- .github/workflows/mlx-project-porting.yml | 47 +- demos/integrations/mlx/README.md | 52 +- .../mlx/contracts/rms_norm.dispatch.json | 867 ++++++++++++++++++ demos/integrations/mlx/expected-gaps.json | 415 ++++++++- demos/integrations/mlx/run_mlx_porting.py | 753 ++++++++++++++- tests/test_ci_workflows.py | 5 +- tests/test_mlx_porting_harness.py | 412 ++++++++- ..._mlx_rms_norm_dispatch_contract_fixture.py | 210 +++++ 8 files changed, 2688 insertions(+), 73 deletions(-) create mode 100644 demos/integrations/mlx/contracts/rms_norm.dispatch.json create mode 100644 tests/test_mlx_rms_norm_dispatch_contract_fixture.py diff --git a/.github/workflows/mlx-project-porting.yml b/.github/workflows/mlx-project-porting.yml index 332958cc5..0364ecc7e 100644 --- a/.github/workflows/mlx-project-porting.yml +++ b/.github/workflows/mlx-project-porting.yml @@ -18,6 +18,7 @@ on: - "tests/test_mlx_copy_opengl_proof.py" - "tests/test_mlx_copy_directx_proof.py" - "tests/test_mlx_dispatch_contract_fixture.py" + - "tests/test_mlx_rms_norm_dispatch_contract_fixture.py" - "tests/test_mlx_layer_norm_directx_proof.py" - "tests/test_mlx_quantized_directx_proof.py" - "tests/test_mlx_quantized_opengl_proof.py" @@ -49,6 +50,7 @@ on: - "tests/test_mlx_copy_opengl_proof.py" - "tests/test_mlx_copy_directx_proof.py" - "tests/test_mlx_dispatch_contract_fixture.py" + - "tests/test_mlx_rms_norm_dispatch_contract_fixture.py" - "tests/test_mlx_layer_norm_directx_proof.py" - "tests/test_mlx_quantized_directx_proof.py" - "tests/test_mlx_quantized_opengl_proof.py" @@ -555,6 +557,10 @@ jobs: MLX_LAYER_NORM_DISPATCH_VARIANTS, MLX_LAYER_NORM_SHA256, MLX_LAYER_NORM_SOURCE, + MLX_RMS_NORM_DISPATCH_CONTENT_IDENTITY, + MLX_RMS_NORM_DISPATCH_VARIANTS, + MLX_RMS_NORM_SHA256, + MLX_RMS_NORM_SOURCE, MLX_OPENGL_DYNAMIC_WORKGROUP_FRONTIER_SOURCES, MLX_OPENGL_FRONTIER_SOURCES, MLX_OPENGL_INDEX_RANGE_ASSERTION_EXPRESSIONS, @@ -818,8 +824,8 @@ jobs: raise SystemExit( f"DirectX workgroup blocker evidence changed for {source}" ) - if directx_blocked_entry_point_count != 94: - raise SystemExit("expected 94 fail-closed DirectX compute entries") + if directx_blocked_entry_point_count != 82: + raise SystemExit("expected 82 fail-closed DirectX compute entries") layer_norm = directx["layerNormDispatchEvidence"] expected_layer_norm_status = ( "translated-dxc-validated" @@ -852,6 +858,43 @@ jobs: or layer_norm["numericalParityClaimed"] is not False ): raise SystemExit("LayerNorm dispatch frontier evidence is incomplete") + rms_norm = directx["rmsNormDispatchEvidence"] + expected_rms_norm_status = ( + "translated-dxc-validated" + if os.environ["RUNNER_OS"] == "Windows" + else "translated" + ) + expected_rms_norm_validated = ( + len(MLX_RMS_NORM_DISPATCH_VARIANTS) + if os.environ["RUNNER_OS"] == "Windows" + else 0 + ) + if ( + rms_norm["status"] != expected_rms_norm_status + or rms_norm["source"] != MLX_RMS_NORM_SOURCE + or rms_norm["sourceSha256"] != MLX_RMS_NORM_SHA256 + or rms_norm["target"] != "directx" + or rms_norm["testSources"] + != [ + "python/tests/test_fast.py::test_rms_norm", + "python/tests/test_fast.py::test_rms_norm_grad", + ] + or rms_norm["dispatchContract"]["contentIdentity"] + != MLX_RMS_NORM_DISPATCH_CONTENT_IDENTITY + or rms_norm["dispatchContract"]["variantCount"] + != len(MLX_RMS_NORM_DISPATCH_VARIANTS) + or rms_norm["dispatchContract"]["resolvedIssue"] + != MLX_HOST_DISPATCH_IMPORT_RESOLVED_ISSUE + or rms_norm["artifactCount"] + != len(MLX_RMS_NORM_DISPATCH_VARIANTS) + or set(rms_norm["variants"]) + != set(MLX_RMS_NORM_DISPATCH_VARIANTS) + or rms_norm["dxcValidatedArtifactCount"] + != expected_rms_norm_validated + or rms_norm["runtimeExecutionAttempted"] is not False + or rms_norm["numericalParityClaimed"] is not False + ): + raise SystemExit("RMSNorm dispatch frontier evidence is incomplete") if os.environ["RUNNER_OS"] == "Windows": if ( directx["directxToolchainRequired"] is not True diff --git a/demos/integrations/mlx/README.md b/demos/integrations/mlx/README.md index b92512f85..bf5b893c1 100644 --- a/demos/integrations/mlx/README.md +++ b/demos/integrations/mlx/README.md @@ -66,9 +66,10 @@ The current harness verifies: `rope.metal`, `scaled_dot_product_attention.metal`, `softmax.metal`, and `ternary.metal`. Vulkan must translate and structurally validate all 11. DirectX emits five aggregate artifacts whose entries do not require a - runtime-selected workgroup size, plus two entry-scoped `layer_norm.metal` - artifacts selected by the checked-in host dispatch contract. It records exact - expected failures for the other five sources. Each blocked report must retain + runtime-selected workgroup size, two entry-scoped `layer_norm.metal` + artifacts, and 12 entry-scoped `rms_norm.metal` artifacts selected by + checked-in host dispatch contracts. It records exact expected failures for + the other four sources. Each blocked report must retain the pinned total specialization count and exactly match its diagnostic entry names to the materialized host names, with no additional diagnostics. Separate configs prevent DirectX @@ -90,13 +91,13 @@ The current harness verifies: dispatch uses runtime axis and pipeline-limit operands unavailable to source materialization; - DirectX HLSL compiler checks with official DXC v1.9.2602.24 on Windows CI for - the seven-artifact frontier representing six pinned sources: `arange.metal`, + the 19-artifact frontier representing seven pinned sources: `arange.metal`, `binary_two.metal`, two bounded `layer_norm.metal` entries, `random.metal`, - `rope.metal`, and `ternary.metal`. At the pinned revision the gate compiles - 11, 225, 2, 2, 18, and 212 entries respectively, for 470 generated compute - entries in total. Each LayerNorm entry is emitted independently with its - host-derived workgroup size, specialization constants, and exact subgroup - width. The + 12 test-derived `rms_norm.metal` entries, `rope.metal`, and `ternary.metal`. + At the pinned revision the gate compiles 11, 225, 2, 2, 12, 18, and 212 + entries respectively, for 482 generated compute entries in total. Each + LayerNorm and RMSNorm artifact is emitted independently with its host-derived + workgroup size, specialization constants, and exact subgroup width. The pinned rope translation supplies required function constant IDs through the quoted `"1"`, `"2"`, and `"3"` selectors in `[project.specialization_constants]` and materializes the concrete DirectX @@ -113,7 +114,7 @@ The current harness verifies: metadata remains tracked by [#1542](https://github.com/CrossGL/crosstl/issues/1542). Host dispatch contract import was completed under [#1793](https://github.com/CrossGL/crosstl/issues/1793). - The five pending aggregate sources cover 94 compute entries and are asserted + The four pending aggregate sources cover 82 compute entries and are asserted as failed artifacts until bounded dispatch manifests are supplied; no placeholder workgroup size is restored. `fence.metal` is excluded because its DirectX translation intentionally fails under @@ -296,7 +297,7 @@ and fails closed if either field is missing or changes. Native-profile bfloat helpers now use exact `uint16_t` boundaries, and the two selected `random.metal` entries compile without the promotion warnings tracked by [#1799](https://github.com/CrossGL/crosstl/issues/1799). DXC reports zero -warnings across all 470 entry-point runs in the seven-artifact emitted frontier. +warnings across all 482 entry-point runs in the 19-artifact emitted frontier. The harness records this as a warning-clean contract and rejects any newly observed warning. Contextual destination conversion under [#1801](https://github.com/CrossGL/crosstl/issues/1801) is resolved for the @@ -854,6 +855,35 @@ evaluation with: tests/test_mlx_dispatch_contract_fixture.py ``` +The checked-in +[`contracts/rms_norm.dispatch.json`](contracts/rms_norm.dispatch.json) fixture +captures 12 distinct dispatch artifacts exercised by the pinned +`python/tests/test_fast.py::test_rms_norm` and `test_rms_norm_grad` workloads. +The forward records cover float32 workgroups of 32, 64, and 128 threads, +float16 and bfloat16 workgroups of 32 threads, and the 1024-thread looped path. +The VJP records cover 32- and 64-thread single-row paths plus the 1024-thread +looped path, each with both concrete values of function constant `20` +(`has_w`). Axis sizes 31, 32, and 33 share the same 32-thread artifact and are +recorded as covered inputs rather than duplicate artifacts. + +The repository harness evaluates those finite records against the unchanged +pinned `rms_norm.metal` source and requires 12 deterministic DirectX artifacts. +Every artifact retains its workload inputs, dispatch workgroup count, +host-derived `numthreads` value, concrete function constants, and +`[WaveSize(32)]` enforcement. Windows CI compiles every artifact with official +DXC `cs_6_6`, `-enable-16bit-types`, and warnings as errors. This is complete +translation and compiler coverage for the listed unit-test dispatch variants; +the MLX runtime is not redirected to these artifacts, the kernels are not +executed, and numerical parity is not claimed. + +Validate the RMSNorm contract schema, provenance, deterministic identities, and +bounded workload set with: + +```bash +.venv/bin/python -m pytest -q -n auto \ + tests/test_mlx_rms_norm_dispatch_contract_fixture.py +``` + The focused `prove_rms_norm_specialization.py` gate fixes the project-level RMSNorm specialization contract to the same upstream commit and to `rms_norm.metal` SHA-256 diff --git a/demos/integrations/mlx/contracts/rms_norm.dispatch.json b/demos/integrations/mlx/contracts/rms_norm.dispatch.json new file mode 100644 index 000000000..591996cab --- /dev/null +++ b/demos/integrations/mlx/contracts/rms_norm.dispatch.json @@ -0,0 +1,867 @@ +{ + "kind": "crosstl-host-dispatch-contract", + "schemaVersion": 1, + "provenance": { + "repository": "https://github.com/ml-explore/mlx", + "commit": "4367c73b60541ddd5a266ce4644fd93d20223b6e", + "sourceReferences": { + "hostDispatch": "mlx/backend/metal/normalization.cpp", + "kernel": "mlx/backend/metal/kernels/rms_norm.metal" + }, + "scope": { + "description": "Pinned RMSNorm unit-test dispatch records.", + "singleRowAndLoopedEntriesIncluded": true, + "runtimeExecutionVerified": false, + "numericalParityVerified": false + } + }, + "inputs": [ + { + "name": "axisSize", + "role": "shape", + "type": "integer", + "provenance": { + "hostSymbol": "axis_size" + } + }, + { + "name": "dtype", + "role": "dtype", + "type": "string" + }, + { + "name": "isVjp", + "role": "feature", + "type": "boolean", + "provenance": { + "hostFunctions": [ + "RMSNorm::eval_gpu", + "RMSNormVJP::eval_gpu" + ] + } + }, + { + "name": "hasW", + "role": "feature", + "type": "boolean", + "provenance": { + "functionConstant": "has_w", + "functionConstantId": 20 + } + } + ], + "workloads": [ + { + "id": "forward-float32-axis-32", + "values": { + "axisSize": 32, + "dtype": "float32", + "isVjp": false, + "hasW": false + }, + "provenance": { + "hostFunction": "RMSNorm::eval_gpu", + "branch": "single-row", + "testSource": "python/tests/test_fast.py::test_rms_norm", + "shape": [ + 2, + 32 + ], + "coveredAxisSizes": [ + 31, + 32, + 33 + ] + } + }, + { + "id": "forward-float32-axis-256", + "values": { + "axisSize": 256, + "dtype": "float32", + "isVjp": false, + "hasW": false + }, + "provenance": { + "hostFunction": "RMSNorm::eval_gpu", + "branch": "single-row", + "testSource": "python/tests/test_fast.py::test_rms_norm", + "shape": [ + 2, + 256 + ] + } + }, + { + "id": "forward-float32-axis-512", + "values": { + "axisSize": 512, + "dtype": "float32", + "isVjp": false, + "hasW": false + }, + "provenance": { + "hostFunction": "RMSNorm::eval_gpu", + "branch": "single-row", + "testSource": "python/tests/test_fast.py::test_rms_norm", + "shape": [ + 2, + 512 + ] + } + }, + { + "id": "forward-float16-axis-32", + "values": { + "axisSize": 32, + "dtype": "float16", + "isVjp": false, + "hasW": false + }, + "provenance": { + "hostFunction": "RMSNorm::eval_gpu", + "branch": "single-row", + "testSource": "python/tests/test_fast.py::test_rms_norm", + "shape": [ + 2, + 32 + ] + } + }, + { + "id": "forward-bfloat16-axis-32", + "values": { + "axisSize": 32, + "dtype": "bfloat16", + "isVjp": false, + "hasW": false + }, + "provenance": { + "hostFunction": "RMSNorm::eval_gpu", + "branch": "single-row", + "testSource": "python/tests/test_fast.py::test_rms_norm", + "shape": [ + 2, + 32 + ] + } + }, + { + "id": "forward-float32-axis-4099", + "values": { + "axisSize": 4099, + "dtype": "float32", + "isVjp": false, + "hasW": false + }, + "provenance": { + "hostFunction": "RMSNorm::eval_gpu", + "branch": "looped", + "testSource": "python/tests/test_fast.py::test_rms_norm", + "shape": [ + 4099 + ] + } + }, + { + "id": "vjp-float32-axis-32-has-w-false", + "values": { + "axisSize": 32, + "dtype": "float32", + "isVjp": true, + "hasW": false + }, + "provenance": { + "hostFunction": "RMSNormVJP::eval_gpu", + "branch": "single-row", + "testSource": "python/tests/test_fast.py::test_rms_norm_grad", + "shape": [ + 8, + 100, + 32 + ] + } + }, + { + "id": "vjp-float32-axis-32-has-w-true", + "values": { + "axisSize": 32, + "dtype": "float32", + "isVjp": true, + "hasW": true + }, + "provenance": { + "hostFunction": "RMSNormVJP::eval_gpu", + "branch": "single-row", + "testSource": "python/tests/test_fast.py::test_rms_norm_grad", + "shape": [ + 8, + 100, + 32 + ] + } + }, + { + "id": "vjp-float32-axis-256-has-w-false", + "values": { + "axisSize": 256, + "dtype": "float32", + "isVjp": true, + "hasW": false + }, + "provenance": { + "hostFunction": "RMSNormVJP::eval_gpu", + "branch": "single-row", + "testSource": "python/tests/test_fast.py::test_rms_norm_grad", + "shape": [ + 8, + 100, + 256 + ] + } + }, + { + "id": "vjp-float32-axis-256-has-w-true", + "values": { + "axisSize": 256, + "dtype": "float32", + "isVjp": true, + "hasW": true + }, + "provenance": { + "hostFunction": "RMSNormVJP::eval_gpu", + "branch": "single-row", + "testSource": "python/tests/test_fast.py::test_rms_norm_grad", + "shape": [ + 8, + 100, + 256 + ] + } + }, + { + "id": "vjp-float32-axis-8192-has-w-false", + "values": { + "axisSize": 8192, + "dtype": "float32", + "isVjp": true, + "hasW": false + }, + "provenance": { + "hostFunction": "RMSNormVJP::eval_gpu", + "branch": "looped", + "testSource": "python/tests/test_fast.py::test_rms_norm_grad", + "shape": [ + 2, + 2, + 8192 + ] + } + }, + { + "id": "vjp-float32-axis-8192-has-w-true", + "values": { + "axisSize": 8192, + "dtype": "float32", + "isVjp": true, + "hasW": true + }, + "provenance": { + "hostFunction": "RMSNormVJP::eval_gpu", + "branch": "looped", + "testSource": "python/tests/test_fast.py::test_rms_norm_grad", + "shape": [ + 2, + 2, + 8192 + ] + } + } + ], + "capabilities": [ + { + "name": "simdWidth", + "type": "integer", + "provenance": { + "kernelSymbol": "SIMD_SIZE" + } + }, + { + "name": "maxThreadsPerWorkgroup", + "type": "integer", + "provenance": { + "requirement": "bounded device record" + } + } + ], + "devices": [ + { + "id": "wave32-max1024", + "values": { + "simdWidth": 32, + "maxThreadsPerWorkgroup": 1024 + }, + "provenance": { + "scope": "bounded wave32 device record" + } + } + ], + "contracts": [ + { + "id": "mlx-rms-norm-unit-tests", + "source": "mlx/backend/metal/kernels/rms_norm.metal", + "provenance": { + "hostSource": "mlx/backend/metal/normalization.cpp", + "kernelSource": "mlx/backend/metal/kernels/rms_norm.metal", + "hostFormula": "32 * ceilDiv(ceilDiv(axisSize, 4), 32), or maxThreadsPerWorkgroup for looped entries", + "dispatchNormalization": "single-row workgroup count derived from dispatch_threads" + }, + "branches": [ + { + "id": "forward-looped-float32", + "when": { + "op": "all", + "args": [ + { + "op": "eq", + "args": [ + { + "input": "dtype" + }, + "float32" + ] + }, + { + "op": "eq", + "args": [ + { + "input": "isVjp" + }, + false + ] + }, + { + "op": "gt", + "args": [ + { + "input": "axisSize" + }, + 4096 + ] + }, + { + "op": "eq", + "args": [ + { + "capability": "simdWidth" + }, + 32 + ] + }, + { + "op": "eq", + "args": [ + { + "capability": "maxThreadsPerWorkgroup" + }, + 1024 + ] + } + ] + }, + "entryPoint": "rms_loopedfloat32", + "workgroupSize": [ + { + "capability": "maxThreadsPerWorkgroup" + }, + 1, + 1 + ], + "subgroupWidth": { + "capability": "simdWidth" + }, + "dispatch": { + "workgroupCount": [ + 1, + 1, + 1 + ] + }, + "provenance": { + "hostFunction": "RMSNorm::eval_gpu", + "branch": "looped", + "functionConstant20Applied": false + } + }, + { + "id": "forward-single-row-float32", + "when": { + "op": "all", + "args": [ + { + "op": "eq", + "args": [ + { + "input": "dtype" + }, + "float32" + ] + }, + { + "op": "eq", + "args": [ + { + "input": "isVjp" + }, + false + ] + }, + { + "op": "le", + "args": [ + { + "input": "axisSize" + }, + 4096 + ] + }, + { + "op": "eq", + "args": [ + { + "capability": "simdWidth" + }, + 32 + ] + }, + { + "op": "eq", + "args": [ + { + "capability": "maxThreadsPerWorkgroup" + }, + 1024 + ] + } + ] + }, + "entryPoint": "rmsfloat32", + "workgroupSize": [ + { + "op": "multiply", + "args": [ + 32, + { + "op": "ceilDiv", + "args": [ + { + "op": "ceilDiv", + "args": [ + { + "input": "axisSize" + }, + 4 + ] + }, + 32 + ] + } + ] + }, + 1, + 1 + ], + "subgroupWidth": { + "capability": "simdWidth" + }, + "dispatch": { + "workgroupCount": [ + 2, + 1, + 1 + ] + }, + "provenance": { + "hostFunction": "RMSNorm::eval_gpu", + "branch": "single-row", + "functionConstant20Applied": false + } + }, + { + "id": "forward-single-row-float16", + "when": { + "op": "all", + "args": [ + { + "op": "eq", + "args": [ + { + "input": "dtype" + }, + "float16" + ] + }, + { + "op": "eq", + "args": [ + { + "input": "isVjp" + }, + false + ] + }, + { + "op": "le", + "args": [ + { + "input": "axisSize" + }, + 4096 + ] + }, + { + "op": "eq", + "args": [ + { + "capability": "simdWidth" + }, + 32 + ] + }, + { + "op": "eq", + "args": [ + { + "capability": "maxThreadsPerWorkgroup" + }, + 1024 + ] + } + ] + }, + "entryPoint": "rmsfloat16", + "workgroupSize": [ + { + "op": "multiply", + "args": [ + 32, + { + "op": "ceilDiv", + "args": [ + { + "op": "ceilDiv", + "args": [ + { + "input": "axisSize" + }, + 4 + ] + }, + 32 + ] + } + ] + }, + 1, + 1 + ], + "subgroupWidth": { + "capability": "simdWidth" + }, + "dispatch": { + "workgroupCount": [ + 2, + 1, + 1 + ] + }, + "provenance": { + "hostFunction": "RMSNorm::eval_gpu", + "branch": "single-row", + "functionConstant20Applied": false + } + }, + { + "id": "forward-single-row-bfloat16", + "when": { + "op": "all", + "args": [ + { + "op": "eq", + "args": [ + { + "input": "dtype" + }, + "bfloat16" + ] + }, + { + "op": "eq", + "args": [ + { + "input": "isVjp" + }, + false + ] + }, + { + "op": "le", + "args": [ + { + "input": "axisSize" + }, + 4096 + ] + }, + { + "op": "eq", + "args": [ + { + "capability": "simdWidth" + }, + 32 + ] + }, + { + "op": "eq", + "args": [ + { + "capability": "maxThreadsPerWorkgroup" + }, + 1024 + ] + } + ] + }, + "entryPoint": "rmsbfloat16", + "workgroupSize": [ + { + "op": "multiply", + "args": [ + 32, + { + "op": "ceilDiv", + "args": [ + { + "op": "ceilDiv", + "args": [ + { + "input": "axisSize" + }, + 4 + ] + }, + 32 + ] + } + ] + }, + 1, + 1 + ], + "subgroupWidth": { + "capability": "simdWidth" + }, + "dispatch": { + "workgroupCount": [ + 2, + 1, + 1 + ] + }, + "provenance": { + "hostFunction": "RMSNorm::eval_gpu", + "branch": "single-row", + "functionConstant20Applied": false + } + }, + { + "id": "vjp-looped-float32", + "when": { + "op": "all", + "args": [ + { + "op": "eq", + "args": [ + { + "input": "dtype" + }, + "float32" + ] + }, + { + "op": "eq", + "args": [ + { + "input": "isVjp" + }, + true + ] + }, + { + "op": "gt", + "args": [ + { + "input": "axisSize" + }, + 4096 + ] + }, + { + "op": "eq", + "args": [ + { + "capability": "simdWidth" + }, + 32 + ] + }, + { + "op": "eq", + "args": [ + { + "capability": "maxThreadsPerWorkgroup" + }, + 1024 + ] + } + ] + }, + "entryPoint": "vjp_rms_loopedfloat32", + "workgroupSize": [ + { + "capability": "maxThreadsPerWorkgroup" + }, + 1, + 1 + ], + "subgroupWidth": { + "capability": "simdWidth" + }, + "specializationConstants": { + "20": { + "input": "hasW", + "provenance": { + "symbol": "has_w" + } + } + }, + "dispatch": { + "workgroupCount": [ + 4, + 1, + 1 + ] + }, + "provenance": { + "hostFunction": "RMSNormVJP::eval_gpu", + "branch": "looped", + "functionConstant20Applied": true + } + }, + { + "id": "vjp-single-row-float32", + "when": { + "op": "all", + "args": [ + { + "op": "eq", + "args": [ + { + "input": "dtype" + }, + "float32" + ] + }, + { + "op": "eq", + "args": [ + { + "input": "isVjp" + }, + true + ] + }, + { + "op": "le", + "args": [ + { + "input": "axisSize" + }, + 4096 + ] + }, + { + "op": "eq", + "args": [ + { + "capability": "simdWidth" + }, + 32 + ] + }, + { + "op": "eq", + "args": [ + { + "capability": "maxThreadsPerWorkgroup" + }, + 1024 + ] + } + ] + }, + "entryPoint": "vjp_rmsfloat32", + "workgroupSize": [ + { + "op": "multiply", + "args": [ + 32, + { + "op": "ceilDiv", + "args": [ + { + "op": "ceilDiv", + "args": [ + { + "input": "axisSize" + }, + 4 + ] + }, + 32 + ] + } + ] + }, + 1, + 1 + ], + "subgroupWidth": { + "capability": "simdWidth" + }, + "specializationConstants": { + "20": { + "input": "hasW", + "provenance": { + "symbol": "has_w" + } + } + }, + "dispatch": { + "workgroupCount": [ + 800, + 1, + 1 + ] + }, + "provenance": { + "hostFunction": "RMSNormVJP::eval_gpu", + "branch": "single-row", + "functionConstant20Applied": true + } + } + ] + } + ] +} diff --git a/demos/integrations/mlx/expected-gaps.json b/demos/integrations/mlx/expected-gaps.json index 106da5d6d..db24f3178 100644 --- a/demos/integrations/mlx/expected-gaps.json +++ b/demos/integrations/mlx/expected-gaps.json @@ -10,13 +10,13 @@ "directx", "vulkan" ], - "artifacts": 23, - "translated_artifacts": 18, - "failed_artifacts": 5, + "artifacts": 34, + "translated_artifacts": 30, + "failed_artifacts": 4, "target_artifacts": { "directx": { - "translated": 7, - "failed": 5 + "translated": 19, + "failed": 4 }, "vulkan": { "translated": 11, @@ -30,7 +30,6 @@ "pending_host_dispatch_sources": [ "mlx/backend/metal/kernels/arg_reduce.metal", "mlx/backend/metal/kernels/logsumexp.metal", - "mlx/backend/metal/kernels/rms_norm.metal", "mlx/backend/metal/kernels/scaled_dot_product_attention.metal", "mlx/backend/metal/kernels/softmax.metal" ], @@ -1480,14 +1479,14 @@ }, "directx_toolchain_status": { "status": "passing-with-bounded-dispatch-and-pending-contracts", - "note": "Windows CI compiles 470 generated compute entries across seven DirectX HLSL artifacts from six pinned sources with zero DXC warnings. LayerNorm contributes two entry-scoped artifacts selected by a replayable host dispatch contract. Five sources covering 94 aggregate entries still fail before emission because no bounded host dispatch manifest is configured; no placeholder workgroup size is applied. This is native compiler validation only and does not establish Direct3D runtime execution or numerical parity.", + "note": "Windows CI compiles 482 generated compute entries across 19 DirectX HLSL artifacts from seven pinned sources with zero DXC warnings. LayerNorm contributes two entry-scoped artifacts, and RMSNorm contributes 12 unit-test-derived entry-scoped artifacts selected by replayable host dispatch contracts. Four sources covering 82 aggregate entries still fail before emission because no bounded host dispatch manifest is configured; no placeholder workgroup size is applied. This is native compiler validation only and does not establish Direct3D runtime execution or numerical parity.", "compiler": { "name": "dxc", "version": "v1.9.2602.24" }, "warning_evidence": { "status": "warning-clean", - "validatedRunCount": 470, + "validatedRunCount": 482, "warningRunCount": 0, "observedWarningCount": 0, "uniqueContractCount": 0, @@ -1770,6 +1769,377 @@ "runtime_execution_attempted": false, "numerical_parity_claimed": false }, + "rms_norm_dispatch_frontier": { + "status": "translated-dxc-validated", + "source": "mlx/backend/metal/kernels/rms_norm.metal", + "source_sha256": "5d411a2350ba7ddf84eb35f9dcac7cde0d441bd55fa1e9e1ccc61d490d428dee", + "target": "directx", + "test_sources": [ + "python/tests/test_fast.py::test_rms_norm", + "python/tests/test_fast.py::test_rms_norm_grad" + ], + "dispatch_contract": { + "path": "demos/integrations/mlx/contracts/rms_norm.dispatch.json", + "normalized_sha256": "ea91c5b5655e776e6db3537bf6539d1f1983c0a4ed74bf90f96a57ad3611e91f", + "content_identity": "sha256:ea238af83b140c33d43b79b5efd1814c398bbc09ac70cbef21375b1e8ff9a1eb", + "variant_count": 12, + "resolved_issue": "https://github.com/CrossGL/crosstl/issues/1793" + }, + "artifact_count": 12, + "variants": { + "vjp-float32-axis-32-has-w-false": { + "entry_point": "vjp_rmsfloat32", + "artifact_id": "sha256:ef832e1ceb8c864a13aee3460d23658f4fffba18db1b800461628ba6ebe38e0a", + "dispatch_variant_id": "sha256:1a25dca51070c3b6fc96f162e6c152049d388833ba02b1f3d10cc1928c5661c4", + "inputs": { + "axisSize": 32, + "dtype": "float32", + "hasW": false, + "isVjp": true + }, + "workgroup_size": [ + 32, + 1, + 1 + ], + "dispatch_workgroup_count": [ + 800, + 1, + 1 + ], + "subgroup_width": 32, + "subgroup_width_enforcement": "WaveSize(32)", + "specialization_constants": { + "20": false + }, + "generated_hlsl": { + "normalized_sha256": "316c6393196ff1ef60ac143cfbb7946635648f296a4cf36df215255cbeadf408", + "size_bytes": 5440 + } + }, + "vjp-float32-axis-32-has-w-true": { + "entry_point": "vjp_rmsfloat32", + "artifact_id": "sha256:a9be06b43a6156fb9ee1f9a6955d03d6bda0940c2a8223b58f564c2d12bd0cd0", + "dispatch_variant_id": "sha256:26177a77e484a56b7b2572516e8e1360714c88b668881237a9cf499129e34f35", + "inputs": { + "axisSize": 32, + "dtype": "float32", + "hasW": true, + "isVjp": true + }, + "workgroup_size": [ + 32, + 1, + 1 + ], + "dispatch_workgroup_count": [ + 800, + 1, + 1 + ], + "subgroup_width": 32, + "subgroup_width_enforcement": "WaveSize(32)", + "specialization_constants": { + "20": true + }, + "generated_hlsl": { + "normalized_sha256": "258a16ae4d2132e8500e2b08180b87322d4f8a223a974fa18b47ab77fa160908", + "size_bytes": 5439 + } + }, + "forward-float32-axis-512": { + "entry_point": "rmsfloat32", + "artifact_id": "sha256:a9b6980b1645e867b2502052d6d5f37a03447258bb4f819535901bf45a01d5da", + "dispatch_variant_id": "sha256:369be00b341be529b913101657bd6a78d8b8c689e9dc7c6d27e694809dd9d098", + "inputs": { + "axisSize": 512, + "dtype": "float32", + "hasW": false, + "isVjp": false + }, + "workgroup_size": [ + 128, + 1, + 1 + ], + "dispatch_workgroup_count": [ + 2, + 1, + 1 + ], + "subgroup_width": 32, + "subgroup_width_enforcement": "WaveSize(32)", + "specialization_constants": {}, + "generated_hlsl": { + "normalized_sha256": "3790adf689bd3460602073ce19312dcc2fc631247c5e75859179e063ae1e38f5", + "size_bytes": 3646 + } + }, + "forward-float32-axis-4099": { + "entry_point": "rms_loopedfloat32", + "artifact_id": "sha256:b81c2043b10bde966cb6f4dbfa198d2b93a3e456f3026030b69557c4a8983729", + "dispatch_variant_id": "sha256:392aee49734fcc2fd3ff4fd232a49f1a97a2bb79f3c22cc80eeabeb1f1ca1959", + "inputs": { + "axisSize": 4099, + "dtype": "float32", + "hasW": false, + "isVjp": false + }, + "workgroup_size": [ + 1024, + 1, + 1 + ], + "dispatch_workgroup_count": [ + 1, + 1, + 1 + ], + "subgroup_width": 32, + "subgroup_width_enforcement": "WaveSize(32)", + "specialization_constants": {}, + "generated_hlsl": { + "normalized_sha256": "0c05c3e676b55f567e18dd9a1d4ede3c84876945a33c22405e442174ac53cc85", + "size_bytes": 4262 + } + }, + "forward-float32-axis-32": { + "entry_point": "rmsfloat32", + "artifact_id": "sha256:00c05fccf276cf11f3fb9b617b8fe0bb3c5f8766c0e4ca1ed990c093e700422e", + "dispatch_variant_id": "sha256:4306831dce3a9a479ef63093a7f2722358caf58b6fcf1a47ed808a1c28dc9ebb", + "inputs": { + "axisSize": 32, + "dtype": "float32", + "hasW": false, + "isVjp": false + }, + "workgroup_size": [ + 32, + 1, + 1 + ], + "dispatch_workgroup_count": [ + 2, + 1, + 1 + ], + "subgroup_width": 32, + "subgroup_width_enforcement": "WaveSize(32)", + "specialization_constants": {}, + "generated_hlsl": { + "normalized_sha256": "bbb3dbbeee01c267ba2b96aa048fe96af45740ad7f88bce5d0f584904b5a9d04", + "size_bytes": 3644 + } + }, + "forward-float16-axis-32": { + "entry_point": "rmsfloat16", + "artifact_id": "sha256:b694e5240f2a87bfae8af862878251a45cfaeaf39fd810bf2df8a5e3724bdad7", + "dispatch_variant_id": "sha256:51de914a20a4defb1d1b79ed26def94a994894228d6c64cae0e1132553809dbb", + "inputs": { + "axisSize": 32, + "dtype": "float16", + "hasW": false, + "isVjp": false + }, + "workgroup_size": [ + 32, + 1, + 1 + ], + "dispatch_workgroup_count": [ + 2, + 1, + 1 + ], + "subgroup_width": 32, + "subgroup_width_enforcement": "WaveSize(32)", + "specialization_constants": {}, + "generated_hlsl": { + "normalized_sha256": "09f5cb40981ae6a29c09a6e69f394657c0668a05721f3af91c30385fa0e083e5", + "size_bytes": 3664 + } + }, + "forward-bfloat16-axis-32": { + "entry_point": "rmsbfloat16", + "artifact_id": "sha256:13655322998b557a5143ac9b871dc898f0ae43c07cf9659aecd505156ed9318b", + "dispatch_variant_id": "sha256:635d28f468bf3e15a0ec7285aacc23d0de32f9596d511beb98adb3ecc2c2abf4", + "inputs": { + "axisSize": 32, + "dtype": "bfloat16", + "hasW": false, + "isVjp": false + }, + "workgroup_size": [ + 32, + 1, + 1 + ], + "dispatch_workgroup_count": [ + 2, + 1, + 1 + ], + "subgroup_width": 32, + "subgroup_width_enforcement": "WaveSize(32)", + "specialization_constants": {}, + "generated_hlsl": { + "normalized_sha256": "4280418858612f83d234445dfb1aa336ba66e215c5a416f134debdb50f8e92ac", + "size_bytes": 4736 + } + }, + "vjp-float32-axis-8192-has-w-true": { + "entry_point": "vjp_rms_loopedfloat32", + "artifact_id": "sha256:345d524ffec14682b6d0325bc97b624b89d83dc257a17ed49bea5e11e24573f3", + "dispatch_variant_id": "sha256:8fc64f93a1be95f7acac67d6595c9f9c66c01b87cf952f29052835c20d4d765b", + "inputs": { + "axisSize": 8192, + "dtype": "float32", + "hasW": true, + "isVjp": true + }, + "workgroup_size": [ + 1024, + 1, + 1 + ], + "dispatch_workgroup_count": [ + 4, + 1, + 1 + ], + "subgroup_width": 32, + "subgroup_width_enforcement": "WaveSize(32)", + "specialization_constants": { + "20": true + }, + "generated_hlsl": { + "normalized_sha256": "7f1712feeb2b01c6e3dfa2ad3384aa675af0e1c39bb883ac24c93eabe6ebb172", + "size_bytes": 6408 + } + }, + "vjp-float32-axis-256-has-w-true": { + "entry_point": "vjp_rmsfloat32", + "artifact_id": "sha256:0944044e2f050bedde1d05e1ae5648628e7144c04752dd49b2e2bb7bcd807b7b", + "dispatch_variant_id": "sha256:a010324f59769fde9d71cc8968852ae1e6c8b0ddc213f035523c2c5f2e12d413", + "inputs": { + "axisSize": 256, + "dtype": "float32", + "hasW": true, + "isVjp": true + }, + "workgroup_size": [ + 64, + 1, + 1 + ], + "dispatch_workgroup_count": [ + 800, + 1, + 1 + ], + "subgroup_width": 32, + "subgroup_width_enforcement": "WaveSize(32)", + "specialization_constants": { + "20": true + }, + "generated_hlsl": { + "normalized_sha256": "cd74fa1af043a346a605922ff79cd97b6d51e9b444ffd8db635c4efc01f87ea4", + "size_bytes": 5439 + } + }, + "vjp-float32-axis-256-has-w-false": { + "entry_point": "vjp_rmsfloat32", + "artifact_id": "sha256:4455fc9204f826fc5d0d7f016bcaf970b75ea31d5f1598d13ceca6a2baa369e7", + "dispatch_variant_id": "sha256:f26ffb357ecfabef6216ff45661f298af4f164a0494da838ca579ea0812e73d2", + "inputs": { + "axisSize": 256, + "dtype": "float32", + "hasW": false, + "isVjp": true + }, + "workgroup_size": [ + 64, + 1, + 1 + ], + "dispatch_workgroup_count": [ + 800, + 1, + 1 + ], + "subgroup_width": 32, + "subgroup_width_enforcement": "WaveSize(32)", + "specialization_constants": { + "20": false + }, + "generated_hlsl": { + "normalized_sha256": "5d43f9c94cf98ef57ecf61339e989cdead9d91d6b25d337042e6b03e03de405b", + "size_bytes": 5440 + } + }, + "forward-float32-axis-256": { + "entry_point": "rmsfloat32", + "artifact_id": "sha256:1ef80b00c1a7a2f7967177bc003a961f3e4448358716d2deb79778fc3cbfb68e", + "dispatch_variant_id": "sha256:f27b102d2ee4f473afef2448e42103223922ff29fd703e6d269d827f280e8bf7", + "inputs": { + "axisSize": 256, + "dtype": "float32", + "hasW": false, + "isVjp": false + }, + "workgroup_size": [ + 64, + 1, + 1 + ], + "dispatch_workgroup_count": [ + 2, + 1, + 1 + ], + "subgroup_width": 32, + "subgroup_width_enforcement": "WaveSize(32)", + "specialization_constants": {}, + "generated_hlsl": { + "normalized_sha256": "5ff7dcbf0f7236769184654424172b3edc34539547fa33b793745667083ca03e", + "size_bytes": 3644 + } + }, + "vjp-float32-axis-8192-has-w-false": { + "entry_point": "vjp_rms_loopedfloat32", + "artifact_id": "sha256:3bd55b546fc00ddf8412f092da4793c0272eec0ad7c130065ad7c1677f60cdce", + "dispatch_variant_id": "sha256:fa4d57c473cd5799c46a982e3ea339debc06de21fb6a255adc11bbb546f9329b", + "inputs": { + "axisSize": 8192, + "dtype": "float32", + "hasW": false, + "isVjp": true + }, + "workgroup_size": [ + 1024, + 1, + 1 + ], + "dispatch_workgroup_count": [ + 4, + 1, + 1 + ], + "subgroup_width": 32, + "subgroup_width_enforcement": "WaveSize(32)", + "specialization_constants": { + "20": false + }, + "generated_hlsl": { + "normalized_sha256": "30d1fb23fc6f6587235a51f8b28d4d1bd520d41051eea0a98cf28a1b62c646b6", + "size_bytes": 6409 + } + } + }, + "dxc_validated_artifact_count": 12, + "runtime_execution_attempted": false, + "numerical_parity_claimed": false + }, "specialization_constants": { "1": false, "2": false, @@ -1787,15 +2157,17 @@ "mlx/backend/metal/kernels/binary_two.metal": 225, "mlx/backend/metal/kernels/layer_norm.metal": 2, "mlx/backend/metal/kernels/random.metal": 2, + "mlx/backend/metal/kernels/rms_norm.metal": 12, "mlx/backend/metal/kernels/rope.metal": 18, "mlx/backend/metal/kernels/ternary.metal": 212 }, - "expected_entry_point_count": 470, + "expected_entry_point_count": 482, "dxc_validated_sources": [ "mlx/backend/metal/kernels/arange.metal", "mlx/backend/metal/kernels/binary_two.metal", "mlx/backend/metal/kernels/layer_norm.metal", "mlx/backend/metal/kernels/random.metal", + "mlx/backend/metal/kernels/rms_norm.metal", "mlx/backend/metal/kernels/rope.metal", "mlx/backend/metal/kernels/ternary.metal" ], @@ -1864,18 +2236,16 @@ "workgroup_blocked_sources": [ "mlx/backend/metal/kernels/arg_reduce.metal", "mlx/backend/metal/kernels/logsumexp.metal", - "mlx/backend/metal/kernels/rms_norm.metal", "mlx/backend/metal/kernels/scaled_dot_product_attention.metal", "mlx/backend/metal/kernels/softmax.metal" ], "workgroup_blocked_entry_point_counts": { "mlx/backend/metal/kernels/arg_reduce.metal": 24, "mlx/backend/metal/kernels/logsumexp.metal": 6, - "mlx/backend/metal/kernels/rms_norm.metal": 12, "mlx/backend/metal/kernels/scaled_dot_product_attention.metal": 42, "mlx/backend/metal/kernels/softmax.metal": 10 }, - "workgroup_blocked_entry_point_count": 94, + "workgroup_blocked_entry_point_count": 82, "workgroup_blocked_diagnostic": "project.translate.workgroup-size-entry-ambiguous", "host_dispatch_import_resolved_by": "https://github.com/CrossGL/crosstl/issues/1793", "dispatch_evidence": { @@ -1920,26 +2290,6 @@ "T" ] }, - "mlx/backend/metal/kernels/rms_norm.metal": { - "specializationCount": 12, - "hostSource": "mlx/backend/metal/normalization.cpp", - "hostLines": "52-91,137-197", - "dispatchFormulas": [ - "block: [32 * ceil_div(ceil_div(axis_size, 4), 32), 1, 1]", - "looped: [maxTotalThreadsPerThreadgroup, 1, 1]" - ], - "dispatchSelection": [ - "block when axis_size <= 4096; looped when axis_size > 4096" - ], - "runtimeOperands": [ - "axis_size", - "maxTotalThreadsPerThreadgroup" - ], - "materializationParameters": [ - "N_READS", - "T" - ] - }, "mlx/backend/metal/kernels/scaled_dot_product_attention.metal": { "specializationCount": 42, "hostSource": "mlx/backend/metal/scaled_dot_product_attention.cpp", @@ -1994,7 +2344,6 @@ "directx_toolchain_gaps": { "mlx/backend/metal/kernels/arg_reduce.metal": "no bounded host dispatch manifest is configured for the runtime-selected axis and pipeline limit; aggregate translation fails with project.translate.workgroup-size-entry-ambiguous (host dispatch import support resolved in https://github.com/CrossGL/crosstl/issues/1793)", "mlx/backend/metal/kernels/logsumexp.metal": "no bounded host dispatch manifest is configured for the runtime-selected axis and pipeline limit; aggregate translation fails with project.translate.workgroup-size-entry-ambiguous (host dispatch import support resolved in https://github.com/CrossGL/crosstl/issues/1793)", - "mlx/backend/metal/kernels/rms_norm.metal": "no bounded host dispatch manifest is configured for the runtime-selected axis and pipeline limit; aggregate translation fails with project.translate.workgroup-size-entry-ambiguous (host dispatch import support resolved in https://github.com/CrossGL/crosstl/issues/1793)", "mlx/backend/metal/kernels/scaled_dot_product_attention.metal": "no bounded host dispatch manifest is configured for the runtime-selected pass, shape, and device architecture; aggregate translation fails with project.translate.workgroup-size-entry-ambiguous (host dispatch import support resolved in https://github.com/CrossGL/crosstl/issues/1793)", "mlx/backend/metal/kernels/softmax.metal": "no bounded host dispatch manifest is configured for the runtime-selected axis and pipeline limit; aggregate translation fails with project.translate.workgroup-size-entry-ambiguous (host dispatch import support resolved in https://github.com/CrossGL/crosstl/issues/1793)", "mlx/backend/metal/kernels/fence.metal": "the requested mem_device, memory_order_seq_cst, thread_scope_system atomic-fence contract is not representable in HLSL and fails with project.translate.directx-atomic-fence-unsupported (https://github.com/CrossGL/crosstl/issues/1537)" diff --git a/demos/integrations/mlx/run_mlx_porting.py b/demos/integrations/mlx/run_mlx_porting.py index 1c84ea525..cdaea2b84 100644 --- a/demos/integrations/mlx/run_mlx_porting.py +++ b/demos/integrations/mlx/run_mlx_porting.py @@ -93,6 +93,234 @@ MLX_QUANTIZED_SOURCE = "mlx/backend/metal/kernels/quantized.metal" MLX_QUANTIZED_SELECTED_ENTRY_POINT = "affine_quantize_float_gs_32_b_2" MLX_RMS_NORM_SOURCE = "mlx/backend/metal/kernels/rms_norm.metal" +MLX_RMS_NORM_SHA256 = "5d411a2350ba7ddf84eb35f9dcac7cde0d441bd55fa1e9e1ccc61d490d428dee" +MLX_RMS_NORM_DISPATCH_CONTRACT_SOURCE = ( + Path(__file__).resolve().parent / "contracts" / "rms_norm.dispatch.json" +) +MLX_RMS_NORM_DISPATCH_NORMALIZED_SHA256 = ( + "ea91c5b5655e776e6db3537bf6539d1f1983c0a4ed74bf90f96a57ad3611e91f" +) +MLX_RMS_NORM_DISPATCH_CONTENT_IDENTITY = ( + "sha256:ea238af83b140c33d43b79b5efd1814c398bbc09ac70cbef21375b1e8ff9a1eb" +) +MLX_RMS_NORM_DISPATCH_VARIANTS = { + "forward-float32-axis-32": { + "entryPoint": "rmsfloat32", + "artifactId": ( + "sha256:00c05fccf276cf11f3fb9b617b8fe0bb3c5f8766c0e4ca1ed990c093e700422e" + ), + "dispatchVariantId": ( + "sha256:4306831dce3a9a479ef63093a7f2722358caf58b6fcf1a47ed808a1c28dc9ebb" + ), + "inputs": { + "axisSize": 32, + "dtype": "float32", + "hasW": False, + "isVjp": False, + }, + "workgroupSize": [32, 1, 1], + "dispatchWorkgroupCount": [2, 1, 1], + "specializationConstants": {}, + }, + "forward-float32-axis-256": { + "entryPoint": "rmsfloat32", + "artifactId": ( + "sha256:1ef80b00c1a7a2f7967177bc003a961f3e4448358716d2deb79778fc3cbfb68e" + ), + "dispatchVariantId": ( + "sha256:f27b102d2ee4f473afef2448e42103223922ff29fd703e6d269d827f280e8bf7" + ), + "inputs": { + "axisSize": 256, + "dtype": "float32", + "hasW": False, + "isVjp": False, + }, + "workgroupSize": [64, 1, 1], + "dispatchWorkgroupCount": [2, 1, 1], + "specializationConstants": {}, + }, + "forward-float32-axis-512": { + "entryPoint": "rmsfloat32", + "artifactId": ( + "sha256:a9b6980b1645e867b2502052d6d5f37a03447258bb4f819535901bf45a01d5da" + ), + "dispatchVariantId": ( + "sha256:369be00b341be529b913101657bd6a78d8b8c689e9dc7c6d27e694809dd9d098" + ), + "inputs": { + "axisSize": 512, + "dtype": "float32", + "hasW": False, + "isVjp": False, + }, + "workgroupSize": [128, 1, 1], + "dispatchWorkgroupCount": [2, 1, 1], + "specializationConstants": {}, + }, + "forward-float16-axis-32": { + "entryPoint": "rmsfloat16", + "artifactId": ( + "sha256:b694e5240f2a87bfae8af862878251a45cfaeaf39fd810bf2df8a5e3724bdad7" + ), + "dispatchVariantId": ( + "sha256:51de914a20a4defb1d1b79ed26def94a994894228d6c64cae0e1132553809dbb" + ), + "inputs": { + "axisSize": 32, + "dtype": "float16", + "hasW": False, + "isVjp": False, + }, + "workgroupSize": [32, 1, 1], + "dispatchWorkgroupCount": [2, 1, 1], + "specializationConstants": {}, + }, + "forward-bfloat16-axis-32": { + "entryPoint": "rmsbfloat16", + "artifactId": ( + "sha256:13655322998b557a5143ac9b871dc898f0ae43c07cf9659aecd505156ed9318b" + ), + "dispatchVariantId": ( + "sha256:635d28f468bf3e15a0ec7285aacc23d0de32f9596d511beb98adb3ecc2c2abf4" + ), + "inputs": { + "axisSize": 32, + "dtype": "bfloat16", + "hasW": False, + "isVjp": False, + }, + "workgroupSize": [32, 1, 1], + "dispatchWorkgroupCount": [2, 1, 1], + "specializationConstants": {}, + }, + "forward-float32-axis-4099": { + "entryPoint": "rms_loopedfloat32", + "artifactId": ( + "sha256:b81c2043b10bde966cb6f4dbfa198d2b93a3e456f3026030b69557c4a8983729" + ), + "dispatchVariantId": ( + "sha256:392aee49734fcc2fd3ff4fd232a49f1a97a2bb79f3c22cc80eeabeb1f1ca1959" + ), + "inputs": { + "axisSize": 4099, + "dtype": "float32", + "hasW": False, + "isVjp": False, + }, + "workgroupSize": [1024, 1, 1], + "dispatchWorkgroupCount": [1, 1, 1], + "specializationConstants": {}, + }, + "vjp-float32-axis-32-has-w-false": { + "entryPoint": "vjp_rmsfloat32", + "artifactId": ( + "sha256:ef832e1ceb8c864a13aee3460d23658f4fffba18db1b800461628ba6ebe38e0a" + ), + "dispatchVariantId": ( + "sha256:1a25dca51070c3b6fc96f162e6c152049d388833ba02b1f3d10cc1928c5661c4" + ), + "inputs": { + "axisSize": 32, + "dtype": "float32", + "hasW": False, + "isVjp": True, + }, + "workgroupSize": [32, 1, 1], + "dispatchWorkgroupCount": [800, 1, 1], + "specializationConstants": {"20": False}, + }, + "vjp-float32-axis-32-has-w-true": { + "entryPoint": "vjp_rmsfloat32", + "artifactId": ( + "sha256:a9be06b43a6156fb9ee1f9a6955d03d6bda0940c2a8223b58f564c2d12bd0cd0" + ), + "dispatchVariantId": ( + "sha256:26177a77e484a56b7b2572516e8e1360714c88b668881237a9cf499129e34f35" + ), + "inputs": { + "axisSize": 32, + "dtype": "float32", + "hasW": True, + "isVjp": True, + }, + "workgroupSize": [32, 1, 1], + "dispatchWorkgroupCount": [800, 1, 1], + "specializationConstants": {"20": True}, + }, + "vjp-float32-axis-256-has-w-false": { + "entryPoint": "vjp_rmsfloat32", + "artifactId": ( + "sha256:4455fc9204f826fc5d0d7f016bcaf970b75ea31d5f1598d13ceca6a2baa369e7" + ), + "dispatchVariantId": ( + "sha256:f26ffb357ecfabef6216ff45661f298af4f164a0494da838ca579ea0812e73d2" + ), + "inputs": { + "axisSize": 256, + "dtype": "float32", + "hasW": False, + "isVjp": True, + }, + "workgroupSize": [64, 1, 1], + "dispatchWorkgroupCount": [800, 1, 1], + "specializationConstants": {"20": False}, + }, + "vjp-float32-axis-256-has-w-true": { + "entryPoint": "vjp_rmsfloat32", + "artifactId": ( + "sha256:0944044e2f050bedde1d05e1ae5648628e7144c04752dd49b2e2bb7bcd807b7b" + ), + "dispatchVariantId": ( + "sha256:a010324f59769fde9d71cc8968852ae1e6c8b0ddc213f035523c2c5f2e12d413" + ), + "inputs": { + "axisSize": 256, + "dtype": "float32", + "hasW": True, + "isVjp": True, + }, + "workgroupSize": [64, 1, 1], + "dispatchWorkgroupCount": [800, 1, 1], + "specializationConstants": {"20": True}, + }, + "vjp-float32-axis-8192-has-w-false": { + "entryPoint": "vjp_rms_loopedfloat32", + "artifactId": ( + "sha256:3bd55b546fc00ddf8412f092da4793c0272eec0ad7c130065ad7c1677f60cdce" + ), + "dispatchVariantId": ( + "sha256:fa4d57c473cd5799c46a982e3ea339debc06de21fb6a255adc11bbb546f9329b" + ), + "inputs": { + "axisSize": 8192, + "dtype": "float32", + "hasW": False, + "isVjp": True, + }, + "workgroupSize": [1024, 1, 1], + "dispatchWorkgroupCount": [4, 1, 1], + "specializationConstants": {"20": False}, + }, + "vjp-float32-axis-8192-has-w-true": { + "entryPoint": "vjp_rms_loopedfloat32", + "artifactId": ( + "sha256:345d524ffec14682b6d0325bc97b624b89d83dc257a17ed49bea5e11e24573f3" + ), + "dispatchVariantId": ( + "sha256:8fc64f93a1be95f7acac67d6595c9f9c66c01b87cf952f29052835c20d4d765b" + ), + "inputs": { + "axisSize": 8192, + "dtype": "float32", + "hasW": True, + "isVjp": True, + }, + "workgroupSize": [1024, 1, 1], + "dispatchWorkgroupCount": [4, 1, 1], + "specializationConstants": {"20": True}, + }, +} MLX_ROPE_SOURCE = "mlx/backend/metal/kernels/rope.metal" MLX_SCALED_DOT_PRODUCT_ATTENTION_SOURCE = ( "mlx/backend/metal/kernels/scaled_dot_product_attention.metal" @@ -145,7 +373,7 @@ MLX_DIRECTX_DYNAMIC_WORKGROUP_FRONTIER_SOURCES = tuple( source for source in MLX_DYNAMIC_WORKGROUP_FRONTIER_SOURCES - if source != MLX_LAYER_NORM_SOURCE + if source not in {MLX_LAYER_NORM_SOURCE, MLX_RMS_NORM_SOURCE} ) MLX_DIRECTX_TRANSLATED_FRONTIER_SOURCES = tuple( source @@ -269,27 +497,33 @@ MLX_TERNARY_SOURCE: 212, } # Aggregate artifacts retain every source entry when one workgroup contract applies. -# LayerNorm contributes two bounded, entry-scoped host-dispatch artifacts. +# LayerNorm and RMSNorm contribute bounded, entry-scoped host-dispatch artifacts. MLX_DIRECTX_TOOLCHAIN_FRONTIER_SOURCES = tuple( source for source in MLX_DIRECTX_VULKAN_FRONTIER_SOURCES if source in MLX_DIRECTX_TRANSLATED_FRONTIER_SOURCES - or source == MLX_LAYER_NORM_SOURCE + or source in {MLX_LAYER_NORM_SOURCE, MLX_RMS_NORM_SOURCE} ) MLX_DIRECTX_TOOLCHAIN_ENTRY_POINT_COUNTS = { source: ( len(MLX_LAYER_NORM_DISPATCH_VARIANTS) if source == MLX_LAYER_NORM_SOURCE - else MLX_DIRECTX_FRONTIER_ENTRY_POINT_COUNTS[source] + else ( + len(MLX_RMS_NORM_DISPATCH_VARIANTS) + if source == MLX_RMS_NORM_SOURCE + else MLX_DIRECTX_FRONTIER_ENTRY_POINT_COUNTS[source] + ) ) for source in MLX_DIRECTX_TOOLCHAIN_FRONTIER_SOURCES } MLX_DIRECTX_TOOLCHAIN_ENTRY_POINT_COUNT = sum( MLX_DIRECTX_TOOLCHAIN_ENTRY_POINT_COUNTS.values() ) -MLX_DIRECTX_TOOLCHAIN_ARTIFACT_COUNT = len( - MLX_DIRECTX_TRANSLATED_FRONTIER_SOURCES -) + len(MLX_LAYER_NORM_DISPATCH_VARIANTS) +MLX_DIRECTX_TOOLCHAIN_ARTIFACT_COUNT = ( + len(MLX_DIRECTX_TRANSLATED_FRONTIER_SOURCES) + + len(MLX_LAYER_NORM_DISPATCH_VARIANTS) + + len(MLX_RMS_NORM_DISPATCH_VARIANTS) +) MLX_DIRECTX_TOOLCHAIN_WARNING_CONTRACTS: tuple[dict[str, Any], ...] = () DIRECTX_TOOLCHAIN_WARNING_TRACKED_ISSUES: tuple[str, ...] = () MLX_DIRECTX_TOOLCHAIN_WARNING_EVIDENCE = { @@ -1470,6 +1704,74 @@ def _prepare_layer_norm_dispatch_contract( } +def _prepare_rms_norm_dispatch_contract( + mlx_root: Path, + work_dir: Path, +) -> dict[str, Any]: + source_path = MLX_RMS_NORM_DISPATCH_CONTRACT_SOURCE + _require( + source_path.is_file(), f"RMSNorm dispatch contract is missing: {source_path}" + ) + _require( + _normalized_text_sha256(source_path) == MLX_RMS_NORM_DISPATCH_NORMALIZED_SHA256, + "RMSNorm dispatch contract normalized content changed", + ) + rms_norm_source = mlx_root / MLX_RMS_NORM_SOURCE + _require( + rms_norm_source.is_file() and _sha256(rms_norm_source) == MLX_RMS_NORM_SHA256, + "pinned RMSNorm source identity changed", + ) + + manifest = load_dispatch_contract(source_path) + content_identity = manifest.content_identity.to_json() + expected_identity = { + "algorithm": "sha256", + "value": MLX_RMS_NORM_DISPATCH_CONTENT_IDENTITY.removeprefix("sha256:"), + } + _require( + content_identity == expected_identity, + "RMSNorm dispatch contract identity changed", + ) + evaluation = manifest.evaluate().to_json() + variants = evaluation.get("variants") + variants_by_workload = { + variant.get("workload", {}).get("id"): variant + for variant in variants or [] + if isinstance(variant, Mapping) and isinstance(variant.get("workload"), Mapping) + } + _require( + isinstance(variants, list) + and len(variants) == len(MLX_RMS_NORM_DISPATCH_VARIANTS) + and set(variants_by_workload) == set(MLX_RMS_NORM_DISPATCH_VARIANTS), + "RMSNorm dispatch contract variant set changed", + ) + for workload_id, expected in MLX_RMS_NORM_DISPATCH_VARIANTS.items(): + variant = variants_by_workload[workload_id] + _require( + variant.get("artifactId") == expected["artifactId"] + and variant.get("variantId") == expected["dispatchVariantId"] + and variant.get("source") == MLX_RMS_NORM_SOURCE + and variant.get("entryPoint") == expected["entryPoint"] + and variant.get("workload", {}).get("inputs") == expected["inputs"] + and variant.get("workgroupSize") == expected["workgroupSize"] + and variant.get("subgroupWidth") == 32 + and variant.get("specializationConstants") + == expected["specializationConstants"] + and variant.get("dispatch", {}).get("workgroupCount") + == expected["dispatchWorkgroupCount"], + f"RMSNorm dispatch contract changed for {workload_id}", + ) + + destination = work_dir / "contracts" / "rms_norm.dispatch.json" + destination.parent.mkdir(parents=True, exist_ok=True) + shutil.copyfile(source_path, destination) + return { + "path": _relpath(destination, mlx_root), + "contentIdentity": content_identity, + "variantCount": len(variants), + } + + def _write_reference_accessor_project_config(path: Path, output_dir: str) -> None: target_list = ", ".join(json.dumps(target) for target in REFERENCE_ACCESSOR_TARGETS) lines = [ @@ -3805,6 +4107,361 @@ def _require_layer_norm_dispatch_frontier_report( return artifacts_by_entry, evidence, toolchain_runs +def _require_rms_norm_dispatch_frontier_report( + mlx_root: Path, + output_dir: Path, + payload: Mapping[str, Any], + *, + contract: Mapping[str, Any], + validated: bool, +) -> tuple[dict[str, Mapping[str, Any]], dict[str, Any], list[Mapping[str, Any]]]: + project = payload.get("project") + contract_records = ( + project.get("dispatchContracts") if isinstance(project, Mapping) else None + ) + artifact_count = len(MLX_RMS_NORM_DISPATCH_VARIANTS) + _require( + isinstance(project, Mapping) + and project.get("includePatterns") == [MLX_RMS_NORM_SOURCE] + and project.get("targets") == ["directx"] + and project.get("dispatchContractFiles") == [contract["path"]] + and project.get("dispatchContractCount") == 1 + and project.get("dispatchVariantCount") == artifact_count + and isinstance(contract_records, list) + and len(contract_records) == 1, + "RMSNorm dispatch frontier project metadata changed", + ) + contract_record = contract_records[0] + expected_manifest_path = str((mlx_root / str(contract["path"])).resolve()) + evaluation = ( + contract_record.get("evaluation") + if isinstance(contract_record, Mapping) + else None + ) + evaluated_variants = ( + evaluation.get("variants") if isinstance(evaluation, Mapping) else None + ) + _require( + isinstance(contract_record, Mapping) + and contract_record.get("path") == contract["path"] + and contract_record.get("schemaVersion") == 1 + and contract_record.get("contentIdentity") == contract["contentIdentity"] + and contract_record.get("manifest", {}).get("provenance", {}).get("commit") + == MLX_COMMIT + and isinstance(evaluation, Mapping) + and evaluation.get("manifestSource") == expected_manifest_path + and evaluation.get("variantCount") == artifact_count + and isinstance(evaluated_variants, list) + and len(evaluated_variants) == artifact_count, + "RMSNorm dispatch frontier did not retain replayable manifest provenance", + ) + evaluation_by_workload = { + variant.get("workload", {}).get("id"): variant + for variant in evaluated_variants + if isinstance(variant, Mapping) and isinstance(variant.get("workload"), Mapping) + } + _require( + set(evaluation_by_workload) == set(MLX_RMS_NORM_DISPATCH_VARIANTS), + "RMSNorm dispatch evaluation workload set changed", + ) + + summary = payload.get("summary") + _require( + isinstance(summary, Mapping) + and summary.get("unitCount") == 1 + and summary.get("artifactCount") == artifact_count + and summary.get("translatedCount") == artifact_count + and summary.get("failedCount") == 0 + and summary.get("diagnosticCounts") == {"error": 0, "note": 0, "warning": 0} + and summary.get("artifactsByTarget", {}).get("directx") + == { + "artifactCount": artifact_count, + "translatedCount": artifact_count, + "failedCount": 0, + } + and payload.get("diagnostics") == [], + "RMSNorm dispatch frontier accounting changed", + ) + units = payload.get("units") + _require( + isinstance(units, list) + and len(units) == 1 + and units[0].get("path") == MLX_RMS_NORM_SOURCE + and units[0].get("sourceBackend") == "metal" + and units[0].get("sourceHash") + == {"algorithm": "sha256", "value": MLX_RMS_NORM_SHA256}, + "RMSNorm dispatch frontier source identity changed", + ) + + dispatch_plan = project.get("dispatchArtifactPlan") + planned_artifacts = ( + dispatch_plan.get("artifacts") if isinstance(dispatch_plan, Mapping) else None + ) + _require( + isinstance(dispatch_plan, Mapping) + and dispatch_plan.get("kind") == "crosstl-dispatch-artifact-plan" + and dispatch_plan.get("schemaVersion") == 1 + and dispatch_plan.get("sourceUnitCount") == 1 + and dispatch_plan.get("artifactCount") == artifact_count + and dispatch_plan.get("dispatchVariantCount") == artifact_count + and isinstance(planned_artifacts, list) + and len(planned_artifacts) == artifact_count, + "RMSNorm dispatch artifact plan changed", + ) + plan_by_artifact_id = { + record.get("artifactId"): record + for record in planned_artifacts + if isinstance(record, Mapping) + } + expected_artifact_ids = { + expected["artifactId"] for expected in MLX_RMS_NORM_DISPATCH_VARIANTS.values() + } + _require( + set(plan_by_artifact_id) == expected_artifact_ids, + "RMSNorm dispatch artifact plan identity set changed", + ) + + artifacts = payload.get("artifacts") + _require( + isinstance(artifacts, list) and len(artifacts) == artifact_count, + "RMSNorm dispatch artifact records are incomplete", + ) + artifacts_by_id = { + artifact.get("dispatchArtifact", {}).get("artifactId"): artifact + for artifact in artifacts + if isinstance(artifact, Mapping) + and isinstance(artifact.get("dispatchArtifact"), Mapping) + } + _require( + set(artifacts_by_id) == expected_artifact_ids, + "RMSNorm dispatch artifact identity set changed", + ) + + output_root = output_dir.resolve() + artifacts_by_workload: dict[str, Mapping[str, Any]] = {} + generated_evidence: dict[str, Any] = {} + for workload_id, expected in MLX_RMS_NORM_DISPATCH_VARIANTS.items(): + artifact_id = expected["artifactId"] + artifact = artifacts_by_id[artifact_id] + plan = plan_by_artifact_id[artifact_id] + evaluated = evaluation_by_workload[workload_id] + entry = artifact.get("entryPoint") + entry_point = entry.get("source") if isinstance(entry, Mapping) else None + _require( + evaluated.get("artifactId") == artifact_id + and evaluated.get("variantId") == expected["dispatchVariantId"] + and evaluated.get("entryPoint") == expected["entryPoint"] + and evaluated.get("workload", {}).get("inputs") == expected["inputs"] + and evaluated.get("workgroupSize") == expected["workgroupSize"] + and evaluated.get("subgroupWidth") == 32 + and evaluated.get("specializationConstants") + == expected["specializationConstants"] + and evaluated.get("dispatch", {}).get("workgroupCount") + == expected["dispatchWorkgroupCount"], + f"RMSNorm evaluated dispatch changed for {workload_id}", + ) + variant_name = "dispatch-" + artifact_id.removeprefix("sha256:") + _require( + artifact.get("source") == MLX_RMS_NORM_SOURCE + and artifact.get("sourceBackend") == "metal" + and artifact.get("sourceHash") + == {"algorithm": "sha256", "value": MLX_RMS_NORM_SHA256} + and artifact.get("target") == "directx" + and artifact.get("status") == "translated" + and artifact.get("variant") == variant_name + and entry_point == expected["entryPoint"] + and entry.get("target") == "CSMain" + and entry.get("stage") == "compute" + and artifact.get("dispatchArtifact") == plan + and plan.get("artifactId") == artifact_id + and plan.get("dispatchVariantIds") == [expected["dispatchVariantId"]] + and plan.get("manifestContentIdentities") + == [MLX_RMS_NORM_DISPATCH_CONTENT_IDENTITY] + and plan.get("source") == MLX_RMS_NORM_SOURCE + and plan.get("entryPoint") == expected["entryPoint"] + and plan.get("workgroupSize") == expected["workgroupSize"] + and plan.get("subgroupWidth") == 32 + and plan.get("specializationConstants") + == expected["specializationConstants"], + f"RMSNorm dispatch artifact contract changed for {workload_id}", + ) + + execution = artifact.get("execution") + execution_entries = ( + execution.get("entryPoints") if isinstance(execution, Mapping) else None + ) + _require( + isinstance(execution, Mapping) + and execution.get("sourceEntryPoints") == [expected["entryPoint"]] + and execution.get("provenance", {}).get("kind") == "host-dispatch-contract" + and execution.get("provenance", {}).get("artifactId") == artifact_id + and execution.get("subgroupWidthProvenance", {}).get("kind") + == "host-dispatch-contract" + and execution.get("subgroupWidthEnforcement") + == { + "mechanism": "hlsl-wave-size-attribute", + "minimumShaderModel": "6.6", + "entryProfiles": [{"entryPoint": "CSMain", "profile": "cs_6_6"}], + } + and isinstance(execution_entries, list) + and len(execution_entries) == 1 + and execution_entries[0].get("sourceEntryPoint") == expected["entryPoint"] + and execution_entries[0].get("materializedEntryPoint") + == expected["entryPoint"] + and execution_entries[0].get("targetEntryPoint") == "CSMain" + and execution_entries[0].get("workgroupSize") == expected["workgroupSize"] + and execution_entries[0].get("subgroupWidth") == 32, + f"RMSNorm execution metadata changed for {workload_id}", + ) + + constants = artifact.get("specializationConstants") or [] + constants_by_id = { + str(record.get("id")): record + for record in constants + if isinstance(record, Mapping) + } + _require( + set(constants_by_id) == set(expected["specializationConstants"]), + f"RMSNorm specialization inputs changed for {workload_id}", + ) + for constant_id, value in expected["specializationConstants"].items(): + record = constants_by_id[constant_id] + _require( + record.get("concreteValue") is value + and record.get("deferred") is False + and record.get("valueProvenance", {}).get("kind") + == "host-dispatch-contract" + and record.get("valueProvenance", {}).get("artifactId") == artifact_id, + f"RMSNorm specialization provenance changed for {workload_id}", + ) + + materialization = artifact.get("templateMaterialization") + specializations = ( + materialization.get("specializations") + if isinstance(materialization, Mapping) + else None + ) + _require( + isinstance(materialization, Mapping) + and materialization.get("status") == "materialized" + and materialization.get("specializationCount") == 1 + and isinstance(specializations, list) + and len(specializations) == 1 + and specializations[0].get("hostName") == expected["entryPoint"] + and materialization.get("unsupported") == [], + f"RMSNorm materialization changed for {workload_id}", + ) + + artifact_path_value = artifact.get("path") + _require( + isinstance(artifact_path_value, str) and artifact_path_value, + f"RMSNorm artifact path is missing for {workload_id}", + ) + artifact_path = (mlx_root / artifact_path_value).resolve() + _require( + _is_relative_to(artifact_path, output_root) and artifact_path.is_file(), + f"RMSNorm artifact is missing or escaped output for {workload_id}", + ) + generated = artifact_path.read_text(encoding="utf-8") + generated_hash = _sha256(artifact_path) + normalized_hash = _normalized_text_sha256(artifact_path) + _require( + artifact.get("generatedHash") + == {"algorithm": "sha256", "value": generated_hash} + and artifact.get("generatedSizeBytes") == artifact_path.stat().st_size + and len(re.findall(r"\bvoid\s+CSMain\s*\(", generated)) == 1 + and re.search(r"\[\s*WaveSize\s*\(\s*32\s*\)\s*\]", generated) is not None + and re.search( + r"\[\s*numthreads\s*\(\s*{}\s*,\s*1\s*,\s*1\s*\)\s*\]".format( + expected["workgroupSize"][0] + ), + generated, + ) + is not None, + f"RMSNorm generated HLSL contract changed for {workload_id}", + ) + if expected["specializationConstants"]: + has_w = str(expected["specializationConstants"]["20"]).lower() + _require( + re.search( + rf"\bstatic\s+const\s+bool\s+has_w\s*=\s*{has_w}\s*;", + generated, + ) + is not None, + f"RMSNorm VJP artifact did not materialize has_w={has_w}", + ) + else: + _require( + re.search(r"\bhas_w\b", generated) is None, + "RMSNorm forward artifact retained an unreachable function constant", + ) + generated_evidence[workload_id] = { + "entryPoint": expected["entryPoint"], + "artifactId": artifact_id, + "dispatchVariantId": expected["dispatchVariantId"], + "inputs": dict(expected["inputs"]), + "workgroupSize": list(expected["workgroupSize"]), + "dispatchWorkgroupCount": list(expected["dispatchWorkgroupCount"]), + "subgroupWidth": 32, + "specializationConstants": dict(expected["specializationConstants"]), + "generatedHlsl": { + "normalizedSha256": normalized_hash, + "contentSha256": generated_hash, + "sizeBytes": artifact_path.stat().st_size, + }, + } + artifacts_by_workload[workload_id] = artifact + + toolchain_runs: list[Mapping[str, Any]] = [] + if validated: + validation = payload.get("validation") + runs = ( + validation.get("toolchainRuns") if isinstance(validation, Mapping) else None + ) + _require( + isinstance(validation, Mapping) + and isinstance(validation.get("summary"), Mapping) + and validation["summary"].get("failedCount") == 0 + and isinstance(runs, list), + "RMSNorm dispatch DXC validation changed", + ) + toolchain_runs = [ + run + for run in runs + if isinstance(run, Mapping) and run.get("target") == "directx" + ] + artifact_paths = {artifact["path"] for artifact in artifacts_by_id.values()} + _require( + len(toolchain_runs) == artifact_count + and all(run.get("status") == "ok" for run in toolchain_runs) + and {run.get("path") for run in toolchain_runs} == artifact_paths, + "RMSNorm dispatch DXC did not validate every artifact", + ) + + evidence = { + "status": "translated-dxc-validated" if validated else "translated", + "source": MLX_RMS_NORM_SOURCE, + "sourceSha256": MLX_RMS_NORM_SHA256, + "target": "directx", + "testSources": [ + "python/tests/test_fast.py::test_rms_norm", + "python/tests/test_fast.py::test_rms_norm_grad", + ], + "dispatchContract": { + "path": contract["path"], + "contentIdentity": MLX_RMS_NORM_DISPATCH_CONTENT_IDENTITY, + "variantCount": artifact_count, + "resolvedIssue": MLX_HOST_DISPATCH_IMPORT_RESOLVED_ISSUE, + }, + "artifactCount": artifact_count, + "variants": generated_evidence, + "dxcValidatedArtifactCount": len(toolchain_runs), + "runtimeExecutionAttempted": False, + "numericalParityClaimed": False, + } + return artifacts_by_workload, evidence, toolchain_runs + + def _require_directx_bfloat16_lowering_evidence( artifacts_by_source: Mapping[str, Mapping[str, Any]], ) -> dict[str, dict[str, Any]]: @@ -4194,6 +4851,37 @@ def _translate_directx_frontier( validated=False, ) + rms_norm_contract = _prepare_rms_norm_dispatch_contract(mlx_root, work_dir) + rms_norm_output_dir = work_dir / "out-directx-rms-norm-dispatch-frontier" + ( + _rms_norm_result, + rms_norm_payload, + _rms_norm_config, + rms_norm_report_path, + ) = _run_frontier_project( + mlx_root=mlx_root, + config_dir=config_dir, + report_dir=report_dir, + log_dir=log_dir, + python=python, + command_name="directx-rms-norm-dispatch-frontier", + target="directx", + sources=(MLX_RMS_NORM_SOURCE,), + output_dir=rms_norm_output_dir, + dispatch_contracts=(str(rms_norm_contract["path"]),), + ) + ( + rms_norm_artifacts, + rms_norm_evidence, + _rms_norm_runs, + ) = _require_rms_norm_dispatch_frontier_report( + mlx_root, + rms_norm_output_dir, + rms_norm_payload, + contract=rms_norm_contract, + validated=False, + ) + blocked_output_dir = work_dir / "out-directx-workgroup-frontier" blocked_result, blocked_payload, _blocked_config, blocked_report_path = ( _run_frontier_project( @@ -4294,10 +4982,41 @@ def _translate_directx_frontier( contract=layer_norm_contract, validated=True, ) - directx_runs = [*aggregate_runs, *layer_norm_runs] + rms_norm_toolchain_output = work_dir / "out-directx-rms-norm-dispatch-toolchain" + ( + _rms_norm_toolchain_result, + rms_norm_toolchain_payload, + _rms_norm_toolchain_config, + _rms_norm_toolchain_report, + ) = _run_frontier_project( + mlx_root=mlx_root, + config_dir=config_dir, + report_dir=report_dir, + log_dir=log_dir, + python=python, + command_name="validate-directx-rms-norm-dispatch-toolchain", + target="directx", + sources=(MLX_RMS_NORM_SOURCE,), + output_dir=rms_norm_toolchain_output, + run_toolchains=True, + dispatch_contracts=(str(rms_norm_contract["path"]),), + ) + ( + rms_norm_artifacts, + rms_norm_evidence, + rms_norm_runs, + ) = _require_rms_norm_dispatch_frontier_report( + mlx_root, + rms_norm_toolchain_output, + rms_norm_toolchain_payload, + contract=rms_norm_contract, + validated=True, + ) + directx_runs = [*aggregate_runs, *layer_norm_runs, *rms_norm_runs] artifact_paths = { *(artifact["path"] for artifact in artifacts.values()), *(artifact["path"] for artifact in layer_norm_artifacts.values()), + *(artifact["path"] for artifact in rms_norm_artifacts.values()), } validated_paths = { run.get("path") for run in directx_runs if run.get("status") == "ok" @@ -4314,6 +5033,10 @@ def _translate_directx_frontier( artifact["path"]: entry_point for entry_point, artifact in layer_norm_artifacts.items() } + rms_norm_workload_by_path = { + artifact["path"]: workload_id + for workload_id, artifact in rms_norm_artifacts.items() + } for run in directx_runs: if run.get("status") != "ok": continue @@ -4322,11 +5045,12 @@ def _translate_directx_frontier( source in directx_entry_points_by_source, f"DirectX toolchain validation reported an unexpected source: {source}", ) - entry_point = ( - layer_norm_entry_by_path.get(run.get("path")) - if source == MLX_LAYER_NORM_SOURCE - else _directx_toolchain_entry_point(run) - ) + if source == MLX_LAYER_NORM_SOURCE: + entry_point = layer_norm_entry_by_path.get(run.get("path")) + elif source == MLX_RMS_NORM_SOURCE: + entry_point = rms_norm_workload_by_path.get(run.get("path")) + else: + entry_point = _directx_toolchain_entry_point(run) _require( entry_point is not None, "DirectX toolchain validation did not record a compute entry command", @@ -4366,12 +5090,14 @@ def _translate_directx_frontier( "scope": "target-split-frontier", "report": _relpath(report_path, mlx_root), "layerNormDispatchReport": _relpath(layer_norm_report_path, mlx_root), + "rmsNormDispatchReport": _relpath(rms_norm_report_path, mlx_root), "workgroupBlockedReport": _relpath(blocked_report_path, mlx_root), "sources": list(MLX_DIRECTX_VULKAN_FRONTIER_SOURCES), "unitCount": len(MLX_DIRECTX_VULKAN_FRONTIER_SOURCES), "artifactCount": ( len(MLX_DIRECTX_TRANSLATED_FRONTIER_SOURCES) + len(MLX_LAYER_NORM_DISPATCH_VARIANTS) + + len(MLX_RMS_NORM_DISPATCH_VARIANTS) + len(MLX_DIRECTX_DYNAMIC_WORKGROUP_FRONTIER_SOURCES) ), "translatedSources": list(MLX_DIRECTX_TOOLCHAIN_FRONTIER_SOURCES), @@ -4409,6 +5135,7 @@ def _translate_directx_frontier( "native16BitArithmeticEvidence": MLX_DIRECTX_NATIVE_16_BIT_ARITHMETIC_EVIDENCE, "bfloat16LoweringEvidence": bfloat16_lowering_evidence, "layerNormDispatchEvidence": layer_norm_evidence, + "rmsNormDispatchEvidence": rms_norm_evidence, "dynamicWorkgroupDispatchEvidence": dispatch_evidence, "semanticReadinessStatus": "not-established", "trackedIssues": [ diff --git a/tests/test_ci_workflows.py b/tests/test_ci_workflows.py index ddb3cdf15..169adefb5 100644 --- a/tests/test_ci_workflows.py +++ b/tests/test_ci_workflows.py @@ -2424,6 +2424,7 @@ def test_mlx_project_porting_workflow_runs_tracked_porting_harness(): mlx_porting ) assert '"tests/test_mlx_porting_harness.py"' in mlx_porting + assert '"tests/test_mlx_rms_norm_dispatch_contract_fixture.py"' in mlx_porting assert '"tests/test_mlx_quantized_directx_proof.py"' in mlx_porting assert '"tests/test_mlx_quantized_opengl_proof.py"' in mlx_porting assert '"tests/test_translator/test_codegen/test_SPIRV_codegen.py"' in mlx_porting @@ -2450,6 +2451,7 @@ def test_mlx_project_porting_workflow_runs_tracked_porting_harness(): assert "MLX_DIRECTX_DYNAMIC_WORKGROUP_FRONTIER_SOURCES" in mlx_porting assert "MLX_HOST_DISPATCH_IMPORT_RESOLVED_ISSUE" in mlx_porting assert "MLX_LAYER_NORM_DISPATCH_VARIANTS" in mlx_porting + assert "MLX_RMS_NORM_DISPATCH_VARIANTS" in mlx_porting assert 'checks["directx-frontier"]' in mlx_porting assert 'checks["vulkan-frontier"]' in mlx_porting assert 'directx["directxToolchainArtifactCount"]' in mlx_porting @@ -2461,8 +2463,9 @@ def test_mlx_project_porting_workflow_runs_tracked_porting_harness(): assert 'directx["bfloat16LoweringEvidence"]' in mlx_porting assert "DirectX bfloat16 lowering evidence is incomplete" in mlx_porting assert "DirectX workgroup blocker evidence changed" in mlx_porting - assert "expected 94 fail-closed DirectX compute entries" in mlx_porting + assert "expected 82 fail-closed DirectX compute entries" in mlx_porting assert "LayerNorm dispatch frontier evidence is incomplete" in mlx_porting + assert "RMSNorm dispatch frontier evidence is incomplete" in mlx_porting assert "matched-materialized-host-names" in mlx_porting assert "DirectX frontier toolchain must validate every configured" in mlx_porting assert "source artifact and compute entry" in mlx_porting diff --git a/tests/test_mlx_porting_harness.py b/tests/test_mlx_porting_harness.py index 33c08daac..7b4ae7047 100644 --- a/tests/test_mlx_porting_harness.py +++ b/tests/test_mlx_porting_harness.py @@ -189,6 +189,77 @@ def test_layer_norm_dispatch_contract_accepts_crlf_checkout(tmp_path, monkeypatc } +def test_rms_norm_dispatch_contract_preparation_copies_verified_manifest( + tmp_path, monkeypatch +): + module = _load_harness() + mlx_root = tmp_path / "mlx" + source_path = mlx_root / module.MLX_RMS_NORM_SOURCE + source_path.parent.mkdir(parents=True) + source_path.write_text("kernel void rms_norm_fixture() {}\n", encoding="utf-8") + monkeypatch.setattr( + module, + "MLX_RMS_NORM_SHA256", + hashlib.sha256(source_path.read_bytes()).hexdigest(), + ) + + contract = module._prepare_rms_norm_dispatch_contract( + mlx_root, + mlx_root / ".crosstl-mlx-porting", + ) + + copied_path = mlx_root / contract["path"] + assert ( + copied_path.read_bytes() + == module.MLX_RMS_NORM_DISPATCH_CONTRACT_SOURCE.read_bytes() + ) + assert contract == { + "path": ".crosstl-mlx-porting/contracts/rms_norm.dispatch.json", + "contentIdentity": { + "algorithm": "sha256", + "value": module.MLX_RMS_NORM_DISPATCH_CONTENT_IDENTITY.removeprefix( + "sha256:" + ), + }, + "variantCount": len(module.MLX_RMS_NORM_DISPATCH_VARIANTS), + } + + +def test_rms_norm_dispatch_contract_accepts_crlf_checkout(tmp_path, monkeypatch): + module = _load_harness() + mlx_root = tmp_path / "mlx" + source_path = mlx_root / module.MLX_RMS_NORM_SOURCE + source_path.parent.mkdir(parents=True) + source_path.write_text("kernel void rms_norm_fixture() {}\n", encoding="utf-8") + monkeypatch.setattr( + module, + "MLX_RMS_NORM_SHA256", + hashlib.sha256(source_path.read_bytes()).hexdigest(), + ) + contract_source = tmp_path / "rms_norm.dispatch.json" + normalized = module.MLX_RMS_NORM_DISPATCH_CONTRACT_SOURCE.read_text( + encoding="utf-8" + ) + contract_source.write_bytes(normalized.replace("\n", "\r\n").encode("utf-8")) + monkeypatch.setattr( + module, + "MLX_RMS_NORM_DISPATCH_CONTRACT_SOURCE", + contract_source, + ) + + contract = module._prepare_rms_norm_dispatch_contract( + mlx_root, + mlx_root / ".crosstl-mlx-porting", + ) + + copied_path = mlx_root / contract["path"] + assert copied_path.read_bytes() == contract_source.read_bytes() + assert contract["contentIdentity"] == { + "algorithm": "sha256", + "value": module.MLX_RMS_NORM_DISPATCH_CONTENT_IDENTITY.removeprefix("sha256:"), + } + + def _load_rms_norm_fixture_metadata(): return json.loads( ( @@ -2023,15 +2094,15 @@ def test_expected_gaps_tracks_current_frontier_and_runtime_fixture_counts(): frontier = expected_gaps["frontier_status"] assert frontier["sources"] == len(module.MLX_DIRECTX_VULKAN_FRONTIER_SOURCES) - assert frontier["artifacts"] == 23 + assert frontier["artifacts"] == 34 assert frontier["status"] == ( "target-split-with-bounded-dispatch-and-pending-contracts" ) assert frontier["scope"] == "target-split-frontier" - assert frontier["translated_artifacts"] == 18 - assert frontier["failed_artifacts"] == 5 + assert frontier["translated_artifacts"] == 30 + assert frontier["failed_artifacts"] == 4 assert frontier["target_artifacts"] == { - "directx": {"translated": 7, "failed": 5}, + "directx": {"translated": 19, "failed": 4}, "vulkan": {"translated": 11, "failed": 0}, } assert frontier["semantic_readiness_status"] == "not-established" @@ -2237,6 +2308,43 @@ def test_expected_gaps_tracks_current_frontier_and_runtime_fixture_counts(): assert layer_norm["dxc_validated_artifact_count"] == 2 assert layer_norm["runtime_execution_attempted"] is False assert layer_norm["numerical_parity_claimed"] is False + rms_norm = directx["rms_norm_dispatch_frontier"] + assert rms_norm["status"] == "translated-dxc-validated" + assert rms_norm["source"] == module.MLX_RMS_NORM_SOURCE + assert rms_norm["source_sha256"] == module.MLX_RMS_NORM_SHA256 + assert rms_norm["test_sources"] == [ + "python/tests/test_fast.py::test_rms_norm", + "python/tests/test_fast.py::test_rms_norm_grad", + ] + assert rms_norm["dispatch_contract"] == { + "path": "demos/integrations/mlx/contracts/rms_norm.dispatch.json", + "normalized_sha256": module.MLX_RMS_NORM_DISPATCH_NORMALIZED_SHA256, + "content_identity": module.MLX_RMS_NORM_DISPATCH_CONTENT_IDENTITY, + "variant_count": len(module.MLX_RMS_NORM_DISPATCH_VARIANTS), + "resolved_issue": module.MLX_HOST_DISPATCH_IMPORT_RESOLVED_ISSUE, + } + assert rms_norm["artifact_count"] == len(module.MLX_RMS_NORM_DISPATCH_VARIANTS) + assert set(rms_norm["variants"]) == set(module.MLX_RMS_NORM_DISPATCH_VARIANTS) + assert rms_norm["dxc_validated_artifact_count"] == 12 + for workload_id, expected in module.MLX_RMS_NORM_DISPATCH_VARIANTS.items(): + variant = rms_norm["variants"][workload_id] + assert variant["entry_point"] == expected["entryPoint"] + assert variant["artifact_id"] == expected["artifactId"] + assert variant["dispatch_variant_id"] == expected["dispatchVariantId"] + assert variant["inputs"] == expected["inputs"] + assert variant["workgroup_size"] == expected["workgroupSize"] + assert variant["dispatch_workgroup_count"] == ( + expected["dispatchWorkgroupCount"] + ) + assert variant["subgroup_width"] == 32 + assert variant["subgroup_width_enforcement"] == "WaveSize(32)" + assert variant["specialization_constants"] == ( + expected["specializationConstants"] + ) + assert len(variant["generated_hlsl"]["normalized_sha256"]) == 64 + assert variant["generated_hlsl"]["size_bytes"] > 0 + assert rms_norm["runtime_execution_attempted"] is False + assert rms_norm["numerical_parity_claimed"] is False assert directx["native_runtime_executed"] is False assert directx["runtime_parity_claimed"] is False @@ -4482,6 +4590,207 @@ def _write_layer_norm_dispatch_report( return report +def _write_rms_norm_dispatch_report( + module, + mlx_root, + output_dir, + report_path, + *, + contract, + toolchain_runs=(), +): + unit = { + "id": module.MLX_RMS_NORM_SOURCE, + "path": module.MLX_RMS_NORM_SOURCE, + "sourceBackend": "metal", + "sourceHash": { + "algorithm": "sha256", + "value": module.MLX_RMS_NORM_SHA256, + }, + "sourceSizeBytes": 12069, + } + artifacts = [] + planned_artifacts = [] + evaluated_variants = [] + for workload_id, expected in module.MLX_RMS_NORM_DISPATCH_VARIANTS.items(): + entry_point = expected["entryPoint"] + variant = "dispatch-" + expected["artifactId"].removeprefix("sha256:") + generated_path = ( + output_dir + / "directx" + / variant + / "mlx/backend/metal/kernels/rms_norm" + / f"{entry_point}.hlsl" + ) + generated_path.parent.mkdir(parents=True, exist_ok=True) + generated_lines = [ + f"[numthreads({expected['workgroupSize'][0]}, 1, 1)]", + "[WaveSize(32)]", + ] + if expected["specializationConstants"]: + has_w = str(expected["specializationConstants"]["20"]).lower() + generated_lines.insert(0, f"static const bool has_w = {has_w};") + generated_lines.extend(("void CSMain() {", "}", "")) + generated_path.write_text("\n".join(generated_lines), encoding="utf-8") + generated_hash = hashlib.sha256(generated_path.read_bytes()).hexdigest() + plan = { + "artifactId": expected["artifactId"], + "dispatchVariantIds": [expected["dispatchVariantId"]], + "entryPoint": entry_point, + "manifestContentIdentities": [ + module.MLX_RMS_NORM_DISPATCH_CONTENT_IDENTITY + ], + "source": module.MLX_RMS_NORM_SOURCE, + "specializationConstants": dict(expected["specializationConstants"]), + "subgroupWidth": 32, + "variant": variant, + "workgroupSize": list(expected["workgroupSize"]), + } + planned_artifacts.append(plan) + evaluated_variants.append( + { + "artifactId": expected["artifactId"], + "variantId": expected["dispatchVariantId"], + "source": module.MLX_RMS_NORM_SOURCE, + "entryPoint": entry_point, + "workload": { + "id": workload_id, + "inputs": dict(expected["inputs"]), + }, + "workgroupSize": list(expected["workgroupSize"]), + "subgroupWidth": 32, + "specializationConstants": dict(expected["specializationConstants"]), + "dispatch": { + "workgroupCount": list(expected["dispatchWorkgroupCount"]) + }, + } + ) + constants = [ + { + "id": int(constant_id), + "concreteValue": value, + "deferred": False, + "valueProvenance": { + "kind": "host-dispatch-contract", + "artifactId": expected["artifactId"], + }, + } + for constant_id, value in expected["specializationConstants"].items() + ] + artifacts.append( + { + "source": module.MLX_RMS_NORM_SOURCE, + "sourceBackend": "metal", + "sourceHash": unit["sourceHash"], + "sourceSizeBytes": unit["sourceSizeBytes"], + "target": "directx", + "status": "translated", + "variant": variant, + "path": generated_path.relative_to(mlx_root).as_posix(), + "generatedHash": { + "algorithm": "sha256", + "value": generated_hash, + }, + "generatedSizeBytes": generated_path.stat().st_size, + "entryPoint": { + "source": entry_point, + "target": "CSMain", + "stage": "compute", + }, + "dispatchArtifact": plan, + "execution": { + "sourceEntryPoints": [entry_point], + "provenance": { + "kind": "host-dispatch-contract", + "artifactId": expected["artifactId"], + }, + "subgroupWidthProvenance": { + "kind": "host-dispatch-contract", + }, + "subgroupWidthEnforcement": { + "mechanism": "hlsl-wave-size-attribute", + "minimumShaderModel": "6.6", + "entryProfiles": [ + {"entryPoint": "CSMain", "profile": "cs_6_6"} + ], + }, + "entryPoints": [ + { + "sourceEntryPoint": entry_point, + "materializedEntryPoint": entry_point, + "targetEntryPoint": "CSMain", + "workgroupSize": list(expected["workgroupSize"]), + "subgroupWidth": 32, + } + ], + }, + "specializationConstants": constants, + "templateMaterialization": { + "status": "materialized", + "specializationCount": 1, + "specializations": [{"hostName": entry_point}], + "unsupported": [], + }, + } + ) + + artifact_count = len(artifacts) + report = { + "project": { + "includePatterns": [module.MLX_RMS_NORM_SOURCE], + "targets": ["directx"], + "dispatchContractFiles": [contract["path"]], + "dispatchContractCount": 1, + "dispatchVariantCount": artifact_count, + "dispatchContracts": [ + { + "path": contract["path"], + "schemaVersion": 1, + "contentIdentity": contract["contentIdentity"], + "manifest": {"provenance": {"commit": module.MLX_COMMIT}}, + "evaluation": { + "manifestSource": str((mlx_root / contract["path"]).resolve()), + "variantCount": artifact_count, + "variants": evaluated_variants, + }, + } + ], + "dispatchArtifactPlan": { + "kind": "crosstl-dispatch-artifact-plan", + "schemaVersion": 1, + "sourceUnitCount": 1, + "artifactCount": artifact_count, + "dispatchVariantCount": artifact_count, + "artifacts": planned_artifacts, + }, + }, + "summary": { + "unitCount": 1, + "artifactCount": artifact_count, + "translatedCount": artifact_count, + "failedCount": 0, + "diagnosticCounts": {"error": 0, "note": 0, "warning": 0}, + "artifactsByTarget": { + "directx": { + "artifactCount": artifact_count, + "translatedCount": artifact_count, + "failedCount": 0, + } + }, + }, + "units": [unit], + "artifacts": artifacts, + "diagnostics": [], + "validation": { + "summary": {"failedCount": 0}, + "toolchainRuns": list(toolchain_runs), + }, + } + report_path.parent.mkdir(parents=True, exist_ok=True) + report_path.write_text(json.dumps(report), encoding="utf-8") + return report + + def _dynamic_workgroup_report( module, mlx_root, @@ -7626,6 +7935,21 @@ def test_reduced_frontier_requires_all_directx_entries_per_artifact( "_prepare_layer_norm_dispatch_contract", lambda *_args: layer_norm_contract, ) + rms_norm_contract = { + "path": ".crosstl-mlx-porting/contracts/rms_norm.dispatch.json", + "contentIdentity": { + "algorithm": "sha256", + "value": module.MLX_RMS_NORM_DISPATCH_CONTENT_IDENTITY.removeprefix( + "sha256:" + ), + }, + "variantCount": len(module.MLX_RMS_NORM_DISPATCH_VARIANTS), + } + monkeypatch.setattr( + module, + "_prepare_rms_norm_dispatch_contract", + lambda *_args: rms_norm_contract, + ) commands = [] def warning_stderr(source, relative_path): @@ -7702,6 +8026,47 @@ def fake_run_command(name, command, *, log_dir, check=True, timeout_seconds=None contract=layer_norm_contract, toolchain_runs=toolchain_runs, ) + elif "rms-norm-dispatch" in name: + is_toolchain = name.startswith("validate-") + output_dir = work_dir / ( + "out-directx-rms-norm-dispatch-toolchain" + if is_toolchain + else "out-directx-rms-norm-dispatch-frontier" + ) + report = _write_rms_norm_dispatch_report( + module, + mlx_root, + output_dir, + report_dir / f"{name}.json", + contract=rms_norm_contract, + ) + if is_toolchain: + toolchain_runs = [ + { + "source": module.MLX_RMS_NORM_SOURCE, + "target": "directx", + "path": artifact["path"], + "command": [ + "dxc", + "-T", + "cs_6_6", + "-E", + "CSMain", + artifact["path"], + ], + "status": "ok", + "stderr": "", + } + for artifact in report["artifacts"] + ] + _write_rms_norm_dispatch_report( + module, + mlx_root, + output_dir, + report_dir / f"{name}.json", + contract=rms_norm_contract, + toolchain_runs=toolchain_runs, + ) else: is_toolchain = name == "validate-directx-frontier-toolchain" output_dir = ( @@ -7809,6 +8174,12 @@ def fake_run_command(name, command, *, log_dir, check=True, timeout_seconds=None assert set(result["layerNormDispatchEvidence"]["variants"]) == set( module.MLX_LAYER_NORM_DISPATCH_VARIANTS ) + assert result["rmsNormDispatchEvidence"]["status"] == ("translated-dxc-validated") + assert result["rmsNormDispatchEvidence"]["artifactCount"] == 12 + assert result["rmsNormDispatchEvidence"]["dxcValidatedArtifactCount"] == 12 + assert set(result["rmsNormDispatchEvidence"]["variants"]) == set( + module.MLX_RMS_NORM_DISPATCH_VARIANTS + ) assert result["workgroupBlockedSources"] == list( module.MLX_DIRECTX_DYNAMIC_WORKGROUP_FRONTIER_SOURCES ) @@ -7819,12 +8190,15 @@ def fake_run_command(name, command, *, log_dir, check=True, timeout_seconds=None assert [name for name, _command in commands] == [ "directx-frontier", "directx-layer-norm-dispatch-frontier", + "directx-rms-norm-dispatch-frontier", "directx-workgroup-frontier", "validate-directx-frontier-toolchain", "validate-directx-layer-norm-dispatch-toolchain", + "validate-directx-rms-norm-dispatch-toolchain", ] - assert "--run-toolchains" in commands[3][1] assert "--run-toolchains" in commands[4][1] + assert "--run-toolchains" in commands[5][1] + assert "--run-toolchains" in commands[6][1] toolchain_config = ( config_dir / "validate-directx-frontier-toolchain.toml" ).read_text(encoding="utf-8") @@ -7840,6 +8214,11 @@ def fake_run_command(name, command, *, log_dir, check=True, timeout_seconds=None ).read_text(encoding="utf-8") assert module.MLX_LAYER_NORM_SOURCE in layer_norm_config assert layer_norm_contract["path"] in layer_norm_config + rms_norm_config = ( + config_dir / "validate-directx-rms-norm-dispatch-toolchain.toml" + ).read_text(encoding="utf-8") + assert module.MLX_RMS_NORM_SOURCE in rms_norm_config + assert rms_norm_contract["path"] in rms_norm_config def test_directx_toolchain_warning_contract_rejects_new_warning(): @@ -7969,6 +8348,7 @@ def test_directx_toolchain_frontier_matches_pinned_dxc_inventory(): module.MLX_BINARY_TWO_SOURCE, module.MLX_LAYER_NORM_SOURCE, module.MLX_RANDOM_SOURCE, + module.MLX_RMS_NORM_SOURCE, module.MLX_ROPE_SOURCE, module.MLX_TERNARY_SOURCE, ) @@ -7983,6 +8363,7 @@ def test_directx_toolchain_frontier_matches_pinned_dxc_inventory(): module.MLX_BINARY_TWO_SOURCE: 225, module.MLX_LAYER_NORM_SOURCE: 2, module.MLX_RANDOM_SOURCE: 2, + module.MLX_RMS_NORM_SOURCE: 12, module.MLX_ROPE_SOURCE: 18, module.MLX_TERNARY_SOURCE: 212, } @@ -7994,12 +8375,12 @@ def test_directx_toolchain_frontier_matches_pinned_dxc_inventory(): module.MLX_SCALED_DOT_PRODUCT_ATTENTION_SOURCE: 42, module.MLX_SOFTMAX_SOURCE: 10, } - assert len(expected_sources) == 6 + assert len(expected_sources) == 7 assert module.MLX_DIRECTX_TOOLCHAIN_ENTRY_POINT_COUNT == sum( module.MLX_DIRECTX_TOOLCHAIN_ENTRY_POINT_COUNTS.values() ) - assert module.MLX_DIRECTX_TOOLCHAIN_ENTRY_POINT_COUNT == 470 - assert module.MLX_DIRECTX_TOOLCHAIN_ARTIFACT_COUNT == 7 + assert module.MLX_DIRECTX_TOOLCHAIN_ENTRY_POINT_COUNT == 482 + assert module.MLX_DIRECTX_TOOLCHAIN_ARTIFACT_COUNT == 19 assert sum(module.MLX_DYNAMIC_WORKGROUP_ENTRY_POINT_COUNTS.values()) == 106 assert { source: evidence["specializationCount"] @@ -8075,12 +8456,12 @@ def test_directx_frontier_readme_records_compile_only_scope_and_current_gaps(): normalized_readme = " ".join(readme.split()) assert "official DXC v1.9.2602.24 on Windows CI" in readme - assert "seven-artifact frontier representing six pinned sources" in ( + assert "19-artifact frontier representing seven pinned sources" in ( normalized_readme ) - assert "11, 225, 2, 2, 18, and 212 entries respectively" in normalized_readme - assert "470 generated compute entries" in normalized_readme - assert "five pending aggregate sources cover 94 compute entries" in ( + assert "11, 225, 2, 2, 12, 18, and 212 entries respectively" in (normalized_readme) + assert "482 generated compute entries" in normalized_readme + assert "four pending aggregate sources cover 82 compute entries" in ( normalized_readme ) assert "no placeholder workgroup size is restored" in normalized_readme @@ -8105,6 +8486,11 @@ def test_directx_frontier_readme_records_compile_only_scope_and_current_gaps(): assert "does not dispatch these kernels or establish numerical parity" in ( normalized_readme ) + assert "captures 12 distinct dispatch artifacts exercised by the pinned" in ( + normalized_readme + ) + assert "test_rms_norm` and `test_rms_norm_grad` workloads" in normalized_readme + assert "the MLX runtime is not redirected to these artifacts" in (normalized_readme) assert "DirectX remains outside the DXC gate" not in readme @@ -8119,7 +8505,7 @@ def test_selected_quantized_frontiers_record_current_target_boundaries(): assert "profile `cs_6_0`, no `-enable-16bit-types`, and `-WX` passes" in ( normalized_readme ) - assert "zero warnings across all 470 entry-point runs" in normalized_readme + assert "zero warnings across all 482 entry-point runs" in normalized_readme assert "records this as a warning-clean contract" in normalized_readme assert "rejects any newly observed warning" in normalized_readme assert "two selected `random.metal` entries compile without" in normalized_readme diff --git a/tests/test_mlx_rms_norm_dispatch_contract_fixture.py b/tests/test_mlx_rms_norm_dispatch_contract_fixture.py new file mode 100644 index 000000000..ecd1a8919 --- /dev/null +++ b/tests/test_mlx_rms_norm_dispatch_contract_fixture.py @@ -0,0 +1,210 @@ +import hashlib +import json +from pathlib import Path + +import crosstl.project as project_api + +MLX_COMMIT = "4367c73b60541ddd5a266ce4644fd93d20223b6e" +MLX_REPOSITORY = "https://github.com/ml-explore/mlx" +MLX_HOST_SOURCE = "mlx/backend/metal/normalization.cpp" +MLX_KERNEL_SOURCE = "mlx/backend/metal/kernels/rms_norm.metal" +FIXTURE = ( + Path(__file__).resolve().parents[1] + / "demos" + / "integrations" + / "mlx" + / "contracts" + / "rms_norm.dispatch.json" +) +EXPECTED_MANIFEST_DIGEST = ( + "ea238af83b140c33d43b79b5efd1814c398bbc09ac70cbef21375b1e8ff9a1eb" +) +EXPECTED_VARIANTS = { + "forward-float32-axis-32": ( + "rmsfloat32", + 32, + {}, + (2, 1, 1), + "4306831dce3a9a479ef63093a7f2722358caf58b6fcf1a47ed808a1c28dc9ebb", + "00c05fccf276cf11f3fb9b617b8fe0bb3c5f8766c0e4ca1ed990c093e700422e", + ), + "forward-float32-axis-256": ( + "rmsfloat32", + 64, + {}, + (2, 1, 1), + "f27b102d2ee4f473afef2448e42103223922ff29fd703e6d269d827f280e8bf7", + "1ef80b00c1a7a2f7967177bc003a961f3e4448358716d2deb79778fc3cbfb68e", + ), + "forward-float32-axis-512": ( + "rmsfloat32", + 128, + {}, + (2, 1, 1), + "369be00b341be529b913101657bd6a78d8b8c689e9dc7c6d27e694809dd9d098", + "a9b6980b1645e867b2502052d6d5f37a03447258bb4f819535901bf45a01d5da", + ), + "forward-float16-axis-32": ( + "rmsfloat16", + 32, + {}, + (2, 1, 1), + "51de914a20a4defb1d1b79ed26def94a994894228d6c64cae0e1132553809dbb", + "b694e5240f2a87bfae8af862878251a45cfaeaf39fd810bf2df8a5e3724bdad7", + ), + "forward-bfloat16-axis-32": ( + "rmsbfloat16", + 32, + {}, + (2, 1, 1), + "635d28f468bf3e15a0ec7285aacc23d0de32f9596d511beb98adb3ecc2c2abf4", + "13655322998b557a5143ac9b871dc898f0ae43c07cf9659aecd505156ed9318b", + ), + "forward-float32-axis-4099": ( + "rms_loopedfloat32", + 1024, + {}, + (1, 1, 1), + "392aee49734fcc2fd3ff4fd232a49f1a97a2bb79f3c22cc80eeabeb1f1ca1959", + "b81c2043b10bde966cb6f4dbfa198d2b93a3e456f3026030b69557c4a8983729", + ), + "vjp-float32-axis-32-has-w-false": ( + "vjp_rmsfloat32", + 32, + {"20": False}, + (800, 1, 1), + "1a25dca51070c3b6fc96f162e6c152049d388833ba02b1f3d10cc1928c5661c4", + "ef832e1ceb8c864a13aee3460d23658f4fffba18db1b800461628ba6ebe38e0a", + ), + "vjp-float32-axis-32-has-w-true": ( + "vjp_rmsfloat32", + 32, + {"20": True}, + (800, 1, 1), + "26177a77e484a56b7b2572516e8e1360714c88b668881237a9cf499129e34f35", + "a9be06b43a6156fb9ee1f9a6955d03d6bda0940c2a8223b58f564c2d12bd0cd0", + ), + "vjp-float32-axis-256-has-w-false": ( + "vjp_rmsfloat32", + 64, + {"20": False}, + (800, 1, 1), + "f26ffb357ecfabef6216ff45661f298af4f164a0494da838ca579ea0812e73d2", + "4455fc9204f826fc5d0d7f016bcaf970b75ea31d5f1598d13ceca6a2baa369e7", + ), + "vjp-float32-axis-256-has-w-true": ( + "vjp_rmsfloat32", + 64, + {"20": True}, + (800, 1, 1), + "a010324f59769fde9d71cc8968852ae1e6c8b0ddc213f035523c2c5f2e12d413", + "0944044e2f050bedde1d05e1ae5648628e7144c04752dd49b2e2bb7bcd807b7b", + ), + "vjp-float32-axis-8192-has-w-false": ( + "vjp_rms_loopedfloat32", + 1024, + {"20": False}, + (4, 1, 1), + "fa4d57c473cd5799c46a982e3ea339debc06de21fb6a255adc11bbb546f9329b", + "3bd55b546fc00ddf8412f092da4793c0272eec0ad7c130065ad7c1677f60cdce", + ), + "vjp-float32-axis-8192-has-w-true": ( + "vjp_rms_loopedfloat32", + 1024, + {"20": True}, + (4, 1, 1), + "8fc64f93a1be95f7acac67d6595c9f9c66c01b87cf952f29052835c20d4d765b", + "345d524ffec14682b6d0325bc97b624b89d83dc257a17ed49bea5e11e24573f3", + ), +} + + +def test_rms_norm_dispatch_fixture_pins_schema_identity_and_provenance(): + manifest = project_api.load_dispatch_contract(FIXTURE) + canonical = json.dumps( + manifest.to_json(), + ensure_ascii=True, + allow_nan=False, + separators=(",", ":"), + sort_keys=True, + ) + + assert manifest.schema_version == project_api.DISPATCH_CONTRACT_SCHEMA_VERSION == 1 + assert manifest.content_identity.to_json() == { + "algorithm": "sha256", + "value": EXPECTED_MANIFEST_DIGEST, + } + assert hashlib.sha256(canonical.encode("utf-8")).hexdigest() == ( + EXPECTED_MANIFEST_DIGEST + ) + assert manifest.provenance["repository"] == MLX_REPOSITORY + assert manifest.provenance["commit"] == MLX_COMMIT + assert manifest.provenance["sourceReferences"] == { + "hostDispatch": MLX_HOST_SOURCE, + "kernel": MLX_KERNEL_SOURCE, + } + assert manifest.contracts[0].provenance["hostSource"] == MLX_HOST_SOURCE + assert manifest.contracts[0].provenance["kernelSource"] == MLX_KERNEL_SOURCE + + +def test_rms_norm_dispatch_fixture_evaluates_unit_test_workloads_exactly(): + result = project_api.load_dispatch_contract(FIXTURE).evaluate() + variants = {variant.workload_id: variant for variant in result} + + assert set(variants) == set(EXPECTED_VARIANTS) + assert len(variants) == 12 + for workload_id, expected in EXPECTED_VARIANTS.items(): + ( + entry_point, + workgroup_x, + specialization_constants, + dispatch_size, + variant_digest, + artifact_digest, + ) = expected + variant = variants[workload_id] + assert variant.entry_point == entry_point + assert variant.contract_id == "mlx-rms-norm-unit-tests" + assert variant.device_id == "wave32-max1024" + assert variant.source == MLX_KERNEL_SOURCE + assert variant.workgroup_size == (workgroup_x, 1, 1) + assert variant.subgroup_width == 32 + assert variant.capabilities == { + "maxThreadsPerWorkgroup": 1024, + "simdWidth": 32, + } + assert variant.specialization_constants == specialization_constants + assert variant.dispatch_field == "workgroupCount" + assert variant.dispatch_size == dispatch_size + assert variant.variant_id == f"sha256:{variant_digest}" + assert variant.artifact_id == f"sha256:{artifact_digest}" + + +def test_rms_norm_dispatch_fixture_is_finite_test_backed_and_non_runtime(): + manifest = project_api.load_dispatch_contract(FIXTURE) + normalized = manifest.to_json() + workloads = {workload["id"]: workload for workload in normalized["workloads"]} + scope = manifest.provenance["scope"] + + assert len(manifest.workloads) == 12 + assert len(manifest.contracts) == 1 + assert len(manifest.contracts[0].branches) == 6 + assert workloads["forward-float32-axis-32"]["provenance"] == { + "hostFunction": "RMSNorm::eval_gpu", + "branch": "single-row", + "testSource": "python/tests/test_fast.py::test_rms_norm", + "shape": [2, 32], + "coveredAxisSizes": [31, 32, 33], + } + assert { + workload["provenance"]["testSource"] for workload in workloads.values() + } == { + "python/tests/test_fast.py::test_rms_norm", + "python/tests/test_fast.py::test_rms_norm_grad", + } + assert scope == { + "description": "Pinned RMSNorm unit-test dispatch records.", + "singleRowAndLoopedEntriesIncluded": True, + "runtimeExecutionVerified": False, + "numericalParityVerified": False, + }