round 5 track b: cached residual norms and fused mtp attention #

deliverables / status #

patch: /tmp/sol_r5_mtp_qsa/sol_r5_mtp_qsa.patch
scratch checkout and evidence: /tmp/sol_r5_mtp_qsa/
source baseline: strata head 296b59d2f73c8f04a805df29c2fa97655adace78; exact read-time file hashes and final hashes in baseline_manifest.json (the shared checkout was already dirty).
patch sha256: 9c5a81a836a837a09955086ffb18dff9922f8b2d998d35ed7995d6c9c0a6689f
git -C vendor/strata apply --check /tmp/sol_r5_mtp_qsa/sol_r5_mtp_qsa.patch: pass. patch is not applied; no tracked strata file was edited.
no access to port 8080; no gpu dispatch, gpu benchmark, ssd benchmark, or subagent spawned.

this is a correctness-screened candidate, not a demonstrated speedup or hardware-bitwise-parity result. the patch switches the window to the new entry points; do not treat those default changes as release-ready before authorized gpu parity and paired timing.

preflight / ownership #

normalization belongs in gr.metal / gr2.metal; launch shape and prepared-plane stride belong in window.cpp. mtp attention belongs next to the qsa reference kernels, not in the general attention path. preserve the original stream-local accumulation/reduction order, token and stream offsets, fp32 residual writes, expert-plan indirection, per-64-cell chunk arithmetic, max-first chunk merge, sigmoid gate, all four prepared planes, and per-16 sums. arbitrary-size generic norm callers and larger configured mtp windows must remain supported.

checks chosen before implementation: run the actual baseline and candidate kernel bodies under a cpu simd/threadgroup emulator, compare complete outputs byte-for-byte, compile actual runtime library combinations with safe math, syntax-check all changed c++ sources, and retain standalone gpu parity tests for a later authorized run. any output mismatch rejects the candidate; no timing claim is possible within the authorized compile/cpu-only budget.

changes #

residual norms #

gr_norm and generic gr_norm_comb retain each thread's residual values in a 10-float local cache when n_embd <= 2560; larger shapes keep the old device reread. no writeback, fma, or reduction ordering changes.
new gr_norm_comb_4s: fixed n_embd=2560, hc=4, grid (1, ntok), 256 threads. form the routed-plus-shared expert output once per dimension, then apply it to all four streams. retain rv[4][10], preserve each stream's element order and eight-simdgroup reduction order, publish part[4][8] across one threadgroup barrier.
window's shared gr_norm launcher selects _4s for deferred combines, with a compile-time fixed-shape assertion; both gr_read and gr_read_m use this helper. other callers of the generic kernel do not change launch shape.
fuse_parity.cpp now covers both generic and four-stream combines, apply off/on, and ntok 1 through 8. syntax checked only; not gpu-executed.

source-level traffic removed at the requested 50+48 norm-call counts: 98 * 40 kib = 3.83 mib of residual rereads per token, plus 48 * 3 * 11 * 2560 * 4 = 15.47 mib of redundant expert/shared-output reads. these are logical load counts, not measured dram traffic; caches and compiler lowering matter. private arrays are not proof of physical register residency or absence of spills.

mtp attention #

important correction: the default host range includes the current cell, so the mature width is 64 sink + 256 previous local + 1 current = 321, i.e. six chunks, not five. environment overrides permit up to 2051 cells.

