CUDA : add warp-per-row WKV7 kernel for single-token decode

[123123213weqw]

Overview

Follow-up to #26080.

This one speeds up the T/B == 1, C/H == 64 case of WKV7 — basically single-token decode on head size 64 models.

Right now each state row is handled by a single thread that keeps its own float state[64] and does the two 64-wide reductions serially. That's a pretty heavy thread, and there's no parallelism in the reduction. So I wrote a new kernel (rwkv_wkv7_f32_t1_warp_row) where a whole warp works on one row — each lane takes care of two columns, and reductions just use warp shuffles.

If the shape doesn't match, it falls back to the existing kernel, nothing changes there.

Additional information

How it's launched:

constexpr int rows_per_block = 4;
block = (32, 4);     // WARP_SIZE x rows_per_block
grid  = (B * H, 16); // CUDA_WKV_BLOCK_SIZE / rows_per_block

The kernel is templated on rows_per_block and carries launch_bounds(WARP_SIZE * rows_per_block, 2).

I tried rows_per_block = 2, 4 and 8 on my 4080 and got 293.1 / 293.3 / 293.4 tok/s, so it's basically noise — stuck with 4.

What changed:

• ggml/src/ggml-cuda/wkv.cu — the new kernel and the dispatch branch
• tests/test-backend-ops.cpp — added wkv7(F32, H=32, D=64, T=1, B=4) since multi-seq decode wasn't covered

Checked against the CPU backend with F32 inputs (NMSE < 1e-7):

head_count=32, head_size=64, n_seq_tokens=1,   n_seqs=1   PASS
head_count=32, head_size=64, n_seq_tokens=1,   n_seqs=4   PASS
head_count=32, head_size=64, n_seq_tokens=32,  n_seqs=1   PASS
head_count=32, head_size=64, n_seq_tokens=32,  n_seqs=4   PASS
head_count=32, head_size=64, n_seq_tokens=128, n_seqs=4   PASS
5/5 passed

The T=32/T=128 ones go through the old kernel, so the fallback path gets a bit of coverage too.

Benchmarks on RWKV7 1.5B Q4_K_M (gen 128, 10 runs, ngl 999, no prompt):

GPU	CUDA	master	this PR
V100 32GB	12.8.93	146.1 tok/s	151.9 tok/s	+4.0%
RTX 4080	13.0.88	283.2 tok/s	293.3 tok/s	+3.5%

HIP shares this .cu file but I don't have any AMD hardware, so no ROCm numbers unfortunately.

Requirements

• I have read and agree with the contributing guidelines
• AI usage disclosure: YES — I used AI to poke around the codebase and sketch early versions of the kernel. The final code and all the benchmarks are mine.

[ggml-gh-bot]

Hi @123123213weqw, thanks for your contribution!

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

• PR Template not respected: Please respect the template when creating a new pull request. Make sure to fill out all required sections.

---

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

[123123213weqw]

Hi @123123213weqw, thanks for your contribution!

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

• PR Template not respected: Please respect the template when creating a new pull request. Make sure to fill out all required sections.

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

done