metal : add MoE and SSM_CONV fusion optimizations (#28948)

* metal : add top-k MoE fusion

Adds a Metal fusion for SOFT_MAX + ARGSORT + GET_ROWS with optional
routing-weight normalization and scale, matching the top-k MoE fusion
available in the CUDA and Vulkan backends. The fused kernel writes the
selected expert ids and routing weights directly, eliding the separate
softmax, argsort, get-rows, sum-rows, clamp, div and scale kernels.

Assisted-by: pi:llama.cpp/DeepSeek-V4-Flash-Vision-Exp

* metal : add MoE weighted reduction fusion

Fuses MUL(experts, weights) plus the expert VIEW/ADD chain into one kernel
that computes the weighted sum directly. The graph_optimize hook keeps the
expert and weight buffers alive until the fused output so the allocator cannot
reuse them while the kernel is still reading them.

Assisted-by: pi:llama.cpp/DeepSeek-V4-Flash-Vision-Exp

* tests : expose MoE weighted reduction in fusion baseline

Use 2 experts per token in the generated MoE test models so the Metal
MoE weighted reduction fusion (MUL + ADD) is exercised by test-fusion.

Assisted-by: pi:llama.cpp/DeepSeek-V4-Flash-Vision-Exp

* metal : fuse RMS_NORM + SCALE

Adds NORM/RMS_NORM + SCALE fusion to the Metal backend by reusing the
norm+mul kernel with a scalar scale flag. Adds test coverage for both
NORM+SCALE and RMS_NORM+SCALE and regenerates the fusion baseline.

Assisted-by: pi:llama.cpp/DeepSeek-V4-Flash-Vision-Exp

* metal : use function constant for RMS_NORM + SCALE

Replaces the runtime use_scale karg with a Metal function constant. The
norm+mul kernel is compiled with FC_norm_use_scale=false for MUL fusion and
FC_norm_use_scale=true for SCALE fusion, so the fused kernel has no runtime
branch.

Assisted-by: pi:llama.cpp/DeepSeek-V4-Flash-Vision-Exp

* metal : use function constant for top-k MoE with_norm

Replaces the runtime with_norm karg with a Metal function constant. The
top-k MoE kernel is compiled separately for the normalized and non-normalized
routing variants, removing the runtime branch.

Assisted-by: pi:llama.cpp/DeepSeek-V4-Flash-Vision-Exp

* metal : rename moe_weighted_reduction suffix to moe_reduce

Shortens the MoE weighted-reduction fusion identifiers, kernel, pipeline,
matcher, args struct, and test op name from moe_weighted_reduction to
moe_reduce.

Assisted-by: pi:llama.cpp/DeepSeek-V4-Flash-Vision-Exp

* metal : add MUL_MAT + UNARY and MUL_MAT + ADD + UNARY fusion

Adds dense mat-vec activation fusion for sigmoid/silu and bias+softplus.
The mat-vec kernels apply the activation/bias epilogue via function
constants, avoiding the separate unary/add passes.

Assisted-by: pi:llama.cpp/DeepSeek-V4-Flash-Vision-Exp

* metal : revert MUL_MAT + UNARY and MUL_MAT + ADD + UNARY fusion

The mat-vec activation fusion regressed decode throughput on Qwen3.6-35B-A3B
by ~8% (tg32 81.5 vs 88.5 t/s). The regression is caused by loss of
concurrency: the standalone unary kernels previously overlapped with other
mat-vec work, while fusing the activation into the mat-vec kernel serializes
it on the critical path.

Assisted-by: pi:llama.cpp/DeepSeek-V4-Flash-Vision-Exp

* metal : add SSM_CONV + UNARY (silu) fusion

The SSM_CONV kernels apply silu directly via a function constant, eliding
the separate unary pass. Regenerates the fusion baseline.

Assisted-by: pi:llama.cpp/DeepSeek-V4-Flash-Vision-Exp

* metal : address fusion review comments

- Fix declaration/table alignment
- Rename top-k MoE kargs fields to val_clamp / val_scale
- Move moe-reduce alloc-deps handling into a general fusion helper
- Remove the public moe-reduce matcher API

Assisted-by: pi:llama.cpp/DeepSeek-V4-Flash-Vision-Exp

* metal : fix unused parameter in top-k MoE fusion check

Assisted-by: pi:llama.cpp/DeepSeek-V4-Flash-Vision-Exp

* metal : guard SSM_CONV fusion lookup behind use_fusion

Assisted-by: pi:llama.cpp/DeepSeek-V4-Flash-Vision-Exp

* metal : track all fused outputs in graph reorder

Assisted-by: pi:llama.cpp/DeepSeek-V4-Flash-Vision-Exp

* metal : keep top-k MoE logits alive until fused output

Assisted-by: pi:llama.cpp/DeepSeek-V4-Flash-Vision-Exp

* metal : refactor alloc deps to pattern-driven approach

Assisted-by: pi:llama.cpp/DeepSeek-V4-Flash-Vision-Exp

* metal : check fused kernel destination in concurrency tracking

Assisted-by: pi:llama.cpp/DeepSeek-V4-Flash-Vision-Exp

* meta : forward graph_optimize to underlying backends

Assisted-by: pi:llama.cpp/DeepSeek-V4-Flash-Vision-Exp

* metal : use vector for fusion table

Assisted-by: pi:llama.cpp/DeepSeek-V4-Flash-Vision-Exp

* meta : keep graph_optimize unimplemented

Assisted-by: pi:llama.cpp/DeepSeek-V4-Flash-Vision-Exp

* parallel : fix non-deterministic prompt selection

Assisted-by: pi:llama.cpp/DeepSeek-V4-Flash-Vision-Exp

* parallel : support dummy models and add global logits run hash

Assisted-by: pi:llama.cpp/DeepSeek-V4-Flash-Vision-Exp

* metal : sync cross-device copies with destination completion event

Assisted-by: pi:llama.cpp/DeepSeek-V4-Flash-Vision-Exp

* metal : avoid const_cast in fusion alloc deps

Assisted-by: pi:llama.cpp/DeepSeek-V4-Flash-Vision-Exp

* metal : skip fusions with aliased sources

Assisted-by: pi:llama.cpp/DeepSeek-V4-Flash-Vision-Exp

* metal : hide fusion pattern definition

Assisted-by: pi:llama.cpp/DeepSeek-V4-Flash-Vision-Exp

* metal : use vector fusion op sequences

Assisted-by: pi:llama.cpp/DeepSeek-V4-Flash-Vision-Exp

* metal : drop redundant struct keywords

Assisted-by: pi:llama.cpp/DeepSeek-V4-Flash-Vision-Exp

* metal : add alloc deps comment separator

Assisted-by: pi:llama.cpp/DeepSeek-V4-Flash-Vision-Exp

* metal : generalize fusion output memory ranges

Assisted-by: pi:llama.cpp/DeepSeek-V4-Flash-Vision-Exp

* metal : rename fusion out_offsets to outs

Assisted-by: pi:llama.cpp/DeepSeek-V4-Flash-Vision-Exp

* metal : avoid dst vector in memory range check

Assisted-by: pi:llama.cpp/DeepSeek-V4-Flash-Vision-Exp

* metal : optimize fusion matching and multi-output handling

- use pointer arithmetic for fusion info count lookup
- avoid heap allocations in top-k MoE and MoE reduce pattern matchers
- use fusion outs for multi-output subgraph checks

Assisted-by: pi:llama.cpp/DeepSeek-V4-Flash-Vision-Exp

* Revert "parallel : support dummy models and add global logits run hash"

This reverts commit 57c7caf941c1b43c270fd5009c9f175063522e96.

* fusion : update MTL.csv

* metal : unroll constant loops in top-k MoE kernel

Assisted-by: pi:llama.cpp/DeepSeek-V4-Flash-Vision-Exp

* metal : use function constants for top-k MoE n_expert and top_k

Assisted-by: pi:llama.cpp/DeepSeek-V4-Flash-Vision-Exp

* metal : rename fusion kargs to scale and clamp

Assisted-by: pi:llama.cpp/DeepSeek-V4-Flash-Vision-Exp

* metal : use function constants for moe_reduce and ssm_conv

Assisted-by: pi:llama.cpp/DeepSeek-V4-Flash-Vision-Exp

* fusion : update MTL.csv
This commit is contained in:
Georgi Gerganov
2026-09-19 13:14:44 +03:00
committed by GitHub
parent 60b06ab9a9
commit 5b59b83f4e
17 changed files with 1419 additions and 175 deletions
+5 -4
View File
@@ -11,6 +11,7 @@
#include <clocale>
#include <cmath>
#include <cstdio>
#include <random>
#include <string>
#include <vector>
#include <ctime>
@@ -156,7 +157,7 @@ static std::vector<std::string> split_string(const std::string& input, char deli
int main(int argc, char ** argv) {
std::setlocale(LC_NUMERIC, "C");
srand(1234);
std::mt19937 rng(1234);
common_params params;
@@ -321,7 +322,7 @@ int main(int argc, char ** argv) {
client.t_start_prompt = ggml_time_us();
client.t_start_gen = 0;
client.input = k_prompts[rand() % k_prompts.size()];
client.input = k_prompts[rng() % k_prompts.size()];
client.response = "";
// construct the prompt:
@@ -334,10 +335,10 @@ int main(int argc, char ** argv) {
client.prompt += k_system;
}
const int n_junk_cur = rand() % n_junk;
const int n_junk_cur = rng() % n_junk;
for (int i = 0; i < n_junk_cur; ++i) {
const int r = rand() % k_questions.size();
const int r = rng() % k_questions.size();
client.prompt += "User:\n" + k_questions[r] + "\nAssistant:\n " + k_answers[r] + "\n";
}
client.prompt += "User:\n" + client.input + "\nAssistant:\n";