-
model: add Laguna MLX support (#17237)
frostbyte_neo 发布于
2026-07-25 01:24:53 +00:00 | 24 次提交 在此版本后已推送到 main- model: add Laguna MLX support
Add Laguna XS 2, XS 2.1, and S 2.1 support to the MLX model and create paths.
Read the source config to apply one quantization policy across dense and routed MoE layers. Keep the tied output head and router at source precision, quantize supported attention and expert projections, selectively promote sensitive expert down projections, and emit per-tensor metadata for mixed quantization blobs.
Correct dense expert loading, BF16 source-layout handling, expert global-scale shapes and dtypes, routing-score scaling, and mixed-precision expert dispatch. Gate/up and down projections select quantized or dense execution independently so promoted BF16 down projections do not force quantized gate/up weights through the dense fallback.
Optimize the forward pass with compatible gate/up fusion, sorted standard GatherMM and GatherQMM operations for larger prefills, model-local mlx.Compile closures for elementwise MoE work, and cache-backed 512-token prefill chunks. This keeps the implementation on maintained MLX operations without custom kernels.
Add focused tests for Laguna configuration variants, quantization policy and metadata, dense and routed expert loading, mixed-precision dispatch, compiled-versus-eager parity, fused projections, routing, and prefill chunking.
- review comments and S 2.1 performance fixes
Address renderer/parser selection and mixed-precision expert quantization review feedback.
Keep Laguna weights resident on Metal to prevent repeated paging of its large, sparsely accessed expert buffers. Scope this policy to Laguna GPU execution.
Remove obsolete 512-token prefill chunking now that the runner's 2048-token path is faster.
-
review comments addressed
-
fix create
下载附件