[Help wanted] Accelerated Laya serving on Apple Silicon

[hsliuustc0106]

Community help wanted

We are looking for community contributors to help build and evaluate faster local Laya serving on Apple Silicon. Contributions to one part of the work are welcome; you do not need to implement the whole backend.

Useful ways to help:

• Profile Laya on PyTorch MPS and identify where inference time is spent.
• Implement and evaluate targeted optimizations; bring MLX, Core ML, or Metal expertise where measurements justify an alternative backend.
• Run reproducible benchmarks on available M-series Macs and share hardware details, commands, raw results, and output-parity checks.
• Help with persistent serving, readiness checks, tests, and setup documentation.

Please comment with the part you would like to tackle, your available hardware if relevant, and any proposed approach so contributors can coordinate before overlapping work. Baseline measurements, unsuccessful optimization results, and small focused PRs are all useful.

Agent-side integration is coordinated in ThinkFlowLab/system1-agents#20.

Problem

We want to serve Laya locally on Apple Silicon Macs with lower decision latency and a model that stays loaded across requests. Issue #1 introduces a Rust HTTP frontend, but leaves tokenization and inference in the backend. That frontend alone does not establish a model inference speedup.

Proposed scope

Add and validate an Apple Silicon inference backend for Laya behind the serving interface from #1. Start with the existing Laya PyTorch MPS implementation as the reference, profile its bottlenecks, and choose the smallest optimization that produces a measured improvement. Evaluate alternatives such as MLX or Core ML only if profiling justifies the additional implementation and maintenance.

• Keep the selected checkpoint resident across requests and perform a representative warmup before reporting readiness.
• Support Laya text decisions (choice, score, and noul) through the common /v1/systemone interface, preserving response fields and probability semantics.
• Make the selected device/backend explicit and observable; do not silently report CPU execution as GPU acceleration.
• Document supported Mac hardware, macOS/runtime versions, checkpoint, installation, startup, and a reproducible request example.
• Keep custom batching and scheduling out of the initial scope unless measurements identify them as necessary for the stated workload.

Validation

Before implementation benchmarking, define the hypothesis, independent variable, fixed controls, success criterion, and stop condition. Compare the reference PyTorch MPS backend with the proposed backend using the same Mac, checkpoint, inputs, question types, concurrency, and warmup/cache policy.

Report separately:

• Checkpoint download, model loading, warmup, and process-to-readiness time.
• Warm inference latency and end-to-end HTTP latency (p50/p95), including the frontend overhead.
• Throughput and memory use at a declared concurrency.
• Decision agreement and probability differences on a fixed set spanning all three decision types, input lengths, and option counts. Define numerical tolerances before comparing results.

Use one feasibility run, then two measured runs per configuration by default, with enough requests per run to support percentile estimates. Record variability, repository SHAs, runtime versions, commands, and raw results. If the run budget does not establish an improvement, report that result rather than claiming acceleration. GPU tests on scheduler-managed hosts must follow their reservation policy.

Acceptance criteria

• A persistent Laya backend runs on an Apple Silicon Mac through the serving API.
• Readiness follows model loading and warmup; the first inference request after readiness succeeds.
• Contract and output-parity checks cover choice, score, and noul.
• Reproducible benchmarks separate backend inference improvements from HTTP frontend overhead.
• Results show whether the predeclared latency target was met without unacceptable output differences; limitations are documented.
• Setup and benchmark instructions are published.

Related: #1. This issue covers local model execution and measured Apple Silicon acceleration; #1 covers the common frontend and transport.

[cacheline999]

Hi @hsliuustc0106, I'd like to help with this. I have an M1 Pro (16 GB) and can start with the MPS baseline.

My plan is to first get Laya running on MPS and measure where the time goes: load and warmup separately from warm requests, plus latency through the Rust frontend from #2 so we can see how much is the model vs HTTP. I'll use a fixed set of inputs covering choice, score and noul with a few different lengths.

Then I'd profile it to see whether short requests are mostly kernel launch overhead or actual compute. I've done similar profiling on CUDA inference before, so I'm curious how MPS compares. I'll try cheap fixes first (fp16, SDPA, torch.compile) and only look at MLX or Metal if the numbers point that way.

I'll share the benchmark setup here before doing the real runs. Is there a specific Laya checkpoint I should use as the reference?

[xiaoyu-xyz]

I'd like to take the CUDA backend (src/backends/cuda/), alongside @cacheline999's MPS work rather than overlapping it — same comparison, NVIDIA arm.

Plan: operator inventory + tolerance contract first (docs-only PR), then a cuBLAS + FA2 + torch.compile baseline, then hand-written kernels only for the measured hotspots. LAYA is prefill-only — no decode, no KV cache — and attention is just 2–10 % of its FLOPs (45.8 GFLOP per decision at S=84, 304.4 at S=512), so the fused GEMM epilogues matter more than attention.

Background: I wrote vLLM's one-shot Triton kernel for MoE LoRA (vllm-project/vllm#55161) and have a merged vLLM-Omni PR (#7465). I can arrange rented NVIDIA GPU access and will report SKU, versions, commands and raw results.

Two questions:

1. CUDA C++ or TileLang for the backend? LAYA already ships laya[fast] built from TileLang kernels doing roughly these operations, so I'd rather settle this in an issue than in a PR — I have a preference for CUDA C++, but either works for me.
2. Which checkpoint and input shapes are the parity reference?

[xiaoyu-xyz]

Correction to my own numbers above, after reading the checkpoint tensors rather than the config.

The MLP is gated: mlp.Wi is [5248, 1024] while intermediate_size is 2624, so the up-projection computes 2F, not F. I under-counted that GEMM. The corrected forward cost per decision is 58.5 / 96.9 / 246.6 / 381.4 GFLOP at S = 84 / 138 / 340 / 512 — about 25 % more than I quoted — and attention is 1.4 % / 7.9 % of it at the two ends, so it matters even less than I said.

Two other things worth fixing in the record while I am here:

• The checkpoint is fp16 throughout (16.0 bits/parameter, 842.6 MB), not fp32. Weights are 0.84 GB on device.
• Arithmetic intensity is 72–421 FLOP/byte over that state range, so this is firmly compute-bound: tensor-core utilisation matters far more than memory bandwidth, which is the opposite emphasis from a decode-bound LLM kernel.

I've written the shapes, the five-kernel set, and the tolerance contract up as a document and will open it as a docs-only PR — which is also where this correction belongs permanently.

[Yunaik]

Follow-up in #9: an RFC for supporting System 1 models beyond LAYA, with CLM (a Qwen3-8B encoder) as the second model. Its question 3 on Apple Silicon memory may affect the Metal backend choices here.