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:
@@ -1,7 +1,11 @@
|
||||
#pragma once
|
||||
#include "attn_common.cuh"
|
||||
#include <cfloat>
|
||||
#include <cuda_bf16.h>
|
||||
#include "attn_common.h"
|
||||
#include "attn_mma_utils.cuh"
|
||||
|
||||
using bf16 = __nv_bfloat16;
|
||||
|
||||
// Tensor-core prefill flash attention (raw mma.sync PTX).
|
||||
// One warp owns BR=16 query rows. S = Q@K^T and O = P@V run on bf16 tensor
|
||||
// cores via mma.sync.m16n8k16 (f32 accumulate). Q fragments are loaded once
|
||||
@@ -38,7 +42,7 @@
|
||||
|
||||
template <int HEAD_DIM, int WARPS, int BC, int MIN_BLOCKS>
|
||||
__global__ __launch_bounds__(WARPS * 32, MIN_BLOCKS)
|
||||
void attn_prefill_split_q_mma_kernel(AttentionParams p) {
|
||||
void attn_prefill_split_q_mma_kernel(AttentionParams<bf16> p) {
|
||||
constexpr int BR = 16;
|
||||
constexpr int KD = HEAD_DIM / 16; // Q/K k-tiles
|
||||
constexpr int NC8 = BC / 8; // S n-tiles (N=8 each)
|
||||
|
||||
Reference in New Issue
Block a user