Repository navigation
perf(cuda): attention_f32 flash-decoding split-KV for depth-dependent decode cost #183
Description
Context
Follow-up to this session's depth/attention-cost investigation (see .docs/handoff.md, "Depth-
dependent attention finding, precisely localized"). Profiling found attn-6c-core
(LaunchAttentionF32 -> attention_f32 kernel, native/kernels/attention_f32.cu) grows +151%
(0.043ms -> 0.108ms/call) from context depth 0 to 256 — the single biggest depth-dependent decode
cost found this session.
attention_f32 is already a properly tiled, online-softmax kernel (TILE_KV=256, running-max/
running-sum accumulation, causal + sliding-window masking) — not naive. The real problem is
structural: blockIdx.x maps to (query_token, head) only (grid = seq_q * num_heads). For
decode, seq_q == 1, so for Bonsai-27B (qwen35.attention.head_count=24) grid == 24 — smaller
than the RTX 3060's 28 SMs, so several SMs get zero work, and the KV-tile loop runs entirely
sequentially within one block per head with zero parallelism across the KV dimension. This is
exactly the scenario the published "Flash-Decoding" technique targets: split KV across additional
blocks, each computing a partial (running_max, running_sum, weighted_output) over its own KV
range, then combine partials with the standard online-softmax merge rule.
Scope
- New split-KV kernel variant in
native/kernels/attention_f32.cu, decode-only (seq_q==1) scope. - Must reuse the existing
fast_exp_negSchraudolph approximation (matches CPU
DotLLM.Cpu.Kernels.FastMath/Softmax.ExecuteFast) for all exponentials in the merge/combine —
no preciseexpf(), to avoid regressing the already-tuned CPU/GPU parity property. - Follow the issue perf(cuda): gdn_scan_step_f32 grid-too-small occupancy — cooperative-groups row-split (ncu-informed follow-up to #173) #180/perf(cuda): gdn_scan_step_f32 cooperative-groups row-split, opt-in default-off #181 (
gdn_scan_step_f32_coop_split4) precedent for shape: reuse the
cuLaunchCooperativeKernel/cuOccupancyMaxActiveBlocksPerMultiprocessorP/Invoke plumbing
already inCudaDriverApi.cs, ship as an opt-in, default-OFF env-var flag with a mandatory
per-shape/per-GPU cooperative-launch co-residency safety check and clean fallback to the exact
kernel when unsafe. - Do NOT confuse with
attention_flash_mma.cu(unrelated prefill-only tensor-core prototype). - Mandatory correctness gate: tolerance-based test (single-step + many-consecutive-decode-step
drift characterization) at REAL depth (256, 1024+), following
CudaGdnScanStepF32CoopSplit4Tests.cs's pattern. - Real end-to-end
dotnet run benchmeasurement at meaningful depth, multiple A/B rounds. - A well-documented negative result (reverted, doc-only PR) is an acceptable outcome if it doesn't
measurably win — same standard as the ~10 prior negative results inpq2_0_gemv.cu's header and
the perf(cuda): verify layer-resid1-norm2/layer-pre-norm1 RmsNorm fusion status (#178) #178/perf(cuda): gdn_scan_step_f32 grid-too-small occupancy — cooperative-groups row-split (ncu-informed follow-up to #173) #180 writeups in.docs/handoff.md.
Acceptance criteria
- Kernel change modeled (launch-count / sync-point / occupancy arithmetic) before implementation.
-
cuobjdump --dump-sass/ptxas -vcheck before and after. - Correctness test at real depth (256, 1024+), tolerance-based, single-step + multi-step drift.
- Real
dotnet run benchA/B, multiple rounds, at depth >= 256. - Full
DotLLM.Tests.Unit.Cudasuite clean. - Shipped as opt-in default-OFF if it's a genuine win, OR reverted with a well-documented
negative-result writeup if it isn't.
Activity
Merge pull request #186 from jamesburton/issue/183-attn-splitkv-flash…
jamesburton commentedon Jul 25, 2026
Merged via PR #186. Shipped as opt-in, default-OFF (
DOTLLM_ATTN_SPLIT_KV=1) -- correctness-validated and zero default-behavior risk, but real bench A/B was inconclusive at the depths safely measurable this session (256/512); depth>=768 hit a pre-existing (not introduced by this change, confirmed on the unmodified baseline too) GPU hang that prevented validating the depth range where the motivating profiling data predicts the clearest win. See PR #186's description for the full writeup and attention_f32.cu's header for the design/literature/verification detail. A future session with more VRAM headroom or a fix for the depth>=768 hang should re-run the same A/B before considering flipping the default.docs(cuda): the mma-decode GQA-split kernel is default-ON, not defaul…