lmcache--lmcache
wehub-resource-sync
42 行
1.8 KiB
C
42 行
1.8 KiB
C
// SPDX-License-Identifier: Apache-2.0
|
|
//
|
|
// Declarations of SYCL/XPU implementations of CacheGen + RoPE kernels.
|
|
// Implementations live in csrc/sycl/*.cpp and are exposed to Python via
|
|
// lmcache.xpu_ops.
|
|
//
|
|
#pragma once
|
|
|
|
#include <torch/all.h>
|
|
|
|
// CacheGen ---------------------------------------------------------------
|
|
//
|
|
// Calculate per-(layer, channel) CDF from an int8/uint8 input tensor.
|
|
// Input: [nlayers, ntokens, nchannels] (uint8)
|
|
// Output: [nlayers, nchannels, max_bins + 1] (int16)
|
|
at::Tensor calculate_cdf_xpu(const at::Tensor& input, int64_t max_bins);
|
|
|
|
// Arithmetic encoder (forward): produces a byte buffer + per-channel
|
|
// length tensor.
|
|
void encode_fast_new_xpu(const at::Tensor& cdf, const at::Tensor& input_sym,
|
|
at::Tensor& output_buffer, at::Tensor& output_lengths);
|
|
|
|
// Arithmetic decoder (per-channel buffer): inverse of encode_fast_new.
|
|
void decode_fast_new_xpu(const at::Tensor& cdf, const at::Tensor& bytestreams,
|
|
const at::Tensor& lengths, at::Tensor& output);
|
|
|
|
// Arithmetic decoder (1-D bytestream with prefix-sum offsets):
|
|
// inverse of encode_fast_new packed via prefix sums.
|
|
void decode_fast_prefsum_xpu(const at::Tensor& cdf,
|
|
const at::Tensor& bytestreams,
|
|
const at::Tensor& lengths_prefsum,
|
|
at::Tensor& output);
|
|
|
|
// Position encoding -------------------------------------------------------
|
|
//
|
|
// Fused undo-then-apply rotary embedding on key tensor (in-place).
|
|
void rotary_embedding_k_fused_xpu(const torch::Tensor& old_positions,
|
|
const torch::Tensor& new_positions,
|
|
torch::Tensor& key, int64_t head_size,
|
|
const torch::Tensor& cos_sin_cache,
|
|
bool is_neox);
|