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 from Wm_ are wired rather than dropped (correct by construction, since MB() routes on mq8t_.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_proj 235.9 MB + experts.gate_up_proj 471.9 MB = 707.8 MB; released 181.2 MB (all BF16 now in mq8_, the 9 F32 norms, the indexer).
  • direct mlock vs WILLNEED + 8-thread parallel touch + mlock:
    • warm page cache: 42-46 ms vs 43-49 ms
    • cold (fresh ditto --nocache copy 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 mlock like 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 lane and lane+32 in registers. simd_max(max(s0,s1)) is the same exact max; p0 + p1 (with p1 = 0 when absent) equals the old 0 + p0 (+ p1) sequence bit for bit, so simd_sum sees identical inputs.
  • qsa_attn V loop: unrolled by 4, so the four vq/vsc loads 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/p4s moved to program-scope constant arrays (qsa_q2s/qsa_p4s). MSL rejects function-scope constant ("automatic variable qualified with an address space"). #pragma clang loop unroll(full) added on the shuffle loop.
  • qsa_merge: tid < nc threads load the chunk maxima into mk[64], then every thread takes M = max over mk in k order. tid < nc threads compute ek[k] = exp(mk[k] - M) once, and tid 0 sums L = fma(l_k, ek[k], L) in k order into L_sh. all 256 threads then do o = fma(*q, ek[k], o) with a running pointer, followed by a barrier and o / L_sh. same values and same order as before, so it's bitwise identical.

3. verification #

  • new qsa_parity --bitwise REF.metal mode: compiles qsa_attn/qsa_merge from this tree and from REF (the pre-edit qsa.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 whole part and xp buffers 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) and QSA_TOPK=1 qsa_parity, P=3000: ALL PASS (attention + gate rel L2 2.44e-07, unchanged).
  • qsa_parity 1000 fails on "indexer q 1.00e+00". that failure is pre-existing: the untouched vendor/strata/build-metal/qsa_parity 1000 fails 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.