64 lines
2.6 KiB
C++
64 lines
2.6 KiB
C++
#pragma once
|
|
|
|
#include <cute/algorithm/copy.hpp>
|
|
|
|
#include <mctlass/mctlass.h>
|
|
#include <mctlass/array.h>
|
|
#include <mctlass/numeric_types.h>
|
|
|
|
#include "block_info.h"
|
|
#include "kernel_traits.h"
|
|
#include "utils.h"
|
|
#include "softmax.h"
|
|
#include "mask.h"
|
|
|
|
#include "xcore1000/flash_fwd_sparse_mla_kernel_k64_64x16_8waves_xcore1000.h"
|
|
|
|
|
|
|
|
namespace flash {
|
|
|
|
using namespace cute;
|
|
|
|
template<typename Kernel_traits, bool Is_causal, bool Is_even_K, typename Params>
|
|
__forceinline__ __device__ void compute_attn_1rowblock_splitkv_sparse_mla(const Params ¶ms, const int m_block_max) {
|
|
constexpr int kBlockN = Kernel_traits::kBlockN;
|
|
const int m_block = blockIdx.x;
|
|
const int bidh = blockIdx.y;
|
|
const int partition_idx = blockIdx.z;
|
|
|
|
extern __shared__ char shared_memory[];
|
|
|
|
int *tile_scheduler_metadata_ptr = params.tile_scheduler_metadata_ptr + partition_idx * TileSchedulerMetaDataSize;
|
|
int4 tile_scheduler_metadata = *reinterpret_cast<int4 *>(tile_scheduler_metadata_ptr);
|
|
int begin_idx = tile_scheduler_metadata.x;
|
|
int sched_begin_block_idx = tile_scheduler_metadata.y;
|
|
int end_idx = tile_scheduler_metadata.z;
|
|
int sched_end_block_idx = tile_scheduler_metadata.w;
|
|
if (begin_idx >= params.b || begin_idx < 0) return;
|
|
int begin_n_split_idx = tile_scheduler_metadata_ptr[4];
|
|
|
|
#pragma unroll 1
|
|
for (int batch_id = begin_idx; batch_id <= end_idx; ++batch_id) {
|
|
const int n_split_idx = batch_id == begin_idx ? begin_n_split_idx : 0;
|
|
const int seqlen_k = params.cu_seqlens_k[batch_id];
|
|
const int n_block_min = batch_id == begin_idx ? sched_begin_block_idx : 0;
|
|
const int n_block_max = batch_id == end_idx ? sched_end_block_idx : cute::ceil_div(params.topk, kBlockN);
|
|
// [n_block_min, n_block_max) need be calculated in kernel
|
|
if (n_block_max <= n_block_min) continue;
|
|
const bool NoSplit = __ldg(params.num_splits_ptr + batch_id + 1) - __ldg(params.num_splits_ptr + batch_id) == 1;
|
|
if (batch_id > begin_idx) {
|
|
__syncthreads(); // Barrier between two tiles.
|
|
}
|
|
|
|
// printf(
|
|
// "batch_id is %d, begin_idx is %d, end_idx is %d, params.b is %d, n_split_idx is %d, seqlen_k is %d, n_block_min is %d, n_block_max is %d, NoSplit is %d\n",
|
|
// batch_id, begin_idx, end_idx, params.b, n_split_idx, seqlen_k, n_block_min, n_block_max, NoSplit);
|
|
|
|
compute_attn_1rowblock_splitkv_sparse_mla_k64_64x16_8waves_xcore1000<Kernel_traits, Is_causal, Is_even_K>(
|
|
params, batch_id, bidh, m_block, n_split_idx, n_block_min, n_block_max, NoSplit);
|
|
}
|
|
}
|
|
|
|
}// namespace flash
|