denseflux: Bỏ qua phần lớn FFN của một LLM dense trên Apple Silicon

Tóm tắt

Phiên bản ngắn

• Vì sao là FFN: decode một luồng bị giới hạn bởi băng thông bộ nhớ, và FFN chiếm khoảng hai phần ba trọng số của một lớp transformer dense. Mỗi hàng FFN không phải đọc là thời gian tiết kiệm được.
• Kết quả: trên một mô hình thưa kiểu ReLU (ProSparse-Llama-2-7B), mỗi token chỉ có 11,2% neuron FFN hoạt động. Bỏ qua phần còn lại bằng gating chính xác làm perplexity tăng +0,44% và giúp FFN nhanh hơn khoảng 2,05×.
• Kernel: kernel Metal của chúng tôi cho FFN 4-bit ở kích thước Qwen2.5-3B hòa vốn ở khoảng 60% neuron hoạt động và đạt 1,93× ở 10% hoạt động trên M4 Pro, khoảng 77% mức tăng tốc lý tưởng suy ra từ số byte tiết kiệm được.
• Điều không hiệu quả: các mô hình SwiGLU thông thường (Qwen2.5) không phù hợp: kích hoạt hiếm khi bằng 0, nên cắt quá ~17% neuron đã làm giảm chất lượng. Gom neuron thành các cụm cố định thất bại trên mọi mô hình chúng tôi thử, và các bộ dự đoán rẻ thất bại ở độ thưa vừa phải.
• Phạm vi: bài này báo cáo số liệu theo từng lớp FFN. Token mỗi giây đầu-cuối, benchmark suy luận và so sánh với baseline sẽ có sau (mục 7).

1 · Động lực

Vì sao là FFN

Khi một LLM sinh từng token một, GPU phải đọc mọi trọng số từ bộ nhớ một lần cho mỗi token nhưng làm rất ít phép tính trên mỗi byte. Vì vậy thời gian decode gần như tỉ lệ thuận với số byte phải đọc. Đây là chế độ mà TEAL, Deja Vu và LLM in a Flash đều nhắm tới.

Một FFN có gate hiện đại tính

Mỗi “neuron” trung gian $i$ sở hữu hàng $i$ của $W_{\mathrm{gate}}$, hàng $i$ của $W_{\mathrm{up}}$ và cột $i$ của $W_{\mathrm{down}}$. Với Llama-2-7B (d_model = 4096, d_ff = 11008), mỗi lớp có khoảng 67M tham số attention (4 × 4096²) và khoảng 135M tham số FFN (3 × 4096 × 11008), nên FFN chiếm khoảng hai phần ba số byte. Các mô hình dùng grouped-query attention còn có tỉ trọng FFN lớn hơn nữa.

Nếu kích hoạt của một neuron bằng 0 với token hiện tại, hàng $W_{\mathrm{up}}$ và cột $W_{\mathrm{down}}$ của nó không đóng góp gì và không cần đọc. Bài viết này trả lời ba câu hỏi: bỏ qua được bao nhiêu neuron, với cái giá chất lượng nào, và một kernel thật trên Apple Silicon có biến điều đó thành tốc độ thực tế được không?

2 · Công trình trước

Những công trình làm nền

3 · Thiết lập

Cách chúng tôi đo

• Phần cứng: Apple M4 Pro cho benchmark kernel. Mọi thí nghiệm Python chạy trên Apple Silicon (MPS).
• Mô hình: Qwen2.5-0.5B và Qwen2.5-3B (SwiGLU); một cặp ReGLU/SwiGLU ~1,3B huấn luyện giống hệt nhau trừ hàm kích hoạt; ProSparse-Llama-2-7B (đã chuyển sang ReLU).
• Độ chính xác số: trọng số 4-bit trừ khi ghi khác.
• Thước đo chất lượng: perplexity trên 16.384 token đánh giá. Prefill giữ dense; độ thưa được áp dụng theo từng token, đúng như lúc decode.
• Thước đo chi phí: tỉ lệ công việc FFN (hoặc byte FFN) giữ lại cho mỗi token. Phép chiếu gate được tính đủ khi nó được tính chính xác, nên gating chính xác giữ lại $\frac{1+2a}{3}$ của FFN, với $a$ là tỉ lệ neuron hoạt động.
• Công cụ: một harness Python (cluster_ffn.py) bọc mọi MLP có gate, và một micro-benchmark Swift/Metal (ffnbench.swift) có đầu ra kernel khớp với bản tham chiếu CPU trong phạm vi 0,04%.

