From 530d280e333164787b6466feadbfa20eb0af6cd3 Mon Sep 17 00:00:00 2001 From: ViperEkura <3081035982@qq.com> Date: Fri, 31 Jul 2026 22:37:44 +0800 Subject: [PATCH] perf: remove split partials memset and overlap decode tile loads - alloc_split_partials now uses torch::empty: the split kernel writes every slot it owns, so the per-call zeros/full memset was pure overhead (2 kernels per layer per step) - decode split-KV MMA kernels now run a true multi-stage cp.async pipeline (wait_group instead of wait_group<0>), keeping STAGES-1 tile loads in flight; the old wait_group<0> serialized load and compute so deeper STAGES made no difference - add a fallback path when ntiles < STAGES to avoid a race on the last tile --- astrai/extension/attention_backend.py | 1 - csrc/kernels/attn_decode_split_kv_mma.cuh | 50 +++++++++++-------- csrc/kernels/attn_entry_utils.cuh | 8 ++- .../attn_paged_decode_split_kv_mma.cuh | 49 ++++++++++-------- 4 files changed, 64 insertions(+), 44 deletions(-) diff --git a/astrai/extension/attention_backend.py b/astrai/extension/attention_backend.py index 2372440..650fb3f 100644 --- a/astrai/extension/attention_backend.py +++ b/astrai/extension/attention_backend.py @@ -30,7 +30,6 @@ Layout convention: all q/k/v are ``[batch, seq_len, n_heads, head_dim]`` import contextvars import enum -import math from abc import ABC, abstractmethod from contextlib import contextmanager from typing import Optional, Union diff --git a/csrc/kernels/attn_decode_split_kv_mma.cuh b/csrc/kernels/attn_decode_split_kv_mma.cuh index 516fbbc..6d45a51 100644 --- a/csrc/kernels/attn_decode_split_kv_mma.cuh +++ b/csrc/kernels/attn_decode_split_kv_mma.cuh @@ -76,26 +76,17 @@ __global__ void attn_decode_split_kv_mma_kernel(AttentionParams p) { cp_async_commit(); }; - constexpr int BUF_MASK = (Traits::STAGES > 1) ? (Traits::STAGES - 1) : 0; - - // Prologue - if (ti_begin < ti_end) { - load_tile(ti_begin, 0); - } - - for (int ti = ti_begin; ti < ti_end; ti++) { - int buf = (ti - ti_begin) & BUF_MASK; - - cp_async_wait_group<0>(); - __syncwarp(); - if constexpr (Traits::STAGES > 1) { - if (ti + 1 < ti_end) - load_tile(ti + 1, (ti + 1 - ti_begin) & BUF_MASK); - } + // ---- Multi-stage cp.async pipeline ---- + // Prologue loads STAGES tiles; each loop iteration waits only for the + // oldest outstanding group (wait_group) so the STAGES-1 newer + // tile loads stay in flight and overlap with the current tile's compute. + constexpr int STAGES = Traits::STAGES; + const int ntiles = ti_end - ti_begin; + auto process_tile = [&](int it, int buf) { const bf16* bK = sK + buf * Traits::BC * Traits::LD; const bf16* bV = sV + buf * Traits::BC * Traits::LD; - int kv0 = ti * Traits::BC; + int kv0 = (ti_begin + it) * Traits::BC; float Sacc[Traits::NC8][4]; mma_compute_scores(Qa, bK, lane, Sacc); @@ -115,12 +106,29 @@ __global__ void attn_decode_split_kv_mma_kernel(AttentionParams p) { Sacc, Oacc, m0, m1, l0, l1, lane); mma_pv_accumulate(Sacc, bV, lane, Oacc); - __syncwarp(); + }; - if constexpr (Traits::STAGES == 1) { - if (ti + 1 < ti_end) - load_tile(ti + 1, 0); + if (ntiles >= STAGES) { + #pragma unroll + for (int i = 0; i < STAGES; i++) + load_tile(ti_begin + i, i); + + for (int it = 0; it < ntiles; it++) { + cp_async_wait_group(); + __syncwarp(); + process_tile(it, it & (STAGES - 1)); + __syncwarp(); + if (it + STAGES < ntiles) + load_tile(ti_begin + it + STAGES, (it + STAGES) & (STAGES - 1)); } + } else { + // Fewer tiles than stages: load all, wait for all, process. + for (int i = 0; i < ntiles; i++) + load_tile(ti_begin + i, i); + cp_async_wait_group<0>(); + __syncwarp(); + for (int it = 0; it < ntiles; it++) + process_tile(it, it); } // ---- write UN-normalised partials for this split ---- diff --git a/csrc/kernels/attn_entry_utils.cuh b/csrc/kernels/attn_entry_utils.cuh index 44d6498..50cb34e 100644 --- a/csrc/kernels/attn_entry_utils.cuh +++ b/csrc/kernels/attn_entry_utils.cuh @@ -21,11 +21,15 @@ using bf16 = __nv_bfloat16; " (supported: 32, 64, 128, 256)"); \ } +// The split kernel unconditionally writes every (batch, q_head, split) slot it +// owns — including empty split ranges, which store m = -FLT_MAX so the combine +// skips them. Allocators are therefore left uninitialized (torch::empty); the +// per-call memset (torch::zeros / torch::full) was pure overhead. template inline void alloc_split_partials(P& p) { auto fopt = torch::TensorOptions().dtype(torch::kFloat32).device(torch::kCUDA); - auto o_part = torch::zeros(at::IntArrayRef{p.batch, p.q_head, MAX_SPLITS, p.head_dim}, fopt); - auto ml_part = torch::full(at::IntArrayRef{p.batch, p.q_head, MAX_SPLITS, 2}, -FLT_MAX, fopt); + auto o_part = torch::empty(at::IntArrayRef{p.batch, p.q_head, MAX_SPLITS, p.head_dim}, fopt); + auto ml_part = torch::empty(at::IntArrayRef{p.batch, p.q_head, MAX_SPLITS, 2}, fopt); p.o_part = (float*)o_part.data_ptr(); p.ml_part = (float*)ml_part.data_ptr(); } diff --git a/csrc/kernels/attn_paged_decode_split_kv_mma.cuh b/csrc/kernels/attn_paged_decode_split_kv_mma.cuh index 92ae734..e692c9b 100644 --- a/csrc/kernels/attn_paged_decode_split_kv_mma.cuh +++ b/csrc/kernels/attn_paged_decode_split_kv_mma.cuh @@ -91,25 +91,17 @@ __global__ void paged_attn_decode_split_kv_mma_kernel(PagedAttentionParams cp_async_commit(); }; - constexpr int BUF_MASK = (Traits::STAGES > 1) ? (Traits::STAGES - 1) : 0; - - if (ti_begin < ti_end) { - load_tile(ti_begin, 0); - } - - for (int ti = ti_begin; ti < ti_end; ti++) { - int buf = (ti - ti_begin) & BUF_MASK; - - cp_async_wait_group<0>(); - __syncwarp(); - if constexpr (Traits::STAGES > 1) { - if (ti + 1 < ti_end) - load_tile(ti + 1, (ti + 1 - ti_begin) & BUF_MASK); - } + // ---- Multi-stage cp.async pipeline ---- + // Prologue loads STAGES tiles; each loop iteration waits only for the + // oldest outstanding group (wait_group) so the STAGES-1 newer + // tile loads stay in flight and overlap with the current tile's compute. + constexpr int STAGES = Traits::STAGES; + const int ntiles = ti_end - ti_begin; + auto process_tile = [&](int it, int buf) { const bf16* bK = sK + buf * Traits::BC * Traits::LD; const bf16* bV = sV + buf * Traits::BC * Traits::LD; - int kv0 = ti * Traits::BC; + int kv0 = (ti_begin + it) * Traits::BC; float Sacc[Traits::NC8][4]; mma_compute_scores(Qa, bK, lane, Sacc); @@ -128,12 +120,29 @@ __global__ void paged_attn_decode_split_kv_mma_kernel(PagedAttentionParams Sacc, Oacc, m0, m1, l0, l1, lane); mma_pv_accumulate(Sacc, bV, lane, Oacc); - __syncwarp(); + }; - if constexpr (Traits::STAGES == 1) { - if (ti + 1 < ti_end) - load_tile(ti + 1, 0); + if (ntiles >= STAGES) { + #pragma unroll + for (int i = 0; i < STAGES; i++) + load_tile(ti_begin + i, i); + + for (int it = 0; it < ntiles; it++) { + cp_async_wait_group(); + __syncwarp(); + process_tile(it, it & (STAGES - 1)); + __syncwarp(); + if (it + STAGES < ntiles) + load_tile(ti_begin + it + STAGES, (it + STAGES) & (STAGES - 1)); } + } else { + // Fewer tiles than stages: load all, wait for all, process. + for (int i = 0; i < ntiles; i++) + load_tile(ti_begin + i, i); + cp_async_wait_group<0>(); + __syncwarp(); + for (int it = 0; it < ntiles; it++) + process_tile(it, it); } auto split_slot = [&](int h) -> size_t {