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 14, 05:30 UTC
[BugFix] Fix triton compile error in `kernel_unified_attention_2/3d` caused by attention sinks (#22368)
Signed-off-by: LucasWilkinson <lwilkinson@neuralmagic.com>
Lucas Wilkinson committed
Aug 6, 2025 at 12:47 UTC
4a6b72c2ab9848af31d51d3105a1992b7d5a01dc
1 file changed
+15
-8
vllm/attention/ops/triton_unified_attention.py
+15
-8
@@ -75,6 +75,7 @@ def kernel_unified_attention_2d(
75
USE_ALIBI_SLOPES: tl.constexpr, # bool
76
USE_QQ_BIAS: tl.constexpr, # bool
77
USE_SOFTCAP: tl.constexpr, # bool
78
+ USE_SINKS: tl.constexpr, # bool
79
SLIDING_WINDOW: tl.constexpr, # int
80
stride_k_cache_0: tl.int64, # int
81
stride_k_cache_1: tl.int64, # int
@@ -132,7 +133,7 @@ def kernel_unified_attention_2d(
133
134
block_table_offset = seq_idx * block_table_stride
135
135
- if sink_ptr is None:
136
+ if not USE_SINKS:
137
M = tl.full([BLOCK_M], float("-inf"), dtype=tl.float32)
138
else:
139
M = tl.load(
@@ -322,6 +323,7 @@ def kernel_unified_attention_3d(
323
USE_ALIBI_SLOPES: tl.constexpr, # bool
324
USE_QQ_BIAS: tl.constexpr, # bool
325
USE_SOFTCAP: tl.constexpr, # bool
326
+ USE_SINKS: tl.constexpr, # bool
327
SLIDING_WINDOW: tl.constexpr, # int
328
stride_k_cache_0: tl.int64, # int
329
stride_k_cache_1: tl.int64, # int
@@ -393,14 +395,17 @@ def kernel_unified_attention_3d(
395
396
block_table_offset = seq_idx * block_table_stride
397
396
- if sink_ptr is None or segm_idx != 0:
397
- M = tl.full([BLOCK_M], float("-inf"), dtype=tl.float32)
398
+ if USE_SINKS:
399
+ if segm_idx == 0:
400
+ M = tl.load(
401
+ sink_ptr + query_offset_1,
402
+ mask=query_mask_1,
403
+ other=float("-inf"),
404
+ ).to(dtype=tl.float32)
405
+ else:
406
+ M = tl.full([BLOCK_M], float("-inf"), dtype=tl.float32)
407
else:
399
- M = tl.load(
400
- sink_ptr + query_offset_1,
401
- mask=query_mask_1,
402
- other=float("-inf"),
403
- ).to(dtype=tl.float32)
408
+ M = tl.full([BLOCK_M], float("-inf"), dtype=tl.float32)
409
410
L = tl.full([BLOCK_M], 1.0, dtype=tl.float32)
411
acc = tl.zeros([BLOCK_M, HEAD_SIZE_PADDED], dtype=tl.float32)
@@ -716,6 +721,7 @@ def unified_attention(
721
USE_ALIBI_SLOPES=use_alibi_slopes,
722
USE_QQ_BIAS=use_qq_bias,
723
USE_SOFTCAP=(softcap > 0),
724
+ USE_SINKS=(sinks is not None),
725
SLIDING_WINDOW=(1 + window_size[0]),
726
stride_k_cache_0=k.stride(0),
727
stride_k_cache_1=k.stride(1),
@@ -787,6 +793,7 @@ def unified_attention(
793
USE_ALIBI_SLOPES=use_alibi_slopes,
794
USE_QQ_BIAS=use_qq_bias,
795
USE_SOFTCAP=(softcap > 0),
796
+ USE_SINKS=(sinks is not None),
797
SLIDING_WINDOW=(1 + window_size[0]),
798
stride_k_cache_0=k.stride(0),
799
stride_k_cache_1=k.stride(1),