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, 06:00 UTC
[Kernel] Add Split-KV Support to Unified Triton Attention Kernel (#19152)
Signed-off-by: Jan van Lunteren <jvl@zurich.ibm.com>
jvlunteren committed
Jun 17, 2025 at 12:45 UTC
ccd7c050898cfea4b6b0de16446dcc47fa02a399
1 file changed
+456
-52
vllm/attention/ops/triton_unified_attention.py
+456
-52
@@ -7,6 +7,7 @@
7
# - Chih-Chieh Yang <chih.chieh.yang@ibm.com>
8
# - Thomas Parnell <tpa@zurich.ibm.com>
9
10
+import torch
11
import triton
12
import triton.language as tl
13
@@ -28,6 +29,24 @@ def apply_softcap(S, x):
29
return x * (p1 - p2) / (p1 + p2)
30
31
32
+@triton.jit
33
+def find_seq_idx(query_start_len_ptr, target_idx, num_seqs,
34
+ BLOCK_Q: tl.constexpr, use_q_block_mode: tl.constexpr):
35
+ left: tl.int32 = 0
36
+ right = num_seqs
37
+ while left < right:
38
+ mid = (left + right) // 2
39
+ val = tl.load(query_start_len_ptr + mid)
40
+ mid_val = val // BLOCK_Q + mid if use_q_block_mode else val
41
+
42
+ if mid_val <= target_idx:
43
+ left = mid + 1
44
+ else:
45
+ right = mid
46
+
47
+ return left - 1
48
+
49
+
50
@triton.jit
51
def kernel_unified_attention_2d(
52
output_ptr, # [num_tokens, num_query_heads, head_size]
@@ -67,21 +86,12 @@ def kernel_unified_attention_2d(
86
num_seqs: tl.int32,
87
BLOCK_M: tl.constexpr, # int
88
):
70
-
89
q_block_global_idx = tl.program_id(0)
90
kv_head_idx = tl.program_id(1)
91
74
- left: tl.int32 = 0
75
- right = num_seqs
76
- while left < right:
77
- mid = (left + right) // 2
78
- mid_val = tl.load(query_start_len_ptr + mid) // BLOCK_Q + mid
79
- if mid_val <= q_block_global_idx:
80
- left = mid + 1
81
- else:
82
- right = mid
92
+ seq_idx = find_seq_idx(query_start_len_ptr, q_block_global_idx, num_seqs,
93
+ BLOCK_Q, True)
94
84
- seq_idx = left - 1
95
q_block_start_idx = tl.load(query_start_len_ptr +
96
seq_idx) // BLOCK_Q + seq_idx
97
@@ -242,6 +252,311 @@ def kernel_unified_attention_2d(
252
)
253
254
255
+@triton.jit
256
+def kernel_unified_attention_3d(
257
+ segm_output_ptr,
258
+ # [num_tokens, num_query_heads, num_segments, head_size]
259
+ segm_max_ptr, # [num_tokens, num_query_heads, num_segments]
260
+ segm_expsum_ptr, # [num_tokens, num_query_heads, num_segments]
261
+ query_ptr, # [num_tokens, num_query_heads, head_size]
262
+ key_cache_ptr, # [num_blks, num_kv_heads, head_size // x, blk_size, x]
263
+ value_cache_ptr, # [num_blks, num_kv_heads, head_size, blk_size]
264
+ block_tables_ptr, # [num_seqs, max_num_blocks_per_seq]
265
+ seq_lens_ptr, # [num_seqs]
266
+ alibi_slopes_ptr, # [num_query_heads]
267
+ scale, # float32
268
+ k_scale, # float32
269
+ v_scale, # float32
270
+ softcap, # float32
271
+ num_query_heads: tl.constexpr, # int
272
+ num_queries_per_kv: tl.constexpr, # int
273
+ block_table_stride: tl.int64, # int
274
+ query_stride_0: tl.int64, # int
275
+ query_stride_1: tl.int64, # int, should be equal to head_size
276
+ BLOCK_SIZE: tl.constexpr, # int
277
+ HEAD_SIZE: tl.constexpr, # int
278
+ HEAD_SIZE_PADDED: tl.constexpr, # int, must be power of 2
279
+ USE_ALIBI_SLOPES: tl.constexpr, # bool
280
+ USE_SOFTCAP: tl.constexpr, # bool
281
+ SLIDING_WINDOW: tl.constexpr, # int
282
+ stride_k_cache_0: tl.int64, # int
283
+ stride_k_cache_1: tl.int64, # int
284
+ stride_k_cache_2: tl.int64, # int
285
+ stride_k_cache_3: tl.constexpr, # int
286
+ stride_v_cache_0: tl.int64, # int
287
+ stride_v_cache_1: tl.int64, # int
288
+ stride_v_cache_2: tl.int64, # int
289
+ stride_v_cache_3: tl.constexpr, # int
290
+ query_start_len_ptr, # [num_seqs+1]
291
+ BLOCK_Q: tl.constexpr, # int
292
+ num_seqs: tl.int32,
293
+ BLOCK_M: tl.constexpr, # int
294
+ NUM_SEGMENTS_PER_SEQ: tl.constexpr, # int
295
+):
296
+ q_block_global_idx = tl.program_id(0)
297
+ kv_head_idx = tl.program_id(1)
298
+ segm_idx = tl.program_id(2)
299
+
300
+ seq_idx = find_seq_idx(query_start_len_ptr, q_block_global_idx, num_seqs,
301
+ BLOCK_Q, True)
302
+
303
+ q_block_start_idx = tl.load(query_start_len_ptr +
304
+ seq_idx) // BLOCK_Q + seq_idx
305
+
306
+ q_block_local_idx = q_block_global_idx - q_block_start_idx
307
+
308
+ cur_batch_in_all_start_index = tl.load(query_start_len_ptr + seq_idx)
309
+ cur_batch_in_all_stop_index = tl.load(query_start_len_ptr + seq_idx + 1)
310
+
311
+ cur_batch_query_len = cur_batch_in_all_stop_index \
312
+ - cur_batch_in_all_start_index
313
+
314
+ if q_block_local_idx * BLOCK_Q >= cur_batch_query_len:
315
+ return
316
+
317
+ # sequence len for this particular sequence
318
+ seq_len = tl.load(seq_lens_ptr + seq_idx)
319
+
320
+ # number of segments for this particular sequence
321
+ num_segments = NUM_SEGMENTS_PER_SEQ
322
+ blocks_per_segment = cdiv_fn(seq_len, num_segments * BLOCK_SIZE)
323
+
324
+ if segm_idx * blocks_per_segment * BLOCK_SIZE >= seq_len:
325
+ return
326
+
327
+ offs_m = tl.arange(0, BLOCK_M)
328
+ offs_d = tl.arange(0, HEAD_SIZE_PADDED)
329
+
330
+ query_pos = q_block_local_idx * BLOCK_Q + offs_m // num_queries_per_kv
331
+
332
+ query_offset_0 = cur_batch_in_all_start_index + query_pos
333
+ query_offset_1 = kv_head_idx * num_queries_per_kv + \
334
+ offs_m % num_queries_per_kv
335
+
336
+ query_offset = (query_offset_0[:, None] * query_stride_0 +
337
+ query_offset_1[:, None] * query_stride_1 + offs_d[None, :])
338
+
339
+ dim_mask = tl.where(offs_d < HEAD_SIZE, 1, 0).to(tl.int1)
340
+ query_mask_0 = tl.where(query_pos < cur_batch_query_len, 1, 0).to(tl.int1)
341
+ query_mask_1 = tl.where(query_offset_1 < num_query_heads, 1, 0).to(tl.int1)
342
+
343
+ # Q : (BLOCK_M, HEAD_SIZE_PADDED)
344
+ Q = tl.load(
345
+ query_ptr + query_offset,
346
+ mask=dim_mask[None, :] & query_mask_0[:, None] & query_mask_1[:, None],
347
+ other=0.0,
348
+ )
349
+
350
+ block_table_offset = seq_idx * block_table_stride
351
+
352
+ M = tl.full([BLOCK_M], float("-inf"), dtype=tl.float32)
353
+ L = tl.full([BLOCK_M], 1.0, dtype=tl.float32)
354
+ acc = tl.zeros([BLOCK_M, HEAD_SIZE_PADDED], dtype=tl.float32)
355
+
356
+ # context length for this particular sequences
357
+ context_len = seq_len - cur_batch_query_len
358
+
359
+ # alibi slope for this head
360
+ if USE_ALIBI_SLOPES:
361
+ alibi_slope = tl.load(alibi_slopes_ptr + query_offset_1,
362
+ mask=query_mask_1,
363
+ other=0.0)
364
+
365
+ num_blocks = cdiv_fn(seq_len, BLOCK_SIZE)
366
+
367
+ # iterate through tiles within current segment
368
+ for j in range(
369
+ segm_idx * blocks_per_segment,
370
+ min((segm_idx + 1) * blocks_per_segment, num_blocks),
371
+ ):
372
+ physical_block_idx = tl.load(block_tables_ptr + block_table_offset + j)
373
+
374
+ offs_n = tl.arange(0, BLOCK_SIZE)
375
+
376
+ v_offset = (physical_block_idx * stride_v_cache_0 +
377
+ kv_head_idx * stride_v_cache_2 +
378
+ offs_d[None, :] * stride_v_cache_3 +
379
+ offs_n[:, None] * stride_v_cache_1)
380
+
381
+ k_offset = (physical_block_idx * stride_k_cache_0 +
382
+ kv_head_idx * stride_k_cache_2 +
383
+ offs_d[:, None] * stride_k_cache_3 +
384
+ offs_n[None, :] * stride_k_cache_1)
385
+
386
+ # K : (HEAD_SIZE, BLOCK_SIZE)
387
+ K_load = tl.load(key_cache_ptr + k_offset,
388
+ mask=dim_mask[:, None],
389
+ other=0.0)
390
+
391
+ if K_load.dtype.is_fp8():
392
+ if Q.dtype.is_fp8():
393
+ K = K_load
394
+ else:
395
+ K = (K_load.to(tl.float32) * tl.load(k_scale)).to(Q.dtype)
396
+ else:
397
+ K = K_load
398
+
399
+ # V : (BLOCK_SIZE, HEAD_SIZE)
400
+ V_load = tl.load(value_cache_ptr + v_offset,
401
+ mask=dim_mask[None, :],
402
+ other=0.0)
403
+
404
+ if V_load.dtype.is_fp8():
405
+ if Q.dtype.is_fp8():
406
+ V = V_load
407
+ else:
408
+ V = (V_load.to(tl.float32) * tl.load(v_scale)).to(Q.dtype)
409
+ else:
410
+ V = V_load
411
+
412
+ seq_offset = j * BLOCK_SIZE + offs_n
413
+
414
+ seq_mask = seq_offset[None, :] < context_len + query_pos[:, None] + 1
415
+
416
+ # S : (BLOCK_M, BLOCK_SIZE)
417
+ S = tl.zeros(shape=(BLOCK_M, BLOCK_SIZE), dtype=tl.float32)
418
+
419
+ S += scale * tl.dot(Q, K)
420
+
421
+ if USE_SOFTCAP:
422
+ S = apply_softcap(S, softcap)
423
+
424
+ S = tl.where(query_mask_1[:, None] & query_mask_0[:, None] & seq_mask,
425
+ S, float("-inf"))
426
+
427
+ if SLIDING_WINDOW > 0:
428
+ S = tl.where((context_len + query_pos[:, None] - seq_offset)
429
+ < SLIDING_WINDOW, S, float("-inf"))
430
+
431
+ if USE_ALIBI_SLOPES:
432
+ S += alibi_slope[:, None] * (seq_offset - context_len)
433
+
434
+ # compute running maximum
435
+ # m_j : (BLOCK_M,)
436
+ m_j = tl.maximum(M, tl.max(S, axis=1))
437
+ # For sliding window there's a chance the max is -inf due to masking of
438
+ # the entire row. In this case we need to set m_j 0 to avoid NaN
439
+ m_j = tl.where(m_j > float("-inf"), m_j, 0.0)
440
+
441
+ # P : (BLOCK_M, BLOCK_SIZE,)
442
+ P = tl.exp(S - m_j[:, None])
443
+
444
+ # l_j : (BLOCK_M,)
445
+ l_j = tl.sum(P, axis=1)
446
+
447
+ # alpha : (BLOCK_M, )
448
+ alpha = tl.exp(M - m_j)
449
+
450
+ # acc : (BLOCK_M, HEAD_SIZE_PADDED)
451
+ acc = acc * alpha[:, None]
452
+
453
+ # update constants
454
+ L = L * alpha + l_j
455
+ M = m_j
456
+
457
+ # acc : (BLOCK_M, HEAD_SIZE_PADDED)
458
+ acc += tl.dot(P.to(V.dtype), V)
459
+
460
+ segm_output_offset = (
461
+ query_offset_0[:, None].to(tl.int64) *
462
+ (num_query_heads * NUM_SEGMENTS_PER_SEQ * HEAD_SIZE_PADDED) +
463
+ query_offset_1[:, None] * (NUM_SEGMENTS_PER_SEQ * HEAD_SIZE_PADDED) +
464
+ segm_idx * HEAD_SIZE_PADDED + tl.arange(0, HEAD_SIZE_PADDED)[None, :])
465
+ tl.store(
466
+ segm_output_ptr + segm_output_offset,
467
+ acc,
468
+ mask=dim_mask[None, :] & query_mask_0[:, None] & query_mask_1[:, None],
469
+ )
470
+ segm_offset = (query_offset_0.to(tl.int64) *
471
+ (num_query_heads * NUM_SEGMENTS_PER_SEQ) +
472
+ query_offset_1 * NUM_SEGMENTS_PER_SEQ + segm_idx)
473
+ tl.store(segm_max_ptr + segm_offset, M, mask=query_mask_0 & query_mask_1)
474
+ tl.store(segm_expsum_ptr + segm_offset,
475
+ L,
476
+ mask=query_mask_0 & query_mask_1)
477
+
478
+
479
+@triton.jit
480
+def reduce_segments(
481
+ output_ptr, # [num_tokens, num_query_heads, head_size]
482
+ segm_output_ptr,
483
+ #[num_tokens, num_query_heads, max_num_segments, head_size]
484
+ segm_max_ptr, # [num_tokens, num_query_heads, max_num_segments]
485
+ segm_expsum_ptr, # [num_tokens, num_query_heads, max_num_segments]
486
+ seq_lens_ptr, # [num_seqs]
487
+ num_seqs, # int
488
+ num_query_heads: tl.constexpr, # int
489
+ output_stride_0: tl.int64, # int
490
+ output_stride_1: tl.int64, # int, should be equal to head_size
491
+ block_table_stride: tl.int64, # int
492
+ BLOCK_SIZE: tl.constexpr, # int
493
+ HEAD_SIZE: tl.constexpr, # int, must be power of 2
494
+ HEAD_SIZE_PADDED: tl.constexpr, # int, must be power of 2
495
+ query_start_len_ptr, # [num_seqs+1]
496
+ BLOCK_Q: tl.constexpr, # int
497
+ NUM_SEGMENTS_PER_SEQ: tl.constexpr, # int
498
+):
499
+ query_token_idx = tl.program_id(0)
500
+ query_head_idx = tl.program_id(1)
501
+
502
+ seq_idx = find_seq_idx(query_start_len_ptr, query_token_idx, num_seqs,
503
+ BLOCK_Q, False)
504
+
505
+ # sequence len for this particular sequence
506
+ seq_len = tl.load(seq_lens_ptr + seq_idx)
507
+
508
+ # number of segments for this particular sequence
509
+ num_segments = NUM_SEGMENTS_PER_SEQ
510
+ blocks_per_segment = cdiv_fn(seq_len, num_segments * BLOCK_SIZE)
511
+
512
+ # create masks for subsequent loads
513
+ act_num_segments = cdiv_fn(seq_len, blocks_per_segment * BLOCK_SIZE)
514
+ segm_mask = tl.arange(0, NUM_SEGMENTS_PER_SEQ) < tl.full(
515
+ [NUM_SEGMENTS_PER_SEQ], act_num_segments, dtype=tl.int32)
516
+ dim_mask = tl.where(tl.arange(0, HEAD_SIZE_PADDED) < HEAD_SIZE, 1,
517
+ 0).to(tl.int1)
518
+
519
+ # load segment maxima
520
+ segm_offset = (query_token_idx.to(tl.int64) *
521
+ (num_query_heads * NUM_SEGMENTS_PER_SEQ) +
522
+ query_head_idx * NUM_SEGMENTS_PER_SEQ +
523
+ tl.arange(0, NUM_SEGMENTS_PER_SEQ))
524
+ segm_max = tl.load(segm_max_ptr + segm_offset,
525
+ mask=segm_mask,
526
+ other=float("-inf"))
527
+ overall_max = tl.max(segm_max)
528
+
529
+ # load and rescale segment exp sums
530
+ segm_expsum = tl.load(segm_expsum_ptr + segm_offset,
531
+ mask=segm_mask,
532
+ other=0.0)
533
+ segm_expsum = segm_expsum * tl.exp(segm_max - overall_max)
534
+ overall_expsum = tl.sum(segm_expsum)
535
+
536
+ # load, rescale, and add segment attention outputs
537
+ segm_output_offset = (
538
+ query_token_idx.to(tl.int64) *
539
+ (num_query_heads * NUM_SEGMENTS_PER_SEQ * HEAD_SIZE_PADDED) +
540
+ query_head_idx * (NUM_SEGMENTS_PER_SEQ * HEAD_SIZE_PADDED) +
541
+ tl.arange(0, NUM_SEGMENTS_PER_SEQ)[:, None] * HEAD_SIZE_PADDED +
542
+ tl.arange(0, HEAD_SIZE_PADDED)[None, :])
543
+ segm_output = tl.load(
544
+ segm_output_ptr + segm_output_offset,
545
+ mask=segm_mask[:, None] & dim_mask[None, :],
546
+ other=0.0,
547
+ )
548
+ segm_output *= tl.exp(segm_max - overall_max)[:, None]
549
+ acc_sum = tl.sum(segm_output, axis=0)
550
+ # safely divide by overall_expsum, returning 0.0 if overall_expsum is 0
551
+ acc = tl.where(overall_expsum == 0.0, 0.0, acc_sum / overall_expsum)
552
+
553
+ # write result
554
+ output_offset = (query_token_idx * output_stride_0 +
555
+ query_head_idx * output_stride_1 +
556
+ tl.arange(0, HEAD_SIZE_PADDED))
557
+ tl.store(output_ptr + output_offset, acc, mask=dim_mask)
558
+
559
+
560
def unified_attention(
561
q,
562
k,
@@ -291,44 +606,133 @@ def unified_attention(
606
# = floor(q.shape[0] / BLOCK_Q) + num_seqs
607
total_num_q_blocks = q.shape[0] // BLOCK_Q + num_seqs
608
294
- kernel_unified_attention_2d[(
295
- total_num_q_blocks,
296
- num_kv_heads,
297
- )](
298
- output_ptr=out,
299
- query_ptr=q,
300
- key_cache_ptr=k,
301
- value_cache_ptr=v,
302
- block_tables_ptr=block_table,
303
- seq_lens_ptr=seqused_k,
304
- alibi_slopes_ptr=alibi_slopes,
305
- scale=softmax_scale,
306
- k_scale=k_descale,
307
- v_scale=v_descale,
308
- softcap=softcap,
309
- num_query_heads=num_query_heads,
310
- num_queries_per_kv=num_queries_per_kv,
311
- block_table_stride=block_table.stride(0),
312
- query_stride_0=q.stride(0),
313
- query_stride_1=q.stride(1),
314
- output_stride_0=out.stride(0),
315
- output_stride_1=out.stride(1),
316
- BLOCK_SIZE=block_size,
317
- HEAD_SIZE=head_size,
318
- HEAD_SIZE_PADDED=triton.next_power_of_2(head_size),
319
- USE_ALIBI_SLOPES=use_alibi_slopes,
320
- USE_SOFTCAP=(softcap > 0),
321
- SLIDING_WINDOW=(1 + window_size[0]),
322
- stride_k_cache_0=k.stride(0),
323
- stride_k_cache_1=k.stride(1),
324
- stride_k_cache_2=k.stride(2),
325
- stride_k_cache_3=k.stride(3),
326
- stride_v_cache_0=v.stride(0),
327
- stride_v_cache_1=v.stride(1),
328
- stride_v_cache_2=v.stride(2),
329
- stride_v_cache_3=v.stride(3),
330
- query_start_len_ptr=cu_seqlens_q,
331
- BLOCK_Q=BLOCK_Q,
332
- num_seqs=num_seqs,
333
- BLOCK_M=BLOCK_M,
334
- )
609
+ # if batch contains a prefill
610
+ if max_seqlen_q > 1 or total_num_q_blocks * num_kv_heads > 128:
611
+ kernel_unified_attention_2d[(
612
+ total_num_q_blocks,
613
+ num_kv_heads,
614
+ )](
615
+ output_ptr=out,
616
+ query_ptr=q,
617
+ key_cache_ptr=k,
618
+ value_cache_ptr=v,
619
+ block_tables_ptr=block_table,
620
+ seq_lens_ptr=seqused_k,
621
+ alibi_slopes_ptr=alibi_slopes,
622
+ scale=softmax_scale,
623
+ k_scale=k_descale,
624
+ v_scale=v_descale,
625
+ softcap=softcap,
626
+ num_query_heads=num_query_heads,
627
+ num_queries_per_kv=num_queries_per_kv,
628
+ block_table_stride=block_table.stride(0),
629
+ query_stride_0=q.stride(0),
630
+ query_stride_1=q.stride(1),
631
+ output_stride_0=out.stride(0),
632
+ output_stride_1=out.stride(1),
633
+ BLOCK_SIZE=block_size,
634
+ HEAD_SIZE=head_size,
635
+ HEAD_SIZE_PADDED=triton.next_power_of_2(head_size),
636
+ USE_ALIBI_SLOPES=use_alibi_slopes,
637
+ USE_SOFTCAP=(softcap > 0),
638
+ SLIDING_WINDOW=(1 + window_size[0]),
639
+ stride_k_cache_0=k.stride(0),
640
+ stride_k_cache_1=k.stride(1),
641
+ stride_k_cache_2=k.stride(2),
642
+ stride_k_cache_3=k.stride(3),
643
+ stride_v_cache_0=v.stride(0),
644
+ stride_v_cache_1=v.stride(1),
645
+ stride_v_cache_2=v.stride(2),
646
+ stride_v_cache_3=v.stride(3),
647
+ query_start_len_ptr=cu_seqlens_q,
648
+ BLOCK_Q=BLOCK_Q,
649
+ num_seqs=num_seqs,
650
+ BLOCK_M=BLOCK_M,
651
+ )
652
+ else:
653
+ # for initial version, NUM_SEGMENTS = 16 is chosen as a default
654
+ # value that showed good performance in tests
655
+ NUM_SEGMENTS = 16
656
+
657
+ segm_output = torch.empty(
658
+ q.shape[0],
659
+ num_query_heads,
660
+ NUM_SEGMENTS,
661
+ triton.next_power_of_2(head_size),
662
+ dtype=torch.float32,
663
+ device=q.device,
664
+ )
665
+ segm_max = torch.empty(
666
+ q.shape[0],
667
+ num_query_heads,
668
+ NUM_SEGMENTS,
669
+ dtype=torch.float32,
670
+ device=q.device,
671
+ )
672
+ segm_expsum = torch.empty(
673
+ q.shape[0],
674
+ num_query_heads,
675
+ NUM_SEGMENTS,
676
+ dtype=torch.float32,
677
+ device=q.device,
678
+ )
679
+
680
+ kernel_unified_attention_3d[(
681
+ total_num_q_blocks, num_kv_heads, NUM_SEGMENTS)](
682
+ segm_output_ptr=segm_output,
683
+ segm_max_ptr=segm_max,
684
+ segm_expsum_ptr=segm_expsum,
685
+ query_ptr=q,
686
+ key_cache_ptr=k,
687
+ value_cache_ptr=v,
688
+ block_tables_ptr=block_table,
689
+ seq_lens_ptr=seqused_k,
690
+ alibi_slopes_ptr=alibi_slopes,
691
+ scale=softmax_scale,
692
+ k_scale=k_descale,
693
+ v_scale=v_descale,
694
+ softcap=softcap,
695
+ num_query_heads=num_query_heads,
696
+ num_queries_per_kv=num_queries_per_kv,
697
+ block_table_stride=block_table.stride(0),
698
+ query_stride_0=q.stride(0),
699
+ query_stride_1=q.stride(1),
700
+ BLOCK_SIZE=block_size,
701
+ HEAD_SIZE=head_size,
702
+ HEAD_SIZE_PADDED=triton.next_power_of_2(head_size),
703
+ USE_ALIBI_SLOPES=use_alibi_slopes,
704
+ USE_SOFTCAP=(softcap > 0),
705
+ SLIDING_WINDOW=(1 + window_size[0]),
706
+ stride_k_cache_0=k.stride(0),
707
+ stride_k_cache_1=k.stride(1),
708
+ stride_k_cache_2=k.stride(2),
709
+ stride_k_cache_3=k.stride(3),
710
+ stride_v_cache_0=v.stride(0),
711
+ stride_v_cache_1=v.stride(1),
712
+ stride_v_cache_2=v.stride(2),
713
+ stride_v_cache_3=v.stride(3),
714
+ query_start_len_ptr=cu_seqlens_q,
715
+ BLOCK_Q=BLOCK_Q,
716
+ num_seqs=num_seqs,
717
+ BLOCK_M=BLOCK_M,
718
+ NUM_SEGMENTS_PER_SEQ=NUM_SEGMENTS,
719
+ )
720
+
721
+ reduce_segments[(q.shape[0], num_query_heads)](
722
+ output_ptr=out,
723
+ segm_output_ptr=segm_output,
724
+ segm_max_ptr=segm_max,
725
+ segm_expsum_ptr=segm_expsum,
726
+ seq_lens_ptr=seqused_k,
727
+ num_seqs=num_seqs,
728
+ num_query_heads=num_query_heads,
729
+ output_stride_0=out.stride(0),
730
+ output_stride_1=out.stride(1),
731
+ block_table_stride=block_table.stride(0),
732
+ BLOCK_SIZE=block_size,
733
+ HEAD_SIZE=head_size,
734
+ HEAD_SIZE_PADDED=triton.next_power_of_2(head_size),
735
+ query_start_len_ptr=cu_seqlens_q,
736
+ BLOCK_Q=BLOCK_Q,
737
+ NUM_SEGMENTS_PER_SEQ=NUM_SEGMENTS,
738
+ )