llama: optimize RWKV7 inference by fusing some graph operators

[MollySophia]

Summary

This PR optimizes RWKV7 inference on CUDA and Vulkan by reducing kernel-launch and dispatch overhead from small elementwise graph operations, and by specializing the Vulkan RWKV7 single-token WKV path.

The optimizations are implemented as conservative backend graph-pattern fusions. No new ggml operators are introduced, and unsupported shapes or backends continue to execute the original unfused graph.

Changes

• Specialize the RWKV7 WKV decode path.

  • CUDA: optimized WKV7 decode kernel. (Merged in CUDA : add warp-per-row WKV7 kernel for single-token decode #26111)
  • Vulkan: add a specialized wkv7_t1 shader for single-token decode.

• Fuse RWKV time-mix lerp patterns on CUDA and Vulkan:

  • SUB -> MUL -> ADD
  • SUB -> REPEAT -> MUL -> ADD
  • Computes cur + (x_prev - cur) * weight.

• Fuse ADD -> MUL patterns on CUDA and Vulkan:

  • Computes (x + y) * scale.

• Fuse affine normalization patterns:

  • CUDA: NORM -> MUL and NORM -> MUL -> ADD.
  • Vulkan: NORM -> MUL -> ADD.
  • Computes norm(x) * weight + bias.

• Fuse the RWKV key-adjust pattern on CUDA and Vulkan:

  • MUL -> MUL -> SUB -> ADD
  • Computes k + (a * (k * k_a) - (k * k_a)).

Performance

Hardware:

• GPU: NVIDIA GeForce RTX 5090
• CPU: AMD Ryzen 9 9950X
• CPU thread count: 16
• Baseline commit: 5d5cb4c3a

RWKV7 1.5B

Backend	Test	Baseline	Current	Delta
CUDA	pp512	22623.58 +/- 2414.14	25363.50 +/- 2652.65	+12.1%
CUDA	tg128	291.06 +/- 0.87	313.48 +/- 1.01	+7.7%
Vulkan	pp512	16818.33 +/- 19.49	18043.84 +/- 28.01	+7.3%
Vulkan	tg128	240.93 +/- 0.52	273.23 +/- 1.03	+13.4%

RWKV7 7.2B

Backend	Test	Baseline	Current	Delta
CUDA	pp512	8330.99 +/- 475.26	9301.38 +/- 535.99	+11.6%
CUDA	tg128	91.46 +/- 0.14	93.91 +/- 0.12	+2.7%
CUDA	tg512	91.67 +/- 0.03	94.00 +/- 0.03	+2.5%
Vulkan	pp512	7328.71 +/- 9.50	8035.70 +/- 8.95	+9.6%
Vulkan	tg128	82.93 +/- 0.03	87.16 +/- 0.02	+5.1%
Vulkan	tg512	82.98 +/- 0.04	87.16 +/- 0.05	+5.0%

Validation

• Fusion backend tests:

  • LERP
  • ADD_MUL
  • NORM_MUL_ADD
  • KEY_ADJUST

• WikiText-2 perplexity is unchanged within noise:

Model	Baseline	Current	Delta
RWKV7 1.5B F16	9.5497 +/- 0.06541	9.5499 +/- 0.06542	+0.0002
RWKV7 7.2B F16	6.4011 +/- 0.03968	6.4014 +/- 0.03969	+0.0003

[MollySophia]

Hi @MollySophia, thanks for your contribution!

Per our contribution guidelines, the automated PR checker found the following issue(s) that need your attention:

• Multiple backend changes in one PR: When adding support for a new model or feature, focus on CPU support only in the initial PR. Add support for other backends like CUDA in follow-up PRs. If you have a good reason to modify multiple backends in one PR, please explain it.
• Large PR: Large changes require prior discussion (e.g. an issue or RFC) and maintainers may not be able to review this PR as-is. Consider splitting it into smaller, focused PRs.

Please note that maintainers reserve the right to make final decisions on PRs. If you believe there is a mistake, please comment below.

Regard PR flags:

1 & 2: This PR makes changes on the existing RWKV7 path, rather than support a new model or feature. The WKV7 op has some semantic changes so that all the supported backends need to be modified at once.

[Copilot]

Pull request overview

This PR targets RWKV7 inference performance by making RWKV7 decode/time-mix semantics explicit in ggml (new ops + fused kernels/shaders), reducing graph-level operator count and backend launch/dispatch overhead across CUDA, Vulkan, SYCL, Metal, and CPU.

Changes:

• Introduce explicit RWKV ops (GGML_OP_RWKV_LERP, GGML_OP_RWKV_RK) and route RWKV7 model graph construction through them.
• Specialize/fuse RWKV7 hot paths across backends (CUDA kernels, Vulkan shaders/pipelines, SYCL kernels, Metal kernel updates) and adjust RWKV7 WKV7 parameterization (kk, a) consistently across backends.
• Add/extend backend-side graph fusions (e.g., CUDA key-update pattern; Vulkan/CPU add-mul and norm-mul-add fusions) and expand backend op test coverage.

Reviewed changes

Copilot reviewed 29 out of 29 changed files in this pull request and generated 1 comment.

Show a summary per file

File	Description
tests/test-backend-ops.cpp	Adds backend-op tests for RWKV_LERP and RWKV_RK; updates RWKV_WKV7 test inputs to new signature.
src/models/rwkv7-base.cpp	Uses new ggml RWKV ops (lerp/rk) and updates RWKV7 WKV7 call to revised inputs.
ggml/src/ggml.c	Adds new ggml ops (RWKV_LERP/RWKV_RK); updates RWKV_WKV7 op signature and metadata strings.
ggml/include/ggml.h	Exposes new RWKV APIs and updates RWKV_WKV7 function signature in the public header.
ggml/include/ggml-rpc.h	Bumps RPC patch version and updates op-count static_assert to match new ops.
ggml/src/ggml-backend-meta.cpp	Marks new RWKV ops as supported for meta backend split-state handling.
ggml/src/ggml-cpu/ops.h	Declares CPU implementations for RWKV_LERP/RWKV_RK and new fused CPU paths.
ggml/src/ggml-cpu/ops.cpp	Implements CPU RWKV_LERP/RWKV_RK plus fused ADD+MUL and NORM+MUL+ADD compute paths.
ggml/src/ggml-cpu/ggml-cpu.c	Wires new RWKV ops into CPU execution and adds CPU-side fusion detection.
ggml/src/ggml-cuda/wkv.cu	Updates CUDA RWKV_WKV7 kernel signature and adds a specialized T=1 decode kernel path.
ggml/src/ggml-cuda/rwkv.cuh	Declares CUDA fused RWKV helper ops (lerp, rk, and elementwise fusions).
ggml/src/ggml-cuda/rwkv.cu	Implements CUDA fused RWKV ops (lerp/rk) and elementwise fusion kernels.
ggml/src/ggml-cuda/norm.cuh	Declares fused CUDA norm entry points (norm+mul, norm+mul+add).
ggml/src/ggml-cuda/norm.cu	Extends CUDA norm kernel to optionally fuse mul/add, and adds fused wrapper functions.
ggml/src/ggml-cuda/ggml-cuda.cu	Dispatches new RWKV ops on CUDA and adds new fusion patterns (incl. RWKV key-update, add-mul, norm fusions).
ggml/src/ggml-vulkan/vulkan-shaders/wkv7.comp	Updates Vulkan WKV7 shader to revised (kk, a) math and inline w transform.
ggml/src/ggml-vulkan/vulkan-shaders/wkv7_t1.comp	Adds a Vulkan subgroup-based specialized WKV7 T=1 decode shader.
ggml/src/ggml-vulkan/vulkan-shaders/rwkv_rk.comp	Adds Vulkan compute shader for fused RWKV_RK.
ggml/src/ggml-vulkan/vulkan-shaders/lerp.comp	Adds Vulkan shader for RWKV_LERP.
ggml/src/ggml-vulkan/vulkan-shaders/add_mul.comp	Adds Vulkan shader for fused (a+b)*scale pattern.
ggml/src/ggml-vulkan/vulkan-shaders/norm_mul_add.comp	Adds Vulkan shader for fused norm*mul+add.
ggml/src/ggml-vulkan/vulkan-shaders/vulkan-shaders-gen.cpp	Registers new Vulkan shaders for SPIR-V generation.
ggml/src/ggml-vulkan/ggml-vulkan.cpp	Adds Vulkan pipelines/dispatch for new RWKV ops, adds WKV7 T=1 pipeline, and implements new fusion paths.
ggml/src/ggml-sycl/wkv.cpp	Updates SYCL WKV7 implementation to revised (kk, a) math and inline w transform.
ggml/src/ggml-sycl/rwkv.hpp	Declares SYCL RWKV_LERP and RWKV_RK ops.
ggml/src/ggml-sycl/rwkv.cpp	Implements SYCL RWKV_LERP and RWKV_RK kernels.
ggml/src/ggml-sycl/backend.hpp	Includes the new SYCL RWKV header into the backend.
ggml/src/ggml-sycl/ggml-sycl.cpp	Wires RWKV_LERP and RWKV_RK into SYCL compute and device support checks.
ggml/src/ggml-metal/ggml-metal.metal	Updates Metal WKV7 kernel signature and math to match revised (kk, a) formulation and inline w transform.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[am17an]

I don't think adding new ops for this architecture is worth it IMO.

[MollySophia]

I don't think adding new ops for this architecture is worth it IMO.

I tried to add some fusing rules in the first version, but I didn't find an really elegant way by myself either. Any suggestions on doing fusion instead of adding ops?

[am17an]

I'm trying to figure out a way to add ops without adding a lot of maintenance burden #24646, perhaps wait for some kind of resolution on that. I agree that fusion is not the right way for complicated patterns

[MollySophia]

I'm trying to figure out a way to add ops without adding a lot of maintenance burden #24646, perhaps wait for some kind of resolution on that. I agree that fusion is not the right way for complicated patterns

That would be nice! I'll wait for that and avoid these ops. They are indeed not really meaningful enough.

[CISC]

I'm trying to figure out a way to add ops without adding a lot of maintenance burden #24646, perhaps wait for some kind of resolution on that. I agree that fusion is not the right way for complicated patterns

That would be nice! I'll wait for that and avoid these ops. They are indeed not really meaningful enough.

I suppose this PR will be reworked?

[MollySophia]

I'm trying to figure out a way to add ops without adding a lot of maintenance burden #24646, perhaps wait for some kind of resolution on that. I agree that fusion is not the right way for complicated patterns

That would be nice! I'll wait for that and avoid these ops. They are indeed not really meaningful enough.

I suppose this PR will be reworked?

Yeah thanks. I'll check this soon

[MollySophia]

I'm trying to figure out a way to add ops without adding a lot of maintenance burden #24646, perhaps wait for some kind of resolution on that. I agree that fusion is not the right way for complicated patterns

That would be nice! I'll wait for that and avoid these ops. They are indeed not really meaningful enough.

I suppose this PR will be reworked?

Hi! Sorry for the long delay. I've reworked on this PR and:

• Reduced the complexity of these PR (e.g. only touched CUDA + Vulkan backends; removed the fusion for a pattern (previously RWKV_RK) that needs further graph rewrite)
• Rewrite previous fused ops with the new fusing mechanism on CUDA and Vulkan backends

The two failed CI tasks seem to be unrelated to this PR (?)

[CISC]

Just curious, why this change?

[MollySophia]

Previously the pattern was ADD -> (logic that computes g) -> MUL
This change should make it to be (logic that computes g) -> ADD -> MUL so that it hits the ADD -> MUL fusion pattern