Skip to content

[Metal] Add TVM_METAL_STORAGE_MODE opt-in for correctness at scale (#20157) - #20391

Open
zacharywhitley wants to merge 1 commit into
apache:mainfrom
zacharywhitley:fix/metal-storage-mode-shared-opt-in
Open

zacharywhitley wants to merge 1 commit into
apache:mainfrom
zacharywhitley:fix/metal-storage-mode-shared-opt-in

Conversation

@zacharywhitley

Copy link
Copy Markdown

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 (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.

Verification (macOS 15.5 / Apple M2 Max)

Workload Private (default) TVM_METAL_STORAGE_MODE=shared
WaveNet-style Relax subgraph cosine = 0.7556 cosine = 1.000
Synthetic N=11 (gated tanh×σ) cosine = 0.9465 cosine = 1.000
Synthetic N=20 cosine = 0.7064 cosine = 1.000
MobileNetV2 native Metal argmax = 915 (✗) argmax = 470, 5.3 ms warm

All match LLVM to fp32 noise (max_delta ≤ 2e-5).

What this doesn't do

Not a root-cause fix. On UMA, Private and Shared refer to the same
physical 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 (MSL
is byte-identical between working and broken sizes). The bug survives
all software-side synchronization and shows up only under Private
mode, pointing at Metal driver behavior rather than TVM codegen.

Suggested follow-ups

  1. Benchmark Shared vs Private on representative GPU workloads to
    decide whether the macOS default should flip.
  2. Consider promoting from env-flag to a target-level config option
    (Target({"kind": "metal", "storage_mode": "shared"})) once perf
    is characterized.
  3. File a companion Apple Metal report with the reproducer for
    driver-side investigation.

…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.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant