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, 22:15 UTC
[V1][Kernel] Refactor the prefix_prefill kernel so that the caller no longer has to pass in the context lengths (#13095)
Sage Moore committed
Feb 22, 2025 at 05:25 UTC
558db8083cfd9b7ee76eafdae32b237d951c8b10
6 files changed
+12
-31
tests/kernels/test_prefix_prefill.py
+2
-6
@@ -100,7 +100,7 @@ def test_contexted_kv_attention(
100
BS, max_block_per_request)
101
b_seq_len = torch.tensor(seq_lens, dtype=torch.long)
102
b_ctx_len = torch.tensor(ctx_lens, dtype=torch.long)
103
- b_start_loc = torch.cumsum(torch.tensor([0] + query_lens[:-1],
103
+ b_start_loc = torch.cumsum(torch.tensor([0] + query_lens,
104
dtype=torch.long),
105
dim=0)
106
max_input_len = MAX_SEQ_LEN
@@ -154,7 +154,6 @@ def test_contexted_kv_attention(
154
block_table,
155
b_start_loc,
156
b_seq_len,
157
- b_ctx_len,
157
max_input_len,
158
k_scale,
159
v_scale,
@@ -171,7 +170,6 @@ def test_contexted_kv_attention(
170
block_table,
171
b_start_loc,
172
b_seq_len,
174
- b_ctx_len,
173
max_input_len,
174
k_scale,
175
v_scale,
@@ -333,7 +331,7 @@ def test_contexted_kv_attention_alibi(
331
BS, max_block_per_request)
332
b_seq_len = torch.tensor(seq_lens, dtype=torch.long)
333
b_ctx_len = torch.tensor(ctx_lens, dtype=torch.long)
336
- b_start_loc = torch.cumsum(torch.tensor([0] + query_lens[:-1],
334
+ b_start_loc = torch.cumsum(torch.tensor([0] + query_lens,
335
dtype=torch.long),
336
dim=0)
337
max_input_len = MAX_SEQ_LEN
@@ -387,7 +385,6 @@ def test_contexted_kv_attention_alibi(
385
block_table,
386
b_start_loc,
387
b_seq_len,
390
- b_ctx_len,
388
max_input_len,
389
k_scale,
390
v_scale,
@@ -404,7 +401,6 @@ def test_contexted_kv_attention_alibi(
401
block_table,
402
b_start_loc,
403
b_seq_len,
407
- b_ctx_len,
404
max_input_len,
405
k_scale,
406
v_scale,
vllm/attention/backends/rocm_flash_attn.py
-1
@@ -753,7 +753,6 @@ class ROCmFlashAttentionImpl(AttentionImpl):
753
prefill_meta.block_tables,
754
prefill_meta.query_start_loc,
755
prefill_meta.seq_lens_tensor,
756
- prefill_meta.context_lens_tensor,
756
prefill_meta.max_query_len,
757
self.alibi_slopes,
758
self.sliding_window[0],
vllm/attention/backends/xformers.py
-1
@@ -580,7 +580,6 @@ class XFormersImpl(AttentionImpl[XFormersMetadata]):
580
prefill_meta.block_tables,
581
prefill_meta.query_start_loc,
582
prefill_meta.seq_lens_tensor,
583
- prefill_meta.context_lens_tensor,
583
prefill_meta.max_query_len,
584
self.alibi_slopes,
585
self.sliding_window,
vllm/attention/ops/paged_attn.py
+1
-3
@@ -202,7 +202,6 @@ class PagedAttention:
202
block_tables: torch.Tensor,
203
query_start_loc: torch.Tensor,
204
seq_lens_tensor: torch.Tensor,
205
- context_lens: torch.Tensor,
205
max_query_len: int,
206
alibi_slopes: Optional[torch.Tensor],
207
sliding_window: Optional[int],
@@ -220,9 +219,8 @@ class PagedAttention:
219
value_cache,
220
block_tables,
221
# query_start_loc is (batch_size + 1,)
223
- query_start_loc[:-1],
222
+ query_start_loc,
223
seq_lens_tensor,
225
- context_lens,
224
max_query_len,
225
k_scale,
226
v_scale,
vllm/attention/ops/prefix_prefill.py
+9
-8
@@ -31,7 +31,6 @@ if triton.__version__ >= "2.1.0":
31
v_scale,
32
B_Start_Loc,
33
B_Seqlen,
34
- B_Ctxlen,
34
block_size,
35
x,
36
Out,
@@ -72,10 +71,12 @@ if triton.__version__ >= "2.1.0":
71
72
cur_kv_head = cur_head // num_queries_per_kv
73
75
- cur_batch_ctx_len = tl.load(B_Ctxlen + cur_batch)
74
cur_batch_seq_len = tl.load(B_Seqlen + cur_batch)
75
cur_batch_in_all_start_index = tl.load(B_Start_Loc + cur_batch)
78
- cur_batch_query_len = cur_batch_seq_len - cur_batch_ctx_len
76
+ cur_batch_in_all_stop_index = tl.load(B_Start_Loc + cur_batch + 1)
77
+ cur_batch_query_len = (cur_batch_in_all_stop_index -
78
+ cur_batch_in_all_start_index)
79
+ cur_batch_ctx_len = cur_batch_seq_len - cur_batch_query_len
80
81
# start position inside of the query
82
# generally, N goes over kv, while M goes over query_len
@@ -466,7 +467,6 @@ if triton.__version__ >= "2.1.0":
467
v_scale,
468
B_Start_Loc,
469
B_Seqlen,
469
- B_Ctxlen,
470
Alibi_slopes,
471
block_size,
472
x,
@@ -511,9 +511,12 @@ if triton.__version__ >= "2.1.0":
511
# cur_batch_seq_len: the length of prompts
512
# cur_batch_ctx_len: the length of prefix
513
# cur_batch_in_all_start_index: the start id of the dim=0
514
- cur_batch_ctx_len = tl.load(B_Ctxlen + cur_batch)
514
cur_batch_seq_len = tl.load(B_Seqlen + cur_batch)
515
cur_batch_in_all_start_index = tl.load(B_Start_Loc + cur_batch)
516
+ cur_batch_in_all_stop_index = tl.load(B_Start_Loc + cur_batch + 1)
517
+ cur_batch_query_len = (cur_batch_in_all_stop_index -
518
+ cur_batch_in_all_start_index)
519
+ cur_batch_ctx_len = cur_batch_seq_len - cur_batch_query_len
520
521
block_start_loc = BLOCK_M * start_m
522
@@ -713,7 +716,6 @@ if triton.__version__ >= "2.1.0":
716
b_loc,
717
b_start_loc,
718
b_seq_len,
716
- b_ctx_len,
719
max_input_len,
720
k_scale: torch.Tensor,
721
v_scale: torch.Tensor,
@@ -765,6 +767,7 @@ if triton.__version__ >= "2.1.0":
767
batch, head = b_seq_len.shape[0], q.shape[1]
768
num_queries_per_kv = q.shape[1] // k.shape[1]
769
770
+ assert batch + 1 == len(b_start_loc)
771
grid = (batch, head, triton.cdiv(max_input_len, BLOCK)) # batch, head,
772
773
# 0 means "disable"
@@ -784,7 +787,6 @@ if triton.__version__ >= "2.1.0":
787
v_scale,
788
b_start_loc,
789
b_seq_len,
787
- b_ctx_len,
790
alibi_slopes,
791
v_cache.shape[3],
792
k_cache.shape[4],
@@ -838,7 +840,6 @@ if triton.__version__ >= "2.1.0":
840
v_scale,
841
b_start_loc,
842
b_seq_len,
841
- b_ctx_len,
843
v_cache.shape[3],
844
k_cache.shape[4],
845
o,
vllm/v1/attention/backends/rocm_attn.py
-12
@@ -150,17 +150,6 @@ class ROCmAttentionImpl(AttentionImpl):
150
layer._v_scale,
151
)
152
153
- # TODO(sage): Refactor the context_attention_fwd kernel so that this
154
- # overhead can be removed
155
- context_lens = torch.empty_like(attn_metadata.seq_lens)
156
- batch_size = len(attn_metadata.query_start_loc) - 1
157
- assert len(context_lens) == batch_size
158
- for i in range(batch_size):
159
- query_start = attn_metadata.query_start_loc[i]
160
- query_end = attn_metadata.query_start_loc[i + 1]
161
- context_lens[i] = attn_metadata.seq_lens[i] - (query_end -
162
- query_start)
163
-
153
# Compute attention and update output up to `num_actual_tokens`.
154
context_attention_fwd(q=query[:num_actual_tokens],
155
k=key[:num_actual_tokens],
@@ -172,7 +161,6 @@ class ROCmAttentionImpl(AttentionImpl):
161
b_loc=attn_metadata.block_table,
162
b_start_loc=attn_metadata.query_start_loc,
163
b_seq_len=attn_metadata.seq_lens,
175
- b_ctx_len=context_lens,
164
max_input_len=attn_metadata.max_query_len,
165
k_scale=layer._k_scale,
166
v_scale=layer._v_scale,