round 19 track b: sample / mtp chain #

result #

keep the mtp_chain.metal changes; leave sample.metal unchanged. the argmax ideas pass exact parity but do not establish a useful speedup on this m5. mtp pick and active slop penalties do.

  • patch: /tmp/sol_r19_sample/sol_r19_sample.patch
  • sha256: 5c4620d311b3e4035081d33e07a62a228785d7d6291363cbabf26408ea577799
  • git -C vendor/strata apply --check /tmp/sol_r19_sample/sol_r19_sample.patch: exit 0.
  • scratch git diff --check: exit 0.
  • patch touches only metal/src/engine/msl/mtp_chain.metal and metal/tests/sample_parity.cpp. no tracked vendor files were edited. no server requests, port access, full-model benchmarks, or subagents.

contract and setup #

question: can shared-reduction gating, loop unrolling, aligned vector loads, or a redundant embedding barrier be removed/changed without changing output bits?

baseline is the current working tree, not clean upstream: vendor head 296b59d2f73c8f04a805df29c2fa97655adace78 had pre-existing changes. the three original owned files were copied before editing into /tmp/sol_r19_sample/baseline/; the patch is relative to those exact contents. all work happened under /tmp/sol_r19_sample/strata. source snapshots, linked static runtime, executable and sha256 manifest are preserved there.

owner stays the existing sampling/draft-chain kernels. no interface, dispatch shape, allocation, state ownership or floating-point accumulation order changes. the compile-time embedding guard removes only empty table loaders and their barrier; table-backed formats still synchronize.

machine: apple m5, macos 27.0 build 26a428; apple clang 21.0.0 (clang-2100.3.27.1). runtime metal safe math; STRATA_METAL_MATH unset. runtime and quantization cpu helpers were statically linked from the existing build-metal/libmtl_rt.a, now preserved as /tmp/sol_r19_sample/libmtl_rt.a.

predeclared rule: exact parity mandatory; require repeatable >3% median speedup for performance-only edits. accept redundant-barrier removal if parity passes and no repeatable >3% regression. five exploratory comparisons then five confirmation comparisons. each timing is seven alternating-order baseline/candidate batches, with a 5ms warmup for each pipeline. sampling/pick/pen batches use 32 repetitions. tests share a gpu, so benchmark processes ran sequentially. no attempt was made to stop or modify the live server; this is not a dedicated-machine or full-model measurement.

retained changes #

  • mtp_slop_pen: unroll_count(4) on the two vocabulary scans, full unrolling of both eight-element threadgroup reductions.
  • mtp_pick: unroll_count(4) on the softmax sum scan.
  • mtp_embed: table initialization/barrier only when fmt<F>::TAB > 1; full unrolling of the per-unit output loop.
  • sample_parity: paired-source comparison plus cpu partial-argmax checks, the requested vocab/window shapes, exceptional logits, draft pick/penalty/state checks, and all 13 exported embedding formats.

important syntax correction: #pragma clang loop unroll(4) is invalid msl on this compiler. both original attempted variants failed at compile time, with “expected enable, full or disable.” #pragma clang loop unroll_count(4) is the accepted numeric-count form. those execution failures are preserved in sample_gate_unroll_1.log / mtp_all_1.log; they are not measurements.

measured results #

medians below are across five final invocations (mtp_long_embed_rep0.log through rep4.log). ratios are medians of each invocation's candidate/baseline ratio, not ratios of the independently aggregated times. times are microseconds per dispatch.

kernel / shape baseline us candidate us candidate / baseline
mtp_pick nv=40525 5.062 3.547 0.7274
mtp_pick nv=248320 24.542 15.055 0.6127
mtp_slop_pen nv=40525 mode=0 25.115 25.055 1.0044
mtp_slop_pen nv=40525 mode=1 45.341 34.977 0.7718
mtp_slop_pen nv=40525 mode=2 55.466 41.454 0.7470
mtp_slop_pen nv=248320 mode=0 24.987 25.043 1.0021
mtp_slop_pen nv=248320 mode=1 123.389 57.797 0.4692
mtp_slop_pen nv=248320 mode=2 217.988 87.898 0.4035
mtp_embed format=8 1.604 1.620 1.0096
mtp_embed format=12 2.634 1.914 0.7266
mtp_embed format=13 2.664 1.919 0.7209

slop modes: 0 = no penalties; 1 = penalties present but hold correction inactive; 2 = hold correction active. tiny timed penalties keep the same branch active across repetitions. parity separately uses normal-sized penalties and resets state/logits between the two pipelines.

pick is about 27–39% faster; active slop paths about 23–60% faster. no-penalty slop is unchanged. final slop ratios at full vocab are tight: mode 1 0.4676–0.4716, mode 2 0.4028–0.4102. the first final invocation was cold/noisy for the draft vocabulary (including a mode-0 outlier ratio 1.8285); five earlier confirmation invocations independently show the same active-path wins. see summary.json for all observed ratio ranges, not only these medians.

