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 3, 04:15 UTC
Update KDA chunk prefill decay to use exp2 semantics (#43195)
Signed-off-by: zexplorerhj <19794632+zexplorerhj@users.noreply.github.com> Co-authored-by: zexplorerhj <19794632+zexplorerhj@users.noreply.github.com>
zexplorerhj committed
May 21, 2026 at 16:16 UTC
b719b1635b4899e2372905def0badf96d4dd242a
3 files changed
+42
-17
vllm/model_executor/layers/fla/ops/chunk_delta_h.py
+26
-7
@@ -13,7 +13,7 @@ import torch
13
from vllm.triton_utils import tl, triton
14
15
from .index import prepare_chunk_indices, prepare_chunk_offsets
16
-from .op import exp
16
+from .op import exp, exp2
17
from .utils import FLA_CHUNK_SIZE, use_cuda_graph
18
19
NUM_WARPS = [2, 4, 8, 16]
@@ -65,6 +65,7 @@ def chunk_gated_delta_rule_fwd_kernel_h_blockdim64(
65
STORE_FINAL_STATE: tl.constexpr,
66
SAVE_NEW_VALUE: tl.constexpr,
67
IS_VARLEN: tl.constexpr,
68
+ USE_EXP2: tl.constexpr,
69
):
70
i_v, i_nh = tl.program_id(0), tl.program_id(1)
71
i_n, i_h = i_nh // H, i_nh % H
@@ -210,8 +211,12 @@ def chunk_gated_delta_rule_fwd_kernel_h_blockdim64(
211
g + bos * H + i_h, (T,), (H,), (i_t * BT,), (BT,), (0,)
212
)
213
b_g = tl.load(p_g, boundary_check=(0,))
213
- b_v = b_v * tl.where(m_t, exp(b_g_last - b_g), 0)[:, None]
214
- b_g_last = exp(b_g_last)
214
+ if USE_EXP2:
215
+ b_v = b_v * tl.where(m_t, exp2(b_g_last - b_g), 0)[:, None]
216
+ b_g_last = exp2(b_g_last)
217
+ else:
218
+ b_v = b_v * tl.where(m_t, exp(b_g_last - b_g), 0)[:, None]
219
+ b_g_last = exp(b_g_last)
220
b_h1 *= b_g_last
221
if K > 64:
222
b_h2 *= b_g_last
@@ -227,7 +232,10 @@ def chunk_gated_delta_rule_fwd_kernel_h_blockdim64(
232
mask=(o_k1 < K),
233
other=0.0,
234
)
230
- b_h1 *= exp(b_gk_last1)[None, :]
235
+ if USE_EXP2:
236
+ b_h1 *= exp2(b_gk_last1)[None, :]
237
+ else:
238
+ b_h1 *= exp(b_gk_last1)[None, :]
239
if K > 64:
240
o_k2 = 64 + o_k1
241
b_gk_last2 = tl.load(
@@ -235,7 +243,10 @@ def chunk_gated_delta_rule_fwd_kernel_h_blockdim64(
243
mask=(o_k2 < K),
244
other=0.0,
245
)
238
- b_h2 *= exp(b_gk_last2)[None, :]
246
+ if USE_EXP2:
247
+ b_h2 *= exp2(b_gk_last2)[None, :]
248
+ else:
249
+ b_h2 *= exp(b_gk_last2)[None, :]
250
if K > 128:
251
o_k3 = 128 + o_k1
252
b_gk_last3 = tl.load(
@@ -243,7 +254,10 @@ def chunk_gated_delta_rule_fwd_kernel_h_blockdim64(
254
mask=(o_k3 < K),
255
other=0.0,
256
)
246
- b_h3 *= exp(b_gk_last3)[None, :]
257
+ if USE_EXP2:
258
+ b_h3 *= exp2(b_gk_last3)[None, :]
259
+ else:
260
+ b_h3 *= exp(b_gk_last3)[None, :]
261
if K > 192:
262
o_k4 = 192 + o_k1
263
b_gk_last4 = tl.load(
@@ -251,7 +265,10 @@ def chunk_gated_delta_rule_fwd_kernel_h_blockdim64(
265
mask=(o_k4 < K),
266
other=0.0,
267
)
254
- b_h4 *= exp(b_gk_last4)[None, :]
268
+ if USE_EXP2:
269
+ b_h4 *= exp2(b_gk_last4)[None, :]
270
+ else:
271
+ b_h4 *= exp(b_gk_last4)[None, :]
272
b_v = b_v.to(k.dtype.element_ty)
273
274
p_k = tl.make_block_ptr(
@@ -311,6 +328,7 @@ def chunk_gated_delta_rule_fwd_h(
328
cu_seqlens: torch.Tensor | None = None,
329
chunk_indices: torch.Tensor | None = None,
330
chunk_offsets: torch.Tensor | None = None,
331
+ use_exp2: bool = False,
332
) -> tuple[torch.Tensor, torch.Tensor]:
333
# This kernel is slightly different from fla to support Q/K with different head numbers.
334
# In fla, Q/K always have the same head number, so Hg is always equal to H.
@@ -357,5 +375,6 @@ def chunk_gated_delta_rule_fwd_h(
375
K=K,
376
V=V,
377
BT=BT,
378
+ USE_EXP2=use_exp2,
379
)
380
return h, v_new, final_state
vllm/model_executor/layers/fla/ops/kda.py
+14
-10
@@ -14,14 +14,14 @@ import torch.nn as nn
14
15
from vllm.model_executor.custom_op import CustomOp
16
from vllm.triton_utils import tl, triton
17
-from vllm.utils.math_utils import cdiv, next_power_of_2
17
+from vllm.utils.math_utils import RCP_LN2, cdiv, next_power_of_2
18
19
from .chunk_delta_h import chunk_gated_delta_rule_fwd_h
20
from .cumsum import chunk_local_cumsum
21
from .fused_recurrent import fused_recurrent_gated_delta_rule_fwd_kernel
22
from .index import prepare_chunk_indices
23
from .l2norm import l2norm_fwd
24
-from .op import exp, log
24
+from .op import exp2, log
25
from .solve_tril import solve_tril
26
from .utils import FLA_CHUNK_SIZE, is_amd
27
@@ -594,16 +594,16 @@ def chunk_kda_scaled_dot_kkt_fwd_kernel_intra_sub_inter(
594
b_gn = tl.load(g + (i_t * BT + i_i * BC) * H * K + o_k, mask=m_k, other=0)
595
# [BC, BK]
596
b_g = tl.load(p_g, boundary_check=(0, 1))
597
- b_k = tl.load(p_k, boundary_check=(0, 1)) * exp(b_g - b_gn[None, :])
597
+ b_k = tl.load(p_k, boundary_check=(0, 1)) * exp2(b_g - b_gn[None, :])
598
# [BK, BC]
599
b_gk = tl.load(p_gk, boundary_check=(0, 1))
600
b_kt = tl.load(b_kt, boundary_check=(0, 1))
601
# [BC, BC]
602
- b_ktg = b_kt * exp(b_gn[:, None] - b_gk)
602
+ b_ktg = b_kt * exp2(b_gn[:, None] - b_gk)
603
b_A += tl.dot(b_k, b_ktg)
604
605
b_q = tl.load(p_q, boundary_check=(0, 1))
606
- b_qg = b_q * exp(b_g - b_gn[None, :]) * scale
606
+ b_qg = b_q * exp2(b_g - b_gn[None, :]) * scale
607
b_Aqk += tl.dot(b_qg, b_ktg)
608
609
b_A *= b_b[:, None]
@@ -703,7 +703,7 @@ def chunk_kda_scaled_dot_kkt_fwd_kernel_intra_sub_intra(
703
for j in range(0, min(BC, T - i_t * BT - i_i * BC)):
704
b_kt = tl.load(p_kt, mask=m_k, other=0).to(tl.float32)
705
b_gk = tl.load(p_gk, mask=m_k, other=0).to(tl.float32)
706
- b_ktg = b_kt[None, :] * exp(b_g - b_gk[None, :])
706
+ b_ktg = b_kt[None, :] * exp2(b_g - b_gk[None, :])
707
b_A = tl.sum(b_k * b_ktg, 1)
708
b_A = tl.where(o_i > j, b_A, 0.0)
709
b_Aqk = tl.sum(b_q * b_ktg, 1)
@@ -912,7 +912,7 @@ def recompute_w_u_fwd_kernel(
912
(1, 0),
913
)
914
b_gk = tl.load(p_gk, boundary_check=(0, 1))
915
- b_kb *= exp(b_gk)
915
+ b_kb *= exp2(b_gk)
916
if STORE_QG:
917
p_q = tl.make_block_ptr(
918
q + (bos * H + i_h) * K,
@@ -931,7 +931,7 @@ def recompute_w_u_fwd_kernel(
931
(1, 0),
932
)
933
b_q = tl.load(p_q, boundary_check=(0, 1))
934
- b_qg = b_q * exp(b_gk)
934
+ b_qg = b_q * exp2(b_gk)
935
tl.store(p_qg, b_qg.to(p_qg.dtype.element_ty), boundary_check=(0, 1))
936
if STORE_KG:
937
last_idx = min(i_t * BT + BT, T) - 1
@@ -941,7 +941,7 @@ def recompute_w_u_fwd_kernel(
941
b_gn = tl.load(
942
gk + ((bos + last_idx) * H + i_h) * K + o_k, mask=m_k, other=0.0
943
)
944
- b_kg = b_k * exp(b_gn - b_gk)
944
+ b_kg = b_k * exp2(b_gn - b_gk)
945
946
p_kg = tl.make_block_ptr(
947
kg + (bos * H + i_h) * K,
@@ -1089,7 +1089,7 @@ def chunk_gla_fwd_kernel_o(
1089
# [BT, BK]
1090
b_g = tl.load(p_g, boundary_check=(0, 1))
1091
# [BT, BK]
1092
- b_qg = (b_q * exp(b_g)).to(b_q.dtype)
1092
+ b_qg = (b_q * exp2(b_g)).to(b_q.dtype)
1093
# [BV, BK]
1094
b_h = tl.load(p_h, boundary_check=(0, 1))
1095
# [BT, BV]
@@ -1179,6 +1179,9 @@ def chunk_kda_fwd(
1179
):
1180
chunk_size = FLA_CHUNK_SIZE
1181
g = chunk_local_cumsum(g, chunk_size=chunk_size, cu_seqlens=cu_seqlens)
1182
+ # KDA evaluates cumulative gate decays with exp2. Convert from natural-log
1183
+ # space so exp(x) is preserved as exp2(x / ln(2)).
1184
+ g = g * RCP_LN2
1185
# the intra Aqk is kept in fp32
1186
# the computation has very marginal effect on the entire throughput
1187
A, Aqk = chunk_kda_scaled_dot_kkt_fwd(
@@ -1209,6 +1212,7 @@ def chunk_kda_fwd(
1212
initial_state=initial_state,
1213
output_final_state=output_final_state,
1214
cu_seqlens=cu_seqlens,
1215
+ use_exp2=True,
1216
)
1217
del w, u, kg
1218
o = chunk_gla_fwd_o_gk(
vllm/model_executor/layers/fla/ops/op.py
+2
@@ -15,10 +15,12 @@ from .utils import is_gather_supported
15
16
if os.environ.get("FLA_USE_FAST_OPS", "0") == "1":
17
exp = tldevice.fast_expf
18
+ exp2 = tl.exp2
19
log = tldevice.fast_logf
20
log2 = tldevice.fast_log2f
21
else:
22
exp = tl.exp
23
+ exp2 = tl.exp2
24
log = tl.log
25
log2 = tl.log2
26