[Kimi-K3] Add FlashInfer speculative KDA backend (#54255)
Signed-off-by: Duncan Moss <dmoss@nvidia.com> Co-authored-by: Roger Wang <hey@rogerw.io> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> Co-authored-by: Thien Tran <gau.nernst@yahoo.com.sg>
Duncan Moss committed
Sep 30, 2026 at 19:28 UTC
766b4ae54c491497306224ee4b316d9c2d57cd30
3 files changed
+383
-9
tests/models/kimi_k3/test_kda.py
+241
@@ -30,13 +30,17 @@ from vllm.models.kimi_k3.amd.ops.third_party.kda import (
30
)
31
from vllm.models.kimi_k3.nvidia import kda as nvidia_kda
32
from vllm.models.kimi_k3.nvidia.kda import (
33
+ KimiK3DeltaAttention,
34
_flashinfer_kda_prefill,
35
_flashkda_prefill,
36
is_flashinfer_fused_kda_decode_supported,
37
+ is_flashinfer_fused_kda_spec_decode_supported,
38
is_flashinfer_recurrent_kda_prefill_supported,
39
is_flashkda_supported,
40
is_fused_kda_decode_supported,
41
+ resolve_kda_spec_decode_backend,
42
)
43
+from vllm.models.kimi_k3.nvidia.kda_metadata import KimiK3KDAMetadata
44
from vllm.models.kimi_k3.nvidia.model import KimiLinearForCausalLM
45
from vllm.models.kimi_k3.nvidia.ops import recoverssm as recoverssm_ops
46
from vllm.models.kimi_k3.nvidia.ops.recoverssm import (
@@ -90,6 +94,25 @@ def test_kda_warmup_skips_missing_metadata(monkeypatch):
94
assert layer._forward(empty, empty, empty, empty, empty) is None
95
96
97
+def test_resolve_kda_spec_decode_backend(monkeypatch: pytest.MonkeyPatch):
98
+ monkeypatch.setattr(
99
+ "vllm.models.kimi_k3.nvidia.kda.is_flashinfer_fused_kda_spec_decode_supported",
100
+ lambda *args: True,
101
+ )
102
+ args = (12, 128, 4, 6, torch.bfloat16, torch.bfloat16, torch.float32, -5.0)
103
+ assert resolve_kda_spec_decode_backend("auto", *args, False) == "flashinfer"
104
+ assert resolve_kda_spec_decode_backend("native", *args, False) == "native"
105
+ assert resolve_kda_spec_decode_backend("flashinfer", *args, False) == "flashinfer"
106
+
107
+ monkeypatch.setattr(
108
+ "vllm.models.kimi_k3.nvidia.kda.is_flashinfer_fused_kda_spec_decode_supported",
109
+ lambda *args: False,
110
+ )
111
+ assert resolve_kda_spec_decode_backend("auto", *args, False) == "native"
112
+ with pytest.raises(RuntimeError, match="packed_fused_kda_decode"):
113
+ resolve_kda_spec_decode_backend("flashinfer", *args, False)
114
+
115
+
116
def test_kda_recoverssm_config_state_layout():
117
vllm_config = SimpleNamespace(
118
model_config=SimpleNamespace(
@@ -1182,6 +1205,224 @@ def test_fused_kda_decode_correctness(
1205
torch.testing.assert_close(state_actual, state_ref, atol=3e-2, rtol=3e-2)
1206
1207
1208
+@torch.inference_mode()
1209
+@pytest.mark.parametrize("state_dtype", [torch.float32, torch.bfloat16])
1210
+def test_flashinfer_fused_kda_spec_decode_integration(
1211
+ monkeypatch: pytest.MonkeyPatch,
1212
+ state_dtype: torch.dtype,
1213
+):
1214
+ if not is_flashinfer_fused_kda_spec_decode_supported(
1215
+ num_heads=12,
1216
+ head_dim=128,
1217
+ conv_width=4,
1218
+ num_spec=6,
1219
+ input_dtype=torch.bfloat16,
1220
+ conv_state_dtype=torch.bfloat16,
1221
+ recurrent_state_dtype=state_dtype,
1222
+ lower_bound=-5.0,
1223
+ use_recoverssm=False,
1224
+ ):
1225
+ pytest.skip("FlashInfer packed fused KDA decode is not supported")
1226
+
1227
+ generator = torch.Generator(device=DEVICE).manual_seed(20260820)
1228
+ H, D, W, T = 12, 128, 4, 7
1229
+ hidden_size = H * D
1230
+ num_rows = 10
1231
+ num_slots = 16
1232
+
1233
+ x = torch.randn(
1234
+ num_rows,
1235
+ 3 * hidden_size,
1236
+ dtype=torch.bfloat16,
1237
+ device=DEVICE,
1238
+ generator=generator,
1239
+ )
1240
+ conv_weight = 0.1 * torch.randn(
1241
+ 3 * hidden_size,
1242
+ W,
1243
+ dtype=torch.float32,
1244
+ device=DEVICE,
1245
+ generator=generator,
1246
+ )
1247
+ fused_weight = conv_weight.view(3, hidden_size, W).transpose(1, 2).contiguous()
1248
+ raw_gate = torch.randn(
1249
+ 1, num_rows, H, D, dtype=torch.bfloat16, device=DEVICE, generator=generator
1250
+ )
1251
+ raw_beta = torch.randn(
1252
+ 1, num_rows, H, dtype=torch.bfloat16, device=DEVICE, generator=generator
1253
+ )
1254
+ output_gate = torch.randn(
1255
+ num_rows, H, D, dtype=torch.bfloat16, device=DEVICE, generator=generator
1256
+ )
1257
+ A_log = torch.log(torch.arange(1, H + 1, dtype=torch.float32, device=DEVICE))
1258
+ dt_bias = torch.empty(hidden_size, dtype=torch.float32, device=DEVICE).uniform_(
1259
+ -7.0, -4.0, generator=generator
1260
+ )
1261
+ norm_weight = torch.randn(
1262
+ D, dtype=torch.float32, device=DEVICE, generator=generator
1263
+ )
1264
+ norm_eps = 1e-5
1265
+ state_indices = torch.tensor(
1266
+ [
1267
+ list(range(1, T + 1)),
1268
+ list(range(T + 1, 2 * T + 1)),
1269
+ [0] * T,
1270
+ ],
1271
+ dtype=torch.int32,
1272
+ device=DEVICE,
1273
+ )
1274
+ query_start_loc = torch.tensor([0, 7, 10, 10], dtype=torch.int32, device=DEVICE)
1275
+ num_accepted_tokens = torch.tensor([1, 4, 1], dtype=torch.int32, device=DEVICE)
1276
+ conv_seed = (
1277
+ 0.5
1278
+ * torch.randn(
1279
+ num_slots,
1280
+ T + W - 2,
1281
+ 3 * hidden_size,
1282
+ dtype=torch.bfloat16,
1283
+ device=DEVICE,
1284
+ generator=generator,
1285
+ )
1286
+ ).transpose(1, 2)
1287
+ state_seed = (
1288
+ 0.5
1289
+ * torch.randn(
1290
+ num_slots,
1291
+ H,
1292
+ D,
1293
+ D,
1294
+ dtype=torch.float32,
1295
+ device=DEVICE,
1296
+ generator=generator,
1297
+ )
1298
+ ).to(state_dtype)
1299
+ metadata = KimiK3KDAMetadata(
1300
+ num_prefills=0,
1301
+ num_prefill_tokens=0,
1302
+ num_decodes=0,
1303
+ num_decode_tokens=0,
1304
+ num_spec_decodes=2,
1305
+ num_spec_decode_tokens=num_rows,
1306
+ num_actual_tokens=num_rows,
1307
+ spec_query_start_loc=query_start_loc,
1308
+ spec_state_indices_tensor=state_indices,
1309
+ num_accepted_tokens=num_accepted_tokens,
1310
+ )
1311
+ monkeypatch.setattr(
1312
+ "vllm.models.kimi_k3.nvidia.kda.get_forward_context",
1313
+ lambda: SimpleNamespace(attn_metadata={"test.layer": metadata}),
1314
+ )
1315
+
1316
+ def clone_strided(tensor: torch.Tensor) -> torch.Tensor:
1317
+ clone = torch.empty_strided(
1318
+ tensor.shape, tensor.stride(), dtype=tensor.dtype, device=tensor.device
1319
+ )
1320
+ clone.copy_(tensor)
1321
+ return clone
1322
+
1323
+ def composed_reference(
1324
+ conv_state: torch.Tensor,
1325
+ recurrent_state: torch.Tensor,
1326
+ ) -> torch.Tensor:
1327
+ mixed_qkv = causal_conv1d_update(
1328
+ x,
1329
+ conv_state,
1330
+ conv_weight,
1331
+ bias=None,
1332
+ activation="silu",
1333
+ conv_state_indices=state_indices[:, 0],
1334
+ num_accepted_tokens=num_accepted_tokens,
1335
+ query_start_loc=query_start_loc,
1336
+ max_query_len=T,
1337
+ validate_data=False,
1338
+ out=torch.empty_like(x),
1339
+ )
1340
+ q, k, v = mixed_qkv.view(num_rows, 3, H, D).unbind(1)
1341
+ recurrent_out, _ = fused_recurrent_kda_nvidia(
1342
+ q=q.unsqueeze(0),
1343
+ k=k.unsqueeze(0),
1344
+ v=v.unsqueeze(0),
1345
+ raw_g=raw_gate,
1346
+ raw_beta=raw_beta,
1347
+ A_log=A_log,
1348
+ dt_bias=dt_bias,
1349
+ lower_bound=-5.0,
1350
+ initial_state=recurrent_state,
1351
+ cu_seqlens=query_start_loc,
1352
+ ssm_state_indices=state_indices,
1353
+ num_accepted_tokens=num_accepted_tokens,
1354
+ )
1355
+ recurrent_float = recurrent_out.float()
1356
+ return (
1357
+ recurrent_float
1358
+ * torch.rsqrt(
1359
+ recurrent_float.square().mean(dim=-1, keepdim=True) + norm_eps
1360
+ )
1361
+ * norm_weight
1362
+ * output_gate.float().sigmoid().unsqueeze(0)
1363
+ ).to(torch.bfloat16)
1364
+
1365
+ def run_flashinfer(
1366
+ conv_state: torch.Tensor,
1367
+ recurrent_state: torch.Tensor,
1368
+ output: torch.Tensor,
1369
+ ) -> torch.Tensor:
1370
+ layer = SimpleNamespace(
1371
+ prefix="test.layer",
1372
+ kv_cache=(conv_state.transpose(-1, -2), recurrent_state),
1373
+ kda_spec_decode_backend="flashinfer",
1374
+ decode_conv1d_weight=fused_weight,
1375
+ decode_norm_weight=norm_weight,
1376
+ gate_lower_bound=-5.0,
1377
+ A_log=A_log,
1378
+ dt_bias=dt_bias,
1379
+ o_norm=SimpleNamespace(eps=norm_eps),
1380
+ )
1381
+ KimiK3DeltaAttention._forward(
1382
+ layer,
1383
+ mixed_qkv=x,
1384
+ g1=raw_gate,
1385
+ g2=output_gate,
1386
+ beta=raw_beta,
1387
+ core_attn_out=output,
1388
+ )
1389
+ return output
1390
+
1391
+ expected_conv = clone_strided(conv_seed)
1392
+ expected_state = state_seed.clone()
1393
+ expected = composed_reference(expected_conv, expected_state)
1394
+ actual_conv = clone_strided(conv_seed)
1395
+ actual_state = state_seed.clone()
1396
+ output = torch.empty_like(expected)
1397
+ actual = run_flashinfer(actual_conv, actual_state, output)
1398
+
1399
+ assert actual.data_ptr() == output.data_ptr()
1400
+ torch.testing.assert_close(actual_conv, expected_conv, atol=0, rtol=0)
1401
+ torch.testing.assert_close(actual_state, expected_state, atol=2e-3, rtol=3e-2)
1402
+ torch.testing.assert_close(actual, expected, atol=2e-2, rtol=3e-2)
1403
+ assert torch.count_nonzero(actual[:, :7]) > 0
1404
+ torch.testing.assert_close(actual_state[0], state_seed[0], atol=0, rtol=0)
1405
+
1406
+ graph = torch.cuda.CUDAGraph()
1407
+ with torch.cuda.graph(graph):
1408
+ captured = run_flashinfer(actual_conv, actual_state, output)
1409
+
1410
+ actual_conv.copy_(conv_seed)
1411
+ actual_state.copy_(state_seed)
1412
+ query_start_loc.copy_(torch.tensor([0, 6, 10, 10], device=DEVICE))
1413
+ num_accepted_tokens.copy_(torch.tensor([3, 2, 1], device=DEVICE))
1414
+ expected_conv.copy_(conv_seed)
1415
+ expected_state.copy_(state_seed)
1416
+ expected = composed_reference(expected_conv, expected_state)
1417
+ graph.replay()
1418
+ torch.accelerator.synchronize()
1419
+
1420
+ assert captured.data_ptr() == output.data_ptr()
1421
+ torch.testing.assert_close(output, expected, atol=2e-2, rtol=3e-2)
1422
+ torch.testing.assert_close(actual_conv, expected_conv, atol=0, rtol=0)
1423
+ torch.testing.assert_close(actual_state, expected_state, atol=2e-3, rtol=3e-2)
1424
+
1425
+
1426
def test_fused_kda_decode_rejects_speculative_conv_state():
1427
assert not is_fused_kda_decode_supported(
1428
num_heads=12,
vllm/models/kimi_k3/nvidia/kda.py
+130
-9
@@ -55,8 +55,10 @@ from vllm.third_party.flash_linear_attention.ops.kda import FusedRMSNormGated
55
from vllm.transformers_utils.configs.kimi_linear import KimiLinearConfig
56
from vllm.utils.flashinfer import (
57
flashinfer_fused_kda_decode,
58
+ flashinfer_packed_fused_kda_decode,
59
flashinfer_recurrent_kda,
60
has_flashinfer_fused_kda_decode,
61
+ has_flashinfer_packed_fused_kda_decode,
62
has_flashinfer_recurrent_kda,
63
)
64
from vllm.v1.attention.backend import AttentionBackend
@@ -214,6 +216,73 @@ def is_flashinfer_fused_kda_decode_supported(
216
)
217
218
219
+def is_flashinfer_fused_kda_spec_decode_supported(
220
+ num_heads: int,
221
+ head_dim: int,
222
+ conv_width: int,
223
+ num_spec: int,
224
+ input_dtype: torch.dtype,
225
+ conv_state_dtype: torch.dtype,
226
+ recurrent_state_dtype: torch.dtype,
227
+ lower_bound: float | None,
228
+ use_recoverssm: bool,
229
+) -> bool:
230
+ return (
231
+ has_flashinfer_packed_fused_kda_decode()
232
+ and current_platform.is_device_capability_family(100)
233
+ and num_heads in (12, 24, 32, 48, 96)
234
+ and head_dim == 128
235
+ and conv_width == 4
236
+ and 1 <= num_spec <= 7
237
+ and input_dtype == torch.bfloat16
238
+ and conv_state_dtype == torch.bfloat16
239
+ and recurrent_state_dtype in (torch.float32, torch.bfloat16)
240
+ and lower_bound is not None
241
+ and lower_bound < 0
242
+ and not use_recoverssm
243
+ and not is_conv_state_dim_first()
244
+ )
245
+
246
+
247
+def resolve_kda_spec_decode_backend(
248
+ backend: str,
249
+ num_heads: int,
250
+ head_dim: int,
251
+ conv_width: int,
252
+ num_spec: int,
253
+ input_dtype: torch.dtype,
254
+ conv_state_dtype: torch.dtype,
255
+ recurrent_state_dtype: torch.dtype,
256
+ lower_bound: float | None,
257
+ use_recoverssm: bool,
258
+) -> str:
259
+ if backend not in ("auto", "native", "flashinfer"):
260
+ raise ValueError(f"Unsupported KDA speculative decode backend: {backend}")
261
+ supported = is_flashinfer_fused_kda_spec_decode_supported(
262
+ num_heads,
263
+ head_dim,
264
+ conv_width,
265
+ num_spec,
266
+ input_dtype,
267
+ conv_state_dtype,
268
+ recurrent_state_dtype,
269
+ lower_bound,
270
+ use_recoverssm,
271
+ )
272
+ if backend == "flashinfer" and not supported:
273
+ raise RuntimeError(
274
+ "FlashInfer packed fused KDA decode requires its "
275
+ "packed_fused_kda_decode API, CUDA SM10x, bfloat16 inputs and "
276
+ "convolution state, float32 or bfloat16 recurrent state, D=128, "
277
+ "W=4, a bounded gate, 1-7 speculative tokens, a supported head "
278
+ "count, the SD convolution-state layout, and RecoverSSM disabled."
279
+ )
280
+ if supported and backend != "native":
281
+ logger.info_once("Using FlashInfer packed fused KDA speculative decode.")
282
+ return "flashinfer"
283
+ return "native"
284
+
285
+
286
def resolve_kda_decode_backend(
287
backend: str,
288
num_heads: int,
@@ -537,6 +606,13 @@ class KimiK3DeltaAttention(GatedDeltaNetAttention):
606
self.projection_size = self.head_dim * self.num_heads
607
self.local_projection_size = divide(self.projection_size, self.tp_size)
608
self.conv_size = kda_config["short_conv_kernel_size"]
609
+ self.gate_lower_bound: float | None = kda_config.get("gate_lower_bound", None)
610
+ if self.gate_lower_bound is not None:
611
+ assert _KDA_GATE_LOGBOUND_MIN <= self.gate_lower_bound < 0, (
612
+ "KDA gate lower bound must be in "
613
+ f"[{_KDA_GATE_LOGBOUND_MIN}, 0). "
614
+ f"Got {self.gate_lower_bound}."
615
+ )
616
assert kda_config.get("use_full_rank_gate", False), (
617
"KimiK3DeltaAttention requires a full-rank gate"
618
)
@@ -618,8 +694,28 @@ class KimiK3DeltaAttention(GatedDeltaNetAttention):
694
conv_state_dtype,
695
recurrent_state_dtype,
696
)
697
+ spec_decode_backend = (
698
+ additional_config.get("kda_spec_decode_backend", "auto")
699
+ if isinstance(additional_config, dict)
700
+ else "auto"
701
+ )
702
+ self.kda_spec_decode_backend = resolve_kda_spec_decode_backend(
703
+ spec_decode_backend,
704
+ self.local_num_heads,
705
+ self.head_dim,
706
+ self.conv_size,
707
+ self.num_spec,
708
+ vllm_config.model_config.dtype,
709
+ conv_state_dtype,
710
+ recurrent_state_dtype,
711
+ self.gate_lower_bound,
712
+ self.use_recoverssm,
713
+ )
714
decode_conv1d_weight = None
622
- if self.kda_decode_backend != "triton":
715
+ if (
716
+ self.kda_decode_backend != "triton"
717
+ or self.kda_spec_decode_backend == "flashinfer"
718
+ ):
719
decode_conv1d_weight = torch.empty(
720
3,
721
self.conv_size,
@@ -648,14 +744,6 @@ class KimiK3DeltaAttention(GatedDeltaNetAttention):
744
)
745
set_weight_attrs(self.A_log, {"weight_loader": a_log_weight_loader(0)})
746
651
- self.gate_lower_bound: float | None = kda_config.get("gate_lower_bound", None)
652
- if self.gate_lower_bound is not None:
653
- assert _KDA_GATE_LOGBOUND_MIN <= self.gate_lower_bound < 0, (
654
- "KDA gate lower bound must be in "
655
- f"[{_KDA_GATE_LOGBOUND_MIN}, 0). "
656
- f"Got {self.gate_lower_bound}."
657
- )
658
-
747
backend = (
748
additional_config.get("kda_prefill_backend", "auto")
749
if isinstance(additional_config, dict)
@@ -860,6 +948,39 @@ class KimiK3DeltaAttention(GatedDeltaNetAttention):
948
if not is_conv_state_dim_first():
949
conv_state = conv_state.transpose(-1, -2)
950
951
+ if (
952
+ self.kda_spec_decode_backend == "flashinfer"
953
+ and has_spec_decode
954
+ and m.num_prefills == 0
955
+ and m.num_decodes == 0
956
+ ):
957
+ assert self.decode_conv1d_weight is not None
958
+ assert self.decode_norm_weight is not None
959
+ assert spec_state_indices_tensor is not None
960
+ assert spec_query_start_loc is not None
961
+ assert num_accepted_tokens is not None
962
+ assert self.gate_lower_bound is not None
963
+ # "Packed" is FlashInfer's name for the varlen T>1 speculative path.
964
+ flashinfer_packed_fused_kda_decode(
965
+ x=mixed_qkv,
966
+ weight=self.decode_conv1d_weight,
967
+ conv_state=conv_state,
968
+ raw_gate=g1,
969
+ raw_beta=beta,
970
+ A_log=self.A_log,
971
+ dt_bias=self.dt_bias,
972
+ state_indices=spec_state_indices_tensor,
973
+ query_start_loc=spec_query_start_loc,
974
+ num_accepted_tokens=num_accepted_tokens,
975
+ state=recurrent_state,
976
+ output_gate=g2[:num_actual_tokens],
977
+ norm_weight=self.decode_norm_weight,
978
+ lower_bound=self.gate_lower_bound,
979
+ norm_eps=self.o_norm.eps,
980
+ output=core_attn_out[:, :num_actual_tokens],
981
+ )
982
+ return
983
+
984
if (
985
self.kda_decode_backend != "triton"
986
and self.decode_conv1d_weight is not None
vllm/utils/flashinfer.py
+12
@@ -313,6 +313,9 @@ flashinfer_xqa_batch_decode_with_kv_cache = _lazy_import_wrapper(
313
"flashinfer.decode",
314
"xqa_batch_decode_with_kv_cache",
315
)
316
+flashinfer_packed_fused_kda_decode = _lazy_import_wrapper(
317
+ "flashinfer", "packed_fused_kda_decode"
318
+)
319
flashinfer_recurrent_kda = _lazy_import_wrapper(
320
"flashinfer.kda",
321
"recurrent_kda",
@@ -337,6 +340,15 @@ def has_flashinfer_comm() -> bool:
340
return has_flashinfer() and importlib.util.find_spec("flashinfer.comm") is not None
341
342
343
+@functools.cache
344
+def has_flashinfer_packed_fused_kda_decode() -> bool:
345
+ """Return whether FlashInfer's packed fused KDA decode API is available."""
346
+ if not has_flashinfer():
347
+ return False
348
+ module = _get_submodule("flashinfer")
349
+ return bool(module and callable(getattr(module, "packed_fused_kda_decode", None)))
350
+
351
+
352
@functools.cache
353
def has_flashinfer_nvlink_two_sided() -> bool:
354
"""Return `True` if FlashInfer mnnvl all2all is available."""