refactor: template AttentionParams, rename .cuh to .h
- Convert AttentionParams to a template struct supporting arbitrary types - Rename attn_common.cuh -> attn_common.h (no CUDA-specific code remains) - Include standard headers explicitly in each .cuh instead of via attn_common.cuh - Allow .h files in csrc/ via .gitignore
This commit is contained in:
@@ -23,7 +23,7 @@ static double now_ms() {
|
||||
// Launch the production prefill path (tensor-core MMA on sm_80+, else the
|
||||
// scalar fallback), mirroring dispatch_prefill() in attn_prefill.cu.
|
||||
template <int HEAD_DIM>
|
||||
static void launch_prefill(AttentionParams& p) {
|
||||
static void launch_prefill(AttentionParams<bf16>& p) {
|
||||
#ifndef ASTRAI_NO_MMA
|
||||
constexpr int WARPS = 4, BR = 16;
|
||||
constexpr int BC = (HEAD_DIM <= 128) ? 32 : 16;
|
||||
@@ -40,7 +40,7 @@ static void launch_prefill(AttentionParams& p) {
|
||||
#endif
|
||||
}
|
||||
|
||||
static void dispatch_prefill(AttentionParams& p) {
|
||||
static void dispatch_prefill(AttentionParams<bf16>& p) {
|
||||
switch (p.head_dim) {
|
||||
case 64: launch_prefill<64>(p); break;
|
||||
case 128: launch_prefill<128>(p); break;
|
||||
@@ -123,7 +123,7 @@ static void bench() {
|
||||
for (size_t i=0;i<nKV;i++) tmp[i]=f2bf(randf());
|
||||
cudaMemcpy(dV,tmp,nKV*2,cudaMemcpyHostToDevice);
|
||||
|
||||
AttentionParams p;
|
||||
AttentionParams<bf16> p;
|
||||
p.batch=B; p.q_head=Hq; p.kv_head=Hk; p.q_len=ql; p.kv_len=kl; p.head_dim=D;
|
||||
p.use_mask=0; p.is_causal=causal; p.causal_offset=0;
|
||||
p.scale=1.0f/sqrtf((float)D);
|
||||
@@ -191,7 +191,7 @@ int main() {
|
||||
for (size_t i=0;i<nKV;i++) tmp[i]=f2bf(hV[i]);
|
||||
cudaMemcpy(dV,tmp,nKV*2,cudaMemcpyHostToDevice);
|
||||
|
||||
AttentionParams p;
|
||||
AttentionParams<bf16> p;
|
||||
p.batch=B; p.q_head=Hq; p.kv_head=Hk; p.q_len=ql; p.kv_len=kl; p.head_dim=D;
|
||||
p.use_mask=0; p.is_causal=causal; p.causal_offset=0;
|
||||
p.scale=1.0f/sqrtf((float)D);
|
||||
|
||||
Reference in New Issue
Block a user