nano-store-kv

nano-vllm/nanovllm/layers/attention.py · bb823b3 · L11–40 · 固定 commit 上游

这是固定版本源码的只读摘录；可在线阅读，也可通过上游链接查看原文件。

11def store_kvcache_kernel(
12    key_ptr,
13    key_stride,
14    value_ptr,
15    value_stride,
16    k_cache_ptr,
17    v_cache_ptr,
18    slot_mapping_ptr,
19    D: tl.constexpr,
20):
21    idx = tl.program_id(0)
22    slot = tl.load(slot_mapping_ptr + idx)
23    if slot == -1: return
24    key_offsets = idx * key_stride + tl.arange(0, D)
25    value_offsets = idx * value_stride + tl.arange(0, D)
26    key = tl.load(key_ptr + key_offsets)
27    value = tl.load(value_ptr + value_offsets)
28    cache_offsets = slot * D + tl.arange(0, D)
29    tl.store(k_cache_ptr + cache_offsets, key)
30    tl.store(v_cache_ptr + cache_offsets, value)
31
32
33def store_kvcache(key: torch.Tensor, value: torch.Tensor, k_cache: torch.Tensor, v_cache: torch.Tensor, slot_mapping: torch.Tensor):
34    N, num_heads, head_dim = key.shape
35    D = num_heads * head_dim
36    assert key.stride(-1) == 1 and value.stride(-1) == 1
37    assert key.stride(1) == head_dim and value.stride(1) == head_dim
38    assert k_cache.stride(1) == D and v_cache.stride(1) == D
39    assert slot_mapping.numel() == N
40    store_kvcache_kernel[(N,)](key, key.stride(0), value, value.stride(0), k_cache, v_cache, slot_mapping, D)