@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / f3b2a819f7
Fork 0

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 2, 14:45 UTC

[Perf][KDA] Fuse gate softplus, chunk-local cumsum, and RCP_LN2 scaling (#43667)

Signed-off-by: haojiangzheng <justineric096@gmail.com> Co-authored-by: haojiangzheng <justineric096@gmail.com>

zexplorerhj committed May 28, 2026 at 21:47 UTC f3b2a819f7f4ee05e0ac46ab72ca1e9f41f27bd5
3 files changed +366 -26
tests/kernels/test_kda.py
+70 -1
@@ -10,7 +10,11 @@ import pytest
10 import torch
11 import torch.nn.functional as F
12
13 -from vllm.model_executor.layers.fla.ops.kda import chunk_kda
13 +from vllm.model_executor.layers.fla.ops.kda import (
14 + chunk_kda,
15 + chunk_kda_with_fused_gate,
16 + fused_kda_gate,
17 +)
18 from vllm.model_executor.layers.fla.ops.l2norm import l2norm_fwd
19
20 DEVICE = "cuda"
@@ -155,3 +159,68 @@ def test_chunk_kda(
159 assert not torch.isnan(tri_ht).any(), "Triton output ht contains NaN"
160 assert_close("o", ref_o, tri_o, 0.005)
161 assert_close("ht", ref_ht, tri_ht.transpose(-1, -2).contiguous(), 0.005)
162 +
163 +
164 +@pytest.mark.parametrize(
165 + ("cu_seqlens", "dtype"),
166 + [
167 + ([0, 64], torch.float16),
168 + ([0, 15, 100, 300], torch.bfloat16),
169 + ],
170 +)
171 +@torch.inference_mode()
172 +def test_chunk_kda_fused_gate_cumsum_matches_unfused(
173 + cu_seqlens: list[int],
174 + dtype: torch.dtype,
175 +):
176 + H, D = 8, 64
177 + T = cu_seqlens[-1]
178 + N = len(cu_seqlens) - 1
179 + torch.manual_seed(123)
180 +
181 + cu_seqlens_t = torch.tensor(cu_seqlens, dtype=torch.int32, device=DEVICE)
182 + q = torch.randn(1, T, H, D, dtype=dtype, device=DEVICE)
183 + k = torch.randn(1, T, H, D, dtype=dtype, device=DEVICE)
184 + v = torch.randn(1, T, H, D, dtype=dtype, device=DEVICE)
185 + raw_g = torch.randn(1, T, H, D, dtype=dtype, device=DEVICE)
186 + beta = torch.rand(1, T, H, dtype=dtype, device=DEVICE).sigmoid()
187 + A_log = (torch.randn(H, dtype=torch.float32, device=DEVICE) * 0.5).contiguous()
188 + dt_bias = (
189 + torch.randn(H * D, dtype=torch.float32, device=DEVICE) * 0.1
190 + ).contiguous()
191 + h0 = torch.randn(N, H, D, D, dtype=torch.float32, device=DEVICE)
192 + initial_state = h0.transpose(-1, -2).contiguous()
193 +
194 + gate = fused_kda_gate(
195 + raw_g.reshape(T, H * D),
196 + A_log,
197 + D,
198 + g_bias=dt_bias,
199 + ).unsqueeze(0)
200 + old_o, old_ht = chunk_kda(
201 + q=q.clone(),
202 + k=k.clone(),
203 + v=v.clone(),
204 + g=gate,
205 + beta=beta.clone(),
206 + initial_state=initial_state.clone(),
207 + output_final_state=True,
208 + cu_seqlens=cu_seqlens_t,
209 + use_qk_l2norm_in_kernel=True,
210 + )
211 + new_o, new_ht = chunk_kda_with_fused_gate(
212 + q=q.clone(),
213 + k=k.clone(),
214 + v=v.clone(),
215 + raw_g=raw_g,
216 + beta=beta.clone(),
217 + A_log=A_log,
218 + g_bias=dt_bias,
219 + initial_state=initial_state.clone(),
220 + output_final_state=True,
221 + cu_seqlens=cu_seqlens_t,
222 + use_qk_l2norm_in_kernel=True,
223 + )
224 +
225 + assert_close("o", old_o, new_o, 1e-3, err_atol=1e-3)
226 + assert_close("ht", old_ht, new_ht, 1e-3, err_atol=1e-3)
vllm/model_executor/layers/fla/ops/kda.py
+282 -18
@@ -721,6 +721,7 @@ def chunk_kda_scaled_dot_kkt_fwd(
721 beta: torch.Tensor | None = None,
722 scale: float | None = None,
723 cu_seqlens: torch.Tensor | None = None,
724 + chunk_indices: torch.Tensor | None = None,
725 chunk_size: int = FLA_CHUNK_SIZE,
726 output_dtype: torch.dtype = torch.float32,
727 ) -> tuple[torch.Tensor, torch.Tensor]:
@@ -748,9 +749,8 @@ def chunk_kda_scaled_dot_kkt_fwd(
749 B, T, H, K = k.shape
750 assert K <= 256
751 BT = chunk_size
751 - chunk_indices = (
752 - prepare_chunk_indices(cu_seqlens, BT) if cu_seqlens is not None else None
753 - )
752 + if chunk_indices is None and cu_seqlens is not None:
753 + chunk_indices = prepare_chunk_indices(cu_seqlens, BT)
754 NT = cdiv(T, BT) if cu_seqlens is None else len(chunk_indices)
755
756 BC = min(16, BT)
@@ -965,15 +965,15 @@ def recompute_w_u_fwd(
965 q: torch.Tensor | None = None,
966 gk: torch.Tensor | None = None,
967 cu_seqlens: torch.Tensor | None = None,
968 + chunk_indices: torch.Tensor | None = None,
969 ) -> tuple[torch.Tensor, torch.Tensor]:
970 B, T, H, K, V = *k.shape, v.shape[-1]
971 BT = A.shape[-1]
972 BK = 64
973 BV = 64
974
974 - chunk_indices = (
975 - prepare_chunk_indices(cu_seqlens, BT) if cu_seqlens is not None else None
976 - )
975 + if chunk_indices is None and cu_seqlens is not None:
976 + chunk_indices = prepare_chunk_indices(cu_seqlens, BT)
977 NT = cdiv(T, BT) if cu_seqlens is None else len(chunk_indices)
978
979 w = torch.empty_like(k)
@@ -1132,16 +1132,14 @@ def chunk_gla_fwd_o_gk(
1132 o: torch.Tensor,
1133 scale: float,
1134 cu_seqlens: torch.Tensor | None = None,
1135 - chunk_size: int = 64,
1135 + chunk_indices: torch.Tensor | None = None,
1136 + chunk_size: int = FLA_CHUNK_SIZE,
1137 ):
1138 B, T, H, K, V = *q.shape, v.shape[-1]
1139 BT = chunk_size
1140
1140 - chunk_indices = (
1141 - prepare_chunk_indices(cu_seqlens, chunk_size)
1142 - if cu_seqlens is not None
1143 - else None
1144 - )
1141 + if chunk_indices is None and cu_seqlens is not None:
1142 + chunk_indices = prepare_chunk_indices(cu_seqlens, chunk_size)
1143 NT = cdiv(T, BT) if cu_seqlens is None else len(chunk_indices)
1144
1145 def grid(meta):
@@ -1166,7 +1164,146 @@ def chunk_gla_fwd_o_gk(
1164 return o
1165
1166
1169 -def chunk_kda_fwd(
1167 +@triton.heuristics(
1168 + {
1169 + "HAS_BIAS": lambda args: args["g_bias"] is not None,
1170 + "IS_VARLEN": lambda args: args["cu_seqlens"] is not None,
1171 + }
1172 +)
1173 +@triton.autotune(
1174 + configs=[
1175 + triton.Config({"BD": BD}, num_warps=num_warps)
1176 + for BD in [32, 64]
1177 + for num_warps in [2, 4, 8]
1178 + ],
1179 + key=["H", "D", "BT", "IS_VARLEN"],
1180 +)
1181 +@triton.jit(do_not_specialize=["T"])
1182 +def kda_gate_cumsum_fwd_kernel(
1183 + g,
1184 + A,
1185 + y,
1186 + g_bias,
1187 + cu_seqlens,
1188 + chunk_indices,
1189 + cumsum_scale,
1190 + beta,
1191 + threshold,
1192 + T,
1193 + H: tl.constexpr,
1194 + D: tl.constexpr,
1195 + BT: tl.constexpr,
1196 + BD: tl.constexpr,
1197 + HAS_BIAS: tl.constexpr,
1198 + IS_VARLEN: tl.constexpr,
1199 +):
1200 + i_d, i_t, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2)
1201 + i_b, i_h = i_bh // H, i_bh % H
1202 + if IS_VARLEN:
1203 + i_n, i_t = (
1204 + tl.load(chunk_indices + i_t * 2).to(tl.int32),
1205 + tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32),
1206 + )
1207 + bos, eos = (
1208 + tl.load(cu_seqlens + i_n).to(tl.int32),
1209 + tl.load(cu_seqlens + i_n + 1).to(tl.int32),
1210 + )
1211 + T = eos - bos
1212 + else:
1213 + bos = i_b * T
1214 +
1215 + p_g = tl.make_block_ptr(
1216 + g + (bos * H + i_h) * D,
1217 + (T, D),
1218 + (H * D, 1),
1219 + (i_t * BT, i_d * BD),
1220 + (BT, BD),
1221 + (1, 0),
1222 + )
1223 + p_y = tl.make_block_ptr(
1224 + y + (bos * H + i_h) * D,
1225 + (T, D),
1226 + (H * D, 1),
1227 + (i_t * BT, i_d * BD),
1228 + (BT, BD),
1229 + (1, 0),
1230 + )
1231 +
1232 + b_g = tl.load(p_g, boundary_check=(0, 1)).to(tl.float32)
1233 + if HAS_BIAS:
1234 + o_d = i_d * BD + tl.arange(0, BD)
1235 + b_bias = tl.load(g_bias + i_h * D + o_d, mask=o_d < D, other=0.0).to(tl.float32)
1236 + b_g = b_g + b_bias[None, :]
1237 +
1238 + b_a = -tl.exp(tl.load(A + i_h).to(tl.float32))
1239 + b_g_scaled = b_g * beta
1240 + b_softplus = tl.where(
1241 + b_g_scaled > threshold,
1242 + b_g,
1243 + (1.0 / beta) * log(1.0 + tl.exp(b_g_scaled)),
1244 + )
1245 + b_gate = b_a * b_softplus
1246 +
1247 + # Out-of-bounds rows (load returns 0, but softplus/bias can still make
1248 + # b_gate non-zero) participate in the dot product. They only contribute to
1249 + # out-of-bounds output rows, which are masked away by `boundary_check` on
1250 + # the store, so visible output matches unfused gate + chunk-local cumsum.
1251 + o_t = tl.arange(0, BT)
1252 + m_cumsum = tl.where(o_t[:, None] >= o_t[None, :], 1.0, 0.0)
1253 + b_y = tl.dot(m_cumsum, b_gate, allow_tf32=False) * cumsum_scale
1254 + tl.store(p_y, b_y.to(p_y.dtype.element_ty), boundary_check=(0, 1))
1255 +
1256 +
1257 +def fused_kda_gate_chunk_cumsum(
1258 + raw_g: torch.Tensor,
1259 + A_log: torch.Tensor,
1260 + g_bias: torch.Tensor | None = None,
1261 + beta: float = 1.0,
1262 + threshold: float = 20.0,
1263 + cu_seqlens: torch.Tensor | None = None,
1264 + chunk_indices: torch.Tensor | None = None,
1265 + chunk_size: int = FLA_CHUNK_SIZE,
1266 + output_dtype: torch.dtype | None = torch.float,
1267 +) -> torch.Tensor:
1268 + if cu_seqlens is not None:
1269 + assert raw_g.shape[0] == 1, (
1270 + "Only batch size 1 is supported when cu_seqlens are provided"
1271 + )
1272 + B, T, H, D = raw_g.shape
1273 + if chunk_indices is None and cu_seqlens is not None:
1274 + chunk_indices = prepare_chunk_indices(cu_seqlens, chunk_size)
1275 + NT = cdiv(T, chunk_size) if cu_seqlens is None else len(chunk_indices)
1276 +
1277 + A_log = A_log.reshape(-1)
1278 + if g_bias is not None:
1279 + g_bias = g_bias.reshape(-1)
1280 + y = torch.empty_like(raw_g, dtype=output_dtype or raw_g.dtype)
1281 +
1282 + def grid(meta):
1283 + return (cdiv(meta["D"], meta["BD"]), NT, B * H)
1284 +
1285 + kda_gate_cumsum_fwd_kernel[grid](
1286 + g=raw_g,
1287 + A=A_log,
1288 + y=y,
1289 + g_bias=g_bias,
1290 + cu_seqlens=cu_seqlens,
1291 + chunk_indices=chunk_indices,
1292 + # RCP_LN2 folds in the natural-log -> log2 conversion so downstream
1293 + # exp2-based kernels reproduce exp(g). Keep this in sync with the
1294 + # `use_exp2=True` path in `_chunk_kda_fwd_with_cumulative_g`.
1295 + cumsum_scale=RCP_LN2,
1296 + beta=beta,
1297 + threshold=threshold,
1298 + T=T,
1299 + H=H,
1300 + D=D,
1301 + BT=chunk_size,
1302 + )
1303 + return y
1304 +
1305 +
1306 +def _chunk_kda_fwd_with_cumulative_g(
1307 q: torch.Tensor,
1308 k: torch.Tensor,
1309 v: torch.Tensor,
@@ -1176,12 +1313,13 @@ def chunk_kda_fwd(
1313 initial_state: torch.Tensor,
1314 output_final_state: bool,
1315 cu_seqlens: torch.Tensor | None = None,
1316 + chunk_indices: torch.Tensor | None = None,
1317 + chunk_size: int = FLA_CHUNK_SIZE,
1318 ):
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
1319 + # `g` must already be chunk-local cumulatively-summed AND scaled by
1320 + # RCP_LN2 (so the downstream exp2-based kernels reproduce exp(g)).
1321 + # Use `chunk_kda_fwd` or `chunk_kda_with_fused_gate_fwd` instead of
1322 + # calling this helper directly unless that invariant is upheld.
1323 # the intra Aqk is kept in fp32
1324 # the computation has very marginal effect on the entire throughput
1325 A, Aqk = chunk_kda_scaled_dot_kkt_fwd(
@@ -1191,6 +1329,7 @@ def chunk_kda_fwd(
1329 beta=beta,
1330 scale=scale,
1331 cu_seqlens=cu_seqlens,
1332 + chunk_indices=chunk_indices,
1333 chunk_size=chunk_size,
1334 output_dtype=torch.float32,
1335 )
@@ -1202,6 +1341,7 @@ def chunk_kda_fwd(
1341 A=A,
1342 gk=g,
1343 cu_seqlens=cu_seqlens,
1344 + chunk_indices=chunk_indices,
1345 )
1346 del A
1347 h, v_new, final_state = chunk_gated_delta_rule_fwd_h(
@@ -1212,6 +1352,7 @@ def chunk_kda_fwd(
1352 initial_state=initial_state,
1353 output_final_state=output_final_state,
1354 cu_seqlens=cu_seqlens,
1355 + chunk_indices=chunk_indices,
1356 use_exp2=True,
1357 )
1358 del w, u, kg
@@ -1224,12 +1365,96 @@ def chunk_kda_fwd(
1365 o=v,
1366 scale=scale,
1367 cu_seqlens=cu_seqlens,
1368 + chunk_indices=chunk_indices,
1369 chunk_size=chunk_size,
1370 )
1371 del Aqk, v_new, h
1372 return o, final_state
1373
1374
1375 +def chunk_kda_fwd(
1376 + q: torch.Tensor,
1377 + k: torch.Tensor,
1378 + v: torch.Tensor,
1379 + g: torch.Tensor,
1380 + beta: torch.Tensor,
1381 + scale: float,
1382 + initial_state: torch.Tensor,
1383 + output_final_state: bool,
1384 + cu_seqlens: torch.Tensor | None = None,
1385 +):
1386 + chunk_size = FLA_CHUNK_SIZE
1387 + chunk_indices = (
1388 + prepare_chunk_indices(cu_seqlens, chunk_size)
1389 + if cu_seqlens is not None
1390 + else None
1391 + )
1392 + g = chunk_local_cumsum(
1393 + g,
1394 + chunk_size=chunk_size,
1395 + cu_seqlens=cu_seqlens,
1396 + chunk_indices=chunk_indices,
1397 + )
1398 + # KDA evaluates cumulative gate decays with exp2. Convert from natural-log
1399 + # space so exp(x) is preserved as exp2(x / ln(2)).
1400 + g = g * RCP_LN2
1401 + return _chunk_kda_fwd_with_cumulative_g(
1402 + q=q,
1403 + k=k,
1404 + v=v,
1405 + g=g,
1406 + beta=beta,
1407 + scale=scale,
1408 + initial_state=initial_state,
1409 + output_final_state=output_final_state,
1410 + cu_seqlens=cu_seqlens,
1411 + chunk_indices=chunk_indices,
1412 + chunk_size=chunk_size,
1413 + )
1414 +
1415 +
1416 +def chunk_kda_with_fused_gate_fwd(
1417 + q: torch.Tensor,
1418 + k: torch.Tensor,
1419 + v: torch.Tensor,
1420 + raw_g: torch.Tensor,
1421 + beta: torch.Tensor,
1422 + A_log: torch.Tensor,
1423 + g_bias: torch.Tensor | None,
1424 + scale: float,
1425 + initial_state: torch.Tensor,
1426 + output_final_state: bool,
1427 + cu_seqlens: torch.Tensor | None = None,
1428 +):
1429 + chunk_size = FLA_CHUNK_SIZE
1430 + chunk_indices = (
1431 + prepare_chunk_indices(cu_seqlens, chunk_size)
1432 + if cu_seqlens is not None
1433 + else None
1434 + )
1435 + g = fused_kda_gate_chunk_cumsum(
1436 + raw_g,
1437 + A_log=A_log,
1438 + g_bias=g_bias,
1439 + cu_seqlens=cu_seqlens,
1440 + chunk_indices=chunk_indices,
1441 + chunk_size=chunk_size,
1442 + )
1443 + return _chunk_kda_fwd_with_cumulative_g(
1444 + q=q,
1445 + k=k,
1446 + v=v,
1447 + g=g,
1448 + beta=beta,
1449 + scale=scale,
1450 + initial_state=initial_state,
1451 + output_final_state=output_final_state,
1452 + cu_seqlens=cu_seqlens,
1453 + chunk_indices=chunk_indices,
1454 + chunk_size=chunk_size,
1455 + )
1456 +
1457 +
1458 def chunk_kda(
1459 q: torch.Tensor,
1460 k: torch.Tensor,
@@ -1264,6 +1489,45 @@ def chunk_kda(
1489 return o, final_state
1490
1491
1492 +def chunk_kda_with_fused_gate(
1493 + q: torch.Tensor,
1494 + k: torch.Tensor,
1495 + v: torch.Tensor,
1496 + raw_g: torch.Tensor,
1497 + beta: torch.Tensor,
1498 + A_log: torch.Tensor,
1499 + g_bias: torch.Tensor | None,
1500 + scale: float | None = None,
1501 + initial_state: torch.Tensor | None = None,
1502 + output_final_state: bool = False,
1503 + use_qk_l2norm_in_kernel: bool = False,
1504 + cu_seqlens: torch.Tensor | None = None,
1505 + **kwargs,
1506 +):
1507 + """Run chunk KDA from raw gate projection using fused gate+cumsum."""
1508 + if scale is None:
1509 + scale = k.shape[-1] ** -0.5
1510 +
1511 + if use_qk_l2norm_in_kernel:
1512 + q = l2norm_fwd(q.contiguous())
1513 + k = l2norm_fwd(k.contiguous())
1514 +
1515 + o, final_state = chunk_kda_with_fused_gate_fwd(
1516 + q=q,
1517 + k=k,
1518 + v=v.contiguous(),
1519 + raw_g=raw_g.contiguous(),
1520 + beta=beta.contiguous(),
1521 + A_log=A_log,
1522 + g_bias=g_bias,
1523 + scale=scale,
1524 + initial_state=initial_state.contiguous() if initial_state is not None else None,
1525 + output_final_state=output_final_state,
1526 + cu_seqlens=cu_seqlens,
1527 + )
1528 + return o, final_state
1529 +
1530 +
1531 @triton.autotune(
1532 configs=[
1533 triton.Config({"BT": bt}, num_warps=nw, num_stages=ns)
vllm/model_executor/layers/mamba/gdn/kimi_gdn_linear_attn.py
+14 -7
@@ -21,7 +21,7 @@ from vllm.v1.attention.backends.gdn_attn import GDNAttentionMetadata
21
22 from ...fla.ops.kda import (
23 FusedRMSNormGated,
24 - chunk_kda,
24 + chunk_kda_with_fused_gate,
25 fused_kda_gate,
26 fused_recurrent_kda,
27 )
@@ -243,9 +243,8 @@ class KimiGatedDeltaNetAttention(GatedDeltaNetAttention):
243
244 beta = self.b_proj(hidden_states)[0].float().sigmoid()
245 g1 = self.f_b_proj(self.f_a_proj(hidden_states)[0])[0]
246 - g1 = fused_kda_gate(g1, self.A_log, self.head_dim, g_bias=self.dt_bias)
246 beta = beta.unsqueeze(0)
248 - g1 = g1.unsqueeze(0)
247 + g1 = rearrange(g1, "n (h d) -> 1 n h d", d=self.head_dim)
248
249 g_proj_states = self.g_b_proj(self.g_a_proj(hidden_states)[0])[0]
250 g2 = rearrange(g_proj_states, "... (h d) -> ... h d", d=self.head_dim)
@@ -298,8 +297,8 @@ class KimiGatedDeltaNetAttention(GatedDeltaNetAttention):
297 q_proj_states = q_proj_states[:num_actual_tokens]
298 k_proj_states = k_proj_states[:num_actual_tokens]
299 v_proj_states = v_proj_states[:num_actual_tokens]
301 - g1 = g1[:num_actual_tokens]
302 - beta = beta[:num_actual_tokens]
300 + g1 = g1[:, :num_actual_tokens]
301 + beta = beta[:, :num_actual_tokens]
302
303 (conv_state_q, conv_state_k, conv_state_v, recurrent_state) = constant_caches
304 # conv_state must be (..., dim, width-1) for the conv kernels.
@@ -401,12 +400,14 @@ class KimiGatedDeltaNetAttention(GatedDeltaNetAttention):
400 (
401 core_attn_out_non_spec,
402 last_recurrent_state,
404 - ) = chunk_kda(
403 + ) = chunk_kda_with_fused_gate(
404 q=q,
405 k=k,
406 v=v,
408 - g=g1,
407 + raw_g=g1,
408 beta=beta,
409 + A_log=self.A_log,
410 + g_bias=self.dt_bias,
411 initial_state=initial_state,
412 output_final_state=True,
413 use_qk_l2norm_in_kernel=True,
@@ -416,6 +417,12 @@ class KimiGatedDeltaNetAttention(GatedDeltaNetAttention):
417 recurrent_state[non_spec_state_indices_tensor] = last_recurrent_state
418 else:
419 assert non_spec_query_start_loc is not None
420 + g1 = fused_kda_gate(
421 + rearrange(g1, "1 n h d -> n (h d)"),
422 + self.A_log,
423 + self.head_dim,
424 + g_bias=self.dt_bias,
425 + ).unsqueeze(0)
426 (
427 core_attn_out_non_spec,
428 last_recurrent_state,