wehub-resource-sync
144 lines
3.7 KiB
C
144 lines
3.7 KiB
C
// SPDX-License-Identifier: Apache-2.0
|
|
|
|
#pragma once
|
|
|
|
// Physical KV-cache memory layout an engine hands to LMCache, plus the
|
|
// classification predicates over it. Vendor-header-free, so every backend and
|
|
// the Python facade (lmc_ops) share one definition. Detection (raw layout ->
|
|
// format) lives in lmcache/v1/gpu_connector/kv_format.
|
|
|
|
/*
|
|
Symbol Reference:
|
|
NL: number of layers
|
|
NB: number of blocks/pages
|
|
BS: block/page size
|
|
NBBS: block/page buffer size = NB * BS
|
|
NH: number of heads
|
|
HS: head size
|
|
TWO: 2
|
|
ONE: 1
|
|
|
|
_ means a dimension within the same tensor
|
|
_X_ means a dimension across a list
|
|
|
|
A_X_B_X_C_D_E means:
|
|
kv_cache: List[List[torch.Tensor]]
|
|
len(kv_cache) = A
|
|
len(kv_cache[0]) = B
|
|
kv_cache[0][0].shape = (C, D, E)
|
|
*/
|
|
enum class EngineKVFormat : int {
|
|
NB_NL_TWO_BS_NH_HS = 0,
|
|
/*
|
|
used by:
|
|
- vLLM CROSS_LAYER mode
|
|
*/
|
|
|
|
NL_X_TWO_NB_BS_NH_HS = 1,
|
|
/*
|
|
used by:
|
|
- vLLM non-MLA flash attention
|
|
*/
|
|
|
|
NL_X_NB_TWO_BS_NH_HS = 2,
|
|
/*
|
|
used by:
|
|
- vLLM non-MLA flash infer
|
|
*/
|
|
|
|
NL_X_NB_BS_HS = 3,
|
|
/*
|
|
used by:
|
|
- vLLM MLA
|
|
*/
|
|
|
|
TWO_X_NL_X_NBBS_NH_HS = 4,
|
|
/*
|
|
used by:
|
|
- SGLang MHA (flash attention and flash infer)
|
|
*/
|
|
|
|
NL_X_NBBS_ONE_HS = 5,
|
|
/*
|
|
used by:
|
|
- SGLang MLA
|
|
*/
|
|
|
|
NL_X_TWO_NB_NH_BS_HS = 6,
|
|
/*
|
|
used by:
|
|
- vLLM non-MLA flash attention (HND layout)
|
|
physical shape per layer: [2, num_blocks, num_heads, block_size, head_size]
|
|
*/
|
|
|
|
NL_X_NB_TWO_NH_BS_HS = 7,
|
|
/*
|
|
used by:
|
|
- vLLM non-MLA flash infer (HND layout)
|
|
physical shape per layer: [num_blocks, 2, num_heads, block_size, head_size]
|
|
*/
|
|
|
|
NB_NL_TWO_NH_BS_HS = 8,
|
|
/*
|
|
used by:
|
|
- TRT-LLM cross-layer (HND layout)
|
|
physical shape: [num_blocks, num_layers, 2, num_heads, block_size, head_size]
|
|
*/
|
|
|
|
TWO_X_NL_X_NB_BS_NH_HS = 9,
|
|
/*
|
|
used by:
|
|
- SGLang MHA via the MP daemon path
|
|
physical shape per layer: [num_blocks, block_size, num_heads, head_size]
|
|
*/
|
|
|
|
NL_X_NB_NH_BS_TWO_HS = 10,
|
|
/*
|
|
used by:
|
|
- vLLM non-MLA blocks-first attention with K/V fused into the trailing dim
|
|
physical shape per layer: [num_blocks, num_heads, block_size, 2, head_size]
|
|
(recovered by splitting the fused trailing [block_size, 2 * head_size]).
|
|
Currently only reached via the host gather/scatter path, not the device
|
|
transfer kernels.
|
|
*/
|
|
};
|
|
|
|
// __host__ __device__ under CUDA/HIP so the kernels can call these; the guard
|
|
// keeps the header vendor-runtime-free.
|
|
#if defined(__CUDACC__) || defined(__HIPCC__)
|
|
#define LMC_KV_FORMAT_HD __host__ __device__
|
|
#else
|
|
#define LMC_KV_FORMAT_HD
|
|
#endif
|
|
|
|
// Structural shape of the normalized kv_caches: exactly one is true per format.
|
|
|
|
// All layers in one fused tensor.
|
|
LMC_KV_FORMAT_HD constexpr bool is_cross_layer(EngineKVFormat f) {
|
|
return f == EngineKVFormat::NB_NL_TWO_BS_NH_HS ||
|
|
f == EngineKVFormat::NB_NL_TWO_NH_BS_HS;
|
|
}
|
|
|
|
// Keys and values in two separate top-level lists: [key_layers, value_layers].
|
|
LMC_KV_FORMAT_HD constexpr bool is_kv_list(EngineKVFormat f) {
|
|
return f == EngineKVFormat::TWO_X_NL_X_NBBS_NH_HS ||
|
|
f == EngineKVFormat::TWO_X_NL_X_NB_BS_NH_HS;
|
|
}
|
|
|
|
// One list entry per layer: kv_caches[layer_idx] is that layer's tensor.
|
|
LMC_KV_FORMAT_HD constexpr bool is_layer_list(EngineKVFormat f) {
|
|
return f == EngineKVFormat::NL_X_TWO_NB_BS_NH_HS ||
|
|
f == EngineKVFormat::NL_X_NB_TWO_BS_NH_HS ||
|
|
f == EngineKVFormat::NL_X_NB_BS_HS ||
|
|
f == EngineKVFormat::NL_X_NBBS_ONE_HS ||
|
|
f == EngineKVFormat::NL_X_TWO_NB_NH_BS_HS ||
|
|
f == EngineKVFormat::NL_X_NB_TWO_NH_BS_HS ||
|
|
f == EngineKVFormat::NL_X_NB_NH_BS_TWO_HS;
|
|
}
|
|
|
|
// Multi-head Latent Attention: a single latent KV head (no separate K/V split).
|
|
LMC_KV_FORMAT_HD constexpr bool is_mla(EngineKVFormat f) {
|
|
return f == EngineKVFormat::NL_X_NB_BS_HS || // vLLM MLA
|
|
f == EngineKVFormat::NL_X_NBBS_ONE_HS; // SGLang MLA
|
|
}
|