A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 15, 16:00 UTC
[Neuron][kernel] Fuse kv cache into a single tensor (#15911)
Signed-off-by: Liangfu Chen <liangfc@amazon.com>
Liangfu Chen committed
Apr 3, 2025 at 09:51 UTC
d2b58ca203fcff18c66e93fc4a2f851090b8bf75
3 files changed
+46
-56
tests/neuron/1_core/test_cache.py
+3
-1
@@ -64,9 +64,11 @@ def test_reshape_and_cache(num_tokens, n_kv_head, d_head, num_blocks,
64
key_cache = torch.zeros_like(key_cache_cpu, device=device)
65
value_cache = torch.zeros_like(value_cache_cpu, device=device)
66
slot_mapping = slot_mapping_cpu.to(device)
67
+ kv_cache = torch.stack([key_cache, value_cache])
68
69
# Run vectorized implementation on XLA device
69
- reshape_and_cache(key, value, key_cache, value_cache, slot_mapping)
70
+ reshape_and_cache(key, value, kv_cache, slot_mapping)
71
+ key_cache, value_cache = torch.unbind(kv_cache, dim=0)
72
73
# Move results back to CPU for comparison
74
key_cache_result = key_cache.cpu()
tests/neuron/1_core/test_prefix_prefill.py
+5
-8
@@ -258,13 +258,13 @@ def sample_inputs(
258
value[start_loc:end_loc])
259
cur_ctx += block_size
260
block_id += 1
261
+ kv_cache = torch.stack([k_cache, v_cache])
262
263
return (
264
query,
265
k,
266
v,
266
- k_cache,
267
- v_cache,
267
+ kv_cache,
268
block_table,
269
key,
270
value,
@@ -361,8 +361,7 @@ def test_contexted_kv_attention(
361
query,
362
k_active,
363
v_active,
364
- k_cache,
365
- v_cache,
364
+ kv_cache,
365
block_table,
366
key,
367
value,
@@ -439,8 +438,7 @@ def test_contexted_kv_attention(
438
query = query.unsqueeze(0).permute(0, 2, 3, 1).contiguous()
439
k = k.unsqueeze(0).permute(0, 2, 3, 1).contiguous()
440
v = v.unsqueeze(0).permute(0, 2, 1, 3).contiguous()
442
- k_cache = k_cache.permute(0, 2, 1, 3).contiguous()
443
- v_cache = v_cache.permute(0, 2, 1, 3).contiguous()
441
+ kv_cache = kv_cache.permute(0, 1, 3, 2, 4).contiguous()
442
443
# transform block table
444
active_block_table = get_active_block_tables(
@@ -487,8 +485,7 @@ def test_contexted_kv_attention(
485
query.to(device=device),
486
k.to(device=device),
487
v.to(device=device),
490
- k_cache.to(device=device),
491
- v_cache.to(device=device),
488
+ kv_cache.to(device=device),
489
active_block_table.to(device=device),
490
attn_mask.to(device=device),
491
)
vllm/attention/ops/nki_flash_attn.py
+38
-47
@@ -144,8 +144,7 @@ def transform_block_tables_for_indirect_load(
144
def load_kv_tile_from_cache(
145
cur_k_tile,
146
cur_v_tile,
147
- key_cache,
148
- value_cache,
147
+ kv_cache,
148
block_tables,
149
large_k_tile_idx,
150
num_blocks_per_large_tile,
@@ -169,8 +168,8 @@ def load_kv_tile_from_cache(
168
for load_idx in nl.affine_range(num_loads):
169
i_p = nl.arange(B_P_SIZE)[:, None]
170
i_f = nl.arange(tiled_block_size * B_D_SIZE)[None, :]
172
- loaded = nl.load(key_cache[block_tables[load_idx, i_p,
173
- large_k_tile_idx], i_f])
171
+ loaded = nl.load(kv_cache[0, block_tables[load_idx, i_p,
172
+ large_k_tile_idx], i_f])
173
if cur_k_tile.dtype != loaded.dtype:
174
loaded = nl.copy(loaded, dtype=cur_k_tile.dtype)
175
# Transpose SBUF tensor using PE
@@ -185,7 +184,7 @@ def load_kv_tile_from_cache(
184
185
# load value cache
186
for load_idx in nl.affine_range(num_loads):
188
- loaded = nl.load(value_cache[block_tables[load_idx, i_p,
187
+ loaded = nl.load(kv_cache[1, block_tables[load_idx, i_p,
188
large_k_tile_idx], i_f])
189
if cur_v_tile.dtype != loaded.dtype:
190
loaded = nl.copy(loaded, dtype=cur_v_tile.dtype)
@@ -418,8 +417,7 @@ def flash_paged_attention(
417
query,
418
key,
419
value,
421
- key_cache,
422
- value_cache,
420
+ kv_cache,
421
block_tables,
422
mask,
423
softmax_scale=None,
@@ -434,8 +432,7 @@ def flash_paged_attention(
432
- query: shape (1, n_heads, d, seq_q)
433
- key: shape (1, n_kv_heads, d, seq_k)
434
- value: shape (1, n_kv_heads, seq_v, d)
437
- - key_cache: (num_blocks, n_kv_heads, block_size, d)
438
- - value_cache: (num_blocks, n_kv_heads, block_size, d)
435
+ - kv_cache: (2, num_blocks, n_kv_heads, block_size, d)
436
- block_tables: (num_active_blocks, )
437
- mask: (seq_q, num_active_blocks * block_size + seq_q)
438
- o: shape (1, n_heads, seq_q, d)
@@ -444,7 +441,7 @@ def flash_paged_attention(
441
- We use continuous batching by default, so the batch dimension is
442
always 1, and different requests are concatenated along sequence
443
dimension.
447
- - We use paged cache blocks (key_cache, value_cache) to store KV cache.
444
+ - We use paged cache blocks (kv_cache) to store KV cache.
445
446
IO tensor dtypes:
447
- This kernel assumes all IO tensors have the same dtype except for
@@ -475,15 +472,13 @@ def flash_paged_attention(
472
b, h, d, seqlen_q = query.shape
473
B_D_SIZE = d
474
n_tile_q = seqlen_q // B_P_SIZE # since q will be loaded on tensor engine
478
- num_blocks, k_h, block_size, _ = key_cache.shape
475
+ _, num_blocks, k_h, block_size, _ = kv_cache.shape
476
q_h_per_k_h = h // k_h
477
assert b == 1, f"invalid batch size {b=}"
478
assert d <= 128, f" we do not support head_dim > 128, got head dim {d=}"
482
- cache_shape = (num_blocks, k_h, block_size, d)
483
- assert (tuple(key_cache.shape) == cache_shape
484
- ), f"{key_cache.shape=} mismatch, expect {cache_shape}"
485
- assert (tuple(value_cache.shape) == cache_shape
486
- ), f"{value_cache.shape=} mismatch, expect {cache_shape}"
479
+ cache_shape = (2, num_blocks, k_h, block_size, d)
480
+ assert (tuple(kv_cache.shape) == cache_shape
481
+ ), f"{kv_cache.shape=} mismatch, expect {cache_shape}"
482
assert key is None or tuple(key.shape) == (
483
1,
484
k_h,
@@ -580,13 +575,13 @@ def flash_paged_attention(
575
head_id=head_id,
576
)
577
583
- # Flatten KV cache to be 2D for loading into SBUF
578
+ # Flatten KV cache to be 3D for loading into SBUF
579
new_cache_shape = (
580
+ 2,
581
num_blocks * k_h * block_size_tiling_factor,
582
tiled_block_size * d,
583
)
588
- key_cache = key_cache.reshape(new_cache_shape)
589
- value_cache = value_cache.reshape(new_cache_shape)
584
+ kv_cache = kv_cache.reshape(new_cache_shape)
585
586
# Global Flash Attention accumulators
587
o_buffer = nl.zeros(
@@ -621,8 +616,7 @@ def flash_paged_attention(
616
load_kv_tile_from_cache(
617
cur_k_tile=cur_k_tile,
618
cur_v_tile=cur_v_tile,
624
- key_cache=key_cache,
625
- value_cache=value_cache,
619
+ kv_cache=kv_cache,
620
block_tables=block_tables_sbuf,
621
large_k_tile_idx=large_k_tile_idx,
622
num_blocks_per_large_tile=num_blocks_per_large_tile,
@@ -821,8 +815,7 @@ def flash_attn_varlen_nkifunc(
815
query,
816
key,
817
value,
824
- key_cache,
825
- value_cache,
818
+ kv_cache,
819
block_table,
820
attn_mask,
821
n_kv_head=None,
@@ -838,8 +831,7 @@ def flash_attn_varlen_nkifunc(
831
- query: (1, n_heads, d, seq_q)
832
- key: (1, n_kv_heads, d, seq_k)
833
- value: (1, n_kv_heads, seq_v, d)
841
- - key_cache: (n_blocks, n_kv_heads, block_size, d)
842
- - value_cache: (n_blocks, n_kv_heads, block_size, d)
834
+ - kv_cache: (2, n_blocks, n_kv_heads, block_size, d)
835
- block_tables: (n_active_blocks, )
836
- attn_mask: (seq_q, n_active_blocks * block_size + seq_q)
837
@@ -849,17 +841,17 @@ def flash_attn_varlen_nkifunc(
841
for better DMA throughput
842
"""
843
if n_kv_head is None:
852
- n_kv_head = key_cache.shape[1]
853
- assert key_cache.shape[1] == n_kv_head
844
+ n_kv_head = kv_cache.shape[2]
845
+ assert kv_cache.shape[0] == 2
846
+ assert kv_cache.shape[2] == n_kv_head
847
if head_size is None:
855
- head_size = key_cache.shape[-1]
848
+ head_size = kv_cache.shape[-1]
849
850
kwargs = dict(
851
query=query,
852
key=key,
853
value=value,
861
- key_cache=key_cache,
862
- value_cache=value_cache,
854
+ kv_cache=kv_cache,
855
block_tables=block_table,
856
mask=attn_mask,
857
softmax_scale=1.0 / (head_size**0.5),
@@ -874,8 +866,7 @@ def flash_attn_varlen_nkifunc(
866
def reshape_and_cache(
867
key: torch.Tensor,
868
value: torch.Tensor,
877
- key_cache: torch.Tensor,
878
- value_cache: torch.Tensor,
869
+ kv_cache: torch.Tensor,
870
slot_mapping: torch.Tensor,
871
) -> None:
872
"""
@@ -886,29 +877,29 @@ def reshape_and_cache(
877
(num_tokens, n_kv_head, d_head)
878
value (torch.Tensor): Value tensor with shape
879
(num_tokens, n_kv_head, d_head)
889
- key_cache (torch.Tensor): Key cache tensor with shape
890
- (num_blocks, n_kv_head, block_size, d_head)
891
- value_cache (torch.Tensor): Value cache tensor with shape
892
- (num_blocks, n_kv_head, block_size, d_head)
880
+ kv_cache (torch.Tensor): Key/value cache tensor with shape
881
+ (2, num_blocks, n_kv_head, block_size, d_head)
882
slot_mapping (torch.Tensor): Mapping tensor indicating cache positions
883
with shape (num_tokens)
884
885
Returns:
897
- None: Updates the key_cache and value_cache tensors in-place
886
+ None: Updates the kv_cache tensor in-place
887
"""
899
- block_size = key_cache.size(2)
888
+ block_size = kv_cache.size(3)
889
+ n_kv_head = key.size(1)
890
891
# Calculate indices with explicit floor division
892
block_indices = torch.div(slot_mapping, block_size, rounding_mode="floor")
893
block_offsets = slot_mapping % block_size
894
895
+ # Create the head indices tensor
896
+ head_indices = torch.arange(n_kv_head, device=key.device)
897
+
898
# Update caches using index_put_
906
- key_cache.index_put_(
907
- (block_indices.unsqueeze(1),
908
- torch.arange(key_cache.size(1),
909
- device=key.device), block_offsets.unsqueeze(1)), key)
910
-
911
- value_cache.index_put_(
912
- (block_indices.unsqueeze(1),
913
- torch.arange(value_cache.size(1),
914
- device=value.device), block_offsets.unsqueeze(1)), value)
899
+ kv_cache.index_put_(
900
+ (torch.tensor([0], device=key.device), block_indices[:, None],
901
+ head_indices[None, :], block_offsets[:, None]), key)
902
+
903
+ kv_cache.index_put_(
904
+ (torch.tensor([1], device=key.device), block_indices[:, None],
905
+ head_indices[None, :], block_offsets[:, None]), value)