From d3cc935658d7b58845fdce2fcd08871e335ecc85 Mon Sep 17 00:00:00 2001 From: Qzong2020 Date: Mon, 14 Sep 2026 10:31:49 +0800 Subject: [PATCH 1/2] Optimize MetaX HD64 paged attention with grouped KV loading --- .../metax/paged_attention_hd64.maca | 312 ++++++++++++++++-- .../metax/paged_attention_metax.maca | 27 +- test/infiniop/paged_attention.py | 85 ++++- 3 files changed, 387 insertions(+), 37 deletions(-) diff --git a/src/infiniop/ops/paged_attention/metax/paged_attention_hd64.maca b/src/infiniop/ops/paged_attention/metax/paged_attention_hd64.maca index 2f8b95b3a..6117be90f 100644 --- a/src/infiniop/ops/paged_attention/metax/paged_attention_hd64.maca +++ b/src/infiniop/ops/paged_attention/metax/paged_attention_hd64.maca @@ -79,6 +79,243 @@ inline int chooseNumSplitsHeuristic(size_t num_heads, size_t num_seqs, size_t se } } // namespace +template +INFINIOP_METAX_KERNEL flashAttentionDecodeHd64GqaSplitKv( + float *partial_acc, // [num_splits, num_seqs, num_heads, head_size] + float *partial_m, // [num_splits, num_seqs, num_heads] + float *partial_l, // [num_splits, num_seqs, num_heads] + const Tdata *q_, + const Tdata *k_cache_, + const Tdata *v_cache_, + const Tindex *block_tables_, + const Tindex *cache_lens_, + const float *alibi_slopes_, + size_t num_kv_heads, + float scale, + size_t max_num_blocks_per_seq, + size_t page_block_size, + ptrdiff_t q_stride, + ptrdiff_t k_batch_stride, + ptrdiff_t k_row_stride, + ptrdiff_t k_head_stride, + ptrdiff_t v_batch_stride, + ptrdiff_t v_row_stride, + ptrdiff_t v_head_stride, + int num_splits, int num_heads, int num_seqs) { + constexpr int HEAD_SIZE = 64; + constexpr int G = 4; + constexpr int T = 8; + // Keep the exact legacy token-wise arithmetic; only KV loading is grouped. + __shared__ __align__(16) Tdata sh_k[T][HEAD_SIZE]; + __shared__ __align__(16) Tdata sh_v[T][HEAD_SIZE]; + __shared__ int physical_page; + + const int seq_idx = blockIdx.y; + const int head_idx = blockIdx.x * G + threadIdx.x / 32; + const int split_idx = static_cast(blockIdx.z); + const int lane = threadIdx.x % 32; + constexpr int kWarpSize = 32; + static_assert(HEAD_SIZE % kWarpSize == 0, "HEAD_SIZE must be divisible by 32."); + constexpr int DIMS_PER_THREAD = HEAD_SIZE / kWarpSize; + + const int seq_len = static_cast(cache_lens_[seq_idx]); + if (seq_len <= 0 || num_splits <= 0) { + return; + } + + // Split the [0, seq_len) range into num_splits contiguous shards. + const int shard = (seq_len + num_splits - 1) / num_splits; + const int start = split_idx * shard; + const int end = min(seq_len, start + shard); + if (start >= end) { + // Empty shard => write neutral element. + const int n = num_seqs * num_heads; + const int idx = (split_idx * n + seq_idx * num_heads + head_idx); + if (lane == 0) { + partial_m[idx] = -INFINITY; + partial_l[idx] = 0.0f; + } +#pragma unroll + for (int i = 0; i < DIMS_PER_THREAD; ++i) { + const int dim = lane * DIMS_PER_THREAD + i; + partial_acc[idx * HEAD_SIZE + dim] = 0.0f; + } + return; + } + + const int num_queries_per_kv = num_heads / static_cast(num_kv_heads); + const int kv_head_idx = head_idx / num_queries_per_kv; + + const float alibi_slope = (alibi_slopes_ == nullptr) ? 0.0f : alibi_slopes_[head_idx]; + constexpr float kLog2e = 1.4426950408889634f; + const float scale_log2 = scale * kLog2e; + + const Tindex *block_table = block_tables_ + seq_idx * static_cast(max_num_blocks_per_seq); + const Tdata *q_ptr = q_ + seq_idx * q_stride + head_idx * HEAD_SIZE; + + float q_reg[DIMS_PER_THREAD]; + float acc[DIMS_PER_THREAD]; +#pragma unroll + for (int i = 0; i < DIMS_PER_THREAD; ++i) { + const int dim = lane * DIMS_PER_THREAD + i; + q_reg[i] = static_cast(q_ptr[dim]); + acc[i] = 0.0f; + } + +#if defined(__CUDA_ARCH__) + float2 q_reg2[DIMS_PER_THREAD / 2]; + if constexpr (std::is_same_v) { + const int dim_base = lane * DIMS_PER_THREAD; + const half2 *q2 = reinterpret_cast(q_ptr + dim_base); +#pragma unroll + for (int j = 0; j < DIMS_PER_THREAD / 2; ++j) { + q_reg2[j] = __half22float2(q2[j]); + } + } + if constexpr (std::is_same_v) { + const int dim_base = lane * DIMS_PER_THREAD; + const __nv_bfloat162 *q2 = reinterpret_cast(q_ptr + dim_base); +#pragma unroll + for (int j = 0; j < DIMS_PER_THREAD / 2; ++j) { + q_reg2[j] = __bfloat1622float2(q2[j]); + } + } +#endif + + float m = -INFINITY; + float l = 0.0f; + const int pbs = static_cast(page_block_size); + + // Scan only [start, end). + int t = start; + int logical_block = t / pbs; + int token_in_block = t - logical_block * pbs; + for (; t < end; ++logical_block) { + if (threadIdx.x == 0) { + physical_page = static_cast(block_table[logical_block]); + } + __syncthreads(); + const Tdata *k_base = k_cache_ + physical_page * k_batch_stride + kv_head_idx * k_head_stride; + const Tdata *v_base = v_cache_ + physical_page * v_batch_stride + kv_head_idx * v_head_stride; + const int token_end = min(pbs, end - logical_block * pbs); + for (; token_in_block < token_end; ) { + const int tile_n = min(T, token_end - token_in_block); + // 128 threads each copy one 16-byte chunk: 64 K chunks, 64 V chunks. + const int chunk_id = threadIdx.x % 64; + const int tok = chunk_id / 8; + const int dim_chunk = (chunk_id % 8) * 8; + Tdata *dst = (threadIdx.x < 64 ? &sh_k[tok][dim_chunk] : &sh_v[tok][dim_chunk]); + if (tok < tile_n) { + const Tdata *src = threadIdx.x < 64 + ? k_base + (token_in_block + tok) * k_row_stride + dim_chunk + : v_base + (token_in_block + tok) * v_row_stride + dim_chunk; + *reinterpret_cast(dst) = *reinterpret_cast(src); + } else { + *reinterpret_cast(dst) = make_uint4(0, 0, 0, 0); + } + __syncthreads(); + for (int j = 0; j < tile_n; ++j, ++t) { + const Tdata *k_ptr = sh_k[j]; + const Tdata *v_ptr = sh_v[j]; + float qk = 0.0f; +#if defined(__CUDA_ARCH__) + if constexpr (std::is_same_v) { + const int dim_base = lane * DIMS_PER_THREAD; + const half2 *k2 = reinterpret_cast(k_ptr + dim_base); +#pragma unroll + for (int j = 0; j < DIMS_PER_THREAD / 2; ++j) { + const float2 qf = q_reg2[j]; + const float2 kf = __half22float2(k2[j]); + qk += qf.x * kf.x + qf.y * kf.y; + } + } else if constexpr (std::is_same_v) { + const int dim_base = lane * DIMS_PER_THREAD; + const __nv_bfloat162 *k2 = reinterpret_cast(k_ptr + dim_base); +#pragma unroll + for (int j = 0; j < DIMS_PER_THREAD / 2; ++j) { + const float2 qf = q_reg2[j]; + const float2 kf = __bfloat1622float2(k2[j]); + qk += qf.x * kf.x + qf.y * kf.y; + } + } else +#endif + { +#pragma unroll + for (int i = 0; i < DIMS_PER_THREAD; ++i) { + const int dim = lane * DIMS_PER_THREAD + i; + qk += q_reg[i] * static_cast(k_ptr[dim]); + } + } + + qk = op::paged_attention::cuda::warpReduceSum(qk); + + float alpha = 1.0f; + float beta = 0.0f; + if (lane == 0) { + float score = qk * scale_log2; + if (alibi_slope != 0.0f) { + score += (alibi_slope * static_cast(t - (seq_len - 1))) * kLog2e; + } + const float m_new = fmaxf(m, score); + alpha = exp2f(m - m_new); + beta = exp2f(score - m_new); + l = l * alpha + beta; + m = m_new; + } + + alpha = __shfl_sync(0xffffffff, alpha, 0); + beta = __shfl_sync(0xffffffff, beta, 0); + +#if defined(__CUDA_ARCH__) + if constexpr (std::is_same_v) { + const int dim_base = lane * DIMS_PER_THREAD; + const half2 *v2 = reinterpret_cast(v_ptr + dim_base); +#pragma unroll + for (int j = 0; j < DIMS_PER_THREAD / 2; ++j) { + const float2 vf = __half22float2(v2[j]); + acc[j * 2 + 0] = acc[j * 2 + 0] * alpha + beta * vf.x; + acc[j * 2 + 1] = acc[j * 2 + 1] * alpha + beta * vf.y; + } + } else if constexpr (std::is_same_v) { + const int dim_base = lane * DIMS_PER_THREAD; + const __nv_bfloat162 *v2 = reinterpret_cast(v_ptr + dim_base); +#pragma unroll + for (int j = 0; j < DIMS_PER_THREAD / 2; ++j) { + const float2 vf = __bfloat1622float2(v2[j]); + acc[j * 2 + 0] = acc[j * 2 + 0] * alpha + beta * vf.x; + acc[j * 2 + 1] = acc[j * 2 + 1] * alpha + beta * vf.y; + } + } else +#endif + { +#pragma unroll + for (int i = 0; i < DIMS_PER_THREAD; ++i) { + const int dim = lane * DIMS_PER_THREAD + i; + const float v_val = static_cast(v_ptr[dim]); + acc[i] = acc[i] * alpha + beta * v_val; + } + } + } + // All query-head warps must finish before producers overwrite the tile. + __syncthreads(); + token_in_block += tile_n; + } + token_in_block = 0; + } + + const int n = num_seqs * num_heads; + const int idx = (split_idx * n + seq_idx * num_heads + head_idx); + if (lane == 0) { + partial_m[idx] = m; + partial_l[idx] = l; + } +#pragma unroll + for (int i = 0; i < DIMS_PER_THREAD; ++i) { + const int dim = lane * DIMS_PER_THREAD + i; + partial_acc[idx * HEAD_SIZE + dim] = acc[i]; + } +} + template INFINIOP_METAX_KERNEL flashAttentionDecodeHd64Warp( Tdata *out, @@ -233,6 +470,7 @@ infiniStatus_t launch_decode_hd64_impl( ptrdiff_t v_row_stride, ptrdiff_t v_head_stride, ptrdiff_t o_stride, + bool gqa_layout_eligible, hcStream_t stream) { dim3 grid(static_cast(num_heads), static_cast(num_seqs), 1); @@ -315,33 +553,62 @@ infiniStatus_t launch_decode_hd64_impl( float *partial_m = partial_acc + acc_elems; float *partial_l = partial_m + m_elems; + // Preserve the existing split policy; only eligible split launches use KV sharing. + const bool aligned_kv = (reinterpret_cast(k_cache) % 16 == 0) + && (reinterpret_cast(v_cache) % 16 == 0) + && k_batch_stride % 8 == 0 && k_head_stride % 8 == 0 && k_row_stride % 8 == 0 + && v_batch_stride % 8 == 0 && v_head_stride % 8 == 0 && v_row_stride % 8 == 0; + const bool grouped = gqa_layout_eligible && (num_splits == 1 || num_splits == 2 || num_splits == 4 || num_splits == 8) + && num_kv_heads > 0 && num_heads == num_kv_heads * 8 + && alibi_slopes == nullptr && aligned_kv + && (page_block_size == 16 || page_block_size == 32 || page_block_size == 64 || page_block_size == 256); dim3 grid_split(static_cast(num_heads), static_cast(num_seqs), static_cast(num_splits)); dim3 block_split(32); if (dtype == INFINI_DTYPE_F16) { - flashAttentionDecodeHd64SplitKv<<>>( - partial_acc, partial_m, partial_l, - static_cast(q), - static_cast(k_cache), - static_cast(v_cache), - block_tables, cache_lens, alibi_slopes, - num_kv_heads, scale, max_num_blocks_per_seq, page_block_size, - q_stride, k_batch_stride, k_row_stride, k_head_stride, - v_batch_stride, v_row_stride, v_head_stride, num_splits); + if (grouped) { + flashAttentionDecodeHd64GqaSplitKv<<>>( + partial_acc, partial_m, partial_l, + static_cast(q), static_cast(k_cache), static_cast(v_cache), + block_tables, cache_lens, alibi_slopes, num_kv_heads, scale, + max_num_blocks_per_seq, page_block_size, q_stride, + k_batch_stride, k_row_stride, k_head_stride, v_batch_stride, v_row_stride, v_head_stride, + num_splits, num_heads, num_seqs); + } else { + flashAttentionDecodeHd64SplitKv<<>>( + partial_acc, partial_m, partial_l, + static_cast(q), + static_cast(k_cache), + static_cast(v_cache), + block_tables, cache_lens, alibi_slopes, + num_kv_heads, scale, max_num_blocks_per_seq, page_block_size, + q_stride, k_batch_stride, k_row_stride, k_head_stride, + v_batch_stride, v_row_stride, v_head_stride, num_splits); + } flashAttentionDecodeHd64SplitKvCombine<<>>( static_cast(out), partial_acc, partial_m, partial_l, num_splits, o_stride); return INFINI_STATUS_SUCCESS; } if (dtype == INFINI_DTYPE_BF16) { - flashAttentionDecodeHd64SplitKv<<>>( - partial_acc, partial_m, partial_l, - static_cast(q), - static_cast(k_cache), - static_cast(v_cache), - block_tables, cache_lens, alibi_slopes, - num_kv_heads, scale, max_num_blocks_per_seq, page_block_size, - q_stride, k_batch_stride, k_row_stride, k_head_stride, - v_batch_stride, v_row_stride, v_head_stride, num_splits); + if (grouped) { + flashAttentionDecodeHd64GqaSplitKv<<>>( + partial_acc, partial_m, partial_l, + static_cast(q), static_cast(k_cache), static_cast(v_cache), + block_tables, cache_lens, alibi_slopes, num_kv_heads, scale, + max_num_blocks_per_seq, page_block_size, q_stride, + k_batch_stride, k_row_stride, k_head_stride, v_batch_stride, v_row_stride, v_head_stride, + num_splits, num_heads, num_seqs); + } else { + flashAttentionDecodeHd64SplitKv<<>>( + partial_acc, partial_m, partial_l, + static_cast(q), + static_cast(k_cache), + static_cast(v_cache), + block_tables, cache_lens, alibi_slopes, + num_kv_heads, scale, max_num_blocks_per_seq, page_block_size, + q_stride, k_batch_stride, k_row_stride, k_head_stride, + v_batch_stride, v_row_stride, v_head_stride, num_splits); + } flashAttentionDecodeHd64SplitKvCombine<__nv_bfloat16><<>>( static_cast<__nv_bfloat16 *>(out), partial_acc, partial_m, partial_l, num_splits, o_stride); return INFINI_STATUS_SUCCESS; @@ -451,12 +718,13 @@ infiniStatus_t launch_decode_hd64_i64( ptrdiff_t v_row_stride, ptrdiff_t v_head_stride, ptrdiff_t o_stride, + bool gqa_layout_eligible, hcStream_t stream) { return launch_decode_hd64_impl( workspace, workspace_size, out, q, k_cache, v_cache, dtype, block_tables, cache_lens, alibi_slopes, num_heads, num_seqs, num_kv_heads, scale, max_num_blocks_per_seq, page_block_size, q_stride, k_batch_stride, k_row_stride, - k_head_stride, v_batch_stride, v_row_stride, v_head_stride, o_stride, stream); + k_head_stride, v_batch_stride, v_row_stride, v_head_stride, o_stride, gqa_layout_eligible, stream); } infiniStatus_t launch_decode_hd64_i32( @@ -484,12 +752,13 @@ infiniStatus_t launch_decode_hd64_i32( ptrdiff_t v_row_stride, ptrdiff_t v_head_stride, ptrdiff_t o_stride, + bool gqa_layout_eligible, hcStream_t stream) { return launch_decode_hd64_impl( workspace, workspace_size, out, q, k_cache, v_cache, dtype, block_tables, cache_lens, alibi_slopes, num_heads, num_seqs, num_kv_heads, scale, max_num_blocks_per_seq, page_block_size, q_stride, k_batch_stride, k_row_stride, - k_head_stride, v_batch_stride, v_row_stride, v_head_stride, o_stride, stream); + k_head_stride, v_batch_stride, v_row_stride, v_head_stride, o_stride, gqa_layout_eligible, stream); } infiniStatus_t launch_decode_hd64_u32( @@ -517,12 +786,13 @@ infiniStatus_t launch_decode_hd64_u32( ptrdiff_t v_row_stride, ptrdiff_t v_head_stride, ptrdiff_t o_stride, + bool gqa_layout_eligible, hcStream_t stream) { return launch_decode_hd64_impl( workspace, workspace_size, out, q, k_cache, v_cache, dtype, block_tables, cache_lens, alibi_slopes, num_heads, num_seqs, num_kv_heads, scale, max_num_blocks_per_seq, page_block_size, q_stride, k_batch_stride, k_row_stride, - k_head_stride, v_batch_stride, v_row_stride, v_head_stride, o_stride, stream); + k_head_stride, v_batch_stride, v_row_stride, v_head_stride, o_stride, gqa_layout_eligible, stream); } } // namespace op::paged_attention::metax diff --git a/src/infiniop/ops/paged_attention/metax/paged_attention_metax.maca b/src/infiniop/ops/paged_attention/metax/paged_attention_metax.maca index 2b9d64b17..59ecbf243 100644 --- a/src/infiniop/ops/paged_attention/metax/paged_attention_metax.maca +++ b/src/infiniop/ops/paged_attention/metax/paged_attention_metax.maca @@ -20,7 +20,7 @@ infiniStatus_t launch_decode_hd64_i64( size_t num_heads, size_t num_seqs, size_t num_kv_heads, float scale, size_t max_num_blocks_per_seq, size_t page_block_size, ptrdiff_t q_stride, ptrdiff_t k_batch_stride, ptrdiff_t k_row_stride, ptrdiff_t k_head_stride, ptrdiff_t v_batch_stride, ptrdiff_t v_row_stride, ptrdiff_t v_head_stride, ptrdiff_t o_stride, - hcStream_t stream); + bool gqa_layout_eligible, hcStream_t stream); infiniStatus_t launch_decode_hd64_i32( void *workspace, size_t workspace_size, @@ -29,7 +29,7 @@ infiniStatus_t launch_decode_hd64_i32( size_t num_heads, size_t num_seqs, size_t num_kv_heads, float scale, size_t max_num_blocks_per_seq, size_t page_block_size, ptrdiff_t q_stride, ptrdiff_t k_batch_stride, ptrdiff_t k_row_stride, ptrdiff_t k_head_stride, ptrdiff_t v_batch_stride, ptrdiff_t v_row_stride, ptrdiff_t v_head_stride, ptrdiff_t o_stride, - hcStream_t stream); + bool gqa_layout_eligible, hcStream_t stream); infiniStatus_t launch_decode_hd64_u32( void *workspace, size_t workspace_size, @@ -38,7 +38,7 @@ infiniStatus_t launch_decode_hd64_u32( size_t num_heads, size_t num_seqs, size_t num_kv_heads, float scale, size_t max_num_blocks_per_seq, size_t page_block_size, ptrdiff_t q_stride, ptrdiff_t k_batch_stride, ptrdiff_t k_row_stride, ptrdiff_t k_head_stride, ptrdiff_t v_batch_stride, ptrdiff_t v_row_stride, ptrdiff_t v_head_stride, ptrdiff_t o_stride, - hcStream_t stream); + bool gqa_layout_eligible, hcStream_t stream); infiniStatus_t launch_decode_hd128_i64( void *workspace, size_t workspace_size, @@ -176,6 +176,7 @@ infiniStatus_t launch_decode_mla_hd576_v512_u32( struct Descriptor::Opaque { std::shared_ptr internal; + bool hd64_gqa_layout_eligible; }; Descriptor::~Descriptor() { @@ -203,8 +204,20 @@ infiniStatus_t Descriptor::create( const size_t per_split = info.num_seqs * info.num_heads * (info.value_size + 2) * sizeof(float); const size_t workspace_bytes = kMaxSplits * per_split; + // Private fast-path guard: the legacy D64 signature assumes packed heads/table rows. + const bool gqa_layout_eligible = info.head_size == 64 && info.value_size == 64 + && q_desc->stride(1) == 64 && out_desc->stride(1) == 64 + && q_desc->stride(0) >= static_cast(info.num_heads * 64) + && out_desc->stride(0) >= static_cast(info.num_heads * 64) + && block_tables_desc->stride(0) == static_cast(info.max_num_blocks_per_seq) + && k_cache_desc->stride(2) == 64 && v_cache_desc->stride(2) == 64 + && k_cache_desc->stride(1) == static_cast(info.page_block_size * 64) + && v_cache_desc->stride(1) == static_cast(info.page_block_size * 64) + && k_cache_desc->stride(0) == static_cast(info.num_kv_heads * info.page_block_size * 64) + && v_cache_desc->stride(0) == static_cast(info.num_kv_heads * info.page_block_size * 64); + *desc_ptr = new Descriptor( - new Opaque{reinterpret_cast(handle)->internal()}, + new Opaque{reinterpret_cast(handle)->internal(), gqa_layout_eligible}, info, workspace_bytes, handle->device, handle->device_id); return INFINI_STATUS_SUCCESS; @@ -256,7 +269,7 @@ infiniStatus_t Descriptor::calculate( _info.max_num_blocks_per_seq, _info.page_block_size, _info.q_stride, _info.k_batch_stride, _info.k_row_stride, _info.k_head_stride, _info.v_batch_stride, _info.v_row_stride, _info.v_head_stride, - _info.o_stride, stream); + _info.o_stride, _opaque->hd64_gqa_layout_eligible, stream); case 128: return launch_decode_hd128_i64( workspace, workspace_size, @@ -326,7 +339,7 @@ infiniStatus_t Descriptor::calculate( _info.max_num_blocks_per_seq, _info.page_block_size, _info.q_stride, _info.k_batch_stride, _info.k_row_stride, _info.k_head_stride, _info.v_batch_stride, _info.v_row_stride, _info.v_head_stride, - _info.o_stride, stream); + _info.o_stride, _opaque->hd64_gqa_layout_eligible, stream); case 128: return launch_decode_hd128_i32( workspace, workspace_size, @@ -396,7 +409,7 @@ infiniStatus_t Descriptor::calculate( _info.max_num_blocks_per_seq, _info.page_block_size, _info.q_stride, _info.k_batch_stride, _info.k_row_stride, _info.k_head_stride, _info.v_batch_stride, _info.v_row_stride, _info.v_head_stride, - _info.o_stride, stream); + _info.o_stride, _opaque->hd64_gqa_layout_eligible, stream); case 128: return launch_decode_hd128_u32( workspace, workspace_size, diff --git a/test/infiniop/paged_attention.py b/test/infiniop/paged_attention.py index d7a6b9679..312b69f93 100644 --- a/test/infiniop/paged_attention.py +++ b/test/infiniop/paged_attention.py @@ -2,9 +2,12 @@ import ctypes from ctypes import c_uint64 import math +import os +from unittest.mock import patch from libinfiniop import ( LIBINFINIOP, TestTensor, + CTensor, get_test_devices, check_error, test_operator, @@ -15,6 +18,7 @@ InfiniDtype, InfiniDtypeNames, InfiniDeviceNames, + InfiniDeviceEnum, infiniopOperatorDescriptor_t, TestWorkspace, ) @@ -135,6 +139,10 @@ def test( max_seq_len, use_alibi, *tail, + index_dtype=InfiniDtype.I64, + lengths=None, + kv_padding=0, + cpu_reference=False, ): if len(tail) == 2: dtype, sync = tail @@ -156,21 +164,40 @@ def test( # Create input tensors q = TestTensor((num_seqs, num_heads, head_size), None, dtype, device) out = TestTensor((num_seqs, num_heads, value_size), None, dtype, device) + kv_strides = None + if kv_padding: + row = block_size + kv_padding + kv_strides = (num_kv_heads * row * head_size, row * head_size, head_size, 1) k_cache = TestTensor( - (num_blocks, num_kv_heads, block_size, head_size), None, dtype, device + (num_blocks, num_kv_heads, block_size, head_size), kv_strides, dtype, device ) v_cache = TestTensor( - (num_blocks, num_kv_heads, block_size, value_size), None, dtype, device + (num_blocks, num_kv_heads, block_size, value_size), kv_strides, dtype, device ) seq_lens_torch = torch.randint(1, max_seq_len, (num_seqs,), dtype=torch.int64) - seq_lens = TestTensor.from_torch(seq_lens_torch, InfiniDtype.I64, device) + if lengths is not None: + seq_lens_torch = torch.tensor(lengths, dtype=torch.int64) + def index_tensor(values): + # Positive I32/U32 indices have identical bits. Keep signed Torch storage + # because some backends cannot clone U32, but exercise a real U32 descriptor. + storage_dtype = InfiniDtype.I32 if index_dtype == InfiniDtype.U32 else index_dtype + tensor = TestTensor.from_torch(values, storage_dtype, device) + if index_dtype == InfiniDtype.U32: + assert bool((values >= 0).all()) + tensor.destroy_desc() + CTensor.__init__(tensor, index_dtype, values.shape, None) + return tensor + + seq_lens = index_tensor(seq_lens_torch) block_tables_py = torch.arange( 0, num_seqs * max_blocks_per_seq, dtype=torch.int64 ).view(num_seqs, max_blocks_per_seq) - block_tables = TestTensor.from_torch(block_tables_py, InfiniDtype.I64, device) + if cpu_reference: + block_tables_py = torch.randperm(num_blocks).view(num_seqs, max_blocks_per_seq) + block_tables = index_tensor(block_tables_py) alibi_slopes_desc = ctypes.c_void_p(0) alibi_slopes_data = ctypes.c_void_p(0) @@ -181,15 +208,20 @@ def test( alibi_slopes_data = alibi_slopes.data() alibi_slopes_torch = alibi_slopes.torch_tensor() + # Targeted split cases use independent CPU FP32 QK, softmax and PV. + def reference_input(tensor): + tensor = tensor.torch_tensor() + return tensor.cpu().float() if cpu_reference else tensor + # Run reference implementation ans = ref_single_query_cached_kv_attention( - q.torch_tensor(), - k_cache.torch_tensor(), - v_cache.torch_tensor(), + reference_input(q), + reference_input(k_cache), + reference_input(v_cache), block_tables.torch_tensor(), seq_lens.torch_tensor(), scale, - alibi_slopes_torch, + alibi_slopes_torch.cpu().float() if cpu_reference and use_alibi else alibi_slopes_torch, ) if sync: @@ -260,7 +292,10 @@ def lib_paged_attention(): atol, rtol = get_tolerance(_TOLERANCE_MAP, dtype) if DEBUG: debug(out.actual_tensor(), ans, atol=atol, rtol=rtol) - assert torch.allclose(out.actual_tensor(), ans, atol=atol, rtol=rtol) + actual = out.actual_tensor() + if cpu_reference: + actual = actual.cpu().float() + assert torch.allclose(actual, ans, atol=atol, rtol=rtol) # Profiling workflow if PROFILE: @@ -277,6 +312,31 @@ def lib_paged_attention(): check_error(LIBINFINIOP.infiniopDestroyPagedAttentionDescriptor(descriptor)) +def test_metax_grouped_split(handle, device, splits, index_dtype, dtype, sync): + # Explicit split settings are scoped to this test; production defaults stay unchanged. + env = { + "INFINIOP_FLASH_DECODE_SPLITKV": "1", + "INFINIOP_FLASH_NUM_SPLITS": str(splits), + } + with patch.dict(os.environ, env), torch.random.fork_rng(devices=[]): + torch.manual_seed(20260910) + for page in (16, 32, 64, 256): + # Empty shards, nondivisible shards, and starts within/across a page. + for lengths in ((1, 3, 7, 9), (page - 1, page, page + 1, 2 * page + 3)): + test( + handle, device, 4, 32, 4, 64, page, max(lengths), False, + dtype, sync, index_dtype=index_dtype, lengths=lengths, + cpu_reference=True, + ) + # Non-G8, ALiBi and padded KV heads exercise the legacy split fallback. + for heads, alibi, padding in ((16, False, 0), (32, True, 0), (32, False, 1)): + test( + handle, device, 2, heads, 4, 64, 16, 35, alibi, + dtype, sync, index_dtype=index_dtype, lengths=(17, 35), + kv_padding=padding, cpu_reference=True, + ) + + if __name__ == "__main__": args = get_args() @@ -288,5 +348,12 @@ def lib_paged_attention(): for device in get_test_devices(args): test_operator(device, test, _TEST_CASES_, _TENSOR_DTYPES) + if device == InfiniDeviceEnum.METAX: + split_cases = [ + (s, i) + for s in (1, 2, 4, 8) + for i in (InfiniDtype.I32, InfiniDtype.I64, InfiniDtype.U32) + ] + test_operator(device, test_metax_grouped_split, split_cases, _TENSOR_DTYPES) print("\033[92mTest passed!\033[0m") From a60fa3fff0fe4583d9cab85fa37f7c475152c043 Mon Sep 17 00:00:00 2001 From: Qzong2020 Date: Tue, 15 Sep 2026 10:51:12 +0800 Subject: [PATCH 2/2] chore: trigger CI