diff --git a/csrc/kernels/attn_decode.cu b/csrc/kernels/attn_decode.cu index 2880784..c704b9d 100644 --- a/csrc/kernels/attn_decode.cu +++ b/csrc/kernels/attn_decode.cu @@ -29,8 +29,15 @@ torch::Tensor attn_decode( TORCH_CHECK(o_part_buf->scalar_type() == torch::kFloat32, "o_part_buf must be f32"); TORCH_CHECK(ml_part_buf->scalar_type() == torch::kFloat32, "ml_part_buf must be f32"); int64_t o_needed = (int64_t)p.batch * p.q_head * MAX_SPLITS * p.head_dim; + int64_t ml_needed = (int64_t)p.batch * p.q_head * MAX_SPLITS * 2; TORCH_CHECK(o_part_buf->numel() >= o_needed, "o_part_buf too small: need ", o_needed, " got ", o_part_buf->numel()); + TORCH_CHECK(ml_part_buf->numel() >= ml_needed, + "ml_part_buf too small: need ", ml_needed, " got ", ml_part_buf->numel()); + TORCH_CHECK(o_part_buf->is_cuda() && ml_part_buf->is_cuda(), + "split buffers must be CUDA tensors"); + TORCH_CHECK(o_part_buf->is_contiguous() && ml_part_buf->is_contiguous(), + "split buffers must be contiguous"); p.o_part = (float*)o_part_buf->data_ptr(); p.ml_part = (float*)ml_part_buf->data_ptr(); } else { diff --git a/csrc/kernels/attn_decode_split_kv_mma.cuh b/csrc/kernels/attn_decode_split_kv_mma.cuh index 853b485..a67aebc 100644 --- a/csrc/kernels/attn_decode_split_kv_mma.cuh +++ b/csrc/kernels/attn_decode_split_kv_mma.cuh @@ -107,10 +107,11 @@ __global__ void attn_decode_split_kv_mma_kernel(AttentionParams p) { int maxc = IsCausal ? KV::decode_attend_len(p, batch) : seq_len; mma_softmax_tile(kv0, maxc, maxc, 0, 0, - p.mask_b_stride, 0, 0, - batch, 0, - p.mask, - Sacc, Oacc, m0, m1, l0, l1, lane); + p.mask_b_stride, p.mask_h_stride, p.mask_q_stride, + batch, q_head0 + gid, q_head0 + gid + 8, + p.mask, + va, vb, + Sacc, Oacc, m0, m1, l0, l1, lane); mma_pv_accumulate(Sacc, bV, lane, Oacc); }; diff --git a/csrc/kernels/attn_dispatchers.cuh b/csrc/kernels/attn_dispatchers.cuh index 38ad145..aad285b 100644 --- a/csrc/kernels/attn_dispatchers.cuh +++ b/csrc/kernels/attn_dispatchers.cuh @@ -80,7 +80,7 @@ template struct PrefillLauncherScalar { template static void launch(AttentionParams& p, cudaStream_t stream) { - constexpr int G = 8, ROWS = 32, P_BC = 32; + constexpr int G = (HEAD_DIM == 32) ? 4 : 8, ROWS = 32, P_BC = 32; int q_len = KV::host_q_len(p); dim3 grid((q_len + ROWS - 1) / ROWS, p.q_head, p.batch); dim3 block(G, ROWS); @@ -161,6 +161,10 @@ struct DecodeLauncherScalar { int g = min(group_size, 32); // cap at 32 to respect 1024-thread limit dim3 grid(p.batch * p.kv_head, 1, p.num_splits); dim3 block(32, g); + cudaFuncSetAttribute( + attn_decode_split_kv_kernel, + cudaFuncAttributeMaxDynamicSharedMemorySize, + smem); attn_decode_split_kv_kernel <<>>(p); } diff --git a/csrc/kernels/attn_entry_utils.cuh b/csrc/kernels/attn_entry_utils.cuh index 90c3618..869a606 100644 --- a/csrc/kernels/attn_entry_utils.cuh +++ b/csrc/kernels/attn_entry_utils.cuh @@ -106,14 +106,17 @@ inline void attn_pack_params( TORCH_CHECK(v.dtype() == torch::kBFloat16); 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"); - extract_q_dims_and_strides(q, layout, p); if (layout == BLHD) k = k.transpose(1, 2), v = v.transpose(1, 2); p.kv_head = (int)k.size(1); p.kv_len = (int)k.size(2); + TORCH_CHECK(p.q_head % p.kv_head == 0, + "q_head must be divisible by kv_head"); TORCH_CHECK(k.size(3) == p.head_dim, "K/V head_dim must match Q"); + TORCH_CHECK(q.stride(3) == 1 && k.stride(3) == 1 && v.stride(3) == 1, + "Q/K/V head_dim must be contiguous"); p.kv_stride_b = (int)k.stride(0); p.kv_stride_h = (int)k.stride(1); @@ -170,6 +173,8 @@ inline void attn_pack_paged_decode_params( p.head_dim = (int)q.size(2); p.kv_head = (int)k_cache.size(1); TORCH_CHECK(k_cache.size(2) == p.head_dim, "k_cache head_dim mismatch"); + TORCH_CHECK(q.stride(2) == 1 && k_cache.stride(2) == 1 && v_cache.stride(2) == 1, + "Q/K/V head_dim must be contiguous"); TORCH_CHECK(p.head_dim % 32 == 0, "head_dim must be multiple of 32"); TORCH_CHECK(p.q_head % p.kv_head == 0, "q_head must be divisible by kv_head"); @@ -252,6 +257,8 @@ inline void attn_pack_paged_prefill_params( p.kv_head = (int)k_cache.size(1); p.batch = (int)req_pool_indices.size(0); TORCH_CHECK(k_cache.size(2) == p.head_dim, "k_cache head_dim mismatch"); + TORCH_CHECK(q.stride(2) == 1 && k_cache.stride(2) == 1 && v_cache.stride(2) == 1, + "Q/K/V head_dim must be contiguous"); TORCH_CHECK(p.head_dim % 16 == 0, "head_dim must be multiple of 16"); TORCH_CHECK(p.q_head % p.kv_head == 0, "q_head must be divisible by kv_head"); TORCH_CHECK(kv_indptr.size(0) == p.batch + 1, "kv_indptr must be [batch+1]"); diff --git a/csrc/kernels/attn_mma_utils.cuh b/csrc/kernels/attn_mma_utils.cuh index 34a2ccb..2463465 100644 --- a/csrc/kernels/attn_mma_utils.cuh +++ b/csrc/kernels/attn_mma_utils.cuh @@ -199,8 +199,9 @@ __device__ inline void mma_softmax_tile( int maxc0, int maxc1, int qrow0, int qrow1, int mask_b_stride, int mask_h_stride, int mask_q_stride, - int mask_batch, int mask_head, + int mask_batch, int mask_head0, int mask_head1, const bool* __restrict__ mask, + bool valid0, bool valid1, float Sacc[Traits::NC8][4], float Oacc[Traits::DN8][4], float& m0, float& m1, @@ -210,16 +211,16 @@ __device__ inline void mma_softmax_tile( int tid4 = lane & 3; float rmax0 = -FLT_MAX, rmax1 = -FLT_MAX; - int mask_base0 = mask_batch * mask_b_stride + mask_head * mask_h_stride + qrow0 * mask_q_stride; - int mask_base1 = mask_batch * mask_b_stride + mask_head * mask_h_stride + qrow1 * mask_q_stride; + int mask_base0 = mask_batch * mask_b_stride + mask_head0 * mask_h_stride + qrow0 * mask_q_stride; + int mask_base1 = mask_batch * mask_b_stride + mask_head1 * mask_h_stride + qrow1 * mask_q_stride; #pragma unroll for (int n8 = 0; n8 < Traits::NC8; n8++) { int cc = kv0 + n8 * 8 + 2 * tid4; int c1 = cc + 1; - bool b0 = (cc >= maxc0) || (HasMask && !mask[mask_base0 + cc]); - bool b1 = (c1 >= maxc0) || (HasMask && !mask[mask_base0 + c1]); - bool b2 = (cc >= maxc1) || (HasMask && !mask[mask_base1 + cc]); - bool b3 = (c1 >= maxc1) || (HasMask && !mask[mask_base1 + c1]); + bool b0 = !valid0 || (cc >= maxc0) || (HasMask && !mask[mask_base0 + cc]); + bool b1 = !valid0 || (c1 >= maxc0) || (HasMask && !mask[mask_base0 + c1]); + bool b2 = !valid1 || (cc >= maxc1) || (HasMask && !mask[mask_base1 + cc]); + bool b3 = !valid1 || (c1 >= maxc1) || (HasMask && !mask[mask_base1 + c1]); float s0 = b0 ? -FLT_MAX : Sacc[n8][0]; float s1 = b1 ? -FLT_MAX : Sacc[n8][1]; float s2 = b2 ? -FLT_MAX : Sacc[n8][2]; diff --git a/csrc/kernels/attn_paged_decode.cu b/csrc/kernels/attn_paged_decode.cu index ba6c963..a0a7e41 100644 --- a/csrc/kernels/attn_paged_decode.cu +++ b/csrc/kernels/attn_paged_decode.cu @@ -27,12 +27,14 @@ torch::Tensor attn_paged_decode( torch::Tensor O; if (out_buf.has_value() && out_buf->defined()) { TORCH_CHECK(out_buf->dtype() == q.dtype(), "out_buf dtype must match q"); + TORCH_CHECK(out_buf->is_cuda() && out_buf->is_contiguous(), + "out_buf must be a contiguous CUDA tensor"); TORCH_CHECK(out_buf->size(0) >= q.size(0), "out_buf batch too small"); - TORCH_CHECK(out_buf->size(1) >= q.size(1), "out_buf heads too small"); - TORCH_CHECK(out_buf->size(2) >= q.size(2), "out_buf head_dim too small"); - O = out_buf.value().slice(0, 0, q.size(0)) - .slice(1, 0, q.size(1)) - .slice(2, 0, q.size(2)); + TORCH_CHECK(out_buf->size(1) == q.size(1), "out_buf heads must match q"); + TORCH_CHECK(out_buf->size(2) == q.size(2), "out_buf head_dim must match q"); + TORCH_CHECK(q.is_contiguous(), + "q must be contiguous when out_buf is provided"); + O = out_buf.value().slice(0, 0, q.size(0)); } else { O = torch::empty({q.size(0), q.size(1), q.size(2)}, q.options()); } @@ -43,8 +45,15 @@ torch::Tensor attn_paged_decode( TORCH_CHECK(o_part_buf->scalar_type() == torch::kFloat32, "o_part_buf must be f32"); TORCH_CHECK(ml_part_buf->scalar_type() == torch::kFloat32, "ml_part_buf must be f32"); int64_t o_needed = (int64_t)p.batch * p.q_head * MAX_SPLITS * p.head_dim; + int64_t ml_needed = (int64_t)p.batch * p.q_head * MAX_SPLITS * 2; TORCH_CHECK(o_part_buf->numel() >= o_needed, "o_part_buf too small: need ", o_needed, " got ", o_part_buf->numel()); + TORCH_CHECK(ml_part_buf->numel() >= ml_needed, + "ml_part_buf too small: need ", ml_needed, " got ", ml_part_buf->numel()); + TORCH_CHECK(o_part_buf->is_cuda() && ml_part_buf->is_cuda(), + "split buffers must be CUDA tensors"); + TORCH_CHECK(o_part_buf->is_contiguous() && ml_part_buf->is_contiguous(), + "split buffers must be contiguous"); p.o_part = (float*)o_part_buf->data_ptr(); p.ml_part = (float*)ml_part_buf->data_ptr(); } else { diff --git a/csrc/kernels/attn_prefill_split_q_mma.cuh b/csrc/kernels/attn_prefill_split_q_mma.cuh index b3b308b..3010825 100644 --- a/csrc/kernels/attn_prefill_split_q_mma.cuh +++ b/csrc/kernels/attn_prefill_split_q_mma.cuh @@ -123,8 +123,9 @@ __global__ void attn_prefill_split_q_mma_kernel(AttentionParams p) { mma_softmax_tile(kv0, maxc0, maxc1, qr0, qr1, p.mask_b_stride, p.mask_h_stride, p.mask_q_stride, - batch, q_head, + batch, q_head, q_head, p.mask, + va, vb, Sacc, Oacc, m0, m1, l0, l1, lane); mma_pv_accumulate(Sacc, bV, lane, Oacc); diff --git a/csrc/tests/attn_test.cu b/csrc/tests/attn_test.cu index c09bf29..9107044 100644 --- a/csrc/tests/attn_test.cu +++ b/csrc/tests/attn_test.cu @@ -323,6 +323,7 @@ int main() { // ---- PREFILL ---- { const int configs[][7] = { + {1,2,1,64,128,32,0}, // scalar fallback D=32 {1,2,1,64,128,64,0}, // tiny: B,Hq,Hk,q,kv,D,causal {1,32,4,512,512,128,0}, // standard {1,32,4,128,256,128,0}, // medium