Investigate 10x performance with mathematical bounds and controlled Metal experiments

Author: mizorewww

README.md

@@ -160,7 +160,15 @@ GitHub Actions runs small-model CPU tests on a macOS arm64 runner. Full checkpoi
 
 ## Performance research
 
-[PERFORMANCE_RESEARCH.md](docs/PERFORMANCE_RESEARCH.md) analyzes the measured bottlenecks and proposes experiments for compilation, projection quantization, actual sparse window attention, batching and exact head pruning. These are research directions, not additional measured speedups in the released runtime.
+The performance investigations include both mathematical analysis and independent local experiments:
+
+- [Initial performance research](docs/PERFORMANCE_RESEARCH.md): implementation bottlenecks, MLX kernel dispatch, and a controlled experiment plan.
+- [Mathematical investigation of a further 10× speedup](docs/MATH_10X_RESEARCH.md): arithmetic budgets, conditional bandwidth bounds, real weight spectra, exact reuse, and smaller-model designs.
+- [Engineering investigation](docs/ENGINEERING_10X_RESEARCH.md): measured compilation, quantization, final-head selection, custom Metal kernels, and representative matrix multiplications.
+
+[experiments/](experiments) contains the research scripts and their raw measurements. The published runtime's performance and validation results are in [BENCHMARKS.md](BENCHMARKS.md); each experimental variant has its own timing and correctness results.
+
+The current investigation does not support a further universal 10× speedup with the same checkpoints. Selected cases show approximately 1.03–1.08× paired median speedups; the engineering report gives the uncertainty intervals, quantization fidelity results, and custom Metal kernel measurements.
 
 To prepare model cards and verified exports for publication, install the reference extras and run:
 

README.zh-CN.md

@@ -94,6 +94,13 @@ uv run laya-mlx convert \
 
 详细的 API、测试和复现命令见 [英文 README](README.md)。[BENCHMARKS.md](BENCHMARKS.md) 包含本机 PyTorch MPS FP32、MLX FP32 与 MLX FP16 的端到端 P50/P95、吞吐量、内存、数值一致性、重复运行和固定抽样分类测试。所有原始测量数据位于 [benchmarks/results](benchmarks/results)，GPU 测试应串行运行。
 
-[性能优化研究](docs/PERFORMANCE_RESEARCH.md) 根据实际实现、测量结果和 MLX 源码，分析编译、量化、稀疏局部注意力、批处理等方向，并给出正确性约束和实验方案。该研究中的优化建议尚未作为加速功能集成到本次发布。
+[初步性能研究](docs/PERFORMANCE_RESEARCH.md) 分析实现、基准和 MLX 源码。针对“能否再快一个数量级”，另有两份深入报告：
+
+- [数学分析](docs/MATH_10X_RESEARCH.md)：计算预算、带宽条件下界、真实权重谱、精确复用，以及蒸馏学生模型的设计空间。
+- [工程实测](docs/ENGINEERING_10X_RESEARCH.md)：编译、量化、最后一层输出裁剪、自定义 Metal 核与矩阵乘法实验。
+
+[experiments/](experiments) 保存研究脚本和原始数据。发布版本的结果见 [BENCHMARKS.md](BENCHMARKS.md)，各实验变体的耗时与数值一致性单独记录。
+
+目前证据不支持相同检查点下普遍再快 10 倍。部分场景的逐轮配对中位加速约为 1.03–1.08 倍；误差区间、量化保真结果和自定义 Metal 核的实测详见工程报告。
 
 这是独立的 MLX 移植，模型能力及其限制来自上游；模型输出概率不等于答案必然正确。采用 Apache-2.0，原作者与移植说明见 [NOTICE](NOTICE)。

docs/ENGINEERING_10X_RESEARCH.md

