opus r20 track A: MTP map mlock + qsa_attn / qsa_merge decode tweaks #
patch: /tmp/opus_r20_mtp_qsa/opus_r20_mtp_qsa.patch (git -C vendor/strata apply --check clean against the current working tree).
scratch tree + build: /tmp/opus_r20_mtp_qsa/strata, /tmp/opus_r20_mtp_qsa/strata/build-metal (qsa_parity, arena_budget_test, strata_window_cli all build).
touched: metal/src/engine/window.cpp (Window::mtp_load only), metal/kernels/qsa.metal (qsa_attn / qsa_write_prepared / qsa_merge only), metal/tests/qsa_parity.cpp.
1. Window::mtp_load: wire what the draft reads, drop what it never reads #
The names-only DONTNEED loop inside the Q8 block is replaced by one pass over fm_.tensors after the Q8 block:
- skip =
mq8t_.count(name) || mnorm_off_.count(name) || name has "self_attn.indexer."-> MADV_DONTNEED on the interior 16 KiB pages (same rounding as before, so no page shared with a kept tensor is dropped). - else, unless
STRATA_METAL_MLOCK_DENSE=0(same switch as stream.cpp): covering 16 KiB pages clamped to the map end, MADV_WILLNEED +mlock. - logs
strata metal: MTP map: X GB wired (mlock; Y GB failed), Z GB unused released. - with
STRATA_METAL_MTP_Q8=0,mq8t_is empty so the BF16 tensors the GPU then reads fromWm_are wired rather than dropped (correct by construction, sinceMB()routes onmq8t_.count).
checked usage: the F32 tensors are only read via NM() -> mnorm_; TQ/MB route quantized names to mq8_; only the two mlp.experts.* tensors are read straight from Wm_ (moe_gu / moe_down).
on models/K512_iq2_xs/mtp-q2_0.gguf (889 MB), a standalone micro-test with the same selection (/tmp/opus_r20_mtp_qsa/micro/mlock_micro.cpp):
- kept/wired: exactly 2 tensors,
experts.down_proj235.9 MB +experts.gate_up_proj471.9 MB = 707.8 MB; released 181.2 MB (all BF16 now in mq8_, the 9 F32 norms, the indexer). - direct
mlockvs WILLNEED + 8-thread parallel touch + mlock:- warm page cache: 42-46 ms vs 43-49 ms
- cold (fresh
ditto --nocachecopy per run): 225/252 ms vs 235/238 ms (touch alone was 219-223 ms) - no measurable win for the parallel touch, so it uses plain
mlocklike stream.cpp. that's +45 ms (warm) / ~0.23 s (cold) one-time load cost.
untested here (rules: no full-model runs, :8080 is live): the actual mtp_ms recovery on 2K/10K prompts. needs a parent-side A/B run with STRATA_METAL_MLOCK_DENSE=1 vs 0.
trade-off to watch: mtp_load runs before stream_arena (window_init -> mtp_load_pending -> prompt_gemm_init / Window::run -> stream_arena), and avail_now() doesn't count wired pages. so on the 24 GB Air the RAM-bound arena budget should come out up to ~0.71 GB smaller than before. before, the arena was effectively "borrowing" those clean MTP pages, and that's exactly the eviction behind the stall. check the arena GB / hit rate line in the A/B too. with MLOCK_DENSE=0 the MTP pages are not added to am.dense_file (stream's dense_unwired_), same as before the patch (that's outside my lines).
2. qsa.metal #
- qsa_attn softmax: a lane holds cells
laneandlane+32in registers.simd_max(max(s0,s1))is the same exact max;p0 + p1(with p1 = 0 when absent) equals the old0 + p0 (+ p1)sequence bit for bit, sosimd_sumsees identical inputs. - qsa_attn V loop: unrolled by 4, so the four
vq/vscloads are issued before the FMAs, and the per-head FMA chain keeps the order c, c+1, c+2, c+3 (plus a scalar tail). - qsa_write_prepared:
q2s/p4smoved to program-scopeconstantarrays (qsa_q2s/qsa_p4s). MSL rejects function-scopeconstant("automatic variable qualified with an address space").#pragma clang loop unroll(full)added on the shuffle loop. - qsa_merge:
tid < ncthreads load the chunk maxima intomk[64], then every thread takes M = max over mk in k order.tid < ncthreads computeek[k] = exp(mk[k] - M)once, andtid 0sumsL = fma(l_k, ek[k], L)in k order intoL_sh. all 256 threads then doo = fma(*q, ek[k], o)with a running pointer, followed by a barrier ando / L_sh. same values and same order as before, so it's bitwise identical.
3. verification #
- new
qsa_parity --bitwise REF.metalmode: compiles qsa_attn/qsa_merge from this tree and from REF (the pre-editqsa.metal, saved at/tmp/opus_r20_mtp_qsa/qsa_base.metal) on random int8 KV, fp16 scales, q/gate and top-k-shaped ids (32768-cell KV). it memcmps the wholepartandxpbuffers across pos0 ∈ {0,1,63,64,512,2047,2048,2050,8192,30000} × ntok ∈ {1,2,4,8}.- result: 0 of 40 cases differ, PASS, in 4 separate runs, ~1.4 s each.
- sensitivity check: against a mutant (L summed without fma, ×1.0000001) it exits 1, so the harness does catch a one-ulp change.
qsa_parity(default topk2) andQSA_TOPK=1 qsa_parity, P=3000: ALL PASS (attention + gate rel L2 2.44e-07, unchanged).qsa_parity 1000fails on "indexer q 1.00e+00". that failure is pre-existing: the untouchedvendor/strata/build-metal/qsa_parity 1000fails the same way (it's the prep stage, outside my lines).arena_budget_test: all passed.
4. timings (M5 Air, GPU time per dispatch, qsa_parity --bench 3000, QSA_TOPK=1, two alternating runs each, old = vendor build's binary on the same pre-edit kernels) #
| old | new | ||
|---|---|---|---|
| ntok 1 attn | 58.9 / 59.2 us | 46.5 / 46.6 us | -21% |
| ntok 1 merge | 7.92 us | 7.70 us | -3% |
| ntok 1 window (5 dispatches) | 79.7 us | 66.4 us | -17% |
| ntok 3 attn | 156.2 us | 125.1 us | -20% |
| ntok 3 merge | 19.8 us | 17.9 us | -10% |
| ntok 3 window | 191.9 us | 156.9 us | -18% |
The --bitwise paired numbers (random scattered ids, noisier since the GPU is shared with the live server) agree. attn at ntok 1 went 59 -> 47 us at pos0 2048/8192; ntok 4 was 263 -> 164 us in 2 of 3 runs, and in one run pos0=8192 measured 205/236. merge at ntok 4 went 28.5 -> 25.0 us.
Per decode window that's about 13 us × the QSA layers per token. merge was a small kernel to begin with, so most of the gain comes from attn.