regolith
strata but optimized for m5 (and maybe more later)
Star
Fork
Watch
regolith audits sol_mtp_fuse.md
6.4 kB · 52 lines
Markdown
at main

round 4 track b: mtp input and hyper-connection fusion #

deliverables #

  • patch: /tmp/sol_mtp_fuse/sol_mtp_fuse.patch
  • apply: git -C vendor/strata apply /tmp/sol_mtp_fuse/sol_mtp_fuse.patch
  • patch sha256: 4316deb086fa556ffd5328f60637d123dc1a2e5539c0eb2dc888c783f10803ae
  • isolated sources, reference harness, command and logs: /tmp/sol_mtp_fuse/
  • no tracked file under vendor/strata/ was edited. no server access, gpu dispatch, or gpu/ssd benchmark was run.

design and invariants #

owner: mtp.metal owns the fused arithmetic; Window::mtp_encode owns layouts and dispatch dependencies. this runs on every draft step and mtp prefill. existing prepared-activation allocations, weights, norm epsilon, residual layout, injection state and gemv call structure are unchanged.

mtp_rms_xprep uses the original 256-thread reduction, the same fma accumulation, simd sums, eight-part summation, and xr[d] * rs * w[d] order. the common mtp_prepare helper preserves all four planes and separate even/odd s16 sums exactly. the generic units.metal helper cannot be called directly: this mtp library is compiled standalone, without quant/units dependencies.

  • embedding: one 2560-wide norm per token, direct prepared output, no men_ writes.
  • hidden, n == 1: one 10240-wide norm, split into four 2560-wide prepared rows, with row count/plane stride still MAXT=8. no mhn_ writes.
  • hidden, n > 1: one 10240-wide norm per token, prepared stream 0 directly; retain normalized mhn_ for the existing stream 1..3 xprep/gemv sequence. deleting that scratch would be incorrect because XHn_ only holds eight prepared 2560-wide rows, not all four streams of eight tokens. the retained scratch is intentional, not a full-prefill intermediate-elimination claim.

mtp_combine_gnorm launches (HC, n) groups of 256 threads. each thread retains its ten combined values across the same reduction as gr_norm(apply=false), so neither rereading mR_ nor recomputing/reloading h/e is necessary. it writes residual mR_ and xn_, preserving v * w[d] * rs (not the other rms kernel's multiply order). a host static assertion binds the ten-register-value layout to N=2560.

gr_read_m(..., already_normed=true) skips only the initial norm. apply=true always runs the original norm regardless of the flag. other callers keep their original default behavior. the normal push() dependency barrier remains before gr_down reads xn_; removing the intermediate norm dispatch removes one barrier, not the producer-consumer dependency. embedding/hidden dispatches remain independent, followed by a barrier before gemv.

obsolete mtp_rms / mtp_combine entry points and their unused argument structs were removed. original snapshots remain available through the temporary repository's baseline commit.

checks actually run #

all final checks exited 0:

  1. clang++ -fsyntax-only on the isolated modified window.cpp, using existing build.ninja defines/includes/flags and the original engine header directory. exact command: /tmp/sol_mtp_fuse/syntax_command.txt; output: cpp_syntax.log (empty).
  2. clang++ -std=c++20 -O2 -ffp-contract=off -fsanitize=address,undefined /tmp/sol_mtp_fuse/reference.cpp -o /tmp/sol_mtp_fuse/reference, then the binary. 40 cases, n=1..8, five input patterns; 21,388,320 float bit-pattern checks passed. checks cover all xp planes/s16 including untouched padding and trailing guards, embedding and hidden nonzero source offsets, untouched draft scratch, retained prefill scratch and later stream preparations, combined residuals and normalized residuals. patterns: seeded random, signed zeros, mixed powers of two, alternating signs, and tiny inputs. no sanitizer failures. output: reference.log.
  3. actual safe-math Metal source compilation using newLibraryWithSource, then compute pipeline creation for both fused entry points and retained xprep. all have SIMD width 32 and max threads 1024, compatible with 256-thread groups. no command encoder or gpu dispatch. source: compile_metal.mm; output: metal_compile.log. the helper build has one expected deprecation warning for its pre-macos-15 fastMathEnabled fallback.
  4. git -C vendor/strata apply --check /tmp/sol_mtp_fuse/sol_mtp_fuse.patch: passed (apply_check.log).
  5. final visual diff review and context review performed on isolated sources and reference harness.

xcrun -sdk macosx metal -c ... is unavailable with this machine's command-line tools (utility not found); the real runtime Metal compiler above was used instead.

what the evidence does and does not establish #

cpu reference results establish exact bit-pattern equivalence of the tested arithmetic/layout transformations. both fused reductions retain the baseline arithmetic order under the engine's default safe-math setting. this is not a gpu-executed bitwise-parity claim: GPU simd_sum / rsqrt were not executed, per task restriction. the reference shares a deterministic simulated SIMD tree and uses host reciprocal square root; neither claims to emulate the device's exact numerical implementation. STRATA_METAL_MATH=fast reassociation is outside the bitwise-equivalence claim.

structurally: three fewer dispatches and two fewer barriers per encode call. for a single-token draft, the eliminated intermediate traffic is nominally 100 kib from normalized embedding/hidden write+read pairs plus 80 kib from the two gr_norm residual reads, or 180 kib/token of explicit device-buffer traffic. caches, register spills, occupancy, and wall-clock speed are unmeasured. scratch allocations remain unchanged. multi-token prefill still writes hidden scratch for later streams, so its traffic reduction is smaller (nominally 30 kib/token from rms/xprep plus 80 kib/token from combine/norm).

recommendation: ready for parent integration/review, then an authorized standalone gpu bitwise-parity test before claiming 100% device equivalence, and a separately authorized benchmark before claiming a speedup.

baseline identity #

vendor revision: 296b59d2f73c8f04a805df29c2fa97655adace78; snapshot included the working-tree source state, not just the revision.

  • original mtp.metal sha256: 564c849df5c36a22baae10b3f11396a5b45cb5835a67a3ee37325e418195fd4c
  • original window.cpp sha256: 8c892efe7441a604695da7bf1badaba09738aa1b5d482e4a46bc6991c6450d31