[Metal] Add TVM_METAL_STORAGE_MODE opt-in for correctness at scale (#20157) - #20391
Open
zacharywhitley wants to merge 1 commit into
Open
zacharywhitley wants to merge 1 commit into
zacharywhitley wants to merge 1 commit into
Conversation
…pache#20157) TVM's Metal runtime allocates every device buffer with MTLResourceStorageModePrivate. On Apple Silicon at graph scale (~11+ chained producer-consumer dispatches with certain fusion patterns, tracked in apache#20157) this exposes a driver-level correctness failure: outputs are silently wrong, cosine(LLVM, Metal) drops to around 0.75-0.95 on the repro. Switching allocation to MTLResourceStorageModeShared (unified memory on UMA — same physical bytes, stronger cache-coherency guarantees) closes the gate to numeric equivalence with LLVM. Adds an env-flag opt-in so callers can select the storage mode without a rebuild. Default remains Private for perf; this is a correctness escape hatch, not a default change. Verified fixes on macOS 15.5 / Apple M2 Max: - WaveNet-style Relax subgraph: cosine 0.7556 -> 1.000 - Synthetic N=11 gated activations: cosine 0.9465 -> 1.000 - Synthetic N=20 gated activations: cosine 0.7064 -> 1.000 - MobileNetV2 Metal: argmax 915 (buggy) -> 470 (gold), 5.3 ms warm inference Twelve TVM-side hypotheses ruled out via prior instrumentation (Relax memory reuse, KillAfterLastUse, TIR storage rewrite, per-encoder barriers, per-dispatch waitUntilCompleted, buffer zero-init, Metal MSL codegen — MSL is byte-identical between working and broken sizes). The bug survives all software-side synchronization and shows up only with Private mode, pointing at Metal driver behavior rather than TVM codegen. Follow-ups: - Benchmark Shared vs Private on typical GPU workloads to decide whether the macOS default should flip. - Consider promoting from env-flag to a target-level config option ({"kind": "metal", "storage_mode": "shared"}) once perf is characterized.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
TVM's Metal runtime allocates every device buffer with
MTLResourceStorageModePrivate. On Apple Silicon at graph scale(~11+ chained producer-consumer dispatches with certain fusion
patterns, tracked in #20157) this exposes a driver-level correctness
failure: outputs are silently wrong,
cosine(LLVM, Metal)drops to~0.75-0.95 on the repro.
Switching the allocation to
MTLResourceStorageModeShared(unifiedmemory on UMA — same physical bytes, stronger cache-coherency
guarantees) closes the gate to numeric equivalence with LLVM.
Adds an env-flag opt-in so callers can select the storage mode
without a rebuild. Default remains
Privatefor perf; this is acorrectness escape hatch, not a default change.
Verification (macOS 15.5 / Apple M2 Max)
Private(default)TVM_METAL_STORAGE_MODE=sharedAll match LLVM to fp32 noise (
max_delta ≤ 2e-5).What this doesn't do
Not a root-cause fix. On UMA,
PrivateandSharedrefer to the samephysical memory; the difference is driver-side cache coherency. Twelve
TVM-side hypotheses were ruled out via C++ instrumentation before
converging on the storage-mode workaround: Relax memory reuse,
KillAfterLastUse, TIR storage rewrite, per-encoder barriers,per-dispatch
waitUntilCompleted, buffer zero-init, MSL codegen (MSLis byte-identical between working and broken sizes). The bug survives
all software-side synchronization and shows up only under
Privatemode, pointing at Metal driver behavior rather than TVM codegen.
Suggested follow-ups
SharedvsPrivateon representative GPU workloads todecide whether the macOS default should flip.
(
Target({"kind": "metal", "storage_mode": "shared"})) once perfis characterized.
driver-side investigation.