4 · Kết quả âm

Những gì không hiệu quả

Các kết quả âm đã định hình thiết kế cuối cùng, nên chúng đứng trước.

4.1 Kích hoạt SwiGLU là dense

Trên Qwen2.5-0.5B chúng tôi đặt ngưỡng cho $|\mathrm{SiLU}(W_{\mathrm{gate}}\cdot x)|$ theo từng neuron ($k=1$, trường hợp thuận lợi nhất để bỏ qua):

Chỉ khoảng 17% neuron có thể bỏ qua (khoảng 11% tính toán FFN) trước khi perplexity tăng 1,4%. Đầu ra SiLU nhỏ nhưng gần như không bao giờ bằng 0, nên mỗi lần bỏ qua là một phép xấp xỉ. Điều này khớp với nhận xét của PowerInfer-2 rằng mô hình SiLU chỉ cho lợi ích khiêm tốn.

Một ngân sách riêng cho từng lớp giúp ích nhiều trên SwiGLU. Trên Qwen2.5-3B, giữ 86,7% neuron làm perplexity tăng +1,39% với ngân sách theo lớp, so với +11,5% với một ngưỡng chung. Dù vậy, 86,7% hoạt động vẫn cao hơn nhiều so với điểm hòa vốn của kernel (mục 5.3), nên trên thực tế mô hình SwiGLU không được lợi gì.

4.2 Cụm cố định gần như luôn hoạt động

Để đọc trọng số theo các khối liền mạch lớn hơn, chúng tôi gom neuron thành các cụm $k$ neuron và tính cả cụm bất cứ khi nào có một neuron trong đó hoạt động, sau khi sắp xếp lại neuron theo mức đồng kích hoạt (một chuỗi tham lam theo kiểu Ripple).

• Trên Qwen2.5-0.5B, 99,4–100% số cụm hoạt động với $k$ = 16, 32 và 64 ở mọi ngưỡng tới 0,2. Sắp xếp theo chuỗi chỉ đưa con số này từ 99,97% xuống 99,42% (ngưỡng 0,2, $k=16$).
• Trên mô hình ReGLU ~1,3B (38,8% neuron hoạt động), 88,7% số cụm $k=8$ hoạt động, so với khoảng 98% khi gom ngẫu nhiên.

Lý do nằm ở phép tính. Nếu một tỉ lệ $a$ neuron hoạt động và chúng độc lập với nhau, một cụm $k$ neuron cần được tính với xác suất

Với $a=0,4$ và $k=8$ con số này đã trên 98%. Cụm chỉ có lợi khi phần lớn neuron bằng 0 và các neuron hoạt động tụ lại với nhau, điều mà chúng tôi không quan sát thấy.

4.3 Bộ dự đoán rẻ thất bại ở độ thưa vừa phải

Gating chính xác cần đọc trọn $W_{\mathrm{gate}}$, cố định một phần ba FFN. Một bộ dự đoán có thể bỏ qua cả phần đó. Trên mô hình ReGLU ~1,3B (38,8% hoạt động, $k=1$):

• layer_in có recall tốt nhất, nhưng nó đọc toàn bộ gate nên không tiết kiệm tính toán trong RAM; nó chỉ giúp prefetch từ SSD. Các lần bỏ sót tập trung ở những lớp đầu (recall 60% ở lớp 0, ~79% ở lớp 1–2, 99,6% ở các lớp cuối), nơi attention thay đổi đầu vào nhiều nhất.
• prev_token cho thấy tập neuron hoạt động thay đổi nhiều giữa các token, nên cache xuyên token (windowing) giúp ít trên mô hình này.
• SparseInfer báo precision trên 99% với các mô hình ProSparse thưa ~89%. Ở độ thưa ~61%, nhiều neuron nằm sát ngưỡng cắt, và tăng α chỉ khôi phục recall bằng cách đọc gần như mọi thứ.

Khi các lần bỏ sót bị bỏ qua luôn, chất lượng giảm:

Không bộ dự đoán nào thắng gating chính xác trên mô hình này. Bộ dự đoán low-rank chưa huấn luyện tốt nhất tiết kiệm khoảng 9 điểm phần trăm byte với perplexity +3,8%. Phiên bản trong bài báo SVDP thêm whitening và hiệu chỉnh bias tham lam theo từng neuron, và được đánh giá trên mô hình thưa ~89%; thử lại đầy đủ công thức đó trên ProSparse là một trong các bước tiếp theo.

