@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 33170081f1
Fork 0

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

[Neuron][Kernel] Vectorize KV cache load in FlashPagedAttention to maximize DMA bandwidth (#13245)

Signed-off-by: Lingfan Yu <lingfany@amazon.com>

Lingfan Yu committed Feb 20, 2025 at 17:45 UTC 33170081f13b4e4453f0fff8df994df8d9aa6c6c
3 files changed +764 -348
tests/neuron/test_block_table.py new
+153
@@ -0,0 +1,153 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +import os
3 +
4 +import neuronxcc.nki.language as nl
5 +import pytest
6 +import torch
7 +import torch.nn.functional as F
8 +from neuronxcc import nki
9 +
10 +from vllm.attention.ops.nki_flash_attn import (
11 + load_block_tables, transform_block_tables_for_indirect_load)
12 +
13 +
14 +def is_power_of_2(n):
15 + return n > 0 and (n & (n - 1) == 0)
16 +
17 +
18 +def nki_load_and_transform_block_tables(
19 + block_tables,
20 + num_tiles,
21 + num_blocks_per_tile,
22 + num_head,
23 + head_id,
24 + block_size_tiling_factor,
25 +):
26 + assert is_power_of_2(
27 + num_blocks_per_tile), f"{num_blocks_per_tile=} must be power of 2"
28 + block_tables_sbuf = load_block_tables(block_tables, num_tiles,
29 + num_blocks_per_tile)
30 +
31 + # we need to pass an Index as head_id
32 + head_id = nl.arange(1)[None, :] + head_id
33 +
34 + block_tables_transposed = transform_block_tables_for_indirect_load(
35 + block_tables_sbuf, block_size_tiling_factor, num_head, head_id)
36 + B_P_SIZE = 128
37 + assert block_tables_transposed.shape[1] == B_P_SIZE
38 +
39 + out = nl.ndarray(
40 + block_tables_transposed.shape,
41 + dtype=nl.int32,
42 + buffer=nl.shared_hbm,
43 + )
44 + for i in nl.affine_range(block_tables_transposed.shape[0]):
45 + nl.store(dst=out[i], value=block_tables_transposed[i])
46 + return out
47 +
48 +
49 +def ref_block_tables_transform(
50 + block_tables,
51 + num_tiles,
52 + num_blocks_per_tile,
53 + num_head,
54 + head_id,
55 + block_size_tiling_factor,
56 +):
57 + assert block_tables.numel() == num_tiles * num_blocks_per_tile
58 + block_tables = block_tables.view(num_tiles, num_blocks_per_tile)
59 + B_F_SIZE = 128
60 + num_tiles_padded = (num_tiles + B_F_SIZE - 1) // B_F_SIZE * B_F_SIZE
61 + block_tables = F.pad(
62 + block_tables,
63 + (0, 0, 0, num_tiles_padded - num_tiles),
64 + "constant",
65 + 0,
66 + )
67 +
68 + block_tables = block_tables * num_head + head_id
69 + block_tables = block_tables.view(num_tiles_padded, num_blocks_per_tile, 1)
70 + offset = torch.arange(0, block_size_tiling_factor).view(1, 1, -1)
71 + block_tables = block_tables * block_size_tiling_factor + offset
72 + block_tables_transposed = block_tables.view(num_tiles_padded, -1).t()
73 +
74 + num_blocks_per_tile = block_tables_transposed.shape[0]
75 + assert num_blocks_per_tile % B_F_SIZE == 0
76 + return block_tables_transposed.view(num_blocks_per_tile // B_F_SIZE,
77 + B_F_SIZE, num_tiles_padded)
78 +
79 +
80 +@pytest.mark.parametrize(
81 + "q_head_per_kv_head,head_id",
82 + [
83 + (1, 0),
84 + (3, 1),
85 + ],
86 +)
87 +@pytest.mark.parametrize(
88 + "num_tiles,num_blocks_per_tile",
89 + [
90 + (1, 1),
91 + (13, 16),
92 + (17, 128),
93 + (35, 512),
94 + (128, 128),
95 + (130, 64),
96 + (280, 256),
97 + (315, 1),
98 + ],
99 +)
100 +@torch.inference_mode()
101 +def test_load_and_transform_block_tables(
102 + num_tiles,
103 + num_blocks_per_tile,
104 + q_head_per_kv_head,
105 + head_id,
106 +) -> None:
107 + import torch_xla.core.xla_model as xm
108 +
109 + device = xm.xla_device()
110 +
111 + compiler_flags = [
112 + "-O1",
113 + "--retry_failed_compilation",
114 + ]
115 + compiler_flags_str = " ".join(compiler_flags)
116 + os.environ["NEURON_CC_FLAGS"] = compiler_flags_str
117 +
118 + torch.manual_seed(10000)
119 + torch.set_printoptions(sci_mode=False)
120 +
121 + # On Neuron, we need B_P_SIZE = 128 blocks to make DMA efficient
122 + B_P_SIZE = 128
123 + if num_blocks_per_tile < B_P_SIZE:
124 + assert B_P_SIZE % num_blocks_per_tile == 0
125 + block_size_tiling_factor = B_P_SIZE // num_blocks_per_tile
126 + else:
127 + block_size_tiling_factor = 1
128 + max_num_blocks = 100000
129 + block_tables = torch.randint(
130 + 0,
131 + max_num_blocks,
132 + (num_tiles * num_blocks_per_tile, ),
133 + dtype=torch.int32,
134 + )
135 + nki_out = nki.jit(nki_load_and_transform_block_tables)[1, 1](
136 + block_tables.to(device=device),
137 + num_tiles,
138 + num_blocks_per_tile,
139 + q_head_per_kv_head,
140 + head_id,
141 + block_size_tiling_factor,
142 + ).cpu()
143 + ref_out = ref_block_tables_transform(
144 + block_tables,
145 + num_tiles,
146 + num_blocks_per_tile,
147 + q_head_per_kv_head,
148 + head_id,
149 + block_size_tiling_factor,
150 + )
151 + assert (nki_out.shape == ref_out.shape
152 + ), f"{nki_out.shape=} != {ref_out.shape=}"
153 + assert torch.all(nki_out == ref_out)
tests/neuron/test_prefix_prefill.py
+183 -149
@@ -107,7 +107,7 @@ def ref_masked_attention(
107 masked_score, dim=-1, return_max_reduce=True)
108 else:
109 norm_score = ref_softmax(masked_score, dim=-1)
110 - out = torch.einsum("hqk,khd->qhd", norm_score, value)
110 + out = torch.einsum("hqk,khd->qhd", norm_score.to(value.dtype), value)
111 if return_max_reduce:
112 return (
113 out,
@@ -118,7 +118,7 @@ def ref_masked_attention(
118 scaled_qk,
119 )
120 else:
121 - return out
121 + return (out, )
122
123
124 def ref_context_attention(
@@ -128,8 +128,6 @@ def ref_context_attention(
128 query_lens,
129 seq_lens,
130 head_size,
131 - num_kv_heads,
132 - num_heads,
131 num_queries_per_kv,
132 return_max_reduce=False,
133 ):
@@ -146,18 +144,19 @@ def ref_context_attention(
144 attn_mask = torch.logical_not(attn_mask)
145 attn_mask = attn_mask.float() * -30000
146
149 - output, cached_max, cached_sum_reciprocal, lse, masked_score, scaled_qk = (
150 - ref_masked_attention(
151 - query,
152 - key,
153 - value,
154 - scale,
155 - attn_mask,
156 - return_max_reduce=return_max_reduce,
157 - ))
147 + output, *debug_tensors = ref_masked_attention(
148 + query,
149 + key,
150 + value,
151 + scale,
152 + attn_mask,
153 + return_max_reduce=return_max_reduce,
154 + )
155
156 output = output.unsqueeze(1)
157 if return_max_reduce:
158 + cached_max, cached_sum_reciprocal, lse, masked_score, scaled_qk = (
159 + debug_tensors)
160 return (
161 output,
162 cached_max,
@@ -170,65 +169,22 @@ def ref_context_attention(
169 return output
170
171
173 -@pytest.mark.parametrize(
174 - "block_size, large_tile_size",
175 - [
176 - (32, 2048), # 64 blocks
177 - (32, 4096), # 128 blocks
178 - (32, 8192), # 256 blocks
179 - (64, 8192), # 128 blocks
180 - ],
181 -)
182 -@pytest.mark.parametrize(
183 - "num_heads,num_queries_per_kv,head_size,mixed_precision",
184 - [
185 - (4, 2, 8, False),
186 - (4, 2, 8, True),
187 - (32, 8, 64, True),
188 - (16, 2, 128, True),
189 - ],
190 -)
191 -@torch.inference_mode()
192 -def test_contexted_kv_attention(
193 - num_heads: int,
194 - num_queries_per_kv: int,
195 - head_size: int,
196 - block_size: int,
197 - large_tile_size,
198 - mixed_precision: bool,
199 -) -> None:
200 - import os
201 -
202 - import torch_xla.core.xla_model as xm
203 -
204 - from vllm.attention.ops.nki_flash_attn import flash_attn_varlen_nkifunc
205 -
206 - assert large_tile_size % block_size == 0
207 -
208 - device = xm.xla_device()
209 -
210 - compiler_flags = [
211 - "--model-type=transformer -O1",
212 - "--internal-hlo2tensorizer-options='--verify-hlo'",
213 - "--retry_failed_compilation",
214 - ]
215 - compiler_flags_str = " ".join(compiler_flags)
216 - os.environ["NEURON_CC_FLAGS"] = compiler_flags_str
217 -
218 - torch.manual_seed(0)
219 - torch.set_printoptions(sci_mode=False)
220 -
221 - min_ctx_len = 32
222 - max_ctx_len = 1024
223 - min_query_len = 16
224 - max_query_len = 512
225 - prefill_batch_size = 4
226 - decode_batch_size = 12
172 +def sample_inputs(
173 + prefill_batch_size,
174 + decode_batch_size,
175 + min_query_len,
176 + max_query_len,
177 + min_ctx_len,
178 + max_ctx_len,
179 + block_size,
180 + num_heads,
181 + num_kv_heads,
182 + head_size,
183 + dtype,
184 +):
185 batch_size = prefill_batch_size + decode_batch_size
186 max_model_len = (max_query_len + max_ctx_len) * 4
229 -
187 max_block_per_request = max_model_len // block_size
231 - dtype = torch.float32
188 cache_size = (batch_size * max_block_per_request) + 2
189 prefill_ctx_lens = torch.randint(min_ctx_len,
190 max_ctx_len + 1, (prefill_batch_size, ),
@@ -244,7 +200,6 @@ def test_contexted_kv_attention(
200 dtype=torch.long,
201 ).tolist() + [1 for _ in range(decode_batch_size)]
202 seq_lens = [a + b for a, b in zip(query_lens, ctx_lens)]
247 - num_kv_heads = num_heads // num_queries_per_kv
203
204 num_tokens = sum(query_lens)
205 query = torch.empty(num_tokens, num_heads, head_size, dtype=dtype)
@@ -304,47 +259,139 @@ def test_contexted_kv_attention(
259 cur_ctx += block_size
260 block_id += 1
261
262 + return (
263 + query,
264 + k,
265 + v,
266 + k_cache,
267 + v_cache,
268 + block_table,
269 + key,
270 + value,
271 + query_lens,
272 + seq_lens,
273 + )
274 +
275 +
276 +def get_active_block_tables(block_tables, query_lens, seq_lens, block_size,
277 + num_blocks):
278 + context_lens = seq_lens - query_lens
279 + blocks_per_seq = (context_lens + block_size - 1) // block_size
280 + num_seqs = len(seq_lens)
281 + active_blocks: list[int] = []
282 + for seq_id in range(num_seqs):
283 + active_blocks = (
284 + active_blocks +
285 + block_tables[seq_id, :blocks_per_seq[seq_id]].tolist())
286 + return F.pad(
287 + torch.tensor(active_blocks, dtype=torch.int32),
288 + (0, num_blocks - len(active_blocks)),
289 + "constant",
290 + 0,
291 + )
292 +
293 +
294 +@pytest.mark.parametrize(
295 + "prefill_batch_size,decode_batch_size,block_size,large_tile_size",
296 + [
297 + (1, 199, 1, 512), # 512 blocks
298 + (4, 12, 256, 2048), # 128 blocks
299 + (4, 12, 16, 2048), # 128 blocks
300 + (4, 12, 4, 1024), # 256 blocks
301 + (4, 12, 32, 2048), # 64 blocks
302 + (4, 12, 32, 4096), # 128 blocks
303 + (4, 12, 32, 8192), # 256 blocks
304 + (4, 12, 64, 8192), # 128 blocks
305 + ],
306 +)
307 +@pytest.mark.parametrize(
308 + "num_heads,num_queries_per_kv,head_size",
309 + [
310 + (4, 2, 8),
311 + (32, 8, 64),
312 + (4, 4, 128),
313 + (8, 1, 32),
314 + ],
315 +)
316 +@pytest.mark.parametrize("mixed_precision", [True, False])
317 +@torch.inference_mode()
318 +def test_contexted_kv_attention(
319 + prefill_batch_size: int,
320 + decode_batch_size: int,
321 + num_heads: int,
322 + num_queries_per_kv: int,
323 + head_size: int,
324 + block_size: int,
325 + large_tile_size,
326 + mixed_precision: bool,
327 +) -> None:
328 + import os
329 +
330 + import torch_xla.core.xla_model as xm
331 +
332 + from vllm.attention.ops.nki_flash_attn import (flash_attn_varlen_nkifunc,
333 + reorder_context_mask)
334 +
335 + assert large_tile_size % block_size == 0
336 +
337 + device = xm.xla_device()
338 +
339 + compiler_flags = [
340 + "-O1",
341 + "--retry_failed_compilation",
342 + ]
343 + compiler_flags_str = " ".join(compiler_flags)
344 + os.environ["NEURON_CC_FLAGS"] = compiler_flags_str
345 +
346 + torch.manual_seed(0)
347 + torch.set_printoptions(sci_mode=False)
348 + dtype = torch.float32
349 +
350 + min_ctx_len = 32
351 + max_ctx_len = 1024
352 + min_query_len = 16
353 + max_query_len = 512
354 + num_kv_heads = num_heads // num_queries_per_kv
355 (
308 - output_ref,
309 - cached_max,
310 - cached_sum_reciprocal,
311 - lse,
312 - masked_score,
313 - scaled_qk,
314 - ) = ref_context_attention(
356 + query,
357 + k_active,
358 + v_active,
359 + k_cache,
360 + v_cache,
361 + block_table,
362 + key,
363 + value,
364 + query_lens,
365 + seq_lens,
366 + ) = sample_inputs(
367 + prefill_batch_size=prefill_batch_size,
368 + decode_batch_size=decode_batch_size,
369 + min_query_len=min_query_len,
370 + max_query_len=max_query_len,
371 + min_ctx_len=min_ctx_len,
372 + max_ctx_len=max_ctx_len,
373 + block_size=block_size,
374 + num_heads=num_heads,
375 + num_kv_heads=num_kv_heads,
376 + head_size=head_size,
377 + dtype=dtype,
378 + )
379 +
380 + output_ref = ref_context_attention(
381 query,
382 key,
383 value,
384 query_lens,
385 seq_lens,
386 head_size,
321 - num_kv_heads,
322 - num_heads,
387 num_queries_per_kv,
324 - return_max_reduce=True,
388 + return_max_reduce=False,
389 )
390
391 # build neuron program
328 - return_debug_tensors = False
392 B_P_SIZE = 128
330 - LARGE_TILE_SZ = large_tile_size
331 -
332 - def get_active_block_tables(block_tables, query_lens, seq_lens, block_size,
333 - num_blocks):
334 - context_lens = seq_lens - query_lens
335 - blocks_per_seq = (context_lens + block_size - 1) // block_size
336 - num_seqs = len(seq_lens)
337 - active_blocks: list[int] = []
338 - for seq_id in range(num_seqs):
339 - active_blocks = (
340 - active_blocks +
341 - block_tables[seq_id, :blocks_per_seq[seq_id]].tolist())
342 - return F.pad(
343 - torch.tensor(active_blocks),
344 - (0, num_blocks - len(active_blocks)),
345 - "constant",
346 - 0,
347 - )
393 + assert (large_tile_size >= B_P_SIZE
394 + ), f"Expect {large_tile_size=} to be larger than {B_P_SIZE=}"
395
396 def ceil_div(a, b):
397 return (a + b - 1) // b
@@ -357,32 +404,27 @@ def test_contexted_kv_attention(
404 return 2**int(a - 1).bit_length()
405
406 # calculate input shapes
360 - max_num_queries = pad_to_multiple(sum(query_lens), block_size)
361 - max_num_queries = pad_to_next_power_of_2(max_num_queries)
362 - head_size_padded = B_P_SIZE
363 - assert head_size_padded >= head_size
407 + max_num_queries = pad_to_next_power_of_2(sum(query_lens))
408 context_lens = torch.tensor(seq_lens) - torch.tensor(query_lens)
409 num_active_blocks = ceil_div(context_lens, block_size).sum().item()
410 num_active_blocks = pad_to_multiple(num_active_blocks,
367 - LARGE_TILE_SZ // block_size)
411 + large_tile_size // block_size)
412 context_kv_len = num_active_blocks * block_size
413 assert (context_kv_len %
370 - LARGE_TILE_SZ == 0), f"invalid context_kv_len={context_kv_len}"
414 + large_tile_size == 0), f"invalid context_kv_len={context_kv_len}"
415
416 # pad QKV tensors
417 pad_dims = (
418 0,
375 - head_size_padded - query.shape[2],
419 + 0,
420 0,
421 0,
422 0,
423 max_num_queries - query.shape[0],
424 )
425 query = F.pad(query, pad_dims, "constant", 0)
382 - k = F.pad(k, pad_dims, "constant", 0)
383 - v = F.pad(v, pad_dims, "constant", 0)
384 - k_cache = F.pad(k_cache, (0, head_size_padded - head_size), "constant", 0)
385 - v_cache = F.pad(v_cache, (0, head_size_padded - head_size), "constant", 0)
426 + k = F.pad(k_active, pad_dims, "constant", 0)
427 + v = F.pad(v_active, pad_dims, "constant", 0)
428
429 # permute QKV tensors
430 # query: (1, n_heads, d, seq_q)
@@ -391,6 +433,8 @@ def test_contexted_kv_attention(
433 query = query.unsqueeze(0).permute(0, 2, 3, 1).contiguous()
434 k = k.unsqueeze(0).permute(0, 2, 3, 1).contiguous()
435 v = v.unsqueeze(0).permute(0, 2, 1, 3).contiguous()
436 + k_cache = k_cache.permute(0, 2, 1, 3).contiguous()
437 + v_cache = v_cache.permute(0, 2, 1, 3).contiguous()
438
439 # transform block table
440 active_block_table = get_active_block_tables(
@@ -405,33 +449,31 @@ def test_contexted_kv_attention(
449 prior_mask, active_mask = (
450 BlockDiagonalCausalFromBottomRightMask.from_seqlens(
451 query_lens, seq_lens, block_size=block_size))
408 - attn_mask = torch.concat(
409 - [
410 - F.pad(
411 - prior_mask,
412 - (
413 - 0,
414 - context_kv_len - prior_mask.shape[1],
415 - 0,
416 - max_num_queries - prior_mask.shape[0],
417 - ),
418 - "constant",
419 - 0,
420 - ).bool(),
421 - F.pad(
422 - active_mask,
423 - (
424 - 0,
425 - max_num_queries - active_mask.shape[1],
426 - 0,
427 - max_num_queries - active_mask.shape[0],
428 - ),
429 - "constant",
430 - 0,
431 - ).bool(),
432 - ],
433 - dim=1,
434 - )
452 + prior_mask_padded = F.pad(
453 + prior_mask,
454 + (
455 + 0,
456 + context_kv_len - prior_mask.shape[1],
457 + 0,
458 + max_num_queries - prior_mask.shape[0],
459 + ),
460 + "constant",
461 + 0,
462 + ).bool()
463 + active_mask_padded = F.pad(
464 + active_mask,
465 + (
466 + 0,
467 + max_num_queries - active_mask.shape[1],
468 + 0,
469 + max_num_queries - active_mask.shape[0],
470 + ),
471 + "constant",
472 + 0,
473 + ).bool()
474 + attn_mask = torch.concat([prior_mask_padded, active_mask_padded], dim=1)
475 +
476 + attn_mask = reorder_context_mask(attn_mask, large_tile_size, block_size)
477
478 input_args = (
479 query.to(device=device),
@@ -439,29 +481,21 @@ def test_contexted_kv_attention(
481 v.to(device=device),
482 k_cache.to(device=device),
483 v_cache.to(device=device),
442 - active_block_table.to(torch.int32).to(device=device),
484 + active_block_table.to(device=device),
485 attn_mask.to(device=device),
486 )
487 input_kwargs = dict(
488 n_kv_head=num_kv_heads,
489 head_size=head_size,
490 mixed_precision=mixed_precision,
449 - LARGE_TILE_SZ=LARGE_TILE_SZ,
450 - return_debug_tensors=return_debug_tensors,
491 + LARGE_TILE_SZ=large_tile_size,
492 )
493
453 - if return_debug_tensors:
454 - output_nki, *debug_tensors = flash_attn_varlen_nkifunc(
455 - *input_args, **input_kwargs)
456 - else:
457 - output_nki = flash_attn_varlen_nkifunc(*input_args, **input_kwargs)
458 - debug_tensors = []
459 -
460 - debug_tensors = [torch.tensor(dt).cpu() for dt in debug_tensors]
494 + output_nki = flash_attn_varlen_nkifunc(*input_args, **input_kwargs)
495
496 num_actual_tokens = sum(query_lens)
497 # - o: shape (bs, n_heads, seq_q, d) -> (bs, seq_q, n_heads, d)
464 - output_nki = output_nki.cpu().permute(0, 2, 1, 3)[:, :, :, :head_size]
498 + output_nki = output_nki.cpu().permute(0, 2, 1, 3)
499 output_nki = output_nki[0, :num_actual_tokens, :, :]
500 output_ref_padded = F.pad(
501 output_ref,
vllm/attention/ops/nki_flash_attn.py
+428 -199
@@ -1,27 +1,203 @@
1 # SPDX-License-Identifier: Apache-2.0
2
3 -from dataclasses import dataclass
4 -
3 import neuronxcc.nki.isa as nisa
4 import neuronxcc.nki.language as nl
5 import numpy as np
6 +import torch
7 from neuronxcc import nki
8 from neuronxcc.nki.language import par_dim
9
10
12 -@dataclass(frozen=True)
13 -class FlashConfig:
11 +def ceil_div(a, b):
12 + return (a + b - 1) // b
13 +
14 +
15 +def is_power_of_2(x):
16 + return x > 0 and (x & (x - 1)) == 0
17 +
18 +
19 +@nki.jit
20 +def load_block_tables(block_tables_hbm, num_tiles, num_blocks_per_tile):
21 + """
22 + Load block tables from HBM into SRAM
23 +
24 + `block_tables_hbm` has shape `(num_tiles * num_blocks_per_tile, )`.
25 + In case `num_tiles > B_P_SIZE`, we need further tile `num_tile` dimension.
26 + """
27 + B_P_SIZE = 128
28 +
29 + # reshape as `(num_tiles, num_blocks_per_tile)`
30 + assert len(block_tables_hbm.shape) == 1
31 + (num_total_blocks, ) = block_tables_hbm.shape
32 + assert num_blocks_per_tile * num_tiles == num_total_blocks
33 + block_tables_hbm = block_tables_hbm.reshape(
34 + (num_tiles, num_blocks_per_tile))
35 +
36 + block_tables_sbuf = nl.zeros(
37 + (ceil_div(num_tiles,
38 + B_P_SIZE), par_dim(B_P_SIZE), num_blocks_per_tile),
39 + dtype=nl.int32,
40 + )
41 + for i in nl.affine_range(ceil_div(num_tiles, B_P_SIZE)):
42 + i_p = nl.arange(B_P_SIZE)[:, None]
43 + i_f = nl.arange(num_blocks_per_tile)[None, :]
44 + block_tables_sbuf[i, i_p, i_f] = nl.load(
45 + block_tables_hbm[i_p + i * B_P_SIZE, i_f],
46 + dtype=nl.int32,
47 + mask=(i_p + i * B_P_SIZE < num_tiles),
48 + )
49 + return block_tables_sbuf
50 +
51 +
52 +@nki.jit
53 +def transform_block_tables_for_indirect_load(
54 + block_tables,
55 + block_size_tiling_factor,
56 + num_head,
57 + head_id,
58 +):
59 """
15 - Config class for flash attention with default values
60 + This function does two things:
61 + 1. calculate new `block_tables` for a `head_id` after flattening
62 + `num_block`, `num_head`, and `block_size_tiling_factor` dimensions
63 + 2. transpose the result so that `block_table` for each tile is mapped to
64 + SBUF Partition dimension for vectorized DMA
65 +
66 + Tiling trick to further improve DMA performance:
67 + Given KV cache shape `(num_block, num_head, block_size, D)`, when loading M
68 + blocks of a given `head_id` from HBM, the load `cache[block_tables,
69 + head_id]` has shape `(M, block_size, D)`. If M < B_P_SIZE = 128, DMA may not
70 + fully utilize hardware parallelization. The solution is to tile `block_size`
71 + into `(block_size_tiling_factor, tiled_block_size)` s.t. `M *
72 + block_size_tiling_factor = B_P_SIZE`. After tiling, KV cache has shape
73 + `(num_block, num_head, block_size_tiling_factor, tiled_block_size, D)`.
74 +
75 + Note:
76 + We don't further tile D dimension as small DMA size also hurts performance.
77 """
78 + B_P_SIZE = 128
79 + num_partitions, num_tiles_per_partition, num_blocks_per_tile = (
80 + block_tables.shape)
81 + assert num_tiles_per_partition == B_P_SIZE
82 + assert is_power_of_2(
83 + num_blocks_per_tile), f"{num_blocks_per_tile=} is not power of 2"
84 +
85 + num_loads = ceil_div(num_blocks_per_tile, B_P_SIZE)
86 + block_tables_transposed = nl.ndarray(
87 + (
88 + num_loads,
89 + par_dim(B_P_SIZE),
90 + num_partitions * num_tiles_per_partition,
91 + ),
92 + dtype=nl.int32,
93 + )
94
18 - seq_tile_size: int = 2048
19 - should_transpose_v: bool = False
95 + # prepare iota ahead of time to avoid repeatedly using Gpsimd
96 + if num_head > 1:
97 + head_id = nisa.iota(head_id, dtype=nl.int32).reshape((1, 1))
98 + head_id = nl.transpose(
99 + head_id.broadcast_to((1, num_tiles_per_partition)))
100 + if num_blocks_per_tile > 1:
101 + head_id = head_id.broadcast_to(
102 + (num_tiles_per_partition, num_blocks_per_tile))
103 +
104 + if block_size_tiling_factor > 1:
105 + broadcast_shape = (
106 + num_tiles_per_partition,
107 + num_blocks_per_tile,
108 + block_size_tiling_factor,
109 + )
110 + offset = nisa.iota(nl.arange(block_size_tiling_factor)[None, None, :],
111 + dtype=nl.int32).broadcast_to(broadcast_shape)
112 +
113 + for partition_id in nl.affine_range(num_partitions):
114 + block_tables_partition = block_tables[partition_id]
115 + if num_head > 1:
116 + # fuse num_block and num_head dimension
117 + block_tables_partition = block_tables_partition * num_head + head_id
118 +
119 + if block_size_tiling_factor > 1:
120 + # need to apply block size tiling trick
121 + assert num_blocks_per_tile * block_size_tiling_factor == B_P_SIZE
122 + block_tables_partition = ((block_tables_partition *
123 + block_size_tiling_factor).reshape(
124 + (num_tiles_per_partition,
125 + num_blocks_per_tile,
126 + 1)).broadcast_to(broadcast_shape))
127 + new_block_tables = block_tables_partition + offset
128 + new_block_tables = new_block_tables.reshape(
129 + (num_tiles_per_partition, B_P_SIZE))
130 + else:
131 + new_block_tables = block_tables_partition
132
21 - __annotations__ = {
22 - "seq_tile_size": int,
23 - "should_transpose_v": bool,
24 - }
133 + # transpose the block table so that it can be used by vector DGE
134 + for i in nl.affine_range(num_loads):
135 + i_p = nl.arange(B_P_SIZE)[:, None]
136 + i_f = (partition_id * num_tiles_per_partition +
137 + nl.arange(num_tiles_per_partition)[None, :])
138 + block_tables_transposed[i, i_p, i_f] = nl.transpose(
139 + new_block_tables[:, nl.ds(i * B_P_SIZE, B_P_SIZE)])
140 + return block_tables_transposed
141 +
142 +
143 +@nki.jit
144 +def load_kv_tile_from_cache(
145 + cur_k_tile,
146 + cur_v_tile,
147 + key_cache,
148 + value_cache,
149 + block_tables,
150 + large_k_tile_idx,
151 + num_blocks_per_large_tile,
152 + tiled_block_size,
153 + B_P_SIZE,
154 + B_D_SIZE,
155 +):
156 + """
157 + Load KV cache and transform Key and Value into layout required by Matmul
158 +
159 + Vectorized DMA Load layout:
160 + Key and Value: (par_dim(B_P_SIZE), seqlen_kv // B_P_SIZE * B_D_SIZE)
161 +
162 + Layout used by attention matmuls:
163 + Key: (par_dim(B_D_SIZE), seqlen_kv)
164 + Value: (seqlen_kv // B_P_SIZE, par_dim(B_P_SIZE), B_D_SIZE)
165 + equivalent to (par_dim(B_P_SIZE), seqlen_kv // B_P_SIZE * B_D_SIZE)
166 + """
167 + # load key cache
168 + num_loads = ceil_div(num_blocks_per_large_tile, B_P_SIZE)
169 + for load_idx in nl.affine_range(num_loads):
170 + i_p = nl.arange(B_P_SIZE)[:, None]
171 + 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])
174 + if cur_k_tile.dtype != loaded.dtype:
175 + loaded = nl.copy(loaded, dtype=cur_k_tile.dtype)
176 + # Transpose SBUF tensor using PE
177 + for tb_i in nl.affine_range(tiled_block_size):
178 + cur_k_tile[
179 + :,
180 + nl.ds(
181 + load_idx * B_P_SIZE * tiled_block_size + tb_i * B_P_SIZE,
182 + B_P_SIZE,
183 + ),
184 + ] = nl.transpose(loaded[:, nl.ds(tb_i * B_D_SIZE, B_D_SIZE)])
185 +
186 + # load value cache
187 + for load_idx in nl.affine_range(num_loads):
188 + loaded = nl.load(value_cache[block_tables[load_idx, i_p,
189 + large_k_tile_idx], i_f])
190 + if cur_v_tile.dtype != loaded.dtype:
191 + loaded = nl.copy(loaded, dtype=cur_v_tile.dtype)
192 + i_p = nl.arange(B_P_SIZE)[:, None]
193 + i_f = nl.arange(tiled_block_size * B_D_SIZE)[None, :]
194 + cur_v_tile[
195 + :,
196 + nl.ds(
197 + load_idx * tiled_block_size * B_D_SIZE,
198 + tiled_block_size * B_D_SIZE,
199 + ),
200 + ] = loaded
201
202
203 @nki.jit
@@ -62,13 +238,13 @@ def _flash_attention_core(
238 o_buffer,
239 l_buffer,
240 m_buffer,
65 - q_tile_idx,
241 kernel_dtype,
242 acc_type,
68 - flash_config: FlashConfig,
69 - use_causal_mask,
243 tile_mask,
244 + use_causal_mask,
245 + q_tile_idx=None,
246 initialize=False,
247 + LARGE_TILE_SZ=2048,
248 B_P_SIZE=128,
249 B_F_SIZE=512,
250 B_D_SIZE=128,
@@ -77,19 +253,19 @@ def _flash_attention_core(
253 """
254 The flash attention core function to calculate self attention between a tile
255 of q and a block of K and V.
80 - The q_local_tile has (B_P_SIZE, B_F_SIZE), which is loaded into the SBUF
81 - already. The block size of K and V
82 - is defined in the seq_tile_size of the flash_config. The results are stored
83 - in the following three buffers
256 + The q_local_tile has (B_P_SIZE, B_D_SIZE)
257 + The K and V have shape (B_D_SIZE, LARGE_TILE_SZ), whose free dimension will
258 + be split into size B_F_SIZE tiles
259 +
260 + The results are stored in the following three buffers
261 o_buffer: (B_P_SIZE, d)
262 l_buffer: (B_P_SIZE, 1)
263 m_buffer: (B_P_SIZE, 1)
264 +
265 + All IO buffers are in SBUF.
266 """
88 - LARGE_TILE_SZ = flash_config.seq_tile_size
267 num_k_tile_per_large_tile = LARGE_TILE_SZ // B_F_SIZE
268
91 - # mask are used to only apply computation to the lower half of the matrix,
92 - # which reduce the arithmetic intensity by half
269 qk_res_buf = nl.ndarray((par_dim(B_P_SIZE), LARGE_TILE_SZ),
270 buffer=nl.sbuf,
271 dtype=acc_type)
@@ -99,6 +275,8 @@ def _flash_attention_core(
275 k_i_b_f_slice = nl.ds(k_i * B_F_SIZE, B_F_SIZE)
276
277 if use_causal_mask:
278 + # mask are used to only apply computation to the lower half of the
279 + # matrix, which reduce the arithmetic intensity by up to 50%
280 multiplication_required_selection = (q_tile_idx * B_P_SIZE
281 >= k_i * B_F_SIZE)
282 else:
@@ -165,7 +343,9 @@ def _flash_attention_core(
343 REDUCTION_TILE = min(2048, LARGE_TILE_SZ // 2)
344
345 p_partial_sum = nl.ndarray(
168 - (par_dim(B_P_SIZE), LARGE_TILE_SZ // REDUCTION_TILE), dtype=acc_type)
346 + (par_dim(B_P_SIZE), LARGE_TILE_SZ // REDUCTION_TILE),
347 + dtype=acc_type,
348 + )
349
350 for k_r_i in nl.affine_range(LARGE_TILE_SZ // REDUCTION_TILE):
351 k_r_i_reduce_slice = nl.ds(k_r_i * REDUCTION_TILE, REDUCTION_TILE)
@@ -194,13 +374,15 @@ def _flash_attention_core(
374 B_F_SIZE=B_F_SIZE,
375 )
376
197 - pv_psum = nl.zeros((par_dim(B_P_SIZE), B_D_SIZE),
198 - dtype=np.float32,
199 - buffer=nl.psum)
377 + pv_psum = nl.zeros(
378 + (par_dim(B_P_SIZE), B_D_SIZE),
379 + dtype=np.float32,
380 + buffer=nl.psum,
381 + )
382 for k_i in nl.affine_range(LARGE_TILE_SZ // B_P_SIZE):
383 pv_psum[:, :] += nl.matmul(
384 p_local_transposed[:, nl.ds(k_i * B_P_SIZE, B_P_SIZE)],
203 - v[k_i, :, :],
385 + v[:, nl.ds(k_i * B_D_SIZE, B_D_SIZE)],
386 transpose_x=True,
387 ) # (128, 128) (p(Br), d)
388
@@ -219,44 +401,16 @@ def _flash_attention_core(
401
402
403 @nki.jit
222 -def load_v_tile(v_hbm_tile, cur_v_tile, j, v_i, config):
223 - LARGE_TILE_SZ = config.seq_tile_size
404 +def load_v_tile(v_hbm_tile, cur_v_tile, large_tile_idx, v_i, LARGE_TILE_SZ):
405 B_P_SIZE = 128
225 -
226 - if not config.should_transpose_v:
227 - cur_v_tile[v_i, :, :] = nl.load(
228 - v_hbm_tile[nl.ds(j * LARGE_TILE_SZ + B_P_SIZE * v_i, B_P_SIZE), :],
229 - dtype=cur_v_tile.dtype,
230 - )
231 - return
232 -
233 - if nisa.get_nc_version() == nisa.nc_version.gen3:
234 - cur_v_tile_transposed = nisa.dma_transpose(
235 - v_hbm_tile[:,
236 - nl.ds(j * LARGE_TILE_SZ + B_P_SIZE * v_i, B_P_SIZE)])
237 - cur_v_tile[v_i, :, :] = nisa.tensor_copy(cur_v_tile_transposed,
238 - dtype=cur_v_tile.dtype)
239 - return
240 -
241 - cur_v_tile[v_i, :, :] = nl.load_transpose2d(
242 - v_hbm_tile[:, nl.ds(j * LARGE_TILE_SZ + B_P_SIZE * v_i, B_P_SIZE)],
243 - dtype=cur_v_tile.dtype,
244 - )
245 -
246 -
247 -@nki.jit
248 -def load_block_tables(block_tables_hbm, num_tiles):
249 - (num_blocks, ) = block_tables_hbm.shape
250 - assert num_blocks % num_tiles == 0
251 - num_blocks_per_tile = num_blocks // num_tiles
252 - block_tables_hbm = block_tables_hbm.reshape(
253 - (num_tiles, num_blocks_per_tile))
254 - block_tables_buffer = nl.load(block_tables_hbm, dtype=nl.int32)
255 - return block_tables_buffer
256 -
257 -
258 -def is_power_of_2(x):
259 - return x > 0 and (x & (x - 1)) == 0
406 + B_D_SIZE = v_hbm_tile.shape[-1]
407 + loaded = nl.load(v_hbm_tile[
408 + nl.ds(large_tile_idx * LARGE_TILE_SZ + B_P_SIZE * v_i, B_P_SIZE),
409 + :,
410 + ])
411 + if cur_v_tile.dtype != loaded.dtype:
412 + loaded = nl.copy(loaded, dtype=cur_v_tile.dtype)
413 + cur_v_tile[:, nl.ds(v_i * B_D_SIZE, B_D_SIZE)] = loaded
414
415
416 @nki.jit
@@ -270,24 +424,21 @@ def flash_paged_attention(
424 mask,
425 softmax_scale=None,
426 mixed_precision=True,
273 - config=None,
427 + LARGE_TILE_SZ=2048,
428 return_debug_tensors=False,
429 ):
430 """
431 Flash PagedAttention Forward Kernel.
278 - - PagedAttention Paper: https://arxiv.org/abs/2309.06180
279 - - Chunked Prefill Paper: https://arxiv.org/abs/2403.02310
432
433 IO tensor layouts:
434 - query: shape (1, n_heads, d, seq_q)
435 - key: shape (1, n_kv_heads, d, seq_k)
436 - value: shape (1, n_kv_heads, seq_v, d)
285 - - key_cache: (num_blocks, block_size, n_kv_heads, d)
286 - - value_cache: (num_blocks, block_size, n_kv_heads, d)
437 + - key_cache: (num_blocks, n_kv_heads, block_size, d)
438 + - value_cache: (num_blocks, n_kv_heads, block_size, d)
439 - block_tables: (num_active_blocks, )
288 - - mask: (seq_q, num_active_blocks * block_size)
440 + - mask: (seq_q, num_active_blocks * block_size + seq_q)
441 - o: shape (1, n_heads, seq_q, d)
290 - - l_m: shape (1, n_heads, seq_q, 2)
442
443 - This kernel requires seq_k == seq_v
444 - We use continuous batching by default, so the batch dimension is
@@ -306,11 +457,8 @@ def flash_paged_attention(
457 - softmax_scale: scaling for softmax, is None, default is `1.0/(d**0.5)`
458 - mixed_precision: flag to set non-matmul ops in fp32 precision, default
459 is set to `true`, if false, we use same precision as input types
309 - - config: Instance of dataclass :class:`nki.kernels.attention.FlashConfig`
310 - with Performance config parameters for flash attention with default
311 - values
312 - seq_tile_size: `default=2048`, size of the kv tile size for attention
313 - computation reduction
460 + - LARGE_TILE_SZ: `default=2048`, size of the kv tile size for attention
461 + computation reduction
462
463 GQA support Notes:
464 the spmd kernel for launching kernel should be on kv_heads instead of
@@ -322,31 +470,65 @@ def flash_paged_attention(
470 GQA: q: [b, h, d, s], k: [b, kv_h, d, s], v: [b, kv_h, s, d]
471 usage: `flash_fwd[b, kv_h](q, k, v, ...)`
472 """
325 - config = config or FlashConfig()
473 B_F_SIZE = 512
474 B_P_SIZE = 128
475 b, h, d, seqlen_q = query.shape
476 B_D_SIZE = d
330 - LARGE_TILE_SZ = config.seq_tile_size
477 n_tile_q = seqlen_q // B_P_SIZE # since q will be loaded on tensor engine
332 - num_blocks, block_size, k_h, _ = key_cache.shape
478 + num_blocks, k_h, block_size, _ = key_cache.shape
479 q_h_per_k_h = h // k_h
334 - assert tuple(key_cache.shape) == (
335 - num_blocks,
336 - block_size,
480 + assert b == 1, f"invalid batch size {b=}"
481 + 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}"
487 + assert key is None or tuple(key.shape) == (
488 + 1,
489 k_h,
490 d,
339 - ), "Input shape mismatch!"
340 - assert tuple(value_cache.shape) == (
341 - num_blocks,
342 - block_size,
491 + seqlen_q,
492 + ), f"key shape {key.shape} mismatch!"
493 + assert value is None or tuple(value.shape) == (
494 + 1,
495 k_h,
496 + seqlen_q,
497 d,
345 - ), "Input shape mismatch!"
346 - assert b == 1, f"invalid batch size {b=}"
347 - assert d <= 128, f" we do not support head_dim > 128, got head dim {d}"
498 + ), f"value shape {value.shape} mismatch!"
499 +
500 + assert (
501 + nl.program_ndim() == 2
502 + ), f"Expect spmd grid with 2 dimensions, got {nl.program_ndim()} instead!"
503 + batch_id = nl.program_id(axis=0)
504 + head_id = nl.program_id(axis=1)
505 +
506 + (num_active_blocks, ) = block_tables.shape
507 + context_kv_len = num_active_blocks * block_size
508 + assert (
509 + LARGE_TILE_SZ % B_F_SIZE == 0
510 + ), f"Need {LARGE_TILE_SZ=} to be divisible by {B_F_SIZE=} in transpose_p"
511 + assert (context_kv_len % LARGE_TILE_SZ == 0
512 + ), f"Need {context_kv_len=} to be divisible by {LARGE_TILE_SZ=}"
513 +
514 + num_blocks_per_large_tile = LARGE_TILE_SZ // block_size
515 + assert is_power_of_2(
516 + num_blocks_per_large_tile
517 + ), f"{num_blocks_per_large_tile=} is expected of be power of 2"
518 + if seqlen_q > B_F_SIZE:
519 + MAX_REDUCTION_TILE = 2048
520 + if seqlen_q // 2 > MAX_REDUCTION_TILE:
521 + assert (
522 + seqlen_q % MAX_REDUCTION_TILE == 0
523 + ), f"{seqlen_q=} should be divisible by {MAX_REDUCTION_TILE=}"
524 + else:
525 + assert (seqlen_q % B_F_SIZE == 0
526 + ), f"{seqlen_q=} should be divisible by {B_F_SIZE=})"
527 +
528 kernel_dtype = nl.bfloat16 if mixed_precision else query.dtype
529 acc_type = np.dtype(np.float32) if mixed_precision else kernel_dtype
530 + softmax_scale = softmax_scale or (1.0 / (d**0.5))
531 + num_large_k_tile = context_kv_len // LARGE_TILE_SZ
532
533 o = nl.ndarray((b, h, seqlen_q, d),
534 dtype=query.dtype,
@@ -373,35 +555,38 @@ def flash_paged_attention(
555 buffer=nl.sbuf,
556 lazy_initialization=True,
557 )
558 + block_tables_sbuf = load_block_tables(
559 + block_tables_hbm=block_tables,
560 + num_tiles=num_large_k_tile,
561 + num_blocks_per_tile=num_blocks_per_large_tile,
562 + )
563
377 - assert (
378 - nl.program_ndim() == 2
379 - ), f"Expect spmd grid with 2 dimensions, got {nl.program_ndim()} instead!"
380 - batch_id = nl.program_id(axis=0)
381 - head_id = nl.program_id(axis=1)
382 -
383 - softmax_scale = softmax_scale or (1.0 / (d**0.5))
384 -
385 - (num_active_blocks, ) = block_tables.shape
386 - context_kv_len = num_active_blocks * block_size
387 - assert (config.seq_tile_size >= 512
388 - ), f" seq tile_size {config.seq_tile_size} cannot be less than 512"
389 - assert (context_kv_len % LARGE_TILE_SZ == 0
390 - ), f"Need {context_kv_len=} to be divisible by {LARGE_TILE_SZ=}"
391 - assert (
392 - LARGE_TILE_SZ % B_P_SIZE == 0
393 - ), f"Need LARGE_TILE_SZ ({LARGE_TILE_SZ}) to be divisible by {B_P_SIZE=}"
394 - assert (B_P_SIZE % block_size == 0
395 - ), f"Need B_P_SIZE ({B_P_SIZE}) to be divisible by {block_size=}"
396 - num_large_k_tile = context_kv_len // LARGE_TILE_SZ
397 - num_blocks_per_large_tile = LARGE_TILE_SZ // block_size
398 - assert block_size % 32 == 0, "block_size is expected to be a multiple of 32"
399 - assert is_power_of_2(
400 - num_blocks_per_large_tile
401 - ), "The number of blocks in each large tile is expected of be power of 2"
402 - assert is_power_of_2(seqlen_q), "seqlen_q is expected to be power of 2"
564 + # On Neuron, we need B_P_SIZE = 128 blocks to make DMA efficient
565 + if num_blocks_per_large_tile < B_P_SIZE:
566 + # we checked num_blocks_per_tile is a power of 2
567 + assert B_P_SIZE % num_blocks_per_large_tile == 0
568 + block_size_tiling_factor = B_P_SIZE // num_blocks_per_large_tile
569 + # We assume block_size >= block_size_tiling_factor
570 + assert block_size % block_size_tiling_factor == 0
571 + else:
572 + block_size_tiling_factor = 1
573 + tiled_block_size = block_size // block_size_tiling_factor
574 +
575 + # Indirect DMA load must be placed along Partition Dimension
576 + block_tables_sbuf = transform_block_tables_for_indirect_load(
577 + block_tables_sbuf,
578 + block_size_tiling_factor=block_size_tiling_factor,
579 + num_head=k_h,
580 + head_id=head_id,
581 + )
582
404 - block_tables_sbuf = load_block_tables(block_tables, num_large_k_tile)
583 + # Flatten KV cache to be 2D for loading into SBUF
584 + new_cache_shape = (
585 + num_blocks * k_h * block_size_tiling_factor,
586 + tiled_block_size * d,
587 + )
588 + key_cache = key_cache.reshape(new_cache_shape)
589 + value_cache = value_cache.reshape(new_cache_shape)
590
591 # Global Flash Attention accumulators
592 o_buffer = nl.zeros(
@@ -411,7 +596,7 @@ def flash_paged_attention(
596 lazy_initialization=True,
597 )
598 l_buffer = nl.zeros(
414 - (par_dim(B_P_SIZE), n_tile_q, q_h_per_k_h),
599 + (n_tile_q, q_h_per_k_h, par_dim(B_P_SIZE), 1),
600 dtype=acc_type,
601 buffer=nl.sbuf,
602 lazy_initialization=True,
@@ -423,50 +608,42 @@ def flash_paged_attention(
608 lazy_initialization=True,
609 )
610
426 - for j in nl.sequential_range(0, num_large_k_tile):
427 - cur_k_tile = nl.ndarray((par_dim(B_D_SIZE), LARGE_TILE_SZ),
428 - dtype=kernel_dtype)
611 + for large_k_tile_idx in nl.sequential_range(0, num_large_k_tile):
612 + num_loads = ceil_div(num_blocks_per_large_tile, B_P_SIZE)
613 + cur_k_tile = nl.ndarray(
614 + (par_dim(B_D_SIZE), LARGE_TILE_SZ),
615 + dtype=kernel_dtype,
616 + )
617 cur_v_tile = nl.ndarray(
430 - (LARGE_TILE_SZ // B_P_SIZE, par_dim(B_P_SIZE), B_D_SIZE),
618 + (par_dim(B_P_SIZE), num_loads * tiled_block_size * B_D_SIZE),
619 dtype=kernel_dtype,
620 )
433 -
434 - for k_i in nl.affine_range(num_blocks_per_large_tile):
435 - loaded = nl.load(key_cache[block_tables_sbuf[j, k_i], :,
436 - head_id, :])
437 - cur_k_tile[:, nl.ds(k_i *
438 - block_size, block_size)] = nl.transpose(loaded)
439 -
440 - load_tile_size = B_P_SIZE
441 - num_blocks_per_partition = load_tile_size // block_size
442 - for partition_idx in nl.affine_range(LARGE_TILE_SZ // load_tile_size):
443 - for block_in_partition in nl.affine_range(
444 - num_blocks_per_partition):
445 - v_i = (partition_idx * num_blocks_per_partition +
446 - block_in_partition)
447 - loaded_v = nl.load(value_cache[block_tables_sbuf[j, v_i], :,
448 - head_id, :])
449 - cur_v_tile[
450 - partition_idx,
451 - nl.ds(block_in_partition * block_size, block_size),
452 - :,
453 - ] = loaded_v
621 + load_kv_tile_from_cache(
622 + cur_k_tile=cur_k_tile,
623 + cur_v_tile=cur_v_tile,
624 + key_cache=key_cache,
625 + value_cache=value_cache,
626 + block_tables=block_tables_sbuf,
627 + large_k_tile_idx=large_k_tile_idx,
628 + num_blocks_per_large_tile=num_blocks_per_large_tile,
629 + tiled_block_size=tiled_block_size,
630 + B_P_SIZE=B_P_SIZE,
631 + B_D_SIZE=B_D_SIZE,
632 + )
633
634 for i in nl.affine_range(n_tile_q):
456 - cur_mask = nl.ndarray((par_dim(B_P_SIZE), LARGE_TILE_SZ),
457 - dtype=mask.dtype)
458 - for m_i in nl.affine_range(LARGE_TILE_SZ // B_F_SIZE):
459 - cur_mask[:, nl.ds(m_i * B_F_SIZE, B_F_SIZE)] = nl.load(mask[
460 - nl.ds(i * B_P_SIZE, B_P_SIZE),
461 - nl.ds(j * LARGE_TILE_SZ + m_i * B_F_SIZE, B_F_SIZE),
462 - ])
635 + cur_mask = nl.load(mask[
636 + nl.ds(i * B_P_SIZE, B_P_SIZE),
637 + nl.ds(large_k_tile_idx * LARGE_TILE_SZ, LARGE_TILE_SZ),
638 + ])
639 for i_q_h in nl.affine_range(q_h_per_k_h):
640 q_tile = nl.ndarray((B_D_SIZE, B_P_SIZE), dtype=kernel_dtype)
641 q_hbm_tile = query[batch_id, head_id * q_h_per_k_h + i_q_h]
466 - q_sbuf_tile = nl.load(
467 - q_hbm_tile[:, nl.ds(i * B_P_SIZE, B_P_SIZE)],
468 - dtype=kernel_dtype,
469 - ) # load (d, 128) tile in SBUF
642 + q_sbuf_tile = nl.load(q_hbm_tile[:,
643 + nl.ds(i *
644 + B_P_SIZE, B_P_SIZE)])
645 + if q_sbuf_tile.dtype != kernel_dtype:
646 + q_sbuf_tile = nl.copy(q_sbuf_tile, dtype=kernel_dtype)
647 q_tile[:, :] = q_sbuf_tile * softmax_scale
648
649 _flash_attention_core(
@@ -474,15 +651,15 @@ def flash_paged_attention(
651 k=cur_k_tile,
652 v=cur_v_tile,
653 o_buffer=o_buffer[i, i_q_h],
477 - l_buffer=l_buffer[:, i, i_q_h],
654 + l_buffer=l_buffer[i, i_q_h],
655 m_buffer=m_buffer[i, i_q_h],
479 - q_tile_idx=i,
656 kernel_dtype=kernel_dtype,
657 acc_type=acc_type,
482 - flash_config=config,
483 - use_causal_mask=False,
658 tile_mask=cur_mask,
485 - initialize=j == 0,
659 + use_causal_mask=False,
660 + q_tile_idx=i,
661 + initialize=large_k_tile_idx == 0,
662 + LARGE_TILE_SZ=LARGE_TILE_SZ,
663 B_P_SIZE=B_P_SIZE,
664 B_F_SIZE=B_F_SIZE,
665 B_D_SIZE=B_D_SIZE,
@@ -492,62 +669,58 @@ def flash_paged_attention(
669 if key is not None and value is not None:
670 B_F_SIZE = min(seqlen_q, B_F_SIZE)
671 LARGE_TILE_SZ = seqlen_q
495 - active_config = FlashConfig(
496 - seq_tile_size=LARGE_TILE_SZ,
497 - should_transpose_v=config.should_transpose_v,
498 - )
672
673 cur_k_tile = nl.ndarray((par_dim(B_D_SIZE), LARGE_TILE_SZ),
674 dtype=kernel_dtype)
675 cur_v_tile = nl.ndarray(
503 - (LARGE_TILE_SZ // B_P_SIZE, par_dim(B_P_SIZE), B_D_SIZE),
676 + (par_dim(B_P_SIZE), LARGE_TILE_SZ // B_P_SIZE * B_D_SIZE),
677 dtype=kernel_dtype,
678 )
679
507 - cur_k_tile[:, :] = nl.load(key[batch_id, head_id, :, :])
680 + loaded = nl.load(key[batch_id, head_id, :, :])
681 + if loaded.dtype != kernel_dtype:
682 + loaded = nl.copy(loaded, dtype=kernel_dtype)
683 + cur_k_tile[:, :] = loaded
684
509 - load_tile_size = B_P_SIZE
685 v_hbm_tile = value[batch_id, head_id]
511 - for v_i in nl.affine_range(LARGE_TILE_SZ // load_tile_size):
686 + for v_i in nl.affine_range(LARGE_TILE_SZ // B_P_SIZE):
687 load_v_tile(
688 v_hbm_tile=v_hbm_tile,
689 cur_v_tile=cur_v_tile,
515 - j=0,
690 + large_tile_idx=0,
691 v_i=v_i,
517 - config=active_config,
692 + LARGE_TILE_SZ=LARGE_TILE_SZ,
693 )
694
695 for i in nl.affine_range(n_tile_q):
521 - cur_mask = nl.load(
522 - mask[
523 - nl.ds(i * B_P_SIZE, B_P_SIZE),
524 - nl.ds(context_kv_len, LARGE_TILE_SZ),
525 - ],
526 - dtype=mask.dtype,
527 - )
696 + cur_mask = nl.load(mask[
697 + nl.ds(i * B_P_SIZE, B_P_SIZE),
698 + nl.ds(context_kv_len, LARGE_TILE_SZ),
699 + ])
700 for i_q_h in nl.affine_range(q_h_per_k_h):
701
702 q_tile = nl.ndarray((B_D_SIZE, B_P_SIZE), dtype=kernel_dtype)
703 q_hbm_tile = query[batch_id, head_id * q_h_per_k_h + i_q_h]
532 - q_sbuf_tile = nl.load(
533 - q_hbm_tile[:, nl.ds(i * B_P_SIZE, B_P_SIZE)],
534 - dtype=kernel_dtype,
535 - ) # load (d, 128) tile in SBUF
704 + q_sbuf_tile = nl.load(q_hbm_tile[:,
705 + nl.ds(i *
706 + B_P_SIZE, B_P_SIZE)])
707 + if q_sbuf_tile.dtype != kernel_dtype:
708 + q_sbuf_tile = nl.copy(q_sbuf_tile, dtype=kernel_dtype)
709 q_tile[:, :] = q_sbuf_tile * softmax_scale
710 _flash_attention_core(
711 q_local_tile=q_tile,
712 k=cur_k_tile,
713 v=cur_v_tile,
714 o_buffer=o_buffer[i, i_q_h],
542 - l_buffer=l_buffer[:, i, i_q_h],
715 + l_buffer=l_buffer[i, i_q_h],
716 m_buffer=m_buffer[i, i_q_h],
544 - q_tile_idx=i,
717 kernel_dtype=kernel_dtype,
718 acc_type=acc_type,
547 - flash_config=active_config,
548 - use_causal_mask=True,
719 tile_mask=cur_mask,
720 + use_causal_mask=True,
721 + q_tile_idx=i,
722 initialize=False,
723 + LARGE_TILE_SZ=LARGE_TILE_SZ,
724 B_P_SIZE=B_P_SIZE,
725 B_F_SIZE=B_F_SIZE,
726 B_D_SIZE=B_D_SIZE,
@@ -559,8 +732,8 @@ def flash_paged_attention(
732 for i_q_h in nl.affine_range(q_h_per_k_h):
733 for i in nl.affine_range(n_tile_q):
734 out = nl.multiply(
562 - o_buffer[i, i_q_h, :, :],
563 - nl.exp(m_buffer[i, i_q_h, :, :] - l_buffer[:, i, i_q_h]),
735 + o_buffer[i, i_q_h],
736 + nl.exp(m_buffer[i, i_q_h] - l_buffer[i, i_q_h]),
737 dtype=kernel_dtype,
738 )
739
@@ -589,7 +762,7 @@ def flash_paged_attention(
762 head_id * q_h_per_k_h + i_q_h,
763 nl.ds(i * B_P_SIZE, B_P_SIZE),
764 ],
592 - l_buffer[:, i, i_q_h],
765 + l_buffer[i, i_q_h],
766 )
767 nl.store(
768 hbm_qk_res[batch_id, head_id * q_h_per_k_h + i_q_h, :, :],
@@ -601,6 +774,49 @@ def flash_paged_attention(
774 return o
775
776
777 +def reorder_context_mask(mask, LARGE_TILE_SZ, block_size):
778 + """
779 + Reorder the mask to make it compatible with the flash attention kernel.
780 +
781 + We vectorize KV cache read to improve DMA utilization. However, the layout
782 + that maximizes DMA bandwidth changes the order tokens are consumed.
783 +
784 + The token layout (inner 2 dimensions) after vectorized load is (B_P_SIZE,
785 + tiled_block_size) in a tile of `B_P_SIZE * tiled_block_size` tokens. And
786 + each step the engine consumes a column (rather than a row) of B_P_SIZE
787 + tokens. Therefore, the tokens are visited in a strided way.
788 +
789 + To make sure mask matches the order tokens are consumed, we need to properly
790 + transpose mask.
791 + """
792 + total_query_len, total_seq_len = mask.shape
793 + context_kv_len = total_seq_len - total_query_len
794 +
795 + B_P_SIZE = 128
796 + assert (LARGE_TILE_SZ
797 + >= B_P_SIZE), f"{LARGE_TILE_SZ=} must be larger than {B_P_SIZE=}"
798 + num_tiled_blocks = max(B_P_SIZE, LARGE_TILE_SZ // block_size)
799 + tiled_block_size = LARGE_TILE_SZ // num_tiled_blocks
800 + if tiled_block_size > 1:
801 + # Mask reordering is needed when tiled_block_size > 1
802 + device = mask.device
803 + mask = mask.cpu()
804 + context_mask = mask[:, :context_kv_len]
805 + context_mask = context_mask.view(
806 + total_query_len,
807 + context_kv_len // LARGE_TILE_SZ,
808 + num_tiled_blocks // B_P_SIZE,
809 + B_P_SIZE,
810 + tiled_block_size,
811 + )
812 + context_mask = context_mask.transpose(3, 4).reshape(
813 + total_query_len, context_kv_len)
814 + new_mask = mask[:, context_kv_len:]
815 + return torch.concat([context_mask, new_mask], dim=1).to(device)
816 + else:
817 + return mask
818 +
819 +
820 def flash_attn_varlen_nkifunc(
821 query,
822 key,
@@ -612,13 +828,32 @@ def flash_attn_varlen_nkifunc(
828 n_kv_head=None,
829 head_size=None,
830 LARGE_TILE_SZ=2048,
615 - return_debug_tensors=False,
831 mixed_precision=True,
832 ):
618 - config = FlashConfig(
619 - seq_tile_size=LARGE_TILE_SZ,
620 - should_transpose_v=False,
621 - )
833 + """
834 + Compute flash paged attention for variable length sequences.
835 +
836 + This function is a wrapper around the flash attention NKI kernel. It takes
837 + in the following arguments:
838 + - query: (1, n_heads, d, seq_q)
839 + - key: (1, n_kv_heads, d, seq_k)
840 + - 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)
843 + - block_tables: (n_active_blocks, )
844 + - attn_mask: (seq_q, n_active_blocks * block_size + seq_q)
845 +
846 + Notes:
847 + - attn_mask must be reordered outside using `reorder_context_mask`
848 + - Key/value cache layout must be (n_blocks, n_kv_heads, block_size, d)
849 + for better DMA throughput
850 + """
851 + if n_kv_head is None:
852 + n_kv_head = key_cache.shape[1]
853 + assert key_cache.shape[1] == n_kv_head
854 + if head_size is None:
855 + head_size = key_cache.shape[-1]
856 +
857 kwargs = dict(
858 query=query,
859 key=key,
@@ -628,15 +863,9 @@ def flash_attn_varlen_nkifunc(
863 block_tables=block_table,
864 mask=attn_mask,
865 softmax_scale=1.0 / (head_size**0.5),
631 - config=config,
866 mixed_precision=mixed_precision,
633 - return_debug_tensors=return_debug_tensors,
867 + LARGE_TILE_SZ=LARGE_TILE_SZ,
868 )
635 - _, n_kv_head, _, _ = key.shape
869
637 - if return_debug_tensors:
638 - o, *debug_tensors = flash_paged_attention[1, n_kv_head](**kwargs)
639 - return o, *debug_tensors
640 - else:
641 - o = flash_paged_attention[1, n_kv_head](**kwargs)
642 - return o
870 + o = flash_paged_attention[1, n_kv_head](**kwargs)
871 + return o