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 -1
View File
@@ -1,5 +1,10 @@
#pragma once
// Pure POD header
namespace astrai {
namespace attention {
// Tensor layout for Q/K/V tensors passed to attention kernels.
// Internally, kernels always operate on BHLD [batch, n_heads, seq_len, head_dim].
// When the caller passes BLHD, dims 1 and 2 are transposed at entry.
@@ -8,6 +13,9 @@ enum TensorLayout : int {
BLHD = 1, // [batch, seq_len, n_heads, head_dim]
};
// Split-KV workspace cap: max decode splits per (batch, q_head).
constexpr int MAX_SPLITS = 32;
// Unified attention params covering BOTH addressing modes:
// - Contiguous K/V: dense [batch, kv_head, kv_len, head_dim] tensors (k/v).
@@ -73,5 +81,7 @@ struct AttentionParams {
int num_splits;
AT* __restrict__ o_part;
AT* __restrict__ ml_part;
};
} // namespace attention
} // namespace astrai