wehub-resource-sync
1024 lines
44 KiB
C++
1024 lines
44 KiB
C++
// SPDX-License-Identifier: Apache-2.0
|
||
//
|
||
// SYCL implementation of LMCache memory kernels for Intel XPU
|
||
// (PVC / Arc / Battlemage).
|
||
//
|
||
// Performance-critical design choices:
|
||
//
|
||
// 1. Work-group size 256 -- keeps Intel EU ALUs fed and hides
|
||
// global-memory latency.
|
||
// 2. [[sycl::reqd_sub_group_size(16)]] -- native SIMD width on all
|
||
// Intel discrete GPUs; prevents IGC from falling back to width 32
|
||
// on PVC (which would halve occupancy).
|
||
// 3. Compile-time template parameters (DIRECTION, USE_MLA) eliminate
|
||
// run-time branches in the innermost loop.
|
||
// 4. Hoisted loop-invariant base offsets -- integer division/modulo
|
||
// (flash_infer block indexing) is computed once per token+layer
|
||
// rather than per inner-loop iteration.
|
||
// 5. 64-bit (int64_t) bulk transfers pack two fp32 / four fp16 /
|
||
// eight int8 values per move.
|
||
// 6. Fused K+V work-groups (non-MLA multi-layer kernels) -- one
|
||
// work-group handles both K and V, halving dispatch count and
|
||
// avoiding redundant slot/pointer/division work.
|
||
|
||
// sycl/accessor.hpp references the deprecated 'host_buffer' internally
|
||
// even when user code only uses USM pointers; suppress the noise.
|
||
#pragma GCC diagnostic push
|
||
#pragma GCC diagnostic ignored "-Wdeprecated-declarations"
|
||
#include <sycl/sycl.hpp>
|
||
#pragma GCC diagnostic pop
|
||
|
||
#include <torch/all.h>
|
||
#include <ATen/ATen.h>
|
||
#include <c10/core/DeviceGuard.h>
|
||
#include <c10/xpu/XPUStream.h>
|
||
|
||
#include "mem_kernels_sycl.h"
|
||
|
||
#include <algorithm>
|
||
#include <cstdint>
|
||
#include <stdexcept>
|
||
#include <string>
|
||
|
||
// ---------------------------------------------------------------------------
|
||
// Tuning constants
|
||
// ---------------------------------------------------------------------------
|
||
// Native SIMD width on PVC, DG2, and BMG.
|
||
constexpr int INTEL_SUB_GROUP_SIZE = 16;
|
||
|
||
// Max work-group size; must be a multiple of INTEL_SUB_GROUP_SIZE.
|
||
constexpr int MAX_WG_SIZE = 256;
|
||
|
||
// ---------------------------------------------------------------------------
|
||
// Round up so the work-group divides evenly into sub-groups.
|
||
// ---------------------------------------------------------------------------
|
||
inline int round_up_to_sg(int n) {
|
||
return ((n + INTEL_SUB_GROUP_SIZE - 1) / INTEL_SUB_GROUP_SIZE) *
|
||
INTEL_SUB_GROUP_SIZE;
|
||
}
|
||
|
||
// ---------------------------------------------------------------------------
|
||
// Namespace lmc -- device-side helper functions
|
||
// ---------------------------------------------------------------------------
|
||
namespace lmc {
|
||
|
||
template <EngineKVFormat format>
|
||
inline int64_t page_buffer_offset(const int k_or_v, const int token_idx,
|
||
const int scalar_offset,
|
||
const int scalars_per_token,
|
||
const int page_buffer_size,
|
||
const int block_size) {
|
||
// vLLM cross layer
|
||
if constexpr (format == EngineKVFormat::NB_NL_TWO_BS_NH_HS) {
|
||
return k_or_v * page_buffer_size * scalars_per_token +
|
||
token_idx * scalars_per_token + scalar_offset;
|
||
}
|
||
// vLLM flash attention
|
||
else if constexpr (format == EngineKVFormat::NL_X_TWO_NB_BS_NH_HS) {
|
||
return k_or_v * page_buffer_size * scalars_per_token +
|
||
token_idx * scalars_per_token + scalar_offset;
|
||
}
|
||
// vLLM flash infer
|
||
else if constexpr (format == EngineKVFormat::NL_X_NB_TWO_BS_NH_HS) {
|
||
const int block_idx = token_idx / block_size;
|
||
const int block_offset = token_idx % block_size;
|
||
return block_idx * 2 * block_size * scalars_per_token +
|
||
k_or_v * block_size * scalars_per_token +
|
||
block_offset * scalars_per_token + scalar_offset;
|
||
}
|
||
// MLA formats: vLLM (NL_X_NB_BS_HS) and SGLang (NL_X_NBBS_ONE_HS)
|
||
else if constexpr (format == EngineKVFormat::NL_X_NB_BS_HS ||
|
||
format == EngineKVFormat::NL_X_NBBS_ONE_HS) {
|
||
return token_idx * scalars_per_token + scalar_offset;
|
||
}
|
||
}
|
||
|
||
/// Loop-invariant base offset for the paged buffer.
|
||
/// page_buffer_offset(k_or_v, slot, i, ...) == base_offset(...) + i.
|
||
/// Hoisting this out of the inner loop avoids re-computing the
|
||
/// integer division / modulo (flash_infer) on every iteration.
|
||
template <EngineKVFormat format>
|
||
inline int64_t page_buffer_base_offset(const int k_or_v, const int token_idx,
|
||
const int scalars_per_token,
|
||
const int page_buffer_size,
|
||
const int block_size) {
|
||
if constexpr (format == EngineKVFormat::NB_NL_TWO_BS_NH_HS) {
|
||
return k_or_v * page_buffer_size * scalars_per_token +
|
||
token_idx * scalars_per_token;
|
||
} else if constexpr (format == EngineKVFormat::NL_X_TWO_NB_BS_NH_HS) {
|
||
return k_or_v * page_buffer_size * scalars_per_token +
|
||
token_idx * scalars_per_token;
|
||
} else if constexpr (format == EngineKVFormat::NL_X_NB_TWO_BS_NH_HS) {
|
||
const int block_idx = token_idx / block_size;
|
||
const int block_offset = token_idx % block_size;
|
||
return block_idx * 2 * block_size * scalars_per_token +
|
||
k_or_v * block_size * scalars_per_token +
|
||
block_offset * scalars_per_token;
|
||
} else if constexpr (format == EngineKVFormat::NL_X_NB_BS_HS ||
|
||
format == EngineKVFormat::NL_X_NBBS_ONE_HS) {
|
||
return token_idx * scalars_per_token;
|
||
}
|
||
}
|
||
|
||
inline int64_t page_buffer_offset_unilateral(const int token_idx,
|
||
const int scalar_offset,
|
||
const int scalars_per_token) {
|
||
return token_idx * scalars_per_token + scalar_offset;
|
||
}
|
||
|
||
inline int64_t key_value_offset(const int k_or_v, const int layer_idx,
|
||
const int token_idx, const int scalar_offset,
|
||
const int scalars_per_token,
|
||
const int num_tokens, const int num_layers) {
|
||
return k_or_v * num_layers * num_tokens * scalars_per_token +
|
||
layer_idx * num_tokens * scalars_per_token +
|
||
token_idx * scalars_per_token + scalar_offset;
|
||
}
|
||
|
||
/// Loop-invariant base offset for the LMCache key_value buffer.
|
||
/// key_value_offset(k_or_v, layer, token, i, ...) == base(...) + i.
|
||
inline int64_t key_value_base_offset(const int k_or_v, const int layer_idx,
|
||
const int token_idx,
|
||
const int scalars_per_token,
|
||
const int num_tokens,
|
||
const int num_layers) {
|
||
return k_or_v * num_layers * num_tokens * scalars_per_token +
|
||
layer_idx * num_tokens * scalars_per_token +
|
||
token_idx * scalars_per_token;
|
||
}
|
||
|
||
} // namespace lmc
|
||
|
||
// ---------------------------------------------------------------------------
|
||
// Pointer helper -- returns a kernel-accessible pointer of the given type.
|
||
// XPU tensors expose their USM device pointer; CPU tensors must be backed
|
||
// by USM host memory (e.g. sycl::malloc_host) to be device-accessible.
|
||
// ---------------------------------------------------------------------------
|
||
template <typename T, typename TENSOR_TYPE>
|
||
T* get_kernel_ptr(TENSOR_TYPE& tensor) {
|
||
torch::Device device = tensor.device();
|
||
if (device.is_xpu()) {
|
||
return static_cast<T*>(tensor.data_ptr());
|
||
} else if (device.is_cpu()) {
|
||
// USM host pointers are device-accessible.
|
||
return static_cast<T*>(tensor.data_ptr());
|
||
} else {
|
||
TORCH_CHECK(false,
|
||
"Invalid device. Device must be xpu or cpu (USM pinned).");
|
||
}
|
||
}
|
||
|
||
// ---------------------------------------------------------------------------
|
||
// Kernel-launch helpers -- multi-layer kernels
|
||
// ---------------------------------------------------------------------------
|
||
|
||
/**
|
||
* Submit the multi-layer KV transfer kernel for MLA formats
|
||
* (k_or_v_size == 1).
|
||
*
|
||
* nd_range layout: group(0)=k_or_v, group(1)=layer, group(2)=token;
|
||
* local_id(2)=tid, local_range(2)=num_threads.
|
||
*
|
||
* Optimizations:
|
||
* - DIRECTION is a compile-time bool (no branch in hot loop)
|
||
* - Loop-invariant base offsets (including flash_infer's integer
|
||
* division) are computed once before the inner loop
|
||
* - Work-group size rounded to a sub-group multiple for full SIMD
|
||
* utilisation
|
||
*/
|
||
template <typename scalar_t, bool DIRECTION, EngineKVFormat format>
|
||
void submit_multi_layer_kernel(sycl::queue& queue, scalar_t* key_value_ptr,
|
||
scalar_t** page_buffer_ptrs,
|
||
const int64_t* slot_mapping_ptr,
|
||
int scalars_per_token, int num_tokens,
|
||
int num_layers, int page_buffer_size,
|
||
int block_size, int skip_prefix_n_tokens,
|
||
int k_or_v_size, int wg_size) {
|
||
int num_transfer_tokens = num_tokens - skip_prefix_n_tokens;
|
||
if (num_transfer_tokens <= 0 || num_layers <= 0) return;
|
||
|
||
sycl::range<3> global_range(
|
||
static_cast<size_t>(k_or_v_size), static_cast<size_t>(num_layers),
|
||
static_cast<size_t>(num_transfer_tokens) * wg_size);
|
||
sycl::range<3> local_range(1, 1, static_cast<size_t>(wg_size));
|
||
|
||
queue.parallel_for(
|
||
sycl::nd_range<3>(global_range, local_range),
|
||
[=](sycl::nd_item<3> item) [[sycl::reqd_sub_group_size(16)]] {
|
||
const int token_id = static_cast<int>(item.get_group(2));
|
||
const int layer_id = static_cast<int>(item.get_group(1));
|
||
const int k_or_v = static_cast<int>(item.get_group(0));
|
||
const int tid = static_cast<int>(item.get_local_id(2));
|
||
const int num_threads = static_cast<int>(item.get_local_range(2));
|
||
|
||
const int kv_token_id = token_id + skip_prefix_n_tokens;
|
||
const int64_t slot_idx = slot_mapping_ptr[kv_token_id];
|
||
scalar_t* paged_buffer_ptr = page_buffer_ptrs[layer_id];
|
||
|
||
if (slot_idx < 0) return;
|
||
|
||
// Hoist loop-invariant base offsets (flash_infer's integer
|
||
// division happens here, once, not per loop iteration).
|
||
const int64_t lmc_base = lmc::key_value_base_offset(
|
||
k_or_v, layer_id, kv_token_id, scalars_per_token, num_tokens,
|
||
num_layers);
|
||
const int64_t vllm_base = lmc::page_buffer_base_offset<format>(
|
||
k_or_v, slot_idx, scalars_per_token, page_buffer_size, block_size);
|
||
|
||
for (int i = tid; i < scalars_per_token; i += num_threads) {
|
||
if constexpr (DIRECTION) {
|
||
key_value_ptr[lmc_base + i] = paged_buffer_ptr[vllm_base + i];
|
||
} else {
|
||
paged_buffer_ptr[vllm_base + i] = key_value_ptr[lmc_base + i];
|
||
}
|
||
}
|
||
});
|
||
}
|
||
|
||
/**
|
||
* Submit a fused K+V multi-layer kernel for non-MLA formats.
|
||
*
|
||
* Processes both key (k_or_v=0) and value (k_or_v=1) within the same
|
||
* work-group, halving work-group count compared to dispatching K and
|
||
* V separately. Slot mapping, pointer-array lookup, and flash_infer's
|
||
* block-index division are each performed once and reused for both.
|
||
*/
|
||
template <typename scalar_t, bool DIRECTION, EngineKVFormat format>
|
||
void submit_multi_layer_kernel_fused_kv(
|
||
sycl::queue& queue, scalar_t* key_value_ptr, scalar_t** page_buffer_ptrs,
|
||
const int64_t* slot_mapping_ptr, int scalars_per_token, int num_tokens,
|
||
int num_layers, int page_buffer_size, int block_size,
|
||
int skip_prefix_n_tokens, int wg_size) {
|
||
int num_transfer_tokens = num_tokens - skip_prefix_n_tokens;
|
||
if (num_transfer_tokens <= 0 || num_layers <= 0) return;
|
||
|
||
// Grid: (1, num_layers, num_transfer_tokens * wg_size);
|
||
// no k_or_v dimension — K and V share one work-group.
|
||
sycl::range<3> global_range(
|
||
1, static_cast<size_t>(num_layers),
|
||
static_cast<size_t>(num_transfer_tokens) * wg_size);
|
||
sycl::range<3> local_range(1, 1, static_cast<size_t>(wg_size));
|
||
|
||
queue.parallel_for(
|
||
sycl::nd_range<3>(global_range, local_range),
|
||
[=](sycl::nd_item<3> item) [[sycl::reqd_sub_group_size(16)]] {
|
||
const int token_id = static_cast<int>(item.get_group(2));
|
||
const int layer_id = static_cast<int>(item.get_group(1));
|
||
const int tid = static_cast<int>(item.get_local_id(2));
|
||
const int num_threads = static_cast<int>(item.get_local_range(2));
|
||
|
||
const int kv_token_id = token_id + skip_prefix_n_tokens;
|
||
const int64_t slot_idx = slot_mapping_ptr[kv_token_id];
|
||
scalar_t* paged_buffer_ptr = page_buffer_ptrs[layer_id];
|
||
|
||
if (slot_idx < 0) return;
|
||
|
||
// Base offsets for K (k_or_v=0) and V (k_or_v=1); the
|
||
// flash_infer division/modulo runs once per token+layer.
|
||
const int64_t lmc_base_k = lmc::key_value_base_offset(
|
||
0, layer_id, kv_token_id, scalars_per_token, num_tokens,
|
||
num_layers);
|
||
const int64_t lmc_base_v = lmc::key_value_base_offset(
|
||
1, layer_id, kv_token_id, scalars_per_token, num_tokens,
|
||
num_layers);
|
||
const int64_t vllm_base_k = lmc::page_buffer_base_offset<format>(
|
||
0, slot_idx, scalars_per_token, page_buffer_size, block_size);
|
||
const int64_t vllm_base_v = lmc::page_buffer_base_offset<format>(
|
||
1, slot_idx, scalars_per_token, page_buffer_size, block_size);
|
||
|
||
for (int i = tid; i < scalars_per_token; i += num_threads) {
|
||
if constexpr (DIRECTION) {
|
||
// paged buffer → LMCache
|
||
key_value_ptr[lmc_base_k + i] = paged_buffer_ptr[vllm_base_k + i];
|
||
key_value_ptr[lmc_base_v + i] = paged_buffer_ptr[vllm_base_v + i];
|
||
} else {
|
||
// LMCache → paged buffer
|
||
paged_buffer_ptr[vllm_base_k + i] = key_value_ptr[lmc_base_k + i];
|
||
paged_buffer_ptr[vllm_base_v + i] = key_value_ptr[lmc_base_v + i];
|
||
}
|
||
}
|
||
});
|
||
}
|
||
|
||
/**
|
||
* Submit the multi-layer unilateral kernel (SGLang MHA).
|
||
*
|
||
* DIRECTION = true → paged buffer → LMCache (D2H)
|
||
* DIRECTION = false → LMCache → paged buffer (H2D)
|
||
*
|
||
* Uses `if constexpr (DIRECTION)` so the compiler can
|
||
* dead-strip the unused branch entirely.
|
||
*/
|
||
template <typename scalar_t, bool DIRECTION>
|
||
void submit_multi_layer_unilateral_kernel(
|
||
sycl::queue& queue, scalar_t* key_value_ptr, scalar_t** page_buffer_ptrs,
|
||
const int64_t* slot_mapping_ptr, int scalars_per_token, int num_tokens,
|
||
int num_layers, int page_buffer_size, int k_or_v_size, int kv_num_tokens,
|
||
int wg_size) {
|
||
if (kv_num_tokens <= 0 || num_layers <= 0) return;
|
||
|
||
sycl::range<3> global_range(static_cast<size_t>(k_or_v_size),
|
||
static_cast<size_t>(num_layers),
|
||
static_cast<size_t>(kv_num_tokens) * wg_size);
|
||
sycl::range<3> local_range(1, 1, static_cast<size_t>(wg_size));
|
||
|
||
queue.parallel_for(
|
||
sycl::nd_range<3>(global_range, local_range),
|
||
[=](sycl::nd_item<3> item) [[sycl::reqd_sub_group_size(16)]] {
|
||
const int token_id = static_cast<int>(item.get_group(2));
|
||
const int layer_id = static_cast<int>(item.get_group(1));
|
||
const int k_or_v = static_cast<int>(item.get_group(0));
|
||
const int tid = static_cast<int>(item.get_local_id(2));
|
||
const int num_threads = static_cast<int>(item.get_local_range(2));
|
||
|
||
const int64_t slot_idx = slot_mapping_ptr[token_id];
|
||
scalar_t* key_ptr = page_buffer_ptrs[layer_id];
|
||
scalar_t* value_ptr = page_buffer_ptrs[layer_id + num_layers];
|
||
|
||
if (slot_idx < 0) return;
|
||
|
||
const int64_t lmc_base = lmc::key_value_base_offset(
|
||
k_or_v, layer_id, token_id, scalars_per_token, num_tokens,
|
||
num_layers);
|
||
const int64_t sgl_base = slot_idx * scalars_per_token;
|
||
|
||
for (int i = tid; i < scalars_per_token; i += num_threads) {
|
||
if (k_or_v == 0) {
|
||
if constexpr (DIRECTION)
|
||
key_value_ptr[lmc_base + i] = key_ptr[sgl_base + i];
|
||
else
|
||
key_ptr[sgl_base + i] = key_value_ptr[lmc_base + i];
|
||
} else {
|
||
if constexpr (DIRECTION)
|
||
key_value_ptr[lmc_base + i] = value_ptr[sgl_base + i];
|
||
else
|
||
value_ptr[sgl_base + i] = key_value_ptr[lmc_base + i];
|
||
}
|
||
}
|
||
});
|
||
}
|
||
|
||
// ---------------------------------------------------------------------------
|
||
// Macros to dispatch multi-layer kernels with a specific EngineKVFormat.
|
||
// MLA formats (k_or_v_size==1) use the per-component kernel.
|
||
// Non-MLA formats (k_or_v_size==2) use the fused K+V kernel.
|
||
// ---------------------------------------------------------------------------
|
||
#define LAUNCH_KERNEL_WITH_FORMAT(T, DIRECTION, FORMAT) \
|
||
submit_multi_layer_kernel<T, DIRECTION, FORMAT>( \
|
||
queue, key_value_ptr, page_buffer_ptrs, slot_mapping_ptr, num_xwords, \
|
||
num_tokens, num_layers, page_buffer_size, block_size, \
|
||
skip_prefix_n_tokens, k_or_v_size, wg_size);
|
||
|
||
#define LAUNCH_FUSED_KV_KERNEL_WITH_FORMAT(T, DIRECTION, FORMAT) \
|
||
submit_multi_layer_kernel_fused_kv<T, DIRECTION, FORMAT>( \
|
||
queue, key_value_ptr, page_buffer_ptrs, slot_mapping_ptr, num_xwords, \
|
||
num_tokens, num_layers, page_buffer_size, block_size, \
|
||
skip_prefix_n_tokens, wg_size);
|
||
|
||
// ---------------------------------------------------------------------------
|
||
// multi_layer_kv_transfer -- templated implementation
|
||
// ---------------------------------------------------------------------------
|
||
template <typename T>
|
||
void multi_layer_kv_transfer_templated(
|
||
torch::Tensor& key_value, const torch::Tensor& key_value_ptrs,
|
||
const torch::Tensor& slot_mapping, const torch::Device& paged_memory_device,
|
||
const int page_buffer_size, const TransferDirection direction,
|
||
const EngineKVFormat engine_kv_format, const int block_size,
|
||
const int skip_prefix_n_tokens) {
|
||
T* key_value_ptr = get_kernel_ptr<T, torch::Tensor>(key_value);
|
||
T** page_buffer_ptrs =
|
||
get_kernel_ptr<T*, const torch::Tensor>(key_value_ptrs);
|
||
const int64_t* slot_mapping_ptr =
|
||
get_kernel_ptr<const int64_t, const torch::Tensor>(slot_mapping);
|
||
|
||
int num_layers = key_value.size(1);
|
||
int num_tokens = key_value.size(2);
|
||
int num_origin_elements = key_value.size(3);
|
||
int elements_per_xword = sizeof(T) / key_value.element_size();
|
||
int num_xwords = num_origin_elements / elements_per_xword;
|
||
|
||
int k_or_v_size = ::is_mla(engine_kv_format) ? 1 : 2;
|
||
|
||
// Round up to a sub-group multiple so every sub-group is full.
|
||
int wg_size = round_up_to_sg(std::min(num_xwords, MAX_WG_SIZE));
|
||
|
||
const c10::OptionalDeviceGuard device_guard(paged_memory_device);
|
||
sycl::queue& queue =
|
||
c10::xpu::getCurrentXPUStream(paged_memory_device.index()).queue();
|
||
|
||
// Non-MLA formats use the fused K+V kernel; MLA formats
|
||
// (k_or_v_size==1) use the per-component kernel.
|
||
if (k_or_v_size == 2) {
|
||
if (direction == TransferDirection::H2D) {
|
||
switch (engine_kv_format) {
|
||
case EngineKVFormat::NB_NL_TWO_BS_NH_HS:
|
||
LAUNCH_FUSED_KV_KERNEL_WITH_FORMAT(
|
||
T, false, EngineKVFormat::NB_NL_TWO_BS_NH_HS);
|
||
break;
|
||
case EngineKVFormat::NL_X_TWO_NB_BS_NH_HS:
|
||
LAUNCH_FUSED_KV_KERNEL_WITH_FORMAT(
|
||
T, false, EngineKVFormat::NL_X_TWO_NB_BS_NH_HS);
|
||
break;
|
||
case EngineKVFormat::NL_X_NB_TWO_BS_NH_HS:
|
||
LAUNCH_FUSED_KV_KERNEL_WITH_FORMAT(
|
||
T, false, EngineKVFormat::NL_X_NB_TWO_BS_NH_HS);
|
||
break;
|
||
default:
|
||
throw std::runtime_error("Unsupported non-MLA EngineKVFormat");
|
||
}
|
||
} else {
|
||
switch (engine_kv_format) {
|
||
case EngineKVFormat::NB_NL_TWO_BS_NH_HS:
|
||
LAUNCH_FUSED_KV_KERNEL_WITH_FORMAT(
|
||
T, true, EngineKVFormat::NB_NL_TWO_BS_NH_HS);
|
||
break;
|
||
case EngineKVFormat::NL_X_TWO_NB_BS_NH_HS:
|
||
LAUNCH_FUSED_KV_KERNEL_WITH_FORMAT(
|
||
T, true, EngineKVFormat::NL_X_TWO_NB_BS_NH_HS);
|
||
break;
|
||
case EngineKVFormat::NL_X_NB_TWO_BS_NH_HS:
|
||
LAUNCH_FUSED_KV_KERNEL_WITH_FORMAT(
|
||
T, true, EngineKVFormat::NL_X_NB_TWO_BS_NH_HS);
|
||
break;
|
||
default:
|
||
throw std::runtime_error("Unsupported non-MLA EngineKVFormat");
|
||
}
|
||
}
|
||
} else {
|
||
// MLA path (k_or_v_size == 1)
|
||
if (direction == TransferDirection::H2D) {
|
||
switch (engine_kv_format) {
|
||
case EngineKVFormat::NL_X_NB_BS_HS:
|
||
LAUNCH_KERNEL_WITH_FORMAT(T, false, EngineKVFormat::NL_X_NB_BS_HS);
|
||
break;
|
||
case EngineKVFormat::NL_X_NBBS_ONE_HS:
|
||
LAUNCH_KERNEL_WITH_FORMAT(T, false, EngineKVFormat::NL_X_NBBS_ONE_HS);
|
||
break;
|
||
default:
|
||
throw std::runtime_error("Unsupported MLA EngineKVFormat");
|
||
}
|
||
} else {
|
||
switch (engine_kv_format) {
|
||
case EngineKVFormat::NL_X_NB_BS_HS:
|
||
LAUNCH_KERNEL_WITH_FORMAT(T, true, EngineKVFormat::NL_X_NB_BS_HS);
|
||
break;
|
||
case EngineKVFormat::NL_X_NBBS_ONE_HS:
|
||
LAUNCH_KERNEL_WITH_FORMAT(T, true, EngineKVFormat::NL_X_NBBS_ONE_HS);
|
||
break;
|
||
default:
|
||
throw std::runtime_error("Unsupported MLA EngineKVFormat");
|
||
}
|
||
}
|
||
}
|
||
}
|
||
|
||
#undef LAUNCH_KERNEL_WITH_FORMAT
|
||
#undef LAUNCH_FUSED_KV_KERNEL_WITH_FORMAT
|
||
|
||
// ---------------------------------------------------------------------------
|
||
// Public API: multi_layer_kv_transfer
|
||
// ---------------------------------------------------------------------------
|
||
void multi_layer_kv_transfer(
|
||
torch::Tensor& key_value, const torch::Tensor& key_value_ptrs,
|
||
const torch::Tensor& slot_mapping, const torch::Device& paged_memory_device,
|
||
const int page_buffer_size, const TransferDirection direction,
|
||
const EngineKVFormat engine_kv_format, const int block_size,
|
||
const int head_size, const int skip_prefix_n_tokens) {
|
||
// head_size is currently unused in the SYCL implementation; accepted to
|
||
// keep ABI parity with the CUDA c_ops binding so callers can pass the
|
||
// same kwargs to either backend.
|
||
(void)head_size;
|
||
int num_origin_elements = key_value.size(3);
|
||
int copy_size = num_origin_elements * key_value.element_size();
|
||
|
||
#define LAUNCH_MULTI_LAYER_KV_TRANSFER(type) \
|
||
do { \
|
||
multi_layer_kv_transfer_templated<type>( \
|
||
key_value, key_value_ptrs, slot_mapping, paged_memory_device, \
|
||
page_buffer_size, direction, engine_kv_format, block_size, \
|
||
skip_prefix_n_tokens); \
|
||
} while (0)
|
||
if (copy_size % 8 == 0) {
|
||
LAUNCH_MULTI_LAYER_KV_TRANSFER(int64_t);
|
||
} else if (copy_size % 4 == 0) {
|
||
LAUNCH_MULTI_LAYER_KV_TRANSFER(int32_t);
|
||
} else if (copy_size % 2 == 0) {
|
||
LAUNCH_MULTI_LAYER_KV_TRANSFER(int16_t);
|
||
} else {
|
||
LAUNCH_MULTI_LAYER_KV_TRANSFER(int8_t);
|
||
}
|
||
#undef LAUNCH_MULTI_LAYER_KV_TRANSFER
|
||
}
|
||
|
||
// ---------------------------------------------------------------------------
|
||
// Public API: multi_layer_kv_transfer_unilateral
|
||
// ---------------------------------------------------------------------------
|
||
void multi_layer_kv_transfer_unilateral(
|
||
torch::Tensor& key_value, const torch::Tensor& key_value_ptrs,
|
||
const torch::Tensor& slot_mapping, const torch::Device& paged_memory_device,
|
||
const int page_buffer_size, const TransferDirection direction,
|
||
const EngineKVFormat engine_kv_format) {
|
||
const bool use_mla = ::is_mla(engine_kv_format);
|
||
// MLA case collapses back to multi_layer_kv_transfer
|
||
if (use_mla) {
|
||
return multi_layer_kv_transfer(key_value, key_value_ptrs, slot_mapping,
|
||
paged_memory_device, page_buffer_size,
|
||
direction, engine_kv_format);
|
||
}
|
||
|
||
int64_t* key_value_ptr = get_kernel_ptr<int64_t, torch::Tensor>(key_value);
|
||
int64_t** page_buffer_ptrs =
|
||
get_kernel_ptr<int64_t*, const torch::Tensor>(key_value_ptrs);
|
||
const int64_t* slot_mapping_ptr =
|
||
get_kernel_ptr<const int64_t, const torch::Tensor>(slot_mapping);
|
||
|
||
int num_layers = key_value.size(1);
|
||
int num_tokens = slot_mapping.size(0);
|
||
int num_origin_elements = key_value.size(3);
|
||
int elements_per_qword = 8 / key_value.element_size();
|
||
int num_qwords = num_origin_elements / elements_per_qword;
|
||
|
||
int k_or_v_size = 2;
|
||
int kv_num_tokens = key_value.size(2);
|
||
int wg_size = round_up_to_sg(std::min(num_qwords, MAX_WG_SIZE));
|
||
|
||
const c10::OptionalDeviceGuard device_guard(paged_memory_device);
|
||
sycl::queue& queue =
|
||
c10::xpu::getCurrentXPUStream(paged_memory_device.index()).queue();
|
||
|
||
if (direction == TransferDirection::H2D) {
|
||
submit_multi_layer_unilateral_kernel<int64_t, false>(
|
||
queue, key_value_ptr, page_buffer_ptrs, slot_mapping_ptr, num_qwords,
|
||
num_tokens, num_layers, page_buffer_size, k_or_v_size, kv_num_tokens,
|
||
wg_size);
|
||
} else {
|
||
submit_multi_layer_unilateral_kernel<int64_t, true>(
|
||
queue, key_value_ptr, page_buffer_ptrs, slot_mapping_ptr, num_qwords,
|
||
num_tokens, num_layers, page_buffer_size, k_or_v_size, kv_num_tokens,
|
||
wg_size);
|
||
}
|
||
}
|
||
|
||
// ---------------------------------------------------------------------------
|
||
// single_layer_kv_transfer — helper template
|
||
// ---------------------------------------------------------------------------
|
||
// USE_MLA and IS_D2H are template parameters so the compiler can
|
||
// dead-strip the unused branch.
|
||
template <bool USE_MLA, bool IS_D2H>
|
||
void single_layer_kv_transfer_impl(sycl::queue& queue, int64_t* lmc_ptr,
|
||
int64_t* vllm_ptr, const int64_t* slot_ptr,
|
||
int num_tokens, int n, int lmc_stride,
|
||
int lmc_value_offset, int block_size,
|
||
int vllm_block_key_stride_in_64bit,
|
||
int vllm_value_offset, int num_heads,
|
||
int head_size_in_64bit, int wg_size) {
|
||
if (num_tokens <= 0) return;
|
||
|
||
sycl::range<1> global_range(static_cast<size_t>(num_tokens) * wg_size);
|
||
sycl::range<1> local_range(static_cast<size_t>(wg_size));
|
||
|
||
queue.parallel_for(
|
||
sycl::nd_range<1>(global_range, local_range),
|
||
[=](sycl::nd_item<1> item) [[sycl::reqd_sub_group_size(16)]] {
|
||
const int64_t token_idx = static_cast<int64_t>(item.get_group(0));
|
||
const int64_t slot_idx = slot_ptr[token_idx];
|
||
if (slot_idx < 0) return;
|
||
|
||
const int64_t block_idx = slot_idx / block_size;
|
||
const int64_t block_offset = slot_idx % block_size;
|
||
|
||
const int tid = static_cast<int>(item.get_local_id(0));
|
||
const int nthreads = static_cast<int>(item.get_local_range(0));
|
||
|
||
for (int i = tid; i < n; i += nthreads) {
|
||
const int64_t lmc_key_idx = token_idx * lmc_stride + i;
|
||
const int head_idx = i / head_size_in_64bit;
|
||
const int head_offset = i % head_size_in_64bit;
|
||
const int64_t vllm_key_idx =
|
||
block_idx * vllm_block_key_stride_in_64bit +
|
||
block_offset * num_heads * head_size_in_64bit +
|
||
head_idx * head_size_in_64bit + head_offset;
|
||
|
||
if constexpr (IS_D2H) {
|
||
lmc_ptr[lmc_key_idx] = vllm_ptr[vllm_key_idx];
|
||
if constexpr (!USE_MLA) {
|
||
const int64_t lmc_value_idx = lmc_key_idx + lmc_value_offset;
|
||
const int64_t vllm_value_idx = vllm_key_idx + vllm_value_offset;
|
||
lmc_ptr[lmc_value_idx] = vllm_ptr[vllm_value_idx];
|
||
}
|
||
} else {
|
||
vllm_ptr[vllm_key_idx] = lmc_ptr[lmc_key_idx];
|
||
if constexpr (!USE_MLA) {
|
||
const int64_t lmc_value_idx = lmc_key_idx + lmc_value_offset;
|
||
const int64_t vllm_value_idx = vllm_key_idx + vllm_value_offset;
|
||
vllm_ptr[vllm_value_idx] = lmc_ptr[lmc_value_idx];
|
||
}
|
||
}
|
||
}
|
||
});
|
||
}
|
||
|
||
// ---------------------------------------------------------------------------
|
||
// Public API: single_layer_kv_transfer
|
||
// ---------------------------------------------------------------------------
|
||
void single_layer_kv_transfer(torch::Tensor& lmc_key_value_cache,
|
||
torch::Tensor& vllm_key_value_cache,
|
||
torch::Tensor& slot_mapping,
|
||
const TransferDirection direction,
|
||
const EngineKVFormat engine_kv_format,
|
||
const bool token_major) {
|
||
int64_t* lmc_key_value_cache_ptr =
|
||
get_kernel_ptr<int64_t, torch::Tensor>(lmc_key_value_cache);
|
||
int64_t* vllm_key_value_cache_ptr =
|
||
get_kernel_ptr<int64_t, torch::Tensor>(vllm_key_value_cache);
|
||
const int64_t* slot_mapping_ptr =
|
||
get_kernel_ptr<const int64_t, const torch::Tensor>(slot_mapping);
|
||
|
||
int elements_per_entry = 8 / vllm_key_value_cache.element_size();
|
||
|
||
int num_tokens = slot_mapping.size(0);
|
||
int num_heads;
|
||
int head_size_in_64bit;
|
||
int block_size;
|
||
|
||
const bool use_mla = ::is_mla(engine_kv_format);
|
||
|
||
if (use_mla) {
|
||
num_heads = 1;
|
||
block_size = vllm_key_value_cache.size(1);
|
||
head_size_in_64bit = vllm_key_value_cache.size(2) / elements_per_entry;
|
||
} else {
|
||
num_heads = vllm_key_value_cache.size(3);
|
||
head_size_in_64bit = vllm_key_value_cache.size(4) / elements_per_entry;
|
||
block_size = vllm_key_value_cache.size(2);
|
||
}
|
||
|
||
int lmc_stride;
|
||
int lmc_value_offset;
|
||
if (use_mla) {
|
||
lmc_stride = lmc_key_value_cache.stride(0) / elements_per_entry;
|
||
lmc_value_offset = 0;
|
||
} else if (token_major) {
|
||
lmc_stride = lmc_key_value_cache.stride(0) / elements_per_entry;
|
||
lmc_value_offset = lmc_key_value_cache.stride(1) / elements_per_entry;
|
||
} else {
|
||
lmc_stride = lmc_key_value_cache.stride(1) / elements_per_entry;
|
||
lmc_value_offset = lmc_key_value_cache.stride(0) / elements_per_entry;
|
||
}
|
||
|
||
int vllm_block_key_stride_in_64bit;
|
||
int vllm_value_offset;
|
||
if (use_mla) {
|
||
vllm_block_key_stride_in_64bit =
|
||
vllm_key_value_cache.stride(0) / elements_per_entry;
|
||
vllm_value_offset = 0;
|
||
} else if (engine_kv_format == EngineKVFormat::NL_X_TWO_NB_BS_NH_HS) {
|
||
vllm_block_key_stride_in_64bit =
|
||
vllm_key_value_cache.stride(1) / elements_per_entry;
|
||
vllm_value_offset = vllm_key_value_cache.stride(0) / elements_per_entry;
|
||
} else if (engine_kv_format == EngineKVFormat::NL_X_NB_TWO_BS_NH_HS) {
|
||
vllm_block_key_stride_in_64bit =
|
||
vllm_key_value_cache.stride(0) / elements_per_entry;
|
||
vllm_value_offset = vllm_key_value_cache.stride(1) / elements_per_entry;
|
||
} else {
|
||
throw std::runtime_error(
|
||
"Unsupported non-MLA EngineKVFormat in single_layer_kv_transfer: " +
|
||
std::to_string(static_cast<int>(engine_kv_format)));
|
||
}
|
||
|
||
int n = num_heads * head_size_in_64bit;
|
||
int wg_size = round_up_to_sg(std::min(n, MAX_WG_SIZE));
|
||
if (num_tokens <= 0) return;
|
||
|
||
const c10::OptionalDeviceGuard device_guard(device_of(vllm_key_value_cache));
|
||
sycl::queue& queue =
|
||
c10::xpu::getCurrentXPUStream(vllm_key_value_cache.device().index())
|
||
.queue();
|
||
|
||
auto lmc_ptr = lmc_key_value_cache_ptr;
|
||
auto vllm_ptr = vllm_key_value_cache_ptr;
|
||
auto slot_ptr = slot_mapping_ptr;
|
||
|
||
// Dispatch to 4 compile-time specialisations (USE_MLA × IS_D2H)
|
||
// so the inner loop is branch-free.
|
||
if (use_mla) {
|
||
if (direction == TransferDirection::D2H)
|
||
single_layer_kv_transfer_impl<true, true>(
|
||
queue, lmc_ptr, vllm_ptr, slot_ptr, num_tokens, n, lmc_stride,
|
||
lmc_value_offset, block_size, vllm_block_key_stride_in_64bit,
|
||
vllm_value_offset, num_heads, head_size_in_64bit, wg_size);
|
||
else
|
||
single_layer_kv_transfer_impl<true, false>(
|
||
queue, lmc_ptr, vllm_ptr, slot_ptr, num_tokens, n, lmc_stride,
|
||
lmc_value_offset, block_size, vllm_block_key_stride_in_64bit,
|
||
vllm_value_offset, num_heads, head_size_in_64bit, wg_size);
|
||
} else {
|
||
if (direction == TransferDirection::D2H)
|
||
single_layer_kv_transfer_impl<false, true>(
|
||
queue, lmc_ptr, vllm_ptr, slot_ptr, num_tokens, n, lmc_stride,
|
||
lmc_value_offset, block_size, vllm_block_key_stride_in_64bit,
|
||
vllm_value_offset, num_heads, head_size_in_64bit, wg_size);
|
||
else
|
||
single_layer_kv_transfer_impl<false, false>(
|
||
queue, lmc_ptr, vllm_ptr, slot_ptr, num_tokens, n, lmc_stride,
|
||
lmc_value_offset, block_size, vllm_block_key_stride_in_64bit,
|
||
vllm_value_offset, num_heads, head_size_in_64bit, wg_size);
|
||
}
|
||
}
|
||
|
||
// ---------------------------------------------------------------------------
|
||
// single_layer_kv_transfer_sgl — helper template
|
||
// ---------------------------------------------------------------------------
|
||
template <bool IS_D2H>
|
||
void single_layer_kv_transfer_sgl_impl(
|
||
sycl::queue& queue, int64_t* lmc_ptr, int64_t* sgl_k_ptr,
|
||
int64_t* sgl_v_ptr, const int64_t* slot_ptr, int num_tokens, int n,
|
||
int lmc_stride, int lmc_value_offset, int block_stride_in_64bit,
|
||
int block_size, int num_heads, int head_size_in_64bit, int wg_size) {
|
||
if (num_tokens <= 0) return;
|
||
|
||
sycl::range<1> global_range(static_cast<size_t>(num_tokens) * wg_size);
|
||
sycl::range<1> local_range(static_cast<size_t>(wg_size));
|
||
|
||
queue.parallel_for(
|
||
sycl::nd_range<1>(global_range, local_range),
|
||
[=](sycl::nd_item<1> item) [[sycl::reqd_sub_group_size(16)]] {
|
||
const int64_t token_idx = static_cast<int64_t>(item.get_group(0));
|
||
const int64_t slot_idx = slot_ptr[token_idx];
|
||
if (slot_idx < 0) return;
|
||
|
||
const int64_t block_idx = slot_idx / block_size;
|
||
const int64_t block_offset = slot_idx % block_size;
|
||
|
||
const int tid = static_cast<int>(item.get_local_id(0));
|
||
const int nthreads = static_cast<int>(item.get_local_range(0));
|
||
|
||
for (int i = tid; i < n; i += nthreads) {
|
||
const int64_t lmc_key_idx = token_idx * lmc_stride + i;
|
||
const int64_t lmc_value_idx = lmc_key_idx + lmc_value_offset;
|
||
|
||
const int head_idx = i / head_size_in_64bit;
|
||
const int head_offset = i % head_size_in_64bit;
|
||
const int64_t sgl_kv_idx =
|
||
block_idx * block_stride_in_64bit +
|
||
block_offset * num_heads * head_size_in_64bit +
|
||
head_idx * head_size_in_64bit + head_offset;
|
||
|
||
if constexpr (IS_D2H) {
|
||
lmc_ptr[lmc_key_idx] = sgl_k_ptr[sgl_kv_idx];
|
||
lmc_ptr[lmc_value_idx] = sgl_v_ptr[sgl_kv_idx];
|
||
} else {
|
||
sgl_k_ptr[sgl_kv_idx] = lmc_ptr[lmc_key_idx];
|
||
sgl_v_ptr[sgl_kv_idx] = lmc_ptr[lmc_value_idx];
|
||
}
|
||
}
|
||
});
|
||
}
|
||
|
||
// ---------------------------------------------------------------------------
|
||
// Public API: single_layer_kv_transfer_sgl (SGLang)
|
||
// ---------------------------------------------------------------------------
|
||
void single_layer_kv_transfer_sgl(torch::Tensor& lmc_key_value_cache,
|
||
torch::Tensor& sgl_key_cache,
|
||
torch::Tensor& sgl_value_cache,
|
||
torch::Tensor& slot_mapping,
|
||
const TransferDirection direction,
|
||
const bool token_major) {
|
||
int64_t* lmc_key_value_cache_ptr =
|
||
get_kernel_ptr<int64_t, torch::Tensor>(lmc_key_value_cache);
|
||
int64_t* sgl_key_cache_ptr =
|
||
get_kernel_ptr<int64_t, torch::Tensor>(sgl_key_cache);
|
||
int64_t* sgl_value_cache_ptr =
|
||
get_kernel_ptr<int64_t, torch::Tensor>(sgl_value_cache);
|
||
const int64_t* slot_mapping_ptr =
|
||
get_kernel_ptr<const int64_t, const torch::Tensor>(slot_mapping);
|
||
|
||
int elements_per_entry = 8 / sgl_key_cache.element_size();
|
||
|
||
int num_tokens = slot_mapping.size(0);
|
||
int num_heads = sgl_key_cache.size(2);
|
||
int head_size_in_64bit = sgl_key_cache.size(3) / elements_per_entry;
|
||
int block_size = sgl_key_cache.size(1);
|
||
|
||
int lmc_stride;
|
||
int lmc_value_offset;
|
||
if (token_major) {
|
||
lmc_stride = lmc_key_value_cache.stride(0) / elements_per_entry;
|
||
lmc_value_offset = lmc_key_value_cache.stride(1) / elements_per_entry;
|
||
} else {
|
||
lmc_stride = lmc_key_value_cache.stride(1) / elements_per_entry;
|
||
lmc_value_offset = lmc_key_value_cache.stride(0) / elements_per_entry;
|
||
}
|
||
|
||
int block_stride_in_64bit = sgl_key_cache.stride(0) / elements_per_entry;
|
||
TORCH_CHECK(sgl_key_cache.stride(0) == sgl_value_cache.stride(0));
|
||
|
||
int n = num_heads * head_size_in_64bit;
|
||
int wg_size = round_up_to_sg(std::min(n, MAX_WG_SIZE));
|
||
if (num_tokens <= 0) return;
|
||
|
||
const c10::OptionalDeviceGuard device_guard(device_of(sgl_key_cache));
|
||
sycl::queue& queue =
|
||
c10::xpu::getCurrentXPUStream(sgl_key_cache.device().index()).queue();
|
||
|
||
if (direction == TransferDirection::D2H)
|
||
single_layer_kv_transfer_sgl_impl<true>(
|
||
queue, lmc_key_value_cache_ptr, sgl_key_cache_ptr, sgl_value_cache_ptr,
|
||
slot_mapping_ptr, num_tokens, n, lmc_stride, lmc_value_offset,
|
||
block_stride_in_64bit, block_size, num_heads, head_size_in_64bit,
|
||
wg_size);
|
||
else
|
||
single_layer_kv_transfer_sgl_impl<false>(
|
||
queue, lmc_key_value_cache_ptr, sgl_key_cache_ptr, sgl_value_cache_ptr,
|
||
slot_mapping_ptr, num_tokens, n, lmc_stride, lmc_value_offset,
|
||
block_stride_in_64bit, block_size, num_heads, head_size_in_64bit,
|
||
wg_size);
|
||
}
|
||
|
||
// ---------------------------------------------------------------------------
|
||
// Public API: load_and_reshape_flash (deprecated -- unit tests only)
|
||
// ---------------------------------------------------------------------------
|
||
void load_and_reshape_flash(torch::Tensor& key_value, torch::Tensor& key_cache,
|
||
torch::Tensor& value_cache,
|
||
torch::Tensor& slot_mapping, const int layer_idx) {
|
||
int64_t* key_value_ptr = get_kernel_ptr<int64_t, torch::Tensor>(key_value);
|
||
int64_t* key_cache_ptr = get_kernel_ptr<int64_t, torch::Tensor>(key_cache);
|
||
int64_t* value_cache_ptr =
|
||
get_kernel_ptr<int64_t, torch::Tensor>(value_cache);
|
||
const int64_t* slot_mapping_ptr =
|
||
get_kernel_ptr<const int64_t, const torch::Tensor>(slot_mapping);
|
||
|
||
int elements_per_entry = 8 / key_cache.element_size();
|
||
int num_tokens = slot_mapping.size(0);
|
||
int num_heads = key_cache.size(2);
|
||
int head_size_in_64bit = key_cache.size(3) / elements_per_entry;
|
||
int block_size = key_cache.size(1);
|
||
int key_value_stride = key_value.stride(2) / elements_per_entry;
|
||
int num_layers = key_value.size(1);
|
||
int key_layer_offset = layer_idx * key_value.stride(1) / elements_per_entry;
|
||
int value_layer_offset =
|
||
(layer_idx + num_layers) * key_value.stride(1) / elements_per_entry;
|
||
int block_stride_in_64bit = key_cache.stride(0) / elements_per_entry;
|
||
TORCH_CHECK(key_cache.stride(0) == value_cache.stride(0));
|
||
|
||
int n = num_heads * head_size_in_64bit;
|
||
int wg_size = round_up_to_sg(std::min(n, MAX_WG_SIZE));
|
||
if (num_tokens <= 0) return;
|
||
|
||
sycl::range<1> global_range(static_cast<size_t>(num_tokens) * wg_size);
|
||
sycl::range<1> local_range(static_cast<size_t>(wg_size));
|
||
|
||
const c10::OptionalDeviceGuard device_guard(device_of(key_cache));
|
||
sycl::queue& queue =
|
||
c10::xpu::getCurrentXPUStream(key_cache.device().index()).queue();
|
||
|
||
auto kv_ptr = key_value_ptr;
|
||
auto k_ptr = key_cache_ptr;
|
||
auto v_ptr = value_cache_ptr;
|
||
auto slot_ptr = slot_mapping_ptr;
|
||
|
||
queue.parallel_for(
|
||
sycl::nd_range<1>(global_range, local_range),
|
||
[=](sycl::nd_item<1> item) [[sycl::reqd_sub_group_size(16)]] {
|
||
const int64_t token_idx = static_cast<int64_t>(item.get_group(0));
|
||
const int64_t slot_idx = slot_ptr[token_idx];
|
||
if (slot_idx < 0) return;
|
||
|
||
const int64_t blk_idx = slot_idx / block_size;
|
||
const int64_t blk_off = slot_idx % block_size;
|
||
|
||
const int tid = static_cast<int>(item.get_local_id(0));
|
||
const int nthreads = static_cast<int>(item.get_local_range(0));
|
||
|
||
for (int i = tid; i < n; i += nthreads) {
|
||
const int64_t tgt_key_idx =
|
||
key_layer_offset + token_idx * key_value_stride + i;
|
||
const int64_t tgt_value_idx =
|
||
value_layer_offset + token_idx * key_value_stride + i;
|
||
|
||
const int head_idx = i / head_size_in_64bit;
|
||
const int head_offset = i % head_size_in_64bit;
|
||
const int64_t src_kv_idx = blk_idx * block_stride_in_64bit +
|
||
blk_off * num_heads * head_size_in_64bit +
|
||
head_idx * head_size_in_64bit +
|
||
head_offset;
|
||
|
||
kv_ptr[tgt_key_idx] = k_ptr[src_kv_idx];
|
||
kv_ptr[tgt_value_idx] = v_ptr[src_kv_idx];
|
||
}
|
||
});
|
||
}
|
||
|
||
// ---------------------------------------------------------------------------
|
||
// Public API: reshape_and_cache_back_flash (deprecated -- unit
|
||
// tests only)
|
||
// ---------------------------------------------------------------------------
|
||
void reshape_and_cache_back_flash(torch::Tensor& key_value,
|
||
torch::Tensor& key_cache,
|
||
torch::Tensor& value_cache,
|
||
torch::Tensor& slot_mapping,
|
||
const int layer_idx) {
|
||
int64_t* key_cache_ptr = get_kernel_ptr<int64_t, torch::Tensor>(key_cache);
|
||
int64_t* value_cache_ptr =
|
||
get_kernel_ptr<int64_t, torch::Tensor>(value_cache);
|
||
int64_t* key_value_ptr = get_kernel_ptr<int64_t, torch::Tensor>(key_value);
|
||
const int64_t* slot_mapping_ptr =
|
||
get_kernel_ptr<const int64_t, const torch::Tensor>(slot_mapping);
|
||
|
||
int elements_per_entry = 8 / key_cache.element_size();
|
||
int num_tokens = slot_mapping.size(0);
|
||
int num_heads = key_cache.size(2);
|
||
int head_size_in_64bit = key_cache.size(3) / elements_per_entry;
|
||
int block_size = key_cache.size(1);
|
||
int key_value_stride = key_value.stride(2) / elements_per_entry;
|
||
int num_layers = key_value.size(1);
|
||
int key_layer_offset = layer_idx * key_value.stride(1) / elements_per_entry;
|
||
int value_layer_offset =
|
||
(layer_idx + num_layers) * key_value.stride(1) / elements_per_entry;
|
||
int block_stride_in_64bit = key_cache.stride(0) / elements_per_entry;
|
||
TORCH_CHECK(key_cache.stride(0) == value_cache.stride(0));
|
||
|
||
int n = num_heads * head_size_in_64bit;
|
||
int wg_size = round_up_to_sg(std::min(n, MAX_WG_SIZE));
|
||
if (num_tokens <= 0) return;
|
||
|
||
sycl::range<1> global_range(static_cast<size_t>(num_tokens) * wg_size);
|
||
sycl::range<1> local_range(static_cast<size_t>(wg_size));
|
||
|
||
const c10::OptionalDeviceGuard device_guard(device_of(key_cache));
|
||
sycl::queue& queue =
|
||
c10::xpu::getCurrentXPUStream(key_cache.device().index()).queue();
|
||
|
||
auto kv_ptr = key_value_ptr;
|
||
auto k_ptr = key_cache_ptr;
|
||
auto v_ptr = value_cache_ptr;
|
||
auto slot_ptr = slot_mapping_ptr;
|
||
|
||
queue.parallel_for(
|
||
sycl::nd_range<1>(global_range, local_range),
|
||
[=](sycl::nd_item<1> item) [[sycl::reqd_sub_group_size(16)]] {
|
||
const int64_t token_idx = static_cast<int64_t>(item.get_group(0));
|
||
const int64_t slot_idx = slot_ptr[token_idx];
|
||
if (slot_idx < 0) return;
|
||
|
||
const int64_t blk_idx = slot_idx / block_size;
|
||
const int64_t blk_off = slot_idx % block_size;
|
||
|
||
const int tid = static_cast<int>(item.get_local_id(0));
|
||
const int nthreads = static_cast<int>(item.get_local_range(0));
|
||
|
||
for (int i = tid; i < n; i += nthreads) {
|
||
const int64_t tgt_key_idx =
|
||
key_layer_offset + token_idx * key_value_stride + i;
|
||
const int64_t tgt_value_idx =
|
||
value_layer_offset + token_idx * key_value_stride + i;
|
||
|
||
const int head_idx = i / head_size_in_64bit;
|
||
const int head_offset = i % head_size_in_64bit;
|
||
const int64_t src_kv_idx = blk_idx * block_stride_in_64bit +
|
||
blk_off * num_heads * head_size_in_64bit +
|
||
head_idx * head_size_in_64bit +
|
||
head_offset;
|
||
|
||
k_ptr[src_kv_idx] = kv_ptr[tgt_key_idx];
|
||
v_ptr[src_kv_idx] = kv_ptr[tgt_value_idx];
|
||
}
|
||
});
|
||
}
|
||
|
||
// ---------------------------------------------------------------------------
|
||
// Public API: lmcache_memcpy_async
|
||
// ---------------------------------------------------------------------------
|
||
void lmcache_memcpy_async(uintptr_t dest, uintptr_t src, size_t nbytes,
|
||
TransferDirection direction,
|
||
size_t host_buffer_offset,
|
||
size_t host_buffer_alignments) {
|
||
TORCH_CHECK((host_buffer_alignments & (host_buffer_alignments - 1)) == 0,
|
||
"host_buffer_alignments must be power of two");
|
||
|
||
// SYCL USM memcpy infers direction from pointer allocation types;
|
||
// the `direction` parameter is retained only for API compatibility.
|
||
(void)direction;
|
||
|
||
sycl::queue& queue = c10::xpu::getCurrentXPUStream().queue();
|
||
|
||
size_t offset = 0;
|
||
const size_t mask = host_buffer_alignments - 1;
|
||
|
||
while (offset < nbytes) {
|
||
size_t current_src = src + offset;
|
||
size_t current_dest = dest + offset;
|
||
|
||
size_t aligned_area_end =
|
||
((offset + host_buffer_offset) & ~mask) + host_buffer_alignments;
|
||
size_t real_end = std::min(host_buffer_offset + nbytes, aligned_area_end);
|
||
size_t max_nbytes = real_end - offset - host_buffer_offset;
|
||
|
||
// USM memcpy is direction-agnostic.
|
||
queue.memcpy(reinterpret_cast<void*>(current_dest),
|
||
reinterpret_cast<const void*>(current_src), max_nbytes);
|
||
|
||
offset += max_nbytes;
|
||
}
|
||
}
|