diff --git a/csrc/jit_kernels/impls/smxx_fp8_fp4_paged_mqa_logits.hpp b/csrc/jit_kernels/impls/smxx_fp8_fp4_paged_mqa_logits.hpp index b6cf3ea0c..6fe087c42 100644 --- a/csrc/jit_kernels/impls/smxx_fp8_fp4_paged_mqa_logits.hpp +++ b/csrc/jit_kernels/impls/smxx_fp8_fp4_paged_mqa_logits.hpp @@ -14,6 +14,7 @@ class SMXXPagedMQALogitsMetadataRuntime final: public LaunchRuntime(&sched::smxx_paged_mqa_logits_metadata< - {}, {}, {}, {} + {}, {}, {}, {}, {} >); }}; -)", args.aligned_batch_size, args.split_kv, args.num_sms, args.is_varlen ? "true" : "false"); +)", args.aligned_batch_size, args.split_kv, args.num_sms, args.num_next_n_atoms, args.is_varlen ? "true" : "false"); } static void launch_impl(const KernelHandle& kernel, const LaunchConfigHandle& config, Args args) { @@ -61,6 +62,10 @@ static void smxx_paged_mqa_logits_metadata(const torch::Tensor& context_lens, constexpr int split_kv = 256; constexpr int num_threads = 32; const int aligned_batch_size = align(batch_size, 32); + const int next_n_atom = (is_varlen or next_n >= 2) ? 2 : 1; + // SM90 pads NextN=3 as one paired atom plus one single-token tail in the kernel. + const int num_next_n_atoms = (device_runtime->get_arch_major() == 9 and next_n == 3 and not is_varlen) + ? 1 : ceil_div(next_n, next_n_atom); DG_HOST_ASSERT(split_kv % block_kv == 0); // Shared memory: prefix_sum[kAlignedBatchSize] plus varlen atom metadata when needed. @@ -74,6 +79,7 @@ static void smxx_paged_mqa_logits_metadata(const torch::Tensor& context_lens, .aligned_batch_size = aligned_batch_size, .split_kv = split_kv, .num_sms = num_sms, + .num_next_n_atoms = num_next_n_atoms, .is_varlen = is_varlen, .batch_size = batch_size, .next_n = next_n, @@ -190,7 +196,8 @@ static void smxx_fp8_paged_mqa_logits(const torch::Tensor& q, const int mma_m = (device_runtime->get_arch_major() == 10 ? 128 : 64); const int num_math_warp_groups = split_kv / mma_m; const int num_math_threads = num_math_warp_groups * 128; - const int num_q_stages = 3, num_kv_stages = (device_runtime->get_arch_major() == 10 ? 4 : 3); + const int num_q_stages = (device_runtime->get_arch_major() == 9 and next_n == 3) ? 4 : 3; + const int num_kv_stages = (device_runtime->get_arch_major() == 10 ? 4 : 3); DG_HOST_ASSERT(split_kv % mma_m == 0 and logits_stride % split_kv == 0); // Construct TMAs @@ -217,8 +224,8 @@ static void smxx_fp8_paged_mqa_logits(const torch::Tensor& q, if (device_runtime->get_arch_major() == 9) { const int swizzle_alignment = head_dim * 8; - const int smem_q_size_per_stage = next_n * num_heads * head_dim * static_cast(q.element_size()); - const int aligned_smem_weight_size_per_stage = align(next_n * num_heads * static_cast(weights.element_size()), swizzle_alignment); + const int smem_q_size_per_stage = next_n_atom * num_heads * head_dim * static_cast(q.element_size()); + const int aligned_smem_weight_size_per_stage = align(next_n_atom * num_heads * static_cast(weights.element_size()), swizzle_alignment); const int smem_q_pipe_size = num_q_stages * (smem_q_size_per_stage + aligned_smem_weight_size_per_stage) + align(num_q_stages * 8 * 2, swizzle_alignment); const int smem_kv_size_per_stage = block_kv * head_dim * static_cast(kv_cache.element_size()); @@ -231,7 +238,7 @@ static void smxx_fp8_paged_mqa_logits(const torch::Tensor& q, smem_size = smem_q_pipe_size + num_math_warp_groups * smem_kv_pipe_size + smem_umma_barriers + smem_tmem_ptr; DG_HOST_ASSERT(smem_size <= SM90ArchSpec::smem_capacity); - DG_HOST_ASSERT(next_n == 1 or next_n == 2); + DG_HOST_ASSERT(next_n >= 1 and next_n <= 3); } else { const int smem_q_size_per_stage = next_n_atom * num_heads * head_dim * static_cast(q.element_size()); const int smem_kv_size_per_stage = split_kv * head_dim * static_cast(kv_cache.element_size()); diff --git a/csrc/jit_kernels/impls/smxx_fp8_paged_mqa_logits.hpp b/csrc/jit_kernels/impls/smxx_fp8_paged_mqa_logits.hpp index 1240aad84..360050e34 100644 --- a/csrc/jit_kernels/impls/smxx_fp8_paged_mqa_logits.hpp +++ b/csrc/jit_kernels/impls/smxx_fp8_paged_mqa_logits.hpp @@ -14,6 +14,7 @@ class SMXXPagedMQALogitsMetadataRuntime final: public LaunchRuntime(&smxx_paged_mqa_logits_metadata< - {}, {}, {} + auto ptr = reinterpret_cast(&sched::smxx_paged_mqa_logits_metadata< + {}, {}, {}, {} >); }}; -)", arch, args.aligned_batch_size, args.split_kv, args.num_sms); +)", arch, args.aligned_batch_size, args.split_kv, args.num_sms, args.num_next_n_atoms); } static void launch_impl(const KernelHandle& kernel, const LaunchConfigHandle& config, Args args) { @@ -46,6 +47,7 @@ static void __instantiate_kernel() {{ args.next_n, args.is_context_lens_2d, args.context_lens, + nullptr, args.schedule_metadata )); } @@ -60,6 +62,8 @@ static void smxx_paged_mqa_logits_metadata(const torch::Tensor& context_lens, constexpr int num_threads = 32; const int aligned_batch_size = align(batch_size, 32); const int split_kv = block_kv * num_math_warpgroups; + const int next_n_atom = (next_n >= 2) ? 2 : 1; + const int num_next_n_atoms = ceil_div(next_n, next_n_atom); // Calculate shared memory size const int smem_size = aligned_batch_size * static_cast(sizeof(int)); @@ -71,6 +75,7 @@ static void smxx_paged_mqa_logits_metadata(const torch::Tensor& context_lens, .aligned_batch_size = aligned_batch_size, .split_kv = split_kv, .num_sms = num_sms, + .num_next_n_atoms = num_next_n_atoms, .batch_size = batch_size, .next_n = next_n, .is_context_lens_2d = is_context_lens_2d, diff --git a/deep_gemm/include/deep_gemm/impls/sm90_fp8_paged_mqa_logits.cuh b/deep_gemm/include/deep_gemm/impls/sm90_fp8_paged_mqa_logits.cuh index cc2592bb4..3851ed8ec 100644 --- a/deep_gemm/include/deep_gemm/impls/sm90_fp8_paged_mqa_logits.cuh +++ b/deep_gemm/include/deep_gemm/impls/sm90_fp8_paged_mqa_logits.cuh @@ -39,7 +39,12 @@ void sm90_fp8_paged_mqa_logits(const uint32_t batch_size, DG_STATIC_ASSERT(not kIsVarlen, "Varlen is not supported for SM90 paged MQA logits"); // Types - using WGMMA = typename mma::sm90::FP8MMASelector::type; + // Odd NextN uses paired-token atoms for the common path and a single-token WGMMA tail. + static constexpr bool kPadOddN = (not kIsVarlen) and (kNextN % 2 == 1) and (kNextN >= 3); + static constexpr uint32_t kNextNAtom = (kIsVarlen or kNextN >= 2) ? 2 : 1; + static constexpr uint32_t kNumNextNAtoms = kPadOddN ? 1 : math::constexpr_ceil_div(kNextN, kNextNAtom); + using WGMMA = typename mma::sm90::FP8MMASelector::type; + using TailWGMMA = typename mma::sm90::FP8MMASelector::type; using Barrier = cutlass::arch::ClusterTransactionBarrier; // NOTES: use `__shfl_sync` to encourage NVCC to use unified registers @@ -61,8 +66,8 @@ void sm90_fp8_paged_mqa_logits(const uint32_t batch_size, // Shared memory configs static constexpr uint32_t kSwizzleAlignment = kHeadDim * 8; - static constexpr uint32_t SMEM_Q_SIZE_PER_STAGE = kNextN * kNumHeads * kHeadDim * sizeof(__nv_fp8_e4m3); - static constexpr uint32_t SMEM_WEIGHT_SIZE_PER_STAGE = kNextN * kNumHeads * sizeof(float); + static constexpr uint32_t SMEM_Q_SIZE_PER_STAGE = kNextNAtom * kNumHeads * kHeadDim * sizeof(__nv_fp8_e4m3); + static constexpr uint32_t SMEM_WEIGHT_SIZE_PER_STAGE = kNextNAtom * kNumHeads * sizeof(float); static constexpr uint32_t ALIGNED_SMEM_WEIGHT_SIZE_PER_STAGE = math::constexpr_align(SMEM_WEIGHT_SIZE_PER_STAGE, kSwizzleAlignment); static constexpr uint32_t SMEM_Q_PIPE_SIZE = kNumQStages * (SMEM_Q_SIZE_PER_STAGE + ALIGNED_SMEM_WEIGHT_SIZE_PER_STAGE) + math::constexpr_align(kNumQStages * 8 * 2, kSwizzleAlignment); @@ -135,8 +140,8 @@ void sm90_fp8_paged_mqa_logits(const uint32_t batch_size, cudaGridDependencySynchronize(); // Scheduler - auto scheduler = sched::PagedMQALogitsScheduler( - blockIdx.x, batch_size, context_lens, schedule_meta, indices); + using Scheduler = sched::PagedMQALogitsScheduler; + auto scheduler = Scheduler(blockIdx.x, batch_size, context_lens, schedule_meta, indices); DG_STATIC_ASSERT(SPLIT_KV % BLOCK_KV == 0, "Unaligned SPLIT_KV"); // Q and KV pipeline @@ -154,46 +159,72 @@ void sm90_fp8_paged_mqa_logits(const uint32_t batch_size, if (kv_group_idx >= kNumMathWarpGroups) return; - const auto issue_tma_q = [&](const uint32_t& stage_idx, const uint32_t& q_idx) { + const auto issue_tma_q = [&](const uint32_t& stage_idx, const uint32_t& q_atom_idx) { if (kv_group_idx == 0 and cute::elect_one_sync()) { - tma::copy(&tensor_map_q, full_q_barriers[stage_idx], smem_q[stage_idx], 0, q_idx * kNextN * kNumHeads); - tma::copy(&tensor_map_weights, full_q_barriers[stage_idx], smem_weights[stage_idx], 0, q_idx * kNextN); + const auto q_token_idx = Scheduler::atom_to_token_idx(q_atom_idx); + tma::copy(&tensor_map_q, full_q_barriers[stage_idx], smem_q[stage_idx], 0, q_token_idx * kNumHeads); + tma::copy(&tensor_map_weights, full_q_barriers[stage_idx], smem_weights[stage_idx], 0, q_token_idx); full_q_barriers[stage_idx]->arrive_and_expect_tx(SMEM_Q_SIZE_PER_STAGE + SMEM_WEIGHT_SIZE_PER_STAGE); } }; + const auto issue_tma_q_pair = [&](const uint32_t& stage_idx, const uint32_t& tail_stage_idx, const uint32_t& q_atom_idx) { + if (kv_group_idx == 0 and cute::elect_one_sync()) { + const auto q_token_idx = Scheduler::atom_to_token_idx(q_atom_idx); + tma::copy(&tensor_map_q, full_q_barriers[stage_idx], smem_q[stage_idx], 0, q_token_idx * kNumHeads); + tma::copy(&tensor_map_weights, full_q_barriers[stage_idx], smem_weights[stage_idx], 0, q_token_idx); + full_q_barriers[stage_idx]->arrive_and_expect_tx(SMEM_Q_SIZE_PER_STAGE + SMEM_WEIGHT_SIZE_PER_STAGE); + tma::copy(&tensor_map_q, full_q_barriers[tail_stage_idx], smem_q[tail_stage_idx], 0, (q_token_idx + 2) * kNumHeads); + tma::copy(&tensor_map_weights, full_q_barriers[tail_stage_idx], smem_weights[tail_stage_idx], 0, q_token_idx + 2); + full_q_barriers[tail_stage_idx]->arrive_and_expect_tx(SMEM_Q_SIZE_PER_STAGE + SMEM_WEIGHT_SIZE_PER_STAGE); + } + }; - // Initialize `q_idx` outside `[0, batch_size)` to indicate it was none - uint32_t q_idx = batch_size, kv_idx, num_kv; - uint32_t next_q_idx, next_kv_idx, next_num_kv; + // Initialize `q_atom_idx` outside the valid range to indicate it was none + uint32_t q_atom_idx = batch_size * kNumNextNAtoms, kv_idx, num_kv; + uint32_t next_q_atom_idx, next_kv_idx, next_num_kv; bool fetched_next_task; - // Prefetch the first Q - if ((fetched_next_task = scheduler.fetch_next_task(next_q_idx, next_kv_idx, next_num_kv))) - issue_tma_q(0, next_q_idx), q_iter_idx = 1; + // Prefetch the first Q atom + if ((fetched_next_task = scheduler.fetch_next_task(next_q_atom_idx, next_kv_idx, next_num_kv))) { + if constexpr (kPadOddN) + issue_tma_q_pair(0, 1, next_q_atom_idx), q_iter_idx = 2; + else + issue_tma_q(0, next_q_atom_idx), q_iter_idx = 1; + } int kv_block_idx_ptr = 32; uint32_t kv_block_idx_storage; while (fetched_next_task) { - // Prefetch next Q when current Q changes - bool prefetch_q = (q_idx != next_q_idx and scheduler.exist_q_atom_idx(next_q_idx + 1)); - q_idx = next_q_idx; + // Prefetch next Q when current atom changes + const auto next_advance = scheduler.get_atom_advance(next_q_atom_idx, batch_size * kNumNextNAtoms); + bool prefetch_q = (q_atom_idx != next_q_atom_idx and scheduler.exist_q_atom_idx(next_q_atom_idx + next_advance)); + q_atom_idx = next_q_atom_idx; kv_idx = next_kv_idx; num_kv = next_num_kv; // Wait Q consumer release and issue TMA Q if (prefetch_q) { - CUTE_TIE_DECL(get_q_pipeline(q_iter_idx ++), q_stage_idx, q_phase); - empty_q_barriers[q_stage_idx]->wait(q_phase ^ 1); - issue_tma_q(q_stage_idx, q_idx + 1); + if constexpr (kPadOddN) { + CUTE_TIE_DECL(get_q_pipeline(q_iter_idx ++), q_stage_idx, q_phase); + CUTE_TIE_DECL(get_q_pipeline(q_iter_idx ++), q_tail_stage_idx, q_tail_phase); + empty_q_barriers[q_stage_idx]->wait(q_phase ^ 1); + empty_q_barriers[q_tail_stage_idx]->wait(q_tail_phase ^ 1); + issue_tma_q_pair(q_stage_idx, q_tail_stage_idx, q_atom_idx + next_advance); + } else { + CUTE_TIE_DECL(get_q_pipeline(q_iter_idx ++), q_stage_idx, q_phase); + empty_q_barriers[q_stage_idx]->wait(q_phase ^ 1); + issue_tma_q(q_stage_idx, q_atom_idx + next_advance); + } } // Read KV block index // TODO: deal with `-1`? if (kv_idx == 0 or kv_block_idx_ptr == 32) { kv_block_idx_ptr = 0; + const auto block_table_offset = Scheduler::atom_to_block_table_row(q_atom_idx) * static_cast(block_table_stride); kv_block_idx_storage = (kv_idx + kv_group_idx + lane_idx * kNumMathWarpGroups < num_kv ? - block_table[q_idx * static_cast(block_table_stride) + (kv_idx + kv_group_idx + lane_idx * kNumMathWarpGroups)] : 0); + block_table[block_table_offset + (kv_idx + kv_group_idx + lane_idx * kNumMathWarpGroups)] : 0); } const auto kv_block_idx = __shfl_sync(0xffffffff, kv_block_idx_storage, kv_block_idx_ptr ++); @@ -211,121 +242,174 @@ void sm90_fp8_paged_mqa_logits(const uint32_t batch_size, } // Fetch next task - fetched_next_task = scheduler.fetch_next_task(next_q_idx, next_kv_idx, next_num_kv); + fetched_next_task = scheduler.fetch_next_task(next_q_atom_idx, next_kv_idx, next_num_kv); } } else { // Math warp-groups for WGMMA cutlass::arch::warpgroup_reg_alloc(); - float accum[WGMMA::kNumAccum], weights[kNextN][kNumHeads / 4]; + float accum[WGMMA::kNumAccum], weights[kNextNAtom][kNumHeads / 4], tail_weights[kNextNAtom][kNumHeads / 4]; const auto sub_warp_offset = (warp_idx % 4) * 16; const auto v_0_offset = lane_idx / 4 + 0; const auto v_1_offset = lane_idx / 4 + 8; - // Initialize `q_idx` outside `[0, batch_size)` to indicate it was none - uint32_t q_idx = batch_size, kv_idx; - uint32_t next_q_idx, next_kv_idx, next_num_kv; - uint32_t q_stage_idx, q_phase; + // Initialize `q_atom_idx` outside the valid range to indicate it was none + uint32_t q_atom_idx = batch_size * kNumNextNAtoms, kv_idx; + uint32_t next_q_atom_idx, next_kv_idx, next_num_kv; + uint32_t q_stage_idx, q_phase, q_tail_stage_idx, q_tail_phase; - while (scheduler.fetch_next_task(next_q_idx, next_kv_idx, next_num_kv)) { - // Current Q changes - if (q_idx != next_q_idx) { + while (scheduler.fetch_next_task(next_q_atom_idx, next_kv_idx, next_num_kv)) { + // Current Q atom changes + if (q_atom_idx != next_q_atom_idx) { // Release Last Q empty - if (q_iter_idx > 0) - empty_q_barriers[(q_iter_idx - 1) % kNumQStages]->arrive(); + if (q_iter_idx > 0) { + if constexpr (kPadOddN) { + empty_q_barriers[(q_iter_idx - 2) % kNumQStages]->arrive(); + empty_q_barriers[(q_iter_idx - 1) % kNumQStages]->arrive(); + } else { + empty_q_barriers[(q_iter_idx - 1) % kNumQStages]->arrive(); + } + } // Wait TMA Q arrival CUTE_TIE(get_q_pipeline(q_iter_idx ++), q_stage_idx, q_phase); full_q_barriers[q_stage_idx]->wait(q_phase); + if constexpr (kPadOddN) { + CUTE_TIE(get_q_pipeline(q_iter_idx ++), q_tail_stage_idx, q_tail_phase); + full_q_barriers[q_tail_stage_idx]->wait(q_tail_phase); + } // Read weights #pragma unroll - for (uint32_t i = 0; i < kNextN; ++ i) { + for (uint32_t i = 0; i < kNextNAtom; ++ i) { #pragma unroll for (uint32_t j = 0; j < kNumHeads / 4; ++ j) weights[i][j] = ptx::ld_shared(smem_weights[q_stage_idx] + i * kNumHeads + (j / 2) * 8 + (j & 1) + (lane_idx % 4) * 2); } + if constexpr (kPadOddN) { + #pragma unroll + for (uint32_t i = 0; i < kNextNAtom; ++ i) { + #pragma unroll + for (uint32_t j = 0; j < kNumHeads / 4; ++ j) + tail_weights[i][j] = ptx::ld_shared(smem_weights[q_tail_stage_idx] + i * kNumHeads + (j / 2) * 8 + (j & 1) + (lane_idx % 4) * 2); + } + } } - // Get current Q and KV index - q_idx = next_q_idx; + // Get current Q atom and KV index + q_atom_idx = next_q_atom_idx; kv_idx = next_kv_idx; // Calculate KV offset in advance - auto kv_offset = q_idx * kNextN * static_cast(logits_stride) + ((kv_idx + kv_group_idx) * BLOCK_KV + sub_warp_offset); + auto kv_offset = Scheduler::atom_to_token_idx(q_atom_idx) * static_cast(logits_stride) + ((kv_idx + kv_group_idx) * BLOCK_KV + sub_warp_offset); - // Compute `[kNextN * kNumHeads, kHeadDim] @ [BLOCK_KV, kHeadDim] -> [kNextN, BLOCK_KV]` + // Compute `[kNextNAtom * kNumHeads, kHeadDim] @ [BLOCK_KV, kHeadDim] -> [kNextNAtom, BLOCK_KV]` // Wait TMA KV arrival CUTE_TIE_DECL(get_kv_pipeline(kv_iter_idx ++), kv_stage_idx, kv_phase); full_kv_barriers[kv_stage_idx]->wait(kv_phase); - // Issue WGMMA DG_STATIC_ASSERT(BLOCK_KV == 64, "Invalid block size"); DG_STATIC_ASSERT(kHeadDim % WGMMA::K == 0, "Invalid head dim"); - #pragma unroll - for (uint32_t i = 0; i < WGMMA::kNumAccum; ++ i) - ptx::warpgroup_fence_operand(accum[i]); - ptx::warpgroup_arrive(); - #pragma unroll - for (uint32_t k = 0; k < kHeadDim / WGMMA::K; ++ k) { - auto desc_a = mma::sm90::make_smem_desc( - smem_kv[kv_stage_idx] + k * WGMMA::K, - mma::sm90::to_swizzle_cute_type(), 0, kHeadDim * 8); - auto desc_b = mma::sm90::make_smem_desc( - smem_q[q_stage_idx] + k * WGMMA::K, - mma::sm90::to_swizzle_cute_type(), 0, kHeadDim * 8); - WGMMA::wgmma(desc_a, desc_b, accum, k); - } - ptx::warpgroup_commit_batch(); - #pragma unroll - for (uint32_t i = 0; i < WGMMA::kNumAccum; ++ i) - ptx::warpgroup_fence_operand(accum[i]); - - // Read per-KV scales - float scale_kv_0 = ptx::ld_shared(smem_kv_scales[kv_stage_idx] + sub_warp_offset + v_0_offset); - float scale_kv_1 = ptx::ld_shared(smem_kv_scales[kv_stage_idx] + sub_warp_offset + v_1_offset); - - // Wait WGMMA - ptx::warpgroup_wait<0>(); - - // Release KV empty - empty_kv_barriers[kv_stage_idx]->arrive(); - - // Reduce over the head dim and store static constexpr uint32_t kNumAccumPerReduce = kNumHeads / 2; DG_STATIC_ASSERT(WGMMA::kNumAccum % kNumAccumPerReduce == 0, "Invalid accumulation"); - DG_STATIC_ASSERT(WGMMA::kNumAccum / kNumAccumPerReduce == kNextN, "Invalid accumulation"); + DG_STATIC_ASSERT(WGMMA::kNumAccum / kNumAccumPerReduce == kNextNAtom, "Invalid accumulation"); + DG_STATIC_ASSERT(TailWGMMA::kNumAccum == kNumAccumPerReduce, "Invalid tail accumulation"); DG_STATIC_ASSERT(kNumHeads % 8 == 0, "Invalid head"); - #pragma unroll - for (uint32_t i = 0; i < kNextN; ++ i) { - auto shifted_accum = accum + i * kNumAccumPerReduce; - const auto transform = [&](const uint32_t& j) { - return fmaxf(shifted_accum[j], 0) * weights[i][(j / 4) * 2 + (j & 1)]; - }; - - // Intra-thread reduction - float sum[4] = {transform(0), transform(1), transform(2), transform(3)}; + + const auto issue_wgmma = [&](const uint32_t& active_q_stage_idx) { #pragma unroll - for (uint32_t j = 1; j < kNumHeads / 8; ++ j) { - #pragma unroll - for (uint32_t k = 0; k < 4; k ++) - sum[k] += transform(j * 4 + k); + for (uint32_t i = 0; i < WGMMA::kNumAccum; ++ i) + ptx::warpgroup_fence_operand(accum[i]); + ptx::warpgroup_arrive(); + #pragma unroll + for (uint32_t k = 0; k < kHeadDim / WGMMA::K; ++ k) { + auto desc_a = mma::sm90::make_smem_desc( + smem_kv[kv_stage_idx] + k * WGMMA::K, + mma::sm90::to_swizzle_cute_type(), 0, kHeadDim * 8); + auto desc_b = mma::sm90::make_smem_desc( + smem_q[active_q_stage_idx] + k * WGMMA::K, + mma::sm90::to_swizzle_cute_type(), 0, kHeadDim * 8); + WGMMA::wgmma(desc_a, desc_b, accum, k); } - float v_0 = (sum[0] + sum[1]) * scale_kv_0; - float v_1 = (sum[2] + sum[3]) * scale_kv_1; - - // Inter-thread reduction + ptx::warpgroup_commit_batch(); + #pragma unroll + for (uint32_t i = 0; i < WGMMA::kNumAccum; ++ i) + ptx::warpgroup_fence_operand(accum[i]); + }; + const auto issue_tail_wgmma = [&](const uint32_t& active_q_stage_idx) { + #pragma unroll + for (uint32_t i = 0; i < TailWGMMA::kNumAccum; ++ i) + ptx::warpgroup_fence_operand(accum[i]); + ptx::warpgroup_arrive(); #pragma unroll - for (uint32_t j = 0; j < 2; ++ j) { - const auto offset = static_cast(1u << j); - v_0 += __shfl_xor_sync(0xffffffffu, v_0, offset); - v_1 += __shfl_xor_sync(0xffffffffu, v_1, offset); + for (uint32_t k = 0; k < kHeadDim / TailWGMMA::K; ++ k) { + auto desc_a = mma::sm90::make_smem_desc( + smem_kv[kv_stage_idx] + k * TailWGMMA::K, + mma::sm90::to_swizzle_cute_type(), 0, kHeadDim * 8); + auto desc_b = mma::sm90::make_smem_desc( + smem_q[active_q_stage_idx] + k * TailWGMMA::K, + mma::sm90::to_swizzle_cute_type(), 0, kHeadDim * 8); + TailWGMMA::wgmma(desc_a, desc_b, accum, k); } + ptx::warpgroup_commit_batch(); + #pragma unroll + for (uint32_t i = 0; i < TailWGMMA::kNumAccum; ++ i) + ptx::warpgroup_fence_operand(accum[i]); + }; + + // Read per-KV scales + float scale_kv_0 = ptx::ld_shared(smem_kv_scales[kv_stage_idx] + sub_warp_offset + v_0_offset); + float scale_kv_1 = ptx::ld_shared(smem_kv_scales[kv_stage_idx] + sub_warp_offset + v_1_offset); - // Store into the global memory - // NOTES: we have redundant writes here, consider more carefully - logits[kv_offset + i * static_cast(logits_stride) + v_0_offset] = static_cast(v_0); - logits[kv_offset + i * static_cast(logits_stride) + v_1_offset] = static_cast(v_1); + const auto reduce_and_store = [&](auto num_iters_c, const auto& active_weights, const uint64_t& active_kv_offset) { + constexpr uint32_t kNumIters = decltype(num_iters_c)::value; + #pragma unroll + for (uint32_t i = 0; i < kNumIters; ++ i) { + auto shifted_accum = accum + i * kNumAccumPerReduce; + const auto transform = [&](const uint32_t& j) { + return fmaxf(shifted_accum[j], 0) * active_weights[i][(j / 4) * 2 + (j & 1)]; + }; + + // Intra-thread reduction + float sum[4] = {transform(0), transform(1), transform(2), transform(3)}; + #pragma unroll + for (uint32_t j = 1; j < kNumHeads / 8; ++ j) { + #pragma unroll + for (uint32_t k = 0; k < 4; k ++) + sum[k] += transform(j * 4 + k); + } + float v_0 = (sum[0] + sum[1]) * scale_kv_0; + float v_1 = (sum[2] + sum[3]) * scale_kv_1; + + // Inter-thread reduction + #pragma unroll + for (uint32_t j = 0; j < 2; ++ j) { + const auto offset = static_cast(1u << j); + v_0 += __shfl_xor_sync(0xffffffffu, v_0, offset); + v_1 += __shfl_xor_sync(0xffffffffu, v_1, offset); + } + + // Store into the global memory + // NOTES: we have redundant writes here, consider more carefully + logits[active_kv_offset + i * static_cast(logits_stride) + v_0_offset] = static_cast(v_0); + logits[active_kv_offset + i * static_cast(logits_stride) + v_1_offset] = static_cast(v_1); + } + }; + + if constexpr (kPadOddN) { + issue_wgmma(q_stage_idx); + ptx::warpgroup_wait<0>(); + reduce_and_store(cute::Int{}, weights, kv_offset); + + issue_tail_wgmma(q_tail_stage_idx); + ptx::warpgroup_wait<0>(); + empty_kv_barriers[kv_stage_idx]->arrive(); + reduce_and_store(cute::Int<1>{}, tail_weights, kv_offset + 2 * static_cast(logits_stride)); + } else { + issue_wgmma(q_stage_idx); + ptx::warpgroup_wait<0>(); + empty_kv_barriers[kv_stage_idx]->arrive(); + reduce_and_store(cute::Int{}, weights, kv_offset); } } } diff --git a/deep_gemm/include/deep_gemm/scheduler/paged_mqa_logits.cuh b/deep_gemm/include/deep_gemm/scheduler/paged_mqa_logits.cuh index 6de3f8ea8..841c61117 100644 --- a/deep_gemm/include/deep_gemm/scheduler/paged_mqa_logits.cuh +++ b/deep_gemm/include/deep_gemm/scheduler/paged_mqa_logits.cuh @@ -6,7 +6,8 @@ namespace deep_gemm::sched { -template +template CUTLASS_GLOBAL __launch_bounds__(32, 1) void smxx_paged_mqa_logits_metadata(const uint32_t batch_size, const uint32_t next_n, const bool is_context_lens_2d, const uint32_t* context_lens, const uint32_t* indices, uint32_t* schedule_metadata) { @@ -94,8 +95,7 @@ void smxx_paged_mqa_logits_metadata(const uint32_t batch_size, const uint32_t ne schedule_metadata[sm_idx * 2 + 1] = kv_split_idx; } } else { - const uint32_t next_n_atom = (next_n >= 2) ? 2 : 1; - const uint32_t num_next_n_atoms = math::ceil_div(next_n, next_n_atom); + const uint32_t num_next_n_atoms = kNumNextNAtoms; const uint32_t total = sum * num_next_n_atoms; const uint32_t q = total / kNumSMs, r = total % kNumSMs; const uint32_t pivot = kNumSMs - r; diff --git a/tests/test_attention.py b/tests/test_attention.py index 769923a5c..3c93eb508 100644 --- a/tests/test_attention.py +++ b/tests/test_attention.py @@ -273,7 +273,13 @@ def enumerate_paged_mqa_logits(): for block_kv in ((32, 64) if arch_major == 10 else (64, )): for use_2d_context_lens, clean_logits in [(True, False)]: for batch_size in (256, 4096): - for next_n in ((1, ) if is_varlen else ((1, 2, 4, 5, 6) if arch_major == 10 else (1, 2))): + if is_varlen: + next_ns = (1, ) + elif arch_major == 10: + next_ns = (1, 2, 4, 5, 6) + else: + next_ns = (1, 2, 3) + for next_n in next_ns: for max_tokens_per_batch in ((1, 4, 10) if is_varlen else (1, )): for num_heads, head_dim in [(64, 128), (32, 128)]: for avg_kv in (8192, 32768):