new qsa_mtp_attn_merge: grid (2, ntok), 256 threads, restricted to width <=384. stage all six chunks' probabilities and per-chunk maxima/denominators in threadgroup memory, keeping the original qsa score-dot, simd reductions, exponentials, value fmas, and final max-first merge order.
compute value accumulators for all 12 heads together, retaining one value load per dimension per cell rather than rereading values per head.
emit prepared planes directly; common qsa_write_prepared retains the old merge's exact per-16 sum order. the fused caller adds the required barrier before reusing the shared sum scratch for its next head.
mtp_encode uses the fused path through width 384, with plane stride MAXT but grid token count 1; larger STRATA_METAL_MTP_WIN values keep the original split/merge path. selection of ids, ids slot offset, projections, and general decode/prompt attention are unchanged.
at width 321, remove one host dispatch, its inter-dispatch barrier, and six chunks * 24 heads * 258 floats = 148608 bytes each of partial writes and reads. that is 290.25 kib of logical round-trip traffic, not measured dram traffic.
scratch footprint from declared arrays: 32320 bytes (under 32768). ids are read directly, avoiding an additional 1536-byte staged array that would exceed the limit. attention heads/chunks are more serial than the split reference; occupancy, barrier overhead, register pressure, and reduced available parallelism may outweigh the saved launch. timing remains necessary.
new metal/tests/qsa_mtp_parity.cpp: checks bitwise prepared planes plus untouched stride rows for width 1,31,64,65,128,319,320,321,383,384; two tokens except at width 384; noncontiguous cache ids; plane stride 8. syntax checked only; deliberately not run.

actual verification #

cpu-emulated baseline/candidate kernel bodies: pass, exit 0 (cpu_emulation.log).
- bodies mechanically extracted from the actual metal sources, with only attributes/address spaces and synchronization intrinsics adapted to c++; explicit fma retained, -ffp-contract=off.
- norms: n=255,2560,2817, apply off/on, two tokens and all four streams. check residual writeback and normalized output byte-for-byte. checks the fixed-shape four-stream path at n=2560 and the >2560 fallback.
- attention: width=1,31,64,65,128,319,320,321,384. check all prepared planes, per-16 sums, and untouched token rows byte-for-byte. two tokens except at width 384; random int8 k/v, fp16 per-64 scales, fp32 queries/gates, scattered ids; seed 20260701.
- the emulator uses the same fixed tree in both implementations, not a claim about apple's exact simd tree, transcendental implementation, or compiler lowering. hardware bitwise parity remains untested.
- repeat with -fsanitize=address,undefined: pass, exit 0, same complete case set, no sanitizer findings (cpu_emulation_asan.log).
actual newLibraryWithSource, MTLMathModeSafe: pass, exit 0 (safe_Metal.log).
- original gr runtime combination: 717 entry points; candidate: 718.
- original qsa+qsa2: 9 entry points; candidate: 10.
- compilation tool creates a device and libraries only; no command queue, command encoder, pipeline creation, dispatch, or synchronization with the running server.
clang++ -std=c++20 -fsyntax-only: pass, exit 0 for final window.cpp, fuse_parity.cpp, qsa_mtp_parity.cpp; logs window_syntax.log, fuse_syntax.log, qsa_mtp_syntax.log.
git diff --check and patch apply check: pass, exit 0 (diff_check.log, apply_check.log).
source review: all diff pages reviewed with view_diff, and changed/new context reviewed with view_code; fixed long cache-store loops, test output formatting, and made the launch's fixed shape explicit.

reproduction #

all commands below are cpu or library compilation only:

clang++ -std=c++20 -O1 -ffp-contract=off -pthread /tmp/sol_r5_mtp_qsa/cpu_emulation.cpp -o /tmp/sol_r5_mtp_qsa/cpu_emulation
/tmp/sol_r5_mtp_qsa/cpu_emulation
clang++ -std=c++20 -fobjc-arc /tmp/sol_r5_mtp_qsa/compile_only.mm -framework Foundation -framework Metal -o /tmp/sol_r5_mtp_qsa/compile_only
/tmp/sol_r5_mtp_qsa/compile_only /tmp/sol_r5_mtp_qsa/gr_baseline.metal /tmp/sol_r5_mtp_qsa/gr_library.metal /tmp/sol_r5_mtp_qsa/qsa_baseline.metal /tmp/sol_r5_mtp_qsa/qsa_library.metal
git -C vendor/strata apply --check /tmp/sol_r5_mtp_qsa/sol_r5_mtp_qsa.patch

next gate #

under a later gpu authorization, run both parity executables first, then compare original/candidate end-to-end decode and mtp draft latency on identical inputs. retain the old entry points and wide-window fallback. no performance acceptance is established here.