Save inference parity implementation and evaluation harness
This commit is contained in:
Georg Bauer
@@ -0,0 +1,551 @@
|
||||
MTPLX Qwen runtime shaders
|
||||
=========================
|
||||
|
||||
mtplx-runtime-0.32.2.metallib is an unchanged copy of mlx/lib/mlx.metallib
|
||||
from the existing MTPLX reference environment (runtime version 0.32.2).
|
||||
|
||||
SHA256: dc59d1cceb1a5c7e578232e6e41e28e2c73c9463ac6dbc3886c3ee17ffc270ed
|
||||
Source tag: v0.32.2
|
||||
Source commit: 1f8e74e3f12f31365464a6867c6579f0e9b29d85
|
||||
Source: https://github.com/ml-explore/mlx/tree/v0.32.2
|
||||
License: MIT, reproduced in MLX-LM-LICENSE.txt (identical license text).
|
||||
|
||||
Only GPU shader code is included. No libmlx.dylib, Python inference code,
|
||||
or C/C++ application/backend code is linked into DS4Server. Rust selects
|
||||
the shader entry points, binds model buffers and owns graph execution.
|
||||
|
||||
The reference's get_quantized_kernel implementation was verified to call
|
||||
Device::get_kernel on the default precompiled library, matching
|
||||
mlx/backend/metal/nojit_kernels.cpp at the pinned commit. The shader library
|
||||
is reused byte-for-byte, including its precompiled specializations, rather
|
||||
than independently rebuilding different kernels or compiler settings.
|
||||
|
||||
Provenance on the evaluation machine:
|
||||
/private/tmp/MTPLX-analysis-20260901/.venv/lib/python3.12/site-packages/mlx/lib/mlx.metallib
|
||||
|
||||
The Rust test mtplx_runtime_shaders_are_pinned enforces the complete file
|
||||
hash. Existing resource packaging includes the metal directory recursively.
|
||||
This artifact adds approximately 174 MiB to the resources; it is not a model
|
||||
download. Loading is lazy on the first runtime-shader dispatch.
|
||||
|
||||
The shader identity is not a claim that the complete Rust model graph,
|
||||
scheduling, cache behavior, or end-to-end performance has reached parity.
|
||||
|
||||
Compiled QSA moving offsets
|
||||
--------------------------
|
||||
|
||||
tests/fixtures/mtplx-qsa-update-jit.json records the actual compiled indexer
|
||||
scalar shaders and compute_dynamic_offset_int32, observed without changing
|
||||
the original compilation by tools/mtplx-jit-reference.py --operation qsa-update.
|
||||
Only the contiguous int32[1] scalar specializations are retained; full observed
|
||||
library source hashes are included. --check reruns the original compiled core.
|
||||
The Dynamic Offset body is used unchanged by the Rust dynamic-copy path;
|
||||
its original full source hash is
|
||||
48a7309664f797e749aa42d2c2c4db0cf3abedf97297f0068d06b7847d988b93.
|
||||
It is Copyright Apple Inc., MIT as reproduced in MLX-LM-LICENSE.txt.
|
||||
The following gg1/gg2_dynamic_copybfloat16bfloat16 kernels are taken directly
|
||||
from the unchanged runtime metallib. No frontiers are read back to the CPU.
|
||||
Clamp/Multiply now feed the connected qsa_compiled_cache_window stage.
|
||||
mtplx-qsa-compiled-scalars.json retains the actual generated scalar kernels,
|
||||
including all three constant-CSE layouts for the Clamp and an independent
|
||||
257/255 specialization check. Only structural integer literals and exported
|
||||
symbols change when Rust specializes a kernel; the computations are unchanged.
|
||||
mtplx-qsa-compiled-header.metal is the unmodified original compiler prefix;
|
||||
SHA256: 2665a76463f3f6ee283c6a50b66e4a527318a114080b31441dfa900042097a39.
|
||||
--qsa-header --check compares that prefix with a fresh original compilation.
|
||||
Both resources contain runtime shader code, not host runtime code. Most is
|
||||
Apple MIT. The unchanged full prefix also retains the Apache-2.0 cexpf.h
|
||||
notice (Apple, NVIDIA, Filipe RNC Maia; license text in MTPLX-LICENSE.txt) and
|
||||
the full BSD-2-Clause expm1f.h notice/disclaimer (Norbert Juffa 2015-2023).
|
||||
Those overloads are retained as original header dependencies, not new Qwen
|
||||
complex/exponential computation in the integer scalar kernels.
|
||||
26 compiled reference calls verify the complete retained-input cache window.
|
||||
They do not establish graph-bank replay, allocation/donation, BFS scheduling
|
||||
or production inference parity.
|
||||
|
||||
tests/fixtures/mtplx-qsa-select-jit.json also retains the actual compiled Add
|
||||
kernels for selector row offsets 0/1/2/3. Rust substitutes only the structural
|
||||
offset literal and exported symbol, preserving the original +0 dispatch.
|
||||
--operation qsa-update --qsa-mode blocks --qsa-score-budget 4096 --qsa-header
|
||||
--check reproduces the connected reference selector and its exact sources.
|
||||
Query preparation and both selector families now accept GPU frontier leaves
|
||||
through the same original kernel dispatch used by the host-frontier entry.
|
||||
66 actual compiled reference calls cover all five output modes, chunked
|
||||
selection and the connected cache state for retained old input leaves. These
|
||||
are functional Q/K-entry checks, not graph-bank, ownership or UI performance
|
||||
acceptance. Host integration and the runtime evaluator remain open.
|
||||
The retained-input Hidden entry has since been connected through the already
|
||||
verified original affine projection kernels to that same Q/K implementation.
|
||||
132 actual select_hidden calls cover 4/8-bit, group32/64 projections and all
|
||||
five output modes. The combined entry test covers 198 reference calls. No new
|
||||
Metal bodies or alternative projection/selection arithmetic are introduced.
|
||||
The installed B1/BF16 cache/phase routing is now connected to backing reserve,
|
||||
explicit GPU frontiers, those same retained-input arithmetic entries and cache
|
||||
commit. 720 actual original host-method decisions and 108 additional complete
|
||||
indexer calls cover routing and ongoing state/output transitions. The combined
|
||||
host-flow test includes the previous 92 non-compiled calls through the same
|
||||
entry. A parameter-bound QSA graph bank now replaces the direct compiled
|
||||
expression chains. It rebinds explicit inputs to cached primitive dependencies
|
||||
and uses the pinned degree/BFS-width algorithm for a single indexer graph.
|
||||
99 optimized original graph contracts, all 198 arithmetic calls, stride-changing
|
||||
replay, parameter invalidation and the original connected-call trace/entry
|
||||
counters pass. Evaluated constants are omitted from structural fingerprints;
|
||||
kernel-source and output checks remain separate. No alternate Metal kernel was
|
||||
introduced. The dtype/shape-generic guard, donation/allocator, early release,
|
||||
global model scheduling and production integration remain open. Last-use
|
||||
graph leaves are now detached after their consumer, separately from explicit
|
||||
completion ownership that protects GPU work until its existing CB finishes.
|
||||
The shared canonical dispatch bridge holds bound Metal resources through
|
||||
completion as well; this fixes four Invalid Resource failures exposed with
|
||||
unretained command buffers. Compile/dispatch use scoped autorelease pools.
|
||||
The normal and strengthened model-free collections cover 34 tests. This is
|
||||
not performance acceptance or evidence of matching whole-model Metal encoder
|
||||
timelines. Donation must still account for both descriptor and shared Data
|
||||
ownership, including outstanding GPU evaluator holds.
|
||||
The canonical test-bound Buffer now separates array/view identity from shared
|
||||
Data ownership, including nested native views and completion holds. QSA COW
|
||||
checks both Rc<Buffer> sharing and underlying Data sharing. A direct pinned
|
||||
QSACache alias/view update and its negative Rust regression check prove that
|
||||
array aliases observe replacement while distinct views retain old values.
|
||||
The QSA graph now applies the pinned primitive input/sibling-minus-primary-output
|
||||
Data retention protocol before the next primitive, rather than re-holding leaves
|
||||
at their last consumer. Its scheduler tracks the actually selected sibling, not
|
||||
just the producer node. CPU scheduling and GPU ownership-count checks cover that
|
||||
distinction, duplicate Data, empty-batch fallback and completion.
|
||||
QSA raw/pool DynamicSliceUpdate now performs actual BF16 vector-copy donation
|
||||
for exclusive mutable cache state. Retained inputs, snapshots, views and GPU
|
||||
Data holds select copying instead. The 16 KiB bound uses root allocation size.
|
||||
All 198 core reference cases also run with state snapshots and exclusive state,
|
||||
checking old/new hashes and actual Data reuse. No kernels or synchronization
|
||||
boundaries changed. Model-wide integration, other primitive donation and generic
|
||||
dtype/layout contracts remain open.
|
||||
|
||||
The allocator policy is now ported from the pinned buffer_cache.h with the same
|
||||
best-fit/oldest-equal-size choice, strict reuse ceiling and age-based/90%-clear
|
||||
eviction. A 101-event trace from the real installed runtime checks allocation
|
||||
identities and active/cached bytes, including cache-limit transitions. Reference page rounding and the
|
||||
device maxBufferLength precheck are connected to the test-bound Buffer methods;
|
||||
logical views preserve tensor bounds while Data records the rounded root size.
|
||||
The explicit Rust Allocator now owns a real 1 MiB untracked/shared Metal heap,
|
||||
uses it for requests below 256 bytes with device-allocation fallback, and
|
||||
recycles native roots only after the final physical allocation hold releases.
|
||||
Its active/cache/peak/resource accounting, cache and memory limits, resource
|
||||
pressure GC (including the original unsigned subtraction), zero/null result
|
||||
and actual cached storage reuse are checked against the 101-event receipt.
|
||||
Residency/wired limits are now connected to this explicit allocator, including
|
||||
heap registration, cache retention and erase-before-release. Set selection and
|
||||
budgets follow resident.cpp: first fit, oversize/empty-set reuse, 32-set ceiling,
|
||||
emptiest fallback and touched-set commits on resize. Native membership and the
|
||||
ten original residency lifecycle scenarios are tested. Queue attachment uses
|
||||
the published set count and is exercised immediately before test commits;
|
||||
The new test-bound Submission owner now attaches automatically at actual native
|
||||
commit boundaries through a scoped encoding-thread callback, including flush,
|
||||
readback, finish and cleanup paths. Its queue cursor persists across batches;
|
||||
scope teardown drains before unregistering the callback. Legacy work outside
|
||||
the scope and other queues do not inherit it. Externally wrapped storage,
|
||||
the process-wide owner and model-wide routing remain open.
|
||||
Existing canonical helpers
|
||||
are not globally switched to untracked buffers before encoder dependencies
|
||||
are ported. No whole-model allocator/performance parity is claimed.
|
||||
|
||||
The test-bound Rust Encoder now owns an independent queue with unretained
|
||||
command buffers and Concurrent compute encoders, following the pinned
|
||||
device.cpp/event.cpp/error.h dependency and completion rules. Access roles,
|
||||
barrier epochs, deferred concurrent outputs, cross-encoder fences, temporary
|
||||
exclusion and shared-event error propagation are managed in Rust. The bridge
|
||||
only issues Metal API calls and reuses the same original kernel dispatcher.
|
||||
Commit thresholds count array.data_size() ELEMENTS (as the reference does),
|
||||
not allocation bytes; counters persist across encoder boundaries. The three
|
||||
new checks include dependent untracked GPU copies, two-queue event transfers
|
||||
and safe synthetic error-completion tests. The 48-test suite passes in both
|
||||
legacy retention modes. Only these new encoder tests use the independent
|
||||
queue. Complete operator access metadata, stream/evaluator integration and
|
||||
production routing are still open; this is not a model-performance receipt.
|
||||
|
||||
Operator scopes now connect the existing normalization and full MoE chain to
|
||||
the independent Concurrent encoder and its pooled untracked allocator. Explicit
|
||||
binding roles/data_size spans cover routing, sorting, gather/unsort, Gate/Up,
|
||||
SwiGLU, casts/norms, affine and gathered quantized projections, Split-K and both
|
||||
stock/fused experts plus the shared expert. Sort and split-reduction scratch is
|
||||
registered as backend temporaries. RMSNorm's default one and GatherSort's divisor
|
||||
are real scalar array bindings, not setBytes replacements. No shaders changed.
|
||||
The nine existing fixture groups (540 cases) execute through BOTH encoders with
|
||||
unchanged reference-output checks. This does not multiply independent fixtures.
|
||||
QSA routing, general array/donation semantics, evaluator/stream integration,
|
||||
production instrumentation and complete model/performance acceptance remain open.
|
||||
|
||||
GDN routing now uses the same typed dispatch/allocator path, including conv,
|
||||
mask/cache, Q/K normalization, compute_g/beta, recurrence, fused step and output.
|
||||
The direct native fused-step bypass is removed. Scalar operands and custom T
|
||||
are original scalar array inputs. Concatenate uses the original concurrent
|
||||
disjoint slice writes and dependency join; checked-input copies retain the
|
||||
original order and backend-temporary registration. No shader bodies changed.
|
||||
Four further existing groups (476 GDN cases) execute in both encoders, bringing
|
||||
the dual-encoder total to 1,016 existing cases. This remains operator-level
|
||||
correctness coverage, not whole-model scheduling or production parity.
|
||||
|
||||
QSA static/dynamic copies now bind explicit array data_size metadata rather
|
||||
than treating the copied region as the whole bound array. Dynamic offset arrays
|
||||
are inputs/backend temporaries. Zero fill, COW/General copies, compiled frontier
|
||||
operations and fused query/pool preparation carry access roles. KV concatenate
|
||||
uses the reference concurrent slice-write region. The existing 12 backing/copy
|
||||
and 26 compiled cache-window cases now run in both encoders (1,054 existing
|
||||
dual-encoder cases in total). Eager preparation, score/select operators and full
|
||||
graph/production integration remain open. Shader sources/geometries unchanged.
|
||||
|
||||
Eager QSA preparation now carries access roles and exact slice spans through
|
||||
RoPE, mean/RMS pooling and projections; all three RoPE concatenates use the
|
||||
reference concurrent writer regions. Scalar arithmetic uses array inputs.
|
||||
The fused selector and chunk concatenation also use the same encoder; its three
|
||||
frontiers are distinct scalar arrays and its private score plane remains a
|
||||
custom-kernel output, not a backend temporary. The existing 210 preparation,
|
||||
88 ongoing cache and 75 selector/chunk cases now execute in both encoders,
|
||||
bringing the existing dual-encoder total to 1,427 cases. Empty-array placeholder
|
||||
storage is unchanged; general zero-storage array semantics and the remaining
|
||||
score/eager-selection, evaluator and production paths are still open.
|
||||
|
||||
Both QSA prefill score producers and the connected score/top-k/chunk route now
|
||||
use the independent encoder. FP32 GEMV/Steel/NAX/Split-K, ReLU and reduction
|
||||
carry explicit array spans; Maximum/Divide scalars are array inputs. Matmul's
|
||||
output is allocated before transpose-copy preparation; those copies and the
|
||||
Split-K plane are registered as backend temporaries after their consumers.
|
||||
The MPP producer retains its original stride-aware inputs. Prefill top-k
|
||||
allocates outputs before checked-input copies, registers those copies after
|
||||
dispatch and uses three separately allocated frontier arrays. Shader bodies,
|
||||
specializations and dispatch geometries are unchanged. The existing 384 FP32
|
||||
score, 146 MPP/top-k/prefill and 96 mixed-producer cases now run in both encoders,
|
||||
bringing dual-encoder coverage to 2,053 existing cases. These manual fixture
|
||||
batches do not establish general evaluator/stream boundaries or production
|
||||
parity; eager selection/output and full graph integration remain open.
|
||||
|
||||
Eager QSA score masking, ranking and all output branches now carry explicit
|
||||
input/output array spans through the independent encoder. Scalar operands are
|
||||
real arrays; per-tile mx.eval(top_t) synchronizes that queue before constructing
|
||||
the next tile. Tiled index concatenation, decode-tail concatenation, rows-gather
|
||||
outputs and dense padding use the original concurrent slice-write regions.
|
||||
Dense padding is materialized before concatenation, and rows-gather builds its
|
||||
two separate Arange expressions; these restore previously collapsed operations.
|
||||
No Metal body or specialization changed.
|
||||
|
||||
The fixed-signature indexer graph now defers input/sibling Data on the active
|
||||
queue rather than retaining fallback descriptors in its returned state. The
|
||||
selected output's Data remains excluded, preserving donation. Compiled chunk
|
||||
concatenation joins concurrent writers and moving frontiers are four separate
|
||||
int32 arrays. Existing empty/nonempty Data-ownership checks run in both queues.
|
||||
The existing 540 eager cases and 200 ongoing indexer calls (including compiled
|
||||
routes) also execute in both encoders: 2,793 existing dual-encoder cases total.
|
||||
Explicit tile eval counts, output/state hashes and graph engagement are checked.
|
||||
These bounded fixture batches still do not apply the complete model-wide
|
||||
evaluator/stream scheduler and its automatic primitive-level commit policy.
|
||||
Production routing/instrumentation and whole-model performance acceptance remain
|
||||
open, as does general zero-storage array handling.
|
||||
|
||||
Primitive submission and task scheduler
|
||||
--------------------------------------
|
||||
|
||||
The fixed indexer tape now checks the original encoder thresholds after a whole
|
||||
primitive, never after an individual kernel dispatch. Input/sibling Data and
|
||||
backend holds are attached before a possible commit. Counted GPU tasks complete
|
||||
from the command callback; failed submissions balance ownership/accounting.
|
||||
The single-stream tape also applies the original active-task/memory pressure
|
||||
condition, finalizes its stream and waits for progress without inserting sleeps.
|
||||
|
||||
The Rust scheduler uses stdlib FIFO workers per CPU stream, earliest-error
|
||||
preservation, non-consuming cross-stream event error propagation and draining
|
||||
shutdown. CPU event waits/signals use the existing Metal shared-event bridge.
|
||||
Explicit tile/final fixture synchronizations are checked separately from
|
||||
automatic commits, with independent expectations for the two full-prefill
|
||||
indexer cases that cross the pinned Max data-size threshold.
|
||||
|
||||
CPU dispatch now counts every tenth operation, with completion as a separate
|
||||
FIFO task so failed work still completes its activity accounting. The CPU
|
||||
primitive cleanup task participates in that count and retains complete backend
|
||||
temporary descriptors until earlier work has run. CPU and GPU now share the
|
||||
input/sibling Data-selection function, excluding donated primary output Data.
|
||||
The private CPU temporary wrapper is Send only for drop-only worker ownership;
|
||||
Buffer itself remains !Send/!Sync. Five scheduler tests and the existing GPU
|
||||
ownership/connected-indexer checks pass, including unretained indexer execution.
|
||||
|
||||
An explicit runtime-owned stream registry now connects CPU/GPU encoders for the
|
||||
installed single-CPU/single-Metal backend. Defaults and template resolution are
|
||||
per thread/device. Local encoders are destroyed at thread exit; global streams
|
||||
allow sequential cross-thread use. Explicit clear preserves the reference's
|
||||
metadata/stale default handles and global CPU/GPU cleanup distinction.
|
||||
Registry locks do not cover encoding or waits. GPU selection reuses the same
|
||||
Rc/TLS encoder and allocator; exclusive ownership is checked before returning
|
||||
an encoder to storage. CPU/GPU events and finalize-all pressure handling use
|
||||
the existing scheduler and native bridge. Two new checks include actual
|
||||
two-queue copies, eleven blocked CPU tasks, thread cleanup and error unwinding.
|
||||
The 200 ongoing indexer cases now use a registered GPU stream in the independent
|
||||
path; the full model-free collection passes 57/57 normal and unretained.
|
||||
|
||||
Graph events now preserve per-copy values/origin streams while sharing native
|
||||
events and errors. Inter-stream fences use either the reference SharedEvent
|
||||
path or its opt-in Metal3/macOS15 fast path. Fast synchronization remains off
|
||||
by default. The existing pinned metallib supplies input_coherent, fence_update
|
||||
and fence_wait unchanged. Rust preserves array output registration, raw timestamp
|
||||
bindings, explicit update barriers, cross-device coherence and CPU SeqCst
|
||||
timestamp operations. Shared fence counts and per-dispatch snapshots are
|
||||
separate; completion/task ownership retains timestamp storage until work ends.
|
||||
The bridge now accepts empty dispatch grids, including the original zero-work
|
||||
coherence dispatch. Three new checks cover events/errors, partial-word/empty
|
||||
coherence and CPU/GPU/GPU/CPU transfer in both modes with a test-only deadlock
|
||||
rescue. All 60 model-free tests pass normal and unretained; no rescue fired.
|
||||
No shader body, metallib, inference default or production routing was changed.
|
||||
|
||||
This is still test-bound. Full array/evaluator cross-stream dependency construction, actual
|
||||
CPU model primitives, compile-cache cleanup integration and the complete model
|
||||
graph remain open. Production Qwen and performance acceptance remain unchanged;
|
||||
stream/operator checks do not establish whole-model parity.
|
||||
|
||||
Canonical dispatch now transfers deduplicated Rust root-allocation ownership
|
||||
to its existing Metal completion callback. This replaces the native resource
|
||||
set in that path and prevents physical release without adding Data aliases
|
||||
that would disable donation. Typed bindings also retain their descriptor/scalar
|
||||
borrow until dispatch. Early failures keep ownership with Rust; registration
|
||||
transfers it even when a submitted standalone command later reports failure.
|
||||
No shader, flush, commit, wait or additional completion-handler change.
|
||||
Normal and unretained collections pass 45/45 tests including native allocation,
|
||||
heap exhaustion, limits, residency and in-flight cache exclusion. This is not model parity.
|
||||
Binding input/output roles, barrier epochs, concurrent contexts, inter-encoder
|
||||
fences and reference commit thresholds still require the full encoder port.
|
||||
|
||||
Runtime-generated indexing shader
|
||||
--------------------------------
|
||||
|
||||
The same pinned runtime JIT-compiles gather_front instead of using the
|
||||
precompiled library. tools/mtplx-kernel-source.py resolves its two headers
|
||||
into metal/mtplx_qwen.metal without editing shader code and supplies the
|
||||
BF16/U32/FP32 template instantiations. Original notices and MIT attribution
|
||||
are retained. Header identities:
|
||||
|
||||
indexing/indexing.h:
|
||||
e820b8ee2b5132a97122780c12433ebb5100d8078d31e211d0429400a11415bb
|
||||
indexing/gather_front.h:
|
||||
64aacebf6576dfcd389383564fa1214bc87f2a091dd33cc64c598c5367ecab96
|
||||
|
||||
The generator requires the pinned runtime source checkout alongside the
|
||||
MTPLX reference checkout, named mtplx-runtime-0.32.2. The product requires
|
||||
only the generated shader resource, not this reference source checkout.
|
||||
|
||||
Runtime-generated SiLU shader
|
||||
----------------------------
|
||||
|
||||
tools/mtplx-jit-reference.py observes original runtime compilation without
|
||||
changing it. tests/fixtures/mtplx-silu-jit.json pins the generated19variants;
|
||||
full observed source SHA256:
|
||||
76cafb45db55a91efba66503dde59b37628ea360f220f35dd860f4ac3c3d0111
|
||||
|
||||
The shader exporter retains the generated computation and its original
|
||||
BF16 math, Sigmoid, Multiply, cast and stride helpers. Only host-name aliases
|
||||
change. All supporting header identities are enforced by the exporter and
|
||||
the emitted unit hashes by the focused Rust test. Copyright Apple Inc.; MIT
|
||||
as reproduced in MLX-LM-LICENSE.txt. This does not link a host runtime.
|
||||
|
||||
Runtime-generated GatherAxis shader
|
||||
-----------------------------------
|
||||
|
||||
The pinned GatherAxis source and generic elem_to_loc helper are retained
|
||||
unchanged, with BF16/U32 index instantiations for contiguous/strided inputs
|
||||
and int/int64 offsets. The router retains the strided last-ten-column view;
|
||||
no replacement top-k kernel is used. Header identities:
|
||||
|
||||
indexing/gather_axis.h:
|
||||
e1a745391ff4990f3f1ad75c5687c3b102dcdc4833d8fbbac38e10f54af29af4
|
||||
utils.h:
|
||||
5e1568e9edde9d05dbf86f68fa0d6c6240f2c32b973c7c6a76166b9c0d91543d
|
||||
|
||||
Copyright Apple Inc.; MIT as reproduced in MLX-LM-LICENSE.txt. Softmax,
|
||||
reduction, binary operations and index copies use the pinned metallib.
|
||||
|
||||
Runtime-generated SwiGLU shader
|
||||
------------------------------
|
||||
|
||||
The original compiled activations.swiglu used by non-sanitize-fused SwitchGLU
|
||||
and Qwen3NextMLP is captured separately from nn.silu. Full observed source SHA:
|
||||
bf78eee5cf96ea7c112c4e61546c12bcacf57fe512e0572182604cb94137510b.
|
||||
tests/fixtures/mtplx-swiglu-jit.json preserves all 19 original generated variants.
|
||||
The exporter changes only host aliases and reuses the already pinned BF16,
|
||||
Sigmoid, Multiply and cast/stride dependencies. Copyright Apple Inc.; MIT.
|
||||
|
||||
Runtime-generated compute_g shader and staged GDN
|
||||
------------------------------------------------
|
||||
|
||||
The original compiled gated_delta.compute_g is captured with the same observer:
|
||||
tests/fixtures/mtplx-compute_g-jit.json. Full observed source SHA256:
|
||||
34143a98046f8af5538767734fc169a5cab22a4920c26f9ba7ea45b8097152de.
|
||||
Receipt SHA256:
|
||||
701e2f54b7cb8bf97616f83256657e6c5e8cc8b46f4b65c559ccc030ab111dbf.
|
||||
All 19 variants retain their original BF16 Add/LogAddExp intermediates and
|
||||
FP32 final exponential. Only host aliases change. The exporter pins the
|
||||
additional Exp, Negative, Add, LogAddExp, Limits and log1p shader dependencies.
|
||||
|
||||
complex.h SHA256:
|
||||
16e8a815b2cbdb6070e0824e64fe33fccb6e918f1b84ea5c792bd89d33e57bf1.
|
||||
cexpf.h SHA256:
|
||||
88b6e15a52a5800d98d9bc6da840ca5cf70bf572fda136409580c1f17b1e0aab.
|
||||
The complex overload dependencies are retained unchanged, not used to add a
|
||||
complex-valued Qwen path. complex.h is Apple MIT. cexpf.h is Apache-2.0,
|
||||
Copyright Apple 2025, NVIDIA 2008-2013 and Filipe RNC Maia 2013. Its full original
|
||||
copyright/license notice remains embedded in the generated shader; the Apache
|
||||
license text is included in MTPLX-LICENSE.txt.
|
||||
|
||||
Stock depthwise Conv1D, copies, casts, reductions and elementary operations
|
||||
use the unchanged runtime metallib. Cache valid-length GatherAxis additionally
|
||||
instantiates the original signed INT32-index template; router indices remain
|
||||
UINT32. No shader body is replaced by a hand-written equivalent.
|
||||
|
||||
Original QSA indexer preparation
|
||||
-------------------------------
|
||||
|
||||
qsa_indexer_prepare.py SHA256:
|
||||
a77f6ca5ae805e729519c4629ae88b455a6dbf473a457a6e1c8219174eb59091.
|
||||
The exporter reads _prepare_queries_kernel and _pool_keys_kernel as AST data;
|
||||
it does not execute the model or kernel module. Both original source strings
|
||||
are unchanged. Header substitutions match the installed geometry: four query
|
||||
heads, width128, rotary64, ratio4, epsilon1e-6, attention scaling1. Includes are
|
||||
resolved at translation-unit scope; separate namespaces avoid collisions among
|
||||
the original header constants. Only entry-point declarations, host aliases and
|
||||
template instantiations are adapted. Stride metadata retains the original
|
||||
constant int64_t address space. Original Metal math and BF16 rounding remain.
|
||||
Copyright MTPLX; Apache-2.0, see MTPLX-LICENSE.txt and MTPLX-NOTICE.txt.
|
||||
|
||||
This is the preparation portion, not the full QSA indexer, selection,
|
||||
attention graph or production integration. Runtime frequencies are input buffers,
|
||||
not host replacements for the model's frequency construction.
|
||||
|
||||
QSACache/KVCache host lifecycle now uses the original scalar/vector/general
|
||||
copy and BF16/FP32 cast entries from this runtime, including positional writes,
|
||||
growth, strided restored state and the derived mirror. Rust distinguishes array
|
||||
object identity (__setitem__ overwrites its descriptor) from shared slice storage.
|
||||
Retained state aliases are checked against actual MTPLX cache operations, not
|
||||
assumed immutable. No additional shader bodies or runtime host library are used.
|
||||
The connected canonical cache remains test-only until product graph integration.
|
||||
|
||||
Original dynamic QSA selector
|
||||
-----------------------------
|
||||
|
||||
qsa_indexer_select.py SHA256:
|
||||
a3c74af27a7045c12f2893a8b7a91724c00d8a4148315c3165f3480c83016cf3.
|
||||
metal/mtplx-qsa-select.json preserves the original header, common body and all
|
||||
three epilogues (blocks, dense_mask, row_tokens), extracted without importing
|
||||
the model. Rust substitutes the original literal header parameters and supplies
|
||||
only the entry-point ABI. Tests additionally compare full generated header/body
|
||||
hashes against the actual MTPLX factory. H4/D128/ratio4 match the installed model;
|
||||
BF16/FP32 operands, backing capacity, top-k and TF32 remain specializations.
|
||||
Native compilation follows runtime 0.32.2 CustomKernel defaults: Safe math and
|
||||
its platform-selected Metal language version. No runtime host library is linked.
|
||||
The original 32MiB score-scratch chunk planner and typed output concatenation
|
||||
are connected to the cache/preparation port. General submission/concurrency,
|
||||
the complete eager indexer and production integration remain open.
|
||||
Copyright MTPLX; Apache-2.0, see MTPLX-LICENSE.txt and MTPLX-NOTICE.txt.
|
||||
|
||||
Original vectorized QSA prefill
|
||||
------------------------------
|
||||
|
||||
qsa_indexer_prefill.py SHA256:
|
||||
4d6fd428243c001746f69f8aed45991356772c2bd4a45586eb3c6813c91998d3.
|
||||
The same JSON export retains _MPP_SCORE_HEADER/_MPP_SCORE_SOURCE, the original
|
||||
top-k body and literal f-string header segments. Rust resolves only their named
|
||||
constants and provides entry-point ABI/type aliases. TensorOps tile layout,
|
||||
ordered per-head ReLU reduction, adaptive radix/insertion and all epilogues are
|
||||
unchanged. The original required General FP32 copy is used for non-contiguous
|
||||
score views; MPP input views keep their strides without added copies.
|
||||
The 128MiB producer-aware planner and score -> top-k -> concat chain are connected
|
||||
for the installed BF16/M5 geometry, including a 2K continuation from live cache.
|
||||
Full indexer branch routing, compiled graph bank and production integration
|
||||
remain open. Copyright MTPLX; Apache-2.0.
|
||||
|
||||
The general FP32 score expression now shares that prefill entry point. Rust
|
||||
ports runtime matmul.cpp's H4/D128 M5 Max routing: AsType Vector/General layout,
|
||||
check_transpose and broadcast copies, batch collapse, GEMV, regular Steel/NAX
|
||||
and both Split-K variants, original per-head Maximum, row/column Sum and Divide.
|
||||
All shader entries come from the unchanged pinned metallib; no shader body or
|
||||
host runtime library was added. The pooled cast is retained once across chunks;
|
||||
producer selection and the H4+1 workspace budget follow the reference.
|
||||
384 score cases and 96 connected selection cases are exact with real runtime
|
||||
MLX_ENABLE_TF32=0/1 in separate processes. These remain canonical correctness
|
||||
fixtures, not production integration or performance-parity evidence.
|
||||
|
||||
Original eager QSA selection
|
||||
---------------------------
|
||||
|
||||
The untiled QSAIndexer._select_eager score/top-k path uses the original
|
||||
runtime Arange, Add, integer Divide, Less, casts, Select, Subtract and
|
||||
ArgPartition (implemented by the pinned runtime as argsort). The chronological
|
||||
flash_prefill block epilogue adds original int32 Sort, int64 index conversion,
|
||||
bool GatherAxis and Select. The exporter adds only the required
|
||||
gather_axis<bool,int64_t,int,true,true> instantiation of the already preserved
|
||||
GatherAxis body; no body is changed. Its original file/unit hashes are unchanged.
|
||||
All untiled output epilogues are connected: dense mask (original bool
|
||||
ScatterAxis, repeat/concatenate and causal/tail mask), rows-gather (argsort-order
|
||||
tokens and validity), decode flash (chronological blocks, host tail bound) and
|
||||
decode gather (chronological tokens and variable-length tail). The flash branch
|
||||
retains precedence; neither decode branch evaluates the dead selected-mask DAG.
|
||||
Shared original cast/sort/vector dispatch helpers do not alter the shader bodies.
|
||||
|
||||
ScatterAxis adds these unchanged pinned Apple MIT runtime source units:
|
||||
- atomic.h, full-file SHA256:
|
||||
4c35ea2798a2335502865247aee878149fc9ada0d7e84c05d771baef0c7fcc60
|
||||
- reduction/ops.h None operation, full-file SHA256:
|
||||
78d06730fc9564a73944e7f1fe3897d25c8789b28a939bf418e1968db311da41
|
||||
- indexing/scatter_axis.h, full-file SHA256:
|
||||
43eabd0216101f8e32f5cdd19ce40b7f954564be27fad98a5e0fe345e7b94ce5
|
||||
Only include/pragma-once placement, namespace and the two required
|
||||
scatter_axis<bool,int64_t,int,None,false/true,true> instantiations are added
|
||||
outside the preserved bodies. The exporter guards full-file and body hashes.
|
||||
No host C/C++ runtime implementation is linked.
|
||||
|
||||
The 408 eager receipts tap actual QSAIndexer calls: 60 score/top-k, 24 prefill
|
||||
blocks, 144 dense masks, 108 rows-gather and 72 decode outputs. They include 2K
|
||||
queries, 65,536 blocks, tails 0/1/3 and separate real TF32-on/off processes.
|
||||
The tiled path shares the same score/rank functions, pooled FP32 input and tie
|
||||
vector. Each original mx.eval(top_t) is a synchronous command completion before
|
||||
the next tile, not an asynchronous flush. Only evaluated index views/backings
|
||||
are retained through the original GeneralGeneral uint32 concatenate. Index
|
||||
stride changes N -> K, independently of the N-strided validity. No new shader
|
||||
source or specialization is needed. Rust shares the original output-branch
|
||||
priority, including the tiled rows-gather exclusion and decode flash precedence.
|
||||
132 further actual-call receipts (108 tiled, 24 tile-off boundaries) bring the
|
||||
eager total to 540. They include observed reference eval row counts, 2K queries,
|
||||
65,536 blocks, tail 0/1, partial tiles and stride-2 FP32 query views.
|
||||
The existing GPU busy counters additionally verify one completed command buffer
|
||||
per observed reference tile eval, plus the final output batch.
|
||||
Full indexer routing, the general scheduling tape/allocator and production
|
||||
integration are not established by these checks.
|
||||
|
||||
Original eager QSA preparation
|
||||
------------------------------
|
||||
|
||||
The installed BF16 H4/D128 query and H1/D128 pool paths now use the original
|
||||
eager preparation expression as well as the fused custom-kernel branch.
|
||||
RMSNorm uses the pinned runtime kernel and its required General-Copy for sliced
|
||||
projection inputs. Pool mean is FP32 sum multiplied by 0.25 and cast to BF16
|
||||
before weighted RMSNorm. RoPE preserves all Arange, casts, concatenations,
|
||||
Cos/Sin, BF16 Negative, FP32 Multiply/Add and final BF16/pass-through stages.
|
||||
The installed rotary64/ratio4/eps1e-6/scaling1 contract is unchanged.
|
||||
All entries come from the unchanged pinned metallib; no shader body was added.
|
||||
|
||||
The existing projection and cache-extension entry points select either branch
|
||||
and retain the eager intermediates through their consuming operations. 105
|
||||
additional actual QSAIndexer receipts cover bare and quantized-projection
|
||||
preparation, including 2K rows, padded/stride-2 inputs and high positions. Another
|
||||
44 real QSACache/KVCache transitions cover eager pooling, capacity growth,
|
||||
reservation, trim, state aliases/restore and FP32 mirror rebuilding. These are
|
||||
not a complete indexer, production-integration or performance-parity receipt.
|
||||
|
||||
Connected non-compiled indexer entry
|
||||
----------------------------------
|
||||
|
||||
The Rust entry after compiled-route rejection now connects projection/supplied
|
||||
QK views, query preparation, raw/pool state and the original large-prefill,
|
||||
legacy-fused and eager selection order. All existing original shader dispatches
|
||||
are reused. Query preparation is dead when dense==sparse; KV.offset advances
|
||||
only in the subsequent Attention step. Shared return variants preserve the
|
||||
model-visible outputs while retaining other encoded kernel outputs.
|
||||
92 actual MTPLX indexer calls in 24 ongoing sequences verify lane choice, 82
|
||||
selection hashes, 268 raw/pool/mirror hashes, capacities/frontiers and command
|
||||
completion counts. They cover 2K rows, 32K history, supplied 704-stride QK views,
|
||||
prefill crossover, gather/flash priority and tiling. No compiled path is silently
|
||||
replaced. Compiled eligibility/core, attention, graph scheduling and production
|
||||
integration are still open; these are correctness, not performance receipts.
|
||||
Reference in New Issue
Block a user