@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 741 lines 20.8 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 import torch
5 from einops import rearrange
6
7 from vllm.platforms import current_platform
8 from vllm.triton_utils import tl, triton
9 from vllm.v1.attention.backends.utils import PAD_SLOT_ID
10
11
12 @triton.jit
13 def _fwd_diag_kernel(
14 Q,
15 K,
16 V,
17 Out,
18 S,
19 b: tl.constexpr,
20 h: tl.constexpr,
21 n,
22 d: tl.constexpr,
23 e: tl.constexpr,
24 BLOCK: tl.constexpr,
25 NUM_BLOCK,
26 CBLOCK: tl.constexpr,
27 ):
28 # This kernel computes the diagonal blocks of the attention matrix
29 # Each diagonal block represents attention
30 # where queries attend to keys in the same block
31 off = tl.program_id(0)
32 off_bh = off // NUM_BLOCK # batch-head index
33 off_block = off % NUM_BLOCK # block index within the sequence
34 off_cblock = tl.program_id(1) # sub-block index within a block
35
36 off_h = off_bh % h # head index
37
38 # Calculate base offsets for the current batch and head
39 qk_offset = off_bh * n * d
40 v_offset = off_bh * n * e
41 o_offset = off_bh * n * e
42
43 # Calculate offsets for the current block
44 block_offset = off_block * BLOCK
45 qk_block_offset = block_offset * d
46 v_block_offset = block_offset * e
47 o_block_offset = block_offset * e
48
49 # Calculate offsets for the current sub-block
50 cblock_offset = off_cblock * CBLOCK
51 q_cblock_offset = cblock_offset * d
52 o_cblock_offset = cblock_offset * e
53
54 # Calculate pointers to the query, key, value, and output tensors
55 Q_block_ptr = (
56 Q
57 + qk_offset
58 + qk_block_offset
59 + q_cblock_offset
60 + tl.arange(0, CBLOCK)[:, None] * d
61 + tl.arange(0, d)[None, :]
62 )
63 K_trans_block_ptr = (
64 K
65 + qk_offset
66 + qk_block_offset
67 + tl.arange(0, CBLOCK)[None, :] * d
68 + tl.arange(0, d)[:, None]
69 )
70 V_block_ptr = (
71 V
72 + v_offset
73 + v_block_offset
74 + tl.arange(0, CBLOCK)[:, None] * e
75 + tl.arange(0, e)[None, :]
76 )
77 O_block_ptr = (
78 Out
79 + o_offset
80 + o_block_offset
81 + o_cblock_offset
82 + tl.arange(0, CBLOCK)[:, None] * e
83 + tl.arange(0, e)[None, :]
84 )
85
86 # Load the decay rate for the current head
87 S_block_ptr = S + off_h
88 s = tl.load(S_block_ptr)
89
90 i = off_cblock
91 q_index = tl.arange(0, CBLOCK) + i * CBLOCK
92
93 # Load query values
94 q = tl.load(Q_block_ptr, mask=block_offset + q_index[:, None] < n, other=0.0).to(
95 tl.float32
96 )
97
98 # Initialize output accumulator
99 qkv = tl.zeros([CBLOCK, e], dtype=tl.float32)
100
101 # Process all sub-blocks up to and
102 # including the current one (causal attention)
103 for j in range(i + 1):
104 kv_index = tl.arange(0, CBLOCK) + j * CBLOCK
105 diff = q_index[:, None] - kv_index[None, :]
106 s_index = s * diff
107 # Apply causal mask: only attend to positions before the current one
108 s_index = tl.where(diff >= 0, -s_index, float("-inf"))
109 decay = tl.exp(s_index)
110
111 # Load key and value
112 k_trans = tl.load(
113 K_trans_block_ptr,
114 mask=block_offset + kv_index[None, :] < n,
115 other=0.0,
116 ).to(tl.float32)
117 v = tl.load(
118 V_block_ptr,
119 mask=block_offset + kv_index[:, None] < n,
120 other=0.0,
121 ).to(tl.float32)
122
123 # Compute attention scores and apply decay
124 qk = tl.dot(q, k_trans) * decay
125
126 # Compute weighted values and accumulate
127 qkv += tl.dot(qk, v)
128
129 # Move to the next sub-block
130 K_trans_block_ptr += CBLOCK * d
131 V_block_ptr += CBLOCK * e
132
133 # Store the result
134 tl.store(
135 O_block_ptr,
136 qkv.to(O_block_ptr.dtype.element_ty),
137 mask=block_offset + q_index[:, None] < n,
138 )
139
140
141 @triton.jit
142 def _fwd_kv_parallel(
143 K,
144 V,
145 K_decay,
146 KV,
147 b: tl.constexpr,
148 h: tl.constexpr,
149 n,
150 d: tl.constexpr,
151 e: tl.constexpr,
152 BLOCK: tl.constexpr,
153 NUM_BLOCK,
154 D_FBLOCK: tl.constexpr,
155 E_FBLOCK: tl.constexpr,
156 NUM_FBLOCK: tl.constexpr,
157 CBLOCK: tl.constexpr,
158 NUM_CBLOCK: tl.constexpr,
159 ):
160 # This kernel computes the key-value outer
161 # products for each block in parallel
162 off_bh = tl.program_id(0) # batch-head index
163 off_block = tl.program_id(1) # block index
164
165 off_h = off_bh % h # head index
166
167 block_offset = off_block * BLOCK
168
169 # Calculate offsets for the current block
170 k_block_offset = block_offset * d
171 v_block_offset = block_offset * e
172 kv_block_offset = off_block * d * e
173
174 # Calculate base offsets for the current batch and head
175 k_offset = off_bh * n * d
176 v_offset = off_bh * n * e
177 kv_offset = off_bh * NUM_BLOCK * d * e
178
179 # Calculate pointers to the key, value, and key-value tensors
180 K_trans_block_ptr = (
181 K
182 + k_offset
183 + k_block_offset
184 + tl.arange(0, CBLOCK)[None, :] * d
185 + tl.arange(0, D_FBLOCK)[:, None]
186 )
187 V_block_ptr = (
188 V
189 + v_offset
190 + v_block_offset
191 + tl.arange(0, CBLOCK)[:, None] * e
192 + tl.arange(0, E_FBLOCK)[None, :]
193 )
194 KV_block_ptr = (
195 KV
196 + kv_offset
197 + kv_block_offset
198 + tl.arange(0, D_FBLOCK)[:, None] * e
199 + tl.arange(0, E_FBLOCK)[None, :]
200 )
201
202 # Load the decay factors for the current head and block
203 k_decay_ptr = K_decay + off_h * BLOCK + tl.arange(0, CBLOCK)
204
205 kv_index = tl.arange(0, CBLOCK)
206
207 # Initialize the key-value outer product accumulator
208 kv = tl.zeros([D_FBLOCK, E_FBLOCK], dtype=tl.float32)
209
210 # Handle the last block which might be smaller than BLOCK
211 split_n = n - (NUM_BLOCK - 1) * BLOCK if off_block == NUM_BLOCK - 1 else BLOCK
212 left_shift = tl.cdiv(split_n, CBLOCK) * CBLOCK - split_n
213 num_blocks = min(tl.cdiv(split_n, CBLOCK), NUM_CBLOCK)
214 k_decay_ptr += (NUM_CBLOCK - num_blocks) * CBLOCK
215
216 # Process all sub-blocks in the current block
217 for j in range(num_blocks):
218 left_bound = (1 - j) * left_shift
219 # Load key and value, handling boundary conditions
220 k_trans = tl.load(
221 K_trans_block_ptr - left_shift * d,
222 mask=kv_index[None, :] >= left_bound,
223 other=0.0,
224 )
225 v = tl.load(
226 V_block_ptr - left_shift * e,
227 mask=kv_index[:, None] >= left_bound,
228 other=0.0,
229 )
230
231 # Load decay factor and compute weighted key-value outer product
232 k_decay = tl.load(k_decay_ptr)
233
234 # NOTE: Need to add the extra dim here due to AMD MLIR lowering error.
235 # Please don't move it back until issue is resolved.
236 # Issue: https://github.com/ROCm/triton/issues/907
237 k_decay = k_decay[None, :]
238
239 kv += tl.dot(k_trans * k_decay, v)
240
241 # Move to the next sub-block
242 K_trans_block_ptr += CBLOCK * d
243 V_block_ptr += CBLOCK * e
244 k_decay_ptr += CBLOCK
245
246 # Store the result
247 tl.store(KV_block_ptr, kv.to(KV_block_ptr.dtype.element_ty))
248
249
250 @triton.jit
251 def _fwd_kv_reduce(
252 S,
253 KV,
254 KV_HISTORY,
255 b: tl.constexpr,
256 h: tl.constexpr,
257 n,
258 d: tl.constexpr,
259 e: tl.constexpr,
260 BLOCK: tl.constexpr,
261 NUM_BLOCK,
262 D_FBLOCK: tl.constexpr,
263 E_FBLOCK: tl.constexpr,
264 ):
265 # This kernel reduces the key-value outer products
266 # across blocks and updates the KV history
267 off_bh = tl.program_id(0) # batch-head index
268 off_h = off_bh % h # head index
269
270 kv_offset = off_bh * NUM_BLOCK * d * e
271
272 # Calculate pointer to the key-value tensor
273 KV_block_ptr = (
274 KV
275 + kv_offset
276 + tl.arange(0, D_FBLOCK)[:, None] * e
277 + tl.arange(0, E_FBLOCK)[None, :]
278 )
279
280 # Load the decay rate for the current head
281 s_ptrs = S + off_h
282 s = tl.load(s_ptrs)
283
284 # Calculate pointer to the key-value history tensor
285 kv_history_offset = off_bh * d * e
286 KV_HISTORY_block_ptr = (
287 KV_HISTORY
288 + kv_history_offset
289 + tl.arange(0, D_FBLOCK)[:, None] * e
290 + tl.arange(0, E_FBLOCK)[None, :]
291 )
292
293 # Load the previous key-value history
294 kv_pre = tl.load(KV_HISTORY_block_ptr).to(tl.float32)
295
296 # Process all blocks in reverse order to compute the prefix sum
297 for i in range(NUM_BLOCK):
298 block_size = min(n - i * BLOCK, BLOCK)
299 # Compute decay factor for the current block
300 block_decay = tl.exp(-s.to(tl.float32) * block_size)
301
302 # Load the current key-value outer product
303 kv_cur = tl.load(KV_block_ptr).to(tl.float32)
304 # Store the previous key-value history to the current block
305 tl.store(KV_block_ptr, kv_pre.to(KV_block_ptr.dtype.element_ty))
306
307 # Update the key-value history with the current block
308 kv_pre = block_decay * kv_pre + kv_cur
309 KV_block_ptr += d * e
310
311 # Store the updated key-value history
312 tl.store(KV_HISTORY_block_ptr, kv_pre)
313
314
315 @triton.jit
316 def _fwd_none_diag_kernel(
317 Q,
318 Out,
319 S,
320 KV,
321 b: tl.constexpr,
322 h: tl.constexpr,
323 n,
324 d: tl.constexpr,
325 e: tl.constexpr,
326 BLOCK: tl.constexpr,
327 NUM_BLOCK,
328 E_FBLOCK: tl.constexpr,
329 CBLOCK: tl.constexpr,
330 NUM_CBLOCK: tl.constexpr,
331 ):
332 # This kernel computes the non-diagonal blocks of the attention matrix
333 # Each non-diagonal block represents attention
334 # where queries attend to keys in different blocks
335 off_bh = tl.program_id(0) # batch-head index
336 off_h = off_bh % h # head index
337
338 off_nc = tl.program_id(1)
339 off_n = off_nc // NUM_CBLOCK # block index
340 off_c = off_nc % NUM_CBLOCK # sub-block index
341 off_e = tl.program_id(2) # output feature block index
342
343 n_offset = off_n * BLOCK
344 c_offset = off_c * CBLOCK
345 e_offset = off_e * E_FBLOCK
346 block_offset = n_offset + c_offset
347
348 # Calculate offsets for the current batch, head, and block
349 q_offset = off_bh * n * d + (n_offset + c_offset) * d
350 o_offset = off_bh * n * e + (n_offset + c_offset) * e + e_offset
351 kv_offset = off_bh * NUM_BLOCK * d * e + off_n * d * e + e_offset
352
353 # Calculate pointers to the query, output, and key-value tensors
354 Q_block_ptr = (
355 Q + q_offset + tl.arange(0, CBLOCK)[:, None] * d + tl.arange(0, d)[None, :]
356 )
357 O_block_ptr = (
358 Out
359 + o_offset
360 + tl.arange(0, CBLOCK)[:, None] * e
361 + tl.arange(0, E_FBLOCK)[None, :]
362 )
363 KV_block_ptr = (
364 KV + kv_offset + tl.arange(0, d)[:, None] * e + tl.arange(0, E_FBLOCK)[None, :]
365 )
366
367 # Load the decay rate for the current head
368 S_block_ptr = S + off_h
369 s = tl.load(S_block_ptr)
370
371 c_array = tl.arange(0, CBLOCK)
372
373 # Load the key-value outer product for the current block
374 kv = tl.load(KV_block_ptr).to(tl.float32)
375 q_index = block_offset + tl.arange(0, CBLOCK)
376
377 # Load query values
378 q = tl.load(Q_block_ptr, mask=q_index[:, None] < n, other=0.0).to(tl.float32)
379
380 # Compute decay factors for the current sub-block
381 q_decay = tl.exp(-s.to(tl.float32) * (off_c * CBLOCK + c_array[:, None]))
382
383 # Compute non-diagonal attention output
384 qkv_none_diag = tl.dot(q, kv) * q_decay
385
386 # Load diagonal attention output (computed by _fwd_diag_kernel)
387 qkv_diag = tl.load(O_block_ptr, mask=q_index[:, None] < n, other=0.0).to(tl.float32)
388
389 # Combine diagonal and non-diagonal attention outputs
390 qkv = qkv_diag + qkv_none_diag
391
392 # Store the result
393 tl.store(
394 O_block_ptr, qkv.to(O_block_ptr.dtype.element_ty), mask=q_index[:, None] < n
395 )
396
397
398 class _attention(torch.autograd.Function):
399 @staticmethod
400 def forward(ctx, q, k, v, s, kv_history):
401 # Forward pass of the lightning attention algorithm
402 q = q.contiguous()
403 k = k.contiguous()
404 v = v.contiguous()
405 s = s.contiguous()
406
407 # Check CUDA compute capability (Ampere+ required for flash attention
408 # path). Other accelerators (ROCm, XPU) rely on their own Triton
409 # backend support and skip this check.
410 if current_platform.is_cuda():
411 capability = torch.cuda.get_device_capability()
412 if capability[0] < 8:
413 raise RuntimeError(
414 "Flash attention currently only supported",
415 "for compute capability >= 80",
416 )
417
418 # Get input dimensions
419 b, h, n, d = q.shape
420 e = v.shape[-1]
421
422 # Initialize output tensor
423 o = torch.empty((b, h, n, e), dtype=q.dtype, device=q.device)
424
425 # Set block sizes
426 BLOCK = 256
427 NUM_BLOCK = triton.cdiv(n, BLOCK)
428
429 CBLOCK = 32
430 NUM_CBLOCK = BLOCK // CBLOCK
431 assert BLOCK % CBLOCK == 0, "BLOCK must be a multiple of CBLOCK"
432
433 # Compute decay factors for keys
434 array = torch.arange(0, BLOCK, device=q.device) + 1
435 k_decay = torch.exp(-s * (BLOCK - array.reshape(1, -1)))
436
437 # Step 1: Compute diagonal blocks of attention
438 grid = (b * h * NUM_BLOCK, NUM_CBLOCK)
439 _fwd_diag_kernel[grid](
440 q,
441 k,
442 v,
443 o,
444 s,
445 b,
446 h,
447 n,
448 d,
449 e,
450 BLOCK=BLOCK,
451 NUM_BLOCK=NUM_BLOCK,
452 CBLOCK=CBLOCK,
453 )
454
455 # Set feature block sizes
456 NUM_FBLOCK = 1
457 D_FBLOCK = d // NUM_FBLOCK
458 assert d % NUM_FBLOCK == 0
459 E_FBLOCK = e // NUM_FBLOCK
460 assert e % NUM_FBLOCK == 0
461
462 CBLOCK = 64
463 NUM_CBLOCK = BLOCK // CBLOCK
464 assert BLOCK % CBLOCK == 0, "BLOCK must be a multiple of CBLOCK"
465
466 # Step 2: Compute key-value outer products for each block in parallel
467 kv = torch.empty((b, h, NUM_BLOCK, d, e), dtype=torch.float32, device=q.device)
468 grid = (b * h, NUM_BLOCK)
469 _fwd_kv_parallel[grid](
470 k,
471 v,
472 k_decay,
473 kv,
474 b,
475 h,
476 n,
477 d,
478 e,
479 BLOCK=BLOCK,
480 NUM_BLOCK=NUM_BLOCK,
481 D_FBLOCK=D_FBLOCK,
482 E_FBLOCK=E_FBLOCK,
483 NUM_FBLOCK=NUM_FBLOCK,
484 CBLOCK=CBLOCK,
485 NUM_CBLOCK=NUM_CBLOCK,
486 )
487
488 # Step 3: Reduce key-value outer products
489 # across blocks and update KV history
490 grid = (b * h, NUM_FBLOCK)
491 _fwd_kv_reduce[grid](
492 s,
493 kv,
494 kv_history,
495 b,
496 h,
497 n,
498 d,
499 e,
500 BLOCK=BLOCK,
501 NUM_BLOCK=NUM_BLOCK,
502 D_FBLOCK=D_FBLOCK,
503 E_FBLOCK=E_FBLOCK,
504 )
505
506 # Step 4: Compute non-diagonal blocks of attention
507 grid = (b * h, NUM_BLOCK * NUM_CBLOCK)
508 _fwd_none_diag_kernel[grid](
509 q,
510 o,
511 s,
512 kv,
513 b,
514 h,
515 n,
516 d,
517 e,
518 BLOCK=BLOCK,
519 NUM_BLOCK=NUM_BLOCK,
520 E_FBLOCK=E_FBLOCK,
521 CBLOCK=CBLOCK,
522 NUM_CBLOCK=NUM_CBLOCK,
523 )
524
525 # Save tensors for backward pass
526 ctx.save_for_backward(q, k, v, s, kv)
527 ctx.BLOCK = BLOCK
528
529 return o, torch.cat([kv, kv_history.unsqueeze(2)], dim=2)
530
531
532 # Apply the lightning attention function
533 lightning_attention_ = _attention.apply
534
535
536 def lightning_attention(
537 q: torch.Tensor,
538 k: torch.Tensor,
539 v: torch.Tensor,
540 ed: torch.Tensor,
541 block_size: int = 256,
542 kv_history: torch.Tensor | None = None,
543 ) -> tuple[torch.Tensor, torch.Tensor]:
544 """Apply lightning attention algorithm
545 to compute attention efficiently.
546
547 Args:
548 q: Query tensor of shape [batch, heads, seq_len, dim]
549 k: Key tensor of shape [batch, heads, seq_len, dim]
550 v: Value tensor of shape [batch, heads, seq_len, dim_v]
551 ed: Decay rate tensor of shape [heads]
552 block_size: Size of blocks for block-sparse attention
553 kv_history: Optional key-value history from previous computations
554
555 Returns:
556 output: Attention output
557 kv: Updated key-value history
558
559 """
560 d = q.shape[-1]
561 e = v.shape[-1]
562
563 if ed.dim() == 1:
564 ed = ed.view(1, -1, 1, 1)
565
566 # Split the computation into chunks for better parallelism
567 m = 128 if d >= 128 else 64
568 assert d % m == 0, f"Dimension d ({d}) must be divisible by m ({m})"
569 arr = [m * i for i in range(d // m + 1)]
570 if arr[-1] != d:
571 arr.append(d)
572 n = len(arr)
573 output = 0
574
575 # Initialize or clone key-value history
576 if kv_history is None:
577 kv_history = torch.zeros(
578 (q.shape[0], q.shape[1], d, e), dtype=torch.float32, device=q.device
579 )
580 else:
581 kv_history = kv_history.clone().contiguous()
582
583 # Process each chunk and accumulate results
584 for i in range(n - 1):
585 s = arr[i]
586 e = arr[i + 1]
587 q1 = q[..., s:e]
588 k1 = k[..., s:e]
589 o, kv = lightning_attention_(q1, k1, v, ed, kv_history)
590 output = output + o
591 return output, kv
592
593
594 @triton.jit
595 def _linear_attn_decode_kernel(
596 q_ptr,
597 k_ptr,
598 v_ptr,
599 kv_cache_ptr,
600 slope_rate,
601 slot_idx,
602 output_ptr,
603 D: tl.constexpr,
604 qkv_b_stride,
605 qkv_h_stride,
606 cache_b_stride,
607 cache_h_stride,
608 cache_d0_stride,
609 cache_d1_stride,
610 pad_slot_id: tl.constexpr,
611 BLOCK_SIZE: tl.constexpr,
612 ):
613 """Kernel for linear attention decoding with KV cache.
614
615 This kernel computes attention for a single token using the KV cache.
616 """
617 pid_b = tl.program_id(0) # batch index
618 pid_h = tl.program_id(1) # head index
619 pid_d = tl.program_id(2) # dimension block index
620
621 # Load slot index for the current batch
622 slot_id = tl.load(slot_idx + pid_b).to(tl.int64)
623
624 # Skip if slot_id is PAD_SLOT_ID (padding)
625 if slot_id == pad_slot_id:
626 return
627
628 batch_id = pid_b
629 head_id = pid_h
630
631 # Load decay rate for the current head
632 ratio = tl.load(slope_rate + pid_h)
633
634 # Calculate offsets for dimensions
635 qk_d_offsets = tl.arange(0, D)
636 v_d_offsets = tl.arange(0, BLOCK_SIZE) + pid_d * BLOCK_SIZE
637 cache_d_offsets = (
638 qk_d_offsets[:, None] * cache_d0_stride + v_d_offsets[None, :] * cache_d1_stride
639 )
640
641 # Calculate offsets for the current batch and head
642 q_offset = batch_id * qkv_b_stride + head_id * qkv_h_stride
643 k_offset = batch_id * qkv_b_stride + head_id * qkv_h_stride
644 v_offset = batch_id * qkv_b_stride + head_id * qkv_h_stride
645
646 cache_offset = slot_id * cache_b_stride + head_id * cache_h_stride
647
648 # Create masks for loading tensors
649 qk_mask = qk_d_offsets < D
650 v_mask = v_d_offsets < D
651
652 # Load query, key, and value tensors
653 q = tl.load(q_ptr + q_offset + qk_d_offsets, mask=qk_mask, other=0.0)
654 k = tl.load(k_ptr + k_offset + qk_d_offsets, mask=qk_mask, other=0.0)
655 v = tl.load(v_ptr + v_offset + v_d_offsets, mask=v_mask, other=0.0)
656
657 # Compute key-value outer product
658 kv_outer = k[:, None] * v[None, :]
659 kv_mask = qk_mask[:, None] & v_mask[None, :]
660
661 # Apply decay to previous KV cache
662 ratio = tl.exp(-ratio)
663 kv_ptr = kv_cache_ptr + cache_offset + cache_d_offsets
664 kv_cache_old = tl.load(kv_ptr, mask=kv_mask, other=0.0)
665 kv_outer = kv_outer + ratio * kv_cache_old
666
667 # Compute attention output
668 output = q[:, None].to(tl.float32) * kv_outer
669 output = tl.sum(output, axis=0)
670
671 # Update KV cache and store output
672 tl.store(kv_ptr, kv_outer, mask=kv_mask)
673 tl.store(output_ptr + q_offset + v_d_offsets, output, mask=v_mask)
674
675
676 def linear_decode_forward_triton(
677 q: torch.Tensor,
678 k: torch.Tensor,
679 v: torch.Tensor,
680 kv_caches: torch.Tensor,
681 slope_rate: torch.Tensor,
682 slot_idx: torch.Tensor,
683 BLOCK_SIZE: int = 32,
684 ) -> torch.Tensor:
685 """Perform linear attention decoding using Triton kernels.
686
687 Args:
688 q: Query tensor of shape [B, H, 1, D]
689 k: Key tensor of shape [B, H, 1, D]
690 v: Value tensor of shape [B, H, 1, D]
691 kv_caches: Key-value cache tensor
692 slope_rate: Decay rate tensor
693 slot_idx: Slot indices for batches
694 BLOCK_SIZE: Size of blocks for processing
695
696 Returns:
697 output: Attention output tensor
698
699 """
700 B, H, _, D = q.shape
701 assert k.shape == (B, H, 1, D)
702 assert v.shape == (B, H, 1, D)
703
704 # Initialize output tensor
705 output = torch.empty_like(q)
706
707 # Set grid dimensions for the kernel
708 grid = (B, H, D // BLOCK_SIZE)
709
710 # Calculate strides for tensors
711 qkv_b_stride = q.stride(0)
712 qkv_h_stride = q.stride(1)
713
714 cache_b_stride = kv_caches.stride(0)
715 cache_h_stride = kv_caches.stride(1)
716 cache_d0_stride = kv_caches.stride(2)
717 cache_d1_stride = kv_caches.stride(3)
718
719 # Launch the kernel
720 _linear_attn_decode_kernel[grid](
721 q,
722 k,
723 v,
724 kv_caches,
725 slope_rate,
726 slot_idx,
727 output,
728 D,
729 qkv_b_stride,
730 qkv_h_stride,
731 cache_b_stride,
732 cache_h_stride,
733 cache_d0_stride,
734 cache_d1_stride,
735 pad_slot_id=PAD_SLOT_ID,
736 BLOCK_SIZE=BLOCK_SIZE,
737 )
738
739 # Reshape output and return
740 output = rearrange(output, "b h n d -> b n (h d)")
741 return output.squeeze(1).contiguous()