4.4 Không có phân tách nóng/lạnh rõ rệt

Giữ các neuron dùng thường xuyên nhất trong RAM chỉ làm giảm khiêm tốn tần suất kích hoạt của phần còn lại:

5 · Kết quả

Những gì hiệu quả

5.1 Độ thưa ReLU là miễn phí

Với cặp ~1,3B huấn luyện giống hệt nhau trừ hàm kích hoạt, mô hình ReGLU chỉ có 38,8% neuron FFN khác 0; bản SwiGLU có 100%. Bỏ qua các giá trị đúng bằng 0 loại bỏ 61% công việc $W_{\mathrm{up}}$/$W_{\mathrm{down}}$ mà không mất chất lượng, giữ lại 59,2% tính toán FFN. Tỉ lệ khác 0 tăng theo độ sâu, từ 17,6% ở lớp 0 tới 53,6% ở lớp cuối.

Cái giá nằm ở chính chất lượng mô hình: perplexity gốc của mô hình ReGLU là 21,78 so với 20,97 của SwiGLU, kém khoảng 4%. Điều này khớp với phát hiện của Sparsing Law rằng hàm kích hoạt, chứ không phải số tham số, quyết định giới hạn độ thưa.

5.2 ProSparse-Llama-2-7B: 11% hoạt động

Trên ProSparse-Llama-2-7B, mỗi token chỉ có 11,2% neuron FFN hoạt động, nhất quán với độ thưa 89,32% được báo trong bài báo ProSparse. Gating chính xác làm perplexity thay đổi +0,44%, và FFN chạy nhanh hơn khoảng 2,05× với kernel của chúng tôi.

Vấn đề còn mở: nếu gating chỉ bỏ qua các giá trị đúng bằng 0, đầu ra phải khớp với bản dense, chỉ sai lệch do số dấu phẩy động. Nguồn gốc của mức +0,44% (ngưỡng kích hoạt trong mô hình, lượng tử hóa 4-bit, hay cách thiết lập đánh giá) chưa được tách riêng.

5.3 Một kernel Metal biến độ thưa thành tốc độ

Chúng tôi đo thời gian một lớp FFN 4-bit ở kích thước Qwen2.5-3B trên M4 Pro, với trọng số ngẫu nhiên (chỉ thời gian là quan trọng). Bản dense chạy 176 µs mỗi lớp, khoảng 216 GB/s. Nếu chỉ một tỉ lệ $a$ neuron hoạt động, số byte phải đọc giảm còn $\frac{1+2a}{3}$ so với dense, và đó là trần tăng tốc:

• Điểm hòa vốn ở khoảng 60% hoạt động. Ở 100% hoạt động, kernel thưa chậm hơn dense (0,79×) vì chi phí quản lý. Qwen 3B ở 87% hoạt động không được lợi gì; mô hình ReGLU ở ~39% sẽ nhanh hơn khoảng 1,3× mỗi lớp FFN.
• Indirect dispatch là thay đổi then chốt. Chỉ khởi chạy bước up cho các hàng hoạt động đã giảm nó từ 36 µs xuống 10 µs ở 10% hoạt động, nâng FFN 3B từ 1,56× lên 1,93×.
• Gate giờ là mức sàn. Ở 10% hoạt động nó chiếm 59 trong 91 µs vì mọi hàng đều phải đọc. Bước down và phép cộng cuối mất tổng cộng khoảng 19 µs và vẫn còn có thể tinh chỉnh.

5.4 Lớp càng lớn càng gần mức lý tưởng

Ở 10% hoạt động, tăng tốc FFN tăng theo kích thước lớp: 1,46× ở kích thước 1B, 1,93× ở 3B và 2,11× ở 7B. Chi phí cố định mỗi lần dispatch ít quan trọng hơn khi mỗi lớp di chuyển nhiều byte hơn.

5.5 Các đoạn liền mạch thay cho cụm cố định

Cho đường SSD về sau, chúng tôi thay cụm cố định bằng các đoạn liền mạch dựng lúc chạy từ tập neuron thực sự hoạt động, gộp các lần đọc kề nhau khi đọc luôn khoảng trống rẻ hơn một yêu cầu mới. Cách này giảm tỉ lệ byte FFN phải đọc từ 41% xuống 27,6%. Quy tắc dựa trên mô hình độ trễ–băng thông,

theo đó hai đoạn được gộp khi khoảng trống nhỏ hơn khoảng $L\cdot B$ byte. $L$ và $B$ phải được đo trên thiết bị đích thay vì giả định.

6 · Giới hạn

Những gì các con số này chưa cho thấy

• Theo lớp, không phải đầu-cuối. Mọi mức tăng tốc chỉ áp dụng cho FFN. Attention, embedding và LM head không đổi. Ước lượng thô: nếu FFN chiếm hai phần ba thời gian decode và chạy nhanh gấp 2×, cả mô hình nhanh hơn khoảng $\frac{1}{\frac{1}{3}+\frac{1}{3}}=1,5\times$. Đây là ước lượng, không phải số đo.
• Một máy. Thời gian kernel đo trên M4 Pro. Mac mini M4 bản thường có băng thông bộ nhớ khác và phải được đo riêng.
• Chỉ có perplexity. TEAL và Sirius cho thấy toán và sinh văn bản dài suy giảm trước khi perplexity thay đổi. Benchmark suy luận chưa được chạy.
• Trọng số ngẫu nhiên trong benchmark kernel. Thời gian không phụ thuộc vào giá trị, nhưng vẫn cần chạy đầu-cuối với mô hình thật.
• Mô hình ReLU tốn chất lượng. Cần mô hình chuyển sang ReLU hoặc huấn luyện với ReLU để có lợi lớn, và chúng có thể yếu hơn bản SwiGLU tương ứng (khoảng 4% perplexity trong cặp ~1,3B của chúng tôi).

7 · Tiếp theo

Các bước tiếp theo

• Đo token mỗi giây đầu-cuối với kernel nằm trong một vòng decode hoàn chỉnh, dense so với thưa, trên ProSparse-Llama-2-7B và ProSparse-MiniCPM-1B.
• Đo tỉ trọng thời gian decode của FFN trên Mac mini M4 bản thường.
• So sánh với llama.cpp và MLX trên cùng một máy.
• Chạy một benchmark suy luận (GSM8K hoặc tương tự) song song với perplexity.
• Giải thích mức thay đổi perplexity +0,44% dưới gating chính xác.
• Thử lại bộ dự đoán SVDP đầy đủ (whitening cộng hiệu chỉnh bias) trên ProSparse, nơi gate chiếm phần lớn chi phí còn lại.
• Thử một tập lõi theo từng prompt, nén một lần sau prefill, cộng với bộ dự đoán theo từng token cho các neuron nằm ngoài lõi.

Tài liệu tham khảo

Nguồn

1. Liu et al., Training-Free Activation Sparsity in Large Language Models (TEAL), ICLR 2025. arxiv.org/abs/2408.14690
2. Lee et al., CATS: Contextually-Aware Thresholding for Sparsity in LLMs, COLM 2024. arxiv.org/abs/2404.08763
3. Liu et al., Deja Vu: Contextual Sparsity for Efficient LLMs at Inference Time, ICML 2023. arxiv.org/abs/2310.17157
4. Alizadeh et al., LLM in a Flash: Efficient LLM Inference with Limited Memory. arxiv.org/abs/2312.11514
5. Song et al., PowerInfer: Fast LLM Serving with a Consumer-grade GPU. arxiv.org/abs/2312.12456
6. Xue et al., PowerInfer-2: Fast LLM Inference on a Smartphone. arxiv.org/abs/2406.06282
7. Song et al., ProSparse: Introducing and Enhancing Intrinsic Activation Sparsity within LLMs. arxiv.org/abs/2402.13516
8. Luo et al., Sparsing Law: Towards LLMs with Greater Activation Sparsity, ICML 2025. arxiv.org/abs/2411.02335
9. SVD Contextual Sparsity Predictors for Fast LLM Inference. arxiv.org/abs/2603.14110
10. SparseInfer: Training-free Prediction of Activation Sparsity for Fast LLM Inference. arxiv.org/abs/2411.12692
11. Dong et al., Prompt-prompted Adaptive Structured Pruning for Efficient LLM Generation (GRIFFIN). arxiv.org/abs/2404.01365
12. CoreInfer: Accelerating LLM Inference with Semantics-Inspired Adaptive Sparse Activation. arxiv.org/abs/2410.18311
13. Ripple / Neuralink: Fast LLM Inference on Smartphones with Neuron Co-Activation Linking. arxiv.org/abs/2410.19274
14. Zhou et al., Sirius: Contextual Sparsity with Correction for Efficient LLMs, NeurIPS 2024. arxiv.org/abs/2409.03856