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:
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user