139 lines
7.4 KiB
Python
139 lines
7.4 KiB
Python
"""Probe projection components selected by MTPLX's Qwen implementation.
|
|
|
|
This is an operator diagnostic, not a full-model or serving benchmark. Only the
|
|
selected projections are materialized; no weights are downloaded.
|
|
"""
|
|
|
|
import argparse
|
|
import hashlib
|
|
import json
|
|
import statistics
|
|
import time
|
|
from pathlib import Path
|
|
|
|
import mlx.core as mx
|
|
import numpy as np
|
|
from mtplx.models import qwen4_exp
|
|
|
|
|
|
def main():
|
|
parser = argparse.ArgumentParser(description=__doc__)
|
|
parser.add_argument("model", type=Path)
|
|
parser.add_argument("--capture", type=Path, help="Capture only the 32-row router GPU operation")
|
|
parser.add_argument("--check-split-k", action="store_true",
|
|
help="Diagnose router rounding with explicitly batched K partitions")
|
|
args = parser.parse_args()
|
|
assert mx.__version__ == "0.32.2", "Use the pinned MTPLX MLX environment"
|
|
index = json.loads((args.model / "model.safetensors.index.json").read_text())["weight_map"]
|
|
config = json.loads((args.model / "config.json").read_text())
|
|
quantization = config["quantization"]
|
|
model_args = qwen4_exp.TextArgs.from_dict(config["text_config"])
|
|
source = Path(qwen4_exp.__file__).resolve()
|
|
print(json.dumps({"reference": "MTPLX Qwen projection components",
|
|
"source": str(source),
|
|
"source_sha256": hashlib.sha256(source.read_bytes()).hexdigest(),
|
|
"mlx_version": mx.__version__}), flush=True)
|
|
values = []
|
|
state = 0x12345678
|
|
for _ in range(2560):
|
|
state = (state * 1664525 + 1013904223) & 0xFFFFFFFF
|
|
values.append(((state >> 24) - 128) / 32)
|
|
for suffix in ("linear_attn.in_proj_qkv", "mlp.gate",
|
|
"linear_attn.in_proj_b", "mlp.shared_expert_gate"):
|
|
name = "language_model.model.layers.0." + suffix
|
|
weights = {}
|
|
for part in ("weight", "scales", "biases"):
|
|
key = name + "." + part
|
|
weights[part] = mx.load(str(args.model / index[key]))[key]
|
|
mx.eval(*weights.values())
|
|
layout = quantization[name]
|
|
module = (qwen4_exp.GatedDeltaNet(model_args) if suffix.startswith("linear_attn.")
|
|
else qwen4_exp.SparseMoeBlock(model_args))
|
|
attribute = suffix.split(".")[-1]
|
|
projection = getattr(module, attribute).to_quantized(**layout)
|
|
projection.load_weights(list(weights.items()), strict=True)
|
|
setattr(module, attribute, projection)
|
|
for rows in (1, 2, 3, 4, 31, 32, 63, 64, 2047, 2048):
|
|
x = mx.array([[values] * rows], dtype=mx.bfloat16)
|
|
mx.eval(x)
|
|
capture = args.capture and suffix == "mlp.gate" and rows == 32
|
|
if capture:
|
|
mx.metal.start_capture(str(args.capture))
|
|
try:
|
|
result = getattr(module, attribute)(x)
|
|
mx.eval(result)
|
|
finally:
|
|
if capture:
|
|
mx.metal.stop_capture()
|
|
raw = np.asarray(result.astype(mx.float32)).astype("<f4").tobytes()
|
|
digest = hashlib.sha256(raw).hexdigest()
|
|
if suffix == "linear_attn.in_proj_qkv" and rows == 1:
|
|
assert digest == "8952fadfdb1ef4450fd23b0060d62ea9b118367733014985a94c2d815a2be203"
|
|
print(json.dumps({"mlx_version": mx.__version__, "projection": name,
|
|
"quantization": layout, "rows": rows,
|
|
"first_row_sha256": hashlib.sha256(raw[:result.shape[-1] * 4]).hexdigest(),
|
|
"sha256": digest}), flush=True)
|
|
if rows in (2, 3, 4, 32, 64, 2048):
|
|
# Operator call + eval wall time, not a serving benchmark or
|
|
# an extra GPU counter probe. Discard the first warmup sample.
|
|
timings = []
|
|
for trial in range(33):
|
|
started = time.perf_counter_ns()
|
|
measured = getattr(module, attribute)(x)
|
|
mx.eval(measured)
|
|
elapsed = (time.perf_counter_ns() - started) / 1000
|
|
if trial:
|
|
timings.append(elapsed)
|
|
print(json.dumps({"diagnostic": "projection_call_eval_wall",
|
|
"projection": name, "rows": rows, "samples": len(timings),
|
|
"median_us": statistics.median(timings)}), flush=True)
|
|
if args.check_split_k and suffix == "mlp.gate" and rows in (32, 64):
|
|
# Counterfactual component operations, not a replacement oracle:
|
|
# a real weight batch bypasses the non-batched split-K dispatch.
|
|
for splits in (1, 2, 4, 5, 8, 10, 20, 40):
|
|
width = x.shape[-1] // splits
|
|
batches = max(2, splits)
|
|
partitioned = {}
|
|
for key, value in weights.items():
|
|
block = value.reshape(value.shape[0], splits, -1).transpose(1, 0, 2)
|
|
partitioned[key] = mx.repeat(block, batches, axis=0) if splits == 1 else block
|
|
inputs = x.reshape(rows, splits, width).transpose(1, 0, 2)
|
|
if splits == 1:
|
|
inputs = mx.repeat(inputs, batches, axis=0)
|
|
elif rows == 32:
|
|
# Small K raises MTPLX's QMV threshold to 33 rows. Pad
|
|
# to retain BF16 QMM dequantization in this control.
|
|
inputs = mx.concatenate([inputs, inputs], axis=1)
|
|
parts = mx.quantized_matmul(
|
|
inputs, partitioned["weight"], partitioned["scales"],
|
|
partitioned["biases"], transpose=True, **layout,
|
|
)
|
|
combined = parts[0] if splits == 1 else parts.sum(axis=0)
|
|
combined = combined[:rows]
|
|
mx.eval(combined)
|
|
diagnostic_hash = hashlib.sha256(
|
|
np.asarray(combined.astype(mx.float32)).astype("<f4").tobytes()
|
|
).hexdigest()
|
|
first_actual = np.asarray(combined.astype(mx.float32))[0]
|
|
first_expected = np.asarray(result.astype(mx.float32))[0, 0]
|
|
differences = np.flatnonzero(first_actual != first_expected)
|
|
print(json.dumps({"diagnostic": "explicit_batched_partitions",
|
|
"rows": rows, "splits": splits,
|
|
"partition_rows": inputs.shape[1],
|
|
"sha256": diagnostic_hash,
|
|
"different_columns": int(len(differences)),
|
|
"first_differences": [[int(i), float(first_actual[i]), float(first_expected[i])] for i in differences[:4]],
|
|
"matches_projection": diagnostic_hash == digest}), flush=True)
|
|
if splits == (20 if rows == 32 else 10):
|
|
assert diagnostic_hash == digest, "Split-K QMM decomposition differs from MTPLX projection"
|
|
if splits == 1:
|
|
expected_unsplit = {
|
|
32: "3da6d25816a695eda420aea0a76ff22ee4c59689857d0ad02ed16112d66625b6",
|
|
64: "8e0702490c8b94ea1c167fdbbca27107d7ecf6ec0c0b4f446ea3f8927d34c0f4",
|
|
}
|
|
assert diagnostic_hash == expected_unsplit[rows], "Unsplit control changed"
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|