Skip to content

feat(cuda): fuse activations and residual add into mmv f/q epilogues#67

Draft
roberteg16 wants to merge 1 commit into
gfx11from
mmv-fusion-epilogues
Draft

feat(cuda): fuse activations and residual add into mmv f/q epilogues#67
roberteg16 wants to merge 1 commit into
gfx11from
mmv-fusion-epilogues

Conversation

@roberteg16

@roberteg16 roberteg16 commented Jul 24, 2026

Copy link
Copy Markdown

Summary

Extends the mul_mat_vec_f / mul_mat_vec_q fusion path to fold several post-matmul patterns into the GEMV epilogue on the decode path (ncols_dst == 1). Each of these previously ran as its own kernel launch. The change is additive and only touches the non-gated branch of the epilogue; the existing GLU / gate fusion path is untouched.

What changed

New fused patterns, detected in ggml_cuda_try_fuse and executed inside the mmv f/q kernels:

  • Non-gated unary activation: mul_mat -> [reshape] -> sigmoid|silu
  • Activation + elementwise mul: mul_mat -> [reshape] -> silu -> mul, carried by a new per-row activation_mul operand; tried before the standalone unary_mul fusion so the whole silu(x) * y collapses into the matmul instead of leaving the mul exposed.
  • Residual add through a view: mul_mat -> reshape -> add, the common case where a projection output is reshaped before its residual add (e.g. the GDN/linear-attention layers), which the existing {MUL_MAT, ADD} fusion could not match because it requires the ops to be directly consecutive.
    Implementation notes:
  • Reuses the existing ggml_cuda_mm_fusion_args_{host,device} plumbing, adding an activation enum (GGML_UNARY_OP_COUNT = none) and an activation_mul operand. The activation is dispatched at runtime inside the existing has_fusion path, so no new template instantiations are added.
  • When a reshape sits between the matmul and the fused tail, detection writes the result using the matmul geometry into the final node's buffer (the intermediate reshape/activation/mul/add nodes are elided).
  • A targeted overlap check replaces the generic memory-range check for these tails: the destination may alias the mul/add operand (an in-place op reads and writes the same element index), but must not alias the matmul inputs (read across the full reduction). Same-shape + contiguous operands are required.

Why

On the decode path these tiny epilogue ops are individual kernel launches whose latency is only partially hidden. Folding them removes launches and, for the n_embd-sized residual add, an extra HBM round-trip per layer.

Measurements

Qwen3.6-35B-A3B, gfx1151 (Strix Halo), Q4_K_M, llama-bench -pg 0,128 -r 5 (decode):

build tg128 (t/s)
baseline 58.73
+ activation fold 58.90
+ activation*mul fold 58.95
+ residual-add fold 59.53
Kernel launches per generated token drop from 1263 to 1103 (161724 -> 141244 total over 128 tokens), i.e. 160 fewer launches/token (~13%):
folded op baseline final per token
SIGMOID (beta + shared-expert gate) 10240 1280 -70
ADD (attn residual, incl. reshaped GDN path) 17920 10240 -60
SILU (silu(z)*norm, absorbed unary_mul) 3840 0 -30
total kernels 161724 141244 -160
Net MUL count is unchanged: the mul after the silu is folded together with the activation rather than left standalone. The MoE ffn_out residual adds (~60/token) are not folded by design.
Roofline per-op attribution (decode, vs baseline, fusion enabled on both):
  • beta_sigmoid + shared_expert_gate_sigmoid: folded into their matmuls.
  • gated-norm silu(z) * norm: folded (no leftover standalone mul).
  • attn_residual adds: fully folded (including the reshaped GDN-layer path).
  • ffn_out (MoE) residual adds are left as-is by design: their first operand is the expert-combine output, not a single GEMV, so there is no matmul to fold into.
    Output verified unchanged (greedy generation coherent).

Extend the mul_mat_vec_f/q fusion path to fold, into the GEMV epilogue,
patterns that previously ran as separate kernels on the decode path:

- non-gated unary activation: mul_mat -> [reshape] -> sigmoid|silu
- activation followed by an elementwise mul: mul_mat -> [reshape] -> silu -> mul
  (a per-row activation_mul operand; wins over the standalone unary_mul fusion)
- residual add through a view: mul_mat -> reshape -> add

The fusions reuse the existing mm_fusion_args plumbing (adding an activation
enum and an activation_mul operand) and only touch the non-gated branch of the
epilogue, leaving the GLU/gate path unchanged. Detection writes the fused result
using the matmul geometry into the final node's buffer; a targeted overlap check
allows aliasing the mul/add operand (in-place, same-index) while forbidding
aliasing the matmul inputs.

On Qwen3.6-35B-A3B (gfx1151, Q4_K_M) decode this folds the beta/shared-expert
sigmoids, the gated-norm silu*x, and all attn residual adds into their matmuls;
tg128 ~58.7 -> ~59.5 t/s with output unchanged.

Co-Authored-By: Claude Opus 4 (1M context) <[email protected]>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant