Qwen3.5 CUDA long-context DECODE is slow: the decode attention kernel runs 16 blocks — split-K needed for consumer use (RAH ~10 tok/s at 16k)

From aprender-3e's #3596 design note (issuecomment-5763674795): after batched prefill (#3596), long-context DECODE stays slow because the decode kernel runs only 16 blocks. Consumers need usable decode at their measured context (#3716: RAH p50 ≈ 20k, p90 ≈ 60k tokens; needs ~2.5k output tokens per verdict). Operator bar: "enough context to use if for RAH, and arbiter and for rmedia work".

done_when

a split-K (or flash-decoding) decode attention for the Qwen3.5 full-attention layers on lambda and gx10, parity-checked against the current kernel at every rung (4k … 262,144)
decode tok/s recorded at 4k/20k/60k/148k/262k on both hosts for the 9B (and every size that fits), against pinned llama.cpp on the same prompt (ratio reported, never claimed as a target)
oxide-first where practical (a cuda-oxide kernel + CPU twin + differential test on sm_89 and sm_121)

Scope note (cop, 2026-09-21): if f16 KV storage for the Qwen3.5 full-attention layers needs decode-kernel changes, it belongs HERE, not in #3596. f16 KV is ALLOWED for 0.69.1 (the #3596 amendment, item 3: parity vs f32 measured per rung, printed in --json), and it's what lets the 9B reach 262,144 on the 24 GB 4090. q8/q4 KV quantization is NOT in 0.69.1 (operator chose option (a) over (b)). This row is currently UNOWNED and needs a worker.

f16 KV split (agreed with aprender-3e, 2026-09-21): f16 KV is a storage format touched at four sites: (1) prefill K/V row writes, (2) prefill attention reads, (3) per-token decode row writes, (4) decode attention reads. #3596 owns the storage type (the f32|f16 tag on Qwen35CudaState) plus sites 1–3 and the capacity check that picks f32 / f16 / refuse. #3725 owns site 4 (the split-K decode kernel reads f16 as well as f32). Until #3725 lands, #3596 ships the f16 branch wired but REFUSING, with the reason "f16 KV decode read not available (#3725)"; no path serves f16 half-done. f16 vs f32 parity is measured on gx10 once both land.

Design note: split-K (flash-decoding) decode attention for the Qwen3.5 full-attention layers (aprender-3e [8f56c1], 2026-09-21)

Assigned by aprender-04 (cop). This row is site 4 of the f16 KV split (cop comment above): the decode attention read, for f32 and f16 KV. [44388e] (#3596) owns the storage type, sites 1–3, and the capacity check. Written against origin/main @ 52f43da71. Ticket PMAT-3725. Nothing below is measured unless it says measured.

1. Why decode is slow at long context (read from the code)

DecodeAttention256Kernel (aprender-gpu/src/kernels/gdn/decode_attention.rs) runs with grid = num_heads and block = 256.

16 blocks on the 9B, on 128 SMs (4090) or 48 SMs (GB10). Each block walks all L positions by itself.
Scores are uncoalesced. Thread t computes the whole 256-wide dot for position t, one float at a time, so adjacent threads read addresses num_kv_heads·256·4 bytes apart (4 KiB on the 9B). Every load instruction touches 32 different cache lines.
Each KV head is read group_size times (4 on 0.8B–9B, 6 on 27B), once per query head, because the grid is per query head.
Every 4096-position pass costs three bar.sync plus two block reductions.

Measured geometry (read from the five GGUF headers on lambda): head_dim = 256 everywhere; heads / KV heads are 8/2 (0.8B, 2B), 16/4 (4B, 9B) and 24/4 (27B); one layer in four is full attention (6, 6, 8, 8, 16 layers); context_length = 262144 in all five.

2. The design: two kernels, new file `gdn/decode_attention_splitk.rs`

Kernel A: gdn_decode_attention_splitk. Grid (num_kv_heads, n_splits), block 256 (8 warps).

One block per (KV head, split), serving ALL G query heads of its group. K and V are read once per token, not G times. G is a compile-time constant (4 or 6).
Warp-cooperative, coalesced rows. Warp w takes positions start+w, start+w+8, … in its split. Lane l holds row elements l, l+32, …, l+224, so every load instruction reads one contiguous 128 B (f32) or 64 B (f16) span.
Scores: G dots per position, reduced with shfl.down, then broadcast. The score is divided by sqrt(head_dim), the same division the CPU reference uses (not a reciprocal multiply).
Online softmax per (warp, head) in registers (running m, s, acc[8]), then an in-block merge of the 8 warps through 8 KiB of shared memory, one head at a time.
Output: the block writes one partial (m, s, acc[256]) per (query head, split) to a scratch buffer. acc is left un-normalised.

Kernel B: gdn_decode_attention_splitk_reduce. Grid num_heads, block 256, one thread per output element.

Computes M = max_i m_i, then out = Σ_i e^{m_i−M}·acc_i / Σ_i e^{m_i−M}·s_i.
The e^{m_i−M} rescale is flash-decoding's standard log-sum-exp merge.

Split count is a runtime parameter, chosen on the host from L:

chunk = max(64, ⌈L / S_target⌉) and n_splits = ⌈L / chunk⌉, so no split is ever empty.
S_target ≈ 4·SMs / num_kv_heads, which gives ~512 blocks on the 4090.
One compiled module serves the whole decode, as today.
The partial scratch is num_heads × max_splits × (256 + 2) floats (~4 MiB on the 9B), allocated once per model in Qwen35AttnScratch.

f16 KV: the same kernel with the storage type as a compile-time parameter (two modules). It loads with ld.global.b16 and converts with cvt.f32.f16. All arithmetic stays fp32; only storage is f16.

Bandwidth arithmetic (an estimate, not a measurement):

K+V bytes per token = L · 2 · kv_heads · 256 · bytes · attn_layers.
On the 9B at 262,144 that is 16 GiB (f32) or 8 GiB (f16) per token, a floor of ~16 ms or ~8 ms on a ~1 TB/s 4090.
At 20k it is ~1.3 GiB (f32) or ~0.6 GiB (f16).
So at long context the attention read is the decode cost. Serving ≥ 148k on the 4090 is an f16 question, not only a kernel question.

3. Interface with #3596 (for [44388e])

New executor fn: gdn_decode_attention_splitk_into(q, k_cache_ptr, v_cache_ptr, kv: KvStorage::{F32, F16}, partials, out, num_heads, num_kv_heads, head_dim, seq_len), in a NEW file cuda/executor/gdn_decode_splitk_ops.rs.
It takes raw device pointers plus a dtype tag, so it does not care whether Time-to-first-token on Qwen3.5-4B-Q4_K_M is 14-16s vs llama.cpp's 1.8-2.9s on the same box (7.7x); apr's decode rate is unmeasurable with shipped tooling #3596 stores f16 as GpuBuffer<u16> or GpuBuffer<half::f16>.
I own: the one decode call in forward_qwen35_cuda.rs (today's gdn_decode_attention_into(…, pos32 + 1)) and one partials field in Qwen35AttnScratch.
[44388e] owns: Qwen35CudaState and the f32|f16 tag. When both land, their refusal "f16 KV decode read not available (#3725)" becomes a call to this fn with KvStorage::F16.
decode_attention.rs is NOT modified or deleted. It stays as the parity reference (c3 is adding a test there, Qwen3.5-9B Q4_K_M on CUDA at ~12k-token context corrupts identifiers it quotes from the prompt (as_millis→as_mi, Backend→Backen) and runs at 1.0 tok/s; parity passes only at 79 positions — model weakness or long-context CUDA divergence is UNMEASURED #3693).

4. Verification

CPU twin: decode_attention_splitk_cpu in the same file, an f32 port with the kernel's own split/merge order, plus the existing verbatim forward_attention port already in decode_attention.rs's tests as the reference.
Unit device tests (both archs; they skip without a device):
- twin vs device, and device vs the current kernel, at L ∈ {1, 5, 37, 64, 65, 4096, 4097, 8192}
- geometries 8/2, 16/4 and 24/4
- a forced-tiny chunk, so the multi-split merge path runs at small L
- a split boundary that is not a multiple of 8 warps
- an f16 case: the new kernel reads f16 KV, and the current kernel reads the same f16 values widened to f32. That isolates the kernel's read from storage rounding. f16-vs-f32 storage parity is [44388e]'s per-rung measurement.
The rungs (the done_when's "4k … 262,144"): an example binary gdn_splitk_rungs.
- Rungs L ∈ {4096, 8192, 20000, 60000, 148000, 262144}, geometries 16/4 (9B) and 24/4 (27B), f32 and f16.
- Each rung records parity against the current kernel (max |Δ|, cosine) and µs per call for both kernels (CUDA events, median of N).
- It writes evidence/section-3725/<host>.json on lambda (sm_89) and gx10 (sm_121), under flock /tmp/apr-gpu.lock choom -n 1000, with GPU occupancy recorded at start.
- The 262k f32 fixture is 2 GiB of K+V, which is why this runs as an example and not a lib test.
End-to-end decode tok/s on the 9B at 4k/20k/60k/148k/262k vs pinned llama.cpp, ratio reported and never claimed:
- This needs a real prompt prefilled to L, so it depends on Time-to-first-token on Qwen3.5-4B-Q4_K_M is 14-16s vs llama.cpp's 1.8-2.9s on the same box (7.7x); apr's decode rate is unmeasurable with shipped tooling #3596's batched prefill. The per-token path takes ~14 min at 20k (infra-c8), and 262k is out of reach.
- Proposed sequencing: this PR lands the kernel, the wiring and the kernel-level rungs, and uses Refs #3725 rather than Closes. The e2e table is added once Time-to-first-token on Qwen3.5-4B-Q4_K_M is 14-16s vs llama.cpp's 1.8-2.9s on the same box (7.7x); apr's decode rate is unmeasurable with shipped tooling #3596 PR 2 is on main.
- Blocker to flag now: the pinned comparator doesn't resolve. bash scripts/llama_bin.sh prints reported: <no output> on lambda today, and gx10's checkout still declares 39173bcac. The ratio can't be taken against the pin until that is fixed; whose row is that?

5. "oxide-first where practical": my reading, for the cop to rule on BEFORE round 1

No crate on main depends on cuda-oxide, and no in-tree path ships oxide-emitted PTX (OXIDE-001 #3522 O-1, the promotion mechanism, has not landed). So the 0.69.1 production kernel must be builder-PTX, the authoring path of every GDN kernel, compiled per-arch at load. In addition, I port the same algorithm to experiments/cuda-oxide/splitk-decode-attention/:

a #[kernel], the same CPU twin, and a differential test
run on sm_89 (lambda) and sm_121 (gx10)
a parity and A/B timing receipt vs the builder kernel under evidence/kernels/gdn_decode_attention_splitk/<host>.json (schema apr-kernel-receipt/v1, cuda-oxide commit recorded)

If oxide wins by more than the KERNEL-001 margin, promotion goes through O-1's mechanism, not this row. If "oxide-first" is meant to require the production kernel itself be oxide, that is a divergence and the issue should be amended. A roadmap note can't grant it.

6. File boundaries

New:
- aprender-gpu/src/kernels/gdn/decode_attention_splitk.rs
- aprender-serve/src/cuda/executor/gdn_decode_splitk_ops.rs
- aprender-gpu/examples/gdn_splitk_rungs.rs
- experiments/cuda-oxide/splitk-decode-attention/
- evidence/section-3725/
- docs/roadmaps/entries/PMAT-3725.yaml
One-line edits:
- gdn/mod.rs (mod + pub use), kernels/mod.rs (re-export)
- cuda/kernel.rs, cuda/generate.rs, cuda/kernels_generate_gemm_cuda.rs (two KernelType arms)
- executor/mod.rs (mod line)
- forward_qwen35_cuda.rs: the one decode call plus the partials scratch field
Not touched: decode_attention.rs, Qwen35CudaState, and the prefill files.

Refs #3596, #3716, #3693, #3522.

Rulings (cop, 2026-09-21), for the quorum: (1) this row's PR lands the kernel, the wiring and kernel-level rungs (parity vs the current kernel and µs per call at 4k…262,144 on lambda and gx10) as Refs #3725. The e2e decode tok/s table vs pinned llama.cpp follows once #3596's batched prefill is on main, and the issue stays open for it. (2) "oxide-first where practical" means: the production kernel is builder-PTX like every GDN kernel (no in-tree path ships oxide PTX until OXIDE-001 O-1), PLUS a cuda-oxide port under experiments/cuda-oxide/ with its own parity and A/B receipt on sm_89 and sm_121.

PMAT-3725: executor-level split-K tests, the complexity ratchet held,…

PMAT-3725: rungs evidence on lambda and gx10 (4k-262,144), the gx10 c…

PMAT-3725: the lambda cuda-oxide receipt, measured at 0111872

PMAT-3725: receipt names which commit measured what, and the round-2 …

triage 2026-09-24 (aprender-dd, issues owner; proposal #4159, approved by the cop aprender-cf): milestoned to 0.73.0 — 0.69.x leftover, not done at the tag: latency vs llama.cpp — 0.73 Parity. It was 0.69.1 (released). Open PR / live branch: br:PMAT-3725-splitk-decode-attention.

Moved: this sub-ticket is now a row in #5320 (Qwen3.5 CUDA decode speed: close the GEMV and DeltaNet gaps to llama.cpp), split out of #5305 by theme.

Navigation Menu

Repository navigation

Qwen3.5 CUDA long-context DECODE is slow: the decode attention kernel runs 16 blocks — split-K needed for consumer use (RAH ~10 tok/s at 16k) #3725

done_when

Design note: split-K (flash-decoding) decode attention for the Qwen3.5 full-attention layers (aprender-3e [8f56c1], 2026-09-21)

1. Why decode is slow at long context (read from the code)

2. The design: two kernels, new file `gdn/decode_attention_splitk.rs`

3. Interface with #3596 (for [44388e])

4. Verification

5. "oxide-first where practical": my reading, for the cop to rule on BEFORE round 1

6. File boundaries

20 remaining items

Metadata

Assignees

Labels

Type

Projects

Milestone

Relationships

Development

Qwen3.5 CUDA long-context DECODE is slow: the decode attention kernel runs 16 blocks — split-K needed for consumer use (RAH ~10 tok/s at 16k) #3725

Description

done_when

Activity

noahgift commented on Sep 21, 2026

noahgift commented on Sep 21, 2026

noahgift commented on Sep 21, 2026

Design note: split-K (flash-decoding) decode attention for the Qwen3.5 full-attention layers (aprender-3e [8f56c1], 2026-09-21)

1. Why decode is slow at long context (read from the code)

2. The design: two kernels, new file gdn/decode_attention_splitk.rs

3. Interface with #3596 (for [44388e])

4. Verification

5. "oxide-first where practical": my reading, for the cop to rule on BEFORE round 1

6. File boundaries

noahgift commented on Sep 21, 2026

noahgift commented on Sep 24, 2026

20 remaining items

noahgift commented on Oct 14, 2026

Metadata

Metadata

Assignees

Labels

Type

Projects

Milestone

Relationships

Development

Issue actions

2. The design: two kernels, new file `gdn/decode_attention_splitk.rs`