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 10, 01:30 UTC
[Performance] Improve Triton prefill attention kernel's performance (#32403)
Signed-off-by: Isotr0py <mozf@mail2.sysu.edu.cn>
Isotr0py committed
Jan 18, 2026 at 12:19 UTC
8cc26acd8b77c8af12d2fbf1f5bd38d09861928c
3 files changed
+33
-48
tests/models/language/pooling/test_token_classification.py
+2
-2
@@ -46,7 +46,7 @@ def test_bert_models(
46
for hf_output, vllm_output in zip(hf_outputs, vllm_outputs):
47
hf_output = hf_output.detach().clone().cpu().float()
48
vllm_output = vllm_output.detach().clone().cpu().float()
49
- torch.testing.assert_close(hf_output, vllm_output, atol=1.2e-2, rtol=1e-3)
49
+ torch.testing.assert_close(hf_output, vllm_output, atol=3.2e-2, rtol=1e-3)
50
51
52
@pytest.mark.parametrize("model", ["disham993/electrical-ner-ModernBERT-base"])
@@ -86,7 +86,7 @@ def test_modernbert_models(
86
for hf_output, vllm_output in zip(hf_outputs, vllm_outputs):
87
hf_output = hf_output.detach().clone().cpu().float()
88
vllm_output = vllm_output.detach().clone().cpu().float()
89
- torch.testing.assert_close(hf_output, vllm_output, atol=1.2e-2, rtol=1e-3)
89
+ torch.testing.assert_close(hf_output, vllm_output, atol=3.2e-2, rtol=1e-3)
90
91
92
@pytest.mark.parametrize("model", ["bd2lcco/Qwen3-0.6B-finetuned"])
vllm/utils/math_utils.py
+4
@@ -2,6 +2,10 @@
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
"""Math utility functions for vLLM."""
4
5
+# Approximate value of 1/ln(2), used for log/exp base conversion
6
+# Best FP32 approximation: 1.4426950216 (hex 0x3FB8AA3B)
7
+RCP_LN2 = 1.4426950216
8
+
9
10
def cdiv(a: int, b: int) -> int:
11
"""Ceiling division."""
vllm/v1/attention/ops/triton_prefill_attention.py
+27
-46
@@ -30,6 +30,7 @@ import torch
30
31
from vllm.platforms import current_platform
32
from vllm.triton_utils import tl, triton
33
+from vllm.utils.math_utils import RCP_LN2
34
35
36
@triton.jit
@@ -110,15 +111,7 @@ def _fwd_kernel(
111
end_n_limit = block_mask * end_n
112
113
for start_n in range(start_n_limit, end_n_limit, BLOCK_N):
113
- start_n = tl.multiple_of(start_n, BLOCK_N)
114
- # -- compute qk ----
115
- k = tl.load(
116
- k_ptrs + (cur_batch_in_all_start_index + start_n) * stride_kbs,
117
- mask=((start_n + offs_n[None, :]) < cur_batch_seq_len) & (mask_d[:, None]),
118
- other=0.0,
119
- )
120
-
121
- # Apply attention mask (causal + bidirectional sliding window)
114
+ # -- prepare attention mask ----
115
# Position indices in the sequence
116
pos_q = offs_m[:, None] # Query positions [BLOCK_M, 1]
117
pos_k = start_n + offs_n[None, :] # Key positions [1, BLOCK_N]
@@ -141,53 +134,38 @@ def _fwd_kernel(
134
if sliding_mask_k is not None:
135
mask &= sliding_mask_k
136
144
- qk = tl.zeros([BLOCK_M, BLOCK_N], dtype=tl.float32)
145
- qk += tl.where(mask, 0, float("-inf"))
146
- qk += tl.dot(q, k)
147
- qk *= sm_scale
148
-
149
- # -- compute m_ij, p, l_ij
150
- m_ij = tl.max(qk, 1)
151
- # For sliding window there's a chance the max is -inf due to masking of
152
- # the entire row. In this case we need to set m_j 0 to avoid NaN
153
- m_ij_valid_mask = m_ij > float("-inf")
154
- m_ij_masked = tl.where(m_ij_valid_mask, m_ij, 0.0)
155
- # -- compute p and l_ij --
156
- p = tl.exp(qk - m_ij_masked[:, None])
137
+ start_n = tl.multiple_of(start_n, BLOCK_N)
138
+ # -- compute qk ----
139
+ k = tl.load(
140
+ k_ptrs + (cur_batch_in_all_start_index + start_n) * stride_kbs,
141
+ mask=(pos_k < cur_batch_seq_len) & (mask_d[:, None]),
142
+ other=0.0,
143
+ )
144
+
145
+ qk = tl.dot(q, k)
146
+ qk = tl.where(mask, qk * sm_scale, -1.0e8)
147
+ m_ij = tl.maximum(m_i, tl.max(qk, 1))
148
+ qk -= m_ij[:, None]
149
+ p = tl.math.exp2(qk)
150
l_ij = tl.sum(p, 1)
151
+
152
# -- update m_i and l_i
159
- m_i_new = tl.maximum(m_i, m_ij)
160
- m_i_new_mask = m_i_new > float("-inf")
161
- alpha = tl.exp(m_i - m_i_new)
162
- beta = tl.exp(m_ij - m_i_new)
163
- # mask alpha and beta for sliding window
164
- alpha = tl.where(m_i_new_mask, alpha, 1.0)
165
- beta = tl.where(m_i_new_mask, beta, 0.0)
166
- l_i_new = alpha * l_i + beta * l_ij
153
+ alpha = tl.math.exp2(m_i - m_ij)
154
+ l_i = l_i * alpha + l_ij
155
# -- update output accumulator --
168
- # scale p
169
- # For sliding window there's a chance the l_i_new is 0 due to masking
170
- # the entire row. We need to set l_i_new 1 to avoid zero division
171
- l_i_new_mask = (l_i_new != 0.0) & (m_i_new_mask > float("-inf"))
172
- l_i_new_safe = tl.where(l_i_new_mask, l_i_new, 1.0)
173
- p_scale = beta / l_i_new_safe
174
- p = p * p_scale[:, None]
175
- # scale acc
176
- acc_scale = l_i / l_i_new_safe * alpha
177
- acc = acc * acc_scale[:, None]
156
+ acc = acc * alpha[:, None]
157
# update acc
158
v = tl.load(
159
v_ptrs + (cur_batch_in_all_start_index + start_n) * stride_vbs,
160
mask=((start_n + offs_n[:, None]) < cur_batch_seq_len) & (mask_d[None, :]),
161
other=0.0,
162
)
184
-
163
p = p.to(v.dtype)
186
- acc += tl.dot(p, v)
187
- # update m_i and l_i
188
- l_i = l_i_new
189
- m_i = m_i_new
190
- # initialize pointers to output
164
+ acc = tl.dot(p, v, acc)
165
+ # update m_i
166
+ m_i = m_ij
167
+
168
+ acc = acc / l_i[:, None]
169
off_o = (
170
(cur_batch_in_all_start_index + offs_m[:, None]) * stride_obs
171
+ cur_head * stride_oh
@@ -234,6 +212,9 @@ def context_attention_fwd(
212
Lq, Lk, _ = q.shape[-1], k.shape[-1], v.shape[-1]
213
214
sm_scale = 1.0 / (Lq**0.5) if softmax_scale is None else softmax_scale
215
+ # rescale with 1/ln(2) for triton exp2
216
+ sm_scale *= RCP_LN2
217
+
218
batch, head = b_seq_len.shape[0], q.shape[1]
219
kv_group_num = q.shape[1] // k.shape[1]
220