From 1722f5a91dc438e6b4557ea0c2dcbeeee3111817 Mon Sep 17 00:00:00 2001 From: Boris Batkin Date: Mon, 27 Jul 2026 11:20:38 -0700 Subject: [PATCH 01/18] lint alignment: complete the @scratch escapes, enable [hot_path] on the whole Metal surface MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The metal rails could not adopt the #3580 hot-path contracts: every kernel dispatch pushes capture/hazard records, and PERF026's escapes could not describe why that is fine. Four checker gaps, each with a fixture case: - @scratch on MODULE GLOBALS (the clear()-recycled capture-rail shape; `var @scratch g : array` — the annotation parses after `var`) via an ExprVar arm in is_scratch_dest, which also FOLLOWS LOCAL REFERENCE BINDINGS (`var lst & = pool.free_bufs[b]` carries the mark). - `pop` joins the sizing set (das pop lowers through resize; a shrink of a reused buffer is a store). - indexing a @scratch TABLE is exempt (the pool / residency-cache shape: insert on first touch, steady-state [] is a hit). - chain suppression honors a nolint ANYWHERE along the chain — the sink often bottoms in daslib/builtin.das, so the honest line is an intermediate call site in a required module; that line now works. dasMetal: metal_command_buffer_failed(cb) : bool — the hot-path probe (status only; the error-string extern allocates its message, success paths should not fetch a string to test emptiness). free_bufs / pipeline cache marked @scratch. Metal surface: [hot_path] on all six roots (upgrading #3580's [no_env, no_io] "no_alloc pending" markers — this is that audit) — metal_decode_forward, metal_batch_decode_forward, metal_prefill_forward, metal_mtp_draft_forward, metal_mtp_seam_row, metal_mtp_spec_eval. 801 PERF026 findings driven to ZERO: - ~20 @scratch marks on genuinely recycled containers (capture/hazard rails, buffer pools, residency + PSO + schedule caches, KV/DN mirrors, per-step staging, the capped trace ring, session rope/mtp scratch). - [cold_path] on the decline recorders, live-reload flushes, arena rebuild/defrag, schedule compile, and the signs-table lazy build (split out of the hot accessor). - six error-check sites converted to the bool probe with a [cold_path] cb_error_text fetch on the failure leg — the per-step string alloc is gone for real. - prefill per-call transients (bks/bvs/cbs) hoisted to reused @scratch globals through reference bindings (their end-of-scope deletes dropped). - parallel_argmax partials became fixed arrays (ARGMAX_MAX_CHUNKS=64 caps chunking, not correctness) — no heap on that path at all. - three reasoned nolints (the batch cb list: two steps' lists are alive at once, a reused global cannot serve). Gates: perf fixtures (expect 7, cases 9-12), tests/metal/test_metal_sdot, dasLLAMA test_kernel_backend / test_forward / test_prefill / test_batch_decode / test_mtp -jit, and the metal decode + prefill parity suites (llama arm, DASLLAMA_METAL_HAZARD_STRICT=1) — all green. Co-Authored-By: Claude Fable 5 --- daslib/perf_lint.das | 35 ++++++-- doc/source/reference/language/lint.rst | 7 +- modules/dasLLAMA/dasllama/dasllama_common.das | 6 +- modules/dasLLAMA/dasllama/dasllama_math.das | 14 ++-- .../dasllama/dasllama_metal_common.das | 81 ++++++++++++------- .../dasllama/dasllama_metal_llama.das | 48 ++++++----- .../dasllama/dasllama_metal_prefill.das | 30 +++---- modules/dasMetal/metal/das_metal_boost.das | 4 +- modules/dasMetal/src/dasMetal.h | 1 + modules/dasMetal/src/dasMetal.mm | 13 +++ skills/perf_lint.md | 21 +++-- utils/lint/tests/perf026_hot_path_alloc.das | 54 ++++++++++++- 12 files changed, 215 insertions(+), 99 deletions(-) diff --git a/daslib/perf_lint.das b/daslib/perf_lint.das index 950ac91213..112a856751 100644 --- a/daslib/perf_lint.das +++ b/daslib/perf_lint.das @@ -189,12 +189,12 @@ def private is_alloc_builtin(n : string) : bool { def private is_container_sizing(n : string) : bool { return (n == "resize" || n == "resize_no_init" || n == "reserve" || n == "clear" || n == "push" || n == "push_clone" || n == "emplace" || n == "push_from" - || n == "push_clone_from" || n == "insert" || n == "erase") + || n == "push_clone_from" || n == "insert" || n == "erase" || n == "pop") } -// True when `expr` reaches a struct field marked @scratch — walking through indexing and deref, -// so `p.kblobs[b]` counts by virtue of kblobs. Sizing such a buffer to the step's geometry is the -// engine's own allocation strategy, declared once at the field, not an accident at each call site. +// True when `expr` reaches a declaration marked @scratch — a struct field or a module global +// (`var @scratch g : array`, the clear()-recycled capture-rail shape) — walking through +// indexing and deref, so `p.kblobs[b]` counts by virtue of kblobs. def private is_scratch_dest(var expr : Expression?) : bool { return false if (expr == null) if (expr is ExprField) { @@ -207,6 +207,19 @@ def private is_scratch_dest(var expr : Expression?) : bool { } return is_scratch_dest(f.value) } + if (expr is ExprVar) { + var v = (expr as ExprVar).variable + if (v != null) { + for (a in v.annotation) { + if (a.name == "scratch") return true + } + // a ref binding (`var lst & = pool.free_bufs[b]`) carries its destination's scratch-ness + if (v._type != null && v._type.flags.ref && v.init != null) { + return is_scratch_dest(v.init) + } + } + return false + } if (expr is ExprRef2Value) return is_scratch_dest((expr as ExprRef2Value.subexpr)) if (expr is ExprAt) return is_scratch_dest((expr as ExprAt.subexpr)) if (expr is ExprSafeAt) return is_scratch_dest((expr as ExprSafeAt.subexpr)) @@ -320,10 +333,11 @@ class HotBodyScan : AstVisitor { } // table[key] inserts a default entry when the key is missing, so it allocates on READ too; - // the ?[] form (ExprSafeAt) has its own visitor hook and never inserts. + // the ?[] form (ExprSafeAt) has its own visitor hook and never inserts. A @scratch table + // (the pool / residency-cache shape) is declared reuse — indexing it is the owner's strategy. def override preVisitExprAt(var expr : ExprAt?) : void { if (expr.subexpr == null || expr.subexpr._type == null) return - if (expr.subexpr._type.baseType == Type.tTable) { + if (expr.subexpr._type.baseType == Type.tTable && !is_scratch_dest(expr.subexpr)) { add_sink(HotBits.alloc, "table index (inserts a default entry when the key is missing)", expr) } } @@ -2083,9 +2097,16 @@ class PerfLintVisitor : AstVisitor { def hot_report(frames : array; fi : int; sink : HotSink; want : HotBits; root_file : string) { // cast each side: AOT emits bitfield & bitfield as a raw C++ &, ambiguous when one is const return if ((uint(sink.bit) & uint(want)) == 0u) - // honor a nolint at EITHER end: a report can be anchored in a different file than the sink + // honor a nolint ANYWHERE along the chain — the sink may bottom in daslib, so the honest line is often an intermediate call site let code = sink.bit.alloc ? "PERF026" : (sink.bit.env ? "PERF027" : "PERF028") return if (is_lint_suppressed(sink.site.at, code)) + var si = fi + while (si >= 0) { + if (frames[si].site != null && is_lint_suppressed(frames[si].site.at, code)) { + return + } + si = frames[si].parent + } // sink's own location, unless another file or a SHARED node — then the caller-side anchor var at = sink.site.at let shared_node = frames[fi].fn.flags.generated || frames[fi].fn.fromGeneric != null diff --git a/doc/source/reference/language/lint.rst b/doc/source/reference/language/lint.rst index c74b3f5338..18f2f606a0 100644 --- a/doc/source/reference/language/lint.rst +++ b/doc/source/reference/language/lint.rst @@ -1399,6 +1399,9 @@ every call site: // a helper that sizes a caller's buffer marks the PARAMETER, since the // destination arrives by reference and the call site cannot see the field def scratch_resize(@scratch var a : array; need : int64) { ... } + + // a clear()-recycled module global is declared the same way (annotation after `var`) + var @scratch g_stage : array **What the scan deliberately ignores.** Arguments to ``panic(...)`` — a panic is fatal in daslang, not an exception, so its interpolated message is on the abort @@ -1408,7 +1411,9 @@ expands into. And indirect calls through a function pointer or lambda, which cannot be resolved statically — annotate the implementations they reach. Escape hatches, in order of preference: ``[cold_path]`` on the callee when the -leg genuinely runs once; ``@scratch`` on a reused destination; ``// nolint`` +leg genuinely runs once; ``@scratch`` on a reused destination (field, by-ref +parameter, or module global — sizing calls, table indexing, and reference +bindings to it all count); ``// nolint`` with a reason (honored at either end of a chain, so a suppression written where the code lives works even when the report anchors elsewhere); and ``DAS_LINT_DISABLE=PERF028`` for a whole run, which needs no source edit and is diff --git a/modules/dasLLAMA/dasllama/dasllama_common.das b/modules/dasLLAMA/dasllama/dasllama_common.das index 25eb878bbb..ca91cf876b 100644 --- a/modules/dasLLAMA/dasllama/dasllama_common.das +++ b/modules/dasLLAMA/dasllama/dasllama_common.das @@ -4294,7 +4294,7 @@ struct Session { mtp_h_pos1 : int64 mtp_cat : array // eh_proj input [enorm(embed(tok)) ; hnorm(h)] (2*dim) mtp_cat_b : array // prompt-warm eh_proj concat image ((npos-1) x 2*dim, scratch-sized per warm) - mtp_xb_save : array // prompt-warm x_b preservation (the warm's eh_proj output rides x_b for the arch blocks; consumers like mtp_verify_row0 / embed_forward need the trunk hiddens back) + @scratch mtp_xb_save : array // prompt-warm x_b preservation (the warm's eh_proj output rides x_b for the arch blocks; consumers like mtp_verify_row0 / embed_forward need the trunk hiddens back) mtp_logits : array // verify-batch row-0 logits (vocab; the draft check reads its argmax) mtp_vbatch : array // the 2-row verify batch [tok, draft] (persistent — no per-step alloc) mtp_norm_b : array // verify-batch final-normed rows (npos x dim), input to the batched classifier @@ -4731,7 +4731,7 @@ def private rope_batch(var q_b : array; var k_b : array; t : Model // Precompute the prefill RoPE cos/sin table once: cos_tab[pi*half + j] = cos((start_pos+pi)*freq[j]), // sin likewise. freq[j] is position-independent (hoisted out of the position loop); mscale folds // in here. Built once before the layer loop, reused across all heads/layers. Bit-identical to classic. -def private build_rope_table(var cos_tab : array; var sin_tab : array; t : Model; theta, fscale : float; start_pos, npos, head_size : int64; use_ff : bool = true) { +def private build_rope_table(@scratch var cos_tab : array; @scratch var sin_tab : array; t : Model; theta, fscale : float; start_pos, npos, head_size : int64; use_ff : bool = true) { let c = t.config let half = head_size / 2l let scale = c.rope_mscale // scales cos/sin — ggml applies mscale to NORM and NEOX alike (mistral3 YaRN) @@ -4759,7 +4759,7 @@ def private build_rope_table(var cos_tab : array; var sin_tab : array; var sin_tab : array; t : Model; +def private build_rope_table_rows(@scratch var cos_tab : array; @scratch var sin_tab : array; t : Model; theta, fscale : float; positions : array; nrows, head_size : int64; use_ff : bool = true) { let c = t.config diff --git a/modules/dasLLAMA/dasllama/dasllama_math.das b/modules/dasLLAMA/dasllama/dasllama_math.das index 316c24becb..98287e1a8c 100644 --- a/modules/dasLLAMA/dasllama/dasllama_math.das +++ b/modules/dasLLAMA/dasllama/dasllama_math.das @@ -3591,11 +3591,13 @@ def exp4(x : float4) : float4 { //! Index of the max element, first-index-wins on ties — the greedy-decode argmax. Chunked team //! scan (serial 151-262k-vocab scan idled every lane 210-284us/token); combine folds candidates //! in index order with strict >, so the result is bit-identical to the serial loop for ANY split. +let ARGMAX_MAX_CHUNKS = 64 // partials cap: enough for any current box's lanes*2 + def public parallel_argmax(x : float const?; n : int64) : int64 { if (n <= 1l) { return 0l } - let nch = int(min(int64(get_dispatch_lanes()) * 2l, n / 4096l)) + let nch = int(min(min(int64(get_dispatch_lanes()) * 2l, n / 4096l), int64(ARGMAX_MAX_CHUNKS))) unsafe { if (nch <= 1) { var best = 0l @@ -3606,10 +3608,10 @@ def public parallel_argmax(x : float const?; n : int64) : int64 { } return best } - var cm : array - var ci : array - cm |> resize(nch) - ci |> resize(nch) + // fixed-size worker partials: no heap on the hot path (fork-pool context clones forbid + // a global scratch); ARGMAX_MAX_CHUNKS caps chunking, not correctness + var cm : float[ARGMAX_MAX_CHUNKS] + var ci : int64[ARGMAX_MAX_CHUNKS] var cmp = addr(cm[0]) var cip = addr(ci[0]) let nn = n @@ -3638,8 +3640,6 @@ def public parallel_argmax(x : float const?; n : int64) : int64 { best = ci[c] } } - delete cm - delete ci return best } } diff --git a/modules/dasLLAMA/dasllama/dasllama_metal_common.das b/modules/dasLLAMA/dasllama/dasllama_metal_common.das index 085917eada..a84b08cd56 100644 --- a/modules/dasLLAMA/dasllama/dasllama_metal_common.das +++ b/modules/dasLLAMA/dasllama/dasllama_metal_common.das @@ -142,7 +142,7 @@ def record_needs(var tab : table; missing : MetalNeed) { var g_dev : MetalDevice? var g_queue : MetalCommandQueue? var g_failed = false -var g_regions : table // base address -> resident weight/scale/norm region +var @scratch g_regions : table // base address -> resident region; upload-once, steady-state [] is a hit var g_pool : MetalBufferPool var g_upool : MetalBufferPool // UNTRACKED pool: uniforms + rope rows (GPU-read-only) var g_decodes = 0l @@ -188,7 +188,7 @@ struct MetalDecodeStepTrace { ge : double } -var g_step_trace : array +var @scratch g_step_trace : array // capacity-capped trace ring (g_step_trace_cap bounds it) var g_step_trace_cap = 0l var g_step_trace_t0 = 0l @@ -207,7 +207,7 @@ struct KVMirror { bytes : uint64 // slice size (k + v planes) } -var g_mirrors : table +var @scratch g_mirrors : table // per-session mirrors; hot [] after key_exists is a hit var g_mirror_bytes = 0ul var g_tick = 0l var g_mirror_cap_mb = -1 // <0 = unset: DASLLAMA_METAL_KV_MIRROR_MB (default 4096) reads at use @@ -218,7 +218,7 @@ var g_mirror_cap_mb = -1 // <0 = unset: DASLLAMA_METAL_KV_MIRROR_MB (default 4 var g_arena : MetalBuffer? var g_arena_bytes = 0ul var g_arena_epoch = 0l // bumps on every grow-rebuild — pre-encoded steps referencing the old buffer are stale -var g_arena_free : array> // sorted by off, coalesced +var @scratch g_arena_free : array> // sorted by off, coalesced // ===== decode PSO + resource globals (declared here, not kernels.das, so [metal_dispatch]-generated // builders — co-located with their kernel class, some in prefill.das — see the pipeline they @@ -392,6 +392,11 @@ def require_or_panic(kind : string; why : string) { } } +// the failure leg's message fetch — hot code checks metal_command_buffer_failed (no alloc) +[cold_path] +def cb_error_text(cb : MetalCommandBuffer?) : string => metal_command_buffer_error(cb) + +[cold_path] // the decline leg IS the fall-off-the-GPU path; its bookkeeping is not hot def decline(why : MetalDecodeDecline) { g_declines++ g_declines_by_reason["{why}"]++ @@ -399,6 +404,7 @@ def decline(why : MetalDecodeDecline) { } // the `feature` flavor: records a need_* key per missing bit, and the required-mode panic names them +[cold_path] // the decline leg IS the fall-off-the-GPU path; its bookkeeping is not hot def decline(why : MetalDecodeDecline; missing : MetalNeed) { g_declines++ g_declines_by_reason["{why}"]++ @@ -406,12 +412,14 @@ def decline(why : MetalDecodeDecline; missing : MetalNeed) { require_or_panic("decode", "{why} {missing}") } +[cold_path] // the decline leg IS the fall-off-the-GPU path; its bookkeeping is not hot def batch_decline(why : MetalDecodeDecline) { g_batch_declines++ g_batch_declines_by_reason["{why}"]++ require_or_panic("batch decode", "{why}") } +[cold_path] // the decline leg IS the fall-off-the-GPU path; its bookkeeping is not hot def batch_decline(why : MetalDecodeDecline; missing : MetalNeed) { g_batch_declines++ g_batch_declines_by_reason["{why}"]++ @@ -724,7 +732,7 @@ def kq_load_gpu_supported(t : Model) : bool { // ===== zero-copy logits: bytesNoCopy wrappers over session-owned storage ===== // Writes land straight into s.logits (no landing memcpy); borrowed, no deallocator. Gated on // 16KB page alignment + page-rounded capacity; the pointer is revalidated on every lookup. -var g_nc_logits : table> +var @scratch g_nc_logits : table> def nc_logits_for(dev : MetalDevice?; s : Session) : MetalBuffer? { let p : uint64 = intptr(unsafe(addr(s.logits[0]))) @@ -912,6 +920,7 @@ def private mirror_ceiling_bytes : uint64 { // release-all rebuild at `want` (clamped to the ceiling): every mirror re-uploads via its // watermark, and the epoch bump invalidates any pre-encoded step whose kernels reference the // OLD arena buffer. Callers must ensure nothing is pinned and no step is in flight. +[cold_path] // arena (re)construction — the grow event, not the per-step path def private arena_rebuild(want : uint64; ceil_bytes : uint64) { for (mr in values(g_mirrors)) { mirror_release(mr) @@ -974,6 +983,7 @@ def mirror_reserve(total : uint64) : bool { // release-all rebuild sized for `total` — the batch driver's fragmentation fallback when the // per-session prepares can't pack a step that DOES fit capacity-wise (exact-fit demand around // older residents can strand contiguity). Callers unpin first; false = past the ceiling. +[cold_path] // arena defrag = the occasional rebuild leg def mirror_defrag(total : uint64) : bool { let ceil_bytes = mirror_ceiling_bytes() if (total > ceil_bytes || !empty(g_pin) || g_arena == null) { @@ -1110,7 +1120,7 @@ struct DnMirror { pos : int64 // next position the GPU state expects (forward-only) tick : int64 } -var g_dn_mirrors : table +var @scratch g_dn_mirrors : table // per-session; hot [] after key_exists is a hit let DN_MIRRORS_MAX = 4 // per-session state is big (35B: ~60MB) — oldest evicts; a returning // evicted session declines unless it restarts at pos 0 or off CPU state @@ -1201,15 +1211,20 @@ def dn_mirror_flip(uid : uint64; pos : int64) { } } +[cold_path] // one-time: fills the 512B FWHT sign table on first touch +def private signs_buffer_build() { + var signs : array + tq4_fill_signs(signs, 128l) + g_signs = metal_new_buffer(g_dev, 512ul) + unsafe { + memcpy(metal_buffer_contents(g_signs), addr < void? >(signs[0]), 512) + } + delete signs +} + def signs_buffer() : MetalBuffer? { if (g_signs == null) { - var signs : array - tq4_fill_signs(signs, 128l) - g_signs = metal_new_buffer(g_dev, 512ul) - unsafe { - memcpy(metal_buffer_contents(g_signs), addr < void? >(signs[0]), 512) - } - delete signs + signs_buffer_build() } return g_signs } @@ -1346,8 +1361,8 @@ var private g_pipe_debug = false // DASLLAMA_METAL_PIPE_DEBUG, resolved once var private g_mirror_cap_env = 4096 // DASLLAMA_METAL_KV_MIRROR_MB, resolved at init; the default stands until then var g_hz_strict = false // DASLLAMA_METAL_HAZARD_STRICT: undeclared dispatch panics var g_hz_dirty = false // an undeclared dispatch ran — the next gate must barrier -var g_hz_stage : array // the next dispatch's declared ranges -var g_hz_pend : array // unbarriered ranges live on the encoder +var @scratch g_hz_stage : array // the next dispatch's declared ranges (clear()-recycled) +var @scratch g_hz_pend : array // unbarriered ranges live on the encoder (clear()-recycled) var g_hz_barriers = 0l // lifetime barrier count (chase prints the per-step delta) var g_hz_gated = 0l // lifetime gated dispatches — the coverage check's twin var g_hz_undeclared = 0l // conservative fallbacks taken (0 once migration completes) @@ -1454,9 +1469,9 @@ struct KBind { } var g_gr_on = false -var g_gr_nodes : array -var g_gr_binds : array -var g_gr_ranges : array +var @scratch g_gr_nodes : array // capture-rail arrays: grown once, clear()-recycled per capture +var @scratch g_gr_binds : array +var @scratch g_gr_ranges : array var g_gr_open = false // a node is being built (kn_pipeline .. kn_dispatch) def kn_pipeline(enc : MetalComputeEncoder?; pso : MetalComputePipeline?) { @@ -1546,7 +1561,7 @@ struct KSched { nodes : int } -var g_gr_scheds : table +var @scratch g_gr_scheds : table var g_gr_sched_mode = -1 // DASLLAMA_METAL_SCHED: 1 = leveled reorder (default), 0 = capture order [cold_path] // one-time env resolution; sched_enabled is a global test after the first call @@ -1576,6 +1591,7 @@ def private gr_shape_hash : uint64 { // level[j] = 1 + max level of any conflicting predecessor; scan partitioned per buffer id. // An UNDECLARED node has no edges and would float freely — refuse and keep capture order. +[cold_path] // schedule compile — once per shape class, inside the g_gr_scheds miss def private gr_schedule(var sched : KSched) : bool { let n = length(g_gr_nodes) sched.nodes = n @@ -1710,27 +1726,27 @@ def public set_metal_decode_hazard(paranoid, strict : bool) { var g_lp_active = false var g_lp_cbs : array -var g_lp_pooled : array> -var g_lp_upooled : array> +var @scratch g_lp_pooled : array> // pending-step stash; clear()-recycled per finish +var @scratch g_lp_upooled : array> var g_lp_blog : MetalBuffer? var g_lp_xb : MetalBuffer? // production sync + CPU-classifier path: x_b writeback var g_lp_xb_dst : void? // ... into the caller's workspace (borrowed, sync-only) var g_lp_xb_bytes = 0l var g_lp_err = false // sticky per-step dispatch-error signal for sync callers -var g_lp_sessions : array -var g_lp_positions : array -var g_lp_koffs : array -var g_lp_voffs : array -var g_lp_caps : array +var @scratch g_lp_sessions : array +var @scratch g_lp_positions : array +var @scratch g_lp_koffs : array +var @scratch g_lp_voffs : array +var @scratch g_lp_caps : array var g_lp_nrows = 0l var g_lp_layers = 0l var g_lp_seq = 0l var g_lp_vocab = 0l // hetero (gemma4): per-layer row bytes + slab bases within a mirror plane (uniform models: // lbase_l == l * cap * rowb_l — captured at encode, one array per step) -var g_lp_lrowb : array -var g_lp_lkbase : array // per (row i, layer l) flattened i * n_layers + l (caps differ per row) -var g_lp_lvbase : array +var @scratch g_lp_lrowb : array +var @scratch g_lp_lkbase : array // per (row i, layer l) flattened i * n_layers + l (caps differ per row) +var @scratch g_lp_lvbase : array var g_lp_prev_gpu_end = 0.0lf // last finished step's gpu_end — the handoff stat's anchor // wait + read back + release the pending step; x_b writes back only on the sync CPU-classifier @@ -1744,8 +1760,8 @@ def finish_pending_step { metal_wait_until_completed(g_lp_cbs[length(g_lp_cbs) - 1]) var err = "" for (c_i in g_lp_cbs) { - if (empty(err)) { - err = metal_command_buffer_error(c_i) + if (empty(err) && metal_command_buffer_failed(c_i)) { + err = cb_error_text(c_i) } g_sched_b_ms += (metal_command_buffer_kernel_end_time(c_i) - metal_command_buffer_kernel_start_time(c_i)) * 1000.0lf } @@ -1874,6 +1890,7 @@ def weight_caches_stale() : bool { return weights_epoch() != g_weight_cache_epoch } +[cold_path] // live-reload only def weight_caches_flush_on_reload { if (!weight_caches_stale()) { return @@ -1976,6 +1993,7 @@ enum MetalPrefillDecline { var g_prefill_declines = 0l var g_prefill_declines_by_reason : table +[cold_path] // the decline leg IS the fall-off-the-GPU path; its bookkeeping is not hot def decline(why : MetalPrefillDecline) { g_prefill_declines++ g_prefill_declines_by_reason["{why}"]++ @@ -1984,6 +2002,7 @@ def decline(why : MetalPrefillDecline) { } } +[cold_path] // the decline leg IS the fall-off-the-GPU path; its bookkeeping is not hot def decline(why : MetalPrefillDecline; missing : MetalNeed) { g_prefill_declines++ g_prefill_declines_by_reason["{why}"]++ diff --git a/modules/dasLLAMA/dasllama/dasllama_metal_llama.das b/modules/dasLLAMA/dasllama/dasllama_metal_llama.das index 2b872f2fbd..ae9010e105 100644 --- a/modules/dasLLAMA/dasllama/dasllama_metal_llama.das +++ b/modules/dasLLAMA/dasllama/dasllama_metal_llama.das @@ -1094,8 +1094,7 @@ def private encode_draft_step(t : Model; var r : StepRes; with_cls : bool = true // a row BELOW the watermark in its OWN slab and never runs the recurrent branch def private finish_draft_step(t : Model; var s : Session; r : StepRes; seam : bool = false) : bool { metal_wait_until_completed(r.cb) - let err = metal_command_buffer_error(r.cb) - let ok = empty(err) + let ok = !metal_command_buffer_failed(r.cb) if (ok) { let c = t.config let l = c.n_layers @@ -1122,7 +1121,7 @@ def private finish_draft_step(t : Model; var s : Session; r : StepRes; seam : bo s.mtp_h_pos1 = 0l // NOT a trunk hidden: the warm's seam row must not consume it } } else { - to_log(LOG_ERROR, "dasLLAMA metal MTP draft: dispatch failed ({err})\n") // nolint:PERF028,PERF026 — error path on a failed draft dispatch + to_log(LOG_ERROR, "dasLLAMA metal MTP draft: dispatch failed ({cb_error_text(r.cb)})\n") // nolint:PERF028,PERF026 — error path on a failed draft dispatch decline(MetalDecodeDecline.gpu_error) } return ok @@ -1131,6 +1130,7 @@ def private finish_draft_step(t : Model; var s : Session; r : StepRes; seam : bo //! forward_mtp's GPU twin: the draft head against a resident blob model. `pos` is the trunk //! position (n_past); the draft row lands at pos-1 (rotary is delta-invariant). Needs a //! GPU-current session (KV mirror holds rows [0, pos)) and GPU logits. false = declined. +[hot_path] def metal_mtp_draft_forward(t : Model; var s : Session; token, pos : int64; seam : bool = false) : bool { let c = t.config if (g_dev == null || !t.metal_blob || c.n_layer_nextn <= 0l || pos < 1l || @@ -1163,13 +1163,14 @@ def metal_mtp_draft_forward(t : Model; var s : Session; token, pos : int64; seam //! forward_mtp(..., false)'s GPU twin — the continuation-window SEAM: write the draft head's //! KV row pos-1 from (embed(token), the pre-window mtp_h) without touching logits or mtp_h. +[hot_path] def metal_mtp_seam_row(t : Model; var s : Session; token, pos : int64) : bool { return metal_mtp_draft_forward(t, s, token, pos, true) } // ===== MTP B=2 same-slab verify (the spec step's trunk forward) ===== -var private g_vrt : array // per-layer 2-row route entries: uint4(koff_el+lbase, voff_el+lbase, cap, cnt_i) +var private @scratch g_vrt : array // per-layer 2-row route entries: uint4(koff_el+lbase, voff_el+lbase, cap, cnt_i) // verify shape gate: f16 KV, silu FFN (dense or routed MoE ± shexp), no epilogue exotics. // Weight sites are fmt-generic (enc_verify_site / enc_moe_gemv — the load gate vets kernel @@ -1460,8 +1461,7 @@ def private encode_verify_step(t : Model; var r : StepRes; pos : int64) { // land the verify: KV rows + row0 logits -> s.mtp_logits, row1 -> s.logits; accept/reject bookkeeping is the caller's def private finish_verify_step(t : Model; var s : Session; r : StepRes; pos : int64) : bool { metal_wait_until_completed(r.cb) - let err = metal_command_buffer_error(r.cb) - let ok = empty(err) + let ok = !metal_command_buffer_failed(r.cb) if (ok) { let c = t.config unsafe { @@ -1483,7 +1483,7 @@ def private finish_verify_step(t : Model; var s : Session; r : StepRes; pos : in memcpy(addr < void? >(s.logits[0]), reinterpret(lg + c.vocab_size * 4l), c.vocab_size * 4l) } } else { - to_log(LOG_ERROR, "dasLLAMA metal MTP verify: dispatch failed ({err})\n") // nolint:PERF028,PERF026 — error path, and the CPU fallback dwarfs the log + to_log(LOG_ERROR, "dasLLAMA metal MTP verify: dispatch failed ({cb_error_text(r.cb)})\n") // nolint:PERF028,PERF026 — error path, and the CPU fallback dwarfs the log decline(MetalDecodeDecline.gpu_error) } return ok @@ -1492,7 +1492,7 @@ def private finish_verify_step(t : Model; var s : Session; r : StepRes; pos : in //! mtp_spec_eval's GPU twin: draft + B=2 same-slab verify, wholly on the metal decode path. //! Cold/decline paths fall back to one plain (GPU) forward — output-invariant either way. //! Reject is ~free: row0 IS the truth decode, the dn shadow region flips in — no re-forward. -[no_env, no_io] // no_alloc pending: the encoder's graph-capture push surface is its own audit +[hot_path] def metal_mtp_spec_eval(t : Model; var s : Session; tok : int64; var accepted : int64&) : bool { let c = t.config let pos = s.n_past @@ -1716,11 +1716,10 @@ def private encode_step(t : Model; var r : StepRes; commit_now : bool; var prev_ def private finish_step(t : Model; var s : Session; r : StepRes) : bool { let ts_wait = ref_time_ticks() metal_wait_until_completed(r.cb) - let err = metal_command_buffer_error(r.cb) + let ok = !metal_command_buffer_failed(r.cb) g_gpu_ms += (metal_command_buffer_gpu_end_time(r.cb) - metal_command_buffer_gpu_start_time(r.cb)) * 1000.0lf let wait_us = int64(get_time_usec(ts_wait)) g_us_wait += wait_us - let ok = empty(err) if (ok) { let ts_rb = ref_time_ticks() let c = t.config @@ -1770,7 +1769,7 @@ def private finish_step(t : Model; var s : Session; r : StepRes) : bool { ge = metal_command_buffer_gpu_end_time(r.cb))) } } else { - to_log(LOG_ERROR, "dasLLAMA metal decode: dispatch failed ({err}) — falling back to the CPU layer loop\n") // nolint:PERF028,PERF026 — error path, and the CPU fallback dwarfs the log + to_log(LOG_ERROR, "dasLLAMA metal decode: dispatch failed ({cb_error_text(r.cb)}) — falling back to the CPU layer loop\n") // nolint:PERF028,PERF026 — error path, and the CPU fallback dwarfs the log decline(MetalDecodeDecline.gpu_error) } return ok @@ -1945,9 +1944,8 @@ def private discard_pre { g_pre.valid = false if (r.committed) { metal_wait_until_completed(r.cb) - let err = metal_command_buffer_error(r.cb) - if (!empty(err)) { - to_log(LOG_ERROR, "dasLLAMA metal decode: discarded pre-step failed ({err})\n") // nolint:PERF028,PERF026 — error path on a discarded pre-step + if (metal_command_buffer_failed(r.cb)) { + to_log(LOG_ERROR, "dasLLAMA metal decode: discarded pre-step failed ({cb_error_text(r.cb)})\n") // nolint:PERF028,PERF026 — error path on a discarded pre-step } } release_step(r) @@ -2030,7 +2028,7 @@ def private metal_decode_slow(t : Model; var s : Session; pos : int64) : bool { //! The DecodeOverrideFn: runs one token step's whole llama layer stack GPU-resident against the //! session's KV mirror. true = done (KV row stored CPU-side too, s.x = final residual, s.logits //! when the GPU classifier gates pass); false = declined, the caller runs the CPU layer loop. -[no_env, no_io] // no_alloc pending: the encoder's graph-capture push surface is its own audit +[hot_path] def metal_decode_forward(t : Model; var s : Session; token, pos : int64) : bool { // mixed batch/single schedules: a pipelined batch step may still own this session's KV row // and logits — land it before reading anything session-side (no-op when nothing pends) @@ -2147,11 +2145,11 @@ def public set_metal_batch_addrms_unfused(v : bool) { // (eval_batch's contract) — OPT-IN for benches only: DASLLAMA_METAL_BATCH_PIPE=1, gpu-logits only. var private g_b_pipe = false // per-step scratch reused across calls (row table image + per-row mirror slices) -var private g_b_rte : array -var private g_b_koffs : array -var private g_b_voffs : array -var private g_b_caps : array -var private g_b_ulays : array // per-layer uniforms, pooled per step +var private @scratch g_b_rte : array +var private @scratch g_b_koffs : array +var private @scratch g_b_voffs : array +var private @scratch g_b_caps : array +var private @scratch g_b_ulays : array // per-layer uniforms, pooled per step // the model half is decode_shape_decline; per batch: uniform f16/f32 KV, every row within the // attention depth cap with a real uid, per-row rope rows built, and dim within the add+rms row @@ -2221,7 +2219,7 @@ def private batch_decode_decline(t : Model; ws : BatchWorkspace; sessions : arra //! The BatchDecodeOverrideFn (P4): one eval_batch step's whole layer stack GPU-resident — B rows //! through ONE weight pass. STACK-ONLY: on true every row's KV is stored (mirror + CPU cache) and //! ws.scr.x_b holds the final residuals; false = declined, CPU loop recomputes from clean state. -[no_env, no_io] // no_alloc pending: the encoder's graph-capture push surface is its own audit +[hot_path] def metal_batch_decode_forward(t : Model; var ws : BatchWorkspace; var sessions : array; nrows : int64) : bool { if (weight_caches_stale()) { // a model (re)load may have recycled cached weight addresses discard_pre() @@ -2572,8 +2570,8 @@ def metal_batch_decode_forward(t : Model; var ws : BatchWorkspace; var sessions let ts_enc = ref_time_ticks() let ncb = min(g_bncb, c.n_layers) let lay_per_cb = (c.n_layers + ncb - 1l) / ncb - var cbs : array - cbs |> reserve(int(ncb)) + var cbs : array // nolint:PERF026 — two steps' cb lists are alive at once (this one + the pending g_lp_cbs), so a reused global cannot serve; bounded by ncb + cbs |> reserve(int(ncb)) // nolint:PERF026 — see above // concurrent-encode rail: a concurrent encoder + explicit barriers ONLY at real buffer hazards // (the serial encoder implicit-barriers every dispatch). Gated to a full pipeline — a knockout // (g_skip) run stays on the serial encoder so a skipped stage can't leave a dependent unbarriered. @@ -2991,7 +2989,7 @@ def metal_batch_decode_forward(t : Model; var ws : BatchWorkspace; var sessions metal_end_encoding(enc) metal_release(enc) metal_commit(cb) - cbs |> push(cb) + cbs |> push(cb) // nolint:PERF026 — see the cbs declaration cb = metal_new_command_buffer_unretained(g_queue) enc = conc ? metal_new_compute_encoder_concurrent(cb) : metal_new_compute_encoder(cb) // fresh encoder + in-order queue serialize cross-cb hazards — start the tracker clean @@ -3050,7 +3048,7 @@ def metal_batch_decode_forward(t : Model; var ws : BatchWorkspace; var sessions metal_end_encoding(enc) metal_release(enc) metal_commit(cb) - cbs |> push(cb) + cbs |> push(cb) // nolint:PERF026 — see the cbs declaration hz_arm(false) // disarm the oracle with the encoder — a leaked arm would eat the next CAPTURE's stage g_us_bencode += int64(get_time_usec(ts_enc)) // optimistic watermark at commit: the in-order queue writes row pos before any later-committed diff --git a/modules/dasLLAMA/dasllama/dasllama_metal_prefill.das b/modules/dasLLAMA/dasllama/dasllama_metal_prefill.das index f4126691b5..f757bcb5ca 100644 --- a/modules/dasLLAMA/dasllama/dasllama_metal_prefill.das +++ b/modules/dasLLAMA/dasllama/dasllama_metal_prefill.das @@ -4434,7 +4434,11 @@ class MetalDnBa { var private g_pf_dev : MetalDevice? var private g_pf_queue : MetalCommandQueue? var private g_pf_failed = false -var private g_pf_regions : table // base address -> resident weight/scale/norm region +var private @scratch g_pf_regions : table // base address -> resident region; upload-once, steady-state [] is a hit +// per-prefill transients, reused across calls (prefill orchestration is serial) +var private @scratch g_pf_bks : array +var private @scratch g_pf_bvs : array +var private @scratch g_pf_cbs : array var private g_pf_pool : MetalBufferPool var private g_pf_upool : MetalBufferPool // UNTRACKED pool: uniforms + rope tables (GPU-read-only) // what the prefill kernels implement beyond the base llama dense block — grows wave by wave @@ -5702,6 +5706,7 @@ def private pf_enc_swiglu_oai(enc : MetalComputeEncoder?; bg, bu, btot, bcnt, bb // own address-keyed region cache when a model (re)load may have recycled the addresses var private g_pf_cache_epoch = -1l +[cold_path] // live-reload only def private pf_weight_caches_flush_on_reload { if (weights_epoch() == g_pf_cache_epoch) { return @@ -5734,7 +5739,7 @@ def private metal_ple_pre_gpu_gate(t : Model; s : Session; npos, start_pos : int return metal_prefill_init() } -[no_env, no_io] // no_alloc pending: the encoder's graph-capture push surface is its own audit +[hot_path] def metal_prefill_forward(t : Model; var s : Session; npos, start_pos : int64) : bool { pf_weight_caches_flush_on_reload() var why = prefill_decline(t, s, npos, start_pos) @@ -5867,8 +5872,10 @@ def metal_prefill_forward(t : Model; var s : Session; npos, start_pos : int64) : let dual_rope = has_dual_rope(c) var bcos_swa = dual_rope ? pool_acquire_untracked(g_pf_upool, g_pf_dev, bytes_tab_sw) : null var bsin_swa = dual_rope ? pool_acquire_untracked(g_pf_upool, g_pf_dev, bytes_tab_sw) : null - var bks : array - var bvs : array + var bks & = unsafe(g_pf_bks) + var bvs & = unsafe(g_pf_bvs) + bks |> clear() + bvs |> clear() bks |> reserve(int(c.n_layers + c.n_layer_nextn)) bvs |> reserve(int(c.n_layers + c.n_layer_nextn)) var no_panel : MetalBuffer? // recurrent layers' null K/V slot @@ -6123,7 +6130,8 @@ def metal_prefill_forward(t : Model; var s : Session; npos, start_pos : int64) : // CONCURRENT encoder (auto-schedule + hz barriers at real hazards). DASLLAMA_METAL_SCHED=0 // keeps capture order; DASLLAMA_METAL_PF_CAPTURE=0 = the direct serial-encode rollback. let cap = g_pf_env_capture - var cbs : array + var cbs & = unsafe(g_pf_cbs) + cbs |> clear() cbs |> reserve(int(ncb)) var cb : MetalCommandBuffer? var enc : MetalComputeEncoder? @@ -6696,10 +6704,9 @@ def metal_prefill_forward(t : Model; var s : Session; npos, start_pos : int64) : let runsp = addr < KVRun const? >(s.kv_runs[0]) for (cbi, i in cbs, count()) { metal_wait_until_completed(cbi) - let e = metal_command_buffer_error(cbi) - if (!empty(e)) { + if (metal_command_buffer_failed(cbi)) { ran = false - err = e + err = cb_error_text(cbi) } let t0 = metal_command_buffer_gpu_start_time(cbi) let t1 = metal_command_buffer_gpu_end_time(cbi) @@ -6746,9 +6753,6 @@ def metal_prefill_forward(t : Model; var s : Session; npos, start_pos : int64) : } } cbs |> clear() - unsafe { - delete cbs - } gpu_ms = (gpu_t1 - gpu_t0) * 1000.0lf // whole-GPU window: first chunk start -> last chunk end let encwait_us = get_time_usec(ts_enc) let ok = ran @@ -6874,10 +6878,6 @@ def metal_prefill_forward(t : Model; var s : Session; npos, start_pos : int64) : } bks |> clear() // non-owning handles — the pool owns them now bvs |> clear() - unsafe { - delete bks - delete bvs - } pool_release(g_pf_upool, u_dim, 4ul) pool_release(g_pf_upool, u_qd, 4ul) pool_release(g_pf_upool, u_kvd, 4ul) diff --git a/modules/dasMetal/metal/das_metal_boost.das b/modules/dasMetal/metal/das_metal_boost.das index f8900b595c..cbfff2810a 100644 --- a/modules/dasMetal/metal/das_metal_boost.das +++ b/modules/dasMetal/metal/das_metal_boost.das @@ -92,7 +92,7 @@ def public metal_log_leaks_if_any : int64 { //! (the source participates so two kernels sharing an entry name cannot collide). Drain releases //! everything — the cache owns its pipelines. struct public MetalPipelineCache { - cache : table + @scratch cache : table // compile-once PSO residency; steady-state [] is a hit } def public pipeline_cached(var pc : MetalPipelineCache; dev : MetalDevice?; source, entry : string; fastmath : bool; var error : string&) : MetalComputePipeline? { @@ -124,7 +124,7 @@ def public cache_drain(var pc : MetalPipelineCache) { //! bucket — deterministic for fixed-shape drivers). Acquired buffers stay owned by the caller //! until released back; drain frees the free lists only. struct public MetalBufferPool { - free_bufs : table> + @scratch free_bufs : table> // size-bucketed free lists; recycle is the design } def private pool_bucket(bytes : uint64) : uint64 { diff --git a/modules/dasMetal/src/dasMetal.h b/modules/dasMetal/src/dasMetal.h index 2cf2f6c5ae..13beadafdd 100644 --- a/modules/dasMetal/src/dasMetal.h +++ b/modules/dasMetal/src/dasMetal.h @@ -65,6 +65,7 @@ namespace das { DAS_MOD_API void metal_commit ( MetalCommandBuffer * cb, Context * ctx, LineInfoArg * at ); DAS_MOD_API void metal_wait_until_completed ( MetalCommandBuffer * cb, Context * ctx, LineInfoArg * at ); DAS_MOD_API char * metal_command_buffer_error ( MetalCommandBuffer * cb, Context * ctx, LineInfoArg * at ); + DAS_MOD_API bool metal_command_buffer_failed ( MetalCommandBuffer * cb, Context * ctx, LineInfoArg * at ); DAS_MOD_API double metal_command_buffer_gpu_start_time ( MetalCommandBuffer * cb, Context * ctx, LineInfoArg * at ); DAS_MOD_API double metal_command_buffer_gpu_end_time ( MetalCommandBuffer * cb, Context * ctx, LineInfoArg * at ); DAS_MOD_API double metal_command_buffer_kernel_start_time ( MetalCommandBuffer * cb, Context * ctx, LineInfoArg * at ); diff --git a/modules/dasMetal/src/dasMetal.mm b/modules/dasMetal/src/dasMetal.mm index 18f728557b..054e4e936b 100644 --- a/modules/dasMetal/src/dasMetal.mm +++ b/modules/dasMetal/src/dasMetal.mm @@ -503,6 +503,16 @@ void metal_wait_until_completed ( MetalCommandBuffer * cb, Context * ctx, LineIn [(__bridge id)(void *) cb waitUntilCompleted]; } + // the hot-path probe: status check only, no string — metal_command_buffer_error allocates + // its message, so per-step success checks use this and fetch the text on the failure leg + bool metal_command_buffer_failed ( MetalCommandBuffer * cb, Context * ctx, LineInfoArg * at ) { + if ( !cb ) ctx->throw_error_at(at, "metal_command_buffer_failed: null command buffer"); + @autoreleasepool { + id c = (__bridge id)(void *) cb; + return c.status == MTLCommandBufferStatusError; + } + } + char * metal_command_buffer_error ( MetalCommandBuffer * cb, Context * ctx, LineInfoArg * at ) { if ( !cb ) ctx->throw_error_at(at, "metal_command_buffer_error: null command buffer"); @autoreleasepool { @@ -702,6 +712,9 @@ int64_t metal_dispatch_call_count () { addExtern(*this, lib, "metal_command_buffer_error", SideEffects::modifyExternal, "metal_command_buffer_error") ->args({"command_buffer", "context", "at"}); + addExtern(*this, lib, "metal_command_buffer_failed", + SideEffects::modifyExternal, "metal_command_buffer_failed") + ->args({"command_buffer", "context", "at"}); addExtern(*this, lib, "metal_command_buffer_gpu_start_time", SideEffects::modifyExternal, "metal_command_buffer_gpu_start_time") ->args({"command_buffer", "context", "at"}); diff --git a/skills/perf_lint.md b/skills/perf_lint.md index 9dc749809a..9904ef2229 100644 --- a/skills/perf_lint.md +++ b/skills/perf_lint.md @@ -144,7 +144,7 @@ transitively, so a sink several frames deep is still reported. **Declaring a contract is free.** The five annotations are registered by the compiler itself as metadata-only markers (the same shape as `[clone]`), so a file under contract requires *nothing* — no module compiles, no build time is paid. `@scratch` is a -field/parameter annotation, which is free-form anyway. The verification lives here, in +field/parameter/global annotation, which is free-form anyway. The verification lives here, in `daslib/perf_lint`, which code under contract does **not** require: the checker is heavy (it pulls `ast_boost`, `lint_config`, `toml`, `json`), and lint already runs where lint belongs — `utils/lint/main.das`, the MCP `lint` subtool, CI — none of which ever needed the target to require @@ -187,8 +187,10 @@ struct Session { def scratch_resize(@scratch var a : array; need : int64) { ... } ``` -A sizing call whose destination reaches a `@scratch` field is not descended into. Field -annotations are free-form (`@name`, no registration), so `@scratch` costs nothing to parse. +A sizing call whose destination reaches a `@scratch` declaration is not descended into — a +struct field, a by-ref helper parameter, or a **module global** (`var @scratch g : array`; +the annotation goes AFTER `var` — before it is a syntax error). Field/variable annotations are +free-form (`@name`, no registration), so `@scratch` costs nothing to parse. **Future — `@scratch` as an optimization hint, not just a lint marker.** Today the declaration only tells the linter "this buffer is reused". The same statement is exactly the precondition a @@ -211,10 +213,15 @@ rather than for the lint, so it can carry the optimization meaning later without 1. `[cold_path]` on the callee — the honest fix when the leg genuinely runs once (lazy init, PSO compile, opt-in bookkeeping, a reference-check path behind a debug flag). -2. `@scratch` on the destination when it is a reused buffer. -3. `// nolint:PERF026` on the line, **with a reason**. Honored at *either* end of a chain: the - report anchors on the caller when the sink is in another file, so a suppression written where - the code actually lives still works. +2. `@scratch` on the destination when it is a reused buffer — a struct field, a by-ref helper + parameter, or a module global. It covers the sizing set (`resize`/`reserve`/`push*`/`emplace`/ + `insert`/`erase`/`pop`/`clear`), **table indexing** (`t[k]` on a `@scratch` table is the + pool / residency-cache shape — insert on first touch, steady-state hit), and it follows local + reference bindings (`var lst & = pool.free_bufs[b]` carries the mark). +3. `// nolint:PERF026` on the line, **with a reason**. Honored **anywhere along the chain** — + the anchor line, the sink line, or any intermediate call site. The sink often bottoms out in + `daslib/builtin.das` (whose lines are nobody's to annotate), so the honest suppression line is + usually the call site in the module that owns the decision, and that works. 4. `DAS_LINT_DISABLE=PERF028` for a whole run — no source edit, which is the point when you are adding a log line to chase a bug. With all three codes disabled the closure walk is skipped entirely, so this buys compile time and not just quiet. diff --git a/utils/lint/tests/perf026_hot_path_alloc.das b/utils/lint/tests/perf026_hot_path_alloc.das index 1bfb980b1d..f2a53b54e1 100644 --- a/utils/lint/tests/perf026_hot_path_alloc.das +++ b/utils/lint/tests/perf026_hot_path_alloc.das @@ -25,7 +25,7 @@ options auto_inline_functions = false // lint fixtures assert SOURCE shapes; s // sink several frames deep is still reported — anchored on the deepest line // written in the root's own file, not on the daslib internal that allocates. -expect 31208:5 +expect 31208:7 require daslib/perf_lint @@ -108,3 +108,55 @@ def bad_new_delete() : int { } return v } + +// 9. a @scratch module GLOBAL is declared intent too — the capture-rail shape: a +// clear()-recycled array that grows once and is reused; its unmarked neighbour still reports +var @scratch g_reused : array +var g_fresh : array + +[no_alloc] +def bad_only_the_unmarked_global() { + g_reused |> push(1) + g_reused |> pop() // pop rides the escape too (shrink of a reused buffer is a store) + g_fresh |> push(2) +} + +// 10. indexing a @scratch TABLE is the declared pool / residency-cache shape (insert-on-first- +// touch, steady-state lookup); an unmarked table's [] still reports its insert-on-read +var @scratch g_pool : table +var g_plain : table + +[no_alloc] +def bad_only_the_unmarked_table(k : int) : int { + unsafe { + return g_pool[k] + g_plain[k] + } +} + +// 11. a local REFERENCE binding carries its destination's scratch-ness (the pool free-list +// idiom: `var lst & = pool.buckets[b]`) +var @scratch g_buckets : table> + +[no_alloc] +def good_scratch_through_ref(b : int) : int { + var lst & = unsafe(g_buckets[b]) + lst |> push(1) + let v = lst[length(lst) - 1] + lst |> pop() + return v +} + +// 12. a nolint ANYWHERE along the chain suppresses — the sink may bottom in daslib, so the +// honest line is often an intermediate call site in a required module +def deep_alloc12(var a : array) { + a |> reserve(64) +} + +def mid_hop12(var a : array) { + deep_alloc12(a) // nolint:PERF026 — chain-middle suppression is honored +} + +[no_alloc] +def good_chain_middle_nolint(var a : array) { + mid_hop12(a) +} From f44e8ef0f4690e60d19e312250b017707fb8d27b Mon Sep 17 00:00:00 2001 From: Boris Batkin Date: Mon, 27 Jul 2026 13:50:13 -0700 Subject: [PATCH 02/18] dasMetal: pin the runtime MSL language version to 4.0 on macOS 26 (Metal-4 tensor lane, step 0) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit metal_new_library_from_source relied on MTLCompileOptions' DEFAULT language version, which tracks the SDK the module binary was linked against — the same OS compiles metal_tensor on a freshly built module and rejects it (`mpp` / `dextents` undeclared) on one built under an older CLT (probe-verified: M1 fresh vs m4's 07-26 build, both macOS 26.5.2). Pin languageVersion to MTLLanguageVersion4_0 under @available(macOS 26) so runtime compilation is deterministic. MSL 4 compiles existing 3.x sources; the metal test suite and the dasLLAMA metal kernels are unaffected (probe + test_metal_sdot green). Step-0 probe results (recorded in plans/metal4-tensor-lane.md): matmul2d over tensor_inline views built in-shader from raw device pointers runs EXACT on the M1 Max fallback — the tensor lane integrates emitter-only, no host MTLTensor. Co-Authored-By: Claude Fable 5 --- modules/dasMetal/src/dasMetal.mm | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/modules/dasMetal/src/dasMetal.mm b/modules/dasMetal/src/dasMetal.mm index 054e4e936b..4a43afaf9c 100644 --- a/modules/dasMetal/src/dasMetal.mm +++ b/modules/dasMetal/src/dasMetal.mm @@ -150,6 +150,14 @@ bool metal_device_has_unified_memory ( MetalDevice * dev, Context * ctx, LineInf opts.mathMode = fastmath ? MTLMathModeFast : MTLMathModeSafe; #else opts.fastMathEnabled = fastmath ? YES : NO; +#endif +#if defined(__MAC_OS_X_VERSION_MAX_ALLOWED) && __MAC_OS_X_VERSION_MAX_ALLOWED >= 260000 + // pin the newest MSL the OS offers: the DEFAULT tracks the SDK the binary was + // LINKED against, so a module built under an older CLT silently loses metal_tensor + // (mpp/dextents undeclared) on the same OS — probe-verified on the m4 box + if (@available(macOS 26.0, *)) { + opts.languageVersion = MTLLanguageVersion4_0; + } #endif NSString * nsSrc = [NSString stringWithUTF8String:src]; if ( nsSrc == nil ) { // invalid UTF-8 — a nil source would raise an ObjC exception below From 8fbf2c2d92fd875ef001433990bae69aaa922382 Mon Sep 17 00:00:00 2001 From: Boris Batkin Date: Mon, 27 Jul 2026 21:11:07 -0700 Subject: [PATCH 03/18] dasMetal: Metal-4 tensor-op lowering (tmm2d builtins + MPP helper emission) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Phase A of the tensor lane: three tmm2d_* builtins (f32*bf16, q8*q8 with f32 or f16 block scales) lower through msl_emit into generated mpp::tensor_ops::matmul2d helpers over tensor_inline views — plain buffer binds, no host MTLTensor. CPU-replay stubs make every kernel exactly replayable; conformance test proves GPU==CPU bit-exact on the M1 fallback. Crown store (runtime.metal_tensor sidecar knob) lands in dasllama_common. Co-Authored-By: Claude Fable 5 --- modules/dasLLAMA/dasllama/dasllama_common.das | 25 +++ modules/dasMetal/metal/metal_builtins.das | 105 ++++++++++ modules/dasMetal/metal/msl_emit.das | 132 +++++++++++- tests/metal/test_metal_tensor_ops.das | 197 ++++++++++++++++++ 4 files changed, 457 insertions(+), 2 deletions(-) create mode 100644 tests/metal/test_metal_tensor_ops.das diff --git a/modules/dasLLAMA/dasllama/dasllama_common.das b/modules/dasLLAMA/dasllama/dasllama_common.das index ca91cf876b..f7922ef6e3 100644 --- a/modules/dasLLAMA/dasllama/dasllama_common.das +++ b/modules/dasLLAMA/dasllama/dasllama_common.das @@ -4534,6 +4534,26 @@ def private apply_i64_knob(rt : JsonValue const?; key : string; blk : block<(v : } } +// ===== Metal-4 tensor-lane crowns (the sidecar's "runtime"."metal_tensor" knob) ===== +// Comma-list of families whose TENSOR twin won this box's tune race; consulted once per +// family at pso build. Default empty = simdgroup; a DAS_TUNE_MANIFEST copy force-crowns. +var private g_tensor_crowns : table + +def public set_metal_tensor_crowns(list : string) { + unsafe { + delete g_tensor_crowns + } + for (f in split(list, ",")) { + if (!empty(f)) { + g_tensor_crowns |> insert(f) + } + } +} + +def public metal_tensor_crowned(family : string) : bool { + return key_exists(g_tensor_crowns, family) +} + //! Apply the runtime knobs from the app tune sidecar's optional "runtime" section (see //! tune_for_this_box.md): token block, L2 budget, threading thresholds, chunk multipliers, and an //! optional kernel-backend pin. Missing/STALE file or missing section = no-op; every applied entry logs. @@ -4583,6 +4603,11 @@ def apply_box_profile_runtime(path : string = "") { apply_i64_knob(rt, "jobque_spin_us") $(v) { set_jobque_spin_us(v) } apply_i64_knob(rt, "jobque_join_poll") $(v) { set_jobque_join_poll(v) } apply_i64_knob(rt, "team_rank_gate") $(v) { set_team_rank_gate(int(min(v, 1l))) } + let mtc = read_json_field(rt, "metal_tensor", "") + if (mtc != "") { + set_metal_tensor_crowns(mtc) + to_log(LOG_INFO, "dasLLAMA: box profile runtime: metal_tensor crowns = {mtc}\n") + } let be = read_json_field(rt, "backend", "") if (be != "") { pin_kernel_backend(be) // warns and keeps the current backend if the name is unknown diff --git a/modules/dasMetal/metal/metal_builtins.das b/modules/dasMetal/metal/metal_builtins.das index 36c8ea752b..5cdee1a7b9 100644 --- a/modules/dasMetal/metal/metal_builtins.das +++ b/modules/dasMetal/metal/metal_builtins.das @@ -11,6 +11,7 @@ options indenting = 4 module metal_builtins shared public require daslib/shader_lingua_franca public +require daslib/math_bits // uint_bits_to_float — the tensor-op CPU stubs decode bf16/f16 halfwords // ===== compute builtins (Metal-only additions) ===== // simdgroup IDs, GLSL subgroup spellings (KHR_shader_subgroup names — portable if dasSpirv grows @@ -244,3 +245,107 @@ def public simdgroup_multiply(var d : simdgroup_half8x8; a, b : simdgroup_half8x } } } + +// ===== Metal-4 tensor ops (MTLTensor / mpp::tensor_ops::matmul2d) ===== +// Statement-form cooperative GEMM markers: on the GPU the MSL emitter replaces each call with +// tensor_inline views + a matmul2d op (see msl_emit); the bodies below are the CPU-replay +// reference. All threads must reach the call convergently (cooperative op); the stubs compute +// the full tile once (thread 0) since replay threads run independently. +// Layout contract (NT, the engine's panels): A = activations, m rows x kk, row stride lda; +// W = weights, n output cols x kk, row stride ldw; C = m x n, row stride ldc, OVERWRITTEN. +// m/n/sgs (tile rows, tile cols, cooperating simdgroups) must be call-site constants. + +//! f32 activations x bf16 weights -> f32 tile (raw bf16 halfwords in `w`). +[unused_argument(sgs), sideeffects] +def public tmm2d_f32_bf16_f32(m, n, sgs : uint; a : float const?; lda : uint; w : uint16 const?; ldw : uint; var c : float?; ldc : uint; kk : uint) { + if (gl_LocalInvocationID.x != 0u) { + return + } + unsafe { + for (r in range(int(m))) { + for (j in range(int(n))) { + var acc = 0.0 + for (k in range(int(kk))) { + acc += a[r * int(lda) + k] * uint_bits_to_float(uint(w[j * int(ldw) + k]) << 16u) + } + c[r * int(ldc) + j] = acc + } + } + } +} + +//! Q8xQ8 tile with the per-32-block scale fold (exact int32 accumulate per block, then +//! facc += float(iacc) * as[r][b] * ws[j][b] — the dot_q8q8 contract). Scale row strides in +//! ELEMENTS (blocks): as at a + r*lda_s, ws at w + j*ldw_s. +[unused_argument(sgs), sideeffects] +def public tmm2d_q8_f32(m, n, sgs : uint; aq : int8 const?; lda : uint; asf : float const?; lda_s : uint; wq : int8 const?; ldw : uint; wsf : float const?; ldw_s : uint; var c : float?; ldc : uint; kk : uint) { + if (gl_LocalInvocationID.x != 0u) { + return + } + unsafe { + let nb = int(kk) / 32 + for (r in range(int(m))) { + for (j in range(int(n))) { + var facc = 0.0 + for (b in range(nb)) { + var iacc = 0 + for (k in range(32)) { + iacc += int(aq[r * int(lda) + b * 32 + k]) * int(wq[j * int(ldw) + b * 32 + k]) + } + facc += float(iacc) * asf[r * int(lda_s) + b] * wsf[j * int(ldw_s) + b] + } + c[r * int(ldc) + j] = facc + } + } + } +} + +// binary16 halfword -> f32 for the CPU stubs (exact; subnormals included) +def private f16_bits_to_f32(h : uint) : float { + let s = (h >> 15u) & 1u + let e = (h >> 10u) & 31u + let mant = h & 1023u + var bits : uint + if (e == 0u) { + if (mant == 0u) { + bits = s << 31u + } else { + var m = mant + var ee = 127u - 15u + 1u + while ((m & 1024u) == 0u) { + m <<= 1u + ee-- + } + bits = (s << 31u) | (ee << 23u) | ((m & 1023u) << 13u) + } + } elif (e == 31u) { + bits = (s << 31u) | (255u << 23u) | (mant << 13u) + } else { + bits = (s << 31u) | ((e + 127u - 15u) << 23u) | (mant << 13u) + } + return uint_bits_to_float(bits) +} + +//! The wscale_f16 twin of tmm2d_q8_f32: weight scales are raw binary16 halfwords. +[unused_argument(sgs), sideeffects] +def public tmm2d_q8_f16s(m, n, sgs : uint; aq : int8 const?; lda : uint; asf : float const?; lda_s : uint; wq : int8 const?; ldw : uint; wsf : uint16 const?; ldw_s : uint; var c : float?; ldc : uint; kk : uint) { + if (gl_LocalInvocationID.x != 0u) { + return + } + unsafe { + let nb = int(kk) / 32 + for (r in range(int(m))) { + for (j in range(int(n))) { + var facc = 0.0 + for (b in range(nb)) { + var iacc = 0 + for (k in range(32)) { + iacc += int(aq[r * int(lda) + b * 32 + k]) * int(wq[j * int(ldw) + b * 32 + k]) + } + facc += float(iacc) * asf[r * int(lda_s) + b] * f16_bits_to_f32(uint(wsf[j * int(ldw_s) + b])) + } + c[r * int(ldc) + j] = facc + } + } + } +} diff --git a/modules/dasMetal/metal/msl_emit.das b/modules/dasMetal/metal/msl_emit.das index ca0654c7c4..77b4932080 100644 --- a/modules/dasMetal/metal/msl_emit.das +++ b/modules/dasMetal/metal/msl_emit.das @@ -49,6 +49,8 @@ struct private MslCtx { renames : table // das identifier -> MSL-safe identifier ptrs : table // pointer local (addr(member[..])) -> its member (das names) helpers : table // preamble helper functions referenced by lowered calls + gen_helpers : table // GENERATED preamble helpers (tensor ops): mangled name -> text + scan_ptrs : table // scan-time pointer local -> member (write-set resolution) self_name : string } @@ -264,6 +266,15 @@ def private mark_write_target(var ctx : MslCtx; e : Expression?) { if (key_exists(ctx.members, name)) { ctx.members[name].written = true } + } elif (e is ExprVar) { + // a pointer local (`var p = unsafe(addr(member[..]))`) writes through to its member + let pn = string((e as ExprVar).name) + if (key_exists(ctx.scan_ptrs, pn)) { + let mem = ctx.scan_ptrs[pn] + if (key_exists(ctx.members, mem)) { + ctx.members[mem].written = true + } + } } } @@ -284,6 +295,12 @@ def private scan_body(var ctx : MslCtx; e : Expression?) { scan_body(ctx, (e as ExprWith).body) } elif (e is ExprLet) { for (v in (e as ExprLet).variables) { + if (v.init != null && v.init is ExprRef2Ptr && (v.init as ExprRef2Ptr).subexpr is ExprAt) { + let pat = (v.init as ExprRef2Ptr).subexpr as ExprAt + if (pat.subexpr is ExprField) { + ctx.scan_ptrs[string(v.name)] = string((pat.subexpr as ExprField).name) + } + } scan_body(ctx, v.init) } } elif (e is ExprCopy) { @@ -324,9 +341,17 @@ def private scan_body(var ctx : MslCtx; e : Expression?) { let ecs = e as ExprCall // simdgroup_store writes through its buffer argument — without this the write-set scan // would lower the parameter `device const` - if (call_base_name("{ecs.name}") == "simdgroup_store" && length(ecs.arguments) == 4) { + let scn = call_base_name("{ecs.name}") + if (scn == "simdgroup_store" && length(ecs.arguments) == 4) { mark_write_target(ctx, ecs.arguments[1]) } + // the tensor GEMMs write through their C pointer (arg 7 for bf16, 11 for the q8 forms) + if (scn == "tmm2d_f32_bf16_f32" && length(ecs.arguments) == 10) { + mark_write_target(ctx, ecs.arguments[7]) + } + if ((scn == "tmm2d_q8_f32" || scn == "tmm2d_q8_f16s") && length(ecs.arguments) == 14) { + mark_write_target(ctx, ecs.arguments[11]) + } for (arg in ecs.arguments) { scan_body(ctx, arg) } @@ -1139,6 +1164,93 @@ def private emit_sgmat_multiply(var ctx : MslCtx; ec : ExprCall?; indent : strin // ===== statement emission ===== +// ===== Metal-4 tensor ops (mpp::tensor_ops::matmul2d) ===== +// Each tmm2d_* call lowers to ONE generated preamble helper (mangled by kind/m/n/sgs) holding +// the tensor_inline views + matmul2d op, and a plain call at the site. m/n/sgs must be +// call-site constants (the simdgroup-lane precedent). Layout contract in metal_builtins.das. + +def private tmm2d_const_uint(var ctx : MslCtx; e : Expression?; what : string) : uint { + var v = e + if (v is ExprRef2Value) { + v = (v as ExprRef2Value).subexpr + } + if (v is ExprConstUInt) { + return (v as ExprConstUInt).value + } + err(ctx, e.at, "tmm2d: {what} must be a compile-time uint constant") + return 0u +} + +def private tmm2d_helper_text(kind : string; m, n, sgs : uint) : string { + let name = "__das_tmm2d_{kind}_{int(m)}_{int(n)}_{int(sgs)}" + let ct = "tensor, tensor_inline>" + if (kind == "bf") { + return build_string() $(var w) { + w |> write("static void {name}(device const float * a, uint lda, device const ushort * bw, uint ldw, device float * c, uint ldc, uint kk) \{\n") + w |> write(" auto A = tensor, tensor_inline>((device float *)a, dextents(int(kk), {int(m)}), metal::array\{1, int(lda)\});\n") + w |> write(" auto W = tensor, tensor_inline>((device bfloat *)bw, dextents(int(kk), {int(n)}), metal::array\{1, int(ldw)\});\n") + w |> write(" auto C = tensor, tensor_inline>(c, dextents({int(n)}, {int(m)}), metal::array\{1, int(ldc)\});\n") + w |> write(" constexpr auto d = matmul2d_descriptor({int(m)}, {int(n)}, static_cast(dynamic_extent), false, true, false);\n") + w |> write(" matmul2d> op;\n") + w |> write(" auto cT = op.get_destination_cooperative_tensor();\n") + w |> write(" for (uint16_t i = 0; i < cT.get_capacity(); ++i) \{ if (cT.is_valid_element(i)) cT[i] = 0.0f; \}\n") + w |> write(" op.run(A, W, cT);\n") + w |> write(" cT.store(C);\n") + w |> write("\}") + } + } + // q8 / q8_f16s: per-32-block exact int accumulate + scale fold in a float cooperative acc + let wst = kind == "q8h" ? "device const ushort" : "device const float" + let wsread = kind == "q8h" ? "float(as_type(ws[uint(ids[0]) * ldws + b]))" : "ws[uint(ids[0]) * ldws + b]" + return build_string() $(var w) { + w |> write("static void {name}(device const char * aq, uint lda, device const float * as, uint ldas, device const char * wq, uint ldw, {wst} * ws, uint ldws, device float * c, uint ldc, uint kk) \{\n") + w |> write(" constexpr auto d = matmul2d_descriptor({int(m)}, {int(n)}, 32, false, true, false);\n") + w |> write(" matmul2d> op;\n") + w |> write(" auto C = tensor, tensor_inline>(c, dextents({int(n)}, {int(m)}), metal::array\{1, int(ldc)\});\n") + w |> write(" auto A0 = {ct}((device int8_t *)aq, dextents(32, {int(m)}), metal::array\{1, int(lda)\});\n") + w |> write(" auto W0 = {ct}((device int8_t *)wq, dextents(32, {int(n)}), metal::array\{1, int(ldw)\});\n") + w |> write(" auto fT = op.get_destination_cooperative_tensor();\n") + w |> write(" for (uint16_t i = 0; i < fT.get_capacity(); ++i) \{ if (fT.is_valid_element(i)) fT[i] = 0.0f; \}\n") + w |> write(" const uint nb = kk / 32u;\n") + w |> write(" for (uint b = 0; b < nb; ++b) \{\n") + w |> write(" auto A = {ct}((device int8_t *)(aq + b * 32u), dextents(32, {int(m)}), metal::array\{1, int(lda)\});\n") + w |> write(" auto W = {ct}((device int8_t *)(wq + b * 32u), dextents(32, {int(n)}), metal::array\{1, int(ldw)\});\n") + w |> write(" auto iT = op.get_destination_cooperative_tensor();\n") + w |> write(" for (uint16_t i = 0; i < iT.get_capacity(); ++i) \{ if (iT.is_valid_element(i)) iT[i] = 0; \}\n") + w |> write(" op.run(A, W, iT);\n") + w |> write(" for (uint16_t i = 0; i < iT.get_capacity(); ++i) \{ if (iT.is_valid_element(i)) \{\n") + w |> write(" auto ids = iT.get_multidimensional_index(i);\n") + w |> write(" fT[i] += float(iT[i]) * as[uint(ids[1]) * ldas + b] * {wsread};\n") + w |> write(" \} \}\n") + w |> write(" \}\n") + w |> write(" fT.store(C);\n") + w |> write("\}") + } +} + +def private emit_tmm2d(var ctx : MslCtx; ec : ExprCall?; indent : string; var lines : array) { + let cname = call_base_name("{ec.name}") + let kind = cname == "tmm2d_f32_bf16_f32" ? "bf" : (cname == "tmm2d_q8_f16s" ? "q8h" : "q8") + let want_args = kind == "bf" ? 10 : 14 + if (length(ec.arguments) != want_args) { + err(ctx, ec.at, "{cname} takes exactly {want_args} arguments") + return + } + let m = tmm2d_const_uint(ctx, ec.arguments[0], "m (tile rows)") + let n = tmm2d_const_uint(ctx, ec.arguments[1], "n (tile cols)") + let sgs = tmm2d_const_uint(ctx, ec.arguments[2], "sgs (simdgroups)") + if (m == 0u || n == 0u || sgs == 0u) { + return + } + let hname = "__das_tmm2d_{kind}_{int(m)}_{int(n)}_{int(sgs)}" + if (!key_exists(ctx.gen_helpers, hname)) { + ctx.gen_helpers[hname] = tmm2d_helper_text(kind, m, n, sgs) + } + var vals <- [for (i in range(3, length(ec.arguments))); emit_value(ctx, ec.arguments[i])] + note(ctx, "call.{cname}") + lines |> push("{indent}{hname}({join(vals, ", ")});") +} + def private emit_stmt(var ctx : MslCtx; e : Expression?; indent : string; var lines : array) { if (e == null) { return @@ -1328,6 +1440,8 @@ def private emit_stmt(var ctx : MslCtx; e : Expression?; indent : string; var li emit_tg_store_float4(ctx, ec, indent, lines) } elif (cname == "simdgroup_multiply_accumulate" || cname == "simdgroup_multiply") { emit_sgmat_multiply(ctx, ec, indent, lines) + } elif (cname == "tmm2d_f32_bf16_f32" || cname == "tmm2d_q8_f32" || cname == "tmm2d_q8_f16s") { + emit_tmm2d(ctx, ec, indent, lines) } else { err(ctx, ec.at, "call `{cname}` as a statement has no MSL lowering — only barrier()/memoryBarrierShared() sync and the simdgroup_matrix ops are supported") } @@ -1569,8 +1683,22 @@ def public generate_msl(fn : FunctionPtr; var errors : array; cfg : MslE hs |> sort() // table iteration order is hash-based — pin the emitted order helper_text = join(hs, "\n") + "\n" } + var mpp_include = "" + if (!empty(ctx.gen_helpers)) { + // Metal-4 tensor ops: the MPP include + namespace, then the generated helpers in + // pinned (sorted) order — table iteration order is hash-based + mpp_include = "#include \nusing namespace mpp::tensor_ops;\n" + var gh <- [for (k in keys(ctx.gen_helpers)); k] + gh |> sort() + helper_text += build_string() $(var w) { + for (k in gh) { + w |> write(ctx.gen_helpers[k]) + w |> write("\n") + } + } + } let chunks = [ - "#include \nusing namespace metal;\n{helper_text}kernel void {kernel_entry_name(fn, cfg)}(", + "#include \nusing namespace metal;\n{mpp_include}{helper_text}kernel void {kernel_entry_name(fn, cfg)}(", join(params, ",\n "), ") \{\n", empty(tg_lines) ? "" : join(tg_lines, "\n") + "\n", diff --git a/tests/metal/test_metal_tensor_ops.das b/tests/metal/test_metal_tensor_ops.das new file mode 100644 index 0000000000..bb5a81cc1e --- /dev/null +++ b/tests/metal/test_metal_tensor_ops.das @@ -0,0 +1,197 @@ +options gen2 +options indenting = 4 + +// Metal-4 tensor ops (mpp::tensor_ops::matmul2d) through the das kernel rail: each tmm2d_* +// call lowers to a generated preamble helper building tensor_inline views over the plain +// buffer binds (no host MTLTensor). Data is integer-valued (and scales are powers of two), +// so every float accumulation order is exact — GPU must equal the CPU-replayed stub EXACTLY. +// The emitted-source asserts run on every platform; the dispatch half is Apple-only. + +require dastest/testing_boost public +require strings +require metal/msl_shader +require daslib/math_bits // float_bits_to_uint — bf16 halfword construction +require llvm/daslib/f16_cvt // f32_to_f16 — the f16s scale plane +require _metal_common // nolint:STYLE030 — used only inside the Apple static_if half; off-Apple lint sees the branch compiled out + +let M = 8 +let N = 16 +let KK = 64 +let NB = KK / 32 +let SGS = 2u + +class TmmKernels { + @ssbo @binding = 0 xa : array // A: M x KK f32 activations (bf16 case) + @ssbo @binding = 1 wb : array // W: N x KK bf16 halfwords + @ssbo @binding = 2 oc : array // C: M x N + @ssbo @binding = 3 aq : array // A: M x KK int8 quants (q8 case) + @ssbo @binding = 4 asf : array // as: M x NB activation block scales + @ssbo @binding = 5 wq : array // W: N x KK int8 quants + @ssbo @binding = 6 wsf : array // ws: N x NB weight block scales + @ssbo @binding = 7 wsh : array // ws as binary16 halfwords (the f16s twin) + + [metal_kernel(name="tmm_bf16_msl"), marker(no_coverage)] + def tmm_bf16_k { + var ap = unsafe(addr(xa[0])) + var wp = unsafe(addr(wb[0])) + var cp = unsafe(addr(oc[0])) + tmm2d_f32_bf16_f32(8u, 16u, 2u, ap, 64u, wp, 64u, cp, 16u, 64u) + } + + [metal_kernel(name="tmm_q8_msl"), marker(no_coverage)] + def tmm_q8_k { + var ap = unsafe(addr(aq[0])) + var sp = unsafe(addr(asf[0])) + var wp = unsafe(addr(wq[0])) + var vp = unsafe(addr(wsf[0])) + var cp = unsafe(addr(oc[0])) + tmm2d_q8_f32(8u, 16u, 2u, ap, 64u, sp, 2u, wp, 64u, vp, 2u, cp, 16u, 64u) + } + + [metal_kernel(name="tmm_q8h_msl"), marker(no_coverage)] + def tmm_q8h_k { + var ap = unsafe(addr(aq[0])) + var sp = unsafe(addr(asf[0])) + var wp = unsafe(addr(wq[0])) + var vp = unsafe(addr(wsh[0])) + var cp = unsafe(addr(oc[0])) + tmm2d_q8_f16s(8u, 16u, 2u, ap, 64u, sp, 2u, wp, 64u, vp, 2u, cp, 16u, 64u) + } +} + +def private fill_oracle(var o : TmmKernels?) { + o.xa |> resize(M * KK) + o.wb |> resize(N * KK) + o.oc |> resize(M * N) + o.aq |> resize(M * KK) + o.asf |> resize(M * NB) + o.wq |> resize(N * KK) + o.wsf |> resize(N * NB) + o.wsh |> resize(N * NB) + for (i in range(M * KK)) { + o.xa[i] = float(i % 7 - 3) // integer-valued -> exact in any sum order + o.aq[i] = int8(i % 17 - 8) + } + for (i in range(N * KK)) { + // bf16 halfword of an integer value: float bits >> 16 (values here are bf16-exact) + o.wb[i] = uint16(float_bits_to_uint(float(i % 5 - 2)) >> 16u) + o.wq[i] = int8((i * 3) % 15 - 7) + } + for (i in range(M * NB)) { + o.asf[i] = float(1 << (i % 3)) // powers of two: scale folds exact + } + for (i in range(N * NB)) { + let v = 0.5 * float(1 << (i % 4)) + o.wsf[i] = v + o.wsh[i] = uint16(f32_to_f16(v)) // powers of two are f16-exact + } +} + +[test] +def test_tensor_ops(t : T?) { + t |> run("tmm2d lowering: generated helpers + MPP preamble") @(t : T?) { + t |> success(find(tmm_bf16_msl, "__das_tmm2d_bf_8_16_2") >= 0, "bf helper generated") + t |> success(find(tmm_q8_msl, "__das_tmm2d_q8_8_16_2") >= 0, "q8 helper generated") + t |> success(find(tmm_q8h_msl, "__das_tmm2d_q8h_8_16_2") >= 0, "q8h helper generated") + t |> success(find(tmm_bf16_msl, "MetalPerformancePrimitives") >= 0, "MPP include present") + t |> success(find(tmm_bf16_msl, "mpp::tensor_ops") >= 0, "namespace present") + t |> success(find(tmm_bf16_msl, "__das_tmm2d_bf_8_16_2(") < find(tmm_bf16_msl, "kernel void"), "helper precedes the kernel") + } + t |> run("tmm2d: GPU == CPU-replay exact") @(t : T?) { + var oracle = new TmmKernels() + fill_oracle(oracle) + gl_LocalInvocationID = uint3(0u, 0u, 0u) + var want_bf : array + var want_q8 : array + var want_q8h : array + oracle->tmm_bf16_k() + want_bf := oracle.oc + oracle->tmm_q8_k() + want_q8 := oracle.oc + oracle->tmm_q8h_k() + want_q8h := oracle.oc + var nz = 0 + for (v in want_q8) { + nz += v != 0.0 ? 1 : 0 + } + t |> success(nz > M * N / 2, "q8 reference is non-trivial") + static_if (typeinfo builtin_module_exists(das_metal)) { + var gpu_bad = -1 + with_metal_device() <| $(dev : MetalDevice?) { + if (dev == null) { + feint("no Metal device on this box; GPU compare skipped\n") + return + } + var queue = metal_new_command_queue(dev) + var bxa = buf_upload(dev, oracle.xa) + var bwb = buf_upload(dev, oracle.wb) + var baq = buf_upload(dev, oracle.aq) + var basf = buf_upload(dev, oracle.asf) + var bwq = buf_upload(dev, oracle.wq) + var bwsf = buf_upload(dev, oracle.wsf) + var bwsh = buf_upload(dev, oracle.wsh) + gpu_bad = 0 + for (which in range(3)) { + let src = which == 0 ? tmm_bf16_msl : (which == 1 ? tmm_q8_msl : tmm_q8h_msl) + let entry = which == 0 ? tmm_bf16_msl_entry : (which == 1 ? tmm_q8_msl_entry : tmm_q8h_msl_entry) + var perr : string + var pso = pipeline_from_source(dev, src, entry, false, perr) + t |> success(pso != null, "pipeline {entry}: {perr}") + if (pso == null) { + gpu_bad++ + continue + } + var boc = buf_fill(dev, M * N, 0.0) + let tew = metal_pipeline_thread_execution_width(pso) + var bufs <- [bxa, bwb, boc, baq, basf, bwq, bwsf, bwsh] + var rerr : string + let tpg = uint(tew) * SGS + let ran = run_compute_1d(queue, pso, bufs, tpg, tpg, rerr) + t |> success(ran, "dispatch {entry}: {rerr}") + if (ran) { + let bad = buf_mismatch_exact(boc, which == 0 ? want_bf : (which == 1 ? want_q8 : want_q8h)) + if (bad != 0) { + var got <- buf_download(boc, M * N, type) + let wref & = unsafe(which == 0 ? want_bf : (which == 1 ? want_q8 : want_q8h)) + var shown = 0 + for (i in range(M * N)) { + if (got[i] != wref[i] && shown < 4) { + print("MISMATCH {entry} [{i / N},{i % N}] gpu={got[i]} want={wref[i]}\n") + shown++ + } + } + delete got + } + t |> equal(bad, 0) + gpu_bad += bad + } else { + gpu_bad++ + } + bufs |> clear() // non-owning handles — clear() skips pointee finalize + unsafe { + delete bufs + } + metal_release(boc) + metal_release(pso) + } + metal_release(bxa) + metal_release(bwb) + metal_release(baq) + metal_release(basf) + metal_release(bwq) + metal_release(bwsf) + metal_release(bwsh) + metal_release(queue) + } + if (gpu_bad >= 0) { + t |> equal(gpu_bad, 0) + } + t |> equal(metal_live_object_count(), 0l) + } else { + feint("das_metal is not built on this platform; CPU-reference half only\n") + } + unsafe { + delete oracle + } + } +} From aeadee1dfcd466b984864e0fbf855de3a8f843ee Mon Sep 17 00:00:00 2001 From: Boris Batkin Date: Mon, 27 Jul 2026 21:31:45 -0700 Subject: [PATCH 04/18] dasLLAMA: Bf16MulMm tensor twin behind the mulmm_bf16 crown MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit MetalBf16MulMmT — the first tune-raced tensor pso twin: same binds, same 32x64 tile and 128-thread dispatch as the simdgroup kernel, body lowered to a single tmm2d_f32_bf16_f32. Crown consulted once at pso build (runtime.metal_tensor); tgmem bind skipped when the twin is selected. The access classifier learns tmm2d pointer operands (buffer-pointer locals resolve to their tracked buffer; C writes, the rest read). Gate: gemma4e-e4b matrix row with a force-crowned manifest — PLE GPU pre-step rides the twin, token-for-token vs CPU truth green. Co-Authored-By: Claude Fable 5 --- .../dasllama/dasllama_kernel_access.das | 66 +++++++++++++++++++ .../dasllama/dasllama_metal_prefill.das | 31 ++++++++- 2 files changed, 95 insertions(+), 2 deletions(-) diff --git a/modules/dasLLAMA/dasllama/dasllama_kernel_access.das b/modules/dasLLAMA/dasllama/dasllama_kernel_access.das index 6e9bbc58b1..e7e2325ae3 100644 --- a/modules/dasLLAMA/dasllama/dasllama_kernel_access.das +++ b/modules/dasLLAMA/dasllama/dasllama_kernel_access.das @@ -83,6 +83,15 @@ def private is_setop2(op : das_string) : bool { def private is_setop1(op : das_string) : bool => op == "++" || op == "--" || op == "+++" || op == "---" +// tmm2d tensor builtins (dasMetal): pointer operands — C is written, the rest are read. +// Returns the C-argument index, or -1 for any other call. +def private tmm2d_c_arg(cname : string) : int { + if (cname == "tmm2d_f32_bf16_f32") { + return 7 + } + return cname == "tmm2d_q8_f32" || cname == "tmm2d_q8_f16s" ? 11 : -1 +} + class private AccessVisitor : AstVisitor { tracked : table fieldMode : bool // tracked names are class FIELDS (metal): match self. @@ -92,6 +101,10 @@ class private AccessVisitor : AstVisitor { callees : table // non-intrinsic call names, for the driver's same-module recursion claimed : table // pure-write target tokens, keyed name:line:column — infer // CLONES subtrees during rewrites, so pointer identity misses + ptrs : table // buffer-pointer local -> its tracked buffer; the init + // (`var p = unsafe(addr(buf[..]))`) is claimed, and each USE + // decides the direction: a modeled intrinsic arg reads or + // writes precisely, anything else restores a conservative read err : string def is_tracked(name : das_string) : bool => key_exists(tracked, string(name)) @@ -113,6 +126,35 @@ class private AccessVisitor : AstVisitor { } } + // record buffer-pointer locals: `var p = unsafe(addr(buf[..]))` (post-infer ExprRef2Ptr; + // pre-infer a bare `addr` call). The buf node is claimed — its direction comes from p's uses. + def override preVisitExprLet(expr : ExprLet?) : void { + for (v in expr.variables) { + var init = v.init + if (init == null) { + continue + } + if (init is ExprRef2Value) { + init = (init as ExprRef2Value).subexpr + } + if (init is ExprUnsafe) { // pre-infer bodies keep the unsafe() wrapper node + init = (init as ExprUnsafe).body + } + var sub : Expression? + if (init is ExprRef2Ptr) { + sub = (init as ExprRef2Ptr).subexpr + } elif (init is ExprCall && (init as ExprCall).name == "addr" && length((init as ExprCall).arguments) == 1) { + sub = (init as ExprCall).arguments[0] + } + continue if (sub == null || !(sub is ExprAt)) + var node : Expression const? + let name = write_root_name(sub, tracked, fieldMode, node) + continue if (empty(name)) + ptrs[string(v.name)] = name + claimed |> insert("{name}:{node.at.line}:{node.at.column}") + } + } + def override preVisitExprCopy(expr : ExprCopy?) : void { note_write(expr.left, false) } @@ -158,6 +200,27 @@ class private AccessVisitor : AstVisitor { || cname == "coopmatLoadTensor" || cname == "coopmatLoadTensorDecode") { return // src argument is a plain read — the generic read pass records it } + let tci = tmm2d_c_arg(cname) + if (tci >= 0) { + // pointer operands resolve through the ptr-map: C writes, the rest read + for (i in range(3, length(expr.arguments))) { + var cur = expr.arguments[i] + if (cur is ExprRef2Value) { + cur = (cur as ExprRef2Value).subexpr + } + continue if (!(cur is ExprVar)) + let pn = string((cur as ExprVar).name) + let bn = ptrs?[pn] ?? (key_exists(tracked, pn) ? pn : "") + continue if (empty(bn)) + claimed |> insert("{pn}:{cur.at.line}:{cur.at.column}") + if (i == tci) { + writes |> insert(bn) + } else { + reads |> insert(bn) + } + } + return + } callees |> insert(cname) // the ratchet: a tracked buffer passed WHOLE to a call we do not model could be written // through a var parameter — refuse to guess. In field mode `self` passed whole is the @@ -185,6 +248,9 @@ class private AccessVisitor : AstVisitor { if (ACCESS_DEBUG) { print("[access-dbg] read {expr.name} at={int(expr.at.line)}:{int(expr.at.column)}\n") } + } elif (key_exists(ptrs, string(expr.name)) && !key_exists(claimed, "{expr.name}:{expr.at.line}:{expr.at.column}")) { + // unmodeled use of a buffer-pointer local — conservative read of its buffer + reads |> insert(ptrs[string(expr.name)]) } return expr } diff --git a/modules/dasLLAMA/dasllama/dasllama_metal_prefill.das b/modules/dasLLAMA/dasllama/dasllama_metal_prefill.das index f757bcb5ca..ff185cfd3c 100644 --- a/modules/dasLLAMA/dasllama/dasllama_metal_prefill.das +++ b/modules/dasLLAMA/dasllama/dasllama_metal_prefill.das @@ -253,6 +253,27 @@ class MetalBf16MulMm { } } +// Metal-4 tensor twin of MetalBf16MulMm: the SAME 32x64 tile per threadgroup off the same +// binds, lowered to mpp matmul2d over tensor_inline views. Raced vs the simdgroup kernel by +// the tuner (runtime.metal_tensor crown "mulmm_bf16"); dispatch shape identical (128 = 4 sg). +class MetalBf16MulMmT { + @ssbo @binding = 0 @role = "weight" wbh : array // bf16 W rows, halfword view + @ssbo @binding = 2 @role = "read" xf : array // raw f32 activations + @ssbo @binding = 3 @role = "write" y : array + @uniform @binding = 4 kdim : uint + @uniform @binding = 5 ndim : uint + + [metal_kernel(name="metal_bf16_mulmm_t_msl")] + def metal_bf16_mulmm_t { + let mBase = gl_WorkGroupID.x * 32u + let nBase = gl_WorkGroupID.y * 64u + var ap = unsafe(addr(xf[mBase * kdim])) + var wp = unsafe(addr(wbh[nBase * kdim])) + var cp = unsafe(addr(y[mBase * ndim + nBase])) + tmm2d_f32_bf16_f32(32u, 64u, 4u, ap, kdim, wp, kdim, cp, ndim, kdim) + } +} + // E-series PLE gather (GPU pre-step, stage 1): dequant token row slices of the Q8_0 // per_layer_token_embd region straight into the LAYER-major side-input panel, sqrt(ple) folded. // One tg per (layer, position) slice; binds ride blob_of's q8 pair at the region base. @@ -4480,6 +4501,7 @@ var private g_pf_pso_cat2 : MetalComputePipeline? var private g_pf_pso_add : MetalComputePipeline? var private g_pf_pso_mm : MetalComputePipeline? // the production das mul_mm (34B q8 blocks + f32 X) var private g_pf_pso_bf16_mm : MetalComputePipeline? // native-BF16 A twin (E-series model_proj) +var private g_pf_bf16_mm_tensor : bool // crowned tensor twin selected (no tgmem bind) var private g_pf_pso_ple_gather : MetalComputePipeline? // PLE pre-step: q8 token-row gather var private g_pf_pso_ple_finish : MetalComputePipeline? // PLE pre-step: rms + combine, in place var private g_pso_qkmm : MetalComputePipeline? @@ -4841,7 +4863,10 @@ def private metal_prefill_init : bool { g_pf_pso_cat2 = pf_compile_pso(metal_pf_cat2_msl, metal_pf_cat2_msl_entry, metal_pf_cat2_msl_fastmath, ok) g_pf_pso_add = pf_compile_pso(metal_add_msl, metal_add_msl_entry, metal_add_msl_fastmath, ok) g_pf_pso_mm = pf_compile_pso(metal_q8_mulmm_msl, metal_q8_mulmm_msl_entry, metal_q8_mulmm_msl_fastmath, ok) - g_pf_pso_bf16_mm = pf_compile_pso(metal_bf16_mulmm_msl, metal_bf16_mulmm_msl_entry, metal_bf16_mulmm_msl_fastmath, ok) + g_pf_bf16_mm_tensor = metal_tensor_crowned("mulmm_bf16") + g_pf_pso_bf16_mm = (g_pf_bf16_mm_tensor + ? pf_compile_pso(metal_bf16_mulmm_t_msl, metal_bf16_mulmm_t_msl_entry, metal_bf16_mulmm_t_msl_fastmath, ok) + : pf_compile_pso(metal_bf16_mulmm_msl, metal_bf16_mulmm_msl_entry, metal_bf16_mulmm_msl_fastmath, ok)) g_pf_pso_ple_gather = pf_compile_pso(metal_ple_gather_q8_msl, metal_ple_gather_q8_msl_entry, metal_ple_gather_q8_msl_fastmath, ok) g_pf_pso_ple_finish = pf_compile_pso(metal_ple_finish_msl, metal_ple_finish_msl_entry, metal_ple_finish_msl_fastmath, ok) g_pso_kq_mm4 = pf_compile_pso(metal_kq_mulmm_k4_msl, metal_kq_mulmm_k4_msl_entry, metal_kq_mulmm_k4_msl_fastmath, ok) @@ -6173,7 +6198,9 @@ def metal_prefill_forward(t : Model; var s : Session; npos, start_pos : int64) : kn_dispatch(enc, uint3(uint(c.n_layers), uint(npos), 1u), uint3(uint(ple_n), 1u, 1u)) let bwp = pf_bf16_of(t, t.ple_model_off) kn_pipeline(enc, g_pf_pso_bf16_mm) - kn_tgmem(enc, metal_bf16_mulmm_msl_tgmem, 0) + if (!g_pf_bf16_mm_tensor) { + kn_tgmem(enc, metal_bf16_mulmm_msl_tgmem, 0) + } kn_buffer(enc, bwp.buf, bwp.boff, 0) kn_buffer(enc, bx, 0ul, 2) kn_buffer(enc, bple_proj, 0ul, 3) From abc07407d6aa17b32bdb335b77dc2524eae95648 Mon Sep 17 00:00:00 2001 From: Boris Batkin Date: Mon, 27 Jul 2026 21:38:05 -0700 Subject: [PATCH 05/18] dasLLAMA: tuner metal-twin race -> runtime.metal_tensor crowns MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit metal_tensor_race() races each tensor pso twin against its shipped simdgroup kernel on this box's GPU (compile + output-envelope to qualify, best GPU time wins); tune_kernels runs it after the CPU sweeps and writes the winner list into the sidecar's runtime.metal_tensor. M1 gate: twin qualifies and loses (0.86ms vs 1.5ms) — empty crown list, simdgroup everywhere, exactly the pre-M5 expectation. Co-Authored-By: Claude Fable 5 --- .../dasllama/dasllama_metal_prefill.das | 150 ++++++++++++++++++ modules/dasLLAMA/harness/tune_kernels.das | 26 ++- 2 files changed, 175 insertions(+), 1 deletion(-) diff --git a/modules/dasLLAMA/dasllama/dasllama_metal_prefill.das b/modules/dasLLAMA/dasllama/dasllama_metal_prefill.das index ff185cfd3c..7b94ae5bd8 100644 --- a/modules/dasLLAMA/dasllama/dasllama_metal_prefill.das +++ b/modules/dasLLAMA/dasllama/dasllama_metal_prefill.das @@ -7006,6 +7006,156 @@ var private g_pf_env_capture = true var private g_pf_env_ncb = 0 // 0 = unset, fall back to the layer-derived default [init] +// ===== Metal-4 tensor twin race (the tuner's crowning section) ===== + +struct public MetalTensorRaceResult { + family : string + winner : string // "tensor" | "simdgroup" | "" when the family could not race + base_ms : double + twin_ms : double + note : string +} + +def private race_buf(dev : MetalDevice?; bytes : uint64; src : void?) : MetalBuffer? { + var b = metal_new_buffer(dev, bytes) + if (src != null) { + unsafe { + memcpy(metal_buffer_contents(b), src, bytes) + } + } + return b +} + +def private race_uniform_u32(dev : MetalDevice?; v : uint) : MetalBuffer? { + var b = metal_new_buffer(dev, 4ul) + unsafe { + var p = reinterpret(metal_buffer_contents(b)) + p[0] = v + } + return b +} + +// best GPU time over `reps` dispatches after one warmup; <0 on dispatch failure +def private race_time_ms(queue : MetalCommandQueue?; reps : int; blk : block<(enc : MetalComputeEncoder?) : void>) : double { + var best = -1.0lf + var err = "" + for (r in range(reps + 1)) { + var gpu_ms = 0.0lf + if (!with_compute_encoder_timed(queue, err, gpu_ms) $(enc) { invoke(blk, enc) }) { + return -1.0lf + } + if (r > 0 && (best < 0.0lf || gpu_ms < best)) { + best = gpu_ms + } + } + return best +} + +def private race_mulmm_bf16(dev : MetalDevice?; queue : MetalCommandQueue?) : MetalTensorRaceResult { + var res = MetalTensorRaceResult(family = "mulmm_bf16", winner = "", base_ms = -1.0lf, twin_ms = -1.0lf, note = "") + let m = 512 + let kdim = 2048 + let ndim = 1024 + var err = "" + var base_pso = pipeline_from_source(dev, metal_bf16_mulmm_msl, metal_bf16_mulmm_msl_entry, metal_bf16_mulmm_msl_fastmath, err) + if (base_pso == null) { + res.note = "base pso: {err}" + return res + } + var twin_pso = pipeline_from_source(dev, metal_bf16_mulmm_t_msl, metal_bf16_mulmm_t_msl_entry, metal_bf16_mulmm_t_msl_fastmath, err) + if (twin_pso == null) { + // no tensor support on this box/toolchain — the simdgroup kernel keeps the crown + res.winner = "simdgroup" + res.note = "twin pso: {err}" + metal_release(base_pso) + return res + } + var xa : array + var wb : array + xa |> resize(m * kdim) + wb |> resize(ndim * kdim) + for (i in range(m * kdim)) { + xa[i] = 0.25 * float(i % 17 - 8) + } + for (i in range(ndim * kdim)) { + wb[i] = uint16(float_bits_to_uint(0.125 * float((i * 3) % 13 - 6)) >> 16u) + } + var bxa = race_buf(dev, uint64(m * kdim * 4), unsafe(addr(xa[0]))) + var bwb = race_buf(dev, uint64(ndim * kdim * 2), unsafe(addr(wb[0]))) + var by_base = race_buf(dev, uint64(m * ndim * 4), null) + var by_twin = race_buf(dev, uint64(m * ndim * 4), null) + var bk = race_uniform_u32(dev, uint(kdim)) + var bn = race_uniform_u32(dev, uint(ndim)) + let grid = uint3(uint(m / 32), uint(ndim / 64), 1u) + let tg = uint3(128u, 1u, 1u) + res.base_ms = race_time_ms(queue, 5) $(enc) { + kn_pipeline(enc, base_pso) + kn_tgmem(enc, metal_bf16_mulmm_msl_tgmem, 0) + kn_buffer(enc, bwb, 0ul, 0) + kn_buffer(enc, bxa, 0ul, 2) + kn_buffer(enc, by_base, 0ul, 3) + kn_buffer(enc, bk, 0ul, 4) + kn_buffer(enc, bn, 0ul, 5) + kn_dispatch(enc, grid, tg) + } + res.twin_ms = race_time_ms(queue, 5) $(enc) { + kn_pipeline(enc, twin_pso) + kn_buffer(enc, bwb, 0ul, 0) + kn_buffer(enc, bxa, 0ul, 2) + kn_buffer(enc, by_twin, 0ul, 3) + kn_buffer(enc, bk, 0ul, 4) + kn_buffer(enc, bn, 0ul, 5) + kn_dispatch(enc, grid, tg) + } + if (res.base_ms < 0.0lf || res.twin_ms < 0.0lf) { + res.winner = res.base_ms < 0.0lf ? "" : "simdgroup" + res.note = "dispatch failed" + } else { + // eligibility envelope: f16-staging noise passes, a garbage twin lands at O(max) + var maxref = 0.0lf + var maxdiff = 0.0lf + unsafe { + let pb = reinterpret(metal_buffer_contents(by_base)) + let pt = reinterpret(metal_buffer_contents(by_twin)) + for (i in range(m * ndim)) { + maxref = max(maxref, abs(double(pb[i]))) + maxdiff = max(maxdiff, abs(double(pt[i]) - double(pb[i]))) + } + } + if (maxdiff > 0.05lf * maxref + 1e-6lf) { + res.winner = "simdgroup" + res.note = "twin OUTPUT MISMATCH maxdiff={maxdiff} maxref={maxref}" + } else { + res.winner = res.twin_ms < res.base_ms ? "tensor" : "simdgroup" + } + } + metal_release(bxa) + metal_release(bwb) + metal_release(by_base) + metal_release(by_twin) + metal_release(bk) + metal_release(bn) + metal_release(base_pso) + metal_release(twin_pso) + return res +} + +//! Tuner entry (the runtime.metal_tensor crown): race every tensor pso twin against its +//! shipped simdgroup kernel on THIS box — compile + output-match to qualify, faster GPU time +//! takes the family's crown. Empty when there is no Metal device. +def metal_tensor_race : array { + var results : array + var dev = metal_create_system_default_device() + if (dev == null) { + return <- results + } + var queue = metal_new_command_queue(dev) + results |> emplace(race_mulmm_bf16(dev, queue)) + metal_release(queue) + metal_release(dev) + return <- results +} + def dasllama_metal_prefill_register() { // registered but dormant: only select_prefill_override("metal") (env rail // DASLLAMA_PIN_PREFILL) activates it, and unsupported shapes decline per call diff --git a/modules/dasLLAMA/harness/tune_kernels.das b/modules/dasLLAMA/harness/tune_kernels.das index 3250a481de..b4cca52b22 100644 --- a/modules/dasLLAMA/harness/tune_kernels.das +++ b/modules/dasLLAMA/harness/tune_kernels.das @@ -21,6 +21,7 @@ require dasllama/dasllama_quant // quantize templates + the Q4/Q8 reference qua require dasllama/dasllama_common // the runtime-knob getters (the profile's "runtime" section snapshots them) require dasllama/dasllama_tune // [dasllama_grid] emitter require llvm/daslib/llvm_tune // tune_sidecar_merge — the app-sidecar upsert writer +require ?das_metal dasllama/dasllama_metal_prefill // nolint:STYLE030 — metal_tensor_race (Apple static_if half) require tuner_cli // --tune-fast, shared with the other half require llvm/daslib/aarch64_neon // nolint:STYLE030 — sdot4_laneq (referenced by the cloned laneq template bodies) require llvm/daslib/f16_cvt // f16 references + the cloned f16-kernel template bodies @@ -28,6 +29,7 @@ require math // mad (referenced by the cloned template bodie require daslib/json require daslib/json_boost require daslib/fio +require daslib/strings_boost // join — the metal_tensor crown list // The shipped-fallback registry: dasllama_tuned_fallbacks() maps every [tuned] kernel to the // perm it falls back to ON THIS BOX, so report() never restates one as a literal. @@ -2044,6 +2046,28 @@ def main { print("note: TUNED_KERNEL_COUNT is {TUNED_KERNEL_COUNT} but {length(g_results)} kernels ran — the progress bar was scaled wrong; update the constant\n") } + // Metal-4 tensor twin race: each tensor pso twin runs against its shipped simdgroup kernel + // on this box's GPU; winners join the runtime "metal_tensor" crown list (empty = simdgroup + // everywhere — the pre-M5 expectation). Skips cleanly with no Metal device. + var metal_crowns : string + static_if (typeinfo builtin_module_exists(das_metal)) { + var mt_wins : array + var mt_ran = false + for (r in metal_tensor_race()) { + mt_ran = true + let mt_note = empty(r.note) ? "" : " ({r.note})" + tune_detail("METAL_TWIN {r.family}: base={r.base_ms}ms tensor={r.twin_ms}ms -> {r.winner}{mt_note}\n") + if (r.winner == "tensor") { + mt_wins |> push(r.family) + } + } + if (!mt_ran) { + tune_detail("METAL_TWIN: no Metal device — tensor race skipped\n") + } + metal_crowns = join(mt_wins, ",") + delete mt_wins + } + // softmax_sink shares softmax's loop shape exactly (one extra scalar joins the reduction), so it // mirrors softmax's winner instead of paying its own sweep. // The "runtime" section snapshots the CURRENT runtime knobs (defaults unless this process changed @@ -2063,7 +2087,7 @@ def main { act_par_threshold = get_act_par_threshold(), q8_chunks_per_job = get_q8_chunks_per_job(), q4_chunks_per_job = get_q4_chunks_per_job(), q8_batch_chunks_per_job = get_q8_batch_chunks_per_job(), jobque_spin_us = get_jobque_spin_us(), jobque_join_poll = get_jobque_join_poll(), - threads = get_total_hw_jobs())) + threads = get_total_hw_jobs(), metal_tensor = metal_crowns)) var winners <- { "dot" => dotw, "axpy" => axpyw, "dot_f16" => dotf16w, "axpy_f16" => axpyf16w, "cvt_f32_to_f16" => cvt16w, "cvt_f16_to_f32" => cvt32w, "add_inplace" => addw, From ac799a1dc604d33d42a7c24d2fb64d658a111f2d Mon Sep 17 00:00:00 2001 From: Boris Batkin Date: Mon, 27 Jul 2026 21:45:32 -0700 Subject: [PATCH 06/18] =?UTF-8?q?dasLLAMA:=20Q8MulMm=20tensor=20twin=20(tm?= =?UTF-8?q?m2d=5Fq8b=5Ff32=20=E2=80=94=20interleaved=20q8=5F0=20W=20x=20f3?= =?UTF-8?q?2=20X)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit New tmm2d_q8b_f32 builtin for the production MulMm dataflow: 34B-block W (two views of one blob) x raw f32 activations, mixed i8xf32 matmul2d chunks with the per-(col,chunk) W-scale fold. Conformance kernel is GPU==CPU exact through the interleaved 34-stride views. MetalQ8MulMmT rides the "mulmm_q8" crown on g_pf_pso_mm; race entry added (loses on M1 fallback: 1.79ms vs 0.71ms, envelope-clean). Also restores the [init] adjacency on dasllama_metal_prefill_register that the race section had split (unknown-override failure), and moves the race to end-of-module. Gate: llama prefill parity base+cont with the crown force-applied — gemm=mulmm rode the twin, token-for-token vs CPU green. Co-Authored-By: Claude Fable 5 --- .../dasllama/dasllama_kernel_access.das | 2 +- .../dasllama/dasllama_metal_prefill.das | 169 ++++++++++++++++-- modules/dasMetal/metal/metal_builtins.das | 28 +++ modules/dasMetal/metal/msl_emit.das | 40 ++++- tests/metal/test_metal_tensor_ops.das | 40 ++++- 5 files changed, 247 insertions(+), 32 deletions(-) diff --git a/modules/dasLLAMA/dasllama/dasllama_kernel_access.das b/modules/dasLLAMA/dasllama/dasllama_kernel_access.das index e7e2325ae3..a9320dd3d3 100644 --- a/modules/dasLLAMA/dasllama/dasllama_kernel_access.das +++ b/modules/dasLLAMA/dasllama/dasllama_kernel_access.das @@ -86,7 +86,7 @@ def private is_setop1(op : das_string) : bool => op == "++" || op == "--" || op // tmm2d tensor builtins (dasMetal): pointer operands — C is written, the rest are read. // Returns the C-argument index, or -1 for any other call. def private tmm2d_c_arg(cname : string) : int { - if (cname == "tmm2d_f32_bf16_f32") { + if (cname == "tmm2d_f32_bf16_f32" || cname == "tmm2d_q8b_f32") { return 7 } return cname == "tmm2d_q8_f32" || cname == "tmm2d_q8_f16s" ? 11 : -1 diff --git a/modules/dasLLAMA/dasllama/dasllama_metal_prefill.das b/modules/dasLLAMA/dasllama/dasllama_metal_prefill.das index 7b94ae5bd8..a1ee6ea923 100644 --- a/modules/dasLLAMA/dasllama/dasllama_metal_prefill.das +++ b/modules/dasLLAMA/dasllama/dasllama_metal_prefill.das @@ -253,6 +253,30 @@ class MetalBf16MulMm { } } +// Metal-4 tensor twin of MetalQ8MulMm: the SAME 32x64 tile per threadgroup off the same binds +// (blob bound twice — half-scale + byte views), lowered to one tmm2d_q8b_f32. Raced vs the +// simdgroup kernel by the tuner (runtime.metal_tensor crown "mulmm_q8"); dispatch identical. +class MetalQ8MulMmT { + @ssbo @binding = 0 @role = "weight" wsh : array // 34B-block W blob, half-scale view + @ssbo @binding = 1 @role = "weight" wqb : array // the SAME blob buffer, byte view + @ssbo @binding = 2 @role = "read" xf : array // raw f32 activations + @ssbo @binding = 3 @role = "write" y : array + @uniform @binding = 4 kdim : uint + @uniform @binding = 5 ndim : uint + + [metal_kernel(name="metal_q8_mulmm_t_msl")] + def metal_q8_mulmm_t { + let mBase = gl_WorkGroupID.x * 32u + let nBase = gl_WorkGroupID.y * 64u + let blk0 = nBase * (kdim / 32u) + var sp = unsafe(addr(wsh[blk0 * 17u])) + var wp = unsafe(addr(wqb[blk0 * 34u + 2u])) + var ap = unsafe(addr(xf[mBase * kdim])) + var cp = unsafe(addr(y[mBase * ndim + nBase])) + tmm2d_q8b_f32(32u, 64u, 4u, sp, wp, ap, kdim, cp, ndim, kdim) + } +} + // Metal-4 tensor twin of MetalBf16MulMm: the SAME 32x64 tile per threadgroup off the same // binds, lowered to mpp matmul2d over tensor_inline views. Raced vs the simdgroup kernel by // the tuner (runtime.metal_tensor crown "mulmm_bf16"); dispatch shape identical (128 = 4 sg). @@ -4502,6 +4526,7 @@ var private g_pf_pso_add : MetalComputePipeline? var private g_pf_pso_mm : MetalComputePipeline? // the production das mul_mm (34B q8 blocks + f32 X) var private g_pf_pso_bf16_mm : MetalComputePipeline? // native-BF16 A twin (E-series model_proj) var private g_pf_bf16_mm_tensor : bool // crowned tensor twin selected (no tgmem bind) +var private g_pf_mm_tensor : bool // ditto for the production q8 mul_mm var private g_pf_pso_ple_gather : MetalComputePipeline? // PLE pre-step: q8 token-row gather var private g_pf_pso_ple_finish : MetalComputePipeline? // PLE pre-step: rms + combine, in place var private g_pso_qkmm : MetalComputePipeline? @@ -4862,7 +4887,10 @@ def private metal_prefill_init : bool { g_pf_pso_copy = pf_compile_pso(metal_pf_copy_msl, metal_pf_copy_msl_entry, metal_pf_copy_msl_fastmath, ok) g_pf_pso_cat2 = pf_compile_pso(metal_pf_cat2_msl, metal_pf_cat2_msl_entry, metal_pf_cat2_msl_fastmath, ok) g_pf_pso_add = pf_compile_pso(metal_add_msl, metal_add_msl_entry, metal_add_msl_fastmath, ok) - g_pf_pso_mm = pf_compile_pso(metal_q8_mulmm_msl, metal_q8_mulmm_msl_entry, metal_q8_mulmm_msl_fastmath, ok) + g_pf_mm_tensor = metal_tensor_crowned("mulmm_q8") + g_pf_pso_mm = (g_pf_mm_tensor + ? pf_compile_pso(metal_q8_mulmm_t_msl, metal_q8_mulmm_t_msl_entry, metal_q8_mulmm_t_msl_fastmath, ok) + : pf_compile_pso(metal_q8_mulmm_msl, metal_q8_mulmm_msl_entry, metal_q8_mulmm_msl_fastmath, ok)) g_pf_bf16_mm_tensor = metal_tensor_crowned("mulmm_bf16") g_pf_pso_bf16_mm = (g_pf_bf16_mm_tensor ? pf_compile_pso(metal_bf16_mulmm_t_msl, metal_bf16_mulmm_t_msl_entry, metal_bf16_mulmm_t_msl_fastmath, ok) @@ -5043,7 +5071,9 @@ def private prefill_decline(t : Model; s : Session; npos, start_pos : int64) : M // yoff places the K/V outputs past a continuation's existing panel rows. def private enc_gemm_mm(enc : MetalComputeEncoder?; bwblob : MetalBuffer?; wboff : uint64; bx, by, bk, bn : MetalBuffer?; mp, d : int64; yoff : uint64 = 0ul) { kn_pipeline(enc, g_pf_pso_mm) - kn_tgmem(enc, metal_q8_mulmm_msl_tgmem, 0) + if (!g_pf_mm_tensor) { + kn_tgmem(enc, metal_q8_mulmm_msl_tgmem, 0) + } kn_buffer(enc, bwblob, wboff, 0) kn_buffer(enc, bwblob, wboff, 1) kn_buffer(enc, bx, 0ul, 2) @@ -7006,6 +7036,23 @@ var private g_pf_env_capture = true var private g_pf_env_ncb = 0 // 0 = unset, fall back to the layer-derived default [init] +def dasllama_metal_prefill_register() { + // registered but dormant: only select_prefill_override("metal") (env rail + // DASLLAMA_PIN_PREFILL) activates it, and unsupported shapes decline per call + register_prefill_override("metal", @@metal_prefill_forward) + register_ple_gpu_gate(@@metal_ple_pre_gpu_gate) + g_pf_env_mulmm = env_flag("DASLLAMA_METAL_MULMM", true) + g_pf_env_logits = env_flag("DASLLAMA_METAL_LOGITS", true) + g_pf_env_attn = env_flag("DASLLAMA_METAL_ATTN", true) + g_pf_env_unretained = env_flag("DASLLAMA_METAL_UNRETAINED", false) + g_pf_env_capture = env_flag("DASLLAMA_METAL_PF_CAPTURE", true) + g_pf_env_ncb = env_int("DASLLAMA_METAL_NCB", 0) + if (has_env_variable("DASLLAMA_METAL_PREFILL_SKIP")) { + g_pf_skip = get_env_variable("DASLLAMA_METAL_PREFILL_SKIP") + to_log(LOG_WARNING, "dasLLAMA metal prefill: SKIP={g_pf_skip} — timing-attribution mode, output is GARBAGE\n") + } +} + // ===== Metal-4 tensor twin race (the tuner's crowning section) ===== struct public MetalTensorRaceResult { @@ -7140,6 +7187,106 @@ def private race_mulmm_bf16(dev : MetalDevice?; queue : MetalCommandQueue?) : Me return res } +def private race_mulmm_q8(dev : MetalDevice?; queue : MetalCommandQueue?) : MetalTensorRaceResult { + var res = MetalTensorRaceResult(family = "mulmm_q8", winner = "", base_ms = -1.0lf, twin_ms = -1.0lf, note = "") + let m = 512 + let kdim = 2048 + let ndim = 1024 + let nkb = kdim / 32 + var err = "" + var base_pso = pipeline_from_source(dev, metal_q8_mulmm_msl, metal_q8_mulmm_msl_entry, metal_q8_mulmm_msl_fastmath, err) + if (base_pso == null) { + res.note = "base pso: {err}" + return res + } + var twin_pso = pipeline_from_source(dev, metal_q8_mulmm_t_msl, metal_q8_mulmm_t_msl_entry, metal_q8_mulmm_t_msl_fastmath, err) + if (twin_pso == null) { + res.winner = "simdgroup" + res.note = "twin pso: {err}" + metal_release(base_pso) + return res + } + var xa : array + var blob : array + xa |> resize(m * kdim) + blob |> resize(ndim * nkb * 34) + for (i in range(m * kdim)) { + xa[i] = 0.25 * float(i % 17 - 8) + } + for (blk in range(ndim * nkb)) { + for (k in range(32)) { + blob[blk * 34 + 2 + k] = int8((blk * 7 + k * 3) % 19 - 9) + } + } + unsafe { + var p8 = addr(blob[0]) + for (blk in range(ndim * nkb)) { + var ph = reinterpret(p8 + blk * 34) + ph[0] = float16(0.5 * float(1 << (blk % 3))) + } + } + var bw = race_buf(dev, uint64(ndim * nkb * 34), unsafe(addr(blob[0]))) + var bxa = race_buf(dev, uint64(m * kdim * 4), unsafe(addr(xa[0]))) + var by_base = race_buf(dev, uint64(m * ndim * 4), null) + var by_twin = race_buf(dev, uint64(m * ndim * 4), null) + var bk = race_uniform_u32(dev, uint(kdim)) + var bn = race_uniform_u32(dev, uint(ndim)) + let grid = uint3(uint(m / 32), uint(ndim / 64), 1u) + let tg = uint3(128u, 1u, 1u) + res.base_ms = race_time_ms(queue, 5) $(enc) { + kn_pipeline(enc, base_pso) + kn_tgmem(enc, metal_q8_mulmm_msl_tgmem, 0) + kn_buffer(enc, bw, 0ul, 0) + kn_buffer(enc, bw, 0ul, 1) + kn_buffer(enc, bxa, 0ul, 2) + kn_buffer(enc, by_base, 0ul, 3) + kn_buffer(enc, bk, 0ul, 4) + kn_buffer(enc, bn, 0ul, 5) + kn_dispatch(enc, grid, tg) + } + res.twin_ms = race_time_ms(queue, 5) $(enc) { + kn_pipeline(enc, twin_pso) + kn_buffer(enc, bw, 0ul, 0) + kn_buffer(enc, bw, 0ul, 1) + kn_buffer(enc, bxa, 0ul, 2) + kn_buffer(enc, by_twin, 0ul, 3) + kn_buffer(enc, bk, 0ul, 4) + kn_buffer(enc, bn, 0ul, 5) + kn_dispatch(enc, grid, tg) + } + if (res.base_ms < 0.0lf || res.twin_ms < 0.0lf) { + res.winner = res.base_ms < 0.0lf ? "" : "simdgroup" + res.note = "dispatch failed" + } else { + // eligibility envelope: f16-staging noise passes, a garbage twin lands at O(max) + var maxref = 0.0lf + var maxdiff = 0.0lf + unsafe { + let pb = reinterpret(metal_buffer_contents(by_base)) + let pt = reinterpret(metal_buffer_contents(by_twin)) + for (i in range(m * ndim)) { + maxref = max(maxref, abs(double(pb[i]))) + maxdiff = max(maxdiff, abs(double(pt[i]) - double(pb[i]))) + } + } + if (maxdiff > 0.05lf * maxref + 1e-6lf) { + res.winner = "simdgroup" + res.note = "twin OUTPUT MISMATCH maxdiff={maxdiff} maxref={maxref}" + } else { + res.winner = res.twin_ms < res.base_ms ? "tensor" : "simdgroup" + } + } + metal_release(bw) + metal_release(bxa) + metal_release(by_base) + metal_release(by_twin) + metal_release(bk) + metal_release(bn) + metal_release(base_pso) + metal_release(twin_pso) + return res +} + //! Tuner entry (the runtime.metal_tensor crown): race every tensor pso twin against its //! shipped simdgroup kernel on THIS box — compile + output-match to qualify, faster GPU time //! takes the family's crown. Empty when there is no Metal device. @@ -7151,24 +7298,8 @@ def metal_tensor_race : array { } var queue = metal_new_command_queue(dev) results |> emplace(race_mulmm_bf16(dev, queue)) + results |> emplace(race_mulmm_q8(dev, queue)) metal_release(queue) metal_release(dev) return <- results } - -def dasllama_metal_prefill_register() { - // registered but dormant: only select_prefill_override("metal") (env rail - // DASLLAMA_PIN_PREFILL) activates it, and unsupported shapes decline per call - register_prefill_override("metal", @@metal_prefill_forward) - register_ple_gpu_gate(@@metal_ple_pre_gpu_gate) - g_pf_env_mulmm = env_flag("DASLLAMA_METAL_MULMM", true) - g_pf_env_logits = env_flag("DASLLAMA_METAL_LOGITS", true) - g_pf_env_attn = env_flag("DASLLAMA_METAL_ATTN", true) - g_pf_env_unretained = env_flag("DASLLAMA_METAL_UNRETAINED", false) - g_pf_env_capture = env_flag("DASLLAMA_METAL_PF_CAPTURE", true) - g_pf_env_ncb = env_int("DASLLAMA_METAL_NCB", 0) - if (has_env_variable("DASLLAMA_METAL_PREFILL_SKIP")) { - g_pf_skip = get_env_variable("DASLLAMA_METAL_PREFILL_SKIP") - to_log(LOG_WARNING, "dasLLAMA metal prefill: SKIP={g_pf_skip} — timing-attribution mode, output is GARBAGE\n") - } -} diff --git a/modules/dasMetal/metal/metal_builtins.das b/modules/dasMetal/metal/metal_builtins.das index 5cdee1a7b9..c9a0b21722 100644 --- a/modules/dasMetal/metal/metal_builtins.das +++ b/modules/dasMetal/metal/metal_builtins.das @@ -349,3 +349,31 @@ def public tmm2d_q8_f16s(m, n, sgs : uint; aq : int8 const?; lda : uint; asf : f } } } + +//! Interleaved-q8_0 W blocks x RAW f32 activations (the production MulMm dataflow). `wsh` is the +//! blob's f16-scale view at the tile's first block (row j chunk b scale at (j*nb+b)*17); `wqb` +//! the byte view pre-offset to that block's quants (+2 done by the caller — row j chunk b quant k +//! at (j*nb+b)*34 + k). f32 chunk accumulate, W scale folded between chunks. +[unused_argument(sgs), sideeffects] +def public tmm2d_q8b_f32(m, n, sgs : uint; wsh : float16 const?; wqb : int8 const?; a : float const?; lda : uint; var c : float?; ldc : uint; kk : uint) { + if (gl_LocalInvocationID.x != 0u) { + return + } + unsafe { + let nb = int(kk) / 32 + for (r in range(int(m))) { + for (j in range(int(n))) { + var facc = 0.0 + for (b in range(nb)) { + let blk = j * nb + b + var ch = 0.0 + for (k in range(32)) { + ch += a[r * int(lda) + b * 32 + k] * float(wqb[blk * 34 + k]) + } + facc += ch * float(wsh[blk * 17]) + } + c[r * int(ldc) + j] = facc + } + } + } +} diff --git a/modules/dasMetal/metal/msl_emit.das b/modules/dasMetal/metal/msl_emit.das index 77b4932080..3615bdf017 100644 --- a/modules/dasMetal/metal/msl_emit.das +++ b/modules/dasMetal/metal/msl_emit.das @@ -345,8 +345,8 @@ def private scan_body(var ctx : MslCtx; e : Expression?) { if (scn == "simdgroup_store" && length(ecs.arguments) == 4) { mark_write_target(ctx, ecs.arguments[1]) } - // the tensor GEMMs write through their C pointer (arg 7 for bf16, 11 for the q8 forms) - if (scn == "tmm2d_f32_bf16_f32" && length(ecs.arguments) == 10) { + // the tensor GEMMs write through their C pointer (arg 7 for the 10-arg forms, 11 for q8xq8) + if ((scn == "tmm2d_f32_bf16_f32" || scn == "tmm2d_q8b_f32") && length(ecs.arguments) == 10) { mark_write_target(ctx, ecs.arguments[7]) } if ((scn == "tmm2d_q8_f32" || scn == "tmm2d_q8_f16s") && length(ecs.arguments) == 14) { @@ -1199,6 +1199,34 @@ def private tmm2d_helper_text(kind : string; m, n, sgs : uint) : string { w |> write("\}") } } + if (kind == "q8b") { + // interleaved-q8_0 W blocks (34B stride, quants pre-offset +2) x raw f32 activations: + // f32 chunk accumulate, per-(col,chunk) W scale folded between chunks + return build_string() $(var w) { + w |> write("static void {name}(device const half * ws, device const char * wq, device const float * a, uint lda, device float * c, uint ldc, uint kk) \{\n") + w |> write(" constexpr auto d = matmul2d_descriptor({int(m)}, {int(n)}, 32, false, true, false);\n") + w |> write(" matmul2d> op;\n") + w |> write(" auto C = tensor, tensor_inline>(c, dextents({int(n)}, {int(m)}), metal::array\{1, int(ldc)\});\n") + w |> write(" const uint nb = kk / 32u;\n") + w |> write(" auto A0 = tensor, tensor_inline>((device float *)a, dextents(32, {int(m)}), metal::array\{1, int(lda)\});\n") + w |> write(" auto W0 = {ct}((device int8_t *)wq, dextents(32, {int(n)}), metal::array\{1, int(nb * 34u)\});\n") + w |> write(" auto fT = op.get_destination_cooperative_tensor();\n") + w |> write(" for (uint16_t i = 0; i < fT.get_capacity(); ++i) \{ if (fT.is_valid_element(i)) fT[i] = 0.0f; \}\n") + w |> write(" for (uint b = 0; b < nb; ++b) \{\n") + w |> write(" auto A = tensor, tensor_inline>((device float *)(a + b * 32u), dextents(32, {int(m)}), metal::array\{1, int(lda)\});\n") + w |> write(" auto W = {ct}((device int8_t *)(wq + b * 34u), dextents(32, {int(n)}), metal::array\{1, int(nb * 34u)\});\n") + w |> write(" auto cT = op.get_destination_cooperative_tensor();\n") + w |> write(" for (uint16_t i = 0; i < cT.get_capacity(); ++i) \{ if (cT.is_valid_element(i)) cT[i] = 0.0f; \}\n") + w |> write(" op.run(A, W, cT);\n") + w |> write(" for (uint16_t i = 0; i < cT.get_capacity(); ++i) \{ if (cT.is_valid_element(i)) \{\n") + w |> write(" auto ids = cT.get_multidimensional_index(i);\n") + w |> write(" fT[i] += cT[i] * float(ws[(uint(ids[0]) * nb + b) * 17u]);\n") + w |> write(" \} \}\n") + w |> write(" \}\n") + w |> write(" fT.store(C);\n") + w |> write("\}") + } + } // q8 / q8_f16s: per-32-block exact int accumulate + scale fold in a float cooperative acc let wst = kind == "q8h" ? "device const ushort" : "device const float" let wsread = kind == "q8h" ? "float(as_type(ws[uint(ids[0]) * ldws + b]))" : "ws[uint(ids[0]) * ldws + b]" @@ -1230,8 +1258,8 @@ def private tmm2d_helper_text(kind : string; m, n, sgs : uint) : string { def private emit_tmm2d(var ctx : MslCtx; ec : ExprCall?; indent : string; var lines : array) { let cname = call_base_name("{ec.name}") - let kind = cname == "tmm2d_f32_bf16_f32" ? "bf" : (cname == "tmm2d_q8_f16s" ? "q8h" : "q8") - let want_args = kind == "bf" ? 10 : 14 + let kind = cname == "tmm2d_f32_bf16_f32" ? "bf" : (cname == "tmm2d_q8b_f32" ? "q8b" : (cname == "tmm2d_q8_f16s" ? "q8h" : "q8")) + let want_args = kind == "bf" || kind == "q8b" ? 10 : 14 if (length(ec.arguments) != want_args) { err(ctx, ec.at, "{cname} takes exactly {want_args} arguments") return @@ -1246,7 +1274,7 @@ def private emit_tmm2d(var ctx : MslCtx; ec : ExprCall?; indent : string; var li if (!key_exists(ctx.gen_helpers, hname)) { ctx.gen_helpers[hname] = tmm2d_helper_text(kind, m, n, sgs) } - var vals <- [for (i in range(3, length(ec.arguments))); emit_value(ctx, ec.arguments[i])] + let vals <- [for (i in range(3, length(ec.arguments))); emit_value(ctx, ec.arguments[i])] note(ctx, "call.{cname}") lines |> push("{indent}{hname}({join(vals, ", ")});") } @@ -1440,7 +1468,7 @@ def private emit_stmt(var ctx : MslCtx; e : Expression?; indent : string; var li emit_tg_store_float4(ctx, ec, indent, lines) } elif (cname == "simdgroup_multiply_accumulate" || cname == "simdgroup_multiply") { emit_sgmat_multiply(ctx, ec, indent, lines) - } elif (cname == "tmm2d_f32_bf16_f32" || cname == "tmm2d_q8_f32" || cname == "tmm2d_q8_f16s") { + } elif (cname == "tmm2d_f32_bf16_f32" || cname == "tmm2d_q8_f32" || cname == "tmm2d_q8_f16s" || cname == "tmm2d_q8b_f32") { emit_tmm2d(ctx, ec, indent, lines) } else { err(ctx, ec.at, "call `{cname}` as a statement has no MSL lowering — only barrier()/memoryBarrierShared() sync and the simdgroup_matrix ops are supported") diff --git a/tests/metal/test_metal_tensor_ops.das b/tests/metal/test_metal_tensor_ops.das index bb5a81cc1e..cf3b1e52c1 100644 --- a/tests/metal/test_metal_tensor_ops.das +++ b/tests/metal/test_metal_tensor_ops.das @@ -29,6 +29,8 @@ class TmmKernels { @ssbo @binding = 5 wq : array // W: N x KK int8 quants @ssbo @binding = 6 wsf : array // ws: N x NB weight block scales @ssbo @binding = 7 wsh : array // ws as binary16 halfwords (the f16s twin) + @ssbo @binding = 8 bsh : array // q8b: 34B-interleaved W blob, half-scale view + @ssbo @binding = 9 bqb : array // q8b: the SAME blob layout, byte view [metal_kernel(name="tmm_bf16_msl"), marker(no_coverage)] def tmm_bf16_k { @@ -57,6 +59,15 @@ class TmmKernels { var cp = unsafe(addr(oc[0])) tmm2d_q8_f16s(8u, 16u, 2u, ap, 64u, sp, 2u, wp, 64u, vp, 2u, cp, 16u, 64u) } + + [metal_kernel(name="tmm_q8b_msl"), marker(no_coverage)] + def tmm_q8b_k { + var sp = unsafe(addr(bsh[0])) + var wp = unsafe(addr(bqb[2])) + var ap = unsafe(addr(xa[0])) + var cp = unsafe(addr(oc[0])) + tmm2d_q8b_f32(8u, 16u, 2u, sp, wp, ap, 64u, cp, 16u, 64u) + } } def private fill_oracle(var o : TmmKernels?) { @@ -85,6 +96,15 @@ def private fill_oracle(var o : TmmKernels?) { o.wsf[i] = v o.wsh[i] = uint16(f32_to_f16(v)) // powers of two are f16-exact } + // q8b blob: block (j*NB+b) = 2B f16 scale + 32 int8 quants, integer-valued and pow2-scaled + o.bsh |> resize(N * NB * 17) + o.bqb |> resize(N * NB * 34) + for (blk in range(N * NB)) { + o.bsh[blk * 17] = float16(0.5 * float(1 << (blk % 3))) + for (k in range(32)) { + o.bqb[blk * 34 + 2 + k] = int8((blk * 7 + k * 3) % 19 - 9) + } + } } [test] @@ -93,6 +113,7 @@ def test_tensor_ops(t : T?) { t |> success(find(tmm_bf16_msl, "__das_tmm2d_bf_8_16_2") >= 0, "bf helper generated") t |> success(find(tmm_q8_msl, "__das_tmm2d_q8_8_16_2") >= 0, "q8 helper generated") t |> success(find(tmm_q8h_msl, "__das_tmm2d_q8h_8_16_2") >= 0, "q8h helper generated") + t |> success(find(tmm_q8b_msl, "__das_tmm2d_q8b_8_16_2") >= 0, "q8b helper generated") t |> success(find(tmm_bf16_msl, "MetalPerformancePrimitives") >= 0, "MPP include present") t |> success(find(tmm_bf16_msl, "mpp::tensor_ops") >= 0, "namespace present") t |> success(find(tmm_bf16_msl, "__das_tmm2d_bf_8_16_2(") < find(tmm_bf16_msl, "kernel void"), "helper precedes the kernel") @@ -110,6 +131,9 @@ def test_tensor_ops(t : T?) { want_q8 := oracle.oc oracle->tmm_q8h_k() want_q8h := oracle.oc + var want_q8b : array + oracle->tmm_q8b_k() + want_q8b := oracle.oc var nz = 0 for (v in want_q8) { nz += v != 0.0 ? 1 : 0 @@ -130,10 +154,12 @@ def test_tensor_ops(t : T?) { var bwq = buf_upload(dev, oracle.wq) var bwsf = buf_upload(dev, oracle.wsf) var bwsh = buf_upload(dev, oracle.wsh) + var bbsh = buf_upload(dev, oracle.bsh) + var bbqb = buf_upload(dev, oracle.bqb) gpu_bad = 0 - for (which in range(3)) { - let src = which == 0 ? tmm_bf16_msl : (which == 1 ? tmm_q8_msl : tmm_q8h_msl) - let entry = which == 0 ? tmm_bf16_msl_entry : (which == 1 ? tmm_q8_msl_entry : tmm_q8h_msl_entry) + for (which in range(4)) { + let src = which == 0 ? tmm_bf16_msl : (which == 1 ? tmm_q8_msl : (which == 2 ? tmm_q8h_msl : tmm_q8b_msl)) + let entry = which == 0 ? tmm_bf16_msl_entry : (which == 1 ? tmm_q8_msl_entry : (which == 2 ? tmm_q8h_msl_entry : tmm_q8b_msl_entry)) var perr : string var pso = pipeline_from_source(dev, src, entry, false, perr) t |> success(pso != null, "pipeline {entry}: {perr}") @@ -143,16 +169,16 @@ def test_tensor_ops(t : T?) { } var boc = buf_fill(dev, M * N, 0.0) let tew = metal_pipeline_thread_execution_width(pso) - var bufs <- [bxa, bwb, boc, baq, basf, bwq, bwsf, bwsh] + var bufs <- [bxa, bwb, boc, baq, basf, bwq, bwsf, bwsh, bbsh, bbqb] var rerr : string let tpg = uint(tew) * SGS let ran = run_compute_1d(queue, pso, bufs, tpg, tpg, rerr) t |> success(ran, "dispatch {entry}: {rerr}") if (ran) { - let bad = buf_mismatch_exact(boc, which == 0 ? want_bf : (which == 1 ? want_q8 : want_q8h)) + let bad = buf_mismatch_exact(boc, which == 0 ? want_bf : (which == 1 ? want_q8 : (which == 2 ? want_q8h : want_q8b))) if (bad != 0) { var got <- buf_download(boc, M * N, type) - let wref & = unsafe(which == 0 ? want_bf : (which == 1 ? want_q8 : want_q8h)) + let wref & = unsafe(which == 0 ? want_bf : (which == 1 ? want_q8 : (which == 2 ? want_q8h : want_q8b))) var shown = 0 for (i in range(M * N)) { if (got[i] != wref[i] && shown < 4) { @@ -181,6 +207,8 @@ def test_tensor_ops(t : T?) { metal_release(bwq) metal_release(bwsf) metal_release(bwsh) + metal_release(bbsh) + metal_release(bbqb) metal_release(queue) } if (gpu_bad >= 0) { From 241de650c312f7cfc74bc984e96b7eb4f009ffea Mon Sep 17 00:00:00 2001 From: Boris Batkin Date: Mon, 27 Jul 2026 22:12:17 -0700 Subject: [PATCH 07/18] dasLLAMA: batch + MoE Q8 GEMM tensor twins; DSE keeps operator-topped stores MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Q8GemmB/Q8GemmBSk/Q8Gemm64B twins (decode side) and the MetalMoeMulMmQ8 twin — contiguous down-site only, picked per dispatch (bkt-gathered X cannot form a strided tensor view, so gate/up keep the simdgroup kernel even when "moe_mulmm_q8" is crowned). tmm2d_q8b_f32 gains an explicit W-row-stride-in-blocks arg (split-K passes a full-row stride with a k-slice kk). Race scaffolding moves to dasllama_metal_common; the tuner aggregates prefill + decode races (six families, all envelope-clean and losing on M1 fallback). Gate: llama batch decode parity with all crowns forced — 3/3 green. Compiler fix uncovered by the move: dead-store elimination stripped a store whose impure rhs is operator-topped, leaving a bare `/` statement that the top-level no-side-effect lint rejects when a reused shared module re-lints its optimized AST. Such stores now stay whole (src/ast/ast_dse.cpp); the optimization_dead_stores arm re-pins it. Co-Authored-By: Claude Fable 5 --- .../dasllama/dasllama_metal_common.das | 99 +++++++ .../dasllama/dasllama_metal_kernels.das | 222 +++++++++++++- .../dasllama/dasllama_metal_prefill.das | 274 +++++++++++------- modules/dasLLAMA/harness/tune_kernels.das | 11 +- modules/dasMetal/metal/metal_builtins.das | 10 +- modules/dasMetal/metal/msl_emit.das | 13 +- src/ast/ast_dse.cpp | 15 +- tests/language/optimization_dead_stores.das | 8 +- tests/metal/test_metal_tensor_ops.das | 2 +- 9 files changed, 522 insertions(+), 132 deletions(-) diff --git a/modules/dasLLAMA/dasllama/dasllama_metal_common.das b/modules/dasLLAMA/dasllama/dasllama_metal_common.das index a84b08cd56..602af025ce 100644 --- a/modules/dasLLAMA/dasllama/dasllama_metal_common.das +++ b/modules/dasLLAMA/dasllama/dasllama_metal_common.das @@ -307,6 +307,9 @@ var g_pso_gemm_b : MetalComputePipeline? var g_pso_gemm_b_sk : MetalComputePipeline? // split-K twin (per-site policy) var g_pso_skred : MetalComputePipeline? // the split-K plane reduce var g_pso_gemm64_b : MetalComputePipeline? // 64-wide N tile (classifier site) +var g_gemm_b_tensor : bool // crowned tensor twins selected (no tgmem bind) +var g_gemm_b_sk_tensor : bool +var g_gemm64_b_tensor : bool var g_pso_mm : MetalComputePipeline? // the prefill module's MetalQ8MulMm (34B blob) var g_pso_attnb16 : MetalComputePipeline? var g_pso_attnb32 : MetalComputePipeline? @@ -2017,3 +2020,99 @@ def public metal_prefill_declines() : table { return clone_to_move(g_prefill_declines_by_reason) } +// ===== Metal-4 tensor twin race scaffolding (shared by the per-module race sections) ===== + +struct public MetalTensorRaceResult { + family : string + winner : string // "tensor" | "simdgroup" | "" when the family could not race + base_ms : double + twin_ms : double + note : string +} + +def race_buf(dev : MetalDevice?; bytes : uint64; src : void?) : MetalBuffer? { + var b = metal_new_buffer(dev, bytes) + if (src != null) { + unsafe { + memcpy(metal_buffer_contents(b), src, bytes) + } + } + return b +} + +def race_uniform_u32(dev : MetalDevice?; v : uint) : MetalBuffer? { + var b = metal_new_buffer(dev, 4ul) + unsafe { + var p = reinterpret(metal_buffer_contents(b)) + p[0] = v + } + return b +} + +// deterministic 34B-interleaved q8_0 W blob for the races: int quants, pow2 f16 scales +def race_q8_blob(nblk : int) : array { + var blob : array + blob |> resize(nblk * 34) + for (blk in range(nblk)) { + for (k in range(32)) { + blob[blk * 34 + 2 + k] = int8((blk * 7 + k * 3) % 19 - 9) + } + } + unsafe { + var p8 = addr(blob[0]) + for (blk in range(nblk)) { + var ph = reinterpret(p8 + blk * 34) + ph[0] = float16(0.5 * float(1 << (blk % 3))) + } + } + return <- blob +} + +// deterministic f32 activation panel for the races +def race_x_f32(count : int) : array { + var xa : array + xa |> resize(count) + for (i in range(count)) { + xa[i] = 0.25 * float(i % 17 - 8) + } + return <- xa +} + +// eligibility envelope: f16-staging noise passes, a garbage twin lands at O(max) +def race_envelope_ok(ba, bb : MetalBuffer?; count : int; var note : string&) : bool { + var maxref = 0.0lf + var maxdiff = 0.0lf + unsafe { + let pb = reinterpret(metal_buffer_contents(ba)) + let pt = reinterpret(metal_buffer_contents(bb)) + for (i in range(count)) { + maxref = max(maxref, abs(double(pb[i]))) + maxdiff = max(maxdiff, abs(double(pt[i]) - double(pb[i]))) + } + } + if (maxdiff > 0.05lf * maxref + 1e-6lf) { + note = "twin OUTPUT MISMATCH maxdiff={maxdiff} maxref={maxref}" + return false + } + return true +} + +// best GPU time over `reps` dispatches after one warmup; <0 on dispatch failure +def race_time_ms(queue : MetalCommandQueue?; reps : int; blk : block<(enc : MetalComputeEncoder?) : void>) : double { + var best = -1.0lf + var err = "" + for (r in range(reps + 1)) { + var gpu_ms = 0.0lf + var okr = false + okr = with_compute_encoder_timed(queue, err, gpu_ms) $(enc) { + invoke(blk, enc) + } + if (!okr) { + return -1.0lf + } + if (r > 0 && (best < 0.0lf || gpu_ms < best)) { + best = gpu_ms + } + } + return best +} diff --git a/modules/dasLLAMA/dasllama/dasllama_metal_kernels.das b/modules/dasLLAMA/dasllama/dasllama_metal_kernels.das index a02461e7e2..07cd3c96d7 100644 --- a/modules/dasLLAMA/dasllama/dasllama_metal_kernels.das +++ b/modules/dasLLAMA/dasllama/dasllama_metal_kernels.das @@ -3161,6 +3161,80 @@ class MetalSkReduce { } } +// Metal-4 tensor twins of the batch GEMMs: same binds/grids, bodies lowered to one +// tmm2d_q8b_f32 each (raced via runtime.metal_tensor — crowns "gemmb_q8" / "gemmb_sk_q8" / +// "gemm64b_q8"). The split-K twin passes the FULL row stride in blocks (ldwb) with a k-slice kk. +class MetalQ8GemmBT { + @ssbo @binding = 0 @role = "weight" wqb : array // 34B-block W blob, byte view + @ssbo @binding = 2 @role = "weight" wsh : array // the same buffer, half-scale view + @ssbo @binding = 3 @role = "read" x : array // X [M' x n] f32 rows + @ssbo @binding = 4 @role = "write" y : array // Y [M' x d] + @uniform @binding = 5 kdim : uint + @uniform @binding = 6 ndim : uint + @uniform @binding = 7 ys : uint + + [metal_kernel(name="metal_q8_gemm_b_t_msl")] + def metal_q8_gemm_b_t { + let ntileN = ndim / 32u + let mBase = (gl_WorkGroupID.x / ntileN) * 32u + let nBase = (gl_WorkGroupID.x % ntileN) * 32u + let nkb = kdim / 32u + var sp = unsafe(addr(wsh[(nBase * nkb) * 17u])) + var wp = unsafe(addr(wqb[(nBase * nkb) * 34u + 2u])) + var ap = unsafe(addr(x[mBase * kdim])) + var cp = unsafe(addr(y[mBase * ys + nBase])) + tmm2d_q8b_f32(32u, 32u, 4u, sp, wp, ap, kdim, cp, ys, kdim, nkb) + } +} + +class MetalQ8GemmBSkT { + @ssbo @binding = 0 @role = "weight" wqb : array // 34B-block W blob, byte view + @ssbo @binding = 2 @role = "weight" wsh : array // the same buffer, half-scale view + @ssbo @binding = 3 @role = "read" x : array // X [32 x n] f32 rows + @ssbo @binding = 4 @role = "write" part : array // [ksplit x 32 x d] partial-C planes + @uniform @binding = 5 kdim : uint + @uniform @binding = 6 ndim : uint + @uniform @binding = 8 ksplit : uint + + [metal_kernel(name="metal_q8_gemm_b_sk_t_msl")] + def metal_q8_gemm_b_sk_t { + let ntileN = ndim / 32u + let mBase = (gl_WorkGroupID.x / ntileN) * 32u + let nBase = (gl_WorkGroupID.x % ntileN) * 32u + let sl = gl_WorkGroupID.y + let nkb = kdim / 32u + let kbn = nkb / ksplit + let kb0 = sl * kbn + var sp = unsafe(addr(wsh[(nBase * nkb + kb0) * 17u])) + var wp = unsafe(addr(wqb[(nBase * nkb + kb0) * 34u + 2u])) + var ap = unsafe(addr(x[mBase * kdim + kb0 * 32u])) + var cp = unsafe(addr(part[sl * 32u * ndim + mBase * ndim + nBase])) + tmm2d_q8b_f32(32u, 32u, 4u, sp, wp, ap, kdim, cp, ndim, kbn * 32u, nkb) + } +} + +class MetalQ8Gemm64BT { + @ssbo @binding = 0 @role = "weight" wqb : array // 34B-block W blob, byte view + @ssbo @binding = 2 @role = "weight" wsh : array // the same buffer, half-scale view + @ssbo @binding = 3 @role = "read" x : array // X [M' x n] f32 rows + @ssbo @binding = 4 @role = "write" y : array // Y [M' x d] + @uniform @binding = 5 kdim : uint + @uniform @binding = 6 ndim : uint + + [metal_kernel(name="metal_q8_gemm64_b_t_msl")] + def metal_q8_gemm64_b_t { + let ntileN = ndim / 64u + let mBase = (gl_WorkGroupID.x / ntileN) * 32u + let nBase = (gl_WorkGroupID.x % ntileN) * 64u + let nkb = kdim / 32u + var sp = unsafe(addr(wsh[(nBase * nkb) * 17u])) + var wp = unsafe(addr(wqb[(nBase * nkb) * 34u + 2u])) + var ap = unsafe(addr(x[mBase * kdim])) + var cp = unsafe(addr(y[mBase * ndim + nBase])) + tmm2d_q8b_f32(32u, 64u, 4u, sp, wp, ap, kdim, cp, ndim, kdim, nkb) + } +} + // The 64-wide twin of MulMm's 32(M) x 64(N) tile: half the threadgroups (d/64) at ~29% fewer // instructions per element — wins where the grid stays fat (e.g. the vocab classifier). Each // simdgroup owns a 16x32 quadrant (2 token x 4 wrow tiles). @@ -5900,10 +5974,19 @@ def metal_decode_init : bool { g_pso_mv_b4 = compile_pso(metal_q8_mv_b4_msl, metal_q8_mv_b4_msl_entry, metal_q8_mv_b4_msl_fastmath, ok) g_pso_w13sw_b2 = compile_pso(metal_gemv_w13sw_b2_msl, metal_gemv_w13sw_b2_msl_entry, metal_gemv_w13sw_b2_msl_fastmath, ok) g_pso_w13sw_b4 = compile_pso(metal_gemv_w13sw_b4_msl, metal_gemv_w13sw_b4_msl_entry, metal_gemv_w13sw_b4_msl_fastmath, ok) - g_pso_gemm_b = compile_pso(metal_q8_gemm_b_msl, metal_q8_gemm_b_msl_entry, metal_q8_gemm_b_msl_fastmath, ok) - g_pso_gemm_b_sk = compile_pso(metal_q8_gemm_b_sk_msl, metal_q8_gemm_b_sk_msl_entry, metal_q8_gemm_b_sk_msl_fastmath, ok) + g_gemm_b_tensor = metal_tensor_crowned("gemmb_q8") + g_pso_gemm_b = (g_gemm_b_tensor + ? compile_pso(metal_q8_gemm_b_t_msl, metal_q8_gemm_b_t_msl_entry, metal_q8_gemm_b_t_msl_fastmath, ok) + : compile_pso(metal_q8_gemm_b_msl, metal_q8_gemm_b_msl_entry, metal_q8_gemm_b_msl_fastmath, ok)) + g_gemm_b_sk_tensor = metal_tensor_crowned("gemmb_sk_q8") + g_pso_gemm_b_sk = (g_gemm_b_sk_tensor + ? compile_pso(metal_q8_gemm_b_sk_t_msl, metal_q8_gemm_b_sk_t_msl_entry, metal_q8_gemm_b_sk_t_msl_fastmath, ok) + : compile_pso(metal_q8_gemm_b_sk_msl, metal_q8_gemm_b_sk_msl_entry, metal_q8_gemm_b_sk_msl_fastmath, ok)) g_pso_skred = compile_pso(metal_sk_reduce_msl, metal_sk_reduce_msl_entry, metal_sk_reduce_msl_fastmath, ok) - g_pso_gemm64_b = compile_pso(metal_q8_gemm64_b_msl, metal_q8_gemm64_b_msl_entry, metal_q8_gemm64_b_msl_fastmath, ok) + g_gemm64_b_tensor = metal_tensor_crowned("gemm64b_q8") + g_pso_gemm64_b = (g_gemm64_b_tensor + ? compile_pso(metal_q8_gemm64_b_t_msl, metal_q8_gemm64_b_t_msl_entry, metal_q8_gemm64_b_t_msl_fastmath, ok) + : compile_pso(metal_q8_gemm64_b_msl, metal_q8_gemm64_b_msl_entry, metal_q8_gemm64_b_msl_fastmath, ok)) g_pso_mm = compile_pso(metal_q8_mulmm_msl, metal_q8_mulmm_msl_entry, metal_q8_mulmm_msl_fastmath, ok) g_pso_attnb16 = compile_pso(metal_sq_attn_b16_msl, metal_sq_attn_b16_msl_entry, metal_sq_attn_b16_msl_fastmath, ok) g_pso_attnb32 = compile_pso(metal_sq_attn_b32_msl, metal_sq_attn_b32_msl_entry, metal_sq_attn_b32_msl_fastmath, ok) @@ -6441,7 +6524,9 @@ def enc_gemm_mm_b(enc : MetalComputeEncoder?; bwblob : MetalBuffer?; wboff : uin def enc_gemm_b(enc : MetalComputeEncoder?; bw : MetalBuffer?; boff : uint64; bx, by, bk, bn, bys : MetalBuffer?; mp32, d : int64; yoff : uint64 = 0ul) { kn_pipeline(enc, g_pso_gemm_b) - kn_tgmem(enc, metal_q8_gemm_b_msl_tgmem, 0) + if (!g_gemm_b_tensor) { + kn_tgmem(enc, metal_q8_gemm_b_msl_tgmem, 0) + } kn_buffer(enc, bw, boff, 0) // byte view kn_buffer(enc, bw, boff, 2) // half-scale view kn_buffer(enc, bx, 0ul, 3) @@ -6461,7 +6546,9 @@ def enc_gemm_b(enc : MetalComputeEncoder?; bw : MetalBuffer?; boff : uint64; def enc_gemm_sk_b(enc : MetalComputeEncoder?; bw : MetalBuffer?; boff : uint64; bx, bpart, bk, bn, bsk : MetalBuffer?; d, nsk : int64) { kn_pipeline(enc, g_pso_gemm_b_sk) - kn_tgmem(enc, metal_q8_gemm_b_sk_msl_tgmem, 0) + if (!g_gemm_b_sk_tensor) { + kn_tgmem(enc, metal_q8_gemm_b_sk_msl_tgmem, 0) + } kn_buffer(enc, bw, boff, 0) // byte view kn_buffer(enc, bw, boff, 2) // half-scale view kn_buffer(enc, bx, 0ul, 3) @@ -6491,7 +6578,9 @@ def enc_sk_reduce_b(enc : MetalComputeEncoder?; bpart, by, btot, bsk : MetalBuff def enc_gemm64_b(enc : MetalComputeEncoder?; bw : MetalBuffer?; boff : uint64; bx, by, bk, bn : MetalBuffer?; mp32, d : int64) { kn_pipeline(enc, g_pso_gemm64_b) - kn_tgmem(enc, metal_q8_gemm64_b_msl_tgmem, 0) + if (!g_gemm64_b_tensor) { + kn_tgmem(enc, metal_q8_gemm64_b_msl_tgmem, 0) + } kn_buffer(enc, bw, boff, 0) // byte view kn_buffer(enc, bw, boff, 2) // half-scale view kn_buffer(enc, bx, 0ul, 3) @@ -7070,3 +7159,124 @@ def metal_kernels_release { release_pso(g_pso_rpstb_tq4) release_pso(g_pso_addrmsb) } + + +// ===== Metal-4 tensor twin races (decode-side batch GEMM families) ===== + +// one q8b batch-GEMM race: shared blob/X geometry, per-family psos + binds via the blocks +def private race_gemmb_family(dev : MetalDevice?; queue : MetalCommandQueue?; family : string; + m, kdim, ndim, out_n : int; grid : uint3; + base_pso, twin_pso : MetalComputePipeline?; base_tgmem : uint64; + bind_extra : block<(enc : MetalComputeEncoder?; slot_y : MetalBuffer?) : void>) : MetalTensorRaceResult { + var res = MetalTensorRaceResult(family = family, winner = "", base_ms = -1.0lf, twin_ms = -1.0lf, note = "") + let nkb = kdim / 32 + var xa <- race_x_f32(m * kdim) + var blob <- race_q8_blob(ndim * nkb) + var bw = race_buf(dev, uint64(ndim * nkb * 34), unsafe(addr(blob[0]))) + var bxa = race_buf(dev, uint64(m * kdim * 4), unsafe(addr(xa[0]))) + var by_base = race_buf(dev, uint64(out_n * 4), null) + var by_twin = race_buf(dev, uint64(out_n * 4), null) + let tg = uint3(128u, 1u, 1u) + res.base_ms = race_time_ms(queue, 5) $(enc) { + kn_pipeline(enc, base_pso) + kn_tgmem(enc, base_tgmem, 0) + kn_buffer(enc, bw, 0ul, 0) + kn_buffer(enc, bw, 0ul, 2) + kn_buffer(enc, bxa, 0ul, 3) + invoke(bind_extra, enc, by_base) + kn_dispatch(enc, grid, tg) + } + res.twin_ms = race_time_ms(queue, 5) $(enc) { + kn_pipeline(enc, twin_pso) + kn_buffer(enc, bw, 0ul, 0) + kn_buffer(enc, bw, 0ul, 2) + kn_buffer(enc, bxa, 0ul, 3) + invoke(bind_extra, enc, by_twin) + kn_dispatch(enc, grid, tg) + } + if (res.base_ms < 0.0lf || res.twin_ms < 0.0lf) { + res.winner = res.base_ms < 0.0lf ? "" : "simdgroup" + res.note = "dispatch failed" + } elif (!race_envelope_ok(by_base, by_twin, out_n, res.note)) { + res.winner = "simdgroup" + } else { + res.winner = res.twin_ms < res.base_ms ? "tensor" : "simdgroup" + } + metal_release(bw) + metal_release(bxa) + metal_release(by_base) + metal_release(by_twin) + return res +} + +def private race_decode_family(dev : MetalDevice?; queue : MetalCommandQueue?; family : string; + base_src, base_entry : string; base_fm : bool; base_tgmem : uint64; + twin_src, twin_entry : string; twin_fm : bool; + m, kdim, ndim, out_n : int; grid : uint3; + bind_extra : block<(enc : MetalComputeEncoder?; slot_y : MetalBuffer?) : void>) : MetalTensorRaceResult { + var err = "" + var base_pso = pipeline_from_source(dev, base_src, base_entry, base_fm, err) + if (base_pso == null) { + return MetalTensorRaceResult(family = family, winner = "", base_ms = -1.0lf, twin_ms = -1.0lf, note = "base pso: {err}") + } + var twin_pso = pipeline_from_source(dev, twin_src, twin_entry, twin_fm, err) + if (twin_pso == null) { + metal_release(base_pso) + return MetalTensorRaceResult(family = family, winner = "simdgroup", base_ms = -1.0lf, twin_ms = -1.0lf, note = "twin pso: {err}") + } + let res = race_gemmb_family(dev, queue, family, m, kdim, ndim, out_n, grid, base_pso, twin_pso, base_tgmem) $(enc, slot_y) { + invoke(bind_extra, enc, slot_y) + } + metal_release(base_pso) + metal_release(twin_pso) + return res +} + +//! Decode-side tuner entry (the runtime.metal_tensor crown): the batch q8 GEMM families. +def metal_tensor_race_decode : array { + var results : array + var dev = metal_create_system_default_device() + if (dev == null) { + return <- results + } + var queue = metal_new_command_queue(dev) + let kdim = 2048 + let ndim = 1024 + var bk = race_uniform_u32(dev, uint(kdim)) + var bn = race_uniform_u32(dev, uint(ndim)) + var bys = race_uniform_u32(dev, uint(ndim)) + var bsk = race_uniform_u32(dev, 4u) + results |> emplace(race_decode_family(dev, queue, "gemmb_q8", + metal_q8_gemm_b_msl, metal_q8_gemm_b_msl_entry, metal_q8_gemm_b_msl_fastmath, metal_q8_gemm_b_msl_tgmem, + metal_q8_gemm_b_t_msl, metal_q8_gemm_b_t_msl_entry, metal_q8_gemm_b_t_msl_fastmath, + 32, kdim, ndim, 32 * ndim, uint3(uint(ndim / 32), 1u, 1u)) $(enc, slot_y) { + kn_buffer(enc, slot_y, 0ul, 4) + kn_buffer(enc, bk, 0ul, 5) + kn_buffer(enc, bn, 0ul, 6) + kn_buffer(enc, bys, 0ul, 7) + }) + results |> emplace(race_decode_family(dev, queue, "gemmb_sk_q8", + metal_q8_gemm_b_sk_msl, metal_q8_gemm_b_sk_msl_entry, metal_q8_gemm_b_sk_msl_fastmath, metal_q8_gemm_b_sk_msl_tgmem, + metal_q8_gemm_b_sk_t_msl, metal_q8_gemm_b_sk_t_msl_entry, metal_q8_gemm_b_sk_t_msl_fastmath, + 32, kdim, ndim, 4 * 32 * ndim, uint3(uint(ndim / 32), 4u, 1u)) $(enc, slot_y) { + kn_buffer(enc, slot_y, 0ul, 4) + kn_buffer(enc, bk, 0ul, 5) + kn_buffer(enc, bn, 0ul, 6) + kn_buffer(enc, bsk, 0ul, 8) + }) + results |> emplace(race_decode_family(dev, queue, "gemm64b_q8", + metal_q8_gemm64_b_msl, metal_q8_gemm64_b_msl_entry, metal_q8_gemm64_b_msl_fastmath, metal_q8_gemm64_b_msl_tgmem, + metal_q8_gemm64_b_t_msl, metal_q8_gemm64_b_t_msl_entry, metal_q8_gemm64_b_t_msl_fastmath, + 64, kdim, ndim, 64 * ndim, uint3(uint(2 * (ndim / 64)), 1u, 1u)) $(enc, slot_y) { + kn_buffer(enc, slot_y, 0ul, 4) + kn_buffer(enc, bk, 0ul, 5) + kn_buffer(enc, bn, 0ul, 6) + }) + metal_release(bk) + metal_release(bn) + metal_release(bys) + metal_release(bsk) + metal_release(queue) + metal_release(dev) + return <- results +} diff --git a/modules/dasLLAMA/dasllama/dasllama_metal_prefill.das b/modules/dasLLAMA/dasllama/dasllama_metal_prefill.das index a1ee6ea923..0e932a6730 100644 --- a/modules/dasLLAMA/dasllama/dasllama_metal_prefill.das +++ b/modules/dasLLAMA/dasllama/dasllama_metal_prefill.das @@ -273,7 +273,7 @@ class MetalQ8MulMmT { var wp = unsafe(addr(wqb[blk0 * 34u + 2u])) var ap = unsafe(addr(xf[mBase * kdim])) var cp = unsafe(addr(y[mBase * ndim + nBase])) - tmm2d_q8b_f32(32u, 64u, 4u, sp, wp, ap, kdim, cp, ndim, kdim) + tmm2d_q8b_f32(32u, 64u, 4u, sp, wp, ap, kdim, cp, ndim, kdim, kdim / 32u) } } @@ -1775,6 +1775,39 @@ class MetalMoeReduce { } } +// Metal-4 tensor twin of MetalMoeMulMmQ8 — the CONTIGUOUS (gather=0, down-site) shape only: +// bkt-indirect X rows cannot form a strided tensor view, so gate/up keep the simdgroup kernel +// even when "moe_mulmm_q8" is crowned (the encode picks per dispatch). +class MetalMoeMulMmQ8T { + @ssbo @binding = 0 @role = "weight" wsh : array // expert STACK blob, half-scale view + @ssbo @binding = 1 @role = "weight" wqb : array // the same blob buffer, byte view + @ssbo @binding = 2 @role = "read" xf : array + @ssbo @binding = 3 @role = "write" y : array + @uniform @binding = 4 kdim : uint + @uniform @binding = 5 ndim : uint + @ssbo @binding = 6 @role = "read" cnt : array + @ssbo @binding = 7 @role = "read" basep : array + @uniform @binding = 9 eblk : uint + + [metal_kernel(name="metal_moe_mulmm_q8_t_msl")] + def metal_moe_mulmm_q8_t { + let e = gl_WorkGroupID.z + let ce = cnt[e] + let mBase = gl_WorkGroupID.x * 32u + if (mBase >= (ce + 31u) / 32u * 32u) { // threadgroup-uniform exit — cooperative-safe + return + } + let rbase = basep[e] + let nBase = gl_WorkGroupID.y * 64u + let blk0 = e * eblk + nBase * (kdim / 32u) + var sp = unsafe(addr(wsh[blk0 * 17u])) + var wp = unsafe(addr(wqb[blk0 * 34u + 2u])) + var ap = unsafe(addr(xf[(rbase + mBase) * kdim])) + var cp = unsafe(addr(y[(rbase + mBase) * ndim + nBase])) + tmm2d_q8b_f32(32u, 64u, 4u, sp, wp, ap, kdim, cp, ndim, kdim, kdim / 32u) + } +} + // Gathered q8 mul_mm — MetalQ8MulMm per expert (grid.z): weight base e*eblk, X rows gathered // via bkt (gather=1: src row = entry/k; gather=0: contiguous bucket rows — the down site), // C tiles land at the expert's padded bucket rows; tiles past the padded count exit. @@ -4527,6 +4560,8 @@ var private g_pf_pso_mm : MetalComputePipeline? // the production das mul_mm var private g_pf_pso_bf16_mm : MetalComputePipeline? // native-BF16 A twin (E-series model_proj) var private g_pf_bf16_mm_tensor : bool // crowned tensor twin selected (no tgmem bind) var private g_pf_mm_tensor : bool // ditto for the production q8 mul_mm +var private g_pf_moe_mm_q8_tensor : bool // "moe_mulmm_q8" crowned (contiguous sites only) +var private g_pf_pso_moe_mm_q8_t : MetalComputePipeline? var private g_pf_pso_ple_gather : MetalComputePipeline? // PLE pre-step: q8 token-row gather var private g_pf_pso_ple_finish : MetalComputePipeline? // PLE pre-step: rms + combine, in place var private g_pso_qkmm : MetalComputePipeline? @@ -4653,6 +4688,10 @@ def public metal_prefill_shutdown { } pool_drain(g_pf_pool) pool_drain(g_pf_upool) + if (g_pf_pso_moe_mm_q8_t != null) { + metal_release(g_pf_pso_moe_mm_q8_t) + g_pf_pso_moe_mm_q8_t = null + } if (g_pso_gemm != null) { metal_release(g_pso_gemm) g_pso_gemm = null @@ -4913,6 +4952,10 @@ def private metal_prefill_init : bool { g_pf_pso_moe_bucket = pf_compile_pso(metal_moe_bucket_msl, metal_moe_bucket_msl_entry, metal_moe_bucket_msl_fastmath, ok) g_pf_pso_moe_reduce = pf_compile_pso(metal_moe_reduce_msl, metal_moe_reduce_msl_entry, metal_moe_reduce_msl_fastmath, ok) g_pf_pso_moe_mm_q8 = pf_compile_pso(metal_moe_mulmm_q8_msl, metal_moe_mulmm_q8_msl_entry, metal_moe_mulmm_q8_msl_fastmath, ok) + g_pf_moe_mm_q8_tensor = metal_tensor_crowned("moe_mulmm_q8") + if (g_pf_moe_mm_q8_tensor) { + g_pf_pso_moe_mm_q8_t = pf_compile_pso(metal_moe_mulmm_q8_t_msl, metal_moe_mulmm_q8_t_msl_entry, metal_moe_mulmm_q8_t_msl_fastmath, ok) + } g_pf_pso_moe_mm_k4 = pf_compile_pso(metal_moe_mulmm_k4_msl, metal_moe_mulmm_k4_msl_entry, metal_moe_mulmm_k4_msl_fastmath, ok) g_pf_pso_moe_mm_k5 = pf_compile_pso(metal_moe_mulmm_k5_msl, metal_moe_mulmm_k5_msl_entry, metal_moe_mulmm_k5_msl_fastmath, ok) g_pf_pso_moe_mm_k6 = pf_compile_pso(metal_moe_mulmm_k6_msl, metal_moe_mulmm_k6_msl_entry, metal_moe_mulmm_k6_msl_fastmath, ok) @@ -5251,13 +5294,17 @@ def private pf_enc_moe_route(enc : MetalComputeEncoder?; t : Model; l : int64; // one gathered expert mul_mm site: grid (ceil(npos/32) tiles, rows/64, ne experts); the twin // exits tiles past the expert's padded count. bgather = 1 (bkt-indirect X) / 0 (bucket rows). def private pf_enc_moe_mm(enc : MetalComputeEncoder?; t : Model; fmt : KqFmt; stack_off, rows, npos : int64; - bx, by, bkdim, bndim, bcnt, bbase, bbkt, bes, bnk, bgather : MetalBuffer?) { + bx, by, bkdim, bndim, bcnt, bbase, bbkt, bes, bnk, bgather : MetalBuffer?; + contiguous : bool = false) { let tiles = (npos + 31l) / 32l let ne = t.config.n_expert if (fmt == KqFmt.q8) { let bw = pf_blob_of(t, stack_off) - kn_pipeline(enc, g_pf_pso_moe_mm_q8) - kn_tgmem(enc, metal_moe_mulmm_q8_msl_tgmem, 0) + let twin = g_pf_moe_mm_q8_tensor && contiguous && g_pf_pso_moe_mm_q8_t != null + kn_pipeline(enc, twin ? g_pf_pso_moe_mm_q8_t : g_pf_pso_moe_mm_q8) + if (!twin) { + kn_tgmem(enc, metal_moe_mulmm_q8_msl_tgmem, 0) + } kn_buffer(enc, bw.buf, bw.boff, 0) kn_buffer(enc, bw.buf, bw.boff, 1) kn_buffer(enc, bx, 0ul, 2) @@ -6605,7 +6652,7 @@ def metal_prefill_forward(t : Model; var s : Session; npos, start_pos : int64) : u_moe_nfe, u_dim, bmcnt, bmbase, bmbkt, u_moe_eblk, u_moe_nk, u_moe_g0, bwb2, bhasb) } else { pf_enc_moe_mm(enc, t, fe2, t.we2_offs[l], dim, npos, bmg, bmdn, - u_moe_nfe, u_dim, bmcnt, bmbase, bmbkt, fe2 == KqFmt.q8 || fe2 == KqFmt.q51 ? u_moe_eblk : u_moe_esb, u_moe_nk, u_moe_g0) + u_moe_nfe, u_dim, bmcnt, bmbase, bmbkt, fe2 == KqFmt.q8 || fe2 == KqFmt.q51 ? u_moe_eblk : u_moe_esb, u_moe_nk, u_moe_g0, contiguous = true) } if (nsh > 0l) { // shexp down into bxb2 (free post-residual) — bxb still feeds the reduce @@ -7055,49 +7102,6 @@ def dasllama_metal_prefill_register() { // ===== Metal-4 tensor twin race (the tuner's crowning section) ===== -struct public MetalTensorRaceResult { - family : string - winner : string // "tensor" | "simdgroup" | "" when the family could not race - base_ms : double - twin_ms : double - note : string -} - -def private race_buf(dev : MetalDevice?; bytes : uint64; src : void?) : MetalBuffer? { - var b = metal_new_buffer(dev, bytes) - if (src != null) { - unsafe { - memcpy(metal_buffer_contents(b), src, bytes) - } - } - return b -} - -def private race_uniform_u32(dev : MetalDevice?; v : uint) : MetalBuffer? { - var b = metal_new_buffer(dev, 4ul) - unsafe { - var p = reinterpret(metal_buffer_contents(b)) - p[0] = v - } - return b -} - -// best GPU time over `reps` dispatches after one warmup; <0 on dispatch failure -def private race_time_ms(queue : MetalCommandQueue?; reps : int; blk : block<(enc : MetalComputeEncoder?) : void>) : double { - var best = -1.0lf - var err = "" - for (r in range(reps + 1)) { - var gpu_ms = 0.0lf - if (!with_compute_encoder_timed(queue, err, gpu_ms) $(enc) { invoke(blk, enc) }) { - return -1.0lf - } - if (r > 0 && (best < 0.0lf || gpu_ms < best)) { - best = gpu_ms - } - } - return best -} - def private race_mulmm_bf16(dev : MetalDevice?; queue : MetalCommandQueue?) : MetalTensorRaceResult { var res = MetalTensorRaceResult(family = "mulmm_bf16", winner = "", base_ms = -1.0lf, twin_ms = -1.0lf, note = "") let m = 512 @@ -7117,13 +7121,9 @@ def private race_mulmm_bf16(dev : MetalDevice?; queue : MetalCommandQueue?) : Me metal_release(base_pso) return res } - var xa : array + var xa <- race_x_f32(m * kdim) var wb : array - xa |> resize(m * kdim) wb |> resize(ndim * kdim) - for (i in range(m * kdim)) { - xa[i] = 0.25 * float(i % 17 - 8) - } for (i in range(ndim * kdim)) { wb[i] = uint16(float_bits_to_uint(0.125 * float((i * 3) % 13 - 6)) >> 16u) } @@ -7157,24 +7157,10 @@ def private race_mulmm_bf16(dev : MetalDevice?; queue : MetalCommandQueue?) : Me if (res.base_ms < 0.0lf || res.twin_ms < 0.0lf) { res.winner = res.base_ms < 0.0lf ? "" : "simdgroup" res.note = "dispatch failed" + } elif (!race_envelope_ok(by_base, by_twin, m * ndim, res.note)) { + res.winner = "simdgroup" } else { - // eligibility envelope: f16-staging noise passes, a garbage twin lands at O(max) - var maxref = 0.0lf - var maxdiff = 0.0lf - unsafe { - let pb = reinterpret(metal_buffer_contents(by_base)) - let pt = reinterpret(metal_buffer_contents(by_twin)) - for (i in range(m * ndim)) { - maxref = max(maxref, abs(double(pb[i]))) - maxdiff = max(maxdiff, abs(double(pt[i]) - double(pb[i]))) - } - } - if (maxdiff > 0.05lf * maxref + 1e-6lf) { - res.winner = "simdgroup" - res.note = "twin OUTPUT MISMATCH maxdiff={maxdiff} maxref={maxref}" - } else { - res.winner = res.twin_ms < res.base_ms ? "tensor" : "simdgroup" - } + res.winner = res.twin_ms < res.base_ms ? "tensor" : "simdgroup" } metal_release(bxa) metal_release(bwb) @@ -7206,25 +7192,8 @@ def private race_mulmm_q8(dev : MetalDevice?; queue : MetalCommandQueue?) : Meta metal_release(base_pso) return res } - var xa : array - var blob : array - xa |> resize(m * kdim) - blob |> resize(ndim * nkb * 34) - for (i in range(m * kdim)) { - xa[i] = 0.25 * float(i % 17 - 8) - } - for (blk in range(ndim * nkb)) { - for (k in range(32)) { - blob[blk * 34 + 2 + k] = int8((blk * 7 + k * 3) % 19 - 9) - } - } - unsafe { - var p8 = addr(blob[0]) - for (blk in range(ndim * nkb)) { - var ph = reinterpret(p8 + blk * 34) - ph[0] = float16(0.5 * float(1 << (blk % 3))) - } - } + var xa <- race_x_f32(m * kdim) + var blob <- race_q8_blob(ndim * nkb) var bw = race_buf(dev, uint64(ndim * nkb * 34), unsafe(addr(blob[0]))) var bxa = race_buf(dev, uint64(m * kdim * 4), unsafe(addr(xa[0]))) var by_base = race_buf(dev, uint64(m * ndim * 4), null) @@ -7257,31 +7226,123 @@ def private race_mulmm_q8(dev : MetalDevice?; queue : MetalCommandQueue?) : Meta if (res.base_ms < 0.0lf || res.twin_ms < 0.0lf) { res.winner = res.base_ms < 0.0lf ? "" : "simdgroup" res.note = "dispatch failed" + } elif (!race_envelope_ok(by_base, by_twin, m * ndim, res.note)) { + res.winner = "simdgroup" } else { - // eligibility envelope: f16-staging noise passes, a garbage twin lands at O(max) - var maxref = 0.0lf - var maxdiff = 0.0lf - unsafe { - let pb = reinterpret(metal_buffer_contents(by_base)) - let pt = reinterpret(metal_buffer_contents(by_twin)) - for (i in range(m * ndim)) { - maxref = max(maxref, abs(double(pb[i]))) - maxdiff = max(maxdiff, abs(double(pt[i]) - double(pb[i]))) - } - } - if (maxdiff > 0.05lf * maxref + 1e-6lf) { - res.winner = "simdgroup" - res.note = "twin OUTPUT MISMATCH maxdiff={maxdiff} maxref={maxref}" - } else { - res.winner = res.twin_ms < res.base_ms ? "tensor" : "simdgroup" - } + res.winner = res.twin_ms < res.base_ms ? "tensor" : "simdgroup" + } + metal_release(bw) + metal_release(bxa) + metal_release(by_base) + metal_release(by_twin) + metal_release(bk) + metal_release(bn) + metal_release(base_pso) + metal_release(twin_pso) + return res +} + +def private race_moe_mulmm_q8(dev : MetalDevice?; queue : MetalCommandQueue?) : MetalTensorRaceResult { + var res = MetalTensorRaceResult(family = "moe_mulmm_q8", winner = "", base_ms = -1.0lf, twin_ms = -1.0lf, note = "") + let ne = 4 + let rows_per = 32 + let m = ne * rows_per + let kdim = 2048 + let ndim = 1024 + let nkb = kdim / 32 + let eblk = ndim * nkb + var err = "" + var base_pso = pipeline_from_source(dev, metal_moe_mulmm_q8_msl, metal_moe_mulmm_q8_msl_entry, metal_moe_mulmm_q8_msl_fastmath, err) + if (base_pso == null) { + res.note = "base pso: {err}" + return res + } + var twin_pso = pipeline_from_source(dev, metal_moe_mulmm_q8_t_msl, metal_moe_mulmm_q8_t_msl_entry, metal_moe_mulmm_q8_t_msl_fastmath, err) + if (twin_pso == null) { + res.winner = "simdgroup" + res.note = "twin pso: {err}" + metal_release(base_pso) + return res + } + var xa <- race_x_f32(m * kdim) + var blob <- race_q8_blob(ne * eblk) + var cnt : array + var basep : array + var bkt : array + cnt |> resize(ne) + basep |> resize(ne) + bkt |> resize(m) + for (e in range(ne)) { + cnt[e] = uint(rows_per) + basep[e] = uint(e * rows_per) + } + for (i in range(m)) { + bkt[i] = uint(i) // identity buckets: gather=0 semantics match + } + var bw = race_buf(dev, uint64(ne * eblk * 34), unsafe(addr(blob[0]))) + var bxa = race_buf(dev, uint64(m * kdim * 4), unsafe(addr(xa[0]))) + var by_base = race_buf(dev, uint64(m * ndim * 4), null) + var by_twin = race_buf(dev, uint64(m * ndim * 4), null) + var bcnt = race_buf(dev, uint64(ne * 4), unsafe(addr(cnt[0]))) + var bbase = race_buf(dev, uint64(ne * 4), unsafe(addr(basep[0]))) + var bbkt = race_buf(dev, uint64(m * 4), unsafe(addr(bkt[0]))) + var bk = race_uniform_u32(dev, uint(kdim)) + var bn = race_uniform_u32(dev, uint(ndim)) + var beblk = race_uniform_u32(dev, uint(eblk)) + var bnk = race_uniform_u32(dev, 1u) + var bg0 = race_uniform_u32(dev, 0u) + let grid = uint3(uint(rows_per / 32), uint(ndim / 64), uint(ne)) + let tg = uint3(128u, 1u, 1u) + res.base_ms = race_time_ms(queue, 5) $(enc) { + kn_pipeline(enc, base_pso) + kn_tgmem(enc, metal_moe_mulmm_q8_msl_tgmem, 0) + kn_buffer(enc, bw, 0ul, 0) + kn_buffer(enc, bw, 0ul, 1) + kn_buffer(enc, bxa, 0ul, 2) + kn_buffer(enc, by_base, 0ul, 3) + kn_buffer(enc, bk, 0ul, 4) + kn_buffer(enc, bn, 0ul, 5) + kn_buffer(enc, bcnt, 0ul, 6) + kn_buffer(enc, bbase, 0ul, 7) + kn_buffer(enc, bbkt, 0ul, 8) + kn_buffer(enc, beblk, 0ul, 9) + kn_buffer(enc, bnk, 0ul, 10) + kn_buffer(enc, bg0, 0ul, 11) + kn_dispatch(enc, grid, tg) + } + res.twin_ms = race_time_ms(queue, 5) $(enc) { + kn_pipeline(enc, twin_pso) + kn_buffer(enc, bw, 0ul, 0) + kn_buffer(enc, bw, 0ul, 1) + kn_buffer(enc, bxa, 0ul, 2) + kn_buffer(enc, by_twin, 0ul, 3) + kn_buffer(enc, bk, 0ul, 4) + kn_buffer(enc, bn, 0ul, 5) + kn_buffer(enc, bcnt, 0ul, 6) + kn_buffer(enc, bbase, 0ul, 7) + kn_buffer(enc, beblk, 0ul, 9) + kn_dispatch(enc, grid, tg) + } + if (res.base_ms < 0.0lf || res.twin_ms < 0.0lf) { + res.winner = res.base_ms < 0.0lf ? "" : "simdgroup" + res.note = "dispatch failed" + } elif (!race_envelope_ok(by_base, by_twin, m * ndim, res.note)) { + res.winner = "simdgroup" + } else { + res.winner = res.twin_ms < res.base_ms ? "tensor" : "simdgroup" } metal_release(bw) metal_release(bxa) metal_release(by_base) metal_release(by_twin) + metal_release(bcnt) + metal_release(bbase) + metal_release(bbkt) metal_release(bk) metal_release(bn) + metal_release(beblk) + metal_release(bnk) + metal_release(bg0) metal_release(base_pso) metal_release(twin_pso) return res @@ -7299,6 +7360,7 @@ def metal_tensor_race : array { var queue = metal_new_command_queue(dev) results |> emplace(race_mulmm_bf16(dev, queue)) results |> emplace(race_mulmm_q8(dev, queue)) + results |> emplace(race_moe_mulmm_q8(dev, queue)) metal_release(queue) metal_release(dev) return <- results diff --git a/modules/dasLLAMA/harness/tune_kernels.das b/modules/dasLLAMA/harness/tune_kernels.das index b4cca52b22..80874bd527 100644 --- a/modules/dasLLAMA/harness/tune_kernels.das +++ b/modules/dasLLAMA/harness/tune_kernels.das @@ -22,6 +22,7 @@ require dasllama/dasllama_common // the runtime-knob getters (the profile's "run require dasllama/dasllama_tune // [dasllama_grid] emitter require llvm/daslib/llvm_tune // tune_sidecar_merge — the app-sidecar upsert writer require ?das_metal dasllama/dasllama_metal_prefill // nolint:STYLE030 — metal_tensor_race (Apple static_if half) +require ?das_metal dasllama/dasllama_metal_kernels // nolint:STYLE030 — metal_tensor_race_decode (Apple static_if half) require tuner_cli // --tune-fast, shared with the other half require llvm/daslib/aarch64_neon // nolint:STYLE030 — sdot4_laneq (referenced by the cloned laneq template bodies) require llvm/daslib/f16_cvt // f16 references + the cloned f16-kernel template bodies @@ -2053,7 +2054,14 @@ def main { static_if (typeinfo builtin_module_exists(das_metal)) { var mt_wins : array var mt_ran = false - for (r in metal_tensor_race()) { + var mt_all <- metal_tensor_race() + var mt_dec <- metal_tensor_race_decode() + mt_all |> reserve(length(mt_all) + length(mt_dec)) + for (r in mt_dec) { + mt_all |> emplace(r) + } + delete mt_dec + for (r in mt_all) { mt_ran = true let mt_note = empty(r.note) ? "" : " ({r.note})" tune_detail("METAL_TWIN {r.family}: base={r.base_ms}ms tensor={r.twin_ms}ms -> {r.winner}{mt_note}\n") @@ -2066,6 +2074,7 @@ def main { } metal_crowns = join(mt_wins, ",") delete mt_wins + delete mt_all } // softmax_sink shares softmax's loop shape exactly (one extra scalar joins the reduction), so it diff --git a/modules/dasMetal/metal/metal_builtins.das b/modules/dasMetal/metal/metal_builtins.das index c9a0b21722..0e3b832b1e 100644 --- a/modules/dasMetal/metal/metal_builtins.das +++ b/modules/dasMetal/metal/metal_builtins.das @@ -351,11 +351,11 @@ def public tmm2d_q8_f16s(m, n, sgs : uint; aq : int8 const?; lda : uint; asf : f } //! Interleaved-q8_0 W blocks x RAW f32 activations (the production MulMm dataflow). `wsh` is the -//! blob's f16-scale view at the tile's first block (row j chunk b scale at (j*nb+b)*17); `wqb` -//! the byte view pre-offset to that block's quants (+2 done by the caller — row j chunk b quant k -//! at (j*nb+b)*34 + k). f32 chunk accumulate, W scale folded between chunks. +//! blob's f16-scale view at the tile's first block (row j chunk b scale at (j*ldwb+b)*17); `wqb` +//! the byte view pre-offset to that block's quants (+2 done by the caller). `ldwb` = W row stride +//! in BLOCKS (kk/32 for a full row; larger under split-K). W scale folded between f32 chunks. [unused_argument(sgs), sideeffects] -def public tmm2d_q8b_f32(m, n, sgs : uint; wsh : float16 const?; wqb : int8 const?; a : float const?; lda : uint; var c : float?; ldc : uint; kk : uint) { +def public tmm2d_q8b_f32(m, n, sgs : uint; wsh : float16 const?; wqb : int8 const?; a : float const?; lda : uint; var c : float?; ldc : uint; kk : uint; ldwb : uint) { if (gl_LocalInvocationID.x != 0u) { return } @@ -365,7 +365,7 @@ def public tmm2d_q8b_f32(m, n, sgs : uint; wsh : float16 const?; wqb : int8 cons for (j in range(int(n))) { var facc = 0.0 for (b in range(nb)) { - let blk = j * nb + b + let blk = j * int(ldwb) + b var ch = 0.0 for (k in range(32)) { ch += a[r * int(lda) + b * 32 + k] * float(wqb[blk * 34 + k]) diff --git a/modules/dasMetal/metal/msl_emit.das b/modules/dasMetal/metal/msl_emit.das index 3615bdf017..8b869fcd3b 100644 --- a/modules/dasMetal/metal/msl_emit.das +++ b/modules/dasMetal/metal/msl_emit.das @@ -346,7 +346,8 @@ def private scan_body(var ctx : MslCtx; e : Expression?) { mark_write_target(ctx, ecs.arguments[1]) } // the tensor GEMMs write through their C pointer (arg 7 for the 10-arg forms, 11 for q8xq8) - if ((scn == "tmm2d_f32_bf16_f32" || scn == "tmm2d_q8b_f32") && length(ecs.arguments) == 10) { + if ((scn == "tmm2d_f32_bf16_f32" && length(ecs.arguments) == 10) + || (scn == "tmm2d_q8b_f32" && length(ecs.arguments) == 11)) { mark_write_target(ctx, ecs.arguments[7]) } if ((scn == "tmm2d_q8_f32" || scn == "tmm2d_q8_f16s") && length(ecs.arguments) == 14) { @@ -1203,24 +1204,24 @@ def private tmm2d_helper_text(kind : string; m, n, sgs : uint) : string { // interleaved-q8_0 W blocks (34B stride, quants pre-offset +2) x raw f32 activations: // f32 chunk accumulate, per-(col,chunk) W scale folded between chunks return build_string() $(var w) { - w |> write("static void {name}(device const half * ws, device const char * wq, device const float * a, uint lda, device float * c, uint ldc, uint kk) \{\n") + w |> write("static void {name}(device const half * ws, device const char * wq, device const float * a, uint lda, device float * c, uint ldc, uint kk, uint ldwb) \{\n") w |> write(" constexpr auto d = matmul2d_descriptor({int(m)}, {int(n)}, 32, false, true, false);\n") w |> write(" matmul2d> op;\n") w |> write(" auto C = tensor, tensor_inline>(c, dextents({int(n)}, {int(m)}), metal::array\{1, int(ldc)\});\n") w |> write(" const uint nb = kk / 32u;\n") w |> write(" auto A0 = tensor, tensor_inline>((device float *)a, dextents(32, {int(m)}), metal::array\{1, int(lda)\});\n") - w |> write(" auto W0 = {ct}((device int8_t *)wq, dextents(32, {int(n)}), metal::array\{1, int(nb * 34u)\});\n") + w |> write(" auto W0 = {ct}((device int8_t *)wq, dextents(32, {int(n)}), metal::array\{1, int(ldwb * 34u)\});\n") w |> write(" auto fT = op.get_destination_cooperative_tensor();\n") w |> write(" for (uint16_t i = 0; i < fT.get_capacity(); ++i) \{ if (fT.is_valid_element(i)) fT[i] = 0.0f; \}\n") w |> write(" for (uint b = 0; b < nb; ++b) \{\n") w |> write(" auto A = tensor, tensor_inline>((device float *)(a + b * 32u), dextents(32, {int(m)}), metal::array\{1, int(lda)\});\n") - w |> write(" auto W = {ct}((device int8_t *)(wq + b * 34u), dextents(32, {int(n)}), metal::array\{1, int(nb * 34u)\});\n") + w |> write(" auto W = {ct}((device int8_t *)(wq + b * 34u), dextents(32, {int(n)}), metal::array\{1, int(ldwb * 34u)\});\n") w |> write(" auto cT = op.get_destination_cooperative_tensor();\n") w |> write(" for (uint16_t i = 0; i < cT.get_capacity(); ++i) \{ if (cT.is_valid_element(i)) cT[i] = 0.0f; \}\n") w |> write(" op.run(A, W, cT);\n") w |> write(" for (uint16_t i = 0; i < cT.get_capacity(); ++i) \{ if (cT.is_valid_element(i)) \{\n") w |> write(" auto ids = cT.get_multidimensional_index(i);\n") - w |> write(" fT[i] += cT[i] * float(ws[(uint(ids[0]) * nb + b) * 17u]);\n") + w |> write(" fT[i] += cT[i] * float(ws[(uint(ids[0]) * ldwb + b) * 17u]);\n") w |> write(" \} \}\n") w |> write(" \}\n") w |> write(" fT.store(C);\n") @@ -1259,7 +1260,7 @@ def private tmm2d_helper_text(kind : string; m, n, sgs : uint) : string { def private emit_tmm2d(var ctx : MslCtx; ec : ExprCall?; indent : string; var lines : array) { let cname = call_base_name("{ec.name}") let kind = cname == "tmm2d_f32_bf16_f32" ? "bf" : (cname == "tmm2d_q8b_f32" ? "q8b" : (cname == "tmm2d_q8_f16s" ? "q8h" : "q8")) - let want_args = kind == "bf" || kind == "q8b" ? 10 : 14 + let want_args = kind == "bf" ? 10 : (kind == "q8b" ? 11 : 14) if (length(ec.arguments) != want_args) { err(ctx, ec.at, "{cname} takes exactly {want_args} arguments") return diff --git a/src/ast/ast_dse.cpp b/src/ast/ast_dse.cpp index 2441cc33a2..00c93d1378 100644 --- a/src/ast/ast_dse.cpp +++ b/src/ast/ast_dse.cpp @@ -338,10 +338,17 @@ namespace das { continue; // the statement vanishes whole: no kills, no uses } // impure rhs: the store dies, the rhs stays as a bare statement - - // no kills (the write is gone), uses stay (the rhs still evaluates) - deadKeepRhs.insert(b->stmts[si]); - for ( int u : SI.uses ) live[u>>6] |= 1ull<<(u&63); - continue; + // no kills (the write is gone), uses stay (the rhs still evaluates). + // NOT for an operator-topped rhs: a bare `a + b` statement trips the + // top-level no-side-effect lint when a reused shared module re-lints + // its optimized AST (the operator node itself is a pure builtin even + // when an operand is not) - that store stays whole + auto rhs = static_cast(b->stmts[si])->right; + if ( !rhs->rtti_isOp2() && !rhs->rtti_isOp1() ) { + deadKeepRhs.insert(b->stmts[si]); + for ( int u : SI.uses ) live[u>>6] |= 1ull<<(u&63); + continue; + } } if ( !SI.letInits.empty() ) { // a later init in the same `let` may read an earlier variable - diff --git a/tests/language/optimization_dead_stores.das b/tests/language/optimization_dead_stores.das index 3fd6b0861e..e5de3cab85 100644 --- a/tests/language/optimization_dead_stores.das +++ b/tests/language/optimization_dead_stores.das @@ -97,7 +97,7 @@ def target_impure_reads(x : int) : int { var u = 0 // nolint:LINT010 -- deliberate dead init u = x * 109 // LIVE: the kept bare rhs below still reads u var t = 0 // nolint:LINT010 -- deliberate dead init - t = bump() + u // dead store, impure rhs: the call survives as a bare statement + t = bump() + u // nolint:LINT010 -- dead store, but operator-topped impure rhs: stays WHOLE return x } @@ -315,10 +315,12 @@ def test_dead_stores_fired(t : T?) { t |> success(has(b, "bump"), "side-effecting rhs kept: {b}") t |> success(!has(b, "= bump"), "dead store stripped to a bare call: {b}") } - t |> run("kept bare rhs still counts as a read") @(t : T?) { + t |> run("operator-topped impure rhs keeps its store whole") @(t : T?) { + // stripping to a bare `bump() + u` would trip the top-level no-side-effect lint + // when a reused shared module re-lints its optimized AST — the store must survive let b = body_of(t, "target_impure_reads") t |> success(has(b, "* 109"), "store read by the kept rhs stays: {b}") - t |> success(!has(b, "= bump"), "dead store stripped to a bare statement: {b}") + t |> success(has(b, "= bump"), "operator-topped dead store kept whole: {b}") } t |> run("addr() disqualifies the variable") @(t : T?) { let b = body_of(t, "target_alias") diff --git a/tests/metal/test_metal_tensor_ops.das b/tests/metal/test_metal_tensor_ops.das index cf3b1e52c1..0a8d9bd947 100644 --- a/tests/metal/test_metal_tensor_ops.das +++ b/tests/metal/test_metal_tensor_ops.das @@ -66,7 +66,7 @@ class TmmKernels { var wp = unsafe(addr(bqb[2])) var ap = unsafe(addr(xa[0])) var cp = unsafe(addr(oc[0])) - tmm2d_q8b_f32(8u, 16u, 2u, sp, wp, ap, 64u, cp, 16u, 64u) + tmm2d_q8b_f32(8u, 16u, 2u, sp, wp, ap, 64u, cp, 16u, 64u, 2u) } } From 9e30029de15f9661a4d067cbf49479215403d483 Mon Sep 17 00:00:00 2001 From: Boris Batkin Date: Mon, 27 Jul 2026 22:23:42 -0700 Subject: [PATCH 08/18] dasMetal: staged-tile cooperative GEMM protocol stubs (Phase D step-0) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The tmm2d_tg_begin/step/store trio for dequant-staged formats (mx4 LUT, k-quant sub-scales): stage flat row-major f16 tiles per K-chunk, step folds each chunk into a persistent accumulator, store writes C. CPU stubs are the full replay; the MSL emitter arms land next. Probe-proven ground truth behind the shape: threadgroup-half tensor_inline views work as matmul2d operands, and run() OVERWRITES a cooperative destination — accumulation must be the explicit per-chunk fold (exact in the chunked probe). Also proven: e2m1-doubled does not fit int4b (+/-8, +/-12 overflow), so LUT->f16 staging supersedes the repack-plane route. Co-Authored-By: Claude Fable 5 --- modules/dasMetal/metal/metal_builtins.das | 51 +++++++++++++++++++++++ 1 file changed, 51 insertions(+) diff --git a/modules/dasMetal/metal/metal_builtins.das b/modules/dasMetal/metal/metal_builtins.das index 0e3b832b1e..0785c0bf0f 100644 --- a/modules/dasMetal/metal/metal_builtins.das +++ b/modules/dasMetal/metal/metal_builtins.das @@ -377,3 +377,54 @@ def public tmm2d_q8b_f32(m, n, sgs : uint; wsh : float16 const?; wqb : int8 cons } } } + +// ===== staged-tile cooperative GEMM protocol (begin / step / store) ===== +// For kernels that DEQUANT-stage weight/activation tiles into @workgroup float16 (mx4 LUT, +// k-quant sub-scales): stage a FLAT row-major [rows][kk] tile pair per K-chunk, then step. +// `acc` is an m*n float @workgroup (or local) scratch — the CPU-replay accumulator; on GPU the +// emitter keys a persistent cooperative tensor to the acc VARIABLE and never touches the array. +// run() OVERWRITES a cooperative destination (probe-proven), so step folds explicitly. + +//! begin: zero the accumulator. m/n/sgs/kk must be call-site constants (kk = chunk K width). +[unused_argument(sgs, kk), sideeffects] +def public tmm2d_tg_begin(var acc; m, n, sgs, kk : uint) { + if (gl_LocalInvocationID.x != 0u) { + return + } + for (i in range(int(m * n))) { + acc[i] = 0.0 + } +} + +//! step: acc[r][j] += ta[r][k] * tb[j][k] over one staged chunk — ta flat [m][kk] f16, +//! tb flat [n][kk] f16, both row-major. Call between the same barriers the staging uses. +[sideeffects] +def public tmm2d_tg_step(var acc; ta; tb; m, n, kk : uint) { + if (gl_LocalInvocationID.x != 0u) { + return + } + for (r in range(int(m))) { + for (j in range(int(n))) { + var facc = 0.0 + for (k in range(int(kk))) { + facc += float(ta[r * int(kk) + k]) * float(tb[j * int(kk) + k]) + } + acc[r * int(n) + j] += facc + } + } +} + +//! store: C[r][j] = acc[r][j] (overwrite, row stride ldc). +[sideeffects] +def public tmm2d_tg_store(acc; var c : float?; m, n, ldc : uint) { + if (gl_LocalInvocationID.x != 0u) { + return + } + unsafe { + for (r in range(int(m))) { + for (j in range(int(n))) { + c[r * int(ldc) + j] = acc[r * int(n) + j] + } + } + } +} From 897116ffa0f4572fabbf9618dd3739abd9b24fa2 Mon Sep 17 00:00:00 2001 From: Boris Batkin Date: Mon, 27 Jul 2026 22:27:36 -0700 Subject: [PATCH 09/18] dasMetal: emit the staged-tile GEMM protocol (tmm2d_tg begin/step/store) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Inline MSL emission keyed by the accumulator variable: begin declares the descriptor/op/threadgroup-half view type/persistent cooperative fold target (zeroed), step builds flat [rows][kk] views over the staged @workgroup f16 tiles, runs the chunk, and folds (run() overwrites a cooperative destination), store writes C strided. The acc local is CPU-replay state only — its declaration is suppressed on GPU; the MPP include now also rides tg-protocol use (no generated helper involved). Conformance gains a staged two-chunk kernel: GPU == width-1 CPU replay exact. Access classifier learns the three calls (store writes arg 1). Co-Authored-By: Claude Fable 5 --- .../dasllama/dasllama_kernel_access.das | 24 +++- modules/dasMetal/metal/msl_emit.das | 113 +++++++++++++++++- tests/metal/test_metal_tensor_ops.das | 48 +++++++- 3 files changed, 174 insertions(+), 11 deletions(-) diff --git a/modules/dasLLAMA/dasllama/dasllama_kernel_access.das b/modules/dasLLAMA/dasllama/dasllama_kernel_access.das index a9320dd3d3..343597cb62 100644 --- a/modules/dasLLAMA/dasllama/dasllama_kernel_access.das +++ b/modules/dasLLAMA/dasllama/dasllama_kernel_access.das @@ -197,8 +197,28 @@ class private AccessVisitor : AstVisitor { return } if (cname == "coopmatLoad" || cname == "simdgroup_load" - || cname == "coopmatLoadTensor" || cname == "coopmatLoadTensorDecode") { - return // src argument is a plain read — the generic read pass records it + || cname == "coopmatLoadTensor" || cname == "coopmatLoadTensorDecode" + || cname == "tmm2d_tg_begin" || cname == "tmm2d_tg_step") { + // plain reads (the generic pass records them) — or, for the staged-GEMM protocol + // begin/step, only untracked locals/@workgroup tiles + return + } + if (cname == "tmm2d_tg_store") { + if (length(expr.arguments) == 5) { + var cur = expr.arguments[1] + if (cur is ExprRef2Value) { + cur = (cur as ExprRef2Value).subexpr + } + if (cur is ExprVar) { + let pn = string((cur as ExprVar).name) + let bn = ptrs?[pn] ?? (key_exists(tracked, pn) ? pn : "") + if (!empty(bn)) { + claimed |> insert("{pn}:{cur.at.line}:{cur.at.column}") + writes |> insert(bn) + } + } + } + return } let tci = tmm2d_c_arg(cname) if (tci >= 0) { diff --git a/modules/dasMetal/metal/msl_emit.das b/modules/dasMetal/metal/msl_emit.das index 8b869fcd3b..edb206e6a2 100644 --- a/modules/dasMetal/metal/msl_emit.das +++ b/modules/dasMetal/metal/msl_emit.das @@ -51,6 +51,7 @@ struct private MslCtx { helpers : table // preamble helper functions referenced by lowered calls gen_helpers : table // GENERATED preamble helpers (tensor ops): mangled name -> text scan_ptrs : table // scan-time pointer local -> member (write-set resolution) + tg_accs : table // staged-GEMM acc local -> "m|n|sgs|kk" (decl suppressed on GPU) self_name : string } @@ -353,6 +354,18 @@ def private scan_body(var ctx : MslCtx; e : Expression?) { if ((scn == "tmm2d_q8_f32" || scn == "tmm2d_q8_f16s") && length(ecs.arguments) == 14) { mark_write_target(ctx, ecs.arguments[11]) } + if (scn == "tmm2d_tg_store" && length(ecs.arguments) == 5) { + mark_write_target(ctx, ecs.arguments[1]) + } + if (scn == "tmm2d_tg_begin" && length(ecs.arguments) == 5) { + var acc0 = ecs.arguments[0] + if (acc0 is ExprRef2Value) { + acc0 = (acc0 as ExprRef2Value).subexpr + } + if (acc0 is ExprVar) { + ctx.tg_accs[string((acc0 as ExprVar).name)] = "" // params fill at emit time + } + } for (arg in ecs.arguments) { scan_body(ctx, arg) } @@ -1257,6 +1270,91 @@ def private tmm2d_helper_text(kind : string; m, n, sgs : uint) : string { } } +// acc arg (position 0 of every tg-protocol call) -> the das variable name, MSL-mangled +def private tg_acc_name(var ctx : MslCtx; e : Expression?) : string { + var a = e + if (a is ExprRef2Value) { + a = (a as ExprRef2Value).subexpr + } + if (a is ExprVar) { + return msl_safe_name(string((a as ExprVar).name)) + } + err(ctx, e.at, "tmm2d_tg_*: acc must be a plain local variable") + return "__bad_acc" +} + +def private emit_tmm2d_tg_begin(var ctx : MslCtx; ec : ExprCall?; indent : string; var lines : array) { + if (length(ec.arguments) != 5) { + err(ctx, ec.at, "tmm2d_tg_begin takes exactly 5 arguments") + return + } + let acc = tg_acc_name(ctx, ec.arguments[0]) + let m = tmm2d_const_uint(ctx, ec.arguments[1], "m (tile rows)") + let n = tmm2d_const_uint(ctx, ec.arguments[2], "n (tile cols)") + let sgs = tmm2d_const_uint(ctx, ec.arguments[3], "sgs (simdgroups)") + let kk = tmm2d_const_uint(ctx, ec.arguments[4], "kk (chunk K width)") + if (m == 0u || n == 0u || sgs == 0u || kk == 0u) { + return + } + var a0 = ec.arguments[0] + if (a0 is ExprRef2Value) { + a0 = (a0 as ExprRef2Value).subexpr + } + if (a0 is ExprVar) { + ctx.tg_accs[string((a0 as ExprVar).name)] = "{int(m)}|{int(n)}|{int(sgs)}|{int(kk)}" + } + note(ctx, "call.tmm2d_tg_begin") + lines |> push("{indent}constexpr auto __tgd_{acc} = matmul2d_descriptor({int(m)}, {int(n)}, {int(kk)}, false, true, false);") + lines |> push("{indent}matmul2d<__tgd_{acc}, execution_simdgroups<{int(sgs)}>> __tgop_{acc};") + lines |> push("{indent}using __tgt_{acc} = tensor, tensor_inline>;") + lines |> push("{indent}auto __tgf_{acc} = __tgop_{acc}.template get_destination_cooperative_tensor<__tgt_{acc}, __tgt_{acc}, float>();") + lines |> push("{indent}for (uint16_t __i = 0; __i < __tgf_{acc}.get_capacity(); ++__i) \{ if (__tgf_{acc}.is_valid_element(__i)) __tgf_{acc}[__i] = 0.0f; \}") +} + +def private emit_tmm2d_tg_step(var ctx : MslCtx; ec : ExprCall?; indent : string; var lines : array) { + if (length(ec.arguments) != 6) { + err(ctx, ec.at, "tmm2d_tg_step takes exactly 6 arguments") + return + } + let acc = tg_acc_name(ctx, ec.arguments[0]) + let ta = emit_value(ctx, ec.arguments[1]) + let tb = emit_value(ctx, ec.arguments[2]) + let m = tmm2d_const_uint(ctx, ec.arguments[3], "m (tile rows)") + let n = tmm2d_const_uint(ctx, ec.arguments[4], "n (tile cols)") + let kk = tmm2d_const_uint(ctx, ec.arguments[5], "kk (chunk K width)") + if (m == 0u || n == 0u || kk == 0u) { + return + } + note(ctx, "call.tmm2d_tg_step") + lines |> push("{indent}\{") + lines |> push("{indent} auto __A = __tgt_{acc}((threadgroup half *){ta}, dextents({int(kk)}, {int(m)}));") + lines |> push("{indent} auto __B = __tgt_{acc}((threadgroup half *){tb}, dextents({int(kk)}, {int(n)}));") + lines |> push("{indent} auto __c = __tgop_{acc}.template get_destination_cooperative_tensor<__tgt_{acc}, __tgt_{acc}, float>();") + lines |> push("{indent} __tgop_{acc}.run(__A, __B, __c);") + lines |> push("{indent} for (uint16_t __i = 0; __i < __tgf_{acc}.get_capacity(); ++__i) \{ if (__tgf_{acc}.is_valid_element(__i)) __tgf_{acc}[__i] += __c[__i]; \}") + lines |> push("{indent}\}") +} + +def private emit_tmm2d_tg_store(var ctx : MslCtx; ec : ExprCall?; indent : string; var lines : array) { + if (length(ec.arguments) != 5) { + err(ctx, ec.at, "tmm2d_tg_store takes exactly 5 arguments") + return + } + let acc = tg_acc_name(ctx, ec.arguments[0]) + let cp = emit_value(ctx, ec.arguments[1]) + let m = tmm2d_const_uint(ctx, ec.arguments[2], "m (tile rows)") + let n = tmm2d_const_uint(ctx, ec.arguments[3], "n (tile cols)") + let ldc = emit_value(ctx, ec.arguments[4]) + if (m == 0u || n == 0u) { + return + } + note(ctx, "call.tmm2d_tg_store") + lines |> push("{indent}\{") + lines |> push("{indent} auto __C = tensor, tensor_inline>({cp}, dextents({int(n)}, {int(m)}), metal::array\{1, int({ldc})\});") + lines |> push("{indent} __tgf_{acc}.store(__C);") + lines |> push("{indent}\}") +} + def private emit_tmm2d(var ctx : MslCtx; ec : ExprCall?; indent : string; var lines : array) { let cname = call_base_name("{ec.name}") let kind = cname == "tmm2d_f32_bf16_f32" ? "bf" : (cname == "tmm2d_q8b_f32" ? "q8b" : (cname == "tmm2d_q8_f16s" ? "q8h" : "q8")) @@ -1295,6 +1393,9 @@ def private emit_stmt(var ctx : MslCtx; e : Expression?; indent : string; var li emit_stmt(ctx, (e as ExprWith).body, indent, lines) } elif (e is ExprLet) { for (v in (e as ExprLet).variables) { + if (key_exists(ctx.tg_accs, string(v.name))) { + continue // staged-GEMM accumulator: CPU-replay state only — the GPU uses a cooperative tensor + } let sgname = sgmat_type_name(v._type) if (!empty(sgname)) { let sgVarName = string(v.name) @@ -1471,6 +1572,12 @@ def private emit_stmt(var ctx : MslCtx; e : Expression?; indent : string; var li emit_sgmat_multiply(ctx, ec, indent, lines) } elif (cname == "tmm2d_f32_bf16_f32" || cname == "tmm2d_q8_f32" || cname == "tmm2d_q8_f16s" || cname == "tmm2d_q8b_f32") { emit_tmm2d(ctx, ec, indent, lines) + } elif (cname == "tmm2d_tg_begin") { + emit_tmm2d_tg_begin(ctx, ec, indent, lines) + } elif (cname == "tmm2d_tg_step") { + emit_tmm2d_tg_step(ctx, ec, indent, lines) + } elif (cname == "tmm2d_tg_store") { + emit_tmm2d_tg_store(ctx, ec, indent, lines) } else { err(ctx, ec.at, "call `{cname}` as a statement has no MSL lowering — only barrier()/memoryBarrierShared() sync and the simdgroup_matrix ops are supported") } @@ -1713,9 +1820,9 @@ def public generate_msl(fn : FunctionPtr; var errors : array; cfg : MslE helper_text = join(hs, "\n") + "\n" } var mpp_include = "" - if (!empty(ctx.gen_helpers)) { - // Metal-4 tensor ops: the MPP include + namespace, then the generated helpers in - // pinned (sorted) order — table iteration order is hash-based + if (!empty(ctx.gen_helpers) || !empty(ctx.tg_accs)) { + // Metal-4 tensor ops: the MPP include + namespace (the tg protocol emits inline — + // no helpers, include still required), then generated helpers in pinned (sorted) order mpp_include = "#include \nusing namespace mpp::tensor_ops;\n" var gh <- [for (k in keys(ctx.gen_helpers)); k] gh |> sort() diff --git a/tests/metal/test_metal_tensor_ops.das b/tests/metal/test_metal_tensor_ops.das index 0a8d9bd947..dcb141e43e 100644 --- a/tests/metal/test_metal_tensor_ops.das +++ b/tests/metal/test_metal_tensor_ops.das @@ -60,6 +60,36 @@ class TmmKernels { tmm2d_q8_f16s(8u, 16u, 2u, ap, 64u, sp, 2u, wp, 64u, vp, 2u, cp, 16u, 64u) } + @workgroup tga : float16[256] // tg protocol: 8 x 32 flat A chunk + @workgroup tgb : float16[512] // 16 x 32 flat B chunk + + [metal_kernel(name="tmm_tg_msl"), marker(no_coverage)] + def tmm_tg_k { + var acc : float[128] + var ap = unsafe(addr(xa[0])) + var wp = unsafe(addr(wb[0])) + var cp = unsafe(addr(oc[0])) + tmm2d_tg_begin(acc, 8u, 16u, 2u, 32u) + var kb = 0u + while (kb < 2u) { + var i = gl_LocalInvocationID.x + while (i < 256u) { + tga[i] = float16(unsafe(ap[(i / 32u) * 64u + kb * 32u + i % 32u])) + i += gl_WorkGroupSize.x + } + i = gl_LocalInvocationID.x + while (i < 512u) { + tgb[i] = float16(uint_bits_to_float(uint(unsafe(wp[(i / 32u) * 64u + kb * 32u + i % 32u])) << 16u)) + i += gl_WorkGroupSize.x + } + barrier() + tmm2d_tg_step(acc, tga, tgb, 8u, 16u, 32u) + barrier() + kb++ + } + tmm2d_tg_store(acc, cp, 8u, 16u, 16u) + } + [metal_kernel(name="tmm_q8b_msl"), marker(no_coverage)] def tmm_q8b_k { var sp = unsafe(addr(bsh[0])) @@ -114,6 +144,8 @@ def test_tensor_ops(t : T?) { t |> success(find(tmm_q8_msl, "__das_tmm2d_q8_8_16_2") >= 0, "q8 helper generated") t |> success(find(tmm_q8h_msl, "__das_tmm2d_q8h_8_16_2") >= 0, "q8h helper generated") t |> success(find(tmm_q8b_msl, "__das_tmm2d_q8b_8_16_2") >= 0, "q8b helper generated") + t |> success(find(tmm_tg_msl, "__tgop_acc.run(") >= 0, "tg protocol op emitted") + t |> success(find(tmm_tg_msl, "float acc[") < 0, "acc local suppressed on GPU") t |> success(find(tmm_bf16_msl, "MetalPerformancePrimitives") >= 0, "MPP include present") t |> success(find(tmm_bf16_msl, "mpp::tensor_ops") >= 0, "namespace present") t |> success(find(tmm_bf16_msl, "__das_tmm2d_bf_8_16_2(") < find(tmm_bf16_msl, "kernel void"), "helper precedes the kernel") @@ -122,6 +154,7 @@ def test_tensor_ops(t : T?) { var oracle = new TmmKernels() fill_oracle(oracle) gl_LocalInvocationID = uint3(0u, 0u, 0u) + gl_WorkGroupSize = uint3(1u, 1u, 1u) // width-1 replay: staging strides cover whole tiles var want_bf : array var want_q8 : array var want_q8h : array @@ -134,6 +167,9 @@ def test_tensor_ops(t : T?) { var want_q8b : array oracle->tmm_q8b_k() want_q8b := oracle.oc + var want_tg : array + oracle->tmm_tg_k() + want_tg := oracle.oc var nz = 0 for (v in want_q8) { nz += v != 0.0 ? 1 : 0 @@ -157,9 +193,9 @@ def test_tensor_ops(t : T?) { var bbsh = buf_upload(dev, oracle.bsh) var bbqb = buf_upload(dev, oracle.bqb) gpu_bad = 0 - for (which in range(4)) { - let src = which == 0 ? tmm_bf16_msl : (which == 1 ? tmm_q8_msl : (which == 2 ? tmm_q8h_msl : tmm_q8b_msl)) - let entry = which == 0 ? tmm_bf16_msl_entry : (which == 1 ? tmm_q8_msl_entry : (which == 2 ? tmm_q8h_msl_entry : tmm_q8b_msl_entry)) + for (which in range(5)) { + let src = which == 0 ? tmm_bf16_msl : (which == 1 ? tmm_q8_msl : (which == 2 ? tmm_q8h_msl : (which == 3 ? tmm_q8b_msl : tmm_tg_msl))) + let entry = which == 0 ? tmm_bf16_msl_entry : (which == 1 ? tmm_q8_msl_entry : (which == 2 ? tmm_q8h_msl_entry : (which == 3 ? tmm_q8b_msl_entry : tmm_tg_msl_entry))) var perr : string var pso = pipeline_from_source(dev, src, entry, false, perr) t |> success(pso != null, "pipeline {entry}: {perr}") @@ -172,13 +208,13 @@ def test_tensor_ops(t : T?) { var bufs <- [bxa, bwb, boc, baq, basf, bwq, bwsf, bwsh, bbsh, bbqb] var rerr : string let tpg = uint(tew) * SGS - let ran = run_compute_1d(queue, pso, bufs, tpg, tpg, rerr) + let ran = run_compute_1d(queue, pso, bufs, tpg, tpg, rerr, which == 4 ? tmm_tg_msl_tgmem : 0ul) t |> success(ran, "dispatch {entry}: {rerr}") if (ran) { - let bad = buf_mismatch_exact(boc, which == 0 ? want_bf : (which == 1 ? want_q8 : (which == 2 ? want_q8h : want_q8b))) + let bad = buf_mismatch_exact(boc, which == 0 ? want_bf : (which == 1 ? want_q8 : (which == 2 ? want_q8h : (which == 3 ? want_q8b : want_tg)))) if (bad != 0) { var got <- buf_download(boc, M * N, type) - let wref & = unsafe(which == 0 ? want_bf : (which == 1 ? want_q8 : (which == 2 ? want_q8h : want_q8b))) + let wref & = unsafe(which == 0 ? want_bf : (which == 1 ? want_q8 : (which == 2 ? want_q8h : (which == 3 ? want_q8b : want_tg)))) var shown = 0 for (i in range(M * N)) { if (got[i] != wref[i] && shown < 4) { From a80b2a14858605331ec695a7bae163628cdcf61f Mon Sep 17 00:00:00 2001 From: Boris Batkin Date: Mon, 27 Jul 2026 22:35:43 -0700 Subject: [PATCH 10/18] dasLLAMA: MoeMulMmMx4 staged-tile twin; @role becomes optional everywhere MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The first tmm2d_tg consumer: flat [rows][32] f16 tiles per K-chunk with per-element doubled-e2m1 LUT dequant, per-expert bias folded as one rank-1 extra chunk (f16-staged — parity-tolerance vs the base's f32 accumulator seed), contiguous down-site only behind the "moe_mulmm_mx4" crown (gate/up keep the simdgroup kernel per dispatch). Gate: gptoss matrix row force-crowned — engage + logits-tolerance cells green. @role on @ssbo fields is now optional on un-lensed classes too (Boris): undeclared = the body derivation is the truth, declared still cross-checks. The new twin classes carry none. The access classifier suffix-matches the tg-protocol trio — generic instances carry a module prefix in post-infer call names that exact matching missed. Co-Authored-By: Claude Fable 5 --- .../dasllama/dasllama_kernel_access.das | 5 +- .../dasllama/dasllama_metal_kernels.das | 24 ++-- .../dasLLAMA/dasllama/dasllama_metal_lens.das | 4 +- .../dasllama/dasllama_metal_prefill.das | 126 +++++++++++++++--- 4 files changed, 125 insertions(+), 34 deletions(-) diff --git a/modules/dasLLAMA/dasllama/dasllama_kernel_access.das b/modules/dasLLAMA/dasllama/dasllama_kernel_access.das index 343597cb62..cbda3406db 100644 --- a/modules/dasLLAMA/dasllama/dasllama_kernel_access.das +++ b/modules/dasLLAMA/dasllama/dasllama_kernel_access.das @@ -6,6 +6,7 @@ module dasllama_kernel_access shared public require daslib/ast require daslib/ast_boost require daslib/rtti +require strings // ends_with — generic-instance call names carry a module prefix // Interprocedural read/write classifier over a declared set of tracked module globals — the // shared analysis half of the GPU dispatch lenses (vulkan today, metal's @role next). Given a @@ -198,12 +199,12 @@ class private AccessVisitor : AstVisitor { } if (cname == "coopmatLoad" || cname == "simdgroup_load" || cname == "coopmatLoadTensor" || cname == "coopmatLoadTensorDecode" - || cname == "tmm2d_tg_begin" || cname == "tmm2d_tg_step") { + || cname |> ends_with("tmm2d_tg_begin") || cname |> ends_with("tmm2d_tg_step")) { // plain reads (the generic pass records them) — or, for the staged-GEMM protocol // begin/step, only untracked locals/@workgroup tiles return } - if (cname == "tmm2d_tg_store") { + if (cname |> ends_with("tmm2d_tg_store")) { if (length(expr.arguments) == 5) { var cur = expr.arguments[1] if (cur is ExprRef2Value) { diff --git a/modules/dasLLAMA/dasllama/dasllama_metal_kernels.das b/modules/dasLLAMA/dasllama/dasllama_metal_kernels.das index 07cd3c96d7..c47099d136 100644 --- a/modules/dasLLAMA/dasllama/dasllama_metal_kernels.das +++ b/modules/dasLLAMA/dasllama/dasllama_metal_kernels.das @@ -3165,10 +3165,10 @@ class MetalSkReduce { // tmm2d_q8b_f32 each (raced via runtime.metal_tensor — crowns "gemmb_q8" / "gemmb_sk_q8" / // "gemm64b_q8"). The split-K twin passes the FULL row stride in blocks (ldwb) with a k-slice kk. class MetalQ8GemmBT { - @ssbo @binding = 0 @role = "weight" wqb : array // 34B-block W blob, byte view - @ssbo @binding = 2 @role = "weight" wsh : array // the same buffer, half-scale view - @ssbo @binding = 3 @role = "read" x : array // X [M' x n] f32 rows - @ssbo @binding = 4 @role = "write" y : array // Y [M' x d] + @ssbo @binding = 0 wqb : array // 34B-block W blob, byte view + @ssbo @binding = 2 wsh : array // the same buffer, half-scale view + @ssbo @binding = 3 x : array // X [M' x n] f32 rows + @ssbo @binding = 4 y : array // Y [M' x d] @uniform @binding = 5 kdim : uint @uniform @binding = 6 ndim : uint @uniform @binding = 7 ys : uint @@ -3188,10 +3188,10 @@ class MetalQ8GemmBT { } class MetalQ8GemmBSkT { - @ssbo @binding = 0 @role = "weight" wqb : array // 34B-block W blob, byte view - @ssbo @binding = 2 @role = "weight" wsh : array // the same buffer, half-scale view - @ssbo @binding = 3 @role = "read" x : array // X [32 x n] f32 rows - @ssbo @binding = 4 @role = "write" part : array // [ksplit x 32 x d] partial-C planes + @ssbo @binding = 0 wqb : array // 34B-block W blob, byte view + @ssbo @binding = 2 wsh : array // the same buffer, half-scale view + @ssbo @binding = 3 x : array // X [32 x n] f32 rows + @ssbo @binding = 4 part : array // [ksplit x 32 x d] partial-C planes @uniform @binding = 5 kdim : uint @uniform @binding = 6 ndim : uint @uniform @binding = 8 ksplit : uint @@ -3214,10 +3214,10 @@ class MetalQ8GemmBSkT { } class MetalQ8Gemm64BT { - @ssbo @binding = 0 @role = "weight" wqb : array // 34B-block W blob, byte view - @ssbo @binding = 2 @role = "weight" wsh : array // the same buffer, half-scale view - @ssbo @binding = 3 @role = "read" x : array // X [M' x n] f32 rows - @ssbo @binding = 4 @role = "write" y : array // Y [M' x d] + @ssbo @binding = 0 wqb : array // 34B-block W blob, byte view + @ssbo @binding = 2 wsh : array // the same buffer, half-scale view + @ssbo @binding = 3 x : array // X [M' x n] f32 rows + @ssbo @binding = 4 y : array // Y [M' x d] @uniform @binding = 5 kdim : uint @uniform @binding = 6 ndim : uint diff --git a/modules/dasLLAMA/dasllama/dasllama_metal_lens.das b/modules/dasLLAMA/dasllama/dasllama_metal_lens.das index cac096c9d2..851bbabf37 100644 --- a/modules/dasLLAMA/dasllama/dasllama_metal_lens.das +++ b/modules/dasLLAMA/dasllama/dasllama_metal_lens.das @@ -547,9 +547,7 @@ class MetalManualDispatchCensus : AstPassMacro { } let role = find_arg(field.annotation, "role") ?as tString ?? "" if (empty(role)) { - if (!lensed) { - macro_error(prog, field.at, "@ssbo field '{field.name}' of {st.name} has no @role (read/write/readwrite/weight/alias) — un-lensed classes declare their hazard intent") - } + pass // optional everywhere: undeclared = the body derivation is the truth } elif (!role_ok(role)) { macro_error(prog, field.at, "@ssbo field '{field.name}' of {st.name} has unknown @role '{role}' (read/write/readwrite/weight/alias)") } diff --git a/modules/dasLLAMA/dasllama/dasllama_metal_prefill.das b/modules/dasLLAMA/dasllama/dasllama_metal_prefill.das index 0e932a6730..2e2095e9de 100644 --- a/modules/dasLLAMA/dasllama/dasllama_metal_prefill.das +++ b/modules/dasLLAMA/dasllama/dasllama_metal_prefill.das @@ -257,10 +257,10 @@ class MetalBf16MulMm { // (blob bound twice — half-scale + byte views), lowered to one tmm2d_q8b_f32. Raced vs the // simdgroup kernel by the tuner (runtime.metal_tensor crown "mulmm_q8"); dispatch identical. class MetalQ8MulMmT { - @ssbo @binding = 0 @role = "weight" wsh : array // 34B-block W blob, half-scale view - @ssbo @binding = 1 @role = "weight" wqb : array // the SAME blob buffer, byte view - @ssbo @binding = 2 @role = "read" xf : array // raw f32 activations - @ssbo @binding = 3 @role = "write" y : array + @ssbo @binding = 0 wsh : array // 34B-block W blob, half-scale view + @ssbo @binding = 1 wqb : array // the SAME blob buffer, byte view + @ssbo @binding = 2 xf : array // raw f32 activations + @ssbo @binding = 3 y : array @uniform @binding = 4 kdim : uint @uniform @binding = 5 ndim : uint @@ -281,9 +281,9 @@ class MetalQ8MulMmT { // binds, lowered to mpp matmul2d over tensor_inline views. Raced vs the simdgroup kernel by // the tuner (runtime.metal_tensor crown "mulmm_bf16"); dispatch shape identical (128 = 4 sg). class MetalBf16MulMmT { - @ssbo @binding = 0 @role = "weight" wbh : array // bf16 W rows, halfword view - @ssbo @binding = 2 @role = "read" xf : array // raw f32 activations - @ssbo @binding = 3 @role = "write" y : array + @ssbo @binding = 0 wbh : array // bf16 W rows, halfword view + @ssbo @binding = 2 xf : array // raw f32 activations + @ssbo @binding = 3 y : array @uniform @binding = 4 kdim : uint @uniform @binding = 5 ndim : uint @@ -1374,6 +1374,86 @@ class MetalMoeGemvQ51 { } } +// Metal-4 staged-tile twin of MetalMoeMulMmMx4 — CONTIGUOUS (down-site) shape: flat f16 tiles +// per K-chunk (per-element LUT dequant), bias folded as one extra rank-1 chunk (f16-staged — +// parity-tolerance vs the base's f32 seed). +class MetalMoeMulMmMx4T { + @ssbo @binding = 0 mxq4 : array // nibble plane, uint4 view + @ssbo @binding = 1 mxe : array // e8m0 scale plane + @ssbo @binding = 2 xf : array + @ssbo @binding = 3 y : array + @uniform @binding = 4 kdim : uint + @uniform @binding = 5 ndim : uint + @ssbo @binding = 6 cnt : array + @ssbo @binding = 7 basep : array + @uniform @binding = 9 eblk : uint + @ssbo @binding = 12 wb : array // [ne x ndim] per-expert bias rows + @ssbo @binding = 13 vtab : array // [16] signed doubled-e2m1 values + @uniform @binding = 14 hasb : uint + @workgroup twa : float16[1024] // X chunk: 32 tokens x 32 k, flat row-major + @workgroup twb : float16[2048] // W chunk: 64 wrows x 32 k, flat row-major + + [metal_kernel(name="metal_moe_mulmm_mx4_t_msl")] + def metal_moe_mulmm_mx4_t { + let e = gl_WorkGroupID.z + let ce = cnt[e] + let mBase = gl_WorkGroupID.x * 32u + if (mBase >= (ce + 31u) / 32u * 32u) { // threadgroup-uniform exit — cooperative-safe + return + } + let rbase = basep[e] + let nBase = gl_WorkGroupID.y * 64u + let nkb = kdim / 32u + let lid = gl_LocalInvocationID.x + var acc : float[2048] + var cp = unsafe(addr(y[(rbase + mBase) * ndim + nBase])) + tmm2d_tg_begin(acc, 32u, 64u, 4u, 32u) + var kb = 0u + while (kb < nkb) { + var i = lid + while (i < 1024u) { + twa[i] = float16(xf[(rbase + mBase + i / 32u) * kdim + kb * 32u + i % 32u]) + i += gl_WorkGroupSize.x + } + i = lid + while (i < 2048u) { + let j = i / 32u + let k = i % 32u + let blk = e * eblk + (nBase + j) * nkb + kb + let sc = uint(mxe[blk]) + let d = uint_bits_to_float(sc < 2u ? (0x00200000u << sc) : ((sc - 1u) << 23u)) + let byi = k % 16u + let u4 = mxq4[blk] + let w32 = byi < 4u ? u4.x : (byi < 8u ? u4.y : (byi < 12u ? u4.z : u4.w)) + let nby = (w32 >> ((byi % 4u) * 8u)) & 255u + twb[i] = float16(vtab[int(k < 16u ? (nby & 15u) : (nby >> 4u))] * d) + i += gl_WorkGroupSize.x + } + barrier() + tmm2d_tg_step(acc, twa, twb, 32u, 64u, 32u) + barrier() + kb++ + } + if (hasb != 0u) { + // rank-1 bias chunk: A col 0 = 1, W col 0 = wb[e][col] + var i = lid + while (i < 1024u) { + twa[i] = float16(i % 32u == 0u ? 1.0 : 0.0) + i += gl_WorkGroupSize.x + } + i = lid + while (i < 2048u) { + twb[i] = float16(i % 32u == 0u ? wb[e * ndim + nBase + i / 32u] : 0.0) + i += gl_WorkGroupSize.x + } + barrier() + tmm2d_tg_step(acc, twa, twb, 32u, 64u, 32u) + barrier() + } + tmm2d_tg_store(acc, cp, 32u, 64u, ndim) + } +} + // Gathered MXFP4 mul_mm — the gathered-GEMM shape with the mx4 A stage: each thread's 16-elem // half IS one nibble half of the 32-block (il0 = 0 low, 1 high). Per-expert bias (hasb) seeds // the accumulator tiles (stride-0 loads of wb[e]) — ggml_add_id order, bias precedes the reduce. @@ -1779,14 +1859,14 @@ class MetalMoeReduce { // bkt-indirect X rows cannot form a strided tensor view, so gate/up keep the simdgroup kernel // even when "moe_mulmm_q8" is crowned (the encode picks per dispatch). class MetalMoeMulMmQ8T { - @ssbo @binding = 0 @role = "weight" wsh : array // expert STACK blob, half-scale view - @ssbo @binding = 1 @role = "weight" wqb : array // the same blob buffer, byte view - @ssbo @binding = 2 @role = "read" xf : array - @ssbo @binding = 3 @role = "write" y : array + @ssbo @binding = 0 wsh : array // expert STACK blob, half-scale view + @ssbo @binding = 1 wqb : array // the same blob buffer, byte view + @ssbo @binding = 2 xf : array + @ssbo @binding = 3 y : array @uniform @binding = 4 kdim : uint @uniform @binding = 5 ndim : uint - @ssbo @binding = 6 @role = "read" cnt : array - @ssbo @binding = 7 @role = "read" basep : array + @ssbo @binding = 6 cnt : array + @ssbo @binding = 7 basep : array @uniform @binding = 9 eblk : uint [metal_kernel(name="metal_moe_mulmm_q8_t_msl")] @@ -4562,6 +4642,8 @@ var private g_pf_bf16_mm_tensor : bool // crowned tensor twin selected (no tgm var private g_pf_mm_tensor : bool // ditto for the production q8 mul_mm var private g_pf_moe_mm_q8_tensor : bool // "moe_mulmm_q8" crowned (contiguous sites only) var private g_pf_pso_moe_mm_q8_t : MetalComputePipeline? +var private g_pf_moe_mm_mx4_tensor : bool // "moe_mulmm_mx4" crowned (contiguous sites only) +var private g_pf_pso_moe_mm_mx4_t : MetalComputePipeline? var private g_pf_pso_ple_gather : MetalComputePipeline? // PLE pre-step: q8 token-row gather var private g_pf_pso_ple_finish : MetalComputePipeline? // PLE pre-step: rms + combine, in place var private g_pso_qkmm : MetalComputePipeline? @@ -4692,6 +4774,10 @@ def public metal_prefill_shutdown { metal_release(g_pf_pso_moe_mm_q8_t) g_pf_pso_moe_mm_q8_t = null } + if (g_pf_pso_moe_mm_mx4_t != null) { + metal_release(g_pf_pso_moe_mm_mx4_t) + g_pf_pso_moe_mm_mx4_t = null + } if (g_pso_gemm != null) { metal_release(g_pso_gemm) g_pso_gemm = null @@ -4963,6 +5049,10 @@ def private metal_prefill_init : bool { g_pf_pso_moe_wscale = pf_compile_pso(metal_moe_wscale_msl, metal_moe_wscale_msl_entry, metal_moe_wscale_msl_fastmath, ok) g_pf_pso_swiglu_oai = pf_compile_pso(metal_swiglu_oai_pf_msl, metal_swiglu_oai_pf_msl_entry, metal_swiglu_oai_pf_msl_fastmath, ok) g_pf_pso_moe_mm_mx4 = pf_compile_pso(metal_moe_mulmm_mx4_msl, metal_moe_mulmm_mx4_msl_entry, metal_moe_mulmm_mx4_msl_fastmath, ok) + g_pf_moe_mm_mx4_tensor = metal_tensor_crowned("moe_mulmm_mx4") + if (g_pf_moe_mm_mx4_tensor) { + g_pf_pso_moe_mm_mx4_t = pf_compile_pso(metal_moe_mulmm_mx4_t_msl, metal_moe_mulmm_mx4_t_msl_entry, metal_moe_mulmm_mx4_t_msl_fastmath, ok) + } g_pf_pso_moe_mm_q51 = pf_compile_pso(metal_moe_mulmm_q51_msl, metal_moe_mulmm_q51_msl_entry, metal_moe_mulmm_q51_msl_fastmath, ok) g_pf_pso_dn_conv = pf_compile_pso(metal_dn_conv_msl, metal_dn_conv_msl_entry, metal_dn_conv_msl_fastmath, ok) g_pf_pso_dn_hist = pf_compile_pso(metal_dn_conv_hist_msl, metal_dn_conv_hist_msl_entry, metal_dn_conv_hist_msl_fastmath, ok) @@ -5388,12 +5478,14 @@ def private pf_enc_moe_mm(enc : MetalComputeEncoder?; t : Model; fmt : KqFmt; st // gemma4 router-input norm rows (weightless rms * 1/sqrt(dim) * rscale, one tg per position) // the MXFP4 gathered expert mul_mm (gpt-oss) def private pf_enc_moe_mm_mx4(enc : MetalComputeEncoder?; t : Model; stack_off, rows, npos : int64; - bx, by, bkdim, bndim, bcnt, bbase, bbkt, beblk, bnk, bgather, bwb, bhasb : MetalBuffer?) { + bx, by, bkdim, bndim, bcnt, bbase, bbkt, beblk, bnk, bgather, bwb, bhasb : MetalBuffer?; + contiguous : bool = false) { let tiles = (npos + 31l) / 32l let ne = t.config.n_expert let mp = mx4_of(g_pf_dev, t, stack_off) - kn_pipeline(enc, g_pf_pso_moe_mm_mx4) - kn_tgmem(enc, metal_moe_mulmm_mx4_msl_tgmem, 0) + let twin = g_pf_moe_mm_mx4_tensor && contiguous && g_pf_pso_moe_mm_mx4_t != null + kn_pipeline(enc, twin ? g_pf_pso_moe_mm_mx4_t : g_pf_pso_moe_mm_mx4) + kn_tgmem(enc, twin ? metal_moe_mulmm_mx4_t_msl_tgmem : metal_moe_mulmm_mx4_msl_tgmem, 0) kn_buffer(enc, mp.qbuf, mp.qoff, 0) kn_buffer(enc, mp.sbuf, mp.soff, 1) kn_buffer(enc, bx, 0ul, 2) @@ -6649,7 +6741,7 @@ def metal_prefill_forward(t : Model; var s : Session; npos, start_pos : int64) : var bhasb = g_pf_skip == "moe_bias" ? u_moe_g0 : u_moe_g1 var bwb2 = pf_upload_region(unsafe(addr < void? >(t.fblob[t.web2_off + l * c.n_expert * dim])), uint64(c.n_expert * dim * 4l)) pf_enc_moe_mm_mx4(enc, t, t.we2_offs[l], dim, npos, bmg, bmdn, - u_moe_nfe, u_dim, bmcnt, bmbase, bmbkt, u_moe_eblk, u_moe_nk, u_moe_g0, bwb2, bhasb) + u_moe_nfe, u_dim, bmcnt, bmbase, bmbkt, u_moe_eblk, u_moe_nk, u_moe_g0, bwb2, bhasb, contiguous = true) } else { pf_enc_moe_mm(enc, t, fe2, t.we2_offs[l], dim, npos, bmg, bmdn, u_moe_nfe, u_dim, bmcnt, bmbase, bmbkt, fe2 == KqFmt.q8 || fe2 == KqFmt.q51 ? u_moe_eblk : u_moe_esb, u_moe_nk, u_moe_g0, contiguous = true) From 6209b2154630fe5e28fa9e3efe0872197dc4ddb4 Mon Sep 17 00:00:00 2001 From: Boris Batkin Date: Mon, 27 Jul 2026 22:43:58 -0700 Subject: [PATCH 11/18] dasLLAMA: k-quant staged-tile twins (K4/K5/K6 mulmm) + kq tuner races MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit KqMulMmK4T/K5T/K6T on the tmm2d_tg rail: flat f16 tiles per K-chunk with per-element kmask / qh-plane / sub-scale dequant — same real products as the base kernels' staged va values. Crowns kq_mulmm_k4/k5/k6 picked in pf_enc_kq_site_mm per format. Races run on random quant planes with constant scales — the twin-vs-base envelope doubles as a layout-agreement proof of the per-element decodes (a missing twin tgmem bind initially zeroed the staged tiles in the race; staged twins bind their own tgmem). Gate: llama prefill kq arm force-crowned (k4 alone, then all three) — token parity green; all nine raced families envelope-clean on M1. Co-Authored-By: Claude Fable 5 --- .../dasllama/dasllama_metal_prefill.das | 362 +++++++++++++++++- 1 file changed, 358 insertions(+), 4 deletions(-) diff --git a/modules/dasLLAMA/dasllama/dasllama_metal_prefill.das b/modules/dasLLAMA/dasllama/dasllama_metal_prefill.das index 2e2095e9de..1f784712a1 100644 --- a/modules/dasLLAMA/dasllama/dasllama_metal_prefill.das +++ b/modules/dasLLAMA/dasllama/dasllama_metal_prefill.das @@ -469,6 +469,70 @@ class MetalKqMulMmK4 { } } +// Metal-4 staged-tile twin of MetalKqMulMmK4: flat f16 tiles per K-chunk, per-element kmask +// scale decode + nibble dequant (same math as the base's va — identical staged values). +class MetalKqMulMmK4T { + @ssbo @binding = 0 ksh : array // 16B scale blocks, half view (d, dmin at 8*blk) + @ssbo @binding = 1 ks4 : array // the same buffer, uint4 view + @ssbo @binding = 2 kqu : array // k4 quant plane, uint view + @ssbo @binding = 3 xf : array + @ssbo @binding = 4 y : array + @uniform @binding = 5 kdim : uint + @uniform @binding = 6 ndim : uint + @workgroup twa : float16[1024] // X chunk: 32 tokens x 32 k, flat row-major + @workgroup twb : float16[2048] // W chunk: 64 wrows x 32 k, flat row-major + + [metal_kernel(name="metal_kq_mulmm_k4_t_msl")] + def metal_kq_mulmm_k4_t { + let mBase = gl_WorkGroupID.x * 32u + let nBase = gl_WorkGroupID.y * 64u + let nkb = kdim / 32u + let nsb = kdim / 256u + let lid = gl_LocalInvocationID.x + var acc : float[2048] + var cp = unsafe(addr(y[mBase * ndim + nBase])) + tmm2d_tg_begin(acc, 32u, 64u, 4u, 32u) + var kb = 0u + while (kb < nkb) { + var i = lid + while (i < 1024u) { + twa[i] = float16(xf[(mBase + i / 32u) * kdim + kb * 32u + i % 32u]) + i += gl_WorkGroupSize.x + } + let sb = kb / 8u + let js = kb % 8u + i = lid + while (i < 2048u) { + let j = i / 32u + let e = i % 32u + let blk = (nBase + j) * nsb + sb + let sv = ks4[blk] + var sc = 0u + var mn = 0u + if (js < 4u) { + sc = (sv.y >> (8u * js)) & 63u + mn = (sv.z >> (8u * js)) & 63u + } else { + let j2 = js - 4u + let hi = sv.w >> (8u * j2) + sc = (hi & 15u) | (((sv.y >> (8u * j2 + 6u)) & 3u) << 4u) + mn = ((hi >> 4u) & 15u) | (((sv.z >> (8u * j2 + 6u)) & 3u) << 4u) + } + let dsc = float(ksh[blk * 8u]) * float(sc) + let dmn = float(ksh[blk * 8u + 1u]) * float(mn) + let q = (kqu[blk * 32u + (js / 2u) * 8u + e / 4u] >> (8u * (e % 4u) + (js & 1u) * 4u)) & 15u + twb[i] = float16(dsc * float(q) - dmn) + i += gl_WorkGroupSize.x + } + barrier() + tmm2d_tg_step(acc, twa, twb, 32u, 64u, 32u) + barrier() + kb++ + } + tmm2d_tg_store(acc, cp, 32u, 64u, ndim) + } +} + class MetalKqMulMmK5 { @ssbo @binding = 0 @role = "weight" ksh : array // 16B compact scale blocks (the Q4_K form) @ssbo @binding = 1 @role = "weight" ks4 : array @@ -570,6 +634,136 @@ class MetalKqMulMmK5 { } } +// Metal-4 staged-tile twin of MetalKqMulMmK5 (K4's scale decode + the qh high-bit plane). +class MetalKqMulMmK5T { + @ssbo @binding = 0 ksh : array + @ssbo @binding = 1 ks4 : array + @ssbo @binding = 2 kqu : array // qs at 40*blk, qh at 40*blk+32 + @ssbo @binding = 3 xf : array + @ssbo @binding = 4 y : array + @uniform @binding = 5 kdim : uint + @uniform @binding = 6 ndim : uint + @workgroup twa : float16[1024] + @workgroup twb : float16[2048] + + [metal_kernel(name="metal_kq_mulmm_k5_t_msl")] + def metal_kq_mulmm_k5_t { + let mBase = gl_WorkGroupID.x * 32u + let nBase = gl_WorkGroupID.y * 64u + let nkb = kdim / 32u + let nsb = kdim / 256u + let lid = gl_LocalInvocationID.x + var acc : float[2048] + var cp = unsafe(addr(y[mBase * ndim + nBase])) + tmm2d_tg_begin(acc, 32u, 64u, 4u, 32u) + var kb = 0u + while (kb < nkb) { + var i = lid + while (i < 1024u) { + twa[i] = float16(xf[(mBase + i / 32u) * kdim + kb * 32u + i % 32u]) + i += gl_WorkGroupSize.x + } + let sb = kb / 8u + let js = kb % 8u + i = lid + while (i < 2048u) { + let j = i / 32u + let e = i % 32u + let blk = (nBase + j) * nsb + sb + let sv = ks4[blk] + var sc = 0u + var mn = 0u + if (js < 4u) { + sc = (sv.y >> (8u * js)) & 63u + mn = (sv.z >> (8u * js)) & 63u + } else { + let j2 = js - 4u + let hi = sv.w >> (8u * j2) + sc = (hi & 15u) | (((sv.y >> (8u * j2 + 6u)) & 3u) << 4u) + mn = ((hi >> 4u) & 15u) | (((sv.z >> (8u * j2 + 6u)) & 3u) << 4u) + } + let dsc = float(ksh[blk * 8u]) * float(sc) + let dmn = float(ksh[blk * 8u + 1u]) * float(mn) + let u = kqu[blk * 40u + (js / 2u) * 8u + e / 4u] + let hu = kqu[blk * 40u + 32u + e / 4u] + let q = ((u >> (8u * (e % 4u) + (js & 1u) * 4u)) & 15u) | (((hu >> (8u * (e % 4u) + js)) & 1u) << 4u) + twb[i] = float16(dsc * float(q) - dmn) + i += gl_WorkGroupSize.x + } + barrier() + tmm2d_tg_step(acc, twa, twb, 32u, 64u, 32u) + barrier() + kb++ + } + tmm2d_tg_store(acc, cp, 32u, 64u, ndim) + } +} + +// Metal-4 staged-tile twin of MetalKqMulMmK6 (per-element ql|qh compose x signed sub-scale — +// same real product as the base's pre-scaled fma form, which is bit-identical by construction). +class MetalKqMulMmK6T { + @ssbo @binding = 0 kdh : array // d plane (scale buffer at byte nsb*16) + @ssbo @binding = 1 ksc : array // 16B sub-scale blocks, same buffer at 0 + @ssbo @binding = 2 kqu : array // ql at 48*blk, qh at 48*blk+32 + @ssbo @binding = 3 xf : array + @ssbo @binding = 4 y : array + @uniform @binding = 5 kdim : uint + @uniform @binding = 6 ndim : uint + @workgroup twa : float16[1024] + @workgroup twb : float16[2048] + + [metal_kernel(name="metal_kq_mulmm_k6_t_msl")] + def metal_kq_mulmm_k6_t { + let mBase = gl_WorkGroupID.x * 32u + let nBase = gl_WorkGroupID.y * 64u + let nkb = kdim / 32u + let nsb = kdim / 256u + let lid = gl_LocalInvocationID.x + var acc : float[2048] + var cp = unsafe(addr(y[mBase * ndim + nBase])) + tmm2d_tg_begin(acc, 32u, 64u, 4u, 32u) + var kb = 0u + while (kb < nkb) { + var i = lid + while (i < 1024u) { + twa[i] = float16(xf[(mBase + i / 32u) * kdim + kb * 32u + i % 32u]) + i += gl_WorkGroupSize.x + } + let sb = kb / 8u + let js = kb % 8u + let half6 = js / 4u + let gg = js % 4u + let nsh = (gg / 2u) * 4u + let hsh = gg * 2u + i = lid + while (i < 2048u) { + let j = i / 32u + let e = i % 32u + let blk = (nBase + j) * nsb + sb + let dall = float(kdh[blk]) + let ile = e / 16u + let k = (e % 16u) / 4u + let c = e % 4u + let si = half6 * 8u + gg * 2u + ile + let sv = ksc[blk] + let scw = si < 8u ? (si < 4u ? sv.x : sv.y) : (si < 12u ? sv.z : sv.w) + let s6 = float(((int(scw >> ((si & 3u) * 8u)) & 255) ^ 128) - 128) + let dsc = dall * s6 + let u = kqu[blk * 48u + half6 * 16u + (gg & 1u) * 8u + ile * 4u + k] + let hu = kqu[blk * 48u + 32u + half6 * 8u + ile * 4u + k] + let q = ((u >> (nsh + 8u * c)) & 15u) | (((hu >> (hsh + 8u * c)) & 3u) << 4u) + twb[i] = float16(dsc * float(q) - dsc * 32.0) + i += gl_WorkGroupSize.x + } + barrier() + tmm2d_tg_step(acc, twa, twb, 32u, 64u, 32u) + barrier() + kb++ + } + tmm2d_tg_store(acc, cp, 32u, 64u, ndim) + } +} + class MetalKqMulMmK6 { @ssbo @binding = 0 @role = "weight" kdh : array // the d plane — the scale buffer bound at byte nsb*16 @ssbo @binding = 1 @role = "weight" ksc : array // 16B sub-scale blocks (16 int8 at blk), same buffer at 0 @@ -4644,6 +4838,12 @@ var private g_pf_moe_mm_q8_tensor : bool // "moe_mulmm_q8" crowned (contiguous s var private g_pf_pso_moe_mm_q8_t : MetalComputePipeline? var private g_pf_moe_mm_mx4_tensor : bool // "moe_mulmm_mx4" crowned (contiguous sites only) var private g_pf_pso_moe_mm_mx4_t : MetalComputePipeline? +var private g_pf_kq_mm4_tensor : bool // "kq_mulmm_k4" crowned +var private g_pso_kq_mm4_t : MetalComputePipeline? +var private g_pf_kq_mm5_tensor : bool // "kq_mulmm_k5" crowned +var private g_pso_kq_mm5_t : MetalComputePipeline? +var private g_pf_kq_mm6_tensor : bool // "kq_mulmm_k6" crowned +var private g_pso_kq_mm6_t : MetalComputePipeline? var private g_pf_pso_ple_gather : MetalComputePipeline? // PLE pre-step: q8 token-row gather var private g_pf_pso_ple_finish : MetalComputePipeline? // PLE pre-step: rms + combine, in place var private g_pso_qkmm : MetalComputePipeline? @@ -4778,6 +4978,18 @@ def public metal_prefill_shutdown { metal_release(g_pf_pso_moe_mm_mx4_t) g_pf_pso_moe_mm_mx4_t = null } + if (g_pso_kq_mm4_t != null) { + metal_release(g_pso_kq_mm4_t) + g_pso_kq_mm4_t = null + } + if (g_pso_kq_mm5_t != null) { + metal_release(g_pso_kq_mm5_t) + g_pso_kq_mm5_t = null + } + if (g_pso_kq_mm6_t != null) { + metal_release(g_pso_kq_mm6_t) + g_pso_kq_mm6_t = null + } if (g_pso_gemm != null) { metal_release(g_pso_gemm) g_pso_gemm = null @@ -5023,8 +5235,20 @@ def private metal_prefill_init : bool { g_pf_pso_ple_gather = pf_compile_pso(metal_ple_gather_q8_msl, metal_ple_gather_q8_msl_entry, metal_ple_gather_q8_msl_fastmath, ok) g_pf_pso_ple_finish = pf_compile_pso(metal_ple_finish_msl, metal_ple_finish_msl_entry, metal_ple_finish_msl_fastmath, ok) g_pso_kq_mm4 = pf_compile_pso(metal_kq_mulmm_k4_msl, metal_kq_mulmm_k4_msl_entry, metal_kq_mulmm_k4_msl_fastmath, ok) + g_pf_kq_mm4_tensor = metal_tensor_crowned("kq_mulmm_k4") + if (g_pf_kq_mm4_tensor) { + g_pso_kq_mm4_t = pf_compile_pso(metal_kq_mulmm_k4_t_msl, metal_kq_mulmm_k4_t_msl_entry, metal_kq_mulmm_k4_t_msl_fastmath, ok) + } g_pso_kq_mm5 = pf_compile_pso(metal_kq_mulmm_k5_msl, metal_kq_mulmm_k5_msl_entry, metal_kq_mulmm_k5_msl_fastmath, ok) + g_pf_kq_mm5_tensor = metal_tensor_crowned("kq_mulmm_k5") + if (g_pf_kq_mm5_tensor) { + g_pso_kq_mm5_t = pf_compile_pso(metal_kq_mulmm_k5_t_msl, metal_kq_mulmm_k5_t_msl_entry, metal_kq_mulmm_k5_t_msl_fastmath, ok) + } g_pso_kq_mm6 = pf_compile_pso(metal_kq_mulmm_k6_msl, metal_kq_mulmm_k6_msl_entry, metal_kq_mulmm_k6_msl_fastmath, ok) + g_pf_kq_mm6_tensor = metal_tensor_crowned("kq_mulmm_k6") + if (g_pf_kq_mm6_tensor) { + g_pso_kq_mm6_t = pf_compile_pso(metal_kq_mulmm_k6_t_msl, metal_kq_mulmm_k6_t_msl_entry, metal_kq_mulmm_k6_t_msl_fastmath, ok) + } g_pf_pso_kq4 = pf_compile_pso(metal_kq_gemv_k4_msl, metal_kq_gemv_k4_msl_entry, metal_kq_gemv_k4_msl_fastmath, ok) g_pf_pso_kq5 = pf_compile_pso(metal_kq_gemv_k5_msl, metal_kq_gemv_k5_msl_entry, metal_kq_gemv_k5_msl_fastmath, ok) g_pf_pso_kq6 = pf_compile_pso(metal_kq_gemv_k6_msl, metal_kq_gemv_k6_msl_entry, metal_kq_gemv_k6_msl_fastmath, ok) @@ -5229,12 +5453,19 @@ def private pf_enc_kq_site_mm(enc : MetalComputeEncoder?; t : Model; fmt : KqFmt let bq = kq_quants_of(g_pf_dev, t, fmt, woff) let bs = kq_scales_of(g_pf_dev, t, fmt, woff) if (fmt == KqFmt.k6) { - kn_pipeline(enc, g_pso_kq_mm6) - kn_tgmem(enc, metal_kq_mulmm_k6_msl_tgmem, 0) + let twin6 = g_pf_kq_mm6_tensor && g_pso_kq_mm6_t != null + kn_pipeline(enc, twin6 ? g_pso_kq_mm6_t : g_pso_kq_mm6) + kn_tgmem(enc, twin6 ? metal_kq_mulmm_k6_t_msl_tgmem : metal_kq_mulmm_k6_msl_tgmem, 0) kn_buffer(enc, bs.buf, bs.doff, 0) + } elif (fmt == KqFmt.k4) { + let twin4 = g_pf_kq_mm4_tensor && g_pso_kq_mm4_t != null + kn_pipeline(enc, twin4 ? g_pso_kq_mm4_t : g_pso_kq_mm4) + kn_tgmem(enc, twin4 ? metal_kq_mulmm_k4_t_msl_tgmem : metal_kq_mulmm_k4_msl_tgmem, 0) + kn_buffer(enc, bs.buf, bs.soff, 0) } else { - kn_pipeline(enc, fmt == KqFmt.k4 ? g_pso_kq_mm4 : g_pso_kq_mm5) - kn_tgmem(enc, fmt == KqFmt.k4 ? metal_kq_mulmm_k4_msl_tgmem : metal_kq_mulmm_k5_msl_tgmem, 0) + let twin5 = g_pf_kq_mm5_tensor && g_pso_kq_mm5_t != null + kn_pipeline(enc, twin5 ? g_pso_kq_mm5_t : g_pso_kq_mm5) + kn_tgmem(enc, twin5 ? metal_kq_mulmm_k5_t_msl_tgmem : metal_kq_mulmm_k5_msl_tgmem, 0) kn_buffer(enc, bs.buf, bs.soff, 0) } kn_buffer(enc, bs.buf, bs.soff, 1) @@ -7440,6 +7671,120 @@ def private race_moe_mulmm_q8(dev : MetalDevice?; queue : MetalCommandQueue?) : return res } +// one kq mulmm race: random quant planes + constant f16 scales (any bit pattern is a valid +// plane — the twin-vs-base envelope doubles as a layout-agreement check). k6 binds the d plane +// at its tail offset; k4/k5 bind the compact scale blocks at 0. +def private race_kq_mulmm(dev : MetalDevice?; queue : MetalCommandQueue?; family : string; + base_src, base_entry : string; base_fm : bool; base_tgmem : uint64; + twin_src, twin_entry : string; twin_fm : bool; twin_tgmem : uint64; + qu_per_sb : int; k6 : bool) : MetalTensorRaceResult { + var res = MetalTensorRaceResult(family = family, winner = "", base_ms = -1.0lf, twin_ms = -1.0lf, note = "") + let m = 512 + let kdim = 2048 + let ndim = 1024 + let nsb = kdim / 256 + var err = "" + var base_pso = pipeline_from_source(dev, base_src, base_entry, base_fm, err) + if (base_pso == null) { + res.note = "base pso: {err}" + return res + } + var twin_pso = pipeline_from_source(dev, twin_src, twin_entry, twin_fm, err) + if (twin_pso == null) { + res.winner = "simdgroup" + res.note = "twin pso: {err}" + metal_release(base_pso) + return res + } + var xa <- race_x_f32(m * kdim) + let nblk = ndim * nsb + var qplane : array + qplane |> resize(nblk * qu_per_sb) + for (i in range(nblk * qu_per_sb)) { + qplane[i] = uint(i) * 2654435761u // deterministic pseudo-random bit pattern + } + // scales: k4/k5 = compact 16B blocks; k6 = int8 sub-scale plane + the f16 d tail + var splane : array + let sbytes = k6 ? nblk * 16 + nblk * 2 : nblk * 16 + splane |> resize(sbytes) + if (k6) { + for (i in range(nblk * 16)) { + splane[i] = uint8(2) + } + unsafe { + var ph = addr(splane[nblk * 16]) + for (b in range(nblk)) { + ph[b] = float16(0.25) + } + } + } else { + unsafe { + var ph = addr(splane[0]) + for (b in range(nblk)) { + ph[b * 8] = float16(0.25) // d + ph[b * 8 + 1] = float16(0.25) // dmin + } + } + for (b in range(nblk)) { + for (kx in range(4, 12)) { + splane[b * 16 + kx] = uint8(17) // kmask bytes: small scales, both halves + } + } + } + var bq = race_buf(dev, uint64(nblk * qu_per_sb * 4), unsafe(addr(qplane[0]))) + var bs = race_buf(dev, uint64(sbytes), unsafe(addr(splane[0]))) + var bxa = race_buf(dev, uint64(m * kdim * 4), unsafe(addr(xa[0]))) + var by_base = race_buf(dev, uint64(m * ndim * 4), null) + var by_twin = race_buf(dev, uint64(m * ndim * 4), null) + var bk = race_uniform_u32(dev, uint(kdim)) + var bn = race_uniform_u32(dev, uint(ndim)) + let doff = k6 ? uint64(nblk * 16) : 0ul + let grid = uint3(uint(m / 32), uint(ndim / 64), 1u) + let tg = uint3(128u, 1u, 1u) + res.base_ms = race_time_ms(queue, 5) $(enc) { + kn_pipeline(enc, base_pso) + kn_tgmem(enc, base_tgmem, 0) + kn_buffer(enc, bs, doff, 0) + kn_buffer(enc, bs, 0ul, 1) + kn_buffer(enc, bq, 0ul, 2) + kn_buffer(enc, bxa, 0ul, 3) + kn_buffer(enc, by_base, 0ul, 4) + kn_buffer(enc, bk, 0ul, 5) + kn_buffer(enc, bn, 0ul, 6) + kn_dispatch(enc, grid, tg) + } + res.twin_ms = race_time_ms(queue, 5) $(enc) { + kn_pipeline(enc, twin_pso) + kn_tgmem(enc, twin_tgmem, 0) + kn_buffer(enc, bs, doff, 0) + kn_buffer(enc, bs, 0ul, 1) + kn_buffer(enc, bq, 0ul, 2) + kn_buffer(enc, bxa, 0ul, 3) + kn_buffer(enc, by_twin, 0ul, 4) + kn_buffer(enc, bk, 0ul, 5) + kn_buffer(enc, bn, 0ul, 6) + kn_dispatch(enc, grid, tg) + } + if (res.base_ms < 0.0lf || res.twin_ms < 0.0lf) { + res.winner = res.base_ms < 0.0lf ? "" : "simdgroup" + res.note = "dispatch failed" + } elif (!race_envelope_ok(by_base, by_twin, m * ndim, res.note)) { + res.winner = "simdgroup" + } else { + res.winner = res.twin_ms < res.base_ms ? "tensor" : "simdgroup" + } + metal_release(bq) + metal_release(bs) + metal_release(bxa) + metal_release(by_base) + metal_release(by_twin) + metal_release(bk) + metal_release(bn) + metal_release(base_pso) + metal_release(twin_pso) + return res +} + //! Tuner entry (the runtime.metal_tensor crown): race every tensor pso twin against its //! shipped simdgroup kernel on THIS box — compile + output-match to qualify, faster GPU time //! takes the family's crown. Empty when there is no Metal device. @@ -7453,6 +7798,15 @@ def metal_tensor_race : array { results |> emplace(race_mulmm_bf16(dev, queue)) results |> emplace(race_mulmm_q8(dev, queue)) results |> emplace(race_moe_mulmm_q8(dev, queue)) + results |> emplace(race_kq_mulmm(dev, queue, "kq_mulmm_k4", + metal_kq_mulmm_k4_msl, metal_kq_mulmm_k4_msl_entry, metal_kq_mulmm_k4_msl_fastmath, metal_kq_mulmm_k4_msl_tgmem, + metal_kq_mulmm_k4_t_msl, metal_kq_mulmm_k4_t_msl_entry, metal_kq_mulmm_k4_t_msl_fastmath, metal_kq_mulmm_k4_t_msl_tgmem, 32, false)) + results |> emplace(race_kq_mulmm(dev, queue, "kq_mulmm_k5", + metal_kq_mulmm_k5_msl, metal_kq_mulmm_k5_msl_entry, metal_kq_mulmm_k5_msl_fastmath, metal_kq_mulmm_k5_msl_tgmem, + metal_kq_mulmm_k5_t_msl, metal_kq_mulmm_k5_t_msl_entry, metal_kq_mulmm_k5_t_msl_fastmath, metal_kq_mulmm_k5_t_msl_tgmem, 40, false)) + results |> emplace(race_kq_mulmm(dev, queue, "kq_mulmm_k6", + metal_kq_mulmm_k6_msl, metal_kq_mulmm_k6_msl_entry, metal_kq_mulmm_k6_msl_fastmath, metal_kq_mulmm_k6_msl_tgmem, + metal_kq_mulmm_k6_t_msl, metal_kq_mulmm_k6_t_msl_entry, metal_kq_mulmm_k6_t_msl_fastmath, metal_kq_mulmm_k6_t_msl_tgmem, 48, true)) metal_release(queue) metal_release(dev) return <- results From 096fab734ed3d60457e1c3607268d8e8c72aae6d Mon Sep 17 00:00:00 2001 From: Boris Batkin Date: Mon, 27 Jul 2026 22:45:16 -0700 Subject: [PATCH 12/18] =?UTF-8?q?dasLLAMA:=20mx4=20tuner=20race=20?= =?UTF-8?q?=E2=80=94=20all=20ten=20twin=20families=20raced?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit race_moe_mulmm_mx4 on the expert scaffolding: random nibble planes, near-1.0 e8m0 scales, the doubled-e2m1 LUT, bias off — twin-vs-base envelope proves the mx4 per-element decode agrees with the base kernel on arbitrary bit patterns. Ten families race on M1, all envelope-clean, all losing to their simdgroup kernels on fallback hardware (empty crown list — the pre-M5 expectation the tuner writes). Co-Authored-By: Claude Fable 5 --- .../dasllama/dasllama_metal_prefill.das | 136 ++++++++++++++++++ 1 file changed, 136 insertions(+) diff --git a/modules/dasLLAMA/dasllama/dasllama_metal_prefill.das b/modules/dasLLAMA/dasllama/dasllama_metal_prefill.das index 1f784712a1..62044e34c7 100644 --- a/modules/dasLLAMA/dasllama/dasllama_metal_prefill.das +++ b/modules/dasLLAMA/dasllama/dasllama_metal_prefill.das @@ -7671,6 +7671,141 @@ def private race_moe_mulmm_q8(dev : MetalDevice?; queue : MetalCommandQueue?) : return res } +def private race_moe_mulmm_mx4(dev : MetalDevice?; queue : MetalCommandQueue?) : MetalTensorRaceResult { + var res = MetalTensorRaceResult(family = "moe_mulmm_mx4", winner = "", base_ms = -1.0lf, twin_ms = -1.0lf, note = "") + let ne = 4 + let rows_per = 32 + let m = ne * rows_per + let kdim = 2048 + let ndim = 1024 + let nkb = kdim / 32 + let eblk = ndim * nkb + var err = "" + var base_pso = pipeline_from_source(dev, metal_moe_mulmm_mx4_msl, metal_moe_mulmm_mx4_msl_entry, metal_moe_mulmm_mx4_msl_fastmath, err) + if (base_pso == null) { + res.note = "base pso: {err}" + return res + } + var twin_pso = pipeline_from_source(dev, metal_moe_mulmm_mx4_t_msl, metal_moe_mulmm_mx4_t_msl_entry, metal_moe_mulmm_mx4_t_msl_fastmath, err) + if (twin_pso == null) { + res.winner = "simdgroup" + res.note = "twin pso: {err}" + metal_release(base_pso) + return res + } + var xa <- race_x_f32(m * kdim) + var nplane : array + nplane |> resize(ne * eblk * 4) // 16B nibbles per block + for (i in range(ne * eblk * 4)) { + nplane[i] = uint(i) * 2654435761u + } + var eplane : array + eplane |> resize(ne * eblk) + for (i in range(ne * eblk)) { + eplane[i] = uint8(125 + i % 4) // e8m0 near 1.0 — no overflow across the fold + } + var vtabh : array + vtabh |> resize(16) + for (i in range(8)) { + let v = float(fixed_array(0, 1, 2, 3, 4, 6, 8, 12)[i]) + vtabh[i] = v + vtabh[i + 8] = -v + } + var cnt : array + var basep : array + var bkt : array + cnt |> resize(ne) + basep |> resize(ne) + bkt |> resize(m) + for (e in range(ne)) { + cnt[e] = uint(rows_per) + basep[e] = uint(e * rows_per) + } + for (i in range(m)) { + bkt[i] = uint(i) + } + var bnq = race_buf(dev, uint64(ne * eblk * 16), unsafe(addr(nplane[0]))) + var bee = race_buf(dev, uint64(ne * eblk), unsafe(addr(eplane[0]))) + var bvt = race_buf(dev, uint64(16 * 4), unsafe(addr(vtabh[0]))) + var bwbb = race_buf(dev, uint64(ne * ndim * 4), null) + var bxa = race_buf(dev, uint64(m * kdim * 4), unsafe(addr(xa[0]))) + var by_base = race_buf(dev, uint64(m * ndim * 4), null) + var by_twin = race_buf(dev, uint64(m * ndim * 4), null) + var bcnt = race_buf(dev, uint64(ne * 4), unsafe(addr(cnt[0]))) + var bbase = race_buf(dev, uint64(ne * 4), unsafe(addr(basep[0]))) + var bbkt = race_buf(dev, uint64(m * 4), unsafe(addr(bkt[0]))) + var bk = race_uniform_u32(dev, uint(kdim)) + var bn = race_uniform_u32(dev, uint(ndim)) + var beblk = race_uniform_u32(dev, uint(eblk)) + var bnk = race_uniform_u32(dev, 1u) + var bg0 = race_uniform_u32(dev, 0u) + let grid = uint3(uint(rows_per / 32), uint(ndim / 64), uint(ne)) + let tg = uint3(128u, 1u, 1u) + res.base_ms = race_time_ms(queue, 5) $(enc) { + kn_pipeline(enc, base_pso) + kn_tgmem(enc, metal_moe_mulmm_mx4_msl_tgmem, 0) + kn_buffer(enc, bnq, 0ul, 0) + kn_buffer(enc, bee, 0ul, 1) + kn_buffer(enc, bxa, 0ul, 2) + kn_buffer(enc, by_base, 0ul, 3) + kn_buffer(enc, bk, 0ul, 4) + kn_buffer(enc, bn, 0ul, 5) + kn_buffer(enc, bcnt, 0ul, 6) + kn_buffer(enc, bbase, 0ul, 7) + kn_buffer(enc, bbkt, 0ul, 8) + kn_buffer(enc, beblk, 0ul, 9) + kn_buffer(enc, bnk, 0ul, 10) + kn_buffer(enc, bg0, 0ul, 11) + kn_buffer(enc, bwbb, 0ul, 12) + kn_buffer(enc, bvt, 0ul, 13) + kn_buffer(enc, bg0, 0ul, 14) + kn_dispatch(enc, grid, tg) + } + res.twin_ms = race_time_ms(queue, 5) $(enc) { + kn_pipeline(enc, twin_pso) + kn_tgmem(enc, metal_moe_mulmm_mx4_t_msl_tgmem, 0) + kn_buffer(enc, bnq, 0ul, 0) + kn_buffer(enc, bee, 0ul, 1) + kn_buffer(enc, bxa, 0ul, 2) + kn_buffer(enc, by_twin, 0ul, 3) + kn_buffer(enc, bk, 0ul, 4) + kn_buffer(enc, bn, 0ul, 5) + kn_buffer(enc, bcnt, 0ul, 6) + kn_buffer(enc, bbase, 0ul, 7) + kn_buffer(enc, beblk, 0ul, 9) + kn_buffer(enc, bwbb, 0ul, 12) + kn_buffer(enc, bvt, 0ul, 13) + kn_buffer(enc, bg0, 0ul, 14) + kn_dispatch(enc, grid, tg) + } + if (res.base_ms < 0.0lf || res.twin_ms < 0.0lf) { + res.winner = res.base_ms < 0.0lf ? "" : "simdgroup" + res.note = "dispatch failed" + } elif (!race_envelope_ok(by_base, by_twin, m * ndim, res.note)) { + res.winner = "simdgroup" + } else { + res.winner = res.twin_ms < res.base_ms ? "tensor" : "simdgroup" + } + metal_release(bnq) + metal_release(bee) + metal_release(bvt) + metal_release(bwbb) + metal_release(bxa) + metal_release(by_base) + metal_release(by_twin) + metal_release(bcnt) + metal_release(bbase) + metal_release(bbkt) + metal_release(bk) + metal_release(bn) + metal_release(beblk) + metal_release(bnk) + metal_release(bg0) + metal_release(base_pso) + metal_release(twin_pso) + return res +} + // one kq mulmm race: random quant planes + constant f16 scales (any bit pattern is a valid // plane — the twin-vs-base envelope doubles as a layout-agreement check). k6 binds the d plane // at its tail offset; k4/k5 bind the compact scale blocks at 0. @@ -7798,6 +7933,7 @@ def metal_tensor_race : array { results |> emplace(race_mulmm_bf16(dev, queue)) results |> emplace(race_mulmm_q8(dev, queue)) results |> emplace(race_moe_mulmm_q8(dev, queue)) + results |> emplace(race_moe_mulmm_mx4(dev, queue)) results |> emplace(race_kq_mulmm(dev, queue, "kq_mulmm_k4", metal_kq_mulmm_k4_msl, metal_kq_mulmm_k4_msl_entry, metal_kq_mulmm_k4_msl_fastmath, metal_kq_mulmm_k4_msl_tgmem, metal_kq_mulmm_k4_t_msl, metal_kq_mulmm_k4_t_msl_entry, metal_kq_mulmm_k4_t_msl_fastmath, metal_kq_mulmm_k4_t_msl_tgmem, 32, false)) From 183bbe21f43c735248c41c4776b3a55cea543754 Mon Sep 17 00:00:00 2001 From: Boris Batkin Date: Mon, 27 Jul 2026 22:52:13 -0700 Subject: [PATCH 13/18] =?UTF-8?q?dasLLAMA:=20attention=20tensor=20twins=20?= =?UTF-8?q?(AttnQKMm/AVMm)=20+=20their=20races=20=E2=80=94=20Phase=20E?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Both ride the staged-tile protocol with no new builtins: QK pre-scales Q at staging (the base's discipline) and copies its causal/window block skips (threadgroup-uniform, cooperative-safe); AV stages V transposed with pad positions >= npos zeroed and walks the same causal klimit. Crowns attn_qkmm/attn_avmm in enc_qk_mm/enc_av_mm. Races pre-zero BOTH output slabs (causal-skipped tiles stay unwritten in both kernels). Gate: llama prefill base+cont+dim+qkv with all TWELVE families force-crowned — token parity green; full race sweep envelope-clean. Co-Authored-By: Claude Fable 5 --- .../dasllama/dasllama_metal_prefill.das | 306 +++++++++++++++++- 1 file changed, 302 insertions(+), 4 deletions(-) diff --git a/modules/dasLLAMA/dasllama/dasllama_metal_prefill.das b/modules/dasLLAMA/dasllama/dasllama_metal_prefill.das index 62044e34c7..3e86c4fa2e 100644 --- a/modules/dasLLAMA/dasllama/dasllama_metal_prefill.das +++ b/modules/dasLLAMA/dasllama/dasllama_metal_prefill.das @@ -4096,6 +4096,62 @@ class MetalAttnAV { // The ggml-geometry attention pair (v21 lessons applied to the trio's QK/AV GEMMs — the old // 32x32 scalar-f32-staged kernels ran ~60x below mul_mm rate, owning 24ms of the 3B window). // QK: S[h] = scale * Q_h . K_h^T, 32x64 C tiles; DASLLAMA_METAL_ATTN=0 pins the trio (hs % 64 != 0 too). + +// Metal-4 staged-tile twin of MetalAttnQKMm: flat f16 tiles per head-dim chunk (Q pre-scaled +// at staging like the base), same causal/window block skips (threadgroup-uniform). +class MetalAttnQKMmT { + @ssbo @binding = 0 q : array + @ssbo @binding = 1 k : array + @ssbo @binding = 2 att : array + @uniform @binding = 3 qd : uint + @uniform @binding = 4 kv_dim : uint + @uniform @binding = 5 head_size : uint + @uniform @binding = 6 kv_mul : uint + @uniform @binding = 7 npos : uint + @uniform @binding = 8 np32 : uint + @uniform @binding = 9 scale : float + @uniform @binding = 10 qoff : uint + @uniform @binding = 11 qrows : uint + @uniform @binding = 12 window : uint + @workgroup twa : float16[1024] // Q chunk: 32 queries x 32 k, flat (pre-scaled) + @workgroup twb : float16[2048] // K chunk: 64 keys x 32 k, flat + + [metal_kernel(name="metal_attn_qk_mm_t_msl")] + def metal_attn_qk_mm_t { + let mBase = gl_WorkGroupID.x * 32u + let nBase = gl_WorkGroupID.y * 64u + if (nBase > qoff + mBase + 31u || + (window > 0u && qoff + mBase + 1u > window && nBase + 64u <= qoff + mBase + 1u - window)) { + return + } + let h = gl_WorkGroupID.z + let kvhoff = (h / kv_mul) * head_size + let qhoff = h * head_size + let lid = gl_LocalInvocationID.x + var acc : float[2048] + var cp = unsafe(addr(att[h * qrows * np32 + mBase * np32 + nBase])) + tmm2d_tg_begin(acc, 32u, 64u, 4u, 32u) + var kb = 0u + while (kb * 32u < head_size) { + var i = lid + while (i < 1024u) { + twa[i] = float16(q[(mBase + i / 32u) * qd + qhoff + kb * 32u + i % 32u] * scale) + i += gl_WorkGroupSize.x + } + i = lid + while (i < 2048u) { + twb[i] = float16(k[(nBase + i / 32u) * kv_dim + kvhoff + kb * 32u + i % 32u]) + i += gl_WorkGroupSize.x + } + barrier() + tmm2d_tg_step(acc, twa, twb, 32u, 64u, 32u) + barrier() + kb++ + } + tmm2d_tg_store(acc, cp, 32u, 64u, np32) + } +} + class MetalAttnQKMm { @ssbo @binding = 0 @role = "read" q : array // [mp x qd/4], roped @ssbo @binding = 1 @role = "read" k : array // [nk64 x kv_dim/4], roped (existing rows + the chunk) @@ -4209,6 +4265,57 @@ class MetalAttnQKMm { // AV: O_h = P_h . V_h — C tile 32 tokens x 64 V-columns, k walks the causally limited position // axis in 32-blocks. V stages contiguous float4 loads, position rows >= npos ZEROED (causal // softmax zeroes P columns >= npos EXACTLY; 0 * NaN from a pad V row would poison the tile). + +// Metal-4 staged-tile twin of MetalAttnAVMm: P x V over flat f16 tiles; V staged TRANSPOSED +// (col-major source) with pad positions >= npos zeroed (0 * stale-NaN would poison the tile). +class MetalAttnAVMmT { + @ssbo @binding = 0 att : array + @ssbo @binding = 1 v : array + @ssbo @binding = 2 xb : array + @uniform @binding = 3 qd : uint + @uniform @binding = 4 kv_dim : uint + @uniform @binding = 5 head_size : uint + @uniform @binding = 6 kv_mul : uint + @uniform @binding = 7 npos : uint + @uniform @binding = 8 np32 : uint + @uniform @binding = 9 qoff : uint + @uniform @binding = 10 qrows : uint + @workgroup twa : float16[1024] // P chunk: 32 tokens x 32 positions, flat + @workgroup twb : float16[2048] // V chunk: 64 v-cols x 32 positions, flat (transposed) + + [metal_kernel(name="metal_attn_av_mm_t_msl")] + def metal_attn_av_mm_t { + let mBase = gl_WorkGroupID.x * 32u + let vBase = gl_WorkGroupID.y * 64u + let h = gl_WorkGroupID.z + let kvhoff = (h / kv_mul) * head_size + let lid = gl_LocalInvocationID.x + let klimit = min(qoff + mBase + 32u, np32) // P rows here are zero past qoff+mBase+31 + var acc : float[2048] + var cp = unsafe(addr(xb[mBase * qd + h * head_size + vBase])) + tmm2d_tg_begin(acc, 32u, 64u, 4u, 32u) + var kb = 0u + while (kb * 32u < klimit) { + var i = lid + while (i < 1024u) { + twa[i] = float16(att[h * qrows * np32 + (mBase + i / 32u) * np32 + kb * 32u + i % 32u]) + i += gl_WorkGroupSize.x + } + i = lid + while (i < 2048u) { + let pos = kb * 32u + i % 32u + twb[i] = pos < npos ? float16(v[pos * kv_dim + kvhoff + vBase + i / 32u]) : float16(0.0) + i += gl_WorkGroupSize.x + } + barrier() + tmm2d_tg_step(acc, twa, twb, 32u, 64u, 32u) + barrier() + kb++ + } + tmm2d_tg_store(acc, cp, 32u, 64u, qd) + } +} + class MetalAttnAVMm { @ssbo @binding = 0 @role = "read" att : array // [n_heads x qrows x np32/4] normalized P @ssbo @binding = 1 @role = "read" v : array // [nk64 x kv_dim/4] raw V (existing rows + the chunk) @@ -4844,6 +4951,10 @@ var private g_pf_kq_mm5_tensor : bool // "kq_mulmm_k5" crowned var private g_pso_kq_mm5_t : MetalComputePipeline? var private g_pf_kq_mm6_tensor : bool // "kq_mulmm_k6" crowned var private g_pso_kq_mm6_t : MetalComputePipeline? +var private g_pf_qkmm_tensor : bool // "attn_qkmm" crowned +var private g_pso_qkmm_t : MetalComputePipeline? +var private g_pf_avmm_tensor : bool // "attn_avmm" crowned +var private g_pso_avmm_t : MetalComputePipeline? var private g_pf_pso_ple_gather : MetalComputePipeline? // PLE pre-step: q8 token-row gather var private g_pf_pso_ple_finish : MetalComputePipeline? // PLE pre-step: rms + combine, in place var private g_pso_qkmm : MetalComputePipeline? @@ -5156,6 +5267,14 @@ def public metal_prefill_shutdown { metal_release(g_pso_qkmm) g_pso_qkmm = null } + if (g_pso_qkmm_t != null) { + metal_release(g_pso_qkmm_t) + g_pso_qkmm_t = null + } + if (g_pso_avmm_t != null) { + metal_release(g_pso_avmm_t) + g_pso_avmm_t = null + } if (g_pso_avmm != null) { metal_release(g_pso_avmm) g_pso_avmm = null @@ -5254,6 +5373,14 @@ def private metal_prefill_init : bool { g_pf_pso_kq6 = pf_compile_pso(metal_kq_gemv_k6_msl, metal_kq_gemv_k6_msl_entry, metal_kq_gemv_k6_msl_fastmath, ok) g_pso_qkmm = pf_compile_pso(metal_attn_qk_mm_msl, metal_attn_qk_mm_msl_entry, metal_attn_qk_mm_msl_fastmath, ok) g_pso_avmm = pf_compile_pso(metal_attn_av_mm_msl, metal_attn_av_mm_msl_entry, metal_attn_av_mm_msl_fastmath, ok) + g_pf_qkmm_tensor = metal_tensor_crowned("attn_qkmm") + if (g_pf_qkmm_tensor) { + g_pso_qkmm_t = pf_compile_pso(metal_attn_qk_mm_t_msl, metal_attn_qk_mm_t_msl_entry, metal_attn_qk_mm_t_msl_fastmath, ok) + } + g_pf_avmm_tensor = metal_tensor_crowned("attn_avmm") + if (g_pf_avmm_tensor) { + g_pso_avmm_t = pf_compile_pso(metal_attn_av_mm_t_msl, metal_attn_av_mm_t_msl_entry, metal_attn_av_mm_t_msl_fastmath, ok) + } g_pf_pso_gemv = pf_compile_pso(metal_q8_gemv_msl, metal_q8_gemv_msl_entry, metal_q8_gemv_msl_fastmath, ok) g_pf_pso_moe_router = pf_compile_pso(metal_moe_router_msl, metal_moe_router_msl_entry, metal_moe_router_msl_fastmath, ok) g_pf_pso_moe_router_b = pf_compile_pso(metal_moe_router_b_msl, metal_moe_router_b_msl_entry, metal_moe_router_b_msl_fastmath, ok) @@ -5834,8 +5961,9 @@ def private enc_rope(enc : MetalComputeEncoder?; bv, bcos, bsin, bn, bhs, bnpair // The key grid spans nk64 (= np32's value — existing rows + the chunk on a continuation); // bqoff/bqrows carry the causal shift and the score slab's row extent. def private enc_qk_mm(enc : MetalComputeEncoder?; bq, bk, batt, bqd, bkvd, bhs, bkvm, bnp, bnp32, bscale, bqoff, bqrows, bwin : MetalBuffer?; mp, nk64, heads : int64) { - kn_pipeline(enc, g_pso_qkmm) - kn_tgmem(enc, metal_attn_qk_mm_msl_tgmem, 0) + let twinq = g_pf_qkmm_tensor && g_pso_qkmm_t != null + kn_pipeline(enc, twinq ? g_pso_qkmm_t : g_pso_qkmm) + kn_tgmem(enc, twinq ? metal_attn_qk_mm_t_msl_tgmem : metal_attn_qk_mm_msl_tgmem, 0) kn_buffer(enc, bq, 0ul, 0) kn_buffer(enc, bk, 0ul, 1) kn_buffer(enc, batt, 0ul, 2) @@ -5857,8 +5985,9 @@ def private enc_qk_mm(enc : MetalComputeEncoder?; bq, bk, batt, bqd, bkvd, bhs, } def private enc_av_mm(enc : MetalComputeEncoder?; batt, bv, bxb, bqd, bkvd, bhs, bkvm, bnp, bnp32, bqoff, bqrows : MetalBuffer?; mp, heads, head_size : int64) { - kn_pipeline(enc, g_pso_avmm) - kn_tgmem(enc, metal_attn_av_mm_msl_tgmem, 0) + let twina = g_pf_avmm_tensor && g_pso_avmm_t != null + kn_pipeline(enc, twina ? g_pso_avmm_t : g_pso_avmm) + kn_tgmem(enc, twina ? metal_attn_av_mm_t_msl_tgmem : metal_attn_av_mm_msl_tgmem, 0) kn_buffer(enc, batt, 0ul, 0) kn_buffer(enc, bv, 0ul, 1) kn_buffer(enc, bxb, 0ul, 2) @@ -7920,6 +8049,172 @@ def private race_kq_mulmm(dev : MetalDevice?; queue : MetalCommandQueue?; family return res } +// attention pair races: synthetic Q/K/V; BOTH att/xb outputs pre-zeroed (causal-skipped tiles +// stay unwritten in both kernels — fresh-buffer garbage would flake the envelope compare). +def private race_attn_pair(dev : MetalDevice?; queue : MetalCommandQueue?) : array { + var out : array + let heads = 8 + let hs = 128 + let qd = heads * hs + let kvd = qd + let mp = 128 + let np32 = 128 + var err = "" + var q_pso = pipeline_from_source(dev, metal_attn_qk_mm_msl, metal_attn_qk_mm_msl_entry, metal_attn_qk_mm_msl_fastmath, err) + var qt_pso = pipeline_from_source(dev, metal_attn_qk_mm_t_msl, metal_attn_qk_mm_t_msl_entry, metal_attn_qk_mm_t_msl_fastmath, err) + var a_pso = pipeline_from_source(dev, metal_attn_av_mm_msl, metal_attn_av_mm_msl_entry, metal_attn_av_mm_msl_fastmath, err) + var at_pso = pipeline_from_source(dev, metal_attn_av_mm_t_msl, metal_attn_av_mm_t_msl_entry, metal_attn_av_mm_t_msl_fastmath, err) + if (q_pso == null || qt_pso == null || a_pso == null || at_pso == null) { + out |> emplace(MetalTensorRaceResult(family = "attn_qkmm", winner = q_pso != null ? "simdgroup" : "", base_ms = -1.0lf, twin_ms = -1.0lf, note = "pso: {err}")) + out |> emplace(MetalTensorRaceResult(family = "attn_avmm", winner = a_pso != null ? "simdgroup" : "", base_ms = -1.0lf, twin_ms = -1.0lf, note = "pso: {err}")) + } else { + var qv <- race_x_f32(mp * qd) + var kv <- race_x_f32(np32 * kvd) + var pv <- race_x_f32(mp * np32 * heads) + var zeros : array + zeros |> resize(heads * mp * np32) + var bq = race_buf(dev, uint64(mp * qd * 4), unsafe(addr(qv[0]))) + var bkk = race_buf(dev, uint64(np32 * kvd * 4), unsafe(addr(kv[0]))) + var bp = race_buf(dev, uint64(heads * mp * np32 * 4), unsafe(addr(pv[0]))) + var batt_b = race_buf(dev, uint64(heads * mp * np32 * 4), unsafe(addr(zeros[0]))) + var batt_t = race_buf(dev, uint64(heads * mp * np32 * 4), unsafe(addr(zeros[0]))) + var bxb_b = race_buf(dev, uint64(mp * qd * 4), unsafe(addr(zeros[0]))) + var bxb_t = race_buf(dev, uint64(mp * qd * 4), unsafe(addr(zeros[0]))) + var uqd = race_uniform_u32(dev, uint(qd)) + var ukvd = race_uniform_u32(dev, uint(kvd)) + var uhs = race_uniform_u32(dev, uint(hs)) + var ukvm = race_uniform_u32(dev, 1u) + var unp = race_uniform_u32(dev, uint(np32)) + var unp32 = race_uniform_u32(dev, uint(np32)) + var uqoff = race_uniform_u32(dev, 0u) + var uqrows = race_uniform_u32(dev, uint(mp)) + var uwin = race_uniform_u32(dev, 0u) + let fscale = 0.125 + var uscale = race_buf(dev, 4ul, unsafe(addr(fscale))) + let qgrid = uint3(uint(mp / 32), uint(np32 / 64), uint(heads)) + let agrid = uint3(uint(mp / 32), uint(hs / 64), uint(heads)) + let tg = uint3(128u, 1u, 1u) + var res = MetalTensorRaceResult(family = "attn_qkmm", winner = "", base_ms = -1.0lf, twin_ms = -1.0lf, note = "") + res.base_ms = race_time_ms(queue, 5) $(enc) { + kn_pipeline(enc, q_pso) + kn_tgmem(enc, metal_attn_qk_mm_msl_tgmem, 0) + kn_buffer(enc, bq, 0ul, 0) + kn_buffer(enc, bkk, 0ul, 1) + kn_buffer(enc, batt_b, 0ul, 2) + kn_buffer(enc, uqd, 0ul, 3) + kn_buffer(enc, ukvd, 0ul, 4) + kn_buffer(enc, uhs, 0ul, 5) + kn_buffer(enc, ukvm, 0ul, 6) + kn_buffer(enc, unp, 0ul, 7) + kn_buffer(enc, unp32, 0ul, 8) + kn_buffer(enc, uscale, 0ul, 9) + kn_buffer(enc, uqoff, 0ul, 10) + kn_buffer(enc, uqrows, 0ul, 11) + kn_buffer(enc, uwin, 0ul, 12) + kn_dispatch(enc, qgrid, tg) + } + res.twin_ms = race_time_ms(queue, 5) $(enc) { + kn_pipeline(enc, qt_pso) + kn_tgmem(enc, metal_attn_qk_mm_t_msl_tgmem, 0) + kn_buffer(enc, bq, 0ul, 0) + kn_buffer(enc, bkk, 0ul, 1) + kn_buffer(enc, batt_t, 0ul, 2) + kn_buffer(enc, uqd, 0ul, 3) + kn_buffer(enc, ukvd, 0ul, 4) + kn_buffer(enc, uhs, 0ul, 5) + kn_buffer(enc, ukvm, 0ul, 6) + kn_buffer(enc, unp, 0ul, 7) + kn_buffer(enc, unp32, 0ul, 8) + kn_buffer(enc, uscale, 0ul, 9) + kn_buffer(enc, uqoff, 0ul, 10) + kn_buffer(enc, uqrows, 0ul, 11) + kn_buffer(enc, uwin, 0ul, 12) + kn_dispatch(enc, qgrid, tg) + } + if (res.base_ms < 0.0lf || res.twin_ms < 0.0lf) { + res.winner = res.base_ms < 0.0lf ? "" : "simdgroup" + res.note = "dispatch failed" + } elif (!race_envelope_ok(batt_b, batt_t, heads * mp * np32, res.note)) { + res.winner = "simdgroup" + } else { + res.winner = res.twin_ms < res.base_ms ? "tensor" : "simdgroup" + } + out |> emplace(res) + var res2 = MetalTensorRaceResult(family = "attn_avmm", winner = "", base_ms = -1.0lf, twin_ms = -1.0lf, note = "") + res2.base_ms = race_time_ms(queue, 5) $(enc) { + kn_pipeline(enc, a_pso) + kn_tgmem(enc, metal_attn_av_mm_msl_tgmem, 0) + kn_buffer(enc, bp, 0ul, 0) + kn_buffer(enc, bkk, 0ul, 1) + kn_buffer(enc, bxb_b, 0ul, 2) + kn_buffer(enc, uqd, 0ul, 3) + kn_buffer(enc, ukvd, 0ul, 4) + kn_buffer(enc, uhs, 0ul, 5) + kn_buffer(enc, ukvm, 0ul, 6) + kn_buffer(enc, unp, 0ul, 7) + kn_buffer(enc, unp32, 0ul, 8) + kn_buffer(enc, uqoff, 0ul, 9) + kn_buffer(enc, uqrows, 0ul, 10) + kn_dispatch(enc, agrid, tg) + } + res2.twin_ms = race_time_ms(queue, 5) $(enc) { + kn_pipeline(enc, at_pso) + kn_tgmem(enc, metal_attn_av_mm_t_msl_tgmem, 0) + kn_buffer(enc, bp, 0ul, 0) + kn_buffer(enc, bkk, 0ul, 1) + kn_buffer(enc, bxb_t, 0ul, 2) + kn_buffer(enc, uqd, 0ul, 3) + kn_buffer(enc, ukvd, 0ul, 4) + kn_buffer(enc, uhs, 0ul, 5) + kn_buffer(enc, ukvm, 0ul, 6) + kn_buffer(enc, unp, 0ul, 7) + kn_buffer(enc, unp32, 0ul, 8) + kn_buffer(enc, uqoff, 0ul, 9) + kn_buffer(enc, uqrows, 0ul, 10) + kn_dispatch(enc, agrid, tg) + } + if (res2.base_ms < 0.0lf || res2.twin_ms < 0.0lf) { + res2.winner = res2.base_ms < 0.0lf ? "" : "simdgroup" + res2.note = "dispatch failed" + } elif (!race_envelope_ok(bxb_b, bxb_t, mp * qd, res2.note)) { + res2.winner = "simdgroup" + } else { + res2.winner = res2.twin_ms < res2.base_ms ? "tensor" : "simdgroup" + } + out |> emplace(res2) + metal_release(bq) + metal_release(bkk) + metal_release(bp) + metal_release(batt_b) + metal_release(batt_t) + metal_release(bxb_b) + metal_release(bxb_t) + metal_release(uqd) + metal_release(ukvd) + metal_release(uhs) + metal_release(ukvm) + metal_release(unp) + metal_release(unp32) + metal_release(uqoff) + metal_release(uqrows) + metal_release(uwin) + metal_release(uscale) + } + if (q_pso != null) { + metal_release(q_pso) + } + if (qt_pso != null) { + metal_release(qt_pso) + } + if (a_pso != null) { + metal_release(a_pso) + } + if (at_pso != null) { + metal_release(at_pso) + } + return <- out +} + //! Tuner entry (the runtime.metal_tensor crown): race every tensor pso twin against its //! shipped simdgroup kernel on THIS box — compile + output-match to qualify, faster GPU time //! takes the family's crown. Empty when there is no Metal device. @@ -7934,6 +8229,9 @@ def metal_tensor_race : array { results |> emplace(race_mulmm_q8(dev, queue)) results |> emplace(race_moe_mulmm_q8(dev, queue)) results |> emplace(race_moe_mulmm_mx4(dev, queue)) + for (r in race_attn_pair(dev, queue)) { + results |> emplace(r) + } results |> emplace(race_kq_mulmm(dev, queue, "kq_mulmm_k4", metal_kq_mulmm_k4_msl, metal_kq_mulmm_k4_msl_entry, metal_kq_mulmm_k4_msl_fastmath, metal_kq_mulmm_k4_msl_tgmem, metal_kq_mulmm_k4_t_msl, metal_kq_mulmm_k4_t_msl_entry, metal_kq_mulmm_k4_t_msl_fastmath, metal_kq_mulmm_k4_t_msl_tgmem, 32, false)) From 2b61c1e5516c0237bc9aee46bc3c85cc974e8f83 Mon Sep 17 00:00:00 2001 From: Boris Batkin Date: Mon, 27 Jul 2026 23:28:50 -0700 Subject: [PATCH 14/18] dasLLAMA: tensor races interleave adjacent base/twin pairs MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Sequential per-side timing blocks read the box's thermal/clock ramp as a phantom winner — measured live on the m4: the standalone probe crowned kq k4/k5/k6 +64% and q8 +23% for tensor, the tuner's identical race in its own process context INVERTED both (k4 base 0.51ms vs the probe's 1.60ms — same binary, same box). race_pair_ms alternates one base and one twin dispatch per rep (the reds-dig2 adjacent-pairs doctrine), best-of per side; every family race site rewritten onto it. M1 verdicts unchanged (all simdgroup, stable margins). Co-Authored-By: Claude Fable 5 --- .../dasllama/dasllama_metal_common.das | 33 +++++++++++++ .../dasllama/dasllama_metal_kernels.das | 7 ++- .../dasllama/dasllama_metal_prefill.das | 49 ++++++++----------- 3 files changed, 57 insertions(+), 32 deletions(-) diff --git a/modules/dasLLAMA/dasllama/dasllama_metal_common.das b/modules/dasLLAMA/dasllama/dasllama_metal_common.das index 602af025ce..8be5d63a12 100644 --- a/modules/dasLLAMA/dasllama/dasllama_metal_common.das +++ b/modules/dasLLAMA/dasllama/dasllama_metal_common.das @@ -2097,6 +2097,39 @@ def race_envelope_ok(ba, bb : MetalBuffer?; count : int; var note : string&) : b return true } +// interleaved adjacent A/B pair — separate per-side blocks read the box's thermal/clock ramp +// as a phantom winner (reds-dig2; sequential blocks INVERTED m4's kq verdicts). Best-of/side. +def race_pair_ms(queue : MetalCommandQueue?; reps : int; var base_ms, twin_ms : double&; + base_blk : block<(enc : MetalComputeEncoder?) : void>; + twin_blk : block<(enc : MetalComputeEncoder?) : void>) { + base_ms = -1.0lf + twin_ms = -1.0lf + var err = "" + var best_b = -1.0lf + var best_t = -1.0lf + for (r in range(reps + 1)) { + var g1 = 0.0lf + var g2 = 0.0lf + let okb = with_compute_encoder_timed(queue, err, g1) $(enc) { + invoke(base_blk, enc) + } + let okt = okb && with_compute_encoder_timed(queue, err, g2) $(enc) { + invoke(twin_blk, enc) + } + return if (!okb || !okt) + if (r > 0) { + if (best_b < 0.0lf || g1 < best_b) { + best_b = g1 + } + if (best_t < 0.0lf || g2 < best_t) { + best_t = g2 + } + } + } + base_ms = best_b + twin_ms = best_t +} + // best GPU time over `reps` dispatches after one warmup; <0 on dispatch failure def race_time_ms(queue : MetalCommandQueue?; reps : int; blk : block<(enc : MetalComputeEncoder?) : void>) : double { var best = -1.0lf diff --git a/modules/dasLLAMA/dasllama/dasllama_metal_kernels.das b/modules/dasLLAMA/dasllama/dasllama_metal_kernels.das index c47099d136..997373f66b 100644 --- a/modules/dasLLAMA/dasllama/dasllama_metal_kernels.das +++ b/modules/dasLLAMA/dasllama/dasllama_metal_kernels.das @@ -7177,7 +7177,7 @@ def private race_gemmb_family(dev : MetalDevice?; queue : MetalCommandQueue?; fa var by_base = race_buf(dev, uint64(out_n * 4), null) var by_twin = race_buf(dev, uint64(out_n * 4), null) let tg = uint3(128u, 1u, 1u) - res.base_ms = race_time_ms(queue, 5) $(enc) { + race_pair_ms(queue, 5, res.base_ms, res.twin_ms, $(enc) { kn_pipeline(enc, base_pso) kn_tgmem(enc, base_tgmem, 0) kn_buffer(enc, bw, 0ul, 0) @@ -7185,15 +7185,14 @@ def private race_gemmb_family(dev : MetalDevice?; queue : MetalCommandQueue?; fa kn_buffer(enc, bxa, 0ul, 3) invoke(bind_extra, enc, by_base) kn_dispatch(enc, grid, tg) - } - res.twin_ms = race_time_ms(queue, 5) $(enc) { + }, $(enc) { kn_pipeline(enc, twin_pso) kn_buffer(enc, bw, 0ul, 0) kn_buffer(enc, bw, 0ul, 2) kn_buffer(enc, bxa, 0ul, 3) invoke(bind_extra, enc, by_twin) kn_dispatch(enc, grid, tg) - } + }) if (res.base_ms < 0.0lf || res.twin_ms < 0.0lf) { res.winner = res.base_ms < 0.0lf ? "" : "simdgroup" res.note = "dispatch failed" diff --git a/modules/dasLLAMA/dasllama/dasllama_metal_prefill.das b/modules/dasLLAMA/dasllama/dasllama_metal_prefill.das index 3e86c4fa2e..ff6ff1b4cc 100644 --- a/modules/dasLLAMA/dasllama/dasllama_metal_prefill.das +++ b/modules/dasLLAMA/dasllama/dasllama_metal_prefill.das @@ -7587,7 +7587,7 @@ def private race_mulmm_bf16(dev : MetalDevice?; queue : MetalCommandQueue?) : Me var bn = race_uniform_u32(dev, uint(ndim)) let grid = uint3(uint(m / 32), uint(ndim / 64), 1u) let tg = uint3(128u, 1u, 1u) - res.base_ms = race_time_ms(queue, 5) $(enc) { + race_pair_ms(queue, 5, res.base_ms, res.twin_ms, $(enc) { kn_pipeline(enc, base_pso) kn_tgmem(enc, metal_bf16_mulmm_msl_tgmem, 0) kn_buffer(enc, bwb, 0ul, 0) @@ -7596,8 +7596,7 @@ def private race_mulmm_bf16(dev : MetalDevice?; queue : MetalCommandQueue?) : Me kn_buffer(enc, bk, 0ul, 4) kn_buffer(enc, bn, 0ul, 5) kn_dispatch(enc, grid, tg) - } - res.twin_ms = race_time_ms(queue, 5) $(enc) { + }, $(enc) { kn_pipeline(enc, twin_pso) kn_buffer(enc, bwb, 0ul, 0) kn_buffer(enc, bxa, 0ul, 2) @@ -7605,7 +7604,7 @@ def private race_mulmm_bf16(dev : MetalDevice?; queue : MetalCommandQueue?) : Me kn_buffer(enc, bk, 0ul, 4) kn_buffer(enc, bn, 0ul, 5) kn_dispatch(enc, grid, tg) - } + }) if (res.base_ms < 0.0lf || res.twin_ms < 0.0lf) { res.winner = res.base_ms < 0.0lf ? "" : "simdgroup" res.note = "dispatch failed" @@ -7654,7 +7653,7 @@ def private race_mulmm_q8(dev : MetalDevice?; queue : MetalCommandQueue?) : Meta var bn = race_uniform_u32(dev, uint(ndim)) let grid = uint3(uint(m / 32), uint(ndim / 64), 1u) let tg = uint3(128u, 1u, 1u) - res.base_ms = race_time_ms(queue, 5) $(enc) { + race_pair_ms(queue, 5, res.base_ms, res.twin_ms, $(enc) { kn_pipeline(enc, base_pso) kn_tgmem(enc, metal_q8_mulmm_msl_tgmem, 0) kn_buffer(enc, bw, 0ul, 0) @@ -7664,8 +7663,7 @@ def private race_mulmm_q8(dev : MetalDevice?; queue : MetalCommandQueue?) : Meta kn_buffer(enc, bk, 0ul, 4) kn_buffer(enc, bn, 0ul, 5) kn_dispatch(enc, grid, tg) - } - res.twin_ms = race_time_ms(queue, 5) $(enc) { + }, $(enc) { kn_pipeline(enc, twin_pso) kn_buffer(enc, bw, 0ul, 0) kn_buffer(enc, bw, 0ul, 1) @@ -7674,7 +7672,7 @@ def private race_mulmm_q8(dev : MetalDevice?; queue : MetalCommandQueue?) : Meta kn_buffer(enc, bk, 0ul, 4) kn_buffer(enc, bn, 0ul, 5) kn_dispatch(enc, grid, tg) - } + }) if (res.base_ms < 0.0lf || res.twin_ms < 0.0lf) { res.winner = res.base_ms < 0.0lf ? "" : "simdgroup" res.note = "dispatch failed" @@ -7745,7 +7743,7 @@ def private race_moe_mulmm_q8(dev : MetalDevice?; queue : MetalCommandQueue?) : var bg0 = race_uniform_u32(dev, 0u) let grid = uint3(uint(rows_per / 32), uint(ndim / 64), uint(ne)) let tg = uint3(128u, 1u, 1u) - res.base_ms = race_time_ms(queue, 5) $(enc) { + race_pair_ms(queue, 5, res.base_ms, res.twin_ms, $(enc) { kn_pipeline(enc, base_pso) kn_tgmem(enc, metal_moe_mulmm_q8_msl_tgmem, 0) kn_buffer(enc, bw, 0ul, 0) @@ -7761,8 +7759,7 @@ def private race_moe_mulmm_q8(dev : MetalDevice?; queue : MetalCommandQueue?) : kn_buffer(enc, bnk, 0ul, 10) kn_buffer(enc, bg0, 0ul, 11) kn_dispatch(enc, grid, tg) - } - res.twin_ms = race_time_ms(queue, 5) $(enc) { + }, $(enc) { kn_pipeline(enc, twin_pso) kn_buffer(enc, bw, 0ul, 0) kn_buffer(enc, bw, 0ul, 1) @@ -7774,7 +7771,7 @@ def private race_moe_mulmm_q8(dev : MetalDevice?; queue : MetalCommandQueue?) : kn_buffer(enc, bbase, 0ul, 7) kn_buffer(enc, beblk, 0ul, 9) kn_dispatch(enc, grid, tg) - } + }) if (res.base_ms < 0.0lf || res.twin_ms < 0.0lf) { res.winner = res.base_ms < 0.0lf ? "" : "simdgroup" res.note = "dispatch failed" @@ -7870,7 +7867,7 @@ def private race_moe_mulmm_mx4(dev : MetalDevice?; queue : MetalCommandQueue?) : var bg0 = race_uniform_u32(dev, 0u) let grid = uint3(uint(rows_per / 32), uint(ndim / 64), uint(ne)) let tg = uint3(128u, 1u, 1u) - res.base_ms = race_time_ms(queue, 5) $(enc) { + race_pair_ms(queue, 5, res.base_ms, res.twin_ms, $(enc) { kn_pipeline(enc, base_pso) kn_tgmem(enc, metal_moe_mulmm_mx4_msl_tgmem, 0) kn_buffer(enc, bnq, 0ul, 0) @@ -7889,8 +7886,7 @@ def private race_moe_mulmm_mx4(dev : MetalDevice?; queue : MetalCommandQueue?) : kn_buffer(enc, bvt, 0ul, 13) kn_buffer(enc, bg0, 0ul, 14) kn_dispatch(enc, grid, tg) - } - res.twin_ms = race_time_ms(queue, 5) $(enc) { + }, $(enc) { kn_pipeline(enc, twin_pso) kn_tgmem(enc, metal_moe_mulmm_mx4_t_msl_tgmem, 0) kn_buffer(enc, bnq, 0ul, 0) @@ -7906,7 +7902,7 @@ def private race_moe_mulmm_mx4(dev : MetalDevice?; queue : MetalCommandQueue?) : kn_buffer(enc, bvt, 0ul, 13) kn_buffer(enc, bg0, 0ul, 14) kn_dispatch(enc, grid, tg) - } + }) if (res.base_ms < 0.0lf || res.twin_ms < 0.0lf) { res.winner = res.base_ms < 0.0lf ? "" : "simdgroup" res.note = "dispatch failed" @@ -8005,7 +8001,7 @@ def private race_kq_mulmm(dev : MetalDevice?; queue : MetalCommandQueue?; family let doff = k6 ? uint64(nblk * 16) : 0ul let grid = uint3(uint(m / 32), uint(ndim / 64), 1u) let tg = uint3(128u, 1u, 1u) - res.base_ms = race_time_ms(queue, 5) $(enc) { + race_pair_ms(queue, 5, res.base_ms, res.twin_ms, $(enc) { kn_pipeline(enc, base_pso) kn_tgmem(enc, base_tgmem, 0) kn_buffer(enc, bs, doff, 0) @@ -8016,8 +8012,7 @@ def private race_kq_mulmm(dev : MetalDevice?; queue : MetalCommandQueue?; family kn_buffer(enc, bk, 0ul, 5) kn_buffer(enc, bn, 0ul, 6) kn_dispatch(enc, grid, tg) - } - res.twin_ms = race_time_ms(queue, 5) $(enc) { + }, $(enc) { kn_pipeline(enc, twin_pso) kn_tgmem(enc, twin_tgmem, 0) kn_buffer(enc, bs, doff, 0) @@ -8028,7 +8023,7 @@ def private race_kq_mulmm(dev : MetalDevice?; queue : MetalCommandQueue?; family kn_buffer(enc, bk, 0ul, 5) kn_buffer(enc, bn, 0ul, 6) kn_dispatch(enc, grid, tg) - } + }) if (res.base_ms < 0.0lf || res.twin_ms < 0.0lf) { res.winner = res.base_ms < 0.0lf ? "" : "simdgroup" res.note = "dispatch failed" @@ -8095,7 +8090,7 @@ def private race_attn_pair(dev : MetalDevice?; queue : MetalCommandQueue?) : arr let agrid = uint3(uint(mp / 32), uint(hs / 64), uint(heads)) let tg = uint3(128u, 1u, 1u) var res = MetalTensorRaceResult(family = "attn_qkmm", winner = "", base_ms = -1.0lf, twin_ms = -1.0lf, note = "") - res.base_ms = race_time_ms(queue, 5) $(enc) { + race_pair_ms(queue, 5, res.base_ms, res.twin_ms, $(enc) { kn_pipeline(enc, q_pso) kn_tgmem(enc, metal_attn_qk_mm_msl_tgmem, 0) kn_buffer(enc, bq, 0ul, 0) @@ -8112,8 +8107,7 @@ def private race_attn_pair(dev : MetalDevice?; queue : MetalCommandQueue?) : arr kn_buffer(enc, uqrows, 0ul, 11) kn_buffer(enc, uwin, 0ul, 12) kn_dispatch(enc, qgrid, tg) - } - res.twin_ms = race_time_ms(queue, 5) $(enc) { + }, $(enc) { kn_pipeline(enc, qt_pso) kn_tgmem(enc, metal_attn_qk_mm_t_msl_tgmem, 0) kn_buffer(enc, bq, 0ul, 0) @@ -8130,7 +8124,7 @@ def private race_attn_pair(dev : MetalDevice?; queue : MetalCommandQueue?) : arr kn_buffer(enc, uqrows, 0ul, 11) kn_buffer(enc, uwin, 0ul, 12) kn_dispatch(enc, qgrid, tg) - } + }) if (res.base_ms < 0.0lf || res.twin_ms < 0.0lf) { res.winner = res.base_ms < 0.0lf ? "" : "simdgroup" res.note = "dispatch failed" @@ -8141,7 +8135,7 @@ def private race_attn_pair(dev : MetalDevice?; queue : MetalCommandQueue?) : arr } out |> emplace(res) var res2 = MetalTensorRaceResult(family = "attn_avmm", winner = "", base_ms = -1.0lf, twin_ms = -1.0lf, note = "") - res2.base_ms = race_time_ms(queue, 5) $(enc) { + race_pair_ms(queue, 5, res2.base_ms, res2.twin_ms, $(enc) { kn_pipeline(enc, a_pso) kn_tgmem(enc, metal_attn_av_mm_msl_tgmem, 0) kn_buffer(enc, bp, 0ul, 0) @@ -8156,8 +8150,7 @@ def private race_attn_pair(dev : MetalDevice?; queue : MetalCommandQueue?) : arr kn_buffer(enc, uqoff, 0ul, 9) kn_buffer(enc, uqrows, 0ul, 10) kn_dispatch(enc, agrid, tg) - } - res2.twin_ms = race_time_ms(queue, 5) $(enc) { + }, $(enc) { kn_pipeline(enc, at_pso) kn_tgmem(enc, metal_attn_av_mm_t_msl_tgmem, 0) kn_buffer(enc, bp, 0ul, 0) @@ -8172,7 +8165,7 @@ def private race_attn_pair(dev : MetalDevice?; queue : MetalCommandQueue?) : arr kn_buffer(enc, uqoff, 0ul, 9) kn_buffer(enc, uqrows, 0ul, 10) kn_dispatch(enc, agrid, tg) - } + }) if (res2.base_ms < 0.0lf || res2.twin_ms < 0.0lf) { res2.winner = res2.base_ms < 0.0lf ? "" : "simdgroup" res2.note = "dispatch failed" From b27b224d1807cc2ecfbaef69c260b8567fa42c8e Mon Sep 17 00:00:00 2001 From: Boris Batkin Date: Tue, 28 Jul 2026 00:21:32 -0700 Subject: [PATCH 15/18] =?UTF-8?q?llvm=5Ftune:=20sidecar=20box-identity=20g?= =?UTF-8?q?ate=20=E2=80=94=20foreign-box=20winners=20read=20as=20stale?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Measurements are a property of the box AND its OS (Metal fallback perf moves with macOS updates), so the sidecar's provenance gains a box identity stamp (platform|arch|hardware|OS-build|cpu) and the staleness check treats a mismatched or missing stamp exactly like an outdated mtime: reads as absent, load_sidecar_doc resets, the tuner regenerates. Closes the cross-box-copy hazard end to end — a foreign-seeded sidecar can no longer satisfy daspkg release completeness or ship wrong crowns, and the reset also drops the foreign file's other-scope keys instead of re-stamping them under a fresh identity. Manifest test gains the foreign-box arm (winners must not stamp; fallback tier engages). Co-Authored-By: Claude Fable 5 --- modules/dasLLVM/daslib/llvm_tune.das | 84 ++++++++++++++++++-- modules/dasLLVM/tests/llvm_tune_manifest.das | 31 ++++++++ 2 files changed, 109 insertions(+), 6 deletions(-) diff --git a/modules/dasLLVM/daslib/llvm_tune.das b/modules/dasLLVM/daslib/llvm_tune.das index afd4773113..17f6e06f2f 100644 --- a/modules/dasLLVM/daslib/llvm_tune.das +++ b/modules/dasLLVM/daslib/llvm_tune.das @@ -136,15 +136,87 @@ def tune_manifest_path() : string { return sidecar_for(empty(app) ? running_binary() : app) } - //! True when the sidecar predates the binary that runs (or builds) the app. A stale sidecar - //! reads as ABSENT everywhere and is regenerated, so measurements never outlive the binary - //! that made them — which is what kills the copy-a-stale-file trap. +var private g_box_identity = "" + + //! This box's tune identity: platform|arch|hardware|OS-build|cpu. Measurements are a + //! property of the box AND its OS (GPU fallback perf moves with macOS updates), so a + //! sidecar carrying a different identity reads as stale — the cross-box-copy hazard. +def tune_box_identity() : string { + if (!empty(g_box_identity)) { + return g_box_identity + } + var hw = "" + let plat = get_platform_name() + if (plat == "windows") { + hw = "{get_env_variable("PROCESSOR_IDENTIFIER")}|{get_env_variable("OS")}" + } elif (plat == "osx") { + unsafe { + popen("sysctl -n hw.model kern.osversion machdep.cpu.brand_string 2>/dev/null") $(f) { + while (!feof(f)) { + let ln = strip(fgets(f)) + if (!empty(ln)) { + hw = empty(hw) ? ln : "{hw}|{ln}" + } + } + } + } + } else { + unsafe { + popen("uname -r 2>/dev/null") $(f) { + hw = strip(fgets(f)) + } + } + fopen("/proc/cpuinfo", "rb") $(f) { + if (f != null) { + while (!feof(f)) { + let ln = fgets(f) + if (ln |> starts_with("model name")) { + let ix = find(ln, ":") + hw = "{hw}|{strip(slice(ln, ix + 1))}" + break + } + } + } + } + } + g_box_identity = "{plat}|{get_architecture_name()}|{hw}" + return g_box_identity +} + +// provenance.box of the sidecar at `path`, "" when absent/unreadable — cached per path+mtime +// (the staleness gate runs per [tuned] compile read; one parse per file per process is enough) +var private g_box_cache : table + +def private sidecar_box(path : string; mkey : string) : string { + let key = "{path}@{mkey}" + if (g_box_cache |> key_exists(key)) { + return g_box_cache[key] + } + var box = "" + let text = fread(path) + if (!empty(text)) { + var err = "" + var doc = read_json(text, err) + if (doc != null) { + box = doc?["provenance"]?["box"] ?? "" + delete_json(doc) + } + } + g_box_cache[key] = box + return box +} + + //! True when the sidecar predates the running binary OR carries a different box identity + //! (missing counts as different — pre-identity files re-tune once); a stale sidecar reads + //! as ABSENT everywhere, killing the copy-a-stale-file and copy-from-another-box traps. def tune_sidecar_stale(path : string) : bool { var sfs : FStat return false if (!stat(path, sfs)) // absent is "missing", not "stale" var bfs : FStat - return false if (!stat(running_binary(), bfs)) - return sfs.mtime < bfs.mtime + if (stat(running_binary(), bfs) && sfs.mtime < bfs.mtime) { + return true + } + return sidecar_box(path, "{sfs.mtime}") != tune_box_identity() } // the sidecar's "kernels" section as a flat { function name : perm suffix } map; @@ -211,7 +283,7 @@ def private load_sidecar_doc(path : string) : JsonValue? { // refresh "provenance" and write the document back def private save_sidecar_doc(path : string; var doc : JsonValue?) : bool { update(doc, "provenance", JV((binary = running_binary(), platform = get_platform_name(), - arch = get_architecture_name()))) + arch = get_architecture_name(), box = tune_box_identity()))) return fwrite(path, write_json(doc)) } diff --git a/modules/dasLLVM/tests/llvm_tune_manifest.das b/modules/dasLLVM/tests/llvm_tune_manifest.das index 51fc67820f..7ea9d83aa3 100644 --- a/modules/dasLLVM/tests/llvm_tune_manifest.das +++ b/modules/dasLLVM/tests/llvm_tune_manifest.das @@ -5,6 +5,8 @@ require daslib/module_path require daslib/strings_boost require daslib/fio require strings +require daslib/json +require daslib/json_boost // Per-app sidecar round-trip, declaration-free: copy the client template into a temp // directory (the sidecar resolves beside the ROOT SCRIPT — .tune.json), spawn it @@ -107,6 +109,35 @@ def test_llvm_tune_manifest_roundtrip(t : T?) { t |> success(!staleStampedK2) // stale winners must not stamp t |> success(staleResult6) // fallback k1 tier: 2+3+1 + // BOX-IDENTITY gate: a sidecar written on another box (foreign provenance.box) reads as + // stale even when fresher than the binary — the cross-box-copy hazard. Rewrite the file + // with the box swapped, touch it fresh, and the client must fall back again. + let sidecarText = fread(sidecarPath) + t |> success(find(sidecarText, "\"box\"") >= 0, "provenance carries the box identity stamp") + var jerr = "" + var jdoc = read_json(sidecarText, jerr) + t |> success(jdoc != null, "sidecar parses: {jerr}") + unsafe { + var prov = jdoc?["provenance"] + if (prov != null) { + update(prov, "box", JV("some-other-box|foreign-arch|foreign-hw")) + } + } + fwrite(sidecarPath, write_json(jdoc)) + delete_json(jdoc) + var flines : array + let rcf = spawn_client("\"{bin}\" -jit \"{clientPath}\"", flines) + t |> equal(rcf, 0) + var foreignStampedK2 = false + var foreignResult6 = false + for (ln in flines) { + foreignStampedK2 ||= ln |> starts_with("llvm_tune: man_add <- k2") + foreignResult6 ||= ln == "RESULT 6" + } + t |> success(!foreignStampedK2) // foreign-box winners must not stamp + t |> success(foreignResult6) // fallback tier again + delete flines + remove(sidecarPath) remove(clientPath) delete wlines From 3643ebdae851a8566679f75774683cab18d8d979 Mon Sep 17 00:00:00 2001 From: Boris Batkin Date: Tue, 28 Jul 2026 00:27:37 -0700 Subject: [PATCH 16/18] dasLLAMA: tmm2d_c_arg suffix-matches like the tg-protocol arms (Copilot) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The direct tmm2d builtins are non-generic so their resolved names stay bare today, but the exact/suffix split in one classifier was a latent under-tracking path — one rule now. Co-Authored-By: Claude Fable 5 --- modules/dasLLAMA/dasllama/dasllama_kernel_access.das | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/modules/dasLLAMA/dasllama/dasllama_kernel_access.das b/modules/dasLLAMA/dasllama/dasllama_kernel_access.das index cbda3406db..22d961ac37 100644 --- a/modules/dasLLAMA/dasllama/dasllama_kernel_access.das +++ b/modules/dasLLAMA/dasllama/dasllama_kernel_access.das @@ -85,12 +85,13 @@ def private is_setop2(op : das_string) : bool { def private is_setop1(op : das_string) : bool => op == "++" || op == "--" || op == "+++" || op == "---" // tmm2d tensor builtins (dasMetal): pointer operands — C is written, the rest are read. -// Returns the C-argument index, or -1 for any other call. +// Returns the C-argument index, or -1 for any other call. Suffix match, like the tg-protocol +// arms: resolved call names may carry a module/instance prefix. def private tmm2d_c_arg(cname : string) : int { - if (cname == "tmm2d_f32_bf16_f32" || cname == "tmm2d_q8b_f32") { + if (cname |> ends_with("tmm2d_f32_bf16_f32") || cname |> ends_with("tmm2d_q8b_f32")) { return 7 } - return cname == "tmm2d_q8_f32" || cname == "tmm2d_q8_f16s" ? 11 : -1 + return cname |> ends_with("tmm2d_q8_f32") || cname |> ends_with("tmm2d_q8_f16s") ? 11 : -1 } class private AccessVisitor : AstVisitor { From 2ef1213a55b0395305b11b6c9d1149531d5dca7c Mon Sep 17 00:00:00 2001 From: Boris Batkin Date: Tue, 28 Jul 2026 01:01:14 -0700 Subject: [PATCH 17/18] tests/metal: gate the tensor-ops GPU half on Metal-4 availability MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The darwin15 CI runner has a Metal device but no MPP / MSL 4.0, so the pso-compile assertions tripped. A trivial MPP-include probe now feints the GPU half where Metal 4 is absent — the production races already degrade the same way (compile-probe eligibility); on a Metal-4 box the real kernels still assert loudly. CPU-replay half runs everywhere. Co-Authored-By: Claude Fable 5 --- tests/metal/test_metal_tensor_ops.das | 18 ++++++++++++++++++ 1 file changed, 18 insertions(+) diff --git a/tests/metal/test_metal_tensor_ops.das b/tests/metal/test_metal_tensor_ops.das index dcb141e43e..a0786d44c4 100644 --- a/tests/metal/test_metal_tensor_ops.das +++ b/tests/metal/test_metal_tensor_ops.das @@ -100,6 +100,20 @@ class TmmKernels { } } +// Metal-4 availability probe: the tensor ops need the MPP framework + MSL 4.0 (macOS 26) — +// absent on older runners, where the GPU half must skip (a compile failure there is not a +// regression; on a Metal-4 box the real kernels still assert loudly). +def private metal4_available(dev : MetalDevice?) : bool { + var err = "" + let src = "#include \n#include \nusing namespace metal;\nkernel void probe() \{\}\n" + var lib = metal_new_library_from_source(dev, src, false, err) + if (lib == null) { + return false + } + metal_release(lib) + return true +} + def private fill_oracle(var o : TmmKernels?) { o.xa |> resize(M * KK) o.wb |> resize(N * KK) @@ -182,6 +196,10 @@ def test_tensor_ops(t : T?) { feint("no Metal device on this box; GPU compare skipped\n") return } + if (!metal4_available(dev)) { + feint("Metal-4 tensor ops unavailable (needs macOS 26 / MSL 4.0); GPU compare skipped\n") + return + } var queue = metal_new_command_queue(dev) var bxa = buf_upload(dev, oracle.xa) var bwb = buf_upload(dev, oracle.wb) From 62583a720c0eb3287589fb98221aff28aa7a5525 Mon Sep 17 00:00:00 2001 From: Boris Batkin Date: Tue, 28 Jul 2026 01:31:28 -0700 Subject: [PATCH 18/18] tests/metal: metal4_available takes its device untyped (off-Apple compile) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The helper named MetalDevice at top level — a type that only exists when das_metal is built, breaking the Windows/Linux compile of the test. An untyped parameter makes it a generic instantiated only inside the Apple static_if half, matching the rest of the file's guard discipline. Co-Authored-By: Claude Fable 5 --- tests/metal/test_metal_tensor_ops.das | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/metal/test_metal_tensor_ops.das b/tests/metal/test_metal_tensor_ops.das index a0786d44c4..25f2c473f4 100644 --- a/tests/metal/test_metal_tensor_ops.das +++ b/tests/metal/test_metal_tensor_ops.das @@ -103,7 +103,7 @@ class TmmKernels { // Metal-4 availability probe: the tensor ops need the MPP framework + MSL 4.0 (macOS 26) — // absent on older runners, where the GPU half must skip (a compile failure there is not a // regression; on a Metal-4 box the real kernels still assert loudly). -def private metal4_available(dev : MetalDevice?) : bool { +def private metal4_available(dev) : bool { // untyped: instantiated only inside the Apple half var err = "" let src = "#include \n#include \nusing namespace metal;\nkernel void probe() \{\}\n" var lib = metal_new_library_from_source(dev, src, false, err)