593 lines
36 KiB
Python
593 lines
36 KiB
Python
"""Extract pinned MTPLX Metal bodies without executing its Python model code.
|
|
|
|
Print the deterministic Metal translation unit; --check compares the checked-in
|
|
copy byte for byte. Only entry-point ABI/template declarations are supplied here.
|
|
The bodies, including whitespace and comments, come unchanged from MTPLX.
|
|
"""
|
|
|
|
import argparse
|
|
import ast
|
|
import hashlib
|
|
import json
|
|
import re
|
|
from pathlib import Path
|
|
import subprocess
|
|
|
|
REVISION = "e652d55e2652137a4abcf1312357abbf3eb9d692"
|
|
GATED_DELTA_SHA256 = "79c8376a51c694b03e54d2f996ced6ea6c8c42868b8571529f97334db165a3e1"
|
|
OUTPUT = Path(__file__).resolve().parents[1] / "metal/mtplx_qwen.metal"
|
|
QSA_SELECT_OUTPUT = OUTPUT.with_name("mtplx-qsa-select.json")
|
|
RUNTIME_REVISION = "1f8e74e3f12f31365464a6867c6579f0e9b29d85"
|
|
|
|
# (module, source variable, entry point, inputs, outputs, builtins, extra template,
|
|
# (entry-point suffix, complete template arguments) specializations)
|
|
KERNELS = [
|
|
("hyper_connection", "_SOURCE", "hyper_read",
|
|
"T:x,gamma,wd,wu,wi", "T:mixed,inject",
|
|
"thread_position_in_grid,thread_position_in_threadgroup", "int HAS_INJECT",
|
|
[("mix_bf16", "bfloat, 0"), ("inj_bf16", "bfloat, 1")]),
|
|
("hyper_connection_v3", "_SRC_R1", "hyper_v3_r1",
|
|
"T:x,wn;uint32_t:qw;T:qs,qb", "T:mix_out,inject_out;float:rms_out",
|
|
"thread_position_in_threadgroup,threadgroup_position_in_grid", "",
|
|
[("bf16", "bfloat")]),
|
|
("hyper_connection_v3", "_SRC_R2", "hyper_v3_r2",
|
|
"T:x,wn,mixv;float:rms_in;uint32_t:qw;T:qs,qb", "T:y",
|
|
"thread_position_in_threadgroup,threadgroup_position_in_grid", "",
|
|
[("bf16", "bfloat")]),
|
|
("gdn_conv_norm", "_SRC", "gdn_conv_norm",
|
|
"T:xnew,state,cw", "T:q_out,k_out,v_out,state_out",
|
|
"thread_position_in_threadgroup,threadgroup_position_in_grid", "",
|
|
[("bf16", "bfloat")]),
|
|
("gdn_conv_norm", "_SRC_ROWS", "gdn_conv_norm_rows",
|
|
"T:xnew,state,cw", "T:q_out,k_out,v_out,state_out",
|
|
"thread_position_in_threadgroup,threadgroup_position_in_grid", "int S",
|
|
[(f"s{s}_bf16", f"bfloat, {s}") for s in range(2, 7)]),
|
|
("gdn_step_fused", "_SRC", "gdn_step_fused",
|
|
"T:xnew,z_row,a_row,b_row,state,cw,A_log,dt_bias;StT:dstate;T:norm_w",
|
|
"T:y,state_out;StT:dstate_out",
|
|
"thread_position_in_threadgroup,threadgroup_position_in_grid", "typename StT",
|
|
[("f32_state_bf16", "bfloat, float")]),
|
|
("gdn_out_fused", "_SRC", "gdn_out_fused",
|
|
"T:x;bfloat:z,wn;uint32_t:qw;bfloat:qs,qb", "T:y",
|
|
"thread_position_in_threadgroup,threadgroup_position_in_grid", "int GS_C",
|
|
[(f"gs{gs}_{tag}", f"{dtype}, {gs}")
|
|
for gs in (32, 64) for tag, dtype in (("bf16", "bfloat"), ("f32", "float"))]),
|
|
("moe_glu_decode", "_SRC_A", "moe_glu_h",
|
|
"T:x;uint32_t:gw;T:gs,gb;uint32_t:experts", "T:h",
|
|
"thread_position_in_threadgroup,threadgroup_position_in_grid",
|
|
"int n_inter_c, int topk_c, int GS_GU",
|
|
[(f"g{gs}_bf16", f"bfloat, 640, 10, {gs}") for gs in (32,64)]),
|
|
("moe_glu_decode", "_SRC_B", "moe_down_y",
|
|
"T:h;uint32_t:dw;T:ds,db;uint32_t:experts;float:rw", "T:y",
|
|
"thread_position_in_threadgroup,threadgroup_position_in_grid",
|
|
"int dmodel_c, int n_inter_c, int topk_c, int GS_DN",
|
|
[(f"g{gs}_bf16", f"bfloat, 2560, 640, 10, {gs}") for gs in (32,64)]),
|
|
("moe_glu_decode", "_SRC_A_M", "moe_glu_h_m",
|
|
"T:x;uint32_t:gw;T:gs,gb;uint32_t:experts", "T:h",
|
|
"thread_position_in_threadgroup,threadgroup_position_in_grid",
|
|
"int n_inter_c, int topk_c, int m_rows_c, int GS_GU",
|
|
[(f"m{m}_g{gs}_bf16", f"bfloat, 640, 10, {m}, {gs}") for m in (2,3,4) for gs in (32,64)]),
|
|
("moe_glu_decode", "_SRC_B_M", "moe_down_y_m",
|
|
"T:h;uint32_t:dw;T:ds,db;uint32_t:experts;float:rw", "T:y",
|
|
"thread_position_in_threadgroup,threadgroup_position_in_grid",
|
|
"int dmodel_c, int n_inter_c, int topk_c, int m_rows_c, int GS_DN",
|
|
[(f"m{m}_g{gs}_bf16", f"bfloat, 2560, 640, 10, {m}, {gs}") for m in (2,3,4) for gs in (32,64)]),
|
|
]
|
|
|
|
|
|
def parameters(groups, output=False):
|
|
return [(f"device {'const ' if not output else ''}{kind.strip()}*", name)
|
|
for group in groups.split(";") for kind, names in [group.split(":")]
|
|
for name in names.split(",")]
|
|
|
|
|
|
def kernel_source(source, raw, body, function, args, builtins, templates, variants):
|
|
signature = [f" {kind} {key} [[buffer({i})]]" for i, (kind, key) in enumerate(args)]
|
|
signature += [f" {'uint' if key == 'thread_index_in_simdgroup' else 'uint3'} "
|
|
f"{key} [[{key}]]" for key in builtins.split(",")]
|
|
output = [f"\n// Source: {source}\n"
|
|
f"// File SHA256: {hashlib.sha256(raw).hexdigest()}\n"
|
|
f"// Body SHA256: {hashlib.sha256(body.encode()).hexdigest()}\n"
|
|
f"template <{templates}>\n"
|
|
f"kernel void {function}(\n" + ",\n".join(signature) + ") {" + body + "}\n"]
|
|
for suffix, argument in variants:
|
|
specialization = f"{function}<{argument}>"
|
|
alias = f"{function}_{suffix}"
|
|
output.append(f"typedef decltype({specialization}) {alias}_type;\n"
|
|
f'template [[host_name("{alias}")]]\n'
|
|
f"kernel {alias}_type {specialization};\n")
|
|
return "".join(output)
|
|
|
|
|
|
def gated_delta_source(path):
|
|
raw = path.read_bytes()
|
|
if hashlib.sha256(raw).hexdigest() != GATED_DELTA_SHA256:
|
|
raise ValueError("Expected the pinned mlx-lm 0.31.3 gated_delta.py used by MTPLX")
|
|
function = next(n for n in ast.parse(raw).body if isinstance(n, ast.FunctionDef)
|
|
and n.name == "_make_gated_delta_kernel")
|
|
output = ["\n// mlx-lm 0.31.3 dependency kernel: Copyright 2023 Apple Inc.\n"
|
|
"// SPDX-License-Identifier: MIT; see MLX-LM-LICENSE.txt.\n"]
|
|
# Qwen compute_g is scalar per head; the masked branch also covers ragged rows.
|
|
for masked in (False, True):
|
|
values = {"has_mask": masked, "vectorized": False}
|
|
|
|
def literal(node):
|
|
if isinstance(node, ast.Constant):
|
|
return node.value
|
|
if isinstance(node, ast.Name):
|
|
return values[node.id]
|
|
if isinstance(node, ast.IfExp):
|
|
return literal(node.body if literal(node.test) else node.orelse)
|
|
if isinstance(node, ast.JoinedStr):
|
|
return "".join(literal(n) for n in node.values)
|
|
if isinstance(node, ast.FormattedValue) and node.format_spec is None and node.conversion == -1:
|
|
return str(literal(node.value))
|
|
raise ValueError(f"Unexpected source expression: {ast.dump(node)}")
|
|
|
|
for statement in function.body:
|
|
if isinstance(statement, ast.If) and isinstance(statement.test, ast.Name):
|
|
for assignment in statement.body if values[statement.test.id] else statement.orelse:
|
|
values[assignment.targets[0].id] = literal(assignment.value)
|
|
elif isinstance(statement, ast.Assign):
|
|
name = statement.targets[0].id
|
|
values[name] = literal(statement.value)
|
|
if name == "source":
|
|
break
|
|
args = parameters("InT:q,k,v;float:g;InT:beta;StT:state_in")
|
|
args += [("constant const int32_t&", "T")]
|
|
if masked:
|
|
args += parameters("bool:mask")
|
|
args += parameters("InT:y;StT:state_out", output=True)
|
|
suffix = "_masked" if masked else ""
|
|
output.append(kernel_source(
|
|
"mlx_lm/models/gated_delta.py::_make_gated_delta_kernel" + suffix,
|
|
raw, values["source"], "kernel_qwen_mtplx_gated_delta" + suffix, args,
|
|
"thread_index_in_simdgroup,thread_position_in_grid,thread_position_in_threadgroup",
|
|
"typename InT, typename StT, int Dk, int Dv, int Hk, int Hv",
|
|
[("bf16", "bfloat, float, 128, 128, 16, 48")]))
|
|
return "".join(output)
|
|
|
|
|
|
def gather_front_source(runtime):
|
|
revision = subprocess.check_output(["git","-C",str(runtime),"rev-parse","HEAD"],text=True).strip()
|
|
if revision != RUNTIME_REVISION:
|
|
raise ValueError(f"Expected MTPLX runtime {RUNTIME_REVISION}, got {revision}")
|
|
output = ["\n// MTPLX runtime JIT shaders: Copyright Apple Inc. SPDX-License-Identifier: MIT\n"
|
|
"// See MLX-LM-LICENSE.txt; only includes are resolved and instantiations supplied.\n"]
|
|
for name, digest in (
|
|
("indexing.h","e820b8ee2b5132a97122780c12433ebb5100d8078d31e211d0429400a11415bb"),
|
|
("gather_front.h","64aacebf6576dfcd389383564fa1214bc87f2a091dd33cc64c598c5367ecab96")):
|
|
raw = (runtime/"mlx/backend/metal/kernels/indexing"/name).read_bytes()
|
|
if hashlib.sha256(raw).hexdigest() != digest:
|
|
raise ValueError(f"Pinned runtime shader changed: {name}")
|
|
body = "".join(line for line in raw.decode().splitlines(keepends=True)
|
|
if not line.startswith(("#include", "#pragma once")))
|
|
output.append(f"// Runtime unit: {name}; file SHA256: {digest}\n"
|
|
f"// Runtime unit SHA256: {hashlib.sha256(body.encode()).hexdigest()}\n"
|
|
"// BEGIN RUNTIME UNIT\n"+body+"// END RUNTIME UNIT\n")
|
|
for tag, kind, n in (("bf16","bfloat",1),("bf16","bfloat",2),("f16","half",1),("f16","half",2),("u32","uint32_t",1),("f32","float",1)):
|
|
for loc in ("int","int64_t"):
|
|
specialization = f"gather_front<{kind}, uint32_t, {loc}, {n}>"
|
|
output.append(f'template [[host_name("kernel_qwen_mtplx_gather_front_{tag}_{loc}_{n}")]]\n'
|
|
f"[[kernel]] decltype({specialization}) {specialization};\n")
|
|
for tag,kind,n in (("bf16","bfloat",1),("bf16","bfloat",2),("f16","half",1),("f16","half",2),("u32","uint32_t",1),("f32","float",1)):
|
|
for loc in ("int","int64_t"):
|
|
specialization=f"gather_front<{kind}, int64_t, {loc}, {n}>"
|
|
output.append(f'template [[host_name("kernel_qwen_mtplx_gather_front_{tag}_idxi64_{loc}_{n}")]]\n'
|
|
f"[[kernel]] decltype({specialization}) {specialization};\n")
|
|
for tag,kind,n in (("bf16","bfloat",1),("bf16","bfloat",2),("f16","half",1),("f16","half",2),("u32","uint32_t",1),("f32","float",1)):
|
|
for loc in ("int","int64_t"):
|
|
specialization=f"gather_front<{kind}, int32_t, {loc}, {n}>"
|
|
output.append(f'template [[host_name("kernel_qwen_mtplx_gather_front_{tag}_idxi32_{loc}_{n}")]]\n'
|
|
f"[[kernel]] decltype({specialization}) {specialization};\n")
|
|
return "".join(output)
|
|
|
|
|
|
def gather_rows_source(runtime):
|
|
output=[]
|
|
path=runtime/"mlx/backend/metal/kernels/indexing/gather.h"
|
|
raw=path.read_bytes()
|
|
digest="3b2f5b21cd2e71427c9641368a457840f5cb131e202095cd1352f983ce7541ce"
|
|
if hashlib.sha256(raw).hexdigest()!=digest: raise ValueError("Pinned gather.h changed")
|
|
body="".join(line for line in raw.decode().splitlines(keepends=True) if not line.startswith(("#include","#pragma once")))
|
|
output.append(f"// Runtime unit: gather.h; file SHA256: {digest}\n"
|
|
f"// Runtime unit SHA256: {hashlib.sha256(body.encode()).hexdigest()}\n"
|
|
"// BEGIN RUNTIME UNIT\n"+body+"// END RUNTIME UNIT\n")
|
|
raw=(runtime/"mlx/backend/metal/jit/indexing.h").read_bytes()
|
|
digest="1b38dbdf3120eca3e5266ffd6d69691591c9432c0b5acece4c5ab2de7a9e802f"
|
|
if hashlib.sha256(raw).hexdigest()!=digest: raise ValueError("Pinned gather JIT wrapper changed")
|
|
template=raw.decode().split('gather_kernels = R"(',1)[1].split(')";',1)[0]
|
|
wrappers=[]
|
|
for tag,kind in (("bf16","bfloat"),("f16","half"),("u32","uint32_t"),("f32","float")):
|
|
for rank in range(5):
|
|
for loc in ("int","int64_t"):
|
|
body=template.format(f"{tag}idxi64",kind,"int64_t",1,
|
|
"const device int64_t *idx0 [[buffer(20)]],","idx0",rank,loc)
|
|
body=body.replace(f"void gather{tag}idxi64_",f"void kernel_qwen_mtplx_gather_{tag}_idxi64_",1)
|
|
wrappers.append(body)
|
|
for tag,kind in (("bf16","bfloat"),("f16","half"),("u32","uint32_t"),("f32","float")):
|
|
for rank in range(5):
|
|
for loc in ("int","int64_t"):
|
|
body=template.format(f"{tag}idxi32",kind,"int32_t",1,
|
|
"const device int32_t *idx0 [[buffer(20)]],","idx0",rank,loc)
|
|
body=body.replace(f"void gather{tag}idxi32_",f"void kernel_qwen_mtplx_gather_{tag}_idxi32_",1)
|
|
wrappers.append(body)
|
|
for tag,kind in (("bf16","bfloat"),("f16","half"),("u32","uint32_t"),("f32","float")):
|
|
for rank in range(5):
|
|
for loc in ("int","int64_t"):
|
|
body=template.format(f"{tag}idxu32",kind,"uint32_t",1,
|
|
"const device uint32_t *idx0 [[buffer(20)]],","idx0",rank,loc)
|
|
body=body.replace(f"void gather{tag}idxu32_",f"void kernel_qwen_mtplx_gather_{tag}_idxu32_",1)
|
|
wrappers.append(body)
|
|
body="".join(wrappers)
|
|
output.append(f"// Runtime unit: gather JIT wrappers; file SHA256: {digest}\n"
|
|
f"// Runtime unit SHA256: {hashlib.sha256(body.encode()).hexdigest()}\n"
|
|
"// BEGIN RUNTIME UNIT\n"+body+"// END RUNTIME UNIT\n")
|
|
return "".join(output)
|
|
|
|
|
|
def silu_source(runtime):
|
|
output = ["\n// Actual runtime-generated SiLU, captured by tools/mtplx-jit-reference.py.\n"
|
|
"// Copyright Apple Inc. SPDX-License-Identifier: MIT. Only host aliases change.\n"]
|
|
def unit(name, body, digest):
|
|
output.append(f"// Runtime unit: {name}; file SHA256: {digest}\n"
|
|
f"// Runtime unit SHA256: {hashlib.sha256(body.encode()).hexdigest()}\n"
|
|
"// BEGIN RUNTIME UNIT\n"+body+"// END RUNTIME UNIT\n")
|
|
for name,digest,sections in (
|
|
("bf16.h","abd87446a310b77ac530ef52a324feae5cb285d03ec9613e3a88ebb71410fdcb",None),
|
|
("bf16_math.h","1f374f8380f756eb89acf6a847741cb8fecbe642945e159fb6208d804cc06496",None),
|
|
("utils.h","5e1568e9edde9d05dbf86f68fa0d6c6240f2c32b973c7c6a76166b9c0d91543d",
|
|
[("template <typename IdxT = int64_t>\nMETAL_FUNC IdxT elem_to_loc_1", "///////////////////////////////////////////////////////////////////////////////"),
|
|
("template <typename U, typename T>\ninline U cast_to", "template <>")]),
|
|
("unary_ops.h","0a5492b65ae39ecb6d8b04e64ea5007e0a4ff60d8d9428559bfbfd03387ece2a",[("struct Sigmoid {","struct Sign {")]),
|
|
("binary_ops.h","2dd13c2496f5d6f0856e4ca99db2ceb5c6d7dc0c344b9ba8e41ef3b7d7ed8a97",[("struct Multiply {","struct NotEqual {")]),
|
|
):
|
|
raw = (runtime/"mlx/backend/metal/kernels"/name).read_bytes()
|
|
if hashlib.sha256(raw).hexdigest()!=digest:
|
|
raise ValueError(f"Pinned SiLU dependency changed: {name}")
|
|
text = raw.decode()
|
|
if sections:
|
|
body = "".join(text[text.index(start):text.index(end,text.index(start))] for start,end in sections)
|
|
else:
|
|
body = "".join(line for line in text.splitlines(keepends=True) if not line.startswith(("#include","#pragma once")))
|
|
unit(name,body,digest)
|
|
raw = (OUTPUT.parents[1]/"tests/fixtures/mtplx-silu-jit.json").read_bytes()
|
|
digest = hashlib.sha256(raw).hexdigest()
|
|
if digest!="ed0ccd2cbbbcdacbb3e1c1fe8eb7f6b4d104a00197c97a43c93904fc0124d0af":
|
|
raise ValueError("Pinned SiLU JIT receipt changed")
|
|
receipt = json.loads(raw)
|
|
prefix = "BV2ISigmoidACV2IBroadcastABDV2IBroadcastBAEV2OMultiplyCD_V_V2_11160318154034397263"
|
|
body = receipt["kernels"].replace(f'host_name("{prefix}', 'host_name("kernel_qwen_mtplx_silu_bf16')
|
|
assert len(re.findall(r'\[\[host_name',body))==19
|
|
unit("captured-silu-jit",body,digest)
|
|
return "".join(output)
|
|
|
|
|
|
def gather_axis_source(runtime):
|
|
output = ["\n// Runtime GatherAxis shader: Copyright Apple Inc. SPDX-License-Identifier: MIT\n"]
|
|
for name,digest,start,end in (
|
|
("utils.h","5e1568e9edde9d05dbf86f68fa0d6c6240f2c32b973c7c6a76166b9c0d91543d",
|
|
"template <typename IdxT = int64_t>\nMETAL_FUNC IdxT elem_to_loc(\n IdxT", "// Non templated version"),
|
|
("indexing/gather_axis.h","e1a745391ff4990f3f1ad75c5687c3b102dcdc4833d8fbbac38e10f54af29af4",None,None),
|
|
):
|
|
raw = (runtime/"mlx/backend/metal/kernels"/name).read_bytes()
|
|
if hashlib.sha256(raw).hexdigest()!=digest:
|
|
raise ValueError(f"Pinned GatherAxis source changed: {name}")
|
|
text = raw.decode()
|
|
body = text[text.index(start):text.index(end,text.index(start))] if start else "".join(
|
|
line for line in text.splitlines(keepends=True) if not line.startswith(("#include","#pragma once")))
|
|
output.append(f"// Runtime unit: {name}; file SHA256: {digest}\n"
|
|
f"// Runtime unit SHA256: {hashlib.sha256(body.encode()).hexdigest()}\n"
|
|
"// BEGIN RUNTIME UNIT\n"+body+"// END RUNTIME UNIT\n")
|
|
for dtype,type_tag in (("bfloat16_t","bf16"),("half","f16"),("float","f32"),("uint32_t","u32")):
|
|
for idx,tag in (("uint32_t",""),("int32_t","idxi32_")):
|
|
for loc in ("int","int64_t"):
|
|
for sc in (False,True):
|
|
for ic in (False,True):
|
|
args = f"{dtype}, {idx}, {loc}, {str(sc).lower()}, {str(ic).lower()}"
|
|
name = f"kernel_qwen_mtplx_gather_axis_{type_tag}_{tag}{loc}_{int(sc)}{int(ic)}"
|
|
output.append(f'template [[host_name("{name}")]]\n[[kernel]] decltype(gather_axis<{args}>) gather_axis<{args}>;\n')
|
|
for loc in ("int","int64_t"):
|
|
for sc in (False,True):
|
|
for ic in (False,True):
|
|
args = f"bool, int64_t, {loc}, {str(sc).lower()}, {str(ic).lower()}"
|
|
name = f"kernel_qwen_mtplx_gather_axis_bool_idxi64_{loc}_{int(sc)}{int(ic)}"
|
|
output.append(f'template [[host_name("{name}")]]\n[[kernel]] decltype(gather_axis<{args}>) gather_axis<{args}>;\n')
|
|
for loc in ("int","int64_t"):
|
|
for sc in (False,True):
|
|
for ic in (False,True):
|
|
args = f"int64_t, int64_t, {loc}, {str(sc).lower()}, {str(ic).lower()}"
|
|
name = f"kernel_qwen_mtplx_gather_axis_i64_idxi64_{loc}_{int(sc)}{int(ic)}"
|
|
output.append(f'template [[host_name("{name}")]]\n[[kernel]] decltype(gather_axis<{args}>) gather_axis<{args}>;\n')
|
|
return "".join(output)
|
|
|
|
|
|
def scatter_axis_source(runtime):
|
|
output = ["\n// Runtime ScatterAxis: Copyright Apple Inc. SPDX-License-Identifier: MIT\n"
|
|
"#include <metal_atomic>\nnamespace mtplx_eager_scatter {\n"]
|
|
for name,digest,start,end in (
|
|
("complex.h","16e8a815b2cbdb6070e0824e64fe33fccb6e918f1b84ea5c792bd89d33e57bf1",None,None),
|
|
("atomic.h","4c35ea2798a2335502865247aee878149fc9ada0d7e84c05d771baef0c7fcc60",None,None),
|
|
("reduction/ops.h","78d06730fc9564a73944e7f1fe3897d25c8789b28a939bf418e1968db311da41",
|
|
"#define DEFINE_SIMD_REDUCE()", "template <typename U>\nstruct Prod"),
|
|
("indexing/indexing.h","e820b8ee2b5132a97122780c12433ebb5100d8078d31e211d0429400a11415bb",None,None),
|
|
("indexing/scatter.h","fa799c286378c59fbb3aeb973e74bf471861e3356ce2817968fe5d6872c3d9ad",None,None),
|
|
("indexing/scatter_axis.h","43eabd0216101f8e32f5cdd19ce40b7f954564be27fad98a5e0fe345e7b94ce5",None,None),
|
|
):
|
|
raw = (runtime/"mlx/backend/metal/kernels"/name).read_bytes()
|
|
if hashlib.sha256(raw).hexdigest()!=digest:
|
|
raise ValueError(f"Pinned ScatterAxis source changed: {name}")
|
|
text = raw.decode()
|
|
body = text[text.index(start):text.index(end,text.index(start))] if start else "".join(
|
|
line for line in text.splitlines(keepends=True) if not line.startswith(("#include","#pragma once")))
|
|
output.append(f"// Runtime unit: {name}; file SHA256: {digest}\n"
|
|
f"// Runtime unit SHA256: {hashlib.sha256(body.encode()).hexdigest()}\n"
|
|
"// BEGIN RUNTIME UNIT\n"+body+"// END RUNTIME UNIT\n")
|
|
for loc in ("int","int64_t"):
|
|
for uc in (False,True):
|
|
for ic in (False,True):
|
|
args = f"bool, int64_t, {loc}, None, {str(uc).lower()}, {str(ic).lower()}"
|
|
name = f"kernel_qwen_mtplx_scatter_axis_bool_idxi64_{loc}_{int(uc)}{int(ic)}"
|
|
output.append(f'template [[host_name("{name}")]]\n[[kernel]] decltype(scatter_axis<{args}>) scatter_axis<{args}>;\n')
|
|
raw = (runtime/"mlx/backend/metal/jit/indexing.h").read_bytes()
|
|
digest = "1b38dbdf3120eca3e5266ffd6d69691591c9432c0b5acece4c5ab2de7a9e802f"
|
|
if hashlib.sha256(raw).hexdigest() != digest:
|
|
raise ValueError("Pinned Scatter JIT wrapper changed")
|
|
template = raw.decode().split('scatter_kernels = R"(',1)[1].split(')";',1)[0]
|
|
wrappers = []
|
|
# One index array/axis, rank-one indices: the original planner always
|
|
# selects nwork=1 for the sparse raw-logit penalty path.
|
|
for tag,kind in (("bf16","bfloat"),("f16","half"),("f32","float")):
|
|
for contiguous in ("false","true"):
|
|
for loc in ("int","int64_t"):
|
|
body = template.format(f"{tag}idxi64_sum",kind,"int64_t",f"Sum<{kind}>",1,
|
|
"const device int64_t *idx0 [[buffer(20)]],","idx0",contiguous,1,loc)
|
|
body = body.replace(f"void scatter{tag}idxi64_sum_",f"void kernel_qwen_mtplx_scatter_{tag}_idxi64_sum_",1)
|
|
# This dependency lives in a namespace alongside ScatterAxis.
|
|
# Preserve the original unqualified entry ABI explicitly.
|
|
name = f"kernel_qwen_mtplx_scatter_{tag}_idxi64_sum_1_updc_{contiguous}_nwork1_{loc}"
|
|
body = body.replace("[[kernel]]",f'[[host_name("{name}")]] [[kernel]]',1)
|
|
wrappers.append(body)
|
|
body = "".join(wrappers)
|
|
output.append(f"// Runtime unit: scatter JIT wrappers; file SHA256: {digest}\n"
|
|
f"// Runtime unit SHA256: {hashlib.sha256(body.encode()).hexdigest()}\n"
|
|
"// BEGIN RUNTIME UNIT\n"+body+"// END RUNTIME UNIT\n")
|
|
output.append("} // namespace mtplx_eager_scatter\n")
|
|
return "".join(output)
|
|
|
|
|
|
def swiglu_source():
|
|
raw = (OUTPUT.parents[1]/"tests/fixtures/mtplx-swiglu-jit.json").read_bytes()
|
|
digest = hashlib.sha256(raw).hexdigest()
|
|
if digest != "8cb2a51a7d9f0cb91a8a7669ef844c3964dc469ec9e4436a719e47f025b8e781":
|
|
raise ValueError("Pinned SwiGLU JIT receipt changed")
|
|
receipt = json.loads(raw)
|
|
prefix = re.search(r'host_name\("([^\"]+)_contiguous"',receipt["kernels"]).group(1)
|
|
body = receipt["kernels"].replace(f'host_name("{prefix}', 'host_name("kernel_qwen_mtplx_swiglu_bf16')
|
|
assert len(re.findall(r'\[\[host_name',body))==19
|
|
return ("\n// Actual runtime-generated SwiGLU: Copyright Apple Inc.; SPDX-License-Identifier: MIT\n"
|
|
f"// Runtime unit: captured-swiglu-jit; file SHA256: {digest}\n"
|
|
f"// Runtime unit SHA256: {hashlib.sha256(body.encode()).hexdigest()}\n"
|
|
"// BEGIN RUNTIME UNIT\n"+body+"// END RUNTIME UNIT\n")
|
|
|
|
|
|
def compute_g_source(runtime):
|
|
output = ["\n// Original compute_g JIT dependencies; Apple MIT except cexpf.h (Apache-2.0).\n"]
|
|
def unit(name,body,digest):
|
|
output.append(f"// Runtime unit: {name}; file SHA256: {digest}\n"
|
|
f"// Runtime unit SHA256: {hashlib.sha256(body.encode()).hexdigest()}\n"
|
|
"// BEGIN RUNTIME UNIT\n"+body+"// END RUNTIME UNIT\n")
|
|
for name,digest,sections in (
|
|
("complex.h","16e8a815b2cbdb6070e0824e64fe33fccb6e918f1b84ea5c792bd89d33e57bf1",None),
|
|
("cexpf.h","88b6e15a52a5800d98d9bc6da840ca5cf70bf572fda136409580c1f17b1e0aab",None),
|
|
("utils.h","5e1568e9edde9d05dbf86f68fa0d6c6240f2c32b973c7c6a76166b9c0d91543d",
|
|
[("template <typename U>\nstruct Limits", "///////////////////////////////////////////////////////////////////////////////"),
|
|
("inline float log1p(float x)", "///////////////////////////////////////////////////////////////////////////////")]),
|
|
("unary_ops.h","0a5492b65ae39ecb6d8b04e64ea5007e0a4ff60d8d9428559bfbfd03387ece2a",
|
|
[("struct Exp {","struct Expm1 {"),("struct Negative {","struct Real {")]),
|
|
("binary_ops.h","2dd13c2496f5d6f0856e4ca99db2ceb5c6d7dc0c344b9ba8e41ef3b7d7ed8a97",
|
|
[("struct Add {","struct FloorDivide {"),("struct LogAddExp {","struct Maximum {")]),
|
|
):
|
|
raw = (runtime/"mlx/backend/metal/kernels"/name).read_bytes()
|
|
if hashlib.sha256(raw).hexdigest()!=digest:
|
|
raise ValueError(f"Pinned compute_g dependency changed: {name}")
|
|
text = raw.decode()
|
|
body = "".join(text[text.index(start):text.index(end,text.index(start))] for start,end in sections) if sections else "".join(
|
|
line for line in text.splitlines(keepends=True) if not line.startswith(("#include","#pragma once")))
|
|
unit(name,body,digest)
|
|
raw = (OUTPUT.parents[1]/"tests/fixtures/mtplx-compute_g-jit.json").read_bytes()
|
|
digest = hashlib.sha256(raw).hexdigest()
|
|
if digest != "701e2f54b7cb8bf97616f83256657e6c5e8cc8b46f4b65c559ccc030ab111dbf":
|
|
raise ValueError("Pinned compute_g JIT receipt changed")
|
|
kernels = json.loads(raw)["kernels"]
|
|
prefix = re.search(r'host_name\("([^\"]+)_contiguous"',kernels).group(1)
|
|
body = kernels.replace(f'host_name("{prefix}', 'host_name("kernel_qwen_mtplx_compute_g_bf16')
|
|
assert len(re.findall(r'\[\[host_name',body))==19
|
|
unit("captured-compute_g-jit",body,digest)
|
|
return "".join(output)
|
|
|
|
|
|
def qsa_prepare_source(reference):
|
|
path = reference/"mtplx/kernels/qsa_indexer_prepare.py"
|
|
raw = path.read_bytes()
|
|
if hashlib.sha256(raw).hexdigest() != "a77f6ca5ae805e729519c4629ae88b455a6dbf473a457a6e1c8219174eb59091":
|
|
raise ValueError("Pinned QSA preparation source changed")
|
|
constants = dict(heads=4,head_dim=128,rotary_dim=64,half_rotary=32,ratio=4,
|
|
eps=1e-6,attention_scaling=1.0)
|
|
def literal(node):
|
|
if isinstance(node,ast.Constant): return node.value
|
|
if isinstance(node,ast.Name): return constants[node.id]
|
|
if isinstance(node,ast.Call) and isinstance(node.func,ast.Name) and node.func.id=="float":
|
|
return float(literal(node.args[0]))
|
|
if isinstance(node,ast.FormattedValue):
|
|
value=literal(node.value)
|
|
return repr(value) if node.conversion==ord("r") else str(value)
|
|
if isinstance(node,ast.JoinedStr): return "".join(literal(v) for v in node.values)
|
|
raise ValueError(f"Unsupported QSA header expression: {ast.dump(node)}")
|
|
output=[]
|
|
for function,name,inputs,outputs in (
|
|
("_prepare_queries_kernel","qsa_prepare_q",
|
|
"T:raw_q;int64_t:raw_q_strides;T:norm_weight;int64_t:norm_weight_strides;float:inv_freq;int64_t:inv_freq_strides;int:pos_start","T:prepared_q"),
|
|
("_pool_keys_kernel","qsa_pool_k",
|
|
"T:raw_keys;int64_t:raw_keys_strides;T:norm_weight;int64_t:norm_weight_strides;float:inv_freq;int64_t:inv_freq_strides;int:block_start","T:pooled"),
|
|
):
|
|
fn=next(n for n in ast.parse(raw).body if isinstance(n,ast.FunctionDef) and n.name==function)
|
|
values={n.targets[0].id:n.value for n in fn.body if isinstance(n,ast.Assign)
|
|
and isinstance(n.targets[0],ast.Name)}
|
|
# metal_stdlib is already included at translation-unit scope; Metal
|
|
# module imports cannot be nested inside our collision-avoiding scope.
|
|
header = "".join(line for line in literal(values["header"]).splitlines(keepends=True)
|
|
if not line.lstrip().startswith("#include"))
|
|
output.append(f"\nnamespace mtplx_{name} {{\n"+header)
|
|
args = [("constant const int64_t*" if key.endswith("_strides") else kind,key)
|
|
for kind,key in parameters(inputs)+parameters(outputs,output=True)]
|
|
output.append(kernel_source(f"mtplx/kernels/qsa_indexer_prepare.py::{function}",raw,
|
|
literal(values["source"]),f"kernel_qwen_mtplx_{name}",
|
|
args,
|
|
"threadgroup_position_in_grid,thread_index_in_simdgroup","typename T",[("bf16","bfloat")]))
|
|
output.append(f"}} // namespace mtplx_{name}\n")
|
|
return "".join(output)
|
|
|
|
|
|
def qsa_select_source(reference):
|
|
path = reference / "mtplx/kernels/qsa_indexer_select.py"
|
|
raw = path.read_bytes()
|
|
digest = hashlib.sha256(raw).hexdigest()
|
|
if digest != "a3c74af27a7045c12f2893a8b7a91724c00d8a4148315c3165f3480c83016cf3":
|
|
raise ValueError("Pinned QSA selector source changed")
|
|
tree = ast.parse(raw)
|
|
header = next(ast.literal_eval(n.value) for n in tree.body
|
|
if isinstance(n, ast.Assign) and isinstance(n.targets[0], ast.Name)
|
|
and n.targets[0].id == "_HEADER")
|
|
fn = next(n for n in tree.body if isinstance(n, ast.FunctionDef) and n.name == "_selector_kernel")
|
|
body = next(n.value for n in fn.body if isinstance(n, ast.Assign)
|
|
and isinstance(n.targets[0], ast.Name) and n.targets[0].id == "source")
|
|
assert isinstance(body, ast.BinOp) and isinstance(body.op, ast.Add) and body.right.id == "epilogue"
|
|
body = ast.literal_eval(body.left)
|
|
epilogues = {}
|
|
ep = next(n for n in tree.body if isinstance(n, ast.FunctionDef) and n.name == "_epilogue")
|
|
for branch in ep.body:
|
|
if isinstance(branch, ast.If):
|
|
mode = ast.literal_eval(branch.test.comparators[0])
|
|
names, source = ast.literal_eval(branch.body[0].value)
|
|
epilogues[mode] = dict(outputs=names, source=source,
|
|
sha256=hashlib.sha256(source.encode()).hexdigest())
|
|
prefill_raw = (reference / "mtplx/kernels/qsa_indexer_prefill.py").read_bytes()
|
|
prefill_digest = hashlib.sha256(prefill_raw).hexdigest()
|
|
if prefill_digest != "4d6fd428243c001746f69f8aed45991356772c2bd4a45586eb3c6813c91998d3":
|
|
raise ValueError("Pinned QSA prefill source changed")
|
|
prefill_tree = ast.parse(prefill_raw)
|
|
literals = {n.targets[0].id: ast.literal_eval(n.value) for n in prefill_tree.body
|
|
if isinstance(n, ast.Assign) and isinstance(n.targets[0], ast.Name)
|
|
and isinstance(n.value, ast.Constant)}
|
|
topk = next(n for n in prefill_tree.body if isinstance(n, ast.FunctionDef) and n.name == "_prefill_topk_kernel")
|
|
assignments = {n.targets[0].id:n.value for n in topk.body if isinstance(n,ast.Assign) and isinstance(n.targets[0],ast.Name)}
|
|
assert assignments["source"].right.id == "epilogue"
|
|
topk_body = ast.literal_eval(assignments["source"].left)
|
|
# Preserve literal f-string segments verbatim; Rust fills only these
|
|
# named header substitutions, including constants and the output mode.
|
|
topk_header = "".join(n.value if isinstance(n,ast.Constant) else "@"+ast.unparse(n.value)+"@"
|
|
for n in assignments["header"].right.values)
|
|
prefill = dict(file_sha256=prefill_digest, mpp_header=literals["_MPP_SCORE_HEADER"],
|
|
mpp_body=literals["_MPP_SCORE_SOURCE"],topk_header=topk_header,topk_body=topk_body)
|
|
for key in ("mpp_header","mpp_body","topk_header","topk_body"):
|
|
prefill[key+"_sha256"] = hashlib.sha256(prefill[key].encode()).hexdigest()
|
|
return json.dumps(dict(revision=REVISION, file_sha256=digest, header=header, body=body,
|
|
header_sha256=hashlib.sha256(header.encode()).hexdigest(),
|
|
body_sha256=hashlib.sha256(body.encode()).hexdigest(),
|
|
epilogues=epilogues,prefill=prefill), indent=2) + "\n"
|
|
|
|
|
|
def generate(reference, gated_delta):
|
|
revision = subprocess.check_output(
|
|
["git", "-C", str(reference), "rev-parse", "HEAD"], text=True).strip()
|
|
if revision != REVISION:
|
|
raise ValueError(f"Expected MTPLX {REVISION}, got {revision}")
|
|
subprocess.run(["git", "-C", str(reference), "diff", "--exit-code", "HEAD"],
|
|
check=True, stdout=subprocess.DEVNULL)
|
|
output = ["// MTPLX bodies: Copyright 2026 MTPLX. SPDX-License-Identifier: Apache-2.0\n"
|
|
"// Dependency bodies carry their own license notice below.\n"
|
|
"// Generated by tools/mtplx-kernel-source.py; do not edit bodies.\n"
|
|
f"// MTPLX revision: {REVISION}\n"
|
|
"// Only entry-point ABI and template instantiations are adapted.\n"
|
|
"#include <metal_stdlib>\nusing namespace metal;\n"]
|
|
for module, variable, name, inputs, outputs, builtins, extra, variants in KERNELS:
|
|
path = reference / f"mtplx/kernels/{module}.py"
|
|
raw = path.read_bytes()
|
|
constants = {node.targets[0].id: node.value.value
|
|
for node in ast.parse(raw).body
|
|
if isinstance(node, ast.Assign) and len(node.targets) == 1
|
|
and isinstance(node.targets[0], ast.Name)
|
|
and isinstance(node.value, ast.Constant)
|
|
and isinstance(node.value.value, str)}
|
|
body = constants[variable]
|
|
args = parameters(inputs) + parameters(outputs, output=True)
|
|
output.append(kernel_source(
|
|
f"mtplx/kernels/{module}.py::{variable}", raw, body,
|
|
f"kernel_qwen_mtplx_{name}", args, builtins,
|
|
f"typename T{', ' + extra if extra else ''}", variants))
|
|
output.append(gated_delta_source(gated_delta))
|
|
output.append(gather_front_source(reference.parent/"mtplx-runtime-0.32.2"))
|
|
output.append(silu_source(reference.parent/"mtplx-runtime-0.32.2"))
|
|
output.append(gather_axis_source(reference.parent/"mtplx-runtime-0.32.2"))
|
|
output.append(gather_rows_source(reference.parent/"mtplx-runtime-0.32.2"))
|
|
output.append(scatter_axis_source(reference.parent/"mtplx-runtime-0.32.2"))
|
|
output.append(swiglu_source())
|
|
output.append(compute_g_source(reference.parent/"mtplx-runtime-0.32.2"))
|
|
output.append(qsa_prepare_source(reference))
|
|
output.append(qsa_attention_source(reference))
|
|
return "".join(output)
|
|
|
|
|
|
def qsa_attention_source(reference):
|
|
output=[]
|
|
for module,variable,digest,inputs,builtins,extra,variants in (
|
|
("qsa_flash_skip","_SRC","27d49eb93f0de82013fdc0b271e77076b19d1972450e1bbc44cdd92984669103",
|
|
"T:q,k,v;int:blocks,params;float:scale",
|
|
"threadgroup_position_in_grid,thread_position_in_threadgroup","int GQA",
|
|
[(f"g{g}_{tag}",f"{kind}, {g}") for g in (1,12) for tag,kind in (("bf16","bfloat"),("f16","half"))]),
|
|
("qsa_prefill_flash","_SOURCE","5ca852ea441810a47b374a1c94a49a636055dab31f7f11c4ce0118f7f70b388d",
|
|
"T:q;int64_t:q_strides;T:k;int64_t:k_strides;T:v;int64_t:v_strides;int:block_ids;int64_t:block_ids_strides;bool:block_valid;int64_t:block_valid_strides;int:params;float:scale",
|
|
"threadgroup_position_in_grid,thread_index_in_simdgroup","",
|
|
[("bf16","bfloat"),("f16","half")]),
|
|
):
|
|
raw=(reference/f"mtplx/kernels/{module}.py").read_bytes()
|
|
if hashlib.sha256(raw).hexdigest()!=digest: raise ValueError(f"Pinned {module} source changed")
|
|
values={n.targets[0].id:ast.literal_eval(n.value) for n in ast.parse(raw).body
|
|
if isinstance(n,ast.Assign) and isinstance(n.targets[0],ast.Name) and isinstance(n.value,ast.Constant)}
|
|
header=values["_HEADER"]
|
|
includes="".join(line for line in header.splitlines(keepends=True) if line.lstrip().startswith("#include"))
|
|
header="".join(line for line in header.splitlines(keepends=True) if not line.lstrip().startswith("#include"))
|
|
output.append(includes+f"\nnamespace mtplx_{module} {{\n"+header)
|
|
args=[("constant const int64_t*" if name.endswith("_strides") else kind,name)
|
|
for kind,name in parameters(inputs)+parameters("T:out",output=True)]
|
|
output.append(kernel_source(f"mtplx/kernels/{module}.py::{variable}",raw,values[variable],
|
|
f"kernel_qwen_mtplx_{module}",args,builtins,"typename T"+(f", {extra}" if extra else ""),variants))
|
|
output.append(f"}} // namespace mtplx_{module}\n")
|
|
return "".join(output)
|
|
|
|
|
|
def main():
|
|
parser = argparse.ArgumentParser(description=__doc__)
|
|
parser.add_argument("reference", type=Path)
|
|
parser.add_argument("--gated-delta-source", type=Path, required=True,
|
|
help="gated_delta.py from MTPLX's pinned mlx-lm 0.31.3 installation")
|
|
parser.add_argument("--check", action="store_true")
|
|
parser.add_argument("--qsa-select", action="store_true", help="Print original dynamic QSA selector source sections")
|
|
args = parser.parse_args()
|
|
source = generate(args.reference, args.gated_delta_source)
|
|
selector = qsa_select_source(args.reference)
|
|
if args.check:
|
|
if OUTPUT.read_text() != source:
|
|
raise SystemExit(f"MTPLX kernel source differs: {OUTPUT}")
|
|
if QSA_SELECT_OUTPUT.read_text() != selector:
|
|
raise SystemExit(f"MTPLX selector source differs: {QSA_SELECT_OUTPUT}")
|
|
print(f"All {len(KERNELS) + 6} Metal bodies and {sum(len(k[-1]) for k in KERNELS) + 10} "
|
|
f"custom entry points and {source.count('// BEGIN RUNTIME UNIT')} runtime units plus dynamic QSA decode/prefill sources match the pinned reference export")
|
|
else:
|
|
print(selector if args.qsa_select else source, end="")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|