@@ -0,0 +1,46 @@
+"""CPU-only analysis of paired experiment files; preserves the raw timings."""
+
+import json
+from pathlib import Path
+
+import numpy as np
+
+from benchmarks.common import save_json
+
+
+def main():
+    rng = np.random.default_rng(20260919)
+    report = {
+        "method": "median of paired per-round eager/candidate ratios; percentile bootstrap of round indices (2000 resamples), exploratory 95% intervals",
+        "results": [],
+    }
+    for path in sorted(Path("experiments/engineering").glob("*-paired.json")):
+        raw = json.loads(path.read_text())
+        for row in raw["results"]:
+            for mode in ("forward", "end_to_end"):
+                baseline = np.asarray(row["candidates"]["eager"][mode]["samples_ms"])
+                for candidate, metrics in row["candidates"].items():
+                    if candidate == "eager":
+                        continue
+                    times = np.asarray(metrics[mode]["samples_ms"])
+                    ratios = baseline / times
+                    resampled = ratios[rng.integers(len(ratios), size=(2000, len(ratios)))]
+                    interval = np.percentile(np.median(resampled, axis=1), [2.5, 97.5])
+                    report["results"].append(
+                        {
+                            "source": path.name,
+                            "case": row["case"],
+                            "candidate": candidate,
+                            "mode": mode,
+                            "rounds": len(ratios),
+                            "paired_speedup_median": float(np.median(ratios)),
+                            "bootstrap_95_interval": interval.tolist(),
+                            "per_round_min_max": [float(ratios.min()), float(ratios.max())],
+                            "ratio_of_p50": float(np.median(baseline) / np.median(times)),
+                        }
+                    )
+    save_json(Path("experiments/engineering/paired_analysis.json"), report)
+
+
+if __name__ == "__main__":
+    main()

docs/MATH_10X_RESEARCH.md

@@ -0,0 +1,63 @@
+"""Experimental Metal GELU/gate kernel, using MLX's same erf implementation.
+
+The vendored erf/expm1 helpers are from MLX v0.32.2, copyright Apple and
+Norbert Juffa. Their original notices are retained in vendor/ and MLX_LICENSE.
+Only this experiment uses them; the production package is unaffected.
+"""
+
+from pathlib import Path
+
+import mlx.core as mx
+import mlx.nn as nn
+
+VENDOR = Path(__file__).with_name("vendor")
+helpers = "\n".join(
+    line
+    for filename in ("mlx_expm1f.h", "mlx_erf.h")
+    for line in (VENDOR / filename).read_text().splitlines()
+    if not line.startswith(("#include", "#pragma"))
+)
+GELU_GATE = mx.fast.metal_kernel(
+    name="laya_experimental_exact_gelu_gate",
+    input_names=["inp"],
+    output_names=["out"],
+    header="namespace laya_erf { using namespace metal;\n" + helpers + "\n}",
+    source="""
+        uint elem = thread_position_in_grid.x;
+        uint row = elem / I;
+        uint col = elem % I;
+        T value = inp[row * (2 * I) + col];
+        T gate = inp[row * (2 * I) + I + col];
+        T scaled = value / T(1.4142135623730951);
+        T erf_value = T(laya_erf::erf(float(scaled)));
+        T gelu = T(T(value * T(T(1) + erf_value)) / T(2));
+        out[elem] = gelu * gate;
+    """,
+    compile_options={"math_mode": "safe"},
+)
+
+
+def metal_gelu_gate(x):
+    if x.dtype != mx.float16:
+        raise ValueError("This research kernel has only been implemented and tested for FP16")
+    if x.shape[-1] % 2:
+        raise ValueError("GELU/gate input needs two equal-width branches")
+    intermediate = x.shape[-1] // 2
+    shape = (*x.shape[:-1], intermediate)
+    return GELU_GATE(
+        inputs=[x],
+        template=[("T", x.dtype), ("I", intermediate)],
+        grid=(x.size // 2, 1, 1),
+        threadgroup=(256, 1, 1),
+        output_shapes=[shape],
+        output_dtypes=[x.dtype],
+    )[0]
+
+
+class MetalMLP(nn.Module):
+    def __init__(self, layer):
+        super().__init__()
+        self.Wi, self.Wo = layer.Wi, layer.Wo
+
+    def __call__(self, x):
+        return self.Wo(metal_gelu_gate(self.Wi(x)))