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] Make condition in triton kernel constexpr (#22370)
Signed-off-by: Gregory Shtrasberg <Gregory.Shtrasberg@amd.com>
Gregory Shtrasberg committed
Aug 6, 2025 at 13:00 UTC
2435ea7ed5c3a7d058cc6f6d649316e96976acaa
2 files changed
+6
-2
vllm/attention/ops/chunked_prefill_paged_decode.py
+3
-1
@@ -60,6 +60,7 @@ def kernel_paged_attention_2d(
60
stride_v_cache_3: tl.int64, # int
61
filter_by_query_len: tl.constexpr, # bool
62
query_start_len_ptr, # [num_seqs+1]
63
+ USE_SINKS: tl.constexpr, # bool
64
):
65
seq_idx = tl.program_id(0)
66
kv_head_idx = tl.program_id(1)
@@ -96,7 +97,7 @@ def kernel_paged_attention_2d(
97
98
block_table_offset = seq_idx * block_table_stride
99
99
- if sink_ptr is None:
100
+ if not USE_SINKS:
101
M = tl.full([num_queries_per_kv_padded],
102
float("-inf"),
103
dtype=tl.float32)
@@ -386,4 +387,5 @@ def chunked_prefill_paged_decode(
387
stride_v_cache_3=value_cache.stride(3),
388
filter_by_query_len=True,
389
query_start_len_ptr=query_start_loc,
390
+ USE_SINKS=sinks is not None,
391
)
vllm/attention/ops/prefix_prefill.py
+3
-1
@@ -81,6 +81,7 @@ def _fwd_kernel(Q,
81
num_unroll_cache: tl.constexpr,
82
num_unroll_request: tl.constexpr,
83
SKIP_DECODE: tl.constexpr,
84
+ USE_SINKS: tl.constexpr,
85
MAX_Q_LEN: tl.constexpr = 0,
86
MAX_CTX_LEN: tl.constexpr = 0):
87
@@ -127,7 +128,7 @@ def _fwd_kernel(Q,
128
other=0.0) # [M,D]
129
130
# initialize pointer to m and l
130
- if sink_ptr is None:
131
+ if not USE_SINKS:
132
m_i = tl.full([BLOCK_M], float("-inf"), dtype=tl.float32)
133
else:
134
m_i = tl.load(
@@ -910,5 +911,6 @@ def context_attention_fwd(q,
911
num_unroll_request=1,
912
num_warps=4,
913
num_stages=1,
914
+ USE_SINKS=sinks is not None,
915
**extra_kargs)
916
return