refactor: namespace csrc kernels and extract common helpers

- attention family -> astrai::attention; fp8 family -> astrai::fp8
- new common/reduce.cuh (warp/group reductions, atomic_max_float)
- new common/cp_async.cuh (predicated cp_async_16, commit/wait group)
- move MAX_SPLITS into attention/common.h; delete warp_utils.cuh
- .cu bindings and pure C tests open family namespaces via using
This commit is contained in:
2026-08-24 14:49:55 +08:00
parent 34471252ab
commit 31ca357c61
21 changed files with 243 additions and 139 deletions
+11 -5
View File
@@ -5,6 +5,9 @@
#include "layout_policies.cuh"
#include "mma_utils.cuh"
namespace astrai {
namespace attention {
// Tensor-core prefill flash attention (raw mma.sync PTX), unified across
// contiguous and paged (SGLang flat-pool) K/V via the KV template parameter.
// One warp owns BR=16 query rows. S = Q@K^T and O = P@V run on bf16 tensor
@@ -85,10 +88,10 @@ __global__ void attn_prefill_split_q_mma_kernel(AttentionParams<bf16> p) {
int token = KV::resolve_token(p, kctx, kc, valid);
KVAddr a = KV::kv_addr_from_token(p, kctx, token, d);
int off = r * Traits::LD + swiz_col(d, r, Traits::SWIZ_MASK);
cp_async_16_pred(&dK[off], a.k, a.valid);
cp_async_16_pred(&dV[off], a.v, a.valid);
astrai::cp_async_16(&dK[off], a.k, a.valid);
astrai::cp_async_16(&dV[off], a.v, a.valid);
}
cp_async_commit();
astrai::cp_async_commit_group();
};
// ---- Prologue: issue first tile load ----
@@ -98,7 +101,7 @@ __global__ void attn_prefill_split_q_mma_kernel(AttentionParams<bf16> p) {
int buf = ti & 1;
// Wait for current tile, then publish cross-warp + guard buffer reuse.
cp_async_wait_group<0>();
astrai::cp_async_wait_group<0>();
__syncthreads();
if (ti < t_end) load_tile(ti + 1, (ti + 1) & 1);
@@ -149,9 +152,12 @@ __global__ void attn_prefill_split_q_mma_kernel(AttentionParams<bf16> p) {
}
if (qr1 < q_len) {
__nv_bfloat162 v = __floats2bfloat162_rn(Oacc[dn8][2] * rl1,
Oacc[dn8][3] * rl1);
Oacc[dn8][3] * rl1);
*reinterpret_cast<__nv_bfloat162*>(
&p.o_ptr[o_base + qr1 * p.q_l_stride + d * p.q_d_stride]) = v;
}
}
}
} // namespace attention
} // namespace astrai