refactor: template combine kernel, fix mask bug, unify dispatch
- Template combine kernel, share macros, extract entry_utils helpers - Fix mask indexing (pass stride not pre-multiplied base) - Remove !p.use_mask — MMA handles mask
This commit is contained in:
@@ -5,10 +5,6 @@
|
||||
|
||||
using bf16 = __nv_bfloat16;
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Shared dispatch helpers — eliminates duplication across .cu entry files.
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
inline int compute_num_splits(int base_blocks, int tiles_total) {
|
||||
int sm_count = 0;
|
||||
cudaDeviceGetAttribute(&sm_count, cudaDevAttrMultiProcessorCount, 0);
|
||||
@@ -16,21 +12,20 @@ inline int compute_num_splits(int base_blocks, int tiles_total) {
|
||||
return std::max(1, std::min(n, std::min(tiles_total, 32)));
|
||||
}
|
||||
|
||||
// Dispatch head_dim to a generic lambda FN (callable as FN.operator()<D>()).
|
||||
template <typename Fn>
|
||||
inline void dispatch_head_dim(int hd, Fn&& fn) {
|
||||
switch (hd) {
|
||||
case 32: fn.template operator()<32>(); break;
|
||||
case 64: fn.template operator()<64>(); break;
|
||||
case 128: fn.template operator()<128>(); break;
|
||||
case 256: fn.template operator()<256>(); break;
|
||||
default:
|
||||
TORCH_CHECK(false, "unsupported head_dim ", hd,
|
||||
" (supported: 32, 64, 128, 256)");
|
||||
// Dispatch head_dim: shared macro — avoids C++20 lambda template syntax.
|
||||
// Usage: DISPATCH_HEAD_DIM(hd, fn, arg)
|
||||
// Expands to: fn<32>(arg); fn<64>(arg); etc.
|
||||
#define DISPATCH_HEAD_DIM(hd, fn, arg) \
|
||||
switch (hd) { \
|
||||
case 32: fn<32>(arg); break; \
|
||||
case 64: fn<64>(arg); break; \
|
||||
case 128: fn<128>(arg); break; \
|
||||
case 256: fn<256>(arg); break; \
|
||||
default: \
|
||||
TORCH_CHECK(false, "unsupported head_dim ", hd, \
|
||||
" (supported: 32, 64, 128, 256)"); \
|
||||
}
|
||||
}
|
||||
|
||||
// Allocate split-KV partial buffers and wire into params.
|
||||
template<typename P>
|
||||
inline void alloc_split_partials(P& p) {
|
||||
auto fopt = torch::TensorOptions().dtype(torch::kFloat32).device(torch::kCUDA);
|
||||
@@ -40,10 +35,48 @@ inline void alloc_split_partials(P& p) {
|
||||
p.ml_part = (float*)ml_part.data_ptr();
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Param packing — fills AttentionParams from torch tensors with validation.
|
||||
// ---------------------------------------------------------------------------
|
||||
// ---- Shared Q-dims + strides extraction ----
|
||||
template <typename P>
|
||||
inline void extract_q_dims_and_strides(torch::Tensor& q, int64_t layout, P& p) {
|
||||
if (layout == 1) q = q.transpose(1, 2);
|
||||
p.batch = (int)q.size(0);
|
||||
p.q_head = (int)q.size(1);
|
||||
p.q_len = (int)q.size(2);
|
||||
p.head_dim = (int)q.size(3);
|
||||
p.q_stride_b = (int)q.stride(0);
|
||||
p.q_stride_h = (int)q.stride(1);
|
||||
p.q_stride_l = (int)q.stride(2);
|
||||
p.q_stride_d = (int)q.stride(3);
|
||||
}
|
||||
|
||||
// ---- Shared mask packing ----
|
||||
template <typename P>
|
||||
inline void pack_mask(const c10::optional<torch::Tensor>& mask, P& p) {
|
||||
if (p.use_mask) {
|
||||
auto m = mask.value();
|
||||
TORCH_CHECK(m.is_cuda(), "mask must be on CUDA");
|
||||
TORCH_CHECK(m.dtype() == torch::kBool, "mask must be bool");
|
||||
TORCH_CHECK(m.size(0) == p.batch, "mask batch mismatch");
|
||||
TORCH_CHECK(m.size(m.dim() - 1) == p.kv_len, "mask kv_len mismatch");
|
||||
if (m.dim() == 2) {
|
||||
p.mask_b_stride = (int)m.stride(0);
|
||||
p.mask_q_stride = 0;
|
||||
} else if (m.dim() == 3) {
|
||||
TORCH_CHECK(m.size(1) == p.q_len, "mask q_len mismatch");
|
||||
p.mask_b_stride = (int)m.stride(0);
|
||||
p.mask_q_stride = (int)m.stride(1);
|
||||
} else {
|
||||
TORCH_CHECK(false, "mask must be 2D [batch, kv_len] or 3D [batch, q_len, kv_len]");
|
||||
}
|
||||
p.mask = m.data_ptr<bool>();
|
||||
} else {
|
||||
p.mask = nullptr;
|
||||
p.mask_b_stride = 0;
|
||||
p.mask_q_stride = 0;
|
||||
}
|
||||
}
|
||||
|
||||
// ---- attn_pack_params (contiguous KV) ----
|
||||
template<typename T>
|
||||
inline void attn_pack_params(
|
||||
torch::Tensor q,
|
||||
@@ -52,7 +85,7 @@ inline void attn_pack_params(
|
||||
c10::optional<torch::Tensor> mask,
|
||||
int64_t causal_offset,
|
||||
double scale,
|
||||
int64_t layout, // 0 = b h l d, 1 = b l h d
|
||||
int64_t layout,
|
||||
AttentionParams<T>& p
|
||||
) {
|
||||
const at::cuda::OptionalCUDAGuard device_guard(device_of(q));
|
||||
@@ -64,26 +97,14 @@ inline void attn_pack_params(
|
||||
TORCH_CHECK(k.sizes() == v.sizes(), "K and V must have identical shapes");
|
||||
TORCH_CHECK(q.dim() == 4 && k.dim() == 4, "Q/K/V must be 4D");
|
||||
|
||||
// Normalize to b h l d view (zero-copy transpose if user passed b l h d)
|
||||
if (layout == 1) {
|
||||
q = q.transpose(1, 2);
|
||||
k = k.transpose(1, 2);
|
||||
v = v.transpose(1, 2);
|
||||
}
|
||||
extract_q_dims_and_strides(q, layout, p);
|
||||
|
||||
if (layout == 1) k = k.transpose(1, 2), v = v.transpose(1, 2);
|
||||
|
||||
p.batch = (int)q.size(0);
|
||||
p.q_head = (int)q.size(1);
|
||||
p.q_len = (int)q.size(2);
|
||||
p.head_dim = (int)q.size(3);
|
||||
p.kv_head = (int)k.size(1);
|
||||
p.kv_len = (int)k.size(2);
|
||||
TORCH_CHECK(k.size(3) == p.head_dim, "K/V head_dim must match Q");
|
||||
|
||||
// Strides (layout-agnostic: works for b h l d and b l h d)
|
||||
p.q_stride_b = (int)q.stride(0);
|
||||
p.q_stride_h = (int)q.stride(1);
|
||||
p.q_stride_l = (int)q.stride(2);
|
||||
p.q_stride_d = (int)q.stride(3);
|
||||
p.kv_stride_b = (int)k.stride(0);
|
||||
p.kv_stride_h = (int)k.stride(1);
|
||||
p.kv_stride_l = (int)k.stride(2);
|
||||
@@ -100,34 +121,10 @@ inline void attn_pack_params(
|
||||
p.o_part = nullptr;
|
||||
p.ml_part = nullptr;
|
||||
|
||||
if (p.use_mask) {
|
||||
auto m = mask.value();
|
||||
TORCH_CHECK(m.is_cuda(), "mask must be on CUDA");
|
||||
TORCH_CHECK(m.dtype() == torch::kBool, "mask must be bool");
|
||||
TORCH_CHECK(m.size(0) == p.batch, "mask batch mismatch");
|
||||
TORCH_CHECK(m.size(m.dim() - 1) == p.kv_len, "mask kv_len mismatch");
|
||||
if (m.dim() == 2) {
|
||||
p.mask_b_stride = (int)m.stride(0);
|
||||
p.mask_q_stride = 0;
|
||||
} else if (m.dim() == 3) {
|
||||
TORCH_CHECK(m.size(1) == p.q_len, "mask q_len mismatch");
|
||||
p.mask_b_stride = (int)m.stride(0);
|
||||
p.mask_q_stride = (int)m.stride(1);
|
||||
} else {
|
||||
TORCH_CHECK(false, "mask must be 2D [batch, kv_len] or 3D [batch, q_len, kv_len]");
|
||||
}
|
||||
p.mask = m.data_ptr<bool>();
|
||||
} else {
|
||||
p.mask = nullptr;
|
||||
p.mask_b_stride = 0;
|
||||
p.mask_q_stride = 0;
|
||||
}
|
||||
pack_mask(mask, p);
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Param packing for paged attention.
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
// ---- attn_pack_paged_params ----
|
||||
template<typename T>
|
||||
inline void attn_pack_paged_params(
|
||||
torch::Tensor q,
|
||||
@@ -139,7 +136,7 @@ inline void attn_pack_paged_params(
|
||||
c10::optional<torch::Tensor> mask,
|
||||
int64_t causal_offset,
|
||||
double scale,
|
||||
int64_t layout, // 0 = b h l d, 1 = b l h d
|
||||
int64_t layout,
|
||||
PagedAttentionParams<T>& p
|
||||
) {
|
||||
const at::cuda::OptionalCUDAGuard device_guard(device_of(q));
|
||||
@@ -151,15 +148,8 @@ inline void attn_pack_paged_params(
|
||||
TORCH_CHECK(page_table.dtype() == torch::kLong, "page_table must be int64");
|
||||
TORCH_CHECK(k_cache.sizes() == v_cache.sizes(), "k_cache and v_cache must have identical shapes");
|
||||
|
||||
// Normalize Q to b h l d view if user passed b l h d
|
||||
if (layout == 1) {
|
||||
q = q.transpose(1, 2);
|
||||
}
|
||||
extract_q_dims_and_strides(q, layout, p);
|
||||
|
||||
p.batch = (int)q.size(0);
|
||||
p.q_head = (int)q.size(1);
|
||||
p.q_len = (int)q.size(2);
|
||||
p.head_dim = (int)q.size(3);
|
||||
p.kv_head = (int)k_cache.size(2);
|
||||
p.kv_len = (int)kv_len;
|
||||
p.page_size = (int)page_size;
|
||||
@@ -171,12 +161,6 @@ inline void attn_pack_paged_params(
|
||||
"k_cache dim 1 must equal page_size, got ",
|
||||
k_cache.size(1), " vs ", page_size);
|
||||
|
||||
// Q strides
|
||||
p.q_stride_b = (int)q.stride(0);
|
||||
p.q_stride_h = (int)q.stride(1);
|
||||
p.q_stride_l = (int)q.stride(2);
|
||||
p.q_stride_d = (int)q.stride(3);
|
||||
|
||||
p.causal_offset = (int)causal_offset;
|
||||
p.use_mask = (mask.has_value() && mask.value().defined()) ? 1 : 0;
|
||||
p.scale = (scale > 0.0) ? (float)scale : 1.0f / sqrtf((float)p.head_dim);
|
||||
@@ -189,26 +173,5 @@ inline void attn_pack_paged_params(
|
||||
p.o_part = nullptr;
|
||||
p.ml_part = nullptr;
|
||||
|
||||
if (p.use_mask) {
|
||||
auto m = mask.value();
|
||||
TORCH_CHECK(m.is_cuda(), "mask must be on CUDA");
|
||||
TORCH_CHECK(m.dtype() == torch::kBool, "mask must be bool");
|
||||
TORCH_CHECK(m.size(0) == p.batch, "mask batch mismatch");
|
||||
TORCH_CHECK(m.size(m.dim() - 1) == p.kv_len, "mask kv_len mismatch");
|
||||
if (m.dim() == 2) {
|
||||
p.mask_b_stride = (int)m.stride(0);
|
||||
p.mask_q_stride = 0;
|
||||
} else if (m.dim() == 3) {
|
||||
TORCH_CHECK(m.size(1) == p.q_len, "mask q_len mismatch");
|
||||
p.mask_b_stride = (int)m.stride(0);
|
||||
p.mask_q_stride = (int)m.stride(1);
|
||||
} else {
|
||||
TORCH_CHECK(false, "mask must be 2D [batch, kv_len] or 3D [batch, q_len, kv_len]");
|
||||
}
|
||||
p.mask = m.data_ptr<bool>();
|
||||
} else {
|
||||
p.mask = nullptr;
|
||||
p.mask_b_stride = 0;
|
||||
p.mask_q_stride = 0;
|
||||
}
|
||||
pack_mask(mask, p);
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user