embedding q4_k/q5_k is about 27–28% faster. q8_0, the live model's embedding format, has no established speedup: final ratio 1.0096 (range 1.0049–1.0156). other formats' final medians range 0.9774–1.0151 except q4_0 at 1.0113 and q3_k at 1.0112; none establishes a >3% regression. barrier-only and unit-unroll-only ablations (mtp_barrier_rep*.log, mtp_unit_unroll_rep*.log) identify unit unrolling as the q4_k/q5_k win; the barrier change is structural, not a measured q8 gain.

measurement repair: initial embedding batches of 32 showed roughly 0.1us alternating timing noise, enough to create apparent 6–7% shifts on 1.6us kernels. final embedding batches use 1024 repetitions; five new baseline/candidate comparisons under the repaired evaluator remove that ambiguity. sampling/pick/pen evaluator batches are unchanged. all final invocations still finished below 1.5s. earlier embedding timing numbers are exploratory, not the final evidence.

rejected argmax experiments #

all candidates passed memcmp on pv/pi/out and cpu-reference checks. five runs each; ratio medians for the requested shapes:

candidate 248320 x 1 x 2 x 4 40525 x 1
shuffle full unroll 1.0000 1.0010 0.9996 1.0012
thread-zero final reduction + shared-loop full unroll 1.0095 1.0172 1.0006 1.0194
part-loop unroll only 0.9734 0.9786 1.0029 1.0716
gate + shuffle unroll 1.0174 1.0191 1.0260 1.0175
gate + shuffle + part unroll 1.0150 0.9859 1.0308 1.0702
gate + shuffle + aligned float4/fallback 0.9989 1.0272 0.9972 1.0833

vector loads were guarded by alignment of the actual token-row pointer plus partition offset, not just partition offset; scalar fallbacks covered unaligned partitions and vector tails. none meets the decision rule. do not apply an argmax directive just because it looks cheaper in source. source variants and raw logs remain in /tmp/sol_r19_sample/.

correctness and final checks #

  • final argmax test: 70 windows, 0 wrong, vocab 248320 x every ntok 1–8, 40525 x 1, and 7 x 1 (empty partials). checks all partial values bitwise, partial ids and final ids against both baseline and cpu. includes all nan, all negative infinity, positive infinity ties, signed-zero ties, mixed nan/finite, and finite exact ties.
  • final mtp test: 0 wrong in each of five final runs. 16 pick cases (both mapping modes; all four output slots; finite, negative infinity, nan), 20 penalty/state cases (j=0 and continued j=1; no-op and active paths; finite, all negative infinity, all nan, positive infinity, tied zero logits), 39 embedding rows (three rows x all 13 exported formats, width 2560). all compared with memcmp to the baseline gpu kernels. finite active penalty cases additionally assert logits really changed.
  • self-compare smoke tests: default argmax and --mtp both exit 0. final baseline argmax comparison took 1.204s, mtp self-compare 0.821s; final paired mtp runs 0.884–1.402s. no invocation reached the 4.9s timeout.
  • final native build exit 0. an optional -Wall -Wextra build found only the pre-existing logical-parentheses warning in slop4.h:194, outside owned files.
  • final visual diff/context review done; duplicate penalty-dispatch construction consolidated in a local helper.

initial penalty timing tests omitted slop4::KNOWN, so their j=0 paths did not exercise penalties. this was caught by reading slop4::penalties, repaired before selection/confirmation, and active-output assertions added. mtp_all_rep*.log is superseded by mtp_all_fixed_rep*.log and the confirmation/final runs. do not use the former as active-penalty evidence.

these are baseline-equivalence smoke tests, not exhaustive proofs over every history or format bit pattern. no end-to-end token/s claim; that measurement was explicitly out of scope.

reproduce #

from the project root (commands shown for a shell; the experiments used equivalent direct run() argv):

clang++ -std=c++20 -O2 \
  -I/tmp/sol_r19_sample/strata/metal/src \
  -I/tmp/sol_r19_sample/strata/third_party/ggml \
  /tmp/sol_r19_sample/strata/metal/tests/sample_parity.cpp \
  /tmp/sol_r19_sample/libmtl_rt.a \
  -framework Metal -framework Foundation -o /tmp/sol_r19_sample/sample_parity
STRATA_METAL_KERNELS=/tmp/sol_r19_sample/strata/metal/kernels \
  /tmp/sol_r19_sample/sample_parity \
  /tmp/sol_r19_sample/baseline/metal/kernels/sample.metal
STRATA_METAL_KERNELS=/tmp/sol_r19_sample/strata/metal/kernels \
  /tmp/sol_r19_sample/sample_parity --mtp \
  /tmp/sol_r19_sample/baseline/metal/src/engine/msl/mtp_chain.metal

baseline/final source hashes: /tmp/sol_r19_sample/manifest.json. final measured sources: /tmp/sol_r19_sample/final/. raw logs, variant sources, executable and summary: /tmp/sol_r19_sample/. stop reason: measured wins confirmed; argmax negative results preserved, no further gpu or model work authorized/needed for this track.