Skip to content

GitLab

"vscode:/vscode.git/clone" did not exist on "8b491162f79282b79292e0f4166194b3f7ccb7b1"
Commit f06d6465 authored by zhushuang's avatar zhushuang
Browse files

issue/1041 - feat: use template to replace int64_t in paged_attention_prefill kernel with test pass

parent abd45713
Expand all Show whitespace changes
Inline Side-by-side
Showing with 185 additions and 163 deletions
src/infiniop/ops/paged_attention_prefill/cuda/kernel_v2.cuh
View file @ f06d6465
...@@ -16,7 +16,8 @@ ...@@ -16,7 +16,8 @@
namespace op::paged_attention_prefill::cuda { namespace op::paged_attention_prefill::cuda {
__device__ __forceinline__ size_t find_seq_id(size_t token_idx, const int64_t *cu_seqlens_q, size_t num_seqs) { template <typename Tindex>
__device__ __forceinline__ size_t find_seq_id(size_t token_idx, const Tindex *cu_seqlens_q, size_t num_seqs) {
size_t low = 0, high = (num_seqs == 0) ? 0 : (num_seqs - 1); size_t low = 0, high = (num_seqs == 0) ? 0 : (num_seqs - 1);
while (low <= high) { while (low <= high) {
size_t mid = (low + high) >> 1; size_t mid = (low + high) >> 1;
...@@ -43,8 +44,8 @@ __device__ void PagedAttentionPrefillWarpKernel( ...@@ -43,8 +44,8 @@ __device__ void PagedAttentionPrefillWarpKernel(
const Tdata *k_cache_, const Tdata *k_cache_,
const Tdata *v_cache_, const Tdata *v_cache_,
const Tindex *block_tables_, const Tindex *block_tables_,
const int64_t *total_kv_lens_, const Tindex *total_kv_lens_,
const int64_t *cu_seqlens_q_, const Tindex *cu_seqlens_q_,
const float *alibi_slopes_, const float *alibi_slopes_,
size_t num_kv_heads, size_t num_kv_heads,
float scale, float scale,
...@@ -73,8 +74,8 @@ __device__ void PagedAttentionPrefillWarpKernel( ...@@ -73,8 +74,8 @@ __device__ void PagedAttentionPrefillWarpKernel(
const int seq_idx = static_cast<int>(blockIdx.y); const int seq_idx = static_cast<int>(blockIdx.y);
const int q_token_local = static_cast<int>(blockIdx.z); const int q_token_local = static_cast<int>(blockIdx.z);
const int64_t q_start = cu_seqlens_q_[seq_idx]; const Tindex q_start = cu_seqlens_q_[seq_idx];
const int64_t q_end = cu_seqlens_q_[seq_idx + 1]; const Tindex q_end = cu_seqlens_q_[seq_idx + 1];
const int q_len = static_cast<int>(q_end - q_start); const int q_len = static_cast<int>(q_end - q_start);
if (q_token_local >= q_len) { if (q_token_local >= q_len) {
return; return;
...@@ -256,8 +257,8 @@ __global__ void PagedAttentionPrefillWarpGlobalKernel( ...@@ -256,8 +257,8 @@ __global__ void PagedAttentionPrefillWarpGlobalKernel(
const Tdata *k_cache_, const Tdata *k_cache_,
const Tdata *v_cache_, const Tdata *v_cache_,
const Tindex *block_tables_, const Tindex *block_tables_,
const int64_t *total_kv_lens_, const Tindex *total_kv_lens_,
const int64_t *cu_seqlens_q_, const Tindex *cu_seqlens_q_,
const float *alibi_slopes_, const float *alibi_slopes_,
size_t num_heads, size_t num_heads,
size_t num_seqs, size_t num_seqs,
...@@ -291,9 +292,9 @@ __global__ void PagedAttentionPrefillWarpGlobalKernel( ...@@ -291,9 +292,9 @@ __global__ void PagedAttentionPrefillWarpGlobalKernel(
return; return;
} }
const size_t seq_idx = find_seq_id(global_token_idx, cu_seqlens_q_, num_seqs); const size_t seq_idx = find_seq_id<Tindex>(global_token_idx, cu_seqlens_q_, num_seqs);
const int64_t q_start = cu_seqlens_q_[seq_idx]; const Tindex q_start = cu_seqlens_q_[seq_idx];
const int64_t q_end = cu_seqlens_q_[seq_idx + 1]; const Tindex q_end = cu_seqlens_q_[seq_idx + 1];
const int q_len = static_cast<int>(q_end - q_start); const int q_len = static_cast<int>(q_end - q_start);
const int q_token_local = static_cast<int>(global_token_idx - static_cast<size_t>(q_start)); const int q_token_local = static_cast<int>(global_token_idx - static_cast<size_t>(q_start));
...@@ -477,8 +478,8 @@ __global__ void PagedAttentionPrefillReferenceKernel( ...@@ -477,8 +478,8 @@ __global__ void PagedAttentionPrefillReferenceKernel(
const Tdata *k_cache_, const Tdata *k_cache_,
const Tdata *v_cache_, const Tdata *v_cache_,
const Tindex *block_tables_, const Tindex *block_tables_,
const int64_t *total_kv_lens_, const Tindex *total_kv_lens_,
const int64_t *cu_seqlens_q_, const Tindex *cu_seqlens_q_,
const float *alibi_slopes_, const float *alibi_slopes_,
size_t num_heads, size_t num_heads,
size_t num_kv_heads, size_t num_kv_heads,
...@@ -506,7 +507,7 @@ __global__ void PagedAttentionPrefillReferenceKernel( ...@@ -506,7 +507,7 @@ __global__ void PagedAttentionPrefillReferenceKernel(
return; return;
} }
const size_t seq_idx = find_seq_id(global_token_idx, cu_seqlens_q_, num_seqs); const size_t seq_idx = find_seq_id<Tindex>(global_token_idx, cu_seqlens_q_, num_seqs);
const size_t q_token_idx = global_token_idx - static_cast<size_t>(cu_seqlens_q_[seq_idx]); const size_t q_token_idx = global_token_idx - static_cast<size_t>(cu_seqlens_q_[seq_idx]);
const size_t q_len = static_cast<size_t>(cu_seqlens_q_[seq_idx + 1] - cu_seqlens_q_[seq_idx]); const size_t q_len = static_cast<size_t>(cu_seqlens_q_[seq_idx + 1] - cu_seqlens_q_[seq_idx]);
...@@ -595,8 +596,8 @@ __device__ void PagedAttentionPrefillWarpCtaKernel( ...@@ -595,8 +596,8 @@ __device__ void PagedAttentionPrefillWarpCtaKernel(
const Tdata *k_cache_, const Tdata *k_cache_,
const Tdata *v_cache_, const Tdata *v_cache_,
const Tindex *block_tables_, const Tindex *block_tables_,
const int64_t *total_kv_lens_, const Tindex *total_kv_lens_,
const int64_t *cu_seqlens_q_, const Tindex *cu_seqlens_q_,
const float *alibi_slopes_, const float *alibi_slopes_,
size_t num_kv_heads, size_t num_kv_heads,
float scale, float scale,
...@@ -632,8 +633,8 @@ __device__ void PagedAttentionPrefillWarpCtaKernel( ...@@ -632,8 +633,8 @@ __device__ void PagedAttentionPrefillWarpCtaKernel(
const int seq_idx = static_cast<int>(blockIdx.y); const int seq_idx = static_cast<int>(blockIdx.y);
const int m_block = static_cast<int>(blockIdx.z); const int m_block = static_cast<int>(blockIdx.z);
const int64_t q_start = cu_seqlens_q_[seq_idx]; const Tindex q_start = cu_seqlens_q_[seq_idx];
const int64_t q_end = cu_seqlens_q_[seq_idx + 1]; const Tindex q_end = cu_seqlens_q_[seq_idx + 1];
const int q_len = static_cast<int>(q_end - q_start); const int q_len = static_cast<int>(q_end - q_start);
if (q_len <= 0) { if (q_len <= 0) {
return; return;
...@@ -865,8 +866,8 @@ __device__ void PagedAttentionPrefillWarpCtaKernelPipelined( ...@@ -865,8 +866,8 @@ __device__ void PagedAttentionPrefillWarpCtaKernelPipelined(
const Tdata *k_cache_, const Tdata *k_cache_,
const Tdata *v_cache_, const Tdata *v_cache_,
const Tindex *block_tables_, const Tindex *block_tables_,
const int64_t *total_kv_lens_, const Tindex *total_kv_lens_,
const int64_t *cu_seqlens_q_, const Tindex *cu_seqlens_q_,
const float *alibi_slopes_, const float *alibi_slopes_,
size_t num_kv_heads, size_t num_kv_heads,
float scale, float scale,
...@@ -904,8 +905,8 @@ __device__ void PagedAttentionPrefillWarpCtaKernelPipelined( ...@@ -904,8 +905,8 @@ __device__ void PagedAttentionPrefillWarpCtaKernelPipelined(
const int seq_idx = static_cast<int>(blockIdx.y); const int seq_idx = static_cast<int>(blockIdx.y);
const int m_block = static_cast<int>(blockIdx.z); const int m_block = static_cast<int>(blockIdx.z);
const int64_t q_start = cu_seqlens_q_[seq_idx]; const Tindex q_start = cu_seqlens_q_[seq_idx];
const int64_t q_end = cu_seqlens_q_[seq_idx + 1]; const Tindex q_end = cu_seqlens_q_[seq_idx + 1];
const int q_len = static_cast<int>(q_end - q_start); const int q_len = static_cast<int>(q_end - q_start);
if (q_len <= 0) { if (q_len <= 0) {
return; return;
...@@ -1312,8 +1313,8 @@ __device__ void PagedAttentionPrefillWarpCtaKernelPipelinedSplitKv( ...@@ -1312,8 +1313,8 @@ __device__ void PagedAttentionPrefillWarpCtaKernelPipelinedSplitKv(
const Tdata *k_cache_, const Tdata *k_cache_,
const Tdata *v_cache_, const Tdata *v_cache_,
const Tindex *block_tables_, const Tindex *block_tables_,
const int64_t *total_kv_lens_, const Tindex *total_kv_lens_,
const int64_t *cu_seqlens_q_, const Tindex *cu_seqlens_q_,
const float *alibi_slopes_, const float *alibi_slopes_,
size_t num_kv_heads, size_t num_kv_heads,
float scale, float scale,
...@@ -1350,8 +1351,8 @@ __device__ void PagedAttentionPrefillWarpCtaKernelPipelinedSplitKv( ...@@ -1350,8 +1351,8 @@ __device__ void PagedAttentionPrefillWarpCtaKernelPipelinedSplitKv(
const int head_idx = static_cast<int>(blockIdx.x); const int head_idx = static_cast<int>(blockIdx.x);
const int seq_idx = static_cast<int>(blockIdx.y); const int seq_idx = static_cast<int>(blockIdx.y);
const int64_t q_start = cu_seqlens_q_[seq_idx]; const Tindex q_start = cu_seqlens_q_[seq_idx];
const int64_t q_end = cu_seqlens_q_[seq_idx + 1]; const Tindex q_end = cu_seqlens_q_[seq_idx + 1];
const int q_len = static_cast<int>(q_end - q_start); const int q_len = static_cast<int>(q_end - q_start);
if (q_len <= 0) { if (q_len <= 0) {
return; return;
...@@ -1778,8 +1779,8 @@ __device__ void PagedAttentionPrefillWarpCtaKernelKOnly( ...@@ -1778,8 +1779,8 @@ __device__ void PagedAttentionPrefillWarpCtaKernelKOnly(
const Tdata *k_cache_, const Tdata *k_cache_,
const Tdata *v_cache_, const Tdata *v_cache_,
const Tindex *block_tables_, const Tindex *block_tables_,
const int64_t *total_kv_lens_, const Tindex *total_kv_lens_,
const int64_t *cu_seqlens_q_, const Tindex *cu_seqlens_q_,
const float *alibi_slopes_, const float *alibi_slopes_,
size_t num_kv_heads, size_t num_kv_heads,
float scale, float scale,
...@@ -1815,8 +1816,8 @@ __device__ void PagedAttentionPrefillWarpCtaKernelKOnly( ...@@ -1815,8 +1816,8 @@ __device__ void PagedAttentionPrefillWarpCtaKernelKOnly(
const int seq_idx = static_cast<int>(blockIdx.y); const int seq_idx = static_cast<int>(blockIdx.y);
const int m_block = static_cast<int>(blockIdx.z); const int m_block = static_cast<int>(blockIdx.z);
const int64_t q_start = cu_seqlens_q_[seq_idx]; const Tindex q_start = cu_seqlens_q_[seq_idx];
const int64_t q_end = cu_seqlens_q_[seq_idx + 1]; const Tindex q_end = cu_seqlens_q_[seq_idx + 1];
const int q_len = static_cast<int>(q_end - q_start); const int q_len = static_cast<int>(q_end - q_start);
if (q_len <= 0) { if (q_len <= 0) {
return; return;
...@@ -2115,12 +2116,12 @@ __device__ __forceinline__ void PagedAttentionPrefillMmaScoreUpdateRow( ...@@ -2115,12 +2116,12 @@ __device__ __forceinline__ void PagedAttentionPrefillMmaScoreUpdateRow(
} }
} }
template <int kWarpSize, int kHeadDim, int kDimsPerThread> template <typename Tindex, int kWarpSize, int kHeadDim, int kDimsPerThread>
__device__ __forceinline__ void PagedAttentionPrefillMmaScoreWriteRow( __device__ __forceinline__ void PagedAttentionPrefillMmaScoreWriteRow(
int lane, int lane,
bool active, bool active,
int q_token_local, int q_token_local,
int64_t q_start, Tindex q_start,
int head_idx, int head_idx,
half *out_, half *out_,
ptrdiff_t o_stride, ptrdiff_t o_stride,
...@@ -2153,8 +2154,8 @@ __device__ void PagedAttentionPrefillWarpCta8MmaHd128Kernel( ...@@ -2153,8 +2154,8 @@ __device__ void PagedAttentionPrefillWarpCta8MmaHd128Kernel(
const half *k_cache_, const half *k_cache_,
const half *v_cache_, const half *v_cache_,
const Tindex *block_tables_, const Tindex *block_tables_,
const int64_t *total_kv_lens_, const Tindex *total_kv_lens_,
const int64_t *cu_seqlens_q_, const Tindex *cu_seqlens_q_,
const float *alibi_slopes_, const float *alibi_slopes_,
size_t num_kv_heads, size_t num_kv_heads,
float scale, float scale,
...@@ -2198,8 +2199,8 @@ __device__ void PagedAttentionPrefillWarpCta8MmaHd128Kernel( ...@@ -2198,8 +2199,8 @@ __device__ void PagedAttentionPrefillWarpCta8MmaHd128Kernel(
const int seq_idx = static_cast<int>(blockIdx.y); const int seq_idx = static_cast<int>(blockIdx.y);
const int m_block = static_cast<int>(blockIdx.z); const int m_block = static_cast<int>(blockIdx.z);
const int64_t q_start = cu_seqlens_q_[seq_idx]; const Tindex q_start = cu_seqlens_q_[seq_idx];
const int64_t q_end = cu_seqlens_q_[seq_idx + 1]; const Tindex q_end = cu_seqlens_q_[seq_idx + 1];
const int q_len = static_cast<int>(q_end - q_start); const int q_len = static_cast<int>(q_end - q_start);
if (q_len <= 0) { if (q_len <= 0) {
return; return;
...@@ -2353,11 +2354,11 @@ __device__ void PagedAttentionPrefillWarpCta8MmaHd128Kernel( ...@@ -2353,11 +2354,11 @@ __device__ void PagedAttentionPrefillWarpCta8MmaHd128Kernel(
// Write outputs. // Write outputs.
if (row0 < kBlockM) { if (row0 < kBlockM) {
PagedAttentionPrefillMmaScoreWriteRow<kWarpSize, kHeadDim, kDimsPerThread>( PagedAttentionPrefillMmaScoreWriteRow<Tindex, kWarpSize, kHeadDim, kDimsPerThread>(
lane, active0, m_start + row0, q_start, head_idx, out_, o_stride, o_head_stride, l0, acc0); lane, active0, m_start + row0, q_start, head_idx, out_, o_stride, o_head_stride, l0, acc0);
} }
if (row1 < kBlockM) { if (row1 < kBlockM) {
PagedAttentionPrefillMmaScoreWriteRow<kWarpSize, kHeadDim, kDimsPerThread>( PagedAttentionPrefillMmaScoreWriteRow<Tindex, kWarpSize, kHeadDim, kDimsPerThread>(
lane, active1, m_start + row1, q_start, head_idx, out_, o_stride, o_head_stride, l1, acc1); lane, active1, m_start + row1, q_start, head_idx, out_, o_stride, o_head_stride, l1, acc1);
} }
} }
......
src/infiniop/ops/paged_attention_prefill/info.h
View file @ f06d6465
...@@ -80,9 +80,13 @@ public: ...@@ -80,9 +80,13 @@ public:
return INFINI_STATUS_BAD_TENSOR_DTYPE; return INFINI_STATUS_BAD_TENSOR_DTYPE;
} }
// Keep it simple: require total_kv_lens + cum_seqlens_q to be int64 for now // Index tensors use int32_t to match mainstream paged-attention implementations
// (matches current paged_attention_prefill signature). We will convert to int32 internally later. // (e.g., vLLM / FlashAttention2). 32-bit indices needed, but now we also support int64_t.
if (total_kv_lens_desc->dtype() != INFINI_DTYPE_I64 || cum_seqlens_q_desc->dtype() != INFINI_DTYPE_I64) { if (!((total_kv_lens_desc->dtype() == INFINI_DTYPE_I64) || (total_kv_lens_desc->dtype() == INFINI_DTYPE_I32) || (total_kv_lens_desc->dtype() == INFINI_DTYPE_U32))) {
return INFINI_STATUS_BAD_TENSOR_DTYPE;
}
if (!((cum_seqlens_q_desc->dtype() == INFINI_DTYPE_I64) || (cum_seqlens_q_desc->dtype() == INFINI_DTYPE_I32) || (cum_seqlens_q_desc->dtype() == INFINI_DTYPE_U32))) {
return INFINI_STATUS_BAD_TENSOR_DTYPE; return INFINI_STATUS_BAD_TENSOR_DTYPE;
} }
......
test/infinicore/ops/paged_attention_prefill.py
View file @ f06d6465
...@@ -31,6 +31,8 @@ _TOLERANCE_MAP = { ...@@ -31,6 +31,8 @@ _TOLERANCE_MAP = {
_TENSOR_DTYPES = [infinicore.float16, infinicore.bfloat16] _TENSOR_DTYPES = [infinicore.float16, infinicore.bfloat16]
_INDEX_DTYPES = [infinicore.int32, infinicore.int64]
class SimpleCacheManager: class SimpleCacheManager:
def __init__(self, num_blocks, block_size): def __init__(self, num_blocks, block_size):
...@@ -72,16 +74,16 @@ def parse_test_cases(): ...@@ -72,16 +74,16 @@ def parse_test_cases():
scale = head_size**-0.5 scale = head_size**-0.5
num_blocks = 8192 num_blocks = 8192
manager = SimpleCacheManager(num_blocks, block_size) manager = SimpleCacheManager(num_blocks, block_size)
kv_lens = torch.zeros(num_seqs, dtype=torch.int64) kv_lens = torch.zeros(num_seqs, dtype=torch.int32)
persistent_k = torch.zeros((num_blocks, num_kv_heads, block_size, head_size)) persistent_k = torch.zeros((num_blocks, num_kv_heads, block_size, head_size))
persistent_v = torch.zeros((num_blocks, num_kv_heads, block_size, head_size)) persistent_v = torch.zeros((num_blocks, num_kv_heads, block_size, head_size))
for r in range(num_rounds): for r in range(num_rounds):
q_lens = torch.randint(1, max_step_len + 1, (num_seqs,), dtype=torch.int64) q_lens = torch.randint(1, max_step_len + 1, (num_seqs,), dtype=torch.int32)
kv_lens = kv_lens + q_lens kv_lens = kv_lens + q_lens
total_q_tokens = q_lens.sum().item() total_q_tokens = q_lens.sum().item()
cum_seqlens_q = torch.zeros(num_seqs + 1, dtype=torch.int64) cum_seqlens_q = torch.zeros(num_seqs + 1, dtype=torch.int32)
cum_seqlens_q[1:] = torch.cumsum(q_lens, dim=0) cum_seqlens_q[1:] = torch.cumsum(q_lens, dim=0)
query_base = torch.randn((total_q_tokens, num_heads, head_size)) query_base = torch.randn((total_q_tokens, num_heads, head_size))
...@@ -106,8 +108,8 @@ def parse_test_cases(): ...@@ -106,8 +108,8 @@ def parse_test_cases():
) )
for dtype in _TENSOR_DTYPES: for dtype in _TENSOR_DTYPES:
for idx_dtype in _INDEX_DTYPES: # Loop through both I32 and I64
tolerance = _TOLERANCE_MAP.get(dtype) tolerance = _TOLERANCE_MAP.get(dtype)
test_cases.append( test_cases.append(
TestCase( TestCase(
inputs=[ inputs=[
...@@ -133,19 +135,19 @@ def parse_test_cases(): ...@@ -133,19 +135,19 @@ def parse_test_cases():
padded_tables.shape, padded_tables.shape,
init_mode=TensorInitializer.MANUAL, init_mode=TensorInitializer.MANUAL,
set_tensor=padded_tables.clone(), set_tensor=padded_tables.clone(),
dtype=infinicore.int64, dtype=idx_dtype,
), ),
TensorSpec.from_tensor( TensorSpec.from_tensor(
kv_lens.shape, kv_lens.shape,
init_mode=TensorInitializer.MANUAL, init_mode=TensorInitializer.MANUAL,
set_tensor=kv_lens.clone(), set_tensor=kv_lens.clone(),
dtype=infinicore.int64, dtype=idx_dtype,
), ),
TensorSpec.from_tensor( TensorSpec.from_tensor(
cum_seqlens_q.shape, cum_seqlens_q.shape,
init_mode=TensorInitializer.MANUAL, init_mode=TensorInitializer.MANUAL,
set_tensor=cum_seqlens_q.clone(), set_tensor=cum_seqlens_q.clone(),
dtype=infinicore.int64, dtype=idx_dtype,
), ),
], ],
kwargs={"scale": scale}, kwargs={"scale": scale},
......
test/infiniop/paged_attention_prefill.py
View file @ f06d6465
...@@ -23,13 +23,20 @@ from libinfiniop import ( ...@@ -23,13 +23,20 @@ from libinfiniop import (
# Configuration (Internal Use Only) # Configuration (Internal Use Only)
# ============================================================================== # ==============================================================================
_TEST_CASES = [ _TEST_CASES = [
# num_seqs, num_heads, num_kv_heads, head_size, block_size, max_step_len, num_rounds # num_seqs, num_heads, num_kv_heads, head_size, block_size, max_step_len, num_rounds, index_dtypes
(1, 1, 1, 128, 8, 16, 1), # index_dtype: The data type used for memory indexing of block_tables, cum_seq_lens and seq_lens
(1, 4, 4, 128, 8, 16, 4), (1, 1, 1, 128, 8, 16, 1, InfiniDtype.I32),
(2, 8, 8, 128, 16, 32, 2), (1, 1, 1, 128, 8, 16, 1, InfiniDtype.I64),
(4, 16, 16, 128, 8, 64, 3), (1, 4, 4, 128, 8, 16, 4, InfiniDtype.I32),
(8, 64, 64, 128, 8, 16, 5), (1, 4, 4, 128, 8, 16, 4, InfiniDtype.I64),
(16, 128, 128, 128, 8, 16, 4), (2, 8, 8, 128, 16, 32, 2, InfiniDtype.I32),
(2, 8, 8, 128, 16, 32, 2, InfiniDtype.I64),
(4, 16, 16, 128, 8, 64, 3, InfiniDtype.I32),
(4, 16, 16, 128, 8, 64, 3, InfiniDtype.I64),
(8, 64, 64, 128, 8, 16, 5, InfiniDtype.I32),
(8, 64, 64, 128, 8, 16, 5, InfiniDtype.I64),
(16, 128, 128, 128, 8, 16, 4, InfiniDtype.I32),
(16, 128, 128, 128, 8, 16, 4, InfiniDtype.I64),
] ]
_TENSOR_DTYPES = [InfiniDtype.BF16, InfiniDtype.F16] _TENSOR_DTYPES = [InfiniDtype.BF16, InfiniDtype.F16]
...@@ -124,13 +131,15 @@ def test( ...@@ -124,13 +131,15 @@ def test(
block_size, block_size,
max_step_len, max_step_len,
num_rounds, num_rounds,
index_dtype=InfiniDtype.I64,
dtype=InfiniDtype.F16, dtype=InfiniDtype.F16,
sync=None, sync=None,
): ):
print( print(
f"Testing PagedAttentionPrefill on {InfiniDeviceNames[device]} with " f"Testing PagedAttentionPrefill on {InfiniDeviceNames[device]} with "
f"seqs:{num_seqs}, heads:{num_heads}, head_size:{head_size}, " f"seqs:{num_seqs}, heads:{num_heads}, head_size:{head_size}, "
f"block:{block_size}, max_step_len:{max_step_len}, num_rounds:{num_rounds}, dtype:{InfiniDtypeNames[dtype]}" f"block:{block_size}, max_step_len:{max_step_len}, num_rounds:{num_rounds}, dtype:{InfiniDtypeNames[dtype]}, "
f"index_dtype:{InfiniDtypeNames[index_dtype]}"
) )
# 1. Initialize persistent resources # 1. Initialize persistent resources
...@@ -194,23 +203,26 @@ def test( ...@@ -194,23 +203,26 @@ def test(
out = TestTensor.from_torch(q_packed_tensors, dtype, device) out = TestTensor.from_torch(q_packed_tensors, dtype, device)
out.actual_tensor().zero_() out.actual_tensor().zero_()
# 3. Referencing index_dtype to set torch dtype
torch_idx_type = torch.int32 if index_dtype == InfiniDtype.I32 else torch.int64
seq_lens = TestTensor.from_torch( seq_lens = TestTensor.from_torch(
torch.tensor(seq_lens_list, dtype=torch.int64), InfiniDtype.I64, device torch.tensor(seq_lens_list, dtype=torch_idx_type), index_dtype, device
) )
cum_seq_lens_q = TestTensor.from_torch( cum_seq_lens_q = TestTensor.from_torch(
torch.tensor(cum_seq_lens_q_list, dtype=torch.int64), torch.tensor(cum_seq_lens_q_list, dtype=torch_idx_type),
InfiniDtype.I64, index_dtype,
device, device,
) )
max_blocks = max(len(t) for t in all_block_tables) max_blocks = max(len(t) for t in all_block_tables)
padded_tables = [t + [0] * (max_blocks - len(t)) for t in all_block_tables] padded_tables = [t + [0] * (max_blocks - len(t)) for t in all_block_tables]
block_tables = TestTensor.from_torch( block_tables = TestTensor.from_torch(
torch.tensor(padded_tables, dtype=torch.int64), InfiniDtype.I64, device torch.tensor(padded_tables, dtype=torch_idx_type), index_dtype, device
) )
# 3. Reference Calculation # 4. Reference Calculation
def torch_paged_attention_multi_turn(): def torch_paged_attention_multi_turn():
return ref_paged_attention_multi_turn( return ref_paged_attention_multi_turn(
q_new.torch_tensor(), q_new.torch_tensor(),
...@@ -224,7 +236,7 @@ def test( ...@@ -224,7 +236,7 @@ def test(
ans = torch_paged_attention_multi_turn() ans = torch_paged_attention_multi_turn()
# 4. Infiniop Operator Execution # 5. Infiniop Operator Execution
descriptor = infiniopOperatorDescriptor_t() descriptor = infiniopOperatorDescriptor_t()
check_error( check_error(
LIBINFINIOP.infiniopCreatePagedAttentionPrefillDescriptor( LIBINFINIOP.infiniopCreatePagedAttentionPrefillDescriptor(
...@@ -272,7 +284,7 @@ def test( ...@@ -272,7 +284,7 @@ def test(
if sync: if sync:
sync() sync()
# 5. Validation # 6. Validation
atol, rtol = get_tolerance(_TOLERANCE_MAP, dtype) atol, rtol = get_tolerance(_TOLERANCE_MAP, dtype)
if DEBUG: if DEBUG:
debug(out.actual_tensor(), ans, atol=atol, rtol=rtol) debug(out.actual_tensor(), ans, atol=atol, rtol=rtol)
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment