Skip to content

perf(cuda): attention_f32 flash-decoding split-KV for depth-dependent decode cost #183

Closed

Description

@jamesburton

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

Acceptance criteria

  • Kernel change modeled (launch-count / sync-point / occupancy arithmetic) before implementation.
  • cuobjdump --dump-sass / ptxas -v check before and after.
  • Correctness test at real depth (256, 1024+), tolerance-based, single-step + multi-step drift.
  • Real dotnet run bench A/B, multiple rounds, at depth >= 256.
  • Full DotLLM.Tests.Unit.Cuda suite 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

  1. added a commit that references this issue on Jul 25, 2026
  2. jamesburton commented on Jul 25, 2026

    @jamesburton
    OwnerAuthor

    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.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Projects

    No projects

      Milestone

      No milestone

      Relationships

      None yet

      Development

      No branches or pull requests

      Issue actions